[MLIR][Presburger] Move IdKind specific insert/append

[Groverkss]

Summary

This patch moves identifier kind specific insert/append functions like
insertDimId, appendSymbolId, etc. from IntegerPolyhedron to
FlatAffineConstraints.

This change allows for a smoother transition to IntegerRelation.

This change is part of a series of patches to introduce Relations in Presburger
library.

[arjunp]

mlir/include/mlir/Analysis/Presburger/PresburgerSpace.h

28

If this is going to be exposed in presburger directly it may be better to make it an enum class

73

Why is this needed?

[Groverkss]

• Address Arjun's comments

[Groverkss]

mlir/include/mlir/Analysis/Presburger/PresburgerSpace.h

73

This was to allow FlatAffineConstraints and it's derivatives (which are out of presburger namespace) to refer to IdKind directly. Moved this to FlatAffineConstraints.

[arjunp]

LGTM

mlir/include/mlir/Analysis/Presburger/IntegerPolyhedron.h

inline ArrayRef<int64_t> getInequality(unsigned idx) const {
     return inequalities.getRow(idx);
   }
 
   /// Insert `num` identifiers of the specified kind at position `pos`.
   /// Positions are relative to the kind of identifier. The coefficient columns
   /// corresponding to the added identifiers are initialized to zero. Return the
   /// absolute column position (i.e., not relative to the kind of identifier)
   /// of the first added identifier.
-  unsigned insertDimId(unsigned pos, unsigned num = 1);
-  unsigned insertSymbolId(unsigned pos, unsigned num = 1);
-  unsigned insertLocalId(unsigned pos, unsigned num = 1);
   unsigned insertId(IdKind kind, unsigned pos, unsigned num = 1) override;
 
   /// Append `num` identifiers of the specified kind after the last identifier.
-  /// of that kind. Return the position of the first appended column. The
-  /// coefficient columns corresponding to the added identifiers are initialized
-  /// to zero.
-  unsigned appendDimId(unsigned num = 1);
-  unsigned appendSymbolId(unsigned num = 1);
-  unsigned appendLocalId(unsigned num = 1);
+  /// of that kind. Return the position of the first appended column relative to
+  /// the kind of identifier. The coefficient columns corresponding to the added
+  /// identifiers are initialized to zero.
+  unsigned appendId(IdKind kind, unsigned num = 1);
 
   /// Adds an inequality (>= 0) from the coefficients specified in `inEq`.
   void addInequality(ArrayRef<int64_t> inEq);
   /// Adds an equality from the coefficients specified in `eq`.
   void addEquality(ArrayRef<int64_t> eq);
 
   /// Eliminate the `posB^th` local identifier, replacing every instance of it
   /// with the `posA^th` local identifier. This should be used when the two

mlir/include/mlir/Analysis/Presburger/PresburgerSpace.h

 #include "llvm/Support/ErrorHandling.h"
 #include "llvm/Support/raw_ostream.h"
 
 namespace mlir {
 namespace presburger {
 
 class PresburgerLocalSpace;
 
+/// Kind of identifier. Implementation wise SetDims are treated as Range
+/// ids, and spaces with no distinction between dimension ids are treated
+/// as relations with zero domain ids.
+enum IdKind { Symbol, Local, Domain, Range, SetDim = Range };
+
 /// PresburgerSpace is the space of all possible values of a tuple of integer
 /// valued variables/identifiers. Each identifier has one of the three types:
 ///
 /// Dimension: Ordinary variables over which the space is represented.
 ///
 /// Symbol: Symbol identifiers correspond to fixed but unknown values.
 /// Mathematically, a space with symbolic identifiers is like a
 /// family of spaces indexed by the symbolic identifiers.

 ///       [SetDim, Symbol, Locals]
 ///
 /// PresburgerSpace does not allow identifiers of kind Local. See
 /// PresburgerLocalSpace for an extension that does allow local identifiers.
 class PresburgerSpace {
   friend PresburgerLocalSpace;
 
 public:
-  /// Kind of identifier. Implementation wise SetDims are treated as Range
-  /// ids, and spaces with no distinction between dimension ids are treated
-  /// as relations with zero domain ids.
-  enum IdKind { Symbol, Local, Domain, Range, SetDim = Range };
+  using IdKind = IdKind;
 
   static PresburgerSpace getRelationSpace(unsigned numDomain, unsigned numRange,
                                           unsigned numSymbols);
 
   static PresburgerSpace getSetSpace(unsigned numDims, unsigned numSymbols);
 
   virtual ~PresburgerSpace() = default;
 

mlir/include/mlir/Dialect/Affine/Analysis/AffineStructures.h

public:
 
   static bool classof(const IntegerPolyhedron *cst) {
     return cst->getKind() == Kind::FlatAffineConstraints;
   }
 
   // Clones this object.
   std::unique_ptr<FlatAffineConstraints> clone() const;
 
+  /// Insert `num` identifiers of the specified kind at position `pos`.
+  /// Positions are relative to the kind of identifier. The coefficient columns
+  /// corresponding to the added identifiers are initialized to zero. Return the
+  /// absolute column position (i.e., not relative to the kind of identifier)
+  /// of the first added identifier.
+  unsigned insertDimId(unsigned pos, unsigned num = 1) {
+    return insertId(IdKind::SetDim, pos, num);
+  }
+  unsigned insertSymbolId(unsigned pos, unsigned num = 1) {
+    return insertId(IdKind::Symbol, pos, num);
+  }
+  unsigned insertLocalId(unsigned pos, unsigned num = 1) {
+    return insertId(IdKind::Local, pos, num);
+  }
+
+  /// Append `num` identifiers of the specified kind after the last identifier.
+  /// of that kind. Return the position of the first appended column relative to
+  /// the kind of identifier. The coefficient columns corresponding to the added
+  /// identifiers are initialized to zero.
+  unsigned appendDimId(unsigned num = 1) {
+    return appendId(IdKind::SetDim, num);
+  }
+  unsigned appendSymbolId(unsigned num = 1) {
+    return appendId(IdKind::Symbol, num);
+  }
+  unsigned appendLocalId(unsigned num = 1) {
+    return appendId(IdKind::Local, num);
+  }
+
   /// Adds a bound for the identifier at the specified position with constraints
   /// being drawn from the specified bound map. In case of an EQ bound, the
   /// bound map is expected to have exactly one result. In case of a LB/UB, the
   /// bound map may have more than one result, for each of which an inequality
   /// is added.
   /// Note: The dimensions/symbols of this FlatAffineConstraints must match the
   /// dimensions/symbols of the affine map.
   LogicalResult addBound(BoundType type, unsigned pos, AffineMap boundMap);

mlir/lib/Analysis/Presburger/IntegerPolyhedron.cpp

IntegerPolyhedron::findIntegerLexMin() const {
   // will be minimized last in the lexmin. So simply truncating out the locals
   // from the end of the answer gives the desired lexmin over the dimensions.
   assert(maybeLexMin->size() == getNumIds() &&
          "Incorrect number of vars in lexMin!");
   maybeLexMin->resize(getNumDimAndSymbolIds());
   return maybeLexMin;
 }
 
-unsigned IntegerPolyhedron::insertDimId(unsigned pos, unsigned num) {
-  return insertId(IdKind::SetDim, pos, num);
-}
-
-unsigned IntegerPolyhedron::insertSymbolId(unsigned pos, unsigned num) {
-  return insertId(IdKind::Symbol, pos, num);
-}
-
-unsigned IntegerPolyhedron::insertLocalId(unsigned pos, unsigned num) {
-  return insertId(IdKind::Local, pos, num);
-}
-
 unsigned IntegerPolyhedron::insertId(IdKind kind, unsigned pos, unsigned num) {
   assert(pos <= getNumIdKind(kind));
 
   unsigned insertPos = PresburgerLocalSpace::insertId(kind, pos, num);
   inequalities.insertColumns(insertPos, num);
   equalities.insertColumns(insertPos, num);
   return insertPos;
 }
 
-unsigned IntegerPolyhedron::appendDimId(unsigned num) {
-  unsigned pos = getNumDimIds();
-  insertId(IdKind::SetDim, pos, num);
-  return pos;
-}
-
-unsigned IntegerPolyhedron::appendSymbolId(unsigned num) {
-  unsigned pos = getNumSymbolIds();
-  insertId(IdKind::Symbol, pos, num);
-  return pos;
-}
-
-unsigned IntegerPolyhedron::appendLocalId(unsigned num) {
-  unsigned pos = getNumLocalIds();
-  insertId(IdKind::Local, pos, num);
-  return pos;
+unsigned IntegerPolyhedron::appendId(IdKind kind, unsigned num) {
+  unsigned pos = getNumIdKind(kind);
+  return insertId(kind, pos, num);
 }
 
 void IntegerPolyhedron::addEquality(ArrayRef<int64_t> eq) {
   assert(eq.size() == getNumCols());
   unsigned row = equalities.appendExtraRow();
   for (unsigned i = 0, e = eq.size(); i < e; ++i)
     equalities(row, i) = eq[i];
 }
void IntegerPolyhedron::mergeLocalIds(IntegerPolyhedron &other) {
 
   IntegerPolyhedron &polyA = *this;
   IntegerPolyhedron &polyB = other;
 
   // Merge local ids of polyA and polyB without using division information,
   // i.e. append local ids of `polyB` to `polyA` and insert local ids of `polyA`
   // to `polyB` at start of its local ids.
   unsigned initLocals = polyA.getNumLocalIds();
-  insertLocalId(polyA.getNumLocalIds(), polyB.getNumLocalIds());
-  polyB.insertLocalId(0, initLocals);
+  insertId(Local, polyA.getNumLocalIds(), polyB.getNumLocalIds());
+  polyB.insertId(Local, 0, initLocals);
 
   // Get division representations from each poly.
   std::vector<SmallVector<int64_t, 8>> divsA, divsB;
   SmallVector<unsigned, 4> denomsA, denomsB;
   polyA.getLocalReprs(divsA, denomsA);
   polyB.getLocalReprs(divsB, denomsB);
 
   // Copy division information for polyB into `divsA` and `denomsA`, so that
void IntegerPolyhedron::convertDimToLocal(unsigned dimStart,
   assert(dimLimit <= getNumDimIds() && "Invalid dim pos range");
 
   if (dimStart >= dimLimit)
     return;
 
   // Append new local variables corresponding to the dimensions to be converted.
   unsigned convertCount = dimLimit - dimStart;
   unsigned newLocalIdStart = getNumIds();
-  appendLocalId(convertCount);
+  appendId(Local, convertCount);
 
   // Swap the new local variables with dimensions.
   for (unsigned i = 0; i < convertCount; ++i)
     swapId(i + dimStart, i + newLocalIdStart);
 
   // Remove dimensions converted to local variables.
   removeIdRange(dimStart, dimLimit);
 }

 /// respect to a positive constant 'divisor'. Two constraints are added to the
 /// system to capture equivalence with the floordiv.
 ///      q = expr floordiv c    <=>   c*q <= expr <= c*q + c - 1.
 void IntegerPolyhedron::addLocalFloorDiv(ArrayRef<int64_t> dividend,
                                          int64_t divisor) {
   assert(dividend.size() == getNumCols() && "incorrect dividend size");
   assert(divisor > 0 && "positive divisor expected");
 
-  appendLocalId();
+  appendId(Local);
 
   // Add two constraints for this new identifier 'q'.
   SmallVector<int64_t, 8> bound(dividend.size() + 1);
 
   // dividend - q * divisor >= 0
   std::copy(dividend.begin(), dividend.begin() + dividend.size() - 1,
             bound.begin());
   bound.back() = dividend.back();

mlir/unittests/Analysis/Presburger/IntegerPolyhedronTest.cpp

TEST(IntegerPolyhedronTest, mergeDivisionsSimple) {
     poly1.addLocalFloorDiv({1, 0, 0}, 2); // y = [x / 2].
     poly1.addEquality({1, 0, -3, 0});     // x = 3y.
     poly1.addInequality({1, 1, 0, 1});    // x + z + 1 >= 0.
 
     // (x) : (exists y = [x / 2], z : x = 5y).
     IntegerPolyhedron poly2(1);
     poly2.addLocalFloorDiv({1, 0}, 2); // y = [x / 2].
     poly2.addEquality({1, -5, 0});     // x = 5y.
-    poly2.appendLocalId();             // Add local id z.
+    poly2.appendId(Local);             // Add local id z.
 
     poly1.mergeLocalIds(poly2);
 
     // Local space should be same.
     EXPECT_EQ(poly1.getNumLocalIds(), poly2.getNumLocalIds());
 
     // 1 division should be matched + 2 unmatched local ids.
     EXPECT_EQ(poly1.getNumLocalIds(), 3u);
TEST(IntegerPolyhedronTest, mergeDivisionsSimple) {
     poly1.addLocalFloorDiv({3, 0, 0}, 6); // y = [3x / 6] -> [x/2].
     poly1.addEquality({1, 0, -3, 0});     // x = 3z.
     poly1.addInequality({1, 1, 0, 1});    // x + y + 1 >= 0.
 
     // (x) : (exists y = [x / 2], z : x = 5y).
     IntegerPolyhedron poly2(1);
     poly2.addLocalFloorDiv({1, 0}, 2); // y = [x / 2].
     poly2.addEquality({1, -5, 0});     // x = 5y.
-    poly2.appendLocalId();             // Add local id z.
+    poly2.appendId(Local);             // Add local id z.
 
     poly1.mergeLocalIds(poly2);
 
     // Local space should be same.
     EXPECT_EQ(poly1.getNumLocalIds(), poly2.getNumLocalIds());
 
     // One division should be matched + 2 unmatched local ids.
     EXPECT_EQ(poly1.getNumLocalIds(), 3u);