From 41481dfdc1a159efe1724712c96f1a680a4d4fa6 Mon Sep 17 00:00:00 2001
From: Maya Amrami <mayaam88@gmail.com>
Date: Thu, 10 Jul 2025 12:05:48 +0300
Subject: [PATCH] [mlir] ViewLikeInterface - verify ranks in
 verifyOffsetSizeAndStrideOp

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().
---
 mlir/lib/Dialect/Utils/StaticValueUtils.cpp |  6 ++++-
 mlir/lib/Interfaces/ViewLikeInterface.cpp   | 26 +++++++++++++++++++++
 mlir/test/Dialect/MemRef/invalid.mlir       | 12 ++++++++++
 3 files changed, 43 insertions(+), 1 deletion(-)

diff --git a/mlir/lib/Dialect/Utils/StaticValueUtils.cpp b/mlir/lib/Dialect/Utils/StaticValueUtils.cpp
index 1cded38c4419e..c3cfc1c1a6fe9 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(dynamicValues.size() ==
+             (unsigned)llvm::count_if(staticValues, ShapedType::isDynamic) &&
+         "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..2f27f6f43b127 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 (op.getSizes().size() != numDynamicDims) {
+    return op->emitError("expected sizes rank to match the number of dynamic "
+                         "dimensions (")
+           << op.getSizes().size() << " vs " << numDynamicDims << ")";
+  }
+  unsigned int numDynamicStrides =
+      llvm::count_if(op.getStaticStrides(), ShapedType::isDynamic);
+  if (op.getStrides().size() != numDynamicStrides) {
+    return op->emitError("expected strides rank to match the number of dynamic "
+                         "strides (")
+           << op.getStrides().size() << " vs " << numDynamicStrides << ")";
+  }
+  unsigned int numDynamicOffsets =
+      llvm::count_if(op.getStaticOffsets(), ShapedType::isDynamic);
+  if (op.getOffsets().size() != numDynamicOffsets) {
+    return op->emitError("expected offsets rank to match the number of dynamic "
+                         "offsets (")
+           << op.getOffsets().size() << " vs " << numDynamicOffsets << ")";
+  }
+
   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..e5479d82edc02 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 offsets (0 vs 1)}}
+  "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}}