Skip to content

Support for tt.dot_scaled operator #2804

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
Merged
merged 13 commits into from
Nov 29, 2024
Merged

Support for tt.dot_scaled operator #2804

merged 13 commits into from
Nov 29, 2024

Conversation

@etiotto
Copy link
Contributor

@etiotto etiotto commented Nov 22, 2024 •
edited
Loading

This PR decomposed a tt.dot_scaled operation into a tt.dot operation where one of the operands (e.g A) is scaled using the triton_gpu_upcast_mxfp operation.

Note: The upcast_mxfp operation is not lowered to LLVM IR in this PR.

@etiotto etiotto self-assigned this Nov 22, 2024
@etiotto etiotto linked an issue Nov 22, 2024 that may be closed by this pull request
Implement support for the tt.dot_scaled operation on XPU #2633
Closed
@etiotto etiotto marked this pull request as ready for review November 25, 2024 14:25
@anmyachev
Copy link
Contributor

anmyachev commented Nov 25, 2024

@etiotto in order to test these changes we need to unskip test_scaled_dot:

intel-xpu-backend-for-triton/python/test/unit/language/test_core.py

Lines 3437 to 3438 in d804502

if is_xpu():
pytest.skip("scaled_dot isn't supported on XPU")

Did you plan to do this here?

@etiotto
Copy link
Contributor Author

etiotto commented Nov 26, 2024 •
edited
Loading

@etiotto in order to test these changes we need to unskip test_scaled_dot:

intel-xpu-backend-for-triton/python/test/unit/language/test_core.py

Lines 3437 to 3438 in d804502

if is_xpu():
pytest.skip("scaled_dot isn't supported on XPU")

Did you plan to do this here?

Yes we need to do that, however the code is not yet fully functional because the lowering code for the triton_gpu.upcast_mxfp operation is not working yet. I will remove that part of the PR and just deal with decomposing the dot_scaled operation in this particular PR.

@etiotto etiotto marked this pull request as draft November 26, 2024 00:38
@etiotto
WIP: tt.scaled_dot
c230c6f 
Signed-off-by: Tiotto, Ettore <[email protected]>
@etiotto
WIP: tt.scaled_dot
5772d3f 
Signed-off-by: Tiotto, Ettore <[email protected]>
@etiotto
Fix precommit
bbf250d 
Signed-off-by: Tiotto, Ettore <[email protected]>
@etiotto etiotto marked this pull request as ready for review November 26, 2024 21:15
whitneywhtsang
whitneywhtsang reviewed Nov 27, 2024
View reviewed changes
@victor-eds
Copy link
Contributor

victor-eds commented Nov 28, 2024

PR approach LGTM. Just some NITs.

@whitneywhtsang
Copy link
Contributor

whitneywhtsang commented Nov 28, 2024

PR approach LGTM. Just some NITs.

@victor-eds Maybe you forgot to submit the NITs?

leonling-ll
leonling-ll approved these changes Nov 28, 2024
View reviewed changes
Copy link
Contributor

@leonling-ll leonling-ll left a comment

Choose a reason for hiding this comment

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

LGTM.

victor-eds
victor-eds reviewed Nov 28, 2024
View reviewed changes
lib/Dialect/TritonGPU/IR/Ops.cpp
newShape[kIdx] *= 2;
retTy = RankedTensorType::get(newShape, FloatType::getBF16(ctx),
newVEncoding);
Type elemType = FloatType::getBF16(ctx);
Copy link
Contributor

@victor-eds victor-eds Nov 28, 2024

Choose a reason for hiding this comment

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

Can we define this inside the if statement below?

Copy link
Contributor Author

@etiotto etiotto Nov 29, 2024

Choose a reason for hiding this comment

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

I rather have it here because elemType is used after the if/else at line 147

third_party/intel/include/Dialect/TritonIntelGPU/IR/Attributes.h
Comment on lines +6 to +8
namespace mlir {
class ModuleOp;
}
Copy link
Contributor

@victor-eds victor-eds Nov 28, 2024

Choose a reason for hiding this comment

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

I'd bet we don't need this

Copy link
Contributor Author

@etiotto etiotto Nov 29, 2024

Choose a reason for hiding this comment

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

ModuleOp is used in "intel/include/Dialect/TritonIntelGPU/IR/TritonIntelGPUAttrDefs.h.inc" now:

static DPASCapability getDPASCapability(mlir::ModuleOp mod);

That is the reason I have put the forward declaration here.

@etiotto
Address code review comments
bd1094c 
Signed-off-by: Tiotto, Ettore <[email protected]>
whitneywhtsang
whitneywhtsang reviewed Nov 29, 2024
View reviewed changes
@etiotto etiotto enabled auto-merge (squash) November 29, 2024 17:18
@etiotto etiotto merged commit 0c70ca3 into main Nov 29, 2024
5 checks passed
@etiotto etiotto deleted the etiotto.add_support_for_scaled_dot branch November 29, 2024 17:48
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Reviewers

@whitneywhtsang whitneywhtsang whitneywhtsang left review comments

@chengjunlu chengjunlu Awaiting requested review from chengjunlu

@victor-eds victor-eds Awaiting requested review from victor-eds

+1 more reviewer

@leonling-ll leonling-ll leonling-ll approved these changes

Reviewers whose approvals may not affect merge requirements

Assignees

@etiotto etiotto

Labels

None yet

Projects

None yet

Milestone

No milestone

Development

Successfully merging this pull request may close these issues.

Implement support for the tt.dot_scaled operation on XPU

6 participants

@etiotto @anmyachev @victor-eds @whitneywhtsang @leonling-ll