跳转至主要内容
Version: develop

运算符

在这里,我们介绍 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✔️✔️✔️

(⭕:需要扩展。)

支持的矩阵运算符

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

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

目前,determinant() 和 inverse() 只能在 Taichi 作用域中运行,且矩阵的大小必须是 1x1、2x2、3x3 或 4x4。