[AffineCFG] Re-arrange sum expression after recreateexpr optimization

[xys-syx]

The minimal reproducer example is:

func.func @forbound_reassoc(%arg0: i32, %init: f64) -> f64 {
  %cst = arith.constant 1.000000e+00 : f64
  %s = arith.index_cast %arg0 : i32 to index
  %r = affine.for %i = 0 to %s iter_args(%acc = %init) -> (f64) {
    %ri = affine.for %j = 0 to affine_map<(d0)[s0] -> (-d0 + s0 - 1)>(%i)[%s]
            iter_args(%a = %acc) -> (f64) {
      %sum = arith.addf %a, %cst : f64
      affine.yield %sum : f64
    }
    affine.yield %ri : f64
  }
  return %r : f64
}

The bug here is:
when CanonicalieForBounds apply on affine.for %j = 0 to affine_map<(d0)[s0] -> (-d0 + s0 - 1)>(%i)[%s]...'s upper bound, the upper bound map and operands affine_map<(d0)[s0] -> (-d0 + s0 - 1)>(%i)[%s] will not change, but recreateExpr

Enzyme-JAX/src/enzyme_ad/jax/Passes/SimplifyAffineExprs.cpp

Line 453 in 6af35e7

AffineExpr mlir::enzyme::recreateExpr(AffineExpr expr) {

change the tree associations of sum expression:
from Add(Add(Mul(d0, C-1), s0), C-1) to Add(Mul(d0, C-1), Add(s0, C-1))
causing --affine-cfg fails to converage.

The fix:
I add a canonicalizeSum to rearrange the sum expression, align with the rule in default mlir simplifyAdd, https://github.com/llvm/llvm-project/blob/5db7e5b483a32f37d1d77a6f18d2bc4c59fa460a/mlir/lib/IR/AffineExpr.cpp#L662

[codecov]

Codecov Report

✅ All modified and coverable lines are covered by tests.
✅ Project coverage is 25.46%. Comparing base (e5fe41b) to head (7658d4d).
⚠️ Report is 42 commits behind head on main.

Additional details and impacted files

@@           Coverage Diff           @@
##             main    #2557   +/-   ##
=======================================
  Coverage   25.46%   25.46%           
=======================================
  Files         223      223           
  Lines       44652    44669   +17     
=======================================
+ Hits        11370    11375    +5     
- Misses      33282    33294   +12     

☔ View full report in Codecov by Harness.
📢 Have feedback on the report? Share it here.

🚀 New features to boost your workflow:

• ❄️ Test Analytics: Detect flaky tests, report on failures, and find test suite problems.

[vimarsh6739]

@codex review

[chatgpt-codex-connector]

Codex Review: Didn't find any major issues. What shall we delve into next?

Reviewed commit: 7658d4d062

ℹ️ About Codex in GitHub

Your team has set up Codex to review pull requests in this repo. Reviews are triggered when you

• Open a pull request for review
• Mark a draft as ready
• Comment "@codex review".

If Codex has suggestions, it will comment; otherwise it will react with 👍.

Codex can also answer questions or update the PR. Try commenting "@codex address that feedback".