Merge commit 'ade4d3ac30dbe4a8de9d3da1441160544beb6d79'

Author: whitneywhtsang

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

@@ -165,6 +165,16 @@ def TritonGPUOptimizePartitionWarps : Pass<"tritongpu-optimize-partition-warps",
   }];
 }
 
+def TritonGPUPartitionScheduling : Pass<"tritongpu-partition-scheduling", "mlir::ModuleOp"> {
+  let summary = "warp specialization partitioning pass";
+
+  let description = [{
+    The `tritongpu-partition-scheduling` analyzes the loads, MMAs, and other
+    operations in a loop that is meant to be warp specialized and determines
+    which partitions to assign to each operation.
+  }];
+}
+
 def TritonGPULoadMMASpecialization : Pass<"tritongpu-load-mma-specialization", "mlir::ModuleOp"> {
   let summary = "load MMA specialization";
 
@@ -219,23 +229,6 @@ 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";
 

include/triton/Dialect/TritonGPU/Transforms/Utility.h

@@ -54,6 +54,10 @@ getNumElementsPerThread(Operation *op, SmallVector<unsigned> order,
 // Returns whether the op is a "view op", i.e. doesn't move any data
 bool isView(Operation *op);
 
+// Returns whether the op is a "noop op", i.e. has one input and one output
+// and lowers to llvm as the identity function (returns the input)
+bool isNoop(Operation *op);
+
 /* Dump Triton IR in graphviz dot format.
  *
  * You can override `onValue` and `onOperation` in a subclass to mark

lib/Conversion/TritonGPUToLLVM/ViewOpToLLVM.cpp

@@ -441,45 +441,7 @@ struct MemDescSubviewOpConversion
     for (int i = rankReduced; i < opOffsetVals.size(); i++) {
       offsetVals.push_back(b.add(opOffsetVals[i], smemObj.getOffsets()[i]));
     }
-    Value offset;
-    auto allocShape = srcTy.getAllocShape();
-    auto nvmmaEnc = dyn_cast<NVMMASharedEncodingAttr>(enc);
-    bool isSimpleSubview =
-        (!nvmmaEnc || allocShape.take_back(destRank) == destTy.getShape() ||
-         nvmmaEnc.getSwizzlingByteWidth() == 0);
-    if (!isSimpleSubview) {
-      assert(destRank >= 2 &&
-             "Shape size should be >= 2 when using NVMMAShared encoding");
-      auto swizzleStride = b.i32_val((nvmmaEnc.getSwizzlingByteWidth() * 8) /
-                                     llvmElemTy.getIntOrFloatBitWidth());
-      offset = b.i32_val(0);
-      for (auto i = 0; i < opOffsetVals.size() - 2; ++i) {
-        offset = b.add(offset, b.mul(opOffsetVals[i], opSmemStrides[i]));
-      }
-      // newOffset = offset - (stridedOff * swizzledStride + contigOff /
-      // swizzledStride * tileSize + contigOff % swizzledStride)
-      // + stridedInc * swizzledStride + contigInc / swizzledStride *
-      // tileSize + contigInc % swizzledStride
-      auto stridedDim = destRank - 1 - layoutOrder[0];
-      auto contigDim = destRank - 1 - layoutOrder[1];
-      auto stridedOff = smemObj.getOffsets()[stridedDim];
-      auto contigOff = smemObj.getOffsets()[contigDim];
-      auto stridedInc = offsetVals[stridedDim];
-      auto contigInc = offsetVals[contigDim];
-      int allocStridedDim = allocShape.size() - 1 - layoutOrder[0];
-      auto tileSize =
-          b.mul(b.i32_val(allocShape[allocStridedDim]), swizzleStride);
-      offset = b.sub(offset, b.mul(stridedOff, swizzleStride));
-      offset = b.sub(offset, b.mul(b.udiv(contigOff, swizzleStride), tileSize));
-      offset = b.sub(offset, b.urem(contigOff, swizzleStride));
-      offset = b.add(offset, b.mul(stridedInc, swizzleStride));
-      offset = b.add(offset, b.mul(b.udiv(contigInc, swizzleStride), tileSize));
-      offset = b.add(offset, b.urem(contigInc, swizzleStride));
-    } else {
-      // Compute the offset based on the original strides of the shared memory
-      // object
-      offset = dot(rewriter, loc, opOffsetVals, opSmemStrides);
-    }
+    Value offset = dot(rewriter, loc, opOffsetVals, opSmemStrides);
     auto base = smemObj.getBase();
     auto elemPtrTy = base.getType();
     smemObj = SharedMemoryObject(b.gep(elemPtrTy, llvmElemTy, base, offset),

lib/Dialect/TritonGPU/Transforms/CMakeLists.txt

@@ -25,7 +25,6 @@ add_triton_library(TritonGPUTransforms
   Pipeliner/PipeliningUtility.cpp
   Pipeliner/Schedule.cpp
   Prefetch.cpp
-  WGMMAPrefetch.cpp
   RemoveLayoutConversions.cpp
   ReorderInstructions.cpp
   CoalesceAsyncCopy.cpp
@@ -35,6 +34,7 @@ add_triton_library(TritonGPUTransforms
   WarpSpecialization/Partition.cpp
   WarpSpecialization/OptimizePartitionWarps.cpp
   WarpSpecialization/PartitionLoops.cpp
+  WarpSpecialization/PartitionScheduling.cpp
   WarpSpecialization/RewritePartitionDependencies.cpp
 
   DEPENDS