xref: /aosp_15_r20/external/pytorch/torch/csrc/jit/mobile/nnc/aot_compiler.h (revision da0073e96a02ea20f0ac840b70461e3646d07c45)
1 #pragma once
2 
3 #include <torch/csrc/Export.h>
4 #include <torch/csrc/jit/ir/ir.h>
5 #include <torch/csrc/jit/mobile/nnc/context.h>
6 
7 namespace torch {
8 namespace jit {
9 namespace mobile {
10 namespace nnc {
11 
12 // Performs Ahead Of Time compilation of a given method in a model
13 // returning the compiled function and LLVM assembly code
14 TORCH_API std::pair<std::unique_ptr<Function>, const std::string> aotCompile(
15     const std::string& method_name,
16     std::shared_ptr<Graph>& subgraph,
17     const std::vector<std::vector<int64_t>>& sizes,
18     const std::vector<at::ScalarType>& types,
19     const std::string& kernel_func_name = "func");
20 
21 } // namespace nnc
22 } // namespace mobile
23 } // namespace jit
24 } // namespace torch
25