[Mlir-commits] [mlir] [MLIR] Fix a crash in CollapseLinalgDimensions (PR #181715)

[llvmlistbot at llvm.org]

llvmbot wrote:


<!--LLVM PR SUMMARY COMMENT-->

@llvm/pr-subscribers-mlir-linalg

Author: Darshan Bhat (darshan-opensource)

<details>
<summary>Changes</summary>

Added a check in areDimSequencesPreserved()
to verify that each map is a projected permutation before calling isDimSequencePreserved().
If a map is not a projected permutation, the
function returns false (dimension sequences cannot be preserved in non-projected-permutation maps).

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


2 Files Affected:

- (modified) mlir/lib/Dialect/Linalg/Transforms/ElementwiseOpFusion.cpp (+4) 
- (modified) mlir/test/Dialect/Linalg/collapse-dim.mlir (+29) 


``````````diff
diff --git a/mlir/lib/Dialect/Linalg/Transforms/ElementwiseOpFusion.cpp b/mlir/lib/Dialect/Linalg/Transforms/ElementwiseOpFusion.cpp
index 1c69eebc808ba..48604a889fed0 100644
--- a/mlir/lib/Dialect/Linalg/Transforms/ElementwiseOpFusion.cpp
+++ b/mlir/lib/Dialect/Linalg/Transforms/ElementwiseOpFusion.cpp
@@ -1348,6 +1348,10 @@ bool mlir::linalg::isDimSequencePreserved(AffineMap indexingMap,
 bool mlir::linalg::areDimSequencesPreserved(
     ArrayRef<AffineMap> maps, ArrayRef<ReassociationIndices> dimSequences) {
   return llvm::all_of(maps, [&](AffineMap map) {
+    // Dimension sequences can only be preserved in projected permutation maps.
+    if (!map.isProjectedPermutation()) {
+      return false;
+    }
     return llvm::all_of(dimSequences, [&](ReassociationIndicesRef dimSequence) {
       return isDimSequencePreserved(map, dimSequence);
     });
diff --git a/mlir/test/Dialect/Linalg/collapse-dim.mlir b/mlir/test/Dialect/Linalg/collapse-dim.mlir
index 2995588e9abca..61c4234c301f8 100644
--- a/mlir/test/Dialect/Linalg/collapse-dim.mlir
+++ b/mlir/test/Dialect/Linalg/collapse-dim.mlir
@@ -199,3 +199,32 @@ func.func private @memref_linalg_copy(%arg0: memref<1x24x32x8xf32, 1>, %arg1: me
   linalg.copy ins(%arg0: memref<1x24x32x8xf32, 1>) outs(%arg1: memref<1x24x32x8xf32, 1>)
   return
 }
+
+// -----
+
+// Test that non-projected-permutation indexing maps don't crash the collapse pass.
+// The indexing map has affine expressions like d1 + d4, which are not projected
+// permutations. The pass should skip collapsing and leave the op unchanged.
+// CHECK-DAG: #[[$MAP:.*]] = affine_map<(d0, d1, d2, d3, d4, d5) -> (d0, d1 + d4, d2 + d5, d3)>
+// CHECK-DAG: #[[$MAP1:.*]] = affine_map<(d0, d1, d2, d3, d4, d5) -> (d4, d5, d3, 0)>
+// CHECK-DAG: #[[$MAP2:.*]] = affine_map<(d0, d1, d2, d3, d4, d5) -> (d0, d1, d2, d3)>
+// CHECK-LABEL: func @non_projected_permutation_no_crash
+//       CHECK:   linalg.generic {indexing_maps = [#[[$MAP]], #[[$MAP1]], #[[$MAP2]]],
+//  CHECK-SAME:       iterator_types = ["parallel", "parallel", "parallel", "parallel", "reduction", "reduction"]
+
+func.func @non_projected_permutation_no_crash(
+    %arg0: tensor<1x5x5x1xf32>, %arg1: tensor<2x2x1x1xf32>) -> tensor<1x4x4x1xf32> {
+  %cst = arith.constant dense<0.0> : tensor<1x4x4x1xf32>
+  %0 = linalg.generic {
+    indexing_maps = [affine_map<(d0, d1, d2, d3, d4, d5) -> (d0, d1 + d4, d2 + d5, d3)>,
+                     affine_map<(d0, d1, d2, d3, d4, d5) -> (d4, d5, d3, 0)>,
+                     affine_map<(d0, d1, d2, d3, d4, d5) -> (d0, d1, d2, d3)>],
+    iterator_types = ["parallel", "parallel", "parallel", "parallel", "reduction", "reduction"]
+  } ins(%arg0, %arg1 : tensor<1x5x5x1xf32>, tensor<2x2x1x1xf32>) outs(%cst : tensor<1x4x4x1xf32>) {
+  ^bb0(%in: f32, %w: f32, %out: f32):
+    %mul = arith.mulf %in, %w : f32
+    %add = arith.addf %out, %mul : f32
+    linalg.yield %add : f32
+  } -> tensor<1x4x4x1xf32>
+  return %0 : tensor<1x4x4x1xf32>
+}

``````````

</details>


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