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: 2aa0795b33]

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<Flag&>(core::Runtime::runtime_singleton_->flag()).parse_masks(maxGpu, maxCu);
 }
 
 bool Load() {

projects/rocr-runtime/runtime/hsa-runtime/core/util/flag.cpp

@@ -67,7 +67,7 @@ static std::vector<std::string> split(std::string& str, char sep) {
 };
 
 // Parse id,id-id,... strings into id lists
-static std::vector<uint32_t> get_elements(std::string& str) {
+static std::vector<uint32_t> get_elements(std::string& str, uint32_t maxElement) {
   std::vector<uint32_t> ret;
   MAKE_NAMED_SCOPE_GUARD(error, [&]() { ret.clear(); });
 
@@ -81,12 +81,13 @@ static std::vector<uint32_t> 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

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<uint32_t, std::vector<uint32_t>> 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);
 };