[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

72

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

72

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

72

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

37

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

44

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

48

Why not use cuCtxPopCurrent here?

61

Doesn't cuCtxCreate already do this?

[csigg]

Switch to primary context.

[csigg]

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

37

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.

44

See comment below.

48

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.

61

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

75

This should say hipCtxGetCurrent will fix later.

[herhut]

Thanks for cleaning this up.

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

76

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 = [] {
+// Context to use in functions below.
+// Can be set with 'mgpuSetContext()', otherwise will be created lazily.
+static CUcontext Context = nullptr;
+
+// Sets the CUDA context for the duration of the instance.
+class ScopedContext {
+public:
+  ScopedContext() {
+    static CUcontext context = MaybeCreateContext();
+    CUDA_REPORT_IF_ERROR(cuCtxGetCurrent(&previous));
+    CUDA_REPORT_IF_ERROR(cuCtxSetCurrent(context));
+  }
+
+  ~ScopedContext() { CUDA_REPORT_IF_ERROR(cuCtxSetCurrent(previous)); }
+
+private:
+  // Initializes 'Context' if necessary and returns it.
+  static CUcontext MaybeCreateContext() {
+    if (Context != nullptr)
+      return Context;
     CUDA_REPORT_IF_ERROR(cuInit(/*flags=*/0));
+    CUcontext previous;
+    CUDA_REPORT_IF_ERROR(cuCtxGetCurrent(&previous));
     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(cuCtxCreate(&Context, /*flags=*/0, device));
+    CUDA_REPORT_IF_ERROR(cuCtxSetCurrent(previous));
+    return Context;
+  }
+
+  CUcontext previous;
+};
+
+// Sets the CUDA context to use in functions below. Not thread safe.
+extern "C" void mgpuSetContext(CUcontext context) { Context = context; }
 
 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 = [] {
+// Context to use in functions below.
+// Can be set with 'mgpuSetContext()', otherwise will be created lazily.
+static hipCtx_t Context = nullptr;
+
+// Sets the HIP context for the duration of the instance.
+class ScopedContext {
+public:
+  ScopedContext() {
+    static hipCtx_t context = MaybeCreateContext();
+    HIP_REPORT_IF_ERROR(hipCtxGetCurrent(&previous));
+    HIP_REPORT_IF_ERROR(hipCtxSetCurrent(context));
+  }
+
+  ~ScopedContext() { HIP_REPORT_IF_ERROR(hipCtxSetCurrent(previous)); }
+
+private:
+  // Initializes 'Context' if necessary and returns it.
+  static hipCtx_t MaybeCreateContext() {
+    if (Context != nullptr)
+      return Context;
     HIP_REPORT_IF_ERROR(hipInit(/*flags=*/0));
+    hipCtx_t previous;
+    HIP_REPORT_IF_ERROR(hipCtxGetCurrent(&previous));
     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;
-}();
+    HIP_REPORT_IF_ERROR(hipCtxCreate(&Context, /*flags=*/0, device));
+    HIP_REPORT_IF_ERROR(hipCtxSetCurrent(previous));
+    return Context;
+  }
+
+  hipCtx_t previous;
+};
+
+// Sets the CUDA context to use in functions below. Not thread safe.
+extern "C" void mgpuSetContext(hipCtx_t context) { Context = context; }
 
 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));
 }
 
 extern "C" hipFunction_t mgpuModuleGetFunction(hipModule_t module,
                                                const char *name) {
   hipFunction_t function = nullptr;
   HIP_REPORT_IF_ERROR(hipModuleGetFunction(&function, module, name));
   return function;
 }
 
 // The wrapper uses intptr_t instead of ROCM'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(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));
 }