[MLIR][Presburger] Implement findSymbolicIntegerLexMin/Max for PresburgerRelation

[iambrj]

Summary

This patch implements findSymbolicIntegerLexMin/Max for PresburgerRelation

[Groverkss]

mlir/lib/Analysis/Presburger/PresburgerRelation.cpp

190–191

Can't you just do getSpace() ?

205–219

These functions are almost the same. Can you make a static function with a boolean flag and call it from both?

mlir/unittests/Analysis/Presburger/PresburgerRelationTest.cpp

194–228

I think some more tests would be nice for this. Maybe something with more dimension variables.

[iambrj]

Address comments

[Groverkss]

mlir/include/mlir/Analysis/Presburger/PresburgerRelation.h

121–123

Sorry for being unclear, I did not mean a static class function. I meant a static function as a local function inside the implementation.

mlir/unittests/Analysis/Presburger/PresburgerRelationTest.cpp

256–258

nit: Can you format these such that each new expression is on a new line?

281

Please remove debug statements.

286–297

It's better if these checks are near the definitions.

[iambrj]

Address comments

[iambrj]

Fix styling

[Groverkss]

LGTM, one small change.

mlir/lib/Analysis/Presburger/PresburgerRelation.cpp

198

Why not use result.lexopt directly instead?

[Groverkss]

mlir/lib/Analysis/Presburger/PresburgerRelation.cpp

198

Maybe you can just do SymbolicLexOpt &s = result.lexopt and PresburgerSet &unboundedDomain = result.unboundedDomain.

[iambrj]

Address comments

mlir/include/mlir/Analysis/Presburger/PresburgerRelation.h

public:
   /// Apply the domain of given relation `rel` to `this` relation.
   ///
   /// Formally, R1.applyDomain(R2) = R2.inverse().compose(R1).
   void applyDomain(const PresburgerRelation &rel);
 
   /// Same as compose, provided for uniformity with applyDomain.
   void applyRange(const PresburgerRelation &rel);
 
+  /// Compute the symbolic integer lexmin of the relation, i.e. for every
+  /// assignment of the symbols and domain the lexicographically minimum value
+  /// attained by the range.
+  SymbolicLexOpt findSymbolicIntegerLexMin() const;
+
+  /// Compute the symbolic integer lexmax of the relation, i.e. for every
+  /// assignment of the symbols and domain the lexicographically maximum value
+  /// attained by the range.
+  SymbolicLexOpt findSymbolicIntegerLexMax() const;
+
   /// Return true if the set contains the given point, and false otherwise.
   bool containsPoint(ArrayRef<MPInt> point) const;
   bool containsPoint(ArrayRef<int64_t> point) const {
     return containsPoint(getMPIntVec(point));
   }
 
   /// Return the complement of this set. All local variables in the set must
   /// correspond to floor divisions.

mlir/lib/Analysis/Presburger/PresburgerRelation.cpp

 //===- PresburgerRelation.cpp - MLIR PresburgerRelation 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/PresburgerRelation.h"
 #include "mlir/Analysis/Presburger/IntegerRelation.h"
+#include "mlir/Analysis/Presburger/PWMAFunction.h"
 #include "mlir/Analysis/Presburger/Simplex.h"
 #include "mlir/Analysis/Presburger/Utils.h"
 #include "llvm/ADT/STLExtras.h"
 #include "llvm/ADT/ScopeExit.h"
 #include "llvm/ADT/SmallBitVector.h"
 #include <optional>
 
 using namespace mlir;
void PresburgerRelation::applyDomain(const PresburgerRelation &rel) {
   compose(rel);
   inverse();
 }
 
 void PresburgerRelation::applyRange(const PresburgerRelation &rel) {
   compose(rel);
 }
 
+static SymbolicLexOpt findSymbolicIntegerLexOpt(const PresburgerRelation &rel,
+                                                bool isMin) {
+  SymbolicLexOpt result(rel.getSpace());
+  PWMAFunction &lexopt = result.lexopt;
+  PresburgerSet &unboundedDomain = result.unboundedDomain;
+  for (const IntegerRelation &cs : rel.getAllDisjuncts()) {
+    SymbolicLexOpt s(rel.getSpace());
+    if (isMin) {
+      s = cs.findSymbolicIntegerLexMin();
+      lexopt = lexopt.unionLexMin(s.lexopt);
+    } else {
+      s = cs.findSymbolicIntegerLexMax();
+      lexopt = lexopt.unionLexMax(s.lexopt);
+    }
+    unboundedDomain = unboundedDomain.intersect(s.unboundedDomain);
+  }
+  return result;
+}
+
+SymbolicLexOpt PresburgerRelation::findSymbolicIntegerLexMin() const {
+  return findSymbolicIntegerLexOpt(*this, true);
+}
+
+SymbolicLexOpt PresburgerRelation::findSymbolicIntegerLexMax() const {
+  return findSymbolicIntegerLexOpt(*this, false);
+}
+
 /// Return the coefficients of the ineq in `rel` specified by  `idx`.
 /// `idx` can refer not only to an actual inequality of `rel`, but also
 /// to either of the inequalities that make up an equality in `rel`.
 ///
 /// When 0 <= idx < rel.getNumInequalities(), this returns the coeffs of the
 /// idx-th inequality of `rel`.
 ///
 /// Otherwise, it is then considered to index into the ineqs corresponding to
 /// eqs of `rel`, and it must hold that
 ///
 /// 0 <= idx - rel.getNumInequalities() < 2*getNumEqualities().
 ///

mlir/unittests/Analysis/Presburger/PresburgerRelationTest.cpp

 //===- PresburgerRelationTest.cpp - Tests for PresburgerRelation 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/PresburgerRelation.h"
 #include "Parser.h"
+#include "mlir/Analysis/Presburger/Simplex.h"
 
 #include <gmock/gmock.h>
 #include <gtest/gtest.h>
 #include <iostream>
 
 using namespace mlir;
 using namespace presburger;
 
PresburgerRelation inverseRel = parsePresburgerRelationFromPresburgerSet(
          "(x, y, a, b)[N] : (y + a == 0, x + b == 0)",
          // (x, y) -> (x - N, y + N)
          "(x, y, a, b)[N] : (x - N - a == 0, y + N - b == 0)"},
         2);
 
     EXPECT_TRUE(rel.isEqual(inverseRel));
   }
 }
+
+TEST(IntegerRelationTest, symbolicLexOpt) {
+  PresburgerRelation rel1 = parsePresburgerRelationFromPresburgerSet(
+      {"(x, y)[N, M] : (x >= 0, y >= 0, N - 1 >= 0, M >= 0, M - 2 * N - 1>= 0, "
+       "2 * N - x >= 0, 2 * N - y >= 0)",
+       "(x, y)[N, M] : (x >= 0, y >= 0, N - 1 >= 0, M >= 0, M - 2 * N - 1>= 0, "
+       "x - N >= 0, M - x >= 0, y - 2 * N >= 0, M - y >= 0)"},
+      1);
+
+  SymbolicLexOpt lexmin1 = rel1.findSymbolicIntegerLexMin();
+
+  PWMAFunction expectedLexMin1 = parsePWMAF({
+      {"(x)[N, M] : (x >= 0, N - 1 >= 0, M >= 0, M - 2 * N - 1 >= 0, "
+       "2 * N - x >= 0)",
+       "(x)[N, M] -> (0)"},
+      {"(x)[N, M] : (x >= 0, N - 1 >= 0, M >= 0, M - 2 * N - 1 >= 0, "
+       "x - 2 * N- 1 >= 0, M - x >= 0)",
+       "(x)[N, M] -> (2 * N)"},
+  });
+
+  SymbolicLexOpt lexmax1 = rel1.findSymbolicIntegerLexMax();
+
+  PWMAFunction expectedLexMax1 = parsePWMAF({
+      {"(x)[N, M] : (x >= 0, N - 1 >= 0, M >= 0, M - 2 * N - 1 >= 0, "
+       "N - 1 - x  >= 0)",
+       "(x)[N, M] -> (2 * N)"},
+      {"(x)[N, M] : (x >= 0, N - 1 >= 0, M >= 0, M - 2 * N - 1 >= 0, "
+       "x - N >= 0, M - x >= 0)",
+       "(x)[N, M] -> (M)"},
+  });
+
+  EXPECT_TRUE(lexmin1.unboundedDomain.isIntegerEmpty());
+  EXPECT_TRUE(lexmin1.lexopt.isEqual(expectedLexMin1));
+  EXPECT_TRUE(lexmax1.unboundedDomain.isIntegerEmpty());
+  EXPECT_TRUE(lexmax1.lexopt.isEqual(expectedLexMax1));
+
+  PresburgerRelation rel2 = parsePresburgerRelationFromPresburgerSet(
+      // x or y or z
+      // lexmin = (x, 0, 1 - x)
+      // lexmax = (x, 1, 1)
+      {"(x, y, z) : (x >= 0, y >= 0, z >= 0, 1 - x >= 0, 1 - y >= 0, "
+       "1 - z >= 0, x + y + z - 1 >= 0)",
+       // (x or y) and (y or z) and (z or x)
+       // lexmin = (x, 1 - x, 1)
+       // lexmax = (x, 1, 1)
+       "(x, y, z) : (x >= 0, y >= 0, z >= 0, 1 - x >= 0, 1 - y >= 0, "
+       "1 - z >= 0, x + y - 1 >= 0, y + z - 1 >= 0, z + x - 1 >= 0)",
+       // x => (not y) or (not z)
+       // lexmin = (x, 0, 0)
+       // lexmax = (x, 1, 1 - x)
+       "(x, y, z) : (x >= 0, y >= 0, z >= 0, 1 - x >= 0, 1 - y >= 0, "
+       "1 - z >= 0, 2 - x - y - z >= 0)"},
+      1);
+
+  SymbolicLexOpt lexmin2 = rel2.findSymbolicIntegerLexMin();
+
+  PWMAFunction expectedLexMin2 =
+      parsePWMAF({{"(x) : (x >= 0, 1 - x >= 0)", "(x) -> (0, 0)"}});
+
+  SymbolicLexOpt lexmax2 = rel2.findSymbolicIntegerLexMax();
+
+  PWMAFunction expectedLexMax2 =
+      parsePWMAF({{"(x) : (x >= 0, 1 - x >= 0)", "(x) -> (1, 1)"}});
+
+  EXPECT_TRUE(lexmin2.unboundedDomain.isIntegerEmpty());
+  EXPECT_TRUE(lexmin2.lexopt.isEqual(expectedLexMin2));
+  EXPECT_TRUE(lexmax2.unboundedDomain.isIntegerEmpty());
+  EXPECT_TRUE(lexmax2.lexopt.isEqual(expectedLexMax2));
+
+  PresburgerRelation rel3 = parsePresburgerRelationFromPresburgerSet(
+      // (x => u or v or w) and (x or v) and (x or (not w))
+      // lexmin = (x, 0, 0, 1 - x)
+      // lexmax = (x, 1, 1 - x, x)
+      {"(x, u, v, w) : (x >= 0, u >= 0, v >= 0, w >= 0, 1 - x >= 0, "
+       "1 - u >= 0, 1 - v >= 0, 1 - w >= 0, -x + u + v + w >= 0, "
+       "x + v - 1 >= 0, x - w >= 0)",
+       // x => (u => (v => w)) and (x or (not v)) and (x or (not w))
+       // lexmin = (x, 0, 0, x)
+       // lexmax = (x, 1, x, x)
+       "(x, u, v, w) : (x >= 0, u >= 0, v >= 0, w >= 0, 1 - x >= 0, "
+       "1 - u >= 0, 1 - v >= 0, 1 - w >= 0, -x - u - v + w + 2 >= 0, "
+       "x - v >= 0, x - w >= 0)",
+       // (x or (u or (not v))) and ((not x) or ((not u) or w))
+       // and (x or (not v)) and (x or (not w))
+       // lexmin = (x, 0, 0, x)
+       // lexmax = (x, 1, x, x)
+       "(x, u, v, w) : (x >= 0, u >= 0, v >= 0, w >= 0, 1 - x >= 0, "
+       "1 - u >= 0, 1 - v >= 0, 1 - w >= 0, x + u - v >= 0, x - u + w >= 0, "
+       "x - v >= 0, x - w >= 0)"},
+      1);
+
+  SymbolicLexOpt lexmin3 = rel3.findSymbolicIntegerLexMin();
+
+  PWMAFunction expectedLexMin3 =
+      parsePWMAF({{"(x) : (x >= 0, 1 - x >= 0)", "(x) -> (0, 0, 0)"}});
+
+  SymbolicLexOpt lexmax3 = rel3.findSymbolicIntegerLexMax();
+
+  PWMAFunction expectedLexMax3 =
+      parsePWMAF({{"(x) : (x >= 0, 1 - x >= 0)", "(x) -> (1, 1, x)"}});
+
+  EXPECT_TRUE(lexmin3.unboundedDomain.isIntegerEmpty());
+  EXPECT_TRUE(lexmin3.lexopt.isEqual(expectedLexMin3));
+  EXPECT_TRUE(lexmax3.unboundedDomain.isIntegerEmpty());
+  EXPECT_TRUE(lexmax3.lexopt.isEqual(expectedLexMax3));
+}