[TLE] barrier & wgmma pipeline with user promising support in hopper

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<TTG_CopyOpType, "", [MemRead<GlobalMemory>, MemWrite<GlobalMemory>]>:$src,
+    Arg<TTG_CopyOpType, "", [MemRead<GlobalMemory>, MemWrite<GlobalMemory>]>:$dst,
+    Variadic<I32>:$indices,
+    Optional<TTG_MemDescType>:$barrier,
+    OptionalAttr<I32Attr>:$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

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

lib/Conversion/TritonToTritonGPU/TritonToTritonGPUPass.cpp

@@ -949,6 +949,7 @@ void populateTleRawPatterns(TritonGPUTypeConverter &typeConverter,
            TleInsertTileOpPattern, GenericOpPattern<tle::LocalPointersOp>,
            GenericOpPattern<tle::RemotePointersOp>,
            GenericOpPattern<tle::ExclusiveCumsumOp>,
+           GenericOpPattern<tle::WGMMAOp>, GenericOpPattern<tle::WGMMAWaitOp>,
            GenericOpPattern<tle::DistributedBarrierOp>,
            GenericOpPattern<tle::YieldOp>,
            GenericOpPattern<tle::ExtractAllocatedPtrOp>,

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

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<StringAttr>(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<scf::ForOp> 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<PipelinePass> {
     // 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());

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
 

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<ttng::WarpGroupDotOp, 8> dots;
+  forOp.getBody()->walk([&](ttng::WarpGroupDotOp dot) {
+    if (dot->getParentOfType<scf::ForOp>() == 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<ttng::WarpGroupDotCommitOp>(next))
+      continue;
+
+    builder.setInsertionPointAfter(dot);
+    ttng::WarpGroupDotCommitOp::create(builder, dot.getLoc());
+  }
+}
+
+} // namespace mlir::triton::gpu::detail
+#endif // __TLE__

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<ttng::WarpGroupDotWaitOp>(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<Value, 8> &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<ttng::WarpGroupDotWaitOp>(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<Value, 8> visited;
+  return isMaterializedByWaitZero(value, visited);
+}
+
+static std::optional<unsigned>
+findYieldOperandForPendingGroup(scf::YieldOp yield,
+                                const PendingSharedWgmmaGroup &pending) {
+  std::optional<unsigned> yieldIndex;
+  for (auto indexed : llvm::enumerate(yield.getOperands())) {
+    unsigned index = static_cast<unsigned>(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<PendingSharedWgmmaGroup> pendingGroups) {
+  auto forOp = dyn_cast<scf::ForOp>(yield->getParentOp());
+  if (!forOp || forOp.getBody() != yield->getBlock())
+    return false;
+  if (pendingGroups.empty())
+    return false;
+
+  llvm::SmallDenseSet<unsigned, 4> carriedYieldIndices;
+  for (const PendingSharedWgmmaGroup &pending : pendingGroups) {
+    std::optional<unsigned> 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<scf::YieldOp>(terminator)) {
-    if (!pendingGroups.empty())
+  if (auto yield = dyn_cast<scf::YieldOp>(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())