[Linear Layouts] Implement LL conversion for DotOperand(version=2) by lezcano · Pull Request #4891 · triton-lang/triton

lezcano · 2024-10-11T15:20:32Z

Note that the current implementation uses DotOperandEncodingAttr::getWarpsPerCTA, which was buggy for cases where the warps are not of the form [numWarps, 1] or [1, numWarps]. This PR bundles a fix for this issue.

We will activate its use for a subset of DotOperandEncodings in a PR coming soon.

ThomasRaoux · 2024-10-11T16:38:16Z

Looks good overall but I'll look more in details after the lit tests are fixed

lezcano · 2024-10-14T13:56:26Z

 SmallVector<unsigned> DotOperandEncodingAttr::getWarpsPerCTA() const {
-  auto parentLayout = getParent();
-  assert(parentLayout && "DotOperandEncodingAttr must have a parent");
-  if (auto distributedLayout =
-          mlir::dyn_cast<DistributedEncodingTrait>(parentLayout)) {
-    return distributedLayout.getWarpsPerCTA();
-  } else {
-    llvm::report_fatal_error(
-        "DotOperandEncodingAttr non-DistributedEncodingAttr parent not "
-        "supported yet");
-  }
+  auto distributedLayout = mlir::cast<DistributedEncodingTrait>(getParent());
+  auto warps = distributedLayout.getWarpsPerCTA();
+  auto rank = warps.size();
+  auto kDim = getOpIdx() == 0 ? rank - 1 : rank - 2;
+  warps[kDim] = 1;
+  return warps;


@Jokeren added this fix as I needed it for the layout conversion to be correct and pass the newly added tests.

lezcano · 2024-10-14T13:56:40Z

Addressed the reviews

Jokeren · 2024-10-15T16:18:54Z


+LinearLayout ampereDotToLinearLayout(ArrayRef<int64_t> shape,
+                                     DotOperandEncodingAttr dot) {
+  // TODO,BE. Implement ampereMMA in terms of this one


What does "BE" mean? Backend?

Better Engineering

@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.

…riton-lang#4891) Note that the current implementation uses `DotOperandEncodingAttr::getWarpsPerCTA`, which was buggy for cases where the warps are not of the form `[numWarps, 1]` or `[1, numWarps]`. This PR bundles a fix for this issue. We will activate its use for a subset of `DotOperandEncoding`s in a PR coming soon.

@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.

…riton-lang#4891) Note that the current implementation uses `DotOperandEncodingAttr::getWarpsPerCTA`, which was buggy for cases where the warps are not of the form `[numWarps, 1]` or `[1, numWarps]`. This PR bundles a fix for this issue. We will activate its use for a subset of `DotOperandEncoding`s in a PR coming soon.

@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.

…riton-lang#4891) Note that the current implementation uses `DotOperandEncodingAttr::getWarpsPerCTA`, which was buggy for cases where the warps are not of the form `[numWarps, 1]` or `[1, numWarps]`. This PR bundles a fix for this issue. We will activate its use for a subset of `DotOperandEncoding`s in a PR coming soon.

@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.

lezcano requested a review from ptillet as a code owner October 11, 2024 15:20

lezcano changed the title ~~Implement LL~~ [Linear Layouts] Implement LL conversion for DotOperand Oct 11, 2024

lezcano requested a review from Jokeren October 11, 2024 15:21

lezcano changed the title ~~[Linear Layouts] Implement LL conversion for DotOperand~~ [Linear Layouts] Implement LL conversion for DotOperand(version=2) Oct 11, 2024

lezcano force-pushed the ll_dot branch from 4a0deb3 to 9a459eb Compare October 11, 2024 15:27

Jokeren reviewed Oct 11, 2024

View reviewed changes

Comment thread lib/Dialect/TritonGPU/IR/LinearLayoutConversions.cpp Outdated

Jokeren reviewed Oct 11, 2024

View reviewed changes

Comment thread unittest/Dialect/TritonGPU/LinearLayoutConversionsTest.cpp Outdated

lezcano mentioned this pull request Oct 14, 2024

[BACKEND] Small fixes for dot operand properties #4895

Merged

lezcano force-pushed the ll_dot branch from 9a459eb to c1c97f3 Compare October 14, 2024 13:55

lezcano commented 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

lezcano requested review from Jokeren and ThomasRaoux October 15, 2024 08:15

lezcano added 3 commits October 15, 2024 17:21

Implement LL

6742df0

Deactivate the LL cast for DotOperand for now

cf47c27

Fix getWarpsPerCTA

1bec1b6

Jokeren reviewed Oct 15, 2024

View reviewed changes

lezcano force-pushed the ll_dot branch from c1c97f3 to 1bec1b6 Compare October 15, 2024 16:21

address review

b0484fc

Jokeren approved these changes Oct 15, 2024

View reviewed changes

lezcano merged commit ec0bd4a into triton-lang:main Oct 15, 2024

lezcano mentioned this pull request Oct 18, 2024

[Backend] Pipeline scale_dot #4950

Merged

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

[Linear Layouts] Implement LL conversion for DotOperand(version=2)#4891

[Linear Layouts] Implement LL conversion for DotOperand(version=2)#4891
lezcano merged 4 commits intotriton-lang:mainfrom
lezcano:ll_dot

lezcano commented Oct 11, 2024 •

edited

Loading

Uh oh!

ThomasRaoux commented Oct 11, 2024

Uh oh!

Uh oh!

Uh oh!

lezcano Oct 14, 2024 •

edited

Loading

Uh oh!

lezcano commented Oct 14, 2024

Uh oh!

Jokeren Oct 15, 2024

Uh oh!

lezcano Oct 15, 2024

Uh oh!

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

3 participants

Conversation

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

Uh oh!

Uh oh!

ThomasRaoux commented Oct 11, 2024

Uh oh!

Uh oh!

Uh oh!

lezcano Oct 14, 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 commented Oct 14, 2024

Uh oh!

Jokeren Oct 15, 2024

Choose a reason for hiding this comment

Uh oh!

lezcano Oct 15, 2024

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

3 participants

lezcano commented Oct 11, 2024 •

edited

Loading

lezcano Oct 14, 2024 •

edited

Loading