Fix the handling for empty GPU tensor (triton-inference-server#187)

krishung5 · web-flow · commit fdf08afea27f · 2022-09-28T17:42:26.000-07:00
* Fix the handling for empty GPU tensor

* Fix start_address initialization
diff --git a/src/pb_memory.cc b/src/pb_memory.cc
@@ -295,11 +295,15 @@ PbMemory::GetGPUStartAddress()
 {
   if (memory_shm_ptr_->memory_type == TRITONSERVER_MEMORY_GPU) {
     CUDAHandler& cuda_api = CUDAHandler::getInstance();
-    CUdeviceptr start_address;
-
-    cuda_api.PointerGetAttribute(
-        &start_address, CU_POINTER_ATTRIBUTE_RANGE_START_ADDR,
-        reinterpret_cast<CUdeviceptr>(data_ptr_));
+    CUdeviceptr start_address = 0;
+
+    // Skip this step for empty tensor as the CUDA API 'cuPointerGetAttribute'
+    // we use in this function does not accept nullptr.
+    if (data_ptr_) {
+      cuda_api.PointerGetAttribute(
+          &start_address, CU_POINTER_ATTRIBUTE_RANGE_START_ADDR,
+          reinterpret_cast<CUdeviceptr>(data_ptr_));
+    }
 
     return reinterpret_cast<void*>(start_address);
   }
