Skip to content

Commit f034f3c

Browse files
sdalvi-quicsdalvi-quic
committed
uncommented
Change-Id: I9c3e4cfb41c846fea5499e4ce1e1828ebb18cb18
1 parent 57fcc04 commit f034f3c

1 file changed

Lines changed: 73 additions & 73 deletions

File tree

‎lib/Conversion/TritonToUnstructured/TritonToUnstructuredPass.cpp‎

Lines changed: 73 additions & 73 deletions
Original file line numberDiff line numberDiff line change
@@ -459,79 +459,79 @@ class TritonToUnstructuredPass
459459
return success();
460460
})
461461
.Case<scf::YieldOp>([](auto) { return success(); })
462-
// .Case<triton::BitcastOp>([&](triton::BitcastOp op) {
463-
// auto res = op.getResult();
464-
// auto resType = res.getType();
465-
466-
// if (!triton::isPtrTypeLike(resType)) {
467-
// return success();
468-
// }
469-
470-
// auto src = op.getSrc();
471-
// auto srcType = src.getType();
472-
473-
// // Extract pointee types, handling both tensor of pointers
474-
// // and scalar pointer cases.
475-
// Type srcPointeeTy, dstPointeeTy;
476-
// if (auto srcTensorTy = dyn_cast<RankedTensorType>(srcType)) {
477-
// srcPointeeTy = cast<triton::PointerType>(
478-
// srcTensorTy.getElementType()).getPointeeType();
479-
// dstPointeeTy = cast<triton::PointerType>(
480-
// cast<RankedTensorType>(resType).getElementType())
481-
// .getPointeeType();
482-
// } else {
483-
// srcPointeeTy = cast<triton::PointerType>(srcType)
484-
// .getPointeeType();
485-
// dstPointeeTy = cast<triton::PointerType>(resType)
486-
// .getPointeeType();
487-
// }
488-
489-
// // Use DataLayout to get the store size in bytes for each
490-
// // pointee type. This correctly handles sub-byte types
491-
// // (e.g., i1 occupies 1 byte in memory).
492-
// auto mod = op->getParentOfType<ModuleOp>();
493-
// mlir::DataLayout dataLayout(mod);
494-
// unsigned srcBytes = dataLayout.getTypeSize(srcPointeeTy);
495-
// unsigned dstBytes = dataLayout.getTypeSize(dstPointeeTy);
496-
497-
// if (srcBytes != dstBytes) {
498-
// op->emitError(
499-
// "bitcast between pointer types with different strides "
500-
// "is not supported in offset propagation (src size: ")
501-
// << srcBytes << " bytes, dst size: " << dstBytes
502-
// << " bytes)";
503-
// return failure();
504-
// }
505-
506-
// // Safe to reuse offset info — both pointer types have the
507-
// // same effective byte stride, so accumulated offsets remain
508-
// // valid after the bitcast.
509-
// auto offsetInfo = offsetMap.at(src);
510-
511-
// // Get the scalar pointer type for the new base pointer.
512-
// Type scalarPtrType;
513-
// if (auto resTensorTy = dyn_cast<RankedTensorType>(resType)) {
514-
// scalarPtrType = resTensorTy.getElementType();
515-
// } else {
516-
// scalarPtrType = resType;
517-
// }
518-
519-
// // Bitcast the scalar base pointer to match the new pointee
520-
// // type.
521-
// OpBuilder b{op};
522-
// Value newBasePtr = triton::BitcastOp::create(
523-
// b, op->getLoc(), scalarPtrType, offsetInfo.ptr);
524-
525-
// PtrOffset newOffsetInfo{newBasePtr, resType,
526-
// offsetInfo.bitWidth,
527-
// offsetInfo.offset};
528-
529-
// offsetMap.insert({res, newOffsetInfo});
530-
// workList.push(res);
531-
// toDelete.push_back(op);
532-
533-
// return success();
534-
// })
462+
.Case<triton::BitcastOp>([&](triton::BitcastOp op) {
463+
auto res = op.getResult();
464+
auto resType = res.getType();
465+
466+
if (!triton::isPtrTypeLike(resType)) {
467+
return success();
468+
}
469+
470+
auto src = op.getSrc();
471+
auto srcType = src.getType();
472+
473+
// Extract pointee types, handling both tensor of pointers
474+
// and scalar pointer cases.
475+
Type srcPointeeTy, dstPointeeTy;
476+
if (auto srcTensorTy = dyn_cast<RankedTensorType>(srcType)) {
477+
srcPointeeTy = cast<triton::PointerType>(
478+
srcTensorTy.getElementType()).getPointeeType();
479+
dstPointeeTy = cast<triton::PointerType>(
480+
cast<RankedTensorType>(resType).getElementType())
481+
.getPointeeType();
482+
} else {
483+
srcPointeeTy = cast<triton::PointerType>(srcType)
484+
.getPointeeType();
485+
dstPointeeTy = cast<triton::PointerType>(resType)
486+
.getPointeeType();
487+
}
488+
489+
// Use DataLayout to get the store size in bytes for each
490+
// pointee type. This correctly handles sub-byte types
491+
// (e.g., i1 occupies 1 byte in memory).
492+
auto mod = op->getParentOfType<ModuleOp>();
493+
mlir::DataLayout dataLayout(mod);
494+
unsigned srcBytes = dataLayout.getTypeSize(srcPointeeTy);
495+
unsigned dstBytes = dataLayout.getTypeSize(dstPointeeTy);
496+
497+
if (srcBytes != dstBytes) {
498+
op->emitError(
499+
"bitcast between pointer types with different strides "
500+
"is not supported in offset propagation (src size: ")
501+
<< srcBytes << " bytes, dst size: " << dstBytes
502+
<< " bytes)";
503+
return failure();
504+
}
505+
506+
// Safe to reuse offset info — both pointer types have the
507+
// same effective byte stride, so accumulated offsets remain
508+
// valid after the bitcast.
509+
auto offsetInfo = offsetMap.at(src);
510+
511+
// Get the scalar pointer type for the new base pointer.
512+
Type scalarPtrType;
513+
if (auto resTensorTy = dyn_cast<RankedTensorType>(resType)) {
514+
scalarPtrType = resTensorTy.getElementType();
515+
} else {
516+
scalarPtrType = resType;
517+
}
518+
519+
// Bitcast the scalar base pointer to match the new pointee
520+
// type.
521+
OpBuilder b{op};
522+
Value newBasePtr = triton::BitcastOp::create(
523+
b, op->getLoc(), scalarPtrType, offsetInfo.ptr);
524+
525+
PtrOffset newOffsetInfo{newBasePtr, resType,
526+
offsetInfo.bitWidth,
527+
offsetInfo.offset};
528+
529+
offsetMap.insert({res, newOffsetInfo});
530+
workList.push(res);
531+
toDelete.push_back(op);
532+
533+
return success();
534+
})
535535
.Case<triton::CatOp>([](triton::CatOp op) {
536536
op->emitError("Do not support gather / scatter with multiple "
537537
"bases yet");

0 commit comments

Comments
 (0)