[MLIR][Affine] Fix bug in simplifySemiAffine utility

[arnab-polymage]

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.

Fixes: #122231

[llvmbot]

@llvm/pr-subscribers-mlir
@llvm/pr-subscribers-mlir-core

@llvm/pr-subscribers-mlir-affine

Author: Arnab Dutta (arnab-polymage)

Changes

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.

---

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

2 Files Affected:

• (modified) mlir/lib/IR/AffineExpr.cpp (+2-1)
• (modified) mlir/test/Dialect/Affine/simplify-structures.mlir (+11)

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 {

[github-actions]

✅ With the latest revision this PR passed the C/C++ code formatter.