[Refactor][CodeGen] Refactor CodeGen part for multi-backend decoupling

CMakeLists.txt

@@ -169,7 +169,9 @@ set(USE_GTEST OFF)
 
 # Include directories for TileLang
 set(TILE_LANG_INCLUDES ${TVM_INCLUDES})
-list(APPEND TILE_LANG_INCLUDES ${CMAKE_CURRENT_SOURCE_DIR}/src)
+# Add TileLang's own src/ to include path so cross-directory includes
+# can use paths relative to src/ (e.g. "target/utils.h", "op/builtin.h").
+list(INSERT TILE_LANG_INCLUDES 0 "${CMAKE_CURRENT_SOURCE_DIR}/src")
 
 # Collect source files
 file(GLOB TILE_LANG_SRCS
@@ -186,8 +188,9 @@ file(GLOB TILE_LANG_SRCS
   src/target/codegen_c_host.cc
   src/target/codegen_c.cc
   src/target/rt_mod_c.cc
-  # intrin_rule doesn't have system dependency
-  src/target/intrin_rule*.cc
+  # intrin_rule doesn't have system dependency; always compiled regardless of backend
+  src/backend/cuda/codegen/intrin_rule_cuda.cc
+  src/backend/rocm/codegen/intrin_rule_hip.cc
 )
 
 # Always include CPU-safe runtime helpers

src/backend/cuda/CMakeLists.txt

@@ -29,7 +29,7 @@ if(TILELANG_USE_CUDA_STUBS)
   #   - cuModuleLoadData, cuLaunchKernel, cuMemsetD32_v2, etc.
   # These can be called directly without any wrapper macros.
   # ============================================================================
-  add_library(cuda_stub SHARED src/target/stubs/cuda.cc)
+  add_library(cuda_stub SHARED src/backend/cuda/codegen/stubs/cuda.cc)
   target_include_directories(cuda_stub PRIVATE ${CUDAToolkit_INCLUDE_DIRS})
   # Export symbols with visibility="default" when building
   target_compile_definitions(cuda_stub PRIVATE TILELANG_CUDA_STUB_EXPORTS)
@@ -53,7 +53,7 @@ if(TILELANG_USE_CUDA_STUBS)
   # The stub exports a minimal set of CUDA Runtime API entrypoints used by TVM
   # and lazily loads libcudart at runtime on first API call.
   # ============================================================================
-  add_library(cudart_stub SHARED src/target/stubs/cudart.cc)
+  add_library(cudart_stub SHARED src/backend/cuda/codegen/stubs/cudart.cc)
   target_include_directories(cudart_stub PRIVATE ${CUDAToolkit_INCLUDE_DIRS})
   target_compile_definitions(cudart_stub PRIVATE TILELANG_CUDART_STUB_EXPORTS)
   target_link_libraries(cudart_stub PRIVATE ${CMAKE_DL_LIBS})
@@ -81,7 +81,7 @@ if(TILELANG_USE_CUDA_STUBS)
   # The stub exports a minimal set of NVRTC C API entrypoints used by TVM and
   # lazily loads libnvrtc at runtime on first API call.
   # ============================================================================
-  add_library(nvrtc_stub SHARED src/target/stubs/nvrtc.cc)
+  add_library(nvrtc_stub SHARED src/backend/cuda/codegen/stubs/nvrtc.cc)
   target_include_directories(nvrtc_stub PRIVATE ${CUDAToolkit_INCLUDE_DIRS})
   target_compile_definitions(nvrtc_stub PRIVATE TILELANG_NVRTC_STUB_EXPORTS)
   target_link_libraries(nvrtc_stub PRIVATE ${CMAKE_DL_LIBS})
@@ -102,13 +102,13 @@ endif()
 
 file(GLOB TILE_LANG_CUDA_SRCS
   src/runtime/runtime.cc
-  src/target/ptx.cc
-  src/target/codegen_cuda.cc
-  src/target/codegen_py.cc
+  src/backend/cuda/codegen/ptx.cc
+  src/backend/cuda/codegen/codegen_cuda.cc
+  src/backend/cuda/codegen/codegen_py.cc
   src/target/codegen_utils.cc
-  src/target/codegen_cutedsl.cc
-  src/target/rt_mod_cuda.cc
-  src/target/rt_mod_cutedsl.cc
+  src/backend/cuda/codegen/codegen_cutedsl.cc
+  src/backend/cuda/codegen/rt_mod_cuda.cc
+  src/backend/cuda/codegen/rt_mod_cutedsl.cc
   src/backend/cuda/op/*.cc
 )
 list(REMOVE_ITEM TILE_LANG_CUDA_SRCS

src/target/codegen_cuda.cc → src/backend/cuda/codegen/codegen_cuda.cc

@@ -15,11 +15,11 @@
 #include <utility>
 #include <vector>
 
-#include "../op/builtin.h"
-#include "../transform/common/attr.h"
-#include "./ptx.h"
-#include "./utils.h"
 #include "arith/pattern_match.h"
+#include "backend/cuda/codegen/ptx.h"
+#include "op/builtin.h"
+#include "target/utils.h"
+#include "transform/common/attr.h"
 
 namespace tvm {
 namespace codegen {

src/target/codegen_cutedsl.cc → src/backend/cuda/codegen/codegen_cutedsl.cc

@@ -3,8 +3,8 @@
  */
 
 #include "codegen_cutedsl.h"
-#include "codegen_utils.h"
-#include "ptx.h"
+#include "backend/cuda/codegen/ptx.h"
+#include "target/codegen_utils.h"
 #include <tvm/arith/analyzer.h>
 #include <tvm/ffi/function.h>
 #include <tvm/ir/transform.h>
@@ -19,8 +19,8 @@
 #include <utility>
 #include <vector>
 
-#include "../op/builtin.h"
 #include "arith/pattern_match.h"
+#include "op/builtin.h"
 
 namespace tvm {
 namespace codegen {

src/target/codegen_py.cc → src/backend/cuda/codegen/codegen_py.cc

@@ -2,7 +2,7 @@
  * \file codegen_py.cc
  */
 #include "codegen_py.h"
-#include "codegen_utils.h"
+#include "target/codegen_utils.h"
 
 #include <tvm/arith/analyzer.h>
 #include <tvm/ir/name_supply.h>

src/target/intrin_rule_cuda.cc → src/backend/cuda/codegen/intrin_rule_cuda.cc

@@ -5,7 +5,7 @@
 #include <tvm/tir/builtin.h>
 #include <tvm/tir/op_attr_types.h>
 
-#include "../support/ffi_aliases.h"
+#include "support/ffi_aliases.h"
 #include "target/intrin_rule.h"
 
 namespace tvm {

src/target/rt_mod_cuda.cc → src/backend/cuda/codegen/rt_mod_cuda.cc

@@ -1,8 +1,8 @@
-#include "../transform/common/attr.h"
 #include "codegen_cuda.h"
 #include "runtime/cuda/cuda_module.h"
 #include "runtime/meta_data.h"
 #include "runtime/pack_args.h"
+#include "transform/common/attr.h"
 #include <tvm/ffi/reflection/registry.h>
 #include <tvm/ir/transform.h>
 

src/backend/metal/CMakeLists.txt

@@ -11,7 +11,7 @@ if(NOT APPLE)
 endif()
 
 file(GLOB TILE_LANG_METAL_SRCS
-  src/target/rt_mod_metal.cc
+  src/backend/metal/codegen/rt_mod_metal.cc
 )
 list(APPEND TILE_LANG_SRCS ${TILE_LANG_METAL_SRCS})
 # FIXME: CIBW failed with backtrace, why???

src/backend/metal/codegen/rt_mod_metal.cc

@@ -0,0 +1,27 @@
+/*!
+ * \file rt_mod_metal.cc
+ * \brief Metal codegen entry point.
+ *
+ * Metal codegen is handled by CodeGenCHost (target/codegen_c_host.cc), which
+ * has built-in Metal context support via the is_in_metal_context flag.
+ * When IR contains AttrStmt with attr_key == "metal_context", the host
+ * codegen emits Metal-specific dispatch_sync / MTLCommandBuffer code.
+ */
+#include "target/codegen_c_host.h"
+
+#include <tvm/ffi/reflection/registry.h>
+
+namespace tvm {
+namespace codegen {
+
+ffi::Module BuildTileLangMetal(IRModule mod, Target target) {
+  return tl::BuildTileLangCHost(mod, target);
+}
+
+TVM_FFI_STATIC_INIT_BLOCK() {
+  namespace refl = tvm::ffi::reflection;
+  refl::GlobalDef().def("target.build.tilelang_metal", BuildTileLangMetal);
+}
+
+} // namespace codegen
+} // namespace tvm

src/backend/rocm/CMakeLists.txt

@@ -23,7 +23,7 @@ if(TILELANG_USE_HIP_STUBS)
   # It also provides minimal HSA wrappers (hsa_init / hsa_shut_down) to avoid a
   # hard DT_NEEDED dependency on libhsa-runtime64 in ROCm-enabled wheels.
   # ============================================================================
-  add_library(hip_stub SHARED src/target/stubs/hip.cc)
+  add_library(hip_stub SHARED src/backend/rocm/codegen/stubs/hip.cc)
   target_include_directories(hip_stub PRIVATE ${ROCM_INCLUDE_DIRS})
   target_compile_definitions(hip_stub PRIVATE TILELANG_HIP_STUB_EXPORTS)
   target_link_libraries(hip_stub PRIVATE ${CMAKE_DL_LIBS})
@@ -40,7 +40,7 @@ if(TILELANG_USE_HIP_STUBS)
   # This library provides a minimal HIPRTC API surface and lazily loads
   # libhiprtc.so at runtime.
   # ============================================================================
-  add_library(hiprtc_stub SHARED src/target/stubs/hiprtc.cc)
+  add_library(hiprtc_stub SHARED src/backend/rocm/codegen/stubs/hiprtc.cc)
   target_include_directories(hiprtc_stub PRIVATE ${ROCM_INCLUDE_DIRS})
   target_compile_definitions(hiprtc_stub PRIVATE TILELANG_HIPRTC_STUB_EXPORTS)
   target_link_libraries(hiprtc_stub PRIVATE ${CMAKE_DL_LIBS})
@@ -66,8 +66,8 @@ if(TILELANG_USE_HIP_STUBS)
 endif()
 
 file(GLOB TILE_LANG_HIP_SRCS
-  src/target/codegen_hip.cc
-  src/target/rt_mod_hip.cc
+  src/backend/rocm/codegen/codegen_hip.cc
+  src/backend/rocm/codegen/rt_mod_hip.cc
   src/backend/rocm/op/*.cc
 )
 list(APPEND TILE_LANG_SRCS ${TILE_LANG_HIP_SRCS})

src/target/codegen_hip.cc → src/backend/rocm/codegen/codegen_hip.cc

@@ -15,9 +15,9 @@
 #include <utility>
 #include <vector>
 
-#include "../op/builtin.h"
+#include "op/builtin.h"
 #include "target/source/ptx.h"
-#include "utils.h"
+#include "target/utils.h"
 
 namespace tvm {
 namespace codegen {

src/target/intrin_rule_hip.cc → src/backend/rocm/codegen/intrin_rule_hip.cc

@@ -5,7 +5,7 @@
 #include <tvm/tir/builtin.h>
 #include <tvm/tir/op_attr_types.h>
 
-#include "../support/ffi_aliases.h"
+#include "support/ffi_aliases.h"
 #include "target/intrin_rule.h"
 
 namespace tvm {

src/op/builtin.cc

@@ -10,8 +10,8 @@
 #include <tvm/tir/op.h>
 #include <tvm/tir/op_attr_types.h>
 
-#include "../target/stubs/cuda.h"
-#include "../target/utils.h"
+#include "backend/cuda/codegen/stubs/cuda.h"
+#include "target/utils.h"
 
 namespace tvm {
 namespace tl {

src/op/utils.h

@@ -6,8 +6,8 @@
 #ifndef TVM_TL_OP_UTILS_H_
 #define TVM_TL_OP_UTILS_H_
 
-#include "../target/stubs/cuda.h"
 #include "./operator.h"
+#include "backend/cuda/codegen/stubs/cuda.h"
 #include "region.h"
 #include "tvm/runtime/base.h"
 #include <tvm/tir/buffer.h>

src/runtime/runtime.cc

@@ -6,7 +6,7 @@
 
 #include "runtime.h"
 
-#include "../target/stubs/cuda.h"
+#include "backend/cuda/codegen/stubs/cuda.h"
 #include <cstdint>
 #include <sstream>
 #include <tvm/ffi/function.h>

src/target/codegen_c_host.h

@@ -127,6 +127,13 @@ class CodeGenCHost : public tvm::codegen::CodeGenC {
   }
 };
 
+/*!
+ * \brief Build a TileLang C host module for the given IRModule and target.
+ * Also handles Metal target through is_in_metal_context flag in IR.
+ */
+::tvm::ffi::Module BuildTileLangCHost(::tvm::IRModule mod,
+                                      ::tvm::Target target);
+
 } // namespace tl
 } // namespace tvm