[Mlir-commits] [mlir] [mlir][affine] Fix crash in `SplitDelinearizeSpanningLastLinearizeArg` (PR #201879)

[Ege Beysel]

https://github.com/egebeysel created https://github.com/llvm/llvm-project/pull/201879

`SplitDelinearizeSpanningLastLinearizeArg` reads the last element of the
source `affine.linearize_index`'s static basis via
`linearizeOp.getStaticBasis().back()` without checking that the basis is
non-empty.

Other canonicalization patterns can legitimately create an
`affine.linearize_index` with no inputs (and therefore an empty basis):
`CancelDelinearizeOfLinearizeDisjointExactTail` peels the matched trailing
basis off both the linearize and the delinearize, and when every linearize
input is part of the matched tail the rewritten linearize is left with no
inputs. When the resulting delinearize is then visited by
`SplitDelinearizeSpanningLastLinearizeArg`, `.back()` is called on an empty
`ArrayRef`, tripping the `assert(!empty())` in `ArrayRef::back()` (or reading
out of bounds in a no-assertions build).

A linearize with no inputs is the constant zero and is not splittable, so
bail out of the pattern when the basis is empty. The empty linearize is
folded away to a constant by its own folder on a later iteration.

Fixes a crash reachable from `-canonicalize`.


Assisted-by: Claude Code

>From ea3d48554cf8ffb9c8ca10bbed1a3ea948158386 Mon Sep 17 00:00:00 2001
From: Ege Beysel <beyselege at gmail.com>
Date: Fri, 5 Jun 2026 18:11:37 +0200
Subject: [PATCH] [mlir][affine] Fix crash in
 SplitDelinearizeSpanningLastLinearizeArg on empty-basis linearize

Signed-off-by: Ege Beysel <beyselege at gmail.com>
---
 mlir/lib/Dialect/Affine/IR/AffineOps.cpp   |  7 +++++++
 mlir/test/Dialect/Affine/canonicalize.mlir | 17 +++++++++++++++++
 2 files changed, 24 insertions(+)

diff --git a/mlir/lib/Dialect/Affine/IR/AffineOps.cpp b/mlir/lib/Dialect/Affine/IR/AffineOps.cpp
index c2d20271a815f..0ecd733842f31 100644
--- a/mlir/lib/Dialect/Affine/IR/AffineOps.cpp
+++ b/mlir/lib/Dialect/Affine/IR/AffineOps.cpp
@@ -5172,6 +5172,13 @@ struct SplitDelinearizeSpanningLastLinearizeArg final
       return rewriter.notifyMatchFailure(linearizeOp,
                                          "linearize isn't disjoint");
 
+    // A linearize with no inputs has an empty basis and folds to a constant
+    // zero; there is nothing to split, and reading its last basis element
+    // below would be out of bounds.
+    if (linearizeOp.getStaticBasis().empty())
+      return rewriter.notifyMatchFailure(
+          linearizeOp, "linearize has no basis elements (no inputs)");
+
     int64_t target = linearizeOp.getStaticBasis().back();
     if (ShapedType::isDynamic(target))
       return rewriter.notifyMatchFailure(
diff --git a/mlir/test/Dialect/Affine/canonicalize.mlir b/mlir/test/Dialect/Affine/canonicalize.mlir
index 347eb64086228..561b021fbc30a 100644
--- a/mlir/test/Dialect/Affine/canonicalize.mlir
+++ b/mlir/test/Dialect/Affine/canonicalize.mlir
@@ -1861,6 +1861,23 @@ func.func @dont_split_delinearize_undershooting_target(%arg0: index, %arg1: inde
 
 // -----
 
+// Regression test: canonicalization can produce a linearize with no inputs
+// (empty basis), which must not be treated as splittable. This previously
+// crashed when reading the (nonexistent) last basis element.
+// CHECK-LABEL: func @split_delinearize_empty_linearize_basis
+//  CHECK-SAME:     %[[ARG0:[a-zA-Z0-9]+]]: index)
+//       CHECK-DAG:   %[[C0:.+]] = arith.constant 0 : index
+//           CHECK:   return %[[C0]], %[[C0]], %[[ARG0]], %[[C0]]
+func.func @split_delinearize_empty_linearize_basis(%arg0: index) -> (index, index, index, index) {
+  %c0 = arith.constant 0 : index
+  %0 = affine.linearize_index disjoint [%c0, %arg0, %c0] by (4, 2, 2) : index
+  %1:4 = affine.delinearize_index %0 into (2, 2, 2, 2)
+      : index, index, index, index
+  return %1#0, %1#1, %1#2, %1#3 : index, index, index, index
+}
+
+// -----
+
 // CHECK-LABEL: @linearize_unit_basis_disjoint
 // CHECK-SAME: (%[[arg0:.+]]: index, %[[arg1:.+]]: index, %[[arg2:.+]]: index, %[[arg3:.+]]: index)
 // CHECK: %[[ret:.+]] = affine.linearize_index disjoint [%[[arg0]], %[[arg2]]] by (3, %[[arg3]]) : index