Implement conversion from FMA dot operand to linear layout #5469
There are no files selected for viewing
| Original file line number | Diff line number | Diff line change |
|---|---|---|
|
|
@@ -239,8 +239,11 @@ LinearLayout sharedToLinearLayoutLeadingOffset(ArrayRef<int64_t> shape, | |
| return combineCtaCgaWithShape(tileLayout, shared.getCTALayout(), shape); | ||
| } | ||
|
|
||
| /// Function to generate lane and warp layout for dot operands. | ||
| LinearLayout broadcastedDotOperandLayout(MLIRContext *ctx, | ||
| ArrayRef<unsigned> shape, | ||
| ArrayRef<unsigned> order, | ||
| unsigned kDim, StringAttr inDimName) { | ||
| // Let warpsPerCTAMma = {2, 2}, then | ||
| // warpsPerCTA = {2, 1} for opA and warpsPerCTA = {1, 2} for opB | ||
| // assume warpOrder = {1, 0} | ||
|
|
@@ -255,24 +258,23 @@ LinearLayout warpsDotOperand(MLIRContext *ctx, ArrayRef<unsigned> warpShape, | |
| // - - | - - - - | - - | ||
| // 2 3 | 2 3 0 2 | 1 3 | ||
| // In other words, we need to broadcast along K | ||
| auto rank = shape.size(); | ||
| auto dimNames = standardOutDimNames(ctx, rank); | ||
| LinearLayout layout = LinearLayout::empty(); | ||
|
|
||
| // We have to broadcast along the inner dimension | ||
| // For A, when moving along M we go from 0 to 2. | ||
| // For B, when moving along N we go from 0 to 1. | ||
| // As such, choosing the order of A {1, 0}, gives us the correct broadcasting | ||
| // Same happens if the warpOrder is {0, 1}, like in Hopper | ||
| for (auto d : order) { | ||
| if (d == kDim) { | ||
| layout *= LinearLayout::zeros1D(shape[d], inDimName, dimNames[d]); | ||
| } else { | ||
| layout *= LinearLayout::identity1D(shape[d], inDimName, dimNames[d]); | ||
| } | ||
| } | ||
| return layout; | ||
| } | ||
|
|
||
| } // anonymous namespace | ||
|
|
@@ -620,7 +622,8 @@ wmmaDotOperandToLinearLayout(DotOperandEncodingAttr dotWmmaLayout, | |
| // Generate warp layout | ||
| auto warpsPerCTA = wmmaLayout.getWarpsPerCTA(); | ||
| auto warpOrder = triton::gpu::getWarpOrder(dotWmmaLayout); | ||
| LinearLayout warpLayout = | ||
| broadcastedDotOperandLayout(ctx, warpsPerCTA, warpOrder, kDim, S("warp")); | ||
|
|
||
| // reorder dim names in rep order, so combineCtaCgaWithShape generate proper | ||
| // extension of layout | ||
|
|
@@ -650,6 +653,48 @@ BlockedEncodingAttr::toLinearLayout(ArrayRef<int64_t> shape) const { | |
| return combineCtaCgaWithShape(ctaLayout, getCTALayout(), shape); | ||
| } | ||
|
|
||
| std::optional<LinearLayout> | ||
| fmaDotToLinearLayout(DotOperandEncodingAttr operandLayout, | ||
| ArrayRef<int64_t> shape) { | ||
| int rank = shape.size(); | ||
| auto blocked = cast<BlockedEncodingAttr>(operandLayout.getParent()); | ||
| MLIRContext *ctx = operandLayout.getContext(); | ||
|
|
||
| // TODO: introduce registerOrder or use getOrder(operandLayout) | ||
| // Currently this order is used in legacy converter, because we do not | ||
| // have access to full dot operand layout, only parent part. | ||
| auto regOrder = blocked.getOrder(); | ||
| // TODO: use operandLayout.getThreadOrder() | ||
| auto threadOrder = blocked.getThreadOrder(); | ||
| auto warpOrder = blocked.getWarpOrder(); | ||
|
There was a problem hiding this comment. I feel like there's something wrong here. Have you tested a warp shape of There was a problem hiding this comment. My point is that warps have to be broadcasted along the k dimension There was a problem hiding this comment. Agreed. I think you have to use There was a problem hiding this comment.
Yes, this converter works with any warp shape. The trick is I create "thread" part of layout using whole K dimension, so any number of warps or threads across k dimension will be squeezed in Let's take an example with dot A operand with shape [m=32, k=32]:
Then I apply There was a problem hiding this comment.
I thought about this, the only reason I choose to go without it for now is aestetic: At this moment cta tile constructions looks like this: with in second variant layout is still extends beyond K dimension due to lane component. to make it uniform I can introduce something like There was a problem hiding this comment. I think the per thread layout would be [r0, r1] There was a problem hiding this comment. Oh, [2*2] is a blocked parent layout property. Dot operand is slightly different. It inherits all attributes of a parent, except k dimension. Dot operand layout implicitly extends per thread size to [2, K] for A operand and [K, 2] for B operand There was a problem hiding this comment. Picture you mentioned is related to intermediate layout, before it is expanded with combineCtaCgaWithShape There was a problem hiding this comment. Got it. I think it's different from the cases where the parent is mma; we don't do implicit broadcasting on only the register dimension. That being said, I think some code clean ups have to happen later. Right now, several methods crash on this dotoperand, like There was a problem hiding this comment. The update looks good to me now. Thanks! |
||
| auto repOrder = blocked.getRepOrder(); | ||
|
There was a problem hiding this comment. Should it be There was a problem hiding this comment. It is not implemented at the moment for blocked parent. We had a conversation about dot operand order functions with @lezcano and did not come to a certain decision. So for now I am just using parent order functions everywhere. |
||
|
|
||
| StringAttr kReg = S("register"); | ||
| StringAttr kLane = S("lane"); | ||
| StringAttr kWarp = S("warp"); | ||
|
|
||
| SmallVector<unsigned> threadSize = blocked.getSizePerThread(); | ||
| auto kDimIdx = operandLayout.getOpIdx() == 0 ? rank - 1 : rank - 2; | ||
| threadSize[kDimIdx] = shape[kDimIdx]; | ||
| auto threadShape = blocked.getThreadsPerWarp(); | ||
| auto warpShape = blocked.getWarpsPerCTA(); | ||
|
|
||
| SmallVector<StringAttr> repDimNames = | ||
| permuteDimNames(standardOutDimNames(ctx, rank), repOrder); | ||
|
|
||
| auto registersLayout = identityStandardND(kReg, threadSize, regOrder); | ||
| auto lanesLayout = broadcastedDotOperandLayout(ctx, threadShape, threadOrder, | ||
| kDimIdx, kLane); | ||
| auto warpsLayout = | ||
| broadcastedDotOperandLayout(ctx, warpShape, warpOrder, kDimIdx, kWarp); | ||
|
|
||
| LinearLayout ctaLayout = registersLayout.transposeOuts(repDimNames) * | ||
| lanesLayout.transposeOuts(repDimNames) * | ||
| warpsLayout.transposeOuts(repDimNames); | ||
|
|
||
| return combineCtaCgaWithShape(ctaLayout, getCTALayout(operandLayout), shape); | ||
| } | ||
|
|
||
| LinearLayout nvidiaMmaTile(MLIRContext *ctx, ArrayRef<unsigned> tileShape, | ||
| unsigned kWidth, ArrayRef<unsigned> order, | ||
| ArrayRef<unsigned> repOrder) { | ||
|
|
@@ -740,19 +785,21 @@ LinearLayout nvidiaDotToLinearLayout(ArrayRef<int64_t> shape, | |
| auto ctaLayout = | ||
| nvidiaMmaTile(ctx, tileShape, kWidth, getOrder(dot), dot.getRepOrder()); | ||
| auto kDim = isA ? rank - 1 : rank - 2; | ||
| ctaLayout *= broadcastedDotOperandLayout(ctx, mma.getWarpsPerCTA(), | ||
| mma.getWarpOrder(), kDim, S("warp")) | ||
| .transposeOuts(llvm::to_vector(ctaLayout.getOutDimNames())); | ||
|
|
||
| return combineCtaCgaWithShape(ctaLayout, getCTALayout(dot), shape); | ||
| } | ||
|
|
||
| std::optional<LinearLayout> | ||
| DotOperandEncodingAttr::toLinearLayout(ArrayRef<int64_t> shape) const { | ||
| auto parent = getParent(); | ||
| if (auto blockedLayout = mlir::dyn_cast<BlockedEncodingAttr>(parent)) { | ||
| return fmaDotToLinearLayout(*this, shape); | ||
| } else if (auto mfmaLayout = mlir::dyn_cast<AMDMfmaEncodingAttr>(parent)) { | ||
| return mfmaDotToLinearLayout(*this, shape); | ||
| } else if (auto wmmaLayout = mlir::dyn_cast<AMDWmmaEncodingAttr>(parent)) { | ||
| return wmmaDotOperandToLinearLayout(*this, shape); | ||
| } else if (auto mma = mlir::dyn_cast<NvidiaMmaEncodingAttr>(parent)) { | ||
| return nvidiaDotToLinearLayout(shape, *this); | ||
|
|
||
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
Does this mean that this path is still preferred over LLs? Could you also make LLs the preferred path, or is there anything blocking us from doing so?
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
I believe LL is already preferred, but I don't want to remove legacy converter yet, just in case.
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
but in which case did you hit the hard error? Can we just revert these changes?
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
I did not hit any errors so far. But want to compare code that is generated by legacy and new converter to see if there are differences in perf/register usage, etc.
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
Perhaps now this can be reverted before merging?