Merge OpenAI Triton commit 0e94b6c (#4392)

This PR change the Triton base from
75fe113
to 0e94b6c (May 26).
Pass rate: 97.23% -> 97.23%

Author: whitneywhtsang

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();
@@ -113,7 +115,7 @@ inline void registerTritonDialects(mlir::DialectRegistry &registry) {
   mlir::registerTritonAMDFoldTrueCmpI();
 
   // NVWS passes
-  mlir::registerNVWSTransformsPasses();
+  mlir::triton::registerNVWSTransformsPasses();
 
   // NVGPU transform passes
   mlir::registerNVHopperTransformsPasses();

include/triton/Conversion/TritonGPUToLLVM/TargetInfoBase.h

@@ -94,7 +94,8 @@ class TargetInfoBase {
 
   virtual bool supportVectorizedAtomics() const = 0;
 
-  virtual bool supportLdStMatrix() const = 0;
+  virtual bool supportLdMatrix() const { return false; }
+  virtual bool supportStMatrix() const { return false; }
 
   // Annotate target specific information to local store operations during
   // lowering to LLVM.

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/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";
 

include/triton/Target/LLVMIR/Passes.h

@@ -5,10 +5,11 @@
 
 namespace mlir {
 
-/// Create a pass to add DIScope
-std::unique_ptr<Pass> createLLVMDIScopePass();
+// Generate the pass class declarations.
+#define GEN_PASS_DECL
+#include "triton/Target/LLVMIR/Passes.h.inc"
 
-/// Generate the code for registering conversion passes.
+// Generate the code for registering conversion passes.
 #define GEN_PASS_REGISTRATION
 #include "triton/Target/LLVMIR/Passes.h.inc"
 

include/triton/Target/LLVMIR/Passes.td

@@ -8,8 +8,6 @@ def LLVMDIScope: Pass<"enable-line-info", "mlir::ModuleOp"> {
   let description = [{
     This pass materializes line mapping information for LLVM IR dialect operations.
   }];
-
-  let constructor = "mlir::createLLVMDIScopePass()";
 }
 
 #endif

include/triton/Tools/LayoutUtils.h

@@ -125,6 +125,11 @@ std::optional<ColumnAction> regPermForDivideLeft(const LinearLayout &A,
 // such that action.apply(A) has the broadcasted registers removed
 ColumnAction actionRemoveBroadcastedRegs(const LinearLayout &layout);
 
+// For a layout A with A.hasInDim(kReg), repeat the values so that they have
+// the same broadcasting as layout
+SmallVector<Value> broadcastAs(const SmallVector<Value> &values,
+                               const LinearLayout &layout);
+
 // Compute the supremum of two lists.
 // Error out if the supremum does not exist (e.g. [a, b] and [b, a]).
 // If the supremum is not unique, we return the first list first

include/triton/Tools/LinearLayout.h

@@ -825,6 +825,9 @@ class ColumnAction {
   //    [x[0], x[4], x[1], x[5], x[2], x[6], x[3], x[7]]
   SmallVector<Value> apply(ValueRange values) const;
 
+  // Inverse of the action
+  ColumnAction inverse() const;
+
   std::string toString() const;
 };