Taichi的数据类型 - 老官童鞋gogo的博客

一、Taichi的架构初始化

Taichi支持多种硬件后端，初始化时通过ti.init(arch=...)指定运行平台。

1
import taichi as ti
2
ti.init(arch = ti.cuda)  # 指定CUDA后端, 用于NVIDIA GPU加速

ti.cpu：在CPU上运行。
ti.gpu：在GPU上运行，自动选择CUDA或OpenGL/Metal。
ti.cuda：显式选择CUDA后端。
ti.arm：ARM架构支持。
ti.x64：x86架构支持。

二、变量作用域：Python域与Taichi域

Taichi @ti.kernel修饰的函数与普通Python函数变量作用域是独立的。

代码示例

1
d = 1
2

3
def foo():
4
    d_python = d
5
    print("d_python:", d_python)
6

7
@ti.kernel
8
def bar():
9
    d_taichi = d
10
    print("d_taichi:", d_taichi)
11

12
d = d + 1
13
foo() # 输出2
14
bar() # 输出2
15
d = d + 1
16
foo() # 输出3
17
bar() # 输出2

Python域（foo）每次访问的都是最新的d值。
Taichi域（bar）在kernel编译时捕获d的值，后续不会自动同步Python域的变化。

三、Taichi的数据类型

Taichi支持丰富的数据类型，包括有符号/无符号整数、浮点数等。

类型类别	Taichi类型
有符号整型	`ti.i8`, `ti.i16`, `ti.i32`, `ti.i64`
无符号整型	`ti.u8`, `ti.u16`, `ti.u32`, `ti.u64`
浮点型	`ti.f16`, `ti.f32`, `ti.f64`

1、默认数据类型

整型：ti.i32
浮点型：ti.f32

可以通过如下语句修改：

1
ti.init(default_ip=ti.i32, default_fp=ti.f32)

2、类型转换规则

不同类型运算，结果类型自动提升。例如：

四、隐式类型转换

1、Python中的类型转换

1
def foo():
2
    a = 1
3
    a = 2.7
4
    print(a)   # 输出2.7
5
foo()

Python的变量类型是动态的。

2、Taichi中的类型转换

1
@ti.kernel
2
def foo():
3
    a = 1
4
    a = 2.7
5
    print(a)   # 输出2
6
foo()

原因： 在Taichi kernel中，a被推断为整型，后续赋值会将浮点部分截断。

3、强制类型转换（cast）

1
@ti.kernel
2
def foo():
3
    a = 1.7
4
    b = ti.cast(a, ti.i32)   # 强制类型转换为整型
5
    c = ti.cast(b, ti.f32)   # 再转换回浮点型
6
    print("b = ", b)         # 输出b = 1
7
    print("c = ", c)         # 输出c = 1.0
8
foo()

五、向量、矩阵、结构体类型

Taichi支持用户自定义向量、矩阵和结构体等类型。

1、类型自定义方式

使用types模块自定义新的数据类型

1
vec3f = ti.types.vector(3, ti.f32)                 # 三维浮点向量
2
mat2f = ti.types.matrix(2, 2, ti.f32)              # 2x2浮点矩阵
3
ray = ti.types.struct(ro=vec3f, rd=vec3f, l=ti.f32) # 光线结构体

2、Taichi内置类型

太极内置的向量、矩阵和结构体类型为：

ti.Vector
ti.Matrix
ti.Struct

3、示例

1
@ti.kernel
2
def foo():
3
    a = vec3f(0.0)
4
    print(a)                 # [0.0, 0.0, 0.0]
5
    b = vec3f(0.0, 1.0, 0.0)
6
    print(b)                 # [0.0, 1.0, 0.0]
7
    C = mat2f([[0.0, 1.0], [2.0, 3.0]])
8
    print(C)                 # [[0.0, 1.0], [2.0, 3.0]]
9
    r = ray(ro=a, rd=b, l=1)
10
    print(r.ro)              # [0.0, 0.0, 0.0]
11
    print(r.rd)              # [0.0, 1.0, 0.0]
12
    print("r =", r)

4、向量、矩阵下标访问

1
@ti.kernel
2
def foo():
3
    a = ti.Vector([0.0, 1.0, 0.0])
4
    print(a[0])              # 0.0
5
    print(a[1])              # 1.0
6
    print(a[2])              # 0.0
7

8
    B = ti.Matrix([[0.0, 1.0], [2.0, 3.0]])
9
    print(B[0, 0])           # 0.0
10
    print(B[0, 1])           # 1.0
11
    print(B[1, 0])           # 2.0
12
    print(B[1, 1])           # 3.0

六、Taichi的场（Field）类型

ti.field是一个全局的N维的包含了各种元素的向量

全局：可以被太极域和python域访问。
N维：可以是0维（标量）、1维（向量）、2维（矩阵）、3维，甚至更高维度。
元素：元素可以装标量、向量、矩阵、结构等等。

1、基本定义

1
heat_field = ti.field(dtype=ti.f32, shape=(256, 256))  # 256x256二维浮点场

dtype：指定元素类型。
shape：指定场的形状（维度）。

2、场的访问

1
pixels = ti.field(dtype=ti.f32, shape=(16, 8))
2
pixels[1, 2] = 42.0
3

4
vf = ti.Vector.field(3, dtype=ti.f32, shape=4)
5
@ti.kernel
6
def foo():
7
    vf[0] = ti.Vector([1.0, 2.0, 3.0])
8
    print(vf[0])             # [1.0, 2.0, 3.0]
9
foo()

3、0维场的访问方法

1
zero_d_scalar = ti.field(ti.f32, shape=())
2
zero_d_scalar[None] = 1.5     # 访问0维场需要用None

4、多维场例子：3D重力场

1
gravitational_field = ti.field(dtype=ti.f32, shape=(256, 256, 256))  # 3D场

这可以用来模拟空间中的物理量分布，比如重力场、温度场等。

老官童鞋gogo