Set GPU context before {cu,hip}MemHostRegister.

Differential Revision: https://reviews.llvm.org/D95856
This commit is contained in:
Christian Sigg 2021-02-02 10:43:32 +01:00
parent 39fbb5c3e3
commit 8a43ec7faa
2 changed files with 2 additions and 0 deletions

View File

@ -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));
}

View File

@ -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));
}