JS快速矩阵计算

有很多方法可以在 JS 中表示矩阵数学。有些方法可读性强，有些方法速度快。我想探索一下这些差异。某些技术实际上能为我节省多少时间？

为此，我将只研究一个操作：逐元素加法以减少总案例数，但差异操作可能会稍微改变整体值，尤其是像矩阵乘法这样需要稍微复杂一些的应用程序规则的运算。这些状态也在我的计算机上，它是稍旧的 i7 8700K，使用 Deno，其底层是 v8。如果有不同的优化，像 Bun 这样的不同运行时可能会表现得非常不同。

本文相关的代码可以从这里下载。

1、简单的函数式方法

我想从这里开始，因为这是我为初稿编写它的方式。它的代码非常优化，但我怀疑性能实际上相当糟糕。不变性对于避免错误非常有用，但对于性能却很糟糕，尤其是当 JS 没有智能副本时。

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

矩阵表示是数组的数组。外部数组是行，内部数组是列。

使用 Deno 的内置基准测试工具，我们可以看到它在不同大小矩阵上的表现。

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);
});

mat100A 和 mat100B 是预先生成的 100x100 矩阵，太大了，无法放入测试文件中。

需要注意的是，我认为 Deno 至少不再允许您设置迭代或预热迭代。我认为它只是寻找数字的收敛。实际运行次数显示在 JSON 输出中，并且每个测试略有不同。

以下是我们的操作方式：

2、循环

所以我认为我们可以改进的第一个方法是循环。函数有开销，所以如果我们去掉它，并且更具有命令性，我们就可以更快一些。

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;
}

请注意，我不会进行严格的边界检查，我们只是假设 a 和 b 的大小相同，因为边界检查只会增加开销。

循环开始时速度更快，但一旦达到 32x32 左右，它们就等于 .map，并且大于 .map 时速度更快。非常令人惊讶！

3、预分配数组

我的下一个想法是预分配数组，因为推入数组可能会导致重新调整大小，也许这就是速度较慢的原因。

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;
}

成功了！我们比开始时快了 1.5 倍！

4、展开循环

如果我们删除所有循环并用手写出来会怎么样？

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]]
    ];

}

这不太灵活，因为需要为要添加的每种矩阵形状都提供一个函数。但是，在某些情况下（例如 3D），情况还不算太糟，因为我们拥有的东西数量非常有限，通常只有 4x4。在机器学习中，这可能会导致问题。

这是一个为展开循环生成 javascript 文本的函数：

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);
}

我也很好奇这种动态生成是否会带来很大的变化：

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);
}

由于我们使用 eval，所以我们应该确保清理输入：

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

对于较小的值，这是一个很大的改进，比预分配的循环快了大约 1.5 到 2 倍，但对于较大的值，速度要慢得多，这可不是件好事。我不确定为什么会这样，也许这与函数本身的大小有关？生成的代码非常庞大。此外，动态生成与写出它们基本相同。因此，如果想节省有效载荷（并且不受 CSP 的限制），你可以动态创建它们而不会受到任何惩罚。

5、展平数组

我认为我们可以节省的另一件事是数组。从技术上讲，我们不需要有很多嵌套数组，它们会增加一些创建开销。所以现在一个 2x2 数组看起来像这样：

[
    4, 7,
    10, 5
]

但是现在你需要知道尺寸才能实现这个功能，因为不同的矩形可以有相同数量的元素。所以也许我们可以把它变成一个对象。

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

形状是一个数组而不是属性，因为我们可以将这个想法扩展到 N 维张量。事实上，这就是 tensorflowjs 等库所做的。为了方便起见，让我们构建一些函数来在格式之间进行转换。

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;
}

到目前为止，我认为预分配数组和循环在扩展到较大值时具有最佳的总体性能，因此我们暂时坚持这一点。这也意味着我将省略平面和循环，因为它们在任何类别中都没有获胜，以及动态，因为它与展开相同。

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;
}

在处理较大的矩阵时，它比我们之前的最佳结果快了约 20%，但在处理 1x1 和 2x2 时，它比展开时慢了 20%。由于这些并不是太重要，我认为这是一个巨大的胜利。

6、行优先还是列优先

如果我们遍历行还是列，这有关系吗？有人可能会怀疑，当涉及到 CPU 缓存时，这可能会有关系，但让我们测试一下。

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;
}

事实证明，列主遍历实际上比行主遍历慢一点。这可能是因为缓存行的读取方式更优化。

但是，由于逐元素添加非常简单，我们实际上可以放弃循环结构，只需使用单个循环线性添加所有元素即可。

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;
}

这比原先快了 20% 以上。

7、展开

我们也可以展开这些，看看会发生什么，也许更简单的结构会有所帮助？使用此代码：

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;
}

我们可以生成如下函数：

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

}

我们可以使用 eval 像这样动态创建它们：

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

它只是在 1x1 和 2x2 时超越了简单循环，而在此之后，它会在较大尺寸时丢失并变得更糟。

8、类型化数组

因此，我能看到下一个可能的优化领域是实际使用类型。我们可以在 Javascript 中使用类型化数组来实现这一点。这将使我们能够分配一个内存块并减少任何数组结构的开销。但这实际上更重要一些。

通过使用类型化数组，我们实际上可以减少转换。WASM、WebGL 和 WebGPU 等 API 处理内存块，我们需要转换的越少，我们结束的速度就越快。所以我认为即使结果有点慢，仍然有充分的理由去追求它。

虽然我们最终得到了不同的路径，一个用于浮点数，一个用于整数，即使如此，如果我们选择不同的位宽，也可能很重要。

此外，由于我们已经表明平面结构总体上表现更好，所以我们不需要考虑嵌套数组。

为了简洁起见，我不会测试所有类型数组组合，因为我们将开始看到一个一般模式。

Float 64

JavaScript 数字是 64 位浮点数。因此，它们的表现比普通的 JavaScript 数组慢确实令人惊讶。处理小数组实际上比 array.map 慢。我猜这与引擎处理它们的方式有关。随着矩阵变大，它们会变得更快，但即使在 100x100 个项目时，它仍然比普通平面数组慢很多。

Float 32

F32 数组与 Float64 数组存在同样的问题。尽管更小，但性能几乎相同，因此单纯从速度角度考虑，选择它们毫无意义。事实上，在 100x100 时，F64 数组的速度相当快。我们获得的唯一好处是内存减少一半，这可能是选择它们的一个原因。

Int 32

I32 再次表现出类似的行为，但在较大的矩阵中开始看到更大的收益。事实上，在 100x100 时，I32 矩阵大约等于平面矩阵。这并不令人惊讶，但如果你正在处理大型整数矩阵，这可能是你的最佳选择。

9、结束语

对于简单的单线程 javascript，我们观察到了一些事情（在 Deno/V8 @ 2023-03-31）：

• 循环的性能大多优于 .map，除非值非常大并且只使用嵌套数组（我尝试了一个平面数组，但它不足以复制粘贴数据）。
• 定制的展开函数在 4x4 或更小的非常小的尺寸上效果很好，但无法击败简单的循环并且会非常非常快地下降。
• 减少结构会带来很大的不同。
• 预分配数组会带来巨大的不同，如果可以的话，请始终这样做。
• 类型化数组没有速度优势（但我们可能会获得更少的转换开销和空间节省）。

我们可以用更多方式来处理矩阵，我可能想看看 WASM 和 WebGPU 是什么样子的，它们开销很大，但由于并行性，实际计算速度可能会大幅提高。Web Workers 也是如此。不同的操作也可能有很大差异。矩阵乘法使用左侧和右侧结构的方式不同，可能需要一些不同的策略。但我认为最大的收获是：

对于广义的元素矩阵操作，最好的选择是对普通 JS 数组进行单一平面循环，因为它速度快，扩展性好

完整的对比记录：

原文链接：Fast Matrix Math in JS

BimAnt翻译整理，转载请标明出处