[PIPELINER] Pipeline RS WGMMA (#6804) (#6812)

This PR allows to pipeline WGMMAs that take the lhs on registers. The
strategy is to wait on the WGMMA from the previous loop to have finished
before executing the next one to avoid overwritting the registers too
early.

Note that this does depend on ptxas handling the register allocation
correctly.

This PR also includes:
- A fix for the WGMMAPrefetch with `swizzlingByteWidth = 128`, which
produced wrong results
- A fix in the way we lower `memdesc_subview` (now it's simpler again)
- A fix in the way we lower mmav3 and mmav5 (which renders the complex
path in `memdesc_subview` unnecessary)
All these are tested end-to-end via the improved `test_cast_matmul.py`

In an 8k x 8k x 8k dense bf16 x mxfp4 matmul we get a speed up of: 2.441
-> 2.039

We might need to split the pointwise computations and interleave them
with the wgmmas similar to how CUTLASS does it, but we don't do that
in this PR.

This PR supersedes WGMMAPrefetch as it drops most of the preconditions
of that pass.

Author: lezcano

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

@@ -219,23 +219,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