[mlir][vector-to-gpu] Fix the Transpose Check in `mma.sync` VectorToGPU Lowering Path

[manishucsd]

Summary

For mma.sync Tensor Core lowering path with split-k reduction and multi-buffering the mainloop's vector.transfer_read op reads from 4d memref instances.

See example:

%27 = vector.transfer_read %alloc_1[%20, %26, %c0, %c0], %cst_0 {in_bounds = [true, true], permutation_map = affine_map<(d0, d1, d2, d3) -> (d1, d3)>} : memref<4x32x1x32xf16, #gpu.address_space<workgroup>>, vector<16x16xf16>
%28 = vector.transfer_read %alloc_1[%20, %26, %c0, %c16], %cst_0 {in_bounds = [true, true], permutation_map = affine_map<(d0, d1, d2, d3) -> (d1, d3)>} : memref<4x32x1x32xf16, #gpu.address_space<workgroup>>, vector<16x16xf16>

%30 = vector.transfer_read %alloc_2[%20, %c0, %c0, %29], %cst_0 {in_bounds = [true, true], permutation_map = affine_map<(d0, d1, d2, d3) -> (d3, d2)>} : memref<4x1x32x32xf16, #gpu.address_space<workgroup>>, vector<16x16xf16>
%31 = vector.transfer_read %alloc_2[%20, %c0, %c16, %29], %cst_0 {in_bounds = [true, true], permutation_map = affine_map<(d0, d1, d2, d3) -> (d3, d2)>} : memref<4x1x32x32xf16, #gpu.address_space<workgroup>>, vector<16x16xf16>

Without this change all four vector.transfered_read are marked to use transpose resulting in ldmatrix.trans for both matrixA and matrixB in row-row matmul, while only matrixB needs a transpose. This patch fixes the issue.

[ThomasRaoux]

Make sure to add a test.

mlir/lib/Conversion/VectorToGPU/VectorToGPU.cpp

661

nit: remove empty line

664–666

either make it an assert or return false. We should not be emitting error and just continuing.

669

the rule is that you should use auto only if this is a cast where the type is explicit set or this is a complex type (iterator kind of stuff) so here you need to spell out the type.

676

Same as above emitting an error here doesn't sound right, either we should have code preventing such case before and then assert or return false or propagate an error.

704

nit /*transpose=*/ doesn't add much anymore I would remove it.

[manishucsd]

Added F16 test

[ThomasRaoux]

mlir/test/Conversion/VectorToGPU/vector-to-mma-ops-mma-sync.mlir

257 ↗

(On Diff #511714)

This comment doesn't really have a meaning in the context of upstream flow. It would be better to comment on what problem it is trying to catch.

271–309 ↗

(On Diff #511714)

Can you minimize the test? It should be as simple as possible while still catching the bug you are fixing.

[manishucsd]

Applied comments.

[manishucsd]

mlir/test/Conversion/VectorToGPU/vector-to-mma-ops-mma-sync.mlir

271–309 ↗

(On Diff #511714)

Updated the comment for the test.

Also, pruned the test to only check the following:
vector.transfer_read for A (transpose=false), vector.transfer_read for B (transpose=true), vector.transfer_read for C (transpose=false) on multi-dimensional vector.transfer_read ops.
mma.sync is coming as we don't want DCE to trigger.

[ThomasRaoux]

Looks good, one last nit

mlir/test/Conversion/VectorToGPU/vector-to-mma-ops-mma-sync.mlir

278 ↗

(On Diff #511743)

replace %c0 to not assume special naming here and below.

mlir/lib/Conversion/VectorToGPU/VectorToGPU.cpp

convertConstantOpMmaSync(RewriterBase &rewriter, arith::ConstantOp op,
 
   Value result = rewriter.create<arith::ConstantOp>(
       op.getLoc(), vectorType,
       DenseElementsAttr::get(vectorType, dense.getSplatValue<Attribute>()));
   valueMapping[op.getResult()] = result;
   return success();
 }
 
+/// Check if the loaded matrix operand requires transposed.
+/// Transposed Map Example:
+/// Example 1   : (..., d0, d1) -> (d1 * 1, d0 * 2)
+/// Example 2   : (d0, d1, d2, d3) -> (d3, d2)
+///
+/// The code below checks if the output 2D is transposed using a generalized
+/// version     : (d0, d1, dn, ..., dm, ...) -> (dm, dn)
+/// Returns     : true; if m > n, false o.w.
+
+static bool isTransposed(vector::TransferReadOp op) {
+  mlir::AffineMap map = op.getPermutationMap();
+  if (map.getNumResults() != 2) {
+    op->emitError("Expected 2D transfer read");
+  }
+
+  // Output 2D matrix dimensions in the order of d0, d1.
+  auto dM = map.getResult(0);
+  auto dN = map.getResult(1);
+
+  //  Find the position of these expressions in the input.
+  auto exprM = dM.dyn_cast<AffineDimExpr>();
+  auto exprN = dN.dyn_cast<AffineDimExpr>();
+  if (!exprM || !exprN) {
+    op->emitError("Expected to find AffineDimExpr in vector::TransferReadOp");
+  }
+  return exprM.getPosition() > exprN.getPosition();
+}
+
 static LogicalResult
 creatLdMatrixCompatibleLoads(RewriterBase &rewriter, vector::TransferReadOp op,
                              llvm::DenseMap<Value, Value> &valueMapping) {
   OpBuilder::InsertionGuard g(rewriter);
   rewriter.setInsertionPoint(op);
   Location loc = op->getLoc();
 
   FailureOr<nvgpu::WarpMatrixInfo> warpMatrixInfo =
       nvgpu::getWarpMatrixInfo(op);
   if (failed(warpMatrixInfo)) {
     LLVM_DEBUG(DBGS() << "no warpMatrixInfo\n");
     return rewriter.notifyMatchFailure(op, "no warpMatrixInfo");
   }
 
   FailureOr<nvgpu::FragmentElementInfo> regInfo =
       nvgpu::getMmaSyncRegisterType(*warpMatrixInfo);
   if (failed(regInfo)) {
     LLVM_DEBUG(DBGS() << "not mma sync reg info\n");
     return rewriter.notifyMatchFailure(op, "not mma sync reg info");
   }
 
-  FailureOr<nvgpu::LdMatrixParams> params = nvgpu::getLdMatrixParams(
-      *warpMatrixInfo,
-      /*transpose=*/!op.getPermutationMap().isMinorIdentity());
+  FailureOr<nvgpu::LdMatrixParams> params =
+      nvgpu::getLdMatrixParams(*warpMatrixInfo,
+                               /*transpose=*/isTransposed(op));
+
   if (failed(params)) {
     LLVM_DEBUG(
         DBGS()
         << "failed to convert vector.transfer_read to ldmatrix. "
         << "Op should likely not be converted to a nvgpu.ldmatrix call.\n");
     return rewriter.notifyMatchFailure(
         op, "failed to convert vector.transfer_read to ldmatrix; this op "
             "likely should not be converted to a nvgpu.ldmatrix call.");
creatLdMatrixCompatibleLoads(RewriterBase &rewriter, vector::TransferReadOp op,
 
   VectorType vectorType = getMmaSyncVectorOperandType(*regInfo);
 
   SmallVector<Value, 4> indices;
   getXferIndices<vector::TransferReadOp>(rewriter, op, *offsets, {laneId},
                                          indices);
   nvgpu::LdMatrixOp newOp = rewriter.create<nvgpu::LdMatrixOp>(
       loc, vectorType, op.getSource(), indices,
-      !op.getPermutationMap().isMinorIdentity(), params->numTiles);
+      /*transpose=*/isTransposed(op), params->numTiles);
   valueMapping[op] = newOp->getResult(0);
   return success();
 }
 
 static LogicalResult
 createNonLdMatrixLoads(RewriterBase &rewriter, vector::TransferReadOp op,
                        llvm::DenseMap<Value, Value> &valueMapping) {
   OpBuilder::InsertionGuard g(rewriter);