mirror of
https://github.com/zebrajr/pytorch.git
synced 2025-12-06 12:20:52 +01:00
3551bd31be
Summary: Pull Request resolved: https://github.com/pytorch/pytorch/pull/54462 Unclean files during sync - Sat Mar 20 04:00:02 PDT 2021 Unclean files during sync - Sun Mar 21 04:00:01 PDT 2021 ghstack-source-id: 124585992 Test Plan: ``` buck run xplat/caffe2/fb/test/delegate:interpreter_test -- --model_file_path=/path/to/mobile_model.ptl ``` Reviewed By: raziel Differential Revision: D27232309 fbshipit-source-id: 8504a3185339d73bfa6e924485c4745acf269cec
60 lines
1.9 KiB
C++
60 lines
1.9 KiB
C++
#include <torch/csrc/jit/backends/backend.h>
|
|
#include <torch/csrc/jit/backends/backend_preprocess.h>
|
|
|
|
namespace torch {
|
|
namespace jit {
|
|
namespace {
|
|
// For this backend, the actual compilation happens in preprocess function AOT.
|
|
// Put here for demonstration of backend
|
|
// as a whole piece. It's used when compilation is required. A dummy function
|
|
// can be passed when there's no usage of compilation in runtime backend lib.
|
|
c10::IValue preprocess(
|
|
const Module& mod,
|
|
const c10::Dict<IValue, IValue>& method_compile_spec) {
|
|
// The output of this process would produce a dictionary
|
|
// Key: method name.
|
|
// Val: compiled blob (represented by a string).
|
|
c10::Dict<IValue, IValue> compiled(StringType::get(), StringType::get());
|
|
for (const auto& method : mod.get_methods()) {
|
|
const auto graph = method.function().graph()->copy();
|
|
auto key = method.name();
|
|
std::stringstream ss;
|
|
for (const auto& node : graph->nodes()) {
|
|
switch (node->kind()) {
|
|
case prim::Constant:
|
|
ss << node->kind().toDisplayString() << "#"
|
|
<< toIValue(node->output()).value();
|
|
break;
|
|
case aten::add:
|
|
ss << node->kind().toQualString();
|
|
break;
|
|
case aten::sub:
|
|
ss << node->kind().toQualString();
|
|
break;
|
|
default:
|
|
TORCH_CHECK(
|
|
false,
|
|
"The node of ",
|
|
node->kind().toQualString(),
|
|
" is not supported in this compiler. Source code: ",
|
|
node->sourceRange().str());
|
|
break;
|
|
}
|
|
ss << ",";
|
|
}
|
|
std::string blob = ss.str();
|
|
if (!blob.empty()) {
|
|
blob.pop_back();
|
|
}
|
|
compiled.insert(method.name(), blob);
|
|
}
|
|
return compiled;
|
|
}
|
|
|
|
constexpr auto backend_name = "backend_with_compiler_demo";
|
|
static auto pre_reg = backend_preprocess_register(backend_name, preprocess);
|
|
} // namespace
|
|
|
|
} // namespace jit
|
|
} // namespace torch
|