[MLIR][Presburger] Remove `spaceKind` from PresburgerSpace

[Groverkss]

Summary

This patch remove spaceKind from PresburgerSpace, making PresburgerSpace only
a space supporting relations.

Sets are still implemented in the same way, i.e. with a zero domain but instead
the asserts to check if the space is still set are added to users of
PresburgerSpace which treat it as a Set space.

[arjunp]

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

154–155

mlir/lib/Analysis/Presburger/PresburgerSpace.cpp

142–144

PresburgerSpace::print(os) would deduplicate this

[Groverkss]

Address Arjun's comments

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

public:
   /// The number of dimensions and symbol ids in `this` and `other` should
   /// match.
   void mergeLocalIds(IntegerRelation &other);
 
   void print(raw_ostream &os) const;
   void dump() const;
 
 protected:
-  /// Constructs a set reserving memory for the specified number
-  /// of constraints and identifiers. This constructor should not be used
-  /// directly to create a relation and should only be used to create Sets.
-  /// Internally this constructs a relation with with no domain and a
-  /// space with no distinction between domain and range identifiers.
-  IntegerRelation(unsigned numReservedInequalities,
-                  unsigned numReservedEqualities, unsigned numReservedCols,
-                  unsigned numDims, unsigned numSymbols, unsigned numLocals)
-      : PresburgerLocalSpace(numDims, numSymbols, numLocals),
-        equalities(0, getNumIds() + 1, numReservedEqualities, numReservedCols),
-        inequalities(0, getNumIds() + 1, numReservedInequalities,
-                     numReservedCols) {
-    assert(numReservedCols >= getNumIds() + 1);
-  }
-
   /// Checks all rows of equality/inequality constraints for trivial
   /// contradictions (for example: 1 == 0, 0 >= 1), which may have surfaced
   /// after elimination. Returns true if an invalid constraint is found;
   /// false otherwise.
   bool hasInvalidConstraint() const;
 
   /// Returns the constant lower bound bound if isLower is true, and the upper
   /// bound if isLower is false.

 class IntegerPolyhedron : public IntegerRelation {
 public:
   /// Constructs a set reserving memory for the specified number
   /// of constraints and identifiers.
   IntegerPolyhedron(unsigned numReservedInequalities,
                     unsigned numReservedEqualities, unsigned numReservedCols,
                     unsigned numDims, unsigned numSymbols, unsigned numLocals)
       : IntegerRelation(numReservedInequalities, numReservedEqualities,
-                        numReservedCols, numDims, numSymbols, numLocals) {}
+                        numReservedCols, /*numDomain=*/0, /*numRange=*/numDims,
+                        numSymbols, numLocals) {}
 
   /// Constructs a relation with the specified number of dimensions and symbols.
   IntegerPolyhedron(unsigned numDims = 0, unsigned numSymbols = 0,
                     unsigned numLocals = 0)
       : IntegerPolyhedron(/*numReservedInequalities=*/0,
                           /*numReservedEqualities=*/0,
                           /*numReservedCols=*/numDims + numSymbols + numLocals +
                               1,
public:
   Kind getKind() const override { return Kind::IntegerPolyhedron; }
 
   static bool classof(const IntegerRelation *cst) {
     return cst->getKind() == Kind::IntegerPolyhedron;
   }
 
   // Clones this object.
   std::unique_ptr<IntegerPolyhedron> clone() const;
+
+  /// Insert `num` identifiers of the specified kind at position `pos`.
+  /// Positions are relative to the kind of identifier. Return the absolute
+  /// column position (i.e., not relative to the kind of identifier) of the
+  /// first added identifier.
+  unsigned insertId(IdKind kind, unsigned pos, unsigned num = 1) override;
 };
 
 } // namespace presburger
 } // namespace mlir
 
 #endif // MLIR_ANALYSIS_PRESBURGER_INTEGERPOLYHEDRON_H

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

 /// symbolic ids.
 ///
 /// Support is provided to compare equality of two such functions as well as
 /// finding the value of the function at a point. Note that local ids in the
 /// piece are not supported for the latter.
 class PWMAFunction : public PresburgerSpace {
 public:
   PWMAFunction(unsigned numDims, unsigned numSymbols, unsigned numOutputs)
-      : PresburgerSpace(numDims, numSymbols), numOutputs(numOutputs) {
+      : PresburgerSpace(/*numDomain=*/0, /*numRange=*/numDims, numSymbols),
+        numOutputs(numOutputs) {
     assert(numOutputs >= 1 && "The function must output something!");
   }
 
   void addPiece(const MultiAffineFunction &piece);
   void addPiece(const IntegerPolyhedron &domain, const Matrix &output);
 
   const MultiAffineFunction &getPiece(unsigned i) const { return pieces[i]; }
   unsigned getNumPieces() const { return pieces.size(); }

mlir/include/mlir/Analysis/Presburger/PresburgerSet.h

public:
   /// Print the set's internal state.
   void print(raw_ostream &os) const;
   void dump() const;
 
 private:
   /// Construct an empty PresburgerSet with the specified number of dimension
   /// and symbols.
   PresburgerSet(unsigned numDims = 0, unsigned numSymbols = 0)
-      : PresburgerSpace(numDims, numSymbols) {}
+      : PresburgerSpace(/*numDomain=*/0, /*numRange=*/numDims, numSymbols) {}
 
   /// The list of polyhedrons that this set is the union of.
   SmallVector<IntegerPolyhedron, 2> integerPolyhedrons;
 };
 
 } // namespace presburger
 } // namespace mlir
 
 #endif // MLIR_ANALYSIS_PRESBURGER_PRESBURGERSET_H

mlir/include/mlir/Analysis/Presburger/PresburgerSpace.h

 /// of projection. In this example, `q` is existentially quantified. This can be
 /// thought of as the result of projecting out `q` from the previous example,
 /// i.e. we obtained {2, 4, 6} by projecting out the second dimension from
 /// {(2, 1), (4, 2), (6, 2)}.
 ///
 /// Dimension identifiers are further divided into Domain and Range identifiers
 /// to support building relations.
 ///
-/// Spaces with distinction between domain and range identifiers should use
-/// IdKind::Domain and IdKind::Range to refer to domain and range identifiers.
-/// Identifiers for such spaces are stored in the following order:
+/// Identifiers are stored in the following order:
 ///       [Domain, Range, Symbols, Locals]
 ///
-/// Spaces with no distinction between domain and range identifiers should use
-/// IdKind::SetDim to refer to dimension identifiers. Identifiers for such
-/// spaces are stored in the following order:
-///       [SetDim, Symbol, Locals]
+/// A space with no distinction between types of dimension identifiers can
+/// be implemented as a space with zero domain. IdKind::SetDim should be used
+/// to refer to dimensions in such spaces.
 ///
 /// PresburgerSpace does not allow identifiers of kind Local. See
 /// PresburgerLocalSpace for an extension that does allow local identifiers.
 class PresburgerSpace {
   friend PresburgerLocalSpace;
 
 public:
-  static PresburgerSpace getRelationSpace(unsigned numDomain, unsigned numRange,
-                                          unsigned numSymbols);
-
-  static PresburgerSpace getSetSpace(unsigned numDims, unsigned numSymbols);
+  PresburgerSpace(unsigned numDomain, unsigned numRange, unsigned numSymbols)
+      : PresburgerSpace(numDomain, numRange, numSymbols, 0) {}
 
   virtual ~PresburgerSpace() = default;
 
   unsigned getNumDomainIds() const { return numDomain; }
   unsigned getNumRangeIds() const { return numRange; }
   unsigned getNumSymbolIds() const { return numSymbols; }
   unsigned getNumSetDimIds() const { return numRange; }
 
public:
   /// symbol count, either a chunk of dimensional identifiers immediately before
   /// the split become symbols, or some of the symbols immediately after the
   /// split become dimensions.
   void setDimSymbolSeparation(unsigned newSymbolCount);
 
   void print(llvm::raw_ostream &os) const;
   void dump() const;
 
-protected:
-  /// Space constructor for Relation space type.
-  PresburgerSpace(unsigned numDomain, unsigned numRange, unsigned numSymbols)
-      : PresburgerSpace(Relation, numDomain, numRange, numSymbols,
-                        /*numLocals=*/0) {}
-
-  /// Space constructor for Set space type.
-  PresburgerSpace(unsigned numDims, unsigned numSymbols)
-      : PresburgerSpace(Set, /*numDomain=*/0, numDims, numSymbols,
-                        /*numLocals=*/0) {}
-
 private:
-  /// Kind of space.
-  enum SpaceKind { Set, Relation };
-
-  PresburgerSpace(SpaceKind spaceKind, unsigned numDomain, unsigned numRange,
-                  unsigned numSymbols, unsigned numLocals)
-      : spaceKind(spaceKind), numDomain(numDomain), numRange(numRange),
-        numSymbols(numSymbols), numLocals(numLocals) {}
-
-  SpaceKind spaceKind;
+  PresburgerSpace(unsigned numDomain, unsigned numRange, unsigned numSymbols,
+                  unsigned numLocals)
+      : numDomain(numDomain), numRange(numRange), numSymbols(numSymbols),
+        numLocals(numLocals) {}
 
   // Number of identifiers corresponding to domain identifiers.
   unsigned numDomain;
 
   // Number of identifiers corresponding to range identifiers.
   unsigned numRange;
 
   /// Number of identifiers corresponding to symbols (unknown but constant for
   /// analysis).
   unsigned numSymbols;
 
   /// Total number of identifiers.
   unsigned numLocals;
 };
 
 /// Extension of PresburgerSpace supporting Local identifiers.
 class PresburgerLocalSpace : public PresburgerSpace {
 public:
-  static PresburgerLocalSpace getRelationSpace(unsigned numDomain,
-                                               unsigned numRange,
-                                               unsigned numSymbols,
-                                               unsigned numLocals);
-
-  static PresburgerLocalSpace getSetSpace(unsigned numDims, unsigned numSymbols,
-                                          unsigned numLocals);
+  PresburgerLocalSpace(unsigned numDomain, unsigned numRange,
+                       unsigned numSymbols, unsigned numLocals)
+      : PresburgerSpace(numDomain, numRange, numSymbols, numLocals) {}
 
   unsigned getNumLocalIds() const { return numLocals; }
 
   /// Insert `num` identifiers of the specified kind at position `pos`.
   /// Positions are relative to the kind of identifier. Return the absolute
   /// column position (i.e., not relative to the kind of identifier) of the
   /// first added identifier.
   unsigned insertId(IdKind kind, unsigned pos, unsigned num = 1) override;
 
   /// Removes identifiers in the column range [idStart, idLimit).
   void removeIdRange(IdKind kind, unsigned idStart, unsigned idLimit) override;
 
   /// Returns true if both the spaces are equal i.e. if both spaces have the
   /// same number of identifiers of each kind.
   bool isEqual(const PresburgerLocalSpace &other) const;
 
   void print(llvm::raw_ostream &os) const;
   void dump() const;
-
-protected:
-  /// Local Space constructor for Relation space type.
-  PresburgerLocalSpace(unsigned numDomain, unsigned numRange,
-                       unsigned numSymbols, unsigned numLocals)
-      : PresburgerSpace(Relation, numDomain, numRange, numSymbols, numLocals) {}
-
-  /// Local Space constructor for Set space type.
-  PresburgerLocalSpace(unsigned numDims, unsigned numSymbols,
-                       unsigned numLocals)
-      : PresburgerSpace(Set, /*numDomain=*/0, numDims, numSymbols, numLocals) {}
 };
 
 } // namespace presburger
 } // namespace mlir
 
 #endif // MLIR_ANALYSIS_PRESBURGER_PRESBURGERSPACE_H

mlir/lib/Analysis/Presburger/IntegerRelation.cpp

for (unsigned j = 0, f = getNumCols(); j < f; ++j) {
       os << atIneq(i, j) << " ";
     }
     os << ">= 0\n";
   }
   os << '\n';
 }
 
 void IntegerRelation::dump() const { print(llvm::errs()); }
+
+unsigned IntegerPolyhedron::insertId(IdKind kind, unsigned pos, unsigned num) {
+  assert((kind != IdKind::Domain || num == 0) &&
+         "Domain has to be zero in a set");
+  return IntegerRelation::insertId(kind, pos, num);
+}

mlir/lib/Analysis/Presburger/PWMAFunction.cpp

 bool MultiAffineFunction::isEqual(const MultiAffineFunction &other) const {
   return PresburgerSpace::isEqual(other) &&
          getDomain().isEqual(other.getDomain()) &&
          isEqualWhereDomainsOverlap(other);
 }
 
 unsigned MultiAffineFunction::insertId(IdKind kind, unsigned pos,
                                        unsigned num) {
+  assert((kind != IdKind::Domain || num == 0) &&
+         "Domain has to be zero in a set");
   unsigned absolutePos = getIdKindOffset(kind) + pos;
   output.insertColumns(absolutePos, num);
   return IntegerPolyhedron::insertId(kind, pos, num);
 }
 
 void MultiAffineFunction::swapId(unsigned posA, unsigned posB) {
   output.swapColumns(posA, posB);
   IntegerPolyhedron::swapId(posA, posB);

mlir/lib/Analysis/Presburger/PresburgerSpace.cpp

 //===- PresburgerSpace.cpp - MLIR PresburgerSpace 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/PresburgerSpace.h"
 #include <algorithm>
 #include <cassert>
 
 using namespace mlir;
 using namespace presburger;
 
-PresburgerSpace PresburgerSpace::getRelationSpace(unsigned numDomain,
-                                                  unsigned numRange,
-                                                  unsigned numSymbols) {
-  return PresburgerSpace(numDomain, numRange, numSymbols);
-}
-
-PresburgerSpace PresburgerSpace::getSetSpace(unsigned numDims,
-                                             unsigned numSymbols) {
-  return PresburgerSpace(numDims, numSymbols);
-}
-
-PresburgerLocalSpace
-PresburgerLocalSpace::getRelationSpace(unsigned numDomain, unsigned numRange,
-                                       unsigned numSymbols,
-                                       unsigned numLocals) {
-  return PresburgerLocalSpace(numDomain, numRange, numSymbols, numLocals);
-}
-
-PresburgerLocalSpace PresburgerLocalSpace::getSetSpace(unsigned numDims,
-                                                       unsigned numSymbols,
-                                                       unsigned numLocals) {
-  return PresburgerLocalSpace(numDims, numSymbols, numLocals);
-}
-
 unsigned PresburgerSpace::getNumIdKind(IdKind kind) const {
   if (kind == IdKind::Domain)
     return getNumDomainIds();
   if (kind == IdKind::Range)
     return getNumRangeIds();
   if (kind == IdKind::Symbol)
     return getNumSymbolIds();
   if (kind == IdKind::Local)
unsigned PresburgerSpace::getIdKindOverlap(IdKind kind, unsigned idStart,
   return overlapEnd - overlapStart;
 }
 
 unsigned PresburgerSpace::insertId(IdKind kind, unsigned pos, unsigned num) {
   assert(pos <= getNumIdKind(kind));
 
   unsigned absolutePos = getIdKindOffset(kind) + pos;
 
-  if (kind == IdKind::Domain) {
-    assert(spaceKind == Relation && "IdKind::Domain is not supported in Set.");
+  if (kind == IdKind::Domain)
     numDomain += num;
-  } else if (kind == IdKind::Range) {
+  else if (kind == IdKind::Range)
     numRange += num;
-  } else if (kind == IdKind::Symbol) {
+  else if (kind == IdKind::Symbol)
     numSymbols += num;
-  } else {
+  else
     llvm_unreachable("PresburgerSpace does not support local identifiers!");
-  }
 
   return absolutePos;
 }
 
 void PresburgerSpace::removeIdRange(IdKind kind, unsigned idStart,
                                     unsigned idLimit) {
   assert(idLimit <= getNumIdKind(kind) && "invalid id limit");
 
   if (idStart >= idLimit)
     return;
 
   unsigned numIdsEliminated = idLimit - idStart;
-  if (kind == IdKind::Domain) {
-    assert(spaceKind == Relation && "IdKind::Domain is not supported in Set.");
+  if (kind == IdKind::Domain)
     numDomain -= numIdsEliminated;
-  } else if (kind == IdKind::Range) {
+  else if (kind == IdKind::Range)
     numRange -= numIdsEliminated;
-  } else if (kind == IdKind::Symbol) {
+  else if (kind == IdKind::Symbol)
     numSymbols -= numIdsEliminated;
-  } else {
+  else
     llvm_unreachable("PresburgerSpace does not support local identifiers!");
-  }
-}
 
 unsigned PresburgerLocalSpace::insertId(IdKind kind, unsigned pos,
                                         unsigned num) {
   if (kind == IdKind::Local) {
     numLocals += num;
     return getIdKindOffset(IdKind::Local) + pos;
   }
   return PresburgerSpace::insertId(kind, pos, num);

 void PresburgerSpace::setDimSymbolSeparation(unsigned newSymbolCount) {
   assert(newSymbolCount <= getNumDimAndSymbolIds() &&
          "invalid separation position");
   numRange = numRange + numSymbols - newSymbolCount;
   numSymbols = newSymbolCount;
 }
 
 void PresburgerSpace::print(llvm::raw_ostream &os) const {
-  if (spaceKind == Relation) {
-    os << "Domain: " << getNumDomainIds() << ", "
-       << "Range: " << getNumRangeIds() << ", ";
-  } else {
-    os << "Dimension: " << getNumDomainIds() << ", ";
-  }
-  os << "Symbols: " << getNumSymbolIds() << "\n";
+     << "Range: " << getNumRangeIds() << ", "
+     << "Symbols: " << getNumSymbolIds() << "\n";
 }
 
 void PresburgerSpace::dump() const { print(llvm::errs()); }
 
 void PresburgerLocalSpace::print(llvm::raw_ostream &os) const {
-  if (spaceKind == Relation) {
-    os << "Domain: " << getNumDomainIds() << ", "
-       << "Range: " << getNumRangeIds() << ", ";
-  } else {
-    os << "Dimension: " << getNumDomainIds() << ", ";
-  }
-  os << "Symbols: " << getNumSymbolIds() << ", "
-     << "Locals: " << getNumLocalIds() << "\n";
+  PresburgerSpace::print(os);
+  os << "Locals: " << getNumLocalIds() << "\n";
 }
 
 void PresburgerLocalSpace::dump() const { print(llvm::errs()); }

mlir/unittests/Analysis/Presburger/PresburgerSpaceTest.cpp

 //===- PresburgerSpaceTest.cpp - Tests for PresburgerSpace ----------------===//
 //
 // 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/PresburgerSpace.h"
 #include <gmock/gmock.h>
 #include <gtest/gtest.h>
 
 using namespace mlir;
 using namespace presburger;
 
 TEST(PresburgerSpaceTest, insertId) {
-  PresburgerSpace space = PresburgerSpace::getRelationSpace(2, 2, 1);
+  PresburgerSpace space(2, 2, 1);
 
   // Try inserting 2 domain ids.
   space.insertId(IdKind::Domain, 0, 2);
   EXPECT_EQ(space.getNumDomainIds(), 4u);
 
   // Try inserting 1 range ids.
   space.insertId(IdKind::Range, 0, 1);
   EXPECT_EQ(space.getNumRangeIds(), 3u);
 }
 
 TEST(PresburgerSpaceTest, insertIdSet) {
-  PresburgerSpace space = PresburgerSpace::getSetSpace(2, 1);
+  PresburgerSpace space(0, 2, 1);
 
   // Try inserting 2 dimension ids. The space should have 4 range ids since
   // spaces which do not distinguish between domain, range are implemented like
   // this.
   space.insertId(IdKind::SetDim, 0, 2);
   EXPECT_EQ(space.getNumRangeIds(), 4u);
 }
 
 TEST(PresburgerSpaceTest, removeIdRange) {
-  PresburgerSpace space = PresburgerSpace::getRelationSpace(2, 1, 3);
+  PresburgerSpace space(2, 1, 3);
 
   // Remove 1 domain identifier.
   space.removeIdRange(IdKind::Domain, 0, 1);
   EXPECT_EQ(space.getNumDomainIds(), 1u);
 
   // Remove 1 symbol and 1 range identifier.
   space.removeIdRange(IdKind::Symbol, 0, 1);
   space.removeIdRange(IdKind::Range, 0, 1);
   EXPECT_EQ(space.getNumDomainIds(), 1u);
   EXPECT_EQ(space.getNumRangeIds(), 0u);
   EXPECT_EQ(space.getNumSymbolIds(), 2u);
 }