feat add shim_npu-pytorch-AtomGit

已合并

feat add shim_npu #35591

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

feat add shim_npu #35591

已合并

zzhongmin创建于 5月13日

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

共 13 个文件变更+1214-30

Atest/_inductor/test_npu_aoti_shim.py+184-0

@@ -0,0 +1,184 @@
	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	+
	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_allocator_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_allocator_output_handle(x), 1)
	140	+
	141	+ def test_null_blob_tensor_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_output_handle(x), 1)
	145	+
	146	+ def test_null_blob_tensor_v2_output_handle_returns_failure_internally(self):
	147	+ x = torch.randn(4, device="npu:0", dtype=torch.float32)
	148	+
	149	+ self.assertEqual(self.module.check_null_blob_tensor_v2_output_handle(x), 1)
	150	+
	151	+ def test_current_device_stream_lookup_uses_negative_one_semantics(self):
	152	+ x = torch.randn(4, device="npu:0", dtype=torch.float32)
	153	+
	154	+ self.assertEqual(self.module.check_current_device_stream_lookup(x), 1)
	155	+
	156	+ def test_default_stream_guard_roundtrip_restores_original_stream(self):
	157	+ x = torch.randn(4, device="npu:0", dtype=torch.float32)
	158	+
	159	+ self.assertEqual(self.module.check_default_stream_guard_roundtrip(x), 1)
	160	+
	161	+ def test_run_npu_shim_checks_restores_device_and_stream(self):
	162	+ device_count = torch.npu.device_count()
	163	+ target_device = 1 if device_count > 1 else 0
	164	+ original_device = 0 if target_device != 0 else target_device
	165	+
	166	+ torch.npu.set_device(original_device)
	167	+ x = torch.randn(16, device=f"npu:{target_device}", dtype=torch.float32)
	168	+ expected = x + 1
	169	+
	170	+ with torch.npu.device(target_device):
	171	+ custom_stream = torch.npu.Stream()
	172	+ with torch.npu.stream(custom_stream):
	173	+ before_stream = torch.npu.current_stream().npu_stream
	174	+ result = self.module.run_npu_shim_checks(x)
	175	+ after_stream = torch.npu.current_stream().npu_stream
	176	+
	177	+ torch.testing.assert_close(result, expected)
	178	+ self.assertEqual(before_stream, custom_stream.npu_stream)
	179	+ self.assertEqual(after_stream, custom_stream.npu_stream)
	180	+ self.assertEqual(torch.npu.current_device(), original_device)
	181	+
	182	+
	183	+if __name__ == "__main__":
	184	+ run_tests()

Mtest/_inductor/test_npu_device.py+9-13

	@@ -1,17 +1,12 @@
1	-import torch
2	-from torch_npu.npu import device_count
3	-from torch_npu.utils._dynamo_device import NpuInterface, current_device, set_device
4	-from torch_npu.utils._inductor import NPUDeviceOpOverrides
5	-from torch_npu._inductor.config import config as npu_config
6	from torch_npu._inductor.npu_device import NewNPUDeviceOpOverrides	1	from torch_npu._inductor.npu_device import NewNPUDeviceOpOverrides
7	-from torch_npu.testing.testcase import TestCase, run_tests	2	+from torch_npu.testing.testcase import run_tests, TestCase
8		3
9		4
10	class TestNpuDevice(TestCase):	5	class TestNpuDevice(TestCase):
11	def test_aoti_get_stream(self):	6	def test_aoti_get_stream(self):
12	overrides = NewNPUDeviceOpOverrides()	7	overrides = NewNPUDeviceOpOverrides()
13	result = overrides.aoti_get_stream()	8	result = overrides.aoti_get_stream()
14	- excepted = "aoti_torch_get_current_cuda_stream"	9	+ excepted = "aoti_torch_get_current_npu_stream"
15	self.assertEqual(result, excepted)	10	self.assertEqual(result, excepted)
16		11
17	def test_cpp_stream_type(self):	12	def test_cpp_stream_type(self):
	@@ -36,18 +31,19 @@ class TestNpuDevice(TestCase):
36	self.assertIn("#include <sys/syscall.h>", result)	31	self.assertIn("#include <sys/syscall.h>", result)
37	self.assertIn("#include <torch_npu/csrc/framework/OpCommand.h>", result)	32	self.assertIn("#include <torch_npu/csrc/framework/OpCommand.h>", result)
38	self.assertIn("#include <torch_npu/csrc/core/npu/NPUStream.h>", result)	33	self.assertIn("#include <torch_npu/csrc/core/npu/NPUStream.h>", result)
39	- self.assertIn("#include \"runtime/runtime/rt.h\"", result)	34	+ self.assertIn('#include "runtime/runtime/rt.h"', result)
40		35
41	def test_cpp_aoti_stream_guard(self):	36	def test_cpp_aoti_stream_guard(self):
42	overrides = NewNPUDeviceOpOverrides()	37	overrides = NewNPUDeviceOpOverrides()
43	result = overrides.cpp_aoti_stream_guard()	38	result = overrides.cpp_aoti_stream_guard()
44	- excepted = "AOTICudaStreamGuard"	39	+ excepted = "AOTINpuStreamGuard"
45	self.assertEqual(result, excepted)	40	self.assertEqual(result, excepted)
46		41
47	- def test_cpp_aoti_device_guard_not_implemented(self):	42	+ def test_cpp_aoti_device_guard(self):
48	overrides = NewNPUDeviceOpOverrides()	43	overrides = NewNPUDeviceOpOverrides()
49	- with self.assertRaises(NotImplementedError):	44	+ result = overrides.cpp_aoti_device_guard()
50	- overrides.cpp_aoti_device_guard()	45	+ excepted = "AOTINpuGuard"
		46	+ self.assertEqual(result, excepted)
51		47
52	def test_device_guard(self):	48	def test_device_guard(self):
53	overrides = NewNPUDeviceOpOverrides()	49	overrides = NewNPUDeviceOpOverrides()
	@@ -78,4 +74,4 @@ class TestNpuDevice(TestCase):
78		74
79		75
80	if __name__ == "__main__":	76	if __name__ == "__main__":
81	- run_tests()	77	+ run_tests()

Atest/cpp_extensions/npu_aoti_shim_extension.cpp+689-0

@@ -0,0 +1,689 @@
	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	+ check_aoti_error(
	224	+ aoti_torch_npu_caching_allocator_raw_delete(nullptr),
	225	+ "aoti_torch_npu_caching_allocator_raw_delete(nullptr)");
	226	+ return 1;
	227	+}
	228	+
	229	+int64_t check_invalid_stream_guard_path_impl(const at::Tensor& tensor)
	230	+{
	231	+ assert_npu_tensor(tensor, "check_invalid_stream_guard_path");
	232	+
	233	+ NPUStreamGuardHandle guard = nullptr;
	234	+ auto err = aoti_torch_create_npu_stream_guard(
	235	+ reinterpret_cast<void*>(0x1),
	236	+ tensor.device().index(),
	237	+ &guard);
	238	+ TORCH_CHECK(
	239	+ err == AOTI_TORCH_FAILURE,
	240	+ "Invalid stream pointer should make aoti_torch_create_npu_stream_guard fail");
	241	+ TORCH_CHECK(guard == nullptr, "Guard should remain nullptr when stream guard creation fails");
	242	+ return 1;
	243	+}
	244	+
	245	+int64_t check_null_stream_guard_path_impl(const at::Tensor& tensor)
	246	+{
	247	+ assert_npu_tensor(tensor, "check_null_stream_guard_path");
	248	+
	249	+ NPUStreamGuardHandle guard = nullptr;
	250	+ auto err =
	251	+ aoti_torch_create_npu_stream_guard(nullptr, tensor.device().index(), &guard);
	252	+ TORCH_CHECK(
	253	+ err == AOTI_TORCH_FAILURE,
	254	+ "Null stream pointer should make aoti_torch_create_npu_stream_guard fail");
	255	+ TORCH_CHECK(guard == nullptr, "Guard should remain nullptr when null stream guard creation fails");
	256	+ return 1;
	257	+}
	258	+
	259	+int64_t check_invalid_device_guard_creation_impl(const at::Tensor& tensor)
	260	+{
	261	+ assert_npu_tensor(tensor, "check_invalid_device_guard_creation");
	262	+
	263	+ const auto invalid_device_index = static_cast<int32_t>(c10_npu::device_count());
	264	+ NPUGuardHandle guard = nullptr;
	265	+ auto err = aoti_torch_create_npu_guard(invalid_device_index, &guard);
	266	+ TORCH_CHECK(
	267	+ err == AOTI_TORCH_FAILURE,
	268	+ "Invalid device index should make aoti_torch_create_npu_guard fail");
	269	+ TORCH_CHECK(guard == nullptr, "Guard should remain nullptr when invalid guard creation fails");
	270	+ return 1;
	271	+}
	272	+
	273	+int64_t check_invalid_device_guard_set_index_impl(const at::Tensor& tensor)
	274	+{
	275	+ assert_npu_tensor(tensor, "check_invalid_device_guard_set_index");
	276	+
	277	+ const auto original_device = c10_npu::current_device();
	278	+ const auto valid_device_index = tensor.device().index();
	279	+ const auto invalid_device_index = static_cast<int32_t>(c10_npu::device_count());
	280	+ NPUGuardHandle guard = nullptr;
	281	+ check_aoti_error(
	282	+ aoti_torch_create_npu_guard(valid_device_index, &guard),
	283	+ "aoti_torch_create_npu_guard(valid for set_index)");
	284	+ TORCH_CHECK(
	285	+ c10_npu::current_device() == valid_device_index,
	286	+ "Valid guard creation should switch to the requested device");
	287	+
	288	+ auto err = aoti_torch_npu_guard_set_index(guard, invalid_device_index);
	289	+ TORCH_CHECK(
	290	+ err == AOTI_TORCH_FAILURE,
	291	+ "Invalid device index should make aoti_torch_npu_guard_set_index fail");
	292	+ TORCH_CHECK(
	293	+ c10_npu::current_device() == valid_device_index,
	294	+ "Failed guard set_index should keep the previously selected device");
	295	+
	296	+ check_aoti_error(
	297	+ aoti_torch_delete_npu_guard(guard),
	298	+ "aoti_torch_delete_npu_guard(valid after failed set_index)");
	299	+ TORCH_CHECK(
	300	+ c10_npu::current_device() == original_device,
	301	+ "Deleting the guard should restore the original device after failed set_index");
	302	+ return 1;
	303	+}
	304	+
	305	+int64_t check_invalid_device_current_stream_impl(const at::Tensor& tensor)
	306	+{
	307	+ assert_npu_tensor(tensor, "check_invalid_device_current_stream");
	308	+
	309	+ const auto invalid_device_index = static_cast<int32_t>(c10_npu::device_count());
	310	+ void* stream = nullptr;
	311	+ auto err = aoti_torch_get_current_npu_stream(invalid_device_index, &stream);
	312	+ TORCH_CHECK(
	313	+ err == AOTI_TORCH_FAILURE,
	314	+ "Invalid device index should make aoti_torch_get_current_npu_stream fail");
	315	+ TORCH_CHECK(stream == nullptr, "Invalid device stream lookup should leave stream as nullptr");
	316	+ return 1;
	317	+}
	318	+
	319	+int64_t check_null_stream_guard_output_handle_impl(const at::Tensor& tensor)
	320	+{
	321	+ assert_npu_tensor(tensor, "check_null_stream_guard_output_handle");
	322	+
	323	+ auto pooled_stream = c10_npu::getNPUStreamFromPool(tensor.device().index());
	324	+ remember_stream(pooled_stream);
	325	+ auto err = aoti_torch_create_npu_stream_guard(
	326	+ reinterpret_cast<void*>(pooled_stream.stream(false)),
	327	+ tensor.device().index(),
	328	+ nullptr);
	329	+ TORCH_CHECK(
	330	+ err == AOTI_TORCH_FAILURE,
	331	+ "Null ret_guard should make aoti_torch_create_npu_stream_guard fail");
	332	+ return 1;
	333	+}
	334	+
	335	+int64_t check_null_current_stream_output_handle_impl(const at::Tensor& tensor)
	336	+{
	337	+ assert_npu_tensor(tensor, "check_null_current_stream_output_handle");
	338	+
	339	+ auto err = aoti_torch_get_current_npu_stream(tensor.device().index(), nullptr);
	340	+ TORCH_CHECK(
	341	+ err == AOTI_TORCH_FAILURE,
	342	+ "Null ret_stream should make aoti_torch_get_current_npu_stream fail");
	343	+ return 1;
	344	+}
	345	+
	346	+int64_t check_null_guard_output_handle_impl(const at::Tensor& tensor)
	347	+{
	348	+ assert_npu_tensor(tensor, "check_null_guard_output_handle");
	349	+
	350	+ auto err = aoti_torch_create_npu_guard(tensor.device().index(), nullptr);
	351	+ TORCH_CHECK(
	352	+ err == AOTI_TORCH_FAILURE,
	353	+ "Null ret_guard should make aoti_torch_create_npu_guard fail");
	354	+ return 1;
	355	+}
	356	+
	357	+int64_t check_null_allocator_output_handle_impl(const at::Tensor& tensor)
	358	+{
	359	+ assert_npu_tensor(tensor, "check_null_allocator_output_handle");
	360	+
	361	+ auto err = aoti_torch_npu_caching_allocator_raw_alloc(64, nullptr);
	362	+ TORCH_CHECK(
	363	+ err == AOTI_TORCH_FAILURE,
	364	+ "Null ret_ptr should make aoti_torch_npu_caching_allocator_raw_alloc fail");
	365	+ return 1;
	366	+}
	367	+
	368	+int64_t check_null_blob_tensor_output_handle_impl(const at::Tensor& tensor)
	369	+{
	370	+ assert_npu_tensor(tensor, "check_null_blob_tensor_output_handle");
	371	+
	372	+ auto contiguous = tensor.contiguous();
	373	+ auto sizes = to_vector(contiguous.sizes());
	374	+ auto strides = to_vector(contiguous.strides());
	375	+ auto err = aoti_torch_create_tensor_from_blob_npu(
	376	+ contiguous.data_ptr(),
	377	+ contiguous.dim(),
	378	+ sizes.data(),
	379	+ strides.data(),
	380	+ contiguous.storage_offset(),
	381	+ aoti_torch_dtype_float32(),
	382	+ aoti_torch_device_type_npu(),
	383	+ tensor.device().index(),
	384	+ nullptr);
	385	+ TORCH_CHECK(
	386	+ err == AOTI_TORCH_FAILURE,
	387	+ "Null ret_new_tensor should make aoti_torch_create_tensor_from_blob_npu fail");
	388	+ return 1;
	389	+}
	390	+
	391	+int64_t check_null_blob_tensor_v2_output_handle_impl(const at::Tensor& tensor)
	392	+{
	393	+ assert_npu_tensor(tensor, "check_null_blob_tensor_v2_output_handle");
	394	+
	395	+ auto contiguous = tensor.contiguous();
	396	+ auto sizes = to_vector(contiguous.sizes());
	397	+ auto strides = to_vector(contiguous.strides());
	398	+ auto err = aoti_torch_create_tensor_from_blob_npu_v2(
	399	+ contiguous.data_ptr(),
	400	+ contiguous.dim(),
	401	+ sizes.data(),
	402	+ strides.data(),
	403	+ contiguous.storage_offset(),
	404	+ aoti_torch_dtype_float32(),
	405	+ aoti_torch_device_type_npu(),
	406	+ tensor.device().index(),
	407	+ nullptr,
	408	+ aoti_torch_layout_strided(),
	409	+ nullptr,
	410	+ 0);
	411	+ TORCH_CHECK(
	412	+ err == AOTI_TORCH_FAILURE,
	413	+ "Null ret_new_tensor should make aoti_torch_create_tensor_from_blob_npu_v2 fail");
	414	+ return 1;
	415	+}
	416	+
	417	+int64_t check_current_device_stream_lookup_impl(const at::Tensor& tensor)
	418	+{
	419	+ assert_npu_tensor(tensor, "check_current_device_stream_lookup");
	420	+
	421	+ const auto original_device = c10_npu::current_device();
	422	+ c10_npu::NPUGuard device_guard(tensor.device().index());
	423	+
	424	+ void* shim_stream = nullptr;
	425	+ check_aoti_error(
	426	+ aoti_torch_get_current_npu_stream(-1, &shim_stream),
	427	+ "aoti_torch_get_current_npu_stream(current device)");
	428	+ auto current_stream =
	429	+ reinterpret_cast<void*>(c10_npu::getCurrentNPUStream(tensor.device().index()).stream(false));
	430	+ TORCH_CHECK(
	431	+ shim_stream == current_stream,
	432	+ "Current-device stream fallback mismatch");
	433	+ TORCH_CHECK(
	434	+ c10_npu::current_device() == tensor.device().index(),
	435	+ "Current-device lookup should keep the selected device");
	436	+ TORCH_CHECK(
	437	+ c10_npu::current_device() == tensor.device().index(),
	438	+ "Device guard should keep the tensor device selected inside the scope");
	439	+ device_guard.set_index(original_device);
	440	+ TORCH_CHECK(
	441	+ c10_npu::current_device() == original_device,
	442	+ "Current-device lookup helper should restore the original device");
	443	+ return 1;
	444	+}
	445	+
	446	+int64_t check_default_stream_guard_roundtrip_impl(const at::Tensor& tensor)
	447	+{
	448	+ assert_npu_tensor(tensor, "check_default_stream_guard_roundtrip");
	449	+
	450	+ const auto device_index = tensor.device().index();
	451	+ auto pooled_stream = c10_npu::getNPUStreamFromPool(device_index);
	452	+ remember_stream(pooled_stream);
	453	+ c10_npu::NPUStreamGuard outer_guard(static_cast<c10::Stream>(pooled_stream));
	454	+
	455	+ auto original_stream =
	456	+ reinterpret_cast<void*>(c10_npu::getCurrentNPUStream(device_index).stream(false));
	457	+ auto default_stream = c10_npu::getDefaultNPUStream(device_index);
	458	+ remember_stream(default_stream);
	459	+
	460	+ NPUStreamGuardHandle stream_guard = nullptr;
	461	+ check_aoti_error(
	462	+ aoti_torch_create_npu_stream_guard(
	463	+ reinterpret_cast<void*>(default_stream.stream(false)),
	464	+ device_index,
	465	+ &stream_guard),
	466	+ "aoti_torch_create_npu_stream_guard(default stream)");
	467	+ auto guarded_stream =
	468	+ reinterpret_cast<void*>(c10_npu::getCurrentNPUStream(device_index).stream(false));
	469	+ TORCH_CHECK(
	470	+ guarded_stream == reinterpret_cast<void*>(default_stream.stream(false)),
	471	+ "Default stream guard did not switch to the default stream");
	472	+ check_aoti_error(
	473	+ aoti_torch_delete_npu_stream_guard(stream_guard),
	474	+ "aoti_torch_delete_npu_stream_guard(default stream)");
	475	+ auto restored_stream =
	476	+ reinterpret_cast<void*>(c10_npu::getCurrentNPUStream(device_index).stream(false));
	477	+ TORCH_CHECK(
	478	+ restored_stream == original_stream,
	479	+ "Default stream guard did not restore the original stream");
	480	+ return 1;
	481	+}
	482	+
	483	+at::Tensor run_npu_shim_checks_impl(const at::Tensor& tensor)
	484	+{
	485	+ assert_npu_tensor(tensor, "run_npu_shim_checks");
	486	+
	487	+ const auto device_index = tensor.device().index();
	488	+ TORCH_CHECK(
	489	+ aoti_torch_device_type_npu() ==
	490	+ static_cast<int32_t>(c10::DeviceType::PrivateUse1),
	491	+ "aoti_torch_device_type_npu should return PrivateUse1");
	492	+
	493	+ void* shim_stream = nullptr;
	494	+ check_aoti_error(
	495	+ aoti_torch_get_current_npu_stream(device_index, &shim_stream),
	496	+ "aoti_torch_get_current_npu_stream");
	497	+ auto current_stream =
	498	+ reinterpret_cast<void*>(c10_npu::getCurrentNPUStream(device_index).stream(false));
	499	+ TORCH_CHECK(
	500	+ shim_stream == current_stream,
	501	+ "Current stream mismatch: shim=",
	502	+ shim_stream,
	503	+ ", direct=",
	504	+ current_stream);
	505	+
	506	+ const auto original_device = c10_npu::current_device();
	507	+ NPUGuardHandle guard = nullptr;
	508	+ check_aoti_error(
	509	+ aoti_torch_create_npu_guard(device_index, &guard),
	510	+ "aoti_torch_create_npu_guard");
	511	+ TORCH_CHECK(
	512	+ c10_npu::current_device() == device_index,
	513	+ "NPU guard did not switch the current device");
	514	+ check_aoti_error(
	515	+ aoti_torch_npu_guard_set_index(guard, device_index),
	516	+ "aoti_torch_npu_guard_set_index");
	517	+ TORCH_CHECK(
	518	+ c10_npu::current_device() == device_index,
	519	+ "NPU guard set_index did not keep the requested device");
	520	+ check_aoti_error(
	521	+ aoti_torch_delete_npu_guard(guard),
	522	+ "aoti_torch_delete_npu_guard");
	523	+ TORCH_CHECK(
	524	+ c10_npu::current_device() == original_device,
	525	+ "Deleting the NPU guard did not restore the original device");
	526	+
	527	+ auto original_stream =
	528	+ reinterpret_cast<void*>(c10_npu::getCurrentNPUStream(device_index).stream(false));
	529	+ auto pooled_stream_obj = c10_npu::getNPUStreamFromPool(device_index);
	530	+ remember_stream(pooled_stream_obj);
	531	+ auto pooled_stream =
	532	+ reinterpret_cast<void*>(pooled_stream_obj.stream(false));
	533	+ NPUStreamGuardHandle stream_guard = nullptr;
	534	+ check_aoti_error(
	535	+ aoti_torch_create_npu_stream_guard(pooled_stream, device_index, &stream_guard),
	536	+ "aoti_torch_create_npu_stream_guard");
	537	+ auto guarded_stream =
	538	+ reinterpret_cast<void*>(c10_npu::getCurrentNPUStream(device_index).stream(false));
	539	+ TORCH_CHECK(
	540	+ guarded_stream == pooled_stream,
	541	+ "NPU stream guard did not switch to the pooled stream");
	542	+ check_aoti_error(
	543	+ aoti_torch_delete_npu_stream_guard(stream_guard),
	544	+ "aoti_torch_delete_npu_stream_guard");
	545	+ auto restored_stream =
	546	+ reinterpret_cast<void*>(c10_npu::getCurrentNPUStream(device_index).stream(false));
	547	+ TORCH_CHECK(
	548	+ restored_stream == original_stream,
	549	+ "Deleting the NPU stream guard did not restore the original stream");
	550	+
	551	+ void* zero_alloc = reinterpret_cast<void*>(0x1);
	552	+ check_aoti_error(
	553	+ aoti_torch_npu_caching_allocator_raw_alloc(0, &zero_alloc),
	554	+ "aoti_torch_npu_caching_allocator_raw_alloc(0)");
	555	+ TORCH_CHECK(
	556	+ zero_alloc == nullptr,
	557	+ "Zero-byte NPU allocation should return nullptr");
	558	+
	559	+ void* alloc_ptr = nullptr;
	560	+ check_aoti_error(
	561	+ aoti_torch_npu_caching_allocator_raw_alloc(64, &alloc_ptr),
	562	+ "aoti_torch_npu_caching_allocator_raw_alloc");
	563	+ TORCH_CHECK(alloc_ptr != nullptr, "NPU allocator returned nullptr");
	564	+ check_aoti_error(
	565	+ aoti_torch_npu_caching_allocator_raw_delete(alloc_ptr),
	566	+ "aoti_torch_npu_caching_allocator_raw_delete");
	567	+
	568	+ auto contiguous = tensor.contiguous();
	569	+ auto sizes = to_vector(contiguous.sizes());
	570	+ auto strides = to_vector(contiguous.strides());
	571	+
	572	+ AtenTensorHandle alias_handle = nullptr;
	573	+ check_aoti_error(
	574	+ aoti_torch_create_tensor_from_blob_npu(
	575	+ contiguous.data_ptr(),
	576	+ contiguous.dim(),
	577	+ sizes.data(),
	578	+ strides.data(),
	579	+ contiguous.storage_offset(),
	580	+ aoti_torch_dtype_float32(),
	581	+ aoti_torch_device_type_npu(),
	582	+ device_index,
	583	+ &alias_handle),
	584	+ "aoti_torch_create_tensor_from_blob_npu");
	585	+ at::Tensor alias =
	586	+ *torch::aot_inductor::tensor_handle_to_tensor_pointer(alias_handle);
	587	+ TORCH_CHECK(alias.device() == contiguous.device(), "Alias tensor device mismatch");
	588	+ TORCH_CHECK(alias.scalar_type() == contiguous.scalar_type(), "Alias tensor dtype mismatch");
	589	+ TORCH_CHECK(alias.data_ptr() == contiguous.data_ptr(), "Alias tensor data_ptr mismatch");
	590	+ TORCH_CHECK(
	591	+ to_vector(alias.sizes()) == sizes,
	592	+ "Alias tensor sizes mismatch");
	593	+ TORCH_CHECK(
	594	+ to_vector(alias.strides()) == strides,
	595	+ "Alias tensor strides mismatch");
	596	+ TORCH_CHECK(
	597	+ alias.storage_offset() == contiguous.storage_offset(),
	598	+ "Alias tensor storage_offset mismatch");
	599	+ TORCH_CHECK(alias.equal(contiguous), "Alias tensor value mismatch");
	600	+ check_aoti_error(
	601	+ aoti_torch_delete_tensor_object(alias_handle),
	602	+ "aoti_torch_delete_tensor_object(alias_handle)");
	603	+
	604	+ auto result = at::add(alias, 1.0);
	605	+ return result;
	606	+}
	607	+
	608	+} // namespace
	609	+
	610	+namespace c10_npu {
	611	+
	612	+NPUStream getNPUStreamFromManagedAclrtStream(aclrtStream stream, c10::DeviceIndex device_index)
	613	+{
	614	+ TORCH_CHECK(stream != nullptr, "stream is nullptr");
	615	+
	616	+ auto find_registered_stream = [stream]() -> std::optional<NPUStream> {
	617	+ std::lock_guard<std::mutex> lock(stream_registry_mutex());
	618	+ auto it = stream_registry().find(reinterpret_cast<void*>(stream));
	619	+ if (it == stream_registry().end()) {
	620	+ return std::nullopt;
	621	+ }
	622	+ return NPUStream(NPUStream::UNCHECKED, it->second);
	623	+ };
	624	+ if (auto registered = find_registered_stream()) {
	625	+ return *registered;
	626	+ }
	627	+
	628	+ auto find_known_stream = [stream](c10::DeviceIndex idx) -> std::optional<NPUStream> {
	629	+ auto current = getCurrentNPUStream(idx);
	630	+ if (current.stream(false) == stream) {
	631	+ remember_stream(current);
	632	+ return current;
	633	+ }
	634	+
	635	+ auto default_stream = getDefaultNPUStream(idx);
	636	+ if (default_stream.stream(false) == stream) {
	637	+ remember_stream(default_stream);
	638	+ return default_stream;
	639	+ }
	640	+
	641	+ return std::nullopt;
	642	+ };
	643	+
	644	+ if (device_index != -1) {
	645	+ if (auto known = find_known_stream(device_index)) {
	646	+ return *known;
	647	+ }
	648	+ } else {
	649	+ const auto device_count = c10_npu::device_count();
	650	+ for (c10::DeviceIndex idx = 0; idx < device_count; ++idx) {
	651	+ if (auto known = find_known_stream(idx)) {
	652	+ return *known;
	653	+ }
	654	+ }
	655	+ }
	656	+
	657	+ TORCH_CHECK(
	658	+ false,
	659	+ "The aclrtStream is not managed by the shim test registry on device ",
	660	+ device_index);
	661	+}
	662	+
	663	+} // namespace c10_npu
	664	+
	665	+PYBIND11_MODULE(TORCH_EXTENSION_NAME, m)
	666	+{
	667	+ m.def("get_npu_raw_stream", &get_npu_raw_stream_impl);
	668	+ m.def("make_zero_size_blob_tensor", &make_zero_size_blob_tensor_impl);
	669	+ m.def("make_zero_size_cpu_blob_tensor", &make_zero_size_cpu_blob_tensor_impl);
	670	+ m.def("check_mkldnn_blob_tensor_v2_rejected", &check_mkldnn_blob_tensor_v2_rejected_impl);
	671	+ m.def(
	672	+ "check_blob_tensor_v2_propagates_invalid_device_failure",
	673	+ &check_blob_tensor_v2_propagates_invalid_device_failure_impl);
	674	+ m.def("check_null_delete_paths", &check_null_delete_paths_impl);
	675	+ m.def("check_invalid_stream_guard_path", &check_invalid_stream_guard_path_impl);
	676	+ m.def("check_null_stream_guard_path", &check_null_stream_guard_path_impl);
	677	+ m.def("check_invalid_device_guard_creation", &check_invalid_device_guard_creation_impl);
	678	+ m.def("check_invalid_device_guard_set_index", &check_invalid_device_guard_set_index_impl);
	679	+ m.def("check_invalid_device_current_stream", &check_invalid_device_current_stream_impl);
	680	+ m.def("check_null_stream_guard_output_handle", &check_null_stream_guard_output_handle_impl);
	681	+ m.def("check_null_current_stream_output_handle", &check_null_current_stream_output_handle_impl);
	682	+ m.def("check_null_guard_output_handle", &check_null_guard_output_handle_impl);
	683	+ m.def("check_null_allocator_output_handle", &check_null_allocator_output_handle_impl);
	684	+ m.def("check_null_blob_tensor_output_handle", &check_null_blob_tensor_output_handle_impl);
	685	+ m.def("check_null_blob_tensor_v2_output_handle", &check_null_blob_tensor_v2_output_handle_impl);
	686	+ m.def("check_current_device_stream_lookup", &check_current_device_stream_lookup_impl);
	687	+ m.def("check_default_stream_guard_roundtrip", &check_default_stream_guard_roundtrip_impl);
	688	+ m.def("run_npu_shim_checks", &run_npu_shim_checks_impl);
	689	+}

Mtest/cpp_extensions/setup.py+8-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,6 +11,10 @@ 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(
	@@ -23,6 +26,10 @@ ext_modules = [
23	NpuExtension(	26	NpuExtension(
24	'torch_test_cpp_extension.stable_libtorch', ['test_stable_libtorch.cpp'],	27	'torch_test_cpp_extension.stable_libtorch', ['test_stable_libtorch.cpp'],
25	extra_compile_args=CXX_FLAGS),	28	extra_compile_args=CXX_FLAGS),
		29	+ NpuExtension(
		30	+ 'torch_test_cpp_extension.npu_aoti_shim',['npu_aoti_shim_extension.cpp', SHIM_SOURCE],
		31	+ include_dirs=[REPO_ROOT],
		32	+ extra_compile_args=CXX_FLAGS),
26	]	33	]
27		34
28	setup(	35	setup(

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
	42	+
	43	+ try:
	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

	@@ -330,6 +330,7 @@ class CppWrapperNpu(CppWrapperCpu):
330	self.header.splice("#include <unistd.h>")	330	self.header.splice("#include <unistd.h>")
331	self.header.splice("#include <filesystem>")	331	self.header.splice("#include <filesystem>")
332	self.header.splice(self.device_codegen.abi_compatible_header())	332	self.header.splice(self.device_codegen.abi_compatible_header())
		333	+ self.header.splice("#include <torch_npu/csrc/inductor/aoti_runtime/utils_npu.h>")
333	self.header.splice(	334	self.header.splice(
334	maybe_hipify_code_wrapper(self.device_codegen.kernel_driver())	335	maybe_hipify_code_wrapper(self.device_codegen.kernel_driver())
335	)	336	)

Mtorch_npu/_inductor/npu_device.py+3-3

	@@ -24,10 +24,10 @@ class NewNPUDeviceOpOverrides(NPUDeviceOpOverrides):
24	return f"torch.npu.utils.device({device_idx})"	24	return f"torch.npu.utils.device({device_idx})"
25		25
26	def cpp_aoti_device_guard(self):	26	def cpp_aoti_device_guard(self):
27	- raise NotImplementedError	27	+ return "AOTINpuGuard"
28		28
29	def cpp_aoti_stream_guard(self):	29	def cpp_aoti_stream_guard(self):
30	- return "AOTICudaStreamGuard"	30	+ return "AOTINpuStreamGuard"
31		31
32	def kernel_driver(self):	32	def kernel_driver(self):
33	source_code = """	33	source_code = """
	@@ -206,7 +206,7 @@ class NewNPUDeviceOpOverrides(NPUDeviceOpOverrides):
206	return "aclrtStream"	206	return "aclrtStream"
207		207
208	def aoti_get_stream(self):	208	def aoti_get_stream(self):
209	- return "aoti_torch_get_current_cuda_stream"	209	+ return "aoti_torch_get_current_npu_stream"
210		210
211	def cpp_kernel_type(self):	211	def cpp_kernel_type(self):
212	return "void *"	212	return "void *"

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+35-5

	@@ -45,6 +45,12 @@ using RAIIDataPtr = std::unique_ptr<void, std::function<void(void*)> >;
45		45
46	RAIIDataPtr RAII_npuMalloc(size_t num_bytes)	46	RAIIDataPtr RAII_npuMalloc(size_t num_bytes)
47	{	47	{
		48	+#ifdef AOT_INDUCTOR_USE_CACHING_ALLOCATOR
		49	+ void* data_ptr = nullptr;
		50	+ AOTI_TORCH_ERROR_CODE_CHECK(aoti_torch_npu_caching_allocator_raw_alloc(num_bytes, &data_ptr));
		51	+ auto deleter = [](void* ptr) { AOTI_TORCH_ERROR_CODE_CHECK(aoti_torch_npu_caching_allocator_raw_delete(ptr)); };
		52	+ return RAIIDataPtr(data_ptr, deleter);
		53	+#else
48	void* data_ptr;	54	void* data_ptr;
49	// aclrtMalloc doesn't support allocate 0-bytes. In this case,	55	// aclrtMalloc doesn't support allocate 0-bytes. In this case,
50	// e.g, model has no weight, we should do padding.	56	// e.g, model has no weight, we should do padding.
	@@ -54,6 +60,7 @@ RAIIDataPtr RAII_npuMalloc(size_t num_bytes)
54	AOTI_RUNTIME_DEVICE_CHECK(aclrtMalloc((void**)&data_ptr, num_bytes, ACL_MEM_MALLOC_HUGE_FIRST));	60	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)); };	61	auto deleter = [](void* ptr) { AOTI_RUNTIME_DEVICE_CHECK(aclrtFree(ptr)); };
56	return RAIIDataPtr(data_ptr, deleter);	62	return RAIIDataPtr(data_ptr, deleter);
		63	+#endif
57	}	64	}
58		65
59	#endif // USE_NPU	66	#endif // USE_NPU
	@@ -149,6 +156,20 @@ public:
149	AOTInductorModelBase(const AOTInductorModelBase&) = delete;	156	AOTInductorModelBase(const AOTInductorModelBase&) = delete;
150	AOTInductorModelBase& operator=(const AOTInductorModelBase&) = delete;	157	AOTInductorModelBase& operator=(const AOTInductorModelBase&) = delete;
151		158
		159	+#if defined(USE_NPU)
		160	+ DeviceStreamType normalize_run_stream(DeviceStreamType stream) const
		161	+ {
		162	+ if (stream != nullptr) {
		163	+ return stream;
		164	+ }
		165	+
		166	+ DeviceStreamType current_stream = nullptr;
		167	+ AOTI_TORCH_ERROR_CODE_CHECK(
		168	+ aoti_torch_get_current_npu_stream(device_idx_, reinterpret_cast<void**>(&current_stream)));
		169	+ return current_stream;
		170	+ }
		171	+#endif
		172	+
152	void run(AtenTensorHandle* input_handles, // array of input AtenTensorHandle; handles	173	void run(AtenTensorHandle* input_handles, // array of input AtenTensorHandle; handles
153	// are stolen; the array itself is borrowed	174	// are stolen; the array itself is borrowed
154	AtenTensorHandle* output_handles, // array for writing output AtenTensorHandle; handles	175	AtenTensorHandle* output_handles, // array for writing output AtenTensorHandle; handles
	@@ -157,20 +178,22 @@ public:
157	DeviceStreamType stream, AOTIProxyExecutorHandle proxy_executor)	178	DeviceStreamType stream, AOTIProxyExecutorHandle proxy_executor)
158	{	179	{
159	#if defined(USE_NPU)	180	#if defined(USE_NPU)
		181	+ auto run_stream = normalize_run_stream(stream);
160	if (!run_finished_) {	182	if (!run_finished_) {
161	aclrtEvent run_finished;	183	aclrtEvent run_finished;
162	AOTI_RUNTIME_DEVICE_CHECK(aclrtCreateEvent(&run_finished));	184	AOTI_RUNTIME_DEVICE_CHECK(aclrtCreateEvent(&run_finished));
163	run_finished_.emplace(run_finished);	185	run_finished_.emplace(run_finished);
164	}	186	}
165	#else	187	#else
		188	+ auto run_stream = stream;
166	run_finished_ = false;	189	run_finished_ = false;
167	#endif	190	#endif
168		191
169	auto* model = static_cast<Model*>(this);	192	auto* model = static_cast<Model*>(this);
170	- model->run_impl(input_handles, output_handles, stream, proxy_executor);	193	+ model->run_impl(input_handles, output_handles, run_stream, proxy_executor);
171		194
172	#if defined(USE_NPU)	195	#if defined(USE_NPU)
173	- AOTI_RUNTIME_DEVICE_CHECK(aclrtRecordEvent(*run_finished_, stream));	196	+ AOTI_RUNTIME_DEVICE_CHECK(aclrtRecordEvent(*run_finished_, run_stream));
174	#else	197	#else
175	run_finished_ = true;	198	run_finished_ = true;
176	#endif	199	#endif
	@@ -185,30 +208,37 @@ public:
185	// borrowed	208	// borrowed
186	DeviceStreamType stream, AOTIProxyExecutorHandle proxy_executor)	209	DeviceStreamType stream, AOTIProxyExecutorHandle proxy_executor)
187	{	210	{
		211	+#if defined(USE_NPU)
		212	+ auto run_stream = normalize_run_stream(stream);
		213	+#else
		214	+ auto run_stream = stream;
		215	+#endif
188	// don't bother with any of the run_finished stuff; this is unsafe to call	216	// don't bother with any of the run_finished stuff; this is unsafe to call
189	// in a threaded context	217	// in a threaded context
190	auto* model = static_cast<Model*>(this);	218	auto* model = static_cast<Model*>(this);
191	- model->run_impl(input_handles, output_handles, stream, proxy_executor);	219	+ model->run_impl(input_handles, output_handles, run_stream, proxy_executor);
192	}	220	}
193		221
194	std::unordered_map<std::string, AtenTensorHandle>	222	std::unordered_map<std::string, AtenTensorHandle>
195	run_const_fold(DeviceStreamType stream, AOTIProxyExecutorHandle proxy_executor, bool initialization = false)	223	run_const_fold(DeviceStreamType stream, AOTIProxyExecutorHandle proxy_executor, bool initialization = false)
196	{	224	{
197	#if defined(USE_NPU)	225	#if defined(USE_NPU)
		226	+ auto run_stream = normalize_run_stream(stream);
198	if (!run_finished_) {	227	if (!run_finished_) {
199	aclrtEvent run_finished;	228	aclrtEvent run_finished;
200	AOTI_RUNTIME_DEVICE_CHECK(aclrtCreateEvent(&run_finished));	229	AOTI_RUNTIME_DEVICE_CHECK(aclrtCreateEvent(&run_finished));
201	run_finished_.emplace(run_finished);	230	run_finished_.emplace(run_finished);
202	}	231	}
203	#else	232	#else
		233	+ auto run_stream = stream;
204	run_finished_ = false;	234	run_finished_ = false;
205	#endif	235	#endif
206		236
207	auto* model = static_cast<Model*>(this);	237	auto* model = static_cast<Model*>(this);
208	- auto folded_constants = model->const_run_impl(stream, proxy_executor, initialization);	238	+ auto folded_constants = model->const_run_impl(run_stream, proxy_executor, initialization);
209		239
210	#if defined(USE_NPU)	240	#if defined(USE_NPU)
211	- AOTI_RUNTIME_DEVICE_CHECK(aclrtRecordEvent(*run_finished_, stream));	241	+ AOTI_RUNTIME_DEVICE_CHECK(aclrtRecordEvent(*run_finished_, run_stream));
212	#else	242	#else
213	run_finished_ = true;	243	run_finished_ = true;
214	#endif	244	#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+31-0

	@@ -264,6 +264,37 @@ 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	+
		291	+AOTI_TORCH_EXPORT AOTITorchError aoti_torch_npu_caching_allocator_raw_alloc(
		292	+ uint64_t nbytes,
		293	+ void** ret_ptr // returns raw NPU memory pointer
		294	+);
		295	+
		296	+AOTI_TORCH_EXPORT AOTITorchError aoti_torch_npu_caching_allocator_raw_delete(void* ptr);
		297	+
267	AOTI_TORCH_EXPORT AOTITorchError aoti_torch__embedding_bag(AtenTensorHandle weight, AtenTensorHandle indices,	298	AOTI_TORCH_EXPORT AOTITorchError aoti_torch__embedding_bag(AtenTensorHandle weight, AtenTensorHandle indices,
268	AtenTensorHandle offsets, int32_t scale_grad_by_freq,	299	AtenTensorHandle offsets, int32_t scale_grad_by_freq,
269	int32_t mode, int32_t sparse,	300	int32_t mode, int32_t sparse,

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

	@@ -3,6 +3,9 @@
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/NPUCachingAllocator.h>
		7	+#include <torch_npu/csrc/core/npu/NPUGuard.h>
		8	+#include <torch_npu/csrc/core/npu/NPUStream.h>
6	#include <torch_npu/csrc/inductor/aoti_torch/c/shim.h>	9	#include <torch_npu/csrc/inductor/aoti_torch/c/shim.h>
7	#include <torch_npu/csrc/inductor/aoti_torch/utils.h>	10	#include <torch_npu/csrc/inductor/aoti_torch/utils.h>
8	#include <torch_npu/csrc/inductor/inductor_ops.h>	11	#include <torch_npu/csrc/inductor/inductor_ops.h>
	@@ -18,6 +21,10 @@ int32_t aoti_torch_device_type_npu() { return (int32_t)c10::DeviceType::PrivateU
18	} // extern "C"	21	} // extern "C"
19	#endif	22	#endif
20		23
		24	+namespace c10_npu {
		25	+NPUStream getNPUStreamFromManagedAclrtStream(aclrtStream stream, c10::DeviceIndex device_index);
		26	+}
		27	+
21	namespace {	28	namespace {
22	static c10::Device c10_device(int32_t device_type, int32_t device_index)	29	static c10::Device c10_device(int32_t device_type, int32_t device_index)
23	{	30	{
	@@ -35,6 +42,8 @@ AOTITorchError aoti_torch_create_tensor_from_blob_npu(void* data, int64_t ndim,
35	AtenTensorHandle* ret_new_tensor)	42	AtenTensorHandle* ret_new_tensor)
36	{	43	{
37	AOTI_TORCH_CONVERT_EXCEPTION_TO_ERROR_CODE({	44	AOTI_TORCH_CONVERT_EXCEPTION_TO_ERROR_CODE({
		45	+ TORCH_CHECK(ret_new_tensor != nullptr, "ret_new_tensor is nullptr");
		46	+ *ret_new_tensor = nullptr;
38	c10::IntArrayRef sizes(sizes_ptr, ndim);	47	c10::IntArrayRef sizes(sizes_ptr, ndim);
39	c10::IntArrayRef strides(strides_ptr, ndim);	48	c10::IntArrayRef strides(strides_ptr, ndim);
40	c10::Device device = c10_device(device_type, device_index);	49	c10::Device device = c10_device(device_type, device_index);
	@@ -53,11 +62,91 @@ 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)	62	const uint8_t* opaque_metadata, int64_t opaque_metadata_size)
54	{	63	{
55	AOTI_TORCH_CONVERT_EXCEPTION_TO_ERROR_CODE({	64	AOTI_TORCH_CONVERT_EXCEPTION_TO_ERROR_CODE({
		65	+ TORCH_CHECK(ret_new_tensor != nullptr, "ret_new_tensor is nullptr");
		66	+ *ret_new_tensor = nullptr;
56	if (layout == static_cast<int32_t>(at::kMkldnn)) {	67	if (layout == static_cast<int32_t>(at::kMkldnn)) {
57	- throw std::runtime_error("do not support mkldnn on npu.");	68	+ TORCH_CHECK(false, "do not support mkldnn on npu.");
58	} else {	69	} else {
59	- aoti_torch_create_tensor_from_blob_npu(data, ndim, sizes_ptr, strides_ptr, storage_offset, dtype,	70	+ 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);	71	+ dtype, device_type, device_index, ret_new_tensor);
		72	+ if (err != AOTI_TORCH_SUCCESS) {
		73	+ return err;
		74	+ }
61	}	75	}
62	});	76	});
63	-}	77	+}
		78	+
		79	+AOTITorchError aoti_torch_create_npu_guard(int32_t device_index, NPUGuardHandle* ret_guard)
		80	+{
		81	+ AOTI_TORCH_CONVERT_EXCEPTION_TO_ERROR_CODE({
		82	+ TORCH_CHECK(ret_guard != nullptr, "ret_guard is nullptr");
		83	+ *ret_guard = nullptr;
		84	+ c10_npu::NPUGuard* guard = new c10_npu::NPUGuard(static_cast<c10::DeviceIndex>(device_index));
		85	+ *ret_guard = reinterpret_cast<NPUGuardHandle>(guard);
		86	+ });
		87	+}
		88	+
		89	+AOTITorchError aoti_torch_delete_npu_guard(NPUGuardHandle guard)
		90	+{
		91	+ AOTI_TORCH_CONVERT_EXCEPTION_TO_ERROR_CODE({ delete reinterpret_cast<c10_npu::NPUGuard*>(guard); });
		92	+}
		93	+
		94	+AOTITorchError aoti_torch_npu_guard_set_index(NPUGuardHandle guard, int32_t device_index)
		95	+{
		96	+ AOTI_TORCH_CONVERT_EXCEPTION_TO_ERROR_CODE({
		97	+ reinterpret_cast<c10_npu::NPUGuard*>(guard)->set_index(static_cast<c10::DeviceIndex>(device_index));
		98	+ });
		99	+}
		100	+
		101	+AOTITorchError aoti_torch_create_npu_stream_guard(void* stream, int32_t device_index, NPUStreamGuardHandle* ret_guard)
		102	+{
		103	+ AOTI_TORCH_CONVERT_EXCEPTION_TO_ERROR_CODE({
		104	+ TORCH_CHECK(ret_guard != nullptr, "ret_guard is nullptr");
		105	+ *ret_guard = nullptr;
		106	+
		107	+ auto raw_stream = static_cast<aclrtStream>(stream);
		108	+ auto managed_stream = c10_npu::getNPUStreamFromManagedAclrtStream(
		109	+ raw_stream, static_cast<c10::DeviceIndex>(device_index));
		110	+ auto* guard = new c10_npu::NPUStreamGuard(static_cast<c10::Stream>(managed_stream));
		111	+ *ret_guard = reinterpret_cast<NPUStreamGuardHandle>(guard);
		112	+ });
		113	+}
		114	+
		115	+AOTITorchError aoti_torch_delete_npu_stream_guard(NPUStreamGuardHandle guard)
		116	+{
		117	+ AOTI_TORCH_CONVERT_EXCEPTION_TO_ERROR_CODE({ delete reinterpret_cast<c10_npu::NPUStreamGuard*>(guard); });
		118	+}
		119	+
		120	+AOTITorchError aoti_torch_get_current_npu_stream(int32_t device_index, void** ret_stream)
		121	+{
		122	+ AOTI_TORCH_CONVERT_EXCEPTION_TO_ERROR_CODE({
		123	+ TORCH_CHECK(ret_stream != nullptr, "ret_stream is nullptr");
		124	+ ret_stream = reinterpret_cast<void>(
		125	+ c10_npu::getCurrentNPUStream(static_cast<c10::DeviceIndex>(device_index)).stream());
		126	+ });
		127	+}
		128	+
		129	+AOTITorchError aoti_torch_npu_caching_allocator_raw_alloc(uint64_t nbytes, void** ret_ptr)
		130	+{
		131	+ AOTI_TORCH_CONVERT_EXCEPTION_TO_ERROR_CODE({
		132	+ TORCH_CHECK(ret_ptr != nullptr, "ret_ptr is nullptr");
		133	+ *ret_ptr = nullptr;
		134	+ if (nbytes == 0) {
		135	+ return AOTI_TORCH_SUCCESS;
		136	+ }
		137	+
		138	+ *ret_ptr = c10_npu::NPUCachingAllocator::raw_alloc(nbytes);
		139	+
		140	+ TORCH_CHECK(
		141	+ *ret_ptr != nullptr, "Failed to allocate ", nbytes, " bytes from NPU caching allocator");
		142	+ });
		143	+}
		144	+
		145	+AOTITorchError aoti_torch_npu_caching_allocator_raw_delete(void* ptr)
		146	+{
		147	+ AOTI_TORCH_CONVERT_EXCEPTION_TO_ERROR_CODE({
		148	+ if (ptr != nullptr) {
		149	+ c10_npu::NPUCachingAllocator::raw_delete(ptr);
		150	+ }
		151	+ });
		152	+}