[Mlir-commits] [mlir] Fix bug in `simplifySemiAffine` utility (PR #129200)

[Arnab Dutta]

https://github.com/arnab-polymage created https://github.com/llvm/llvm-project/pull/129200

Whenever `symbolicDivide` returns nullptr when called from inside `simplifySemiAffine` we substitute the result with the original expression (`expr`). nullptr simply indicates that the floordiv expression cannot be simplified further.

>From 4ddb1833beb8c33c3cc2ee74eb37fb16bee814ee Mon Sep 17 00:00:00 2001
From: Arnab Dutta <arnab at polymagelabs.com>
Date: Fri, 28 Feb 2025 12:15:04 +0530
Subject: [PATCH] Fix bug in `simplifySemiAffine` utility

Whenever `symbolicDivide` returns nullptr when called from inside
`simplifySemiAffine` we substitute the result with the original
expression (`expr`). nullptr simply indicates that the floordiv
expression cannot be simplified further.
---
 mlir/lib/IR/AffineExpr.cpp                        |  3 ++-
 mlir/test/Dialect/Affine/simplify-structures.mlir | 11 +++++++++++
 2 files changed, 13 insertions(+), 1 deletion(-)

diff --git a/mlir/lib/IR/AffineExpr.cpp b/mlir/lib/IR/AffineExpr.cpp
index 59df0cd6833db..0077a869c014e 100644
--- a/mlir/lib/IR/AffineExpr.cpp
+++ b/mlir/lib/IR/AffineExpr.cpp
@@ -597,7 +597,8 @@ static AffineExpr simplifySemiAffine(AffineExpr expr, unsigned numDims,
       return getAffineBinaryOpExpr(expr.getKind(), sLHS, sRHS);
     if (expr.getKind() == AffineExprKind::Mod)
       return getAffineConstantExpr(0, expr.getContext());
-    return symbolicDivide(sLHS, symbolPos, expr.getKind());
+    AffineExpr simplifiedQuotient = symbolicDivide(sLHS, symbolPos, expr.getKind());
+    return simplifiedQuotient ? simplifiedQuotient : expr;
   }
   }
   llvm_unreachable("Unknown AffineExpr");
diff --git a/mlir/test/Dialect/Affine/simplify-structures.mlir b/mlir/test/Dialect/Affine/simplify-structures.mlir
index d1f34f20fa5da..723bd366e6ed4 100644
--- a/mlir/test/Dialect/Affine/simplify-structures.mlir
+++ b/mlir/test/Dialect/Affine/simplify-structures.mlir
@@ -282,6 +282,7 @@ func.func @simplify_zero_dim_map(%in : memref<f32>) -> f32 {
 // Tests the simplification of a semi-affine expression in various cases.
 // CHECK-DAG: #[[$map0:.*]] = affine_map<()[s0, s1] -> (-(s1 floordiv s0) + 2)>
 // CHECK-DAG: #[[$map1:.*]] = affine_map<()[s0, s1] -> (-(s1 floordiv s0) + 42)>
+// CHECK-DAG: #[[$NESTED_FLOORDIV:.*]] = affine_map<()[s0] -> ((s0 floordiv s0) floordiv s0)>
 
 // Tests the simplification of a semi-affine expression with a modulo operation on a floordiv and multiplication.
 // CHECK-LABEL: func @semiaffine_mod
@@ -299,6 +300,16 @@ func.func @semiaffine_floordiv(%arg0: index, %arg1: index) -> index {
   return %a : index
 }
 
+// The following semi-affine expression with nested floordiv cannot be simplified.
+// CHECK-LABEL: func @semiaffine_floordiv_no_simplification
+// CHECK-SAME:  %[[ARG0:.*]]: index
+func.func @semiaffine_floordiv_no_simplification(%arg0: index) -> index {
+  %a = affine.apply affine_map<()[s0] ->((s0 floordiv s0) floordiv s0)> ()[%arg0]
+  return %a : index
+  // CHECK:       %[[RES:.*]] = affine.apply #[[$NESTED_FLOORDIV]]()[%[[ARG0]]]
+  // CHECK-NEXT:  return %[[RES]] : index
+}
+
 // Tests the simplification of a semi-affine expression with a ceildiv operation and a division of arith.constant 0 by a symbol.
 // CHECK-LABEL: func @semiaffine_ceildiv
 func.func @semiaffine_ceildiv(%arg0: index, %arg1: index) -> index {