[BACKEND] Small fixes for dot operand properties by Jokeren · Pull Request #4895 · triton-lang/triton

Jokeren · 2024-10-11T20:25:54Z

Co-authored-by: Mario Lezcano Casado lezcano@openai.com

Jokeren · 2024-10-11T20:28:05Z

@lezcano I haven't committed all fixes yet since I noticed some differences between my PR and yours. Most changes in this PR are probably consistent except for getContigPerThread

zhanglx13 · 2024-10-12T16:24:08Z

+    std::iota(order.rbegin(), order.rend(), 0);
+    if (dotOpLayout.getOpIdx() == 0) {
+      std::swap(order[0], order[1]);
+    }


This will result in
opIdx=0: [1, 2, 0]
opIdx=1: [2, 1, 0]

And I assume batch is dim0, so maybe you want the following?
opIdx=0: [1, 0, 2]
opIdx=1: [0, 1, 2]

Yep. What I was trying to say is order[rank - 2] and order[rank - 1]. Thanks for the review.

Because this is swapping the order array, I think the above code is correct.
order[0] refers to the leading dimension.
For opIdx=0, [/*dim0*/batch, /*dim1=*/m, /*dim2=*/k], the leading dimension should be dim1=m
For opIdx=1, [/*dim0*/batch, /*dim1=*/k, /*dim2=*/n], the leading dimension should be dim2=n

lezcano

Thank you for the clean-up!

Here's a list of other bugs I found. Feel free to bundle the fixes into this PR, or we can land another round of fixes after my PR:

getShapePerCTATileForDotOperands should be ForOperand. Should take a kWidth as argument, and should return kWidth * 2 * 4 in the K dimension (or equiv. for AMD layouts)
getWarpsPerCTA should clamp the K dimension to 1 as per

auto warps = distributedLayout.getWarpsPerCTA();
auto kDim = getOpIdx() == 0 ? 1 : 0;
warps[kDim] = 1;

All the ForOperand ops should be moved to DotOperandLayoutAttr class, to be able to call class-dependent ops, like the getWarpsPerCTA defined above (this bit me when getMMAv2Rep calls getWarpsPerCTA, and in some other place as well).
getThreadOrder should be modified the same way as getWarpOrder

lezcano · 2024-10-14T09:32:33Z

-      std::iota(order.rbegin(), order.rend(), 0);
-    }
-    return order;
+    return getOrderForDotOperand(dotLayout.getOpIdx(), rank);


This is change is correct for consistency, but will conflict with #4891. We will have to fix there (and perhaps the LL lowering) if this one lands after the LL one.

Jokeren · 2024-10-15T04:32:39Z

@lezcano ready for another round of review. Also added you as a co-author

lezcano

LGTM. We are missing fixing getWarpsPerCTA, but we can fix that one in a different PR, as it may affect the way we lower MMA ops.

@Jokeren

This PR includes #4891 and #4895. I will rebase once those have landed. It includes a number of hacks to work around bugs in `DotOperandEncodingAttr`. All these are marked as `FIXME [Dot LL]` to be easy to grep for. @Jokeren is working on a comprehensive revamp of `DotOperandEncodingAttr` which will get rid of all these. #4895 is the first step in this direction.

@Jokeren

This PR includes triton-lang#4891 and triton-lang#4895. I will rebase once those have landed. It includes a number of hacks to work around bugs in `DotOperandEncodingAttr`. All these are marked as `FIXME [Dot LL]` to be easy to grep for. @Jokeren is working on a comprehensive revamp of `DotOperandEncodingAttr` which will get rid of all these. triton-lang#4895 is the first step in this direction.

Co-authored-by: Mario Lezcano Casado <lezcano@openai.com>

@Jokeren

This PR includes triton-lang#4891 and triton-lang#4895. I will rebase once those have landed. It includes a number of hacks to work around bugs in `DotOperandEncodingAttr`. All these are marked as `FIXME [Dot LL]` to be easy to grep for. @Jokeren is working on a comprehensive revamp of `DotOperandEncodingAttr` which will get rid of all these. triton-lang#4895 is the first step in this direction.

Co-authored-by: Mario Lezcano Casado <lezcano@openai.com>

@Jokeren

This PR includes triton-lang#4891 and triton-lang#4895. I will rebase once those have landed. It includes a number of hacks to work around bugs in `DotOperandEncodingAttr`. All these are marked as `FIXME [Dot LL]` to be easy to grep for. @Jokeren is working on a comprehensive revamp of `DotOperandEncodingAttr` which will get rid of all these. triton-lang#4895 is the first step in this direction.

Co-authored-by: Mario Lezcano Casado <lezcano@openai.com>

@Jokeren

This PR includes triton-lang#4891 and triton-lang#4895. I will rebase once those have landed. It includes a number of hacks to work around bugs in `DotOperandEncodingAttr`. All these are marked as `FIXME [Dot LL]` to be easy to grep for. @Jokeren is working on a comprehensive revamp of `DotOperandEncodingAttr` which will get rid of all these. triton-lang#4895 is the first step in this direction.

@Jokeren

This PR includes triton-lang/triton#4891 and triton-lang/triton#4895. I will rebase once those have landed. It includes a number of hacks to work around bugs in `DotOperandEncodingAttr`. All these are marked as `FIXME [Dot LL]` to be easy to grep for. @Jokeren is working on a comprehensive revamp of `DotOperandEncodingAttr` which will get rid of all these. triton-lang/triton#4895 is the first step in this direction.

Jokeren added 5 commits October 11, 2024 16:08

Update

3bf9218

Update

c7ee428

Update

5428fb2

Update

55bae07

Update

5a8a9ec

Jokeren marked this pull request as ready for review October 12, 2024 12:32

Jokeren requested review from antiagainst, ptillet and zhanglx13 as code owners October 12, 2024 12:32

Jokeren requested review from ThomasRaoux and lezcano October 12, 2024 12:32

zhanglx13 reviewed Oct 12, 2024

View reviewed changes

Jokeren added 2 commits October 12, 2024 19:15

Update Dialect.cpp

dc9588f

Update Dialect.cpp

a3ef441

lezcano mentioned this pull request Oct 14, 2024

[AMD] Add basics to allow bypass LDS for dot RHS #4856

Closed

lezcano reviewed Oct 14, 2024

View reviewed changes

Jokeren mentioned this pull request Oct 14, 2024

Requirements to pass WGMMA LHS operand in registers #4785

Open

lezcano mentioned this pull request Oct 14, 2024

[Backend] Implement scaled_dot(mxfp4, fp8) #4904

Merged

Jokeren added 9 commits October 14, 2024 20:48

Update

f9159e9

Update

7789dde

Update

78061cb

Update

539e815

Update

e01a3ca

Update

8f4f760

Update

d5b26f2

Update

dfa69e4

Update

7e90ac8

Update comments

5706933

Jokeren and others added 4 commits October 15, 2024 07:18

Update

b4856ae

Update

17378ed

Update

d7ea5d3

Merge branch 'main' into keren/dot-op-fix

e264941

lezcano approved these changes Oct 15, 2024

View reviewed changes

Jokeren merged commit f9688ab into main Oct 15, 2024

Jokeren deleted the keren/dot-op-fix branch October 15, 2024 16:16

lezcano mentioned this pull request Oct 18, 2024

[Backend] Pipeline scale_dot #4950

Merged

Luosuu pushed a commit to Luosuu/triton that referenced this pull request Nov 13, 2024

[BACKEND] Small fixes for dot operand properties (triton-lang#4895)

2f7efd8

Co-authored-by: Mario Lezcano Casado <lezcano@openai.com>

guacamoleo pushed a commit to guacamoleo/triton that referenced this pull request Nov 14, 2024

[BACKEND] Small fixes for dot operand properties (triton-lang#4895)

0fd5f19

Co-authored-by: Mario Lezcano Casado <lezcano@openai.com>

bertmaher pushed a commit to bertmaher/triton that referenced this pull request Dec 10, 2024

[BACKEND] Small fixes for dot operand properties (triton-lang#4895)

4cc2aaf

Co-authored-by: Mario Lezcano Casado <lezcano@openai.com>

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

[BACKEND] Small fixes for dot operand properties#4895

[BACKEND] Small fixes for dot operand properties#4895
Jokeren merged 21 commits intomainfrom
keren/dot-op-fix

Jokeren commented Oct 11, 2024 •

edited

Loading

Uh oh!

Jokeren commented Oct 11, 2024

Uh oh!

zhanglx13 Oct 12, 2024

Uh oh!

Jokeren Oct 13, 2024

Uh oh!

Jokeren Oct 15, 2024 •

edited

Loading

Uh oh!

lezcano left a comment

Uh oh!

Uh oh!

lezcano Oct 14, 2024

Uh oh!

Uh oh!

Jokeren commented Oct 15, 2024

Uh oh!

lezcano left a comment

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

3 participants

Conversation

Jokeren commented Oct 11, 2024 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

Jokeren commented Oct 11, 2024

Uh oh!

zhanglx13 Oct 12, 2024

Choose a reason for hiding this comment

Uh oh!

Jokeren Oct 13, 2024

Choose a reason for hiding this comment

Uh oh!

Jokeren Oct 15, 2024 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

lezcano left a comment

Choose a reason for hiding this comment

Uh oh!

Uh oh!

lezcano Oct 14, 2024

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Jokeren commented Oct 15, 2024

Uh oh!

lezcano left a comment

Choose a reason for hiding this comment

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

3 participants

Jokeren commented Oct 11, 2024 •

edited

Loading

Jokeren Oct 15, 2024 •

edited

Loading