252 lines
10 KiB
JavaScript
252 lines
10 KiB
JavaScript
|
import { Worker, isMainThread, parentPort, workerData } from 'worker_threads';
|
||
|
|
import { WASI } from 'wasi';
|
||
|
|
import { fileURLToPath } from 'url';
|
||
|
|
import path from 'path';
|
||
|
|
import fs from 'fs/promises'; // For async file reading
|
||
|
|
import os from 'os'; // For optimal NUM_WORKERS
|
||
|
|
import { performance } from 'perf_hooks'; // For benchmarking
|
||
|
|
|
||
|
|
// Helper to get __dirname equivalent in ES Modules for this file
|
||
|
|
const __filename = fileURLToPath(import.meta.url);
|
||
|
|
const __dirname = path.dirname(__filename);
|
||
|
|
|
||
|
|
// --- Configuration Constants (shared by main and workers where relevant) ---
|
||
|
|
const ARRAY_SIZE = 15000576; // Must match C code
|
||
|
|
const NUM_WORKERS = os.cpus().length; // Optimal number of workers based on CPU cores
|
||
|
|
|
||
|
|
// WASM Memory configuration (consistent with Emscripten compile flags)
|
||
|
|
const INITIAL_MEMORY_BYTES = 180224000;
|
||
|
|
const MAXIMUM_MEMORY_BYTES = 268435456;
|
||
|
|
const PAGE_SIZE = 65536; // 64KB
|
||
|
|
|
||
|
|
// --- WASI Import Version (based on previous successful tests) ---
|
||
|
|
const WASI_VERSION = 'preview1'; // Or 'wasi_snapshot_preview1', 'wasi_unstable' based on your Node.js's specific requirement.
|
||
|
|
|
||
|
|
// --- Main Thread Logic ---
|
||
|
|
if (isMainThread) {
|
||
|
|
|
||
|
|
async function main() {
|
||
|
|
const wasmPath = './binaries/wasm_add.wasm'; // Path to your compiled WASM binary
|
||
|
|
|
||
|
|
try {
|
||
|
|
// Read the WebAssembly binary file once for all threads
|
||
|
|
const wasmBuffer = await fs.readFile(wasmPath);
|
||
|
|
|
||
|
|
// Call the setupWebAssembly function to initialize and run the benchmarks
|
||
|
|
await setupWebAssembly(wasmBuffer);
|
||
|
|
|
||
|
|
} catch (error) {
|
||
|
|
console.error("Error during WebAssembly execution:", error);
|
||
|
|
process.exit(1); // Exit with an error code on failure
|
||
|
|
}
|
||
|
|
}
|
||
|
|
|
||
|
|
// Call the main function to start the application
|
||
|
|
main();
|
||
|
|
|
||
|
|
} else { // --- Worker Thread Logic ---
|
||
|
|
|
||
|
|
// Data passed from the main thread when the worker was created
|
||
|
|
const { wasmBinary, wasmMemory, functionName, args } = workerData;
|
||
|
|
|
||
|
|
// Setup WASI imports for this worker thread's WASM instance
|
||
|
|
const wasi = new WASI({
|
||
|
|
version: WASI_VERSION,
|
||
|
|
args: process.argv,
|
||
|
|
env: process.env,
|
||
|
|
preopens: { '.': '.' }
|
||
|
|
});
|
||
|
|
|
||
|
|
// Minimal import object for standalone WASM in a worker
|
||
|
|
const importObject = {
|
||
|
|
wasi_snapshot_preview1: wasi.wasiImports, // Provide WASI imports
|
||
|
|
env: {
|
||
|
|
memory: wasmMemory, // The shared WebAssembly.Memory object
|
||
|
|
// Common stubs for functions WASM might try to call (if not handled by WASI)
|
||
|
|
_abort: () => { console.error("WASM called abort from worker!"); process.exit(1); },
|
||
|
|
_sbrk: (increment) => {
|
||
|
|
console.warn(`Worker: _sbrk called with increment ${increment}. Not fully implemented for custom runtime.`);
|
||
|
|
return 0; // Return dummy value
|
||
|
|
},
|
||
|
|
__setThrew: (ptr, val) => {},
|
||
|
|
__resumeException: (ptr) => {},
|
||
|
|
},
|
||
|
|
};
|
||
|
|
|
||
|
|
async function initAndRunWorker() {
|
||
|
|
let instance;
|
||
|
|
try {
|
||
|
|
// Instantiate the WASM module within this worker thread
|
||
|
|
instance = (await WebAssembly.instantiate(wasmBinary, importObject)).instance;
|
||
|
|
} catch (e) {
|
||
|
|
console.error("Worker: Failed to instantiate WASM module:", e);
|
||
|
|
parentPort.postMessage({ type: 'error', message: `WASM instantiation failed: ${e.message}` });
|
||
|
|
process.exit(1);
|
||
|
|
}
|
||
|
|
|
||
|
|
const exports = instance.exports;
|
||
|
|
|
||
|
|
// Call the specified WASM function with its arguments
|
||
|
|
if (typeof exports[functionName] === 'function') {
|
||
|
|
try {
|
||
|
|
exports[functionName](...args);
|
||
|
|
parentPort.postMessage({ type: 'done' }); // Signal successful completion
|
||
|
|
} catch (e) {
|
||
|
|
console.error(`Worker: Error executing ${functionName}:`, e);
|
||
|
|
parentPort.postMessage({ type: 'error', message: `Execution failed: ${e.message}` });
|
||
|
|
}
|
||
|
|
} else {
|
||
|
|
console.error(`Worker: Function '${functionName}' not found in WASM exports.`);
|
||
|
|
parentPort.postMessage({ type: 'error', message: `Function '${functionName}' not found.` });
|
||
|
|
}
|
||
|
|
|
||
|
|
process.exit(0); // Exit the worker process cleanly
|
||
|
|
}
|
||
|
|
|
||
|
|
// Start execution in the worker thread
|
||
|
|
initAndRunWorker();
|
||
|
|
}
|
||
|
|
|
||
|
|
// --- The setupWebAssembly Function (available to both, but called only by main) ---
|
||
|
|
// This function orchestrates the WASM module loading and parallel execution.
|
||
|
|
export async function setupWebAssembly(wasmBuffer) {
|
||
|
|
|
||
|
|
// Create the shared WebAssembly.Memory object
|
||
|
|
const memory = new WebAssembly.Memory({
|
||
|
|
initial: INITIAL_MEMORY_BYTES / PAGE_SIZE,
|
||
|
|
maximum: MAXIMUM_MEMORY_BYTES / PAGE_SIZE,
|
||
|
|
shared: true
|
||
|
|
});
|
||
|
|
|
||
|
|
// Setup WASI for the main thread's WASM instance
|
||
|
|
const wasi = new WASI({
|
||
|
|
version: WASI_VERSION,
|
||
|
|
args: process.argv,
|
||
|
|
env: process.env,
|
||
|
|
preopens: { '.': '.' }
|
||
|
|
});
|
||
|
|
|
||
|
|
// Instantiate the WASM module on the main thread
|
||
|
|
const { instance } = await WebAssembly.instantiate(wasmBuffer, {
|
||
|
|
wasi_snapshot_preview1: wasi.wasiImports, // Provide WASI imports
|
||
|
|
env: {
|
||
|
|
memory: memory, // Pass the shared memory
|
||
|
|
emscripten_notify_memory_growth: (index) => {
|
||
|
|
console.log(`Memory grew to ${index / 65536} pages`);
|
||
|
|
},
|
||
|
|
_abort: () => { console.error("WASM called abort from main!"); process.exit(1); },
|
||
|
|
_sbrk: (increment) => {
|
||
|
|
console.warn(`Main: _sbrk called with increment ${increment}. Not fully implemented.`);
|
||
|
|
return 0;
|
||
|
|
},
|
||
|
|
__setThrew: (ptr, val) => {},
|
||
|
|
__resumeException: (ptr) => {},
|
||
|
|
}
|
||
|
|
});
|
||
|
|
|
||
|
|
const exports = instance.exports; // Original WASM exports from the main instance
|
||
|
|
|
||
|
|
const ARRAY_SIZE = 15000576; // Keep consistent with C code
|
||
|
|
|
||
|
|
// Allocate arrays using the main thread's WASM instance (malloc/free are C exports)
|
||
|
|
const a = exports._alloc_float_array(ARRAY_SIZE);
|
||
|
|
const b = exports._alloc_float_array(ARRAY_SIZE);
|
||
|
|
const c = exports._alloc_float_array(ARRAY_SIZE);
|
||
|
|
if (!a || !b || !c) throw new Error("Failed to allocate memory");
|
||
|
|
|
||
|
|
console.log(`Allocated arrays in shared memory at: a=${a}, b=${b}, c=${c}`);
|
||
|
|
|
||
|
|
// --- Worker Orchestration Helper (Internal to setupWebAssembly) ---
|
||
|
|
// This helper function handles spawning workers and waiting for their completion.
|
||
|
|
async function runRangeTaskInWorkers(wasmFunctionName, arrayPointers) {
|
||
|
|
const workers = [];
|
||
|
|
const promises = [];
|
||
|
|
const chunkSize = Math.ceil(ARRAY_SIZE / NUM_WORKERS);
|
||
|
|
|
||
|
|
for (let i = 0; i < NUM_WORKERS; i++) {
|
||
|
|
const startIdx = i * chunkSize;
|
||
|
|
const endIdx = Math.min(startIdx + chunkSize, ARRAY_SIZE);
|
||
|
|
|
||
|
|
if (startIdx >= ARRAY_SIZE) break; // Stop if no more work chunks
|
||
|
|
|
||
|
|
const worker = new Worker(__filename, { // Crucially, the worker executes *this same file*
|
||
|
|
workerData: {
|
||
|
|
wasmBinary: wasmBuffer, // Pass the WASM binary data to workers
|
||
|
|
wasmMemory: memory, // Pass the shared memory object
|
||
|
|
functionName: wasmFunctionName, // The C function to call in the worker
|
||
|
|
args: [...arrayPointers, startIdx, endIdx] // Arguments for the C function
|
||
|
|
}
|
||
|
|
});
|
||
|
|
|
||
|
|
const workerPromise = new Promise((resolve, reject) => {
|
||
|
|
worker.on('message', (msg) => {
|
||
|
|
if (msg.type === 'done') {
|
||
|
|
resolve();
|
||
|
|
} else if (msg.type === 'error') {
|
||
|
|
reject(new Error(`Worker error for ${wasmFunctionName}: ${msg.message}`));
|
||
|
|
}
|
||
|
|
});
|
||
|
|
worker.on('error', reject); // Catch errors emitted by the worker process
|
||
|
|
worker.on('exit', (code) => {
|
||
|
|
if (code !== 0 && code !== undefined) {
|
||
|
|
reject(new Error(`Worker exited with code ${code} for ${wasmFunctionName}`));
|
||
|
|
} else {
|
||
|
|
resolve(); // Resolve if exited cleanly (code 0 or undefined)
|
||
|
|
}
|
||
|
|
});
|
||
|
|
});
|
||
|
|
workers.push(worker);
|
||
|
|
promises.push(workerPromise);
|
||
|
|
}
|
||
|
|
|
||
|
|
// Wait for all workers to complete their assigned tasks
|
||
|
|
await Promise.all(promises);
|
||
|
|
}
|
||
|
|
|
||
|
|
// --- Override/Wrap original C function calls with worker orchestration ---
|
||
|
|
// This makes the subsequent benchmark calls look exactly like your original script.
|
||
|
|
const wrappedExports = {
|
||
|
|
init_arrays: async (aPtr, bPtr) => {
|
||
|
|
await runRangeTaskInWorkers('_init_arrays_range', [aPtr, bPtr]);
|
||
|
|
},
|
||
|
|
calculate: async (aPtr, bPtr, cPtr) => {
|
||
|
|
await runRangeTaskInWorkers('_calculate_range', [aPtr, bPtr, cPtr]);
|
||
|
|
},
|
||
|
|
// Keep original alloc/dealloc directly from WASM as they are not parallelized
|
||
|
|
alloc_float_array: exports._alloc_float_array,
|
||
|
|
dealloc_float_array: exports._dealloc_float_array,
|
||
|
|
alloc: exports._alloc,
|
||
|
|
dealloc: exports._dealloc,
|
||
|
|
// Expose memory for direct access like in your original code
|
||
|
|
memory: exports.memory, // This points to the same shared memory object
|
||
|
|
};
|
||
|
|
|
||
|
|
|
||
|
|
// --- Benchmark Execution (This part now perfectly matches your desired input!) ---
|
||
|
|
|
||
|
|
// Benchmark initialization
|
||
|
|
const initStart = performance.now();
|
||
|
|
await wrappedExports.init_arrays(a, b); // Calls the JS wrapper
|
||
|
|
const initTime = performance.now() - initStart;
|
||
|
|
|
||
|
|
// Benchmark calculation
|
||
|
|
const calcStart = performance.now();
|
||
|
|
await wrappedExports.calculate(a, b, c); // Calls the JS wrapper
|
||
|
|
const calcTime = performance.now() - calcStart;
|
||
|
|
|
||
|
|
console.log(`Array initialization time: ${initTime.toFixed(3)} ms`);
|
||
|
|
console.log(`Calculation time: ${calcTime.toFixed(3)} ms`);
|
||
|
|
console.log(`Total time: ${(initTime + calcTime).toFixed(3)} ms`);
|
||
|
|
|
||
|
|
// Sample output
|
||
|
|
// Access memory via the original instance's exports.memory (or the shared 'memory' var)
|
||
|
|
const mem = new Float32Array(wrappedExports.memory.buffer);
|
||
|
|
for (let i = 0; i < 10; i++) {
|
||
|
|
console.log(`c[${i}] = ${mem[c/4 + i]}`);
|
||
|
|
}
|
||
|
|
|
||
|
|
// Clean up
|
||
|
|
wrappedExports.dealloc_float_array(a);
|
||
|
|
wrappedExports.dealloc_float_array(b);
|
||
|
|
wrappedExports.dealloc_float_array(c);
|
||
|
|
}
|