[dx12] Only use llvm to compile dx12.

cmake/TaichiCore.cmake

@@ -232,10 +232,12 @@ if(TI_WITH_LLVM)
         add_subdirectory(taichi/rhi/dx12)
         add_subdirectory(taichi/runtime/dx12)
         add_subdirectory(taichi/codegen/dx12)
+        add_subdirectory(taichi/runtime/program_impls/dx12)
 
         target_include_directories(${CORE_LIBRARY_NAME} PRIVATE external/DirectX-Headers/include)
         target_link_libraries(${CORE_LIBRARY_NAME} PRIVATE dx12_codegen)
         target_link_libraries(${CORE_LIBRARY_NAME} PRIVATE dx12_runtime)
+        target_link_libraries(${CORE_LIBRARY_NAME} PRIVATE dx12_program_impl)
     endif()
 
     add_subdirectory(taichi/rhi/llvm)

taichi/codegen/dx12/codegen_dx12.cpp

@@ -232,8 +232,6 @@ static std::vector<uint8_t> generate_dxil_from_llvm(
 
 KernelCodeGenDX12::CompileResult KernelCodeGenDX12::compile() {
   TI_AUTO_PROF;
-  auto *llvm_prog = get_llvm_program(prog);
-  auto *tlctx = llvm_prog->get_llvm_context(kernel->arch);
   auto &config = prog->this_thread_config();
   std::string kernel_key = get_hashed_offline_cache_key(&config, kernel);
   kernel->set_kernel_key_for_cache(kernel_key);

taichi/program/program.cpp

@@ -39,6 +39,10 @@
 #include "taichi/runtime/program_impls/dx/dx_program.h"
 #include "taichi/rhi/dx/dx_api.h"
 #endif
+#ifdef TI_WITH_DX12
+#include "taichi/runtime/program_impls/dx12/dx12_program.h"
+#include "taichi/rhi/dx12/dx12_api.h"
+#endif
 
 #if defined(_M_X64) || defined(__x86_64)
 // For _MM_SET_FLUSH_ZERO_MODE
@@ -81,7 +85,17 @@ Program::Program(Arch desired_arch) : snode_rw_accessors_bank_(this) {
   profiler = make_profiler(config.arch, config.kernel_profiler);
   if (arch_uses_llvm(config.arch)) {
 #ifdef TI_WITH_LLVM
-    program_impl_ = std::make_unique<LlvmProgramImpl>(config, profiler.get());
+    if (config.arch != Arch::dx12) {
+      program_impl_ = std::make_unique<LlvmProgramImpl>(config, profiler.get());
+    } else {
+      // NOTE: use Dx12ProgramImpl to avoid using LlvmRuntimeExecutor for dx12.
+#ifdef TI_WITH_DX12
+      TI_ASSERT(directx12::is_dx12_api_available());
+      program_impl_ = std::make_unique<Dx12ProgramImpl>(config);
+#else
+      TI_ERROR("This taichi is not compiled with DX12");
+#endif
+    }
 #else
     TI_ERROR("This taichi is not compiled with LLVM");
 #endif
@@ -188,7 +202,8 @@ void Program::materialize_runtime() {
 void Program::destroy_snode_tree(SNodeTree *snode_tree) {
   TI_ASSERT(arch_uses_llvm(this_thread_config().arch) ||
             this_thread_config().arch == Arch::vulkan ||
-            this_thread_config().arch == Arch::dx11);
+            this_thread_config().arch == Arch::dx11 ||
+            this_thread_config().arch == Arch::dx12);
   program_impl_->destroy_snode_tree(snode_tree);
   free_snode_tree_ids_.push(snode_tree->id());
 }
@@ -225,7 +240,8 @@ void Program::synchronize() {
   if (arch_uses_llvm(this_thread_config().arch) ||
       this_thread_config().arch == Arch::metal ||
       this_thread_config().arch == Arch::vulkan ||
-      this_thread_config().arch == Arch::opengl) {
+      this_thread_config().arch == Arch::opengl ||
+      this_thread_config().arch == Arch::dx12) {
     program_impl_->synchronize();
   }
 }
@@ -461,7 +477,8 @@ std::size_t Program::get_snode_num_dynamically_allocated(SNode *snode) {
   TI_ASSERT(arch_uses_llvm(this_thread_config().arch) ||
             this_thread_config().arch == Arch::metal ||
             this_thread_config().arch == Arch::vulkan ||
-            this_thread_config().arch == Arch::opengl);
+            this_thread_config().arch == Arch::opengl ||
+            this_thread_config().arch == Arch::dx12);
   return program_impl_->get_snode_num_dynamically_allocated(snode,
                                                             result_buffer);
 }
@@ -531,7 +548,8 @@ std::unique_ptr<AotModuleBuilder> Program::make_aot_module_builder(Arch arch) {
   if (arch_uses_llvm(this_thread_config().arch) ||
       this_thread_config().arch == Arch::metal ||
       this_thread_config().arch == Arch::vulkan ||
-      this_thread_config().arch == Arch::opengl) {
+      this_thread_config().arch == Arch::opengl ||
+      this_thread_config().arch == Arch::dx12) {
     return program_impl_->make_aot_module_builder();
   }
   return nullptr;

taichi/runtime/program_impls/dx12/CMakeLists.txt

@@ -0,0 +1,19 @@
+# ./taichi/runtime/program_impls/dx12/CMakeLists.txt
+
+add_library(dx12_program_impl)
+target_sources(dx12_program_impl
+  PRIVATE
+    dx12_program.cpp
+  )
+
+target_include_directories(dx12_program_impl
+  PRIVATE
+    ${PROJECT_SOURCE_DIR}
+    ${PROJECT_SOURCE_DIR}/external/eigen
+    ${PROJECT_SOURCE_DIR}/external/spdlog/include
+
+    ${PROJECT_SOURCE_DIR}/external/DirectX-Headers/include
+    ${LLVM_INCLUDE_DIRS}
+  )
+
+target_link_libraries(dx12_program_impl PRIVATE dx12_rhi)

taichi/runtime/program_impls/dx12/dx12_program.cpp

@@ -0,0 +1,21 @@
+#ifdef TI_WITH_DX12
+
+#include "taichi/runtime/program_impls/dx12/dx12_program.h"
+#include "taichi/runtime/dx12/aot_module_builder_impl.h"
+#include "taichi/rhi/dx12/dx12_api.h"
+
+namespace taichi {
+namespace lang {
+
+Dx12ProgramImpl::Dx12ProgramImpl(CompileConfig &config)
+    : LlvmProgramImpl(config, nullptr) {
+}
+
+std::unique_ptr<AotModuleBuilder> Dx12ProgramImpl::make_aot_module_builder() {
+  return std::make_unique<directx12::AotModuleBuilderImpl>(this);
+}
+
+}  // namespace lang
+}  // namespace taichi
+
+#endif

taichi/runtime/program_impls/dx12/dx12_program.h

@@ -0,0 +1,20 @@
+#pragma once
+
+#ifdef TI_WITH_DX12
+
+#include "taichi/runtime/program_impls/llvm/llvm_program.h"
+
+namespace taichi {
+namespace lang {
+
+class Dx12ProgramImpl : public LlvmProgramImpl {
+ public:
+  Dx12ProgramImpl(CompileConfig &config);
+
+  std::unique_ptr<AotModuleBuilder> make_aot_module_builder() override;
+};
+
+}  // namespace lang
+}  // namespace taichi
+
+#endif

taichi/util/offline_cache.cpp

@@ -26,6 +26,8 @@ std::string get_cache_path_by_arch(const std::string &base_path, Arch arch) {
     subdir = "gfx";
   } else if (arch == Arch::metal) {
     subdir = "metal";
+  } else if (arch == Arch::dx12) {
+    subdir = "dx12";
   } else {
     return base_path;
   }