[Mlir-commits] [mlir] [mlir][vector] Add verification for incorrect vector.extract (PR #115824)

[llvmlistbot at llvm.org]

llvmbot wrote:


<!--LLVM PR SUMMARY COMMENT-->

@llvm/pr-subscribers-mlir-vector

Author: Diego Caballero (dcaballe)

<details>
<summary>Changes</summary>

This PR fixes the `vector.extract` verifier so that we have to provide as many indices as vector dimensions to extract a scalar. I.e., the following example is incorrect:

```
  %1 = vector.extract %arg0[0, 0] : f32 from vector<4x8x1xf32>
```

A similar check already exists for `vector.insert`.

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


3 Files Affected:

- (modified) mlir/lib/Dialect/Vector/IR/VectorOps.cpp (+22-2) 
- (modified) mlir/test/Dialect/Vector/invalid.mlir (+7) 
- (modified) mlir/test/Dialect/Vector/ops.mlir (+7) 


``````````diff
diff --git a/mlir/lib/Dialect/Vector/IR/VectorOps.cpp b/mlir/lib/Dialect/Vector/IR/VectorOps.cpp
index db199a46e1637c..9b19b5cd0cd221 100644
--- a/mlir/lib/Dialect/Vector/IR/VectorOps.cpp
+++ b/mlir/lib/Dialect/Vector/IR/VectorOps.cpp
@@ -1349,13 +1349,33 @@ LogicalResult vector::ExtractOp::verify() {
         "corresponding dynamic position) -- this can only happen due to an "
         "incorrect fold/rewrite");
   auto position = getMixedPosition();
-  if (position.size() > static_cast<unsigned>(getSourceVectorType().getRank()))
+  VectorType srcVecType = getSourceVectorType();
+  int64_t srcRank = srcVecType.getRank();
+  if (position.size() > static_cast<unsigned>(srcRank))
     return emitOpError(
         "expected position attribute of rank no greater than vector rank");
+
+  VectorType dstVecType = dyn_cast<VectorType>(getResult().getType());
+  if (dstVecType) {
+    int64_t srcRankMinusIndices = srcRank - getNumIndices();
+    int64_t dstRank = dstVecType.getRank();
+    if ((srcRankMinusIndices == 0 && dstRank != 1) ||
+        (srcRankMinusIndices != 0 && srcRankMinusIndices != dstRank)) {
+      return emitOpError(
+          "expected source rank minus number of indices to match "
+          "destination vector rank");
+    }
+  } else {
+    // Scalar result.
+    if (srcRank != getNumIndices())
+      return emitOpError("expected source rank to match number of indices "
+                         "for scalar result");
+  }
+
   for (auto [idx, pos] : llvm::enumerate(position)) {
     if (pos.is<Attribute>()) {
       int64_t constIdx = cast<IntegerAttr>(pos.get<Attribute>()).getInt();
-      if (constIdx < 0 || constIdx >= getSourceVectorType().getDimSize(idx)) {
+      if (constIdx < 0 || constIdx >= srcVecType.getDimSize(idx)) {
         return emitOpError("expected position attribute #")
                << (idx + 1)
                << " to be a non-negative integer smaller than the "
diff --git a/mlir/test/Dialect/Vector/invalid.mlir b/mlir/test/Dialect/Vector/invalid.mlir
index d591c60acb64e7..e09b031dd49648 100644
--- a/mlir/test/Dialect/Vector/invalid.mlir
+++ b/mlir/test/Dialect/Vector/invalid.mlir
@@ -192,6 +192,13 @@ func.func @extract_position_overflow(%arg0: vector<4x8x16xf32>) {
 
 // -----
 
+func.func @extract_scalar_missing_indices(%arg0: vector<4x8x1xf32>) {
+  // expected-error at +1 {{expected source rank to match number of indices for scalar result}}
+  %1 = vector.extract %arg0[0, 0] : f32 from vector<4x8x1xf32>
+}
+
+// -----
+
 func.func @insert_element(%arg0: f32, %arg1: vector<f32>) {
   %c = arith.constant 3 : i32
   // expected-error at +1 {{expected position to be empty with 0-D vector}}
diff --git a/mlir/test/Dialect/Vector/ops.mlir b/mlir/test/Dialect/Vector/ops.mlir
index 3baacba9b61243..8595b0635cc94a 100644
--- a/mlir/test/Dialect/Vector/ops.mlir
+++ b/mlir/test/Dialect/Vector/ops.mlir
@@ -240,6 +240,13 @@ func.func @extract_0d(%a: vector<f32>) -> f32 {
   return %0 : f32
 }
 
+// CHECK-LABEL: @extract_single_element_vector
+func.func @extract_single_element_vector(%arg0: vector<4x8x3xf32>) -> vector<1xf32> {
+  // CHECK: vector.extract {{.*}}[0, 0, 0] : vector<1xf32> from vector<4x8x3xf32>
+  %1 = vector.extract %arg0[0, 0, 0] : vector<1xf32> from vector<4x8x3xf32>
+  return %1 : vector<1xf32>
+}
+
 // CHECK-LABEL: @insert_element_0d
 func.func @insert_element_0d(%a: f32, %b: vector<f32>) -> vector<f32> {
   // CHECK-NEXT: vector.insertelement %{{.*}}, %{{.*}}[] : vector<f32>

``````````

</details>


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