[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

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

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

60

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

39

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

43

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.

49

Hoist to match and rewrite.

53

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

57

Hoist to match and rewrite.

62

sg!

66

Hoist to match and rewrite.

116

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

135

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 ?

180

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 ?

200

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/Passes.h

 
 /// Creates a pass which folds arith ops on induction variable into
 /// loop range.
 std::unique_ptr<Pass> createForLoopRangeFoldingPass();
 
 // Creates a pass which lowers for loops into while loops.
 std::unique_ptr<Pass> createForToWhileLoopPass();
 
+/// Creates a pass to pull ops before and after an scf.if op into both scf.if op
+/// regions.
+std::unique_ptr<Pass> createIfRegionExpansionPass();
+
 //===----------------------------------------------------------------------===//
 // Registration
 //===----------------------------------------------------------------------===//
 
 /// Generate the code for registering passes.
 #define GEN_PASS_REGISTRATION
 #include "mlir/Dialect/SCF/Passes.h.inc"
 
 } // namespace mlir
 
 #endif // MLIR_DIALECT_SCF_PASSES_H_

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

# After:
         %2 = arith.addi %arg2, %arg2 : i32
         memref.store %2, %arg0[%i] : memref<?xi32>
         scf.yield %1 : index
       }
     ```
   }];
 }
 
+def SCFIfRegionExpansion
+  : FunctionPass<"if-region-expansion"> {
+  let summary = "Pulls ops before and after scf.if into both scf.if regions";
+  let constructor = "mlir::createIfRegionExpansionPass()";
+  let description = [{
+    This pass expands an scf.if op's regions by pulling in ops before and after
+    scf.if op into both regions of the scf.if op. This can be helpful as a
+    prelimiary step to enable further optimizations in both regions, which now
+    does not need to cross region boundaries.
+
+    The scf.if op's parent region should only contain one block.
+  }];
+}
+
 #endif // MLIR_DIALECT_SCF_PASSES

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
+/// }
+/// ```
+void populateIfRegionExpansionPatterns(RewritePatternSet &patterns);
 } // namespace scf
 } // namespace mlir
 
 #endif // MLIR_DIALECT_SCF_TRANSFORMS_H_

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

 add_mlir_dialect_library(MLIRSCFTransforms
   AffineCanonicalizationUtils.cpp
   Bufferize.cpp
   ForToWhile.cpp
+  IfRegionExpansion.cpp
   LoopCanonicalization.cpp
   LoopPipelining.cpp
   LoopRangeFolding.cpp
   LoopSpecialization.cpp
   ParallelLoopFusion.cpp
   ParallelLoopTiling.cpp
   StructuralTypeConversions.cpp
   Utils.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/AffineStructures.h"
+#include "mlir/Dialect/Affine/IR/AffineOps.h"
+#include "mlir/Dialect/SCF/AffineCanonicalizationUtils.h"
+#include "mlir/Dialect/SCF/Passes.h"
+#include "mlir/Dialect/SCF/SCF.h"
+#include "mlir/Dialect/SCF/Transforms.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;
+
+static constexpr char kExpandedIfMarker[] = "__expanded_if_regions__";
+
+/// 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,
+                                                 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();
+
+  auto prevOps = llvm::makeArrayRef(allOps).take_while(
+      [&ifOp](Operation *op) { return op != ifOp.getOperation(); });
+  auto nextOps = llvm::makeArrayRef(allOps).drop_front(prevOps.size() + 1);
+
+  // Require all previous ops to have on side effects, so that after cloning
+  // them into both regions, we can rely on DCE to remove them.
+  if (llvm::any_of(prevOps, [](Operation *op) {
+        return !MemoryEffectOpInterface::hasNoEffect(op);
+      }))
+    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);
+  }
+
+  return newIfOp;
+}
+
+namespace {
+
+struct IfRegionExpansionPattern final : public OpRewritePattern<scf::IfOp> {
+  using OpRewritePattern::OpRewritePattern;
+
+  LogicalResult matchAndRewrite(scf::IfOp ifOp,
+                                PatternRewriter &rewriter) const override {
+    if (ifOp->hasAttr(kExpandedIfMarker))
+      return failure();
+
+    auto newOp = pullOpsIntoIfRegions(ifOp, rewriter);
+    if (failed(newOp))
+      return failure();
+
+    newOp.getValue()->setAttr(kExpandedIfMarker, rewriter.getUnitAttr());
+    return success();
+  }
+};
+
+struct IfRegionExpansion : public SCFIfRegionExpansionBase<IfRegionExpansion> {
+  void runOnFunction() override {
+    FuncOp funcOp = getFunction();
+    MLIRContext *ctx = funcOp.getContext();
+    RewritePatternSet patterns(ctx);
+    patterns.add<IfRegionExpansionPattern>(ctx);
+    (void)applyPatternsAndFoldGreedily(funcOp, std::move(patterns));
+  }
+};
+
+} // namespace
+
+void scf::populateIfRegionExpansionPatterns(RewritePatternSet &patterns) {
+  patterns.insert<IfRegionExpansionPattern>(patterns.getContext());
+}
+
+std::unique_ptr<Pass> mlir::createIfRegionExpansionPass() {
+  return std::make_unique<IfRegionExpansion>();
+}

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

+// RUN: mlir-opt -split-input-file -allow-unregistered-dialect -if-region-expansion %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-NOT: __expanded_if_regions__
+
+// -----
+
+// CHECK-LABEL: func @side_effect_op_before_if_op
+func @side_effect_op_before_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>
+  }
+  %mul = arith.muli %v0, %v1 : i32
+  memref.store %mul, %buffer[%c2] : memref<3xi32>
+  return
+}
+
+// CHECK-NOT: __expanded_if_regions__
+
+// -----
+
+// 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-NOT: __expanded_if_regions__
+
+// -----
+
+// CHECK-LABEL: func @ops_before_and_after_if_op
+//  CHECK-SAME: (%[[COND:.+]]: i1, %[[V0:.+]]: i32, %[[V1:.+]]: i32)
+func @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 @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