[MLIR][Presburger] Support lexicographic max/min union of two PWMAFunction

[Groverkss]

Summary

This patch implements a lexicographic max/min union of two PWMAFunctions.

The lexmax/lexmin union of two functions is defined as a function defined on
the union of the input domains of both functions, such that when only one of the
functions are defined, it outputs the same as that function, and if both are
defined, it outputs the lexmax/lexmin of the two outputs. On points where
neither function is defined, the union is not defined either.

[arjunp]

mlir/include/mlir/Analysis/Presburger/PWMAFunction.h

58–59

I would call these getOutputMatrix and getOutputExpr

174–177

I prefer to explain it as:

Return a function defined on the union of the domains of this and func,
such that when only one of the functions is defined, it outputs the same as that function,
and if both are defined, it outputs the lexmax/lexmin of the two outputs.
On points where neither function is defined, the returned function is not defined either.

190–202

I prefer explaining it like this:

Return a function defined on the union of the domains of this and func,
such that when only one of the functions is defined, it outputs the same as that function,
and if neither is defined, the returned function is not defined either.

The provided `tiebreak` function determines which of the two functions' output should be used
on inputs where both the functions are defined. More precisely, given two `MultiAffineFunction`s
`mafA` and `mafB`, `tiebreak` returns the subset of the intersection of the two functions' domains
where the output of `mafA` should be used.

Also, what about just calling the function argument tiebreak?

205–206

Please use mafA and mafB

mlir/lib/Analysis/Presburger/PWMAFunction.cpp

237

Nits: Please use funcA and funcB below.

238–239

240

Please write a high-level description of the algorithm used here, and explain why all the output domains will be disjoint

241

nit: I prefer = here

243–248

These comments seem redundant, as they just repeat the same thing that's written in the code?

256–259

265–267

Doc?

272–273

Why is this a reference?

276

278–283

Why are you recomputing this every time in the loop? Just maintain the set incrementally, add and remove the last inequality every time

285

Please describe what inequality is being added in the doc.

286–288

I'm confused. For lexmin this should be maf1 < maf2 right? You seem to have written maf1 > maf2? Please write doc to clarify.

[Groverkss]

mlir/lib/Analysis/Presburger/PWMAFunction.cpp

241

This will not work since the copy constructor is explcit.

[Groverkss]

Address Arjun's comments

[arjunp]

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

587–590

Locals need not be divs anymroe

mlir/lib/Analysis/Presburger/PWMAFunction.cpp

238–242

252

259–260

It has to be disjoint from the other pieces already added. I don't understand what you've written here?

AFAIU, you need to write that all pieces that have been added so far are either

a) subsets of the domains of other MAFs in this, which are guaranteed to be disjoint from dom, or
b) they are one of the pieces added for funcA, and we have been subtracting all such pieces from dom so dom is disjoint from those pieces as well.

273

274

282

assert that the spaces are compatible, have same number of outputs

294

303

same as above

314

This still seems incorrect. the complement of outA > outB is outA <= outB. We want outA < outB. Please also add regression tests for these bug fixes

323

[Groverkss]

Address Arjun's comments

[arjunp]

mlir/lib/Analysis/Presburger/PWMAFunction.cpp

314

Again, this doesn't look correct.

ineq is outA - outB - 1 >= 0.
We want outB - outA - 1 >= 0.
You seem to be adding outB - outA + 1 >= 0.

Also, can you please add some regression tests?

[Groverkss]

Fix bug and add regression tests for tiebreakLex

[arjunp]

mlir/unittests/Analysis/Presburger/PWMAFunctionTest.cpp

208

func3 is redundant?

314

func3 is redundant

377

Please write short explanations of what's going on in the more complex test cases

[bondhugula]

Given a piecewise multi affine function func, the lexicographic
maximum/minimum union of func and this is defined as the maximum/minimum
output at the points where both functions are defined and the output of the
defined function where one of the functions is defined but the
other is not.

This is really convoluted. Can you please rewrite clearly referring to the right functions?

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

802

Missing doc comment.

mlir/include/mlir/Analysis/Presburger/PWMAFunction.h

58–59

These are missing doc comments as well -- the non-trivial ones above as well.

[Groverkss]

• Address Arjun's comments
• Address Uday's comments

[arjunp]

Can you write a little bit more detailed explanations for the complex test cases?

mlir/unittests/Analysis/Presburger/PWMAFunctionTest.cpp

464

Can you call all the func3's result?

[arjunp]

mlir/unittests/Analysis/Presburger/PWMAFunctionTest.cpp

372

explain the other two pieces in the output?

[Groverkss]

• Add more docs for tests

[Groverkss]

mlir/unittests/Analysis/Presburger/PWMAFunctionTest.cpp

372

I think they are already explained, right?

[arjunp]

mlir/unittests/Analysis/Presburger/PWMAFunctionTest.cpp

495

should be y - 1 < y ?

[Groverkss]

• Fix doc in tests

[arjunp]

mlir/unittests/Analysis/Presburger/PWMAFunctionTest.cpp

495

I think this is always true, not just when x == 0

[Groverkss]

• Fix doc

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

return PresburgerSpace::getRelationSpace(space.getNumDomainVars(),
                                              space.getNumRangeVars(),
                                              space.getNumSymbolVars());
   }
 
   /// Appends constraints from `other` into `this`. This is equivalent to an
   /// intersection with no simplification of any sort attempted.
   void append(const IntegerRelation &other);
 
-  /// Return the intersection of the two sets.
+  /// Return the intersection of the two relations.
   /// If there are locals, they will be merged.
   IntegerRelation intersect(IntegerRelation other) const;
 
   /// Return whether `this` and `other` are equal. This is integer-exact
   /// and somewhat expensive, since it uses the integer emptiness check
   /// (see IntegerRelation::findIntegerSample()).
   bool isEqual(const IntegerRelation &other) const;
 
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 relation.
   ///
   /// 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;
 
+  /// Return the set difference of this set and the given set, i.e.,
+  /// return `this \ set`.
+  PresburgerRelation subtract(const PresburgerRelation &set) 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:
   // 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;
+
+  /// Return the intersection of the two relations.
+  /// If there are locals, they will be merged.
+  IntegerPolyhedron intersect(const IntegerPolyhedron &other) const;
+
+  /// Return the set difference of this set and the given set, i.e.,
+  /// return `this \ set`.
+  PresburgerSet subtract(const PresburgerSet &other) const;
 };
 
 } // namespace presburger
 } // namespace mlir
 
 #endif // MLIR_ANALYSIS_PRESBURGER_INTEGERRELATION_H

mlir/include/mlir/Analysis/Presburger/PWMAFunction.h

public:
   MultiAffineFunction(const Matrix &output, const PresburgerSpace &space)
       : domainSet(space), output(output) {}
 
   unsigned getNumInputs() const { return domainSet.getNumDimAndSymbolVars(); }
   unsigned getNumOutputs() const { return output.getNumRows(); }
   bool isConsistent() const {
     return output.getNumColumns() == domainSet.getNumVars() + 1;
   }
-  const IntegerPolyhedron &getDomain() const { return domainSet; }
+
+  /// Get the space of the input domain of this function.
   const PresburgerSpace &getDomainSpace() const { return domainSet.getSpace(); }
+  /// Get the input domain of this function.
+  const IntegerPolyhedron &getDomain() const { return domainSet; }
+  /// Get a matrix with each row representing row^th output expression.
+  const Matrix &getOutputMatrix() const { return output; }
+  /// Get the `i^th` output expression.
+  ArrayRef<int64_t> getOutputExpr(unsigned i) const { return output.getRow(i); }
 
   /// Insert `num` variables of the specified kind at position `pos`.
   /// Positions are relative to the kind of variable. The coefficient columns
   /// corresponding to the added variables are initialized to zero. 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);
 
assert(space.getNumLocalVars() == 0 &&
            "PWMAFunction cannot have local vars.");
     assert(numOutputs >= 1 && "The function must output something!");
   }
 
   const PresburgerSpace &getSpace() const { return space; }
 
   void addPiece(const MultiAffineFunction &piece);
   void addPiece(const IntegerPolyhedron &domain, const Matrix &output);
+  void addPiece(const PresburgerSet &domain, const Matrix &output);
 
   const MultiAffineFunction &getPiece(unsigned i) const { return pieces[i]; }
   unsigned getNumPieces() const { return pieces.size(); }
   unsigned getNumOutputs() const { return numOutputs; }
   unsigned getNumInputs() const { return space.getNumVars(); }
   MultiAffineFunction &getPiece(unsigned i) { return pieces[i]; }
 
   /// Return the domain of this piece-wise MultiAffineFunction. This is the
public:
   /// value at every point in the domain.
   bool isEqual(const PWMAFunction &other) const;
 
   /// Truncate the output dimensions to the first `count` dimensions.
   ///
   /// TODO: refactor so that this can be accomplished through removeVarRange.
   void truncateOutput(unsigned count);
 
+  /// Return a function defined on the union of the domains of this and func,
+  /// such that when only one of the functions is defined, it outputs the same
+  /// as that function, and if both are defined, it outputs the lexmax/lexmin of
+  /// the two outputs. On points where neither function is defined, the returned
+  /// function is not defined either.
+  ///
+  /// Currently this does not support PWMAFunctions which have pieces containing
+  /// local variables.
+  /// TODO: Support local variables in peices.
+  PWMAFunction unionLexMin(const PWMAFunction &func);
+  PWMAFunction unionLexMax(const PWMAFunction &func);
+
   void print(raw_ostream &os) const;
   void dump() const;
 
 private:
+  /// Return a function defined on the union of the domains of `this` and
+  /// `func`, such that when only one of the functions is defined, it outputs
+  /// the same as that function, and if neither is defined, the returned
+  /// function is not defined either.
+  ///
+  /// The provided `tiebreak` function determines which of the two functions'
+  /// output should be used on inputs where both the functions are defined. More
+  /// precisely, given two `MultiAffineFunction`s `mafA` and `mafB`, `tiebreak`
+  /// returns the subset of the intersection of the two functions' domains where
+  /// the output of `mafA` should be used.
+  ///
+  /// The PresburgerSet returned by `tiebreak` should be disjoint.
+  /// TODO: Remove this constraint of returning disjoint set.
+  PWMAFunction
+  unionFunction(const PWMAFunction &func,
+                llvm::function_ref<PresburgerSet(MultiAffineFunction mafA,
+                                                 MultiAffineFunction mafB)>
+                    tiebreak) const;
+
   PresburgerSpace space;
 
   /// The list of pieces in this piece-wise MultiAffineFunction.
   SmallVector<MultiAffineFunction, 4> pieces;
 
   /// The number of output vars.
   unsigned numOutputs;
 };
 
 } // namespace presburger
 } // namespace mlir
 
 #endif // MLIR_ANALYSIS_PRESBURGER_PWMAFUNCTION_H

mlir/lib/Analysis/Presburger/IntegerRelation.cpp

SymbolicLexMin IntegerRelation::findSymbolicIntegerLexMin() const {
 
   // 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;
 }
 
+PresburgerRelation
+IntegerRelation::subtract(const PresburgerRelation &set) const {
+  return PresburgerRelation(*this).subtract(set);
+}
+
 unsigned IntegerRelation::insertVar(VarKind kind, unsigned pos, unsigned num) {
   assert(pos <= getNumVarKind(kind));
 
   unsigned insertPos = space.insertVar(kind, pos, num);
   inequalities.insertColumns(insertPos, num);
   equalities.insertColumns(insertPos, num);
   return insertPos;
 }

 void IntegerRelation::dump() const { print(llvm::errs()); }
 
 unsigned IntegerPolyhedron::insertVar(VarKind kind, unsigned pos,
                                       unsigned num) {
   assert((kind != VarKind::Domain || num == 0) &&
          "Domain has to be zero in a set");
   return IntegerRelation::insertVar(kind, pos, num);
 }
+IntegerPolyhedron
+IntegerPolyhedron::intersect(const IntegerPolyhedron &other) const {
+  return IntegerPolyhedron(IntegerRelation::intersect(other));
+}
+
+PresburgerSet IntegerPolyhedron::subtract(const PresburgerSet &other) const {
+  return PresburgerSet(IntegerRelation::subtract(other));
+}

mlir/lib/Analysis/Presburger/PWMAFunction.cpp

void PWMAFunction::addPiece(const MultiAffineFunction &piece) {
   pieces.push_back(piece);
 }
 
 void PWMAFunction::addPiece(const IntegerPolyhedron &domain,
                             const Matrix &output) {
   addPiece(MultiAffineFunction(domain, output));
 }
 
+void PWMAFunction::addPiece(const PresburgerSet &domain, const Matrix &output) {
+  for (const IntegerRelation &newDom : domain.getAllDisjuncts())
+    addPiece(IntegerPolyhedron(newDom), output);
+}
+
 void PWMAFunction::print(raw_ostream &os) const {
   os << pieces.size() << " pieces:\n";
   for (const MultiAffineFunction &piece : pieces)
     piece.print(os);
 }
 
 void PWMAFunction::dump() const { print(llvm::errs()); }
+
+PWMAFunction PWMAFunction::unionFunction(
+    const PWMAFunction &func,
+    llvm::function_ref<PresburgerSet(MultiAffineFunction maf1,
+                                     MultiAffineFunction maf2)>
+        tiebreak) const {
+  assert(getNumOutputs() == func.getNumOutputs() &&
+         "Number of outputs of functions should be same.");
+  assert(getSpace().isCompatible(func.getSpace()) &&
+         "Space is not compatible.");
+
+  // The algorithm used here is as follows:
+  // - Add the output of funcB for the part of the domain where both funcA and
+  //   funcB are defined, and `tiebreak` chooses the output of funcB.
+  // - Add the output of funcA, where funcB is not defined or `tiebreak` chooses
+  //   funcA over funcB.
+  // - Add the output of funcB, where funcA is not defined.
+
+  // Add parts of the common domain where funcB's output is used. Also
+  // add all the parts where funcA's output is used, both common and non-common.
+  PWMAFunction result(getSpace(), getNumOutputs());
+  for (const MultiAffineFunction &funcA : pieces) {
+    PresburgerSet dom(funcA.getDomain());
+    for (const MultiAffineFunction &funcB : func.pieces) {
+      PresburgerSet better = tiebreak(funcB, funcA);
+      // Add the output of funcB, where it is better than output of funcA.
+      // The disjuncts in "better" will be disjoint as tiebreak should gurantee
+      // that.
+      result.addPiece(better, funcB.getOutputMatrix());
+      dom = dom.subtract(better);
+    }
+    // Add output of funcA, where it is better than funcB, or funcB is not
+    // defined.
+    //
+    // `dom` here is guranteed to be disjoint from already added pieces
+    // because because the pieces added before are either:
+    // - Subsets of the domain of other MAFs in `this`, which are guranteed
+    //   to be disjoint from `dom`, or
+    // - They are one of the pieces added for `funcB`, and we have been
+    //   subtracting all such pieces from `dom`, so `dom` is disjoint from those
+    //   pieces as well.
+    result.addPiece(dom, funcA.getOutputMatrix());
+  }
+
+  // Add parts of funcB which are not shared with funcA.
+  PresburgerSet dom = getDomain();
+  for (const MultiAffineFunction &funcB : func.pieces)
+    result.addPiece(funcB.getDomain().subtract(dom), funcB.getOutputMatrix());
+
+  return result;
+}
+
+/// A tiebreak function which breaks ties by comparing the outputs
+/// lexicographically. If `lexMin` is true, then the ties are broken by
+/// taking the lexicographically smaller output and otherwise, by taking the
+/// lexicographically larger output.
+template <bool lexMin>
+static PresburgerSet tiebreakLex(const MultiAffineFunction &mafA,
+                                 const MultiAffineFunction &mafB) {
+  // TODO: Support local variables here.
+  assert(mafA.getDomainSpace().isCompatible(mafB.getDomainSpace()) &&
+         "Domain spaces should be compatible.");
+  assert(mafA.getNumOutputs() == mafB.getNumOutputs() &&
+         "Number of outputs of both functions should be same.");
+  assert(mafA.getDomain().getNumLocalVars() == 0 &&
+         "Local variables are not supported yet.");
+
+  PresburgerSpace compatibleSpace = mafA.getDomain().getSpaceWithoutLocals();
+  const PresburgerSpace &space = mafA.getDomain().getSpace();
+
+  // We first create the set `result`, corresponding to the set where output
+  // of mafA is lexicographically larger/smaller than mafB. This is done by
+  // creating a PresburgerSet with the following constraints:
+  //
+  //    (outA[0] > outB[0]) U
+  //    (outA[0] = outB[0], outA[1] > outA[1]) U
+  //    (outA[0] = outB[0], outA[1] = outA[1], outA[2] > outA[2]) U
+  //    ...
+  //    (outA[0] = outB[0], ..., outA[n-2] = outB[n-2], outA[n-1] > outB[n-1])
+  //
+  // where `n` is the number of outputs.
+  // If `lexMin` is set, the complement inequality is used:
+  //
+  //    (outA[0] < outB[0]) U
+  //    (outA[0] = outB[0], outA[1] < outA[1]) U
+  //    (outA[0] = outB[0], outA[1] = outA[1], outA[2] < outA[2]) U
+  //    ...
+  //    (outA[0] = outB[0], ..., outA[n-2] = outB[n-2], outA[n-1] < outB[n-1])
+  PresburgerSet result = PresburgerSet::getEmpty(compatibleSpace);
+  IntegerPolyhedron levelSet(/*numReservedInequalities=*/1,
+                             /*numReservedEqualities=*/mafA.getNumOutputs(),
+                             /*numReservedCols=*/space.getNumVars() + 1, space);
+  for (unsigned level = 0; level < mafA.getNumOutputs(); ++level) {
+
+    // Create the expression `outA - outB` for this level.
+    SmallVector<int64_t, 8> subExpr =
+        subtract(mafA.getOutputExpr(level), mafB.getOutputExpr(level));
+
+    if (lexMin) {
+      // For lexMin, we add an upper bound of -1:
+      //        outA - outB <= -1
+      //        outA <= outB - 1
+      //        outA < outB
+      levelSet.addBound(IntegerPolyhedron::BoundType::UB, subExpr, -1);
+    } else {
+      // For lexMax, we add a lower bound of 1:
+      //        outA - outB >= 1
+      //        outA > outB + 1
+      //        outA > outB
+      levelSet.addBound(IntegerPolyhedron::BoundType::LB, subExpr, 1);
+    }
+
+    // Union the set with the result.
+    result.unionInPlace(levelSet);
+    // There is only 1 inequality in `levelSet`, so the index is always 0.
+    levelSet.removeInequality(0);
+    // Add equality `outA - outB == 0` for this level for next iteration.
+    levelSet.addEquality(subExpr);
+  }
+
+  // We then intersect `result` with the domain of mafA and mafB, to only
+  // tiebreak on the domain where both are defined.
+  result = result.intersect(PresburgerSet(mafA.getDomain()))
+               .intersect(PresburgerSet(mafB.getDomain()));
+
+  return result;
+}
+
+PWMAFunction PWMAFunction::unionLexMin(const PWMAFunction &func) {
+  return unionFunction(func, tiebreakLex</*lexMin=*/true>);
+}
+
+PWMAFunction PWMAFunction::unionLexMax(const PWMAFunction &func) {
+  return unionFunction(func, tiebreakLex</*lexMin=*/false>);
+}

mlir/unittests/Analysis/Presburger/PWMAFunctionTest.cpp

TEST(PWMAFunction, eliminateRedundantLocalIdRegressionTest) {
   PWMAFunction pwmafB = parsePWMAF(
       /*numInputs=*/2, /*numOutputs=*/1,
       {
           {"(x, y) : (x - 2*(x floordiv 2) == 0, x - 2*y == 0)",
            {{1, -1, 0, 0}}} // (0, 0)
       });
   EXPECT_TRUE(pwmafA.isEqual(pwmafB));
 }
+
+TEST(PWMAFunction, unionLexMaxSimple) {
+  // func2 is better than func1, but func2's domain is empty.
+  {
+    PWMAFunction func1 = parsePWMAF(
+        /*numInputs=*/1, /*numOutputs=*/1,
+        {
+            {"(x) : ()", {{0, 1}}},
+        });
+
+    PWMAFunction func2 = parsePWMAF(
+        /*numInputs=*/1, /*numOutputs=*/1,
+        {
+            {"(x) : (1 == 0)", {{0, 2}}},
+        });
+
+    EXPECT_TRUE(func1.unionLexMax(func2).isEqual(func1));
+    EXPECT_TRUE(func2.unionLexMax(func1).isEqual(func1));
+  }
+
+  // func2 is better than func1 on a subset of func1.
+  {
+    PWMAFunction func1 = parsePWMAF(
+        /*numInputs=*/1, /*numOutputs=*/1,
+        {
+            {"(x) : ()", {{0, 1}}},
+        });
+
+    PWMAFunction func2 = parsePWMAF(
+        /*numInputs=*/1, /*numOutputs=*/1,
+        {
+            {"(x) : (x >= 0, 10 - x >= 0)", {{0, 2}}},
+        });
+
+    PWMAFunction result = parsePWMAF(
+        /*numInputs=*/1, /*numOutputs=*/1,
+        {
+            {"(x) : (-1 - x >= 0)", {{0, 1}}},
+            {"(x) : (x >= 0, 10 - x >= 0)", {{0, 2}}},
+            {"(x) : (x - 11 >= 0)", {{0, 1}}},
+        });
+
+    EXPECT_TRUE(func1.unionLexMax(func2).isEqual(result));
+    EXPECT_TRUE(func2.unionLexMax(func1).isEqual(result));
+  }
+
+  // func1 and func2 are defined over the whole domain with different outputs.
+  {
+    PWMAFunction func1 = parsePWMAF(
+        /*numInputs=*/1, /*numOutputs=*/1,
+        {
+            {"(x) : ()", {{1, 0}}},
+        });
+
+    PWMAFunction func2 = parsePWMAF(
+        /*numInputs=*/1, /*numOutputs=*/1,
+        {
+            {"(x) : ()", {{-1, 0}}},
+        });
+
+    PWMAFunction result = parsePWMAF(
+        /*numInputs=*/1, /*numOutputs=*/1,
+        {
+            {"(x) : (x >= 0)", {{1, 0}}},
+            {"(x) : (-1 - x >= 0)", {{-1, 0}}},
+        });
+
+    EXPECT_TRUE(func1.unionLexMax(func2).isEqual(result));
+    EXPECT_TRUE(func2.unionLexMax(func1).isEqual(result));
+  }
+
+  // func1 and func2 have disjoint domains.
+  {
+    PWMAFunction func1 = parsePWMAF(
+        /*numInputs=*/1, /*numOutputs=*/1,
+        {
+            {"(x) : (x >= 0, 10 - x >= 0)", {{0, 1}}},
+            {"(x) : (x - 71 >= 0, 80 - x >= 0)", {{0, 1}}},
+        });
+
+    PWMAFunction func2 = parsePWMAF(
+        /*numInputs=*/1, /*numOutputs=*/1,
+        {
+            {"(x) : (x - 20 >= 0, 41 - x >= 0)", {{0, 2}}},
+            {"(x) : (x - 101 >= 0, 120 - x >= 0)", {{0, 2}}},
+        });
+
+    PWMAFunction result = parsePWMAF(
+        /*numInputs=*/1, /*numOutputs=*/1,
+        {
+            {"(x) : (x >= 0, 10 - x >= 0)", {{0, 1}}},
+            {"(x) : (x - 71 >= 0, 80 - x >= 0)", {{0, 1}}},
+            {"(x) : (x - 20 >= 0, 41 - x >= 0)", {{0, 2}}},
+            {"(x) : (x - 101 >= 0, 120 - x >= 0)", {{0, 2}}},
+        });
+
+    EXPECT_TRUE(func1.unionLexMin(func2).isEqual(result));
+    EXPECT_TRUE(func2.unionLexMin(func1).isEqual(result));
+  }
+}
+
+TEST(PWMAFunction, unionLexMinSimple) {
+  // func2 is better than func1, but func2's domain is empty.
+  {
+    PWMAFunction func1 = parsePWMAF(
+        /*numInputs=*/1, /*numOutputs=*/1,
+        {
+            {"(x) : ()", {{0, -1}}},
+        });
+
+    PWMAFunction func2 = parsePWMAF(
+        /*numInputs=*/1, /*numOutputs=*/1,
+        {
+            {"(x) : (1 == 0)", {{0, -2}}},
+        });
+
+    EXPECT_TRUE(func1.unionLexMin(func2).isEqual(func1));
+    EXPECT_TRUE(func2.unionLexMin(func1).isEqual(func1));
+  }
+
+  // func2 is better than func1 on a subset of func1.
+  {
+    PWMAFunction func1 = parsePWMAF(
+        /*numInputs=*/1, /*numOutputs=*/1,
+        {
+            {"(x) : ()", {{0, -1}}},
+        });
+
+    PWMAFunction func2 = parsePWMAF(
+        /*numInputs=*/1, /*numOutputs=*/1,
+        {
+            {"(x) : (x >= 0, 10 - x >= 0)", {{0, -2}}},
+        });
+
+    PWMAFunction result = parsePWMAF(
+        /*numInputs=*/1, /*numOutputs=*/1,
+        {
+            {"(x) : (-1 - x >= 0)", {{0, -1}}},
+            {"(x) : (x >= 0, 10 - x >= 0)", {{0, -2}}},
+            {"(x) : (x - 11 >= 0)", {{0, -1}}},
+        });
+
+    EXPECT_TRUE(func1.unionLexMin(func2).isEqual(result));
+    EXPECT_TRUE(func2.unionLexMin(func1).isEqual(result));
+  }
+
+  // func1 and func2 are defined over the whole domain with different outputs.
+  {
+    PWMAFunction func1 = parsePWMAF(
+        /*numInputs=*/1, /*numOutputs=*/1,
+        {
+            {"(x) : ()", {{-1, 0}}},
+        });
+
+    PWMAFunction func2 = parsePWMAF(
+        /*numInputs=*/1, /*numOutputs=*/1,
+        {
+            {"(x) : ()", {{1, 0}}},
+        });
+
+    PWMAFunction result = parsePWMAF(
+        /*numInputs=*/1, /*numOutputs=*/1,
+        {
+            {"(x) : (x >= 0)", {{-1, 0}}},
+            {"(x) : (-1 - x >= 0)", {{1, 0}}},
+        });
+
+    EXPECT_TRUE(func1.unionLexMin(func2).isEqual(result));
+    EXPECT_TRUE(func2.unionLexMin(func1).isEqual(result));
+  }
+}
+
+TEST(PWMAFunction, unionLexMaxComplex) {
+  // Union of function containing 4 different pieces of output.
+  //
+  // x >= 21               --> func1 (func2 not defined)
+  // x <= 0                --> func2 (func1 not defined)
+  // 10 <= x <= 20, y >  0 --> func1 (x + y  > x - y for y >  0)
+  // 10 <= x <= 20, y <= 0 --> func2 (x + y <= x - y for y <= 0)
+  {
+    PWMAFunction func1 = parsePWMAF(
+        /*numInputs=*/2, /*numOutputs=*/1,
+        {
+            {"(x, y) : (x >= 10)", {{1, 1, 0}}},
+        });
+
+    PWMAFunction func2 = parsePWMAF(
+        /*numInputs=*/2, /*numOutputs=*/1,
+        {
+            {"(x, y) : (x <= 20)", {{1, -1, 0}}},
+        });
+
+    PWMAFunction result = parsePWMAF(/*numInputs=*/2, /*numOutputs=*/1,
+                                     {{"(x, y) : (x >= 10, x <= 20, y >= 1)",
+                                       {
+                                           {1, 1, 0},
+                                       }},
+                                      {"(x, y) : (x >= 21)",
+                                       {
+                                           {1, 1, 0},
+                                       }},
+                                      {"(x, y) : (x <= 9)",
+                                       {
+                                           {1, -1, 0},
+                                       }},
+                                      {"(x, y) : (x >= 10, x <= 20, y <= 0)",
+                                       {
+                                           {1, -1, 0},
+                                       }}});
+
+    EXPECT_TRUE(func1.unionLexMax(func2).isEqual(result));
+  }
+
+  // Functions with more than one output, with contribution from both functions.
+  //
+  // If y >= 1, func1 is better because in the first output,
+  // x + y (func1) > x (func2), when y >= 1
+  //
+  // If y == 0, the first output is same for both functions, so we look at the
+  // second output. -2x + 4 (func1) > 2x - 2 (func2) when 0 <= x <= 1, so we
+  // take func1 for this domain and func2 for the remaining.
+  {
+    PWMAFunction func1 = parsePWMAF(
+        /*numInputs=*/2, /*numOutputs=*/2,
+        {
+            {"(x, y) : (x >= 0, y >= 0)", {{1, 1, 0}, {-2, 0, 4}}},
+        });
+
+    PWMAFunction func2 = parsePWMAF(
+        /*numInputs=*/2, /*numOutputs=*/2,
+        {
+            {"(x, y) : (x >= 0, y >= 0)", {{1, 0, 0}, {2, 0, -2}}},
+        });
+
+    PWMAFunction result = parsePWMAF(/*numInputs=*/2, /*numOutputs=*/2,
+                                     {{"(x, y) : (x >= 0, y >= 1)",
+                                       {
+                                           {1, 1, 0},
+                                           {-2, 0, 4},
+                                       }},
+                                      {"(x, y) : (x >= 0, x <= 1, y == 0)",
+                                       {
+                                           {1, 1, 0},
+                                           {-2, 0, 4},
+                                       }},
+                                      {"(x, y) : (x >= 2, y == 0)",
+                                       {
+                                           {1, 0, 0},
+                                           {2, 0, -2},
+                                       }}});
+
+    EXPECT_TRUE(func1.unionLexMax(func2).isEqual(result));
+    EXPECT_TRUE(func2.unionLexMax(func1).isEqual(result));
+  }
+
+  // Function with three boolean variables `a, b, c` used to control which
+  // output will be taken lexicographically.
+  //
+  // a == 1                 --> Take func2
+  // a == 0, b == 1         --> Take func1
+  // a == 0, b == 0, c == 1 --> Take func2
+  {
+    PWMAFunction func1 = parsePWMAF(
+        /*numInputs=*/3, /*numOutputs=*/3,
+        {
+            {"(a, b, c) : (a >= 0, 1 - a >= 0, b >= 0, 1 - b >= 0, c "
+             ">= 0, 1 - c >= 0)",
+             {{0, 0, 0, 0}, {0, 1, 0, 0}, {0, 0, 0, 0}}},
+        });
+
+    PWMAFunction func2 = parsePWMAF(
+        /*numInputs=*/3, /*numOutputs=*/3,
+        {
+            {"(a, b, c) : (a >= 0, 1 - a >= 0, b >= 0, 1 - b >= 0, c >= 0, 1 - "
+             "c >= 0)",
+             {{1, 0, 0, 0}, {0, 0, 0, 0}, {0, 0, 1, 0}}},
+        });
+
+    PWMAFunction result = parsePWMAF(
+        /*numInputs=*/3, /*numOutputs=*/3,
+        {
+            {"(a, b, c) : (a - 1 == 0, b >= 0, 1 - b >= 0, c >= 0, 1 - c >= 0)",
+             {{1, 0, 0, 0}, {0, 0, 0, 0}, {0, 0, 1, 0}}},
+            {"(a, b, c) : (a == 0, b - 1 == 0, c >= 0, 1 - c >= 0)",
+             {{0, 0, 0, 0}, {0, 1, 0, 0}, {0, 0, 0, 0}}},
+            {"(a, b, c) : (a == 0, b == 0, c >= 0, 1 - c >= 0)",
+             {{1, 0, 0, 0}, {0, 0, 0, 0}, {0, 0, 1, 0}}},
+        });
+
+    EXPECT_TRUE(func1.unionLexMax(func2).isEqual(result));
+    EXPECT_TRUE(func2.unionLexMax(func1).isEqual(result));
+  }
+}
+
+TEST(PWMAFunction, unionLexMinComplex) {
+  // Regression test checking if lexicographic tiebreak produces disjoint
+  // domains.
+  //
+  // If x == 1, func1 is better since in the first output,
+  // -x (func1) is < 0 (func2) when x == 1.
+  //
+  // If x == 0, func1 and func2 both have the same first output. So we take a
+  // look at the second output. func2 is better since in the second output,
+  // x - 1 (func2) is < x (func1) when x == 0.
+  PWMAFunction func1 = parsePWMAF(
+      /*numInputs=*/2, /*numOutputs=*/2,
+      {
+          {"(x, y) : (x >= 0, x <= 1, y >= 0, y <= 1)",
+           {{-1, 0, 0}, {0, 1, 0}}},
+      });
+
+  PWMAFunction func2 = parsePWMAF(
+      /*numInputs=*/2, /*numOutputs=*/2,
+      {
+          {"(x, y) : (x >= 0, x <= 1, y >= 0, y <= 1)",
+           {{0, 0, 0}, {0, 1, -1}}},
+      });
+
+  PWMAFunction result = parsePWMAF(
+      /*numInputs=*/2, /*numOutputs=*/2,
+      {
+          {"(x, y) : (x == 1, y >= 0, y <= 1)", {{-1, 0, 0}, {0, 1, 0}}},
+          {"(x, y) : (x == 0, y >= 0, y <= 1)", {{0, 0, 0}, {0, 1, -1}}},
+      });
+
+  EXPECT_TRUE(func1.unionLexMin(func2).isEqual(result));
+  EXPECT_TRUE(func2.unionLexMin(func1).isEqual(result));
+}