如何用Java 17 Vector API优化矩阵乘法及边界处理

1. 处理非向量长度倍数的矩阵/数组 #

当矩阵维度或数组长度不是向量宽度的整数倍时，有两种可靠的处理方式：

• 分块+标量收尾：先处理能被向量长度整除的元素块，用Vector API批量计算；剩余不足一个向量长度的元素，用传统标量循环处理，避免越界。
• 掩码向量操作：利用Vector API的掩码功能，创建仅覆盖有效元素的掩码，在向量加载、计算时只对掩码标记的元素执行操作，自动忽略超出范围的位置。

2. Vector API优化矩阵乘法示例 #

Java 17的Vector API属于孵化器模块，编译运行时需添加--add-modules jdk.incubator.vector参数。核心优化思路是将矩阵点积中的乘累加操作转为向量指令，减少循环迭代次数，同时优化缓存访问模式。

优化代码实现 #

import jdk.incubator.vector.IntVector;
import jdk.incubator.vector.VectorSpecies;

public class VectorizedMatrixMult {
    // 自动适配当前CPU最优的int向量宽度
    private static final VectorSpecies<Integer> INT_SPECIES = IntVector.SPECIES_PREFERRED;

    public static void main(String[] args) {
        int rowA = 4, colA = 3, rowB = 3, colB = 4;
        int[][] A = {{1,1,1}, {2,2,2}, {3,3,3}, {4,4,4}};
        int[][] B = {{1,1,1,1}, {2,2,2,2}, {3,3,3,3}};

        if (rowB != colA) {
            System.out.println("Multiplication Not Possible");
            return;
        }

        int[][] result = new int[rowA][colB];
        vectorizedMultiply(A, B, result, rowA, colA, colB);

        // 输出结果
        for (int[] row : result) {
            for (int val : row) System.out.print(val + " ");
            System.out.println();
        }
    }

    private static void vectorizedMultiply(int[][] A, int[][] B, int[][] result, int rowA, int colA, int colB) {
        int vecLen = INT_SPECIES.length();
        // 预先转置B矩阵，将列访问转为行访问，优化缓存命中率
        int[][] transposedB = transposeMatrix(B);

        for (int i = 0; i < rowA; i++) {
            int[] rowAData = A[i];
            for (int j = 0; j < colB; j++) {
                int sum = 0;
                int k = 0;

                // 批量处理向量长度倍数的元素
                for (; k <= colA - vecLen; k += vecLen) {
                    IntVector vecA = IntVector.fromArray(INT_SPECIES, rowAData, k);
                    IntVector vecB = IntVector.fromArray(INT_SPECIES, transposedB[j], k);
                    sum += vecA.mul(vecB).reduceLanesToInt(Integer::sum);
                }

                // 处理剩余元素（标量方式）
                for (; k < colA; k++) {
                    sum += rowAData[k] * B[k][j];
                }

                // 或者用掩码方式处理剩余元素，替代上面的标量循环
                /*
                if (k < colA) {
                    var mask = INT_SPECIES.indexInRange(k, colA);
                    IntVector vecA = IntVector.fromArray(INT_SPECIES, rowAData, k, mask);
                    IntVector vecB = IntVector.fromArray(INT_SPECIES, transposedB[j], k, mask);
                    sum += vecA.mul(vecB).reduceLanesToInt(Integer::sum, mask);
                }
                */

                result[i][j] = sum;
            }
        }
    }

    // 矩阵转置辅助方法：将B的列转为行，优化缓存访问
    private static int[][] transposeMatrix(int[][] matrix) {
        int rows = matrix.length;
        int cols = matrix[0].length;
        int[][] transposed = new int[cols][rows];
        for (int i = 0; i < rows; i++) {
            for (int j = 0; j < cols; j++) {
                transposed[j][i] = matrix[i][j];
            }
        }
        return transposed;
    }
}

关键优化点说明 #

• 矩阵转置：原代码中频繁访问B矩阵的列会导致缓存行失效，转置后将列访问转为连续的行访问，大幅提升缓存利用率。
• 向量乘累加：用IntVector.mul()完成批量乘法，再通过reduceLanesToInt()完成向量内元素的求和，替代原有的标量乘累加循环。
• 掩码处理：注释中的掩码方式可以统一向量和标量处理逻辑，无需分支判断剩余元素。

编译运行命令 #

# 编译
javac --add-modules jdk.incubator.vector VectorizedMatrixMult.java
# 运行
java --add-modules jdk.incubator.vector VectorizedMatrixMult

内容的提问来源于stack exchange，提问作者Isuru Perera