Skip to content

[DAE2] Refinalize when changing types #9290

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

Open
tlively wants to merge 1 commit into
base: main
Choose a base branch
from
Open

[DAE2] Refinalize when changing types #9290

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
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
79 changes: 60 additions & 19 deletions src/passes/DeadArgumentElimination2.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -50,6 +50,7 @@
#include "ir/local-graph.h"
#include "ir/module-utils.h"
#include "ir/type-updating.h"
#include "ir/utils.h"
#include "pass.h"
#include "support/index.h"
#include "support/mixed_arena.h"
Expand Down Expand Up @@ -204,6 +205,11 @@ struct FunctionInfo {

// Root function types that this function indirectly tail-calls.
std::vector<HeapType> tailCalleeTypes;

// Replacing a removed expression with its kept children can change its type
// from none to unreachable, requiring enclosing expressions to be
// refinalized.
bool refinalize = false;
};

// Analysis results and call graph information for a tree of related function
Expand Down Expand Up @@ -992,10 +998,10 @@ struct Optimizer
// types to match. Function types are updated after this pass runs.
bool requiresNonNullableLocalFixups() override { return false; }

const DAE2& parent;
DAE2& parent;

// The info for the function we are running on.
const FunctionInfo* funcInfo = nullptr;
FunctionInfo* funcInfo = nullptr;

// Map old local indices to new local indices for the function we are
// currently optimizing. Kept parameters and locals may need to have their
Expand All @@ -1015,7 +1021,16 @@ struct Optimizer
// We will need to generate fresh labels for trampoline blocks.
std::optional<LabelUtils::LabelManager> labels;

Optimizer(const DAE2& parent) : parent(parent) {}
Optimizer(DAE2& parent) : parent(parent) {}

void replaceCurrent(Expression* rep) {
// If we introduce any new unreachability, we will have to refinalize.
if (rep->type == Type::unreachable &&
getCurrent()->type != Type::unreachable) {
funcInfo->refinalize = true;
}
Super::replaceCurrent(rep);
}

std::unique_ptr<Pass> create() override {
return std::make_unique<Optimizer>(parent);
Expand Down Expand Up @@ -1573,37 +1588,63 @@ Expression* Optimizer::getReplacement(Expression* curr) {
return collector.collectExprsOrNull();
}

struct Refinalizer : public ReFinalize {
const DAE2& parent;

Refinalizer(const DAE2& parent) : parent(parent) {}

std::unique_ptr<Pass> create() override {
return std::make_unique<Refinalizer>(parent);
}

void runOnFunction(Module* wasm, Function* func) override {
if (parent.funcInfos[parent.funcIndices.at(func->name)].refinalize) {
ReFinalize::runOnFunction(wasm, func);
}
}
};

void DAE2::optimize() {
TIME(Timer timer);

Optimizer optimizer(*this);
if (!optimizeReferencedFuncs) {
// Cannot globally rewrite types or optimize referenced functions, so just
// run the optimizer to optimize unreferenced functions.
optimizer.run(getPassRunner(), wasm);
return;
}

TIME(Timer timer);
TIME(std::cerr << " optimize: " << timer.lastElapsed() << "\n");
} else {
// Global type rewriting
DAETypeUpdater typeUpdater(*this);
auto [oldTypes, newTypes] = typeUpdater.rebuildTypes();

// Global type rewriting
DAETypeUpdater typeUpdater(*this);
auto [oldTypes, newTypes] = typeUpdater.rebuildTypes();
TIME(std::cerr << " rebuild types: " << timer.lastElapsed() << "\n");

TIME(std::cerr << " rebuild types: " << timer.lastElapsed() << "\n");
makeUnreferencedFunctionTypes(oldTypes, newTypes);

makeUnreferencedFunctionTypes(oldTypes, newTypes);
TIME(std::cerr << " make replacements: " << timer.lastElapsed() << "\n");

TIME(std::cerr << " make replacements: " << timer.lastElapsed() << "\n");
optimizer.run(getPassRunner(), wasm);

TIME(std::cerr << " optimize: " << timer.lastElapsed() << "\n");

optimizer.run(getPassRunner(), wasm);
// Update the types for referenced functions and all the locations that
// might hold them. The types of non-referenced functions have already been
// updated separately.
typeUpdater.mapTypes(newTypes);

TIME(std::cerr << " optimize: " << timer.lastElapsed() << "\n");
TIME(std::cerr << " update types: " << timer.lastElapsed() << "\n");
}

// Update the types for referenced functions and all the locations that might
// hold them. The types of non-referenced functions have already been updated
// separately.
typeUpdater.mapTypes(newTypes);
// Refinalize after global type rewriting rather than in
// Optimizer::runOnFunction so that indirect calls whose results were removed
// do not have their types reset to the original signature results, and so
// that ref.func instructions see the final rewritten function types rather
// than temporary scratch types (or race on concurrent function type updates).
Refinalizer(*this).run(getPassRunner(), wasm);

TIME(std::cerr << " update types: " << timer.lastElapsed() << "\n");
TIME(std::cerr << " refinalize: " << timer.lastElapsed() << "\n");
}

// Generate new (possibly optimized) types for each unreferenced function so
Expand Down
160 changes: 160 additions & 0 deletions test/lit/passes/dae2-results-control-flow.wast
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -279,3 +279,163 @@
(global.set $g (call $test))
)
)

(module
;; CHECK: (func $callee (type $0)
;; CHECK-NEXT: (unreachable)
;; CHECK-NEXT: )
(func $callee (result f32)
(unreachable)
)
;; CHECK: (func $caller (type $2) (param $x i32)
;; CHECK-NEXT: (drop
;; CHECK-NEXT: (block
;; CHECK-NEXT: (if
;; CHECK-NEXT: (local.get $x)
;; CHECK-NEXT: (then
;; CHECK-NEXT: (unreachable)
;; CHECK-NEXT: (call $callee)
;; CHECK-NEXT: )
;; CHECK-NEXT: (else
;; CHECK-NEXT: (unreachable)
;; CHECK-NEXT: )
;; CHECK-NEXT: )
;; CHECK-NEXT: (unreachable)
;; CHECK-NEXT: )
;; CHECK-NEXT: )
;; CHECK-NEXT: )
(func $caller (param $x i32)
;; Removing the unused result of the inner block makes it unreachable, which
;; in turn makes the enclosing if unreachable. Ensure we refinalize the
;; function so enclosing expressions have updated types.
(drop
(block (result i32)
(if
(local.get $x)
(then
(drop
(block (result f32)
(unreachable)
(call $callee)
)
)
)
(else
(unreachable)
)
)
(unreachable)
)
)
)
)

(module
;; CHECK: (type $sig (func))
(type $sig (func (result f32)))
;; CHECK: (func $callee (type $0)
;; CHECK-NEXT: (unreachable)
;; CHECK-NEXT: )
(func $callee (result f32)
(unreachable)
)
;; CHECK: (func $caller (type $2) (param $x i32) (param $ref (ref $sig))
;; CHECK-NEXT: (call_ref $sig
;; CHECK-NEXT: (local.get $ref)
;; CHECK-NEXT: )
;; CHECK-NEXT: (drop
;; CHECK-NEXT: (block
;; CHECK-NEXT: (if
;; CHECK-NEXT: (local.get $x)
;; CHECK-NEXT: (then
;; CHECK-NEXT: (unreachable)
;; CHECK-NEXT: (call $callee)
;; CHECK-NEXT: )
;; CHECK-NEXT: (else
;; CHECK-NEXT: (unreachable)
;; CHECK-NEXT: )
;; CHECK-NEXT: )
;; CHECK-NEXT: (unreachable)
;; CHECK-NEXT: )
;; CHECK-NEXT: )
;; CHECK-NEXT: )
(func $caller (param $x i32) (param $ref (ref $sig))
;; Refinalizing must happen after global type rewriting so that the call_ref
;; whose result was removed does not have its type reset to the old concrete
;; result type.
(drop
(call_ref $sig
(local.get $ref)
)
)
(drop
(block (result i32)
(if
(local.get $x)
(then
(drop
(block (result f32)
(unreachable)
(call $callee)
)
)
)
(else
(unreachable)
)
)
(unreachable)
)
)
)
)

(module
;; CHECK: (rec
;; CHECK-NEXT: (type $super (sub (func)))
(type $super (sub (func (result f32))))
;; CHECK: (type $sub (sub $super (func)))
(type $sub (sub $super (func (result f32))))
;; CHECK: (func $f (type $sub)
;; CHECK-NEXT: (local $x (ref null $super))
;; CHECK-NEXT: (drop
;; CHECK-NEXT: (block
;; CHECK-NEXT: (local.set $x
;; CHECK-NEXT: (ref.func $f)
;; CHECK-NEXT: )
;; CHECK-NEXT: (if
;; CHECK-NEXT: (i32.const 0)
;; CHECK-NEXT: (then
;; CHECK-NEXT: (unreachable)
;; CHECK-NEXT: (call $f)
;; CHECK-NEXT: )
;; CHECK-NEXT: (else
;; CHECK-NEXT: (unreachable)
;; CHECK-NEXT: )
;; CHECK-NEXT: )
;; CHECK-NEXT: (unreachable)
;; CHECK-NEXT: )
;; CHECK-NEXT: )
;; CHECK-NEXT: )
(func $f (type $sub) (result f32)
;; Refinalizing must happen after global type rewriting so that the ref.func
;; gets the final rewritten subtype rather than a temporary scratch type.
Comment on lines +421 to +422

@aheejin aheejin Oct 24, 2026

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

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

It would be nice to elaborate what "scratch type" here is and why it is necessary. (Grepping the codebase with "scratch type" doesn't yield anything.
(Gemini told me it was this

binaryen/src/passes/DeadArgumentElimination2.cpp

Lines 1085 to 1090 in f1f3a48

// Replace the function's type with a new type using only the kept
// parameters and result. This ensures that newly added locals get the
// right indices. Preserve sharedness in case this becomes the permanent
// new type (when GC is disabled, see below).
Signature sig(Type(keptParams),
funcInfo->resultUsage ? func->getResults() : Type::none);

But it would have been nice if could be read without Gemini's help when you don't have a detailed knowledge about the pass)

Also it would be nice if the comment describes what goes wrong if refinalizing doesn't happen at the end in this test.

(local $x (ref null $super))
(local.set $x (ref.func $f))
(if
(i32.const 0)
(then
(drop
(block (result f32)
(unreachable)
(call $f)
)
)
)
(else
(unreachable)
)
)
(unreachable)
)
)
Loading