Gitea Mirror Bot
116 lines
2.3 KiB
C++
116 lines
2.3 KiB
C++
/*++
|
|
|
|
Copyright (c) Microsoft Corporation. All rights reserved.
|
|
|
|
Licensed under the MIT License.
|
|
|
|
Module Name:
|
|
|
|
sgemm_pack_b_rvv.cpp
|
|
|
|
Abstract:
|
|
|
|
This module implements an RVV packing helper for the single precision
|
|
matrix/matrix multiply operation (SGEMM) on riscv64.
|
|
|
|
--*/
|
|
|
|
#include "mlasi.h"
|
|
|
|
#if defined(MLAS_USE_RVV)
|
|
|
|
#include <riscv_vector.h>
|
|
|
|
namespace {
|
|
|
|
// Keep MLAS packing in 16-column tiles, but let RVV decide the actual chunk
|
|
// size at runtime via vsetvl so the same code works across different VLENs.
|
|
constexpr size_t kPackedCountN = 16;
|
|
|
|
MLAS_FORCEINLINE
|
|
void
|
|
MlasStoreZeroPaddedBlock(
|
|
float* D,
|
|
const float* B,
|
|
size_t CountX
|
|
)
|
|
{
|
|
size_t remaining = kPackedCountN;
|
|
size_t offset = 0;
|
|
|
|
while (remaining > 0) {
|
|
const size_t vl = __riscv_vsetvl_e32m4(remaining);
|
|
__riscv_vse32_v_f32m4(D + offset, __riscv_vfmv_v_f_f32m4(0.0f, vl), vl);
|
|
offset += vl;
|
|
remaining -= vl;
|
|
}
|
|
|
|
remaining = CountX;
|
|
offset = 0;
|
|
|
|
while (remaining > 0) {
|
|
const size_t vl = __riscv_vsetvl_e32m4(remaining);
|
|
__riscv_vse32_v_f32m4(D + offset, __riscv_vle32_v_f32m4(B + offset, vl), vl);
|
|
offset += vl;
|
|
remaining -= vl;
|
|
}
|
|
}
|
|
|
|
MLAS_FORCEINLINE
|
|
void
|
|
MlasStoreFullBlock(
|
|
float* D,
|
|
const float* B
|
|
)
|
|
{
|
|
size_t remaining = kPackedCountN;
|
|
size_t offset = 0;
|
|
|
|
while (remaining > 0) {
|
|
const size_t vl = __riscv_vsetvl_e32m4(remaining);
|
|
__riscv_vse32_v_f32m4(D + offset, __riscv_vle32_v_f32m4(B + offset, vl), vl);
|
|
offset += vl;
|
|
remaining -= vl;
|
|
}
|
|
}
|
|
|
|
} // namespace
|
|
|
|
void
|
|
MlasSgemmCopyPackBRvv(
|
|
float* D,
|
|
const float* B,
|
|
size_t ldb,
|
|
size_t CountX,
|
|
size_t CountY
|
|
)
|
|
{
|
|
while (CountX >= kPackedCountN) {
|
|
const float* b = B;
|
|
size_t y = CountY;
|
|
|
|
do {
|
|
MlasStoreFullBlock(D, b);
|
|
D += kPackedCountN;
|
|
b += ldb;
|
|
y--;
|
|
} while (y > 0);
|
|
|
|
B += kPackedCountN;
|
|
CountX -= kPackedCountN;
|
|
}
|
|
|
|
if (CountX > 0) {
|
|
size_t y = CountY;
|
|
|
|
do {
|
|
MlasStoreZeroPaddedBlock(D, B, CountX);
|
|
D += kPackedCountN;
|
|
B += ldb;
|
|
y--;
|
|
} while (y > 0);
|
|
}
|
|
}
|
|
|
|
#endif // defined(MLAS_USE_RVV)
|