Pure WASM loop 2.5x slower than calling WASM from within JS loop

[mhart]

Describe the Bug

I suspect I might just be missing a compiler optimization here.

Essentially, I'm seeing the same function calls being slower from within WASM than I am from JS to WASM. This is especially pronounced if I have a for loop, or nested for loops. If I create the for loops in JS and call into WASM, it's 2.5x faster than if I do the entire thing in WASM (for the exact same function), which is very counter-intuitive (as it's not needing to cross the JS-WASM boundary).

Perhaps there's some sort of loop or memory checking being done that I can turn off?

Steps to Reproduce

use std::slice;
use wasm_bindgen::prelude::wasm_bindgen;

#[wasm_bindgen]
pub fn alloc(len: usize) -> *mut u8 {
    let mut buf = Vec::with_capacity(len);
    let ptr = buf.as_mut_ptr();
    std::mem::forget(buf);
    ptr
}

#[wasm_bindgen]
pub unsafe fn l2(p_x: *mut f32, p_y: *mut f32, size: usize) -> f32 {
    slice::from_raw_parts(p_x, size)
        .iter()
        .zip(slice::from_raw_parts(p_y, size))
        .fold(0., |sum, (x, y)| {
            let diff = x - y;
            sum + diff * diff
        })
}

#[wasm_bindgen]
pub unsafe fn l2_loop(ptr: *mut f32, num_vectors: usize, size: usize) -> f32 {
    let mut distance: f32 = 0.;
    for i in 0..num_vectors {
        for j in 0..num_vectors {
            distance += l2(
                ptr.offset((i * size) as isize),
                ptr.offset((j * size) as isize),
                size,
            );
        }
    }
    distance
}

Compiling with:

wasm-pack build -t web

Called like so:

import { readFileSync } from "node:fs";
import { initSync, alloc, l2, l2_loop } from "./pkg/hello_world.js";

const { memory } = initSync(
  await WebAssembly.compile(
    readFileSync(new URL("pkg/hello_world_bg.wasm", import.meta.url))
  )
);
memory.grow((128 * 1024 * 1024) / (64 * 1024));

const dim = 1536;
const numVectors = 1600;

const ptr = alloc(dim * numVectors * 4);

const arr = new Float32Array(memory.buffer, ptr, dim * numVectors);
arr.set(Array.from({ length: arr.length }, () => Math.random()));

let start, end, distance;

// JS Loop
// Calls the `l2` wasm function. Thought it would be slower due to many JS->WASM boundary calls

start = performance.now();
distance = 0;
for (let i = 0; i < numVectors; i++) {
  for (let j = 0; j < numVectors; j++) {
    distance += l2(ptr + i * dim * 4, ptr + j * dim * 4, dim);
  }
}
end = performance.now();
console.log("l2 from JS took:", end - start);
console.log("distance:", distance);

// WASM Loop
// Calls the same `l2` function in a loop from within WASM. Should be faster, or at least not noticeably slower

start = performance.now();
distance = l2_loop(ptr, numVectors, dim);
end = performance.now();
console.log("l2_loop took:", end - start);
console.log("distance:", distance);

(Have tried all the wasm-opt options, including turning it off, none make a noticeable difference. Also tried nightly, with no difference)

Expected Behavior

Looping from within WASM should take roughly the same time as from JS, or even faster.

Actual Behavior

Observed on a Mac M1 Pro, Node.js v20.6.1

l2 from JS took: 3651.1732909679413
distance: 654785779.2068481

l2_loop took: 8687.912124991417
distance: 654949120

ie, l2_loop, which is doing the work entirely in WASM, and calling the same function, is 2.5x slower.

[daxpedda]

This is some very interesting code you got here.

Just to make sure, did you build in release mode?
wasm-pack build uses the dev profile by default, see the documentation.

[daxpedda]

I reproduced the same results on Chrome as well. Whats interesting is that Firefox with or without profiling was producing the same improved timings as Chrome during profiling, so I'm assuming that this is some kind of bug in V8s Wasm engine.

Firefox with or without profiling:

l2 from JS took: 1442 [localhost:8000:37:17](http://localhost:8000/)
distance: 655488059.1537781 [localhost:8000:38:17](http://localhost:8000/)
l2_loop took: 1372 [localhost:8000:46:17](http://localhost:8000/)
distance: 655488059.1537781

Chrome without profiling:

l2 from JS took: 3578.60000000149
distance: 654943093.3420715
l2_loop took: 3585
distance: 654943093.3420715

Chrome with profiling:

l2 from JS took: 1476.7999999970198
distance: 655310011.2146912
l2_loop took: 3999.5
distance: 655310011.2146912

I did adjust the l2_loop() function to use f64 for calculations to more accurately match the JS results, but it changed nothing.

---

Your best bet is probably to build a really minimal Wasm example, without Rust or a shim at all, and make a bug report to Chromium.

[Liamolucko]

I'm also able to reproduce this in Chrome. However, if I close the devtools, I get the same results as running while profiling.

Weirdly, if I reload the page and run the benchmarks again (still with the devtools closed), l2_loop gets faster and matches the speed of l2 being called from JS. If I open and close the devtools and reload again, l2_loop then goes back to being slower than l2 from JS; and I reload again, it goes back to being the same speed.

Where things get really weird is if I open/close the devtools a second time. It then requires reloading the page twice to get l2_loop to return to the speed of l2, rather than once. If I keep on opening/closing the devtools, the number of reloads required keeps on increasing. I got these results, although they varied a bit between my Mac laptop and Linux PC:

Number of times devtools opened	Number of reloads required (Mac)	Number of reloads required (Linux)
0	1	1
1	1	1
2	2	2
3	4	4
4	7	8
5	15	16
6	30	31
7	59	62
8	117	124
9	234	247

Opening the page in a new tab or restarting the browser resets everything.

Also, for some reason after the 8th open/close of devtools, after the first reload things run at devtools-open speed, before returning to the Node.js speeds for the next 117/124 reloads. Then after the 9th open/close, it takes two reloads to get out of the devtools-open speed... seems like a similar pattern.

I also had a weird issue where, for some reason, if I manually changed the URL instead of pressing reload, things would reset... but only on Linux for some reason, not on my Mac?

This is the HTML file I used to automatically count how many reloads were required, as well as write to the page instead of using the console (Rust code was unchanged):

<!DOCTYPE html>
<pre id="output">
</pre>
<script type="module">
    let priorReloads = parseInt(new URL(location.href).searchParams.get("reloads") ?? "0");
    const output = document.getElementById("output");

    import init, {alloc, l2, l2_loop} from "./pkg/hello_world.js";

    const {memory} = await init();
    memory.grow((128 * 1024 * 1024) / (64 * 1024));

    const dim = 1536;
    const numVectors = 200;

    const ptr = alloc(dim * numVectors * 4);

    const arr = new Float32Array(memory.buffer, ptr, dim * numVectors);
    arr.set(Array.from({length: arr.length}, () => Math.random()));

    let start, end, distance;

    output.textContent += "Starting benchmark...\n";
    output.textContent += `${priorReloads} prior reloads\n`;
    await new Promise(resolve => setTimeout(resolve, 0));

    // JS Loop
    // Calls the `l2` wasm function. Thought it would be slower due to many JS->WASM boundary calls

    start = performance.now();
    distance = 0;
    for (let i = 0; i < numVectors; i++) {
        for (let j = 0; j < numVectors; j++) {
            distance += l2(ptr + i * dim * 4, ptr + j * dim * 4, dim);
        }
    }
    end = performance.now();
    const l2Time = end - start;
    output.textContent += `l2 from JS took: ${l2Time}\n`;
    output.textContent += `distance: ${distance}\n`;
    await new Promise(resolve => setTimeout(resolve, 0));

    // WASM Loop
    // Calls the same `l2` function in a loop from within WASM. Should be faster, or at least not noticeably slower

    start = performance.now();
    distance = l2_loop(ptr, numVectors, dim);
    end = performance.now();
    const l2LoopTime = end - start;
    output.textContent += `l2_loop took: ${l2LoopTime}\n`;
    output.textContent += `distance: ${distance}\n`;

    if (l2LoopTime > 1.2 * l2Time) {
        location.href = `?reloads=${priorReloads + 1}`;
    } else {
        history.replaceState(null, "", location.href.slice(0, location.href.length - location.search.length));
    }
</script>

I also reduced numVectors so it didn't take as long.

[mhart]

Oh that's interesting – so you're seeing the "fast" results, so long as dev tools isn't open (and so long as it's not running a perf profile).

So that suggests that it is the norm that JS calling into WASM is faster – for whatever reason

[Liamolucko]

So that suggests that it is the norm that JS calling into WASM is faster – for whatever reason

No, I mustn't have been clear enough - if devtools is closed and I reload, WASM speeds up to the same speed as JS (in fact, a little bit faster).

If I increase numVectors back to 1600, I get these results (also on an M1 Pro):

Devtools open:
l2 from JS took: 8572.30000000028
distance: 655308372.9647827
l2_loop took: 8540.599999999627
distance: 655347520

Devtools closed, 1 reload:
l2 from JS took: 3633.2000000001863
distance: 654743623.7380676
l2_loop took: 8555.100000000093
distance: 654882496

Devtools closed, 2 reloads:
l2 from JS took: 3608.3000000002794
distance: 654895268.1144104
l2_loop took: 3533.600000000093
distance: 654961792

[mhart]

Oh wow, that's even weirder than I thought lol. Quite the heisenbug