空间稀疏数据结构
note
先决条件:请先阅读 Field、Field(进阶版) 和 SNode。
图:使用粒子及网格的三维流体仿真。 从左至右:粒子、1x1x1 体素、4x4x4 块、16x16x16 块。
动机
在大规模的空间计算中,例如物理建模、图形学和三维重建,经常用到高分辨率 2D/3D 网格。 然而,如果我们使用稠密数据结构, 这些网格会消耗大量内存空间和处理能力(见 field 和 field(进阶版))。 虽然编程人员可能会分配大规模稠密网格来存储空间数据(特别是物理量,如密度或速度场),他们可能只关心这个稠密网格中的一小部分,其他的部分可能表示空的空间(如真空或空气)。
为解释这一理念,下图稀疏网格中重要的区域只会占用整个边界盒的一小部分。 如果我们能够利用这种“空间稀疏性”,将计算的重点放在我们关心的区域,我们将节省大量的存储空间和算力。
note
利用空间稀疏性的关键是将稠密网格替换为稀疏网格。
以往,稀疏数据结构基于 四叉树(2D)和 八叉树(3D)。 考虑到在现代计算机架构中,解引用指针的成本相对较高, 四叉树和八叉树对性能的友好程度不如分支因子更大的浅层树,例如 VDB 和 SPGrid。 在 Taichi 中,你可以使用 SNode 组成类似于 VDB 和 SPGrid 的数据结构。 Taichi 的空间稀疏数据结构的优点包括:
- 使用索引访问,和访问稠密数据结构的方式一样。
- 迭代时自动并行。
- 自动内存访问优化。
note
后端兼容:基于 LLVM 的后端(CPU/CUDA)为在空间稀疏数据结构上进行计算提供了完整的功能。 在 Metal 后端,稀疏数据结构现在已被弃用。 对动态 SNode(dynamic SNode)的支持已在 v1.3.0 中移除,对指针(pointer)/位掩码(bitmasked)SNode 的支持将在 v1.4.0 中移除。
note
Taichi 中的稀疏矩阵通常不是通过空间稀疏数据结构来实现的。 请参阅 稀疏矩阵。
Taichi 中的空间稀疏数据结构
Taichi 中的空间稀疏数据结构由 pointer
、bitmasked
、dynamic
,和 dense
SNode 组成。 一个只由 dense
SNode 组成的 SNode 树不是空间稀疏数据结构。
在空间稀疏数据结构中,如果一个像素、 体素或一个网格节点被分配参与计算,就认为它是活跃的。 网格其他部分就是未激活的。 用 SNode 术语解释,叶结点或中间单元是否激活用一个布尔值表示。 当且仅当一个单元活跃时,此单元的激活值为 True
。 当写入未激活的单元格时,Taichi 会自动激活它。 Taichi 还提供对单元活跃度的手动操作:查看 显式操作和查询稀疏性。
note
读取未激活的像素将返回 0。
指针 SNode
下面的代码片段创建了一个 8x8 大小的稀疏网格,其顶层是 4x4 的指针阵列(pointer.py
的第二行),每一个指针都指向一个 2x2 的稠密块。 就和稠密 field 一样,你可以使用索引来读写稀疏 field。 下图以绿色显示活跃的块和像素。
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()
函数自动激活 block[1,1]
,其中包括 x[2,3]
,和 block[1,2]
,其中包括 x[2,4]
。 block[1,1]
(x[2,2], x[3,2], x[3,3]
)和 block[1,2]
(x[2,5], x[3,4], x[3,5]
)的其他像素也被隐式激活,因为稠密块中的所有像素共享相同的激活值。
事实上,稀疏 field 是如下图所示的一棵 SNode 树。 你可以使用 for
循环来遍历 SNode 树的不同层级,就像上述例子里 print_active()
函数一样。 遍历一个块的并行循环 for i, j in block
将遍历所有活跃的 pointer
SNode。 遍历一个像素的并行循环 for i, j in pixel
将遍历所有活跃的 dense
SNode。
位掩码 SNode
虽然空指针可以很有效地表示空的子树,但是在叶节点使用 64 位来表示一个单独的像素的活跃状态会消耗过多的空间。
例如,如果每个像素包含单个的 f32
值(4 字节),一个 64 位指向该值的指针就需要占用 8 字节的空间。 指针的存储开销高于存储值本身所需要的空间,这种情况违背了我们使用空间稀疏数据结构来节省存储空间的目标。
为抵消指针的存储开销,你可以将像素按块排布,让指针直接指向块,就像 pointer.py
中定义的数据结构一样。
这种设计中需要注意的一点是,处于同一个 dense
块中的像素将不能再自由地改变其活跃状态。 而是拥有同一个活跃标识。 要解决这个问题,bitmasked
SNode 额外为每像素分配 1 比特数据来表示每个像素的活跃状态。
下面的代码片段使用 8x8 网格解释了这个做法。 bitmasked.py
和 pointer.py
之间唯一的区别是,位掩码 SNode 取代了稠密 SNode(第 3 行)。
x = 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
相同。 然而,块中的位掩码像素并非全部激活,因为每个像素都有一个激活值。
位掩码 SNode 就像带有辅助激活值的稠密 SNode。
动态 SNode
为支持可变长度的 field,Taichi 提供动态 SNode。
dynamic
的第一个参数是轴,第二个参数是最大长度,第三个参数(可选)是 chunk_size
。
chunk_size
指定了动态 SNode 在先前分配的空间用尽时分配多少空间。 例如,chunk_size=4
,动态 SNode 在添加第 1 个元素时分配四个元素的空间,在添加第 5(第 9、第 13...)个元素时再分配四个元素的空间。
你可以使用 x[i].append(...)
来添加元素,使用 x[i].length()
获取长度,使用 x[i].deactivate()
清除列表。
下面的代码片段创建了一个结构体 field,储存 (i16, i64)
。 i
轴是一个稠密 SNode,j
轴是一个动态 SNode。
pair = ti.types.struct(a=ti.i16, b=ti.i64)
pair_field = pair.field()
block = ti.root.dense(ti.i, 4)
pixel = block.dynamic(ti.j, 100, chunk_size=4)
pixel.place(pair_field)
l = ti.field(ti.i32)
ti.root.dense(ti.i, 5).place(l)
@ti.kernel
def dynamic_pair():
for i in range(4):
pair_field[i].deactivate()
for j in range(i * i):
pair_field[i].append(pair(i, j + 1))
# pair_field = [[],
# [(1, 1)],
# [(2, 1), (2, 2), (2, 3), (2, 4)],
# [(3, 1), (3, 2), ... , (3, 8), (3, 9)]]
l[i] = pair_field[i].length() # l = [0, 1, 4, 9]
note
动态 SNode 必须只有一个轴,且该轴必须是最后一个轴。 动态 SNode 下不可再放置其他 SNode。 换言之,动态 SNode 必须直接放置在 field 中。 沿着从动态 SNode 到 SNode 树根节点的路径,其他 SNode 不能 具有与动态 SNode 相同的轴。
空间稀疏数据结构的计算
稀疏的结构 for 循环
在不规则的稀疏网格上进行有效的循环是非常具有挑战性的,尤其是在如 GPU 等并行设备上。 在 Taichi, for
循环具有对空间稀疏数据结构的原生支持,并且通过高效的自动并行机制,只遍历当前活跃的像素。
显式操作和查询稀疏性
Taichi 还提供可以显示操作数据结构稀疏性的 API。 你可以手动检查 SNode 的活跃度,激活 SNode,或者使 SNode 失活。 我们现在基于如下定义的 field 来说明这些功能。
x = ti.field(dtype=ti.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)
1. 活跃度检查
使用 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)
2. 激活
使用 ti.activate(snode, [i, j, ...])
显式激活 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: 1
3. 失活
- 使用
ti.deactivate(snode, [i, j, ...])
显式让snode[i, j, ...]
的一个单元转为不活跃。 - 使用
snode.deactivate_all()
来取消激活 SNodesnode
的所有单元。 这个操作也会递归地使其所有的子节点不再活跃。 - 使用
ti.disactivate_all_snodes()
来取消激活所有稀疏 SNode 的所有单元。
当取消激活发生时,Taichi 运行时将自动回收并使用 0 来填充不再活跃的容器的内存。
note
出于性能原因,ti.activate(snode, index)
仅激活 snode[index]
。 编程人员必须确保 snode[index]
的所有祖先容器已经激活。 否则,此操作会引起未定义的行为。
类似地,ti.disactivate
...
- 不会递归取消激活一个单元的所有后代。
- 不会触发取消激活其父容器,即使父容器的所有子容器都已取消激活。
4. 祖先索引查询
使用 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 行,你也可以根据给定的 pixel
索引 [7, 3]
,计算 block1
索引,即 [7//2//2, 3//2//2]
。 然而,这样做将使计算代码和数据结构的内部配置(在这种情况下,block1
容器的大小)耦合。 使用 ti.rescale_index()
,你可以避免硬编码数据结构的内部信息。
扩展阅读
请阅读这篇 SIGGRAPH Asia 2019 论文,或观看 介绍视频,同时查看 幻灯片 来了解更多关于空间稀疏数据结构的计算的信息。
Taichi elements 是一个在 Taichi 稀疏网格上实现的高性能 MLS-MPM 求解器。