[mlir] ViewLikeInterface - verify ranks in verifyOffsetSizeAndStrideOp

[amrami]

getMixedOffsets() calls getMixedValues() with static_offsets and offsets. It is assumed that the number of dynamic offsets in static_offsets equals the rank of offsets. Otherwise, we fail on assert when trying to access an array out of its bounds.
The same applies to getMixedStrides() and getMixedOffsets().

A verification of this assumption is added to verifyOffsetSizeAndStrideOp() and a clear assert is added in getMixedValues().

[llvmbot]

@llvm/pr-subscribers-mlir-memref

@llvm/pr-subscribers-mlir

Author: Maya Amrami (amrami)

Changes

getMixedOffsets() calls getMixedValues() with static_offsets and offsets. It is assumed that the number of dynamic offsets in static_offsets equals the rank of offsets. Otherwise, we fail on assert when trying to access an array out of its bounds.
The same applies to getMixedStrides() and getMixedOffsets().

A verification of this assumption is added to verifyOffsetSizeAndStrideOp() and a clear assert is added in getMixedValues().

---

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

3 Files Affected:

• (modified) mlir/lib/Dialect/Utils/StaticValueUtils.cpp (+5-1)
• (modified) mlir/lib/Interfaces/ViewLikeInterface.cpp (+26)
• (modified) mlir/test/Dialect/MemRef/invalid.mlir (+12)

diff --git a/mlir/lib/Dialect/Utils/StaticValueUtils.cpp b/mlir/lib/Dialect/Utils/StaticValueUtils.cpp
index 1cded38c4419e..af494b99cb4cb 100644
--- a/mlir/lib/Dialect/Utils/StaticValueUtils.cpp
+++ b/mlir/lib/Dialect/Utils/StaticValueUtils.cpp
@@ -181,12 +181,16 @@ bool isEqualConstantIntOrValueArray(ArrayRef<OpFoldResult> ofrs1,
   return true;
 }
 
-/// Return a vector of OpFoldResults with the same size a staticValues, but all
+/// Return a vector of OpFoldResults with the same size as staticValues, but all
 /// elements for which ShapedType::isDynamic is true, will be replaced by
 /// dynamicValues.
 SmallVector<OpFoldResult> getMixedValues(ArrayRef<int64_t> staticValues,
                                          ValueRange dynamicValues,
                                          MLIRContext *context) {
+  assert(llvm::count_if(staticValues, ShapedType::isDynamic) ==
+                 dynamicValues.size() &&
+             "expected the rank of dynamic values to match the number of"
+             " values known to be dynamic";);
   SmallVector<OpFoldResult> res;
   res.reserve(staticValues.size());
   unsigned numDynamic = 0;
diff --git a/mlir/lib/Interfaces/ViewLikeInterface.cpp b/mlir/lib/Interfaces/ViewLikeInterface.cpp
index 3112da9ef182a..677d7dcc1b135 100644
--- a/mlir/lib/Interfaces/ViewLikeInterface.cpp
+++ b/mlir/lib/Interfaces/ViewLikeInterface.cpp
@@ -94,6 +94,32 @@ SliceBoundsVerificationResult mlir::verifyInBoundsSlice(
 
 LogicalResult
 mlir::detail::verifyOffsetSizeAndStrideOp(OffsetSizeAndStrideOpInterface op) {
+  // A dynamic size is represented as ShapedType::kDynamic in `static_sizes`.
+  // Its corresponding Value appears in `sizes`. Thus, the number of dynamic
+  // dimensions in `static_sizes` must equal the rank of `sizes`.
+  // The same applies to strides and offsets.
+  unsigned int numDynamicDims =
+      llvm::count_if(op.getStaticSizes(), ShapedType::isDynamic);
+  if (numDynamicDims != op.getSizes().size()) {
+    return op->emitError("expected sizes rank to match the number of dynamic "
+                         "dimensions (")
+           << numDynamicDims << " vs " << op.getSizes().size() << ")";
+  }
+  unsigned int numDynamicStrides =
+      llvm::count_if(op.getStaticStrides(), ShapedType::isDynamic);
+  if (numDynamicStrides != op.getStrides().size()) {
+    return op->emitError("expected strides rank to match the number of dynamic "
+                         "dimensions (")
+           << numDynamicStrides << " vs " << op.getStrides().size() << ")";
+  }
+  unsigned int numDynamicOffsets =
+      llvm::count_if(op.getStaticOffsets(), ShapedType::isDynamic);
+  if (numDynamicOffsets != op.getOffsets().size()) {
+    return op->emitError("expected offsets rank to match the number of dynamic "
+                         "dimensions (")
+           << numDynamicOffsets << " vs " << op.getOffsets().size() << ")";
+  }
+
   std::array<unsigned, 3> maxRanks = op.getArrayAttrMaxRanks();
   // Offsets can come in 2 flavors:
   //   1. Either single entry (when maxRanks == 1).
diff --git a/mlir/test/Dialect/MemRef/invalid.mlir b/mlir/test/Dialect/MemRef/invalid.mlir
index 704cdaf838f45..43d0362b78e1d 100644
--- a/mlir/test/Dialect/MemRef/invalid.mlir
+++ b/mlir/test/Dialect/MemRef/invalid.mlir
@@ -658,6 +658,18 @@ func.func @invalid_subview(%arg0 : index, %arg1 : index, %arg2 : index) {
 
 // -----
 
+// This test is not written in the op's assembly format, to reproduce a mismatch
+// between the rank of static_offsets and the number of Values sent as the
+// dynamic offsets.
+func.func @invalid_subview(%arg0 : memref<?x128xi8, 1>) {
+  %0 = memref.alloc() :memref<1xf32>
+  // expected-error@+1 {{expected offsets rank to match the number of dynamic dimensions (1 vs 0)}}
+  "memref.subview"(%0) <{operandSegmentSizes = array<i32: 1, 0, 0, 0>, static_offsets = array<i64: -9223372036854775808>, static_sizes = array<i64: 1>, static_strides = array<i64: 1>}> : (memref<1xf32>) -> memref<1xf32, strided<[1], offset: ?>>
+  return
+}
+
+// -----
+
 func.func @invalid_subview(%arg0 : index, %arg1 : index, %arg2 : index) {
   %0 = memref.alloc() : memref<8x16x4xf32>
   // expected-error@+1 {{expected mixed sizes rank to match mixed strides rank (3 vs 2) so the rank of the result type is well-formed}}