[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?

57–58

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

mlir/lib/Analysis/Presburger/PresburgerSpace.cpp

26

exist

40

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

68

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

68

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

68

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

68

Ok makes sense. WIll add the enum.

[arjunp]

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

42

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?

124

These enums are usually named like SpaceKind I think

mlir/lib/Analysis/Presburger/PresburgerSpace.cpp

68

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

42

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

[arjunp]

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

46

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

111–115

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

129

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

146

special case for SetDim spaceKind?

156

"Locals: " missing?

[Groverkss]

• Address Arjun's comments

[arjunp]

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

129

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

129

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

129

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"
+
+using namespace llvm;
 
 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 no distinction between domain and
+/// range identifiers are implemented as Spaces with only range and zero domain
+/// and should use Dimension to refer to dimensions, while spaces with a
+/// distinction between domain and range should use Domain, Range to distinguish
+/// between different types of dimensions and Dimension to collectively refer to
+/// both Domain and Range.
+///
+/// PresburgerSpace does not support identifiers of Local type.
 class PresburgerSpace {
   friend PresburgerLocalSpace;
 
 public:
-  /// Kind of identifier (column).
-  enum IdKind { Dimension, Symbol, Local };
+  /// Kind of identifier.
+  enum IdKind { Dimension, Symbol, Local, Domain, Range };
+
+  PresburgerSpace(unsigned numDomain, unsigned numRange, unsigned numSymbols)
+      : PresburgerSpace(numDomain, numRange, numSymbols, /*numLocals=*/0) {}
 
   PresburgerSpace(unsigned numDims, unsigned numSymbols)
-      : numDims(numDims), numSymbols(numSymbols), numLocals(0) {}
+      : PresburgerSpace(/*numDomain=*/0, numDims, numSymbols, /*numLocals=*/0) {
+  }
 
   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 getNumDimIds() const { return numDomain + numRange; }
   unsigned getNumSymbolIds() const { return numSymbols; }
-  unsigned getNumDimAndSymbolIds() const { return numDims + numSymbols; }
+  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(raw_ostream &os) const;
+  void dump() const;
+
 private:
-  PresburgerSpace(unsigned numDims, unsigned numSymbols, unsigned numLocals)
-      : numDims(numDims), numSymbols(numSymbols), numLocals(numLocals) {}
+  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 real dimensions.
-  unsigned numDims;
+  // 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 numDomain, unsigned numRange,
+                       unsigned numSymbols, unsigned numLocals)
+      : PresburgerSpace(numDomain, numRange, numSymbols, numLocals) {}
+
   PresburgerLocalSpace(unsigned numDims, unsigned numSymbols,
                        unsigned numLocals)
-      : PresburgerSpace(numDims, numSymbols, numLocals) {}
+      : PresburgerLocalSpace(/*numDomain=*/0, numDims, 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(unsigned idStart, unsigned idLimit) override;
+
+  void print(raw_ostream &os) const;
+  void dump() const;
 };
 
 } // namespace mlir
 
 #endif // MLIR_ANALYSIS_PRESBURGER_PRESBURGERSPACE_H

mlir/lib/Analysis/Presburger/IntegerPolyhedron.cpp

 
 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);
   return pos;
 }
 

mlir/lib/Analysis/Presburger/PresburgerSpace.cpp

 #include <algorithm>
 #include <cassert>
 
 using namespace mlir;
 
 unsigned PresburgerSpace::getNumIdKind(IdKind kind) const {
   if (kind == IdKind::Dimension)
     return getNumDimIds();
+  if (kind == IdKind::Domain)
+    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)
     return 0;
+  if (kind == IdKind::Domain)
+    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)
-    numSymbols += num;
+  if (kind == IdKind::Dimension) {
+    // If the identifier lies on the boundary of domain and range, treat it
+    // as a range identifier.
+    if (pos < getNumDomainIds())
+      numDomain += num;
     else
+      numRange += num;
+  } else if (kind == IdKind::Range) {
+    numRange += num;
+  } else if (kind == IdKind::Domain) {
+    numDomain += num;
+  } else if (kind == IdKind::Symbol) {
+    numSymbols += num;
+  } else {
     llvm_unreachable(
         "PresburgerSpace only supports Dimensions and Symbol 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 =
+      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(raw_ostream &os) const {
+  os << "Domain: " << getNumDomainIds() << ", "
+     << "Range: " << getNumRangeIds() << ", "
+     << "Symbols: " << getNumSymbolIds() << "\n";
+}
+
+void PresburgerSpace::dump() const { print(llvm::errs()); }
+
+void PresburgerLocalSpace::print(raw_ostream &os) const {
+  os << "Domain: " << getNumDomainIds() << ", "
+     << "Range: " << getNumRangeIds() << ", "
+     << "Symbols: " << getNumSymbolIds() << ", " << getNumLocalIds() << "\n";
+}
+
+void PresburgerLocalSpace::dump() const { print(llvm::errs()); }

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/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(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(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::Dimension, 0, 2);
+  EXPECT_EQ(space.getNumRangeIds(), 4u);
+}
+
+TEST(PresburgerSpaceTest, removeIdRange) {
+  PresburgerSpace space(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);
+}