feat add shim_npu-pytorch-AtomGit

已合并

feat add shim_npu #35590

zzhongmin创建于 5月13日从zzhongmin/pytorch： v2.9.0_shim 合入到 Ascend/pytorch： v2.9.0

feat add shim_npu #35590

已合并

zzhongmin创建于 5月13日

从zzhongmin/pytorch v2.9.0_shim合入到Ascend/pytorch v2.9.0

共 12 个文件变更+1128-16

Atest/_inductor/test_npu_aoti_shim.py+179-0

@@ -0,0 +1,179 @@
	1	+import os
	2	+import sys
	3	+from pathlib import Path
	4	+
	5	+
	6	+REPO_ROOT = Path(__file__).resolve().parents[2]
	7	+BUILD_PACKAGES_DIR = REPO_ROOT / "build" / "packages"
	8	+
	9	+# Avoid backend autoload importing the source-tree torch_npu package before
	10	+# torch is fully initialized. Prefer the freshly built package output when it
	11	+# exists, and otherwise fall back to the installed package.
	12	+os.environ["TORCH_DEVICE_BACKEND_AUTOLOAD"] = "0"
	13	+for path in ("", str(REPO_ROOT)):
	14	+ while path in sys.path:
	15	+ sys.path.remove(path)
	16	+if BUILD_PACKAGES_DIR.exists():
	17	+ sys.path.insert(0, str(BUILD_PACKAGES_DIR))
	18	+
	19	+import torch
	20	+
	21	+from torch_npu.testing.testcase import run_tests, TestCase # noqa: F401
	22	+
	23	+
	O O openLiBingCI5月18日此条代码评论区间+19至+23 【openlibing.ci】识别到代码检查告警抑制注释，匹配工具：flake8，请Committer检视其合理性。
	24	+CPP_EXTENSIONS_DIR = REPO_ROOT / "test" / "cpp_extensions"
	25	+sys.path.insert(0, str(CPP_EXTENSIONS_DIR))
	26	+
	27	+from torch_test_cpp_extension.load_npu_aoti_shim import load_npu_aoti_shim_extension
	28	+
	29	+
	30	+class TestNpuAOTIShim(TestCase):
	31	+ module = None
	32	+
	33	+ @classmethod
	34	+ def setUpClass(cls):
	35	+ super().setUpClass()
	36	+ if not torch.npu.is_available():
	37	+ raise RuntimeError("test_npu_aoti_shim requires an available NPU")
	38	+ cls.module = load_npu_aoti_shim_extension()
	39	+
	40	+ def test_npu_raw_stream_matches_python_stream(self):
	41	+ device_index = 0
	42	+ x = torch.randn(8, device=f"npu:{device_index}", dtype=torch.float32)
	43	+
	44	+ with torch.npu.device(device_index):
	45	+ expected_default = torch.npu.current_stream().npu_stream
	46	+ actual_default = self.module.get_npu_raw_stream(x)
	47	+ self.assertEqual(actual_default, expected_default)
	48	+
	49	+ custom_stream = torch.npu.Stream()
	50	+ with torch.npu.stream(custom_stream):
	51	+ actual_custom = self.module.get_npu_raw_stream(x)
	52	+ self.assertEqual(actual_custom, custom_stream.npu_stream)
	53	+
	54	+ def test_zero_size_tensor_from_blob_uses_null_data_path_on_npu(self):
	55	+ x = torch.randn(4, device="npu:0", dtype=torch.float32)
	56	+
	57	+ result = self.module.make_zero_size_blob_tensor(x)
	58	+
	59	+ self.assertEqual(result.device, x.device)
	60	+ self.assertEqual(result.dtype, torch.float32)
	61	+ self.assertEqual(result.layout, torch.strided)
	62	+ self.assertEqual(tuple(result.shape), (0,))
	63	+ self.assertEqual(tuple(result.stride()), (1,))
	64	+ self.assertEqual(result.numel(), 0)
	65	+
	66	+ def test_zero_size_tensor_from_blob_uses_cpu_device_branch(self):
	67	+ x = torch.randn(4, device="npu:0", dtype=torch.float32)
	68	+
	69	+ result = self.module.make_zero_size_cpu_blob_tensor(x)
	70	+
	71	+ self.assertEqual(result.device.type, "cpu")
	72	+ self.assertEqual(result.dtype, torch.float32)
	73	+ self.assertEqual(result.layout, torch.strided)
	74	+ self.assertEqual(tuple(result.shape), (0,))
	75	+ self.assertEqual(tuple(result.stride()), (1,))
	76	+ self.assertEqual(result.numel(), 0)
	77	+
	78	+ def test_mkldnn_blob_tensor_v2_is_rejected_by_original_npu_logic(self):
	79	+ x = torch.randn(4, device="npu:0", dtype=torch.float32)
	80	+
	81	+ self.assertEqual(self.module.check_mkldnn_blob_tensor_v2_rejected(x), 1)
	82	+
	83	+ def test_blob_tensor_v2_propagates_inner_failure(self):
	84	+ x = torch.randn(4, device="npu:0", dtype=torch.float32)
	85	+
	86	+ self.assertEqual(
	87	+ self.module.check_blob_tensor_v2_propagates_invalid_device_failure(x),
	88	+ 1,
	89	+ )
	90	+
	91	+ def test_null_delete_paths_accept_nullptr_inputs(self):
	92	+ x = torch.randn(4, device="npu:0", dtype=torch.float32)
	93	+
	94	+ self.assertEqual(self.module.check_null_delete_paths(x), 1)
	95	+
	96	+ def test_invalid_stream_guard_path_returns_failure_internally(self):
	97	+ x = torch.randn(4, device="npu:0", dtype=torch.float32)
	98	+
	99	+ self.assertEqual(self.module.check_invalid_stream_guard_path(x), 1)
	100	+
	101	+ def test_null_stream_guard_path_returns_failure_internally(self):
	102	+ x = torch.randn(4, device="npu:0", dtype=torch.float32)
	103	+
	104	+ self.assertEqual(self.module.check_null_stream_guard_path(x), 1)
	105	+
	106	+ def test_invalid_device_guard_creation_returns_failure_internally(self):
	107	+ x = torch.randn(4, device="npu:0", dtype=torch.float32)
	108	+
	109	+ self.assertEqual(self.module.check_invalid_device_guard_creation(x), 1)
	110	+
	111	+ def test_invalid_device_guard_set_index_returns_failure_internally(self):
	112	+ x = torch.randn(4, device="npu:0", dtype=torch.float32)
	113	+
	114	+ self.assertEqual(self.module.check_invalid_device_guard_set_index(x), 1)
	115	+
	116	+ def test_invalid_device_current_stream_returns_failure_internally(self):
	117	+ x = torch.randn(4, device="npu:0", dtype=torch.float32)
	118	+
	119	+ self.assertEqual(self.module.check_invalid_device_current_stream(x), 1)
	120	+
	121	+ def test_null_stream_guard_output_handle_returns_failure_internally(self):
	122	+ x = torch.randn(4, device="npu:0", dtype=torch.float32)
	123	+
	124	+ self.assertEqual(self.module.check_null_stream_guard_output_handle(x), 1)
	125	+
	126	+ def test_null_current_stream_output_handle_returns_failure_internally(self):
	127	+ x = torch.randn(4, device="npu:0", dtype=torch.float32)
	128	+
	129	+ self.assertEqual(self.module.check_null_current_stream_output_handle(x), 1)
	130	+
	131	+ def test_null_guard_output_handle_returns_failure_internally(self):
	132	+ x = torch.randn(4, device="npu:0", dtype=torch.float32)
	133	+
	134	+ self.assertEqual(self.module.check_null_guard_output_handle(x), 1)
	135	+
	136	+ def test_null_blob_tensor_output_handle_returns_failure_internally(self):
	137	+ x = torch.randn(4, device="npu:0", dtype=torch.float32)
	138	+
	139	+ self.assertEqual(self.module.check_null_blob_tensor_output_handle(x), 1)
	140	+
	141	+ def test_null_blob_tensor_v2_output_handle_returns_failure_internally(self):
	142	+ x = torch.randn(4, device="npu:0", dtype=torch.float32)
	143	+
	144	+ self.assertEqual(self.module.check_null_blob_tensor_v2_output_handle(x), 1)
	145	+
	146	+ def test_current_device_stream_lookup_uses_negative_one_semantics(self):
	147	+ x = torch.randn(4, device="npu:0", dtype=torch.float32)
	148	+
	149	+ self.assertEqual(self.module.check_current_device_stream_lookup(x), 1)
	150	+
	151	+ def test_default_stream_guard_roundtrip_restores_original_stream(self):
	152	+ x = torch.randn(4, device="npu:0", dtype=torch.float32)
	153	+
	154	+ self.assertEqual(self.module.check_default_stream_guard_roundtrip(x), 1)
	155	+
	156	+ def test_run_npu_shim_checks_restores_device_and_stream(self):
	157	+ device_count = torch.npu.device_count()
	158	+ target_device = 1 if device_count > 1 else 0
	159	+ original_device = 0 if target_device != 0 else target_device
	160	+
	161	+ torch.npu.set_device(original_device)
	162	+ x = torch.randn(16, device=f"npu:{target_device}", dtype=torch.float32)
	163	+ expected = x + 1
	164	+
	165	+ with torch.npu.device(target_device):
	166	+ custom_stream = torch.npu.Stream()
	167	+ with torch.npu.stream(custom_stream):
	168	+ before_stream = torch.npu.current_stream().npu_stream
	169	+ result = self.module.run_npu_shim_checks(x)
	170	+ after_stream = torch.npu.current_stream().npu_stream
	171	+
	172	+ torch.testing.assert_close(result, expected)
	173	+ self.assertEqual(before_stream, custom_stream.npu_stream)
	174	+ self.assertEqual(after_stream, custom_stream.npu_stream)
	175	+ self.assertEqual(torch.npu.current_device(), original_device)
	176	+
	177	+
	178	+if __name__ == "__main__":
	179	+ run_tests()

Mtest/_inductor/test_npu_device.py+2-2

	@@ -31,7 +31,7 @@ class TestNpuDevice(TestCase):
31	self.assertIn("#include <sys/syscall.h>", result)	31	self.assertIn("#include <sys/syscall.h>", result)
32	self.assertIn("#include <torch_npu/csrc/framework/OpCommand.h>", result)	32	self.assertIn("#include <torch_npu/csrc/framework/OpCommand.h>", result)
33	self.assertIn("#include <torch_npu/csrc/core/npu/NPUStream.h>", result)	33	self.assertIn("#include <torch_npu/csrc/core/npu/NPUStream.h>", result)
34	- self.assertIn("#include \"runtime/runtime/rt.h\"", result)	34	+ self.assertIn('#include "runtime/runtime/rt.h"', result)
35		35
36	def test_cpp_aoti_stream_guard(self):	36	def test_cpp_aoti_stream_guard(self):
37	overrides = NewNPUDeviceOpOverrides()	37	overrides = NewNPUDeviceOpOverrides()
	@@ -74,4 +74,4 @@ class TestNpuDevice(TestCase):
74		74
75		75
76	if __name__ == "__main__":	76	if __name__ == "__main__":
77	- run_tests()	77	+ run_tests()

Atest/cpp_extensions/npu_aoti_shim_extension.cpp+657-0

@@ -0,0 +1,657 @@
	1	+#include <ATen/Tensor.h>
	2	+#include <ATen/ops/add.h>
	3	+#include <c10/core/DeviceType.h>
	4	+#include <torch/extension.h>
	5	+
	6	+#include <torch_npu/csrc/core/npu/NPUFunctions.h>
	7	+#include <torch_npu/csrc/core/npu/NPUGuard.h>
	8	+#include <torch_npu/csrc/core/npu/NPUStream.h>
	9	+#include <torch_npu/csrc/inductor/aoti_torch/c/shim.h>
	10	+#include <torch_npu/csrc/inductor/aoti_torch/utils.h>
	11	+
	12	+#include <mutex>
	13	+#include <optional>
	14	+#include <unordered_map>
	15	+#include <vector>
	16	+
	17	+namespace {
	18	+
	19	+std::mutex& stream_registry_mutex()
	20	+{
	21	+ static std::mutex mutex;
	22	+ return mutex;
	23	+}
	24	+
	25	+std::unordered_map<void*, c10::Stream>& stream_registry()
	26	+{
	27	+ static std::unordered_map<void*, c10::Stream> registry;
	28	+ return registry;
	29	+}
	30	+
	31	+void remember_stream(const c10_npu::NPUStream& stream)
	32	+{
	33	+ auto* raw_stream = reinterpret_cast<void*>(stream.stream(false));
	34	+ std::lock_guard<std::mutex> lock(stream_registry_mutex());
	35	+ stream_registry().insert_or_assign(raw_stream, static_cast<c10::Stream>(stream));
	36	+}
	37	+
	38	+void check_aoti_error(AOTITorchError err, const char* call)
	39	+{
	40	+ TORCH_CHECK(err == AOTI_TORCH_SUCCESS, call, " failed with error code ", err);
	41	+}
	42	+
	43	+std::vector<int64_t> to_vector(c10::IntArrayRef values)
	44	+{
	45	+ return std::vector<int64_t>(values.begin(), values.end());
	46	+}
	47	+
	48	+void assert_npu_tensor(const at::Tensor& tensor, const char* op_name)
	49	+{
	50	+ TORCH_CHECK(
	51	+ tensor.device().type() == c10::DeviceType::PrivateUse1,
	52	+ op_name,
	53	+ " expects an NPU tensor");
	54	+ TORCH_CHECK(
	55	+ tensor.scalar_type() == at::kFloat,
	56	+ op_name,
	57	+ " currently tests float32 tensors only");
	58	+ TORCH_CHECK(
	59	+ tensor.layout() == at::kStrided,
	60	+ op_name,
	61	+ " expects a strided tensor");
	62	+}
	63	+
	64	+int64_t get_npu_raw_stream_impl(const at::Tensor& tensor)
	65	+{
	66	+ assert_npu_tensor(tensor, "get_npu_raw_stream");
	67	+
	68	+ void* shim_stream = nullptr;
	69	+ const auto device_index = tensor.device().index();
	70	+ check_aoti_error(
	71	+ aoti_torch_get_current_npu_stream(device_index, &shim_stream),
	72	+ "aoti_torch_get_current_npu_stream");
	73	+
	74	+ auto direct_stream =
	75	+ reinterpret_cast<void*>(c10_npu::getCurrentNPUStream(device_index).stream(false));
	76	+ TORCH_CHECK(
	77	+ shim_stream == direct_stream,
	78	+ "aoti_torch_get_current_npu_stream returned ",
	79	+ shim_stream,
	80	+ ", but c10_npu returned ",
	81	+ direct_stream);
	82	+ return reinterpret_cast<int64_t>(shim_stream);
	83	+}
	84	+
	85	+at::Tensor make_zero_size_blob_tensor_impl(const at::Tensor& tensor)
	86	+{
	87	+ assert_npu_tensor(tensor, "make_zero_size_blob_tensor");
	88	+
	89	+ const auto device_index = tensor.device().index();
	90	+ std::vector<int64_t> sizes = {0};
	91	+ std::vector<int64_t> strides = {1};
	92	+ AtenTensorHandle handle = nullptr;
	93	+ check_aoti_error(
	94	+ aoti_torch_create_tensor_from_blob_npu(
	95	+ nullptr,
	96	+ static_cast<int64_t>(sizes.size()),
	97	+ sizes.data(),
	98	+ strides.data(),
	99	+ 0,
	100	+ aoti_torch_dtype_float32(),
	101	+ aoti_torch_device_type_npu(),
	102	+ device_index,
	103	+ &handle),
	104	+ "aoti_torch_create_tensor_from_blob_npu(nullptr, npu)");
	105	+
	106	+ at::Tensor result =
	107	+ *torch::aot_inductor::tensor_handle_to_tensor_pointer(handle);
	108	+ TORCH_CHECK(result.device() == tensor.device(), "Zero-size NPU tensor device mismatch");
	109	+ TORCH_CHECK(result.scalar_type() == at::kFloat, "Zero-size NPU tensor dtype mismatch");
	110	+ TORCH_CHECK(result.layout() == at::kStrided, "Zero-size NPU tensor layout mismatch");
	111	+ TORCH_CHECK(result.sizes().vec() == sizes, "Zero-size NPU tensor sizes mismatch");
	112	+ TORCH_CHECK(result.strides().vec() == strides, "Zero-size NPU tensor strides mismatch");
	113	+ TORCH_CHECK(result.numel() == 0, "Zero-size NPU tensor should be empty");
	114	+ check_aoti_error(
	115	+ aoti_torch_delete_tensor_object(handle),
	116	+ "aoti_torch_delete_tensor_object(zero_size_npu)");
	117	+ return result;
	118	+}
	119	+
	120	+at::Tensor make_zero_size_cpu_blob_tensor_impl(const at::Tensor& tensor)
	121	+{
	122	+ assert_npu_tensor(tensor, "make_zero_size_cpu_blob_tensor");
	123	+
	124	+ std::vector<int64_t> sizes = {0};
	125	+ std::vector<int64_t> strides = {1};
	126	+ AtenTensorHandle handle = nullptr;
	127	+ check_aoti_error(
	128	+ aoti_torch_create_tensor_from_blob_npu(
	129	+ nullptr,
	130	+ static_cast<int64_t>(sizes.size()),
	131	+ sizes.data(),
	132	+ strides.data(),
	133	+ 0,
	134	+ aoti_torch_dtype_float32(),
	135	+ aoti_torch_device_type_cpu(),
	136	+ 0,
	137	+ &handle),
	138	+ "aoti_torch_create_tensor_from_blob_npu(nullptr, cpu)");
	139	+
	140	+ at::Tensor result =
	141	+ *torch::aot_inductor::tensor_handle_to_tensor_pointer(handle);
	142	+ TORCH_CHECK(result.device().type() == c10::DeviceType::CPU, "Zero-size CPU tensor device mismatch");
	143	+ TORCH_CHECK(result.scalar_type() == at::kFloat, "Zero-size CPU tensor dtype mismatch");
	144	+ TORCH_CHECK(result.layout() == at::kStrided, "Zero-size CPU tensor layout mismatch");
	145	+ TORCH_CHECK(result.sizes().vec() == sizes, "Zero-size CPU tensor sizes mismatch");
	146	+ TORCH_CHECK(result.strides().vec() == strides, "Zero-size CPU tensor strides mismatch");
	147	+ TORCH_CHECK(result.numel() == 0, "Zero-size CPU tensor should be empty");
	148	+ check_aoti_error(
	149	+ aoti_torch_delete_tensor_object(handle),
	150	+ "aoti_torch_delete_tensor_object(zero_size_cpu)");
	151	+ return result;
	152	+}
	153	+
	154	+int64_t check_mkldnn_blob_tensor_v2_rejected_impl(const at::Tensor& tensor)
	155	+{
	156	+ assert_npu_tensor(tensor, "check_mkldnn_blob_tensor_v2_rejected");
	157	+
	158	+ auto contiguous = tensor.contiguous();
	159	+ auto sizes = to_vector(contiguous.sizes());
	160	+ auto strides = to_vector(contiguous.strides());
	161	+ AtenTensorHandle handle = nullptr;
	162	+ auto err = aoti_torch_create_tensor_from_blob_npu_v2(
	163	+ contiguous.data_ptr(),
	164	+ contiguous.dim(),
	165	+ sizes.data(),
	166	+ strides.data(),
	167	+ contiguous.storage_offset(),
	168	+ aoti_torch_dtype_float32(),
	169	+ aoti_torch_device_type_npu(),
	170	+ tensor.device().index(),
	171	+ &handle,
	172	+ static_cast<int32_t>(at::kMkldnn),
	173	+ nullptr,
	174	+ 0);
	175	+ TORCH_CHECK(
	176	+ err == AOTI_TORCH_FAILURE,
	177	+ "mkldnn layout should make aoti_torch_create_tensor_from_blob_npu_v2 fail");
	178	+ TORCH_CHECK(handle == nullptr, "Handle should remain nullptr when mkldnn layout is rejected");
	179	+ return 1;
	180	+}
	181	+
	182	+int64_t check_blob_tensor_v2_propagates_invalid_device_failure_impl(const at::Tensor& tensor)
	183	+{
	184	+ assert_npu_tensor(tensor, "check_blob_tensor_v2_propagates_invalid_device_failure");
	185	+
	186	+ auto contiguous = tensor.contiguous();
	187	+ auto sizes = to_vector(contiguous.sizes());
	188	+ auto strides = to_vector(contiguous.strides());
	189	+ const auto invalid_device_index = static_cast<int32_t>(c10_npu::device_count());
	190	+ AtenTensorHandle handle = nullptr;
	191	+ auto err = aoti_torch_create_tensor_from_blob_npu_v2(
	192	+ contiguous.data_ptr(),
	193	+ contiguous.dim(),
	194	+ sizes.data(),
	195	+ strides.data(),
	196	+ contiguous.storage_offset(),
	197	+ aoti_torch_dtype_float32(),
	198	+ aoti_torch_device_type_npu(),
	199	+ invalid_device_index,
	200	+ &handle,
	201	+ aoti_torch_layout_strided(),
	202	+ nullptr,
	203	+ 0);
	204	+ TORCH_CHECK(
	205	+ err == AOTI_TORCH_FAILURE,
	206	+ "Invalid device should make aoti_torch_create_tensor_from_blob_npu_v2 fail");
	207	+ TORCH_CHECK(
	208	+ handle == nullptr,
	209	+ "Handle should remain nullptr when the inner blob tensor path fails");
	210	+ return 1;
	211	+}
	212	+
	213	+int64_t check_null_delete_paths_impl(const at::Tensor& tensor)
	214	+{
	215	+ assert_npu_tensor(tensor, "check_null_delete_paths");
	216	+
	217	+ check_aoti_error(
	218	+ aoti_torch_delete_npu_guard(nullptr),
	219	+ "aoti_torch_delete_npu_guard(nullptr)");
	220	+ check_aoti_error(
	221	+ aoti_torch_delete_npu_stream_guard(nullptr),
	222	+ "aoti_torch_delete_npu_stream_guard(nullptr)");
	223	+ return 1;
	224	+}
	225	+
	226	+int64_t check_invalid_stream_guard_path_impl(const at::Tensor& tensor)
	227	+{
	228	+ assert_npu_tensor(tensor, "check_invalid_stream_guard_path");
	229	+
	230	+ NPUStreamGuardHandle guard = nullptr;
	231	+ auto err = aoti_torch_create_npu_stream_guard(
	232	+ reinterpret_cast<void*>(0x1),
	233	+ tensor.device().index(),
	234	+ &guard);
	235	+ TORCH_CHECK(
	236	+ err == AOTI_TORCH_FAILURE,
	237	+ "Invalid stream pointer should make aoti_torch_create_npu_stream_guard fail");
	238	+ TORCH_CHECK(guard == nullptr, "Guard should remain nullptr when stream guard creation fails");
	239	+ return 1;
	240	+}
	241	+
	242	+int64_t check_null_stream_guard_path_impl(const at::Tensor& tensor)
	243	+{
	244	+ assert_npu_tensor(tensor, "check_null_stream_guard_path");
	245	+
	246	+ NPUStreamGuardHandle guard = nullptr;
	247	+ auto err =
	248	+ aoti_torch_create_npu_stream_guard(nullptr, tensor.device().index(), &guard);
	249	+ TORCH_CHECK(
	250	+ err == AOTI_TORCH_FAILURE,
	251	+ "Null stream pointer should make aoti_torch_create_npu_stream_guard fail");
	252	+ TORCH_CHECK(guard == nullptr, "Guard should remain nullptr when null stream guard creation fails");
	253	+ return 1;
	254	+}
	255	+
	256	+int64_t check_invalid_device_guard_creation_impl(const at::Tensor& tensor)
	257	+{
	258	+ assert_npu_tensor(tensor, "check_invalid_device_guard_creation");
	259	+
	260	+ const auto invalid_device_index = static_cast<int32_t>(c10_npu::device_count());
	261	+ NPUGuardHandle guard = nullptr;
	262	+ auto err = aoti_torch_create_npu_guard(invalid_device_index, &guard);
	263	+ TORCH_CHECK(
	264	+ err == AOTI_TORCH_FAILURE,
	265	+ "Invalid device index should make aoti_torch_create_npu_guard fail");
	266	+ TORCH_CHECK(guard == nullptr, "Guard should remain nullptr when invalid guard creation fails");
	267	+ return 1;
	268	+}
	269	+
	270	+int64_t check_invalid_device_guard_set_index_impl(const at::Tensor& tensor)
	271	+{
	272	+ assert_npu_tensor(tensor, "check_invalid_device_guard_set_index");
	273	+
	274	+ const auto original_device = c10_npu::current_device();
	275	+ const auto valid_device_index = tensor.device().index();
	276	+ const auto invalid_device_index = static_cast<int32_t>(c10_npu::device_count());
	277	+ NPUGuardHandle guard = nullptr;
	278	+ check_aoti_error(
	279	+ aoti_torch_create_npu_guard(valid_device_index, &guard),
	280	+ "aoti_torch_create_npu_guard(valid for set_index)");
	281	+ TORCH_CHECK(
	282	+ c10_npu::current_device() == valid_device_index,
	283	+ "Valid guard creation should switch to the requested device");
	284	+
	285	+ auto err = aoti_torch_npu_guard_set_index(guard, invalid_device_index);
	286	+ TORCH_CHECK(
	287	+ err == AOTI_TORCH_FAILURE,
	288	+ "Invalid device index should make aoti_torch_npu_guard_set_index fail");
	289	+ TORCH_CHECK(
	290	+ c10_npu::current_device() == valid_device_index,
	291	+ "Failed guard set_index should keep the previously selected device");
	292	+
	293	+ check_aoti_error(
	294	+ aoti_torch_delete_npu_guard(guard),
	295	+ "aoti_torch_delete_npu_guard(valid after failed set_index)");
	296	+ TORCH_CHECK(
	297	+ c10_npu::current_device() == original_device,
	298	+ "Deleting the guard should restore the original device after failed set_index");
	299	+ return 1;
	300	+}
	301	+
	302	+int64_t check_invalid_device_current_stream_impl(const at::Tensor& tensor)
	303	+{
	304	+ assert_npu_tensor(tensor, "check_invalid_device_current_stream");
	305	+
	306	+ const auto invalid_device_index = static_cast<int32_t>(c10_npu::device_count());
	307	+ void* stream = nullptr;
	308	+ auto err = aoti_torch_get_current_npu_stream(invalid_device_index, &stream);
	309	+ TORCH_CHECK(
	310	+ err == AOTI_TORCH_FAILURE,
	311	+ "Invalid device index should make aoti_torch_get_current_npu_stream fail");
	312	+ TORCH_CHECK(stream == nullptr, "Invalid device stream lookup should leave stream as nullptr");
	313	+ return 1;
	314	+}
	315	+
	316	+int64_t check_null_stream_guard_output_handle_impl(const at::Tensor& tensor)
	317	+{
	318	+ assert_npu_tensor(tensor, "check_null_stream_guard_output_handle");
	319	+
	320	+ auto pooled_stream = c10_npu::getNPUStreamFromPool(tensor.device().index());
	321	+ remember_stream(pooled_stream);
	322	+ auto err = aoti_torch_create_npu_stream_guard(
	323	+ reinterpret_cast<void*>(pooled_stream.stream(false)),
	324	+ tensor.device().index(),
	325	+ nullptr);
	326	+ TORCH_CHECK(
	327	+ err == AOTI_TORCH_FAILURE,
	328	+ "Null ret_guard should make aoti_torch_create_npu_stream_guard fail");
	329	+ return 1;
	330	+}
	331	+
	332	+int64_t check_null_current_stream_output_handle_impl(const at::Tensor& tensor)
	333	+{
	334	+ assert_npu_tensor(tensor, "check_null_current_stream_output_handle");
	335	+
	336	+ auto err = aoti_torch_get_current_npu_stream(tensor.device().index(), nullptr);
	337	+ TORCH_CHECK(
	338	+ err == AOTI_TORCH_FAILURE,
	339	+ "Null ret_stream should make aoti_torch_get_current_npu_stream fail");
	340	+ return 1;
	341	+}
	342	+
	343	+int64_t check_null_guard_output_handle_impl(const at::Tensor& tensor)
	344	+{
	345	+ assert_npu_tensor(tensor, "check_null_guard_output_handle");
	346	+
	347	+ auto err = aoti_torch_create_npu_guard(tensor.device().index(), nullptr);
	348	+ TORCH_CHECK(
	349	+ err == AOTI_TORCH_FAILURE,
	350	+ "Null ret_guard should make aoti_torch_create_npu_guard fail");
	351	+ return 1;
	352	+}
	353	+
	354	+int64_t check_null_blob_tensor_output_handle_impl(const at::Tensor& tensor)
	355	+{
	356	+ assert_npu_tensor(tensor, "check_null_blob_tensor_output_handle");
	357	+
	358	+ auto contiguous = tensor.contiguous();
	359	+ auto sizes = to_vector(contiguous.sizes());
	360	+ auto strides = to_vector(contiguous.strides());
	361	+ auto err = aoti_torch_create_tensor_from_blob_npu(
	362	+ contiguous.data_ptr(),
	363	+ contiguous.dim(),
	364	+ sizes.data(),
	365	+ strides.data(),
	366	+ contiguous.storage_offset(),
	367	+ aoti_torch_dtype_float32(),
	368	+ aoti_torch_device_type_npu(),
	369	+ tensor.device().index(),
	370	+ nullptr);
	371	+ TORCH_CHECK(
	372	+ err == AOTI_TORCH_FAILURE,
	373	+ "Null ret_new_tensor should make aoti_torch_create_tensor_from_blob_npu fail");
	374	+ return 1;
	375	+}
	376	+
	377	+int64_t check_null_blob_tensor_v2_output_handle_impl(const at::Tensor& tensor)
	378	+{
	379	+ assert_npu_tensor(tensor, "check_null_blob_tensor_v2_output_handle");
	380	+
	381	+ auto contiguous = tensor.contiguous();
	382	+ auto sizes = to_vector(contiguous.sizes());
	383	+ auto strides = to_vector(contiguous.strides());
	384	+ auto err = aoti_torch_create_tensor_from_blob_npu_v2(
	385	+ contiguous.data_ptr(),
	386	+ contiguous.dim(),
	387	+ sizes.data(),
	388	+ strides.data(),
	389	+ contiguous.storage_offset(),
	390	+ aoti_torch_dtype_float32(),
	391	+ aoti_torch_device_type_npu(),
	392	+ tensor.device().index(),
	393	+ nullptr,
	394	+ aoti_torch_layout_strided(),
	395	+ nullptr,
	396	+ 0);
	397	+ TORCH_CHECK(
	398	+ err == AOTI_TORCH_FAILURE,
	399	+ "Null ret_new_tensor should make aoti_torch_create_tensor_from_blob_npu_v2 fail");
	400	+ return 1;
	401	+}
	402	+
	403	+int64_t check_current_device_stream_lookup_impl(const at::Tensor& tensor)
	404	+{
	405	+ assert_npu_tensor(tensor, "check_current_device_stream_lookup");
	406	+
	407	+ const auto original_device = c10_npu::current_device();
	408	+ c10_npu::NPUGuard device_guard(tensor.device().index());
	409	+
	410	+ void* shim_stream = nullptr;
	411	+ check_aoti_error(
	412	+ aoti_torch_get_current_npu_stream(-1, &shim_stream),
	413	+ "aoti_torch_get_current_npu_stream(current device)");
	414	+ auto current_stream =
	415	+ reinterpret_cast<void*>(c10_npu::getCurrentNPUStream(tensor.device().index()).stream(false));
	416	+ TORCH_CHECK(
	417	+ shim_stream == current_stream,
	418	+ "Current-device stream fallback mismatch");
	419	+ TORCH_CHECK(
	420	+ c10_npu::current_device() == tensor.device().index(),
	421	+ "Current-device lookup should keep the selected device");
	422	+ TORCH_CHECK(
	423	+ c10_npu::current_device() == tensor.device().index(),
	424	+ "Device guard should keep the tensor device selected inside the scope");
	425	+ device_guard.set_index(original_device);
	426	+ TORCH_CHECK(
	427	+ c10_npu::current_device() == original_device,
	428	+ "Current-device lookup helper should restore the original device");
	429	+ return 1;
	430	+}
	431	+
	432	+int64_t check_default_stream_guard_roundtrip_impl(const at::Tensor& tensor)
	433	+{
	434	+ assert_npu_tensor(tensor, "check_default_stream_guard_roundtrip");
	435	+
	436	+ const auto device_index = tensor.device().index();
	437	+ auto pooled_stream = c10_npu::getNPUStreamFromPool(device_index);
	438	+ remember_stream(pooled_stream);
	439	+ c10_npu::NPUStreamGuard outer_guard(static_cast<c10::Stream>(pooled_stream));
	440	+
	441	+ auto original_stream =
	442	+ reinterpret_cast<void*>(c10_npu::getCurrentNPUStream(device_index).stream(false));
	443	+ auto default_stream = c10_npu::getDefaultNPUStream(device_index);
	444	+ remember_stream(default_stream);
	445	+
	446	+ NPUStreamGuardHandle stream_guard = nullptr;
	447	+ check_aoti_error(
	448	+ aoti_torch_create_npu_stream_guard(
	449	+ reinterpret_cast<void*>(default_stream.stream(false)),
	450	+ device_index,
	451	+ &stream_guard),
	452	+ "aoti_torch_create_npu_stream_guard(default stream)");
	453	+ auto guarded_stream =
	454	+ reinterpret_cast<void*>(c10_npu::getCurrentNPUStream(device_index).stream(false));
	455	+ TORCH_CHECK(
	456	+ guarded_stream == reinterpret_cast<void*>(default_stream.stream(false)),
	457	+ "Default stream guard did not switch to the default stream");
	458	+ check_aoti_error(
	459	+ aoti_torch_delete_npu_stream_guard(stream_guard),
	460	+ "aoti_torch_delete_npu_stream_guard(default stream)");
	461	+ auto restored_stream =
	462	+ reinterpret_cast<void*>(c10_npu::getCurrentNPUStream(device_index).stream(false));
	463	+ TORCH_CHECK(
	464	+ restored_stream == original_stream,
	465	+ "Default stream guard did not restore the original stream");
	466	+ return 1;
	467	+}
	468	+
	469	+at::Tensor run_npu_shim_checks_impl(const at::Tensor& tensor)
	470	+{
	471	+ assert_npu_tensor(tensor, "run_npu_shim_checks");
	472	+
	473	+ const auto device_index = tensor.device().index();
	474	+ TORCH_CHECK(
	475	+ aoti_torch_device_type_npu() ==
	476	+ static_cast<int32_t>(c10::DeviceType::PrivateUse1),
	477	+ "aoti_torch_device_type_npu should return PrivateUse1");
	478	+
	479	+ void* shim_stream = nullptr;
	480	+ check_aoti_error(
	481	+ aoti_torch_get_current_npu_stream(device_index, &shim_stream),
	482	+ "aoti_torch_get_current_npu_stream");
	483	+ auto current_stream =
	484	+ reinterpret_cast<void*>(c10_npu::getCurrentNPUStream(device_index).stream(false));
	485	+ TORCH_CHECK(
	486	+ shim_stream == current_stream,
	487	+ "Current stream mismatch: shim=",
	488	+ shim_stream,
	489	+ ", direct=",
	490	+ current_stream);
	491	+
	492	+ const auto original_device = c10_npu::current_device();
	493	+ NPUGuardHandle guard = nullptr;
	494	+ check_aoti_error(
	495	+ aoti_torch_create_npu_guard(device_index, &guard),
	496	+ "aoti_torch_create_npu_guard");
	497	+ TORCH_CHECK(
	498	+ c10_npu::current_device() == device_index,
	499	+ "NPU guard did not switch the current device");
	500	+ check_aoti_error(
	501	+ aoti_torch_npu_guard_set_index(guard, device_index),
	502	+ "aoti_torch_npu_guard_set_index");
	503	+ TORCH_CHECK(
	504	+ c10_npu::current_device() == device_index,
	505	+ "NPU guard set_index did not keep the requested device");
	506	+ check_aoti_error(
	507	+ aoti_torch_delete_npu_guard(guard),
	508	+ "aoti_torch_delete_npu_guard");
	509	+ TORCH_CHECK(
	510	+ c10_npu::current_device() == original_device,
	511	+ "Deleting the NPU guard did not restore the original device");
	512	+
	513	+ auto original_stream =
	514	+ reinterpret_cast<void*>(c10_npu::getCurrentNPUStream(device_index).stream(false));
	515	+ auto pooled_stream_obj = c10_npu::getNPUStreamFromPool(device_index);
	516	+ remember_stream(pooled_stream_obj);
	517	+ auto pooled_stream =
	518	+ reinterpret_cast<void*>(pooled_stream_obj.stream(false));
	519	+ NPUStreamGuardHandle stream_guard = nullptr;
	520	+ check_aoti_error(
	521	+ aoti_torch_create_npu_stream_guard(pooled_stream, device_index, &stream_guard),
	522	+ "aoti_torch_create_npu_stream_guard");
	523	+ auto guarded_stream =
	524	+ reinterpret_cast<void*>(c10_npu::getCurrentNPUStream(device_index).stream(false));
	525	+ TORCH_CHECK(
	526	+ guarded_stream == pooled_stream,
	527	+ "NPU stream guard did not switch to the pooled stream");
	528	+ check_aoti_error(
	529	+ aoti_torch_delete_npu_stream_guard(stream_guard),
	530	+ "aoti_torch_delete_npu_stream_guard");
	531	+ auto restored_stream =
	532	+ reinterpret_cast<void*>(c10_npu::getCurrentNPUStream(device_index).stream(false));
	533	+ TORCH_CHECK(
	534	+ restored_stream == original_stream,
	535	+ "Deleting the NPU stream guard did not restore the original stream");
	536	+
	537	+ auto contiguous = tensor.contiguous();
	538	+ auto sizes = to_vector(contiguous.sizes());
	539	+ auto strides = to_vector(contiguous.strides());
	540	+
	541	+ AtenTensorHandle alias_handle = nullptr;
	542	+ check_aoti_error(
	543	+ aoti_torch_create_tensor_from_blob_npu(
	544	+ contiguous.data_ptr(),
	545	+ contiguous.dim(),
	546	+ sizes.data(),
	547	+ strides.data(),
	548	+ contiguous.storage_offset(),
	549	+ aoti_torch_dtype_float32(),
	550	+ aoti_torch_device_type_npu(),
	551	+ device_index,
	552	+ &alias_handle),
	553	+ "aoti_torch_create_tensor_from_blob_npu");
	554	+ at::Tensor alias =
	555	+ *torch::aot_inductor::tensor_handle_to_tensor_pointer(alias_handle);
	556	+ TORCH_CHECK(alias.device() == contiguous.device(), "Alias tensor device mismatch");
	557	+ TORCH_CHECK(alias.scalar_type() == contiguous.scalar_type(), "Alias tensor dtype mismatch");
	558	+ TORCH_CHECK(alias.data_ptr() == contiguous.data_ptr(), "Alias tensor data_ptr mismatch");
	559	+ TORCH_CHECK(
	560	+ to_vector(alias.sizes()) == sizes,
	561	+ "Alias tensor sizes mismatch");
	562	+ TORCH_CHECK(
	563	+ to_vector(alias.strides()) == strides,
	564	+ "Alias tensor strides mismatch");
	565	+ TORCH_CHECK(
	566	+ alias.storage_offset() == contiguous.storage_offset(),
	567	+ "Alias tensor storage_offset mismatch");
	568	+ TORCH_CHECK(alias.equal(contiguous), "Alias tensor value mismatch");
	569	+ check_aoti_error(
	570	+ aoti_torch_delete_tensor_object(alias_handle),
	571	+ "aoti_torch_delete_tensor_object(alias_handle)");
	572	+
	573	+ auto result = at::add(alias, 1.0);
	574	+ return result;
	575	+}
	576	+
	577	+} // namespace
	578	+
	579	+namespace c10_npu {
	580	+
	581	+NPUStream getNPUStreamFromManagedAclrtStream(aclrtStream stream, c10::DeviceIndex device_index)
	582	+{
	583	+ TORCH_CHECK(stream != nullptr, "stream is nullptr");
	584	+
	585	+ auto find_registered_stream = [stream]() -> std::optional<NPUStream> {
	586	+ std::lock_guard<std::mutex> lock(stream_registry_mutex());
	587	+ auto it = stream_registry().find(reinterpret_cast<void*>(stream));
	588	+ if (it == stream_registry().end()) {
	589	+ return std::nullopt;
	590	+ }
	591	+ return NPUStream(NPUStream::UNCHECKED, it->second);
	592	+ };
	593	+ if (auto registered = find_registered_stream()) {
	594	+ return *registered;
	595	+ }
	596	+
	597	+ auto find_known_stream = [stream](c10::DeviceIndex idx) -> std::optional<NPUStream> {
	598	+ auto current = getCurrentNPUStream(idx);
	599	+ if (current.stream(false) == stream) {
	600	+ remember_stream(current);
	601	+ return current;
	602	+ }
	603	+
	604	+ auto default_stream = getDefaultNPUStream(idx);
	605	+ if (default_stream.stream(false) == stream) {
	606	+ remember_stream(default_stream);
	607	+ return default_stream;
	608	+ }
	609	+
	610	+ return std::nullopt;
	611	+ };
	612	+
	613	+ if (device_index != -1) {
	614	+ if (auto known = find_known_stream(device_index)) {
	615	+ return *known;
	616	+ }
	617	+ } else {
	618	+ const auto device_count = c10_npu::device_count();
	619	+ for (c10::DeviceIndex idx = 0; idx < device_count; ++idx) {
	620	+ if (auto known = find_known_stream(idx)) {
	621	+ return *known;
	622	+ }
	623	+ }
	624	+ }
	625	+
	626	+ TORCH_CHECK(
	627	+ false,
	628	+ "The aclrtStream is not managed by the shim test registry on device ",
	629	+ device_index);
	630	+}
	631	+
	632	+} // namespace c10_npu
	633	+
	634	+PYBIND11_MODULE(TORCH_EXTENSION_NAME, m)
	635	+{
	636	+ m.def("get_npu_raw_stream", &get_npu_raw_stream_impl);
	637	+ m.def("make_zero_size_blob_tensor", &make_zero_size_blob_tensor_impl);
	638	+ m.def("make_zero_size_cpu_blob_tensor", &make_zero_size_cpu_blob_tensor_impl);
	639	+ m.def("check_mkldnn_blob_tensor_v2_rejected", &check_mkldnn_blob_tensor_v2_rejected_impl);
	640	+ m.def(
	641	+ "check_blob_tensor_v2_propagates_invalid_device_failure",
	642	+ &check_blob_tensor_v2_propagates_invalid_device_failure_impl);
	643	+ m.def("check_null_delete_paths", &check_null_delete_paths_impl);
	644	+ m.def("check_invalid_stream_guard_path", &check_invalid_stream_guard_path_impl);
	645	+ m.def("check_null_stream_guard_path", &check_null_stream_guard_path_impl);
	646	+ m.def("check_invalid_device_guard_creation", &check_invalid_device_guard_creation_impl);
	647	+ m.def("check_invalid_device_guard_set_index", &check_invalid_device_guard_set_index_impl);
	648	+ m.def("check_invalid_device_current_stream", &check_invalid_device_current_stream_impl);
	649	+ m.def("check_null_stream_guard_output_handle", &check_null_stream_guard_output_handle_impl);
	650	+ m.def("check_null_current_stream_output_handle", &check_null_current_stream_output_handle_impl);
	651	+ m.def("check_null_guard_output_handle", &check_null_guard_output_handle_impl);
	652	+ m.def("check_null_blob_tensor_output_handle", &check_null_blob_tensor_output_handle_impl);
	653	+ m.def("check_null_blob_tensor_v2_output_handle", &check_null_blob_tensor_v2_output_handle_impl);
	654	+ m.def("check_current_device_stream_lookup", &check_current_device_stream_lookup_impl);
	655	+ m.def("check_default_stream_guard_roundtrip", &check_default_stream_guard_roundtrip_impl);
	656	+ m.def("run_npu_shim_checks", &run_npu_shim_checks_impl);
	657	+}

Mtest/cpp_extensions/setup.py+9-1

	@@ -1,4 +1,3 @@
1	-import sys
2	import os	1	import os
3		2
4	from setuptools import setup	3	from setuptools import setup
	@@ -12,11 +11,20 @@ set_npu_device()
12	CXX_FLAGS = ['-g']	11	CXX_FLAGS = ['-g']
13		12
14	USE_NINJA = os.getenv('USE_NINJA') == '1'	13	USE_NINJA = os.getenv('USE_NINJA') == '1'
		14	+REPO_ROOT = os.path.dirname(os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
		15	+SHIM_SOURCE = os.path.join(
		16	+ REPO_ROOT, "torch_npu", "csrc", "inductor", "aoti_torch", "shim_npu.cpp"
		17	+)
15		18
16	ext_modules = [	19	ext_modules = [
17	NpuExtension(	20	NpuExtension(
18	'torch_test_cpp_extension.npu', ['extension.cpp'],	21	'torch_test_cpp_extension.npu', ['extension.cpp'],
19	extra_compile_args=CXX_FLAGS),	22	extra_compile_args=CXX_FLAGS),
		23	+ NpuExtension(
		24	+ 'torch_test_cpp_extension.npu_aoti_shim',
		25	+ ['npu_aoti_shim_extension.cpp', SHIM_SOURCE],
		26	+ include_dirs=[REPO_ROOT],
		27	+ extra_compile_args=CXX_FLAGS),
20	NpuExtension(	28	NpuExtension(
21	'torch_test_cpp_extension.npu_from_blob', ['test_from_blob.cpp'],	29	'torch_test_cpp_extension.npu_from_blob', ['test_from_blob.cpp'],
22	extra_compile_args=CXX_FLAGS),	30	extra_compile_args=CXX_FLAGS),

Atest/cpp_extensions/torch_test_cpp_extension/load_npu_aoti_shim.py+49-0

@@ -0,0 +1,49 @@
	1	+import importlib
	2	+import os
	3	+import subprocess
	4	+import sys
	5	+from pathlib import Path
	6	+
	7	+
	8	+REPO_ROOT = Path(__file__).resolve().parents[3]
	9	+BUILD_PACKAGES_DIR = REPO_ROOT / "build" / "packages"
	10	+CPP_EXTENSIONS_DIR = REPO_ROOT / "test" / "cpp_extensions"
	11	+MODULE_NAME = "torch_test_cpp_extension.npu_aoti_shim"
	12	+
	13	+
	14	+def _build_pythonpath_for_subprocess():
	15	+ parts = []
	16	+ if BUILD_PACKAGES_DIR.exists():
	17	+ parts.append(str(BUILD_PACKAGES_DIR))
	18	+ existing = os.environ.get("PYTHONPATH")
	19	+ if existing:
	20	+ parts.append(existing)
	21	+ return os.pathsep.join(parts)
	22	+
	23	+
	24	+def _build_extension_inplace():
	25	+ env = os.environ.copy()
	26	+ pythonpath = _build_pythonpath_for_subprocess()
	27	+ if pythonpath:
	28	+ env["PYTHONPATH"] = pythonpath
	29	+
	30	+ subprocess.run(
	31	+ [sys.executable, "setup.py", "build_ext", "--inplace"],
	32	+ cwd=str(CPP_EXTENSIONS_DIR),
	33	+ check=True,
	34	+ env=env,
	35	+ )
	36	+
	37	+
	38	+def load_npu_aoti_shim_extension():
	39	+ import torch # noqa: F401
	40	+
	41	+ import torch_npu # noqa: F401
	O O openLiBingCI5月18日此条代码评论区间+37至+41 【openlibing.ci】识别到代码检查告警抑制注释，匹配工具：flake8，请Committer检视其合理性。
	42	+
	43	+ try:
	O O openLiBingCI5月18日此条代码评论区间+39至+43 【openlibing.ci】识别到代码检查告警抑制注释，匹配工具：flake8，请Committer检视其合理性。
	44	+ return importlib.import_module(MODULE_NAME)
	45	+ except ImportError:
	46	+ _build_extension_inplace()
	47	+ sys.modules.pop(MODULE_NAME, None)
	48	+ importlib.invalidate_caches()
	49	+ return importlib.import_module(MODULE_NAME)

Mtorch_npu/_inductor/codegen/cpp_wrapper.py+1-0

	@@ -243,6 +243,7 @@ class CppWrapperNpu(CppWrapperGpu):
243	self.header.splice("#include <unistd.h>")	243	self.header.splice("#include <unistd.h>")
244	self.header.splice("#include <filesystem>")	244	self.header.splice("#include <filesystem>")
245	self.header.splice(self.device_codegen.abi_compatible_header())	245	self.header.splice(self.device_codegen.abi_compatible_header())
		246	+ self.header.splice("#include <torch_npu/csrc/inductor/aoti_runtime/utils_npu.h>")
246	self.header.splice(	247	self.header.splice(
247	maybe_hipify_code_wrapper(self.device_codegen.kernel_driver())	248	maybe_hipify_code_wrapper(self.device_codegen.kernel_driver())
248	)	249	)

Mtorch_npu/csrc/core/npu/NPUStream.cpp+51-0

	@@ -437,6 +437,57 @@ NPUStream getCurrentNPUStream(c10::DeviceIndex device_index)
437	return NPUStream_fromInternals(current_streams[device_index]);	437	return NPUStream_fromInternals(current_streams[device_index]);
438	}	438	}
439		439
		440	+NPUStream getNPUStreamFromManagedAclrtStream(aclrtStream stream, c10::DeviceIndex device_index)
		441	+{
		442	+ initNPUStreamsOnce();
		443	+ TORCH_CHECK(stream != nullptr, "stream is nullptr", PTA_ERROR(ErrCode::PTR));
		444	+
		445	+ auto find_stream_in_device = [stream](c10::DeviceIndex idx) -> const LeakyStreamInternals* {
		446	+ if (default_streams[idx].stream == stream) {
		447	+ return &default_streams[idx];
		448	+ }
		449	+ if (secondary_streams[idx].stream == stream) {
		450	+ return &secondary_streams[idx];
		451	+ }
		452	+ for (const auto priority : c10::irange(kMaxStreamPriorities)) {
		453	+ for (const auto stream_idx : c10::irange(GetStreamsPerPool())) {
		454	+ if (npu_streams[priority][idx][stream_idx].stream == stream) {
		455	+ return &npu_streams[priority][idx][stream_idx];
		456	+ }
		457	+ }
		458	+ }
		459	+ for (const auto stream_idx : c10::irange(kSyncLaunchStreamsPerPool)) {
		460	+ if (sync_launch_streams[idx][stream_idx].stream == stream) {
		461	+ return &sync_launch_streams[idx][stream_idx];
		462	+ }
		463	+ }
		464	+ return nullptr;
		465	+ };
		466	+
		467	+ if (device_index != -1) {
		468	+ check_npu(device_index);
		469	+ if (auto ptr = find_stream_in_device(device_index)) {
		470	+ return NPUStream_fromInternals(ptr);
		471	+ }
		472	+ } else {
		473	+ for (const auto idx : c10::irange(num_npus)) {
		474	+ if (auto ptr = find_stream_in_device(idx)) {
		475	+ return NPUStream_fromInternals(ptr);
		476	+ }
		477	+ }
		478	+ }
		479	+
		480	+ if (device_index != -1) {
		481	+ TORCH_CHECK(
		482	+ false,
		483	+ "The aclrtStream is not managed by torch_npu on device ",
		484	+ device_index,
		485	+ PTA_ERROR(ErrCode::VALUE));
		486	+ }
		487	+
		488	+ TORCH_CHECK(false, "The aclrtStream is not managed by torch_npu", PTA_ERROR(ErrCode::VALUE));
		489	+}
		490	+
440	NPUStream getCurrentSecondaryStream(c10::DeviceIndex device_index)	491	NPUStream getCurrentSecondaryStream(c10::DeviceIndex device_index)
441	{	492	{
442	initNPUStreamsOnce();	493	initNPUStreamsOnce();

Mtorch_npu/csrc/inductor/aoti_runner/model_container_runner_npu.cpp+7-4

	@@ -56,16 +56,19 @@ std::vector<at::Tensor> AOTIModelContainerRunnerNpu::run_impl(std::vector<AtenTe
56	void* stream_handle)	56	void* stream_handle)
57	{	57	{
58	init_proxy_executor();	58	init_proxy_executor();
59	- c10_npu::NPUStream npu_stream = c10_npu::getCurrentNPUStream();	59	+ void* effective_stream_handle = stream_handle;
60	- return AOTIModelContainerRunner::run_impl(input_handles, reinterpret_cast<void*>(npu_stream.stream()));	60	+ if (effective_stream_handle == nullptr) {
		61	+ c10_npu::NPUStream npu_stream = c10_npu::getCurrentNPUStream();
		62	+ effective_stream_handle = reinterpret_cast<void*>(npu_stream.stream());
		63	+ }
		64	+ return AOTIModelContainerRunner::run_impl(input_handles, effective_stream_handle);
61	}	65	}
62		66
63	std::vector<at::Tensor> AOTIModelContainerRunnerNpu::run_with_npu_stream(const std::vector<at::Tensor>& inputs,	67	std::vector<at::Tensor> AOTIModelContainerRunnerNpu::run_with_npu_stream(const std::vector<at::Tensor>& inputs,
64	const c10_npu::NPUStream& npu_stream)	68	const c10_npu::NPUStream& npu_stream)
65	{	69	{
66	init_proxy_executor();	70	init_proxy_executor();
67	- c10_npu::NPUStream cur_npu_stream = c10_npu::getCurrentNPUStream();	71	+ return run(inputs, reinterpret_cast<void*>(npu_stream.stream()));
68	- return run(inputs, reinterpret_cast<void*>(cur_npu_stream.stream()));
69	}	72	}
70		73
71	namespace {	74	namespace {

Mtorch_npu/csrc/inductor/aoti_runtime/model.h+28-5

	@@ -149,6 +149,20 @@ public:
149	AOTInductorModelBase(const AOTInductorModelBase&) = delete;	149	AOTInductorModelBase(const AOTInductorModelBase&) = delete;
150	AOTInductorModelBase& operator=(const AOTInductorModelBase&) = delete;	150	AOTInductorModelBase& operator=(const AOTInductorModelBase&) = delete;
151		151
		152	+#if defined(USE_NPU)
		153	+ DeviceStreamType normalize_run_stream(DeviceStreamType stream) const
		154	+ {
		155	+ if (stream != nullptr) {
		156	+ return stream;
		157	+ }
		158	+
		159	+ DeviceStreamType current_stream = nullptr;
		160	+ AOTI_TORCH_ERROR_CODE_CHECK(
		161	+ aoti_torch_get_current_npu_stream(device_idx_, reinterpret_cast<void**>(&current_stream)));
		162	+ return current_stream;
		163	+ }
		164	+#endif
		165	+
152	void run(AtenTensorHandle* input_handles, // array of input AtenTensorHandle; handles	166	void run(AtenTensorHandle* input_handles, // array of input AtenTensorHandle; handles
153	// are stolen; the array itself is borrowed	167	// are stolen; the array itself is borrowed
154	AtenTensorHandle* output_handles, // array for writing output AtenTensorHandle; handles	168	AtenTensorHandle* output_handles, // array for writing output AtenTensorHandle; handles
	@@ -157,20 +171,22 @@ public:
157	DeviceStreamType stream, AOTIProxyExecutorHandle proxy_executor)	171	DeviceStreamType stream, AOTIProxyExecutorHandle proxy_executor)
158	{	172	{
159	#if defined(USE_NPU)	173	#if defined(USE_NPU)
		174	+ auto run_stream = normalize_run_stream(stream);
160	if (!run_finished_) {	175	if (!run_finished_) {
161	aclrtEvent run_finished;	176	aclrtEvent run_finished;
162	AOTI_RUNTIME_DEVICE_CHECK(aclrtCreateEvent(&run_finished));	177	AOTI_RUNTIME_DEVICE_CHECK(aclrtCreateEvent(&run_finished));
163	run_finished_.emplace(run_finished);	178	run_finished_.emplace(run_finished);
164	}	179	}
165	#else	180	#else
		181	+ auto run_stream = stream;
166	run_finished_ = false;	182	run_finished_ = false;
167	#endif	183	#endif
168		184
169	auto* model = static_cast<Model*>(this);	185	auto* model = static_cast<Model*>(this);
170	- model->run_impl(input_handles, output_handles, stream, proxy_executor);	186	+ model->run_impl(input_handles, output_handles, run_stream, proxy_executor);
171		187
172	#if defined(USE_NPU)	188	#if defined(USE_NPU)
173	- AOTI_RUNTIME_DEVICE_CHECK(aclrtRecordEvent(*run_finished_, stream));	189	+ AOTI_RUNTIME_DEVICE_CHECK(aclrtRecordEvent(*run_finished_, run_stream));
174	#else	190	#else
175	run_finished_ = true;	191	run_finished_ = true;
176	#endif	192	#endif
	@@ -185,30 +201,37 @@ public:
185	// borrowed	201	// borrowed
186	DeviceStreamType stream, AOTIProxyExecutorHandle proxy_executor)	202	DeviceStreamType stream, AOTIProxyExecutorHandle proxy_executor)
187	{	203	{
		204	+#if defined(USE_NPU)
		205	+ auto run_stream = normalize_run_stream(stream);
		206	+#else
		207	+ auto run_stream = stream;
		208	+#endif
188	// don't bother with any of the run_finished stuff; this is unsafe to call	209	// don't bother with any of the run_finished stuff; this is unsafe to call
189	// in a threaded context	210	// in a threaded context
190	auto* model = static_cast<Model*>(this);	211	auto* model = static_cast<Model*>(this);
191	- model->run_impl(input_handles, output_handles, stream, proxy_executor);	212	+ model->run_impl(input_handles, output_handles, run_stream, proxy_executor);
192	}	213	}
193		214
194	std::unordered_map<std::string, AtenTensorHandle>	215	std::unordered_map<std::string, AtenTensorHandle>
195	run_const_fold(DeviceStreamType stream, AOTIProxyExecutorHandle proxy_executor, bool initialization = false)	216	run_const_fold(DeviceStreamType stream, AOTIProxyExecutorHandle proxy_executor, bool initialization = false)
196	{	217	{
197	#if defined(USE_NPU)	218	#if defined(USE_NPU)
		219	+ auto run_stream = normalize_run_stream(stream);
198	if (!run_finished_) {	220	if (!run_finished_) {
199	aclrtEvent run_finished;	221	aclrtEvent run_finished;
200	AOTI_RUNTIME_DEVICE_CHECK(aclrtCreateEvent(&run_finished));	222	AOTI_RUNTIME_DEVICE_CHECK(aclrtCreateEvent(&run_finished));
201	run_finished_.emplace(run_finished);	223	run_finished_.emplace(run_finished);
202	}	224	}
203	#else	225	#else
		226	+ auto run_stream = stream;
204	run_finished_ = false;	227	run_finished_ = false;
205	#endif	228	#endif
206		229
207	auto* model = static_cast<Model*>(this);	230	auto* model = static_cast<Model*>(this);
208	- auto folded_constants = model->const_run_impl(stream, proxy_executor, initialization);	231	+ auto folded_constants = model->const_run_impl(run_stream, proxy_executor, initialization);
209		232
210	#if defined(USE_NPU)	233	#if defined(USE_NPU)
211	- AOTI_RUNTIME_DEVICE_CHECK(aclrtRecordEvent(*run_finished_, stream));	234	+ AOTI_RUNTIME_DEVICE_CHECK(aclrtRecordEvent(*run_finished_, run_stream));
212	#else	235	#else
213	run_finished_ = true;	236	run_finished_ = true;
214	#endif	237	#endif

Atorch_npu/csrc/inductor/aoti_runtime/utils_npu.h+54-0

@@ -0,0 +1,54 @@
	1	+#pragma once
	2	+
	3	+#if defined(USE_NPU)
	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_npu/csrc/inductor/aoti_torch/c/shim.h. The same rule
	7	+// applies to other files under torch_npu/csrc/inductor/aoti_runtime/.
	8	+#include <torch_npu/csrc/inductor/aoti_runtime/utils.h>
	9	+
	10	+namespace torch::aot_inductor {
	11	+
	12	+inline void delete_npu_guard(void* ptr)
	13	+{
	14	+ AOTI_TORCH_ERROR_CODE_CHECK(aoti_torch_delete_npu_guard(reinterpret_cast<NPUGuardHandle>(ptr)));
	15	+}
	16	+
	17	+inline void delete_npu_stream_guard(void* ptr)
	18	+{
	19	+ AOTI_TORCH_ERROR_CODE_CHECK(aoti_torch_delete_npu_stream_guard(reinterpret_cast<NPUStreamGuardHandle>(ptr)));
	20	+}
	21	+
	22	+class AOTINpuGuard {
	23	+public:
	24	+ explicit AOTINpuGuard(int32_t device_index) : guard_(nullptr, delete_npu_guard)
	25	+ {
	26	+ NPUGuardHandle ptr = nullptr;
	27	+ AOTI_TORCH_ERROR_CODE_CHECK(aoti_torch_create_npu_guard(device_index, &ptr));
	28	+ guard_.reset(ptr);
	29	+ }
	30	+
	31	+ void set_index(int32_t device_index)
	32	+ {
	33	+ AOTI_TORCH_ERROR_CODE_CHECK(aoti_torch_npu_guard_set_index(guard_.get(), device_index));
	34	+ }
	35	+
	36	+private:
	37	+ std::unique_ptr<NPUGuardOpaque, DeleterFnPtr> guard_;
	38	+};
	39	+
	40	+class AOTINpuStreamGuard {
	41	+public:
	42	+ AOTINpuStreamGuard(void* stream, int32_t device_index) : guard_(nullptr, delete_npu_stream_guard)
	43	+ {
	44	+ NPUStreamGuardHandle ptr = nullptr;
	45	+ AOTI_TORCH_ERROR_CODE_CHECK(aoti_torch_create_npu_stream_guard(stream, device_index, &ptr));
	46	+ guard_.reset(ptr);
	47	+ }
	48	+
	49	+private:
	50	+ std::unique_ptr<NPUStreamGuardOpaque, DeleterFnPtr> guard_;
	51	+};
	52	+
	53	+} // namespace torch::aot_inductor
	54	+#endif // USE_NPU

Mtorch_npu/csrc/inductor/aoti_torch/c/shim.h+24-0

	@@ -264,6 +264,30 @@ AOTI_TORCH_EXPORT AOTITorchError aoti_torch_create_tensor_from_blob_npu_v2(
264	AtenTensorHandle* ret, // returns new reference	264	AtenTensorHandle* ret, // returns new reference
265	int32_t layout, const uint8_t* opaque_metadata, int64_t opaque_metadata_size);	265	int32_t layout, const uint8_t* opaque_metadata, int64_t opaque_metadata_size);
266		266
		267	+struct NPUGuardOpaque;
		268	+using NPUGuardHandle = NPUGuardOpaque*;
		269	+
		270	+AOTI_TORCH_EXPORT AOTITorchError aoti_torch_create_npu_guard(
		271	+ int32_t device_index,
		272	+ NPUGuardHandle* ret_guard // returns new reference
		273	+);
		274	+
		275	+AOTI_TORCH_EXPORT AOTITorchError aoti_torch_delete_npu_guard(NPUGuardHandle guard);
		276	+
		277	+AOTI_TORCH_EXPORT AOTITorchError aoti_torch_npu_guard_set_index(NPUGuardHandle guard, int32_t device_index);
		278	+
		279	+struct NPUStreamGuardOpaque;
		280	+using NPUStreamGuardHandle = NPUStreamGuardOpaque*;
		281	+
		282	+AOTI_TORCH_EXPORT AOTITorchError aoti_torch_create_npu_stream_guard(
		283	+ void* stream, int32_t device_index,
		284	+ NPUStreamGuardHandle* ret_guard // returns new reference
		285	+);
		286	+
		287	+AOTI_TORCH_EXPORT AOTITorchError aoti_torch_delete_npu_stream_guard(NPUStreamGuardHandle guard);
		288	+
		289	+AOTI_TORCH_EXPORT AOTITorchError aoti_torch_get_current_npu_stream(int32_t device_index, void** ret_stream);
		290	+
267	AOTI_TORCH_EXPORT AOTITorchError aoti_torch__embedding_bag(AtenTensorHandle weight, AtenTensorHandle indices,	291	AOTI_TORCH_EXPORT AOTITorchError aoti_torch__embedding_bag(AtenTensorHandle weight, AtenTensorHandle indices,
268	AtenTensorHandle offsets, int32_t scale_grad_by_freq,	292	AtenTensorHandle offsets, int32_t scale_grad_by_freq,
269	int32_t mode, int32_t sparse,	293	int32_t mode, int32_t sparse,

Mtorch_npu/csrc/inductor/aoti_torch/shim_npu.cpp+67-4

	@@ -3,6 +3,8 @@
3	#include <c10/core/DeviceType.h>	3	#include <c10/core/DeviceType.h>
4	#include <c10/core/ScalarType.h>	4	#include <c10/core/ScalarType.h>
5	#include <torch_npu/csrc/aten/common/from_blob.h>	5	#include <torch_npu/csrc/aten/common/from_blob.h>
		6	+#include <torch_npu/csrc/core/npu/NPUGuard.h>
		7	+#include <torch_npu/csrc/core/npu/NPUStream.h>
6	#include <torch_npu/csrc/inductor/aoti_torch/c/shim.h>	8	#include <torch_npu/csrc/inductor/aoti_torch/c/shim.h>
7	#include <torch_npu/csrc/inductor/aoti_torch/utils.h>	9	#include <torch_npu/csrc/inductor/aoti_torch/utils.h>
8	#include <torch_npu/csrc/inductor/inductor_ops.h>	10	#include <torch_npu/csrc/inductor/inductor_ops.h>
	@@ -18,6 +20,10 @@ int32_t aoti_torch_device_type_npu() { return (int32_t)c10::DeviceType::PrivateU
18	} // extern "C"	20	} // extern "C"
19	#endif	21	#endif
20		22
		23	+namespace c10_npu {
		24	+NPUStream getNPUStreamFromManagedAclrtStream(aclrtStream stream, c10::DeviceIndex device_index);
		25	+}
		26	+
21	namespace {	27	namespace {
22	static c10::Device c10_device(int32_t device_type, int32_t device_index)	28	static c10::Device c10_device(int32_t device_type, int32_t device_index)
23	{	29	{
	@@ -35,6 +41,8 @@ AOTITorchError aoti_torch_create_tensor_from_blob_npu(void* data, int64_t ndim,
35	AtenTensorHandle* ret_new_tensor)	41	AtenTensorHandle* ret_new_tensor)
36	{	42	{
37	AOTI_TORCH_CONVERT_EXCEPTION_TO_ERROR_CODE({	43	AOTI_TORCH_CONVERT_EXCEPTION_TO_ERROR_CODE({
		44	+ TORCH_CHECK(ret_new_tensor != nullptr, "ret_new_tensor is nullptr");
		45	+ *ret_new_tensor = nullptr;
38	c10::IntArrayRef sizes(sizes_ptr, ndim);	46	c10::IntArrayRef sizes(sizes_ptr, ndim);
39	c10::IntArrayRef strides(strides_ptr, ndim);	47	c10::IntArrayRef strides(strides_ptr, ndim);
40	c10::Device device = c10_device(device_type, device_index);	48	c10::Device device = c10_device(device_type, device_index);
	@@ -53,11 +61,66 @@ AOTITorchError aoti_torch_create_tensor_from_blob_npu_v2(void* data, int64_t ndi
53	const uint8_t* opaque_metadata, int64_t opaque_metadata_size)	61	const uint8_t* opaque_metadata, int64_t opaque_metadata_size)
54	{	62	{
55	AOTI_TORCH_CONVERT_EXCEPTION_TO_ERROR_CODE({	63	AOTI_TORCH_CONVERT_EXCEPTION_TO_ERROR_CODE({
		64	+ TORCH_CHECK(ret_new_tensor != nullptr, "ret_new_tensor is nullptr");
		65	+ *ret_new_tensor = nullptr;
56	if (layout == static_cast<int32_t>(at::kMkldnn)) {	66	if (layout == static_cast<int32_t>(at::kMkldnn)) {
57	- throw std::runtime_error("do not support mkldnn on npu.");	67	+ TORCH_CHECK(false, "do not support mkldnn on npu.");
58	} else {	68	} else {
59	- aoti_torch_create_tensor_from_blob_npu(data, ndim, sizes_ptr, strides_ptr, storage_offset, dtype,	69	+ auto err = aoti_torch_create_tensor_from_blob_npu(data, ndim, sizes_ptr, strides_ptr, storage_offset,
60	- device_type, device_index, ret_new_tensor);	70	+ dtype, device_type, device_index, ret_new_tensor);
		71	+ if (err != AOTI_TORCH_SUCCESS) {
		72	+ return err;
		73	+ }
61	}	74	}
62	});	75	});
63	-}	76	+}
		77	+
		78	+AOTITorchError aoti_torch_create_npu_guard(int32_t device_index, NPUGuardHandle* ret_guard)
		79	+{
		80	+ AOTI_TORCH_CONVERT_EXCEPTION_TO_ERROR_CODE({
		81	+ TORCH_CHECK(ret_guard != nullptr, "ret_guard is nullptr");
		82	+ *ret_guard = nullptr;
		83	+ c10_npu::NPUGuard* guard = new c10_npu::NPUGuard(static_cast<c10::DeviceIndex>(device_index));
		84	+ *ret_guard = reinterpret_cast<NPUGuardHandle>(guard);
		85	+ });
		86	+}
		87	+
		88	+AOTITorchError aoti_torch_delete_npu_guard(NPUGuardHandle guard)
		89	+{
		90	+ AOTI_TORCH_CONVERT_EXCEPTION_TO_ERROR_CODE({ delete reinterpret_cast<c10_npu::NPUGuard*>(guard); });
		91	+}
		92	+
		93	+AOTITorchError aoti_torch_npu_guard_set_index(NPUGuardHandle guard, int32_t device_index)
		94	+{
		95	+ AOTI_TORCH_CONVERT_EXCEPTION_TO_ERROR_CODE({
		96	+ reinterpret_cast<c10_npu::NPUGuard*>(guard)->set_index(static_cast<c10::DeviceIndex>(device_index));
		97	+ });
		98	+}
		99	+
		100	+AOTITorchError aoti_torch_create_npu_stream_guard(void* stream, int32_t device_index, NPUStreamGuardHandle* ret_guard)
		101	+{
		102	+ AOTI_TORCH_CONVERT_EXCEPTION_TO_ERROR_CODE({
		103	+ TORCH_CHECK(ret_guard != nullptr, "ret_guard is nullptr");
		104	+ *ret_guard = nullptr;
		105	+
		106	+ auto raw_stream = static_cast<aclrtStream>(stream);
		107	+ auto managed_stream = c10_npu::getNPUStreamFromManagedAclrtStream(
		108	+ raw_stream, static_cast<c10::DeviceIndex>(device_index));
		109	+ auto* guard = new c10_npu::NPUStreamGuard(static_cast<c10::Stream>(managed_stream));
		110	+ *ret_guard = reinterpret_cast<NPUStreamGuardHandle>(guard);
		111	+ });
		112	+}
		113	+
		114	+AOTITorchError aoti_torch_delete_npu_stream_guard(NPUStreamGuardHandle guard)
		115	+{
		116	+ AOTI_TORCH_CONVERT_EXCEPTION_TO_ERROR_CODE({ delete reinterpret_cast<c10_npu::NPUStreamGuard*>(guard); });
		117	+}
		118	+
		119	+AOTITorchError aoti_torch_get_current_npu_stream(int32_t device_index, void** ret_stream)
		120	+{
		121	+ AOTI_TORCH_CONVERT_EXCEPTION_TO_ERROR_CODE({
		122	+ TORCH_CHECK(ret_stream != nullptr, "ret_stream is nullptr");
		123	+ ret_stream = reinterpret_cast<void>(
		124	+ c10_npu::getCurrentNPUStream(static_cast<c10::DeviceIndex>(device_index)).stream());
		125	+ });
		126	+}