Support using JSPI to load the secondary wasm split module. (#5431)

When using JSPI with wasm-split, any calls to secondary module functions
will now first check a global to see  if the module is loaded. If not
loaded it will call a JSPI'ed function that will handle loading module.

The setup is split into the JSPI pass and wasm-split tool since the JSPI
pass is first run by emscripten and we need to JSPI'ify the load secondary
module function. wasm-split then injects all the checks and calls to the
load function.

Author: brendandahl

src/ir/export-utils.h

@@ -24,6 +24,16 @@ namespace wasm::ExportUtils {
 std::vector<Function*> getExportedFunctions(Module& wasm);
 std::vector<Global*> getExportedGlobals(Module& wasm);
 
+inline bool isExported(const Module& module, const Function& func) {
+  for (auto& exportFunc : module.exports) {
+    if (exportFunc->kind == ExternalKind::Function &&
+        exportFunc->value == func.name) {
+      return true;
+    }
+  }
+  return false;
+};
+
 } // namespace wasm::ExportUtils
 
 #endif // wasm_ir_export_h

src/ir/module-splitting.cpp

@@ -69,6 +69,7 @@
 
 #include "ir/module-splitting.h"
 #include "ir/element-utils.h"
+#include "ir/export-utils.h"
 #include "ir/manipulation.h"
 #include "ir/module-utils.h"
 #include "ir/names.h"
@@ -80,6 +81,8 @@ namespace wasm::ModuleSplitting {
 
 namespace {
 
+static const Name LOAD_SECONDARY_STATUS = "load_secondary_module_status";
+
 template<class F> void forEachElement(Module& module, F f) {
   ModuleUtils::iterActiveElementSegments(module, [&](ElementSegment* segment) {
     Name base = "";
@@ -273,6 +276,9 @@ struct ModuleSplitter {
   // Map placeholder indices to the names of the functions they replace.
   std::map<size_t, Name> placeholderMap;
 
+  // Internal name of the LOAD_SECONDARY_MODULE function.
+  Name internalLoadSecondaryModule;
+
   // Initialization helpers
   static std::unique_ptr<Module> initSecondary(const Module& primary);
   static std::pair<std::set<Name>, std::set<Name>>
@@ -281,8 +287,10 @@ struct ModuleSplitter {
 
   // Other helpers
   void exportImportFunction(Name func);
+  Expression* maybeLoadSecondary(Builder& builder, Expression* callIndirect);
 
   // Main splitting steps
+  void setupJSPI();
   void moveSecondaryFunctions();
   void thunkExportedSecondaryFunctions();
   void indirectCallsToSecondaryFunctions();
@@ -297,6 +305,9 @@ struct ModuleSplitter {
       primaryFuncs(classifiedFuncs.first),
       secondaryFuncs(classifiedFuncs.second), tableManager(primary),
       exportedPrimaryFuncs(initExportedPrimaryFuncs(primary)) {
+    if (config.jspi) {
+      setupJSPI();
+    }
     moveSecondaryFunctions();
     thunkExportedSecondaryFunctions();
     indirectCallsToSecondaryFunctions();
@@ -306,6 +317,23 @@ struct ModuleSplitter {
   }
 };
 
+void ModuleSplitter::setupJSPI() {
+  assert(primary.getExportOrNull(LOAD_SECONDARY_MODULE) &&
+         "The load secondary module function must exist");
+  // Remove the exported LOAD_SECONDARY_MODULE function since it's only needed
+  // internally.
+  internalLoadSecondaryModule = primary.getExport(LOAD_SECONDARY_MODULE)->value;
+  primary.removeExport(LOAD_SECONDARY_MODULE);
+  Builder builder(primary);
+  // Add a global to track whether the secondary module has been loaded yet.
+  primary.addGlobal(builder.makeGlobal(LOAD_SECONDARY_STATUS,
+                                       Type::i32,
+                                       builder.makeConst(int32_t(0)),
+                                       Builder::Mutable));
+  primary.addExport(builder.makeExport(
+    LOAD_SECONDARY_STATUS, LOAD_SECONDARY_STATUS, ExternalKind::Global));
+}
+
 std::unique_ptr<Module> ModuleSplitter::initSecondary(const Module& primary) {
   // Create the secondary module and copy trivial properties.
   auto secondary = std::make_unique<Module>();
@@ -318,7 +346,12 @@ std::pair<std::set<Name>, std::set<Name>>
 ModuleSplitter::classifyFunctions(const Module& primary, const Config& config) {
   std::set<Name> primaryFuncs, secondaryFuncs;
   for (auto& func : primary.functions) {
-    if (func->imported() || config.primaryFuncs.count(func->name)) {
+    // In JSPI mode exported functions cannot be moved to the secondary
+    // module since that would make them async when they may not have the JSPI
+    // wrapper. Exported JSPI functions can still benefit from splitting though
+    // since only the JSPI wrapper stub will remain in the primary module.
+    if (func->imported() || config.primaryFuncs.count(func->name) ||
+        (config.jspi && ExportUtils::isExported(primary, *func))) {
       primaryFuncs.insert(func->name);
     } else {
       assert(func->name != primary.start && "The start function must be kept");
@@ -409,6 +442,20 @@ void ModuleSplitter::thunkExportedSecondaryFunctions() {
   }
 }
 
+Expression* ModuleSplitter::maybeLoadSecondary(Builder& builder,
+                                               Expression* callIndirect) {
+  if (!config.jspi) {
+    return callIndirect;
+  }
+  // Check if the secondary module is loaded and if it isn't, call the
+  // function to load it.
+  auto* loadSecondary = builder.makeIf(
+    builder.makeUnary(EqZInt32,
+                      builder.makeGlobalGet(LOAD_SECONDARY_STATUS, Type::i32)),
+    builder.makeCall(internalLoadSecondaryModule, {}, Type::none));
+  return builder.makeSequence(loadSecondary, callIndirect);
+}
+
 void ModuleSplitter::indirectCallsToSecondaryFunctions() {
   // Update direct calls of secondary functions to be indirect calls of their
   // corresponding table indices instead.
@@ -425,12 +472,14 @@ void ModuleSplitter::indirectCallsToSecondaryFunctions() {
       }
       auto* func = parent.secondary.getFunction(curr->target);
       auto tableSlot = parent.tableManager.getSlot(curr->target, func->type);
-      replaceCurrent(
+
+      replaceCurrent(parent.maybeLoadSecondary(
+        builder,
         builder.makeCallIndirect(tableSlot.tableName,
                                  tableSlot.makeExpr(parent.primary),
                                  curr->operands,
                                  func->type,
-                                 curr->isReturn));
+                                 curr->isReturn)));
     }
     void visitRefFunc(RefFunc* curr) {
       assert(false && "TODO: handle ref.func as well");

src/ir/module-splitting.h

@@ -44,6 +44,8 @@
 
 namespace wasm::ModuleSplitting {
 
+static const Name LOAD_SECONDARY_MODULE("__load_secondary_module");
+
 struct Config {
   // The set of functions to keep in the primary module. All others are split
   // out into the new secondary module. Must include the start function if it
@@ -64,6 +66,9 @@ struct Config {
   // false, the original function names will be used (after `newExportPrefix`)
   // as the new export names.
   bool minimizeNewExportNames = false;
+  // When JSPI support is enabled the secondary module loading is handled by an
+  // imported function.
+  bool jspi = false;
 };
 
 struct Results {

src/passes/JSPI.cpp

@@ -18,6 +18,7 @@
 #include "ir/element-utils.h"
 #include "ir/import-utils.h"
 #include "ir/literal-utils.h"
+#include "ir/module-splitting.h"
 #include "ir/names.h"
 #include "ir/utils.h"
 #include "pass.h"
@@ -49,6 +50,12 @@
 //      Wrap each export in the comma-separated list. Similar to jspi-imports,
 //      wildcards and separate files are supported.
 //
+//   --pass-arg=jspi-split-module
+//
+//      Enables integration with wasm-split. A JSPI'ed function named
+//      `__load_secondary_module` will be injected that is used by wasm-split to
+//      load a secondary module.
+//
 
 namespace wasm {
 
@@ -87,6 +94,22 @@ struct JSPI : public Pass {
       options.getArgumentOrDefault("jspi-exports", "")));
     String::Split listedExports(stateChangingExports, ",");
 
+    bool wasmSplit = options.hasArgument("jspi-split-module");
+    if (wasmSplit) {
+      // Make an import for the load secondary module function so a JSPI wrapper
+      // version will be created.
+      auto import =
+        Builder::makeFunction(ModuleSplitting::LOAD_SECONDARY_MODULE,
+                              Signature(Type::none, Type::none),
+                              {});
+      import->module = ENV;
+      import->base = ModuleSplitting::LOAD_SECONDARY_MODULE;
+      module->addFunction(std::move(import));
+      listedImports.push_back(
+        ENV.toString() + "." +
+        ModuleSplitting::LOAD_SECONDARY_MODULE.toString());
+    }
+
     // Create a global to store the suspender that is passed into exported
     // functions and will then need to be passed out to the imported functions.
     Name suspender = Names::getValidGlobalName(*module, "suspender");
@@ -128,7 +151,7 @@ struct JSPI : public Pass {
       if (im->imported() &&
           canChangeState(getFullFunctionName(im->module, im->base),
                          listedImports)) {
-        makeWrapperForImport(im, module, suspender);
+        makeWrapperForImport(im, module, suspender, wasmSplit);
       }
     }
   }
@@ -181,7 +204,10 @@ struct JSPI : public Pass {
     return module->addFunction(std::move(wrapperFunc))->name;
   }
 
-  void makeWrapperForImport(Function* im, Module* module, Name suspender) {
+  void makeWrapperForImport(Function* im,
+                            Module* module,
+                            Name suspender,
+                            bool wasmSplit) {
     Builder builder(*module);
     auto wrapperIm = make_unique<Function>();
     wrapperIm->name = Names::getValidFunctionName(
@@ -234,6 +260,15 @@ struct JSPI : public Pass {
     stub->body = block;
     wrapperIm->type = Signature(Type(params), call->type);
 
+    if (wasmSplit && im->name == ModuleSplitting::LOAD_SECONDARY_MODULE) {
+      // In non-debug builds the name of the JSPI wrapper function for loading
+      // the secondary module will be removed. Create an export of it so that
+      // wasm-split can find it.
+      module->addExport(
+        builder.makeExport(ModuleSplitting::LOAD_SECONDARY_MODULE,
+                           ModuleSplitting::LOAD_SECONDARY_MODULE,
+                           ExternalKind::Function));
+    }
     module->removeFunction(im->name);
     module->addFunction(std::move(stub));
     module->addFunction(std::move(wrapperIm));

src/tools/wasm-split/split-options.cpp

@@ -205,15 +205,14 @@ WasmSplitOptions::WasmSplitOptions()
          [&](Options* o, const std::string& argument) {
            placeholderNamespace = argument;
          })
-    .add(
-      "--asyncify",
-      "",
-      "Transform the module to support unwinding the stack from placeholder "
-      "functions and rewinding it once the secondary module has been loaded.",
-      WasmSplitOption,
-      {Mode::Split},
-      Options::Arguments::Zero,
-      [&](Options* o, const std::string& argument) { asyncify = true; })
+    .add("--jspi",
+         "",
+         "Transform the module to support asynchronously loading the secondary "
+         "module before any placeholder functions have been called.",
+         WasmSplitOption,
+         {Mode::Split},
+         Options::Arguments::Zero,
+         [&](Options* o, const std::string& argument) { jspi = true; })
     .add(
       "--export-prefix",
       "",

src/tools/wasm-split/split-options.h

@@ -46,7 +46,7 @@ struct WasmSplitOptions : ToolOptions {
   bool emitBinary = true;
   bool symbolMap = false;
   bool placeholderMap = false;
-  bool asyncify = false;
+  bool jspi = false;
 
   // TODO: Remove this. See the comment in wasm-binary.h.
   bool emitModuleNames = false;

src/tools/wasm-split/wasm-split.cpp

@@ -248,6 +248,11 @@ void splitModule(const WasmSplitOptions& options) {
     }
   }
 
+  if (options.jspi) {
+    // The load secondary module function must be kept in the main module.
+    keepFuncs.insert(ModuleSplitting::LOAD_SECONDARY_MODULE);
+  }
+
   if (!options.quiet && keepFuncs.size() == 0) {
     std::cerr << "warning: not keeping any functions in the primary module\n";
   }
@@ -300,22 +305,13 @@ void splitModule(const WasmSplitOptions& options) {
     config.newExportPrefix = options.exportPrefix;
   }
   config.minimizeNewExportNames = !options.passOptions.debugInfo;
+  config.jspi = options.jspi;
   auto splitResults = ModuleSplitting::splitFunctions(wasm, config);
   auto& secondary = splitResults.secondary;
 
   adjustTableSize(wasm, options.initialTableSize);
   adjustTableSize(*secondary, options.initialTableSize);
 
-  // Run asyncify on the primary module
-  if (options.asyncify) {
-    PassOptions passOptions;
-    passOptions.optimizeLevel = 1;
-    passOptions.arguments.insert({"asyncify-ignore-imports", ""});
-    PassRunner runner(&wasm, passOptions);
-    runner.add("asyncify");
-    runner.run();
-  }
-
   if (options.symbolMap) {
     writeSymbolMap(wasm, options.primaryOutput + ".symbols");
     writeSymbolMap(*secondary, options.secondaryOutput + ".symbols");

test/lit/help/wasm-split.test

@@ -66,10 +66,10 @@
 ;; CHECK-NEXT:                                        import placeholder functions into the
 ;; CHECK-NEXT:                                        primary module.
 ;; CHECK-NEXT:
-;; CHECK-NEXT:   --asyncify                           [split] Transform the module to support
-;; CHECK-NEXT:                                        unwinding the stack from placeholder
-;; CHECK-NEXT:                                        functions and rewinding it once the
-;; CHECK-NEXT:                                        secondary module has been loaded.
+;; CHECK-NEXT:   --jspi                               [split] Transform the module to support
+;; CHECK-NEXT:                                        asynchronously loading the secondary
+;; CHECK-NEXT:                                        module before any placeholder functions
+;; CHECK-NEXT:                                        have been called.
 ;; CHECK-NEXT:
 ;; CHECK-NEXT:   --export-prefix                      [split] An identifying prefix to prepend
 ;; CHECK-NEXT:                                        to new export names created by module

test/lit/passes/jspi-split-module.wast

@@ -0,0 +1,30 @@
+;; NOTE: Assertions have been generated by update_lit_checks.py --all-items and should not be edited.
+;; RUN: wasm-opt %s --jspi --pass-arg=jspi-split-module -all -S -o - | filecheck %s
+
+;; The following should be generated besides the usual JSPI wasm:
+;;  - function import
+;;  - JSPI'ed version of the import
+;;  - export of the above
+(module)
+;; CHECK:      (type $none_=>_none (func))
+
+;; CHECK:      (type $externref_=>_none (func (param externref)))
+
+;; CHECK:      (import "env" "__load_secondary_module" (func $import$__load_secondary_module (param externref)))
+
+;; CHECK:      (global $suspender (mut externref) (ref.null noextern))
+
+;; CHECK:      (export "__load_secondary_module" (func $__load_secondary_module))
+
+;; CHECK:      (func $__load_secondary_module (type $none_=>_none)
+;; CHECK-NEXT:  (local $0 externref)
+;; CHECK-NEXT:  (local.set $0
+;; CHECK-NEXT:   (global.get $suspender)
+;; CHECK-NEXT:  )
+;; CHECK-NEXT:  (call $import$__load_secondary_module
+;; CHECK-NEXT:   (global.get $suspender)
+;; CHECK-NEXT:  )
+;; CHECK-NEXT:  (global.set $suspender
+;; CHECK-NEXT:   (local.get $0)
+;; CHECK-NEXT:  )
+;; CHECK-NEXT: )