`taichi.lang.kernel_impl`¶

Module Contents¶

Classes¶

`Func`
`TaichiCallableTemplateMapper`
`Kernel`

Functions¶

`func`(fn)	Marks a function as callable in Taichi-scope.
`pyfunc`(fn)	Marks a function as callable in both Taichi and Python scopes.
`kernel`(fn)	Marks a function as a Taichi kernel.
`data_oriented`(cls)	Marks a class as Taichi compatible.

Attributes¶

`classfunc`
`classkernel`

taichi.lang.kernel_impl.func(fn)¶

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
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.pyfunc(fn)¶

Marks a function as callable in both Taichi and Python scopes.

When called inside the Taichi scope, Taichi will JIT compile it into native instructions. Otherwise it will be invoked directly as a Python function.

See also func().

Parameters: fn (Callable) – The Python function to be decorated
Returns: The decorated function
Return type: Callable

class taichi.lang.kernel_impl.Func(_func, _classfunc=False, _pyfunc=False)¶

function_counter = 0¶

func_call_rvalue(self, key, args)¶

do_compile(self, key, args)¶

extract_arguments(self)¶

class taichi.lang.kernel_impl.TaichiCallableTemplateMapper(annotations, template_slot_locations)¶

static extract_arg(arg, anno)¶

extract(self, args)¶

lookup(self, args)¶

exception taichi.lang.kernel_impl.KernelDefError¶

Bases: Exception

Common base class for all non-exit exceptions.

exception taichi.lang.kernel_impl.KernelArgError(pos, needed, provided)¶

Bases: Exception

Common base class for all non-exit exceptions.

class taichi.lang.kernel_impl.Kernel(_func, is_grad, _classkernel=False)¶

counter = 0¶

reset(self)¶

extract_arguments(self)¶

materialize(self, key=None, args=None, arg_features=None)¶

get_torch_callbacks(self, v, has_torch, is_ndarray=True)¶

get_function_body(self, t_kernel)¶

static match_ext_arr(v)¶

ensure_compiled(self, *args)¶

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.

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

taichi.lang.kernel_impl.classfunc¶

taichi.lang.kernel_impl.classkernel¶

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.

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¶

Module Contents¶

Classes¶

Functions¶

Attributes¶

`taichi.lang.kernel_impl`¶