跳到主要内容

taichi.lang.kernel_impl#

taichi.lang.kernel_impl.data_oriented(cls)#

Marks a class as Taichi compatible.

To allow for modularized code, Taichi provides this decorator so that Taichi kernels can be defined inside a class.

See also https://docs.taichi-lang.org/docs/odop

Example:

>>> @ti.data_oriented
>>> class TiArray:
>>>     def __init__(self, n):
>>>         self.x = ti.field(ti.f32, shape=n)
>>>
>>>     @ti.kernel
>>>     def inc(self):
>>>         for i in self.x:
>>>             self.x[i] += 1.0
>>>
>>> a = TiArray(32)
>>> a.inc()
Parameters:

cls (Class) – the class to be decorated

Returns:

The decorated class.

taichi.lang.kernel_impl.func(fn, is_real_function=False)#

Marks a function as callable in Taichi-scope.

This decorator transforms a Python function into a Taichi one. Taichi will JIT compile it into native instructions.

Parameters:
  • fn (Callable) – The Python function to be decorated

  • is_real_function (bool) – Whether the function is a real function

Returns:

The decorated function

Return type:

Callable

Example:

>>> @ti.func
>>> def foo(x):
>>>     return x + 2
>>>
>>> @ti.kernel
>>> def run():
>>>     print(foo(40))  # 42
taichi.lang.kernel_impl.kernel(fn)#

Marks a function as a Taichi kernel.

A Taichi kernel is a function written in Python, and gets JIT compiled by Taichi into native CPU/GPU instructions (e.g. a series of CUDA kernels). The top-level for loops are automatically parallelized, and distributed to either a CPU thread pool or massively parallel GPUs.

Kernel’s gradient kernel would be generated automatically by the AutoDiff system.

See also https://docs.taichi-lang.org/docs/syntax#kernel.

Parameters:

fn (Callable) – the Python function to be decorated

Returns:

The decorated function

Return type:

Callable

Example:

>>> x = ti.field(ti.i32, shape=(4, 8))
>>>
>>> @ti.kernel
>>> def run():
>>>     # Assigns all the elements of `x` in parallel.
>>>     for i in x:
>>>         x[i] = i