Merge OpenAI Triton commit e1162ee (#3913)

This PR change the Triton base from
a0cc214 to
e1162ee (Apr 10).
Pass rate: 88.42%->88.52%

Author: whitneywhtsang

bin/RegisterTritonDialects.h

@@ -101,6 +101,7 @@ inline void registerTritonDialects(mlir::DialectRegistry &registry) {
   mlir::registerTritonAMDGPUConvertToBufferOps();
   mlir::registerTritonAMDGPUInThreadTranspose();
   mlir::registerTritonAMDGPUCoalesceAsyncCopy();
+  mlir::registerTritonAMDGPUUpdateAsyncWaitCount();
   mlir::triton::registerTritonAMDGPUInsertInstructionSchedHints();
   mlir::triton::registerTritonAMDGPULowerInstructionSchedHints();
   mlir::registerTritonAMDFoldTrueCmpI();

include/triton/Analysis/Alias.h

@@ -89,6 +89,11 @@ class SharedMemoryAliasAnalysis
   visitOperation(Operation *op,
                  ArrayRef<const dataflow::Lattice<AliasInfo> *> operands,
                  ArrayRef<dataflow::Lattice<AliasInfo> *> results) override;
+
+  void visitNonControlFlowArguments(
+      Operation *op, const RegionSuccessor &successor,
+      ArrayRef<dataflow::Lattice<AliasInfo> *> argLattices,
+      unsigned firstIndex) override;
 };
 
 } // namespace mlir

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

@@ -218,6 +218,23 @@ def TritonGPUPrefetch : Pass<"tritongpu-prefetch", "mlir::ModuleOp"> {
                            "mlir::arith::ArithDialect"];
 }
 
+def TritonGPUWGMMAPrefetch : Pass<"tritongpu-wgmma-prefetch", "mlir::ModuleOp"> {
+   let summary = "prefetch for wgmma mixed precision";
+
+   let description = [{
+       This pass attempts to prefetch from shared memory for mixed-precision
+       wgmma when operand A is in the shared memory and needs to be loaded
+       to the local registers.
+   }];
+
+   let dependentDialects = [ "mlir::triton::gpu::TritonGPUDialect",
+                             "mlir::triton::nvidia_gpu::TritonNvidiaGPUDialect",
+                             "mlir::scf::SCFDialect",
+                             "mlir::arith::ArithDialect"];
+}
+
+
+
 def TritonGPUAccelerateMatmul : Pass<"tritongpu-accelerate-matmul", "mlir::ModuleOp"> {
   let summary = "accelerate matmul";
 

lib/Analysis/Alias.cpp

@@ -58,6 +58,30 @@ LogicalResult SharedMemoryAliasAnalysis::visitOperation(
   return success();
 }
 
+void SharedMemoryAliasAnalysis::visitNonControlFlowArguments(
+    Operation *op, const RegionSuccessor &successor,
+    ArrayRef<dataflow::Lattice<AliasInfo> *> argLattices, unsigned firstIndex) {
+  auto wsOp = dyn_cast<triton::gpu::WarpSpecializePartitionsOp>(op);
+  if (!wsOp) {
+    setAllToEntryStates(argLattices.take_front(firstIndex));
+    setAllToEntryStates(argLattices.drop_front(
+        firstIndex + successor.getSuccessorInputs().size()));
+    return;
+  }
+
+  // Propagate aliases from the parent operation's operands to the block
+  // arguments.
+  assert(!successor.isParent());
+  ProgramPoint *point = getProgramPointAfter(wsOp);
+
+  for (auto [capture, argLattice] :
+       llvm::zip(wsOp.getParentOp().getExplicitCaptures(), argLattices)) {
+    propagateIfChanged(
+        argLattice,
+        argLattice->join(getLatticeElementFor(point, capture)->getValue()));
+  }
+}
+
 AliasResult SharedMemoryAliasAnalysis::alias(Value lhs, Value rhs) {
   // TODO: implement
   return AliasResult::MayAlias;

lib/Analysis/Allocation.cpp

@@ -332,7 +332,7 @@ class AllocationAnalysis {
         solver->load<SharedMemoryAliasAnalysis>();
     // Run the analysis rooted at every isolated from above operation, including
     // the top-level function but also any nested regions.
-    operation->walk([&](Operation *op) {
+    operation->walk<mlir::WalkOrder::PreOrder>([&](Operation *op) {
       if (op->hasTrait<OpTrait::IsIsolatedFromAbove>() &&
           failed(solver->initializeAndRun(op))) {
         // TODO: return error instead of bailing out..

lib/Dialect/TritonGPU/Transforms/CMakeLists.txt

@@ -26,6 +26,7 @@ add_triton_library(TritonGPUTransforms
   Pipeliner/PipeliningUtility.cpp
   Pipeliner/Schedule.cpp
   Prefetch.cpp
+  WGMMAPrefetch.cpp
   RemoveLayoutConversions.cpp
   ReorderInstructions.cpp
   CoalesceAsyncCopy.cpp