[Mlir-commits] [mlir] 686acf6 - [mlir] Make [tensor|memref]::ExpandShapeOp verifier stricter. (#181020)
llvmlistbot at llvm.org
llvmlistbot at llvm.org
Sun Feb 22 16:42:47 PST 2026
Author: Han-Chung Wang
Date: 2026-02-22T16:42:43-08:00
New Revision: 686acf63823734b38f55dd8ae6fcd8aa27930f15
URL: https://github.com/llvm/llvm-project/commit/686acf63823734b38f55dd8ae6fcd8aa27930f15
DIFF: https://github.com/llvm/llvm-project/commit/686acf63823734b38f55dd8ae6fcd8aa27930f15.diff
LOG: [mlir] Make [tensor|memref]::ExpandShapeOp verifier stricter. (#181020)
The number of dynamic dims in output_shape is expected to be as the same
as the result type.
The revision also trims double whitespaces from the doc, because it also
updates the op description.
---------
Signed-off-by: hanhanW <hanhan0912 at gmail.com>
Added:
Modified:
mlir/include/mlir/Dialect/MemRef/IR/MemRefOps.td
mlir/include/mlir/Dialect/Tensor/IR/TensorOps.td
mlir/lib/Dialect/MemRef/IR/MemRefOps.cpp
mlir/lib/Dialect/Tensor/IR/TensorOps.cpp
mlir/test/Dialect/Linalg/data-layout-propagation.mlir
mlir/test/Dialect/MemRef/invalid.mlir
mlir/test/Dialect/Tensor/canonicalize.mlir
mlir/test/Dialect/Tensor/invalid.mlir
mlir/test/Dialect/Tosa/tosa-infer-shapes.mlir
Removed:
################################################################################
diff --git a/mlir/include/mlir/Dialect/MemRef/IR/MemRefOps.td b/mlir/include/mlir/Dialect/MemRef/IR/MemRefOps.td
index 6f8f1481725fc..70180c101407a 100644
--- a/mlir/include/mlir/Dialect/MemRef/IR/MemRefOps.td
+++ b/mlir/include/mlir/Dialect/MemRef/IR/MemRefOps.td
@@ -1824,10 +1824,11 @@ 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`.
+ argument. A special sentinel value `ShapedType::kDynamic` encodes that the
+ 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 455b5541d550f..c9b858519a592 100644
--- a/mlir/include/mlir/Dialect/Tensor/IR/TensorOps.td
+++ b/mlir/include/mlir/Dialect/Tensor/IR/TensorOps.td
@@ -1113,16 +1113,17 @@ 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
- inputs in `output_shape` as there are `ShapedType::kDynamic` entries in
- `static_output_shape`.
+ argument. A special sentinel value `ShapedType::kDynamic` encodes that the
+ 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:
diff --git a/mlir/lib/Dialect/MemRef/IR/MemRefOps.cpp b/mlir/lib/Dialect/MemRef/IR/MemRefOps.cpp
index 496cad5219cae..844e6183cff06 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 d837947e0dc3b..4c0ab7c9ec8a0 100644
--- a/mlir/lib/Dialect/Tensor/IR/TensorOps.cpp
+++ b/mlir/lib/Dialect/Tensor/IR/TensorOps.cpp
@@ -2078,6 +2078,19 @@ 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();
+ 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/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/canonicalize.mlir b/mlir/test/Dialect/Tensor/canonicalize.mlir
index e125ea3b62b44..fc499da5422fc 100644
--- a/mlir/test/Dialect/Tensor/canonicalize.mlir
+++ b/mlir/test/Dialect/Tensor/canonicalize.mlir
@@ -1364,7 +1364,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
diff --git a/mlir/test/Dialect/Tensor/invalid.mlir b/mlir/test/Dialect/Tensor/invalid.mlir
index 0483dcaa3c6f0..6ee2f9911663f 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
+}
// -----
diff --git a/mlir/test/Dialect/Tosa/tosa-infer-shapes.mlir b/mlir/test/Dialect/Tosa/tosa-infer-shapes.mlir
index 9213fdd636813..a56c5b1702baf 100644
--- a/mlir/test/Dialect/Tosa/tosa-infer-shapes.mlir
+++ b/mlir/test/Dialect/Tosa/tosa-infer-shapes.mlir
@@ -1542,8 +1542,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
}
More information about the Mlir-commits
mailing list