Set GPU context before {cu,hip}MemHostRegister.

Differential Revision: https://reviews.llvm.org/D95856
llvm · Feb 3, 2021 · 8a43ec7 · 8a43ec7
1 parent 39fbb5c
commit 8a43ec7
Show file tree

Hide file tree

Showing 2 changed files with 2 additions and 0 deletions.
diff --git a/mlir/tools/mlir-cuda-runner/cuda-runtime-wrappers.cpp b/mlir/tools/mlir-cuda-runner/cuda-runtime-wrappers.cpp
@@ -148,6 +148,7 @@ extern "C" void mgpuMemcpy(void *dst, void *src, uint64_t sizeBytes,
 // Allows to register byte array with the CUDA runtime. Helpful until we have
 // transfer functions implemented.
 extern "C" void mgpuMemHostRegister(void *ptr, uint64_t sizeBytes) {
+  ScopedContext scopedContext;
   CUDA_REPORT_IF_ERROR(cuMemHostRegister(ptr, sizeBytes, /*flags=*/0));
 }
 

diff --git a/mlir/tools/mlir-rocm-runner/rocm-runtime-wrappers.cpp b/mlir/tools/mlir-rocm-runner/rocm-runtime-wrappers.cpp
@@ -148,6 +148,7 @@ extern "C" void mgpuMemcpy(void *dst, void *src, uint64_t sizeBytes,
 // Allows to register byte array with the ROCM runtime. Helpful until we have
 // transfer functions implemented.
 extern "C" void mgpuMemHostRegister(void *ptr, uint64_t sizeBytes) {
+  ScopedContext scopedContext;
   HIP_REPORT_IF_ERROR(hipHostRegister(ptr, sizeBytes, /*flags=*/0));
 }