Refactor: move method to func compilation work to compileMethod, add option to specify method name (#66726)

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

@@ -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<c10::IValue, c10::IValue> 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<int64_t> getInputSizesForMethod(
-    const c10::Dict<c10::IValue, c10::IValue>& method_compile_spec,
-    const std::string& method_name) {
-  return method_compile_spec.at(method_name)
+    const c10::Dict<c10::IValue, c10::IValue>& 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<c10::IValue, c10::IValue>& 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 name>"
       " --model_version=<model version>"
       " --input_dims='1,3,224,224'"
+      " [--method_name=<mehhod name>]"
       " [--output_llvm=<llvm assembly output file path>]"
       " [--output_model=<output model file path>]");
 
@@ -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();

torch/csrc/jit/mobile/nnc/aot_compiler.cpp

@@ -33,9 +33,18 @@ std::vector<int64_t> getConstSizes(const BufPtr b) {
   return r;
 }
 
-void compileFunction(
+std::unique_ptr<Function> compileMethod(
     std::shared_ptr<tensorexpr::TensorExprKernel> kernel,
-    Function* func) {
+    const std::string& method_name,
+    const std::vector<int64_t>& sizes) {
+  auto func = std::make_unique<Function>();
+  func->set_name(method_name);
+
+  InputSpec input;
+  input.sizes_ = sizes;
+  input.dtype_ = c10::ScalarType::Float;
+  func->set_input_specs({input});
+
   std::vector<at::Tensor> 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<std::unique_ptr<Function>, const std::string> aotCompile(
     const std::string& method_name,
     std::shared_ptr<Graph>& g,
     const std::vector<int64_t>& 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<std::unique_ptr<Function>, const std::string> aotCompile(
       std::make_shared<tensorexpr::TensorExprKernel>(g);
   const std::string compiled_assembly = kernel->getCodeText();
 
-  g = g2;
-
-  auto func = std::make_unique<Function>();
-  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);
 }