[mlir][Vector] Significantly improve VectorToGPU.cpp

[nicolasvasilache]

Summary

This revision performs a bunch of cleanups and tracks free-flowing IR mutations.
APIs are systematized around RewriterBase and relevant debug messages are added.
Deliberate use of OpBuilder::InsertionGuard is added where needed.

[ThomasRaoux]

Thanks!

[nicolasvasilache]

Rebase

[stella.stamenova]

Looks like this broke the windows mlir buildbot: https://lab.llvm.org/buildbot/#/builders/13/builds/32020

mlir/include/mlir/Conversion/VectorToGPU/VectorToGPU.h

 //===- VectorToGPU.h - Convert vector to GPU dialect ------------*- C++ -*-===//
 //
 // Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
 // See https://llvm.org/LICENSE.txt for license information.
 // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
 //
 //===----------------------------------------------------------------------===//
 
 #ifndef MLIR_CONVERSION_VECTORTOGPU_VECTORTOGPU_H
 #define MLIR_CONVERSION_VECTORTOGPU_VECTORTOGPU_H
 
 #include "mlir/IR/PatternMatch.h"
 
 namespace mlir {
+class LogicalResult;
 class MLIRContext;
 class Pass;
 class RewritePatternSet;
 
 #define GEN_PASS_DECL_CONVERTVECTORTOGPU
 #include "mlir/Conversion/Passes.h.inc"
 
 /// Patterns to transform vector ops into a canonical form to convert to MMA
 /// matrix operations. If `useNvGpu` is true, then the patterns will populated
 /// will prepare for conversion to `nvgpu` mma operations rather than the `gpu`
 /// dialect WMMA operations.
 void populatePrepareVectorToMMAPatterns(RewritePatternSet &patterns,
                                         bool useNvGpu = false);
 
 /// Convert vector ops to MMA matrix operations nested under `rootOp`. This will
 /// convert slice of operations that can be legally converted to MMA operations.
 /// The rest of the vector operations are left untouched.
-void convertVectorToMMAOps(Operation *rootOp);
+LogicalResult convertVectorToMMAOps(RewriterBase &rewriter, Operation *rootOp);
 
 /// Convert vector ops ops nested under `rootOp` to vector and GPU operaitons
 /// compatible with the `nvvm.mma.sync` lowering path. This will convert a slice
 /// of operations that can be legally lowered on this path while the rest of
 /// the vector operations are left untouched.
-LogicalResult convertVectorToNVVMCompatibleMMASync(Operation *rootOp);
+LogicalResult convertVectorToNVVMCompatibleMMASync(RewriterBase &rewriter,
+                                                   Operation *rootOp);
 
 /// Convert from vector to GPU ops.
 std::unique_ptr<Pass> createConvertVectorToGPUPass(bool useNvGpu = false);
 
 } // namespace mlir
 
 #endif // MLIR_CONVERSION_VECTORTOGPU_VECTORTOGPU_H

mlir/include/mlir/Dialect/NVGPU/Utils/MMAUtils.h

 /// Returns an AffineMap which maps a two dimensions representing (laneId,
 /// logicalValueId) and returns two results representing offsets within a
 /// matrix operand. The offsets point to the values the thread is responsible
 /// for (AKA the matrix fragment values) during a warp-collective matrix
 /// operation. For a visual reference of this LaneId -> (row, col) mapping,
 /// please see NVIDIA's PTX documentation:
 /// https://docs.nvidia.com/cuda/parallel-thread-execution/index.html#warp-level-matrix-instructions-for-mma
 FailureOr<AffineMap>
-getLaneIdAndValueIdToOperandCoord(Location loc, OpBuilder &builder,
+getLaneIdAndValueIdToOperandCoord(OpBuilder &builder, Location loc,
                                   const WarpMatrixInfo &fragmentType);
 
 /// Encapsulates the parameters needed to lower a `nvgpu.ldmatrix` operation to
 /// `nvvm.ldmatrix`.
 struct LdMatrixParams {
   VectorType fragmentType;
   bool isAccum;
   int64_t numTiles;
   vector::IteratorType contiguousDimType;
   NVVM::MMALayout targetLayout;
 };
 
 /// Given `type` that contains info for a warp-matrix operand and whether or not
 /// the load is a transposed load, return the LdMatrixParams.
 FailureOr<LdMatrixParams> getLdMatrixParams(const WarpMatrixInfo &type,
                                             bool transpose);
 /// Returns an AffineMap which maps a single dimension representing the laneId
 /// to two results representing offsets within the matrix operand that should
 /// be the pointer locations a thread should pass to the ldmatrix instruction.
 FailureOr<AffineMap>
-getLaneIdToLdMatrixMatrixCoord(Location loc, OpBuilder &builder,
+getLaneIdToLdMatrixMatrixCoord(OpBuilder &builder, Location loc,
                                const LdMatrixParams &params);
 
 /// Transform `vector.contract` into (m,k)x(n,k)x(m,n) form so that it can be
 /// converted to `nvgpu.mma.sync`. This specific form is meant to indicate that
 /// the vector operands are organized such that the reduction dimension is
 /// contiguous.
 struct PrepareContractToGPUMMASync
     : public OpRewritePattern<vector::ContractionOp> {

mlir/lib/Conversion/VectorToGPU/VectorToGPU.cpp

 #include "mlir/Dialect/MemRef/IR/MemRef.h"
 #include "mlir/Dialect/NVGPU/IR/NVGPUDialect.h"
 #include "mlir/Dialect/NVGPU/Utils/MMAUtils.h"
 #include "mlir/Dialect/SCF/IR/SCF.h"
 #include "mlir/Dialect/Utils/StructuredOpsUtils.h"
 #include "mlir/Dialect/Vector/IR/VectorOps.h"
 #include "mlir/Dialect/Vector/Utils/VectorUtils.h"
 #include "mlir/IR/Builders.h"
+#include "mlir/IR/BuiltinOps.h"
+#include "mlir/IR/Region.h"
 #include "mlir/Pass/Pass.h"
+#include "mlir/Support/LogicalResult.h"
 #include "mlir/Transforms/GreedyPatternRewriteDriver.h"
 #include "mlir/Transforms/Passes.h"
+#include "llvm/ADT/STLExtras.h"
 #include "llvm/ADT/TypeSwitch.h"
 
+#define DEBUG_TYPE "vector-to-gpu"
+#define DBGS() (llvm::dbgs() << "[" DEBUG_TYPE "]: ")
+#define DBGSNL() (llvm::dbgs() << "\n")
+
 namespace mlir {
 #define GEN_PASS_DEF_CONVERTVECTORTOGPU
 #include "mlir/Conversion/Passes.h.inc"
 } // namespace mlir
 
 using namespace mlir;
 
 /// For a vector TransferOpType `xferOp`, an empty `indices` vector, and an
 /// AffineMap representing offsets to apply to indices, the function fills
 /// `indices` with the original indices plus the offsets. The offsets are
 /// applied by taking into account the permutation map of the transfer op. If
 /// the `offsetMap` has dimension placeholders, those should be provided in
 /// `dimValues`.
 template <typename TransferOpType>
-static void getXferIndices(OpBuilder &b, TransferOpType xferOp,
+static void getXferIndices(RewriterBase &rewriter, TransferOpType xferOp,
                            AffineMap offsetMap, ArrayRef<Value> dimValues,
                            SmallVector<Value, 4> &indices) {
   indices.append(xferOp.getIndices().begin(), xferOp.getIndices().end());
   Location loc = xferOp.getLoc();
   unsigned offsetsIdx = 0;
   for (auto expr : xferOp.getPermutationMap().getResults()) {
     if (auto dim = expr.template dyn_cast<AffineDimExpr>()) {
       Value prevIdx = indices[dim.getPosition()];
       SmallVector<Value, 3> dims(dimValues.begin(), dimValues.end());
       dims.push_back(prevIdx);
-      AffineExpr d0 = b.getAffineDimExpr(offsetMap.getNumDims());
+      AffineExpr d0 = rewriter.getAffineDimExpr(offsetMap.getNumDims());
       indices[dim.getPosition()] = makeComposedAffineApply(
-          b, loc, d0 + offsetMap.getResult(offsetsIdx++), dims);
+          rewriter, loc, d0 + offsetMap.getResult(offsetsIdx++), dims);
       continue;
     }
   }
 }
 
 // Return true if the contract op can be convert to MMA matmul.
 static bool contractSupportsMMAMatrixType(vector::ContractionOp contract,
                                           bool useNvGpu) {
if (useNvGpu &&
       contract.getIndexingMapsArray() != infer({{m, k}, {n, k}, {m, n}}))
     return false;
 
   return true;
 }
 
 // Return true if the given map represents a transposed matrix load,
 // i.e. (d0, d1, ...) -> (dn-1, dn-2).
-static bool isTransposeMatrixLoadMap(OpBuilder &b, AffineMap permutationMap) {
-  MLIRContext *ctx = b.getContext();
+static bool isTransposeMatrixLoadMap(AffineMap permutationMap) {
+  MLIRContext *ctx = permutationMap.getContext();
+  // Local OpBuilder is fine here, we just build attributes.
+  OpBuilder b(ctx);
   auto nDim = permutationMap.getNumDims();
   AffineExpr zero = b.getAffineConstantExpr(0);
   if (nDim < 2) {
     // Support transposed+broadcasted cases: affine_map<(d0) -> (d0, 0)>.
     AffineExpr dim0 = b.getAffineDimExpr(0);
     return permutationMap == AffineMap::get(1, 0, {dim0, zero}, ctx);
   }
 
static bool transferReadSupportsMMAMatrixType(vector::TransferReadOp readOp,
 
   // Only allow integer types if the signedness can be inferred.
   if (!useNvGpu && readOp.getVectorType().getElementType().isInteger(8))
     if (!readOp->hasOneUse() || (!isa<arith::ExtSIOp>(*readOp->user_begin()) &&
                                  !isa<arith::ExtUIOp>(*readOp->user_begin())))
       return false;
 
   AffineMap map = readOp.getPermutationMap();
-  OpBuilder b(readOp.getContext());
-  AffineExpr innerDim = b.getAffineDimExpr(map.getNumDims() - 1);
-  AffineExpr zero = b.getAffineConstantExpr(0);
-  auto broadcastInnerDim = AffineMap::get(map.getNumDims(), 0, {zero, innerDim},
-                                          readOp.getContext());
+
+  MLIRContext *ctx = readOp.getContext();
+  AffineExpr innerDim = getAffineDimExpr(map.getNumDims() - 1, ctx);
+  AffineExpr zero = getAffineConstantExpr(0, ctx);
+  auto broadcastInnerDim =
+      AffineMap::get(map.getNumDims(), 0, {zero, innerDim}, ctx);
 
   if (!useNvGpu) {
     bool result = map.isMinorIdentity() || map == broadcastInnerDim ||
-                  isTransposeMatrixLoadMap(b, map);
+                  isTransposeMatrixLoadMap(map);
     return result;
   }
 
   return true;
 }
 
 // Return true if the transfer op can be converted to a MMA matrix store.
 static bool
LogicalResult matchAndRewrite(vector::ContractionOp op,
     AffineExpr m, n, k;
     bindDims(rewriter.getContext(), m, n, k);
     static constexpr std::array<int64_t, 2> perm = {1, 0};
     auto iteratorTypes = op.getIteratorTypes().getValue();
     SmallVector<AffineMap, 4> maps = op.getIndexingMapsArray();
     if (!(vector::isParallelIterator(iteratorTypes[0]) &&
           vector::isParallelIterator(iteratorTypes[1]) &&
           vector::isReductionIterator(iteratorTypes[2])))
-      return failure();
+      return rewriter.notifyMatchFailure(op, "not a gemm contraction");
     //
     // Two outer parallel, one inner reduction (matmat flavor).
     //
-    if (maps == infer({{m, k}, {k, n}, {m, n}})) {
-      // This is the classical row-major matmul, nothing to do.
-      return failure();
-    }
+    if (maps == infer({{m, k}, {k, n}, {m, n}}))
+      return rewriter.notifyMatchFailure(op, "contraction already prepared");
     if (maps == infer({{m, k}, {n, k}, {m, n}})) {
       rhs = rewriter.create<vector::TransposeOp>(loc, rhs, perm);
     } else if (maps == infer({{k, m}, {k, n}, {m, n}})) {
       lhs = rewriter.create<vector::TransposeOp>(loc, lhs, perm);
     } else if (maps == infer({{k, m}, {n, k}, {m, n}})) {
       rhs = rewriter.create<vector::TransposeOp>(loc, rhs, perm);
       lhs = rewriter.create<vector::TransposeOp>(loc, lhs, perm);
     } else if (maps == infer({{m, k}, {k, n}, {n, m}})) {
       std::swap(rhs, lhs);
       rhs = rewriter.create<vector::TransposeOp>(loc, rhs, perm);
       lhs = rewriter.create<vector::TransposeOp>(loc, lhs, perm);
     } else if (maps == infer({{m, k}, {n, k}, {n, m}})) {
       std::swap(rhs, lhs);
       rhs = rewriter.create<vector::TransposeOp>(loc, rhs, perm);
     } else if (maps == infer({{k, m}, {k, n}, {n, m}})) {
       std::swap(lhs, rhs);
       lhs = rewriter.create<vector::TransposeOp>(loc, lhs, perm);
     } else if (maps == infer({{k, m}, {n, k}, {n, m}})) {
       std::swap(lhs, rhs);
     } else {
-      return failure();
+      // TODO: llvm_unreachable ?
+      return rewriter.notifyMatchFailure(op, "unexpected contraction case");
     }
     rewriter.replaceOpWithNewOp<vector::ContractionOp>(
         op, lhs, rhs, res,
         rewriter.getAffineMapArrayAttr(infer({{m, k}, {k, n}, {m, n}})),
         op.getIteratorTypes());
     return success();
   }
 };
if ((extOp = source.getDefiningOp<arith::ExtSIOp>()) ||
       source = extOp->getOperand(0);
       resultType =
           VectorType::get(resultType.getShape(),
                           source.getType().cast<VectorType>().getElementType());
     }
 
     auto transferReadOp = source.getDefiningOp<vector::TransferReadOp>();
     if (!transferReadOp)
-      return failure();
+      return rewriter.notifyMatchFailure(op, "no transfer read");
 
     // TODO: support 0-d corner case.
     if (transferReadOp.getTransferRank() == 0)
-      return failure();
+      return rewriter.notifyMatchFailure(op, "0-D transfer read");
 
     if (transferReadOp.getMask() || transferReadOp.hasOutOfBoundsDim())
-      return failure();
+      return rewriter.notifyMatchFailure(op, "not inbounds transfer read");
+
     SmallVector<int64_t, 2> perm;
     op.getTransp(perm);
     SmallVector<unsigned, 2> permU;
     for (int64_t o : perm)
       permU.push_back(unsigned(o));
     AffineMap permutationMap =
         AffineMap::getPermutationMap(permU, op.getContext());
     AffineMap newMap =
for (Operation *users : op->getUsers()) {
     if (contract.getLhs() == op->getResult(0))
       return "AOp";
     if (contract.getRhs() == op->getResult(0))
       return "BOp";
   }
   return "COp";
 }
 
-static void convertTransferReadOp(vector::TransferReadOp op,
-                                  llvm::DenseMap<Value, Value> &valueMapping) {
+static LogicalResult
+convertTransferReadOp(RewriterBase &rewriter, vector::TransferReadOp op,
+  OpBuilder::InsertionGuard g(rewriter);
+  rewriter.setInsertionPoint(op);
+
   assert(op.getTransferRank() > 0 && "unexpected 0-d transfer");
   assert(transferReadSupportsMMAMatrixType(op, /*useNvGpu=*/false));
 
   std::optional<int64_t> stride =
       getMemrefConstantHorizontalStride(op.getShapedType());
+  if (!stride.has_value()) {
+    LLVM_DEBUG(DBGS() << "no stride\n");
+    return rewriter.notifyMatchFailure(op, "no stride");
+  }
 
   AffineMap map = op.getPermutationMap();
-  OpBuilder b(op);
-  bool isTranspose = isTransposeMatrixLoadMap(b, map);
+  bool isTranspose = isTransposeMatrixLoadMap(map);
 
   // Handle broadcast by setting the stride to 0.
   if (auto cstExpr =
           map.getResult(isTranspose).dyn_cast<AffineConstantExpr>()) {
     assert(cstExpr.getValue() == 0);
     stride = 0;
   }
-  assert(stride);
+
   Value mappingResult = op.getResult();
   auto elType = op.getVectorType().getElementType();
   const char *fragType = inferFragType(op);
   if (op->hasOneUse()) {
     auto user = *op->user_begin();
     // Infer the signedness of the mma type from the integer extend.
     bool isSignedExtend = isa<arith::ExtSIOp>(user);
     if (isSignedExtend || isa<arith::ExtUIOp>(user)) {
       elType = IntegerType::get(
           op.getContext(), elType.cast<IntegerType>().getWidth(),
           isSignedExtend ? IntegerType::Signed : IntegerType::Unsigned);
       mappingResult = user->getResult(0);
       fragType = inferFragType(user);
     }
   }
   gpu::MMAMatrixType type =
       gpu::MMAMatrixType::get(op.getVectorType().getShape(), elType, fragType);
-  Value load = b.create<gpu::SubgroupMmaLoadMatrixOp>(
+  Value load = rewriter.create<gpu::SubgroupMmaLoadMatrixOp>(
       op.getLoc(), type, op.getSource(), op.getIndices(),
-      b.getIndexAttr(*stride), isTranspose ? b.getUnitAttr() : UnitAttr());
+      rewriter.getIndexAttr(*stride),
+      isTranspose ? rewriter.getUnitAttr() : UnitAttr());
   valueMapping[mappingResult] = load;
+
+  LLVM_DEBUG(DBGS() << "transfer read to: " << load << "\n");
+  return success();
 }
 
-static void convertTransferWriteOp(vector::TransferWriteOp op,
-                                   llvm::DenseMap<Value, Value> &valueMapping) {
+static LogicalResult
+convertTransferWriteOp(RewriterBase &rewriter, vector::TransferWriteOp op,
+  OpBuilder::InsertionGuard g(rewriter);
+  rewriter.setInsertionPoint(op);
+
   assert(transferWriteSupportsMMAMatrixType(op));
   std::optional<int64_t> stride =
       getMemrefConstantHorizontalStride(op.getShapedType());
-  assert(stride);
-  OpBuilder b(op);
-  Value matrix = valueMapping.find(op.getVector())->second;
-  b.create<gpu::SubgroupMmaStoreMatrixOp>(
+  if (!stride.has_value()) {
+    LLVM_DEBUG(DBGS() << "no stride\n");
+    return rewriter.notifyMatchFailure(op, "no stride");
+  }
+
+  auto it = valueMapping.find(op.getVector());
+  if (it == valueMapping.end()) {
+    LLVM_DEBUG(DBGS() << "no mapping\n");
+    return rewriter.notifyMatchFailure(op, "no mapping");
+  }
+
+  Value matrix = it->second;
+  auto store = rewriter.create<gpu::SubgroupMmaStoreMatrixOp>(
       op.getLoc(), matrix, op.getSource(), op.getIndices(),
-      b.getIndexAttr(*stride), /*transpose=*/UnitAttr());
-  op.erase();
+      rewriter.getIndexAttr(*stride), /*transpose=*/UnitAttr());
+  (void)store;
+
+  LLVM_DEBUG(DBGS() << "transfer write to: " << store << "\n");
+
+  LLVM_DEBUG(DBGS() << "erase: " << op << "\n");
+  rewriter.eraseOp(op);
+  return success();
 }
 
 /// Returns the vector type which represents a matrix fragment.
 static VectorType
 getMmaSyncVectorOperandType(const nvgpu::FragmentElementInfo &regInfo) {
   SmallVector<int64_t> shape{regInfo.numRegistersPerFragment,
                              regInfo.elementsPerRegister};
   Type elType = regInfo.registerLLVMType;
   if (auto vecType = elType.dyn_cast<VectorType>())
     elType = vecType.getElementType();
   return VectorType::get(shape, elType);
 }
 
 /// Convert a 2D splat ConstantOp to a SubgroupMmaConstantMatrix op.
 static LogicalResult
-convertConstantOpMmaSync(arith::ConstantOp op,
+convertConstantOpMmaSync(RewriterBase &rewriter, arith::ConstantOp op,
                          llvm::DenseMap<Value, Value> &valueMapping) {
-  OpBuilder b(op);
+  OpBuilder::InsertionGuard g(rewriter);
+  rewriter.setInsertionPoint(op);
+
   FailureOr<nvgpu::WarpMatrixInfo> warpMatrixInfo =
       nvgpu::getWarpMatrixInfo(op);
-  if (failed(warpMatrixInfo))
-    return failure();
+  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))
-    return failure();
+  if (failed(regInfo)) {
+    LLVM_DEBUG(DBGS() << "not mma sync reg info\n");
+    return rewriter.notifyMatchFailure(op, "not mma sync reg info");
+  }
 
   VectorType vectorType = getMmaSyncVectorOperandType(*regInfo);
   auto dense = op.getValue().dyn_cast<SplatElementsAttr>();
-  if (!dense)
-    return failure();
-  Value result = b.create<arith::ConstantOp>(
+  if (!dense) {
+    LLVM_DEBUG(DBGS() << "not a splat\n");
+    return rewriter.notifyMatchFailure(op, "not a splat");
+  }
+
+  Value result = rewriter.create<arith::ConstantOp>(
       op.getLoc(), vectorType,
       DenseElementsAttr::get(vectorType, dense.getSplatValue<Attribute>()));
   valueMapping[op.getResult()] = result;
   return success();
 }
 
 static LogicalResult
-creatLdMatrixCompatibleLoads(vector::TransferReadOp op, OpBuilder &builder,
+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))
-    return failure();
+  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))
-    return failure();
+  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());
   if (failed(params)) {
-    return op->emitError()
-           << "failed to convert vector.transfer_read to ldmatrix; this op "
-              "likely "
-              "should not be converted to a nvgpu.ldmatrix call.";
+    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.");
   }
 
   // Adjust the load offset.
-  auto laneId = builder.create<gpu::LaneIdOp>(loc);
+  auto laneId = rewriter.create<gpu::LaneIdOp>(loc);
   FailureOr<AffineMap> offsets =
-      nvgpu::getLaneIdToLdMatrixMatrixCoord(loc, builder, *params);
-  if (failed(offsets))
-    return failure();
+      nvgpu::getLaneIdToLdMatrixMatrixCoord(rewriter, loc, *params);
+  if (failed(offsets)) {
+    LLVM_DEBUG(DBGS() << "no offsets\n");
+    return rewriter.notifyMatchFailure(op, "no offsets");
+  }
 
   VectorType vectorType = getMmaSyncVectorOperandType(*regInfo);
 
   SmallVector<Value, 4> indices;
-  getXferIndices<vector::TransferReadOp>(builder, op, *offsets, {laneId},
+  getXferIndices<vector::TransferReadOp>(rewriter, op, *offsets, {laneId},
                                          indices);
-  nvgpu::LdMatrixOp newOp = builder.create<nvgpu::LdMatrixOp>(
+  nvgpu::LdMatrixOp newOp = rewriter.create<nvgpu::LdMatrixOp>(
       loc, vectorType, op.getSource(), indices,
       !op.getPermutationMap().isMinorIdentity(), params->numTiles);
   valueMapping[op] = newOp->getResult(0);
   return success();
 }
 
 static LogicalResult
-createNonLdMatrixLoads(vector::TransferReadOp op, OpBuilder &builder,
+createNonLdMatrixLoads(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))
-    return failure();
+    return rewriter.notifyMatchFailure(op, "no warpMatrixInfo");
   FailureOr<nvgpu::FragmentElementInfo> regInfo =
       nvgpu::getMmaSyncRegisterType(*warpMatrixInfo);
   if (failed(regInfo)) {
-    op->emitError() << "Failed to deduce register fragment type during "
-                       "conversion to distributed non-ldmatrix compatible load";
-    return failure();
+    rewriter.notifyMatchFailure(
+        op, "Failed to deduce register fragment type during "
+            "conversion to distributed non-ldmatrix compatible load");
   }
 
-  Value laneId = builder.create<gpu::LaneIdOp>(loc);
+  Value laneId = rewriter.create<gpu::LaneIdOp>(loc);
   SmallVector<Value, 4> elements;
 
   // This is the individual element type.
   Type loadedElType = regInfo->registerLLVMType;
   VectorType vectorType = getMmaSyncVectorOperandType(*regInfo);
 
-  Value fill = builder.create<arith::ConstantOp>(
+  Value fill = rewriter.create<arith::ConstantOp>(
       op.getLoc(), vectorType.getElementType(),
-      builder.getZeroAttr(vectorType.getElementType()));
-  Value result = builder.create<vector::SplatOp>(op.getLoc(), fill, vectorType);
+      rewriter.getZeroAttr(vectorType.getElementType()));
+  Value result =
+      rewriter.create<vector::SplatOp>(op.getLoc(), fill, vectorType);
 
   bool isTransposeLoad = !op.getPermutationMap().isMinorIdentity();
 
   // If we are not transposing, then we can use vectorized loads. Otherwise, we
   // must load each element individually.
   if (!isTransposeLoad) {
     if (!loadedElType.isa<VectorType>()) {
       loadedElType = VectorType::get({1}, loadedElType);
     }
 
     for (int i = 0; i < vectorType.getShape()[0]; i++) {
       FailureOr<AffineMap> coords = nvgpu::getLaneIdAndValueIdToOperandCoord(
-          op.getLoc(), builder, *warpMatrixInfo);
+          rewriter, op.getLoc(), *warpMatrixInfo);
       if (failed(coords))
-        return failure();
-      Value logicalValueId = builder.create<arith::ConstantOp>(
-          loc, builder.getIndexType(),
-          builder.getIndexAttr(i * regInfo->elementsPerRegister));
+        return rewriter.notifyMatchFailure(op, "no coords");
+
+      Value logicalValueId = rewriter.create<arith::ConstantOp>(
+          loc, rewriter.getIndexType(),
+          rewriter.getIndexAttr(i * regInfo->elementsPerRegister));
       SmallVector<Value, 4> newIndices;
       getXferIndices<vector::TransferReadOp>(
-          builder, op, *coords, {laneId, logicalValueId}, newIndices);
+          rewriter, op, *coords, {laneId, logicalValueId}, newIndices);
 
-      Value el = builder.create<vector::LoadOp>(loc, loadedElType,
+      Value el = rewriter.create<vector::LoadOp>(loc, loadedElType,
                                                  op.getSource(), newIndices);
-      result = builder.create<vector::InsertOp>(loc, el, result,
-                                                builder.getI64ArrayAttr(i));
+      result = rewriter.create<vector::InsertOp>(loc, el, result,
+                                                 rewriter.getI64ArrayAttr(i));
     }
   } else {
     if (auto vecType = loadedElType.dyn_cast<VectorType>()) {
       loadedElType = vecType.getElementType();
     }
     for (int i = 0; i < vectorType.getShape()[0]; i++) {
       for (unsigned innerIdx = 0; innerIdx < vectorType.getShape()[1];
            innerIdx++) {
 
-        Value logicalValueId = builder.create<arith::ConstantOp>(
-            loc, builder.getIndexType(),
-            builder.getIndexAttr(i * regInfo->elementsPerRegister + innerIdx));
+        Value logicalValueId = rewriter.create<arith::ConstantOp>(
+            loc, rewriter.getIndexType(),
+            rewriter.getIndexAttr(i * regInfo->elementsPerRegister + innerIdx));
         FailureOr<AffineMap> coords = nvgpu::getLaneIdAndValueIdToOperandCoord(
-            op.getLoc(), builder, *warpMatrixInfo);
+            rewriter, op.getLoc(), *warpMatrixInfo);
         if (failed(coords))
-          return failure();
+          return rewriter.notifyMatchFailure(op, "no coords");
 
         SmallVector<Value, 4> newIndices;
         getXferIndices<vector::TransferReadOp>(
-            builder, op, *coords, {laneId, logicalValueId}, newIndices);
-        Value el = builder.create<memref::LoadOp>(op.getLoc(), loadedElType,
+            rewriter, op, *coords, {laneId, logicalValueId}, newIndices);
+        Value el = rewriter.create<memref::LoadOp>(op.getLoc(), loadedElType,
                                                    op.getSource(), newIndices);
-        result = builder.create<vector::InsertOp>(
-            op.getLoc(), el, result, builder.getI64ArrayAttr({i, innerIdx}));
+        result = rewriter.create<vector::InsertOp>(
+            op.getLoc(), el, result, rewriter.getI64ArrayAttr({i, innerIdx}));
       }
     }
   }
 
   valueMapping[op.getResult()] = result;
   return success();
 }
 
 /// Return true if this is a shared memory memref type.
 static bool isSharedMemory(MemRefType type) {
   auto addressSpace =
       type.getMemorySpace().dyn_cast_or_null<gpu::AddressSpaceAttr>();
   if (addressSpace &&
       addressSpace.getValue() == gpu::GPUDialect::getWorkgroupAddressSpace())
     return true;
   return false;
 }
 
 /// Converts a `vector.transfer_read` operation directly to either a
 /// `vector.load` or a `nvgpu.ldmatrix` operation. This function should only be
 /// used when converting to `nvgpu.mma.sync` operations.
 static LogicalResult
-convertTransferReadToLoads(vector::TransferReadOp op,
+convertTransferReadToLoads(RewriterBase &rewriter, vector::TransferReadOp op,
                            llvm::DenseMap<Value, Value> &valueMapping) {
-  OpBuilder b(op);
+  OpBuilder::InsertionGuard g(rewriter);
+  rewriter.setInsertionPoint(op);
 
   FailureOr<nvgpu::WarpMatrixInfo> warpMatrixInfo =
       nvgpu::getWarpMatrixInfo(op);
   if (failed(warpMatrixInfo))
-    return failure();
+    return rewriter.notifyMatchFailure(op, "no warpMatrixInfo");
 
   bool isLdMatrixCompatible =
       isSharedMemory(op.getSource().getType().cast<MemRefType>()) &&
       nvgpu::inferTileWidthInBits(*warpMatrixInfo) == 128;
 
   VectorType vecTy = op.getVectorType();
   int64_t bitWidth = vecTy.getElementType().getIntOrFloatBitWidth();
 
   // When we are transposing the B operand, ldmatrix will only work if we have
   // at least 8 rows to read and the width to read for the transpose is 128
   // bits.
   if (!op.getPermutationMap().isMinorIdentity() &&
       (bitWidth != 16 || vecTy.getDimSize(1) < 8 ||
        vecTy.getDimSize(0) * bitWidth < 128))
     isLdMatrixCompatible = false;
 
   if (!isLdMatrixCompatible)
-    return createNonLdMatrixLoads(op, b, valueMapping);
+    return createNonLdMatrixLoads(rewriter, op, valueMapping);
 
-  return creatLdMatrixCompatibleLoads(op, b, valueMapping);
+  return creatLdMatrixCompatibleLoads(rewriter, op, valueMapping);
 }
 
 static LogicalResult
-convertTransferWriteToStores(vector::TransferWriteOp op,
+convertTransferWriteToStores(RewriterBase &rewriter, vector::TransferWriteOp op,
                              llvm::DenseMap<Value, Value> &valueMapping) {
-  OpBuilder b(op);
+  OpBuilder::InsertionGuard g(rewriter);
+  rewriter.setInsertionPoint(op);
+
   Location loc = op->getLoc();
-  Value matrix = valueMapping.find(op.getVector())->second;
+  auto it = valueMapping.find(op.getVector());
+  if (it == valueMapping.end())
+    return rewriter.notifyMatchFailure(op, "no mapping");
+  Value matrix = it->second;
 
   FailureOr<nvgpu::WarpMatrixInfo> warpMatrixInfo =
       nvgpu::getWarpMatrixInfo(op);
   if (failed(warpMatrixInfo))
-    return failure();
+    return rewriter.notifyMatchFailure(op, "no warpMatrixInfo");
   FailureOr<nvgpu::FragmentElementInfo> regInfo =
       nvgpu::getMmaSyncRegisterType(*warpMatrixInfo);
   if (failed(regInfo))
-    return failure();
+    return rewriter.notifyMatchFailure(op, "not mma sync reg info");
 
   VectorType vectorType = getMmaSyncVectorOperandType(*regInfo);
-  Value laneId = b.create<gpu::LaneIdOp>(loc);
+  Value laneId = rewriter.create<gpu::LaneIdOp>(loc);
 
   for (unsigned i = 0; i < vectorType.getShape()[0]; i++) {
-    Value logicalValueId = b.create<arith::ConstantOp>(
-        loc, b.getIndexType(),
-        b.getIndexAttr(i * regInfo->elementsPerRegister));
+    Value logicalValueId = rewriter.create<arith::ConstantOp>(
+        loc, rewriter.getIndexType(),
+        rewriter.getIndexAttr(i * regInfo->elementsPerRegister));
     FailureOr<AffineMap> coords = nvgpu::getLaneIdAndValueIdToOperandCoord(
-        op.getLoc(), b, *warpMatrixInfo);
+        rewriter, op.getLoc(), *warpMatrixInfo);
     if (failed(coords))
-      return failure();
+      return rewriter.notifyMatchFailure(op, "no coords");
 
-    Value el = b.create<vector::ExtractOp>(loc, matrix, ArrayRef<int64_t>{i});
+    Value el =
+        rewriter.create<vector::ExtractOp>(loc, matrix, ArrayRef<int64_t>{i});
     SmallVector<Value, 4> newIndices;
     getXferIndices<vector::TransferWriteOp>(
-        b, op, *coords, {laneId, logicalValueId}, newIndices);
-    b.create<vector::StoreOp>(loc, el, op.getSource(), newIndices);
+        rewriter, op, *coords, {laneId, logicalValueId}, newIndices);
+    rewriter.create<vector::StoreOp>(loc, el, op.getSource(), newIndices);
   }
-  op->erase();
+
+  LLVM_DEBUG(DBGS() << "erase: " << op << "\n");
+  rewriter.eraseOp(op);
   return success();
 }
 
 static void populateFromInt64AttrArray(ArrayAttr arrayAttr,
                                        SmallVectorImpl<int64_t> &results) {
   for (auto attr : arrayAttr)
     results.push_back(attr.cast<IntegerAttr>().getInt());
 }
 
 static LogicalResult
-convertExtractStridedSlice(vector::ExtractStridedSliceOp op,
+convertExtractStridedSlice(RewriterBase &rewriter,
+                           vector::ExtractStridedSliceOp op,
                            llvm::DenseMap<Value, Value> &valueMapping) {
+  OpBuilder::InsertionGuard g(rewriter);
+  rewriter.setInsertionPoint(op);
 
-  OpBuilder b(op);
   Location loc = op->getLoc();
 
   FailureOr<nvgpu::WarpMatrixInfo> warpMatrixInfo =
       nvgpu::getWarpMatrixInfo(op);
   if (failed(warpMatrixInfo))
-    return failure();
+    return rewriter.notifyMatchFailure(op, "no warpMatrixInfo");
 
   FailureOr<nvgpu::FragmentElementInfo> mmaSyncFragmentInfo =
       nvgpu::getMmaSyncRegisterType(*warpMatrixInfo);
   if (failed(mmaSyncFragmentInfo))
-    return failure();
+    return rewriter.notifyMatchFailure(op, "no mmaSyncFragmentInfo");
 
   // Find the vector.transer_read whose result vector is being sliced.
   auto transferReadOp = op.getVector().getDefiningOp<vector::TransferReadOp>();
   if (!transferReadOp)
-    return failure();
+    return rewriter.notifyMatchFailure(op, "no transfer read");
 
   warpMatrixInfo = nvgpu::getWarpMatrixInfo(transferReadOp);
   if (failed(warpMatrixInfo))
-    return failure();
+    return rewriter.notifyMatchFailure(op, "no warpMatrixInfo");
 
   FailureOr<nvgpu::FragmentElementInfo> ldFragmentInfo =
       nvgpu::getMmaSyncRegisterType(*warpMatrixInfo);
   if (failed(ldFragmentInfo))
-    return failure();
+    return rewriter.notifyMatchFailure(op, "no ldFragmentInfo");
 
   assert(
       (mmaSyncFragmentInfo->elementsPerRegister ==
        ldFragmentInfo->elementsPerRegister) &&
       "Number of elements per register should be same for load and mma.sync");
 
   // Create vector.extract_strided_slice op for thread-owned fragments.
   std::array<int64_t, 2> strides = {1,
                                     1}; // stride for extract slice is always 1.
   std::array<int64_t, 2> sliceShape = {
       mmaSyncFragmentInfo->numRegistersPerFragment,
       mmaSyncFragmentInfo->elementsPerRegister};
-  auto sourceVector = valueMapping.find(transferReadOp)->second;
+  auto it = valueMapping.find(transferReadOp);
+  if (it == valueMapping.end())
+    return rewriter.notifyMatchFailure(op, "no mapping");
+  auto sourceVector = it->second;
 
   // offset and sizes at warp-level of onwership.
   SmallVector<int64_t> offsets;
   populateFromInt64AttrArray(op.getOffsets(), offsets);
 
   SmallVector<int64_t> sizes;
   populateFromInt64AttrArray(op.getSizes(), sizes);
   ArrayRef<int64_t> warpVectorShape = op.getVectorType().getShape();
 
   // Compute offset in vector registers. Note that the mma.sync vector registers
   // are shaped as numberOfFragments x numberOfRegistersPerfFragment. The vector
   // registers can only be sliced along numberOfFragments, i.e., sliceOffset[0].
   std::array<int64_t, 2> sliceOffset = {0, 0};
 
   if (offsets[0] && offsets[1])
     return op->emitError() << "Slicing fragments in 2D is not supported. ";
   if (offsets[0])
     sliceOffset[0] = (warpVectorShape[0] / offsets[0]);
   else if (offsets[1])
     sliceOffset[0] = (warpVectorShape[1] / offsets[1]);
 
-  Value newOp = b.create<vector::ExtractStridedSliceOp>(
+  Value newOp = rewriter.create<vector::ExtractStridedSliceOp>(
       loc, sourceVector, sliceOffset, sliceShape, strides);
 
   valueMapping[op] = newOp;
   return success();
 }
 
-static void convertContractOp(vector::ContractionOp op,
-                              llvm::DenseMap<Value, Value> &valueMapping) {
-  OpBuilder b(op);
-  Value opA = valueMapping.find(op.getLhs())->second;
-  Value opB = valueMapping.find(op.getRhs())->second;
-  Value opC = valueMapping.find(op.getAcc())->second;
-  Value matmul = b.create<gpu::SubgroupMmaComputeOp>(
+static LogicalResult
+convertContractOp(RewriterBase &rewriter, vector::ContractionOp op,
+  OpBuilder::InsertionGuard g(rewriter);
+  rewriter.setInsertionPoint(op);
+
+  auto itA = valueMapping.find(op.getLhs());
+  auto itB = valueMapping.find(op.getRhs());
+  auto itC = valueMapping.find(op.getAcc());
+  if (itA == valueMapping.end() || itB == valueMapping.end() ||
+      itC == valueMapping.end())
+    return rewriter.notifyMatchFailure(op, "no mapping");
+  Value opA = itA->second, opB = itB->second, opC = itC->second;
+  Value matmul = rewriter.create<gpu::SubgroupMmaComputeOp>(
       op.getLoc(), opC.getType(), opA, opB, opC, /*a_transpose=*/UnitAttr(),
       /*b_transpose=*/UnitAttr());
   valueMapping[op.getResult()] = matmul;
+  return success();
 }
 
 static LogicalResult
-convertContractOpToMmaSync(vector::ContractionOp op,
+convertContractOpToMmaSync(RewriterBase &rewriter, vector::ContractionOp op,
                            llvm::DenseMap<Value, Value> &valueMapping) {
-  OpBuilder b(op);
-  Value opA = valueMapping.find(op.getLhs())->second;
-  Value opB = valueMapping.find(op.getRhs())->second;
-  Value opC = valueMapping.find(op.getAcc())->second;
+  OpBuilder::InsertionGuard g(rewriter);
+  rewriter.setInsertionPoint(op);
+
+  auto itA = valueMapping.find(op.getLhs());
+  auto itB = valueMapping.find(op.getRhs());
+  auto itC = valueMapping.find(op.getAcc());
+  if (itA == valueMapping.end() || itB == valueMapping.end() ||
+      itC == valueMapping.end())
+    return rewriter.notifyMatchFailure(op, "no mapping");
+  Value opA = itA->second, opB = itB->second, opC = itC->second;
   int64_t m = op.getLhs().getType().cast<VectorType>().getShape()[0];
   int64_t n = op.getRhs().getType().cast<VectorType>().getShape()[0];
   int64_t k = op.getLhs().getType().cast<VectorType>().getShape()[1];
-  Value matmul = b.create<nvgpu::MmaSyncOp>(op.getLoc(), opA, opB, opC,
-                                            b.getI64ArrayAttr({m, n, k}));
+  Value matmul = rewriter.create<nvgpu::MmaSyncOp>(
+      op.getLoc(), opA, opB, opC, rewriter.getI64ArrayAttr({m, n, k}));
   valueMapping[op.getResult()] = matmul;
   return success();
 }
 
 /// Convert a 2D splat ConstantOp to a SubgroupMmaConstantMatrix op.
-static void convertConstantOp(arith::ConstantOp op,
-                              llvm::DenseMap<Value, Value> &valueMapping) {
+static LogicalResult
+convertConstantOp(RewriterBase &rewriter, arith::ConstantOp op,
+  OpBuilder::InsertionGuard g(rewriter);
+  rewriter.setInsertionPoint(op);
+
   assert(constantSupportsMMAMatrixType(op));
-  OpBuilder b(op);
+
   auto splat =
       op.getValue().cast<SplatElementsAttr>().getSplatValue<TypedAttr>();
   auto scalarConstant =
-      b.create<arith::ConstantOp>(op.getLoc(), splat.getType(), splat);
+      rewriter.create<arith::ConstantOp>(op.getLoc(), splat.getType(), splat);
   const char *fragType = inferFragType(op);
   auto vecType = op.getType().cast<VectorType>();
   gpu::MMAMatrixType type = gpu::MMAMatrixType::get(
       vecType.getShape(), vecType.getElementType(), llvm::StringRef(fragType));
-  auto matrix = b.create<gpu::SubgroupMmaConstantMatrixOp>(op.getLoc(), type,
-                                                           scalarConstant);
+  auto matrix = rewriter.create<gpu::SubgroupMmaConstantMatrixOp>(
+      op.getLoc(), type, scalarConstant);
   valueMapping[op.getResult()] = matrix;
+  return success();
 }
 
 /// Convert a vector.broadcast from scalar to a SubgroupMmaConstantMatrix op.
-static void convertBroadcastOp(vector::BroadcastOp op,
-                               llvm::DenseMap<Value, Value> &valueMapping) {
+static LogicalResult
+convertBroadcastOp(RewriterBase &rewriter, vector::BroadcastOp op,
+  OpBuilder::InsertionGuard g(rewriter);
+  rewriter.setInsertionPoint(op);
+
   assert(broadcastSupportsMMAMatrixType(op));
-  OpBuilder b(op);
+
   const char *fragType = inferFragType(op);
   auto vecType = op.getVectorType();
   gpu::MMAMatrixType type = gpu::MMAMatrixType::get(
       vecType.getShape(), vecType.getElementType(), llvm::StringRef(fragType));
-  auto matrix = b.create<gpu::SubgroupMmaConstantMatrixOp>(op.getLoc(), type,
-                                                           op.getSource());
+  auto matrix = rewriter.create<gpu::SubgroupMmaConstantMatrixOp>(
+      op.getLoc(), type, op.getSource());
   valueMapping[op.getResult()] = matrix;
+  return success();
 }
 
 // Replace ForOp with a new ForOp with extra operands. The YieldOp is not
 // updated and needs to be updated separatly for the loop to be correct.
-static scf::ForOp replaceForOpWithNewSignature(OpBuilder &b, scf::ForOp loop,
+static scf::ForOp replaceForOpWithNewSignature(RewriterBase &rewriter,
+                                               scf::ForOp loop,
                                                ValueRange newIterOperands) {
+  OpBuilder::InsertionGuard g(rewriter);
+  rewriter.setInsertionPoint(loop);
+
   // Create a new loop before the existing one, with the extra operands.
-  OpBuilder::InsertionGuard g(b);
-  b.setInsertionPoint(loop);
+  rewriter.setInsertionPoint(loop);
   auto operands = llvm::to_vector<4>(loop.getIterOperands());
   operands.append(newIterOperands.begin(), newIterOperands.end());
-  scf::ForOp newLoop =
-      b.create<scf::ForOp>(loop.getLoc(), loop.getLowerBound(),
-                           loop.getUpperBound(), loop.getStep(), operands);
+  scf::ForOp newLoop = rewriter.create<scf::ForOp>(
+      loop.getLoc(), loop.getLowerBound(), loop.getUpperBound(), loop.getStep(),
+      operands);
   newLoop.getBody()->erase();
+
   newLoop.getLoopBody().getBlocks().splice(
       newLoop.getLoopBody().getBlocks().begin(),
       loop.getLoopBody().getBlocks());
   for (Value operand : newIterOperands)
     newLoop.getBody()->addArgument(operand.getType(), operand.getLoc());
 
   for (auto it : llvm::zip(loop.getResults(), newLoop.getResults().take_front(
                                                   loop.getNumResults())))
-    std::get<0>(it).replaceAllUsesWith(std::get<1>(it));
-  loop.erase();
+    rewriter.replaceAllUsesWith(std::get<0>(it), std::get<1>(it));
+
+  LLVM_DEBUG(DBGS() << "newLoop now: " << newLoop << "\n");
+  LLVM_DEBUG(DBGS() << "stripped scf.for: " << loop << "\n");
+  LLVM_DEBUG(DBGS() << "erase: " << loop);
+
+  rewriter.eraseOp(loop);
   return newLoop;
 }
 
-static void convertForOp(scf::ForOp op,
+static LogicalResult convertForOp(RewriterBase &rewriter, scf::ForOp op,
                                   llvm::DenseMap<Value, Value> &valueMapping) {
+  OpBuilder::InsertionGuard g(rewriter);
+  rewriter.setInsertionPoint(op);
+
   SmallVector<Value> newOperands;
   SmallVector<std::pair<size_t, size_t>> argMapping;
   for (const auto &operand : llvm::enumerate(op.getIterOperands())) {
     auto it = valueMapping.find(operand.value());
-    if (it == valueMapping.end())
+    if (it == valueMapping.end()) {
+      LLVM_DEBUG(DBGS() << "no value mapping for: " << operand.value() << "\n");
       continue;
+    }
     argMapping.push_back(std::make_pair(
         operand.index(), op.getNumIterOperands() + newOperands.size()));
     newOperands.push_back(it->second);
   }
-  OpBuilder b(op);
-  scf::ForOp newForOp = replaceForOpWithNewSignature(b, op, newOperands);
+
+  scf::ForOp newForOp = replaceForOpWithNewSignature(rewriter, op, newOperands);
   Block &loopBody = *newForOp.getBody();
   for (auto mapping : argMapping) {
     valueMapping[newForOp.getResult(mapping.first)] =
         newForOp.getResult(mapping.second);
     valueMapping[loopBody.getArgument(mapping.first +
                                       newForOp.getNumInductionVars())] =
         loopBody.getArgument(mapping.second + newForOp.getNumInductionVars());
   }
+
+  LLVM_DEBUG(DBGS() << "scf.for to: " << newForOp << "\n");
+  return success();
 }
 
-static void convertYieldOp(scf::YieldOp op,
-                           llvm::DenseMap<Value, Value> &valueMapping) {
-  OpBuilder b(op);
+static LogicalResult
+convertYieldOp(RewriterBase &rewriter, scf::YieldOp op,
+  OpBuilder::InsertionGuard g(rewriter);
+  rewriter.setInsertionPoint(op);
+
   auto loop = cast<scf::ForOp>(op->getParentOp());
   auto yieldOperands = llvm::to_vector<4>(op.getOperands());
   for (const auto &operand : llvm::enumerate(op.getOperands())) {
     auto it = valueMapping.find(operand.value());
     if (it == valueMapping.end())
       continue;
     // Replace the yield of old value with the for op argument to make it easier
     // to remove the dead code.
     yieldOperands[operand.index()] = loop.getIterOperands()[operand.index()];
     yieldOperands.push_back(it->second);
   }
-  b.create<scf::YieldOp>(op.getLoc(), yieldOperands);
-  op.erase();
+  rewriter.create<scf::YieldOp>(op.getLoc(), yieldOperands);
+
+  LLVM_DEBUG(DBGS() << "erase: " << op << "\n");
+  rewriter.eraseOp(op);
+  return success();
 }
 
 /// Convert an elementwise op to the equivalent elementwise op on MMA matrix.
-static void convertElementwiseOp(Operation *op, gpu::MMAElementwiseOp opType,
-                                 llvm::DenseMap<Value, Value> &valueMapping) {
-  OpBuilder b(op);
+static LogicalResult
+convertElementwiseOp(RewriterBase &rewriter, Operation *op,
+                     gpu::MMAElementwiseOp opType,
+  OpBuilder::InsertionGuard g(rewriter);
+  rewriter.setInsertionPoint(op);
+
   SmallVector<Value> matrixOperands;
-  for (Value operand : op->getOperands())
-    matrixOperands.push_back(valueMapping.find(operand)->second);
-  Value newOp = b.create<gpu::SubgroupMmaElementwiseOp>(
+  for (Value operand : op->getOperands()) {
+    auto it = valueMapping.find(operand);
+    if (it == valueMapping.end())
+      return rewriter.notifyMatchFailure(op, "no mapping");
+    matrixOperands.push_back(it->second);
+  }
+  Value newOp = rewriter.create<gpu::SubgroupMmaElementwiseOp>(
       op->getLoc(), matrixOperands[0].getType(), matrixOperands, opType);
   valueMapping[op->getResult(0)] = newOp;
+  return success();
 }
 
 void mlir::populatePrepareVectorToMMAPatterns(RewritePatternSet &patterns,
                                               bool useNvGpu) {
   if (!useNvGpu) {
     patterns.add<PrepareContractToGPUMMA, CombineTransferReadOpTranspose>(
         patterns.getContext());
     return;
   }
   patterns
       .add<nvgpu::PrepareContractToGPUMMASync, CombineTransferReadOpTranspose>(
           patterns.getContext());
 }
 
-void mlir::convertVectorToMMAOps(Operation *rootOp) {
+LogicalResult mlir::convertVectorToMMAOps(RewriterBase &rewriter,
+                                          Operation *rootOp) {
   SetVector<Operation *> ops = getOpToConvert(rootOp, /*useNvGpu=*/false);
   llvm::DenseMap<Value, Value> valueMapping;
+
+  auto globalRes = LogicalResult::success();
   for (Operation *op : ops) {
+    LLVM_DEBUG(DBGS() << "Process op: " << *op << "\n");
+    // Apparently callers do not want to early exit on failure here.
+    auto res = LogicalResult::success();
     if (auto transferRead = dyn_cast<vector::TransferReadOp>(op)) {
-      convertTransferReadOp(transferRead, valueMapping);
+      res = convertTransferReadOp(rewriter, transferRead, valueMapping);
     } else if (auto transferWrite = dyn_cast<vector::TransferWriteOp>(op)) {
-      convertTransferWriteOp(transferWrite, valueMapping);
+      res = convertTransferWriteOp(rewriter, transferWrite, valueMapping);
     } else if (auto contractOp = dyn_cast<vector::ContractionOp>(op)) {
-      convertContractOp(contractOp, valueMapping);
+      res = convertContractOp(rewriter, contractOp, valueMapping);
     } else if (auto constantOp = dyn_cast<arith::ConstantOp>(op)) {
-      convertConstantOp(constantOp, valueMapping);
+      res = convertConstantOp(rewriter, constantOp, valueMapping);
     } else if (auto broadcastOp = dyn_cast<vector::BroadcastOp>(op)) {
-      convertBroadcastOp(broadcastOp, valueMapping);
+      res = convertBroadcastOp(rewriter, broadcastOp, valueMapping);
     } else if (auto forOp = dyn_cast<scf::ForOp>(op)) {
-      convertForOp(forOp, valueMapping);
-    } else if (auto yiledOp = dyn_cast<scf::YieldOp>(op)) {
-      convertYieldOp(yiledOp, valueMapping);
+      res = convertForOp(rewriter, forOp, valueMapping);
+    } else if (auto yieldOp = dyn_cast<scf::YieldOp>(op)) {
+      res = convertYieldOp(rewriter, yieldOp, valueMapping);
     } else if (auto elementwiseType = convertElementwiseOpToMMA(op)) {
-      convertElementwiseOp(op, *elementwiseType, valueMapping);
+      res = convertElementwiseOp(rewriter, op, *elementwiseType, valueMapping);
     }
+    if (failed(res))
+      globalRes = failure();
   }
+  return globalRes;
 }
 
-LogicalResult mlir::convertVectorToNVVMCompatibleMMASync(Operation *rootOp) {
+LogicalResult mlir::convertVectorToNVVMCompatibleMMASync(RewriterBase &rewriter,
+                                                         Operation *rootOp) {
   SetVector<Operation *> ops = getOpToConvert(rootOp, /*useNvGpu=*/true);
   llvm::DenseMap<Value, Value> valueMapping;
   for (Operation *op : ops) {
     if (llvm::TypeSwitch<Operation *, LogicalResult>(op)
             .Case([&](vector::TransferReadOp transferReadOp) {
-              return convertTransferReadToLoads(transferReadOp, valueMapping);
+              return convertTransferReadToLoads(rewriter, transferReadOp,
+                                                valueMapping);
             })
             .Case([&](vector::TransferWriteOp transferWriteOp) {
-              return convertTransferWriteToStores(transferWriteOp,
+              return convertTransferWriteToStores(rewriter, transferWriteOp,
                                                   valueMapping);
             })
             .Case([&](vector::ExtractStridedSliceOp extractStridedSliceOp) {
-              return convertExtractStridedSlice(extractStridedSliceOp,
+              return convertExtractStridedSlice(rewriter, extractStridedSliceOp,
                                                 valueMapping);
             })
             .Case([&](vector::ContractionOp contractionOp) {
-              return convertContractOpToMmaSync(contractionOp, valueMapping);
+              return convertContractOpToMmaSync(rewriter, contractionOp,
+                                                valueMapping);
             })
             .Case([&](scf::ForOp forOp) {
-              convertForOp(forOp, valueMapping);
-              return success();
+              return convertForOp(rewriter, forOp, valueMapping);
             })
             .Case([&](scf::YieldOp yieldOp) {
-              convertYieldOp(yieldOp, valueMapping);
-              return success();
+              return convertYieldOp(rewriter, yieldOp, valueMapping);
             })
             .Case([&](arith::ConstantOp constOp) {
-              return convertConstantOpMmaSync(constOp, valueMapping);
+              return convertConstantOpMmaSync(rewriter, constOp, valueMapping);
             })
             .Default([&](Operation *op) {
-              op->emitError() << "unhandled vector to mma type: " << *op;
-              return failure();
+              return op->emitError() << "unhandled vector to mma type: " << *op;
             })
             .failed()) {
-      op->emitError() << "Failed to convert op " << *op;
-      return failure();
+      return op->emitError() << "Failed to convert op " << *op;
     }
   }
   return success();
 }
 
 namespace {
 
 struct ConvertVectorToGPUPass
     : public impl::ConvertVectorToGPUBase<ConvertVectorToGPUPass> {
 
   explicit ConvertVectorToGPUPass(bool useNvGpu_) {
     useNvGpu.setValue(useNvGpu_);
   }
 
   void runOnOperation() override {
     RewritePatternSet patterns(&getContext());
     populatePrepareVectorToMMAPatterns(patterns, useNvGpu.getValue());
     if (failed(
             applyPatternsAndFoldGreedily(getOperation(), std::move(patterns))))
       return signalPassFailure();
 
+    IRRewriter rewriter(&getContext());
     if (useNvGpu.getValue()) {
-      if (failed(convertVectorToNVVMCompatibleMMASync(getOperation())))
+      if (failed(
+              convertVectorToNVVMCompatibleMMASync(rewriter, getOperation())))
         return signalPassFailure();
     }
-
-    (void)convertVectorToMMAOps(getOperation());
+    (void)convertVectorToMMAOps(rewriter, getOperation());
   }
 };
 
 } // namespace
 
 std::unique_ptr<Pass> mlir::createConvertVectorToGPUPass(bool useNvGpu) {
   return std::make_unique<ConvertVectorToGPUPass>(useNvGpu);
 }

mlir/lib/Dialect/NVGPU/Utils/MMAUtils.cpp

static AffineMap getRegisterIndexToTileOffsetMap(int64_t lineSize,
   return AffineMap::get(
       2, 0,
       {(registerIdx % num8x128bTiles[0]) * 8,
        (registerIdx.floorDiv(num8x128bTiles[0])) * elementsPerLine},
       elementType.getContext());
 }
 
 FailureOr<AffineMap>
-nvgpu::getLaneIdAndValueIdToOperandCoord(Location loc, OpBuilder &builder,
+nvgpu::getLaneIdAndValueIdToOperandCoord(OpBuilder &builder, Location loc,
                                          const WarpMatrixInfo &fragmentType) {
   Type elementType = fragmentType.vectorType.getElementType();
   ArrayRef<int64_t> operandShape = fragmentType.vectorType.getShape();
   FailureOr<nvgpu::FragmentElementInfo> regInfo =
       getMmaSyncRegisterType(fragmentType);
   if (failed(regInfo))
     return failure();
 
nvgpu::getLdMatrixParams(const WarpMatrixInfo &type, bool transpose) {
 
   if (params.numTiles == 0)
     return failure();
 
   return params;
 }
 
 FailureOr<AffineMap>
-nvgpu::getLaneIdToLdMatrixMatrixCoord(Location loc, OpBuilder &builder,
+nvgpu::getLaneIdToLdMatrixMatrixCoord(OpBuilder &builder, Location loc,
                                       const LdMatrixParams &params) {
   // One thread per 128b row.
   const int bitsPerElement = static_cast<int>(
       params.fragmentType.getElementType().getIntOrFloatBitWidth());
   const int kElementsPer128b = (128 / bitsPerElement);
   ArrayRef<int64_t> operandShape = params.fragmentType.getShape();
   AffineExpr d0 = getAffineDimExpr(0, builder.getContext());
 

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

-// RUN: mlir-opt %s -pass-pipeline="builtin.module(func.func(convert-vector-to-gpu),canonicalize)" | FileCheck %s
+// RUN: mlir-opt %s -pass-pipeline="builtin.module(func.func(convert-vector-to-gpu),canonicalize)" --split-input-file | FileCheck %s
 
 #map0 = affine_map<(d0, d1) -> (d1, d0)>
 #map1 = affine_map<(d0, d1, d2) -> (d0, d2)>
 #map2 = affine_map<(d0, d1, d2) -> (d1, d2)>
 #map3 = affine_map<(d0, d1, d2) -> (d0, d1)>
 #map4 = affine_map<(d0) -> (d0, 0)>
 #map5 = affine_map<(d0, d1) -> (d0, d1)>
 
func.func @matmul(%arg0: memref<16x16xf16>, %arg1: memref<16x16xf16>, %arg2: memref<16x16xf16>) {
   %A = vector.transfer_read %arg0[%c0, %c0], %cst {in_bounds = [true, true]} : memref<16x16xf16>, vector<16x16xf16>
   %B = vector.transfer_read %arg1[%c0, %c0], %cst {permutation_map = #map0, in_bounds = [true, true]} : memref<16x16xf16>, vector<16x16xf16>
   %C = vector.transfer_read %arg2[%c0, %c0], %cst {in_bounds = [true, true]} : memref<16x16xf16>, vector<16x16xf16>
   %D = vector.contract {indexing_maps = [#map1, #map2, #map3], iterator_types = ["parallel", "parallel", "reduction"], kind = #vector.kind<add>} %A, %B, %C : vector<16x16xf16>, vector<16x16xf16> into vector<16x16xf16>
   vector.transfer_write %D, %arg2[%c0, %c0] {in_bounds = [true, true]} : vector<16x16xf16>, memref<16x16xf16>
   return
 }
 
+// -----
+
+#map0 = affine_map<(d0, d1) -> (d1, d0)>
+#map1 = affine_map<(d0, d1, d2) -> (d0, d2)>
+#map2 = affine_map<(d0, d1, d2) -> (d1, d2)>
+#map3 = affine_map<(d0, d1, d2) -> (d0, d1)>
+#map4 = affine_map<(d0) -> (d0, 0)>
+#map5 = affine_map<(d0, d1) -> (d0, d1)>
+
 // CHECK-LABEL: func @matmul_cst
 //   CHECK-DAG:   %[[CST:.+]] = arith.constant 0.000000e+00 : f16
 //   CHECK-DAG:   %[[A:.+]] = gpu.subgroup_mma_load_matrix %{{.*}}[%{{.*}}, %{{.*}}] {leadDimension = 16 : index} : memref<16x16xf16> -> !gpu.mma_matrix<16x16xf16, "AOp">
 //   CHECK-DAG:   %[[B:.+]] = gpu.subgroup_mma_load_matrix %{{.*}}[%{{.*}}, %{{.*}}] {leadDimension = 16 : index} : memref<16x16xf16> -> !gpu.mma_matrix<16x16xf16, "BOp">
 //   CHECK-DAG:   %[[C:.+]] = gpu.subgroup_mma_constant_matrix %[[CST]] : !gpu.mma_matrix<16x16xf16, "COp">
 //       CHECK:   %[[D:.+]] = gpu.subgroup_mma_compute %[[A]], %[[B]], %[[C]] : !gpu.mma_matrix<16x16xf16, "AOp">, !gpu.mma_matrix<16x16xf16, "BOp"> -> !gpu.mma_matrix<16x16xf16, "COp">
 //       CHECK:   gpu.subgroup_mma_store_matrix %[[D]], %{{.*}}[%{{.*}}, %{{.*}}] {leadDimension = 16 : index} : !gpu.mma_matrix<16x16xf16, "COp">, memref<16x16xf16>
 func.func @matmul_cst(%arg0: memref<16x16xf16>, %arg1: memref<16x16xf16>, %arg2: memref<16x16xf16>) {
   %cst_0 = arith.constant dense<0.000000e+00> : vector<16x16xf16>
   %c0 = arith.constant 0 : index
   %cst = arith.constant 0.000000e+00 : f16
   %A = vector.transfer_read %arg0[%c0, %c0], %cst {in_bounds = [true, true]} : memref<16x16xf16>, vector<16x16xf16>
   %B = vector.transfer_read %arg1[%c0, %c0], %cst {permutation_map = #map0, in_bounds = [true, true]} : memref<16x16xf16>, vector<16x16xf16>
   %D = vector.contract {indexing_maps = [#map1, #map2, #map3], iterator_types = ["parallel", "parallel", "reduction"], kind = #vector.kind<add>} %A, %B, %cst_0 : vector<16x16xf16>, vector<16x16xf16> into vector<16x16xf16>
   vector.transfer_write %D, %arg2[%c0, %c0] {in_bounds = [true, true]} : vector<16x16xf16>, memref<16x16xf16>
   return
 }
 
+// -----
+
+#map0 = affine_map<(d0, d1) -> (d1, d0)>
+#map1 = affine_map<(d0, d1, d2) -> (d0, d2)>
+#map2 = affine_map<(d0, d1, d2) -> (d1, d2)>
+#map3 = affine_map<(d0, d1, d2) -> (d0, d1)>
+#map4 = affine_map<(d0) -> (d0, 0)>
+#map5 = affine_map<(d0, d1) -> (d0, d1)>
+
 // CHECK-LABEL: func @matmul_broadcast
 //  CHECK-SAME:   (%{{.*}}: memref<16x16xf16>, %{{.*}}: memref<16x16xf16>, %{{.*}}: memref<16x16xf16>, %[[F:.*]]: f16)
 //   CHECK-DAG:   %[[C:.+]] = gpu.subgroup_mma_constant_matrix %[[F]] : !gpu.mma_matrix<16x16xf16, "COp">
 //   CHECK-DAG:   %[[A:.+]] = gpu.subgroup_mma_load_matrix %{{.*}}[%{{.*}}, %{{.*}}] {leadDimension = 16 : index} : memref<16x16xf16> -> !gpu.mma_matrix<16x16xf16, "AOp">
 //   CHECK-DAG:   %[[B:.+]] = gpu.subgroup_mma_load_matrix %{{.*}}[%{{.*}}, %{{.*}}] {leadDimension = 16 : index} : memref<16x16xf16> -> !gpu.mma_matrix<16x16xf16, "BOp">
 //       CHECK:   %[[D:.+]] = gpu.subgroup_mma_compute %[[A]], %[[B]], %[[C]] : !gpu.mma_matrix<16x16xf16, "AOp">, !gpu.mma_matrix<16x16xf16, "BOp"> -> !gpu.mma_matrix<16x16xf16, "COp">
 //       CHECK:   gpu.subgroup_mma_store_matrix %[[D]], %{{.*}}[%{{.*}}, %{{.*}}] {leadDimension = 16 : index} : !gpu.mma_matrix<16x16xf16, "COp">, memref<16x16xf16>
 func.func @matmul_broadcast(%arg0: memref<16x16xf16>, %arg1: memref<16x16xf16>, %arg2: memref<16x16xf16>, %f: f16) {
   %C = vector.broadcast %f : f16 to vector<16x16xf16>
   %c0 = arith.constant 0 : index
   %cst = arith.constant 0.000000e+00 : f16
   %A = vector.transfer_read %arg0[%c0, %c0], %cst {in_bounds = [true, true]} : memref<16x16xf16>, vector<16x16xf16>
   %B = vector.transfer_read %arg1[%c0, %c0], %cst {permutation_map = #map0, in_bounds = [true, true]} : memref<16x16xf16>, vector<16x16xf16>
   %D = vector.contract {indexing_maps = [#map1, #map2, #map3], iterator_types = ["parallel", "parallel", "reduction"], kind = #vector.kind<add>} %A, %B, %C : vector<16x16xf16>, vector<16x16xf16> into vector<16x16xf16>
   vector.transfer_write %D, %arg2[%c0, %c0] {in_bounds = [true, true]} : vector<16x16xf16>, memref<16x16xf16>
   return
 }
 
+// -----
+
+#map0 = affine_map<(d0, d1) -> (d1, d0)>
+#map1 = affine_map<(d0, d1, d2) -> (d0, d2)>
+#map2 = affine_map<(d0, d1, d2) -> (d1, d2)>
+#map3 = affine_map<(d0, d1, d2) -> (d0, d1)>
+#map4 = affine_map<(d0) -> (d0, 0)>
+#map5 = affine_map<(d0, d1) -> (d0, d1)>
+
 // CHECK-LABEL: func @matmul_loop
 //       CHECK:   %[[C:.+]] = gpu.subgroup_mma_load_matrix %{{.*}}[%{{.*}}, %{{.*}}] {leadDimension = 128 : index} : memref<128x128xf16> -> !gpu.mma_matrix<16x16xf16, "COp">
 //       CHECK:   %[[ACC:.+]] = scf.for {{.*}} iter_args(%[[ACC1:.+]] = %[[C]]) -> (!gpu.mma_matrix<16x16xf16, "COp">) {
 //   CHECK-DAG:     %[[A:.+]] = gpu.subgroup_mma_load_matrix %{{.*}}[%{{.*}}, %{{.*}}] {leadDimension = 128 : index} : memref<128x128xf16> -> !gpu.mma_matrix<16x16xf16, "AOp">
 //   CHECK-DAG:     %[[B:.+]] = gpu.subgroup_mma_load_matrix %{{.*}}[%{{.*}}, %{{.*}}] {leadDimension = 128 : index} : memref<128x128xf16> -> !gpu.mma_matrix<16x16xf16, "BOp">
 //  CHECK-NEXT:     %[[D:.+]] = gpu.subgroup_mma_compute %[[A]], %[[B]], %[[ACC1]] : !gpu.mma_matrix<16x16xf16, "AOp">, !gpu.mma_matrix<16x16xf16, "BOp"> -> !gpu.mma_matrix<16x16xf16, "COp">
 //  CHECK-NEXT:     scf.yield %[[D]] : !gpu.mma_matrix<16x16xf16, "COp">
 //  CHECK-NEXT:   }
%14 = scf.for %arg17 = %c0 to %c128 step %c32 iter_args(%arg18 = %C) -> (vector<16x16xf16>) {
     %18 = vector.transfer_read %arg1[%arg17, %c0], %cst {permutation_map = #map0, in_bounds = [true, true]} : memref<128x128xf16>, vector<16x16xf16>
     %19 = vector.contract {indexing_maps = [#map1, #map2, #map3], iterator_types = ["parallel", "parallel", "reduction"], kind = #vector.kind<add>} %17, %18, %arg18 : vector<16x16xf16>, vector<16x16xf16> into vector<16x16xf16>
     scf.yield %19 : vector<16x16xf16>
   }
   vector.transfer_write %14, %arg2[%c0, %c0] {in_bounds = [true, true]} : vector<16x16xf16>, memref<128x128xf16>
   return
 }
 
+// -----
+
+#map0 = affine_map<(d0, d1) -> (d1, d0)>
+#map1 = affine_map<(d0, d1, d2) -> (d0, d2)>
+#map2 = affine_map<(d0, d1, d2) -> (d1, d2)>
+#map3 = affine_map<(d0, d1, d2) -> (d0, d1)>
+#map4 = affine_map<(d0) -> (d0, 0)>
+#map5 = affine_map<(d0, d1) -> (d0, d1)>
+
 // CHECK-LABEL: func @matmul_fused_elementwise
 //   CHECK-DAG:   %[[CST_0:.+]] = arith.constant 0.000000e+00 : f16
 //   CHECK-DAG:   %[[CST_1:.+]] = arith.constant 1.000000e+00 : f16
 //   CHECK-DAG:   %[[A:.+]] = gpu.subgroup_mma_load_matrix %{{.*}}[%{{.*}}, %{{.*}}] {leadDimension = 16 : index} : memref<16x16xf16> -> !gpu.mma_matrix<16x16xf16, "AOp">
 //   CHECK-DAG:   %[[B:.+]] = gpu.subgroup_mma_load_matrix %{{.*}}[%{{.*}}, %{{.*}}] {leadDimension = 16 : index} : memref<16x16xf16> -> !gpu.mma_matrix<16x16xf16, "BOp">
 //   CHECK-DAG:   %[[C0:.+]] = gpu.subgroup_mma_constant_matrix %[[CST_0]] : !gpu.mma_matrix<16x16xf16, "COp">
 //   CHECK-DAG:   %[[C1:.+]] = gpu.subgroup_mma_constant_matrix %[[CST_1]] : !gpu.mma_matrix<16x16xf16, "COp">
 //       CHECK:   %[[D:.+]] = gpu.subgroup_mma_compute %[[A]], %[[B]], %[[C0]] : !gpu.mma_matrix<16x16xf16, "AOp">, !gpu.mma_matrix<16x16xf16, "BOp"> -> !gpu.mma_matrix<16x16xf16, "COp">
 //       CHECK:   %[[E:.+]] = gpu.subgroup_mma_elementwise addf %[[D]], %[[C1]] : (!gpu.mma_matrix<16x16xf16, "COp">, !gpu.mma_matrix<16x16xf16, "COp">) -> !gpu.mma_matrix<16x16xf16, "COp">
 //       CHECK:   gpu.subgroup_mma_store_matrix %[[E]], %{{.*}}[%{{.*}}, %{{.*}}] {leadDimension = 16 : index} : !gpu.mma_matrix<16x16xf16, "COp">, memref<16x16xf16>
 func.func @matmul_fused_elementwise(%arg0: memref<16x16xf16>, %arg1: memref<16x16xf16>, %arg2: memref<16x16xf16>) {
   %cst_0 = arith.constant dense<0.000000e+00> : vector<16x16xf16>
   %cst_1 = arith.constant dense<1.000000e+00> : vector<16x16xf16>
   %c0 = arith.constant 0 : index
   %cst = arith.constant 0.000000e+00 : f16
   %A = vector.transfer_read %arg0[%c0, %c0], %cst {in_bounds = [true, true]} : memref<16x16xf16>, vector<16x16xf16>
   %B = vector.transfer_read %arg1[%c0, %c0], %cst {permutation_map = #map0, in_bounds = [true, true]} : memref<16x16xf16>, vector<16x16xf16>
   %D = vector.contract {indexing_maps = [#map1, #map2, #map3], iterator_types = ["parallel", "parallel", "reduction"], kind = #vector.kind<add>} %A, %B, %cst_0 : vector<16x16xf16>, vector<16x16xf16> into vector<16x16xf16>
   %E = arith.addf %D, %cst_1 : vector<16x16xf16>
   vector.transfer_write %E, %arg2[%c0, %c0] {in_bounds = [true, true]} : vector<16x16xf16>, memref<16x16xf16>
   return
 }
 
+// -----
+
+#map0 = affine_map<(d0, d1) -> (d1, d0)>
+#map1 = affine_map<(d0, d1, d2) -> (d0, d2)>
+#map2 = affine_map<(d0, d1, d2) -> (d1, d2)>
+#map3 = affine_map<(d0, d1, d2) -> (d0, d1)>
+#map4 = affine_map<(d0) -> (d0, 0)>
+#map5 = affine_map<(d0, d1) -> (d0, d1)>
+
 // CHECK-LABEL: func @matmul_fused_broadcast
 //   CHECK-DAG:   %[[CST_0:.+]] = arith.constant 0.000000e+00 : f16
 //   CHECK-DAG:   %[[A:.+]] = gpu.subgroup_mma_load_matrix %{{.*}}[%{{.*}}, %{{.*}}] {leadDimension = 16 : index} : memref<16x16xf16> -> !gpu.mma_matrix<16x16xf16, "AOp">
 //   CHECK-DAG:   %[[B:.+]] = gpu.subgroup_mma_load_matrix %{{.*}}[%{{.*}}, %{{.*}}] {leadDimension = 16 : index} : memref<16x16xf16> -> !gpu.mma_matrix<16x16xf16, "BOp">
 //   CHECK-DAG:   %[[C0:.+]] = gpu.subgroup_mma_constant_matrix %[[CST_0]] : !gpu.mma_matrix<16x16xf16, "COp">
 //       CHECK:   %[[D:.+]] = gpu.subgroup_mma_compute %[[A]], %[[B]], %[[C0]] : !gpu.mma_matrix<16x16xf16, "AOp">, !gpu.mma_matrix<16x16xf16, "BOp"> -> !gpu.mma_matrix<16x16xf16, "COp">
 //       CHECK:   %[[E:.+]] = gpu.subgroup_mma_load_matrix %{{.*}}[%{{.*}}, %{{.*}}, %{{.*}}, %{{.*}}] {leadDimension = 0 : index} : memref<16x16x16x16xf16> -> !gpu.mma_matrix<16x16xf16, "COp">
 //       CHECK:   %[[F:.+]] = gpu.subgroup_mma_elementwise divf %[[D]], %[[E]] : (!gpu.mma_matrix<16x16xf16, "COp">, !gpu.mma_matrix<16x16xf16, "COp">) -> !gpu.mma_matrix<16x16xf16, "COp">
func.func @matmul_fused_broadcast(%arg0: memref<16x16xf16>, %arg1: memref<16x16xf16>,
   %E = vector.transfer_read %arg3[%c0, %c0, %c0, %c0], %cst
     {in_bounds = [true, true], permutation_map = affine_map<(d0, d1, d2, d3)->(0, d3)>}
     : memref<16x16x16x16xf16>, vector<16x16xf16>
   %F = arith.divf %D, %E : vector<16x16xf16>
   vector.transfer_write %F, %arg2[%c0, %c0] {in_bounds = [true, true]} : vector<16x16xf16>, memref<16x16xf16>
   return
 }
 
+// -----
+
+#map0 = affine_map<(d0, d1) -> (d1, d0)>
+#map1 = affine_map<(d0, d1, d2) -> (d0, d2)>
+#map2 = affine_map<(d0, d1, d2) -> (d1, d2)>
+#map3 = affine_map<(d0, d1, d2) -> (d0, d1)>
+#map4 = affine_map<(d0) -> (d0, 0)>
+#map5 = affine_map<(d0, d1) -> (d0, d1)>
+
 // CHECK-LABEL: func @matmul_3Dmemref
 //   CHECK-DAG:   %[[C0:.+]] = arith.constant 0 : index
 //   CHECK-DAG:   %[[A:.+]] = gpu.subgroup_mma_load_matrix %{{.*}}[%[[C0]], %[[C0]], %[[C0]]] {leadDimension = 16 : index} : memref<2x16x16xf16> -> !gpu.mma_matrix<16x16xf16, "AOp">
 //   CHECK-DAG:   %[[B:.+]] = gpu.subgroup_mma_load_matrix %{{.*}}[%[[C0]]] {leadDimension = 0 : index} : memref<16xf16> -> !gpu.mma_matrix<16x16xf16, "BOp">
 //   CHECK-DAG:   %[[C:.+]] = gpu.subgroup_mma_load_matrix %{{.*}}[%[[C0]], %[[C0]], %[[C0]]] {leadDimension = 16 : index} : memref<2x16x16xf16> -> !gpu.mma_matrix<16x16xf16, "COp">
 //       CHECK:   %[[D:.+]] = gpu.subgroup_mma_compute %[[A]], %[[B]], %[[C]] : !gpu.mma_matrix<16x16xf16, "AOp">, !gpu.mma_matrix<16x16xf16, "BOp"> -> !gpu.mma_matrix<16x16xf16, "COp">
 //       CHECK:   gpu.subgroup_mma_store_matrix %[[D]], %{{.*}}[%[[C0]], %[[C0]], %[[C0]]] {leadDimension = 16 : index} : !gpu.mma_matrix<16x16xf16, "COp">, memref<2x16x16xf16>
 func.func @matmul_3Dmemref(%arg0: memref<2x16x16xf16>, %arg1: memref<16xf16>, %arg2: memref<2x16x16xf16>) {
   %cst_0 = arith.constant dense<0.000000e+00> : vector<16x16xf16>
   %c0 = arith.constant 0 : index
   %cst = arith.constant 0.000000e+00 : f16
   %A = vector.transfer_read %arg0[%c0, %c0, %c0], %cst {in_bounds = [true, true]} : memref<2x16x16xf16>, vector<16x16xf16>
   %B = vector.transfer_read %arg1[%c0], %cst {permutation_map = #map4, in_bounds = [true, true]} : memref<16xf16>, vector<16x16xf16>
   %C = vector.transfer_read %arg2[%c0, %c0, %c0], %cst {in_bounds = [true, true]} : memref<2x16x16xf16>, vector<16x16xf16>
   %D = vector.contract {indexing_maps = [#map1, #map2, #map3], iterator_types = ["parallel", "parallel", "reduction"], kind = #vector.kind<add>} %A, %B, %C : vector<16x16xf16>, vector<16x16xf16> into vector<16x16xf16>
   vector.transfer_write %D, %arg2[%c0, %c0, %c0] {in_bounds = [true, true]} : vector<16x16xf16>, memref<2x16x16xf16>
   return
 }
 
+// -----
+
+#map0 = affine_map<(d0, d1) -> (d1, d0)>
+#map1 = affine_map<(d0, d1, d2) -> (d0, d2)>
+#map2 = affine_map<(d0, d1, d2) -> (d1, d2)>
+#map3 = affine_map<(d0, d1, d2) -> (d0, d1)>
+#map4 = affine_map<(d0) -> (d0, 0)>
+#map5 = affine_map<(d0, d1) -> (d0, d1)>
+
 // CHECK-LABEL: func @matmul_memref_strided
 //   CHECK-DAG:   %[[C0:.+]] = arith.constant 0 : index
 //   CHECK-DAG:   %[[A:.+]] = gpu.subgroup_mma_load_matrix %{{.*}}[%[[C0]], %[[C0]], %[[C0]]] {leadDimension = 32 : index} : memref<2x16x16xf16, #{{.*}}> -> !gpu.mma_matrix<16x16xf16, "AOp">
 //   CHECK-DAG:   %[[B:.+]] = gpu.subgroup_mma_load_matrix %{{.*}}[%[[C0]]] {leadDimension = 0 : index} : memref<16xf16> -> !gpu.mma_matrix<16x16xf16, "BOp">
 //   CHECK-DAG:   %[[C:.+]] = gpu.subgroup_mma_load_matrix %{{.*}}[%[[C0]], %[[C0]], %[[C0]]] {leadDimension = 16 : index} : memref<2x16x16xf16> -> !gpu.mma_matrix<16x16xf16, "COp">
 //       CHECK:   %[[D:.+]] = gpu.subgroup_mma_compute %[[A]], %[[B]], %[[C]] : !gpu.mma_matrix<16x16xf16, "AOp">, !gpu.mma_matrix<16x16xf16, "BOp"> -> !gpu.mma_matrix<16x16xf16, "COp">
 //       CHECK:   gpu.subgroup_mma_store_matrix %[[D]], %{{.*}}[%[[C0]], %[[C0]], %[[C0]]] {leadDimension = 16 : index} : !gpu.mma_matrix<16x16xf16, "COp">, memref<2x16x16xf16>
 func.func @matmul_memref_strided(%arg0: memref<2x16x16xf16, affine_map<(d0, d1, d2) -> (d0 * 512 + d1 * 32 + d2)>>, %arg1: memref<16xf16>, %arg2: memref<2x16x16xf16>) {
   %cst_0 = arith.constant dense<0.000000e+00> : vector<16x16xf16>
   %c0 = arith.constant 0 : index
   %cst = arith.constant 0.000000e+00 : f16
   %A = vector.transfer_read %arg0[%c0, %c0, %c0], %cst {in_bounds = [true, true]} : memref<2x16x16xf16, affine_map<(d0, d1, d2) -> (d0 * 512 + d1 * 32 + d2)>>, vector<16x16xf16>
   %B = vector.transfer_read %arg1[%c0], %cst {permutation_map = #map4, in_bounds = [true, true]} : memref<16xf16>, vector<16x16xf16>
   %C = vector.transfer_read %arg2[%c0, %c0, %c0], %cst {in_bounds = [true, true]} : memref<2x16x16xf16>, vector<16x16xf16>
   %D = vector.contract {indexing_maps = [#map1, #map2, #map3], iterator_types = ["parallel", "parallel", "reduction"], kind = #vector.kind<add>} %A, %B, %C : vector<16x16xf16>, vector<16x16xf16> into vector<16x16xf16>
   vector.transfer_write %D, %arg2[%c0, %c0, %c0] {in_bounds = [true, true]} : vector<16x16xf16>, memref<2x16x16xf16>
   return
 }
 
+// -----
+
+#map0 = affine_map<(d0, d1) -> (d1, d0)>
+#map1 = affine_map<(d0, d1, d2) -> (d0, d2)>
+#map2 = affine_map<(d0, d1, d2) -> (d1, d2)>
+#map3 = affine_map<(d0, d1, d2) -> (d0, d1)>
+#map4 = affine_map<(d0) -> (d0, 0)>
+#map5 = affine_map<(d0, d1) -> (d0, d1)>
+
 // CHECK-LABEL: func @matmul_transposed
 //   CHECK-DAG:   %[[A:.+]] = gpu.subgroup_mma_load_matrix %{{.*}}[%{{.*}}, %{{.*}}] {leadDimension = 16 : index} : memref<16x16xf16> -> !gpu.mma_matrix<16x16xf16, "AOp">
 //   CHECK-DAG:   %[[B:.+]] = gpu.subgroup_mma_load_matrix %{{.*}}[%{{.*}}, %{{.*}}] {leadDimension = 16 : index, transpose} : memref<16x16xf16> -> !gpu.mma_matrix<16x16xf16, "BOp">
 //   CHECK-DAG:   %[[C:.+]] = gpu.subgroup_mma_load_matrix %{{.*}}[%{{.*}}, %{{.*}}] {leadDimension = 16 : index} : memref<16x16xf16> -> !gpu.mma_matrix<16x16xf16, "COp">
 //       CHECK:   %[[D:.+]] = gpu.subgroup_mma_compute %[[A]], %[[B]], %[[C]] : !gpu.mma_matrix<16x16xf16, "AOp">, !gpu.mma_matrix<16x16xf16, "BOp"> -> !gpu.mma_matrix<16x16xf16, "COp">
 //       CHECK:   gpu.subgroup_mma_store_matrix %[[D]], %{{.*}}[%{{.*}}, %{{.*}}] {leadDimension = 16 : index} : !gpu.mma_matrix<16x16xf16, "COp">, memref<16x16xf16>
 func.func @matmul_transposed(%arg0: memref<16x16xf16>, %arg1: memref<16x16xf16>, %arg2: memref<16x16xf16>) {
   %cst_0 = arith.constant dense<0.000000e+00> : vector<16x16xf16>
   %c0 = arith.constant 0 : index
   %cst = arith.constant 0.000000e+00 : f16
   %A = vector.transfer_read %arg0[%c0, %c0], %cst {in_bounds = [true, true]} : memref<16x16xf16>, vector<16x16xf16>
   %B = vector.transfer_read %arg1[%c0, %c0], %cst {permutation_map = #map5, in_bounds = [true, true]} : memref<16x16xf16>, vector<16x16xf16>
   %C = vector.transfer_read %arg2[%c0, %c0], %cst {in_bounds = [true, true]} : memref<16x16xf16>, vector<16x16xf16>
   %D = vector.contract {indexing_maps = [#map1, #map2, #map3], iterator_types = ["parallel", "parallel", "reduction"], kind = #vector.kind<add>} %A, %B, %C : vector<16x16xf16>, vector<16x16xf16> into vector<16x16xf16>
   vector.transfer_write %D, %arg2[%c0, %c0] {in_bounds = [true, true]} : vector<16x16xf16>, memref<16x16xf16>
   return
 }
 
+// -----
+
+#map0 = affine_map<(d0, d1) -> (d1, d0)>
+#map1 = affine_map<(d0, d1, d2) -> (d0, d2)>
+#map2 = affine_map<(d0, d1, d2) -> (d1, d2)>
+#map3 = affine_map<(d0, d1, d2) -> (d0, d1)>
+#map4 = affine_map<(d0) -> (d0, 0)>
+#map5 = affine_map<(d0, d1) -> (d0, d1)>
+
 // CHECK-LABEL: func @matmul_transposed_broadcasted_1d
 //   CHECK-DAG:   %[[A:.+]] = gpu.subgroup_mma_load_matrix %{{.*}}[%{{.*}}] {leadDimension = 0 : index, transpose} : memref<16xf16> -> !gpu.mma_matrix<16x16xf16, "AOp">
 //   CHECK-DAG:   %[[B:.+]] = gpu.subgroup_mma_load_matrix %{{.*}}[%{{.*}}] {leadDimension = 0 : index} : memref<16xf16> -> !gpu.mma_matrix<16x16xf16, "BOp">
 //   CHECK-DAG:   %[[C:.+]] = gpu.subgroup_mma_load_matrix %{{.*}}[%{{.*}}, %{{.*}}] {leadDimension = 16 : index} : memref<16x16xf16> -> !gpu.mma_matrix<16x16xf16, "COp">
 //       CHECK:   %[[D:.+]] = gpu.subgroup_mma_compute %[[A]], %[[B]], %[[C]] : !gpu.mma_matrix<16x16xf16, "AOp">, !gpu.mma_matrix<16x16xf16, "BOp"> -> !gpu.mma_matrix<16x16xf16, "COp">
 //       CHECK:   gpu.subgroup_mma_store_matrix %[[D]], %{{.*}}[%{{.*}}, %{{.*}}] {leadDimension = 16 : index} : !gpu.mma_matrix<16x16xf16, "COp">, memref<16x16xf16>
 func.func @matmul_transposed_broadcasted_1d(%arg0: memref<16xf16>, %arg1: memref<16xf16>, %arg2: memref<16x16xf16>) {
   %cst_0 = arith.constant dense<0.000000e+00> : vector<16x16xf16>
   %c0 = arith.constant 0 : index
   %cst = arith.constant 0.000000e+00 : f16
   %A = vector.transfer_read %arg0[%c0], %cst {in_bounds = [true, true], permutation_map = affine_map<(d0) -> (d0, 0)>} : memref<16xf16>, vector<16x16xf16>
   %B = vector.transfer_read %arg1[%c0], %cst {in_bounds = [true, true], permutation_map = affine_map<(d0) -> (d0, 0)>} : memref<16xf16>, vector<16x16xf16>
   %C = vector.transfer_read %arg2[%c0, %c0], %cst {in_bounds = [true, true]} : memref<16x16xf16>, vector<16x16xf16>
   %D = vector.contract {indexing_maps = [#map1, #map2, #map3], iterator_types = ["parallel", "parallel", "reduction"], kind = #vector.kind<add>} %A, %B, %C : vector<16x16xf16>, vector<16x16xf16> into vector<16x16xf16>
   vector.transfer_write %D, %arg2[%c0, %c0] {in_bounds = [true, true]} : vector<16x16xf16>, memref<16x16xf16>
   return
 }
 
+// -----
+
+#map0 = affine_map<(d0, d1) -> (d1, d0)>
+#map1 = affine_map<(d0, d1, d2) -> (d0, d2)>
+#map2 = affine_map<(d0, d1, d2) -> (d1, d2)>
+#map3 = affine_map<(d0, d1, d2) -> (d0, d1)>
+#map4 = affine_map<(d0) -> (d0, 0)>
+#map5 = affine_map<(d0, d1) -> (d0, d1)>
+
 // CHECK-LABEL: func @matmul_transposed_broadcasted_2d
 //   CHECK-DAG:   %[[A:.+]] = gpu.subgroup_mma_load_matrix %{{.*}}[%{{.*}}] {leadDimension = 0 : index, transpose} : memref<32x32xf16> -> !gpu.mma_matrix<16x16xf16, "AOp">
 //   CHECK-DAG:   %[[B:.+]] = gpu.subgroup_mma_load_matrix %{{.*}}[%{{.*}}] {leadDimension = 0 : index} : memref<32x32xf16> -> !gpu.mma_matrix<16x16xf16, "BOp">
 //   CHECK-DAG:   %[[C:.+]] = gpu.subgroup_mma_load_matrix %{{.*}}[%{{.*}}, %{{.*}}] {leadDimension = 16 : index} : memref<16x16xf16> -> !gpu.mma_matrix<16x16xf16, "COp">
 //       CHECK:   %[[D:.+]] = gpu.subgroup_mma_compute %[[A]], %[[B]], %[[C]] : !gpu.mma_matrix<16x16xf16, "AOp">, !gpu.mma_matrix<16x16xf16, "BOp"> -> !gpu.mma_matrix<16x16xf16, "COp">
 //       CHECK:   gpu.subgroup_mma_store_matrix %[[D]], %{{.*}}[%{{.*}}, %{{.*}}] {leadDimension = 16 : index} : !gpu.mma_matrix<16x16xf16, "COp">, memref<16x16xf16>
 func.func @matmul_transposed_broadcasted_2d(%arg0: memref<32x32xf16>, %arg1: memref<32x32xf16>, %arg2: memref<16x16xf16>) {
   %cst_0 = arith.constant dense<0.000000e+00> : vector<16x16xf16>
   %c0 = arith.constant 0 : index
   %cst = arith.constant 0.000000e+00 : f16
   %A = vector.transfer_read %arg0[%c0, %c0], %cst {in_bounds = [true, true], permutation_map = affine_map<(d0, d1) -> (d1, 0)>} : memref<32x32xf16>, vector<16x16xf16>
   %B = vector.transfer_read %arg1[%c0, %c0], %cst {in_bounds = [true, true], permutation_map = affine_map<(d0, d1) -> (d1, 0)>} : memref<32x32xf16>, vector<16x16xf16>
   %C = vector.transfer_read %arg2[%c0, %c0], %cst {in_bounds = [true, true]} : memref<16x16xf16>, vector<16x16xf16>
   %D = vector.contract {indexing_maps = [#map1, #map2, #map3], iterator_types = ["parallel", "parallel", "reduction"], kind = #vector.kind<add>} %A, %B, %C : vector<16x16xf16>, vector<16x16xf16> into vector<16x16xf16>
   vector.transfer_write %D, %arg2[%c0, %c0] {in_bounds = [true, true]} : vector<16x16xf16>, memref<16x16xf16>
   return
 }
 
+// -----
+
+#map0 = affine_map<(d0, d1) -> (d1, d0)>
+#map1 = affine_map<(d0, d1, d2) -> (d0, d2)>
+#map2 = affine_map<(d0, d1, d2) -> (d1, d2)>
+#map3 = affine_map<(d0, d1, d2) -> (d0, d1)>
+#map4 = affine_map<(d0) -> (d0, 0)>
+#map5 = affine_map<(d0, d1) -> (d0, d1)>
+
+//   CHECK-DAG: #[[$map:.*]] = affine_map<(d0, d1, d2) -> (d0, d2)>
+//   CHECK-DAG: #[[$map1:.*]] = affine_map<(d0, d1, d2) -> (d2, d1)>
+//   CHECK-DAG: #[[$map2:.*]] = affine_map<(d0, d1, d2) -> (d0, d1)>
+
 // Do not convert to subgroup_mma ops with integer types if signedness cannot be inferred.
 // CHECK-LABEL: func @matmul_no_extend_int8
 //   CHECK-DAG:   %[[A:.+]] = vector.transfer_read %{{.*}}[%{{.*}}, %{{.*}}], %{{.*}} {in_bounds = [true, true]} : memref<16x16xi8>, vector<16x16xi8>
 //   CHECK-DAG:   %[[B:.+]] = vector.transfer_read %{{.*}}[%{{.*}}, %{{.*}}], %{{.*}} {in_bounds = [true, true]} : memref<16x16xi8>, vector<16x16xi8>
 //   CHECK-DAG:   %[[C:.+]] = vector.transfer_read %{{.*}}[%{{.*}}, %{{.*}}], %{{.*}} {in_bounds = [true, true]} : memref<16x16xi32>, vector<16x16xi32>
-//       CHECK:   %[[D:.+]] = vector.contract {indexing_maps = [#map1, #map2, #map3], iterator_types = ["parallel", "parallel", "reduction"], kind = #vector.kind<add>} %[[A]], %[[B]], %[[C]] : vector<16x16xi8>, vector<16x16xi8> into vector<16x16xi32>
+//       CHECK:   %[[D:.+]] = vector.contract {indexing_maps = [#[[$map]], #[[$map1]], #[[$map2]]], iterator_types = ["parallel", "parallel", "reduction"], kind = #vector.kind<add>} %[[A]], %[[B]], %[[C]] : vector<16x16xi8>, vector<16x16xi8> into vector<16x16xi32>
 //       CHECK:   vector.transfer_write %{{.*}}, %{{.*}}[%{{.*}}, %{{.*}}] {in_bounds = [true, true]} : vector<16x16xi32>, memref<16x16xi32>
 func.func @matmul_no_extend_int8(%arg0: memref<16x16xi8>, %arg1: memref<16x16xi8>, %arg2: memref<16x16xi32>) {
   %cst_0 = arith.constant dense<0> : vector<16x16xi8>
   %c0 = arith.constant 0 : index
   %cst_i8 = arith.constant 0 : i8
   %cst_i32 = arith.constant 0 : i32
   %A = vector.transfer_read %arg0[%c0, %c0], %cst_i8 {in_bounds = [true, true]} : memref<16x16xi8>, vector<16x16xi8>
   %B = vector.transfer_read %arg1[%c0, %c0], %cst_i8 {permutation_map = #map0, in_bounds = [true, true]} : memref<16x16xi8>, vector<16x16xi8>
   %C = vector.transfer_read %arg2[%c0, %c0], %cst_i32 {in_bounds = [true, true]} : memref<16x16xi32>, vector<16x16xi32>
   %D = vector.contract {indexing_maps = [#map1, #map2, #map3], iterator_types = ["parallel", "parallel", "reduction"], kind = #vector.kind<add>} %A, %B, %C : vector<16x16xi8>, vector<16x16xi8> into vector<16x16xi32>
   vector.transfer_write %D, %arg2[%c0, %c0] {in_bounds = [true, true]} : vector<16x16xi32>, memref<16x16xi32>
   return
 }
 
+// -----
+
+#map0 = affine_map<(d0, d1) -> (d1, d0)>
+#map1 = affine_map<(d0, d1, d2) -> (d0, d2)>
+#map2 = affine_map<(d0, d1, d2) -> (d1, d2)>
+#map3 = affine_map<(d0, d1, d2) -> (d0, d1)>
+#map4 = affine_map<(d0) -> (d0, 0)>
+#map5 = affine_map<(d0, d1) -> (d0, d1)>
+
 // CHECK-LABEL: func @matmul_int8
 //   CHECK-DAG:   %[[A:.+]] = gpu.subgroup_mma_load_matrix %{{.*}}[%{{.*}}, %{{.*}}] {leadDimension = 16 : index} : memref<16x16xi8> -> !gpu.mma_matrix<16x16xsi8, "AOp">
 //   CHECK-DAG:   %[[B:.+]] = gpu.subgroup_mma_load_matrix %{{.*}}[%{{.*}}, %{{.*}}] {leadDimension = 16 : index} : memref<16x16xi8> -> !gpu.mma_matrix<16x16xsi8, "BOp">
 //   CHECK-DAG:   %[[C:.+]] = gpu.subgroup_mma_load_matrix %{{.*}}[%{{.*}}, %{{.*}}] {leadDimension = 16 : index} : memref<16x16xi32> -> !gpu.mma_matrix<16x16xi32, "COp">
 //       CHECK:   %[[D:.+]] = gpu.subgroup_mma_compute %[[A]], %[[B]], %[[C]] : !gpu.mma_matrix<16x16xsi8, "AOp">, !gpu.mma_matrix<16x16xsi8, "BOp"> -> !gpu.mma_matrix<16x16xi32, "COp">
 //       CHECK:   gpu.subgroup_mma_store_matrix %[[D]], %{{.*}}[%{{.*}}, %{{.*}}] {leadDimension = 16 : index} : !gpu.mma_matrix<16x16xi32, "COp">, memref<16x16xi32>
 func.func @matmul_int8(%arg0: memref<16x16xi8>, %arg1: memref<16x16xi8>, %arg2: memref<16x16xi32>) {
   %cst_0 = arith.constant dense<0> : vector<16x16xi8>
   %c0 = arith.constant 0 : index
   %cst_i8 = arith.constant 0 : i8
   %cst_i32 = arith.constant 0 : i32
   %Ar = vector.transfer_read %arg0[%c0, %c0], %cst_i8 {in_bounds = [true, true]} : memref<16x16xi8>, vector<16x16xi8>
   %Br = vector.transfer_read %arg1[%c0, %c0], %cst_i8 {permutation_map = #map0, in_bounds = [true, true]} : memref<16x16xi8>, vector<16x16xi8>
   %C = vector.transfer_read %arg2[%c0, %c0], %cst_i32 {in_bounds = [true, true]} : memref<16x16xi32>, vector<16x16xi32>
   %Ae = arith.extsi %Ar : vector<16x16xi8> to vector<16x16xi32>
   %Be = arith.extsi %Br : vector<16x16xi8> to vector<16x16xi32>
   %D = vector.contract {indexing_maps = [#map1, #map2, #map3], iterator_types = ["parallel", "parallel", "reduction"], kind = #vector.kind<add>} %Ae, %Be, %C : vector<16x16xi32>, vector<16x16xi32> into vector<16x16xi32>
   vector.transfer_write %D, %arg2[%c0, %c0] {in_bounds = [true, true]} : vector<16x16xi32>, memref<16x16xi32>
   return
 }
 
+// -----
+
+#map0 = affine_map<(d0, d1) -> (d1, d0)>
+#map1 = affine_map<(d0, d1, d2) -> (d0, d2)>
+#map2 = affine_map<(d0, d1, d2) -> (d1, d2)>
+#map3 = affine_map<(d0, d1, d2) -> (d0, d1)>
+#map4 = affine_map<(d0) -> (d0, 0)>
+#map5 = affine_map<(d0, d1) -> (d0, d1)>
+
 // CHECK-LABEL: func @matmul_mixed_signedness_int8
 //   CHECK-DAG:   %[[A:.+]] = gpu.subgroup_mma_load_matrix %{{.*}}[%{{.*}}, %{{.*}}] {leadDimension = 16 : index} : memref<16x16xi8> -> !gpu.mma_matrix<16x16xui8, "AOp">
 //   CHECK-DAG:   %[[B:.+]] = gpu.subgroup_mma_load_matrix %{{.*}}[%{{.*}}, %{{.*}}] {leadDimension = 16 : index} : memref<16x16xi8> -> !gpu.mma_matrix<16x16xsi8, "BOp">
 //   CHECK-DAG:   %[[C:.+]] = gpu.subgroup_mma_load_matrix %{{.*}}[%{{.*}}, %{{.*}}] {leadDimension = 16 : index} : memref<16x16xi32> -> !gpu.mma_matrix<16x16xi32, "COp">
 //       CHECK:   %[[D:.+]] = gpu.subgroup_mma_compute %[[A]], %[[B]], %[[C]] : !gpu.mma_matrix<16x16xui8, "AOp">, !gpu.mma_matrix<16x16xsi8, "BOp"> -> !gpu.mma_matrix<16x16xi32, "COp">
 //       CHECK:   gpu.subgroup_mma_store_matrix %[[D]], %{{.*}}[%{{.*}}, %{{.*}}] {leadDimension = 16 : index} : !gpu.mma_matrix<16x16xi32, "COp">, memref<16x16xi32>
 func.func @matmul_mixed_signedness_int8(%arg0: memref<16x16xi8>, %arg1: memref<16x16xi8>, %arg2: memref<16x16xi32>) {
   %cst_0 = arith.constant dense<0> : vector<16x16xi8>