高级数据布局

张量（Tensors of scalars）可以 放置(place) 在特定的形状和 布局(layout) 中。构造适当的数据布局对性能来说非常关键，特别是对内存密集型的应用程序而言。精心设计的数据布局可以显著提高缓存/ 旁路转换缓冲(TLB)命中率和缓存行(CacheLine)利用率。不过某些情况下性能不是最优先要考虑的因素，因此你可能不需要去担心它。

在 Taichi 中，布局是以递归的方式定义。请参阅 Structural nodes (SNodes) 获得更多关于其工作方式的细节。我们建议从默认的布局规范开始（通过在 ti.var/Vector/Matrix 中指定 shape 来创建张量），如果需要的话，之后可以再使用 ti.root.X 语法迁移到更高级的布局。

Taichi 将算法与数据布局解耦，并且 Taichi 编译器可以自动优化特定数据布局上的数据访问。这些 Taichi 特性使得程序员可以快速尝试不同的数据布局，并找出针对特定任务和计算机体系结构的最有效布局。

由 shape 到 ti.root.X

例如，这里声明了一个零维张量：

  1. x = ti.var(ti.f32)
  2. ti.root.place(x)
  3. # 相当于:
  4. x = ti.var(ti.f32, shape=())

这里声明了一个尺寸为 3 的一维张量：

  1. x = ti.var(ti.f32)
  2. ti.root.dense(ti.i, 3).place(x)
  3. # 相当于:
  4. x = ti.var(ti.f32, shape=3)

这里声明了一个尺寸为 (3, 4) 的二维张量：

  1. x = ti.var(ti.f32)
  2. ti.root.dense(ti.ij, (3, 4)).place(x)
  3. # 相当于:
  4. x = ti.var(ti.f32, shape=(3, 4))

你可能会有些疑问，单纯的指定张量的 尺寸(shape) 不就行了? 为什么还要使用更为复杂的放置方式? 这是个相当好的问题，接着读下去让我们一起找出原因。

行优先 vs 列优先

让我们先从最简单的布局开始。

由于地址空间在现代计算机结构中是线性排列的，所以对于 Taichi 中的一维张量，第 i 个元素的地址就是简单的处于第 i 号位置上.

为了存储一个多维张量，必须将它扁平化(flatten)，以适应一维地址空间。例如，要存储一个大小为 (3, 2) 的二维张量，有两种方法:

  1. 第 (i, j) 位置的地址是 起始位置 + i * 2 + j （行优先）。
  2. 第 (i, j) 位置的地址是 起始位置 + j * 3 + i （列优先）。

下面是在 Taichi 中指定使用以上哪种布局的方式:

  1. ti.root.dense(ti.i, 3).dense(ti.j, 2).place(x)    # 默认行优先
  2. ti.root.dense(ti.j, 2).dense(ti.i, 3).place(y)    # 列优先

x 和 y 的形状都是 (3, 2) ，访问它们的下标都满足 0 <= i < 3 && 0 <= j < 2 的约束。当然也可以通过相同的下标访问它们: x[i, j] 和 y[i, j] 。不过它们有着非常不同的内存布局:

  1. # 存储地址低 ——————————————————————>  存储地址高
  2. # x:  x[0,0]   x[0,1]   x[0,2] | x[1,0]   x[1,1]   x[1,2]
  3. # y:  y[0,0]   y[1,0] | y[0,1]   y[1,1] | y[0,2]   y[1,2]

由此可见， x 的存储地址首先根据第一个索引下标（即行优先）增加，而 y 首先根据第二个索引下标（即列优先）增加。

注解

对于熟悉 C/C++ 的人来说，这可能看起来像是：

  1. int x[3][2];  // 行优先
  2. int y[2][3];  // 列优先
  4. for (int i = 0; i < 3; i++) {
  5.     for (int j = 0; j < 2; j++) {
  6.         do_something ( x[i][j] );
  7.         do_something ( y[j][i] );
  8.     }
  9. }

数组结构体(AoS)，结构体数组(SoA)

同样大小的张量可以被放置到一起。

例如，这里在尺寸为3的一层中放置了两个一维张量（数组结构体，AoS）:

  1. ti.root.dense(ti.i, 3).place(x, y)

他们的内存布局：

  1. # 存储地址低 ————————————————————>  存储地址高
  2. #  x[0]   y[0] | x[1]  y[1] | x[2]   y[2]

相反的，下列两个张量则被分开放置（结构体数组，SoA）:

  1. ti.root.dense(ti.i, 3).place(x)
  2. ti.root.dense(ti.i, 3).place(y)

与之对应，它们的内存布局是：

  1. # 存储地址低 ————————————————————>  存储地址高
  2. #  x[0]  x[1]   x[2] | y[0]   y[1]   y[2]

通常情况下，您不必担心不同布局之间的性能差别，可以从定义最简单的布局开始。然而，局部性(locality)有时会对性能产生重大影响，尤其是当张量很大的时。

为了改善内存访问的空间局部性(即缓存命中率/ 缓存行利用率)，有时将数据元素放置在相对较近的存储位置（如果它们经常一起被访问的话）会很有帮助。 以一个简单的一维波动方程的求解为例：

  1. N = 200000
  2. pos = ti.var(ti.f32)
  3. vel = ti.var(ti.f32)
  4. ti.root.dense(ti.i, N).place(pos)
  5. ti.root.dense(ti.i, N).place(vel)
  7. @ti.kernel
  8. def step():
  9.     pos[i] += vel[i] * dt
  10.     vel[i] += -k * pos[i] * dt

这里，我们将 pos 和 vel 分开放置。由此 pos[i] 和 vel[i] 之间地址空间的距离是 200000 。这将导致糟糕的空间局部性和大量的缓存缺失(Cache-Misses)，会在很大程度上降低性能。一个更好的放置方案是把它们放置在一起:

  1. ti.root.dense(ti.i, N).place(pos, vel)

将 vel[i] 放在 pos[i] 旁边，这样就可以提高缓存命中率，从而提高性能。

平面布局 vs 层次布局

默认情况下，当分配一个 ti.var 时，它遵循的是最简单的数据布局。

  1. val = ti.var(ti.f32, shape=(32, 64, 128))
  2. # 相当于 C++ 中的: float val[32][64][128]

但是，对于计算机图形任务而言，有些时候这种数据布局不是最理想的。 例如，val [i, j, k] 和 val [i + 1, j, k] 彼此之间的距离非常远 (32 KB)，这导致对于某些计算任务会有着低效的地址访问。 具体而言，在诸如纹理的三线性插值之类任务中，这两个元素甚至不在同一 4KB 的页面内，这将产生巨大的缓存/转址旁路缓存压力。

此时更好的布局可能是

  1. val = ti.var(ti.f32)
  2. ti.root.dense(ti.ijk, (8, 16, 32)).dense(ti.ijk, (4, 4, 4)).place(val)

这会在 4x4x4 模块中放置 val ，因此很有可能有 val [i, j, k] 及其邻居在存储上彼此靠近（即，在同一高速缓存行或内存页中） ）。

对高级稠密数据布局进行结构 for 循环

在嵌套稠密数据结构上的结构 for 循环将会自动地遵循它们在内存中的数据顺序。例如，如果二维标量张量 A 是以行为主的顺序存储的，

  1. for i, j in A:
  2.   A[i, j] += 1

将按照行优先的顺序对 A 中的元素进行遍历。 如果 A 是列优先的，则按照列优先进行遍历。

如果 A 是分层的，则迭代将在层级之间发生。 在大多数情况下，这可以最大化内存带宽利用率。

稀疏张量的结构 for 循环遵循相同的原理，这将在 Sparse computation (WIP) 中进一步讨论。

示例

二维矩阵，行优先

  1. A = ti.var(ti.f32)
  2. ti.root.dense(ti.ij, (256, 256)).place(A)

二维矩阵，列优先

  1. A = ti.var(ti.f32)
  2. ti.root.dense(ti.ji, (256, 256)).place(A) # 注意是 ti.ji 而不是ti.ij

按 8x8 的大小将 1024x1024 的二维数组分块

  1. density = ti.var(ti.f32)
  2. ti.root.dense(ti.ij, (128, 128)).dense(ti.ij, (8, 8)).place(density)

三维粒子位置和速度，数组结构体(AoS)

  1. pos = ti.Vector(3, dt=ti.f32)
  2. vel = ti.Vector(3, dt=ti.f32)
  3. ti.root.dense(ti.i, 1024).place(pos, vel)
  4. # 相当于
  5. ti.root.dense(ti.i, 1024).place(pos(0), pos(1), pos(2), vel(0), vel(1), vel(2))

三维粒子位置和速度，结构体数组(SoA)

  1. pos = ti.Vector(3, dt=ti.f32)
  2. vel = ti.Vector(3, dt=ti.f32)
  3. for i in range(3):
  4.   ti.root.dense(ti.i, 1024).place(pos(i))
  5. for i in range(3):
  6.   ti.root.dense(ti.i, 1024).place(vel(i))