Skip to content

Commit 841b59f

Browse files
kripkenkripken
authored
[Branch Hinting] Support branch hints on If and BrOn (#7578)
With this, support for the feature is complete.
1 parent 1bd4267 commit 841b59f

File tree

4 files changed

+116
-15
lines changed

4 files changed

+116
-15
lines changed

‎src/parser/contexts.h

Lines changed: 5 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -1914,8 +1914,10 @@ struct ParseDefsCtx : TypeParserCtx<ParseDefsCtx> {
19141914
if (!type.isSignature()) {
19151915
return in.err(pos, "expected function type");
19161916
}
1917+
auto likely = getBranchHint(annotations);
19171918
return withLoc(
1918-
pos, irBuilder.makeIf(label ? *label : Name{}, type.getSignature()));
1919+
pos,
1920+
irBuilder.makeIf(label ? *label : Name{}, type.getSignature(), likely));
19191921
}
19201922

19211923
Result<> visitElse() { return withLoc(irBuilder.visitElse()); }
@@ -2546,7 +2548,8 @@ struct ParseDefsCtx : TypeParserCtx<ParseDefsCtx> {
25462548
BrOnOp op,
25472549
Type in = Type::none,
25482550
Type out = Type::none) {
2549-
return withLoc(pos, irBuilder.makeBrOn(label, op, in, out));
2551+
auto likely = getBranchHint(annotations);
2552+
return withLoc(pos, irBuilder.makeBrOn(label, op, in, out, likely));
25502553
}
25512554

25522555
Result<> makeStructNew(Index pos,

‎src/wasm-ir-builder.h

Lines changed: 10 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -116,7 +116,8 @@ class IRBuilder : public UnifiedExpressionVisitor<IRBuilder, Result<>> {
116116
// signatures ensure that there are no missing fields.
117117
Result<> makeNop();
118118
Result<> makeBlock(Name label, Signature sig);
119-
Result<> makeIf(Name label, Signature sig);
119+
Result<>
120+
makeIf(Name label, Signature sig, std::optional<bool> likely = std::nullopt);
120121
Result<> makeLoop(Name label, Signature sig);
121122
Result<> makeBreak(Index label,
122123
bool isConditional,
@@ -201,8 +202,11 @@ class IRBuilder : public UnifiedExpressionVisitor<IRBuilder, Result<>> {
201202
Result<> makeCallRef(HeapType type, bool isReturn);
202203
Result<> makeRefTest(Type type);
203204
Result<> makeRefCast(Type type);
204-
Result<>
205-
makeBrOn(Index label, BrOnOp op, Type in = Type::none, Type out = Type::none);
205+
Result<> makeBrOn(Index label,
206+
BrOnOp op,
207+
Type in = Type::none,
208+
Type out = Type::none,
209+
std::optional<bool> likely = std::nullopt);
206210
Result<> makeStructNew(HeapType type);
207211
Result<> makeStructNewDefault(HeapType type);
208212
Result<>
@@ -694,6 +698,9 @@ class IRBuilder : public UnifiedExpressionVisitor<IRBuilder, Result<>> {
694698
Expression* fixExtraOutput(ScopeCtx& scope, Name label, Expression* expr);
695699
void fixLoopWithInput(Loop* loop, Type inputType, Index scratch);
696700

701+
// Add a branch hint, if |likely| is present.
702+
void addBranchHint(Expression* expr, std::optional<bool> likely);
703+
697704
void dump();
698705
};
699706

‎src/wasm/wasm-ir-builder.cpp

Lines changed: 20 additions & 10 deletions
Original file line numberDiff line numberDiff line change
@@ -1380,9 +1380,11 @@ Result<> IRBuilder::makeBlock(Name label, Signature sig) {
13801380
return visitBlockStart(block, sig.params);
13811381
}
13821382

1383-
Result<> IRBuilder::makeIf(Name label, Signature sig) {
1383+
Result<>
1384+
IRBuilder::makeIf(Name label, Signature sig, std::optional<bool> likely) {
13841385
auto* iff = wasm.allocator.alloc<If>();
13851386
iff->type = sig.results;
1387+
addBranchHint(iff, likely);
13861388
return visitIfStart(iff, label, sig.params);
13871389
}
13881390

@@ -1407,14 +1409,9 @@ Result<> IRBuilder::makeBreak(Index label,
14071409
curr.condition = isConditional ? &curr : nullptr;
14081410
CHECK_ERR(ChildPopper{*this}.visitBreak(&curr, *labelType));
14091411
auto* br = builder.makeBreak(curr.name, curr.value, curr.condition);
1412+
addBranchHint(br, likely);
14101413
push(br);
14111414

1412-
if (likely) {
1413-
// Branches are only possible inside functions.
1414-
assert(func);
1415-
func->codeAnnotations[br].branchLikely = likely;
1416-
}
1417-
14181415
return Ok{};
14191416
}
14201417

@@ -1979,7 +1976,8 @@ Result<> IRBuilder::makeRefCast(Type type) {
19791976
return Ok{};
19801977
}
19811978

1982-
Result<> IRBuilder::makeBrOn(Index label, BrOnOp op, Type in, Type out) {
1979+
Result<> IRBuilder::makeBrOn(
1980+
Index label, BrOnOp op, Type in, Type out, std::optional<bool> likely) {
19831981
BrOn curr;
19841982
curr.op = op;
19851983
curr.castType = out;
@@ -2036,7 +2034,9 @@ Result<> IRBuilder::makeBrOn(Index label, BrOnOp op, Type in, Type out) {
20362034
auto name = getLabelName(label);
20372035
CHECK_ERR(name);
20382036

2039-
push(builder.makeBrOn(op, *name, curr.ref, out));
2037+
auto* br = builder.makeBrOn(op, *name, curr.ref, out);
2038+
addBranchHint(br, likely);
2039+
push(br);
20402040
return Ok{};
20412041
}
20422042

@@ -2058,7 +2058,9 @@ Result<> IRBuilder::makeBrOn(Index label, BrOnOp op, Type in, Type out) {
20582058

20592059
// Perform the branch.
20602060
CHECK_ERR(visitBrOn(&curr));
2061-
push(builder.makeBrOn(op, extraLabel, curr.ref, out));
2061+
auto* br = builder.makeBrOn(op, extraLabel, curr.ref, out);
2062+
addBranchHint(br, likely);
2063+
push(br);
20622064

20632065
// If the branch wasn't taken, we need to leave the extra values on the
20642066
// stack. For all instructions except br_on_non_null the extra values need
@@ -2531,4 +2533,12 @@ Result<> IRBuilder::makeStackSwitch(HeapType ct, Name tag) {
25312533
return Ok{};
25322534
}
25332535

2536+
void IRBuilder::addBranchHint(Expression* expr, std::optional<bool> likely) {
2537+
if (likely) {
2538+
// Branches are only possible inside functions.
2539+
assert(func);
2540+
func->codeAnnotations[expr].branchLikely = likely;
2541+
}
2542+
}
2543+
25342544
} // namespace wasm

‎test/lit/wat-annotations.wast

Lines changed: 81 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -7,6 +7,8 @@
77

88
;; CHECK: (type $0 (func (param i32)))
99

10+
;; CHECK: (type $1 (func (param anyref)))
11+
1012
;; CHECK: (func $no-annotations (type $0) (param $x i32)
1113
;; CHECK-NEXT: (block $out
1214
;; CHECK-NEXT: (br_if $out
@@ -16,6 +18,8 @@
1618
;; CHECK-NEXT: )
1719
;; RTRIP: (type $0 (func (param i32)))
1820

21+
;; RTRIP: (type $1 (func (param anyref)))
22+
1923
;; RTRIP: (func $no-annotations (type $0) (param $x i32)
2024
;; RTRIP-NEXT: (block $block
2125
;; RTRIP-NEXT: (br_if $block
@@ -148,4 +152,81 @@
148152
)
149153
)
150154
)
155+
156+
;; CHECK: (func $branch-hints-if (type $0) (param $x i32)
157+
;; CHECK-NEXT: (@metadata.code.branch_hint "\00")
158+
;; CHECK-NEXT: (if
159+
;; CHECK-NEXT: (local.get $x)
160+
;; CHECK-NEXT: (then
161+
;; CHECK-NEXT: (@metadata.code.branch_hint "\01")
162+
;; CHECK-NEXT: (if
163+
;; CHECK-NEXT: (local.get $x)
164+
;; CHECK-NEXT: (then
165+
;; CHECK-NEXT: (nop)
166+
;; CHECK-NEXT: )
167+
;; CHECK-NEXT: )
168+
;; CHECK-NEXT: )
169+
;; CHECK-NEXT: )
170+
;; CHECK-NEXT: )
171+
;; RTRIP: (func $branch-hints-if (type $0) (param $x i32)
172+
;; RTRIP-NEXT: (@metadata.code.branch_hint "\00")
173+
;; RTRIP-NEXT: (if
174+
;; RTRIP-NEXT: (local.get $x)
175+
;; RTRIP-NEXT: (then
176+
;; RTRIP-NEXT: (@metadata.code.branch_hint "\01")
177+
;; RTRIP-NEXT: (if
178+
;; RTRIP-NEXT: (local.get $x)
179+
;; RTRIP-NEXT: (then
180+
;; RTRIP-NEXT: (nop)
181+
;; RTRIP-NEXT: )
182+
;; RTRIP-NEXT: )
183+
;; RTRIP-NEXT: )
184+
;; RTRIP-NEXT: )
185+
;; RTRIP-NEXT: )
186+
(func $branch-hints-if (param $x i32)
187+
(@metadata.code.branch_hint "\00")
188+
(if
189+
(local.get $x)
190+
(then
191+
(@metadata.code.branch_hint "\01")
192+
(if
193+
(local.get $x)
194+
(then
195+
(nop)
196+
)
197+
)
198+
)
199+
)
200+
)
201+
202+
;; CHECK: (func $branch-hints-br_on (type $1) (param $x anyref)
203+
;; CHECK-NEXT: (block $out
204+
;; CHECK-NEXT: (drop
205+
;; CHECK-NEXT: (@metadata.code.branch_hint "\00")
206+
;; CHECK-NEXT: (br_on_null $out
207+
;; CHECK-NEXT: (local.get $x)
208+
;; CHECK-NEXT: )
209+
;; CHECK-NEXT: )
210+
;; CHECK-NEXT: )
211+
;; CHECK-NEXT: )
212+
;; RTRIP: (func $branch-hints-br_on (type $1) (param $x anyref)
213+
;; RTRIP-NEXT: (block $block
214+
;; RTRIP-NEXT: (drop
215+
;; RTRIP-NEXT: (@metadata.code.branch_hint "\00")
216+
;; RTRIP-NEXT: (br_on_null $block
217+
;; RTRIP-NEXT: (local.get $x)
218+
;; RTRIP-NEXT: )
219+
;; RTRIP-NEXT: )
220+
;; RTRIP-NEXT: )
221+
;; RTRIP-NEXT: )
222+
(func $branch-hints-br_on (param $x anyref)
223+
(block $out
224+
(drop
225+
(@metadata.code.branch_hint "\00")
226+
(br_on_null $out
227+
(local.get $x)
228+
)
229+
)
230+
)
231+
)
151232
)

0 commit comments

Comments
 (0)