[Mlir-commits] [mlir] [mlir][vector] support unrolling vector.contract with scalar result (PR #232203)

[llvmlistbot at llvm.org]

================
@@ -343,21 +343,28 @@ struct UnrollContractionPattern
       auto *accIt = accCache.find(accOffets);
       if (accIt != accCache.end())
         slicesOperands[2] = accIt->second;
+      else if (scalarResult)
+        slicesOperands[2] = contractOp.getAcc();
       else
         extractOperand(2, contractOp.getAcc(), accPermutationMap, accOffets);
-
-      SmallVector<int64_t> dstShape =
-          applyPermutationMap(dstAffineMap, ArrayRef<int64_t>(*targetShape));
-      auto targetType = VectorType::get(dstShape, dstVecType.getElementType());
+      Type targetType = contractOp.getResultType();
+      if (!scalarResult) {
+        SmallVector<int64_t> dstShape =
+            applyPermutationMap(dstAffineMap, ArrayRef<int64_t>(*targetShape));
+        targetType = VectorType::get(dstShape, dstVecType.getElementType());
+      }
       Operation *newOp = cloneOpWithOperandsAndTypes(
           rewriter, loc, contractOp, slicesOperands, targetType);
-
       SmallVector<int64_t> dstOffets =
           applyPermutationMap(dstAffineMap, ArrayRef<int64_t>(offsets));
       // Save the accumulated value untill all the loops are unrolled since
       // reduction loop keep updating the accumulator.
       accCache[dstOffets] = newOp->getResult(0);
     }
+    if (scalarResult) {
+      rewriter.replaceOp(contractOp, accCache.begin()->second);
+      return success();
+    }
----------------
StarOne01 wrote:

couldn't get it to one line, hence adding a bonus line hehe

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