Simplification of AffineMaps

In one of my unit tests, I am seeing affine maps that could be further simplified. E.g., in https://github.com/llvm/llvm-project/blob/main/mlir/test/Dialect/SCF/for-loop-peeling.mlir#L4:

#map = affine_map<()[s0, s1, s2] -> (-(s0 - (s0 - s1) mod s2) + s0)>

Note: The test case is running with -canonicalize.

Looking at AffineExpr.cpp, there are helper functions that simplify various affine exprs, e.g.:

/// Simplify add expression. Return nullptr if it can't be simplified.
static AffineExpr simplifyAdd(AffineExpr lhs, AffineExpr rhs) {
  /* ... */
}

It may be possible to just add a new rule there. However, as AffineExprs are trees, eliminating two exprs such as s0 in the example above is difficult, if they appear in two different subtrees (maybe deeply nested). That and the fact that stuff such as ... - s0 is actually represented as ... + (-1 * s0).

This made me wonder, why do we have simplifyAdd etc. in the first place? (Performance, maybe?) Could this entire simplification logic be replaced by:

  1. Convert the AffineExpr into its flattened form.
  2. Convert the flattened form back into an AffineExpr tree form.

Going to the flattened form and back would be the most comprehensive form of simplification IMO, while simplifyAdd etc. can handle only a few cases that we match for explicitly.

simplifyAffineExpr already does that:

/// Simplify the affine expression by flattening it and reconstructing it.
AffineExpr mlir::simplifyAffineExpr(AffineExpr expr, unsigned numDims,
                                    unsigned numSymbols) {
  // Simplify semi-affine expressions separately.
  if (!expr.isPureAffine())
    expr = simplifySemiAffine(expr);
  if (!expr.isPureAffine())
    return expr;

  SimpleAffineExprFlattener flattener(numDims, numSymbols);
  flattener.walkPostOrder(expr);
  ArrayRef<int64_t> flattenedExpr = flattener.operandExprStack.back();
  auto simplifiedExpr =
      getAffineExprFromFlatForm(flattenedExpr, numDims, numSymbols,
                                flattener.localExprs, expr.getContext());
  flattener.operandExprStack.pop_back();
  assert(flattener.operandExprStack.empty());

  return simplifiedExpr;
}

Those simplify<AffineExprTy> methods were inexpensive methods to do simplification by construction itself – to avoid such trivially simplifiable expressions to be constructed in the first place (avoids MLIRContext resident lifetime storage).

Looking at this function, the problem here is the expr.isPureAffine(). The affine expr above has a ... mod s2, so it cannot be flattened. I guess the only way to simplify this affine map is to extend simplifyAdd etc. then…

No, we can actually extend simplifyAffineExpr and the associated flattening to handle this one as well. If you prefer, could you create a bugzilla issue and assign it to me? (bugs.llvm.org) simplifyAdd is only for trivial local simplification - so it won’t be powerful enough. For eg., you’ll be stuck if you had:
(((-(s0 - (s0 - s1) mod s2) + s3) + (s1 + s0))
and even here s0 would cancel out. The flattening based simplification will catch this as well.

Makes sense, if the flattening can be extended to support semi-affine exprs, then that’s definitely the way to go.

Filed this as 51573 – Support semi-affine exprs in expr flattener

Hi, @matthias-springer please review the following revisions related to simplification of semi-affine exprs.

https://reviews.llvm.org/D112808
https://reviews.llvm.org/D112985

CC - @bondhugula