From dbba14f823691cdb996daeed8f645f14bb356988 Mon Sep 17 00:00:00 2001 From: Sean Keely Date: Thu, 29 Jul 2021 02:50:45 -0500 Subject: [PATCH] Improve HSA_CU_MASK parsing efficiency. Delay parsing until after GPU discovery. Use the surfaced GPU count and maximum phyiscal CU count to limit parsed bit masks. This prevents pathological input such as HSA_CU_MASK=0-8000000:0-8000000 from attempting to consume 7TiB. Change-Id: I3773d2db3740c2023b0f6275d1818b69119b0495 [ROCm/ROCR-Runtime commit: 2aa0795b3351a4c0ab7a34b64f4562abf7c47a6a] --- .../hsa-runtime/core/runtime/amd_topology.cpp | 10 ++++++++++ .../runtime/hsa-runtime/core/util/flag.cpp | 15 ++++++++++----- .../runtime/hsa-runtime/core/util/flag.h | 8 +++++--- 3 files changed, 25 insertions(+), 8 deletions(-) diff --git a/projects/rocr-runtime/runtime/hsa-runtime/core/runtime/amd_topology.cpp b/projects/rocr-runtime/runtime/hsa-runtime/core/runtime/amd_topology.cpp index 001b58feb0..db6a71dc19 100644 --- a/projects/rocr-runtime/runtime/hsa-runtime/core/runtime/amd_topology.cpp +++ b/projects/rocr-runtime/runtime/hsa-runtime/core/runtime/amd_topology.cpp @@ -355,6 +355,16 @@ void BuildTopology() { // Instantiate ROCr objects to encapsulate Gpu devices SurfaceGpuList(gpu_usr_list, xnack_mode); + + // Parse HSA_CU_MASK with GPU and CU count limits. + uint32_t maxGpu = core::Runtime::runtime_singleton_->gpu_agents().size(); + uint32_t maxCu = 0; + uint32_t cus; + for (auto& gpu : core::Runtime::runtime_singleton_->gpu_agents()) { + gpu->GetInfo((hsa_agent_info_t)HSA_AMD_AGENT_INFO_COMPUTE_UNIT_COUNT, &cus); + maxCu = Max(maxCu, cus); + } + const_cast(core::Runtime::runtime_singleton_->flag()).parse_masks(maxGpu, maxCu); } bool Load() { diff --git a/projects/rocr-runtime/runtime/hsa-runtime/core/util/flag.cpp b/projects/rocr-runtime/runtime/hsa-runtime/core/util/flag.cpp index 8a8c6e6e00..f45ea9992d 100644 --- a/projects/rocr-runtime/runtime/hsa-runtime/core/util/flag.cpp +++ b/projects/rocr-runtime/runtime/hsa-runtime/core/util/flag.cpp @@ -67,7 +67,7 @@ static std::vector split(std::string& str, char sep) { }; // Parse id,id-id,... strings into id lists -static std::vector get_elements(std::string& str) { +static std::vector get_elements(std::string& str, uint32_t maxElement) { std::vector ret; MAKE_NAMED_SCOPE_GUARD(error, [&]() { ret.clear(); }); @@ -81,12 +81,13 @@ static std::vector get_elements(std::string& str) { uint32_t index = strtoul(range[0].c_str(), &end, 10); // Invalid syntax - id's must be base 10 digits only. if (*end != '\0') return ret; - ret.push_back(index); + if (index <= maxElement) ret.push_back(index); if (range.size() == 2) { uint32_t secondindex = strtoul(range[1].c_str(), &end, 10); if (*end != '\0') return ret; // bad syntax if (secondindex < index) return ret; // inverted range + secondindex = Min(secondindex, maxElement); for (uint32_t i = index + 1; i < secondindex + 1; i++) ret.push_back(i); } } @@ -122,7 +123,7 @@ Specifying a mask with no usable CUs (CU_list is 0x0) is a syntax error. Users should use ROCM_VISIBLE_DEVICES if they want to exclude use of a particular GPU. */ -void Flag::parse_masks(std::string& var) { +void Flag::parse_masks(std::string& var, uint32_t maxGpu, uint32_t maxCU) { if (var.empty()) return; // Remove whitespace @@ -165,6 +166,10 @@ void Flag::parse_masks(std::string& var) { mask.push_back(chunk); } + // Trim dwords beyond maxCUs + uint32_t maxDwords = maxCU / 32 + 1; + if (maxDwords < mask.size()) mask.resize(maxDwords); + // Trim leading zeros while (!mask.empty() && mask.back() == 0) mask.pop_back(); @@ -173,7 +178,7 @@ void Flag::parse_masks(std::string& var) { } else { // parse cu lists - auto cu_indices = get_elements(parts[1]); + auto cu_indices = get_elements(parts[1], maxCU); if (cu_indices.empty()) return; uint32_t maxdword = cu_indices.back() / 32 + 1; mask.resize(maxdword, 0); @@ -186,7 +191,7 @@ void Flag::parse_masks(std::string& var) { } // parse device list - gpu_index = get_elements(parts[0]); + gpu_index = get_elements(parts[0], maxGpu); if (gpu_index.empty()) return; // Ensure that no GPU was repeated across cu_sets diff --git a/projects/rocr-runtime/runtime/hsa-runtime/core/util/flag.h b/projects/rocr-runtime/runtime/hsa-runtime/core/util/flag.h index b8d996c1b8..cd395c667d 100644 --- a/projects/rocr-runtime/runtime/hsa-runtime/core/util/flag.h +++ b/projects/rocr-runtime/runtime/hsa-runtime/core/util/flag.h @@ -151,9 +151,11 @@ class Flag { var = os::GetEnvVar("HSA_ENABLE_DEBUG"); debug_ = (var == "1") ? true : false; + } - var = os::GetEnvVar("HSA_CU_MASK"); - parse_masks(var); + void parse_masks(uint32_t maxGpu, uint32_t maxCU) { + std::string var = os::GetEnvVar("HSA_CU_MASK"); + parse_masks(var, maxGpu, maxCU); } bool check_flat_scratch() const { return check_flat_scratch_; } @@ -258,7 +260,7 @@ class Flag { // Map GPU index post RVD to its default cu mask. std::map> cu_mask_; - void parse_masks(std::string& args); + void parse_masks(std::string& args, uint32_t maxGpu, uint32_t maxCU); DISALLOW_COPY_AND_ASSIGN(Flag); };