Refactor kernel compilation

[PGZXB]

(This is part of Refactor kernel compilation and launch, which aims to split kernel compilation and kernel launching totally.)

Brief Introduction

1. Isolate the stuff which is not related with kernel compilation from lang::Kernel.
2. lang::CompileConfig per compilation instead of lang::CompileConfig per lang::Program.
3. Introduce lang::CompiledKernelData, lang::KernelCompiler to unify the compilation result and the interfaces of kernel compilation.
4. Introduce lang::KernelCompilationManager to unify the implementation of cache on ALL backends, including online cache and offline cache
5. Re-impl the JIT in Python frontend and remove the online cache in frontend

Prototype

• An incomplete implementation of this proposal: see [refactor] Refactor kernel compilation (1st prototype, WIP, Draft, Don't merge) #6819 or PGZXB:dev-refactor-kernelCR-1st-prototype

Provide @ti.pkernel

NOTE: The @ti.pkernel is only an attempt based on the prototype in order to verify the effect of the proposal, which is NOT a "feature request".

• @ti.pkernel = @ti.kernel + CompileConfig, which allows the user to specify different compile configuration options for each kernel.

Implementation

• See PGZXB:dev-pkernel-based_on-refactorKCaL_1st_prototype

Example

import taichi as ti

ti.init(arch=ti.vulkan, offline_cache=True, print_preprocessed_ir=False)

@ti.kernel
def K1(a: ti.i32, b: ti.i32) -> ti.i32:
  return a + b


@ti.pkernel(offline_cache=False, print_preprocessed_ir=True)
def PK1(a: ti.i32, b: ti.i32) -> ti.i32:
  return a * b

# Not allowed before finishing "Refactor kernel launch"
# @ti.pkernel(arch=ti.cpu) # ti.cpu != ti.vulkan
# def PK2(a: ti.i32, b: ti.i32) -> ti.i32:
#   return a * b

K1(10, 10)
PK1(10, 10)
# PK2(10, 100)

To-Do List

• Remove dependencies on Program::current_ast_builder() in C++ side, see prototype: 73e5b71
• Remove Program::current_ast_builder(), see prototype: 0b30140
• Remove Program::current_callable, see prototype: f6d5a55; fdd230f
• Remove Kernel::arch
• Remove Kernel::lower & Add irpass::ast_to_ir, see prototype: 6a83c69; a77e46a
• Remove CompileConfig::lazy_compilation, see prototype: b82cd7d
• Move Kernel::compile and Kernel::operator() outside, see prototype: b5d6beb
• Rename arg to param in class Callable, see prototype: b943d3d
• Refactor Kernel::get_ret*, see prototype: d87a8ba; 80bea89; 732b4a3; 0920501; 0f2a0b5; 7eb13d5
• Remove dependencies on Program::this_thread_config() in metal::CacheManager, see prototype: e595513
• Remove offloaded parameter of Program*::compile, see prototype: 582f287
• Remove ir parameter in KernelCodeGen::KernelCodeGen(Kernel *kernel, IRNode *ir), see prototype: a2f1e15
• Remove dependencies on Program::this_thread_config() in KernelCodeGen, see prototype: 187f206
• Remove dependencies on Program::this_thread_config() in codgen on llvm backends, see prototype: 423a29c; 8ecf111; b8deaf0
• Remove dependencies on Program::this_thread_config() in irpass::scalarize, see prototype: 40d4b33
• Remove dependencies on Program::this_thread_config() in gfx/runtime.cpp, see prototype: 2120e43; 53d9e8e
• Remove dependencies on Program::this_thread_config() in spirv::lower, see prototype: 79b4b9a
• Remove dependencies on Program::this_thread_config() in codegen_cc.cpp, see prototype: 19d6e11
• Remove IRNode::get_config, IRNode::get_kernel and IRNode::kernel, see prototype: 592f798; 72b23df
• Remove dependencies on Program::this_thread_config() in lang::launch_kernel(), see prototype: 1f958a3; f9de0fe
• Remove dependencies on Program::this_thread_config() in constant folding & rm possible race condition, see prototype: 2571bf9; 3164640
• Remove Program::this_thread_config() & Add Program::global_compile_compileconfig(), see prototype: abcff07
• Split lang::Function, see prototype: d87e6d9
• Implement irpass::clone for AST
• Introduce CompiledKernelData & KernelCompiler & KernelCompilationManager, impl them for spirv/gfx backends and re-impl VulkanProgramImpl::compile() and gfx::AotModuleBuilder, see prototype: 40e9751; c8e6285; 0697032
• compile - runtime split for TaichiLLVMContext #7311
• Intorduce CompiledKernelData & KernelCompiler for LLVM backends (Hard)
• Intorduce CompiledKernelData & KernelCompiler for Metal backends
• Fix implementation of JIT in Python-frontend: t_kernel(ctx) -> load_or_compile(...) && prog.launch_kernel(...)
• Remove online cache in Python side (PyTaichi.compiled_functions)
• Redesign the algorithm to generate the offline cache key for real_funcs (to fix The offline cache doesn't work properly on the program that contains indirect recursive calls to real functions #6877)
• Remove lang::Callable::program #7286
• Fix all const CompileConfig * to const CompileConfig &`
• Remove Kernel::from_cache_
• Remove Kernel::task_counter_
• Remove useless Kernel::lowered_

Related work

• Maybe start replacing pybind11 with XXXXXX from the APIs about kernel compilation and launch. Some of them will be refactored and standardized in this proposal.