From b983c1972978583b1519fa8e8725bdda378c429e Mon Sep 17 00:00:00 2001
From: Konstantin Zhuravlyov <kzhuravl_dev@outlook.com>
Date: Wed, 3 Apr 2024 16:58:23 -0400
Subject: [PATCH] Switch to per-executable contexts in the loader

  - Per-executable contexts should be used from now on
  - Global contexts are left as is for now for backwards
    compatibility and will be phased out in follow up
    patches.

Change-Id: I6291abf865c7ed24ee71f5065e539afc23f5ce64
---
 .../hsa-runtime/core/inc/amd_hsa_loader.hpp   |  8 ++++++
 runtime/hsa-runtime/core/runtime/hsa.cpp      |  1 +
 runtime/hsa-runtime/loader/executable.cpp     | 28 +++++++++++++++++++
 runtime/hsa-runtime/loader/executable.hpp     | 13 +++++++++
 4 files changed, 50 insertions(+)

diff --git a/runtime/hsa-runtime/core/inc/amd_hsa_loader.hpp b/runtime/hsa-runtime/core/inc/amd_hsa_loader.hpp
index 72d55afd67..c63b7a9612 100644
--- a/runtime/hsa-runtime/core/inc/amd_hsa_loader.hpp
+++ b/runtime/hsa-runtime/core/inc/amd_hsa_loader.hpp
@@ -50,6 +50,7 @@
 #include "inc/hsa_ven_amd_loader.h"
 #include "inc/amd_hsa_elf.h"
 #include <string>
+#include <memory>
 #include <mutex>
 #include <vector>
 
@@ -457,6 +458,13 @@ public:
       const char *options,
       hsa_default_float_rounding_mode_t default_float_rounding_mode = HSA_DEFAULT_FLOAT_ROUNDING_MODE_DEFAULT) = 0;
 
+  /// @brief Creates empty AMD HSA Executable with specified @p profile,
+  /// @p options and @p isolated_context that is isolated from the runtime.
+  virtual Executable* CreateExecutable(
+      std::unique_ptr<Context> isolated_context,
+      hsa_profile_t profile,
+      const char *options,
+      hsa_default_float_rounding_mode_t default_float_rounding_mode = HSA_DEFAULT_FLOAT_ROUNDING_MODE_DEFAULT) = 0;
 
   /// @brief Freezes @p executable
   virtual hsa_status_t FreezeExecutable(Executable *executable, const char *options) = 0;
diff --git a/runtime/hsa-runtime/core/runtime/hsa.cpp b/runtime/hsa-runtime/core/runtime/hsa.cpp
index c509fda5ec..6c161cd0a1 100644
--- a/runtime/hsa-runtime/core/runtime/hsa.cpp
+++ b/runtime/hsa-runtime/core/runtime/hsa.cpp
@@ -2195,6 +2195,7 @@ hsa_status_t hsa_executable_create_alt(
   IS_BAD_PTR(executable);
 
   Executable *exec = GetLoader()->CreateExecutable(
+      std::unique_ptr<amd::LoaderContext>(new amd::LoaderContext()),
       profile, options, default_float_rounding_mode);
   CHECK_ALLOC(exec);
 
diff --git a/runtime/hsa-runtime/loader/executable.cpp b/runtime/hsa-runtime/loader/executable.cpp
index e251a05448..150a263650 100644
--- a/runtime/hsa-runtime/loader/executable.cpp
+++ b/runtime/hsa-runtime/loader/executable.cpp
@@ -188,6 +188,18 @@ Executable* AmdHsaCodeLoader::CreateExecutable(
   return executables.back();
 }
 
+Executable* AmdHsaCodeLoader::CreateExecutable(
+      std::unique_ptr<Context> isolated_context,
+      hsa_profile_t profile,
+      const char *options,
+      hsa_default_float_rounding_mode_t default_float_rounding_mode)
+{
+  WriterLockGuard<ReaderWriterLock> writer_lock(rw_lock_);
+
+  executables.push_back(new ExecutableImpl(profile, std::move(isolated_context), executables.size(), default_float_rounding_mode));
+  return executables.back();
+}
+
 static void AddCodeObjectInfoIntoDebugMap(link_map* map) {
   if (r_debug_tail) {
       r_debug_tail->l_next = map;
@@ -740,6 +752,22 @@ ExecutableImpl::ExecutableImpl(
 {
 }
 
+ExecutableImpl::ExecutableImpl(
+    const hsa_profile_t &_profile,
+    std::unique_ptr<Context> unique_context,
+    size_t id,
+    hsa_default_float_rounding_mode_t default_float_rounding_mode)
+  : Executable()
+  , profile_(_profile)
+  , unique_context_(std::move(unique_context))
+  , id_(id)
+  , default_float_rounding_mode_(default_float_rounding_mode)
+  , state_(HSA_EXECUTABLE_STATE_UNFROZEN)
+  , program_allocation_segment(nullptr)
+{
+  context_ = unique_context_.get();
+}
+
 ExecutableImpl::~ExecutableImpl() {
   for (ExecutableObject* o : objects) {
     o->Destroy();
diff --git a/runtime/hsa-runtime/loader/executable.hpp b/runtime/hsa-runtime/loader/executable.hpp
index cb985ca3e1..9429ff9486 100644
--- a/runtime/hsa-runtime/loader/executable.hpp
+++ b/runtime/hsa-runtime/loader/executable.hpp
@@ -423,6 +423,12 @@ public:
       size_t id,
       hsa_default_float_rounding_mode_t default_float_rounding_mode);
 
+  ExecutableImpl(
+      const hsa_profile_t &_profile,
+      std::unique_ptr<Context> unique_context,
+      size_t id,
+      hsa_default_float_rounding_mode_t default_float_rounding_mode);
+
   ~ExecutableImpl();
 
   hsa_status_t GetInfo(hsa_executable_info_t executable_info, void *value) override;
@@ -549,6 +555,7 @@ private:
   amd::hsa::common::ReaderWriterLock rw_lock_;
   hsa_profile_t profile_;
   Context *context_;
+  std::unique_ptr<Context> unique_context_;
   Logger logger_;
   const size_t id_;
   hsa_default_float_rounding_mode_t default_float_rounding_mode_;
@@ -578,6 +585,12 @@ public:
       const char *options,
       hsa_default_float_rounding_mode_t default_float_rounding_mode = HSA_DEFAULT_FLOAT_ROUNDING_MODE_DEFAULT) override;
 
+  Executable* CreateExecutable(
+      std::unique_ptr<Context> isolated_context,
+      hsa_profile_t profile,
+      const char *options,
+      hsa_default_float_rounding_mode_t default_float_rounding_mode = HSA_DEFAULT_FLOAT_ROUNDING_MODE_DEFAULT) override;
+
   hsa_status_t FreezeExecutable(Executable *executable, const char *options) override;
   void DestroyExecutable(Executable *executable) override;