[MLIR][Presburger] Add getDomainSet, getRangeSet to IntegerRelation

[Groverkss]

Summary

This patch adds support for obtaining a set corresponding to the domain/range
of the relation.

[ftynse]

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

657

Should this be marked explicit? Given this can assert, it sounds preferable for the user to explicitly request the conversion rather than have it fire in some implicit conversion sequence.

mlir/lib/Analysis/Presburger/IntegerRelation.cpp

2089

Looks like you only need to provide IntegerPolyhedron(IntegerRelation &&) for this, so why not do it immediately rather than leaving a todo?

2090–2091

Doing return IntegerPolyhedron(copyRel); (or return IntegerPolyhedron(std::move(copyRel)); will guarantee copy/move elision in C++17. Naming the value will not, although the compiler is likely to figure it out on its own.

[Groverkss]

• Make IntegerPolyhedron constructor from IntegerRelation explicit.
• Add move constructor for IntegerPolyhedron and use it.

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

 #include "mlir/Analysis/Presburger/Matrix.h"
 #include "mlir/Analysis/Presburger/PresburgerSpace.h"
 #include "mlir/Analysis/Presburger/Utils.h"
 #include "mlir/Support/LogicalResult.h"
 
 namespace mlir {
 namespace presburger {
 
+class IntegerRelation;
+class IntegerPolyhedron;
+
 /// An IntegerRelation represents the set of points from a PresburgerSpace that
 /// satisfy a list of affine constraints. Affine constraints can be inequalities
 /// or equalities in the form:
 ///
 /// Inequality: c_0*x_0 + c_1*x_1 + .... + c_{n-1}*x_{n-1} + c_n >= 0
 /// Equality  : c_0*x_0 + c_1*x_1 + .... + c_{n-1}*x_{n-1} + c_n == 0
 ///
 /// where c_0, c_1, ..., c_n are integers and n is the total number of
public:
   /// 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) {
     space.setDimSymbolSeparation(newSymbolCount);
   }
 
+  /// Return a set corresponding to all points in the domain of the relation.
+  IntegerPolyhedron getDomainSet() const;
+
+  /// Return a set corresponding to all points in the range of the relation.
+  IntegerPolyhedron getRangeSet() const;
+
   void print(raw_ostream &os) const;
   void dump() const;
 
 protected:
   /// 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.
public:
 
   /// Constructs a relation with the specified number of dimensions and
   /// symbols.
   explicit IntegerPolyhedron(const PresburgerSpace &space)
       : IntegerPolyhedron(/*numReservedInequalities=*/0,
                           /*numReservedEqualities=*/0,
                           /*numReservedCols=*/space.getNumIds() + 1, space) {}
 
+  /// Construct a set from an IntegerRelation. The relation should have
+  /// no domain ids.
+  explicit IntegerPolyhedron(const IntegerRelation &rel)
+      : IntegerRelation(rel) {
+    assert(space.getNumDomainIds() == 0 &&
+           "Number of domain id's should be zero in Set kind space.");
+  }
+
+  /// Construct a set from an IntegerRelation, but instead of creating a copy,
+  /// use move constructor. The relation should have no domain ids.
+  explicit IntegerPolyhedron(IntegerRelation &&rel) : IntegerRelation(rel) {
+    assert(space.getNumDomainIds() == 0 &&
+           "Number of domain id's should be zero in Set kind space.");
+  }
+
   /// Return a system with no constraints, i.e., one which is satisfied by all
   /// points.
   static IntegerPolyhedron getUniverse(const PresburgerSpace &space) {
     return IntegerPolyhedron(space);
   }
 
   /// Return the kind of this IntegerRelation.
   Kind getKind() const override { return Kind::IntegerPolyhedron; }

mlir/lib/Analysis/Presburger/IntegerRelation.cpp

void IntegerRelation::removeIndependentConstraints(unsigned pos, unsigned num) {
   // TODO: This method can be made more efficient (because removal of each
   // inequality leads to much shifting/copying in the underlying buffer).
   for (auto nbIndex : llvm::reverse(nbIneqIndices))
     removeInequality(nbIndex);
   for (auto nbIndex : llvm::reverse(nbEqIndices))
     removeEquality(nbIndex);
 }
 
+IntegerPolyhedron IntegerRelation::getDomainSet() const {
+  IntegerRelation copyRel = *this;
+
+  // Convert Range variables to Local variables.
+  copyRel.convertIdKind(IdKind::Range, 0, getNumIdKind(IdKind::Range),
+                        IdKind::Local);
+
+  // Convert Domain variables to SetDim(Range) variables.
+  copyRel.convertIdKind(IdKind::Domain, 0, getNumIdKind(IdKind::Domain),
+                        IdKind::SetDim);
+
+  return IntegerPolyhedron(std::move(copyRel));
+}
+
+IntegerPolyhedron IntegerRelation::getRangeSet() const {
+  IntegerRelation copyRel = *this;
+
+  // Convert Domain variables to Local variables.
+  copyRel.convertIdKind(IdKind::Domain, 0, getNumIdKind(IdKind::Domain),
+                        IdKind::Local);
+
+  // We do not need to do anything to Range variables since they are already in
+  // SetDim position.
+
+  return IntegerPolyhedron(std::move(copyRel));
+}
+
 void IntegerRelation::printSpace(raw_ostream &os) const {
   space.print(os);
   os << getNumConstraints() << " constraints\n";
 }
 
 void IntegerRelation::print(raw_ostream &os) const {
   assert(hasConsistentState());
   printSpace(os);

mlir/unittests/Analysis/Presburger/CMakeLists.txt

 add_mlir_unittest(MLIRPresburgerTests
   IntegerPolyhedronTest.cpp
+  IntegerRelationTest.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/IntegerRelationTest.cpp

+//===- IntegerRelationTest.cpp - Tests for IntegerRelation 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/IntegerRelation.h"
+#include "./Utils.h"
+
+#include <gmock/gmock.h>
+#include <gtest/gtest.h>
+
+using namespace mlir;
+using namespace presburger;
+
+static IntegerRelation parseRelationFromSet(StringRef set, unsigned numDomain) {
+  IntegerRelation rel = parsePoly(set);
+
+  rel.convertIdKind(IdKind::SetDim, 0, numDomain, IdKind::Domain);
+
+  return rel;
+}
+
+TEST(IntegerRelationTest, getDomainAndRangeSet) {
+  IntegerRelation rel = parseRelationFromSet(
+      "(x, xr)[N] : (xr - x - 10 == 0, xr >= 0, N - xr >= 0)", 1);
+
+  IntegerPolyhedron domainSet = rel.getDomainSet();
+
+  IntegerPolyhedron expectedDomainSet =
+      parsePoly("(x)[N] : (x + 10 >= 0, N - x - 10 >= 0)");
+
+  EXPECT_TRUE(domainSet.isEqual(expectedDomainSet));
+
+  IntegerPolyhedron rangeSet = rel.getRangeSet();
+
+  IntegerPolyhedron expectedRangeSet =
+      parsePoly("(x)[N] : (x >= 0, N - x >= 0)");
+
+  EXPECT_TRUE(rangeSet.isEqual(expectedRangeSet));
+}