gaspersic/triton
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

[Triton-MLIR] fix a tiny bug in coalesce pass (#782)

Browse Source
This commit is contained in:
goostavz
2022-10-17 11:29:55 +08:00
committed by GitHub
parent 5898352f97
commit e948a618b3
1 changed files with 7 additions and 3 deletions
Download Patch File Download Diff File Expand all files Collapse all files

10
lib/Dialect/TritonGPU/Transforms/Coalesce.cpp
View File

@@ -63,9 +63,13 @@ struct CoalescePass : public TritonGPUCoalesceBase<CoalescePass> {
auto convertType = getTypeConverter(axisInfo, ptr, numWarps);
// convert operands
SmallVector<Value, 4> newArgs;
for (auto v : op->getOperands())
newArgs.push_back(builder.create<triton::gpu::ConvertLayoutOp>(
op->getLoc(), convertType(v.getType()), v));
for (auto v : op->getOperands()) {
if (v.getType().isa<RankedTensorType>())
newArgs.push_back(builder.create<triton::gpu::ConvertLayoutOp>(
op->getLoc(), convertType(v.getType()), v));
else
newArgs.push_back(v);
}
// convert output types
SmallVector<Type, 4> newTypes;
for (auto t : op->getResultTypes()) {