Skip to content

[Outlining] Fixes break reconstruction #6352

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
merged 7 commits into from
Feb 28, 2024
Merged

[Outlining] Fixes break reconstruction #6352

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
7 commits
Select commit Hold shift + click to select a range
2a00196
label fix
ashleynh Jan 30, 2024
a8de1ee
comments
ashleynh Feb 27, 2024
c84d16b
assert
ashleynh Feb 27, 2024
ae7ab8a
check type none
ashleynh Feb 27, 2024
d7a60e3
assert comment
ashleynh Feb 28, 2024
c9ac6eb
assert comment2
ashleynh Feb 28, 2024
ed481b5
oops
ashleynh Feb 28, 2024
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
14 changes: 13 additions & 1 deletion src/passes/Outlining.cpp
View file
Open in desktop
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));
Copy link
Member

@tlively tlively Feb 27, 2024

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

It would be good to add an assertion somewhere to future-proof ourselves against the possibility that a new kind of Expression starts to need a visit*WithType variant. This could either be an assert(!Properties::isBranch(curr)) here guarding the default builder->visit(curr) case, or it could be an assertion that curr is a Break or Switch inside IRBuilder::getBranchValue.

Copy link
Collaborator Author

@ashleynh ashleynh Feb 27, 2024

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

done

}
}
DBG(printVisitExpression(curr));

Expand Down
20 changes: 18 additions & 2 deletions src/wasm-ir-builder.h
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -220,10 +220,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 @@ -528,8 +544,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
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -404,8 +404,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 @@ -425,23 +431,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
Open in desktop
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
Copy link
Member

@tlively tlively Feb 27, 2024

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

It would be good to add a test for the br_table case as well.

Copy link
Collaborator Author

@ashleynh ashleynh Feb 27, 2024

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

done

(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