AI/rocm-systems
2
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

SWDEV-465602 - Refactor kernel arg pool allocation for graphs

Browse Source 
- Allocate additional argument space to accomodate for kernel node
param updates

Change-Id: I2d4ea8bddd716f1191f3cbea807920d0248f8c4e
This commit is contained in:
Saleel Kudchadker
2024-06-20 19:37:04 +00:00
parent f309d49b32
commit 17313ec99d
2 changed files with 92 additions and 82 deletions
Download Patch File Download Diff File Expand all files Collapse all files
hipamd/src/hip_graph_internal.cpp
+57 -55
View File
@@ -237,6 +237,7 @@ void Graph::GetRunList(std::vector<std::vector<Node>>& parallelLists,
}
}
}
bool Graph::TopologicalOrder(std::vector<Node>& TopoOrder) {
std::queue<Node> q;
std::unordered_map<Node, int> inDegree;
@@ -263,6 +264,7 @@ bool Graph::TopologicalOrder(std::vector<Node>& TopoOrder) {
}
return false;
}
Graph* Graph::clone(std::unordered_map<Node, Node>& clonedNodes) const {
Graph* newGraph = new Graph(device_, this);
for (auto entry : vertices_) {
@@ -352,6 +354,9 @@ hipError_t GraphExec::Init() {
return status;
}
//! Chunk size to add to kern arg pool
constexpr uint32_t kKernArgChunkSize = 128 * Ki;
// ================================================================================================
void GetKernelArgSizeForGraph(std::vector<std::vector<Node>>& parallelLists,
size_t& kernArgSizeForGraph) {
// GPU packet capture is enabled for kernel nodes. Calculate the kernel
@@ -361,7 +366,8 @@ void GetKernelArgSizeForGraph(std::vector<std::vector<Node>>& parallelLists,
if (node->GetType() == hipGraphNodeTypeKernel) {
kernArgSizeForGraph += reinterpret_cast<hip::GraphKernelNode*>(node)->GetKerArgSize();
} else if (node->GetType() == hipGraphNodeTypeGraph) {
auto& childParallelLists = reinterpret_cast<hip::ChildGraphNode*>(node)->GetParallelLists();
auto& childParallelLists =
reinterpret_cast<hip::ChildGraphNode*>(node)->GetParallelLists();
if (childParallelLists.size() == 1) {
GetKernelArgSizeForGraph(childParallelLists, kernArgSizeForGraph);
}
@@ -370,8 +376,9 @@ void GetKernelArgSizeForGraph(std::vector<std::vector<Node>>& parallelLists,
}
}
hipError_t AllocKernelArgForGraph(std::vector<hip::Node>& topoOrder, hip::Stream* capture_stream,
hip::GraphExec* graphExec) {
// ================================================================================================
hipError_t AllocKernelArgForGraphNode(std::vector<hip::Node>& topoOrder,
hip::Stream* capture_stream, hip::GraphExec* graphExec) {
hipError_t status = hipSuccess;
for (auto& node : topoOrder) {
if (node->GetType() == hipGraphNodeTypeKernel) {
@@ -399,7 +406,8 @@ hipError_t AllocKernelArgForGraph(std::vector<hip::Node>& topoOrder, hip::Stream
if (childParallelLists.size() == 1) {
childNode->SetGraphCaptureStatus(true);
status =
AllocKernelArgForGraph(childNode->GetChildGraphNodeOrder(), capture_stream, graphExec);
AllocKernelArgForGraphNode(childNode->GetChildGraphNodeOrder(),
capture_stream, graphExec);
if (status != hipSuccess) {
return status;
}
@@ -409,28 +417,44 @@ hipError_t AllocKernelArgForGraph(std::vector<hip::Node>& topoOrder, hip::Stream
return status;
}
// ================================================================================================
hipError_t GraphExec::AllocGraphKernargPool(size_t pool_size) {
hipError_t status = hipSuccess;
assert(pool_size > 0);
address graph_kernarg_base;
auto device = g_devices[ihipGetDevice()]->devices()[0];
if (device->info().largeBar_) {
graph_kernarg_base =
reinterpret_cast<address>(device->deviceLocalAlloc(pool_size));
device_kernarg_pool_ = true;
} else {
graph_kernarg_base = reinterpret_cast<address>(
device->hostAlloc(pool_size, 0,
amd::Device::MemorySegment::kKernArg));
}
if (graph_kernarg_base == nullptr) {
return hipErrorMemoryAllocation;
}
kernarg_graph_.push_back(KernelArgPoolGraph(graph_kernarg_base, pool_size));
return status;
}
// ================================================================================================
hipError_t GraphExec::CaptureAQLPackets() {
hipError_t status = hipSuccess;
if (parallelLists_.size() == 1) {
size_t kernArgSizeForGraph = 0;
GetKernelArgSizeForGraph(parallelLists_, kernArgSizeForGraph);
auto device = g_devices[ihipGetDevice()]->devices()[0];
if (kernArgSizeForGraph != 0) {
if (device->info().largeBar_) {
kernarg_pool_graph_ =
reinterpret_cast<address>(device->deviceLocalAlloc(kernArgSizeForGraph));
device_kernarg_pool_ = true;
} else {
kernarg_pool_graph_ = reinterpret_cast<address>(
device->hostAlloc(kernArgSizeForGraph, 0, amd::Device::MemorySegment::kKernArg));
}
if (kernarg_pool_graph_ == nullptr) {
return hipErrorMemoryAllocation;
}
kernarg_pool_size_graph_ = kernArgSizeForGraph;
// Add a larger initial pool to accomodate for any updates to kernel args
status = AllocGraphKernargPool(kernArgSizeForGraph + kKernArgChunkSize);
if (status != hipSuccess) {
return status;
}
status = AllocKernelArgForGraph(topoOrder_, capture_stream_, this);
status = AllocKernelArgForGraphNode(topoOrder_, capture_stream_, this);
if (status != hipSuccess) {
return status;
}
@@ -442,8 +466,9 @@ hipError_t GraphExec::CaptureAQLPackets() {
*device->info().hdpMemFlushCntl = 1u;
auto kSentinel = *reinterpret_cast<volatile int*>(device->info().hdpMemFlushCntl);
} else if (kernArgImpl == KernelArgImpl::DeviceKernelArgsReadback &&
kernarg_pool_size_graph_ != 0) {
address dev_ptr = kernarg_pool_graph_ + kernarg_pool_size_graph_;
kernarg_graph_.back().kernarg_pool_addr_ != 0) {
address dev_ptr = kernarg_graph_.back().kernarg_pool_addr_ +
kernarg_graph_.back().kernarg_pool_size_;
auto kSentinel = *reinterpret_cast<volatile address>(dev_ptr - 1);
_mm_sfence();
*(dev_ptr - 1) = kSentinel;
@@ -455,46 +480,23 @@ hipError_t GraphExec::CaptureAQLPackets() {
return status;
}
// ================================================================================================
hipError_t GraphExec::UpdateAQLPacket(hip::GraphKernelNode* node) {
hipError_t status = hipSuccess;
if (parallelLists_.size() == 1) {
size_t pool_new_usage = 0;
address result = nullptr;
if (!kernarg_graph_.empty()) {
// 1. Allocate memory for the kernel args
size_t kernArgSizeForNode = 0;
kernArgSizeForNode = node->GetKerArgSize();
address kernArgOffset = nullptr;
kernArgOffset = allocKernArg(node->GetKerArgSize(),
node->GetKernargSegmentAlignment());
result = amd::alignUp(kernarg_graph_.back() + kernarg_graph_cur_offset_,
node->GetKernargSegmentAlignment());
pool_new_usage = (result + kernArgSizeForNode) - kernarg_graph_.back();
}
if (pool_new_usage != 0 && pool_new_usage <= kernarg_graph_size_) {
kernarg_graph_cur_offset_ = pool_new_usage;
} else {
address kernarg_graph;
auto device = g_devices[ihipGetDevice()]->devices()[0];
if (device->info().largeBar_) {
kernarg_graph = reinterpret_cast<address>(device->deviceLocalAlloc(kernarg_graph_size_));
} else {
kernarg_graph = reinterpret_cast<address>(
device->hostAlloc(kernarg_graph_size_, 0, amd::Device::MemorySegment::kKernArg));
}
kernarg_graph_.push_back(kernarg_graph);
kernarg_graph_cur_offset_ = 0;
// 1. Allocate memory for the kernel args
size_t kernArgSizeForNode = 0;
kernArgSizeForNode = node->GetKerArgSize();
result = amd::alignUp(kernarg_graph_.back() + kernarg_graph_cur_offset_,
node->GetKernargSegmentAlignment());
const size_t pool_new_usage = (result + kernArgSizeForNode) - kernarg_graph_.back();
if (pool_new_usage <= kernarg_graph_size_) {
kernarg_graph_cur_offset_ = pool_new_usage;
}
if (kernArgOffset == nullptr ) {
// Allocate new pool for kernarg and get the offset
status = AllocGraphKernargPool(kKernArgChunkSize);
kernArgOffset = allocKernArg(node->GetKerArgSize(),
node->GetKernargSegmentAlignment());
}
// 2. copy kernel args / create new AQL packet
node->CaptureAndFormPacket(capture_stream_, result);
node->CaptureAndFormPacket(capture_stream_, kernArgOffset);
}
return hipSuccess;
}