Taichi provides metaprogramming infrastructures. There are many benefits of metaprogramming in Taichi:
- Enabling the development of dimensionality-independent code, e.g., code which is adaptive for both 2D/3D physical simulations.
- Improving runtime performance by moving computations from runtime to compile time.
- Simplifying the development of Taichi standard library.
Taichi kernels are lazily instantiated and large amounts of computation can be executed at compile-time. Every kernel in Taichi is a template kernel, even if it has no template arguments.
ti.template() as a argument type hint, a Taichi field can be passed into a kernel. Template programming also enables the code to be reused for fields with different shapes:
@ti.kerneldef copy_1D(x: ti.template(), y: ti.template()): for i in x: y[i] = x[i] a = ti.field(ti.f32, 4)b = ti.field(ti.f32, 4)c = ti.field(ti.f32, 12)d = ti.field(ti.f32, 12) # Pass field a and b as arguments of the kernel `copy_1D`:copy_1D(a, b) # Reuse the kernel for field c and d:copy_1D(c, d)
The template parameters are inlined into the generated kernel after compilation.
ti.grouped syntax which supports grouping loop indices into a
ti.Vector. It enables dimensionality-independent programming, i.e., code are adaptive to scenarios of different dimensionalities automatically:
@ti.kerneldef copy_1D(x: ti.template(), y: ti.template()): for i in x: y[i] = x[i] @ti.kerneldef copy_2d(x: ti.template(), y: ti.template()): for i, j in x: y[i, j] = x[i, j] @ti.kerneldef copy_3d(x: ti.template(), y: ti.template()): for i, j, k in x: y[i, j, k] = x[i, j, k] # Kernels listed above can be unified into one kernel using `ti.grouped`:@ti.kerneldef copy(x: ti.template(), y: ti.template()): for I in ti.grouped(y): # I is a vector with dimensionality same to y # If y is 0D, then I = ti.Vector(), which is equivalent to `None` used in x[I] # If y is 1D, then I = ti.Vector([i]) # If y is 2D, then I = ti.Vector([i, j]) # If y is 3D, then I = ti.Vector([i, j, k]) # ... x[I] = y[I]
The two attributes data type and shape of fields can be accessed by
field.shape, in both Taichi-scope and Python-scope:
x = ti.field(dtype=ti.f32, shape=(3, 3)) # Print field metadata in Python-scopeprint("Field dimensionality is ", x.shape)print("Field data type is ", x.dtype) # Print field metadata in Taichifirstname.lastname@example.org print_field_metadata(x: ti.template()): print("Field dimensionality is ", len(x.shape)) for i in ti.static(range(len(x.shape))): print("Size along dimension ", i, "is", x.shape[i]) ti.static_print("Field data type is ", x.dtype)
For sparse fields, the full domain shape will be returned.
matrix.n returns the number of columns and rows, respectively. For vectors, they are treated as matrices with one column in Taichi, where
vector.n is the number of elements of the vector.
@ti.kerneldef foo(): matrix = ti.Matrix([[1, 2], [3, 4], [5, 6]]) print(matrix.n) # number of row: 3 print(matrix.m) # number of column: 2 vector = ti.Vector([7, 8, 9]) print(vector.n) # number of elements: 3 print(vector.m) # always equals to 1 for a vector
Using compile-time evaluation allows for some computation to be executed when kernels are instantiated. This helps the compiler to conduct optimization and reduce computational overhead at runtime:
ti.staticfor compile-time branching (for those who are familiar with C++17, this is similar to if constexpr.):
enable_projection = True @ti.kerneldef static(): if ti.static(enable_projection): # No runtime overhead x = 1
One of the two branches of the
static if will be discarded after compilation.
ti.staticfor forced loop unrolling:
@ti.kerneldef func(): for i in ti.static(range(4)): print(i) # The code snippet above is equivalent to: print(0) print(1) print(2) print(3)
When to use
ti.static with for loops#
There are two reasons to use
ti.static with for loops:
- Loop unrolling for improving runtime performance (see section [Compile-time evaluations](##Compile-time evaluations)).
- Accessing elements of Taichi matrices/vectors. Indices for accessing Taichi fields can be runtime variables, while indices for Taichi matrices/vectors must be a compile-time constant.
For example, when accessing a vector field
field_index can be a runtime variable, while the
vector_component_index must be a compile-time constant:
# Here we declare a field contains 3 vector. Each vector contains 8 elements.x = ti.Vector.field(8, ti.f32, shape=(3))@ti.kerneldef reset(): for i in x: for j in ti.static(range(x.n)): # The inner loop must be unrolled since j is an index for accessing a vector x[i][j] = 0