Move get_threads_per_warp to init_triton_intel (#2932)

Signed-off-by: Whitney Tsang <[email protected]>

Author: whitneywhtsang

python/src/ir.cc

@@ -138,17 +138,6 @@ void outputWarning(Location loc, const std::string &msg) {
 
 } // anonymous namespace
 
-/*****************************************************************************/
-/* Python bindings for triton::ir::ttgir                                     */
-/*****************************************************************************/
-
-void init_triton_ttgpuir(py::module &&m) {
-  m.def("get_threads_per_warp", [](mlir::ModuleOp &mod) -> py::object {
-    auto ret = mlir::triton::gpu::TritonGPUDialect::getThreadsPerWarp(mod);
-    return py::int_(ret);
-  });
-}
-
 /*****************************************************************************/
 /* Python bindings for ir                                                    */
 /*****************************************************************************/
@@ -1794,9 +1783,6 @@ void init_triton_ir(py::module &&m) {
         if (failed(self.run(mod.getOperation())))
           throw std::runtime_error("PassManager::run failed");
       });
-
-  // ttgpu dialect bindings.
-  init_triton_ttgpuir(m.def_submodule("ttgpuir"));
 }
 
 void init_triton_env_vars(py::module &m) {

third_party/intel/backend/compiler.py

@@ -228,7 +228,7 @@ def make_ttgir(mod, metadata, opt, properties):
         pm.run(mod)
 
         # Overwrite the threads_per_warp option with the module annotation.
-        opt.threads_per_warp = ir.ttgpuir.get_threads_per_warp(mod)
+        opt.threads_per_warp = intel.get_threads_per_warp(mod)
 
         # Run the TTIR -> TTGIR pass pipeline.
         pm = ir.pass_manager(mod.context)
@@ -271,7 +271,7 @@ def make_llir(src, metadata, options):
         num_warp_groups = src.get_int_attr("ttg.num-warp-groups-per-cta")
         if num_warp_groups is not None:
             metadata["num_warps"] *= num_warp_groups
-        threads_per_warp = ir.ttgpuir.get_threads_per_warp(src)
+        threads_per_warp = intel.get_threads_per_warp(src)
         metadata["threads_per_warp"] = threads_per_warp
         mod = src
         # TritonGPU -> LLVM-IR (MLIR)

third_party/intel/triton_xpu.cc

@@ -234,6 +234,11 @@ void init_triton_intel(py::module &&m) {
     context.loadAllAvailableDialects();
   });
 
+  m.def("get_threads_per_warp", [](mlir::ModuleOp &mod) -> py::object {
+    auto ret = mlir::triton::gpu::TritonGPUDialect::getThreadsPerWarp(mod);
+    return py::int_(ret);
+  });
+
   // May do this after llvm ir according to user fmath flag.
   m.def("set_fast_math", [](mlir::ModuleOp mod) {
     using namespace mlir;