[mlir] Set CUDA/ROCm context before creating resources.

[csigg]

Summary

The current context is thread-local state, and in preparation of GPU async execution (on multiple threads) we need to set the context before calling API that create resources.

[herhut]

mlir/tools/mlir-cuda-runner/cuda-runtime-wrappers.cpp

61

Should this rather use push/pop in case there is some external (to the gpu dialect) use of the context, too? Like if this runs inside of some other runtime.

[csigg]

mlir/tools/mlir-cuda-runner/cuda-runtime-wrappers.cpp

61

It certainly could, but it seems a little over-engineered at this stage. But happy to add it if you think it makes sense.

[herhut]

mlir/tools/mlir-cuda-runner/cuda-runtime-wrappers.cpp

61

CUDA context issues are annoying to debug and why not if we can avoid creating that issue. I will forget this and then be puzzled :)

[csigg]

Restore context.

[herhut]

Hmm, this turned out more complex than I had thought. I had a simple push/pop in mind. If that is not enough, lets keep it at the simple version for now.

mlir/tools/mlir-cuda-runner/cuda-runtime-wrappers.cpp

45

Creating it always as before would make this less complex. What is the drawback?

52

This might no longer be the current one, if it was just created.

56

Why not use cuCtxPopCurrent here?

69

Doesn't cuCtxCreate already do this?

[csigg]

Switch to primary context.

[csigg]

mlir/tools/mlir-cuda-runner/cuda-runtime-wrappers.cpp

45

Setting a specific context allows running on a different device, for example. The use is quite limited though because mgpuSetContext() is not thread safe.

We will probably need to expose the per-thread context per thread, or per function that needs one.

I switched it to the primary context, which is the simplest.

52

See comment below.

56

The CUDA context stack is from early CUDA days. I have not seen anyone using it in years, and the HIP equivalent is marked deprecated.

69

cuCtxCreate sets the current context, this restores it so that the c'tor can grab it. It's a bit of a back and forth, but there is no call_once-else.

[csigg]

Rebase.

[csigg]

mlir/tools/mlir-rocm-runner/rocm-runtime-wrappers.cpp

49

This should say hipCtxGetCurrent will fix later.

[herhut]

Thanks for cleaning this up.

mlir/tools/mlir-rocm-runner/rocm-runtime-wrappers.cpp

50

context -> Context

[csigg]

Fixing typos.

mlir/tools/mlir-cuda-runner/cuda-runtime-wrappers.cpp

if (!result)                                                               \
       return;                                                                  \
     const char *name = nullptr;                                                \
     cuGetErrorName(result, &name);                                             \
     if (!name)                                                                 \
       name = "<unknown>";                                                      \
     llvm::errs() << "'" << #expr << "' failed with '" << name << "'\n";        \
   }(expr)
 
-// Static initialization of CUDA context for device ordinal 0.
-static auto InitializeCtx = [] {
+// Static reference to CUDA primary context for device ordinal 0.
+static CUcontext Context = [] {
   CUDA_REPORT_IF_ERROR(cuInit(/*flags=*/0));
   CUdevice device;
   CUDA_REPORT_IF_ERROR(cuDeviceGet(&device, /*ordinal=*/0));
   CUcontext context;
-  CUDA_REPORT_IF_ERROR(cuCtxCreate(&context, /*flags=*/0, device));
-  return 0;
+  CUDA_REPORT_IF_ERROR(cuDevicePrimaryCtxRetain(&context, device));
+  return context;
 }();
 
+// Sets the `Context` for the duration of the instance and restores the previous
+// context on destruction.
+class ScopedContext {
+public:
+  ScopedContext() {
+    CUDA_REPORT_IF_ERROR(cuCtxGetCurrent(&previous));
+    CUDA_REPORT_IF_ERROR(cuCtxSetCurrent(Context));
+  }
+
+  ~ScopedContext() { CUDA_REPORT_IF_ERROR(cuCtxSetCurrent(previous)); }
+
+private:
+  CUcontext previous;
+};
+
 extern "C" CUmodule mgpuModuleLoad(void *data) {
+  ScopedContext scopedContext;
   CUmodule module = nullptr;
   CUDA_REPORT_IF_ERROR(cuModuleLoadData(&module, data));
   return module;
 }
 
 extern "C" void mgpuModuleUnload(CUmodule module) {
   CUDA_REPORT_IF_ERROR(cuModuleUnload(module));
 }
 
 extern "C" CUfunction mgpuModuleGetFunction(CUmodule module, const char *name) {
   CUfunction function = nullptr;
   CUDA_REPORT_IF_ERROR(cuModuleGetFunction(&function, module, name));
   return function;
 }
 
 // The wrapper uses intptr_t instead of CUDA's unsigned int to match
 // the type of MLIR's index type. This avoids the need for casts in the
 // generated MLIR code.
 extern "C" void mgpuLaunchKernel(CUfunction function, intptr_t gridX,
                                  intptr_t gridY, intptr_t gridZ,
                                  intptr_t blockX, intptr_t blockY,
                                  intptr_t blockZ, int32_t smem, CUstream stream,
                                  void **params, void **extra) {
+  ScopedContext scopedContext;
   CUDA_REPORT_IF_ERROR(cuLaunchKernel(function, gridX, gridY, gridZ, blockX,
                                       blockY, blockZ, smem, stream, params,
                                       extra));
 }
 
 extern "C" CUstream mgpuStreamCreate() {
+  ScopedContext scopedContext;
   CUstream stream = nullptr;
   CUDA_REPORT_IF_ERROR(cuStreamCreate(&stream, CU_STREAM_NON_BLOCKING));
   return stream;
 }
 
 extern "C" void mgpuStreamDestroy(CUstream stream) {
   CUDA_REPORT_IF_ERROR(cuStreamDestroy(stream));
 }
 
 extern "C" void mgpuStreamSynchronize(CUstream stream) {
   CUDA_REPORT_IF_ERROR(cuStreamSynchronize(stream));
 }
 
 extern "C" void mgpuStreamWaitEvent(CUstream stream, CUevent event) {
   CUDA_REPORT_IF_ERROR(cuStreamWaitEvent(stream, event, /*flags=*/0));
 }
 
 extern "C" CUevent mgpuEventCreate() {
+  ScopedContext scopedContext;
   CUevent event = nullptr;
   CUDA_REPORT_IF_ERROR(cuEventCreate(&event, CU_EVENT_DISABLE_TIMING));
   return event;
 }
 
 extern "C" void mgpuEventDestroy(CUevent event) {
   CUDA_REPORT_IF_ERROR(cuEventDestroy(event));
 }
 
 extern "C" void mgpuEventSynchronize(CUevent event) {
   CUDA_REPORT_IF_ERROR(cuEventSynchronize(event));
 }
 
 extern "C" void mgpuEventRecord(CUevent event, CUstream stream) {
   CUDA_REPORT_IF_ERROR(cuEventRecord(event, stream));
 }
 
 extern "C" void *mgpuMemAlloc(uint64_t sizeBytes, CUstream /*stream*/) {
+  ScopedContext scopedContext;
   CUdeviceptr ptr;
   CUDA_REPORT_IF_ERROR(cuMemAlloc(&ptr, sizeBytes));
   return reinterpret_cast<void *>(ptr);
 }
 
 extern "C" void mgpuMemFree(void *ptr, CUstream /*stream*/) {
   CUDA_REPORT_IF_ERROR(cuMemFree(reinterpret_cast<CUdeviceptr>(ptr)));
 }

mlir/tools/mlir-rocm-runner/rocm-runtime-wrappers.cpp

[](hipError_t result) {                                                      \
     if (!result)                                                               \
       return;                                                                  \
     const char *name = hipGetErrorName(result);                                \
     if (!name)                                                                 \
       name = "<unknown>";                                                      \
     llvm::errs() << "'" << #expr << "' failed with '" << name << "'\n";        \
   }(expr)
 
-// Static initialization of HIP context for device ordinal 0.
-static auto InitializeCtx = [] {
+// Static reference to HIP primary context for device ordinal 0.
+static hipCtx_t Context = [] {
   HIP_REPORT_IF_ERROR(hipInit(/*flags=*/0));
   hipDevice_t device;
   HIP_REPORT_IF_ERROR(hipDeviceGet(&device, /*ordinal=*/0));
-  hipContext_t context;
-  HIP_REPORT_IF_ERROR(hipCtxCreate(&context, /*flags=*/0, device));
-  return 0;
+  hipCtx_t context;
+  HIP_REPORT_IF_ERROR(hipDevicePrimaryCtxRetain(&context, device));
+  return context;
 }();
 
+// Sets the `Context` for the duration of the instance and restores the previous
+// context on destruction.
+class ScopedContext {
+public:
+  ScopedContext() {
+    HIP_REPORT_IF_ERROR(hipCtxPushCurrent(&previous));
+    HIP_REPORT_IF_ERROR(hipCtxSetCurrent(context));
+  }
+
+  ~ScopedContext() { HIP_REPORT_IF_ERROR(hipCtxSetCurrent(previous)); }
+
+private:
+  hipCtx_t previous;
+};
+
 extern "C" hipModule_t mgpuModuleLoad(void *data) {
+  ScopedContext scopedContext;
   hipModule_t module = nullptr;
   HIP_REPORT_IF_ERROR(hipModuleLoadData(&module, data));
   return module;
 }
 
 extern "C" void mgpuModuleUnload(hipModule_t module) {
   HIP_REPORT_IF_ERROR(hipModuleUnload(module));
 }

 // the type of MLIR's index type. This avoids the need for casts in the
 // generated MLIR code.
 extern "C" void mgpuLaunchKernel(hipFunction_t function, intptr_t gridX,
                                  intptr_t gridY, intptr_t gridZ,
                                  intptr_t blockX, intptr_t blockY,
                                  intptr_t blockZ, int32_t smem,
                                  hipStream_t stream, void **params,
                                  void **extra) {
+  ScopedContext scopedContext;
   HIP_REPORT_IF_ERROR(hipModuleLaunchKernel(function, gridX, gridY, gridZ,
                                             blockX, blockY, blockZ, smem,
                                             stream, params, extra));
 }
 
 extern "C" hipStream_t mgpuStreamCreate() {
+  ScopedContext scopedContext;
   hipStream_t stream = nullptr;
   HIP_REPORT_IF_ERROR(hipStreamCreate(&stream));
   return stream;
 }
 
 extern "C" void mgpuStreamDestroy(hipStream_t stream) {
   HIP_REPORT_IF_ERROR(hipStreamDestroy(stream));
 }
 
 extern "C" void mgpuStreamSynchronize(hipStream_t stream) {
   return HIP_REPORT_IF_ERROR(hipStreamSynchronize(stream));
 }
 
 extern "C" void mgpuStreamWaitEvent(hipStream_t stream, hipEvent_t event) {
   HIP_REPORT_IF_ERROR(hipStreamWaitEvent(stream, event, /*flags=*/0));
 }
 
 extern "C" hipEvent_t mgpuEventCreate() {
+  ScopedContext scopedContext;
   hipEvent_t event = nullptr;
   HIP_REPORT_IF_ERROR(hipEventCreateWithFlags(&event, hipEventDisableTiming));
   return event;
 }
 
 extern "C" void mgpuEventDestroy(hipEvent_t event) {
   HIP_REPORT_IF_ERROR(hipEventDestroy(event));
 }
 
 extern "C" void mgpuEventSynchronize(hipEvent_t event) {
   HIP_REPORT_IF_ERROR(hipEventSynchronize(event));
 }
 
 extern "C" void mgpuEventRecord(hipEvent_t event, hipStream_t stream) {
   HIP_REPORT_IF_ERROR(hipEventRecord(event, stream));
 }
 
 extern "C" void *mgpuMemAlloc(uint64_t sizeBytes, hipStream_t /*stream*/) {
+  ScopedContext scopedContext;
   void *ptr;
   HIP_REPORT_IF_ERROR(hipMemAlloc(&ptr, sizeBytes));
   return ptr;
 }
 
 extern "C" void mgpuMemFree(void *ptr, hipStream_t /*stream*/) {
   HIP_REPORT_IF_ERROR(hipMemFree(ptr));
 }