Skip to content

[Custom Descriptors] Optimize RefCast to RefCastDesc in GlobalStructInference #7906

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

Draft
wants to merge 10 commits into
base: main
Choose a base branch
from
Draft

[Custom Descriptors] Optimize RefCast to RefCastDesc in GlobalStructInference #7906

Show file tree
Hide file tree
Changes from 9 commits
Commits
Show all changes
10 commits
Select commit Hold shift + click to select a range
ccc5be0
start
kripken Sep 18, 2025
f4b9e45
test
kripken Sep 18, 2025
7691673
test
kripken Sep 18, 2025
10d80a7
test
kripken Sep 18, 2025
8bebd6b
test
kripken Sep 18, 2025
4475e50
test
kripken Sep 18, 2025
891f616
test
kripken Sep 18, 2025
c3d3c5a
test
kripken Sep 18, 2025
46b37d3
test
kripken Sep 18, 2025
9a9c122
TODO for exact casts
kripken Sep 19, 2025
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
54 changes: 54 additions & 0 deletions src/passes/GlobalStructInference.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -79,6 +79,8 @@ struct GlobalStructInference : public Pass {
// optimizable it will have an entry here, and not if not.
std::unordered_map<HeapType, std::vector<Name>> typeGlobals;

std::unique_ptr<SubTypes> subTypes;

void run(Module* module) override {
if (!module->features.hasGC()) {
return;
Expand Down Expand Up @@ -208,6 +210,12 @@ struct GlobalStructInference : public Pass {
return;
}

// When CD is enabled, we can optimize to ref.get_desc, depending on the
// presence of subtypes.
if (module->features.hasCustomDescriptors()) {
subTypes = std::make_unique<SubTypes>(*module);
}

// The above loop on typeGlobalsCopy is on an unsorted data structure, and
// that can lead to nondeterminism in typeGlobals. Sort the vectors there to
// ensure determinism.
Expand Down Expand Up @@ -528,6 +536,52 @@ struct GlobalStructInference : public Pass {
right));
}

void visitRefCast(RefCast* curr) {
// When we see (ref.cast $T), and the type has a descriptor, and that
// desceriptor only has a single global, then we can do (ref.cast_desc)
// using the descriptor. Descriptor XXX
// casts are usually more efficient than normal ones (and even more so
// if we get lucky and are in a loop, where the global.get of the
// descriptor can be hoisted).
// TODO: only do this when shrinkLevel == 0?

// Check if we have a descriptor.
auto type = curr->type;
if (type == Type::unreachable) {
return;
}
auto heapType = type.getHeapType();
auto desc = heapType.getDescriptorType();
if (!desc) {
return;
}

// Check if the type has no subtypes, as a ref.cast_desc will find
// precisely that type and nothing else. TODO: exact types too
if (!parent.subTypes->getStrictSubTypes(heapType).empty()) {
return;
}

// Check if we have a single global for the descriptor.
auto iter = parent.typeGlobals.find(*desc);
if (iter == parent.typeGlobals.end()) {
return;
}
const auto& globals = iter->second;
if (globals.size() != 1) {
return;
}

// We can optimize!
auto global = globals[0];
auto& wasm = *getModule();
Builder builder(wasm);
auto* getGlobal =
builder.makeGlobalGet(global, wasm.getGlobal(global)->type);
auto* castDesc = builder.makeRefCast(curr->ref, getGlobal, curr->type);
replaceCurrent(castDesc);
}

void visitFunction(Function* func) {
if (refinalize) {
ReFinalize().walkFunctionInModule(func, getModule());
Expand Down
291 changes: 291 additions & 0 deletions test/lit/passes/gsi-desc.wast
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -186,3 +186,294 @@
)
)

;; Two types with descriptors and subtyping between them.
(module
(rec
;; CHECK: (rec
;; CHECK-NEXT: (type $A (sub (descriptor $A.desc (struct))))
(type $A (sub (descriptor $A.desc (struct))))
;; CHECK: (type $A.desc (sub (describes $A (struct))))
(type $A.desc (sub (describes $A (struct))))

;; CHECK: (type $B (sub $A (descriptor $B.desc (struct))))
(type $B (sub $A (descriptor $B.desc (struct))))
;; CHECK: (type $B.desc (sub $A.desc (describes $B (struct))))
(type $B.desc (sub $A.desc (describes $B (struct))))
)

;; CHECK: (type $4 (func (param anyref)))

;; CHECK: (global $A.desc (ref $A.desc) (struct.new_default $A.desc))
(global $A.desc (ref $A.desc) (struct.new $A.desc))

;; CHECK: (global $B.desc (ref $B.desc) (struct.new_default $B.desc))
(global $B.desc (ref $B.desc) (struct.new $B.desc))

;; CHECK: (func $test (type $4) (param $any anyref)
;; CHECK-NEXT: (drop
;; CHECK-NEXT: (ref.cast (ref $A)
;; CHECK-NEXT: (local.get $any)
;; CHECK-NEXT: )
;; CHECK-NEXT: )
;; CHECK-NEXT: (drop
;; CHECK-NEXT: (ref.cast_desc (ref $B)
;; CHECK-NEXT: (local.get $any)
;; CHECK-NEXT: (global.get $B.desc)
;; CHECK-NEXT: )
;; CHECK-NEXT: )
;; CHECK-NEXT: )
(func $test (param $any anyref)
;; The second cast here is optimizable: it can only be to a single type with
;; no subtypes, so we can use ref.cast_desc.
(drop
(ref.cast (ref $A)
(local.get $any)
)
)
(drop
(ref.cast (ref $B)
(local.get $any)
)
)
)
)

;; As above, but without subtyping between $A and $B.
(module
(rec
;; CHECK: (rec
;; CHECK-NEXT: (type $A (sub (descriptor $A.desc (struct))))
(type $A (sub (descriptor $A.desc (struct))))
;; CHECK: (type $A.desc (sub (describes $A (struct))))
(type $A.desc (sub (describes $A (struct))))

;; CHECK: (type $B (sub (descriptor $B.desc (struct))))
(type $B (sub (descriptor $B.desc (struct))))
;; CHECK: (type $B.desc (sub $A.desc (describes $B (struct))))
(type $B.desc (sub $A.desc (describes $B (struct))))
)

;; CHECK: (type $4 (func (param anyref)))

;; CHECK: (global $A.desc (ref $A.desc) (struct.new_default $A.desc))
(global $A.desc (ref $A.desc) (struct.new $A.desc))

;; CHECK: (global $B.desc (ref $B.desc) (struct.new_default $B.desc))
(global $B.desc (ref $B.desc) (struct.new $B.desc))

;; CHECK: (func $test (type $4) (param $any anyref)
;; CHECK-NEXT: (drop
;; CHECK-NEXT: (ref.cast (ref $A)
;; CHECK-NEXT: (local.get $any)
;; CHECK-NEXT: )
;; CHECK-NEXT: )
;; CHECK-NEXT: (drop
;; CHECK-NEXT: (ref.cast_desc (ref $B)
;; CHECK-NEXT: (local.get $any)
;; CHECK-NEXT: (global.get $B.desc)
;; CHECK-NEXT: )
;; CHECK-NEXT: )
;; CHECK-NEXT: )
(func $test (param $any anyref)
;; We still cannot optimize $A: while $A has no subtypes, the descriptor
;; $A.desc has a subtype. We could optimize this TODO
(drop
(ref.cast (ref $A)
(local.get $any)
)
)
(drop
(ref.cast (ref $B)
(local.get $any)
)
)
)
)

;; As above, but without subtyping between $A.desc and $B.desc.
(module
(rec
;; CHECK: (rec
;; CHECK-NEXT: (type $A (sub (descriptor $A.desc (struct))))
(type $A (sub (descriptor $A.desc (struct))))
;; CHECK: (type $A.desc (sub (describes $A (struct))))
(type $A.desc (sub (describes $A (struct))))

;; CHECK: (type $B (sub (descriptor $B.desc (struct))))
(type $B (sub (descriptor $B.desc (struct))))
;; CHECK: (type $B.desc (sub (describes $B (struct))))
(type $B.desc (sub (describes $B (struct))))
)

;; CHECK: (type $4 (func (param anyref)))

;; CHECK: (global $A.desc (ref $A.desc) (struct.new_default $A.desc))
(global $A.desc (ref $A.desc) (struct.new $A.desc))

;; CHECK: (global $B.desc (ref $B.desc) (struct.new_default $B.desc))
(global $B.desc (ref $B.desc) (struct.new $B.desc))

;; CHECK: (func $test (type $4) (param $any anyref)
;; CHECK-NEXT: (drop
;; CHECK-NEXT: (ref.cast_desc (ref $A)
;; CHECK-NEXT: (local.get $any)
;; CHECK-NEXT: (global.get $A.desc)
;; CHECK-NEXT: )
;; CHECK-NEXT: )
;; CHECK-NEXT: (drop
;; CHECK-NEXT: (ref.cast_desc (ref $B)
;; CHECK-NEXT: (local.get $any)
;; CHECK-NEXT: (global.get $B.desc)
;; CHECK-NEXT: )
;; CHECK-NEXT: )
;; CHECK-NEXT: )
(func $test (param $any anyref)
;; We can fully optimize these two independent cases.
(drop
(ref.cast (ref $A)
(local.get $any)
)
)
(drop
(ref.cast (ref $B)
(local.get $any)
)
)
)
)

;; Zero descriptor instances in globals.
(module
(rec
;; CHECK: (rec
;; CHECK-NEXT: (type $A (sub (descriptor $A.desc (struct))))
(type $A (sub (descriptor $A.desc (struct))))
;; CHECK: (type $A.desc (sub (describes $A (struct))))
(type $A.desc (sub (describes $A (struct))))
)

;; CHECK: (type $2 (func (param anyref)))

;; CHECK: (func $test (type $2) (param $any anyref)
;; CHECK-NEXT: (drop
;; CHECK-NEXT: (ref.cast (ref $A)
;; CHECK-NEXT: (local.get $any)
;; CHECK-NEXT: )
;; CHECK-NEXT: )
;; CHECK-NEXT: )
(func $test (param $any anyref)
;; We do not optimize here. TODO: we could make this trap
(drop
(ref.cast (ref $A)
(local.get $any)
)
)
)
)

;; Two descriptor instances in globals.
(module
(rec
;; CHECK: (rec
;; CHECK-NEXT: (type $A (sub (descriptor $A.desc (struct))))
(type $A (sub (descriptor $A.desc (struct))))
;; CHECK: (type $A.desc (sub (describes $A (struct))))
(type $A.desc (sub (describes $A (struct))))
)

;; CHECK: (type $2 (func (param anyref)))

;; CHECK: (global $A.desc (ref $A.desc) (struct.new_default $A.desc))
(global $A.desc (ref $A.desc) (struct.new $A.desc))

;; CHECK: (global $A.desc2 (ref $A.desc) (struct.new_default $A.desc))
(global $A.desc2 (ref $A.desc) (struct.new $A.desc))

;; CHECK: (func $test (type $2) (param $any anyref)
;; CHECK-NEXT: (drop
;; CHECK-NEXT: (ref.cast (ref $A)
;; CHECK-NEXT: (local.get $any)
;; CHECK-NEXT: )
;; CHECK-NEXT: )
;; CHECK-NEXT: (drop
;; CHECK-NEXT: (block ;; (replaces unreachable RefCast we can't emit)
;; CHECK-NEXT: (drop
;; CHECK-NEXT: (unreachable)
;; CHECK-NEXT: )
;; CHECK-NEXT: (unreachable)
;; CHECK-NEXT: )
;; CHECK-NEXT: )
;; CHECK-NEXT: )
(func $test (param $any anyref)
;; We do not optimize here. TODO: we could with a select
(drop
(ref.cast (ref $A)
(local.get $any)
)
)
;; We do not error on unreachable casts.
(drop
(ref.cast (ref $A)
(unreachable)
)
)
)
)

(module
;; CHECK: (type $0 (func (param anyref)))

;; CHECK: (type $A (sub (struct)))
(type $A (sub (struct)))

;; CHECK: (func $test (type $0) (param $any anyref)
;; CHECK-NEXT: (drop
;; CHECK-NEXT: (ref.cast (ref $A)
;; CHECK-NEXT: (local.get $any)
;; CHECK-NEXT: )
;; CHECK-NEXT: )
;; CHECK-NEXT: )
(func $test (param $any anyref)
;; We do not handle casts to types without descriptors.
(drop
(ref.cast (ref $A)
(local.get $any)
)
)
)
)

;; Nullable cast.
(module
(rec
;; CHECK: (rec
;; CHECK-NEXT: (type $A (sub (descriptor $A.desc (struct))))
(type $A (sub (descriptor $A.desc (struct))))
;; CHECK: (type $A.desc (sub (describes $A (struct))))
(type $A.desc (sub (describes $A (struct))))
)

;; CHECK: (type $2 (func (param anyref)))

;; CHECK: (global $A.desc (ref $A.desc) (struct.new_default $A.desc))
(global $A.desc (ref $A.desc) (struct.new $A.desc))

;; CHECK: (func $test (type $2) (param $any anyref)
;; CHECK-NEXT: (drop
;; CHECK-NEXT: (ref.cast_desc (ref null $A)
;; CHECK-NEXT: (local.get $any)
;; CHECK-NEXT: (global.get $A.desc)
;; CHECK-NEXT: )
;; CHECK-NEXT: )
;; CHECK-NEXT: )
(func $test (param $any anyref)
;; The cast is nullable, which we can still optimize: null will succeed as
;; expected.
(drop
(ref.cast (ref null $A)
(local.get $any)
)
)
)
)

Loading