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

Add compiler support for offloading matmul ops to custom dispatch plugin. #71

Closed
wants to merge 43 commits into from

Conversation

monorimet
Copy link
Collaborator

Add lit test and fix build.

Fixes to LowerToAccelUKernelPass

Tweaks to LowerToAccelUKernelsPass.

Add AccelMatmulExpert pass pipeline

Apply clang-format to new C++ files. (#3)

  • Apply clangformat.

Use parameter struct calling convention

Tweaks to KernelDispatch and AccelMatmulExpert pipeline

@MaheshRavishankar
Copy link

Probably needs a rebase.

@monorimet
Copy link
Collaborator Author

Auto-rebase 🫠

Copy link

@MaheshRavishankar MaheshRavishankar left a comment

Choose a reason for hiding this comment

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

This looks mostly fine. Have a few comments (also marked draft). Would be good to add an e2e test with the flag enabled to check for results on a matmul calling the plugin.

linalg::MatmulOp matmulOp) {
assert(!getLoweringConfig(matmulOp) && "expected lowering_config is not set");
SmallVector<int64_t> tileSizes;
tileSizes.push_back(1);

Choose a reason for hiding this comment

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

Why do you need {1} here? Are you actually tiling it by 1?

auto targetAttr = IREE::HAL::ExecutableTargetAttr::lookup(op);
auto fn = getFnNameAndDefAttrs("accel_matmul_f32", rewriter, targetAttr);
auto genericMicroKernelOp = rewriter.create<IREE::Codegen::UKernelGenericOp>(
loc, outType, fn.name, ValueRange{lhs, rhs}, out, ValueRange{m, n, k},

Choose a reason for hiding this comment

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

Are you sure the plugin is taking {m, n, k} last I saw it was {m, k, n}. Just need to be consistent.

IREE::HAL::ExecutableTargetAttr targetAttr) {
FnNameAndDefAttrs result;
result.name = ukernelName;
result.defAttrs.emplace_back(

Choose a reason for hiding this comment

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

Only add these if the plugin needs it. If it doesnt, just lesser things to deal with.

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.