Support JIT Offline Cache for Taichi

[PGZXB]

Solution

Workflow (on llvm backends)

... → Python Code → Taichi AST → Trans AST to string as key → Hash it(hashed key) → Try find offline cache file by hashed-key:

• Found: Load cache data from disk → Create kernel form cache → Run kernel
• Not Found: ( Continue to compile ) ... → Get llvm::Module + offloaded_task_list → Cache them → Run kernel → ... → ( Before exiting ) Dump cache data to disk

Todo & Memo

• Support for cpu
• Support for cuda
• Add ASTKeyGenerator to generator key of Taichi AST instead of IRPrinter, which will holds more information compared with IRPrinter.
• Fix bugs that some global-vars' changes will not cause re-compiling (Maybe let result of IRPrinter and Expression::serialize hold more information).
• Fix IRPrinter to generate offline-cache-key more correctly
• Add tests
• Consider compile-config's change
• Trace useless cache-files and delete them. Current implementation causes "cache-file leak"
• Impl binary ticache file format
• Run on multi thread/process
• Support on vulkan
• Support on opengl
• Support on metal
• Refactor (see Support JIT Offline Cache for Taichi #4401 (comment)). see Refactor kernel compilation #7002

- [ ] Support on dx11
- [ ] Support on dx12
- [ ] Handle hash collisions
- [ ] Allow to set/unset offline_cache per kernel ( Optional

Usage

Just set offline_cache=True The feature is enabled by default.

import taichi as ti

# ti.init(arch=ti.cpu, offline_cache=True)
ti.init(arch=ti.cpu)

@kernel
def f():
    print("Hello ticache")

f()

Supported backends

See #4401 (comment)

For more, see Offline Cache

Potential Bugs

• The offline cache doesn't work properly on the program that contains indirect recursive calls to real functions #6877
• test_print.py with offline_cache=True (https://github.com/taichi-dev/taichi/actions/runs/6176139423/job/16764496768)

[bobcao3]

Requesting for extending this to all back ends, considering the huge range of users on Mac or non-cuda laptops

[PGZXB]

Requesting for extending this to all back ends, considering the huge range of users on Mac or non-cuda laptops

Yes, impl the feature for all backends is my goal.

[PGZXB]

#4401 (comment)

[bobcao3]

I argue strongly against this solution. We have profiles showing the LLVM codegen takes about 30% of the entire JIT codegen time, it would be much wiser to spend time figuring out AST->CHI-IR caching first.

[bobcao3]

A two staged caching gives a major JIT time boost to all backends, I'd argue this is a lot cleaner to implement as well compared to having one stage caching for each individual backend, which will cause maintenance problems down the line

[PGZXB]

A two staged caching gives a major JIT time boost to all backends, I'd argue this is a lot cleaner to implement as well compared to having one stage caching for each individual backend, which will cause maintenance problems down the line
I argue strongly against this solution. We have profiles showing the LLVM codegen takes about 30% of the entire JIT codegen time, it would be much wiser to spend time figuring out AST->CHI-IR caching first.

Step by step. Maybe temporary solution. We don't have serialization of CHI IR now. After CHI IR's serialization is implemented, maybe 2-level cache is better, especially for multi backends...

ps. I think CHI IR's serialization is very important for standardizing CHI IR, which needs a feasible efficient standard (more .adj to show the importance I think) solution, like llvm-ir, IL, Java bytecode, intel-asm, which is not easy...

[k-ye]

Is there a middle ground we can find out? E.g. how easy is it for us to migrate the implementation from caching LLVM to caching CHI IR? If most users don't care about the internal implementation of the cache, I expect the following scenario to happen:

1. At first, they can only benefit from the caching behavior for CUDA/CPU backends
2. Then after release X, they find out the caching is working for all the backends automatically.

In addition, IMHO the complexity still comes from the cache key part (considering all the involved global states). The cached contents can be adjusted fairly easily, provided that CHI IR serialization is implemented.

[PGZXB]

Is there a middle ground we can find out? E.g. how easy is it for us to migrate the implementation from caching LLVM to caching CHI IR? If most users don't care about the internal implementation of the cache, I expect the following scenario to happen:

1. At first, they can only benefit from the caching behavior for CUDA/CPU backends
2. Then after release X, they find out the caching is working for all the backends automatically.

The (new) implementation of offline-cache is transparent. All logic is in C++ side. Frontend only see the offline_cache: bool and offline_cache_file_path: str options. If we have serialization and deserialization of CHI IR, implementing caching CHI IR will be simple. Maybe doing it after standardizing CHI IR is better. After release X, users can also use it by simply set options without any migration cost. And maybe multilevel cache is optional(even better) solution, running backend lang directly is fastest.

[PGZXB]

In addition, IMHO the complexity still comes from the cache key part (considering all the involved global states). The cached contents can be adjusted fairly easily, provided that CHI IR serialization is implemented.

Can't agree more. Because taichi's kernels depends on global vars/states, generating a key which can uniquely identifies a kernel is difficult and the key of implementing caching a kernel. And, at present, before we have a standardized de/serializable CHI IR, dumping and loading & running backend-language is more simple than CHI IR because they have mature/standard solution.

ps. Overhead of generating key is what we should consider. Python -> Taichi AST -> CHI IR -> Backend lang. From left to right:

• overhead of generating key ↑ ,
• overhead of loading & running offline-cache-file ↓ ,
• difficulty of generating cache key which can uniquely identifies a kernel ↓

[PGZXB]

Supported or not

Backend	Supported or not	Overhead (running Cornell Box)
CPU	✔	393.25ms
CUDA	✔	882.426ms
Vulkan	✔	218.030ms
OpenGL	✔
Metal	✔
AMDGPU	✔
Microsoft DirectX 11	✔
Microsoft DirectX 12	❌	N/A

P.S.

1. The "overhead" is the time spent on loading cached compiled data and converting it to a callable object.
2. Testing environment:
   • OS: Windows 11, CPU: Intel(R) Core(TM) i7-10710U CPU @ 1.10GHz 1.61 GHz, RAM: 16GB for CPU, CUDA, OpenGL and Vulkan
3. ⏩: Working in progress