Hook into the creaky lazy-reinit machinery. Try to minimise race-risk.

[ROCm/clr commit: 3470692998]

projects/clr/hipamd/include/hip/hcc_detail/functional_grid_launch.hpp

@@ -103,18 +103,23 @@ inline std::vector<std::uint8_t> make_kernarg(
 
     if (sizeof...(Formals) == 0) return {};
 
-    const auto it = function_names().find(
-        reinterpret_cast<std::uintptr_t>(kernel));
-
+    auto it = function_names().find(reinterpret_cast<std::uintptr_t>(kernel));
     if (it == function_names().cend()) {
-        throw std::runtime_error{"Undefined __global__ function."};
+        it =
+            function_names(true).find(reinterpret_cast<std::uintptr_t>(kernel));
+        if (it == function_names().cend()) {
+            throw std::runtime_error{"Undefined __global__ function."};
+        }
     }
 
-    const auto it1 = kernargs().find(it->second);
-
+    auto it1 = kernargs().find(it->second);
     if (it1 == kernargs().end()) {
-        throw std::runtime_error{
-            "Missing metadata for __global__ function: " + it->second};
+        it1 = kernargs(true).find(it->second);
+
+        if (it1 == kernargs().end()) {
+            throw std::runtime_error{
+                "Missing metadata for __global__ function: " + it->second};
+        }
     }
 
     std::tuple<Formals...> to_formals{std::move(actuals)};

projects/clr/hipamd/include/hip/hcc_detail/program_state.hpp

@@ -99,8 +99,9 @@ const std::unordered_map<std::uintptr_t, std::vector<std::pair<hsa_agent_t, Kern
 functions(bool rebuild = false);
 const std::unordered_map<std::uintptr_t, std::string>& function_names(bool rebuild = false);
 std::unordered_map<std::string, void*>& globals(bool rebuild = false);
-std::unordered_map<
-    std::string, std::vector<std::pair<std::size_t, std::size_t>>>& kernargs();
+const std::unordered_map<
+    std::string, std::vector<std::pair<std::size_t, std::size_t>>>&
+    kernargs(bool rebuild = false);
 
 hsa_executable_t load_executable(const std::string& file, hsa_executable_t executable,
                                  hsa_agent_t agent);

projects/clr/hipamd/src/program_state.cpp

@@ -538,6 +538,7 @@ const unordered_map<uintptr_t, vector<pair<hsa_agent_t, Kernel_descriptor>>>& fu
             // created previously
 
             function_names(rebuild);
+            kernargs(rebuild);
             kernels(rebuild);
             globals(rebuild);
         }
@@ -585,12 +586,12 @@ unordered_map<string, void*>& globals(bool rebuild) {
     return r;
 }
 
-
-unordered_map<string, vector<pair<size_t, size_t>>>& kernargs() {
+const unordered_map<string, vector<pair<size_t, size_t>>>& kernargs(
+    bool rebuild) {
     static unordered_map<string, vector<pair<size_t, size_t>>> r;
     static once_flag f;
 
-    call_once(f, []() {
+    static const auto build_map = [](decltype(r)& x) {
         for (auto&& isa_blobs : code_object_blobs()) {
             for (auto&& blob : isa_blobs.second) {
                 stringstream tmp{std::string{blob.cbegin(), blob.cend()}};
@@ -598,10 +599,28 @@ unordered_map<string, vector<pair<size_t, size_t>>>& kernargs() {
                 elfio reader;
                 if (!reader.load(tmp)) continue;
 
-                read_kernarg_metadata(reader, r);
+                read_kernarg_metadata(reader, x);
             }
         }
-    });
+    };
+    call_once(f, []() { r.reserve(function_names().size()); build_map(r); });
+
+    if (rebuild) {
+        static mutex mtx;
+        thread_local static decltype(r) tmp;
+
+        {
+            lock_guard<mutex> lck{mtx};
+
+            tmp.insert(r.cbegin(), r.cend()); // Should use merge in C++17.
+        }
+
+        build_map(tmp);
+
+        lock_guard<mutex> lck{mtx};
+
+        r.insert(tmp.cbegin(), tmp.cend());
+    }
 
     return r;
 }