Skip to content

Improve inlining of return_call* #6477

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 4 commits into from
Apr 10, 2024
Merged

Improve inlining of return_call* #6477

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
4 commits
Select commit Hold shift + click to select a range
ffa6180
Improve inlining of `return_call*`
vouillon Mar 28, 2024
816d80f
Comment + assertion
vouillon Apr 5, 2024
f7ea6f7
Updated examples
vouillon Apr 5, 2024
2e62ffe
Switch from try_table to try blocks
vouillon Apr 9, 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
88 changes: 60 additions & 28 deletions src/passes/Inlining.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -242,9 +242,10 @@ struct FunctionInfoScanner
struct InliningAction {
Expression** callSite;
Function* contents;
bool insideATry;

InliningAction(Expression** callSite, Function* contents)
: callSite(callSite), contents(contents) {}
InliningAction(Expression** callSite, Function* contents, bool insideATry)
: callSite(callSite), contents(contents), insideATry(insideATry) {}
};

struct InliningState {
Expand All @@ -254,7 +255,7 @@ struct InliningState {
std::unordered_map<Name, std::vector<InliningAction>> actionsForFunction;
};

struct Planner : public WalkerPass<PostWalker<Planner>> {
struct Planner : public WalkerPass<TryDepthWalker<Planner>> {
bool isFunctionParallel() override { return true; }

Planner(InliningState* state) : state(state) {}
Expand Down Expand Up @@ -287,15 +288,15 @@ struct Planner : public WalkerPass<PostWalker<Planner>> {
// can't add a new element in parallel
assert(state->actionsForFunction.count(getFunction()->name) > 0);
state->actionsForFunction[getFunction()->name].emplace_back(
&block->list[0], getModule()->getFunction(curr->target));
&block->list[0], getModule()->getFunction(curr->target), tryDepth > 0);
}
}

private:
InliningState* state;
};

struct Updater : public PostWalker<Updater> {
struct Updater : public TryDepthWalker<Updater> {
Module* module;
std::map<Index, Index> localMapping;
Name returnName;
Expand Down Expand Up @@ -335,19 +336,38 @@ struct Updater : public PostWalker<Updater> {
return;
}

// Set the children to locals as necessary, then add a branch out of the
// inlined body. The branch label will be set later when we create branch
// targets for the calls.
Block* childBlock = ChildLocalizer(curr, getFunction(), *module, options)
.getChildrenReplacement();
Break* branch = builder->makeBreak(Name());
childBlock->list.push_back(branch);
childBlock->type = Type::unreachable;
replaceCurrent(childBlock);

curr->isReturn = false;
curr->type = sig.results;
returnCallInfos.push_back({curr, branch});
if (tryDepth == 0) {
// Return calls in inlined functions should only break out of
// the scope of the inlined code, not the entire function they
// are being inlined into. To achieve this, make the call a
// non-return call and add a break. This does not cause
// unbounded stack growth because inlining and return calling
// both avoid creating a new stack frame.
curr->isReturn = false;
curr->type = sig.results;
// There might still be unreachable children causing this to be
// unreachable.
curr->finalize();
if (sig.results.isConcrete()) {
replaceCurrent(builder->makeBreak(returnName, curr));
} else {
replaceCurrent(builder->blockify(curr, builder->makeBreak(returnName)));
}
} else {
// Set the children to locals as necessary, then add a branch out of the
// inlined body. The branch label will be set later when we create branch
// targets for the calls.
Block* childBlock = ChildLocalizer(curr, getFunction(), *module, options)
.getChildrenReplacement();
Break* branch = builder->makeBreak(Name());
childBlock->list.push_back(branch);
childBlock->type = Type::unreachable;
replaceCurrent(childBlock);

curr->isReturn = false;
curr->type = sig.results;
returnCallInfos.push_back({curr, branch});
}
}

void visitCall(Call* curr) {
Expand Down Expand Up @@ -464,14 +484,15 @@ static Expression* doInlining(Module* module,
//
// (In this case we could use a second block and define the named block $X
// after the call's parameters, but that adds work for an extremely rare
// situation.) The latter case does not apply if the call is a return_call,
// because in that case the call's children do not appear inside the same
// block as the inlined body.
// situation.) The latter case does not apply if the call is a
// return_call inside a try, because in that case the call's
// children do not appear inside the same block as the inlined body.
bool hoistCall = call->isReturn && action.insideATry;
if (BranchUtils::hasBranchTarget(from->body, block->name) ||
(!call->isReturn && BranchUtils::BranchSeeker::has(call, block->name))) {
(!hoistCall && BranchUtils::BranchSeeker::has(call, block->name))) {
auto fromNames = BranchUtils::getBranchTargets(from->body);
auto callNames = call->isReturn ? BranchUtils::NameSet{}
: BranchUtils::BranchAccumulator::get(call);
auto callNames = hoistCall ? BranchUtils::NameSet{}
: BranchUtils::BranchAccumulator::get(call);
block->name = Names::getValidName(block->name, [&](Name test) {
return !fromNames.count(test) && !callNames.count(test);
});
Expand All @@ -490,7 +511,7 @@ static Expression* doInlining(Module* module,
updater.localMapping[i] = builder.addVar(into, from->getLocalType(i));
}

if (call->isReturn) {
if (hoistCall) {
// Wrap the existing function body in a block we can branch out of before
// entering the inlined function body. This block must have a name that is
// different from any other block name above the branch.
Expand Down Expand Up @@ -544,7 +565,16 @@ static Expression* doInlining(Module* module,
builder.makeLocalSet(updater.localMapping[from->getVarIndexBase() + i],
LiteralUtils::makeZero(type, *module)));
}
*action.callSite = block;
if (call->isReturn) {
assert(!action.insideATry);
if (retType.isConcrete()) {
*action.callSite = builder.makeReturn(block);
} else {
*action.callSite = builder.makeSequence(block, builder.makeReturn());
}
} else {
*action.callSite = block;
}
}

// Generate and update the inlined contents
Expand Down Expand Up @@ -1396,8 +1426,10 @@ struct InlineMainPass : public Pass {
// No call at all.
return;
}
doInlining(
module, main, InliningAction(callSite, originalMain), getPassOptions());
doInlining(module,
main,
InliningAction(callSite, originalMain, true),
getPassOptions());
}
};

Expand Down
47 changes: 47 additions & 0 deletions src/wasm-traversal.h
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -537,6 +537,53 @@ struct ExpressionStackWalker : public PostWalker<SubType, VisitorType> {
}
};

// Traversal keeping track of try depth

// This is used to keep track of whether we are in the scope of an
// exception handler. This matters since return_call is not equivalent
// to return + call within an exception handler. If another kind of
// handler scope is added, this code will need to be updated.
template<typename SubType, typename VisitorType = Visitor<SubType>>
struct TryDepthWalker : public PostWalker<SubType, VisitorType> {
TryDepthWalker() = default;

size_t tryDepth = 0;

static void doEnterTry(SubType* self, Expression** currp) {
self->tryDepth++;
}

static void doLeaveTry(SubType* self, Expression** currp) {
self->tryDepth--;
}

static void scan(SubType* self, Expression** currp) {
auto* curr = *currp;

if (curr->is<Try>()) {
self->pushTask(SubType::doVisitTry, currp);
auto& catchBodies = curr->cast<Try>()->catchBodies;
for (int i = int(catchBodies.size()) - 1; i >= 0; i--) {
self->pushTask(SubType::scan, &catchBodies[i]);
}
self->pushTask(SubType::doLeaveTry, currp);
self->pushTask(SubType::scan, &curr->cast<Try>()->body);
self->pushTask(SubType::doEnterTry, currp);
return;
}

if (curr->is<TryTable>()) {
self->pushTask(SubType::doLeaveTry, currp);
}

PostWalker<SubType, VisitorType>::scan(self, currp);

if (curr->is<TryTable>()) {
self->pushTask(SubType::doEnterTry, currp);
}
}
};

} // namespace wasm

#endif // wasm_wasm_traversal_h
58 changes: 45 additions & 13 deletions test/lit/passes/inlining-unreachable.wast
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -66,22 +66,51 @@
)

(module
;; CHECK: (type $0 (func (param i32) (result i32)))
;; CHECK: (type $0 (func))

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

;; CHECK: (import "env" "imported" (func $imported (type $0) (param i32) (result i32)))
;; CHECK: (import "env" "imported" (func $imported (type $1) (param i32) (result i32)))
(import "env" "imported" (func $imported (param i32) (result i32)))

;; CHECK: (func $caller (type $1)
;; CHECK: (func $caller (type $0)
;; CHECK-NEXT: (drop
;; CHECK-NEXT: (block
;; CHECK-NEXT: (block $__inlined_func$callee
;; CHECK-NEXT: (call $imported
;; CHECK-NEXT: (unreachable)
;; CHECK-NEXT: )
;; CHECK-NEXT: )
;; CHECK-NEXT: )
;; CHECK-NEXT: )
;; CHECK-NEXT: )
(func $caller
(drop
(call $callee)
)
)

;; After inlining, this return_call will turn into a call, but should still be
;; unreachable. Validation will fail if it is not.
(func $callee (result i32)
(return_call $imported
(unreachable)
)
)

;; CHECK: (func $caller-2 (type $0)
;; CHECK-NEXT: (drop
;; CHECK-NEXT: (block
;; CHECK-NEXT: (block $__inlined_func$callee-2$1
;; CHECK-NEXT: (block
;; CHECK-NEXT: (block $__return_call
;; CHECK-NEXT: (block
;; CHECK-NEXT: (unreachable)
;; CHECK-NEXT: (br $__return_call)
;; CHECK-NEXT: (try $try
;; CHECK-NEXT: (do
;; CHECK-NEXT: (unreachable)
;; CHECK-NEXT: (br $__return_call)
;; CHECK-NEXT: )
;; CHECK-NEXT: )
;; CHECK-NEXT: )
;; CHECK-NEXT: )
;; CHECK-NEXT: (call $imported
Expand All @@ -92,17 +121,20 @@
;; CHECK-NEXT: )
;; CHECK-NEXT: )
;; CHECK-NEXT: )
(func $caller
(func $caller-2
(drop
(call $callee)
(call $callee-2)
)
)

;; After inlining, this return_call will turn into a call, but should still be
;; unreachable. Validation will fail if it is not.
(func $callee (result i32)
(return_call $imported
(unreachable)
;; Same as above, but with a return_call with a try block
(func $callee-2 (result i32)
(try
(do
(return_call $imported
(unreachable)
)
)
)
)
)
Expand Down
Loading