[MLIR] Fix mlir-opt crash in ReshapeOpsUtils.cpp when collapse_shape index is invalid #173791
Conversation
|
Thank you for submitting a Pull Request (PR) to the LLVM Project! This PR will be automatically labeled and the relevant teams will be notified. If you wish to, you can add reviewers by using the "Reviewers" section on this page. If this is not working for you, it is probably because you do not have write permissions for the repository. In which case you can instead tag reviewers by name in a comment by using If you have received no comments on your PR for a week, you can request a review by "ping"ing the PR by adding a comment “Ping”. The common courtesy "ping" rate is once a week. Please remember that you are asking for valuable time from other developers. If you have further questions, they may be answered by the LLVM GitHub User Guide. You can also ask questions in a comment on this PR, on the LLVM Discord or on the forums. |
|
@llvm/pr-subscribers-mlir-tensor Author: Arjun Parmar (akparmar004) ChangesThis 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. Full diff: https://github.com/llvm/llvm-project/pull/173791.diff 3 Files Affected:
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@+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>
+}
+
|
|
@llvm/pr-subscribers-mlir Author: Arjun Parmar (akparmar004) ChangesThis 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. Full diff: https://github.com/llvm/llvm-project/pull/173791.diff 3 Files Affected:
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@+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>
+}
+
|
2 participants
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