fix(hotswap): gate entry trampoline rewrites for gfx12.5

projects/rocr-runtime/rocrtst/suites/functional/hotswap/hotswap_gfx_query_test.cc

@@ -33,6 +33,8 @@ const char* kGfx1250IsaWithFeatures =
     "amdgcn-amd-amdhsa--gfx1250:sramecc+:xnack-";
 const char* kGfx942Isa = "amdgcn-amd-amdhsa--gfx942";
 const char* kGfx1251Isa = "amdgcn-amd-amdhsa--gfx1251";
+const char* kGfx12_5GenericIsaWithFeatures =
+    "amdgcn-amd-amdhsa--gfx12-5-generic:sramecc+";
 
 void ResetTestEnv() {
   g_fake_hsa_env = FakeHsaEnv{};
@@ -97,6 +99,7 @@ namespace {
 
 using rocr::hotswap::AgentGfxRevision;
 using rocr::hotswap::GetAgentGfxRevision;
+using rocr::hotswap::IsGfx12_5Target;
 using rocr::hotswap::IsHotswapSupportedGfxRevision;
 
 TEST(HotswapGfxQuery, Gfx1250A0Passes) {
@@ -145,6 +148,26 @@ TEST(HotswapGfxQuery, NearMissTargetBlocks) {
   EXPECT_FALSE(IsHotswapSupportedGfxRevision(revision));
 }
 
+TEST(HotswapGfxQuery, Gfx12_5GenericFeatureSuffixParsed) {
+  ResetTestEnv();
+  g_fake_hsa_env.isa_name = kGfx12_5GenericIsaWithFeatures;
+  g_fake_hsa_env.asic_revision = 0;
+
+  const AgentGfxRevision revision = GetAgentGfxRevision(MakeFreshAgent());
+
+  EXPECT_EQ(revision.gfx_target, "gfx12-5-generic");
+  EXPECT_FALSE(IsHotswapSupportedGfxRevision(revision));
+}
+
+TEST(HotswapGfxQuery, Gfx12_5TargetPredicateIsStrict) {
+  EXPECT_TRUE(IsGfx12_5Target("gfx1250"));
+  EXPECT_TRUE(IsGfx12_5Target("gfx1251"));
+  EXPECT_TRUE(IsGfx12_5Target("gfx12-5-generic"));
+  EXPECT_FALSE(IsGfx12_5Target("gfx125"));
+  EXPECT_FALSE(IsGfx12_5Target("gfx125foo"));
+  EXPECT_FALSE(IsGfx12_5Target("gfx942"));
+}
+
 TEST(HotswapGfxQuery, Gfx1250NonA0Blocks) {
   ResetTestEnv();
   g_fake_hsa_env.isa_name = kGfx1250Isa;

projects/rocr-runtime/rocrtst/suites/functional/hotswap/hotswap_rewrite_test.cc

@@ -31,6 +31,14 @@
 namespace {
 
 constexpr const char* kGfx1250Isa = "amdgcn-amd-amdhsa--gfx1250";
+constexpr const char* kGfx1251Isa = "amdgcn-amd-amdhsa--gfx1251";
+constexpr const char* kGfx12_5GenericIsa =
+    "amdgcn-amd-amdhsa--gfx12-5-generic";
+constexpr const char* kGfx942Isa = "amdgcn-amd-amdhsa--gfx942";
+constexpr const char* kGfx1250B0Isa =
+    "amdgcn-amd-amdhsa--gfx1250:gfx1250-b0-specific+";
+constexpr const char* kGfx1250A0Isa =
+    "amdgcn-amd-amdhsa--gfx1250:gfx1250-b0-specific-";
 
 struct FakeHsaEnv {
   std::string isa_name = kGfx1250Isa;
@@ -50,7 +58,11 @@ LibHandle LoadLib(std::string filename) {
 #if defined(_WIN32) || defined(_WIN64)
   return LoadLibraryA(filename.c_str());
 #else
-  return dlopen(filename.c_str(), RTLD_NOW | RTLD_LOCAL);
+  int flags = RTLD_LAZY;
+#ifdef RTLD_NODELETE
+  flags |= RTLD_NODELETE;
+#endif
+  return dlopen(filename.c_str(), flags);
 #endif
 }
 
@@ -210,6 +222,77 @@ rocr::hotswap::CodeObjectView MakeRealCodeObjectView() {
   return code_object;
 }
 
+rocr::hotswap::AgentGfxRevision MakeRevision(const std::string& gfx_target,
+                                             uint32_t asic_revision,
+                                             bool has_asic_revision = true) {
+  rocr::hotswap::AgentGfxRevision revision;
+  revision.gfx_target = gfx_target;
+  revision.asic_revision = asic_revision;
+  revision.has_asic_revision = has_asic_revision;
+  return revision;
+}
+
+TEST(HotswapRewriteDecision, A0RetargetsWhenEntryTrampolinesDisabled) {
+  const auto decision = rocr::hotswap::DecideHotswapRewriteForTesting(
+      MakeRevision("gfx1250", 0), kGfx1250Isa, kGfx1250Isa, {false});
+
+  ASSERT_TRUE(decision.has_value());
+  EXPECT_EQ(decision->source_isa, kGfx1250B0Isa);
+  EXPECT_EQ(decision->target_isa, kGfx1250A0Isa);
+  EXPECT_FALSE(decision->request_entry_trampolines);
+}
+
+TEST(HotswapRewriteDecision, EntryTrampolinesDefaultOnRoutesNonA0Gfx1250) {
+  const auto decision = rocr::hotswap::DecideHotswapRewriteForTesting(
+      MakeRevision("gfx1250", 1), kGfx1250Isa, kGfx1250Isa, {});
+
+  ASSERT_TRUE(decision.has_value());
+  EXPECT_EQ(decision->source_isa, kGfx1250B0Isa);
+  EXPECT_EQ(decision->target_isa, kGfx1250B0Isa);
+  EXPECT_TRUE(decision->request_entry_trampolines);
+}
+
+TEST(HotswapRewriteDecision, EntryTrampolinesDisabledBlocksNonA0Gfx1250) {
+  const auto decision = rocr::hotswap::DecideHotswapRewriteForTesting(
+      MakeRevision("gfx1250", 1), kGfx1250Isa, kGfx1250Isa, {false});
+
+  EXPECT_FALSE(decision.has_value());
+}
+
+TEST(HotswapRewriteDecision, EntryTrampolinesRouteGfx12_5Family) {
+  const auto concrete = rocr::hotswap::DecideHotswapRewriteForTesting(
+      MakeRevision("gfx1251", 1), kGfx1251Isa, kGfx1251Isa, {});
+  const auto generic = rocr::hotswap::DecideHotswapRewriteForTesting(
+      MakeRevision("gfx12-5-generic", 1), kGfx12_5GenericIsa,
+      kGfx12_5GenericIsa, {});
+
+  ASSERT_TRUE(concrete.has_value());
+  EXPECT_EQ(concrete->source_isa, kGfx1251Isa);
+  EXPECT_EQ(concrete->target_isa, kGfx1251Isa);
+  EXPECT_TRUE(concrete->request_entry_trampolines);
+  ASSERT_TRUE(generic.has_value());
+  EXPECT_EQ(generic->source_isa, kGfx12_5GenericIsa);
+  EXPECT_EQ(generic->target_isa, kGfx12_5GenericIsa);
+  EXPECT_TRUE(generic->request_entry_trampolines);
+}
+
+TEST(HotswapRewriteDecision, EntryTrampolinesUseGenericSourceAsTarget) {
+  const auto decision = rocr::hotswap::DecideHotswapRewriteForTesting(
+      MakeRevision("gfx1251", 1), kGfx12_5GenericIsa, kGfx1251Isa, {});
+
+  ASSERT_TRUE(decision.has_value());
+  EXPECT_EQ(decision->source_isa, kGfx12_5GenericIsa);
+  EXPECT_EQ(decision->target_isa, kGfx12_5GenericIsa);
+  EXPECT_TRUE(decision->request_entry_trampolines);
+}
+
+TEST(HotswapRewriteDecision, EntryTrampolinesBlockNonGfx12_5) {
+  const auto decision = rocr::hotswap::DecideHotswapRewriteForTesting(
+      MakeRevision("gfx942", 0), kGfx942Isa, kGfx942Isa, {});
+
+  EXPECT_FALSE(decision.has_value());
+}
+
 TEST(HotswapRewrite, GetIsaNameRealCodeObject) {
   const std::string isa =
       rocr::hotswap::GetCodeObjectIsaName(kGfx1250MinCo, sizeof(kGfx1250MinCo));
@@ -321,12 +404,36 @@ TEST(HotswapRewrite, RuntimeLoadUsesRewrittenCodeObject) {
       rocr::hotswap::RetainedRewrittenElfBufferCountForTesting(executable), 0u);
 }
 
-TEST(HotswapRewrite, RuntimeLoadNonA0FallsBackToOriginal) {
+TEST(HotswapRewrite, RuntimeLoadNonA0UsesDefaultEntryTrampolines) {
   ResetRuntimeTestEnv();
   g_fake_hsa_env.asic_revision = 1;
   LoadRecorder load;
   const hsa_executable_t executable = MakeTestExecutable(0x502);
 
+  const hsa_status_t status = rocr::hotswap::LoadAgentCodeObjectWithHotswap(
+      executable, MakeTestAgent(), MakeRealCodeObjectView(), nullptr, nullptr,
+      MakeLoadCallbacks(&load));
+
+  EXPECT_EQ(status, HSA_STATUS_SUCCESS);
+  ASSERT_EQ(load.calls.size(), 1u);
+  EXPECT_EQ(load.calls[0].path, LoadPath::kRewritten);
+  EXPECT_NE(load.calls[0].code_object, static_cast<const void*>(kGfx1250MinCo));
+  EXPECT_GT(load.calls[0].code_object_size, 0u);
+  EXPECT_EQ(
+      rocr::hotswap::RetainedRewrittenElfBufferCountForTesting(executable), 1u);
+
+  rocr::hotswap::ReleaseRetainedRewrittenElfBuffers(executable);
+  EXPECT_EQ(
+      rocr::hotswap::RetainedRewrittenElfBufferCountForTesting(executable), 0u);
+}
+
+TEST(HotswapRewrite, RuntimeLoadNonA0FallsBackWhenEntryTrampolinesAreZero) {
+  ResetRuntimeTestEnv();
+  g_fake_hsa_env.asic_revision = 1;
+  g_fake_env_vars["AMD_COMGR_HOTSWAP_ENTRY_TRAMPOLINES"] = "0";
+  LoadRecorder load;
+  const hsa_executable_t executable = MakeTestExecutable(0x503);
+
   const hsa_status_t status = rocr::hotswap::LoadAgentCodeObjectWithHotswap(
       executable, MakeTestAgent(), MakeRealCodeObjectView(), nullptr, nullptr,
       MakeLoadCallbacks(&load));
@@ -339,11 +446,41 @@ TEST(HotswapRewrite, RuntimeLoadNonA0FallsBackToOriginal) {
       rocr::hotswap::RetainedRewrittenElfBufferCountForTesting(executable), 0u);
 }
 
+TEST(HotswapRewrite, RuntimeLoadNonA0UsesEntryTrampolinesUnlessEnvIsZero) {
+  const char* const env_values[] = {"false", ""};
+  uint64_t executable_handle = 0x504;
+  for (const char* env_value : env_values) {
+    SCOPED_TRACE(env_value);
+    ResetRuntimeTestEnv();
+    g_fake_hsa_env.asic_revision = 1;
+    g_fake_env_vars["AMD_COMGR_HOTSWAP_ENTRY_TRAMPOLINES"] = env_value;
+    LoadRecorder load;
+    const hsa_executable_t executable =
+        MakeTestExecutable(executable_handle++);
+
+    const hsa_status_t status = rocr::hotswap::LoadAgentCodeObjectWithHotswap(
+        executable, MakeTestAgent(), MakeRealCodeObjectView(), nullptr, nullptr,
+        MakeLoadCallbacks(&load));
+
+    EXPECT_EQ(status, HSA_STATUS_SUCCESS);
+    ASSERT_EQ(load.calls.size(), 1u);
+    EXPECT_EQ(load.calls[0].path, LoadPath::kRewritten);
+    EXPECT_EQ(
+        rocr::hotswap::RetainedRewrittenElfBufferCountForTesting(executable),
+        1u);
+
+    rocr::hotswap::ReleaseRetainedRewrittenElfBuffers(executable);
+    EXPECT_EQ(
+        rocr::hotswap::RetainedRewrittenElfBufferCountForTesting(executable),
+        0u);
+  }
+}
+
 TEST(HotswapRewrite, RuntimeLoadDisableEnvFallsBackToOriginal) {
   ResetRuntimeTestEnv();
   g_fake_env_vars["HSA_HOTSWAP_DISABLE"] = "1";
   LoadRecorder load;
-  const hsa_executable_t executable = MakeTestExecutable(0x503);
+  const hsa_executable_t executable = MakeTestExecutable(0x506);
 
   const hsa_status_t status = rocr::hotswap::LoadAgentCodeObjectWithHotswap(
       executable, MakeTestAgent(), MakeRealCodeObjectView(), nullptr, nullptr,
@@ -366,7 +503,7 @@ TEST(HotswapRewrite, RuntimeLoadRewriteFailureFallsBackToOriginal) {
   code_object.size = sizeof(fake_elf);
   code_object.uri = "memory://invalid.hsaco";
   LoadRecorder load;
-  const hsa_executable_t executable = MakeTestExecutable(0x504);
+  const hsa_executable_t executable = MakeTestExecutable(0x507);
 
   const hsa_status_t status = rocr::hotswap::LoadAgentCodeObjectWithHotswap(
       executable, MakeTestAgent(), code_object, nullptr, nullptr,
@@ -384,7 +521,7 @@ TEST(HotswapRewrite, RuntimeLoadRewrittenLoadFailureFallsBackToOriginal) {
   ResetRuntimeTestEnv();
   LoadRecorder load;
   load.rewritten_status = HSA_STATUS_ERROR_INVALID_CODE_OBJECT;
-  const hsa_executable_t executable = MakeTestExecutable(0x505);
+  const hsa_executable_t executable = MakeTestExecutable(0x508);
 
   const hsa_status_t status = rocr::hotswap::LoadAgentCodeObjectWithHotswap(
       executable, MakeTestAgent(), MakeRealCodeObjectView(), nullptr, nullptr,

projects/rocr-runtime/rocrtst/suites/test_common/CMakeLists.txt

@@ -802,7 +802,7 @@ if(CMAKE_SYSTEM_NAME STREQUAL "Linux")
             OUTPUT_VARIABLE _comgr_symbols
             ERROR_QUIET)
           if(_comgr_symbol_result EQUAL 0 AND
-             _comgr_symbols MATCHES "amd_comgr_hotswap_rewrite")
+             _comgr_symbols MATCHES "amd_comgr_hotswap_rewrite_with_options")
             list(APPEND HOTSWAP_BUILD_RPATH "${_prefix}/${_libdir}")
             break()
           endif()

projects/rocr-runtime/runtime/docs/data/env_variables.rst

@@ -105,6 +105,12 @@
       - | 0, ``false``, ``off``, ``no``, ``n``, or ``f``: Disable HotSwap diagnostic logging.
         | Any other non-empty value: Enable HotSwap diagnostic logging.
 
+    * - | ``AMD_COMGR_HOTSWAP_ENTRY_TRAMPOLINES``
+        | Controls whether ROCr requests COMGR entry-trampoline HotSwap rewriting for gfx12.5 targets.
+      - ``1``
+      - | 0: Disable entry-trampoline rewrites.
+        | Unset or any other value, including empty: Enable entry-trampoline rewrites for gfx125* and ``gfx12-5-generic`` targets.
+
     * - | ``HSA_ENABLE_DXG_DETECTION``
         | Controls detection of the DXG driver (/dev/dxg) on WSL2.
       - ``1``

projects/rocr-runtime/runtime/hsa-runtime/core/inc/hotswap.hpp

@@ -46,6 +46,7 @@
 #include <cstddef>
 #include <cstdlib>
 #include <memory>
+#include <optional>
 #include <string>
 
 #include "core/inc/amd_hsa_loader.hpp"
@@ -54,6 +55,8 @@
 namespace rocr {
 namespace hotswap {
 
+struct AgentGfxRevision;
+
 using OwnedElfBuffer = std::unique_ptr<void, decltype(&std::free)>;
 
 struct CodeObjectView {
@@ -62,6 +65,16 @@ struct CodeObjectView {
   std::string uri;
 };
 
+struct RewriteOptions {
+  bool gfx12_5_rewrite_enabled = true;
+};
+
+struct RewriteDecision {
+  std::string source_isa;
+  std::string target_isa;
+  bool request_entry_trampolines = false;
+};
+
 using LoadOriginalCodeObjectFn = hsa_status_t (*)(
     void* context, hsa_agent_t agent, hsa_code_object_t code_object,
     const char* options, const std::string& uri,
@@ -82,7 +95,8 @@ std::string GetCodeObjectIsaName(const void* elf_data, size_t elf_size);
 
 bool RetargetCodeObject(const void* elf_data, size_t elf_size,
                         const char* source_isa, const char* target_isa,
-                        OwnedElfBuffer* out_elf_buffer, size_t* out_elf_size);
+                        OwnedElfBuffer* out_elf_buffer, size_t* out_elf_size,
+                        bool request_entry_trampolines = false);
 
 bool TryRetargetCodeObject(const CodeObjectView& code_object, hsa_agent_t agent,
                            OwnedElfBuffer* out_elf_buffer,
@@ -103,6 +117,9 @@ void RetainRewrittenElfBuffer(hsa_executable_t executable,
 void ReleaseRetainedRewrittenElfBuffers(hsa_executable_t executable);
 
 #ifdef ROCR_HOTSWAP_TESTING
+std::optional<RewriteDecision> DecideHotswapRewriteForTesting(
+    const AgentGfxRevision& gfx, const std::string& source_isa,
+    const std::string& target_isa, const RewriteOptions& options);
 size_t RetainedRewrittenElfBufferCountForTesting(hsa_executable_t executable);
 #endif
 

projects/rocr-runtime/runtime/hsa-runtime/core/inc/hotswap_gfx_query.hpp

@@ -62,6 +62,7 @@ struct AgentGfxRevision {
 
 std::string GetAgentIsaName(hsa_agent_t agent);
 std::string ExtractGfxTarget(const std::string& isa_name);
+bool IsGfx12_5Target(const std::string& gfx_target);
 AgentGfxRevision GetAgentGfxRevision(hsa_agent_t agent);
 void ResetAgentGfxRevisionCache();