企业官网建设流程全解析

• 编译器
• 深度学习
• 模型优化

本文基于 tvm.runtime.ndarray API 参考文档 编写，系统讲解 TVM（Apache TVM）Python 前端中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 到 NDArray：tvm.nd.array(numpy_array, device=...)会执行一次数据拷贝（见下文copyfrom的实现）；
• 从 NDArray 到 NumPy：arr.numpy()同样通过内存拷贝完成；
• 在 Python 侧__repr__与__str__的实现中，最终输出内容也是通过self.numpy().__repr__()得到的，参见 ndarray.py。

二、创建 NDArray：tvm.nd.array 与 tvm.nd.empty

API 文档中给出了两个核心创建函数，均在 python/tvm/runtime/ndarray.py 中实现。

2.1 tvm.nd.array —— 从源数据创建

tvm.nd.array(arr, device=tvm.cpu(0), mem_scope=None)

其核心实现逻辑（ndarray.py）：

if 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.Array（IR 容器），会直接抛出AttributeError，防止误用。

2.2 tvm.nd.empty —— 只分配不初始化

tvm.nd.empty(shape, dtype="float32", device=tvm.cpu(0), mem_scope=None)

实现中先将 shape 规范化为tvm.runtime.ShapeTuple，再通过 FFI 调用底层分配：

arr = _ffi_api.TVMArrayAllocWithScope(shape, dtype, device, mem_scope)

TVMArrayAllocWithScope在 C++ 侧注册于 src/runtime/ndarray.cc：

TVM_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 支持的 dtype

NDArray 支持 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 类型（如float32x4），dtype 字符串中的x4后缀会体现在DataType.lanes字段上。

三、NDArray 核心属性与方法

NDArray的完整 API 通过autoclass指令自动收集到 ndarray.rst 中，以下逐一讲解。

3.1 属性：shape、dtype、device

arr = tvm.nd.array(np.random.rand(3, 4)) print(arr.shape) # (3, 4) print(arr.dtype) # float32 print(arr.device) # Device(0, 0) （device_type=1 表示 CPU）

• dtype属性直接读取self.handle.contents.dtype，其中handle是TVMArrayHandle（指向 DLTensor 结构）；
• device属性返回self.handle.contents.device，即 DLTensor 中记录的DLDevice结构（device_type + device_id）；
• shape则由 C 扩展层的NDArrayBase提供。

这三个属性与 DLPack 规范中的DLTensor元信息一一对应，是数组身份的核心标识。

3.2 转换到 NumPy：numpy() 与 asnumpy()

np_arr = arr.numpy() # 推荐用法 np_arr = arr.asnumpy() # 旧 API，会触发 DeprecationWarning

numpy()的实现（ndarray.py）要点：

1. 若 dtype 是向量化类型（lanes > 1），会先把lanes维度展开到 shape 末尾；
2. 对特殊类型做归一化：int4映射为int8，bfloat16映射为uint16，e4m3_float8/e5m2_float8在安装ml_dtypes时映射为对应 float8 类型，否则抛RuntimeError；
3. 分配 C 连续的 NumPy 数组后，调用TVMArrayCopyToBytes从设备内存拷贝数据；
4. int4类型还需额外做高低半字节的重新排列才能还原真实数据。

asnumpy()在 ndarray.py 中明确标注"将在 TVM v0.8 起废弃，请改用numpy()"。

3.3 数据拷贝：copyfrom 与 copyto

copyfrom（外部 → 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。

copyto（NDArray → 目标）

target_arr = tvm.nd.empty((3, 4), device=tvm.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。

切片赋值（setitem）：arr[:] = numpy_array或arr[:] = other_ndarray被重载为数据拷贝操作。注意它只支持全切片[:]，如果传入带 start/stop 的切片会抛ValueError("Array only support set from numpy array")，这是有意设计的简化约束，参见 ndarray.py。

3.4 身份比较：same_as /eq/hash

NDArray 的相等性比较是基于对象身份（底层句柄指针）而非数据内容的：

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_view

NDArray._create_view允许基于现有数组共享底层分配、创建不同逻辑 shape 的视图（ndarray.py），通过 FFI 的TVMArrayCreateView（C++ 注册见 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：

设备名称到类型掩码的映射表定义在 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, stream=None)（ndarray.py）返回封装了DLManagedTensor的PyCapsule，stream参数用于多流场景下由消费者告知生产者同步语义；
• __dlpack_device__()（ndarray.py）按 DLPack 约定返回(device_type, device_id)元组。

实现了这两个协议方法后，NDArray 即可直接被 PyTorch、JAX 等支持 DLPack 的框架零拷贝接收。

5.2 导入（其他框架 → NDArray）

torch_tensor = torch.arange(6).reshape(2, 3).cuda() tvm_arr = tvm.nd.from_dlpack(torch_tensor) # 零拷贝视图

from_dlpack（ndarray.py）接受两类输入：

1. 具有__dlpack__方法的任意对象（如 PyTorch Tensor）；
2. 已经是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 Allocator，CUDA 走 cudaMalloc 等） → 返回引用计数管理的 Container（ObjectRef）

NDArray 在 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 API（src/runtime/ndarray.cc），Python 层的copyfrom/numpy分别通过TVMArrayCopyFromBytes/TVMArrayCopyToBytes与运行时交互；
• 对齐要求：NDArray::IsAligned（ndarray.cc）要求data + byte_offset满足kAllocAlignment对齐，拷贝时CopyFromTo也会校验IsAligned（ndarray.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. 创建 NDArray（CPU） x_np = np.random.randint(0, 10, size=(3, 4)).astype("float32") x = tvm.nd.array(x_np, device=tvm.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,), dtype="float32") 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,), dtype="float32") 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.rst（automodule/autoclass/autofunction自动生成的模块级 API 页面）
• Python 实现：python/tvm/runtime/ndarray.py（NDArray 类、array/empty/from_dlpack、全部设备构造函数）
• FFI 类型定义：python/tvm/_ffi/runtime_ctypes.py（DataType、Device、DataTypeCode、TVMArray、STR2DTYPE、STR2MASK）
• C++ 头文件：include/tvm/runtime/ndarray.h（NDArray类与Container定义、Empty/CreateView/ToDLPack声明）
• C++ 实现：src/runtime/ndarray.cc（NDArray::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 推理与测试代码的第一步。

• 编译器
• 深度学习
• 模型优化