SWDEV-465602 - Refactor kernel arg pool allocation for graphs

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

Change-Id: I2d4ea8bddd716f1191f3cbea807920d0248f8c4e
Этот коммит содержится в:
Saleel Kudchadker
2024-06-20 19:37:04 +00:00
родитель f309d49b32
Коммит 17313ec99d
2 изменённых файлов: 92 добавлений и 82 удалений
+57 -55
Просмотреть файл
@@ -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;
}
+35 -27
Просмотреть файл
@@ -546,8 +546,19 @@ struct Graph {
};
struct GraphKernelNode;
struct GraphExec : public amd::ReferenceCountedObject {
struct KernelArgPoolGraph {
KernelArgPoolGraph(address base_addr, size_t size)
: kernarg_pool_addr_(base_addr),
kernarg_pool_size_(size),
kernarg_pool_offset_(0)
{}
address kernarg_pool_addr_; //! Base address of the kernel arg pool
size_t kernarg_pool_size_; //! Size of the pool
size_t kernarg_pool_offset_; //! Current offset in the kernel arg alloc
};
std::vector<std::vector<Node>> parallelLists_;
// Topological order of the graph doesn't include nodes embedded as part of the child graph
//! Topological order of the graph doesn't include nodes embedded as part of the child graph
std::vector<Node> topoOrder_;
std::unordered_map<Node, std::vector<Node>> nodeWaitLists_;
struct Graph* clonedGraph_;
@@ -560,16 +571,13 @@ struct GraphExec : public amd::ReferenceCountedObject {
static amd::Monitor graphExecSetLock_;
uint64_t flags_ = 0;
bool repeatLaunch_ = false;
// Graph Kernel arg vars
bool device_kernarg_pool_ = false;
address kernarg_pool_graph_ = nullptr;
uint32_t kernarg_pool_size_graph_ = 0;
uint32_t kernarg_pool_cur_graph_offset_ = 0;
std::vector<address> kernarg_graph_;
uint32_t kernarg_graph_cur_offset_ = 0;
uint32_t kernarg_graph_size_ = 128 * Ki;
bool device_kernarg_pool_ = false; //! Indicate if kernel pool in device mem
std::vector<KernelArgPoolGraph> kernarg_graph_; //! Vector of allocated kernarg pool
uint32_t kernarg_graph_cur_offset_ = 0; //! Current offset in kernarg pool
int instantiateDeviceId_ = -1;
bool hasHiddenHeap_ = false; //!< Kernel has hidden heap(device side allocation)
bool hasHiddenHeap_ = false; //!< Hidden heap indicator for Kernel node
public:
GraphExec(std::vector<Node>& topoOrder, std::vector<std::vector<Node>>& lists,
@@ -589,21 +597,17 @@ struct GraphExec : public amd::ReferenceCountedObject {
}
~GraphExec() {
// new commands are launched for every launch they are destroyed as and when command is
// terminated after it complete execution
for (auto stream : parallel_streams_) {
if (stream != nullptr) {
hip::Stream::Destroy(stream);
}
}
// Release the kernel arg memory.
auto device = g_devices[ihipGetDevice()]->devices()[0];
//! Release the kernel arg pools
if (DEBUG_CLR_GRAPH_PACKET_CAPTURE) {
if (kernarg_pool_size_graph_ != 0) {
device->hostFree(kernarg_pool_graph_, kernarg_pool_size_graph_);
for (auto& element : kernarg_graph_) {
device->hostFree(element, kernarg_graph_size_);
}
auto device = g_devices[ihipGetDevice()]->devices()[0];
for (auto& element : kernarg_graph_) {
device->hostFree(element.kernarg_pool_addr_, element.kernarg_pool_size_);
}
}
amd::ScopedLock lock(graphExecSetLock_);
@@ -620,25 +624,29 @@ struct GraphExec : public amd::ReferenceCountedObject {
}
return clonedNode;
}
// returns if graph has nodes that require hidden heap/not
//! Check if kernel node has hidden heap
bool HasHiddenHeap() const { return hasHiddenHeap_; }
// Graph has nodes that require hidden heap.
//! Graph has nodes that require hidden heap.
void SetHiddenHeap() { hasHiddenHeap_ = true; }
address allocKernArg(size_t size, size_t alignment) {
assert(alignment != 0);
address result = nullptr;
result = amd::alignUp(kernarg_pool_graph_ + kernarg_pool_cur_graph_offset_, alignment);
const size_t pool_new_usage = (result + size) - kernarg_pool_graph_;
if (pool_new_usage <= kernarg_pool_size_graph_) {
kernarg_pool_cur_graph_offset_ = pool_new_usage;
result = amd::alignUp(kernarg_graph_.back().kernarg_pool_addr_ +
kernarg_graph_.back().kernarg_pool_offset_,
alignment);
const size_t pool_new_usage = (result + size) - kernarg_graph_.back().kernarg_pool_addr_;
if (pool_new_usage <= kernarg_graph_.back().kernarg_pool_size_) {
kernarg_graph_.back().kernarg_pool_offset_ = pool_new_usage;
} else {
return nullptr;
}
return result;
}
// check executable graphs validity
//! Check executable graphs validity
static bool isGraphExecValid(GraphExec* pGraphExec);
hipError_t AllocGraphKernargPool(size_t pool_size);
std::vector<Node>& GetNodes() { return topoOrder_; }
hip::Stream* GetAvailableStreams() {