运算符

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

基本类型支持的运算符

算数运算符

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

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.5，1 / 2 = 0.5，1 // 2 = 0，4.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 == b	若 a 等于 b，则为 True，否则为 False
a != b	若 a 不等于 b，则为 True，否则为 False
a > b	若 a 严格大于 b，则为 True，否则为 False
a < b	若 a 严格小于 b，则为 True，否则为 False
a >= b	若 a 大于或等于 b，则为 True，否则为 False
a <= b	若 a 小于或等于 b，则为 True，否则为 False

逻辑运算符

运算	结果
not a	若 a 为 False，则为 True，否则为 False
a or b	若 a 为 False，则为 b 的值，否则为 a 的值
a and b	若 a 为 False，则为 a 的值，否则为 b 的值

条件运算符

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

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

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])

位运算符

运算	结果
~a	a 逐位取反
a & b	a 和 b 按位与
a ^ b	a 和 b 按位异或
a | b	a 和 b 按位或
a << b	a 左移 b 位
a >> b	a 右移 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`

round、floor 和 ceil 函数中的 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
u32	0 至 4,294,967,295
i64	-9,223,372,036,854,775,808 至 9,223,372,036,854,775,807
u64	0 至 18,446,744,073,709,551,615
floating point	0.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

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

类型	CPU	CUDA	OpenGL	Metal	C 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

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