diff options
4 files changed, 317 insertions, 25 deletions
diff --git a/src/passes/RemoveUnusedModuleElements.cpp b/src/passes/RemoveUnusedModuleElements.cpp index 8f8fc24a3..2466d6252 100644 --- a/src/passes/RemoveUnusedModuleElements.cpp +++ b/src/passes/RemoveUnusedModuleElements.cpp @@ -26,6 +26,7 @@ #include "ir/module-utils.h" #include "ir/utils.h" #include "pass.h" +#include "wasm-builder.h" #include "wasm.h" namespace wasm { @@ -43,6 +44,27 @@ struct ReachabilityAnalyzer : public PostWalker<ReachabilityAnalyzer> { std::set<ModuleElement> reachable; bool usesMemory = false; + // The signatures that we have seen a call_ref for. When we see a RefFunc of a + // signature in here, we know it is reachable. + std::unordered_set<HeapType> calledSignatures; + + // All the RefFuncs we've seen, grouped by heap type. When we see a CallRef of + // one of the types here, we know all the RefFuncs corresponding to it are + // reachable. This is the reverse side of calledSignatures: for a function to + // be reached via a reference, we need the combination of a RefFunc of it as + // well as a CallRef of that, and we may see them in any order. (Or, if the + // RefFunc is in a table, we need a CallIndirect, which is handled in the + // table logic.) + // + // After we see a call for a type, we can clear out the entry here for it, as + // we'll have that type in calledSignatures, and so this contains only + // RefFuncs that we have not seen a call for yet, hence "uncalledRefFuncMap." + // + // TODO: We assume a closed world in the GC space atm, but eventually should + // have a flag for that, and when the world is not closed we'd need to + // check for RefFuncs that flow out to exports. + std::unordered_map<HeapType, std::vector<Name>> uncalledRefFuncMap; + ReachabilityAnalyzer(Module* module, const std::vector<ModuleElement>& roots) : module(module) { queue = roots; @@ -105,6 +127,33 @@ struct ReachabilityAnalyzer : public PostWalker<ReachabilityAnalyzer> { } void visitCallIndirect(CallIndirect* curr) { maybeAddTable(curr->table); } + void visitCallRef(CallRef* curr) { + // Ignore unreachable code. + if (!curr->target->type.isRef()) { + return; + } + + auto type = curr->target->type.getHeapType(); + + // Call all the functions of that signature. We can then forget about + // them, as this signature will be marked as called. + auto iter = uncalledRefFuncMap.find(type); + if (iter != uncalledRefFuncMap.end()) { + // We must not have a type in both calledSignatures and + // uncalledRefFuncMap: once it is called, we do not track RefFuncs for + // it any more. + assert(calledSignatures.count(type) == 0); + + for (Name target : iter->second) { + maybeAdd(ModuleElement(ModuleElementKind::Function, target)); + } + + uncalledRefFuncMap.erase(iter); + } + + calledSignatures.insert(type); + } + void visitGlobalGet(GlobalGet* curr) { maybeAdd(ModuleElement(ModuleElementKind::Global, curr->name)); } @@ -126,7 +175,19 @@ struct ReachabilityAnalyzer : public PostWalker<ReachabilityAnalyzer> { void visitMemorySize(MemorySize* curr) { usesMemory = true; } void visitMemoryGrow(MemoryGrow* curr) { usesMemory = true; } void visitRefFunc(RefFunc* curr) { - maybeAdd(ModuleElement(ModuleElementKind::Function, curr->func)); + auto type = curr->type.getHeapType(); + if (calledSignatures.count(type)) { + // We must not have a type in both calledSignatures and + // uncalledRefFuncMap: once it is called, we do not track RefFuncs for it + // any more. + assert(uncalledRefFuncMap.count(type) == 0); + + // We've seen a RefFunc for this, so it is reachable. + maybeAdd(ModuleElement(ModuleElementKind::Function, curr->func)); + } else { + // We've never seen a CallRef for this, but might see one later. + uncalledRefFuncMap[type].push_back(curr->func); + } } void visitTableGet(TableGet* curr) { maybeAddTable(curr->table); } void visitTableSet(TableSet* curr) { maybeAddTable(curr->table); } @@ -199,15 +260,42 @@ struct RemoveUnusedModuleElements : public Pass { importsMemory = true; } // For now, all functions that can be called indirectly are marked as roots. + // TODO: Compute this based on which ElementSegments are actually reachable, + // and which functions have a call_indirect of the proper type. ElementUtils::iterAllElementFunctionNames(module, [&](Name& name) { roots.emplace_back(ModuleElementKind::Function, name); }); // Compute reachability starting from the root set. ReachabilityAnalyzer analyzer(module, roots); + + // RefFuncs that are never called are a special case: We cannot remove the + // function, since then (ref.func $foo) would not validate. But if we know + // it is never called, at least the contents do not matter, so we can + // empty it out. + std::unordered_set<Name> uncalledRefFuncs; + for (auto& [type, targets] : analyzer.uncalledRefFuncMap) { + for (auto target : targets) { + uncalledRefFuncs.insert(target); + } + } + // Remove unreachable elements. module->removeFunctions([&](Function* curr) { - return analyzer.reachable.count( - ModuleElement(ModuleElementKind::Function, curr->name)) == 0; + if (analyzer.reachable.count( + ModuleElement(ModuleElementKind::Function, curr->name))) { + return false; + } + + if (uncalledRefFuncs.count(curr->name)) { + // See comment above on uncalledRefFuncs. + if (!curr->imported()) { + curr->body = Builder(*module).makeUnreachable(); + } + return false; + } + + // The function is not reached and has no references; remove it. + return true; }); module->removeGlobals([&](Global* curr) { return analyzer.reachable.count( diff --git a/test/lit/passes/remove-unused-module-elements-refs.wast b/test/lit/passes/remove-unused-module-elements-refs.wast new file mode 100644 index 000000000..a31854d43 --- /dev/null +++ b/test/lit/passes/remove-unused-module-elements-refs.wast @@ -0,0 +1,226 @@ +;; NOTE: Assertions have been generated by update_lit_checks.py --all-items and should not be edited. +;; RUN: foreach %s %t wasm-opt --remove-unused-module-elements --nominal -all -S -o - | filecheck %s + +(module + ;; CHECK: (type $A (func_subtype func)) + (type $A (func)) + ;; CHECK: (type $B (func_subtype func)) + (type $B (func)) + + ;; CHECK: (elem declare func $target-A $target-B) + + ;; CHECK: (export "foo" (func $foo)) + + ;; CHECK: (func $foo (type $A) + ;; CHECK-NEXT: (drop + ;; CHECK-NEXT: (ref.func $target-A) + ;; CHECK-NEXT: ) + ;; CHECK-NEXT: (drop + ;; CHECK-NEXT: (ref.func $target-B) + ;; CHECK-NEXT: ) + ;; CHECK-NEXT: (call_ref + ;; CHECK-NEXT: (ref.null $A) + ;; CHECK-NEXT: ) + ;; CHECK-NEXT: (block + ;; CHECK-NEXT: (unreachable) + ;; CHECK-NEXT: ) + ;; CHECK-NEXT: ) + (func $foo (export "foo") + ;; This export has two RefFuncs, and one CallRef. + (drop + (ref.func $target-A) + ) + (drop + (ref.func $target-B) + ) + (call_ref + (ref.null $A) + ) + ;; Verify that we do not crash on an unreachable call_ref, which has no + ;; heap type for us to analyze. + (call_ref + (unreachable) + ) + ) + + ;; CHECK: (func $target-A (type $A) + ;; CHECK-NEXT: (nop) + ;; CHECK-NEXT: ) + (func $target-A (type $A) + ;; This function is reachable from the export "foo": there is a RefFunc and + ;; a CallRef for it there. + ) + + (func $target-A-noref (type $A) + ;; This function is not reachable. We have a CallRef of the right type, but + ;; no RefFunc. + ) + + ;; CHECK: (func $target-B (type $B) + ;; CHECK-NEXT: (unreachable) + ;; CHECK-NEXT: ) + (func $target-B (type $B) + ;; This function is not reachable. We have a RefFunc in "foo" but no + ;; suitable CallRef. + ;; + ;; Note that we cannot remove the function, as the RefFunc must refer to + ;; something in order to validate. But we can clear out the body of this + ;; function with an unreachable. + ) +) + +;; As above, but reverse the order inside $foo, so we see the CallRef first. +(module + ;; CHECK: (type $A (func_subtype func)) + (type $A (func)) + (type $B (func)) + + ;; CHECK: (elem declare func $target-A) + + ;; CHECK: (export "foo" (func $foo)) + + ;; CHECK: (func $foo (type $A) + ;; CHECK-NEXT: (call_ref + ;; CHECK-NEXT: (ref.null $A) + ;; CHECK-NEXT: ) + ;; CHECK-NEXT: (drop + ;; CHECK-NEXT: (ref.func $target-A) + ;; CHECK-NEXT: ) + ;; CHECK-NEXT: ) + (func $foo (export "foo") + (call_ref + (ref.null $A) + ) + (drop + (ref.func $target-A) + ) + ) + + ;; CHECK: (func $target-A (type $A) + ;; CHECK-NEXT: (nop) + ;; CHECK-NEXT: ) + (func $target-A (type $A) + ;; This function is reachable. + ) + + (func $target-A-noref (type $A) + ;; This function is not reachable. + ) +) + +;; As above, but interleave CallRefs with RefFuncs. +(module + ;; CHECK: (type $A (func_subtype func)) + (type $A (func)) + (type $B (func)) + + ;; CHECK: (elem declare func $target-A-1 $target-A-2) + + ;; CHECK: (export "foo" (func $foo)) + + ;; CHECK: (func $foo (type $A) + ;; CHECK-NEXT: (call_ref + ;; CHECK-NEXT: (ref.null $A) + ;; CHECK-NEXT: ) + ;; CHECK-NEXT: (drop + ;; CHECK-NEXT: (ref.func $target-A-1) + ;; CHECK-NEXT: ) + ;; CHECK-NEXT: (call_ref + ;; CHECK-NEXT: (ref.null $A) + ;; CHECK-NEXT: ) + ;; CHECK-NEXT: (drop + ;; CHECK-NEXT: (ref.func $target-A-2) + ;; CHECK-NEXT: ) + ;; CHECK-NEXT: ) + (func $foo (export "foo") + (call_ref + (ref.null $A) + ) + (drop + (ref.func $target-A-1) + ) + (call_ref + (ref.null $A) + ) + (drop + (ref.func $target-A-2) + ) + ) + + ;; CHECK: (func $target-A-1 (type $A) + ;; CHECK-NEXT: (nop) + ;; CHECK-NEXT: ) + (func $target-A-1 (type $A) + ;; This function is reachable. + ) + + ;; CHECK: (func $target-A-2 (type $A) + ;; CHECK-NEXT: (nop) + ;; CHECK-NEXT: ) + (func $target-A-2 (type $A) + ;; This function is reachable. + ) + + (func $target-A-3 (type $A) + ;; This function is not reachable. + ) +) + +;; As above, with the order reversed inside $foo. The results should be the +;; same. +(module + ;; CHECK: (type $A (func_subtype func)) + (type $A (func)) + (type $B (func)) + + ;; CHECK: (elem declare func $target-A-1 $target-A-2) + + ;; CHECK: (export "foo" (func $foo)) + + ;; CHECK: (func $foo (type $A) + ;; CHECK-NEXT: (drop + ;; CHECK-NEXT: (ref.func $target-A-1) + ;; CHECK-NEXT: ) + ;; CHECK-NEXT: (call_ref + ;; CHECK-NEXT: (ref.null $A) + ;; CHECK-NEXT: ) + ;; CHECK-NEXT: (drop + ;; CHECK-NEXT: (ref.func $target-A-2) + ;; CHECK-NEXT: ) + ;; CHECK-NEXT: (call_ref + ;; CHECK-NEXT: (ref.null $A) + ;; CHECK-NEXT: ) + ;; CHECK-NEXT: ) + (func $foo (export "foo") + (drop + (ref.func $target-A-1) + ) + (call_ref + (ref.null $A) + ) + (drop + (ref.func $target-A-2) + ) + (call_ref + (ref.null $A) + ) + ) + + ;; CHECK: (func $target-A-1 (type $A) + ;; CHECK-NEXT: (nop) + ;; CHECK-NEXT: ) + (func $target-A-1 (type $A) + ;; This function is reachable. + ) + + ;; CHECK: (func $target-A-2 (type $A) + ;; CHECK-NEXT: (nop) + ;; CHECK-NEXT: ) + (func $target-A-2 (type $A) + ;; This function is reachable. + ) + + (func $target-A-3 (type $A) + ;; This function is not reachable. + ) +) diff --git a/test/passes/remove-unused-module-elements_all-features.txt b/test/passes/remove-unused-module-elements_all-features.txt index 864624c7e..c2806d6e4 100644 --- a/test/passes/remove-unused-module-elements_all-features.txt +++ b/test/passes/remove-unused-module-elements_all-features.txt @@ -295,16 +295,3 @@ ) ) ) -(module - (type $none_=>_none (func)) - (elem declare func $foo) - (export "test" (func $test)) - (func $foo - (nop) - ) - (func $test - (drop - (ref.func $foo) - ) - ) -) diff --git a/test/passes/remove-unused-module-elements_all-features.wast b/test/passes/remove-unused-module-elements_all-features.wast index 09f3ab0e3..fbccda8b8 100644 --- a/test/passes/remove-unused-module-elements_all-features.wast +++ b/test/passes/remove-unused-module-elements_all-features.wast @@ -286,12 +286,3 @@ ) ) ) -(module ;; functions referenced by ref.func cannot be removed - (export "test" $test) - (func $foo) - (func $test - (drop - (ref.func $foo) - ) - ) -) |