Kernels and functions¶
Taichi-scope vs Python-scope¶
Code decorated by @ti.kernel
or @ti.func
is in the Taichi-scope.
They are to be compiled and executed on CPU or GPU devices with high parallelization performance, on the cost of less flexibility.
Note
For people from CUDA, Taichi-scope = device side.
Code outside @ti.kernel
or @ti.func
is in the Python-scope.
They are not compiled by the Taichi compiler and have lower performance but with a richer type system and better flexibility.
Note
For people from CUDA, Python-scope = host side.
Kernels¶
A Python function decorated by @ti.kernel
is a Taichi kernel:
@ti.kernel
def my_kernel():
...
my_kernel()
Kernels should be called from Python-scope.
Note
For people from CUDA, Taichi kernels = __global__
functions.
Arguments¶
Kernels can have at most 8 parameters so that you can pass values from Python-scope to Taichi-scope easily.
Kernel arguments must be type-hinted:
@ti.kernel
def my_kernel(x: ti.i32, y: ti.f32):
print(x + y)
my_kernel(2, 3.3) # prints: 5.3
Note
For now, we only support scalars as arguments. Specifying ti.Matrix
or ti.Vector
as argument is not supported. For example:
@ti.kernel
def bad_kernel(v: ti.Vector):
...
@ti.kernel
def good_kernel(vx: ti.f32, vy: ti.f32):
v = ti.Vector([vx, vy])
...
Return value¶
A kernel may or may not have a scalar return value. If it does, the type of return value must be hinted:
@ti.kernel
def my_kernel() -> ti.f32:
return 233.33
print(my_kernel()) # 233.33
The return value will be automatically cast into the hinted type. e.g.,
@ti.kernel
def add_xy() -> ti.i32: # int32
return 233.33
print(my_kernel()) # 233, since return type is ti.i32
Note
For now, a kernel can only have one scalar return value. Returning ti.Matrix
or ti.Vector
is not supported. Python-style tuple return is not supported either. For example:
@ti.kernel
def bad_kernel() -> ti.Matrix:
return ti.Matrix([[1, 0], [0, 1]]) # Error
@ti.kernel
def bad_kernel() -> (ti.i32, ti.f32):
x = 1
y = 0.5
return x, y # Error
Advanced arguments¶
We also support template arguments (see Template metaprogramming) and external array arguments (see Interacting with external arrays) in Taichi kernels. Use ti.template()
or ti.ext_arr()
as their type-hints respectively.
Note
When using differentiable programming, there are a few more constraints on kernel structures. See the Kernel Simplicity Rule in Differentiable programming.
Also, please do not use kernel return values in differentiable programming, since the return value will not be tracked by automatic differentiation. Instead, store the result into a global variable (e.g. loss[None]
).
Functions¶
A Python function decorated by @ti.func
is a Taichi function:
@ti.func
def my_func():
...
@ti.kernel
def my_kernel():
...
my_func() # call functions from Taichi-scope
...
my_kernel() # call kernels from Python-scope
Taichi functions should be called from Taichi-scope.
Note
For people from CUDA, Taichi functions = __device__
functions.
Note
Taichi functions can be nested.
Warning
Currently, all functions are force-inlined. Therefore, no recursion is allowed.
Arguments and return values¶
Functions can have multiple arguments and return values. Unlike kernels, arguments in functions don’t need to be type-hinted:
@ti.func
def my_add(x, y):
return x + y
@ti.kernel
def my_kernel():
...
ret = my_add(2, 3.3)
print(ret) # 5.3
...
Function arguments are passed by value. So changes made inside function scope won’t affect the outside value in the caller:
@ti.func
def my_func(x):
x = x + 1 # won't change the original value of x
@ti.kernel
def my_kernel():
...
x = 233
my_func(x)
print(x) # 233
...
Advanced arguments¶
You may use ti.template()
as type-hint to force arguments to be passed by
reference:
@ti.func
def my_func(x: ti.template()):
x = x + 1 # will change the original value of x
@ti.kernel
def my_kernel():
...
x = 233
my_func(x)
print(x) # 234
...
Note
Unlike kernels, functions do support vectors or matrices as arguments and return values:
@ti.func
def sdf(u): # functions support matrices and vectors as arguments. No type-hints needed.
return u.norm() - 1
@ti.kernel
def render(d_x: ti.f32, d_y: ti.f32): # kernels do not support vector/matrix arguments yet. We have to use a workaround.
d = ti.Vector([d_x, d_y])
p = ti.Vector([0.0, 0.0])
t = sdf(p)
p += d * t
...
Warning
Functions with multiple return
statements are not supported for now. Use a local variable to store the results, so that you end up with only one return
statement:
# Bad function - two return statements
@ti.func
def safe_sqrt(x):
if x >= 0:
return ti.sqrt(x)
else:
return 0.0
# Good function - single return statement
@ti.func
def safe_sqrt(x):
ret = 0.0
if x >= 0:
ret = ti.sqrt(x)
else:
ret = 0.0
return ret