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

[llvmlistbot at llvm.org]

Author: Arnab Dutta
Date: 2025-03-01T19:47:15+05:30
New Revision: 60b44d31afe53556ec707dc855983b93971d83c4

URL: https://github.com/llvm/llvm-project/commit/60b44d31afe53556ec707dc855983b93971d83c4
DIFF: https://github.com/llvm/llvm-project/commit/60b44d31afe53556ec707dc855983b93971d83c4.diff

LOG: [MLIR][Affine] Fix bug in `simplifySemiAffine` utility (#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.

Fixes: https://github.com/llvm/llvm-project/issues/122231

Added: 
    

Modified: 
    mlir/lib/IR/AffineExpr.cpp
    mlir/test/Dialect/Affine/simplify-structures.mlir

Removed: 
    


################################################################################
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..e4a8512b002ee 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_nested_floordiv
+// CHECK-SAME:  %[[ARG0:.*]]: index
+func.func @semiaffine_nested_floordiv(%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 {