torch_npu support dlpack-pytorch-AtomGit

Atest/npu/test_dlpack.py+120-0

@@ -0,0 +1,120 @@
	1	+import torch
	2	+from torch.utils.dlpack import to_dlpack, from_dlpack
	3	+import torch_npu
	4	+from torch_npu.testing.testcase import TestCase, run_tests
	5	+from torch_npu.testing.decorator import Dtypes, instantiate_tests
	6	+
	7	+
	8	+@instantiate_tests
	9	+class TestDLPack(TestCase):
	10	+
	11	+ @Dtypes(torch.float, torch.double, torch.int8, torch.int16, torch.int32, torch.int64, torch.uint8, torch.bool)
	12	+ def test_dlpack_roundtrip_basic(self, dtype, device="npu"):
	13	+ """Test basic dlpack roundtrip: torch_npu tensor -> dlpack -> torch_npu tensor"""
	14	+ # Create original tensor
	15	+ if dtype == torch.bool:
	16	+ original = torch.randint(0, 2, (2, 3, 4), dtype=dtype, device=device)
	17	+ elif dtype in [torch.int8, torch.int16, torch.int32, torch.int64, torch.uint8]:
	18	+ original = torch.randint(-10, 10, (2, 3, 4), dtype=dtype, device=device)
	19	+ else:
	20	+ original = torch.randn(2, 3, 4, dtype=dtype, device=device)
	21	+
	22	+ # Convert to dlpack
	23	+ dlpack_tensor = to_dlpack(original)
	24	+
	25	+ # Convert back to torch_npu tensor
	26	+ restored = from_dlpack(dlpack_tensor)
	27	+
	28	+ # Verify the roundtrip
	29	+ self.assertEqual(original, restored)
	30	+ self.assertEqual(original.dtype, restored.dtype)
	31	+ self.assertEqual(original.device, restored.device)
	32	+ self.assertEqual(original.shape, restored.shape)
	33	+ self.assertEqual(original.stride(), restored.stride())
	34	+
	35	+ @Dtypes(torch.half, torch.float, torch.double)
	36	+ def test_dlpack_roundtrip_different_shapes(self, dtype, device="npu"):
	37	+ """Test dlpack roundtrip with different tensor shapes"""
	38	+ shapes = [
	39	+ (1,), # 1D tensor
	40	+ (5, 5), # 2D tensor
	41	+ (2, 3, 4), # 3D tensor
	42	+ (2, 2, 2, 2), # 4D tensor
	43	+ (1, 1, 1, 1, 1) # 5D tensor
	44	+ ]
	45	+
	46	+ for shape in shapes:
	47	+ with self.subTest(shape=shape):
	48	+ original = torch.randn(shape, dtype=dtype, device=device)
	49	+ dlpack_tensor = to_dlpack(original)
	50	+ restored = from_dlpack(dlpack_tensor)
	51	+
	52	+ self.assertEqual(original, restored)
	53	+ self.assertEqual(original.shape, restored.shape)
	54	+
	55	+ @Dtypes(torch.float)
	56	+ def test_dlpack_roundtrip_contiguous(self, dtype, device="npu"):
	57	+ """Test dlpack roundtrip with contiguous and non-contiguous tensors"""
	58	+ # Test contiguous tensor
	59	+ original_contiguous = torch.randn(4, 4, dtype=dtype, device=device)
	60	+ self.assertTrue(original_contiguous.is_contiguous())
	61	+
	62	+ dlpack_tensor = to_dlpack(original_contiguous)
	63	+ restored = from_dlpack(dlpack_tensor)
	64	+
	65	+ self.assertEqual(original_contiguous, restored)
	66	+ self.assertTrue(restored.is_contiguous())
	67	+
	68	+ # Test non-contiguous tensor (transpose)
	69	+ original_non_contiguous = original_contiguous.t()
	70	+ self.assertFalse(original_non_contiguous.is_contiguous())
	71	+
	72	+ dlpack_tensor = to_dlpack(original_non_contiguous)
	73	+ restored = from_dlpack(dlpack_tensor)
	74	+
	75	+ self.assertEqual(original_non_contiguous, restored)
	76	+ self.assertEqual(original_non_contiguous.stride(), restored.stride())
	77	+
	78	+ @Dtypes(torch.float)
	79	+ def test_dlpack_memory_sharing(self, dtype, device="npu"):
	80	+ """Test that dlpack shares memory with original tensor"""
	81	+ original = torch.randn(3, 3, dtype=dtype, device=device)
	82	+ original_data_ptr = original.data_ptr()
	83	+
	84	+ # Convert to dlpack and back
	85	+ dlpack_tensor = to_dlpack(original)
	86	+ restored = from_dlpack(dlpack_tensor)
	87	+
	88	+ # Check if memory is shared (data_ptr should be the same)
	89	+ self.assertEqual(original_data_ptr, restored.data_ptr())
	90	+
	91	+ # Modify original tensor and check if restored tensor is also modified
	92	+ original.fill_(42.0)
	93	+ self.assertEqual(original, restored)
	94	+
	95	+ @Dtypes(torch.bfloat16)
	96	+ def test_dlpack_bfloat16_support(self, dtype, device="npu"):
	97	+ """Test dlpack with bfloat16 data type"""
	98	+ original = torch.randn(3, 4, dtype=dtype, device=device)
	99	+ dlpack_tensor = to_dlpack(original)
	100	+ restored = from_dlpack(dlpack_tensor)
	101	+
	102	+ self.assertEqual(original, restored)
	103	+ self.assertEqual(original.dtype, restored.dtype)
	104	+
	105	+ @Dtypes(torch.float)
	106	+ def test_dlpack_cpu(self, dtype, device="cpu"):
	107	+ """Test that dlpack shares memory with original cpu tensor"""
	108	+ original = torch.randn(3, 3, dtype=dtype, device=device)
	109	+ original_data_ptr = original.data_ptr()
	110	+
	111	+ # Convert to dlpack and back
	112	+ dlpack_tensor = to_dlpack(original)
	113	+ restored = from_dlpack(dlpack_tensor)
	114	+
	115	+ # Check if memory is shared (data_ptr should be the same)
	116	+ self.assertEqual(original_data_ptr, restored.data_ptr())
	117	+
	118	+
	119	+if __name__ == '__main__':
	120	+ run_tests()

Athird_party/dlpack/dlpack.h+226-0

@@ -0,0 +1,226 @@
	1	+/*!
	2	+ * Copyright (c) 2017 by Contributors
	3	+ * \file dlpack.h
	4	+ * \brief The common header of DLPack.
	5	+ */
	6	+#ifndef DLPACK_DLPACK_H_
	7	+#define DLPACK_DLPACK_H_
	8	+
	9	+/**
	10	+ * \brief Compatibility with C++
	11	+ */
	12	+#ifdef __cplusplus
	13	+#define DLPACK_EXTERN_C extern "C"
	14	+#else
	15	+#define DLPACK_EXTERN_C
	16	+#endif
	17	+
	18	+/! \brief The current version of dlpack /
	19	+#define DLPACK_VERSION 80
	20	+
	21	+/! \brief The current ABI version of dlpack /
	22	+#define DLPACK_ABI_VERSION 1
	23	+
	24	+/! \brief DLPACK_DLL prefix for windows /
	25	+#ifdef _WIN32
	26	+#ifdef DLPACK_EXPORTS
	27	+#define DLPACK_DLL __declspec(dllexport)
	28	+#else
	29	+#define DLPACK_DLL __declspec(dllimport)
	30	+#endif
	31	+#else
	32	+#define DLPACK_DLL
	33	+#endif
	34	+
	35	+// NOLINTNEXTLINE(modernize-deprecated-headers)
	36	+#include <stdint.h>
	37	+// NOLINTNEXTLINE(modernize-deprecated-headers)
	38	+#include <stddef.h>
	39	+
	40	+#ifdef __cplusplus
	41	+extern "C" {
	42	+#endif
	43	+/*!
	44	+ * \brief The device type in DLDevice.
	45	+ */
	46	+#ifdef __cplusplus
	47	+typedef enum : int32_t {
	48	+#else
	49	+typedef enum {
	50	+#endif
	51	+ /! \brief CPU device /
	52	+ kDLCPU = 1,
	53	+ /! \brief CUDA GPU device /
	54	+ kDLCUDA = 2,
	55	+ kDLCUDAHost = 3,
	56	+ /! \brief OpenCL devices. /
	57	+ kDLOpenCL = 4,
	58	+ /! \brief Vulkan buffer for next generation graphics. /
	59	+ kDLVulkan = 7,
	60	+ /! \brief Metal for Apple GPU. /
	61	+ kDLMetal = 8,
	62	+ /! \brief Verilog simulator buffer /
	63	+ kDLVPI = 9,
	64	+ /! \brief ROCm GPUs for AMD GPUs /
	65	+ kDLROCM = 10,
	66	+ kDLROCMHost = 11,
	67	+ /*!
	68	+ * \brief Reserved extension device type,
	69	+ * used for quickly test extension device
	70	+ * The semantics can differ depending on the implementation.
	71	+ */
	72	+ kDLExtDev = 12,
	73	+ /*!
	74	+ * \brief CUDA managed/unified memory allocated by cudaMallocManaged
	75	+ */
	76	+ kDLCUDAManaged = 13,
	77	+ /*!
	78	+ * \brief Unified shared memory allocated on a oneAPI non-partititioned
	79	+ * device. Call to oneAPI runtime is required to determine the device
	80	+ * type, the USM allocation type and the sycl context it is bound to.
	81	+ *
	82	+ */
	83	+ kDLOneAPI = 14,
	84	+ /! \brief GPU support for next generation WebGPU standard. /
	85	+ kDLWebGPU = 15,
	86	+ /! \brief Qualcomm Hexagon DSP /
	87	+ kDLHexagon = 16,
	88	+ /! \brief Microsoft AI Accelerator /
	89	+ kDLMAIA = 17,
	90	+} DLDeviceType;
	91	+
	92	+/*!
	93	+ * \brief A Device for Tensor and operator.
	94	+ */
	95	+typedef struct {
	96	+ /! \brief The device type used in the device. /
	97	+ DLDeviceType device_type;
	98	+ int32_t device_id;
	99	+} DLDevice;
	100	+
	101	+/*!
	102	+ * \brief The type code options DLDataType.
	103	+ */
	104	+typedef enum {
	105	+ /! \brief signed integer /
	106	+ kDLInt = 0U,
	107	+ /! \brief unsigned integer /
	108	+ kDLUInt = 1U,
	109	+ /! \brief IEEE floating point /
	110	+ kDLFloat = 2U,
	111	+ /*!
	112	+ * \brief Opaque handle type, reserved for testing purposes.
	113	+ * Frameworks need to agree on the handle data type for the exchange to be well-defined.
	114	+ */
	115	+ kDLOpaqueHandle = 3U,
	116	+ /! \brief bfloat16 /
	117	+ kDLBfloat = 4U,
	118	+ /*!
	119	+ * \brief complex number
	120	+ * (C/C++/Python layout: compact struct per complex number)
	121	+ */
	122	+ kDLComplex = 5U,
	123	+ /! \brief boolean /
	124	+ kDLBool = 6U,
	125	+} DLDataTypeCode;
	126	+
	127	+/*!
	128	+ * \brief The data type the tensor can hold. The data type is assumed to follow the
	129	+ * native endian-ness. An explicit error message should be raised when attempting to
	130	+ * export an array with non-native endianness
	131	+ *
	132	+ * Examples
	133	+ * - float: type_code = 2, bits = 32, lanes = 1
	134	+ * - float4(vectorized 4 float): type_code = 2, bits = 32, lanes = 4
	135	+ * - int8: type_code = 0, bits = 8, lanes = 1
	136	+ * - std::complex<float>: type_code = 5, bits = 64, lanes = 1
	137	+ * - bool: type_code = 6, bits = 8, lanes = 1 (as per common array library convention, the underlying storage size of bool is 8 bits)
	138	+ */
	139	+typedef struct {
	140	+ /*!
	141	+ * \brief Type code of base types.
	142	+ * We keep it uint8_t instead of DLDataTypeCode for minimal memory
	143	+ * footprint, but the value should be one of DLDataTypeCode enum values.
	144	+ * */
	145	+ uint8_t code;
	146	+ /*!
	147	+ * \brief Number of bits, common choices are 8, 16, 32.
	148	+ */
	149	+ uint8_t bits;
	150	+ /! \brief Number of lanes in the type, used for vector types. /
	151	+ uint16_t lanes;
	152	+} DLDataType;
	153	+
	154	+/*!
	155	+ * \brief Plain C Tensor object, does not manage memory.
	156	+ */
	157	+typedef struct {
	158	+ /*!
	159	+ * \brief The data pointer points to the allocated data. This will be CUDA
	160	+ * device pointer or cl_mem handle in OpenCL. It may be opaque on some device
	161	+ * types. This pointer is always aligned to 256 bytes as in CUDA. The
	162	+ * `byte_offset` field should be used to point to the beginning of the data.
	163	+ *
	164	+ * Note that as of Nov 2021, multiply libraries (CuPy, PyTorch, TensorFlow,
	165	+ * TVM, perhaps others) do not adhere to this 256 byte aligment requirement
	166	+ * on CPU/CUDA/ROCm, and always use `byte_offset=0`. This must be fixed
	167	+ * (after which this note will be updated); at the moment it is recommended
	168	+ * to not rely on the data pointer being correctly aligned.
	169	+ *
	170	+ * For given DLTensor, the size of memory required to store the contents of
	171	+ * data is calculated as follows:
	172	+ *
	173	+ * \code{.c}
	174	+ * static inline size_t GetDataSize(const DLTensor* t) {
	175	+ * size_t size = 1;
	176	+ * for (tvm_index_t i = 0; i < t->ndim; ++i) {
	177	+ * size *= t->shape[i];
	178	+ * }
	179	+ * size = (t->dtype.bits t->dtype.lanes + 7) / 8;
	180	+ * return size;
	181	+ * }
	182	+ * \endcode
	183	+ */
	184	+ void* data;
	185	+ /! \brief The device of the tensor /
	186	+ DLDevice device;
	187	+ /! \brief Number of dimensions /
	188	+ int32_t ndim;
	189	+ /! \brief The data type of the pointer /
	190	+ DLDataType dtype;
	191	+ /! \brief The shape of the tensor /
	192	+ const int64_t* shape;
	193	+ /*!
	194	+ * \brief strides of the tensor (in number of elements, not bytes)
	195	+ * can be NULL, indicating tensor is compact and row-majored.
	196	+ */
	197	+ const int64_t* strides;
	198	+ /! \brief The offset in bytes to the beginning pointer to data /
	199	+ uint64_t byte_offset;
	200	+} DLTensor;
	201	+
	202	+/*!
	203	+ * \brief C Tensor object, manage memory of DLTensor. This data structure is
	204	+ * intended to facilitate the borrowing of DLTensor by another framework. It is
	205	+ * not meant to transfer the tensor. When the borrowing framework doesn't need
	206	+ * the tensor, it should call the deleter to notify the host that the resource
	207	+ * is no longer needed.
	208	+ */
	209	+typedef struct DLManagedTensor {
	210	+ /! \brief DLTensor which is being memory managed /
	211	+ DLTensor dl_tensor;
	212	+ /*! \brief the context of the original host framework of DLManagedTensor in
	213	+ * which DLManagedTensor is used in the framework. It can also be NULL.
	214	+ */
	215	+ void* manager_ctx;
	216	+ /! \brief Destructor signature void ()(void*) - this should be called
	217	+ * to destruct manager_ctx which holds the DLManagedTensor. It can be NULL
	218	+ * if there is no way for the caller to provide a reasonable destructor.
	219	+ * The destructors deletes the argument self as well.
	220	+ */
	221	+ void (deleter)(struct DLManagedTensor self);
	222	+} DLManagedTensor;
	223	+#ifdef __cplusplus
	224	+} // DLPACK_EXTERN_C
	225	+#endif
	226	+#endif // DLPACK_DLPACK_H_

Mtorch_npu/__init__.py+2-1

	@@ -73,7 +73,7 @@ from torch_npu.contrib.module import npu_modules
73	from torch_npu.utils import _apply_module_patch, _add_tensor_methods, _add_collect_env_methods, \	73	from torch_npu.utils import _apply_module_patch, _add_tensor_methods, _add_collect_env_methods, \
74	_add_storage_methods, _add_serialization_methods, add_dynamo_methods, add_perf_dump_patch, \	74	_add_storage_methods, _add_serialization_methods, add_dynamo_methods, add_perf_dump_patch, \
75	add_optim_method, _inductor_register_device_op_overrides, \	75	add_optim_method, _inductor_register_device_op_overrides, \
76	- _apply_npu_show_warning, _apply_npugraph_tree_methods	76	+ _apply_npu_show_warning, _apply_npugraph_tree_methods, _apply_dlpack_patch
77	from torch_npu.utils._dynamo_device import _dynamo_register_interface_for_device	77	from torch_npu.utils._dynamo_device import _dynamo_register_interface_for_device
78	from torch_npu.npu._format import _apply_npu_format_patch	78	from torch_npu.npu._format import _apply_npu_format_patch
79	import torch_npu.utils.custom_ops	79	import torch_npu.utils.custom_ops
	@@ -164,6 +164,7 @@ def _apply_sharded_grad_scaler_patch():
164	def _apply_class_patches():	164	def _apply_class_patches():
165	_apply_npu_show_warning()	165	_apply_npu_show_warning()
166	_add_storage_methods()	166	_add_storage_methods()
		167	+ _apply_dlpack_patch()
167	_apply_module_patch()	168	_apply_module_patch()
168	_add_tensor_methods()	169	_add_tensor_methods()
169	_add_serialization_methods()	170	_add_serialization_methods()

Atorch_npu/csrc/aten/common/DLConvertor.cpp+309-0

@@ -0,0 +1,309 @@
	1	+#include <ATen/Functions.h>
	2	+
	3	+#include "torch_npu/csrc/aten/common/from_blob.h"
	4	+#include "torch_npu/csrc/core/npu/NPUException.h"
	5	+#include "torch_npu/csrc/aten/common/DLConvertor.h"
	6	+
	7	+using namespace std;
	8	+namespace at {
	9	+
	10	+DLDataType getDLDataType(const Tensor& t)
	11	+{
	12	+ DLDataType dtype;
	13	+ dtype.lanes = 1;
	14	+ // Convert element size to bits by multiplying with bits per byte
	15	+ constexpr int BITS_PER_BYTE = 8;
	16	+ dtype.bits = t.element_size() * BITS_PER_BYTE;
	17	+ switch (t.scalar_type()) {
	18	+ case ScalarType::UInt1:
	19	+ case ScalarType::UInt2:
	20	+ case ScalarType::UInt3:
	21	+ case ScalarType::UInt4:
	22	+ case ScalarType::UInt5:
	23	+ case ScalarType::UInt6:
	24	+ case ScalarType::UInt7:
	25	+ case ScalarType::Byte:
	26	+ case ScalarType::UInt16:
	27	+ case ScalarType::UInt32:
	28	+ case ScalarType::UInt64:
	29	+ dtype.code = DLDataTypeCode::kDLUInt;
	30	+ break;
	31	+ case ScalarType::Int1:
	32	+ case ScalarType::Int2:
	33	+ case ScalarType::Int3:
	34	+ case ScalarType::Int4:
	35	+ case ScalarType::Int5:
	36	+ case ScalarType::Int6:
	37	+ case ScalarType::Int7:
	38	+ case ScalarType::Char:
	39	+ dtype.code = DLDataTypeCode::kDLInt;
	40	+ break;
	41	+ // NOLINTNEXTLINE(bugprone-branch-clone)
	42	+ case ScalarType::Double:
	43	+ dtype.code = DLDataTypeCode::kDLFloat;
	44	+ break;
	45	+ case ScalarType::Float:
	46	+ dtype.code = DLDataTypeCode::kDLFloat;
	47	+ break;
	48	+ // NOLINTNEXTLINE(bugprone-branch-clone)
	49	+ case ScalarType::Int:
	50	+ dtype.code = DLDataTypeCode::kDLInt;
	51	+ break;
	52	+ case ScalarType::Long:
	53	+ dtype.code = DLDataTypeCode::kDLInt;
	54	+ break;
	55	+ case ScalarType::Short:
	56	+ dtype.code = DLDataTypeCode::kDLInt;
	57	+ break;
	58	+ case ScalarType::Half:
	59	+ dtype.code = DLDataTypeCode::kDLFloat;
	60	+ break;
	61	+ case ScalarType::Bool:
	62	+ dtype.code = DLDataTypeCode::kDLBool;
	63	+ break;
	64	+ case ScalarType::ComplexHalf:
	65	+ case ScalarType::ComplexFloat:
	66	+ case ScalarType::ComplexDouble:
	67	+ dtype.code = DLDataTypeCode::kDLComplex;
	68	+ break;
	69	+ case ScalarType::BFloat16:
	70	+ dtype.code = DLDataTypeCode::kDLBfloat;
	71	+ break;
	72	+ case ScalarType::Float8_e5m2:
	73	+ case ScalarType::Float8_e5m2fnuz:
	74	+ case ScalarType::Float8_e4m3fn:
	75	+ case ScalarType::Float8_e4m3fnuz:
	76	+ case ScalarType::Float8_e8m0fnu:
	77	+ TORCH_CHECK(false, "float8 types are not supported by dlpack", PTA_ERROR(ErrCode::TYPE));
	78	+ break;
	79	+ case ScalarType::QInt8:
	80	+ case ScalarType::QUInt8:
	81	+ case ScalarType::QInt32:
	82	+ case ScalarType::QUInt4x2:
	83	+ case ScalarType::QUInt2x4:
	84	+ TORCH_CHECK(false, "QUInt/QInt types are not supported by dlpack", PTA_ERROR(ErrCode::TYPE));
	85	+ break;
	86	+ case ScalarType::Bits1x8:
	87	+ case ScalarType::Bits2x4:
	88	+ case ScalarType::Bits4x2:
	89	+ case ScalarType::Bits8:
	90	+ case ScalarType::Bits16:
	91	+ TORCH_CHECK(false, "Bit types are not supported by dlpack", PTA_ERROR(ErrCode::TYPE));
	92	+ break;
	93	+ case ScalarType::Undefined:
	94	+ TORCH_CHECK(false, "Undefined is not a valid ScalarType", PTA_ERROR(ErrCode::TYPE));
	95	+ case ScalarType::NumOptions:
	96	+ TORCH_CHECK(false, "NumOptions is not a valid ScalarType", PTA_ERROR(ErrCode::TYPE));
	97	+ }
	98	+ return dtype;
	99	+}
	100	+
	101	+static DLDevice getDLDevice(const Tensor& tensor, c10::DeviceIndex device_id)
	102	+{
	103	+ DLDevice ctx;
	104	+ ctx.device_id = static_cast<int32_t>(static_cast<unsigned char>(device_id));
	105	+ switch (tensor.device().type()) {
	106	+ case DeviceType::CPU:
	107	+ ctx.device_type = DLDeviceType::kDLCPU;
	108	+ break;
	109	+ case DeviceType::PrivateUse1:
	110	+ ctx.device_type = DLDeviceType::kDLExtDev;
	111	+ break;
	112	+ default:
	113	+ TORCH_CHECK(false, "Cannot pack tensors on " + tensor.device().str(), PTA_ERROR(ErrCode::NOT_SUPPORT));
	114	+ }
	115	+ return ctx;
	116	+}
	117	+
	118	+static Device getATenDevice(const DLDevice& ctx, void* data)
	119	+{
	120	+ switch (ctx.device_type) {
	121	+ case DLDeviceType::kDLCPU:
	122	+ return at::Device(DeviceType::CPU);
	123	+ case DLDeviceType::kDLExtDev:
	124	+ return at::Device(DeviceType::PrivateUse1, static_cast<c10::DeviceIndex>(ctx.device_id));
	125	+ default:
	126	+ TORCH_CHECK(false, "Unsupported device_type: ", std::to_string(ctx.device_type), PTA_ERROR(ErrCode::NOT_SUPPORT));
	127	+ }
	128	+}
	129	+
	130	+namespace {
	131	+constexpr int k8Bits = 8;
	132	+constexpr int k16Bits = 16;
	133	+constexpr int k32Bits = 32;
	134	+constexpr int k64Bits = 64;
	135	+constexpr int k128Bits = 128;
	136	+} // namespace
	137	+
	138	+ScalarType toScalarType(const DLDataType& dtype)
	139	+{
	140	+ ScalarType stype = ScalarType::Undefined;
	141	+ TORCH_CHECK(dtype.lanes == 1, "ATen does not support lanes != 1", PTA_ERROR(ErrCode::NOT_SUPPORT));
	142	+ switch (dtype.code) {
	143	+ case DLDataTypeCode::kDLUInt:
	144	+ switch (dtype.bits) {
	145	+ case k8Bits:
	146	+ stype = ScalarType::Byte;
	147	+ break;
	148	+ case k16Bits:
	149	+ stype = ScalarType::UInt16;
	150	+ break;
	151	+ case k32Bits:
	152	+ stype = ScalarType::UInt32;
	153	+ break;
	154	+ case k64Bits:
	155	+ stype = ScalarType::UInt64;
	156	+ break;
	157	+ default:
	158	+ TORCH_CHECK(false, "Unsupported kUInt bits ", std::to_string(dtype.bits), PTA_ERROR(ErrCode::NOT_SUPPORT));
	159	+ }
	160	+ break;
	161	+ case DLDataTypeCode::kDLInt:
	162	+ switch (dtype.bits) {
	163	+ case k8Bits:
	164	+ stype = ScalarType::Char;
	165	+ break;
	166	+ case k16Bits:
	167	+ stype = ScalarType::Short;
	168	+ break;
	169	+ case k32Bits:
	170	+ stype = ScalarType::Int;
	171	+ break;
	172	+ case k64Bits:
	173	+ stype = ScalarType::Long;
	174	+ break;
	175	+ default:
	176	+ TORCH_CHECK(false, "Unsupported kInt bits ", std::to_string(dtype.bits), PTA_ERROR(ErrCode::NOT_SUPPORT));
	177	+ }
	178	+ break;
	179	+ case DLDataTypeCode::kDLFloat:
	180	+ switch (dtype.bits) {
	181	+ case k16Bits:
	182	+ stype = ScalarType::Half;
	183	+ break;
	184	+ case k32Bits:
	185	+ stype = ScalarType::Float;
	186	+ break;
	187	+ case k64Bits:
	188	+ stype = ScalarType::Double;
	189	+ break;
	190	+ default:
	191	+ TORCH_CHECK(false, "Unsupported kFloat bits ", std::to_string(dtype.bits), PTA_ERROR(ErrCode::NOT_SUPPORT));
	192	+ }
	193	+ break;
	194	+ case DLDataTypeCode::kDLBfloat:
	195	+ switch (dtype.bits) {
	196	+ case k16Bits:
	197	+ stype = ScalarType::BFloat16;
	198	+ break;
	199	+ default:
	200	+ TORCH_CHECK(false, "Unsupported kFloat bits ", std::to_string(dtype.bits), PTA_ERROR(ErrCode::NOT_SUPPORT));
	201	+ }
	202	+ break;
	203	+ case DLDataTypeCode::kDLComplex:
	204	+ switch (dtype.bits) {
	205	+ case k32Bits:
	206	+ stype = ScalarType::ComplexHalf;
	207	+ break;
	208	+ case k64Bits:
	209	+ stype = ScalarType::ComplexFloat;
	210	+ break;
	211	+ case k128Bits:
	212	+ stype = ScalarType::ComplexDouble;
	213	+ break;
	214	+ default:
	215	+ TORCH_CHECK(false, "Unsupported kFloat bits ", std::to_string(dtype.bits), PTA_ERROR(ErrCode::NOT_SUPPORT));
	216	+ }
	217	+ break;
	218	+ case DLDataTypeCode::kDLBool:
	219	+ switch (dtype.bits) {
	220	+ case k8Bits:
	221	+ stype = ScalarType::Bool;
	222	+ break;
	223	+ default:
	224	+ TORCH_CHECK(false, "Unsupported kDLBool bits ", std::to_string(dtype.bits), PTA_ERROR(ErrCode::NOT_SUPPORT));
	225	+ }
	226	+ break;
	227	+ default:
	228	+ TORCH_CHECK(false, "Unsupported code ", std::to_string(dtype.code), PTA_ERROR(ErrCode::NOT_SUPPORT));
	229	+ }
	230	+ return stype;
	231	+}
	232	+
	233	+namespace {
	234	+struct ATenDLMTensor {
	235	+ Tensor handle;
	236	+ DLManagedTensor tensor{};
	237	+};
	238	+} // namespace
	239	+
	240	+static void deleter(DLManagedTensor* arg)
	241	+{
	242	+ delete static_cast<ATenDLMTensor*>(arg->manager_ctx);
	243	+}
	244	+
	245	+// This function returns a shared_ptr to memory managed DLpack tensor
	246	+// constructed out of ATen tensor
	247	+DLManagedTensor* toDLPack(const Tensor& src)
	248	+{
	249	+ // create a new tensor with possibly normalized strides
	250	+ // gh-83069
	251	+ auto shape = src.sizes();
	252	+ auto strides = src.strides().vec();
	253	+ static constexpr int64_t kMinDimForStride = 2;
	254	+ for (int i = 0; i < src.dim(); i++) {
	255	+ if (shape[i] < kMinDimForStride) {
	256	+ strides[i] = 1;
	257	+ }
	258	+ }
	259	+
	260	+ auto view = src.as_strided(shape, strides, src.storage_offset());
	261	+ ATenDLMTensor* atDLMTensor(new ATenDLMTensor);
	262	+ atDLMTensor->handle = view;
	263	+ atDLMTensor->tensor.manager_ctx = atDLMTensor;
	264	+ atDLMTensor->tensor.deleter = &deleter;
	265	+ atDLMTensor->tensor.dl_tensor.data = view.data_ptr();
	266	+ c10::DeviceIndex device_id = 0;
	267	+ if (src.is_cuda() \|\| src.is_privateuseone()) {
	268	+ device_id = src.get_device();
	269	+ }
	270	+ atDLMTensor->tensor.dl_tensor.device = getDLDevice(src, device_id);
	271	+ atDLMTensor->tensor.dl_tensor.ndim = static_cast<int32_t>(src.dim());
	272	+ atDLMTensor->tensor.dl_tensor.dtype = getDLDataType(src);
	273	+ atDLMTensor->tensor.dl_tensor.shape = view.sizes().data();
	274	+ atDLMTensor->tensor.dl_tensor.strides = view.strides().data();
	275	+ atDLMTensor->tensor.dl_tensor.byte_offset = 0;
	276	+ return &(atDLMTensor->tensor);
	277	+}
	278	+
	279	+Tensor fromDLPack(DLManagedTensor* src)
	280	+{
	281	+ auto deleter = [src](void* self [[maybe_unused]]) {
	282	+ if (src->deleter) {
	283	+ src->deleter(src);
	284	+ }
	285	+ };
	286	+ return fromDLPack(src, std::move(deleter));
	287	+}
	288	+
	289	+Tensor fromDLPack(DLManagedTensor* src, std::function<void(void*)> deleter)
	290	+{
	291	+ Device device = getATenDevice(src->dl_tensor.device, src->dl_tensor.data);
	292	+ ScalarType stype = toScalarType(src->dl_tensor.dtype);
	293	+ if (!src->dl_tensor.strides) {
	294	+ return at_npu::native::from_blob(
	295	+ src->dl_tensor.data,
	296	+ IntArrayRef(src->dl_tensor.shape, src->dl_tensor.ndim),
	297	+ std::move(deleter),
	298	+ at::device(device).dtype(stype),
	299	+ {device});
	300	+ }
	301	+ return at_npu::native::from_blob(
	302	+ src->dl_tensor.data,
	303	+ IntArrayRef(src->dl_tensor.shape, src->dl_tensor.ndim),
	304	+ IntArrayRef(src->dl_tensor.strides, src->dl_tensor.ndim),
	305	+ deleter,
	306	+ at::device(device).dtype(stype),
	307	+ {device});
	308	+}
	309	+} // namespace at

Atorch_npu/csrc/aten/common/DLConvertor.h+20-0

@@ -0,0 +1,20 @@
	1	+#pragma once
	2	+
	3	+#include <ATen/ATen.h>
	4	+#include <ATen/Tensor.h>
	5	+#include "third_party/dlpack/dlpack.h"
	6	+
	7	+// this convertor will:
	8	+// 1) take a Tensor object and wrap it in the DLPack tensor
	9	+// 2) take a dlpack tensor and convert it to the ATen Tensor
	10	+
	11	+namespace at {
	12	+
	13	+TORCH_API ScalarType toScalarType(const DLDataType& dtype);
	14	+TORCH_API DLManagedTensor* toDLPack(const Tensor& src);
	15	+TORCH_API Tensor fromDLPack(DLManagedTensor* src);
	16	+TORCH_API Tensor fromDLPack(DLManagedTensor* src, std::function<void(void*)> deleter);
	17	+TORCH_API DLDataType getDLDataType(const Tensor& t);
	18	+TORCH_API DLDevice getDLContext(const Tensor& tensor, const int64_t& device_id);
	19	+
	20	+} // namespace at

Mtorch_npu/csrc/npu/Module.cpp+67-0

	@@ -17,10 +17,12 @@
17	#include <torch/csrc/utils/pybind.h>	17	#include <torch/csrc/utils/pybind.h>
18	#include <torch/csrc/utils/python_strings.h>	18	#include <torch/csrc/utils/python_strings.h>
19	#include <torch/csrc/utils/python_numbers.h>	19	#include <torch/csrc/utils/python_numbers.h>
		20	+#include <torch/csrc/utils/tensor_numpy.h>
20	#include <torch/csrc/profiler/python/combined_traceback.h>	21	#include <torch/csrc/profiler/python/combined_traceback.h>
21		22
22	#include "torch_npu/csrc/aten/NPUGeneratorImpl.h"	23	#include "torch_npu/csrc/aten/NPUGeneratorImpl.h"
23	#include "torch_npu/csrc/aten/NPUNativeFunctions.h"	24	#include "torch_npu/csrc/aten/NPUNativeFunctions.h"
		25	+#include "torch_npu/csrc/aten/common/DLConvertor.h"
24	#include "torch_npu/csrc/aten/common/SetNpu.h"	26	#include "torch_npu/csrc/aten/common/SetNpu.h"
25	#include "torch_npu/csrc/core/npu/NPUException.h"	27	#include "torch_npu/csrc/core/npu/NPUException.h"
26	#include "torch_npu/csrc/core/npu/NPUFunctions.h"	28	#include "torch_npu/csrc/core/npu/NPUFunctions.h"
	@@ -1855,6 +1857,69 @@ static PyObject* THNPModule_reset_device_res_limit(PyObject* self, PyObject *arg
1855	END_HANDLE_TH_ERRORS	1857	END_HANDLE_TH_ERRORS
1856	}	1858	}
1857		1859
		1860	+static void DLPack_Capsule_Destructor(PyObject* data)
		1861	+{
		1862	+ if (C10_LIKELY(!PyCapsule_IsValid(data, "dltensor"))) {
		1863	+ // early out, see DLPack spec: if a consuming library sets the capsule
		1864	+ // name to something else, they own it and we don't need to do anything
		1865	+ return;
		1866	+ }
		1867	+ HANDLE_TH_ERRORS
		1868	+ // Causes overheads for validity checks again, but this case is rare
		1869	+ // since consuming libraries should rename the capsule according to spec.
		1870	+ // Note that this cannot set a python error (we checked validity above),
		1871	+ // so we don't need to handle python error state here.
		1872	+ DLManagedTensor* dlMTensor =
		1873	+ (DLManagedTensor*)PyCapsule_GetPointer(data, "dltensor");
		1874	+ // the dlMTensor has not been consumed, call deleter ourselves.
		1875	+ // DLPack spec mentions that deleter may be NULL, but deleter from
		1876	+ // `at::toDLPack` is never NULL, so no need for an additional check here.
		1877	+ dlMTensor->deleter(dlMTensor);
		1878	+ END_HANDLE_TH_ERRORS_RET()
		1879	+}
		1880	+
		1881	+static PyObject* THPModule_toDLPack(PyObject* _unused, PyObject* data)
		1882	+{
		1883	+ HANDLE_TH_ERRORS
		1884	+ TORCH_CHECK(THPVariable_Check(data), "data must be a Tensor");
		1885	+ DLManagedTensor* dlMTensor = at::toDLPack(THPVariable_Unpack(data));
		1886	+ return PyCapsule_New(dlMTensor, "dltensor", DLPack_Capsule_Destructor);
		1887	+ END_HANDLE_TH_ERRORS
		1888	+}
		1889	+
		1890	+static PyObject* THPModule_fromDLPack(PyObject* _unused, PyObject* data)
		1891	+{
		1892	+ using namespace torch::autograd;
		1893	+ HANDLE_TH_ERRORS
		1894	+ DLManagedTensor* dlMTensor =
		1895	+ (DLManagedTensor*)PyCapsule_GetPointer(data, "dltensor");
		1896	+ TORCH_CHECK(
		1897	+ dlMTensor,
		1898	+ "from_dlpack received an invalid capsule. "
		1899	+ "Note that DLTensor capsules can be consumed only once, "
		1900	+ "so you might have already constructed a tensor from it once.");
		1901	+
		1902	+ auto deleter_with_gil = [dlMTensor](void*) {
		1903	+ if (dlMTensor->deleter) {
		1904	+ pybind11::gil_scoped_acquire gil;
		1905	+ dlMTensor->deleter(dlMTensor);
		1906	+ }
		1907	+ };
		1908	+
		1909	+ // atensor steals the ownership of the underlying storage. It also passes a
		1910	+ // destructor function that will be called when the underlying storage goes
		1911	+ // out of scope. When the destructor is called, the dlMTensor is destructed
		1912	+ // too.
		1913	+ // HACK: Ensure that we hold the GIL here just in case the
		1914	+ // managed tensor originating from a buggy NumPy build.
		1915	+ auto atensor = at::fromDLPack(dlMTensor);
		1916	+
		1917	+ // Make sure this capsule will never be used again.
		1918	+ PyCapsule_SetName(data, "used_dltensor");
		1919	+ return THPVariable_Wrap(atensor);
		1920	+ END_HANDLE_TH_ERRORS
		1921	+}
		1922	+
1858	static PyObject* THNPModule_set_stream_res_limit(PyObject* self, PyObject args, PyObject kwargs)	1923	static PyObject* THNPModule_set_stream_res_limit(PyObject* self, PyObject args, PyObject kwargs)
1859	{	1924	{
1860	HANDLE_TH_ERRORS	1925	HANDLE_TH_ERRORS
	@@ -2019,6 +2084,8 @@ static struct PyMethodDef THNPModule_methods[] = {
2019	{"_npu_set_stream_res_limit", (PyCFunction)THNPModule_set_stream_res_limit, METH_VARARGS \| METH_KEYWORDS, nullptr},	2084	{"_npu_set_stream_res_limit", (PyCFunction)THNPModule_set_stream_res_limit, METH_VARARGS \| METH_KEYWORDS, nullptr},
2020	{"_npu_reset_stream_res_limit", (PyCFunction)THNPModule_reset_stream_res_limit, METH_VARARGS \| METH_KEYWORDS, nullptr},	2085	{"_npu_reset_stream_res_limit", (PyCFunction)THNPModule_reset_stream_res_limit, METH_VARARGS \| METH_KEYWORDS, nullptr},
2021	{"_npu_get_stream_res_limit", (PyCFunction)THNPModule_get_stream_res_limit, METH_VARARGS \| METH_KEYWORDS, nullptr},	2086	{"_npu_get_stream_res_limit", (PyCFunction)THNPModule_get_stream_res_limit, METH_VARARGS \| METH_KEYWORDS, nullptr},
		2087	+ {"_npu_to_dlpack", (PyCFunction)THPModule_toDLPack, METH_O, nullptr},
		2088	+ {"_npu_from_dlpack", (PyCFunction)THPModule_fromDLPack, METH_O, nullptr},
2022	{nullptr}};	2089	{nullptr}};
2023		2090
2024	TORCH_NPU_API PyMethodDef* THNPModule_get_methods()	2091	TORCH_NPU_API PyMethodDef* THNPModule_get_methods()

Mtorch_npu/utils/__init__.py+2-0

	@@ -21,6 +21,8 @@ from .flops_count import _FlopsCounter as FlopsCounter
21	from .affinity import _set_thread_affinity as set_thread_affinity	21	from .affinity import _set_thread_affinity as set_thread_affinity
22	from .affinity import _reset_thread_affinity as reset_thread_affinity	22	from .affinity import _reset_thread_affinity as reset_thread_affinity
23	from ._graph_tree import _apply_npugraph_tree_methods	23	from ._graph_tree import _apply_npugraph_tree_methods
		24	+from .dlpack import _apply_dlpack_patch
		25	+
24		26
25		27
26	# init flopcount	28	# init flopcount

Atorch_npu/utils/dlpack.py+85-0

@@ -0,0 +1,85 @@
	1	+import enum
	2	+import torch
	3	+
	4	+from torch_npu._C import _npu_from_dlpack
	5	+from torch_npu._C import _npu_to_dlpack
	6	+
	7	+
	8	+def _to_dlpack(tensor):
	9	+ return _npu_to_dlpack(tensor)
	10	+
	11	+
	12	+def _from_dlpack(ext_tensor) -> 'torch.Tensor':
	13	+ if hasattr(ext_tensor, '__dlpack__'):
	14	+ dlpack = ext_tensor.__dlpack__()
	15	+ else:
	16	+ # Old versions just call the converter
	17	+ dlpack = ext_tensor
	18	+ return _npu_from_dlpack(dlpack)
	19	+
	20	+
	21	+def _apply_dlpack_patch():
	22	+ """Patch torch.utils.dlpack and torch.utils to use torch_npu implementation for NPU tensors"""
	23	+ import torch.utils.dlpack as torch_dlpack
	24	+
	25	+ # Store original functions
	26	+ _original_to_dlpack = torch_dlpack.to_dlpack
	27	+ _original_from_dlpack = torch_dlpack.from_dlpack
	28	+
	29	+ def create_patched_to_dlpack(module_name):
	30	+ """Create a patched to_dlpack function with proper __module__ attribute"""
	31	+ def patched_to_dlpack(tensor):
	32	+ """Patched to_dlpack that uses torch_npu implementation for NPU tensors"""
	33	+ if hasattr(tensor, 'device') and tensor.device.type == 'npu':
	34	+ return _to_dlpack(tensor)
	35	+ return _original_to_dlpack(tensor)
	36	+ patched_to_dlpack.__module__ = module_name
	37	+ return patched_to_dlpack
	38	+
	39	+ def create_patched_from_dlpack(module_name):
	40	+ """Create a patched from_dlpack function with proper __module__ attribute"""
	41	+ def patched_from_dlpack(ext_tensor):
	42	+ """Patched from_dlpack that uses torch_npu implementation when appropriate"""
	43	+ # For NPU tensors or when torch_npu is available, use our implementation
	44	+ try:
	45	+ return _from_dlpack(ext_tensor)
	46	+ except Exception:
	47	+ # Fallback to original implementation
	48	+ return _original_from_dlpack(ext_tensor)
	49	+ patched_from_dlpack.__module__ = module_name
	50	+ return patched_from_dlpack
	51	+
	52	+ # Apply patches to torch.utils.dlpack
	53	+ torch_dlpack.to_dlpack = create_patched_to_dlpack('torch.utils.dlpack')
	54	+ torch_dlpack.from_dlpack = create_patched_from_dlpack('torch.utils.dlpack')
	55	+
	56	+ # Also patch torch.utils.to_dlpack and torch.utils.from_dlpack if they exist
	57	+ if hasattr(torch.utils, 'to_dlpack'):
	58	+ _original_torch_utils_to_dlpack = torch.utils.to_dlpack
	59	+ torch.utils.to_dlpack = create_patched_to_dlpack('torch.utils')
	60	+
	61	+ if hasattr(torch.utils, 'from_dlpack'):
	62	+ _original_torch_utils_from_dlpack = torch.utils.from_dlpack
	63	+ torch.utils.from_dlpack = create_patched_from_dlpack('torch.utils')
	64	+
	65	+ # Also patch torch.from_dlpack and torch.to_dlpack if they exist
	66	+ if hasattr(torch, 'from_dlpack'):
	67	+ _original_torch_from_dlpack = torch.from_dlpack
	68	+ torch.from_dlpack = create_patched_from_dlpack('torch')
	69	+
	70	+ if hasattr(torch, 'to_dlpack'):
	71	+ _original_torch_to_dlpack = torch.to_dlpack
	72	+ torch.to_dlpack = create_patched_to_dlpack('torch')
	73	+
	74	+ # Add to_dlpack to torch.__all__ if it exists, otherwise create it
	75	+ if not hasattr(torch, '__all__'):
	76	+ torch.__all__ = []
	77	+ if 'to_dlpack' not in torch.__all__:
	78	+ torch.__all__.append('to_dlpack')
	79	+
	80	+ # Also ensure from_dlpack is in torch.__all__ if it exists
	81	+ if hasattr(torch, 'from_dlpack'):
	82	+ if not hasattr(torch, '__all__'):
	83	+ torch.__all__ = []
	84	+ if 'from_dlpack' not in torch.__all__:
	85	+ torch.__all__.append('from_dlpack')