From c22876b550f5905b1cb7a4a0f329e76fac3e0232 Mon Sep 17 00:00:00 2001
From: Stephan Herhut <herhut@google.com>
Date: Tue, 21 Apr 2020 21:21:37 +0200
Subject: [PATCH] [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
---
 mlir/lib/Conversion/GPUToCUDA/ConvertKernelFuncToCubin.cpp | 7 +++++++
 1 file changed, 7 insertions(+)

diff --git a/mlir/lib/Conversion/GPUToCUDA/ConvertKernelFuncToCubin.cpp b/mlir/lib/Conversion/GPUToCUDA/ConvertKernelFuncToCubin.cpp
index 7cdb0dda4454..3f99c56c4716 100644
--- a/mlir/lib/Conversion/GPUToCUDA/ConvertKernelFuncToCubin.cpp
+++ b/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 @@ public:
   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();