[mlir] Initialize CUDA context lazily.

[csigg]

Summary

So we can remove the ignore-warning pragma again.

[mehdi_amini]

mlir/lib/ExecutionEngine/CudaRuntimeWrappers.cpp

44–45

This has to be done before the initialization of static CUcontext Context below? It wasn't the case before this revision?
(it seems like the only reason to pull the cuInit out and move it above though.

[csigg]

Simplify.

mlir/lib/ExecutionEngine/CudaRuntimeWrappers.cpp

44–45

Before this revision, cuInit() was called during static initialization.

The current context and behavior cuCtxCreate vs cuDevicePrimaryPrimaryRetain always confuse me, so I took the safe route. It's not very well documented, but here is how it works:

There is a stack of contexts, the top one is current. A context needs to be active (ref-count > 0) to be usable.

cuDevicePrimaryPrimaryRetain/Release: increment/decrement primary.ref-count
cuCtxSetCurrent(ctx): stack.pop(), if (ctx != null) stack.push(ctx)
cuCtxCreate: create ctx, ctx.ref-count = 1, stack.push(ctx)
cuCtxDestroy(ctx): ctx.ref-count = 0, if (stack.top() == ctx) stack.pop()

Before calling cuCtxGetCurrent(), we need to call cuInit() and according to the explanation above we also /can/ retain the primary context because it does not affect the stack. I've changed the code now because it's indeed simpler and easier to read.

Thanks a lot for noticing/asking.

[csigg]

Switch to context push/pop.

[herhut]

So we get to use the context stack after all. Looks very clean now.

mlir/lib/ExecutionEngine/CudaRuntimeWrappers.cpp

if (!result)                                                               \
       return;                                                                  \
     const char *name = nullptr;                                                \
     cuGetErrorName(result, &name);                                             \
     if (!name)                                                                 \
       name = "<unknown>";                                                      \
     fprintf(stderr, "'%s' failed with '%s'\n", #expr, name);                   \
   }(expr)
 
-#pragma clang diagnostic push
-#pragma clang diagnostic ignored "-Wglobal-constructors"
+// Make the primary context of device 0 current for the duration of the instance
+// and restore the previous context on destruction.
+class ScopedContext {
+public:
+  ScopedContext() {
     // Static reference to CUDA primary context for device ordinal 0.
-static CUcontext Context = [] {
+    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(cuDevicePrimaryCtxRetain(&context, device));
-  return context;
+      CUcontext ctx;
+      // Note: this does not affect the current context.
+      CUDA_REPORT_IF_ERROR(cuDevicePrimaryCtxRetain(&ctx, device));
+      return ctx;
     }();
-#pragma clang diagnostic pop
 
-// 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));
+    CUDA_REPORT_IF_ERROR(cuCtxPushCurrent(context));
   }
 
-  ~ScopedContext() { CUDA_REPORT_IF_ERROR(cuCtxSetCurrent(previous)); }
-
-private:
-  CUcontext previous;
+  ~ScopedContext() { CUDA_REPORT_IF_ERROR(cuCtxPopCurrent(nullptr)); }
 };
 
 extern "C" MLIR_CUDA_WRAPPERS_EXPORT CUmodule mgpuModuleLoad(void *data) {
   ScopedContext scopedContext;
   CUmodule module = nullptr;
   CUDA_REPORT_IF_ERROR(cuModuleLoadData(&module, data));
   return module;
 }