[wasm-split] Split immutable globals

src/ir/module-splitting.cpp

@@ -77,6 +77,7 @@
 #include "ir/find_all.h"
 #include "ir/module-utils.h"
 #include "ir/names.h"
+#include "ir/properties.h"
 #include "support/small_vector.h"
 #include "support/stdckdint.h"
 #include "wasm-builder.h"
@@ -336,6 +337,7 @@
   struct ItemInfo {
     UsedNames* owner = nullptr;
     SmallVector<Module*, 2> usingSecondaries;
+    bool usedByPrimary = false;
   };
 
   std::unordered_map<Name, ItemInfo> tables;
@@ -376,6 +378,9 @@
     // used by the primary module or multiple secondary modules, the primary
     // module is the owner.
     auto [it, inserted] = (this->*mapField).insert({name, ItemInfo{owner, {}}});
+    if (owner == &primaryUsed) {
+      it->second.usedByPrimary = true;
+    }
     Module* secondary = nullptr;
     if (owner != &primaryUsed) {
       size_t index = owner - secondaryUsed.data();
@@ -420,13 +425,30 @@
   }
 
   bool isUnused(Name name, const std::unordered_map<Name, ItemInfo>& map) {
-    return getOwner(name, map) == nullptr;
+    auto it = map.find(name);
+    if (it != map.end()) {
+      return !it->second.usedByPrimary && it->second.usingSecondaries.empty();
+    }
+    return true;
+  }
+
+  bool isUsedByPrimary(Name name,
+                       const std::unordered_map<Name, ItemInfo>& map) {
+    auto it = map.find(name);
+    if (it != map.end()) {
+      return it->second.usedByPrimary;
+    }
+    return false;
   }
 
   bool usedBySingleSecondary(Name name,
                              const std::unordered_map<Name, ItemInfo>& map) {
-    auto* owner = getOwner(name, map);
-    return owner != nullptr && owner != &primaryUsed;
+    auto it = map.find(name);
+    if (it != map.end()) {
+      return !it->second.usedByPrimary &&
+             it->second.usingSecondaries.size() == 1;
+    }
+    return false;
   }
 };
 
@@ -758,6 +780,18 @@
   }
 }
 
+static bool canCopyGlobal(Global* global, Module& wasm) {
+  // Mutable globals cannot be copied to multiple modules because mutations in
+  // one module would not be visible in others. Immutable globals with
+  // generative initializers (e.g. struct.new) also cannot be copied because
+  // there would be an observable difference (e.g. struct identity tested with
+  // ref.eq) between the values in the copies. Note that imported globals can be
+  // copied / moved.
+  return !global->mutable_ &&
+         (!global->init ||
+          !Properties::isGenerative(global->init, nullptr, wasm));
+}
+
 void ModuleSplitter::computeUsedNames() {
   tracker.secondaries = &secondaries;
   tracker.secondaryUsed.resize(secondaries.size());
@@ -1047,9 +1081,41 @@
     if (!global->init) {
       continue;
     }
-    if (UsedNames* owner = tracker.getOwner(global->name, tracker.globals)) {
-      for (auto* get : FindAll<GlobalGet>(global->init).list) {
-        tracker.insert<Global>(get->name, owner);
+    auto gets = FindAll<GlobalGet>(global->init).list;
+    if (gets.empty()) {
+      continue;
+    }
+
+    // In case of mutable globals or immutable globals with generative
+    // initializers (e.g. struct.new), we cannot have multiple copies. Compute
+    // the 'owner' of the global and insert its dependent globals there.
+    if (!canCopyGlobal(global.get(), primary)) {
+      if (UsedNames* owner = tracker.getOwner(global->name, tracker.globals)) {
+        for (auto* get : gets) {
+          tracker.insert<Global>(get->name, owner);
+        }
+      }
+      continue;
+    }
+
+    // In case of copyable immutable globals, we can have multiple copies of
+    // them. To reduce the primary module size, we just copy the globals to all
+    // secondary modules using them. So here we insert the dependent globals in
+    // all using modules.
+    if (tracker.isUsedByPrimary(global->name, tracker.globals)) {
+      for (auto* get : gets) {
+        tracker.insert<Global>(get->name, &primaryUsed);
+      }
+    }
+    auto& usingSecs =
+      tracker.getUsingSecondaries(global->name, tracker.globals);
+    for (size_t i = 0; i < secondaries.size(); ++i) {
+      Module* sec = secondaries[i].get();
+      if (std::find(usingSecs.begin(), usingSecs.end(), sec) !=
+          usingSecs.end()) {
+        for (auto* get : gets) {
+          tracker.insert<Global>(get->name, &secondaryUsed[i]);
+        }
       }
     }
   }
@@ -1063,8 +1129,8 @@
   // 2. If an item is used by only a single secondary module, move the item to
   //    that secondary module. If an item is used by multiple modules (including
   //    the primary and secondary modules), export the item from the primary and
-  //    import it from the using secondary modules.
-
+  //    import it from the using secondary modules. (except for non-generative
+  //    immutable globals, which we just copy to using secondary modules)
   auto shareElements = [&](auto& elements,
                            auto& trackerElements,
                            auto copyElement,
@@ -1073,18 +1139,41 @@
                            ExternalKind kind = ExternalKind::Invalid) {
     std::unordered_set<Name> elementsToRemove;
     for (auto& element : elements) {
+      using T = std::remove_pointer_t<decltype(element.get())>;
+      bool isCopyableGlobal = false;
+      if constexpr (std::is_same_v<T, Global>) {
+        isCopyableGlobal = canCopyGlobal(element.get(), primary);
+      }
+
       if (tracker.isUnused(element->name, trackerElements)) {
+        // If this element is not used anywhere, just remove it
         elementsToRemove.insert(element->name);
       } else if (tracker.usedBySingleSecondary(element->name,
                                                trackerElements)) {
+        // If this element is used in a single secondary module, move it to the
+        // secondary module
         auto* secondary =
           tracker.getUsingSecondaries(element->name, trackerElements)[0];
         copyElement(element.get(), *secondary);
         elementsToRemove.insert(element->name);
+      } else if (isCopyableGlobal) {
+        // If this element is a copyable immutable global, copy it to all using
+        // secondary modules to reduce the primary module size, because we can
+        // have multiple copies. If it is not used in the primary module, remove
+        // it from there.
+        if (!tracker.isUsedByPrimary(element->name, trackerElements)) {
+          elementsToRemove.insert(element->name);
+        }
+        for (auto* secondary :
+             tracker.getUsingSecondaries(element->name, trackerElements)) {
+          copyElement(element.get(), *secondary);
+        }
       } else {
+        // If this element is used by multiple modules (primary + multiple
+        // secondaries or just multiple secondaries), export them from the
+        // primary module and import them from using secondary modules.
         // We only import and export Importables, i.e., we don't do this for
         // segments.
-        using T = std::remove_pointer_t<decltype(element.get())>;
         if constexpr (std::is_base_of_v<Importable, T>) {
           for (auto* secondary :
                tracker.getUsingSecondaries(element->name, trackerElements)) {

test/lit/wasm-split/dispatch-table-base-global-used-elsewhere.wast

@@ -25,8 +25,6 @@
 
  ;; PRIMARY:      (export "table" (table $table))
 
- ;; PRIMARY:      (export "global" (global $base))
-
  ;; PRIMARY:      (export "keep" (func $keep))
 
  ;; PRIMARY:      (func $keep
@@ -42,7 +40,7 @@
  )
  ;; SECONDARY:      (import "primary" "table" (table $table 2 2 funcref))
 
- ;; SECONDARY:      (import "primary" "global" (global $base i32))
+ ;; SECONDARY:      (import "env" "base" (global $base i32))
 
  ;; SECONDARY:      (import "primary" "keep" (func $keep))
 

test/lit/wasm-split/passive-deps.wast

@@ -22,15 +22,14 @@
 ;; PRIMARY-NEXT:  (type $0 (func))
 ;; PRIMARY-NEXT:  (global $g funcref (ref.null nofunc))
 ;; PRIMARY-NEXT:  (elem $passive-elem funcref (item (global.get $g)))
-;; PRIMARY-NEXT:  (export "global" (global $g))
 ;; PRIMARY-NEXT:  (func $keep (type $0)
 ;; PRIMARY-NEXT:   (elem.drop $passive-elem)
 ;; PRIMARY-NEXT:  )
 ;; PRIMARY-NEXT: )
 
 ;; SECONDARY:      (module
 ;; SECONDARY-NEXT:  (type $0 (func))
-;; SECONDARY-NEXT:  (import "primary" "global" (global $g funcref))
+;; SECONDARY-NEXT:  (global $g funcref (ref.null nofunc))
 ;; SECONDARY-NEXT:  (func $split (type $0)
 ;; SECONDARY-NEXT:   (drop
 ;; SECONDARY-NEXT:    (global.get $g)

test/lit/wasm-split/split-module-items.wast

@@ -5,21 +5,27 @@
 ;; Check that
 ;; 1. Items only used in the primary module stay in the primary module
 ;; 2. Items only used in the secondary module are moved to the secondary module
-;; 3. Items used in both modules are exported from the primary and imported from
-;;    the secondary module
+;; 3. Items (that are not copyable immutable globals) used in both modules are
+;;    exported from the primary and imported from the secondary module
+;; 3-1. Non-generative immutable globals used in both modules stay in the
+;;      primary module and also are copied to the secondary module
+;; 3-2. Generative immutable globals (e.g. struct.new) used in both modules are
+;;      exported from the primary and imported from the secondary module
 
 (module
  (rec
   (type $struct (descriptor $desc) (struct))
   (type $desc (describes $struct) (struct))
  )
+ (type $simple-struct (struct))
  (global $null-desc (ref null none) (ref.null none))
 
  (memory $keep-memory 1 1)
  ;; This is only used in the secondary module, but segments $keep-data2 can trap
  ;; so it is pinned to the primary, so this will be too.
  (memory $keep-memory2 1 1)
  (global $keep-immutable-global i32 (i32.const 20))
+ (global $keep-immutable-generative-global (ref $simple-struct) (struct.new_default $simple-struct))
  (global $keep-mutable-global (mut i32) (i32.const 20))
  (table $keep-table 1 1 funcref)
  ;; This is only used in the secondary module, but segments $keep-elem2 and
@@ -37,6 +43,7 @@
 
  (memory $split-memory 1 1)
  (global $split-immutable-global i32 (i32.const 20))
+ (global $split-immutable-generative-global (ref $simple-struct) (struct.new_default $simple-struct))
  (global $split-mutable-global (mut i32) (i32.const 20))
  (table $split-table 1 1 funcref)
  (tag $split-tag (param i32))
@@ -45,15 +52,18 @@
 
  (memory $shared-memory 1 1)
  (global $shared-immutable-global i32 (i32.const 20))
+ (global $shared-immutable-generative-global (ref $simple-struct) (struct.new_default $simple-struct))
  (global $shared-mutable-global (mut i32) (i32.const 20))
  (table $shared-table 1 1 funcref)
  (tag $shared-tag (param i32))
  (elem $shared-elem (table $shared-table) (i32.const 0) funcref (item (ref.null nofunc)))
  (data $shared-data (memory $shared-memory) (i32.const 0) "a")
 
  ;; PRIMARY:      (global $keep-immutable-global i32 (i32.const 20))
+ ;; PRIMARY-NEXT: (global $keep-immutable-generative-global (ref $simple-struct) (struct.new_default $simple-struct))
  ;; PRIMARY-NEXT: (global $keep-mutable-global (mut i32) (i32.const 20))
  ;; PRIMARY-NEXT: (global $shared-immutable-global i32 (i32.const 20))
+ ;; PRIMARY-NEXT: (global $shared-immutable-generative-global (ref $simple-struct) (struct.new_default $simple-struct))
  ;; PRIMARY-NEXT: (global $shared-mutable-global (mut i32) (i32.const 20))
  ;; PRIMARY-NEXT: (memory $keep-memory 1 1)
  ;; PRIMARY-NEXT: (memory $keep-memory2 1 1)
@@ -71,36 +81,38 @@
  ;; PRIMARY-NEXT:  (global.get $null-desc)
  ;; PRIMARY-NEXT: )))
  ;; PRIMARY-NEXT: (elem $shared-elem (table $shared-table) (i32.const 0) funcref (item (ref.null nofunc)))
- ;; PRIMARY:      (tag $keep-tag (type $3) (param i32))
- ;; PRIMARY-NEXT: (tag $shared-tag (type $3) (param i32))
+ ;; PRIMARY:      (tag $keep-tag (type $4) (param i32))
+ ;; PRIMARY-NEXT: (tag $shared-tag (type $4) (param i32))
 
  ;; PRIMARY:      (export "memory" (memory $keep-memory2))
  ;; PRIMARY-NEXT: (export "memory_1" (memory $shared-memory))
  ;; PRIMARY-NEXT: (export "table" (table $keep-table2))
  ;; PRIMARY-NEXT: (export "table_3" (table $shared-table))
- ;; PRIMARY-NEXT: (export "global" (global $shared-immutable-global))
+ ;; PRIMARY-NEXT: (export "global" (global $shared-immutable-generative-global))
  ;; PRIMARY-NEXT: (export "global_5" (global $shared-mutable-global))
  ;; PRIMARY-NEXT: (export "tag" (tag $shared-tag))
  ;; PRIMARY-NEXT: (export "keep" (func $keep))
  ;; PRIMARY-NEXT: (export "table_8" (table $3))
 
  ;; SECONDARY:      (import "primary" "memory" (memory $keep-memory2 1 1))
  ;; SECONDARY-NEXT: (import "primary" "memory_1" (memory $shared-memory 1 1))
- ;; SECONDARY-NEXT: (import "primary" "table" (table $keep-table2 1 1 (ref null $2)))
+ ;; SECONDARY-NEXT: (import "primary" "table" (table $keep-table2 1 1 (ref null $3)))
  ;; SECONDARY-NEXT: (import "primary" "table_3" (table $shared-table 1 1 funcref))
  ;; SECONDARY-NEXT: (import "primary" "table_8" (table $timport$2 1 funcref))
- ;; SECONDARY-NEXT: (import "primary" "global" (global $shared-immutable-global i32))
+ ;; SECONDARY-NEXT: (import "primary" "global" (global $shared-immutable-generative-global (ref $0)))
  ;; SECONDARY-NEXT: (import "primary" "global_5" (global $shared-mutable-global (mut i32)))
  ;; SECONDARY-NEXT: (import "primary" "keep" (func $keep (exact (param i32) (result i32))))
- ;; SECONDARY-NEXT: (import "primary" "tag" (tag $shared-tag (type $1) (param i32)))
+ ;; SECONDARY-NEXT: (import "primary" "tag" (tag $shared-tag (type $2) (param i32)))
 
  ;; SECONDARY:      (global $split-immutable-global i32 (i32.const 20))
+ ;; SECONDARY-NEXT: (global $split-immutable-generative-global (ref $0) (struct.new_default $0))
  ;; SECONDARY-NEXT: (global $split-mutable-global (mut i32) (i32.const 20))
+ ;; SECONDARY-NEXT: (global $shared-immutable-global i32 (i32.const 20))
  ;; SECONDARY-NEXT: (memory $split-memory 1 1)
  ;; SECONDARY-NEXT: (data $split-data (memory $split-memory) (i32.const 0) "a")
  ;; SECONDARY-NEXT: (table $split-table 1 1 funcref)
  ;; SECONDARY-NEXT: (elem $split-elem (table $split-table) (i32.const 0) funcref (item (ref.null nofunc)))
- ;; SECONDARY:      (tag $split-tag (type $1) (param i32))
+ ;; SECONDARY:      (tag $split-tag (type $2) (param i32))
 
  (func $keep (param i32) (result i32)
   (call $split (i32.const 0))
@@ -117,10 +129,14 @@
     (i32.const 0)
    )
   )
-  ;; Uses $keep-immutable-global and $keep-mutable-global
+  ;; Uses $keep-immutable-global, $keep-immutable-generative-global, and
+  ;; $keep-mutable-global
   (drop
    (global.get $keep-immutable-global)
   )
+  (drop
+   (global.get $keep-immutable-generative-global)
+  )
   (drop
    (global.get $keep-mutable-global)
   )
@@ -141,10 +157,14 @@
     (i32.const 0)
    )
   )
-  ;; Uses $shared-immutable-global and $shared-mutable-global
+  ;; Uses $shared-immutable-global, $shared-immutable-generative-global, and
+  ;; $shared-mutable-global
   (drop
    (global.get $shared-immutable-global)
   )
+  (drop
+   (global.get $shared-immutable-generative-global)
+  )
   (drop
    (global.get $shared-mutable-global)
   )
@@ -180,10 +200,14 @@
   (drop
    (table.get $keep-table2 (i32.const 0))
   )
-  ;; Uses $split-immutable-global and $shared-mutable-global
+  ;; Uses $split-immutable-global, $split-immutable-generative-global, and
+  ;; $split-mutable-global
   (drop
    (global.get $split-immutable-global)
   )
+  (drop
+   (global.get $split-immutable-generative-global)
+  )
   (drop
    (global.get $split-mutable-global)
   )
@@ -204,10 +228,14 @@
     (i32.const 0)
    )
   )
-  ;; Uses $shared-immutable-global and $shared-mutable-global
+  ;; Uses $shared-immutable-global, $shared-immutable-generative-global, and
+  ;; $shared-mutable-global
   (drop
    (global.get $shared-immutable-global)
   )
+  (drop
+   (global.get $shared-immutable-generative-global)
+  )
   (drop
    (global.get $shared-mutable-global)
   )