Experimental WebAssembly multithreading support

ictrobot · ictrobot · commit f18220e7fe4c · 2024-09-28T17:40:27.000+01:00
diff --git a/Cargo.lock b/Cargo.lock
diff --git a/crates/aoc_wasm/Cargo.toml b/crates/aoc_wasm/Cargo.toml
@@ -9,6 +9,10 @@ rust-version = { workspace = true }
 
 [dependencies]
 aoc = { path = "../aoc" }
+utils = { path = "../utils", optional = true }
+
+[features]
+multithreading = ["utils/wasm-multithreading"]
 
 [lints]
 workspace = true
diff --git a/crates/aoc_wasm/src/lib.rs b/crates/aoc_wasm/src/lib.rs
@@ -1,6 +1,8 @@
 //! Simple WebAssembly interface without external libraries.
 
 mod custom_sections;
+#[cfg(feature = "multithreading")]
+mod multithreading;
 
 use aoc::all_puzzles;
 use aoc::utils::input::InputType;
diff --git a/crates/aoc_wasm/src/multithreading.rs b/crates/aoc_wasm/src/multithreading.rs
@@ -0,0 +1,17 @@
+use aoc::utils::wasm::scoped_tasks::worker;
+use std::alloc::{alloc_zeroed, Layout};
+
+/// Allocate stack for worker threads.
+///
+/// **WARNING**: Stack overflows on worker threads will corrupt other parts of the linear memory.
+#[no_mangle]
+extern "C" fn allocate_stack(size: usize, align: usize) -> *mut u8 {
+    let layout = Layout::from_size_align(size, align).unwrap();
+    unsafe { alloc_zeroed(layout) }
+}
+
+/// Run worker thread.
+#[no_mangle]
+extern "C" fn worker_thread() {
+    worker();
+}
diff --git a/crates/aoc_wasm/web/aoc.mjs b/crates/aoc_wasm/web/aoc.mjs
@@ -6,6 +6,12 @@
  * @property {WebAssembly.Global} PART1
  * @property {WebAssembly.Global} PART2
  * @property {number} PUZZLES
+ * @property {WebAssembly.Global} [__tls_size]                          If multithreaded
+ * @property {WebAssembly.Global} [__tls_align]                         If multithreaded
+ * @property {WebAssembly.Global} [__tls_base]                          If multithreaded
+ * @property {WebAssembly.Global} [__stack_pointer]                     If multithreaded
+ * @property {(size: number, align: number) => number} [allocate_stack] If multithreaded
+ * @property {() => void} [worker_thread]                               If multithreaded
  */
 
 /**
@@ -22,12 +28,16 @@
 const BUFFER_SIZE = 1024 * 1024;
 
 export class Aoc {
+    /** @type {boolean} */
+    #multithreaded;
     /** @type {WebAssembly.Module} */
     #module;
     /** @type {WebAssembly.Instance} */
     #instance;
-    /** @type {Puzzles} */
-    #puzzles;
+    /** @type {WebAssembly.Memory} */
+    #memory;
+    /** @type {Worker[]} */
+    #workers;
 
     /**
      * @param {WebAssembly.Module} module
@@ -90,14 +100,66 @@ export class Aoc {
      * @param {WebAssembly.Instance} [instance]
      */
     constructor(module, instance) {
-        this.#module = module;
-        this.#instance = instance ?? new WebAssembly.Instance(module);
+        const imports = WebAssembly.Module.imports(module);
+        if (imports.length === 0) {
+            this.#multithreaded = false;
+            this.#module = module;
+            this.#instance = instance ?? new WebAssembly.Instance(module);
+            this.#memory = this.#exports.memory;
+        } else if (imports.length === 1 && imports[0].module === "env" && imports[0].name === "memory" && imports[0].kind === "memory") {
+            this.#multithreaded = true;
+            this.#module = module;
+            if (instance) throw new Error("Instance cannot be provided for multithreaded modules");
+            this.newInstance();
+        } else {
+            throw new Error("Unsupported module");
+        }
     }
 
-    /** @return {Puzzles} */
-    get puzzles() {
-        this.#puzzles ??= Aoc.puzzleList(this.#module);
-        return this.#puzzles;
+    /** @param {number} [numWorkers] */
+    newInstance(numWorkers) {
+        if (this.#multithreaded) {
+            if (this.#workers?.length > 0) {
+                // Stop existing workers
+                for (const worker of this.#workers) {
+                    worker.terminate();
+                }
+                numWorkers ??= this.#workers.length;
+                this.#workers = [];
+            }
+            numWorkers ??= navigator.hardwareConcurrency;
+
+            this.#memory = new WebAssembly.Memory({initial: 96, maximum: 2048, shared: true});
+            this.#instance = new WebAssembly.Instance(this.#module, {env: {memory: this.#memory}});
+
+            // Stack alignment must be at least 16 bytes.
+            //
+            // Only aligning the stack to 8 bytes (this.#exports.__tls_align.value at the time of writing) causes 2016
+            // day 14 to inconsistently return wrong answers in release builds as the optimizer uses `i32.or` instead of
+            // `i32.add` when adding on small array indexes.
+            let align = Math.max(16, this.#exports.__tls_align.value);
+            let tlsSize = Math.ceil(this.#exports.__tls_size.value / align) * align;
+            let stackSize = Math.ceil(this.#exports.__stack_pointer.value / align) * align;
+
+            // Use a single allocation for stack & tls, using the first stackSize bytes for the stack and the remaining
+            // tlsSize bytes for thread local storage. This makes __tls_base and __stack_pointer the same value (as
+            // the stack grows downwards and TLS is above __tls_base), similar to the main thread.
+            //
+            // Allocate all the stacks at once to avoid memory growing as workers start, which seems to cause problems.
+            const stacks = [];
+            for (let i = 0; i < numWorkers; i++) {
+                stacks.push(this.#exports.allocate_stack(stackSize + tlsSize, align));
+            }
+
+            this.#workers = [];
+            for (let i = 0; i < numWorkers; i++) {
+                const worker = new Worker("./worker.mjs", {type: "module"});
+                worker.postMessage(["thread", this.#module, this.#memory, stacks[i] + stackSize]);
+                this.#workers.push(worker);
+            }
+        } else {
+            this.#instance = new WebAssembly.Instance(this.#module);
+        }
     }
 
     /**
@@ -115,7 +177,7 @@ export class Aoc {
             this.#write(input);
             success = this.#exports.run_puzzle(year, day, isExample, part1, part2);
         } catch (e) {
-            this.#instance = new WebAssembly.Instance(this.#module);
+            this.newInstance();
             return {
                 success: false,
                 error: "Unexpected error: " + e.toString() + (e.stack ? "\n\n" + e.stack : ""),
@@ -147,18 +209,29 @@ export class Aoc {
      */
     #buffer(type) {
         const address = this.#exports[type].value;
-        return new Uint8Array(this.#exports.memory.buffer)
+        return new Uint8Array(this.#memory.buffer)
             .subarray(address, address + BUFFER_SIZE);
     }
 
     /** @param {string} input */
     #write(input) {
         const buffer = this.#buffer("INPUT");
-        const result = new TextEncoder().encodeInto(input, buffer);
-        if (result.read < input.length || result.written === buffer.length) {
-            throw new Error("Input string is too long");
+        if (this.#multithreaded) {
+            // Can't encode directly into SharedArrayBuffer
+            const temp = new Uint8Array(BUFFER_SIZE);
+            const result = new TextEncoder().encodeInto(input, temp);
+            if (result.read < input.length || result.written === buffer.length) {
+                throw new Error("Input string is too long");
+            }
+            buffer.set(temp.subarray(0, result.written));
+            buffer[result.written] = 0;
+        } else {
+            const result = new TextEncoder().encodeInto(input, buffer);
+            if (result.read < input.length || result.written === buffer.length) {
+                throw new Error("Input string is too long");
+            }
+            buffer[result.written] = 0;
         }
-        buffer[result.written] = 0;
     }
 
     /**
@@ -173,6 +246,13 @@ export class Aoc {
             buffer = buffer.subarray(0, end);
         }
 
+        if (this.#multithreaded) {
+            // Can't decode directly from SharedArrayBuffer
+            const temp = new Uint8Array(buffer.length);
+            temp.set(buffer);
+            buffer = temp;
+        }
+
         return (new TextDecoder()).decode(buffer);
     }
 }
diff --git a/crates/aoc_wasm/web/cross-origin-isolation-service-worker.js b/crates/aoc_wasm/web/cross-origin-isolation-service-worker.js
@@ -0,0 +1,38 @@
+// This is a hack to set Cross-Origin-Embedder-Policy/Cross-Origin-Opener-Policy (required to send SharedArrayBuffer
+// instances across workers/use multithreaded WebAssembly) on GitHub pages
+
+if (typeof window === "undefined") {
+    self.addEventListener("install", () => self.skipWaiting());
+    self.addEventListener("activate", event => event.waitUntil(self.clients.claim()));
+
+    self.addEventListener("fetch", event => {
+        event.respondWith(fetch(event.request).then(response =>  {
+            if (response.status === 0) return response;
+
+            const headers = new Headers(response.headers);
+            headers.set('Cross-Origin-Embedder-Policy', 'require-corp');
+            headers.set('Cross-Origin-Opener-Policy', 'same-origin');
+
+            return new Response(response.body, {status: response.status, statusText: response.statusText, headers});
+        }));
+    });
+} else if (window.crossOriginIsolated && !navigator.serviceWorker?.controller) {
+    console.log("cross-origin isolated without service worker"); // No service worker workaround needed
+} else {
+    if (window.crossOriginIsolated) {
+        console.log("cross-origin isolated using service worker"); // Still re-register
+    } else {
+        console.log("not cross-origin isolated, trying service worker");
+    }
+
+    navigator.serviceWorker.register(document.currentScript.src).then(registration => {
+        console.log(`service worker registered with scope ${registration.scope}`);
+
+        registration.addEventListener("updatefound", () => {
+            console.log("reloading page due to service worker update");
+            window.location.reload();
+        });
+    }, error => {
+        console.error("error registering service worker", error);
+    });
+}
diff --git a/crates/aoc_wasm/web/index.html b/crates/aoc_wasm/web/index.html
@@ -68,10 +68,11 @@
             }
         }
     </style>
+    <script src="cross-origin-isolation-service-worker.js" defer></script>
     <script src="./web.mjs" type="module"></script>
     <link rel="modulepreload" href="./aoc.mjs" />
     <link rel="modulepreload" href="./worker.mjs" />
-    <link rel="preload" href="./aoc-simd128.wasm" as="fetch" crossorigin />
+    <link rel="preload" href="./aoc-threads.wasm" as="fetch" crossorigin />
 </head>
 <body class="is-flex is-flex-direction-column">
     <nav class="navbar" role="navigation" aria-label="main navigation">
diff --git a/crates/aoc_wasm/web/web.mjs b/crates/aoc_wasm/web/web.mjs
@@ -4,6 +4,9 @@ const MODULE_PATHS = [
     "./aoc-simd128.wasm",
     "./aoc.wasm",
 ];
+if (window.crossOriginIsolated) {
+    MODULE_PATHS.unshift("./aoc-threads.wasm");
+}
 
 let module;
 for (const path of MODULE_PATHS) {
diff --git a/crates/aoc_wasm/web/worker.mjs b/crates/aoc_wasm/web/worker.mjs
@@ -15,5 +15,12 @@ onmessage = (e) => {
             postMessage(result);
             console.log(result);
             break;
+        case "thread":
+            const [module, memory, ptr] = e.data;
+            instance = new WebAssembly.Instance(module, {env: {memory}});
+            instance.exports.__stack_pointer.value = ptr; // Stack uses storage below the provided pointer
+            instance.exports.__wasm_init_tls(ptr); // TLS uses storage above the provided pointer
+            instance.exports.worker_thread();
+            throw new Error("unreachable");
     }
 };
diff --git a/crates/utils/Cargo.toml b/crates/utils/Cargo.toml
@@ -11,6 +11,7 @@ rust-version = { workspace = true }
 
 [features]
 unsafe = []
+wasm-multithreading = ["unsafe"]
 
 [lints]
 workspace = true
diff --git a/crates/utils/src/lib.rs b/crates/utils/src/lib.rs
@@ -7,15 +7,19 @@ mod framework;
 pub mod grid;
 pub mod input;
 pub mod md5;
-#[cfg_attr(target_family = "wasm", path = "multithreading_wasm.rs")]
+#[cfg(not(target_family = "wasm"))]
 pub mod multithreading;
 pub mod multiversion;
 pub mod number;
 pub mod parser;
 pub mod point;
 pub mod simd;
+#[cfg(target_family = "wasm")]
+pub mod wasm;
 
 pub use framework::{Puzzle, PuzzleExamples};
+#[cfg(target_family = "wasm")]
+pub use wasm::multithreading;
 
 /// Standard imports for puzzle solutions.
 pub mod prelude {
diff --git a/crates/utils/src/multithreading.rs b/crates/utils/src/multithreading.rs
@@ -12,6 +12,7 @@ static NUM_THREADS: AtomicUsize = AtomicUsize::new(0);
 /// Get the number of worker threads to use.
 ///
 /// Defaults to [`std::thread::available_parallelism`] unless set by [`set_thread_count`].
+#[must_use]
 pub fn get_thread_count() -> NonZeroUsize {
     if let Some(threads) = NonZeroUsize::new(NUM_THREADS.load(Relaxed)) {
         threads
diff --git a/crates/utils/src/wasm/mod.rs b/crates/utils/src/wasm/mod.rs
@@ -0,0 +1,8 @@
+//! WebAssembly helpers.
+
+#[cfg(feature = "wasm-multithreading")]
+pub mod scoped_tasks;
+
+#[cfg_attr(feature = "wasm-multithreading", path = "multithreading_wasm.rs")]
+#[cfg_attr(not(feature = "wasm-multithreading"), path = "multithreading_stub.rs")]
+pub mod multithreading;
diff --git a/crates/utils/src/wasm/multithreading_stub.rs b/crates/utils/src/wasm/multithreading_stub.rs
@@ -2,6 +2,7 @@
 
 use std::num::NonZeroUsize;
 
+#[must_use]
 pub fn get_thread_count() -> NonZeroUsize {
     NonZeroUsize::new(1).unwrap()
 }
diff --git a/crates/utils/src/wasm/multithreading_wasm.rs b/crates/utils/src/wasm/multithreading_wasm.rs
@@ -0,0 +1,23 @@
+//! Experimental threading helpers for WebAssembly.
+
+use super::scoped_tasks::{scope, worker_count};
+use std::num::{NonZero, NonZeroUsize};
+
+#[must_use]
+pub fn get_thread_count() -> NonZeroUsize {
+    // If there are no workers, `scoped_task` will fall back to running tasks on the current thread.
+    NonZeroUsize::new(worker_count()).unwrap_or(NonZero::new(1).unwrap())
+}
+
+pub fn set_thread_count(_: NonZeroUsize) {
+    unreachable!();
+}
+
+pub fn worker_pool(worker: impl Fn() + Copy + Send) {
+    let threads = get_thread_count().get();
+    scope(|scope| {
+        for _ in 0..threads {
+            scope.spawn(worker);
+        }
+    });
+}
diff --git a/crates/utils/src/wasm/scoped_tasks.rs b/crates/utils/src/wasm/scoped_tasks.rs
diff --git a/crates/xtask/src/cmd/web.rs b/crates/xtask/src/cmd/web.rs

Original file line number	Diff line number	Diff line change
`@@ -68,10 +68,11 @@`
`68`	`68`	`}`
`69`	`69`	`}`
`70`	`70`	`</style>`
	`71`	`+ <script src="cross-origin-isolation-service-worker.js" defer></script>`
`71`	`72`	`<script src="./web.mjs" type="module"></script>`
`72`	`73`	`<link rel="modulepreload" href="./aoc.mjs" />`
`73`	`74`	`<link rel="modulepreload" href="./worker.mjs" />`
`74`		`- <link rel="preload" href="./aoc-simd128.wasm" as="fetch" crossorigin />`
	`75`	`+ <link rel="preload" href="./aoc-threads.wasm" as="fetch" crossorigin />`
`75`	`76`	`</head>`
`76`	`77`	`<body class="is-flex is-flex-direction-column">`
`77`	`78`	`<nav class="navbar" role="navigation" aria-label="main navigation">`