[MLIR][Presburger] PWMAFunction::valueAt: support local ids

[arjunp]

Summary

Add a baseline implementation of support for local ids for PWMAFunction::valueAt. This can be made more efficient later if needed by handling locals with known div representations separately.

[Groverkss]

Please add a small patch description specifying the reasoning.

mlir/lib/Analysis/Presburger/IntegerRelation.cpp

789–790

I think this comment should be done in setAndEliminate, not here.

[arjunp]

Elaborate on doc

mlir/lib/Analysis/Presburger/IntegerRelation.cpp

789–790

Elaborated further what I mean.

[Groverkss]

LGTM.

mlir/include/mlir/Analysis/Presburger/IntegerRelation.h

public:
   Optional<SmallVector<int64_t, 8>> findIntegerSample() const;
 
   /// Compute an overapproximation of the number of integer points in the
   /// relation. Symbol ids are currently not supported. If the computed
   /// overapproximation is infinite, an empty optional is returned.
   Optional<uint64_t> computeVolume() const;
 
   /// Returns true if the given point satisfies the constraints, or false
-  /// otherwise.
-  ///
-  /// Note: currently, if the relation contains local ids, the values of
-  /// the local ids must also be provided.
+  /// otherwise. Takes the values of all ids including locals.
   bool containsPoint(ArrayRef<int64_t> point) const;
+  /// Given the values of non-local ids, return a satisfying assignment to the
+  /// local if one exists, or an empty optional otherwise.
+  Optional<SmallVector<int64_t, 8>>
+  containsPointNoLocal(ArrayRef<int64_t> point) const;
 
   /// Find equality and pairs of inequality contraints identified by their
   /// position indices, using which an explicit representation for each local
   /// variable can be computed. The indices of the constraints are stored in
   /// `MaybeLocalRepr` struct. If no such pair can be found, the kind attribute
   /// in `MaybeLocalRepr` is set to None.
   ///
   /// The dividends of the explicit representations are stored in `dividends`

mlir/include/mlir/Analysis/Presburger/PWMAFunction.h

 ///
 /// In this way every point in the polyhedron has a tuple of integers associated
 /// with it. If the integer polyhedron has local ids, then the output
 /// expressions can use them as well. The output expressions are represented as
 /// a matrix with one row for every element in the output vector one column for
 /// each id, and an extra column at the end for the constant term.
 ///
 /// Checking equality of two such functions is supported, as well as finding the
-/// value of the function at a specified point. Note that local ids in the
-/// domain are not yet supported for finding the value at a point.
+/// value of the function at a specified point.
 class MultiAffineFunction : protected IntegerPolyhedron {
 public:
   /// We use protected inheritance to avoid inheriting the whole public
   /// interface of IntegerPolyhedron. These using declarations explicitly make
   /// only the relevant functions part of the public interface.
   using IntegerPolyhedron::getNumDimAndSymbolIds;
   using IntegerPolyhedron::getNumDimIds;
   using IntegerPolyhedron::getNumIds;
public:
 
   /// Return whether the `this` and `other` are equal. This is the case if
   /// they lie in the same space, i.e. have the same dimensions, and their
   /// domains are identical and their outputs are equal on their domain.
   bool isEqual(const MultiAffineFunction &other) const;
 
   /// Get the value of the function at the specified point. If the point lies
   /// outside the domain, an empty optional is returned.
-  ///
-  /// Note: domains with local ids are not yet supported, and will assert-fail.
   Optional<SmallVector<int64_t, 8>> valueAt(ArrayRef<int64_t> point) const;
 
   void print(raw_ostream &os) const;
   void dump() const;
 
 private:
   /// The function's output is a tuple of integers, with the ith element of the
   /// tuple defined by the affine expression given by the ith row of this output

 /// Note that the domains all have to be *disjoint*. Otherwise, the behaviour of
 /// this class is undefined. The domains need not cover all possible points;
 /// this represents a partial function and so could be undefined at some points.
 ///
 /// As in PresburgerSets, the input ids are partitioned into dimension ids and
 /// symbolic ids.
 ///
 /// Support is provided to compare equality of two such functions as well as
-/// finding the value of the function at a point. Note that local ids in the
-/// piece are not supported for the latter.
+/// finding the value of the function at a point.
 class PWMAFunction : public PresburgerSpace {
 public:
   PWMAFunction(unsigned numDims, unsigned numSymbols, unsigned numOutputs)
       : PresburgerSpace(/*numDomain=*/0, /*numRange=*/numDims, numSymbols),
         numOutputs(numOutputs) {
     assert(numOutputs >= 1 && "The function must output something!");
   }
 
   void addPiece(const MultiAffineFunction &piece);
   void addPiece(const IntegerPolyhedron &domain, const Matrix &output);
 
   const MultiAffineFunction &getPiece(unsigned i) const { return pieces[i]; }
   unsigned getNumPieces() const { return pieces.size(); }
   unsigned getNumOutputs() const { return numOutputs; }
   unsigned getNumInputs() const { return getNumIds(); }
   MultiAffineFunction &getPiece(unsigned i) { return pieces[i]; }
 
   /// Return the domain of this piece-wise MultiAffineFunction. This is the
   /// union of the domains of all the pieces.
   PresburgerSet getDomain() const;
 
   /// Return the value at the specified point and an empty optional if the
   /// point does not lie in the domain.
-  ///
-  /// Note: domains with local ids are not yet supported, and will assert-fail.
   Optional<SmallVector<int64_t, 8>> valueAt(ArrayRef<int64_t> point) const;
 
   /// Return whether `this` and `other` are equal as PWMAFunctions, i.e. whether
   /// they have the same dimensions, the same domain and they take the same
   /// value at every point in the domain.
   bool isEqual(const PWMAFunction &other) const;
 
   void print(raw_ostream &os) const;

mlir/lib/Analysis/Presburger/IntegerRelation.cpp

bool IntegerRelation::containsPoint(ArrayRef<int64_t> point) const {
   }
   for (unsigned i = 0, e = getNumInequalities(); i < e; ++i) {
     if (valueAt(getInequality(i), point) < 0)
       return false;
   }
   return true;
 }
 
+/// Just substitute the values given and check if an integer sample exists for
+/// the local ids.
+///
+/// TODO: this could be made more efficient by handling divisions separately.
+/// Instead of finding an integer sample over all the locals, we can first
+/// compute the values of the locals that have division representations and
+/// only use the integer emptiness check for the locals that don't have this.
+/// Handling this correctly requires ordering the divs, though.
+Optional<SmallVector<int64_t, 8>>
+IntegerRelation::containsPointNoLocal(ArrayRef<int64_t> point) const {
+  assert(point.size() == getNumIds() - getNumLocalIds() &&
+         "Point should contain all ids except locals!");
+  assert(getIdKindOffset(IdKind::Local) == getNumIds() - getNumLocalIds() &&
+         "This function depends on locals being stored last!");
+  IntegerRelation copy = *this;
+  copy.setAndEliminate(0, point);
+  return copy.findIntegerSample();
+}
+
 void IntegerRelation::getLocalReprs(std::vector<MaybeLocalRepr> &repr) const {
   std::vector<SmallVector<int64_t, 8>> dividends(getNumLocalIds());
   SmallVector<unsigned, 4> denominators(getNumLocalIds());
   getLocalReprs(dividends, denominators, repr);
 }
 
 void IntegerRelation::getLocalReprs(
     std::vector<SmallVector<int64_t, 8>> &dividends,

mlir/lib/Analysis/Presburger/PWMAFunction.cpp

PresburgerSet domain =
       PresburgerSet::getEmpty(getNumDimIds(), getNumSymbolIds());
   for (const MultiAffineFunction &piece : pieces)
     domain.unionInPlace(piece.getDomain());
   return domain;
 }
 
 Optional<SmallVector<int64_t, 8>>
 MultiAffineFunction::valueAt(ArrayRef<int64_t> point) const {
-  assert(getNumLocalIds() == 0 && "Local ids are not yet supported!");
-  assert(point.size() == getNumIds() && "Point has incorrect dimensionality!");
+  assert(point.size() == getNumDimAndSymbolIds() &&
+         "Point has incorrect dimensionality!");
 
-  if (!getDomain().containsPoint(point))
+  Optional<SmallVector<int64_t, 8>> maybeLocalValues =
+      getDomain().containsPointNoLocal(point);
+  if (!maybeLocalValues)
     return {};
 
   // The point lies in the domain, so we need to compute the output value.
+  SmallVector<int64_t, 8> pointHomogenous{llvm::to_vector(point)};
+  // The given point didn't include the values of locals which the output is a
+  // function of; we have computed one possible set of values and use them
+  // here. The function is not allowed to have local ids that take more than
+  // one possible value.
+  pointHomogenous.append(*maybeLocalValues);
   // The matrix `output` has an affine expression in the ith row, corresponding
   // to the expression for the ith value in the output vector. The last column
   // of the matrix contains the constant term. Let v be the input point with
   // a 1 appended at the end. We can see that output * v gives the desired
   // output vector.
-  SmallVector<int64_t, 8> pointHomogenous{llvm::to_vector(point)};
   pointHomogenous.push_back(1);
   SmallVector<int64_t, 8> result =
       output.postMultiplyWithColumn(pointHomogenous);
   assert(result.size() == getNumOutputs());
   return result;
 }
 
 Optional<SmallVector<int64_t, 8>>

mlir/unittests/Analysis/Presburger/IntegerPolyhedronTest.cpp

expectComputedVolumeIsValidOverapprox(
                 "y - 2*(y floordiv 2) == 0)"),
       /*trueVolume=*/61ull, /*resultBound=*/441ull);
 
   // Unbounded polytope.
   expectComputedVolumeIsValidOverapprox(
       parsePoly("(x, y) : (2*x - y >= 0, y - 3*x >= 0)"),
       /*trueVolume=*/{}, /*resultBound=*/{});
 }
+
+TEST(IntegerPolyhedronTest, containsPointNoLocal) {
+  IntegerPolyhedron poly1 = parsePoly("(x) : ((x floordiv 2) - x == 0)");
+  EXPECT_TRUE(poly1.containsPointNoLocal({0}));
+  EXPECT_FALSE(poly1.containsPointNoLocal({1}));
+
+  IntegerPolyhedron poly2 = parsePoly(
+      "(x) : (x - 2*(x floordiv 2) == 0, x - 4*(x floordiv 4) - 2 == 0)");
+  EXPECT_TRUE(poly2.containsPointNoLocal({6}));
+  EXPECT_FALSE(poly2.containsPointNoLocal({4}));
+
+  IntegerPolyhedron poly3 = parsePoly("(x, y) : (2*x - y >= 0, y - 3*x >= 0)");
+  EXPECT_TRUE(poly3.containsPointNoLocal({0, 0}));
+  EXPECT_FALSE(poly3.containsPointNoLocal({1, 0}));
+}

mlir/unittests/Analysis/Presburger/PWMAFunctionTest.cpp

PWMAFunction mul5 = parsePWMAF(
       /*numInputs=*/1, /*numOutputs=*/1,
       {
           {"(x) : (x - 5*(x floordiv 5) == 0)", {{0, 26, 0}}}, // 26(x/5).
       });
   EXPECT_FALSE(mul2AndMul3.isEqual(mul5));
 }
 
 TEST(PWMAFunction, valueAt) {
-  PWMAFunction nonNegPWAF = parsePWMAF(
+  PWMAFunction nonNegPWMAF = parsePWMAF(
       /*numInputs=*/2, /*numOutputs=*/2,
       {
           {"(x, y) : (x >= 0)", {{1, 2, 3}, {3, 4, 5}}}, // (x, y).
           {"(x, y) : (y >= 0, -x - 1 >= 0)", {{-1, 2, 3}, {-3, 4, 5}}} // (x, y)
       });
-  EXPECT_THAT(*nonNegPWAF.valueAt({2, 3}), ElementsAre(11, 23));
-  EXPECT_THAT(*nonNegPWAF.valueAt({-2, 3}), ElementsAre(11, 23));
-  EXPECT_THAT(*nonNegPWAF.valueAt({2, -3}), ElementsAre(-1, -1));
-  EXPECT_FALSE(nonNegPWAF.valueAt({-2, -3}).hasValue());
+  EXPECT_THAT(*nonNegPWMAF.valueAt({2, 3}), ElementsAre(11, 23));
+  EXPECT_THAT(*nonNegPWMAF.valueAt({-2, 3}), ElementsAre(11, 23));
+  EXPECT_THAT(*nonNegPWMAF.valueAt({2, -3}), ElementsAre(-1, -1));
+  EXPECT_FALSE(nonNegPWMAF.valueAt({-2, -3}).hasValue());
+
+  PWMAFunction divPWMAF = parsePWMAF(
+      /*numInputs=*/2, /*numOutputs=*/2,
+      {
+          {"(x, y) : (x >= 0, x - 2*(x floordiv 2) == 0)",
+           {{0, 2, 1, 3}, {0, 4, 3, 5}}}, // (x, y).
+          {"(x, y) : (y >= 0, -x - 1 >= 0)", {{-1, 2, 3}, {-3, 4, 5}}} // (x, y)
+      });
+  EXPECT_THAT(*divPWMAF.valueAt({4, 3}), ElementsAre(11, 23));
+  EXPECT_THAT(*divPWMAF.valueAt({4, -3}), ElementsAre(-1, -1));
+  EXPECT_FALSE(divPWMAF.valueAt({3, 3}).hasValue());
+  EXPECT_FALSE(divPWMAF.valueAt({3, -3}).hasValue());
+
+  EXPECT_THAT(*divPWMAF.valueAt({-2, 3}), ElementsAre(11, 23));
+  EXPECT_FALSE(divPWMAF.valueAt({-2, -3}).hasValue());
 }
 
 TEST(PWMAFunction, removeIdRangeRegressionTest) {
   PWMAFunction pwafA = parsePWMAF(
       /*numInputs=*/2, /*numOutputs=*/1,
       {
           {"(x, y) : (x == 0, y == 0, x - 2*(x floordiv 2) == 0, y - 2*(y "
            "floordiv 2) == 0)",