[wasm-split] Share active table with caller modules

src/ir/module-splitting.cpp

@@ -342,6 +342,8 @@ struct ModuleSplitter {
   // Map from original secondary function name to its trampoline
   std::unordered_map<Name, Name> trampolineMap;
 
+  void shareActiveTable(Module* secondary);
+
   // Initialization helpers
   static std::unique_ptr<Module> initSecondary(const Module& primary);
   static std::unordered_map<Name, Name>
@@ -590,6 +592,32 @@ Name ModuleSplitter::getTrampoline(Name funcName) {
   return trampoline;
 }
 
+void ModuleSplitter::shareActiveTable(Module* secondary) {
+  assert(tableManager.activeTable);
+  auto secondaryTable =
+    secondary->getTableOrNull(tableManager.activeTable->name);
+  if (secondaryTable) {
+    // In case it's already in the secondary module, sync the initial/max
+    secondaryTable->initial = tableManager.activeTable->initial;
+    secondaryTable->max = tableManager.activeTable->max;
+  } else {
+    secondaryTable =
+      ModuleUtils::copyTable(tableManager.activeTable, *secondary);
+    makeImportExport(
+      *tableManager.activeTable, *secondaryTable, "table", ExternalKind::Table);
+  }
+  if (tableManager.activeBase.global) {
+    auto* primaryGlobal = primary.getGlobal(tableManager.activeBase.global);
+    auto* secondaryGlobal =
+      secondary->getGlobalOrNull(tableManager.activeBase.global);
+    if (!secondaryGlobal) {
+      secondaryGlobal = ModuleUtils::copyGlobal(primaryGlobal, *secondary);
+      makeImportExport(
+        *primaryGlobal, *secondaryGlobal, "global", ExternalKind::Global);
+    }
+  }
+}
+
 void ModuleSplitter::thunkExportedSecondaryFunctions() {
   // Update exports of secondary functions in the primary module to export
   // wrapper functions that indirectly call the secondary functions. We are
@@ -988,6 +1016,7 @@ void ModuleSplitter::indirectCallsToSecondaryFunctions() {
   // corresponding table indices instead.
   struct CallIndirector : public PostWalker<CallIndirector> {
     ModuleSplitter& parent;
+    std::unordered_set<Module*> activeTableUsingSecondaries;
     CallIndirector(ModuleSplitter& parent) : parent(parent) {}
     void visitCall(Call* curr) {
       // Return if the call's target is not in one of the secondary module.
@@ -1014,13 +1043,23 @@ void ModuleSplitter::indirectCallsToSecondaryFunctions() {
                                  curr->operands,
                                  func->type.getHeapType(),
                                  curr->isReturn));
+
+      // Share the active table with the current module (caller). We share the
+      // active table with with calleeModule later in setupTablePathing.
+      if (currModule != &parent.primary) {
+        activeTableUsingSecondaries.insert(currModule);
+      }
     }
   };
   CallIndirector callIndirector(*this);
   callIndirector.walkModule(&primary);
   for (auto& secondaryPtr : secondaries) {
     callIndirector.walkModule(secondaryPtr.get());
   }
+
+  for (auto* secondary : callIndirector.activeTableUsingSecondaries) {
+    shareActiveTable(secondary);
+  }
 }
 
 void ModuleSplitter::exportImportCalledPrimaryFunctions() {
@@ -1120,40 +1159,10 @@ void ModuleSplitter::setupTablePatching() {
 
   for (auto& [secondaryPtr, replacedElems] : moduleToReplacedElems) {
     Module& secondary = *secondaryPtr;
-    // Import and export the active table if necessary. Unless we use an
-    // existing table as an active table (e.g. because reference-types is
-    // disabled) and that table was already being used by an existing indirect
-    // call, shareImportableItems wasn't able to mark it as used in secondaries,
-    // so we should export and import the active table here.
-    auto secondaryTable =
-      secondary.getTableOrNull(tableManager.activeTable->name);
-    if (secondaryTable) {
-      // In case it's already in the secondary module, sync the initial/max
-      secondaryTable->initial = tableManager.activeTable->initial;
-      secondaryTable->max = tableManager.activeTable->max;
-    } else {
-      secondaryTable =
-        ModuleUtils::copyTable(tableManager.activeTable, secondary);
-      makeImportExport(*tableManager.activeTable,
-                       *secondaryTable,
-                       "table",
-                       ExternalKind::Table);
-    }
+    shareActiveTable(&secondary);
+    auto* secondaryTable = secondary.getTable(tableManager.activeTable->name);
 
     if (tableManager.activeBase.global) {
-      // Import and export the active table's base global if necessary. Unless
-      // the base global was already being used elsewhere in secondaries,
-      // shareImportableItems wasn't able to mark it as used in secondaries, so
-      // we should export and import it here.
-      auto* primaryGlobal = primary.getGlobal(tableManager.activeBase.global);
-      auto* secondaryGlobal =
-        secondary.getGlobalOrNull(tableManager.activeBase.global);
-      if (!secondaryGlobal) {
-        secondaryGlobal = ModuleUtils::copyGlobal(primaryGlobal, secondary);
-        makeImportExport(
-          *primaryGlobal, *secondaryGlobal, "global", ExternalKind::Global);
-      }
-
       assert(tableManager.activeTableSegments.size() == 1 &&
              "Unexpected number of segments with non-const base");
       assert(secondary.tables.size() == 1 && secondary.elementSegments.empty());

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

@@ -0,0 +1,67 @@
+;; NOTE: Assertions have been generated by update_lit_checks.py --all-items and should not be edited.
+
+;; RUN: wasm-split -all -g --multi-split %s --manifest %S/multi-split.wast.manifest --out-prefix=%t -o %t.wasm
+;; RUN: wasm-dis %t.wasm | filecheck %s --check-prefix=PRIMARY
+;; RUN: wasm-dis %t1.wasm | filecheck %s --check-prefix=MOD1
+;; RUN: wasm-dis %t2.wasm | filecheck %s --check-prefix=MOD2
+;; RUN: wasm-dis %t3.wasm | filecheck %s --check-prefix=MOD3
+
+;; A regresion test for the case the active table was not correctly shared with
+;; MOD1. Because func $A is not called by any function, MOD1 is a secondary
+;; module who only acts as a caller.
+
+(module
+ ;; MOD1:      (type $0 (func))
+
+ ;; MOD1:      (import "primary" "table" (table $timport$0 2 funcref))
+
+ ;; MOD1:      (func $A
+ ;; MOD1-NEXT:  (call_indirect (type $0)
+ ;; MOD1-NEXT:   (i32.const 0)
+ ;; MOD1-NEXT:  )
+ ;; MOD1-NEXT:  (call_indirect (type $0)
+ ;; MOD1-NEXT:   (i32.const 1)
+ ;; MOD1-NEXT:  )
+ ;; MOD1-NEXT: )
+ (func $A
+  (call $B)
+  (call $C)
+ )
+
+ ;; MOD2:      (type $0 (func))
+
+ ;; MOD2:      (import "primary" "table" (table $timport$0 2 funcref))
+
+ ;; MOD2:      (elem $0 (i32.const 0) $B)
+
+ ;; MOD2:      (func $B
+ ;; MOD2-NEXT:  (call_indirect (type $0)
+ ;; MOD2-NEXT:   (i32.const 1)
+ ;; MOD2-NEXT:  )
+ ;; MOD2-NEXT: )
+ (func $B
+  (call $C)
+ )
+
+ ;; MOD3:      (type $0 (func))
+
+ ;; MOD3:      (import "primary" "table" (table $timport$0 2 funcref))
+
+ ;; MOD3:      (elem $0 (i32.const 1) $C)
+
+ ;; MOD3:      (func $C
+ ;; MOD3-NEXT: )
+ (func $C
+ )
+)
+;; PRIMARY:      (type $0 (func))
+
+;; PRIMARY:      (import "placeholder.2" "0" (func $placeholder_0))
+
+;; PRIMARY:      (import "placeholder.3" "1" (func $placeholder_1))
+
+;; PRIMARY:      (table $0 2 funcref)
+
+;; PRIMARY:      (elem $0 (i32.const 0) $placeholder_0 $placeholder_1)
+
+;; PRIMARY:      (export "table" (table $0))