[MLIR] Add extra locking during cubin generation.

We also need to lock the LLVMDialect mutex when initializing
LLVM targets or destroying llvm modules concurrently. Added another
scoped lock to that effect.

Differential Revision: https://reviews.llvm.org/D78580

Author: Stephan Herhut

mlir/lib/Conversion/GPUToCUDA/ConvertKernelFuncToCubin.cpp

@@ -31,6 +31,7 @@
 #include "llvm/IR/LegacyPassManager.h"
 #include "llvm/IR/Module.h"
 #include "llvm/Support/Error.h"
+#include "llvm/Support/Mutex.h"
 #include "llvm/Support/TargetRegistry.h"
 #include "llvm/Support/TargetSelect.h"
 #include "llvm/Target/TargetMachine.h"
@@ -57,6 +58,12 @@ class GpuKernelToCubinPass
   void runOnOperation() override {
     gpu::GPUModuleOp module = getOperation();
 
+    // Lock access to the llvm context.
+    llvm::sys::SmartScopedLock<true> scopedLock(
+        module.getContext()
+            ->getRegisteredDialect<LLVM::LLVMDialect>()
+            ->getLLVMContextMutex());
+
     // Make sure the NVPTX target is initialized.
     LLVMInitializeNVPTXTarget();
     LLVMInitializeNVPTXTargetInfo();