Star882
1390
代码
代码
Issues1739
Pull Requests2328
Actions
项目讨论
Wiki
分析
Star882
1390
已合并
add aot inductor for v2.8.0 #25446
kkjocker创建于 2025年9月29日
add aot inductor for v2.8.0 #25446
已合并
kkjocker创建于 2025年9月29日
从kkjocker/pytorch_280aoti_280合入到Ascend/pytorchv2.8.0
共 30 个文件变更+4630-21
MCMakeLists.txt+5-2
@@ -203,6 +203,7 @@ endif()
203 203 
204include_directories(${PROJECT_SOURCE_DIR})204include_directories(${PROJECT_SOURCE_DIR})
205include_directories(${PROJECT_SOURCE_DIR}/torch_npu/csrc/aten)205include_directories(${PROJECT_SOURCE_DIR}/torch_npu/csrc/aten)
206+include_directories(${PROJECT_SOURCE_DIR}/torch_npu/csrc/inductor)
206include_directories(${PROJECT_SOURCE_DIR}/third_party/hccl/inc)207include_directories(${PROJECT_SOURCE_DIR}/third_party/hccl/inc)
207include_directories(${PROJECT_SOURCE_DIR}/third_party/acl/inc)208include_directories(${PROJECT_SOURCE_DIR}/third_party/acl/inc)
208include_directories(${PROJECT_SOURCE_DIR}/third_party/Tensorpipe)209include_directories(${PROJECT_SOURCE_DIR}/third_party/Tensorpipe)
@@ -229,6 +230,7 @@ set(ATEN_SRCS)
229set(CORE_SRCS)230set(CORE_SRCS)
230set(FRAMEWORK_SRCS)231set(FRAMEWORK_SRCS)
231set(LOGGING_SRCS)232set(LOGGING_SRCS)
233+set(INDUCTOR_SRCS)
232 234 
233if (NOT DEFINED BUILD_LIBTORCH)235if (NOT DEFINED BUILD_LIBTORCH)
234 set(DIST_SRCS)236 set(DIST_SRCS)
@@ -251,6 +253,7 @@ add_subdirectory(${TORCHNPU_ROOT}/framework)
251add_subdirectory(${TORCHNPU_ROOT}/flopcount)253add_subdirectory(${TORCHNPU_ROOT}/flopcount)
252add_subdirectory(${TORCHNPU_ROOT}/logging)254add_subdirectory(${TORCHNPU_ROOT}/logging)
253add_subdirectory(${TORCHNPU_ROOT}/custom_dtype)255add_subdirectory(${TORCHNPU_ROOT}/custom_dtype)
256+add_subdirectory(${TORCHNPU_ROOT}/inductor)
254 257 
255if (NOT DEFINED BUILD_LIBTORCH)258if (NOT DEFINED BUILD_LIBTORCH)
256 add_subdirectory(${TORCHNPU_ROOT}/distributed)259 add_subdirectory(${TORCHNPU_ROOT}/distributed)
@@ -286,10 +289,10 @@ if (DEFINED BUILD_TENSORPIPE)
286endif()289endif()
287 290 
288if (DEFINED BUILD_LIBTORCH)291if (DEFINED BUILD_LIBTORCH)
289- set(CPP_SRCS ${ATEN_SRCS} ${CORE_SRCS} ${OPS_PLUGIN_SRCS} ${FLOP_SRCS} ${CUS_DTYPE_SRCS} ${FRAMEWORK_SRCS} ${LOGGING_SRCS} ${NPU_CPP_LIBS_SRCS})292+ set(CPP_SRCS ${ATEN_SRCS} ${INDUCTOR_SRCS} ${CORE_SRCS} ${OPS_PLUGIN_SRCS} ${FLOP_SRCS} ${CUS_DTYPE_SRCS} ${FRAMEWORK_SRCS} ${LOGGING_SRCS} ${NPU_CPP_LIBS_SRCS})
290else()293else()
291# Compile code with pybind11294# Compile code with pybind11
292- set(CPP_SRCS ${ATEN_SRCS} ${CORE_SRCS} ${OPS_PLUGIN_SRCS} ${DIST_SRCS} ${FLOP_SRCS} ${CUS_DTYPE_SRCS} ${LOGGING_SRCS} ${FRAMEWORK_SRCS} ${NPU_SRCS} ${PROF_SRCS} ${IPC_SRCS} ${UTILS_SRCS} ${SAN_SRCS} ${AFD_SRCS})295+ set(CPP_SRCS ${ATEN_SRCS} ${INDUCTOR_SRCS} ${CORE_SRCS} ${OPS_PLUGIN_SRCS} ${DIST_SRCS} ${FLOP_SRCS} ${CUS_DTYPE_SRCS} ${LOGGING_SRCS} ${FRAMEWORK_SRCS} ${NPU_SRCS} ${PROF_SRCS} ${IPC_SRCS} ${UTILS_SRCS} ${SAN_SRCS} ${AFD_SRCS})
293endif()296endif()
294 297 
295add_library(${PLUGIN_NAME} SHARED ${CPP_SRCS})298add_library(${PLUGIN_NAME} SHARED ${CPP_SRCS})
Mbuild_libtorch_npu.py+2-1
@@ -225,7 +225,8 @@ def copy_hpp():
225 "torch_npu/csrc/framework/*.h",225 "torch_npu/csrc/framework/*.h",
226 "torch_npu/csrc/framework/*/*.h",226 "torch_npu/csrc/framework/*/*.h",
227 "torch_npu/csrc/framework/*/*/*.h",227 "torch_npu/csrc/framework/*/*/*.h",
228- "torch_npu/csrc/libs/*.h"228+ "torch_npu/csrc/libs/*.h",
229+ "torch_npu/csrc/inductor/**/*.h",
229 ]230 ]
230 glob_header_files = []231 glob_header_files = []
231 for regex_pattern in header_files:232 for regex_pattern in header_files:
Msetup.py+22-0
@@ -481,6 +481,28 @@ def get_src_py_and_dst():
481 os.path.relpath(src, os.path.join(BASE_DIR, "patch/include")))481 os.path.relpath(src, os.path.join(BASE_DIR, "patch/include")))
482 os.makedirs(os.path.dirname(dst), exist_ok=True)482 os.makedirs(os.path.dirname(dst), exist_ok=True)
483 ret.append((src, dst))483 ret.append((src, dst))
484+ 
485+ aot_inductor_files = [
486+ # Follow torch v2.6.0.
487+ # These aoti_runtime/*.cpp don't compile to libtorch_npu,
488+ # but act like header files when generate cppwrapper in aot-inductor.
489+ "torch_npu/_inductor/codegen/aoti_runtime/*.cpp"
490+ ]
491+ glob_aoti_files = []
492+ for regex_pattern in aot_inductor_files:
493+ glob_aoti_files += glob.glob(
494+ os.path.join(BASE_DIR, regex_pattern), recursive=True
495+ )
496+ 
497+ for src in glob_aoti_files:
498+ # Dst: torch_npu/_inductor/codegen/aoti_runtime/*.cpp
499+ dst = os.path.join(
500+ os.path.join(BASE_DIR, "build/packages/torch_npu/"),
501+ os.path.relpath(src, os.path.join(BASE_DIR, "torch_npu")),
502+ )
503+ os.makedirs(os.path.dirname(dst), exist_ok=True)
504+ ret.append((src, dst))
505+ 
484 return ret506 return ret
485 507 
486 508 
Mtorch_npu/_inductor/codegen/aoti_runtime/interface.cpp+66-18
@@ -1,3 +1,4 @@
1+// Licensed under the BSD 3-Clause License (the "License");
1// Definition of AOTI runtime interface functions2// Definition of AOTI runtime interface functions
2 3 
3#include <torch_npu/csrc/inductor/aoti_runtime/interface.h>4#include <torch_npu/csrc/inductor/aoti_runtime/interface.h>
@@ -12,10 +13,10 @@
12 try { \13 try { \
13 __VA_ARGS__ \14 __VA_ARGS__ \
14 } catch (const std::exception& e) { \15 } catch (const std::exception& e) { \
15- std::cerr << "Error: " << e.what() << std::endl; \16+ std::cerr << "Error: " << e.what() << '\n'; \
16 return AOTI_RUNTIME_FAILURE; \17 return AOTI_RUNTIME_FAILURE; \
17 } catch (...) { \18 } catch (...) { \
18- std::cerr << "Unknown exception occurred." << std::endl; \19+ std::cerr << "Unknown exception occurred.\n"; \
19 return AOTI_RUNTIME_FAILURE; \20 return AOTI_RUNTIME_FAILURE; \
20 } \21 } \
21 return AOTI_RUNTIME_SUCCESS;22 return AOTI_RUNTIME_SUCCESS;
@@ -34,9 +35,13 @@
34// A RAII, thread local (!) guard that enables or disables grad mode upon35// A RAII, thread local (!) guard that enables or disables grad mode upon
35// construction, and sets it back to the original value upon destruction.36// construction, and sets it back to the original value upon destruction.
36struct AOTINoGradGuard {37struct AOTINoGradGuard {
37- AOTINoGradGuard() : prev_mode(aoti_torch_grad_mode_is_enabled()) { aoti_torch_grad_mode_set_enabled(false); }38+ AOTINoGradGuard() { aoti_torch_grad_mode_set_enabled(false); }
39+ AOTINoGradGuard(const AOTINoGradGuard&) = delete;
40+ AOTINoGradGuard(AOTINoGradGuard&&) noexcept = delete;
38 ~AOTINoGradGuard() { aoti_torch_grad_mode_set_enabled(prev_mode); }41 ~AOTINoGradGuard() { aoti_torch_grad_mode_set_enabled(prev_mode); }
39- bool prev_mode;42+ AOTINoGradGuard& operator=(const AOTINoGradGuard&) = delete;
43+ AOTINoGradGuard& operator=(AOTINoGradGuard&&) noexcept = delete;
44+ bool prev_mode{aoti_torch_grad_mode_is_enabled()};
40};45};
41 46 
42extern "C" {47extern "C" {
@@ -51,10 +56,9 @@ AOTIRuntimeError AOTInductorModelContainerCreateWithDevice(AOTInductorModelConta
51 const char* cubin_dir)56 const char* cubin_dir)
52{57{
53 if (num_models == 0) {58 if (num_models == 0) {
54- std::cerr << "Error: num_models must be positive, but got 0" << std::endl;59+ std::cerr << "Error: num_models must be positive, but got 0\n";
55 return AOTI_RUNTIME_FAILURE;60 return AOTI_RUNTIME_FAILURE;
56 }61 }
57- 
58 CONVERT_EXCEPTION_TO_ERROR_CODE({62 CONVERT_EXCEPTION_TO_ERROR_CODE({
59 std::optional<std::string> cubin_dir_opt;63 std::optional<std::string> cubin_dir_opt;
60 if (cubin_dir != nullptr) {64 if (cubin_dir != nullptr) {
@@ -75,9 +79,13 @@ AOTIRuntimeError AOTInductorModelContainerDelete(AOTInductorModelContainerHandle
75}79}
76 80 
77AOTIRuntimeError AOTInductorModelContainerRun(AOTInductorModelContainerHandle container_handle,81AOTIRuntimeError AOTInductorModelContainerRun(AOTInductorModelContainerHandle container_handle,
78- AtenTensorHandle* input_handles, size_t num_inputs,82+ AtenTensorHandle* input_handles, // array of input AtenTensorHandle; handles
79- AtenTensorHandle* output_handles, size_t num_outputs,83+ // are stolen; the array itself is borrowed
80- AOTInductorStreamHandle stream_handle,84+ size_t num_inputs,
85+ AtenTensorHandle* output_handles, // array for writing output AtenTensorHandle; handles
86+ // will be stolen by the caller; the array itself is
87+ // borrowed
88+ size_t num_outputs, AOTInductorStreamHandle stream_handle,
81 AOTIProxyExecutorHandle proxy_executor_handle)89 AOTIProxyExecutorHandle proxy_executor_handle)
82{90{
83 auto* container = reinterpret_cast<torch::aot_inductor::AOTInductorModelContainer*>(container_handle);91 auto* container = reinterpret_cast<torch::aot_inductor::AOTInductorModelContainer*>(container_handle);
@@ -91,11 +99,15 @@ AOTIRuntimeError AOTInductorModelContainerRun(AOTInductorModelContainerHandle co
91 })99 })
92}100}
93 101 
94-AOTIRuntimeError AOTInductorModelContainerRunSingleThreaded(AOTInductorModelContainerHandle container_handle,102+AOTIRuntimeError AOTInductorModelContainerRunSingleThreaded(
95- AtenTensorHandle* input_handles, size_t num_inputs,103+ AOTInductorModelContainerHandle container_handle,
96- AtenTensorHandle* output_handles, size_t num_outputs,104+ AtenTensorHandle* input_handles, // array of input AtenTensorHandle; handles
97- AOTInductorStreamHandle stream_handle,105+ // are stolen; the array itself is borrowed
98- AOTIProxyExecutorHandle proxy_executor_handle)106+ size_t num_inputs,
107+ AtenTensorHandle* output_handles, // array for writing output AtenTensorHandle; handles
108+ // will be stolen by the caller; the array itself is
109+ // borrowed
110+ size_t num_outputs, AOTInductorStreamHandle stream_handle, AOTIProxyExecutorHandle proxy_executor_handle)
99{111{
100 auto* container = reinterpret_cast<torch::aot_inductor::AOTInductorModelContainer*>(container_handle);112 auto* container = reinterpret_cast<torch::aot_inductor::AOTInductorModelContainer*>(container_handle);
101 AOTI_VECTOR_SIZE_CHECK(num_inputs, container->num_inputs(), "inputs");113 AOTI_VECTOR_SIZE_CHECK(num_inputs, container->num_inputs(), "inputs");
@@ -150,6 +162,38 @@ AOTIRuntimeError AOTInductorModelContainerGetConstantDtype(AOTInductorModelConta
150 CONVERT_EXCEPTION_TO_ERROR_CODE({ *dtype = container->constant_dtype(idx); })162 CONVERT_EXCEPTION_TO_ERROR_CODE({ *dtype = container->constant_dtype(idx); })
151}163}
152 164 
165+AOTIRuntimeError AOTInductorModelContainerGetConstantDataSize(AOTInductorModelContainerHandle container_handle,
166+ size_t idx, size_t* data_size)
167+{
168+ auto* container = reinterpret_cast<torch::aot_inductor::AOTInductorModelContainer*>(container_handle);
169+ CONVERT_EXCEPTION_TO_ERROR_CODE({ *data_size = container->constant_data_size(idx); })
170+}
171+ 
172+AOTIRuntimeError AOTInductorModelContainerExtractConstantsMap(AOTInductorModelContainerHandle container_handle,
173+ AOTInductorConstantMapHandle constant_map_handle,
174+ bool use_inactive)
175+{
176+ auto* container = reinterpret_cast<torch::aot_inductor::AOTInductorModelContainer*>(container_handle);
177+ auto constants_map = reinterpret_cast<std::unordered_map<std::string, AtenTensorHandle>*>(constant_map_handle);
178+ CONVERT_EXCEPTION_TO_ERROR_CODE({
179+ const auto ret = container->extract_constants_map(use_inactive);
180+ for (const auto& pair : ret) {
181+ constants_map->emplace(pair.first, pair.second);
182+ }
183+ })
184+}
185+ 
186+AOTIRuntimeError AOTInductorModelContainerUpdateUserManagedConstantBuffer(AOTInductorModelContainerHandle container_handle,
187+ AOTInductorConstantMapHandle constant_map_handle,
188+ bool use_inactive, bool validate_full_update)
189+{
190+ auto* container = reinterpret_cast<torch::aot_inductor::AOTInductorModelContainer*>(container_handle);
191+ auto input_map = reinterpret_cast<std::unordered_map<std::string, AtenTensorHandle>*>(constant_map_handle);
192+ CONVERT_EXCEPTION_TO_ERROR_CODE({
193+ container->update_constant_buffer(*input_map, use_inactive, validate_full_update, true);
194+ })
195+}
196+ 
153AOTIRuntimeError AOTInductorModelContainerUpdateConstantBuffer(AOTInductorModelContainerHandle container_handle,197AOTIRuntimeError AOTInductorModelContainerUpdateConstantBuffer(AOTInductorModelContainerHandle container_handle,
154 AOTInductorConstantMapHandle constant_map_handle,198 AOTInductorConstantMapHandle constant_map_handle,
155 bool use_inactive, bool validate_full_update)199 bool use_inactive, bool validate_full_update)
@@ -166,6 +210,12 @@ AOTIRuntimeError AOTInductorModelContainerUpdateInactiveConstantBuffer(AOTInduct
166 return AOTInductorModelContainerUpdateConstantBuffer(container_handle, constant_map_handle, true, true);210 return AOTInductorModelContainerUpdateConstantBuffer(container_handle, constant_map_handle, true, true);
167}211}
168 212 
213+AOTIRuntimeError AOTInductorModelContainerFreeInactiveConstantBuffer(AOTInductorModelContainerHandle container_handle)
214+{
215+ auto* container = reinterpret_cast<torch::aot_inductor::AOTInductorModelContainer*>(container_handle);
216+ CONVERT_EXCEPTION_TO_ERROR_CODE({ container->free_inactive_constant_buffer(); })
217+}
218+ 
169AOTIRuntimeError AOTInductorModelContainerRunConstantFolding(AOTInductorModelContainerHandle container_handle,219AOTIRuntimeError AOTInductorModelContainerRunConstantFolding(AOTInductorModelContainerHandle container_handle,
170 bool use_inactive, AOTInductorStreamHandle stream_handle,220 bool use_inactive, AOTInductorStreamHandle stream_handle,
171 AOTIProxyExecutorHandle proxy_executor_handle)221 AOTIProxyExecutorHandle proxy_executor_handle)
@@ -259,8 +309,7 @@ AOTIRuntimeError AOTInductorModelRun(AOTInductorModelHandle model_handle, AtenTe
259 309 
260AOTIRuntimeError AOTInductorModelDelete(AOTInductorModelHandle model_handle)310AOTIRuntimeError AOTInductorModelDelete(AOTInductorModelHandle model_handle)
261{311{
262- CONVERT_EXCEPTION_TO_ERROR_CODE(312+ CONVERT_EXCEPTION_TO_ERROR_CODE({
263- {
264 auto model = reinterpret_cast<torch::aot_inductor::AOTInductorModel*>(model_handle);313 auto model = reinterpret_cast<torch::aot_inductor::AOTInductorModel*>(model_handle);
265 delete model;314 delete model;
266 })315 })
@@ -269,8 +318,7 @@ AOTIRuntimeError AOTInductorModelDelete(AOTInductorModelHandle model_handle)
269AOTIRuntimeError AOTInductorModelGetNumOutputs(AOTInductorModelHandle model_handle,318AOTIRuntimeError AOTInductorModelGetNumOutputs(AOTInductorModelHandle model_handle,
270 size_t* ret_num_outputs)319 size_t* ret_num_outputs)
271{320{
272- CONVERT_EXCEPTION_TO_ERROR_CODE(321+ CONVERT_EXCEPTION_TO_ERROR_CODE({
273- {
274 auto model = reinterpret_cast<torch::aot_inductor::AOTInductorModel*>(model_handle);322 auto model = reinterpret_cast<torch::aot_inductor::AOTInductorModel*>(model_handle);
275 *ret_num_outputs = model->num_outputs();323 *ret_num_outputs = model->num_outputs();
276 })324 })
Atorch_npu/csrc/inductor/CMakeLists.txt+9-0
@@ -0,0 +1,9 @@
1+FILE(GLOB _INDUCTOR_SRCS
2+ *.cpp
3+ aoti_runner/*.cpp
4+ aoti_torch/*.cpp)
5+ 
6+LIST(APPEND INDUCTOR_SRCS ${_INDUCTOR_SRCS})
7+ 
8+# Pass to parent
9+set(INDUCTOR_SRCS ${INDUCTOR_SRCS} PARENT_SCOPE)
Atorch_npu/csrc/inductor/aoti_package/model_package_loader.h+40-0
@@ -0,0 +1,40 @@
1+#if !defined(C10_MOBILE) && !defined(ANDROID)
2+#pragma once
3+ 
4+#include <ATen/Tensor.h>
5+#include <torch_npu/csrc/inductor/aoti_runner/model_container_runner.h>
6+ 
7+namespace torch::inductor {
8+class TORCH_API AOTIModelPackageLoader {
9+public:
10+ AOTIModelPackageLoader(const std::string& model_package_path, const std::string& model_name = "model",
11+ const bool run_single_threaded = false, const size_t num_runners = 1,
12+ const c10::DeviceIndex device_index = -1);
13+ ~AOTIModelPackageLoader();
14+ 
15+ AOTIModelContainerRunner* get_runner();
16+ std::unordered_map<std::string, std::string> get_metadata();
17+ 
18+ std::vector<at::Tensor> run(const std::vector<at::Tensor>& inputs, void* stream_handle = nullptr);
19+ 
20+ // boxed_run will steal the ownership of the input tensors
21+ std::vector<at::Tensor> boxed_run(std::vector<at::Tensor>&& inputs, void* stream_handle = nullptr);
22+ 
23+ std::vector<std::string> get_call_spec();
24+ void load_constants(std::unordered_map<std::string, at::Tensor>& constants_map, bool use_inactive,
25+ bool check_full_update, bool user_managed = false);
26+ std::vector<std::string> get_constant_fqns();
27+ 
28+ void update_constant_buffer(std::unordered_map<std::string, at::Tensor>& tensor_map, bool use_inactive,
29+ bool validate_full_updates, bool user_managed = false);
30+ 
31+private:
32+ std::string temp_dir_;
33+ std::unique_ptr<AOTIModelContainerRunner> runner_;
34+ std::unordered_map<std::string, std::string> metadata_;
35+ 
36+ void load_metadata(const std::string& cpp_filename);
37+};
38+ 
39+} // namespace torch::inductor
40+#endif
Atorch_npu/csrc/inductor/aoti_package/pybind.h+7-0
@@ -0,0 +1,7 @@
1+#include <torch/csrc/python_headers.h>
2+ 
3+namespace torch::inductor {
4+ 
5+void initAOTIPackageBindings(PyObject* module);
6+ 
7+} // namespace torch::inductor
Atorch_npu/csrc/inductor/aoti_runner/model_container_runner.h+101-0
@@ -0,0 +1,101 @@
1+// Licensed under the BSD 3-Clause License (the "License");
2+#if !defined(C10_MOBILE) && !defined(ANDROID)
3+#pragma once
4+ 
5+#include <ATen/Tensor.h>
6+#include <torch_npu/csrc/inductor/aoti_runtime/interface.h>
7+#include <torch_npu/csrc/inductor/aoti_torch/proxy_executor.h>
8+ 
9+// Forward declare DynamicLibrary
10+namespace at {
11+struct DynamicLibrary;
12+}
13+ 
14+namespace torch::inductor {
15+using TensorConstantMap = std::unordered_map<std::string, at::Tensor*>;
16+ 
17+class TORCH_API AOTIModelContainerRunner {
18+public:
19+ AOTIModelContainerRunner() = delete;
20+ AOTIModelContainerRunner(const AOTIModelContainerRunner& other) = delete;
21+ AOTIModelContainerRunner(AOTIModelContainerRunner&& other) = delete;
22+ AOTIModelContainerRunner& operator=(const AOTIModelContainerRunner& other) = delete;
23+ AOTIModelContainerRunner& operator=(AOTIModelContainerRunner&& other) = delete;
24+ virtual ~AOTIModelContainerRunner();
25+ 
26+ std::vector<at::Tensor> run(const std::vector<at::Tensor>& inputs, void* stream_handle = nullptr);
27+ 
28+ // boxed_run will steal the ownership of the input tensors
29+ std::vector<at::Tensor> boxed_run(std::vector<at::Tensor>&& inputs, void* stream_handle = nullptr);
30+ 
31+ std::unordered_map<std::string, std::string> getConstantNamesToOriginalFQNs() const;
32+ std::unordered_map<std::string, int32_t> getConstantNamesToDtypes() const;
33+ 
34+ const std::unordered_map<std::string, at::Tensor> extract_constants_map(bool use_inactive) const;
35+ void update_inactive_constant_buffer(const TensorConstantMap& const_map);
36+ void update_constant_buffer(std::unordered_map<std::string, at::Tensor>& tensor_map, bool use_inactive,
37+ bool validate_full_updates, bool user_managed = false);
38+ void update_constant_buffer(const TensorConstantMap& const_map, bool use_inactive, bool validate_full_updates,
39+ bool user_managed = false);
40+ void run_const_fold(bool use_inactive, AOTInductorStreamHandle cuda_stream_handle = nullptr);
41+ void swap_constant_buffer();
42+ void free_inactive_constant_buffer();
43+ 
44+ std::vector<std::string> get_call_spec();
45+ 
46+protected:
47+ AOTIModelContainerRunner(const std::string& model_so_path, size_t num_models, const std::string& device_str,
48+ const std::string& cubin_dir, const bool run_single_threaded);
49+ 
50+ virtual std::vector<at::Tensor> run_impl(std::vector<AtenTensorHandle>& input_handles, void* stream_handle);
51+ 
52+ std::unique_ptr<at::DynamicLibrary> model_so_;
53+ decltype(&AOTInductorModelContainerCreateWithDevice) create_func_ { nullptr };
54+ decltype(&AOTInductorModelContainerDelete) delete_func_ { nullptr };
55+ decltype(&AOTInductorModelContainerGetNumOutputs) get_num_outputs_func_ { nullptr };
56+ decltype(&AOTInductorModelContainerRun) run_func_ { nullptr };
57+ decltype(&AOTInductorModelContainerGetNumConstants) get_num_constants_func_ { nullptr };
58+ decltype(&AOTInductorModelContainerGetConstantName) get_constant_name_func_ { nullptr };
59+ decltype(&AOTInductorModelContainerGetConstantOriginalFQN) get_constant_original_fqn_func_ { nullptr };
60+ decltype(&AOTInductorModelContainerGetConstantDtype) get_constant_dtype_func_ { nullptr };
61+ decltype(&AOTInductorModelContainerExtractConstantsMap) extract_constants_map_func_ { nullptr };
62+ decltype(&AOTInductorModelContainerUpdateUserManagedConstantBuffer) update_user_managed_constant_buffer_func_ {
63+ nullptr
64+ };
65+ decltype(&AOTInductorModelContainerUpdateConstantBuffer) update_constant_buffer_func_ { nullptr };
66+ decltype(&AOTInductorModelContainerUpdateInactiveConstantBuffer) update_inactive_constant_buffer_func_ { nullptr };
67+ decltype(&AOTInductorModelContainerRunConstantFolding) run_const_fold_func_ { nullptr };
68+ decltype(&AOTInductorModelContainerSwapConstantBuffer) swap_constant_buffer_func_ { nullptr };
69+ decltype(&AOTInductorModelContainerFreeInactiveConstantBuffer) free_inactive_constant_buffer_func_ { nullptr };
70+ decltype(&AOTInductorModelContainerGetCallSpec) get_call_spec_func_ { nullptr };
71+ 
72+ AOTInductorModelContainerHandle container_handle_ = nullptr;
73+ 
74+ AOTIProxyExecutorHandle proxy_executor_handle_;
75+ 
76+private:
77+ std::unique_ptr<torch::aot_inductor::ProxyExecutor> proxy_executor_;
78+};
79+ 
80+using CreateAOTIModelRunnerFunc = std::unique_ptr<AOTIModelContainerRunner> (*)(const std::string& model_so_path,
81+ size_t num_models,
82+ const std::string& device_str,
83+ const std::string& bin_dir,
84+ const bool run_single_threaded);
85+ 
86+// Return a global map "device name" -> "aoti model runner create function" for
87+// all registered in AOTI external backends
88+TORCH_API std::unordered_map<std::string, CreateAOTIModelRunnerFunc>& getAOTIModelRunnerRegistry();
89+ 
90+// To register a new external backend in AOTI one needs to create an instance of
91+// this struct. It is not thread-safe. Because it is expected to be called
92+// during the initialization of the program.
93+struct TORCH_API RegisterAOTIModelRunner {
94+ RegisterAOTIModelRunner(const std::string& name, CreateAOTIModelRunnerFunc create_aoti_model_runner_fn)
95+ {
96+ getAOTIModelRunnerRegistry()[name] = create_aoti_model_runner_fn;
97+ } // namespace torch::inductor
98+};
99+ 
100+} // namespace torch::inductor
101+#endif
Atorch_npu/csrc/inductor/aoti_runner/model_container_runner_npu.cpp+85-0
@@ -0,0 +1,85 @@
1+#if !defined(C10_MOBILE) && !defined(ANDROID)
2+#include <torch_npu/csrc/inductor/aoti_runner/model_container_runner_npu.h>
3+#include <torch_npu/csrc/inductor/aoti_torch/oss_proxy_executor_npu.h>
4+ 
5+#include <iostream>
6+ 
7+#ifndef _WIN32
8+#include <sys/stat.h>
9+#else
10+#include <filesystem>
11+namespace fs = std::filesystem;
12+#endif
13+ 
14+namespace {
15+bool file_exists(std::string& path)
16+{
17+#ifdef _WIN32
18+ return fs::exists(path);
19+#else
20+ struct stat rc{};
21+ return lstat(path.c_str(), &rc) == 0;
22+#endif
23+}
24+} // namespace
25+ 
26+namespace torch::inductor {
27+ 
28+AOTIModelContainerRunnerNpu::AOTIModelContainerRunnerNpu(const std::string& model_so_path, size_t num_models,
29+ const std::string& device_str, const std::string& cubin_dir,
30+ const bool run_single_threaded)
31+ : AOTIModelContainerRunner(model_so_path, num_models, device_str, cubin_dir, run_single_threaded)
32+{
33+ model_so_path_ = model_so_path;
34+ init_flag_ = false;
35+}
36+ 
37+AOTIModelContainerRunnerNpu::~AOTIModelContainerRunnerNpu() = default;
38+ 
39+void AOTIModelContainerRunnerNpu::init_proxy_executor()
40+{
41+ if (init_flag_)
42+ return;
43+ 
44+ init_flag_ = true;
45+ size_t lastindex = model_so_path_.find_last_of('.');
46+ std::string json_filename = model_so_path_.substr(0, lastindex) + "_npu.json";
47+ if (file_exists(json_filename)) {
48+ proxy_executor_npu_ = std::make_unique<torch::aot_inductor::OSSProxyExecutorNpu>(json_filename, false);
49+ proxy_executor_handle_ = reinterpret_cast<AOTIProxyExecutorHandle>(proxy_executor_npu_.get());
50+ } else {
51+ proxy_executor_handle_ = nullptr;
52+ }
53+}
54+ 
55+std::vector<at::Tensor> AOTIModelContainerRunnerNpu::run_impl(std::vector<AtenTensorHandle>& input_handles,
56+ void* stream_handle)
57+{
58+ init_proxy_executor();
59+ c10_npu::NPUStream npu_stream = c10_npu::getCurrentNPUStream();
60+ return AOTIModelContainerRunner::run_impl(input_handles, reinterpret_cast<void*>(npu_stream.stream()));
61+}
62+ 
63+std::vector<at::Tensor> AOTIModelContainerRunnerNpu::run_with_npu_stream(const std::vector<at::Tensor>& inputs,
64+ const c10_npu::NPUStream& npu_stream)
65+{
66+ init_proxy_executor();
67+ c10_npu::NPUStream cur_npu_stream = c10_npu::getCurrentNPUStream();
68+ return run(inputs, reinterpret_cast<void*>(cur_npu_stream.stream()));
69+}
70+ 
71+namespace {
72+std::unique_ptr<AOTIModelContainerRunner> create_aoti_runner_npu(const std::string& model_so_path, size_t num_models,
73+ const std::string& device_str,
74+ const std::string& cubin_dir,
75+ const bool run_single_threaded)
76+{
77+ return std::make_unique<AOTIModelContainerRunnerNpu>(model_so_path, num_models, device_str, cubin_dir,
78+ run_single_threaded);
79+}
80+} // namespace
81+ 
82+RegisterAOTIModelRunner register_npu_runner("npu", &create_aoti_runner_npu);
83+ 
84+} // namespace torch::inductor
85+#endif
Atorch_npu/csrc/inductor/aoti_runner/model_container_runner_npu.h+36-0
@@ -0,0 +1,36 @@
1+#if !defined(C10_MOBILE) && !defined(ANDROID)
2+#pragma once
3+ 
4+#include <torch_npu/csrc/core/npu/NPUStream.h>
5+#include <torch_npu/csrc/inductor/aoti_runner/model_container_runner.h>
6+ 
7+namespace torch::inductor {
8+ 
9+// NOTICE: Following APIs are subject to change due to active development
10+// We provide NO BC guarantee for these APIs
11+// NOLINTNEXTLINE(cppcoreguidelines-special-member-functions)
12+class AOTIModelContainerRunnerNpu : public AOTIModelContainerRunner {
13+public:
14+ // @param device_str: npu device string, e.g. "npu", "npu:0"
15+ AOTIModelContainerRunnerNpu(const std::string& model_so_path, size_t num_models = 1,
16+ const std::string& device_str = "npu", const std::string& cubin_dir = "",
17+ const bool run_single_threaded = false);
18+ 
19+ ~AOTIModelContainerRunnerNpu() override;
20+ 
21+ std::vector<at::Tensor> run_impl(std::vector<AtenTensorHandle>& input_handles, void* stream_handle) override;
22+ 
23+ std::vector<at::Tensor> run_with_npu_stream(const std::vector<at::Tensor>& inputs,
24+ const c10_npu::NPUStream& npu_stream);
25+ void init_proxy_executor();
26+ 
27+ void set_proxy_executor(AOTIProxyExecutorHandle handle);
28+ 
29+private:
30+ std::string model_so_path_;
31+ bool init_flag_;
32+ std::unique_ptr<torch::aot_inductor::ProxyExecutor> proxy_executor_npu_;
33+};
34+ 
35+} // namespace torch::inductor
36+#endif
Atorch_npu/csrc/inductor/aoti_runner/pybind.h+7-0
@@ -0,0 +1,7 @@
1+#include <torch/csrc/python_headers.h>
2+ 
3+namespace torch::inductor {
4+ 
5+void initAOTIRunnerBindings(PyObject* module);
6+ 
7+} // namespace torch::inductor
Atorch_npu/csrc/inductor/aoti_runtime/arrayref_tensor.h+284-0
@@ -0,0 +1,284 @@
1+#pragma once
2+ 
3+#include <torch_npu/csrc/inductor/aoti_runtime/utils.h>
4+#include <torch_npu/csrc/inductor/aoti_torch/c/shim.h>
5+ 
6+#include <cassert>
7+#include <cstdint>
8+#include <cstring>
9+ 
10+namespace torch::aot_inductor {
11+constexpr int64_t kInt32Size = 4;
12+ 
13+// Can't use c10::ArrayRef because it's not truly header-only and
14+// pulls in other c10 headers. This is (sadly) copy-pasted and
15+// adapted.
16+template <typename T> class MiniArrayRef final {
17+public:
18+ using iterator = T*;
19+ using const_iterator = const T*;
20+ using size_type = size_t;
21+ using value_type = T;
22+ using reverse_iterator = std::reverse_iterator<iterator>;
23+ 
24+ /// @name Constructors
25+ /// @{
26+ 
27+ /// Construct an empty MiniArrayRef.
28+ constexpr MiniArrayRef() : Data(nullptr), Length(0) {}
29+ 
30+ /// Construct an MiniArrayRef from a single element.
31+ constexpr MiniArrayRef(const T& OneElt) : Data(&OneElt), Length(1) {}
32+ 
33+ /// Construct an MiniArrayRef from a pointer and length.
34+ constexpr MiniArrayRef(T* data, size_t length) : Data(data), Length(length) {}
35+ 
36+ /// Construct an MiniArrayRef from a range.
37+ constexpr MiniArrayRef(T* begin, T* end) : Data(begin), Length(end - begin) {}
38+ 
39+ template <typename Container, typename = std::enable_if_t<std::is_same_v<
40+ std::remove_const_t<decltype(std::declval<Container>().data())>, T*> > >
41+ MiniArrayRef(Container& container) : Data(container.data()), Length(container.size())
42+ {
43+ }
44+ 
45+ /// Construct an MiniArrayRef from a std::vector.
46+ // The enable_if stuff here makes sure that this isn't used for
47+ // std::vector<bool>, because MiniArrayRef can't work on a std::vector<bool>
48+ // bitfield.
49+ template <typename A> MiniArrayRef(const std::vector<T, A>& Vec) : Data(Vec.data()), Length(Vec.size())
50+ {
51+ static_assert(!std::is_same_v<T, bool>, "MiniArrayRef<bool> cannot be constructed from a "
52+ "std::vector<bool> bitfield.");
53+ }
54+ 
55+ /// Construct an MiniArrayRef from a std::array
56+ template <size_t N> constexpr MiniArrayRef(std::array<T, N>& Arr) : Data(Arr.data()), Length(N) {}
57+ 
58+ /// Construct an MiniArrayRef from a C array.
59+ template <size_t N>
60+ // NOLINTNEXTLINE(*c-array*)
61+ constexpr MiniArrayRef(T (&Arr)[N]) : Data(Arr), Length(N)
62+ {
63+ }
64+ 
65+ // /// Construct an MiniArrayRef from an empty C array.
66+ constexpr MiniArrayRef(const volatile void* Arr) : Data(nullptr), Length(0) {}
67+ 
68+ /// Construct an MiniArrayRef from a std::initializer_list.
69+ constexpr MiniArrayRef(const std::initializer_list<T>& Vec)
70+ : Data(std::begin(Vec) == std::end(Vec) ? static_cast<T*>(nullptr) : std::begin(Vec)), Length(Vec.size())
71+ {
72+ }
73+ 
74+ /// @}
75+ /// @name Simple Operations
76+ /// @{
77+ 
78+ constexpr iterator begin() const { return Data; }
79+ constexpr iterator end() const { return Data + Length; }
80+ 
81+ // These are actually the same as iterator, since MiniArrayRef only
82+ // gives you const iterators.
83+ constexpr const_iterator cbegin() const { return Data; }
84+ constexpr const_iterator cend() const { return Data + Length; }
85+ 
86+ constexpr reverse_iterator rbegin() const { return reverse_iterator(end()); }
87+ constexpr reverse_iterator rend() const { return reverse_iterator(begin()); }
88+ 
89+ /// empty - Check if the array is empty.
90+ constexpr bool empty() const { return Length == 0; }
91+ 
92+ constexpr T* data() const { return Data; }
93+ 
94+ /// size - Get the array size.
95+ constexpr size_t size() const { return Length; }
96+ 
97+ /// equals - Check for element-wise equality.
98+ constexpr bool equals(MiniArrayRef RHS) const
99+ {
100+ return Length == RHS.Length && std::equal(begin(), end(), RHS.begin());
101+ }
102+ 
103+ /// @}
104+ /// @name Operator Overloads
105+ /// @{
106+ constexpr const T& operator[](size_t Index) const { return Data[Index]; }
107+ 
108+ /// Disallow accidental assignment from a temporary.
109+ ///
110+ /// The declaration here is extra complicated so that "arrayRef = {}"
111+ /// continues to select the move assignment operator.
112+ template <typename U>
113+ std::enable_if_t<std::is_same_v<U, T>, MiniArrayRef<T> >& operator=(
114+ // NOLINTNEXTLINE(cppcoreguidelines-missing-std-forward)
115+ U&& Temporary) = delete;
116+ 
117+ /// Disallow accidental assignment from a temporary.
118+ ///
119+ /// The declaration here is extra complicated so that "arrayRef = {}"
120+ /// continues to select the move assignment operator.
121+ template <typename U>
122+ std::enable_if_t<std::is_same_v<U, T>, MiniArrayRef<T> >& operator=(std::initializer_list<U>) = delete;
123+ 
124+private:
125+ /// The start of the array, in an external buffer.
126+ T* Data;
127+ 
128+ /// The number of elements.
129+ size_type Length;
130+};
131+ 
132+using MiniIntArrayRef = MiniArrayRef<int64_t>;
133+ 
134+static_assert(sizeof(MiniIntArrayRef) == sizeof(void*) + sizeof(size_t),
135+ "changing the size of MiniArrayRef breaks ABI compatibility!");
136+ 
137+inline bool is_contiguous_strides_for_shape(int64_t ndim, const int64_t* strides_ptr, const int64_t* sizes_ptr)
138+{
139+ int64_t z = 1;
140+ for (int64_t d = ndim - 1; d >= 0; d--) {
141+ const auto& size_d = sizes_ptr[d];
142+ if (size_d != 1) {
143+ if (strides_ptr[d] == z) {
144+ z *= size_d;
145+ } else {
146+ return false;
147+ }
148+ }
149+ }
150+ return true;
151+}
152+ 
153+// Shim for AOTI generated code to pretend a raw array works like an
154+// AtenTensorHandle.
155+template <typename T> class ArrayRefTensor {
156+public:
157+ ArrayRefTensor() = default;
158+ 
159+ explicit ArrayRefTensor(MiniArrayRef<T> arr, MiniArrayRef<const int64_t> sizes, MiniArrayRef<const int64_t> strides,
160+ int32_t device_type, int32_t device_idx)
161+ : arrayRef_(arr), sizes_(sizes), strides_(strides), device_type_(device_type), device_idx_(device_idx)
162+ {
163+ }
164+ 
165+ AtenTensorHandle expensiveCopyToTensor() const
166+ {
167+ AtenTensorHandle result = nullptr;
168+ AOTI_TORCH_ERROR_CODE_CHECK(aoti_torch_empty_strided(sizes_.size(), sizes_.data(), strides_.data(),
169+ aoti_torch_dtype<std::remove_const_t<T> >(), device_type_,
170+ device_idx_, &result));
171+ void* dataPtr = nullptr;
172+ AOTI_TORCH_ERROR_CODE_CHECK(aoti_torch_get_data_ptr(result, &dataPtr));
173+ std::memcpy(dataPtr, data(), numel() * sizeof(T));
174+ return result;
175+ }
176+ 
177+ // We need to look the same as RAIIAtenTensorHandle, which returns
178+ // an owning AtenTensorHandle from release(). So, we allocate one!
179+ AtenTensorHandle release() { return expensiveCopyToTensor(); }
180+ 
181+ AtenTensorHandle borrowAsTensor() const
182+ {
183+ AtenTensorHandle result = nullptr;
184+ AOTI_TORCH_ERROR_CODE_CHECK(aoti_torch_create_tensor_from_blob_v2(
185+ data(), sizes_.size(), sizes_.data(), strides_.data(), 0, aoti_torch_dtype<std::remove_const_t<T> >(),
186+ device_type_, device_idx_, &result, aoti_torch_layout_strided(), nullptr, 0));
187+ return result;
188+ }
189+ 
190+ // We don't need to free any memory.
191+ void reset() {}
192+ 
193+ auto sizes() const { return sizes_; }
194+ 
195+ auto strides() const { return strides_; }
196+ 
197+ auto device_type() const { return device_type_; }
198+ 
199+ auto device_idx() const { return device_idx_; }
200+ 
201+ T* data() const { return arrayRef_.data(); }
202+ 
203+ auto numel() const { return arrayRef_.size(); }
204+ 
205+ void set_arrayref(MiniArrayRef<T> new_arrayref) { arrayRef_ = new_arrayref; }
206+ 
207+private:
208+ MiniArrayRef<T> arrayRef_;
209+ // We expect generated code to have statically available sizes &
210+ // strides for us.
211+ MiniArrayRef<const int64_t> sizes_;
212+ MiniArrayRef<const int64_t> strides_;
213+ int32_t device_type_ = 0;
214+ int32_t device_idx_ = 0;
215+ // We continue to zero-initialize this field in case we repurpose
216+ // the space later; having predictable contents can only help.
217+ int32_t unusedDoNotRemoveForABICompatibility_ = 0;
218+};
219+ 
220+static_assert(sizeof(ArrayRefTensor<int>) == 3 * sizeof(MiniIntArrayRef) + 3 * sizeof(int32_t) +
221+ (alignof(ArrayRefTensor<int>) > kInt32Size ? sizeof(int32_t) : 0),
222+ "changing the size of ArrayRefTensor breaks ABI compatibility!");
223+ 
224+template <typename T>
225+inline ArrayRefTensor<T> reinterpret_tensor_wrapper(const ArrayRefTensor<T>& self, int64_t ndim,
226+ const int64_t* sizes_ptr, const int64_t* strides_ptr,
227+ int64_t storage_offset)
228+{
229+ return ArrayRefTensor<T>(MiniArrayRef<T>(self.data() + storage_offset, self.numel() - storage_offset),
230+ MiniArrayRef<const int64_t>(sizes_ptr, ndim),
231+ MiniArrayRef<const int64_t>(strides_ptr, ndim), self.device_type(), self.device_idx());
232+}
233+ 
234+template <typename T> inline T* get_data_ptr_wrapper(ArrayRefTensor<T>& tensor) { return tensor.data(); }
235+ 
236+template <typename T> inline T* get_data_ptr_wrapper(const MiniArrayRef<T>& arr) { return arr.data(); }
237+ 
238+template <typename T> inline const ArrayRefTensor<T>& unwrap_raii_handle_if_needed(const ArrayRefTensor<T>& tensor)
239+{
240+ return tensor;
241+}
242+ 
243+template <typename T> inline ArrayRefTensor<T>& unwrap_raii_handle_if_needed(ArrayRefTensor<T>& tensor)
244+{
245+ return tensor;
246+}
247+ 
248+template <typename T> inline const ArrayRefTensor<T>& wrap_with_raii_handle_if_needed(const ArrayRefTensor<T>& tensor)
249+{
250+ return tensor;
251+}
252+ 
253+template <typename T> inline ArrayRefTensor<T>& wrap_with_raii_handle_if_needed(ArrayRefTensor<T>& tensor)
254+{
255+ return tensor;
256+}
257+ 
258+template <typename T> inline ArrayRefTensor<T> wrap_with_raii_handle_if_needed(ArrayRefTensor<T>&& tensor)
259+{
260+ return std::move(tensor);
261+}
262+ 
263+template <typename T> inline RAIIAtenTensorHandle expensive_copy_to_tensor_if_needed(const ArrayRefTensor<T>& tensor)
264+{
265+ return tensor.expensiveCopyToTensor();
266+}
267+ 
268+inline AtenTensorHandle expensive_copy_to_tensor_if_needed(AtenTensorHandle handle) { return handle; }
269+ 
270+template <typename T> const T& copy_arrayref_tensor_to_tensor(const T& t) { return t; }
271+ 
272+template <typename T> RAIIAtenTensorHandle copy_arrayref_tensor_to_tensor(const ArrayRefTensor<T>& art)
273+{
274+ return art.expensiveCopyToTensor();
275+}
276+ 
277+template <typename T> const T& borrow_arrayref_tensor_as_tensor(const T& t) { return t; }
278+ 
279+template <typename T> RAIIAtenTensorHandle borrow_arrayref_tensor_as_tensor(const ArrayRefTensor<T>& art)
280+{
281+ return art.borrowAsTensor();
282+}
283+ 
284+} // namespace torch::aot_inductor
Atorch_npu/csrc/inductor/aoti_runtime/device_utils.h+41-0
@@ -0,0 +1,41 @@
1+#pragma once
2+ 
3+#if defined(USE_NPU)
4+ 
5+#include "third_party/acl/inc/acl/acl_base.h"
6+#include "third_party/acl/inc/acl/acl_rt.h"
7+ 
8+typedef void* NPUdeviceptr;
9+ 
10+typedef void* NPUfunction;
11+ 
12+#define AOTI_RUNTIME_DEVICE_CHECK(EXPR) \
13+ do { \
14+ const aclError code = EXPR; \
15+ if (code != ACL_SUCCESS) { \
16+ throw std::runtime_error(std::string("NPU error core: ") + std::to_string(code) + std::string(" ") + \
17+ std::string(__FILE__) + std::string(":") + std::to_string(__LINE__)); \
18+ } \
19+ } while (0)
20+ 
21+namespace torch::aot_inductor {
22+ 
23+using DeviceStreamType = aclrtStream;
24+ 
25+} // namespace torch::aot_inductor
26+ 
27+#else
28+ 
29+#define AOTI_RUNTIME_DEVICE_CHECK(EXPR) \
30+ bool ok = EXPR; \
31+ if (!ok) { \
32+ throw std::runtime_error("CPU runtime error"); \
33+ }
34+ 
35+namespace torch::aot_inductor {
36+ 
37+using DeviceStreamType = void*;
38+ 
39+} // namespace torch::aot_inductor
40+ 
41+#endif // USE_NPU
Atorch_npu/csrc/inductor/aoti_runtime/interface.h+178-0
@@ -0,0 +1,178 @@
1+// Licensed under the BSD 3-Clause License (the "License");
2+#pragma once
3+ 
4+// WARNING: Be careful when adding new includes here. This header will be used
5+// in model.so, and should not refer to any aten/c10 headers except the stable
6+// C ABI defined in torch/csrc/inductor/aoti_torch/c/shim.h. The same rule
7+// applies to other files under torch/csrc/inductor/aoti_runtime/.
8+#include <torch_npu/csrc/inductor/aoti_runtime/utils.h>
9+extern "C" {
10+struct AOTInductorModelOpaque;
11+using AOTInductorModelHandle = AOTInductorModelOpaque*;
12+ 
13+struct AOTInductorModelContainerOpaque;
14+using AOTInductorModelContainerHandle = AOTInductorModelContainerOpaque*;
15+ 
16+struct AOTInductorStreamOpaque;
17+using AOTInductorStreamHandle = AOTInductorStreamOpaque*;
18+ 
19+struct AOTInductorConstantMap;
20+using AOTInductorConstantMapHandle = AOTInductorConstantMap*;
21+ 
22+// Please use AOTInductorModelContainerCreateWithDevice instead.
23+AOTIRuntimeError AOTInductorModelContainerCreate(AOTInductorModelContainerHandle* container_handle, size_t num_models,
24+ bool is_cpu, const char* cubin_dir);
25+ 
26+// Creates an AOTInductor model container. The parameter num_models
27+// specifies the number of model instances that may be run concurrently for
28+// the same input model.
29+// `device_str` MUST NOT be nullptr. It must be a valid device string, e.g.
30+// "cpu", "cuda", "cuda:0", etc. If the device index is not specified for CUDA
31+// device, runtime will use the device index returned by
32+AOTIRuntimeError AOTInductorModelContainerCreateWithDevice(AOTInductorModelContainerHandle* container_handle,
33+ size_t num_models, const char* device_str,
34+ const char* cubin_dir);
35+ 
36+// Deletes the AOTInductor model container.
37+AOTIRuntimeError AOTInductorModelContainerDelete(AOTInductorModelContainerHandle container_handle);
38+ 
39+// Runs the inference.
40+AOTIRuntimeError AOTInductorModelContainerRun(AOTInductorModelContainerHandle container_handle,
41+ AtenTensorHandle* input_handles, // array of input AtenTensorHandle; handles
42+ // are stolen; the array itself is borrowed
43+ size_t num_inputs,
44+ AtenTensorHandle* output_handles, // array for writing output AtenTensorHandle; handles
45+ // will be stolen by the caller; the array itself is
46+ // borrowed
47+ size_t num_outputs, AOTInductorStreamHandle stream_handle,
48+ AOTIProxyExecutorHandle proxy_executor_handle);
49+ 
50+// Single-threaded variant of previous.
51+AOTIRuntimeError AOTInductorModelContainerRunSingleThreaded(
52+ AOTInductorModelContainerHandle container_handle,
53+ AtenTensorHandle* input_handles, // array of input AtenTensorHandle; handles
54+ // are stolen; the array itself is borrowed
55+ size_t num_inputs,
56+ AtenTensorHandle* output_handles, // array for writing output AtenTensorHandle; handles
57+ // will be stolen by the caller; the array itself is
58+ // borrowed
59+ size_t num_outputs, AOTInductorStreamHandle stream_handle, AOTIProxyExecutorHandle proxy_executor_handle);
60+ 
61+// Retrieves the number of constants for the model.
62+AOTIRuntimeError AOTInductorModelContainerGetNumConstants(AOTInductorModelContainerHandle container_handle,
63+ size_t* num_constants);
64+ 
65+// Retrieves a constant's name.
66+// idx is the index of the internal's constants.
67+// Need idx < num_constants from AOTInductorModelContainerGetNumConstants
68+AOTIRuntimeError AOTInductorModelContainerGetConstantName(AOTInductorModelContainerHandle container_handle, size_t idx,
69+ const char** name);
70+ 
71+// Retrieves a constant's original FQN.
72+// idx is the index of the internal's constants.
73+// Need idx < num_constants from AOTInductorModelContainerGetNumConstants
74+AOTIRuntimeError AOTInductorModelContainerGetConstantOriginalFQN(AOTInductorModelContainerHandle container_handle,
75+ size_t idx, const char** original_fqn);
76+ 
77+// Retrieves whether a constant is from folded.
78+// idx is the index of the internal's constants.
79+// Need idx < num_constants from AOTInductorModelContainerGetNumConstants
80+AOTIRuntimeError AOTInductorModelContainerGetConstantFromFolded(AOTInductorModelContainerHandle container_handle,
81+ size_t idx, bool* from_folded);
82+ 
83+// Retrieves the inductor constant type.
84+// idx is the index of the internal's constants.
85+// Need idx < num_constants from AOTInductorModelContainerGetNumConstants
86+AOTIRuntimeError AOTInductorModelContainerGetConstantType(AOTInductorModelContainerHandle container_handle, size_t idx,
87+ int32_t* type);
88+ 
89+// Retrieves a constant's dtype.
90+// idx is the index of the internal's constants.
91+// Need idx < num_constants from AOTInductorModelContainerGetNumConstants
92+AOTIRuntimeError AOTInductorModelContainerGetConstantDtype(AOTInductorModelContainerHandle container_handle, size_t idx,
93+ int32_t* dtype);
94+ 
95+// Retrieves a constant's data size.
96+// idx is the index of the internal's constants.
97+// Need idx < num_constants from AOTInductorModelContainerGetNumConstants
98+AOTIRuntimeError AOTInductorModelContainerGetConstantDataSize(AOTInductorModelContainerHandle container_handle,
99+ size_t idx, size_t* data_size);
100+ 
101+// Extract the constants that is being used in the container.
102+AOTIRuntimeError AOTInductorModelContainerExtractConstantsMap(AOTInductorModelContainerHandle container_handle,
103+ AOTInductorConstantMapHandle constant_map_handle,
104+ bool use_inactive);
105+ 
106+// Setup the constant buffer in model container with provided ConstantMap.
107+// The ConstantMap is user managed, and the user would retain ownership.
108+AOTIRuntimeError AOTInductorModelContainerUpdateUserManagedConstantBuffer(AOTInductorModelContainerHandle container_handle,
109+ AOTInductorConstantMapHandle constant_map_handle,
110+ bool use_inactive, bool validate_full_update);
111+ 
112+// Setup the constant buffer in model container with provided ConstantMap
113+// use_inactive should be set as true if the inactive buffer is to be updated.
114+// validate_full_update checks if all constants are included in the ConstantMap
115+AOTIRuntimeError AOTInductorModelContainerUpdateConstantBuffer(AOTInductorModelContainerHandle container_handle,
116+ AOTInductorConstantMapHandle constant_map_handle,
117+ bool use_inactive, bool validate_full_update);
118+ 
119+// Setup the inactive constant buffer in model container with provided
120+// ConstantMap
121+AOTIRuntimeError AOTInductorModelContainerUpdateInactiveConstantBuffer(AOTInductorModelContainerHandle container_handle,
122+ AOTInductorConstantMapHandle constant_map_handle);
123+ 
124+// Free the inactive constant buffer in model container.
125+AOTIRuntimeError AOTInductorModelContainerFreeInactiveConstantBuffer(AOTInductorModelContainerHandle container_handle);
126+ 
127+// Run constant folding on constant buffer.
128+AOTIRuntimeError AOTInductorModelContainerRunConstantFolding(AOTInductorModelContainerHandle container_handle,
129+ bool use_inactive, AOTInductorStreamHandle stream_handle,
130+ AOTIProxyExecutorHandle proxy_executor_handle);
131+ 
132+// Swap the constant buffer being used to the inactive one.
133+AOTIRuntimeError AOTInductorModelContainerSwapConstantBuffer(AOTInductorModelContainerHandle container_handle);
134+ 
135+// Retrieves the number of inputs for the model.
136+AOTIRuntimeError AOTInductorModelContainerGetNumInputs(AOTInductorModelContainerHandle container_handle,
137+ size_t* ret_num_inputs);
138+ 
139+// Retrieves the input name at the given index.
140+AOTIRuntimeError AOTInductorModelContainerGetInputName(AOTInductorModelContainerHandle container_handle,
141+ size_t input_idx, const char** ret_input_names);
142+ 
143+// Retrieves the number of outputs for the model.
144+AOTIRuntimeError AOTInductorModelContainerGetNumOutputs(AOTInductorModelContainerHandle container_handle,
145+ size_t* ret_num_outputs);
146+ 
147+// Retrieves the output name at the given index.
148+AOTIRuntimeError AOTInductorModelContainerGetOutputName(AOTInductorModelContainerHandle container_handle,
149+ size_t output_idx, const char** ret_output_names);
150+ 
151+// Creates an AOTInductorModel instance. This is a thin and light wrapper
152+// around the compiled model; it doesn't handle concurrency, queueing, device
153+// management, etc. Use this if bare-metal performance is needed and you are
154+// willing to handle other "management" aspects yourself.
155+// constant_map_handle is an opaque type to satisfy the C ABI. It should be a
156+// std::unordered_map<std::string, at::Tensor*>*.
157+AOTIRuntimeError AOTInductorModelCreate(AOTInductorModelHandle* model_handle,
158+ AOTInductorConstantMapHandle constant_map_handle);
159+ 
160+// Run an AOTInductorModel (see AOTInductorModelCreate for when one should use
161+// this function versus AOTInductorModelContainerRun).
162+AOTIRuntimeError AOTInductorModelRun(AOTInductorModelHandle model_handle, AtenTensorHandle* input_handles,
163+ AtenTensorHandle* output_handles);
164+ 
165+// Replace AOTInductorModel's constant map. Note it doesn't handle concurrency
166+// so be sure to handle ordering if AOTInductorModelRun is ran concurrently.
167+AOTIRuntimeError AOTInductorModelUpdateConstantsMap(AOTInductorModelHandle model_handle,
168+ AOTInductorConstantMapHandle constant_map_handle);
169+ 
170+// Delete an AOTInductorModel created by AOTInductorModelCreate.
171+AOTIRuntimeError AOTInductorModelDelete(AOTInductorModelHandle model_handle);
172+ 
173+AOTIRuntimeError AOTInductorModelGetNumOutputs(AOTInductorModelHandle model_handle, size_t* ret_num_outputs);
174+ 
175+AOTIRuntimeError AOTInductorModelContainerGetCallSpec(AOTInductorModelContainerHandle container_handle,
176+ const char** in_spec, const char** out_spec);
177+ 
178+} // extern "C"
Atorch_npu/csrc/inductor/aoti_runtime/model.h+548-0
@@ -0,0 +1,548 @@
1+#pragma once
2+ 
3+#include <dlfcn.h>
4+#include <fcntl.h>
5+#include <sys/mman.h>
6+#include <unistd.h>
7+ 
8+#include <optional>
9+#include <regex>
10+#include <stdexcept>
11+#include <unordered_map>
12+#include <utility>
13+ 
14+// WARNING: Be careful when adding new includes here. This header will be used
15+// in model.so, and should not refer to any aten/c10 headers except the stable
16+// C ABI defined in torch/csrc/inductor/aoti_torch/c/shim.h. The same rule
17+// applies to other files under torch/csrc/inductor/aoti_runtime/.
18+#include <torch_npu/csrc/inductor/aoti_runtime/device_utils.h>
19+#include <torch_npu/csrc/inductor/aoti_runtime/utils.h>
20+ 
21+#define AOTI_RUNTIME_CHECK(EXPR, MSG) \
22+ do { \
23+ bool ok = EXPR; \
24+ if (!ok) { \
25+ throw std::runtime_error(MSG); \
26+ } \
27+ } while (0)
28+ 
29+// At codegen time, we write out a binary file called constants.bin.
30+// We then turn the raw binary to an object file that exposes this
31+// symbol and link it into the final .so.
32+// The constants are NOT readonly because they may be mutated.
33+// NOLINTNEXTLINE(*array*)
34+extern uint8_t _binary_constants_bin_start[];
35+// NOLINTNEXTLINE(*array*)
36+extern uint8_t _binary_constants_bin_end[];
37+ 
38+#define AOTI_CONST_ALIGNMENT 64
39+ 
40+namespace {
41+ 
42+using RAIIDataPtr = std::unique_ptr<void, std::function<void(void*)> >;
43+ 
44+#ifdef USE_NPU
45+ 
46+RAIIDataPtr RAII_npuMalloc(size_t num_bytes)
47+{
48+ void* data_ptr;
49+ // aclrtMalloc doesn't support allocate 0-bytes. In this case,
50+ // e.g, model has no weight, we should do padding.
51+ size_t padding_bytes = 32;
52+ if (num_bytes == 0)
53+ num_bytes = padding_bytes;
54+ AOTI_RUNTIME_DEVICE_CHECK(aclrtMalloc((void**)&data_ptr, num_bytes, ACL_MEM_MALLOC_HUGE_FIRST));
55+ auto deleter = [](void* ptr) { AOTI_RUNTIME_DEVICE_CHECK(aclrtFree(ptr)); };
56+ return RAIIDataPtr(data_ptr, deleter);
57+}
58+ 
59+#endif // USE_NPU
60+ 
61+RAIIDataPtr RAII_cpuMalloc(size_t num_bytes)
62+{
63+ void* data_ptr = std::malloc(num_bytes);
64+ if (!data_ptr) {
65+ throw std::bad_alloc();
66+ }
67+ auto deleter = [](void* ptr) { std::free(ptr); };
68+ return RAIIDataPtr(data_ptr, deleter);
69+}
70+ 
71+} // anonymous namespace
72+ 
73+namespace torch::aot_inductor {
74+enum ConstantType : uint8_t {
75+ Unknown = 0,
76+ Parameter = 1,
77+ Buffer = 2,
78+ TensorConstant = 3,
79+ FoldedConstant = 4,
80+};
81+ 
82+using ConstantMap = std::unordered_map<std::string, MaybeOwningAtenTensorHandle>;
83+ 
84+// valid device strs are: cpu, npu, npu:0, npu:1, ...
85+// Update the list here if more devices are supported in the future
86+inline void parse_device_str(const std::string& device_str, int32_t& device_type, int32_t& device_idx)
87+{
88+ std::regex re("(cpu|npu)(:([0-9]+))?");
89+ std::smatch sm;
90+ bool matched = std::regex_match(device_str, sm, re);
91+ AOTI_RUNTIME_CHECK(matched, "Invalid device: " + device_str);
92+ 
93+ if (sm[1].str() == "cpu") {
94+ device_type = aoti_torch_device_type_cpu();
95+#ifdef USE_NPU
96+ } else if (sm[1].str() == "npu") {
97+ device_type = aoti_torch_device_type_npu();
98+#endif
99+ } else {
100+ AOTI_RUNTIME_CHECK(false, "Invalid device: " + device_str);
101+ }
102+ const size_t default_sm = 3;
103+ if (sm[default_sm].matched) {
104+ device_idx = stoi(sm[default_sm].str());
105+ } else {
106+ device_idx = -1;
107+ }
108+}
109+ 
110+// Defines the base class for AOTInductorModel, which is generated by the
111+// AOTInductor cpp codegen. Since we do not need dynamic dispatch, we rely
112+// on curiously recurring template pattern (CRTP) to save some runtime
113+// v-table overhead. The generated AOTInductorModel is specialized with
114+// methods such as run_impl.
115+template <typename Model> class AOTInductorModelBase {
116+public:
117+ AOTInductorModelBase(size_t num_inputs, size_t num_outputs, size_t num_constants, const std::string& device_str,
118+ std::optional<std::string> cubin_dir, bool include_weights = true)
119+ : inputs_info_(num_inputs), outputs_info_(num_outputs), constants_info_(num_constants),
120+ cubin_dir_(std::move(cubin_dir)), include_weights(include_weights)
121+ {
122+ parse_device_str(device_str, device_type_, device_idx_);
123+ 
124+#ifdef USE_NPU
125+ if (device_idx_ == -1) {
126+ AOTI_RUNTIME_DEVICE_CHECK(aclrtSetDevice(0));
127+ AOTI_RUNTIME_DEVICE_CHECK(aclrtGetDevice(&device_idx_));
128+ } else {
129+ AOTI_RUNTIME_DEVICE_CHECK(aclrtSetDevice(device_idx_));
130+ }
131+#endif // USE_NPU
132+ }
133+ 
134+ // NOLINTNEXTLINE(modernize-use-equals-default)
135+ ~AOTInductorModelBase()
136+ {
137+#ifdef USE_NPU
138+ if (run_finished_) {
139+ auto code = aclrtDestroyEvent(*run_finished_);
140+ if (code != ACL_SUCCESS) {
141+ std::cerr << "Failed to destroy NPU event in AOTInductor model erorr code: " << code << std::endl;
142+ }
143+ }
144+#endif // USE_NPU
145+ }
146+ 
147+ AOTInductorModelBase(AOTInductorModelBase&&) = delete;
148+ AOTInductorModelBase& operator=(AOTInductorModelBase&&) = delete;
149+ AOTInductorModelBase(const AOTInductorModelBase&) = delete;
150+ AOTInductorModelBase& operator=(const AOTInductorModelBase&) = delete;
151+ 
152+ void run(AtenTensorHandle* input_handles, // array of input AtenTensorHandle; handles
153+ // are stolen; the array itself is borrowed
154+ AtenTensorHandle* output_handles, // array for writing output AtenTensorHandle; handles
155+ // will be stolen by the caller; the array itself is
156+ // borrowed
157+ DeviceStreamType stream, AOTIProxyExecutorHandle proxy_executor)
158+ {
159+#if defined(USE_NPU)
160+ if (!run_finished_) {
161+ aclrtEvent run_finished;
162+ AOTI_RUNTIME_DEVICE_CHECK(aclrtCreateEvent(&run_finished));
163+ run_finished_.emplace(run_finished);
164+ }
165+#else
166+ run_finished_ = false;
167+#endif
168+ 
169+ auto* model = static_cast<Model*>(this);
170+ model->run_impl(input_handles, output_handles, stream, proxy_executor);
171+ 
172+#if defined(USE_NPU)
173+ AOTI_RUNTIME_DEVICE_CHECK(aclrtRecordEvent(*run_finished_, stream));
174+#else
175+ run_finished_ = true;
176+#endif
177+ }
178+ 
179+ // Non-thread-aware variant of run(). Obviously unsafe to use in a threaded
180+ // environment :)
181+ void run_single_threaded(AtenTensorHandle* input_handles, // array of input AtenTensorHandle; handles
182+ // are stolen; the array itself is borrowed
183+ AtenTensorHandle* output_handles, // array for writing output AtenTensorHandle; handles
184+ // will be stolen by the caller; the array itself is
185+ // borrowed
186+ DeviceStreamType stream, AOTIProxyExecutorHandle proxy_executor)
187+ {
188+ // don't bother with any of the run_finished stuff; this is unsafe to call
189+ // in a threaded context
190+ auto* model = static_cast<Model*>(this);
191+ model->run_impl(input_handles, output_handles, stream, proxy_executor);
192+ }
193+ 
194+ std::unordered_map<std::string, AtenTensorHandle>
195+ run_const_fold(DeviceStreamType stream, AOTIProxyExecutorHandle proxy_executor, bool initialization = false)
196+ {
197+#if defined(USE_NPU)
198+ if (!run_finished_) {
199+ aclrtEvent run_finished;
200+ AOTI_RUNTIME_DEVICE_CHECK(aclrtCreateEvent(&run_finished));
201+ run_finished_.emplace(run_finished);
202+ }
203+#else
204+ run_finished_ = false;
205+#endif
206+ 
207+ auto* model = static_cast<Model*>(this);
208+ auto folded_constants = model->const_run_impl(stream, proxy_executor, initialization);
209+ 
210+#if defined(USE_NPU)
211+ AOTI_RUNTIME_DEVICE_CHECK(aclrtRecordEvent(*run_finished_, stream));
212+#else
213+ run_finished_ = true;
214+#endif
215+ return folded_constants;
216+ }
217+ 
218+ void load_constants()
219+ {
220+ size_t num_constants = this->num_constants();
221+ size_t num_folded_constants = this->num_folded_constants();
222+ constants_map_->reserve(num_constants);
223+ 
224+ std::vector<size_t> constants_internal_offset(num_constants - num_folded_constants);
225+ size_t blob_size = 0;
226+ compute_constant_blob(blob_size, constants_internal_offset);
227+#if defined(USE_NPU)
228+ constant_blob_ = RAII_npuMalloc(blob_size);
229+#else
230+ constant_blob_ = RAII_cpuMalloc(blob_size);
231+#endif
232+ if (!include_weights) {
233+ return;
234+ }
235+ 
236+ size_t bytes_read = 0;
237+ for (size_t i = 0; i < num_constants; i++) {
238+ bool from_folded = this->constant_from_folded(i);
239+ if (from_folded) {
240+ continue;
241+ }
242+ std::string name = this->constant_name(i);
243+ size_t data_size = this->constant_data_size(i);
244+ uint8_t* internal_ptr = (data_size != 0) ? constant_ptr(constants_internal_offset[i], bytes_read, data_size,
245+ false)
246+ : nullptr;
247+ bytes_read += data_size;
248+ 
249+ // Create at::Tensor from copied memory.
250+ auto dtype = this->constant_dtype(i);
251+ auto ndim = this->constant_ndim(i);
252+ auto size = this->constant_shape(i);
253+ auto stride = this->constant_stride(i);
254+ auto offset = this->constant_offset(i);
255+ auto layout = this->constant_layout(i);
256+ auto opaque_metadata_ptr = this->opaque_metadata(i);
257+ auto opaque_metadata_size = this->opaque_metadata_size(i);
258+ 
259+ AtenTensorHandle tensor_handle = nullptr;
260+ AOTI_TORCH_ERROR_CODE_CHECK(aoti_torch_create_tensor_from_blob_npu_v2(
261+ internal_ptr, ndim, size, stride, offset, dtype, device_type_, device_idx_, &tensor_handle, layout,
262+ opaque_metadata_ptr, opaque_metadata_size));
263+ constants_map_->emplace(std::move(name), tensor_handle);
264+ }
265+ if (constants_map_) {
266+ this->update_constants_array_from_map();
267+ }
268+ }
269+ 
270+ RAIIDataPtr&& release_constant_blob() { return std::move(constant_blob_); }
271+ 
272+ std::shared_ptr<std::vector<ConstantHandle> > get_constants_array() { return constants_; }
273+ 
274+ int32_t get_device_type() const { return device_type_; }
275+ 
276+ int32_t get_device_idx() const { return device_idx_; }
277+ 
278+ uint8_t* constant_ptr(size_t constant_offset, size_t bytes_read, size_t data_size, bool skip_copy)
279+ {
280+ auto* constants_ptr = static_cast<uint8_t*>(constant_blob_.get());
281+ uint8_t* internal_ptr = constants_ptr + constant_offset;
282+ if (!skip_copy) {
283+#if defined(USE_NPU)
284+ AOTI_RUNTIME_DEVICE_CHECK(aclrtMemcpy(internal_ptr, data_size, _get_constants_start() + bytes_read,
285+ data_size, ACL_MEMCPY_HOST_TO_DEVICE));
286+#else
287+ memcpy(internal_ptr, _get_constants_start() + bytes_read, data_size);
288+#endif
289+ }
290+ return internal_ptr;
291+ }
292+ 
293+ void compute_constant_blob(size_t& blob_size, std::vector<size_t>& constants_internal_offset)
294+ {
295+ size_t num_constants = this->num_constants();
296+ blob_size = 0;
297+ size_t curr_idx = 0;
298+ for (size_t i = 0; i < num_constants; i++) {
299+ if (this->constant_from_folded(i)) {
300+ continue;
301+ }
302+ size_t data_size = this->constant_data_size(i);
303+ if (data_size % AOTI_CONST_ALIGNMENT) {
304+ data_size = AOTI_CONST_ALIGNMENT + (data_size / AOTI_CONST_ALIGNMENT) * AOTI_CONST_ALIGNMENT;
305+ }
306+ constants_internal_offset[curr_idx++] = blob_size;
307+ blob_size += data_size;
308+ }
309+ }
310+ 
311+ size_t num_inputs() const { return inputs_info_.size(); }
312+ 
313+ size_t num_outputs() const { return outputs_info_.size(); }
314+ 
315+ size_t num_constants() const { return constants_info_.size(); }
316+ 
317+ size_t num_folded_constants() const
318+ {
319+ size_t total_consts = this->num_constants();
320+ size_t folded_consts = 0;
321+ for (size_t i = 0; i < total_consts; i++) {
322+ if (this->constant_from_folded(i)) {
323+ folded_consts++;
324+ }
325+ }
326+ return folded_consts;
327+ }
328+ 
329+ const char* input_name(int64_t idx) const { return inputs_info_.at(idx).name; }
330+ 
331+ const char* output_name(int64_t idx) const { return outputs_info_.at(idx).name; }
332+ 
333+ const char* constant_name(int64_t idx) const { return constants_info_.at(idx).name; }
334+ 
335+ size_t constant_ndim(int64_t idx) { return constants_info_.at(idx).shape.size(); }
336+ 
337+ const int64_t* constant_shape(int64_t idx) const { return constants_info_.at(idx).shape.data(); }
338+ 
339+ const int64_t* constant_stride(int64_t idx) const { return constants_info_.at(idx).stride.data(); }
340+ 
341+ int32_t constant_dtype(int64_t idx) const { return constants_info_.at(idx).dtype; }
342+ 
343+ int32_t constant_layout(int64_t idx) const { return constants_info_.at(idx).layout; }
344+ 
345+ size_t constant_offset(int64_t idx) const { return constants_info_.at(idx).offset; }
346+ 
347+ size_t constant_data_size(int64_t idx) const { return constants_info_.at(idx).data_size; }
348+ 
349+ const char* constant_original_fqn(int64_t idx) const { return constants_info_.at(idx).original_fqn; }
350+ 
351+ const uint8_t* opaque_metadata(int64_t idx) const { return constants_info_.at(idx).opaque_metadata.data(); }
352+ 
353+ size_t opaque_metadata_size(int64_t idx) { return constants_info_.at(idx).opaque_metadata.size(); }
354+ 
355+ bool constant_from_folded(int64_t idx) const { return constants_info_.at(idx).from_folded; }
356+ 
357+ int32_t constant_type(int64_t idx) const { return constants_info_.at(idx).type; }
358+ 
359+ const char* get_in_spec() const { return in_spec_.c_str(); }
360+ 
361+ const char* get_out_spec() const { return out_spec_.c_str(); }
362+ 
363+ void update_constants_array_from_map()
364+ {
365+ if (!constants_map_) {
366+ throw std::runtime_error{
367+ "constants_map_ was not ready when constants_ is trying to be constructed from it!"};
368+ }
369+ if (!constants_) {
370+ constants_ = std::make_shared<std::vector<ConstantHandle> >(constants_info_.size());
371+ } else {
372+ constants_->resize(constants_info_.size());
373+ }
374+ int idx = 0;
375+ for (const auto& info : constants_info_) {
376+ const auto it = constants_map_->find(info.name);
377+ if (it != constants_map_->end()) {
378+ constants_->at(idx) = ConstantHandle(it->second);
379+ }
380+ idx++;
381+ }
382+ }
383+ 
384+ void update_constants_map(std::shared_ptr<ConstantMap> constants_map, bool remap_constants_array = true)
385+ {
386+ constants_map_ = std::move(constants_map);
387+ if (remap_constants_array) {
388+ update_constants_array_from_map();
389+ }
390+ }
391+ 
392+ // This function allows us to update the constants_ that is used to look up
393+ // the corresponding constant tensor during runtime.
394+ void update_constants_array(std::shared_ptr<std::vector<ConstantHandle> > constants_array)
395+ {
396+ constants_ = std::move(constants_array);
397+ }
398+ 
399+ /// Returns true if the model is complete.
400+ bool is_finished()
401+ {
402+#if defined(USE_NPU)
403+ if (!run_finished_) {
404+ throw std::runtime_error{"Model NPU event was not initialized"};
405+ }
406+ aclrtEventRecordedStatus recordStatus = ACL_EVENT_RECORDED_STATUS_NOT_READY;
407+ AOTI_RUNTIME_DEVICE_CHECK(aclrtQueryEventStatus(*run_finished_, &recordStatus));
408+ 
409+ if (recordStatus == ACL_EVENT_RECORDED_STATUS_COMPLETE) {
410+ return true;
411+ } else {
412+ return false;
413+ }
414+#else
415+ return run_finished_;
416+#endif
417+ }
418+ 
419+ /// Synchronizes completion event.
420+ void wait_for_completion() {}
421+ 
422+protected:
423+ uint8_t* _get_constants_start()
424+ {
425+#ifndef USE_MMAP_SELF
426+ // NOLINTNEXTLINE(*const-cast*)
427+ return const_cast<uint8_t*>(_binary_constants_bin_start);
428+#else
429+ if (self_mmap) {
430+ return self_mmap;
431+ }
432+ Dl_info dl_info;
433+ // get pointer to constant which are appended to the binary
434+ AOTI_RUNTIME_CHECK(dladdr(__func__, &dl_info), "Can't find shared library name");
435+ int fd = open(dl_info.dli_fname, O_RDONLY);
436+ AOTI_RUNTIME_CHECK(fd >= 0, "Shared library file cannot be opened");
437+ auto fsize = lseek(fd, 0, SEEK_END);
438+ auto weights_size = reinterpret_cast<const uint64_t*>(_binary_constants_bin_start)[0];
439+ auto magic_number = reinterpret_cast<const uint64_t*>(_binary_constants_bin_start)[1];
440+ auto weights_offset = fsize - weights_size;
441+ AOTI_RUNTIME_CHECK((weights_offset & 0x3fff) == 0, "weights_offset must be aligned to 16K boundary");
442+ auto ptr = mmap(NULL, weights_size, PROT_READ | PROT_WRITE, MAP_PRIVATE, fd, weights_offset);
443+ close(fd);
444+ AOTI_RUNTIME_CHECK(ptr != MAP_FAILED, "mmap() failed");
445+ self_mmap = static_cast<uint8_t*>(ptr);
446+ AOTI_RUNTIME_CHECK(reinterpret_cast<uint64_t*>(self_mmap + weights_size - sizeof(uint64_t))[0] == magic_number,
447+ "Weights data seems corrupt");
448+ return self_mmap;
449+#endif
450+ }
451+ struct ParamInfo {
452+ const char* name = nullptr;
453+ };
454+ 
455+ struct ConstInfo {
456+ const char* name = nullptr;
457+ std::vector<int64_t> shape;
458+ std::vector<int64_t> stride;
459+ int32_t dtype{};
460+ int64_t offset{};
461+ size_t data_size{};
462+ int32_t layout{};
463+ std::vector<uint8_t> opaque_metadata;
464+ int64_t opaque_metadata_size{};
465+ const char* original_fqn = nullptr;
466+ bool from_folded{};
467+ int32_t type{};
468+ };
469+ 
470+ std::vector<ParamInfo> inputs_info_;
471+ std::vector<ParamInfo> outputs_info_;
472+ std::vector<ConstInfo> constants_info_;
473+ std::string in_spec_;
474+ std::string out_spec_;
475+ 
476+ std::shared_ptr<ConstantMap> constants_map_;
477+ std::shared_ptr<std::vector<ConstantHandle> > constants_;
478+ 
479+ // Holds the blob storage for constants' at::Tensor.
480+ RAIIDataPtr constant_blob_;
481+ 
482+#ifdef USE_MMAP_SELF
483+ uint8_t* self_mmap = NULL;
484+#endif
485+ 
486+ // A directory with CUDA binary files, e.g. compiled kernels, etc.
487+ const std::optional<std::string> cubin_dir_;
488+ 
489+ // This is the flag that implies whether the weight is included in the model.
490+ // If True, we would prepare the weight when loading the model, otherwise the
491+ // model will be loaded without weights, and need to be provided by the user.
492+ bool include_weights;
493+ 
494+ // Record if the model finishes an inference run so that its owning
495+ // AOTModelContainer can re-use this instance.
496+#if defined(USE_NPU)
497+ std::optional<aclrtEvent> run_finished_;
498+#else
499+ bool run_finished_{};
500+#endif
501+ 
502+ // Generated model uses this device index to create CUDA guards.
503+ int32_t device_type_{};
504+ int32_t device_idx_{};
505+};
506+ 
507+// Codegen-ed classes can derive from this to keep pointers to loaded kernels.
508+class AOTInductorModelKernelsBase {
509+public:
510+ virtual ~AOTInductorModelKernelsBase() = default;
511+};
512+ 
513+class AOTInductorModel : public AOTInductorModelBase<AOTInductorModel> {
514+public:
515+ AOTInductorModel(std::shared_ptr<ConstantMap> constants_map,
516+ std::shared_ptr<std::vector<ConstantHandle>> constants_array, const std::string &device_str,
517+ std::optional<std::string> cubin_dir);
518+ 
519+ std::unordered_map<std::string, AtenTensorHandle> const_run_impl(
520+ DeviceStreamType stream, AOTIProxyExecutorHandle proxy_executor, bool initialization = false);
521+ 
522+ void _const_run_impl(
523+ std::vector<AtenTensorHandle> &output_handles, DeviceStreamType stream, AOTIProxyExecutorHandle proxy_executor);
524+ 
525+ void run_impl(AtenTensorHandle *input_handles, // array of input AtenTensorHandle; handles
526+ // are stolen; the array itself is borrowed
527+ AtenTensorHandle *output_handles, // array for writing output AtenTensorHandle; handles
528+ // will be stolen by the caller; the array itself is
529+ // borrowed
530+ DeviceStreamType stream, AOTIProxyExecutorHandle proxy_executor);
531+ 
532+ template <typename Inputs, typename Outputs>
533+ Outputs run_impl_minimal_arrayref_interface(
534+ const Inputs &inputs, DeviceStreamType stream, AOTIProxyExecutorHandle proxy_executor);
535+ 
536+ static std::unique_ptr<AOTInductorModel> Create(std::shared_ptr<ConstantMap> constants_map,
537+ std::shared_ptr<std::vector<ConstantHandle>> constants_array, const std::string &device_str,
538+ std::optional<std::string> cubin_dir)
539+ {
540+ return std::make_unique<AOTInductorModel>(
541+ std::move(constants_map), std::move(constants_array), device_str, std::move(cubin_dir));
542+ }
543+ 
544+private:
545+ std::unique_ptr<AOTInductorModelKernelsBase> kernels_;
546+};
547+ 
548+} // namespace torch::aot_inductor
Atorch_npu/csrc/inductor/aoti_runtime/model_container.h+673-0
@@ -0,0 +1,673 @@
1+// Licensed under the BSD 3-Clause License (the "License");
2+#pragma once
3+ 
4+#include <algorithm>
5+#include <condition_variable>
6+#include <deque>
7+#include <mutex>
8+#include <shared_mutex>
9+ 
10+// WARNING: Be careful when adding new includes here. This header will be used
11+// in model.so, and should not refer to any aten/c10 headers except the stable
12+// C ABI defined in torch/csrc/inductor/aoti_torch/c/shim.h. The same rule
13+// applies to other files under torch/csrc/inductor/aoti_runtime/.
14+#include <torch_npu/csrc/inductor/aoti_runtime/model.h>
15+ 
16+namespace torch::aot_inductor {
17+// The state transition is done by:
18+// (1) NONE state: The default state when created. This state should only exist
19+// when model_container is created and no constants are being loaded or updated.
20+// (2) INITIALIZED state: This state get set whenever we load the constants into
21+// the buffer. This could be done by load_constants or update_constants_buffer.
22+// (3) FOLDED state: This state should transition from INITIALIZED after
23+// const_fold is being invoked.
24+enum class ConstantState : uint8_t { NONE, INITIALIZED, FOLDED, UNKNOWN };
25+ 
26+inline std::string toStringConstantState(ConstantState state)
27+{
28+ switch (state) {
29+ case ConstantState::NONE:
30+ return "ConstantState::NONE";
31+ case ConstantState::INITIALIZED:
32+ return "ConstantState::INITIALIZED";
33+ case ConstantState::FOLDED:
34+ return "ConstantState::FOLDED";
35+ case ConstantState::UNKNOWN:
36+ return "ConstantState::UNKNOWN";
37+ default:
38+ return "Unknown enum class state for ConstantState";
39+ }
40+}
41+ 
42+class AOTInductorModelContainer {
43+public:
44+ AOTInductorModelContainer(size_t num_models, const std::string& device_str,
45+ const std::optional<std::string>& cubin_dir = std::nullopt)
46+ {
47+ constants_map_ = std::make_shared<ConstantMap>();
48+ constants_array_ = std::make_shared<std::vector<ConstantHandle> >();
49+ 
50+ models_.reserve(num_models);
51+ available_models_.reserve(num_models);
52+ for (size_t i = 0; i < num_models; ++i) {
53+ models_.push_back(AOTInductorModel::Create(constants_map_, constants_array_, device_str, cubin_dir));
54+ available_models_.push_back(models_.back().get());
55+ }
56+ 
57+ // Note that the all following fields (input_names_, output_names,
58+ // etc) can be filled in by the AOT
59+ // codegen. However, we choose to query such information from
60+ // the owned AOTInductorModel for a couple of reasons:
61+ // * simplify the codegen templates
62+ // * reduce information fragmentation and duplication
63+ // * the initialization process below is done only once when the container
64+ // is constructed, so it would have little performance impact
65+ auto* model = available_models_[0];
66+ size_t num_inputs = model->num_inputs();
67+ input_names_.reserve(num_inputs);
68+ for (size_t i = 0; i < num_inputs; i++) {
69+ input_names_.emplace_back(model->input_name(static_cast<int64_t>(i)));
70+ }
71+ 
72+ size_t num_outputs = model->num_outputs();
73+ output_names_.reserve(num_outputs);
74+ for (size_t i = 0; i < num_outputs; i++) {
75+ output_names_.emplace_back(model->output_name(static_cast<int64_t>(i)));
76+ }
77+ model->load_constants();
78+ constant_blob_ = model->release_constant_blob();
79+ constants_internal_offset_.resize(model->num_constants() - model->num_folded_constants());
80+ model->compute_constant_blob(blob_size_, constants_internal_offset_);
81+ constant_folded_ = ConstantState::INITIALIZED;
82+ 
83+ for (auto& model : models_) {
84+ model->update_constants_map(constants_map_);
85+ }
86+ 
87+ in_spec_ = model->get_in_spec();
88+ out_spec_ = model->get_out_spec();
89+ }
90+ 
91+ void run(AtenTensorHandle* input_handles, // array of input AtenTensorHandle; handles
92+ // are stolen; the array itself is borrowed
93+ AtenTensorHandle* output_handles, // array for writing output AtenTensorHandle; handles
94+ // will be stolen by the caller; the array itself is
95+ // borrowed
96+ DeviceStreamType stream, AOTIProxyExecutorHandle proxy_executor)
97+ {
98+ std::shared_lock model_lk(model_exec_mutex_);
99+ auto* model = get_available_model();
100+ 
101+ ConstantState& const_folded = use_secondary_ ? constant_folded_secondary_ : constant_folded_;
102+ if (const_folded == ConstantState::INITIALIZED) {
103+ // At this point, constant is not ready yet. We need to call constant
104+ // folding before we execute the model. We obtain a unique lock at this
105+ // point to make sure constant is ready for all.
106+ model_lk.unlock();
107+ std::unique_lock constants_folding_lk(model_exec_mutex_);
108+ // Double locking to make sure constant folding is only ran once.
109+ if (const_folded == ConstantState::INITIALIZED) {
110+ auto folded_const_map = model->run_const_fold(stream, proxy_executor, true);
111+ update_constant_buffer(std::move(folded_const_map), false, false);
112+ const_folded = ConstantState::FOLDED;
113+ }
114+ constants_folding_lk.unlock();
115+ model_lk.lock();
116+ } else if (const_folded != ConstantState::FOLDED) {
117+ throw std::runtime_error("Unknown constant state: " + toStringConstantState(constant_folded_));
118+ }
119+ 
120+ try {
121+ model->run(input_handles, output_handles, stream, proxy_executor);
122+ } catch (...) {
123+ std::lock_guard lk(models_mutex_);
124+ available_models_.push_back(model);
125+ throw;
126+ }
127+ 
128+ {
129+ std::lock_guard lk(models_mutex_);
130+ pending_models_.push_back(model);
131+ }
132+ pending_models_available_.notify_one();
133+ }
134+ 
135+ // Non-thread-aware variant of run(). Obviously unsafe to use in a threaded
136+ // environment :)
137+ void run_single_threaded(AtenTensorHandle* input_handles, // array of input AtenTensorHandle; handles
138+ // are stolen; the array itself is borrowed
139+ AtenTensorHandle* output_handles, // array for writing output AtenTensorHandle; handles
140+ // will be stolen by the caller; the array itself is
141+ // borrowed
142+ DeviceStreamType stream, AOTIProxyExecutorHandle proxy_executor)
143+ {
144+ auto* model = available_models_[0];
145+ 
146+ ConstantState& const_folded = use_secondary_ ? constant_folded_secondary_ : constant_folded_;
147+ if (const_folded == ConstantState::INITIALIZED) {
148+ auto folded_const_map = model->run_const_fold(stream, proxy_executor, true);
149+ update_constant_buffer(std::move(folded_const_map), false, false);
150+ const_folded = ConstantState::FOLDED;
151+ } else if (constant_folded_ != ConstantState::FOLDED) {
152+ throw std::runtime_error("Unknown constant state: " + toStringConstantState(constant_folded_));
153+ }
154+ 
155+ model->run_single_threaded(input_handles, output_handles, stream, proxy_executor);
156+ }
157+ 
158+ const std::unordered_map<std::string, AtenTensorHandle> extract_constants_map(bool use_inactive) const
159+ {
160+ size_t n_consts = this->num_constants();
161+ std::unordered_map<std::string, AtenTensorHandle> ret;
162+ ret.reserve(n_consts);
163+ 
164+ std::shared_ptr<ConstantMap> extract_map = constants_map_;
165+ // Essentially a XOR
166+ if (use_inactive != use_secondary_) {
167+ extract_map = constants_map_secondary_;
168+ }
169+ for (size_t idx = 0; idx < n_consts; idx++) {
170+ if (this->constant_from_folded(idx)) {
171+ continue;
172+ }
173+ 
174+ auto it = extract_map->find(this->constant_name(idx));
175+ if (it != extract_map->end()) {
176+ ret.emplace(this->constant_original_fqn(idx), it->second);
177+ continue;
178+ }
179+ }
180+ 
181+ return ret;
182+ }
183+ 
184+ size_t num_constants() const
185+ {
186+ if (this->num_models() == 0) {
187+ throw std::runtime_error("No available models in container!");
188+ }
189+ return models_[0]->num_constants();
190+ }
191+ 
192+ // retrieve the constant name of constants_info_[idx]
193+ const char* constant_name(size_t idx) const
194+ {
195+ if (this->num_models() == 0) {
196+ throw std::runtime_error("No available models in container!");
197+ }
198+ return models_[0]->constant_name(static_cast<int64_t>(idx));
199+ }
200+ 
201+ // retrieve original FQN of constants_info_[idx]
202+ const char* constant_original_fqn(size_t idx) const
203+ {
204+ if (this->num_models() == 0) {
205+ throw std::runtime_error("No available models in container!");
206+ }
207+ return models_[0]->constant_original_fqn(static_cast<int64_t>(idx));
208+ }
209+ 
210+ // retrieve whether constant is from folded of constants_info_[idx]
211+ bool constant_from_folded(size_t idx) const
212+ {
213+ if (this->num_models() == 0) {
214+ throw std::runtime_error("No available models in container!");
215+ }
216+ return models_[0]->constant_from_folded(static_cast<int64_t>(idx));
217+ }
218+ 
219+ size_t constant_data_size(size_t idx) const
220+ {
221+ if (this->num_models() == 0) {
222+ throw std::runtime_error("No available models in container!");
223+ }
224+ return models_[0]->constant_data_size(static_cast<int64_t>(idx));
225+ }
226+ 
227+ // retrieve type of constants_info_[idx]
228+ int32_t constant_type(size_t idx) const
229+ {
230+ if (this->num_models() == 0) {
231+ throw std::runtime_error("No available models in container!");
232+ }
233+ return models_[0]->constant_type(static_cast<int64_t>(idx));
234+ }
235+ 
236+ // retrieve dtype of constants_info_[idx]
237+ int32_t constant_dtype(size_t idx) const
238+ {
239+ if (this->num_models() == 0) {
240+ throw std::runtime_error("No available models in container!");
241+ }
242+ return models_[0]->constant_dtype(static_cast<int64_t>(idx));
243+ }
244+ 
245+ void run_const_fold(bool inactive_buffer, DeviceStreamType stream, AOTIProxyExecutorHandle proxy_executor)
246+ {
247+ AOTInductorModel* model;
248+ ConstantState& const_folded = inactive_buffer == use_secondary_ ? constant_folded_ : constant_folded_secondary_;
249+ if (!inactive_buffer) {
250+ // We would need to acquire a unique lock if we want to run constant
251+ // folding on the active buffer.
252+ std::unique_lock constants_folding_lk(model_exec_mutex_);
253+ model = get_available_model();
254+ try {
255+ auto folded_const_map = model->run_const_fold(stream, proxy_executor);
256+ update_constant_buffer(std::move(folded_const_map), false, false);
257+ const_folded = ConstantState::FOLDED;
258+ } catch (...) {
259+ std::lock_guard lk(models_mutex_);
260+ available_models_.push_back(model);
261+ throw;
262+ }
263+ } else {
264+ std::shared_lock model_lk(model_exec_mutex_);
265+ model = get_available_model();
266+ 
267+ // const run.
268+ auto constants_map = get_constants_map(true);
269+ auto constants_array = get_constants_array(true);
270+ 
271+ try {
272+ model->update_constants_map(constants_map, false);
273+ model->update_constants_array(constants_array);
274+ 
275+ auto folded_const_map = model->run_const_fold(stream, proxy_executor);
276+ update_constant_buffer(std::move(folded_const_map), true, false);
277+ 
278+ constants_map = get_constants_map(false);
279+ constants_array = get_constants_array(false);
280+ model->update_constants_map(constants_map, false);
281+ model->update_constants_array(constants_array);
282+ const_folded = ConstantState::FOLDED;
283+ } catch (...) {
284+ std::lock_guard lk(models_mutex_);
285+ available_models_.push_back(model);
286+ throw;
287+ }
288+ }
289+ 
290+ {
291+ std::lock_guard lk(models_mutex_);
292+ pending_models_.push_back(model);
293+ }
294+ pending_models_available_.notify_one();
295+ }
296+ 
297+ bool _is_tensor_constant_type(const size_t idx) const
298+ {
299+ auto constant_type = models_[0]->constant_type(static_cast<int64_t>(idx));
300+ // We should skip constants
301+ return constant_type == ConstantType::TensorConstant;
302+ }
303+ 
304+ bool _is_buffer_type(const size_t idx) const
305+ {
306+ auto constant_type = models_[0]->constant_type(static_cast<int64_t>(idx));
307+ // Buffer can be optionally skipped, so if it not provided by upstream
308+ // services, it is OK to relax the check.
309+ return constant_type == ConstantType::Buffer;
310+ }
311+ 
312+ bool _is_tensor_constant_or_buffer_type(const size_t idx) const
313+ {
314+ return _is_tensor_constant_type(idx) || _is_buffer_type(idx);
315+ }
316+ 
317+ void assert_all_constants(const std::unordered_map<std::string, AtenTensorHandle>& constants_map)
318+ {
319+ auto num_constants = models_[0]->num_constants();
320+ for (size_t idx = 0; idx < num_constants; idx++) {
321+ if (models_[0]->constant_from_folded(static_cast<int64_t>(idx))) {
322+ continue;
323+ }
324+ 
325+ auto constant_name = std::string(models_[0]->constant_name(static_cast<int64_t>(idx)));
326+ auto it = constants_map.find(constant_name);
327+ if (it == constants_map.end()) {
328+ if (_is_tensor_constant_or_buffer_type(idx)) {
329+ // tracing sometimes creates tensors that are non-existent in
330+ // original graph. We could skip those and do a direct copy.
331+ std::cerr << "[WARNING] Found constant or module state buffer " << constant_name
332+ << " in model, but not provided by user!\n";
333+ continue;
334+ }
335+ throw std::runtime_error(std::string("Cannot find constants ") + constant_name +
336+ std::string(" in constants_map!"));
337+ }
338+ }
339+ }
340+ 
341+ // We directly take ownership from AtenTensorHandle if constants are moved.
342+ void update_constant_buffer(std::unordered_map<std::string, AtenTensorHandle>&& constants_map, bool use_inactive,
343+ bool validate_full_update)
344+ {
345+ if (this->num_models() == 0) {
346+ throw std::runtime_error("No model available in container!");
347+ }
348+ if (validate_full_update) {
349+ assert_all_constants(constants_map);
350+ }
351+ 
352+ ConstantState& const_folded = use_inactive == use_secondary_ ? constant_folded_ : constant_folded_secondary_;
353+ const_folded = ConstantState::INITIALIZED;
354+ 
355+ auto original_constants_map = get_constants_map(!use_inactive);
356+ auto constants_map_to_update = get_constants_map(use_inactive);
357+ 
358+ auto num_constants = models_[0]->num_constants();
359+ for (size_t idx = 0; idx < num_constants; idx++) {
360+ auto constant_name = std::string(models_[0]->constant_name(static_cast<int64_t>(idx)));
361+ auto it = constants_map.find(constant_name);
362+ if (it == constants_map.end() && !(use_inactive && _is_tensor_constant_type(idx))) {
363+ continue;
364+ }
365+ 
366+ AtenTensorHandle tensor;
367+ if (it == constants_map.end()) {
368+ aoti_torch_clone(original_constants_map->find(constant_name)->second.get(), &tensor);
369+ } else {
370+ tensor = it->second;
371+ }
372+ 
373+ constants_map_to_update->insert_or_assign(constant_name, RAIIAtenTensorHandle(tensor));
374+ }
375+ // Update the inactive constant array.
376+ update_array_from_map(get_constants_array(use_inactive), constants_map_to_update);
377+ }
378+ 
379+ // This function updates the buffer for storing constants.
380+ void update_constant_buffer(const std::unordered_map<std::string, AtenTensorHandle>& constants_map,
381+ bool use_inactive, bool validate_full_update, bool user_managed = false)
382+ {
383+ if (this->num_models() == 0) {
384+ throw std::runtime_error("No model available in container!");
385+ }
386+ if (validate_full_update) {
387+ assert_all_constants(constants_map);
388+ }
389+ 
390+ ConstantState& const_folded = use_inactive == use_secondary_ ? constant_folded_ : constant_folded_secondary_;
391+ const_folded = ConstantState::INITIALIZED;
392+ 
393+ auto original_constants_map = get_constants_map(!use_inactive);
394+ auto constants_map_to_update = get_constants_map(use_inactive);
395+ 
396+ auto num_constants = models_[0]->num_constants();
397+ for (size_t idx = 0; idx < num_constants; idx++) {
398+ auto constant_name = std::string(models_[0]->constant_name(static_cast<int64_t>(idx)));
399+ auto it = constants_map.find(constant_name);
400+ if (it == constants_map.end() && !(use_inactive && _is_tensor_constant_or_buffer_type(idx))) {
401+ continue;
402+ }
403+ 
404+ AtenTensorHandle tensor;
405+ if (it == constants_map.end()) {
406+ tensor = original_constants_map->find(constant_name)->second.get();
407+ } else {
408+ tensor = it->second;
409+ }
410+ 
411+ if (user_managed) {
412+ // If user managed, we pass in the pointer directly, and skip the
413+ // copy.
414+ constants_map_to_update->insert_or_assign(
415+ constant_name, MaybeOwningAtenTensorHandle(tensor, true));
416+ continue;
417+ }
418+ 
419+ auto* constants_blob_ptr = static_cast<uint8_t*>(get_constant_blob_ptr(use_inactive));
420+ 
421+ // Move the data to container handled blob.
422+ uint8_t* internal_constants_ptr = constants_blob_ptr + constants_internal_offset_[idx];
423+ void* user_constant_ptr;
424+ int64_t constant_size;
425+ aoti_torch_get_data_ptr(tensor, &user_constant_ptr);
426+ aoti_torch_get_storage_size(tensor, &constant_size);
427+#if defined(USE_NPU)
428+ AOTI_RUNTIME_DEVICE_CHECK(aclrtMemcpy(internal_constants_ptr, constant_size, user_constant_ptr,
429+ constant_size, ACL_MEMCPY_HOST_TO_DEVICE));
430+#else
431+ memcpy(internal_constants_ptr, user_constant_ptr, constant_size);
432+#endif
433+ // Generate Tensor from container handled blob.
434+ // We extract stride and offset from provided Tensor since we do not
435+ // guarantee that the tensor is contiguous.
436+ AtenTensorHandle tensor_handle;
437+ int64_t* stride;
438+ int64_t offset;
439+ int device_type = models_[0]->get_device_type();
440+ int device_idx = models_[0]->get_device_idx();
441+ AOTI_TORCH_ERROR_CODE_CHECK(aoti_torch_get_strides(tensor, &stride));
442+ AOTI_TORCH_ERROR_CODE_CHECK(aoti_torch_get_storage_offset(tensor, &offset));
443+ AOTI_TORCH_ERROR_CODE_CHECK(aoti_torch_create_tensor_from_blob(
444+ internal_constants_ptr, models_[0]->constant_ndim(idx), models_[0]->constant_shape(idx), stride, offset,
445+ models_[0]->constant_dtype(idx), device_type, device_idx, &tensor_handle));
446+ 
447+ // Now place the tensor to constants_map. Note at this point the
448+ // ownership of the tensor_handle will be taken over.
449+ constants_map_to_update->insert_or_assign(constant_name, RAIIAtenTensorHandle(tensor_handle));
450+ }
451+ // Update the inactive constant array.
452+ update_array_from_map(get_constants_array(use_inactive), constants_map_to_update);
453+ }
454+ 
455+ void update_array_from_map(const std::shared_ptr<std::vector<ConstantHandle> >& constants_array,
456+ const std::shared_ptr<ConstantMap>& constants_map)
457+ {
458+ auto num_constants = models_[0]->num_constants();
459+ for (size_t idx = 0; idx < num_constants; idx++) {
460+ if (constants_map->find(models_[0]->constant_name(static_cast<int64_t>(idx))) != constants_map->end()) {
461+ constants_array->at(idx) =
462+ ConstantHandle(constants_map->find(models_[0]->constant_name(static_cast<int64_t>(idx)))->second);
463+ }
464+ }
465+ }
466+ 
467+ void swap_constant_buffer()
468+ {
469+ std::lock_guard unique_lk(model_exec_mutex_);
470+ 
471+ auto constants_map = get_constants_map(true);
472+ auto constants_array = get_constants_array(true);
473+ 
474+ for (auto& model : models_) {
475+ model->update_constants_map(constants_map, false);
476+ model->update_constants_array(constants_array);
477+ }
478+ 
479+ use_secondary_ = !use_secondary_;
480+ }
481+ 
482+ void free_inactive_constant_buffer()
483+ {
484+ if (use_secondary_) {
485+ constant_folded_ = ConstantState::NONE;
486+ constant_blob_.reset();
487+ } else {
488+ constant_folded_secondary_ = ConstantState::NONE;
489+ constant_blob_secondary_.reset();
490+ }
491+ // Free the internally held constants
492+ int num_constants = static_cast<int>(models_[0]->num_constants());
493+ std::shared_ptr<ConstantMap> to_free_map = use_secondary_ ? constants_map_ : constants_map_secondary_;
494+ 
495+ for (int i = 0; i < num_constants; i++) {
496+ if (models_[0]->constant_from_folded(i)) {
497+ auto it = to_free_map->find(models_[0]->constant_name(i));
498+ if (it != to_free_map->end()) {
499+ it->second.reset();
500+ }
501+ }
502+ }
503+ }
504+ 
505+ size_t num_inputs() const { return input_names_.size(); }
506+ 
507+ size_t num_outputs() const { return output_names_.size(); }
508+ 
509+ const char* input_name(size_t idx) const { return input_names_.at(idx).c_str(); }
510+ 
511+ const char* output_name(size_t idx) const { return output_names_.at(idx).c_str(); }
512+ 
513+ size_t num_models() const { return models_.size(); }
514+ 
515+ const char* get_in_spec() const { return in_spec_; }
516+ 
517+ const char* get_out_spec() const { return out_spec_; }
518+ 
519+private:
520+ std::vector<std::string> input_names_;
521+ std::vector<std::string> output_names_;
522+ const char* in_spec_;
523+ const char* out_spec_;
524+ 
525+ // Holds the blob storage for constants' at::Tensor within the container.
526+ // This blob of memory will be managed by the container.
527+ RAIIDataPtr constant_blob_;
528+ RAIIDataPtr constant_blob_secondary_;
529+ 
530+ size_t blob_size_;
531+ std::vector<size_t> constants_internal_offset_;
532+ 
533+ // Determine which constants is being used for the model.
534+ // If true,
535+ // constants_map_secondary/constant_blob_secondary/constants_array_secondary
536+ // is being used.
537+ bool use_secondary_{false};
538+ 
539+ // Determine whether we have ran constant folding
540+ ConstantState constant_folded_{ConstantState::NONE};
541+ ConstantState constant_folded_secondary_{ConstantState::NONE};
542+ 
543+ // The underlying data of at::Tensor is in either constant_blob_ (for CUDA).
544+ // or _binary_constants_bin_start (for CPU).
545+ std::shared_ptr<ConstantMap> constants_map_;
546+ std::shared_ptr<ConstantMap> constants_map_secondary_;
547+ 
548+ // Holds the indexed array of constant for faster lookup during runtime.
549+ std::shared_ptr<std::vector<ConstantHandle> > constants_array_;
550+ std::shared_ptr<std::vector<ConstantHandle> > constants_array_secondary_;
551+ 
552+ // Holds all the AOTInductorModel instances owned by this container.
553+ std::vector<std::unique_ptr<AOTInductorModel> > models_;
554+ 
555+ // Holds the AOTInductorModel instances available for inference.
556+ std::vector<AOTInductorModel*> available_models_;
557+ 
558+ // Holds the AOTInductorModel instances that have started running
559+ // inference and can be placed onto available_models_ upon their
560+ // completion.
561+ std::deque<AOTInductorModel*> pending_models_;
562+ 
563+ // Protects available_models_ and pending_models_.
564+ std::mutex models_mutex_;
565+ 
566+ // Notified whenever a model is placed onto pending_models_.
567+ std::condition_variable pending_models_available_;
568+ 
569+ AOTInductorModel* get_available_model()
570+ {
571+ std::unique_lock lk(models_mutex_);
572+ if (available_models_.empty()) {
573+ reclaim_finished_models(lk);
574+ }
575+ auto* result = available_models_.back();
576+ available_models_.pop_back();
577+ return result;
578+ }
579+ 
580+ // This mutex is used to protect execution of model.
581+ // We acquire the mutex in shared mode if we allow concurrent execution.
582+ // We acquire the mutex in unique mode when we want exclusive access of the
583+ // make sure no one is executing the model.
584+ std::shared_mutex model_exec_mutex_;
585+ 
586+ RAIIDataPtr allocate_constant_blob()
587+ {
588+#if defined(USE_NPU)
589+ return RAII_npuMalloc(blob_size_);
590+#else
591+ return RAII_cpuMalloc(blob_size_);
592+#endif // USE_CUDA
593+ }
594+ 
595+ void* get_constant_blob_ptr(bool get_inactive)
596+ {
597+ if ((get_inactive && use_secondary_) || (!get_inactive && !use_secondary_)) {
598+ if (!constant_blob_) {
599+ constant_blob_ = allocate_constant_blob();
600+ }
601+ return constant_blob_.get();
602+ } else {
603+ if (!constant_blob_secondary_) {
604+ constant_blob_secondary_ = allocate_constant_blob();
605+ }
606+ return constant_blob_secondary_.get();
607+ }
608+ }
609+ 
610+ std::shared_ptr<ConstantMap> get_constants_map(bool get_inactive)
611+ {
612+ if ((get_inactive && use_secondary_) || (!get_inactive && !use_secondary_)) {
613+ return constants_map_;
614+ } else {
615+ if (!constants_map_secondary_) {
616+ constants_map_secondary_ = std::make_shared<ConstantMap>();
617+ }
618+ return constants_map_secondary_;
619+ }
620+ }
621+ 
622+ std::shared_ptr<std::vector<ConstantHandle> > get_constants_array(bool get_inactive)
623+ {
624+ if ((get_inactive && use_secondary_) || (!get_inactive && !use_secondary_)) {
625+ return constants_array_;
626+ } else {
627+ if (!constants_array_secondary_) {
628+ constants_array_secondary_ =
629+ std::make_shared<std::vector<ConstantHandle> >(models_[0]->num_constants());
630+ }
631+ return constants_array_secondary_;
632+ }
633+ }
634+ 
635+ void reclaim_finished_models(std::unique_lock<std::mutex>& lk)
636+ {
637+#ifdef __aarch64__
638+ // push finished model instances to the end of pending_models_
639+ auto it = std::partition(pending_models_.begin(), pending_models_.end(),
640+ [](AOTInductorModel* m) { return !m->is_finished(); });
641+#else
642+ // push finished model instances to the end of pending_models_
643+ auto it = std::stable_partition(pending_models_.begin(), pending_models_.end(),
644+ [](AOTInductorModel* m) { return !m->is_finished(); });
645+#endif
646+ if (it != pending_models_.end()) {
647+ // We have finished model instances that can be pushed into
648+ // available_models_ so that we don't have to be blocked on waiting
649+ // the pending_models_available_ condition.
650+ available_models_.insert(available_models_.end(), it, pending_models_.end());
651+ pending_models_.erase(it, pending_models_.end());
652+ return;
653+ }
654+ 
655+ pending_models_available_.wait(lk, [this]() { return !pending_models_.empty(); });
656+ // Let's make the schedule simple first. We always wait on the first
657+ // pending_models_ to be complete.
658+ auto* model = pending_models_.front();
659+ pending_models_.pop_front();
660+ lk.unlock();
661+ try {
662+ model->wait_for_completion();
663+ } catch (...) {
664+ lk.lock();
665+ available_models_.push_back(model);
666+ throw;
667+ }
668+ lk.lock();
669+ available_models_.push_back(model);
670+ }
671+};
672+ 
673+} // namespace torch::aot_inductor
Atorch_npu/csrc/inductor/aoti_runtime/scalar_to_tensor.h+37-0
@@ -0,0 +1,37 @@
1+#pragma once
2+ 
3+#include <c10/util/complex.h>
4+#include <torch_npu/csrc/inductor/aoti_runtime/utils.h>
5+ 
6+namespace torch::aot_inductor {
7+ 
8+template <typename T> inline RAIIAtenTensorHandle scalar_to_tensor_handle(T value)
9+{
10+ throw std::runtime_error("Unsupported scalar_to_tensor_handle");
11+}
12+ 
13+// Specialize for supported C++ primitive types
14+#define AOTI_RUNTIME_SCALAR_TO_TENSOR(dtype, ctype) \
15+ template <> inline RAIIAtenTensorHandle scalar_to_tensor_handle<ctype>(ctype value) \
16+ { \
17+ AtenTensorHandle tensor_handle; \
18+ AOTI_TORCH_ERROR_CODE_CHECK(aoti_torch_scalar_to_tensor_##dtype(value, &tensor_handle)); \
19+ return RAIIAtenTensorHandle(tensor_handle); \
20+ }
21+ 
22+AOTI_RUNTIME_SCALAR_TO_TENSOR(float32, float)
23+AOTI_RUNTIME_SCALAR_TO_TENSOR(float64, double)
24+AOTI_RUNTIME_SCALAR_TO_TENSOR(uint8, uint8_t)
25+AOTI_RUNTIME_SCALAR_TO_TENSOR(uint16, uint16_t)
26+AOTI_RUNTIME_SCALAR_TO_TENSOR(uint32, uint32_t)
27+AOTI_RUNTIME_SCALAR_TO_TENSOR(uint64, uint64_t)
28+AOTI_RUNTIME_SCALAR_TO_TENSOR(int8, int8_t)
29+AOTI_RUNTIME_SCALAR_TO_TENSOR(int16, int16_t)
30+AOTI_RUNTIME_SCALAR_TO_TENSOR(int32, int32_t)
31+AOTI_RUNTIME_SCALAR_TO_TENSOR(int64, int64_t)
32+AOTI_RUNTIME_SCALAR_TO_TENSOR(bool, bool)
33+AOTI_RUNTIME_SCALAR_TO_TENSOR(complex64, c10::complex<float>)
34+AOTI_RUNTIME_SCALAR_TO_TENSOR(complex128, c10::complex<double>)
35+#undef AOTI_RUNTIME_SCALAR_TO_TENSOR
36+ 
37+} // namespace torch::aot_inductor
Atorch_npu/csrc/inductor/aoti_runtime/thread_local.h+115-0
@@ -0,0 +1,115 @@
1+#pragma once
2+ 
3+#include <torch_npu/csrc/inductor/aoti_runtime/arrayref_tensor.h>
4+ 
5+namespace torch::aot_inductor {
6+ 
7+template <typename T> struct ThreadLocalCachedOutputTensor;
8+ 
9+template <> struct ThreadLocalCachedOutputTensor<RAIIAtenTensorHandle> {
10+ explicit ThreadLocalCachedOutputTensor(const RAIIAtenTensorHandle&) {}
11+ void copy_data_from(const RAIIAtenTensorHandle& handle) { throw std::runtime_error("can't happen"); }
12+ 
13+ AtenTensorHandle tensor() const { throw std::runtime_error("can't happen"); }
14+};
15+ 
16+template <> struct ThreadLocalCachedOutputTensor<AtenTensorHandle> {
17+ explicit ThreadLocalCachedOutputTensor(const AtenTensorHandle&) {}
18+ void copy_data_from(const AtenTensorHandle& handle) { throw std::runtime_error("can't happen"); }
19+ 
20+ AtenTensorHandle tensor() const { throw std::runtime_error("can't happen"); }
21+};
22+ 
23+template <> struct ThreadLocalCachedOutputTensor<ConstantHandle> {
24+ explicit ThreadLocalCachedOutputTensor(const ConstantHandle&) {}
25+ void copy_data_from(const ConstantHandle& handle) { throw std::runtime_error("can't happen"); }
26+ 
27+ AtenTensorHandle tensor() const { throw std::runtime_error("can't happen"); }
28+};
29+ 
30+template <typename T> struct ThreadLocalCachedOutputTensor<ArrayRefTensor<T> > {
31+ explicit ThreadLocalCachedOutputTensor(const ArrayRefTensor<T>& t) { re_alloc_tensor(t); }
32+ 
33+ void copy_data_from(const ArrayRefTensor<T>& t)
34+ {
35+ if (t.numel() > capacity_) {
36+ re_alloc_tensor(t);
37+ }
38+ std::copy(t.data(), t.data() + t.numel(), storage_.get());
39+ }
40+ 
41+ AtenTensorHandle tensor() const { return tensor_.get(); }
42+ 
43+private:
44+ void re_alloc_tensor(const ArrayRefTensor<T>& t)
45+ {
46+ capacity_ = t.numel();
47+ // NOLINTNEXTLINE(*arrays*)
48+ storage_ = std::make_unique<T[]>(t.numel());
49+ AtenTensorHandle handle = nullptr;
50+ AOTI_TORCH_ERROR_CODE_CHECK(aoti_torch_create_tensor_from_blob_npu(
51+ storage_.get(), t.sizes().size(), t.sizes().data(), t.strides().data(), 0,
52+ aoti_torch_dtype<std::remove_const_t<T> >(), t.device_type(), t.device_idx(), &handle));
53+ tensor_ = handle;
54+ }
55+ 
56+ // NOLINTNEXTLINE(*arrays*)
57+ std::unique_ptr<T[]> storage_;
58+ int64_t capacity_ = 0;
59+ RAIIAtenTensorHandle tensor_;
60+};
61+ 
62+template <typename T> struct ThreadLocalCachedOutputArray;
63+ 
64+// Just needs to compile, doesn't need to do anything.
65+template <> struct ThreadLocalCachedOutputArray<RAIIAtenTensorHandle> {
66+ explicit ThreadLocalCachedOutputArray(const RAIIAtenTensorHandle&) { throw std::runtime_error("can't happen"); }
67+ 
68+ // Not supported yet! We would need to put contiguous() or
69+ // expect_contiguous() into the ABI.
70+ void copy_data_from(const RAIIAtenTensorHandle&) { throw std::runtime_error("can't happen"); }
71+ 
72+ template <typename U> ArrayRefTensor<U> arrayref_tensor() const { throw std::runtime_error("can't happen"); }
73+};
74+ 
75+// Just needs to compile, doesn't need to do anything.
76+template <> struct ThreadLocalCachedOutputArray<ConstantHandle> {
77+ explicit ThreadLocalCachedOutputArray(const ConstantHandle&) { throw std::runtime_error("can't happen"); }
78+ 
79+ // Not supported yet! We would need to put contiguous() or
80+ // expect_contiguous() into the ABI.
81+ void copy_data_from(const ConstantHandle&) { throw std::runtime_error("can't happen"); }
82+ 
83+ template <typename U> ArrayRefTensor<U> arrayref_tensor() const { throw std::runtime_error("can't happen"); }
84+};
85+ 
86+template <typename T> struct ThreadLocalCachedOutputArray<ArrayRefTensor<T> > {
87+ explicit ThreadLocalCachedOutputArray(const ArrayRefTensor<T>& t) {}
88+ 
89+ template <typename U,
90+ std::enable_if_t<std::is_same_v<std::remove_const_t<T>, std::remove_const_t<U> >, bool> = true>
91+ ArrayRefTensor<T> arrayref_tensor() const
92+ {
93+ return tensor_;
94+ }
95+ 
96+ void copy_data_from(const ArrayRefTensor<T>& t)
97+ {
98+ if (t.numel() > capacity_) {
99+ capacity_ = t.numel();
100+ // NOLINTNEXTLINE(*arrays*)
101+ storage_ = std::make_unique<T[]>(capacity_);
102+ }
103+ std::copy(t.data(), t.data() + t.numel(), storage_.get());
104+ tensor_ = t;
105+ tensor_.set_arrayref(MiniArrayRef<T>(storage_.get(), t.numel()));
106+ }
107+ 
108+private:
109+ // NOLINTNEXTLINE(*arrays*)
110+ std::unique_ptr<T[]> storage_;
111+ uint32_t capacity_ = 0;
112+ ArrayRefTensor<T> tensor_;
113+};
114+ 
115+} // namespace torch::aot_inductor
Atorch_npu/csrc/inductor/aoti_runtime/utils.h+338-0
@@ -0,0 +1,338 @@
1+#pragma once
2+ 
3+#include <iostream>
4+#include <memory>
5+#include <sstream>
6+#include <stdexcept>
7+#include <string>
8+#include <vector>
9+ 
10+// WARNING: Be careful when adding new includes here. This header will be used
11+// in model.so, and should not refer to any aten/c10 headers except the stable
12+// C ABI defined in torch/csrc/inductor/aoti_torch/c/shim.h. The same rule
13+// applies to other files under torch/csrc/inductor/aoti_runtime/.
14+#include <torch_npu/csrc/inductor/aoti_torch/c/shim.h>
15+ 
16+#if defined(__GNUC__) || defined(__clang__)
17+#define AOTI_NOINLINE __attribute__((noinline))
18+#elif _MSC_VER
19+#define AOTI_NOINLINE __declspec(noinline)
20+#else
21+#define AOTI_NOINLINE
22+#endif
23+ 
24+AOTI_NOINLINE static void throw_exception(const char* call, const char* file, int64_t line)
25+{
26+ std::stringstream ss;
27+ ss << call << " API call failed at " << file << ", line " << line;
28+ throw std::runtime_error(ss.str());
29+}
30+ 
31+#define AOTI_TORCH_ERROR_CODE_CHECK(call) \
32+ if ((call) != AOTI_TORCH_SUCCESS) { \
33+ throw_exception(#call, __FILE__, __LINE__); \
34+ }
35+ 
36+using AOTIRuntimeError = int32_t;
37+#define AOTI_RUNTIME_SUCCESS 0
38+#define AOTI_RUNTIME_FAILURE 1
39+ 
40+#define AOTI_RUNTIME_ERROR_CODE_CHECK(call) \
41+ if ((call) != AOTI_RUNTIME_SUCCESS) { \
42+ throw_exception(#call, __FILE__, __LINE__); \
43+ }
44+ 
45+namespace torch::aot_inductor {
46+ 
47+using DeleterFnPtr = void (*)(void*);
48+ 
49+inline void noop_deleter(void*) {}
50+ 
51+inline void delete_tensor_object(void* ptr)
52+{
53+ AOTI_TORCH_ERROR_CODE_CHECK(aoti_torch_delete_tensor_object(reinterpret_cast<AtenTensorHandle>(ptr)));
54+}
55+ 
56+// RAIIAtenTensorHandle steals the tensor objects created by the libtorch C ABI
57+class RAIIAtenTensorHandle {
58+public:
59+ RAIIAtenTensorHandle() : handle_(nullptr, noop_deleter) {}
60+ RAIIAtenTensorHandle(const RAIIAtenTensorHandle& other) = delete;
61+ RAIIAtenTensorHandle& operator=(const RAIIAtenTensorHandle& other) = delete;
62+ 
63+ // Steal the ownership from another RAIIAtenTensorHandle using std::move
64+ RAIIAtenTensorHandle(RAIIAtenTensorHandle&& other) = default;
65+ RAIIAtenTensorHandle& operator=(RAIIAtenTensorHandle&& other) = default;
66+ 
67+ // Steal the ownership from raw AtenTensorHandle
68+ RAIIAtenTensorHandle(AtenTensorHandle handle) : handle_(handle, delete_tensor_object) {}
69+ 
70+ ~RAIIAtenTensorHandle() { handle_.reset(); }
71+ 
72+ // Return a raw AtenTensorHandle to be used by aoti_torch functions
73+ // Note: this function does NOT transfer the ownership of the handle
74+ operator AtenTensorHandle() const { return handle_.get(); }
75+ 
76+ AtenTensorHandle release() { return handle_.release(); }
77+ 
78+ AtenTensorHandle get() const { return handle_.get(); }
79+ 
80+ void reset() { handle_.reset(); }
81+ 
82+ int64_t size(int64_t d)
83+ {
84+ int64_t size = 0;
85+ AOTI_TORCH_ERROR_CODE_CHECK(aoti_torch_get_size(handle_.get(), d, &size));
86+ return size;
87+ }
88+ 
89+ int64_t stride(int64_t d)
90+ {
91+ int64_t stride = 0;
92+ AOTI_TORCH_ERROR_CODE_CHECK(aoti_torch_get_stride(handle_.get(), d, &stride));
93+ return stride;
94+ }
95+ 
96+ int64_t storage_offset()
97+ {
98+ int64_t storage_offset = 0;
99+ AOTI_TORCH_ERROR_CODE_CHECK(aoti_torch_get_storage_offset(handle_.get(), &storage_offset));
100+ return storage_offset;
101+ }
102+ 
103+ void* data_ptr() const
104+ {
105+ void* result = nullptr;
106+ AOTI_TORCH_ERROR_CODE_CHECK(aoti_torch_get_data_ptr(handle_.get(), &result));
107+ return result;
108+ }
109+ 
110+ int64_t* sizes() const
111+ {
112+ int64_t* result = nullptr;
113+ AOTI_TORCH_ERROR_CODE_CHECK(aoti_torch_get_sizes(handle_.get(), &result));
114+ return result;
115+ }
116+ 
117+ int64_t* strides() const
118+ {
119+ int64_t* result = nullptr;
120+ AOTI_TORCH_ERROR_CODE_CHECK(aoti_torch_get_strides(handle_.get(), &result));
121+ return result;
122+ }
123+ 
124+private:
125+ std::unique_ptr<AtenTensorOpaque, DeleterFnPtr> handle_;
126+};
127+ 
128+class MaybeOwningAtenTensorHandle {
129+public:
130+ MaybeOwningAtenTensorHandle() : handle_(nullptr), raii_handle_() {}
131+ // We skip copy constructor as MaybeOwningAtenTensorHandle might be RAII which
132+ // makes it undefined.
133+ MaybeOwningAtenTensorHandle(const MaybeOwningAtenTensorHandle& other) = delete;
134+ MaybeOwningAtenTensorHandle& operator=(const MaybeOwningAtenTensorHandle& other) = delete;
135+ 
136+ // Move constructor and move assignment operator
137+ MaybeOwningAtenTensorHandle(MaybeOwningAtenTensorHandle&& other) = default;
138+ MaybeOwningAtenTensorHandle& operator=(MaybeOwningAtenTensorHandle&& other) = default;
139+ 
140+ // Steal the ownership from another RAIIAtenTensorHandle using std::move
141+ MaybeOwningAtenTensorHandle(RAIIAtenTensorHandle&& other) : raii_handle_(std::move(other))
142+ {
143+ handle_ = raii_handle_.get();
144+ }
145+ MaybeOwningAtenTensorHandle& operator=(RAIIAtenTensorHandle&& other)
146+ {
147+ raii_handle_ = std::move(other);
148+ handle_ = raii_handle_.get();
149+ return *this;
150+ }
151+ 
152+ // By default, steal the ownership from raw AtenTensorHandle
153+ MaybeOwningAtenTensorHandle(AtenTensorHandle handle) : raii_handle_(handle) { handle_ = raii_handle_.get(); }
154+ 
155+ // If user_managed is true, we do not steal the ownership.
156+ MaybeOwningAtenTensorHandle(AtenTensorHandle handle, bool user_managed)
157+ {
158+ if (user_managed) {
159+ aoti_torch_new_tensor_handle(handle, &handle_);
160+ } else {
161+ raii_handle_ = RAIIAtenTensorHandle(handle);
162+ handle_ = raii_handle_.get();
163+ }
164+ }
165+ 
166+ ~MaybeOwningAtenTensorHandle()
167+ {
168+ // This is no-op if we don't hold raii_handle with the
169+ // MaybeOwningAtenTensorHandle.
170+ raii_handle_.reset();
171+ }
172+ 
173+ // Return a raw AtenTensorHandle to be used by aoti_torch functions
174+ // Note: this function does NOT transfer the ownership of the handle
175+ operator AtenTensorHandle() const { return handle_; }
176+ 
177+ AtenTensorHandle release()
178+ {
179+ if (raii_handle_) {
180+ return raii_handle_.release();
181+ } else {
182+ AtenTensorHandle handle = handle_;
183+ handle_ = nullptr;
184+ return handle;
185+ }
186+ }
187+ 
188+ AtenTensorHandle get() const { return handle_; }
189+ 
190+ void reset()
191+ {
192+ handle_ = nullptr;
193+ raii_handle_.reset();
194+ }
195+ 
196+ int64_t size(int64_t d)
197+ {
198+ int64_t size = 0;
199+ AOTI_TORCH_ERROR_CODE_CHECK(aoti_torch_get_size(handle_, d, &size));
200+ return size;
201+ }
202+ 
203+ int64_t stride(int64_t d)
204+ {
205+ int64_t stride = 0;
206+ AOTI_TORCH_ERROR_CODE_CHECK(aoti_torch_get_stride(handle_, d, &stride));
207+ return stride;
208+ }
209+ 
210+ int64_t storage_offset()
211+ {
212+ int64_t storage_offset = 0;
213+ AOTI_TORCH_ERROR_CODE_CHECK(aoti_torch_get_storage_offset(handle_, &storage_offset));
214+ return storage_offset;
215+ }
216+ 
217+ void* data_ptr() const
218+ {
219+ void* result = nullptr;
220+ AOTI_TORCH_ERROR_CODE_CHECK(aoti_torch_get_data_ptr(handle_, &result));
221+ return result;
222+ }
223+ 
224+ int64_t* sizes() const
225+ {
226+ int64_t* result = nullptr;
227+ AOTI_TORCH_ERROR_CODE_CHECK(aoti_torch_get_sizes(handle_, &result));
228+ return result;
229+ }
230+ 
231+ int64_t* strides() const
232+ {
233+ int64_t* result = nullptr;
234+ AOTI_TORCH_ERROR_CODE_CHECK(aoti_torch_get_strides(handle_, &result));
235+ return result;
236+ }
237+ 
238+private:
239+ // handle_ is the underlying AtenTensorHandle of raii_handle_ if raii_handle_
240+ // exists. Otherwise it would just be the AtenTensorHandle passed in by users.
241+ AtenTensorHandle handle_;
242+ RAIIAtenTensorHandle raii_handle_;
243+};
244+ 
245+// Steal the ownership from raw AtenTensorHandle to RAIIAtenTensorHandle
246+inline std::vector<RAIIAtenTensorHandle> steal_from_raw_handles_to_raii_handles(AtenTensorHandle* handles, size_t size)
247+{
248+ std::vector<RAIIAtenTensorHandle> result;
249+ result.reserve(size);
250+ for (size_t i = 0; i < size; i++) {
251+ result.emplace_back(handles[i]);
252+ handles[i] = nullptr;
253+ }
254+ return result;
255+}
256+ 
257+inline AtenTensorHandle reinterpret_tensor_wrapper(AtenTensorHandle self, int64_t ndim, const int64_t* sizes_ptr,
258+ const int64_t* strides_ptr, int64_t storage_offset)
259+{
260+ AtenTensorHandle result = nullptr;
261+ AOTI_TORCH_ERROR_CODE_CHECK(
262+ aoti_torch__reinterpret_tensor(self, ndim, sizes_ptr, strides_ptr, storage_offset, &result));
263+ return result;
264+}
265+ 
266+inline void* get_data_ptr_wrapper(AtenTensorHandle tensor)
267+{
268+ void* result = nullptr;
269+ AOTI_TORCH_ERROR_CODE_CHECK(aoti_torch_get_data_ptr(tensor, &result));
270+ return result;
271+}
272+ 
273+inline AtenTensorHandle unwrap_raii_handle_if_needed(const RAIIAtenTensorHandle& handle) { return handle.get(); }
274+ 
275+inline RAIIAtenTensorHandle wrap_with_raii_handle_if_needed(AtenTensorHandle handle)
276+{
277+ return RAIIAtenTensorHandle(handle);
278+}
279+ 
280+class ConstantHandle {
281+public:
282+ ConstantHandle() = default;
283+ 
284+ explicit ConstantHandle(AtenTensorHandle handle) : handle_(handle)
285+ {
286+ AOTI_TORCH_ERROR_CODE_CHECK(aoti_torch_get_data_ptr(handle_, &data_));
287+ }
288+ 
289+ operator AtenTensorHandle() const { return handle_; }
290+ 
291+ AtenTensorHandle tensor() const { return handle_; }
292+ 
293+ AtenTensorHandle get() const { return handle_; }
294+ 
295+ void* data_ptr() const { return data_; }
296+ 
297+ int64_t* sizes() const
298+ {
299+ int64_t* result = nullptr;
300+ AOTI_TORCH_ERROR_CODE_CHECK(aoti_torch_get_sizes(handle_, &result));
301+ return result;
302+ }
303+ 
304+ int64_t* strides() const
305+ {
306+ int64_t* result = nullptr;
307+ AOTI_TORCH_ERROR_CODE_CHECK(aoti_torch_get_strides(handle_, &result));
308+ return result;
309+ }
310+ 
311+private:
312+ AtenTensorHandle handle_{};
313+ void* data_ = nullptr;
314+};
315+ 
316+inline void* get_data_ptr_wrapper(const ConstantHandle& constant) { return constant.data_ptr(); }
317+ 
318+inline const ConstantHandle& unwrap_raii_handle_if_needed(const ConstantHandle& handle) { return handle; }
319+ 
320+// Shouldn't be called.
321+inline AtenTensorHandle wrap_with_raii_handle_if_needed(const ConstantHandle& handle) = delete;
322+ 
323+// DANGEROUS. Do not call unless you explicitly intend to get a reference to a
324+// temporary value, which will expire at the end of the current expression.
325+// This should only be called in cases where the C-shim API expects an optional
326+// input argument (passed by pointer), and a temporary needs to be passed to it.
327+template <class T> T& temporary_reference(T&& t) { return t; }
328+ 
329+#define CACHE_TORCH_DTYPE(typename) static auto cached_torch_dtype_##typename = aoti_torch_dtype_##typename()
330+ 
331+#define CACHE_TORCH_DEVICE(device) static auto cached_torch_device_type_##device = aoti_torch_device_type_##device()
332+ 
333+#define CACHE_TORCH_LAYOUT(layout) static auto cached_torch_layout_##layout = aoti_torch_layout_##layout()
334+ 
335+#define CACHE_TORCH_MEMORY_FORMAT(format) \
336+ static auto cached_torch_memory_format_##format = aoti_torch_memory_format_##format()
337+ 
338+} // namespace torch::aot_inductor