I like writing this simple blog entries to not forget about stuff I have learned, this time I want to explain how to share memory between Javascript and Webassembly.

This is not something that I came up by my self, but read about it somewhere on the internet, but I have lost the link to it, I'll update this entry if I can find it again. In any case, I was looking for a way to transfer large amounts od data between Javascript and a Webassembly module, with as little overhead as possible.

Usually this is taken care of by libraries like wasm-bindgen, but wanted the extra flexibility and the knowhow. And this is what I settled on.

From a high level view, the steps are as follow:

1. Have a function inside the webassembly module that allocates a block of memory and returns it address
2. Call that function from Javascript and save the address
3. From Javascript use that address to create a typed array.
4. Update the array from either Javscript or the Webassembly module
5. Release the memory when done

I like using Rust to write webassembly modules, so the examples will be done using that.

Create a webassembly library

$ cargo new memory-sharing --lib

The library will be compiled in two different modes cdylib which is a dynamic library and rlib which is a Rust library file (more information here), as such, change Cargo.toml to look like this:

[lib]
crate-type = ["cdylib", "rlib"]

Implement a function that allocates a memory block

This function has to do three things:

1. Allocate the memory, in this case using Vec::with_capacity.
2. Get a pointer to the newly allocated memory, the resuling pointer from calling .as_mut_ptr() will address the vector buffer directly.
3. Forget about the memory, so it will not be freed after the function finishes.
4. Return the pointer.

This is the function code:

#[no_mangle]
pub unsafe fn wasm_alloc_bytes(size_in_bytes: usize) -> *mut u8 {
    let mut memory = Vec::with_capacity(size_in_bytes);
    let ptr = memory.as_mut_ptr();
    std::mem::forget(memory);
    ptr
}

Since this is basically manual memory management, freeing memory is achived with:

#[no_mangle]
pub unsafe fn free_f32_vec(n: usize, ptr: *mut f32) {
    let _bytes: Vec<f32> = Vec::from_raw_parts(ptr, n, n);
}

Get the data back and do something with it

As an example, this function takes threww vectors, two which will be multiplied and one to store the results. Which is an example I took from here, this will make sense in a moment.

#[no_mangle]
pub unsafe fn multiply_f32_v(n: usize, output: *mut f32, input_a: *mut f32, input_b: *mut f32) {
    let mut out: Vec<f32> = Vec::from_raw_parts(output, n, n);
    let a: Vec<f32> = Vec::from_raw_parts(input_a, n, n);
    let b: Vec<f32> = Vec::from_raw_parts(input_b, n, n);
    a.iter()
        .zip(b.iter())
        .zip(out.iter_mut())
        .for_each(|((a, b), out)| {
            *out = a * b;
        });
    std::mem::forget(out);
    std::mem::forget(a);
    std::mem::forget(b);
}

Build it

$ cargo build --target wasm32-unknown-unknown --release

Call the function from Javascript

WebAssembly.instantiateStreaming(fetch("./wasm/memory_sharing.wasm"), {}).then(
  (wasmBinary) => {
    const {instance: { exports: { memory, free_f32_vec, multiply_f32_v, wasm_alloc_f32 }}} = wasmBinary;
    const totalElements = 16;
    const sourceArray1Ptr = wasm_alloc_f32(totalElements);
    const sourceArray2Ptr = wasm_alloc_f32(totalElements);
    const resultsArrayPtr = wasm_alloc_f32(totalElements);

    const sourceArray1 = new Float32Array(memory.buffer, sourceArray1Ptr, totalElements);
    for(let i = 0; i < sourceArray1.length; i++) {
      sourceArray1[i] = i + 1;
    }

    const sourceArray2 = new Float32Array(memory.buffer, sourceArray2Ptr, totalElements);
    for(let i = 0; i < sourceArray2.length; i++) {
      sourceArray2[i] = i + 1;
    }

    multiply_f32_v(totalElements, resultsArrayPtr, sourceArray1Ptr, sourceArray2Ptr);

    const resultsArray = new Float32Array(memory.buffer, resultsArrayPtr, totalElements);

    console.log(resultsArray);
    free_f32_vec(totalElements, sourceArray1Ptr);
    free_f32_vec(totalElements, sourceArray2Ptr);
    free_f32_vec(totalElements, resultsArrayPtr);
  }
);

Here is a quick bulletpoint list of what is going on there:

• Get and instantiate the webassembly module.
• Allocate three blocks of memory, using the wasm_alloc_f32 function.
• Using the returned memory pointers, create javascript typed arrays.
• Update the data inside those arrays.
• Call the multiply_f32_v function.
• Reap the results!
• Free the memory when done and live to tell the story.

Now for the fun tidbit

The array multiplication implemented in Rust has the property of being very easy to vectorize, and given that webassembly supports SIMD instructions, by setting the target-feature flag, the compiler can be nudged to use them while building the wasm module:

RUSTFLAGS="-Ctarget-feature=+simd128" cargo build --target wasm32-unknown-unknown --release

The ouput of this is a SIMD enabled implementation of the multiplication algorithm, this can be verified by looking at the text version of the wasm module, (excerpt):

(func $multiply_f32_v (type 4) (param i32 i32 i32 i32)
    ;; ...
          loop  ;; label = @4
            local.get 5
            local.get 10
            v128.load align=4
            local.get 6
            v128.load align=4
            f32x4.mul
            v128.store align=4
    ;; ...
    end)

Compare with the non-SIMD version in the equivalent section of the multiply_f32_v function, (another excerpt):

(func $multiply_f32_v (type 4) (param i32 i32 i32 i32)
    ;; ...
        loop  ;; label = @3
          local.get 0
          local.get 8
          f32.load
          local.get 7
          f32.load
          f32.mul
          f32.store
    ;; ...
    end)

SIMD instructions in webassembly are not supported in all browsers, for instance at this point Safary will fail while parsing the module, but Firefox and Chrome can execute it just fine. To try it, change the memory_sharing.wasm file name to memory_sharing_simd.wasm, and see what happens.

The Rust project is in Github.