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

Support matmul 4d inputs on pack-peel-4-level-tiling #1098

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

Conversation

yzhang93
Copy link
Contributor

@yzhang93 yzhang93 commented Feb 12, 2025

This PR adds basic support for matmul with 4d inputs and output on pack-peel-4-level-tiling.

In theory, the tiling strategy and pipeline should support different layouts of matmul 4d operations. However, the order of the input dims (both inner and outer dims) are crucial for the correct compilation and results.

To ensure the correctness and for comparison purpose, this PR only adds a test that corresponds to the standard matmul, which means the order of the input dims for this operation corresponds to the L2 shapes of the matmul op after the first level packing, i.e.,

C += matmul4d(A,B) where A:MxKxM0xK0, B:NxKxK0xN0, C:NxMxM0xN0

The test class and instance added in run.py is preliminary and for experimental purpose. Generalization of the test class will be addressed as follow-ups.

Runtime comparison on Phoenix CI:
matmul_512_4096_512_bf16_f32 : 1141us vs matmul4d_16_128_8_bf16_f32: 998us

@yzhang93 yzhang93 force-pushed the support_mmt4d_on_4level_tiling branch from 11d610c to 11ed3b9 Compare February 12, 2025 03:59
Copy link
Contributor

@Abhishek-Varma Abhishek-Varma left a comment

Choose a reason for hiding this comment

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

This is nice @yzhang93 !

But I think it'd be better to have a separate pipeline for targeting MMT4D - a notion similar to what IREE LLVMCPU path takes.

I'm not too sure though if having an attribute that relays "number of loops" through an operation only to be used by the pipeline to decide which all set of passes to include is a nice design.

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment
Labels
None yet
Projects
None yet
Development

Successfully merging this pull request may close these issues.

2 participants