`taichi.aot`#

class taichi.aot.Module(arch)#

An AOT module to save and load Taichi kernels.

This module serializes the Taichi kernels for a specific arch. The serialized module can later be loaded to run on that backend, without the Python environment.

Example

Usage:

m = ti.aot.Module(ti.metal)
m.add_kernel(foo)
m.add_kernel(bar)

m.save('/path/to/module')

# Now the module file '/path/to/module' contains the Metal kernels
# for running ``foo`` and ``bar``.

add_field(self, name, field)#

Add a taichi field to the AOT module.

Parameters

name – name of taichi field
field – taichi field

Example:

>>> a = ti.field(ti.f32, shape=(4,4))
>>> b = ti.field("something")
>>>
>>> m.add_field(a)
>>> m.add_field(b)
>>>
>>> # Must add in sequence

add_kernel(self, kernel_fn, example_any_arrays=None, name=None)#

Add a taichi kernel to the AOT module.

Parameters

kernel_fn (Function) – the function decorated by taichi kernel.
example_any_arrays (Dict[int, ti.ndarray]) – a dict where key is arg_id and key is example any_arr input.
name (str) – Name to identify this kernel in the module. If not provided, uses the built-in __name__ attribute of kernel_fn.

add_kernel_template(self, kernel_fn)#

Add a taichi kernel (with template parameters) to the AOT module.

Parameters: kernel_fn (Function) – the function decorated by taichi kernel.

Example:

>>> @ti.kernel
>>> def bar_tmpl(a: ti.template()):
>>>   x = a
>>>   # or y = a
>>>   # do something with `x` or `y`
>>>
>>> m = ti.aot.Module(arch)
>>> with m.add_kernel_template(bar_tmpl) as kt:
>>>   kt.instantiate(a=x)
>>>   kt.instantiate(a=y)
>>>
>>> @ti.kernel
>>> def bar_tmpl_multiple_args(a: ti.template(), b: ti.template())
>>>   x = a
>>>   y = b
>>>   # do something with `x` and `y`
>>>
>>> with m.add_kernel_template(bar_tmpl) as kt:
>>>   kt.instantiate(a=x, b=y)

save(self, filepath, filename)#

Parameters

filepath (str) – path to a folder to store aot files.
filename (str) – filename prefix for stored aot files.

taichi.aot.start_recording(filename)#

taichi.aot.stop_recording()#

taichi.aot#

`taichi.aot`#