Skip to content

[jspi] - Support new version of JSPI for module splitting. #6546

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
brendandahl merged 1 commit into from
Apr 29, 2024
Merged

[jspi] - Support new version of JSPI for module splitting. #6546

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
26 changes: 20 additions & 6 deletions src/ir/module-splitting.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -68,6 +68,7 @@
// not yet support passive table segments anyway).

#include "ir/module-splitting.h"
#include "asmjs/shared-constants.h"
#include "ir/element-utils.h"
#include "ir/export-utils.h"
#include "ir/manipulation.h"
Expand Down Expand Up @@ -318,12 +319,25 @@ 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);
// Support the first version of JSPI, where the JSPI pass added the load
// secondary module export.
// TODO: remove this when the new JSPI API is only supported.
if (primary.getExportOrNull(LOAD_SECONDARY_MODULE)) {
internalLoadSecondaryModule =
primary.getExport(LOAD_SECONDARY_MODULE)->value;
// Remove the exported LOAD_SECONDARY_MODULE function since it's only needed
// internally.
primary.removeExport(LOAD_SECONDARY_MODULE);
} else {
// Add an imported function to load the secondary module.
auto import = Builder::makeFunction(ModuleSplitting::LOAD_SECONDARY_MODULE,
Signature(Type::none, Type::none),
{});
import->module = ENV;
import->base = ModuleSplitting::LOAD_SECONDARY_MODULE;
primary.addFunction(std::move(import));
internalLoadSecondaryModule = ModuleSplitting::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,
Expand Down
107 changes: 41 additions & 66 deletions test/lit/wasm-split/jspi-secondary-export.wast
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -1,86 +1,61 @@
;; 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-exports@foo --pass-arg=jspi-split-module -all -S -o %t.jspi.wast
;; RUN: wasm-split %t.jspi.wast --export-prefix='%' -g -o1 %t.1.wasm -o2 %t.2.wasm --jspi --enable-reference-types
;; RUN: wasm-split %s --export-prefix='%' -g -o1 %t.1.wasm -o2 %t.2.wasm --jspi
;; RUN: wasm-dis %t.1.wasm | filecheck %s --check-prefix PRIMARY
;; RUN: wasm-dis %t.2.wasm | filecheck %s --check-prefix SECONDARY

;; Ensure exported functions are not moved to secondary module when JSPI is
;; enabled.

(module
;; PRIMARY: (type $0 (func (param i32) (result i32)))

;; PRIMARY: (type $1 (func))

;; PRIMARY: (import "env" "__load_secondary_module" (func $fimport$0))

;; PRIMARY: (import "placeholder" "0" (func $placeholder_0 (param i32) (result i32)))

;; PRIMARY: (global $global$0 (mut i32) (i32.const 0))

;; PRIMARY: (table $0 1 funcref)

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

;; PRIMARY: (export "foo" (func $foo))
(export "foo" (func $foo))
;; PRIMARY: (export "load_secondary_module_status" (global $global$0))

;; PRIMARY: (export "%table" (table $0))

;; PRIMARY: (func $foo (param $0 i32) (result i32)
;; PRIMARY-NEXT: (if
;; PRIMARY-NEXT: (i32.eqz
;; PRIMARY-NEXT: (global.get $global$0)
;; PRIMARY-NEXT: )
;; PRIMARY-NEXT: (then
;; PRIMARY-NEXT: (call $fimport$0)
;; PRIMARY-NEXT: )
;; PRIMARY-NEXT: )
;; PRIMARY-NEXT: (call_indirect (type $0)
;; PRIMARY-NEXT: (local.get $0)
;; PRIMARY-NEXT: (i32.const 0)
;; PRIMARY-NEXT: )
;; PRIMARY-NEXT: )
(func $foo (param i32) (result i32)
(call $bar (local.get $0))
Copy link
Member

@tlively tlively Apr 29, 2024

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

This should be local.get 0

)
;; SECONDARY: (type $0 (func (param i32) (result i32)))

;; SECONDARY: (import "primary" "%table" (table $timport$0 1 funcref))

;; SECONDARY: (import "primary" "%global" (global $suspender (mut externref)))

;; SECONDARY: (import "primary" "load_secondary_module_status" (global $gimport$1 (mut i32)))
;; SECONDARY: (import "primary" "load_secondary_module_status" (global $gimport$0 (mut i32)))

;; SECONDARY: (elem $0 (i32.const 0) $foo)
;; SECONDARY: (elem $0 (i32.const 0) $bar)

;; SECONDARY: (func $foo (param $0 i32) (result i32)
;; SECONDARY: (func $bar (param $0 i32) (result i32)
;; SECONDARY-NEXT: (i32.const 0)
;; SECONDARY-NEXT: )
(func $foo (param i32) (result i32)
(func $bar (param i32) (result i32)
(i32.const 0)
)
)
;; PRIMARY: (type $0 (func (param i32) (result i32)))

;; PRIMARY: (type $3 (func (param externref)))

;; PRIMARY: (type $1 (func (param externref i32) (result i32)))

;; PRIMARY: (type $2 (func))

;; PRIMARY: (import "env" "__load_secondary_module" (func $import$__load_secondary_module (param externref)))

;; PRIMARY: (import "placeholder" "0" (func $placeholder_0 (param i32) (result i32)))

;; PRIMARY: (global $suspender (mut externref) (ref.null noextern))

;; PRIMARY: (global $global$1 (mut i32) (i32.const 0))

;; PRIMARY: (table $0 1 funcref)

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

;; PRIMARY: (export "foo" (func $export$foo))

;; PRIMARY: (export "load_secondary_module_status" (global $global$1))

;; PRIMARY: (export "%table" (table $0))

;; PRIMARY: (export "%global" (global $suspender))

;; PRIMARY: (func $export$foo (param $susp externref) (param $0 i32) (result i32)
;; PRIMARY-NEXT: (global.set $suspender
;; PRIMARY-NEXT: (local.get $susp)
;; PRIMARY-NEXT: )
;; PRIMARY-NEXT: (if
;; PRIMARY-NEXT: (i32.eqz
;; PRIMARY-NEXT: (global.get $global$1)
;; PRIMARY-NEXT: )
;; PRIMARY-NEXT: (then
;; PRIMARY-NEXT: (call $__load_secondary_module)
;; PRIMARY-NEXT: )
;; PRIMARY-NEXT: )
;; PRIMARY-NEXT: (call_indirect (type $0)
;; PRIMARY-NEXT: (local.get $0)
;; PRIMARY-NEXT: (i32.const 0)
;; PRIMARY-NEXT: )
;; PRIMARY-NEXT: )

;; PRIMARY: (func $__load_secondary_module
;; PRIMARY-NEXT: (local $0 externref)
;; PRIMARY-NEXT: (local.set $0
;; PRIMARY-NEXT: (global.get $suspender)
;; PRIMARY-NEXT: )
;; PRIMARY-NEXT: (call $import$__load_secondary_module
;; PRIMARY-NEXT: (global.get $suspender)
;; PRIMARY-NEXT: )
;; PRIMARY-NEXT: (global.set $suspender
;; PRIMARY-NEXT: (local.get $0)
;; PRIMARY-NEXT: )
;; PRIMARY-NEXT: )
56 changes: 11 additions & 45 deletions test/lit/wasm-split/jspi.wast
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -1,6 +1,5 @@
;; 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-exports@foo --pass-arg=jspi-split-module -all -S -o %t.jspi.wast
;; RUN: wasm-split %t.jspi.wast --export-prefix='%' -g -o1 %t.1.wasm -o2 %t.2.wasm --keep-funcs=foo --jspi --enable-reference-types
;; RUN: wasm-split %s --export-prefix='%' -g -o1 %t.1.wasm -o2 %t.2.wasm --keep-funcs=foo --jspi
;; RUN: wasm-dis %t.1.wasm | filecheck %s --check-prefix PRIMARY
;; RUN: wasm-dis %t.2.wasm | filecheck %s --check-prefix SECONDARY

Expand All @@ -10,41 +9,31 @@
(module
;; PRIMARY: (type $0 (func (param i32) (result i32)))

;; PRIMARY: (type $3 (func (param externref)))
;; PRIMARY: (type $1 (func))

;; PRIMARY: (type $1 (func (param externref i32) (result i32)))

;; PRIMARY: (type $2 (func))

;; PRIMARY: (import "env" "__load_secondary_module" (func $import$__load_secondary_module (param externref)))
;; PRIMARY: (import "env" "__load_secondary_module" (func $fimport$0))

;; PRIMARY: (import "placeholder" "0" (func $placeholder_0 (param i32) (result i32)))

;; PRIMARY: (global $suspender (mut externref) (ref.null noextern))

;; PRIMARY: (global $global$1 (mut i32) (i32.const 0))
;; PRIMARY: (global $global$0 (mut i32) (i32.const 0))

;; PRIMARY: (table $0 1 funcref)

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

;; PRIMARY: (export "foo" (func $export$foo))

;; PRIMARY: (export "load_secondary_module_status" (global $global$1))

;; PRIMARY: (export "%foo" (func $foo))
;; PRIMARY: (export "foo" (func $foo))
(export "foo" (func $foo))
;; PRIMARY: (export "%table" (table $0))
;; PRIMARY: (export "load_secondary_module_status" (global $global$0))

;; PRIMARY: (export "%global" (global $suspender))
;; PRIMARY: (export "%table" (table $0))

;; PRIMARY: (func $foo (param $0 i32) (result i32)
;; PRIMARY-NEXT: (if
;; PRIMARY-NEXT: (i32.eqz
;; PRIMARY-NEXT: (global.get $global$1)
;; PRIMARY-NEXT: (global.get $global$0)
;; PRIMARY-NEXT: )
;; PRIMARY-NEXT: (then
;; PRIMARY-NEXT: (call $__load_secondary_module)
;; PRIMARY-NEXT: (call $fimport$0)
;; PRIMARY-NEXT: )
;; PRIMARY-NEXT: )
;; PRIMARY-NEXT: (call_indirect (type $0)
Expand All @@ -59,11 +48,9 @@

;; SECONDARY: (import "primary" "%table" (table $timport$0 1 funcref))

;; SECONDARY: (import "primary" "%global" (global $suspender (mut externref)))

;; SECONDARY: (import "primary" "load_secondary_module_status" (global $gimport$1 (mut i32)))
;; SECONDARY: (import "primary" "load_secondary_module_status" (global $gimport$0 (mut i32)))

;; SECONDARY: (import "primary" "%foo" (func $foo (param i32) (result i32)))
;; SECONDARY: (import "primary" "foo" (func $foo (param i32) (result i32)))

;; SECONDARY: (elem $0 (i32.const 0) $bar)

Expand All @@ -77,24 +64,3 @@
)
)

;; PRIMARY: (func $export$foo (param $susp externref) (param $0 i32) (result i32)
;; PRIMARY-NEXT: (global.set $suspender
;; PRIMARY-NEXT: (local.get $susp)
;; PRIMARY-NEXT: )
;; PRIMARY-NEXT: (call $foo
;; PRIMARY-NEXT: (local.get $0)
;; PRIMARY-NEXT: )
;; PRIMARY-NEXT: )

;; PRIMARY: (func $__load_secondary_module
;; PRIMARY-NEXT: (local $0 externref)
;; PRIMARY-NEXT: (local.set $0
;; PRIMARY-NEXT: (global.get $suspender)
;; PRIMARY-NEXT: )
;; PRIMARY-NEXT: (call $import$__load_secondary_module
;; PRIMARY-NEXT: (global.get $suspender)
;; PRIMARY-NEXT: )
;; PRIMARY-NEXT: (global.set $suspender
;; PRIMARY-NEXT: (local.get $0)
;; PRIMARY-NEXT: )
;; PRIMARY-NEXT: )