search

taichi-dev / taichi

Productive, portable, and performant GPU programming in Python.
https://taichi-lang.org
Apache License 2.0
25.58k stars 2.29k forks source link

Ahead-of-time (AOT) Feature Roadmap #4615

Open k-ye opened 2 years ago

k-ye commented 2 years ago

Untitled Diagram (1)

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:

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

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

  1. Quantized type formalization: For saving memory resource and atomic acceleration https://github.com/taichi-dev/taichi/issues/4857
  2. LLVM AOT implementation https://github.com/taichi-dev/taichi/issues/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. https://github.com/taichi-dev/taichi/issues/4786
  5. Texture types (e.g. OpenGL texture, imageND, ...) https://github.com/taichi-dev/taichi/issues/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. https://github.com/taichi-dev/taichi/issues/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. https://github.com/taichi-dev/taichi/issues/4808
k-ye commented 2 years ago

Some updates about our immediate efforts:

  1. Clean up Taichi's C++ codebase, make it more modular and easily linkable with other C++ projects: https://github.com/taichi-dev/taichi/issues/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 https://github.com/taichi-dev/taichi/issues/4786#issuecomment-1127214195. TLDR, we will start with supporting sequential kernel invocations. This is led by @ailzhang , and would be tested by @YuCrazing .
  3. LLVM AOT: https://github.com/taichi-dev/taichi/issues/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.