[MLIR][Presburger] support symbolicLexMin for IntegerRelation

[arjunp]

Summary

This also changes the space of the returned lexmin for IntegerPolyhedrons;
the symbols in the poly now correspond to symbols in the result rather than dims.

[Groverkss]

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

587

[Groverkss]

LGTM.

[arjunp]

• address kunwar's comment

[arjunp]

• address kunwar's comment
• rebase, remove the varKind flag stuff and use an isSymbol bitvector.

[arjunp]

Actually rebase properly this time

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

 
 namespace mlir {
 namespace presburger {
 
 class IntegerRelation;
 class IntegerPolyhedron;
 class PresburgerSet;
 class PresburgerRelation;
+struct SymbolicLexMin;
 
 /// 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
 ///
public:
   void applyRange(const IntegerRelation &rel);
 
   /// Compute an equivalent representation of the same set, such that all local
   /// vars in all disjuncts have division representations. This representation
   /// may involve local vars that correspond to divisions, and may also be a
   /// union of convex disjuncts.
   PresburgerRelation computeReprWithOnlyDivLocals() const;
 
+  /// Compute the symbolic integer lexmin of the polyhedron.
+  ///
+  /// This finds, for every assignment to the symbols and domain,
+  /// the lexicographically minimum value attained by the range.
+  ///
+  /// For example, the symbolic lexmin of the set
+  ///
+  /// (x, y)[a, b, c] : (a <= x, b <= x, x <= c)
+  ///
+  /// can be written as
+  ///
+  /// x = a if b <= a, a <= c
+  /// x = b if a <  b, b <= c
+  ///
+  /// This function is stored in the `lexmin` function in the result.
+  /// Some assignments to the symbols might make the set empty.
+  /// Such points are not part of the function's domain.
+  /// In the above example, this happens when max(a, b) > c.
+  ///
+  /// For some values of the symbols, the lexmin may be unbounded.
+  /// `SymbolicLexMin` stores these parts of the symbolic domain in a separate
+  /// `PresburgerSet`, `unboundedDomain`.
+  SymbolicLexMin findSymbolicIntegerLexMin() 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.
protected:
 
   /// Coefficients of affine equalities (in == 0 form).
   Matrix equalities;
 
   /// Coefficients of affine inequalities (in >= 0 form).
   Matrix inequalities;
 };
 
-struct SymbolicLexMin;
-
 /// An IntegerPolyhedron 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:
   // Clones this object.
   std::unique_ptr<IntegerPolyhedron> clone() const;
 
   /// Insert `num` variables of the specified kind at position `pos`.
   /// Positions are relative to the kind of variable. Return the absolute
   /// column position (i.e., not relative to the kind of variable) of the
   /// first added variable.
   unsigned insertVar(VarKind kind, unsigned pos, unsigned num = 1) override;
-
-  /// Compute the symbolic integer lexmin of the polyhedron.
-  /// This finds, for every assignment to the symbols, the lexicographically
-  /// minimum value attained by the dimensions. For example, the symbolic lexmin
-  /// of the set
-  ///
-  /// (x, y)[a, b, c] : (a <= x, b <= x, x <= c)
-  ///
-  /// can be written as
-  ///
-  /// x = a if b <= a, a <= c
-  /// x = b if a <  b, b <= c
-  ///
-  /// This function is stored in the `lexmin` function in the result.
-  /// Some assignments to the symbols might make the set empty.
-  /// Such points are not part of the function's domain.
-  /// In the above example, this happens when max(a, b) > c.
-  ///
-  /// For some values of the symbols, the lexmin may be unbounded.
-  /// `SymbolicLexMin` stores these parts of the symbolic domain in a separate
-  /// `PresburgerSet`, `unboundedDomain`.
-  SymbolicLexMin findSymbolicIntegerLexMin() const;
 };
 
 } // namespace presburger
 } // namespace mlir
 
 #endif // MLIR_ANALYSIS_PRESBURGER_INTEGERRELATION_H

mlir/lib/Analysis/Presburger/IntegerRelation.cpp

PresburgerRelation IntegerRelation::computeReprWithOnlyDivLocals() const {
   // The result set might lie in the wrong space -- all its ids are dims.
   // Set it to the desired space and return.
   PresburgerSpace space = getSpace();
   space.removeVarRange(VarKind::Local, 0, getNumLocalVars());
   result.setSpace(space);
   return result;
 }
 
-SymbolicLexMin IntegerPolyhedron::findSymbolicIntegerLexMin() const {
+SymbolicLexMin IntegerRelation::findSymbolicIntegerLexMin() const {
   // Compute the symbolic lexmin of the dims and locals, with the symbols being
   // the actual symbols of this set.
   SymbolicLexMin result =
-      SymbolicLexSimplex(*this, IntegerPolyhedron(PresburgerSpace::getSetSpace(
-                                    /*numDims=*/getNumSymbolVars())))
+      SymbolicLexSimplex(*this,
+                         IntegerPolyhedron(PresburgerSpace::getSetSpace(
+                             /*numDims=*/getNumDomainVars(),
+                             /*numSymbols=*/getNumSymbolVars())),
+                         varKindFlagSymbol | varKindFlagDomain)
           .computeSymbolicIntegerLexMin();
 
   // We want to return only the lexmin over the dims, so strip the locals from
   // the computed lexmin.
   result.lexmin.truncateOutput(result.lexmin.getNumOutputs() -
                                getNumLocalVars());
   return result;
 }

mlir/unittests/Analysis/Presburger/IntegerPolyhedronTest.cpp

void expectSymbolicIntegerLexMin(
     ArrayRef<StringRef> expectedUnboundedDomainRepr) {
   IntegerPolyhedron poly = parsePoly(polyStr);
 
   ASSERT_NE(poly.getNumDimVars(), 0u);
   ASSERT_NE(poly.getNumSymbolVars(), 0u);
 
   PWMAFunction expectedLexmin =
       parsePWMAF(/*numInputs=*/poly.getNumSymbolVars(),
-                 /*numOutputs=*/poly.getNumDimVars(), expectedLexminRepr);
+                 /*numOutputs=*/poly.getNumDimVars(), expectedLexminRepr,
+                 /*numSymbols=*/poly.getNumSymbolVars());
 
   PresburgerSet expectedUnboundedDomain = parsePresburgerSetFromPolyStrings(
-      poly.getNumSymbolVars(), expectedUnboundedDomainRepr);
+      /*numDims=*/0, expectedUnboundedDomainRepr, poly.getNumSymbolVars());
 
   SymbolicLexMin result = poly.findSymbolicIntegerLexMin();
 
   EXPECT_TRUE(result.lexmin.isEqual(expectedLexmin));
   if (!result.lexmin.isEqual(expectedLexmin)) {
     llvm::errs() << "got:\n";
     result.lexmin.dump();
     llvm::errs() << "expected:\n";
void expectSymbolicIntegerLexMin(
     ArrayRef<std::pair<StringRef, SmallVector<SmallVector<int64_t, 8>, 8>>>
         result) {
   expectSymbolicIntegerLexMin(polyStr, result, {});
 }
 
 TEST(IntegerPolyhedronTest, findSymbolicIntegerLexMin) {
   expectSymbolicIntegerLexMin("(x)[a] : (x - a >= 0)",
                               {
-                                  {"(a) : ()", {{1, 0}}}, // a
+                                  {"()[a] : ()", {{1, 0}}}, // a
                               });
 
   expectSymbolicIntegerLexMin(
       "(x)[a, b] : (x - a >= 0, x - b >= 0)",
       {
-          {"(a, b) : (a - b >= 0)", {{1, 0, 0}}},     // a
-          {"(a, b) : (b - a - 1 >= 0)", {{0, 1, 0}}}, // b
+          {"()[a, b] : (a - b >= 0)", {{1, 0, 0}}},     // a
+          {"()[a, b] : (b - a - 1 >= 0)", {{0, 1, 0}}}, // b
       });
 
   expectSymbolicIntegerLexMin(
       "(x)[a, b, c] : (x -a >= 0, x - b >= 0, x - c >= 0)",
       {
-          {"(a, b, c) : (a - b >= 0, a - c >= 0)", {{1, 0, 0, 0}}},         // a
-          {"(a, b, c) : (b - a - 1 >= 0, b - c >= 0)", {{0, 1, 0, 0}}},     // b
-          {"(a, b, c) : (c - a - 1 >= 0, c - b - 1 >= 0)", {{0, 0, 1, 0}}}, // c
+          {"()[a, b, c] : (a - b >= 0, a - c >= 0)", {{1, 0, 0, 0}}},     // a
+          {"()[a, b, c] : (b - a - 1 >= 0, b - c >= 0)", {{0, 1, 0, 0}}}, // b
+          {"()[a, b, c] : (c - a - 1 >= 0, c - b - 1 >= 0)",
+           {{0, 0, 1, 0}}}, // c
       });
 
   expectSymbolicIntegerLexMin("(x, y)[a] : (x - a >= 0, x + y >= 0)",
                               {
-                                  {"(a) : ()", {{1, 0}, {-1, 0}}}, // (a, -a)
+                                  {"()[a] : ()", {{1, 0}, {-1, 0}}}, // (a, -a)
                               });
 
   expectSymbolicIntegerLexMin(
       "(x, y)[a] : (x - a >= 0, x + y >= 0, y >= 0)",
       {
-          {"(a) : (a >= 0)", {{1, 0}, {0, 0}}},       // (a, 0)
-          {"(a) : (-a - 1 >= 0)", {{1, 0}, {-1, 0}}}, // (a, -a)
+          {"()[a] : (a >= 0)", {{1, 0}, {0, 0}}},       // (a, 0)
+          {"()[a] : (-a - 1 >= 0)", {{1, 0}, {-1, 0}}}, // (a, -a)
       });
 
   expectSymbolicIntegerLexMin(
       "(x, y)[a, b, c] : (x - a >= 0, y - b >= 0, c - x - y >= 0)",
       {
-          {"(a, b, c) : (c - a - b >= 0)",
+          {"()[a, b, c] : (c - a - b >= 0)",
            {{1, 0, 0, 0}, {0, 1, 0, 0}}}, // (a, b)
       });
 
   expectSymbolicIntegerLexMin(
       "(x, y, z)[a, b, c] : (c - z >= 0, b - y >= 0, x + y + z - a == 0)",
       {
-          {"(a, b, c) : ()",
+          {"()[a, b, c] : ()",
            {{1, -1, -1, 0}, {0, 1, 0, 0}, {0, 0, 1, 0}}}, // (a - b - c, b, c)
       });
 
   expectSymbolicIntegerLexMin(
       "(x)[a, b] : (a >= 0, b >= 0, x >= 0, a + b + x - 1 >= 0)",
       {
-          {"(a, b) : (a >= 0, b >= 0, a + b - 1 >= 0)", {{0, 0, 0}}}, // 0
-          {"(a, b) : (a == 0, b == 0)", {{0, 0, 1}}},                 // 1
+          {"()[a, b] : (a >= 0, b >= 0, a + b - 1 >= 0)", {{0, 0, 0}}}, // 0
+          {"()[a, b] : (a == 0, b == 0)", {{0, 0, 1}}},                 // 1
       });
 
   expectSymbolicIntegerLexMin(
       "(x)[a, b] : (1 - a >= 0, a >= 0, 1 - b >= 0, b >= 0, 1 - x >= 0, x >= "
       "0, a + b + x - 1 >= 0)",
       {
-          {"(a, b) : (1 - a >= 0, a >= 0, 1 - b >= 0, b >= 0, a + b - 1 >= 0)",
-           {{0, 0, 0}}},                              // 0
-          {"(a, b) : (a == 0, b == 0)", {{0, 0, 1}}}, // 1
+          {"()[a, b] : (1 - a >= 0, a >= 0, 1 - b >= 0, b >= 0, a + b - 1 >= "
+           "0)",
+           {{0, 0, 0}}},                                // 0
+          {"()[a, b] : (a == 0, b == 0)", {{0, 0, 1}}}, // 1
       });
 
   expectSymbolicIntegerLexMin(
       "(x, y, z)[a, b] : (x - a == 0, y - b == 0, x >= 0, y >= 0, z >= 0, x + "
       "y + z - 1 >= 0)",
       {
-          {"(a, b) : (a >= 0, b >= 0, 1 - a - b >= 0)",
+          {"()[a, b] : (a >= 0, b >= 0, 1 - a - b >= 0)",
            {{1, 0, 0}, {0, 1, 0}, {-1, -1, 1}}}, // (a, b, 1 - a - b)
-          {"(a, b) : (a >= 0, b >= 0, a + b - 2 >= 0)",
+          {"()[a, b] : (a >= 0, b >= 0, a + b - 2 >= 0)",
            {{1, 0, 0}, {0, 1, 0}, {0, 0, 0}}}, // (a, b, 0)
       });
 
   expectSymbolicIntegerLexMin("(x)[a, b] : (x - a == 0, x - b >= 0)",
                               {
-                                  {"(a, b) : (a - b >= 0)", {{1, 0, 0}}}, // a
+                                  {"()[a, b] : (a - b >= 0)", {{1, 0, 0}}}, // a
                               });
 
   expectSymbolicIntegerLexMin(
       "(q)[a] : (a - 1 - 3*q == 0, q >= 0)",
       {
-          {"(a) : (a - 1 - 3*(a floordiv 3) == 0, a >= 0)",
+          {"()[a] : (a - 1 - 3*(a floordiv 3) == 0, a >= 0)",
            {{0, 1, 0}}}, // a floordiv 3
       });
 
   expectSymbolicIntegerLexMin(
       "(r, q)[a] : (a - r - 3*q == 0, q >= 0, 1 - r >= 0, r >= 0)",
       {
-          {"(a) : (a - 0 - 3*(a floordiv 3) == 0, a >= 0)",
+          {"()[a] : (a - 0 - 3*(a floordiv 3) == 0, a >= 0)",
            {{0, 0, 0}, {0, 1, 0}}}, // (0, a floordiv 3)
-          {"(a) : (a - 1 - 3*(a floordiv 3) == 0, a >= 0)",
+          {"()[a] : (a - 1 - 3*(a floordiv 3) == 0, a >= 0)",
            {{0, 0, 1}, {0, 1, 0}}}, // (1 a floordiv 3)
       });
 
   expectSymbolicIntegerLexMin(
       "(r, q)[a] : (a - r - 3*q == 0, q >= 0, 2 - r >= 0, r - 1 >= 0)",
       {
-          {"(a) : (a - 1 - 3*(a floordiv 3) == 0, a >= 0)",
+          {"()[a] : (a - 1 - 3*(a floordiv 3) == 0, a >= 0)",
            {{0, 0, 1}, {0, 1, 0}}}, // (1, a floordiv 3)
-          {"(a) : (a - 2 - 3*(a floordiv 3) == 0, a >= 0)",
+          {"()[a] : (a - 2 - 3*(a floordiv 3) == 0, a >= 0)",
            {{0, 0, 2}, {0, 1, 0}}}, // (2, a floordiv 3)
       });
 
   expectSymbolicIntegerLexMin(
       "(r, q)[a] : (a - r - 3*q == 0, q >= 0, r >= 0)",
       {
-          {"(a) : (a - 3*(a floordiv 3) == 0, a >= 0)",
+          {"()[a] : (a - 3*(a floordiv 3) == 0, a >= 0)",
            {{0, 0, 0}, {0, 1, 0}}}, // (0, a floordiv 3)
-          {"(a) : (a - 1 - 3*(a floordiv 3) == 0, a >= 0)",
+          {"()[a] : (a - 1 - 3*(a floordiv 3) == 0, a >= 0)",
            {{0, 0, 1}, {0, 1, 0}}}, // (1, a floordiv 3)
-          {"(a) : (a - 2 - 3*(a floordiv 3) == 0, a >= 0)",
+          {"()[a] : (a - 2 - 3*(a floordiv 3) == 0, a >= 0)",
            {{0, 0, 2}, {0, 1, 0}}}, // (2, a floordiv 3)
       });
 
   expectSymbolicIntegerLexMin(
       "(x, y, z, w)[g] : ("
       // x, y, z, w are boolean variables.
       "1 - x >= 0, x >= 0, 1 - y >= 0, y >= 0,"
       "1 - z >= 0, z >= 0, 1 - w >= 0, w >= 0,"
       // We have some constraints on them:
       "x + y + z - 1 >= 0,"             // x or y or z
       "x + y + w - 1 >= 0,"             // x or y or w
       "1 - x + 1 - y + 1 - w - 1 >= 0," // ~x or ~y or ~w
       // What's the lexmin solution using exactly g true vars?
       "g - x - y - z - w == 0)",
       {
-          {"(g) : (g - 1 == 0)",
+          {"()[g] : (g - 1 == 0)",
            {{0, 0}, {0, 1}, {0, 0}, {0, 0}}}, // (0, 1, 0, 0)
-          {"(g) : (g - 2 == 0)",
+          {"()[g] : (g - 2 == 0)",
            {{0, 0}, {0, 0}, {0, 1}, {0, 1}}}, // (0, 0, 1, 1)
-          {"(g) : (g - 3 == 0)",
+          {"()[g] : (g - 3 == 0)",
            {{0, 0}, {0, 1}, {0, 1}, {0, 1}}}, // (0, 1, 1, 1)
       });
 
   // Bezout's lemma: if a, b are constants,
   // the set of values that ax + by can take is all multiples of gcd(a, b).
   expectSymbolicIntegerLexMin(
       // If (x, y) is a solution for a given [a, r], then so is (x - 5, y + 2).
       // So the lexmin is unbounded if it exists.
       "(x, y)[a, r] : (a >= 0, r - a + 14*x + 35*y == 0)", {},
       // According to Bezout's lemma, 14x + 35y can take on all multiples
       // of 7 and no other values. So the solution exists iff r - a is a
       // multiple of 7.
-      {"(a, r) : (a >= 0, r - a - 7*((r - a) floordiv 7) == 0)"});
+      {"()[a, r] : (a >= 0, r - a - 7*((r - a) floordiv 7) == 0)"});
 
   // The lexmins are unbounded.
   expectSymbolicIntegerLexMin("(x, y)[a] : (9*x - 4*y - 2*a >= 0)", {},
-                              {"(a) : ()"});
+                              {"()[a] : ()"});
 
   // Test cases adapted from isl.
   expectSymbolicIntegerLexMin(
       // a = 2b - 2(c - b), c - b >= 0.
       // So b is minimized when c = b.
       "(b, c)[a] : (a - 4*b + 2*c == 0, c - b >= 0)",
       {
-          {"(a) : (a - 2*(a floordiv 2) == 0)",
+          {"()[a] : (a - 2*(a floordiv 2) == 0)",
            {{0, 1, 0}, {0, 1, 0}}}, // (a floordiv 2, a floordiv 2)
       });
 
   expectSymbolicIntegerLexMin(
       // 0 <= b <= 255, 1 <= a - 512b <= 509,
       // b + 8 >= 1 + 16*(b + 8 floordiv 16) // i.e. b % 16 != 8
       "(b)[a] : (255 - b >= 0, b >= 0, a - 512*b - 1 >= 0, 512*b -a + 509 >= "
       "0, b + 7 - 16*((8 + b) floordiv 16) >= 0)",
       {
-          {"(a) : (255 - (a floordiv 512) >= 0, a >= 0, a - 512*(a floordiv "
+          {"()[a] : (255 - (a floordiv 512) >= 0, a >= 0, a - 512*(a floordiv "
            "512) - 1 >= 0, 512*(a floordiv 512) - a + 509 >= 0, (a floordiv "
            "512) + 7 - 16*((8 + (a floordiv 512)) floordiv 16) >= 0)",
            {{0, 1, 0, 0}}}, // (a floordiv 2, a floordiv 2)
       });
 
   expectSymbolicIntegerLexMin(
       "(a, b)[K, N, x, y] : (N - K - 2 >= 0, K + 4 - N >= 0, x - 4 >= 0, x + 6 "
       "- 2*N >= 0, K+N - x - 1 >= 0, a - N + 1 >= 0, K+N-1-a >= 0,a + 6 - b - "
       "N >= 0, 2*N - 4 - a >= 0,"
       "2*N - 3*K + a - b >= 0, 4*N - K + 1 - 3*b >= 0, b - N >= 0, a - x - 1 "
       ">= 0)",
       {{
-          "(K, N, x, y) : (x + 6 - 2*N >= 0, 2*N - 5 - x >= 0, x + 1 -3*K + N "
+          "()[K, N, x, y] : (x + 6 - 2*N >= 0, 2*N - 5 - x >= 0, x + 1 -3*K + "
+          "N "
           ">= 0, N + K - 2 - x >= 0, x - 4 >= 0)",
           {{0, 0, 1, 0, 1}, {0, 1, 0, 0, 0}} // (1 + x, N)
       }});
 }
 
 static void
 expectComputedVolumeIsValidOverapprox(const IntegerPolyhedron &poly,
                                       Optional<uint64_t> trueVolume,

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 "mlir/Analysis/Presburger/Simplex.h"
 
 #include <gmock/gmock.h>
 #include <gtest/gtest.h>
 
 using namespace mlir;
 using namespace presburger;
 
 static IntegerRelation parseRelationFromSet(StringRef set, unsigned numDomain) {
TEST(IntegerRelationTest, applyDomainAndRange) {
     IntegerRelation map3 =
         parseRelationFromSet("(x, y, a, b)[N] : (x - N == 0, y - N == 0)", 2);
 
     map1.applyDomain(map2);
 
     EXPECT_TRUE(map1.isEqual(map3));
   }
 }
+
+TEST(IntegerRelationTest, symbolicLexmin) {
+  SymbolicLexMin lexmin =
+      parseRelationFromSet("(a, x)[b] : (x - a >= 0, x - b >= 0)", 1)
+          .findSymbolicIntegerLexMin();
+
+  PWMAFunction expectedLexmin =
+      parsePWMAF(/*numInputs=*/2,
+                 /*numOutputs=*/1,
+                 {
+                     {"(a)[b] : (a - b >= 0)", {{1, 0, 0}}},     // a
+                     {"(a)[b] : (b - a - 1 >= 0)", {{0, 1, 0}}}, // b
+                 },
+                 /*numSymbols=*/1);
+  EXPECT_TRUE(lexmin.unboundedDomain.isIntegerEmpty());
+  EXPECT_TRUE(lexmin.lexmin.isEqual(expectedLexmin));
+}

mlir/unittests/Analysis/Presburger/Utils.h

 
 #ifndef MLIR_UNITTESTS_ANALYSIS_PRESBURGER_UTILS_H
 #define MLIR_UNITTESTS_ANALYSIS_PRESBURGER_UTILS_H
 
 #include "../../Dialect/Affine/Analysis/AffineStructuresParser.h"
 #include "mlir/Analysis/Presburger/IntegerRelation.h"
 #include "mlir/Analysis/Presburger/PWMAFunction.h"
 #include "mlir/Analysis/Presburger/PresburgerRelation.h"
+#include "mlir/Analysis/Presburger/Simplex.h"
 #include "mlir/IR/MLIRContext.h"
 #include "mlir/Support/LLVM.h"
 
 #include <gtest/gtest.h>
 
 namespace mlir {
 namespace presburger {
 
 /// Parses a IntegerPolyhedron from a StringRef. It is expected that the
 /// string represents a valid IntegerSet, otherwise it will violate a gtest
 /// assertion.
 inline IntegerPolyhedron parsePoly(StringRef str) {
   MLIRContext context(MLIRContext::Threading::DISABLED);
   FailureOr<IntegerPolyhedron> poly = parseIntegerSetToFAC(str, &context);
   EXPECT_TRUE(succeeded(poly));
   return *poly;
 }
 
 /// Parse a list of StringRefs to IntegerRelation and combine them into a
 /// PresburgerSet be using the union operation. It is expected that the strings
 /// are all valid IntegerSet representation and that all of them have the same
 /// number of dimensions as is specified by the numDims argument.
 inline PresburgerSet
-parsePresburgerSetFromPolyStrings(unsigned numDims, ArrayRef<StringRef> strs) {
-  PresburgerSet set =
-      PresburgerSet::getEmpty(PresburgerSpace::getSetSpace(numDims));
+parsePresburgerSetFromPolyStrings(unsigned numDims, ArrayRef<StringRef> strs,
+                                  unsigned numSymbols = 0) {
+  PresburgerSet set = PresburgerSet::getEmpty(
+      PresburgerSpace::getSetSpace(numDims, numSymbols));
   for (StringRef str : strs)
     set.unionInPlace(parsePoly(str));
   return set;
 }
 
 inline Matrix makeMatrix(unsigned numRow, unsigned numColumns,
                          ArrayRef<SmallVector<int64_t, 8>> matrix) {
   Matrix results(numRow, numColumns);

 /// and a 2D array that represents the output.
 inline PWMAFunction parsePWMAF(
     unsigned numInputs, unsigned numOutputs,
     ArrayRef<std::pair<StringRef, SmallVector<SmallVector<int64_t, 8>, 8>>>
         data,
     unsigned numSymbols = 0) {
   static MLIRContext context;
 
-  PWMAFunction result(
-      PresburgerSpace::getSetSpace(numInputs - numSymbols, numSymbols),
+  PWMAFunction result(PresburgerSpace::getSetSpace(
+                          /*numDims=*/numInputs - numSymbols, numSymbols),
                       numOutputs);
   for (const auto &pair : data) {
     IntegerPolyhedron domain = parsePoly(pair.first);
 
     result.addPiece(
         domain, makeMatrix(numOutputs, domain.getNumVars() + 1, pair.second));
   }
   return result;
 }