`taichi.aot`#

Taichi’s AOT (ahead of time) module.

Users can use Taichi as a GPU compute shader/kernel compiler by compiling their Taichi kernels into an AOT module.

class taichi.aot.Module(arch=None, caps=None)#

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_graph(self, name, graph)#

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

Add a taichi kernel to the AOT module.

Parameters:

kernel_fn (Function) – the function decorated by taichi kernel.
template_args (Dict[str, Any]) – a dict where key is the template parameter name, and value is the instantiating arg. Note that this works for both template and for :class:`~taichi.types.ndarray.
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)

archive(self, filepath: str)#

Parameters:: filepath (str) – path to the stored archive of aot artifacts, MUST end with .tcm.

save(self, filepath)#

Parameters:: filepath (str) – path to a folder to store aot files.

taichi.aot.start_recording(filename)#

Starts recording kernel information to a yml file.

Parameters:: filename (str) – output yml file.

Example:

>>> ti.aot.start_recording('record.yml')
>>> ti.init(arch=ti.cc)
>>> loss = ti.field(float, (), needs_grad=True)
>>> x = ti.field(float, 233, needs_grad=True)
>>>
>>> @ti.kernel
>>> def compute_loss():
>>>     for i in x:
>>>         loss[None] += x[i]**2
>>>
>>> @ti.kernel
>>> def do_some_works():
>>>     for i in x:
>>>         x[i] -= x.grad[i]
>>>
>>> with ti.ad.Tape(loss):
>>>     compute_loss()
>>> do_some_works()

taichi.aot.stop_recording()#

Stops recording kernel information.

This function should be called in pair with start_recording().

taichi.aot#

`taichi.aot`#