[Mlir-commits] [mlir] [mlir] Extract RHS rows once when lowering vector.contract to dot (PR #130130)

[Andrzej Warzyński]

================
@@ -758,12 +758,20 @@ FailureOr<Value> ContractionOpToDotLowering::matchAndRewriteMaskableOp(
   Value res = rewriter.create<arith::ConstantOp>(loc, dstType,
                                                  rewriter.getZeroAttr(dstType));
   bool isInt = isa<IntegerType>(dstType.getElementType());
+  llvm::SmallVector<Value> extractedCols;
+  extractedCols.reserve(dstColumns);
   for (unsigned r = 0; r < dstRows; ++r) {
     Value a = rewriter.create<vector::ExtractOp>(op.getLoc(), lhs, r);
     for (unsigned c = 0; c < dstColumns; ++c) {
-      Value b = rank == 1
-                    ? rhs
-                    : rewriter.create<vector::ExtractOp>(op.getLoc(), rhs, c);
+      if (r == 0) {
+        // We only need to extract the columns of the RHS once
+        // and then re-use them later.
+        Value b = rank == 1
+                      ? rhs
+                      : rewriter.create<vector::ExtractOp>(op.getLoc(), rhs, c);
+        extractedCols.push_back(b);
+      }
+      Value b = extractedCols[c];
----------------
banach-space wrote:

[nit] Avoid re-using/re-defining variable names. Btw, you could use `lhs` and `rhs` for `a` and `b`, respectively.


https://github.com/llvm/llvm-project/pull/130130