core: (affine) Add Affine.binary

[watermelonwolverine]

Cherry picked from #3650. Prerequisite for implementing vector.transfer_read and vector.transfer_write ops

[codecov]

Codecov Report

All modified and coverable lines are covered by tests ✅

Project coverage is 90.19%. Comparing base (3d4d51d) to head (01518ad).
Report is 10 commits behind head on main.

Additional details and impacted files

@@            Coverage Diff             @@
##             main    #4006      +/-   ##
==========================================
+ Coverage   90.17%   90.19%   +0.02%     
==========================================
  Files         458      458              
  Lines       58304    58358      +54     
  Branches     5694     5694              
==========================================
+ Hits        52573    52637      +64     
+ Misses       4339     4331       -8     
+ Partials     1392     1390       -2     

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

[watermelonwolverine]

While implementing this I realized that when composing two affine maps the expressions are not automatically simplified/constant folded.
So, a resulting AffineBinaryOpExpr where both lhs and rhs are AffineConstantExpr is not simplified to a AffineConstantExpr.
That's why I added the AffineExpr.as_constant function. Is this by design or should the constant folding happen while composing?

[superlopuh]

What does MLIR do here? I have a suspicion that they constant fold when creating the affine expr, it would be worth doing the same in xDSL. I can't find these helpers in the AffineExpr and AffineMap header files.

[watermelonwolverine]

MLIR simplifies while composing. This is the call stack, the simplification happens at the last call:

AffineExpr AffineExpr::compose(AffineMap map)
AffineExpr AffineExpr::replaceDimsAndSymbols(ArrayRef<AffineExpr> dimReplacements,ArrayRef<AffineExpr> symReplacements)
AffineExpr mlir::getAffineBinaryOpExpr(AffineExprKind kind, AffineExpr lhs, AffineExpr rhs)
AffineExpr AffineExpr::operator+(AffineExpr other)
static AffineExpr simplifyAdd(AffineExpr lhs, AffineExpr rhs)

[watermelonwolverine]

Added AffineExpr.binary which mimics AffineExpr mlir::getAffineBinaryOpExpr from AffineExpr.cpp.
Now it should be pretty close to how MLIR is doing it

[superlopuh]

I've been thinking about this, I think you've found an important deficiency in the API. I have a feeling that this API isn't quite right, but it addresses an important problem. How do you feel about adding a method AffineBinaryOpKind.simplified(lhs: AffineExpr, rhs: AffineExpr) that would perform the simplification in a generic way? That helper could then be used in the __ operators on AffineExpr, so the logic would already be tested.

[watermelonwolverine]

I'm not quite sure I understand what you mean. Do you mean adding something like this to AffineBinaryOpExpr (or did you really mean AffineBinaryOpKind?:

    @staticmethod
    def simplified(
        kind: AffineBinaryOpKind,
        lhs: AffineExpr,
        rhs: AffineExpr,
    ) -> AffineExpr:
        match kind:
            case AffineBinaryOpKind.Add:
                return AffineBinaryOpExpr.simplified_add(lhs, rhs)
            case AffineBinaryOpKind.Mul:
                return AffineBinaryOpExpr.simplified_mul(lhs, rhs)
            case AffineBinaryOpKind.Mod:
                return AffineBinaryOpExpr.simplified_mod(lhs, rhs)
            case AffineBinaryOpKind.FloorDiv:
                return AffineBinaryOpExpr.simplified_floor_div(lhs, rhs)
            case AffineBinaryOpKind.CeilDiv:
                return AffineBinaryOpExpr.simplified_ceil_div(lhs, rhs)

That would work, but really wouldn't do much other than moving and consolidating all the logic that basically already exists (like _simplify_add) from AffineExpr to AffineBinaryOpExpr/AffineBinaryOpKind. The way it is done right now is actually quite generic and is also pretty similar to how MLIR does it as far as I can tell. Can you help me understand what the goal is? Should AffineExpr.binary not return anything but a AffineBinaryOpExpr?