kernel 与函数

Taichi and Python share a similar syntax, but they are not identical. To distinguish Taichi code from native Python code, we utilize two decorators, @ti.kernel and @ti.func:

• Functions decorated with @ti.kernel are known as Taichi kernels or simply kernels. These functions are the entry points where Taichi's runtime takes over the tasks, and they must be directly invoked by Python code. You can use native Python to prepare tasks, such as reading data from disk and pre-processing, before calling the kernel to offload computation-intensive tasks to Taichi.
• Functions decorated with @ti.func are known as Taichi functions. These functions are building blocks of kernels and can only be invoked by another Taichi function or a kernel. Like normal Python functions, you can divide your tasks into multiple Taichi functions to enhance readability and reuse them across different kernels.

In the following example, inv_square() is decorated with @ti.func and is a Taichi function. partial_sum() is decorated with @ti.kernel and is a kernel. 前者（inv_square()）被后者（partial_sum()）调用。 The arguments and return value in partial_sum() are type hinted, while those in the Taichi function inv_square() are not.

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

partial_sum(1000)

Here comes a significant difference between Python and Taichi - type hinting:

• Type hinting in Python is recommended, but not compulsory.
• Taichi mandates that the arguments and return value of a kernel are type hinted, unless it has neither an argument nor a return statement.

WARNING

Calling a Taichi function from within the native Python code (the Python scope) results in a syntax error raised by Taichi. 例如：

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

You must call Taichi functions from within the Taichi scope, a concept as opposed to the Python scope.

Let's introduce two important concepts: Taichi scope and Python scope.

• The code inside a kernel or a Taichi function is part of the Taichi scope. Taichi's runtime compiles and executes this code in parallel on multi-core CPU or GPU devices for high-performance computation. Taichi 作用域相当于 CUDA 的设备端。

• Code outside of the Taichi scope belongs to the Python scope. This code is written in native Python and executed by Python's virtual machine, not by Taichi's runtime. Python 作用域相当于 CUDA 的主机端。

It is important to distinguish between kernels and Taichi functions as they have slightly different syntax. The following sections explain their respective usages.

kernel

A kernel is the basic unit of execution in Taichi, and serves as the entry point for Taichi's runtime, which takes over from Python's virtual machine. Kernels are called in the same way as Python functions, and allow for switching between Taichi's runtime and Python's virtual machine.

For instance, the partial_sum() kernel can be called from within a Python function:

@ti.kernel
def partial_sum(n: int) -> float:
    ...

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

main()

Multiple kernels can be defined in a single Taichi program. These kernels are independent of each other, and are compiled and executed in the same order in which they are first called. The compiled kernels are cached to reduce the launch overhead for subsequent calls.

WARNING

Kernels in Taichi can be called either directly or from inside a native Python function. However, calling a kernel from inside another kernel or from inside a Taichi function is not allowed. In other words, kernels can only be called from the Python scope.

参数

A kernel can accept multiple arguments. However, it's important to note that you can't pass arbitrary Python objects to a kernel. This is because Python objects can be dynamic and may contain data that the Taichi compiler cannot recognize.

The kernel can accept various argument types, including scalars, ti.types.matrix(), ti.types.vector(), ti.types.struct(), ti.types.ndarray(), and ti.template(). These argument types make it easy to pass data from the Python scope to the Taichi scope. You can find the supported types in the ti.types module. For more information on this, see the Type System.

Scalars, ti.types.matrix(), ti.types.vector(), and ti.types.struct() are passed by value, which means that the kernel receives a copy of the argument. However, ti.types.ndarray() and ti.template() are passed by reference, which means that any changes made to the argument inside the kernel will affect the original value as well.

Note that we won't cover ti.template() here as it is a more advanced topic and is discussed in Metaprogramming.

Here is an example of passing arguments x and y 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

Here is another example of passing a nested struct argument with a matrix to a kernel by value, in which we created a struct type transform_type that contains two members: a rotation matrix R and a translation vector T. We then created another struct type pos_type that has transform_type as its member and passed an instance of pos_type to a kernel.

transform_type = ti.types.struct(R=ti.math.mat3, T=ti.math.vec3)
pos_type = ti.types.struct(x=ti.math.vec3, trans=transform_type)

@ti.kernel
def kernel_with_nested_struct_arg(p: pos_type) -> ti.math.vec3:
    return p.trans.R @ p.x + p.trans.T

trans = transform_type(ti.math.mat3(1), [1, 1, 1])
p = pos_type(x=[1, 1, 1], trans=trans)
print(kernel_with_nested_struct_arg(p))  # [4., 4., 4.]

You can use ti.types.ndarray() as a type hint to pass a ndarray from NumPy or a tensor from PyTorch to a kernel. Taichi recognizes the shape and data type of these data structures, which allows you to access their attributes in a kernel.

In the example below, x is updated after my_kernel() is called since it is passed by reference:

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()):
    # Taichi recognizes the shape of the array x and allows you to access it in a kernel
    for i in range(x.shape[0]):
        x[i] += y[i]

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

返回值

In Taichi, a kernel is allowed to have a maximum of one return value, which could either be a scalar, ti.types.matrix(), or ti.types.vector(). Moreover, in the LLVM-based backends (CPU and CUDA backends), a return value could also be a ti.types.struct().

Here is an example of a kernel that returns a ti.Struct:

s0 = ti.types.struct(a=ti.math.vec3, b=ti.i16)
s1 = ti.types.struct(a=ti.f32, b=s0)

@ti.kernel
def foo() -> s1:
    return s1(a=1, b=s0(a=ti.math.vec3(100, 0.2, 3), b=1))

print(foo())  # {'a': 1.0, 'b': {'a': [100.0, 0.2, 3.0], 'b': 1}}

When defining the return value of a kernel in Taichi, it is important to follow these rules:

• Use type hint to specify the return value of a kernel.
• Make sure that you have at most one return value in a kernel.
• Make sure that you have at most one return statement in a kernel.
• If the return value is a vector or matrix, please ensure that it contains no more than 32 elements. In case it contains more than 32 elements, the kernel will still compile, but a warning will be raised.

最多一个返回值

In this code snippet, the test() kernel cannot have more than one return value:

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)  # Fine

自动类型转换

在以下代码片段中，返回值被自动转换为提示的类型：

@ti.kernel
def my_kernel() -> ti.i32:  # int32
    return 128.32
# The return value is cast into the hinted type ti.i32
print(my_kernel())  # 128

最多一条 return 语句

In this code snippet, Taichi raises an error because the kernel test_sign() has more than one return statement:

@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

全局变量是编译时常量

In Taichi, a kernel treats global variables as compile-time constants. This means that it takes in the current values of the global variables at the time it is compiled and does not track changes to them afterwards. Consider the following example:

import taichi as ti
ti.init()

a = 1

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

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

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

Here, kernel_1 and kernel_2 both access the global variable a. The first call to kernel_1 prints 1, which is the value of a at the time the kernel was compiled. When a is updated to 2, the second call to kernel_1 still prints 1 because the kernel does not track changes to a after it is compiled.

On the other hand, kernel_2 is compiled after a is updated, so it takes in the current value of a and prints 2.

Taichi 函数

Taichi functions are fundamental units of a kernel and can only be called from within a kernel or another Taichi function.

In the code snippet below, Taichi will raise an error because the function foo_1() is called from the Python scope, not the Taichi scope:

# A normal Python function
def foo_py():
    print("This is a Python function.")

@ti.func
def foo_1():
    print("This is a Taichi function to be called by another Taichi function, foo_2().")

@ti.func
def foo_2():
    print("This is a Taichi function to be called by a kernel.")
    foo_1()

@ti.kernel
def foo_kernel():
    print("This is a kernel calling a Taichi function, foo_2().")
    foo_2()

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

WARNING

所有的 Taichi 函数都被强制内联。 This means that if you call a Taichi function from another Taichi function, the calling function is fully expanded, or inlined, into the called function at compile time. This process continues until there are no more function calls to inline, resulting in a single, large function. This means that runtime recursion is not allowed in Taichi, because it would cause an infinite expansion of the function call stack at compile time.

参数

A Taichi function can accept multiple arguments, which may include scalar, ti.types.matrix(), ti.types.vector(), ti.types.struct(), ti.types.ndarray(), ti.field(), and ti.template() types. Note that some of the restrictions on kernel arguments do not apply to Taichi functions:

• It is not strictly required to type hint the function arguments (but it is still recommended).
• You can pass an unlimited number of elements in the function arguments.

返回值

Return values of a Taichi function can be scalars, ti.types.matrix(), ti.types.vector(), ti.types.struct(), or other types. Note the following:

• 不同于 kernel，Taichi 函数可以有多个返回值。
• It is not required (but recommended) to type hint the return values of a Taichi function.
• A Taichi function cannot have more than one return statement.

小结：Taichi kernel 与Taichi 函数对比

	kernel	Taichi 函数
调用范围	Python 作用域	Taichi 作用域
参数的类型提示	必需	推荐
返回值的类型提示	必需	推荐
返回类型	标量 ti.types.matrix() ti.types.vector() ti.types.struct()(Only on LLVM-based backends)	标量 ti.types.matrix() ti.types.vector() ti.types.struct() ...
参数中元素数量上限	32（适用于 OpenGL） 64（适用于其他后端）	无限制
return 语句中返回值数量上限	1	无限制

关键术语

后端

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

编译时递归

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

强制内联

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

元编程

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

运行时递归

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

类型提示

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

常见问题

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

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