[mlir][vector] Disable vector.matrix_multiply for scalable vectors

[banach-space]

Disables vector.matrix_multiply for scalable vectors. As per the docs:

This is the counterpart of llvm.matrix.multiply in MLIR

I'm not aware of any use of matrix-multiply intrinsics in the context of
scalable vectors, hence disabling.

[llvmbot]

@llvm/pr-subscribers-mlir-core
@llvm/pr-subscribers-mlir
@llvm/pr-subscribers-mlir-vector

@llvm/pr-subscribers-mlir-ods

Author: Andrzej Warzyński (banach-space)

Changes

Disables vector.matrix_multiply for scalable vectors. As per the docs:

> This is the counterpart of llvm.matrix.multiply in MLIR

I'm not aware of any use of matrix-multiply intrinsics in the context of
scalable vectors, hence disabling.

---

Full diff: https://github.com/llvm/llvm-project/pull/102573.diff

3 Files Affected:

• (modified) mlir/include/mlir/Dialect/Vector/IR/VectorOps.td (+7-4)
• (modified) mlir/include/mlir/IR/CommonTypeConstraints.td (+8)
• (modified) mlir/test/Dialect/Vector/invalid.mlir (+13)

diff --git a/mlir/include/mlir/Dialect/Vector/IR/VectorOps.td b/mlir/include/mlir/Dialect/Vector/IR/VectorOps.td
index bc97a5ae7d2f70..b8559efda13e99 100644
--- a/mlir/include/mlir/Dialect/Vector/IR/VectorOps.td
+++ b/mlir/include/mlir/Dialect/Vector/IR/VectorOps.td
@@ -2683,13 +2683,13 @@ def Vector_MatmulOp : Vector_Op<"matrix_multiply", [Pure,
                     TCresVTEtIsSameAsOpBase<0, 1>>]>,
       Arguments<(
         // TODO: tighten vector element types that make sense.
-        ins VectorOfRankAndType<[1],
+        ins FixedVectorOfRankAndType<[1],
               [AnySignlessInteger, AnySignedInteger, Index, AnyFloat]>:$lhs,
-            VectorOfRankAndType<[1],
+            FixedVectorOfRankAndType<[1],
               [AnySignlessInteger, AnySignedInteger, Index, AnyFloat]>:$rhs,
             I32Attr:$lhs_rows, I32Attr:$lhs_columns, I32Attr:$rhs_columns)>,
       Results<(
-        outs VectorOfRankAndType<[1],
+        outs FixedVectorOfRankAndType<[1],
                [AnySignlessInteger, AnySignedInteger, Index, AnyFloat]>:$res)>
 {
   let summary = "Vector matrix multiplication op that operates on flattened 1-D"
@@ -2707,7 +2707,10 @@ def Vector_MatmulOp : Vector_Op<"matrix_multiply", [Pure,
     <rhs_columns> and multiplies them. The result matrix is returned embedded in
     the result vector.
 
-    Also see:
+    Note, the semantics of the corresponding LLVM intrinsic,
+    `@llvm.matrix.multiply.*`, are not clear in the context of scalable
+    vectors. Hence, this Op is only available for fixed-width vectors. Also
+    see:
 
     http://llvm.org/docs/LangRef.html#llvm-matrix-multiply-intrinsic
 
diff --git a/mlir/include/mlir/IR/CommonTypeConstraints.td b/mlir/include/mlir/IR/CommonTypeConstraints.td
index 5b6ec167fa2420..2eec2c6073bbf2 100644
--- a/mlir/include/mlir/IR/CommonTypeConstraints.td
+++ b/mlir/include/mlir/IR/CommonTypeConstraints.td
@@ -494,6 +494,14 @@ class VectorOfRankAndType<list<int> allowedRanks,
   VectorOf<allowedTypes>.summary # VectorOfRank<allowedRanks>.summary,
   "::mlir::VectorType">;
 
+// Fixed-width vector where the rank is from the given `allowedRanks` list and
+// the type is from the given `allowedTypes` list
+class FixedVectorOfRankAndType<list<int> allowedRanks,
+                          list<Type> allowedTypes> : AllOfType<
+  [FixedVectorOf<allowedTypes>, VectorOfRank<allowedRanks>],
+  FixedVectorOf<allowedTypes>.summary # VectorOfRank<allowedRanks>.summary,
+  "::mlir::VectorType">;
+
 // Whether the number of elements of a vector is from the given
 // `allowedLengths` list
 class IsVectorOfLengthPred<list<int> allowedLengths> :
diff --git a/mlir/test/Dialect/Vector/invalid.mlir b/mlir/test/Dialect/Vector/invalid.mlir
index ba1efe8b3c2d38..6e077a2fb4cee4 100644
--- a/mlir/test/Dialect/Vector/invalid.mlir
+++ b/mlir/test/Dialect/Vector/invalid.mlir
@@ -1862,3 +1862,16 @@ func.func @invalid_step_2d() {
   vector.step : vector<2x4xf32>
   return
 }
+
+// -----
+
+func.func @matrix_matmul_scalable(%a: vector<[4]xf64>, %b: vector<4xf64>) {
+  // expected-error @+1 {{'vector.matrix_multiply' op operand #0 must be fixed-length vector of signless integer or signed integer or index or floating-point values of ranks 1, but got 'vector<[4]xf64>'}}
+  %c = vector.matrix_multiply %a, %b {
+    lhs_rows = 2: i32,
+    lhs_columns = 2: i32 ,
+    rhs_columns = 2: i32 }
+  : (vector<[4]xf64>, vector<4xf64>) -> vector<4xf64>
+
+  return
+}

[c-rhodes]

Thanks for addressing this, left a few minor comments, also the commit header needs fixing s/nlir/mlir

[c-rhodes]

LGTM cheers