From 58aa4f61c2c36615eb1184a355ce274a3c93657b Mon Sep 17 00:00:00 2001
From: Rahul Garg <rahul.garg@amd.com>
Date: Fri, 26 Aug 2016 19:03:23 +0530
Subject: [PATCH] Added logic to update primary ctx when ctx stack is empty,
 updated hipCtxDestroy and ctxGetCurrent functions

Change-Id: Ia0a8943c121bc1279788a1cfa9be59af614b04a6


[ROCm/hip commit: 1211cc931ce75da600ea05aca5a919797a86a1ff]
---
 projects/hip/include/hcc_detail/hip_hcc.h |  2 ++
 projects/hip/src/hip_context.cpp          | 38 ++++++++++++++++++++---
 2 files changed, 35 insertions(+), 5 deletions(-)

diff --git a/projects/hip/include/hcc_detail/hip_hcc.h b/projects/hip/include/hcc_detail/hip_hcc.h
index e2a4d709f2..48f7dac318 100644
--- a/projects/hip/include/hcc_detail/hip_hcc.h
+++ b/projects/hip/include/hcc_detail/hip_hcc.h
@@ -166,6 +166,7 @@ class ihipCtx_t;
 // generate trace string that can be output to stderr or to ATP file.
 #define HIP_INIT_API(...) \
 	std::call_once(hip_initialized, ihipInit);\
+    ihipCtxStackUpdate();\
     API_TRACE(__VA_ARGS__);
 
 #define ihipLogStatus(hipStatus) \
@@ -666,6 +667,7 @@ extern const char *ihipErrorString(hipError_t);
 extern ihipCtx_t    *ihipGetTlsDefaultCtx();
 extern void          ihipSetTlsDefaultCtx(ihipCtx_t *ctx);
 extern hipError_t    ihipSynchronize(void);
+extern hipError_t    ihipCtxStackUpdate();
 
 extern ihipDevice_t *ihipGetDevice(int);
 ihipCtx_t * ihipGetPrimaryCtx(unsigned deviceIndex);
diff --git a/projects/hip/src/hip_context.cpp b/projects/hip/src/hip_context.cpp
index 33caf610bc..c0c98de827 100644
--- a/projects/hip/src/hip_context.cpp
+++ b/projects/hip/src/hip_context.cpp
@@ -29,6 +29,18 @@ THE SOFTWARE.
 // Stack of contexts
 thread_local std::stack<ihipCtx_t *>  tls_ctxStack;
 
+hipError_t ihipCtxStackUpdate()
+{
+    //HIP_INIT_API();
+    printf("Inside ihipCtxStackUpdate\n");
+    hipError_t e = hipSuccess;
+
+    if(tls_ctxStack.empty()) {
+            tls_ctxStack.push(ihipGetTlsDefaultCtx());
+    }
+
+    return ihipLogStatus(e);
+}
 
 hipError_t hipInit(unsigned int flags)
 {
@@ -92,11 +104,20 @@ hipError_t hipCtxDestroy(hipCtx_t ctx)
 {
     hipError_t e = hipSuccess;
     ihipCtx_t* currentCtx= ihipGetTlsDefaultCtx();
-    if(currentCtx == ctx) {
-        //need to destroy the ctx associated with calling thread
-        tls_ctxStack.pop();
+    ihipCtx_t* primaryCtx= ((ihipDevice_t*)ctx->getDevice())->_primaryCtx;
+    if(primaryCtx== ctx)
+    {
+        e = hipErrorInvalidValue;
     }
-    delete ctx; //As per CUDA docs , attempting to access ctx from those threads which has this ctx as current, will result in the error HIP_ERROR_CONTEXT_IS_DESTROYED.
+    else
+    {
+        if(currentCtx == ctx) {
+            //need to destroy the ctx associated with calling thread
+            tls_ctxStack.pop();
+        }
+        delete ctx; //As per CUDA docs , attempting to access ctx from those threads which has this ctx as current, will result in the error HIP_ERROR_CONTEXT_IS_DESTROYED.
+    }
+
     return ihipLogStatus(e);
 }
 
@@ -135,11 +156,18 @@ hipError_t hipCtxPushCurrent(hipCtx_t ctx)
 hipError_t hipCtxGetCurrent(hipCtx_t* ctx)
 {
     hipError_t e = hipSuccess;
-
+#if 0
     *ctx = ihipGetTlsDefaultCtx();
     if(*ctx == nullptr) {
         *ctx = NULL;    //TODO - is it required? Can return nullptr?
     }
+#endif
+    if(!tls_ctxStack.empty()) {
+        *ctx= tls_ctxStack.top();
+    }
+    else {
+        *ctx = NULL;
+    }
     return ihipLogStatus(e);
 }