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

[llvmlistbot at llvm.org]

llvmbot wrote:


<!--LLVM PR SUMMARY COMMENT-->

@llvm/pr-subscribers-mlir

Author: Arjun Parmar (akparmar004)

<details>
<summary>Changes</summary>

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


---
Full diff: https://github.com/llvm/llvm-project/pull/173791.diff


3 Files Affected:

- (modified) mlir/lib/Dialect/Tensor/IR/TensorOps.cpp (+5) 
- (modified) mlir/lib/Dialect/Utils/ReshapeOpsUtils.cpp (+3-1) 
- (added) mlir/test/Dialect/Tensor/invalid-collapse-shape.mlir (+13) 


``````````diff
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>
+}
+

``````````

</details>


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