diff --git a/src/cpu/rv64/rvv_matmul.cpp b/src/cpu/rv64/rvv_matmul.cpp index 02ea41dc8c8..8a9889a8fe4 100644 --- a/src/cpu/rv64/rvv_matmul.cpp +++ b/src/cpu/rv64/rvv_matmul.cpp @@ -50,6 +50,7 @@ status_t rvv_matmul_t::execute(const exec_ctx_t &ctx) const { const dim_t M = pd()->M_; const dim_t K = pd()->K_; const dim_t N = pd()->N_; + const bool weights_col_major = pd()->weights_col_major_; // row-major [M, K] / [M, N] are reinterpreted as column-major matrices and // mapped to GEMM as: @@ -62,12 +63,14 @@ status_t rvv_matmul_t::execute(const exec_ctx_t &ctx) const { // column-major [K, M] with leading dim K, so transb = 'N', ldb = K. // - C is viewed as column-major [N, M] with leading dim N, which matches // row-major [M, N] in memory. - char transa = 'N'; + char transa = weights_col_major ? 'T' : 'N'; char transb = 'N'; dim_t M_gemm = N; dim_t N_gemm = M; dim_t K_gemm = K; - dim_t lda = N; // weights: row-major [K, N] -> col-major [N, K] + // weights: col-major [K, N] uses leading K; row-major [K, N] is equivalent + // to col-major [N, K] so leading dim is N. + dim_t lda = weights_col_major ? K : N; dim_t ldb = K; // src: row-major [M, K] -> col-major [K, M] dim_t ldc = N; // dst: row-major [M, N] -> col-major [N, M] float alpha = 1.0f; diff --git a/src/cpu/rv64/rvv_matmul.hpp b/src/cpu/rv64/rvv_matmul.hpp index ec5cbd6b016..815881df5a1 100644 --- a/src/cpu/rv64/rvv_matmul.hpp +++ b/src/cpu/rv64/rvv_matmul.hpp @@ -170,6 +170,7 @@ struct rvv_matmul_t : public primitive_t { M_ = src_dims[ndims - 2]; K_ = src_dims[ndims - 1]; N_ = wei_dims[wei_ndims - 1]; + weights_col_major_ = is_col_major(weights_mdw); dim_t weights_batch_size = 1; for (int i = 0; i < wei_ndims - 2; ++i) @@ -182,6 +183,7 @@ struct rvv_matmul_t : public primitive_t { dim_t K_ = 0; dim_t batch_ = 0; bool weights_are_broadcast_ = false; + bool weights_col_major_ = false; }; rvv_matmul_t(const pd_t *apd) : primitive_t(apd) {}