Taichi Ndarray 完全指南稠密数组、外部互操作与内核模板编译【免费下载链接】taichiProductive, portable, and performant GPU programming in Python.项目地址: https://gitcode.com/GitHub_Trending/ta/taichiTaichi 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_fillTrue即 ndarray 默认被初始化为 0这一点与ti.field一致见 python/taichi/lang/_ndarray.py。Python scope 使用构造、填充、读写与拷贝ndarray只能在 Python scope 中构造不能在 Taichi kernel 或函数内部构造这一点与 field 相同。基本构造语句如下arr ti.ndarray(dtypeti.math.vec3, shape(4, 4))其中dtype可以是标量类型如ti.f32、ti.i32也可以是向量/矩阵类型如ti.math.vec2、ti.math.mat2shape表示相对该数据类型而言的数组尺寸。除此之外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 copyndarray 同时支持浅拷贝与深拷贝# 从另一个同尺寸 ndarray 拷贝数据 b ti.ndarray(dtypeti.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_frompython/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.0to_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(dtypeti.f32, ndim2)): do_something()dtype 与 ndim 可选运行时推断与校验需要特别注意的是dtype与ndim参数在实例化类型注解时是可选的。如果省略数据类型与维度会从传入的数组在运行时推断如果指定Taichi 会校验传入数组的数据类型与维度是否匹配不匹配即抛出错误。这一校验逻辑在NdarrayType.check_matched中有完整实现dtype 不一致抛TypeError/ValueError维度不一致抛ValueError例如required ndim2, 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(archti.cuda) arr_ty ti.types.ndarray(dtypeti.math.vec3, ndim2) 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(dtypeti.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()实际支持以下参数参数含义说明dtype元素类型PrimitiveType、VectorType、MatrixType或None运行时推断ndim数组维度数None表示运行时推断对外部数组当前暂被忽略element_dim元素维度0标量1向量2矩阵与dtype组合烹饪出向量/矩阵类型element_shape每个元素的 shape向量为 1 维 tuple矩阵为 2 维 tupleneeds_grad是否需要梯度校验时允许运行时传 needs_gradTrue 给标注为 False 的参数反之则抛错boundary边界处理默认unsafe其中element_dim/element_shape若与dtype同时指定了复合类型会抛TypeError维度数超过 2 会抛ValueError只允许标量、向量、矩阵作为元素。使用外部数据容器NumPy 与 PyTorch 零拷贝互操作引入 Taichi ndarray 后kernel 通过ti.types.ndarray注解不仅能接收 Taichi ndarray还能直接接收外部数组。当前支持的外部数组为NumPy ndarray 与 PyTorch tensor。下面的 kernel 为数组每个元素加上1.0ti.init(archti.cuda) ti.kernel def add_one(arr : ti.types.ndarray(dtypeti.f32, ndim2)): for i in ti.grouped(arr): arr[i] arr[i] 1.0外部数组无需任何额外类型转换即可喂入 kernel# 喂入 NumPy ndarray arr_np np.ones((3, 3), dtypenp.float32) add_one(arr_np) # arr_np 被 taichi kernel 原地更新 # 喂入 PyTorch tensor arr_torch torch.tensor([[1, 2, 3], [4, 5, 6], [7, 8, 9]], dtypetorch.float, devicecuda:0) add_one(arr_torch) # arr_torch 被 taichi kernel 原地更新kernel 结束后arr_np与arr_torch的每个元素都加上了1.0。同设备零开销跨设备自动搬运同设备场景当外部数据容器与 Taichi 使用同一设备时传参不产生额外开销。例如上例中 PyTorch tensor 分配在 CUDA 上而 Taichi 也使用 CUDAkernel 可以直接读写 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(dtypeti.math.mat3, ndim2)): for i in ti.grouped(arr): arr[i] arr[i] 1.0此时(2, 2, 3, 3)被解读为2×2 的 3×3 矩阵数组前两维是数组 shape后两维是每个元素的 shape。这与NdarrayType文档字符串中的描述一致给定dtypeti.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)输入时无需为每种组合重复定义 kernelti.kernel def test(arr: ti.types.ndarray()): for I in ti.grouped(arr): arr[I] 2可以把ti.types.ndarray()看作参数为dtype和ndim的模板类型。借助 JIT 编译器带模板化 ndarray 参数的 kernel 在被调用时经历以下两步Taichi 首先检查是否已编译过相同dtype与ndim输入的 kernel若已存在直接加载并启动编译好的 kernel若从未编译过或本次调用传入的(dtype, ndim)与之前不同则自动触发 kernel 编译编译结果同样会被缓存供后续复用。下面的例子演示了缓存行为a ti.ndarray(dtypeti.math.vec3, shape(4, 4)) b ti.ndarray(dtypeti.math.vec3, shape(5, 5)) c ti.ndarray(dtypeti.f32, shape(4, 4)) d ti.ndarray(dtypeti.f32, shape(8, 6)) e ti.ndarray(dtypeti.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 下的复用能力。FAQndarray 与自动微分目前 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密集计算应移入单个 kernelkernel 内通过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.pyPython 层实现见 python/taichi/lang/_ndarray.py底层内存管理与读写逻辑见 taichi/program/ndarray.cpp。【免费下载链接】taichiProductive, portable, and performant GPU programming in Python.项目地址: https://gitcode.com/GitHub_Trending/ta/taichi创作声明：本文部分内容由AI辅助生成（AIGC），仅供参考