diff --git a/include/triton/Dialect/TritonGPU/IR/TritonGPUOps.td b/include/triton/Dialect/TritonGPU/IR/TritonGPUOps.td index 43dfb56240..3eacf4ac79 100644 --- a/include/triton/Dialect/TritonGPU/IR/TritonGPUOps.td +++ b/include/triton/Dialect/TritonGPU/IR/TritonGPUOps.td @@ -598,6 +598,40 @@ def TTG_LocalBarrierOp : TTG_Op<"local_barrier"> { } // flagtree tle +#ifdef __TLE__ +def TTG_TMACopyOp : TTG_Op<"tma_copy", [ + Pure, + MemoryEffects<[MemWrite]>, + AttrSizedOperandSegments +]> { + let summary = "self-defined tma_copy operation"; + let description = [{ + 'ttg.tma_copy' triton copy op is designed to copy data between memory regions. + Example: + ```mlir + ttg.tma_copy %src, %dst, %shape : tensor<128xf32> + ``` + + A global-to-shared TMA load may optionally carry an explicit completion + barrier. In that form, lowering emits barrier_expect and async TMA copy + against the provided barrier, while the user controls the wait point. + }]; + let arguments = (ins + Arg, MemWrite]>:$src, + Arg, MemWrite]>:$dst, + Variadic:$indices, + Optional:$barrier, + OptionalAttr:$expect_bytes + ); + //assemble + let assemblyFormat = [{ + $src `,` $dst `,` `[` $indices `]` + (`,` `barrier` $barrier^)? + attr-dict `:` type($src) `,` type($dst) + (`,` qualified(type($barrier))^)? + }]; +} +#else def TTG_TMACopyOp : TTG_Op<"tma_copy", [Pure, MemoryEffects<[MemWrite]>]> { let summary = "self-defined tma_copy operation"; let description = [{ @@ -615,5 +649,6 @@ def TTG_TMACopyOp : TTG_Op<"tma_copy", [Pure, MemoryEffects<[MemWrite]>]> { //assemble let assemblyFormat = "$src `,` $dst `,` `[` $indices `]` attr-dict `:` type($src) `,` type($dst)"; } +#endif #endif // TRITONGPU_OPS diff --git a/include/triton/Dialect/TritonNvidiaGPU/IR/TritonNvidiaGPUOps.td b/include/triton/Dialect/TritonNvidiaGPU/IR/TritonNvidiaGPUOps.td index c9d89eefa1..44526516c2 100644 --- a/include/triton/Dialect/TritonNvidiaGPU/IR/TritonNvidiaGPUOps.td +++ b/include/triton/Dialect/TritonNvidiaGPU/IR/TritonNvidiaGPUOps.td @@ -353,6 +353,24 @@ def TTNG_ArriveBarrierOp : TTNG_Op<"arrive_barrier"> { } #endif +#ifdef __TLE__ +def TTNG_NamedBarrierArriveOp : TTNG_Op<"arrive_barrier_named", []> { + let summary = "named barrier arrive"; + + let arguments = (ins I32:$bar, I32:$numThreads); + + let assemblyFormat = "$bar `,` $numThreads attr-dict `:` type(operands)"; +} + +def TTNG_NamedBarrierWaitOp : TTNG_Op<"wait_barrier_named", []> { + let summary = "named barrier wait"; + + let arguments = (ins I32:$bar, I32:$numThreads); + + let assemblyFormat = "$bar `,` $numThreads attr-dict `:` type(operands)"; +} +#endif + def TTNG_AsyncCopyMbarrierArriveOp : TTNG_Op<"async_copy_mbarrier_arrive"> { let summary = "arrive on mbarrier once all previously issued copies are completed"; let arguments = (ins diff --git a/lib/Conversion/TritonToTritonGPU/TritonToTritonGPUPass.cpp b/lib/Conversion/TritonToTritonGPU/TritonToTritonGPUPass.cpp index cd81ac070f..249931358b 100644 --- a/lib/Conversion/TritonToTritonGPU/TritonToTritonGPUPass.cpp +++ b/lib/Conversion/TritonToTritonGPU/TritonToTritonGPUPass.cpp @@ -949,6 +949,7 @@ void populateTleRawPatterns(TritonGPUTypeConverter &typeConverter, TleInsertTileOpPattern, GenericOpPattern, GenericOpPattern, GenericOpPattern, + GenericOpPattern, GenericOpPattern, GenericOpPattern, GenericOpPattern, GenericOpPattern, diff --git a/lib/Dialect/TritonGPU/Transforms/CMakeLists.txt b/lib/Dialect/TritonGPU/Transforms/CMakeLists.txt index 6cdd688943..6684daaef6 100644 --- a/lib/Dialect/TritonGPU/Transforms/CMakeLists.txt +++ b/lib/Dialect/TritonGPU/Transforms/CMakeLists.txt @@ -3,6 +3,7 @@ if(FLAGTREE_TLE) list(APPEND TritonGPUTransformsTleSources ${CMAKE_SOURCE_DIR}/third_party/tle/dialect/lib/Transforms/EncodingRematerialization.cpp Pipeliner/TleWGMMAAnalysis.cpp + Pipeliner/TleWGMMAUserPromisePipeline.cpp ) endif() diff --git a/lib/Dialect/TritonGPU/Transforms/Pipeliner/SoftwarePipeliner.cpp b/lib/Dialect/TritonGPU/Transforms/Pipeliner/SoftwarePipeliner.cpp index d36a115bf0..812512b164 100644 --- a/lib/Dialect/TritonGPU/Transforms/Pipeliner/SoftwarePipeliner.cpp +++ b/lib/Dialect/TritonGPU/Transforms/Pipeliner/SoftwarePipeliner.cpp @@ -1,5 +1,8 @@ #include "mlir/Dialect/Tensor/IR/Tensor.h" #include "mlir/Dialect/UB/IR/UBOps.h" +#ifdef __TLE__ +#include "TleWGMMAAnalysis.h" +#endif #include "mlir/IR/TypeUtilities.h" #include "mlir/IR/Verifier.h" #include "mlir/Interfaces/SideEffectInterfaces.h" @@ -34,14 +37,44 @@ namespace gpu { #define GEN_PASS_DEF_TRITONGPUPIPELINE #include "triton/Dialect/TritonGPU/Transforms/Passes.h.inc" -static void pipelineWgmma(ModuleOp moduleOp, unsigned numStages) { +#ifdef __TLE__ +static constexpr llvm::StringLiteral + kTleWgmmaPipelineModeAttr("tle.wgmma_pipeline_mode"); +static constexpr llvm::StringLiteral kTleWgmmaCompilerAutoMode("compiler_auto"); +static constexpr llvm::StringLiteral kTleWgmmaUserPromiseMode("user_promise"); +#endif + +static LogicalResult pipelineWgmma(ModuleOp moduleOp, unsigned numStages) { +#ifdef __TLE__ + StringRef mode = kTleWgmmaCompilerAutoMode; + if (auto attr = + moduleOp->getAttrOfType(kTleWgmmaPipelineModeAttr)) + mode = attr.getValue(); + + if (mode != kTleWgmmaCompilerAutoMode && mode != kTleWgmmaUserPromiseMode) { + moduleOp.emitError("TLE WGMMA pipeline mode module attribute '") + << kTleWgmmaPipelineModeAttr << "' must be '" + << kTleWgmmaCompilerAutoMode << "' or '" << kTleWgmmaUserPromiseMode + << "', got '" << mode << "'"; + return failure(); + } +#endif + SmallVector loops; moduleOp->walk([&](scf::ForOp forOp) { loops.push_back(forOp); }); for (scf::ForOp forOp : loops) { - if (getNumStagesOrDefault(forOp, numStages) >= 1) + if (getNumStagesOrDefault(forOp, numStages) >= 1) { +#ifdef __TLE__ + if (mode == kTleWgmmaUserPromiseMode) { + mlir::triton::gpu::detail::scheduleTleWgmmaUserPromisePipeline(forOp); + continue; + } +#endif mlir::triton::asyncLaunchDots(forOp); + } } + return success(); } static bool hasMMAv5WaitsInLastStage(scf::ForOp forOp, @@ -194,7 +227,8 @@ struct PipelinePass : public impl::TritonGPUPipelineBase { // Cleanup the IR from the pipeline attributes. removePipeliningAttributes(moduleOp); - pipelineWgmma(moduleOp, numStages); + if (failed(pipelineWgmma(moduleOp, numStages))) + return signalPassFailure(); // schedule the waits mlir::triton::updateWaits(getOperation()); diff --git a/lib/Dialect/TritonGPU/Transforms/Pipeliner/TleWGMMAAnalysis.h b/lib/Dialect/TritonGPU/Transforms/Pipeliner/TleWGMMAAnalysis.h index e6bbdfd56b..39a30d5b6c 100644 --- a/lib/Dialect/TritonGPU/Transforms/Pipeliner/TleWGMMAAnalysis.h +++ b/lib/Dialect/TritonGPU/Transforms/Pipeliner/TleWGMMAAnalysis.h @@ -60,6 +60,7 @@ class TleWgmmaScheduleAnalysis { }; void scheduleTleWgmmaAsyncLaunch(scf::ForOp forOp); +void scheduleTleWgmmaUserPromisePipeline(scf::ForOp forOp); } // namespace mlir::triton::gpu::detail diff --git a/lib/Dialect/TritonGPU/Transforms/Pipeliner/TleWGMMAUserPromisePipeline.cpp b/lib/Dialect/TritonGPU/Transforms/Pipeliner/TleWGMMAUserPromisePipeline.cpp new file mode 100644 index 0000000000..e155d88f63 --- /dev/null +++ b/lib/Dialect/TritonGPU/Transforms/Pipeliner/TleWGMMAUserPromisePipeline.cpp @@ -0,0 +1,40 @@ +#ifdef __TLE__ +#include "TleWGMMAAnalysis.h" +#include "mlir/Dialect/SCF/IR/SCF.h" +#include "mlir/IR/Builders.h" +#include "mlir/IR/PatternMatch.h" +#include "mlir/Support/LLVM.h" +#include "triton/Dialect/TritonNvidiaGPU/IR/Dialect.h" +#include "llvm/ADT/STLExtras.h" + +using namespace mlir; +namespace ttng = mlir::triton::nvidia_gpu; + +namespace mlir::triton::gpu::detail { + +static constexpr llvm::StringLiteral + kTleExplicitWgmmaCommitAttr("tle.explicit_wgmma_commit"); + +void scheduleTleWgmmaUserPromisePipeline(scf::ForOp forOp) { + IRRewriter builder(forOp.getContext()); + SmallVector dots; + forOp.getBody()->walk([&](ttng::WarpGroupDotOp dot) { + if (dot->getParentOfType() == forOp) + dots.push_back(dot); + }); + + for (ttng::WarpGroupDotOp dot : llvm::make_early_inc_range(dots)) { + dot.setIsAsync(true); + dot->setAttr(kTleExplicitWgmmaCommitAttr, builder.getUnitAttr()); + + Operation *next = dot->getNextNode(); + if (next && isa(next)) + continue; + + builder.setInsertionPointAfter(dot); + ttng::WarpGroupDotCommitOp::create(builder, dot.getLoc()); + } +} + +} // namespace mlir::triton::gpu::detail +#endif // __TLE__ diff --git a/lib/Dialect/TritonGPU/Transforms/Pipeliner/WGMMAPipeline.cpp b/lib/Dialect/TritonGPU/Transforms/Pipeliner/WGMMAPipeline.cpp index 44903ac156..d8a99a20a3 100644 --- a/lib/Dialect/TritonGPU/Transforms/Pipeliner/WGMMAPipeline.cpp +++ b/lib/Dialect/TritonGPU/Transforms/Pipeliner/WGMMAPipeline.cpp @@ -671,6 +671,105 @@ findPendingGroupForValue(Value value, return std::nullopt; } +static ttng::WarpGroupDotWaitOp getDefiningPositivePendingWait(Value value) { + while (Operation *def = value.getDefiningOp()) { + if (auto wait = dyn_cast(def)) { + if (wait.getPendings() > 0) + return wait; + return {}; + } + if (!isNoop(def) || def->getNumOperands() != 1 || def->getNumResults() != 1) + return {}; + value = def->getOperand(0); + } + return {}; +} + +static bool isMaterializedByWaitZero(Value value, + llvm::SmallDenseSet &visited) { + if (!visited.insert(value).second) + return true; + + bool hasUse = false; + for (OpOperand &use : value.getUses()) { + hasUse = true; + Operation *user = use.getOwner(); + + if (auto wait = dyn_cast(user)) { + if (wait.getPendings() == 0) + continue; + return false; + } + + if (isNoop(user) && user->getNumResults() == 1) { + if (!isMaterializedByWaitZero(user->getResult(0), visited)) + return false; + continue; + } + + return false; + } + + return hasUse; +} + +static bool isMaterializedByWaitZero(Value value) { + llvm::SmallDenseSet visited; + return isMaterializedByWaitZero(value, visited); +} + +static std::optional +findYieldOperandForPendingGroup(scf::YieldOp yield, + const PendingSharedWgmmaGroup &pending) { + std::optional yieldIndex; + for (auto indexed : llvm::enumerate(yield.getOperands())) { + unsigned index = static_cast(indexed.index()); + Value yielded = indexed.value(); + ttng::WarpGroupDotWaitOp wait = getDefiningPositivePendingWait(yielded); + if (!wait) + continue; + + bool carriesPending = + llvm::any_of(pending.dots, [&](ttng::WarpGroupDotOp dot) { + return valueDependsOn(yielded, dot.getResult()); + }); + if (!carriesPending) + continue; + + if (yieldIndex && *yieldIndex != index) + return std::nullopt; + yieldIndex = index; + } + return yieldIndex; +} + +static bool canCarryPendingGroupsThroughForYield( + scf::YieldOp yield, ArrayRef pendingGroups) { + auto forOp = dyn_cast(yield->getParentOp()); + if (!forOp || forOp.getBody() != yield->getBlock()) + return false; + if (pendingGroups.empty()) + return false; + + llvm::SmallDenseSet carriedYieldIndices; + for (const PendingSharedWgmmaGroup &pending : pendingGroups) { + std::optional yieldIndex = + findYieldOperandForPendingGroup(yield, pending); + if (!yieldIndex) + return false; + carriedYieldIndices.insert(*yieldIndex); + } + + for (unsigned yieldIndex : carriedYieldIndices) { + if (yieldIndex >= forOp.getNumResults()) + return false; + if (!isMaterializedByWaitZero(forOp.getResult(yieldIndex))) + return false; + } + + return true; +} + static bool isAllowedAccumulatorChainUse(Operation *op, OpOperand &use, const TleWgmmaScheduleAnalysis &analysis) { @@ -909,10 +1008,14 @@ static void scheduleTleWgmmaWaitsInBlock( Operation *terminator = block->getTerminator(); if (!terminator) return; - if (deferLoopCarriedYield && isa(terminator)) { - if (!pendingGroups.empty()) + if (auto yield = dyn_cast(terminator)) { + // Preserve explicit TLE accumulator pipelining across the loop backedge + // when the loop result is closed by a matching wait_group 0 outside. + if (deferLoopCarriedYield && + canCarryPendingGroupsThroughForYield(yield, pendingGroups)) { insertWgmmaDepthWaitAfterLastPendingDot(pendingGroups); - return; + return; + } } drainForMaterializedOperands(terminator, analysis, pendingGroups); if (!pendingGroups.empty()) diff --git a/python/test/tle/integration/test_tle_tma_copy.py b/python/test/tle/integration/test_tle_tma_copy.py index 46c6bc9c05..ff4583c9ab 100644 --- a/python/test/tle/integration/test_tle_tma_copy.py +++ b/python/test/tle/integration/test_tle_tma_copy.py @@ -21,6 +21,11 @@ def _is_enflame_backend(): return target.backend == "gcu" +def _is_nvidia_backend(): + target = triton.runtime.driver.active.get_current_target() + return target.backend == "cuda" + + def _has_hopper_gpu() -> bool: if _is_enflame_backend(): # Assume Enflame backend has Hopper support for testing purposes @@ -97,6 +102,55 @@ def elementwise_add(A, B, C, XBLOCK=32, YBLOCK=64): return elementwise_tma_add_kernel[grid](A, B, C, xnumel, ynumel, XBLOCK, YBLOCK) +@triton.jit +def elementwise_tma_add_explicit_barrier_kernel( + a_desc, + b_desc, + c_desc, + xnumel, + ynumel, + XBLOCK: tl.constexpr, + YBLOCK: tl.constexpr, + TILE_BYTES: tl.constexpr, +): + pid = tl.program_id(0) + + a_smem = tle.gpu.alloc([XBLOCK, YBLOCK], dtype=tl.float32, layout=None, scope=tle.gpu.smem) + b_smem = tle.gpu.alloc([XBLOCK, YBLOCK], dtype=tl.float32, layout=None, scope=tle.gpu.smem) + c_smem = tle.gpu.alloc([XBLOCK, YBLOCK], dtype=tl.float32, layout=None, scope=tle.gpu.smem) + + row_ids = tl.arange(0, XBLOCK)[:, None] + col_ids = tl.arange(0, YBLOCK)[None, :] + row_ids = tl.broadcast_to(row_ids, (XBLOCK, YBLOCK)) + col_ids = tl.broadcast_to(col_ids, (XBLOCK, YBLOCK)) + a_smem_ptrs = tle.gpu.local_ptr(a_smem, (row_ids, col_ids)) + b_smem_ptrs = tle.gpu.local_ptr(b_smem, (row_ids, col_ids)) + c_smem_ptrs = tle.gpu.local_ptr(c_smem, (row_ids, col_ids)) + + a_bar = tle.gpu.alloc_barrier(expect_bytes=TILE_BYTES) + b_bar = tle.gpu.alloc_barrier(expect_bytes=TILE_BYTES) + + for yoff in range(0, ynumel, YBLOCK): + phase = (yoff // YBLOCK) & 1 + tle.gpu.copy(a_desc, a_smem, [XBLOCK, YBLOCK], [pid * XBLOCK, yoff], barrier=a_bar) + tle.gpu.copy(b_desc, b_smem, [XBLOCK, YBLOCK], [pid * XBLOCK, yoff], barrier=b_bar) + tle.gpu.barrier_wait(a_bar, phaseIdx=phase) + tle.gpu.barrier_wait(b_bar, phaseIdx=phase) + + aval = tl.load(a_smem_ptrs) + bval = tl.load(b_smem_ptrs) + c_val = aval + bval + tl.store(c_smem_ptrs, c_val) + tle.gpu.copy(c_smem, c_desc, [XBLOCK, YBLOCK], [pid * XBLOCK, yoff]) + + +def elementwise_add_explicit_barrier(A, B, C, XBLOCK=32, YBLOCK=64): + xnumel, ynumel = 512, 512 + grid = (triton.cdiv(xnumel, XBLOCK), ) + return elementwise_tma_add_explicit_barrier_kernel[grid](A, B, C, xnumel, ynumel, XBLOCK, YBLOCK, + XBLOCK * YBLOCK * 4) + + class TestTLETmaCopy: """TLE TMA Copy Integration Tests""" @@ -122,6 +176,28 @@ def test_tma_copy_basic(self): expected = a + b torch.testing.assert_close(c, expected, atol=1e-5, rtol=1e-5) + @pytest.mark.skipif(not _is_nvidia_backend(), reason="Explicit TMA completion barriers require NVIDIA backend") + def test_tma_copy_explicit_barrier(self): + """Test TMA load with user-provided completion barriers.""" + torch.manual_seed(43) + + xnumel, ynumel = 512, 512 + XBLOCK, YBLOCK = 32, 64 + + a = torch.randn(xnumel, ynumel, device="cuda", dtype=torch.float32) + b = torch.randn(xnumel, ynumel, device="cuda", dtype=torch.float32) + c = torch.empty_like(a, device="cuda", dtype=torch.float32) + + from triton.tools.tensor_descriptor import TensorDescriptor + a_tma = TensorDescriptor.from_tensor(a, block_shape=[XBLOCK, YBLOCK]) + b_tma = TensorDescriptor.from_tensor(b, block_shape=[XBLOCK, YBLOCK]) + c_tma = TensorDescriptor.from_tensor(c, block_shape=[XBLOCK, YBLOCK]) + + elementwise_add_explicit_barrier(a_tma, b_tma, c_tma, XBLOCK, YBLOCK) + + expected = a + b + torch.testing.assert_close(c, expected, atol=1e-5, rtol=1e-5) + def test_tma_copy_different_block_sizes(self): """Test TMA copy with different block sizes""" torch.manual_seed(123) diff --git a/python/test/tle/integration/test_tle_ws_tma_gemm.py b/python/test/tle/integration/test_tle_ws_tma_gemm.py new file mode 100644 index 0000000000..77721e1087 --- /dev/null +++ b/python/test/tle/integration/test_tle_ws_tma_gemm.py @@ -0,0 +1,258 @@ +# flagtree tle +""" +Smoke tests for TLE warp_specialize GEMM with staged TMA completion barriers. + +This intentionally keeps the GEMM to a single output tile while splitting K +across multiple shared-memory slots. The goal is to validate the basic +producer/consumer protocol: + +- default partition waits on per-slot "empty" barriers and issues TMA loads +- TMA completion signals "full" barriers +- worker partition waits on per-slot "full" barriers, computes WGMMA, stores C +- worker arrives on "empty" barriers to release the smem buffers +""" + +import pytest +import torch +import triton +import triton.language as tl +import triton.experimental.tle.language as tle + +torch.backends.cuda.matmul.allow_tf32 = False +torch.backends.cudnn.allow_tf32 = False + + +def _is_nvidia_backend() -> bool: + target = triton.runtime.driver.active.get_current_target() + return target.backend == "cuda" + + +def _has_nvidia_hopper_gpu() -> bool: + return _is_nvidia_backend() and torch.cuda.is_available() and torch.cuda.get_device_capability()[0] >= 9 + + +pytestmark = pytest.mark.skipif( + not _has_nvidia_hopper_gpu(), + reason="warp_specialize TMA WGMMA GEMM requires NVIDIA Hopper (sm90+)", +) + + +@triton.jit +def _slot_phase(k_iter, num_slots: tl.constexpr): + slot = k_iter % num_slots + phase = k_iter // num_slots + return slot, phase + + +@triton.jit +def _ws_tma_multi_slot_producer( + a_desc, + b_desc, + a_smem, + b_smem, + a_empty, + b_empty, + a_full, + b_full, + BLOCK_M: tl.constexpr, + BLOCK_N: tl.constexpr, + BLOCK_K: tl.constexpr, + K_TILES: tl.constexpr, + NUM_SLOTS: tl.constexpr, +): + for k_iter in range(0, K_TILES): + slot, phase = _slot_phase(k_iter, NUM_SLOTS) + k_offset = k_iter * BLOCK_K + + tle.gpu.barrier_wait(a_empty[slot], phaseIdx=phase) + tle.gpu.barrier_wait(b_empty[slot], phaseIdx=phase) + + tle.gpu.copy(a_desc, a_smem.slot(slot), [BLOCK_M, BLOCK_K], [0, k_offset], barrier=a_full[slot]) + tle.gpu.copy(b_desc, b_smem.slot(slot), [BLOCK_K, BLOCK_N], [k_offset, 0], barrier=b_full[slot]) + + for k_iter in range(K_TILES - NUM_SLOTS, K_TILES): + slot, phase = _slot_phase(k_iter, NUM_SLOTS) + release_phase = phase + 1 + + tle.gpu.barrier_wait(a_empty[slot], phaseIdx=release_phase) + tle.gpu.barrier_wait(b_empty[slot], phaseIdx=release_phase) + + +@triton.jit +def _ws_tma_multi_slot_consumer( + a_smem, + b_smem, + a_empty, + b_empty, + a_full, + b_full, + c_ptr, + stride_cm: tl.constexpr, + stride_cn: tl.constexpr, + BLOCK_M: tl.constexpr, + BLOCK_N: tl.constexpr, + K_TILES: tl.constexpr, + NUM_SLOTS: tl.constexpr, +): + acc = tl.zeros((BLOCK_M, BLOCK_N), dtype=tl.float32) + + slot, phase = _slot_phase(0, NUM_SLOTS) + tle.gpu.barrier_wait(a_full[slot], phaseIdx=phase) + tle.gpu.barrier_wait(b_full[slot], phaseIdx=phase) + acc = tle.gpu.wgmma(a_smem.slot(slot), b_smem.slot(slot), acc) + last_slot = slot + last_phase = phase + + for k_iter in range(1, K_TILES): + slot, phase = _slot_phase(k_iter, NUM_SLOTS) + + tle.gpu.barrier_wait(a_full[slot], phaseIdx=phase) + tle.gpu.barrier_wait(b_full[slot], phaseIdx=phase) + + acc = tle.gpu.wgmma(a_smem.slot(slot), b_smem.slot(slot), acc) + acc = tle.gpu.wgmma_wait(1, acc) + tle.gpu.barrier_arrive(a_empty[last_slot], phaseIdx=last_phase) + tle.gpu.barrier_arrive(b_empty[last_slot], phaseIdx=last_phase) + + last_slot = slot + last_phase = phase + + acc = tle.gpu.wgmma_wait(0, acc) + tle.gpu.barrier_arrive(a_empty[last_slot], phaseIdx=last_phase) + tle.gpu.barrier_arrive(b_empty[last_slot], phaseIdx=last_phase) + + offs_m = tl.arange(0, BLOCK_M) + offs_n = tl.arange(0, BLOCK_N) + + c_ptrs = c_ptr + offs_m[:, None] * stride_cm + offs_n[None, :] * stride_cn + tl.store(c_ptrs, acc) + + +@triton.jit +def ws_tma_multi_slot_gemm_kernel( + a_desc, + b_desc, + c_ptr, + stride_cm: tl.constexpr, + stride_cn: tl.constexpr, + BLOCK_M: tl.constexpr, + BLOCK_N: tl.constexpr, + BLOCK_K: tl.constexpr, + A_TILE_BYTES: tl.constexpr, + B_TILE_BYTES: tl.constexpr, + K_TILES: tl.constexpr, + NUM_SLOTS: tl.constexpr, +): + a_smem = tle.gpu.alloc( + [NUM_SLOTS, BLOCK_M, BLOCK_K], + dtype=tl.float16, + layout=None, + scope=tle.gpu.smem, + ) + b_smem = tle.gpu.alloc( + [NUM_SLOTS, BLOCK_K, BLOCK_N], + dtype=tl.float16, + layout=None, + scope=tle.gpu.smem, + ) + + a_empty = tle.gpu.alloc_barriers(num_barriers=NUM_SLOTS, init=tle.gpu.READY) + b_empty = tle.gpu.alloc_barriers(num_barriers=NUM_SLOTS, init=tle.gpu.READY) + a_full = tle.gpu.alloc_barriers(num_barriers=NUM_SLOTS, expect_bytes=A_TILE_BYTES) + b_full = tle.gpu.alloc_barriers(num_barriers=NUM_SLOTS, expect_bytes=B_TILE_BYTES) + + tle.gpu.warp_specialize( + [ + ( + _ws_tma_multi_slot_producer, + ( + a_desc, + b_desc, + a_smem, + b_smem, + a_empty, + b_empty, + a_full, + b_full, + BLOCK_M, + BLOCK_N, + BLOCK_K, + K_TILES, + NUM_SLOTS, + ), + ), + ( + _ws_tma_multi_slot_consumer, + ( + a_smem, + b_smem, + a_empty, + b_empty, + a_full, + b_full, + c_ptr, + stride_cm, + stride_cn, + BLOCK_M, + BLOCK_N, + K_TILES, + NUM_SLOTS, + ), + ), + ], + [4], + [168], + ) + + +def ws_tma_multi_slot_gemm(A, B, C, launch_num_warps, block_k, num_slots): + assert A.ndim == 2 and B.ndim == 2 and C.ndim == 2 + assert A.shape[1] == B.shape[0] + assert C.shape == (A.shape[0], B.shape[1]) + assert A.dtype == torch.float16 and B.dtype == torch.float16 and C.dtype == torch.float32 + + block_m, total_k = A.shape + total_k_b, block_n = B.shape + assert total_k == total_k_b + assert total_k % block_k == 0 + k_tiles = total_k // block_k + assert k_tiles >= num_slots + + from triton.tools.tensor_descriptor import TensorDescriptor + + a_desc = TensorDescriptor.from_tensor(A, block_shape=[block_m, block_k]) + b_desc = TensorDescriptor.from_tensor(B, block_shape=[block_k, block_n]) + return ws_tma_multi_slot_gemm_kernel[(1, )]( + a_desc, + b_desc, + C, + C.stride(0), + C.stride(1), + block_m, + block_n, + block_k, + block_m * block_k * A.element_size(), + block_k * block_n * B.element_size(), + k_tiles, + num_slots, + num_warps=launch_num_warps, + ) + + +class TestTLEWarpSpecializeTmaGemm: + + @pytest.mark.parametrize("launch_num_warps", [4]) + def test_multi_slot_producer_consumer_wgmma(self, launch_num_warps): + torch.manual_seed(2026 + launch_num_warps) + block_m, block_n, block_k = 64, 16, 16 + k_tiles, num_slots = 4, 2 + + a = torch.randn(block_m, block_k * k_tiles, device="cuda", dtype=torch.float16).contiguous() + b = torch.randn(block_k * k_tiles, block_n, device="cuda", dtype=torch.float16).contiguous() + c = torch.empty((block_m, block_n), device="cuda", dtype=torch.float32).contiguous() + + kernel = ws_tma_multi_slot_gemm(a, b, c, launch_num_warps, block_k, num_slots) + assert "tt.call" not in kernel.asm["ttgir"] + + expected = torch.matmul(a.float(), b.float()) + torch.testing.assert_close(c, expected, atol=5e-2, rtol=5e-2) diff --git a/python/test/tle/unit/test_tle.py b/python/test/tle/unit/test_tle.py index adbdcef101..014593a603 100644 --- a/python/test/tle/unit/test_tle.py +++ b/python/test/tle/unit/test_tle.py @@ -12,6 +12,7 @@ import pytest import torch +import inspect import triton.language as tl import triton.experimental.tle.language as tle from triton.language.core import base_value @@ -157,6 +158,7 @@ class _FakeTensor: def __init__(self, handle, ty): self.handle = handle self.type = ty + self.dtype = ty class _FakeBlockType: @@ -173,11 +175,15 @@ def __init__(self): self.memdesc_index_args = None self.swizzled_encoding_args = None self.pipe_create_args = None + self.tma_copy_args = None self.pipe_ops = [] def get_half_ty(self): return "fp16" + def get_int64_ty(self): + return "i64" + def make_swizzled_shared_encoding_attr(self, vector_size, per_phase, max_phase, order, ctas_per_cga, cta_split_num, cta_order): self.swizzled_encoding_args = ( @@ -200,6 +206,9 @@ def create_memdesc_index(self, result_ty, src, index): self.memdesc_index_args = (result_ty, src, index) return "slot_handle" + def create_tma_copy(self, src, dst, offsets, barrier=None, expect_bytes=-1): + self.tma_copy_args = (src, dst, list(offsets), barrier, expect_bytes) + def create_pipe_create(self, fields, capacity, scope, pipe_name, field_names, reader_names, one_shot): self.pipe_create_args = (list(fields), capacity, scope, pipe_name, list(field_names), list(reader_names), one_shot) @@ -241,6 +250,9 @@ def to_tensor(self, value): return TestBufferedTensor._FakeTensor(f"stage_{value}", tl.int32) raise TypeError(f"unsupported fake tensor input: {value!r}") + def _convert_to_ir_values(self, values, require_i64=False): + return [self.to_tensor(value).handle for value in values] + def _make_buffer(self, shape): semantic = self._FakeSemantic() layout = tle.gpu.swizzled_shared_layout.make_default(len(shape)) @@ -308,6 +320,77 @@ def test_buffered_tensor_slot_rejects_non_int32_stage(self): buffer.slot(stage, _semantic=semantic) +class TestTmaCopyBarrierFrontend: + """Test TMA copy explicit completion barrier validation.""" + + def _make_desc_buffer_semantic(self, shape): + semantic = TestBufferedTensor._FakeSemantic() + layout = tle.gpu.swizzled_shared_layout.make_default(len(shape)) + buffer = tle.gpu.buffered_tensor("smem", tl.float16, shape, tle.gpu.smem, layout, semantic) + desc = object.__new__(tl.tensor_descriptor) + desc.handle = "desc" + desc.shape = list(shape) + return desc, buffer, semantic + + def _make_barrier(self, semantic, expect_bytes, num_barriers=1, shape=None): + layout = tle.gpu.swizzled_shared_layout.make_default(2) + return tle.gpu.barrier( + "bar", + num_barriers, + 1, + tle.gpu.PENDING, + expect_bytes, + layout, + semantic, + shape=list(shape if shape is not None else [num_barriers, 1]), + allocation_key="bar", + ) + + def test_copy_accepts_explicit_tma_completion_barrier(self): + desc, buffer, semantic = self._make_desc_buffer_semantic([16, 16]) + barrier = self._make_barrier(semantic, 512, shape=[1, 1]) + + tle.gpu.copy(desc, buffer, (16, 16), (0, 0), barrier=barrier, _semantic=semantic) + + assert semantic.builder.tma_copy_args == ("desc", "smem", ["stage_0", "stage_0"], "slot_handle", 512) + assert semantic._tle_barrier_backend_uses == {("bar", 0): "mbarrier"} + + def test_copy_barrier_requires_barrier_value(self): + desc, buffer, semantic = self._make_desc_buffer_semantic([16, 16]) + + with pytest.raises(ValueError, match="expects tle.gpu barrier"): + tle.gpu.copy(desc, buffer, (16, 16), (0, 0), barrier="bar", _semantic=semantic) + + def test_copy_barrier_requires_expect_bytes(self): + desc, buffer, semantic = self._make_desc_buffer_semantic([16, 16]) + barrier = self._make_barrier(semantic, None) + + with pytest.raises(ValueError, match="expect_bytes"): + tle.gpu.copy(desc, buffer, (16, 16), (0, 0), barrier=barrier, _semantic=semantic) + + def test_copy_barrier_requires_indexed_barrier_array(self): + desc, buffer, semantic = self._make_desc_buffer_semantic([16, 16]) + barrier = self._make_barrier(semantic, 512, num_barriers=2, shape=[2, 1]) + + with pytest.raises(ValueError, match="arrays must be indexed"): + tle.gpu.copy(desc, buffer, (16, 16), (0, 0), barrier=barrier, _semantic=semantic) + + def test_copy_barrier_rejects_tma_store(self): + desc, buffer, semantic = self._make_desc_buffer_semantic([16, 16]) + barrier = self._make_barrier(semantic, 512) + + with pytest.raises(ValueError, match="global-to-shared"): + tle.gpu.copy(buffer, desc, (16, 16), (0, 0), barrier=barrier, _semantic=semantic) + + def test_copy_barrier_rejects_mixed_backend(self): + desc, buffer, semantic = self._make_desc_buffer_semantic([16, 16]) + barrier = self._make_barrier(semantic, 512) + semantic._tle_barrier_backend_uses = {("bar", 0): "named"} + + with pytest.raises(ValueError, match="cannot mix"): + tle.gpu.copy(desc, buffer, (16, 16), (0, 0), barrier=barrier, _semantic=semantic) + + class TestPipeFrontend: """Test strict front-end validation for tle.pipe.""" @@ -515,6 +598,14 @@ def test_tle_module_import(self): assert hasattr(tle.gpu, 'pipeline') assert hasattr(tle.gpu, 'storage_kind') assert hasattr(tle.gpu, 'buffered_tensor') + assert hasattr(tle.gpu, 'alloc_barriers') + assert hasattr(tle.gpu, 'alloc_barrier') + assert hasattr(tle.gpu, 'barrier_wait') + assert hasattr(tle.gpu, 'barrier_arrive') + assert hasattr(tle.gpu, 'PENDING') + assert hasattr(tle.gpu, 'READY') + assert not hasattr(tle.gpu, 'barrier_expect_bytes') + assert "barrier" in inspect.signature(tle.gpu.copy).parameters def test_tle_functions_have_docstrings(self): """Test TLE functions have docstrings""" diff --git a/python/test/tle/unit/test_tle_wgmma_pipeline_routing.py b/python/test/tle/unit/test_tle_wgmma_pipeline_routing.py new file mode 100644 index 0000000000..c7b2c07f98 --- /dev/null +++ b/python/test/tle/unit/test_tle_wgmma_pipeline_routing.py @@ -0,0 +1,153 @@ +import pytest +import triton +import triton.language as tl +import triton.experimental.tle.language as tle +from triton._C.libtriton import ir +from triton.backends.compiler import GPUTarget +from triton.compiler.compiler import ASTSource, make_backend + +_WGMMA_PIPELINE_MODE_ATTR = "tle.wgmma_pipeline_mode" +_USER_PROMISE_MODE = "user_promise" +_USER_PROMISE_MODE_ATTRS = ( + f'"{_WGMMA_PIPELINE_MODE_ATTR}" = "{_USER_PROMISE_MODE}"', + f'{_WGMMA_PIPELINE_MODE_ATTR} = "{_USER_PROMISE_MODE}"', +) +_HOPPER_TARGET = GPUTarget("cuda", 90, 32) + + +def _require_cuda(): + try: + import torch + + target = triton.runtime.driver.active.get_current_target() + if target.backend != "cuda": + pytest.skip(f"CUDA Hopper backend is required, got {target.backend}") + if int(target.arch) < 90: + pytest.skip(f"CUDA Hopper backend is required, got sm{target.arch}") + torch.cuda.init() + except Exception as exc: + pytest.skip(f"CUDA init failed: {exc}") + + +@pytest.fixture(scope="module", autouse=True) +def _cuda_guard(): + _require_cuda() + + +@triton.jit +def _no_trigger_kernel(out): + tl.store(out, 0) + + +@triton.jit +def _alloc_barriers_trigger_kernel(): + tle.gpu.alloc_barriers(2) + + +@triton.jit +def _alloc_barrier_trigger_kernel(): + tle.gpu.alloc_barrier() + + +@triton.jit +def _wgmma_wait_trigger_kernel(): + acc = tl.zeros((16, 16), tl.float32) + tle.gpu.wgmma_wait(0, acc) + + +@triton.jit +def _ws_call_default(out): + tl.store(out, 1) + + +@triton.jit +def _ws_call_worker(out): + tl.store(out, 2) + + +@triton.jit +def _warp_specialize_call_lower_kernel(out): + tle.gpu.warp_specialize( + [ + (_ws_call_default, (out, )), + (_ws_call_worker, (out, )), + ], + [4], + [168], + ) + + +@triton.jit +def _ws_barrier_default(bar): + tle.gpu.barrier_arrive(bar, phaseIdx=0) + + +@triton.jit +def _ws_barrier_worker(bar, out): + tle.gpu.barrier_wait(bar, phaseIdx=0) + tl.store(out, 1) + + +@triton.jit +def _warp_specialize_barrier_inline_kernel(out): + bar = tle.gpu.alloc_barrier() + tle.gpu.warp_specialize( + [ + (_ws_barrier_default, (bar, )), + (_ws_barrier_worker, (bar, out)), + ], + [4], + [168], + ) + + +def _make_ttir(kernel, signature=None): + backend = make_backend(_HOPPER_TARGET) + options = backend.parse_options({"num_warps": 4}) + context = ir.context() + ir.load_dialects(context) + backend.load_dialects(context) + src = ASTSource(fn=kernel, signature=signature or {}, constexprs={}) + module = src.make_ir( + _HOPPER_TARGET, + options, + backend.get_codegen_implementation(options), + backend.get_module_map(), + context, + ) + return str(module) + + +def _has_user_promise_mode_attr(ttir): + return any(attr in ttir for attr in _USER_PROMISE_MODE_ATTRS) + + +@pytest.mark.parametrize( + ("kernel", "signature", "routes_user_promise"), + [ + (_no_trigger_kernel, {"out": "*i32"}, False), + (_alloc_barriers_trigger_kernel, {}, True), + (_alloc_barrier_trigger_kernel, {}, True), + (_wgmma_wait_trigger_kernel, {}, True), + ], +) +def test_tle_wgmma_pipeline_route_marker_exact_api_list(kernel, signature, routes_user_promise): + ttir = _make_ttir(kernel, signature) + + assert _has_user_promise_mode_attr(ttir) is routes_user_promise + + +def test_tle_warp_specialize_keeps_call_lowering_without_user_promise_marker(): + ttir = _make_ttir(_warp_specialize_call_lower_kernel, {"out": "*i32"}) + + assert "ttg.warp_specialize" in ttir + assert "tt.call" in ttir + assert not _has_user_promise_mode_attr(ttir) + + +def test_tle_warp_specialize_inlines_with_user_promise_marker(): + ttir = _make_ttir(_warp_specialize_barrier_inline_kernel, {"out": "*i32"}) + + assert "ttg.warp_specialize" in ttir + assert _has_user_promise_mode_attr(ttir) + assert "tt.call" not in ttir diff --git a/python/triton/compiler/code_generator.py b/python/triton/compiler/code_generator.py index 250e14f54e..e79845f51f 100644 --- a/python/triton/compiler/code_generator.py +++ b/python/triton/compiler/code_generator.py @@ -397,6 +397,7 @@ def __init__(self, context, prototype, gscope, function_name, jit_fn: JITFunctio self.local_defs: Dict[str, tensor] = {} self.dereference_name: Callable[[str], Any] = self._define_name_lookup() self.fn = None + self.used_vars = set() # Are we currently visiting an ast.arg's default value? These have some # special handling. self.visiting_arg_default_value = False @@ -808,6 +809,7 @@ def visit_Lambda(self, node: ast.Lambda): def visit_Name(self, node): if type(node.ctx) is ast.Store: return node.id + self.used_vars.add(node.id) return self.dereference_name(node.id) def visit_Store(self, node): @@ -1435,6 +1437,70 @@ def call_JitFunction(self, fn: JITFunction, args, kwargs, caller_context=None): handles = [call_op.get_result(i) for i in range(call_op.get_num_results())] return next(unflatten_ir_values(handles, [callee_ret_type])) + def inline_JitFunction(self, fn: JITFunction, args, kwargs, caller_context=None): + """Inline a JITFunction body into the current insertion block. + + This is intentionally narrower than a general inliner: it is used by + TLE warp-specialize regions so partition-local lowering can see the + body directly instead of a helper ``tt.call`` boundary. + """ + bound_args = inspect.getcallargs(fn.fn, *args, **kwargs) + ordered_args = [bound_args[name] for name in fn.arg_names] + for i, arg in enumerate(ordered_args): + if isinstance(arg, (language.dtype, float, int, bool, JITFunction)): + ordered_args[i] = language.core.constexpr(arg) + + parsed = fn.parse() + if isinstance(parsed, ast.Module): + if len(parsed.body) != 1 or not isinstance(parsed.body[0], ast.FunctionDef): + raise ValueError("inline_JitFunction expects a single function definition") + fn_def = parsed.body[0] + else: + fn_def = parsed + if not isinstance(fn_def, ast.FunctionDef): + raise ValueError("inline_JitFunction expects a function definition") + + mapped_gscope = {} + for k, v in fn.get_capture_scope().items(): + if isinstance(v, ModuleType): + mapped_gscope[k] = self.builder.module_map.get(v.__name__, v) + continue + module_name = getattr(v, "__module__", "") + if module_name in self.builder.module_map: + mapped_gscope[k] = getattr(self.builder.module_map[module_name], v.__name__) + else: + mapped_gscope[k] = v + + prev_gscope = self.gscope + prev_lscope = self.lscope + prev_defs = self.local_defs + prev_caller_context = self.caller_context + try: + self.gscope = mapped_gscope + self.lscope = {} + self.local_defs = {} + self.caller_context = caller_context or self.caller_context + for arg_name, arg_value in zip(fn.arg_names, ordered_args): + self.set_value(arg_name, arg_value) + + def decay_return(value): + if isinstance(value, language.tuple): + return _apply_to_tuple_values(value, decay_return) + if isinstance(value, (language.constexpr, int, float)): + return self.semantic.to_tensor(value) + return value + + for stmt in fn_def.body: + if isinstance(stmt, ast.Return): + return decay_return(self.visit(stmt.value)) if stmt.value is not None else None + self.visit(stmt) + return None + finally: + self.gscope = prev_gscope + self.lscope = prev_lscope + self.local_defs = prev_defs + self.caller_context = prev_caller_context + def call_Function(self, node, fn, args, kws): # 4. Get current line number and hints flagtree_hints = hint_trigger("get_node_hints", self, node) diff --git a/python/triton/experimental/tle/language/gpu/__init__.py b/python/triton/experimental/tle/language/gpu/__init__.py index 939adb04f2..e1ecb4a6ea 100644 --- a/python/triton/experimental/tle/language/gpu/__init__.py +++ b/python/triton/experimental/tle/language/gpu/__init__.py @@ -2,13 +2,19 @@ from .core import ( pipeline, alloc, + alloc_barrier, + alloc_barriers, + barrier_arrive, + barrier_wait, copy, memory_space, local_ptr, warp_specialize, + wgmma, + wgmma_wait, ) from .types import (layout, shared_layout, swizzled_shared_layout, tensor_memory_layout, nv_mma_shared_layout, scope, - buffered_tensor, buffered_tensor_type, smem, tmem) + buffered_tensor, buffered_tensor_type, barrier, barrier_type, smem, tmem, PENDING, READY) # Backward-compat alias expected by existing tests/tutorials. storage_kind = memory_space @@ -16,9 +22,15 @@ __all__ = [ "pipeline", "alloc", + "alloc_barrier", + "alloc_barriers", + "barrier_arrive", + "barrier_wait", "copy", "local_ptr", "warp_specialize", + "wgmma", + "wgmma_wait", "storage_kind", "layout", "memory_space", @@ -29,6 +41,10 @@ "scope", "buffered_tensor", "buffered_tensor_type", + "barrier", + "barrier_type", + "PENDING", + "READY", "smem", "tmem", ] diff --git a/python/triton/experimental/tle/language/gpu/core.py b/python/triton/experimental/tle/language/gpu/core.py index 127bc305a0..486297586f 100644 --- a/python/triton/experimental/tle/language/gpu/core.py +++ b/python/triton/experimental/tle/language/gpu/core.py @@ -17,6 +17,25 @@ # Address space 3 matches the shared-memory space used in TritonGPU lowering. SHARED_MEMORY_ADDRESS_SPACE = 3 +_WGMMA_PIPELINE_MODE_ATTR = "tle.wgmma_pipeline_mode" +_WGMMA_PIPELINE_MODE_USER_PROMISE = "user_promise" + + +def _mark_wgmma_user_promise(_semantic, _generator): + if _generator is None or _semantic is None: + return + _generator.module.set_attr( + _WGMMA_PIPELINE_MODE_ATTR, + _semantic.builder.get_string_attr(_WGMMA_PIPELINE_MODE_USER_PROMISE), + ) + + +def _is_wgmma_user_promise_marked(_generator): + if _generator is None: + return False + mode = _generator.module.get_operation().get_str_attr(_WGMMA_PIPELINE_MODE_ATTR) + return mode == _WGMMA_PIPELINE_MODE_USER_PROMISE + class pipeline(range): """ @@ -62,6 +81,19 @@ def normalize(arg): return tuple(normalize(arg) for arg in args) +def _as_warp_specialize_entry(entry, index: int): + entry = tl._unwrap_if_constexpr(entry) + if isinstance(entry, tl.tuple): + entry = tuple(entry.values) + if not isinstance(entry, tuple): + raise ValueError(f"warp_specialize entry {index} must be a tuple, got {type(entry).__name__}") + if len(entry) != 2: + raise ValueError(f"warp_specialize entry {index} must be (fn, args); pass cid as a normal constexpr arg") + if not isinstance(tl._unwrap_if_constexpr(entry[1]), (tuple, tl.tuple)): + raise ValueError(f"warp_specialize entry {index} args must be a tuple") + return entry + + def _as_result_values(results): if results is None: return tuple() @@ -125,19 +157,25 @@ def warp_specialize(functions_and_args, worker_num_warps, worker_num_regs, _sema builder = _semantic.builder insert_pt = builder.get_insertion_point() + inline_user_promise = _is_wgmma_user_promise_marked(_generator) + if inline_user_promise: + call_jit_function = _generator.inline_JitFunction + else: + call_jit_function = _generator.call_JitFunction - default_fn, default_args = functions_and_args[0] + default_fn, default_args = _as_warp_specialize_entry(functions_and_args[0], 0) default_args = _as_call_args(default_args) - default_block = builder.new_block() + default_block = builder.create_block() builder.set_insertion_point_to_start(default_block) - default_results = _generator.call_JitFunction(default_fn, default_args, kwargs={}) + default_results = call_jit_function(default_fn, default_args, kwargs={}) default_result_values = _as_result_values(default_results) default_result_handles = flatten_values_to_ir(default_result_values) - builder.create_warp_yield(default_result_handles) + builder.set_insertion_point_to_end(default_block) result_types = [result.get_type() for result in default_result_handles] worker_items = [] - for worker_fn, worker_args in functions_and_args[1:]: + for idx, entry in enumerate(functions_and_args[1:], start=1): + worker_fn, worker_args = _as_warp_specialize_entry(entry, idx) worker_args = _as_call_args(worker_args) flattened = flatten_values_to_ir(worker_args) worker_items.append((worker_fn, worker_args, flattened)) @@ -145,7 +183,11 @@ def warp_specialize(functions_and_args, worker_num_warps, worker_num_regs, _sema builder.restore_insertion_point(insert_pt) ws_op = builder.create_warp_specialize(result_types, worker_arg_handles, worker_num_warps) - ws_op.get_default_region().push_back(default_block) + real_default_block = builder.create_block_with_parent(ws_op.get_default_region(), []) + default_block.merge_block_before(real_default_block) + default_block = real_default_block + builder.set_insertion_point_to_end(default_block) + builder.create_warp_yield(default_result_handles) ws_op.set_requested_registers(worker_num_regs) builder.create_block_with_parent(ws_op.get_partition_op_holder(), []) @@ -156,7 +198,8 @@ def warp_specialize(functions_and_args, worker_num_warps, worker_num_regs, _sema block_args = [block.get_argument(remapped[j]) for j in builtins.range(len(flattened))] block_values = tuple(unflatten_ir_values(block_args, [arg.type for arg in worker_args])) caller_context = WarpSpecializeCallerContext(worker_num_warps[idx]) - _generator.call_JitFunction(worker_fn, block_values, kwargs={}, caller_context=caller_context) + call_jit_function(worker_fn, block_values, kwargs={}, caller_context=caller_context) + builder.set_insertion_point_to_end(block) builder.create_warp_return() builder.set_insertion_point_after(ws_op.get_operation()) @@ -306,6 +349,471 @@ def alloc( raise RuntimeError(f"Memory allocation failed: {str(e)}") from e +def _unwrap_barrier_constexpr(value): + value = tl._unwrap_if_constexpr(value) + if isinstance(value, tl.constexpr): + return value.value + return value + + +def _require_barrier_int(value, name: str) -> int: + value = _unwrap_barrier_constexpr(value) + if not isinstance(value, int): + raise ValueError(f"{name} must be a compile-time integer") + return value + + +def _normalize_barrier_init(init) -> str: + init = _unwrap_barrier_constexpr(init) + if isinstance(init, str): + normalized = init.lower() + if normalized in (tle.PENDING, "pending"): + return tle.PENDING + if normalized in (tle.READY, "ready"): + return tle.READY + raise ValueError("barrier init must be tle.gpu.PENDING or tle.gpu.READY") + + +# NVIDIA hardware named barriers are ids 0..15. TLE frontend assigns virtual +# ids from 16 upward so each source barrier is distinguishable until the late +# compiler pass remaps them to conflict-free physical ids. +_FIRST_VIRTUAL_NAMED_BARRIER_ID = 16 + + +def _reserve_named_barrier_ids(_semantic, count: int) -> int: + next_id = getattr(_semantic, "_tle_next_named_barrier_id", _FIRST_VIRTUAL_NAMED_BARRIER_ID) + last_id = next_id + count - 1 + setattr(_semantic, "_tle_next_named_barrier_id", last_id + 1) + return next_id + + +def _barrier_handle_key(handle): + try: + hash(handle) + return handle + except TypeError: + return id(handle) + + +def _ensure_named_barrier_ids(slot: tle.barrier, _semantic) -> None: + if slot.named_base_id > 0: + return + key = slot.allocation_key + if key is None: + key = _barrier_handle_key(slot.handle) + named_bases = getattr(_semantic, "_tle_barrier_named_bases", None) + if named_bases is None: + named_bases = {} + setattr(_semantic, "_tle_barrier_named_bases", named_bases) + base_id = named_bases.get(key) + if base_id is None: + base_id = _reserve_named_barrier_ids(_semantic, slot.num_barriers) + named_bases[key] = base_id + slot.named_base_id = base_id + slot.type.named_base_id = base_id + + +def _barrier_phase_tensor(phaseIdx, init: str, _semantic) -> tl.tensor: + init_polarity = 1 if init == tle.READY else 0 + raw_phase = _unwrap_barrier_constexpr(phaseIdx) + if isinstance(raw_phase, int): + return _semantic.to_tensor((raw_phase & 1) ^ init_polarity).to(tl.int32, _semantic=_semantic) + + phase = _semantic.to_tensor(phaseIdx) + if getattr(phase.type, "is_block", lambda: False)(): + raise ValueError("barrier phaseIdx must be a scalar integer") + if not getattr(phase.type, "is_int", lambda: False)(): + raise ValueError(f"barrier phaseIdx must be integer, got {phase.type}") + if phase.dtype != tl.int32: + phase = phase.to(tl.int32, _semantic=_semantic) + phase = phase.__and__(1, _semantic=_semantic) + if init_polarity: + phase = phase.__xor__(1, _semantic=_semantic) + return phase + + +def _barrier_slot(value: tle.barrier, _semantic) -> tle.barrier: + if not isinstance(value, tle.barrier): + raise ValueError(f"barrier operation expects tle.gpu barrier, got {type(value).__name__}") + if value.is_slot: + return value + return value.__getitem__(0, _semantic=_semantic) + + +def _record_barrier_backend(slot: tle.barrier, backend: str, _semantic) -> None: + if slot.allocation_key is not None and slot.static_index is not None: + key = (slot.allocation_key, slot.static_index) + elif slot.named_base_id > 0 and slot.static_index is not None: + key = (slot.named_base_id, slot.static_index) + else: + key = _barrier_handle_key(slot.handle) + uses = getattr(_semantic, "_tle_barrier_backend_uses", None) + if uses is None: + uses = {} + setattr(_semantic, "_tle_barrier_backend_uses", uses) + previous = uses.get(key) + if previous is not None and previous != backend: + raise ValueError("cannot mix named and mbarrier backends for the same barrier slot") + uses[key] = backend + + +def _tma_completion_barrier_slot(value, _semantic) -> tle.barrier: + if not isinstance(value, tle.barrier): + raise ValueError(f"TMA copy barrier expects tle.gpu barrier, got {type(value).__name__}") + if not value.is_slot and value.num_barriers != 1: + raise ValueError("TMA copy barrier arrays must be indexed, e.g. bars[i]") + + slot = _barrier_slot(value, _semantic) + if slot.expect_bytes is None: + raise ValueError("TMA copy barrier must be allocated with expect_bytes") + if not isinstance(slot.expect_bytes, int) or slot.expect_bytes <= 0: + raise ValueError("TMA copy barrier expect_bytes must be a positive compile-time integer") + + _record_barrier_backend(slot, "mbarrier", _semantic) + return slot + + +@tl.builtin +def alloc_barriers( + num_barriers, + arrive_count=1, + init=tle.PENDING, + expect_bytes=None, + _semantic=None, + _generator=None, +) -> tle.barrier: + """Allocate a TLE GPU barrier array.""" + _mark_wgmma_user_promise(_semantic, _generator) + + num_barriers = _require_barrier_int(num_barriers, "num_barriers") + arrive_count = _require_barrier_int(arrive_count, "arrive_count") + init = _normalize_barrier_init(init) + if num_barriers <= 0: + raise ValueError("num_barriers must be positive") + if arrive_count <= 0: + raise ValueError("arrive_count must be positive") + + expect_bytes = _unwrap_barrier_constexpr(expect_bytes) + if expect_bytes is not None: + if not isinstance(expect_bytes, int): + raise ValueError("expect_bytes must be a compile-time integer or None") + if expect_bytes <= 0: + raise ValueError("expect_bytes must be positive when provided") + + layout = tle.swizzled_shared_layout.make_default(rank=2) + named_base_id = 0 + barrier_ty = tle.barrier_type(num_barriers, arrive_count, init, expect_bytes, layout, _semantic, + shape=[num_barriers, 1], named_base_id=named_base_id) + handle = _semantic.builder.create_barrier_alloc( + barrier_ty.to_ir(_semantic.builder), + num_barriers, + arrive_count, + 1 if init == tle.READY else 0, + -1 if expect_bytes is None else expect_bytes, + ) + allocation_key = _barrier_handle_key(handle) + barrier_ty.allocation_key = allocation_key + return tle.barrier(handle, num_barriers, arrive_count, init, expect_bytes, layout, _semantic, + shape=[num_barriers, 1], named_base_id=named_base_id, allocation_key=allocation_key) + + +@tl.builtin +def alloc_barrier( + arrive_count=1, + init=tle.PENDING, + expect_bytes=None, + _semantic=None, + _generator=None, +) -> tle.barrier: + """Allocate a single TLE GPU barrier.""" + _mark_wgmma_user_promise(_semantic, _generator) + return alloc_barriers( + 1, + arrive_count, + init, + expect_bytes, + _semantic=_semantic, + ) + + +@tl.builtin +def barrier_wait(barr, phaseIdx=None, _semantic=None) -> None: + """Wait on a TLE GPU barrier slot.""" + slot = _barrier_slot(barr, _semantic) + if phaseIdx is None: + if slot.expect_bytes is not None: + raise ValueError("barrier_wait on a barrier with expect_bytes requires phaseIdx") + if slot.init == tle.READY: + raise ValueError("barrier_wait without phaseIdx selects named barrier, which does not support READY") + if slot.static_index is None: + raise ValueError("named barrier backend requires a static barrier slot index") + _ensure_named_barrier_ids(slot, _semantic) + _record_barrier_backend(slot, "named", _semantic) + _semantic.builder.create_barrier_wait_named(slot.handle, slot.named_base_id + slot.static_index, + slot.arrive_count) + return + + phase = _barrier_phase_tensor(phaseIdx, slot.init, _semantic) + _record_barrier_backend(slot, "mbarrier", _semantic) + _semantic.builder.create_barrier_wait_mbarrier(slot.handle, phase.handle) + + +@tl.builtin +def barrier_arrive(barr, arrive_count=1, phaseIdx=None, _semantic=None) -> None: + """Arrive on a TLE GPU barrier slot.""" + slot = _barrier_slot(barr, _semantic) + arrive_count = _require_barrier_int(arrive_count, "arrive_count") + if arrive_count <= 0: + raise ValueError("arrive_count must be positive") + + if phaseIdx is None: + if slot.expect_bytes is not None: + raise ValueError("barrier_arrive on a barrier with expect_bytes requires phaseIdx") + if slot.init == tle.READY: + raise ValueError("barrier_arrive without phaseIdx selects named barrier, which does not support READY") + if arrive_count != 1: + raise ValueError("named barrier backend requires barrier_arrive arrive_count = 1") + if slot.static_index is None: + raise ValueError("named barrier backend requires a static barrier slot index") + _ensure_named_barrier_ids(slot, _semantic) + _record_barrier_backend(slot, "named", _semantic) + _semantic.builder.create_barrier_arrive_named(slot.handle, slot.named_base_id + slot.static_index, + slot.arrive_count) + return + + phase = _barrier_phase_tensor(phaseIdx, slot.init, _semantic) + _record_barrier_backend(slot, "mbarrier", _semantic) + _semantic.builder.create_barrier_arrive_mbarrier(slot.handle, arrive_count, phase.handle) + + +def _require_wgmma_int(value, name: str) -> int: + value = tl._unwrap_if_constexpr(value) + if isinstance(value, tl.constexpr): + value = value.value + if not isinstance(value, int): + raise ValueError(f"{name} must be a compile-time integer") + return value + + +def _require_wgmma_smem_operand(value, name: str) -> tle.buffered_tensor: + if not isinstance(value, tle.buffered_tensor): + raise ValueError(f"{name} must be a tle.gpu buffered_tensor in shared memory") + if value.type.storage is not tle.smem: + raise ValueError(f"{name} must live in shared memory") + if not isinstance(value.type.layout, tle.nv_mma_shared_layout): + raise ValueError(f"{name} must use nv_mma_shared_layout; allocate it with the default tle.gpu.alloc layout") + if len(value.type.shape) != 2: + raise ValueError(f"{name} must be a rank-2 shared tile") + return value + + +def _require_wgmma_bool(value, name: str) -> bool: + value = tl._unwrap_if_constexpr(value) + if isinstance(value, tl.constexpr): + value = value.value + if not isinstance(value, bool): + raise ValueError(f"{name} must be a compile-time bool") + return value + + +def _require_rank2_wgmma_operand(value, name: str): + if len(value.type.shape) != 2: + raise ValueError(f"{name} transpose currently supports only rank-2 operands") + + +def _require_transpose_order(order, rank: int, name: str): + if len(order) != rank or sorted(order) != list(builtins.range(rank)): + raise ValueError(f"{name} transpose order must be a permutation of rank {rank}") + + +def _transpose_wgmma_smem_operand(value: tle.buffered_tensor, name: str, _semantic) -> tle.buffered_tensor: + _require_rank2_wgmma_operand(value, name) + order = [1, 0] + _require_transpose_order(order, len(value.type.shape), name) + handle = _semantic.builder.create_memdesc_trans(value.handle, order) + shape = [value.type.shape[i] for i in order] + + alloc_shape = value.type.alloc_shape + leading_rank = len(alloc_shape) - len(value.type.shape) + alloc_tail = alloc_shape[leading_rank:] + transposed_alloc_shape = alloc_shape[:leading_rank] + [alloc_tail[i] for i in order] + + layout = value.type.layout.make_permute(order) + return tle.buffered_tensor( + handle, + value.dtype, + shape, + value.type.storage, + layout, + _semantic, + alloc_shape=transposed_alloc_shape, + ) + + +_WGMMA_ALLOWED_OPERAND_TYPE_PAIRS = ( + (tle.buffered_tensor, tle.buffered_tensor), + (tl.tensor, tle.buffered_tensor), +) + + +def _canonicalize_wgmma_operands(a, b, trans_a: bool, trans_b: bool, _semantic): + a = tl._unwrap_if_constexpr(a) + b = tl._unwrap_if_constexpr(b) + + if not any(isinstance(a, a_ty) and isinstance(b, b_ty) for a_ty, b_ty in _WGMMA_ALLOWED_OPERAND_TYPE_PAIRS): + if isinstance(b, tl.tensor): + raise ValueError( + "wgmma b currently supports only shared-memory buffered_tensor operands; tensor B is unsupported") + raise ValueError("wgmma operands must be one of: " + "(shared-memory buffered_tensor, shared-memory buffered_tensor) or " + "(tl.tensor, shared-memory buffered_tensor)") + + if isinstance(a, tle.buffered_tensor): + a = _require_wgmma_smem_operand(a, "wgmma a") + if trans_a: + a = _transpose_wgmma_smem_operand(a, "wgmma a", _semantic) + elif trans_a: + _require_rank2_wgmma_operand(a, "wgmma a") + a = tl.trans(a, _semantic=_semantic) + + b = _require_wgmma_smem_operand(b, "wgmma b") + if trans_b: + b = _transpose_wgmma_smem_operand(b, "wgmma b", _semantic) + return a, b + + +def _wgmma_ret_scalar_ty(lhs_dtype, out_dtype): + if lhs_dtype.is_int(): + if lhs_dtype != tl.int8: + raise ValueError("wgmma integer operands currently support only tl.int8") + return tl.int32 + if out_dtype.is_bf16(): + raise ValueError("wgmma out_dtype=bfloat16 is unsupported; use float32/float16 and cast afterward") + if lhs_dtype.is_fp32() or lhs_dtype.is_bf16(): + return tl.float32 + return out_dtype + + +def _wgmma_zero_value(builder, scalar_ty): + if scalar_ty.is_int(): + return builder.get_int32(0) + if scalar_ty.is_fp64(): + return builder.get_fp64(0) + if scalar_ty.is_fp16(): + return builder.get_fp16(0) + return builder.get_fp32(0) + + +@tl.builtin +def wgmma( + a, + b, + acc=None, + input_precision=None, + max_num_imprecise_acc=None, + out_dtype=tl.float32, + trans_a: tl.constexpr = False, + trans_b: tl.constexpr = False, + _semantic=None, +) -> tl.tensor: + """ + Issue an asynchronous Hopper WGMMA. + + A may be a shared-memory TLE buffer or a register tensor. B must currently + be a shared-memory TLE buffer. ``trans_a`` and ``trans_b`` are rank-2 + descriptor/tensor transpose requests; shared-memory transposes are emitted + as descriptor-only views. + + The returned accumulator is an async WGMMA dependency value. Use + ``tle.gpu.wgmma_wait(pendings, acc)`` before consuming it with ordinary + tensor operations or storing it. + """ + trans_a = _require_wgmma_bool(trans_a, "trans_a") + trans_b = _require_wgmma_bool(trans_b, "trans_b") + a, b = _canonicalize_wgmma_operands(a, b, trans_a, trans_b, _semantic) + + m, k = [int(tl._unwrap_if_constexpr(dim)) for dim in a.type.shape] + k_b, n = [int(tl._unwrap_if_constexpr(dim)) for dim in b.type.shape] + if k != k_b: + raise ValueError(f"wgmma shape mismatch: a is {a.type.shape}, b is {b.type.shape}") + if m < 64 or m % 64 != 0: + raise ValueError("wgmma result M dimension must be divisible by 64") + if n < 8 or n % 8 != 0: + raise ValueError("wgmma result N dimension must be divisible by 8") + if k < 16: + raise ValueError("wgmma K dimension must be at least 16") + + if not (a.dtype.is_fp8() and b.dtype.is_fp8()): + if a.dtype != b.dtype: + raise ValueError(f"wgmma operands must have the same dtype, got {a.dtype} and {b.dtype}") + if a.dtype not in (tl.int8, tl.float16, tl.bfloat16, tl.float32): + raise ValueError(f"unsupported wgmma operand dtype {a.dtype}") + + out_dtype = tl._unwrap_if_constexpr(out_dtype) + if not isinstance(out_dtype, tl.dtype): + raise ValueError(f"wgmma out_dtype must be a Triton dtype, got {type(out_dtype).__name__}") + + if input_precision is None: + input_precision = _semantic.builder.options.default_dot_input_precision + input_precision = _semantic._str_to_dot_input_precision(tl._unwrap_if_constexpr(input_precision)) + + max_num_imprecise_acc = tl._unwrap_if_constexpr(max_num_imprecise_acc) + if isinstance(max_num_imprecise_acc, tl.constexpr): + max_num_imprecise_acc = max_num_imprecise_acc.value + if max_num_imprecise_acc is None: + if a.dtype.is_fp8() and b.dtype.is_fp8(): + max_num_imprecise_acc = _semantic.builder.options.max_num_imprecise_acc_default + else: + max_num_imprecise_acc = 0 + else: + max_num_imprecise_acc = _require_wgmma_int(max_num_imprecise_acc, "max_num_imprecise_acc") + if max_num_imprecise_acc < 0: + raise ValueError("max_num_imprecise_acc must be non-negative") + + ret_scalar_ty = _wgmma_ret_scalar_ty(a.dtype, out_dtype) + ret_ty = tl.block_type(ret_scalar_ty, [m, n]) + builder = _semantic.builder + + acc = tl._unwrap_if_constexpr(acc) + if acc is None: + zero = _wgmma_zero_value(builder, ret_scalar_ty) + acc_handle = builder.create_splat(ret_ty.to_ir(builder), zero) + else: + if not isinstance(acc, tl.tensor): + raise ValueError(f"wgmma acc must be a tl.tensor or None, got {type(acc).__name__}") + if tuple(int(tl._unwrap_if_constexpr(dim)) for dim in acc.type.shape) != (m, n): + raise ValueError(f"wgmma acc shape must be {(m, n)}, got {acc.type.shape}") + if acc.dtype != ret_scalar_ty: + raise ValueError(f"wgmma acc dtype must be {ret_scalar_ty}, got {acc.dtype}") + acc_handle = acc.handle + + result = builder.create_tle_wgmma( + a.handle, + b.handle, + acc_handle, + input_precision, + max_num_imprecise_acc, + True, + ) + return tensor(result, ret_ty) + + +@tl.builtin +def wgmma_wait(pendings, acc=None, _semantic=None, _generator=None) -> tl.tensor: + """Wait until ``pendings`` or fewer async WGMMA groups remain outstanding.""" + _mark_wgmma_user_promise(_semantic, _generator) + if acc is None and isinstance(pendings, tl.tensor): + acc = pendings + pendings = 0 + pendings = _require_wgmma_int(pendings, "pendings") + if pendings < 0: + raise ValueError("wgmma_wait pendings must be non-negative") + if not isinstance(acc, tl.tensor): + raise ValueError(f"wgmma_wait acc must be a tl.tensor, got {type(acc).__name__}") + result = _semantic.builder.create_tle_wgmma_wait(acc.handle, pendings) + return tensor(result, acc.type) + + class CopyDirection(Enum): """Copy direction enum for data transfer operations""" GM_TO_LOCAL = "GMTOLOCAL" # Global memory to local memory @@ -318,6 +826,7 @@ def copy( dst, shape, offsets: Sequence[constexpr | tensor] = None, + barrier=None, _semantic=None, ) -> None: """ @@ -348,6 +857,8 @@ def copy( shape: Tuple specifying the dimensions of the data to copy offsets: Sequence of offsets for multi-dimensional addressing. Used with TMA operations to specify the starting coordinates within the tensor. Required for TMA copy. + barrier: Optional TLE GPU mbarrier completion barrier for global-to-shared TMA copy. + The barrier must come from ``tle.gpu.alloc_barrier(s)(expect_bytes=...)``. _semantic: Internal semantic analyzer for validation and compilation (user-provided) Raises: @@ -361,6 +872,11 @@ def copy( TMA copy with offsets: tle.copy(tma_desc, local_buf, [64, 64], [x_offset, y_offset]) + + TMA copy with explicit completion barrier: + bar = tle.gpu.alloc_barrier(expect_bytes=64 * 64 * 2) + tle.copy(tma_desc, local_buf, [64, 64], [x_offset, y_offset], barrier=bar) + tle.gpu.barrier_wait(bar, phaseIdx=0) """ mthreads_enabled = mthreads_copy.enabled() iluvatar_enabled = iluvatar_copy.enabled() @@ -415,6 +931,7 @@ def tmacopy( direction, shape: tuple, offsets: Sequence[constexpr | tensor], + barrier=None, _semantic=None, ) -> None: # Parameter validation @@ -454,11 +971,20 @@ def tmacopy( else: raise ValueError(f"Shape parameter must be tuple or list, but got {type(shape)}") + barrier_slot = None + expect_bytes = -1 + if barrier is not None: + if direction != CopyDirection.GM_TO_LOCAL: + raise ValueError("TMA copy barrier is only supported for global-to-shared TMA copy") + barrier_slot = _tma_completion_barrier_slot(barrier, _semantic) + expect_bytes = barrier_slot.expect_bytes + # 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) + _semantic.builder.create_tma_copy(src.handle, dst.handle, offsets, + None if barrier_slot is None else barrier_slot.handle, expect_bytes) return # Parameter validation @@ -499,11 +1025,15 @@ def tmacopy( else: raise ValueError(f"Shape parameter must be tuple or list, but got {type(shape)}") if is_normcopy: + if barrier is not None: + raise ValueError("copy barrier is only supported for TMA global-to-shared copy") return normcopy(src, dst, shape, direction, _semantic) if mthreads_enabled: + if barrier is not None: + raise ValueError("TMA copy barrier is only supported on NVIDIA backend") return mthreads_copy.tmacopy(src, dst, direction, shape, offsets, _semantic) else: - return tmacopy(src, dst, direction, shape, offsets, _semantic) + return tmacopy(src, dst, direction, shape, offsets, barrier, _semantic) def _expand_index_to_shape(index: tl.tensor, shape: Sequence[int], axis: int, _semantic) -> tl.tensor: diff --git a/python/triton/experimental/tle/language/gpu/types.py b/python/triton/experimental/tle/language/gpu/types.py index 12db66f7fb..7d095e1bc7 100644 --- a/python/triton/experimental/tle/language/gpu/types.py +++ b/python/triton/experimental/tle/language/gpu/types.py @@ -25,6 +25,9 @@ def to_ir(self, builder: ir.builder) -> None: smem = scope('share_memory') tmem = scope('tensor_memory') +PENDING = "pending" +READY = "ready" + def _storage_to_memdesc_space(storage: scope) -> str: if storage is smem: @@ -188,14 +191,18 @@ def make_default(cls, shape, elemType): """ def make_permute(self, dims): - permuted_order = tuple(self.order[d] for d in dims) + permuted_shape = [self.shape[d] for d in dims] + permuted_order = [self.order[d] for d in dims] + permuted_num_ctas_per_cga = [self.numCTAsPerCGA[d] for d in dims] + permuted_num_cta_split = [self.numCTASplit[d] for d in dims] + permuted_num_cta_order = [self.numCTAOrder[d] for d in dims] return nv_mma_shared_layout( - self.shape, + permuted_shape, permuted_order, self.elemType, - self.numCTAsPerCGA, - self.numCTASplit, - self.numCTAOrder, + permuted_num_ctas_per_cga, + permuted_num_cta_split, + permuted_num_cta_order, self.fp4Padded, self.swizzled, ) @@ -420,6 +427,148 @@ def _flatten_ir(self, handles) -> None: handles.append(self.handle) +class barrier(tl.base_value): + """A TLE GPU barrier array or a single indexed barrier slot.""" + + def __init__( + self, + handle, + num_barriers: int, + arrive_count: int, + init: str, + expect_bytes: Optional[int], + layout: shared_layout, + semantic: TritonSemantic, + *, + shape: Optional[List[int]] = None, + static_index: Optional[int] = None, + named_base_id: int = 0, + allocation_key=None, + ): + super().__init__() + self.handle = handle + self.num_barriers = num_barriers + self.arrive_count = arrive_count + self.init = init + self.expect_bytes = expect_bytes + self.layout = layout + self.static_index = static_index + self.named_base_id = named_base_id + self.allocation_key = allocation_key + self.shape = list(shape if shape is not None else [num_barriers, 1]) + self.type = barrier_type(num_barriers, arrive_count, init, expect_bytes, layout, semantic, shape=self.shape, + static_index=static_index, named_base_id=named_base_id, allocation_key=allocation_key) + + @property + def is_slot(self) -> bool: + return self.shape == [1] + + def _flatten_ir(self, handles) -> None: + handles.append(self.handle) + + @tl.builtin + def __getitem__(self, index, _semantic=None): + if self.is_slot: + raise ValueError("tle.gpu barrier slot cannot be indexed again") + + raw_index = tl._unwrap_if_constexpr(index) + static_index = raw_index if isinstance(raw_index, int) else None + if static_index is not None and (static_index < 0 or static_index >= self.num_barriers): + raise ValueError(f"barrier index {static_index} out of bounds for {self.num_barriers} barriers") + + index_tensor = _semantic.to_tensor(index) + if getattr(index_tensor.type, "is_block", lambda: False)(): + raise ValueError("barrier index must be a scalar integer") + if not getattr(index_tensor.type, "is_int", lambda: False)(): + raise ValueError(f"barrier index must be integer, got {index_tensor.type}") + if index_tensor.dtype != tl.int32: + index_tensor = index_tensor.to(tl.int32, _semantic=_semantic) + + slot_layout = _make_slot_layout(self.layout, [1]) + slot_ty = barrier_type(self.num_barriers, self.arrive_count, self.init, self.expect_bytes, slot_layout, + _semantic, shape=[1], static_index=static_index, named_base_id=self.named_base_id, + allocation_key=self.allocation_key) + slot_handle = _semantic.builder.create_memdesc_index(slot_ty.to_ir(_semantic.builder), self.handle, + index_tensor.handle) + return barrier(slot_handle, self.num_barriers, self.arrive_count, self.init, self.expect_bytes, slot_layout, + _semantic, shape=[1], static_index=static_index, named_base_id=self.named_base_id, + allocation_key=self.allocation_key) + + +class barrier_type(tl.block_type): + + def __init__( + self, + num_barriers: int, + arrive_count: int, + init: str, + expect_bytes: Optional[int], + layout: shared_layout, + semantic: TritonSemantic, + *, + shape: Optional[List[int]] = None, + static_index: Optional[int] = None, + named_base_id: int = 0, + allocation_key=None, + ): + self.num_barriers = num_barriers + self.arrive_count = arrive_count + self.init = init + self.expect_bytes = expect_bytes + self.layout = layout + self.semantic = semantic + self.shape = list(shape if shape is not None else [num_barriers, 1]) + self.static_index = static_index + self.named_base_id = named_base_id + self.allocation_key = allocation_key + super().__init__(tl.int64, self.shape) + + def _unflatten_ir(self, handles: List[ir.value], cursor: int) -> Tuple[barrier, int]: + value = barrier( + handles[cursor], + self.num_barriers, + self.arrive_count, + self.init, + self.expect_bytes, + self.layout, + self.semantic, + shape=self.shape, + static_index=self.static_index, + named_base_id=self.named_base_id, + allocation_key=self.allocation_key, + ) + return value, cursor + 1 + + def _flatten_ir_types(self, builder: ir.builder, out: List[ir.type]) -> None: + out.append(self.to_ir(builder)) + + def to_ir(self, builder: ir.builder) -> None: + builder = self.semantic.builder + return builder.get_memdesc_type( + self.shape, + tl.int64.to_ir(builder), + self.layout.to_ir(builder), + _storage_to_memdesc_space(smem), + self.shape, + ) + + def mangle(self) -> str: + expect = "none" if self.expect_bytes is None else str(self.expect_bytes) + index = "array" if self.static_index is None else str(self.static_index) + shape = "_".join(map(str, self.shape)) + return f"barrier_{shape}_{self.num_barriers}_{self.arrive_count}_{self.init}_{expect}_{index}_{self.named_base_id}" + + def __str__(self) -> str: + return (f"barrier<{self.shape}, arrive_count={self.arrive_count}, init={self.init}, " + f"expect_bytes={self.expect_bytes}>") + + def __eq__(self, other) -> bool: + return (type(self) is type(other) and self.shape == other.shape and self.num_barriers == other.num_barriers + and self.arrive_count == other.arrive_count and self.init == other.init + and self.expect_bytes == other.expect_bytes and self.layout == other.layout + and self.static_index == other.static_index and self.named_base_id == other.named_base_id) + + class pipe_slot_type(tl.base_type): def __init__(self, fields): diff --git a/third_party/mthreads/python/test/unit/tle/test_tle.py b/third_party/mthreads/python/test/unit/tle/test_tle.py index ea7c09cacf..e538d729a8 100644 --- a/third_party/mthreads/python/test/unit/tle/test_tle.py +++ b/third_party/mthreads/python/test/unit/tle/test_tle.py @@ -42,6 +42,7 @@ def test_tle_language_import_exports_load_signature(): "dst", "shape", "offsets", + "barrier", "_semantic", ] assert hasattr(tle.gpu, "copy") diff --git a/third_party/nvidia/backend/compiler.py b/third_party/nvidia/backend/compiler.py index db24f7c452..fe4ecde0c7 100644 --- a/third_party/nvidia/backend/compiler.py +++ b/third_party/nvidia/backend/compiler.py @@ -309,11 +309,13 @@ def make_ttgir(mod, metadata, opt, capability): tle.passes.add_optimize_local_pointer_stores(pm) # end flagtree tle passes.ttgpuir.add_accelerate_matmul(pm) + tle.passes.add_lower_wgmma(pm) passes.ttgpuir.add_remove_layout_conversions(pm) passes.ttgpuir.add_optimize_dot_operands(pm, capability >= 80) tle.passes.add_promote_local_store_staging(pm) nvidia.passes.ttnvgpuir.add_optimize_descriptor_encoding(pm) passes.ttir.add_loop_aware_cse(pm) + tle.passes.add_lower_barriers(pm) if capability // 10 in [8, 9]: passes.ttgpuir.add_fuse_nested_loops(pm) passes.common.add_canonicalizer(pm) @@ -369,6 +371,9 @@ def make_ttgir(mod, metadata, opt, capability): passes.common.add_symbol_dce(pm) nvidia.passes.ttnvgpuir.add_fence_insertion(pm, capability) nvidia.passes.ttnvgpuir.add_lower_mma(pm) + # Materialize physical named barrier ids in the user-visible TTGIR + # after warp-specialization structure is known. + tle.passes.add_allocate_named_barriers(pm) passes.common.add_sccp(pm) passes.common.add_cse(pm) passes.common.add_canonicalizer(pm) @@ -432,6 +437,8 @@ def make_llir(self, src, metadata, options, capability): # Inline TLE DSL regions before TritonGPU->LLVM lowering so no # `tle.dsl_region` op survives into the conversion pipeline. tle.raw_passes.add_tle_dsl_region_inline(pm) + # Keep this as an idempotent guard for externally provided TTGIR. + tle.passes.add_allocate_named_barriers(pm) # instrumentation point here so we can override IRs above (e.g., ttir and ttgir) if CUDABackend.instrumentation: CUDABackend.instrumentation.patch("ttgpuir_to_llvmir", pm, mod.context) diff --git a/third_party/nvidia/lib/TritonNVIDIAGPUToLLVM/BarrierOpToLLVM.cpp b/third_party/nvidia/lib/TritonNVIDIAGPUToLLVM/BarrierOpToLLVM.cpp index f4651e8fee..881af595ef 100644 --- a/third_party/nvidia/lib/TritonNVIDIAGPUToLLVM/BarrierOpToLLVM.cpp +++ b/third_party/nvidia/lib/TritonNVIDIAGPUToLLVM/BarrierOpToLLVM.cpp @@ -279,6 +279,41 @@ struct ArriveBarrierOpConversion return success(); } }; + +#if defined(__TLE__) && !defined(__HCU__) +struct NamedBarrierArriveOpConversion + : public ConvertOpToLLVMPattern { + using ConvertOpToLLVMPattern< + triton::nvidia_gpu::NamedBarrierArriveOp>::ConvertOpToLLVMPattern; + + LogicalResult + matchAndRewrite(triton::nvidia_gpu::NamedBarrierArriveOp op, + OpAdaptor adaptor, + ConversionPatternRewriter &rewriter) const override { + LLVM::createLLVMIntrinsicCallOp( + rewriter, op.getLoc(), "llvm.nvvm.barrier.cta.arrive.aligned.count", + TypeRange{}, {adaptor.getBar(), adaptor.getNumThreads()}); + rewriter.eraseOp(op); + return success(); + } +}; + +struct NamedBarrierWaitOpConversion + : public ConvertOpToLLVMPattern { + using ConvertOpToLLVMPattern< + triton::nvidia_gpu::NamedBarrierWaitOp>::ConvertOpToLLVMPattern; + + LogicalResult + matchAndRewrite(triton::nvidia_gpu::NamedBarrierWaitOp op, OpAdaptor adaptor, + ConversionPatternRewriter &rewriter) const override { + LLVM::createLLVMIntrinsicCallOp( + rewriter, op.getLoc(), "llvm.nvvm.barrier.cta.sync.aligned.count", + TypeRange{}, {adaptor.getBar(), adaptor.getNumThreads()}); + rewriter.eraseOp(op); + return success(); + } +}; +#endif } // namespace void mlir::triton::NVIDIA::populateBarrierOpToLLVMPatterns( @@ -290,4 +325,8 @@ void mlir::triton::NVIDIA::populateBarrierOpToLLVMPatterns( patterns.add(typeConverter, benefit, targetInfo); patterns.add(typeConverter, benefit); patterns.add(typeConverter, benefit); +#if defined(__TLE__) && !defined(__HCU__) + patterns.add(typeConverter, benefit); + patterns.add(typeConverter, benefit); +#endif } diff --git a/third_party/tle/dialect/include/IR/TleOps.td b/third_party/tle/dialect/include/IR/TleOps.td index c72513de42..3fffa4d472 100644 --- a/third_party/tle/dialect/include/IR/TleOps.td +++ b/third_party/tle/dialect/include/IR/TleOps.td @@ -7,6 +7,7 @@ include "mlir/Interfaces/SideEffectInterfaces.td" include "mlir/IR/CommonTypeConstraints.td" include "tle/dialect/include/IR/TleDialect.td" include "tle/dialect/include/IR/TleAttrDefs.td" +include "triton/Dialect/Triton/IR/TritonAttrDefs.td" include "triton/Dialect/Triton/IR/TritonInterfaces.td" include "triton/Dialect/Triton/IR/TritonTypes.td" include "triton/Dialect/TritonGPU/IR/TritonGPUAttrDefs.td" @@ -63,6 +64,7 @@ def Tle_TensorType : AnyTypeOf<[TT_Type, TTG_MemDescType]>; def Tle_ArgType : AnyTypeOf<[Tle_TensorType, LLVMPointerType, LLVMStructType]>; def Tle_LocalPointerResultType : AnyTypeOf<[TT_Tensor, TT_Ptr, TT_Int]>; def Tle_LocalPointerIndexType : AnyTypeOf<[TT_Tensor, TT_Int]>; +def Tle_WGMMAAType : AnyTypeOf<[TT_FpIntTensor, TTG_MemDescType]>; def Tle_LocalPointersOp : Tle_Op<"local_pointers", [Pure]> { let arguments = (ins TTG_MemDescType:$src, @@ -210,6 +212,96 @@ def Tle_WGMMASharedOperandFenceOp : Tle_Op<"wgmma_shared_operand_fence"> { let hasVerifier = 1; } +def Tle_WGMMAOp : Tle_Op<"wgmma", [ + TypesMatchWith<"result's type matches accumulator's type", "d", "c", "$_self"> +]> { + let summary = "TLE dialect-level Hopper WGMMA"; + + let description = [{ + Represents a user-level asynchronous Hopper WGMMA before TritonGPU layout + selection for the accumulator is materialized. The `a` operand may be a + shared-memory descriptor or an in-register tensor, `b` is a shared-memory + descriptor, and the accumulator/result are ordinary tensor values in the + surrounding IR layout. `triton-tle-lower-wgmma` rewrites this op to + `ttg.convert_layout` plus `ttng.warp_group_dot`. + }]; + + let arguments = (ins + Tle_WGMMAAType:$a, + TTG_MemDescType:$b, + TT_FpIntTensor:$c, + DefaultValuedAttr:$inputPrecision, + DefaultValuedAttr:$maxNumImpreciseAcc, + DefaultValuedAttr:$isAsync + ); + + let results = (outs TT_FpIntTensor:$d); + + let assemblyFormat = [{ + $a`,` $b`,` $c attr-dict `:` type($a) `*` + qualified(type($b)) `,` type($c) `->` type($d) + }]; + + let hasVerifier = 1; +} + +def Tle_WGMMAWaitOp : Tle_Op<"wgmma_wait", [ + TypesMatchWith<"result's type matches input's type", "output", "input", "$_self"> +]> { + let summary = "wait for TLE dialect-level asynchronous Hopper WGMMA"; + + let arguments = (ins + TT_FpIntTensor:$input, + I32Attr:$pendings + ); + + let results = (outs TT_FpIntTensor:$output); + + let assemblyFormat = "$input attr-dict `:` type($input) `->` type($output)"; + + let hasVerifier = 1; +} + +def Tle_BarrierAllocOp : Tle_Op<"barrier.alloc"> { + let summary = "create a typed TLE GPU barrier array"; + let arguments = (ins + I32Attr:$num_barriers, + I32Attr:$arrive_count, + I32Attr:$init_polarity, + OptionalAttr:$expect_bytes + ); + let results = (outs TTG_MemDescType:$result); + let assemblyFormat = "attr-dict `:` qualified(type($result))"; + let hasVerifier = 1; +} + +def Tle_BarrierWaitOp : Tle_Op<"barrier.wait"> { + let summary = "wait on a TLE GPU barrier slot"; + let arguments = (ins + TTG_MemDescType:$barrier, + Optional:$phase, + StrAttr:$backend, + I32Attr:$named_id, + I32Attr:$named_num_threads + ); + let assemblyFormat = "$barrier (`,` $phase^)? attr-dict `:` qualified(type($barrier))"; + let hasVerifier = 1; +} + +def Tle_BarrierArriveOp : Tle_Op<"barrier.arrive"> { + let summary = "arrive on a TLE GPU barrier slot"; + let arguments = (ins + TTG_MemDescType:$barrier, + Optional:$phase, + StrAttr:$backend, + I32Attr:$arrive_count, + I32Attr:$named_id, + I32Attr:$named_num_threads + ); + let assemblyFormat = "$barrier (`,` $phase^)? attr-dict `:` qualified(type($barrier))"; + let hasVerifier = 1; +} + def Tle_TMAStoreCommitGroupOp : Tle_Op<"tma_store.commit_group", [MemoryEffects<[MemRead, MemWrite]>]> { let summary = "commit the current TMA store bulk group"; diff --git a/third_party/tle/dialect/include/Transforms/Passes.td b/third_party/tle/dialect/include/Transforms/Passes.td index bef5701756..2ddd3b1bf0 100644 --- a/third_party/tle/dialect/include/Transforms/Passes.td +++ b/third_party/tle/dialect/include/Transforms/Passes.td @@ -254,6 +254,24 @@ def TritonTleLowerAsyncLoad "mlir::triton::TritonDialect"]; } +def TritonTleLowerWGMMA + : Pass<"triton-tle-lower-wgmma", "mlir::ModuleOp"> { + let summary = "lower TLE dialect-level WGMMA ops to native TritonGPU WGMMA"; + + let description = [{ + This pass rewrites `tle.wgmma` and `tle.wgmma_wait` to native + `ttng.warp_group_dot` and `ttng.warp_group_dot_wait` operations after + TritonGPU encodings are available. It materializes accumulator layout + conversions at the last responsible point so the frontend does not inject + TTG layout ops into the TTIR pipeline. + }]; + + let dependentDialects = ["mlir::triton::gpu::TritonGPUDialect", + "mlir::triton::nvidia_gpu::TritonNvidiaGPUDialect", + "mlir::triton::TritonDialect", + "mlir::triton::tle::TleDialect"]; +} + def TritonTleLowerTmaCopy : Pass { @@ -273,6 +291,41 @@ def TritonTleLowerTmaCopy "mlir::arith::ArithDialect"]; } +def TritonTleLowerBarriers + : Pass<"triton-tle-lower-barriers", "mlir::ModuleOp"> { + let summary = "lower TLE GPU barrier operations to NVIDIA barrier ops"; + + let description = [{ + This pass lowers the public ``tle.gpu`` barrier API IR. Mbarrier-backed + barriers become shared i64 allocations plus ttng init/wait/arrive ops. + Named barriers become statically numbered ttng named barrier ops. + }]; + + let dependentDialects = ["mlir::arith::ArithDialect", + "mlir::triton::gpu::TritonGPUDialect", + "mlir::triton::nvidia_gpu::TritonNvidiaGPUDialect", + "mlir::triton::tle::TleDialect"]; +} + +def TritonTleAllocateNamedBarriers + : Pass<"triton-tle-allocate-named-barriers", "mlir::ModuleOp"> { + let summary = "allocate physical NVIDIA named barrier ids for TLE virtual ids"; + + let description = [{ + This pass rewrites TLE virtual named barrier ids into physical NVIDIA + hardware named barrier ids. TLE frontend ids start at 16 so each user + barrier remains distinguishable in IR before this late pass. The pass + reserves the barrier ids used by warp-specialization lowering, preserves + existing physical ids in the 0..15 range, and assigns virtual ids to free + physical ids. + }]; + + let dependentDialects = ["mlir::arith::ArithDialect", + "mlir::triton::gpu::TritonGPUDialect", + "mlir::triton::nvidia_gpu::TritonNvidiaGPUDialect", + "mlir::triton::TritonDialect"]; +} + def TritonTleScheduleTmaStoreSync : Pass<"triton-tle-schedule-tma-store-sync", "mlir::ModuleOp"> { let summary = "schedule explicit TLE TMA store commit and wait operations"; diff --git a/third_party/tle/dialect/lib/IR/Ops.cpp b/third_party/tle/dialect/lib/IR/Ops.cpp index 669ea0bcf4..e3a868b9b6 100644 --- a/third_party/tle/dialect/lib/IR/Ops.cpp +++ b/third_party/tle/dialect/lib/IR/Ops.cpp @@ -198,6 +198,138 @@ LogicalResult WGMMASharedOperandFenceOp::verify() { return success(); } +LogicalResult WGMMAOp::verify() { + auto aType = cast(getA().getType()); + auto bType = cast(getB().getType()); + auto cType = cast(getC().getType()); + auto dType = cast(getD().getType()); + + if (aType.getRank() != 2 || bType.getRank() != 2 || cType.getRank() != 2) + return emitOpError("expects rank-2 A, B, and accumulator operands"); + if (!isa(bType.getMemorySpace())) + return emitOpError("expects shared-memory B descriptor"); + if (auto aMemDescType = + dyn_cast(getA().getType())) { + if (!isa(aMemDescType.getMemorySpace())) + return emitOpError("expects shared-memory A descriptor"); + } + + ArrayRef aShape = aType.getShape(); + ArrayRef bShape = bType.getShape(); + ArrayRef cShape = cType.getShape(); + ArrayRef dShape = dType.getShape(); + if (aShape[1] != bShape[0]) + return emitOpError("expects A and B K dimensions to match"); + if (cShape[0] != aShape[0] || cShape[1] != bShape[1]) + return emitOpError("expects accumulator shape to be MxN from A and B"); + if (dShape != cShape) + return emitOpError("expects result shape to match accumulator shape"); + + if (aShape[0] < 64 || aShape[0] % 64 != 0) + return emitOpError("expects M dimension to be divisible by 64"); + if (bShape[1] < 8 || bShape[1] % 8 != 0) + return emitOpError("expects N dimension to be divisible by 8"); + if (aShape[1] < 16) + return emitOpError("expects K dimension to be at least 16"); + return success(); +} + +LogicalResult WGMMAWaitOp::verify() { + auto pendings = getOperation()->getAttrOfType("pendings"); + if (pendings.getInt() < 0) + return emitOpError("expects non-negative pendings"); + return success(); +} + +static LogicalResult +verifyBarrierMemDesc(Operation *op, triton::gpu::MemDescType type, bool array) { + if (!type.getElementType().isInteger(64)) + return op->emitOpError("expects i64 barrier storage"); + if (!isa(type.getMemorySpace())) + return op->emitOpError("expects shared-memory barrier storage"); + if (!type.getMutableMemory()) + return op->emitOpError("expects mutable barrier storage"); + + ArrayRef shape = type.getShape(); + if (array) { + if (shape.size() != 2 || shape[0] <= 0 || shape[1] != 1) + return op->emitOpError("expects barrier array type shaped Nx1xi64"); + } else { + if (shape != ArrayRef({1})) + return op->emitOpError("expects barrier slot type shaped 1xi64"); + } + return success(); +} + +static LogicalResult verifyBarrierBackend(Operation *op, StringRef backend, + bool hasPhase, int64_t namedId, + int64_t namedNumThreads) { + static constexpr int64_t kFirstVirtualNamedBarrierId = 16; + if (backend == "mbarrier") { + if (!hasPhase) + return op->emitOpError("mbarrier backend requires phase"); + return success(); + } + if (backend != "named") + return op->emitOpError("backend must be 'mbarrier' or 'named'"); + if (hasPhase) + return op->emitOpError("named barrier backend does not accept phase"); + bool isPhysicalNamedId = namedId >= 1 && namedId <= 15; + bool isVirtualNamedId = namedId >= kFirstVirtualNamedBarrierId; + if (!isPhysicalNamedId && !isVirtualNamedId) + return op->emitOpError("named barrier id must be a physical id in range " + "[1, 15] or a TLE virtual id >= ") + << kFirstVirtualNamedBarrierId; + if (namedNumThreads <= 0) + return op->emitOpError("named barrier thread count must be positive"); + return success(); +} + +LogicalResult BarrierAllocOp::verify() { + if (failed(verifyBarrierMemDesc(getOperation(), getResult().getType(), + /*array=*/true))) + return failure(); + auto resultTy = getResult().getType(); + if (getNumBarriers() <= 0) + return emitOpError("num_barriers must be positive"); + if (getArriveCount() <= 0) + return emitOpError("arrive_count must be positive"); + if (getInitPolarity() != 0 && getInitPolarity() != 1) + return emitOpError("init_polarity must be 0 or 1"); + if (getNumBarriers() != resultTy.getShape()[0]) + return emitOpError("num_barriers must match result leading dimension"); + if (auto expectBytes = + getOperation()->getAttrOfType("expect_bytes")) { + if (expectBytes.getInt() <= 0) + return emitOpError("expect_bytes must be positive when present"); + } + return success(); +} + +LogicalResult BarrierWaitOp::verify() { + if (failed(verifyBarrierMemDesc(getOperation(), getBarrier().getType(), + /*array=*/false))) + return failure(); + StringRef backend = + getOperation()->getAttrOfType("backend").getValue(); + return verifyBarrierBackend(getOperation(), backend, bool(getPhase()), + getNamedId(), getNamedNumThreads()); +} + +LogicalResult BarrierArriveOp::verify() { + if (failed(verifyBarrierMemDesc(getOperation(), getBarrier().getType(), + /*array=*/false))) + return failure(); + if (getArriveCount() <= 0) + return emitOpError("arrive_count must be positive"); + StringRef backend = + getOperation()->getAttrOfType("backend").getValue(); + if (backend == "named" && getArriveCount() != 1) + return emitOpError("named barrier arrive requires arrive_count = 1"); + return verifyBarrierBackend(getOperation(), backend, bool(getPhase()), + getNamedId(), getNamedNumThreads()); +} + static bool isAsciiIdentStart(char c) { return (c >= 'A' && c <= 'Z') || (c >= 'a' && c <= 'z'); } diff --git a/third_party/tle/dialect/lib/Transforms/CMakeLists.txt b/third_party/tle/dialect/lib/Transforms/CMakeLists.txt index 7d2e6c7f7f..f77a3b9372 100644 --- a/third_party/tle/dialect/lib/Transforms/CMakeLists.txt +++ b/third_party/tle/dialect/lib/Transforms/CMakeLists.txt @@ -12,8 +12,11 @@ add_triton_library(TritonTLETransforms TleDowngradeInvalidAsyncCopy.cpp TleOptimizeExclusiveCumsumLayouts.cpp TleLowerExclusiveCumsum.cpp + TleLowerWGMMA.cpp TleLowerAsyncLoad.cpp TleLowerPipeToNvws.cpp + TleLowerBarriers.cpp + TleAllocateNamedBarriers.cpp TleTileToLLVMUtils.cpp ExtractTileToLLVM.cpp InsertTileToLLVM.cpp diff --git a/third_party/tle/dialect/lib/Transforms/TleAllocateNamedBarriers.cpp b/third_party/tle/dialect/lib/Transforms/TleAllocateNamedBarriers.cpp new file mode 100644 index 0000000000..6b3e7bcce7 --- /dev/null +++ b/third_party/tle/dialect/lib/Transforms/TleAllocateNamedBarriers.cpp @@ -0,0 +1,202 @@ +#include "mlir/Dialect/Arith/IR/Arith.h" +#include "mlir/IR/BuiltinOps.h" +#include "mlir/IR/PatternMatch.h" +#include "mlir/Pass/Pass.h" +#include "tle/dialect/include/Transforms/Passes.h" +#include "triton/Dialect/Triton/IR/Dialect.h" +#include "triton/Dialect/TritonGPU/IR/Dialect.h" +#include "triton/Dialect/TritonNvidiaGPU/IR/Dialect.h" +#include "llvm/ADT/MapVector.h" +#include "llvm/ADT/SmallBitVector.h" +#include +#include + +namespace mlir::triton::tle { + +namespace ttg = mlir::triton::gpu; +namespace ttng = mlir::triton::nvidia_gpu; + +#define GEN_PASS_DEF_TRITONTLEALLOCATENAMEDBARRIERS +#include "tle/dialect/include/Transforms/Passes.h.inc" + +#if defined(__HCU__) +namespace { + +struct TritonTleAllocateNamedBarriers + : public impl::TritonTleAllocateNamedBarriersBase< + TritonTleAllocateNamedBarriers> { + void runOnOperation() override {} +}; + +} // namespace +#else +namespace { + +constexpr int64_t kFirstVirtualNamedBarrierId = 16; +constexpr int64_t kNumPhysicalNamedBarriers = 16; +// Keep virtual TLE allocations in the historical user-visible range. Existing +// physical id 0 ops are preserved and still block future virtual allocations. +constexpr int64_t kFirstTleAllocatedPhysicalNamedBarrierId = 1; +constexpr int64_t kDefaultWarpGroupBarrierIdx = 0; +constexpr int64_t kSwitchLoopBarrierIdx = 1; +constexpr int64_t kNumWarpSpecializeReservedBarriers = 2; + +static std::optional getConstantI32(Value value) { + if (auto op = value.getDefiningOp()) + return op.value(); + if (auto op = value.getDefiningOp()) { + if (auto attr = dyn_cast(op.getValue())) + return attr.getInt(); + } + return std::nullopt; +} + +static void reserveWarpSpecializeBarrierIds(triton::FuncOp func, + llvm::SmallBitVector &reserved) { + int64_t maxPartitions = 0; + func.walk([&](ttg::WarpSpecializeOp ws) { + maxPartitions = + std::max(maxPartitions, ws.getPartitionRegions().size()); + }); + + if (maxPartitions == 0) + return; + + reserved.set(kDefaultWarpGroupBarrierIdx); + reserved.set(kSwitchLoopBarrierIdx); + for (int64_t i = 0; i < maxPartitions; ++i) { + int64_t id = kNumWarpSpecializeReservedBarriers + i; + if (id < kNumPhysicalNamedBarriers) + reserved.set(id); + } +} + +static LogicalResult +collectNamedBarrierId(Operation *op, Value idValue, + llvm::SmallBitVector &reserved, + llvm::SmallBitVector &used, + llvm::MapVector &virtualToPhysical) { + std::optional id = getConstantI32(idValue); + if (!id) + return op->emitOpError("requires a constant named barrier id before " + "physical allocation"); + + if (*id >= kFirstVirtualNamedBarrierId) { + virtualToPhysical.insert({*id, -1}); + return success(); + } + + if (*id < 0 || *id >= kNumPhysicalNamedBarriers) + return op->emitOpError("has invalid physical named barrier id ") + << *id << "; expected 0..15 or a TLE virtual id >= " + << kFirstVirtualNamedBarrierId; + + // Physical ids are treated as fixed allocations. They may have been created + // by lower-level NVIDIA passes, while TLE virtual ids are remapped below. + used.set(*id); + return success(); +} + +static LogicalResult +allocateVirtualIds(triton::FuncOp func, llvm::SmallBitVector &reserved, + llvm::SmallBitVector &used, + llvm::MapVector &virtualToPhysical) { + for (int64_t i = 0; i < kNumPhysicalNamedBarriers; ++i) + if (reserved.test(i)) + used.set(i); + for (auto &entry : virtualToPhysical) { + int64_t physicalId = -1; + for (int64_t candidate = kFirstTleAllocatedPhysicalNamedBarrierId; + candidate < kNumPhysicalNamedBarriers; ++candidate) { + if (!used.test(candidate)) { + physicalId = candidate; + break; + } + } + + if (physicalId < 0) + return func.emitError("cannot allocate physical NVIDIA named barrier " + "id for virtual TLE named barrier ") + << entry.first << "; all ids in allocatable range [" + << kFirstTleAllocatedPhysicalNamedBarrierId << ", " + << (kNumPhysicalNamedBarriers - 1) + << "] are already reserved or used"; + + entry.second = physicalId; + used.set(physicalId); + } + return success(); +} + +static void rewriteNamedBarrierId( + Operation *op, Value idValue, + const llvm::MapVector &virtualToPhysical) { + std::optional id = getConstantI32(idValue); + if (!id || *id < kFirstVirtualNamedBarrierId) + return; + + auto it = virtualToPhysical.find(*id); + if (it == virtualToPhysical.end()) + return; + + OpBuilder builder(op); + Value physicalId = + builder.create(op->getLoc(), it->second, 32); + op->setOperand(0, physicalId); +} + +static LogicalResult allocateForFunction(triton::FuncOp func) { + llvm::SmallBitVector reserved(kNumPhysicalNamedBarriers, false); + llvm::SmallBitVector used(kNumPhysicalNamedBarriers, false); + llvm::MapVector virtualToPhysical; + + reserveWarpSpecializeBarrierIds(func, reserved); + + WalkResult collectResult = func.walk([&](Operation *op) -> WalkResult { + if (auto wait = dyn_cast(op)) { + if (failed(collectNamedBarrierId(op, wait.getBar(), reserved, used, + virtualToPhysical))) + return WalkResult::interrupt(); + } else if (auto arrive = dyn_cast(op)) { + if (failed(collectNamedBarrierId(op, arrive.getBar(), reserved, used, + virtualToPhysical))) + return WalkResult::interrupt(); + } + return WalkResult::advance(); + }); + + if (collectResult.wasInterrupted()) + return failure(); + + if (failed(allocateVirtualIds(func, reserved, used, virtualToPhysical))) + return failure(); + + func.walk([&](Operation *op) { + if (auto wait = dyn_cast(op)) + rewriteNamedBarrierId(op, wait.getBar(), virtualToPhysical); + else if (auto arrive = dyn_cast(op)) + rewriteNamedBarrierId(op, arrive.getBar(), virtualToPhysical); + }); + + return success(); +} + +struct TritonTleAllocateNamedBarriers + : public impl::TritonTleAllocateNamedBarriersBase< + TritonTleAllocateNamedBarriers> { + void runOnOperation() override { + ModuleOp module = getOperation(); + WalkResult result = module.walk([&](triton::FuncOp func) -> WalkResult { + if (failed(allocateForFunction(func))) + return WalkResult::interrupt(); + return WalkResult::advance(); + }); + if (result.wasInterrupted()) + signalPassFailure(); + } +}; + +} // namespace +#endif + +} // namespace mlir::triton::tle diff --git a/third_party/tle/dialect/lib/Transforms/TleLowerBarriers.cpp b/third_party/tle/dialect/lib/Transforms/TleLowerBarriers.cpp new file mode 100644 index 0000000000..07db3d8cad --- /dev/null +++ b/third_party/tle/dialect/lib/Transforms/TleLowerBarriers.cpp @@ -0,0 +1,152 @@ +#include "mlir/Dialect/Arith/IR/Arith.h" +#include "mlir/IR/BuiltinOps.h" +#include "mlir/IR/PatternMatch.h" +#include "mlir/Pass/Pass.h" +#include "tle/dialect/include/IR/Dialect.h" +#include "tle/dialect/include/Transforms/Passes.h" +#include "triton/Dialect/TritonGPU/IR/Dialect.h" +#include "triton/Dialect/TritonNvidiaGPU/IR/Dialect.h" + +namespace mlir::triton::tle { + +namespace ttg = mlir::triton::gpu; +namespace ttng = mlir::triton::nvidia_gpu; + +#define GEN_PASS_DEF_TRITONTLELOWERBARRIERS +#include "tle/dialect/include/Transforms/Passes.h.inc" + +namespace { + +static int64_t getI32Attr(Operation *op, StringRef name) { + return op->getAttrOfType(name).getInt(); +} + +static ttg::MemDescType getBarrierSlotType(ttg::MemDescType arrayTy) { + auto context = arrayTy.getContext(); + auto ctaLayout = ttg::CTAEncodingAttr::getDefault(context, 1); + Attribute slotEncoding = + ttg::SwizzledSharedEncodingAttr::get(context, 1, 1, 1, {0}, ctaLayout); + return ttg::MemDescType::get({1}, arrayTy.getElementType(), slotEncoding, + arrayTy.getMemorySpace(), + arrayTy.getMutableMemory()); +} + +static Value createBarrierSlot(OpBuilder &builder, Location loc, Value array, + int64_t index) { + auto arrayTy = cast(array.getType()); + auto slotTy = getBarrierSlotType(arrayTy); + Value idx = builder.create(loc, index, 32); + return builder.create(loc, slotTy, array, idx); +} + +#if !defined(__HCU__) +static std::pair +createNamedBarrierOperands(OpBuilder &builder, Location loc, Operation *op) { + Value id = + builder.create(loc, getI32Attr(op, "named_id"), 32); + Value threads = builder.create( + loc, getI32Attr(op, "named_num_threads"), 32); + return {id, threads}; +} +#endif + +struct TritonTleLowerBarriers + : public impl::TritonTleLowerBarriersBase { + void runOnOperation() override { + ModuleOp module = getOperation(); + + SmallVector waits; + SmallVector arrives; + SmallVector allocs; + module.walk([&](Operation *op) { + if (auto wait = dyn_cast(op)) + waits.push_back(wait); + else if (auto arrive = dyn_cast(op)) + arrives.push_back(arrive); + else if (auto alloc = dyn_cast(op)) + allocs.push_back(alloc); + }); + + for (BarrierWaitOp op : waits) { + OpBuilder builder(op); + Location loc = op.getLoc(); + StringRef backend = op->getAttrOfType("backend").getValue(); + if (backend == "mbarrier") { + builder.create(loc, op.getBarrier(), + op.getPhase()); + } else { +#if defined(__HCU__) + op.emitOpError("named barrier lowering is only supported on NVIDIA " + "backend"); + signalPassFailure(); + return; +#else + auto [id, threads] = + createNamedBarrierOperands(builder, loc, op.getOperation()); + builder.create(loc, id, threads); +#endif + } + op.erase(); + } + + for (BarrierArriveOp op : arrives) { + OpBuilder builder(op); + Location loc = op.getLoc(); + StringRef backend = op->getAttrOfType("backend").getValue(); + if (backend == "mbarrier") { + int64_t count = getI32Attr(op.getOperation(), "arrive_count"); + builder.create(loc, op.getBarrier(), + static_cast(count)); + } else { +#if defined(__HCU__) + op.emitOpError("named barrier lowering is only supported on NVIDIA " + "backend"); + signalPassFailure(); + return; +#else + auto [id, threads] = + createNamedBarrierOperands(builder, loc, op.getOperation()); + builder.create(loc, id, threads); +#endif + } + op.erase(); + } + + for (BarrierAllocOp op : allocs) { + OpBuilder builder(op); + Location loc = op.getLoc(); + auto arrayTy = op.getResult().getType(); + bool onlyUnusedViews = true; + SmallVector deadViews; + for (OpOperand &use : op.getResult().getUses()) { + auto view = dyn_cast(use.getOwner()); + if (!view || !view->use_empty()) { + onlyUnusedViews = false; + break; + } + deadViews.push_back(view); + } + if (onlyUnusedViews) { + for (auto view : deadViews) + view.erase(); + op.erase(); + continue; + } + + Value alloc = builder.create(loc, arrayTy); + int64_t numBarriers = getI32Attr(op.getOperation(), "num_barriers"); + int64_t arriveCount = getI32Attr(op.getOperation(), "arrive_count"); + for (int64_t i = 0; i < numBarriers; ++i) { + Value slot = createBarrierSlot(builder, loc, alloc, i); + builder.create(loc, slot, + static_cast(arriveCount)); + } + op.getResult().replaceAllUsesWith(alloc); + op.erase(); + } + } +}; + +} // namespace + +} // namespace mlir::triton::tle diff --git a/third_party/tle/dialect/lib/Transforms/TleLowerTmaCopy.cpp b/third_party/tle/dialect/lib/Transforms/TleLowerTmaCopy.cpp index 8af0de72fa..bf4e0e6bba 100644 --- a/third_party/tle/dialect/lib/Transforms/TleLowerTmaCopy.cpp +++ b/third_party/tle/dialect/lib/Transforms/TleLowerTmaCopy.cpp @@ -149,48 +149,76 @@ class TMACopyLowering : public OpRewritePattern { auto tensorType = RankedTensorType::get( dstType.getShape(), dstType.getElementType(), dstType.getEncoding()); - // Create minimal mbarrier allocation with #shared2 encoding (similar to - // our current implementation) - auto mbarrierCTALayout = gpu::CTAEncodingAttr::fromSplitParams( - tensorType.getContext(), {1}, {1}, {0}); - auto mbarrierEncoding = gpu::SwizzledSharedEncodingAttr::get( - tensorType.getContext(), 1, 1, 1, {0}, mbarrierCTALayout); - Attribute sharedMemorySpace = - triton::gpu::SharedMemorySpaceAttr::get(op.getContext()); - - gpu::MemDescType mbarrierMemDescType = - gpu::MemDescType::get({1}, rewriter.getI64Type(), mbarrierEncoding, - sharedMemorySpace, /*mutableMemory=*/true); - - Value mbarrierAlloc = - rewriter.create(loc, mbarrierMemDescType); - rewriter.create(loc, mbarrierAlloc, 1); - - // Calculate size in bytes - auto encoding = getEncodingFromDescriptor(op, tensorType, op.getSrc()); - auto shapePerCTA = getShapePerCTA(encoding, tensorType.getShape()); - int sizeInBytes = product(shapePerCTA) * - tensorType.getElementType().getIntOrFloatBitWidth() / 8; - Value pred = rewriter.create(loc, 1, 1); - rewriter.create(loc, mbarrierAlloc, - sizeInBytes, pred); // Create TMA indices auto indices = translateTMAIndices(rewriter, op.getLoc(), srcType.getBlockType().getEncoding(), op.getIndices()); - // Perform async TMA copy from global to existing shared memory - rewriter.create( - op.getLoc(), op.getSrc(), indices, mbarrierAlloc, dstMemDesc, pred); - - // Wait for completion and invalidate barrier - Value phase = rewriter.create(loc, 0, 32); - rewriter.create(loc, mbarrierAlloc, phase); - rewriter.create(loc, mbarrierAlloc); +#if !defined(__HCU__) + if (Value userBarrier = op.getBarrier()) { + auto expectBytes = + op->getAttrOfType(op.getExpectBytesAttrName()); + if (!expectBytes || expectBytes.getInt() <= 0) + return op.emitOpError("with explicit completion barrier requires " + "positive expect_bytes"); + rewriter.create( + loc, userBarrier, static_cast(expectBytes.getInt()), pred); + rewriter.create( + op.getLoc(), op.getSrc(), indices, userBarrier, dstMemDesc, pred); + } else { + if (op->hasAttr(op.getExpectBytesAttrName())) + return op.emitOpError("expect_bytes requires an explicit completion " + "barrier"); +#endif + + // Create minimal mbarrier allocation with #shared2 encoding (similar to + // our current implementation) + auto mbarrierCTALayout = gpu::CTAEncodingAttr::fromSplitParams( + tensorType.getContext(), {1}, {1}, {0}); + auto mbarrierEncoding = gpu::SwizzledSharedEncodingAttr::get( + tensorType.getContext(), 1, 1, 1, {0}, mbarrierCTALayout); + Attribute sharedMemorySpace = + triton::gpu::SharedMemorySpaceAttr::get(op.getContext()); + + gpu::MemDescType mbarrierMemDescType = + gpu::MemDescType::get({1}, rewriter.getI64Type(), mbarrierEncoding, + sharedMemorySpace, /*mutableMemory=*/true); + + Value mbarrierAlloc = + rewriter.create(loc, mbarrierMemDescType); + rewriter.create(loc, mbarrierAlloc, 1); + + // Calculate size in bytes + auto encoding = getEncodingFromDescriptor(op, tensorType, op.getSrc()); + auto shapePerCTA = getShapePerCTA(encoding, tensorType.getShape()); + int sizeInBytes = product(shapePerCTA) * + tensorType.getElementType().getIntOrFloatBitWidth() / + 8; + + rewriter.create(loc, mbarrierAlloc, + sizeInBytes, pred); + + // Perform async TMA copy from global to existing shared memory + rewriter.create( + op.getLoc(), op.getSrc(), indices, mbarrierAlloc, dstMemDesc, pred); + + // Wait for completion and invalidate barrier + Value phase = rewriter.create(loc, 0, 32); + rewriter.create(loc, mbarrierAlloc, phase); + rewriter.create(loc, mbarrierAlloc); +#if !defined(__HCU__) + } +#endif } else { +#if !defined(__HCU__) + if (op.getBarrier()) + return op.emitOpError( + "barrier is only supported for global-to-shared TMA copy"); +#endif + // Store from shared memory to global memory auto dstType = cast(op.getDst().getType()); auto srcType = cast(op.getSrc().getType()); diff --git a/third_party/tle/dialect/lib/Transforms/TleLowerWGMMA.cpp b/third_party/tle/dialect/lib/Transforms/TleLowerWGMMA.cpp new file mode 100644 index 0000000000..a9cb0259a2 --- /dev/null +++ b/third_party/tle/dialect/lib/Transforms/TleLowerWGMMA.cpp @@ -0,0 +1,276 @@ +#include "mlir/Dialect/SCF/IR/SCF.h" +#include "mlir/IR/BuiltinOps.h" +#include "mlir/IR/PatternMatch.h" +#include "mlir/Pass/Pass.h" +#include "tle/dialect/include/IR/Dialect.h" +#include "tle/dialect/include/Transforms/Passes.h" +#include "triton/Dialect/Triton/IR/Dialect.h" +#include "triton/Dialect/TritonGPU/IR/Dialect.h" +#include "triton/Dialect/TritonGPU/Transforms/Utility.h" +#include "triton/Dialect/TritonNvidiaGPU/IR/Dialect.h" +#include "llvm/ADT/DenseMap.h" +#include "llvm/ADT/STLExtras.h" + +namespace mlir::triton::tle { + +namespace ttg = mlir::triton::gpu; +namespace ttng = mlir::triton::nvidia_gpu; + +#define GEN_PASS_DEF_TRITONTLELOWERWGMMA +#include "tle/dialect/include/Transforms/Passes.h.inc" + +namespace { + +static constexpr unsigned kScfForControlOperands = 3; + +static std::optional getForInitArgIndex(OpOperand &use) { + auto forOp = dyn_cast(use.getOwner()); + if (!forOp) + return std::nullopt; + unsigned operandNumber = use.getOperandNumber(); + if (operandNumber < kScfForControlOperands) + return std::nullopt; + return operandNumber - kScfForControlOperands; +} + +static bool isForYield(OpOperand &use) { + auto yieldOp = dyn_cast(use.getOwner()); + return yieldOp && isa(yieldOp->getParentOp()); +} + +static bool isAsyncAccumulatorUse(OpOperand &use) { + Operation *user = use.getOwner(); + if (isa(user)) + return use.getOperandNumber() == 2; + if (isa(user)) + return true; + if (getForInitArgIndex(use)) + return true; + return isForYield(use); +} + +static LogicalResult verifyWGMMAUses(WGMMAOp op) { + for (OpOperand &use : op.getD().getUses()) { + Operation *user = use.getOwner(); + if (auto next = dyn_cast(user)) { + if (use.getOperandNumber() == 2) + continue; + return op.emitOpError("result may only feed the accumulator operand of " + "another tle.wgmma or tle.wgmma_wait"); + } + if (isa(user)) + continue; + if (getForInitArgIndex(use)) + continue; + return op.emitOpError("async result must be consumed by tle.wgmma_wait " + "before ordinary tensor use"); + } + return success(); +} + +static RankedTensorType getMMAType(WGMMAOp op) { + auto context = op.getContext(); + auto accType = cast(op.getC().getType()); + auto aElemType = + cast(op.getA().getType()).getElementType(); + std::optional maybeNumWarps = + ttg::maybeLookupNumWarps(op.getOperation()); + if (!maybeNumWarps) { + op.emitOpError("requires a contextual `ttg.num-warps` to select the " + "WGMMA accumulator layout"); + return {}; + } + + unsigned numWarps = *maybeNumWarps; + SmallVector retShapePerCTA = + accType.getEncoding() ? ttg::getShapePerCTA(accType) + : SmallVector(accType.getShape()); + SmallVector instrShape = + mmaVersionToInstrShape(3, retShapePerCTA, aElemType, numWarps); + SmallVector warpsPerCTA = {numWarps, 1}; + SmallVector CTAsPerCGA = {1, 1}; + SmallVector CTASplitNum = {1, 1}; + SmallVector CTAOrder = {1, 0}; + auto CTALayout = ttg::CTAEncodingAttr::fromSplitParams(context, CTAsPerCGA, + CTASplitNum, CTAOrder); + auto mmaEncoding = ttg::NvidiaMmaEncodingAttr::get(context, 3, 0, warpsPerCTA, + CTALayout, instrShape); + return RankedTensorType::get(accType.getShape(), accType.getElementType(), + mmaEncoding); +} + +static bool isMMAEncoded(Value value) { + auto type = dyn_cast(value.getType()); + if (!type) + return false; + Attribute encoding = type.getEncoding(); + return encoding && isa(encoding); +} + +static Value lookupEncodedAccumulator(Value value, + DenseMap &encodedAccs) { + if (Value mapped = encodedAccs.lookup(value)) + return mapped; + if (isMMAEncoded(value)) + return value; + return {}; +} + +static void convertLoopCarriedAccumulator(OpOperand &use, Value encodedInit, + DenseMap &encodedAccs) { + std::optional maybeIndex = getForInitArgIndex(use); + if (!maybeIndex) + return; + + auto forOp = cast(use.getOwner()); + unsigned initIndex = *maybeIndex; + Type encodedType = encodedInit.getType(); + forOp->setOperand(use.getOperandNumber(), encodedInit); + + BlockArgument regionArg = + forOp.getBody()->getArgument(forOp.getNumInductionVars() + initIndex); + regionArg.setType(encodedType); + forOp.getResult(initIndex).setType(encodedType); + + encodedAccs[regionArg] = regionArg; + encodedAccs[forOp.getResult(initIndex)] = forOp.getResult(initIndex); +} + +static Value materializeMMAAccumulator(OpBuilder &builder, WGMMAOp op, + RankedTensorType mmaType, + DenseMap &encodedAccs) { + Value acc = op.getC(); + if (Value mapped = lookupEncodedAccumulator(acc, encodedAccs)) + return mapped; + + auto accType = cast(acc.getType()); + if (accType.getEncoding() == mmaType.getEncoding()) + return acc; + return ttg::ConvertLayoutOp::create(builder, op.getLoc(), mmaType, acc); +} + +static Value materializeAOperand(OpBuilder &builder, WGMMAOp op, + RankedTensorType mmaType) { + Value a = op.getA(); + if (isa(a.getType())) + return a; + + auto aType = cast(a.getType()); + Attribute dotEncoding = ttg::DotOperandEncodingAttr::get( + op.getContext(), /*opIdx=*/0, mmaType.getEncoding(), + aType.getElementType()); + auto dotType = aType.cloneWithEncoding(dotEncoding); + if (aType == dotType) + return a; + return ttg::ConvertLayoutOp::create(builder, op.getLoc(), dotType, a); +} + +struct TritonTleLowerWGMMAPass + : public impl::TritonTleLowerWGMMABase { + void runOnOperation() override { + ModuleOp module = getOperation(); + bool failed = false; + + module.walk([&](triton::FuncOp func) { + if (failed) + return; + + SmallVector worklist; + SmallVector wgmmas; + func.walk([&](Operation *op) { + if (isa(op)) + worklist.push_back(op); + if (auto wgmma = dyn_cast(op)) + wgmmas.push_back(wgmma); + }); + + for (WGMMAOp op : wgmmas) { + if (mlir::failed(verifyWGMMAUses(op))) { + failed = true; + return; + } + } + + DenseMap encodedAccs; + for (Operation *op : worklist) { + OpBuilder builder(op); + if (auto wgmma = dyn_cast(op)) { + RankedTensorType mmaType = getMMAType(wgmma); + if (!mmaType) { + failed = true; + return; + } + Value acc = + materializeMMAAccumulator(builder, wgmma, mmaType, encodedAccs); + Value a = materializeAOperand(builder, wgmma, mmaType); + if (!acc) { + failed = true; + return; + } + auto nativeDot = ttng::WarpGroupDotOp::create( + builder, wgmma.getLoc(), acc.getType(), a, wgmma.getB(), acc, + Value(), wgmma.getInputPrecision(), wgmma.getMaxNumImpreciseAcc(), + wgmma.getIsAsync()); + encodedAccs[wgmma.getD()] = nativeDot.getD(); + for (OpOperand &use : + llvm::make_early_inc_range(wgmma.getD().getUses())) + convertLoopCarriedAccumulator(use, nativeDot.getD(), encodedAccs); + continue; + } + + auto wait = cast(op); + Value encodedInput = + lookupEncodedAccumulator(wait.getInput(), encodedAccs); + if (!encodedInput) { + wait.emitOpError("input must be the async result of tle.wgmma"); + failed = true; + return; + } + + SmallVector waitInputs{encodedInput}; + auto nativeWait = ttng::WarpGroupDotWaitOp::create( + builder, wait.getLoc(), waitInputs, wait.getPendings()); + Value waited = nativeWait.getResult(0); + if (wait.getPendings() > 0) { + // Rewrite every async accumulator use directly before erasing this + // wait. Keeping wait.getOutput() in encodedAccs would leave a map + // entry keyed by a soon-to-be-dangling Value; later rewrites may then + // accidentally pick an accumulator from another isolated region. + Value released; + for (OpOperand &use : + llvm::make_early_inc_range(wait.getOutput().getUses())) { + if (isAsyncAccumulatorUse(use)) { + if (getForInitArgIndex(use)) + convertLoopCarriedAccumulator(use, waited, encodedAccs); + else + use.set(waited); + continue; + } + + if (!released) + released = ttg::ConvertLayoutOp::create( + builder, wait.getLoc(), wait.getOutput().getType(), waited); + use.set(released); + } + wait.erase(); + continue; + } + + Value released = ttg::ConvertLayoutOp::create( + builder, wait.getLoc(), wait.getOutput().getType(), waited); + wait.getOutput().replaceAllUsesWith(released); + wait.erase(); + } + + for (WGMMAOp op : llvm::reverse(wgmmas)) + op.erase(); + }); + + if (failed) + signalPassFailure(); + } +}; + +} // namespace + +} // namespace mlir::triton::tle diff --git a/third_party/tle/test/GPU/test_tle_allocate_named_barriers.mlir b/third_party/tle/test/GPU/test_tle_allocate_named_barriers.mlir new file mode 100644 index 0000000000..64b27094dc --- /dev/null +++ b/third_party/tle/test/GPU/test_tle_allocate_named_barriers.mlir @@ -0,0 +1,76 @@ +// RUN: triton-opt %s -triton-tle-allocate-named-barriers -split-input-file | FileCheck %s + +module attributes {"ttg.num-ctas" = 1 : i32, "ttg.num-warps" = 4 : i32, ttg.target = "cuda:90", "ttg.threads-per-warp" = 32 : i32} { + // CHECK-LABEL: tt.func @remap_virtual_ids_without_warpspec + tt.func @remap_virtual_ids_without_warpspec() { + %v0 = arith.constant 16 : i32 + %v1 = arith.constant 17 : i32 + %threads = arith.constant 256 : i32 + + // CHECK: %[[ID0:.+]] = arith.constant 1 : i32 + // CHECK: ttng.wait_barrier_named %[[ID0]], {{.*}} : i32, i32 + ttng.wait_barrier_named %v0, %threads : i32, i32 + + // CHECK: %[[ID1:.+]] = arith.constant 2 : i32 + // CHECK: ttng.arrive_barrier_named %[[ID1]], {{.*}} : i32, i32 + ttng.arrive_barrier_named %v1, %threads : i32, i32 + + // CHECK: %[[ID0_AGAIN:.+]] = arith.constant 1 : i32 + // CHECK: ttng.arrive_barrier_named %[[ID0_AGAIN]], {{.*}} : i32, i32 + ttng.arrive_barrier_named %v0, %threads : i32, i32 + tt.return + } +} + +// ----- + +module attributes {"ttg.num-ctas" = 1 : i32, "ttg.num-warps" = 8 : i32, ttg.target = "cuda:90", "ttg.threads-per-warp" = 32 : i32} { + // CHECK-LABEL: tt.func @remap_virtual_ids_with_warpspec + tt.func @remap_virtual_ids_with_warpspec() { + ttg.warp_specialize() attributes {requestedRegisters = array} + default { + ttg.warp_yield + } + partition0() num_warps(4) { + %v0 = arith.constant 16 : i32 + %threads = arith.constant 256 : i32 + // CHECK: %[[ID4:.+]] = arith.constant 4 : i32 + // CHECK: ttng.wait_barrier_named %[[ID4]], {{.*}} : i32, i32 + ttng.wait_barrier_named %v0, %threads : i32, i32 + ttg.warp_return + } + partition1() num_warps(4) { + %v1 = arith.constant 17 : i32 + %threads = arith.constant 256 : i32 + // CHECK: %[[ID5:.+]] = arith.constant 5 : i32 + // CHECK: ttng.arrive_barrier_named %[[ID5]], {{.*}} : i32, i32 + ttng.arrive_barrier_named %v1, %threads : i32, i32 + ttg.warp_return + } : () -> () + tt.return + } +} + +// ----- + +module attributes {"ttg.num-ctas" = 1 : i32, "ttg.num-warps" = 8 : i32, ttg.target = "cuda:90", "ttg.threads-per-warp" = 32 : i32} { + // CHECK-LABEL: tt.func @preserve_existing_physical_ids_with_warpspec + tt.func @preserve_existing_physical_ids_with_warpspec() { + ttg.warp_specialize() attributes {requestedRegisters = array} + default { + ttg.warp_yield + } + partition0() num_warps(4) { + %physical0 = arith.constant 0 : i32 + %threads = arith.constant 256 : i32 + // CHECK: %[[PHYSICAL0:.+]] = arith.constant 0 : i32 + // CHECK: ttng.wait_barrier_named %[[PHYSICAL0]], {{.*}} : i32, i32 + ttng.wait_barrier_named %physical0, %threads : i32, i32 + ttg.warp_return + } + partition1() num_warps(4) { + ttg.warp_return + } : () -> () + tt.return + } +} diff --git a/third_party/tle/test/GPU/test_tle_allocate_named_barriers_from_tle.mlir b/third_party/tle/test/GPU/test_tle_allocate_named_barriers_from_tle.mlir new file mode 100644 index 0000000000..c1c6c03d2f --- /dev/null +++ b/third_party/tle/test/GPU/test_tle_allocate_named_barriers_from_tle.mlir @@ -0,0 +1,20 @@ +// RUN: triton-opt %s -triton-tle-lower-barriers -triton-tle-allocate-named-barriers -split-input-file | FileCheck %s + +#slot = #ttg.swizzled_shared<{vec = 1, perPhase = 1, maxPhase = 1, order = [0]}> +#smem = #ttg.shared_memory + +module attributes {"ttg.num-ctas" = 1 : i32, "ttg.num-warps" = 4 : i32} { + // CHECK-LABEL: tt.func @lower_and_allocate_tle_virtual_named_barriers + tt.func @lower_and_allocate_tle_virtual_named_barriers(%slot: !ttg.memdesc<1xi64, #slot, #smem, mutable>) { + tle.barrier.wait %slot {backend = "named", named_id = 16 : i32, named_num_threads = 256 : i32} : !ttg.memdesc<1xi64, #slot, #smem, mutable> + tle.barrier.arrive %slot {arrive_count = 1 : i32, backend = "named", named_id = 17 : i32, named_num_threads = 256 : i32} : !ttg.memdesc<1xi64, #slot, #smem, mutable> + tt.return + } +} + +// CHECK-DAG: %[[ID1:.+]] = arith.constant 1 : i32 +// CHECK-DAG: %[[THREADS:.+]] = arith.constant 256 : i32 +// CHECK: ttng.wait_barrier_named %[[ID1]], %[[THREADS]] +// CHECK: %[[ID2:.+]] = arith.constant 2 : i32 +// CHECK: ttng.arrive_barrier_named %[[ID2]], {{.*}} +// CHECK-NOT: tle.barrier diff --git a/third_party/tle/test/GPU/test_tle_lower_barriers.mlir b/third_party/tle/test/GPU/test_tle_lower_barriers.mlir new file mode 100644 index 0000000000..075b7ed6ad --- /dev/null +++ b/third_party/tle/test/GPU/test_tle_lower_barriers.mlir @@ -0,0 +1,68 @@ +// RUN: triton-opt %s -triton-tle-lower-barriers -split-input-file | FileCheck %s + +#shared = #ttg.swizzled_shared<{vec = 1, perPhase = 1, maxPhase = 1, order = [1, 0]}> +#slot = #ttg.swizzled_shared<{vec = 1, perPhase = 1, maxPhase = 1, order = [0]}> +#smem = #ttg.shared_memory + +module attributes {"ttg.num-ctas" = 1 : i32, "ttg.num-warps" = 4 : i32} { + // CHECK-LABEL: tt.func @lower_mbarrier + tt.func @lower_mbarrier() { + %c0 = arith.constant 0 : i32 + %bars = tle.barrier.alloc {arrive_count = 2 : i32, init_polarity = 0 : i32, num_barriers = 2 : i32} : !ttg.memdesc<2x1xi64, #shared, #smem, mutable> + %slot = ttg.memdesc_index %bars[%c0] : !ttg.memdesc<2x1xi64, #shared, #smem, mutable> -> !ttg.memdesc<1xi64, #slot, #smem, mutable> + tle.barrier.wait %slot, %c0 {backend = "mbarrier", named_id = 0 : i32, named_num_threads = 0 : i32} : !ttg.memdesc<1xi64, #slot, #smem, mutable> + tle.barrier.arrive %slot, %c0 {arrive_count = 2 : i32, backend = "mbarrier", named_id = 0 : i32, named_num_threads = 0 : i32} : !ttg.memdesc<1xi64, #slot, #smem, mutable> + tt.return + } +} + +// CHECK: %[[ALLOC:.*]] = ttg.local_alloc +// CHECK: ttng.init_barrier +// CHECK: ttng.init_barrier +// CHECK: ttng.wait_barrier +// CHECK: ttng.arrive_barrier +// CHECK-NOT: tle.barrier + +// ----- + +#shared = #ttg.swizzled_shared<{vec = 1, perPhase = 1, maxPhase = 1, order = [1, 0]}> +#slot = #ttg.swizzled_shared<{vec = 1, perPhase = 1, maxPhase = 1, order = [0]}> +#smem = #ttg.shared_memory + +module attributes {"ttg.num-ctas" = 1 : i32, "ttg.num-warps" = 4 : i32} { + // CHECK-LABEL: tt.func @lower_ready_mbarrier_no_initial_arrive + tt.func @lower_ready_mbarrier_no_initial_arrive() { + %c0 = arith.constant 0 : i32 + %c1 = arith.constant 1 : i32 + %bars = tle.barrier.alloc {arrive_count = 2 : i32, init_polarity = 1 : i32, num_barriers = 1 : i32} : !ttg.memdesc<1x1xi64, #shared, #smem, mutable> + %slot = ttg.memdesc_index %bars[%c0] : !ttg.memdesc<1x1xi64, #shared, #smem, mutable> -> !ttg.memdesc<1xi64, #slot, #smem, mutable> + tle.barrier.wait %slot, %c1 {backend = "mbarrier", named_id = 0 : i32, named_num_threads = 0 : i32} : !ttg.memdesc<1xi64, #slot, #smem, mutable> + tt.return + } +} + +// CHECK: ttng.init_barrier +// CHECK-NOT: ttng.arrive_barrier +// CHECK: ttng.wait_barrier +// CHECK-NOT: tle.barrier + +// ----- + +#shared = #ttg.swizzled_shared<{vec = 1, perPhase = 1, maxPhase = 1, order = [1, 0]}> +#slot = #ttg.swizzled_shared<{vec = 1, perPhase = 1, maxPhase = 1, order = [0]}> +#smem = #ttg.shared_memory + +module attributes {"ttg.num-ctas" = 1 : i32, "ttg.num-warps" = 4 : i32} { + // CHECK-LABEL: tt.func @lower_named_barrier + tt.func @lower_named_barrier(%slot: !ttg.memdesc<1xi64, #slot, #smem, mutable>) { + tle.barrier.wait %slot {backend = "named", named_id = 3 : i32, named_num_threads = 256 : i32} : !ttg.memdesc<1xi64, #slot, #smem, mutable> + tle.barrier.arrive %slot {arrive_count = 1 : i32, backend = "named", named_id = 3 : i32, named_num_threads = 256 : i32} : !ttg.memdesc<1xi64, #slot, #smem, mutable> + tt.return + } +} + +// CHECK: %[[ID:.*]] = arith.constant 3 : i32 +// CHECK: %[[THREADS:.*]] = arith.constant 256 : i32 +// CHECK: ttng.wait_barrier_named %[[ID]], %[[THREADS]] +// CHECK: ttng.arrive_barrier_named +// CHECK-NOT: tle.barrier diff --git a/third_party/tle/test/GPU/test_tle_lower_wgmma.mlir b/third_party/tle/test/GPU/test_tle_lower_wgmma.mlir new file mode 100644 index 0000000000..5f97e0be34 --- /dev/null +++ b/third_party/tle/test/GPU/test_tle_lower_wgmma.mlir @@ -0,0 +1,144 @@ +// RUN: triton-opt %s -split-input-file -triton-tle-lower-wgmma | FileCheck %s + +#blocked = #ttg.blocked<{sizePerThread = [1, 1], threadsPerWarp = [1, 32], warpsPerCTA = [4, 1], order = [1, 0]}> +#shared = #ttg.nvmma_shared<{swizzlingByteWidth = 32, transposed = false, elementBitWidth = 16}> +#smem = #ttg.shared_memory + +module attributes {"ttg.target" = "cuda:90", "ttg.num-ctas" = 1 : i32, "ttg.num-warps" = 4 : i32, "ttg.threads-per-warp" = 32 : i32} { + // CHECK-LABEL: tt.func @lower_single_wgmma + tt.func @lower_single_wgmma( + %a: !ttg.memdesc<64x16xf16, #shared, #smem, mutable>, + %b: !ttg.memdesc<16x16xf16, #shared, #smem, mutable>) { + %zero = arith.constant dense<0.000000e+00> : tensor<64x16xf32, #blocked> + // CHECK: %[[ZERO:.+]] = arith.constant + // CHECK-NEXT: %[[ACC:.+]] = ttg.convert_layout %[[ZERO]] + // CHECK-NEXT: %[[DOT:.+]] = ttng.warp_group_dot %a, %b, %[[ACC]] + // CHECK-NEXT: %[[WAIT:.+]] = ttng.warp_group_dot_wait %[[DOT]] {pendings = 0 : i32} + // CHECK-NEXT: ttg.convert_layout %[[WAIT]] + %dot = tle.wgmma %a, %b, %zero {inputPrecision = 0 : i32, isAsync = true, maxNumImpreciseAcc = 0 : i32} : !ttg.memdesc<64x16xf16, #shared, #smem, mutable> * !ttg.memdesc<16x16xf16, #shared, #smem, mutable>, tensor<64x16xf32, #blocked> -> tensor<64x16xf32, #blocked> + %wait = tle.wgmma_wait %dot {pendings = 0 : i32} : tensor<64x16xf32, #blocked> -> tensor<64x16xf32, #blocked> + tt.return + } +} + +// ----- + +#blocked = #ttg.blocked<{sizePerThread = [1, 1], threadsPerWarp = [1, 32], warpsPerCTA = [4, 1], order = [1, 0]}> +#shared = #ttg.nvmma_shared<{swizzlingByteWidth = 32, transposed = false, elementBitWidth = 16}> +#smem = #ttg.shared_memory + +module attributes {"ttg.target" = "cuda:90", "ttg.num-ctas" = 1 : i32, "ttg.num-warps" = 4 : i32, "ttg.threads-per-warp" = 32 : i32} { + // CHECK-LABEL: tt.func @lower_chained_wgmma + tt.func @lower_chained_wgmma( + %a0: !ttg.memdesc<64x16xf16, #shared, #smem, mutable>, + %b0: !ttg.memdesc<16x16xf16, #shared, #smem, mutable>, + %a1: !ttg.memdesc<64x16xf16, #shared, #smem, mutable>, + %b1: !ttg.memdesc<16x16xf16, #shared, #smem, mutable>) { + %zero = arith.constant dense<0.000000e+00> : tensor<64x16xf32, #blocked> + // CHECK: %[[ZERO:.+]] = arith.constant + // CHECK-NEXT: %[[ACC:.+]] = ttg.convert_layout %[[ZERO]] + // CHECK-NEXT: %[[DOT0:.+]] = ttng.warp_group_dot %a0, %b0, %[[ACC]] + // CHECK-NEXT: %[[DOT1:.+]] = ttng.warp_group_dot %a1, %b1, %[[DOT0]] + // CHECK-NEXT: %[[WAIT:.+]] = ttng.warp_group_dot_wait %[[DOT1]] {pendings = 0 : i32} + // CHECK-NEXT: ttg.convert_layout %[[WAIT]] + %dot0 = tle.wgmma %a0, %b0, %zero {inputPrecision = 0 : i32, isAsync = true, maxNumImpreciseAcc = 0 : i32} : !ttg.memdesc<64x16xf16, #shared, #smem, mutable> * !ttg.memdesc<16x16xf16, #shared, #smem, mutable>, tensor<64x16xf32, #blocked> -> tensor<64x16xf32, #blocked> + %dot1 = tle.wgmma %a1, %b1, %dot0 {inputPrecision = 0 : i32, isAsync = true, maxNumImpreciseAcc = 0 : i32} : !ttg.memdesc<64x16xf16, #shared, #smem, mutable> * !ttg.memdesc<16x16xf16, #shared, #smem, mutable>, tensor<64x16xf32, #blocked> -> tensor<64x16xf32, #blocked> + %wait = tle.wgmma_wait %dot1 {pendings = 0 : i32} : tensor<64x16xf32, #blocked> -> tensor<64x16xf32, #blocked> + tt.return + } +} + +// ----- + +#blocked = #ttg.blocked<{sizePerThread = [1, 1], threadsPerWarp = [1, 32], warpsPerCTA = [4, 1], order = [1, 0]}> +#shared = #ttg.nvmma_shared<{swizzlingByteWidth = 32, transposed = false, elementBitWidth = 16}> +#smem = #ttg.shared_memory + +module attributes {"ttg.target" = "cuda:90", "ttg.num-ctas" = 1 : i32, "ttg.num-warps" = 4 : i32, "ttg.threads-per-warp" = 32 : i32} { + // CHECK-LABEL: tt.func @lower_register_a_wgmma + tt.func @lower_register_a_wgmma( + %a: tensor<64x16xf16, #blocked>, + %b: !ttg.memdesc<16x16xf16, #shared, #smem, mutable>) { + %zero = arith.constant dense<0.000000e+00> : tensor<64x16xf32, #blocked> + // CHECK: %[[ZERO:.+]] = arith.constant + // CHECK-NEXT: %[[ACC:.+]] = ttg.convert_layout %[[ZERO]] + // CHECK-NEXT: %[[A:.+]] = ttg.convert_layout %a + // CHECK-NEXT: %[[DOT:.+]] = ttng.warp_group_dot %[[A]], %b, %[[ACC]] + // CHECK-NEXT: %[[WAIT:.+]] = ttng.warp_group_dot_wait %[[DOT]] {pendings = 0 : i32} + // CHECK-NEXT: ttg.convert_layout %[[WAIT]] + %dot = tle.wgmma %a, %b, %zero {inputPrecision = 0 : i32, isAsync = true, maxNumImpreciseAcc = 0 : i32} : tensor<64x16xf16, #blocked> * !ttg.memdesc<16x16xf16, #shared, #smem, mutable>, tensor<64x16xf32, #blocked> -> tensor<64x16xf32, #blocked> + %wait = tle.wgmma_wait %dot {pendings = 0 : i32} : tensor<64x16xf32, #blocked> -> tensor<64x16xf32, #blocked> + tt.return + } +} + +// ----- + +#blocked = #ttg.blocked<{sizePerThread = [1, 1], threadsPerWarp = [1, 32], warpsPerCTA = [4, 1], order = [1, 0]}> +#shared = #ttg.nvmma_shared<{swizzlingByteWidth = 32, transposed = false, elementBitWidth = 16}> +#smem = #ttg.shared_memory + +module attributes {"ttg.target" = "cuda:90", "ttg.num-ctas" = 1 : i32, "ttg.num-warps" = 4 : i32, "ttg.threads-per-warp" = 32 : i32} { + // CHECK-LABEL: tt.func @lower_loop_carried_wgmma + tt.func @lower_loop_carried_wgmma( + %a0: !ttg.memdesc<64x16xf16, #shared, #smem, mutable>, + %b0: !ttg.memdesc<16x16xf16, #shared, #smem, mutable>, + %a1: !ttg.memdesc<64x16xf16, #shared, #smem, mutable>, + %b1: !ttg.memdesc<16x16xf16, #shared, #smem, mutable>) { + %lb = arith.constant 0 : index + %ub = arith.constant 4 : index + %step = arith.constant 1 : index + %zero = arith.constant dense<0.000000e+00> : tensor<64x16xf32, #blocked> + // CHECK: %[[ACC:.+]] = ttg.convert_layout %{{.+}} + // CHECK-NEXT: %[[DOT0:.+]] = ttng.warp_group_dot %a0, %b0, %[[ACC]] + // CHECK-NEXT: %[[LOOP:.+]] = scf.for {{.*}} iter_args({{.*}} = %[[DOT0]]) + // CHECK: %[[DOT1:.+]] = ttng.warp_group_dot %a1, %b1, %{{.+}} + // CHECK-NEXT: %[[WAIT1:.+]] = ttng.warp_group_dot_wait %[[DOT1]] {pendings = 1 : i32} + // CHECK-NOT: ttg.convert_layout + // CHECK: scf.yield %[[WAIT1]] + // CHECK: %[[WAIT0:.+]] = ttng.warp_group_dot_wait %[[LOOP]] {pendings = 0 : i32} + // CHECK-NEXT: ttg.convert_layout %[[WAIT0]] + %dot0 = tle.wgmma %a0, %b0, %zero {inputPrecision = 0 : i32, isAsync = true, maxNumImpreciseAcc = 0 : i32} : !ttg.memdesc<64x16xf16, #shared, #smem, mutable> * !ttg.memdesc<16x16xf16, #shared, #smem, mutable>, tensor<64x16xf32, #blocked> -> tensor<64x16xf32, #blocked> + %loop = scf.for %i = %lb to %ub step %step iter_args(%acc = %dot0) -> (tensor<64x16xf32, #blocked>) { + %dot1 = tle.wgmma %a1, %b1, %acc {inputPrecision = 0 : i32, isAsync = true, maxNumImpreciseAcc = 0 : i32} : !ttg.memdesc<64x16xf16, #shared, #smem, mutable> * !ttg.memdesc<16x16xf16, #shared, #smem, mutable>, tensor<64x16xf32, #blocked> -> tensor<64x16xf32, #blocked> + %wait1 = tle.wgmma_wait %dot1 {pendings = 1 : i32} : tensor<64x16xf32, #blocked> -> tensor<64x16xf32, #blocked> + scf.yield %wait1 : tensor<64x16xf32, #blocked> + } + %wait0 = tle.wgmma_wait %loop {pendings = 0 : i32} : tensor<64x16xf32, #blocked> -> tensor<64x16xf32, #blocked> + tt.return + } +} + +// ----- + +#blocked = #ttg.blocked<{sizePerThread = [1, 1], threadsPerWarp = [1, 32], warpsPerCTA = [4, 1], order = [1, 0]}> +#shared = #ttg.nvmma_shared<{swizzlingByteWidth = 32, transposed = false, elementBitWidth = 16}> +#smem = #ttg.shared_memory + +module attributes {"ttg.target" = "cuda:90", "ttg.num-ctas" = 1 : i32, "ttg.num-warps" = 4 : i32, "ttg.threads-per-warp" = 32 : i32} { + // CHECK-LABEL: tt.func @lower_wait1_ordinary_use + tt.func @lower_wait1_ordinary_use( + %a0: !ttg.memdesc<64x16xf16, #shared, #smem, mutable>, + %b0: !ttg.memdesc<16x16xf16, #shared, #smem, mutable>, + %a1: !ttg.memdesc<64x16xf16, #shared, #smem, mutable>, + %b1: !ttg.memdesc<16x16xf16, #shared, #smem, mutable>) { + %zero = arith.constant dense<0.000000e+00> : tensor<64x16xf32, #blocked> + // CHECK: %[[ACC0:.+]] = ttg.convert_layout %{{.+}} + // CHECK-NEXT: %[[DOT0:.+]] = ttng.warp_group_dot %a0, %b0, %[[ACC0]] + // CHECK-NEXT: %[[ACC1:.+]] = ttg.convert_layout %{{.+}} + // CHECK-NEXT: %{{.+}} = ttng.warp_group_dot %a1, %b1, %[[ACC1]] + // CHECK-NEXT: %[[WAIT:.+]] = ttng.warp_group_dot_wait %[[DOT0]] {pendings = 1 : i32} + // CHECK-NEXT: %[[RELEASED:.+]] = ttg.convert_layout %[[WAIT]] + // CHECK-NEXT: "tt.reduce"(%[[RELEASED]]) + %dot0 = tle.wgmma %a0, %b0, %zero {inputPrecision = 0 : i32, isAsync = true, maxNumImpreciseAcc = 0 : i32} : !ttg.memdesc<64x16xf16, #shared, #smem, mutable> * !ttg.memdesc<16x16xf16, #shared, #smem, mutable>, tensor<64x16xf32, #blocked> -> tensor<64x16xf32, #blocked> + %dot1 = tle.wgmma %a1, %b1, %zero {inputPrecision = 0 : i32, isAsync = true, maxNumImpreciseAcc = 0 : i32} : !ttg.memdesc<64x16xf16, #shared, #smem, mutable> * !ttg.memdesc<16x16xf16, #shared, #smem, mutable>, tensor<64x16xf32, #blocked> -> tensor<64x16xf32, #blocked> + %wait = tle.wgmma_wait %dot0 {pendings = 1 : i32} : tensor<64x16xf32, #blocked> -> tensor<64x16xf32, #blocked> + %red = "tt.reduce"(%wait) <{axis = 1 : i32}> ({ + ^bb0(%lhs: f32, %rhs: f32): + %max = arith.maxnumf %lhs, %rhs : f32 + tt.reduce.return %max : f32 + }) : (tensor<64x16xf32, #blocked>) -> tensor<64xf32, #ttg.slice<{dim = 1, parent = #blocked}>> + %wait1 = tle.wgmma_wait %dot1 {pendings = 0 : i32} : tensor<64x16xf32, #blocked> -> tensor<64x16xf32, #blocked> + tt.return + } +} diff --git a/third_party/tle/test/GPU/test_tle_lower_wgmma_errors.mlir b/third_party/tle/test/GPU/test_tle_lower_wgmma_errors.mlir new file mode 100644 index 0000000000..d614e69110 --- /dev/null +++ b/third_party/tle/test/GPU/test_tle_lower_wgmma_errors.mlir @@ -0,0 +1,30 @@ +// RUN: triton-opt %s -split-input-file -triton-tle-lower-wgmma -verify-diagnostics + +#blocked = #ttg.blocked<{sizePerThread = [1, 1], threadsPerWarp = [1, 32], warpsPerCTA = [4, 1], order = [1, 0]}> +#shared = #ttg.nvmma_shared<{swizzlingByteWidth = 32, transposed = false, elementBitWidth = 16}> +#smem = #ttg.shared_memory + +module attributes {"ttg.target" = "cuda:90", "ttg.num-ctas" = 1 : i32, "ttg.num-warps" = 4 : i32, "ttg.threads-per-warp" = 32 : i32} { + tt.func @direct_use_without_wait( + %a: !ttg.memdesc<64x16xf16, #shared, #smem, mutable>, + %b: !ttg.memdesc<16x16xf16, #shared, #smem, mutable>, + %out: tensor<64x16x!tt.ptr, #blocked>) { + %zero = arith.constant dense<0.000000e+00> : tensor<64x16xf32, #blocked> + // expected-error @+1 {{async result must be consumed by tle.wgmma_wait before ordinary tensor use}} + %dot = tle.wgmma %a, %b, %zero {inputPrecision = 0 : i32, isAsync = true, maxNumImpreciseAcc = 0 : i32} : !ttg.memdesc<64x16xf16, #shared, #smem, mutable> * !ttg.memdesc<16x16xf16, #shared, #smem, mutable>, tensor<64x16xf32, #blocked> -> tensor<64x16xf32, #blocked> + tt.store %out, %dot : tensor<64x16x!tt.ptr, #blocked> + tt.return + } +} + +// ----- + +#blocked = #ttg.blocked<{sizePerThread = [1, 1], threadsPerWarp = [1, 32], warpsPerCTA = [4, 1], order = [1, 0]}> + +module attributes {"ttg.target" = "cuda:90", "ttg.num-ctas" = 1 : i32, "ttg.num-warps" = 4 : i32, "ttg.threads-per-warp" = 32 : i32} { + tt.func @wait_without_wgmma(%acc: tensor<64x16xf32, #blocked>) { + // expected-error @+1 {{input must be the async result of tle.wgmma}} + %wait = tle.wgmma_wait %acc {pendings = 0 : i32} : tensor<64x16xf32, #blocked> -> tensor<64x16xf32, #blocked> + tt.return + } +} diff --git a/third_party/tle/test/GPU/test_tle_tma_copy_barrier.mlir b/third_party/tle/test/GPU/test_tle_tma_copy_barrier.mlir new file mode 100644 index 0000000000..0da731dce6 --- /dev/null +++ b/third_party/tle/test/GPU/test_tle_tma_copy_barrier.mlir @@ -0,0 +1,55 @@ +// RUN: triton-opt %s -split-input-file -triton-tle-lower-barriers -triton-tle-lower-tma-copy | FileCheck %s + +#nvmma = #ttg.nvmma_shared<{swizzlingByteWidth = 128, transposed = false, elementBitWidth = 32}> +#bar_shared = #ttg.swizzled_shared<{vec = 1, perPhase = 1, maxPhase = 1, order = [1, 0]}> +#bar_slot = #ttg.swizzled_shared<{vec = 1, perPhase = 1, maxPhase = 1, order = [0]}> +#smem = #ttg.shared_memory + +module attributes {"ttg.num-ctas" = 1 : i32, "ttg.num-warps" = 4 : i32, "ttg.threads-per-warp" = 32 : i32, ttg.target = "cuda:90"} { + // CHECK-LABEL: tt.func @explicit_tma_barrier + tt.func @explicit_tma_barrier(%desc: !tt.tensordesc>) { + %c0 = arith.constant 0 : i32 + %dst = ttg.local_alloc : () -> !ttg.memdesc<32x64xf32, #nvmma, #smem, mutable> + // CHECK: %[[BARS:.*]] = ttg.local_alloc : () -> !ttg.memdesc<1x1xi64 + // CHECK: ttng.init_barrier + %bars = tle.barrier.alloc {arrive_count = 1 : i32, expect_bytes = 8192 : i32, init_polarity = 0 : i32, num_barriers = 1 : i32} : !ttg.memdesc<1x1xi64, #bar_shared, #smem, mutable> + %slot = ttg.memdesc_index %bars[%c0] : !ttg.memdesc<1x1xi64, #bar_shared, #smem, mutable> -> !ttg.memdesc<1xi64, #bar_slot, #smem, mutable> + + // CHECK-NOT: ttg.local_alloc : () -> !ttg.memdesc<1xi64 + // CHECK: ttng.barrier_expect %[[SLOT:.*]], 8192 + // CHECK: ttng.async_tma_copy_global_to_local {{.*}} %[[SLOT]], %true + // CHECK-NOT: ttng.inval_barrier %[[SLOT]] + // CHECK: ttng.wait_barrier %[[SLOT]] + ttg.tma_copy %desc, %dst, [%c0, %c0], barrier %slot {expect_bytes = 8192 : i32} : !tt.tensordesc>, !ttg.memdesc<32x64xf32, #nvmma, #smem, mutable>, !ttg.memdesc<1xi64, #bar_slot, #smem, mutable> + tle.barrier.wait %slot, %c0 {backend = "mbarrier", named_id = 0 : i32, named_num_threads = 0 : i32} : !ttg.memdesc<1xi64, #bar_slot, #smem, mutable> + tt.return + } +} + +// ----- + +#nvmma = #ttg.nvmma_shared<{swizzlingByteWidth = 128, transposed = false, elementBitWidth = 32}> +#bar_shared = #ttg.swizzled_shared<{vec = 1, perPhase = 1, maxPhase = 1, order = [1, 0]}> +#bar_slot = #ttg.swizzled_shared<{vec = 1, perPhase = 1, maxPhase = 1, order = [0]}> +#smem = #ttg.shared_memory + +module attributes {"ttg.num-ctas" = 1 : i32, "ttg.num-warps" = 4 : i32, "ttg.threads-per-warp" = 32 : i32, ttg.target = "cuda:90"} { + // CHECK-LABEL: tt.func @explicit_tma_barrier_two_slots + tt.func @explicit_tma_barrier_two_slots(%a_desc: !tt.tensordesc>, %b_desc: !tt.tensordesc>) { + %c0 = arith.constant 0 : i32 + %c1 = arith.constant 1 : i32 + %a = ttg.local_alloc : () -> !ttg.memdesc<32x64xf32, #nvmma, #smem, mutable> + %b = ttg.local_alloc : () -> !ttg.memdesc<32x64xf32, #nvmma, #smem, mutable> + %bars = tle.barrier.alloc {arrive_count = 1 : i32, expect_bytes = 8192 : i32, init_polarity = 0 : i32, num_barriers = 2 : i32} : !ttg.memdesc<2x1xi64, #bar_shared, #smem, mutable> + %slot0 = ttg.memdesc_index %bars[%c0] : !ttg.memdesc<2x1xi64, #bar_shared, #smem, mutable> -> !ttg.memdesc<1xi64, #bar_slot, #smem, mutable> + %slot1 = ttg.memdesc_index %bars[%c1] : !ttg.memdesc<2x1xi64, #bar_shared, #smem, mutable> -> !ttg.memdesc<1xi64, #bar_slot, #smem, mutable> + + // CHECK: ttng.barrier_expect %[[SLOT0:.*]], 8192 + // CHECK: ttng.async_tma_copy_global_to_local {{.*}} %[[SLOT0]], %true + ttg.tma_copy %a_desc, %a, [%c0, %c0], barrier %slot0 {expect_bytes = 8192 : i32} : !tt.tensordesc>, !ttg.memdesc<32x64xf32, #nvmma, #smem, mutable>, !ttg.memdesc<1xi64, #bar_slot, #smem, mutable> + // CHECK: ttng.barrier_expect %[[SLOT1:.*]], 8192 + // CHECK: ttng.async_tma_copy_global_to_local {{.*}} %[[SLOT1]], %true + ttg.tma_copy %b_desc, %b, [%c0, %c0], barrier %slot1 {expect_bytes = 8192 : i32} : !tt.tensordesc>, !ttg.memdesc<32x64xf32, #nvmma, #smem, mutable>, !ttg.memdesc<1xi64, #bar_slot, #smem, mutable> + tt.return + } +} diff --git a/third_party/tle/test/GPU/test_tle_tma_copy_barrier_errors.mlir b/third_party/tle/test/GPU/test_tle_tma_copy_barrier_errors.mlir new file mode 100644 index 0000000000..f9698278c1 --- /dev/null +++ b/third_party/tle/test/GPU/test_tle_tma_copy_barrier_errors.mlir @@ -0,0 +1,33 @@ +// RUN: triton-opt %s -split-input-file -triton-tle-lower-tma-copy -verify-diagnostics + +#nvmma = #ttg.nvmma_shared<{swizzlingByteWidth = 128, transposed = false, elementBitWidth = 32}> +#bar_shared = #ttg.swizzled_shared<{vec = 1, perPhase = 1, maxPhase = 1, order = [1, 0]}> +#bar_slot = #ttg.swizzled_shared<{vec = 1, perPhase = 1, maxPhase = 1, order = [0]}> +#smem = #ttg.shared_memory + +module attributes {"ttg.num-ctas" = 1 : i32, "ttg.num-warps" = 4 : i32, "ttg.threads-per-warp" = 32 : i32, ttg.target = "cuda:90"} { + tt.func @reject_tma_store_completion_barrier(%desc: !tt.tensordesc>, %bar: !ttg.memdesc<1xi64, #bar_slot, #smem, mutable>) { + %c0 = arith.constant 0 : i32 + %src = ttg.local_alloc : () -> !ttg.memdesc<32x64xf32, #nvmma, #smem, mutable> + // expected-error @+1 {{barrier is only supported for global-to-shared TMA copy}} + ttg.tma_copy %src, %desc, [%c0, %c0], barrier %bar {expect_bytes = 8192 : i32} : !ttg.memdesc<32x64xf32, #nvmma, #smem, mutable>, !tt.tensordesc>, !ttg.memdesc<1xi64, #bar_slot, #smem, mutable> + tt.return + } +} + +// ----- + +#nvmma = #ttg.nvmma_shared<{swizzlingByteWidth = 128, transposed = false, elementBitWidth = 32}> +#bar_shared = #ttg.swizzled_shared<{vec = 1, perPhase = 1, maxPhase = 1, order = [1, 0]}> +#bar_slot = #ttg.swizzled_shared<{vec = 1, perPhase = 1, maxPhase = 1, order = [0]}> +#smem = #ttg.shared_memory + +module attributes {"ttg.num-ctas" = 1 : i32, "ttg.num-warps" = 4 : i32, "ttg.threads-per-warp" = 32 : i32, ttg.target = "cuda:90"} { + tt.func @reject_missing_expect_bytes(%desc: !tt.tensordesc>, %bar: !ttg.memdesc<1xi64, #bar_slot, #smem, mutable>) { + %c0 = arith.constant 0 : i32 + %dst = ttg.local_alloc : () -> !ttg.memdesc<32x64xf32, #nvmma, #smem, mutable> + // expected-error @+1 {{with explicit completion barrier requires positive expect_bytes}} + ttg.tma_copy %desc, %dst, [%c0, %c0], barrier %bar : !tt.tensordesc>, !ttg.memdesc<32x64xf32, #nvmma, #smem, mutable>, !ttg.memdesc<1xi64, #bar_slot, #smem, mutable> + tt.return + } +} diff --git a/third_party/tle/test/GPU/test_tle_wgmma_pipeline_accumulator_chain.mlir b/third_party/tle/test/GPU/test_tle_wgmma_pipeline_accumulator_chain.mlir index 5a482522d7..cb215f4f8b 100644 --- a/third_party/tle/test/GPU/test_tle_wgmma_pipeline_accumulator_chain.mlir +++ b/third_party/tle/test/GPU/test_tle_wgmma_pipeline_accumulator_chain.mlir @@ -89,6 +89,82 @@ module attributes {"ttg.target" = "cuda:90", "ttg.num-ctas" = 1 : i32, "ttg.num- #shared1 = #ttg.nvmma_shared<{swizzlingByteWidth = 128, transposed = true, elementBitWidth = 16}> #smem = #ttg.shared_memory +module attributes {"ttg.target" = "cuda:90", "ttg.num-ctas" = 1 : i32, "ttg.num-warps" = 4 : i32, "ttg.threads-per-warp" = 32 : i32} { + // CHECK-LABEL: tt.func @carry_pending_wgmma_to_loop_result_wait0 + tt.func @carry_pending_wgmma_to_loop_result_wait0( + %a: !ttg.memdesc<64x64xbf16, #shared, #smem>, + %b: !ttg.memdesc<64x64xbf16, #shared1, #smem, mutable>, + %out: tensor<64x64x!tt.ptr, #mma>) { + %c0 = arith.constant 0 : index + %c1 = arith.constant 1 : index + %c8 = arith.constant 8 : index + %zero = arith.constant dense<0.000000e+00> : tensor<64x64xf32, #mma> + %b_view = tle.memdesc_wgmma_view %b {order = array} : !ttg.memdesc<64x64xbf16, #shared1, #smem, mutable> -> !ttg.memdesc<64x64xbf16, #shared1, #smem> + // CHECK: %[[RES:.+]] = scf.for + %res = scf.for %iv = %c0 to %c8 step %c1 iter_args(%acc = %zero) -> (tensor<64x64xf32, #mma>) { + // CHECK: %[[DOT:.+]] = ttng.warp_group_dot + %dot = ttng.warp_group_dot %a, %b_view, %acc {inputPrecision = 0 : i32, isAsync = true} : !ttg.memdesc<64x64xbf16, #shared, #smem> * !ttg.memdesc<64x64xbf16, #shared1, #smem> -> tensor<64x64xf32, #mma> + // CHECK-NEXT: ttng.warp_group_dot_commit + // CHECK-NEXT: %[[WAIT1:.+]] = ttng.warp_group_dot_wait %[[DOT]] + // CHECK-SAME: {pendings = 1 : i32} + %wait1 = ttng.warp_group_dot_wait %dot {pendings = 1 : i32} : tensor<64x64xf32, #mma> + // CHECK-NOT: ttng.warp_group_dot_wait + // CHECK: scf.yield %[[WAIT1]] + scf.yield %wait1 : tensor<64x64xf32, #mma> + } + // CHECK: %[[WAIT0:.+]] = ttng.warp_group_dot_wait %[[RES]] + // CHECK-SAME: {pendings = 0 : i32} + %wait0 = ttng.warp_group_dot_wait %res {pendings = 0 : i32} : tensor<64x64xf32, #mma> + // CHECK: tt.store %{{.*}}, %[[WAIT0]] + tt.store %out, %wait0 : tensor<64x64x!tt.ptr, #mma> + tt.return + } +} + +// ----- + +#mma = #ttg.nvidia_mma<{versionMajor = 3, versionMinor = 0, warpsPerCTA = [4, 1], instrShape = [16, 64, 16]}> +#shared = #ttg.nvmma_shared<{swizzlingByteWidth = 128, transposed = false, elementBitWidth = 16}> +#shared1 = #ttg.nvmma_shared<{swizzlingByteWidth = 128, transposed = true, elementBitWidth = 16}> +#smem = #ttg.shared_memory + +module attributes {"ttg.target" = "cuda:90", "ttg.num-ctas" = 1 : i32, "ttg.num-warps" = 4 : i32, "ttg.threads-per-warp" = 32 : i32} { + // CHECK-LABEL: tt.func @pending_loop_result_without_wait0_drains_before_yield + tt.func @pending_loop_result_without_wait0_drains_before_yield( + %a: !ttg.memdesc<64x64xbf16, #shared, #smem>, + %b: !ttg.memdesc<64x64xbf16, #shared1, #smem, mutable>, + %out: tensor<64x64x!tt.ptr, #mma>) { + %c0 = arith.constant 0 : index + %c1 = arith.constant 1 : index + %c8 = arith.constant 8 : index + %zero = arith.constant dense<0.000000e+00> : tensor<64x64xf32, #mma> + %b_view = tle.memdesc_wgmma_view %b {order = array} : !ttg.memdesc<64x64xbf16, #shared1, #smem, mutable> -> !ttg.memdesc<64x64xbf16, #shared1, #smem> + // CHECK: %[[RES:.+]] = scf.for + %res = scf.for %iv = %c0 to %c8 step %c1 iter_args(%acc = %zero) -> (tensor<64x64xf32, #mma>) { + // CHECK: %[[DOT:.+]] = ttng.warp_group_dot + %dot = ttng.warp_group_dot %a, %b_view, %acc {inputPrecision = 0 : i32, isAsync = true} : !ttg.memdesc<64x64xbf16, #shared, #smem> * !ttg.memdesc<64x64xbf16, #shared1, #smem> -> tensor<64x64xf32, #mma> + // CHECK-NEXT: ttng.warp_group_dot_commit + // CHECK-NEXT: %[[WAIT1:.+]] = ttng.warp_group_dot_wait %[[DOT]] + // CHECK-SAME: {pendings = 1 : i32} + %wait1 = ttng.warp_group_dot_wait %dot {pendings = 1 : i32} : tensor<64x64xf32, #mma> + // CHECK: %[[YIELD_WAIT:.+]]:{{.*}} = ttng.warp_group_dot_wait %[[WAIT1]] + // CHECK-SAME: {pendings = 0 : i32} + // CHECK: scf.yield %[[YIELD_WAIT]]#0 + scf.yield %wait1 : tensor<64x64xf32, #mma> + } + // CHECK: tt.store %{{.*}}, %[[RES]] + tt.store %out, %res : tensor<64x64x!tt.ptr, #mma> + tt.return + } +} + +// ----- + +#mma = #ttg.nvidia_mma<{versionMajor = 3, versionMinor = 0, warpsPerCTA = [4, 1], instrShape = [16, 64, 16]}> +#shared = #ttg.nvmma_shared<{swizzlingByteWidth = 128, transposed = false, elementBitWidth = 16}> +#shared1 = #ttg.nvmma_shared<{swizzlingByteWidth = 128, transposed = true, elementBitWidth = 16}> +#smem = #ttg.shared_memory + module attributes {"ttg.target" = "cuda:90", "ttg.num-ctas" = 1 : i32, "ttg.num-warps" = 4 : i32, "ttg.threads-per-warp" = 32 : i32} { // CHECK-LABEL: tt.func @defer_wait_for_three_dot_wgmma_c_chain tt.func @defer_wait_for_three_dot_wgmma_c_chain( diff --git a/third_party/tle/test/GPU/test_tle_wgmma_user_promise_pipeline.mlir b/third_party/tle/test/GPU/test_tle_wgmma_user_promise_pipeline.mlir new file mode 100644 index 0000000000..895b208681 --- /dev/null +++ b/third_party/tle/test/GPU/test_tle_wgmma_user_promise_pipeline.mlir @@ -0,0 +1,68 @@ +// RUN: triton-opt %s -split-input-file -tritongpu-pipeline -canonicalize | FileCheck %s + +#mma = #ttg.nvidia_mma<{versionMajor = 3, versionMinor = 0, warpsPerCTA = [4, 1], instrShape = [16, 64, 16]}> +#shared = #ttg.nvmma_shared<{swizzlingByteWidth = 128, transposed = false, elementBitWidth = 16}> +#shared1 = #ttg.nvmma_shared<{swizzlingByteWidth = 128, transposed = true, elementBitWidth = 16}> +#smem = #ttg.shared_memory + +module attributes {"ttg.target" = "cuda:90", "ttg.num-ctas" = 1 : i32, "ttg.num-warps" = 4 : i32, "ttg.threads-per-warp" = 32 : i32, "tle.wgmma_pipeline_mode" = "user_promise"} { + // CHECK-LABEL: tt.func @explicit_wait_is_preserved_without_extra_drain + tt.func @explicit_wait_is_preserved_without_extra_drain( + %a: !ttg.memdesc<64x64xbf16, #shared, #smem>, + %b: !ttg.memdesc<64x64xbf16, #shared1, #smem>, + %out: tensor<64x64x!tt.ptr, #mma>) { + %c0 = arith.constant 0 : index + %c1 = arith.constant 1 : index + %c8 = arith.constant 8 : index + %zero = arith.constant dense<0.000000e+00> : tensor<64x64xf32, #mma> + // CHECK: %[[RES:.+]] = scf.for + %res = scf.for %iv = %c0 to %c8 step %c1 iter_args(%acc = %zero) -> (tensor<64x64xf32, #mma>) { + // CHECK: %[[DOT:.+]] = ttng.warp_group_dot + // CHECK-SAME: {inputPrecision = 0 : i32, isAsync = true, tle.explicit_wgmma_commit} + %dot = ttng.warp_group_dot %a, %b, %acc {inputPrecision = 0 : i32} : !ttg.memdesc<64x64xbf16, #shared, #smem> * !ttg.memdesc<64x64xbf16, #shared1, #smem> -> tensor<64x64xf32, #mma> + // CHECK-NEXT: ttng.warp_group_dot_commit + // CHECK-NEXT: %[[WAIT1:.+]] = ttng.warp_group_dot_wait %[[DOT]] + // CHECK-SAME: {pendings = 1 : i32} + %wait1 = ttng.warp_group_dot_wait %dot {pendings = 1 : i32} : tensor<64x64xf32, #mma> + // CHECK-NOT: ttng.warp_group_dot_wait + // CHECK: scf.yield %[[WAIT1]] + scf.yield %wait1 : tensor<64x64xf32, #mma> + } + // CHECK: %[[WAIT0:.+]] = ttng.warp_group_dot_wait %[[RES]] + // CHECK-SAME: {pendings = 0 : i32} + %wait0 = ttng.warp_group_dot_wait %res {pendings = 0 : i32} : tensor<64x64xf32, #mma> + // CHECK: tt.store %{{.*}}, %[[WAIT0]] + tt.store %out, %wait0 : tensor<64x64x!tt.ptr, #mma> + tt.return + } +} + +// ----- + +#mma = #ttg.nvidia_mma<{versionMajor = 3, versionMinor = 0, warpsPerCTA = [4, 1], instrShape = [16, 64, 16]}> +#shared = #ttg.nvmma_shared<{swizzlingByteWidth = 128, transposed = false, elementBitWidth = 16}> +#shared1 = #ttg.nvmma_shared<{swizzlingByteWidth = 128, transposed = true, elementBitWidth = 16}> +#smem = #ttg.shared_memory + +module attributes {"ttg.target" = "cuda:90", "ttg.num-ctas" = 1 : i32, "ttg.num-warps" = 4 : i32, "ttg.threads-per-warp" = 32 : i32, "tle.wgmma_pipeline_mode" = "user_promise"} { + // CHECK-LABEL: tt.func @ordinary_use_gets_no_auto_wait + tt.func @ordinary_use_gets_no_auto_wait( + %a: !ttg.memdesc<64x64xbf16, #shared, #smem>, + %b: !ttg.memdesc<64x64xbf16, #shared1, #smem>, + %out: tensor<64x64x!tt.ptr, #mma>) { + %c0 = arith.constant 0 : index + %c1 = arith.constant 1 : index + %c8 = arith.constant 8 : index + %zero = arith.constant dense<0.000000e+00> : tensor<64x64xf32, #mma> + scf.for %iv = %c0 to %c8 step %c1 { + // CHECK: %[[DOT:.+]] = ttng.warp_group_dot + // CHECK-SAME: {inputPrecision = 0 : i32, isAsync = true, tle.explicit_wgmma_commit} + %dot = ttng.warp_group_dot %a, %b, %zero {inputPrecision = 0 : i32} : !ttg.memdesc<64x64xbf16, #shared, #smem> * !ttg.memdesc<64x64xbf16, #shared1, #smem> -> tensor<64x64xf32, #mma> + // CHECK-NEXT: ttng.warp_group_dot_commit + // CHECK-NOT: ttng.warp_group_dot_wait + // CHECK: tt.store %{{.*}}, %[[DOT]] + tt.store %out, %dot : tensor<64x64x!tt.ptr, #mma> + } + tt.return + } +} diff --git a/third_party/tle/triton_tle.cc b/third_party/tle/triton_tle.cc index 6de9b1056b..7709e90b24 100644 --- a/third_party/tle/triton_tle.cc +++ b/third_party/tle/triton_tle.cc @@ -154,6 +154,51 @@ void init_triton_tle_ir(py::module &&m) { context, blockM, blockN, colStride, CTASplitM, CTASplitN, /*twoCTAs=*/false)); }) + .def("make_nv_mma_encoding_attr", + [](TritonOpBuilder &self, Value opndA, Value opndAcc, + unsigned versionMajor, unsigned versionMinor, + unsigned moduleNumWarps) { + auto context = self.getBuilder().getContext(); + auto dtypeA = + cast(opndA.getType()).getElementType(); + auto retType = cast(opndAcc.getType()); + Operation *parentOp = + self.getBuilder().getInsertionBlock()->getParentOp(); + unsigned numWarps = + ttg::maybeLookupNumWarps(parentOp).value_or(moduleNumWarps); + auto instrShape = mmaVersionToInstrShape( + versionMajor, retType.getShape(), dtypeA, numWarps); + + // Match the current Hopper WGMMA lowering convention: partition + // the accumulator rows across the warp group. + SmallVector warpsPerCTA = {numWarps, 1}; + SmallVector CTAsPerCGA = {1, 1}; + SmallVector CTASplitNum = {1, 1}; + SmallVector CTAOrder = {1, 0}; + auto CTALayout = ttg::CTAEncodingAttr::fromSplitParams( + context, CTAsPerCGA, CTASplitNum, CTAOrder); + return mlir::cast(ttg::NvidiaMmaEncodingAttr::get( + context, versionMajor, versionMinor, warpsPerCTA, CTALayout, + instrShape)); + }) + .def("make_dot_operand_encoding_attr", + [](TritonOpBuilder &self, Value opnd, unsigned opIdx, + Attribute parentEnc) -> Attribute { + auto context = self.getBuilder().getContext(); + auto eltType = + cast(opnd.getType()).getElementType(); + return ttg::DotOperandEncodingAttr::get(context, opIdx, parentEnc, + eltType); + }) + .def("get_block_ty_with_encoding", + [](TritonOpBuilder &self, Type &elementType, + std::vector &shape, Attribute &encoding) -> Type { + return RankedTensorType::get(shape, elementType, encoding); + }) + .def("create_convert_layout", + [](TritonOpBuilder &self, Type resultTy, Value value) -> Value { + return self.create(resultTy, value); + }) .def("create_local_alloc", [](TritonOpBuilder &self, std::vector shape, Type &elementType, Attribute &encoding) -> mlir::Value { @@ -171,9 +216,37 @@ void init_triton_tle_ir(py::module &&m) { .def("create_tma_copy", [](TritonOpBuilder &self, Value src, Value dst, std::vector &indices) { +#ifdef __HCU__ self.create(src, dst, indices); +#else + self.create(src, dst, indices, Value(), + IntegerAttr()); +#endif return; }) + .def( + "create_tma_copy", + [](TritonOpBuilder &self, Value src, Value dst, + std::vector &indices, py::object barrier, + int32_t expectBytes) { +#ifdef __HCU__ + if (!barrier.is_none() || expectBytes > 0) + throw py::value_error( + "TMA completion barrier is only supported on NVIDIA backend"); + self.create(src, dst, indices); +#else + auto &builder = self.getBuilder(); + Value barrierValue; + if (!barrier.is_none()) + barrierValue = py::cast(barrier); + IntegerAttr expectBytesAttr; + if (expectBytes > 0) + expectBytesAttr = builder.getI32IntegerAttr(expectBytes); + self.create(src, dst, indices, barrierValue, + expectBytesAttr); +#endif + return; + }) .def("create_local_load", [](TritonOpBuilder &self, Type resultTy, Value memDesc) -> Value { return self.create(resultTy, memDesc); @@ -182,6 +255,41 @@ void init_triton_tle_ir(py::module &&m) { [](TritonOpBuilder &self, Value &dst, Value ®Values) -> void { self.create(regValues, dst); }) + .def("create_tle_wgmma", + [](TritonOpBuilder &self, mlir::Value &a, mlir::Value &b, + mlir::Value &c, triton::InputPrecision inputPrecision, + int maxNumImpreciseAcc, bool isAsync) -> mlir::Value { + return self.create(c.getType(), a, b, c, + inputPrecision, + maxNumImpreciseAcc, isAsync); + }) + .def("create_tle_wgmma_wait", + [](TritonOpBuilder &self, mlir::Value &input, + unsigned pendings) -> mlir::Value { + auto pendingsAttr = self.getBuilder().getI32IntegerAttr(pendings); + return self + .create(input.getType(), input, pendingsAttr) + .getOutput(); + }) + .def("create_warp_group_dot", + [](TritonOpBuilder &self, mlir::Value &a, mlir::Value &b, + mlir::Value &c, triton::InputPrecision inputPrecision, + int maxNumImpreciseAcc, bool isAsync) -> mlir::Value { + return self.create( + c.getType(), a, b, c, Value(), inputPrecision, + maxNumImpreciseAcc, isAsync); + }) + .def("create_warp_group_dot_wait", + [](TritonOpBuilder &self, std::vector inputs, + unsigned pendings) -> std::vector { + auto waitOp = + self.create(inputs, pendings); + std::vector outputs; + outputs.reserve(waitOp->getNumResults()); + for (Value result : waitOp->getResults()) + outputs.push_back(result); + return outputs; + }) .def("create_local_pointers", [](TritonOpBuilder &self, Type resultTy, Value memDesc, py::args args) -> OpState { @@ -198,6 +306,62 @@ void init_triton_tle_ir(py::module &&m) { Value index) -> Value { return self.create(resultType, src, index); }) + .def("create_memdesc_trans", + [](TritonOpBuilder &self, Value src, + std::vector &order) -> Value { + return self.create(src, order); + }) + .def("create_barrier_alloc", + [](TritonOpBuilder &self, Type resultType, int32_t numBarriers, + int32_t arriveCount, int32_t initPolarity, + int32_t expectBytes) -> Value { + auto &builder = self.getBuilder(); + IntegerAttr expectBytesAttr; + if (expectBytes > 0) + expectBytesAttr = builder.getI32IntegerAttr(expectBytes); + return self.create( + resultType, builder.getI32IntegerAttr(numBarriers), + builder.getI32IntegerAttr(arriveCount), + builder.getI32IntegerAttr(initPolarity), expectBytesAttr); + }) + .def("create_barrier_wait_mbarrier", + [](TritonOpBuilder &self, Value barrier, Value phase) -> void { + auto &builder = self.getBuilder(); + self.create( + barrier, phase, builder.getStringAttr("mbarrier"), + builder.getI32IntegerAttr(0), builder.getI32IntegerAttr(0)); + }) + .def("create_barrier_wait_named", + [](TritonOpBuilder &self, Value barrier, int32_t namedId, + int32_t numThreads) -> void { + auto &builder = self.getBuilder(); + self.create( + barrier, Value(), builder.getStringAttr("named"), + builder.getI32IntegerAttr(namedId), + builder.getI32IntegerAttr(numThreads)); + }) + .def("create_barrier_arrive_mbarrier", + [](TritonOpBuilder &self, Value barrier, int32_t arriveCount, + py::object phase) -> void { + auto &builder = self.getBuilder(); + Value phaseValue; + if (!phase.is_none()) + phaseValue = py::cast(phase); + self.create( + barrier, phaseValue, builder.getStringAttr("mbarrier"), + builder.getI32IntegerAttr(arriveCount), + builder.getI32IntegerAttr(0), builder.getI32IntegerAttr(0)); + }) + .def("create_barrier_arrive_named", + [](TritonOpBuilder &self, Value barrier, int32_t namedId, + int32_t numThreads) -> void { + auto &builder = self.getBuilder(); + self.create( + barrier, Value(), builder.getStringAttr("named"), + builder.getI32IntegerAttr(1), + builder.getI32IntegerAttr(namedId), + builder.getI32IntegerAttr(numThreads)); + }) .def("create_memdesc_subslice", [](TritonOpBuilder &self, Type resultType, Value src, std::vector &offsets) -> Value { @@ -497,8 +661,12 @@ void init_triton_tle_passes(py::module &&m) { tle::createTritonTleLowerExclusiveCumsum); ADD_PASS_WRAPPER_0("add_lower_async_load", tle::createTritonTleLowerAsyncLoad); + ADD_PASS_WRAPPER_0("add_lower_wgmma", tle::createTritonTleLowerWGMMA); ADD_PASS_WRAPPER_0("add_lower_pipe_to_nvws", tle::createTritonTleLowerPipeToNvws); + ADD_PASS_WRAPPER_0("add_lower_barriers", tle::createTritonTleLowerBarriers); + ADD_PASS_WRAPPER_0("add_allocate_named_barriers", + tle::createTritonTleAllocateNamedBarriers); ADD_PASS_WRAPPER_0("add_lower_tma_copy", tle::createTritonTleLowerTmaCopy); ADD_PASS_WRAPPER_0("add_schedule_tma_store_sync", tle::createTritonTleScheduleTmaStoreSync); diff --git a/third_party/tle/tutorials/test/test_fa.sh b/third_party/tle/tutorials/test/test_fa.sh new file mode 100755 index 0000000000..dce91556f7 --- /dev/null +++ b/third_party/tle/tutorials/test/test_fa.sh @@ -0,0 +1,25 @@ +#!/usr/bin/env sh +set -eu + +SCRIPT_DIR=$(CDPATH= cd -- "$(dirname -- "$0")" && pwd) +PARENT_DIR=$(CDPATH= cd -- "$(dirname -- "$0")/.." && pwd) + +echo "$SCRIPT_DIR" + +python "${PARENT_DIR}/tle_hopper_fa_ws_pipelined_pingpong_persistent.py" \ + --warmup 25 \ + --rep 100 \ + --block-m 128 \ + --block-n 128 \ + --cuda-graph \ + --out "${SCRIPT_DIR}/tle_fa_user_promise_benchmark.csv" \ + --problem 4x32x1024x128 \ + --problem 4x32x2048x128 \ + --problem 4x32x4096x128 \ + --problem 4x32x8192x128 \ + --check \ + --include-sdpa \ + --sdpa-requires-grad \ + --sm-scale 1.3 \ + --dump-summary \ + "$@" diff --git a/third_party/tle/tutorials/test/test_gemm.sh b/third_party/tle/tutorials/test/test_gemm.sh new file mode 100755 index 0000000000..101fd63bb3 --- /dev/null +++ b/third_party/tle/tutorials/test/test_gemm.sh @@ -0,0 +1,23 @@ +#!/usr/bin/env sh +set -eu + +SCRIPT_DIR=$(CDPATH= cd -- "$(dirname -- "$0")" && pwd) +PARENT_DIR=$(CDPATH= cd -- "$(dirname -- "$0")/.." && pwd) + +echo "$SCRIPT_DIR" + +python "${PARENT_DIR}/tle_hopper_gemm_ws_persistent.py" \ + --compare \ + --warmup 25 \ + --rep 100 \ + --bm 128 \ + --bn 128 \ + --bk 64 \ + --cuda-graph \ + --out "${SCRIPT_DIR}/tle_gemm_user_promise_benchmark.csv" \ + --shape 2048x2048x2048 \ + --shape 4096x4096x4096 \ + --shape 8192x8192x512 \ + --shape 8192x8192x8192 \ + --check \ + "$@" diff --git a/third_party/tle/tutorials/tle_hopper_fa_ws_pipelined_pingpong_persistent.py b/third_party/tle/tutorials/tle_hopper_fa_ws_pipelined_pingpong_persistent.py new file mode 100755 index 0000000000..dc7298d81f --- /dev/null +++ b/third_party/tle/tutorials/tle_hopper_fa_ws_pipelined_pingpong_persistent.py @@ -0,0 +1,816 @@ +#!/usr/bin/env python3 +"""Experimental TLE FA3-style forward attention. + +* one producer partition issues TMA loads for Q/K/V into shared memory +* two explicit consumer partitions split the query block along M +* mbarriers coordinate producer/consumer ownership of TMA buffers +* named barriers serialize the two consumer WGMMA issue streams in a ping-pong pattern +* K uses WGMMA descriptor transposition via ``wgmma(..., trans_b=True)`` +* PV uses a register/tensor A operand, avoiding an intermediate P shared tile +""" + +from __future__ import annotations + +import argparse +import csv +from dataclasses import dataclass +from typing import Callable, Iterable, Optional + +import torch +import triton +import triton.language as tl +import triton.experimental.tle.language as tle +from triton.tools.tensor_descriptor import TensorDescriptor + +DEVICE = triton.runtime.driver.active.get_active_torch_device() + + +def alloc_fn(size: int, align: int, stream: Optional[int]): + return torch.empty(size, dtype=torch.int8, device=DEVICE) + + +@triton.jit +def _buf_phase(count, num_buffers: tl.constexpr): + buf = count % num_buffers + phase_idx = count // num_buffers + return buf, phase_idx + + +@triton.jit +def _compute_offsets(tile_idx, H, N_CTX, BLOCK_M: tl.constexpr): + num_pid_m = tl.cdiv(N_CTX, BLOCK_M) + off_hz = tile_idx // num_pid_m + start_m = tile_idx % num_pid_m + off_z = off_hz // H + off_h = off_hz % H + offset_y = off_z * (N_CTX * H) + off_h * N_CTX + qo_offset_y = offset_y + start_m * BLOCK_M + kv_offset_y = offset_y + return start_m, off_hz, qo_offset_y, kv_offset_y + + +@triton.jit +def _attn_fwd_tle_ws_pipelined_pingpong_producer( + Z, + H, + desc_q, + desc_k, + desc_v, + N_CTX, + q_smem, + k_smem, + v_smem, + q_empties, + q_fulls, + k_empties, + k_fulls, + v_empties, + v_fulls, + HEAD_DIM: tl.constexpr, + BLOCK_M: tl.constexpr, + BLOCK_N: tl.constexpr, + NUM_BUFFERS_Q: tl.constexpr, + NUM_BUFFERS_KV: tl.constexpr, + BM_SPLIT: tl.constexpr, + CID: tl.constexpr, +): + prog_id = tl.program_id(0) + num_progs = tl.num_programs(0) + num_pid_m = tl.cdiv(N_CTX, BLOCK_M) + total_tiles = num_pid_m * Z * H + + tile_idx = prog_id + tile_count = 0 + accum_cnt_kv = 0 + while tile_idx < total_tiles: + start_m, off_hz, qo_offset_y, kv_offset_y = _compute_offsets(tile_idx, H, N_CTX, BLOCK_M) + + q_buf, q_phase_idx = _buf_phase(tile_count, NUM_BUFFERS_Q) + q0_idx = q_buf + q1_idx = q_buf + NUM_BUFFERS_Q + + tle.gpu.barrier_wait(q_empties[q0_idx], phaseIdx=q_phase_idx) + tle.gpu.copy( + desc_q, + q_smem.slot(q0_idx), + [BM_SPLIT, HEAD_DIM], + [qo_offset_y, 0], + barrier=q_fulls[q0_idx], + ) + + kv_buf, kv_phase_idx = _buf_phase(accum_cnt_kv, NUM_BUFFERS_KV) + tle.gpu.barrier_wait(k_empties[kv_buf], phaseIdx=kv_phase_idx) + tle.gpu.copy( + desc_k, + k_smem.slot(kv_buf), + [BLOCK_N, HEAD_DIM], + [kv_offset_y, 0], + barrier=k_fulls[kv_buf], + ) + + tle.gpu.barrier_wait(q_empties[q1_idx], phaseIdx=q_phase_idx) + tle.gpu.copy( + desc_q, + q_smem.slot(q1_idx), + [BM_SPLIT, HEAD_DIM], + [qo_offset_y + BM_SPLIT, 0], + barrier=q_fulls[q1_idx], + ) + + tle.gpu.barrier_wait(v_empties[kv_buf], phaseIdx=kv_phase_idx) + tle.gpu.copy( + desc_v, + v_smem.slot(kv_buf), + [BLOCK_N, HEAD_DIM], + [kv_offset_y, 0], + barrier=v_fulls[kv_buf], + ) + accum_cnt_kv += 1 + + for kv_idx in range(BLOCK_N, N_CTX, BLOCK_N): + kv_buf, kv_phase_idx = _buf_phase(accum_cnt_kv, NUM_BUFFERS_KV) + kv_offset = kv_offset_y + kv_idx + + tle.gpu.barrier_wait(k_empties[kv_buf], phaseIdx=kv_phase_idx) + tle.gpu.copy( + desc_k, + k_smem.slot(kv_buf), + [BLOCK_N, HEAD_DIM], + [kv_offset, 0], + barrier=k_fulls[kv_buf], + ) + + tle.gpu.barrier_wait(v_empties[kv_buf], phaseIdx=kv_phase_idx) + tle.gpu.copy( + desc_v, + v_smem.slot(kv_buf), + [BLOCK_N, HEAD_DIM], + [kv_offset, 0], + barrier=v_fulls[kv_buf], + ) + accum_cnt_kv += 1 + + tile_idx += num_progs + tile_count += 1 + + +@triton.jit +def _attn_fwd_tle_ws_pipelined_pingpong_consumer( + sm_scale, + m_ptr, + Z, + H, + desc_o, + N_CTX, + q_smem, + k_smem, + v_smem, + q_empties, + q_fulls, + k_empties, + k_fulls, + v_empties, + v_fulls, + ping_to_c0, + ping_to_c1, + HEAD_DIM: tl.constexpr, + BLOCK_M: tl.constexpr, + BLOCK_N: tl.constexpr, + NUM_BUFFERS_Q: tl.constexpr, + NUM_BUFFERS_KV: tl.constexpr, + BM_SPLIT: tl.constexpr, + CID: tl.constexpr, +): + prog_id = tl.program_id(0) + num_progs = tl.num_programs(0) + num_pid_m = tl.cdiv(N_CTX, BLOCK_M) + total_tiles = num_pid_m * Z * H + consumer_idx: tl.constexpr = CID - 1 + + if consumer_idx == 1: + tle.gpu.barrier_arrive(ping_to_c0) + + tile_idx = prog_id + tile_count = 0 + accum_cnt_kv = 0 + while tile_idx < total_tiles: + start_m, off_hz, qo_offset_y, kv_offset_y = _compute_offsets(tile_idx, H, N_CTX, BLOCK_M) + + m_i = tl.zeros([BM_SPLIT], dtype=tl.float32) - float("inf") + l_i = tl.zeros([BM_SPLIT], dtype=tl.float32) + 1.0 + acc = tl.zeros([BM_SPLIT, HEAD_DIM], dtype=tl.float32) + qk_scale = sm_scale * 1.44269504 + + q_buf, q_phase_idx = _buf_phase(tile_count, NUM_BUFFERS_Q) + q_idx = q_buf + consumer_idx * NUM_BUFFERS_Q + tle.gpu.barrier_wait(q_fulls[q_idx], phaseIdx=q_phase_idx) + + kv_buf, kv_phase_idx = _buf_phase(accum_cnt_kv, NUM_BUFFERS_KV) + tle.gpu.barrier_wait(k_fulls[kv_buf], phaseIdx=kv_phase_idx) + + if consumer_idx == 0: + tle.gpu.barrier_wait(ping_to_c0) + else: + tle.gpu.barrier_wait(ping_to_c1) + qk = tle.gpu.wgmma( + q_smem.slot(q_idx), + k_smem.slot(kv_buf), + out_dtype=tl.float32, + trans_b=True, + ) + if consumer_idx == 0: + tle.gpu.barrier_arrive(ping_to_c1) + else: + tle.gpu.barrier_arrive(ping_to_c0) + qk = tle.gpu.wgmma_wait(0, qk) + tle.gpu.barrier_arrive(k_empties[kv_buf], phaseIdx=kv_phase_idx) + + m_ij = tl.maximum(m_i, tl.max(qk, 1) * qk_scale) + qk = qk * qk_scale - m_ij[:, None] + p = tl.math.exp2(qk) + alpha = tl.math.exp2(m_i - m_ij) + l_ij = tl.sum(p, 1) + l_i = l_i * alpha + l_ij + m_i = m_ij + accum_cnt_kv += 1 + + for _ in range(BLOCK_N, N_CTX, BLOCK_N): + kv_buf, kv_phase_idx = _buf_phase(accum_cnt_kv, NUM_BUFFERS_KV) + tle.gpu.barrier_wait(k_fulls[kv_buf], phaseIdx=kv_phase_idx) + + if consumer_idx == 0: + tle.gpu.barrier_wait(ping_to_c0) + else: + tle.gpu.barrier_wait(ping_to_c1) + qk = tle.gpu.wgmma( + q_smem.slot(q_idx), + k_smem.slot(kv_buf), + out_dtype=tl.float32, + trans_b=True, + ) + if consumer_idx == 0: + tle.gpu.barrier_arrive(ping_to_c1) + else: + tle.gpu.barrier_arrive(ping_to_c0) + + v_buf, v_phase_idx = _buf_phase(accum_cnt_kv - 1, NUM_BUFFERS_KV) + tle.gpu.barrier_wait(v_fulls[v_buf], phaseIdx=v_phase_idx) + acc = tle.gpu.wgmma(p.to(tl.float16), v_smem.slot(v_buf), acc) + + qk = tle.gpu.wgmma_wait(1, qk) + tle.gpu.barrier_arrive(k_empties[kv_buf], phaseIdx=kv_phase_idx) + + m_ij = tl.maximum(m_i, tl.max(qk, 1) * qk_scale) + qk = qk * qk_scale - m_ij[:, None] + p = tl.math.exp2(qk) + alpha = tl.math.exp2(m_i - m_ij) + l_ij = tl.sum(p, 1) + l_i = l_i * alpha + l_ij + m_i = m_ij + + acc = tle.gpu.wgmma_wait(0, acc) + tle.gpu.barrier_arrive(v_empties[v_buf], phaseIdx=v_phase_idx) + acc = acc * alpha[:, None] + accum_cnt_kv += 1 + + v_buf, v_phase_idx = _buf_phase(accum_cnt_kv - 1, NUM_BUFFERS_KV) + tle.gpu.barrier_wait(v_fulls[v_buf], phaseIdx=v_phase_idx) + acc = tle.gpu.wgmma(p.to(tl.float16), v_smem.slot(v_buf), acc) + + acc = tle.gpu.wgmma_wait(1, acc) + tle.gpu.barrier_arrive(q_empties[q_idx], phaseIdx=q_phase_idx) + + acc = tle.gpu.wgmma_wait(0, acc) + tle.gpu.barrier_arrive(v_empties[v_buf], phaseIdx=v_phase_idx) + + m_i += tl.math.log2(l_i) + acc = acc / l_i[:, None] + + offs_m = start_m * BLOCK_M + consumer_idx * BM_SPLIT + tl.arange(0, BM_SPLIT) + m_ptrs = m_ptr + off_hz * N_CTX + offs_m + tl.store(m_ptrs, m_i) + + desc_o.store((qo_offset_y + consumer_idx * BM_SPLIT, 0), acc.to(tl.float16)) + + tile_idx += num_progs + tile_count += 1 + + +@triton.jit +def _attn_fwd_tle_ws_pipelined_pingpong_persistent( + sm_scale, + M, + Z, + H, + desc_q, + desc_k, + desc_v, + desc_o, + N_CTX, + HEAD_DIM: tl.constexpr, + BLOCK_M: tl.constexpr, + BLOCK_N: tl.constexpr, + NUM_BUFFERS_Q: tl.constexpr, + NUM_BUFFERS_KV: tl.constexpr, + NUM_MMA_WARPS: tl.constexpr, + NUM_MMA_GROUPS: tl.constexpr, + Q_STAGE_CAPACITY: tl.constexpr, + KV_STAGE_CAPACITY: tl.constexpr, +): + BM_SPLIT: tl.constexpr = BLOCK_M // NUM_MMA_GROUPS + THREADS_IN_MMA_GROUPS: tl.constexpr = NUM_MMA_WARPS * 32 + + q_smem = tle.gpu.alloc( + [Q_STAGE_CAPACITY, BM_SPLIT, HEAD_DIM], + dtype=tl.float16, + layout=None, + scope=tle.gpu.smem, + ) + k_smem = tle.gpu.alloc( + [KV_STAGE_CAPACITY, BLOCK_N, HEAD_DIM], + dtype=tl.float16, + layout=None, + scope=tle.gpu.smem, + ) + v_smem = tle.gpu.alloc( + [KV_STAGE_CAPACITY, BLOCK_N, HEAD_DIM], + dtype=tl.float16, + layout=None, + scope=tle.gpu.smem, + ) + q_empties = tle.gpu.alloc_barriers(num_barriers=Q_STAGE_CAPACITY, arrive_count=1, init=tle.gpu.READY) + q_fulls = tle.gpu.alloc_barriers( + num_barriers=Q_STAGE_CAPACITY, + arrive_count=1, + expect_bytes=BM_SPLIT * HEAD_DIM * 2, + ) + k_empties = tle.gpu.alloc_barriers( + num_barriers=KV_STAGE_CAPACITY, + arrive_count=NUM_MMA_GROUPS, + init=tle.gpu.READY, + ) + k_fulls = tle.gpu.alloc_barriers( + num_barriers=KV_STAGE_CAPACITY, + arrive_count=1, + expect_bytes=BLOCK_N * HEAD_DIM * 2, + ) + v_empties = tle.gpu.alloc_barriers( + num_barriers=KV_STAGE_CAPACITY, + arrive_count=NUM_MMA_GROUPS, + init=tle.gpu.READY, + ) + v_fulls = tle.gpu.alloc_barriers( + num_barriers=KV_STAGE_CAPACITY, + arrive_count=1, + expect_bytes=BLOCK_N * HEAD_DIM * 2, + ) + + # Named barriers for the consumer ping-pong issue protocol. They are used + # without phaseIdx, selecting TLE's named-barrier backend. + pingpong = tle.gpu.alloc_barriers(num_barriers=2, arrive_count=THREADS_IN_MMA_GROUPS) + ping_to_c0 = pingpong[0] + ping_to_c1 = pingpong[1] + + mma_warps: tl.constexpr = NUM_MMA_WARPS // NUM_MMA_GROUPS + # Pass partition CIDs explicitly: producer is 0, consumers start at 1. + tle.gpu.warp_specialize( + [ + ( + _attn_fwd_tle_ws_pipelined_pingpong_producer, + ( + Z, + H, + desc_q, + desc_k, + desc_v, + N_CTX, + q_smem, + k_smem, + v_smem, + q_empties, + q_fulls, + k_empties, + k_fulls, + v_empties, + v_fulls, + HEAD_DIM, + BLOCK_M, + BLOCK_N, + NUM_BUFFERS_Q, + NUM_BUFFERS_KV, + BM_SPLIT, + 0, + ), + ), + ( + _attn_fwd_tle_ws_pipelined_pingpong_consumer, + ( + sm_scale, + M, + Z, + H, + desc_o, + N_CTX, + q_smem, + k_smem, + v_smem, + q_empties, + q_fulls, + k_empties, + k_fulls, + v_empties, + v_fulls, + ping_to_c0, + ping_to_c1, + HEAD_DIM, + BLOCK_M, + BLOCK_N, + NUM_BUFFERS_Q, + NUM_BUFFERS_KV, + BM_SPLIT, + 1, + ), + ), + ( + _attn_fwd_tle_ws_pipelined_pingpong_consumer, + ( + sm_scale, + M, + Z, + H, + desc_o, + N_CTX, + q_smem, + k_smem, + v_smem, + q_empties, + q_fulls, + k_empties, + k_fulls, + v_empties, + v_fulls, + ping_to_c0, + ping_to_c1, + HEAD_DIM, + BLOCK_M, + BLOCK_N, + NUM_BUFFERS_Q, + NUM_BUFFERS_KV, + BM_SPLIT, + 2, + ), + ), + ], + [mma_warps, mma_warps], + [232, 232], + ) + + +def tle_attention( + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + sm_scale: float, + *, + block_m: int = 128, + block_n: int = 128, + num_buffers_q: int = 1, + num_buffers_kv: int = 2, + num_mma_warps: int = 8, + num_mma_groups: int = 2, + producer_num_warps: int = 4, + return_kernel: bool = False, + out: torch.Tensor | None = None, + m_out: torch.Tensor | None = None, +): + assert q.is_cuda and k.is_cuda and v.is_cuda + assert q.dtype == torch.float16 and k.dtype == torch.float16 and v.dtype == torch.float16 + assert q.is_contiguous() and k.is_contiguous() and v.is_contiguous() + assert q.shape == k.shape == v.shape + assert q.ndim == 4 + assert num_mma_groups == 2 + assert block_m % num_mma_groups == 0 + assert producer_num_warps % 4 == 0 + + z, h, n_ctx, head_dim = q.shape + if head_dim not in (16, 32, 64, 128, 256): + raise ValueError("HEAD_DIM must be one of 16, 32, 64, 128, 256") + if n_ctx % block_m != 0 or n_ctx % block_n != 0: + raise ValueError("this prototype expects N_CTX to be a multiple of BLOCK_M and BLOCK_N") + + triton.set_allocator(alloc_fn) + + if out is None: + o = torch.empty_like(q) + else: + assert out.shape == q.shape + assert out.device == q.device and out.dtype == q.dtype + assert out.is_contiguous() + o = out + if m_out is None: + m = torch.empty((z, h, n_ctx), device=q.device, dtype=torch.float32) + else: + assert m_out.shape == (z, h, n_ctx) + assert m_out.device == q.device and m_out.dtype == torch.float32 + assert m_out.is_contiguous() + m = m_out + y_dim = z * h * n_ctx + block_m_split = block_m // num_mma_groups + + desc_q = TensorDescriptor(q, shape=[y_dim, head_dim], strides=[head_dim, 1], block_shape=[block_m_split, head_dim]) + desc_k = TensorDescriptor(k, shape=[y_dim, head_dim], strides=[head_dim, 1], block_shape=[block_n, head_dim]) + desc_v = TensorDescriptor(v, shape=[y_dim, head_dim], strides=[head_dim, 1], block_shape=[block_n, head_dim]) + desc_o = TensorDescriptor(o, shape=[y_dim, head_dim], strides=[head_dim, 1], block_shape=[block_m_split, head_dim]) + + num_sms = torch.cuda.get_device_properties(q.device).multi_processor_count + total_tiles = triton.cdiv(n_ctx, block_m) * z * h + grid = (min(num_sms, total_tiles), ) + q_stage_capacity = _next_power_of_2(num_buffers_q * num_mma_groups) + kv_stage_capacity = _next_power_of_2(num_buffers_kv) + + kernel = _attn_fwd_tle_ws_pipelined_pingpong_persistent[grid]( + sm_scale, + m, + z, + h, + desc_q, + desc_k, + desc_v, + desc_o, + n_ctx, + HEAD_DIM=head_dim, + BLOCK_M=block_m, + BLOCK_N=block_n, + NUM_BUFFERS_Q=num_buffers_q, + NUM_BUFFERS_KV=num_buffers_kv, + NUM_MMA_WARPS=num_mma_warps, + NUM_MMA_GROUPS=num_mma_groups, + Q_STAGE_CAPACITY=q_stage_capacity, + KV_STAGE_CAPACITY=kv_stage_capacity, + num_warps=producer_num_warps, + ) + if return_kernel: + return o, m, kernel + return o + + +def _next_power_of_2(value: int) -> int: + assert value > 0 + return 1 << (value - 1).bit_length() + + +def reference_attention(q: torch.Tensor, k: torch.Tensor, v: torch.Tensor, sm_scale: float) -> torch.Tensor: + scores = torch.matmul(q.float(), k.transpose(-2, -1).float()) * sm_scale + probs = torch.softmax(scores, dim=-1) + return torch.matmul(probs, v.float()).to(q.dtype) + + +def sdpa_attention(q: torch.Tensor, k: torch.Tensor, v: torch.Tensor, sm_scale: float) -> torch.Tensor: + return torch.nn.functional.scaled_dot_product_attention( + q, + k, + v, + scale=sm_scale, + is_causal=False, + ) + + +def bench_ms(fn: Callable[[], object], warmup: int, rep: int, *, + cuda_graph: bool = False) -> tuple[float, float, float]: + if cuda_graph: + result = triton.testing.do_bench_cudagraph(fn, rep=rep, quantiles=(0.5, 0.2, 0.8)) + else: + result = triton.testing.do_bench(fn, warmup=warmup, rep=rep, quantiles=(0.5, 0.2, 0.8)) + if isinstance(result, (tuple, list)): + return float(result[0]), float(result[1]), float(result[2]) + ms = float(result) + return ms, ms, ms + + +@dataclass(frozen=True) +class AttentionProblem: + z: int + h: int + n_ctx: int + head_dim: int + + @property + def flops(self) -> int: + return 4 * self.z * self.h * self.n_ctx * self.n_ctx * self.head_dim + + +def parse_problem(text: str) -> AttentionProblem: + z, h, n_ctx, head_dim = [int(x) for x in text.lower().replace(",", "x").split("x")] + if min(z, h, n_ctx, head_dim) <= 0: + raise argparse.ArgumentTypeError("Z, H, N_CTX, and HEAD_DIM must be positive") + return AttentionProblem(z, h, n_ctx, head_dim) + + +def make_row( + variant: str, + problem: AttentionProblem, + ms: float, + p20: float, + p80: float, + block_m: int, + block_n: int, + extra: dict[str, object] | None = None, + cuda_graph: bool = False, +) -> dict[str, object]: + row: dict[str, object] = { + "variant": variant, + "Z": problem.z, + "H": problem.h, + "N_CTX": problem.n_ctx, + "HEAD_DIM": problem.head_dim, + "BLOCK_M": block_m, + "BLOCK_N": block_n, + "cuda_graph": cuda_graph, + "ms": f"{ms:.6f}", + "p20_ms": f"{p20:.6f}", + "p80_ms": f"{p80:.6f}", + "tflops_approx": f"{problem.flops / (ms * 1e-3) / 1e12:.3f}", + } + if extra: + row.update(extra) + return row + + +def make_error_row( + variant: str, + problem: AttentionProblem, + block_m: int, + block_n: int, + exc: Exception, +) -> dict[str, object]: + return { + "variant": variant, + "Z": problem.z, + "H": problem.h, + "N_CTX": problem.n_ctx, + "HEAD_DIM": problem.head_dim, + "BLOCK_M": block_m, + "BLOCK_N": block_n, + "ms": "", + "p20_ms": "", + "p80_ms": "", + "tflops_approx": "", + "status": "error", + "error_type": type(exc).__name__, + "error": str(exc), + } + + +def write_rows(rows: Iterable[dict[str, object]], out: str | None) -> None: + rows = list(rows) + fields: list[str] = [] + for row in rows: + for key in row: + if key not in fields: + fields.append(key) + if out: + with open(out, "w", newline="") as f: + writer = csv.DictWriter(f, fieldnames=fields) + writer.writeheader() + writer.writerows(rows) + writer = csv.DictWriter(__import__("sys").stdout, fieldnames=fields) + writer.writeheader() + writer.writerows(rows) + + +def main() -> None: + parser = argparse.ArgumentParser() + parser.add_argument("--problem", action="append", type=parse_problem, default=[], + help="attention problem ZxHxN_CTXxHEAD_DIM; may be repeated") + parser.add_argument("--sm-scale", type=float, default=None) + parser.add_argument("--block-m", type=int, default=128) + parser.add_argument("--block-n", type=int, default=128) + parser.add_argument("--warmup", type=int, default=25) + parser.add_argument("--rep", type=int, default=100) + parser.add_argument("--cuda-graph", action="store_true", + help="benchmark by capturing the workload once and timing CUDA Graph replay") + parser.add_argument("--out", default=None) + parser.add_argument("--check", action="store_true") + parser.add_argument("--include-sdpa", action="store_true", help="also benchmark torch SDPA on the same inputs") + parser.add_argument("--sdpa-requires-grad", action="store_true", + help="set q/k/v.requires_grad_() to match the TLX SDPA perf test") + parser.add_argument("--continue-on-tle-error", action="store_true", + help="record a CSV error row instead of aborting if the TLE kernel fails") + parser.add_argument("--dump-summary", action="store_true") + args = parser.parse_args() + + if not torch.cuda.is_available() or torch.cuda.get_device_capability()[0] < 9: + raise RuntimeError("Hopper or newer CUDA GPU is required") + + problems = args.problem or [AttentionProblem(1, 1, 1024, 64)] + rows = [] + for problem in problems: + q = torch.randn((problem.z, problem.h, problem.n_ctx, problem.head_dim), device=DEVICE, dtype=torch.float16) + k = torch.randn_like(q) + v = torch.randn_like(q) + if args.sdpa_requires_grad: + q.requires_grad_() + k.requires_grad_() + v.requires_grad_() + sm_scale = args.sm_scale + if sm_scale is None: + sm_scale = problem.head_dim**-0.5 + + if args.include_sdpa: + sdpa_out = sdpa_attention(q, k, v, sm_scale) + torch.cuda.synchronize() + + def run_sdpa(): + sdpa_attention(q, k, v, sm_scale) + + ms, p20, p80 = bench_ms(run_sdpa, args.warmup, args.rep, cuda_graph=args.cuda_graph) + rows.append( + make_row( + "SDPA", + problem, + ms, + p20, + p80, + args.block_m, + args.block_n, + { + "output_dtype": str(sdpa_out.dtype).replace("torch.", ""), + "has_warp_specialize": False, + "has_wgmma": False, + "baseline_source": "torch.nn.functional.scaled_dot_product_attention", + "requires_grad": args.sdpa_requires_grad, + "status": "ok", + }, + cuda_graph=args.cuda_graph, + )) + + try: + o, m, kernel = tle_attention( + q, + k, + v, + sm_scale, + block_m=args.block_m, + block_n=args.block_n, + return_kernel=True, + ) + torch.cuda.synchronize() + + if args.check: + ref = reference_attention(q, k, v, sm_scale) + torch.testing.assert_close(o, ref, atol=5e-2, rtol=5e-2) + + bench_o = torch.empty_like(q) + bench_m = torch.empty((problem.z, problem.h, problem.n_ctx), device=q.device, dtype=torch.float32) + + def run(): + tle_attention( + q, + k, + v, + sm_scale, + block_m=args.block_m, + block_n=args.block_n, + out=bench_o, + m_out=bench_m, + ) + + ms, p20, p80 = bench_ms(run, args.warmup, args.rep, cuda_graph=args.cuda_graph) + extra = { + "output_dtype": str(o.dtype).replace("torch.", ""), + "has_warp_specialize": "ttg.warp_specialize" in kernel.asm["ttgir"], + "has_wgmma": "ttng.warp_group_dot" in kernel.asm["ttgir"], + "status": "ok", + } + if args.dump_summary: + extra["ttgir_len"] = len(kernel.asm.get("ttgir", "")) + extra["ptx_len"] = len(kernel.asm.get("ptx", "")) + rows.append( + make_row( + "flagtree.tle.fa3.ws_pipelined_pingpong_persistent", + problem, + ms, + p20, + p80, + args.block_m, + args.block_n, + extra, + cuda_graph=args.cuda_graph, + )) + except Exception as exc: + if not args.continue_on_tle_error: + raise + rows.append( + make_error_row( + "flagtree.tle.fa3.ws_pipelined_pingpong_persistent", + problem, + args.block_m, + args.block_n, + exc, + )) + + write_rows(rows, args.out) + + +if __name__ == "__main__": + main() diff --git a/third_party/tle/tutorials/tle_hopper_gemm_ws_persistent.py b/third_party/tle/tutorials/tle_hopper_gemm_ws_persistent.py new file mode 100755 index 0000000000..564a01e101 --- /dev/null +++ b/third_party/tle/tutorials/tle_hopper_gemm_ws_persistent.py @@ -0,0 +1,1432 @@ +#!/usr/bin/env python3 +"""Experimental TLE persistent warp-specialized GEMM. + +This is a performance-oriented prototype that mirrors the simpler TLX +``hopper_gemm_ws.py`` structure: + +* one producer partition issues TMA loads into staged shared-memory buffers +* two explicit consumer partitions split the output tile along M +* full/empty mbarriers coordinate producer/consumer ownership of each stage +* WGMMA computes each half tile and descriptor stores write C + +It intentionally avoids the more advanced ping-pong epilogue scheduling in +``hopper-persistent-gemm-ws-pingpong.py`` so it can stay within today's TLE API. +""" + +from __future__ import annotations + +import argparse +import csv +import importlib.util +import pathlib +import sys +from dataclasses import dataclass +from typing import Callable, Iterable, Optional + +import torch +import triton +import triton.language as tl +import triton.experimental.tle.language as tle +from triton.tools.tensor_descriptor import TensorDescriptor + +DEVICE = triton.runtime.driver.active.get_active_torch_device() + + +def alloc_fn(size: int, align: int, stream: Optional[int]): + return torch.empty(size, dtype=torch.int8, device=DEVICE) + + +@triton.jit +def _compute_pid(tile_id, num_pid_n, num_pid_m, group_size_m: tl.constexpr): + num_pid_in_group = group_size_m * num_pid_n + group_id = tile_id // num_pid_in_group + first_pid_m = group_id * group_size_m + cur_group_size_m = min(num_pid_m - first_pid_m, group_size_m) + pid_m = first_pid_m + (tile_id % cur_group_size_m) + pid_n = (tile_id % num_pid_in_group) // cur_group_size_m + return pid_m, pid_n + + +@triton.jit +def _buf_phase(count, num_stages: tl.constexpr): + buf = count % num_stages + phase = (count // num_stages) & 1 + return buf, phase + + +@triton.jit +def _tle_ws_persistent_gemm_producer( + a_desc, + b_desc, + a_smem, + b_smem, + empty_a, + empty_b, + full_a, + full_b, + M, + N, + K, + NUM_SMS: tl.constexpr, + BM: tl.constexpr, + BN: tl.constexpr, + BK: tl.constexpr, + GROUP_SIZE_M: tl.constexpr, + NUM_STAGES: tl.constexpr, + BM_SPLIT: tl.constexpr, + CID: tl.constexpr, +): + sm_id = tl.program_id(0) + num_pid_m = tl.cdiv(M, BM) + num_pid_n = tl.cdiv(N, BN) + num_tiles = num_pid_m * num_pid_n + + tile_id = sm_id + smem_count = 0 + while tile_id < num_tiles: + pid_m, pid_n = _compute_pid(tile_id, num_pid_n, num_pid_m, GROUP_SIZE_M) + off_m = pid_m * BM + off_n = pid_n * BN + + for k_iter in range(0, tl.cdiv(K, BK)): + buf, phase = _buf_phase(smem_count, NUM_STAGES) + off_k = k_iter * BK + + a0_idx = buf + a1_idx = buf + NUM_STAGES + + tle.gpu.barrier_wait(empty_a[a0_idx], phaseIdx=phase) + tle.gpu.copy( + a_desc, + a_smem.slot(a0_idx), + [BM_SPLIT, BK], + [off_m, off_k], + barrier=full_a[a0_idx], + ) + + tle.gpu.barrier_wait(empty_b[buf], phaseIdx=phase) + tle.gpu.copy( + b_desc, + b_smem.slot(buf), + [BK, BN], + [off_k, off_n], + barrier=full_b[buf], + ) + + tle.gpu.barrier_wait(empty_a[a1_idx], phaseIdx=phase) + tle.gpu.copy( + a_desc, + a_smem.slot(a1_idx), + [BM_SPLIT, BK], + [off_m + BM_SPLIT, off_k], + barrier=full_a[a1_idx], + ) + + smem_count += 1 + + tile_id += NUM_SMS + + +@triton.jit +def _tle_ws_persistent_gemm_consumer( + a_smem, + b_smem, + empty_a, + empty_b, + full_a, + full_b, + c_desc, + M, + N, + K, + NUM_SMS: tl.constexpr, + BM: tl.constexpr, + BN: tl.constexpr, + BK: tl.constexpr, + GROUP_SIZE_M: tl.constexpr, + NUM_STAGES: tl.constexpr, + BM_SPLIT: tl.constexpr, + WGMMA_PIPELINE: tl.constexpr, + CID: tl.constexpr, +): + sm_id = tl.program_id(0) + num_pid_m = tl.cdiv(M, BM) + num_pid_n = tl.cdiv(N, BN) + num_tiles = num_pid_m * num_pid_n + consumer_idx: tl.constexpr = CID - 1 + + tile_id = sm_id + smem_count = 0 + while tile_id < num_tiles: + pid_m, pid_n = _compute_pid(tile_id, num_pid_n, num_pid_m, GROUP_SIZE_M) + off_m = pid_m * BM + consumer_idx * BM_SPLIT + off_n = pid_n * BN + + acc = tl.zeros((BM_SPLIT, BN), dtype=tl.float32) + if WGMMA_PIPELINE: + buf, phase = _buf_phase(smem_count, NUM_STAGES) + a_idx = buf + consumer_idx * NUM_STAGES + + tle.gpu.barrier_wait(full_a[a_idx], phaseIdx=phase) + tle.gpu.barrier_wait(full_b[buf], phaseIdx=phase) + acc = tle.gpu.wgmma(a_smem.slot(a_idx), b_smem.slot(buf), acc) + + last_buf = buf + last_phase = phase + last_a_idx = a_idx + smem_count += 1 + + for k_iter in range(1, tl.cdiv(K, BK)): + buf, phase = _buf_phase(smem_count, NUM_STAGES) + a_idx = buf + consumer_idx * NUM_STAGES + + tle.gpu.barrier_wait(full_a[a_idx], phaseIdx=phase) + tle.gpu.barrier_wait(full_b[buf], phaseIdx=phase) + acc = tle.gpu.wgmma(a_smem.slot(a_idx), b_smem.slot(buf), acc) + acc = tle.gpu.wgmma_wait(1, acc) + + tle.gpu.barrier_arrive(empty_a[last_a_idx], phaseIdx=last_phase) + tle.gpu.barrier_arrive(empty_b[last_buf], phaseIdx=last_phase) + last_buf = buf + last_phase = phase + last_a_idx = a_idx + smem_count += 1 + + acc = tle.gpu.wgmma_wait(0, acc) + tle.gpu.barrier_arrive(empty_a[last_a_idx], phaseIdx=last_phase) + tle.gpu.barrier_arrive(empty_b[last_buf], phaseIdx=last_phase) + else: + for k_iter in range(0, tl.cdiv(K, BK)): + buf, phase = _buf_phase(smem_count, NUM_STAGES) + a_idx = buf + consumer_idx * NUM_STAGES + + tle.gpu.barrier_wait(full_a[a_idx], phaseIdx=phase) + tle.gpu.barrier_wait(full_b[buf], phaseIdx=phase) + + acc = tle.gpu.wgmma(a_smem.slot(a_idx), b_smem.slot(buf), acc) + acc = tle.gpu.wgmma_wait(0, acc) + + tle.gpu.barrier_arrive(empty_a[a_idx], phaseIdx=phase) + tle.gpu.barrier_arrive(empty_b[buf], phaseIdx=phase) + smem_count += 1 + + c_desc.store((off_m, off_n), acc.to(tl.float16)) + tile_id += NUM_SMS + + +@triton.jit +def _tle_ws_nonpersistent_gemm_producer( + a_desc, + b_desc, + a_smem, + b_smem, + empty_a, + empty_b, + full_a, + full_b, + M, + N, + K, + BM: tl.constexpr, + BN: tl.constexpr, + BK: tl.constexpr, + GROUP_SIZE_M: tl.constexpr, + NUM_STAGES: tl.constexpr, + BM_SPLIT: tl.constexpr, + CID: tl.constexpr, +): + tile_id = tl.program_id(0) + num_pid_m = tl.cdiv(M, BM) + num_pid_n = tl.cdiv(N, BN) + pid_m, pid_n = _compute_pid(tile_id, num_pid_n, num_pid_m, GROUP_SIZE_M) + off_m = pid_m * BM + off_n = pid_n * BN + + for k_iter in range(0, tl.cdiv(K, BK)): + buf, phase = _buf_phase(k_iter, NUM_STAGES) + off_k = k_iter * BK + + a0_idx = buf + a1_idx = buf + NUM_STAGES + + tle.gpu.barrier_wait(empty_a[a0_idx], phaseIdx=phase) + tle.gpu.copy( + a_desc, + a_smem.slot(a0_idx), + [BM_SPLIT, BK], + [off_m, off_k], + barrier=full_a[a0_idx], + ) + + tle.gpu.barrier_wait(empty_b[buf], phaseIdx=phase) + tle.gpu.copy( + b_desc, + b_smem.slot(buf), + [BK, BN], + [off_k, off_n], + barrier=full_b[buf], + ) + + tle.gpu.barrier_wait(empty_a[a1_idx], phaseIdx=phase) + tle.gpu.copy( + a_desc, + a_smem.slot(a1_idx), + [BM_SPLIT, BK], + [off_m + BM_SPLIT, off_k], + barrier=full_a[a1_idx], + ) + + +@triton.jit +def _tle_ws_nonpersistent_gemm_consumer( + a_smem, + b_smem, + empty_a, + empty_b, + full_a, + full_b, + c_desc, + M, + N, + K, + BM: tl.constexpr, + BN: tl.constexpr, + BK: tl.constexpr, + GROUP_SIZE_M: tl.constexpr, + NUM_STAGES: tl.constexpr, + BM_SPLIT: tl.constexpr, + WGMMA_PIPELINE: tl.constexpr, + CID: tl.constexpr, +): + tile_id = tl.program_id(0) + num_pid_m = tl.cdiv(M, BM) + num_pid_n = tl.cdiv(N, BN) + pid_m, pid_n = _compute_pid(tile_id, num_pid_n, num_pid_m, GROUP_SIZE_M) + consumer_idx: tl.constexpr = CID - 1 + off_m = pid_m * BM + consumer_idx * BM_SPLIT + off_n = pid_n * BN + + acc = tl.zeros((BM_SPLIT, BN), dtype=tl.float32) + if WGMMA_PIPELINE: + buf, phase = _buf_phase(0, NUM_STAGES) + a_idx = buf + consumer_idx * NUM_STAGES + + tle.gpu.barrier_wait(full_a[a_idx], phaseIdx=phase) + tle.gpu.barrier_wait(full_b[buf], phaseIdx=phase) + acc = tle.gpu.wgmma(a_smem.slot(a_idx), b_smem.slot(buf), acc) + + last_buf = buf + last_phase = phase + last_a_idx = a_idx + + for k_iter in range(1, tl.cdiv(K, BK)): + buf, phase = _buf_phase(k_iter, NUM_STAGES) + a_idx = buf + consumer_idx * NUM_STAGES + + tle.gpu.barrier_wait(full_a[a_idx], phaseIdx=phase) + tle.gpu.barrier_wait(full_b[buf], phaseIdx=phase) + acc = tle.gpu.wgmma(a_smem.slot(a_idx), b_smem.slot(buf), acc) + acc = tle.gpu.wgmma_wait(1, acc) + + tle.gpu.barrier_arrive(empty_a[last_a_idx], phaseIdx=last_phase) + tle.gpu.barrier_arrive(empty_b[last_buf], phaseIdx=last_phase) + last_buf = buf + last_phase = phase + last_a_idx = a_idx + + acc = tle.gpu.wgmma_wait(0, acc) + tle.gpu.barrier_arrive(empty_a[last_a_idx], phaseIdx=last_phase) + tle.gpu.barrier_arrive(empty_b[last_buf], phaseIdx=last_phase) + else: + for k_iter in range(0, tl.cdiv(K, BK)): + buf, phase = _buf_phase(k_iter, NUM_STAGES) + a_idx = buf + consumer_idx * NUM_STAGES + + tle.gpu.barrier_wait(full_a[a_idx], phaseIdx=phase) + tle.gpu.barrier_wait(full_b[buf], phaseIdx=phase) + + acc = tle.gpu.wgmma(a_smem.slot(a_idx), b_smem.slot(buf), acc) + acc = tle.gpu.wgmma_wait(0, acc) + + tle.gpu.barrier_arrive(empty_a[a_idx], phaseIdx=phase) + tle.gpu.barrier_arrive(empty_b[buf], phaseIdx=phase) + + c_desc.store((off_m, off_n), acc.to(tl.float16)) + + +@triton.jit +def _tle_ws_persistent_gemm_kernel( + a_desc, + b_desc, + c_desc, + M, + N, + K, + NUM_SMS: tl.constexpr, + BM: tl.constexpr, + BN: tl.constexpr, + BK: tl.constexpr, + GROUP_SIZE_M: tl.constexpr, + NUM_STAGES: tl.constexpr, + A_STAGE_CAPACITY: tl.constexpr, + B_STAGE_CAPACITY: tl.constexpr, + WGMMA_PIPELINE: tl.constexpr, +): + # v1 mirrors TLX's two-MMA-group split in M. + BM_SPLIT: tl.constexpr = BM // 2 + + a_smem = tle.gpu.alloc( + [A_STAGE_CAPACITY, BM_SPLIT, BK], + dtype=tl.float16, + layout=None, + scope=tle.gpu.smem, + ) + b_smem = tle.gpu.alloc( + [B_STAGE_CAPACITY, BK, BN], + dtype=tl.float16, + layout=None, + scope=tle.gpu.smem, + ) + + empty_a = tle.gpu.alloc_barriers(num_barriers=A_STAGE_CAPACITY, arrive_count=1, init=tle.gpu.READY) + empty_b = tle.gpu.alloc_barriers(num_barriers=B_STAGE_CAPACITY, arrive_count=2, init=tle.gpu.READY) + full_a = tle.gpu.alloc_barriers( + num_barriers=A_STAGE_CAPACITY, + arrive_count=1, + expect_bytes=BM_SPLIT * BK * 2, + ) + full_b = tle.gpu.alloc_barriers( + num_barriers=B_STAGE_CAPACITY, + arrive_count=1, + expect_bytes=BK * BN * 2, + ) + + # Pass partition CIDs explicitly: producer is 0, consumers start at 1. + tle.gpu.warp_specialize( + [ + ( + _tle_ws_persistent_gemm_producer, + ( + a_desc, + b_desc, + a_smem, + b_smem, + empty_a, + empty_b, + full_a, + full_b, + M, + N, + K, + NUM_SMS, + BM, + BN, + BK, + GROUP_SIZE_M, + NUM_STAGES, + BM_SPLIT, + 0, + ), + ), + ( + _tle_ws_persistent_gemm_consumer, + ( + a_smem, + b_smem, + empty_a, + empty_b, + full_a, + full_b, + c_desc, + M, + N, + K, + NUM_SMS, + BM, + BN, + BK, + GROUP_SIZE_M, + NUM_STAGES, + BM_SPLIT, + WGMMA_PIPELINE, + 1, + ), + ), + ( + _tle_ws_persistent_gemm_consumer, + ( + a_smem, + b_smem, + empty_a, + empty_b, + full_a, + full_b, + c_desc, + M, + N, + K, + NUM_SMS, + BM, + BN, + BK, + GROUP_SIZE_M, + NUM_STAGES, + BM_SPLIT, + WGMMA_PIPELINE, + 2, + ), + ), + ], + [4, 4], + [168, 168], + ) + + +@triton.jit +def _tle_ws_nonpersistent_gemm_kernel( + a_desc, + b_desc, + c_desc, + M, + N, + K, + BM: tl.constexpr, + BN: tl.constexpr, + BK: tl.constexpr, + GROUP_SIZE_M: tl.constexpr, + NUM_STAGES: tl.constexpr, + A_STAGE_CAPACITY: tl.constexpr, + B_STAGE_CAPACITY: tl.constexpr, + WGMMA_PIPELINE: tl.constexpr, +): + BM_SPLIT: tl.constexpr = BM // 2 + + a_smem = tle.gpu.alloc( + [A_STAGE_CAPACITY, BM_SPLIT, BK], + dtype=tl.float16, + layout=None, + scope=tle.gpu.smem, + ) + b_smem = tle.gpu.alloc( + [B_STAGE_CAPACITY, BK, BN], + dtype=tl.float16, + layout=None, + scope=tle.gpu.smem, + ) + + empty_a = tle.gpu.alloc_barriers(num_barriers=A_STAGE_CAPACITY, arrive_count=1, init=tle.gpu.READY) + empty_b = tle.gpu.alloc_barriers(num_barriers=B_STAGE_CAPACITY, arrive_count=2, init=tle.gpu.READY) + full_a = tle.gpu.alloc_barriers( + num_barriers=A_STAGE_CAPACITY, + arrive_count=1, + expect_bytes=BM_SPLIT * BK * 2, + ) + full_b = tle.gpu.alloc_barriers( + num_barriers=B_STAGE_CAPACITY, + arrive_count=1, + expect_bytes=BK * BN * 2, + ) + + # Pass partition CIDs explicitly: producer is 0, consumers start at 1. + tle.gpu.warp_specialize( + [ + ( + _tle_ws_nonpersistent_gemm_producer, + ( + a_desc, + b_desc, + a_smem, + b_smem, + empty_a, + empty_b, + full_a, + full_b, + M, + N, + K, + BM, + BN, + BK, + GROUP_SIZE_M, + NUM_STAGES, + BM_SPLIT, + 0, + ), + ), + ( + _tle_ws_nonpersistent_gemm_consumer, + ( + a_smem, + b_smem, + empty_a, + empty_b, + full_a, + full_b, + c_desc, + M, + N, + K, + BM, + BN, + BK, + GROUP_SIZE_M, + NUM_STAGES, + BM_SPLIT, + WGMMA_PIPELINE, + 1, + ), + ), + ( + _tle_ws_nonpersistent_gemm_consumer, + ( + a_smem, + b_smem, + empty_a, + empty_b, + full_a, + full_b, + c_desc, + M, + N, + K, + BM, + BN, + BK, + GROUP_SIZE_M, + NUM_STAGES, + BM_SPLIT, + WGMMA_PIPELINE, + 2, + ), + ), + ], + [4, 4], + [168, 168], + ) + + +def tle_ws_persistent_matmul( + a: torch.Tensor, + b: torch.Tensor, + *, + bm: int = 128, + bn: int = 128, + bk: int = 64, + group_size_m: int = 8, + num_stages: int = 3, + wgmma_pipeline: bool = True, + producer_num_warps: int = 4, + out: torch.Tensor | None = None, +): + assert a.is_cuda and b.is_cuda + assert a.dtype == torch.float16 and b.dtype == torch.float16 + assert a.is_contiguous() and b.is_contiguous() + assert a.shape[1] == b.shape[0] + assert bm % 2 == 0 + if wgmma_pipeline and num_stages < 2: + raise ValueError("wgmma_pipeline requires at least two logical smem stages") + + triton.set_allocator(alloc_fn) + + m, k = a.shape + _, n = b.shape + if out is None: + c = torch.empty((m, n), device=a.device, dtype=torch.float16) + else: + assert out.shape == (m, n) + assert out.device == a.device and out.dtype == torch.float16 + assert out.is_contiguous() + c = out + + a_desc = TensorDescriptor.from_tensor(a, block_shape=[bm // 2, bk]) + b_desc = TensorDescriptor.from_tensor(b, block_shape=[bk, bn]) + c_desc = TensorDescriptor.from_tensor(c, block_shape=[bm // 2, bn]) + num_sms = torch.cuda.get_device_properties(a.device).multi_processor_count + grid = (min(num_sms, triton.cdiv(m, bm) * triton.cdiv(n, bn)), ) + a_stage_capacity = _next_power_of_2(num_stages * 2) + b_stage_capacity = _next_power_of_2(num_stages) + + kernel = _tle_ws_persistent_gemm_kernel[grid]( + a_desc, + b_desc, + c_desc, + m, + n, + k, + NUM_SMS=num_sms, + BM=bm, + BN=bn, + BK=bk, + GROUP_SIZE_M=group_size_m, + NUM_STAGES=num_stages, + A_STAGE_CAPACITY=a_stage_capacity, + B_STAGE_CAPACITY=b_stage_capacity, + WGMMA_PIPELINE=wgmma_pipeline, + num_warps=producer_num_warps, + ) + return c, kernel + + +def tle_ws_nonpersistent_matmul( + a: torch.Tensor, + b: torch.Tensor, + *, + bm: int = 128, + bn: int = 128, + bk: int = 64, + group_size_m: int = 8, + num_stages: int = 3, + wgmma_pipeline: bool = True, + producer_num_warps: int = 4, + out: torch.Tensor | None = None, +): + assert a.is_cuda and b.is_cuda + assert a.dtype == torch.float16 and b.dtype == torch.float16 + assert a.is_contiguous() and b.is_contiguous() + assert a.shape[1] == b.shape[0] + assert bm % 2 == 0 + if wgmma_pipeline and num_stages < 2: + raise ValueError("wgmma_pipeline requires at least two logical smem stages") + + triton.set_allocator(alloc_fn) + + m, k = a.shape + _, n = b.shape + if m % bm != 0 or n % bn != 0 or k % bk != 0: + raise ValueError("non-persistent TLE benchmark currently expects M, N, K to be exact tile multiples") + if out is None: + c = torch.empty((m, n), device=a.device, dtype=torch.float16) + else: + assert out.shape == (m, n) + assert out.device == a.device and out.dtype == torch.float16 + assert out.is_contiguous() + c = out + + a_desc = TensorDescriptor.from_tensor(a, block_shape=[bm // 2, bk]) + b_desc = TensorDescriptor.from_tensor(b, block_shape=[bk, bn]) + c_desc = TensorDescriptor.from_tensor(c, block_shape=[bm // 2, bn]) + grid = (triton.cdiv(m, bm) * triton.cdiv(n, bn), ) + a_stage_capacity = _next_power_of_2(num_stages * 2) + b_stage_capacity = _next_power_of_2(num_stages) + + kernel = _tle_ws_nonpersistent_gemm_kernel[grid]( + a_desc, + b_desc, + c_desc, + m, + n, + k, + BM=bm, + BN=bn, + BK=bk, + GROUP_SIZE_M=group_size_m, + NUM_STAGES=num_stages, + A_STAGE_CAPACITY=a_stage_capacity, + B_STAGE_CAPACITY=b_stage_capacity, + WGMMA_PIPELINE=wgmma_pipeline, + num_warps=producer_num_warps, + ) + return c, kernel + + +def _next_power_of_2(value: int) -> int: + assert value > 0 + return 1 << (value - 1).bit_length() + + +@dataclass(frozen=True) +class Problem: + m: int + n: int + k: int + + @property + def flops(self) -> int: + return 2 * self.m * self.n * self.k + + +@dataclass(frozen=True) +class GemmConfig: + bm: int + bn: int + bk: int + num_stages: int + group_size_m: int + wgmma_pipeline: bool + producer_num_warps: int + + def label(self) -> str: + wait = "pipe" if self.wgmma_pipeline else "sync" + return (f"BM{self.bm}.BN{self.bn}.BK{self.bk}.S{self.num_stages}." + f"G{self.group_size_m}.P{self.producer_num_warps}.{wait}") + + +def parse_problem(text: str) -> Problem: + m, n, k = [int(x) for x in text.lower().replace(",", "x").split("x")] + return Problem(m, n, k) + + +def parse_gemm_config(text: str) -> GemmConfig: + dims_text, _, mode_text = text.partition(":") + try: + dims = [int(x) for x in dims_text.lower().replace(",", "x").split("x")] + except Exception as exc: + raise argparse.ArgumentTypeError( + f"invalid config '{text}', expected BMxBNxBKxSTAGES[xGROUP_SIZE_M[xPRODUCER_WARPS]][:pipe|sync]") from exc + if len(dims) not in (4, 5, 6): + raise argparse.ArgumentTypeError( + f"invalid config '{text}', expected BMxBNxBKxSTAGES[xGROUP_SIZE_M[xPRODUCER_WARPS]][:pipe|sync]") + mode = mode_text.lower() or "pipe" + if mode in ("pipe", "pipeline", "pipelined", "wait1"): + wgmma_pipeline = True + elif mode in ("sync", "serial", "wait0", "nopipe", "no-pipe"): + wgmma_pipeline = False + else: + raise argparse.ArgumentTypeError("config mode must be pipe or sync") + + group_size_m = dims[4] if len(dims) >= 5 else 8 + producer_num_warps = dims[5] if len(dims) == 6 else 4 + if producer_num_warps % 4 != 0: + raise argparse.ArgumentTypeError("PRODUCER_WARPS must be a multiple of 4") + return GemmConfig(dims[0], dims[1], dims[2], dims[3], group_size_m, wgmma_pipeline, producer_num_warps) + + +def default_compare_configs(args: argparse.Namespace) -> list[GemmConfig]: + configs = [ + GemmConfig(args.bm, args.bn, args.bk, args.num_stages, args.group_size_m, False, args.producer_num_warps), + GemmConfig(args.bm, args.bn, args.bk, args.num_stages, args.group_size_m, True, args.producer_num_warps), + ] + aligned_bn = 256 if args.bn == 128 else args.bn + aligned_stages = 2 if args.num_stages != 2 else args.num_stages + configs.append( + GemmConfig(args.bm, aligned_bn, args.bk, aligned_stages, args.group_size_m, True, args.producer_num_warps)) + + unique: list[GemmConfig] = [] + for cfg in configs: + if cfg not in unique: + unique.append(cfg) + return unique + + +def bench_ms(fn: Callable[[], object], warmup: int, rep: int, *, + cuda_graph: bool = False) -> tuple[float, float, float]: + if cuda_graph: + result = triton.testing.do_bench_cudagraph(fn, rep=rep, quantiles=(0.5, 0.2, 0.8)) + else: + result = triton.testing.do_bench(fn, warmup=warmup, rep=rep, quantiles=(0.5, 0.2, 0.8)) + if not isinstance(result, (tuple, list)): + ms = float(result) + return ms, ms, ms + return float(result[0]), float(result[1]), float(result[2]) + + +def make_row( + variant: str, + problem: Problem, + ms: float, + p20: float, + p80: float, + bm: int, + bn: int, + bk: int, + num_stages: int, + group_size_m: int, + wgmma_pipeline: bool, + producer_num_warps: int, + *, + has_warp_specialize: bool | None = None, + has_wgmma: bool | None = None, + cuda_graph: bool = False, +) -> dict[str, object]: + row: dict[str, object] = { + "variant": variant, + "M": problem.m, + "N": problem.n, + "K": problem.k, + "BM": bm, + "BN": bn, + "BK": bk, + "NUM_STAGES": num_stages, + "GROUP_SIZE_M": group_size_m, + "PRODUCER_NUM_WARPS": producer_num_warps, + "wgmma_pipeline": wgmma_pipeline, + "producer_num_warps": producer_num_warps, + "cuda_graph": cuda_graph, + "ms": f"{ms:.6f}", + "p20_ms": f"{p20:.6f}", + "p80_ms": f"{p80:.6f}", + "tflops": f"{problem.flops / (ms * 1e-3) / 1e12:.3f}", + } + if has_warp_specialize is not None: + row["has_warp_specialize"] = has_warp_specialize + if has_wgmma is not None: + row["has_wgmma"] = has_wgmma + return row + + +def make_baseline_row( + variant: str, + problem: Problem, + ms: float, + p20: float, + p80: float, + *, + baseline_source: str | None = None, + cuda_graph: bool = False, +) -> dict[str, object]: + row: dict[str, object] = { + "variant": variant, + "M": problem.m, + "N": problem.n, + "K": problem.k, + "cuda_graph": cuda_graph, + "ms": f"{ms:.6f}", + "p20_ms": f"{p20:.6f}", + "p80_ms": f"{p80:.6f}", + "tflops": f"{problem.flops / (ms * 1e-3) / 1e12:.3f}", + } + if baseline_source is not None: + row["baseline_source"] = baseline_source + return row + + +def make_error_row(variant: str, problem: Problem, note: str, + extra: dict[str, object] | None = None) -> dict[str, object]: + row: dict[str, object] = { + "variant": variant, + "M": problem.m, + "N": problem.n, + "K": problem.k, + "ms": "", + "p20_ms": "", + "p80_ms": "", + "tflops": "", + "note": note, + } + if extra: + row.update(extra) + return row + + +def make_native_row( + variant: str, + problem: Problem, + ms: float, + p20: float, + p80: float, + cfg: dict[str, int], + *, + has_warp_specialize: bool, + has_wgmma: bool, + cuda_graph: bool = False, +) -> dict[str, object]: + return { + "variant": variant, + "M": problem.m, + "N": problem.n, + "K": problem.k, + "ms": f"{ms:.6f}", + "p20_ms": f"{p20:.6f}", + "p80_ms": f"{p80:.6f}", + "tflops": f"{problem.flops / (ms * 1e-3) / 1e12:.3f}", + **cfg, + "cuda_graph": cuda_graph, + "has_warp_specialize": has_warp_specialize, + "has_wgmma": has_wgmma, + } + + +def run_torch_matmul_baseline( + variant: str, + problem: Problem, + a: torch.Tensor, + b: torch.Tensor, + warmup: int, + rep: int, + *, + use_out: bool, + cuda_graph: bool = False, +) -> dict[str, object]: + if use_out or cuda_graph: + c = torch.empty((problem.m, problem.n), device=a.device, dtype=torch.float16) + + def run(): + torch.matmul(a, b, out=c) + else: + + def run(): + torch.matmul(a, b) + + ms, p20, p80 = bench_ms(run, warmup, rep, cuda_graph=cuda_graph) + baseline_source = "tlx.torch.matmul" if variant == "cuBLAS" else None + return make_baseline_row(variant, problem, ms, p20, p80, baseline_source=baseline_source, cuda_graph=cuda_graph) + + +def add_cublas_speedups(rows: list[dict[str, object]]) -> None: + cublas_ms: dict[tuple[int, int, int], float] = {} + for row in rows: + if row.get("variant") == "cuBLAS" and row.get("ms"): + key = (int(row["M"]), int(row["N"]), int(row["K"])) + cublas_ms[key] = float(row["ms"]) + + for row in rows: + key = (int(row["M"]), int(row["N"]), int(row["K"])) + if key not in cublas_ms or not row.get("ms"): + continue + row["speedup_vs_cublas"] = f"{cublas_ms[key] / float(row['ms']):.3f}" + + +def import_from_path(module_name: str, path: pathlib.Path): + spec = importlib.util.spec_from_file_location(module_name, path) + if spec is None or spec.loader is None: + raise RuntimeError(f"cannot load module from {path}") + module = importlib.util.module_from_spec(spec) + sys.modules[module_name] = module + spec.loader.exec_module(module) + return module + + +def find_flagtree_repo_root() -> pathlib.Path | None: + marker = pathlib.Path("python") / "test" / "unit" / "language" / "test_warp_specialization.py" + for parent in pathlib.Path(__file__).resolve().parents: + if (parent / marker).exists(): + return parent + return None + + +def cfg_extra(cfg: GemmConfig) -> dict[str, object]: + return { + "BM": cfg.bm, + "BN": cfg.bn, + "BK": cfg.bk, + "NUM_STAGES": cfg.num_stages, + "GROUP_SIZE_M": cfg.group_size_m, + "PRODUCER_NUM_WARPS": cfg.producer_num_warps, + "wgmma_pipeline": cfg.wgmma_pipeline, + } + + +def run_tle_compare_variant( + problem: Problem, + a: torch.Tensor, + b: torch.Tensor, + ref: torch.Tensor | None, + warmup: int, + rep: int, + cfg: GemmConfig, + *, + persistent: bool, + dump_summary: bool, + cuda_graph: bool = False, +) -> dict[str, object]: + split = "persistent_split_m" if persistent else "nonpersistent_split_m" + variant = f"flagtree.tle.warp_specialize.{split}.{cfg.label()}" + matmul_fn = tle_ws_persistent_matmul if persistent else tle_ws_nonpersistent_matmul + try: + c, kernel = matmul_fn( + a, + b, + bm=cfg.bm, + bn=cfg.bn, + bk=cfg.bk, + group_size_m=cfg.group_size_m, + num_stages=cfg.num_stages, + wgmma_pipeline=cfg.wgmma_pipeline, + producer_num_warps=cfg.producer_num_warps, + ) + torch.cuda.synchronize() + ttgir = kernel.asm["ttgir"] + has_warp_specialize = "ttg.warp_specialize" in ttgir + has_wgmma = "ttng.warp_group_dot" in ttgir + if ref is not None: + torch.testing.assert_close(c, ref, atol=2e-2, rtol=2e-2) + if dump_summary: + print( + f"{problem.m}x{problem.n}x{problem.k} {variant}: " + f"has_warp_specialize={has_warp_specialize} has_wgmma={has_wgmma} " + f"asm_keys={','.join(kernel.asm.keys())}", + file=sys.stderr, + ) + + bench_out = torch.empty((problem.m, problem.n), device=a.device, dtype=torch.float16) + + def run(): + matmul_fn( + a, + b, + bm=cfg.bm, + bn=cfg.bn, + bk=cfg.bk, + group_size_m=cfg.group_size_m, + num_stages=cfg.num_stages, + wgmma_pipeline=cfg.wgmma_pipeline, + producer_num_warps=cfg.producer_num_warps, + out=bench_out, + )[0] + + ms, p20, p80 = bench_ms(run, warmup, rep, cuda_graph=cuda_graph) + return make_row( + variant, + problem, + ms, + p20, + p80, + cfg.bm, + cfg.bn, + cfg.bk, + cfg.num_stages, + cfg.group_size_m, + cfg.wgmma_pipeline, + cfg.producer_num_warps, + has_warp_specialize=has_warp_specialize, + has_wgmma=has_wgmma, + cuda_graph=cuda_graph, + ) + except Exception as exc: + return make_error_row( + variant, + problem, + f"compile/run failed: {type(exc).__name__}: {exc}", + cfg_extra(cfg), + ) + + +def run_native_ws_tma_variants(problem: Problem, warmup: int, rep: int, *, + cuda_graph: bool = False) -> list[dict[str, object]]: + repo_root = find_flagtree_repo_root() + if repo_root is None: + return [ + make_error_row( + "flagtree.native_tl_range_ws_tma.skip", + problem, + "cannot find FlagTree python/test/unit/language/test_warp_specialization.py", + ) + ] + + sys.path.insert(0, str(repo_root / "python")) + try: + module = import_from_path( + "flagtree_native_ws_test", + repo_root / "python" / "test" / "unit" / "language" / "test_warp_specialization.py", + ) + except Exception as exc: + return [ + make_error_row( + "flagtree.native_tl_range_ws_tma.skip", + problem, + f"cannot import native WS test kernel: {type(exc).__name__}: {exc}", + ) + ] + + triton.set_allocator(alloc_fn) + rows: list[dict[str, object]] = [] + configs = [ + {"BLOCK_SIZE_M": 128, "BLOCK_SIZE_N": 128, "BLOCK_SIZE_K": 64, "num_stages": 3, "num_warps": 4}, + {"BLOCK_SIZE_M": 128, "BLOCK_SIZE_N": 256, "BLOCK_SIZE_K": 64, "num_stages": 3, "num_warps": 4}, + ] + for cfg in configs: + smem_bytes = (cfg["num_stages"] * cfg["BLOCK_SIZE_K"] * + (cfg["BLOCK_SIZE_M"] + cfg["BLOCK_SIZE_N"]) + cfg["BLOCK_SIZE_M"] * cfg["BLOCK_SIZE_N"]) * 2 + if smem_bytes > 228 * 1024: + continue + + try: + a = torch.randn((problem.m, problem.k), device=DEVICE, dtype=torch.float16).contiguous() + b = torch.randn((problem.n, problem.k), device=DEVICE, dtype=torch.float16).contiguous() + c = torch.empty((problem.m, problem.n), device=DEVICE, dtype=torch.float16) + grid = lambda meta: (triton.cdiv(problem.m, meta["BLOCK_SIZE_M"]) * triton.cdiv( + problem.n, meta["BLOCK_SIZE_N"]), ) + + def launch(): + return module.matmul_tma_ws_kernel[grid]( + a, + b, + c, + *a.stride(), + *b.stride(), + *c.stride(), + problem.m, + problem.n, + problem.k, + cfg["num_stages"], + cfg["BLOCK_SIZE_M"], + cfg["BLOCK_SIZE_N"], + cfg["BLOCK_SIZE_K"], + 8, + num_warps=cfg["num_warps"], + USE_FP8=False, + ) + + kernel = launch() + torch.cuda.synchronize() + ttgir = kernel.asm["ttgir"] + + def run(): + launch() + + ms, p20, p80 = bench_ms(run, warmup, rep, cuda_graph=cuda_graph) + rows.append( + make_native_row( + "flagtree.native_tl_range_ws_tma", + problem, + ms, + p20, + p80, + cfg, + has_warp_specialize="ttg.warp_specialize" in ttgir, + has_wgmma="ttng.warp_group_dot" in ttgir, + cuda_graph=cuda_graph, + )) + except Exception as exc: + rows.append( + make_error_row( + "flagtree.native_tl_range_ws_tma", + problem, + f"compile/run failed: {type(exc).__name__}: {exc}", + cfg, + )) + return rows + + +def run_compare(args: argparse.Namespace, problems: list[Problem]) -> list[dict[str, object]]: + configs = args.compare_config or default_compare_configs(args) + rows: list[dict[str, object]] = [] + for problem in problems: + a = torch.randn((problem.m, problem.k), device=DEVICE, dtype=torch.float16).contiguous() + b = torch.randn((problem.k, problem.n), device=DEVICE, dtype=torch.float16).contiguous() + problem_rows: list[dict[str, object]] = [] + + if not args.no_cublas: + problem_rows.append( + run_torch_matmul_baseline( + "cuBLAS", + problem, + a, + b, + args.warmup, + args.rep, + use_out=False, + cuda_graph=args.cuda_graph, + )) + if args.include_torch: + problem_rows.append( + run_torch_matmul_baseline( + "torch.matmul.out", + problem, + a, + b, + args.warmup, + args.rep, + use_out=True, + cuda_graph=args.cuda_graph, + )) + + ref = torch.matmul(a, b) if args.check else None + for cfg in configs: + problem_rows.append( + run_tle_compare_variant( + problem, + a, + b, + ref, + args.warmup, + args.rep, + cfg, + persistent=True, + dump_summary=args.dump_summary, + cuda_graph=args.cuda_graph, + )) + if not args.no_nonpersistent: + for cfg in configs: + problem_rows.append( + run_tle_compare_variant( + problem, + a, + b, + ref, + args.warmup, + args.rep, + cfg, + persistent=False, + dump_summary=args.dump_summary, + cuda_graph=args.cuda_graph, + )) + + if not args.no_native: + problem_rows.extend(run_native_ws_tma_variants(problem, args.warmup, args.rep, cuda_graph=args.cuda_graph)) + + if not args.no_cublas: + add_cublas_speedups(problem_rows) + rows.extend(problem_rows) + return rows + + +def write_rows(rows: Iterable[dict[str, object]], out: str | None) -> None: + rows = list(rows) + fields = [] + for row in rows: + for key in row: + if key not in fields: + fields.append(key) + if out: + with open(out, "w", newline="") as f: + writer = csv.DictWriter(f, fieldnames=fields) + writer.writeheader() + writer.writerows(rows) + writer = csv.DictWriter(__import__("sys").stdout, fieldnames=fields) + writer.writeheader() + writer.writerows(rows) + + +def main() -> None: + parser = argparse.ArgumentParser() + parser.add_argument("--shape", action="append", type=parse_problem, default=[]) + parser.add_argument("--warmup", type=int, default=25) + parser.add_argument("--rep", type=int, default=100) + parser.add_argument("--cuda-graph", action="store_true", + help="benchmark by capturing the workload once and timing CUDA Graph replay") + parser.add_argument("--out", default=None) + parser.add_argument("--check", action="store_true") + parser.add_argument("--compare", action="store_true", + help="run cuBLAS, TLE persistent/non-persistent configs, and native WS TMA configs") + parser.add_argument("--compare-config", action="append", type=parse_gemm_config, default=[], + help="compare config BMxBNxBKxSTAGES[xGROUP_SIZE_M[xPRODUCER_WARPS]][:pipe|sync]") + parser.add_argument("--include-cublas", action="store_true", + help="include torch.matmul(a, b), matching the TLX cuBLAS baseline") + parser.add_argument("--include-torch", action="store_true", + help="include a preallocated torch.matmul(a, b, out=c) baseline") + parser.add_argument("--no-cublas", action="store_true", help="omit cuBLAS baseline in --compare mode") + parser.add_argument("--no-native", action="store_true", help="omit native tl.range WS TMA rows in --compare mode") + parser.add_argument("--no-nonpersistent", action="store_true", + help="omit TLE non-persistent split-M rows in --compare mode") + parser.add_argument("--dump-summary", action="store_true", + help="print a compact kernel TTGIR feature summary to stderr") + parser.add_argument("--bm", type=int, default=128) + parser.add_argument("--bn", type=int, default=128) + parser.add_argument("--bk", type=int, default=64) + parser.add_argument("--group-size-m", type=int, default=8) + parser.add_argument("--num-stages", type=int, default=3) + parser.add_argument("--no-wgmma-pipeline", action="store_true") + parser.add_argument("--producer-num-warps", type=int, default=4) + parser.add_argument("--non-persistent", action="store_true") + args = parser.parse_args() + + if not torch.cuda.is_available() or torch.cuda.get_device_capability()[0] < 9: + raise RuntimeError("Hopper or newer CUDA GPU is required") + if args.producer_num_warps % 4 != 0: + parser.error("--producer-num-warps must be a multiple of 4 for TLE warp_specialize") + + problems = args.shape or [Problem(4096, 4096, 4096), Problem(8192, 8192, 512)] + if args.compare: + write_rows(run_compare(args, problems), args.out) + return + + rows = [] + for problem in problems: + a = torch.randn((problem.m, problem.k), device=DEVICE, dtype=torch.float16).contiguous() + b = torch.randn((problem.k, problem.n), device=DEVICE, dtype=torch.float16).contiguous() + problem_rows: list[dict[str, object]] = [] + + if args.compare or args.include_cublas: + problem_rows.append( + run_torch_matmul_baseline( + "cuBLAS", + problem, + a, + b, + args.warmup, + args.rep, + use_out=False, + cuda_graph=args.cuda_graph, + )) + if args.include_torch: + problem_rows.append( + run_torch_matmul_baseline( + "torch.matmul.out", + problem, + a, + b, + args.warmup, + args.rep, + use_out=True, + cuda_graph=args.cuda_graph, + )) + + wgmma_pipeline = not args.no_wgmma_pipeline + matmul_fn = tle_ws_nonpersistent_matmul if args.non_persistent else tle_ws_persistent_matmul + c, kernel = matmul_fn( + a, + b, + bm=args.bm, + bn=args.bn, + bk=args.bk, + group_size_m=args.group_size_m, + num_stages=args.num_stages, + wgmma_pipeline=wgmma_pipeline, + producer_num_warps=args.producer_num_warps, + ) + torch.cuda.synchronize() + ttgir = kernel.asm["ttgir"] + has_warp_specialize = "ttg.warp_specialize" in ttgir + has_wgmma = "ttng.warp_group_dot" in ttgir + assert has_warp_specialize + assert has_wgmma + + if args.dump_summary: + variant = "nonpersistent" if args.non_persistent else "persistent" + print( + f"{problem.m}x{problem.n}x{problem.k} {variant}: " + f"has_warp_specialize={has_warp_specialize} has_wgmma={has_wgmma} " + f"asm_keys={','.join(kernel.asm.keys())}", + file=sys.stderr, + ) + + if args.check: + torch.testing.assert_close(c, torch.matmul(a, b), atol=2e-2, rtol=2e-2) + + bench_out = torch.empty((problem.m, problem.n), device=a.device, dtype=torch.float16) + + def run(): + matmul_fn( + a, + b, + bm=args.bm, + bn=args.bn, + bk=args.bk, + group_size_m=args.group_size_m, + num_stages=args.num_stages, + wgmma_pipeline=wgmma_pipeline, + producer_num_warps=args.producer_num_warps, + out=bench_out, + )[0] + + ms, p20, p80 = bench_ms(run, args.warmup, args.rep, cuda_graph=args.cuda_graph) + variant = "flagtree.tle.warp_specialize.nonpersistent_split_m" if args.non_persistent else \ + "flagtree.tle.warp_specialize.persistent_split_m" + problem_rows.append( + make_row( + variant, + problem, + ms, + p20, + p80, + args.bm, + args.bn, + args.bk, + args.num_stages, + args.group_size_m, + wgmma_pipeline, + args.producer_num_warps, + has_warp_specialize=has_warp_specialize, + has_wgmma=has_wgmma, + cuda_graph=args.cuda_graph, + )) + if args.compare or args.include_cublas: + add_cublas_speedups(problem_rows) + rows.extend(problem_rows) + + write_rows(rows, args.out) + + +if __name__ == "__main__": + main()