[Mlir-commits] [mlir] [MLIR] Fix mlir-opt crash in ReshapeOpsUtils.cpp when collapse_shape index is invalid (PR #173791)

[Arjun Parmar]

https://github.com/akparmar004 updated https://github.com/llvm/llvm-project/pull/173791

>From 113138b864db8d0075f805cf437499a8d5be0162 Mon Sep 17 00:00:00 2001
From: Bazinga! <akparmar004>
Date: Mon, 29 Dec 2025 00:04:34 +0530
Subject: [PATCH 1/2] [mlir][tensor] Fix crash on empty collapse_shape
 reassociation

---
 mlir/lib/Dialect/Tensor/IR/TensorOps.cpp            |  5 +++++
 mlir/lib/Dialect/Utils/ReshapeOpsUtils.cpp          |  4 +++-
 .../test/Dialect/Tensor/invalid-collapse-shape.mlir | 13 +++++++++++++
 3 files changed, 21 insertions(+), 1 deletion(-)
 create mode 100644 mlir/test/Dialect/Tensor/invalid-collapse-shape.mlir

diff --git a/mlir/lib/Dialect/Tensor/IR/TensorOps.cpp b/mlir/lib/Dialect/Tensor/IR/TensorOps.cpp
index 204e9bb73e12c..466791ee41761 100644
--- a/mlir/lib/Dialect/Tensor/IR/TensorOps.cpp
+++ b/mlir/lib/Dialect/Tensor/IR/TensorOps.cpp
@@ -2072,6 +2072,11 @@ LogicalResult ExpandShapeOp::verify() {
 }
 
 LogicalResult CollapseShapeOp::verify() {
+  CollapseShapeOp op = *this;
+  if (llvm::any_of(op.getReassociationIndices(),
+                   [](auto &group) { return group.empty(); })) {
+    return op.emitOpError("reassociation indices must not be empty");
+  }
   return verifyTensorReshapeOp(*this, getSrcType(), getResultType());
 }
 
diff --git a/mlir/lib/Dialect/Utils/ReshapeOpsUtils.cpp b/mlir/lib/Dialect/Utils/ReshapeOpsUtils.cpp
index 6e9118e1f7b0b..52f1004ccb6a2 100644
--- a/mlir/lib/Dialect/Utils/ReshapeOpsUtils.cpp
+++ b/mlir/lib/Dialect/Utils/ReshapeOpsUtils.cpp
@@ -451,7 +451,9 @@ mlir::getSymbolLessAffineMaps(ArrayRef<ReassociationExprs> reassociation) {
   SmallVector<AffineMap, 4> maps;
   maps.reserve(reassociation.size());
   for (const auto &exprs : reassociation) {
-    assert(!exprs.empty());
+    if (exprs.empty()) {
+      return {};
+    }
     maps.push_back(AffineMap::get(maxDim + 1, 0, exprs, exprs[0].getContext()));
   }
   return maps;
diff --git a/mlir/test/Dialect/Tensor/invalid-collapse-shape.mlir b/mlir/test/Dialect/Tensor/invalid-collapse-shape.mlir
new file mode 100644
index 0000000000000..920a7496abb9b
--- /dev/null
+++ b/mlir/test/Dialect/Tensor/invalid-collapse-shape.mlir
@@ -0,0 +1,13 @@
+// RUN: mlir-opt %s -split-input-file -verify-diagnostics
+
+// This test checks that an empty reassociation group in `tensor.collapse_shape`
+// results in a proper error instead of an assert/crash.
+
+// -----
+
+func.func @test_empty_reassociation(%arg0: tensor<1x?xf32>) -> tensor<?x10xf32> {
+  // expected-error at +1 {{'tensor.collapse_shape' op reassociation indices must not be empty}}
+  %0 = tensor.collapse_shape %arg0 [[0, 1], []] : tensor<1x?xf32> into tensor<?x10xf32>
+  return %0 : tensor<?x10xf32>
+}
+

>From 2a5edb4ade86178f47f4e243f880e563d4e82093 Mon Sep 17 00:00:00 2001
From: Bazinga! <akparmar004>
Date: Tue, 30 Dec 2025 00:12:57 +0530
Subject: [PATCH 2/2] Address review feedback: replace auto with explicit type
 and restore assertion

---
 mlir/lib/Dialect/Tensor/IR/TensorOps.cpp   | 2 +-
 mlir/lib/Dialect/Utils/ReshapeOpsUtils.cpp | 4 +---
 mlir/test/Dialect/Tensor/invalid.mlir      | 9 +++++++++
 3 files changed, 11 insertions(+), 4 deletions(-)

diff --git a/mlir/lib/Dialect/Tensor/IR/TensorOps.cpp b/mlir/lib/Dialect/Tensor/IR/TensorOps.cpp
index 466791ee41761..bed30d29db047 100644
--- a/mlir/lib/Dialect/Tensor/IR/TensorOps.cpp
+++ b/mlir/lib/Dialect/Tensor/IR/TensorOps.cpp
@@ -2074,7 +2074,7 @@ LogicalResult ExpandShapeOp::verify() {
 LogicalResult CollapseShapeOp::verify() {
   CollapseShapeOp op = *this;
   if (llvm::any_of(op.getReassociationIndices(),
-                   [](auto &group) { return group.empty(); })) {
+                   [](ReassociationIndices group) { return group.empty(); })) {
     return op.emitOpError("reassociation indices must not be empty");
   }
   return verifyTensorReshapeOp(*this, getSrcType(), getResultType());
diff --git a/mlir/lib/Dialect/Utils/ReshapeOpsUtils.cpp b/mlir/lib/Dialect/Utils/ReshapeOpsUtils.cpp
index 52f1004ccb6a2..6e9118e1f7b0b 100644
--- a/mlir/lib/Dialect/Utils/ReshapeOpsUtils.cpp
+++ b/mlir/lib/Dialect/Utils/ReshapeOpsUtils.cpp
@@ -451,9 +451,7 @@ mlir::getSymbolLessAffineMaps(ArrayRef<ReassociationExprs> reassociation) {
   SmallVector<AffineMap, 4> maps;
   maps.reserve(reassociation.size());
   for (const auto &exprs : reassociation) {
-    if (exprs.empty()) {
-      return {};
-    }
+    assert(!exprs.empty());
     maps.push_back(AffineMap::get(maxDim + 1, 0, exprs, exprs[0].getContext()));
   }
   return maps;
diff --git a/mlir/test/Dialect/Tensor/invalid.mlir b/mlir/test/Dialect/Tensor/invalid.mlir
index 665657a67dc61..f36678c3d7589 100644
--- a/mlir/test/Dialect/Tensor/invalid.mlir
+++ b/mlir/test/Dialect/Tensor/invalid.mlir
@@ -681,3 +681,12 @@ func.func @bitcast_index_1(%arg0 : tensor<?xindex>) -> tensor<?xi64> {
   %0 = tensor.bitcast %arg0 : tensor<?xindex> to tensor<?xi64>
   return %0 : tensor<?xi64>
 }
+
+// -----
+
+func.func @test_empty_reassociation(%arg0: tensor<1x?xf32>) -> tensor<?x10xf32> {
+  // expected-error at below {{'tensor.collapse_shape' op reassociation indices must not be empty}}
+  %0 = tensor.collapse_shape %arg0 [[0, 1], []] : tensor<1x?xf32> into tensor<?x10xf32>
+  return %0 : tensor<?x10xf32>
+}
+