[MLIR][Affine] Reject fully consumed bounded delinearize split

[joker-eph]

Do not apply SplitDelinearizeSpanningLastLinearizeArg when the split would consume an entire outer-bounded basis. Rewriting that case can discard earlier linearization inputs and previously built an invalid zero-result prefix operation.

Assisted-by: Codex

[llvmorg-github-actions]

@llvm/pr-subscribers-mlir

@llvm/pr-subscribers-mlir-affine

Author: Mehdi Amini (joker-eph)

Changes

Do not apply SplitDelinearizeSpanningLastLinearizeArg when the split would consume an entire outer-bounded basis. Rewriting that case can discard earlier linearization inputs and previously built an invalid zero-result prefix operation.

Assisted-by: Codex

---

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

2 Files Affected:

• (modified) mlir/lib/Dialect/Affine/IR/AffineOps.cpp (+7)
• (modified) mlir/test/Dialect/Affine/canonicalize.mlir (+21-2)

diff --git a/mlir/lib/Dialect/Affine/IR/AffineOps.cpp b/mlir/lib/Dialect/Affine/IR/AffineOps.cpp
index b1f7b987703a5..9f0734dab3b31 100644
--- a/mlir/lib/Dialect/Affine/IR/AffineOps.cpp
+++ b/mlir/lib/Dialect/Affine/IR/AffineOps.cpp
@@ -5233,6 +5233,9 @@ struct CancelDelinearizeOfLinearizeDisjointExactTail
 /// last k > 1 components of the delinearization basis multiply to the
 /// last component of the linearization basis, break the linearization and
 /// delinearization into two parts, peeling off the last input to linearization.
+/// The split does not apply when it would consume an entire outer-bounded
+/// delinearization basis because earlier linearization inputs still contribute
+/// to the first delinearized result.
 ///
 /// For example:
 ///    %0 = affine.linearize_index [%z, %y, %x] by (3, 2, 32) : index
@@ -5297,6 +5300,10 @@ struct SplitDelinearizeSpanningLastLinearizeArg final
           delinearizeOp,
           "need at least two elements to form the basis product");
 
+    if (elemsToSplit == basis.size() && delinearizeOp.hasOuterBound())
+      return rewriter.notifyMatchFailure(
+          delinearizeOp, "split would consume entire bounded basis");
+
     Value linearizeWithoutBack = affine::AffineLinearizeIndexOp::create(
         rewriter, linearizeOp.getLoc(), linearizeOp.getLinearIndex().getType(),
         linearizeOp.getMultiIndex().drop_back(), linearizeOp.getDynamicBasis(),
diff --git a/mlir/test/Dialect/Affine/canonicalize.mlir b/mlir/test/Dialect/Affine/canonicalize.mlir
index 1b13d29335523..57b92e79403fe 100644
--- a/mlir/test/Dialect/Affine/canonicalize.mlir
+++ b/mlir/test/Dialect/Affine/canonicalize.mlir
@@ -1,8 +1,6 @@
 // RUN: mlir-opt -allow-unregistered-dialect %s -split-input-file -canonicalize="test-convergence" | FileCheck %s
 // RUN: mlir-opt -allow-unregistered-dialect %s -split-input-file -canonicalize="test-convergence top-down=0" | FileCheck %s --check-prefix=CHECK-BOTTOM-UP
 
-// XFAIL: mlir-expensive-checks
-
 // -----
 
 // CHECK-DAG: #[[$MAP0:.*]] = affine_map<(d0) -> (d0 - 1)>
@@ -1881,6 +1879,27 @@ func.func @split_delinearize_empty_linearize_basis(%arg0: index) -> (index, inde
 
 // -----
 
+// A split that consumes an entire bounded delinearization basis would lose the
+// contribution of earlier linearization inputs to the first result.
+// CHECK-LABEL: func @dont_split_fully_consumed_bounded_basis
+// CHECK-SAME:    (%[[A:.+]]: index, %[[B:.+]]: index)
+// CHECK:         %[[LIN:.+]] = affine.linearize_index disjoint [%[[A]], %[[B]]] by (2, 4) : index
+// CHECK:         %[[DELIN:.+]]:2 = affine.delinearize_index %[[LIN]] into (2, 2) : index, index
+// CHECK:         return %[[DELIN]]#0, %[[DELIN]]#1
+// CHECK-BOTTOM-UP-LABEL: func @dont_split_fully_consumed_bounded_basis
+// CHECK-BOTTOM-UP-SAME:    (%[[A:.+]]: index, %[[B:.+]]: index)
+// CHECK-BOTTOM-UP:         %[[LIN:.+]] = affine.linearize_index disjoint [%[[A]], %[[B]]] by (2, 4) : index
+// CHECK-BOTTOM-UP:         %[[DELIN:.+]]:2 = affine.delinearize_index %[[LIN]] into (2, 2) : index, index
+// CHECK-BOTTOM-UP:         return %[[DELIN]]#0, %[[DELIN]]#1
+func.func @dont_split_fully_consumed_bounded_basis(%a: index, %b: index)
+    -> (index, index) {
+  %0 = affine.linearize_index disjoint [%a, %b] by (2, 4) : index
+  %1:2 = affine.delinearize_index %0 into (2, 2) : index, index
+  return %1#0, %1#1 : 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

[banach-space]

LGTM, thank you!