fix bug for uva when getting the device pointer

f7b4c93d · sangwzh · 50c27a8e · f7b4c93d
Commit f7b4c93d authored Oct 14, 2024 by sangwzh
Show whitespace changes
Inline Side-by-side

Showing with 9 additions and 1 deletion

graphbolt/src/cuda/index_select_impl.hip graphbolt/src/cuda/index_select_impl.hip +9 -1

No files found.
--- a/graphbolt/src/cuda/index_select_impl.hip
+++ b/graphbolt/src/cuda/index_select_impl.hip
@@ -112,7 +112,15 @@ torch::Tensor UVAIndexSelectImpl_(torch::Tensor input, torch::Tensor index) {
      {return_len, original_feature_size}, torch::TensorOptions()
                                               .dtype(input.dtype())
                                               .device(c10::DeviceType::CUDA));
-  DType* input_ptr = reinterpret_cast<DType*>(input.data_ptr());
+  DType* input_ptr = nullptr;                                               
+  if(input.is_pinned())
+  {
+    CUDA_CALL(hipHostGetDevicePointer((void**)&input_ptr, input.data_ptr(), 0));
+  }
+  else{
+    input_ptr= reinterpret_cast<DType*>(input.data_ptr());
+
+  }
  DType* ret_ptr = reinterpret_cast<DType*>(ret.data_ptr());

  // Sort the index to improve the memory access pattern.