move all topo sorts into the utils file

Author: Dinistro

mlir/include/mlir/Analysis/SliceAnalysis.h

@@ -223,12 +223,6 @@ SetVector<Operation *>
 getSlice(Operation *op, const BackwardSliceOptions &backwardSliceOptions = {},
          const ForwardSliceOptions &forwardSliceOptions = {});
 
-/// Multi-root DAG topological sort.
-/// Performs a topological sort of the Operation in the `toSort` SetVector.
-/// Returns a topologically sorted SetVector.
-/// Does not support multi-sets.
-SetVector<Operation *> topologicalSort(const SetVector<Operation *> &toSort);
-
 /// Utility to match a generic reduction given a list of iteration-carried
 /// arguments, `iterCarriedArgs` and the position of the potential reduction
 /// argument within the list, `redPos`. If a reduction is matched, returns the

mlir/include/mlir/Analysis/TopologicalSortUtils.h

@@ -108,6 +108,10 @@ bool computeTopologicalSorting(
 /// stable.
 SetVector<Block *> getBlocksSortedByDominance(Region &region);
 
+/// Sorts all operation in `toSort` topologically while also region semantics.
+/// Does not support multi-sets.
+SetVector<Operation *> topologicalSort(const SetVector<Operation *> &toSort);
+
 } // end namespace mlir
 
 #endif // MLIR_ANALYSIS_TOPOLOGICALSORTUTILS_H

mlir/lib/Analysis/SliceAnalysis.cpp

@@ -165,80 +165,6 @@ mlir::getSlice(Operation *op, const BackwardSliceOptions &backwardSliceOptions,
   return topologicalSort(slice);
 }
 
-/// Computes the common ancestor region of all operations in `ops`. Remembers
-/// all the traversed regions in `traversedRegions`.
-static Region *findCommonParentRegion(const SetVector<Operation *> &ops,
-                                      DenseSet<Region *> &traversedRegions) {
-  // Map to count the number of times a region was encountered.
-  llvm::DenseMap<Region *, size_t> regionCounts;
-  size_t expectedCount = ops.size();
-
-  // Walk the region tree for each operation towards the root and add to the
-  // region count.
-  Region *res = nullptr;
-  for (Operation *op : ops) {
-    Region *current = op->getParentRegion();
-    while (current) {
-      // Insert or get the count.
-      auto it = regionCounts.try_emplace(current, 0).first;
-      size_t count = ++it->getSecond();
-      if (count == expectedCount) {
-        res = current;
-        break;
-      }
-      current = current->getParentRegion();
-    }
-  }
-  auto firstRange = llvm::make_first_range(regionCounts);
-  traversedRegions.insert(firstRange.begin(), firstRange.end());
-  return res;
-}
-
-/// Topologically traverses `region` and insers all encountered operations in
-/// `toSort` into the result. Recursively traverses regions when they are
-/// present in `relevantRegions`.
-static void topoSortRegion(Region &region,
-                           const DenseSet<Region *> &relevantRegions,
-                           const SetVector<Operation *> &toSort,
-                           SetVector<Operation *> &result) {
-  SetVector<Block *> sortedBlocks = getBlocksSortedByDominance(region);
-  for (Block *block : sortedBlocks) {
-    for (Operation &op : *block) {
-      if (toSort.contains(&op))
-        result.insert(&op);
-      for (Region &subRegion : op.getRegions()) {
-        // Skip regions that do not contain operations from `toSort`.
-        if (!relevantRegions.contains(&region))
-          continue;
-        topoSortRegion(subRegion, relevantRegions, toSort, result);
-      }
-    }
-  }
-}
-
-SetVector<Operation *>
-mlir::topologicalSort(const SetVector<Operation *> &toSort) {
-  if (toSort.size() <= 1)
-    return toSort;
-
-  assert(llvm::all_of(toSort,
-                      [&](Operation *op) { return toSort.count(op) == 1; }) &&
-         "expected only unique set entries");
-
-  // First, find the root region to start the recursive traversal through the
-  // IR.
-  DenseSet<Region *> relevantRegions;
-  Region *rootRegion = findCommonParentRegion(toSort, relevantRegions);
-  assert(rootRegion && "expected all ops to have a common ancestor");
-
-  // Sort all element in `toSort` by recursively traversing the IR.
-  SetVector<Operation *> result;
-  topoSortRegion(*rootRegion, relevantRegions, toSort, result);
-  assert(result.size() == toSort.size() &&
-         "expected all operations to be present in the result");
-  return result;
-}
-
 /// Returns true if `value` (transitively) depends on iteration-carried values
 /// of the given `ancestorOp`.
 static bool dependsOnCarriedVals(Value value,

mlir/lib/Analysis/TopologicalSortUtils.cpp

@@ -167,3 +167,77 @@ SetVector<Block *> mlir::getBlocksSortedByDominance(Region &region) {
 
   return blocks;
 }
+
+/// Computes the common ancestor region of all operations in `ops`. Remembers
+/// all the traversed regions in `traversedRegions`.
+static Region *findCommonParentRegion(const SetVector<Operation *> &ops,
+                                      DenseSet<Region *> &traversedRegions) {
+  // Map to count the number of times a region was encountered.
+  llvm::DenseMap<Region *, size_t> regionCounts;
+  size_t expectedCount = ops.size();
+
+  // Walk the region tree for each operation towards the root and add to the
+  // region count.
+  Region *res = nullptr;
+  for (Operation *op : ops) {
+    Region *current = op->getParentRegion();
+    while (current) {
+      // Insert or get the count.
+      auto it = regionCounts.try_emplace(current, 0).first;
+      size_t count = ++it->getSecond();
+      if (count == expectedCount) {
+        res = current;
+        break;
+      }
+      current = current->getParentRegion();
+    }
+  }
+  auto firstRange = llvm::make_first_range(regionCounts);
+  traversedRegions.insert(firstRange.begin(), firstRange.end());
+  return res;
+}
+
+/// Topologically traverses `region` and insers all encountered operations in
+/// `toSort` into the result. Recursively traverses regions when they are
+/// present in `relevantRegions`.
+static void topoSortRegion(Region &region,
+                           const DenseSet<Region *> &relevantRegions,
+                           const SetVector<Operation *> &toSort,
+                           SetVector<Operation *> &result) {
+  SetVector<Block *> sortedBlocks = getBlocksSortedByDominance(region);
+  for (Block *block : sortedBlocks) {
+    for (Operation &op : *block) {
+      if (toSort.contains(&op))
+        result.insert(&op);
+      for (Region &subRegion : op.getRegions()) {
+        // Skip regions that do not contain operations from `toSort`.
+        if (!relevantRegions.contains(&region))
+          continue;
+        topoSortRegion(subRegion, relevantRegions, toSort, result);
+      }
+    }
+  }
+}
+
+SetVector<Operation *>
+mlir::topologicalSort(const SetVector<Operation *> &toSort) {
+  if (toSort.size() <= 1)
+    return toSort;
+
+  assert(llvm::all_of(toSort,
+                      [&](Operation *op) { return toSort.count(op) == 1; }) &&
+         "expected only unique set entries");
+
+  // First, find the root region to start the recursive traversal through the
+  // IR.
+  DenseSet<Region *> relevantRegions;
+  Region *rootRegion = findCommonParentRegion(toSort, relevantRegions);
+  assert(rootRegion && "expected all ops to have a common ancestor");
+
+  // Sort all element in `toSort` by recursively traversing the IR.
+  SetVector<Operation *> result;
+  topoSortRegion(*rootRegion, relevantRegions, toSort, result);
+  assert(result.size() == toSort.size() &&
+         "expected all operations to be present in the result");
+  return result;
+}

mlir/lib/Conversion/VectorToGPU/VectorToGPU.cpp

@@ -15,6 +15,7 @@
 #include <type_traits>
 
 #include "mlir/Analysis/SliceAnalysis.h"
+#include "mlir/Analysis/TopologicalSortUtils.h"
 #include "mlir/Dialect/Affine/IR/AffineOps.h"
 #include "mlir/Dialect/Arith/IR/Arith.h"
 #include "mlir/Dialect/GPU/IR/GPUDialect.h"

mlir/lib/Dialect/Affine/Utils/LoopFusionUtils.cpp

@@ -12,6 +12,7 @@
 
 #include "mlir/Dialect/Affine/LoopFusionUtils.h"
 #include "mlir/Analysis/SliceAnalysis.h"
+#include "mlir/Analysis/TopologicalSortUtils.h"
 #include "mlir/Dialect/Affine/Analysis/AffineAnalysis.h"
 #include "mlir/Dialect/Affine/Analysis/LoopAnalysis.h"
 #include "mlir/Dialect/Affine/Analysis/Utils.h"

mlir/lib/Transforms/SROA.cpp

@@ -9,6 +9,7 @@
 #include "mlir/Transforms/SROA.h"
 #include "mlir/Analysis/DataLayoutAnalysis.h"
 #include "mlir/Analysis/SliceAnalysis.h"
+#include "mlir/Analysis/TopologicalSortUtils.h"
 #include "mlir/Interfaces/MemorySlotInterfaces.h"
 #include "mlir/Transforms/Passes.h"
 

mlir/test/lib/Analysis/TestSlice.cpp

@@ -6,7 +6,7 @@
 //
 //===----------------------------------------------------------------------===//
 
-#include "mlir/Analysis/SliceAnalysis.h"
+#include "mlir/Analysis/TopologicalSortUtils.h"
 #include "mlir/IR/BuiltinTypes.h"
 #include "mlir/IR/SymbolTable.h"
 #include "mlir/Pass/Pass.h"