Version: develop

运算符

在这里，我们介绍 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

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