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

[Arnab Dutta]

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

>From 83c034dedaaa2f054158f6e15b82a45bd950dc94 Mon Sep 17 00:00:00 2001
From: Arnab Dutta <arnab at polymagelabs.com>
Date: Fri, 28 Feb 2025 13:55:54 +0530
Subject: [PATCH] [MLIR][Affine] Fix bug in `simplifySemiAffine` utility

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

diff --git a/mlir/lib/IR/AffineExpr.cpp b/mlir/lib/IR/AffineExpr.cpp
index 59df0cd6833db..c8d9761511bec 100644
--- a/mlir/lib/IR/AffineExpr.cpp
+++ b/mlir/lib/IR/AffineExpr.cpp
@@ -597,7 +597,11 @@ 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
+               : getAffineBinaryOpExpr(expr.getKind(), sLHS, sRHS);
   }
   }
   llvm_unreachable("Unknown AffineExpr");
diff --git a/mlir/test/Dialect/Affine/simplify-structures.mlir b/mlir/test/Dialect/Affine/simplify-structures.mlir
index d1f34f20fa5da..304972cf0efa7 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: #[[$FLOORDIV:.*]] = affine_map<()[s0] -> (1 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 #[[$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 {