SWDEV-255979 - Add support for dynamic __managed__ variables

Change-Id: I62b790853ea3ab3b7ac57bab389046c627fdecce


[ROCm/hip commit: d9fffacfb3]

projects/hip/rocclr/hip_code_object.cpp

@@ -33,6 +33,9 @@ THE SOFTWARE.
 
 hipError_t ihipMemcpy(void* dst, const void* src, size_t sizeBytes, hipMemcpyKind kind,
                       amd::HostQueue& queue, bool isAsync = false);
+hipError_t ihipFree(void* ptr);
+//forward declaration of methods required for managed variables
+hipError_t ihipMallocManaged(void** ptr, size_t size, unsigned int align = 0);
 namespace {
 size_t constexpr strLiteralLength(char const* str) {
   return *str ? 1 + strLiteralLength(str + 1) : 0;
@@ -514,6 +517,9 @@ DynCO::~DynCO() {
   amd::ScopedLock lock(dclock_);
 
   for (auto& elem : vars_) {
+    if(elem.second->getVarKind() == Var::DVK_Managed) {
+      ihipFree(elem.second->getManagedVarPtr());
+    }
     delete elem.second;
   }
   vars_.clear();
@@ -560,16 +566,69 @@ hipError_t DynCO::getDynFunc(hipFunction_t* hfunc, std::string func_name) {
   return it->second->getDynFunc(hfunc, module());
 }
 
+hipError_t DynCO::initDynManagedVars(const std::string& managedVar) {
+  amd::ScopedLock lock(dclock_);
+  DeviceVar* dvar;
+  void* pointer = nullptr;
+  hipError_t status = hipSuccess;
+  // To get size of the managed variable
+  status = getDeviceVar(&dvar, managedVar + ".managed");
+  if (status != hipSuccess) {
+    ClPrint(amd::LOG_ERROR, amd::LOG_API, "Status %d, failed to get .managed device variable:%s",
+            status, managedVar.c_str());
+    return status;
+  }
+  // Allocate managed memory for these symbols
+  status = ihipMallocManaged(&pointer, dvar->size());
+  if (status != hipSuccess) {
+    ClPrint(amd::LOG_ERROR, amd::LOG_API, "Status %d, failed to allocate managed memory", status);
+    guarantee(false, "Error during allocation of managed memory!");
+  }
+  // update as manager variable and set managed memory pointer and size
+  auto it = vars_.find(managedVar);
+  it->second->setManagedVarInfo(pointer, dvar->size());
+
+  // copy initial value to the managed variable to the managed memory allocated
+  amd::HostQueue* queue = hip::getNullStream();
+  if (queue != nullptr) {
+    status = ihipMemcpy(pointer, reinterpret_cast<address>(dvar->device_ptr()), dvar->size(),
+                        hipMemcpyDeviceToDevice, *queue);
+    if (status != hipSuccess) {
+      ClPrint(amd::LOG_ERROR, amd::LOG_API, "Status %d, failed to copy device ptr:%s", status,
+              managedVar.c_str());
+      return status;
+    }
+  } else {
+    ClPrint(amd::LOG_ERROR, amd::LOG_API, "Host Queue is NULL");
+    return hipErrorInvalidResourceHandle;
+  }
+
+  // Get deivce ptr to initialize with managed memory pointer
+  status = getDeviceVar(&dvar, managedVar);
+  if (status != hipSuccess) {
+    ClPrint(amd::LOG_ERROR, amd::LOG_API, "Status %d, failed to get managed device variable:%s",
+            status, managedVar.c_str());
+    return status;
+  }
+  // copy managed memory pointer to the managed device variable
+  status = ihipMemcpy(reinterpret_cast<address>(dvar->device_ptr()), &pointer, dvar->size(),
+                      hipMemcpyHostToDevice, *queue);
+  if (status != hipSuccess) {
+    ClPrint(amd::LOG_ERROR, amd::LOG_API, "Status %d, failed to copy device ptr:%s", status,
+            managedVar.c_str());
+    return status;
+  }
+  return status;
+}
+
 hipError_t DynCO::populateDynGlobalVars() {
   amd::ScopedLock lock(dclock_);
-
   std::vector<std::string> var_names;
   std::vector<std::string> undef_var_names;
-
-  //For Dynamic Modules there is only one hipFatBinaryDevInfo_
-  device::Program* dev_program
-    = fb_info_->GetProgram(ihipGetDevice())->getDeviceProgram
-                          (*hip::getCurrentDevice()->devices()[0]);
+  std::string managedVarExt = ".managed";
+  // For Dynamic Modules there is only one hipFatBinaryDevInfo_
+  device::Program* dev_program = fb_info_->GetProgram(ihipGetDevice())
+                                     ->getDeviceProgram(*hip::getCurrentDevice()->devices()[0]);
 
   if (!dev_program->getGlobalVarFromCodeObj(&var_names)) {
     LogPrintfError("Could not get Global vars from Code Obj for Module: 0x%x \n", module());
@@ -577,9 +636,17 @@ hipError_t DynCO::populateDynGlobalVars() {
   }
 
   for (auto& elem : var_names) {
-    vars_.insert(std::make_pair(elem, new Var(elem, Var::DeviceVarKind::DVK_Variable, 0, 0, 0, nullptr)));
+    vars_.insert(
+        std::make_pair(elem, new Var(elem, Var::DeviceVarKind::DVK_Variable, 0, 0, 0, nullptr)));
   }
 
+  for (auto& elem : var_names) {
+    if (elem.find(managedVarExt) != std::string::npos) {
+      std::string managedVar = elem;
+      managedVar.erase(managedVar.length() - managedVarExt.length(), managedVarExt.length());
+      initDynManagedVars(managedVar);
+    }
+  }
   return hipSuccess;
 }
 
@@ -661,6 +728,7 @@ hipError_t StatCO::removeFatBinary(FatBinaryInfo** module) {
   auto it = managedVars_.begin();
   while (it != managedVars_.end()) {
     if ((*it)->moduleInfo() == module) {
+      ihipFree((*it)->getManagedVarPtr());
       delete *it;
       managedVars_.erase(it);
     } else {

projects/hip/rocclr/hip_code_object.hpp

@@ -92,6 +92,14 @@ public:
   hipError_t getDynFunc(hipFunction_t* hfunc, std::string func_name);
   hipError_t getDeviceVar(DeviceVar** dvar, std::string var_name);
 
+  hipError_t getManagedVarPointer(std::string name, void** pointer, size_t* size_ptr) const {
+    auto it = vars_.find(name);
+    if (it != vars_.end() && it->second->getVarKind() == Var::DVK_Managed) {
+      *pointer = it->second->getManagedVarPtr();
+      *size_ptr = it->second->getSize();
+    }
+    return hipSuccess;
+  }
   // Device ID Check to check if module is launched in the same device it was loaded.
   inline void CheckDeviceIdMatch() {
     if (device_id_ != ihipGetDevice()) {
@@ -110,6 +118,7 @@ private:
   //Populate Global Vars/Funcs from an code object(@ module_load)
   hipError_t populateDynGlobalFuncs();
   hipError_t populateDynGlobalVars();
+  hipError_t initDynManagedVars(const std::string& managedVar);
 };
 
 //Static Code Object

projects/hip/rocclr/hip_global.hpp

@@ -101,7 +101,15 @@ public:
   void resize_dVar(size_t size) { dVar_.resize(size); }
 
   FatBinaryInfo** moduleInfo() { return modules_; };
+  DeviceVarKind getVarKind() const { return dVarKind_; }
+  size_t getSize() const { return size_; }
+  
   void* getManagedVarPtr() { return managedVarPtr_; };
+  void setManagedVarInfo(void* pointer, size_t size) {
+    managedVarPtr_ = pointer;
+    size_ = size;
+    dVarKind_ = DVK_Managed;
+  }
 private:
   std::vector<DeviceVar*> dVar_;   // DeviceVarObj per Device
   std::string name_;               // Variable name (not unique identifier)

projects/hip/rocclr/hip_platform.cpp

@@ -820,12 +820,15 @@ hipError_t PlatformState::getDynGlobalVar(const char* hostVar, hipModule_t hmod,
     LogPrintfError("Cannot find the module: 0x%x", hmod);
     return hipErrorNotFound;
   }
-
-  hip::DeviceVar* dvar = nullptr;
-  IHIP_RETURN_ONFAIL(it->second->getDeviceVar(&dvar, hostVar));
-  *dev_ptr = dvar->device_ptr();
-  *size_ptr = dvar->size();
-
+  *dev_ptr = nullptr;
+  it->second->getManagedVarPointer(hostVar, dev_ptr, size_ptr);
+  // if dev_ptr is nullptr, hostvar is not in managed variable list
+  if (*dev_ptr == nullptr) {
+    hip::DeviceVar* dvar = nullptr;
+    IHIP_RETURN_ONFAIL(it->second->getDeviceVar(&dvar, hostVar));
+    *dev_ptr = dvar->device_ptr();
+    *size_ptr = dvar->size();
+  }
   return hipSuccess;
 }
 

projects/hip/tests/src/runtimeApi/module/hipManagedKeyword.cpp

@@ -0,0 +1,69 @@
+/*
+Copyright (c) 2021-present Advanced Micro Devices, Inc. All rights reserved.
+
+Permission is hereby granted, free of charge, to any person obtaining a copy
+of this software and associated documentation files (the "Software"), to deal
+in the Software without restriction, including without limitation the rights
+to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+copies of the Software, and to permit persons to whom the Software is
+furnished to do so, subject to the following conditions:
+
+The above copyright notice and this permission notice shall be included in
+all copies or substantial portions of the Software.
+
+THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT.  IN NO EVENT SHALL THE
+AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN
+THE SOFTWARE.
+*/
+
+/* HIT_START
+ * BUILD_CMD: managed_kernel.code %hc --genco %S/managed_kernel.cpp -o managed_kernel.code EXCLUDE_HIP_PLATFORM amd
+ * BUILD: %t %s ../../test_common.cpp EXCLUDE_HIP_PLATFORM nvidia EXCLUDE_HIP_PLATFORM amd
+ * TEST: %t
+ * HIT_END
+ */
+
+#include "hip/hip_runtime.h"
+#include <iostream>
+#include "test_common.h"
+
+#define MANAGED_VAR_INIT_VALUE 10
+#define fileName "managed_kernel.code"
+
+bool managedMultiGPUTest() {
+  int numDevices = 0;
+  hipDeviceptr_t x;
+  size_t xSize;
+  int data;
+  hipGetDeviceCount(&numDevices);
+  for (int i = 0; i < numDevices; i++) {
+    hipSetDevice(i);
+    hipModule_t Module;
+    HIPCHECK(hipModuleLoad(&Module, fileName));
+    hipFunction_t Function;
+    HIPCHECK(hipModuleGetFunction(&Function, Module, "GPU_func"));
+    HIPCHECK(hipModuleLaunchKernel(Function, 1, 1, 1, 1, 1, 1, 0, 0, NULL, NULL));
+    hipDeviceSynchronize();
+    HIPCHECK(hipModuleGetGlobal((hipDeviceptr_t*)&x, &xSize, Module, "x"));
+    HIPCHECK(hipMemcpyDtoH(&data, hipDeviceptr_t(x), xSize));
+    if (data != (1 + MANAGED_VAR_INIT_VALUE)) {
+      HIPCHECK(hipModuleUnload(Module));
+      return false;
+    }
+    HIPCHECK(hipModuleUnload(Module));
+  }
+  return true;
+}
+
+int main(int argc, char** argv) {
+  hipInit(0);
+  bool testStatus = managedMultiGPUTest();
+  if (!testStatus) {
+    failed("Managed keyword module test failed!");
+  }
+  passed();
+}

projects/hip/tests/src/runtimeApi/module/managed_kernel.cpp

@@ -0,0 +1,27 @@
+/*
+Copyright (c) 2021-present Advanced Micro Devices, Inc. All rights reserved.
+
+Permission is hereby granted, free of charge, to any person obtaining a copy
+of this software and associated documentation files (the "Software"), to deal
+in the Software without restriction, including without limitation the rights
+to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+copies of the Software, and to permit persons to whom the Software is
+furnished to do so, subject to the following conditions:
+
+The above copyright notice and this permission notice shall be included in
+all copies or substantial portions of the Software.
+
+THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT.  IN NO EVENT SHALL THE
+AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN
+THE SOFTWARE.
+*/
+#include "hip/hip_runtime.h"
+__managed__ int x = 10;
+
+extern "C" __global__ void GPU_func() {
+  x++;
+}