mirror of
https://github.com/zebrajr/pytorch.git
synced 2025-12-06 12:20:52 +01:00
b5ae8e69a7
Summary: Pull Request resolved: https://github.com/pytorch/pytorch/pull/52870 Add the missing parts to support to_backend modules by lite interpreter. 1. Add ISINSTANCE instruction support, which is used in to_backend for output type check. 2. Bypass lite interpreter's type parser by checking the qualified name. If it starts with "torch.jit", use the same type resolver as nn module (starting with "__torch__"). Tests Mobile module is serialized and loaded in ```BackendTest.TestCompiler```. The results are compared to those from original torchscript module. Test Plan: Imported from OSS Reviewed By: raziel Differential Revision: D26715351 Pulled By: iseeyuan fbshipit-source-id: ad9d74ee81c6aa692ab9e5dd7a9003bae5d4f01f
91 lines
2.0 KiB
C++
91 lines
2.0 KiB
C++
#include <torch/csrc/jit/runtime/instruction.h>
|
|
#include <cstring>
|
|
#include <iostream>
|
|
|
|
namespace torch {
|
|
namespace jit {
|
|
std::ostream& operator<<(std::ostream& out, OpCode op) {
|
|
switch (op) {
|
|
#define OP_STRING(x, _) \
|
|
case x: \
|
|
return out << #x;
|
|
FORALL_OPCODES(OP_STRING)
|
|
#undef OP_STRING
|
|
}
|
|
return out;
|
|
}
|
|
|
|
char const* toString(OpCode op) {
|
|
switch (op) {
|
|
#define OP_STRING(x, _) \
|
|
case x: \
|
|
return #x;
|
|
FORALL_OPCODES(OP_STRING)
|
|
#undef OP_STRING
|
|
}
|
|
return nullptr;
|
|
}
|
|
|
|
const char* OpInfo(OpCode op) {
|
|
switch (op) {
|
|
#define OP_INFO(x, info) \
|
|
case x: \
|
|
return info;
|
|
FORALL_OPCODES(OP_INFO)
|
|
#undef OP_INFO
|
|
}
|
|
return nullptr;
|
|
}
|
|
|
|
static constexpr size_t instruction_size = 8;
|
|
static_assert(
|
|
sizeof(Instruction) == instruction_size,
|
|
"Instructions should be 8 bytes");
|
|
std::ostream& operator<<(std::ostream& out, Instruction inst) {
|
|
// TODO: use op info to print out the op in a more user-friendly way
|
|
int nargs = std::strlen(OpInfo(inst.op));
|
|
out << inst.op;
|
|
if (nargs > 0) {
|
|
out << " " << inst.X;
|
|
}
|
|
if (nargs > 1) {
|
|
out << " " << inst.N;
|
|
}
|
|
return out;
|
|
}
|
|
|
|
static constexpr char* strOpCode[] = {
|
|
#define STR_OP(x, _) #x,
|
|
FORALL_OPCODES(STR_OP)
|
|
#undef STR_OP
|
|
};
|
|
|
|
OpCode parseOpCode(const char* str) {
|
|
const int n = sizeof(strOpCode) / sizeof(strOpCode[0]);
|
|
for (int i = 0; i < n; ++i) {
|
|
if (strcmp(strOpCode[i], str) == 0)
|
|
return (OpCode)i;
|
|
}
|
|
return OP;
|
|
}
|
|
|
|
bool isOpSupportedInMobile(OpCode op) {
|
|
// clang-format off
|
|
static constexpr OpCode supported_ops_in_mobile[] {
|
|
OP, OPN, LOAD, MOVE, STOREN, STORE, DROP, DROPR, LOADC, JF, JMP, LOOP,
|
|
RET, GET_ATTR, SET_ATTR, LIST_CONSTRUCT, TUPLE_CONSTRUCT, WARN,
|
|
INTERFACE_CALL, LIST_UNPACK, TUPLE_SLICE, DICT_CONSTRUCT,
|
|
NAMED_TUPLE_CONSTRUCT, CREATE_OBJECT, ISINSTANCE
|
|
};
|
|
// clang-format on
|
|
|
|
for (auto sop : supported_ops_in_mobile) {
|
|
if (op == sop)
|
|
return true;
|
|
}
|
|
return false;
|
|
}
|
|
|
|
} // namespace jit
|
|
} // namespace torch
|