跳转至主要内容
Version: v1.1.3

运算符

在这里,我们介绍 Taichi 中支持的基本类型和复合类型(如矩阵)的运算符。

基本类型支持的运算符

算数运算符

运算结果
-aa 取反
+aa 不变
a + bab 求和
a - bab 的差
a * bab 的乘积
a / bab 的商
a // bab 的商向下取整
a % ba / b 的余数
a ** bab 次幂
note

Taichi 的 % 操作符遵循 Python 风格而不是 C 语言风格,例如:

# In Taichi-scope or Python-scope:
print(2 % 3) # 2
print(-2 % 3) # 1

要想使用 C 语言风格的取模运算(%),请使用 ti.raw_mod。 此函数也可以接收浮点数作为参数。

ti.raw_mod(a, b) 返回 a - b * int(float(a) / b)

print(ti.raw_mod(2, 3))      # 2
print(ti.raw_mod(-2, 3)) # -2
print(ti.raw_mod(3.5, 1.5)) # 0.5
note

Python3 区分了 /(标准除法)和 //(整数除法),例如:1.0 / 2.0 = 0.51 / 2 = 0.51 // 2 = 04.2 / 2 = 2。 Taichi 遵循相同的设计:

  • 标准除法首先会将整型的操作数转换为默认的浮点类型。
  • 整数除法首先会将浮点类型的操作数转换为默认的整型。

为了避免这种隐式的类型转换,你可以使用 ti.cast 手动将操作数转换为需要的类型。 请参阅 默认精度 了解更多关于默认数值类型的详细信息。

Taichi 也提供 ti.raw_div 函数,如果其中一个操作数是浮点类型,则执行标准除法,如果两个操作数都是整数类型,则执行整数除法。

print(ti.raw_div(5, 2))    # 2
print(ti.raw_div(5, 2.0)) # 2.5

比较运算符

运算结果
a == ba 等于 b,则为 True,否则为 False
a != ba 不等于 b,则为 True,否则为 False
a > ba 严格大于 b,则为 True,否则为 False
a < ba 严格小于 b,则为 True,否则为 False
a >= ba 大于或等于 b,则为 True,否则为 False
a <= ba 小于或等于 b,则为 True,否则为 False

逻辑运算符

运算结果
not aa 为 False,则为 True,否则为 False
a or ba 为 False,则为 b 的值,否则为 a 的值
a and ba 为 False,则为 a 的值,否则为 b 的值

条件运算符

条件表达式 a if cond else b,当 cond 为 True 时,值为 a,否则值为 bab 必须具有相同的类型。

条件表达式进行短路求值,这意味着不会对未选定的分支进行求值。

a = ti.field(ti.i32, shape=(10,))
for i in range(10):
a[i] = i

@ti.kernel
def cond_expr(ind: ti.i32) -> ti.i32:
return a[ind] if ind < 10 else 0

cond_expr(3) # returns 3
cond_expr(10) # returns 0, a[10] is not evaluated

对于 Taichi 向量和矩阵的逐元素条件运算,Taichi 提供了 ti.select(cond, a, b) 函数,其中 进行短路求值。

cond = ti.Vector([1, 0])
a = ti.Vector([2, 3])
b = ti.Vector([4, 5])
ti.select(cond, a, b) # ti.Vector([2, 5])

位运算符

运算结果
~aa 逐位取反
a & bab 按位与
a ^ bab 按位异或
a | bab 按位或
a << ba 左移 b
a >> ba 右移 b
note

>> 表示 算术移位 右移(SAR)运算。 对于 逻辑移位 右移(SHR)运算,请考虑使用 ti.bit_shr()。 对于左移运算,SAL 与 SHL 是等同的。

三角函数

ti.sin(x)
ti.cos(x)
ti.tan(x)
ti.asin(x)
ti.acos(x)
ti.atan2(x, y)
ti.tanh(x)

其他数学函数

ti.sqrt(x)
ti.rsqrt(x) # A fast version for `1 / ti.sqrt(x)`.
ti.exp(x)
ti.log(x)
ti.round(x, dtype=None)
ti.floor(x, dtype=None)
ti.ceil(x, dtype=None)
ti.sum(x)
ti.max(x, y, ...)
ti.min(x, y, ...)
ti.abs(x) # Same as `abs(x)`
ti.pow(x, y) # Same as `pow(x, y)` and `x ** y`

roundfloorceil 函数中的 dtype 参数指定返回值的数据类型。 默认值 None 表示返回的类型与输入 x 相同。

内置函数

abs(x)  # Same as `ti.abs(x, y)`
pow(x, y) # Same as `ti.pow(x, y)` and `x ** y`.

随机数生成器

ti.random(dtype=float)
note

ti.random 支持 u32, i32, u64, i64 以及所有浮点型。 返回值的范围视类型而定。

类型范围
i32-2,147,483,648 至 2,147,483,647
u320 至 4,294,967,295
i64-9,223,372,036,854,775,808 至 9,223,372,036,854,775,807
u640 至 18,446,744,073,709,551,615
floating point0.0 至 1.0

支持的原子操作

在 Taichi 中,增强赋值语句(如:x[i] += 1)自动是 原子的

warning

在并行模式下修改全局变量时,请确保使用原子操作。 例如,要计算 x 所有元素的总和:

@ti.kernel
def sum():
for i in x:
# Approach 1: OK
total[None] += x[i]

# Approach 2: OK
ti.atomic_add(total[None], x[i])

# Approach 3: Wrong result since the operation is not atomic.
total[None] = total[None] + x[i]
note

当试图在局部变量上应用原子操作时,Taichi 编译器将尝试将这些操作转换为对应的非原子操作。

除了增强赋值语句以外,显式原子操作(如:ti.atomic_add)也会以原子方式执行读-改-写。 这些操作会额外返回第一个参数的 旧值。 例如,

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

下面是所有显式原子操作的列表:

运算行为
ti.atomic_add(x, y)原子地计算 x + y,将结果存储在 x 中,并返回 x 的旧值
ti.atomic_sub(x, y)原子地计算 x - y,将结果存储在 x 中,并返回 x 的旧值
ti.atomic_and(x, y)原子地计算 x & y,将结果存储在 x 中,并返回 x 的旧值
ti.atomic_or(x, y)原子地计算 x | y,将结果存储在 x 中,并返回 x 的旧值
ti.atomic_xor(x, y)原子地计算 x ^ y,将结果存储在 x 中,并返回 x 的旧值
ti.atomic_max(x, y)原子地计算 max(x, y),将结果存储在 x 中,并返回 x 的旧值
ti.atomic_min(x, y)原子地计算 min(x, y),将结果存储在 x 中,并返回 x 的旧值
note

不同后端的原子操作支持情况:

类型CPUCUDAOpenGLMetalC source
i32✔️✔️✔️✔️✔️
f32✔️✔️✔️✔️✔️
i64✔️✔️✔️
f64✔️✔️✔️

(⭕:需要扩展。)

支持的矩阵运算符

上面提到的基本类型运算也适用于复合类型(如矩阵)。 在这种情形下,运算是逐元素进行的。 For example:

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 = B ** 2
# is equivalent to
for i in ti.static(range(2)):
for j in ti.static(range(3)):
A[i, j] = B[i, j] ** 2

A = B ** C
# is equivalent to
for i in ti.static(range(2)):
for j in ti.static(range(3)):
A[i, j] = 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]

此外,Taichi 支持以下矩阵运算:

a = ti.Matrix([[2, 3], [4, 5]])
a.transpose() # the transposed matrix of `a`, will not effect the data in `a`.
a.trace() # the trace of matrix `a`, the returned scalar value can be computed as `a[0, 0] + a[1, 1] + ...`.
a.determinant() # the determinant of matrix `a`.
a.inverse() # (ti.Matrix) the inverse of matrix `a`.
a@a # @ denotes matrix multiplication
note

For now, determinant() and inverse() only works in Taichi-scope, and the size of the matrix must be 1x1, 2x2, 3x3 or 4x4.

支持的 SIMT 内建函数

对于 CUDA 后端,Taichi 现在支持编写高性能的 SIMT 内核需要的 warp 级和 block 级内建函数。 你可以在 Taichi 中以类似于 CUDA 内核中的用法 使用它们。 目前支持以下函数:

运算映射的 CUDA 内建函数
ti.simt.warp.all_nonzero__all_sync
ti.simt.warp.any_nonzero__any_sync
ti.simt.warp.unique__uni_sync
ti.simt.warp.ballot__ballot_sync
ti.simt.warp.shfl_sync_i32__shfl_sync
ti.simt.warp.shfl_sync_f32__shfl_sync
ti.simt.warp.shfl_up_i32__shfl_up_sync
ti.simt.warp.shfl_up_f32__shfl_up_sync
ti.simt.warp.shfl_down_i32__shfl_down_sync
ti.simt.warp.shfl_down_f32__shfl_down_sync
ti.simt.warp.shfl_xor_i32__shfl_xor_sync
ti.simt.warp.match_any__match_any_sync
ti.simt.warp.match_all__match_all_sync
ti.simt.warp.active_mask__activemask
ti.simt.warp.sync__syncwarp

请参考我们的 [API 文档](https://docs. taichi. graphics/api/taichi/lang/simt/warp/#module-taichi. lang. simt. warp) 了解每个函数的更多信息。

这是一个在 Taichi 的 warp 内进行数据交换的示例:

a = ti.field(dtype=ti.i32, shape=32)

@ti.kernel
def foo():
ti.loop_config(block_dim=32)
for i in range(32):
a[i] = ti.simt.warp.shfl_up_i32(ti.u32(0xFFFFFFFF), a[i], 1)

for i in range(32):
a[i] = i * i

foo()

for i in range(1, 32):
assert a[i] == (i - 1) * (i - 1)