[mlir][VectorToGPU] Add conversion for scf::For op with Matrix operands

[ThomasRaoux]

ping :)

mlir/lib/Conversion/VectorToGPU/VectorToGPU.cpp

 #include <type_traits>
 
 #include "mlir/Conversion/VectorToGPU/VectorToGPU.h"
 
 #include "../PassDetail.h"
 #include "mlir/Analysis/SliceAnalysis.h"
 #include "mlir/Dialect/GPU/GPUDialect.h"
 #include "mlir/Dialect/MemRef/IR/MemRef.h"
+#include "mlir/Dialect/SCF/SCF.h"
 #include "mlir/Dialect/Utils/StructuredOpsUtils.h"
 #include "mlir/Dialect/Vector/VectorOps.h"
 #include "mlir/Dialect/Vector/VectorUtils.h"
 #include "mlir/IR/Builders.h"
 #include "mlir/Pass/Pass.h"
 #include "mlir/Transforms/GreedyPatternRewriteDriver.h"
 #include "mlir/Transforms/Passes.h"
 

 static bool constantSupportsMMAMatrixType(ConstantOp constantOp) {
   auto vecType = constantOp.getType().dyn_cast<VectorType>();
   if (!vecType || vecType.getRank() != 2)
     return false;
   return constantOp.value().isa<SplatElementsAttr>();
 }
 
 static bool supportsMMaMatrixType(Operation *op) {
+  if (isa<scf::ForOp, scf::YieldOp>(op))
+    return true;
   if (auto transferRead = dyn_cast<vector::TransferReadOp>(op))
     return transferReadSupportsMMAMatrixType(transferRead);
   if (auto transferWrite = dyn_cast<vector::TransferWriteOp>(op))
     return transferWriteSupportsMMAMatrixType(transferWrite);
   if (auto contract = dyn_cast<vector::ContractionOp>(op))
     return contractSupportsMMAMatrixType(contract);
   if (auto constant = dyn_cast<ConstantOp>(op))
     return constantSupportsMMAMatrixType(constant);
static void convertConstantOp(ConstantOp 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);
   valueMapping[op.getResult()] = matrix;
 }
 
+// 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.
+  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.lowerBound(), loop.upperBound(),
+                           loop.step(), operands);
+  newLoop.getBody()->erase();
+  newLoop.getLoopBody().getBlocks().splice(
+      newLoop.getLoopBody().getBlocks().begin(),
+      loop.getLoopBody().getBlocks());
+  for (auto operand : newIterOperands)
+    newLoop.getBody()->addArgument(operand.getType());
+
+  for (auto it : llvm::zip(loop.getResults(), newLoop.getResults().take_front(
+                                                  loop.getNumResults())))
+    std::get<0>(it).replaceAllUsesWith(std::get<1>(it));
+  loop.erase();
+  return newLoop;
+}
+
+static void convertForOp(scf::ForOp op,
+                         llvm::DenseMap<Value, Value> &valueMapping) {
+  SmallVector<Value> newOperands;
+  SmallVector<std::pair<size_t, size_t>> argMapping;
+  for (auto operand : llvm::enumerate(op.getIterOperands())) {
+    auto it = valueMapping.find(operand.value());
+    if (it == valueMapping.end())
+      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);
+  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());
+  }
+}
+
+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 (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();
+}
+
 namespace mlir {
 
 void populatePrepareVectorToMMAPatterns(RewritePatternSet &patterns) {
   patterns.add<PrepareContractToGPUMMA, CombineTransferReadOpTranspose>(
       patterns.getContext());
 }
 
 void convertVectorToMMAOps(FuncOp funcOp) {
   SetVector<Operation *> ops = getOpToConvert(funcOp);
   llvm::DenseMap<Value, Value> valueMapping;
   for (Operation *op : ops) {
     if (auto transferRead = dyn_cast<vector::TransferReadOp>(op)) {
       convertTransferReadOp(transferRead, valueMapping);
     } 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<ConstantOp>(op)) {
       convertConstantOp(constantOp, 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);
     }
   }
 }
 
 } // namespace mlir
 namespace {
 
 struct ConvertVectorToGPUPass

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

func @matmul_cst(%arg0: memref<16x16xf16>, %arg1: memref<16x16xf16>, %arg2: memref<16x16xf16>) {
   %cst = 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
 }
 
-// Negative test until scf.for support is added.
 // CHECK-LABEL: func @matmul_loop
-//       CHECK:   vector.contract
+//       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:   }
+//  CHECK-NEXT:   gpu.subgroup_mma_store_matrix %[[ACC]], %{{.*}}[%{{.*}}, %{{.*}}] {leadDimension = 128 : index} : !gpu.mma_matrix<16x16xf16, "COp">, memref<128x128xf16>
 func @matmul_loop(%arg0: memref<128x128xf16>, %arg1: memref<128x128xf16>, %arg2: memref<128x128xf16>) {
   %c0 = constant 0 : index
   %c128 = constant 128 : index
   %c32 = constant 32 : index
   %cst = constant 0.000000e+00 : f16
   %C = vector.transfer_read %arg2[%c0, %c0], %cst {in_bounds = [true, true]} : memref<128x128xf16>, vector<16x16xf16>
   %14 = scf.for %arg17 = %c0 to %c128 step %c32 iter_args(%arg18 = %C) -> (vector<16x16xf16>) {
     %17 = vector.transfer_read %arg0[%c0, %arg17], %cst {in_bounds = [true, true]} : memref<128x128xf16>, 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
 }