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

[Darshan Bhat]

https://github.com/darshan-opensource updated https://github.com/llvm/llvm-project/pull/181715

>From 3eba8d1c3c1c74ffe4758a460533b9662637309e Mon Sep 17 00:00:00 2001
From: Darshan Bhat <darshanbhatsirsi at gmail.com>
Date: Mon, 16 Feb 2026 19:09:49 +0100
Subject: [PATCH] [MLIR] Fix a crash in CollapseLinalgDimensions

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).
---
 .../Linalg/Transforms/ElementwiseOpFusion.cpp |  7 +++--
 mlir/test/Dialect/Linalg/collapse-dim.mlir    | 29 +++++++++++++++++++
 2 files changed, 34 insertions(+), 2 deletions(-)

diff --git a/mlir/lib/Dialect/Linalg/Transforms/ElementwiseOpFusion.cpp b/mlir/lib/Dialect/Linalg/Transforms/ElementwiseOpFusion.cpp
index 1c69eebc808ba..db46de75abd1a 100644
--- a/mlir/lib/Dialect/Linalg/Transforms/ElementwiseOpFusion.cpp
+++ b/mlir/lib/Dialect/Linalg/Transforms/ElementwiseOpFusion.cpp
@@ -1308,8 +1308,11 @@ bool mlir::linalg::isDimSequencePreserved(AffineMap indexingMap,
                                           ReassociationIndicesRef dimSequence) {
   assert(!dimSequence.empty() &&
          "expected non-empty list for dimension sequence");
-  assert(indexingMap.isProjectedPermutation() &&
-         "expected indexing map to be projected permutation");
+
+  // Dimension sequences can only be preserved in projected permutation maps.
+  if (!indexingMap.isProjectedPermutation()) {
+    return false;
+  }
 
   llvm::SmallDenseSet<unsigned, 4> sequenceElements;
   sequenceElements.insert_range(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>
+}