Skip to content

Commit 1ef95ae

Browse files
ashleynhashleynh
authored
[Outlining] Fixes break reconstruction (#6352)
Adds new visitBreakWithType and visitSwitchWithType functions to the IRBuilder API. These functions work around an assumption in IRBuilder that the module is being traversed in the fully nested format, i.e., that the destination scope of a break or switch has been visited before visiting the break or switch. Instead, the type of the destination scope is passed to IRBuilder.
1 parent c62a0c9 commit 1ef95ae

File tree

4 files changed

+174
-6
lines changed

4 files changed

+174
-6
lines changed

‎src/passes/Outlining.cpp‎

Lines changed: 13 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -132,7 +132,19 @@ struct ReconstructStringifyWalker
132132
: state == NotInSeq ? &existingBuilder
133133
: nullptr;
134134
if (builder) {
135-
ASSERT_OK(builder->visit(curr));
135+
if (auto* expr = curr->dynCast<Break>()) {
136+
Type type = expr->value ? expr->value->type : Type::none;
137+
ASSERT_OK(builder->visitBreakWithType(expr, type));
138+
} else if (auto* expr = curr->dynCast<Switch>()) {
139+
Type type = expr->value ? expr->value->type : Type::none;
140+
ASSERT_OK(builder->visitSwitchWithType(expr, type));
141+
} else {
142+
// Assert ensures new unhandled branch instructions
143+
// will quickly cause an error. Serves as a reminder to
144+
// implement a new special-case visit*WithType.
145+
assert(curr->is<BrOn>() || !Properties::isBranch(curr));
146+
ASSERT_OK(builder->visit(curr));
147+
}
136148
}
137149
DBG(printVisitExpression(curr));
138150

‎src/wasm-ir-builder.h‎

Lines changed: 18 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -224,10 +224,26 @@ class IRBuilder : public UnifiedExpressionVisitor<IRBuilder, Result<>> {
224224
[[nodiscard]] Result<> visitStructNew(StructNew*);
225225
[[nodiscard]] Result<> visitArrayNew(ArrayNew*);
226226
[[nodiscard]] Result<> visitArrayNewFixed(ArrayNewFixed*);
227+
// Used to visit break exprs when traversing the module in the fully nested
228+
// format. Break label destinations are assumed to have already been visited,
229+
// with a corresponding push onto the scope stack. As a result, an error will
230+
// return if a corresponding scope is not found for the break.
227231
[[nodiscard]] Result<> visitBreak(Break*,
228232
std::optional<Index> label = std::nullopt);
233+
// Used to visit break nodes when traversing a single block without its
234+
// context. The type indicates how many values the break carries to its
235+
// destination.
236+
[[nodiscard]] Result<> visitBreakWithType(Break*, Type);
229237
[[nodiscard]] Result<>
238+
// Used to visit switch exprs when traversing the module in the fully nested
239+
// format. Switch label destinations are assumed to have already been visited,
240+
// with a corresponding push onto the scope stack. As a result, an error will
241+
// return if a corresponding scope is not found for the switch.
230242
visitSwitch(Switch*, std::optional<Index> defaultLabel = std::nullopt);
243+
// Used to visit switch nodes when traversing a single block without its
244+
// context. The type indicates how many values the switch carries to its
245+
// destination.
246+
[[nodiscard]] Result<> visitSwitchWithType(Switch*, Type);
231247
[[nodiscard]] Result<> visitCall(Call*);
232248
[[nodiscard]] Result<> visitCallIndirect(CallIndirect*);
233249
[[nodiscard]] Result<> visitCallRef(CallRef*);
@@ -535,8 +551,8 @@ class IRBuilder : public UnifiedExpressionVisitor<IRBuilder, Result<>> {
535551
[[nodiscard]] Result<> packageHoistedValue(const HoistedVal&,
536552
size_t sizeHint = 1);
537553

538-
[[nodiscard]] Result<Expression*> getBranchValue(Name labelName,
539-
std::optional<Index> label);
554+
[[nodiscard]] Result<Expression*>
555+
getBranchValue(Expression* curr, Name labelName, std::optional<Index> label);
540556

541557
void dump();
542558
};

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

Lines changed: 43 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -419,8 +419,14 @@ Result<> IRBuilder::visitArrayNewFixed(ArrayNewFixed* curr) {
419419
return Ok{};
420420
}
421421

422-
Result<Expression*> IRBuilder::getBranchValue(Name labelName,
422+
Result<Expression*> IRBuilder::getBranchValue(Expression* curr,
423+
Name labelName,
423424
std::optional<Index> label) {
425+
// As new branch instructions are added, one of the existing branch visit*
426+
// functions is likely to be copied, along with its call to getBranchValue().
427+
// This assert serves as a reminder to also add an implementation of
428+
// visit*WithType() for new branch instructions.
429+
assert(curr->is<Break>() || curr->is<Switch>());
424430
if (!label) {
425431
auto index = getLabelIndex(labelName);
426432
CHECK_ERR(index);
@@ -440,23 +446,57 @@ Result<> IRBuilder::visitBreak(Break* curr, std::optional<Index> label) {
440446
CHECK_ERR(cond);
441447
curr->condition = *cond;
442448
}
443-
auto value = getBranchValue(curr->name, label);
449+
auto value = getBranchValue(curr, curr->name, label);
444450
CHECK_ERR(value);
445451
curr->value = *value;
446452
return Ok{};
447453
}
448454

455+
Result<> IRBuilder::visitBreakWithType(Break* curr, Type type) {
456+
if (curr->condition) {
457+
auto cond = pop();
458+
CHECK_ERR(cond);
459+
curr->condition = *cond;
460+
}
461+
if (type == Type::none) {
462+
curr->value = nullptr;
463+
} else {
464+
auto value = pop(type.size());
465+
CHECK_ERR(value)
466+
curr->value = *value;
467+
}
468+
curr->finalize();
469+
push(curr);
470+
return Ok{};
471+
}
472+
449473
Result<> IRBuilder::visitSwitch(Switch* curr,
450474
std::optional<Index> defaultLabel) {
451475
auto cond = pop();
452476
CHECK_ERR(cond);
453477
curr->condition = *cond;
454-
auto value = getBranchValue(curr->default_, defaultLabel);
478+
auto value = getBranchValue(curr, curr->default_, defaultLabel);
455479
CHECK_ERR(value);
456480
curr->value = *value;
457481
return Ok{};
458482
}
459483

484+
Result<> IRBuilder::visitSwitchWithType(Switch* curr, Type type) {
485+
auto cond = pop();
486+
CHECK_ERR(cond);
487+
curr->condition = *cond;
488+
if (type == Type::none) {
489+
curr->value = nullptr;
490+
} else {
491+
auto value = pop(type.size());
492+
CHECK_ERR(value)
493+
curr->value = *value;
494+
}
495+
curr->finalize();
496+
push(curr);
497+
return Ok{};
498+
}
499+
460500
Result<> IRBuilder::visitCall(Call* curr) {
461501
auto numArgs = wasm.getFunction(curr->target)->getNumParams();
462502
curr->operands.resize(numArgs);

‎test/lit/passes/outlining.wast‎

Lines changed: 100 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -614,6 +614,106 @@
614614
)
615615
)
616616

617+
;; Tests branch with condition is reconstructed without error.
618+
(module
619+
;; CHECK: (type $0 (func))
620+
621+
;; CHECK: (func $outline$ (type $0)
622+
;; CHECK-NEXT: (drop
623+
;; CHECK-NEXT: (i32.const 2)
624+
;; CHECK-NEXT: )
625+
;; CHECK-NEXT: (drop
626+
;; CHECK-NEXT: (i32.const 1)
627+
;; CHECK-NEXT: )
628+
;; CHECK-NEXT: )
629+
630+
;; CHECK: (func $a (type $0)
631+
;; CHECK-NEXT: (block $label1
632+
;; CHECK-NEXT: (call $outline$)
633+
;; CHECK-NEXT: (loop $loop-in
634+
;; CHECK-NEXT: (br $label1)
635+
;; CHECK-NEXT: )
636+
;; CHECK-NEXT: (call $outline$)
637+
;; CHECK-NEXT: )
638+
;; CHECK-NEXT: )
639+
(func $a
640+
(block $label1
641+
(drop
642+
(i32.const 2)
643+
)
644+
(drop
645+
(i32.const 1)
646+
)
647+
(loop
648+
(br $label1)
649+
)
650+
(drop
651+
(i32.const 2)
652+
)
653+
(drop
654+
(i32.const 1)
655+
)
656+
)
657+
)
658+
)
659+
660+
;; Tests br_table instruction is reconstructed without error.
661+
(module
662+
;; CHECK: (type $0 (func))
663+
664+
;; CHECK: (type $1 (func (param i32) (result i32)))
665+
666+
;; CHECK: (func $outline$ (type $0)
667+
;; CHECK-NEXT: (drop
668+
;; CHECK-NEXT: (i32.const 2)
669+
;; CHECK-NEXT: )
670+
;; CHECK-NEXT: (drop
671+
;; CHECK-NEXT: (i32.const 1)
672+
;; CHECK-NEXT: )
673+
;; CHECK-NEXT: )
674+
675+
;; CHECK: (func $a (type $1) (param $0 i32) (result i32)
676+
;; CHECK-NEXT: (call $outline$)
677+
;; CHECK-NEXT: (block $block
678+
;; CHECK-NEXT: (block $block0
679+
;; CHECK-NEXT: (br_table $block $block0
680+
;; CHECK-NEXT: (local.get $0)
681+
;; CHECK-NEXT: )
682+
;; CHECK-NEXT: (return
683+
;; CHECK-NEXT: (i32.const 21)
684+
;; CHECK-NEXT: )
685+
;; CHECK-NEXT: )
686+
;; CHECK-NEXT: (return
687+
;; CHECK-NEXT: (i32.const 20)
688+
;; CHECK-NEXT: )
689+
;; CHECK-NEXT: )
690+
;; CHECK-NEXT: (call $outline$)
691+
;; CHECK-NEXT: (i32.const 22)
692+
;; CHECK-NEXT: )
693+
(func $a (param i32) (result i32)
694+
(drop
695+
(i32.const 2)
696+
)
697+
(drop
698+
(i32.const 1)
699+
)
700+
(block
701+
(block
702+
(br_table 1 0 (local.get $0))
703+
(return (i32.const 21))
704+
)
705+
(return (i32.const 20))
706+
)
707+
(drop
708+
(i32.const 2)
709+
)
710+
(drop
711+
(i32.const 1)
712+
)
713+
(i32.const 22)
714+
)
715+
)
716+
617717
;; Tests return instructions are correctly filtered from being outlined.
618718
(module
619719
;; CHECK: (type $0 (func (result i32)))

0 commit comments

Comments
 (0)