🐛 Fix device id

yhs0602 · Jan 1, 2025 · a547031 · a547031
1 parent 6342f3a
commit a547031
Show file tree

Hide file tree

Showing 4 changed files with 37 additions and 19 deletions.
diff --git a/src/cpp/ipc_cuda.cpp b/src/cpp/ipc_cuda.cpp
@@ -1,5 +1,6 @@
 #include "ipc_cuda.h"
 #include "dlpack.h"
+#include <cassert>
 #include <cstdlib>
 #include <stdexcept>
 #include <string>
@@ -33,20 +34,18 @@ mtl_tensor_from_cuda_ipc_handle(void *cuda_ipc_handle, int width, int height) {
             std::string(cudaGetErrorString(err))
         );
     }
+    // cuda_ipc_handle : sizeof(cudaIpcMemHandle_t)
+    // after the data (sizeof(int)): device_id concatenated
+    cudaIpcMemHandle_t *handle = (cudaIpcMemHandle_t *)cuda_ipc_handle;
+    int device_id = *(int *)(handle + 1);
 
     void *device_ptr = nullptr;
-    int success_device_id = 0;
-    for (; success_device_id < deviceCount; ++success_device_id) {
-        cudaSetDevice(success_device_id);
-        err= cudaIpcOpenMemHandle(
-            &device_ptr,
-            *reinterpret_cast<cudaIpcMemHandle_t *>(cuda_ipc_handle),
-            cudaIpcMemLazyEnablePeerAccess
-        );
-        if (err == cudaSuccess) {
-            break;
-        }
-    }
+    cudaSetDevice(device_id);
+    err= cudaIpcOpenMemHandle(
+        &device_ptr,
+        *reinterpret_cast<cudaIpcMemHandle_t *>(cuda_ipc_handle),
+        cudaIpcMemLazyEnablePeerAccess
+    );
     if (err != cudaSuccess || device_ptr == nullptr) {
         throw std::runtime_error(
             "Failed to open CUDA IPC handle: " +
@@ -94,9 +93,21 @@ mtl_tensor_from_cuda_ipc_handle(void *cuda_ipc_handle, int width, int height) {
         );
     }
 
-    int device_id;
     if (attributes.devicePointer != nullptr) {
-        device_id = attributes.device;
+        if (attributes.type == cudaMemoryTypeDevice) {
+            printf("Memory type: Device\n");
+        } else if (attributes.type == cudaMemoryTypeHost) {
+            printf("Memory type: Host\n");
+        } else if (attributes.type == cudaMemoryTypeManaged) {
+            printf("Memory type: Managed\n");
+        } else {
+            printf("Memory type: Unknown\n");
+        }
+        if (attributes.device != device_id) {
+            free(tensor->dl_tensor.shape);
+            free(tensor);
+            throw std::runtime_error("Device ID mismatch: Attribute=" + std::to_string(attributes.device) + "!= actual=" + std::to_string(device_id));
+        }
         printf("\nOpen tensor from ipc handle: Device ID: %d\n", device_id);
         fflush(stdout);
     } else {
@@ -108,7 +119,7 @@ mtl_tensor_from_cuda_ipc_handle(void *cuda_ipc_handle, int width, int height) {
 
     tensor->dl_tensor.dtype =
         (DLDataType){kDLUInt, 8, 1}; // Unsigned 8-bit integer
-    tensor->dl_tensor.device = (DLDevice){kDLCUDA, device_id}; // cuda gpu  + 1
+    tensor->dl_tensor.device = (DLDevice){kDLCUDA, device_id};
 
     tensor->deleter = deleteDLManagedTensor;
     return tensor;

diff --git a/src/craftground/MinecraftEnv/src/main/cpp/framebuffer_capturer.cpp b/src/craftground/MinecraftEnv/src/main/cpp/framebuffer_capturer.cpp
@@ -484,8 +484,9 @@ Java_com_kyhsgeekcode_minecraftenv_FramebufferCapturer_initializeZerocopyImpl(
     }
 
     cudaIpcMemHandle_t memHandle;
+    int deviceId = -1;
     int size = initialize_cuda_ipc(
-        width, height, colorAttachment, depthAttachment, &memHandle
+        width, height, colorAttachment, depthAttachment, &memHandle, &deviceId
     );
 
     if (size < 0) {
@@ -496,7 +497,7 @@ Java_com_kyhsgeekcode_minecraftenv_FramebufferCapturer_initializeZerocopyImpl(
         return nullptr;
     }
 
-    jbyteArray byteArray = env->NewByteArray(size);
+    jbyteArray byteArray = env->NewByteArray(size + sizeof(int));
     if (byteArray == nullptr || env->ExceptionCheck()) {
         // Handle error
         return nullptr;
@@ -505,6 +506,9 @@ Java_com_kyhsgeekcode_minecraftenv_FramebufferCapturer_initializeZerocopyImpl(
     env->SetByteArrayRegion(
         byteArray, 0, size, reinterpret_cast<jbyte *>(&memHandle)
     );
+    env->SetByteArrayRegion(
+        byteArray, size, sizeof(int), reinterpret_cast<jbyte *>(&deviceId)
+    );
     jobject byteStringObject =
         env->CallStaticObjectMethod(byteStringClass, copyFromMethod, byteArray);
     if (byteStringObject == nullptr || env->ExceptionCheck()) {

diff --git a/src/craftground/MinecraftEnv/src/main/cpp/framebuffer_capturer_cuda.cpp b/src/craftground/MinecraftEnv/src/main/cpp/framebuffer_capturer_cuda.cpp
@@ -14,7 +14,8 @@ int initialize_cuda_ipc(
     int height,
     int colorAttachment,
     int depthAttachment,
-    cudaIpcMemHandle_t *memHandlePtr
+    cudaIpcMemHandle_t *memHandlePtr,
+    int *deviceId
 ) {
     if (initialized) {
         fprintf(stderr, "CUDA IPC already initialized\n");
@@ -74,6 +75,7 @@ int initialize_cuda_ipc(
         return -1;
     }
     initialized = true;
+    *deviceId = rendering_gpu;
     fprintf(stdout, "\n\nInitialized CUDA IPC: %p\n\n", sharedCudaColorMem);
     fflush(stdout);
     return sizeof(cudaIpcMemHandle_t);

diff --git a/src/craftground/MinecraftEnv/src/main/cpp/framebuffer_capturer_cuda.h b/src/craftground/MinecraftEnv/src/main/cpp/framebuffer_capturer_cuda.h
@@ -12,7 +12,8 @@ int initialize_cuda_ipc(
     int height,
     int colorAttachment,
     int depthAttachment,
-    cudaIpcMemHandle_t *memHandlePtr
+    cudaIpcMemHandle_t *memHandlePtr,
+    int *deviceId
 );
 
 void copyFramebufferToCudaSharedMemory(int width, int height);