[MLIR][Presburger] Refactor duplicate division merging to Utils

[Groverkss]

Summary

This patch moves merging of duplicate divisions to presburger utility
functions. This is required to support division merging in structures other
than IntegerPolyhedron.

[arjunp]

mlir/include/mlir/Analysis/Presburger/Utils.h

60–61

I think it'll be helpful if you mention clearly that division j thus gets eliminated, and its coefficients get added to i.

61

68

Why is this return value needed?

68

I think you can specify parameter names here

mlir/lib/Analysis/Presburger/Utils.cpp

197

can mention parameter names here too

[Groverkss]

mlir/include/mlir/Analysis/Presburger/Utils.h

68

This is in case some merge does not want to merge divisions belonging to the same set, which I think we wanted to implement at some point.

[Groverkss]

• Address arjun's comments

[arjunp]

mlir/include/mlir/Analysis/Presburger/Utils.h

68

I think that can be resolved by just having a separate function to merge within a set and calling that earlier, and I prefer that solution.

[arjunp]

mlir/include/mlir/Analysis/Presburger/Utils.h

69

I think we should use llvm::function_ref here https://llvm.org/doxygen/classllvm_1_1function__ref_3_01Ret_07Params_8_8_8_08_4.html

[Groverkss]

• Use llvm::function_ref instead of std::function &

mlir/include/mlir/Analysis/Presburger/Utils.h

68

I do not completely understand the way you propose. Could you elaborate more?

I think the returning a bool way makes it more general and its a small addition only.

[arjunp]

mlir/include/mlir/Analysis/Presburger/Utils.h

68

Sorry, I see that was unclear. What I meant is just write a separate simplify function in IntegerPolyhedron to just simplify divisions within the set. If we call that initially before we run whatever algorithm is using merging, then we know that internal simplifications won't be a problem after that.

I'm fine with keeping the bool return as well, but I'm not sure this generality is needed.

mlir/include/mlir/Analysis/Presburger/Utils.h

 //===- Utils.h - General utilities for Presburger library ------*- C++ -*-===//
 //
 // 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
 //
 //===----------------------------------------------------------------------===//
 //
 // Utility functions required by the Presburger Library.
 //
 //===----------------------------------------------------------------------===//
 
 #ifndef MLIR_ANALYSIS_PRESBURGER_UTILS_H
 #define MLIR_ANALYSIS_PRESBURGER_UTILS_H
 
 #include "mlir/Support/LLVM.h"
+#include "llvm/ADT/STLExtras.h"
 
 namespace mlir {
 
 class IntegerPolyhedron;
 
 namespace presburger_utils {
 
 /// `ReprKind` enum is used to set the constraint type in `MaybeLocalRepr`.

 /// can be computed. If the representation could be computed, `dividend` and
 /// `denominator` are set. If the representation could not be computed,
 /// `llvm::None` is returned.
 MaybeLocalRepr computeSingleVarRepr(const IntegerPolyhedron &cst,
                                     ArrayRef<bool> foundRepr, unsigned pos,
                                     SmallVector<int64_t, 8> &dividend,
                                     unsigned &divisor);
 
+/// Given dividends of divisions `divs` and denominators `denoms`, detects and
+/// removes duplicate divisions. `localOffset` is the offset in dividend of a
+/// division from where local identifiers start.
+///
+/// On every possible duplicate division found, `merge(i, j)`, where `i`, `j`
+/// are current index of the duplicate divisions, is called and division at
+/// index `j` is merged into division at index `i`. If `merge(i, j)` returns
+/// `true`, the divisions are merged i.e. `j^th` division gets eliminated and
+/// it's each instance is replaced by `i^th` division. If it returns `false`,
+/// the divisions are not merged. `merge` can also do side effects, For example
+/// it can merge the local identifiers in IntegerPolyhedron.
+void removeDuplicateDivs(
+    std::vector<SmallVector<int64_t, 8>> &divs,
+    SmallVectorImpl<unsigned> &denoms, unsigned localOffset,
+    llvm::function_ref<bool(unsigned i, unsigned j)> merge);
+
 } // namespace presburger_utils
 } // namespace mlir
 
 #endif // MLIR_ANALYSIS_PRESBURGER_UTILS_H

mlir/lib/Analysis/Presburger/IntegerPolyhedron.cpp

void IntegerPolyhedron::mergeLocalIds(IntegerPolyhedron &other) {
   // these have the combined division information of both polys. Since newly
   // added local variables in polyA and polyB have no constraints, they will not
   // have any division representation.
   std::copy(divsB.begin() + initLocals, divsB.end(),
             divsA.begin() + initLocals);
   std::copy(denomsB.begin() + initLocals, denomsB.end(),
             denomsA.begin() + initLocals);
 
-  // Find and merge duplicate divisions.
-  // TODO: Add division normalization to support divisions that differ by
-  // a constant.
-  // TODO: Add division ordering such that a division representation for local
-  // identifier at position `i` only depends on local identifiers at position <
-  // `i`. This would make sure that all divisions depending on other local
-  // variables that can be merged, are merged.
-  unsigned localOffset = getIdKindOffset(IdKind::Local);
-  for (unsigned i = 0; i < divsA.size(); ++i) {
-    // Check if a division representation exists for the `i^th` local id.
-    if (denomsA[i] == 0)
-      continue;
-    // Check if a division exists which is a duplicate of the division at `i`.
-    for (unsigned j = i + 1; j < divsA.size(); ++j) {
-      // Check if a division representation exists for the `j^th` local id.
-      if (denomsA[j] == 0)
-        continue;
-      // Check if the denominators match.
-      if (denomsA[i] != denomsA[j])
-        continue;
-      // Check if the representations are equal.
-      if (divsA[i] != divsA[j])
-        continue;
-
-      // Merge divisions at position `j` into division at position `i`.
-      eliminateRedundantLocalId(polyA, i, j);
-      eliminateRedundantLocalId(polyB, i, j);
-      for (unsigned k = 0, g = divsA.size(); k < g; ++k) {
-        SmallVector<int64_t, 8> &div = divsA[k];
-        if (denomsA[k] != 0) {
-          div[localOffset + i] += div[localOffset + j];
-          div.erase(div.begin() + localOffset + j);
-        }
-      }
+  // Merge function that merges the local variables in both sets by treating
+  // them as the same identifier.
+  auto merge = [&polyA, &polyB](unsigned i, unsigned j) -> bool {
+    return true;
+  };
 
-      divsA.erase(divsA.begin() + j);
-      denomsA.erase(denomsA.begin() + j);
-      // Since `j` can never be zero, we do not need to worry about overflows.
-      --j;
-    }
-  }
+  // Merge all divisions by removing duplicate divisions.
+  unsigned localOffset = getIdKindOffset(IdKind::Local);
+  removeDuplicateDivs(divsA, denomsA, localOffset, merge);
 }
 
 /// Removes local variables using equalities. Each equality is checked if it
 /// can be reduced to the form: `e = affine-expr`, where `e` is a local
 /// variable and `affine-expr` is an affine expression not containing `e`.
 /// If an equality satisfies this form, the local variable is replaced in
 /// each constraint and then removed. The equality used to replace this local
 /// variable is also removed.

mlir/lib/Analysis/Presburger/Utils.cpp

for (unsigned lbPos : lbIndices) {
 
       repr.kind = ReprKind::Inequality;
       repr.repr.inEqualityPair = {ubPos, lbPos};
       return repr;
     }
   }
   return repr;
 }
+
+void presburger_utils::removeDuplicateDivs(
+    std::vector<SmallVector<int64_t, 8>> &divs,
+    SmallVectorImpl<unsigned> &denoms, unsigned localOffset,
+    llvm::function_ref<bool(unsigned i, unsigned j)> merge) {
+
+  // Find and merge duplicate divisions.
+  // TODO: Add division normalization to support divisions that differ by
+  // a constant.
+  // TODO: Add division ordering such that a division representation for local
+  // identifier at position `i` only depends on local identifiers at position <
+  // `i`. This would make sure that all divisions depending on other local
+  // variables that can be merged, are merged.
+  for (unsigned i = 0; i < divs.size(); ++i) {
+    // Check if a division representation exists for the `i^th` local id.
+    if (denoms[i] == 0)
+      continue;
+    // Check if a division exists which is a duplicate of the division at `i`.
+    for (unsigned j = i + 1; j < divs.size(); ++j) {
+      // Check if a division representation exists for the `j^th` local id.
+      if (denoms[j] == 0)
+        continue;
+      // Check if the denominators match.
+      if (denoms[i] != denoms[j])
+        continue;
+      // Check if the representations are equal.
+      if (divs[i] != divs[j])
+        continue;
+
+      // Merge divisions at position `j` into division at position `i`. If
+      // merge fails, do not merge these divs.
+      bool mergeResult = merge(i, j);
+      if (!mergeResult)
+        continue;
+
+      // Update division information to reflect merging.
+      for (unsigned k = 0, g = divs.size(); k < g; ++k) {
+        SmallVector<int64_t, 8> &div = divs[k];
+        if (denoms[k] != 0) {
+          div[localOffset + i] += div[localOffset + j];
+          div.erase(div.begin() + localOffset + j);
+        }
+      }
+
+      divs.erase(divs.begin() + j);
+      denoms.erase(denoms.begin() + j);
+      // Since `j` can never be zero, we do not need to worry about overflows.
+      --j;
+    }
+  }
+}