fix: fix from blob bug-pytorch-AtomGit

Mtest/cpp_extensions/extension.cpp+0-62

	@@ -2,7 +2,6 @@
2	#include <chrono>	2	#include <chrono>
3	#include <torch/extension.h>	3	#include <torch/extension.h>
4	#include "torch_npu/csrc/core/npu/NPUFormat.h"	4	#include "torch_npu/csrc/core/npu/NPUFormat.h"
5	-#include "torch_npu/csrc/aten/common/from_blob.h"
6	#include "torch_npu/csrc/framework/OpCommand.h"	5	#include "torch_npu/csrc/framework/OpCommand.h"
7	#include "torch_npu/csrc/framework/OpHook.h"	6	#include "torch_npu/csrc/framework/OpHook.h"
8	// test in .setup with relative path	7	// test in .setup with relative path
	@@ -33,64 +32,6 @@ bool check_storage_sizes(const Tensor &tensor, const c10::IntArrayRef &sizes)
33	return false;	32	return false;
34	}	33	}
35		34
36	-bool check_from_blob()
37	-{
38	- auto data = torch::tensor({1.0, 2.0, 3.0}, torch::kFloat).to(at::Device("npu:0"));
39	- auto tensor = at_npu::native::from_blob(data.data_ptr(), data.sizes(), torch::dtype(torch::kFloat));
40	-
41	- bool dtype_same = (tensor.dtype() == torch::kFloat);
42	- bool num_same = (tensor.numel() == 3);
43	- bool pos1_same = (tensor[0].item<float>() == 1);
44	- bool pos2_same = (tensor[1].item<float>() == 2);
45	- bool pos3_same = (tensor[2].item<float>() == 3);
46	- tensor = tensor -1;
47	- bool sub_same = ((tensor[2].item<float>() == 2));
48	- return dtype_same && num_same && pos1_same && pos2_same && pos3_same && sub_same;
49	-}
50	-
51	-bool check_from_blob_delete()
52	-{
53	- int isgone = 0;
54	- {
55	- auto data = torch::tensor({1.0, 2.0, 3.0}, torch::kFloat).to(at::Device("npu:0"));
56	- auto res = at_npu::native::from_blob(data.data_ptr(), data.sizes(), [&](void*) { isgone++; });
57	- }
58	- bool is_deleted = (isgone == 1);
59	- return is_deleted;
60	-}
61	-
62	-bool check_from_blob_strides()
63	-{
64	- auto data = torch::tensor({1, 2, 3, 4, 5, 6, 7, 8, 9}, torch::kInt32).to(at::Device("npu:0"));
65	- auto tensor = at_npu::native::from_blob(data.data_ptr(), {3, 3}, {1, 3}, torch::kInt32); // sizes = {3,3}, strides = {1,3}
66	-
67	- bool dtype_same = (tensor.dtype() == torch::kInt32);
68	- bool num_same = (tensor.numel() == data.numel());
69	- const std::vector<int64_t> expected_strides = {1, 3};
70	- auto result_strides = tensor.strides();
71	- bool stride_same = std::equal(result_strides.begin(), result_strides.end(), expected_strides.begin());
72	- bool pos_same = true;
73	- for (const auto i : c10::irange(tensor.size(0))) {
74	- for (const auto j : c10::irange(tensor.size(1))) {
75	- // NOTE: This is column major because the strides are swapped.
76	- if (tensor[i][j].item<int32_t>() != (1 + (j * tensor.size(1)) + i))
77	- pos_same = false;
78	- }
79	- }
80	- auto tensor_clone = tensor.clone();
81	- bool clone_same = at::equal(tensor_clone, tensor);
82	- auto tensor_add = tensor + 1;
83	- bool add_same = true;
84	- for (const auto i : c10::irange(tensor_add.size(0))) {
85	- for (const auto j : c10::irange(tensor_add.size(1))) {
86	- // NOTE: This is column major because the strides are swapped.
87	- if (tensor_add[i][j].item<int32_t>() != (2 + (j * tensor_add.size(1)) + i))
88	- add_same = false;
89	- }
90	- }
91	- return dtype_same && num_same && pos_same && stride_same && clone_same && add_same;
92	-}
93	-
94	Tensor blocking_ops(Tensor x)	35	Tensor blocking_ops(Tensor x)
95	{	36	{
96	auto blocking_call = []() -> int {	37	auto blocking_call = []() -> int {
	@@ -140,9 +81,6 @@ PYBIND11_MODULE(TORCH_EXTENSION_NAME, m)
140	m.def("tanh_add", &tanh_add, "tanh(x) + tanh(y)");	81	m.def("tanh_add", &tanh_add, "tanh(x) + tanh(y)");
141	m.def("npu_add", &npu_add, "x + y");	82	m.def("npu_add", &npu_add, "x + y");
142	m.def("check_storage_sizes", &check_storage_sizes, "check_storage_sizes");	83	m.def("check_storage_sizes", &check_storage_sizes, "check_storage_sizes");
143	- m.def("check_from_blob", &check_from_blob, "check_from_blob");
144	- m.def("check_from_blob_strides", &check_from_blob_strides, "check_from_blob_strides");
145	- m.def("check_from_blob_delete", &check_from_blob_delete, "check_from_blob_delete");
146	m.def("blocking_ops", &blocking_ops, "blocking_ops");	84	m.def("blocking_ops", &blocking_ops, "blocking_ops");
147	m.def("register_op_hook", &register_op_hook, "register_op_hook");	85	m.def("register_op_hook", &register_op_hook, "register_op_hook");
148	m.def("get_op_hook_call_count", &get_op_hook_call_count, "get_op_hook_call_count");	86	m.def("get_op_hook_call_count", &get_op_hook_call_count, "get_op_hook_call_count");

Mtest/cpp_extensions/setup.py+3-0

	@@ -17,6 +17,9 @@ ext_modules = [
17	NpuExtension(	17	NpuExtension(
18	'torch_test_cpp_extension.npu', ['extension.cpp'],	18	'torch_test_cpp_extension.npu', ['extension.cpp'],
19	extra_compile_args=CXX_FLAGS),	19	extra_compile_args=CXX_FLAGS),
		20	+ NpuExtension(
		21	+ 'torch_test_cpp_extension.npu_from_blob', ['test_from_blob.cpp'],
		22	+ extra_compile_args=CXX_FLAGS),
20	]	23	]
21		24
22	setup(	25	setup(

Mtest/cpp_extensions/test/test_cpp_extensions_aot.py+31-18

	@@ -15,12 +15,12 @@ from torch_npu.testing.common_utils import create_common_tensor
15		15
16	try:	16	try:
17	import torch_test_cpp_extension.npu as npu_extension	17	import torch_test_cpp_extension.npu as npu_extension
		18	+ import torch_test_cpp_extension.npu_from_blob as from_blob_ext
18	except ImportError as e:	19	except ImportError as e:
19	raise RuntimeError(	20	raise RuntimeError(
20	"test_cpp_extensions_aot.py cannot be invoked directly. Run "	21	"test_cpp_extensions_aot.py cannot be invoked directly. Run "
21	"`python run_cpp_test.py` instead.") from e	22	"`python run_cpp_test.py` instead.") from e
22		23
23	-
24	class TestCppExtensionAOT(TestCase):	24	class TestCppExtensionAOT(TestCase):
25	"""Tests ahead-of-time cpp extensions	25	"""Tests ahead-of-time cpp extensions
26	"""	26	"""
	@@ -39,22 +39,6 @@ class TestCppExtensionAOT(TestCase):
39	npu_z = npu_extension.npu_add(x.npu(), y.npu())	39	npu_z = npu_extension.npu_add(x.npu(), y.npu())
40	self.assertEqual(npu_z.cpu(), (x + y))	40	self.assertEqual(npu_z.cpu(), (x + y))
41		41
42	- def test_storage_sizes(self):
43	- t = torch_npu.npu_format_cast(torch.ones(128, 512, dtype=torch.int8).npu(), 29)
44	- self.assertTrue(npu_extension.check_storage_sizes(t, (16, 8, 16, 32)))
45	- t = torch_npu.npu_format_cast(torch.ones(31, 127, 511, dtype=torch.int8).npu(), 29)
46	- self.assertTrue(npu_extension.check_storage_sizes(t, (31, 16, 8, 16, 32)))
47	- t = torch_npu.npu_format_cast(torch.ones(128, 512, dtype=torch.float16).npu(), 29)
48	- self.assertTrue(npu_extension.check_storage_sizes(t, (32, 8, 16, 16)))
49	- # float32 will cast to float16 before calculate
50	- t = torch_npu.npu_format_cast(torch.ones(128, 512, dtype=torch.float32).npu(), 29)
51	- self.assertTrue(npu_extension.check_storage_sizes(t, (32, 8, 16, 16)))
52	-
53	- def test_from_blob(self):
54	- self.assertTrue(npu_extension.check_from_blob())
55	- self.assertTrue(npu_extension.check_from_blob_strides())
56	- self.assertTrue(npu_extension.check_from_blob_delete())
57	-
58	def test_dispatch_allreduce(self):	42	def test_dispatch_allreduce(self):
59	flags = os.O_WRONLY \| os.O_RDONLY \| os.O_CREAT	43	flags = os.O_WRONLY \| os.O_RDONLY \| os.O_CREAT
60	modes = stat.S_IWUSR \| stat.S_IRUSR	44	modes = stat.S_IWUSR \| stat.S_IRUSR
	@@ -65,7 +49,7 @@ class TestCppExtensionAOT(TestCase):
65	cmd = ["torchrun", "--nproc-per-node=1", code_file]	49	cmd = ["torchrun", "--nproc-per-node=1", code_file]
66	p = subprocess.Popen(cmd, stderr=subprocess.STDOUT, stdout=f)	50	p = subprocess.Popen(cmd, stderr=subprocess.STDOUT, stdout=f)
67	p.wait()	51	p.wait()
68	-	52	+
69	timeout = 0	53	timeout = 0
70	with open(log_pth, 'r', encoding='utf-8') as f:	54	with open(log_pth, 'r', encoding='utf-8') as f:
71	tmp = f.readlines()	55	tmp = f.readlines()
	@@ -207,6 +191,35 @@ class TestCppExtensionAOT(TestCase):
207	os.remove(dump_pth)	191	os.remove(dump_pth)
208	os.remove(dump_pth + "_py_traceback")	192	os.remove(dump_pth + "_py_traceback")
209		193
		194	+class TestFromBlob(TestCase):
		195	+ """Tests for at_npu::native::from_blob interface"""
		196	+
		197	+ def test_from_blob_basic(self):
		198	+ self.assertTrue(from_blob_ext.test_from_blob_basic())
		199	+
		200	+ def test_from_blob_deleter(self):
		201	+ self.assertTrue(from_blob_ext.test_from_blob_deleter())
		202	+
		203	+ def test_from_blob_strides(self):
		204	+ self.assertTrue(from_blob_ext.test_from_blob_strides())
		205	+
		206	+ def test_from_blob_storage_offset(self):
		207	+ self.assertTrue(from_blob_ext.test_from_blob_storage_offset())
		208	+
		209	+ def test_from_blob_storage_offset_2d(self):
		210	+ self.assertTrue(from_blob_ext.test_from_blob_storage_offset_2d())
		211	+
		212	+ def test_from_blob_storage_offset_dtype(self):
		213	+ self.assertTrue(from_blob_ext.test_from_blob_storage_offset_dtype())
		214	+
		215	+ def test_from_blob_storage_offset_contiguous(self):
		216	+ self.assertTrue(from_blob_ext.test_from_blob_storage_offset_contiguous())
		217	+
		218	+ def test_from_blob_non_owning(self):
		219	+ self.assertTrue(from_blob_ext.test_from_blob_non_owning())
		220	+
		221	+ def test_from_blob_clone(self):
		222	+ self.assertTrue(from_blob_ext.test_from_blob_clone())
210		223
211	if __name__ == "__main__":	224	if __name__ == "__main__":
212	run_tests()	225	run_tests()

Atest/cpp_extensions/test_from_blob.cpp+255-0

@@ -0,0 +1,255 @@
	1	+#include <torch/extension.h>
	2	+#include <c10/util/irange.h>
	3	+#include "torch_npu/csrc/aten/common/from_blob.h"
	4	+
	5	+using namespace at;
	6	+
	7	+static const auto npu_device = at::Device("npu:0");
	8	+
	9	+// Helper: create a 1D float tensor on NPU
	10	+static Tensor make_npu_float(std::vector<float> vals)
	11	+{
	12	+ return torch::tensor(vals, torch::kFloat).to(npu_device);
	13	+}
	14	+
	15	+// Helper: create a 1D int32 tensor on NPU
	16	+static Tensor make_npu_int(std::vector<int32_t> vals)
	17	+{
	18	+ return torch::tensor(vals, torch::kInt32).to(npu_device);
	19	+}
	20	+
	21	+// Test 1: Basic from_blob with sizes and dtype.
	22	+// Verifies: dtype, element count, per-element values match source,
	23	+// and storage context is null (non-owning, no deleter).
	24	+bool test_from_blob_basic()
	25	+{
	26	+ auto data = make_npu_float({1.0, 2.0, 3.0});
	27	+ auto tensor = at_npu::native::from_blob(
	28	+ data.data_ptr(), data.sizes(), torch::dtype(torch::kFloat));
	29	+ if (tensor.dtype() != torch::kFloat) return false;
	30	+ if (tensor.numel() != data.numel()) return false;
	31	+ for (const auto i : c10::irange(data.numel())) {
	32	+ if (tensor[i].item<float>() != data[i].item<float>()) return false;
	33	+ }
	34	+ if (tensor.storage().data_ptr().get_context() != nullptr) return false;
	35	+ return true;
	36	+}
	37	+
	38	+// Test 2: Deleter is called exactly once when tensor goes out of scope.
	39	+bool test_from_blob_deleter()
	40	+{
	41	+ int called = 0;
	42	+ {
	43	+ auto data = make_npu_float({1.0, 2.0, 3.0});
	44	+ auto tensor = at_npu::native::from_blob(
	45	+ data.data_ptr(), data.sizes(), [&](void*) { called++; });
	46	+ }
	47	+ return (called == 1);
	48	+}
	49	+
	50	+// Test 3: from_blob with strides (column-major layout).
	51	+// sizes={3,3}, strides={1,3}: element [i][j] = 1 + j*cols + i (column major).
	52	+bool test_from_blob_strides()
	53	+{
	54	+ auto data = make_npu_int({1, 2, 3, 4, 5, 6, 7, 8, 9});
	55	+ auto tensor = at_npu::native::from_blob(
	56	+ data.data_ptr(), {3, 3}, {1, 3}, torch::kInt32);
	57	+ if (tensor.dtype() != torch::kInt32) return false;
	58	+ if (tensor.numel() != data.numel()) return false;
	59	+
	60	+ const std::vector<int64_t> expected_strides = {1, 3};
	61	+ auto result_strides = tensor.strides();
	62	+ if (!std::equal(result_strides.begin(), result_strides.end(),
	63	+ expected_strides.begin()))
	64	+ return false;
	65	+
	66	+ for (const auto i : c10::irange(tensor.size(0))) {
	67	+ for (const auto j : c10::irange(tensor.size(1))) {
	68	+ if (tensor[i][j].item<int32_t>() != (1 + (j * tensor.size(1)) + i))
	69	+ return false;
	70	+ }
	71	+ }
	72	+ return true;
	73	+}
	74	+
	75	+// Test 4: from_blob with storage_offset, verifies storage size is correct.
	76	+// This is the bug-fix test: storage_offset must be multiplied by itemsize.
	77	+// Base tensor: arange(kBaseSize) as float32 (4 bytes each).
	78	+// View: sizes=[kViewSize], strides=[kStride], offset=kOffset (simulates base[3:8]).
	79	+// Expected storage bytes = (kOffset + kStride + kStride(kViewSize-1)) sizeof(float)
	80	+// = (3 + 1 + 14) 4 = 32 bytes.
	81	+// Before the fix: storage was computed in mixed units (bytes + elements), giving wrong result.
	82	+// Verified: element i of the view equals float(i + kOffset).
	83	+bool test_from_blob_storage_offset()
	84	+{
	85	+ constexpr int64_t kBaseSize = 10;
	86	+ constexpr int64_t kViewSize = 5;
	87	+ constexpr int64_t kStride = 1;
	88	+ constexpr int64_t kOffset = 3;
	89	+ auto base = torch::arange(kBaseSize, torch::kFloat).to(npu_device);
	90	+ auto tensor = at_npu::native::from_blob(
	91	+ base.storage().mutable_data(),
	92	+ {kViewSize},
	93	+ {kStride},
	94	+ kOffset,
	95	+ torch::dtype(torch::kFloat));
	96	+ if (tensor.storage_offset() != kOffset) return false;
	97	+ for (int i = 0; i < kViewSize; i++) {
	98	+ if (tensor[i].item<float>() != static_cast<float>(i + kOffset))
	99	+ return false;
	100	+ }
	101	+ size_t expected_nbytes = (kOffset + kStride + kStride * (kViewSize - 1)) * sizeof(float);
	102	+ if (tensor.storage().nbytes() != expected_nbytes) return false;
	103	+ return true;
	104	+}
	105	+
	106	+// Test 5: from_blob with storage_offset for 2D tensor.
	107	+// Base: arange(kBaseSize) float32. View: kRows x kCols submatrix of a 4x4 matrix.
	108	+// strides=[kRowStride, kColStride], offset=kOffset (row 1, col 1 of 4x4 -> offset=5).
	109	+// Expected values: [kOffset, kOffset+kColStride], [kOffset+kRowStride, kOffset+kRowStride+kColStride].
	110	+// Expected storage bytes = (kOffset + kColStride + kRowStride(kRows-1) + kColStride(kCols-1)) * sizeof(float)
	111	+// = (5 + 1 + 41 + 11) * 4 = 44 bytes.
	112	+bool test_from_blob_storage_offset_2d()
	113	+{
	114	+ constexpr int64_t kBaseSize = 16;
	115	+ constexpr int64_t kRows = 2;
	116	+ constexpr int64_t kCols = 2;
	117	+ constexpr int64_t kRowStride = 4;
	118	+ constexpr int64_t kColStride = 1;
	119	+ constexpr int64_t kOffset = 5;
	120	+ auto base = torch::arange(kBaseSize, torch::kFloat).to(npu_device);
	121	+ auto tensor = at_npu::native::from_blob(
	122	+ base.storage().mutable_data(),
	123	+ {kRows, kCols},
	124	+ {kRowStride, kColStride},
	125	+ kOffset,
	126	+ torch::dtype(torch::kFloat));
	127	+ if (tensor.storage_offset() != kOffset) return false;
	128	+ if (tensor.sizes() != c10::IntArrayRef({kRows, kCols})) return false;
	129	+ if (tensor.strides() != c10::IntArrayRef({kRowStride, kColStride})) return false;
	130	+ if (tensor[0][0].item<float>() != static_cast<float>(kOffset)) return false;
	131	+ if (tensor[0][1].item<float>() != static_cast<float>(kOffset + kColStride)) return false;
	132	+ if (tensor[1][0].item<float>() != static_cast<float>(kOffset + kRowStride)) return false;
	133	+ if (tensor[1][1].item<float>() != static_cast<float>(kOffset + kRowStride + kColStride)) return false;
	134	+ size_t expected_nbytes =
	135	+ (kOffset + kColStride + kRowStride * (kRows - 1) + kColStride * (kCols - 1)) * sizeof(float);
	136	+ if (tensor.storage().nbytes() != expected_nbytes) return false;
	137	+ return true;
	138	+}
	139	+
	140	+// Test 6: from_blob with storage_offset for different dtypes.
	141	+// Verifies the itemsize multiplier is correct across float64, float16, and int32.
	142	+// All sub-tests: sizes=[kViewSize], strides=[kStride], offset=kDtypeOffset.
	143	+// float64 (8 bytes): expected = (kDtypeOffset+kStride+kStride(kViewSize-1))sizeof(double) = 56.
	144	+// float16 (2 bytes): expected = (kDtypeOffset+kStride+kStride(kViewSize-1))sizeof(at::Half) = 14.
	145	+// int32 (4 bytes): expected = (kDtypeOffset+kStride+kStride(kViewSize-1))sizeof(int32_t) = 28.
	146	+// The first element of each view should equal kDtypeOffset.
	147	+bool test_from_blob_storage_offset_dtype()
	148	+{
	149	+ bool all_pass = true;
	150	+ constexpr int64_t kBaseSize = 10;
	151	+ constexpr int64_t kViewSize = 5;
	152	+ constexpr int64_t kStride = 1;
	153	+ constexpr int64_t kDtypeOffset = 2;
	154	+
	155	+ {
	156	+ auto base = torch::arange(kBaseSize, torch::kFloat64).to(npu_device);
	157	+ auto tensor = at_npu::native::from_blob(
	158	+ base.storage().mutable_data(),
	159	+ {kViewSize}, {kStride}, kDtypeOffset,
	160	+ torch::dtype(torch::kFloat64));
	161	+ size_t expected = (kDtypeOffset + kStride + kStride * (kViewSize - 1)) * sizeof(double);
	162	+ if (tensor.storage().nbytes() != expected) all_pass = false;
	163	+ if (tensor[0].item<double>() != static_cast<double>(kDtypeOffset)) all_pass = false;
	164	+ }
	165	+
	166	+ {
	167	+ auto base = torch::arange(kBaseSize, torch::kFloat16).to(npu_device);
	168	+ auto tensor = at_npu::native::from_blob(
	169	+ base.storage().mutable_data(),
	170	+ {kViewSize}, {kStride}, kDtypeOffset,
	171	+ torch::dtype(torch::kFloat16));
	172	+ size_t expected = (kDtypeOffset + kStride + kStride * (kViewSize - 1)) * sizeof(at::Half);
	173	+ if (tensor.storage().nbytes() != expected) all_pass = false;
	174	+ }
	175	+
	176	+ {
	177	+ auto base = torch::arange(kBaseSize, torch::kInt32).to(npu_device);
	178	+ auto tensor = at_npu::native::from_blob(
	179	+ base.storage().mutable_data(),
	180	+ {kViewSize}, {kStride}, kDtypeOffset,
	181	+ torch::dtype(torch::kInt32));
	182	+ size_t expected = (kDtypeOffset + kStride + kStride * (kViewSize - 1)) * sizeof(int32_t);
	183	+ if (tensor.storage().nbytes() != expected) all_pass = false;
	184	+ if (tensor[0].item<int32_t>() != static_cast<int32_t>(kDtypeOffset)) all_pass = false;
	185	+ }
	186	+
	187	+ return all_pass;
	188	+}
	189	+
	190	+// Test 7: from_blob without explicit strides (contiguous) with storage_offset.
	191	+// Verifies storage_offset, element count, first and last element values.
	192	+bool test_from_blob_storage_offset_contiguous()
	193	+{
	194	+ constexpr int64_t kBaseSize = 10;
	195	+ constexpr int64_t kViewSize = 5;
	196	+ constexpr int64_t kStride = 1;
	197	+ constexpr int64_t kOffset = 3;
	198	+ auto base = torch::arange(kBaseSize, torch::kFloat).to(npu_device);
	199	+ auto tensor = at_npu::native::from_blob(
	200	+ base.storage().mutable_data(),
	201	+ {kViewSize},
	202	+ {kStride},
	203	+ kOffset,
	204	+ torch::dtype(torch::kFloat));
	205	+ if (tensor.storage_offset() != kOffset) return false;
	206	+ if (tensor.numel() != kViewSize) return false;
	207	+ if (tensor[0].item<float>() != static_cast<float>(kOffset)) return false;
	208	+ if (tensor[kViewSize - 1].item<float>() != static_cast<float>(kOffset + kViewSize - 1)) return false;
	209	+ return true;
	210	+}
	211	+
	212	+// Test 8: from_blob produces a non-owning reference; original tensor stays valid after view is destroyed.
	213	+// data={10.0, 20.0, 30.0}. The view (weak) shares the same data pointer.
	214	+// After weak goes out of scope, data[0] and data[kLastDataIdx] must still be accessible.
	215	+bool test_from_blob_non_owning()
	216	+{
	217	+ constexpr int64_t kLastDataIdx = 2;
	218	+ auto data = make_npu_float({10.0, 20.0, 30.0});
	219	+ void* original_ptr = data.data_ptr();
	220	+ {
	221	+ auto weak = at_npu::native::from_blob(
	222	+ data.data_ptr(), data.sizes(), torch::dtype(torch::kFloat));
	223	+ if (weak.data_ptr() != original_ptr) return false;
	224	+ if (weak[1].item<float>() != 20.0f) return false;
	225	+ }
	226	+ if (data[0].item<float>() != 10.0f) return false;
	227	+ if (data[kLastDataIdx].item<float>() != 30.0f) return false;
	228	+ return true;
	229	+}
	230	+
	231	+// Test 9: from_blob clone produces an independent copy.
	232	+// Cloned tensor must be equal in value but reside in different storage.
	233	+bool test_from_blob_clone()
	234	+{
	235	+ auto data = make_npu_float({1.0, 2.0, 3.0});
	236	+ auto tensor = at_npu::native::from_blob(
	237	+ data.data_ptr(), data.sizes(), torch::dtype(torch::kFloat));
	238	+ auto cloned = tensor.clone();
	239	+ if (!at::equal(tensor, cloned)) return false;
	240	+ if (cloned.data_ptr() == tensor.data_ptr()) return false;
	241	+ return true;
	242	+}
	243	+
	244	+PYBIND11_MODULE(TORCH_EXTENSION_NAME, m)
	245	+{
	246	+ m.def("test_from_blob_basic", &test_from_blob_basic);
	247	+ m.def("test_from_blob_deleter", &test_from_blob_deleter);
	248	+ m.def("test_from_blob_strides", &test_from_blob_strides);
	249	+ m.def("test_from_blob_storage_offset", &test_from_blob_storage_offset);
	250	+ m.def("test_from_blob_storage_offset_2d", &test_from_blob_storage_offset_2d);
	251	+ m.def("test_from_blob_storage_offset_dtype", &test_from_blob_storage_offset_dtype);
	252	+ m.def("test_from_blob_storage_offset_contiguous", &test_from_blob_storage_offset_contiguous);
	253	+ m.def("test_from_blob_non_owning", &test_from_blob_non_owning);
	254	+ m.def("test_from_blob_clone", &test_from_blob_clone);
	255	+}

Mtest/npu/test_npu_format.py+22-0

	@@ -56,6 +56,28 @@ class TestNPUFormat(TestCase):
56	self.assertEqual(fmt1, torch_npu.Format.FRACTAL_NZ)	56	self.assertEqual(fmt1, torch_npu.Format.FRACTAL_NZ)
57	self.assertEqual(x1.data_ptr(), weak_x1.data_ptr())	57	self.assertEqual(x1.data_ptr(), weak_x1.data_ptr())
58		58
		59	+ def test_weak_ref_tensor_with_storage_offset(self):
		60	+ """test _weak_ref_tensor preserves shape, strides, offset and data"""
		61	+ view_shape = [2, 1, 8, 64]
		62	+ view_strides = [1536, 0, 192, 1]
		63	+ view_offset = 128
		64	+
		65	+ max_offset = view_offset
		66	+ for i in range(len(view_shape)):
		67	+ max_offset += (view_shape[i] - 1) * view_strides[i]
		68	+ storage_size = max_offset + 1
		69	+
		70	+ base = torch.arange(storage_size, dtype=torch.float32).npu()
		71	+ view = torch.as_strided(base, size=view_shape, stride=view_strides,
		72	+ storage_offset=view_offset)
		73	+
		74	+ weak = _weak_ref_tensor(view)
		75	+ self.assertEqual(weak.size(), view.size())
		76	+ self.assertEqual(weak.stride(), view.stride())
		77	+ self.assertEqual(weak.storage_offset(), view.storage_offset())
		78	+ self.assertEqual(weak.storage().nbytes(), view.storage().nbytes())
		79	+ self.assertTrue(torch.equal(weak, view))
		80	+
59		81
60	if __name__ == "__main__":	82	if __name__ == "__main__":
61	run_tests()	83	run_tests()

Mtorch_npu/csrc/aten/common/from_blob.cpp+2-2

	@@ -75,7 +75,7 @@ std::size_t TensorMaker::computeStorageSize() const noexcept
75	if (strides_) {	75	if (strides_) {
76	auto storage_size = at::detail::computeStorageNbytes(sizes_, *strides_, itemsize);	76	auto storage_size = at::detail::computeStorageNbytes(sizes_, *strides_, itemsize);
77	if (storage_offset_) {	77	if (storage_offset_) {
78	- storage_size += storage_offset_.value();	78	+ storage_size += storage_offset_.value() * itemsize;
79	}	79	}
80	return storage_size;	80	return storage_size;
81	}	81	}
	@@ -86,7 +86,7 @@ std::size_t TensorMaker::computeStorageSize() const noexcept
86	}	86	}
87	auto storage_size = size * itemsize;	87	auto storage_size = size * itemsize;
88	if (storage_offset_) {	88	if (storage_offset_) {
89	- storage_size += storage_offset_.value();	89	+ storage_size += storage_offset_.value() * itemsize;
90	}	90	}
91	return storage_size;	91	return storage_size;
92	}	92	}

Mtorch_npu/csrc/npu/Module.cpp+10-7

	@@ -555,7 +555,7 @@ void RegisterNpuPluggableAllocator(PyObject* module)
555	}	555	}
556	auto delta = c10_npu::NPUCachingAllocator::setCheckpointPoolState(device, std::move(pps));	556	auto delta = c10_npu::NPUCachingAllocator::setCheckpointPoolState(device, std::move(pps));
557	auto& freed_pointers = delta.ptrs_freed;	557	auto& freed_pointers = delta.ptrs_freed;
558	-	558	+
559	std::unordered_set<void*> allocd_set;	559	std::unordered_set<void*> allocd_set;
560	for (auto& data_ptr : delta.dataptrs_allocd) {	560	for (auto& data_ptr : delta.dataptrs_allocd) {
561	allocd_set.insert(data_ptr.get());	561	allocd_set.insert(data_ptr.get());
	@@ -576,7 +576,7 @@ void RegisterNpuPluggableAllocator(PyObject* module)
576	ptr_set.size() >= definite_freed_count,	576	ptr_set.size() >= definite_freed_count,
577	"Any stale tensors which are being manually freed"	577	"Any stale tensors which are being manually freed"
578	" must be passed to set checkpoint", PTA_ERROR(ErrCode::PARAM));	578	" must be passed to set checkpoint", PTA_ERROR(ErrCode::PARAM));
579	-	579	+
580	removeStorageDeleterFns(ptrs, freed_pointer_set);	580	removeStorageDeleterFns(ptrs, freed_pointer_set);
581	std::vector<c10::StorageImpl*> storages_to_add_deleters_to;	581	std::vector<c10::StorageImpl*> storages_to_add_deleters_to;
582	storages_to_add_deleters_to.reserve(storages_to_add_deleters_to_ptr.size());	582	storages_to_add_deleters_to.reserve(storages_to_add_deleters_to_ptr.size());
	@@ -584,7 +584,7 @@ void RegisterNpuPluggableAllocator(PyObject* module)
584	// NOLINTNEXTLINE(performance-no-int-to-ptr)	584	// NOLINTNEXTLINE(performance-no-int-to-ptr)
585	storages_to_add_deleters_to.push_back((c10::StorageImpl*)ptr_int);	585	storages_to_add_deleters_to.push_back((c10::StorageImpl*)ptr_int);
586	}	586	}
587	-	587	+
588	addStorageDeleterFns(storages_to_add_deleters_to, delta);	588	addStorageDeleterFns(storages_to_add_deleters_to, delta);
589	});	589	});
590	m.def(	590	m.def(
	@@ -677,11 +677,14 @@ void RegisterNpuPluggableAllocator(PyObject* module)
677	m.def(	677	m.def(
678	"_weak_ref_tensor",	678	"_weak_ref_tensor",
679	[](const at::Tensor& t) {	679	[](const at::Tensor& t) {
680	- void* data_ptr = t.data_ptr();	680	+ void* storage_data_ptr = t.storage().mutable_data();
681	- std::vector<int64_t> sizes = t.sizes().vec();	681	+ int64_t storage_numel = static_cast<int64_t>(t.storage().nbytes()) / t.element_size();
682	- std::vector<int64_t> strides = t.strides().vec();
683	auto options = t.options();	682	auto options = t.options();
684	- auto new_tensor = at_npu::native::from_blob(data_ptr, sizes, strides, options);	683	+
		684	+ auto new_tensor = at_npu::native::from_blob(storage_data_ptr, {storage_numel}, options);
		685	+ auto* impl = new_tensor.unsafeGetTensorImpl();
		686	+ impl->set_sizes_and_strides(t.sizes(), t.strides());
		687	+ impl->set_storage_offset(t.storage_offset());
685		688
686	auto dst_desc = torch_npu::NPUBridge::GetNpuStorageImpl(t)->npu_desc_;	689	auto dst_desc = torch_npu::NPUBridge::GetNpuStorageImpl(t)->npu_desc_;
687	torch_npu::NPUBridge::GetNpuStorageImpl(new_tensor)->npu_desc_ = dst_desc;	690	torch_npu::NPUBridge::GetNpuStorageImpl(new_tensor)->npu_desc_ = dst_desc;