[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: #122231

Author: arnab-polymage

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");

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 {