[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 created https://github.com/llvm/llvm-project/pull/173791

This patch fixes a crash occurring in mlir-opt when running collapse_shape with an invalid index configuration. Instead of crashing, an error message is returned to the user.
Fixes: #173567


>From 113138b864db8d0075f805cf437499a8d5be0162 Mon Sep 17 00:00:00 2001
From: Bazinga! <akparmar004>
Date: Mon, 29 Dec 2025 00:04:34 +0530
Subject: [PATCH] [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>
+}
+