[Mlir-commits] [mlir] [nlir][vector] Disable `vector.matrix_multiply` for scalable vectors (PR #102573)

[Andrzej Warzyński]

https://github.com/banach-space created https://github.com/llvm/llvm-project/pull/102573

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.


>From c0334a95251ca6ef700647a85e5de96c1c2cd12e Mon Sep 17 00:00:00 2001
From: Andrzej Warzynski <andrzej.warzynski at arm.com>
Date: Fri, 9 Aug 2024 06:58:58 +0100
Subject: [PATCH] [nlir][vector] Disable `vector.matrix_multiply` for scalable
 vectors

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.
---
 mlir/include/mlir/Dialect/Vector/IR/VectorOps.td | 11 +++++++----
 mlir/include/mlir/IR/CommonTypeConstraints.td    |  8 ++++++++
 mlir/test/Dialect/Vector/invalid.mlir            | 13 +++++++++++++
 3 files changed, 28 insertions(+), 4 deletions(-)

diff --git a/mlir/include/mlir/Dialect/Vector/IR/VectorOps.td b/mlir/include/mlir/Dialect/Vector/IR/VectorOps.td
index bc97a5ae7d2f7..b8559efda13e9 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 5b6ec167fa242..2eec2c6073bbf 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 ba1efe8b3c2d3..6e077a2fb4cee 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
+}