Poc gpu tileir conversion main sync同步 Triton 3.6 主线并完善 TLE 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,6 +265,21 @@ endif()
 include_directories(${PROJECT_SOURCE_DIR}/third_party)
 include_directories(${PROJECT_BINARY_DIR}/third_party) # Tablegen'd files
 
+# 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/CommonIR)
+  add_subdirectory(third_party/flir/lib/Dialect/CommonIR)
+endif()
+
 # link_directories(${LLVM_LIBRARY_DIR})
 if (FLAGTREE_BACKEND MATCHES "^(cambricon|aipu|tsingmicro|enflame|rpu|thrive)$")
   include_directories(${PROJECT_SOURCE_DIR}/include)

bin/RegisterTritonDialects.h

@@ -14,6 +14,12 @@
 #include "tle/dialect/include/IR/Dialect.h" // flagtree tle raw
 #include "tle/dialect/include/Transforms/Passes.h"
 #endif
+#ifdef __FLIR_TILEIR__
+#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"
 #include "triton/Dialect/TritonGPU/IR/Dialect.h"
@@ -74,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();
@@ -156,6 +165,9 @@ inline void registerTritonDialects(mlir::DialectRegistry &registry) {
       mlir::triton::proton::gpu::ProtonGPUDialect, mlir::ROCDL::ROCDLDialect,
 #ifdef __TLE__
       mlir::triton::tle::TleDialect, // flagtree tle raw
+#endif
+#ifdef __FLIR_TILEIR__
+      mlir::triton::tile::CommonIRDialect,
 #endif
       mlir::triton::gluon::GluonDialect>();
 }

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()

python/triton/experimental/tle/language/gpu/__init__.py

@@ -3,6 +3,8 @@
     pipeline,
     alloc,
     copy,
+    store_tensor,
+    to_tensor,
     memory_space,
     local_ptr,
     warp_specialize,
@@ -17,6 +19,8 @@
     "pipeline",
     "alloc",
     "copy",
+    "to_tensor",
+    "store_tensor",
     "local_ptr",
     "warp_specialize",
     "storage_kind",

python/triton/experimental/tle/language/gpu/core.py

@@ -1,10 +1,10 @@
 # flagtree tle
 import builtins
 import triton.language.core as tl
-from typing import Optional, Sequence
 from enum import Enum
+from typing import Optional, Sequence
 from . import types as tle
-from .mthreads import copy as mthreads_copy
+from . import semantic as tle_semantic
 from triton.compiler.code_generator import flatten_values_to_ir, unflatten_ir_values
 
 from triton.language.core import (
@@ -247,9 +247,7 @@ def alloc(
 
     try:
         unwrapped_shape = [tl._unwrap_if_constexpr(dim) for dim in shape]
-        full_shape = unwrapped_shape
         dtype = tl._unwrap_if_constexpr(dtype)
-        elem_type = dtype.to_ir(_semantic.builder)
 
         if layout is None:
             if storage == tle.smem:
@@ -290,25 +288,25 @@ def alloc(
             # Use provided layout
             layout_handle = layout.to_ir(_semantic.builder)
 
-        if storage == tle.smem:
-            if init_value is not None:
-                mutable_ty = _semantic.builder.get_memdesc_type(full_shape, elem_type, layout_handle, "smem")
-                tensor_handle = _semantic.builder.create_local_alloc(mutable_ty, init_value.handle)
-            else:
-                tensor_handle = _semantic.builder.create_local_alloc(full_shape, elem_type, layout_handle)
-        else:
-            raise ValueError(f"Storage type {storage} not yet supported")
-
-        return tle.buffered_tensor(tensor_handle, dtype, unwrapped_shape, storage, layout, _semantic)
+        return tle_semantic.alloc(
+            unwrapped_shape,
+            dtype,
+            storage,
+            layout,
+            layout_handle,
+            init_value,
+            _semantic,
+        )
 
     except Exception as e:
         raise RuntimeError(f"Memory allocation failed: {str(e)}") from e
 
 
 class CopyDirection(Enum):
-    """Copy direction enum for data transfer operations"""
-    GM_TO_LOCAL = "GMTOLOCAL"  # Global memory to local memory
-    LOCAL_TO_GM = "LOCALTOGM"  # Local memory to global memory
+    """Copy direction for global/local transfers."""
+
+    GM_TO_LOCAL = "GMTOLOCAL"
+    LOCAL_TO_GM = "LOCALTOGM"
 
 
 @tl.builtin
@@ -361,104 +359,6 @@ def copy(
         TMA copy with offsets:
             tle.copy(tma_desc, local_buf, [64, 64], [x_offset, y_offset])
     """
-    mthreads_enabled = mthreads_copy.enabled()
-
-    def normcopy(
-        src: tl.tensor,
-        dst: tle.buffered_tensor,
-        shape: tuple,
-        direction,
-        _semantic=None,
-    ) -> None:
-        if mthreads_enabled:
-            mthreads_copy.validate_normal_copy(src, dst, shape, direction)
-
-        # Semantic analysis
-        try:
-            from .semantic import TLESemantic
-            if isinstance(_semantic, TLESemantic):
-                _semantic.analyze_copy_operation(src, dst, shape)
-        except ImportError:
-            # If semantic analysis module is not available, continue with warning
-            import warnings
-            warnings.warn("TLE semantic analysis module not available, skipping validation", UserWarning)
-
-        mask = None
-        other = None
-        boundary_check = ()
-        padding_option = ""
-        cache_modifier = ""
-        eviction_policy = ""
-        volatile = False
-
-        try:
-            if direction == CopyDirection.GM_TO_LOCAL:
-                # None fills the FlagTree hints slot; TLE copy has no hints to pass.
-                load_extra_args = () if mthreads_enabled else (None, )
-                tt_load = _semantic.load(src, mask, other, boundary_check, padding_option, cache_modifier,
-                                         eviction_policy, volatile, *load_extra_args)
-                local_ptrs = local_ptr(dst, _make_full_indices(dst, _semantic), _semantic=_semantic)
-                _semantic.store(local_ptrs, tt_load, mask, boundary_check, cache_modifier, eviction_policy)
-            else:
-                local_ptrs = local_ptr(src, _make_full_indices(src, _semantic), _semantic=_semantic)
-                load = tl.load(local_ptrs, _semantic=_semantic)
-                _semantic.store(dst, load, mask, boundary_check, cache_modifier, eviction_policy)
-        except Exception as e:
-            raise RuntimeError(f"copy operation failed: {str(e)}") from e
-
-    # this api is use for tma copy
-    def tmacopy(
-        src: tle.buffered_tensor | tl.tensor_descriptor,
-        dst: tle.buffered_tensor | tl.tensor_descriptor,
-        direction,
-        shape: tuple,
-        offsets: Sequence[constexpr | tensor],
-        _semantic=None,
-    ) -> None:
-        # Parameter validation
-        valid_types = (tle.buffered_tensor, tl.tensor_descriptor)
-
-        if not isinstance(src, valid_types):
-            raise ValueError(
-                f"Source parameter must be tle.buffered_tensor or tl.tensor_descriptor, but got {type(src).__name__}")
-
-        if not isinstance(dst, valid_types):
-            raise ValueError(
-                f"Destination parameter must be tle.buffered_tensor or tl.tensor_descriptor, but got {type(dst).__name__}"
-            )
-
-        # Auto-determine copy direction based on operand types
-        if isinstance(src, tle.buffered_tensor) and isinstance(dst, tl.tensor_descriptor):
-            desc = dst
-        elif isinstance(src, tl.tensor_descriptor) and isinstance(dst, tle.buffered_tensor):
-            desc = src
-        else:
-            raise ValueError(
-                f"Invalid copy combination: src={type(src).__name__}, dst={type(dst).__name__}. "
-                "One operand must be tl.tensor_descriptor (global memory) and the other must be tle.buffered_tensor (local memory)"
-            )
-
-        if not isinstance(shape, (tuple, list)):
-            # Try to handle Triton tuple-like objects
-            if hasattr(shape, '__iter__'):
-                shape = tuple(shape)
-            else:
-                raise ValueError(f"Shape parameter must be tuple or list, but got {type(shape)}")
-
-        if not isinstance(offsets, (tuple, list)):
-            # Try to handle Triton tuple-like objects
-            if hasattr(offsets, '__iter__'):
-                offsets = tuple(offsets)
-            else:
-                raise ValueError(f"Shape parameter must be tuple or list, but got {type(shape)}")
-
-        # Note: Skip shape assertion at this level since it requires _semantic context
-        # assert desc.shape == shape, "Shape mismatch between descriptor and provided shape"
-        assert len(offsets) == len(desc.shape), "Offsets and shape must have the same length"
-        offsets = _semantic._convert_to_ir_values(offsets, require_i64=False)
-        _semantic.builder.create_tma_copy(src.handle, dst.handle, offsets)
-        return
-
     # Parameter validation
     valid_types = (tl.tensor, tle.buffered_tensor, tl.tensor_descriptor)
 
@@ -472,55 +372,57 @@ def tmacopy(
             f"Destination parameter must be tl.tensor or tle.buffered_tensor  tl.tensor_descriptor, but got {type(dst).__name__}"
         )
 
-    # Auto-determine copy direction based on operand types
-    if isinstance(src, tle.buffered_tensor) and isinstance(dst, tl.tensor):
-        direction = CopyDirection.LOCAL_TO_GM
-        is_normcopy = True
-    elif isinstance(src, tl.tensor) and isinstance(dst, tle.buffered_tensor):
-        direction = CopyDirection.GM_TO_LOCAL
-        is_normcopy = True
-    elif isinstance(src, tle.buffered_tensor) and isinstance(dst, tl.tensor_descriptor):
-        direction = CopyDirection.LOCAL_TO_GM
-        is_normcopy = False
-    elif isinstance(src, tl.tensor_descriptor) and isinstance(dst, tle.buffered_tensor):
-        direction = CopyDirection.GM_TO_LOCAL
-        is_normcopy = False
-    else:
-        raise ValueError(
-            f"Invalid copy combination: src={type(src).__name__}, dst={type(dst).__name__}. "
-            "One operand must be tl.tensor (global memory) and the other must be tle.buffered_tensor (local memory)")
+    local_to_global = isinstance(src, tle.buffered_tensor) and isinstance(dst, (tl.tensor, tl.tensor_descriptor))
+    global_to_local = isinstance(dst, tle.buffered_tensor) and isinstance(src, (tl.tensor, tl.tensor_descriptor))
+    if not (local_to_global or global_to_local):
+        raise ValueError(f"Invalid copy combination: src={type(src).__name__}, dst={type(dst).__name__}. "
+                         "One operand must be a global tensor/descriptor and the other must be tle.buffered_tensor")
+    direction = CopyDirection.LOCAL_TO_GM if local_to_global else CopyDirection.GM_TO_LOCAL
 
     if not isinstance(shape, (tuple, list)):
         # Try to handle Triton tuple-like objects
         if hasattr(shape, '__iter__'):
             shape = tuple(shape)
         else:
             raise ValueError(f"Shape parameter must be tuple or list, but got {type(shape)}")
-    if is_normcopy:
-        return normcopy(src, dst, shape, direction, _semantic)
-    if mthreads_enabled:
-        return mthreads_copy.tmacopy(src, dst, direction, shape, offsets, _semantic)
-    else:
-        return tmacopy(src, dst, direction, shape, offsets, _semantic)
 
+    descriptor = src if isinstance(src,
+                                   tl.tensor_descriptor) else dst if isinstance(dst, tl.tensor_descriptor) else None
+    if descriptor is not None:
+        if offsets is None:
+            raise ValueError("TMA copy requires offsets")
+        if not isinstance(offsets, (tuple, list)):
+            if hasattr(offsets, '__iter__'):
+                offsets = tuple(offsets)
+            else:
+                raise ValueError(f"Offsets must be tuple or list, but got {type(offsets)}")
+        if len(offsets) != len(descriptor.shape):
+            raise ValueError("Offsets and descriptor shape must have the same rank")
 
-def _expand_index_to_shape(index: tl.tensor, shape: Sequence[int], axis: int, _semantic) -> tl.tensor:
-    idx = index
-    for _ in builtins.range(axis):
-        idx = tl.expand_dims(idx, 0, _semantic=_semantic)
-    for _ in builtins.range(len(shape) - axis - 1):
-        idx = tl.expand_dims(idx, len(idx.shape), _semantic=_semantic)
-    return tl.broadcast_to(idx, *shape, _semantic=_semantic)
+    tle_semantic.copy(src, dst, shape, offsets if descriptor is not None else None, direction, _semantic)
 
 
-def _make_full_indices(buffer: tle.buffered_tensor, _semantic) -> tuple[tl.tensor, ...]:
-    shape = tuple(int(tl._unwrap_if_constexpr(dim)) for dim in buffer.type.shape)
-    indices = []
-    for axis, dim in enumerate(shape):
-        idx = tl.arange(0, dim, _semantic=_semantic)
-        idx = _expand_index_to_shape(idx, shape, axis, _semantic)
-        indices.append(idx)
-    return tuple(indices)
+@tl.builtin
+def to_tensor(
+    memref: tle.buffered_tensor,
+    writable: bool = True,
+    target_shape=None,
+    _semantic=None,
+) -> tl.tensor:
+    if not isinstance(memref, tle.buffered_tensor):
+        raise ValueError(f"memref must be tle.gpu.buffered_tensor, got {type(memref).__name__}")
+
+    target_shape = tl._unwrap_if_constexpr(target_shape)
+    shape = list(memref.shape if target_shape is None else target_shape)
+    writable = tl._unwrap_if_constexpr(writable)
+    return tle_semantic.to_tensor(memref, bool(writable), tl.block_type(memref.dtype, shape), _semantic)
+
+
+@tl.builtin
+def store_tensor(tensor_value: tl.tensor, dst: tle.buffered_tensor, _semantic=None) -> None:
+    if not isinstance(dst, tle.buffered_tensor):
+        raise ValueError(f"dst must be tle.gpu.buffered_tensor, got {type(dst).__name__}")
+    tle_semantic.store_tensor(tensor_value, dst, _semantic)
 
 
 @tl.builtin
@@ -628,8 +530,7 @@ def local_ptr(
     else:
         result_ty = tl.block_type(ptr_dtype, list(view_shape))
         result_ir = result_ty.to_ir(_semantic.builder)
-    handles = [idx.handle for idx in idx_tensors]
-    local_ptr_op = _semantic.builder.create_local_pointers(result_ir, buffer.handle, *handles)
+    local_ptr_op = tle_semantic.local_ptr(result_ir, buffer, idx_tensors, _semantic)
 
     result_tensor = tl.tensor(local_ptr_op.get_result(0), result_ty)