企业官网建设流程全解析

Taichi Ndarray 完全指南：稠密数组、外部互操作与内核模板编译

【免费下载链接】taichiProductive, portable, and performant GPU programming in Python.项目地址: https://gitcode.com/GitHub_Trending/ta/taichi

Taichi ndarray 是 Taichi 提供的一种持有连续多维数据的数组对象，其角色与 NumPy 中的numpy.ndarray类似，但底层内存分配在用户通过ti.init指定的 Taichi arch（如ti.cpu、ti.cuda）上，并由 Taichi runtime 统一管理。本文以 Taichi 仓库 docs/lang/articles/basic/ndarray.md 为主线，结合python/taichi/lang/_ndarray.py、python/taichi/types/ndarray_type.py、taichi/program/ndarray.cpp等源码实现与 tests/python/test_ndarray.py 测试用例，系统讲解 ndarray 的构造、Python scope 数据交互、ti.kernel内类型注解、与 NumPy/PyTorch 外部容器的互操作，以及基于ti.types.ndarray()的模板化内核编译机制。读完本文，你将掌握在 Python scope 与 Taichi kernel 中安全、高效地使用 ndarray 的完整实践。

何时使用 ndarray：与 ti.field 的取舍

在多数场景下，你可以使用ti.field作为数据容器。但 field 可能拥有非常复杂的树形结构布局，甚至包含稀疏性（sparse）——外部库很难直接解读或使用存储在ti.field中的计算结果。而 ndarray总是分配一块连续的内存，从而能够与外部库直接进行数据交换。

用一句话概括官方文档的结论：

• fields主要用于借助复杂数据布局追求极致性能（如稀疏结构、按位打包等场景）；
• ndarray面向纯稠密数据处理，或需要与外部库（NumPy、PyTorch）互操作的需求。

从源码看 ndarray 的"连续内存"保证

在 C++ 侧，Ndarray的构造函数中通过std::accumulate将各维 shape 相乘得到总元素数nelement_，再乘以元素字节大小element_size_，最后调用prog->allocate_memory_on_device(nelement_ * element_size_, prog->result_buffer)一次性分配整块设备内存，参见 taichi/program/ndarray.cpp。Python 侧构造时还显式传入zero_fill=True，即 ndarray 默认被初始化为 0，这一点与ti.field一致（见 python/taichi/lang/_ndarray.py）。

Python scope 使用：构造、填充、读写与拷贝

ndarray只能在 Python scope 中构造，不能在 Taichi kernel 或函数内部构造，这一点与 field 相同。基本构造语句如下：

arr = ti.ndarray(dtype=ti.math.vec3, shape=(4, 4))

其中dtype可以是标量类型（如ti.f32、ti.i32），也可以是向量/矩阵类型（如ti.math.vec2、ti.math.mat2）；shape表示相对该数据类型而言的数组尺寸。除此之外，Taichi 还提供便捷构造器ti.Vector.ndarray(n, dtype, shape)与ti.Matrix.ndarray(n, m, dtype, shape)，测试用例 tests/python/test_ndarray.py 中即分别验证了这两种构造方式以及element_shape属性。

fill：标量填充

arr.fill(1.0)

从实现上看，fill并非总是直接写内存：对于 CUDA/x64 上的标量f32/i32/u32ndarray，会走prog.fill_float/fill_int/fill_uint的原生快速路径；其余情况（其他 arch、张量元素类型）则退化为调用fill_ndarray内核逐元素填充，见 python/taichi/lang/_ndarray.py 与 python/taichi/_kernels.py。

Python scope 读写元素

# 返回 ti.Vector，是元素的一份拷贝 print(arr[0, 0]) # [1.0, 1.0, 1.0] # 写入一个元素 arr[0, 0] = [1.0, 2.0, 3.0] # arr[0, 0] 现在是 [1.0, 2.0, 3.0] # 写入向量元素内部的某个标量分量 arr[0, 0][1] = 2.2 # arr[0, 0] 现在是 [1.0, 2.2, 3.0]

标量 ndarray 的__getitem__/__setitem__通过NdarrayHostAccessor完成，其底层最终调用 C++ 侧Ndarray::read/Ndarray::write：先分配一个 host staging buffer，通过memcpy_internal在设备与主机之间搬运单个元素数据，再进行 map/unmap 与std::memcpy，参见 taichi/program/ndarray.cpp。这意味着逐元素访问的开销并不低。

注意（官方文档原话）：从 Python scope 访问 ndarray 元素虽然方便，但会导致多个小型 Taichi kernel 的创建与启动，从性能角度看并非高效做法。建议将计算密集型任务放在单个 Taichi kernel 内完成，而不是在 Python scope 中逐个操作数组元素。

拷贝：shallow copy 与 deep copy

ndarray 同时支持浅拷贝与深拷贝：

# 从另一个同尺寸 ndarray 拷贝数据 b = ti.ndarray(dtype=ti.math.vec3, shape=(4, 4)) b.copy_from(arr) # 把 arr 的全部数据拷入 b import copy # 深拷贝 c = copy.deepcopy(b) # c 是一个新的 ndarray，持有 b 数据的拷贝 # 浅拷贝 d = copy.copy(b) # d 是 b 的浅拷贝；二者共享底层内存 d[0, 0][0] = 1.2 # 这会同时修改 b，因此 b[0, 0][0] 也变成 1.2

实现上，copy_from会校验源与目标 shape 完全一致（assert tuple(self.arr.shape) == tuple(other.arr.shape)），然后调用_kernels.py中的ndarray_to_ndarray内核完成设备端到设备端的拷贝；而__deepcopy__则先构造同 dtype、同 shape 的新 ndarray 再copy_from（python/taichi/lang/_ndarray.py）。深拷贝是真正的新内存，浅拷贝则共享底层DeviceAllocation。

与 NumPy 的双向数据交换

# to_numpy 返回与 d 同 shape 的 NumPy 数组，内容为 d 的拷贝 e = d.to_numpy() # from_numpy 将 NumPy 数组 e 的数据拷入 Taichi ndarray d e.fill(10.0) # 把 NumPy 数组填成 10.0 d.from_numpy(e) # 现在 d 内也全部是 10.0

to_numpy底层调用ndarray_to_ext_arr内核，from_numpy则调用ext_arr_to_ndarray内核（python/taichi/_kernels.py），二者本质上仍是"以 Taichi kernel 方式逐元素搬运"。from_numpy还会校验 shape 一致性，并且若传入的 NumPy 数组不是 C 连续布局（arr.flags.c_contiguous为 False），会自动调用np.ascontiguousarray先行转换（python/taichi/lang/_ndarray.py）。

在 ti.kernel 中使用 ndarray：类型注解与传引用语义

在 Taichi kernel 中使用 ndarray，需要在 kernel 定义中正确注解其类型，并在运行时把Ndarray对象传入。ndarray 按引用传递，因此可以在 kernel 内部修改其内容。类型注解示例如下：

@ti.kernel def foo(A: ti.types.ndarray(dtype=ti.f32, ndim=2)): do_something()

dtype 与 ndim 可选：运行时推断与校验

需要特别注意的是，dtype与ndim参数在实例化类型注解时是可选的。如果省略，数据类型与维度会从传入的数组在运行时推断；如果指定，Taichi 会校验传入数组的数据类型与维度是否匹配，不匹配即抛出错误。

这一校验逻辑在NdarrayType.check_matched中有完整实现：dtype 不一致抛TypeError/ValueError，维度不一致抛ValueError（例如required ndim=2, but 3d ndarray with shape (4, 4, 4) is provided），详见 python/taichi/types/ndarray_type.py。测试文件 tests/python/test_ndarray.py 中test_ndarray_1d、test_ndarray_2d、test_ndarray_compound_element等用例覆盖了标量、向量、矩阵元素的 kernel 内读写与往返一致性。

向量/矩阵元素：RGB 像素图示例

处理向量或矩阵元素数组（如 RGB 像素图）时，可使用向量/矩阵数据类型。下面是一个完整的 vec3 像素图处理示例（源自官方文档）：

import taichi as ti ti.init(arch=ti.cuda) arr_ty = ti.types.ndarray(dtype=ti.math.vec3, ndim=2) @ti.kernel def proc(rgb_map : arr_ty): for I in ti.grouped(rgb_map): rgb_map[I] = [0.1, 0.2, 0.3] # do something rgb = ti.ndarray(dtype=ti.types.vector(3, ti.f32), shape=(8,8)) proc(rgb)

提示：这里用arr_ty作为 2 维 vec3 ndarray 类型的别名，使类型注解更短、更易读。ti.types.vector(3, ti.f32)与ti.math.vec3是等价的向量类型。

遍历 ndarray 时，使用 range-for 还是 struct-for（如ti.grouped）都没有关系。

NdarrayType 的完整参数面

从 python/taichi/types/ndarray_type.py 可以看到，ti.types.ndarray()实际支持以下参数：

其中element_dim/element_shape若与dtype同时指定了复合类型会抛TypeError；维度数超过 2 会抛ValueError（只允许标量、向量、矩阵作为元素）。

使用外部数据容器：NumPy 与 PyTorch 零拷贝互操作

引入 Taichi ndarray 后，kernel 通过ti.types.ndarray注解，不仅能接收 Taichi ndarray，还能直接接收外部数组。当前支持的外部数组为NumPy ndarray 与 PyTorch tensor。

下面的 kernel 为数组每个元素加上1.0：

ti.init(arch=ti.cuda) @ti.kernel def add_one(arr : ti.types.ndarray(dtype=ti.f32, ndim=2)): for i in ti.grouped(arr): arr[i] = arr[i] + 1.0

外部数组无需任何额外类型转换即可喂入 kernel：

# 喂入 NumPy ndarray arr_np = np.ones((3, 3), dtype=np.float32) add_one(arr_np) # arr_np 被 taichi kernel 原地更新 # 喂入 PyTorch tensor arr_torch = torch.tensor([[1, 2, 3], [4, 5, 6], [7, 8, 9]], dtype=torch.float, device='cuda:0') add_one(arr_torch) # arr_torch 被 taichi kernel 原地更新

kernel 结束后，arr_np与arr_torch的每个元素都加上了1.0。

同设备零开销，跨设备自动搬运

• 同设备场景：当外部数据容器与 Taichi 使用同一设备时，传参不产生额外开销。例如上例中 PyTorch tensor 分配在 CUDA 上，而 Taichi 也使用 CUDA，kernel 可以直接读写 PyTorch 分配的原始 CUDA buffer，无需拷贝。
• 跨设备场景：例如 NumPy 使用 CPU、Taichi 使用 CUDA 时，Taichi 会自动管理设备间的数据传输，用户无需手动干预。

提示：NumPy 默认数据精度是 64 位，对大多数桌面 GPU 而言效率不高，建议显式指定 32 位数据类型。

连续性约束：只支持连续数组

只有连续的 NumPy 数组和 PyTorch tensor 受支持。转置等操作返回的视图（view）是非连续的，直接传入会报错：

# 转置 tensor 返回的是视图，不连续 p = arr_torch.T # add_one(p) # Error! z = p.clone() add_one(z) # 正确 k = p.contiguous() add_one(k) # 正确

标量外部数组的向量/矩阵解释

当标量类型的 NumPy ndarray 或 PyTorch tensor 作为参数传入 kernel 时，它可以被解释为标量数组、向量数组或矩阵数组——具体由类型注解中的dtype与ndim决定。例如把 shape 为(2, 2, 3, 3)的 NumPy ndarray 安全地作为mat3元素数组传入：

@ti.kernel def add_one(arr : ti.types.ndarray(dtype=ti.math.mat3, ndim=2)): for i in ti.grouped(arr): arr[i] = arr[i] + 1.0

此时(2, 2, 3, 3)被解读为"2×2 的 3×3 矩阵数组"，前两维是数组 shape，后两维是每个元素的 shape。这与NdarrayType文档字符串中的描述一致：给定dtype=ti.math.vec3时，np.zeros(10, 10, 3)会被识别为"由 vec3 元素组成的 10×10 矩阵"（见 python/taichi/types/ndarray_type.py）。从 C++ 侧看，这是通过data_type_shape(dtype)取得元素形状，再依据 AOS/SOA 布局拼接到total_shape_中实现的（taichi/program/ndarray.cpp）。

Kernel 模板化编译：ti.types.ndarray() 与 JIT 缓存

前面例子都在 kernel 类型注解中显式指定了dtype与ndim，但 Taichi 也允许完全省略，仅用ti.types.ndarray()注解。当同一个ti.kernel需要处理不同的(dtype, ndim)输入时，无需为每种组合重复定义 kernel：

@ti.kernel def test(arr: ti.types.ndarray()): for I in ti.grouped(arr): arr[I] += 2

可以把ti.types.ndarray()看作参数为dtype和ndim的模板类型。借助 JIT 编译器，带模板化 ndarray 参数的 kernel 在被调用时经历以下两步：

1. Taichi 首先检查是否已编译过相同dtype与ndim输入的 kernel：若已存在，直接加载并启动编译好的 kernel；
2. 若从未编译过，或本次调用传入的(dtype, ndim)与之前不同，则自动触发 kernel 编译，编译结果同样会被缓存供后续复用。

下面的例子演示了缓存行为：

a = ti.ndarray(dtype=ti.math.vec3, shape=(4, 4)) b = ti.ndarray(dtype=ti.math.vec3, shape=(5, 5)) c = ti.ndarray(dtype=ti.f32, shape=(4, 4)) d = ti.ndarray(dtype=ti.f32, shape=(8, 6)) e = ti.ndarray(dtype=ti.math.vec3, shape=(4, 4, 4)) test(a) # 触发新的 kernel 编译 test(b) # 复用为 a 编译的 kernel test(c) # 触发新的 kernel 编译 test(d) # 复用为 c 编译的 kernel test(e) # 触发新的 kernel 编译

可以看出：改变 shape 不会触发重新编译，改变数据类型或数组维度数才会。这一规则同样适用于来自 NumPy 或 PyTorch 的外部数组——例如 tests/python/test_ndarray.py 中test_ndarray_1d/test_ndarray_2d均以ti.types.ndarray()模板化注解运行，验证了同一 kernel 定义在不同 dtype/ndim 下的复用能力。

FAQ：ndarray 与自动微分

目前 Taichi 对 ndarray 的自动微分（autodiff）支持仍不完整，官方正在持续改进该功能，并承诺后续提供更详细的教程。如果需要使用 autodiff 与 ndarray 结合的场景，建议参考官方维护的 taichi-nerfs 示例项目中的 autodiff notebook（内容为外部链接，本文不展开）。

小结

• ndarray 是连续内存的稠密多维数组，内存分配在用户指定的 arch 上、由 Taichi runtime 管理，默认零初始化；
• Python scope 中可用fill、下标读写、copy_from、copy.deepcopy/copy.copy、to_numpy/from_numpy完成全套数据操作，但逐元素访问会启动多个小 kernel，密集计算应移入单个 kernel；
• kernel 内通过ti.types.ndarray(dtype=..., ndim=...)注解并按引用传递，dtype/ndim可省略并在运行时推断校验；
• 外部数组（NumPy、PyTorch）可直接喂入 kernel：同设备零开销、跨设备自动搬运，但只支持连续数组，标量外部数组可被解释为标量/向量/矩阵元素数组；
• ti.types.ndarray()是(dtype, ndim)上的模板类型，JIT 按(dtype, ndim)缓存编译结果，改 shape 不触发重编译。

延伸阅读：可在 tests/python/test_ndarray.py 查看 ndarray 的完整测试矩阵（覆盖ti.cpu/cuda/opengl/vulkan/metal/amdgpu六种 arch 与标量/向量/矩阵元素）；类型系统细节见 python/taichi/types/ndarray_type.py，Python 层实现见 python/taichi/lang/_ndarray.py，底层内存管理与读写逻辑见 taichi/program/ndarray.cpp。

【免费下载链接】taichiProductive, portable, and performant GPU programming in Python.项目地址: https://gitcode.com/GitHub_Trending/ta/taichi