Taichi: A refactoring plan

[masahi]

See the comment below for the context.
https://github.com/taichi-dev/taichi/issues/412#issuecomment-582147182

Starting from the easy ones, I'll add more items as I learn more about the codebase. Feel free to add what you find as well.

• [x] Remove SNodeAllocator in struct.h as it is not used anywhere
• [x] Clarify remove_useless_libdevice_functions(...) in taichi_llvm_context.cpp
• [x] Remove the legacy cuda source backend to avoid confusion
• [x] Clean up compiler command stuff in tlang_util.cpp
• [x] Remove KernelCodeGen::codegen()
• [x] CodeGenBase is the base class of StructCompiler and KernelCodeGen, but it doesn't seem useful. Consider restructuring the class hierarchy or just remove it
• [x] Remove cuda_context stuff from Program https://github.com/taichi-dev/taichi/blob/master/taichi/program.cpp#L36-L43
• [x] Clean up the cuda_context global var (also marked as "TODO")
  https://github.com/taichi-dev/taichi/blob/master/taichi/backends/llvm_jit_ptx.cpp#L24
  https://github.com/taichi-dev/taichi/blob/master/taichi/backends/cuda_context.h#L72
• [x] Move compile_module_to_ptx from llvm_jit_ptx.cpp to codegen_llvm_ptx.cpp. Rename the rest of llvm_jit_ptx.cpp to cuda_context.cpp
• [x] LoopGenerator is not used and it generates c or cuda source, so probably this is legacy code. Also loopgen in StructCompiler can be removed. https://github.com/taichi-dev/taichi/blob/master/taichi/backends/struct.h#L18
• [x] Remove generate_types and generate_leaf_accessors from StructCompiler https://github.com/taichi-dev/taichi/blob/master/taichi/backends/struct.cpp#L122-L130
• [x] Rename pointer_activate etc into Pointer_activate
• [x] Remove Stmt::adjoint and Stmt::value.
• [x] Use std::unique_ptr instead of std::shared_ptr in SNode for clear children ownership.

[yuanming-hu]

Thanks for proposing these changes!

Remove SNodeAllocator in struct.h as it is not used anywhere

That's a good point. I think the whole struct.h can be removed sometime. Currently it's just used as a reference of the old data structure implementations. Some data structures in the old source2source backend are not yet implemented in the LLVM backend, so I would keep this file. I have moved it to misc for now.

Remove remove_useless_libdevice_functions(...) in taichi_llvm_context.cpp and use llvm::GlobalValue::AvailableExternallyLinkage for linkage instead

I know this function looks messy. However, its presence deletes a bunch of huge CUDA runtime functions, immediately after libdevice.bc is loaded. This prevents us from further wasting any time from them. This hacky optimization saves ~0.3s (if I remember correctly) when linking libdevice against our modules.

Remove the legacy cuda source backend to avoid confusion

I removed the legacy CUDA source backend a few hours ago :-)

[yuanming-hu]

Meanwhile, I'm also actively working on removing the dependency on CUDA unified memory to make way for AMDGPU backend.

[masahi]

I know this function looks messy. However, its presence deletes a bunch of huge CUDA runtime functions, immediately after libdevice.bc is loaded. This prevents us from further wasting any time from them. This hacky optimization saves ~0.3s (if I remember correctly) when linking libdevice against our modules.

I think even with "remove_useless_libdevice_functions" there are many useless functions in libdevice that will become part of the final module. I implemented libdevice linking in TVM, and I remember using AvailableExternallyLinkage it removed all __nv_ intrinsics that are not used in the llvm module tvm generates before linking. To be sure I need to see a taichi llvm module dump before it gets sent to codegen.

[yuanming-hu]

Yeah, ultimately all the unused nv intrinsics will be eliminated during the PTX emission process. I deleted the huge libdevice functions just to accelerate this single linking step (which is even before we can analyze which nv intrinsics are used):
https://github.com/taichi-dev/taichi/blob/794e923c1370aa74073cc63ccbb4bbfd6ae14900/taichi/taichi_llvm_context.cpp#L344

I could be wrong though, and I'm definitely open to any better designs :-)

[masahi]

do we still need this function now that legacy codegen is removed? Or is it used for compiling runtime.cpp?
https://github.com/taichi-dev/taichi/blob/master/taichi/tlang_util.cpp#L243-L290

[yuanming-hu]

No, it's no longer used. I considered removing this, however, there is still a chance that we need a C backend in the future...

[masahi]

we can at least remove nvcc bits.

[yuanming-hu]

Yeah, I agree. I think the above discussions are all centered around a question: Does "demonstrating how a source2source compiler should be implemented in the future" justify the existence of certain legacy code?

[masahi]

I think such legacy code can exist but should be separated from the rest of the code that is actually used. Current codebase has both legacy and current code and for newcomers it is not clear which is which.

[yuanming-hu]

Yeah, that sounds like a good solution. How about just moving these to misc?

I agree that having legacy code can be very confusing.

[masahi]

misc is not a part of build right (it has some python scripts too)? I think it is ok.

[yuanming-hu]

Yeah, that's just for holding nonessential random and experimental files. Thanks for the inputs. I'll gradually move the legacy code there for code clarity :-)

[masahi]

regarding libdevice linking, is it possible to ship custom libdevice.bc in the taichi repo? It seems Halide does this https://github.com/halide/Halide/tree/master/src/runtime/nvidia_libdevice_bitcode. I'm not sure if there is a license or other issues. It will also enable removing libdevice.bc path finding.

[yuanming-hu]

Yeah, that sounds good! I think now we ship two bc files (runtime_cuda.bc and libdevice.bc), and they should be merged with the huge nv intrinsics removed. Not sure about licenses though. I can confirm with the Halide people (e.g. Andrew Adams) to learn about what he thinks about the libdevice license issues.

[yuanming-hu]

@masahi I've done several rounds of aggressive refactorings to improve code readability/maintainability. Most improvements are not actually listed in this list. Now I'll first rip off unified memory dependency and then come back to the remaining refactoring items in your list. Please feel free to propose new refactoring opportunities in the new codebase!

(@k-ye @archibate my apologies in advance if I have broken any code of you guys...)

[masahi]

@yuanming-hu thanks, I also want to start contributing code changes, the only problem being finding time to work on it.