Taichi Lang

Taichi Lang is an open-source domain-specific language (DSL) designed for high-performance parallel computing.. It is embedded in Python and uses its own intermediate representation (IR) to just-in-time (JIT) compile compute-intense Python code to native GPU or CPU instructions. It adopts the imperative programming paradigm.

The language was invented by Yuanming Hu in 2019 to reduce the burden of learning close-to-hardware languages for computer graphics researchers. The design philosophies center around friendly syntax (similar to Python), high performance, and strong portability.

Three years into its inception, the open-source project has grown into a general parallel programming language and found applications in many fields, including physical simulation, image processing, and visual effects.

History

In January 2019, Hu, who was a doctoral student at MIT at the time, began to work on a brand new programming language. He noticed that the traditional close-to-hardware languages face a dissatisfying trade-off between performance and productivity when performing visual computing workloads. He set two high-level goals for his project later known as Taichi Lang: To simplify the process of high-performance visual computing system development and deployment, and to explore novel language abstractions and compilation approaches for visual computing. He introduced Taichi in a paper published on SIGGRAPH Asia 2019, which proves that Taichi achieves 4.55x higher performance in 1/10 lines of code than hand-optimized reference implementations.

In 2020, Hu moved Taichi's frontend from C++ to Python to take advantage of the latter's friendly syntax and lower the learning barrier to suit a wider user base. He also enabled differentiable programming in Taichi (DiffTaichi) and shared the results of the DiffTaichi project at ICLR 2020.

On April 13, 2022, Taichi v1.0.0, the first milestone, was released. Starting from the v1.2.0 release announced on October 25, 2022, Taichi follows semantic versioning.

As of November 2022, Taichi had around 200 contributors from around the world.

Features

JIT compilation

The JIT compilation pipeline first transforms the Python Abstract Syntax Tree (AST) of a given Python function into the Taichi "Front-end AST", which is then auto-parallelized and transformed into an IR developed by Taichi. It is under CHI-IR that the Taichi compiler performs all the domain-specific optimizations. Afterwards, the optimized CHI-IR is sent to the "Taichi backend", where it is compiled into executable GPU kernels.

High-performance parallel computing

Taichi offers a decorator ti.kernel as the entry point from which Taichi's runtime takes over compute-intensive tasks from Python.

The high performance of Taichi programs boils down to three reasons:

• Taichi's compiler compilers kernels directly to machine code.
• The outermost for loops in a kernel are automatically parallelized.
• The same code can be deployed to various backends.

Integration with Python ecosystem

Users can import Taichi (import taichi as ti) as they do with other Python libraries and write Taichi code the same way as they write Python except for a few extra rules. Taichi is part of the Python ecosystem and can interact with mainstream computing frameworks embedded in Python, such as NumPy and PyTorch.

Taichi provides various interfaces for data transfer between Taichi fields and NumPy arrays/PyTorch tensors/Paddle tensors.

Spatially sparse data structures

High-resolution 2D/3D grids are often used in large-scale spatial computation, such as physical simulation, rendering, and 3D reconstruction. However, dense data structures often consume a huge amount of memory space and computation. While a programmer may allocate a large dense grid to store spatial data, they often only care about a small fraction of this dense grid since the rest may be empty space (vacuum or air). Taichi provides spatially sparse data structures to avoid wasted computation for empty space. Users can access a spatially sparse data structure in the same way they do with a dense data structure. This feature decouples computation and data structures, allowing users to experiment with different data arrangements and tailor high-performance data structures for specific computational tasks.

As spatially sparse programming imposes analysis and optimization complexities, Taichi provides a state-flow formulation for performance analysis and an asynchronous execution engine for optimization. The optimized Taichi kernels are then just-in-time compiled in parallel, and dispatched to parallel devices such as multithreaded CPU and massively parallel GPUs.

SNode system

Taichi adopts a Structural Node (SNode) tree system, which is the key to decoupling computation from data structures. SNodes can be used to compose a hierarchical multi-dimensional field (Taichi's data container). The system allows users to experiment with different memory layouts without rewriting code.

GUI & GGUI

Taichi has a built-in GUI system for visualizing simulation data in data containers like Taichi fields or NumPy ndarrays. It also has limited support for drawing primitive geometries.

Starting from v0.8.0, Taichi has added a new UI system GGUI. The name was derived from the fact that the system uses GPU for rendering, making it faster to render 3D scenes.

Objective data-oriented programming

Taichi adopts a data-oriented programming (DOP) scheme that separates computation from data layouts so that users can experiment with different layouts for memory access optimization without rewriting the code. Considering that DOP makes modularization difficult, Taichi borrows some concepts from the object-oriented programming (OOP) design and calls the hybrid scheme objective data-oriented programming (ODOP).

ODOP allows users to organize data and methods in a class and call the methods to manipulate the data in the Taichi scope.

Quantization

High-resolution simulations can deliver great visual quality, but are often limited by the capacity of the onboard memory, GPU memory in particular. To reduce the memory footprint of programs, Taichi provides quantized data types. This feature allows users to store data of arbitrarily low precision bits without changing the code for computation. In some cases, Taichi's quantized data types can reduce the use of GPU memory by up to 8x.

Autodiff

Differentiable programming is an important advanced feature of Taichi (DiffTaichi). Taichi supports two modes of automatic differentiation (autodiff): reverse mode and forward mode.

When implementing reverse-mode autodiff, Taichi generates gradient kernels at compile time with the built-in Source Code Transformation (SCT) mechanism and then uses the lightweight tape in the Python scope to record the launched kernels and replay them in reverse order. Since v.1.1.0, Taichi has enabled forward-mode autodiff, which computes Jacobian-vector product (Jvp) when evaluating derivatives. Forward-mode autodiff is more efficient than reverse mode if the number of a function's outputs is greater than its inputs.

Efficient mesh-based operations

Mesh-based operations are usually slow due to unstructured memory access patterns. Taichi designs a compiler MeshTaichi to handle such operations in a more efficient way. The compiler provides an intuitive programming model that allows users to write mesh-based operations using reference-style neighborhood queries. It also decouples optimization from computation, making it possible to explore different localized data attributes and different memory orderings without modification of the computation code.

Ahead-of-time (AOT) module

Though choosing Python as the frontend, Taichi is not subject to the restrictions of Python's heavy virtual machine design when it comes to deployment. To facilitate deployment in real-life industrial scenarios, Taichi provides a built-in ahead-of-time (AOT) system for exporting Python code as binary/shader files, which can then be called in C/C++ and run without the Python environment.

In August 2022, a series of 3D physical simulation wallpapers co-developed by OPPO and Taichi Lang was officially launched, and Taichi AOT provides the key technical support behind.

Applications

Accelerate Python

Taichi differs from Python in the following aspects:

• Taichi is a statically-type language, while Python is dynamic.
• Taichi is a compiled language, while Python is interpreted.
• Taichi can automatically execute your code in parallel, while Python is usually single-threaded.
• Taichi can run on both CPU and GPU, while Python runs on CPU.

Taichi can take over compute-intensive tasks from Python with two decorators @ti.func and @ti.kernel. Its JIT compiler compiles the decorated functions to machine code, and all subsequent calls to these functions are executed on multi-core CPUs or GPUs.

Accelerate PyTorch

PyTorch is an open-source deep learning framework. Taichi can complement PyTorch in data preprocessing and ML operator customization in that it provides finer control over parallelization and enables element-level operations.

Physical simulation

A typical example showing Taichi's advantages in physical simulation is Taichi Elements, which is a high-performance multi-material continuum physics engine applicable to multi-threaded CPUs and massively parallel GPUs.

In Fall 2021, the Computer Graphics Laboratory (CGL) at ETH Zurich offered a course on physically-based simulation, and most of the student groups chose Taichi for real-time physical simulation.

Differentiable programming in Taichi also optimizes neural network controllers efficiently with brute-force gradient descent, significantly boosting the performance and productivity of differentiable physical simulators. Given that the simulation process normally does not provide gradients for planning and control optimizations, a differentiable physics benchmark, PasticineLab, was devised, which supports differentiable elastic and plastic deformation using the DiffTaichi system.

Image processing

Python is widely used for image processing due to fast algorithm prototyping and verification. OpenCV, a popular image processing library, also provides a Python encapsulation so that users can write Python code to call underlying C++ algorithms and get performance gains. However, some image processing algorithms are not supported by OpenCV, and iterating over pixels in Python can be slow. Taichi can complement OpenCV and accelerate image processing in Python based on its automatic parallelization mechanism and ability to switch between CPU and GPU execution.

Computational fluid dynamics (CFD)

CFD is a branch of fluid mechanics that endeavors to precisely reproduce the behavior of liquid/gas flows and their interaction with solid boundaries. The most compute-intensive step in CFD is the discretization of differential equations, which involves massive operations of 2D or 3D data arrays. Taichi's data container makes the representation of a flow field intuitive, and the automatic parallelization mechanism saves computational time.

Rendering

Taichi is also used for rendering, including path tracing, because of its cross-platform compatibility. For example, bsavery implemented Ray Tracing in One Weekend using Python and Taichi

References

This article "Taichi Lang" is from Wikipedia. The list of its authors can be seen in its historical and/or the page Edithistory:Taichi Lang. Articles copied from Draft Namespace on Wikipedia could be seen on the Draft Namespace of Wikipedia and not main one.

---