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

[llvmlistbot at llvm.org]

llvmorg-github-actions[bot] wrote:


<!--LLVM PR SUMMARY COMMENT-->

@llvm/pr-subscribers-mlir-tensor

Author: Akun (AkunCube)

<details>
<summary>Changes</summary>

`tensor.expand_shape` with an empty reassociation group (e.g. `[[]]`) used to abort the compiler: ExpandShapeOp::verify forwarded the group to getSymbolLessAffineMaps, which asserts on missing affine expressions. The sibling `tensor.collapse_shape` already rejects this input with a clean diagnostic.

Extract the empty-group check into a helper shared by both tensor reshape verifiers, invoke it from `ExpandShapeOp::verify`, and add an invalid.mlir regression test that used to SIGABRT.

Fixes: #<!-- -->222602

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


2 Files Affected:

- (modified) mlir/lib/Dialect/Tensor/IR/TensorOps.cpp (+17-4) 
- (modified) mlir/test/Dialect/Tensor/invalid.mlir (+8) 


``````````diff
diff --git a/mlir/lib/Dialect/Tensor/IR/TensorOps.cpp b/mlir/lib/Dialect/Tensor/IR/TensorOps.cpp
index c4f85bea14b49..3a7d01e7d5f9d 100644
--- a/mlir/lib/Dialect/Tensor/IR/TensorOps.cpp
+++ b/mlir/lib/Dialect/Tensor/IR/TensorOps.cpp
@@ -2089,7 +2089,21 @@ static LogicalResult verifyTensorReshapeOp(TensorReshapeOp op,
   return success();
 }
 
+/// Verify that none of the reassociation groups is empty.
+template <typename TensorReshapeOp>
+static LogicalResult verifyReassociationIndicesNotEmpty(TensorReshapeOp 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();
+
   RankedTensorType srcType = getSrc().getType();
   RankedTensorType resultType = getResult().getType();
 
@@ -2125,10 +2139,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>

``````````

</details>


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