[MLIR] Add support for extracting an integer sample point (if one exists) from an unbounded FlatAffineConstraints.

[arjunp]

Summary

With this, we have complete support for finding integer sample points in FlatAffineConstraints.

[bollu]

mlir/lib/Analysis/AffineStructures.cpp

1193–1194

Consider moving the line

cpp
  FlatAffineConstraints boundedSet = transformedSet;

below the line

// 2) Remove the unbounded dimensions and constraints involving them to
// obtain a bounded set.

?

That way, it's a little less jarring to see boundedSet = transformedSet as the comment explains what's going on.

1198–1259

consider

const unsigned numBoundedDims = result.first;
const unsigned numUnboundedDims = getNumIds() - numBoundedDims;

to aid the reader that these are not modified.

1312

consider:

SmallVector<int64_t, 8> coneSample(llvm::map_range(shrunkenConeSample, ceil));

I believe the llvm::map_range lives in

llvm/ADT/STLExtras.h

2286

Consider adding a comment about how setting a variable is equivalent to updating the constant term and then removing the Id?

mlir/lib/Analysis/LinearTransform.cpp

113–114

I believe this should be

result.reserve(matrix.getNumColumns());

consider:

SmallVector<int64_t, 8> result(0, matrix.getNumColumns());
  for (unsigned col = 0, e = matrix.getNumColumns(); col < e; ++col)
    for (unsigned i = 0, e = matrix.getNumRows(); i < e; ++i)
      result[col] += rowVec[i] * matrix(i, col);

132

I believe the reserve should be:

result.reserve(rowVec.size());

Consider rewriting the body as:

// show reserve rowVec.size() ?
SmallVector<int64_t, 8> result(0, rowVec.size());
// eliminate intermediate `elem`
for (unsigned row = 0, e = matrix.getNumRows(); row < e; row++)
  for (unsigned i = 0, e = matrix.getNumColumns(); i < e; i++)
    result[row] += matrix(row, i) * colVec[i];

[arjunp]

Addressed most of Siddharth's comments.

[arjunp]

Address another comment

[ftynse]

mlir/lib/Analysis/AffineStructures.cpp

1149–1150

I don't understand what it means for a vector v to be a solution to a set. Do you mean the vector belongs to the set?

1198–1259

Like I said in another review, MLIR does not follow this pattern, and many core contributors request the inverse change in their reviews. It is not a codified rule, but we have a strong preference for using a consistent style across the code base. If you prefer we used this pattern in MLIR, please make your case for it and post it on https://llvm.discourse.group/c/mlir/31, folks will be happy to discuss it and reach consensus. Until then, it might be preferable to keep the status quo, i.e., not using this pattern.

1227–1231

Nit: took me some time to convince myself this was okay. In particular, why is this considering e_i = 1 after saying 0 <= e_i < 1 in the previous line (handling the case of e_i = 1 automatically handles the entirety of possible e_i). An example could help IMO

1262

Nit: sample.append. I would also move-construct sample from boundedSample to spare a copy.

2280–2295

ArrayRef<T> is implicitly constructible from T, no need for this magic incantation here. I suppose you can even drop the setAndEliminate(unsigned, int64_t) overload entirely and everything will keep working.

[bollu]

mlir/lib/Analysis/AffineStructures.cpp

1198–1259

(Unfortunately, the link above does not work @ftynse; I'd be happy to read a case for disallowing local constants).

I know the larger case that is made here: https://mlir.llvm.org/docs/Rationale/UsageOfConst/

but that argues about const members of classes; not about local primitive const values. I don't see the problem with const-ing unsigned variables *in local scope* ?

[arjunp]

Addressed Alex's comments.

mlir/lib/Analysis/AffineStructures.cpp

1149–1150

Yes, changed it to say that.

1227–1231

Added an example and a more detailed explanation.

1262

Just made sample a reference to boundedSample.

2280–2295

Dropped the overload.

[ftynse]

mlir/lib/Analysis/AffineStructures.cpp

1198–1259

Strange, the link does work for me. Anyway, it's the usual MLIR discussion forum, there are links to it from mlir.llvm.org.

I generally dislike unnecessary syntactic verbosity, but that's beyond the point here. MLIR code base may or may not follow some patterns that are not codified in the style description. We don't use const on local variables if it is not part of the type (i.e., const reference), we don't use explicit inline on free functions in source files, and we attach * and & to the name rather than the type, to name a few. We use the range abstraction when possible and prefix getters with get. We prefer getters to public const fields. All these are implicit, and that's one of the reasons why we do code review. There is a general rule of following the style of the surrounding code, so if there somehow is a file in the code base that has different patterns, those may prevail. Now, for these relatively new files, the surrounding code is the rest of the code base. Introducing the patterns that locally diverge from the rest of the code base makes life harder for anybody who has to work a lot across the code base, usually the people who will review and maintain this code, and perform large-scale changes later. So I care a lot about consistency of this code with the rest of the code base. I am open to discussing the patterns we prefer to use or not as long as this discussion applies to the entire code base and not a single file, directory, subsystem or patch author. Since you are the one who suggests to diverge from the established practice, it sounds reasonable that the burden of starting the discussion and reaching consensus is on you.

[bondhugula]

mlir/lib/Analysis/AffineStructures.cpp

1198–1259

The * and & attachment to the name comes directly from LLVM's style which we aren't overriding. clang-tidy/format is already set to catch and fix these. The MLIR coding style only says where we diverge from the LLVM style. For the rest that don't find a mention in the LLVM as well as MLIR style guides, it's normally expected that a developer will use what they think is better. We should include more in the style file if a specific pattern is being consistently used and is preferred for consistency -- or else, it just leads to more review-time burden for both reviewers and contributors.

mlir/include/mlir/Analysis/AffineStructures.h

public:
   bool isEmpty() const;
 
   /// Runs the GCD test on all equality constraints. Returns 'true' if this test
   /// fails on any equality. Returns 'false' otherwise.
   /// This test can be used to disprove the existence of a solution. If it
   /// returns true, no integer solution to the equality constraints can exist.
   bool isEmptyByGCDTest() const;
 
-  /// Runs the GCD test heuristic. If it proves inconclusive, falls back to
-  /// generalized basis reduction if the set is bounded.
-  ///
   /// Returns true if the set of constraints is found to have no solution,
-  /// false if a solution exists or all tests were inconclusive.
+  /// false if a solution exists. Uses the same algorithm as findIntegerSample.
   bool isIntegerEmpty() const;
 
   // Returns a matrix where each row is a vector along which the polytope is
   // bounded. The span of the returned vectors is guaranteed to contain all
   // such vectors. The returned vectors are NOT guaranteed to be linearly
   // independent. This function should not be called on empty sets.
   Matrix getBoundedDirections() const;
 
-  /// Find a sample point satisfying the constraints. This uses a branch and
-  /// bound algorithm with generalized basis reduction, which always works if
-  /// the set is bounded. This should not be called for unbounded sets.
+  /// Find an integer sample point satisfying the constraints using a
+  /// branch and bound algorithm with generalized basis reduction, with some
+  /// additional processing using Simplex for unbounded sets.
   ///
-  /// Returns such a point if one exists, or an empty Optional otherwise.
+  /// 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.
   bool containsPoint(ArrayRef<int64_t> point) const;
 
   // Clones this object.
   std::unique_ptr<FlatAffineConstraints> clone() const;
public:
   /// Changes the partition between dimensions and symbols. Depending on the new
   /// symbol count, either a chunk of trailing dimensional identifiers becomes
   /// symbols, or some of the leading symbols become dimensions.
   void setDimSymbolSeparation(unsigned newSymbolCount);
 
   /// Changes all symbol identifiers which are loop IVs to dim identifiers.
   void convertLoopIVSymbolsToDims();
 
-  /// Sets the specified identifier to a constant and removes it.
-  void setAndEliminate(unsigned pos, int64_t constVal);
+  /// Sets the values.size() identifiers starting at pos to the specified values
+  /// and removes them.
+  void setAndEliminate(unsigned pos, ArrayRef<int64_t> values);
 
   /// Tries to fold the specified identifier to a constant using a trivial
   /// equality detection; if successful, the constant is substituted for the
   /// identifier everywhere in the constraint system and then removed from the
   /// system.
   LogicalResult constantFoldId(unsigned pos);
 
   /// This method calls constantFoldId for the specified range of identifiers,

mlir/include/mlir/Analysis/LinearTransform.h

public:
   // Specifically, T is such that in every column the first non-zero row is
   // strictly below that of the previous column, and all columns which have only
   // zeros are at the end.
   static std::pair<unsigned, LinearTransform>
   makeTransformToColumnEchelon(Matrix m);
 
   // Returns a FlatAffineConstraints having a constraint vector vT for every
   // constraint vector v in fac, where T is this transform.
-  FlatAffineConstraints applyTo(const FlatAffineConstraints &fac);
+  FlatAffineConstraints applyTo(const FlatAffineConstraints &fac) const;
 
-  // Post-multiply the given vector v with this transform, say T, returning vT.
-  SmallVector<int64_t, 8> applyTo(ArrayRef<int64_t> v);
+  // The given vector is interpreted as a row vector v. Post-multiply v with
+  // this transform, say T, and return vT.
+  SmallVector<int64_t, 8> postMultiplyRow(ArrayRef<int64_t> rowVec) const;
+
+  // The given vector is interpreted as a column vector v. Pre-multiply v with
+  // this transform, say T, and return Tv.
+  SmallVector<int64_t, 8> preMultiplyColumn(ArrayRef<int64_t> colVec) const;
 
 private:
   Matrix matrix;
 };
 
 } // namespace mlir
 #endif // MLIR_ANALYSIS_LINEARTRANSFORM_H

mlir/include/mlir/Analysis/Presburger/Simplex.h

public:
   /// Returns true if the polytope is unbounded, i.e., extends to infinity in
   /// some direction. Otherwise, returns false.
   bool isUnbounded();
 
   /// Make a tableau to represent a pair of points in the given tableaus, one in
   /// tableau A and one in B.
   static Simplex makeProduct(const Simplex &a, const Simplex &b);
 
-  /// Returns the current sample point if it is integral. Otherwise, returns an
+  /// Returns a rational sample point. This should not be called when Simplex is
+  /// empty.
+  SmallVector<Fraction, 8> getRationalSample() const;
+
+  /// Returns the current sample point if it is integral. Otherwise, returns
   /// None.
   Optional<SmallVector<int64_t, 8>> getSamplePointIfIntegral() const;
 
   /// Returns an integer sample point if one exists, or None
   /// otherwise. This should only be called for bounded sets.
   Optional<SmallVector<int64_t, 8>> findIntegerSample();
 
   /// Print the tableau's internal state.

mlir/lib/Analysis/AffineStructures.cpp

 #include "mlir/Analysis/Presburger/Simplex.h"
 #include "mlir/Dialect/Affine/IR/AffineOps.h"
 #include "mlir/Dialect/Affine/IR/AffineValueMap.h"
 #include "mlir/Dialect/StandardOps/IR/Ops.h"
 #include "mlir/IR/AffineExprVisitor.h"
 #include "mlir/IR/IntegerSet.h"
 #include "mlir/Support/LLVM.h"
 #include "mlir/Support/MathExtras.h"
+#include "llvm/ADT/STLExtras.h"
 #include "llvm/ADT/SmallPtrSet.h"
 #include "llvm/ADT/SmallVector.h"
 #include "llvm/Support/Debug.h"
 #include "llvm/Support/raw_ostream.h"
 
 #define DEBUG_TYPE "affine-structures"
 
 using namespace mlir;
void removeConstraintsInvolvingSuffixDims(FlatAffineConstraints &fac,
   for (unsigned i = fac.getNumEqualities(); i > 0; i--)
     if (eqInvolvesSuffixDims(fac, i - 1, unboundedDims))
       fac.removeEquality(i - 1);
   for (unsigned i = fac.getNumInequalities(); i > 0; i--)
     if (ineqInvolvesSuffixDims(fac, i - 1, unboundedDims))
       fac.removeInequality(i - 1);
 }
 
+bool FlatAffineConstraints::isIntegerEmpty() const {
+  return !findIntegerSample().hasValue();
+}
+
 /// Let this set be S. If S is bounded then we directly call into the GBR
 /// sampling algorithm. Otherwise, there are some unbounded directions, i.e.,
-/// vectors v such that S extends to infininty along v or -v. In this case we
+/// vectors v such that S extends to infinity along v or -v. In this case we
 /// use an algorithm described in the integer set library (isl) manual and used
 /// by the isl_set_sample function in that library. The algorithm is:
 ///
 /// 1) Apply a unimodular transform T to S to obtain S*T, such that all
 /// dimensions in which S*T is bounded lie in the linear span of a prefix of the
 /// dimensions.
 ///
-/// 2) Construct a set transformedSet by removing all constraints that involve
-/// the unbounded dimensions and also deleting the unbounded dimensions. Note
-/// that this is a bounded set.
+/// 2) Construct a set B by removing all constraints that involve
+/// the unbounded dimensions and then deleting the unbounded dimensions. Note
+/// that B is a Bounded set.
 ///
-/// 3) Check if transformedSet is empty using the GBR sampling algorithm.
+/// 3) Try to obtain a sample from B using the GBR sampling
+/// algorithm. If no sample is found, return that S is empty.
 ///
-/// 4) return S is empty iff transformedSet is empty.
+/// 4) Otherwise, substitute the obtained sample into S*T to obtain a set
+/// C. C is a full-dimensional Cone and always contains a sample.
 ///
-/// Since T is unimodular, a vector v is a solution to S*T iff T*v is a
-/// solution to S. The following is a sketch of a proof that S*T is empty
-/// iff transformedSet is empty:
+/// 5) Obtain an integer sample from C.
 ///
-/// If transformedSet is empty, then S*T is certainly empty since transformedSet
-/// was obtained by removing constraints and deleting dimensions from S*T.
+/// 6) Return T*v, where v is the concatenation of the samples from B and C.
 ///
-/// If transformedSet contains a sample, consider the set C obtained by
-/// substituting the sample for the bounded dimensions of S*T. All the
-/// constraints of S*T that did not involve unbounded dimensions are
-/// satisfied by this substitution.
+/// The following is a sketch of a proof that
+/// a) If the algorithm returns empty, then S is empty.
+/// b) If the algorithm returns a sample, it is a valid sample in S.
 ///
-/// In step 1, all dimensions in the linear span of the dimensions outside the
-/// prefix are unbounded in S*T. Substituting values for the bounded dimensions
-/// cannot makes these dimensions bounded, and these are the only remaining
-/// dimensions in C, so C is unbounded along every vector. C is hence a
-/// full-dimensional cone and therefore always contains an integer point, which
-/// we can then substitute to get a full solution to S*T.
-bool FlatAffineConstraints::isIntegerEmpty() const {
+/// The algorithm returns empty only if B is empty, in which case S*T is
+/// certainly empty since B was obtained by removing constraints and then
+/// deleting unconstrained dimensions from S*T. Since T is unimodular, a vector
+/// v is in S*T iff T*v is in S. So in this case, since
+/// S*T is empty, S is empty too.
+///
+/// Otherwise, the algorithm substitutes the sample from B into S*T. All the
+/// constraints of S*T that did not involve unbounded dimensions are satisfied
+/// by this substitution. All dimensions in the linear span of the dimensions
+/// outside the prefix are unbounded in S*T (step 1). Substituting values for
+/// the bounded dimensions cannot make these dimensions bounded, and these are
+/// the only remaining dimensions in C, so C is unbounded along every vector (in
+/// the positive or negative direction, or both). C is hence a full-dimensional
+/// cone and therefore always contains an integer point.
+///
+/// Concatenating the samples from B and C gives a sample v in S*T, so the
+/// returned sample T*v is a sample in S.
+Optional<SmallVector<int64_t, 8>>
+FlatAffineConstraints::findIntegerSample() const {
   // First, try the GCD test heuristic.
   if (isEmptyByGCDTest())
-    return true;
+    return {};
 
   Simplex simplex(*this);
   if (simplex.isEmpty())
-    return true;
+    return {};
 
   // For a bounded set, we directly call into the GBR sampling algorithm.
   if (!simplex.isUnbounded())
-    return !simplex.findIntegerSample().hasValue();
+    return simplex.findIntegerSample();
 
   // The set is unbounded. We cannot directly use the GBR algorithm.
   //
   // m is a matrix containing, in each row, a vector in which S is
   // bounded, such that the linear span of all these dimensions contains all
   // bounded dimensions in S.
   Matrix m = getBoundedDirections();
   // In column echelon form, each row of m occupies only the first rank(m)
   // columns and has zeros on the other columns. The transform T that brings S
   // to column echelon form is unimodular as well, so this is a suitable
   // transform to use in step 1 of the algorithm.
   std::pair<unsigned, LinearTransform> result =
       LinearTransform::makeTransformToColumnEchelon(std::move(m));
-  FlatAffineConstraints transformedSet = result.second.applyTo(*this);
-
+  const LinearTransform &transform = result.second;
+  // 1) Apply T to S to obtain S*T.
+  FlatAffineConstraints transformedSet = transform.applyTo(*this);
+
+  // 2) Remove the unbounded dimensions and constraints involving them to
+  // obtain a bounded set.
+  FlatAffineConstraints boundedSet = transformedSet;
   unsigned numBoundedDims = result.first;
   unsigned numUnboundedDims = getNumIds() - numBoundedDims;
-  removeConstraintsInvolvingSuffixDims(transformedSet, numUnboundedDims);
+  removeConstraintsInvolvingSuffixDims(boundedSet, numUnboundedDims);
+  boundedSet.removeIdRange(numBoundedDims, boundedSet.getNumIds());
 
-  // Remove all the unbounded dimensions.
-  transformedSet.removeIdRange(numBoundedDims, transformedSet.getNumIds());
+  // 3) Try to obtain a sample from the bounded set.
+  Optional<SmallVector<int64_t, 8>> boundedSample =
+      Simplex(boundedSet).findIntegerSample();
+  if (!boundedSample)
+    return {};
+  assert(boundedSet.containsPoint(*boundedSample) &&
+         "Simplex returned an invalid sample!");
+
+  // 4) Substitute the values of the bounded dimensions into S*T to obtain a
+  // full-dimensional cone, which necessarily contains an integer sample.
+  transformedSet.setAndEliminate(0, *boundedSample);
+  FlatAffineConstraints &cone = transformedSet;
 
-  return !Simplex(transformedSet).findIntegerSample().hasValue();
+  // 5) Obtain an integer sample from the cone.
+  //
+  // We shrink the cone such that for any rational point in the shrunken cone,
+  // rounding up each of the point's coordinates produces a point that still
+  // lies in the original cone.
+  //
+  // Rounding up a point x adds a number e_i in [0, 1) to each coordinate x_i.
+  // For each inequality sum_i a_i x_i + c >= 0 in the original cone, the
+  // shrunken cone will have the inequality tightened by some amount s, such
+  // that if x satisfies the shrunken cone's tightened inequality, then x + e
+  // satisfies the original inequality, i.e.,
+  //
+  // sum_i a_i x_i + c + s >= 0 implies sum_i a_i (x_i + e_i) + c >= 0
+  //
+  // for any e_i values in [0, 1). In fact, we will handle the slightly more
+  // general case where e_i can be in [0, 1]. For example, consider the
+  // inequality 2x_1 - 3x_2 - 7x_3 - 6 >= 0, and let x = (3, 0, 0). How low
+  // could the LHS go if we added a number in [0, 1] to each coordinate? The LHS
+  // is minimized when we add 1 to the x_i with negative coefficient a_i and
+  // keep the other x_i the same. In the example, we would get x = (3, 1, 1),
+  // changing the value of the LHS by -3 + -7 = -10.
+  //
+  // In general, the value of the LHS can change by at most the sum of the
+  // negative a_i, so we accomodate this by shifting the inequality by this
+  // amount for the shrunken cone.
+  for (unsigned i = 0, e = cone.getNumInequalities(); i < e; ++i) {
+    for (unsigned j = 0; j < cone.numIds; ++j) {
+      int64_t coeff = cone.atIneq(i, j);
+      if (coeff < 0)
+        cone.atIneq(i, cone.numIds) += coeff;
+    }
   }
 
-Optional<SmallVector<int64_t, 8>>
-FlatAffineConstraints::findIntegerSample() const {
-  return Simplex(*this).findIntegerSample();
+  // Obtain an integer sample in the cone by rounding up a rational point from
+  // the shrunken cone. Shrinking the cone amounts to shifting its apex
+  // "inwards" without changing its "shape"; the shrunken cone is still a
+  // full-dimensional cone and is hence non-empty.
+  Simplex shrunkenConeSimplex(cone);
+  assert(!shrunkenConeSimplex.isEmpty() && "Shrunken cone cannot be empty!");
+  SmallVector<Fraction, 8> shrunkenConeSample =
+      shrunkenConeSimplex.getRationalSample();
+
+  SmallVector<int64_t, 8> coneSample(llvm::map_range(shrunkenConeSample, ceil));
+
+  // 6) Return transform * concat(boundedSample, coneSample).
+  SmallVector<int64_t, 8> &sample = boundedSample.getValue();
+  sample.append(coneSample.begin(), coneSample.end());
+  return transform.preMultiplyColumn(sample);
 }
 
 /// Helper to evaluate an affine expression at a point.
 /// The expression is a list of coefficients for the dimensions followed by the
 /// constant term.
 static int64_t valueAt(ArrayRef<int64_t> expr, ArrayRef<int64_t> point) {
   assert(expr.size() == 1 + point.size() &&
          "Dimensionalities of point and expression don't match!");
for (unsigned i = 0, e = getNumInequalities(); i < e; ++i) {
     for (unsigned j = 1; j < numCols - 1; ++j) {
       gcd = llvm::GreatestCommonDivisor64(gcd, std::abs(atIneq(i, j)));
     }
     if (gcd > 0 && gcd != 1) {
       int64_t gcdI = static_cast<int64_t>(gcd);
       // Tighten the constant term and normalize the constraint by the GCD.
       atIneq(i, numCols - 1) = mlir::floorDiv(atIneq(i, numCols - 1), gcdI);
       for (unsigned j = 0, e = numCols - 1; j < e; ++j)
         atIneq(i, j) /= gcdI;
     }
   }
 }
 
 // Eliminates all identifier variables in column range [posStart, posLimit).
 // Returns the number of variables eliminated.
 unsigned FlatAffineConstraints::gaussianEliminateIds(unsigned posStart,
                                                      unsigned posLimit) {
for (unsigned r = 0, e = cst.getNumEqualities(); r < e; r++) {
     }
     if (c == f)
       // Equality is free of other identifiers.
       return r;
   }
   return -1;
 }
 
-void FlatAffineConstraints::setAndEliminate(unsigned pos, int64_t constVal) {
-  assert(pos < getNumIds() && "invalid position");
-  for (unsigned r = 0, e = getNumInequalities(); r < e; r++) {
-    atIneq(r, getNumCols() - 1) += atIneq(r, pos) * constVal;
-  }
-  for (unsigned r = 0, e = getNumEqualities(); r < e; r++) {
-    atEq(r, getNumCols() - 1) += atEq(r, pos) * constVal;
-  }
-  removeId(pos);
+void FlatAffineConstraints::setAndEliminate(unsigned pos,
+                                            ArrayRef<int64_t> values) {
+  if (values.empty())
+    return;
+  assert(pos + values.size() <= getNumIds() &&
+         "invalid position or too many values");
+  // Setting x_j = p in sum_i a_i x_i + c is equivalent to adding p*a_j to the
+  // constant term and removing the id x_j. We do this for all the ids
+  // pos, pos + 1, ... pos + values.size() - 1.
+  for (unsigned r = 0, e = getNumInequalities(); r < e; r++)
+    for (unsigned i = 0, numVals = values.size(); i < numVals; ++i)
+      atIneq(r, getNumCols() - 1) += atIneq(r, pos + i) * values[i];
+  for (unsigned r = 0, e = getNumEqualities(); r < e; r++)
+    for (unsigned i = 0, numVals = values.size(); i < numVals; ++i)
+      atEq(r, getNumCols() - 1) += atEq(r, pos + i) * values[i];
+  removeIdRange(pos, pos + values.size());
 }
 
 LogicalResult FlatAffineConstraints::constantFoldId(unsigned pos) {
   assert(pos < getNumIds() && "invalid position");
   int rowIdx;
   if ((rowIdx = findEqualityToConstant(*this, pos)) == -1)
     return failure();
 

mlir/lib/Analysis/LinearTransform.cpp

for (unsigned i = echelonCol + 1, e = m.getNumColumns(); i < e; ++i) {
       }
     }
 
     ++echelonCol;
   }
 
   return {echelonCol, LinearTransform(std::move(resultMatrix))};
 }
 
-SmallVector<int64_t, 8> LinearTransform::applyTo(ArrayRef<int64_t> v) {
-  assert(v.size() == matrix.getNumRows() &&
-         "vector dimension should be matrix output dimension");
-
-  SmallVector<int64_t, 8> result;
-  result.reserve(v.size());
-  for (unsigned col = 0, e = matrix.getNumColumns(); col < e; ++col) {
-    int64_t elem = 0;
+SmallVector<int64_t, 8>
+LinearTransform::postMultiplyRow(ArrayRef<int64_t> rowVec) const {
+  assert(rowVec.size() == matrix.getNumRows() &&
+         "row vector dimension should match transform output dimension");
+
+  SmallVector<int64_t, 8> result(matrix.getNumColumns(), 0);
+  for (unsigned col = 0, e = matrix.getNumColumns(); col < e; ++col)
     for (unsigned i = 0, e = matrix.getNumRows(); i < e; ++i)
-      elem += v[i] * matrix(i, col);
-    result.push_back(elem);
-  }
+      result[col] += rowVec[i] * matrix(i, col);
+  return result;
+
+SmallVector<int64_t, 8>
+LinearTransform::preMultiplyColumn(ArrayRef<int64_t> colVec) const {
+  assert(matrix.getNumColumns() == colVec.size() &&
+         "column vector dimension should match transform input dimension");
+
+  SmallVector<int64_t, 8> result(matrix.getNumRows(), 0);
+  for (unsigned row = 0, e = matrix.getNumRows(); row < e; row++)
+    for (unsigned i = 0, e = matrix.getNumColumns(); i < e; i++)
+      result[row] += matrix(row, i) * colVec[i];
   return result;
 }
 
 FlatAffineConstraints
-LinearTransform::applyTo(const FlatAffineConstraints &fac) {
+LinearTransform::applyTo(const FlatAffineConstraints &fac) const {
   FlatAffineConstraints result(fac.getNumDimIds());
 
   for (unsigned i = 0, e = fac.getNumEqualities(); i < e; ++i) {
     ArrayRef<int64_t> eq = fac.getEquality(i);
 
     int64_t c = eq.back();
 
-    SmallVector<int64_t, 8> newEq = applyTo(eq.drop_back());
+    SmallVector<int64_t, 8> newEq = postMultiplyRow(eq.drop_back());
     newEq.push_back(c);
     result.addEquality(newEq);
   }
 
   for (unsigned i = 0, e = fac.getNumInequalities(); i < e; ++i) {
     ArrayRef<int64_t> ineq = fac.getInequality(i);
 
     int64_t c = ineq.back();
 
-    SmallVector<int64_t, 8> newIneq = applyTo(ineq.drop_back());
+    SmallVector<int64_t, 8> newIneq = postMultiplyRow(ineq.drop_back());
     newIneq.push_back(c);
     result.addInequality(newIneq);
   }
 
   return result;
 }
 
 } // namespace mlir

mlir/lib/Analysis/Presburger/Simplex.cpp

Simplex Simplex::makeProduct(const Simplex &a, const Simplex &b) {
   for (unsigned row = a.nRedundant; row < a.nRow; ++row)
     appendRowFromA(row);
   for (unsigned row = b.nRedundant; row < b.nRow; ++row)
     appendRowFromB(row);
 
   return result;
 }
 
-Optional<SmallVector<int64_t, 8>> Simplex::getSamplePointIfIntegral() const {
-  // The tableau is empty, so no sample point exists.
-  if (empty)
-    return {};
+SmallVector<Fraction, 8> Simplex::getRationalSample() const {
+  assert(!empty && "This should not be called when Simplex is empty.");
 
-  SmallVector<int64_t, 8> sample;
+  SmallVector<Fraction, 8> sample;
+  sample.reserve(var.size());
   // Push the sample value for each variable into the vector.
   for (const Unknown &u : var) {
     if (u.orientation == Orientation::Column) {
       // If the variable is in column position, its sample value is zero.
-      sample.push_back(0);
+      sample.emplace_back(0, 1);
     } else {
       // If the variable is in row position, its sample value is the entry in
       // the constant column divided by the entry in the common denominator
-      // column. If this is not an integer, then the sample point is not
-      // integral so we return None.
-      if (tableau(u.pos, 1) % tableau(u.pos, 0) != 0)
-        return {};
-      sample.push_back(tableau(u.pos, 1) / tableau(u.pos, 0));
+      // column.
+      sample.emplace_back(tableau(u.pos, 1), tableau(u.pos, 0));
     }
   }
   return sample;
 }
 
+Optional<SmallVector<int64_t, 8>> Simplex::getSamplePointIfIntegral() const {
+  // If the tableau is empty, no sample point exists.
+  if (empty)
+    return {};
+  SmallVector<Fraction, 8> rationalSample = getRationalSample();
+  SmallVector<int64_t, 8> integerSample;
+  integerSample.reserve(var.size());
+  for (const Fraction &coord : rationalSample) {
+    // If the sample is non-integral, return None.
+    if (coord.num % coord.den != 0)
+      return {};
+    integerSample.push_back(coord.num / coord.den);
+  }
+  return integerSample;
+}
+
 /// Given a simplex for a polytope, construct a new simplex whose variables are
 /// identified with a pair of points (x, y) in the original polytope. Supports
 /// some operations needed for generalized basis reduction. In what follows,
 /// dotProduct(x, y) = x_1 * y_1 + x_2 * y_2 + ... x_n * y_n where n is the
 /// dimension of the original polytope.
 ///
 /// This supports adding equality constraints dotProduct(dir, x - y) == 0. It
 /// also supports rolling back this addition, by maintaining a snapshot stack

mlir/unittests/Analysis/AffineStructuresTest.cpp

TEST(FlatAffineConstraintsTest, FindSampleTest) {
   checkSample(true, makeFACFromConstraints(2,
                                            {
                                                {2, 0, 0},   // 2x >= 1.
                                                {-2, 0, 99}, // 2x <= 99.
                                                {0, 2, 0},   // 2y >= 0.
                                                {0, -2, 99}, // 2y <= 99.
                                            },
                                            {}));
-}
-
-TEST(FlatAffineConstraintsTest, IsIntegerEmptyTest) {
-  // 1 <= 5x and 5x <= 4 (no solution).
-  EXPECT_TRUE(
-      makeFACFromConstraints(1, {{5, -1}, {-5, 4}}, {}).isIntegerEmpty());
-  // 1 <= 5x and 5x <= 9 (solution: x = 1).
-  EXPECT_FALSE(
-      makeFACFromConstraints(1, {{5, -1}, {-5, 9}}, {}).isIntegerEmpty());
-
-  // Unbounded sets.
-  EXPECT_TRUE(makeFACFromConstraints(3,
-                                     {
-                                         {2, 0, 0, -1}, // 2x >= 1
-                                         {-2, 0, 0, 1}, // 2x <= 1
-                                         {0, 2, 0, -1}, // 2y >= 1
-                                         {0, -2, 0, 1}, // 2y <= 1
-                                         {0, 0, 2, -1}, // 2z >= 1
-                                     },
-                                     {})
-                  .isIntegerEmpty());
-
-  EXPECT_FALSE(makeFACFromConstraints(3,
-                                      {
-                                          {2, 0, 0, -1},  // 2x >= 1
-                                          {-3, 0, 0, 3},  // 3x <= 3
-                                          {0, 0, 5, -6},  // 5z >= 6
-                                          {0, 0, -7, 17}, // 7z <= 17
-                                          {0, 3, 0, -2},  // 3y >= 2
-                                      },
-                                      {})
-                   .isIntegerEmpty());
-
   // 2D cone with apex at (10000, 10000) and
   // edges passing through (1/3, 0) and (2/3, 0).
-  EXPECT_FALSE(
+  checkSample(
+      true,
       makeFACFromConstraints(
-          2, {{300000, -299999, -100000}, {-300000, 299998, 200000}}, {})
-          .isIntegerEmpty());
+          2, {{300000, -299999, -100000}, {-300000, 299998, 200000}}, {}));
 
   // Cartesian product of a tetrahedron and a 2D cone.
   // The tetrahedron has vertices at
   // (1/3, 0, 0), (2/3, 0, 0), (2/3, 0, 10000), and (10000, 10000, 10000).
   // The first three points form a triangular base on the xz plane with the
   // apex at the fourth point, which is the only integer point.
   // The cone has apex at (10000, 10000) and
   // edges passing through (1/3, 0) and (2/3, 0).
checkPermutationsSample(
           {-150000, 149999, 0, 0, 0, 100000},
 
           // Triangle constraints:
           // 300000p - 299999q >= 100000
           {0, 0, 0, 300000, -299999, -100000},
           // -300000p + 299998q + 200000 >= 0
           {0, 0, 0, -300000, 299998, 200000},
       },
-      {}, TestFunction::Empty);
+      {});
 
   // Cartesian product of same tetrahedron as above and {(p, q) : 1/3 <= p <=
   // 2/3}. Since the second set is empty, the whole set is too.
   checkPermutationsSample(
       false /* empty */, 5,
       {
           // Tetrahedron contraints:
           {0, 1, 0, 0, 0, 0},  // y >= 0
           {0, -1, 1, 0, 0, 0}, // z >= y
                                // -300000x + 299998y + 100000 + z <= 0.
           {300000, -299998, -1, 0, 0, -100000},
           // -150000x + 149999y + 100000 >= 0.
           {-150000, 149999, 0, 0, 0, 100000},
 
           // Second set constraints:
           // 3p >= 1
           {0, 0, 0, 3, 0, -1},
           // 3p <= 2
           {0, 0, 0, -3, 0, 2},
       },
-      {}, TestFunction::Empty);
+      {});
 
   // Cartesian product of same tetrahedron as above and
   // {(p, q, r) : 1 <= p <= 2 and p = 3q + 3r}.
   // Since the second set is empty, the whole set is too.
   checkPermutationsSample(
       false /* empty */, 5,
       {
           // Tetrahedron contraints:
           {0, 1, 0, 0, 0, 0, 0},  // y >= 0
           {0, -1, 1, 0, 0, 0, 0}, // z >= y
                                   // -300000x + 299998y + 100000 + z <= 0.
           {300000, -299998, -1, 0, 0, 0, -100000},
           // -150000x + 149999y + 100000 >= 0.
           {-150000, 149999, 0, 0, 0, 0, 100000},
 
           // Second set constraints:
           // p >= 1
           {0, 0, 0, 1, 0, 0, -1},
           // p <= 2
           {0, 0, 0, -1, 0, 0, 2},
       },
       {
           {0, 0, 0, 1, -3, -3, 0}, // p = 3q + 3r
-      },
-      TestFunction::Empty);
+      });
 
   // Cartesian product of a tetrahedron and a 2D cone.
   // The tetrahedron is empty and has vertices at
   // (1/3, 0, 0), (2/3, 0, 0), (2/3, 0, 100), and (100, 100 - 1/3, 100).
   // The cone has apex at (10000, 10000) and
   // edges passing through (1/3, 0) and (2/3, 0).
   // Since the tetrahedron is empty, the Cartesian product is too.
   checkPermutationsSample(false /* empty */, 5,
                           {
                               // Tetrahedron contraints:
                               {0, 1, 0, 0, 0, 0},
                               {0, -300, 299, 0, 0, 0},
                               {300 * 299, -89400, -299, 0, 0, -100 * 299},
                               {-897, 894, 0, 0, 0, 598},
 
                               // Triangle constraints:
                               // 300000p - 299999q >= 100000
                               {0, 0, 0, 300000, -299999, -100000},
                               // -300000p + 299998q + 200000 >= 0
                               {0, 0, 0, -300000, 299998, 200000},
                           },
-                          {}, TestFunction::Empty);
+                          {});
 
   // Cartesian product of same tetrahedron as above and
   // {(p, q) : 1/3 <= p <= 2/3}.
   checkPermutationsSample(false /* empty */, 5,
                           {
                               // Tetrahedron contraints:
                               {0, 1, 0, 0, 0, 0},
                               {0, -300, 299, 0, 0, 0},
                               {300 * 299, -89400, -299, 0, 0, -100 * 299},
                               {-897, 894, 0, 0, 0, 598},
 
                               // Second set constraints:
                               // 3p >= 1
                               {0, 0, 0, 3, 0, -1},
                               // 3p <= 2
                               {0, 0, 0, -3, 0, 2},
                           },
-                          {}, TestFunction::Empty);
+                          {});
+
+  checkSample(true, makeFACFromConstraints(3,
+                                           {
+                                               {2, 0, 0, -1}, // 2x >= 1
+                                           },
+                                           {{
+                                               {1, -1, 0, -1}, // y = x - 1
+                                               {0, 1, -1, 0},  // z = y
+                                           }}));
+}
+
+TEST(FlatAffineConstraintsTest, IsIntegerEmptyTest) {
+  // 1 <= 5x and 5x <= 4 (no solution).
+  EXPECT_TRUE(
+      makeFACFromConstraints(1, {{5, -1}, {-5, 4}}, {}).isIntegerEmpty());
+  // 1 <= 5x and 5x <= 9 (solution: x = 1).
+  EXPECT_FALSE(
+      makeFACFromConstraints(1, {{5, -1}, {-5, 9}}, {}).isIntegerEmpty());
+
+  // Unbounded sets.
+  EXPECT_TRUE(makeFACFromConstraints(3,
+                                     {
+                                         {0, 2, 0, -1}, // 2y >= 1
+                                         {0, -2, 0, 1}, // 2y <= 1
+                                         {0, 0, 2, -1}, // 2z >= 1
+                                     },
+                                     {{2, 0, 0, -1}} // 2x = 1
+                                     )
+                  .isIntegerEmpty());
+
+  EXPECT_FALSE(makeFACFromConstraints(3,
+                                      {
+                                          {2, 0, 0, -1},  // 2x >= 1
+                                          {-3, 0, 0, 3},  // 3x <= 3
+                                          {0, 0, 5, -6},  // 5z >= 6
+                                          {0, 0, -7, 17}, // 7z <= 17
+                                          {0, 3, 0, -2},  // 3y >= 2
+                                      },
+                                      {})
+                   .isIntegerEmpty());
 
   EXPECT_FALSE(makeFACFromConstraints(3,
                                       {
                                           {2, 0, 0, -1}, // 2x >= 1
                                       },
                                       {{
                                           {1, -1, 0, -1}, // y = x - 1
                                           {0, 1, -1, 0},  // z = y

mlir/unittests/Analysis/LinearTransformTest.cpp

void testColumnEchelonForm(const Matrix &m, unsigned expectedRank) {
   std::pair<unsigned, LinearTransform> result =
       LinearTransform::makeTransformToColumnEchelon(m);
   unsigned rank = result.first;
   EXPECT_EQ(rank, expectedRank);
   LinearTransform transform = result.second;
   // In column echelon form, each row's last non-zero value can be at most one
   // column to the right of the last non-zero column among the previous rows.
   for (unsigned row = 0, nRows = m.getNumRows(); row < nRows; ++row) {
-    SmallVector<int64_t, 8> rowVec = transform.applyTo(m.getRow(row));
+    SmallVector<int64_t, 8> rowVec = transform.postMultiplyRow(m.getRow(row));
     for (unsigned col = lastAllowedNonZeroCol + 1, nCols = m.getNumColumns();
          col < nCols; ++col) {
       EXPECT_EQ(rowVec[col], 0);
       if (rowVec[col] != 0) {
         llvm::errs() << "Failed at input matrix:\n";
         m.dump();
       }
     }