From 8a43ec7faa274257325823f312ca2e7657c79785 Mon Sep 17 00:00:00 2001 From: Christian Sigg Date: Tue, 2 Feb 2021 10:43:32 +0100 Subject: [PATCH] Set GPU context before {cu,hip}MemHostRegister. Differential Revision: https://reviews.llvm.org/D95856 --- mlir/tools/mlir-cuda-runner/cuda-runtime-wrappers.cpp | 1 + mlir/tools/mlir-rocm-runner/rocm-runtime-wrappers.cpp | 1 + 2 files changed, 2 insertions(+) diff --git a/mlir/tools/mlir-cuda-runner/cuda-runtime-wrappers.cpp b/mlir/tools/mlir-cuda-runner/cuda-runtime-wrappers.cpp index d4360de76ed9..b8554bbc256b 100644 --- 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 index cf3c75720807..361ba8f8529d 100644 --- 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)); }