[Mlir-commits] [mlir] [mlir][reducer] Separate Reduction Steps in `findOptimal` and `applyPatterns` (PR #190560)

[llvmlistbot at llvm.org]

================
@@ -125,86 +75,118 @@ static LogicalResult findOptimal(ModuleOp module, Region &region,
   // the path and apply the reducer to it.
   SmallVector<ReductionNode *> trace;
   ReductionNode *curNode = smallestNode;
-  trace.push_back(curNode);
   while (curNode != root) {
-    curNode = curNode->getParent();
     trace.push_back(curNode);
+    curNode = curNode->getParent();
   }
 
+  if (trace.empty())
+    // If trace is empty, then the smallestNode == root and therefore we were
+    // not successful in reducing the module
+    return failure();
+
   // Reduce the region through the optimal path.
   while (!trace.empty()) {
     ReductionNode *top = trace.pop_back_val();
-    applyPatterns(region, patterns, top->getStartRanges(), eraseOpNotInRange);
+    applyFn(region, top->getStartRanges());
   }
 
-  if (test.isInteresting(module).first != Tester::Interestingness::True)
+  std::pair<Tester::Interestingness, size_t> finalStatus =
+      test.isInteresting(module);
+
+  if (finalStatus.first != Tester::Interestingness::True)
     llvm::report_fatal_error("Reduced module is not interesting");
-  if (test.isInteresting(module).second != smallestNode->getSize())
+  if (finalStatus.second != smallestNode->getSize())
     llvm::report_fatal_error(
         "Reduced module doesn't have consistent size with smallestNode");
   return success();
 }
 
-/// This function attempts to erase all operations within the region currently
-/// being processed.
-static LogicalResult eraseAllOpsInRegion(ModuleOp module, Region &region,
-                                         const Tester &test) {
-  std::pair<Tester::Interestingness, size_t> initStatus =
-      test.isInteresting(module);
-
-  // While exploring the reduction tree, we always branch from an interesting
-  // node. Thus the root node must be interesting.
-  if (initStatus.first != Tester::Interestingness::True)
-    return module.emitWarning() << "uninterested module will not be reduced";
-  llvm::SpecificBumpPtrAllocator<ReductionNode> allocator;
-
-  // Setting the ranges to {{0, 0}} will result in the deletion of all ops
-  // within the region.
-  std::vector<ReductionNode::Range> ranges{{0, 0}};
-
-  // We allocate memory on the stack, and the 'allocator' is only used to
-  // construct the 'root node'. Since we won't be constructing any child nodes
-  // for emptyRegionNode, it is only used within the current scope.
-  ReductionNode emptyRegionNode(nullptr, ranges, allocator);
-  ReductionNode *root = &emptyRegionNode;
+/// We implicitly number each operation in the region and if an operation's
+/// number falls into rangeToKeep, we'll keep it.
+static void eliminateOperations(Region &region,
+                                ArrayRef<ReductionNode::Range> rangeToKeep) {
+  std::vector<Operation *> opsNotInRange;
+  size_t keepIndex = 0;
+  for (const auto &op : enumerate(region.getOps())) {
+    int index = op.index();
+    if (keepIndex < rangeToKeep.size() &&
+        index == rangeToKeep[keepIndex].second)
+      ++keepIndex;
+    if (keepIndex == rangeToKeep.size() || index < rangeToKeep[keepIndex].first)
+      opsNotInRange.push_back(&op.value());
+  }
 
-  // Create a copy of the current IR.
-  if (failed(root->initialize(module, region)))
-    llvm_unreachable("unexpected initialization failure");
+  for (Operation *op : opsNotInRange) {
+    op->dropAllUses();
+    op->erase();
+  }
+}
 
-  // Erase all operations within the corresponding region of the clone.
-  applyPatterns(root->getRegion(), {}, root->getRanges(), true);
-  root->update(test.isInteresting(root->getModule()));
-  if (root->isInteresting() == Tester::Interestingness::True) {
-    // If we can successfully remove all ops in the region, we apply the same
-    // transformation to the original IR and return success.
-    applyPatterns(region, {}, root->getRanges(), true);
-    return success();
+/// We implicitly number each operation in the region and if an operation's
+/// number falls into rangeToApply, we'll apply the given rewrite patterns on
+/// it.
+static void applyPatterns(Region &region,
+                          const FrozenRewritePatternSet &patterns,
+                          ArrayRef<ReductionNode::Range> rangeToApply) {
+  size_t rangeIndex = 0;
+  for (const auto &op : enumerate(region.getOps())) {
+    int index = op.index();
+    if (rangeIndex < rangeToApply.size() &&
+        index == rangeToApply[rangeIndex].second)
+      ++rangeIndex;
+    if (rangeIndex < rangeToApply.size() &&
+        index >= rangeToApply[rangeIndex].first)
+      // `applyOpPatternsGreedily` with folding returns whether the op is
+      // converted. Omit it because we don't have expectation this reduction
+      // will be success or not.
+      (void)applyOpPatternsGreedily(&op.value(), patterns,
+                                    GreedyRewriteConfig().setStrictness(
+                                        GreedyRewriteStrictness::ExistingOps));
   }
-  return failure();
 }
 
 template <typename IteratorType>
 static LogicalResult findOptimal(ModuleOp module, Region &region,
                                  const FrozenRewritePatternSet &patterns,
                                  const Tester &test) {
-  // We separate the reduction process into 3 steps, the first one is to erase
-  // redundant operations and the second one is to apply the reducer patterns.
 
-  // In the first phase, we attempt to erase all operations within the entire
-  // region.
-  if (succeeded(eraseAllOpsInRegion(module, region, test)))
+  // We first test the interstingness of the module passed to findOptimal.
+  std::pair<Tester::Interestingness, size_t> initStatus =
----------------
aidint wrote:

This is also not a new variable, I just moved it from [here](https://github.com/llvm/llvm-project/pull/190560/changes#diff-39aaa77be68211deee1769c8727551cd68613e7f0e4f1ce835560a3d95aad054L87) to this place, to avoid running the tester on a module multiple times.

https://github.com/llvm/llvm-project/pull/190560