Exact integer emptiness checks for FlatAffineConstraints

[arjunp]

Summary

This patch adds the capability to perform exact integer emptiness checks for FlatAffineConstraints using the General Basis Reduction algorithm (GBR). Previously, only a heuristic was available for emptiness checks, which was not guaranteed to always give a conclusive result.

This patch adds a Simplex class, which can be constructed using a FlatAffineConstraints, and can find an integer sample point (if one exists) using the GBR algorithm. Additionally, it adds two classes Matrix and Fraction, which are used by Simplex.

The integer emptiness check functionality can be accessed through the new FlatAffineConstraints::isIntegerEmpty() function, which runs the existing heuristic first and, if that proves to be inconclusive, runs the GBR algorithm to produce a conclusive result.

[Kayjukh]

It looks like there are clang-format errors (https://results.llvm-merge-guard.org/amd64_debian_testing_clang-1651/summary.html). Can you format your changes using git-clang-format HEAD^ and update the diff? This would allow us to have a clean build.

[arjunp]

Fix some whitespace and clang-format.

[ftynse]

First of all, thanks for your contribution! This is an important and hard piece of polyhedral machinery.

The sheer size of the patch is huge and this will take some time to review. I only looked at the +/-style and simple issues and did not go into details of the algorithm itself. It would be helpful to know if you ran some code coverage tools on it.

Here is a list of major observations, which I did not always repeat:

• I am concerned by 32 as default number of stack elements in SmallVector, do you have evidence for this specific number?
• Please drop most of llvm:: prefixes, common classes are re-exported in the mlir:: namespace.
• In general, all top-level entities as well as class methods must have documentation.
• Consider making the API surface, both public and private, of classes smaller and using auxiliary functions with static linkage instead; same goes for multiple in-line lambdas that capture the world by-reference. These functions only need to be methods if they access class state, everything else can take that state as (const when possible) arguments and be cleaner about mutating it.
• Some additional documentation on the practical implementation details is necessary, in particular, the ownership model on the tableau and the bitwise negation in indexing.

Please bear with us.

mlir/include/mlir/Analysis/AffineStructures.h

150

Multiple common LLVM ADT structures are re-exported in mlir namespace, see the full list here https://github.com/llvm/llvm-project/blob/master/mlir/include/mlir/Support/LLVM.h. We generally avoid prefixing these with llvm:: or mlir:: for brevity.

mlir/include/mlir/Analysis/Presburger/Fraction.h

18

Prefer <cstdint> in C++ code for namespacing reasons

22

This needs a doc

34

This may still be subject to overflows. I'm fine with documenting this behavior somewhere and leaving as is, but we need to be aware of the problem. Have you considered performance implications of basic Fraction on APInt instead of int64?

mlir/include/mlir/Analysis/Presburger/Matrix.h

36

All public methods should have at least some documentation.

67

This rationale should be documented in the code. The interplay between resizes and cache behavior in case of consecutive reads (which I suppose are much more common) is non-trivial, especially if we allow reserve-ing space. I'm also slightly concerned about using 8k bytes on stack by default, we rarely use more than 8 stack elements in small vectors in MLIR.

149

Please keep the print. It allows one to print objects into streams different than llvm::errs().

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

37–41

This comment lost me about what are unknowns, variables and constraints, and what is the relation between them, even though I know the theory behind this.

47

Could you please elaborate a little about this bit-completent indexing, e.g. that -1 gets transformed into 0 and -2 into 1 and so on, sparing you an addition. This is a clever trick, but clever tricks are very expensive to maintain

56

Nit: suffix Op is pervasively used for IR Operations, please consider a different name

96

Nit: null value -> None

113

Nit: empty llvm::Optional -> None, here and below

127

Extra points: use a bit field or a uint8_t for flags

131

Nit: -1u (or UINT_MAX if you prefer)

133

Please, provide documentation about this class. Especially the ownership model and how it is maintained.

187

Prefer mlir::LogicalResult for success/failure conditions to avoid deadling with LLVM's changing mapping of boolean values to success status

203

Nit: true _if_ value is positive

218

General design question: how much of these functions need to be private methods, i.e. have access to class members, as opposed to be just free functions, static to the compilation unit, that actual methods can call ?

237

Why is this int when pretty much everything else uses either unsigned or int64_t?

mlir/lib/Analysis/Presburger/Matrix.cpp

19

unsigned is used everywhere else

77

nit: use single quotes for one character

mlir/lib/Analysis/Presburger/Simplex.cpp

2

Drop "C++" in .cpp files, it's only necessary for headers (.h) so that tooling can tell between C and C++ that share the extension

15

Is this include necessary?

417

Please use braces symmetrically, i.e. use them if at least one if/else block needs them

534

Why does this copy the first two columns from A?

563

Please use braces symmetrically; here and below

591

Use SmallVectorImpl with references so you don't need to hardcode the size

607

Nit: consider a ternary expression here

614

This class needs a better description of what are the invariants on the snapshot stack. This method rolls back one snapshot, which happens to be be before an equality was added, but without proper documentation, there is a risk this might be broken in future.

628

Don't evaluate the exit condition on every iteration, use

for (unsigned i = 0, e = dir.size(); i < e; ++i)

697

Please consider a lengthier name for this and all derived objects, e.g. fAtI

758

Nit: asserts are for things that should never happen, no need to mention it in the description

834

This function would be cleaner if it took opt as argument and did not capture the world by-reference. It could then also go outside as a helper with static linkage

867

MLIR generally discourages recursive functions on anything that runs on pieces of IR, stack overflows in the compiler are bad. Any chance of making this iterative? Otherwise, please provide some textual argument on how fast the recursion stops.

883

Nit: print may be used outside dump, so I'd abstain from refering to "Dumping" here

[arjunp]

Add two tests of "long" sets. Fix some bugs that this revealed.

[bondhugula]

I am happy to see this submitted for review. An exact integer emptiness check had been missing and having one would be a really useful contribution. Although expensive in theory, for many of the use cases here (which are typically low / fixed dimensionality ones), things may be practical and fast enough. It just has be carefully used so as to not be detrimental to compile time for "standard" enough use cases - since there may be alternative faster approaches. It's nevertheless useful to have this technique at our disposal and have this maintained and optimized for use cases of interest, including for experimental/research purposes.

I strongly agree with all of @ftynse's comments; they already cover pretty much all of what I may have wanted to say. Just one additional comment below on the storage for Matrix.

mlir/include/mlir/Analysis/Presburger/Matrix.h

67

Please avoid using vector of vector or SmallVector of SmallVector here. (A SmallVector isn't really optimized to hold other vectors or SmallVectors as its elements.) All of these would be bad enough for performance to warrant using a 1-d storage to start with. Using a single (contiguous) buffer for all the coefficients (a single allocation holding it) will give you cache locality (spatial), fewer/no conflict misses, better prefetching, and fewer TLB misses (just one in nearly all cases). You'll also get completely inline / stack allocations for small matrices. Your API won't be really impacted. FlatAffineConstraints uses such a single buffer of coefficients via a SmallVector (you can actually copy over its relevant methods - instead of equalities and inequalities, there's just coefficients here.) Resizing isn't really a downside with this. Adding rows is straightforward since they are at the tail, and adding columns will require movement within the contiguous buffer, which is fine. FAC reserves columns to reduce movement most of the time for column addition. With that, I don't see any additional resizing or movement penalty with a 1-d buffer over vector of vectors - you have too many allocations with the latter.

[bondhugula]

mlir/include/mlir/Analysis/Presburger/Matrix.h

67

Esp. for mid-size matrices like 10x10 or 20x20, all of the allocs and deallocs from a smallvector of smallvector will be a heap killer!

[mehdi_amini]

mlir/lib/Analysis/Presburger/Simplex.cpp

628

or better with range-style: for (int i : llvm::seq<int>(dir.size())) :)

[arjunp]

Thanks for your review. Apologies for the delayed response.

• Here is a code coverage report.
• The lengths of these SmallVectors corresponds to the dimensionality. I've changed the length parameter to 8 for vectors and 64 for the 2D matrix, as it is in FlatAffineConstraints.
• The lambdas in makeProduct and reduceBasis are very specific to these functions and access a lot of the local variables present in those functions. Does it still make sense to pull these out into free functions?
• Added more documentation for the Simplex class.

mlir/include/mlir/Analysis/Presburger/Fraction.h

34

Added documentation.

56

I didn't document these overloads as it's unclear to me what the documentation could add. Is it necessary to write something for these?

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

218

I move flippedDirection and signMatchesDirection out into free functions. reduceBasis does need access to this, but only to construct the GBRSimplex.

237

Added documentation. It's not unsigned because it's convenient to use sign comparisons to determine if an index corresponds to a constraint or a variable.

mlir/lib/Analysis/Presburger/Simplex.cpp

498

Pulling this out would require making Simplex::Unknown public or templating on the store type.

534

Nice catch. Fixed.

607

(Not applicable anymore.)

697

Changed to width. It's not strictly speaking the width but the width under some constraints, but I guess this is a good enough approximation. (and it is already referred to as "width" in computeWidthAndDuals).

834

It needs access to this for computeOptimum. Factored out into a member function computeIntegerBounds.

867

The depth of the recursion is at most the dimensionality of the set, which is not expected to be greater than say 32. This could be made iterative, albeit with some loss in readability.

[arjunp]

This comment was removed by arjunp.

[mehdi_amini]

mlir/lib/Analysis/Presburger/Simplex.cpp

867

Please document clearly the bound on the recursion. What impacts the dimensionality of the set here? When you "is not expected to be greater than say 32", how can it grow larger?

[mehdi_amini]

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

118

s/particlar/particular/

123

Nice doc for the overall class! Thanks :)

mlir/lib/Analysis/Presburger/Simplex.cpp

474

Please avoid auto (here and everywhere applicable) when it doesn't improve readability (unsigned here right?)

[bondhugula]

mlir/lib/Analysis/Presburger/Simplex.cpp

908

Please avoid auto here and similarly below.

945

Please do not return SmallVectors this way - return value optimization here will not happen I believe leading to too many copies. Instead pass SmallVectorImpl<int64_t> by reference as an output argument. @mehdi_amini or @rriddle can confirm on the style part and RVO.

mlir/unittests/Analysis/AffineStructuresTest.cpp

35

Avoid auto here for readability.

39

Avoid auto.

115

Please terminate comments with a period when it ends.

[bondhugula]

mlir/lib/Analysis/Presburger/Simplex.cpp

74

Nit: You could leave the first three lines as the doc comment and move the ones below inline into the method as regular comments.

[arjunp]

Address comments. Made findIntegerSample iterative.

[arjunp]

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

123

:)

mlir/lib/Analysis/Presburger/Simplex.cpp

867

It could theoretically grow arbitrarily large, although it may not happen in practice. I've made findIntegerSample iterative.

945

(I think this is no longer applicable since findIntegerSample is iterative now)

[bondhugula]

Thanks very much for all the changes. A few more requests - a couple on test cases and other minor ones.

mlir/unittests/Analysis/AffineStructuresTest.cpp

179

I think you are missing a test case involving equalities that is integer empty but not rational empty - a flat one along the lines of the test case added with purely inequalities. It could just be a 1-d line segment in a 2-d space that passes in between and thus misses all integer points - it could be potentially slanted as well.

mlir/unittests/Analysis/Presburger/SimplexTest.cpp

17

It may be good to add a line on what the test is testing.

87–88

Missing a test case here that involves equalities?

[bondhugula]

mlir/lib/Analysis/Presburger/Simplex.cpp

455–459

A short comment here please.

576

Comment here please.

mlir/unittests/Analysis/Presburger/SimplexTest.cpp

23

Avoid auto here please.

[bondhugula]

Not for this patch, but if you want to just quickly try and test it on more cases, you could just replace mlir::boundCheckLoadOrStoreOp's isEmpty() with isIntegerEmpty() and see if -test-memref-bound-check test cases pass. The -test-memref-bound-check pass tests for out of bound memref accesses and actually needs an integer exact method (whenever it determines non-empty) to be accurate. This way you would also have a speed comparison, and it would also be a practical application of this technique for an analysis pass.

tests/Transforms/memref-bound-check.mlir

[arjunp]

Addressed comments.

[arjunp]

I've added some more tests and documentation as requested.

[bondhugula]

mlir/unittests/Analysis/AffineStructuresTest.cpp

187

I'm sorry this would be caught by the GCD test itself. Instead, could you please replace the equality with something to the effect:

3x - 3y >= -2
3x - 3y  <= -1

This should be empty as well but would exercise your method given how isIntegerEmpty() is structured now.

[bondhugula]

mlir/include/mlir/Analysis/Presburger/Fraction.h

66

Are you assuming f.den is always positive? If both num and den are both negative, you'd want to add one.

[arjunp]

Added the requested test.

[arjunp]

Fix a comment in Fraction.

[arjunp]

mlir/include/mlir/Analysis/Presburger/Fraction.h

66

Yes, den is always positive. There was an incomplete comment for the member declaration that I've fixed. (This is also documented in the top-level class documentation)

mlir/unittests/Analysis/AffineStructuresTest.cpp

187

Since checkSample directly calls findIntegerSample() rather than isIntegerEmpty(), this shouldn't be an issue. Nonetheless, I've also added the test you requested.

[ftynse]

Thanks for providing more detailed documentation!

I did not extensively analyze the algorithms. At a glance, the implementation looks plausible. Again, given the complexity, I would encourage you to get 100% test coverage on the core functions. Another way to stress-test this would be to generate a bunch of arbitrary sets and run this + gcd emptiness checks and analyze the differences manually to make sure they only happen when gcd is not supposed to be correct.

I don't have any further comments.

@andydavis1 had a bug which might have been related to integer exactness, but I don't know if it is still active.

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

38

Nit: SImplex -> simplex

46

Nit: 2x - 3y >= 0 (currently uses z)

[bondhugula]

mlir/include/mlir/Analysis/Presburger/Fraction.h

66

If den is positive, this can be simplified to avoid an extra check.

f.num % f.den > 0 ? f.num / f.den + 1 : f.num / f.den;

(one comparison instead of two comparisons)

Also, if there is a doc comment here, please mention that den is positive. Actually, you don't even need this logic; mlir::ceilDiv does exactly this.

[bondhugula]

In D80860#2106414, @ftynse wrote:

Thanks for providing more detailed documentation!

I did not extensively analyze the algorithms. At a glance, the implementation looks plausible. Again, given the complexity, I would encourage you to get 100% test coverage

+1 I don't think it's practical to have a review checking each of the steps of the algorithm. I would just rely on (1) good test case coverage, and (2) sufficient documentation/comments for now. Other implementation aspects can be brought up and revised later when there are other users with a specialized interest in digging deeper.

[andydavis1]

Thanks for your contribution!

It would be great to have tests on integer sets which resulting in an over approximation using the existing FM algorithm, but which returned integer exact results using GBR from this change.

Also wondering what the plan is to check for integer overflow of coefficient values for some of these algorithms?

[arjunp]

Improved test coverage. There is still one uncovered line; it seems quite difficult to find a case that covers this. I ran hundreds of tests with randomly generated sets intersected with a 1000 x 1000 x 1000 cube. Whenever GBR and the FM/GCD-based algorithm disagreed, I checked against a brute-force that checks every point in this cube. All these tests passed, as well as tests from some other types of generated cases.

Added some tests that GBR detects as empty but the existing FM/GCD-based algorithm does not.

Regarding overflow issues, it does not seem to be something specific to this patch. It looks to me like there already exists other functionality in FlatAffineConstraints, like FM, that performs operations without complete overflow checks. I feel that all these need to be addressed together by, say, introducing arbitrary-precision arithmetic, which is beyond the scope of this patch.

Thanks again for all your help with improving this patch!

[bondhugula]

mlir/include/mlir/Analysis/Presburger/Fraction.h

39

This has to be a /// comment.

[bondhugula]

In D80860#2115848, @arjunp wrote:

Added some tests that GBR detects as empty but the existing FM/GCD-based algorithm does not.

It'll be useful if you marked these tests with a comment to that effect. I might have missed it if you already did.

[bondhugula]

A bunch of mostly minor comments.

mlir/lib/Analysis/Presburger/Simplex.cpp

707

Can you rephrase this? <dir, x-y> is unclear.

711

Please use coeffs.reserve since you know big it's going to be.

712–714

Why emplace_back instead of push_back for a POD?

719–723

Triple comments please.

866–867

Reflow - use the width?

875

Please make this a complete sentence.

961

You can use llvm::to_vector.

1011

llvm::to_vector

mlir/unittests/Analysis/AffineStructuresTest.cpp

30

Nit: hasValue -> 'hasValue'

56

SmallVector -> ArrayRef

57

Likewise.

68

Typo: least

72–73

Please pass ArrayRef instead of SmallVector for inputs. (If these were output args or inouts, SmallVectorImpl is to be used.)

mlir/unittests/Analysis/Presburger/SimplexTest.cpp

54

SmallVector 4?

60

Likewise.

[arjunp]

Addressed comments.

[arjunp]

In D80860#2118522, @bondhugula wrote:

In D80860#2115848, @arjunp wrote:

Added some tests that GBR detects as empty but the existing FM/GCD-based algorithm does not.

It'll be useful if you marked these tests with a comment to that effect. I might have missed it if you already did.

This is marked at AffineStructuresTest.cpp:247.

mlir/lib/Analysis/Presburger/Simplex.cpp

707

I've replaced it with the more explicit dotProduct everywhere now. (this was originally explained in the top-level class comment)

mlir/unittests/Analysis/AffineStructuresTest.cpp

65

Is there a better way? This is needed since we moved to ArrayRefs for arguments but in checkPermutationsSample below we need to pass SmallVectors of SmallVectors to makeFACFromConstraints.

[bondhugula]

mlir/unittests/Analysis/AffineStructuresTest.cpp

65

We don't need to switch to ArrayRef of ArrayRef if you have SmallVector of SmallVector, but only to ArrayRef of SmallVector.

[bondhugula]

mlir/unittests/Analysis/AffineStructuresTest.cpp

65

I actually meant ArrayRef<SmallVector,...> as the function argument instead of SmallVector<SmallVector<.... Make sure you don't run into underlying storage lifetime issues. ArrayRef<SmallVector..> should be fine.

[arjunp]

Switched from ArrayRef of ArrayRef to ArrayRef of SmallVectors.

[ftynse]

@Kayjukh your review is blocking this revision, have your concerns been addressed?

[Kayjukh]

Yes, sorry about blocking the revision. Note that there is still a clang-format issue in one of the files.

mlir/unittests/Analysis/AffineStructuresTest.cpp

89

Please fix formatting

[arjunp]

clang-format.

[arjunp]

I don't have commit access yet, so I think someone else will need to land this.

[bondhugula]

In D80860#2127981, @arjunp wrote:

I don't have commit access yet, so I think someone else will need to land this.

I've committed this.

[mehdi_amini]

mlir/lib/Analysis/Presburger/Simplex.cpp

131

What guarantees that gcd can't be 0 and produce an invalid division here?
Should we guard this loop?

(flagged by Coverity)

[mehdi_amini]

mlir/lib/Analysis/Presburger/Simplex.cpp

325

If retRow is true, then retConst is used uninitialized right? Can you look into this?
(flagged by Coverity)

325

(Same with retElem)

[arjunp]

mlir/lib/Analysis/Presburger/Simplex.cpp

131

One of the elements being GCD'd is a "denominator" value that is always non-zero. I will add a comment noting this.

325

If retRow has a value then lines 319-322 were executed in the iteration that that values was set in, so the values of retElem and retConst were also set in that iteration already.

[grosser]

Hi @arjunp, I am wondering if there are ways to prevent future coverity warnings on this code?

mlir/lib/Analysis/Presburger/Simplex.cpp

131

Given that these are coverity errors, it might be useful to document this with an assert to ensure that coverity does not warn about this anymore.

325

Again, as these are coverity errors the question is how we can avoid coverity to flag these issues. I am almost tempted to just suggest to set retElement and retConst to 0 and add a comment stating that this is done to avoid the coverity warning.

[arjunp]

mlir/lib/Analysis/Presburger/Simplex.cpp

131

Done: https://reviews.llvm.org/rG0f925339e52e361946a1e894cfa6dbe494b39290

325

Done: https://reviews.llvm.org/rGef8351598ef33b906fd1962420c8b464812cea85

mlir/include/mlir/Analysis/AffineStructures.h

public:
 
   void reset(unsigned numDims = 0, unsigned numSymbols = 0,
              unsigned numLocals = 0, ArrayRef<Value> idArgs = {});
 
   /// Appends constraints from 'other' into this. This is equivalent to an
   /// intersection with no simplification of any sort attempted.
   void append(const FlatAffineConstraints &other);
 
-  // Checks for emptiness by performing variable elimination on all identifiers,
-  // running the GCD test on each equality constraint, and checking for invalid
-  // constraints.
-  // Returns true if the GCD test fails for any equality, or if any invalid
-  // constraints are discovered on any row. Returns false otherwise.
+  /// Checks for emptiness by performing variable elimination on all
+  /// identifiers, running the GCD test on each equality constraint, and
+  /// checking for invalid constraints. Returns true if the GCD test fails for
+  /// any equality, or if any invalid constraints are discovered on any row.
+  /// Returns false otherwise.
   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.
+  /// 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.
+  bool isIntegerEmpty() 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.
+  ///
+  /// Returns such a point if one exists, or an empty Optional otherwise.
+  Optional<SmallVector<int64_t, 8>> findIntegerSample() const;
+
   // Clones this object.
   std::unique_ptr<FlatAffineConstraints> clone() const;
 
   /// Returns the value at the specified equality row and column.
   inline int64_t atEq(unsigned i, unsigned j) const {
     return equalities[i * numReservedCols + j];
   }
   inline int64_t &atEq(unsigned i, unsigned j) {

mlir/include/mlir/Analysis/Presburger/Fraction.h

+//===- Fraction.h - MLIR Fraction 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
+//
+//===----------------------------------------------------------------------===//
+//
+// This is a simple class to represent fractions. It supports multiplication,
+// comparison, floor, and ceiling operations.
+//
+//===----------------------------------------------------------------------===//
+
+#ifndef MLIR_ANALYSIS_PRESBURGER_FRACTION_H
+#define MLIR_ANALYSIS_PRESBURGER_FRACTION_H
+
+#include "mlir/Support/MathExtras.h"
+
+namespace mlir {
+
+/// A class to represent fractions. The sign of the fraction is represented
+/// in the sign of the numerator; the denominator is always positive.
+///
+/// Note that overflows may occur if the numerator or denominator are not
+/// representable by 64-bit integers.
+struct Fraction {
+  /// Default constructor initializes the represented rational number to zero.
+  Fraction() : num(0), den(1) {}
+
+  /// Construct a Fraction from a numerator and denominator.
+  Fraction(int64_t oNum, int64_t oDen) : num(oNum), den(oDen) {
+    if (den < 0) {
+      num = -num;
+      den = -den;
+    }
+  }
+
+  /// The numerator and denominator, respectively. The denominator is always
+  /// positive.
+  int64_t num, den;
+};
+
+/// Three-way comparison between two fractions.
+/// Returns +1, 0, and -1 if the first fraction is greater than, equal to, or
+/// less than the second fraction, respectively.
+inline int compare(Fraction x, Fraction y) {
+  int64_t diff = x.num * y.den - y.num * x.den;
+  if (diff > 0)
+    return +1;
+  if (diff < 0)
+    return -1;
+  return 0;
+}
+
+inline int64_t floor(Fraction f) { return floorDiv(f.num, f.den); }
+
+inline int64_t ceil(Fraction f) { return ceilDiv(f.num, f.den); }
+
+inline Fraction operator-(Fraction x) { return Fraction(-x.num, x.den); }
+
+inline bool operator<(Fraction x, Fraction y) { return compare(x, y) < 0; }
+
+inline bool operator<=(Fraction x, Fraction y) { return compare(x, y) <= 0; }
+
+inline bool operator==(Fraction x, Fraction y) { return compare(x, y) == 0; }
+
+inline bool operator>(Fraction x, Fraction y) { return compare(x, y) > 0; }
+
+inline bool operator>=(Fraction x, Fraction y) { return compare(x, y) >= 0; }
+
+inline Fraction operator*(Fraction x, Fraction y) {
+  return Fraction(x.num * y.num, x.den * y.den);
+}
+
+} // namespace mlir
+
+#endif // MLIR_ANALYSIS_PRESBURGER_FRACTION_H

mlir/include/mlir/Analysis/Presburger/Matrix.h

+//===- Matrix.h - MLIR Matrix 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
+//
+//===----------------------------------------------------------------------===//
+//
+// This is a simple 2D matrix class that supports reading, writing, resizing,
+// swapping rows, and swapping columns.
+//
+//===----------------------------------------------------------------------===//
+
+#ifndef MLIR_ANALYSIS_PRESBURGER_MATRIX_H
+#define MLIR_ANALYSIS_PRESBURGER_MATRIX_H
+
+#include "mlir/Support/LLVM.h"
+#include "llvm/ADT/ArrayRef.h"
+#include "llvm/Support/raw_ostream.h"
+
+#include <cassert>
+
+namespace mlir {
+
+/// This is a simple class to represent a resizable matrix.
+///
+/// The data is stored in the form of a vector of vectors.
+class Matrix {
+public:
+  Matrix() = delete;
+
+  /// Construct a matrix with the specified number of rows and columns.
+  /// Initially, the values are default initialized.
+  Matrix(unsigned rows, unsigned columns);
+
+  /// Return the identity matrix of the specified dimension.
+  static Matrix identity(unsigned dimension);
+
+  /// Access the element at the specified row and column.
+  int64_t &at(unsigned row, unsigned column);
+  int64_t at(unsigned row, unsigned column) const;
+  int64_t &operator()(unsigned row, unsigned column);
+  int64_t operator()(unsigned row, unsigned column) const;
+
+  /// Swap the given columns.
+  void swapColumns(unsigned column, unsigned otherColumn);
+
+  /// Swap the given rows.
+  void swapRows(unsigned row, unsigned otherRow);
+
+  unsigned getNumRows() const;
+
+  unsigned getNumColumns() const;
+
+  /// Get an ArrayRef corresponding to the specified row.
+  ArrayRef<int64_t> getRow(unsigned row) const;
+
+  /// Add `scale` multiples of the source row to the target row.
+  void addToRow(unsigned sourceRow, unsigned targetRow, int64_t scale);
+
+  /// Resize the matrix to the specified dimensions. If a dimension is smaller,
+  /// the values are truncated; if it is bigger, the new values are default
+  /// initialized.
+  void resizeVertically(unsigned newNRows);
+
+  /// Print the matrix.
+  void print(raw_ostream &os) const;
+  void dump() const;
+
+private:
+  unsigned nRows, nColumns;
+
+  /// Stores the data. data.size() is equal to nRows * nColumns.
+  SmallVector<int64_t, 64> data;
+};
+
+} // namespace mlir
+
+#endif // MLIR_ANALYSIS_PRESBURGER_MATRIX_H

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

+//===- Simplex.h - MLIR Simplex 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
+//
+//===----------------------------------------------------------------------===//
+//
+// Functionality to perform analysis on FlatAffineConstraints. In particular,
+// support for performing emptiness checks.
+//
+//===----------------------------------------------------------------------===//
+
+#ifndef MLIR_ANALYSIS_PRESBURGER_SIMPLEX_H
+#define MLIR_ANALYSIS_PRESBURGER_SIMPLEX_H
+
+#include "mlir/Analysis/AffineStructures.h"
+#include "mlir/Analysis/Presburger/Fraction.h"
+#include "mlir/Analysis/Presburger/Matrix.h"
+#include "mlir/Support/LogicalResult.h"
+#include "llvm/ADT/ArrayRef.h"
+#include "llvm/ADT/Optional.h"
+#include "llvm/ADT/SmallVector.h"
+#include "llvm/Support/raw_ostream.h"
+
+namespace mlir {
+
+class GBRSimplex;
+
+/// This class implements a version of the Simplex and Generalized Basis
+/// Reduction algorithms, which can perform analysis of integer sets with affine
+/// inequalities and equalities. A Simplex can be constructed
+/// by specifying the dimensionality of the set. It supports adding affine
+/// inequalities and equalities, and can perform emptiness checks, i.e., it can
+/// find a solution to the set of constraints if one exists, or say that the
+/// set is empty if no solution exists. Currently, this only works for bounded
+/// sets. Simplex can also be constructed from a FlatAffineConstraints object.
+///
+/// The implementation of this Simplex class, other than the functionality
+/// for sampling, is based on the paper
+/// "Simplify: A Theorem Prover for Program Checking"
+/// by D. Detlefs, G. Nelson, J. B. Saxe.
+///
+/// We define variables, constraints, and unknowns. Consider the example of a
+/// two-dimensional set defined by 1 + 2x + 3y >= 0 and 2x - 3y >= 0. Here,
+/// x, y, are variables while 1 + 2x + 3y >= 0, 2x - 3y >= 0 are
+/// constraints. Unknowns are either variables or constraints, i.e., x, y,
+/// 1 + 2x + 3y >= 0, 2x - 3y >= 0 are all unknowns.
+///
+/// The implementation involves a matrix called a tableau, which can be thought
+/// of as a 2D matrix of rational numbers having number of rows equal to the
+/// number of constraints and number of columns equal to one plus the number of
+/// variables. In our implementation, instead of storing rational numbers, we
+/// store a common denominator for each row, so it is in fact a matrix of
+/// integers with number of rows equal to number of constraints and number of
+/// columns equal to _two_ plus the number of variables. For example, instead of
+/// storing a row of three rationals [1/2, 2/3, 3], we would store [6, 3, 4, 18]
+/// since 3/6 = 1/2, 4/6 = 2/3, and 18/6 = 3.
+///
+/// Every row and column except the first and second columns is associated with
+/// an unknown and every unknown is associated with a row or column. The second
+/// column represents the constant, explained in more detail below. An unknown
+/// associated with a row or column is said to be in row or column position
+/// respectively.
+///
+/// The vectors var and con store information about the variables and
+/// constraints respectively, namely, whether they are in row or column
+/// position, which row or column they are associated with, and whether they
+/// correspond to a variable or a constraint.
+///
+/// An unknown is addressed by its index. If the index i is non-negative, then
+/// the variable var[i] is being addressed. If the index i is negative, then
+/// the constraint con[~i] is being addressed. Effectively this maps
+/// 0 -> var[0], 1 -> var[1], -1 -> con[0], -2 -> con[1], etc. rowUnknown[r] and
+/// colUnknown[c] are the indexes of the unknowns associated with row r and
+/// column c, respectively.
+///
+/// The unknowns in column position are together called the basis. Initially the
+/// basis is the set of variables -- in our example above, the initial basis is
+/// x, y.
+///
+/// The unknowns in row position are represented in terms of the basis unknowns.
+/// If the basis unknowns are u_1, u_2, ... u_m, and a row in the tableau is
+/// d, c, a_1, a_2, ... a_m, this representats the unknown for that row as
+/// (c + a_1*u_1 + a_2*u_2 + ... + a_m*u_m)/d. In our running example, if the
+/// basis is the initial basis of x, y, then the constraint 1 + 2x + 3y >= 0
+/// would be represented by the row [1, 1, 2, 3].
+///
+/// The association of unknowns to rows and columns can be changed by a process
+/// called pivoting, where a row unknown and a column unknown exchange places
+/// and the remaining row variables' representation is changed accordingly
+/// by eliminating the old column unknown in favour of the new column unknown.
+/// If we had pivoted the column for x with the row for 2x - 3y >= 0,
+/// the new row for x would be [2, 1, 3] since x = (1*(2x - 3y) + 3*y)/2.
+/// See the documentation for the pivot member function for details.
+///
+/// The association of unknowns to rows and columns is called the _tableau
+/// configuration_. The _sample value_ of an unknown in a particular tableau
+/// configuration is its value if all the column unknowns were set to zero.
+/// Concretely, for unknowns in column position the sample value is zero and
+/// for unknowns in row position the sample value is the constant term divided
+/// by the common denominator.
+///
+/// The tableau configuration is called _consistent_ if the sample value of all
+/// restricted unknowns is non-negative. Initially there are no constraints, and
+/// the tableau is consistent. When a new constraint is added, its sample value
+/// in the current tableau configuration may be negative. In that case, we try
+/// to find a series of pivots to bring us to a consistent tableau
+/// configuration, i.e. we try to make the new constraint's sample value
+/// non-negative without making that of any other constraints negative. (See
+/// findPivot and findPivotRow for details.) If this is not possible, then the
+/// set of constraints is mutually contradictory and the tableau is marked
+/// _empty_, which means the set of constraints has no solution.
+///
+/// This Simplex class also supports taking snapshots of the current state
+/// and rolling back to prior snapshots. This works by maintaing an undo log
+/// of operations. Snapshots are just pointers to a particular location in the
+/// log, and rolling back to a snapshot is done by reverting each log entry's
+/// operation from the end until we reach the snapshot's location.
+///
+/// Finding an integer sample is done with the Generalized Basis Reduction
+/// algorithm. See the documentation for findIntegerSample and reduceBasis.
+class Simplex {
+public:
+  enum class Direction { Up, Down };
+
+  Simplex() = delete;
+  explicit Simplex(unsigned nVar);
+  explicit Simplex(const FlatAffineConstraints &constraints);
+
+  /// Returns true if the tableau is empty (has conflicting constraints),
+  /// false otherwise.
+  bool isEmpty() const;
+
+  /// Add an inequality to the tableau. If coeffs is c_0, c_1, ... c_n, where n
+  /// is the current number of variables, then the corresponding inequality is
+  /// c_n + c_0*x_0 + c_1*x_1 + ... + c_{n-1}*x_{n-1} >= 0.
+  void addInequality(ArrayRef<int64_t> coeffs);
+
+  /// Returns the number of variables in the tableau.
+  unsigned numVariables() const;
+
+  /// Returns the number of constraints in the tableau.
+  unsigned numConstraints() const;
+
+  /// Add an equality to the tableau. If coeffs is c_0, c_1, ... c_n, where n
+  /// is the current number of variables, then the corresponding equality is
+  /// c_n + c_0*x_0 + c_1*x_1 + ... + c_{n-1}*x_{n-1} == 0.
+  void addEquality(ArrayRef<int64_t> coeffs);
+
+  /// Mark the tableau as being empty.
+  void markEmpty();
+
+  /// Get a snapshot of the current state. This is used for rolling back.
+  unsigned getSnapshot() const;
+
+  /// Rollback to a snapshot. This invalidates all later snapshots.
+  void rollback(unsigned snapshot);
+
+  /// Compute the maximum or minimum value of the given row, depending on
+  /// direction.
+  ///
+  /// Returns a (num, den) pair denoting the optimum, or None if no
+  /// optimum exists, i.e., if the expression is unbounded in this direction.
+  Optional<Fraction> computeRowOptimum(Direction direction, unsigned row);
+
+  /// Compute the maximum or minimum value of the given expression, depending on
+  /// direction.
+  ///
+  /// Returns a (num, den) pair denoting the optimum, or a null value if no
+  /// optimum exists, i.e., if the expression is unbounded in this direction.
+  Optional<Fraction> computeOptimum(Direction direction,
+                                    ArrayRef<int64_t> coeffs);
+
+  /// Returns a (min, max) pair denoting the minimum and maximum integer values
+  /// of the given expression.
+  std::pair<int64_t, int64_t> computeIntegerBounds(ArrayRef<int64_t> coeffs);
+
+  /// 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
+  /// 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.
+  void print(raw_ostream &os) const;
+  void dump() const;
+
+private:
+  friend class GBRSimplex;
+
+  enum class Orientation { Row, Column };
+
+  /// An Unknown is either a variable or a constraint. It is always associated
+  /// with either a row or column. Whether it's a row or a column is specified
+  /// by the orientation and pos identifies the specific row or column it is
+  /// associated with. If the unknown is restricted, then it has a
+  /// non-negativity constraint associated with it, i.e., its sample value must
+  /// always be non-negative and if it cannot be made non-negative without
+  /// violating other constraints, the tableau is empty.
+  struct Unknown {
+    Unknown(Orientation oOrientation, bool oRestricted, unsigned oPos)
+        : pos(oPos), orientation(oOrientation), restricted(oRestricted) {}
+    unsigned pos;
+    Orientation orientation;
+    bool restricted : 1;
+
+    void print(raw_ostream &os) const {
+      os << (orientation == Orientation::Row ? "r" : "c");
+      os << pos;
+      if (restricted)
+        os << " [>=0]";
+    }
+  };
+
+  struct Pivot {
+    unsigned row, column;
+  };
+
+  /// Find a pivot to change the sample value of row in the specified
+  /// direction. The returned pivot row will be row if and only
+  /// if the unknown is unbounded in the specified direction.
+  ///
+  /// Returns a (row, col) pair denoting a pivot, or an empty Optional if
+  /// no valid pivot exists.
+  Optional<Pivot> findPivot(int row, Direction direction) const;
+
+  /// Swap the row with the column in the tableau's data structures but not the
+  /// tableau itself. This is used by pivot.
+  void swapRowWithCol(unsigned row, unsigned col);
+
+  /// Pivot the row with the column.
+  void pivot(unsigned row, unsigned col);
+  void pivot(Pivot pair);
+
+  /// Returns the unknown associated with index.
+  const Unknown &unknownFromIndex(int index) const;
+  /// Returns the unknown associated with col.
+  const Unknown &unknownFromColumn(unsigned col) const;
+  /// Returns the unknown associated with row.
+  const Unknown &unknownFromRow(unsigned row) const;
+  /// Returns the unknown associated with index.
+  Unknown &unknownFromIndex(int index);
+  /// Returns the unknown associated with col.
+  Unknown &unknownFromColumn(unsigned col);
+  /// Returns the unknown associated with row.
+  Unknown &unknownFromRow(unsigned row);
+
+  /// Add a new row to the tableau and the associated data structures.
+  unsigned addRow(ArrayRef<int64_t> coeffs);
+
+  /// Normalize the given row by removing common factors between the numerator
+  /// and the denominator.
+  void normalizeRow(unsigned row);
+
+  /// Swap the two rows in the tableau and associated data structures.
+  void swapRows(unsigned i, unsigned j);
+
+  /// Restore the unknown to a non-negative sample value.
+  ///
+  /// Returns true if the unknown was successfully restored to a non-negative
+  /// sample value, false otherwise.
+  LogicalResult restoreRow(Unknown &u);
+
+  enum class UndoLogEntry { RemoveLastConstraint, UnmarkEmpty };
+
+  /// Undo the operation represented by the log entry.
+  void undo(UndoLogEntry entry);
+
+  /// Find a row that can be used to pivot the column in the specified
+  /// direction. If skipRow is not null, then this row is excluded
+  /// from consideration. The returned pivot will maintain all constraints
+  /// except the column itself and skipRow, if it is set. (if these unknowns
+  /// are restricted).
+  ///
+  /// Returns the row to pivot to, or an empty Optional if the column
+  /// is unbounded in the specified direction.
+  Optional<unsigned> findPivotRow(Optional<unsigned> skipRow,
+                                  Direction direction, unsigned col) const;
+
+  /// Reduce the given basis, starting at the specified level, using general
+  /// basis reduction.
+  void reduceBasis(Matrix &basis, unsigned level);
+
+  /// The number of rows in the tableau.
+  unsigned nRow;
+
+  /// The number of columns in the tableau, including the common denominator
+  /// and the constant column.
+  unsigned nCol;
+
+  /// The matrix representing the tableau.
+  Matrix tableau;
+
+  /// This is true if the tableau has been detected to be empty, false
+  /// otherwise.
+  bool empty;
+
+  /// Holds a log of operations, used for rolling back to a previous state.
+  SmallVector<UndoLogEntry, 8> undoLog;
+
+  /// These hold the indexes of the unknown at a given row or column position.
+  /// We keep these as signed integers since that makes it convenient to check
+  /// if an index corresponds to a variable or a constraint by checking the
+  /// sign.
+  ///
+  /// colUnknown is padded with two null indexes at the front since the first
+  /// two columns don't correspond to any unknowns.
+  SmallVector<int, 8> rowUnknown, colUnknown;
+
+  /// These hold information about each unknown.
+  SmallVector<Unknown, 8> con, var;
+};
+
+} // namespace mlir
+
+#endif // MLIR_ANALYSIS_PRESBURGER_SIMPLEX_H

mlir/lib/Analysis/AffineStructures.cpp

 //===- AffineStructures.cpp - MLIR Affine Structures 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
 //
 //===----------------------------------------------------------------------===//
 //
 // Structures for affine/polyhedral analysis of affine dialect ops.
 //
 //===----------------------------------------------------------------------===//
 
 #include "mlir/Analysis/AffineStructures.h"
+#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/SmallPtrSet.h"
for (unsigned i = 0, e = getNumEqualities(); i < e; ++i) {
     int64_t v = std::abs(atEq(i, numCols - 1));
     if (gcd > 0 && (v % gcd != 0)) {
       return true;
     }
   }
   return false;
 }
 
+// First, try the GCD test heuristic.
+//
+// If that doesn't find the set empty, check if the set is unbounded. If it is,
+// we cannot use the GBR algorithm and we conservatively return false.
+//
+// If the set is bounded, we use the complete emptiness check for this case
+// provided by Simplex::findIntegerSample(), which gives a definitive answer.
+bool FlatAffineConstraints::isIntegerEmpty() const {
+  if (isEmptyByGCDTest())
+    return true;
+
+  Simplex simplex(*this);
+  if (simplex.isUnbounded())
+    return false;
+  return !simplex.findIntegerSample().hasValue();
+}
+
+Optional<SmallVector<int64_t, 8>>
+FlatAffineConstraints::findIntegerSample() const {
+  return Simplex(*this).findIntegerSample();
+}
+
 /// Tightens inequalities given that we are dealing with integer spaces. This is
 /// analogous to the GCD test but applied to inequalities. The constant term can
 /// be reduced to the preceding multiple of the GCD of the coefficients, i.e.,
 ///  64*i - 100 >= 0  =>  64*i - 128 >= 0 (since 'i' is an integer). This is a
 /// fast method - linear in the number of coefficients.
 // Example on how this affects practical cases: consider the scenario:
 // 64*i >= 100, j = 64*i; without a tightening, elimination of i would yield
 // j >= 100 instead of the tighter (exact) j >= 128.

mlir/lib/Analysis/CMakeLists.txt

add_mlir_library(MLIRAnalysis
   DEPENDS
   mlir-headers
 
   LINK_LIBS PUBLIC
   MLIRAffineOps
   MLIRCallInterfaces
   MLIRControlFlowInterfaces
   MLIRInferTypeOpInterface
+  MLIRPresburger
   MLIRSCF
   )
 
 add_mlir_library(MLIRLoopAnalysis
   AffineAnalysis.cpp
   AffineStructures.cpp
   LoopAnalysis.cpp
   NestedMatcher.cpp
   Utils.cpp
 
   ADDITIONAL_HEADER_DIRS
   ${MLIR_MAIN_INCLUDE_DIR}/mlir/Analysis
 
   DEPENDS
   mlir-headers
 
   LINK_LIBS PUBLIC
   MLIRAffineOps
   MLIRCallInterfaces
   MLIRControlFlowInterfaces
   MLIRInferTypeOpInterface
+  MLIRPresburger
   MLIRSCF
   )
+  
+add_subdirectory(Presburger)

mlir/lib/Analysis/Presburger/CMakeLists.txt

+add_mlir_library(MLIRPresburger
+  Simplex.cpp
+  Matrix.cpp
+  )

mlir/lib/Analysis/Presburger/Matrix.cpp

+//===- Matrix.cpp - MLIR Matrix 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/Matrix.h"
+
+namespace mlir {
+
+Matrix::Matrix(unsigned rows, unsigned columns)
+    : nRows(rows), nColumns(columns), data(nRows * nColumns) {}
+
+Matrix Matrix::identity(unsigned dimension) {
+  Matrix matrix(dimension, dimension);
+  for (unsigned i = 0; i < dimension; ++i)
+    matrix(i, i) = 1;
+  return matrix;
+}
+
+int64_t &Matrix::at(unsigned row, unsigned column) {
+  assert(row < getNumRows() && "Row outside of range");
+  assert(column < getNumColumns() && "Column outside of range");
+  return data[row * nColumns + column];
+}
+
+int64_t Matrix::at(unsigned row, unsigned column) const {
+  assert(row < getNumRows() && "Row outside of range");
+  assert(column < getNumColumns() && "Column outside of range");
+  return data[row * nColumns + column];
+}
+
+int64_t &Matrix::operator()(unsigned row, unsigned column) {
+  return at(row, column);
+}
+
+int64_t Matrix::operator()(unsigned row, unsigned column) const {
+  return at(row, column);
+}
+
+unsigned Matrix::getNumRows() const { return nRows; }
+
+unsigned Matrix::getNumColumns() const { return nColumns; }
+
+void Matrix::resizeVertically(unsigned newNRows) {
+  nRows = newNRows;
+  data.resize(nRows * nColumns);
+}
+
+void Matrix::swapRows(unsigned row, unsigned otherRow) {
+  assert((row < getNumRows() && otherRow < getNumRows()) &&
+         "Given row out of bounds");
+  if (row == otherRow)
+    return;
+  for (unsigned col = 0; col < nColumns; col++)
+    std::swap(at(row, col), at(otherRow, col));
+}
+
+void Matrix::swapColumns(unsigned column, unsigned otherColumn) {
+  assert((column < getNumColumns() && otherColumn < getNumColumns()) &&
+         "Given column out of bounds");
+  if (column == otherColumn)
+    return;
+  for (unsigned row = 0; row < nRows; row++)
+    std::swap(at(row, column), at(row, otherColumn));
+}
+
+ArrayRef<int64_t> Matrix::getRow(unsigned row) const {
+  return {&data[row * nColumns], nColumns};
+}
+
+void Matrix::addToRow(unsigned sourceRow, unsigned targetRow, int64_t scale) {
+  if (scale == 0)
+    return;
+  for (unsigned col = 0; col < nColumns; ++col)
+    at(targetRow, col) += scale * at(sourceRow, col);
+  return;
+}
+
+void Matrix::print(raw_ostream &os) const {
+  for (unsigned row = 0; row < nRows; ++row) {
+    for (unsigned column = 0; column < nColumns; ++column)
+      os << at(row, column) << ' ';
+    os << '\n';
+  }
+}
+
+void Matrix::dump() const { print(llvm::errs()); }
+
+} // namespace mlir

mlir/lib/Analysis/Presburger/Simplex.cpp

+//===- Simplex.cpp - MLIR Simplex 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/Simplex.h"
+#include "mlir/Analysis/Presburger/Matrix.h"
+#include "mlir/Support/MathExtras.h"
+
+namespace mlir {
+using Direction = Simplex::Direction;
+
+const int nullIndex = std::numeric_limits<int>::max();
+
+/// Construct a Simplex object with `nVar` variables.
+Simplex::Simplex(unsigned nVar)
+    : nRow(0), nCol(2), tableau(0, 2 + nVar), empty(false) {
+  colUnknown.push_back(nullIndex);
+  colUnknown.push_back(nullIndex);
+  for (unsigned i = 0; i < nVar; ++i) {
+    var.emplace_back(Orientation::Column, /*restricted=*/false, /*pos=*/nCol);
+    colUnknown.push_back(i);
+    nCol++;
+  }
+}
+
+Simplex::Simplex(const FlatAffineConstraints &constraints)
+    : Simplex(constraints.getNumIds()) {
+  for (unsigned i = 0, numIneqs = constraints.getNumInequalities();
+       i < numIneqs; ++i)
+    addInequality(constraints.getInequality(i));
+  for (unsigned i = 0, numEqs = constraints.getNumEqualities(); i < numEqs; ++i)
+    addEquality(constraints.getEquality(i));
+}
+
+const Simplex::Unknown &Simplex::unknownFromIndex(int index) const {
+  assert(index != nullIndex && "nullIndex passed to unknownFromIndex");
+  return index >= 0 ? var[index] : con[~index];
+}
+
+const Simplex::Unknown &Simplex::unknownFromColumn(unsigned col) const {
+  assert(col < nCol && "Invalid column");
+  return unknownFromIndex(colUnknown[col]);
+}
+
+const Simplex::Unknown &Simplex::unknownFromRow(unsigned row) const {
+  assert(row < nRow && "Invalid row");
+  return unknownFromIndex(rowUnknown[row]);
+}
+
+Simplex::Unknown &Simplex::unknownFromIndex(int index) {
+  assert(index != nullIndex && "nullIndex passed to unknownFromIndex");
+  return index >= 0 ? var[index] : con[~index];
+}
+
+Simplex::Unknown &Simplex::unknownFromColumn(unsigned col) {
+  assert(col < nCol && "Invalid column");
+  return unknownFromIndex(colUnknown[col]);
+}
+
+Simplex::Unknown &Simplex::unknownFromRow(unsigned row) {
+  assert(row < nRow && "Invalid row");
+  return unknownFromIndex(rowUnknown[row]);
+}
+
+/// Add a new row to the tableau corresponding to the given constant term and
+/// list of coefficients. The coefficients are specified as a vector of
+/// (variable index, coefficient) pairs.
+unsigned Simplex::addRow(ArrayRef<int64_t> coeffs) {
+  assert(coeffs.size() == 1 + var.size() &&
+         "Incorrect number of coefficients!");
+
+  ++nRow;
+  // If the tableau is not big enough to accomodate the extra row, we extend it.
+  if (nRow >= tableau.getNumRows())
+    tableau.resizeVertically(nRow);
+  rowUnknown.push_back(~con.size());
+  con.emplace_back(Orientation::Row, false, nRow - 1);
+
+  tableau(nRow - 1, 0) = 1;
+  tableau(nRow - 1, 1) = coeffs.back();
+  for (unsigned col = 2; col < nCol; ++col)
+    tableau(nRow - 1, col) = 0;
+
+  // Process each given variable coefficient.
+  for (unsigned i = 0; i < var.size(); ++i) {
+    unsigned pos = var[i].pos;
+    if (coeffs[i] == 0)
+      continue;
+
+    if (var[i].orientation == Orientation::Column) {
+      // If a variable is in column position at column col, then we just add the
+      // coefficient for that variable (scaled by the common row denominator) to
+      // the corresponding entry in the new row.
+      tableau(nRow - 1, pos) += coeffs[i] * tableau(nRow - 1, 0);
+      continue;
+    }
+
+    // If the variable is in row position, we need to add that row to the new
+    // row, scaled by the coefficient for the variable, accounting for the two
+    // rows potentially having different denominators. The new denominator is
+    // the lcm of the two.
+    int64_t lcm = mlir::lcm(tableau(nRow - 1, 0), tableau(pos, 0));
+    int64_t nRowCoeff = lcm / tableau(nRow - 1, 0);
+    int64_t idxRowCoeff = coeffs[i] * (lcm / tableau(pos, 0));
+    tableau(nRow - 1, 0) = lcm;
+    for (unsigned col = 1; col < nCol; ++col)
+      tableau(nRow - 1, col) =
+          nRowCoeff * tableau(nRow - 1, col) + idxRowCoeff * tableau(pos, col);
+  }
+
+  normalizeRow(nRow - 1);
+  // Push to undo log along with the index of the new constraint.
+  undoLog.push_back(UndoLogEntry::RemoveLastConstraint);
+  return con.size() - 1;
+}
+
+/// Normalize the row by removing factors that are common between the
+/// denominator and all the numerator coefficients.
+void Simplex::normalizeRow(unsigned row) {
+  int64_t gcd = 0;
+  for (unsigned col = 0; col < nCol; ++col) {
+    if (gcd == 1)
+      break;
+    gcd = llvm::greatestCommonDivisor(gcd, std::abs(tableau(row, col)));
+  }
+  for (unsigned col = 0; col < nCol; ++col)
+    tableau(row, col) /= gcd;
+}
+
+namespace {
+bool signMatchesDirection(int64_t elem, Direction direction) {
+  assert(elem != 0 && "elem should not be 0");
+  return direction == Direction::Up ? elem > 0 : elem < 0;
+}
+
+Direction flippedDirection(Direction direction) {
+  return direction == Direction::Up ? Direction::Down : Simplex::Direction::Up;
+}
+} // anonymous namespace
+
+/// Find a pivot to change the sample value of the row in the specified
+/// direction. The returned pivot row will involve `row` if and only if the
+/// unknown is unbounded in the specified direction.
+///
+/// To increase (resp. decrease) the value of a row, we need to find a live
+/// column with a non-zero coefficient. If the coefficient is positive, we need
+/// to increase (decrease) the value of the column, and if the coefficient is
+/// negative, we need to decrease (increase) the value of the column. Also,
+/// we cannot decrease the sample value of restricted columns.
+///
+/// If multiple columns are valid, we break ties by considering a lexicographic
+/// ordering where we prefer unknowns with lower index.
+Optional<Simplex::Pivot> Simplex::findPivot(int row,
+                                            Direction direction) const {
+  Optional<unsigned> col;
+  for (unsigned j = 2; j < nCol; ++j) {
+    int64_t elem = tableau(row, j);
+    if (elem == 0)
+      continue;
+
+    if (unknownFromColumn(j).restricted &&
+        !signMatchesDirection(elem, direction))
+      continue;
+    if (!col || colUnknown[j] < colUnknown[*col])
+      col = j;
+  }
+
+  if (!col)
+    return {};
+
+  Direction newDirection =
+      tableau(row, *col) < 0 ? flippedDirection(direction) : direction;
+  Optional<unsigned> maybePivotRow = findPivotRow(row, newDirection, *col);
+  return Pivot{maybePivotRow.getValueOr(row), *col};
+}
+
+/// Swap the associated unknowns for the row and the column.
+///
+/// First we swap the index associated with the row and column. Then we update
+/// the unknowns to reflect their new position and orientation.
+void Simplex::swapRowWithCol(unsigned row, unsigned col) {
+  std::swap(rowUnknown[row], colUnknown[col]);
+  Unknown &uCol = unknownFromColumn(col);
+  Unknown &uRow = unknownFromRow(row);
+  uCol.orientation = Orientation::Column;
+  uRow.orientation = Orientation::Row;
+  uCol.pos = col;
+  uRow.pos = row;
+}
+
+void Simplex::pivot(Pivot pair) { pivot(pair.row, pair.column); }
+
+/// Pivot pivotRow and pivotCol.
+///
+/// Let R be the pivot row unknown and let C be the pivot col unknown.
+/// Since initially R = a*C + sum b_i * X_i
+/// (where the sum is over the other column's unknowns, x_i)
+/// C = (R - (sum b_i * X_i))/a
+///
+/// Let u be some other row unknown.
+/// u = c*C + sum d_i * X_i
+/// So u = c*(R - sum b_i * X_i)/a + sum d_i * X_i
+///
+/// This results in the following transform:
+///            pivot col    other col                   pivot col    other col
+/// pivot row     a             b       ->   pivot row     1/a         -b/a
+/// other row     c             d            other row     c/a        d - bc/a
+///
+/// Taking into account the common denominators p and q:
+///
+///            pivot col    other col                    pivot col   other col
+/// pivot row     a/p          b/p     ->   pivot row      p/a         -b/a
+/// other row     c/q          d/q          other row     cp/aq    (da - bc)/aq
+///
+/// The pivot row transform is accomplished be swapping a with the pivot row's
+/// common denominator and negating the pivot row except for the pivot column
+/// element.
+void Simplex::pivot(unsigned pivotRow, unsigned pivotCol) {
+  assert(pivotCol >= 2 && "Refusing to pivot invalid column");
+
+  swapRowWithCol(pivotRow, pivotCol);
+  std::swap(tableau(pivotRow, 0), tableau(pivotRow, pivotCol));
+  // We need to negate the whole pivot row except for the pivot column.
+  if (tableau(pivotRow, 0) < 0) {
+    // If the denominator is negative, we negate the row by simply negating the
+    // denominator.
+    tableau(pivotRow, 0) = -tableau(pivotRow, 0);
+    tableau(pivotRow, pivotCol) = -tableau(pivotRow, pivotCol);
+  } else {
+    for (unsigned col = 1; col < nCol; ++col) {
+      if (col == pivotCol)
+        continue;
+      tableau(pivotRow, col) = -tableau(pivotRow, col);
+    }
+  }
+  normalizeRow(pivotRow);
+
+  for (unsigned row = 0; row < nRow; ++row) {
+    if (row == pivotRow)
+      continue;
+    if (tableau(row, pivotCol) == 0) // Nothing to do.
+      continue;
+    tableau(row, 0) *= tableau(pivotRow, 0);
+    for (unsigned j = 1; j < nCol; ++j) {
+      if (j == pivotCol)
+        continue;
+      // Add rather than subtract because the pivot row has been negated.
+      tableau(row, j) = tableau(row, j) * tableau(pivotRow, 0) +
+                        tableau(row, pivotCol) * tableau(pivotRow, j);
+    }
+    tableau(row, pivotCol) *= tableau(pivotRow, pivotCol);
+    normalizeRow(row);
+  }
+}
+
+/// Perform pivots until the unknown has a non-negative sample value or until
+/// no more upward pivots can be performed. Return the sign of the final sample
+/// value.
+LogicalResult Simplex::restoreRow(Unknown &u) {
+  assert(u.orientation == Orientation::Row &&
+         "unknown should be in row position");
+
+  while (tableau(u.pos, 1) < 0) {
+    Optional<Pivot> maybePivot = findPivot(u.pos, Direction::Up);
+    if (!maybePivot)
+      break;
+
+    pivot(*maybePivot);
+    if (u.orientation == Orientation::Column)
+      return LogicalResult::Success; // the unknown is unbounded above.
+  }
+  return success(tableau(u.pos, 1) >= 0);
+}
+
+/// Find a row that can be used to pivot the column in the specified direction.
+/// This returns an empty optional if and only if the column is unbounded in the
+/// specified direction (ignoring skipRow, if skipRow is set).
+///
+/// If skipRow is set, this row is not considered, and (if it is restricted) its
+/// restriction may be violated by the returned pivot. Usually, skipRow is set
+/// because we don't want to move it to column position unless it is unbounded,
+/// and we are either trying to increase the value of skipRow or explicitly
+/// trying to make skipRow negative, so we are not concerned about this.
+///
+/// If the direction is up (resp. down) and a restricted row has a negative
+/// (positive) coefficient for the column, then this row imposes a bound on how
+/// much the sample value of the column can change. Such a row with constant
+/// term c and coefficient f for the column imposes a bound of c/|f| on the
+/// change in sample value (in the specified direction). (note that c is
+/// non-negative here since the row is restricted and the tableau is consistent)
+///
+/// We iterate through the rows and pick the row which imposes the most
+/// stringent bound, since pivoting with a row changes the row's sample value to
+/// 0 and hence saturates the bound it imposes. We break ties between rows that
+/// impose the same bound by considering a lexicographic ordering where we
+/// prefer unknowns with lower index value.
+Optional<unsigned> Simplex::findPivotRow(Optional<unsigned> skipRow,
+                                         Direction direction,
+                                         unsigned col) const {
+  Optional<unsigned> retRow;
+  int64_t retElem, retConst;
+  for (unsigned row = 0; row < nRow; ++row) {
+    if (skipRow && row == *skipRow)
+      continue;
+    int64_t elem = tableau(row, col);
+    if (elem == 0)
+      continue;
+    if (!unknownFromRow(row).restricted)
+      continue;
+    if (signMatchesDirection(elem, direction))
+      continue;
+    int64_t constTerm = tableau(row, 1);
+
+    if (!retRow) {
+      retRow = row;
+      retElem = elem;
+      retConst = constTerm;
+      continue;
+    }
+
+    int64_t diff = retConst * elem - constTerm * retElem;
+    if ((diff == 0 && rowUnknown[row] < rowUnknown[*retRow]) ||
+        (diff != 0 && !signMatchesDirection(diff, direction))) {
+      retRow = row;
+      retElem = elem;
+      retConst = constTerm;
+    }
+  }
+  return retRow;
+}
+
+bool Simplex::isEmpty() const { return empty; }
+
+void Simplex::swapRows(unsigned i, unsigned j) {
+  if (i == j)
+    return;
+  tableau.swapRows(i, j);
+  std::swap(rowUnknown[i], rowUnknown[j]);
+  unknownFromRow(i).pos = i;
+  unknownFromRow(j).pos = j;
+}
+
+/// Mark this tableau empty and push an entry to the undo stack.
+void Simplex::markEmpty() {
+  undoLog.push_back(UndoLogEntry::UnmarkEmpty);
+  empty = true;
+}
+
+/// Add an inequality to the tableau. If coeffs is c_0, c_1, ... c_n, where n
+/// is the curent number of variables, then the corresponding inequality is
+/// c_n + c_0*x_0 + c_1*x_1 + ... + c_{n-1}*x_{n-1} >= 0.
+///
+/// We add the inequality and mark it as restricted. We then try to make its
+/// sample value non-negative. If this is not possible, the tableau has become
+/// empty and we mark it as such.
+void Simplex::addInequality(ArrayRef<int64_t> coeffs) {
+  unsigned conIndex = addRow(coeffs);
+  Unknown &u = con[conIndex];
+  u.restricted = true;
+  LogicalResult result = restoreRow(u);
+  if (failed(result))
+    markEmpty();
+}
+
+/// Add an equality to the tableau. If coeffs is c_0, c_1, ... c_n, where n
+/// is the curent number of variables, then the corresponding equality is
+/// c_n + c_0*x_0 + c_1*x_1 + ... + c_{n-1}*x_{n-1} == 0.
+///
+/// We simply add two opposing inequalities, which force the expression to
+/// be zero.
+void Simplex::addEquality(ArrayRef<int64_t> coeffs) {
+  addInequality(coeffs);
+  SmallVector<int64_t, 8> negatedCoeffs;
+  for (int64_t coeff : coeffs)
+    negatedCoeffs.emplace_back(-coeff);
+  addInequality(negatedCoeffs);
+}
+
+unsigned Simplex::numVariables() const { return var.size(); }
+unsigned Simplex::numConstraints() const { return con.size(); }
+
+/// Return a snapshot of the curent state. This is just the current size of the
+/// undo log.
+unsigned Simplex::getSnapshot() const { return undoLog.size(); }
+
+void Simplex::undo(UndoLogEntry entry) {
+  if (entry == UndoLogEntry::RemoveLastConstraint) {
+    Unknown &constraint = con.back();
+    if (constraint.orientation == Orientation::Column) {
+      unsigned column = constraint.pos;
+      Optional<unsigned> row;
+
+      // Try to find any pivot row for this column that preserves tableau
+      // consistency (except possibly the column itself, which is going to be
+      // deallocated anyway).
+      //
+      // If no pivot row is found in either direction, then the unknown is
+      // unbounded in both directions and we are free to
+      // perform any pivot at all. To do this, we just need to find any row with
+      // a non-zero coefficient for the column.
+      if (Optional<unsigned> maybeRow =
+              findPivotRow({}, Direction::Up, column)) {
+        row = *maybeRow;
+      } else if (Optional<unsigned> maybeRow =
+                     findPivotRow({}, Direction::Down, column)) {
+        row = *maybeRow;
+      } else {
+        // The loop doesn't find a pivot row only if the column has zero
+        // coefficients for every row. But the unknown is a constraint,
+        // so it was added initially as a row. Such a row could never have been
+        // pivoted to a column. So a pivot row will always be found.
+        for (unsigned i = 0; i < nRow; ++i) {
+          if (tableau(i, column) != 0) {
+            row = i;
+            break;
+          }
+        }
+      }
+      assert(row.hasValue() && "No pivot row found!");
+      pivot(*row, column);
+    }
+
+    // Move this unknown to the last row and remove the last row from the
+    // tableau.
+    swapRows(constraint.pos, nRow - 1);
+    // It is not strictly necessary to shrink the tableau, but for now we
+    // maintain the invariant that the tableau has exactly nRow rows.
+    tableau.resizeVertically(nRow - 1);
+    nRow--;
+    rowUnknown.pop_back();
+    con.pop_back();
+  } else if (entry == UndoLogEntry::UnmarkEmpty) {
+    empty = false;
+  }
+}
+
+/// Rollback to the specified snapshot.
+///
+/// We undo all the log entries until the log size when the snapshot was taken
+/// is reached.
+void Simplex::rollback(unsigned snapshot) {
+  while (undoLog.size() > snapshot) {
+    undo(undoLog.back());
+    undoLog.pop_back();
+  }
+}
+
+Optional<Fraction> Simplex::computeRowOptimum(Direction direction,
+                                              unsigned row) {
+  // Keep trying to find a pivot for the row in the specified direction.
+  while (Optional<Pivot> maybePivot = findPivot(row, direction)) {
+    // If findPivot returns a pivot involving the row itself, then the optimum
+    // is unbounded, so we return None.
+    if (maybePivot->row == row)
+      return {};
+    pivot(*maybePivot);
+  }
+
+  // The row has reached its optimal sample value, which we return.
+  // The sample value is the entry in the constant column divided by the common
+  // denominator for this row.
+  return Fraction(tableau(row, 1), tableau(row, 0));
+}
+
+/// Compute the optimum of the specified expression in the specified direction,
+/// or None if it is unbounded.
+Optional<Fraction> Simplex::computeOptimum(Direction direction,
+                                           ArrayRef<int64_t> coeffs) {
+  assert(!empty && "Tableau should not be empty");
+
+  unsigned snapshot = getSnapshot();
+  unsigned conIndex = addRow(coeffs);
+  unsigned row = con[conIndex].pos;
+  Optional<Fraction> optimum = computeRowOptimum(direction, row);
+  rollback(snapshot);
+  return optimum;
+}
+
+bool Simplex::isUnbounded() {
+  if (empty)
+    return false;
+
+  SmallVector<int64_t, 8> dir(var.size() + 1);
+  for (unsigned i = 0; i < var.size(); ++i) {
+    dir[i] = 1;
+
+    Optional<Fraction> maybeMax = computeOptimum(Direction::Up, dir);
+    if (!maybeMax)
+      return true;
+
+    Optional<Fraction> maybeMin = computeOptimum(Direction::Down, dir);
+    if (!maybeMin)
+      return true;
+
+    dir[i] = 0;
+  }
+  return false;
+}
+
+/// Make a tableau to represent a pair of points in the original tableau.
+///
+/// The product constraints and variables are stored as: first A's, then B's.
+///
+/// The product tableau has row layout:
+///   A's rows, B's rows.
+///
+/// It has column layout:
+///   denominator, constant, A's columns, B's columns.
+Simplex Simplex::makeProduct(const Simplex &a, const Simplex &b) {
+  unsigned numVar = a.numVariables() + b.numVariables();
+  unsigned numCon = a.numConstraints() + b.numConstraints();
+  Simplex result(numVar);
+
+  result.tableau.resizeVertically(numCon);
+  result.empty = a.empty || b.empty;
+
+  auto concat = [](ArrayRef<Unknown> v, ArrayRef<Unknown> w) {
+    SmallVector<Unknown, 8> result;
+    result.reserve(v.size() + w.size());
+    result.insert(result.end(), v.begin(), v.end());
+    result.insert(result.end(), w.begin(), w.end());
+    return result;
+  };
+  result.con = concat(a.con, b.con);
+  result.var = concat(a.var, b.var);
+
+  auto indexFromBIndex = [&](int index) {
+    return index >= 0 ? a.numVariables() + index
+                      : ~(a.numConstraints() + ~index);
+  };
+
+  result.colUnknown.assign(2, nullIndex);
+  for (unsigned i = 2; i < a.nCol; ++i) {
+    result.colUnknown.push_back(a.colUnknown[i]);
+    result.unknownFromIndex(result.colUnknown.back()).pos =
+        result.colUnknown.size() - 1;
+  }
+  for (unsigned i = 2; i < b.nCol; ++i) {
+    result.colUnknown.push_back(indexFromBIndex(b.colUnknown[i]));
+    result.unknownFromIndex(result.colUnknown.back()).pos =
+        result.colUnknown.size() - 1;
+  }
+
+  auto appendRowFromA = [&](unsigned row) {
+    for (unsigned col = 0; col < a.nCol; ++col)
+      result.tableau(result.nRow, col) = a.tableau(row, col);
+    result.rowUnknown.push_back(a.rowUnknown[row]);
+    result.unknownFromIndex(result.rowUnknown.back()).pos =
+        result.rowUnknown.size() - 1;
+    result.nRow++;
+  };
+
+  // Also fixes the corresponding entry in rowUnknown and var/con (as the case
+  // may be).
+  auto appendRowFromB = [&](unsigned row) {
+    result.tableau(result.nRow, 0) = b.tableau(row, 0);
+    result.tableau(result.nRow, 1) = b.tableau(row, 1);
+
+    unsigned offset = a.nCol - 2;
+    for (unsigned col = 2; col < b.nCol; ++col)
+      result.tableau(result.nRow, offset + col) = b.tableau(row, col);
+    result.rowUnknown.push_back(indexFromBIndex(b.rowUnknown[row]));
+    result.unknownFromIndex(result.rowUnknown.back()).pos =
+        result.rowUnknown.size() - 1;
+    result.nRow++;
+  };
+
+  for (unsigned row = 0; row < a.nRow; ++row)
+    appendRowFromA(row);
+  for (unsigned row = 0; 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<int64_t, 8> sample;
+  // 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);
+    } 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));
+    }
+  }
+  return sample;
+}
+
+/// 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
+/// that contains a snapshot of the Simplex's state for each equality, just
+/// before that equality was added.
+class GBRSimplex {
+  using Orientation = Simplex::Orientation;
+
+public:
+  GBRSimplex(const Simplex &originalSimplex)
+      : simplex(Simplex::makeProduct(originalSimplex, originalSimplex)),
+        simplexConstraintOffset(simplex.numConstraints()) {}
+
+  /// Add an equality dotProduct(dir, x - y) == 0.
+  /// First pushes a snapshot for the current simplex state to the stack so
+  /// that this can be rolled back later.
+  void addEqualityForDirection(ArrayRef<int64_t> dir) {
+    assert(
+        std::any_of(dir.begin(), dir.end(), [](int64_t x) { return x != 0; }) &&
+        "Direction passed is the zero vector!");
+    snapshotStack.push_back(simplex.getSnapshot());
+    simplex.addEquality(getCoeffsForDirection(dir));
+  }
+
+  /// Compute max(dotProduct(dir, x - y)) and save the dual variables for only
+  /// the direction equalities to `dual`.
+  Fraction computeWidthAndDuals(ArrayRef<int64_t> dir,
+                                SmallVectorImpl<int64_t> &dual,
+                                int64_t &dualDenom) {
+    unsigned snap = simplex.getSnapshot();
+    unsigned conIndex = simplex.addRow(getCoeffsForDirection(dir));
+    unsigned row = simplex.con[conIndex].pos;
+    Optional<Fraction> maybeWidth =
+        simplex.computeRowOptimum(Simplex::Direction::Up, row);
+    assert(maybeWidth.hasValue() && "Width should not be unbounded!");
+    dualDenom = simplex.tableau(row, 0);
+    dual.clear();
+    // The increment is i += 2 because equalities are added as two inequalities,
+    // one positive and one negative. Each iteration processes one equality.
+    for (unsigned i = simplexConstraintOffset; i < conIndex; i += 2) {
+      // The dual variable is the negative of the coefficient of the new row
+      // in the column of the constraint, if the constraint is in a column.
+      // Note that the second inequality for the equality is negated.
+      //
+      // We want the dual for the original equality. If the positive inequality
+      // is in column position, the negative of its row coefficient is the
+      // desired dual. If the negative inequality is in column position, its row
+      // coefficient is the desired dual. (its coefficients are already the
+      // negated coefficients of the original equality, so we don't need to
+      // negate it now.)
+      //
+      // If neither are in column position, we move the negated inequality to
+      // column position. Since the inequality must have sample value zero
+      // (since it corresponds to an equality), we are free to pivot with
+      // any column. Since both the unknowns have sample value before and after
+      // pivoting, no other sample values will change and the tableau will
+      // remain consistent. To pivot, we just need to find a column that has a
+      // non-zero coefficient in this row. There must be one since otherwise the
+      // equality would be 0 == 0, which should never be passed to
+      // addEqualityForDirection.
+      //
+      // After finding a column, we pivot with the column, after which we can
+      // get the dual from the inequality in column position as explained above.
+      if (simplex.con[i].orientation == Orientation::Column) {
+        dual.push_back(-simplex.tableau(row, simplex.con[i].pos));
+      } else {
+        if (simplex.con[i + 1].orientation == Orientation::Row) {
+          unsigned ineqRow = simplex.con[i + 1].pos;
+          // Since it is an equality, the the sample value must be zero.
+          assert(simplex.tableau(ineqRow, 1) == 0 &&
+                 "Equality's sample value must be zero.");
+          for (unsigned col = 2; col < simplex.nCol; ++col) {
+            if (simplex.tableau(ineqRow, col) != 0) {
+              simplex.pivot(ineqRow, col);
+              break;
+            }
+          }
+          assert(simplex.con[i + 1].orientation == Orientation::Column &&
+                 "No pivot found. Equality has all-zeros row in tableau!");
+        }
+        dual.push_back(simplex.tableau(row, simplex.con[i + 1].pos));
+      }
+    }
+    simplex.rollback(snap);
+    return *maybeWidth;
+  }
+
+  /// Remove the last equality that was added through addEqualityForDirection.
+  ///
+  /// We do this by rolling back to the snapshot at the top of the stack, which
+  /// should be a snapshot taken just before the last equality was added.
+  void removeLastEquality() {
+    assert(!snapshotStack.empty() && "Snapshot stack is empty!");
+    simplex.rollback(snapshotStack.back());
+    snapshotStack.pop_back();
+  }
+
+private:
+  /// Returns coefficients of the expression 'dot_product(dir, x - y)',
+  /// i.e.,   dir_1 * x_1 + dir_2 * x_2 + ... + dir_n * x_n
+  ///       - dir_1 * y_1 - dir_2 * y_2 - ... - dir_n * y_n,
+  /// where n is the dimension of the original polytope.
+  SmallVector<int64_t, 8> getCoeffsForDirection(ArrayRef<int64_t> dir) {
+    assert(2 * dir.size() == simplex.numVariables() &&
+           "Direction vector has wrong dimensionality");
+    SmallVector<int64_t, 8> coeffs(dir.begin(), dir.end());
+    coeffs.reserve(2 * dir.size());
+    for (int64_t coeff : dir)
+      coeffs.push_back(-coeff);
+    coeffs.push_back(0); // constant term
+    return coeffs;
+  }
+
+  Simplex simplex;
+  /// The first index of the equality constraints, the index immediately after
+  /// the last constraint in the initial product simplex.
+  unsigned simplexConstraintOffset;
+  /// A stack of snapshots, used for rolling back.
+  SmallVector<unsigned, 8> snapshotStack;
+};
+
+/// Reduce the basis to try and find a direction in which the polytope is
+/// "thin". This only works for bounded polytopes.
+///
+/// This is an implementation of the algorithm described in the paper
+/// "An Implementation of Generalized Basis Reduction for Integer Programming"
+/// by W. Cook, T. Rutherford, H. E. Scarf, D. Shallcross.
+///
+/// Let b_{level}, b_{level + 1}, ... b_n be the current basis.
+/// Let width_i(v) = max <v, x - y> where x and y are points in the original
+/// polytope such that <b_j, x - y> = 0 is satisfied for all level <= j < i.
+///
+/// In every iteration, we first replace b_{i+1} with b_{i+1} + u*b_i, where u
+/// is the integer such that width_i(b_{i+1} + u*b_i) is minimized. Let dual_i
+/// be the dual variable associated with the constraint <b_i, x - y> = 0 when
+/// computing width_{i+1}(b_{i+1}). It can be shown that dual_i is the
+/// minimizing value of u, if it were allowed to be fractional. Due to
+/// convexity, the minimizing integer value is either floor(dual_i) or
+/// ceil(dual_i), so we just need to check which of these gives a lower
+/// width_{i+1} value. If dual_i turned out to be an integer, then u = dual_i.
+///
+/// Now if width_i(b_{i+1}) < 0.75 * width_i(b_i), we swap b_i and (the new)
+/// b_{i + 1} and decrement i (unless i = level, in which case we stay at the
+/// same i). Otherwise, we increment i.
+///
+/// We keep f values and duals cached and invalidate them when necessary.
+/// Whenever possible, we use them instead of recomputing them. We implement the
+/// algorithm as follows.
+///
+/// In an iteration at i we need to compute:
+///   a) width_i(b_{i + 1})
+///   b) width_i(b_i)
+///   c) the integer u that minimizes width_i(b_{i + 1} + u*b_i)
+///
+/// If width_i(b_i) is not already cached, we compute it.
+///
+/// If the duals are not already cached, we compute width_{i+1}(b_{i+1}) and
+/// store the duals from this computation.
+///
+/// We call updateBasisWithUAndGetFCandidate, which finds the minimizing value
+/// of u as explained before, caches the duals from this computation, sets
+/// b_{i+1} to b_{i+1} + u*b_i, and returns the new value of width_i(b_{i+1}).
+///
+/// Now if width_i(b_{i+1}) < 0.75 * width_i(b_i), we swap b_i and b_{i+1} and
+/// decrement i, resulting in the basis
+/// ... b_{i - 1}, b_{i + 1} + u*b_i, b_i, b_{i+2}, ...
+/// with corresponding f values
+/// ... width_{i-1}(b_{i-1}), width_i(b_{i+1} + u*b_i), width_{i+1}(b_i), ...
+/// The values up to i - 1 remain unchanged. We have just gotten the middle
+/// value from updateBasisWithUAndGetFCandidate, so we can update that in the
+/// cache. The value at width_{i+1}(b_i) is unknown, so we evict this value from
+/// the cache. The iteration after decrementing needs exactly the duals from the
+/// computation of width_i(b_{i + 1} + u*b_i), so we keep these in the cache.
+///
+/// When incrementing i, no cached f values get invalidated. However, the cached
+/// duals do get invalidated as the duals for the higher levels are different.
+void Simplex::reduceBasis(Matrix &basis, unsigned level) {
+  const Fraction epsilon(3, 4);
+
+  if (level == basis.getNumRows() - 1)
+    return;
+
+  GBRSimplex gbrSimplex(*this);
+  SmallVector<Fraction, 8> width;
+  SmallVector<int64_t, 8> dual;
+  int64_t dualDenom;
+
+  // Finds the value of u that minimizes width_i(b_{i+1} + u*b_i), caches the
+  // duals from this computation, sets b_{i+1} to b_{i+1} + u*b_i, and returns
+  // the new value of width_i(b_{i+1}).
+  //
+  // If dual_i is not an integer, the minimizing value must be either
+  // floor(dual_i) or ceil(dual_i). We compute the expression for both and
+  // choose the minimizing value.
+  //
+  // If dual_i is an integer, we don't need to perform these computations. We
+  // know that in this case,
+  //   a) u = dual_i.
+  //   b) one can show that dual_j for j < i are the same duals we would have
+  //      gotten from computing width_i(b_{i + 1} + u*b_i), so the correct duals
+  //      are the ones already in the cache.
+  //   c) width_i(b_{i+1} + u*b_i) = min_{alpha} width_i(b_{i+1} + alpha * b_i),
+  //   which
+  //      one can show is equal to width_{i+1}(b_{i+1}). The latter value must
+  //      be in the cache, so we get it from there and return it.
+  auto updateBasisWithUAndGetFCandidate = [&](unsigned i) -> Fraction {
+    assert(i < level + dual.size() && "dual_i is not known!");
+
+    int64_t u = floorDiv(dual[i - level], dualDenom);
+    basis.addToRow(i, i + 1, u);
+    if (dual[i - level] % dualDenom != 0) {
+      SmallVector<int64_t, 8> candidateDual[2];
+      int64_t candidateDualDenom[2];
+      Fraction widthI[2];
+
+      // Initially u is floor(dual) and basis reflects this.
+      widthI[0] = gbrSimplex.computeWidthAndDuals(
+          basis.getRow(i + 1), candidateDual[0], candidateDualDenom[0]);
+
+      // Now try ceil(dual), i.e. floor(dual) + 1.
+      ++u;
+      basis.addToRow(i, i + 1, 1);
+      widthI[1] = gbrSimplex.computeWidthAndDuals(
+          basis.getRow(i + 1), candidateDual[1], candidateDualDenom[1]);
+
+      unsigned j = widthI[0] < widthI[1] ? 0 : 1;
+      if (j == 0)
+        // Subtract 1 to go from u = ceil(dual) back to floor(dual).
+        basis.addToRow(i, i + 1, -1);
+      dual = std::move(candidateDual[j]);
+      dualDenom = candidateDualDenom[j];
+      return widthI[j];
+    }
+    assert(i + 1 - level < width.size() && "width_{i+1} wasn't saved");
+    // When dual minimizes f_i(b_{i+1} + dual*b_i), this is equal to
+    // width_{i+1}(b_{i+1}).
+    return width[i + 1 - level];
+  };
+
+  // In the ith iteration of the loop, gbrSimplex has constraints for directions
+  // from `level` to i - 1.
+  unsigned i = level;
+  while (i < basis.getNumRows() - 1) {
+    if (i >= level + width.size()) {
+      // We don't even know the value of f_i(b_i), so let's find that first.
+      // We have to do this first since later we assume that width already
+      // contains values up to and including i.
+
+      assert((i == 0 || i - 1 < level + width.size()) &&
+             "We are at level i but we don't know the value of width_{i-1}");
+
+      // We don't actually use these duals at all, but it doesn't matter
+      // because this case should only occur when i is level, and there are no
+      // duals in that case anyway.
+      assert(i == level && "This case should only occur when i == level");
+      width.push_back(
+          gbrSimplex.computeWidthAndDuals(basis.getRow(i), dual, dualDenom));
+    }
+
+    if (i >= level + dual.size()) {
+      assert(i + 1 >= level + width.size() &&
+             "We don't know dual_i but we know width_{i+1}");
+      // We don't know dual for our level, so let's find it.
+      gbrSimplex.addEqualityForDirection(basis.getRow(i));
+      width.push_back(gbrSimplex.computeWidthAndDuals(basis.getRow(i + 1), dual,
+                                                      dualDenom));
+      gbrSimplex.removeLastEquality();
+    }
+
+    // This variable stores width_i(b_{i+1} + u*b_i).
+    Fraction widthICandidate = updateBasisWithUAndGetFCandidate(i);
+    if (widthICandidate < epsilon * width[i - level]) {
+      basis.swapRows(i, i + 1);
+      width[i - level] = widthICandidate;
+      // The values of width_{i+1}(b_{i+1}) and higher may change after the
+      // swap, so we remove the cached values here.
+      width.resize(i - level + 1);
+      if (i == level) {
+        dual.clear();
+        continue;
+      }
+
+      gbrSimplex.removeLastEquality();
+      i--;
+      continue;
+    }
+
+    // Invalidate duals since the higher level needs to recompute its own duals.
+    dual.clear();
+    gbrSimplex.addEqualityForDirection(basis.getRow(i));
+    i++;
+  }
+}
+
+/// Search for an integer sample point using a branch and bound algorithm.
+///
+/// Each row in the basis matrix is a vector, and the set of basis vectors
+/// should span the space. Initially this is the identity matrix,
+/// i.e., the basis vectors are just the variables.
+///
+/// In every level, a value is assigned to the level-th basis vector, as
+/// follows. Compute the minimum and maximum rational values of this direction.
+/// If only one integer point lies in this range, constrain the variable to
+/// have this value and recurse to the next variable.
+///
+/// If the range has multiple values, perform generalized basis reduction via
+/// reduceBasis and then compute the bounds again. Now we try constraining
+/// this direction in the first value in this range and "recurse" to the next
+/// level. If we fail to find a sample, we try assigning the direction the next
+/// value in this range, and so on.
+///
+/// If no integer sample is found from any of the assignments, or if the range
+/// contains no integer value, then of course the polytope is empty for the
+/// current assignment of the values in previous levels, so we return to
+/// the previous level.
+///
+/// If we reach the last level where all the variables have been assigned values
+/// already, then we simply return the current sample point if it is integral,
+/// and go back to the previous level otherwise.
+///
+/// To avoid potentially arbitrarily large recursion depths leading to stack
+/// overflows, this algorithm is implemented iteratively.
+Optional<SmallVector<int64_t, 8>> Simplex::findIntegerSample() {
+  if (empty)
+    return {};
+
+  unsigned nDims = var.size();
+  Matrix basis = Matrix::identity(nDims);
+
+  unsigned level = 0;
+  // The snapshot just before constraining a direction to a value at each level.
+  SmallVector<unsigned, 8> snapshotStack;
+  // The maximum value in the range of the direction for each level.
+  SmallVector<int64_t, 8> upperBoundStack;
+  // The next value to try constraining the basis vector to at each level.
+  SmallVector<int64_t, 8> nextValueStack;
+
+  snapshotStack.reserve(basis.getNumRows());
+  upperBoundStack.reserve(basis.getNumRows());
+  nextValueStack.reserve(basis.getNumRows());
+  while (level != -1u) {
+    if (level == basis.getNumRows()) {
+      // We've assigned values to all variables. Return if we have a sample,
+      // or go back up to the previous level otherwise.
+      if (auto maybeSample = getSamplePointIfIntegral())
+        return maybeSample;
+      level--;
+      continue;
+    }
+
+    if (level >= upperBoundStack.size()) {
+      // We haven't populated the stack values for this level yet, so we have
+      // just come down a level ("recursed"). Find the lower and upper bounds.
+      // If there is more than one integer point in the range, perform
+      // generalized basis reduction.
+      SmallVector<int64_t, 8> basisCoeffs =
+          llvm::to_vector<8>(basis.getRow(level));
+      basisCoeffs.push_back(0);
+
+      int64_t minRoundedUp, maxRoundedDown;
+      std::tie(minRoundedUp, maxRoundedDown) =
+          computeIntegerBounds(basisCoeffs);
+
+      // Heuristic: if the sample point is integral at this point, just return
+      // it.
+      if (auto maybeSample = getSamplePointIfIntegral())
+        return *maybeSample;
+
+      if (minRoundedUp < maxRoundedDown) {
+        reduceBasis(basis, level);
+        basisCoeffs = llvm::to_vector<8>(basis.getRow(level));
+        basisCoeffs.push_back(0);
+        std::tie(minRoundedUp, maxRoundedDown) =
+            computeIntegerBounds(basisCoeffs);
+      }
+
+      snapshotStack.push_back(getSnapshot());
+      // The smallest value in the range is the next value to try.
+      nextValueStack.push_back(minRoundedUp);
+      upperBoundStack.push_back(maxRoundedDown);
+    }
+
+    assert((snapshotStack.size() - 1 == level &&
+            nextValueStack.size() - 1 == level &&
+            upperBoundStack.size() - 1 == level) &&
+           "Mismatched variable stack sizes!");
+
+    // Whether we "recursed" or "returned" from a lower level, we rollback
+    // to the snapshot of the starting state at this level. (in the "recursed"
+    // case this has no effect)
+    rollback(snapshotStack.back());
+    int64_t nextValue = nextValueStack.back();
+    nextValueStack.back()++;
+    if (nextValue > upperBoundStack.back()) {
+      // We have exhausted the range and found no solution. Pop the stack and
+      // return up a level.
+      snapshotStack.pop_back();
+      nextValueStack.pop_back();
+      upperBoundStack.pop_back();
+      level--;
+      continue;
+    }
+
+    // Try the next value in the range and "recurse" into the next level.
+    SmallVector<int64_t, 8> basisCoeffs(basis.getRow(level).begin(),
+                                        basis.getRow(level).end());
+    basisCoeffs.push_back(-nextValue);
+    addEquality(basisCoeffs);
+    level++;
+  }
+
+  return {};
+}
+
+/// Compute the minimum and maximum integer values the expression can take. We
+/// compute each separately.
+std::pair<int64_t, int64_t>
+Simplex::computeIntegerBounds(ArrayRef<int64_t> coeffs) {
+  int64_t minRoundedUp;
+  if (Optional<Fraction> maybeMin =
+          computeOptimum(Simplex::Direction::Down, coeffs))
+    minRoundedUp = ceil(*maybeMin);
+  else
+    llvm_unreachable("Tableau should not be unbounded");
+
+  int64_t maxRoundedDown;
+  if (Optional<Fraction> maybeMax =
+          computeOptimum(Simplex::Direction::Up, coeffs))
+    maxRoundedDown = floor(*maybeMax);
+  else
+    llvm_unreachable("Tableau should not be unbounded");
+
+  return {minRoundedUp, maxRoundedDown};
+}
+
+void Simplex::print(raw_ostream &os) const {
+  os << "rows = " << nRow << ", columns = " << nCol << "\n";
+  if (empty)
+    os << "Simplex marked empty!\n";
+  os << "var: ";
+  for (unsigned i = 0; i < var.size(); ++i) {
+    if (i > 0)
+      os << ", ";
+    var[i].print(os);
+  }
+  os << "\ncon: ";
+  for (unsigned i = 0; i < con.size(); ++i) {
+    if (i > 0)
+      os << ", ";
+    con[i].print(os);
+  }
+  os << '\n';
+  for (unsigned row = 0; row < nRow; ++row) {
+    if (row > 0)
+      os << ", ";
+    os << "r" << row << ": " << rowUnknown[row];
+  }
+  os << '\n';
+  os << "c0: denom, c1: const";
+  for (unsigned col = 2; col < nCol; ++col)
+    os << ", c" << col << ": " << colUnknown[col];
+  os << '\n';
+  for (unsigned row = 0; row < nRow; ++row) {
+    for (unsigned col = 0; col < nCol; ++col)
+      os << tableau(row, col) << '\t';
+    os << '\n';
+  }
+  os << '\n';
+}
+
+void Simplex::dump() const { print(llvm::errs()); }
+
+} // namespace mlir

mlir/unittests/Analysis/AffineStructuresTest.cpp

+//===- AffineStructuresTest.cpp - Tests for AffineStructures ----*- 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
+//
+//===----------------------------------------------------------------------===//
+
+#include "mlir/Analysis/AffineStructures.h"
+
+#include <gmock/gmock.h>
+#include <gtest/gtest.h>
+
+#include <numeric>
+
+namespace mlir {
+
+/// Evaluate the value of the given affine expression at the specified point.
+/// The expression is a list of coefficients for the dimensions followed by the
+/// constant term.
+int64_t valueAt(ArrayRef<int64_t> expr, ArrayRef<int64_t> point) {
+  assert(expr.size() == 1 + point.size());
+  int64_t value = expr.back();
+  for (unsigned i = 0; i < point.size(); ++i)
+    value += expr[i] * point[i];
+  return value;
+}
+
+/// If 'hasValue' is true, check that findIntegerSample returns a valid sample
+/// for the FlatAffineConstraints fac.
+///
+/// If hasValue is false, check that findIntegerSample does not return None.
+void checkSample(bool hasValue, const FlatAffineConstraints &fac) {
+  Optional<SmallVector<int64_t, 8>> maybeSample = fac.findIntegerSample();
+  if (!hasValue) {
+    EXPECT_FALSE(maybeSample.hasValue());
+    if (maybeSample.hasValue()) {
+      for (auto x : *maybeSample)
+        llvm::errs() << x << ' ';
+      llvm::errs() << '\n';
+    }
+  } else {
+    ASSERT_TRUE(maybeSample.hasValue());
+    for (unsigned i = 0; i < fac.getNumEqualities(); ++i)
+      EXPECT_EQ(valueAt(fac.getEquality(i), *maybeSample), 0);
+    for (unsigned i = 0; i < fac.getNumInequalities(); ++i)
+      EXPECT_GE(valueAt(fac.getInequality(i), *maybeSample), 0);
+  }
+}
+
+/// Construct a FlatAffineConstraints from a set of inequality and
+/// equality constraints.
+FlatAffineConstraints
+makeFACFromConstraints(unsigned dims, ArrayRef<SmallVector<int64_t, 4>> ineqs,
+                       ArrayRef<SmallVector<int64_t, 4>> eqs) {
+  FlatAffineConstraints fac(ineqs.size(), eqs.size(), dims + 1, dims);
+  for (const auto &eq : eqs)
+    fac.addEquality(eq);
+  for (const auto &ineq : ineqs)
+    fac.addInequality(ineq);
+  return fac;
+}
+
+/// Check sampling for all the permutations of the dimensions for the given
+/// constraint set. Since the GBR algorithm progresses dimension-wise, different
+/// orderings may cause the algorithm to proceed differently. At least some of
+///.these permutations should make it past the heuristics and test the
+/// implementation of the GBR algorithm itself.
+void checkPermutationsSample(bool hasValue, unsigned nDim,
+                             ArrayRef<SmallVector<int64_t, 4>> ineqs,
+                             ArrayRef<SmallVector<int64_t, 4>> eqs) {
+  SmallVector<unsigned, 4> perm(nDim);
+  std::iota(perm.begin(), perm.end(), 0);
+  auto permute = [&perm](ArrayRef<int64_t> coeffs) {
+    SmallVector<int64_t, 4> permuted;
+    for (unsigned id : perm)
+      permuted.push_back(coeffs[id]);
+    permuted.push_back(coeffs.back());
+    return permuted;
+  };
+  do {
+    SmallVector<SmallVector<int64_t, 4>, 4> permutedIneqs, permutedEqs;
+    for (const auto &ineq : ineqs)
+      permutedIneqs.push_back(permute(ineq));
+    for (const auto &eq : eqs)
+      permutedEqs.push_back(permute(eq));
+
+    checkSample(hasValue,
+                makeFACFromConstraints(nDim, permutedIneqs, permutedEqs));
+  } while (std::next_permutation(perm.begin(), perm.end()));
+}
+
+TEST(FlatAffineConstraintsTest, FindSampleTest) {
+  // Bounded sets with only inequalities.
+
+  // 0 <= 7x <= 5
+  checkSample(true, makeFACFromConstraints(1, {{7, 0}, {-7, 5}}, {}));
+
+  // 1 <= 5x and 5x <= 4 (no solution).
+  checkSample(false, makeFACFromConstraints(1, {{5, -1}, {-5, 4}}, {}));
+
+  // 1 <= 5x and 5x <= 9 (solution: x = 1).
+  checkSample(true, makeFACFromConstraints(1, {{5, -1}, {-5, 9}}, {}));
+
+  // Bounded sets with equalities.
+  // x >= 8 and 40 >= y and x = y.
+  checkSample(
+      true, makeFACFromConstraints(2, {{1, 0, -8}, {0, -1, 40}}, {{1, -1, 0}}));
+
+  // x <= 10 and y <= 10 and 10 <= z and x + 2y = 3z.
+  // solution: x = y = z = 10.
+  checkSample(true, makeFACFromConstraints(
+                        3, {{-1, 0, 0, 10}, {0, -1, 0, 10}, {0, 0, 1, -10}},
+                        {{1, 2, -3, 0}}));
+
+  // x <= 10 and y <= 10 and 11 <= z and x + 2y = 3z.
+  // This implies x + 2y >= 33 and x + 2y <= 30, which has no solution.
+  checkSample(false, makeFACFromConstraints(
+                         3, {{-1, 0, 0, 10}, {0, -1, 0, 10}, {0, 0, 1, -11}},
+                         {{1, 2, -3, 0}}));
+
+  // 0 <= r and r <= 3 and 4q + r = 7.
+  // Solution: q = 1, r = 3.
+  checkSample(true,
+              makeFACFromConstraints(2, {{0, 1, 0}, {0, -1, 3}}, {{4, 1, -7}}));
+
+  // 4q + r = 7 and r = 0.
+  // Solution: q = 1, r = 3.
+  checkSample(false, makeFACFromConstraints(2, {}, {{4, 1, -7}, {0, 1, 0}}));
+
+  // The next two sets are large sets that should take a long time to sample
+  // with a naive branch and bound algorithm but can be sampled efficiently with
+  // the GBR algroithm.
+  //
+  // This is a triangle with vertices at (1/3, 0), (2/3, 0) and (10000, 10000).
+  checkSample(
+      true,
+      makeFACFromConstraints(
+          2, {{0, 1, 0}, {300000, -299999, -100000}, {-300000, 299998, 200000}},
+          {}));
+
+  // This is a tetrahedron with 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.
+  checkPermutationsSample(
+      true, 3,
+      {
+          {0, 1, 0, 0},  // y >= 0
+          {0, -1, 1, 0}, // z >= y
+          {300000, -299998, -1,
+           -100000},                    // -300000x + 299998y + 100000 + z <= 0.
+          {-150000, 149999, 0, 100000}, // -150000x + 149999y + 100000 >= 0.
+      },
+      {});
+
+  // Same thing with some spurious extra dimensions equated to constants.
+  checkSample(true,
+              makeFACFromConstraints(
+                  5,
+                  {
+                      {0, 1, 0, 1, -1, 0},
+                      {0, -1, 1, -1, 1, 0},
+                      {300000, -299998, -1, -9, 21, -112000},
+                      {-150000, 149999, 0, -15, 47, 68000},
+                  },
+                  {{0, 0, 0, 1, -1, 0},       // p = q.
+                   {0, 0, 0, 1, 1, -2000}})); // p + q = 20000 => p = q = 10000.
+
+  // This is a tetrahedron with vertices at
+  // (1/3, 0, 0), (2/3, 0, 0), (2/3, 0, 100), (100, 100 - 1/3, 100).
+  checkPermutationsSample(false, 3,
+                          {
+                              {0, 1, 0, 0},
+                              {0, -300, 299, 0},
+                              {300 * 299, -89400, -299, -100 * 299},
+                              {-897, 894, 0, 598},
+                          },
+                          {});
+
+  // Two tests involving equalities that are integer empty but not rational
+  // empty.
+
+  // This is a line segment from (0, 1/3) to (100, 100 + 1/3).
+  checkSample(false, makeFACFromConstraints(
+                         2,
+                         {
+                             {1, 0, 0},   // x >= 0.
+                             {-1, 0, 100} // -x + 100 >= 0, i.e., x <= 100.
+                         },
+                         {
+                             {3, -3, 1} // 3x - 3y + 1 = 0, i.e., y = x + 1/3.
+                         }));
+
+  // A thin parallelogram. 0 <= x <= 100 and x + 1/3 <= y <= x + 2/3.
+  checkSample(false, makeFACFromConstraints(2,
+                                            {
+                                                {1, 0, 0},    // x >= 0.
+                                                {-1, 0, 100}, // x <= 100.
+                                                {3, -3, 2},   // 3x - 3y >= -2.
+                                                {-3, 3, -1},  // 3x - 3y <= -1.
+                                            },
+                                            {}));
+
+  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());
+
+  // An unbounded set, which isIntegerEmpty should detect as unbounded and
+  // return without calling findIntegerSample.
+  EXPECT_FALSE(makeFACFromConstraints(3,
+                                      {
+                                          {2, 0, 0, -1},
+                                          {-2, 0, 0, 1},
+                                          {0, 2, 0, -1},
+                                          {0, -2, 0, 1},
+                                          {0, 0, 2, -1},
+                                      },
+                                      {})
+                   .isIntegerEmpty());
+
+  // FlatAffineConstraints::isEmpty() does not detect the following sets to be
+  // empty.
+
+  // 3x + 7y = 1 and 0 <= x, y <= 10.
+  // Since x and y are non-negative, 3x + 7y can never be 1.
+  EXPECT_TRUE(
+      makeFACFromConstraints(
+          2, {{1, 0, 0}, {-1, 0, 10}, {0, 1, 0}, {0, -1, 10}}, {{3, 7, -1}})
+          .isIntegerEmpty());
+
+  // 2x = 3y and y = x - 1 and x + y = 6z + 2 and 0 <= x, y <= 100.
+  // Substituting y = x - 1 in 3y = 2x, we obtain x = 3 and hence y = 2.
+  // Since x + y = 5 cannot be equal to 6z + 2 for any z, the set is empty.
+  EXPECT_TRUE(
+      makeFACFromConstraints(3,
+                             {
+                                 {1, 0, 0, 0},
+                                 {-1, 0, 0, 100},
+                                 {0, 1, 0, 0},
+                                 {0, -1, 0, 100},
+                             },
+                             {{2, -3, 0, 0}, {1, -1, 0, -1}, {1, 1, -6, -2}})
+          .isIntegerEmpty());
+
+  // 2x = 3y and y = x - 1 + 6z and x + y = 6q + 2 and 0 <= x, y <= 100.
+  // 2x = 3y implies x is a multiple of 3 and y is even.
+  // Now y = x - 1 + 6z implies y = 2 mod 3. In fact, since y is even, we have
+  // y = 2 mod 6. Then since x = y + 1 + 6z, we have x = 3 mod 6, implying
+  // x + y = 5 mod 6, which contradicts x + y = 6q + 2, so the set is empty.
+  EXPECT_TRUE(makeFACFromConstraints(
+                  4,
+                  {
+                      {1, 0, 0, 0, 0},
+                      {-1, 0, 0, 0, 100},
+                      {0, 1, 0, 0, 0},
+                      {0, -1, 0, 0, 100},
+                  },
+                  {{2, -3, 0, 0, 0}, {1, -1, 6, 0, -1}, {1, 1, 0, -6, -2}})
+                  .isIntegerEmpty());
+}
+
+} // namespace mlir

mlir/unittests/Analysis/CMakeLists.txt

+add_mlir_unittest(MLIRAnalysisTests
+  AffineStructuresTest.cpp
+)
+
+target_link_libraries(MLIRAnalysisTests
+  PRIVATE MLIRLoopAnalysis)
+
+add_subdirectory(Presburger)

mlir/unittests/Analysis/Presburger/CMakeLists.txt

+add_mlir_unittest(MLIRPresburgerTests
+  MatrixTest.cpp
+  SimplexTest.cpp
+)
+
+target_link_libraries(MLIRPresburgerTests
+  PRIVATE MLIRPresburger)

mlir/unittests/Analysis/Presburger/MatrixTest.cpp

+//===- MatrixTest.cpp - Tests for Matrix ----------------------------------===//
+//
+// 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/Matrix.h"
+#include <gmock/gmock.h>
+#include <gtest/gtest.h>
+
+namespace mlir {
+
+TEST(MatrixTest, ReadWrite) {
+  Matrix mat(5, 5);
+  for (unsigned row = 0; row < 5; ++row)
+    for (unsigned col = 0; col < 5; ++col)
+      mat(row, col) = 10 * row + col;
+  for (unsigned row = 0; row < 5; ++row)
+    for (unsigned col = 0; col < 5; ++col)
+      EXPECT_EQ(mat(row, col), int(10 * row + col));
+}
+
+TEST(MatrixTest, SwapColumns) {
+  Matrix mat(5, 5);
+  for (unsigned row = 0; row < 5; ++row)
+    for (unsigned col = 0; col < 5; ++col)
+      mat(row, col) = col == 3 ? 1 : 0;
+  mat.swapColumns(3, 1);
+  for (unsigned row = 0; row < 5; ++row)
+    for (unsigned col = 0; col < 5; ++col)
+      EXPECT_EQ(mat(row, col), col == 1 ? 1 : 0);
+
+  // swap around all the other columns, swap (1, 3) twice for no effect.
+  mat.swapColumns(3, 1);
+  mat.swapColumns(2, 4);
+  mat.swapColumns(1, 3);
+  mat.swapColumns(0, 4);
+  mat.swapColumns(2, 2);
+
+  for (unsigned row = 0; row < 5; ++row)
+    for (unsigned col = 0; col < 5; ++col)
+      EXPECT_EQ(mat(row, col), col == 1 ? 1 : 0);
+}
+
+TEST(MatrixTest, SwapRows) {
+  Matrix mat(5, 5);
+  for (unsigned row = 0; row < 5; ++row)
+    for (unsigned col = 0; col < 5; ++col)
+      mat(row, col) = row == 2 ? 1 : 0;
+  mat.swapRows(2, 0);
+  for (unsigned row = 0; row < 5; ++row)
+    for (unsigned col = 0; col < 5; ++col)
+      EXPECT_EQ(mat(row, col), row == 0 ? 1 : 0);
+
+  // swap around all the other rows, swap (2, 0) twice for no effect.
+  mat.swapRows(3, 4);
+  mat.swapRows(1, 4);
+  mat.swapRows(2, 0);
+  mat.swapRows(1, 1);
+  mat.swapRows(0, 2);
+
+  for (unsigned row = 0; row < 5; ++row)
+    for (unsigned col = 0; col < 5; ++col)
+      EXPECT_EQ(mat(row, col), row == 0 ? 1 : 0);
+}
+
+TEST(MatrixTest, resizeVertically) {
+  Matrix mat(5, 5);
+  EXPECT_EQ(mat.getNumRows(), 5u);
+  EXPECT_EQ(mat.getNumColumns(), 5u);
+  for (unsigned row = 0; row < 5; ++row)
+    for (unsigned col = 0; col < 5; ++col)
+      mat(row, col) = 10 * row + col;
+
+  mat.resizeVertically(3);
+  EXPECT_EQ(mat.getNumRows(), 3u);
+  EXPECT_EQ(mat.getNumColumns(), 5u);
+  for (unsigned row = 0; row < 3; ++row)
+    for (unsigned col = 0; col < 5; ++col)
+      EXPECT_EQ(mat(row, col), int(10 * row + col));
+
+  mat.resizeVertically(5);
+  EXPECT_EQ(mat.getNumRows(), 5u);
+  EXPECT_EQ(mat.getNumColumns(), 5u);
+  for (unsigned row = 0; row < 5; ++row)
+    for (unsigned col = 0; col < 5; ++col)
+      EXPECT_EQ(mat(row, col), row >= 3 ? 0 : int(10 * row + col));
+}
+
+} // namespace mlir

mlir/unittests/Analysis/Presburger/SimplexTest.cpp

+//===- SimplexTest.cpp - Tests for Simplex --------------------------------===//
+//
+// 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/Simplex.h"
+
+#include <gmock/gmock.h>
+#include <gtest/gtest.h>
+
+namespace mlir {
+
+/// Take a snapshot, add constraints making the set empty, and rollback.
+/// The set should not be empty after rolling back.
+TEST(SimplexTest, emptyRollback) {
+  Simplex simplex(2);
+  // (u - v) >= 0
+  simplex.addInequality({1, -1, 0});
+  EXPECT_FALSE(simplex.isEmpty());
+
+  unsigned snapshot = simplex.getSnapshot();
+  // (u - v) <= -1
+  simplex.addInequality({-1, 1, -1});
+  EXPECT_TRUE(simplex.isEmpty());
+  simplex.rollback(snapshot);
+  EXPECT_FALSE(simplex.isEmpty());
+}
+
+/// Check that the set gets marked as empty when we add contradictory
+/// constraints.
+TEST(SimplexTest, addEquality_separate) {
+  Simplex simplex(1);
+  simplex.addInequality({1, -1}); // x >= 1.
+  ASSERT_FALSE(simplex.isEmpty());
+  simplex.addEquality({1, 0}); // x == 0.
+  EXPECT_TRUE(simplex.isEmpty());
+}
+
+void expectInequalityMakesSetEmpty(Simplex &simplex, ArrayRef<int64_t> coeffs,
+                                   bool expect) {
+  ASSERT_FALSE(simplex.isEmpty());
+  unsigned snapshot = simplex.getSnapshot();
+  simplex.addInequality(coeffs);
+  EXPECT_EQ(simplex.isEmpty(), expect);
+  simplex.rollback(snapshot);
+}
+
+TEST(SimplexTest, addInequality_rollback) {
+  Simplex simplex(3);
+  SmallVector<int64_t, 4> coeffs[]{{1, 0, 0, 0},   // u >= 0.
+                                   {-1, 0, 0, 0},  // u <= 0.
+                                   {1, -1, 1, 0},  // u - v + w >= 0.
+                                   {1, 1, -1, 0}}; // u + v - w >= 0.
+  // The above constraints force u = 0 and v = w.
+  // The constraints below violate v = w.
+  SmallVector<int64_t, 4> checkCoeffs[]{{0, 1, -1, -1},  // v - w >= 1.
+                                        {0, -1, 1, -1}}; // v - w <= -1.
+
+  for (int run = 0; run < 4; run++) {
+    unsigned snapshot = simplex.getSnapshot();
+
+    expectInequalityMakesSetEmpty(simplex, checkCoeffs[0], false);
+    expectInequalityMakesSetEmpty(simplex, checkCoeffs[1], false);
+
+    for (int i = 0; i < 4; i++)
+      simplex.addInequality(coeffs[(run + i) % 4]);
+
+    expectInequalityMakesSetEmpty(simplex, checkCoeffs[0], true);
+    expectInequalityMakesSetEmpty(simplex, checkCoeffs[1], true);
+
+    simplex.rollback(snapshot);
+    EXPECT_EQ(simplex.numConstraints(), 0u);
+
+    expectInequalityMakesSetEmpty(simplex, checkCoeffs[0], false);
+    expectInequalityMakesSetEmpty(simplex, checkCoeffs[1], false);
+  }
+}
+
+Simplex simplexFromConstraints(unsigned nDim,
+                               SmallVector<SmallVector<int64_t, 8>, 8> ineqs,
+                               SmallVector<SmallVector<int64_t, 8>, 8> eqs) {
+  Simplex simplex(nDim);
+  for (const auto &ineq : ineqs)
+    simplex.addInequality(ineq);
+  for (const auto &eq : eqs)
+    simplex.addEquality(eq);
+  return simplex;
+}
+
+TEST(SimplexTest, isUnbounded) {
+  EXPECT_FALSE(simplexFromConstraints(
+                   2, {{1, 1, 0}, {-1, -1, 0}, {1, -1, 5}, {-1, 1, -5}}, {})
+                   .isUnbounded());
+
+  EXPECT_TRUE(
+      simplexFromConstraints(2, {{1, 1, 0}, {1, -1, 5}, {-1, 1, -5}}, {})
+          .isUnbounded());
+
+  EXPECT_TRUE(
+      simplexFromConstraints(2, {{-1, -1, 0}, {1, -1, 5}, {-1, 1, -5}}, {})
+          .isUnbounded());
+
+  EXPECT_TRUE(simplexFromConstraints(2, {}, {}).isUnbounded());
+
+  EXPECT_FALSE(simplexFromConstraints(3,
+                                      {
+                                          {2, 0, 0, -1},
+                                          {-2, 0, 0, 1},
+                                          {0, 2, 0, -1},
+                                          {0, -2, 0, 1},
+                                          {0, 0, 2, -1},
+                                          {0, 0, -2, 1},
+                                      },
+                                      {})
+                   .isUnbounded());
+
+  EXPECT_TRUE(simplexFromConstraints(3,
+                                     {
+                                         {2, 0, 0, -1},
+                                         {-2, 0, 0, 1},
+                                         {0, 2, 0, -1},
+                                         {0, -2, 0, 1},
+                                         {0, 0, -2, 1},
+                                     },
+                                     {})
+                  .isUnbounded());
+
+  EXPECT_TRUE(simplexFromConstraints(3,
+                                     {
+                                         {2, 0, 0, -1},
+                                         {-2, 0, 0, 1},
+                                         {0, 2, 0, -1},
+                                         {0, -2, 0, 1},
+                                         {0, 0, 2, -1},
+                                     },
+                                     {})
+                  .isUnbounded());
+
+  // Bounded set with equalities.
+  EXPECT_FALSE(simplexFromConstraints(2,
+                                      {{1, 1, 1},    // x + y >= -1.
+                                       {-1, -1, 1}}, // x + y <=  1.
+                                      {{1, -1, 0}}   // x = y.
+                                      )
+                   .isUnbounded());
+
+  // Unbounded set with equalities.
+  EXPECT_TRUE(simplexFromConstraints(3,
+                                     {{1, 1, 1, 1},     // x + y + z >= -1.
+                                      {-1, -1, -1, 1}}, // x + y + z <=  1.
+                                     {{1, -1, -1, 0}}   // x = y + z.
+                                     )
+                  .isUnbounded());
+
+  // Rational empty set.
+  EXPECT_FALSE(simplexFromConstraints(3,
+                                      {
+                                          {2, 0, 0, -1},
+                                          {-2, 0, 0, 1},
+                                          {0, 2, 2, -1},
+                                          {0, -2, -2, 1},
+                                          {3, 3, 3, -4},
+                                      },
+                                      {})
+                   .isUnbounded());
+}
+
+TEST(SimplexTest, getSamplePointIfIntegral) {
+  // Empty set.
+  EXPECT_FALSE(simplexFromConstraints(3,
+                                      {
+                                          {2, 0, 0, -1},
+                                          {-2, 0, 0, 1},
+                                          {0, 2, 2, -1},
+                                          {0, -2, -2, 1},
+                                          {3, 3, 3, -4},
+                                      },
+                                      {})
+                   .getSamplePointIfIntegral()
+                   .hasValue());
+
+  auto maybeSample = simplexFromConstraints(2,
+                                            {// x = y - 2.
+                                             {1, -1, 2},
+                                             {-1, 1, -2},
+                                             // x + y = 2.
+                                             {1, 1, -2},
+                                             {-1, -1, 2}},
+                                            {})
+                         .getSamplePointIfIntegral();
+
+  EXPECT_TRUE(maybeSample.hasValue());
+  EXPECT_THAT(*maybeSample, testing::ElementsAre(0, 2));
+
+  auto maybeSample2 = simplexFromConstraints(2,
+                                             {
+                                                 {1, 0, 0},  // x >= 0.
+                                                 {-1, 0, 0}, // x <= 0.
+                                             },
+                                             {
+                                                 {0, 1, -2} // y = 2.
+                                             })
+                          .getSamplePointIfIntegral();
+  EXPECT_TRUE(maybeSample2.hasValue());
+  EXPECT_THAT(*maybeSample2, testing::ElementsAre(0, 2));
+
+  EXPECT_FALSE(simplexFromConstraints(1,
+                                      {// 2x = 1. (no integer solutions)
+                                       {2, -1},
+                                       {-2, +1}},
+                                      {})
+                   .getSamplePointIfIntegral()
+                   .hasValue());
+}
+
+} // namespace mlir

mlir/unittests/CMakeLists.txt

 add_custom_target(MLIRUnitTests)
 set_target_properties(MLIRUnitTests PROPERTIES FOLDER "MLIR Tests")
 
 function(add_mlir_unittest test_dirname)
   add_unittest(MLIRUnitTests ${test_dirname} ${ARGN})
 endfunction()
 
+add_subdirectory(Analysis)
 add_subdirectory(Dialect)
 add_subdirectory(IR)
 add_subdirectory(Pass)
 add_subdirectory(SDBM)
 add_subdirectory(TableGen)