Presburger / Affine Analysis: Modeling domain / range dimensions, FlatAffineRelation::compose bug

Hi everyone,

when working with the Presburger library / Affine Analysis, I noticed that there seem to be two places for a FlatAffineRelation to keep track of the number of Domain and Range dimensions/variables:

  1. llvm-project/PresburgerSpace.h at main · llvm/llvm-project · GitHub
  2. llvm-project/AffineStructures.h at main · llvm/llvm-project · GitHub

My best understanding of where this comes from is that at the time when the FlatAffineRelation was created ([MLIR] Generalize Affine dependence analysis using Affine Relations · llvm/llvm-project@52d6c5d · GitHub), the Presburger library didn’t exist and the underlying representation of the relation space was therefore different.
Another possible reason is that FlatAffineRelation is built on top of the FlatLinearValueConstraints, which (as far as I understand) is itself built indirectly on top of the IntegerRelation/PresburgerSpace but uses them to model a set and therefore gets rid of the separation between domain and range.
To summarize, it seems to me that we are building a relation on top of a set, which is built on top of a relation; losing and reintroducing the domain/range separation along the way.

Or is this maybe intentional and there is a difference in semantics between the “domain/range vars” in the Presburger lib’s IntegerRelation/PresburgerSpace and the “domain/range dims” in the FlatAffineRelation?

It definitely leads to some strange results when removing dimensions/variables because FlatAffineRelation sometimes “forgets” to reduce the number of dimensions. Example:

  // rel2_1 is a relation mapping from (d0, d1) -> (r0)
  mlir::FlatAffineRelation rel2_1;
  rel2_1.appendDomainVar(2);
  rel2_1.appendRangeVar(1);

  std::cout << "rel2_1 domain vars: " << rel2_1.getNumDomainVars() << "\n" // PresburgerSpace
            << "rel2_1 domain dims: " << rel2_1.getNumDomainDims() << "\n" // FlatAffineRelation
            << "rel2_1 range vars: " << rel2_1.getNumRangeVars() << "\n"   // PresburgerSpace
            << "rel2_1 range dims: " << rel2_1.getNumRangeDims() << "\n";  // FlatAffineRelation

  // rel3_2 is a relation mapping from (d2, d3, d4) -> (r1, r2)
  mlir::FlatAffineRelation rel3_2;
  rel3_2.appendDomainVar(3);
  rel3_2.appendRangeVar(2);

  // We want to get the relation (d2, d3, d4) -> (r0) by composing rel2_1 with rel3_2
  mlir::FlatAffineRelation rel3_1 = rel2_1;
  rel3_1.compose(rel3_2);

  std::cout << "rel3_1 domain vars: " << rel3_1.getNumDomainVars() << "\n" // PresburgerSpace
            << "rel3_1 domain dims: " << rel3_1.getNumDomainDims() << "\n" // FlatAffineRelation
            << "rel3_1 range vars: " << rel3_1.getNumRangeVars() << "\n"   // PresburgerSpace
            << "rel3_1 range dims: " << rel3_1.getNumRangeDims() << "\n";  // FlatAffineRelation

  rel3_1.dump();

Output:

rel2_1 domain vars: 0
rel2_1 domain dims: 2
rel2_1 range vars: 3
rel2_1 range dims: 1
rel3_1 domain vars: 0
rel3_1 domain dims: 3
rel3_1 range vars: 4
rel3_1 range dims: 2
Domain: 0, Range: 4, Symbols: 0, Locals: 2
0 constraints
(None   None    None    None    Local   Local   const)

According to the FlatAffineRelation, the result of the composition has 3 domain dimensions and 2 range dimensions. The correct result would have 3 domain dimensions and 1 range dimension. When dumping the underlying constraint representation, we can see that it seems to be correct there: 4 range variables (3 domain + 1 range).

The reason for the wrong count in FlatAffineRelation seems to be that the two ranges get concatenated in the beginning of the composition, increasing the range counter in FlatAffineRelation:

But then afterwards, when getting rid of the extra range dimensions by converting them to locals, the functions of the underlying IntegerRelation are used, therefore not decreasing the range counter in FlatAffineRelation:

I would be very interested to hear your opinions on this. Is it a bug? If so, would it be appropriate to fix it just by decreasing the dimension counter in compose? It seems to me like a better fix would be to use the underlying representations fully, or is there a reason for modeling the domain/range twice?

Hi! This problem of having Domain and Range dimensions/variables at two different places is known. I currently have some patches in review that plan to slowly fix this: ⚙ D146912 [MLIR][Analysis] Use Identifiers outside Presburger library.

After this patch, you can just directly use IntegerRelation instead of FlatAffineRelation. I have other patches ready to improve this interface and I will be working on upstreaming them in the coming weeks.