taichi.aot.module

Module Contents

Classes

KernelTemplate

Module

An AOT module to save and load Taichi kernels.
class taichi.aot.module.KernelTemplate(kernel_fn, aot_module)
static keygen(v, key_p, fields)
instantiate(self, **kwargs)
class taichi.aot.module.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.