taichi:嵌入在Python中的DSL(动态脚本语言),其编译器将被 @ti.kernel 装饰的函数编译到硬件上,包括CPU和GPU,然后进行高性能计算;kernel包裹的函数参数需要指定类型,如果有返回也要指定返回类型。
安装:pip install taichi -i https://pypi.douban.com/simple,不要使用带-nightly的包
Taichi 是一种面向数据的编程语言,其中密集或空间稀疏的字段是一等公民。
pixels = ti.field(dtype=float, shape=(n * 2, n))分配了一个名为pixels,size为(640, 320)和 element data type的二维密集字段float。
计算发生在 Taichi 的内核和函数中;Taichi的内核由装饰器@ti.kernel来定义,从 Python 中调用 kernel 来进行计算, 如果 kernel 有参数的话,必须指定参数类型。
Taichi 函数由装饰器 @ti.func 定义。 只能从 Taichi kernel 或其他 Taichi 函数中调用。
Taichi 的 kernel 与函数所用语法与 Python 语法看起来全然一致,Taichi 前端编译器会将其转换为编译型、静态类型、有词法作用域、并行执行且可微分的语言。
Taichi 作用域与 Python 作用域:
任何被 @ti.kernel 和 @ti.func 修饰的部分都处于 Taichi 作用域中,由 Taichi 编译器编译。
其余部分处于 Python 作用域中,是 Python 原生代码。
Taichi 自动并行执行 kernel 里位于最外层作用域的 for 循环。 有两种形式的 for 循环:
Range-for 循环与 Python for 循环没有什么不同,只是在最外层范围内使用时是并行化的。
Range-for 循环可以嵌套,ti.loop_config(serialize=True) 禁用自动的外层循环并行运算。
被并行化的是最外层范围(作用域)的循环,而不是最外层的循环,并行循环不支持 break。
struct-for 循环在遍历(稀疏)field 中的元素时尤其有用;这种循环只会遍历稀疏 field 中的活跃元素。 在稠密 field 中,所有元素都是活跃元素,结构 for 循环只能使用在内核的最外层作用域。
kernel 是标志 Taichi 的运行时接管的入口,也是运行时执行的最小单位。
一个程序中可定义多个 kernel,kernel 之间仙湖独立。 调用 kernel 的方式与调用 Python 函数相同,可以在 Taichi 的运行时和 Python 的虚拟机之间来回切换。Taichi 的运行时按照调用顺序编译并执行 kernel。将已编译的 kernel 存入缓存,下一次调用到同一 kernel 时就不需要再重新编译。
不能从一个 kernel 中调用另一个 kernel,或是从 Taichi 函数中调用 kernel。
只能直接调用或从 Python 原生函数中调用 kernel。只能从Python 作用域调用 kernel,kernel 相当于 CUDA 中的__global__
函数。
一个 kernel 可以接收多个参数,支持标量、ti.Matrix
,以及 ti.Vector
作为参数类型。从 Python 作用域传递值到 Taichi 作用域变得更轻松、灵活。
向标量、ti.Matrix
, 及 ti.Vector
中传递参数采用值传递的方式,所以改变 kernel 中的参数不会影响调用者函数中的原始变量。
定义参数时须遵循如下规则:
输入 kernel 参数的类型提示。
确保 kernel 参数中的元素总数不超过某一上限
元素数量的上限与后端相关:
OpenGL 后端支持 8 个元素
CPU、Vulkan、CUDA,和 Metal 均支持 64 个元素
标量参数中的元素数量总是 1。
ti.Matrix 或 ti.Vector 中的元素数量是其所含标量的实际个数。
Taichi 函数的返回值可以是标量、ti.Matrix、ti.Vector、ti.Struct等。
注意:
不同于 kernel,Taichi 函数可以有多个返回值。
无需(但仍然推荐)给出 Taichi 函数返回值的类型提示。
返回值中的元素数量没有限制。
Taichi 函数中不能包含超过一条的 return
语句。
Taichi kernel与函数对比:
kernel 是 Taichi 运行时执行的最小单位。 不能从 Taichi 函数内(即 Taichi 作用域内)调用 kernel,只能从 Python 作用域调用 kernel。
初始化 Taichi 时通过明确指定默认类型更改默认基本类型:
ti.init(default_ip=ti.i64) # the default integer type to ti.i64
ti.init(default_fp=ti.f64) # Set default floating-point type to ti.f64
读取 0D 向量 field
x = ti.Vector.field(n=3, dtype=ti.f32, shape=())
第p个组件:x[None][p]
(0 ≤ p < n).
参考:
胡渊鸣:import一个“太极”库,让Python代码提速100倍!
Accelerate Python code 100x by import taichi as ti | Taichi Docs