何为物理引擎
“A physics engine is computer software that provides an
approximate simulation of certain physical systems, such as
rigid body dynamics (including collision detection), soft body
dynamics, and fluid dynamics, of use in the domains of
computer graphics, video games and film.”
定义
- 计算机软件
- 对物理系统进行近似模拟
- 物理系统:刚体、软体、流体
“在计算机里模拟世界”
应用
- 计算机图像
- 游戏
- 电影视觉特效
- 工业软件仿真分析
- VR/AR
- 通用计算
虽然 Taichi 初衷是便利物理模拟,但借助 GPU 通用计算,Taichi 实际在多个领域得到应用。
taichimd:基于 Taichi 编程语言的 GPU 加速交互式(宏观)分子动态仿真。
TaichiSLAM:用 Taichi 语言构建 SLAM(即时定位和地图构建)的后端 3D 稠密建图库,可用于无人机集群。
Stannum:将 Taichi 与 PyTorch 结合。
What is Taichi
嵌入于 Python,使用即时编译(JIT)架构(如 LLVM、SPIR-V),将 Python 源代码转化为 GPU 或 CPU 的原生指令,在开发时和运行时均提供优越性能。
Taichi 起步于 MIT 的计算机科学与人工智能实验室(CSAIL),
设计初衷是便利计算机图形学研究人员的日常工作,
帮助他们快速实现适用于 GPU 的视觉计算和物理模拟算法。
High-performance domain-specific language embedded in Python
• Designed for computer graphics applications with productivity and
portability in mind(高效、可移植性)
• Data-oriented, parallel, megakernels(面向对象、并行、mega-kernel)
• Decouple data structures from computation (将数据结构与计算分离)
• Access spatially sparse tensors as if they are dense (提供稀疏数据结构)
• Differentiable programming support(可微编程支持)
传统深度学习框架使用的运算符都是简单的数学表达式,需要在计算图层面融合运算符,以实现更高的算法强度。但 Taichi 的指令式编程让用户能在一个 kernel 中轻松完成大量计算。 我们将这样的 kernel 命名为 mega-kernel。
安装
使用流程
1、导入Taichi
import taichi as ti
2、初始化 init
# Choose any of the following backend when initializing Taichi
# - ti.cpu
# - ti.gpu
# - ti.cuda
# - ti.vulkan
# - ti.metal
# - ti.opengl
ti.init(arch=ti.gpu)
The most useful argument: arch, i.e., the backend (architecture) to use
• ti.x64/arm/cuda/opengl/metal: stick to a certain backend.
• ti.cpu (default), automatically detects x64/arm CPUs.
• ti.gpu, try cuda/metal/opengl. If none is detected, Taichi falls back on CPUs.
3、使用数据结构定义变量
Taichi特性:
- 强类型:形参要指定数据类型,返回数据时也要指定数据类型
- 静态类型编程:
静态类型语言要求在编译时明确定义变量的类型,编译器会检查类型匹配,从而减少运行时错误。这种类型系统在C++、Java、C#等语言中广泛使用。
一旦声明了一个变量,就不能给它赋不同类型的值。 - 面向数据对象式编程(DOP):
@ti.data_oriented: 当使用的数据在 Python 作用域中活跃更新(例如当前时间和用户输入事件)并在 Taichi kernel 中追踪,建议使用面向数据的类。
@ti.dataclass: 是 ti.types.struct 的封装器,允许在数据类中定义成员函数,以实现面向对象编程(OOP)的功能:可以把 Taichi 函数定义为它的方法(成员函数),并在 Taichi 作用域内调用这些方法。【不支持继承,有更方便的做法:在ti.types.struct中__struct_methods参数指定函数】
@ti.data_oriented 示例:
定义一个 Taichi kernel 为 Python 类成员函数:
使用 @ti.data_oriented 来新修饰该类。
在面向数据的 Python 类中定义 ti.kernel 和 ti.func。
@ti.data_oriented
class TiArray:
def __init__(self, n):
self.x = ti.field(dtype=ti.i32, shape=n)
@ti.kernel
def inc(self):
for i in self.x:
self.x[i] += 1
a = TiArray(32)
a.inc()
@ti.dataclass示例:
vec3 = ti.math.vec3
@ti.dataclass
class Sphere:
center: vec3
radius: ti.f32
@ti.func
def area(self):
# a function to run in taichi scope
return 4 * math.pi * self.radius * self.radius
def is_zero_sized(self):
# a python scope function
return self.radius == 0.0
虽然通过 @ti.dataclass 定义的结构体关联函数很方便且推荐做法,但通过 __struct_methods 参数的帮助, ti.types.struct 也可以实现相同的功能。
@ti.func
def area(self):
# a function to run in taichi scope
return 4 * math.pi * self.radius * self.radius
Sphere = ti.types.struct(center=ti.math.vec3, radius=ti.f32,
__struct_methods={'area': area})
如上所述,定义结构体类型的两种方法产生的输出是相同的。
数据类型
(1)基础类型:
类别字母可以是 i 表示有符号整数, u 表示无符号整数,或 f 表示浮点数。
位精度可以是 8,16,32 或 64。
最常用的两种原始类型是:
i32: 32 位符号整型
f32 : 32 位浮点型
The support of Taichi’s primitive types by various backends may vary. Consult the following table for detailed information, and note that some backends may require extensions for complete support of a specific primitive type.
在不同后端的支持可能有所不同。请参阅下表以获取详细信息,并注意某些后端可能需要扩展(红圆圈)才能完全支持特定的原始类型。
数据类型别名:int、float
此外,在 Python 环境中,使用 ti.field 、 ti.Vector 、 ti.Matrix 、 ti.ndarray 声明 Taichi 的数据容器时,这两个名称也作为默认整数和浮点类型别名。
x = ti.field(float, 5) # Is equivalent to:
x = ti.field(ti.f64, 5)
显示类型转换
ti.cast() 函数允许您将给定的值转换为特定的目标类型。例如,您可以使用 ti.cast(x, float) 将变量 x 转换为浮点类型。
截至Taichi v1.1.0 版本,已经引入了使用原始类型如 ti.f32 和 ti.i64 对标量变量进行类型转换的能力。这使得您可以轻松地将标量变量转换为不同的标量类型。
@ti.kernel
def foo():
a = 3.14
b = ti.cast(a, ti.i32) # 3
c = ti.cast(b, ti.f32) # 3.0
x = int(a) # 3
y = float(a) # 3.14
z = ti.i32(a) # 3
w = ti.f64(a) # 3.14
隐式类型转换
隐式类型转换发生在将值放置或分配到预期数据类型不同的位置时。
二元运算中的隐式类型转换:
在Taichi中,隐式类型转换在二进制运算或赋值操作期间可能发生。
这些转换规则是为Taichi专门实现的,与 C 编程语言中的规则略有不同。这些规则的优先级如下:
1. 整数+浮点-> 浮点数
i32 + f32 -> f32
i16 + f16 -> f16
2. 低精度位 + 高精度位 -> 高精度位
i16 + i32 -> i32
f16 + f32 -> f32
u8 + u16 -> u16
3. 带符号整数 + 无符号整数 -> 无符号整数
u32 + i32 -> u32
u8 + i8 -> u8
出现规则冲突时,最高优先级的规则适用:
u8 + i16 -> i16 (当规则#2 与规则#3 冲突时,规则#2 适用。)
f16 + i32 -> f16 (当规则#1 与规则#2 冲突时,规则#1 适用。)
几个例外:
a. 位移运算返回 lhs (左侧) 数据类型:
u8 << i32 -> u8
i16 << i8 -> i16
b. 逻辑运算返回 i32。
c. 比较运算返回 i32。
(2)复合类型
复合类型是用户自定义的数据类型,由多个元素组成。
支持的复合类型包括向量、矩阵、ndarray 和结构体。
Taichi 允许将 ti.types 模块中提供的所有类型作为脚手架来自定义更高等级的复合类型。
ti.types.matrix((m,n), dataType):创建矩阵类型
ti.types.vector((…), dataType) :创建向量类型
ti.types.struct(center=dataType, radius=dataType) :创建一种结构类型,用于在三维空间中表示球体,通过其中心和半径
向量&矩阵:
vec4d = ti.types.vector(4, ti.f64) # a 64-bit floating-point 4D vector type
mat4x3i = ti.types.matrix(4, 3, int) # a 4x3 integer matrix type
Taichi 支持类型转换的复合数据类型仅限于向量和矩阵。在转换向量或矩阵的类型时,是逐元素进行的,从而生成新的向量和矩阵。
@ti.kernel
def foo():
u = ti.Vector([2.3, 4.7])
v = int(u) # ti.Vector([2, 4])
# If you are using ti.i32 as default_ip, this is equivalent to:
v = ti.cast(u, ti.i32) # ti.Vector([2, 4])
结构体:
例子-如何定义:中心为(0,0,0)半径为1.0的球体sphere1、中心为(1,1,1)半径为1.0的球体sphere2
思路:可以调用 ti.types.vector() 和 ti.types.struct() 创建两个更高级的复合类型:vec3 和 sphere_type。然后,这些类型可以用作模板来初始化两个局部变量 sphere1 和 sphere2 ,分别代表两个球体的实例。
vec3 = ti.types.vector(3, float) # Define a compound type vec3 to represent a sphere's center
sphere_type = ti.types.struct(center=vec3, radius=float) # Define a compound type sphere_type to represent a sphere
sphere1 = sphere_type(center=vec3([0, 0, 0]), radius=1.0) # Initialize sphere1, whose center is at [0,0,0] and whose radius is 1.0
sphere2 = sphere_type(center=vec3([1, 1, 1]), radius=1.0) # Initialize sphere2, whose center is at [1,1,1] and whose radius is 1.0
定义包含大量成员的结构体时,使用 ti.types.struct 可能会导致代码混乱且无序。
Taichi 提供了一个更优雅的解决方案,即 @ti.dataclass修饰器 ,它作为结构体类型的轻量级包装器使用。
@ti.dataclass
class Sphere:
center: vec3
radius: float
利用 @ti.dataclass 而不是 ti.types.struct 的另一个好处是在数据类中定义成员函数的能力,这使得可以实现面向对象编程(OOP)的功能。
(3)数据容器
标量:Field
数组:Taichi Ndarray
空间稀疏数据结构:spare
偏移坐标描述Taichi场:offset
4、处理、运算(Taichi函数)
4.1 函数修饰器 @ti.kernel Vs @ti.func
@ti.kernel
-
被@ti.kernel修饰的函数被称为Taichi kernels。在单个 Taichi 程序中可以定义多个内核。这些内核彼此独立,按照首次调用的顺序进行编译和执行。编译后的内核被缓存以减少后续调用的启动开销。
-
是Taichi运行时接管任务的入口点,必须只能从 Python 作用域中调用。
-
不允许从一个Kenel函数调用另一个Kenel函数,也不允许从@ti.func内部调用Kenel函数。
-
Tai要求Kenel函数指定参数和返回值的数据类型,除非它既没有参数也没有返回语句。
-
内核可以接受多个参数,但不能将任意的 Python 对象传递给内核,参数类型可接受:标量、 ti.types.matrix() 、 ti.types.vector() 、 ti.types.struct() 、 ti.types.ndarray() 和 ti.template() (可理解为范型T)。
-
由Python 的虚拟机接管。内核以与 Python 函数相同的方式被调用,允许在太初运行时与 Python 的虚拟机之间进行切换。
-
应用场景:可以使用原生 Python 来准备任务,例如从磁盘读取数据和预处理,然后调用内核来卸载计算密集型任务到Taichi。
@ti.kernel
def partial_sum(n: int) -> float:
...
def main():
print(partial_sum(100))
main()
@ti.func
- 被@ti.func修饰的函数被称为Taichi functions,是Kenel函数的构建块
- 只能由另一个Taichi function或Taichi kernel调用。
- 强制内联(forced-inlined):直接迁入到调用的位置,不通过压栈来跳转函数。更快而不支持递归
- 与正常的 Python 函数一样,可以将任务分解为多个太初函数以提高可读性,并在不同的内核中重用它们。
4.2 作用域 Taichi scope Vs Python scope
Taichi scope
Everything decorated with ti.kernel and ti.func.
@ti.kernel 、 @ti.func. 修饰的函数代码属于Taichi scope,
Taichi 的运行时在多核 CPU 或 GPU 设备上并行编译和执行此代码以实现高性能计算。Taichi scope相当于 CUDA 的设备端。
Python scope
Code outside the Taichi-scope.
不在Taichi scope的代码属于 Python scope,使用原生 Python 编写,并由 Python 的虚拟机执行,而不是由Taichi的运行时执行。Python scope相当于 CUDA 的主机端。
例子:可以在Python scope 用Pytorch计算,把得到的结果输入给Taichi scope中
4.3 数学函数库
具体见官方文档
示例
import taichi as ti
import taichi.math as tm
ti.init()
@ti.kernel
def test():
a = 1.0
x = tm.sin(a)
y = tm.floor(a)
z = tm.degrees(a)
w = tm.log2(a)
...
taichi.math模块提供了一些小型vector和matrix类型:
- vec2/vec3/vec4 : 二维/三维/四维浮点向量类型。
- ivec2/ivec3/ivec4 :二维/三维/四维整数向量类型。
- uvec2/uvec3/uvec4 : 二维/三维/四维无符号整数向量类型。
- mat2/mat3/mat4 : 二维/三维/四维浮点平方矩阵类型。
矩阵运算:(转置T、逆-1、迹、行列式、规范化)
4.4 For loops
for 1
for 2
for 3
并行运算时:
- 只有最外层For Loop会自动并行,里面的for2、for3在外层for1并行的线程分支里
- 如果for1外侧有if分支,就不会自动并行了(因为加了if,无法自动分配并行资源)
Range-for loops: for i in range(n) ( i : 0… n-1)
Struct-for loops: (稀疏计算)遍历所有稀疏张量里面所有active元素
4.5 原子操作
并行操作修改程序变量时,要确保使用原子操作
total [None] = total[None] + x[i]: 可能会发生一个thread读出total[None]并计算、要写回的过程中,被另一个thread修改了
支持的原子操作:
- x[i] += 1 ( +=运算:自动原子化)
- ti. atomic_add( sum, x[i]) 以原子方式执行读取-修改-写入操作。这些操作还会返回第一个参数的旧值。
x[i] = 3
y[i] = 4
z[i] = ti.atomic_add(x[i], y[i])
# now x[i] = 7, y[i] = 4, z[i] = 3
5、调试
项目文件中调试Taichi
另:提供Debug mode可以做数组越界检查
图中现在版本标量定义应为:a=ti.field(ti.i32, shape=(10))
Tips:有很多点还是需要多看官网文档:https://docs.taichi-lang.org/