[TLE][BACKEND] Add AMD tile and cumsum lowering

include/triton/Conversion/TritonGPUToLLVM/TargetInfoBase.h

@@ -137,6 +137,10 @@ class TargetInfoBase {
   virtual bool isCuda() const { return false; }
   virtual bool isHCU() const { return false; }
 
+  virtual bool usesTargetThreadIdLowering() const { return false; }
+  virtual bool usesTargetBarrierLowering() const { return false; }
+  virtual bool usesTargetShuffleLowering() const { return false; }
+
   // Annotate target specific information to local load operations during
   // lowering to LLVM. `llLoadOp` is the generated LLVM load op.
   virtual void localLoadOpAnnotation(triton::gpu::LocalLoadOp localLoadOp,

python/test/tle/unit/test_tle_cumsum.py

@@ -10,7 +10,6 @@
 import triton
 import triton.language as tl
 import triton.experimental.tle.language as tle
-from triton._flagtree_backend import FLAGTREE_BACKEND
 
 
 def _is_enflame_backend():
@@ -19,14 +18,34 @@ def _is_enflame_backend():
 
 
 def _is_hcu_backend():
-    target = triton.runtime.driver.active.get_current_target()
-    return target.backend == "hip"
+    try:
+        driver = triton.runtime.driver.active
+        return type(driver).__module__.startswith("triton.backends.hcu")
+    except Exception:
+        return False
 
 
 _nv_mma_shared_layout = tl.constexpr(False if _is_hcu_backend() else True)
 threads_per_warp = 64 if _is_hcu_backend() else 32
 
 
+def _is_nvidia_cuda_backend():
+    try:
+        driver = triton.runtime.driver.active
+        target = driver.get_current_target()
+        return (target.backend == "cuda" and type(driver).__module__.startswith("triton.backends.nvidia"))
+    except Exception:
+        return False
+
+
+def _is_amd_hip_backend():
+    try:
+        driver = triton.runtime.driver.active
+        return type(driver).__module__.startswith("triton.backends.amd")
+    except Exception:
+        return False
+
+
 def _require_cuda():
     try:
         if _is_enflame_backend():
@@ -203,9 +222,10 @@ def test_tle_cumsum_exclusive_and_total(dtype, n, block, reverse, num_warps):
         torch.testing.assert_close(total[0], expected_total)
 
 
-@pytest.mark.skipif(_is_enflame_backend(), reason="PTX-specific regression guard not applicable on Enflame GCU")
-@pytest.mark.skipif(_is_hcu_backend(), reason="PTX-specific regression guard not applicable on HCU")
-@pytest.mark.skipif(FLAGTREE_BACKEND == "ppu", reason="PTX-specific regression guard not applicable on PPU")
+@pytest.mark.skipif(
+    not _is_nvidia_cuda_backend(),
+    reason="PTX-specific regression guard requires NVIDIA CUDA backend",
+)
 def test_tle_cumsum_ptx_fastpath_regression_guard():
     block = 512
     x = torch.randint(-1024, 1024, (block, ), device="cuda", dtype=torch.int32)
@@ -275,6 +295,7 @@ def test_tle_cumsum_amdgcn_fastpath_regression_guard():
         "Detected predicated ds_write: possible regression to generic path"
 
 
+@pytest.mark.skipif(_is_amd_hip_backend(), reason="requires AMD local-pointer lowering")
 def test_tle_cumsum_helper_preserves_adjacent_sentinel():
     block = 512
     num_warps = block // threads_per_warp
@@ -296,6 +317,7 @@ def test_tle_cumsum_helper_preserves_adjacent_sentinel():
     torch.testing.assert_close(sentinel, expected_sentinel)
 
 
+@pytest.mark.skipif(_is_amd_hip_backend(), reason="requires AMD local-pointer lowering")
 def test_tle_cumsum_scalar_base_addptr_alias_regression():
     block = 512
     num_warps = block // threads_per_warp

test/Conversion/amd/tle_tile_ops_to_llvm.mlir

@@ -0,0 +1,69 @@
+// Copyright 2025-     FlagOS Contributors
+//
+// Permission is hereby granted, free of charge, to any person obtaining
+// a copy of this software and associated documentation files
+// (the "Software"), to deal in the Software without restriction,
+// including without limitation the rights to use, copy, modify, merge,
+// publish, distribute, sublicense, and/or sell copies of the Software,
+// and to permit persons to whom the Software is furnished to do so,
+// subject to the following conditions:
+//
+// The above copyright notice and this permission notice shall be
+// included in all copies or substantial portions of the Software.
+//
+// THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND,
+// EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF
+// MERCHANTABILITY, FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT.
+// IN NO EVENT SHALL THE AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY
+// CLAIM, DAMAGES OR OTHER LIABILITY, WHETHER IN AN ACTION OF CONTRACT,
+// TORT OR OTHERWISE, ARISING FROM, OUT OF OR IN CONNECTION WITH THE
+// SOFTWARE OR THE USE OR OTHER DEALINGS IN THE SOFTWARE.
+
+// RUN: triton-opt %s -split-input-file --allocate-amdgpu-shared-memory --convert-triton-amdgpu-to-llvm=arch=gfx1201 --convert-builtin-func-to-llvm | FileCheck %s
+
+#blocked = #ttg.blocked<{sizePerThread = [1, 1], threadsPerWarp = [32, 1], warpsPerCTA = [1, 1], order = [1, 0]}>
+
+module attributes {"ttg.num-ctas" = 1 : i32, "ttg.num-warps" = 1 : i32, ttg.target = "hip:gfx1201", "ttg.threads-per-warp" = 32 : i32} {
+  // CHECK-LABEL: llvm.func @extract_tile_smem
+  // CHECK: rocdl.workitem.id.x
+  // CHECK-COUNT-2: rocdl.barrier
+  // CHECK-NOT: tle.extract_tile
+  // CHECK-NOT: nvvm.barrier
+  tt.func @extract_tile_smem(%src: tensor<32x32xf32, #blocked>, %idx: i32) {
+    %tile = tle.extract_tile %src[%idx] {tile_shape = array<i64: 16, 16>} : tensor<32x32xf32, #blocked>, i32 -> tensor<16x16xf32, #blocked>
+    tt.return
+  }
+}
+
+// -----
+
+#blocked = #ttg.blocked<{sizePerThread = [1, 1], threadsPerWarp = [32, 1], warpsPerCTA = [1, 1], order = [1, 0]}>
+
+module attributes {"ttg.num-ctas" = 1 : i32, "ttg.num-warps" = 1 : i32, ttg.target = "hip:gfx1201", "ttg.threads-per-warp" = 32 : i32} {
+  // CHECK-LABEL: llvm.func @insert_tile_smem
+  // CHECK: rocdl.workitem.id.x
+  // CHECK-COUNT-2: rocdl.barrier
+  // CHECK-NOT: tle.insert_tile
+  // CHECK-NOT: nvvm.barrier
+  tt.func @insert_tile_smem(%src: tensor<32x32xf32, #blocked>, %tile: tensor<16x16xf32, #blocked>, %idx: i32) {
+    %result = tle.insert_tile %src[%idx] = %tile {tile_shape = array<i64: 16, 16>} : tensor<32x32xf32, #blocked>, i32, tensor<16x16xf32, #blocked> -> tensor<32x32xf32, #blocked>
+    tt.return
+  }
+}
+
+// -----
+
+#blocked = #ttg.blocked<{sizePerThread = [1], threadsPerWarp = [32], warpsPerCTA = [4], order = [0]}>
+
+module attributes {"ttg.num-ctas" = 1 : i32, "ttg.num-warps" = 4 : i32, ttg.target = "hip:gfx1201", "ttg.threads-per-warp" = 32 : i32} {
+  // CHECK-LABEL: llvm.func @exclusive_cumsum
+  // CHECK-COUNT-6: rocdl.ds_bpermute
+  // CHECK-COUNT-2: rocdl.barrier
+  // CHECK-NOT: tle.exclusive_cumsum
+  // CHECK-NOT: nvvm.shfl
+  tt.func public @exclusive_cumsum(%arg0: tensor<128xi32, #blocked>, %out: !tt.ptr<i32>) {
+    %exclusive, %total = "tle.exclusive_cumsum"(%arg0) {axis = 0 : i32, reverse = false} : (tensor<128xi32, #blocked>) -> (tensor<128xi32, #blocked>, i32)
+    tt.store %out, %total : !tt.ptr<i32>
+    tt.return
+  }
+}

third_party/amd/backend/compiler.py

@@ -27,6 +27,21 @@ def is_in_thread_transpose_enabled(arch):
     return (arch == "gfx942") if knobs.amd.use_in_thread_transpose is None else knobs.amd.use_in_thread_transpose
 
 
+def reject_residual_tle_ops(mod, arch):
+    residual_ops = set()
+
+    def visit(op):
+        name = op.get_name()
+        if name.startswith("tle."):
+            residual_ops.add(name)
+
+    mod.walk(visit)
+    if residual_ops:
+        unsupported = ", ".join(sorted(residual_ops))
+        raise NotImplementedError(f"TLE op(s) remain after AMD lowering for arch {arch}: {unsupported}. "
+                                  "These operations do not have an AMD TritonGPU-to-LLVM lowering.")
+
+
 @dataclass(frozen=True)
 class HIPOptions:
     num_warps: int = 4
@@ -353,6 +368,8 @@ def make_llir(src, metadata, options):
             passes.llvmir.add_di_local_variable(pm)
             pm.run(mod, 'make_llir.dump_ir_extract_di_local_variables')
 
+        reject_residual_tle_ops(mod, options.arch)
+
         # LLVM-IR (MLIR) -> LLVM-IR (LLVM)
         llvm.init_targets()
         context = llvm.context()

third_party/amd/lib/Analysis/AMDGPUAllocation.cpp

@@ -6,6 +6,11 @@
 
 #include "third_party/amd/include/Dialect/TritonAMDGPU/Utility/CommonUtils.h"
 
+#ifdef __TLE__
+#include "tle/dialect/include/IR/Dialect.h"
+#include <limits>
+#endif
+
 namespace mlir::triton::AMD {
 
 // Max shmem instruction in bits
@@ -139,6 +144,43 @@ unsigned AMDAllocationAnalysisScratchSizeFn(Operation *op) {
                                           op->hasAttr(AttrSharedMemPadded));
   }
 
+#ifdef __TLE__
+  // Tile-level extension (TLE) ops stage data through shared memory; register
+  // their scratch sizes so attachAllocationSizeAndOffsetAttr assigns an
+  // allocation.offset (mirrors the NVIDIA scratch-size function).
+  if (auto cumsumOp = dyn_cast<mlir::triton::tle::ExclusiveCumsumOp>(op)) {
+    auto srcTy = dyn_cast<RankedTensorType>(cumsumOp.getSrc().getType());
+    if (!srcTy || srcTy.getRank() != 1)
+      return 0;
+    int64_t axisExtent = srcTy.getShape()[0];
+    if (ShapedType::isDynamic(axisExtent) || axisExtent <= 0)
+      return 0;
+    unsigned elemBytes =
+        static_cast<unsigned>(std::max<int>(1, getBitwidth(srcTy) / 8));
+    int64_t numWarps = std::max<int64_t>(1, triton::gpu::lookupNumWarps(op));
+    uint64_t totalBytes = (static_cast<uint64_t>(axisExtent) +
+                           static_cast<uint64_t>(numWarps) + 1ull) *
+                          elemBytes;
+    if (totalBytes > std::numeric_limits<unsigned>::max())
+      return 0;
+    return static_cast<unsigned>(totalBytes);
+  }
+  if (auto extractTileOp = dyn_cast<mlir::triton::tle::ExtractTileOp>(op)) {
+    auto dstTy = dyn_cast<RankedTensorType>(extractTileOp.getType());
+    if (!dstTy)
+      return 0;
+    return static_cast<unsigned>(dstTy.getNumElements() *
+                                 (getBitwidth(dstTy) / 8));
+  }
+  if (auto insertTileOp = dyn_cast<mlir::triton::tle::InsertTileOp>(op)) {
+    auto tileTy = dyn_cast<RankedTensorType>(insertTileOp.getTile().getType());
+    if (!tileTy)
+      return 0;
+    return static_cast<unsigned>(tileTy.getNumElements() *
+                                 (getBitwidth(tileTy) / 8));
+  }
+#endif
+
   return defaultAllocationAnalysisScratchSizeFn(op);
 }
 

third_party/amd/lib/TritonAMDGPUToLLVM/CMakeLists.txt

@@ -1,3 +1,13 @@
+if(FLAGTREE_TLE)
+    set(_TLE_LIBS
+        TritonTLEAnalysis
+        TleToLLVM
+        TritonTLETransforms
+    )
+else()
+    set(_TLE_LIBS "")
+endif()
+
 add_triton_library(TritonAMDGPUToLLVM
     AsyncUtility.cpp
     AtomicRMWOpsEmitter.cpp
@@ -40,4 +50,5 @@ add_triton_library(TritonAMDGPUToLLVM
     LLVMCore
     LLVMPasses
     LLVMSupport
+    ${_TLE_LIBS}
 )

third_party/amd/lib/TritonAMDGPUToLLVM/TargetInfo.h

@@ -24,6 +24,10 @@ class TargetInfo : public mlir::triton::TargetInfoBase {
 
   bool supportMaximumMinimum() const override;
 
+  bool usesTargetThreadIdLowering() const override { return true; }
+  bool usesTargetBarrierLowering() const override { return true; }
+  bool usesTargetShuffleLowering() const override { return true; }
+
   Value getClusterCTAId(RewriterBase &rewriter, Location loc) const override;
 
   Value ballot(RewriterBase &rewriter, Location loc, Type type,

third_party/amd/lib/TritonAMDGPUToLLVM/TritonGPUToLLVM.cpp

@@ -27,6 +27,12 @@
 #include "triton/Dialect/TritonGPU/IR/Dialect.h"
 #include "triton/Dialect/TritonNvidiaGPU/IR/Dialect.h"
 
+#ifdef __TLE__
+#include "tle/dialect/include/Conversion/TleToLLVM/ExclusiveCumsumOpToLLVM.h"
+#include "tle/dialect/include/IR/Dialect.h"
+#include "tle/dialect/include/Transforms/PatternTleToLLVM.h"
+#endif
+
 namespace mlir::triton {
 #define GEN_PASS_DEF_CONVERTTRITONAMDGPUTOLLVM
 #include "TritonAMDGPUToLLVM/Passes.h.inc"
@@ -63,6 +69,25 @@ class TritonLLVMConversionTarget : public ConversionTarget {
   }
 };
 
+#ifdef __TLE__
+// Conversion target for the dedicated TLE-to-LLVM partial conversion. Only
+// the tile-level extension ops that have AMD lowering patterns are illegal
+// (must be converted); every other op (including unsupported TLE ops) stays
+// legal and is reported later by the make_llir residual-op guard.
+class TleLLVMConversionTarget : public ConversionTarget {
+public:
+  explicit TleLLVMConversionTarget(MLIRContext &ctx) : ConversionTarget(ctx) {
+    addLegalDialect<LLVM::LLVMDialect, ROCDL::ROCDLDialect,
+                    NVVM::NVVMDialect>();
+    addIllegalOp<mlir::triton::tle::ExtractTileOp,
+                 mlir::triton::tle::InsertTileOp,
+                 mlir::triton::tle::ExclusiveCumsumOp>();
+    addLegalOp<mlir::UnrealizedConversionCastOp>();
+    markUnknownOpDynamicallyLegal([](Operation *) -> bool { return true; });
+  }
+};
+#endif
+
 class TritonAMDGPUToLLVMTypeConverter : public TritonGPUToLLVMTypeConverter {
 public:
   TritonAMDGPUToLLVMTypeConverter(MLIRContext *ctx,
@@ -176,6 +201,28 @@ struct ConvertTritonAMDGPUToLLVM
     // Make benefit for AMD specific patterns higher so they apply before common
     // patterns
     int AMDBenefit = commonBenefit + 1;
+
+#ifdef __TLE__
+    // Lower the supported tile-level extension (TLE) ops (extract_tile /
+    // insert_tile / exclusive_cumsum) to LLVM via the backend-agnostic
+    // conversion patterns, in a dedicated partial conversion (mirrors the
+    // NVIDIA / HCU path). Unsupported TLE ops pass through and are reported
+    // by the make_llir residual-op guard.
+    {
+      TleLLVMConversionTarget tleTarget(*context);
+      RewritePatternSet tlePatterns(context);
+      mlir::triton::tle::populateExtractTileOpToLLVMPatterns(
+          typeConverter, tlePatterns, targetInfo, commonBenefit);
+      mlir::triton::tle::populateInsertTileOpToLLVMPatterns(
+          typeConverter, tlePatterns, targetInfo, commonBenefit);
+      mlir::triton::tle::populateExclusiveCumsumOpToLLVMPatterns(
+          typeConverter, targetInfo, tlePatterns, commonBenefit);
+      if (failed(
+              applyPartialConversion(mod, tleTarget, std::move(tlePatterns))))
+        return signalPassFailure();
+    }
+#endif
+
     auto populatePatterns1 = [&](auto populateFunc, int benefit) {
       populateFunc(typeConverter, patterns, axisInfoAnalysis, allocation,
                    benefit);

third_party/amd/python/test/test_tle_compiler.py

@@ -0,0 +1,36 @@
+import pytest
+
+from triton.backends.amd.compiler import reject_residual_tle_ops
+
+
+class FakeOp:
+
+    def __init__(self, name):
+        self.name = name
+
+    def get_name(self):
+        return self.name
+
+
+class FakeModule:
+
+    def __init__(self, *names):
+        self.ops = [FakeOp(name) for name in names]
+
+    def walk(self, callback):
+        for op in self.ops:
+            callback(op)
+
+
+def test_reject_residual_tle_ops_ignores_non_tle_ops():
+    reject_residual_tle_ops(FakeModule("tt.load", "llvm.store"), "gfx1201")
+
+
+def test_reject_residual_tle_ops_reports_sorted_unique_names():
+    mod = FakeModule("tle.remote_pointers", "tt.load", "tle.foo", "tle.remote_pointers")
+
+    with pytest.raises(
+            NotImplementedError,
+            match=r"arch gfx1201: tle\.foo, tle\.remote_pointers",
+    ):
+        reject_residual_tle_ops(mod, "gfx1201")

third_party/hcu/include/triton/Conversion/TritonGPUToLLVM/TargetInfoBase.h

@@ -112,6 +112,10 @@ class TargetInfoBase {
   virtual bool isCuda() const { return false; }
   virtual bool isHCU() const { return false; }
 
+  virtual bool usesTargetThreadIdLowering() const { return false; }
+  virtual bool usesTargetBarrierLowering() const { return false; }
+  virtual bool usesTargetShuffleLowering() const { return false; }
+
   // Annotate target specific information to local load operations during
   // lowering to LLVM. `llLoadOp` is the generated LLVM load op.
   virtual void localLoadOpAnnotation(triton::gpu::LocalLoadOp localLoadOp,

third_party/hcu/lib/TritonHCUGPUToLLVM/TargetInfo.h

@@ -90,6 +90,9 @@ class TargetInfo : public mlir::triton::TargetInfoBase {
   bool supportVectorizedAtomics() const override;
 
   bool isHCU() const override { return true; }
+  bool usesTargetThreadIdLowering() const override { return true; }
+  bool usesTargetBarrierLowering() const override { return true; }
+  bool usesTargetShuffleLowering() const override { return true; }
 
   // Returns true if the target supports per lane addresses into LDS for
   // direct-to-lds loads. Some architectures (e.g. GFX9) do not support