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

[Andrzej Warzyński]

================
@@ -1339,6 +1339,50 @@ bool ExtractOp::isCompatibleReturnTypes(TypeRange l, TypeRange r) {
   return l == r;
 }
 
+// Common verification rules for `InsertOp` and `ExtractOp` involving indices.
+// `indexedType` is the vector type being indexed in the operation, i.e., the
+// destination type in InsertOp and the source type in ExtractOp.
+// `vecOrScalarType` is the type that is not indexed in the op and can be
+// either a scalar or a vector, i.e., the source type in InsertOp and the
+// return type in ExtractOp.
+static LogicalResult verifyInsertExtractIndices(Operation *op,
+                                                VectorType indexedType,
+                                                int64_t numIndices,
+                                                Type vecOrScalarType) {
+  int64_t indexedRank = indexedType.getRank();
+  if (numIndices > indexedRank)
+    return op->emitOpError(
+        "expected a number of indices no greater than the indexed vector rank");
+
+  if (auto nonIndexedVecType = dyn_cast<VectorType>(vecOrScalarType)) {
+    // Vector case, including:
+    //  * 0-D vector:
+    //    * vector.extract %src[2]: vector<f32> from vector<8xf32)
+    //    * vector.insert %src, %dst[3]: vector<f32> into vector<8xf32>
+    //  * One-element vector:
+    //    * vector.extract %src[4]: vector<1xf32> from vector<8xf32>
+    //    * vector.insert %src, %dst[1]: vector<1xf32> into vector<8xf32>
+    //    * vector.extract %src[7]: vector<1xf32> from vector<8x1xf32>
+    //    * vector.insert %src, %dst[5]: vector<1xf32> into vector<8x1xf32>
+    int64_t indexedRankMinusIndices = indexedRank - numIndices;
----------------
banach-space wrote:

Looks like you've already simplified the way I had in mind, thanks!

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