基于python与CUDA的N卡GPU并行程序——taichi语言笔记（二）

基于结构节点的数据布局

自定义不同的数据结构的布局，可以提升缓存命中率，提升运行速度。有好几种Snode，如root,dense,pointer,bitmasked,dynamic,place。

声明不同大小的张量

# 零维张量
x = ti.field(ti.f32)
ti.root.place(x)
# 相当于:
x = ti.field(ti.f32, shape=())

# 一维张量
x = ti.field(ti.f32)
ti.root.dense(ti.i, 3).place(x)
# is equivalent to:
x = ti.field(ti.f32, shape=3)

# 二维张量
x = ti.field(ti.f32)
ti.root.dense(ti.ij, (3, 4)).place(x)
# is equivalent to:
x = ti.field(ti.f32, shape=(3, 4))

# 一次定义二维张量
ti.root.dense(ti.i, 3).dense(ti.j, 2).place(x)    # row-major (default)
ti.root.dense(ti.j, 2).dense(ti.i, 3).place(y)    # column-major

# 二维张量
ti.root.dense(ti.ij,(1,2)).dense(ti.ij,2).dense(ti.ij,2).place(2)

下标索引的偏移

对下标索引进行偏移之后，可以使其小于零

# -512 <= i <= 512 and -256 <= j <= 768
a = ti.field(dtype=ti.i32, shape=(1024, 1024), offset=(-512, -256))

AoS和SoA结构

同样大小的张量还可以放到一起，其中有两种防止方式，分别是AoS和SoA

# AoS
ti.root.dense(ti.i, 3).place(x, y)
# 存储地址低 ————————————————————>  存储地址高
#  x[0]   y[0] | x[1]  y[1] | x[2]   y[2]

# SoA
ti.root.dense(ti.i, 3).place(x)
ti.root.dense(ti.i, 3).place(y)
# 存储地址低 ————————————————————>  存储地址高
#  x[0]  x[1]   x[2] | y[0]   y[1]   y[2]

对于一维波动方程来说，计算位置时要用到速度，计算速度时要用到位置，所以把位置和速度放到一起，可以提升内存访问效率。

N = 200000
pos = ti.field(ti.f32)
vel = ti.field(ti.f32)
ti.root.dense(ti.i, N).place(pos, vel)
@ti.kernel
def step():
    pos[i] += vel[i] * dt
    vel[i] += -k * pos[i] * dt

层次结构

如果要计算纹理的三线性插值，val[i,j,k]和val[i+1,j,k]的距离比较远，所以更改数据布局，使相邻位置的数据在内存上也是相近的。

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

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

for循环

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

# 三维粒子的位置和速度
# AoS
pos = ti.Vector(3, dt=ti.f32)
vel = ti.Vector(3, dt=ti.f32)
ti.root.dense(ti.i, 1024).place(pos, vel)
# 相当于
ti.root.dense(ti.i, 1024).place(pos(0), pos(1), pos(2), vel(0), vel(1), vel(2))

# SoA
pos = ti.Vector(3, dt=ti.f32)
vel = ti.Vector(3, dt=ti.f32)
for i in range(3):
  ti.root.dense(ti.i, 1024).place(pos(i))
for i in range(3):
  ti.root.dense(ti.i, 1024).place(vel(i))

for循环遍历中间分层的节点

a = ti.field(dtype=ti.i32)

blk = ti.root.dense(ti.ij, 2)
blk.dense(ti.ij, (2, 4)).place(a)

@ti.kernel
def task():
    for i, j in blk:
        print(i, j)

task()

# output
0 0
0 4
2 0
2 4

稀疏结构

使用pointer Snode结构节点，定义稀疏数据，与dense稠密数据不同的是，pointer里面的元素可以是空的。

block = ti.root.pointer(ti.ij, 2).dense(ti.ij, 4)

block = ti.root.pointer(ti.ij, 2).pointer(ti.ij, 2).dense(ti.ij, 2)

当读取非激活的数据时，返回零，当写一个非激活数据时，会将其激活并进行赋值。for循环时，会自动跳过非激活元素，只循环已经激活的元素。

# 稀疏数据在循环时，只遍历已经激活的元素
for i, j in a:
    a[i, j] += 1

除了使用写数据自动激活，还可以手动指定激活过程，当将其元素去掉之后再次激活时，其元素会自动变成零。

# 激活稀疏数据的元素
ti.activate/deactivate(snode, indices)
snode.deactivate_all()

编译时求值

编译时将一部分代码进行求值计算，这样可以节省运行时的计算开销，使用ti.static可以在编译时进行判断语句的分支展开，也可以进行循环展开。

# 判断分支展开
enable_projection = True

@ti.kernel
def static():
  if ti.static(enable_projection): # 没有运行时开销
    x[0] = 1

# 循环展开
@ti.kernel
def func():
  for i in ti.static(range(4)):
      print(i)

  # 相当于:
  print(0)
  print(1)
  print(2)
  print(3)

可以直接对filed场中的向量或矩阵元素，进行循环展开以提升性能，张量的索引可以为变量，而向量或矩阵的索引在编译时必须是常量，若x是由3维向量组成的1维张量。

# 循环展开
@ti.kernel
def reset():
  for i in x:
    for j in ti.static(range(x.n)):
      # 内部循环必须被展开， 因为 j 是向量索引
      # 而不是全局张量索引
      x[i][j] = 0