高性能矩阵乘法库中Tensor转置适配matmul函数的问题求助

你的matmul函数默认输入矩阵是连续的行优先物理存储，但当前Tensor.transpose()只是做了逻辑转置——仅交换shape和stride，并未改变数据的物理存储顺序。因此B.transpose().data()返回的还是原B的连续存储地址，和matmul预期的[K,N]连续矩阵结构不匹配，导致计算错误。

以下是三种实用的修改方向，按高性能库的设计优先级排序：

1. 修改matmul以支持带步长（stride）的矩阵乘法 #

这是最合理的方案——高性能矩阵库本就应该支持非连续张量，避免不必要的数据拷贝。

步骤1：扩展matmul参数，接收矩阵的步长信息 #

// 支持非连续存储的matmul实现
void matmul(float* A, const std::vector<uint32_t>& A_stride,
            float* B, const std::vector<uint32_t>& B_stride,
            float* C, int M, int N, int K) {
    // 按逻辑索引计算：C[i][j] = sum_{k=0 to K-1} A[i][k] * B[k][j]
    for (int i = 0; i < M; ++i) {
        for (int j = 0; j < N; ++j) {
            float sum = 0.0f;
            for (int k = 0; k < K; ++k) {
                // 计算A[i][k]的物理地址：基地址 + i*行步长 + k*列步长
                float a_val = *(A + i * A_stride[0] + k * A_stride[1]);
                // 计算B[k][j]的物理地址（转置后的B逻辑上是[K,N]）
                float b_val = *(B + k * B_stride[0] + j * B_stride[1]);
                sum += a_val * b_val;
            }
            // 假设C是连续行优先存储，直接按顺序写入
            C[i * N + j] = sum;
        }
    }
}

步骤2：修改transpose为非原地版本（避免修改原Tensor） #

把原void类型的transpose改成返回新Tensor的版本，防止意外修改原对象的shape和stride：

class Tensor {
    void* data;
    std::vector<uint32_t> stride;
    std::vector<uint32_t> shape;

public:
    // 返回转置后的新Tensor，原Tensor不变
    Tensor transpose(uint32_t dim0, uint32_t dim1) const {
        Tensor res = *this;
        std::swap(res.stride[dim0], res.stride[dim1]);
        std::swap(res.shape[dim0], res.shape[dim1]);
        return res;
    }

    // 暴露必要的成员访问接口
    void* data() const { return data; }
    const std::vector<uint32_t>& stride() const { return stride; }
};

调用方式 #

// 得到B的逻辑转置Tensor（无数据拷贝）
auto B_t = B.transpose(0, 1);
// 调用带步长的matmul
matmul(static_cast<float*>(A.data()), A.stride(),
       static_cast<float*>(B_t.data()), B_t.stride(),
       static_cast<float*>(C.data()), M, N, K);

2. 实现物理转置（数据重排） #

如果暂时无法修改matmul，可以做一个物理转置函数，把逻辑转置后的Tensor数据重排成连续的行优先存储：

Tensor physical_transpose(const Tensor& t, uint32_t dim0, uint32_t dim1) {
    Tensor res;
    // 设置转置后的shape和连续存储的stride
    res.shape = t.shape;
    std::swap(res.shape[dim0], res.shape[dim1]);
    res.stride = {res.shape[1], 1}; // 行优先连续存储的步长

    // 分配连续内存
    size_t elem_count = res.shape[0] * res.shape[1];
    res.data = malloc(elem_count * sizeof(float));

    float* src = static_cast<float*>(t.data);
    float* dst = static_cast<float*>(res.data);
    int dim0_size = t.shape[dim0];
    int dim1_size = t.shape[dim1];

    // 按逻辑转置的顺序拷贝数据
    for (int i = 0; i < dim0_size; ++i) {
        for (int j = 0; j < dim1_size; ++j) {
            float val = *(src + i * t.stride[dim0] + j * t.stride[dim1]);
            *(dst + j * res.stride[0] + i * res.stride[1]) = val;
        }
    }
    return res;
}

调用方式 #

// 得到物理上连续的转置矩阵（会产生数据拷贝）
auto B_t = physical_transpose(B, 0, 1);
// 直接调用原matmul
matmul(static_cast<float*>(A.data()), static_cast<float*>(B_t.data()),
       static_cast<float*>(C.data()), M, N, K);

注意：该方案会引入数据拷贝，适合小矩阵或无法修改matmul的场景，性能低于方案1。

3. 封装Tensor专属的matmul重载 #

为了简化调用，可以封装一个直接接收Tensor参数的matmul，自动处理shape和stride的校验与传递：

void matmul(const Tensor& A, const Tensor& B, Tensor& C) {
    // 校验矩阵乘法的shape合法性
    assert(A.shape.back() == B.shape[B.shape.size() - 2]);
    int M = A.shape[0];
    int K = A.shape[1];
    int N = B.shape[1];

    // 调用带步长的底层实现
    matmul(static_cast<float*>(A.data()), A.stride(),
           static_cast<float*>(B.data()), B.stride(),
           static_cast<float*>(C.data()), M, N, K);
}

调用方式 #

auto B_t = B.transpose(0, 1);
matmul(A, B_t, C);

内容的提问来源于stack exchange，提问作者Frank Ngwee