如果要做并行程序,就要使用CUDA,这个原本是要用C语言来写的,但是C语言的开发稍微有点麻烦,经常出现内存报错之类的bug,如果可以使用语法更加简单的python语言来开发,就会更加快捷方便,这时可以有一个选择,就是使用taichi语言,这里记录一些零散的笔记。
简单Julia集示例
首先来看一个简单的例子,生成一个分形Julia集。
# fractal.py
#导入taichi语言
import taichi as ti
#在gpu上运行,自动选择后端,也可以是arch=ti.cuda,ti.opengl,cpu,metal
ti.init(arch=ti.gpu)
n = 320
#创建了一个二维变量数组,size为640*320,元素为浮点类型
pixels = ti.field(dtype=float, shape=(n * 2, n))
#ti.func这个修饰符表示taichi函数,只能被taichi语言域的代码调用
@ti.func
def complex_sqr(z):
return ti.Vector([z[0]**2 - z[1]**2, z[1] * z[0] * 2])
#ti.kernel这个修饰符表示taichi核函数,只能被python语言域的代码调用
#ti.kernel和ti.func类似于CUDA中的__device__和__global__
@ti.kernel
def paint(t: float):
#最外层作用域的循环,是并行执行的,内部的循环则不是,如果了解CUDA并行编程,那么可以自然而然理解这点,在遍历稀疏数组时注意这种并行循环
for i, j in pixels: # Parallized over all pixels
c = ti.Vector([-0.8, ti.cos(t) * 0.2])
z = ti.Vector([i / n - 1, j / n - 0.5]) * 2
iterations = 0
while z.norm() < 20 and iterations < 50:
z = complex_sqr(z) + c
iterations += 1
pixels[i, j] = 1 - iterations * 0.02
gui = ti.GUI("Julia Set", res=(n * 2, n))
for i in range(1000000):
paint(i * 0.03)
gui.set_image(pixels)
gui.show()
@ti.kernel
def fill():
for i in range(10): # Parallelized
x[i] += i
s = 0
for j in range(5): # Serialized in each parallel thread
s += j
y[i] = s
@ti.kernel
def fill_3d():
# Parallelized for all 3 <= i < 8, 1 <= j < 6, 0 <= k < 9
for i, j, k in ti.ndrange((3, 8), (1, 6), 9):
x[i, j, k] = i + j + k
可以很方便的给taichi数组赋值
import taichi as ti
pixels = ti.field(ti.f32, (1024, 512))
pixels[42, 11] = 0.7 # store data into pixels
print(pixels[42, 11]) # prints 0.7
还可以很方便的转换成numpy数组格式
import taichi as ti
pixels = ti.field(ti.f32, (1024, 512))
import numpy as np
arr = np.random.rand(1024, 512)
pixels.from_numpy(arr) # load numpy data into taichi fields
import matplotlib.pyplot as plt
arr = pixels.to_numpy() # store taichi data into numpy arrays
plt.imshow(arr)
plt.show()
import matplotlib.cm as cm
cmap = cm.get_cmap('magma')
gui = ti.GUI('Color map')
while gui.running:
render_pixels()
arr = pixels.to_numpy()
gui.set_image(cmap(arr))
gui.show()
最外层作用域的 for 循环是被 自动并行执行 的。Taichi 的 for 循环具有两种形式, 区间 for 循环,和 结构 for 循环。区间 for 循环 和普通的 Python for 循环没多大区别,只是 Taichi 最外层的 for 会并行执行而已。区间 for 循环可以嵌套。
@ti.kernel
def fill():
for i in range(10): # 并行执行
x[i] += i
s = 0
for j in range(5): # 在每个并行的线程中顺序执行
s += j
y[i] = s
@ti.kernel
def fill_3d():
# 在区间 3 <= i < 8, 1 <= j < 6, 0 <= k < 9 上展开并行
for i, j, k in ti.ndrange((3, 8), (1, 6), 9):
x[i, j, k] = i + j + k
核函数以及类型
taichi是一个伪装成python语言的高性能语言,为了更高的执行速度,所以跟C++一样是一个强类型语言,ti.kernel核函数的输入输出参数要写上类型,不过一般来说ti.kernel核函数不会有返回值,大量的并行数据都是通过内存引用进行赋值和传递的,在CUDA的global核函数中一般也是如此。所以实际上,taichi语言的ti.kernel核函数只能支持返回一个标量值,不支持返回并行的数组,因为这样会影响性能,也没必要。需要注意的是,taichi函数不支持递归,在做一些程序比如渲染器时,光线的多次重复反射一般用递归,这时要修改成循环。
@ti.kernel
def add_xy(x: ti.f32, y: ti.f32) -> ti.i32:
return x + y # 等价于: ti.cast(x + y, ti.i32)
res = add_xy(2.3, 1.1)
print(res) # 3,因为返回值类型是 ti.i32
也可以使用模板ti.template()来指定类型
@ti.func
def my_func(x: ti.template()):
x = x + 1 # will change the original value of x
@ti.kernel
def my_kernel():
...
x = 233
my_func(x)
print(x) # 234
...
kernel核函数不支持向量和矩阵作为参数,但是func函数可以支持
@ti.func
def sdf(u): # functions support matrices and vectors as arguments. No type-hints needed.
return u.norm() - 1
@ti.kernel
def render(d_x: ti.f32, d_y: ti.f32): # kernels do not support vector/matrix arguments yet. We have to use a workaround.
d = ti.Vector([d_x, d_y])
p = ti.Vector([0.0, 0.0])
t = sdf(p)
p += d * t
...
func函数不支持多个返回值return,所以这里需要稍微转换一下
# Bad function - two return statements
@ti.func
def safe_sqrt(x):
if x >= 0:
return ti.sqrt(x)
else:
return 0.0
# Good function - single return statement
@ti.func
def safe_sqrt(x):
ret = 0.0
if x >= 0:
ret = ti.sqrt(x)
else:
ret = 0.0
return ret
有几个简单方便的标量算术函数,当这些标量函数被作用在矩阵或向量上时,会被逐个作用到所有元素
B = ti.Matrix([[1.0, 2.0, 3.0], [4.0, 5.0, 6.0]])
C = ti.Matrix([[3.0, 4.0, 5.0], [6.0, 7.0, 8.0]])
A = ti.sin(B)
# is equivalent to
for i in ti.static(range(2)):
for j in ti.static(range(3)):
A[i, j] = ti.sin(B[i, j])
A = ti.pow(B, 2)
# is equivalent to
for i in ti.static(range(2)):
for j in ti.static(range(3)):
A[i, j] = ti.pow(B[i, j], 2)
A = ti.pow(B, C)
# is equivalent to
for i in ti.static(range(2)):
for j in ti.static(range(3)):
A[i, j] = ti.pow(B[i, j], C[i, j])
A += 2
# is equivalent to
for i in ti.static(range(2)):
for j in ti.static(range(3)):
A[i, j] += 2
A += B
# is equivalent to
for i in ti.static(range(2)):
for j in ti.static(range(3)):
A[i, j] += B[i, j]
数据位数可以是8,16,32,64位,一般常用的是32位,类型可以是
- i 用于有符号整数,例如233,-666
- u 用于无符号整数,例如233,666
- f 用于浮点数,例如2.33, 1e-4
创建field张量场,大小为12864,每个元素是32的小矩阵
A = ti.Matrix.field(3, 2, dtype=ti.f32, shape=(128, 64))
#获取小矩阵
mat = A[i, j]
#获取小矩阵元素
a0 = mat[0, 1]
a0 = A[i, j][0, 1]
在taichi中,增量赋值自动作为原子操作
@ti.kernel
def sum():
for i in x:
# 方式 1: 正确
total[None] += x[i]
# 方式 2: 正确
ti.atomic_add(total[None], x[i])
# 方式 3: 非原子操作因而会得到错误结果
total[None] = total[None] + x[i]
也可以使用命令进行显式原子操作
#这些操作会返回一个参数的旧值
ti.atomic_add(x, y)
ti.atomic_sub(x, y)
ti.atomic_and(x, y)
ti.atomic_or(x, y)
ti.atomic_xor(x, y)
将taichi数据转换成numpy矩阵的形式
@ti.kernel
def my_kernel():
for i in x:
x[i] = i * 2
x = ti.field(ti.f32, 4)
my_kernel()
x_np = x.to_numpy()
print(x_np) # np.array([0, 2, 4, 6])
x = ti.field(ti.f32, 4)
x_np = np.array([1, 7, 3, 5])
x.from_numpy(x_np)
- field.to_numpy()
参数:field可以是ti.field,ti.Vector.field或ti.Matrix.field
返回值:np.array数组 - field.from_numpy(array)
参数:field可以是ti.field,ti.Vector.field或ti.Matrix.field
array为np.array数组
转换成torch的过程和numpy的命令是类似的。在转换时需要注意到尺寸大小的对应关系
field = ti.field(ti.i32, shape=(233, 666))
field.shape # (233, 666)
array = field.to_numpy()
array.shape # (233, 666)
field.from_numpy(array) # the input array must be of shape (233, 666)
field = ti.Vector.field(3, ti.i32, shape=(233, 666))
field.shape # (233, 666)
field.n # 3
array = field.to_numpy()
array.shape # (233, 666, 3)
field.from_numpy(array) # the input array must be of shape (233, 666, 3)
field = ti.Matrix.field(3, 4, ti.i32, shape=(233, 666))
field.shape # (233, 666)
field.n # 3
field.m # 4
array = field.to_numpy()
array.shape # (233, 666, 3, 4)
field.from_numpy(array) # the input array must be of shape (233, 666, 3, 4)
# Matrix
arr = np.ones(shape=(n, m, 3, 4), dtype=np.int32)
mat.from_numpy(arr)
arr = mat.to_numpy()
assert arr.shape == (n, m, 3, 4)
标量场
taichi场的元素可以是一个标量、向量或矩阵,首先介绍标量场,field场的维度最高为8,可以是稠密或者是稀疏的,首先看稠密场的定义.
ti.field(dtype, shape = None, offset = None)
参数:
dtype – (DataType) type of the field element
shape – (optional, scalar or tuple) the shape of field
offset – (optional, scalar or tuple) see Coordinate offsets
For example, this creates a dense field with four int32 as elements:
x = ti.field(ti.i32, shape=4)
This creates a 4x3 dense field with float32 elements:
x = ti.field(ti.f32, shape=(4, 3))
If shape is () (empty tuple), then a 0-D field (scalar) is created:
x = ti.field(ti.f32, shape=())
Then access it by passing None as index:
x[None] = 2
If shape is not provided or None, the user must manually place it afterwards:
x = ti.field(ti.f32)
ti.root.dense(ti.ij, (4, 3)).place(x)
# equivalent to: x = ti.field(ti.f32, shape=(4, 3))
如果没有设置shape参数,后面必须要手动指定,这时可以按照自己的场景,自定义更加灵活的field场的布局。需要注意的是,在核函数计算之前,要定义好所有的数据变量。
x = ti.field(ti.i32, (6, 5))
x.shape # (6, 5)
x = ti.field(ti.i32, (2, 3))
x.dtype # ti.i32
# 设置field场的布局
x = ti.field(ti.i32)
y = ti.field(ti.i32)
blk1 = ti.root.dense(ti.ij, (6, 5))
blk2 = blk1.dense(ti.ij, (3, 2))
blk1.place(x)
blk2.place(y)
# 查询布局节点的层次
x.parent() # blk1
y.parent() # blk2
y.parent(2) # blk1
向量场
定义一个长度为3的全局的向量场
a = ti.Vector.field(3, dtype=ti.f32, shape=(5, 4))
x = a[6, 3][0]
# or
vec = a[6, 3]
x = vec[0]
定义一个局部向量场
a = ti.Vector([2, 3, 4])
向量场支持xyzw四轴表示
a.x
Same as a[0].
a.y
Same as a[1].
a.z
Same as a[2].
a.w
Same as a[3].
使用taichi_glsl进行表示
import taichi_glsl as tl
v = tl.vec(2, 3, 4)
print(v.xy) # [2 3]
print(v._xYzX_z) # [0 2 -3 4 -2 0 4]
有几个常用的方法函数
# 模长
a = ti.Vector([3, 4])
a.norm() # sqrt(3*3 + 4*4 + 0) = 5
# 模长平方
a.norm_sqr() # 3*3 + 4*4 = 25
# 归一化
a.normalized() # [3 / 5, 4 / 5]
# 点积
b = ti.Vector([2, 4])
a.dot(b) # 1*2 + 3*4 = 14
# 叉积
c = ti.cross(a, b)
# 外积
c = ti.outer_product(a, b) # NOTE: c[i, j] = a[i] * b[j]
# 转换类型
a = ti.Vector([1.6, 2.3])
a.cast(ti.i32) # [2, 3]
# 长度
a = ti.Vector([1, 2, 3])
a.n # 3
# 对于全局向量场,获取大小和类型
a = ti.Vector.field(3, dtype=ti.f32, shape=(4, 5))
a.shape # (4, 5)
a.dtype # ti.f32