特定领域语言(Domain-specific language, DSL),嵌入Python,高效并行
Taichi的高级内容还需要整理,RootDense
Taichi是嵌入Python的静态语言,所以它像Python的包一样易用,同时又可以具有很快的速度。在运行之前需要进行编译(特定的部分,这部分会被放到目标设备上运行),所以在Taichi的作用域内部需要遵守静态语言的规则。
包括kernel/内核(对应CUDA中的__global__,使用修饰器ti.kernel定义)和func/函数(对应CUDA中的__device__,使用修饰器ti.func定义),那么对应的就有,kernel可以被host端的程序调用,func可以被kernel和其他的func调用。普通的Python函数可以看作是__host__函数。在kernel和func类型函数的外部,可以不收任何约束。
在作用域最外层的for循环(不是最外层的循环结构,而是当最外层的结构是for循环时)会自动并行执行,并行的for循环不支持break。for循环分为:区间循环和结构循环,结构循环会选择性的迭代活跃元素(对于稀疏的场),这样会更快。结构循环的例子:
# 在区间 3 <= i < 8, 1 <= j < 6, 0 <= k < 9 上展开并行 for i, j, k in ti.ndrange((3, 8), (1, 6), 9):
既可以在CPU上运行也可以在GPU上运行,在GPU上运行时可以指定后端:
ti.init(arch=ti.gpu) ti.init(arch=ti.cuda)
Taichi是面向数据的,Field是第一类公民。分配一个稠密的场:
pixels = ti.field(dtype=float, shape=(n * 2, n))
要分配每一个元素都是一个 3 x 2 矩阵的形状为 128 x 64 的矩阵场:
A = ti.Matrix.field(3, 2, dtype=ti.f32, shape=(128, 64))
访问矩阵场时请使用:A[i, j][0, 1]。当从全局矩阵场加载矩阵元素时会有两个索引运算符 []:第一个用于场索引,而第二个则用于矩阵索引。
当标量运算作用于矩阵或者向量的时候,是作用于其中的左右元素,地板除的方式依然可用
在Taichi的作用域外部,可以通过一般的索引语法访问Field的元素:
pixels[42, 11] = 0.7 # 将数据存储到像素点中 print(pixels[42, 11]) # 打印 0.7
使用from_numpy和to_numpy就可以使ndarray和filed交互,对于向量场和矩阵场,其对应的ndarray的shape分别是(*field_shape, vector_n)和(*field_shape, matrix_n, matrix_m)
使用@ti.kernel修饰器来定义Taichi内核。内核如果有参数的话,则参数必须显式指定类型,最多有8个参数,(目前)仅支持标量作为参数,最多有一个标量返回值,如果有返回值的话,同样必须有类型提示,参数返回值都有类型提示,那么在传递的过程中自然就会强制转换。内核不支持嵌套。
使用@ti.func修饰器来定义Taichi函数。函数可以嵌套,(目前)不支持递归,都是强制内联的。函数的参数、返回值不需要类型提示,同时也可以传入、返回多个,也不要求是标量。(目前)函数只能有一个返回值语句(只能有一个返回点),多个返回点的情况要存在临时变量中统一最后返回。函数的参数默认都是按值传递的,强制按引用传递的方法:
def my_func(x: ti.template()):
在Taichi的作用域内,以及参数、返回值的类型说明,都请使用Taichi规定的类型。每种类型都由:一个字符指明它的类别和一个数字指明它的精度位数。数据的类别可以是:i用于有符号整数,u用于无符号整数,f用于浮点数 数据的精度位数可以是:8、16、32、64。
不同类型的数据进行运算时会发生类型提升,运算结果类型(u<i<f)和数据位数会分别提升(类似C)
变量的类型会在其被赋初值时确定(即使没有显式的指定),修改变量时数据类型并不会改变(静态语言的特性),这时的类型转换是隐式的,不会发出警告。强制类型转换使用ti.cast:
b = ti.cast(a, ti.i32)
以上的类型转换也可用于向量、矩阵,转换是逐元素的
ti.atomic_add ti.atomic_xxx
场是全局变量(或者是@data_oriented的成员变量)。有稠密稀疏之分。场的元素可以是标量、向量或矩阵。标量被看作是一个0维的场,访问场都要使用索引,像x[i, j, k]这样,如果是标量请使用x[None],场的元素默认初始值是0。
# 创建场的方式 # 最基础的一种 self.particles = Particle.field(shape=(particles_cnt,)) # Particle 是 @dataclass 或者说是 struct # 使用 ti.root.dense self.particles = Particle.field() ti.root.dense(ti.i, particles_cnt).place(self.particles) # 普通的数据类型,使用 ti.root.dense self.id_to_index = ti.field(ti.int32) ti.root.dense(ti.i, particles_cnt).place(self.id_to_index) # 使用 ti.root.dense 可以更详细地控制内存布局方式 # 具体看文档
# 可以在 ti.algorithms 里面找到 # 比如:并行计算前缀和 self.prefix_sum_executor = ti.algorithms.PrefixSumExecutor(self.grid_cnt_sum) self.prefix_sum_executor.run(self.particles_cnt_in_every_grid)