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

[Triton-MLIR] Keren/code gen for extract slice and alloc tensor (#692)

Browse Source 
Co-authored-by: gzhu <goostavz@outlook.com>
This commit is contained in:
Keren Zhou
2022-09-23 12:38:14 -07:00
committed by GitHub
parent c56f0198dd
commit ecd1bc33df
5 changed files with 134 additions and 58 deletions
Download Patch File Download Diff File Expand all files Collapse all files

14
lib/Analysis/Allocation.cpp
View File

@@ -43,7 +43,6 @@ getScratchConfigForCvtLayout(triton::gpu::ConvertLayoutOp op, unsigned &inVec,
return 0;
}
};
// blocked -> blocked
if (srcLayout.isa<BlockedEncodingAttr>() &&
dstLayout.isa<BlockedEncodingAttr>()) {
auto srcBlockedLayout = srcLayout.cast<BlockedEncodingAttr>();
@@ -66,14 +65,6 @@ getScratchConfigForCvtLayout(triton::gpu::ConvertLayoutOp op, unsigned &inVec,
}
paddedRepShape[outOrd[0]] += pad;
}
// blocked -> shared
if (srcLayout.isa<BlockedEncodingAttr>() &&
dstLayout.isa<SharedEncodingAttr>()) {
auto sharedLayout = dstLayout.cast<SharedEncodingAttr>();
for (int v : dstTy.getShape())
paddedRepShape.push_back(v);
}
return paddedRepShape;
}
@@ -140,8 +131,9 @@ private:
auto dstTy = cvtLayout.result().getType().cast<RankedTensorType>();
auto srcEncoding = srcTy.getEncoding();
auto dstEncoding = dstTy.getEncoding();
if (srcEncoding.isa<SharedEncodingAttr>()) {
// only block->block and block->shared is supported now
if (srcEncoding.isa<SharedEncodingAttr>() ||
dstEncoding.isa<SharedEncodingAttr>()) {
// Only blocked -> blocked conversion requires for scratch allocation
return;
}
// ConvertLayoutOp with both input/output non-shared_layout