Fast Matrix Math 2

Last time we looked a number of ways to add two matrices starting from naive and applying simplifications until we basically just formed a flat loop over the data which turned out to be the fastest, beating out even TypedArrays. In some ways I'm regretting choosing such a simple case of element-wise addition because it did just simplify into testing the speed of flat loops and that maybe matrix multiplication would have been better. Oh well, maybe next time.

There were some shortcomings because of the sheer amount of permutations I couldn't get to everything so we'll start off really looking at byte representations as a source of speed. In addition I've removed the 100x100 for a 128x128 sized tests and added a final 256x256 to really challenge differences with "big" matrices.

This actually changes some strategies as things like the simple flat loop actually get a lot worse than allocated loops. Typed arrays really pull ahead. My guess is because the compiler doesn't trust allocating large JS arrays.

Simple Typed Array Loops

I neglected to test typed arrays with flat loops since we know we don't need the structured row-column loops we have before. There's not much to these:

export function addMatrixFloat64Flat(a, b) {
    const out = {
        shape: a.shape,
        data: new Float64Array(a.data.length)
    };
    for (let i = 0; i < a.data.length; i++) {
        out.data[i] = a.data[i] + b.data[i];
    }
    return out;
}

As expected they are a tad faster except strangely for I32 which is slightly faster with the inner loop bound checks, weird.

Fixing Tests for Float32 Arrays

One issue that came up that I simply ignored last time was that I could write tests for Float32Arrays because the precision differences between them and 64-bit floats. I found that there's actually a function in the Javascript Math library called Math.fround that can do floating point precision rounding. While generating data I can fround the arguments and fround the result to keep it accurate for 32-bit floats.

WASM

Another way we can look to speed things up is WASM. WASM is often discribed as a math co-processor for the web. At least, people think that if you move your code to WASM you'll see a big speed up. This however is much more complicated that we might think. The overhead of calling WASM is high because memory has to be moved between the host and WASM module, this means that we might not see big gains, but hopefully the low internal overhead might help speed things up.

I built a module in pure WAT, the WebAssembly Text Format. While normally we don't do this and instead use a language like Rust or C to compile to WASM there's a lot of complexity to the compilation and weird stuff might make it into the module and change our results and it's hard to optimize. To keep it as close to the metal as possible WAT is the best option. Thankfully what we are doing is not complicated and can be reasonably written in WAT.

;;mat.wat
(module
    (func $matrix_add_f64 (param $lhs_ptr i32) (param $rhs_ptr i32) (param $length i32) (param $out_ptr i32)
        (local $i i32)
        (local $offset i32)
        (local.set $i (i32.const 0))
        (block
            (loop               
                (local.set $offset (i32.mul (local.get $i) (i32.const 8)))

                (i32.add (local.get $out_ptr) (local.get $offset)) ;;out addr on stack
                (f64.add ;;result on stack
                    (f64.load (i32.add (local.get $lhs_ptr) (local.get $offset))) 
                    (f64.load (i32.add (local.get $rhs_ptr) (local.get $offset))))

                (f64.store)

                (local.set $i (i32.add (local.get $i) (i32.const 1)))
                (br_if 1 (i32.eq (local.get $i) (local.get $length)))
                (br 0)
            )
        )
    )
    (export "addMatrix" (func $matrix_add))
    (memory (export "memory") 1)
)

I'm not going to explain this much as it's out of scope, but this is basically the same as the flat simple version. It treats data as a singular array and iterates over it while adding. The parameters are the pointer to the left-hand array, the right-hand array, the total length and where to store the output. We could have technically just used length and if we start at 0 know exactly where everything is but to start I wanted explicit to start.

Compiling WAT

To actually compile this we can use a tool called WABT (WebAssembly Binary Toolkit). It's basically a mess that requires CMake and I couldn't get it to run on WSL and I wasn't going to install MinGW. Instead there's a nice tool called WAPM from Wasmer which works like npm for webassembly packages and since it's been compiled down to webassembly we can run it in any environment. In fact we don't even need to add configuration so long as wapm is installed. We can run wax wat2wasm -- wat/mat.wat -o wasm/mat.wasm. wax is like npx for npm. If you're wondering the command we give wax is defined by the wasmer/wabt package: https://wapm.io/wasmer/wabt. Also for some reason you can't prefix local paths with ./ so wax wat2wasm -- ./wat/mat.wat doesn't work which tool me a while to figure out. Anyway this provides a nice simple compile environment if you want to work on raw WAT files.

This will work up until we try with a 128x128 matrix. It will give us out-of-bounds access. The reason is because the memory we allocated can't fit all the data we need! We can however grow the pages in response to larger inputs.

With page growing:

export function addMatrixWasmF64(a,b){
    const lhsElementOffset = 0;
    const rhsElementOffset = lhsElementOffset + a.data.length;
    const rhsByteOffset = rhsElementOffset * 8;
    const resultElementOffset = lhsElementOffset + a.data.length + b.data.length;
    const resultByteOffset = resultElementOffset * 8;
    const elementLength = a.data.length;

    //grow memory if needed
    const spaceLeftover = matInstance.exports.memory.buffer.byteLength - (a.data.length * 3 * 8)
    if (spaceLeftover < 0){
        const pagesNeeded = Math.ceil((a.data.length * 3 * 8) / (64 * 1024));
        const pagesHave = matInstance.exports.memory.buffer.byteLength / (64 * 1024);
        matInstance.exports.memory.grow(pagesNeeded - pagesHave);
    }

    const f64View = new Float64Array(matInstance.exports.memory.buffer);
    f64View.set(a.data, lhsElementOffset);
    f64View.set(b.data, rhsElementOffset);

    matInstance.exports.addMatrix(0, rhsByteOffset, elementLength, resultByteOffset);

    return {
        shape: a.shape,
        data: f64View.slice(resultElementOffset, resultElementOffset + elementLength)
    };
}

This might be how we'd do a production version so things don't explode but for performance reasons I'm going to leave that out and make it static, we'll only do up-to 256x256 sized matrixes so we need a total of 24 pages allocated.

(memory (export "memory") 24)

The benchmark results are tad disappointing. It expectedly starts out at the bottom of the pack, loading data into the wasm module is expensive overhead. It slowly climbs the latter but even at 128x128 is still slower than a simple loop. The lack of allocation make a big difference. This means unless more of the data can be stored on the WASM side, it's unlikely WASM will provide much of a benefit. Still, WASM has one more trick.

SIMD

Unlike Javascript WASM has access to SIMD instructions which can be used to parallelize the iteration and these are pretty well supported. SIMD stands for Single Instruction Multiple Data. The idea is that it let's us do add/load etc more than one value at a time provided everything lines up correctly. WASM at least as of the time of writing only has a 128-bit type. The way it works is that we load a 128-bit value and at the operation level we can choose how to interpret it, which in our case will correspond to 2 float-64s. When we increment $i we now do so by 2.

It's actually not too hard to modify the existing code. Instead of stepping be 8-bytes, we step by 16-bytes loading 128-bit values instead of 64. Then we use the f64x2.add to add the two floats at the same time and store it back to memory as a 128-bit value.

(func $matrix_add_simd_f64 (param $lhs_ptr i32) (param $rhs_ptr i32) (param $length i32) (param $out_ptr i32)
    (local $i i32)
    (local $offset i32)
    (local.set $i (i32.const 0))
    (block
        (loop               
            (local.set $offset (i32.mul (local.get $i) (i32.const 8))) ;;moving in 128-bit blocks
            (i32.add (local.get $out_ptr) (local.get $offset)) ;;out addr on stack
            (f64x2.add 
                (v128.load (i32.add (local.get $lhs_ptr) (local.get $offset))) 
                (v128.load (i32.add (local.get $rhs_ptr) (local.get $offset)))) ;;result on stack
            (v128.store)
            (local.set $i (i32.add (local.get $i) (i32.const 2)))
            (br_if 1 (i32.ge_u (local.get $i) (local.get $length)))
            (br 0)
        )
    )
)

Another thing to think about is what happens if we don't have exactly 2 floats to add? We need one modification which is to check the branch condition to be greater-than-or-equal (i32.ge) rather than equal (i32.eq) to account for possible overrun by one. But let's also think about the memory bounds. If we have just one f64 then the second f64 will load garbage from outside our bounds but we never write (v128.store) in such a way it'll bump into the next segment (assuming the pointers are not overlapping), only the final array at $out_ptr is ever written to. Since we only read back the exact number of elements we want the odd value is just ignored. We also can't read outside the bounds of allocated memory because they are always given by even pages.

I have to do a little debugging to get this right. A simple way is to import little functions for printing values:

const { instance: matInstance } = await WebAssembly.instantiate(wasm, {
    lib: {
        print_int: function(num){
            console.log(num);
        },
        print_float: function(num){
            console.log(num);
        },
        print_brk: function(){
            console.log("---")
        }
    }
});

(import "lib" "print_int" (func $print_int (param i32)))
(import "lib" "print_float" (func $print_float (param f32)))
(import "lib" "print_brk" (func $print_brk))

;;---stuff

(call $print_int (local.get $offset)) ;;print offset
(call $print_int (local.get $i)) ;;print $i
(call $print_brk) ;;at a visual barrier between element logs

We claw our way up the latter a bit further. We're still not beating a preallocated loop at 256x256. So still not very usable.

SIMD F32 and I32

We can still mess with the precision. 32-bit values would give us an advantage of being able to add 4 at a time. We also know that integers generally work a bit faster.

(func $matrix_add_simd_f32 (param $lhs_ptr i32) (param $rhs_ptr i32) (param $length i32) (param $out_ptr i32)
    (local $i i32)
    (local $offset i32)
    (local.set $i (i32.const 0))
    (block
        (loop               
            (local.set $offset (i32.mul (local.get $i) (i32.const 8))) ;;moving in 128-bit blocks
            (i32.add (local.get $out_ptr) (local.get $offset)) ;;out addr on stack
            (f32x4.add 
                (v128.load (i32.add (local.get $lhs_ptr) (local.get $offset))) 
                (v128.load (i32.add (local.get $rhs_ptr) (local.get $offset)))) ;;result on stack
            (v128.store)
            (local.set $i (i32.add (local.get $i) (i32.const 2)))
            (br_if 1 (i32.ge_u (local.get $i) (local.get $length)))
            (br 0)
        )
    )
)
(func $matrix_add_simd_i32 (param $lhs_ptr i32) (param $rhs_ptr i32) (param $length i32) (param $out_ptr i32)
    (local $i i32)
    (local $offset i32)
    (local.set $i (i32.const 0))
    (block
        (loop               
            (local.set $offset (i32.mul (local.get $i) (i32.const 8))) ;;moving in 128-bit blocks
            (i32.add (local.get $out_ptr) (local.get $offset)) ;;out addr on stack
            (i32x4.add 
                (v128.load (i32.add (local.get $lhs_ptr) (local.get $offset))) 
                (v128.load (i32.add (local.get $rhs_ptr) (local.get $offset)))) ;;result on stack
            (v128.store)
            (local.set $i (i32.add (local.get $i) (i32.const 2)))
            (br_if 1 (i32.ge_u (local.get $i) (local.get $length)))
            (br 0)
        )
    )
)

export function addMatrixWasmSimdF32(a, b) {
    const lhsElementOffset = 0;
    const rhsElementOffset = lhsElementOffset + a.data.length;
    const rhsByteOffset = rhsElementOffset * 4;
    const resultElementOffset = lhsElementOffset + a.data.length + b.data.length;
    const resultByteOffset = resultElementOffset * 4;
    const elementLength = a.data.length;

    //grow memory if needed
    // const spaceLeftover = matInstance.exports.memory.buffer.byteLength - (a.data.length * 3 * 8)
    // if (spaceLeftover < 0){
    //  const pagesNeeded = Math.ceil((a.data.length * 3 * 8) / (64 * 1024));
    //  const pagesHave = matInstance.exports.memory.buffer.byteLength / (64 * 1024);
    //  matInstance.exports.memory.grow(pagesNeeded - pagesHave);
    // }

    const f32View = new Float32Array(matInstance.exports.memory.buffer);
    f32View.set(a.data, lhsElementOffset);
    f32View.set(b.data, rhsElementOffset);

    matInstance.exports.addMatrixSimdF32(0, rhsByteOffset, elementLength, resultByteOffset);

    return {
        shape: a.shape,
        data: f32View.slice(resultElementOffset, resultElementOffset + elementLength)
    };
}

export function addMatrixWasmSimdI32(a, b) {
    const lhsElementOffset = 0;
    const rhsElementOffset = lhsElementOffset + a.data.length;
    const rhsByteOffset = rhsElementOffset * 4;
    const resultElementOffset = lhsElementOffset + a.data.length + b.data.length;
    const resultByteOffset = resultElementOffset * 4;
    const elementLength = a.data.length;

    //grow memory if needed
    // const spaceLeftover = matInstance.exports.memory.buffer.byteLength - (a.data.length * 3 * 8)
    // if (spaceLeftover < 0){
    //  const pagesNeeded = Math.ceil((a.data.length * 3 * 8) / (64 * 1024));
    //  const pagesHave = matInstance.exports.memory.buffer.byteLength / (64 * 1024);
    //  matInstance.exports.memory.grow(pagesNeeded - pagesHave);
    // }

    const i32View = new Int32Array(matInstance.exports.memory.buffer);
    i32View.set(a.data, lhsElementOffset);
    i32View.set(b.data, rhsElementOffset);

    matInstance.exports.addMatrixSimdI32(0, rhsByteOffset, elementLength, resultByteOffset);

    return {
        shape: a.shape,
        data: i32View.slice(resultElementOffset, resultElementOffset + elementLength)
    };
}

These are just basic substitutions of the original code but with 4-byte sized elements. The WASM still jumps 128-bit values it's just treating them as 4 values instead of 2. For arrays larger than 32x32 I32 and F32 are nearly equal and out new speed kings! Still, it's not really a huge difference versus an I32Array in plain javascript (~+10%) so it's a far bit of complexity.

Conclusion

What we've observed this time is that WASM by itself does not yield great performance gains until it actually overcomes the overhead of memory copy, which in our case is around the 64x64 element mark. In these cases using typed arrays can still beat it by not introducing that complexity, the compiler seems smart enough to optimize these. However, once we start introducing SIMD to the picture WASM can overtake plain JS starting around 64x64 and pulling away as sizes get bigger.

Code

https://github.com/ndesmic/fast-mat/tree/v1.1

Data