diff --git a/cmake/FlagTreeOptions.cmake b/cmake/FlagTreeOptions.cmake index 6b0c9ca59f..531ae89af7 100644 --- a/cmake/FlagTreeOptions.cmake +++ b/cmake/FlagTreeOptions.cmake @@ -468,6 +468,7 @@ function(flagtree_add_tle_generated_header_dependencies) # parallel build can compile those libraries before the generated .inc files. foreach(_flagtree_tle_header_target IN ITEMS TritonAnalysis + TritonAMDAnalysis TritonToTritonGPU TritonGPUTransforms TritonNvidiaGPUTransforms diff --git a/python/test/tle/unit/test_tle_cumsum.py b/python/test/tle/unit/test_tle_cumsum.py index 60e2f88dea..83bb68a6b7 100644 --- a/python/test/tle/unit/test_tle_cumsum.py +++ b/python/test/tle/unit/test_tle_cumsum.py @@ -11,6 +11,7 @@ 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(): @@ -19,12 +20,15 @@ def _is_enflame_backend(): def _is_hcu_backend(): - target = triton.runtime.driver.active.get_current_target() - return target.backend == "hip" + 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 +threads_per_warp = get_current_target().warp_size if is_hip() else 32 + + +def _is_amd_hip_backend(): + return is_hip() and not FLAGTREE_BACKEND def _require_cuda(): @@ -203,9 +207,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_cuda(), + 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 +280,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 +302,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 diff --git a/test/Conversion/amd/tle_tile_ops_to_llvm.mlir b/test/Conversion/amd/tle_tile_ops_to_llvm.mlir new file mode 100644 index 0000000000..de680d5468 --- /dev/null +++ b/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} : 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} : 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) { + %exclusive, %total = "tle.exclusive_cumsum"(%arg0) {axis = 0 : i32, reverse = false} : (tensor<128xi32, #blocked>) -> (tensor<128xi32, #blocked>, i32) + tt.store %out, %total : !tt.ptr + tt.return + } +} diff --git a/third_party/amd/lib/Analysis/AMDGPUAllocation.cpp b/third_party/amd/lib/Analysis/AMDGPUAllocation.cpp index 8df56815ce..d203631386 100644 --- a/third_party/amd/lib/Analysis/AMDGPUAllocation.cpp +++ b/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 +#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(op)) { + auto srcTy = dyn_cast(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(std::max(1, getBitwidth(srcTy) / 8)); + int64_t numWarps = std::max(1, triton::gpu::lookupNumWarps(op)); + uint64_t totalBytes = (static_cast(axisExtent) + + static_cast(numWarps) + 1ull) * + elemBytes; + if (totalBytes > std::numeric_limits::max()) + return 0; + return static_cast(totalBytes); + } + if (auto extractTileOp = dyn_cast(op)) { + auto dstTy = dyn_cast(extractTileOp.getType()); + if (!dstTy) + return 0; + return static_cast(dstTy.getNumElements() * + (getBitwidth(dstTy) / 8)); + } + if (auto insertTileOp = dyn_cast(op)) { + auto tileTy = dyn_cast(insertTileOp.getTile().getType()); + if (!tileTy) + return 0; + return static_cast(tileTy.getNumElements() * + (getBitwidth(tileTy) / 8)); + } +#endif + return defaultAllocationAnalysisScratchSizeFn(op); } diff --git a/third_party/amd/lib/TritonAMDGPUToLLVM/CMakeLists.txt b/third_party/amd/lib/TritonAMDGPUToLLVM/CMakeLists.txt index af01fc73f9..5fc468754f 100644 --- a/third_party/amd/lib/TritonAMDGPUToLLVM/CMakeLists.txt +++ b/third_party/amd/lib/TritonAMDGPUToLLVM/CMakeLists.txt @@ -1,3 +1,14 @@ +if(FLAGTREE_TLE) + set(_TLE_LIBS + TritonTLEAnalysis + TritonNVIDIAGPUToLLVM + TleToLLVM + TritonTLETransforms + ) +else() + set(_TLE_LIBS "") +endif() + add_triton_library(TritonAMDGPUToLLVM AsyncUtility.cpp AtomicRMWOpsEmitter.cpp @@ -40,4 +51,5 @@ add_triton_library(TritonAMDGPUToLLVM LLVMCore LLVMPasses LLVMSupport + ${_TLE_LIBS} ) diff --git a/third_party/amd/lib/TritonAMDGPUToLLVM/TritonGPUToLLVM.cpp b/third_party/amd/lib/TritonAMDGPUToLLVM/TritonGPUToLLVM.cpp index c4bb2dab1c..f2d8a6219e 100644 --- a/third_party/amd/lib/TritonAMDGPUToLLVM/TritonGPUToLLVM.cpp +++ b/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,20 @@ class TritonLLVMConversionTarget : public ConversionTarget { } }; +#ifdef __TLE__ +// 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(); + addIllegalDialect(); + addLegalOp(); + markUnknownOpDynamicallyLegal([](Operation *) -> bool { return true; }); + } +}; +#endif + class TritonAMDGPUToLLVMTypeConverter : public TritonGPUToLLVMTypeConverter { public: TritonAMDGPUToLLVMTypeConverter(MLIRContext *ctx, @@ -176,6 +196,27 @@ 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) via the backend-agnostic conversion + // patterns. The dedicated partial conversion rejects unsupported TLE ops + // and accidental NVVM emission. + { + 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); diff --git a/third_party/tle/dialect/lib/Conversion/TleToLLVM/ExclusiveCumsumOpToLLVM.cpp b/third_party/tle/dialect/lib/Conversion/TleToLLVM/ExclusiveCumsumOpToLLVM.cpp index 0e49d07d66..5b5943483f 100644 --- a/third_party/tle/dialect/lib/Conversion/TleToLLVM/ExclusiveCumsumOpToLLVM.cpp +++ b/third_party/tle/dialect/lib/Conversion/TleToLLVM/ExclusiveCumsumOpToLLVM.cpp @@ -122,7 +122,7 @@ static Value createWarpScanStepI32(Location loc, ConversionPatternRewriter &rewriter, const TargetInfoBase &targetInfo, Value val, int offset, Value laneId, Type elemTy) { - if (targetInfo.isHCU()) { + if (!targetInfo.isCuda()) { auto b = TritonLLVMOpBuilder(loc, rewriter); Value shuffled = targetInfo.shuffleUp(rewriter, loc, val, offset); Value pred = b.icmp_sge(laneId, b.i32_val(offset)); diff --git a/third_party/tle/dialect/lib/Transforms/ExtractTileToLLVM.cpp b/third_party/tle/dialect/lib/Transforms/ExtractTileToLLVM.cpp index 5fca3f036f..b4869af1b9 100644 --- a/third_party/tle/dialect/lib/Transforms/ExtractTileToLLVM.cpp +++ b/third_party/tle/dialect/lib/Transforms/ExtractTileToLLVM.cpp @@ -304,7 +304,7 @@ lowerExtractTileViaSMEM(ExtractTileOp op, ExtractTileOp::Adaptor adaptor, // ------------------------------------------------------------------ // Step 4: __syncthreads() -- ensure all writes are visible // ------------------------------------------------------------------ - if (targetInfo.isHCU()) + if (!targetInfo.isCuda()) targetInfo.barrier(loc, rewriter, /*isWarpSync=*/false); else rewriter.create(loc); @@ -372,7 +372,7 @@ lowerExtractTileViaSMEM(ExtractTileOp op, ExtractTileOp::Adaptor adaptor, // ------------------------------------------------------------------ // Step 6: __syncthreads() -- allow SMEM reuse after reads complete // ------------------------------------------------------------------ - if (targetInfo.isHCU()) + if (!targetInfo.isCuda()) targetInfo.barrier(loc, rewriter, /*isWarpSync=*/false); else rewriter.create(loc); diff --git a/third_party/tle/dialect/lib/Transforms/InsertTileToLLVM.cpp b/third_party/tle/dialect/lib/Transforms/InsertTileToLLVM.cpp index 792ff32deb..5493530127 100644 --- a/third_party/tle/dialect/lib/Transforms/InsertTileToLLVM.cpp +++ b/third_party/tle/dialect/lib/Transforms/InsertTileToLLVM.cpp @@ -309,7 +309,7 @@ lowerInsertTileViaSMEMDynamic(InsertTileOp op, InsertTileOp::Adaptor adaptor, rewriter.create(loc, tileVals[i], sp, elemBytes); } // Synchronize threads after tile store. - if (targetInfo.isHCU()) + if (!targetInfo.isCuda()) targetInfo.barrier(loc, rewriter, /*isWarpSync=*/false); else rewriter.create(loc); @@ -362,7 +362,7 @@ lowerInsertTileViaSMEMDynamic(InsertTileOp op, InsertTileOp::Adaptor adaptor, resultVals.push_back(merged); } - if (targetInfo.isHCU()) + if (!targetInfo.isCuda()) targetInfo.barrier(loc, rewriter, /*isWarpSync=*/false); else rewriter.create(loc); diff --git a/third_party/tle/dialect/lib/Transforms/TleTileToLLVMUtils.cpp b/third_party/tle/dialect/lib/Transforms/TleTileToLLVMUtils.cpp index ae423db13c..bbff0d45d2 100644 --- a/third_party/tle/dialect/lib/Transforms/TleTileToLLVMUtils.cpp +++ b/third_party/tle/dialect/lib/Transforms/TleTileToLLVMUtils.cpp @@ -115,7 +115,7 @@ SmallVector computeThreadOffsets(Location loc, auto i32Ty = rewriter.getIntegerType(32); Value threadId; - if (targetInfo.isHCU()) + if (!targetInfo.isCuda()) threadId = getThreadId(rewriter, loc); else threadId = rewriter.create(loc, i32Ty);