From 962c6476da625d168e835f558e8e9936ddcb1d2b Mon Sep 17 00:00:00 2001 From: Priya Ramani Date: Sat, 16 Oct 2021 20:01:19 -0700 Subject: [PATCH] Refactor: move method to func compilation work to compileMethod, add option to specify method name (#66726) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Summary: Pull Request resolved: https://github.com/pytorch/pytorch/pull/66726 Move method to func compilation work to compileMethod Test Plan: Mobilenetv3 compiles and runs successfully ``` (pytorch) ~/fbsource/fbcode/caffe2/fb/nnc └─ $ buck run //caffe2/binaries:aot_model_compiler -- --model mobilenetv3.pt --model_name=pytorch_dev_mobilenetv3 --model_version=v1 --input_dims="1,3,224,224" Downloaded 0/4 artifacts, 0.00 bytes, 100.0% cache miss (for updated rules) Building: finished in 13.2 sec (100%) 18719/18719 jobs, 2/18719 updated Total time: 13.5 sec BUILD SUCCEEDED The compiled llvm assembly code was saved to mobilenetv3.compiled.ll The compiled model was saved to mobilenetv3.compiled.pt ``` Reviewed By: ljk53, IvanKobzarev Differential Revision: D31624342 fbshipit-source-id: 233a6e94ea05ba8d6fc166d2414034c9e58cb076 --- binaries/aot_model_compiler.cc | 33 ++++++++++++++-------- torch/csrc/jit/mobile/nnc/aot_compiler.cpp | 29 ++++++++++--------- 2 files changed, 36 insertions(+), 26 deletions(-) diff --git a/binaries/aot_model_compiler.cc b/binaries/aot_model_compiler.cc index 7ca7f187ae30..eab9938e19e9 100644 --- a/binaries/aot_model_compiler.cc +++ b/binaries/aot_model_compiler.cc @@ -21,6 +21,7 @@ C10_DEFINE_string( "For input float TensorCPUs, specify the dimension using comma " "separated numbers. If multiple inputs needed, use semicolon " "to separate the dimension of different tensors."); +C10_DEFINE_string(method_name, "forward", "The name of the method."); C10_DEFINE_string( output_llvm, "", @@ -71,23 +72,22 @@ c10::Dict createCompileSpec() { "Wrong # of input shapes: ", input_shapes.size()); method_spec.insert("sizes", input_shapes[0]); // TODO: support multiple inputs - compile_spec.insert("forward", method_spec); + compile_spec.insert(FLAGS_method_name, method_spec); return compile_spec; } std::vector getInputSizesForMethod( - const c10::Dict& method_compile_spec, - const std::string& method_name) { - return method_compile_spec.at(method_name) + const c10::Dict& method_compile_spec) { + return method_compile_spec.at(FLAGS_method_name) .toGenericDict() .at("sizes") .toIntVector(); } -std::string getNncKernelId(const std::string& method_name) { +std::string getNncKernelId() { // TODO: calculate the version_token. const std::string version_token = "VERTOKEN"; - return FLAGS_model_name + ":" + FLAGS_model_version + ":" + method_name + + return FLAGS_model_name + ":" + FLAGS_model_version + ":" + FLAGS_method_name + ":" + version_token; } @@ -100,23 +100,31 @@ void writeOutputLlvmAssembly(const std::string& asm_code) { std::ofstream output(output_llvm_file_name); output << asm_code; + std::cout << "The compiled llvm assembly code was saved to " << output_llvm_file_name + << std::endl; } c10::IValue preprocess( const torch::jit::Module& mod, const c10::Dict& method_compile_spec, const torch::jit::BackendDebugHandleGenerator& generate_debug_handles) { - const std::string& method_name = "forward"; - auto method = mod.get_method(method_name); + + std::string output_llvm_file_name = FLAGS_output_llvm; + if (output_llvm_file_name.empty()) { + output_llvm_file_name = + FLAGS_model.substr(0, FLAGS_model.find('.')) + ".compiled.ll"; + } + + auto method = mod.get_method(FLAGS_method_name); auto graph = method.function().graph()->copy(); - auto sizes = getInputSizesForMethod(method_compile_spec, method_name); + auto sizes = getInputSizesForMethod(method_compile_spec); std::string llvm_asm_code; - auto compiled = torch::jit::mobile::nnc::aotCompile(method_name, graph, sizes); + auto compiled = torch::jit::mobile::nnc::aotCompile(FLAGS_method_name, graph, sizes); writeOutputLlvmAssembly(compiled.second); auto func = std::move(compiled.first); - func->set_nnc_kernel_id(getNncKernelId(method_name)); + func->set_nnc_kernel_id(getNncKernelId()); torch::jit::mobile::nnc::CompilationUnit cu; cu.register_function(std::move(func)); @@ -135,6 +143,7 @@ int main(int argc, char** argv) { " --model_name=" " --model_version=" " --input_dims='1,3,224,224'" + " [--method_name=]" " [--output_llvm=]" " [--output_model=]"); @@ -155,7 +164,7 @@ int main(int argc, char** argv) { auto m = torch::jit::load(FLAGS_model); m.eval(); auto frozen_m = torch::jit::freeze_module(m.clone()); - auto graph = frozen_m.get_method("forward").graph(); + auto graph = frozen_m.get_method(FLAGS_method_name).graph(); torch::jit::OptimizeFrozenGraph(graph, true); auto compile_spec = createCompileSpec(); diff --git a/torch/csrc/jit/mobile/nnc/aot_compiler.cpp b/torch/csrc/jit/mobile/nnc/aot_compiler.cpp index de7db04be332..234594c4d827 100644 --- a/torch/csrc/jit/mobile/nnc/aot_compiler.cpp +++ b/torch/csrc/jit/mobile/nnc/aot_compiler.cpp @@ -33,9 +33,18 @@ std::vector getConstSizes(const BufPtr b) { return r; } -void compileFunction( +std::unique_ptr compileMethod( std::shared_ptr kernel, - Function* func) { + const std::string& method_name, + const std::vector& sizes) { + auto func = std::make_unique(); + func->set_name(method_name); + + InputSpec input; + input.sizes_ = sizes; + input.dtype_ = c10::ScalarType::Float; + func->set_input_specs({input}); + std::vector parameters; auto const_descriptors = kernel->getConstantDescriptors(); @@ -64,14 +73,16 @@ void compileFunction( out_spec.push_back(output); } func->set_output_specs(out_spec); + + return func; } std::pair, const std::string> aotCompile( const std::string& method_name, std::shared_ptr& g, const std::vector& sizes) { - auto g2 = g->copy(); GRAPH_DEBUG("Input sizes ", sizes); + GRAPH_DEBUG("Method name ", method_name); RemoveTensorMutation(g); EliminateDeadCode(g->block()); @@ -91,17 +102,7 @@ std::pair, const std::string> aotCompile( std::make_shared(g); const std::string compiled_assembly = kernel->getCodeText(); - g = g2; - - auto func = std::make_unique(); - func->set_name(method_name); - - InputSpec input; - input.sizes_ = sizes; - input.dtype_ = c10::ScalarType::Float; - func->set_input_specs({input}); - - compileFunction(kernel, func.get()); + auto func = compileMethod(kernel, method_name, sizes); return std::make_pair(std::move(func), compiled_assembly); }