Skip to content

Commit b6d15d5

Browse files
elstainnileselstainniles
committed
fix(t2u): scope scalar address carriers to markers
Use complete i64 scalar-pointer addresses only for for/while slots listed by PointerDescriptorBoundary. Restore unmarked scalar loops to the established base-plus-relative-offset analysis while preserving descriptor-owned and tensor-pointer lowering.
1 parent 471302d commit b6d15d5

2 files changed

Lines changed: 65 additions & 34 deletions

File tree

‎third_party/ascend/lib/TritonToUnstructure/OffsetAnalysis.cpp‎

Lines changed: 3 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -390,7 +390,8 @@ void parseLoopRegionIterArg(LoopLikeOpInterface loopOp, const Location &loc,
390390
RewriterBase &rewriter,
391391
llvm::DenseMap<Value, PtrOffsetInfo> &offsetMap,
392392
BlockArgument regionIterArg) {
393-
if (isScalarPointer(regionIterArg)) {
393+
if (isScalarPointer(regionIterArg) &&
394+
isPointerDescriptorBoundaryRegionArgument(loopOp, regionIterArg)) {
394395
recordOpaqueScalarPointer(regionIterArg, offsetMap);
395396
return;
396397
}
@@ -1284,7 +1285,7 @@ void parseYield(scf::YieldOp op, const Location &loc, RewriterBase &rewriter,
12841285
void parseLoopOp(LoopLikeOpInterface op, const Location &loc,
12851286
RewriterBase &rewriter,
12861287
llvm::DenseMap<Value, PtrOffsetInfo> &offsetMap, Value dst) {
1287-
if (isScalarPointer(dst)) {
1288+
if (isScalarPointer(dst) && isPointerDescriptorBoundaryResult(op, dst)) {
12881289
recordOpaqueScalarPointer(dst, offsetMap);
12891290
return;
12901291
}

‎third_party/ascend/lib/TritonToUnstructure/ReplaceArguments.cpp‎

Lines changed: 62 additions & 32 deletions
Original file line numberDiff line numberDiff line change
@@ -20,6 +20,7 @@
2020
* THE SOFTWARE.
2121
*/
2222

23+
#include "TritonControlFlowOpt/ControlFlowRewrite.h"
2324
#include "TritonToUnstructure/UnstructureConversionPass.h"
2425
#include "Utils/Utils.h"
2526

@@ -37,6 +38,25 @@ static bool isScalarPointerType(Type type) {
3738
return pointerType && !isa<ShapedType>(pointerType.getPointeeType());
3839
}
3940

41+
static bool isPointerDescriptorBoundarySlot(Operation *loop, unsigned slot) {
42+
if (!loop)
43+
return false;
44+
auto slots = dyn_cast_or_null<DenseI32ArrayAttr>(
45+
loop->getAttr(controlflow::kPointerDescriptorBoundaryAttr));
46+
return slots &&
47+
llvm::is_contained(slots.asArrayRef(), static_cast<int32_t>(slot));
48+
}
49+
50+
// Only scalar pointers owned by the CFO descriptor schema cross a loop as a
51+
// complete i64 address. Ordinary scalar pointers retain the established
52+
// base-plus-relative-offset representation.
53+
static bool useCompleteScalarAddress(Operation *loop, unsigned slot,
54+
Type type) {
55+
return loop && isa<scf::ForOp, scf::WhileOp>(loop) &&
56+
isScalarPointerType(type) &&
57+
isPointerDescriptorBoundarySlot(loop, slot);
58+
}
59+
4060
// A scalar-pointer SCF slot is represented by one complete i64 address on all
4161
// structural edges. Region-local pointer users are rebuilt immediately, so no
4262
// pointer type crosses the control-flow boundary.
@@ -221,11 +241,12 @@ int getPtrTensorRank(Type type) {
221241
}
222242

223243
SmallVector<Value> constructOperands(ValueRange operands, Value tempVar,
224-
IRMapping mapping, OpBuilder &builder) {
244+
IRMapping mapping, OpBuilder &builder,
245+
Operation *loop) {
225246
SmallVector<Value> newOperands;
226-
for (Value originalOperand : operands) {
247+
for (auto [slot, originalOperand] : llvm::enumerate(operands)) {
227248
Value mappedOperand = mapping.lookupOrDefault(originalOperand);
228-
if (isScalarPointerType(originalOperand.getType()))
249+
if (useCompleteScalarAddress(loop, slot, originalOperand.getType()))
229250
mappedOperand = materializeScalarPointerAddress(mappedOperand, builder,
230251
originalOperand.getLoc());
231252
newOperands.push_back(mappedOperand);
@@ -236,10 +257,10 @@ SmallVector<Value> constructOperands(ValueRange operands, Value tempVar,
236257
return newOperands;
237258
}
238259

239-
SmallVector<Type> constructTypes(TypeRange types) {
260+
SmallVector<Type> constructTypes(TypeRange types, Operation *loop) {
240261
SmallVector<Type> newTypes;
241-
for (auto type : types) {
242-
newTypes.push_back(isScalarPointerType(type)
262+
for (auto [slot, type] : llvm::enumerate(types)) {
263+
newTypes.push_back(useCompleteScalarAddress(loop, slot, type)
243264
? IntegerType::get(type.getContext(), 64)
244265
: type);
245266
if (auto ptrType = dyn_cast<triton::PointerType>(type)) {
@@ -269,17 +290,18 @@ void replacePtrArguments(triton::FuncOp funcOp,
269290
// replacement immediately before the old op and erase the old op last.
270291
rewriter.setInsertionPoint(op);
271292
if (auto forOp = dyn_cast<scf::ForOp>(op)) {
272-
SmallVector<Value> newInitArgs =
273-
constructOperands(forOp.getInitArgs(), tempVar, mapping, rewriter);
293+
SmallVector<Value> newInitArgs = constructOperands(
294+
forOp.getInitArgs(), tempVar, mapping, rewriter, forOp);
274295
newOp = rewriter.create<scf::ForOp>(
275296
forOp.getLoc(), forOp.getLowerBound(), forOp.getUpperBound(),
276297
forOp.getStep(), newInitArgs,
277298
[&](OpBuilder &b, Location loc, Value iv, ValueRange args) {
278299
mapping.map(forOp.getInductionVar(), iv);
279300
auto newArgIter = args.begin();
280-
for (auto oldArg : forOp.getRegionIterArgs()) {
301+
for (auto [slot, oldArg] :
302+
llvm::enumerate(forOp.getRegionIterArgs())) {
281303
Value mappedArg = *newArgIter;
282-
if (isScalarPointerType(oldArg.getType()))
304+
if (useCompleteScalarAddress(forOp, slot, oldArg.getType()))
283305
mappedArg =
284306
rebuildScalarPointer(mappedArg, oldArg.getType(), b, loc);
285307
mapping.map(oldArg, mappedArg);
@@ -290,20 +312,23 @@ void replacePtrArguments(triton::FuncOp funcOp,
290312
b.clone(bodyOp, mapping);
291313
}
292314
auto yieldOp = cast<scf::YieldOp>(forOp.getBody()->getTerminator());
293-
b.create<scf::YieldOp>(
294-
yieldOp.getLoc(),
295-
constructOperands(yieldOp.getOperands(), tempVar, mapping, b));
315+
b.create<scf::YieldOp>(yieldOp.getLoc(),
316+
constructOperands(yieldOp.getOperands(),
317+
tempVar, mapping, b,
318+
forOp));
296319
});
297320
} else if (auto whileOp = dyn_cast<scf::WhileOp>(op)) {
298-
SmallVector<Value> newInits =
299-
constructOperands(whileOp.getInits(), tempVar, mapping, rewriter);
321+
SmallVector<Value> newInits = constructOperands(
322+
whileOp.getInits(), tempVar, mapping, rewriter, whileOp);
300323
newOp = rewriter.create<scf::WhileOp>(
301-
whileOp.getLoc(), constructTypes(whileOp->getResultTypes()), newInits,
324+
whileOp.getLoc(), constructTypes(whileOp->getResultTypes(), whileOp),
325+
newInits,
302326
[&](OpBuilder &b, Location loc, ValueRange args) {
303327
auto newArgIter = args.begin();
304-
for (auto oldArg : whileOp.getBeforeArguments()) {
328+
for (auto [slot, oldArg] :
329+
llvm::enumerate(whileOp.getBeforeArguments())) {
305330
Value mappedArg = *newArgIter;
306-
if (isScalarPointerType(oldArg.getType()))
331+
if (useCompleteScalarAddress(whileOp, slot, oldArg.getType()))
307332
mappedArg =
308333
rebuildScalarPointer(mappedArg, oldArg.getType(), b, loc);
309334
mapping.map(oldArg, mappedArg);
@@ -317,13 +342,15 @@ void replacePtrArguments(triton::FuncOp funcOp,
317342
b.create<scf::ConditionOp>(
318343
conditionOp.getLoc(),
319344
mapping.lookup(conditionOp.getCondition()),
320-
constructOperands(conditionOp.getArgs(), tempVar, mapping, b));
345+
constructOperands(conditionOp.getArgs(), tempVar, mapping, b,
346+
whileOp));
321347
},
322348
[&](OpBuilder &b, Location loc, ValueRange args) {
323349
auto newArgIter = args.begin();
324-
for (auto oldArg : whileOp.getAfterArguments()) {
350+
for (auto [slot, oldArg] :
351+
llvm::enumerate(whileOp.getAfterArguments())) {
325352
Value mappedArg = *newArgIter;
326-
if (isScalarPointerType(oldArg.getType()))
353+
if (useCompleteScalarAddress(whileOp, slot, oldArg.getType()))
327354
mappedArg =
328355
rebuildScalarPointer(mappedArg, oldArg.getType(), b, loc);
329356
mapping.map(oldArg, mappedArg);
@@ -334,9 +361,10 @@ void replacePtrArguments(triton::FuncOp funcOp,
334361
b.clone(bodyOp, mapping);
335362
}
336363
auto yieldOp = whileOp.getYieldOp();
337-
b.create<scf::YieldOp>(
338-
yieldOp.getLoc(),
339-
constructOperands(yieldOp.getOperands(), tempVar, mapping, b));
364+
b.create<scf::YieldOp>(yieldOp.getLoc(),
365+
constructOperands(yieldOp.getOperands(),
366+
tempVar, mapping, b,
367+
whileOp));
340368
});
341369
} else if (auto ifOp = dyn_cast<scf::IfOp>(op);
342370
ifOp && ifOp->getNumResults() > 0) {
@@ -347,18 +375,20 @@ void replacePtrArguments(triton::FuncOp funcOp,
347375
b.clone(bodyOp, mapping);
348376
}
349377
auto yieldOp = ifOp.thenYield();
350-
b.create<scf::YieldOp>(
351-
yieldOp.getLoc(),
352-
constructOperands(yieldOp.getOperands(), tempVar, mapping, b));
378+
b.create<scf::YieldOp>(yieldOp.getLoc(),
379+
constructOperands(yieldOp.getOperands(),
380+
tempVar, mapping, b,
381+
/*loop=*/nullptr));
353382
},
354383
[&](OpBuilder &b, Location loc) {
355384
for (auto &bodyOp : ifOp.elseBlock()->without_terminator()) {
356385
b.clone(bodyOp, mapping);
357386
}
358387
auto yieldOp = ifOp.elseYield();
359-
b.create<scf::YieldOp>(
360-
yieldOp.getLoc(),
361-
constructOperands(yieldOp.getOperands(), tempVar, mapping, b));
388+
b.create<scf::YieldOp>(yieldOp.getLoc(),
389+
constructOperands(yieldOp.getOperands(),
390+
tempVar, mapping, b,
391+
/*loop=*/nullptr));
362392
});
363393
} else if (auto loopOp = dyn_cast<LoopLikeOpInterface>(op)) {
364394
llvm_unreachable("Unsupported loop op");
@@ -370,9 +400,9 @@ void replacePtrArguments(triton::FuncOp funcOp,
370400
os << "Converting\n" << *op << "\nto\n" << *newOp << "\n";
371401
});
372402
auto resIter = newOp->result_begin();
373-
for (auto res : op->getResults()) {
403+
for (auto [slot, res] : llvm::enumerate(op->getResults())) {
374404
Value replacement = *resIter;
375-
if (isScalarPointerType(res.getType())) {
405+
if (useCompleteScalarAddress(op, slot, res.getType())) {
376406
rewriter.setInsertionPointAfter(newOp);
377407
replacement = rebuildScalarPointer(replacement, res.getType(),
378408
rewriter, res.getLoc());

0 commit comments

Comments
 (0)