[MLIR][Presburger] Optimize for union & subtract

[gilsaia]

Summary

Added a series of optimization to the Subtract & Union function of PresburgerRelation, referring to the ISL implementation.
Add isPlainEqual to Subtract & union,also some basic check to union.
Tested it on a simple Benchmark implemented by myself to see that it can speed up the Subtract operation and Union operation, also decrease the result size.

The Benchmark can be found here: benchmark

The overall results for Union & Subtract are as follows (previous benchmark has a bug,after fix that,the figure below is new)

The results for each case are as follows

[Groverkss]

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

136

I don't think this documentation is accurate. The first line is inaccurate, it should be something like

"Perform a quick equality check on this and other. The relations are equal if the check return true, but may or may not be equal if the check returns false. The equality check is ..."

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

143

Can you make the documentation of this similar to IntegerRelation::isPlainEmpty?

mlir/lib/Analysis/Presburger/IntegerRelation.cpp

83–85

You don't need to repeat documentation here. The header file already specifies this. You can omit this.

87

I think you should use space.isEqual here. Local variables also need to be the same.

95–104

You can just write "unsigned" instead of "unsigned int"

mlir/lib/Analysis/Presburger/PresburgerRelation.cpp

56–57

No need to add this line.

507–508

Sentences must end with a full stop.

509–511

You can remove braces here.

[gilsaia]

• [fix] change comment and style

[Groverkss]

Nice! LGTM.

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

140

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

144

Full stop at end of sentence.

mlir/lib/Analysis/Presburger/PresburgerRelation.cpp

532–535

No need to repeat what is already in the header. In general, we use header files for API documentation and the documentation in the cpp file explains the implementation. You can drop this.

[gilsaia]

• [fix] change comment

[gilsaia]

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

144

Sorry,forgot about that.

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

public:
   /// 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;
 
+  /// Return whether `this` and `other` are equal. The equality check is
+  /// performed in a plain manner, by comparing if all the equalities and
+  /// inequalities in `this` and `other` IntegerRelations are the same.
+  bool isPlainEqual(const IntegerRelation &other) const;
+
   /// Return whether this is a subset of the given IntegerRelation. This is
   /// integer-exact and somewhat expensive, since it uses the integer emptiness
   /// check (see IntegerRelation::findIntegerSample()).
   bool isSubsetOf(const IntegerRelation &other) const;
 
   /// Returns the value at the specified equality row and column.
   inline MPInt atEq(unsigned i, unsigned j) const { return equalities(i, j); }
   /// The same, but casts to int64_t. This is unsafe and will assert-fail if the

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

public:
 
   /// Return true if there is no disjunct, false otherwise.
   bool isPlainEmpty() const;
 
   /// Return true if the set is known to have one unconstrained disjunct, false
   /// otherwise.
   bool isPlainUniverse() const;
 
+  /// Return true if the set is obviously to be equal with the other set,
+  /// directly comparing whether each internal disjunct is the same
+  bool isPlainEqual(const PresburgerRelation &set) const;
+
   /// Return true if the set is consist of a single disjunct, without any local
   /// variables, false otherwise.
   bool isConvexNoLocals() const;
 
   /// Find an integer sample from the given set. This should not be called if
   /// any of the disjuncts in the union are unbounded.
   bool findIntegerSample(SmallVectorImpl<MPInt> &sample);
 

mlir/lib/Analysis/Presburger/IntegerRelation.cpp

IntegerRelation IntegerRelation::intersect(IntegerRelation other) const {
   return result;
 }
 
 bool IntegerRelation::isEqual(const IntegerRelation &other) const {
   assert(space.isCompatible(other.getSpace()) && "Spaces must be compatible.");
   return PresburgerRelation(*this).isEqual(PresburgerRelation(other));
 }
 
+/// The equality check is performed by comparing the number of equalities and
+/// inequalities in `this` and `other` IntegerRelations, and then comparing the
+/// coefficients of each corresponding equality and inequality.
+bool IntegerRelation::isPlainEqual(const IntegerRelation &other) const {
+  if (!space.isCompatible(other.getSpace()))
+    return false;
+  if (getNumEqualities() != other.getNumEqualities())
+    return false;
+  if (getNumInequalities() != other.getNumInequalities())
+    return false;
+
+  unsigned int cols = getNumCols();
+  for (unsigned int i = 0, eqs = getNumEqualities(); i < eqs; ++i) {
+    for (unsigned int j = 0; j < cols; ++j) {
+      if (atEq(i, j) != other.atEq(i, j))
+        return false;
+    }
+  }
+  for (unsigned int i = 0, ineqs = getNumInequalities(); i < ineqs; ++i) {
+    for (unsigned int j = 0; j < cols; ++j) {
+      if (atIneq(i, j) != other.atIneq(i, j))
+        return false;
+    }
+  }
+  return true;
+}
+
 bool IntegerRelation::isSubsetOf(const IntegerRelation &other) const {
   assert(space.isCompatible(other.getSpace()) && "Spaces must be compatible.");
   return PresburgerRelation(*this).isSubsetOf(PresburgerRelation(other));
 }
 
 MaybeOptimum<SmallVector<Fraction, 8>>
 IntegerRelation::findRationalLexMin() const {
   assert(getNumSymbolVars() == 0 && "Symbols are not supported!");

mlir/lib/Analysis/Presburger/PresburgerRelation.cpp

 void PresburgerRelation::unionInPlace(const IntegerRelation &disjunct) {
   assert(space.isCompatible(disjunct.getSpace()) && "Spaces should match");
   disjuncts.push_back(disjunct);
 }
 
 /// Mutate this set, turning it into the union of this set and the given set.
 ///
 /// This is accomplished by simply adding all the disjuncts of the given set
-/// to this set.
+/// to this set, except for cases where it's obvious that no processing is
+/// needed for.
 void PresburgerRelation::unionInPlace(const PresburgerRelation &set) {
   assert(space.isCompatible(set.getSpace()) && "Spaces should match");
+
+  if (isPlainEqual(set))
+    return;
+
+  if (isPlainEmpty()) {
+    disjuncts = set.disjuncts;
+    return;
+  }
+  if (set.isPlainEmpty())
+    return;
+
+  if (isPlainUniverse())
+    return;
+  if (set.isPlainUniverse()) {
+    disjuncts = set.disjuncts;
+    return;
+  }
+
   for (const IntegerRelation &disjunct : set.disjuncts)
     unionInPlace(disjunct);
 }
 
 /// Return the union of this set and the given set.
 PresburgerRelation
 PresburgerRelation::unionSet(const PresburgerRelation &set) const {
   assert(space.isCompatible(set.getSpace()) && "Spaces should match");

 }
 
 /// Return the result of subtract the given set from this set, i.e.,
 /// return `this \ set`.
 PresburgerRelation
 PresburgerRelation::subtract(const PresburgerRelation &set) const {
   assert(space.isCompatible(set.getSpace()) && "Spaces should match");
   PresburgerRelation result(getSpace());
+
+  // If we know that the two sets are clearly equal, we can simply return the
+  // empty set
+  if (isPlainEqual(set)) {
+    return result;
+  }
+
   // We compute (U_i t_i) \ (U_i set_i) as U_i (t_i \ V_i set_i).
   for (const IntegerRelation &disjunct : disjuncts)
     result.unionInPlace(getSetDifference(disjunct, set));
   return result;
 }
 
 /// T is a subset of S iff T \ S is empty, since if T \ S contains a
 /// point then this is a point that is contained in T but not S, and
 /// if T contains a point that is not in S, this also lies in T \ S.
 bool PresburgerRelation::isSubsetOf(const PresburgerRelation &set) const {
   return this->subtract(set).isIntegerEmpty();
 }
 
 /// Two sets are equal iff they are subsets of each other.
 bool PresburgerRelation::isEqual(const PresburgerRelation &set) const {
   assert(space.isCompatible(set.getSpace()) && "Spaces should match");
   return this->isSubsetOf(set) && set.isSubsetOf(*this);
 }
 
+/// Check if this PresburgerRelation is equal to another given
+/// PresburgerRelation. Two PresburgerRelations are considered equal if they
+/// have the same space, and each disjunct in this PresburgerRelation is equal
+/// to the corresponding disjunct in the other PresburgerRelation.
+bool PresburgerRelation::isPlainEqual(const PresburgerRelation &set) const {
+  if (!space.isCompatible(set.getSpace()))
+    return false;
+
+  if (getNumDisjuncts() != set.getNumDisjuncts())
+    return false;
+
+  // Compare each disjunct in this PresburgerRelation with the corresponding
+  // disjunct in the other PresburgerRelation.
+  for (unsigned int i = 0, n = getNumDisjuncts(); i < n; ++i) {
+    if (!getDisjunct(i).isPlainEqual(set.getDisjunct(i)))
+      return false;
+  }
+  return true;
+}
+
 /// Return true if the Presburger relation represents the universe set, false
 /// otherwise. It is a simple check that only check if the relation has at least
 /// one unconstrained disjunct, indicating the absence of constraints or
 /// conditions.
 bool PresburgerRelation::isPlainUniverse() const {
   for (auto &disjunct : getAllDisjuncts()) {
     if (disjunct.getNumConstraints() == 0)
       return true;