[Mlir-commits] [mlir] [mlir][memref] Fix runtime verification of expand_shape (PR #206125)
llvmlistbot at llvm.org
llvmlistbot at llvm.org
Fri Jun 26 09:28:51 PDT 2026
llvmorg-github-actions[bot] wrote:
<!--LLVM PR SUMMARY COMMENT-->
@llvm/pr-subscribers-mlir
Author: Longsheng Mou (CoTinker)
<details>
<summary>Changes</summary>
The original runtime verification only multiplied static output dimensions per reassociation group and checked `srcDim % staticProduct` == 0, which failed to catch invalid dynamic dimension values.
The new verification uses getMixedOutputShape() to get all output dimensions (static and dynamic), computes their product, and asserts `product == srcDim`. Fixes #<!-- -->205981.
Assisted-by: MiMo-V2.5-Pro
---
Full diff: https://github.com/llvm/llvm-project/pull/206125.diff
3 Files Affected:
- (modified) mlir/lib/Dialect/MemRef/Transforms/RuntimeOpVerification.cpp (+15-25)
- (modified) mlir/test/Dialect/MemRef/runtime-verification.mlir (+6-6)
- (added) mlir/test/Integration/Dialect/MemRef/expand-shape-runtime-verification.mlir (+86)
``````````diff
diff --git a/mlir/lib/Dialect/MemRef/Transforms/RuntimeOpVerification.cpp b/mlir/lib/Dialect/MemRef/Transforms/RuntimeOpVerification.cpp
index d7fdedce415d9..8f49ed471fbf7 100644
--- a/mlir/lib/Dialect/MemRef/Transforms/RuntimeOpVerification.cpp
+++ b/mlir/lib/Dialect/MemRef/Transforms/RuntimeOpVerification.cpp
@@ -362,37 +362,27 @@ struct ExpandShapeOpInterface
generateErrorMessage) const {
auto expandShapeOp = cast<ExpandShapeOp>(op);
- // Verify that the expanded dim sizes are a product of the collapsed dim
- // size.
+ SmallVector<OpFoldResult> outputShape = expandShapeOp.getMixedOutputShape();
+
+ // Verify that the product of output dim sizes in each reassociation group
+ // equals the corresponding input dim size.
for (const auto &it :
llvm::enumerate(expandShapeOp.getReassociationIndices())) {
Value srcDimSz =
DimOp::create(builder, loc, expandShapeOp.getSrc(), it.index());
- int64_t groupSz = 1;
- bool foundDynamicDim = false;
- for (int64_t resultDim : it.value()) {
- if (expandShapeOp.getResultType().isDynamicDim(resultDim)) {
- // Keep this assert here in case the op is extended in the future.
- assert(!foundDynamicDim &&
- "more than one dynamic dim found in reassoc group");
- (void)foundDynamicDim;
- foundDynamicDim = true;
- continue;
- }
- groupSz *= expandShapeOp.getResultType().getDimSize(resultDim);
+ Value groupProduct = getValueOrCreateConstantIndexOp(
+ builder, loc, outputShape[it.value().front()]);
+ for (int64_t resultDim : llvm::drop_begin(it.value())) {
+ Value dimSz = getValueOrCreateConstantIndexOp(builder, loc,
+ outputShape[resultDim]);
+ groupProduct = arith::MulIOp::create(builder, loc, groupProduct, dimSz);
}
- Value staticResultDimSz =
- arith::ConstantIndexOp::create(builder, loc, groupSz);
- // staticResultDimSz must divide srcDimSz evenly.
- Value mod =
- arith::RemSIOp::create(builder, loc, srcDimSz, staticResultDimSz);
- Value isModZero = arith::CmpIOp::create(
- builder, loc, arith::CmpIPredicate::eq, mod,
- arith::ConstantIndexOp::create(builder, loc, 0));
+ Value isEqual = arith::CmpIOp::create(
+ builder, loc, arith::CmpIPredicate::eq, groupProduct, srcDimSz);
cf::AssertOp::create(
- builder, loc, isModZero,
- generateErrorMessage(op, "static result dims in reassoc group do not "
- "divide src dim evenly"));
+ builder, loc, isEqual,
+ generateErrorMessage(op, "product of output dims in reassoc group "
+ "does not equal input dim"));
}
}
};
diff --git a/mlir/test/Dialect/MemRef/runtime-verification.mlir b/mlir/test/Dialect/MemRef/runtime-verification.mlir
index 2198f0032aff5..f449ac6e3b30c 100644
--- a/mlir/test/Dialect/MemRef/runtime-verification.mlir
+++ b/mlir/test/Dialect/MemRef/runtime-verification.mlir
@@ -2,12 +2,12 @@
// CHECK-LABEL: func @expand_shape(
// CHECK-SAME: %[[m:.*]]: memref<?xf32>
-// CHECK-SAME: %[[sz0:.*]]: index
-// CHECK-DAG: %[[c0:.*]] = arith.constant 0 : index
-// CHECK-DAG: %[[c5:.*]] = arith.constant 5 : index
-// CHECK-DAG: %[[dim:.*]] = memref.dim %[[m]], %[[c0]]
-// CHECK: %[[mod:.*]] = arith.remsi %[[dim]], %[[c5]]
-// CHECK: %[[cmpi:.*]] = arith.cmpi eq, %[[mod]], %[[c0]]
+// CHECK-SAME: %[[sz0:.*]]: index
+// CHECK: %[[c0:.*]] = arith.constant 0 : index
+// CHECK: %[[dim:.*]] = memref.dim %[[m]], %[[c0]]
+// CHECK: %[[c5:.*]] = arith.constant 5 : index
+// CHECK: %[[prod:.*]] = arith.muli %[[sz0]], %[[c5]]
+// CHECK: %[[cmpi:.*]] = arith.cmpi eq, %[[prod]], %[[dim]]
// CHECK: cf.assert %[[cmpi]], "ERROR: Runtime op verification failed
func.func @expand_shape(%m: memref<?xf32>, %sz0: index) -> memref<?x5xf32> {
%0 = memref.expand_shape %m [[0, 1]] output_shape [%sz0, 5] : memref<?xf32> into memref<?x5xf32>
diff --git a/mlir/test/Integration/Dialect/MemRef/expand-shape-runtime-verification.mlir b/mlir/test/Integration/Dialect/MemRef/expand-shape-runtime-verification.mlir
new file mode 100644
index 0000000000000..8ef8ade1808b3
--- /dev/null
+++ b/mlir/test/Integration/Dialect/MemRef/expand-shape-runtime-verification.mlir
@@ -0,0 +1,86 @@
+// RUN: mlir-opt %s -generate-runtime-verification \
+// RUN: -lower-affine \
+// RUN: -test-cf-assert \
+// RUN: -convert-scf-to-cf \
+// RUN: -convert-to-llvm | \
+// RUN: mlir-runner -e main -entry-point-result=void \
+// RUN: -shared-libs=%mlir_runner_utils 2>&1 | \
+// RUN: FileCheck %s
+
+// RUN: mlir-opt %s -generate-runtime-verification \
+// RUN: -lower-affine \
+// RUN: -test-cf-assert \
+// RUN: -convert-scf-to-cf \
+// RUN: -convert-to-llvm="allow-pattern-rollback=0" \
+// RUN: -reconcile-unrealized-casts | \
+// RUN: mlir-runner -e main -entry-point-result=void \
+// RUN: -shared-libs=%mlir_runner_utils 2>&1 | \
+// RUN: FileCheck %s
+
+// Dynamic source into dynamic target.
+func.func @expand_shape_dynamic(%m: memref<?xf32>, %sz0: index, %sz1: index) {
+ %0 = memref.expand_shape %m [[0, 1]] output_shape [%sz0, %sz1]
+ : memref<?xf32> into memref<?x?xf32>
+ return
+}
+
+// Dynamic source into mixed static/dynamic target.
+func.func @expand_shape_mixed(%m: memref<?xf32>, %sz0: index) {
+ %0 = memref.expand_shape %m [[0, 1]] output_shape [%sz0, 5]
+ : memref<?xf32> into memref<?x5xf32>
+ return
+}
+
+// Multiple reassociation groups: 2D -> 3D.
+func.func @expand_shape_multi_group(%m: memref<?x?xf32>, %sz0: index, %sz1: index) {
+ %0 = memref.expand_shape %m [[0], [1, 2]] output_shape [%sz0, %sz1, 4]
+ : memref<?x?xf32> into memref<?x?x4xf32>
+ return
+}
+
+func.func @main() {
+ %alloca_10 = memref.alloca() : memref<10xf32>
+ %alloca_10_dyn = memref.cast %alloca_10 : memref<10xf32> to memref<?xf32>
+
+ %alloca_3x20 = memref.alloca() : memref<3x20xf32>
+ %alloca_3x20_dyn = memref.cast %alloca_3x20 : memref<3x20xf32> to memref<?x?xf32>
+
+ %2 = arith.constant 2 : index
+ %3 = arith.constant 3 : index
+ %4 = arith.constant 4 : index
+ %5 = arith.constant 5 : index
+
+ // Product 3*5=15 does not equal input dim 10.
+ // CHECK: ERROR: Runtime op verification failed
+ // CHECK-NEXT: memref.expand_shape %{{.*}} {{\[}}[0, 1]{{\]}} output_shape [%{{.*}}, %{{.*}}] : memref<?xf32> into memref<?x?xf32>
+ // CHECK-NEXT: ^ product of output dims in reassoc group does not equal input dim
+ // CHECK-NEXT: Location: loc({{.*}})
+ func.call @expand_shape_dynamic(%alloca_10_dyn, %3, %5)
+ : (memref<?xf32>, index, index) -> ()
+
+ // Product 2*5=10 equals input dim 10. No error.
+ // CHECK-NOT: ERROR: Runtime op verification failed
+ func.call @expand_shape_dynamic(%alloca_10_dyn, %2, %5)
+ : (memref<?xf32>, index, index) -> ()
+
+ // Product 4*5=20 does not equal input dim 10.
+ // CHECK: ERROR: Runtime op verification failed
+ // CHECK-NEXT: memref.expand_shape %{{.*}} {{\[}}[0, 1]{{\]}} output_shape [%{{.*}}, 5] : memref<?xf32> into memref<?x5xf32>
+ // CHECK-NEXT: ^ product of output dims in reassoc group does not equal input dim
+ // CHECK-NEXT: Location: loc({{.*}})
+ func.call @expand_shape_mixed(%alloca_10_dyn, %4)
+ : (memref<?xf32>, index) -> ()
+
+ // Product 2*5=10 equals input dim 10. No error.
+ // CHECK-NOT: ERROR: Runtime op verification failed
+ func.call @expand_shape_mixed(%alloca_10_dyn, %2)
+ : (memref<?xf32>, index) -> ()
+
+ // Group 0: dim 0 -> [dim 0] product 3, input dim 0 = 3 -> OK.
+ // Group 1: dim 1 -> [dim 1, dim 2] product 5*4=20, input dim 1 = 20 -> OK.
+ // CHECK-NOT: ERROR: Runtime op verification failed
+ func.call @expand_shape_multi_group(%alloca_3x20_dyn, %3, %5)
+ : (memref<?x?xf32>, index, index) -> ()
+
+ return
+}
``````````
</details>
https://github.com/llvm/llvm-project/pull/206125
More information about the Mlir-commits
mailing list