Skip to content

Commit d502a7b

Browse files
whitneywhtsangwhitneywhtsang
authored
Reland "[LAYOUTS] Kill getWarpsPerCTA(Attribute) and prefer LinearLayout-based impl (#6252)" (#3786)
Closes #3757
2 parents 32dd479 + eeed3f5 commit d502a7b

File tree

17 files changed

+49
-111
lines changed

17 files changed

+49
-111
lines changed

‎include/triton/Dialect/TritonGPU/IR/Dialect.h‎

Lines changed: 10 additions & 11 deletions
Original file line numberDiff line numberDiff line change
@@ -99,9 +99,16 @@ unsigned getTotalElemsPerThread(Attribute layout, ArrayRef<int64_t> shape);
9999

100100
SmallVector<unsigned> getElemsPerThread(Type type);
101101

102-
// Returns the number of warps per CTA that may have access to replicated
103-
// elements. If you want non-replicated warps, use getWarpsPerCTAWithUniqueData.
104-
SmallVector<unsigned> getWarpsPerCTA(Attribute layout);
102+
// Returns the number of warps per CTA that have access to non-replicated
103+
// elements of the tensor. E.g. for a blocked layout with sizePerThread = [1,
104+
// 1], threadsPerWarp = [2, 16], warpsPerCTA = [1, 4] and tensor shape = [2, 2],
105+
// returns [1, 1], since the first warp has access to the full tensor, whereas
106+
// the other warps have access to replicated elements.
107+
SmallVector<unsigned> getWarpsPerCTA(Attribute layout,
108+
ArrayRef<int64_t> tensorShape);
109+
inline SmallVector<unsigned> getWarpsPerCTA(RankedTensorType type) {
110+
return getWarpsPerCTA(type.getEncoding(), type.getShape());
111+
}
105112

106113
// Returns the number of contiguous elements of the logical tensor that each
107114
// thread has access to, on each dimension of the tensor. For a blocked layout
@@ -122,14 +129,6 @@ inline SmallVector<unsigned> getThreadsPerWarp(RankedTensorType type) {
122129
return getThreadsPerWarp(type.getEncoding(), type.getShape());
123130
}
124131

125-
// Returns the number of warps per CTA that have access to non-replicated
126-
// elements of the tensor. E.g. for a blocked layout with sizePerThread = [1,
127-
// 1], threadsPerWarp = [2, 16], warpsPerCTA = [1, 4] and tensor shape = [2, 2],
128-
// returns [1, 1], since the first warp has access to the full tensor, whereas
129-
// the other warps have access to replicated elements.
130-
SmallVector<unsigned>
131-
getWarpsPerCTAWithUniqueData(Attribute layout, ArrayRef<int64_t> tensorShape);
132-
133132
// Returns the dimensions of the tensor from minor (fast-varying) to
134133
// major (slow-varying). For distributed layouts, this represents
135134
// the order of the elements within a thread.

‎include/triton/Dialect/TritonGPU/IR/TritonGPUAttrDefs.td‎

Lines changed: 5 additions & 10 deletions
Original file line numberDiff line numberDiff line change
@@ -597,11 +597,6 @@ We call each individual tile "rep".
597597
/*defaultImplementation=*/[{
598598
return toLinearEncoding($_self, shape).getElemsPerThread(shape);
599599
}]>,
600-
// Interface for the meta information about the multiple thread hierarchy.
601-
InterfaceMethod<"Get the shape of the warps per CTA.",
602-
"SmallVector<unsigned>",
603-
"getWarpsPerCTA">,
604-
605600
InterfaceMethod<"Convert to LinearLayout.",
606601
"LinearLayout",
607602
"toLinearLayout",
@@ -653,7 +648,6 @@ L(T) = [ {0,8} , {1,9} , {2,10}, {3,11}, {0,8} , {1, 9} , {2, 10}, {3, 11},
653648
SmallVector<unsigned> getCTAsPerCGA() const;
654649
SmallVector<unsigned> getCTAOrder() const;
655650
SmallVector<unsigned> getCTASplitNum() const;
656-
SmallVector<unsigned> getWarpsPerCTA() const;
657651

658652
LinearLayout toLinearLayout(ArrayRef<int64_t> shape) const;
659653
}];
@@ -703,6 +697,7 @@ def LinearEncodingAttr : DistributedEncoding<"LinearEncoding", "linear_encoding"
703697
SmallVector<unsigned> basesPerDim(StringAttr dimName,
704698
bool skipBroadcast = true) const;
705699
SmallVector<unsigned> getThreadsPerWarp() const;
700+
SmallVector<unsigned> getWarpsPerCTA() const;
706701

707702
// [FIXME LL] Supports legacy behaviour. We should remove these functions
708703
SmallVector<unsigned> getShapePerCTATile() const;
@@ -813,7 +808,7 @@ for
813808
ins
814809
ArrayRefParameter<"unsigned">:$sizePerThread,
815810
ArrayRefParameter<"unsigned">:$threadsPerWarp,
816-
ArrayRefParameter<"unsigned">:$warpsPerCTA__,
811+
ArrayRefParameter<"unsigned">:$warpsPerCTA,
817812
ArrayRefParameter<"unsigned">:$order, // the fastest-changing axis first
818813

819814
// CTALayout is optional in the textual IR. If omitted, we infer it to be a
@@ -1012,7 +1007,7 @@ V [ 0,4,8...60 1,5...61 2,6...62 3,7...63 ] [ 128,132...188 129,
10121007
ins
10131008
"unsigned": $versionMajor,
10141009
"unsigned": $versionMinor,
1015-
ArrayRefParameter<"unsigned">:$warpsPerCTA__,
1010+
ArrayRefParameter<"unsigned">:$warpsPerCTA,
10161011
"unsigned":$MDim,
10171012
"unsigned":$NDim,
10181013
"bool":$isTransposed,
@@ -1132,7 +1127,7 @@ Row |
11321127
ins
11331128
"unsigned": $version,
11341129
"bool":$isTransposed,
1135-
ArrayRefParameter<"unsigned">:$warpsPerCTA__,
1130+
ArrayRefParameter<"unsigned">:$warpsPerCTA,
11361131
"CTALayoutAttr":$CTALayout
11371132
);
11381133

@@ -1237,7 +1232,7 @@ For example, the matrix L corresponding to blockTileSize=[32,16] is:
12371232
ins
12381233
"unsigned":$versionMajor,
12391234
"unsigned":$versionMinor,
1240-
ArrayRefParameter<"unsigned">:$warpsPerCTA__,
1235+
ArrayRefParameter<"unsigned">:$warpsPerCTA,
12411236
"CTALayoutAttr":$CTALayout,
12421237
ArrayRefParameter<"unsigned">:$instrShape
12431238
);

‎lib/Analysis/Utility.cpp‎

Lines changed: 5 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -100,7 +100,7 @@ bool shouldUseDistSmem(Attribute srcLayout, Attribute dstLayout) {
100100
}
101101

102102
unsigned ReduceOpHelper::getInterWarpSizeWithUniqueData() {
103-
return getWarpsPerCTAWithUniqueData(srcEncoding, srcShape)[axis];
103+
return getWarpsPerCTA(srcEncoding, srcShape)[axis];
104104
}
105105

106106
unsigned ReduceOpHelper::getIntraWarpSizeWithUniqueData() {
@@ -113,7 +113,7 @@ bool ReduceOpHelper::isWarpSynchronous() {
113113
// in order to remove this change.
114114
if (!srcEncoding)
115115
return true;
116-
return getWarpsPerCTAWithUniqueData(srcEncoding, srcShape)[axis] == 1;
116+
return getWarpsPerCTA(srcEncoding, srcShape)[axis] == 1;
117117
}
118118

119119
SmallVector<unsigned> ReduceOpHelper::getScratchRepShape() {
@@ -181,7 +181,7 @@ unsigned ScanLoweringHelper::getAxisNumWarpsWithUniqueData() {
181181
unsigned ScanLoweringHelper::getAxisNumBlocks() {
182182
auto contigPerThread = getEncoding().getContigPerThread();
183183
auto threadsPerWarp = getEncoding().getThreadsPerWarp();
184-
auto warpsPerCTA = getWarpsPerCTA(getEncoding());
184+
auto warpsPerCTA = getEncoding().getWarpsPerCTA();
185185
unsigned axis = getAxis();
186186
return ceil<unsigned>(
187187
getShape()[axis],
@@ -191,7 +191,7 @@ unsigned ScanLoweringHelper::getAxisNumBlocks() {
191191
unsigned ScanLoweringHelper::getNonAxisNumBlocks() {
192192
auto contigPerThread = getEncoding().getContigPerThread();
193193
auto threadsPerWarp = getEncoding().getThreadsPerWarp();
194-
auto warpsPerCTA = getWarpsPerCTA(getEncoding());
194+
auto warpsPerCTA = getEncoding().getWarpsPerCTA();
195195
auto rank = contigPerThread.size();
196196
unsigned axis = getAxis();
197197
unsigned numBlocks = 1;
@@ -528,7 +528,7 @@ unsigned ScanLoweringHelper::getAxisBlockStride() {
528528
unsigned stride = 1;
529529
auto contigPerThread = getEncoding().getContigPerThread();
530530
auto threadsPerWarp = getEncoding().getThreadsPerWarp();
531-
auto warpsPerCTA = getWarpsPerCTA(getEncoding());
531+
auto warpsPerCTA = getEncoding().getWarpsPerCTA();
532532
for (unsigned dim : order) {
533533
if (dim == getAxis())
534534
return stride;

‎lib/Conversion/TritonGPUToLLVM/HistogramOpToLLVM.cpp‎

Lines changed: 2 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -87,9 +87,8 @@ static SmallVector<Value> computeCrossWarpHistogram(
8787
Value threadId, int numWarps) {
8888
auto b = TritonLLVMOpBuilder(loc, rewriter);
8989
SmallVector<Value> histogramValues;
90-
unsigned numWarpsWithUniqueData =
91-
mlir::triton::gpu::getWarpsPerCTAWithUniqueData(srcType.getEncoding(),
92-
srcType.getShape())[0];
90+
unsigned numWarpsWithUniqueData = mlir::triton::gpu::getWarpsPerCTA(
91+
srcType.getEncoding(), srcType.getShape())[0];
9392
Value laneId = b.and_(threadId, b.i32_val(numThreadPerWarp - 1));
9493
// Initialize the shared memory with zeros.
9594
int64_t numElementPerThread =

‎lib/Conversion/TritonGPUToLLVM/ScanOpToLLVM.cpp‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -393,7 +393,7 @@ ScanOpConversion::getDelinearizedIds(ConversionPatternRewriter &rewriter,
393393
auto srcEncoding = helper.getEncoding();
394394

395395
auto threadsPerWarp = srcEncoding.getThreadsPerWarp();
396-
auto warpsPerCTA = triton::gpu::getWarpsPerCTA(srcEncoding);
396+
auto warpsPerCTA = srcEncoding.getWarpsPerCTA();
397397
auto [multiDimLaneId, isRepresentativeLane] =
398398
getMultiDimLaneId(rewriter, helper, laneId);
399399
auto [multiDimWarpId, isRepresentativeWarp] =

‎lib/Conversion/TritonToTritonGPU/TritonToTritonGPUPass.cpp‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -155,7 +155,7 @@ struct TritonExpandDimsPattern
155155
retSizePerThread.insert(retSizePerThread.begin() + op.getAxis(), 1);
156156
auto retThreadsPerWarp = to_vector(argEncoding.getThreadsPerWarp());
157157
retThreadsPerWarp.insert(retThreadsPerWarp.begin() + op.getAxis(), 1);
158-
auto retWarpsPerCTA = argEncoding.getWarpsPerCTA();
158+
auto retWarpsPerCTA = to_vector(argEncoding.getWarpsPerCTA());
159159
retWarpsPerCTA.insert(retWarpsPerCTA.begin() + op.getAxis(), 1);
160160
SmallVector<unsigned, 4> retOrder(retShape.size());
161161
std::iota(retOrder.begin(), retOrder.end(), 0);

‎lib/Dialect/TritonGPU/IR/Dialect.cpp‎

Lines changed: 6 additions & 45 deletions
Original file line numberDiff line numberDiff line change
@@ -88,18 +88,8 @@ SmallVector<unsigned> getThreadsPerWarp(Attribute layout,
8888
return toLinearEncoding(layout, shape).getThreadsPerWarp();
8989
}
9090

91-
SmallVector<unsigned> getWarpsPerCTA(Attribute layout) {
92-
if (auto distributedLayout =
93-
mlir::dyn_cast<DistributedEncodingTrait>(layout)) {
94-
return distributedLayout.getWarpsPerCTA();
95-
}
96-
97-
llvm::report_fatal_error("getWarpsPerCTA not implemented");
98-
return SmallVector<unsigned>();
99-
}
100-
101-
SmallVector<unsigned> getWarpsPerCTAWithUniqueData(Attribute layout,
102-
ArrayRef<int64_t> shape) {
91+
SmallVector<unsigned> getWarpsPerCTA(Attribute layout,
92+
ArrayRef<int64_t> shape) {
10393
return toLinearEncoding(layout, shape).getWarpsPerCTA();
10494
}
10595

@@ -581,9 +571,6 @@ SmallVector<unsigned> BlockedEncodingAttr::getCTAOrder() const {
581571
SmallVector<unsigned> BlockedEncodingAttr::getCTASplitNum() const {
582572
return SmallVector<unsigned>(getCTALayout().getCTASplitNum());
583573
}
584-
SmallVector<unsigned> BlockedEncodingAttr::getWarpsPerCTA() const {
585-
return SmallVector<unsigned>(getWarpsPerCTA__());
586-
}
587574

588575
template <class T>
589576
SmallVector<T> SliceEncodingAttr::paddedShape(ArrayRef<T> shape) const {
@@ -640,15 +627,6 @@ SmallVector<unsigned> SliceEncodingAttr::getCTAsPerCGA() const {
640627
llvm::report_fatal_error(
641628
"getCTAsPerCGA for SliceEncodingAttr is not well-defined");
642629
}
643-
SmallVector<unsigned> SliceEncodingAttr::getWarpsPerCTA() const {
644-
auto parent = getParent();
645-
auto parentWarpsPerCTA = ::getWarpsPerCTA(parent);
646-
SmallVector<unsigned> warpsPerCTA = parentWarpsPerCTA;
647-
warpsPerCTA.erase(warpsPerCTA.begin() + getDim());
648-
int32_t nextDim = getDim() < warpsPerCTA.size() ? getDim() : getDim() - 1;
649-
warpsPerCTA[nextDim] *= parentWarpsPerCTA[getDim()];
650-
return warpsPerCTA;
651-
}
652630

653631
// Wmma encoding
654632

@@ -704,14 +682,6 @@ SmallVector<unsigned> DotOperandEncodingAttr::getCTASplitNum() const {
704682
res[kDim] = 1;
705683
return res;
706684
}
707-
SmallVector<unsigned> DotOperandEncodingAttr::getWarpsPerCTA() const {
708-
auto distributedLayout = mlir::cast<DistributedEncodingTrait>(getParent());
709-
auto warps = distributedLayout.getWarpsPerCTA();
710-
auto rank = warps.size();
711-
auto kDim = getOpIdx() == 0 ? rank - 1 : rank - 2;
712-
warps[kDim] = 1;
713-
return warps;
714-
}
715685

716686
LogicalResult DotOperandEncodingAttr::verify(
717687
::llvm::function_ref<::mlir::InFlightDiagnostic()> emitError,
@@ -1339,7 +1309,7 @@ void NvidiaMmaEncodingAttr::print(AsmPrinter &printer) const {
13391309
<< ", warpsPerCTA = [" << ArrayRef(getWarpsPerCTA()) << "]";
13401310

13411311
maybePrintCTALayout(getContext(), printer, getCTALayout(),
1342-
/*rank=*/getWarpsPerCTA().size());
1312+
/*rank=*/getRank());
13431313

13441314
printer << ", instrShape = [" << getInstrShape() << "]}>";
13451315
}
@@ -1419,11 +1389,11 @@ void AMDMfmaEncodingAttr::print(AsmPrinter &printer) const {
14191389
printer << "<{"
14201390
<< "versionMajor = " << getVersionMajor() //
14211391
<< ", versionMinor = " << getVersionMinor() //
1422-
<< ", warpsPerCTA = [" << ArrayRef(getWarpsPerCTA()) << "]" //
1392+
<< ", warpsPerCTA = [" << getWarpsPerCTA() << "]" //
14231393
<< ", instrShape = [" << ArrayRef{getMDim(), getNDim()} << "]" //
14241394
<< ", isTransposed = " << getIsTransposed();
14251395
maybePrintCTALayout(getContext(), printer, getCTALayout(),
1426-
/*rank=*/getWarpsPerCTA().size());
1396+
/*rank=*/getRank());
14271397
printer << "}>";
14281398
}
14291399

@@ -1754,9 +1724,6 @@ SmallVector<unsigned> AMDMfmaEncodingAttr::getCTAOrder() const {
17541724
SmallVector<unsigned> AMDMfmaEncodingAttr::getCTASplitNum() const {
17551725
return SmallVector<unsigned>(getCTALayout().getCTASplitNum());
17561726
}
1757-
SmallVector<unsigned> AMDMfmaEncodingAttr::getWarpsPerCTA() const {
1758-
return SmallVector<unsigned>(getWarpsPerCTA__());
1759-
}
17601727

17611728
SmallVector<int64_t>
17621729
AMDMfmaEncodingAttr::getInstrShapeForOperand(int kWidth, int opIdx) const {
@@ -1875,9 +1842,6 @@ SmallVector<unsigned> AMDWmmaEncodingAttr::getCTAOrder() const {
18751842
SmallVector<unsigned> AMDWmmaEncodingAttr::getCTASplitNum() const {
18761843
return SmallVector<unsigned>(getCTALayout().getCTASplitNum());
18771844
}
1878-
SmallVector<unsigned> AMDWmmaEncodingAttr::getWarpsPerCTA() const {
1879-
return SmallVector<unsigned>(getWarpsPerCTA__());
1880-
}
18811845

18821846
SmallVector<int64_t> AMDWmmaEncodingAttr::getElemsPerInstrForOperands() const {
18831847
return {16, 16};
@@ -1949,9 +1913,6 @@ SmallVector<unsigned> NvidiaMmaEncodingAttr::getCTAOrder() const {
19491913
SmallVector<unsigned> NvidiaMmaEncodingAttr::getCTASplitNum() const {
19501914
return SmallVector<unsigned>(getCTALayout().getCTASplitNum());
19511915
}
1952-
SmallVector<unsigned> NvidiaMmaEncodingAttr::getWarpsPerCTA() const {
1953-
return SmallVector<unsigned>(getWarpsPerCTA__());
1954-
}
19551916

19561917
SmallVector<unsigned>
19571918
NvidiaMmaEncodingAttr::getRepOrderForOperand(int opIdx) const {
@@ -1966,7 +1927,7 @@ NvidiaMmaEncodingAttr::getRepForOperand(ArrayRef<int64_t> shape, int bitwidth,
19661927
"kWidth must be >= 32 / bitwidth for this function to be well-defined");
19671928
auto rank = shape.size();
19681929
// Broadcast long K
1969-
auto warpsPerCTA = getWarpsPerCTA();
1930+
auto warpsPerCTA = to_vector(getWarpsPerCTA());
19701931
auto kDim = opIdx == 0 ? rank - 1 : rank - 2;
19711932
warpsPerCTA[kDim] = 1;
19721933

‎lib/Dialect/TritonGPU/IR/LinearLayoutConversions.cpp‎

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -365,7 +365,7 @@ static LinearLayout broadcastedDotOperandLayout(MLIRContext *ctx,
365365
LinearLayout
366366
AMDMfmaEncodingAttr::toLinearLayout(ArrayRef<int64_t> shape) const {
367367
int rank = shape.size();
368-
assert(rank == getWarpsPerCTA().size());
368+
assert(rank == getRank());
369369

370370
bool hasBatchDim = rank == 3;
371371
int mIndex = 0 + hasBatchDim;
@@ -713,7 +713,7 @@ LinearLayout mfmaDotToLinearLayout(DotOperandEncodingAttr dotMfmaLayout,
713713
LinearLayout
714714
AMDWmmaEncodingAttr::toLinearLayout(ArrayRef<int64_t> shape) const {
715715
int rank = shape.size();
716-
assert(rank == getWarpsPerCTA().size());
716+
assert(rank == getRank());
717717

718718
bool hasBatchDim = rank == 3;
719719
int mIndex = 0 + hasBatchDim;

‎lib/Dialect/TritonGPU/Transforms/AccelerateMatmul.cpp‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -72,7 +72,7 @@ SmallVector<unsigned> warpsPerTileV2(DotOp dotOp, const ArrayRef<int64_t> shape,
7272
}
7373
if (auto mmaEncoding =
7474
dyn_cast<NvidiaMmaEncodingAttr>(resTy.getEncoding())) {
75-
return getWarpsPerCTA(mmaEncoding);
75+
return to_vector(mmaEncoding.getWarpsPerCTA());
7676
}
7777
hasChainedDot = true;
7878
}

‎third_party/amd/lib/TritonAMDGPUToLLVM/OptimizeLDSUtility.cpp‎

Lines changed: 3 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -72,11 +72,10 @@ createTmpLayout(triton::gpu::DistributedEncodingTrait layout,
7272
src.getKWidth());
7373
}
7474
if (auto src = dyn_cast<triton::gpu::SliceEncodingAttr>(layout)) {
75-
// TODO: think of a way to construct slice layouts based on warpsPerCTA
76-
// argument
77-
auto parentWarpsPerCTA = triton::gpu::getWarpsPerCTA(src.getParent());
75+
auto warps = to_vector(warpsPerCTA);
76+
warps.insert(warps.begin() + src.getDim(), 1);
7877
return triton::gpu::SliceEncodingAttr::get(
79-
ctx, src.getDim(), createTmpLayout(src.getParent(), parentWarpsPerCTA));
78+
ctx, src.getDim(), createTmpLayout(src.getParent(), warps));
8079
}
8180
// TODO: support linear layout if needed.
8281
if (isa<triton::gpu::LinearEncodingAttr>(layout))

0 commit comments

Comments
 (0)