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

[BACKEND] Keren/shared memory barrier (#59)

Browse Source
This commit is contained in:
Keren Zhou
2022-08-18 12:32:57 -07:00
committed by GitHub
parent 8776ad1a0e
commit e0bedeb44c
13 changed files with 904 additions and 278 deletions
Download Patch File Download Diff File Expand all files Collapse all files

27
test/lib/Analysis/TestAllocation.cpp
View File

@@ -19,24 +19,29 @@ struct TestAllocationPass
void runOnOperation() override {
Operation *operation = getOperation();
auto &os = llvm::errs();
os << "Testing: " << operation->getName() << "\n";
AllocationAnalysis analysis(operation);
// Convert to std::string can remove quotes from op_name
auto op_name = SymbolTable::getSymbolName(operation).getValue().str();
os << op_name << "\n";
Allocation allocation(operation);
operation->walk([&](Operation *op) {
auto scratchBufferId = allocation.getBufferId(op);
if (scratchBufferId != Allocation::InvalidBufferId) {
size_t offset = allocation.getOffset(scratchBufferId);
size_t size = allocation.getAllocatedSize(scratchBufferId);
os << "scratch offset = " << offset << ", size = " << size << "\n";
}
if (op->getNumResults() < 1)
return;
for (Value result : op->getResults()) {
Type type = result.getType();
if (auto tensorType = type.dyn_cast<RankedTensorType>()) {
Attribute encoding = tensorType.getEncoding();
if (encoding.isa<triton::gpu::TritonGPUSharedEncodingAttr>()) {
size_t offset = analysis.getOffset(result);
size_t size = analysis.getAllocatedSize(result);
os << "offset = " << offset << ", size = " << size << "\n";
}
auto bufferId = allocation.getBufferId(result);
if (bufferId != Allocation::InvalidBufferId) {
size_t offset = allocation.getOffset(bufferId);
size_t size = allocation.getAllocatedSize(bufferId);
os << "offset = " << offset << ", size = " << size << "\n";
}
}
});
os << "size = " << analysis.getSharedMemorySize() << "\n";
os << "size = " << allocation.getSharedMemorySize() << "\n";
}
};