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

finish coding

Browse Source
This commit is contained in:
Superjomn
2022-11-04 16:54:05 +08:00
parent 1f552308c4
commit 1ed6ee34ba
3 changed files with 135 additions and 44 deletions
Download Patch File Download Diff File Expand all files Collapse all files

10
lib/Dialect/TritonGPU/IR/Dialect.cpp
View File

@@ -117,10 +117,12 @@ SmallVector<unsigned> getShapePerCTA(const Attribute &layout) {
"BlockedEncodingAttr not implemented");
}
} else if (auto mmaLayout = layout.dyn_cast<MmaEncodingAttr>()) {
assert(mmaLayout.getVersion() == 2 &&
"mmaLayout version = 1 is not implemented yet");
return {16 * mmaLayout.getWarpsPerCTA()[0],
8 * mmaLayout.getWarpsPerCTA()[1]};
if (mmaLayout.getVersion() == 2)
return {16 * mmaLayout.getWarpsPerCTA()[0],
8 * mmaLayout.getWarpsPerCTA()[1]};
if (mmaLayout.getVersion() == 1)
return {16 * mmaLayout.getWarpsPerCTA()[0],
16 * mmaLayout.getWarpsPerCTA()[1]};
} else {
assert(0 && "Unimplemented usage of getShapePerCTA");
}