From b380a95597c8a2ab9d15cf5e7c590eeb44b17d47 Mon Sep 17 00:00:00 2001
From: Ajay <ajay.gunashekar@amd.com>
Date: Thu, 4 Aug 2022 00:14:55 +0000
Subject: [PATCH] SWDEV-348678 - hipManagedVar to initialize on all devices

Change-Id: Id714b56e89e4930ee067f3a7472a0029057d2711
---
 hipamd/src/hip_code_object.cpp |  2 +-
 hipamd/src/hip_global.cpp      | 17 +++++++++++------
 hipamd/src/hip_global.hpp      |  2 +-
 hipamd/src/hip_module.cpp      |  5 ++++-
 4 files changed, 17 insertions(+), 9 deletions(-)

diff --git a/hipamd/src/hip_code_object.cpp b/hipamd/src/hip_code_object.cpp
index 3442eb42a5..ef92c8b96b 100644
--- a/hipamd/src/hip_code_object.cpp
+++ b/hipamd/src/hip_code_object.cpp
@@ -843,7 +843,7 @@ hipError_t StatCO::initStatManagedVarDevicePtr(int deviceId) {
       DeviceVar* dvar = nullptr;
       IHIP_RETURN_ONFAIL(var->getStatDeviceVar(&dvar, deviceId));
 
-      amd::HostQueue* queue = hip::getNullStream();
+      amd::HostQueue* queue = g_devices.at(deviceId)->NullStream();
       if (queue != nullptr) {
         err = ihipMemcpy(reinterpret_cast<address>(dvar->device_ptr()), var->getManagedVarPtr(),
                          dvar->size(), hipMemcpyHostToDevice, *queue);
diff --git a/hipamd/src/hip_global.cpp b/hipamd/src/hip_global.cpp
index 75a6822fe5..5759cb715e 100644
--- a/hipamd/src/hip_global.cpp
+++ b/hipamd/src/hip_global.cpp
@@ -40,11 +40,16 @@ void __hipGetPCH(const char** pch, unsigned int *size) {
 namespace hip {
 
 //Device Vars
-DeviceVar::DeviceVar(std::string name, hipModule_t hmod) : shadowVptr(nullptr), name_(name),
-                                                           amd_mem_obj_(nullptr), device_ptr_(nullptr),
-                                                           size_(0) {
+DeviceVar::DeviceVar(std::string name,
+                     hipModule_t hmod,
+                     int deviceId) :
+                     shadowVptr(nullptr), name_(name),
+                     amd_mem_obj_(nullptr), device_ptr_(nullptr),
+                     size_(0) {
   amd::Program* program = as_amd(reinterpret_cast<cl_program>(hmod));
-  device::Program* dev_program = program->getDeviceProgram(*hip::getCurrentDevice()->devices()[0]);
+  device::Program* dev_program =
+                   program->getDeviceProgram(*g_devices.at(deviceId)->devices()[0]);
+
   if (dev_program == nullptr) {
     LogPrintfError("Cannot get Device Program for module: 0x%x \n", hmod);
     guarantee(false, "Cannot get Device Program");
@@ -206,7 +211,7 @@ hipError_t Var::getDeviceVar(DeviceVar** dvar, int deviceId, hipModule_t hmod) {
              "Device Var not initialized to size");
 
   if (dVar_[deviceId] == nullptr) {
-    dVar_[deviceId] = new DeviceVar(name_, hmod);
+    dVar_[deviceId] = new DeviceVar(name_, hmod, deviceId);
   }
 
   *dvar = dVar_[deviceId];
@@ -221,7 +226,7 @@ hipError_t Var::getStatDeviceVar(DeviceVar** dvar, int deviceId) {
     hipModule_t hmod = nullptr;
     IHIP_RETURN_ONFAIL((*modules_)->BuildProgram(deviceId));
     IHIP_RETURN_ONFAIL((*modules_)->GetModule(deviceId, &hmod));
-    dVar_[deviceId] = new DeviceVar(name_, hmod);
+    dVar_[deviceId] = new DeviceVar(name_, hmod, deviceId);
   }
   *dvar = dVar_[deviceId];
   return hipSuccess;
diff --git a/hipamd/src/hip_global.hpp b/hipamd/src/hip_global.hpp
index 00fffd6a1d..f769e85af2 100644
--- a/hipamd/src/hip_global.hpp
+++ b/hipamd/src/hip_global.hpp
@@ -18,7 +18,7 @@ class CodeObject;
 //Device Structures
 class DeviceVar {
 public:
-  DeviceVar(std::string name, hipModule_t hmod);
+  DeviceVar(std::string name, hipModule_t hmod, int deviceId);
   ~DeviceVar();
 
   //Accessors for device ptr and size, populated during constructor.
diff --git a/hipamd/src/hip_module.cpp b/hipamd/src/hip_module.cpp
index 4518f4589e..7921de3ac9 100644
--- a/hipamd/src/hip_module.cpp
+++ b/hipamd/src/hip_module.cpp
@@ -343,7 +343,10 @@ hipError_t ihipModuleLaunchKernel(hipFunction_t f, uint32_t globalWorkSizeX,
                startEvent, stopEvent, flags, params);
 
   int deviceId = hip::Stream::DeviceId(hStream);
-  HIP_RETURN_ONFAIL(PlatformState::instance().initStatManagedVarDevicePtr(deviceId));
+  for (size_t dev = 0; dev < g_devices.size(); ++dev) {
+    HIP_RETURN_ONFAIL(PlatformState::instance().initStatManagedVarDevicePtr(dev));
+  }
+  
   if (f == nullptr) {
     LogPrintfError("%s", "Function passed is null");
     return hipErrorInvalidImage;