[mlir][scf] Add a pass to expand scf.if regions

[antiagainst]

Summary

This pass expands the scf.if op regions by pulling
in ops before and after the scf.if op into both
regions. This is an useful transformation to put
ops within the same region so that we can optimze
together better, e.g., when we have a fast-slow
path in the same kernel we'd want to optimize
each path separately.

Depends On D117018

[mehdi_amini]

I wonder if there isn't a name in the literature for this transformation?

Also I would think that a canonicalization could undo this transformation right?

mlir/include/mlir/Dialect/SCF/Passes.td

127 ↗

(On Diff #398928)

This seems like an unnecessary restriction to me: can't you just process operation in the same block as the scf.if without further restrictions?

[antiagainst]

Rebase

[antiagainst]

In D117019#3234808, @mehdi_amini wrote:

I wonder if there isn't a name in the literature for this transformation?

Good question! I didn't find one. Was tempted to call it sinking, but that only pulls in ops before the region I think..

Also I would think that a canonicalization could undo this transformation right?

Canonicalization typically won't move ops across region boundaries. (Constant is an exception though.) But yes generally the application of this pattern should be controlled precisely w.r.t. other passes, like CSE, which may reverse some of the changes.

mlir/include/mlir/Dialect/SCF/Passes.td

127 ↗

(On Diff #398928)

Good point. Though I'm not sure how likely we are gonna see the case where we have multiple basic blocks containing scf.if ops. To me scf.if ops are structured control flow, while basic blocks are normal control flow; they typically don't operate at the same level. But can certainly drop if there is such a case and need though.

[nicolasvasilache]

mlir/include/mlir/Dialect/SCF/Passes.h

64 ↗

(On Diff #406202)

As a transformation and a pattern this makes sense to me.
But I would turn this into a test pass: this is unlikely to grow standalone heuristics and control beyond: just apply greedily.

mlir/lib/Dialect/SCF/Transforms/IfRegionExpansion.cpp

63

typo: no

155

This is unnecessary and an effect of running with a greedy rewriter.

As we move away from this towards better compositions of patterns, this goes against the flow.
How about having a test pass that just walks the scf.if ops in your IR and applies the pattern (e.g. with applyOpPatternsAndFold).

[antiagainst]

Address comments

mlir/lib/Dialect/SCF/Transforms/IfRegionExpansion.cpp

155

SG! Done.

[nicolasvasilache]

mlir/lib/Dialect/SCF/Transforms/IfRegionExpansion.cpp

62

Can you explain what DCE has to do here?
I am unclear why you want to limit prevOps to all being noEffect; copying to all branches should be ok no?

Also, why only prevOps and no nextOps ?

72

I don't get this, are you assuming that all ops that follow scf.if in the block will get pulled in and that will be the final type?

First, I am not sure why nextOps let you do that while prevOps don't.

Second, you're also doing this in the context of GPU right?
If so, what about divergent control flow?
Such ops shouldn't be pulled in IIUC?

84

What happens with all the backward slice involved in computing the condition?
Your code and tests seem to assume this is always a bbarg of an enclosing block, which sounds super restrictive.
In any case, this looks like a bug to me atm.

103

This looks fishy, should this be newYieldOp ?
Otherwise you're plugging in old values that would not exist in the cloned code.

[antiagainst]

Address comments

mlir/lib/Dialect/SCF/Transforms/IfRegionExpansion.cpp

62

The original idea of requiring no side effecting for all prevOps is that we can then free to clone them into both regions without problem. And we can leave the original prevOps just there---given they are not side effecting and after cloning and rewrite users, they will be dead code. So DCE can kick in later to remove them.

Yes indeed copying them to both regions are still fine. So I changed the impl. to not do side effecting analysis here. Instead, we just move them all ops to both regions, as long as they can be moved (according to the filter function) and does not cause invalid IR def-use relationship.

72

Yes this is assuming all ops following the scf.if can be pulled into both regions. Therefore, values escaping the new scf.if regions are can only be used by the original terminator.

The planned use for this pattern is early in the CodeGen flow with tensor level abstractions. So such assumption is fine. It help to simplify the impl. Otherwise we need to analyze the ops following the scf.if op and see what values are now generated from the new scf.if regions and also yield those. I don't have a concrete use case here; it's better to support that when needed.

Good point regarding divergent control flow. For the case at my hand it's not an issue because tensor level abstractions and value semantics. But indeed, if it's used at a lower level, it can be problematic to also pull in ops like barriers. To address this, I introduced a filter function to the pattern to check whether it's possible to pull in ops before/after into both regions. This gives pattern users control. The filter function will be used to decide with prev/next ops to not pull in with the backward/forward slices.

84

The original idea was to require all prevOps to have no side effects, including the backward slice for the scf.if condition. So we can feel free to clone them into scf.if regions, but they will be dead code if not used in scf.if regions and will be removed. The previous impl. also don't explicitly remove prevOps and let DCE did that, so backward slice of the scf.if condition will be kept as-is.

I changed the impl. to explicitly filter the backward slice for the scf.if condition now to make it clear. Also added a test where we have multiple ops computing the scf.if condition.

103

oldYieldOp is intended here. Thew newYieldOp is not created until at the end of this pullIntoBlock lambda. We clone the old block's operations *without* the old yield op when there are nextOps to pull into both regions. The newYieldOp will only be created after pulling in all nextOps too.

[nicolasvasilache]

mlir/lib/Dialect/SCF/Transforms/IfRegionExpansion.cpp

38

This is much too large, it should be broken down with the use of a bunch of helper functions.

42

Logically, all checks that operate only on the ifOp should be pulled out of this helper function and become early exits in the matchAndRewrite method.

48

Hoist to match and rewrite.

52

Hoist to match and rewrite and pass ops as argument to helper function.

56

Hoist to match and rewrite.

62

sg!

65

Hoist to match and rewrite.

115

There should be much more idiomatic constructs in llvm, in either SetOperations.h or SetVector.h to do intersection, union, difference.

134

It seems fishy that the terminator of the op enclosing the ifOp is involved in the logic.
Can you explain with doc comments and pseudo-IR ?

179

It seems fishy that the terminator of the op enclosing the ifOp is involved in the logic.
Can you explain with doc comments and pseudo-IR ?

199

It seems fishy that the terminator of the op enclosing the ifOp is involved in the logic.
Can you explain with doc comments and pseudo-IR ?

mlir/include/mlir/Dialect/SCF/Transforms.h

 void populateSCFLoopPipeliningPatterns(RewritePatternSet &patterns,
                                        const PipeliningOption &options);
 
 /// Populate patterns for canonicalizing operations inside SCF loop bodies.
 /// At the moment, only affine.min/max computations with iteration variables,
 /// loop bounds and loop steps are canonicalized.
 void populateSCFForLoopCanonicalizationPatterns(RewritePatternSet &patterns);
 
+/// Expands an scf.if op's regions by pulling ops before and after scf.if op
+/// into both regions of the scf.if op.
+//
+/// For example, it converts the following IR:
+/// ```
+/// %0 = opA ..
+/// %1 = opB ..
+/// %2 = scf.if .. {
+///   %3 = opC %0 ..
+///   scf.yield %3
+/// } else {
+///   %4 = opD ..
+///   scf.yield %4
+/// }
+/// %5 = opE %2 ..
+/// ```
+/// Into:
+/// ```
+/// %2 = scf.if .. {
+///   %0 = opA ..
+///   %1 = opB ..
+///   %3 = opC %0 ..
+///   %5 = opE %3 ..
+///   scf.yield %5
+/// } else {
+///   %0 = opA ..
+///   %1 = opB ..
+///   %4 = opD ..
+///   %5 = opE %4 ..
+///   scf.yield %5
+/// }
+/// ```
+///
+/// `canMoveToRegion` is called for each op before and after the scf.if op to
+/// decide whether it can be moved into the region. If not, the op and its
+/// backfoward/forward slices are kept as-is.
+void populateIfRegionExpansionPatterns(
+    RewritePatternSet &patterns,
+    function_ref<bool(Operation *)> canMoveToRegion);
 } // namespace scf
 } // namespace mlir
 
 #endif // MLIR_DIALECT_SCF_TRANSFORMS_H_

mlir/lib/Dialect/SCF/Transforms/CMakeLists.txt

 add_mlir_dialect_library(MLIRSCFTransforms
   BufferizableOpInterfaceImpl.cpp
   Bufferize.cpp
   ForToWhile.cpp
+  IfRegionExpansion.cpp
   LoopCanonicalization.cpp
   LoopPipelining.cpp
   LoopRangeFolding.cpp
   LoopSpecialization.cpp
   ParallelLoopCollapsing.cpp
   ParallelLoopFusion.cpp
   ParallelLoopTiling.cpp
   StructuralTypeConversions.cpp

mlir/lib/Dialect/SCF/Transforms/IfRegionExpansion.cpp

+//===- IfRegionExpansion.cpp - Pull ops into scf.if Region ----------------===//
+//
+// 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 file implements patterns and passes for expanding scf.if's regions
+// by pulling in ops before and after the scf.if op into both regions of the
+// scf.if op.
+//
+//===----------------------------------------------------------------------===//
+
+#include "PassDetail.h"
+#include "mlir/Analysis/SliceAnalysis.h"
+#include "mlir/Dialect/Affine/Analysis/AffineStructures.h"
+#include "mlir/Dialect/Affine/IR/AffineOps.h"
+#include "mlir/Dialect/SCF/Passes.h"
+#include "mlir/Dialect/SCF/SCF.h"
+#include "mlir/Dialect/SCF/Transforms.h"
+#include "mlir/Dialect/SCF/Utils/AffineCanonicalizationUtils.h"
+#include "mlir/IR/BlockAndValueMapping.h"
+#include "mlir/IR/Matchers.h"
+#include "mlir/IR/PatternMatch.h"
+#include "mlir/Transforms/GreedyPatternRewriteDriver.h"
+#include "mlir/Transforms/RegionUtils.h"
+#include "llvm/ADT/SetVector.h"
+#include "llvm/Support/Debug.h"
+
+#define DEBUG_TYPE "mlir-scf-expand-if-region"
+
+using namespace mlir;
+
+/// Pulls ops at the same nest level as the given `ifOp` into both regions of
+/// the if `ifOp`.
+static FailureOr<scf::IfOp>
+pullOpsIntoIfRegions(scf::IfOp ifOp,
+                     function_ref<bool(Operation *)> canMoveToRegion,
+                     RewriterBase &rewriter) {
+  // Need to pull ops into both regions.
+  if (!ifOp.elseBlock())
+    return failure();
+
+  // Expect to only have one block in the enclosing region. This is the common
+  // case for the level where we have structured control flows and it avoids
+  // traditional control flow and simplifies the analysis.
+  if (!llvm::hasSingleElement(ifOp->getParentRegion()->getBlocks()))
+    return failure();
+
+  SmallVector<Operation *> allOps;
+  for (Operation &op : ifOp->getBlock()->without_terminator())
+    allOps.push_back(&op);
+
+  // If no ops before or after the if op, there is nothing to do.
+  if (allOps.size() == 1)
+    return failure();
+
+  // Return true if the given `op` is in the same region as the scf.if op.
+  auto isFromSameRegion = [ifOp](Operation *op) {
+    return op->getParentRegion() == ifOp->getParentRegion();
+  };
+
+  // Collect ops before and after the scf.if op.
+  auto allPrevOps = llvm::makeArrayRef(allOps).take_while(
+      [&ifOp](Operation *op) { return op != ifOp.getOperation(); });
+  auto allNextOps =
+      llvm::makeArrayRef(allOps).drop_front(allPrevOps.size() + 1);
+
+  // Find previous ops that cannot be moved into the regions.
+  SetVector<Operation *> stickyPrevOps;
+  // We cannot move the op into the region if
+  // - The op is part of the backward slice for computing conditions.
+  // - The op is part of the backward slice for a op that cannot move.
+  if (Operation *condOp = ifOp.getCondition().getDefiningOp()) {
+    getBackwardSlice(condOp, &stickyPrevOps, isFromSameRegion);
+    stickyPrevOps.insert(condOp);
+  }
+  for (Operation *op : llvm::reverse(allPrevOps)) {
+    if (stickyPrevOps.contains(op))
+      continue;
+    if (!canMoveToRegion(op)) {
+      getBackwardSlice(op, &stickyPrevOps, isFromSameRegion);
+      stickyPrevOps.insert(op); // Add the current op back.
+    }
+  }
+
+  // Find out previous ops that cannot be moved into the regions.
+  SetVector<Operation *> stickyNextOps;
+  // We cannot move the op into the region if
+  // - The op is part of the forward slice for a op that cannot move.
+  for (Operation *op : allNextOps) {
+    if (!canMoveToRegion(op)) {
+      getForwardSlice(op, &stickyNextOps, isFromSameRegion);
+      stickyNextOps.insert(op); // Add the current op back.
+    }
+  }
+
+  LLVM_DEBUG({
+    llvm::dbgs() << "sticky previous ops:\n";
+    for (Operation *op : stickyPrevOps)
+      llvm::dbgs() << "  " << *op << "\n";
+    llvm::dbgs() << "sticky next ops:\n";
+    for (Operation *op : stickyNextOps)
+      llvm::dbgs() << "  " << *op << "\n";
+  });
+
+  // NYI support for the case where we have sticky next ops. For such cases we
+  // need to analyze their operands and figure out which are later coming from
+  // the if region. It can be complicated; support this only truly needed.
+  if (!stickyNextOps.empty())
+    return failure();
+
+  // Now get the ops that can be moved into the regions.
+  SmallVector<Operation *> prevOps, nextOps;
+  for (Operation *op : allPrevOps)
+    if (!stickyPrevOps.contains(op))
+      prevOps.push_back(op);
+  for (Operation *op : allNextOps)
+    if (!stickyNextOps.contains(op))
+      nextOps.push_back(op);
+
+  LLVM_DEBUG({
+    llvm::dbgs() << "previous ops to move:\n";
+    for (Operation *op : prevOps)
+      llvm::dbgs() << "  " << *op << "\n";
+    llvm::dbgs() << "next ops to move:\n";
+    for (Operation *op : nextOps)
+      llvm::dbgs() << "  " << *op << "\n";
+  });
+  if (prevOps.empty() && nextOps.empty())
+    return failure();
+
+  Operation *parentTerminator = ifOp->getBlock()->getTerminator();
+  TypeRange resultTypes = ifOp.getResultTypes();
+  if (!nextOps.empty()) {
+    // The if op should yield the values used by the terminator.
+    resultTypes = parentTerminator->getOperandTypes();
+  }
+
+  auto newIfOp = rewriter.create<scf::IfOp>(
+      ifOp.getLoc(), resultTypes, ifOp.getCondition(), ifOp.elseBlock());
+
+  auto pullIntoBlock = [&](Block *newblock, Block *oldBlock) {
+    OpBuilder::InsertionGuard guard(rewriter);
+    rewriter.setInsertionPointToStart(newblock);
+    BlockAndValueMapping bvm;
+
+    // Clone all ops defined before the original if op.
+    for (Operation *prevOp : prevOps)
+      rewriter.clone(*prevOp, bvm);
+
+    // Clone all ops defined inside the original if block.
+    for (Operation &blockOp : oldBlock->without_terminator())
+      rewriter.clone(blockOp, bvm);
+
+    if (nextOps.empty()) {
+      // If the if op needs to return value, its builder won't automatically
+      // insert terminators. Just clone the old one here.
+      if (newIfOp->getNumResults())
+        rewriter.clone(*oldBlock->getTerminator(), bvm);
+      return;
+    }
+
+    // There are ops after the old if op. Uses of the old if op should be
+    // replaced by the cloned yield value.
+    auto oldYieldOp = cast<scf::YieldOp>(oldBlock->back());
+    for (int i = 0, e = ifOp->getNumResults(); i < e; ++i) {
+      bvm.map(ifOp->getResult(i), bvm.lookup(oldYieldOp.getOperand(i)));
+    }
+
+    // Clone all ops defined after the original if op. While doing that, we need
+    // to check whether the op is used by the terminator. If so, we need to
+    // yield its result value at the proper index.
+    SmallVector<Value> yieldValues(newIfOp.getNumResults());
+    for (Operation *nextOp : nextOps) {
+      rewriter.clone(*nextOp, bvm);
+      for (OpOperand &use : nextOp->getUses()) {
+        if (use.getOwner() == parentTerminator) {
+          unsigned index = use.getOperandNumber();
+          yieldValues[index] = bvm.lookup(use.get());
+        }
+      }
+    }
+
+    if (!yieldValues.empty()) {
+      // Again the if builder won't insert terminators automatically.
+      rewriter.create<scf::YieldOp>(ifOp.getLoc(), yieldValues);
+    }
+  };
+
+  pullIntoBlock(newIfOp.thenBlock(), ifOp.thenBlock());
+  pullIntoBlock(newIfOp.elseBlock(), ifOp.elseBlock());
+
+  if (nextOps.empty()) {
+    rewriter.replaceOp(ifOp, newIfOp->getResults());
+  } else {
+    // Update the terminator to use the new if op's results.
+    rewriter.updateRootInPlace(parentTerminator, [&]() {
+      parentTerminator->setOperands(newIfOp->getResults());
+    });
+    // We have pulled in all ops following the if op into both regions. Now
+    // remove them all. Do this in the reverse order.
+    for (Operation *op : llvm::reverse(nextOps))
+      rewriter.eraseOp(op);
+    rewriter.eraseOp(ifOp);
+  }
+  for (Operation *op : llvm::reverse(prevOps))
+    rewriter.eraseOp(op);
+
+  return newIfOp;
+}
+
+namespace {
+
+class IfRegionExpansionPattern final : public OpRewritePattern<scf::IfOp> {
+public:
+  IfRegionExpansionPattern(MLIRContext *context,
+                           function_ref<bool(Operation *)> canMoveToRegion,
+                           PatternBenefit benefit = 1)
+      : OpRewritePattern(context, benefit), canMoveToRegion(canMoveToRegion) {}
+
+  LogicalResult matchAndRewrite(scf::IfOp ifOp,
+                                PatternRewriter &rewriter) const override {
+    return pullOpsIntoIfRegions(ifOp, canMoveToRegion, rewriter);
+  }
+
+private:
+  std::function<bool(Operation *)> canMoveToRegion;
+};
+
+} // namespace
+
+void scf::populateIfRegionExpansionPatterns(
+    RewritePatternSet &patterns,
+    function_ref<bool(Operation *)> canMoveToRegion) {
+  patterns.insert<IfRegionExpansionPattern>(patterns.getContext(),
+                                            canMoveToRegion);
+}

mlir/test/Dialect/SCF/if-region-expansion.mlir

+// RUN: mlir-opt -split-input-file -allow-unregistered-dialect -test-scf-transform-patterns=test-expand-if-region-patterns %s | FileCheck %s
+
+// CHECK-LABEL: func @parent_region_only_contains_if_op
+func @parent_region_only_contains_if_op(%cond: i1, %val: i32) -> i32 {
+  %if = scf.if %cond -> i32 {
+    scf.yield %val: i32
+  } else {
+    scf.yield %val: i32
+  }
+  return %if: i32
+}
+
+// CHECK: scf.if
+// CHECK:   scf.yield
+// CHECK: else
+// CHECK:   scf.yield
+
+// -----
+
+// CHECK-LABEL: func @if_op_without_else_branch
+func @if_op_without_else_branch(%cond: i1, %v0: i32, %v1: i32, %buffer: memref<i32>) {
+  %add = arith.addi %v0, %v1 : i32
+  scf.if %cond {
+    memref.store %add, %buffer[] : memref<i32>
+  }
+  %mul = arith.muli %v0, %v1 : i32
+  memref.store %mul, %buffer[] : memref<i32>
+  return
+}
+
+// CHECK: arith.addi
+// CHECK: scf.if
+// CHECK:   memref.store
+// CHECK: arith.muli
+// CHECK: memref.store
+
+// -----
+
+// CHECK-LABEL: func @non_side_effecting_ops_before_and_after_if_op
+//  CHECK-SAME: (%[[COND:.+]]: i1, %[[V0:.+]]: i32, %[[V1:.+]]: i32)
+func @non_side_effecting_ops_before_and_after_if_op(%cond: i1, %v0: i32, %v1: i32) -> (i32, i32) {
+  %add = arith.addi %v0, %v1 : i32
+  %sub = arith.subi %v0, %v1 : i32
+  %if = scf.if %cond -> i32 {
+    scf.yield %add: i32
+  } else {
+    scf.yield %sub: i32
+  }
+  %mul = arith.muli %if, %v0 : i32
+  %div = arith.divsi %if, %v1 : i32
+  return %mul, %div: i32, i32
+}
+
+// CHECK: %[[IF:.+]]:2 = scf.if %[[COND]] -> (i32, i32)
+// CHECK:   %[[ADD:.+]] = arith.addi %[[V0]], %[[V1]]
+// CHECK:   %[[MUL:.+]] = arith.muli %[[ADD]], %[[V0]]
+// CHECK:   %[[DIV:.+]] = arith.divsi %[[ADD]], %[[V1]]
+// CHECK:   scf.yield %[[MUL]], %[[DIV]]
+// CHECK: } else {
+// CHECK:   %[[SUB:.+]] = arith.subi %[[V0]], %[[V1]]
+// CHECK:   %[[MUL:.+]] = arith.muli %[[SUB]], %[[V0]]
+// CHECK:   %[[DIV:.+]] = arith.divsi %[[SUB]], %[[V1]]
+// CHECK:   scf.yield %[[MUL]], %[[DIV]]
+// CHECK: return %[[IF]]#0, %[[IF]]#1
+
+// -----
+
+// CHECK-LABEL: func @side_effect_op_before_after_if_op
+func @side_effect_op_before_after_if_op(%cond: i1, %v0: i32, %v1: i32, %buffer: memref<3xi32>) {
+  %c0 = arith.constant 0 : index
+  %c1 = arith.constant 1 : index
+  %c2 = arith.constant 2 : index
+  memref.store %v0, %buffer[%c0] : memref<3xi32>
+  %add = arith.addi %v0, %v1 : i32
+  scf.if %cond {
+    memref.store %add, %buffer[%c1] : memref<3xi32>
+  } else {
+    memref.store %add, %buffer[%c1] : memref<3xi32>
+  }
+  %mul = arith.muli %v0, %v1 : i32
+  memref.store %mul, %buffer[%c2] : memref<3xi32>
+  return
+}
+
+// The control in the test pass allows moving side effecting ops.
+
+// CHECK: scf.if
+// CHECK:   memref.store
+// CHECK:   arith.addi
+// CHECK:   memref.store
+// CHECK:   arith.muli
+// CHECK:   memref.store
+// CHECK: } else {
+// CHECK:   memref.store
+// CHECK:   arith.addi
+// CHECK:   memref.store
+// CHECK:   arith.muli
+// CHECK:   memref.store
+
+// -----
+
+// CHECK-LABEL: func @zero_result_if_op
+func @zero_result_if_op(%cond: i1, %v0: i32, %v1: i32, %buffer: memref<i32>) {
+  %add = arith.addi %v0, %v1 : i32
+  %sub = arith.subi %v0, %v1 : i32
+  scf.if %cond {
+    memref.store %add, %buffer[] : memref<i32>
+  } else {
+    memref.store %sub, %buffer[] : memref<i32>
+  }
+  %mul = arith.muli %v0, %v1 : i32
+  memref.store %mul, %buffer[] : memref<i32>
+  return
+}
+
+// CHECK: scf.if
+// CHECK:   %[[ADD:.+]] = arith.addi
+// CHECK:   memref.store %[[ADD]]
+// CHECK:   %[[MUL:.+]] = arith.muli
+// CHECK:   memref.store %[[MUL]]
+// CHECK: } else {
+// CHECK:   %[[SUB:.+]] = arith.subi
+// CHECK:   memref.store %[[SUB]]
+// CHECK:   %[[MUL:.+]] = arith.muli
+// CHECK:   memref.store %[[MUL]]
+
+// -----
+
+// CHECK-LABEL: func @multi_result_if_op
+//  CHECK-SAME: (%[[COND:.+]]: i1, %[[V0:.+]]: i32, %[[V1:.+]]: i32)
+func @multi_result_if_op(%cond: i1, %v0: i32, %v1: i32) -> (i32, i32) {
+  %add = arith.addi %v0, %v1 : i32
+  %sub = arith.subi %v0, %v1 : i32
+  %if:2 = scf.if %cond -> (i32, i32) {
+    scf.yield %add, %sub: i32, i32
+  } else {
+    scf.yield %sub, %add: i32, i32
+  }
+  %mul = arith.muli %if#0, %v0 : i32
+  %div = arith.divsi %if#1, %v1 : i32
+  return %mul, %div: i32, i32
+}
+
+// CHECK: %[[IF:.+]]:2 = scf.if
+// CHECK:   %[[ADD:.+]] = arith.addi %[[V0]], %[[V1]]
+// CHECK:   %[[SUB:.+]] = arith.subi %[[V0]], %[[V1]]
+// CHECK:   %[[MUL:.+]] = arith.muli %[[ADD]], %[[V0]]
+// CHECK:   %[[DIV:.+]] = arith.divsi %[[SUB]], %[[V1]]
+// CHECK:   scf.yield %[[MUL]], %[[DIV]]
+// CHECK: } else {
+// CHECK:   %[[ADD:.+]] = arith.addi %[[V0]], %[[V1]]
+// CHECK:   %[[SUB:.+]] = arith.subi %[[V0]], %[[V1]]
+// CHECK:   %[[MUL:.+]] = arith.muli %[[SUB]], %[[V0]]
+// CHECK:   %[[DIV:.+]] = arith.divsi %[[ADD]], %[[V1]]
+// CHECK:   scf.yield %[[MUL]], %[[DIV]]
+// CHECK: return %[[IF]]#0, %[[IF]]#1
+
+// -----
+
+// CHECK-LABEL: func @multi_use_in_terminator
+//  CHECK-SAME: (%[[COND:.+]]: i1, %[[V0:.+]]: i32, %[[V1:.+]]: i32)
+func @multi_use_in_terminator(%cond: i1, %v0: i32, %v1: i32) -> (i32, i32, i32) {
+  %add = arith.addi %v0, %v1 : i32
+  %sub = arith.subi %v0, %v1 : i32
+  %if = scf.if %cond -> i32 {
+    scf.yield %add: i32
+  } else {
+    scf.yield %sub: i32
+  }
+  %mul = arith.muli %if, %if : i32
+  return %mul, %mul, %mul: i32, i32, i32
+}
+
+// CHECK: %[[IF:.+]]:3 = scf.if
+// CHECK:   %[[ADD:.+]] = arith.addi %[[V0]], %[[V1]]
+// CHECK:   %[[MUL:.+]] = arith.muli %[[ADD]], %[[ADD]]
+// CHECK:   scf.yield %[[MUL]], %[[MUL]], %[[MUL]]
+// CHECK: } else {
+// CHECK:   %[[SUB:.+]] = arith.subi %[[V0]], %[[V1]]
+// CHECK:   %[[MUL:.+]] = arith.muli %[[SUB]], %[[SUB]]
+// CHECK:   scf.yield %[[MUL]], %[[MUL]], %[[MUL]]
+// CHECK: return %[[IF]]#0, %[[IF]]#1, %[[IF]]#2
+
+// -----
+
+// CHECK-LABEL: func @sticky_op_before_if_op
+func @sticky_op_before_if_op(%cond: i1, %index: index, %v0: i32, %v1: i32, %buffer: memref<3xi32>) {
+  %c0 = arith.constant 0 : index
+  %c1 = arith.constant 1 : index
+  %c2 = arith.constant 2 : index
+  %loc = arith.muli %index, %index : index
+  memref.store %v0, %buffer[%loc] {sticky} : memref<3xi32>
+  %add = arith.addi %v0, %v1 : i32
+  scf.if %cond {
+    memref.store %add, %buffer[%c1] : memref<3xi32>
+  } else {
+    memref.store %add, %buffer[%c1] : memref<3xi32>
+  }
+  %mul = arith.muli %v0, %v1 : i32
+  memref.store %mul, %buffer[%c2] : memref<3xi32>
+  return
+}
+
+// The control in the test pass disallows moving previous ops with "sticky" attribute and its backward slice.
+
+// CHECK: arith.muli
+// CHECK: memref.store {{.*}} {sticky}
+// CHECK: scf.if
+// CHECK:   arith.constant
+// CHECK:   arith.constant
+// CHECK:   arith.constant
+// CHECK:   arith.addi
+// CHECK:   memref.store
+// CHECK:   arith.muli
+// CHECK:   memref.store
+// CHECK: } else {
+// CHECK:   arith.constant
+// CHECK:   arith.constant
+// CHECK:   arith.constant
+// CHECK:   arith.addi
+// CHECK:   memref.store
+// CHECK:   arith.muli
+// CHECK:   memref.store
+
+// -----
+
+// CHECK-LABEL: func @sticky_op_after_if_op
+func @sticky_op_after_if_op(%cond: i1, %index: index, %v0: i32, %v1: i32, %buffer: memref<3xi32>) {
+  %c0 = arith.constant 0 : index
+  %c1 = arith.constant 1 : index
+  %c2 = arith.constant 2 : index
+  %loc = arith.muli %index, %index : index
+  memref.store %v0, %buffer[%loc] : memref<3xi32>
+  %add = arith.addi %v0, %v1 : i32
+  scf.if %cond {
+    memref.store %add, %buffer[%c1] : memref<3xi32>
+  } else {
+    memref.store %add, %buffer[%c1] : memref<3xi32>
+  }
+  %mul = arith.muli %v0, %v1 : i32
+  memref.store %mul, %buffer[%c2] {sticky} : memref<3xi32>
+  return
+}
+
+// NYI case for "sticky" next ops and its backward slice.
+
+// CHECK: arith.muli
+// CHECK: memref.store
+// CHECK: arith.addi
+// CHECK: scf.if {{.*}} {
+// CHECK:   memref.store
+// CHECK: } else {
+// CHECK:   memref.store
+// CHECK: }
+// CHECK: arith.muli
+// CHECK: memref.store {{.*}} {sticky}
+
+// -----
+
+// CHECK-LABEL: func @condition_back_slice
+func @condition_back_slice(%cond0: i1, %cond1: i1, %cond2 : i1, %v0: i32, %v1: i32, %buffer: memref<3xi32>) {
+  %c0 = arith.constant 0 : index
+  %c1 = arith.constant 1 : index
+  %c2 = arith.constant 2 : index
+  %add = arith.addi %v0, %v1 : i32
+  %and = arith.andi %cond0, %cond1 : i1
+  %or = arith.ori %and, %cond2 : i1
+  scf.if %or {
+    memref.store %add, %buffer[%c1] : memref<3xi32>
+  } else {
+    memref.store %add, %buffer[%c1] : memref<3xi32>
+  }
+  return
+}
+
+// CHECK: %[[AND:.+]] = arith.andi
+// CHECK: %[[OR:.+]] = arith.ori %[[AND]]
+// CHECK: scf.if %[[OR]]
+// CHECK:   arith.constant
+// CHECK:   arith.constant
+// CHECK:   arith.constant
+// CHECK:   arith.addi
+// CHECK:   memref.store
+// CHECK: } else {
+// CHECK:   arith.constant
+// CHECK:   arith.constant
+// CHECK:   arith.constant
+// CHECK:   arith.addi
+// CHECK:   memref.store

mlir/test/lib/Dialect/SCF/CMakeLists.txt

 # Exclude tests from libMLIR.so
 add_mlir_library(MLIRSCFTestPasses
   TestLoopParametricTiling.cpp
   TestLoopUnrolling.cpp
+  TestSCFTransforms.cpp
   TestSCFUtils.cpp
 
   EXCLUDE_FROM_LIBMLIR
 
   LINK_LIBS PUBLIC
   MLIRPass
   MLIRSCF
   MLIRSCFTransforms
   MLIRSCFUtils
   MLIRTransformUtils
   )

mlir/test/lib/Dialect/SCF/TestSCFTransforms.cpp

+//===- TestSCFTransforms.cpp - Test SCF transformation patterns -----------===//
+//
+// 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 file implements logic for testing SCF transformations.
+//
+//===----------------------------------------------------------------------===//
+
+#include "mlir/Dialect/Arithmetic/IR/Arithmetic.h"
+#include "mlir/Dialect/SCF/SCF.h"
+#include "mlir/Dialect/SCF/Transforms.h"
+#include "mlir/IR/BuiltinAttributes.h"
+#include "mlir/Pass/Pass.h"
+#include "mlir/Transforms/GreedyPatternRewriteDriver.h"
+
+using namespace mlir;
+
+namespace {
+struct TestSCFTransforms
+    : public PassWrapper<TestSCFTransforms, OperationPass<FuncOp>> {
+  TestSCFTransforms() = default;
+  TestSCFTransforms(const TestSCFTransforms &pass) : PassWrapper(pass) {}
+
+  void getDependentDialects(DialectRegistry &registry) const override {
+    registry.insert<arith::ArithmeticDialect, scf::SCFDialect>();
+  }
+
+  StringRef getArgument() const final { return "test-scf-transform-patterns"; }
+  StringRef getDescription() const final {
+    return "Test SCF patterns by applying them selectively or greedily.";
+  }
+
+  void runOnOperation() override;
+
+  Option<bool> testExpandIfRegionPatterns{
+      *this, "test-expand-if-region-patterns",
+      llvm::cl::desc("Test patterns to expand if regions"),
+      llvm::cl::init(false)};
+};
+} // namespace
+
+static void applyExpandIfRegionPatterns(FuncOp funcOp) {
+  SmallVector<scf::IfOp> candidates;
+  funcOp.walk([&](scf::IfOp ifOp) { candidates.push_back(ifOp); });
+
+  RewritePatternSet patterns(funcOp.getContext());
+  auto canMoveToRegion = [](Operation *op) {
+    return !op->hasAttrOfType<UnitAttr>("sticky");
+  };
+  scf::populateIfRegionExpansionPatterns(patterns, canMoveToRegion);
+  FrozenRewritePatternSet frozenPatterns(std::move(patterns));
+  for (scf::IfOp ifOp : candidates) {
+    // Apply transforms per op to avoid recursive behavior.
+    (void)applyOpPatternsAndFold(ifOp, frozenPatterns, /*erased=*/nullptr);
+  }
+}
+
+void TestSCFTransforms::runOnOperation() {
+  FuncOp func = getOperation();
+  if (testExpandIfRegionPatterns)
+    applyExpandIfRegionPatterns(func);
+}
+
+namespace mlir {
+namespace test {
+void registerTestSCFTransforms() { PassRegistration<TestSCFTransforms>(); }
+} // namespace test
+} // namespace mlir

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

 void registerTestMathPolynomialApproximationPass();
 void registerTestMemRefDependenceCheck();
 void registerTestMemRefStrideCalculation();
 void registerTestOpaqueLoc();
 void registerTestPadFusion();
 void registerTestPDLByteCodePass();
 void registerTestPreparationPassWithAllowedMemrefResults();
 void registerTestRecursiveTypesPass();
+void registerTestSCFTransforms();
 void registerTestSCFUtilsPass();
 void registerTestSliceAnalysisPass();
 void registerTestTensorTransforms();
 void registerTestVectorLowerings();
 } // namespace test
 } // namespace mlir
 
 namespace test {
#endif
   mlir::test::registerTestMathAlgebraicSimplificationPass();
   mlir::test::registerTestMathPolynomialApproximationPass();
   mlir::test::registerTestMemRefDependenceCheck();
   mlir::test::registerTestMemRefStrideCalculation();
   mlir::test::registerTestOpaqueLoc();
   mlir::test::registerTestPadFusion();
   mlir::test::registerTestPDLByteCodePass();
   mlir::test::registerTestRecursiveTypesPass();
+  mlir::test::registerTestSCFTransforms();
   mlir::test::registerTestSCFUtilsPass();
   mlir::test::registerTestSliceAnalysisPass();
   mlir::test::registerTestTensorTransforms();
   mlir::test::registerTestVectorLowerings();
 }
 #endif
 
 int main(int argc, char **argv) {