Fast Matrix Math in JS

There's a lot of ways to represent matrix math in JS. Some are readable and some are fast. I wanted to explore the differences a bit. How much are certain techniques actually saving me?

For this I'll just be looking at a single operation: element-wise addition to reduce the total cases, but difference operations could slightly change the overall values especially ones like matrix multiplication which require slightly more complex application rules. The states are also on my computer which is a slightly older i7 8700K using Deno which is v8 under-the-hood. A different runtime like Bun might behave very differently if there are different optimizations.

The simple functional way

I thought I'd start here as this is how I might write it for a first draft. It's very code optimized but I'd suspect the performance is actually pretty bad. Immutability is great for avoiding bugs but terrible for perf, especially when JS doesn't have intelligent copies.

//mat.js
export function addMatrixFunc(a, b) {
    return a.map((row, ri) => row.map((val, ci) => b[ri][ci] + val));
}

The matrix representation is an array of arrays. The outer array is rows, the inner arrays are columns.

Using Deno's built-in benchmarking tool, we can see how this performs on different size matrices.

import { addMatrixFunc } from "./mat.js";
import { mat100A, mat100B } from "./data/mat-data.js";

Deno.bench("Add 1x1", () => {
    addMatrixFunc([[44]], [[65]]);
});

Deno.bench("Add 2x2", () => {
    addMatrixFunc([[44]], [[65]]);
});

Deno.bench("Add 4x4", () => {
    addMatrixFunc([[44]], [[65]]);
});

/* ... */

Deno.bench("Add 100x100", () => {
    addMatrixFunc(mat100A, mat100B);
});

The mat100A and mat100B are pre-generated 100x100 matrices which are way too big to put in the test file.

Some notes here is that I don't think Deno lets you set the iterations or warmup iterations, at least anymore. I think it just looks for convergence of numbers. The actual number of runs is displayed in the JSON output and it's slightly different per test.

Here's how we do:

Loops

So the first way I think we can improve is loops. Functions have overhead so if we remove that and be a little more imperative we can be a bit faster.

export function addMatrixLoop(a, b) {
    const out = [];
    for (let row = 0; row < a.length; row++) {
        const arrayRow = [];
        for (let col = 0; col < a[0].length; col++) {
            arrayRow.push(a[row][col] + b[row][col])
        }
        out.push(arrayRow);
    }
    return out;
}

Note that I'm not going to do strict bounds checking, we're just assuming a and b are the same size as the bounds checks just add overhead.

Loops start out faster but once we hit around 32x32 the they are equal to .map and larger than that .map is faster. Very surprising!

Pre-allocating arrays

My next idea was to pre-allocate the arrays since pushing into an array can cause re-sizes and maybe these are why this is slower.

export function addMatrixLoopPreAlloc(a, b) {
    const out = new Array(a.length);
    for (let row = 0; row < a.length; row++) {
        const arrayRow = new Array(a[0].length);
        for (let col = 0; col < a[0].length; col++) {
            arrayRow[col] = a[row][col] + b[row][col];
        }
        out[row] = arrayRow;
    }
    return out;
}

That did the trick! We're about 1.5x faster than where we started!

Unrolling the loops

What if we just removed all the loops and wrote it out long-hand?

export function addMatrix4x4(a, b) {
    return [
        [a[0][0] + b[0][0], a[0][1] + b[0][1], a[0][2] + b[0][2], a[0][3] + b[0][3]],
        [a[1][0] + b[1][0], a[1][1] + b[1][1], a[1][2] + b[1][2], a[1][3] + b[1][3]],
        [a[2][0] + b[2][0], a[2][1] + b[2][1], a[2][2] + b[2][2], a[2][3] + b[2][3]],
        [a[3][0] + b[3][0], a[3][1] + b[3][1], a[3][2] + b[3][2], a[3][3] + b[3][3]]
    ];

}

This is not very flexible as you need a function for every shape of matrix you want to add. However, in some cases like 3D this isn't too bad because you have a very finite amount of things, usually only ever 4x4. In machine learning this might cause problems.

Here's a function that generates the javascript text for the unrolled loop:

export function genMatAddBody(rows, cols) {
    let funcBody = "return [\n";

    for (let r = 0; r < rows; r++) {
        funcBody += "\t\t["
        for (let c = 0; c < cols; c++) {
            funcBody += `a[${r}][${c}] + b[${r}][${c}]${c < cols - 1 ? ", " : ""}`
        }
        funcBody += `]${r < rows - 1 ? ", " : ""}\n`
    }

    funcBody += `\t];\n`
    return funcBody;
}
export function genMatAddFunc(rows, cols) {
    rows = Number(rows);
    cols = Number(cols);
    const body = genMatAddBody(rows, cols);
    return new Function("a", "b", body);
}

I was also curious if making this dynamic generation changes much:

export function genMatAddFunc(rows, cols) {
    rows = Number(rows); //prevents code injection
    cols = Number(cols);
    const body = genMatAddBody(rows, cols);
    return new Function("a", "b", body);
}

Since we are using eval we should make sure to sanitize the input.

const addMatrix1x1Dyn = genMatAddFunc(1,1);
const addMatrix2x2Dyn = genMatAddFunc(2,2);
const addMatrix4x4Dyn = genMatAddFunc(4,4);
// etc.
const addMatrix100x100Dyn = genMatAddFunc(100,100);

It's a big improvement for small values beating the pre-allocated loop by about 1.5 to 2x but not good for large ones being considerably slower. I'm not sure why that's the case, maybe it has to due with the size of the function itself? The code generated is massive. Also dynamic generation is basically the same as writing them out. So if you want to save payload (and aren't limited by CSP) you can dynamically create these at no penalty.

Flattening the arrays

Another thing where I think we can save is the arrays. We technically don't need to have a lot of nested arrays and they'll add some overhead to create. So now a 2x2 array looks like this:

[
    4, 7,
    10, 5
]

However you now need to know the dimensions for this to work because different rectangular shapes can have the same number of elements. So maybe let's make it an object.

{
    shape: [2,2],
    data: [
        4, 7,
        10, 5
    ]
}

The shape is an array rather than properties because we could scale this idea up into N-dimensional tensors. In fact, this is how libraries like tensorflowjs do it. For convenience let's build some function to convert between the formats.

export function nestedArrayToFlat(nested){
    return {
        shape: [nested.length, nested[0].length],
        data: nested.flat(Infinity)
    }
}

export function flatToNestedArray(flat){
    const data = new Array(flat.shape[0]);
    for(let row = 0; row < flat.shape[0]; row++){
        const rowArray = new Array(flat.shape[1]);
        for(let col = 0; col < flat.shape[1]; col++){
            rowArray[col] = flat.data[row * flat.shape[1] + col]; 
        }
        data[row] = rowArray;
    }
    return data;
}

So far I think pre-allocated arrays and loops have the best general performance scaling to larger values so we'll stick to that for now. This also means I'll be omitting flat and loop since they aren't winning in any category as well as dynamic because it's the same as unrolled.

export function addMatrixFlat(a, b) {
    const out = {
        shape: a.shape,
        data: new Array(a.data.length)
    };
    for (let row = 0; row < a.shape[0]; row++) {
        for (let col = 0; col < a.shape[1]; col++) {
            const index = (row * a.shape[1]) + col;
            out.data[index] = a.data[index] + b.data[index];
        }
    }
    return out;
}

It's about 20% faster on larger matrices than our previous best but 20% slower on 1x1 and 2x2 than unrolled. Since those aren't too important I'd say this is a big win.

Row vs column major

Does it matter if we traverse over rows versus columns? One might suspect that it could when CPU caching an stuff gets involved, but let's test.

export function addMatrixFlatColMajor(a, b) {
    const out = {
        shape: a.shape,
        data: new Array(a.data.length)
    };
    for (let col = 0; col < a.shape[1]; col++) {
        for (let row = 0; row < a.shape[0]; row++) {
            const index = (row * a.shape[1]) + col;
            out.data[index] = a.data[index] + b.data[index];
        }
    }
    return out;
}

As it turn out, column major traversal is actually a tiny bit slower than row major traversal. This is likely because the cache lines are being read more optimally.

However, since element-wise add is so simple we can actually just ditch the loop structure and just add all the elements linearly with a single loop.

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

This is like 20%+ faster.

Unrolled

We can also unroll these too and see what happens, maybe the simpler structure helps? Using this code:

export function genMatAddFlatBody(rows, cols){
    let funcBody = "return [\n";

    for (let r = 0; r < rows; r++) {
        for (let c = 0; c < cols; c++) {
            funcBody += `a[${r * cols + c}] + b[${r * cols + c}]${(c * r) < ((rows - 1) * (cols - 1)) ? ", " : ""}`
        }
    }

    funcBody += `];\n`
    return funcBody;
}

We can generate functions like this:

export function addMatrixFlat2x2(a, b) {
    return [
        a[0] + b[0], a[1] + b[1], a[2] + b[2], a[3] + b[3]];

}

We can dynamically create them with eval like this:

export function genMatAddFlatFunc(rows, cols) {
    rows = Number(rows);
    cols = Number(cols);
    const body = genMatAddFlatBody(rows, cols);
    return new Function("a", "b", body);
}

It just edges the simple loop out at 1x1 and 2x2 and beyond that it loses and get much, much worse at large sizes.

Typed arrays

So the next area of possible optimization I can see is to actually use types. We can do this in Javascript using typed-arrays. This will allow us to allocate a block of memory and reduce the overhead of any array structure. This is actually a little more important though. By using typed-arrays we can actually reduce conversions. APIs like WASM, WebGL and WebGPU deal with blocks of memory and the less we have to convert the faster we're going to wind up. So I think even if it turns out this is a little bit slower, there's still good reasons to pursue it. Although we end up with different paths, one for floats and one for integers and even then it might matter if we choose different bit-widths. Also, since we've already shown that flat structures perform better overall we don't need to consider nested arrays. For brevity I'm not going to test all Typed array combos as we'll start to see a general pattern.

Float 64

Javascript numbers are float 64s. So it's really surprising that these behave slower than a normal javascript array. Doing small arrays is actually slower than array.map. I'm guessing this has something to do with how the engine treats them. As the matrices get larger these get faster but even at 100x100 items it's still quite bit slower than a normal flat array.

Float 32

F32 arrays have the same problem as Float64s. The performance is near identical despite being smaller so for pure speed there's no point in choosing them. In fact at 100x100 F64 arrays are decently faster. The only benefit we get is halfing our memory which might be a reason to choose these.

Int 32

I32s again behave similarly but start to see much bigger gains at larger matrices. In fact at 100x100 the I32 matrix is about equal to a flat matrix. Not amazing but if you are dealing with large integer matrices this is probably your best choice.

Conclusion

For simple, single-threaded javascript we've observed a few things (in Deno/V8 @ 2023-03-31):

• Loops will mostly perform better than .map except at very large values and only with nested arrays (I tried a flat array and it wasn't notable enough to copy-paste the data).
• Bespoke unrolled functions work well on very small sizes 4x4 or less but doesn't beat a simple loop and falls off very, very quickly.
• Reducing structure makes a big difference.
• Pre-allocating arrays makes a huge different, always do this if you can.
• Typed-arrays provide no speed advantage (but we might get less conversion overhead and space-space saving).

There's more ways we can deal matrices, I'd like to maybe look at what WASM and WebGPU look like with high overhead but potentially massive speed increases in actual calculation due to parallelism. Web Workers as well. Also different ops can vary widely. Matrix multiplication uses the left-hand and right-hand structures differently and might require some different strategies. But I think the biggest take-away:

Your best bet for a generalized element-wise matrix op is a single flat loop over a normal JS array as it's fast and scales well

The Data

Code

https://github.com/ndesmic/fast-mat