SIMD快速矩阵乘法：1D float[]代码适配2D float[][]相关疑问

适配核心思路 #

原来的1D存储是将矩阵按行优先（也有列优先的场景，看原实现约定）平铺为一维数组，索引规则为矩阵[i行j列] = 1D数组[i * 列数 + j]，适配时只需要把索引规则替换为2D数组的原生下标访问即可，核心适配点如下：

• 入参替换：将原方法的入参从float[] a, float[] b, float[] res改为float[][] a, float[][] b, float[][] res，调用前需先做维度合法性校验：a的列数等于b的行数，res的行数等于a的行数、列数等于b的列数，且每一行的长度统一，避免不规则2D数组导致的越界问题。
• 索引逻辑替换：直接用a[i][k]、b[k][j]、res[i][j]替换原来的1D偏移计算表达式即可，无需再手动计算全局偏移量。
• SIMD加载逻辑适配：绝大多数语言的2D数组单一行的内存是连续存储的，加载SIMD向量时可以直接从行数组的对应下标加载，不用再计算全局偏移，代码可读性更强。比如Java VectorAPI的加载逻辑可以从FloatVector.fromArray(SPECIES, a, i * aCols + k)简化为FloatVector.fromArray(SPECIES, a[i], k)。
• 性能优化小技巧：循环中可以提前把当前行的1D数组引用取出来（比如float[] aRow = a[i]），避免每次访问都重复做2D数组的行寻址，性能和1D版本基本持平。

核心代码示例 #

// 原1D版本核心循环片段
for (int i = 0; i < aRows; i++) {
  for (int k = 0; k < aCols; k++) {
    float aVal = a[i * aCols + k];
    for (int j = 0; j < bCols; j += SPECIES.length()) {
      var bVec = FloatVector.fromArray(SPECIES, b, k * bCols + j);
      var resVec = FloatVector.fromArray(SPECIES, res, i * bCols + j);
      resVec = bVec.mul(aVal).add(resVec);
      resVec.intoArray(res, i * bCols + j);
    }
  }
}

// 适配后2D版本核心循环片段
// 前置维度校验：a[0].length == b.length && res.length == a.length && res[0].length == b[0].length
int aRows = a.length, aCols = a[0].length, bCols = b[0].length;
for (int i = 0; i < aRows; i++) {
  float[] aRow = a[i];
  float[] resRow = res[i];
  for (int k = 0; k < aCols; k++) {
    float aVal = aRow[k];
    float[] bRow = b[k];
    for (int j = 0; j < bCols; j += SPECIES.length()) {
      var bVec = FloatVector.fromArray(SPECIES, bRow, j);
      var resVec = FloatVector.fromArray(SPECIES, resRow, j);
      resVec = bVec.mul(aVal).add(resVec);
      resVec.intoArray(resRow, j);
    }
  }
}

是高性能矩阵运算场景下的标准通用做法，核心原因有三点：

• 内存连续性更优：1D数组是整块连续内存，CPU缓存预取命中率远高于可能存在行间隙的2D数组（尤其是数组的数组实现的2D结构），也更方便做内存对齐、边界填充等SIMD友好的优化，适配CPU、GPU等各类并行计算硬件的要求。
• 索引计算效率更高：i * 列数 + j的索引计算可以被编译器深度优化，比2D数组先取行指针、再取下标的两次寻址开销更低，多层循环嵌套下累积的性能收益非常明显。
• 生态兼容性更强：所有主流高性能矩阵运算库（BLAS、OpenBLAS、MKL、各类深度学习框架）的底层实现都是用1D连续内存存储矩阵，上层暴露的2D下标访问只是语法糖，本质还是调用1D索引计算。

如果是小矩阵运算、追求代码可读性优先的业务场景，直接用2D数组完全没问题，二者的性能差异可以忽略。

内容的提问来源于stack exchange，提问作者Ondrej Sotolar