[Mlir-commits] [mlir] [MLIR][AArch64] Check indexing maps before checking for dimensions compatibility (PR #145702)

[Andrzej Warzyński]

================
@@ -136,11 +136,26 @@ class LowerContractionToSVEI8MMPattern
   LogicalResult matchAndRewrite(vector::ContractionOp op,
                                 PatternRewriter &rewriter) const override {
 
-    Location loc = op.getLoc();
+    // Check permutation maps. For now only accept
+    //   lhs: (d0, d1, d2) -> (d0, d2)
+    //   rhs: (d0, d1, d2) -> (d1, d2)
+    //   acc: (d0, d1, d2) -> (d0, d1)
+    // This corresponds to matrix multiplication with transposed RHS.
+    if (op.getIndexingMapsArray()[0] !=
+            AffineMap::getMultiDimMapWithTargets(3, ArrayRef{0u, 2u},
+                                                 op.getContext()) ||
+        op.getIndexingMapsArray()[1] !=
+            AffineMap::getMultiDimMapWithTargets(3, ArrayRef{1u, 2u},
+                                                 op.getContext()) ||
+        op.getIndexingMapsArray()[2] !=
+            AffineMap::getMultiDimMapWithTargets(3, ArrayRef{0u, 1u},
+                                                 op.getContext()))
+      return rewriter.notifyMatchFailure(op, "non-matching permutation maps");
----------------
banach-space wrote:

Don't these permutation maps imply the dimensionality of `lhs` and `rhs`? It feels like there's no need to check the rank below?

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