[mlir][cf-sink] Accept a callback for sinking operations

[Mogball]

Summary

(This was a TODO from the initial patch).

The control-flow sink utility accepts a callback that is used to sink an operation into a region.
The moveIntoRegion is called on the same operation and region that return true for shouldMoveIntoRegion.
The callback must preserve the dominance of the operation within the region. In the default control-flow
sink implementation, this is moving the operation to the start of the entry block.

[mehdi_amini]

Can you figure out a test for this?

[Mogball]

Add a test control-flow sink pass for testing the utilities

[rriddle]

The failures look real here, missing CMake update?

[Mogball]

Yes probably. Let me fix it.

[Mogball]

fix cmake

[Mogball]

rebase (to fix bazel build)

mlir/include/mlir/Transforms/ControlFlowSinkUtils.h

 ///   scf.yield %0
 /// } else {
 ///   scf.yield %arg2
 /// }
 /// ```
 ///
 /// Users must supply a callback `shouldMoveIntoRegion` that determines whether
 /// the given operation that only has users in the given operation should be
-/// moved into that region.
+/// moved into that region. If this returns true, `moveIntoRegion` is called on
+/// the same operation and region.
+///
+/// `moveIntoRegion` must move the operation into the region such that dominance
+/// of the operation is preserved; for example, by moving the operation to the
+/// start of the entry block. This ensures the preservation of SSA dominance of
+/// the operation's results.
 ///
 /// Returns the number of operations sunk.
 size_t
 controlFlowSink(ArrayRef<Region *> regions, DominanceInfo &domInfo,
-                function_ref<bool(Operation *, Region *)> shouldMoveIntoRegion);
+                function_ref<bool(Operation *, Region *)> shouldMoveIntoRegion,
+                function_ref<void(Operation *, Region *)> moveIntoRegion);
 
 /// Populates `regions` with regions of the provided region branch op that are
 /// executed at most once at that are reachable given the current operands of
 /// the op. These regions can be passed to `controlFlowSink` to perform sinking
 /// on the regions of the operation.
 void getSinglyExecutedRegionsToSink(RegionBranchOpInterface branch,
                                     SmallVectorImpl<Region *> &regions);
 
 } // namespace mlir
 
 #endif // MLIR_TRANSFORMS_CONTROLFLOWSINKUTILS_H

mlir/lib/Transforms/ControlFlowSink.cpp

 
 void ControlFlowSink::runOnOperation() {
   auto &domInfo = getAnalysis<DominanceInfo>();
   getOperation()->walk([&](RegionBranchOpInterface branch) {
     SmallVector<Region *> regionsToSink;
     // Get the regions are that known to be executed at most once.
     getSinglyExecutedRegionsToSink(branch, regionsToSink);
     // Sink side-effect free operations.
-    numSunk =
-        controlFlowSink(regionsToSink, domInfo, [](Operation *op, Region *) {
-          return isSideEffectFree(op);
+    numSunk = controlFlowSink(
+        regionsToSink, domInfo,
+        [](Operation *op, Region *) { return isSideEffectFree(op); },
+        [](Operation *op, Region *region) {
+          // Move the operation to the beginning of the region's entry block.
+          // This guarantees the preservation of SSA dominance of all of the
+          // operation's uses are in the region.
+          op->moveBefore(&region->front(), region->front().begin());
         });
   });
 }
 
 std::unique_ptr<Pass> mlir::createControlFlowSinkPass() {
   return std::make_unique<ControlFlowSink>();
 }

mlir/lib/Transforms/Utils/ControlFlowSinkUtils.cpp

 using namespace mlir;
 
 namespace {
 /// A helper struct for control-flow sinking.
 class Sinker {
 public:
   /// Create an operation sinker with given dominance info.
   Sinker(function_ref<bool(Operation *, Region *)> shouldMoveIntoRegion,
+         function_ref<void(Operation *, Region *)> moveIntoRegion,
          DominanceInfo &domInfo)
-      : shouldMoveIntoRegion(shouldMoveIntoRegion), domInfo(domInfo) {}
+      : shouldMoveIntoRegion(shouldMoveIntoRegion),
+        moveIntoRegion(moveIntoRegion), domInfo(domInfo), numSunk(0) {}
 
   /// Given a list of regions, find operations to sink and sink them. Return the
   /// number of operations sunk.
   size_t sinkRegions(ArrayRef<Region *> regions);
 
 private:
   /// Given a region and an op which dominates the region, returns true if all
   /// users of the given op are dominated by the entry block of the region, and
void tryToSinkPredecessors(Operation *user, Region *region,
                              std::vector<Operation *> &stack);
 
   /// Iterate over all the ops in a region and try to sink their predecessors.
   /// Recurse on subgraphs using a work queue.
   void sinkRegion(Region *region);
 
   /// The callback to determine whether an op should be moved in to a region.
   function_ref<bool(Operation *, Region *)> shouldMoveIntoRegion;
+  /// The calback to move an operation into the region.
+  function_ref<void(Operation *, Region *)> moveIntoRegion;
   /// Dominance info to determine op user dominance with respect to regions.
   DominanceInfo &domInfo;
   /// The number of operations sunk.
   size_t numSunk = 0;
 };
 } // end anonymous namespace
 
 bool Sinker::allUsersDominatedBy(Operation *op, Region *region) {
for (Value value : user->getOperands()) {
     Operation *op = value.getDefiningOp();
     // Ignore block arguments and ops that are already inside the region.
     if (!op || op->getParentRegion() == region)
       continue;
     LLVM_DEBUG(op->print(llvm::dbgs() << "\nTry to sink:\n"));
 
     // If the op's users are all in the region and it can be moved, then do so.
     if (allUsersDominatedBy(op, region) && shouldMoveIntoRegion(op, region)) {
-      // Move the op into the region's entry block. If the op is part of a
-      // subgraph, dependee ops would have been moved first, so inserting before
-      // the start of the block will ensure SSA dominance is preserved locally
-      // in the subgraph. Ops can only be safely moved into the entry block as
-      // the region's other blocks may for a loop.
-      op->moveBefore(&region->front(), region->front().begin());
+      moveIntoRegion(op, region);
       ++numSunk;
       // Add the op to the work queue.
       stack.push_back(op);
     }
   }
 }
 
 void Sinker::sinkRegion(Region *region) {
size_t Sinker::sinkRegions(ArrayRef<Region *> regions) {
   for (Region *region : regions)
     if (!region->empty())
       sinkRegion(region);
   return numSunk;
 }
 
 size_t mlir::controlFlowSink(
     ArrayRef<Region *> regions, DominanceInfo &domInfo,
-    function_ref<bool(Operation *, Region *)> shouldMoveIntoRegion) {
-  return Sinker(shouldMoveIntoRegion, domInfo).sinkRegions(regions);
+    function_ref<bool(Operation *, Region *)> shouldMoveIntoRegion,
+    function_ref<void(Operation *, Region *)> moveIntoRegion) {
+  return Sinker(shouldMoveIntoRegion, moveIntoRegion, domInfo)
+      .sinkRegions(regions);
 }
 
 void mlir::getSinglyExecutedRegionsToSink(RegionBranchOpInterface branch,
                                           SmallVectorImpl<Region *> &regions) {
   // Collect constant operands.
   SmallVector<Attribute> operands(branch->getNumOperands(), Attribute());
   for (auto &it : llvm::enumerate(branch->getOperands()))
     (void)matchPattern(it.value(), m_Constant(&operands[it.index()]));

mlir/test/Transforms/control-flow-sink-test.mlir

+// Invoke the test control-flow sink pass to test the utilities.
+// RUN: mlir-opt -test-control-flow-sink %s | FileCheck %s
+
+// CHECK-LABEL: func @test_sink
+func @test_sink() {
+  %0 = "test.sink_me"() : () -> i32
+  // CHECK-NEXT: test.sink_target
+  "test.sink_target"() ({
+    // CHECK-NEXT: %[[V0:.*]] = "test.sink_me"() {was_sunk = 0 : i32}
+    // CHECK-NEXT: "test.use"(%[[V0]])
+    "test.use"(%0) : (i32) -> ()
+  }) : () -> ()
+  return
+}
+
+// CHECK-LABEL: func @test_sink_first_region_only
+func @test_sink_first_region_only() {
+  %0 = "test.sink_me"() {first} : () -> i32
+  // CHECK-NEXT: %[[V1:.*]] = "test.sink_me"() {second}
+  %1 = "test.sink_me"() {second} : () -> i32
+  // CHECK-NEXT: test.sink_target
+  "test.sink_target"() ({
+    // CHECK-NEXT: %[[V0:.*]] = "test.sink_me"() {first, was_sunk = 0 : i32}
+    // CHECK-NEXT: "test.use"(%[[V0]])
+    "test.use"(%0) : (i32) -> ()
+  }, {
+    "test.use"(%1) : (i32) -> ()
+  }) : () -> ()
+  return
+}
+
+// CHECK-LABEL: func @test_sink_targeted_op_only
+func @test_sink_targeted_op_only() {
+  %0 = "test.sink_me"() : () -> i32
+  // CHECK-NEXT: %[[V1:.*]] = "test.dont_sink_me"
+  %1 = "test.dont_sink_me"() : () -> i32
+  // CHECK-NEXT: test.sink_target
+  "test.sink_target"() ({
+    // CHECK-NEXT: %[[V0:.*]] = "test.sink_me"
+    // CHECK-NEXT: "test.use"(%[[V0]], %[[V1]])
+    "test.use"(%0, %1) : (i32, i32) -> ()
+  }) : () -> ()
+  return
+}

mlir/test/Transforms/control-flow-sink.mlir

+// Test the default control-flow sink pass.
 // RUN: mlir-opt -control-flow-sink %s | FileCheck %s
 
 // Test that operations can be sunk.
 
 // CHECK-LABEL: @test_simple_sink
 // CHECK-SAME:  (%[[ARG0:.*]]: i32, %[[ARG1:.*]]: i32, %[[ARG2:.*]]: i32)
 // CHECK-NEXT: %[[V0:.*]] = arith.subi %[[ARG2]], %[[ARG1]]
 // CHECK-NEXT: %[[V1:.*]] = test.region_if %[[ARG0]]: i32 -> i32 then {

mlir/test/lib/Transforms/TestControlFlowSink.cpp

+//===- TestControlFlowSink.cpp - Test control-flow sink pass --------------===//
+//
+// 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
+//
+//===----------------------------------------------------------------------===//
+//
+// This pass tests the control-flow sink utilities by implementing an example
+// control-flow sink pass.
+//
+//===----------------------------------------------------------------------===//
+
+#include "mlir/Dialect/Func/IR/FuncOps.h"
+#include "mlir/IR/Dominance.h"
+#include "mlir/Pass/Pass.h"
+#include "mlir/Transforms/ControlFlowSinkUtils.h"
+
+using namespace mlir;
+
+namespace {
+/// An example control-flow sink pass to test the control-flow sink utilites.
+/// This pass will sink ops named `test.sink_me` and tag them with an attribute
+/// `was_sunk` into the first region of `test.sink_target` ops.
+struct TestControlFlowSinkPass
+    : public PassWrapper<TestControlFlowSinkPass, OperationPass<FuncOp>> {
+  /// Get the command-line argument of the test pass.
+  StringRef getArgument() const final { return "test-control-flow-sink"; }
+  /// Get the description of the test pass.
+  StringRef getDescription() const final {
+    return "Test control-flow sink pass";
+  }
+
+  /// Runs the pass on the function.
+  void runOnOperation() override {
+    auto &domInfo = getAnalysis<DominanceInfo>();
+    auto shouldMoveIntoRegion = [](Operation *op, Region *region) {
+      return region->getRegionNumber() == 0 &&
+             op->getName().getStringRef() == "test.sink_me";
+    };
+    auto moveIntoRegion = [](Operation *op, Region *region) {
+      Block &entry = region->front();
+      op->moveBefore(&entry, entry.begin());
+      op->setAttr("was_sunk",
+                  Builder(op).getI32IntegerAttr(region->getRegionNumber()));
+    };
+
+    getOperation()->walk([&](Operation *op) {
+      if (op->getName().getStringRef() != "test.sink_target")
+        return;
+      SmallVector<Region *> regions =
+          llvm::to_vector(RegionRange(op->getRegions()));
+      controlFlowSink(regions, domInfo, shouldMoveIntoRegion, moveIntoRegion);
+    });
+  }
+};
+} // end anonymous namespace
+
+namespace mlir {
+namespace test {
+void registerTestControlFlowSink() {
+  PassRegistration<TestControlFlowSinkPass>();
+};
+} // end namespace test
+} // end namespace mlir

mlir/tools/mlir-opt/mlir-opt.cpp

 void registerMemRefBoundCheck();
 void registerPatternsTestPass();
 void registerSimpleParametricTilingPass();
 void registerTestAffineLoopParametricTilingPass();
 void registerTestAliasAnalysisPass();
 void registerTestBuiltinAttributeInterfaces();
 void registerTestCallGraphPass();
 void registerTestConstantFold();
+void registerTestControlFlowSink();
 void registerTestGpuSerializeToCubinPass();
 void registerTestGpuSerializeToHsacoPass();
 void registerTestDataLayoutQuery();
 void registerTestDecomposeCallGraphTypes();
 void registerTestDiagnosticsPass();
 void registerTestDominancePass();
 void registerTestDynamicPipelinePass();
 void registerTestExpandTanhPass();
void registerTestPasses() {
   mlir::test::registerMemRefBoundCheck();
   mlir::test::registerPatternsTestPass();
   mlir::test::registerSimpleParametricTilingPass();
   mlir::test::registerTestAffineLoopParametricTilingPass();
   mlir::test::registerTestAliasAnalysisPass();
   mlir::test::registerTestBuiltinAttributeInterfaces();
   mlir::test::registerTestCallGraphPass();
   mlir::test::registerTestConstantFold();
+  mlir::test::registerTestControlFlowSink();
   mlir::test::registerTestDiagnosticsPass();
 #if MLIR_CUDA_CONVERSIONS_ENABLED
   mlir::test::registerTestGpuSerializeToCubinPass();
 #endif
 #if MLIR_ROCM_CONVERSIONS_ENABLED
   mlir::test::registerTestGpuSerializeToHsacoPass();
 #endif
   mlir::test::registerTestDecomposeCallGraphTypes();