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 { // 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 {