Merge commit '9f88c7f8ba9b7c56f2a1aff55163ecacf820df36'

Author: whitneywhtsang

.github/workflows/build-macos.yml

@@ -13,7 +13,7 @@ jobs:
     strategy:
       matrix:
         runner: ${{ fromJson(inputs.matrix) }}
-    timeout-minutes: 40
+    timeout-minutes: 60
     env:
       RUNNER_TYPE: ${{ matrix.runner[0] }}
     name: Build MacOS

bin/CMakeLists.txt

@@ -13,6 +13,7 @@ target_link_libraries(triton-opt PRIVATE
   ${triton_libs}
   # tests
   TritonTestAnalysis
+  TritonTestDialect
   TritonAMDGPUTestAnalysis
   # MLIR core
   MLIROptLib
@@ -32,6 +33,7 @@ target_link_libraries(triton-reduce PRIVATE
   ${triton_libs}
   # tests
   TritonTestAnalysis
+  TritonTestDialect
   TritonAMDGPUTestAnalysis
   # MLIR core
   MLIRReduceLib
@@ -50,6 +52,7 @@ target_link_libraries(triton-lsp PRIVATE
   ${triton_libs}
   # tests
   TritonTestAnalysis
+  TritonTestDialect
   TritonAMDGPUTestAnalysis
   # MLIR core
   MLIRLspServerLib
@@ -88,5 +91,6 @@ target_link_libraries(triton-tensor-layout PRIVATE
   ${conversion_libs}
   ${dialect_libs}
   TritonTestAnalysis
+  TritonTestDialect
   TritonAMDGPUTestAnalysis
   )

bin/RegisterTritonDialects.h

@@ -54,6 +54,7 @@ void registerTestLivenessPass();
 void registerTestMembarPass();
 void registerTestAMDGPUMembarPass();
 void registerTestTritonAMDGPURangeAnalysis();
+void registerTestLoopPeelingPass();
 } // namespace test
 } // namespace mlir
 
@@ -68,6 +69,7 @@ inline void registerTritonDialects(mlir::DialectRegistry &registry) {
   mlir::test::registerTestAllocationPass();
   mlir::test::registerTestLivenessPass();
   mlir::test::registerTestMembarPass();
+  mlir::test::registerTestLoopPeelingPass();
   mlir::test::registerTestAMDGPUMembarPass();
   mlir::test::registerTestTritonAMDGPURangeAnalysis();
   mlir::triton::registerConvertTritonToTritonGPUPass();

include/triton/Dialect/Triton/IR/Utility.h

@@ -177,6 +177,9 @@ template <typename T> auto seq(T start, T end, T step) {
 Value getPredMask(RewriterBase &rewriter, Type typeLike, Value currentMask,
                   Value pred);
 
+// Get the value of the induction variable at the end of the loop.
+Value getLastInductionValue(OpBuilder &b, scf::ForOp loop);
+
 MakeTensorPtrOp getMakeTensorPtrOp(Value v);
 
 } // namespace triton

include/triton/Dialect/Triton/Transforms/LoopPeeling.h

@@ -0,0 +1,18 @@
+#ifndef TRITON_DIALECT_TRITON_TRANSFORMS_LOOP_PEELING_H_
+#define TRITON_DIALECT_TRITON_TRANSFORMS_LOOP_PEELING_H_
+
+#include "mlir/Dialect/SCF/IR/SCF.h"
+
+namespace mlir {
+namespace triton {
+
+// Peel the single last iteration of the loop.
+void peelLoopEpilogue(
+    scf::ForOp forOp,
+    function_ref<Operation *(RewriterBase &, Operation *, bool)>
+        processPeeledOp = nullptr);
+
+} // namespace triton
+} // namespace mlir
+
+#endif // TRITON_DIALECT_TRITON_TRANSFORMS_LOOP_PEELING_H_

include/triton/Dialect/TritonGPU/IR/TritonGPUAttrDefs.td

@@ -437,7 +437,11 @@ def NVMMASharedEncodingAttr :
         } else {
           swizzlingByteWidth = 0;
         }
-        if (shapePerCTA.size() < 2 || shapePerCTA[order[1]] < 8) {
+        int flattenOutterDim = 1;
+        for (int i = 1; i < shapePerCTA.size(); i++) {
+          flattenOutterDim *= shapePerCTA[order[i]];
+        }
+        if (shapePerCTA.size() < 2 || flattenOutterDim < 8) {
           swizzlingByteWidth = 0;
         }
         bool transposed = order[0] == 0;

include/triton/Dialect/TritonGPU/IR/TritonGPUOps.td

@@ -337,6 +337,24 @@ def TTG_PredicateStageOp: TTG_Op<"predicate_stage",
   let assemblyFormat = "$iv `,` $ub `,` $step `maxStage` $maxStage `stage` $stage attr-dict `:` type($iv) `->` type($result)";
 }
 
+def TTG_MaskOp: TTG_Op<"mask",
+                       [SingleBlock]> {
+    let summary = "mask op for pipelining";
+    let arguments = (ins I1:$pred);
+    let results = (outs Variadic<AnyType>:$result);
+    let regions = (region SizedRegion<1>:$region);
+    let builders = [
+        OpBuilder<(ins "Value":$pred)>,
+    ];
+}
+
+def TTG_MaskReturnOp: TTG_Op<"mask.return",
+                             [HasParent<"MaskOp">, Pure, Terminator, ReturnLike]> {
+    let summary = "terminator for mask operator";
+    let arguments = (ins Variadic<AnyType>:$result);
+    let assemblyFormat = "$result attr-dict `:` type($result)";
+}
+
 def TTG_Fp4ToFpOp : TTG_Op<"fp4_to_fp", [Pure]> {
   let summary = "Upcast fp4 (e2m1) to fp";
 
@@ -450,7 +468,9 @@ def TTG_WarpSpecializeOp : TTG_Op<"warp_specialize", [
   let builders = [
     OpBuilder<(ins "TypeRange":$resultTypes,
                    "ArrayRef<int32_t>":$partitionNumWarps,
-                   "unsigned":$numPartitionRegions)>
+                   "unsigned":$numPartitionRegions)>,
+    OpBuilder<(ins "TypeRange":$resultTypes, "ValueRange":$explicitCaptures,
+                   "ArrayRef<int32_t>":$partitionNumWarps)>,
   ];
 
   let hasVerifier = 1;

lib/Dialect/Triton/IR/Utility.cpp

@@ -1,5 +1,6 @@
 #include "triton/Dialect/Triton/IR/Utility.h"
 #include "mlir/Dialect/ControlFlow/IR/ControlFlowOps.h"
+#include "mlir/Dialect/SCF/IR/SCF.h"
 #include "triton/Dialect/Triton/IR/Dialect.h"
 
 using namespace mlir;
@@ -90,3 +91,15 @@ tt::MakeTensorPtrOp tt::getMakeTensorPtrOp(Value v) {
   }
   llvm_unreachable("Unable to getMakeTensorPtr()");
 }
+
+Value tt::getLastInductionValue(OpBuilder &b, scf::ForOp loop) {
+  Location loc = loop.getLoc();
+  // (ub - lb -1) // step * step + lb
+  Value diff =
+      b.create<arith::SubIOp>(loc, loop.getUpperBound(), loop.getLowerBound());
+  diff = b.create<arith::SubIOp>(
+      loc, diff, b.create<arith::ConstantOp>(loc, b.getI32IntegerAttr(1)));
+  Value ceilStep = b.create<arith::MulIOp>(
+      loc, b.create<arith::DivSIOp>(loc, diff, loop.getStep()), loop.getStep());
+  return b.create<arith::AddIOp>(loc, ceilStep, loop.getLowerBound());
+}

lib/Dialect/Triton/Transforms/CMakeLists.txt

@@ -6,6 +6,7 @@ add_triton_library(TritonTransforms
   Combine.cpp
   LoopAwareCSE.cpp
   LoopInvariantCodeMotion.cpp
+  LoopPeeling.cpp
   LoopUnroll.cpp
   ReorderBroadcast.cpp
   RewriteTensorPointer.cpp
@@ -20,5 +21,7 @@ add_triton_library(TritonTransforms
   LINK_LIBS PUBLIC
   MLIRPass
   MLIRTransformUtils
+  MLIRTransforms
+  MLIRSCFToControlFlow
   TritonIR
 )

lib/Dialect/Triton/Transforms/LoopPeeling.cpp

@@ -0,0 +1,68 @@
+#include "triton/Dialect/Triton/Transforms/LoopPeeling.h"
+#include "mlir/Dialect/SCF/IR/SCF.h"
+#include "mlir/Pass/Pass.h"
+#include "triton/Dialect/Triton/IR/Utility.h"
+
+using namespace mlir;
+
+namespace mlir {
+namespace triton {
+
+void peelLoopEpilogue(
+    scf::ForOp forOp,
+    function_ref<Operation *(RewriterBase &, Operation *, bool)>
+        processPeeledOp) {
+  SmallVector<Operation *> loopBodyOps;
+  IRRewriter rewriter(forOp);
+  Location loc = forOp.getLoc();
+  Type type = forOp.getStep().getType();
+
+  // Fetch loop bounds and step
+  Value lowerBound = forOp.getLowerBound();
+  Value upperBound = forOp.getUpperBound();
+  Value step = forOp.getStep();
+  Value newUpperBound = rewriter.create<arith::SubIOp>(loc, upperBound, step);
+
+  rewriter.setInsertionPointAfter(forOp);
+  Value lastIV = getLastInductionValue(rewriter, forOp);
+
+  auto cond = rewriter.create<arith::CmpIOp>(loc, arith::CmpIPredicate::slt,
+                                             lowerBound, upperBound);
+
+  // Create an if op to execute the peeled iteration
+  IRMapping map;
+  map.map(forOp.getRegionIterArgs(), forOp.getResults());
+  map.map(forOp.getInductionVar(), lastIV);
+  auto ifOp = rewriter.create<scf::IfOp>(loc, forOp.getResultTypes(), cond,
+                                         /*hasElse=*/true);
+  ifOp.getThenRegion().front().erase();
+  forOp.getBodyRegion().cloneInto(&ifOp.getThenRegion(), map);
+  rewriter.setInsertionPointToStart(&ifOp.getElseRegion().front());
+  rewriter.create<scf::YieldOp>(loc, forOp.getResults());
+
+  forOp->replaceUsesWithIf(ifOp, [&](OpOperand &operand) {
+    return !ifOp->isAncestor(operand.getOwner());
+  });
+
+  forOp.getUpperBoundMutable().assign(newUpperBound);
+
+  if (processPeeledOp) {
+    for (auto &op :
+         llvm::make_early_inc_range(forOp.getBody()->without_terminator())) {
+      Operation *newOp = processPeeledOp(rewriter, &op, /*isEpilogue=*/false);
+      if (newOp && newOp != &op) {
+        op.replaceAllUsesWith(newOp);
+      }
+    }
+    for (auto &op : llvm::make_early_inc_range(
+             ifOp.getThenRegion().front().without_terminator())) {
+      Operation *newOp = processPeeledOp(rewriter, &op, /*isEpilogue=*/true);
+      if (newOp && newOp != &op) {
+        op.replaceAllUsesWith(newOp);
+      }
+    }
+  }
+}
+
+} // namespace triton
+} // namespace mlir