Skip to content

Commit 7a1bf05

Browse files
ph0375flagtree-bot
ph0375
and
flagtree-bot
authored
[TLE] Add tle.gpu.set_layout primitive (#908)
* tle.gpu.set_layout * Apply code-format changes * Resolve merge conflicts with main_2 * Sync set_layout helpers to ppu/sunrise spec_cpp Utility files * Apply code-format changes * tle: guard explicit layout set_layout/tests behind __TLE__ capability check enflame/iluvatar CI use different TLE variants (or no cuda backend registered), so ir.builder.ensure_ttg_layout_attrs is unavailable there. Add a runtime capability check in set_layout and skip the two new frontend tests when the current build doesn't support it. * Apply code-format changes * fix: guard get_backend import failure in _cuda_backend_available * Add code comments * Apply code-format changes --------- Co-authored-by: flagtree-bot <flagtree_ai@163.com>
1 parent fd3ea29 commit 7a1bf05

18 files changed

Lines changed: 1143 additions & 22 deletions

File tree

‎include/triton/Dialect/TritonGPU/Transforms/Utility.h‎

Lines changed: 18 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -33,6 +33,9 @@
3333
#include "triton/Dialect/TritonGPU/IR/Dialect.h"
3434
#include <algorithm>
3535
#include <numeric>
36+
#ifdef __TLE__
37+
#include <string>
38+
#endif
3639

3740
namespace mlir {
3841
class DominanceInfo;
@@ -84,6 +87,21 @@ bool isView(Operation *op);
8487
// and lowers to llvm as the identity function (returns the input)
8588
bool isNoop(Operation *op);
8689

90+
#ifdef __TLE__
91+
std::string getTleExplicitEncodingAttrName(unsigned resultNumber);
92+
const char *getTleExplicitMemoryEncodingAttrName();
93+
Attribute getTleExplicitResultEncoding(Operation *op, unsigned resultNumber);
94+
void setTleExplicitResultEncoding(Operation *op, unsigned resultNumber,
95+
Attribute encoding);
96+
void setTleExplicitResultEncoding(OpResult result, Attribute encoding);
97+
Attribute getTleExplicitMemoryEncoding(Operation *op);
98+
void setTleExplicitMemoryEncoding(Operation *op, Attribute encoding);
99+
Attribute getTleExplicitValueEncoding(Value value);
100+
LogicalResult inferTleExplicitMemoryEncoding(Operation *op,
101+
Attribute &encoding);
102+
bool isTleExplicitConvertLayoutOp(Operation *op);
103+
#endif
104+
87105
/* Dump Triton IR in graphviz dot format.
88106
*
89107
* You can override `onValue` and `onOperation` in a subclass to mark

‎lib/Conversion/TritonToTritonGPU/TritonGPUConversion.cpp‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -212,6 +212,7 @@ TritonGPUConversionTarget::TritonGPUConversionTarget(
212212

213213
#ifdef __TLE__
214214
// flagtree tle raw
215+
addIllegalOp<triton::tle::SetLayoutOp>();
215216
addDynamicallyLegalDialect<triton::tle::TleDialect>([&](Operation *op) {
216217
bool hasLegalRegions = true;
217218
for (auto &region : op->getRegions()) {

‎lib/Conversion/TritonToTritonGPU/TritonToTritonGPUPass.cpp‎

Lines changed: 290 additions & 15 deletions
Original file line numberDiff line numberDiff line change
@@ -31,6 +31,8 @@
3131
#include "mlir/Transforms/DialectConversion.h"
3232
#ifdef __TLE__
3333
#include "tle/dialect/include/IR/Dialect.h"
34+
#include "llvm/ADT/MapVector.h"
35+
#include "llvm/ADT/PriorityWorklist.h"
3436
#endif
3537
#include "triton/Conversion/TritonToTritonGPU/Passes.h"
3638
#include "triton/Dialect/Triton/IR/Dialect.h"
@@ -58,6 +60,241 @@ static void addNamedAttrs(Operation *op, DictionaryAttr dictAttrs) {
5860
op->setAttr(attr.getName(), attr.getValue());
5961
}
6062

63+
#ifdef __TLE__
64+
struct TleEncodingInfo {
65+
Attribute encoding;
66+
bool mayVary = false;
67+
68+
explicit operator bool() const { return bool(encoding); }
69+
};
70+
71+
static bool tleEncodingsMayVary(Operation *op) {
72+
return isa<triton::JoinOp, triton::SplitOp, triton::ReshapeOp, triton::CatOp,
73+
triton::TransOp>(op);
74+
}
75+
76+
static LogicalResult mergeTleEncodingInfo(TleEncodingInfo oldInfo,
77+
TleEncodingInfo newInfo,
78+
Operation *op,
79+
TleEncodingInfo &merged) {
80+
if (!oldInfo) {
81+
merged = newInfo;
82+
return success();
83+
}
84+
if (!newInfo) {
85+
merged = oldInfo;
86+
return success();
87+
}
88+
if (oldInfo.encoding == newInfo.encoding) {
89+
merged = oldInfo;
90+
merged.mayVary = oldInfo.mayVary && newInfo.mayVary;
91+
return success();
92+
}
93+
if (oldInfo.mayVary && !newInfo.mayVary) {
94+
merged = newInfo;
95+
return success();
96+
}
97+
if (!oldInfo.mayVary && newInfo.mayVary) {
98+
merged = oldInfo;
99+
return success();
100+
}
101+
if (oldInfo.mayVary && newInfo.mayVary) {
102+
merged = oldInfo;
103+
return success();
104+
}
105+
106+
op->emitOpError("found conflicting TLE encoding hints for value:\n ")
107+
<< oldInfo.encoding << "\nand\n " << newInfo.encoding;
108+
return failure();
109+
}
110+
111+
static LogicalResult
112+
updateTleEncoding(ArrayRef<Value> values, TleEncodingInfo info, FuncOp func,
113+
llvm::MapVector<Value, TleEncodingInfo> &valueToEncoding,
114+
llvm::PriorityWorklist<Value> &worklist) {
115+
for (Value value : values) {
116+
if (!isa<RankedTensorType>(value.getType()))
117+
continue;
118+
119+
auto [it, inserted] = valueToEncoding.insert({value, info});
120+
if (!inserted) {
121+
Operation *defOp = value.getDefiningOp();
122+
Operation *diagOp = defOp ? defOp : func.getOperation();
123+
TleEncodingInfo merged;
124+
if (failed(mergeTleEncodingInfo(it->second, info, diagOp, merged)))
125+
return failure();
126+
if (merged.encoding == it->second.encoding &&
127+
merged.mayVary == it->second.mayVary)
128+
continue;
129+
it->second = merged;
130+
}
131+
worklist.insert(value);
132+
}
133+
return success();
134+
}
135+
136+
static LogicalResult propagateTleEncodingHints(FuncOp func) {
137+
llvm::SmallVector<std::pair<Value, TleEncodingInfo>> seedEncodings;
138+
func.walk([&](tle::SetLayoutOp op) {
139+
seedEncodings.push_back(
140+
{op.getSrc(), TleEncodingInfo{op.getTargetEncoding(), true}});
141+
seedEncodings.push_back(
142+
{op.getResult(), TleEncodingInfo{op.getTargetEncoding(), false}});
143+
});
144+
if (seedEncodings.empty())
145+
return success();
146+
147+
llvm::MapVector<Value, TleEncodingInfo> valueToEncoding;
148+
llvm::PriorityWorklist<Value> worklist;
149+
for (auto &[value, info] : seedEncodings) {
150+
if (failed(
151+
updateTleEncoding({value}, info, func, valueToEncoding, worklist)))
152+
return failure();
153+
}
154+
155+
while (!worklist.empty()) {
156+
Value value = worklist.pop_back_val();
157+
TleEncodingInfo info = valueToEncoding[value];
158+
assert(info && "worklist value must have an encoding");
159+
160+
for (OpOperand &use : value.getUses()) {
161+
Operation *op = use.getOwner();
162+
if (isa<scf::ForOp, scf::WhileOp>(op)) {
163+
int offset = 3 * isa<scf::ForOp>(op);
164+
auto tiedArgs = getTiedArgs(op, use.getOperandNumber() - offset);
165+
if (failed(updateTleEncoding(tiedArgs, info, func, valueToEncoding,
166+
worklist)))
167+
return failure();
168+
continue;
169+
}
170+
if (isa<scf::YieldOp>(op)) {
171+
auto tiedArgs = getTiedArgs(op, use.getOperandNumber());
172+
if (failed(updateTleEncoding(tiedArgs, info, func, valueToEncoding,
173+
worklist)))
174+
return failure();
175+
continue;
176+
}
177+
if (auto dot = dyn_cast<triton::DotOpInterface>(op)) {
178+
// A/B use dot-operand encodings, while C and D use the parent MMA
179+
// encoding. Only C and D are layout-equivalent across a dot.
180+
if (use.getOperandNumber() == 2 &&
181+
failed(updateTleEncoding({dot.getD()}, info, func, valueToEncoding,
182+
worklist)))
183+
return failure();
184+
continue;
185+
}
186+
if (isa<tle::SetLayoutOp>(op))
187+
continue;
188+
189+
if (isa<tle::LocalPointersOp>(op)) {
190+
if (failed(
191+
updateTleEncoding(llvm::to_vector_of<Value>(op->getResults()),
192+
info, func, valueToEncoding, worklist)))
193+
return failure();
194+
continue;
195+
}
196+
197+
Attribute dstEncoding = inferDstEncoding(op, info.encoding);
198+
if (!dstEncoding)
199+
continue;
200+
TleEncodingInfo dstInfo{dstEncoding,
201+
info.mayVary || tleEncodingsMayVary(op)};
202+
if (failed(updateTleEncoding(llvm::to_vector_of<Value>(op->getResults()),
203+
dstInfo, func, valueToEncoding, worklist)))
204+
return failure();
205+
}
206+
207+
if (auto opResult = dyn_cast<OpResult>(value)) {
208+
Operation *definingOp = opResult.getOwner();
209+
if (isa<scf::ForOp, scf::WhileOp, scf::IfOp>(definingOp)) {
210+
auto tiedArgs = getTiedArgs(definingOp, opResult.getResultNumber());
211+
if (failed(updateTleEncoding(tiedArgs, info, func, valueToEncoding,
212+
worklist)))
213+
return failure();
214+
} else if (isa<triton::DotOpInterface>(definingOp)) {
215+
if (failed(updateTleEncoding({definingOp->getOperand(2)}, info, func,
216+
valueToEncoding, worklist)))
217+
return failure();
218+
} else if (auto localPointers =
219+
dyn_cast<tle::LocalPointersOp>(definingOp)) {
220+
llvm::SmallVector<Value> tensorIndices;
221+
for (Value index : localPointers.getIndices())
222+
if (isa<RankedTensorType>(index.getType()))
223+
tensorIndices.push_back(index);
224+
if (failed(updateTleEncoding(tensorIndices, info, func, valueToEncoding,
225+
worklist)))
226+
return failure();
227+
} else if (!isa<tle::SetLayoutOp>(definingOp)) {
228+
Attribute srcEncoding = inferSrcEncoding(definingOp, info.encoding);
229+
if (srcEncoding) {
230+
TleEncodingInfo srcInfo{
231+
srcEncoding, info.mayVary || tleEncodingsMayVary(definingOp)};
232+
llvm::SmallVector<Value> tensorOperands;
233+
for (Value operand : definingOp->getOperands())
234+
if (isa<RankedTensorType>(operand.getType()))
235+
tensorOperands.push_back(operand);
236+
if (failed(updateTleEncoding(tensorOperands, srcInfo, func,
237+
valueToEncoding, worklist)))
238+
return failure();
239+
}
240+
}
241+
} else if (auto blockArg = dyn_cast<BlockArgument>(value)) {
242+
Operation *parentOp = blockArg.getOwner()->getParentOp();
243+
if (isa<scf::ForOp, scf::WhileOp>(parentOp)) {
244+
int offset = isa<scf::ForOp>(parentOp);
245+
auto tiedArgs = getTiedArgs(parentOp, blockArg.getArgNumber() - offset);
246+
if (failed(updateTleEncoding(tiedArgs, info, func, valueToEncoding,
247+
worklist)))
248+
return failure();
249+
}
250+
}
251+
}
252+
253+
for (auto &[value, info] : valueToEncoding) {
254+
auto existingTy = cast<RankedTensorType>(value.getType());
255+
if (existingTy.getEncoding() != info.encoding) {
256+
auto newTy = existingTy.cloneWithEncoding(info.encoding);
257+
value.setType(newTy);
258+
259+
if (auto opResult = dyn_cast<OpResult>(value)) {
260+
if (auto constant = dyn_cast<arith::ConstantOp>(opResult.getOwner())) {
261+
if (auto elements =
262+
dyn_cast<DenseElementsAttr>(constant.getValueAttr()))
263+
constant.setValueAttr(elements.reshape(newTy));
264+
}
265+
}
266+
}
267+
268+
if (auto opResult = dyn_cast<OpResult>(value))
269+
setTleExplicitResultEncoding(opResult, info.encoding);
270+
}
271+
272+
WalkResult memoryWalk = func.walk([&](Operation *op) {
273+
if (!getMemAccessPtr(op))
274+
return WalkResult::advance();
275+
276+
Attribute explicitEncoding;
277+
if (failed(inferTleExplicitMemoryEncoding(op, explicitEncoding)))
278+
return WalkResult::interrupt();
279+
280+
if (explicitEncoding)
281+
setTleExplicitMemoryEncoding(op, explicitEncoding);
282+
return WalkResult::advance();
283+
});
284+
if (memoryWalk.wasInterrupted())
285+
return failure();
286+
287+
return success();
288+
}
289+
290+
static LogicalResult applyTleEncodingHints(ModuleOp mod) {
291+
for (FuncOp func : mod.getOps<FuncOp>())
292+
if (failed(propagateTleEncodingHints(func)))
293+
return failure();
294+
return success();
295+
}
296+
#endif
297+
61298
template <class Op> struct GenericOpPattern : public OpConversionPattern<Op> {
62299
using OpConversionPattern<Op>::OpConversionPattern;
63300

@@ -959,25 +1196,58 @@ class TleInsertTileOpPattern : public OpConversionPattern<tle::InsertTileOp> {
9591196
}
9601197
};
9611198

1199+
#ifdef __TLE__
1200+
class TleSetLayoutOpPattern : public OpConversionPattern<tle::SetLayoutOp> {
1201+
public:
1202+
using OpConversionPattern::OpConversionPattern;
1203+
1204+
LogicalResult
1205+
matchAndRewrite(tle::SetLayoutOp op, tle::SetLayoutOp::Adaptor adaptor,
1206+
ConversionPatternRewriter &rewriter) const override {
1207+
Type convertedType = getTypeConverter()->convertType(op.getResult());
1208+
auto resultType = dyn_cast<RankedTensorType>(convertedType);
1209+
if (!resultType)
1210+
return rewriter.notifyMatchFailure(op, "expected ranked tensor result");
1211+
1212+
Value src = adaptor.getSrc();
1213+
if (src.getType() == resultType) {
1214+
if (auto srcResult = dyn_cast<OpResult>(src))
1215+
setTleExplicitResultEncoding(srcResult, op.getTargetEncoding());
1216+
rewriter.replaceOp(op, src);
1217+
return success();
1218+
}
1219+
1220+
auto convert = rewriter.replaceOpWithNewOp<triton::gpu::ConvertLayoutOp>(
1221+
op, resultType, src);
1222+
convert->setAttr(getTleExplicitEncodingAttrName(0), op.getTargetEncoding());
1223+
return success();
1224+
}
1225+
};
1226+
1227+
#endif
1228+
9621229
// flagtree tle raw
9631230
void populateTleRawPatterns(TritonGPUTypeConverter &typeConverter,
9641231
RewritePatternSet &patterns) {
9651232
MLIRContext *context = patterns.getContext();
966-
patterns
967-
.add<TleDSLRegionOpPattern, TleExtractTileOpPattern,
968-
TleInsertTileOpPattern, GenericOpPattern<tle::LocalPointersOp>,
969-
GenericOpPattern<tle::RemotePointersOp>,
970-
GenericOpPattern<tle::ExclusiveCumsumOp>,
971-
GenericOpPattern<tle::WGMMAOp>, GenericOpPattern<tle::WGMMAWaitOp>,
972-
GenericOpPattern<tle::DistributedBarrierOp>,
973-
GenericOpPattern<tle::YieldOp>,
974-
GenericOpPattern<tle::ExtractAllocatedPtrOp>,
975-
GenericOpPattern<tle::ExtractAlignedPtrOp>,
976-
GenericOpPattern<tle::ExtractOffsetOp>,
977-
GenericOpPattern<tle::ExtractSizesOp>,
978-
GenericOpPattern<tle::ExtractStridesOp>,
979-
GenericOpPattern<tle::ExtractPtrOp>, GenericOpPattern<tle::PackOp>>(
980-
typeConverter, context);
1233+
patterns.add<
1234+
TleDSLRegionOpPattern, TleExtractTileOpPattern, TleInsertTileOpPattern,
1235+
#ifdef __TLE__
1236+
TleSetLayoutOpPattern,
1237+
#endif
1238+
GenericOpPattern<tle::LocalPointersOp>,
1239+
GenericOpPattern<tle::RemotePointersOp>,
1240+
GenericOpPattern<tle::ExclusiveCumsumOp>, GenericOpPattern<tle::WGMMAOp>,
1241+
GenericOpPattern<tle::WGMMAWaitOp>,
1242+
GenericOpPattern<tle::DistributedBarrierOp>,
1243+
GenericOpPattern<tle::YieldOp>,
1244+
GenericOpPattern<tle::ExtractAllocatedPtrOp>,
1245+
GenericOpPattern<tle::ExtractAlignedPtrOp>,
1246+
GenericOpPattern<tle::ExtractOffsetOp>,
1247+
GenericOpPattern<tle::ExtractSizesOp>,
1248+
GenericOpPattern<tle::ExtractStridesOp>,
1249+
GenericOpPattern<tle::ExtractPtrOp>, GenericOpPattern<tle::PackOp>>(
1250+
typeConverter, context);
9811251
}
9821252
#endif
9831253

@@ -1022,6 +1292,11 @@ class ConvertTritonToTritonGPU
10221292
mod->setAttr(AttrNumCTAsName, b.getI32IntegerAttr(numCTAs));
10231293
mod->setAttr(AttrTargetName, b.getStringAttr(this->target.getValue()));
10241294

1295+
#ifdef __TLE__
1296+
if (failed(applyTleEncodingHints(mod)))
1297+
return signalPassFailure();
1298+
#endif
1299+
10251300
if (failed(applyPartialConversion(mod, target, std::move(patterns))))
10261301
return signalPassFailure();
10271302
}

0 commit comments

Comments
 (0)