[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.
----------------
banach-space wrote:

[nit] This implies that RHS is somehow special, but that's not really what you mean, right? Instead, you are making sure that for _both_ LHS and RHS we extract the rows/cols only once. Could you rephrase?

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