From 6611cc015d3b1a26449b38632ac4fdfbfea5b35d Mon Sep 17 00:00:00 2001
From: Saleel Kudchadker <Saleel.Kudchadker@amd.com>
Date: Tue, 5 Nov 2024 22:52:57 +0000
Subject: [PATCH] SWDEV-494149 - Improve hipGet/Set Device

Change-Id: If8975687a3ba9caadafc48a0066f19a4ebaab9e2
---
 hipamd/src/hip_context.cpp        |  7 ++++---
 hipamd/src/hip_device_runtime.cpp | 23 +++++++++++++++--------
 hipamd/src/hip_internal.hpp       | 12 ++++++------
 3 files changed, 25 insertions(+), 17 deletions(-)

diff --git a/hipamd/src/hip_context.cpp b/hipamd/src/hip_context.cpp
index 2a10a2aa58..784805c277 100644
--- a/hipamd/src/hip_context.cpp
+++ b/hipamd/src/hip_context.cpp
@@ -1,4 +1,4 @@
-/* Copyright (c) 2015 - 2021 Advanced Micro Devices, Inc.
+/* Copyright (c) 2015 - 2024 Advanced Micro Devices, Inc.
 
  Permission is hereby granted, free of charge, to any person obtaining a copy
  of this software and associated documentation files (the "Software"), to deal
@@ -48,10 +48,11 @@ void init(bool* status) {
   }
   ClPrint(amd::LOG_INFO, amd::LOG_INIT, "Direct Dispatch: %d", AMD_DIRECT_DISPATCH);
 
-
   const std::vector<amd::Device*>& devices = amd::Device::getDevices(CL_DEVICE_TYPE_GPU, false);
+  const size_t deviceCount = devices.size();
+  g_devices.reserve(deviceCount);  // Pre-allocate space for better performance
 
-  for (unsigned int i=0; i<devices.size(); i++) {
+  for (unsigned int i = 0; i < deviceCount; i++) {
     // Enable active wait on the device by default
     devices[i]->SetActiveWait(true);
     // use the eternal contexts that already exist for new hip::Device's here
diff --git a/hipamd/src/hip_device_runtime.cpp b/hipamd/src/hip_device_runtime.cpp
index bf757fe0f1..1f07c1b4e9 100644
--- a/hipamd/src/hip_device_runtime.cpp
+++ b/hipamd/src/hip_device_runtime.cpp
@@ -656,16 +656,17 @@ int ihipGetDevice() {
 hipError_t hipGetDevice(int* deviceId) {
   HIP_INIT_API(hipGetDevice, deviceId);
 
-  if (deviceId != nullptr) {
-    int dev = ihipGetDevice();
-    if (dev == -1) {
-      HIP_RETURN(hipErrorNoDevice);
-    }
-    *deviceId = dev;
-    HIP_RETURN(hipSuccess, *deviceId);
-  } else {
+  if (deviceId == nullptr) {
     HIP_RETURN(hipErrorInvalidValue);
   }
+
+  Device* device = hip::getCurrentDevice();
+  if (device == nullptr) {
+    HIP_RETURN(hipErrorNoDevice);
+  }
+
+  *deviceId = device->deviceId();
+  HIP_RETURN(hipSuccess, *deviceId);
 }
 
 hipError_t hipGetDeviceCount(int* count) {
@@ -685,6 +686,12 @@ hipError_t hipGetDeviceFlags(unsigned int* flags) {
 
 hipError_t hipSetDevice(int device) {
   HIP_INIT_API_NO_RETURN(hipSetDevice, device);
+
+  // Check if the device is already set
+  if (hip::tls.device_ != nullptr && hip::tls.device_->deviceId() == device) {
+    HIP_RETURN(hipSuccess);
+  }
+
   if (static_cast<unsigned int>(device) < g_devices.size()) {
     hip::setCurrentDevice(device);
 
diff --git a/hipamd/src/hip_internal.hpp b/hipamd/src/hip_internal.hpp
index 8d7bea1ae5..6a30b19635 100644
--- a/hipamd/src/hip_internal.hpp
+++ b/hipamd/src/hip_internal.hpp
@@ -112,7 +112,7 @@ const char* ihipGetErrorName(hipError_t hip_error);
 #define HIP_INIT(noReturn)                                                                         \
   {                                                                                                \
     bool status = true;                                                                            \
-    std::call_once(hip::g_ihipInitialized, hip::init, &status);                                         \
+    std::call_once(hip::g_ihipInitialized, hip::init, &status);                                    \
     if (!status && !noReturn) {                                                                    \
       HIP_RETURN(hipErrorInvalidDevice);                                                           \
     }                                                                                              \
@@ -125,16 +125,16 @@ const char* ihipGetErrorName(hipError_t hip_error);
 #define HIP_INIT_VOID()                                                                            \
   {                                                                                                \
     bool status = true;                                                                            \
-    std::call_once(hip::g_ihipInitialized, hip::init, &status);                                         \
-    if (hip::tls.device_ == nullptr && hip::g_devices.size() > 0) {                           \
-      hip::tls.device_ = hip::g_devices[0];                                                   \
-      amd::Os::setPreferredNumaNode(hip::g_devices[0]->devices()[0]->getPreferredNumaNode()); \
+    std::call_once(hip::g_ihipInitialized, hip::init, &status);                                    \
+    if (hip::tls.device_ == nullptr && hip::g_devices.size() > 0) {                                \
+      hip::tls.device_ = hip::g_devices[0];                                                        \
+      amd::Os::setPreferredNumaNode(hip::g_devices[0]->devices()[0]->getPreferredNumaNode());      \
     }                                                                                              \
   }
 
 
 #define HIP_API_PRINT(...)                                          \
-  uint64_t startTimeUs=0;                                           \
+  uint64_t startTimeUs = 0;                                         \
   HIPPrintDuration(amd::LOG_INFO, amd::LOG_API, &startTimeUs,       \
                   "%s %s ( %s ) %s", KGRN,                          \
                   __func__, ToString( __VA_ARGS__ ).c_str(), KNRM);