[TLE][AMD] Tighten backend-specific lowering checks

Use FlagTree backend identity helpers in the shared cumsum tests and derive the warp size from the active target. Reject unsupported TLE and accidental NVVM operations during AMD LLVM conversion.

Author: zihaomu

python/test/tle/unit/test_tle_cumsum.py

@@ -10,6 +10,8 @@
 import triton
 import triton.language as tl
 import triton.experimental.tle.language as tle
+from triton._flagtree_backend import FLAGTREE_BACKEND
+from triton._internal_testing import get_current_target, is_cuda, is_hip
 
 
 def _is_enflame_backend():
@@ -18,32 +20,15 @@ def _is_enflame_backend():
 
 
 def _is_hcu_backend():
-    try:
-        driver = triton.runtime.driver.active
-        return type(driver).__module__.startswith("triton.backends.hcu")
-    except Exception:
-        return False
+    return FLAGTREE_BACKEND == "hcu"
 
 
 _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
+threads_per_warp = get_current_target().warp_size
 
 
 def _is_amd_hip_backend():
-    try:
-        driver = triton.runtime.driver.active
-        return type(driver).__module__.startswith("triton.backends.amd")
-    except Exception:
-        return False
+    return is_hip() and not FLAGTREE_BACKEND
 
 
 def _require_cuda():
@@ -223,7 +208,7 @@ def test_tle_cumsum_exclusive_and_total(dtype, n, block, reverse, num_warps):
 
 
 @pytest.mark.skipif(
-    not _is_nvidia_cuda_backend(),
+    not is_cuda(),
     reason="PTX-specific regression guard requires NVIDIA CUDA backend",
 )
 def test_tle_cumsum_ptx_fastpath_regression_guard():

third_party/amd/lib/TritonAMDGPUToLLVM/TritonGPUToLLVM.cpp

@@ -70,18 +70,13 @@ 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.
+// Reject unsupported TLE operations and accidental CUDA lowering in the
+// dedicated AMD TLE-to-LLVM partial conversion.
 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>();
+    addLegalDialect<LLVM::LLVMDialect, ROCDL::ROCDLDialect>();
+    addIllegalDialect<NVVM::NVVMDialect, mlir::triton::tle::TleDialect>();
     addLegalOp<mlir::UnrealizedConversionCastOp>();
     markUnknownOpDynamicallyLegal([](Operation *) -> bool { return true; });
   }
@@ -204,10 +199,9 @@ struct ConvertTritonAMDGPUToLLVM
 
 #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 remain legal in this focused
-    // conversion.
+    // insert_tile / exclusive_cumsum) via the backend-agnostic conversion
+    // patterns. The dedicated partial conversion rejects unsupported TLE ops
+    // and accidental NVVM emission.
     {
       TleLLVMConversionTarget tleTarget(*context);
       RewritePatternSet tlePatterns(context);