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