Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Jump to bottom

Matmul nbits to optimize memory layout for avx instructions #22203

Draft
wants to merge 14 commits into
base: main
Choose a base branch
from
Draft

Matmul nbits to optimize memory layout for avx instructions #22203

wants to merge 14 commits into from

Conversation

liqunfu
Copy link
Contributor

@liqunfu liqunfu commented Sep 24, 2024

Description

Motivation and Context

@liqunfu
matmul nbits to optimize memory layout for avx instructions
555e951 
Signed-off-by: liqunfu <liqun.fu@microsoft.com>
@liqunfu liqunfu requested a review from a team as a code owner September 24, 2024 15:50
@liqunfu liqunfu marked this pull request as draft September 24, 2024 15:50
@liqunfu
Merge branch 'main' into liqun/avx-layout
076998c
@liqunfu
intermediate push
99aec95 
Signed-off-by: liqunfu <liqun.fu@microsoft.com>
github-actions[bot]
github-actions bot reviewed Nov 18, 2024
View reviewed changes
Copy link
Contributor

@github-actions github-actions bot left a comment

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

You can commit the suggested changes from lintrunner.

onnxruntime/test/contrib_ops/matmul_4bits_test.cc Outdated
Comment on lines 367 to 394
TEST(MatMulNBits, Float32_Accuracy4) {
TestMatMulNBitsTyped<float, 1, 1, 16, 16, 4>();
TestMatMulNBitsTyped<float, 1, 2, 16, 16, 4>();
TestMatMulNBitsTyped<float, 1, 32, 16, 16, 4>();
TestMatMulNBitsTyped<float, 1, 32, 32, 16, 4>();
TestMatMulNBitsTyped<float, 1, 32, 16, 128, 4>();
TestMatMulNBitsTyped<float, 1, 288, 16, 16, 4>();
TestMatMulNBitsTyped<float, 1, 288, 1024, 16, 4>();
TestMatMulNBitsTyped<float, 1, 288, 1024, 128, 4>();
TestMatMulNBitsTyped<float, 1, 288, 93, 32, 4>();
TestMatMulNBitsTyped<float, 1, 288, 93, 128, 4>();
TestMatMulNBitsTyped<float, 1, 288, 1234, 16, 4>();
TestMatMulNBitsTyped<float, 2, 1, 16, 16, 4>();
TestMatMulNBitsTyped<float, 2, 2, 16, 16, 4>();
TestMatMulNBitsTyped<float, 100, 1, 16, 16, 4>();
TestMatMulNBitsTyped<float, 100, 2, 16, 16, 4>();
TestMatMulNBitsTyped<float, 100, 32, 16, 16, 4>();
TestMatMulNBitsTyped<float, 100, 32, 32, 16, 4>();
TestMatMulNBitsTyped<float, 100, 32, 16, 128, 4>();
TestMatMulNBitsTyped<float, 100, 288, 16, 16, 4>();
TestMatMulNBitsTyped<float, 100, 288, 1024, 16, 4>();
TestMatMulNBitsTyped<float, 100, 288, 1024, 128, 4>();
TestMatMulNBitsTyped<float, 100, 288, 93, 32, 4>();
TestMatMulNBitsTyped<float, 100, 288, 93, 128, 4>();
TestMatMulNBitsTyped<float, 100, 288, 1234, 16, 4>();
//TestMatMulNBitsTyped<float, 1, 1, 16, 16, 4>();
//TestMatMulNBitsTyped<float, 1, 2, 16, 16, 4>();
//TestMatMulNBitsTyped<float, 1, 32, 16, 16, 4>();
//TestMatMulNBitsTyped<float, 1, 32, 32, 16, 4>();
//TestMatMulNBitsTyped<float, 1, 32, 16, 128, 4>();
//TestMatMulNBitsTyped<float, 1, 288, 16, 16, 4>();
//TestMatMulNBitsTyped<float, 1, 288, 1024, 16, 4>();
//TestMatMulNBitsTyped<float, 1, 288, 1024, 128, 4>();
//TestMatMulNBitsTyped<float, 1, 288, 93, 32, 4>();
//TestMatMulNBitsTyped<float, 1, 288, 93, 128, 4>();
//TestMatMulNBitsTyped<float, 1, 288, 1234, 16, 4>();
//TestMatMulNBitsTyped<float, 2, 1, 16, 16, 4>();
//TestMatMulNBitsTyped<float, 2, 2, 16, 16, 4>();
//TestMatMulNBitsTyped<float, 100, 1, 16, 16, 4>();
//TestMatMulNBitsTyped<float, 100, 2, 16, 16, 4>();
//TestMatMulNBitsTyped<float, 100, 32, 16, 16, 4>();
//TestMatMulNBitsTyped<float, 100, 32, 32, 16, 4>();
//TestMatMulNBitsTyped<float, 100, 32, 16, 128, 4>();
//TestMatMulNBitsTyped<float, 100, 288, 16, 16, 4>();
//TestMatMulNBitsTyped<float, 100, 288, 1024, 16, 4>();
//TestMatMulNBitsTyped<float, 100, 288, 1024, 128, 4>();
//TestMatMulNBitsTyped<float, 100, 288, 93, 32, 4>();
//TestMatMulNBitsTyped<float, 100, 288, 93, 128, 4>();
//TestMatMulNBitsTyped<float, 100, 288, 1234, 16, 4>();
TestMatMulNBitsTyped<float, 2, 4, 128, 32, 4>();
//TestMatMulNBitsTyped<float, 100, 288, 1234, 32, 4>();
}
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Suggested change
TEST(MatMulNBits, Float32_Accuracy4) {
TestMatMulNBitsTyped<float, 1, 1, 16, 16, 4>();
TestMatMulNBitsTyped<float, 1, 2, 16, 16, 4>();
TestMatMulNBitsTyped<float, 1, 32, 16, 16, 4>();
TestMatMulNBitsTyped<float, 1, 32, 32, 16, 4>();
TestMatMulNBitsTyped<float, 1, 32, 16, 128, 4>();
TestMatMulNBitsTyped<float, 1, 288, 16, 16, 4>();
TestMatMulNBitsTyped<float, 1, 288, 1024, 16, 4>();
TestMatMulNBitsTyped<float, 1, 288, 1024, 128, 4>();
TestMatMulNBitsTyped<float, 1, 288, 93, 32, 4>();
TestMatMulNBitsTyped<float, 1, 288, 93, 128, 4>();
TestMatMulNBitsTyped<float, 1, 288, 1234, 16, 4>();
TestMatMulNBitsTyped<float, 2, 1, 16, 16, 4>();
TestMatMulNBitsTyped<float, 2, 2, 16, 16, 4>();
TestMatMulNBitsTyped<float, 100, 1, 16, 16, 4>();
TestMatMulNBitsTyped<float, 100, 2, 16, 16, 4>();
TestMatMulNBitsTyped<float, 100, 32, 16, 16, 4>();
TestMatMulNBitsTyped<float, 100, 32, 32, 16, 4>();
TestMatMulNBitsTyped<float, 100, 32, 16, 128, 4>();
TestMatMulNBitsTyped<float, 100, 288, 16, 16, 4>();
TestMatMulNBitsTyped<float, 100, 288, 1024, 16, 4>();
TestMatMulNBitsTyped<float, 100, 288, 1024, 128, 4>();
TestMatMulNBitsTyped<float, 100, 288, 93, 32, 4>();
TestMatMulNBitsTyped<float, 100, 288, 93, 128, 4>();
TestMatMulNBitsTyped<float, 100, 288, 1234, 16, 4>();
//TestMatMulNBitsTyped<float, 1, 1, 16, 16, 4>();
//TestMatMulNBitsTyped<float, 1, 2, 16, 16, 4>();
//TestMatMulNBitsTyped<float, 1, 32, 16, 16, 4>();
//TestMatMulNBitsTyped<float, 1, 32, 32, 16, 4>();
//TestMatMulNBitsTyped<float, 1, 32, 16, 128, 4>();
//TestMatMulNBitsTyped<float, 1, 288, 16, 16, 4>();
//TestMatMulNBitsTyped<float, 1, 288, 1024, 16, 4>();
//TestMatMulNBitsTyped<float, 1, 288, 1024, 128, 4>();
//TestMatMulNBitsTyped<float, 1, 288, 93, 32, 4>();
//TestMatMulNBitsTyped<float, 1, 288, 93, 128, 4>();
//TestMatMulNBitsTyped<float, 1, 288, 1234, 16, 4>();
//TestMatMulNBitsTyped<float, 2, 1, 16, 16, 4>();
//TestMatMulNBitsTyped<float, 2, 2, 16, 16, 4>();
//TestMatMulNBitsTyped<float, 100, 1, 16, 16, 4>();
//TestMatMulNBitsTyped<float, 100, 2, 16, 16, 4>();
//TestMatMulNBitsTyped<float, 100, 32, 16, 16, 4>();
//TestMatMulNBitsTyped<float, 100, 32, 32, 16, 4>();
//TestMatMulNBitsTyped<float, 100, 32, 16, 128, 4>();
//TestMatMulNBitsTyped<float, 100, 288, 16, 16, 4>();
//TestMatMulNBitsTyped<float, 100, 288, 1024, 16, 4>();
//TestMatMulNBitsTyped<float, 100, 288, 1024, 128, 4>();
//TestMatMulNBitsTyped<float, 100, 288, 93, 32, 4>();
//TestMatMulNBitsTyped<float, 100, 288, 93, 128, 4>();
//TestMatMulNBitsTyped<float, 100, 288, 1234, 16, 4>();
TestMatMulNBitsTyped<float, 2, 4, 128, 32, 4>();
//TestMatMulNBitsTyped<float, 100, 288, 1234, 32, 4>();
}
TEST(MatMulNBits, Float32_Accuracy4) {
// TestMatMulNBitsTyped<float, 1, 1, 16, 16, 4>();
// TestMatMulNBitsTyped<float, 1, 2, 16, 16, 4>();
// TestMatMulNBitsTyped<float, 1, 32, 16, 16, 4>();
// TestMatMulNBitsTyped<float, 1, 32, 32, 16, 4>();
// TestMatMulNBitsTyped<float, 1, 32, 16, 128, 4>();
// TestMatMulNBitsTyped<float, 1, 288, 16, 16, 4>();
// TestMatMulNBitsTyped<float, 1, 288, 1024, 16, 4>();
// TestMatMulNBitsTyped<float, 1, 288, 1024, 128, 4>();
// TestMatMulNBitsTyped<float, 1, 288, 93, 32, 4>();
// TestMatMulNBitsTyped<float, 1, 288, 93, 128, 4>();
// TestMatMulNBitsTyped<float, 1, 288, 1234, 16, 4>();
// TestMatMulNBitsTyped<float, 2, 1, 16, 16, 4>();
// TestMatMulNBitsTyped<float, 2, 2, 16, 16, 4>();
// TestMatMulNBitsTyped<float, 100, 1, 16, 16, 4>();
// TestMatMulNBitsTyped<float, 100, 2, 16, 16, 4>();
// TestMatMulNBitsTyped<float, 100, 32, 16, 16, 4>();
// TestMatMulNBitsTyped<float, 100, 32, 32, 16, 4>();
// TestMatMulNBitsTyped<float, 100, 32, 16, 128, 4>();
// TestMatMulNBitsTyped<float, 100, 288, 16, 16, 4>();
// TestMatMulNBitsTyped<float, 100, 288, 1024, 16, 4>();
// TestMatMulNBitsTyped<float, 100, 288, 1024, 128, 4>();
// TestMatMulNBitsTyped<float, 100, 288, 93, 32, 4>();
// TestMatMulNBitsTyped<float, 100, 288, 93, 128, 4>();
// TestMatMulNBitsTyped<float, 100, 288, 1234, 16, 4>();
TestMatMulNBitsTyped<float, 2, 4, 128, 32, 4>();
// TestMatMulNBitsTyped<float, 100, 288, 1234, 32, 4>();
}

github-advanced-security[bot]
github-advanced-security bot found potential problems Nov 27, 2024
View reviewed changes
onnxruntime/core/mlas/lib/sqnbitgemm_kernel_avx512.cpp
float* AScaledBlkSum // scale_k * Sum_blklen(a_i)
)
{
const size_t BlkLen = 32;

Check warning

Code scanning / PREfast

The const variable 'BlkLen' can be computed at compile-time. Consider using constexpr (con.5). Warning

The const variable 'BlkLen' can be computed at compile-time. Consider using constexpr (con.5).
onnxruntime/core/mlas/lib/sqnbitgemm_kernel_avx512.cpp
)
{
const size_t BlkLen = 32;
const int64_t SubBlkLen = 4 * BlkLen; // process 128 weights at a time and then process the remaining weights

Check warning

Code scanning / PREfast

The const variable 'SubBlkLen' can be computed at compile-time. Consider using constexpr (con.5). Warning

The const variable 'SubBlkLen' can be computed at compile-time. Consider using constexpr (con.5).
github-advanced-security[bot]
github-advanced-security bot found potential problems Nov 27, 2024
View reviewed changes
onnxruntime/core/mlas/lib/sqnbitgemm_kernel_avx512.cpp

// Convert int32 to int8
i_16_epi8[index] = _mm512_cvtepi32_epi8(i0);
//_mm_storeu_si128(dst++, i0_8);

Check notice

Code scanning / CodeQL

Commented-out code

This comment appears to contain commented-out code.
onnxruntime/core/mlas/lib/sqnbitgemm_kernel_avx512.cpp
}

while (k_remaining > 0) {
// for (size_t k = 0; k < CountK; k += BlkLen) {

Check notice

Code scanning / CodeQL

Commented-out code

This comment appears to contain commented-out code.
onnxruntime/test/contrib_ops/matmul_4bits_test.cc
}

TEST(MatMulNBits, LongTestFloat32) {
// onnxruntime::profiling::Profiler::Profiler::Instance().StartProfiling<char>("profile.json");

Check notice

Code scanning / CodeQL

Commented-out code

This comment appears to contain commented-out code.
@liqunfu
pass avx512/vnni-blklen32
d371c59 
Signed-off-by: liqunfu <liqun.fu@microsoft.com>
@liqunfu
pass avx512vnni-blklen128. plan to compute blksum in different loop t…
790b03f 
…hus not to implement avx512

Signed-off-by: liqunfu <liqun.fu@microsoft.com>
@liqunfu
attmpt to make blklen256 work. failed because blksum computation need…
557fbb0 
… to be in a separate loop. defer this work later

Signed-off-by: liqunfu <liqun.fu@microsoft.com>
@liqunfu
avx512 blklen64 to compute blksum in a separate loop
6b28657 
Signed-off-by: liqunfu <liqun.fu@microsoft.com>
@liqunfu
avx512 scaled_zp compute in a separate loop except blklen16
0b867f8 
Signed-off-by: liqunfu <liqun.fu@microsoft.com>
@liqunfu
avx512, all blklens, scaled_zp compute in a separate loop
2e74f56 
Signed-off-by: liqunfu <liqun.fu@microsoft.com>
@liqunfu
Merge branch 'main' into liqun/avx-layout
0bf47f7
github-advanced-security[bot]
github-advanced-security bot found potential problems Dec 12, 2024
View reviewed changes
onnxruntime/core/mlas/lib/sqnbitgemm_kernel_avx512_int8_blklen32.h
static MLAS_FORCEINLINE
__m512 load_1blksum_512(const float* BlksumPtr) {
// Create a mask to set only the lowest element
const __mmask16 mask = 0x01; // Binary: 0000 0000 0000 0001

Check warning

Code scanning / PREfast

The const variable 'mask' can be computed at compile-time. Consider using constexpr (con.5). Warning

The const variable 'mask' can be computed at compile-time. Consider using constexpr (con.5).
onnxruntime/core/mlas/lib/sqnbitgemm_kernel_avx512_int8_blklen32.h
// Function to load a single float value into the lowest element of a __m256 register
static MLAS_FORCEINLINE
__m256 load_1blksum_256(const float* BlksumPtr) {
const __mmask8 mask = 0x01; // Binary: 0000 0001

Check warning

Code scanning / PREfast

The const variable 'mask' can be computed at compile-time. Consider using constexpr (con.5). Warning

The const variable 'mask' can be computed at compile-time. Consider using constexpr (con.5).
onnxruntime/core/mlas/lib/sqnbitgemm_kernel_avx512_int8_blklen64.h
@@ -32,6 +32,42 @@
return result;
}

static MLAS_FORCEINLINE
__m512 load_broadcast_512(const float& combined_scale) {
const __mmask16 mask = 00000001; // Binary: 0000 0000 0000 0001, lowest element

Check warning

Code scanning / PREfast

The const variable 'mask' can be computed at compile-time. Consider using constexpr (con.5). Warning

The const variable 'mask' can be computed at compile-time. Consider using constexpr (con.5).
onnxruntime/core/mlas/lib/sqnbitgemm_kernel_avx512_int8_blklen64.h
// const __m512 scale_b_16_ps = _mm512_broadcast_f32x8(scale_b_ps);
// return;

const __mmask8 mask = 0xff; // Binary: 1111 1111, to set all elements

Check warning

Code scanning / PREfast

The const variable 'mask' can be computed at compile-time. Consider using constexpr (con.5). Warning

The const variable 'mask' can be computed at compile-time. Consider using constexpr (con.5).
liqunfu and others added 2 commits December 13, 2024 10:09
@liqunfu
avx2 passes
c19ae9e 
Signed-off-by: Liqun Fu <liqun.fu@microsoft.com>
@liqunfu
avxvnni, matmul_nbit kernel
b26b075 
Signed-off-by: Liqun Fu <liqun_fu@hotmail.com>
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment
Reviewers

@github-actions github-actions[bot] github-actions[bot] left review comments

At least 1 approving review is required to merge this pull request.

Assignees
No one assigned
Labels
None yet
Projects
None yet
Milestone
No milestone
Development

Successfully merging this pull request may close these issues.
1 participant
@liqunfu