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

kernel 与函数

Embedded in Python, Taichi resembles Python in language syntax. To differentiate Taichi code from native Python code, we use the two decorators @ti.kernel and @ti.func:

  • Functions decorated with @ti.kernel are called Taichi kernels (or kernels for short). They serve as the entry points where Taichi begins to take over the tasks, and they must be called directly by Python code.
  • @ti.func 装饰的函数被称为 Taichi 函数。 They serve as the building blocks of kernels and can only be called by kernels or other Taichi functions.

Let's see an example:

import taichi as ti
ti.init(arch=ti.cpu)

@ti.func
def inv_square(x): # a Taichi function
return 1.0 / (x * x)

@ti.kernel
def partial_sum(n: int) -> float: # a kernel
total = 0.0
for i in range(1, n + 1):
total += inv_square(n)
return total

In the code above, inv_square() is a Taichi function because it is decorated by @ti.func, while partial_sum() is a kernel because it is decorated by @ti.kernel. The Taichi function (former) is called by the kernel (latter).

You may have noticed that the argument and the return value in the kernel partial_sum() are both type-hinted, but those in the Taichi function inv_square() are not. Here comes an important difference between Python and Taichi. In native Python code, type hinting is recommended but not mandatory. But Taichi makes it compulsory that kernels must take type-hinted arguments and return type-hinted values. The only exception where you can leave out a type hint in a kernel is that the kernel does not have an argument or a return statement.

It is worth your attention that Taichi will raise a syntax error if you try to call inv_square() directly from the native Python code (i.e., out of the Taichi scope). For example:

import taichi as ti
ti.init(arch=ti.cpu)

@ti.func
def inv_square(x):
return 1.0 / (x * x)

print(inv_square(1.0)) # Syntax error!

The Taichi function should have fallen in the Taichi scope, a concept as opposed to the "Python scope".

IMPORTANT

We give the following definitions:

  1. The code inside a kernel or a Taichi function is in the Taichi scope. Taichi 作用域内的代码由 Taichi 的运行时编译,并在 CPU 或 GPU 设备上并行执行,以实现高性能计算。

    The Taichi scope corresponds to the device side in CUDA.

  2. Code outside of the Taichi scope is in the Python scope. The code in the Python scope is native Python and executed by Python's virtual machine, not by Taichi's runtime.

    The Python scope corresponds to the host side in CUDA.

We should not confuse kernels with Taichi functions. Though they belong to the Taichi scope, the syntax rules applied to them are not exactly the same. We now dive into their usages and the roles they play in detail.

kernel

As the smallest execution unit in Taichi, a kernel is the entry point from which Taichi's runtime takes control. You call a kernel the same way you call a Python function and can switch back and forth between Taichi's runtime and Python's virtual machine.

For example, you can call the kernel partial_sum() as defined in the above section from inside a Python function:

def main():
print(partial_sum(100))
print(partial_sum(1000))

main()

You can define multiple kernels in your program. They are mutually independent of each other and are compiled and executed in the same order as they are first called. The compiled kernels are stored in the cache to save the launch overhead for subsequent calls.

WARNING

You must not call a kernel from inside another kernel or from inside a Taichi function. 只能直接调用或从 Python 原生函数中调用 kernel。 In other words, you can only call a kernel from inside the Python scope.

参数

A kernel can take multiple arguments. However, you cannot pass any arbitrary Python object to a kernel because Python objects can be highly dynamic and may hold data unrecognized by Taichi's compiler.

The argument types accepted by kernels are scalars, ti.Matrix/ti.Vector (In Taichi, vectors are essentially matrices), ti.types.ndarray() and ti.template(). You can easily pass data from the Python scope to the Taichi scope.

It should be noted that scalars and ti.Matrix are passed by value, while ti.types.ndarray() and ti.template() are passed by reference. In the latter case, any modification to the arguments in the called function also affects the original values.

In the following example, the arguments x and y are passed to my_kernel by value:

@ti.kernel
def my_kernel(x: int, y: float):
print(x + y)

my_kernel(1, 1.0) # prints 2.0

Using ti.types.ndarray() as the type hint, you can pass a NumPy's ndarray or a PyTorch's tensor to a kernel. Taichi recognizes the shape and data type of such a data structure and allows you to access these attributes in a kernel. For example:

import numpy as np
import taichi as ti
ti.init(arch=ti.cpu)

x = np.array([1, 2, 3])
y = np.array([4, 5, 6])

@ti.kernel
def my_kernel(x: ti.types.ndarray(), y: ti.types.ndarray()):
for i in range(x.shape[0]):
x[i] += y[i]

my_kernel(x, y)
print(x) # prints [5, 7, 9]

x is modified by my_kernel() because it is passed by reference.

note

We skip ti.template() here and leave it for a more advanced topic: Meta-programming. Refer to Metaprogramming for more information.

返回值

一个 kernel 最多有一个返回值,这个值可以是标量,也可以是 ti.Matrixti.Vector。 定义 kernel 的返回值时须遵循如下规则:

  • 输入 kernel 返回值的类型提示。
  • 确保一个 kernel 最多有一个返回值。
  • 确保一个 kernel 最多包含一个 return 语句。
  • 确保返回值中的元素数量不超过 30。

Let's see an exmaple:

vec2 = ti.math.vec2

@ti.kernel
def test(x: float, y: float) -> vec2: # Return value must be type hinted
# Return x, y # Compilation error: Only one return value is allowed
return vec2(x, y) # OK!

此外,返回值会自动转换为提示的类型。

@ti.kernel
def my_kernel() -> ti.i32: # int32
return 128.32

print(my_kernel()) # 128, the return value is cast into ti.i32

一个 kernel 最多包含一条 return 语句

@ti.kernel
def test_sign(x: float) -> float:
if x >= 0:
return 1.0
else:
return -1.0
# Error: multiple return statements

要绕开这一限制,可以用一个局部变量保存结果,然后在结尾返回:

@ti.kernel
def test_sign(x: float) -> float:
sign = 1.0
if x < 0:
sign = -1.0
return sign
# One return statement works fine

全局变量是编译时常量

kernel 将全局变量视为编译时常量。 这意味着 kernel 在编译时接收全局变量的当前值,之后就不再追踪这些变量的变化。 如果同一个 kernel 被调用两次,期间一个全局变量的值发生变化,那么第二次调用该 kernel 时不使用变量更新后的值。

请看以下示例,其中全局变量 a 的值在 kernel_1 第一次被调用后更新。

  • 因为 kernel_1 在编译完成后不会追踪 a 的变化,所以第二次调用 kernel_1 时仍会打印 1
  • kernerl_2 编译时,a 已更新,所以 kernerl_2 接收的是 a 的当前值,并打印 2
import taichi as ti
ti.init()

a = 1

@ti.kernel
def kernel_1():
print(a)


@ti.kernel
def kernel_2():
print(a)

kernel_1() # 1
a = 2
kernel_1() # 1
kernel_2() # 2

Taichi 函数

Taichi 函数是 kernel 的构件。 You must call a Taichi function from inside a kernel or from inside another Taichi function. 所有的 Taichi 函数都被强制内联。 因此,不允许运行时递归。

Let's see an example:

# a normal python function
def foo_py():
print("I'm a python function")

@ti.func
def foo_1():
print("I'm a taichi function called by another taichi function")

@ti.func
def foo_2():
print("I'm a taichi function called by a kernel")
foo_1()

@ti.kernel
def foo_kernel():
print("I'm a kernel calling a taichi function")
foo_2()

foo_py()
#foo_func() # You cannot call a Taichi function from within the Python scope
foo_kernel()

参数

A Taichi function can have multiple arguments, supporting scalar, ti.Matrix/ti.Vector, ti.types.ndarray(), ti.template(), ti.field and ti.Struct as argument types. 请注意,适用于 kernel 参数的限制在此不适用:

  • 无需(但仍然推荐)输入参数的类型提示。
  • 参数中可以包含不限数量的元素。

返回值

The return values of a Taichi function can be scalars, ti.Matrix, ti.Vector, ti.Struct, or others. 请注意:

  • 不同于 kernel,Taichi 函数可以有多个返回值。
  • 无需(但仍然推荐)给出 Taichi 函数返回值的类型提示。

However, you still cannot have more than one return statement in a Taichi function.

小结:Taichi kernel 与Taichi 函数对比

kernelTaichi 函数
调用范围Python 作用域Taichi 作用域
参数的类型提示必需推荐
返回值的类型提示必需推荐
返回类型标量/ti.Vector/ti.Matrix标量/ti.Vector/ti.Matrix/ti.Struct/......
参数中元素数量上限
  • 32 (for OpenGL)
  • 64(适用于其他后端)
无限制
return 语句中返回值数量上限1无限制

关键术语

后端

在计算机领域,术语后端可能根据上下文有不同的含义。 一般指所有不与用户直接交互的软件程序的部分。 在使用 Taichi 时,后端是代码执行的地方,例如 cpuopenglcudavulkan

编译时递归

编译时递归是元编程的一种技巧。 此种递归由 Taichi 编译器处理,并被扩展、编译为不递归的串行函数。 在编译时,递归条件必须保持不变,递归的深度必须是常量。

强制内联

强制内联意味着用户无法选择是否内联某个函数。 被强制内联的函数总是由编译器扩展到调用方。

元编程

元编程通常是指用程序操纵程序。 就 Taichi 而言,原编程意味着使用编译时计算生成实际运行程序。 在许多情况下,这让开发者能够用尽量少的代码行来表达解决方案。

运行时递归

运行时递归是运行时发生的递归类型。 此种递归不会被编译器拓展,而是被编译为递归式调用自身的函数。 递归条件在运行时评估,深度也不一定是常量。

类型提示

类型提示是在代码中静态显示值的类型的正式解决方案。

常见问题

我可以从 Taichi 函数中调用 kernel 吗?

不可以。 要记住,kernel 是 Taichi 运行时执行的最小单位。 你不能从 Taichi 函数内(即 Taichi 作用域内)调用 kernel。 你只能从 Python 作用域调用 kernel。