mlir/lib/Dialect/GPU/Transforms: improve context management in SerializeToCubin

[rohany]

Summary

This commit adjusts the CUDA context management in the SerializeToCubin pass.
In particular, it uses the device 0 primary context instead of creating a new
CUDA context on each invocation of SerializeToCubin. This yields very large
improvements in compile time, especially if an application (like a JIT compiler)
is calling SerializeToCubin repeatedly.

[rohany]

Sorry, didn't see the github migration. Will move this to github.

mlir/lib/Dialect/GPU/Transforms/SerializeToCubin.cpp

SerializeToCubinPass::serializeISA(const std::string &isa) {
   char jitErrorBuffer[4096] = {0};
 
   RETURN_ON_CUDA_ERROR(cuInit(0));
 
   // Linking requires a device context.
   CUdevice device;
   RETURN_ON_CUDA_ERROR(cuDeviceGet(&device, 0));
   CUcontext context;
-  RETURN_ON_CUDA_ERROR(cuCtxCreate(&context, 0, device));
+  // Use the primary context.
+  RETURN_ON_CUDA_ERROR(cuDevicePrimaryCtxRetain(&context, device));
+  // Push the primary context so that the next CUDA operations
+  // actually use it.
+  RETURN_ON_CUDA_ERROR(cuCtxPushCurrent(context));
   CUlinkState linkState;
 
   CUjit_option jitOptions[] = {CU_JIT_ERROR_LOG_BUFFER,
                                CU_JIT_ERROR_LOG_BUFFER_SIZE_BYTES};
   void *jitOptionsVals[] = {jitErrorBuffer,
                             reinterpret_cast<void *>(sizeof(jitErrorBuffer))};
 
   RETURN_ON_CUDA_ERROR(cuLinkCreate(2,              /* number of jit options */
SerializeToCubinPass::serializeISA(const std::string &isa) {
   RETURN_ON_CUDA_ERROR(cuLinkComplete(linkState, &cubinData, &cubinSize));
 
   char *cubinAsChar = static_cast<char *>(cubinData);
   auto result =
       std::make_unique<std::vector<char>>(cubinAsChar, cubinAsChar + cubinSize);
 
   // This will also destroy the cubin data.
   RETURN_ON_CUDA_ERROR(cuLinkDestroy(linkState));
-  RETURN_ON_CUDA_ERROR(cuCtxDestroy(context));
+  // Pop and release the primary context.
+  CUcontext poppedContext;
+  RETURN_ON_CUDA_ERROR(cuCtxPopCurrent(&poppedContext));
+  RETURN_ON_CUDA_ERROR(cuDevicePrimaryCtxRelease(device));
 
   return result;
 }
 
 // Register pass to serialize GPU kernel functions to a CUBIN binary annotation.
 void mlir::registerGpuSerializeToCubinPass() {
   PassRegistration<SerializeToCubinPass> registerSerializeToCubin(
       [] { return std::make_unique<SerializeToCubinPass>(); });