[mlir] Initialize CUDA context lazily.

[csigg]

Summary

So we can remove the ignore-warning pragma again.

[mehdi_amini]

mlir/lib/ExecutionEngine/CudaRuntimeWrappers.cpp

49

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

49

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"
+// Sets the `Context` for the duration of the instance and restores the previous
+// context on destruction.
+class ScopedContext {
+public:
+  ScopedContext() {
+    static std::once_flag CuInitFlag;
+    std::call_once(CuInitFlag,
+                   [] { CUDA_REPORT_IF_ERROR(cuInit(/*flags=*/0)); });
+
+    CUDA_REPORT_IF_ERROR(cuCtxGetCurrent(&previous));
+
     // 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(cuDevicePrimaryCtxRetain(&context, device));
       return context;
     }();
-#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));
   }
 
   ~ScopedContext() { CUDA_REPORT_IF_ERROR(cuCtxSetCurrent(previous)); }
 
 private:
   CUcontext previous;
 };