mirror of
https://github.com/saymrwulf/onnxruntime.git
synced 2026-07-08 17:17:15 +00:00
MLAS: add SSE 4.1 u8s8 kernel (#7490)
This commit is contained in:
Tracy Sharpe
GitHub
parent
e73c3e0651
commit
2b0bbfd1a8
5 changed files with 463 additions and 29 deletions
|
|
@ -655,6 +655,7 @@ MlasSgemmOperation(
|
|||
struct MLAS_GEMM_U8X8_DISPATCH;
|
||||
|
||||
extern const MLAS_GEMM_U8X8_DISPATCH MlasGemmU8X8DispatchSse;
|
||||
extern const MLAS_GEMM_U8X8_DISPATCH MlasGemmU8S8DispatchSse41;
|
||||
extern const MLAS_GEMM_U8X8_DISPATCH MlasGemmU8S8DispatchAvx2;
|
||||
extern const MLAS_GEMM_U8X8_DISPATCH MlasGemmU8U8DispatchAvx2;
|
||||
extern const MLAS_GEMM_U8X8_DISPATCH MlasGemmU8X8DispatchNeon;
|
||||
|
|
@ -770,7 +771,7 @@ MlasExecuteThreaded(
|
|||
|
||||
/**
|
||||
* @brief Distribute multiple iterations of work over a thread pool if supported
|
||||
*
|
||||
*
|
||||
* @param ThreadPool [IN] Optional thread pool. Ignored when using OpenMP
|
||||
* @param Iterations [IN] Total number of iterations
|
||||
* @param Work [IN] Logic for computing a range of iterations [begin, end)
|
||||
|
|
|
|||
|
|
@ -163,10 +163,6 @@ Return Value:
|
|||
|
||||
#endif
|
||||
|
||||
//
|
||||
// Check if the processor supports the AVX and OSXSAVE features.
|
||||
//
|
||||
|
||||
unsigned Cpuid1[4];
|
||||
#if defined(_WIN32)
|
||||
__cpuid((int*)Cpuid1, 1);
|
||||
|
|
@ -174,6 +170,22 @@ Return Value:
|
|||
__cpuid(1, Cpuid1[0], Cpuid1[1], Cpuid1[2], Cpuid1[3]);
|
||||
#endif
|
||||
|
||||
#if defined(MLAS_TARGET_AMD64) && defined(_MSC_VER)
|
||||
|
||||
//
|
||||
// Check if the processor supports SSE 4.1 instructions.
|
||||
//
|
||||
|
||||
if ((Cpuid1[2] & 0x80000) != 0) {
|
||||
this->GemmU8S8Dispatch = &MlasGemmU8S8DispatchSse41;
|
||||
}
|
||||
|
||||
#endif
|
||||
|
||||
//
|
||||
// Check if the processor supports the AVX and OSXSAVE features.
|
||||
//
|
||||
|
||||
if ((Cpuid1[2] & 0x18000000) == 0x18000000) {
|
||||
|
||||
//
|
||||
|
|
|
|||
|
|
@ -154,7 +154,12 @@ int32_t
|
|||
MlasGemmU8X8FixupZeroPointB(
|
||||
int32_t ZeroPointB,
|
||||
bool BIsSigned
|
||||
);
|
||||
)
|
||||
{
|
||||
MLAS_UNREFERENCED_PARAMETER(BIsSigned);
|
||||
|
||||
return ZeroPointB;
|
||||
}
|
||||
|
||||
template<typename KernelType>
|
||||
MLAS_FORCEINLINE
|
||||
|
|
@ -741,7 +746,7 @@ MlasGemmU8X8CopyPackA<MLAS_GEMM_U8X8_KERNEL_SSE>(
|
|||
|
||||
while (k >= 8) {
|
||||
|
||||
__m128i Bytes = _mm_loadl_epi64((__m128i*)&a[0]);
|
||||
__m128i Bytes = _mm_loadl_epi64((const __m128i*)&a[0]);
|
||||
__m128i Words = _mm_unpacklo_epi8(Bytes, ZeroVector);
|
||||
|
||||
ReductionVector = _mm_add_epi16(ReductionVector, Words);
|
||||
|
|
@ -864,8 +869,8 @@ MlasGemmU8X8CopyPackB<MLAS_GEMM_U8X8_KERNEL_SSE>(
|
|||
|
||||
while (k >= MLAS_GEMM_U8X8_KERNEL_SSE::PackedK) {
|
||||
|
||||
__m128i BytesRow0 = _mm_loadl_epi64((__m128i*)&b[0]);
|
||||
__m128i BytesRow1 = _mm_loadl_epi64((__m128i*)&b[ldb]);
|
||||
__m128i BytesRow0 = _mm_loadl_epi64((const __m128i*)&b[0]);
|
||||
__m128i BytesRow1 = _mm_loadl_epi64((const __m128i*)&b[ldb]);
|
||||
|
||||
MlasGemmU8X8CopyPackBProcessSse(D, BytesRow0, BytesRow1, BitFlipVector, ColumnSums);
|
||||
|
||||
|
|
@ -876,7 +881,7 @@ MlasGemmU8X8CopyPackB<MLAS_GEMM_U8X8_KERNEL_SSE>(
|
|||
|
||||
if (k > 0) {
|
||||
|
||||
__m128i BytesRow0 = _mm_loadl_epi64((__m128i*)&b[0]);
|
||||
__m128i BytesRow0 = _mm_loadl_epi64((const __m128i*)&b[0]);
|
||||
|
||||
MlasGemmU8X8CopyPackBProcessSse(D, BytesRow0, BitFlipVector, BitFlipVector, ColumnSums);
|
||||
|
||||
|
|
@ -1026,8 +1031,8 @@ MlasGemmU8X8Kernel<MLAS_GEMM_U8X8_KERNEL_SSE>(
|
|||
Accumulators[1] = Accumulators[0];
|
||||
}
|
||||
|
||||
Accumulators[0] = _mm_add_epi32(Accumulators[0], _mm_loadu_si128((__m128i*)&ColumnSumBuffer[0]));
|
||||
Accumulators[1] = _mm_add_epi32(Accumulators[1], _mm_loadu_si128((__m128i*)&ColumnSumBuffer[4]));
|
||||
Accumulators[0] = _mm_add_epi32(Accumulators[0], _mm_loadu_si128((const __m128i*)&ColumnSumBuffer[0]));
|
||||
Accumulators[1] = _mm_add_epi32(Accumulators[1], _mm_loadu_si128((const __m128i*)&ColumnSumBuffer[4]));
|
||||
ColumnSumBuffer += 8;
|
||||
|
||||
//
|
||||
|
|
@ -1147,6 +1152,435 @@ const MLAS_GEMM_U8X8_DISPATCH MlasGemmU8X8DispatchSse = {
|
|||
|
||||
#endif
|
||||
|
||||
// N.B. MSVC does not require turning on SSE 4.1 intrinsics and the current use
|
||||
// for this code is Windows only, so restrict this kernel to that environment.
|
||||
#if defined(MLAS_SSE2_INTRINSICS) && defined(_MSC_VER)
|
||||
|
||||
struct MLAS_GEMM_U8S8_KERNEL_SSE41
|
||||
{
|
||||
typedef uint8_t PackedAType;
|
||||
typedef uint8_t PackedBType;
|
||||
typedef int8_t OffsetBType;
|
||||
|
||||
static constexpr size_t PackedK = 4;
|
||||
static constexpr MLAS_GEMM_U8X8_STRIDES Strides{24, 128, 128};
|
||||
static constexpr MLAS_GEMM_U8X8_STRIDES PackedStrides{24, 128, 128};
|
||||
};
|
||||
|
||||
constexpr size_t MLAS_GEMM_U8S8_KERNEL_SSE41::PackedK;
|
||||
constexpr MLAS_GEMM_U8X8_STRIDES MLAS_GEMM_U8S8_KERNEL_SSE41::Strides;
|
||||
constexpr MLAS_GEMM_U8X8_STRIDES MLAS_GEMM_U8S8_KERNEL_SSE41::PackedStrides;
|
||||
|
||||
template<>
|
||||
void
|
||||
MlasGemmU8X8CopyPackA<MLAS_GEMM_U8S8_KERNEL_SSE41>(
|
||||
MLAS_GEMM_U8S8_KERNEL_SSE41::PackedAType* D,
|
||||
const uint8_t* A,
|
||||
size_t lda,
|
||||
size_t CountM,
|
||||
size_t CountK,
|
||||
int32_t* RowSumBuffer
|
||||
)
|
||||
{
|
||||
const __m128i ZeroVector = _mm_setzero_si128();
|
||||
const __m128i OnesWordBroadcast = _mm_set1_epi16(1);
|
||||
|
||||
//
|
||||
// Process a single row of matrix A in a loop.
|
||||
//
|
||||
|
||||
while (CountM > 0) {
|
||||
|
||||
const uint8_t* a = A;
|
||||
size_t k = CountK;
|
||||
__m128i ReductionVector = ZeroVector;
|
||||
|
||||
//
|
||||
// Copy the source bytes to the packed buffer.
|
||||
//
|
||||
// The packed buffer has the same data ordering as the source bytes,
|
||||
// but CountK is aligned up to a multiple of 4 to maintain 32-bit
|
||||
// alignment. All extra bytes are zero-padded.
|
||||
//
|
||||
|
||||
while (k >= 8) {
|
||||
|
||||
__m128i Bytes = _mm_loadl_epi64((const __m128i*)&a[0]);
|
||||
|
||||
__m128i Words = _mm_unpacklo_epi8(Bytes, ZeroVector);
|
||||
ReductionVector = _mm_add_epi32(ReductionVector, _mm_madd_epi16(Words, OnesWordBroadcast));
|
||||
|
||||
_mm_storel_epi64((__m128i*)&D[0], Bytes);
|
||||
|
||||
a += 8;
|
||||
D += 8;
|
||||
k -= 8;
|
||||
}
|
||||
|
||||
if (k > 0) {
|
||||
|
||||
//
|
||||
// Copy the remaining bytes to the zero padded stack buffer.
|
||||
//
|
||||
|
||||
_mm_storel_epi64((__m128i*)&D[0], ZeroVector);
|
||||
|
||||
std::copy_n(&a[0], k, &D[0]);
|
||||
|
||||
__m128i Bytes = _mm_loadl_epi64((__m128i*)&D[0]);
|
||||
D += (k + 3) & ~3;
|
||||
|
||||
__m128i Words = _mm_unpacklo_epi8(Bytes, ZeroVector);
|
||||
ReductionVector = _mm_add_epi32(ReductionVector, _mm_madd_epi16(Words, OnesWordBroadcast));
|
||||
}
|
||||
|
||||
//
|
||||
// Reduce the partial accumulators.
|
||||
//
|
||||
|
||||
ReductionVector = _mm_hadd_epi32(ReductionVector, ReductionVector);
|
||||
ReductionVector = _mm_hadd_epi32(ReductionVector, ReductionVector);
|
||||
|
||||
*RowSumBuffer++ = _mm_cvtsi128_si32(ReductionVector);
|
||||
|
||||
A += lda;
|
||||
CountM -= 1;
|
||||
}
|
||||
}
|
||||
|
||||
MLAS_FORCEINLINE
|
||||
void
|
||||
MlasGemmU8X8CopyPackBProcessSse41(
|
||||
MLAS_GEMM_U8S8_KERNEL_SSE41::PackedBType* D,
|
||||
__m128i BytesRows[4],
|
||||
__m128i OnesByteBroadcast,
|
||||
__m128i OnesWordBroadcast,
|
||||
__m128i ColumnSums[2]
|
||||
)
|
||||
{
|
||||
__m128i PairsInterleaved0 = _mm_unpacklo_epi8(BytesRows[0], BytesRows[1]);
|
||||
__m128i PairsInterleaved1 = _mm_unpacklo_epi8(BytesRows[2], BytesRows[3]);
|
||||
|
||||
__m128i QuadsInterleaved0 = _mm_unpacklo_epi16(PairsInterleaved0, PairsInterleaved1);
|
||||
__m128i QuadsInterleaved1 = _mm_unpackhi_epi16(PairsInterleaved0, PairsInterleaved1);
|
||||
|
||||
__m128i PairwiseAdd0 = _mm_maddubs_epi16(OnesByteBroadcast, QuadsInterleaved0);
|
||||
__m128i PairwiseAdd1 = _mm_maddubs_epi16(OnesByteBroadcast, QuadsInterleaved1);
|
||||
|
||||
PairwiseAdd0 = _mm_madd_epi16(PairwiseAdd0, OnesWordBroadcast);
|
||||
PairwiseAdd1 = _mm_madd_epi16(PairwiseAdd1, OnesWordBroadcast);
|
||||
|
||||
ColumnSums[0] = _mm_add_epi32(ColumnSums[0], PairwiseAdd0);
|
||||
ColumnSums[1] = _mm_add_epi32(ColumnSums[1], PairwiseAdd1);
|
||||
|
||||
_mm_storeu_si128((__m128i*)&D[0], QuadsInterleaved0);
|
||||
_mm_storeu_si128((__m128i*)&D[16], QuadsInterleaved1);
|
||||
}
|
||||
|
||||
template<>
|
||||
void
|
||||
MlasGemmU8X8CopyPackB<MLAS_GEMM_U8S8_KERNEL_SSE41>(
|
||||
MLAS_GEMM_U8S8_KERNEL_SSE41::PackedBType* D,
|
||||
const uint8_t* B,
|
||||
size_t ldb,
|
||||
size_t CountN,
|
||||
size_t CountK,
|
||||
int32_t* ColumnSumBuffer,
|
||||
bool BIsSigned
|
||||
)
|
||||
{
|
||||
const __m128i OnesByteBroadcast = _mm_set1_epi8(1);
|
||||
const __m128i OnesWordBroadcast = _mm_set1_epi16(1);
|
||||
__m128i BytesRows[4];
|
||||
|
||||
MLAS_UNREFERENCED_PARAMETER(BIsSigned);
|
||||
|
||||
//
|
||||
// Process 8 columns of matrix B in a loop.
|
||||
//
|
||||
|
||||
while (CountN >= 8) {
|
||||
|
||||
const uint8_t* b = B;
|
||||
size_t k = CountK;
|
||||
__m128i ColumnSums[2];
|
||||
|
||||
ColumnSums[0] = _mm_setzero_si128();
|
||||
ColumnSums[1] = _mm_setzero_si128();
|
||||
|
||||
//
|
||||
// Interleave rows of matrix B and write to the packed buffer.
|
||||
//
|
||||
|
||||
while (k >= MLAS_GEMM_U8S8_KERNEL_SSE41::PackedK) {
|
||||
|
||||
BytesRows[0] = _mm_loadl_epi64((const __m128i*)&b[ldb * 0]);
|
||||
BytesRows[1] = _mm_loadl_epi64((const __m128i*)&b[ldb * 1]);
|
||||
BytesRows[2] = _mm_loadl_epi64((const __m128i*)&b[ldb * 2]);
|
||||
BytesRows[3] = _mm_loadl_epi64((const __m128i*)&b[ldb * 3]);
|
||||
|
||||
MlasGemmU8X8CopyPackBProcessSse41(D, BytesRows, OnesByteBroadcast, OnesWordBroadcast, ColumnSums);
|
||||
|
||||
b += ldb * 4;
|
||||
D += 32;
|
||||
k -= 4;
|
||||
}
|
||||
|
||||
if (k > 0) {
|
||||
|
||||
BytesRows[0] = _mm_loadl_epi64((const __m128i*)&b[ldb * 0]);
|
||||
BytesRows[1] = _mm_setzero_si128();
|
||||
BytesRows[2] = _mm_setzero_si128();
|
||||
BytesRows[3] = _mm_setzero_si128();
|
||||
|
||||
if (k >= 2) {
|
||||
BytesRows[1] = _mm_loadl_epi64((const __m128i*)&b[ldb * 1]);
|
||||
}
|
||||
|
||||
if (k >= 3) {
|
||||
BytesRows[2] = _mm_loadl_epi64((const __m128i*)&b[ldb * 2]);
|
||||
}
|
||||
|
||||
MlasGemmU8X8CopyPackBProcessSse41(D, BytesRows, OnesByteBroadcast, OnesWordBroadcast, ColumnSums);
|
||||
|
||||
D += 32;
|
||||
}
|
||||
|
||||
_mm_storeu_si128((__m128i*)&ColumnSumBuffer[0], ColumnSums[0]);
|
||||
_mm_storeu_si128((__m128i*)&ColumnSumBuffer[4], ColumnSums[1]);
|
||||
ColumnSumBuffer += 8;
|
||||
|
||||
B += 8;
|
||||
CountN -= 8;
|
||||
}
|
||||
|
||||
//
|
||||
// Process the remaining columns of matrix B.
|
||||
//
|
||||
|
||||
if (CountN > 0) {
|
||||
|
||||
const __m128i ZeroVector = _mm_setzero_si128();
|
||||
|
||||
__m128i ColumnSums[2];
|
||||
uint8_t PaddedMatrixBData[32];
|
||||
|
||||
ColumnSums[0] = _mm_setzero_si128();
|
||||
ColumnSums[1] = _mm_setzero_si128();
|
||||
|
||||
while (CountK > 0) {
|
||||
|
||||
size_t k = std::min(CountK, MLAS_GEMM_U8S8_KERNEL_SSE41::PackedK);
|
||||
CountK -= k;
|
||||
|
||||
_mm_storeu_si128((__m128i*)&PaddedMatrixBData[0], ZeroVector);
|
||||
_mm_storeu_si128((__m128i*)&PaddedMatrixBData[16], ZeroVector);
|
||||
|
||||
uint8_t* padded = PaddedMatrixBData;
|
||||
|
||||
do {
|
||||
|
||||
std::copy_n(B, CountN, padded);
|
||||
|
||||
padded += 8;
|
||||
B += ldb;
|
||||
k -= 1;
|
||||
|
||||
} while (k > 0);
|
||||
|
||||
BytesRows[0] = _mm_loadl_epi64((__m128i*)&PaddedMatrixBData[0]);
|
||||
BytesRows[1] = _mm_loadl_epi64((__m128i*)&PaddedMatrixBData[8]);
|
||||
BytesRows[2] = _mm_loadl_epi64((__m128i*)&PaddedMatrixBData[16]);
|
||||
BytesRows[3] = _mm_loadl_epi64((__m128i*)&PaddedMatrixBData[24]);
|
||||
|
||||
MlasGemmU8X8CopyPackBProcessSse41(D, BytesRows, OnesByteBroadcast, OnesWordBroadcast, ColumnSums);
|
||||
|
||||
D += 32;
|
||||
}
|
||||
|
||||
_mm_storeu_si128((__m128i*)&ColumnSumBuffer[0], ColumnSums[0]);
|
||||
_mm_storeu_si128((__m128i*)&ColumnSumBuffer[4], ColumnSums[1]);
|
||||
}
|
||||
}
|
||||
|
||||
MLAS_FORCEINLINE
|
||||
void
|
||||
MlasGemmU8X8MultiplyAccumulateRowSse41(
|
||||
__m128i ABroadcast,
|
||||
const MLAS_GEMM_U8S8_KERNEL_SSE41::PackedBType* B,
|
||||
__m128i OnesWordBroadcast,
|
||||
__m128i Accumulators[2]
|
||||
)
|
||||
{
|
||||
__m128i BElements0 = _mm_load_si128((__m128i*)&B[0]);
|
||||
__m128i BElements1 = _mm_load_si128((__m128i*)&B[16]);
|
||||
|
||||
__m128i Intermediate0 = _mm_maddubs_epi16(ABroadcast, BElements0);
|
||||
__m128i Intermediate1 = _mm_maddubs_epi16(ABroadcast, BElements1);
|
||||
|
||||
Accumulators[0] = _mm_add_epi32(Accumulators[0], _mm_madd_epi16(Intermediate0, OnesWordBroadcast));
|
||||
Accumulators[1] = _mm_add_epi32(Accumulators[1], _mm_madd_epi16(Intermediate1, OnesWordBroadcast));
|
||||
}
|
||||
|
||||
template<>
|
||||
size_t
|
||||
MlasGemmU8X8Kernel<MLAS_GEMM_U8S8_KERNEL_SSE41>(
|
||||
const MLAS_GEMM_U8S8_KERNEL_SSE41::PackedAType* A,
|
||||
const MLAS_GEMM_U8S8_KERNEL_SSE41::PackedBType* B,
|
||||
int32_t* C,
|
||||
size_t PackedCountK,
|
||||
size_t CountM,
|
||||
size_t CountN,
|
||||
size_t ldc,
|
||||
const int32_t* RowSumBuffer,
|
||||
const int32_t* ColumnSumBuffer,
|
||||
const int32_t* ZeroPointB,
|
||||
bool ZeroMode
|
||||
)
|
||||
{
|
||||
const __m128i OnesWordBroadcast = _mm_set1_epi16(1);
|
||||
|
||||
MLAS_UNREFERENCED_PARAMETER(CountM);
|
||||
MLAS_UNREFERENCED_PARAMETER(ldc);
|
||||
|
||||
while (CountN > 0) {
|
||||
|
||||
__m128i Accumulators[2];
|
||||
|
||||
//
|
||||
// Initialize the accumulators with the row and column sums.
|
||||
//
|
||||
|
||||
Accumulators[0] = _mm_set1_epi32(RowSumBuffer[0]);
|
||||
Accumulators[1] = Accumulators[0];
|
||||
|
||||
if (ZeroPointB != nullptr) {
|
||||
Accumulators[0] = _mm_mullo_epi32(Accumulators[0], _mm_loadu_si128((const __m128i*)&ZeroPointB[0]));
|
||||
Accumulators[1] = _mm_mullo_epi32(Accumulators[1], _mm_loadu_si128((const __m128i*)&ZeroPointB[4]));
|
||||
ZeroPointB += 8;
|
||||
}
|
||||
|
||||
Accumulators[0] = _mm_add_epi32(Accumulators[0], _mm_loadu_si128((const __m128i*)&ColumnSumBuffer[0]));
|
||||
Accumulators[1] = _mm_add_epi32(Accumulators[1], _mm_loadu_si128((const __m128i*)&ColumnSumBuffer[4]));
|
||||
ColumnSumBuffer += 8;
|
||||
|
||||
//
|
||||
// Broadcast each quad of 8-bit values from the matrix A and multiply
|
||||
// with the quad of 8-bit values from matrix B, and add the 32-bit
|
||||
// intermediate into the accumulator registers.
|
||||
//
|
||||
|
||||
const uint8_t* a = A;
|
||||
size_t k = PackedCountK;
|
||||
|
||||
while (k >= 4) {
|
||||
|
||||
__m128i AElements = _mm_loadu_si128((__m128i*)a);
|
||||
__m128i ABroadcast;
|
||||
|
||||
ABroadcast = _mm_shuffle_epi32(AElements, _MM_SHUFFLE(0, 0, 0, 0));
|
||||
MlasGemmU8X8MultiplyAccumulateRowSse41(ABroadcast, &B[0], OnesWordBroadcast, Accumulators);
|
||||
|
||||
ABroadcast = _mm_shuffle_epi32(AElements, _MM_SHUFFLE(1, 1, 1, 1));
|
||||
MlasGemmU8X8MultiplyAccumulateRowSse41(ABroadcast, &B[32], OnesWordBroadcast, Accumulators);
|
||||
|
||||
ABroadcast = _mm_shuffle_epi32(AElements, _MM_SHUFFLE(2, 2, 2, 2));
|
||||
MlasGemmU8X8MultiplyAccumulateRowSse41(ABroadcast, &B[64], OnesWordBroadcast, Accumulators);
|
||||
|
||||
ABroadcast = _mm_shuffle_epi32(AElements, _MM_SHUFFLE(3, 3, 3, 3));
|
||||
MlasGemmU8X8MultiplyAccumulateRowSse41(ABroadcast, &B[96], OnesWordBroadcast, Accumulators);
|
||||
|
||||
a += 4 * 4;
|
||||
B += 4 * 32;
|
||||
k -= 4;
|
||||
}
|
||||
|
||||
while (k > 0) {
|
||||
|
||||
__m128i ABroadcast = _mm_set1_epi32(*((int32_t*)a));
|
||||
MlasGemmU8X8MultiplyAccumulateRowSse41(ABroadcast, &B[0], OnesWordBroadcast, Accumulators);
|
||||
|
||||
a += 4;
|
||||
B += 32;
|
||||
k -= 1;
|
||||
}
|
||||
|
||||
//
|
||||
// Output the accumulator block after optionally accumulating the values
|
||||
// from matrix C.
|
||||
//
|
||||
|
||||
if (CountN >= 8) {
|
||||
|
||||
if (!ZeroMode) {
|
||||
Accumulators[0] = _mm_add_epi32(Accumulators[0], _mm_loadu_si128((__m128i*)&C[0]));
|
||||
Accumulators[1] = _mm_add_epi32(Accumulators[1], _mm_loadu_si128((__m128i*)&C[4]));
|
||||
}
|
||||
|
||||
_mm_storeu_si128((__m128i*)&C[0], Accumulators[0]);
|
||||
_mm_storeu_si128((__m128i*)&C[4], Accumulators[1]);
|
||||
|
||||
C += 8;
|
||||
CountN -= 8;
|
||||
|
||||
} else {
|
||||
|
||||
//
|
||||
// Output the remaining partial output block.
|
||||
//
|
||||
|
||||
if ((CountN & 4) != 0) {
|
||||
|
||||
if (!ZeroMode) {
|
||||
Accumulators[0] = _mm_add_epi32(Accumulators[0], _mm_loadu_si128((__m128i*)&C[0]));
|
||||
}
|
||||
|
||||
_mm_storeu_si128((__m128i*)&C[0], Accumulators[0]);
|
||||
C += 4;
|
||||
|
||||
Accumulators[0] = Accumulators[1];
|
||||
}
|
||||
|
||||
if ((CountN & 2) != 0) {
|
||||
|
||||
if (!ZeroMode) {
|
||||
Accumulators[0] = _mm_add_epi32(Accumulators[0], _mm_loadl_epi64((__m128i*)&C[0]));
|
||||
}
|
||||
|
||||
_mm_storel_epi64((__m128i*)&C[0], Accumulators[0]);
|
||||
C += 2;
|
||||
|
||||
Accumulators[0] = _mm_shuffle_epi32(Accumulators[0], _MM_SHUFFLE(3, 2, 3, 2));
|
||||
}
|
||||
|
||||
if ((CountN & 1) != 0) {
|
||||
|
||||
int32_t AccumulatorValue = _mm_cvtsi128_si32(Accumulators[0]);
|
||||
|
||||
if (!ZeroMode) {
|
||||
AccumulatorValue += C[0];
|
||||
}
|
||||
|
||||
C[0] = AccumulatorValue;
|
||||
}
|
||||
|
||||
CountN = 0;
|
||||
}
|
||||
}
|
||||
|
||||
return 1;
|
||||
}
|
||||
|
||||
const MLAS_GEMM_U8X8_DISPATCH MlasGemmU8S8DispatchSse41 = {
|
||||
MlasGemmU8X8Operation<MLAS_GEMM_U8S8_KERNEL_SSE41>,
|
||||
MlasGemmU8X8PackedOperation<MLAS_GEMM_U8S8_KERNEL_SSE41>,
|
||||
MlasGemmU8X8CopyPackB<MLAS_GEMM_U8S8_KERNEL_SSE41>,
|
||||
MLAS_GEMM_U8S8_KERNEL_SSE41::PackedK,
|
||||
MLAS_GEMM_U8S8_KERNEL_SSE41::PackedStrides.K,
|
||||
};
|
||||
|
||||
#endif
|
||||
|
||||
#if defined(MLAS_TARGET_AMD64)
|
||||
|
||||
//
|
||||
|
|
@ -1334,19 +1768,6 @@ constexpr size_t MLAS_GEMM_U8U8_KERNEL_AVX2::PackedK;
|
|||
constexpr MLAS_GEMM_U8X8_STRIDES MLAS_GEMM_U8U8_KERNEL_AVX2::Strides;
|
||||
constexpr MLAS_GEMM_U8X8_STRIDES MLAS_GEMM_U8U8_KERNEL_AVX2::PackedStrides;
|
||||
|
||||
template<>
|
||||
MLAS_FORCEINLINE
|
||||
int32_t
|
||||
MlasGemmU8X8FixupZeroPointB<MLAS_GEMM_U8U8_KERNEL_AVX2>(
|
||||
int32_t ZeroPointB,
|
||||
bool BIsSigned
|
||||
)
|
||||
{
|
||||
MLAS_UNREFERENCED_PARAMETER(BIsSigned);
|
||||
|
||||
return ZeroPointB;
|
||||
}
|
||||
|
||||
template<>
|
||||
MLAS_FORCEINLINE
|
||||
void
|
||||
|
|
|
|||
|
|
@ -120,7 +120,7 @@ Status QLinearConv::PrePack(const Tensor& tensor, int input_idx, bool& is_packed
|
|||
|
||||
// Don't pack the filter buffer if the MlasConvDepthwise path is used.
|
||||
if (group_input_channels != 1 && group_output_channels != 1) {
|
||||
packed_W_size_ = MlasGemmPackBSize(group_output_channels, kernel_dim, true);
|
||||
packed_W_size_ = MlasGemmPackBSize(group_output_channels, kernel_dim, is_W_signed_);
|
||||
|
||||
if (packed_W_size_ != 0) {
|
||||
auto* packed_W = static_cast<uint8_t*>(alloc->Alloc(SafeInt<size_t>(group_count) * packed_W_size_));
|
||||
|
|
|
|||
|
|
@ -210,7 +210,7 @@ class MlasQgemmU8X8Test<xint8_t, int32_t, Packed, Threaded> : public MlasQgemmU8
|
|||
for (size_t n = 0; n < N; n++, f++) {
|
||||
ASSERT_EQ(C[f], CReference[f]) << "@[" << batch << "x" << m << "x" << n << "], "
|
||||
<< "Batch=" << BatchSize << "M=" << M << ", N=" << N << ", K=" << K
|
||||
<< ", offa=" << int(offa) << ", offb=" << offb;
|
||||
<< ", offa=" << int(offa) << ", offb=" << int(offb);
|
||||
}
|
||||
}
|
||||
}
|
||||
|
|
@ -240,7 +240,7 @@ class MlasQgemmU8X8Test<xint8_t, int32_t, Packed, Threaded> : public MlasQgemmU8
|
|||
for (size_t n = 0; n < N; n++, f++) {
|
||||
ASSERT_EQ(C[f], CReference[f]) << "@[" << batch << "x" << m << "x" << n << "], "
|
||||
<< "Batch=" << BatchSize << "M=" << M << ", N=" << N << ", K=" << K
|
||||
<< ", offa=" << int(offa) << ", offb=" << offb;
|
||||
<< ", offa=" << int(offa) << ", offb=--";
|
||||
}
|
||||
}
|
||||
}
|
||||
|
|
@ -444,7 +444,7 @@ class MlasQgemmU8X8Test<xint8_t, float, Packed, Threaded> : public MlasQgemmU8X8
|
|||
MlasGemm(CblasNoTrans, CblasNoTrans, M, N, K, 1.0f,
|
||||
AFloat + K * M * b, lda,
|
||||
BFloat + N * K * b, ldb, 0.0f,
|
||||
CReference + N * M * b, ldc,
|
||||
CReference + N * M * b, ldc,
|
||||
MlasQgemmU8X8U8X8TestBase<Packed, Threaded>::threadpool_);
|
||||
}
|
||||
|
||||
|
|
|
|||
Loading…
Reference in a new issue