Star864
1362
代码
代码
Issues1734
Pull Requests2257
项目讨论
Wiki
分析
Star864
1362
已合并
aoti_load_package接口实现 #35687
qiaoyaodan创建于 5月14日
aoti_load_package接口实现 #35687
已合并
qiaoyaodan创建于 5月14日
从qiaoyaodan888/pytorchaoti_qyd合入到Ascend/pytorchv2.7.1
共 5 个文件变更+218-13
Mbuild_libtorch_npu.py+5-2
@@ -284,11 +284,14 @@ def copy_hpp():
284 "third_party/acl/inc/*/*/*.h",284 "third_party/acl/inc/*/*/*.h",
285 "third_party/hccl/inc/*/*.h",285 "third_party/hccl/inc/*/*.h",
286 ]286 ]
287- glob_header_files = []287+ glob_header_files_old = []
288 for regex_pattern in header_files:288 for regex_pattern in header_files:
289- glob_header_files += glob.glob(289+ glob_header_files_old += glob.glob(
290 os.path.join(BASE_DIR, regex_pattern), recursive=True290 os.path.join(BASE_DIR, regex_pattern), recursive=True
291 )291 )
292+
293+ suffixes = ("/torch_npu/csrc/inductor/aoti_package/pybind.h", "/torch_npu/csrc/inductor/aoti_runner/pybind.h")
294+ glob_header_files = [path for path in glob_header_files_old if not path.endswith(suffixes)]
292 295 
293 for src in glob_header_files:296 for src in glob_header_files:
294 dst = os.path.join(297 dst = os.path.join(
Atorch_npu/csrc/inductor/aoti_package/pybind.cpp+70-0
@@ -0,0 +1,70 @@
1+#ifndef BUILD_LIBTORCH
2+#ifdef USE_NPU
Z
Zzhucehw5月16日

完善aoti的测试用例，用来看护该特性

likedislike
qiaoyaodan
qiaoyaodan
5月19日 评论：
3+#include <torch_npu/csrc/inductor/aoti_runner/model_container_runner_npu.h>
4+#endif
5+ 
6+#include <torch/csrc/autograd/python_variable.h>
7+ 
8+#include <torch_npu/csrc/inductor/aoti_package/model_package_loader.h>
9+#include <torch_npu/csrc/inductor/aoti_runner/model_container_runner.h>
10+#include <torch_npu/csrc/inductor/aoti_runner/pybind.h>
11+ 
12+namespace torch::inductor {
13+ 
14+class AOTIModelPackageLoaderPybind : public AOTIModelPackageLoader {
15+ public:
16+ AOTIModelPackageLoaderPybind(
17+ const std::string& model_package_path,
18+ const std::string& model_name,
19+ const bool run_single_threaded)
20+ : AOTIModelPackageLoader(
21+ model_package_path,
22+ model_name,
23+ run_single_threaded) {}
24+ 
25+ py::list boxed_run(py::list& inputs, void* stream_handle = nullptr) {
26+ std::vector<at::Tensor> input_tensors;
27+ input_tensors.reserve(inputs.size());
28+ for (auto& item : inputs) {
29+ input_tensors.emplace_back(py::cast<at::Tensor>(item));
30+ }
31+ // Explicitly clear the passed-in Python list
32+ inputs.attr("clear")();
33+ 
34+ std::vector<at::Tensor> result_tensors = AOTIModelPackageLoader::boxed_run(
35+ std::move(input_tensors), stream_handle);
36+ 
37+ py::list outputs;
38+ for (const auto& tensor : result_tensors) {
39+ outputs.append(
40+ py::reinterpret_steal<py::object>(THPVariable_Wrap(tensor)));
41+ }
42+ return outputs;
43+ }
44+};
45+ 
46+void initAOTIPackageBindings(PyObject* module) {
47+ auto rootModule = py::handle(module).cast<py::module>();
48+ auto m = rootModule.def_submodule("_aoti");
49+ 
50+ py::class_<AOTIModelPackageLoaderPybind>(m, "AOTIModelPackageLoader")
51+ .def(py::init<const std::string&, const std::string&, const bool>())
52+ .def("get_metadata", &AOTIModelPackageLoaderPybind::get_metadata)
53+ .def(
54+ "run",
55+ &AOTIModelPackageLoaderPybind::run,
56+ py::arg("inputs"),
57+ py::arg("stream_handle") = nullptr)
58+ .def(
59+ "boxed_run",
60+ &AOTIModelPackageLoaderPybind::boxed_run,
61+ py::arg("inputs"),
62+ py::arg("stream_handle") = nullptr)
63+ .def("get_call_spec", &AOTIModelPackageLoaderPybind::get_call_spec)
64+ .def("load_constants", &AOTIModelPackageLoaderPybind::load_constants)
65+ .def(
66+ "get_constant_fqns",
67+ &AOTIModelPackageLoaderPybind::get_constant_fqns);
68+}
69+} // namespace torch::inductor
70+#endif
Atorch_npu/csrc/inductor/aoti_runner/model_container_runner_cpu.h+18-0
@@ -0,0 +1,18 @@
1+#if !defined(C10_MOBILE) && !defined(ANDROID)
2+#pragma once
3+ 
4+#include <torch_npu/csrc/inductor/aoti_runner/model_container_runner.h>
5+ 
6+namespace torch::inductor {
7+class TORCH_API AOTIModelContainerRunnerCpu : public AOTIModelContainerRunner {
8+ public:
9+ AOTIModelContainerRunnerCpu(
10+ const std::string& model_so_path,
11+ size_t num_models = 1,
12+ const bool run_single_threaded = false);
13+ 
14+ ~AOTIModelContainerRunnerCpu() override;
15+};
16+ 
17+} // namespace torch::inductor
18+#endif
Atorch_npu/csrc/inductor/aoti_runner/pybind.cpp+93-0
@@ -0,0 +1,93 @@
1+#ifndef BUILD_LIBTORCH
2+#include <torch/csrc/utils/pybind.h>
3+ 
4+#ifdef USE_NPU
5+#include <torch_npu/csrc/inductor/aoti_runner/model_container_runner_npu.h>
6+#endif
7+#include <torch_npu/csrc/inductor/aoti_runner/model_container_runner_cpu.h>
8+#include <torch_npu/csrc/inductor/aoti_runner/pybind.h>
9+#include <torch_npu/csrc/inductor/aoti_torch/utils.h>
10+#include <torch_npu/csrc/inductor/aoti_torch/tensor_converter.h>
11+ 
12+namespace torch::inductor {
13+ 
14+void initAOTIRunnerBindings(PyObject* module) {
15+ auto rootModule = py::handle(module).cast<py::module>();
16+ auto m = rootModule.def_submodule("_aoti");
17+ 
18+ py::class_<AOTIModelContainerRunnerCpu>(m, "AOTIModelContainerRunnerCpu")
19+ .def(py::init<const std::string&, int>())
20+ .def(
21+ "run",
22+ &AOTIModelContainerRunnerCpu::run,
23+ py::arg("inputs"),
24+ py::arg("stream_handle") = nullptr)
25+ .def("get_call_spec", &AOTIModelContainerRunnerCpu::get_call_spec)
26+ .def(
27+ "get_constant_names_to_original_fqns",
28+ &AOTIModelContainerRunnerCpu::getConstantNamesToOriginalFQNs)
29+ .def(
30+ "get_constant_names_to_dtypes",
31+ &AOTIModelContainerRunnerCpu::getConstantNamesToDtypes)
32+ .def(
33+ "update_constant_buffer",
34+ static_cast<void (AOTIModelContainerRunnerCpu::*)(
35+ std::unordered_map<std::string, at::Tensor>&, bool, bool)>(
36+ &AOTIModelContainerRunnerCpu::update_constant_buffer));
37+ 
38+#ifdef USE_NPU
39+ py::class_<AOTIModelContainerRunnerNpu>(m, "AOTIModelContainerRunnerNpu")
40+ .def(py::init<const std::string&, int>())
41+ .def(py::init<const std::string&, int, const std::string&>())
42+ .def(py::init<
43+ const std::string&,
44+ int,
45+ const std::string&,
46+ const std::string&>())
47+ .def(
48+ "run",
49+ &AOTIModelContainerRunnerNpu::run,
50+ py::arg("inputs"),
51+ py::arg("stream_handle") = nullptr)
52+ .def("get_call_spec", &AOTIModelContainerRunnerNpu::get_call_spec)
53+ .def(
54+ "get_constant_names_to_original_fqns",
55+ &AOTIModelContainerRunnerNpu::getConstantNamesToOriginalFQNs)
56+ .def(
57+ "get_constant_names_to_dtypes",
58+ &AOTIModelContainerRunnerNpu::getConstantNamesToDtypes)
59+ .def(
60+ "update_constant_buffer",
61+ static_cast<void (AOTIModelContainerRunnerNpu::*)(
62+ std::unordered_map<std::string, at::Tensor>&, bool, bool)>(
63+ &AOTIModelContainerRunnerNpu::update_constant_buffer));
64+#endif
65+ 
66+ m.def(
67+ "unsafe_alloc_void_ptrs_from_tensors",
68+ [](const std::vector<at::Tensor>& tensors) {
69+ std::vector<AtenTensorHandle> handles =
70+ torch::aot_inductor::unsafe_alloc_new_handles_from_tensors(tensors);
71+ std::vector<void*> result(
72+ reinterpret_cast<void**>(handles.data()),
73+ reinterpret_cast<void**>(handles.data()) + handles.size());
74+ return result;
75+ });
76+ m.def("unsafe_alloc_void_ptr_from_tensor", [](at::Tensor& tensor) {
77+ return reinterpret_cast<void*>(
78+ torch::aot_inductor::new_tensor_handle(std::move(tensor)));
79+ });
80+ m.def(
81+ "alloc_tensors_by_stealing_from_void_ptrs",
82+ [](std::vector<void*>& raw_handles) {
83+ return torch::aot_inductor::alloc_tensors_by_stealing_from_handles(
84+ reinterpret_cast<AtenTensorHandle*>(raw_handles.data()),
85+ raw_handles.size());
86+ });
87+ m.def("alloc_tensor_by_stealing_from_void_ptr", [](void* raw_handle) {
88+ return *torch::aot_inductor::tensor_handle_to_tensor_pointer(
89+ reinterpret_cast<AtenTensorHandle>(raw_handle));
90+ });
91+}
92+} // namespace torch::inductor
93+#endif
Mtorch_npu/csrc/inductor/aoti_runtime/model.h+32-11
@@ -87,25 +87,46 @@ using ConstantMap = std::unordered_map<std::string, RAIIAtenTensorHandle>;
87// Update the list here if more devices are supported in the future87// Update the list here if more devices are supported in the future
88inline void parse_device_str(const std::string& device_str, int32_t& device_type, int32_t& device_idx)88inline void parse_device_str(const std::string& device_str, int32_t& device_type, int32_t& device_idx)
89{89{
90- std::regex re("(cpu|npu)(:([0-9]+))?");90+ if (device_str.empty()) {
91- std::smatch sm;91+ AOTI_RUNTIME_CHECK(false, "Invalid device: " + device_str);
92- bool matched = std::regex_match(device_str, sm, re);92+ }
93- AOTI_RUNTIME_CHECK(matched, "Invalid device: " + device_str);
94 93 
95- if (sm[1].str() == "cpu") {94+ size_t colon_pos = device_str.find(':');
95+ std::string device_type_str;
96+ if (colon_pos == std::string::npos) {
97+ device_type_str = device_str;
98+ } else {
99+ device_type_str = device_str.substr(0, colon_pos);
100+ }
101+ 
102+ if (device_type_str == "cpu") {
96 device_type = aoti_torch_device_type_cpu();103 device_type = aoti_torch_device_type_cpu();
97#ifdef USE_NPU104#ifdef USE_NPU
98- } else if (sm[1].str() == "npu") {105+ } else if (device_type_str == "npu") {
99 device_type = aoti_torch_device_type_npu();106 device_type = aoti_torch_device_type_npu();
100#endif107#endif
101 } else {108 } else {
102 AOTI_RUNTIME_CHECK(false, "Invalid device: " + device_str);109 AOTI_RUNTIME_CHECK(false, "Invalid device: " + device_str);
103 }110 }
104- const size_t default_sm = 3;111+ 
105- if (sm[default_sm].matched) {112+ device_idx = -1;
106- device_idx = stoi(sm[default_sm].str());113+ if (colon_pos != std::string::npos) {
107- } else {114+ std::string index_str = device_str.substr(colon_pos + 1);
108- device_idx = -1;115+ if (index_str.empty()) {
116+ AOTI_RUNTIME_CHECK(false, "Invalid device: " + device_str);
117+ }
118+ for (char c : index_str) {
119+ if (!std::isdigit(c)) {
120+ AOTI_RUNTIME_CHECK(false, "Invalid device: " + device_str);
121+ }
122+ }
123+ try {
124+ device_idx = std::stoi(index_str);
125+ } catch (const std::out_of_range& e) {
126+ AOTI_RUNTIME_CHECK(false, "Invalid device: " + device_str);
127+ } catch (const std::invalid_argument& e) {
128+ AOTI_RUNTIME_CHECK(false, "Invalid device: " + device_str);
129+ }
109 }130 }
110}131}
111 132