From b4d8586a1134ffc45c22e787d9cbad5c53f86845 Mon Sep 17 00:00:00 2001
From: Rakesh Roy <rakesh.roy@amd.com>
Date: Thu, 24 Nov 2022 19:43:40 +0530
Subject: [PATCH] SWDEV-368881 - Enable hipLaunchCooperativeKernel for hipRTC
 compiled function

- Add helper function to validate whether a dynamic function is present or not
- If getStatFunc fails, then check whether the function passed is a valid dynamic function

Change-Id: I4141172f42dce7ff689709560c483e602b5dad4b
---
 hipamd/src/hip_code_object.cpp | 14 ++++++++++++++
 hipamd/src/hip_code_object.hpp |  3 +++
 hipamd/src/hip_global.cpp      |  8 ++++++++
 hipamd/src/hip_global.hpp      |  2 ++
 hipamd/src/hip_module.cpp      | 10 +++++++++-
 hipamd/src/hip_platform.cpp    | 12 ++++++++++++
 hipamd/src/hip_platform.hpp    |  2 ++
 7 files changed, 50 insertions(+), 1 deletion(-)

diff --git a/hipamd/src/hip_code_object.cpp b/hipamd/src/hip_code_object.cpp
index 7070c4e5aa..8e76f0a947 100644
--- a/hipamd/src/hip_code_object.cpp
+++ b/hipamd/src/hip_code_object.cpp
@@ -562,6 +562,20 @@ hipError_t DynCO::getDynFunc(hipFunction_t* hfunc, std::string func_name) {
   return it->second->getDynFunc(hfunc, module());
 }
 
+bool DynCO::isValidDynFunc(const hipFunction_t& hfunc) {
+  amd::ScopedLock lock(dclock_);
+
+  CheckDeviceIdMatch();
+
+  for (auto it : functions_) {
+    if (it.second->isValidDynFunc(hfunc)) {
+      return true;
+    }
+  }
+
+  return false;
+}
+
 hipError_t DynCO::initDynManagedVars(const std::string& managedVar) {
   amd::ScopedLock lock(dclock_);
   DeviceVar* dvar;
diff --git a/hipamd/src/hip_code_object.hpp b/hipamd/src/hip_code_object.hpp
index efbfbc4d19..e79ec2ce34 100644
--- a/hipamd/src/hip_code_object.hpp
+++ b/hipamd/src/hip_code_object.hpp
@@ -92,6 +92,9 @@ public:
   hipError_t getDynFunc(hipFunction_t* hfunc, std::string func_name);
   hipError_t getDeviceVar(DeviceVar** dvar, std::string var_name);
 
+  // Verify whether the dynamic function is present or not
+  bool isValidDynFunc(const hipFunction_t& hfunc);
+
   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) {
diff --git a/hipamd/src/hip_global.cpp b/hipamd/src/hip_global.cpp
index 5759cb715e..31943769fc 100644
--- a/hipamd/src/hip_global.cpp
+++ b/hipamd/src/hip_global.cpp
@@ -135,6 +135,14 @@ hipError_t Function::getDynFunc(hipFunction_t* hfunc, hipModule_t hmod) {
   return hipSuccess;
 }
 
+bool Function::isValidDynFunc(const hipFunction_t& hfunc) {
+  guarantee((dFunc_.size() == g_devices.size()), "dFunc Size mismatch");
+  if (dFunc_[ihipGetDevice()] == nullptr) {
+    return false;
+  }
+  return (dFunc_[ihipGetDevice()]->asHipFunction() == hfunc);
+}
+
 hipError_t Function::getStatFunc(hipFunction_t* hfunc, int deviceId) {
   guarantee(modules_ != nullptr, "Module not initialized");
 
diff --git a/hipamd/src/hip_global.hpp b/hipamd/src/hip_global.hpp
index f769e85af2..93f5cf2d99 100644
--- a/hipamd/src/hip_global.hpp
+++ b/hipamd/src/hip_global.hpp
@@ -62,6 +62,8 @@ public:
 
   //Return DeviceFunc for this this dynamically loaded module
   hipError_t getDynFunc(hipFunction_t* hfunc, hipModule_t hmod);
+  // Verify whether the dynamic device function is present or not
+  bool isValidDynFunc(const hipFunction_t& hfunc);
 
   //Return Device Func & attr . Generate/build if not already done so.
   hipError_t getStatFunc(hipFunction_t *hfunc, int deviceId);
diff --git a/hipamd/src/hip_module.cpp b/hipamd/src/hip_module.cpp
index 5deb099e57..1840830025 100644
--- a/hipamd/src/hip_module.cpp
+++ b/hipamd/src/hip_module.cpp
@@ -497,7 +497,15 @@ hipError_t hipLaunchCooperativeKernel_common(const void* f, dim3 gridDim, dim3 b
 
   hipFunction_t func = nullptr;
   int deviceId = hip::Stream::DeviceId(hStream);
-  HIP_RETURN_ONFAIL(PlatformState::instance().getStatFunc(&func, f, deviceId));
+  hipError_t status = PlatformState::instance().getStatFunc(&func, f, deviceId);
+  if (status != hipSuccess) {
+    // Check if its a dynamic function
+    if (!PlatformState::instance().isValidDynFunc(hipFunction_t(f))) {
+      return status;
+    }
+    func = (hipFunction_t)f;
+  }
+
   size_t globalWorkSizeX = static_cast<size_t>(gridDim.x) * blockDim.x;
   size_t globalWorkSizeY = static_cast<size_t>(gridDim.y) * blockDim.y;
   size_t globalWorkSizeZ = static_cast<size_t>(gridDim.z) * blockDim.z;
diff --git a/hipamd/src/hip_platform.cpp b/hipamd/src/hip_platform.cpp
index 8f532361bc..8f221b48e7 100644
--- a/hipamd/src/hip_platform.cpp
+++ b/hipamd/src/hip_platform.cpp
@@ -747,6 +747,18 @@ hipError_t PlatformState::getDynFunc(hipFunction_t* hfunc, hipModule_t hmod,
   return it->second->getDynFunc(hfunc, func_name);
 }
 
+bool PlatformState::isValidDynFunc(const hipFunction_t& hfunc) {
+  amd::ScopedLock lock(lock_);
+
+  for (auto it : dynCO_map_) {
+    if (it.second->isValidDynFunc(hfunc)) {
+      return true;
+    }
+  }
+
+  return false;
+}
+
 hipError_t PlatformState::getDynGlobalVar(const char* hostVar, hipModule_t hmod,
                                           hipDeviceptr_t* dev_ptr, size_t* size_ptr) {
   amd::ScopedLock lock(lock_);
diff --git a/hipamd/src/hip_platform.hpp b/hipamd/src/hip_platform.hpp
index 109a921547..600c1a9178 100644
--- a/hipamd/src/hip_platform.hpp
+++ b/hipamd/src/hip_platform.hpp
@@ -47,6 +47,8 @@ class PlatformState {
   hipError_t unloadModule(hipModule_t hmod);
 
   hipError_t getDynFunc(hipFunction_t* hfunc, hipModule_t hmod, const char* func_name);
+  // Verify whether the dynamic function is present or not
+  bool isValidDynFunc(const hipFunction_t& hfunc);
   hipError_t getDynGlobalVar(const char* hostVar, hipModule_t hmod, hipDeviceptr_t* dev_ptr,
                              size_t* size_ptr);
   hipError_t getDynTexRef(const char* hostVar, hipModule_t hmod, textureReference** texRef);