Skip to content

Commit 48be223

Browse files
aheejinaheejin
authored
[wasm-split] Do not create duplicate trampolines (#7807)
When a function is exported and also referenced somewhere, like in an `elem` section or in a `ref.func` instruction, we create a trampoline for it twice, once for its export in `thunkExportedSecondaryFunctions` and again for its reference in `thunkExportedSecondaryFunctions`. This makes it to be generated only once in the case.
1 parent 0e68ed9 commit 48be223

File tree

7 files changed

+132
-109
lines changed

7 files changed

+132
-109
lines changed

‎src/ir/module-splitting.cpp

Lines changed: 33 additions & 31 deletions
Original file line numberDiff line numberDiff line change
@@ -308,6 +308,9 @@ struct ModuleSplitter {
308308
// Internal name of the LOAD_SECONDARY_MODULE function.
309309
Name internalLoadSecondaryModule;
310310

311+
// Map from original secondary function name to its trampoline
312+
std::unordered_map<Name, Name> trampolineMap;
313+
311314
// Initialization helpers
312315
static std::unique_ptr<Module> initSecondary(const Module& primary);
313316
static std::pair<std::set<Name>, std::set<Name>>
@@ -317,6 +320,7 @@ struct ModuleSplitter {
317320
// Other helpers
318321
void exportImportFunction(Name func);
319322
Expression* maybeLoadSecondary(Builder& builder, Expression* callIndirect);
323+
Name getTrampoline(Name funcName);
320324

321325
// Main splitting steps
322326
void setupJSPI();
@@ -517,6 +521,31 @@ void ModuleSplitter::moveSecondaryFunctions() {
517521
}
518522
}
519523

524+
Name ModuleSplitter::getTrampoline(Name funcName) {
525+
auto [it, inserted] = trampolineMap.insert({funcName, Name()});
526+
if (!inserted) {
527+
return it->second;
528+
}
529+
530+
Builder builder(primary);
531+
auto* oldFunc = secondary.getFunction(funcName);
532+
auto trampoline = Names::getValidFunctionName(
533+
primary, std::string("trampoline_") + funcName.toString());
534+
it->second = trampoline;
535+
536+
// Generate the call and the function.
537+
std::vector<Expression*> args;
538+
for (Index i = 0; i < oldFunc->getNumParams(); i++) {
539+
args.push_back(builder.makeLocalGet(i, oldFunc->getLocalType(i)));
540+
}
541+
auto* call = builder.makeCall(funcName, args, oldFunc->getResults());
542+
543+
auto func = builder.makeFunction(trampoline, oldFunc->type, {}, call);
544+
func->hasExplicitName = oldFunc->hasExplicitName;
545+
primary.addFunction(std::move(func));
546+
return trampoline;
547+
}
548+
520549
void ModuleSplitter::thunkExportedSecondaryFunctions() {
521550
// Update exports of secondary functions in the primary module to export
522551
// wrapper functions that indirectly call the secondary functions. We are
@@ -529,21 +558,8 @@ void ModuleSplitter::thunkExportedSecondaryFunctions() {
529558
!secondaryFuncs.count(*ex->getInternalName())) {
530559
continue;
531560
}
532-
Name secondaryFunc = *ex->getInternalName();
533-
if (primary.getFunctionOrNull(secondaryFunc)) {
534-
// We've already created a thunk for this function
535-
continue;
536-
}
537-
auto* func = primary.addFunction(Builder::makeFunction(
538-
secondaryFunc, secondary.getFunction(secondaryFunc)->type, {}));
539-
std::vector<Expression*> args;
540-
Type params = func->getParams();
541-
for (size_t i = 0, size = params.size(); i < size; ++i) {
542-
args.push_back(builder.makeLocalGet(i, params[i]));
543-
}
544-
auto tableSlot = tableManager.getSlot(secondaryFunc, func->type);
545-
func->body = builder.makeCallIndirect(
546-
tableSlot.tableName, tableSlot.makeExpr(primary), args, func->type);
561+
Name trampoline = getTrampoline(*ex->getInternalName());
562+
ex->setInternalName(trampoline);
547563
}
548564
}
549565

@@ -614,24 +630,10 @@ void ModuleSplitter::indirectReferencesToSecondaryFunctions() {
614630
continue;
615631
}
616632

617-
auto* oldFunc = secondary.getFunction(name);
618-
auto newName = Names::getValidFunctionName(
619-
primary, std::string("trampoline_") + name.toString());
620-
621-
// Generate the call and the function.
622-
std::vector<Expression*> args;
623-
for (Index i = 0; i < oldFunc->getNumParams(); i++) {
624-
args.push_back(builder.makeLocalGet(i, oldFunc->getLocalType(i)));
625-
}
626-
auto* call = builder.makeCall(name, args, oldFunc->getResults());
627-
628-
auto trampoline = builder.makeFunction(newName, oldFunc->type, {}, call);
629-
trampoline->hasExplicitName = oldFunc->hasExplicitName;
630-
primary.addFunction(std::move(trampoline));
631-
633+
Name trampoline = getTrampoline(name);
632634
// Update RefFuncs to refer to it.
633635
for (auto* refFunc : relevantRefFuncs) {
634-
refFunc->func = newName;
636+
refFunc->func = trampoline;
635637
}
636638
}
637639
}

‎src/wasm.h

Lines changed: 4 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -2378,6 +2378,10 @@ class Export {
23782378
assert(std::get_if<Name>(&value));
23792379
}
23802380
Name* getInternalName() { return std::get_if<Name>(&value); }
2381+
void setInternalName(Name name) {
2382+
assert(std::holds_alternative<Name>(value));
2383+
value = name;
2384+
}
23812385
};
23822386

23832387
class ElementSegment : public Named {

‎test/example/module-splitting.txt

Lines changed: 11 additions & 35 deletions
Original file line numberDiff line numberDiff line change
@@ -336,9 +336,9 @@ After:
336336
(import "placeholder" "0" (func $placeholder_0 (type $0) (param i32) (result i32)))
337337
(table $0 1 funcref)
338338
(elem $0 (i32.const 0) $placeholder_0)
339-
(export "foo" (func $foo))
339+
(export "foo" (func $trampoline_foo))
340340
(export "%table" (table $0))
341-
(func $foo (type $0) (param $0 i32) (result i32)
341+
(func $trampoline_foo (type $0) (param $0 i32) (result i32)
342342
(call_indirect $0 (type $0)
343343
(local.get $0)
344344
(i32.const 0)
@@ -451,15 +451,9 @@ After:
451451
(table $0 1 funcref)
452452
(elem $0 (table $table) (i32.const 42) func $trampoline_foo)
453453
(elem $0_1 (table $0) (i32.const 0) func $placeholder_0)
454-
(export "foo" (func $foo))
454+
(export "foo" (func $trampoline_foo))
455455
(export "%table" (table $table))
456456
(export "%table_2" (table $0))
457-
(func $foo (type $0) (param $0 i32) (result i32)
458-
(call_indirect $0 (type $0)
459-
(local.get $0)
460-
(i32.const 0)
461-
)
462-
)
463457
(func $trampoline_foo (type $0) (param $0 i32) (result i32)
464458
(call_indirect $0 (type $0)
465459
(local.get $0)
@@ -499,16 +493,10 @@ After:
499493
(table $0 1 funcref)
500494
(elem $0 (table $table) (global.get $base) func $trampoline_foo)
501495
(elem $0_1 (table $0) (i32.const 0) func $placeholder_0)
502-
(export "foo" (func $foo))
496+
(export "foo" (func $trampoline_foo))
503497
(export "%table" (table $table))
504498
(export "%table_2" (table $0))
505499
(export "%global" (global $base))
506-
(func $foo (type $0) (param $0 i32) (result i32)
507-
(call_indirect $0 (type $0)
508-
(local.get $0)
509-
(i32.const 0)
510-
)
511-
)
512500
(func $trampoline_foo (type $0) (param $0 i32) (result i32)
513501
(call_indirect $0 (type $0)
514502
(local.get $0)
@@ -548,16 +536,10 @@ After:
548536
(table $0 1 funcref)
549537
(elem $0 (table $table) (global.get $base) func $trampoline_foo $trampoline_foo)
550538
(elem $0_1 (table $0) (i32.const 0) func $placeholder_0)
551-
(export "foo" (func $foo))
539+
(export "foo" (func $trampoline_foo))
552540
(export "%table" (table $table))
553541
(export "%table_2" (table $0))
554542
(export "%global" (global $base))
555-
(func $foo (type $0) (param $0 i32) (result i32)
556-
(call_indirect $0 (type $0)
557-
(local.get $0)
558-
(i32.const 0)
559-
)
560-
)
561543
(func $trampoline_foo (type $0) (param $0 i32) (result i32)
562544
(call_indirect $0 (type $0)
563545
(local.get $0)
@@ -601,18 +583,12 @@ After:
601583
(table $0 1 funcref)
602584
(elem $0 (table $table) (global.get $base) func $null $trampoline_foo)
603585
(elem $0_1 (table $0) (i32.const 0) func $placeholder_0)
604-
(export "foo" (func $foo))
586+
(export "foo" (func $trampoline_foo))
605587
(export "%table" (table $table))
606588
(export "%table_2" (table $0))
607589
(export "%global" (global $base))
608590
(func $null (type $0)
609591
)
610-
(func $foo (type $1) (param $0 i32) (result i32)
611-
(call_indirect $0 (type $1)
612-
(local.get $0)
613-
(i32.const 0)
614-
)
615-
)
616592
(func $trampoline_foo (type $1) (param $0 i32) (result i32)
617593
(call_indirect $0 (type $1)
618594
(local.get $0)
@@ -753,13 +729,13 @@ After:
753729
(import "placeholder" "0" (func $placeholder_0 (type $0)))
754730
(table $0 1 funcref)
755731
(elem $0 (i32.const 0) $placeholder_0)
756-
(export "%foo" (func $bar))
732+
(export "%foo" (func $trampoline_bar))
757733
(export "%foo_1" (func $foo))
758734
(export "%table" (table $0))
759735
(func $foo (type $0)
760736
(nop)
761737
)
762-
(func $bar (type $0)
738+
(func $trampoline_bar (type $0)
763739
(call_indirect $0 (type $0)
764740
(i32.const 0)
765741
)
@@ -1151,10 +1127,10 @@ After:
11511127
(import "placeholder" "0" (func $placeholder_0 (type $0)))
11521128
(table $0 1 funcref)
11531129
(elem $0 (i32.const 0) $placeholder_0)
1154-
(export "foo1" (func $foo))
1155-
(export "foo2" (func $foo))
1130+
(export "foo1" (func $trampoline_foo))
1131+
(export "foo2" (func $trampoline_foo))
11561132
(export "%table" (table $0))
1157-
(func $foo (type $0)
1133+
(func $trampoline_foo (type $0)
11581134
(call_indirect $0 (type $0)
11591135
(i32.const 0)
11601136
)

‎test/lit/wasm-split/multi-split.wast

Lines changed: 18 additions & 36 deletions
Original file line numberDiff line numberDiff line change
@@ -14,13 +14,11 @@
1414
;; RUN: wasm-dis %t3.wasm | filecheck %s --check-prefix=MOD3-OPTIONS
1515

1616
(module
17-
;; PRIMARY: (type $ret-i64 (func (result i64)))
18-
1917
;; PRIMARY: (type $ret-i32 (func (result i32)))
20-
;; PRIMARY-OPTIONS: (type $ret-i64 (func (result i64)))
21-
2218
;; PRIMARY-OPTIONS: (type $ret-i32 (func (result i32)))
2319
(type $ret-i32 (func (result i32)))
20+
;; PRIMARY: (type $ret-i64 (func (result i64)))
21+
;; PRIMARY-OPTIONS: (type $ret-i64 (func (result i64)))
2422
(type $ret-i64 (func (result i64)))
2523
;; PRIMARY: (type $ret-f32 (func (result f32)))
2624
;; PRIMARY-OPTIONS: (type $ret-f32 (func (result f32)))
@@ -206,11 +204,11 @@
206204

207205
;; MOD3: (type $2 (func (result f32)))
208206

209-
;; MOD3: (import "" "table_6" (table $timport$0 1 funcref))
207+
;; MOD3: (import "" "table_5" (table $timport$0 1 funcref))
210208

211209
;; MOD3: (import "" "trampoline_A" (func $trampoline_A (type $0) (result i32)))
212210

213-
;; MOD3: (import "" "trampoline_B" (func $trampoline_B (type $1) (result i64)))
211+
;; MOD3: (import "" "B" (func $trampoline_B (type $1) (result i64)))
214212

215213
;; MOD3: (elem $0 (i32.const 0) $C)
216214

@@ -240,11 +238,11 @@
240238

241239
;; MOD3-OPTIONS: (type $2 (func (result f32)))
242240

243-
;; MOD3-OPTIONS: (import "custom_env" "table_6" (table $timport$0 1 funcref))
241+
;; MOD3-OPTIONS: (import "custom_env" "table_5" (table $timport$0 1 funcref))
244242

245243
;; MOD3-OPTIONS: (import "custom_env" "trampoline_A" (func $trampoline_A (type $0) (result i32)))
246244

247-
;; MOD3-OPTIONS: (import "custom_env" "trampoline_B" (func $trampoline_B (type $1) (result i64)))
245+
;; MOD3-OPTIONS: (import "custom_env" "B" (func $trampoline_B (type $1) (result i64)))
248246

249247
;; MOD3-OPTIONS: (elem $0 (i32.const 0) $C)
250248

@@ -289,9 +287,9 @@
289287
)
290288
;; PRIMARY: (import "placeholder" "0" (func $placeholder_0 (type $ret-i32) (result i32)))
291289

292-
;; PRIMARY: (import "placeholder" "0" (func $placeholder_0_5 (type $ret-i64) (result i64)))
290+
;; PRIMARY: (import "placeholder" "0" (func $placeholder_0_4 (type $ret-i64) (result i64)))
293291

294-
;; PRIMARY: (import "placeholder" "0" (func $placeholder_0_6 (type $ret-f32) (result f32)))
292+
;; PRIMARY: (import "placeholder" "0" (func $placeholder_0_5 (type $ret-f32) (result f32)))
295293

296294
;; PRIMARY: (table $0 1 funcref)
297295

@@ -301,43 +299,35 @@
301299

302300
;; PRIMARY: (elem $0 (table $0) (i32.const 0) func $placeholder_0)
303301

304-
;; PRIMARY: (elem $1 (table $1) (i32.const 0) func $placeholder_0_5)
302+
;; PRIMARY: (elem $1 (table $1) (i32.const 0) func $placeholder_0_4)
305303

306-
;; PRIMARY: (elem $2 (table $2) (i32.const 0) func $placeholder_0_6)
304+
;; PRIMARY: (elem $2 (table $2) (i32.const 0) func $placeholder_0_5)
307305

308-
;; PRIMARY: (export "B" (func $1))
306+
;; PRIMARY: (export "B" (func $trampoline_B))
309307

310-
;; PRIMARY: (export "C" (func $3))
308+
;; PRIMARY: (export "C" (func $trampoline_C))
311309

312310
;; PRIMARY: (export "table" (table $0))
313311

314312
;; PRIMARY: (export "trampoline_A" (func $trampoline_A))
315313

316314
;; PRIMARY: (export "table_4" (table $1))
317315

318-
;; PRIMARY: (export "trampoline_B" (func $trampoline_B))
319-
320-
;; PRIMARY: (export "table_6" (table $2))
316+
;; PRIMARY: (export "table_5" (table $2))
321317

322318
;; PRIMARY: (func $trampoline_A (type $ret-i32) (result i32)
323319
;; PRIMARY-NEXT: (call_indirect $0 (type $ret-i32)
324320
;; PRIMARY-NEXT: (i32.const 0)
325321
;; PRIMARY-NEXT: )
326322
;; PRIMARY-NEXT: )
327323

328-
;; PRIMARY: (func $1 (type $ret-i64) (result i64)
329-
;; PRIMARY-NEXT: (call_indirect $1 (type $ret-i64)
330-
;; PRIMARY-NEXT: (i32.const 0)
331-
;; PRIMARY-NEXT: )
332-
;; PRIMARY-NEXT: )
333-
334324
;; PRIMARY: (func $trampoline_B (type $ret-i64) (result i64)
335325
;; PRIMARY-NEXT: (call_indirect $1 (type $ret-i64)
336326
;; PRIMARY-NEXT: (i32.const 0)
337327
;; PRIMARY-NEXT: )
338328
;; PRIMARY-NEXT: )
339329

340-
;; PRIMARY: (func $3 (type $ret-f32) (result f32)
330+
;; PRIMARY: (func $trampoline_C (type $ret-f32) (result f32)
341331
;; PRIMARY-NEXT: (call_indirect $2 (type $ret-f32)
342332
;; PRIMARY-NEXT: (i32.const 0)
343333
;; PRIMARY-NEXT: )
@@ -355,39 +345,31 @@
355345

356346
;; PRIMARY-OPTIONS: (elem $2 (table $2) (i32.const 0) funcref (item (ref.null nofunc)))
357347

358-
;; PRIMARY-OPTIONS: (export "B" (func $1))
348+
;; PRIMARY-OPTIONS: (export "B" (func $trampoline_B))
359349

360-
;; PRIMARY-OPTIONS: (export "C" (func $3))
350+
;; PRIMARY-OPTIONS: (export "C" (func $trampoline_C))
361351

362352
;; PRIMARY-OPTIONS: (export "table" (table $0))
363353

364354
;; PRIMARY-OPTIONS: (export "trampoline_A" (func $trampoline_A))
365355

366356
;; PRIMARY-OPTIONS: (export "table_4" (table $1))
367357

368-
;; PRIMARY-OPTIONS: (export "trampoline_B" (func $trampoline_B))
369-
370-
;; PRIMARY-OPTIONS: (export "table_6" (table $2))
358+
;; PRIMARY-OPTIONS: (export "table_5" (table $2))
371359

372360
;; PRIMARY-OPTIONS: (func $trampoline_A (type $ret-i32) (result i32)
373361
;; PRIMARY-OPTIONS-NEXT: (call_indirect $0 (type $ret-i32)
374362
;; PRIMARY-OPTIONS-NEXT: (i32.const 0)
375363
;; PRIMARY-OPTIONS-NEXT: )
376364
;; PRIMARY-OPTIONS-NEXT: )
377365

378-
;; PRIMARY-OPTIONS: (func $1 (type $ret-i64) (result i64)
379-
;; PRIMARY-OPTIONS-NEXT: (call_indirect $1 (type $ret-i64)
380-
;; PRIMARY-OPTIONS-NEXT: (i32.const 0)
381-
;; PRIMARY-OPTIONS-NEXT: )
382-
;; PRIMARY-OPTIONS-NEXT: )
383-
384366
;; PRIMARY-OPTIONS: (func $trampoline_B (type $ret-i64) (result i64)
385367
;; PRIMARY-OPTIONS-NEXT: (call_indirect $1 (type $ret-i64)
386368
;; PRIMARY-OPTIONS-NEXT: (i32.const 0)
387369
;; PRIMARY-OPTIONS-NEXT: )
388370
;; PRIMARY-OPTIONS-NEXT: )
389371

390-
;; PRIMARY-OPTIONS: (func $3 (type $ret-f32) (result f32)
372+
;; PRIMARY-OPTIONS: (func $trampoline_C (type $ret-f32) (result f32)
391373
;; PRIMARY-OPTIONS-NEXT: (call_indirect $2 (type $ret-f32)
392374
;; PRIMARY-OPTIONS-NEXT: (i32.const 0)
393375
;; PRIMARY-OPTIONS-NEXT: )

‎test/lit/wasm-split/no-active-segment.wast

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -30,11 +30,11 @@
3030

3131
;; PRIMARY: (elem $0 (i32.const 0) $placeholder_0)
3232

33-
;; PRIMARY: (export "foo" (func $0))
33+
;; PRIMARY: (export "foo" (func $trampoline_foo))
3434

3535
;; PRIMARY: (export "table" (table $0))
3636

37-
;; PRIMARY: (func $0 (type $0)
37+
;; PRIMARY: (func $trampoline_foo (type $0)
3838
;; PRIMARY-NEXT: (call_indirect $0 (type $0)
3939
;; PRIMARY-NEXT: (i32.const 0)
4040
;; PRIMARY-NEXT: )

0 commit comments

Comments
 (0)