Skip to content

Commit 835215f

Browse files
bgogulbgogul
authored
Simplify and refactor contractUncondEdges as a utility in TFCanonicalizeCFG. (#19546)
This PR also refactors the contractUncondBranches to reuse the mergeBasicBlockWithSuccessor(...) function from Utils/CFG.cpp
1 parent ad7def2 commit 835215f

File tree

4 files changed

+145
-37
lines changed

4 files changed

+145
-37
lines changed

‎lib/SILOptimizer/Mandatory/TFCanonicalizeCFG.cpp

Lines changed: 24 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -118,6 +118,29 @@ void ConditionalSESERegion::print(llvm::raw_ostream &OS, unsigned indent) const
118118
// CFG Canonicalization Implementation
119119
//===----------------------------------------------------------------------===//
120120

121+
// Our partitioning and other transformations can leave around lots of
122+
// unconditional branches between blocks that formerly had control edges. Go
123+
// through and merge those to make later passes simpler.
124+
bool tf::contractUncondBranches(SILFunction *fn, DominanceInfo* DI, SILLoopInfo *LI) {
125+
bool changed = false;
126+
// Iterate carefully to avoid invalidating iterators: we mutate the block list
127+
// while we walk it.
128+
for (auto bbi = fn->begin(), e = fn->end(); bbi != e;) {
129+
auto *bb = &*bbi;
130+
if (mergeBasicBlockWithSuccessor(bb, DI, LI)) {
131+
// The block was merged with this successor. Therefore, revisit this node:
132+
// we have new successor(s) and may need to contract them as well. Also,
133+
// bbi may be invalidated at this point.
134+
changed = true;
135+
bbi = SILFunction::iterator(bb);
136+
} else {
137+
// Move to the next block if this was not merged.
138+
++bbi;
139+
}
140+
}
141+
return changed;
142+
}
143+
121144
namespace {
122145
class SESERegionBuilder {
123146
DominanceInfo DI;
@@ -1156,8 +1179,8 @@ namespace {
11561179
/// The entry point to the transformation.
11571180
void run() override {
11581181
auto fn = getFunction();
1182+
contractUncondBranches(fn, /*DI*/nullptr, /*LI*/ nullptr);
11591183
auto region = canonicalizeCFGForXLA(fn);
1160-
11611184
llvm::outs() << "--- XLA CFG Canonicalize: " << fn->getName() << "\n";
11621185
region->print(llvm::outs());
11631186
llvm::outs() << "\n--- XLA CFG Canonicalize end\n";

‎lib/SILOptimizer/Mandatory/TFCanonicalizeCFG.h

Lines changed: 6 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -20,6 +20,7 @@
2020
#define SWIFT_SILOPTIMIZER_TFCANONICALIZECFG_H
2121

2222
#include "swift/SIL/SILBasicBlock.h"
23+
#include "swift/SIL/LoopInfo.h"
2324

2425
namespace swift {
2526
namespace tf {
@@ -157,6 +158,11 @@ namespace tf {
157158
}
158159
};
159160

161+
// Our partitioning and other transformations can leave around lots of
162+
// unconditional branches between blocks that formerly had control edges. Go
163+
// through and merge those to make later passes simpler.
164+
bool contractUncondBranches(SILFunction *fn, DominanceInfo *DI,
165+
SILLoopInfo *LI);
160166

161167
/// Transform the function into a properly nested series of
162168
/// single-entry-single-exit regions and return the data structure that

‎lib/SILOptimizer/Mandatory/TFPartition.cpp

Lines changed: 2 additions & 36 deletions
Original file line numberDiff line numberDiff line change
@@ -16,6 +16,7 @@
1616
//===----------------------------------------------------------------------===//
1717

1818
#define DEBUG_TYPE "tf-partition"
19+
#include "TFCanonicalizeCFG.h"
1920
#include "TFUtilities.h"
2021
#include "swift/AST/DiagnosticsSIL.h"
2122
#include "swift/AST/Expr.h"
@@ -4038,41 +4039,6 @@ bool TFFunctionPartition::partitionAndLowerGraph(bool isTest) {
40384039
return partition(isTest) || lowerGraph(isTest);
40394040
}
40404041

4041-
// Our partitioning can leave around lots of unconditional branches between
4042-
// blocks that formerly had control edges. Go through and merge those to make
4043-
// later passes simpler.
4044-
static void contractUncondBranches(SILFunction *fn) {
4045-
// Iterate carefully to avoid invalidating iterators: we mutate the block list
4046-
// while we walk it.
4047-
for (auto bbi = fn->begin(), e = fn->end(); bbi != e;) {
4048-
auto *bb = &*bbi;
4049-
++bbi; // Increment the iterator in case we do no transformation.
4050-
4051-
if (auto succ = bb->getSingleSuccessorBlock()) {
4052-
if (succ != bb && succ->getSinglePredecessorBlock()) {
4053-
if (auto *BI = dyn_cast<BranchInst>(bb->getTerminator())) {
4054-
// If there are any BB arguments in the destination, replace them with
4055-
// the branch operands, since they must dominate the dest block.
4056-
for (unsigned i = 0, e = BI->getArgs().size(); i != e; ++i) {
4057-
assert(succ->getArgument(i) != BI->getArg(i) &&
4058-
"Cloned code regions are always reachable");
4059-
succ->getArgument(i)->replaceAllUsesWith(BI->getArg(i));
4060-
}
4061-
4062-
// Zap BI and move all of the instructions from DestBB into this one.
4063-
BI->eraseFromParent();
4064-
bb->spliceAtEnd(succ);
4065-
succ->eraseFromParent();
4066-
4067-
// Revisit this node: we have new successor(s) and may need to
4068-
// contract them as well. Also, bbi may be invalidated at this point.
4069-
bbi = SILFunction::iterator(bb);
4070-
}
4071-
}
4072-
}
4073-
}
4074-
}
4075-
40764042
/// Return true if a user instruction is returning or forming a return value.
40774043
static bool isReturning(SILInstruction *user) {
40784044
if (isa<ReturnInst>(user)) return true;
@@ -4263,7 +4229,7 @@ bool TFFunctionPartition::partition(bool isTest) {
42634229
// Our partitioning can leave around lots of unconditional branches between
42644230
// blocks that formerly had control edges. Go through and merge those to
42654231
// make later passes simpler.
4266-
contractUncondBranches(acceleratorFn);
4232+
contractUncondBranches(acceleratorFn, /*DI*/ nullptr, /*LI*/ nullptr);
42674233

42684234
if (auto outs = getTFDumpIntermediateStream()) {
42694235
*outs << "--- TFPartition Accelerator Result: " << acceleratorFn->getName()

‎test/TensorFlow/contract_uncond_edges_test.sil

Lines changed: 113 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,113 @@
1+
// RUN: %target-sil-opt -tf-xla-cfg-canonicalize -tf-ensure-single-loop-exit -assume-parsing-unqualified-ownership-sil %s -o /dev/null | %FileCheck %s
2+
3+
import Builtin
4+
import Swift
5+
import TensorFlow
6+
7+
// Straight line blocks
8+
sil @$straightLine : $@convention(thin) (Builtin.Int32) -> Builtin.Int32 {
9+
bb0(%0 : $Builtin.Int32):
10+
br bb1(%0 : $Builtin.Int32)
11+
12+
bb1(%1: $Builtin.Int32):
13+
br bb2(%1 : $Builtin.Int32)
14+
15+
bb2(%2: $Builtin.Int32):
16+
br bb3(%2 : $Builtin.Int32)
17+
18+
bb3(%3 : $Builtin.Int32):
19+
return %3 : $Builtin.Int32
20+
}
21+
// CHECK-LABEL: --- XLA CFG Canonicalize: $straightLine
22+
// CHECK: block bb0
23+
// CHECK: --- XLA CFG Canonicalize end
24+
// CHECK: sil @$straightLine : $@convention(thin) (Builtin.Int32) -> Builtin.Int32 {
25+
// CHECK: bb0(%0 : $Builtin.Int32):
26+
// CHECK-NEXT: return %0 : $Builtin.Int32
27+
// CHECK: }
28+
29+
30+
// Conditionals
31+
sil @$condRegion : $@convention(thin) (Builtin.Int32, Builtin.Int32) -> Builtin.Int32 {
32+
bb0(%0 : $Builtin.Int32, %1 : $Builtin.Int32):
33+
%2 = builtin "cmp_slt_Int32"(%0 : $Builtin.Int32, %1 : $Builtin.Int32) : $Builtin.Int1
34+
cond_br %2, bb1, bb4
35+
36+
bb1:
37+
br bb2(%0 : $Builtin.Int32)
38+
39+
bb2(%3 : $Builtin.Int32):
40+
br bb3(%3 : $Builtin.Int32)
41+
42+
bb3(%4 : $Builtin.Int32):
43+
br bb5(%4 : $Builtin.Int32)
44+
45+
bb4:
46+
br bb5(%1 : $Builtin.Int32)
47+
48+
bb5(%5 : $Builtin.Int32):
49+
return %5 : $Builtin.Int32
50+
}
51+
// CHECK-LABEL: --- XLA CFG Canonicalize: $condRegion
52+
// CHECK:[sequence
53+
// CHECK: {condition Header: bb0
54+
// CHECK: block bb1
55+
// CHECK: block bb2}
56+
// CHECK: block bb3]
57+
// CHECK: --- XLA CFG Canonicalize end
58+
// CHECK: sil @$condRegion : $@convention(thin) {{.*}} {
59+
// CHECK: bb0(%0 : $Builtin.Int32, %1 : $Builtin.Int32):
60+
// CHECK: cond_br %2, bb1, bb2
61+
// CHECK: bb1:
62+
// CHECK: br bb3(%0 : $Builtin.Int32)
63+
// CHECK: bb2:
64+
// CHECK: br bb3(%1 : $Builtin.Int32)
65+
// CHECK: bb3([[A:%.*]] : $Builtin.Int32):
66+
// CHECK: return [[A]] : $Builtin.Int32
67+
// CHECK: }
68+
69+
// Simple loop
70+
sil @$simpleLoop : $@convention(thin) (Builtin.Int32, Builtin.Int32) -> Builtin.Int32 {
71+
bb0(%0 : $Builtin.Int32, %1 : $Builtin.Int32):
72+
br bb1(%0 : $Builtin.Int32)
73+
74+
bb1(%2 : $Builtin.Int32):
75+
%3 = builtin "cmp_slt_Int32"(%2 : $Builtin.Int32, %1 : $Builtin.Int32) : $Builtin.Int1
76+
cond_br %3, bb3, bb2
77+
78+
bb2:
79+
br bb6(%2 : $Builtin.Int32)
80+
81+
bb3:
82+
br bb4(%2 : $Builtin.Int32)
83+
84+
bb4(%4 : $Builtin.Int32):
85+
br bb5(%4 : $Builtin.Int32)
86+
87+
bb5(%5 : $Builtin.Int32):
88+
br bb1(%5 : $Builtin.Int32)
89+
90+
bb6(%6 : $Builtin.Int32):
91+
return %6 : $Builtin.Int32
92+
}
93+
// CHECK-LABEL:--- XLA CFG Canonicalize: $simpleLoop
94+
// CHECK:[sequence
95+
// CHECK: <while Preheader: bb0, Header: bb1, exit: bb2
96+
// CHECK: block bb3>
97+
// CHECK: block bb2]
98+
// CHECK:--- XLA CFG Canonicalize end
99+
// CHECK:sil @$simpleLoop : $@convention(thin) (Builtin.Int32, Builtin.Int32) -> Builtin.Int32 {
100+
// CHECK:bb0(%0 : $Builtin.Int32, %1 : $Builtin.Int32):
101+
// CHECK: br bb1(%0 : $Builtin.Int32)
102+
// CHECK:bb1(%3 : $Builtin.Int32):
103+
// CHECK: %4 = builtin "cmp_slt_Int32"(%3 : $Builtin.Int32, %1 : $Builtin.Int32) : $Builtin.Int1
104+
// CHECK: cond_br %4, bb3, bb2
105+
// CHECK:bb2:
106+
// CHECK: return %3 : $Builtin.Int32
107+
// CHECK:bb3:
108+
// CHECK: br bb1(%3 : $Builtin.Int32)
109+
// CHECK:}
110+
111+
112+
113+

0 commit comments

Comments
 (0)