Taichi 空间稀疏数据结构Spatially Sparse Data Structures完全指南Pointer、Bitmasked 与 Dynamic SNode 实战【免费下载链接】taichiProductive, portable, and performant GPU programming in Python.项目地址: https://gitcode.com/GitHub_Trending/ta/taichi空间稀疏数据结构Spatially Sparse Data Structures是 Taichi 面向大规模物理仿真、图形学与三维重建场景的核心能力它允许你用「按需分配」的方式承载高分辨率 2D/3D 网格仅对活跃区域付出存储与算力开销。本文以 docs/lang/articles/basic/sparse.md 为主体结合仓库中的 Python 侧实现python/taichi/lang/snode.py、python/taichi/sparse/_sparse_grid.py与示例程序python/taichi/examples/features/sparse系统讲解pointer、bitmasked、dynamic三类稀疏 SNode 的构建、遍历、显式激活/查询 API以及ti.sparse.grid()简化接口。读完本文你将掌握在 Taichi 中组合 VDB / SPGrid 风格稀疏网格的完整方法并能写出可自动并行化、自动优化内存访问的稀疏计算内核。为什么要使用稀疏数据结构空间稀疏性Motivation在大规模空间计算物理建模、图形学、三维重建中高分辨率 2D/3D 网格经常是必需的。但如果使用稠密dense数据结构这些网格会消耗大量内存空间与处理资源参见 Field 基础 与 Field 进阶布局。程序员可能分配了很大的稠密网格来存储空间数据特别是密度、速度场等物理量但真正感兴趣的往往只是其中很小一部分——其余部分可能是真空或空气等空区域。这类「感兴趣区域只占整个包围盒一小部分」的分布正是空间稀疏性spatial sparsity。利用空间稀疏性的关键是用稀疏网格替代稠密网格。这样做能显著节省存储与算力。传统上稀疏数据结构基于四叉树Quadtree2D和八叉树Octree3D。但在现代计算机架构上指针解引用dereferencing pointers相对昂贵因此四叉树/八叉树在性能上不如分支因子更大、树更浅的结构例如 VDB 和 SPGrid。在 Taichi 中可以用 SNode 组合出与 VDB、SPGrid 类似的稀疏数据结构其优势包括按下标访问访问稀疏场与访问稠密数据结构一样直观迭代自动并行化遍历稀疏数据时自动获得高效并行自动内存访问优化运行时自动优化访存模式。后端兼容性说明以仓库当前版本为准基于 LLVM 的后端CPU/CUDA提供在空间稀疏数据结构上进行计算的完整功能。在 Metal 后端上使用稀疏数据结构已被弃用Dynamic SNode 的支持已在 v1.3.0 移除Pointer/Bitmasked SNode 的支持将在 v1.4.0 移除。此外稀疏矩阵sparse matrix通常不是通过空间稀疏数据结构实现的请参见 稀疏矩阵。Taichi 中的空间稀疏数据结构构成Taichi 的空间稀疏数据结构由pointer、bitmasked、dynamic与dense四类 SNode 组合而成。注意仅由denseSNode 组成的 SNode 树并不是空间稀疏数据结构参见 python/taichi/lang/snode.py 中dense的定义。在空间稀疏数据结构上我们把一个像素pixel、体素voxel或网格节点称为活跃active当且仅当它已被分配且参与计算其余网格单元则是非活跃inactive。在 SNode 术语中叶子或中间单元的活跃性用一个布尔值表示单元活跃当且仅当其活跃值为True。当写入一个非活跃单元时Taichi 会自动将其激活Taichi 也提供手动操控活跃性的 API见下文「显式操控与查询稀疏性」一节。读取非活跃像素返回零值。这是稀疏场最重要的语义之一非活跃单元并不占用内存但读取它会得到对应数据类型的零值例如f32字段返回0.0。从 Python 侧源码看pointer/bitmasked/dynamic在创建时都会检查当前后端是否支持Extension.sparse扩展见 python/taichi/lang/snode.py不支持时直接抛出TaichiRuntimeError这从实现层面印证了文档中的后端兼容性说明。Pointer SNode按块管理空子树下面的代码创建一个 8x8 的稀疏网格顶层是 4x4 的 pointer 数组pointer.py第 2 行每个指针指向一个 2x2 的 dense 块x ti.field(ti.f32) block ti.root.pointer(ti.ij, (4,4)) pixel block.dense(ti.ij, (2,2)) pixel.place(x) ti.kernel def activate(): x[2,3] 1.0 x[2,4] 2.0 ti.kernel def print_active(): for i, j in block: print(Active block, i, j) # output: Active block 1 1 # Active block 1 2 for i, j in x: print(field x[{}, {}] {}.format(i, j, x[i, j])) # output: field x[2, 2] 0.000000 # field x[2, 3] 1.000000 # field x[3, 2] 0.000000 # field x[3, 3] 0.000000 # field x[2, 4] 2.000000 # field x[2, 5] 0.000000 # field x[3, 4] 0.000000 # field x[3, 5] 0.000000执行activate()会自动激活包含x[2,3]的block[1,1]以及包含x[2,4]的block[1,2]。由于同一 dense 块内的所有像素共享同一个活跃值block[1,1]的其余像素x[2,2]、x[3,2]、x[3,3]与block[1,2]的其余像素x[2,5]、x[3,4]、x[3,5]也被隐式激活——这也是上述输出中那些值为0.000000的单元会出现在遍历结果里的原因。这个稀疏场本质上是一棵 SNode 树。可以用for循环遍历 SNode 树的不同层级就像前面print_active()所做的那样for i, j in block并行遍历所有活跃的pointerSNodefor i, j in pixel并行遍历所有活跃的denseSNode。从 API 语义上理解pointer用「空指针」表示一整棵空子树因此一个 64 位指针即可代表「该区域完全未分配」这是它节省空间的核心机制。Bitmasked SNode为叶子单元分配 1-bit 活跃标志虽然空指针可以高效表示空子树但在叶子层用 64 位表示单个像素的活跃性会占用过多空间。例如每个像素只存一个f324 字节时指向它的 64 位指针反而要占 8 字节——指针的存储成本高于数值本身的存储成本这违背了用稀疏结构省空间的初衷。一种折中是像pointer.py那样**分块blocked**组织像素让指针直接指向块从而摊薄指针存储成本。但这样做的缺陷是同一 dense 块内的像素共享一个活跃标志无法灵活地单独改变活跃性。为了解决这个问题bitmaskedSNode 额外为每个像素分配1 位数据来表示其活跃性。下面的代码用 8x8 网格说明这一思想。bitmasked.py与pointer.py的唯一区别在于第 3 行用bitmaskedSNode 替换了denseSNodex ti.field(ti.f32) block ti.root.pointer(ti.ij, (4,4)) pixel block.bitmasked(ti.ij, (2,2)) pixel.place(x) ti.kernel def activate(): x[2,3] 1.0 x[2,4] 2.0 ti.kernel def print_active(): for i, j in block: print(Active block, i, j) for i, j in x: print(field x[{}, {}] {}.format(i, j, x[i, j]))在这个版本中活跃块与pointer.py相同但块内的 bitmasked 像素不会全部被激活因为每个像素各自拥有独立的活跃值。可以这样理解bitmasked SNode 就像带辅助活跃值的 dense SNode——它保留了按块分配内存带来的空间效率同时把「哪个像素活跃」的粒度细化到每个单元1 bit/像素。Dynamic SNode可变长度列表Taichi 自 v1.4.0 起正式支持动态数据结构Dynamic SNode。你可以把它理解为一个只能存储固定类型数据的List支持的元素类型包括标量、向量/矩阵和结构体。它提供三个 APIappend动态添加一个元素等价于 Python list 的appenddeactivate清空所有已存元素等价于 Python list 的clearlength获取当前实际存储的元素个数等价于 Python list 的__len__。这三个方法都必须在 Taichi 作用域kernel 或ti.func内调用。需要说明的是Dynamic SNode不支持像pop、remove那样动态删除单个元素——因为在并行计算中难以高效实现这类操作。使用 Dynamic SNode 时必须遵守以下规则Dynamic SNode 只能在 CPU 和 CUDA 后端使用Dynamic SNode 必须只有一个轴且该轴必须是最后一个轴Dynamic SNode 之下不能再放置其他 SNode即 Dynamic SNode 必须直接place字段从 Dynamic SNode 到 SNode 树根的路径上其他 SNode不得与 Dynamic SNode 使用相同的轴。从 Python 侧源码看dynamic接口python/taichi/lang/snode.py同样会先检查后端是否支持Extension.sparse并断言轴的数量为 1assert len(axis) 1且chunk_size缺省时取值为dimension即整段一次性分配。一维动态列表声明一个存储整数的一维动态列表xS ti.root.dynamic(ti.i, 1024, chunk_size32) x ti.field(int) S.place(x)逐行解释ti.root.dynamic表示S的直接父节点是ti.root。一般地对某个 SNodeP调用S P.dynamic()即把S在 SNode 树中的直接父节点设为P从而确定S在树中的位置dynamic的第一个参数是S所在的轴该轴必须是一维的且不能被S的任何父节点占用。这里用ti.i等价于 NumPy 的axis0第二个参数是S的最大长度。Dynamic SNode 按需动态分配内存无数据时不占空间但该最大长度有上限32 位 int 类型的最大值不能随意设为天文数字。超出最大长度也可以继续append超出范围的下标也能正常访问但建议将列表大小保持在最大长度范围内第三个参数chunk_size的含义见下文得到 Dynamic SNodeS后声明整数字段x ti.field(int)并调用S.place(x)把x变成由S描述的数据结构。调用place之前x不能存储数据调用之后x即可作为一个int类型的可变列表使用。在 kernel 内用append添加数据、用length获取实际长度ti.kernel def add_data(): for i in range(1000): x.append(i) print(x.length()) add_data()调用deactivate清空整个列表等价于恢复未初始化状态ti.kernel def clear_data(): x.deactivate() print(x.length()) # will print 0chunk_size的底层含义Dynamic SNode 的内部实现使用链表多个元素被紧凑地打包进链表的一个节点chunk中每个 chunk 包含chunk_size个元素元素分配与释放均以 chunk 为单位进行。因此实际分配的 chunk 数量为ceil(x.length() / chunk_size)。chunk_size是在内存占用chunk 数量与链表遍历开销之间做权衡的关键参数。多维动态列表数组的每个元素是变长列表可以定义更复杂的变长列表例如长度为n 10的数组x其中每个元素是一个一维变长列表S ti.root.dense(ti.i, 10).dynamic(ti.j, 1024, chunk_size32) x ti.field(int) S.place(x)这里ti.root.dense(ti.i, 10)是沿ti.i轴长度为 10 的 Dense SNode稠密数组S ti.root.dense(ti.i, 10).dynamic(ti.j, ...)表示该 Dense SNode 的子节点占据ti.j轴与父节点的轴不同。注意这正好满足前面「Dynamic SNode 的轴必须是最后一个轴、且路径上的父节点不得占用同轴」的约束。与一维情形类似可以动态操作第 i 个列表ti.kernel def add_data(): for i in range(10): for j in range(i): x[i].append(j) print(x[i].length()) # will print i for i in range(10): x[i].deactivate() print(x[i].length()) # will print 0仓库中的示例 python/taichi/examples/features/sparse/taichi_dynamic.py 展示了同一模式它对每个i追加j * j并用ti.length(x.parent(), i)与ti.append(x.parent(), i, ...)的底层接口形式操作之后在主循环中逐项断言l[i] i且x[i, j] (j * j if j i else 0)——这是对 Dynamic SNode 语义追加顺序、长度、越界读取为零的实证校验。结构体元素上述讨论同样适用于其他数值类型对向量/矩阵和结构体类型步骤完全一致。例如存储结构体的动态列表S ti.root.dynamic(ti.i, 1024, chunk_size32) SphereType ti.types.struct(centerti.math.vec3, radiusfloat) x SphereType.field() S.place(x)这里x是一维变长列表每个元素是SphereType结构体。在空间稀疏数据结构上进行计算稀疏 struct-fors只遍历活跃单元在并行设备尤其是 GPU上高效遍历分布不规则的稀疏网格单元是很有挑战性的。Taichi 的for循环原生支持空间稀疏数据结构它只遍历当前活跃的像素并自动进行高效的并行化。这意味着你无需手动维护「哪些单元活跃」的列表for i, j in x这样的循环天然只访问活跃数据。显式操控与查询稀疏性Explicitly manipulating and querying sparsity除了写入时自动激活Taichi 还提供显式 API 来检查check、**激活activate与去激活deactivate**SNode 的活跃性。以下基于一个多层稀疏字段说明x ti.field(dtypeti.i32) block1 ti.root.pointer(ti.ij, (3, 3)) block2 block1.pointer(ti.ij, (2, 2)) pixel block2.bitmasked(ti.ij, (2, 2)) pixel.place(x)这棵 SNode 树为root→pointer(3,3)→pointer(2,2)→bitmasked(2,2)→ 字段x最终覆盖一个 12x12 的逻辑网格。1. 活跃性检查Activity checking用ti.is_active(snode, [i, j, ...])显式查询snode[i, j, ...]是否活跃ti.kernel def activity_checking(snode: ti.template(), i: ti.i32, j: ti.i32): print(ti.is_active(snode, [i, j])) for i in range(3): for j in range(3): activity_checking(block1, i, j) for i in range(6): for j in range(6): activity_checking(block2, i, j) for i in range(12): for j in range(12): activity_checking(pixel, i, j)Python 侧is_active的实现会生成expr_snode_is_active前端表达式python/taichi/lang/snode.py其参数要求是pointer、hash或bitmasked节点。2. 激活Activation用ti.activate(snode, [i, j, ...])显式激活某个单元ti.kernel def activate_snodes(): ti.activate(block1, [1, 0]) ti.activate(block2, [3, 1]) ti.activate(pixel, [7, 3]) activity_checking(block1, 1, 0) # output: 1 activity_checking(block2, 3, 1) # output: 1 activity_checking(pixel, 7, 3) # output: 13. 去激活Deactivation三种去激活方式ti.deactivate(snode, [i, j, ...])显式去激活snode[i, j, ...]snode.deactivate_all()去激活snode的所有单元该操作会递归去激活其所有子节点ti.deactivate_all_snodes()去激活所有具有稀疏性的 SNode 的所有单元。去激活发生时Taichi 运行时**自动回收并清零recycle and zero-fill**被去激活容器的内存。在 Python 侧deactivate_allpython/taichi/lang/snode.py会递归遍历子树对pointer/bitmasked调用snode_deactivate对dynamic则调用snode_deactivate_dynamic——正如源码注释所解释的dynamic 节点与其它稀疏节点不同只需去激活其父节点该父节点的 chunk 链表会被整体删除即可。ti.deactivate_all_snodes()则遍历所有已 finalize 的根节点逐个执行deactivate_allpython/taichi/lang/impl.py。性能语义注意ti.activate(snode, index)出于性能考虑只激活snode[index]。程序员必须保证snode[index]的所有祖先容器已经活跃否则行为未定义undefined behaviorti.deactivate不会递归去激活一个单元的所有后代ti.deactivate不会触发父容器的去激活——即使父容器的所有子单元都已去激活父容器本身仍保持活跃。4. 祖先索引查询Ancestor index query用ti.rescale_index(descendant_snode/field, ancestor_snode, index)根据后代索引计算祖先索引print(ti.rescale_index(x, block1, ti.Vector([7, 3]))) # output: [1, 0] print(ti.rescale_index(x, block2, [7, 3])) # output: [3, 1] print(ti.rescale_index(x, pixel, [7, 3])) # output: [7, 3] print(ti.rescale_index(block2, block1, [3, 1])) # output: [1, 0]以第 1 行为例也可以手动计算给定像素索引[7, 3]block1索引为[7//2//2, 3//2//2] [1, 0]。但这种写法把计算代码与数据结构内部配置本例中block1容器的大小耦合在一起使用ti.rescale_index()则可以避免硬编码数据结构内部信息。其实现python/taichi/lang/snode.py正是按照后代与祖先各轴的 shape 比例做整除//或放大*运算。仓库示例 python/taichi/examples/features/sparse/taichi_sparse.py 中paint内核用ti.rescale_index同时求得三层 pointer 的索引并配合ti.is_active判断每层活跃性进而画出不同灰度——展示了多层稀疏结构 祖先索引查询的典型组合用法。简化 APIti.sparse.grid()与ti.sparse.usage()Taichi 还提供创建稀疏网格的简化 APIti.sparse.grid()以及用ti.sparse.usage()打印使用率内存/单元占用比例。用法如下import taichi as ti # create a 2D sparse grid grid ti.sparse.grid( { pos: ti.math.vec2, mass: ti.f32, grid2particles: ti.types.vector(20, ti.i32), }, shape(10, 10), ) # access grid[0, 0].pos ti.math.vec2(1, 2) grid[0, 0].mass 1.0 grid[0, 0].grid2particles[2] 123 # print the usage of the sparse grid, which is in [0,1] ti.sparse.usage(grid)可能输出Grid usage: 0.010000其实现位于 python/taichi/sparse/_sparse_grid.pygrid(field_dict, shape)接收一个name: type形式的字典如pos: ti.math.vec2、mass: ti.f32、grid2particles: ti.types.vector(20, ti.i32)以及 2D/3D 的 shape。它把每个元素声明为一个结构体字段并在底层将结构体放置到root.bitmasked(ij/ijk, shape)之上——即每个单元用 bitmasked 的 1-bit 活跃标志管理稀疏性python/taichi/sparse/_sparse_grid.py。shape 维度只支持 2 或 3否则抛出异常usage(x)是一个ti.kernel它通过grouped(x.parent())遍历父容器用ti.is_active统计活跃单元数除以总单元数shape[0]*shape[1]或shape[0]*shape[1]*shape[2]得到[0, 1]范围内的使用率python/taichi/sparse/_sparse_grid.py。这个简化接口特别适合快速搭建粒子-网格耦合如 MPM场景中的稀疏背景网格grid2particles这类向量字段可以记录映射到同一网格单元的粒子索引列表。更多示例与延伸阅读仓库中的完整可运行示例位于 python/taichi/examples/features/sparsetaichi_sparse.py在 CUDA 后端上用三层 pointer dense结构渲染旋转的 Taichi Logo并在每一帧调用block1.deactivate_all()后重新激活演示「每帧重建稀疏结构」的动画式用法taichi_bitmasked.pybitmasked SNode 的独立示例taichi_dynamic.pyDynamic SNode 的多列表用法与正确性断言explicit_activation.py 与 tutorial.py显式激活与入门教程。关于空间稀疏数据结构的更多原理细节可阅读 SIGGRAPH Asia 2019 关于 Taichi 语言的论文及其配套视频/幻灯片Taichi elements 项目则基于 Taichi 稀疏网格实现了一个高性能 MLS-MPM 求解器。此外动态 SNode 的append/length底层对应ti.append/ti.length接口python/taichi/lang/snode.py进阶用户可以直接使用这些底层形式。小结空间稀疏数据结构是 Taichi 处理「大而空」网格场景的核心武器。本文从稀疏性动机出发依次讲解了pointer指针按块表示空子树、bitmasked每个单元 1-bit 活跃标志与dynamic变长列表支持append/deactivate/length三类稀疏 SNode 的构建规则与使用细节覆盖了稀疏 struct-for 遍历、ti.is_active/ti.activate/ti.deactivate/snode.deactivate_all/ti.deactivate_all_snodes/ti.rescale_index等显式稀疏性 API并给出了ti.sparse.grid()简化接口与仓库内的完整示例。在 CPU/CUDA 后端上这些组合既能让你按下标直接读写又能获得自动并行化与内存访问优化是构建高分辨率物理仿真与图形学应用的重要基础。【免费下载链接】taichiProductive, portable, and performant GPU programming in Python.项目地址: https://gitcode.com/GitHub_Trending/ta/taichi创作声明：本文部分内容由AI辅助生成（AIGC），仅供参考