Ahead-of-time (AOT) Feature Roadmap

[k-ye]

I'd like to share our vision and the initial roadmap for the AOT feature. AOT is our solution to make Taichi kernels and programs runnable in non-Python environments. Mobile and game engines, for example.

The ongoing efforts include:

• defining the base AOT API
• implementing this AOT API for the Vulkan backend

The API design, its features and limitations are all documented in

• [RFC] Taichi's Ahead of Time (AOT) Module #3642

---

Starting from here, we can immediately explore the following directions:

1. Quantized type formalization: For saving memory resource and atomic acceleration Polish and officially release quantized types #4857
2. LLVM AOT implementation AOT for the LLVM backend #4800
3. SPIR-V based codegen: Cross compiling the SPIR-V shader to other platforms, including OpenGL, Apple Metal (OpenCL?)
4. AOT compute graph: This would significantly reduce the efforts to porting a demo written in Taichi to non-Python environments. In essence, the AOT module will save not only the Taichi kernels, but also the host logic invoking them. [AOT] Serialize compute graph in python and load in C++ runtime #4786
5. Texture types (e.g. OpenGL texture, imageND, ...) [lang] Texture support #5109

---

1. The LLVM AOT impl would allow us to understand how to support the more advanced sparse SNodes, e.g. pointer.
2. The SPIR-V codegen would also need to support these sparse SNodes. When the SPIR-V cross-compilation is ready, all SPIR-V-based backends will gain sparsity support for free.
3. LLVM + SPIR-V means that the AOT solution can finally cover all the backends in Taichi.
4. Support all kinds of SNode types, making Taichi's field system feature complete in AOT. Support all kinds of SNodes in AOT #4777
5. Supporting texture means we can better support signed-distance-field (SDF) and plug in more customized geometries.
6. AOT solution would help us more accurately benchmark the kernel performance. Meanwhile, we can make quant types tunable to find better type layout.
7. Integrate Taichi AOT with other platforms or game engines, Unity, for example. [AOT] Integrating Taichi AOT with Unity #4808

[k-ye]

Some updates about our immediate efforts:

1. Clean up Taichi's C++ codebase, make it more modular and easily linkable with other C++ projects: Taichi CMake Overhaul #4832. @PENGUINLIONG is working on the Windows part, while @qiao-bo (and possibly @turbo0628) is working on the Linux part.
   a. This would help integrating Taichi with Unity, or any other kind of C/C++ projects. In the end, it's just a nicely encapsulated CPP library.
2. Compute graph to make Taichi scripts easily deployable: see this update [AOT] Serialize compute graph in python and load in C++ runtime #4786 (comment). TLDR, we will start with supporting sequential kernel invocations. This is led by @ailzhang , and would be tested by @YuCrazing .
3. LLVM AOT: AOT for the LLVM backend #4800. Currently @k-ye is making this work on CPU, @jim19930609 would help cover the CUDA part later.
4. Better AOT runtime API: Right now this is kind of messy. Different backends have their own runtime implementation, but are not unified. So too much implementation details will be leaked to AOT users as of now. @turbo0628 and I will find a way to refactor this.