Skip to content

Commit

Permalink
[Outlining] Fixes break reconstruction (WebAssembly#6352)
Browse files Browse the repository at this point in the history 
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.
  • Loading branch information
@ashleynh
ashleynh authored Feb 28, 2024
1 parent c62a0c9 commit 1ef95ae
Show file tree
Hide file tree
Showing 4 changed files with 174 additions and 6 deletions.
14 changes: 13 additions & 1 deletion src/passes/Outlining.cpp
View file
Original file line number Diff line number Diff line change
Expand Up @@ -132,7 +132,19 @@ struct ReconstructStringifyWalker
: state == NotInSeq ? &existingBuilder
: nullptr;
if (builder) {
ASSERT_OK(builder->visit(curr));
if (auto* expr = curr->dynCast<Break>()) {
Type type = expr->value ? expr->value->type : Type::none;
ASSERT_OK(builder->visitBreakWithType(expr, type));
} else if (auto* expr = curr->dynCast<Switch>()) {
Type type = expr->value ? expr->value->type : Type::none;
ASSERT_OK(builder->visitSwitchWithType(expr, type));
} else {
// Assert ensures new unhandled branch instructions
// will quickly cause an error. Serves as a reminder to
// implement a new special-case visit*WithType.
assert(curr->is<BrOn>() || !Properties::isBranch(curr));
ASSERT_OK(builder->visit(curr));
}
}
DBG(printVisitExpression(curr));

Expand Down
20 changes: 18 additions & 2 deletions src/wasm-ir-builder.h
View file
Original file line number Diff line number Diff line change
Expand Up @@ -224,10 +224,26 @@ class IRBuilder : public UnifiedExpressionVisitor<IRBuilder, Result<>> {
[[nodiscard]] Result<> visitStructNew(StructNew*);
[[nodiscard]] Result<> visitArrayNew(ArrayNew*);
[[nodiscard]] Result<> visitArrayNewFixed(ArrayNewFixed*);
// Used to visit break exprs when traversing the module in the fully nested
// format. Break label destinations are assumed to have already been visited,
// with a corresponding push onto the scope stack. As a result, an error will
// return if a corresponding scope is not found for the break.
[[nodiscard]] Result<> visitBreak(Break*,
std::optional<Index> label = std::nullopt);
// Used to visit break nodes when traversing a single block without its
// context. The type indicates how many values the break carries to its
// destination.
[[nodiscard]] Result<> visitBreakWithType(Break*, Type);
[[nodiscard]] Result<>
// Used to visit switch exprs when traversing the module in the fully nested
// format. Switch label destinations are assumed to have already been visited,
// with a corresponding push onto the scope stack. As a result, an error will
// return if a corresponding scope is not found for the switch.
visitSwitch(Switch*, std::optional<Index> defaultLabel = std::nullopt);
// Used to visit switch nodes when traversing a single block without its
// context. The type indicates how many values the switch carries to its
// destination.
[[nodiscard]] Result<> visitSwitchWithType(Switch*, Type);
[[nodiscard]] Result<> visitCall(Call*);
[[nodiscard]] Result<> visitCallIndirect(CallIndirect*);
[[nodiscard]] Result<> visitCallRef(CallRef*);
Expand Down Expand Up @@ -535,8 +551,8 @@ class IRBuilder : public UnifiedExpressionVisitor<IRBuilder, Result<>> {
[[nodiscard]] Result<> packageHoistedValue(const HoistedVal&,
size_t sizeHint = 1);

[[nodiscard]] Result<Expression*> getBranchValue(Name labelName,
std::optional<Index> label);
[[nodiscard]] Result<Expression*>
getBranchValue(Expression* curr, Name labelName, std::optional<Index> label);

void dump();
};
Expand Down
46 changes: 43 additions & 3 deletions src/wasm/wasm-ir-builder.cpp
View file
Original file line number Diff line number Diff line change
Expand Up @@ -419,8 +419,14 @@ Result<> IRBuilder::visitArrayNewFixed(ArrayNewFixed* curr) {
return Ok{};
}

Result<Expression*> IRBuilder::getBranchValue(Name labelName,
Result<Expression*> IRBuilder::getBranchValue(Expression* curr,
Name labelName,
std::optional<Index> label) {
// As new branch instructions are added, one of the existing branch visit*
// functions is likely to be copied, along with its call to getBranchValue().
// This assert serves as a reminder to also add an implementation of
// visit*WithType() for new branch instructions.
assert(curr->is<Break>() || curr->is<Switch>());
if (!label) {
auto index = getLabelIndex(labelName);
CHECK_ERR(index);
Expand All @@ -440,23 +446,57 @@ Result<> IRBuilder::visitBreak(Break* curr, std::optional<Index> label) {
CHECK_ERR(cond);
curr->condition = *cond;
}
auto value = getBranchValue(curr->name, label);
auto value = getBranchValue(curr, curr->name, label);
CHECK_ERR(value);
curr->value = *value;
return Ok{};
}

Result<> IRBuilder::visitBreakWithType(Break* curr, Type type) {
if (curr->condition) {
auto cond = pop();
CHECK_ERR(cond);
curr->condition = *cond;
}
if (type == Type::none) {
curr->value = nullptr;
} else {
auto value = pop(type.size());
CHECK_ERR(value)
curr->value = *value;
}
curr->finalize();
push(curr);
return Ok{};
}

Result<> IRBuilder::visitSwitch(Switch* curr,
std::optional<Index> defaultLabel) {
auto cond = pop();
CHECK_ERR(cond);
curr->condition = *cond;
auto value = getBranchValue(curr->default_, defaultLabel);
auto value = getBranchValue(curr, curr->default_, defaultLabel);
CHECK_ERR(value);
curr->value = *value;
return Ok{};
}

Result<> IRBuilder::visitSwitchWithType(Switch* curr, Type type) {
auto cond = pop();
CHECK_ERR(cond);
curr->condition = *cond;
if (type == Type::none) {
curr->value = nullptr;
} else {
auto value = pop(type.size());
CHECK_ERR(value)
curr->value = *value;
}
curr->finalize();
push(curr);
return Ok{};
}

Result<> IRBuilder::visitCall(Call* curr) {
auto numArgs = wasm.getFunction(curr->target)->getNumParams();
curr->operands.resize(numArgs);
Expand Down
100 changes: 100 additions & 0 deletions test/lit/passes/outlining.wast
View file
Original file line number Diff line number Diff line change
Expand Up @@ -614,6 +614,106 @@
)
)

;; Tests branch with condition is reconstructed without error.
(module
;; CHECK: (type $0 (func))

;; CHECK: (func $outline$ (type $0)
;; CHECK-NEXT: (drop
;; CHECK-NEXT: (i32.const 2)
;; CHECK-NEXT: )
;; CHECK-NEXT: (drop
;; CHECK-NEXT: (i32.const 1)
;; CHECK-NEXT: )
;; CHECK-NEXT: )

;; CHECK: (func $a (type $0)
;; CHECK-NEXT: (block $label1
;; CHECK-NEXT: (call $outline$)
;; CHECK-NEXT: (loop $loop-in
;; CHECK-NEXT: (br $label1)
;; CHECK-NEXT: )
;; CHECK-NEXT: (call $outline$)
;; CHECK-NEXT: )
;; CHECK-NEXT: )
(func $a
(block $label1
(drop
(i32.const 2)
)
(drop
(i32.const 1)
)
(loop
(br $label1)
)
(drop
(i32.const 2)
)
(drop
(i32.const 1)
)
)
)
)

;; Tests br_table instruction is reconstructed without error.
(module
;; CHECK: (type $0 (func))

;; CHECK: (type $1 (func (param i32) (result i32)))

;; CHECK: (func $outline$ (type $0)
;; CHECK-NEXT: (drop
;; CHECK-NEXT: (i32.const 2)
;; CHECK-NEXT: )
;; CHECK-NEXT: (drop
;; CHECK-NEXT: (i32.const 1)
;; CHECK-NEXT: )
;; CHECK-NEXT: )

;; CHECK: (func $a (type $1) (param $0 i32) (result i32)
;; CHECK-NEXT: (call $outline$)
;; CHECK-NEXT: (block $block
;; CHECK-NEXT: (block $block0
;; CHECK-NEXT: (br_table $block $block0
;; CHECK-NEXT: (local.get $0)
;; CHECK-NEXT: )
;; CHECK-NEXT: (return
;; CHECK-NEXT: (i32.const 21)
;; CHECK-NEXT: )
;; CHECK-NEXT: )
;; CHECK-NEXT: (return
;; CHECK-NEXT: (i32.const 20)
;; CHECK-NEXT: )
;; CHECK-NEXT: )
;; CHECK-NEXT: (call $outline$)
;; CHECK-NEXT: (i32.const 22)
;; CHECK-NEXT: )
(func $a (param i32) (result i32)
(drop
(i32.const 2)
)
(drop
(i32.const 1)
)
(block
(block
(br_table 1 0 (local.get $0))
(return (i32.const 21))
)
(return (i32.const 20))
)
(drop
(i32.const 2)
)
(drop
(i32.const 1)
)
(i32.const 22)
)
)

;; Tests return instructions are correctly filtered from being outlined.
(module
;; CHECK: (type $0 (func (result i32)))
Expand Down

0 comments on commit 1ef95ae

Please sign in to comment.