[Mlir-commits] [mlir] [mlir][memref] Reject empty collapse_shape and expand_shape reassociation groups (PR #225348)

[llvmlistbot at llvm.org]

llvmorg-github-actions[bot] wrote:


<!--LLVM PR SUMMARY COMMENT-->

@llvm/pr-subscribers-mlir

Author: Wijkqwe

<details>
<summary>Changes</summary>

`memref.collapse_shape` and `memref.expand_shape` both accept an empty
reassociation group, which is not a valid reassociation: every group must be a
non-empty, contiguous segment of dimensions.

* `memref.collapse_shape` with an empty group aborts the compiler for a
  non-identity source layout. `CollapseShapeOp::verify` forwards the
  reassociation to `computeCollapsedLayoutMap`, which calls `ArrayRef::back()`
  on each group and asserts on the empty one:

      mlir-opt: llvm/include/llvm/ADT/ArrayRef.h:151:
      const T& llvm::ArrayRef<T>::back() const [with T = long int]:
      Assertion `!empty()' failed.

  With an identity source layout the op is instead silently accepted.

* `memref.expand_shape` with an empty group is silently accepted: the empty
  group maps a source dimension to no result dimension, so a real dimension is
  dropped without any diagnostic.

The tensor reshape verifiers already reject this input (https://github.com/llvm/llvm-project/pull/173791, https://github.com/llvm/llvm-project/issues/222602 / https://github.com/llvm/llvm-project/pull/222648).
Extract the empty-group check into a helper shared by both memref verifiers,
mirroring https://github.com/llvm/llvm-project/pull/222648, and invoke it from `ExpandShapeOp::verify` and
`CollapseShapeOp::verify`.

Add regression tests to `mlir/test/Dialect/MemRef/invalid.mlir`.

Assisted-by: AI.

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


2 Files Affected:

- (modified) mlir/lib/Dialect/MemRef/IR/MemRefOps.cpp (+17) 
- (modified) mlir/test/Dialect/MemRef/invalid.mlir (+31) 


``````````diff
diff --git a/mlir/lib/Dialect/MemRef/IR/MemRefOps.cpp b/mlir/lib/Dialect/MemRef/IR/MemRefOps.cpp
index 4b5a4c1780507..e4bfe20f8307d 100644
--- a/mlir/lib/Dialect/MemRef/IR/MemRefOps.cpp
+++ b/mlir/lib/Dialect/MemRef/IR/MemRefOps.cpp
@@ -2637,7 +2637,21 @@ void ExpandShapeOp::build(OpBuilder &builder, OperationState &result,
   build(builder, result, *resultType, src, reassociation, outputShape);
 }
 
+/// Verify that none of the reassociation groups is empty.
+template <typename MemrefReshapeOp>
+static LogicalResult verifyReassociationIndicesNotEmpty(MemrefReshapeOp 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();
+
   MemRefType srcType = getSrcType();
   MemRefType resultType = getResultType();
 
@@ -2903,6 +2917,9 @@ void CollapseShapeOp::build(OpBuilder &b, OperationState &result, Value src,
 }
 
 LogicalResult CollapseShapeOp::verify() {
+  if (failed(verifyReassociationIndicesNotEmpty(*this)))
+    return failure();
+
   MemRefType srcType = getSrcType();
   MemRefType resultType = getResultType();
 
diff --git a/mlir/test/Dialect/MemRef/invalid.mlir b/mlir/test/Dialect/MemRef/invalid.mlir
index 7f2861985f6dc..c0105290aa6dc 100644
--- a/mlir/test/Dialect/MemRef/invalid.mlir
+++ b/mlir/test/Dialect/MemRef/invalid.mlir
@@ -1220,3 +1220,34 @@ func.func @distinct_objects_0_operands() {
   "memref.distinct_objects"() : () -> ()
   return
 }
+
+// -----
+
+func.func @collapse_shape_empty_reassociation_group(
+    %arg0: memref<2x3xf32, strided<[6, 1], offset: 0>>) {
+  // expected-error at +1 {{'memref.collapse_shape' op reassociation indices must not be empty}}
+  %0 = memref.collapse_shape %arg0 [[], [0, 1]]
+      : memref<2x3xf32, strided<[6, 1], offset: 0>>
+        into memref<1x6xf32, strided<[6, 1], offset: 0>>
+  return
+}
+
+// -----
+
+func.func @collapse_shape_empty_reassociation_group_identity(
+    %arg0: memref<2x3xf32>) {
+  // expected-error at +1 {{'memref.collapse_shape' op reassociation indices must not be empty}}
+  %0 = memref.collapse_shape %arg0 [[], [0, 1]]
+      : memref<2x3xf32> into memref<1x6xf32>
+  return
+}
+
+// -----
+
+func.func @expand_shape_empty_reassociation_group(
+    %arg0: memref<1x6xf32>) {
+  // expected-error at +1 {{'memref.expand_shape' op reassociation indices must not be empty}}
+  %0 = memref.expand_shape %arg0 [[], [0, 1]] output_shape [2, 3]
+      : memref<1x6xf32> into memref<2x3xf32>
+  return
+}

``````````

</details>


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