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

[KernelDispatch] Refactor the pass for generic ops and utils for matmul_transpose_b ops #991

Merged
merged 3 commits into from
Dec 17, 2024

Conversation

yzhang93
Copy link
Contributor

This PR refactored and simplified the root config setting in KernelDispatch.cpp for the following aspects:

  • Removed setTransposeLikeOpRootConfig and flag isMatmulTranspseB. Instead, directly go through setRootConfig for generic version as well as the named op version of matmul_transpse_b. This will be easier when we extend the support to cover matmul_transpose_a.
  • Moved the isMatmulTransposeB utils to Utils.cpp and reuse/merge with the existing utils for matmul ops.
  • Added missing tests for generic version of matmul and matmul_transpose_b.

Copy link
Collaborator

@jtuyls jtuyls left a comment

Choose a reason for hiding this comment

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

LGTM, just a few nits

Comment on lines 110 to 112
if (!bodyMatcherForMatmulLikeOps(yieldVal, body)) {
return false;
}
Copy link
Collaborator

Choose a reason for hiding this comment

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

Suggested change
if (!bodyMatcherForMatmulLikeOps(yieldVal, body)) {
return false;
}
if (!bodyMatcherForMatmulLikeOps(yieldVal, body)) return false;

@@ -215,6 +234,22 @@ bool isMatmul(linalg::LinalgOp linalgOp) {
match6DLinalgGenericMatmul(linalgOp);
}

/// Utility to indentify whether a linalg op is a matmul_transpose_b op.
Copy link
Collaborator

Choose a reason for hiding this comment

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

Suggested change
/// Utility to indentify whether a linalg op is a matmul_transpose_b op.
/// Utility to identify whether a linalg op is a matmul_transpose_b op.

@@ -48,6 +48,9 @@ FailureOr<unsigned> getTilingScaleFactor(Type elemType);
/// Utility to indentify whether a linalg op is a matmul op.
bool isMatmul(linalg::LinalgOp linalgOp);

/// Utility to indentify whether a linalg op is a matmul_transpose_b op.
Copy link
Collaborator

Choose a reason for hiding this comment

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

Suggested change
/// Utility to indentify whether a linalg op is a matmul_transpose_b op.
/// Utility to identify whether a linalg op is a matmul_transpose_b op.

@jtuyls jtuyls changed the title [KernelDispatch] Refactor the pass for generic ops and utils for mamtul_transpose_b ops [KernelDispatch] Refactor the pass for generic ops and utils for matmul_transpose_b ops Dec 17, 2024
@yzhang93 yzhang93 force-pushed the refactor_kernel_dispatch branch from 24e0320 to 10a33ad Compare December 17, 2024 21:10
@yzhang93 yzhang93 enabled auto-merge (squash) December 17, 2024 21:33
@yzhang93 yzhang93 merged commit d24568f into nod-ai:main Dec 17, 2024
7 checks passed
@yzhang93 yzhang93 deleted the refactor_kernel_dispatch branch January 3, 2025 05:53
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