[MLIR][Presburger] Add support for piece-wise multi-affine functions

[arjunp]

Summary

Add the class MultiAffineFunction which represents functions whose domain is an
IntegerPolyhedron and which produce an output given by a tuple of affine
expressions in the IntegerPolyhedron's ids.

Also add support for piece-wise MultiAffineFunctions, which are defined on a
union of IntegerPolyhedrons, and may have different output affine expressions
on each IntegerPolyhedron. Thus the function is affine on each individual
IntegerPolyhedron piece in the domain.

This is part of a series of patches leading up to parametric integer programming.

Depends on D118778.

[arjunp]

Rebase and add forgotten classof

[Groverkss]

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

46

I think there is a better way to do this. The only way identifiers can be moved/removed/inserted in IntegerPolyhedron is through 3 functions:

1. IntegerPolyhedron::insertId
2. IntegerPolyhedron::removeIdRange
3. IntegerPolyhedron::swapId

(I'm very sure that these are the only three functions that do this. In case there are any other, we should make them use these functions as these are the only operations required to modify the identifier space)

You could make these functions virtual in IntegerPolyhedron and just override these functions here to carry over the output information. After this, you can just inherit from public IntegerPolyhedron this and it should carry over the information properly for every method.

173–177

I think you can make these inline.

183–193

I don't think these should be in the header file. The only implementations that should be in the header files are functions that are supposed to be inlined.

mlir/lib/Analysis/Presburger/IntegerPolyhedron.cpp

1073

You can use IntegerPolyhedron::getIdKindOffset here.

mlir/lib/Analysis/Presburger/PWMAFunction.cpp

15–16

Please make this function static. Also a doc comment explaining what it does would be nice.

17

Please add a string to explain what this assert does.

[arjunp]

Thanks for the review!

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

46

The reason I made the inheritance protected is because I don't want this class to inherit the public interface of IntegerPolyhedron, since many of its functions probably don't make sense here.

This class already works basically in the way that you're suggesting. I didn't override removeIdRange earlier though, which I have done now.

173–177

They are already implicitly inline right?

[arjunp]

Address Kunwar's comments

[Groverkss]

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

46

Right. I now understand what you want to do better.

59–71

You can replace all of these except getNumInputs and getNumOutputs by:

using IntegerPolyhedron::getNumIds;
using IntegerPolyhedron::getNumSymbolIds;
.
.

[arjunp]

Incorporate Kunwar's suggestion and fix the build

[arjunp]

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

59–71

Good point! Done.

[Groverkss]

mlir/lib/Analysis/Presburger/PWMAFunction.cpp

112

The usage of the object this and the word "this" is a bit confusing. Putting backquotes should help. Please fix in other places also in this function implementation.

157–159

Please put a string explaining what these asserts do with the assert.

mlir/unittests/Analysis/Presburger/PWMAFunctionTest.cpp

2–14

I think this header needs to be changed.

63

This comment seems to be unrelated here and seems to be unfinished.

67

This seems incomplete to me.

166

I think this should be at top of file.

[arjunp]

Address Kunwar's comments

[Groverkss]

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

158

Is this implemented? Also, a doc comment would be good.

159

Doc comment, please.

161

Same for this. Is this implemented? Also, a doc comment would be good.

163–164

Doc comment please.

mlir/lib/Analysis/Presburger/PWMAFunction.cpp

157–160

Can you add an assert which just checks if the newly added domain intersects (non-empty intersection) with the existing domain? While the check is expensive, it should be fine inside an assert.

mlir/unittests/Analysis/Presburger/PWMAFunctionTest.cpp

10–14

I think this doc comment is for PresburgerSet. Please change it.

[arjunp]

Address Kunwar's comments

[arjunp]

Fix tests

[Groverkss]

LGTM. Thank you for this patch!

mlir/lib/Analysis/Presburger/PWMAFunction.cpp

127–129

I think you should add a comment here that you are adding constraints such that the output of both functions should be equal.

[arjunp]

Address kunwar's comment

[arjunp]

Rebase

[arjunp]

Rebase

[arjunp]

@Groverkss thanks for the review!

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

 /// result of projecting out `q` from the previous example, i.e. we obtained {2,
 /// 4, 6} by projecting out the second dimension from {(2, 1), (4, 2), (6, 2)}.
 class IntegerPolyhedron {
 public:
   /// All derived classes of IntegerPolyhedron.
   enum class Kind {
     FlatAffineConstraints,
     FlatAffineValueConstraints,
+    MultiAffineFunction,
     IntegerPolyhedron
   };
 
   /// Kind of identifier (column).
   enum IdKind { Dimension, Symbol, Local };
 
   /// Constructs a constraint system reserving memory for the specified number
   /// of constraints and identifiers.
public:
   unsigned appendSymbolId(unsigned num = 1);
   unsigned appendLocalId(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
+  /// local variables are known to always take the same values.
+  virtual void eliminateRedundantLocalId(unsigned posA, unsigned posB);
+
   /// 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);
public:
   /// additional processing using Simplex for unbounded sets.
   ///
   /// Returns an integer sample point if one exists, or an empty Optional
   /// otherwise.
   Optional<SmallVector<int64_t, 8>> findIntegerSample() const;
 
   /// Returns true if the given point satisfies the constraints, or false
   /// otherwise.
+  ///
+  /// Note: currently, if the polyhedron contains local ids, the values of
+  /// the local ids must also be provided.
   bool containsPoint(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.
   ///

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

+//===- PWMAFunction.h - MLIR PWMAFunction Class------------------*- C++ -*-===//
+//
+// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
+// See https://llvm.org/LICENSE.txt for license information.
+// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+//
+//===----------------------------------------------------------------------===//
+//
+// Support for piece-wise multi-affine functions. These are functions that are
+// defined on a domain that is a union of IntegerPolyhedrons, and on each domain
+// the value of the function is a tuple of integers, with each value in the
+// tuple being an affine expression in the ids of the IntegerPolyhedron.
+//
+//===----------------------------------------------------------------------===//
+
+#ifndef MLIR_ANALYSIS_PRESBURGER_PWMAFUNCTION_H
+#define MLIR_ANALYSIS_PRESBURGER_PWMAFUNCTION_H
+
+#include "mlir/Analysis/Presburger/IntegerPolyhedron.h"
+#include "mlir/Analysis/Presburger/PresburgerSet.h"
+
+namespace mlir {
+
+/// This class represents a multi-affine function whose domain is given by an
+/// IntegerPolyhedron. This can be thought of as an IntegerPolyhedron with a
+/// tuple of integer values attached to every point in the polyhedron, with the
+/// value of each element of the tuple given by an affine expression in the ids
+/// of the polyhedron. For example we could have the domain
+///
+/// (x, y) : (x >= 5, y >= x)
+///
+/// and a tuple of three integers defined at every point in the polyhedron:
+///
+/// (x, y) -> (x + 2, 2*x - 3y + 5, 2*x + y).
+///
+/// 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.
+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;
+  using IntegerPolyhedron::getNumLocalIds;
+  using IntegerPolyhedron::getNumSymbolIds;
+
+  MultiAffineFunction(const IntegerPolyhedron &domain, const Matrix &output)
+      : IntegerPolyhedron(domain), output(output) {}
+  MultiAffineFunction(const Matrix &output, unsigned numDims,
+                      unsigned numSymbols = 0, unsigned numLocals = 0)
+      : IntegerPolyhedron(numDims, numSymbols, numLocals), output(output) {}
+
+  ~MultiAffineFunction() override = default;
+  Kind getKind() const override { return Kind::MultiAffineFunction; }
+  bool classof(const IntegerPolyhedron *poly) const {
+    return poly->getKind() == Kind::MultiAffineFunction;
+  }
+
+  unsigned getNumInputs() const { return getNumDimAndSymbolIds(); }
+  unsigned getNumOutputs() const { return output.getNumRows(); }
+  bool isConsistent() const { return output.getNumColumns() == numIds + 1; }
+  const IntegerPolyhedron &getDomain() const { return *this; }
+
+  bool hasCompatibleDimensions(const MultiAffineFunction &f) 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 insertId(IdKind kind, unsigned pos, unsigned num = 1) override;
+
+  /// Swap the posA^th identifier with the posB^th identifier.
+  void swapId(unsigned posA, unsigned posB) override;
+
+  /// Remove the specified range of ids.
+  void removeIdRange(unsigned idStart, unsigned idLimit) override;
+
+  /// Eliminate the `posB^th` local identifier, replacing every instance of it
+  /// with the `posA^th` local identifier. This should be used when the two
+  /// local variables are known to always take the same values.
+  void eliminateRedundantLocalId(unsigned posA, unsigned posB) override;
+
+  /// Return whether the outputs of `this` and `other` agree wherever both
+  /// functions are defined, i.e., the outputs should be equal for all points in
+  /// the intersection of the domains.
+  bool isEqualWhereDomainsOverlap(MultiAffineFunction other) const;
+
+  /// 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
+  /// matrix.
+  Matrix output;
+};
+
+/// This class represents a piece-wise MultiAffineFunction. This can be thought
+/// of as a list of MultiAffineFunction with disjoint domains, with each having
+/// their own affine expressions for their output tuples. For example, we could
+/// have a function with two input variables (x, y), defined as
+///
+/// f(x, y) = (2*x + y, y - 4)  if x >= 0, y >= 0
+///         = (-2*x + y, y + 4) if x < 0,  y < 0
+///         = (4, 1)            if x < 0,  y >= 0
+///
+/// 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.
+class PWMAFunction {
+public:
+  PWMAFunction(unsigned numDims, unsigned numSymbols, unsigned numOutputs)
+      : numDims(numDims), numSymbols(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 numDims + numSymbols; }
+  unsigned getNumDimIds() const { return numDims; }
+  unsigned getNumSymbolIds() const { return numSymbols; }
+  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;
+
+  /// Check whether the `this` and the given function have compatible
+  /// dimensions, i.e., the same number of dimension inputs, symbol inputs, and
+  /// outputs.
+  bool hasCompatibleDimensions(const MultiAffineFunction &f) const;
+  bool hasCompatibleDimensions(const PWMAFunction &f) 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;
+  void dump() const;
+
+private:
+  /// The list of pieces in this piece-wise MultiAffineFunction.
+  SmallVector<MultiAffineFunction, 4> pieces;
+
+  /// The number of dimensions ids in the domains.
+  unsigned numDims;
+  /// The number of symbol ids in the domains.
+  unsigned numSymbols;
+  /// The number of output ids.
+  unsigned numOutputs;
+};
+
+} // namespace mlir
+
+#endif // MLIR_ANALYSIS_PRESBURGER_PWMAFUNCTION_H

mlir/lib/Analysis/Presburger/CMakeLists.txt

 add_mlir_library(MLIRPresburger
   IntegerPolyhedron.cpp
   LinearTransform.cpp
   Matrix.cpp
   PresburgerSet.cpp
+  PWMAFunction.cpp
   Simplex.cpp
   Utils.cpp
 
   LINK_LIBS PUBLIC
   MLIRSupport
   )

mlir/lib/Analysis/Presburger/IntegerPolyhedron.cpp

void IntegerPolyhedron::removeRedundantConstraints() {
   for (unsigned r = 0, e = getNumEqualities(); r < e; r++) {
     if (!(simplex.isMarkedRedundant(numIneqs + 2 * r) &&
           simplex.isMarkedRedundant(numIneqs + 2 * r + 1)))
       copyEquality(r, pos++);
   }
   equalities.resizeVertically(pos);
 }
 
-/// Eliminate `pos2^th` local identifier, replacing its every instance with
-/// `pos1^th` local identifier. This function is intended to be used to remove
-/// redundancy when local variables at position `pos1` and `pos2` are restricted
-/// to have the same value.
-static void eliminateRedundantLocalId(IntegerPolyhedron &poly, unsigned pos1,
-                                      unsigned pos2) {
-
-  assert(pos1 < poly.getNumLocalIds() && "Invalid local id position");
-  assert(pos2 < poly.getNumLocalIds() && "Invalid local id position");
-
-  unsigned localOffset = poly.getNumDimAndSymbolIds();
-  pos1 += localOffset;
-  pos2 += localOffset;
-  for (unsigned i = 0, e = poly.getNumInequalities(); i < e; ++i)
-    poly.atIneq(i, pos1) += poly.atIneq(i, pos2);
-  for (unsigned i = 0, e = poly.getNumEqualities(); i < e; ++i)
-    poly.atEq(i, pos1) += poly.atEq(i, pos2);
-  poly.removeId(pos2);
+void IntegerPolyhedron::eliminateRedundantLocalId(unsigned posA,
+                                                  unsigned posB) {
+  assert(posA < getNumLocalIds() && "Invalid local id position");
+  assert(posB < getNumLocalIds() && "Invalid local id position");
+
+  unsigned localOffset = getIdKindOffset(IdKind::Local);
+  posA += localOffset;
+  posB += localOffset;
+  inequalities.addToColumn(posB, posA, 1);
+  equalities.addToColumn(posB, posA, 1);
+  removeId(posB);
 }
 
 /// Adds additional local ids to the sets such that they both have the union
 /// of the local ids in each set, without changing the set of points that
 /// lie in `this` and `other`.
 ///
 /// To detect local ids that always take the same in both sets, each local id is
 /// represented as a floordiv with constant denominator in terms of other ids.
void IntegerPolyhedron::mergeLocalIds(IntegerPolyhedron &other) {
   std::copy(divsB.begin() + initLocals, divsB.end(),
             divsA.begin() + initLocals);
   std::copy(denomsB.begin() + initLocals, denomsB.end(),
             denomsA.begin() + initLocals);
 
   // Merge function that merges the local variables in both sets by treating
   // them as the same identifier.
   auto merge = [&polyA, &polyB](unsigned i, unsigned j) -> bool {
-    eliminateRedundantLocalId(polyA, i, j);
-    eliminateRedundantLocalId(polyB, i, j);
+    polyA.eliminateRedundantLocalId(i, j);
+    polyB.eliminateRedundantLocalId(i, j);
     return true;
   };
 
   // Merge all divisions by removing duplicate divisions.
   unsigned localOffset = getIdKindOffset(IdKind::Local);
   removeDuplicateDivs(divsA, denomsA, localOffset, merge);
 }
 

mlir/lib/Analysis/Presburger/PWMAFunction.cpp

+//===- PWMAFunction.cpp - MLIR PWMAFunction Class -------------------------===//
+//
+// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
+// See https://llvm.org/LICENSE.txt for license information.
+// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+//
+//===----------------------------------------------------------------------===//
+
+#include "mlir/Analysis/Presburger/PWMAFunction.h"
+#include "mlir/Analysis/Presburger/Simplex.h"
+
+using namespace mlir;
+
+// Return the result of subtracting the two given vectors pointwise.
+// The vectors must be of the same size.
+// e.g., [3, 4, 6] - [2, 5, 1] = [1, -1, 5].
+static SmallVector<int64_t, 8> subtract(ArrayRef<int64_t> vecA,
+                                        ArrayRef<int64_t> vecB) {
+  assert(vecA.size() == vecB.size() &&
+         "Cannot subtract vectors of differing lengths!");
+  SmallVector<int64_t, 8> result;
+  result.reserve(vecA.size());
+  for (unsigned i = 0, e = vecA.size(); i < e; ++i)
+    result.push_back(vecA[i] - vecB[i]);
+  return result;
+}
+
+PresburgerSet PWMAFunction::getDomain() const {
+  PresburgerSet domain =
+      PresburgerSet::getEmptySet(getNumDimIds(), getNumSymbolIds());
+  for (const MultiAffineFunction &piece : pieces)
+    domain.unionPolyInPlace(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!");
+
+  if (!getDomain().containsPoint(point))
+    return {};
+
+  // The point lies in the domain, so we need to compute the output value.
+  // 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>>
+PWMAFunction::valueAt(ArrayRef<int64_t> point) const {
+  assert(point.size() == getNumInputs() &&
+         "Point has incorrect dimensionality!");
+  for (const MultiAffineFunction &piece : pieces)
+    if (Optional<SmallVector<int64_t, 8>> output = piece.valueAt(point))
+      return output;
+  return {};
+}
+
+void MultiAffineFunction::print(raw_ostream &os) const {
+  os << "Domain:";
+  IntegerPolyhedron::print(os);
+  os << "Output:\n";
+  output.print(os);
+  os << "\n";
+}
+
+void MultiAffineFunction::dump() const { print(llvm::errs()); }
+
+bool MultiAffineFunction::isEqual(const MultiAffineFunction &other) const {
+  return hasCompatibleDimensions(other) &&
+         getDomain().isEqual(other.getDomain()) &&
+         isEqualWhereDomainsOverlap(other);
+}
+
+unsigned MultiAffineFunction::insertId(IdKind kind, unsigned pos,
+                                       unsigned num) {
+  unsigned absolutePos = getIdKindOffset(kind) + pos;
+  output.insertColumns(absolutePos, num);
+  return IntegerPolyhedron::insertId(kind, pos, num);
+}
+
+void MultiAffineFunction::swapId(unsigned posA, unsigned posB) {
+  output.swapColumns(posA, posB);
+  IntegerPolyhedron::swapId(posA, posB);
+}
+
+void MultiAffineFunction::removeIdRange(unsigned idStart, unsigned idLimit) {
+  output.removeColumns(idStart, idLimit - idStart);
+  IntegerPolyhedron::removeIdRange(idStart, idLimit);
+}
+
+void MultiAffineFunction::eliminateRedundantLocalId(unsigned posA,
+                                                    unsigned posB) {
+  output.addToColumn(posB, posA, /*scale=*/1);
+  IntegerPolyhedron::eliminateRedundantLocalId(posA, posB);
+}
+
+bool MultiAffineFunction::isEqualWhereDomainsOverlap(
+    MultiAffineFunction other) const {
+  if (!hasCompatibleDimensions(other))
+    return false;
+
+  // `commonFunc` has the same output as `this`.
+  MultiAffineFunction commonFunc = *this;
+  // After this merge, `commonFunc` and `other` have the same local ids; they
+  // are merged.
+  commonFunc.mergeLocalIds(other);
+  // After this, the domain of `commonFunc` will be the intersection of the
+  // domains of `this` and `other`.
+  commonFunc.IntegerPolyhedron::append(other);
+
+  // `commonDomainMatching` contains the subset of the common domain
+  // where the outputs of `this` and `other` match.
+  IntegerPolyhedron commonDomainMatching = commonFunc.getDomain();
+
+  // We can't use `this->output` below because its locals aren't merged with
+  // those of `other`.
+  for (unsigned row = 0, e = getNumOutputs(); row < e; ++row)
+    commonDomainMatching.addEquality(
+        subtract(commonFunc.output.getRow(row), other.output.getRow(row)));
+
+  // If the whole common domain is a subset of commonDomainMatching, then they
+  // are equal and the two functions match on the whole common domain.
+  return commonFunc.getDomain().isSubsetOf(commonDomainMatching);
+}
+
+/// Two PWMAFunctions are equal if they have the same dimensionalities,
+/// the same domain, and take the same value at every point in the domain.
+bool PWMAFunction::isEqual(const PWMAFunction &other) const {
+  if (!hasCompatibleDimensions(other))
+    return false;
+
+  if (!this->getDomain().isEqual(other.getDomain()))
+    return false;
+
+  // Check if, whenever the domains of a piece of `this` and a piece of `other`
+  // overlap, they take the same output value. If `this` and `other` have the
+  // same domain (checked above), then this check passes iff the two functions
+  // have the same output at every point in the domain.
+  for (const MultiAffineFunction &aPiece : this->pieces)
+    for (const MultiAffineFunction &bPiece : other.pieces)
+      if (!aPiece.isEqualWhereDomainsOverlap(bPiece))
+        return false;
+  return true;
+}
+
+void PWMAFunction::addPiece(const MultiAffineFunction &piece) {
+  assert(hasCompatibleDimensions(piece) &&
+         "Piece to be added is not compatible with this PWMAFunction!");
+  assert(piece.isConsistent() && "Piece is internally inconsistent!");
+  assert(this->getDomain()
+             .intersect(PresburgerSet(piece.getDomain()))
+             .isIntegerEmpty() &&
+         "New piece's domain overlaps with that of existing pieces!");
+  pieces.push_back(piece);
+}
+
+void PWMAFunction::addPiece(const IntegerPolyhedron &domain,
+                            const Matrix &output) {
+  addPiece(MultiAffineFunction(domain, output));
+}
+
+void PWMAFunction::print(raw_ostream &os) const {
+  os << pieces.size() << " pieces:\n";
+  for (const MultiAffineFunction &piece : pieces)
+    piece.print(os);
+}
+
+/// The hasCompatibleDimensions functions don't check the number of local ids;
+/// functions are still compatible if they have differing number of locals.
+bool MultiAffineFunction::hasCompatibleDimensions(
+    const MultiAffineFunction &f) const {
+  return getNumDimIds() == f.getNumDimIds() &&
+         getNumSymbolIds() == f.getNumSymbolIds() &&
+         getNumOutputs() == f.getNumOutputs();
+}
+bool PWMAFunction::hasCompatibleDimensions(const MultiAffineFunction &f) const {
+  return getNumDimIds() == f.getNumDimIds() &&
+         getNumSymbolIds() == f.getNumSymbolIds() &&
+         getNumOutputs() == f.getNumOutputs();
+}
+bool PWMAFunction::hasCompatibleDimensions(const PWMAFunction &f) const {
+  return getNumDimIds() == f.getNumDimIds() &&
+         getNumSymbolIds() == f.getNumSymbolIds() &&
+         getNumOutputs() == f.getNumOutputs();
+}

mlir/unittests/Analysis/Presburger/CMakeLists.txt

 add_mlir_unittest(MLIRPresburgerTests
   IntegerPolyhedronTest.cpp
   LinearTransformTest.cpp
   MatrixTest.cpp
   PresburgerSetTest.cpp
+  PWMAFunctionTest.cpp
   SimplexTest.cpp
   ../../Dialect/Affine/Analysis/AffineStructuresParser.cpp
 )
 
 target_link_libraries(MLIRPresburgerTests
   PRIVATE MLIRPresburger
   MLIRAffineAnalysis
   MLIRParser
   )

mlir/unittests/Analysis/Presburger/PWMAFunctionTest.cpp

+//===- PWMAFunctionTest.cpp - Tests for PWMAFunction ----------------------===//
+//
+// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
+// See https://llvm.org/LICENSE.txt for license information.
+// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+//
+//===----------------------------------------------------------------------===//
+//
+// This file contains tests for PWMAFunction.
+//
+//===----------------------------------------------------------------------===//
+
+#include "mlir/Analysis/Presburger/PWMAFunction.h"
+#include "../../Dialect/Affine/Analysis/AffineStructuresParser.h"
+#include "mlir/Analysis/Presburger/PresburgerSet.h"
+
+#include <gmock/gmock.h>
+#include <gtest/gtest.h>
+
+namespace mlir {
+using testing::ElementsAre;
+
+/// Parses an IntegerPolyhedron from a StringRef. It is expected that the
+/// string represents a valid IntegerSet, otherwise it will violate a gtest
+/// assertion.
+static IntegerPolyhedron parsePoly(StringRef str, MLIRContext *context) {
+  FailureOr<IntegerPolyhedron> poly = parseIntegerSetToFAC(str, context);
+  EXPECT_TRUE(succeeded(poly));
+  return *poly;
+}
+
+static Matrix makeMatrix(unsigned numRow, unsigned numColumns,
+                         ArrayRef<SmallVector<int64_t, 8>> matrix) {
+  Matrix results(numRow, numColumns);
+  assert(matrix.size() == numRow);
+  for (unsigned i = 0; i < numRow; ++i) {
+    assert(matrix[i].size() == numColumns &&
+           "Output expression has incorrect dimensionality!");
+    for (unsigned j = 0; j < numColumns; ++j)
+      results(i, j) = matrix[i][j];
+  }
+  return results;
+}
+
+/// Construct a PWMAFunction given the dimensionalities and an array describing
+/// the list of pieces. Each piece is given by a string describing the domain
+/// and a 2D array that represents the output.
+static PWMAFunction parsePWMAF(
+    unsigned numInputs, unsigned numOutputs,
+    ArrayRef<std::pair<StringRef, SmallVector<SmallVector<int64_t, 8>, 8>>>
+        data,
+    unsigned numSymbols = 0) {
+  static MLIRContext context;
+
+  PWMAFunction result(numInputs - numSymbols, numSymbols, numOutputs);
+  for (const auto &pair : data) {
+    IntegerPolyhedron domain = parsePoly(pair.first, &context);
+    result.addPiece(
+        domain, makeMatrix(numOutputs, domain.getNumIds() + 1, pair.second));
+  }
+  return result;
+}
+
+TEST(PWAFunctionTest, isEqual) {
+  MLIRContext context;
+
+  // The output expressions are different but it doesn't matter because they are
+  // equal in this domain.
+  PWMAFunction idAtZeros = parsePWMAF(
+      /*numInputs=*/2, /*numOutputs=*/2,
+      {
+          {"(x, y) : (y == 0)", {{1, 0, 0}, {0, 1, 0}}},             // (x, y).
+          {"(x, y) : (y - 1 >= 0, x == 0)", {{1, 0, 0}, {0, 1, 0}}}, // (x, y).
+          {"(x, y) : (-y - 1 >= 0, x == 0)", {{1, 0, 0}, {0, 1, 0}}} // (x, y).
+      });
+  PWMAFunction idAtZeros2 = parsePWMAF(
+      /*numInputs=*/2, /*numOutputs=*/2,
+      {
+          {"(x, y) : (y == 0)", {{1, 0, 0}, {0, 20, 0}}}, // (x, 20y).
+          {"(x, y) : (y - 1 >= 0, x == 0)", {{30, 0, 0}, {0, 1, 0}}}, //(30x, y)
+          {"(x, y) : (-y - 1 > =0, x == 0)", {{30, 0, 0}, {0, 1, 0}}} //(30x, y)
+      });
+  EXPECT_TRUE(idAtZeros.isEqual(idAtZeros2));
+
+  PWMAFunction notIdAtZeros = parsePWMAF(
+      /*numInputs=*/2, /*numOutputs=*/2,
+      {
+          {"(x, y) : (y == 0)", {{1, 0, 0}, {0, 1, 0}}},              // (x, y).
+          {"(x, y) : (y - 1 >= 0, x == 0)", {{1, 0, 0}, {0, 2, 0}}},  // (x, 2y)
+          {"(x, y) : (-y - 1 >= 0, x == 0)", {{1, 0, 0}, {0, 2, 0}}}, // (x, 2y)
+      });
+  EXPECT_FALSE(idAtZeros.isEqual(notIdAtZeros));
+
+  // These match at their intersection but one has a bigger domain.
+  PWMAFunction idNoNegNegQuadrant = parsePWMAF(
+      /*numInputs=*/2, /*numOutputs=*/2,
+      {
+          {"(x, y) : (x >= 0)", {{1, 0, 0}, {0, 1, 0}}},             // (x, y).
+          {"(x, y) : (-x - 1 >= 0, y >= 0)", {{1, 0, 0}, {0, 1, 0}}} // (x, y).
+      });
+  PWMAFunction idOnlyPosX =
+      parsePWMAF(/*numInputs=*/2, /*numOutputs=*/2,
+                 {
+                     {"(x, y) : (x >= 0)", {{1, 0, 0}, {0, 1, 0}}}, // (x, y).
+                 });
+  EXPECT_FALSE(idNoNegNegQuadrant.isEqual(idOnlyPosX));
+
+  // Different representations of the same domain.
+  PWMAFunction sumPlusOne = parsePWMAF(
+      /*numInputs=*/2, /*numOutputs=*/1,
+      {
+          {"(x, y) : (x >= 0)", {{1, 1, 1}}},                   // x + y + 1.
+          {"(x, y) : (-x - 1 >= 0, -y - 1 >= 0)", {{1, 1, 1}}}, // x + y + 1.
+          {"(x, y) : (-x - 1 >= 0, y >= 0)", {{1, 1, 1}}}       // x + y + 1.
+      });
+  PWMAFunction sumPlusOne2 =
+      parsePWMAF(/*numInputs=*/2, /*numOutputs=*/1,
+                 {
+                     {"(x, y) : ()", {{1, 1, 1}}}, // x + y + 1.
+                 });
+  EXPECT_TRUE(sumPlusOne.isEqual(sumPlusOne2));
+
+  // Functions with zero input dimensions.
+  PWMAFunction noInputs1 = parsePWMAF(/*numInputs=*/0, /*numOutputs=*/1,
+                                      {
+                                          {"() : ()", {{1}}}, // 1.
+                                      });
+  PWMAFunction noInputs2 = parsePWMAF(/*numInputs=*/0, /*numOutputs=*/1,
+                                      {
+                                          {"() : ()", {{2}}}, // 1.
+                                      });
+  EXPECT_TRUE(noInputs1.isEqual(noInputs1));
+  EXPECT_FALSE(noInputs1.isEqual(noInputs2));
+
+  // Mismatched dimensionalities.
+  EXPECT_FALSE(noInputs1.isEqual(sumPlusOne));
+  EXPECT_FALSE(idOnlyPosX.isEqual(sumPlusOne));
+
+  // Divisions.
+  // Domain is only multiples of 6; x = 6k for some k.
+  // x + 4(x/2) + 4(x/3) == 26k.
+  PWMAFunction mul2AndMul3 = parsePWMAF(
+      /*numInputs=*/1, /*numOutputs=*/1,
+      {
+          {"(x) : (x - 2*(x floordiv 2) == 0, x - 3*(x floordiv 3) == 0)",
+           {{1, 4, 4, 0}}}, // x + 4(x/2) + 4(x/3).
+      });
+  PWMAFunction mul6 = parsePWMAF(
+      /*numInputs=*/1, /*numOutputs=*/1,
+      {
+          {"(x) : (x - 6*(x floordiv 6) == 0)", {{0, 26, 0}}}, // 26(x/6).
+      });
+  EXPECT_TRUE(mul2AndMul3.isEqual(mul6));
+
+  PWMAFunction mul6diff = parsePWMAF(
+      /*numInputs=*/1, /*numOutputs=*/1,
+      {
+          {"(x) : (x - 5*(x floordiv 5) == 0)", {{0, 52, 0}}}, // 52(x/6).
+      });
+  EXPECT_FALSE(mul2AndMul3.isEqual(mul6diff));
+
+  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(
+      /*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());
+}
+
+} // namespace mlir