From cc63d2d7b86304ea78ca44d709b03ee2fb74a807 Mon Sep 17 00:00:00 2001
From: Christophe Paquot <christophe.paquot@amd.com>
Date: Thu, 13 Feb 2020 10:44:10 -0800
Subject: [PATCH] hipLaunchByPtr and hipLaunchKernel deviceId potential issue

Those APIs should look at the device associated with the stream first.
If that stream is null then get the current device ID.

Change-Id: Iedde1d1644818ba64f128b988f0bd9674f5b8ad6
---
 vdi/hip_internal.hpp |  3 ++-
 vdi/hip_platform.cpp | 12 +++++++-----
 vdi/hip_stream.cpp   |  6 +++---
 3 files changed, 12 insertions(+), 9 deletions(-)

diff --git a/vdi/hip_internal.hpp b/vdi/hip_internal.hpp
index e39d8f3f4c..377d5bef9d 100644
--- a/vdi/hip_internal.hpp
+++ b/vdi/hip_internal.hpp
@@ -102,8 +102,9 @@ namespace hip {
     amd::Context* context;
     amd::CommandQueue::Priority priority;
     unsigned int flags;
+    int deviceId;
 
-    Stream(amd::Device* dev, amd::Context* ctx, amd::CommandQueue::Priority p, unsigned int f);
+    Stream(amd::Device* dev, amd::Context* ctx, amd::CommandQueue::Priority p, unsigned int f, int d);
     void create();
     amd::HostQueue* asHostQueue();
     void destroy();
diff --git a/vdi/hip_platform.cpp b/vdi/hip_platform.cpp
index a33390091c..11a5c043c6 100644
--- a/vdi/hip_platform.cpp
+++ b/vdi/hip_platform.cpp
@@ -536,7 +536,11 @@ extern "C" hipError_t hipLaunchByPtr(const void *hostFunction)
 {
   HIP_INIT_API(NONE, hostFunction);
 
-  int deviceId = ihipGetDevice();
+  ihipExec_t exec;
+  PlatformState::instance().popExec(exec);
+
+  hip::Stream* stream = reinterpret_cast<hip::Stream*>(exec.hStream_);
+  int deviceId = (stream != nullptr)? stream->deviceId : ihipGetDevice();
   if (deviceId == -1) {
     HIP_RETURN(hipErrorNoDevice);
   }
@@ -545,9 +549,6 @@ extern "C" hipError_t hipLaunchByPtr(const void *hostFunction)
     HIP_RETURN(hipErrorInvalidDeviceFunction);
   }
 
-  ihipExec_t exec;
-  PlatformState::instance().popExec(exec);
-
   size_t size = exec.arguments_.size();
   void *extra[] = {
       HIP_LAUNCH_PARAM_BUFFER_POINTER, &exec.arguments_[0],
@@ -940,7 +941,8 @@ extern "C" hipError_t hipLaunchKernel(const void *hostFunction,
   HIP_INIT_API(NONE, hostFunction, gridDim, blockDim, args, sharedMemBytes,
                stream);
 
-  int deviceId = ihipGetDevice();
+  hip::Stream* s = reinterpret_cast<hip::Stream*>(stream);
+  int deviceId = (s != nullptr)? s->deviceId : ihipGetDevice();
   if (deviceId == -1) {
     HIP_RETURN(hipErrorNoDevice);
   }
diff --git a/vdi/hip_stream.cpp b/vdi/hip_stream.cpp
index c182fc6de9..bfc47e5748 100644
--- a/vdi/hip_stream.cpp
+++ b/vdi/hip_stream.cpp
@@ -50,8 +50,8 @@ void syncStreams() {
   }
 }
 
-Stream::Stream(amd::Device* dev, amd::Context* ctx, amd::CommandQueue::Priority p, unsigned int f) :
-  queue(nullptr), device(dev), context(ctx), priority(p), flags(f) {}
+Stream::Stream(amd::Device* dev, amd::Context* ctx, amd::CommandQueue::Priority p, unsigned int f, int d) :
+  queue(nullptr), device(dev), context(ctx), priority(p), flags(f), deviceId(d) {}
 
 void Stream::create() {
   cl_command_queue_properties properties = CL_QUEUE_PROFILING_ENABLE;
@@ -95,7 +95,7 @@ void CL_CALLBACK ihipStreamCallback(cl_event event, cl_int command_exec_status,
 static hipError_t ihipStreamCreate(hipStream_t *stream, unsigned int flags, amd::CommandQueue::Priority priority) {
   amd::Device* device = hip::getCurrentContext()->devices()[0];
 
-  hip::Stream* hStream = new hip::Stream(device, hip::getCurrentContext(), priority, flags);
+  hip::Stream* hStream = new hip::Stream(device, hip::getCurrentContext(), priority, flags, ihipGetDevice());
 
   if (hStream == nullptr) {
     return hipErrorOutOfMemory;