Skip to content

[mlir][vector] Restrict vector.shape_cast (scalable vectors) #100331

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
merged 6 commits into from
Jul 25, 2024

Conversation

banach-space
Copy link
Contributor

Updates the verifier for vector.shape_cast so that the following
incorrect cases are immediately rejected:

  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 (*):

(*) Note, that's just one specific case that I've identified so far.

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 (*):
* llvm#100325

(*) Note, that's just one specific case that I've identified so far.
@llvmbot
Copy link
Member

llvmbot commented Jul 24, 2024

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

@llvm/pr-subscribers-mlir-vector

Author: Andrzej Warzyński (banach-space)

Changes

Updates the verifier for vector.shape_cast so that the following
incorrect cases are immediately rejected:

  vector.shape_cast %vec : vector&lt;1x1x[4]xindex&gt; to vector&lt;4xindex&gt;

Seperately, here's a fix for the Linalg vectorizer to prevent the
vectorizer from generating such shape casts (*):

(*) Note, that's just one specific case that I've identified so far.


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

2 Files Affected:

  • (modified) mlir/lib/Dialect/Vector/IR/VectorOps.cpp (+5)
  • (modified) mlir/test/Dialect/Vector/invalid.mlir (+7)
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@+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@+1 {{'vector.bitcast' invalid kind of type specified}}
   %0 = vector.bitcast %arg0 : vector<5x1x3x2xf32> to f32

@llvmbot llvmbot added mlir:core MLIR Core Infrastructure mlir:ods labels Jul 24, 2024
Copy link
Member

@kuhar kuhar left a comment

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Thanks. LGTM % nit.

@banach-space banach-space merged commit 98c73d5 into llvm:main Jul 25, 2024
4 of 7 checks passed
@banach-space banach-space deleted the andrzej/restrict_shape_cast branch July 25, 2024 11:39
yuxuanchen1997 pushed a commit that referenced this pull request Jul 25, 2024
Summary:
Updates the verifier for `vector.shape_cast` so that incorrect cases
where "scalability" is dropped are immediately rejected. For example:
```mlir
  vector.shape_cast %vec : vector<1x1x[4]xindex> to vector<4xindex>
```

Also, as a separate PR, I've prepared a fix for the Linalg vectorizer to
avoid generating such shape casts (*):
* #100325

(*) Note, that's just one specific case that I've identified so far.

Test Plan: 

Reviewers: 

Subscribers: 

Tasks: 

Tags: 


Differential Revision: https://phabricator.intern.facebook.com/D60250589
banach-space added a commit to banach-space/llvm-project that referenced this pull request Jul 26, 2024
Updates the codebase to use the recently introduced VectorType helper:
  * `getNumScalableDims` (introduced in llvm#100331).
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment
Projects
None yet
Development

Successfully merging this pull request may close these issues.

4 participants