Question about FlatAffineConstraints + Possible Refactoring

I have a question about the affine analysis (FlatAffineConstraints) in MLIR.

Dimensions can be associated with values, but this seems to be optional:

void addDimId(unsigned pos, Value id = nullptr);

At the same time, some functions assume that every dimension has a value, e.g., mergeAndAlignIds:

  assert(std::all_of(a->getIds().begin() + offset,
                     a->getIds().begin() + a->getNumDimAndSymbolIds(),
                     [](Optional<Value> id) { return id.hasValue(); }));

This looks inconsistent to me. Is it mandatory to provide values when adding dimensions? Or is this a bug in mergeAndAlignIds (and/or how it is used)?

In my case, mergeAndAlignIds is called via FlatAffineConstraints::addLowerOrUpperBound. I use FlatAffineConstraints in my SCF loop peeling pattern and associated canonicalizations of affine.min ops (⚙ D107222 [mlir][scf] Simplify affine.min ops after loop peeling).

Unrelated to this, there’s a possible refactoring that I was thinking of. The current FlatAffineConstraints seems quite tied to the affine dialect. But the analysis is also useful in other places such as the loop peeling pattern or certain places in linalg. What I propose is:

  • Rename FlatAffineConstraints to ConstraintSet.
  • Create a new subclass FlatAffineConstraints of ConstraintSet and move all affine-specific functions into that class (stuff such as addAffineForOpDomain). There are no assumptions about dimensions in ConstraintSet such as the ones in addInductionVarOrTerminalSymbol. Nothing changes for code that uses FlatAffineConstraints, but users that need only Fourier-Motzkin can use the more lightweight ConstraintSet.
  • Not sure yet how to best incorporate this into the design: In the loop peeling case, addLowerOrUpperBound was very useful, but I don’t want to compose the affine map and operands. Basically, I would like all constraints to be expressed in terms of the operands. E.g., I don’t care if the step size of the loop is defined in terms of another affine.apply op. Just use “step” as a dimension. (Also, some operands may be semi-affine/non-affine and I don’t want to fail in such cases.) Related commit: ⚙ D107221 [mlir][affine] addLowerOrUpperBound: Make map+operand composing optional
  • Not sure what to do with IntegerSet stuff etc. Maybe put it in FlatAffineConstraints instead of ConstraintSet.

What do you think about this? Is this a reasonable refactoring?

The alignment here is done based on unique Values. So, the presence of Values has been a pre-condition for this method. Without those values, there isn’t any merging or alignment we can define within that meaning.

For your use case, it’s really addLowerOrUpperBound to focus on which is expecting the constraint system to have been associated with Values whenever there are local variables. There may be something to fix or improve here.

Reg. FlatAffineConstraints and ConstraintSet, note that the “affine” isn’t because it’s used by the affine dialect but because the constraints themselves (each of the rows of the buffer) are affine. So, from a naming standpoint, it would be inaccurate to separate FlatAffineConstraints and ConstraintSet unless your goal is to have non-affine (non-linear) constraints which isn’t the case here. Also, the “flat” in the name is because unlike AffineMap or AffineExpr there is no tree structure/hierarchy in the expressions.

Yes, but the Values associated with the system are exactly those operands. The real improvement/extension you desire here is dropping any assumptions on results coming from affine ops. I’ll take a look at the loop peeling revision for better context.

Step is actually always a constant for affine.for. Perhaps you meant a Value used in the lower or upper bound operands.

There is perhaps a typo here. IntegerSet is an MLIR attribute. FlatAffineConstraints is a temporary in-memory data structure for analysis.

There are two quick fixes I can think of:

a) Modify mergeAndAlignIds, so that it always adds dimensions without a value, potentially duplicating such dimensions. I think this is how local columns are handled.

b) In my pattern implementation, use local columns (addLocalId) instead of dimensions. However, this would violate assumptions about local columns, potentially breaking other stuff: Each local identifier is always (by construction) a floordiv of a pure add/mul affine function of dimensional, symbolic, and other local identifiers, in a non-mutually recursive way.

Otherwise, getFlattenedAffineExprs would have to be modified to add local IDs to an existing FlatAffineConstraints instead of a brand new one. That’s why column merging is required in the first place.

FlatAffineConstraints is the right name then. I was thinking of isolating the Fourier-Motzkin part + useful helper functions. And putting everything that is affine dialect-specific into a subclass. In particular, the functions that take an AffineIfOp or AffineForOp as a parameter. Ideally, the dependency on mlir/Dialect/Affine/IR/AffineOps.h could be dropped (AffineExpr.h is OK).

One could even go so far as removing addLowerOrUpperBound etc. from the base class, with the only functionality left being Fourier-Motzkin and helper functions for adding adding equalities/inequalities and querying lower/upper bounds. But addLowerOrUpperBound is actually pretty useful…

Correct, that’s the main part. These “assumptions on results coming from affine ops” seem to be useful in the context of the affine dialect mostly, right? In my case, I don’t care where values are coming from, I just want to compute an upper/lower bound by solving a system of inequalities. So maybe, coming back to the aforementioned refactoring, the base class could drop such assumptions, and the subclass (being used in the context of the affine dialect) could have them.

It can be non-constant in scf.for loops. But you are right, in many cases step is constant and only the lower/upper bounds are non-constant.

I was referring to functions such as:

  /// Returns the constraint system as an integer set. Returns a null integer
  /// set if the system has no constraints, or if an integer set couldn't be
  /// constructed as a result of a local variable's explicit representation not
  /// being known and such a local variable appearing in any of the constraints.
  IntegerSet getAsIntegerSet(MLIRContext *context) const;

Not sure, where this is used. Anyway, I am not really concerned about these…

On a general note, FlatAffineConstraints has been around for a long time and some design decisions have been revised since it was introduced. You may notice references to “ML values”, which were a special kind of SSA values usable in an affine context and always subject to what we currently know as restrictions on affine symbols and dimensions. (By the way, cleaning up those references is very welcome!)

On a more or less cursory glance, the refactoring that makes the most sense to me that should also satisfy your needs is separating the purely “constraint set” data structure agnostic of values from the “constraint set where dimensions may be associated with values” data structure. This is similar to what we have for AffineMap and AffineValueMap. The former will lose some simplification capability, e.g. it will have to concatenate parameter lists by default unless it is guaranteed by the caller that the lists are aligned, but can be used in contexts where values are either irrelevant (such as yours) or even not available. This doesn’t look very deep, we just need to make sure that implementations of core FlatAffineConstraints methods can handle the absence of Value association conservatively, e.g., again concatenate parameter lists if they can’t be aligned. Creating a separate FlatAffineValueConstaints class that derives FlatAffineConstraints and moving the ids field into it looks like a good way to shake off all the issues. At FlatAffineConstraints level, we can provide merge-like APIs that additionally take a map indicating which dimensions of one constraint set must be merged with which dimensions of another constraint set.

I agree with Uday’s reasoning about keeping the name. The set is indeed composed of affine constraints that are expressed in a flat format. We may and likely will have other kinds of constraint sets.

Local (or existentially quantified) variables have a special semantics in affine sets and serve to express divisibility: there exists a local integer variable L such that d0 = 3*L is the way of saying “d0 must be divisible by 3”. We don’t want to diverge from that.

You can still have it by providing a different mechanism that indicates how to merge dimensions and symbols, e.g., by providing an explicit mapping, than the current one based on SSA values. It may be sufficient to check that the SSA values are different, not necessarily traverse their affine provenance. Again, it may lose some simplification power, but should work.

IntegerSet is a built-in concept, not a part of the Affine dialect. And it shouldn’t need the association with Values to be constructed

All of what @ftynse suggests makes a lot of sense to me. I am +1 to any refactoring that allows reuse of more base functionality.

I’m happy to clean these stale references up.

On a side note, I wanted to mention that as scf.for can always be normalized to have stride one (as you know) and as such you’ll always have a linear constraint:

%lb + %i <= %ub

Otherwise, with a step that is not a constant SSA value, you don’t have a linear constraint.

This sounds good! I’ll prepare a revision and send it for review when it’s ready.

Thanks everyone for the discussion and reviews. The FlatAffineConstraints splitting and a large part of the SCF loop peeling pattern has landed.

Interestingly, in the final version of the affine.min canonicalization (after peeling), I ended up using FlatAffineValueConstraints because it was tedious to keep track of id → value mappings manually. (But also using many of the refactored member functions.) As it turns out, associating ids with SSA values is pretty useful: When adding a new constraint drawn from an affine map (and operands), I have to know which dim/sym column a given dim/sym of the map corresponds to, in order to align the map with the constraint system. This is done by checking for SSA value equality.

I think there may be an additional refactoring opportunity here: Extract affine-dialect related stuff from FlatAffineValueConstraints into yet another subclass, name TBD. In total, there would be 3 classes:

  • FlatAffineConstraints: Should be used for solving a set of (in)equations, when no SSA values are associated with the identifiers.
  • FlatAffineValueConstraints: Same as FlatAffineConstraints, but maintains an optional id → SSA value mapping and a few helper functions for aligning AffineValueMaps with the constraint system. Nothing more.
  • FlatAffineOpValueConstraints: Name TBD, not sure what to call this. This is the same as FlatAffineValueConstraints, but has additional methods for dealing with Affine dialect ops. By “Affine dialect ops” I mean affine.for and affine.if. I do not mean affine.min, affine.max and affine.apply. (The latter ones are often used to do mathematical computations in lieu of std.addi etc.).

Stuff that would go into FlatAffineOpValueConstraints:

  • addAffineForOpDomain
  • addDomainFromSliceMaps
  • addAffineIfOpDomain
  • addInductionVarOrTerminalSymbol

In particular addInductionVarOrTerminalSymbol, which is currently called from one of the addBound overloadings.

The benefit of this refactorings would be IMO a cleaner separation of FlatAffineConstraints from the Affine dialect. It could also make the API safer to use, because a user of FlatAffineValueConstraints could be calling the problematic overloading of addBound, which calls addInductionVarOrTerminalSymbol, and then raises an assertion because there are assumptions about how certain associated SSA values are defined.

Here is a WIP revision that shows the proposed changes (users of FlatAffineValueConstraints not updated yet): ⚙ D108442 [WIP][mlir][Analysis] Split FlatAffineValueConstraints

What do you think? Is it worth the refactoring?