Merge OpenAI Triton commit 22b1a44 (#5229)

This PR change the Triton base from
407b8a3 to
22b1a44 (Sep 25).
Pass rate: 97.26%->92.66%

Author: anmyachev

.github/workflows/integration-tests-amd.yml

@@ -173,9 +173,7 @@ jobs:
           LD_PRELOAD=$CLANG_ASAN_LIB:$HIP_ASAN_LIB python3 -m pytest -s test_address_sanitizer.py
       - name: Run regression tests
         run: |
-          # Reenable test_functional_regression.py once it's fixed
-          cd python/test/regression
-          python3 -m pytest -s -n 8 ./test_cast_matmul.py
+          make test-regression
       - name: Run microbenchmark tests
         run: |
           python3 python/test/microbenchmark/launch_overhead.py

include/triton/Analysis/AxisInfo.h

@@ -161,10 +161,13 @@ class AxisInfo {
 
   std::optional<int64_t> getConstantValue() const { return constantValue; }
 
-  template <class T>
-  static void
-  initPessimisticStateFromFunc(int argNumber, T funcOp, DimVectorT *contiguity,
-                               DimVectorT *divisibility, DimVectorT *constancy);
+  static void initPessimisticStateFromFunc(int argNumber,
+                                           FunctionOpInterface funcOp,
+                                           DimVectorT *contiguity,
+                                           DimVectorT *divisibility,
+                                           DimVectorT *constancy);
+
+  static void initDimVectorFromHint(Attribute attr, DimVectorT *vec);
 
   bool operator==(const AxisInfo &other) const {
     return contiguity == other.contiguity &&

include/triton/Dialect/TritonGPU/Transforms/Passes.td

@@ -176,24 +176,6 @@ def TritonGPUPartitionScheduling : Pass<"tritongpu-partition-scheduling", "mlir:
   }];
 }
 
-def TritonGPULoadMMASpecialization : Pass<"tritongpu-load-mma-specialization", "mlir::ModuleOp"> {
-  let summary = "load MMA specialization";
-
-  let description = [{
-    The `tritongpu-load-mma-specialization` pass looks for matmul loops in the
-    module and attempts to create a partition schedule, separating async loads
-    and async MMAs into separate partitions.
-  }];
-
-  let dependentDialects = ["mlir::triton::gpu::TritonGPUDialect",
-                           "mlir::triton::nvidia_gpu::TritonNvidiaGPUDialect"];
-
-  let options = [
-    Option<"numStages", "num-stages", "int32_t", /*default*/"3",
-           "number of pipeline stages">
-  ];
-}
-
 def TritonGPUF32DotTC : Pass<"tritongpu-F32DotTC", "mlir::ModuleOp"> {
   let summary = "3xTF32 trick";
 

lib/Analysis/AxisInfo.cpp

@@ -1084,18 +1084,12 @@ LogicalResult AxisInfoAnalysis::visitOperation(
   auto newContiguity = curr.getContiguity();
   auto newDivisibility = curr.getDivisibility();
   auto newConstancy = curr.getConstancy();
-  if (Attribute attr = op->getDiscardableAttr("tt.contiguity")) {
-    auto vals = cast<DenseElementsAttr>(attr).getValues<int>();
-    newContiguity = AxisInfo::DimVectorT(vals.begin(), vals.end());
-  }
-  if (Attribute attr = op->getDiscardableAttr("tt.divisibility")) {
-    auto vals = cast<DenseElementsAttr>(attr).getValues<int>();
-    newDivisibility = AxisInfo::DimVectorT(vals.begin(), vals.end());
-  }
-  if (Attribute attr = op->getDiscardableAttr("tt.constancy")) {
-    auto vals = cast<DenseElementsAttr>(attr).getValues<int>();
-    newConstancy = AxisInfo::DimVectorT(vals.begin(), vals.end());
-  }
+  AxisInfo::initDimVectorFromHint(op->getDiscardableAttr("tt.contiguity"),
+                                  &newContiguity);
+  AxisInfo::initDimVectorFromHint(op->getDiscardableAttr("tt.divisibility"),
+                                  &newDivisibility);
+  AxisInfo::initDimVectorFromHint(op->getDiscardableAttr("tt.constancy"),
+                                  &newConstancy);
   curr = AxisInfo(newContiguity, newDivisibility, newConstancy,
                   curr.getConstantValue());
   // join all lattice elements
@@ -1140,25 +1134,29 @@ void AxisInfoAnalysis::visitWarpSpecializeExplicitCaptures(
 
 } // anonymous namespace
 
-template <class T>
-void AxisInfo::initPessimisticStateFromFunc(int argNumber, T funcOp,
+void AxisInfo::initPessimisticStateFromFunc(int argNumber,
+                                            FunctionOpInterface funcOp,
                                             DimVectorT *contiguity,
                                             DimVectorT *divisibility,
                                             DimVectorT *constancy) {
-  // liast of attributes that we care about
+  // list of attributes that we care about
   SmallVector<std::pair<DimVectorT *, std::string>> retVecs;
   retVecs.push_back({contiguity, "tt.contiguity"});
   retVecs.push_back({divisibility, "tt.divisibility"});
   retVecs.push_back({constancy, "tt.constancy"});
   // initialize attributes one by one
   for (auto [vec, attrName] : retVecs) {
     Attribute attr = funcOp.getArgAttr(argNumber, attrName);
-    if (auto int_attr = dyn_cast_or_null<IntegerAttr>(attr))
-      *vec = DimVectorT(contiguity->size(), int_attr.getValue().getZExtValue());
-    if (auto dense_attr = dyn_cast_or_null<DenseElementsAttr>(attr)) {
-      auto vals = dense_attr.getValues<int>();
-      *vec = DimVectorT(vals.begin(), vals.end());
-    }
+    AxisInfo::initDimVectorFromHint(attr, vec);
+  }
+}
+
+void AxisInfo::initDimVectorFromHint(Attribute attr, DimVectorT *vec) {
+  if (auto int_attr = dyn_cast_or_null<IntegerAttr>(attr))
+    *vec = DimVectorT(1, int_attr.getValue().getZExtValue());
+  if (auto dense_attr = dyn_cast_or_null<DenseElementsAttr>(attr)) {
+    auto vals = dense_attr.getValues<int>();
+    *vec = DimVectorT(vals.begin(), vals.end());
   }
 }
 
@@ -1202,18 +1200,12 @@ void AxisInfo::initPessimisticStateFromFunc(int argNumber, T funcOp,
     }
     // Other operations are conservatively initialized with the lowest possible
     // divisibility, contiguity, and constancy unless they have specified.
-    if (Attribute attr = op->getDiscardableAttr("tt.divisibility")) {
-      auto vals = cast<DenseElementsAttr>(attr).getValues<int>();
-      knownDivisibility = DimVectorT(vals.begin(), vals.end());
-    }
-    if (Attribute attr = op->getDiscardableAttr("tt.contiguity")) {
-      auto vals = cast<DenseElementsAttr>(attr).getValues<int>();
-      knownContiguity = DimVectorT(vals.begin(), vals.end());
-    }
-    if (Attribute attr = op->getDiscardableAttr("tt.constancy")) {
-      auto vals = cast<DenseElementsAttr>(attr).getValues<int>();
-      knownConstancy = DimVectorT(vals.begin(), vals.end());
-    }
+    AxisInfo::initDimVectorFromHint(op->getDiscardableAttr("tt.divisibility"),
+                                    &knownDivisibility);
+    AxisInfo::initDimVectorFromHint(op->getDiscardableAttr("tt.contiguity"),
+                                    &knownContiguity);
+    AxisInfo::initDimVectorFromHint(op->getDiscardableAttr("tt.constancy"),
+                                    &knownConstancy);
   }
 
   return AxisInfo(knownContiguity, knownDivisibility, knownConstancy);

lib/Dialect/TritonGPU/Transforms/CMakeLists.txt

@@ -29,7 +29,6 @@ add_triton_library(TritonGPUTransforms
   Utility.cpp
   LayoutPropagationUtility.cpp
   WarpSpecialization/AutomaticWarpSpecialization.cpp
-  WarpSpecialization/LoadMMASpecialization.cpp
   WarpSpecialization/Partition.cpp
   WarpSpecialization/OptimizePartitionWarps.cpp
   WarpSpecialization/PartitionBuilder.cpp

lib/Dialect/TritonGPU/Transforms/WarpSpecialization/AutomaticWarpSpecialization.cpp

@@ -36,11 +36,7 @@ void AutomaticWarpSpecialization::runOnOperation() {
   OpPassManager pm;
   pm.addPass(createTritonGPUPartitionScheduling());
   pm.addPass(createNVWSInsertAref());
-#if 0
   pm.addPass(createNVWSInsertTmemAref());
-#else
-  pm.addPass(createTritonGPULoadMMASpecialization({numStages}));
-#endif
   pm.addPass(createTritonGPURewritePartitionDependencies());
   // `int-range-optimizations` and SCCP are good at cleaning up loop arithmetic.
   // FIXME: Re-enable integer range analysis once it is fixed.
@@ -50,19 +46,6 @@ void AutomaticWarpSpecialization::runOnOperation() {
   pm.addPass(createNVWSLowerAref({numStages}));
   pm.addPass(createTritonGPUPartitionLoops());
   pm.addPass(createNVWSLowerWarpGroup());
-  if (failed(runPipeline(pm, getOperation())))
-    return signalPassFailure();
-
-  // Cleanup code generated by warp specialization.
-  RewritePatternSet patterns(&getContext());
-  populateForOpDeadArgumentElimination(patterns);
-  scf::ForOp::getCanonicalizationPatterns(patterns, &getContext());
-  scf::IfOp::getCanonicalizationPatterns(patterns, &getContext());
-  WarpSpecializeOp::getCanonicalizationPatterns(patterns, &getContext());
-  if (failed(applyPatternsGreedily(getOperation(), std::move(patterns))))
-    return signalPassFailure();
-
-  pm.clear();
   pm.addPass(createTritonGPUOptimizePartitionWarps());
   pm.addPass(createTritonGPUScheduleLoops());
   if (failed(runPipeline(pm, getOperation())))