Fl0w3r

carpe diem

Advanced Data Layouts

Posted on In
7.1k 13 mins.

引言

太极图形课会分成「太极」语言与计算机「图形学」两部分。前半部分主讲太极编程语言的基础语法，高级用法，以及调试和优化太极程序的方式。后半部分侧重于图形学知识，我们将以实践者的角度来聊一聊基于太极语言的渲染和仿真。

高级稠密数据结构

Before we go: packed mode

在 Taichi 的初始化 ti.init() 中可以加入 packed=True/False，它会更喜欢二的幂次方的数据诸如 2、4、6、8、16……这样的数据长度，那么如果我这样定义：

ti.init() # default: packed=False
a = ti.field(ti.i32, shape=(18, 65)) # padded to (32, 128)

它会自动补全到最近的二的幂次方，显然它浪费了一些内存，但运算效率会更高。

Taichi: optimized for data-access

x = ti.field(ti.i32, shape=16)

@ti.kernel
def fill():
    for i in x:
        x[i] = i
        
fill()

我们定义了一个 field，类型为 ti.i32。对这个 field 我们可以使用 fill() 做一个赋值操作，而 Taichi 会尽量运用并行计算硬件，所以它可能会分成四块：

这些灰色的格子记在内存中，橘色代表 prefetch 缓存中的数据：

你的程序要访问一次内存是一件比较慢的事情，通常来说我们会做一件这样的事情：

填完第一个格子的时候（比如图中有四个线程在跑并行的 for）并不满足只填的格子内容，还会预缓存一段相邻的东西。

这样我们就不用到内存中找数据，而是直接在缓存中找数据：

我们除了并行之外把内存访问的顺序和内存真实分布的数据合在一起，保持一致。

How about multi-dimensional fields?

那么问题就来了，一维的还好说，万一你的 field 是二维的呢？

x = ti.field(ti.i32, shape = (4, 4))

@ti.kernel
def fill():
    for i, j in x:
        x[i, j] = 10*i + j
        
fill()

这种情况下该如何定义我的访问顺序呢？

N-D fields are stored in our 1-D memory…

在物理内存里排布是一维的，但你的 field 可能是一维的、二维的或者三维的，那怎么存放这个东西呢？

However the access pattern is not determined…

What we want:

  • 行主序
  • 列主序
  • 块主序

Ideal memory layout of an N-D field:

我们要尽可能把我们的访问顺序和我们的内存顺序匹配起来。

问题来了，那该怎么“内存友好”地存这个 shape = (4, 4) 的 field 呢？

Access row/col-majored arrays/fields in C/C++

int x[3][2];	// row-major
int y[2][3];	// column-major

foo(){
    for (int i = 0; i < 3; i++) {
        for (int j = 0; j < 2; j++) {
            do_something(x[i][j]);
        }
    }
    
    for (int j = 0; j < 2; j++) {
        for (int i = 0; i < 3; i++) {
            do_something(y[j][i]);
        }
    }
}

如果你精通 C/C++，你可能会以行主序或列主序遍历，但这需要耗费心智。

Upgrade your ti.field()

在 Taichi 中给出的解决方法是升级你的 ti.field()，如果之前是能跑，那么现在就是要让它跑起来更加舒服。

Layout 101: from shape to ti.root

我们第一步要做的是不要再用 shape 描述 field 了，例如这里生成了一个 16*1 的 field，这个 field 每一个元素都是 3*1、类型为 ti.f32 的小向量：

x = ti.Vector.field(3, ti.f32, shape = 16)

现在我们要做的就是把这个代码翻译成下面的代码：

x = ti.Vector.field(3, ti.f32)
ti.root.dense(ti.i, 16).place(x)

它首先定义了 x 是 ti.Vector.field，但我不知道这个 field 的形状是什么样子的。它的形状由 ti.root 来描述，.dense(...).place(x) 用于 field 最后的实例化。

重新说人话就是首先 root 挂了一个描述方向和大小（也就是形状）的东西 root.dense()，而 ti.i 代表的是行的方向，在这个例子中也就是 16 行。

而 .palce(x) 则是 3*1 的 ti.Vector.field() 向量。

ti.root: more examples:

-x = ti.field(ti.f32, shape=())
+x = ti.field(ti.f32)
+ti.root.place(x)

-x = ti.field(ti.f32, shape=3)
+x = ti.field(ti.f32)
+ti.root.dense(ti.i, 3).place(x)

-x = ti.field(ti.f32, shape=(3, 4))
+x = ti.field(ti.f32)
+ti.root.dense(ti.ij, (3, 4)).place(x)

-x = ti.Matrix.field(2, 2, ti.f32, shape=5)
+x = ti.Matrix.field(2, 2, ti.f32)
+ti.root.dense(ti.i, 5).place(x)

上面是我们常用 field 的转变示例。

ti.root: the root of a SNode-tree

很显然这是一个树的形状，ti.root 为树根，上面长了 .dense() 的东西，叶节点放着 ti.field()。

The SNode-tree

这样做的好处在于我们可以更进一步拆成另外一种形式：

x = ti.field(ti.i32, shape = (4, 4))

x = ti.field(ti.i32)
ti.root.dense(ti.ij, (4, 4)).place(x)

x = ti.field(ti.i32)
ti.root.dense(ti.i, 4).dense(ti.j, 4).place(x)

Row-major v.s. column-major

现在我们就可以声明标量场，在 field 实例化时按照不同方式实例化。

Row-major access

Access row/col-majored fields

使用时就是 Taichi struct-for 就可以了，只是行主序与列主序存在区别。

Access row/col-majored arrays/fields in C/C++ v.s. in Taichi

A special case:

那如果我在 dense(ti.i) 后面再连一个 dense(ti.i) 又会怎么样呢？

Hierarchical layouts

Access a hierarchical 1-D field

import taichi as ti
ti.init(arch = ti.cpu)

x = ti.field(ti.i32)
ti.root.dense(ti.i, 4).dense(ti.i, 4).place(x)
# A hierarchical 1-D field

@ti.kernel
def print_id():
	for i in x:
		print(i, end = ' ')

print_id()

层级定义的好处在于虽然存放的时候像一个二维数组，实际访问时仍然是一维的。

Block-majored access?

同样的，这个方式也可以帮助我们用块来访问。

Block-majored access using hierarchical fields

x = ti.field(ti.i32)
ti.root.dense(ti.ij, (2,2)).dense(ti.ij, (2,2)).place(x) # block-major

Flat layouts v.s. hierarchical layouts

我们可以看到平衡的数据分布和层级的数据分布它们维度的区别也仅仅在数据定义上面，调用方式是完全一样的。

Why do we need block-majored access?

在墨戏中需要频繁访问九宫格，此时行/列主序都不是特别好的访问方法，而块主序是优于前面两者的。

Array of structures (AoS) v.s. structure of arrays (SoA) in C/C++

AoS 和 SoA 的存放能力是一样的，但是内存排布方式截然不同。

AoS v.s. SoA, which one is better?

那么它俩谁更好呢？根据访问的具体情况具体分析，做出更好的选择。

SoA in Taichi

x = ti.field(ti.i32)
y = ti.field(ti.i32)
ti.root.dense(ti.i, 8).place(x)
ti.root.dense(ti.i, 8).place(y)

这样 Taichi 就会先存 x 再存 y，得到最后的内存排布。

AoS in Taichi

x = ti.field(ti.i32)
y = ti.field(ti.i32)
ti.root.dense(ti.i, 8).place(x, y)

如果转成 AoS，只需把它们 place() 到相同的 block。

x = ti.field(ti.i32)
y = ti.field(ti.i32)
ti.root.dense(ti.i, 8).place(x)
ti.root.dense(ti.i, 16).place(y)

值得一提的是只有相同形状的 field 才可以被 place() 到 AoS 中去。

x = ti.field(ti.i32)
y = ti.Vector.field(2, ti.i32)
ti.root.dense(ti.i, 8).dense(ti.j, 8).place(x, y)

就算 x 是标量场，y 是矢量场也没有关系，我们仍旧可以 place() 到同一个 array 中去。

Switching between AoS and SoA in Taichi

Loop over advanced data layouts

高级数据结构的访问方式没有任何变化，只需要更改数据的组织结构。

Before moving to the next topic…

Generates advanced dense data layouts using ti.root

  • Tree-structured: SNode-trees.

    • The SNode stands for “Structural Nodes”

  • All fields in Taichi are built using SNode-trees

    • ti.root is actuall “the root” of an SNode-tree
    • x=ti.field(ti.f32, shape=N) <==> x=ti.field(ti.f32) + ti.root.dense(ti.i, N).place(x)
    • ti.root.dense(ti.ij, (N, M)) <==> ti.root.dense(ti.i, N).dense(ti.j, M)

  • You can append (multiple) dense cells to other dense cells

    • Row/col-major: ti.root.dense(ti.i,N).dense(ti.j,M)
    • Hierarchical layouts: ti.root.dense(ti.i, N).dense(ti.i, M)
    • SoA/AoS: ti.root.dense(ti.i,N).place(x,y,z)

  • You do not need to worry about the access of your data layouts

    • The Taichi struct-for handles it for you

高级稀疏数据结构

The SNode-tree

Sparse computation! but why?

MPM 模拟占用了 256x256 个像素，很多空白并不需要分配空间。我们可以把这个空间先分割成 16x16 的空间，然后在这个 16x16 的空间再分割出 16x16 的空间。

Sparse computation! Then how?

x = ti.field(ti.i32)
block1 = ti.root.dense(ti.i, 3)
block2 = block1.root.dense(ti.j, 3)
block2.place(x)
# equivalent to ti.root.pointer(ti.i,3).dense(ti.j,3).place(x)

比如这里定义了一个 ti.field()，它先挂了一个 i，再挂了一个 j，可以等同于：

ti.root.pointer(ti.i, 3).dense(ti.j, 3).place(x)

但很有可能空间利用率会非常低，比如这些空间中大部分都是 0。那么这些 Dense 中大部分都是没有用的：

From .dense() to .pointer()

在这种情况下我们可以换一个定义方式。我们不再定义一个稠密的数据结构来连接整个空间，而是定义一个指针 Pointer，这个指针会告诉我们下面稠密的空间在什么地方。

定义指针的好处在于如果指针发现下面的 Dense Block 都没有被激活，那么可以直接把一整块空间全部释放。定义的修改也非常简单：

x = ti.field(ti.i32)

block1 = ti.root.pinter(ti.i, 3)
block2 = block1.root.dense(ti.j, 3)
block2.place(x)
# equivalent to ti.root.pointer(ti.i,3).pointer(ti.j,3).place(x)

我们把一个稠密的数据结构改为一个稀疏的数据结构就能达到目的。

Activation

x = ti.field(ti.i32)

block1 = ti.root.pointer(ti.i, 3)
block2 = block1.dense(ti.j, 3)
block2.place(x)
# equivalent to ti.root.pointer(ti.i,3) .dense(ti.j,3).place(x)

在这种情况下这颗树初始化时不会分配任何内存，所有的 Pointer 都是 nullptr 的，它们指向一个零空间，并不会把稠密空间分配给这些 Pointer。

x[0,0] = 1
# activates block1[0]

一旦将一个空的地方写入数据，我们就会激活它的 Parent Pointer。

x[0,0] = 1
# activates block1[0] and thereby block2[0], block2[1] and block2[2]

激活父节点后会把这个节点下面所有的 Dense 都激活了，因为实现的时候只是用了一个指针，这个指针指向数据的头，它们之间的内存分布是连续的。

Data access in a sparse field (a sparse SNode-tree)

@ti.kernel
def access_all():
    for i,j in x:
        print(x[i, j]) # 1, 0, 0
        
print(x[2, 2]) # 0

我们依然可以通过 struct-for 完成对稀疏数据结构的访问，如果发现指针没有被激活会自动跳过访问，在刚才的代码中只会打印有值的 1，0，0。

如果手动强制访问，它会跳过访问直接告诉你为 0，因为父节点未激活。

Why not using pointer everywhere?

如果下面一层也定义为 pointer 也是可以的，这种情况下看似节约了空间，其实反而是浪费了。因为 Pointer 所指的 field 很可能仅仅是一个标量，而指针的开销不可忽视，并且访问不再具有空间连续性，访问的也就更慢。

Why activating x[0, 1] and x[0, 2] as well?

Use bitmasks if you really want to flag leaf cells one at a time…

x = ti.field(ti.i32)

block1 = ti.root.pointer(ti.i, 3)
block2 = block1.bitmasked(ti.j, 3)
block2.place(x)
# equivalent to ti.root.pointer(ti.i,3) .bitmasked(ti.j,3).place(x)

这样我们依旧可以让每一个 field 拥有一个独立的 activation flag。

@ti.kernel
def access_all():
    for i,j in x:
        print(x[i, j]) # 1

当用 struct-for 调用时不会访问 bitmask 为 0 的空间，所以只会打印出 1。

Manual sparse field manipulation

Putting things together

Previous section:

  • Row-major v.s. col-major, flat v.s. hierarchical layouts

This section:

  • .dense() v.s. .pointer()/.bitmasked()

A rolling Taichi [Code]

Sparse data layouts

Append more types to your SNode-tree:

  • .pointer() to represent sparse cells
  • .bitmasked() to represent sparse leaf cells

Activate cells (and its ancestors) by writing

  • x[0,0] = 1

Use Taichi struct-for(s) to access sparse fields

  • as if they were dense☺
Related Posts