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

[Ege Beysel]

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

>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 1/2] [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

>From e87ac3423b54537204d41b8de660f5becaa6f46c Mon Sep 17 00:00:00 2001
From: Ege Beysel <beyselege at gmail.com>
Date: Mon, 15 Jun 2026 23:54:23 +0200
Subject: [PATCH 2/2] add comment

Co-authored-by: Artem Gindinson <gindinson at roofline.ai>
---
 mlir/test/Dialect/Affine/canonicalize.mlir | 3 ++-
 1 file changed, 2 insertions(+), 1 deletion(-)

diff --git a/mlir/test/Dialect/Affine/canonicalize.mlir b/mlir/test/Dialect/Affine/canonicalize.mlir
index 561b021fbc30a..7d236ef3c2421 100644
--- a/mlir/test/Dialect/Affine/canonicalize.mlir
+++ b/mlir/test/Dialect/Affine/canonicalize.mlir
@@ -1863,7 +1863,8 @@ 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.
+// Canonicalization can produce a linearize with no inputs (empty basis),
+// which must not be treated as splittable.
 // CHECK-LABEL: func @split_delinearize_empty_linearize_basis
 //  CHECK-SAME:     %[[ARG0:[a-zA-Z0-9]+]]: index)
 //       CHECK-DAG:   %[[C0:.+]] = arith.constant 0 : index