[Mlir-commits] [mlir] [mlir][memref] Don't fold casts into expand_shape with negative sizes (PR #233174)

[Chang Yong Lik]

https://github.com/ahlag created https://github.com/llvm/llvm-project/pull/233174

Fixes #233050, a fuzzer-discovered crash in `-canonicalize`. The crash occurs because `ExpandShapeOpMemRefCastFolder` folds constant `output_shape` values into static dimensions without validating them. When a dimension folds to a negative value (in this case, `-1`), it constructs an invalid memref type, causing `MemRefType::get` to assert with `invalid memref size`. This change makes the pattern bail out on negative dimensions, consistent with the existing tensor folder in `TensorOps.cpp`. Since negative dimensions are already invalid IR, valid inputs are unaffected. A regression test, `@no_fold_memref_expand_cast_negative_size`, is added to `mlir/test/Dialect/MemRef/canonicalize.mlir`. The test reproduces the crash without the fix and passes with it. All 34 MemRef lit tests pass.


>From bc94db89c9808b297bc266a62b098597f947f6dc Mon Sep 17 00:00:00 2001
From: ahlag <theeahlag at gmail.com>
Date: Sun, 18 Oct 2026 11:48:06 +0900
Subject: [PATCH] [mlir][memref] Don't fold casts into expand_shape with
 negative sizes

ExpandShapeOpMemRefCastFolder turned constant output_shape operands into
static sizes without checking them. A negative constant produced an
invalid MemRefType and hit the verifyInvariants assertion. Bail out on
negative sizes, as other folders do for invalid IR.

Fixes #233050.

Assisted-by: Claude
---
 mlir/lib/Dialect/MemRef/IR/MemRefOps.cpp   |  4 ++++
 mlir/test/Dialect/MemRef/canonicalize.mlir | 14 ++++++++++++++
 2 files changed, 18 insertions(+)

diff --git a/mlir/lib/Dialect/MemRef/IR/MemRefOps.cpp b/mlir/lib/Dialect/MemRef/IR/MemRefOps.cpp
index 245b2ba6babf04..9495a98eab2051 100644
--- a/mlir/lib/Dialect/MemRef/IR/MemRefOps.cpp
+++ b/mlir/lib/Dialect/MemRef/IR/MemRefOps.cpp
@@ -2736,6 +2736,10 @@ struct ExpandShapeOpMemRefCastFolder : public OpRewritePattern<ExpandShapeOp> {
         continue;
       }
 
+      // Bail on invalid IR.
+      if (sizeOpt.value() < 0)
+        return rewriter.notifyMatchFailure(op, "negative output shape size");
+
       newOutputShapeSizes.push_back(sizeOpt.value());
       newOutputShape[dimIdx] = rewriter.getIndexAttr(sizeOpt.value());
     }
diff --git a/mlir/test/Dialect/MemRef/canonicalize.mlir b/mlir/test/Dialect/MemRef/canonicalize.mlir
index b0261c665754a9..04083027ce4134 100644
--- a/mlir/test/Dialect/MemRef/canonicalize.mlir
+++ b/mlir/test/Dialect/MemRef/canonicalize.mlir
@@ -568,6 +568,20 @@ func.func @fold_memref_expand_with_static_to_dynamic_cast(%arg0: memref<8x4xf32>
 
 // -----
 
+// Negative output shape sizes must not be folded into the static type.
+// CHECK-LABEL:   func.func @no_fold_memref_expand_cast_negative_size(
+// CHECK:           %[[CAST:.*]] = memref.cast
+// CHECK:           memref.expand_shape %[[CAST]] {{.*}} : memref<?x4xf32> into memref<?x1x4x4xf32>
+func.func @no_fold_memref_expand_cast_negative_size(%arg0: memref<8x4xf32>) -> memref<?x1x4x4xf32> {
+  %c-1 = arith.constant -1 : index
+  %cast = memref.cast %arg0 : memref<8x4xf32> to memref<?x4xf32>
+  %expand_shape = memref.expand_shape %cast [[0, 1, 2], [3]] output_shape [%c-1, 1, 4, 4]
+      : memref<?x4xf32> into memref<?x1x4x4xf32>
+  return %expand_shape : memref<?x1x4x4xf32>
+}
+
+// -----
+
 // CHECK-LABEL:   func.func @fold_memref_expand_static_to_dynamic_partial(
 // CHECK-SAME:      %[[ARG0:.*]]: memref<8x?xf32>) -> memref<1x8x1x?xf32> {
 // CHECK:           %[[C1:.*]] = arith.constant 1 : index