[mlir][Presburger] Introduce Domain and Range identifiers in PresburgerSpace

[Groverkss]

Summary

This patch introducing seperating dimensions into two types: Domain and Range.
This allows building relations over PresburgerSpace.

This patch is part of a series of patches to introduce relations in Presburger
library.

[bondhugula]

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

17

Do you need this header?

55–58

Use /*...=*/... format when passing constant values.

mlir/lib/Analysis/Presburger/PresburgerSpace.cpp

50

exist

64

exist

[Groverkss]

• Address Uday's comments

[Groverkss]

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

17

Removed the header.

[Groverkss]

• Remove accidently added file

[Groverkss]

• Fix compilation issue

[arjunp]

mlir/lib/Analysis/Presburger/PresburgerSpace.cpp

92

This shouldn't be reached when getNumDomainIds != 0, right? I would add an assert for that and just increase range.

(In fact, maybe you can store an enum denoting whether this is a relation or not and assert throughout that the right IdKind is being used.

You can set this enum in the corresponding constructor, or maybe subclass to make a PresburgerSetSpace that has that constructor.)

[Groverkss]

mlir/lib/Analysis/Presburger/PresburgerSpace.cpp

92

I think I can add the assert.

As for the enum, I don't think I want the implementation detail that a set is a relation with only range to be hidden. It should be rather visible that a set is a relation with only range.

[Groverkss]

• Address Arjun's comments

[arjunp]

mlir/lib/Analysis/Presburger/PresburgerSpace.cpp

92

It's not really about hiding the implementation detail. I just meant that we can then use asserts to check at runtime that dimension idkinds is used when there is no distinction and range/domain idkinds are used only when there is distinction

[Groverkss]

• Add enum to distinguish between Set and Relation.

mlir/lib/Analysis/Presburger/PresburgerSpace.cpp

92

Ok makes sense. WIll add the enum.

[arjunp]

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

40

I see, so we want to support using Dimension IdKind in relation spaces as well to consider domain + range as one block? This seems a little weird, is it needed?

142

These enums are usually named like SpaceKind I think

mlir/lib/Analysis/Presburger/PresburgerSpace.cpp

92

You can add these asserts throughout, wherever idkinds are being used. You can make a member function isSupportedIdKind

[Groverkss]

• Add asserts to ensure domain is not used when a space is a Set.
• Use IdKind::SetDim instead of IdKind::Dimension.

[Groverkss]

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

40

It does seem a bit weird maybe. I tried a different design now.

[arjunp]

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

44

I would mention that users should use PresburgerLocalSpace if they need that

51

I would mention that in the implementation of the class, SetDims are treated as range ids, and spaces with no distinction are treated as relations with zero domain ids

127–133

Personally, I prefer spacekind to be the first param and the first thing initialized

147

It's hard to see the different between the two constructors at the callsite. I would instead make two static functions, PresburgerLocalSpace::getRelationSpace() and getSetSpace() or something

Similarly for PresburgerSpace above

mlir/lib/Analysis/Presburger/PresburgerSpace.cpp

164

special case for SetDim spaceKind?

174

"Locals: " missing?

[Groverkss]

• Address Arjun's comments

[arjunp]

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

147

I would just remove the two extra constructors, they aren't needed and we can keep only the static functions for constructing sets, which would call into the single non-public constructor.

[Groverkss]

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

147

These constructors are supposed to be called by IntegerPolyhedron, IntegerRelation (in future) and any other derived classes of PresburgerSpace. I don't think we can use static functions for that.

[arjunp]

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

147

Right.

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

 //===- PresburgerSpace.h - MLIR PresburgerSpace 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
 //
 //===----------------------------------------------------------------------===//
 //
 // Classes representing space information like number of identifiers and kind of
 // identifiers.
 //
 //===----------------------------------------------------------------------===//
 
 #ifndef MLIR_ANALYSIS_PRESBURGER_PRESBURGERSPACE_H
 #define MLIR_ANALYSIS_PRESBURGER_PRESBURGERSPACE_H
 
 #include "llvm/Support/ErrorHandling.h"
+#include "llvm/Support/raw_ostream.h"
 
 namespace mlir {
 
 class PresburgerLocalSpace;
 
 /// PresburgerSpace is a tuple of identifiers with information about what kind
 /// they correspond to. The identifiers can be split into three types:
 ///
 /// Dimension: Ordinary variables over which the space is represented.
 ///
 /// Symbol: Symbol identifiers correspond to fixed but unknown values.
 /// Mathematically, a space with symbolic identifiers is like a
 /// family of spaces indexed by the symbolic identifiers.
 ///
 /// Local: Local identifiers correspond to existentially quantified variables.
 ///
-/// PresburgerSpace only supports identifiers of kind Dimension and Symbol.
+/// 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.
+///
+/// Spaces with no distinction between domain and range identifiers should use
+/// IdKind::SetDim to refer to dimension identifiers.
+///
+/// PresburgerSpace does not support identifiers of kind Local. See
+/// PresburgerLocalSpace for an extension that supports Local ids.
 class PresburgerSpace {
   friend PresburgerLocalSpace;
 
 public:
-  /// Kind of identifier (column).
-  enum IdKind { Dimension, Symbol, Local };
+  /// Kind of identifier. Implementation wise SetDims are treated as Range
+  /// ids, and spaces with no distinction between dimension ids are treated
+  /// as relations with zero domain ids.
+  enum IdKind { Symbol, Local, Domain, Range, SetDim = Range };
 
-  PresburgerSpace(unsigned numDims, unsigned numSymbols)
-      : numDims(numDims), numSymbols(numSymbols), numLocals(0) {}
+  static PresburgerSpace getRelationSpace(unsigned numDomain, unsigned numRange,
+                                          unsigned numSymbols);
+
+  static PresburgerSpace getSetSpace(unsigned numDims, unsigned numSymbols);
 
   virtual ~PresburgerSpace() = default;
 
-  unsigned getNumIds() const { return numDims + numSymbols + numLocals; }
-  unsigned getNumDimIds() const { return numDims; }
+  unsigned getNumDomainIds() const { return numDomain; }
+  unsigned getNumRangeIds() const { return numRange; }
   unsigned getNumSymbolIds() const { return numSymbols; }
-  unsigned getNumDimAndSymbolIds() const { return numDims + numSymbols; }
+  unsigned getNumSetDimIds() const { return numRange; }
+
+  unsigned getNumDimIds() const { return numDomain + numRange; }
+  unsigned getNumDimAndSymbolIds() const {
+    return numDomain + numRange + numSymbols;
+  }
+  unsigned getNumIds() const {
+    return numDomain + numRange + numSymbols + numLocals;
+  }
 
   /// Get the number of ids of the specified kind.
   unsigned getNumIdKind(IdKind kind) const;
 
   /// Return the index at which the specified kind of id starts.
   unsigned getIdKindOffset(IdKind kind) const;
 
   /// Return the index at Which the specified kind of id ends.
public:
   virtual void removeIdRange(unsigned idStart, unsigned idLimit);
 
   /// Changes the partition between dimensions and symbols. Depending on the new
   /// 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:
-  PresburgerSpace(unsigned numDims, unsigned numSymbols, unsigned numLocals)
-      : numDims(numDims), numSymbols(numSymbols), numLocals(numLocals) {}
+  /// 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) {}
 
-  /// Number of identifiers corresponding to real dimensions.
-  unsigned numDims;
+  SpaceKind spaceKind;
+
+  // 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:
-  PresburgerLocalSpace(unsigned numDims, unsigned numSymbols,
-                       unsigned numLocals)
-      : PresburgerSpace(numDims, numSymbols, numLocals) {}
+  static PresburgerLocalSpace getRelationSpace(unsigned numDomain,
+                                               unsigned numRange,
+                                               unsigned numSymbols,
+                                               unsigned numLocals);
+
+  static PresburgerLocalSpace getSetSpace(unsigned numDims, unsigned numSymbols,
+                                          unsigned 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(unsigned idStart, unsigned idLimit) override;
+
+  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 mlir
 
 #endif // MLIR_ANALYSIS_PRESBURGER_PRESBURGERSPACE_H

mlir/lib/Analysis/Presburger/IntegerPolyhedron.cpp

IntegerPolyhedron::getRationalLexMin() const {
   // from the end of the answer gives the desired lexmin over the dimensions.
   assert(maybeLexMin->size() == getNumIds() &&
          "Incorrect number of vars in lexMin!");
   maybeLexMin->resize(getNumDimAndSymbolIds());
   return maybeLexMin;
 }
 
 unsigned IntegerPolyhedron::insertDimId(unsigned pos, unsigned num) {
-  return insertId(IdKind::Dimension, pos, num);
+  return insertId(IdKind::SetDim, pos, num);
 }
 
 unsigned IntegerPolyhedron::insertSymbolId(unsigned pos, unsigned num) {
   return insertId(IdKind::Symbol, pos, num);
 }
 
 unsigned IntegerPolyhedron::insertLocalId(unsigned pos, unsigned num) {
   return insertId(IdKind::Local, pos, num);
 }
 
 unsigned IntegerPolyhedron::insertId(IdKind kind, unsigned pos, unsigned num) {
   assert(pos <= getNumIdKind(kind));
 
-  unsigned absolutePos = getIdKindOffset(kind) + pos;
-  inequalities.insertColumns(absolutePos, num);
-  equalities.insertColumns(absolutePos, num);
-
-  return PresburgerLocalSpace::insertId(kind, pos, num);
+  unsigned insertPos = PresburgerLocalSpace::insertId(kind, pos, num);
+  inequalities.insertColumns(insertPos, num);
+  equalities.insertColumns(insertPos, num);
+  return insertPos;
 }
 
 unsigned IntegerPolyhedron::appendDimId(unsigned num) {
   unsigned pos = getNumDimIds();
-  insertId(IdKind::Dimension, pos, num);
+  insertId(IdKind::SetDim, pos, num);
   return pos;
 }
 
 unsigned IntegerPolyhedron::appendSymbolId(unsigned num) {
   unsigned pos = getNumSymbolIds();
   insertId(IdKind::Symbol, pos, num);
   return pos;
 }

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;
 
+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::Dimension)
-    return getNumDimIds();
+  if (kind == IdKind::Domain) {
+    assert(spaceKind == Relation && "IdKind::Domain is not supported in Set.");
+    return getNumDomainIds();
+  }
+  if (kind == IdKind::Range)
+    return getNumRangeIds();
   if (kind == IdKind::Symbol)
     return getNumSymbolIds();
   if (kind == IdKind::Local)
     return numLocals;
-  llvm_unreachable("IdKind does not exit!");
+  llvm_unreachable("IdKind does not exist!");
 }
 
 unsigned PresburgerSpace::getIdKindOffset(IdKind kind) const {
-  if (kind == IdKind::Dimension)
+  if (kind == IdKind::Domain) {
+    assert(spaceKind == Relation && "IdKind::Domain is not supported in Set.");
     return 0;
+  }
+  if (kind == IdKind::Range)
+    return getNumDomainIds();
   if (kind == IdKind::Symbol)
     return getNumDimIds();
   if (kind == IdKind::Local)
     return getNumDimAndSymbolIds();
-  llvm_unreachable("IdKind does not exit!");
+  llvm_unreachable("IdKind does not exist!");
 }
 
 unsigned PresburgerSpace::getIdKindEnd(IdKind kind) const {
   return getIdKindOffset(kind) + getNumIdKind(kind);
 }
 
 unsigned PresburgerSpace::getIdKindOverlap(IdKind kind, unsigned idStart,
                                            unsigned idLimit) const {
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::Dimension)
-    numDims += num;
-  else if (kind == IdKind::Symbol)
+  if (kind == IdKind::Domain) {
+    assert(spaceKind == Relation && "IdKind::Domain is not supported in Set.");
+    numDomain += num;
+  } else if (kind == IdKind::Range) {
+    numRange += num;
+  } else if (kind == IdKind::Symbol) {
     numSymbols += num;
-  else
-    llvm_unreachable(
-        "PresburgerSpace only supports Dimensions and Symbol identifiers!");
+  } else {
+    llvm_unreachable("PresburgerSpace does not support local identifiers!");
+  }
 
   return absolutePos;
 }
 
 void PresburgerSpace::removeIdRange(unsigned idStart, unsigned idLimit) {
   assert(idLimit <= getNumIds() && "invalid id limit");
 
   if (idStart >= idLimit)
     return;
 
   // We are going to be removing one or more identifiers from the range.
   assert(idStart < getNumIds() && "invalid idStart position");
 
-  // Update members numDims, numSymbols and numIds.
-  unsigned numDimsEliminated =
-      getIdKindOverlap(IdKind::Dimension, idStart, idLimit);
+  // Update members numDomain, numRange, numSymbols and numIds.
+  unsigned numDomainEliminated = 0;
+  if (spaceKind == Relation)
+    numDomainEliminated = getIdKindOverlap(IdKind::Domain, idStart, idLimit);
+  unsigned numRangeEliminated =
+      getIdKindOverlap(IdKind::Range, idStart, idLimit);
   unsigned numSymbolsEliminated =
       getIdKindOverlap(IdKind::Symbol, idStart, idLimit);
 
-  numDims -= numDimsEliminated;
+  numDomain -= numDomainEliminated;
+  numRange -= numRangeEliminated;
   numSymbols -= numSymbolsEliminated;
 }
 
 unsigned PresburgerLocalSpace::insertId(IdKind kind, unsigned pos,
                                         unsigned num) {
   if (kind == IdKind::Local) {
     numLocals += num;
     return getIdKindOffset(IdKind::Local) + pos;
void PresburgerLocalSpace::removeIdRange(unsigned idStart, unsigned idLimit) {
 
   // We are going to be removing one or more identifiers from the range.
   assert(idStart < getNumIds() && "invalid idStart position");
 
   unsigned numLocalsEliminated =
       getIdKindOverlap(IdKind::Local, idStart, idLimit);
 
   // Update space parameters.
-  PresburgerSpace::removeIdRange(
-      idStart, std::min(idLimit, PresburgerSpace::getNumIds()));
+  PresburgerSpace::removeIdRange(idStart, idLimit);
 
   // Update local ids.
   numLocals -= numLocalsEliminated;
 }
 
 void PresburgerSpace::setDimSymbolSeparation(unsigned newSymbolCount) {
   assert(newSymbolCount <= getNumDimAndSymbolIds() &&
          "invalid separation position");
-  numDims = numDims + numSymbols - newSymbolCount;
+  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";
+}
+
+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";
+}
+
+void PresburgerLocalSpace::dump() const { print(llvm::errs()); }

mlir/lib/Dialect/Affine/Analysis/AffineStructures.cpp

void FlatAffineValueConstraints::reset(unsigned newNumDims,
                                        unsigned newNumLocals,
                                        ArrayRef<Value> valArgs) {
   reset(0, 0, newNumDims + newNumSymbols + newNumLocals + 1, newNumDims,
         newNumSymbols, newNumLocals, valArgs);
 }
 
 unsigned FlatAffineValueConstraints::appendDimId(ValueRange vals) {
   unsigned pos = getNumDimIds();
-  insertId(IdKind::Dimension, pos, vals);
+  insertId(IdKind::SetDim, pos, vals);
   return pos;
 }
 
 unsigned FlatAffineValueConstraints::appendSymbolId(ValueRange vals) {
   unsigned pos = getNumSymbolIds();
   insertId(IdKind::Symbol, pos, vals);
   return pos;
 }
 
 unsigned FlatAffineValueConstraints::insertDimId(unsigned pos,
                                                  ValueRange vals) {
-  return insertId(IdKind::Dimension, pos, vals);
+  return insertId(IdKind::SetDim, pos, vals);
 }
 
 unsigned FlatAffineValueConstraints::insertSymbolId(unsigned pos,
                                                     ValueRange vals) {
   return insertId(IdKind::Symbol, pos, vals);
 }
 
 unsigned FlatAffineValueConstraints::insertId(IdKind kind, unsigned pos,
areIdsUnique(const FlatAffineConstraints &cst) {
   return areIdsUnique(cst, 0, cst.getNumIds());
 }
 
 /// Checks if the SSA values associated with `cst`'s identifiers of kind `kind`
 /// are unique.
 static bool LLVM_ATTRIBUTE_UNUSED areIdsUnique(
     const FlatAffineValueConstraints &cst, FlatAffineConstraints::IdKind kind) {
 
-  if (kind == FlatAffineConstraints::IdKind::Dimension)
+  if (kind == FlatAffineConstraints::IdKind::SetDim)
     return areIdsUnique(cst, 0, cst.getNumDimIds());
   if (kind == FlatAffineConstraints::IdKind::Symbol)
     return areIdsUnique(cst, cst.getNumDimIds(), cst.getNumDimAndSymbolIds());
   if (kind == FlatAffineConstraints::IdKind::Local)
     return areIdsUnique(cst, cst.getNumDimAndSymbolIds(), cst.getNumIds());
   llvm_unreachable("Unexpected IdKind");
 }
 
#ifndef NDEBUG
   SmallVector<Value> newSyms;
   SmallVector<Value> *newSymsPtr = &newSyms;
 #else
   SmallVector<Value> *newSymsPtr = nullptr;
 #endif // NDEBUG
 
   dims.reserve(getNumDimIds());
   syms.reserve(getNumSymbolIds());
-  for (unsigned i = getIdKindOffset(IdKind::Dimension),
-                e = getIdKindEnd(IdKind::Dimension);
+  for (unsigned i = getIdKindOffset(IdKind::SetDim),
+                e = getIdKindEnd(IdKind::SetDim);
        i < e; ++i)
     dims.push_back(values[i] ? *values[i] : Value());
   for (unsigned i = getIdKindOffset(IdKind::Symbol),
                 e = getIdKindEnd(IdKind::Symbol);
        i < e; ++i)
     syms.push_back(values[i] ? *values[i] : Value());
 
   AffineMap alignedMap =

mlir/unittests/Analysis/Presburger/CMakeLists.txt

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

mlir/unittests/Analysis/Presburger/IntegerPolyhedronTest.cpp

 TEST(IntegerPolyhedronTest, removeIdRange) {
   IntegerPolyhedron set(3, 2, 1);
 
   set.addInequality({10, 11, 12, 20, 21, 30, 40});
   set.removeId(IntegerPolyhedron::IdKind::Symbol, 1);
   EXPECT_THAT(set.getInequality(0),
               testing::ElementsAre(10, 11, 12, 20, 30, 40));
 
-  set.removeIdRange(IntegerPolyhedron::IdKind::Dimension, 0, 2);
+  set.removeIdRange(IntegerPolyhedron::IdKind::SetDim, 0, 2);
   EXPECT_THAT(set.getInequality(0), testing::ElementsAre(12, 20, 30, 40));
 
   set.removeIdRange(IntegerPolyhedron::IdKind::Local, 1, 1);
   EXPECT_THAT(set.getInequality(0), testing::ElementsAre(12, 20, 30, 40));
 
   set.removeIdRange(IntegerPolyhedron::IdKind::Local, 0, 1);
   EXPECT_THAT(set.getInequality(0), testing::ElementsAre(12, 20, 40));
 }

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 IdKind = PresburgerSpace::IdKind;
+
+TEST(PresburgerSpaceTest, insertId) {
+  PresburgerSpace space = PresburgerSpace::getRelationSpace(2, 2, 1);
+
+  // Try inserting 2 domain ids.
+  space.insertId(PresburgerSpace::IdKind::Domain, 0, 2);
+  EXPECT_EQ(space.getNumDomainIds(), 4u);
+
+  // Try inserting 1 range ids.
+  space.insertId(PresburgerSpace::IdKind::Range, 0, 1);
+  EXPECT_EQ(space.getNumRangeIds(), 3u);
+}
+
+TEST(PresburgerSpaceTest, insertIdSet) {
+  PresburgerSpace space = PresburgerSpace::getSetSpace(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(PresburgerSpace::IdKind::SetDim, 0, 2);
+  EXPECT_EQ(space.getNumRangeIds(), 4u);
+}
+
+TEST(PresburgerSpaceTest, removeIdRange) {
+  PresburgerSpace space = PresburgerSpace::getRelationSpace(2, 1, 3);
+
+  // Remove 1 domain identifier.
+  space.removeIdRange(0, 1);
+  EXPECT_EQ(space.getNumDomainIds(), 1u);
+
+  // Remove 1 symbol and 1 range identifier.
+  space.removeIdRange(1, 3);
+  EXPECT_EQ(space.getNumDomainIds(), 1u);
+  EXPECT_EQ(space.getNumRangeIds(), 0u);
+  EXPECT_EQ(space.getNumSymbolIds(), 2u);
+}