[BACKEND] Triton 3.6：重建 GPU CommonIR 并直接转换为 TTGIR

CMakeLists.txt

@@ -16,6 +16,13 @@ if(FLAGTREE_BACKEND)
     add_definitions(-DFLAGTREE_BACKEND=\"${FLAGTREE_BACKEND}\")
 endif()
 
+# CommonIR is an explicit build-time feature. Keep the public contract as the
+# C/C++ FLAGTREE_COMMON_IR macro rather than adding a CMake cache option.
+set(FLAGTREE_COMMON_IR_ENABLED "$ENV{FLAGTREE_COMMON_IR}")
+if(FLAGTREE_COMMON_IR_ENABLED)
+  add_compile_definitions(FLAGTREE_COMMON_IR)
+endif()
+
 set(FLAGTREE_TLE ON)
 if(FLAGTREE_TLE)
   add_definitions(-D__TLE__)
@@ -258,15 +265,19 @@ endif()
 include_directories(${PROJECT_SOURCE_DIR}/third_party)
 include_directories(${PROJECT_BINARY_DIR}/third_party) # Tablegen'd files
 
-# TileIR dialect from FLIR is used by the TLE GPU -> Common IR POC path.
-# Keep it optional at configure time so non-POC builds without FLIR are not
-# forced to download the dependency.
-if(EXISTS "${CMAKE_CURRENT_SOURCE_DIR}/third_party/flir/CMakeLists.txt" AND NOT TARGET TileIRIR)
+# CommonIR is only part of builds that explicitly enable CommonIR.
+if(FLAGTREE_COMMON_IR_ENABLED AND NOT TARGET CommonIRIR)
+  if(FLAGTREE_BACKEND)
+    message(FATAL_ERROR "FLAGTREE_COMMON_IR on triton_v3.6.x requires the default NVIDIA backend")
+  endif()
+  if(NOT EXISTS "${CMAKE_CURRENT_SOURCE_DIR}/third_party/flir/CMakeLists.txt")
+    message(FATAL_ERROR "FLAGTREE_COMMON_IR requires FLIR at third_party/flir")
+  endif()
   add_definitions(-D__FLIR_TILEIR__)
   include_directories(${CMAKE_CURRENT_SOURCE_DIR}/third_party/flir/include)
   include_directories(${CMAKE_CURRENT_BINARY_DIR}/third_party/flir/include)
-  add_subdirectory(third_party/flir/include/mlir-ext/Dialect/TileIR)
-  add_subdirectory(third_party/flir/lib/Dialect/TileIR)
+  add_subdirectory(third_party/flir/include/mlir-ext/Dialect/CommonIR)
+  add_subdirectory(third_party/flir/lib/Dialect/CommonIR)
 endif()
 
 # link_directories(${LLVM_LIBRARY_DIR})

bin/RegisterTritonDialects.h

@@ -15,7 +15,10 @@
 #include "tle/dialect/include/Transforms/Passes.h"
 #endif
 #ifdef __FLIR_TILEIR__
-#include "mlir-ext/Dialect/TileIR/IR/TileIRDialect.h"
+#include "mlir-ext/Dialect/CommonIR/IR/CommonIRDialect.h"
+#endif
+#ifdef FLAGTREE_COMMON_IR
+#include "nvidia/include/CommonIRToTTGIR/Passes.h"
 #endif
 #include "triton/Dialect/Gluon/Transforms/Passes.h"
 #include "triton/Dialect/Triton/IR/Dialect.h"
@@ -77,6 +80,9 @@ inline void registerTritonDialects(mlir::DialectRegistry &registry) {
   mlir::triton::gluon::registerGluonPasses();
 #ifdef __TLE__
   mlir::triton::tle::registerPasses(); // flagtree tle
+#endif
+#ifdef FLAGTREE_COMMON_IR
+  mlir::triton::registerCommonIRToTTGIRPasses();
 #endif
   mlir::test::registerTestAliasPass();
   mlir::test::registerTestAlignmentPass();
@@ -161,7 +167,7 @@ inline void registerTritonDialects(mlir::DialectRegistry &registry) {
       mlir::triton::tle::TleDialect, // flagtree tle raw
 #endif
 #ifdef __FLIR_TILEIR__
-      mlir::triton::tile::TileIRDialect,
+      mlir::triton::tile::CommonIRDialect,
 #endif
       mlir::triton::gluon::GluonDialect>();
 }

lib/Conversion/TritonToTritonGPU/TritonToTritonGPUPass.cpp

@@ -606,6 +606,9 @@ void populateTritonPatterns(TritonGPUTypeConverter &typeConverter,
       GenericOpPattern<triton::gpu::LocalAllocOp>,
       GenericOpPattern<triton::gpu::LocalStoreOp>,
       GenericOpPattern<triton::gpu::LocalLoadOp>,
+#endif
+#ifdef FLAGTREE_COMMON_IR
+      GenericOpPattern<triton::gpu::AsyncCopyGlobalToLocalOp>,
 #endif
       TritonExpandDimsPattern,
       TritonTransPattern,
@@ -975,6 +978,10 @@ class ConvertTritonToTritonGPU
     TritonGPUTypeConverter typeConverter(context, numWarps, threadsPerWarp,
                                          numCTAs, enableSourceRemat);
     TritonGPUConversionTarget target(*context, typeConverter);
+#ifdef FLAGTREE_COMMON_IR
+    target.addDynamicallyLegalOp<triton::gpu::AsyncCopyGlobalToLocalOp>(
+        [&](Operation *op) { return typeConverter.isLegal(op); });
+#endif
     // rewrite patterns
     RewritePatternSet patterns(context);
     // add rules

python/setup_tools/utils/__init__.py

@@ -38,7 +38,11 @@ def _register_submodule(self, submodule):
         "name": "triton_shared", "url": "https://github.com/microsoft/triton-shared.git", "commit_id":
         "5842469a16b261e45a2c67fbfc308057622b03ee"
     },
-    {"name": "flir", "url": "https://github.com/FlagTree/flir.git"},
+    {
+        "name": "flir",
+        "url": "https://github.com/kateyijian/flir.git",
+        "commit_id": "a884fd3c608db2203f08ec5a281327f5dc0153b0",
+    },
     {"name": "flagcx", "url": "https://github.com/flagos-ai/FlagCX.git", "relative_path": "tle/third_party/flagcx"},
 ))
 

python/triton/_common_ir.py

@@ -0,0 +1,3 @@
+from triton._C.libtriton import tle as tle_ir
+
+ENABLED = tle_ir.is_common_ir_enabled()