taichi
taichi copied to clipboard
Link modules instead of cloning modules in the codegen in LLVM backends
Concisely describe the proposed feature As described in #5282, the compilation of Taichi kernel becomes slow when there are many SNodeTrees. After profiling, we found that the process of cloning the modules take up most of the time. The cloning happen in three places. Two of them happen in the construction process of the struct module as described in #3146. https://github.com/taichi-dev/taichi/blob/77398de87e635e4de583de808137e8efe484e8a9/taichi/llvm/llvm_program.cpp#L138-L145 https://github.com/taichi-dev/taichi/blob/77398de87e635e4de583de808137e8efe484e8a9/taichi/llvm/llvm_context.cpp#L455 The other one happens when the LLVM codegen gets the struct module as the initial module of kernel module. https://github.com/taichi-dev/taichi/blob/d00348c1acccda88c1823e60e151a78b1c4f1b8a/taichi/codegen/llvm/codegen_llvm.cpp#L318-L322
Describe the solution you'd like (if any) We can store the runtime functions, the struct functions of the SNodeTrees, and the kernel functions in different modules, and link the needed modules at the end. I plan to do the split as below:
- [ ] Split the runtime module out of the kernel module (keep the struct functions in the kernel module)
- [ ] Split the struct module out of the kernel module (the struct functions of all the SNodeTrees are stored in a single struct module)
- [ ] Let struct functions of different SNodeTrees store in different modules.