[Mlir-commits] [mlir] [mlir][Tensor] Reject empty expand_shape reassociation groups (PR #222648)

[llvmlistbot at llvm.org]

https://github.com/AkunCube updated https://github.com/llvm/llvm-project/pull/222648

>From d150fba515bb9dd2ab07fdfd6e52a33b8fc6e4c3 Mon Sep 17 00:00:00 2001
From: cubehan3 <cubehan3 at gmail.com>
Date: Thu, 10 Sep 2026 13:21:51 +0000
Subject: [PATCH] [mlir][Tensor] Reject empty expand_shape reassociation groups

---
 mlir/include/mlir/Dialect/Utils/ReshapeOpsUtils.h | 12 ++++++++++++
 mlir/lib/Dialect/MemRef/IR/MemRefOps.cpp          | 12 +-----------
 mlir/lib/Dialect/Tensor/IR/TensorOps.cpp          | 10 ++++++----
 mlir/test/Dialect/Tensor/invalid.mlir             |  8 ++++++++
 4 files changed, 27 insertions(+), 15 deletions(-)

diff --git a/mlir/include/mlir/Dialect/Utils/ReshapeOpsUtils.h b/mlir/include/mlir/Dialect/Utils/ReshapeOpsUtils.h
index 6f58ed48fd64f..4df9a05da3a91 100644
--- a/mlir/include/mlir/Dialect/Utils/ReshapeOpsUtils.h
+++ b/mlir/include/mlir/Dialect/Utils/ReshapeOpsUtils.h
@@ -19,6 +19,7 @@
 #include "mlir/IR/OpImplementation.h"
 #include "mlir/IR/PatternMatch.h"
 #include "mlir/Support/LLVM.h"
+#include "llvm/ADT/STLExtras.h"
 #include "llvm/ADT/StringRef.h"
 #include <optional>
 
@@ -83,6 +84,17 @@ getReassociationIndicesForCollapse(ArrayRef<int64_t> sourceShape,
 bool isReassociationValid(ArrayRef<AffineMap> reassociation,
                           int *invalidIndex = nullptr);
 
+/// Verify that none of the reassociation groups is empty.
+template <typename ReshapeOpTy>
+LogicalResult verifyReassociationIndicesNotEmpty(ReshapeOpTy op) {
+  if (llvm::any_of(
+          op.getReassociationIndices(),
+          [](const ReassociationIndices &group) { return group.empty(); })) {
+    return op.emitOpError("reassociation indices must not be empty");
+  }
+  return success();
+}
+
 template <typename ReshapeOpTy, typename InverseReshapeOpTy>
 OpFoldResult foldReshapeOp(ReshapeOpTy reshapeOp,
                            ArrayRef<Attribute> operands) {
diff --git a/mlir/lib/Dialect/MemRef/IR/MemRefOps.cpp b/mlir/lib/Dialect/MemRef/IR/MemRefOps.cpp
index e4bfe20f8307d..217e5ce841292 100644
--- a/mlir/lib/Dialect/MemRef/IR/MemRefOps.cpp
+++ b/mlir/lib/Dialect/MemRef/IR/MemRefOps.cpp
@@ -9,6 +9,7 @@
 #include "mlir/Dialect/Arith/IR/Arith.h"
 #include "mlir/Dialect/Arith/Utils/Utils.h"
 #include "mlir/Dialect/MemRef/IR/MemRef.h"
+#include "mlir/Dialect/Utils/ReshapeOpsUtils.h"
 #include "mlir/Dialect/Utils/StaticValueUtils.h"
 #include "mlir/Dialect/Utils/VerificationUtils.h"
 #include "mlir/IR/AffineMap.h"
@@ -2637,17 +2638,6 @@ void ExpandShapeOp::build(OpBuilder &builder, OperationState &result,
   build(builder, result, *resultType, src, reassociation, outputShape);
 }
 
-/// Verify that none of the reassociation groups is empty.
-template <typename MemrefReshapeOp>
-static LogicalResult verifyReassociationIndicesNotEmpty(MemrefReshapeOp op) {
-  if (llvm::any_of(
-          op.getReassociationIndices(),
-          [](const ReassociationIndices &group) { return group.empty(); })) {
-    return op.emitOpError("reassociation indices must not be empty");
-  }
-  return success();
-}
-
 LogicalResult ExpandShapeOp::verify() {
   if (failed(verifyReassociationIndicesNotEmpty(*this)))
     return failure();
diff --git a/mlir/lib/Dialect/Tensor/IR/TensorOps.cpp b/mlir/lib/Dialect/Tensor/IR/TensorOps.cpp
index 43bb5cabb0f4f..53ee57fa72a17 100644
--- a/mlir/lib/Dialect/Tensor/IR/TensorOps.cpp
+++ b/mlir/lib/Dialect/Tensor/IR/TensorOps.cpp
@@ -2092,6 +2092,9 @@ static LogicalResult verifyTensorReshapeOp(TensorReshapeOp op,
 }
 
 LogicalResult ExpandShapeOp::verify() {
+  if (failed(verifyReassociationIndicesNotEmpty(*this)))
+    return failure();
+
   RankedTensorType srcType = getSrc().getType();
   RankedTensorType resultType = getResult().getType();
 
@@ -2127,10 +2130,9 @@ LogicalResult ExpandShapeOp::verify() {
 
 LogicalResult CollapseShapeOp::verify() {
   CollapseShapeOp op = *this;
-  if (llvm::any_of(op.getReassociationIndices(),
-                   [](ReassociationIndices group) { return group.empty(); })) {
-    return op.emitOpError("reassociation indices must not be empty");
-  }
+  if (failed(verifyReassociationIndicesNotEmpty(op)))
+    return failure();
+
   RankedTensorType srcType = op.getSrc().getType();
   RankedTensorType resultType = op.getResult().getType();
 
diff --git a/mlir/test/Dialect/Tensor/invalid.mlir b/mlir/test/Dialect/Tensor/invalid.mlir
index a526d7ed61722..19db4680b6a63 100644
--- a/mlir/test/Dialect/Tensor/invalid.mlir
+++ b/mlir/test/Dialect/Tensor/invalid.mlir
@@ -699,6 +699,14 @@ func.func @test_empty_reassociation(%arg0: tensor<1x?xf32>) -> tensor<?x10xf32>
 
 // -----
 
+func.func @expand_shape_empty_reassociation(%arg0: tensor<4xf32>) -> tensor<4xf32> {
+  // expected-error at below {{'tensor.expand_shape' op reassociation indices must not be empty}}
+  %0 = tensor.expand_shape %arg0 [[]] output_shape [4] : tensor<4xf32> into tensor<4xf32>
+  return %0 : tensor<4xf32>
+}
+
+// -----
+
 func.func @collapse_shape_requires_ranked_tensor(%arg0: tensor<*xf32>) {
   // expected-error at +1 {{custom op 'tensor.collapse_shape' invalid kind of type specified: expected builtin.tensor, but found 'tensor<*xf32>'}}
   %0 = tensor.collapse_shape %arg0 [[0]] : tensor<*xf32> into tensor<f32>