[Mlir-commits] [mlir] [mlir][vector] Restrict vector.shape_cast (scalable vectors) (PR #100331)
Andrzej Warzyński
llvmlistbot at llvm.org
Wed Jul 24 10:23:47 PDT 2024
https://github.com/banach-space updated https://github.com/llvm/llvm-project/pull/100331
>From 858c6a7e4b41b1b05da77dec177de33dfdf82c6d Mon Sep 17 00:00:00 2001
From: Andrzej Warzynski <andrzej.warzynski at arm.com>
Date: Wed, 24 Jul 2024 09:54:11 +0100
Subject: [PATCH 1/3] [mlir][vector] Restrict vector.shape_cast (scalable
vectors)
Updates the verifier for `vector.shape_cast` so that the following
incorrect cases are immediately rejected:
```mlir
vector.shape_cast %vec : vector<1x1x[4]xindex> to vector<4xindex>
```
Seperately, here's a fix for the Linalg vectorizer to prevent the
vectorizer from generating such shape casts (*):
* https://github.com/llvm/llvm-project/pull/100325
(*) Note, that's just one specific case that I've identified so far.
---
mlir/lib/Dialect/Vector/IR/VectorOps.cpp | 5 +++++
mlir/test/Dialect/Vector/invalid.mlir | 7 +++++++
2 files changed, 12 insertions(+)
diff --git a/mlir/lib/Dialect/Vector/IR/VectorOps.cpp b/mlir/lib/Dialect/Vector/IR/VectorOps.cpp
index df3a59ed80ad4..f145dc9f8817d 100644
--- a/mlir/lib/Dialect/Vector/IR/VectorOps.cpp
+++ b/mlir/lib/Dialect/Vector/IR/VectorOps.cpp
@@ -5238,6 +5238,11 @@ static LogicalResult verifyVectorShapeCast(Operation *op,
if (!isValidShapeCast(resultShape, sourceShape))
return op->emitOpError("invalid shape cast");
}
+
+ // Check that (non-)scalability is preserved
+ if (sourceVectorType.isScalable() != resultVectorType.isScalable())
+ return op->emitOpError("non-matching scalability flags");
+
return success();
}
diff --git a/mlir/test/Dialect/Vector/invalid.mlir b/mlir/test/Dialect/Vector/invalid.mlir
index 208982a3e0e7b..3fad61198b474 100644
--- a/mlir/test/Dialect/Vector/invalid.mlir
+++ b/mlir/test/Dialect/Vector/invalid.mlir
@@ -1182,6 +1182,13 @@ func.func @shape_cast_invalid_rank_expansion(%arg0 : vector<15x2xf32>) {
// -----
+func.func @shape_cast_scalability_flag_is_dropped(%arg0 : vector<15x[2]xf32>) {
+ // expected-error at +1 {{non-matching scalability flags}}
+ %0 = vector.shape_cast %arg0 : vector<15x[2]xf32> to vector<30xf32>
+}
+
+// -----
+
func.func @bitcast_not_vector(%arg0 : vector<5x1x3x2xf32>) {
// expected-error at +1 {{'vector.bitcast' invalid kind of type specified}}
%0 = vector.bitcast %arg0 : vector<5x1x3x2xf32> to f32
>From f67115a0805144c907a04d7e4fd1dd981b81833c Mon Sep 17 00:00:00 2001
From: Andrzej Warzynski <andrzej.warzynski at arm.com>
Date: Wed, 24 Jul 2024 17:21:20 +0100
Subject: [PATCH 2/3] fixup! [mlir][vector] Restrict vector.shape_cast
(scalable vectors)
Make the verifier even stricter
---
mlir/include/mlir/IR/BuiltinTypes.td | 6 ++++++
mlir/lib/Dialect/Vector/IR/VectorOps.cpp | 6 ++++--
mlir/test/Dialect/Vector/invalid.mlir | 9 ++++++++-
3 files changed, 18 insertions(+), 3 deletions(-)
diff --git a/mlir/include/mlir/IR/BuiltinTypes.td b/mlir/include/mlir/IR/BuiltinTypes.td
index 0b3532dcc7d4f..079d1b4921645 100644
--- a/mlir/include/mlir/IR/BuiltinTypes.td
+++ b/mlir/include/mlir/IR/BuiltinTypes.td
@@ -1168,6 +1168,12 @@ def Builtin_Vector : Builtin_Type<"Vector", "vector",
return !llvm::is_contained(getScalableDims(), false);
}
+ /// Get the number of scalable dimension.
+ int64_t getNumScalableDims() const {
+ return llvm::count(getScalableDims(), true);
+ }
+
+
/// Get or create a new VectorType with the same shape as `this` and an
/// element type of bitwidth scaled by `scale`.
/// Return null if the scaled element type cannot be represented.
diff --git a/mlir/lib/Dialect/Vector/IR/VectorOps.cpp b/mlir/lib/Dialect/Vector/IR/VectorOps.cpp
index f145dc9f8817d..d52411a54c36e 100644
--- a/mlir/lib/Dialect/Vector/IR/VectorOps.cpp
+++ b/mlir/lib/Dialect/Vector/IR/VectorOps.cpp
@@ -5240,8 +5240,10 @@ static LogicalResult verifyVectorShapeCast(Operation *op,
}
// Check that (non-)scalability is preserved
- if (sourceVectorType.isScalable() != resultVectorType.isScalable())
- return op->emitOpError("non-matching scalability flags");
+ if (sourceVectorType.getNumScalableDims() !=
+ resultVectorType.getNumScalableDims())
+ return op->emitOpError("non-matching scalable dims");
+ sourceVectorType.getNumDynamicDims();
return success();
}
diff --git a/mlir/test/Dialect/Vector/invalid.mlir b/mlir/test/Dialect/Vector/invalid.mlir
index 3fad61198b474..79fbdc55ed426 100644
--- a/mlir/test/Dialect/Vector/invalid.mlir
+++ b/mlir/test/Dialect/Vector/invalid.mlir
@@ -1183,12 +1183,19 @@ func.func @shape_cast_invalid_rank_expansion(%arg0 : vector<15x2xf32>) {
// -----
func.func @shape_cast_scalability_flag_is_dropped(%arg0 : vector<15x[2]xf32>) {
- // expected-error at +1 {{non-matching scalability flags}}
+ // expected-error at +1 {{non-matching scalable dims}}
%0 = vector.shape_cast %arg0 : vector<15x[2]xf32> to vector<30xf32>
}
// -----
+func.func @shape_cast_scalability_flag_is_dropped(%arg0 : vector<2x[15]x[2]xf32>) {
+ // expected-error at +1 {{non-matching scalable dims}}
+ %0 = vector.shape_cast %arg0 : vector<2x[15]x[2]xf32> to vector<30x[2]f32>
+}
+
+// -----
+
func.func @bitcast_not_vector(%arg0 : vector<5x1x3x2xf32>) {
// expected-error at +1 {{'vector.bitcast' invalid kind of type specified}}
%0 = vector.bitcast %arg0 : vector<5x1x3x2xf32> to f32
>From 946ff5b7d58c68004e1ff28e949e8205d087580c Mon Sep 17 00:00:00 2001
From: Andrzej Warzynski <andrzej.warzynski at arm.com>
Date: Wed, 24 Jul 2024 18:23:15 +0100
Subject: [PATCH 3/3] fixup! fixup! [mlir][vector] Restrict vector.shape_cast
(scalable vectors)
Fix typos
---
mlir/include/mlir/IR/BuiltinTypes.td | 2 +-
mlir/test/Dialect/Vector/invalid.mlir | 2 +-
2 files changed, 2 insertions(+), 2 deletions(-)
diff --git a/mlir/include/mlir/IR/BuiltinTypes.td b/mlir/include/mlir/IR/BuiltinTypes.td
index 079d1b4921645..38419dfa50106 100644
--- a/mlir/include/mlir/IR/BuiltinTypes.td
+++ b/mlir/include/mlir/IR/BuiltinTypes.td
@@ -1168,7 +1168,7 @@ def Builtin_Vector : Builtin_Type<"Vector", "vector",
return !llvm::is_contained(getScalableDims(), false);
}
- /// Get the number of scalable dimension.
+ /// Get the number of scalable dimensions.
int64_t getNumScalableDims() const {
return llvm::count(getScalableDims(), true);
}
diff --git a/mlir/test/Dialect/Vector/invalid.mlir b/mlir/test/Dialect/Vector/invalid.mlir
index 79fbdc55ed426..4fb9f7b7d9854 100644
--- a/mlir/test/Dialect/Vector/invalid.mlir
+++ b/mlir/test/Dialect/Vector/invalid.mlir
@@ -1191,7 +1191,7 @@ func.func @shape_cast_scalability_flag_is_dropped(%arg0 : vector<15x[2]xf32>) {
func.func @shape_cast_scalability_flag_is_dropped(%arg0 : vector<2x[15]x[2]xf32>) {
// expected-error at +1 {{non-matching scalable dims}}
- %0 = vector.shape_cast %arg0 : vector<2x[15]x[2]xf32> to vector<30x[2]f32>
+ %0 = vector.shape_cast %arg0 : vector<2x[15]x[2]xf32> to vector<30x[2]xf32>
}
// -----
More information about the Mlir-commits
mailing list