[Mlir-commits] [mlir] [mlir] Make [tensor|memref]::ExpandShapeOp verifier stricter. (PR #181020)
Han-Chung Wang
llvmlistbot at llvm.org
Sun Feb 22 05:34:07 PST 2026
https://github.com/hanhanW updated https://github.com/llvm/llvm-project/pull/181020
>From e5b2fa852c475db16fc530dc7e596d35a979595a Mon Sep 17 00:00:00 2001
From: hanhanW <hanhan0912 at gmail.com>
Date: Wed, 11 Feb 2026 13:11:09 -0800
Subject: [PATCH 1/4] [mlir] Make [tensor|memref]::ExpandShapeOp verifier
stricker.
The number of dynamic dims in output_shape is expected to be as the same
as the result type.
Signed-off-by: hanhanW <hanhan0912 at gmail.com>
---
mlir/include/mlir/Dialect/MemRef/IR/MemRefOps.td | 6 ++++--
mlir/include/mlir/Dialect/Tensor/IR/TensorOps.td | 10 ++++++----
mlir/lib/Dialect/MemRef/IR/MemRefOps.cpp | 6 ++++++
mlir/lib/Dialect/Tensor/IR/TensorOps.cpp | 6 ++++++
mlir/test/Dialect/MemRef/invalid.mlir | 8 ++++++++
mlir/test/Dialect/Tensor/invalid.mlir | 7 +++++++
6 files changed, 37 insertions(+), 6 deletions(-)
diff --git a/mlir/include/mlir/Dialect/MemRef/IR/MemRefOps.td b/mlir/include/mlir/Dialect/MemRef/IR/MemRefOps.td
index 6f8f1481725fc..6453c11e475ae 100644
--- a/mlir/include/mlir/Dialect/MemRef/IR/MemRefOps.td
+++ b/mlir/include/mlir/Dialect/MemRef/IR/MemRefOps.td
@@ -1824,10 +1824,12 @@ def MemRef_ExpandShapeOp : MemRef_ReassociativeReshapeOp<"expand_shape", [
The representation for the output shape supports a partially-static
specification via attributes specified through the `static_output_shape`
- argument. A special sentinel value `ShapedType::kDynamic` encodes that the
+ argument. A special sentinel value `ShapedType::kDynamic` encodes that the
corresponding entry has a dynamic value. There must be exactly as many SSA
inputs in `output_shape` as there are `ShapedType::kDynamic` entries in
- `static_output_shape`.
+ `static_output_shape`. Additionally, the number of `ShapedType::kDynamic`
+ entries in `static_output_shape` must match the number of dynamic dimensions
+ in the result type.
Note: This op currently assumes that the inner strides are of the
source/result layout map are the faster-varying ones.
diff --git a/mlir/include/mlir/Dialect/Tensor/IR/TensorOps.td b/mlir/include/mlir/Dialect/Tensor/IR/TensorOps.td
index 455b5541d550f..57c1ea63fd1bc 100644
--- a/mlir/include/mlir/Dialect/Tensor/IR/TensorOps.td
+++ b/mlir/include/mlir/Dialect/Tensor/IR/TensorOps.td
@@ -1113,16 +1113,18 @@ def Tensor_ExpandShapeOp : Tensor_ReassociativeReshapeOp<"expand_shape"> {
`src`.
A reassociation is defined as a continuous grouping of dimensions and is
- represented with an array of DenseI64ArrayAttr attribute. The reassociation
+ represented with an array of DenseI64ArrayAttr attribute. The reassociation
maps applied to the result tensor with the higher rank must result in the
operand tensor with the smaller rank.
The representation for the output shape supports a partially-static
specification via attributes specified through the `static_output_shape`
- argument. A special sentinel value `ShapedType::kDynamic` encodes that the
- corresponding entry has a dynamic value. There must be exactly as many SSA
+ argument. A special sentinel value `ShapedType::kDynamic` encodes that the
+ corresponding entry has a dynamic value. There must be exactly as many SSA
inputs in `output_shape` as there are `ShapedType::kDynamic` entries in
- `static_output_shape`.
+ `static_output_shape`. Additionally, the number of `ShapedType::kDynamic`
+ entries in `static_output_shape` must match the number of dynamic dimensions
+ in the result type.
Example:
diff --git a/mlir/lib/Dialect/MemRef/IR/MemRefOps.cpp b/mlir/lib/Dialect/MemRef/IR/MemRefOps.cpp
index 4ac8505c1223a..7c3d9faf0af78 100644
--- a/mlir/lib/Dialect/MemRef/IR/MemRefOps.cpp
+++ b/mlir/lib/Dialect/MemRef/IR/MemRefOps.cpp
@@ -2522,6 +2522,12 @@ LogicalResult ExpandShapeOp::verify() {
<< " dynamic dims while output_shape has " << getOutputShape().size()
<< " values";
+ // Verify that the number of dynamic dims in output_shape matches the number
+ // of dynamic dims in the result type.
+ if (failed(verifyDynamicDimensionCount(getOperation(), resultType,
+ getOutputShape())))
+ return failure();
+
// Verify if provided output shapes are in agreement with output type.
DenseI64ArrayAttr staticOutputShapes = getStaticOutputShapeAttr();
ArrayRef<int64_t> resShape = getResult().getType().getShape();
diff --git a/mlir/lib/Dialect/Tensor/IR/TensorOps.cpp b/mlir/lib/Dialect/Tensor/IR/TensorOps.cpp
index fc42aa90a257d..170172e07df76 100644
--- a/mlir/lib/Dialect/Tensor/IR/TensorOps.cpp
+++ b/mlir/lib/Dialect/Tensor/IR/TensorOps.cpp
@@ -2068,6 +2068,12 @@ LogicalResult ExpandShapeOp::verify() {
<< " dynamic dims while output_shape has " << getOutputShape().size()
<< " values";
+ // Verify that the number of dynamic dims in output_shape matches the number
+ // of dynamic dims in the result type.
+ if (failed(verifyDynamicDimensionCount(getOperation(), resultType,
+ getOutputShape())))
+ return failure();
+
return verifyTensorReshapeOp(*this, resultType, srcType);
}
diff --git a/mlir/test/Dialect/MemRef/invalid.mlir b/mlir/test/Dialect/MemRef/invalid.mlir
index 46e010fc878fe..af068d8ca8e95 100644
--- a/mlir/test/Dialect/MemRef/invalid.mlir
+++ b/mlir/test/Dialect/MemRef/invalid.mlir
@@ -425,6 +425,14 @@ func.func @expand_shape_illegal_output_shape(%arg0: memref<2xf32>) {
// -----
+func.func @expand_shape_output_shape_dynamic_dim_mismatch(%arg0: memref<?xf32>) {
+ // expected-error @+1 {{incorrect number of dynamic sizes, has 0, expected 2}}
+ %0 = memref.expand_shape %arg0 [[0, 1]] output_shape [2, 3] : memref<?xf32> into memref<?x?xf32>
+ return
+}
+
+// -----
+
func.func @collapse_shape_out_of_bounds(%arg0: memref<?x?xf32>) {
// expected-error @+1 {{op reassociation index 2 is out of bounds}}
%0 = memref.collapse_shape %arg0 [[0, 1, 2]] : memref<?x?xf32> into memref<?xf32>
diff --git a/mlir/test/Dialect/Tensor/invalid.mlir b/mlir/test/Dialect/Tensor/invalid.mlir
index c149c39f99dce..26e2143c80256 100644
--- a/mlir/test/Dialect/Tensor/invalid.mlir
+++ b/mlir/test/Dialect/Tensor/invalid.mlir
@@ -374,6 +374,13 @@ func.func @expand_shape_illegal_output_shape(%arg0: tensor<2xf32>) {
return
}
+// -----
+
+func.func @expand_shape_output_shape_dynamic_dim_mismatch(%arg0: tensor<6xf32>) {
+ // expected-error @+1 {{incorrect number of dynamic sizes, has 0, expected 2}}
+ %0 = tensor.expand_shape %arg0 [[0, 1]] output_shape [2, 3] : tensor<6xf32> into tensor<?x?xf32>
+ return
+}
// -----
>From 9bcbc7b28ce0e3e464a961170c1d317ff07d123e Mon Sep 17 00:00:00 2001
From: hanhanW <hanhan0912 at gmail.com>
Date: Wed, 11 Feb 2026 15:50:04 -0800
Subject: [PATCH 2/4] Check shapes and fix invalid ops.
Signed-off-by: hanhanW <hanhan0912 at gmail.com>
---
mlir/lib/Dialect/Tensor/IR/TensorOps.cpp | 7 +++++++
mlir/test/Dialect/Linalg/data-layout-propagation.mlir | 2 +-
mlir/test/Dialect/Tensor/canonicalize.mlir | 2 +-
3 files changed, 9 insertions(+), 2 deletions(-)
diff --git a/mlir/lib/Dialect/Tensor/IR/TensorOps.cpp b/mlir/lib/Dialect/Tensor/IR/TensorOps.cpp
index 170172e07df76..89564843e7a3c 100644
--- a/mlir/lib/Dialect/Tensor/IR/TensorOps.cpp
+++ b/mlir/lib/Dialect/Tensor/IR/TensorOps.cpp
@@ -2074,6 +2074,13 @@ LogicalResult ExpandShapeOp::verify() {
getOutputShape())))
return failure();
+ // Verify if provided output shapes are in agreement with output type.
+ DenseI64ArrayAttr staticOutputShapes = getStaticOutputShapeAttr();
+ ArrayRef<int64_t> resShape = getResult().getType().getShape();
+ for (auto [pos, shape] : llvm::enumerate(resShape))
+ if (ShapedType::isStatic(shape) && shape != staticOutputShapes[pos])
+ return emitOpError("invalid output shape provided at pos ") << pos;
+
return verifyTensorReshapeOp(*this, resultType, srcType);
}
diff --git a/mlir/test/Dialect/Linalg/data-layout-propagation.mlir b/mlir/test/Dialect/Linalg/data-layout-propagation.mlir
index 7a16bc0a4faee..ec34a02096d5f 100644
--- a/mlir/test/Dialect/Linalg/data-layout-propagation.mlir
+++ b/mlir/test/Dialect/Linalg/data-layout-propagation.mlir
@@ -1132,7 +1132,7 @@ func.func @bubble_up_pack_through_expand_dynamic(%arg0: tensor<?x64xf32>) -> ten
func.func @bubble_up_pack_non_expanded_padding_through_expand(%arg0: tensor<32x60xf32>) -> tensor<4x2x8x4x8xf32> {
%cst = arith.constant 3.000000e+00 : f32
%empty = tensor.empty() : tensor<4x2x8x4x8xf32>
- %expanded = tensor.expand_shape %arg0 [[0, 1], [2]] output_shape [4, 8, 64] : tensor<32x60xf32> into tensor<4x8x60xf32>
+ %expanded = tensor.expand_shape %arg0 [[0, 1], [2]] output_shape [4, 8, 60] : tensor<32x60xf32> into tensor<4x8x60xf32>
%pack = linalg.pack %expanded padding_value(%cst : f32) inner_dims_pos = [1, 2] inner_tiles = [4, 8] into %empty : tensor<4x8x60xf32> -> tensor<4x2x8x4x8xf32>
return %pack : tensor<4x2x8x4x8xf32>
}
diff --git a/mlir/test/Dialect/Tensor/canonicalize.mlir b/mlir/test/Dialect/Tensor/canonicalize.mlir
index 7a2d53c0c5850..67d1b2bc32b78 100644
--- a/mlir/test/Dialect/Tensor/canonicalize.mlir
+++ b/mlir/test/Dialect/Tensor/canonicalize.mlir
@@ -1351,7 +1351,7 @@ func.func @compose_expand_of_collapse_dynamic(%arg0 : tensor<4x?x10x64x2xf16>, %
func.func @no_compose_collapse_of_expand_dynamic(%arg0 : tensor<?x8x128x?xf16>, %arg1: index) -> tensor<?x128x?xf16> {
%collapse = tensor.collapse_shape %arg0 [[0, 1, 2, 3]] : tensor<?x8x128x?xf16> into tensor<?xf16>
- %expanded_19 = tensor.expand_shape %collapse [[0, 1, 2]] output_shape [%arg1, 8, %arg1] : tensor<?xf16> into tensor<?x128x?xf16>
+ %expanded_19 = tensor.expand_shape %collapse [[0, 1, 2]] output_shape [%arg1, 128, %arg1] : tensor<?xf16> into tensor<?x128x?xf16>
return %expanded_19 : tensor<?x128x?xf16>
}
// CHECK-LABEL: func @no_compose_collapse_of_expand_dynamic
>From f2f4c2267ba75d5a0c9739e450a67c00e6db4c72 Mon Sep 17 00:00:00 2001
From: hanhanW <hanhan0912 at gmail.com>
Date: Wed, 11 Feb 2026 16:12:22 -0800
Subject: [PATCH 3/4] Fix one more file
Signed-off-by: hanhanW <hanhan0912 at gmail.com>
---
mlir/test/Dialect/Tosa/tosa-infer-shapes.mlir | 4 ++--
1 file changed, 2 insertions(+), 2 deletions(-)
diff --git a/mlir/test/Dialect/Tosa/tosa-infer-shapes.mlir b/mlir/test/Dialect/Tosa/tosa-infer-shapes.mlir
index 0deb31dab1c6f..f96beedcf7b10 100644
--- a/mlir/test/Dialect/Tosa/tosa-infer-shapes.mlir
+++ b/mlir/test/Dialect/Tosa/tosa-infer-shapes.mlir
@@ -1578,8 +1578,8 @@ func.func @test_multiple_non_inferrable_consumers(%arg0: tensor<1x2x8xf32>) {
// CHECK: tensor.expand_shape %[[TENSOR_CAST]]
// CHECK: tensor.expand_shape %[[TENSOR_CAST]]
- %expanded_0 = tensor.expand_shape %0 [[0], [1, 2], [3]] output_shape [%dim, 1, 4, 8] : tensor<?x2x8xf32> into tensor<?x1x2x8xf32>
- %expanded_1 = tensor.expand_shape %0 [[0], [1, 2], [3]] output_shape [%dim, 1, 4, 8] : tensor<?x2x8xf32> into tensor<?x1x2x8xf32>
+ %expanded_0 = tensor.expand_shape %0 [[0], [1, 2], [3]] output_shape [%dim, 1, 2, 8] : tensor<?x2x8xf32> into tensor<?x1x2x8xf32>
+ %expanded_1 = tensor.expand_shape %0 [[0], [1, 2], [3]] output_shape [%dim, 1, 2, 8] : tensor<?x2x8xf32> into tensor<?x1x2x8xf32>
return
}
>From f3b564d6fda58b5e3e984e3844b5cc86d345da63 Mon Sep 17 00:00:00 2001
From: hanhanW <hanhan0912 at gmail.com>
Date: Sun, 22 Feb 2026 05:33:45 -0800
Subject: [PATCH 4/4] Improve doc
Signed-off-by: hanhanW <hanhan0912 at gmail.com>
---
mlir/include/mlir/Dialect/MemRef/IR/MemRefOps.td | 9 ++++-----
mlir/include/mlir/Dialect/Tensor/IR/TensorOps.td | 9 ++++-----
2 files changed, 8 insertions(+), 10 deletions(-)
diff --git a/mlir/include/mlir/Dialect/MemRef/IR/MemRefOps.td b/mlir/include/mlir/Dialect/MemRef/IR/MemRefOps.td
index 6453c11e475ae..70180c101407a 100644
--- a/mlir/include/mlir/Dialect/MemRef/IR/MemRefOps.td
+++ b/mlir/include/mlir/Dialect/MemRef/IR/MemRefOps.td
@@ -1825,11 +1825,10 @@ def MemRef_ExpandShapeOp : MemRef_ReassociativeReshapeOp<"expand_shape", [
The representation for the output shape supports a partially-static
specification via attributes specified through the `static_output_shape`
argument. A special sentinel value `ShapedType::kDynamic` encodes that the
- corresponding entry has a dynamic value. There must be exactly as many SSA
- inputs in `output_shape` as there are `ShapedType::kDynamic` entries in
- `static_output_shape`. Additionally, the number of `ShapedType::kDynamic`
- entries in `static_output_shape` must match the number of dynamic dimensions
- in the result type.
+ corresponding entry has a dynamic value. Both the number of SSA inputs in
+ `output_shape` and the number of `ShapedType::kDynamic` entries in
+ `static_output_shape` match the number of dynamic dimensions in the result
+ type.
Note: This op currently assumes that the inner strides are of the
source/result layout map are the faster-varying ones.
diff --git a/mlir/include/mlir/Dialect/Tensor/IR/TensorOps.td b/mlir/include/mlir/Dialect/Tensor/IR/TensorOps.td
index 57c1ea63fd1bc..c9b858519a592 100644
--- a/mlir/include/mlir/Dialect/Tensor/IR/TensorOps.td
+++ b/mlir/include/mlir/Dialect/Tensor/IR/TensorOps.td
@@ -1120,11 +1120,10 @@ def Tensor_ExpandShapeOp : Tensor_ReassociativeReshapeOp<"expand_shape"> {
The representation for the output shape supports a partially-static
specification via attributes specified through the `static_output_shape`
argument. A special sentinel value `ShapedType::kDynamic` encodes that the
- corresponding entry has a dynamic value. There must be exactly as many SSA
- inputs in `output_shape` as there are `ShapedType::kDynamic` entries in
- `static_output_shape`. Additionally, the number of `ShapedType::kDynamic`
- entries in `static_output_shape` must match the number of dynamic dimensions
- in the result type.
+ corresponding entry has a dynamic value. Both the number of SSA inputs in
+ `output_shape` and the number of `ShapedType::kDynamic` entries in
+ `static_output_shape` match the number of dynamic dimensions in the result
+ type.
Example:
More information about the Mlir-commits
mailing list