[Mlir-commits] [mlir] [mlir][ArmSME] Support 4-way widening outer products (PR #79288)

[Benjamin Maxwell]

================
@@ -225,37 +271,189 @@ class OuterProductFusion2Way
 
     return success();
   }
+};
 
-  // An outer product is compatible if all of the following are true:
-  // - the result type matches `resultType`.
-  // - the defining operations of the inputs are identical and of the type
-  //   `ExtOp`.
-  // - the input types of the defining operations are identical and match
-  //   `inputType`.
-  template <typename ExtOp>
-  LogicalResult isCompatible(PatternRewriter &rewriter,
-                             arm_sme::OuterProductOp op, VectorType resultType,
-                             VectorType inputType) const {
-    if (op.getResultType() != resultType)
-      return rewriter.notifyMatchFailure(op.getLoc(), [&](Diagnostic &diag) {
-        diag << "unsupported result type, expected " << resultType;
-      });
-
-    auto lhsDefOp = op.getLhs().getDefiningOp<ExtOp>();
-    auto rhsDefOp = op.getRhs().getDefiningOp<ExtOp>();
-
-    if (!lhsDefOp || !rhsDefOp)
-      return rewriter.notifyMatchFailure(
-          op, "defining op of outerproduct operands must be one of: "
-              "'arith.extf' or 'arith.extsi' or 'arith.extui'");
+// Fuse four 'arm_sme.outerproduct' operations that are chained via the
+// accumulator into 4-way outer product operation.
+class OuterProductFusion4Way
+    : public OpRewritePattern<arm_sme::OuterProductOp> {
+public:
+  using OpRewritePattern::OpRewritePattern;
 
-    auto lhsInType = cast<VectorType>(lhsDefOp.getIn().getType());
-    auto rhsInType = cast<VectorType>(rhsDefOp.getIn().getType());
+  LogicalResult matchAndRewrite(arm_sme::OuterProductOp op,
+                                PatternRewriter &rewriter) const override {
+    SmallVector<arm_sme::OuterProductOp, 4> outerProductChain;
+    outerProductChain.push_back(op);
+
+    for (int i = 0; i < 3; ++i) {
+      auto currentOp = outerProductChain.back();
+      auto acc = currentOp.getAcc();
+      if (!acc)
+        return rewriter.notifyMatchFailure(op, MATCH_FAILURE_NO_ACCUMULATOR);
+      auto previousOp = acc.getDefiningOp<arm_sme::OuterProductOp>();
+      if (!previousOp)
+        return rewriter.notifyMatchFailure(
+            op, MATCH_FAILURE_EXPECTED_OUTERPRODUCT_DEF_OP);
+      if (!previousOp->hasOneUse())
+        return rewriter.notifyMatchFailure(
+            op, MATCH_FAILURE_OUTERPRODUCT_NOT_SINGLE_USE);
+      if (previousOp.getKind() != currentOp.getKind())
+        return rewriter.notifyMatchFailure(
+            op, MATCH_FAILURE_INCONSISTENT_COMBINING_KIND);
+      if (bool(previousOp.getLhsMask()) != bool(currentOp.getLhsMask()))
+        return rewriter.notifyMatchFailure(
+            op, MATCH_FAILURE_INCONSISTENT_COMBINING_KIND);
+      outerProductChain.push_back(previousOp);
+    }
 
-    if (lhsInType != inputType || rhsInType != inputType)
-      return rewriter.notifyMatchFailure(op.getLoc(), [&](Diagnostic &diag) {
-        diag << "unsupported input type, expected " << inputType;
-      });
+    if (failed(canFuseOuterProducts(rewriter, outerProductChain)))
+      return failure();
+
+    arm_sme::OuterProductOp op1 = outerProductChain[3];
+    arm_sme::OuterProductOp op2 = outerProductChain[2];
+    arm_sme::OuterProductOp op3 = outerProductChain[1];
+    arm_sme::OuterProductOp op4 = outerProductChain[0];
+
+    auto loc = op.getLoc();
+
+    auto packInputs = [&](Value lhs, Value rhs) {
+      auto inputType = cast<VectorType>(lhs.getType());
+      VectorType inputTypeX2 =
+          VectorType::Builder(inputType).setDim(0, inputType.getShape()[0] * 2);
+      return rewriter.create<LLVM::experimental_vector_interleave2>(
+          loc, inputTypeX2, lhs, rhs);
+    };
----------------
MacDue wrote:

Move to helper? This identical to the copy of this in the 2-way code.

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