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

kernel 与函数

Taichi 是一个嵌入在 Python 中的领域特定语言(DSL)。 “嵌入 Python”有两层含义:你用 Taichi 写的代码是有效的 Python 代码,但它是由 Taichi 的运行时编译和执行的;其余的代码则按 Python 原生代码处理,并在 Python 的虚拟机中执行。

我们用两个装饰器:@ti.kernel@ti.func 来区分 Taichi 代码和 Python 代码:

  • @ti.kernel 装饰的函数被称为 kernel。
  • @ti.func 装饰的函数被称为 Taichi 函数。

Taichi 作用域与 Python 作用域

Taichi 引入了两个术语:“Taichi 作用域”和“Python 作用域”,以便于定义函数的调用范围或指定变量的有效范围。 为更好地理解后文,请先了解这两个术语。

Taichi 作用域

kernel 或 Taichi 函数内的代码处在 Taichi 作用域内。 Taichi 作用域内的代码由 Taichi 的运行时编译,并在 CPU 或 GPU 设备上并行执行,以实现高性能计算。

note

Taichi 作用域相当于 CUDA 的设备端

Python 作用域

处在 Taichi 作用域之外的代码即属于 Python 作用域。 Python 作用域中的代码为 Python 原生代码,由 Python 虚拟机而非Taichi 的运行时执行。

note

Python 作用域相当于 CUDA 的主机端

kernel

kernel 是标志 Taichi 的运行时接管的入口,也是运行时执行的最小单位。 一个程序中可定义多个 kernel,kernel 之间相互独立。 调用 kernel 的方式与调用 Python 函数相同,并且你可以在 Taichi 的运行时和 Python 的虚拟机之间来回切换。

Taichi 的运行时按照调用顺序编译并执行 kernel。 它将已编译的 kernel 存入缓存,下一次调用到同一 kernel 时就不需要再重新编译。

WARNING

不能从一个 kernel 中调用另一个 kernel,或是从 Taichi 函数中调用 kernel。 只能直接调用或从 Python 原生函数中调用 kernel。 换言之,只能从 Python 作用域调用 kernel。

note

kernel 相当于 CUDA 中的__global__函数。

参数

一个 kernel 可以接收多个参数,支持标量、ti.Matrix,以及 ti.Vector 作为参数类型。 这使得从 Python 作用域传递值到 Taichi 作用域变得更轻松、灵活。

WARNING

向标量、ti.Matrix, 及 ti.Vector 中传递参数采用值传递的方式,所以改变 kernel 中的参数不会影响调用者函数中的原始变量。

定义参数时须遵循如下规则:

  • 输入 kernel 参数的类型提示。
  • 确保 kernel 参数中的元素总数不超过某一上限(见下文)。

kernel 参数的类型提示

@ti.kernel
def my_kernel(x: ti.i32, y: ti.f32):
print(x + y)

my_kernel(24, 3.2) # The system prints 27.2

确保 kernel 参数中的元素总数不超过某一上限

元素数量的上限与后端相关:

  • OpenGL 后端支持 8 个元素
  • CPU、Vulkan、CUDA,和 Metal 均支持 64 个元素
note
  • 标量参数中的元素数量总是 1。
  • ti.Matrixti.Vector 中的元素数量是其所含标量的实际个数。
@ti.kernel
def valid_scalar_argument(vx: ti.f32, vy: ti.f32):
v = ti.Vector([vx, vy])
...

@ti.kernel
def valid_matrix_argument(u: ti.i32, v: ti.types.matrix(2, 2, ti.i32)): # OK: has five elements in total
...

@ti.kernel
def error_too_many_arguments(u: ti.i32, v: ti.i64, w: ti.types.matrix(7, 9, ti.i64)): # Error: has 65 elements in total
...
高级参数

初学者可以跳过这一部分。

kernel 也可以接收以下两种类型的高级参数:

返回值

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

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

Kernel 返回值的类型提示

@ti.kernel
def test(x: ti.f32) -> ti.f32: # The return value is type hinted
return 1.0

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

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

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

一个 kernel 最多有一个返回值。

@ti.kernel
def error_multiple_return() -> (ti.i32, ti.f32):
x = 1
y = 0.5
return x, y # Compilation error: more than one return value

一个 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

返回值最多包含 30 个元素

N = 6
matN = ti.types.matrix(N, N, ti.i32)

@ti.kernel
def test_kernel() -> matN:
return matN([[0] * N for _ in range(N)])
# Compilation error: The number of elements is 36 > 30

全局变量是编译时常量

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 的构件。 所有的 Taichi 函数都被强制内联。 因此,不允许运行时递归。

WARNING

必须从 kernel 或另一个 Taichi 函数内调用 Taichi 函数。 换言之,只能从 Taichi 作用域内调用 Taichi 函数,而不能从 Python 作用域调用。

note

Taichi 函数相当于 CUDA 中的 __device__ 函数。

下面的示例显示了 kernel 和 Taichi 函数之间的差异:

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

参数

一个 Taichi 函数可以接收多个参数,支持标量、ti.Matrix,以及 ti.Vector 作为参数类型。 请注意,适用于 kernel 参数的限制在此不适用:

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

向标量、ti.Matrix, 及 ti.Vector 中传递参数采用值传递的方式,所以改变 Taichi 函数中的参数不会影响调用者函数中的原始变量。 见以下示例:

@ti.func
def my_func(x):
x = x + 1 # Will not change the original value of x


@ti.kernel
def my_kernel():
x = 24
my_func(x)
print(x) # 24
高级参数

初学者可以跳过这一部分。

Taichi 函数也可以接收模板参数。 使用 ti.template() 作为类型提示。 ti.template() 将强制参数进行引用传递。 例如:

@ti.func
def my_func(x: ti.template()): # x is forced to pass by reference
x = x + 1 # This line changes the original value of x

@ti.kernel
def my_kernel():
x = 24
my_func(x)
print(x) # will print 25

返回值

Taichi 函数的返回值可以是标量、ti.Matrixti.Vectorti.Struct,或其它类型。 请注意:

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

但是,Taichi 函数中不能包含超过一条的 return 语句。

最多一条 return 语句

Taichi 函数中只能有一条 return 语句。

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

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

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

小结:Taichi kernel 与Taichi 函数对比

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

关键术语

后端

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

编译时递归

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

强制内联

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

元编程

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

运行时递归

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

类型提示

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

常见问题

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

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