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

kernel 与函数

Taichi 和 Python 语法相似,但并不完全相同。 我们用两个装饰器:@ti.kernel@ti.func 来区分 Taichi 代码和 Python 原生代码:

  • 使用 @ti.kernel 装饰的函数是 Taichi kernel,简称为 kernel。 kernel 是 Taichi 运行时接管任务的入口,且必须直接被 Python 代码调用。 你可以用 Python 原生代码为任务作准备,如从磁盘读取数据并进行预处理,然后调用 kernel,让 Taichi 接管计算密集任务。

  • @ti.func 装饰的函数是 Taichi 函数。 它们是 kernel 的构成部分,只能被 kernel 或另一个 Taichi 函数调用。 正如使用普通的 Python 函数一样,你可以将任务分解成多个 Taichi 函数,让代码更易读,并在不同的 kernel 中复用这些函数。

在下面的示例中,inv_square()@ti. func 装饰,因而是一个 Taichi 函数;partial_sum()@ti.kernel 装饰,因而是一个 kernel。 前者(inv_square())被后者(partial_sum())调用。 partial_sum() 中的参数和返回值都带有类型提示,而 Taichi 函数 inv_square() 中的参数和返回值没有。

这涉及到 Python 和 Taichi 之间的一个重要区别——类型提示:

  • 在 Python 中,推荐使用类型提示,但并非强制要求。
  • Taichi 规定必须提示 kernel 的参数和返回值的类型,除非该 kernel 没有参数或 return 语句。
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
WARNING

如果你从 Python 原生代码(即 Python 作用域)中调用 Taichi 函数,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

你必须在 Taichi 作用域中调用 Taichi 函数。Taichi 作用域是与 Python 作用域 相对的概念。

IMPORTANT

为方便起见,我们引入两个概念:Taichi 作用域Python 作用域

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

  • Taichi 作用域外的代码处在 Python 作用域中。 Python 作用域中的代码为 Python 原生代码,由 Python 虚拟机而非Taichi 的运行时执行。 Python 作用域相当于 CUDA 的主机端。

注意避免混淆 kernel 和 Taichi 函数。 二者语法略有不同。 下文描述了它们的用法。

kernel

kernel 是在 Taichi 的基础执行单元,也是 Taichi 运行时从 Python 虚拟机接管任务的入口。 调用 kernel 的方式与调用 Python 函数相同,并且你可以在 Taichi 的运行时和 Python 的虚拟机之间来回切换。

例如,你可以在一个 Python 函数中调用 kernel partial_sum()

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

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

main()

你可以在程序中定义多个 kernel。 kernel 之间相互独立 ,并且按照首次被调用的顺序进行编译和执行(编译后的 kernel 存储在缓存中,以节省后续调用的启动开销)。

WARNING

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

参数

一个 kernel 可以接收多个参数。 需注意,不能将任意的 Python 对象传入 kernel,因为 Python 对象可能是高度动态的,并且可能含有无法被 Taichi 编译器识别的数据。

kernel 接收的参数类型有标量、ti.Matrixti.Vector(向量本质也是矩阵)、ti.types.ndarray()ti.template(),让你可以轻松将数据从 Python 作用域传入 Taichi 作用域。 ti.types 模块中定义了支持的类型(详情请参阅 类型系统)。

  • 标量和 ti.Matrix 按值传递
  • ti.types.ndarray()ti.template() 均按引用传递。 这意味着对被调用 kernel 中参数的任何修改也会影响到原始值。

我们在此跳过 ti.template() ,留待在更高阶的话题:元编程 中探讨。

在下面的例子中,参数 xy 通过传递到 my_kernel

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

my_kernel(1, 1.0) # Prints 2.0

使用 ti.types.ndarray() 作为类型提示,你可以将 NumPy 的 ndarray 或 PyTorch 的 tensor 传到 kernel 中。 Taichi 能识别这种数据结构的形状和数据类型,并允许你在 kernel 中访问这些属性。 在下面的示例中, xmy_kernel() 被调用后更新,因为它是通过引用传递的。

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]

返回值

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

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

最多一个返回值

在以下代码片段中,kernel test() 不能有一个以上的返回值:

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 语句

在以下代码片段中,Taichi 抛出了一个错误,因为 kernel test_sign() 有不止一个 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 的第二次调用仍然打印 1,因为 kernel_1 在编译完成后不会跟踪对 a 的更改。
  • kernerl_2 接收 a 的当前值,并打印 2,因为该 kernel 在 a 更新后编译。
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

Taichi 函数

Taichi 函数是 kernel 的构件。 必须从 kernel 或另一个 Taichi 函数内调用 Taichi 函数

在以下代码片段中,Taichi 抛出了一个错误,因为 Taichi 函数 foo_1() 必须从 Taichi 作用域调用:

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

参数

一个 Taichi 函数可以有多个参数,支持标量、ti.Matrixti.Vectorti.types.ndarray()ti.template()ti.fieldti.Struct 作为参数类型。 请注意,适用于 kernel 参数的限制在此不适用:

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

返回值

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

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

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

小结:Taichi kernel 与Taichi 函数对比

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

关键术语

后端

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

编译时递归

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

强制内联

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

元编程

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

运行时递归

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

类型提示

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

常见问题

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

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