Skip to content

Commit 97986a6

Browse files
aheejinaheejin
committed
[wasm-split] Pre-create trampolines for global initializers
Before #8443, we scanned `ref.func`s in global initizliers early in `indirectReferencesToSecondaryFunctions` and created trampolines for them and replaced `ref.func $func`s with `ref.func $trampoline_func` if `func` was set to move to a secondary module. But in case the global containing `ref.func $trampoline_func` also ends up moving to the same secondary module, creating trampoline and using it was not necessary, because the global can simply use `ref.func $func` because `func` is in the same secondary module. To fix this, in #8443, we postponed creating trampolines for `ref.func`s in global initializers until `shareImportableItems`. This had a problem, because we end up creating new trampolines late in `shareImportableItems`. But trampolines were designed to go through `indirectCallsToSecondaryFunctions` and `setupTablePatching`, so those late trampolines were invalid, like ```wast (func $trampoline_foo (call $foo) ) ``` when `foo` is in a secondary module. This was supposed to be converted to a `call_indirect` in `indirectCallsToSecondaryFunctions` and the table elements were supposed to set up in `setupTablePatching`. --- To fix the issue, this PR pre-creates trampolines for `ref.func`s in all global initializers in `indirectReferencesToSecondaryFunctions`, before `indirectCallsToSecondaryFunctions` and `setupTablePatching`, but doesn't make `ref.func`s reference them just yet. We make `ref.func`s in global initializers reference the trampolines only when it is decided they are not moving to the same secondary module, in `shareImportableItems`. But this can leave unused trampolines in the primary module. So in `cleanupUnusedTrampolines`, we remove unused trampolines. This increases acx_gallery code size only by 0.6%. If we didn't clean up the unnecessary trampolines, the increase was 3.6%. This only removes unused trampoline functions and not placeholder imports and elements yet, but they don't affect acx_gallery because it uses `--no-placeholders`. Fixes #8510. (Both tests fail for the same reason)
1 parent 6d47684 commit 97986a6

5 files changed

Lines changed: 176 additions & 48 deletions

File tree

‎src/ir/module-splitting.cpp‎

Lines changed: 69 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -73,7 +73,6 @@
7373
// from the IR before splitting.
7474
//
7575
#include "ir/module-splitting.h"
76-
#include "ir/export-utils.h"
7776
#include "ir/find_all.h"
7877
#include "ir/module-utils.h"
7978
#include "ir/names.h"
@@ -336,6 +335,7 @@ struct ModuleSplitter {
336335
void exportImportCalledPrimaryFunctions();
337336
void setupTablePatching();
338337
void shareImportableItems();
338+
void cleanupUnusedTrampolines();
339339

340340
ModuleSplitter(Module& primary, const Config& config)
341341
: config(config), primary(primary), tableManager(primary),
@@ -348,6 +348,7 @@ struct ModuleSplitter {
348348
exportImportCalledPrimaryFunctions();
349349
setupTablePatching();
350350
shareImportableItems();
351+
cleanupUnusedTrampolines();
351352
}
352353
};
353354

@@ -508,7 +509,9 @@ Name ModuleSplitter::getTrampoline(Name funcName) {
508509
primary, std::string("trampoline_") + funcName.toString());
509510
it->second = trampoline;
510511

511-
// Generate the call and the function.
512+
// Generate the call and the function. We generate a direct call here, but
513+
// this will be converted to a call_indirect in
514+
// indirectCallsToSecondaryFunctions.
512515
std::vector<Expression*> args;
513516
for (Index i = 0; i < oldFunc->getNumParams(); i++) {
514517
args.push_back(builder.makeLocalGet(i, oldFunc->getLocalType(i)));
@@ -559,6 +562,21 @@ static void walkSegments(Walker& walker, Module* module) {
559562
}
560563

561564
void ModuleSplitter::indirectReferencesToSecondaryFunctions() {
565+
// Pre-create trampolines for secondary functions referenced in global
566+
// initializers. We don't mutate the globals yet (that happens later in
567+
// shareImportableItems), but we must create the trampolines now so they
568+
// are processed correctly by indirectCallsToSecondaryFunctions and
569+
// setupTablePatching.
570+
for (auto& global : primary.globals) {
571+
if (global->init) {
572+
for (auto* ref : FindAll<RefFunc>(global->init).list) {
573+
if (allSecondaryFuncs.count(ref->func)) {
574+
getTrampoline(ref->func);
575+
}
576+
}
577+
}
578+
}
579+
562580
// Turn references to secondary functions into references to thunks that
563581
// perform a direct call to the original referent. The direct calls in the
564582
// thunks will be handled like all other cross-module calls later, in
@@ -1222,6 +1240,55 @@ void ModuleSplitter::shareImportableItems() {
12221240
}
12231241
}
12241242

1243+
void ModuleSplitter::cleanupUnusedTrampolines() {
1244+
// We pre-created trampolines for all functions referenced in global
1245+
// initializers' ref.funcs in indirectReferencesToSecondaryFunctions.
1246+
// But after splitting module items, if a trampoline is not used in the
1247+
// primary module, it means it was created for a global initializer but that
1248+
// global ended up moving to the same secondary module as the function
1249+
// referenced in the global initializer's ref.func. We can remove them here.
1250+
// TODO Remove placeholders and table elements created by unused trampolines
1251+
std::unordered_set<Name> usedFuncs;
1252+
1253+
struct FuncrefCollector : public PostWalker<FuncrefCollector> {
1254+
std::vector<Name>& used;
1255+
FuncrefCollector(std::vector<Name>& used) : used(used) {}
1256+
void visitRefFunc(RefFunc* curr) { used.push_back(curr->func); }
1257+
};
1258+
1259+
ModuleUtils::ParallelFunctionAnalysis<std::vector<Name>> analysis(
1260+
primary, [&](Function* func, std::vector<Name>& used) {
1261+
if (!func->imported()) {
1262+
FuncrefCollector(used).walkFunction(func);
1263+
}
1264+
});
1265+
1266+
for (auto& [_, usedList] : analysis.map) {
1267+
usedFuncs.insert(usedList.begin(), usedList.end());
1268+
}
1269+
1270+
std::vector<Name> moduleCodeUsed;
1271+
FuncrefCollector(moduleCodeUsed).walkModuleCode(&primary);
1272+
usedFuncs.insert(moduleCodeUsed.begin(), moduleCodeUsed.end());
1273+
1274+
for (auto& ex : primary.exports) {
1275+
if (ex->kind == ExternalKind::Function) {
1276+
usedFuncs.insert(*ex->getInternalName());
1277+
}
1278+
}
1279+
1280+
std::vector<Name> trampolinesToRemove;
1281+
for (auto& [_, trampoline] : trampolineMap) {
1282+
if (!usedFuncs.count(trampoline)) {
1283+
trampolinesToRemove.push_back(trampoline);
1284+
}
1285+
}
1286+
for (auto& name : trampolinesToRemove) {
1287+
primary.removeFunction(name);
1288+
primaryFuncs.erase(name);
1289+
}
1290+
}
1291+
12251292
} // anonymous namespace
12261293

12271294
Results splitFunctions(Module& primary, const Config& config) {

‎test/lit/wasm-split/global-funcref.wast‎

Lines changed: 0 additions & 43 deletions
This file was deleted.

‎test/lit/wasm-split/global-reffunc.wast‎

Lines changed: 51 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,51 @@
1+
;; RUN: wasm-split %s -all -g -o1 %t.1.wasm -o2 %t.2.wasm --keep-funcs=keep
2+
;; RUN: wasm-dis %t.1.wasm | filecheck %s --check-prefix PRIMARY
3+
;; RUN: wasm-dis %t.2.wasm | filecheck %s --check-prefix SECONDARY
4+
5+
;; When a split global ($a here)'s initializer contains a ref.func of a split
6+
;; function, we should NOT create any trampolines, and the split global should
7+
;; direclty refer to the function.
8+
9+
(module
10+
(global $a funcref (ref.func $split))
11+
(global $b funcref (ref.func $keep))
12+
13+
(func $keep)
14+
15+
(func $split
16+
(drop
17+
(global.get $a)
18+
)
19+
(drop
20+
(global.get $b)
21+
)
22+
)
23+
)
24+
25+
;; PRIMARY: (module
26+
;; PRIMARY-NEXT: (type $0 (func))
27+
;; PRIMARY-NEXT: (import "placeholder.deferred" "0" (func $placeholder_0))
28+
;; PRIMARY-NEXT: (table $0 1 funcref)
29+
;; PRIMARY-NEXT: (elem $0 (i32.const 0) $placeholder_0)
30+
;; PRIMARY-NEXT: (export "table" (table $0))
31+
;; PRIMARY-NEXT: (export "keep" (func $keep))
32+
;; PRIMARY-NEXT: (func $keep
33+
;; PRIMARY-NEXT: )
34+
;; PRIMARY-NEXT: )
35+
36+
;; SECONDARY: (module
37+
;; SECONDARY-NEXT: (type $0 (func))
38+
;; SECONDARY-NEXT: (import "primary" "table" (table $timport$0 1 funcref))
39+
;; SECONDARY-NEXT: (import "primary" "keep" (func $keep (exact)))
40+
;; SECONDARY-NEXT: (global $a funcref (ref.func $split))
41+
;; SECONDARY-NEXT: (global $b funcref (ref.func $keep))
42+
;; SECONDARY-NEXT: (elem $0 (i32.const 0) $split)
43+
;; SECONDARY-NEXT: (func $split
44+
;; SECONDARY-NEXT: (drop
45+
;; SECONDARY-NEXT: (global.get $a)
46+
;; SECONDARY-NEXT: )
47+
;; SECONDARY-NEXT: (drop
48+
;; SECONDARY-NEXT: (global.get $b)
49+
;; SECONDARY-NEXT: )
50+
;; SECONDARY-NEXT: )
51+
;; SECONDARY-NEXT: )

‎test/lit/wasm-split/global-reffunc2.wast‎

Lines changed: 53 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,53 @@
1+
;; RUN: wasm-split %s -all -g -o1 %t.1.wasm -o2 %t.2.wasm --split-funcs=split1,split2
2+
;; RUN: wasm-dis %t.1.wasm | filecheck %s --check-prefix PRIMARY
3+
;; RUN: wasm-dis %t.2.wasm | filecheck %s --check-prefix SECONDARY
4+
5+
;; Global $g1 is used (exported) in the primary module so it can't move, and
6+
;; global $g2 is only used in the secondary module so it will move there.
7+
8+
(module
9+
(global $g1 funcref (ref.func $split1))
10+
(global $g2 funcref (ref.func $split2))
11+
(export "g1" (global $g1))
12+
13+
(func $split1
14+
(unreachable)
15+
)
16+
17+
(func $split2
18+
(drop
19+
(global.get $g2)
20+
)
21+
)
22+
)
23+
24+
;; PRIMARY-NEXT: (module
25+
;; PRIMARY-NEXT: (type $0 (func))
26+
;; PRIMARY-NEXT: (import "placeholder.deferred" "0" (func $placeholder_0))
27+
;; PRIMARY-NEXT: (import "placeholder.deferred" "1" (func $placeholder_1))
28+
;; PRIMARY-NEXT: (global $g1 funcref (ref.func $trampoline_split1))
29+
;; PRIMARY-NEXT: (table $0 2 funcref)
30+
;; PRIMARY-NEXT: (elem $0 (i32.const 0) $placeholder_0 $placeholder_1)
31+
;; PRIMARY-NEXT: (export "g1" (global $g1))
32+
;; PRIMARY-NEXT: (export "table" (table $0))
33+
;; PRIMARY-NEXT: (func $trampoline_split1
34+
;; PRIMARY-NEXT: (call_indirect (type $0)
35+
;; PRIMARY-NEXT: (i32.const 0)
36+
;; PRIMARY-NEXT: )
37+
;; PRIMARY-NEXT: )
38+
;; PRIMARY-NEXT: )
39+
40+
;; SECONDARY-NEXT: (module
41+
;; SECONDARY-NEXT: (type $0 (func))
42+
;; SECONDARY-NEXT: (import "primary" "table" (table $timport$0 2 funcref))
43+
;; SECONDARY-NEXT: (global $g2 funcref (ref.func $split2))
44+
;; SECONDARY-NEXT: (elem $0 (i32.const 0) $split1 $split2)
45+
;; SECONDARY-NEXT: (func $split1
46+
;; SECONDARY-NEXT: (unreachable)
47+
;; SECONDARY-NEXT: )
48+
;; SECONDARY-NEXT: (func $split2
49+
;; SECONDARY-NEXT: (drop
50+
;; SECONDARY-NEXT: (global.get $g2)
51+
;; SECONDARY-NEXT: )
52+
;; SECONDARY-NEXT: )
53+
;; SECONDARY-NEXT: )

‎test/lit/wasm-split/ref.func.wast‎

Lines changed: 3 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -73,7 +73,7 @@
7373

7474
;; SECONDARY: (import "primary" "prime" (func $prime (exact (type $0))))
7575

76-
;; SECONDARY: (elem $0 (i32.const 0) $second-in-table $second)
76+
;; SECONDARY: (elem $0 (i32.const 0) $second $second-in-table)
7777

7878
;; SECONDARY: (elem declare func $prime)
7979

@@ -109,13 +109,13 @@
109109
;; (but we will get a placeholder, as all split-out functions do).
110110
)
111111
)
112-
;; PRIMARY: (func $trampoline_second-in-table (type $0)
112+
;; PRIMARY: (func $trampoline_second (type $0)
113113
;; PRIMARY-NEXT: (call_indirect $1 (type $0)
114114
;; PRIMARY-NEXT: (i32.const 0)
115115
;; PRIMARY-NEXT: )
116116
;; PRIMARY-NEXT: )
117117

118-
;; PRIMARY: (func $trampoline_second (type $0)
118+
;; PRIMARY: (func $trampoline_second-in-table (type $0)
119119
;; PRIMARY-NEXT: (call_indirect $1 (type $0)
120120
;; PRIMARY-NEXT: (i32.const 1)
121121
;; PRIMARY-NEXT: )

0 commit comments

Comments
 (0)