[MLIR] AffineStructures::removeIdRange: support specifying a range within an IdKind

[Groverkss]

Added some comments.

mlir/include/mlir/Analysis/AffineStructures.h

285–287

Please use triple / docs comments here.

mlir/lib/Analysis/AffineStructures.cpp

347

nit: You could use getNumDimAndSymbolds() here for more readability.

386

An assert message here would be useful.

[arjunp]

Address Kunwar's comments.

[arjunp]

Thanks for the review!

mlir/include/mlir/Analysis/AffineStructures.h

285–287

Thanks. I wish this was checked by the linter.

mlir/lib/Analysis/AffineStructures.cpp

386

I took a look at the existing removeIdRange behaviour again and it seems to support idStart >= idLimit cases, it just returns without doing anything. So I'm removing the assert.

[Groverkss]

LGTM.

[grosser]

This looks good from my side as well. Given that these are rather straightforward extensions to the API as well as much appreciated polishing of some simple functions, you can probably go ahead with committing them.

mlir/include/mlir/Analysis/AffineStructures.h

public:
   /// Projects out (aka eliminates) `num` identifiers starting at position
   /// `pos`. The resulting constraint system is the shadow along the dimensions
   /// that still exist. This method may not always be integer exact.
   // TODO: deal with integer exactness when necessary - can return a value to
   // mark exactness for example.
   void projectOut(unsigned pos, unsigned num);
   inline void projectOut(unsigned pos) { return projectOut(pos, 1); }
 
+  /// Removes identifiers of the specified kind with the specified pos (or
+  /// within the specified range) from the system. The specified location is
+  /// relative to the first identifier of the specified kind.
+  void removeId(IdKind kind, unsigned pos);
+  void removeIdRange(IdKind kind, unsigned idStart, unsigned idLimit);
+
   /// Removes the specified identifier from the system.
   void removeId(unsigned pos);
 
   void removeEquality(unsigned pos);
   void removeInequality(unsigned pos);
 
   /// Sets the `values.size()` identifiers starting at `po`s to the specified
   /// values and removes them.
public:
 
   /// Removes all equalities and inequalities.
   void clearConstraints();
 
   void print(raw_ostream &os) const;
   void dump() const;
 
 protected:
+  /// Return the index at which the specified kind of id starts.
+  unsigned getIdKindOffset(IdKind kind) const;
+
+  /// Assert that `value` is at most the number of ids of the specified kind.
+  void assertAtMostNumIdKind(unsigned value, IdKind kind) const;
+
   /// Returns false if the fields corresponding to various identifier counts, or
   /// equality/inequality buffer sizes aren't consistent; true otherwise. This
   /// is meant to be used within an assert internally.
   virtual bool hasConsistentState() const;
 
   /// Checks all rows of equality/inequality constraints for trivial
   /// contradictions (for example: 1 == 0, 0 >= 1), which may have surfaced
   /// after elimination. Returns true if an invalid constraint is found;

mlir/lib/Analysis/AffineStructures.cpp

 }
 
 unsigned FlatAffineConstraints::insertLocalId(unsigned pos, unsigned num) {
   return insertId(IdKind::Local, pos, num);
 }
 
 unsigned FlatAffineConstraints::insertId(IdKind kind, unsigned pos,
                                          unsigned num) {
-  if (kind == IdKind::Dimension)
-    assert(pos <= getNumDimIds());
-  else if (kind == IdKind::Symbol)
-    assert(pos <= getNumSymbolIds());
-  else
-    assert(pos <= getNumLocalIds());
+  assertAtMostNumIdKind(pos, kind);
 
-  unsigned absolutePos;
-  if (kind == IdKind::Dimension) {
-    absolutePos = pos;
+  unsigned absolutePos = getIdKindOffset(kind) + pos;
+  if (kind == IdKind::Dimension)
     numDims += num;
-  } else if (kind == IdKind::Symbol) {
-    absolutePos = pos + getNumDimIds();
+  else if (kind == IdKind::Symbol)
     numSymbols += num;
-  } else {
-    absolutePos = pos + getNumDimIds() + getNumSymbolIds();
-  }
   numIds += num;
 
   inequalities.insertColumns(absolutePos, num);
   equalities.insertColumns(absolutePos, num);
 
   return absolutePos;
 }
 
+void FlatAffineConstraints::assertAtMostNumIdKind(unsigned val,
+                                                  IdKind kind) const {
+  if (kind == IdKind::Dimension)
+    assert(val <= getNumDimIds());
+  else if (kind == IdKind::Symbol)
+    assert(val <= getNumSymbolIds());
+  else if (kind == IdKind::Local)
+    assert(val <= getNumLocalIds());
+  else
+    llvm_unreachable("IdKind expected to be Dimension, Symbol or Local!");
+}
+
+unsigned FlatAffineConstraints::getIdKindOffset(IdKind kind) const {
+  if (kind == IdKind::Dimension)
+    return 0;
+  if (kind == IdKind::Symbol)
+    return getNumDimIds();
+  if (kind == IdKind::Local)
+    return getNumDimAndSymbolIds();
+  llvm_unreachable("IdKind expected to be Dimension, Symbol or Local!");
+}
+
 unsigned FlatAffineValueConstraints::insertId(IdKind kind, unsigned pos,
                                               unsigned num) {
   unsigned absolutePos = FlatAffineConstraints::insertId(kind, pos, num);
   values.insert(values.begin() + absolutePos, num, None);
   assert(values.size() == getNumIds());
   return absolutePos;
 }
 

 }
 
 bool FlatAffineValueConstraints::hasValues() const {
   return llvm::find_if(values, [](Optional<Value> id) {
            return id.hasValue();
          }) != values.end();
 }
 
+void FlatAffineConstraints::removeId(IdKind kind, unsigned pos) {
+  removeIdRange(kind, pos, pos + 1);
+}
+
+void FlatAffineConstraints::removeIdRange(IdKind kind, unsigned idStart,
+                                          unsigned idLimit) {
+  assertAtMostNumIdKind(idLimit, kind);
+  removeIdRange(getIdKindOffset(kind) + idStart,
+                getIdKindOffset(kind) + idLimit);
+}
+
 /// Checks if two constraint systems are in the same space, i.e., if they are
 /// associated with the same set of identifiers, appearing in the same order.
 static bool areIdsAligned(const FlatAffineValueConstraints &a,
                           const FlatAffineValueConstraints &b) {
   return a.getNumDimIds() == b.getNumDimIds() &&
          a.getNumSymbolIds() == b.getNumSymbolIds() &&
          a.getNumIds() == b.getNumIds() &&
          a.getMaybeValues().equals(b.getMaybeValues());

mlir/unittests/Analysis/AffineStructuresTest.cpp

std::vector<SmallVector<int64_t, 8>> divisions = {{0, -2, 7, 2, 0, 0, 0, 10},
                                                     {3, 0, 9, 2, 2, 0, 0, 10},
                                                     {0, 1, -123, 2, 0, -4, 10}};
   SmallVector<int64_t, 8> denoms = {3, 5, 3};
 
   // Check if floordivs which may depend on other floordivs can be computed.
   checkDivisionRepresentation(fac, divisions, denoms);
 }
 
+TEST(FlatAffineConstraintsTest, removeIdRange) {
+  FlatAffineConstraints fac(3, 2, 1);
+
+  fac.addInequality({10, 11, 12, 20, 21, 30, 40});
+  fac.removeId(FlatAffineConstraints::IdKind::Symbol, 1);
+  EXPECT_THAT(fac.getInequality(0),
+              testing::ElementsAre(10, 11, 12, 20, 30, 40));
+
+  fac.removeIdRange(FlatAffineConstraints::IdKind::Dimension, 0, 2);
+  EXPECT_THAT(fac.getInequality(0), testing::ElementsAre(12, 20, 30, 40));
+
+  fac.removeIdRange(FlatAffineConstraints::IdKind::Local, 1, 1);
+  EXPECT_THAT(fac.getInequality(0), testing::ElementsAre(12, 20, 30, 40));
+
+  fac.removeIdRange(FlatAffineConstraints::IdKind::Local, 0, 1);
+  EXPECT_THAT(fac.getInequality(0), testing::ElementsAre(12, 20, 40));
+}
+
 } // namespace mlir