[refactor] Remove offloaded parameter of Program::compile() (#7045)

Issue: #7002
taichi-dev · Jan 5, 2023 · 875c14a · 875c14a
1 parent a7b7a2f
commit 875c14a
Show file tree

Hide file tree

Showing 16 changed files with 18 additions and 26 deletions.
diff --git a/taichi/codegen/cc/cc_program.cpp b/taichi/codegen/cc/cc_program.cpp
@@ -15,7 +15,7 @@ CCProgramImpl::CCProgramImpl(CompileConfig &config) : ProgramImpl(config) {
   context_ = std::make_unique<CCContext>();
 }
 
-FunctionType CCProgramImpl::compile(Kernel *kernel, OffloadedStmt *) {
+FunctionType CCProgramImpl::compile(Kernel *kernel) {
   CCKernelGen codegen(kernel, this);
   auto ker = codegen.compile();
   auto ker_ptr = ker.get();

diff --git a/taichi/codegen/cc/cc_program.h b/taichi/codegen/cc/cc_program.h
@@ -29,7 +29,7 @@ class CCProgramImpl : public ProgramImpl {
  public:
   explicit CCProgramImpl(CompileConfig &config);
 
-  FunctionType compile(Kernel *kernel, OffloadedStmt *) override;
+  FunctionType compile(Kernel *kernel) override;
 
   std::size_t get_snode_num_dynamically_allocated(
       SNode *snode,

diff --git a/taichi/program/program.cpp b/taichi/program/program.cpp
@@ -191,10 +191,10 @@ Function *Program::create_function(const FunctionKey &func_key) {
   return functions_.back().get();
 }
 
-FunctionType Program::compile(Kernel &kernel, OffloadedStmt *offloaded) {
+FunctionType Program::compile(Kernel &kernel) {
   auto start_t = Time::get_time();
   TI_AUTO_PROF;
-  auto ret = program_impl_->compile(&kernel, offloaded);
+  auto ret = program_impl_->compile(&kernel);
   TI_ASSERT(ret);
   total_compilation_time_ += Time::get_time() - start_t;
   return ret;

diff --git a/taichi/program/program.h b/taichi/program/program.h
@@ -208,8 +208,7 @@ class TI_DLL_EXPORT Program {
 
   // TODO: This function is doing two things: 1) compiling CHI IR, and 2)
   // offloading them to each backend. We should probably separate the logic?
-  // TODO(Lin): remove the offloaded parameter
-  FunctionType compile(Kernel &kernel, OffloadedStmt *offloaded = nullptr);
+  FunctionType compile(Kernel &kernel);
 
   void check_runtime_error();
 

diff --git a/taichi/program/program_impl.h b/taichi/program/program_impl.h
@@ -35,7 +35,7 @@ class ProgramImpl {
   /**
    * Codegen to specific backend
    */
-  virtual FunctionType compile(Kernel *kernel, OffloadedStmt *offloaded) = 0;
+  virtual FunctionType compile(Kernel *kernel) = 0;
 
   /**
    * Allocate runtime buffer, e.g result_buffer or backend specific runtime

diff --git a/taichi/runtime/program_impls/dx/dx_program.cpp b/taichi/runtime/program_impls/dx/dx_program.cpp
@@ -26,8 +26,7 @@ FunctionType compile_to_executable(Kernel *kernel,
 Dx11ProgramImpl::Dx11ProgramImpl(CompileConfig &config) : ProgramImpl(config) {
 }
 
-FunctionType Dx11ProgramImpl::compile(Kernel *kernel,
-                                      OffloadedStmt *offloaded) {
+FunctionType Dx11ProgramImpl::compile(Kernel *kernel) {
   spirv::lower(kernel);
   return directx11::compile_to_executable(kernel, runtime_.get(),
                                           snode_tree_mgr_.get());

diff --git a/taichi/runtime/program_impls/dx/dx_program.h b/taichi/runtime/program_impls/dx/dx_program.h
@@ -11,7 +11,7 @@ namespace taichi::lang {
 class Dx11ProgramImpl : public ProgramImpl {
  public:
   Dx11ProgramImpl(CompileConfig &config);
-  FunctionType compile(Kernel *kernel, OffloadedStmt *offloaded) override;
+  FunctionType compile(Kernel *kernel) override;
 
   std::size_t get_snode_num_dynamically_allocated(
       SNode *snode,

diff --git a/taichi/runtime/program_impls/llvm/llvm_program.cpp b/taichi/runtime/program_impls/llvm/llvm_program.cpp
@@ -35,9 +35,8 @@ LlvmProgramImpl::LlvmProgramImpl(CompileConfig &config_,
   }
 }
 
-FunctionType LlvmProgramImpl::compile(Kernel *kernel,
-                                      OffloadedStmt *offloaded) {
-  auto codegen = KernelCodeGen::create(kernel->arch, kernel, offloaded);
+FunctionType LlvmProgramImpl::compile(Kernel *kernel) {
+  auto codegen = KernelCodeGen::create(kernel->arch, kernel);
   return codegen->compile_to_function();
 }
 

diff --git a/taichi/runtime/program_impls/llvm/llvm_program.h b/taichi/runtime/program_impls/llvm/llvm_program.h
@@ -40,7 +40,7 @@ class LlvmProgramImpl : public ProgramImpl {
 
   // TODO(zhanlue): compile-time runtime split for LLVM::CodeGen
   // For now, compile = codegen + convert
-  FunctionType compile(Kernel *kernel, OffloadedStmt *offloaded) override;
+  FunctionType compile(Kernel *kernel) override;
 
   void compile_snode_tree_types(SNodeTree *tree) override;
 

diff --git a/taichi/runtime/program_impls/metal/metal_program.cpp b/taichi/runtime/program_impls/metal/metal_program.cpp
@@ -45,9 +45,7 @@ MetalProgramImpl::MetalProgramImpl(CompileConfig &config_)
     : ProgramImpl(config_) {
 }
 
-FunctionType MetalProgramImpl::compile(Kernel *kernel,
-                                       OffloadedStmt *offloaded) {
-  TI_ASSERT(offloaded == nullptr);
+FunctionType MetalProgramImpl::compile(Kernel *kernel) {
   return metal::compiled_kernel_to_metal_executable(
       get_cache_manager()->load_or_compile(config, kernel),
       metal_kernel_mgr_.get());

diff --git a/taichi/runtime/program_impls/metal/metal_program.h b/taichi/runtime/program_impls/metal/metal_program.h
@@ -20,7 +20,7 @@ class MetalProgramImpl : public ProgramImpl {
  public:
   explicit MetalProgramImpl(CompileConfig &config);
 
-  FunctionType compile(Kernel *kernel, OffloadedStmt *offloaded) override;
+  FunctionType compile(Kernel *kernel) override;
 
   std::size_t get_snode_num_dynamically_allocated(
       SNode *snode,

diff --git a/taichi/runtime/program_impls/opengl/opengl_program.cpp b/taichi/runtime/program_impls/opengl/opengl_program.cpp
@@ -24,8 +24,7 @@ OpenglProgramImpl::OpenglProgramImpl(CompileConfig &config)
     : ProgramImpl(config) {
 }
 
-FunctionType OpenglProgramImpl::compile(Kernel *kernel,
-                                        OffloadedStmt *offloaded) {
+FunctionType OpenglProgramImpl::compile(Kernel *kernel) {
   return register_params_to_executable(
       get_cache_manager()->load_or_compile(config, kernel), runtime_.get());
 }

diff --git a/taichi/runtime/program_impls/opengl/opengl_program.h b/taichi/runtime/program_impls/opengl/opengl_program.h
@@ -11,8 +11,7 @@ class OpenglProgramImpl : public ProgramImpl {
  public:
   explicit OpenglProgramImpl(CompileConfig &config);
   ~OpenglProgramImpl() override;
-
-  FunctionType compile(Kernel *kernel, OffloadedStmt *offloaded) override;
+  FunctionType compile(Kernel *kernel) override;
 
   std::size_t get_snode_num_dynamically_allocated(
       SNode *snode,

diff --git a/taichi/runtime/program_impls/vulkan/vulkan_program.cpp b/taichi/runtime/program_impls/vulkan/vulkan_program.cpp
@@ -78,8 +78,7 @@ VulkanProgramImpl::VulkanProgramImpl(CompileConfig &config)
     : ProgramImpl(config) {
 }
 
-FunctionType VulkanProgramImpl::compile(Kernel *kernel,
-                                        OffloadedStmt *offloaded) {
+FunctionType VulkanProgramImpl::compile(Kernel *kernel) {
   return register_params_to_executable(
       get_cache_manager()->load_or_compile(config, kernel),
       vulkan_runtime_.get());

diff --git a/taichi/runtime/program_impls/vulkan/vulkan_program.h b/taichi/runtime/program_impls/vulkan/vulkan_program.h
@@ -31,7 +31,7 @@ class VulkanDeviceCreator;
 class VulkanProgramImpl : public ProgramImpl {
  public:
   explicit VulkanProgramImpl(CompileConfig &config);
-  FunctionType compile(Kernel *kernel, OffloadedStmt *offloaded) override;
+  FunctionType compile(Kernel *kernel) override;
 
   std::size_t get_snode_num_dynamically_allocated(
       SNode *snode,

diff --git a/tests/cpp/backends/dx11_device_test.cpp b/tests/cpp/backends/dx11_device_test.cpp
@@ -151,7 +151,7 @@ TEST(Dx11ProgramTest, MaterializeRuntimeTest) {
   auto block = builder.extract_ir();
   test_prog.prog()->this_thread_config().arch = Arch::dx11;
   auto ker = std::make_unique<Kernel>(*test_prog.prog(), std::move(block));
-  program->compile(ker.get(), nullptr);
+  program->compile(ker.get());
 }
 
 }  // namespace directx11