Taichi Lang

Taichi is a domain-specific language (DSL) embedded in Python designed for high-performance visual computing and physics simulation. It uses a Just-In-Time (JIT) compilation framework (utilizing technologies like LLVM and SPIR-V) to offload Python source code to native GPU or CPU instructions.

Key characteristics include:

• Imperative Programming Paradigm: Unlike graph-based DL frameworks, Taichi allows writing large amounts of computation in a single "mega-kernel", providing flexibility for non-standard computing patterns.
• Decoupled Computation and Data Structures: Uses a mechanism called SNode to compose hierarchical, dense, or sparse multi-dimensional fields, allowing easy switching between memory layouts (e.g., Array-of-Structures vs. Structure-of-Arrays).
• Customizable Quantized Types: Supports defining fixed-point or floating-point numbers with arbitrary bit widths (up to 64 bits) to optimize GPU memory capacity and bandwidth.
• Automatic Differentiation: Provides an automatic differentiation module via source code transformation at the Intermediate Representation (IR) level, suitable for differentiable simulation in machine learning.

Taichi is a domain-specific language embedded in Python. It uses a multi-stage programming model with two distinct types of evaluation and values:

1. Python Values (Compile-time)

Python values exist only during the compilation phase. They are sourced from:

• Literals
• Arguments passed via ti.template()
• Free variables

When an operation involves only Python values, Taichi performs compile-time evaluation, resulting in a Python value. You can use ti.static() to access an advanced environment for compile-time meta-programming.

2. Taichi Values (Runtime)

Taichi values are evaluated at runtime. They have Taichi types, which include:

• Primitive types
• Compound types
• ndarray types
• Sparse matrix builder types

Evaluation Rules Summary

• Python value + Python value = Python value (Compile-time)
• Python value + Taichi value = Taichi value (Runtime)
• Taichi value + Taichi value = Taichi value (Runtime)

The taichi.math module provides specialized small vector and matrix types inspired by GLSL. These are useful for graphics and physics computations.

Available Types:

• vec2/vec3/vec4: 2D/3D/4D floating-point vectors.
• ivec2/ivec3/ivec4: 2D/3D/4D integer vectors.
• uvec2/uvec3/uvec4: 2D/3D/4D unsigned integer vectors.
• mat2/mat3/mat4: 2D/3D/4D floating-point square matrices.

Initialization: These types have flexible initialization routines. You can initialize them with scalars, lists, or other vectors/matrices.

Vector Swizzling: Vector types created via ti.types.vector() support GLSL-style swizzling (e.g., .xyz, .rgba, .wzyx) to access or rearrange elements.

import taichi as ti
import taichi.math as tm

# Using pre-defined types from taichi.math
vec3 = tm.vec3
mat2 = tm.mat2

v = vec3(1, 2, 3)
u = v.xyz
w = v.xxx  # vec3(1, 1, 1)

m = mat2(1, 2, 3, 4)  # [[1., 2.], [3., 4.]]

Atoms are the most basic elements of expressions in Taichi. They include:

• Identifiers (Names): Follow Python lexical rules. If a name is visible in Taichi, it evaluates to the runtime variable value. If it is only visible in Python (outside Taichi), it triggers compile-time evaluation of the Python value. If invisible, a TaichiNameError is thrown.
• Literals: Integer and floating-point literals (following Python) are evaluated to Python values at compile time.
• Parenthesized forms: (expression_list) evaluates to the result of the list. An empty pair () evaluates to an empty tuple at compile time.
• List and Dictionary displays: Taichi supports list [...] and dictionary {...} construction via explicit items or comprehensions. These are evaluated at compile time, meaning all expressions in a comprehension or keys in a dictionary must be evaluatable to Python values.

@ti.kernel
def test(p: ti.i32):
    # Valid: the range is a Python value at compile time
    a = ti.Matrix([i * p for i in range(10)]) 
    
    # Compile error: p is a Taichi runtime value, not a Python value
    b = ti.Matrix([i * p for i in range(p)]) 

The ListManager is a low-level data structure used by the runtime to manage memory in chunks, behaving similarly to std::deque in C++. It is designed to be 'infinitely long' by allocating new chunks on demand.

Key Characteristics:

• Chunked Storage: Data is organized into chunks, where each chunk holds max_num_elements_per_chunk elements.
• On-demand Allocation: New chunks are allocated via touch_chunk() when the current capacity is exceeded.
• Indexing: Elements are accessed via an index, which is converted to a pointer using bitwise operations based on the chunk size (which must be a power of two).

Core Methods:

• reserve_new_element(): Increments the element count and ensures the corresponding chunk is allocated.
• get_element_ptr(i32 i): Calculates the memory address of the element at index i.
• ptr2index(): Performs the reverse operation, finding the index associated with a given pointer by checking chunk ranges.
• append(void *data_ptr): Allocates a new element and copies data into it.
• clear(): Resets the element count to zero without clearing the actual memory contents.

// Example of how ListManager calculates element address
Ptr get_element_ptr(i32 i) {
  return chunks[i >> log2chunk_num_elements] +                       // chunk base
          element_size * (i & ((1 << log2chunk_num_elements) - 1));  // slot within the chunk
}

To reduce memory, use ti.BitpackedFields to pack multiple fields (whose dtypes are quantized) into a single primitive type. You can place a BitpackedFields instance under an SNode just like a regular field.

Shared Exponent

When packing multiple quantized floating-point fields, you can use shared_exponent=True in the .place() method. This allows fields to share a common exponent, leaving more bits available for the components (e.g., in a 3D velocity vector where components have similar magnitudes).

a = ti.field(float_type_a)  # 15 bits
b = ti.field(fixed_type_b)  # 5 bits
c = ti.field(fixed_type_c)  # 6 bits
d = ti.field(u5)            # 5 bits

bitpack = ti.BitpackedFields(max_num_bits=32)
bitpack.place(a, b, c, d)  # 31 bits used
ti.root.dense(ti.i, 10).place(bitpack)

# Example with shared exponent for floating point components
bitpack.place(velocity, shared_exponent=True)