TVM Runtime NDArray 完全指南：tvm.nd.array / tvm.nd.empty / NDArray 核心 API 与底层实现解析

编译器深度学习模型优化【免费下载链接】tvmOpen deep learning compiler stack for cpu, gpu and specialized accelerators项目地址https://gitcode.com/gh_mirrors/tvm7/tvm点击查看免费下载本文基于 tvm.runtime.ndarray API 参考文档 编写系统讲解 TVMApache TVMPython 前端中tvm.nd.NDArray这一运行时核心数据容器的完整用法。NDArray 是连接前端Relay / Relax / TE与后端运行时CPU / GPU / 各类加速器的数据交接桥掌握tvm.nd.array、tvm.nd.empty以及copyto、numpy、DLPack 互操作等 API是编写可执行的 TVM 推理脚本、自定义算子测试与多设备部署代码的必备技能。读完本文你将能够熟练创建、转换、迁移 NDArray并理解其背后的引用计数与 FFI 调用链。一、NDArray 是什么轻量数据容器而非新的数组库在 python/tvm/runtime/ndarray.py 的类定义中NDArray被明确定义为Lightweight NDArray class of TVM runtime. Strictly this is only an Array Container (a buffer object). No arithmetic operations are defined. All operations are performed by TVM functions.这段话揭示了 TVM 的设计哲学NDArray 只是一个缓冲区对象负责承载数据内存与元信息shape / dtype / device不定义任何算术运算如加法、矩阵乘等。所有计算都由 TVM 的编译产物通过tvm.build生成的函数或 TVM 算子库来执行NDArray 仅作为输入输出的载体。TVM 的目标不是再造一个数组库如 NumPy而是提供一个极简的数据结构便于在已有项目可能已有自己的数组容器中集成 TVM。与 NumPy 的关系NDArray 与 NumPy 是互相转换而非互相替代的关系从 NumPy 到 NDArraytvm.nd.array(numpy_array, device...)会执行一次数据拷贝见下文copyfrom的实现从 NDArray 到 NumPyarr.numpy()同样通过内存拷贝完成在 Python 侧__repr__与__str__的实现中最终输出内容也是通过self.numpy().__repr__()得到的参见 ndarray.py。二、创建 NDArraytvm.nd.array 与 tvm.nd.emptyAPI 文档中给出了两个核心创建函数均在 python/tvm/runtime/ndarray.py 中实现。2.1 tvm.nd.array —— 从源数据创建tvm.nd.array(arr, devicetvm.cpu(0), mem_scopeNone)参数类型说明arrnumpy.ndarray 或 array_like待拷贝的数据源deviceDevice, 可选数组所在的设备默认 CPU 设备 0mem_scopestr, 可选数组的内存作用域memory scope返回NDArray创建的数组其核心实现逻辑ndarray.pyif not isinstance(arr, (np.ndarray, NDArray)): arr np.array(arr) return empty(arr.shape, arr.dtype, device, mem_scope).copyfrom(arr)可见tvm.nd.array本质上是一个先empty分配、再copyfrom拷贝的组合操作如果传入的不是 NumPy 数组会先通过np.array(arr)自动转换因此可以直接传入 Python list如果传入的已经是NDArray则跳过转换注意如果传入的是tvm.ir.container.ArrayIR 容器会直接抛出AttributeError防止误用。2.2 tvm.nd.empty —— 只分配不初始化tvm.nd.empty(shape, dtypefloat32, devicetvm.cpu(0), mem_scopeNone)参数类型说明shapeShapeTuple 或 Sequence数组形状如[1024, 1024]dtypestr数据类型默认float32deviceDevice数组所在设备mem_scopestr, 可选内存作用域返回NDArray未初始化数据的空数组实现中先将 shape 规范化为tvm.runtime.ShapeTuple再通过 FFI 调用底层分配arr _ffi_api.TVMArrayAllocWithScope(shape, dtype, device, mem_scope)TVMArrayAllocWithScope在 C 侧注册于 src/runtime/ndarray.ccTVM_REGISTER_GLOBAL(runtime.TVMArrayAllocWithScope).set_body_typed(NDArray::Empty);即最终调用NDArray::Empty声明于 include/tvm/runtime/ndarray.h由各设备CPU、CUDA、OpenCL 等的分配器完成实际内存申请。测试 test_runtime_ndarray.py 中的test_memory_usage还验证了empty分配后设备可用内存会相应减少且del arr释放最后一个引用后内存会恢复——这正是下一节要讲的引用计数机制。2.3 支持的 dtypeNDArray 支持 TVM 完整的数据类型体系。在 python/tvm/_ffi/runtime_ctypes.py 的STR2DTYPE中登记了void、bool、int8/16/32/64、uint8/16/32/64、float16/32/64以及面向量化与训练场景的e4m3_float8、e5m2_float8二者转换回 NumPy 时需要安装ml_dtypes库。此外 TVM 还支持向量化 lane 类型如float32x4dtype 字符串中的x4后缀会体现在DataType.lanes字段上。三、NDArray 核心属性与方法NDArray的完整 API 通过autoclass指令自动收集到 ndarray.rst 中以下逐一讲解。3.1 属性shape、dtype、devicearr tvm.nd.array(np.random.rand(3, 4)) print(arr.shape) # (3, 4) print(arr.dtype) # float32 print(arr.device) # Device(0, 0) device_type1 表示 CPUdtype属性直接读取self.handle.contents.dtype其中handle是TVMArrayHandle指向 DLTensor 结构device属性返回self.handle.contents.device即 DLTensor 中记录的DLDevice结构device_type device_idshape则由 C 扩展层的NDArrayBase提供。这三个属性与 DLPack 规范中的DLTensor元信息一一对应是数组身份的核心标识。3.2 转换到 NumPynumpy() 与 asnumpy()np_arr arr.numpy() # 推荐用法 np_arr arr.asnumpy() # 旧 API会触发 DeprecationWarningnumpy()的实现ndarray.py要点若 dtype 是向量化类型lanes 1会先把lanes维度展开到 shape 末尾对特殊类型做归一化int4映射为int8bfloat16映射为uint16e4m3_float8/e5m2_float8在安装ml_dtypes时映射为对应 float8 类型否则抛RuntimeError分配 C 连续的 NumPy 数组后调用TVMArrayCopyToBytes从设备内存拷贝数据int4类型还需额外做高低半字节的重新排列才能还原真实数据。asnumpy()在 ndarray.py 中明确标注将在 TVM v0.8 起废弃请改用numpy()。3.3 数据拷贝copyfrom 与 copytocopyfrom外部 → NDArray同步拷贝arr.copyfrom(numpy_array) # 返回 arr 自身便于链式调用实现ndarray.py会做多重校验与适配若源是 NDArray直接调用其copyto若源不是 NumPy 数组先按self.dtype转换检查 shape 严格匹配否则抛ValueError对于bfloat16先将数据视为uint16字节流拷贝若源不是 C 连续先通过np.ascontiguousarray转换最终调用 C 函数TVMArrayCopyFromBytes(self.handle, data, nbytes)C 实现见 src/runtime/ndarray.cc该拷贝总是同步的会触发一次TVMSynchronize。copytoNDArray → 目标target_arr tvm.nd.empty((3, 4), devicetvm.cuda(0)) arr.copyto(target_arr) # 拷贝到已有数组跨设备也可以 gpu_arr arr.copyto(tvm.cuda(0)) # 传 Device自动分配目标数组实现ndarray.py目标是NDArrayBase调用 C 层的_copyto对应TVMArrayCopyFromTo见 src/runtime/ndarray.cc要求两侧字节数严格相等ICHECK_EQ目标是Device先empty分配同 shape/dtype 的数组再拷贝其他类型抛ValueError。切片赋值setitemarr[:] numpy_array或arr[:] other_ndarray被重载为数据拷贝操作。注意它只支持全切片[:]如果传入带 start/stop 的切片会抛ValueError(Array only support set from numpy array)这是有意设计的简化约束参见 ndarray.py。3.4 身份比较same_as /eq/hashNDArray 的相等性比较是基于对象身份底层句柄指针而非数据内容的a tvm.nd.array([1, 2, 3]) b tvm.nd.array([1, 2, 3]) print(a b) # False不同分配 print(a.same_as(a)) # True实现见 ndarray.py__hash__返回handle指针的值__eq__委托给same_as后者仅当两对象是同一底层引用时才返回 True。若需要比较内容应使用np.testing.assert_equal(a.numpy(), b.numpy())这也是测试代码中的通行写法。3.5 创建视图_create_viewNDArray._create_view允许基于现有数组共享底层分配、创建不同逻辑 shape 的视图ndarray.py通过 FFI 的TVMArrayCreateViewC 注册见 src/runtime/ndarray.cc声明见 include/tvm/runtime/ndarray.h实现。文档注释明确警告该接口只应用于底层内存操作会破坏 TVM 的非别名non-aliasing假设未来可能被移除普通用户不应使用。四、设备管理与 Device 对象tvm.nd.empty/tvm.nd.array的device参数接收tvm.runtime.Device。构造 Device 有两种途径4.1 tvm.device(dev_type, dev_id) 通用构造函数assert tvm.device(cpu, 1) tvm.cpu(1) assert tvm.device(cuda, 0) tvm.cuda(0) assert tvm.device(cuda:2) tvm.cuda(2) # 字符串内联设备号实现见 ndarray.py支持int类型掩码或字符串两种输入字符串支持cuda、cuda:2等格式冒号后的数字会覆盖dev_id参数无法识别的设备名抛ValueError。4.2 便捷设备构造函数模块内为每种设备提供了专用工厂函数全部位于 ndarray.py函数对应设备备注tvm.cpu(dev_id0)CPU默认设备tvm.cuda(dev_id0)NVIDIA GPU0.9.0 起推荐替代tvm.gputvm.gpu(dev_id0)CUDA GPU已废弃调用会发 DeprecationWarningtvm.rocm(dev_id0)AMD ROCm GPUtvm.opencl(dev_id0)/tvm.clOpenCL 设备cl opencl别名tvm.metal(dev_id0)/tvm.mtlApple Metalmtl metal别名tvm.vulkan(dev_id0)Vulkan 设备tvm.vpi(dev_id0)VPI 模拟设备用于 RTL 仿真tvm.hexagon(dev_id0)Qualcomm Hexagontvm.webgpu(dev_id0)WebGPUtvm.ext_dev(dev_id0)扩展设备保留给插件设备 API 快速试验设备名称到类型掩码的映射表定义在 python/tvm/_ffi/runtime_ctypes.py 的Device.STR2MASK中其中llvm、stackvm、c、test、hybrid、composite等编译目标名称也统一映射到 CPU。Device结构体runtime_ctypes.py还提供exist设备是否存在且可访问、max_threads_per_block、available_global_memory等运行时属性后者被内存测试用例直接使用。五、DLPack 互操作零拷贝跨框架数据交换DLPack 是 TVM 支持的标准张量内存共享协议NDArray 原生支持 DLPack 导出与导入实现零拷贝的数据交换。5.1 导出NDArray → 其他框架capsule arr.__dlpack__() # 返回 DLPack capsule device_type, device_id arr.__dlpack_device__()__dlpack__(self, streamNone)ndarray.py返回封装了DLManagedTensor的PyCapsulestream参数用于多流场景下由消费者告知生产者同步语义__dlpack_device__()ndarray.py按 DLPack 约定返回(device_type, device_id)元组。实现了这两个协议方法后NDArray 即可直接被 PyTorch、JAX 等支持 DLPack 的框架零拷贝接收。5.2 导入其他框架 → NDArraytorch_tensor torch.arange(6).reshape(2, 3).cuda() tvm_arr tvm.nd.from_dlpack(torch_tensor) # 零拷贝视图from_dlpackndarray.py接受两类输入具有__dlpack__方法的任意对象如 PyTorch Tensor已经是PyCapsule形式的 DLPack capsule。其实现会取出 DLPack 张量的指针创建数组视图并接管原 DLPack tensor 的析构职责移除原始 destructor从而保证生命周期安全。这也是test_runtime_packed_func.py等测试中tvm.nd.array([1, 2, 3])直接用于 PackedFunc 参数的基础。六、底层原理从 Python 到 C 的调用链以tvm.nd.empty为例一次分配经历如下完整链路tvm.nd.empty(shape, dtype, device) → tvm.runtime.ndarray.empty → _ffi_api.TVMArrayAllocWithScope(shape, dtype, device, mem_scope) → NDArray::Empty(ShapeTuple, DLDataType, Device, mem_scope) [src/runtime/ndarray.cc:385] → 设备分配器CPU 走 kDLCPU AllocatorCUDA 走 cudaMalloc 等 → 返回引用计数管理的 ContainerObjectRefNDArray 在 C 侧是ObjectRef的封装include/tvm/runtime/ndarray.h底层由NDArray::Container继承Object与ContainerBase见 ndarray.h持有DLTensor数据指针与引用计数。因此引用计数管理del arr后若引用计数归零底层分配的内存会被立即释放测试test_memory_usage已验证该行为C 接口兼容底层还导出TVMArrayAlloc、TVMArrayFree、TVMArrayCopyFromTo、TVMArrayCopyFromBytes、TVMArrayCopyToBytes等 C APIsrc/runtime/ndarray.ccPython 层的copyfrom/numpy分别通过TVMArrayCopyFromBytes/TVMArrayCopyToBytes与运行时交互对齐要求NDArray::IsAlignedndarray.cc要求data byte_offset满足kAllocAlignment对齐拷贝时CopyFromTo也会校验IsAlignedndarray.cc与字节数一致ICHECK_EQ。七、实战完整可运行的 NDArray 工作流结合 test_runtime_ndarray.py 与 test_minimal_target_codegen_llvm.py 中的通行用法给出一个端到端示例import numpy as np import tvm from tvm import te # 1. 创建 NDArrayCPU x_np np.random.randint(0, 10, size(3, 4)).astype(float32) x tvm.nd.array(x_np, devicetvm.cpu(0)) # 拷贝进 TVM assert isinstance(x, tvm.nd.NDArray) assert x.shape x_np.shape and x.dtype x_np.dtype # 2. 跨设备迁移 cuda_arr x.copyto(tvm.cuda(0)) # 若机器有 GPU back_cpu cuda_arr.copyto(tvm.cpu(0)) # 3. 与编译产物配合构建一个简单的 te 算子 n 100 A te.placeholder((n,), dtypefloat32) B te.compute((n,), lambda i: A[i] 1.0) func tvm.build(te.create_schedule([B.op]), [A, B], llvm) in_arr tvm.nd.array(np.random.randn(n).astype(float32)) out_arr tvm.nd.empty((n,), dtypefloat32) func(in_arr, out_arr) # NDArray 直接作为 PackedFunc 参数 # 4. 转回 NumPy 验证 np.testing.assert_allclose(out_arr.numpy(), in_arr.numpy() 1.0) # 5. DLPack 零拷贝互操作以 torch 为例需已安装 torch # torch_t torch.arange(6).reshape(2, 3) # tvm_arr tvm.nd.from_dlpack(torch_t) # tvm_arr[:] torch_t.numpy() * 2 # 支持全切片赋值注意事项同步语义copyfrom、copyto、numpy()涉及设备的拷贝总是同步的底层会触发TVMSynchronize涉及 GPU 上下文时这是可靠获取结果的前提不要用比较内容NDArray 的是引用比较内容校验请用np.testing.assert_*asnumpy()已废弃新代码统一使用numpy()gpu()已废弃新代码统一使用cuda()。八、相关文档与源码索引API 参考文档docs/reference/api/python/ndarray.rstautomodule/autoclass/autofunction自动生成的模块级 API 页面Python 实现python/tvm/runtime/ndarray.pyNDArray 类、array/empty/from_dlpack、全部设备构造函数FFI 类型定义python/tvm/_ffi/runtime_ctypes.pyDataType、Device、DataTypeCode、TVMArray、STR2DTYPE、STR2MASKC 头文件include/tvm/runtime/ndarray.hNDArray类与Container定义、Empty/CreateView/ToDLPack声明C 实现src/runtime/ndarray.ccNDArray::Empty、CopyFromTo、C API 与 PackedFunc 注册单元测试tests/python/all-platform-minimal-test/test_runtime_ndarray.py覆盖创建、跨设备拷贝、内存生命周期、fp16 转换综上tvm.runtime.ndarray虽然 API 表面简洁却是 TVM 运行时体系的地基向上承接 NumPy / DLPack 生态的数据进出向下通过 FFI 直通各设备分配器。理解 NDArray 的容器定位、引用计数生命周期与拷贝/转换语义是写出正确、高效 TVM 推理与测试代码的第一步。赞分享编译器深度学习模型优化【免费下载链接】tvmOpen deep learning compiler stack for cpu, gpu and specialized accelerators项目地址https://gitcode.com/gh_mirrors/tvm7/tvm点击查看免费下载相关推荐MXNet 稀疏 NDArray 完全指南CSRNDArray 与 RowSparseNDArray 的构造、算子与底层实现MXNet 稀疏 NDArray 完全指南CSRNDArray 与 RowSparseNDArray 的构造、算子与底层实现 MXNet 的 mxnet.nd深度学习人工智能机器学习分布式训练TVM4J 实战指南在 JVM 中接入 Apache TVM Runtime实现 NDArray、PackedFunc、共享库加载与 RPCTVM4J 实战指南在 JVM 中接入 Apache TVM Runtime实现 NDArray、PackedFunc、共享库加载与 RPC TVM4J 是编译器深度学习模型优化Apache MXNet NDArray 完全指南mxnet.ndarray 命令式张量核心解析Apache MXNet NDArray 完全指南mxnet.ndarray 命令式张量核心解析 导读 mxnet.ndarray 是 Apache MXNe深度学习人工智能机器学习分布式训练创作声明：本文部分内容由AI辅助生成（AIGC），仅供参考