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.

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

brgemm: support arbitrary K on AMX #2319

Open
wants to merge 5 commits into
base: main
Choose a base branch
from

Conversation

ankalinin
Copy link
Contributor

@ankalinin ankalinin commented Dec 27, 2024

There are two problems regarding brgemm K value on AMX:

brgemm doesn't support K not divisible by vnni granularity
for K not divisible by tile width (32 or 64) the blocking by K dimension may be not optimal.
To get around this limitation brgemm primitives are forced to either transform the matrix A or call many small brgemm kernels with different tile configurations. Both ways leads to performance lost.
This PR implements support of arbitrary K in brgemm on AMX and implements corresponding updates in 1x1 convolutions and in matmul to use this new ability.

Performance update for convolutions:
image

Performance update for matmul
image

@ankalinin ankalinin requested review from a team as code owners December 27, 2024 00:25
@github-actions github-actions bot added platform:cpu-x64 Intel64/AMD64 processors. Codeowner: @oneapi-src/onednn-cpu-x64 component:tests Codeowner: @oneapi-src/onednn-arch labels Dec 27, 2024
@ankalinin ankalinin removed the component:tests Codeowner: @oneapi-src/onednn-arch label Dec 27, 2024
@ankalinin
Copy link
Contributor Author

make test
disable device_gpu

@ankalinin ankalinin force-pushed the akalinin/brgemm_k_tail_pr branch from 527c2a0 to 4bd35fa Compare January 2, 2025 20:55
@ankalinin ankalinin requested a review from a team as a code owner January 2, 2025 20:55
@github-actions github-actions bot added platform:cpu-aarch64 Codeowner: @oneapi-src/onednn-cpu-aarch64 component:tests Codeowner: @oneapi-src/onednn-arch labels Jan 2, 2025
@ankalinin ankalinin force-pushed the akalinin/brgemm_k_tail_pr branch from 4bd35fa to 280499f Compare January 2, 2025 22:15
@theComputeKid
Copy link
Member

I restarted the AArch64 CI because the encountered failure is a known sporadic bug on the c7g.

@theComputeKid
Copy link
Member

cc: @Radu2k

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment
Labels
component:tests Codeowner: @oneapi-src/onednn-arch platform:cpu-aarch64 Codeowner: @oneapi-src/onednn-cpu-aarch64 platform:cpu-x64 Intel64/AMD64 processors. Codeowner: @oneapi-src/onednn-cpu-x64
Projects
None yet
Development

Successfully merging this pull request may close these issues.

4 participants