[WIP] Add support for MMA conversion for 1-D vector.transfer followed by a broadcast to 2-D

mlir/lib/Conversion/VectorToGPU/VectorToGPU.cpp

if (!useNvGpu) {
     bool result = map.isMinorIdentity() || map == broadcastInnerDim ||
                   isTransposeMatrixLoadMap(b, map);
     return result;
   }
 
   return true;
 }
 
+// Return true if the transfer op can be converted to a MMA matrix load.
+static bool transferReadFollowedByBroadcastSupportsMMAMatrixType(
+    vector::TransferReadOp readOp, bool useNvGpu) {
+  bool res = true;
+  if (readOp.getMask() || readOp.hasOutOfBoundsDim() ||
+      readOp.getVectorType().getRank() != 1)
+    res = false;
+  if (!getMemrefConstantHorizontalStride(readOp.getShapedType()))
+    res = false;
+  AffineMap map = readOp.getPermutationMap();
+  OpBuilder b(readOp.getContext());
+
+  if (res && !useNvGpu)
+    return map.isMinorIdentity() || isTransposeMatrixLoadMap(b, map);
+
+  llvm::errs() << "RES transferReadFollowedByBroadcastSupportsMMAMatrixType: "
+               << res << "\n";
+  return res;
+}
+
 // Return true if the transfer op can be converted to a MMA matrix store.
 static bool
 transferWriteSupportsMMAMatrixType(vector::TransferWriteOp writeOp) {
   // TODO: support 0-d corner case.
   if (writeOp.getTransferRank() == 0)
     return false;
 
   if (writeOp.getMask() || writeOp.hasOutOfBoundsDim() ||
static bool constantSupportsMMAMatrixType(arith::ConstantOp constantOp) {
   auto vecType = constantOp.getType().dyn_cast<VectorType>();
   if (!vecType || vecType.getRank() != 2)
     return false;
   return constantOp.getValue().isa<SplatElementsAttr>();
 }
 
 /// Return true if this is a broadcast from scalar to a 2D vector.
 static bool broadcastSupportsMMAMatrixType(vector::BroadcastOp broadcastOp) {
-  return broadcastOp.getVectorType().getRank() == 2 &&
+  auto res = broadcastOp.getVectorType().getRank() == 2 &&
              broadcastOp.getSource().getType().isa<FloatType>();
+  llvm::errs() << "RES broadcastSupportsMMAMatrixType: " << res << "\n";
+  return res;
+}
+
+/// Return true if this is a broadcast from 1-D to a 2-D vector and the 1-D
+/// vector comes from a TransferReadOp.
+static bool
+broadcastFromTransferReadSupportsMMAMatrixType(vector::BroadcastOp broadcastOp,
+                                               bool useNvGpu) {
+  auto readOp = broadcastOp.getSource().getDefiningOp<vector::TransferReadOp>();
+  auto sourceVectorType =
+      broadcastOp.getSource().getType().dyn_cast<VectorType>();
+  auto res =
+      !broadcastSupportsMMAMatrixType(broadcastOp) && sourceVectorType &&
+      sourceVectorType.getRank() == 1 &&
+      transferReadFollowedByBroadcastSupportsMMAMatrixType(readOp, useNvGpu);
+  llvm::errs() << "RES broadcastFromTransferReadSupportsMMAMatrixType: " << res
+               << "\n";
+  return res;
 }
 
 /// Return the MMA elementwise enum associated with `op` if it is supported.
 /// Return `llvm::None` otherwise.
 static std::optional<gpu::MMAElementwiseOp>
 convertElementwiseOpToMMA(Operation *op) {
   if (isa<arith::AddFOp>(op))
     return gpu::MMAElementwiseOp::ADDF;
FailureOr<nvgpu::WarpMatrixInfo> warpMatrixInfo =
       nvgpu::getWarpMatrixInfo(op);
   if (failed(warpMatrixInfo))
     return false;
 
   FailureOr<vector::ContractionOp> contractOp = nvgpu::getUserContract(op);
   if (failed(contractOp))
     return false;
 
-  // Handle vector.extract_strided_slice on registers containing
-  // matrixB and matrixC operands. vector.extract_strided_slice op
-  // is not supported on registers containing matrixA operands.
+  // Handle vector.extract_strided_slice on registers
+  // containing matrixB and matrixC operands.
+  // vector.extract_strided_slice op is not supported on
+  // registers containing matrixA operands.
   if (warpMatrixInfo->operandRole == nvgpu::MatMulOperandRole::B)
     return (op->getResult(0).getType().cast<VectorType>() ==
             (*contractOp).getRhs().getType().cast<VectorType>());
   if (warpMatrixInfo->operandRole == nvgpu::MatMulOperandRole::C)
     return (op->getResult(0).getType().cast<VectorType>() ==
             (*contractOp).getAcc().getType().cast<VectorType>());
 
   return false;
 }
 
 static bool supportsMMaMatrixType(Operation *op, bool useNvGpu) {
   if (isa<scf::ForOp, scf::YieldOp>(op))
     return true;
   if (auto transferRead = dyn_cast<vector::TransferReadOp>(op))
-    return transferReadSupportsMMAMatrixType(transferRead, useNvGpu);
+    return transferReadSupportsMMAMatrixType(transferRead, useNvGpu) ||
+           transferReadFollowedByBroadcastSupportsMMAMatrixType(transferRead,
+                                                                useNvGpu);
   if (auto transferWrite = dyn_cast<vector::TransferWriteOp>(op))
     return transferWriteSupportsMMAMatrixType(transferWrite);
   if (auto extractStridedSlice = dyn_cast<vector::ExtractStridedSliceOp>(op))
     return useNvGpu &&
            extractStridedSliceSupportsMMAMatrixType(extractStridedSlice);
   if (auto contract = dyn_cast<vector::ContractionOp>(op))
     return contractSupportsMMAMatrixType(contract, useNvGpu);
   if (auto constant = dyn_cast<arith::ConstantOp>(op))
     return constantSupportsMMAMatrixType(constant);
-  if (auto broadcast = dyn_cast<vector::BroadcastOp>(op))
-    return broadcastSupportsMMAMatrixType(broadcast);
+  if (auto broadcast = dyn_cast<vector::BroadcastOp>(op)) {
+    return broadcastSupportsMMAMatrixType(broadcast) ||
+           broadcastFromTransferReadSupportsMMAMatrixType(broadcast, useNvGpu);
+  }
   return elementwiseSupportsMMAMatrixType(op);
 }
 
 /// Return an unsorted slice handling scf.for region differently than
 /// `getSlice`. In scf.for we only want to include as part of the slice elements
 /// that are part of the use/def chain.
 static SetVector<Operation *> getSliceContract(Operation *op,
                                                TransitiveFilter backwardFilter,
                                                TransitiveFilter forwardFilter) {
   SetVector<Operation *> slice;
   slice.insert(op);
   unsigned currentIndex = 0;
   SetVector<Operation *> backwardSlice;
   SetVector<Operation *> forwardSlice;
   while (currentIndex != slice.size()) {
     auto *currentOp = (slice)[currentIndex];
-    // Compute and insert the backwardSlice starting from currentOp.
+    // Compute and insert the backwardSlice starting from
+    // currentOp.
     backwardSlice.clear();
     getBackwardSlice(currentOp, &backwardSlice, backwardFilter);
     slice.insert(backwardSlice.begin(), backwardSlice.end());
 
-    // Compute and insert the forwardSlice starting from currentOp.
+    // Compute and insert the forwardSlice starting from
+    // currentOp.
     forwardSlice.clear();
-    // Special case for ForOp, we don't want to include the whole region but
-    // only the value using the region arguments.
-    // TODO: We should refine this to only care about the region arguments being
-    // converted to matrix type.
+    // Special case for ForOp, we don't want to include the
+    // whole region but only the value using the region
+    // arguments.
+    // TODO: We should refine this to only care about the
+    // region arguments being converted to matrix type.
     if (auto forOp = dyn_cast<scf::ForOp>(currentOp)) {
       for (Value forOpResult : forOp.getResults())
         getForwardSlice(forOpResult, &forwardSlice, forwardFilter);
       for (BlockArgument &arg : forOp.getRegionIterArgs())
         getForwardSlice(arg, &forwardSlice, forwardFilter);
     } else {
       getForwardSlice(currentOp, &forwardSlice, forwardFilter);
     }
return llvm::any_of(op->getOperandTypes(),
                         [](Type t) { return t.isa<VectorType>(); });
   };
   SetVector<Operation *> opToConvert;
   op->walk([&](vector::ContractionOp contract) {
     if (opToConvert.contains(contract.getOperation()))
       return;
     SetVector<Operation *> dependentOps =
         getSliceContract(contract, hasVectorDest, hasVectorSrc);
-    // If any instruction cannot use MMA matrix type drop the whole
-    // chain. MMA matrix are stored in an opaque type so they cannot be used
-    // by all operations.
+    // If any instruction cannot use MMA matrix type drop the
+    // whole chain. MMA matrix are stored in an opaque type so
+    // they cannot be used by all operations.
     if (llvm::any_of(dependentOps, [useNvGpu](Operation *op) {
-          return !supportsMMaMatrixType(op, useNvGpu);
+          auto res = !supportsMMaMatrixType(op, useNvGpu);
+          if (res)
+            llvm::errs() << "DOES NOT SUPPORT: " << *op << "\n";
+          return res;
         }))
       return;
     opToConvert.insert(dependentOps.begin(), dependentOps.end());
   });
-  // Sort the operations so that we can convert them in topological order.
+  // Sort the operations so that we can convert them in
+  // topological order.
   return topologicalSort(opToConvert);
 }
 
 namespace {
 // Transform contract into (m, k)x(k, n)x(m, n) form so that it can be converted
 // to MMA matmul.
 struct PrepareContractToGPUMMA
     : public OpRewritePattern<vector::ContractionOp> {
if (contract.getRhs() == op.getResult())
       return "BOp";
   }
   return "COp";
 }
 
 static void convertTransferReadOp(vector::TransferReadOp op,
                                   llvm::DenseMap<Value, Value> &valueMapping) {
   assert(op.getTransferRank() > 0 && "unexpected 0-d transfer");
-  assert(transferReadSupportsMMAMatrixType(op, /*useNvGpu=*/false));
+  if (!transferReadSupportsMMAMatrixType(op,
+                                         /*useNvGpu=*/false))
+    return;
+  // Only transfers that return 2-D vectors are supported.
+  if (op.getVectorType().getRank() != 2)
+    return;
   std::optional<int64_t> stride =
       getMemrefConstantHorizontalStride(op.getShapedType());
   AffineMap map = op.getPermutationMap();
   // Handle broadcast by setting the stride to 0.
   if (map.getResult(0).isa<AffineConstantExpr>()) {
     assert(map.getResult(0).cast<AffineConstantExpr>().getValue() == 0);
     stride = 0;
   }
FailureOr<nvgpu::FragmentElementInfo> regInfo =
       nvgpu::getMmaSyncRegisterType(*warpMatrixInfo);
   if (failed(regInfo))
     return failure();
 
   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 "
+    return op->emitError() << "failed to convert vector.transfer_read to "
+                              "ldmatrix; this op "
                               "likely "
-              "should not be converted to a nvgpu.ldmatrix call.";
+                              "should not be converted to a nvgpu.ldmatrix "
+                              "call.";
   }
 
   // Adjust the load offset.
   auto laneId = builder.create<gpu::LaneIdOp>(loc);
   FailureOr<AffineMap> offsets =
       nvgpu::getLaneIdToLdMatrixMatrixCoord(loc, builder, *params);
   if (failed(offsets))
     return failure();
createNonLdMatrixLoads(vector::TransferReadOp op, OpBuilder &builder,
   FailureOr<nvgpu::WarpMatrixInfo> warpMatrixInfo =
       nvgpu::getWarpMatrixInfo(op);
   if (failed(warpMatrixInfo))
     return failure();
   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";
+                       "conversion to distributed non-ldmatrix compatible "
+                       "load";
     return failure();
   }
 
   Value laneId = builder.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>(
       op.getLoc(), vectorType.getElementType(),
       builder.getZeroAttr(vectorType.getElementType()));
   Value result = builder.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 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);
convertTransferReadToLoads(vector::TransferReadOp op,
 
   bool isLdMatrixCompatible =
       op.getSource().getType().cast<MemRefType>().getMemorySpaceAsInt() == 3 &&
       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.
+  // 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);
 
convertExtractStridedSlice(vector::ExtractStridedSliceOp op,
   if (failed(warpMatrixInfo))
     return failure();
 
   FailureOr<nvgpu::FragmentElementInfo> mmaSyncFragmentInfo =
       nvgpu::getMmaSyncRegisterType(*warpMatrixInfo);
   if (failed(mmaSyncFragmentInfo))
     return failure();
 
-  // Find the vector.transer_read whose result vector is being sliced.
+  // Find the vector.transer_read whose result vector is being
+  // sliced.
   auto transferReadOp = op.getVector().getDefiningOp<vector::TransferReadOp>();
   if (!transferReadOp)
     return failure();
 
   warpMatrixInfo = nvgpu::getWarpMatrixInfo(transferReadOp);
   if (failed(warpMatrixInfo))
     return failure();
 
   FailureOr<nvgpu::FragmentElementInfo> ldFragmentInfo =
       nvgpu::getMmaSyncRegisterType(*warpMatrixInfo);
   if (failed(ldFragmentInfo))
     return failure();
 
-  assert(
-      (mmaSyncFragmentInfo->elementsPerRegister ==
+  assert((mmaSyncFragmentInfo->elementsPerRegister ==
           ldFragmentInfo->elementsPerRegister) &&
-      "Number of elements per register should be same for load and mma.sync");
+         "Number of elements per register should be same for "
+         "load and mma.sync");
 
-  // Create vector.extract_strided_slice op for thread-owned fragments.
+  // 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;
 
   // 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].
+  // 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]);
static void convertConstantOp(arith::ConstantOp op,
   auto matrix = b.create<gpu::SubgroupMmaConstantMatrixOp>(op.getLoc(), type,
                                                            scalarConstant);
   valueMapping[op.getResult()] = matrix;
 }
 
 /// Convert a vector.broadcast from scalar to a SubgroupMmaConstantMatrix op.
 static void convertBroadcastOp(vector::BroadcastOp op,
                                llvm::DenseMap<Value, Value> &valueMapping) {
-  assert(broadcastSupportsMMAMatrixType(op));
+  // This op only catches the broadcasts that can directly
+  // convert to an MMA op.
+  if (!broadcastSupportsMMAMatrixType(op))
+    return;
   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());
   valueMapping[op.getResult()] = matrix;
 }
 
+/// Convert a vector.broadcast from scalar to a SubgroupMmaConstantMatrix op.
+static void
+convertBroadcastFromTransferReadOp(vector::BroadcastOp broadcastOp,
+                                   llvm::DenseMap<Value, Value> &valueMapping) {
+  // This op catches the broadcasts that cannot directly convert to an MMA
+  // op.
+  if (broadcastSupportsMMAMatrixType(broadcastOp))
+    return;
+  if (!broadcastFromTransferReadSupportsMMAMatrixType(broadcastOp,
+                                                      /*useNvGpu=*/false))
+    return;
+  auto readOp = broadcastOp.getSource().getDefiningOp<vector::TransferReadOp>();
+  assert(readOp && readOp.getVectorType().getRank() == 1);
+  // Handle broadcast by setting the stride to 0, unconditionally.
+  int64_t stride = 0;
+  const char *fragType = inferFragType(readOp);
+  gpu::MMAMatrixType type = gpu::MMAMatrixType::get(
+      broadcastOp.getVectorType().getShape(),
+      broadcastOp.getVectorType().getElementType(), fragType);
+  OpBuilder b(readOp);
+  bool isTranspose = false;
+  Value load = b.create<gpu::SubgroupMmaLoadMatrixOp>(
+      readOp.getLoc(), type, readOp.getSource(), readOp.getIndices(),
+      b.getIndexAttr(stride), isTranspose ? b.getUnitAttr() : UnitAttr());
+  valueMapping[broadcastOp.getResult()] = load;
+}
+
 // 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,
                                                ValueRange newIterOperands) {
-  // Create a new loop before the existing one, with the extra operands.
+  // Create a new loop before the existing one, with the extra
+  // operands.
   OpBuilder::InsertionGuard g(b);
   b.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);
   newLoop.getBody()->erase();
static void convertYieldOp(scf::YieldOp op,
                            llvm::DenseMap<Value, Value> &valueMapping) {
   OpBuilder b(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.
+    // 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();
 }
 
 /// Convert an elementwise op to the equivalent elementwise op on MMA matrix.
for (Operation *op : ops) {
     } else if (auto transferWrite = dyn_cast<vector::TransferWriteOp>(op)) {
       convertTransferWriteOp(transferWrite, valueMapping);
     } else if (auto contractOp = dyn_cast<vector::ContractionOp>(op)) {
       convertContractOp(contractOp, valueMapping);
     } else if (auto constantOp = dyn_cast<arith::ConstantOp>(op)) {
       convertConstantOp(constantOp, valueMapping);
     } else if (auto broadcastOp = dyn_cast<vector::BroadcastOp>(op)) {
       convertBroadcastOp(broadcastOp, valueMapping);
+      convertBroadcastFromTransferReadOp(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);
     } else if (auto elementwiseType = convertElementwiseOpToMMA(op)) {
       convertElementwiseOp(op, *elementwiseType, valueMapping);
     }
   }
struct ConvertVectorToGPUPass
 
   void runOnOperation() override {
     RewritePatternSet patterns(&getContext());
     populatePrepareVectorToMMAPatterns(patterns, useNvGpu.getValue());
     if (failed(
             applyPatternsAndFoldGreedily(getOperation(), std::move(patterns))))
       return signalPassFailure();
 
+    getOperation()->dump();
+
     if (useNvGpu.getValue()) {
       if (failed(convertVectorToNVVMCompatibleMMASync(getOperation())))
         return signalPassFailure();
     }
 
     (void)convertVectorToMMAOps(getOperation());
   }
 };
 
 } // namespace
 
 std::unique_ptr<Pass> mlir::createConvertVectorToGPUPass(bool useNvGpu) {
   return std::make_unique<ConvertVectorToGPUPass>(useNvGpu);
 }

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
 
 #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
 //   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_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>

 //   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">
 //       CHECK:   gpu.subgroup_mma_store_matrix %[[F]], %{{.*}}[%{{.*}}, %{{.*}}] {leadDimension = 16 : index} : !gpu.mma_matrix<16x16xf16, "COp">, memref<16x16xf16>
+// func.func @matmul_fused_broadcast(%arg0: memref<16x16xf16>, %arg1: memref<16x16xf16>,
+//   %arg2: memref<16x16xf16>, %arg3: memref<16x16x16x16xf16>) {
+//   %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>
+//   %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
+// }
 func.func @matmul_fused_broadcast(%arg0: memref<16x16xf16>, %arg1: memref<16x16xf16>,
   %arg2: memref<16x16xf16>, %arg3: memref<16x16x16x16xf16>) {
   %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>
-  %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>
+  %Eread = vector.transfer_read %arg3[%c0, %c0, %c0, %c0], %cst
+    {in_bounds = [true], permutation_map = affine_map<(d0, d1, d2, d3)->(d3)>}
+    : memref<16x16x16x16xf16>, vector<16xf16>
+  %E = vector.broadcast %Eread: vector<16xf16> to 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
 }
 
 // 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 = affine_map<(d0) -> (d0, 0)>, 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
+// }
 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>
+  %Bread = vector.transfer_read %arg1[%c0], %cst {permutation_map = affine_map<(d0) -> (d0)>, in_bounds = [true]} : memref<16xf16>, vector<16xf16>
+  %B = vector.broadcast %Bread: vector<16xf16> to 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
 }
 
 // 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 = affine_map<(d0) -> (d0, 0)>, 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
+// }
 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>
+  %Bread = vector.transfer_read %arg1[%c0], %cst {permutation_map = affine_map<(d0) -> (d0)>, in_bounds = [true]} : memref<16xf16>, vector<16xf16>
+  %B = vector.broadcast %Bread: vector<16xf16> to 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
 }
 
 // CHECK-LABEL: func @matmul_transposed
 //   CHECK-DAG:   %[[A:.+]] = gpu.subgroup_mma_load_matrix %{{.*}}[%{{.*}}, %{{.*}}] {leadDimension = 16 : index} : memref<16x16xf16> -> !gpu.mma_matrix<16x16xf16, "AOp">