Introduce subtraction for FlatAffineConstraints

Subtraction is a foundational arithmetic operation that is often used when computing, for example, data transfer sets or cache hits. Since the result of subtraction need not be a convex polytope, a new class `PresburgerSet` is introduced to represent unions of convex polytopes.

Reviewed By: ftynse, bondhugula

Differential Revision: https://reviews.llvm.org/D87068

mlir/include/mlir/Analysis/AffineStructures.h

@@ -97,6 +97,13 @@ class FlatAffineConstraints {
       ids.append(idArgs.begin(), idArgs.end());
   }
 
+  /// Return a system with no constraints, i.e., one which is satisfied by all
+  /// points.
+  static FlatAffineConstraints getUniverse(unsigned numDims = 0,
+                                           unsigned numSymbols = 0) {
+    return FlatAffineConstraints(numDims, numSymbols);
+  }
+
   /// Create a flat affine constraint system from an AffineValueMap or a list of
   /// these. The constructed system will only include equalities.
   explicit FlatAffineConstraints(const AffineValueMap &avm);
@@ -153,6 +160,10 @@ class FlatAffineConstraints {
   /// Returns such a point if one exists, or an empty Optional otherwise.
   Optional<SmallVector<int64_t, 8>> findIntegerSample() const;
 
+  /// Returns true if the given point satisfies the constraints, or false
+  /// otherwise.
+  bool containsPoint(ArrayRef<int64_t> point) const;
+
   // Clones this object.
   std::unique_ptr<FlatAffineConstraints> clone() const;
 

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

@@ -169,6 +169,9 @@ class Simplex {
   /// Rollback to a snapshot. This invalidates all later snapshots.
   void rollback(unsigned snapshot);
 
+  /// Add all the constraints from the given FlatAffineConstraints.
+  void intersectFlatAffineConstraints(const FlatAffineConstraints &fac);
+
   /// Compute the maximum or minimum value of the given row, depending on
   /// direction. The specified row is never pivoted.
   ///

mlir/include/mlir/Analysis/PresburgerSet.h

@@ -0,0 +1,112 @@
+//===- Set.h - MLIR PresburgerSet 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
+//
+//===----------------------------------------------------------------------===//
+//
+// A class to represent unions of FlatAffineConstraints.
+//
+//===----------------------------------------------------------------------===//
+
+#ifndef MLIR_ANALYSIS_PRESBURGERSET_H
+#define MLIR_ANALYSIS_PRESBURGERSET_H
+
+#include "mlir/Analysis/AffineStructures.h"
+
+namespace mlir {
+
+/// This class can represent a union of FlatAffineConstraints, with support for
+/// union, intersection, subtraction and complement operations, as well as
+/// sampling.
+///
+/// The FlatAffineConstraints (FACs) are stored in a vector, and the set
+/// represents the union of these FACs. An empty list corresponds to the empty
+/// set.
+///
+/// Note that there are no invariants guaranteed on the list of FACs other than
+/// that they are all in the same space, i.e., they all have the same number of
+/// dimensions and symbols. For example, the FACs may overlap each other.
+class PresburgerSet {
+public:
+  explicit PresburgerSet(const FlatAffineConstraints &fac);
+
+  /// Return the number of FACs in the union.
+  unsigned getNumFACs() const;
+
+  /// Return the number of real dimensions.
+  unsigned getNumDims() const;
+
+  /// Return the number of symbolic dimensions.
+  unsigned getNumSyms() const;
+
+  /// Return a reference to the list of FlatAffineConstraints.
+  ArrayRef<FlatAffineConstraints> getAllFlatAffineConstraints() const;
+
+  /// Return the FlatAffineConstraints at the specified index.
+  const FlatAffineConstraints &getFlatAffineConstraints(unsigned index) const;
+
+  /// Mutate this set, turning it into the union of this set and the given
+  /// FlatAffineConstraints.
+  void unionFACInPlace(const FlatAffineConstraints &fac);
+
+  /// Mutate this set, turning it into the union of this set and the given set.
+  void unionSetInPlace(const PresburgerSet &set);
+
+  /// Return the union of this set and the given set.
+  PresburgerSet unionSet(const PresburgerSet &set) const;
+
+  /// Return the intersection of this set and the given set.
+  PresburgerSet intersect(const PresburgerSet &set) const;
+
+  /// Return true if the set contains the given point, or false otherwise.
+  bool containsPoint(ArrayRef<int64_t> point) const;
+
+  /// Print the set's internal state.
+  void print(raw_ostream &os) const;
+  void dump() const;
+
+  /// Return the complement of this set.
+  PresburgerSet complement() const;
+
+  /// Return the set difference of this set and the given set, i.e.,
+  /// return `this \ set`.
+  PresburgerSet subtract(const PresburgerSet &set) const;
+
+  /// Return a universe set of the specified type that contains all points.
+  static PresburgerSet getUniverse(unsigned nDim = 0, unsigned nSym = 0);
+  /// Return an empty set of the specified type that contains no points.
+  static PresburgerSet getEmptySet(unsigned nDim = 0, unsigned nSym = 0);
+
+  /// Return true if all the sets in the union are known to be integer empty
+  /// false otherwise.
+  bool isIntegerEmpty() const;
+
+  /// Find an integer sample from the given set. This should not be called if
+  /// any of the FACs in the union are unbounded.
+  bool findIntegerSample(SmallVectorImpl<int64_t> &sample);
+
+private:
+  /// Construct an empty PresburgerSet.
+  PresburgerSet(unsigned nDim = 0, unsigned nSym = 0)
+      : nDim(nDim), nSym(nSym) {}
+
+  /// Return the set difference fac \ set.
+  static PresburgerSet getSetDifference(FlatAffineConstraints fac,
+                                        const PresburgerSet &set);
+
+  /// Number of identifiers corresponding to real dimensions.
+  unsigned nDim;
+
+  /// Number of symbolic dimensions, unknown but constant for analysis, as in
+  /// FlatAffineConstraints.
+  unsigned nSym;
+
+  /// The list of flatAffineConstraints that this set is the union of.
+  SmallVector<FlatAffineConstraints, 2> flatAffineConstraints;
+};
+
+} // namespace mlir
+
+#endif // MLIR_ANALYSIS_PRESBURGERSET_H

mlir/lib/Analysis/AffineStructures.cpp

@@ -1056,6 +1056,33 @@ FlatAffineConstraints::findIntegerSample() const {
   return Simplex(*this).findIntegerSample();
 }
 
+/// Helper to evaluate an affine expression at a point.
+/// The expression is a list of coefficients for the dimensions followed by the
+/// constant term.
+static int64_t valueAt(ArrayRef<int64_t> expr, ArrayRef<int64_t> point) {
+  assert(expr.size() == 1 + point.size() &&
+         "Dimensionalities of point and expresion don't match!");
+  int64_t value = expr.back();
+  for (unsigned i = 0; i < point.size(); ++i)
+    value += expr[i] * point[i];
+  return value;
+}
+
+/// A point satisfies an equality iff the value of the equality at the
+/// expression is zero, and it satisfies an inequality iff the value of the
+/// inequality at that point is non-negative.
+bool FlatAffineConstraints::containsPoint(ArrayRef<int64_t> point) const {
+  for (unsigned i = 0, e = getNumEqualities(); i < e; ++i) {
+    if (valueAt(getEquality(i), point) != 0)
+      return false;
+  }
+  for (unsigned i = 0, e = getNumInequalities(); i < e; ++i) {
+    if (valueAt(getInequality(i), point) < 0)
+      return false;
+  }
+  return true;
+}
+
 /// 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.,

mlir/lib/Analysis/CMakeLists.txt

@@ -5,6 +5,7 @@ set(LLVM_OPTIONAL_SOURCES
   Liveness.cpp
   LoopAnalysis.cpp
   NestedMatcher.cpp
+  PresburgerSet.cpp
   SliceAnalysis.cpp
   Utils.cpp
   )
@@ -25,7 +26,6 @@ add_mlir_library(MLIRAnalysis
   MLIRCallInterfaces
   MLIRControlFlowInterfaces
   MLIRInferTypeOpInterface
-  MLIRPresburger
   MLIRSCF
   )
 
@@ -34,6 +34,7 @@ add_mlir_library(MLIRLoopAnalysis
   AffineStructures.cpp
   LoopAnalysis.cpp
   NestedMatcher.cpp
+  PresburgerSet.cpp
   Utils.cpp
 
   ADDITIONAL_HEADER_DIRS
@@ -51,4 +52,4 @@ add_mlir_library(MLIRLoopAnalysis
   MLIRSCF
   )
 
-add_subdirectory(Presburger)
+add_subdirectory(Presburger)

mlir/lib/Analysis/Presburger/CMakeLists.txt

@@ -1,4 +1,4 @@
 add_mlir_library(MLIRPresburger
   Simplex.cpp
   Matrix.cpp
-  )
+  )

mlir/lib/Analysis/Presburger/Simplex.cpp

@@ -451,6 +451,16 @@ void Simplex::rollback(unsigned snapshot) {
   }
 }
 
+/// Add all the constraints from the given FlatAffineConstraints.
+void Simplex::intersectFlatAffineConstraints(const FlatAffineConstraints &fac) {
+  assert(fac.getNumIds() == numVariables() &&
+         "FlatAffineConstraints must have same dimensionality as simplex");
+  for (unsigned i = 0, e = fac.getNumInequalities(); i < e; ++i)
+    addInequality(fac.getInequality(i));
+  for (unsigned i = 0, e = fac.getNumEqualities(); i < e; ++i)
+    addEquality(fac.getEquality(i));
+}
+
 Optional<Fraction> Simplex::computeRowOptimum(Direction direction,
                                               unsigned row) {
   // Keep trying to find a pivot for the row in the specified direction.