跳转至主要内容
Version: develop

空间稀疏数据结构

note

先决条件:请先阅读 FieldField(进阶版)SNode

image 图:使用粒子及网格的三维流体仿真。 从左至右:粒子、1x1x1 体素、4x4x4 块、16x16x16 块。

动机

在大规模的空间计算中,例如物理建模、图形学和三维重建,经常用到高分辨率 2D/3D 网格。 然而,如果我们使用稠密数据结构, 这些网格会消耗大量内存空间和处理能力(见 fieldfield(进阶版))。 虽然编程人员可能会分配大规模稠密网格来存储空间数据(特别是物理量,如密度或速度场),他们可能只关心这个稠密网格中的一小部分,其他的部分可能表示空的空间(如真空或空气)。

为解释这一理念,下图稀疏网格中重要的区域只会占用整个边界盒的一小部分。 如果我们能够利用这种“空间稀疏性”,将计算的重点放在我们关心的区域,我们将节省大量的存储空间和算力。

BVH

note

利用空间稀疏性的关键是将稠密网格替换为稀疏网格。

以往,稀疏数据结构基于 四叉树(2D)和 八叉树(3D)。 考虑到在现代计算机架构中,解引用指针的成本相对较高, 四叉树和八叉树对性能的友好程度不如分支因子更大的浅层树,例如 VDBSPGrid。 在 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 中的空间稀疏数据结构由 pointerbitmaskeddynamic,和 dense SNode 组成。 一个只由 dense SNode 组成的 SNode 树不是空间稀疏数据结构。

在空间稀疏数据结构中,如果一个像素、 体素或一个网格节点被分配参与计算,就认为它是活跃的。 网格其他部分就是未激活的。 用 SNode 术语解释,叶结点或中间单元是否激活用一个布尔值表示。 当且仅当一个单元活跃时,此单元的激活值为 True。 当写入未激活的单元格时,Taichi 会自动激活它。 Taichi 还提供对单元活跃度的手动操作:查看 显式操作和查询稀疏性

note

读取未激活的像素将返回 0。

指针 SNode

下面的代码片段创建了一个 8x8 大小的稀疏网格,其顶层是 4x4 的指针阵列(pointer.py 的第二行),每一个指针都指向一个 2x2 的稠密块。 就和稠密 field 一样,你可以使用索引来读写稀疏 field。 下图以绿色显示活跃的块和像素。

pointer.py
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

Pointer

执行 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。

Pointer SNode Tree

位掩码 SNode

虽然空指针可以很有效地表示空的子树,但是在叶节点使用 64 位来表示一个单独的像素的活跃状态会消耗过多的空间。

例如,如果每个像素包含单个的 f32 值(4 字节),一个 64 位指向该值的指针就需要占用 8 字节的空间。 指针的存储开销高于存储值本身所需要的空间,这种情况违背了我们使用空间稀疏数据结构来节省存储空间的目标。

为抵消指针的存储开销,你可以将像素按块排布,让指针直接指向块,就像 pointer.py 中定义的数据结构一样。

这种设计中需要注意的一点是,处于同一个 dense 块中的像素将不能再自由地改变其活跃状态。 而是拥有同一个活跃标识。 要解决这个问题,bitmasked SNode 额外为每像素分配 1 比特数据来表示每个像素的活跃状态。

下面的代码片段使用 8x8 网格解释了这个做法。 bitmasked.pypointer.py 之间唯一的区别是,位掩码 SNode 取代了稠密 SNode(第 3 行)。

bitmasked.py
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 相同。 然而,块中的位掩码像素并非全部激活,因为每个像素都有一个激活值。

Bitmasked

位掩码 SNode 就像带有辅助激活值的稠密 SNode。

Bitmasked SNode Tree

动态 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。

dynamic.py
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]

Dynamic

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

Activation

3. 失活

  • 使用 ti.deactivate(snode, [i, j, ...]) 显式让 snode[i, j, ...] 的一个单元转为不活跃。
  • 使用 snode.deactivate_all() 来取消激活 SNode snode 的所有单元。 这个操作也会递归地使其所有的子节点不再活跃。
  • 使用 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 求解器。