[mlir] take LLVMContext in MLIR-to-LLVM-IR translation

[ftynse]

Summary

Due to the original type system implementation, LLVMDialect in MLIR contains an
LLVMContext in which the relevant objects (types, metadata) are created. When
an MLIR module using the LLVM dialect (and related intrinsic-based dialects
NVVM, ROCDL, AVX512) is converted to LLVM IR, it could only live in the
LLVMContext owned by the dialect. The type system no longer relies on the
LLVMContext, so this limitation can be removed. Instead, translation functions
now take a reference to an LLVMContext in which the LLVM IR module should be
constructed. The caller of the translation functions is responsible for
ensuring the same LLVMContext is not used concurrently as the translation no
longer uses a dialect-wide context lock.

As an additional bonus, this change removes the need to recreate the LLVM IR
module in a different LLVMContext through printing and parsing back, decreasing
the compilation overhead in JIT and GPU-kernel-to-blob passes.

[mehdi_amini]

Nice! I'm so glad to see this change :)

What state is still present in the dialect after this?

[ftynse]

What state is still present in the dialect after this?

A follow-up commit removes the context and the module, and only keeps an llvm::DataLayout. We need a proper model for that in MLIR before we can remove it completely, but we can move it from dialect to module as an intermediate step after a minor refactoring in LLVM.

[rriddle]

mlir/include/mlir/Target/ROCDLIR.h

18

Can we trim the context include here?

[ftynse]

Address review

mlir/docs/Tutorials/Toy/Ch-6.md

define void @main()
   %putchar.1 = tail call i32 @putchar(i32 10)
   %4 = tail call i32 (i8*, ...) @printf(i8* nonnull dereferenceable(1) getelementptr inbounds ([4 x i8], [4 x i8]* @frmt_spec, i64 0, i64 0), double 9.000000e+00)
   %5 = tail call i32 (i8*, ...) @printf(i8* nonnull dereferenceable(1) getelementptr inbounds ([4 x i8], [4 x i8]* @frmt_spec, i64 0, i64 0), double 3.600000e+01)
   %putchar.2 = tail call i32 @putchar(i32 10)
   ret void
 }
 ```
 
-The full code listing for dumping LLVM IR can be found in 
+The full code listing for dumping LLVM IR can be found in
 `examples/toy/Ch6/toy.cpp` in the `dumpLLVMIR()` function:
 
 ```c++
 
 int dumpLLVMIR(mlir::ModuleOp module) {
-  // Translate the module, that contains the LLVM dialect, to LLVM IR.
-  auto llvmModule = mlir::translateModuleToLLVMIR(module);
+  // Translate the module, that contains the LLVM dialect, to LLVM IR. Use a
+  // fresh LLVM IR context. (Note that LLVM is not thread-safe and any
+  // concurrent use of a context requires external locking.)
+  llvm::LLVMContext llvmContext;
+  auto llvmModule = mlir::translateModuleToLLVMIR(module, llvmContext);
   if (!llvmModule) {
     llvm::errs() << "Failed to emit LLVM IR\n";
     return -1;
   }
 
   // Initialize LLVM targets.
   llvm::InitializeNativeTarget();
   llvm::InitializeNativeTargetAsmPrinter();

mlir/examples/toy/Ch6/toyc.cpp

int dumpAST() {
   if (!moduleAST)
     return 1;
 
   dump(*moduleAST);
   return 0;
 }
 
 int dumpLLVMIR(mlir::ModuleOp module) {
-  auto llvmModule = mlir::translateModuleToLLVMIR(module);
+  // Convert the module to LLVM IR in a new LLVM IR context.
+  llvm::LLVMContext llvmContext;
+  auto llvmModule = mlir::translateModuleToLLVMIR(module, llvmContext);
   if (!llvmModule) {
     llvm::errs() << "Failed to emit LLVM IR\n";
     return -1;
   }
 
   // Initialize LLVM targets.
   llvm::InitializeNativeTarget();
   llvm::InitializeNativeTargetAsmPrinter();

mlir/examples/toy/Ch7/toyc.cpp

int dumpAST() {
   if (!moduleAST)
     return 1;
 
   dump(*moduleAST);
   return 0;
 }
 
 int dumpLLVMIR(mlir::ModuleOp module) {
-  auto llvmModule = mlir::translateModuleToLLVMIR(module);
+  // Convert the module to LLVM IR in a new LLVM IR context.
+  llvm::LLVMContext llvmContext;
+  auto llvmModule = mlir::translateModuleToLLVMIR(module, llvmContext);
   if (!llvmModule) {
     llvm::errs() << "Failed to emit LLVM IR\n";
     return -1;
   }
 
   // Initialize LLVM targets.
   llvm::InitializeNativeTarget();
   llvm::InitializeNativeTargetAsmPrinter();

mlir/include/mlir/Conversion/GPUCommon/GPUCommonPass.h

 
 namespace LLVM {
 class LLVMDialect;
 } // namespace LLVM
 
 using OwnedBlob = std::unique_ptr<std::vector<char>>;
 using BlobGenerator =
     std::function<OwnedBlob(const std::string &, Location, StringRef)>;
-using LoweringCallback =
-    std::function<std::unique_ptr<llvm::Module>(Operation *)>;
+using LoweringCallback = std::function<std::unique_ptr<llvm::Module>(
+    Operation *, llvm::LLVMContext &, StringRef)>;
 
 /// Creates a pass to convert a gpu.launch_func operation into a sequence of
 /// GPU runtime calls.
 ///
 /// This pass does not generate code to call GPU runtime APIs directly but
 /// instead uses a small wrapper library that exports a stable and conveniently
 /// typed ABI on top of GPU runtimes such as CUDA or ROCm (HIP).
 std::unique_ptr<OperationPass<ModuleOp>>

mlir/include/mlir/Dialect/LLVMIR/LLVMDialect.h

 /// string (operations inserted at the builder insertion point).
 Value createGlobalString(Location loc, OpBuilder &builder, StringRef name,
                          StringRef value, LLVM::Linkage linkage);
 
 /// LLVM requires some operations to be inside of a Module operation. This
 /// function confirms that the Operation has the desired properties.
 bool satisfiesLLVMModule(Operation *op);
 
-/// Clones the given module into the provided context. This is implemented by
-/// transforming the module into bitcode and then reparsing the bitcode in the
-/// provided context.
-std::unique_ptr<llvm::Module>
-cloneModuleIntoNewContext(llvm::LLVMContext *context, llvm::Module *module);
-
 } // end namespace LLVM
 } // end namespace mlir
 
 #endif // MLIR_DIALECT_LLVMIR_LLVMDIALECT_H_

mlir/include/mlir/Target/LLVMIR.h

 //===- LLVMIR.h - MLIR to LLVM IR conversion --------------------*- C++ -*-===//
 //
 // Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
 // See https://llvm.org/LICENSE.txt for license information.
 // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
 //
 //===----------------------------------------------------------------------===//
 //
 // This file declares the entry point for the MLIR to LLVM IR conversion.
 //
 //===----------------------------------------------------------------------===//
 
 #ifndef MLIR_TARGET_LLVMIR_H
 #define MLIR_TARGET_LLVMIR_H
 
+#include "mlir/Support/LLVM.h"
+#include "llvm/ADT/StringRef.h"
 #include <memory>
 
 // Forward-declare LLVM classes.
 namespace llvm {
 class LLVMContext;
 class Module;
 } // namespace llvm
 
 namespace mlir {
 
 class OwningModuleRef;
 class MLIRContext;
 class ModuleOp;
 
 /// Convert the given MLIR module into LLVM IR.  The LLVM context is extracted
 /// from the registered LLVM IR dialect.  In case of error, report it
 /// to the error handler registered with the MLIR context, if any (obtained from
 /// the MLIR module), and return `nullptr`.
-std::unique_ptr<llvm::Module> translateModuleToLLVMIR(ModuleOp m);
+std::unique_ptr<llvm::Module>
+translateModuleToLLVMIR(ModuleOp m, llvm::LLVMContext &llvmContext,
+                        StringRef name = "LLVMDialectModule");
 
 /// Convert the given LLVM module into MLIR's LLVM dialect.  The LLVM context is
 /// extracted from the registered LLVM IR dialect. In case of error, report it
 /// to the error handler registered with the MLIR context, if any (obtained from
 /// the MLIR module), and return `{}`.
 OwningModuleRef
 translateLLVMIRToModule(std::unique_ptr<llvm::Module> llvmModule,
                         MLIRContext *context);
 
 } // namespace mlir
 
 #endif // MLIR_TARGET_LLVMIR_H

mlir/include/mlir/Target/LLVMIR/ModuleTranslation.h

 /// Implementation class for module translation. Holds a reference to the module
 /// being translated, and the mappings between the original and the translated
 /// functions, basic blocks and values. It is practically easier to hold these
 /// mappings in one class since the conversion of control flow operations
 /// needs to look up block and function mappings.
 class ModuleTranslation {
 public:
   template <typename T = ModuleTranslation>
-  static std::unique_ptr<llvm::Module> translateModule(Operation *m) {
+  static std::unique_ptr<llvm::Module>
+  translateModule(Operation *m, llvm::LLVMContext &llvmContext,
+                  StringRef name = "LLVMDialectModule") {
     if (!satisfiesLLVMModule(m))
       return nullptr;
     if (failed(checkSupportedModuleOps(m)))
       return nullptr;
-    auto llvmModule = prepareLLVMModule(m);
-    if (!llvmModule)
-      return nullptr;
+    std::unique_ptr<llvm::Module> llvmModule =
+        prepareLLVMModule(m, llvmContext, name);
 
     LLVM::ensureDistinctSuccessors(m);
 
     T translator(m, std::move(llvmModule));
     if (failed(translator.convertFunctionSignatures()))
       return nullptr;
     if (failed(translator.convertGlobals()))
       return nullptr;
protected:
   virtual LogicalResult convertOmpOperation(Operation &op,
                                             llvm::IRBuilder<> &builder);
   virtual LogicalResult convertOmpParallel(Operation &op,
                                            llvm::IRBuilder<> &builder);
 
   /// Converts the type from MLIR LLVM dialect to LLVM.
   llvm::Type *convertType(LLVMType type);
 
-  static std::unique_ptr<llvm::Module> prepareLLVMModule(Operation *m);
+  static std::unique_ptr<llvm::Module>
+  prepareLLVMModule(Operation *m, llvm::LLVMContext &llvmContext,
+                    StringRef name);
 
   /// A helper to look up remapped operands in the value remapping table.
   SmallVector<llvm::Value *, 8> lookupValues(ValueRange values);
 
 private:
   /// Check whether the module contains only supported ops directly in its body.
   static LogicalResult checkSupportedModuleOps(Operation *m);
 
private:
   std::unique_ptr<llvm::Module> llvmModule;
   /// A converter for translating debug information.
   std::unique_ptr<detail::DebugTranslation> debugTranslation;
 
   /// Builder for LLVM IR generation of OpenMP constructs.
   std::unique_ptr<llvm::OpenMPIRBuilder> ompBuilder;
   /// Precomputed pointer to OpenMP dialect.
   const Dialect *ompDialect;
-  /// Pointer to the llvmDialect;
-  LLVMDialect *llvmDialect;
 
   /// Mappings between llvm.mlir.global definitions and corresponding globals.
   DenseMap<Operation *, llvm::GlobalValue *> globalsMapping;
 
   /// A stateful object used to translate types.
   TypeToLLVMIRTranslator typeTranslator;
 
 protected:

mlir/include/mlir/Target/NVVMIR.h

 //===- NVVMIR.h - MLIR to LLVM + NVVM IR conversion -------------*- C++ -*-===//
 //
 // Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
 // See https://llvm.org/LICENSE.txt for license information.
 // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
 //
 //===----------------------------------------------------------------------===//
 //
 // This file declares the entry point for the MLIR to LLVM + NVVM IR conversion.
 //
 //===----------------------------------------------------------------------===//
 
 #ifndef MLIR_TARGET_NVVMIR_H
 #define MLIR_TARGET_NVVMIR_H
 
+#include "llvm/ADT/StringRef.h"
 #include <memory>
 
 // Forward-declare LLVM classes.
 namespace llvm {
+class LLVMContext;
 class Module;
 } // namespace llvm
 
 namespace mlir {
 class Operation;
 
 /// Convert the given LLVM-module-like operation into NVVM IR. This conversion
 /// requires the registration of the LLVM IR dialect and will extract the LLVM
 /// context from the registered LLVM IR dialect.  In case of error, report it to
 /// the error handler registered with the MLIR context, if any (obtained from
 /// the MLIR module), and return `nullptr`.
-std::unique_ptr<llvm::Module> translateModuleToNVVMIR(Operation *m);
+std::unique_ptr<llvm::Module>
+translateModuleToNVVMIR(Operation *m, llvm::LLVMContext &llvmContext,
+                        llvm::StringRef name = "LLVMDialectModule");
 
 } // namespace mlir
 
 #endif // MLIR_TARGET_NVVMIR_H

mlir/include/mlir/Target/ROCDLIR.h

 //===- ROCDLIR.h - MLIR to LLVM + ROCDL IR conversion -----------*- C++ -*-===//
 //
 // Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
 // See https://llvm.org/LICENSE.txt for license information.
 // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
 //
 //===----------------------------------------------------------------------===//
 //
 // This file declares the entry point for the MLIR to LLVM + ROCDL IR
 // conversion.
 //
 //===----------------------------------------------------------------------===//
 
 #ifndef MLIR_TARGET_ROCDLIR_H
 #define MLIR_TARGET_ROCDLIR_H
 
+#include "llvm/ADT/StringRef.h"
 #include <memory>
 
 // Forward-declare LLVM classes.
 namespace llvm {
+class LLVMContext;
 class Module;
 } // namespace llvm
 
 namespace mlir {
 class Operation;
 
 /// Convert the given LLVM-module-like operation into ROCDL IR. This conversion
 /// requires the registration of the LLVM IR dialect and will extract the LLVM
 /// context from the registered LLVM IR dialect.  In case of error, report it to
 /// the error handler registered with the MLIR context, if any (obtained from
 /// the MLIR module), and return `nullptr`.
-std::unique_ptr<llvm::Module> translateModuleToROCDLIR(Operation *m);
+std::unique_ptr<llvm::Module>
+translateModuleToROCDLIR(Operation *m, llvm::LLVMContext &llvmContext,
+                         llvm::StringRef name = "LLVMDialectModule");
 
 } // namespace mlir
 
 #endif // MLIR_TARGET_ROCDLIR_H

mlir/lib/Conversion/GPUCommon/ConvertKernelFuncToBlob.cpp

GpuKernelToBlobPass(LoweringCallback loweringCallback,
                       StringRef gpuBinaryAnnotation)
       : loweringCallback(loweringCallback), blobGenerator(blobGenerator),
         triple(triple), targetChip(targetChip), features(features),
         blobAnnotation(gpuBinaryAnnotation) {}
 
   void runOnOperation() override {
     gpu::GPUModuleOp module = getOperation();
 
-    // Lock access to the llvm context.
-    llvm::sys::SmartScopedLock<true> scopedLock(
-        module.getContext()
-            ->getRegisteredDialect<LLVM::LLVMDialect>()
-            ->getLLVMContextMutex());
-
-    // Lower the module to a llvm module.
-    std::unique_ptr<llvm::Module> llvmModule = loweringCallback(module);
+    // Lower the module to an LLVM IR module using a separate context to enable
+    // multi-threaded processing.
+    llvm::LLVMContext llvmContext;
+    std::unique_ptr<llvm::Module> llvmModule =
+        loweringCallback(module, llvmContext, "LLVMDialectModule");
     if (!llvmModule)
       return signalPassFailure();
 
     // Translate the llvm module to a target blob and attach the result as
     // attribute to the module.
     if (auto blobAttr = translateGPUModuleToBinaryAnnotation(
             *llvmModule, module.getLoc(), module.getName()))
       module.setAttr(blobAnnotation, blobAttr);

 
 } // anonymous namespace
 
 std::string
 GpuKernelToBlobPass::translateModuleToISA(llvm::Module &module,
                                           llvm::TargetMachine &targetMachine) {
   std::string targetISA;
   {
-    // Clone the llvm module into a new context to enable concurrent compilation
-    // with multiple threads.
-    llvm::LLVMContext llvmContext;
-    auto clone = LLVM::cloneModuleIntoNewContext(&llvmContext, &module);
-
     llvm::raw_string_ostream stream(targetISA);
     llvm::buffer_ostream pstream(stream);
     llvm::legacy::PassManager codegenPasses;
     targetMachine.addPassesToEmitFile(codegenPasses, pstream, nullptr,
                                       llvm::CGFT_AssemblyFile);
-    codegenPasses.run(*clone);
+    codegenPasses.run(module);
   }
 
   return targetISA;
 }
 
 OwnedBlob GpuKernelToBlobPass::convertModuleToBlob(llvm::Module &llvmModule,
                                                    Location loc,
                                                    StringRef name) {

mlir/lib/Dialect/LLVMIR/IR/LLVMDialect.cpp

Value mlir::LLVM::createGlobalString(Location loc, OpBuilder &builder,
   return builder.create<LLVM::GEPOp>(loc, LLVM::LLVMType::getInt8PtrTy(ctx),
                                      globalPtr, ArrayRef<Value>({cst0, cst0}));
 }
 
 bool mlir::LLVM::satisfiesLLVMModule(Operation *op) {
   return op->hasTrait<OpTrait::SymbolTable>() &&
          op->hasTrait<OpTrait::IsIsolatedFromAbove>();
 }
-
-std::unique_ptr<llvm::Module>
-mlir::LLVM::cloneModuleIntoNewContext(llvm::LLVMContext *context,
-                                      llvm::Module *module) {
-  SmallVector<char, 1> buffer;
-  {
-    llvm::raw_svector_ostream os(buffer);
-    WriteBitcodeToFile(*module, os);
-  }
-  llvm::MemoryBufferRef bufferRef(StringRef(buffer.data(), buffer.size()),
-                                  "cloned module buffer");
-  return cantFail(parseBitcodeFile(bufferRef, *context));
-}

mlir/lib/ExecutionEngine/ExecutionEngine.cpp

Expected<std::unique_ptr<ExecutionEngine>> ExecutionEngine::create(
     Optional<llvm::CodeGenOpt::Level> jitCodeGenOptLevel,
     ArrayRef<StringRef> sharedLibPaths, bool enableObjectCache,
     bool enableGDBNotificationListener, bool enablePerfNotificationListener) {
   auto engine = std::make_unique<ExecutionEngine>(
       enableObjectCache, enableGDBNotificationListener,
       enablePerfNotificationListener);
 
   std::unique_ptr<llvm::LLVMContext> ctx(new llvm::LLVMContext);
-  auto llvmModule = translateModuleToLLVMIR(m);
+  auto llvmModule = translateModuleToLLVMIR(m, *ctx);
   if (!llvmModule)
     return make_string_error("could not convert to LLVM IR");
   // FIXME: the triple should be passed to the translation or dialect conversion
   // instead of this.  Currently, the LLVM module created above has no triple
   // associated with it.
   setupTargetTriple(llvmModule.get());
   packFunctionArguments(llvmModule.get());
 
-  // Clone module in a new LLVMContext since translateModuleToLLVMIR buries
-  // ownership too deeply.
-  // TODO: Reevaluate model of ownership of LLVMContext in LLVMDialect.
-  std::unique_ptr<Module> deserModule =
-      LLVM::cloneModuleIntoNewContext(ctx.get(), llvmModule.get());
-  auto dataLayout = deserModule->getDataLayout();
+  auto dataLayout = llvmModule->getDataLayout();
 
   // Callback to create the object layer with symbol resolution to current
   // process and dynamically linked libraries.
   auto objectLinkingLayerCreator = [&](ExecutionSession &session,
                                        const Triple &TT) {
     auto objectLayer = std::make_unique<RTDyldObjectLinkingLayer>(
         session, []() { return std::make_unique<SectionMemoryManager>(); });
 
Expected<std::unique_ptr<ExecutionEngine>> ExecutionEngine::create(
   // Create the LLJIT by calling the LLJITBuilder with 2 callbacks.
   auto jit =
       cantFail(llvm::orc::LLJITBuilder()
                    .setCompileFunctionCreator(compileFunctionCreator)
                    .setObjectLinkingLayerCreator(objectLinkingLayerCreator)
                    .create());
 
   // Add a ThreadSafemodule to the engine and return.
-  ThreadSafeModule tsm(std::move(deserModule), std::move(ctx));
+  ThreadSafeModule tsm(std::move(llvmModule), std::move(ctx));
   if (transformer)
     cantFail(tsm.withModuleDo(
         [&](llvm::Module &module) { return transformer(&module); }));
   cantFail(jit->addIRModule(std::move(tsm)));
   engine->jit = std::move(jit);
 
   // Resolve symbols that are statically linked in the current process.
   llvm::orc::JITDylib &mainJD = engine->jit->getMainJITDylib();

mlir/lib/Target/LLVMIR/ConvertToLLVMIR.cpp

 #include "mlir/Translation.h"
 
 #include "llvm/ADT/StringRef.h"
 #include "llvm/IR/Module.h"
 #include "llvm/Support/ToolOutputFile.h"
 
 using namespace mlir;
 
-std::unique_ptr<llvm::Module> mlir::translateModuleToLLVMIR(ModuleOp m) {
-  return LLVM::ModuleTranslation::translateModule<>(m);
+std::unique_ptr<llvm::Module>
+mlir::translateModuleToLLVMIR(ModuleOp m, llvm::LLVMContext &llvmContext,
+                              StringRef name) {
+  return LLVM::ModuleTranslation::translateModule<>(m, llvmContext, name);
 }
 
 namespace mlir {
 void registerToLLVMIRTranslation() {
   TranslateFromMLIRRegistration registration(
       "mlir-to-llvmir", [](ModuleOp module, raw_ostream &output) {
-        auto llvmModule = LLVM::ModuleTranslation::translateModule<>(module);
+        llvm::LLVMContext llvmContext;
+        auto llvmModule = LLVM::ModuleTranslation::translateModule<>(
+            module, llvmContext, "LLVMDialectModule");
         if (!llvmModule)
           return failure();
 
         llvmModule->print(output, nullptr);
         return success();
       });
 }
 } // namespace mlir

mlir/lib/Target/LLVMIR/ConvertToNVVMIR.cpp

#include "mlir/Dialect/LLVMIR/NVVMConversions.inc"
     return LLVM::ModuleTranslation::convertOperation(opInst, builder);
   }
 
   /// Allow access to the constructor.
   friend LLVM::ModuleTranslation;
 };
 } // namespace
 
-std::unique_ptr<llvm::Module> mlir::translateModuleToNVVMIR(Operation *m) {
-  auto llvmModule =
-      LLVM::ModuleTranslation::translateModule<ModuleTranslation>(m);
+std::unique_ptr<llvm::Module>
+mlir::translateModuleToNVVMIR(Operation *m, llvm::LLVMContext &llvmContext,
+                              StringRef name) {
+  auto llvmModule = LLVM::ModuleTranslation::translateModule<ModuleTranslation>(
+      m, llvmContext, name);
   if (!llvmModule)
     return llvmModule;
 
   // Insert the nvvm.annotations kernel so that the NVVM backend recognizes the
   // function as a kernel.
   for (auto func :
        ModuleTranslation::getModuleBody(m).getOps<LLVM::LLVMFuncOp>()) {
     if (!gpu::GPUDialect::isKernel(func))
mlir::translateModuleToNVVMIR(Operation *m, llvm::LLVMContext &llvmContext,
 
   return llvmModule;
 }
 
 namespace mlir {
 void registerToNVVMIRTranslation() {
   TranslateFromMLIRRegistration registration(
       "mlir-to-nvvmir", [](ModuleOp module, raw_ostream &output) {
-        auto llvmModule = mlir::translateModuleToNVVMIR(module);
+        llvm::LLVMContext llvmContext;
+        auto llvmModule = mlir::translateModuleToNVVMIR(module, llvmContext);
         if (!llvmModule)
           return failure();
 
         llvmModule->print(output, nullptr);
         return success();
       });
 }
 } // namespace mlir

mlir/lib/Target/LLVMIR/ConvertToROCDLIR.cpp

#include "mlir/Dialect/LLVMIR/ROCDLConversions.inc"
     return LLVM::ModuleTranslation::convertOperation(opInst, builder);
   }
 
   /// Allow access to the constructor.
   friend LLVM::ModuleTranslation;
 };
 } // namespace
 
-std::unique_ptr<llvm::Module> mlir::translateModuleToROCDLIR(Operation *m) {
+std::unique_ptr<llvm::Module>
+mlir::translateModuleToROCDLIR(Operation *m, llvm::LLVMContext &llvmContext,
+                               StringRef name) {
   // lower MLIR (with RODL Dialect) to LLVM IR (with ROCDL intrinsics)
-  auto llvmModule =
-      LLVM::ModuleTranslation::translateModule<ModuleTranslation>(m);
+  auto llvmModule = LLVM::ModuleTranslation::translateModule<ModuleTranslation>(
+      m, llvmContext, name);
 
   // foreach GPU kernel
   // 1. Insert AMDGPU_KERNEL calling convention.
   // 2. Insert amdgpu-flat-workgroup-size(1, 1024) attribute.
   for (auto func :
        ModuleTranslation::getModuleBody(m).getOps<LLVM::LLVMFuncOp>()) {
     if (!func.getAttrOfType<UnitAttr>(gpu::GPUDialect::getKernelFuncAttrName()))
       continue;
 
     auto *llvmFunc = llvmModule->getFunction(func.getName());
 
     llvmFunc->setCallingConv(llvm::CallingConv::AMDGPU_KERNEL);
 
     llvmFunc->addFnAttr("amdgpu-flat-work-group-size", "1, 1024");
   }
 
   return llvmModule;
 }
 
 namespace mlir {
 void registerToROCDLIRTranslation() {
   TranslateFromMLIRRegistration registration(
       "mlir-to-rocdlir", [](ModuleOp module, raw_ostream &output) {
-        auto llvmModule = mlir::translateModuleToROCDLIR(module);
+        llvm::LLVMContext llvmContext;
+        auto llvmModule = mlir::translateModuleToROCDLIR(module, llvmContext);
         if (!llvmModule)
           return failure();
 
         llvmModule->print(output, nullptr);
         return success();
       });
 }
 } // namespace mlir

mlir/lib/Target/LLVMIR/LLVMAVX512Intr.cpp

protected:
   LogicalResult convertOperation(Operation &opInst,
                                  llvm::IRBuilder<> &builder) override {
 #include "mlir/Dialect/LLVMIR/LLVMAVX512Conversions.inc"
 
     return LLVM::ModuleTranslation::convertOperation(opInst, builder);
   }
 };
 
-std::unique_ptr<llvm::Module> translateLLVMAVX512ModuleToLLVMIR(Operation *m) {
+std::unique_ptr<llvm::Module>
+translateLLVMAVX512ModuleToLLVMIR(Operation *m, llvm::LLVMContext &llvmContext,
+                                  StringRef name) {
   return LLVM::ModuleTranslation::translateModule<LLVMAVX512ModuleTranslation>(
-      m);
+      m, llvmContext, name);
 }
 } // end namespace
 
 namespace mlir {
 void registerAVX512ToLLVMIRTranslation() {
   TranslateFromMLIRRegistration reg(
       "avx512-mlir-to-llvmir", [](ModuleOp module, raw_ostream &output) {
-        auto llvmModule = translateLLVMAVX512ModuleToLLVMIR(module);
+        llvm::LLVMContext llvmContext;
+        auto llvmModule = translateLLVMAVX512ModuleToLLVMIR(
+            module, llvmContext, "LLVMDialectModule");
         if (!llvmModule)
           return failure();
 
         llvmModule->print(output, nullptr);
         return success();
       });
 }
 } // namespace mlir

mlir/lib/Target/LLVMIR/ModuleTranslation.cpp

 
 ModuleTranslation::ModuleTranslation(Operation *module,
                                      std::unique_ptr<llvm::Module> llvmModule)
     : mlirModule(module), llvmModule(std::move(llvmModule)),
       debugTranslation(
           std::make_unique<DebugTranslation>(module, *this->llvmModule)),
       ompDialect(
           module->getContext()->getRegisteredDialect<omp::OpenMPDialect>()),
-      llvmDialect(module->getContext()->getRegisteredDialect<LLVMDialect>()),
       typeTranslator(this->llvmModule->getContext()) {
   assert(satisfiesLLVMModule(mlirModule) &&
          "mlirModule should honor LLVM's module semantics.");
 }
 ModuleTranslation::~ModuleTranslation() {
   if (ompBuilder)
     ompBuilder->finalize();
 }
LogicalResult ModuleTranslation::convertBlock(Block &bb, bool ignoreArguments) {
   }
 
   return success();
 }
 
 /// Create named global variables that correspond to llvm.mlir.global
 /// definitions.
 LogicalResult ModuleTranslation::convertGlobals() {
-  // Lock access to the llvm context.
-  llvm::sys::SmartScopedLock<true> scopedLock(
-      llvmDialect->getLLVMContextMutex());
   for (auto op : getModuleBody(mlirModule).getOps<LLVM::GlobalOp>()) {
     llvm::Type *type = convertType(op.getType());
     llvm::Constant *cst = llvm::UndefValue::get(type);
     if (op.getValueOrNull()) {
       // String attributes are treated separately because they cannot appear as
       // in-function constants and are thus not supported by getLLVMConstant.
       if (auto strAttr = op.getValueOrNull().dyn_cast_or_null<StringAttr>()) {
         cst = llvm::ConstantDataArray::getString(

 LogicalResult ModuleTranslation::checkSupportedModuleOps(Operation *m) {
   for (Operation &o : getModuleBody(m).getOperations())
     if (!isa<LLVM::LLVMFuncOp, LLVM::GlobalOp>(&o) && !o.isKnownTerminator())
       return o.emitOpError("unsupported module-level operation");
   return success();
 }
 
 LogicalResult ModuleTranslation::convertFunctionSignatures() {
-  // Lock access to the llvm context.
-  llvm::sys::SmartScopedLock<true> scopedLock(
-      llvmDialect->getLLVMContextMutex());
-
   // Declare all functions first because there may be function calls that form a
   // call graph with cycles, or global initializers that reference functions.
   for (auto function : getModuleBody(mlirModule).getOps<LLVMFuncOp>()) {
     llvm::FunctionCallee llvmFuncCst = llvmModule->getOrInsertFunction(
         function.getName(),
         cast<llvm::FunctionType>(convertType(function.getType())));
     llvm::Function *llvmFunc = cast<llvm::Function>(llvmFuncCst.getCallee());
     llvmFunc->setLinkage(convertLinkageToLLVM(function.linkage()));
     functionMapping[function.getName()] = llvmFunc;
 
     // Forward the pass-through attributes to LLVM.
     if (failed(forwardPassthroughAttributes(function.getLoc(),
                                             function.passthrough(), llvmFunc)))
       return failure();
   }
 
   return success();
 }
 
 LogicalResult ModuleTranslation::convertFunctions() {
-  // Lock access to the llvm context.
-  llvm::sys::SmartScopedLock<true> scopedLock(
-      llvmDialect->getLLVMContextMutex());
-
   // Convert functions.
   for (auto function : getModuleBody(mlirModule).getOps<LLVMFuncOp>()) {
     // Ignore external functions.
     if (function.isExternal())
       continue;
 
     if (failed(convertOneFunction(function)))
       return failure();
   }
 
   return success();
 }
 
 llvm::Type *ModuleTranslation::convertType(LLVMType type) {
-  // Lock the LLVM context as we create types in it.
-  llvm::sys::SmartScopedLock<true> lock(llvmDialect->getLLVMContextMutex());
   return typeTranslator.translateType(type);
 }
 
 /// A helper to look up remapped operands in the value remapping table.`
 SmallVector<llvm::Value *, 8>
 ModuleTranslation::lookupValues(ValueRange values) {
   SmallVector<llvm::Value *, 8> remapped;
   remapped.reserve(values.size());
   for (Value v : values) {
     assert(valueMapping.count(v) && "referencing undefined value");
     remapped.push_back(valueMapping.lookup(v));
   }
   return remapped;
 }
 
-std::unique_ptr<llvm::Module>
-ModuleTranslation::prepareLLVMModule(Operation *m) {
+std::unique_ptr<llvm::Module> ModuleTranslation::prepareLLVMModule(
+    Operation *m, llvm::LLVMContext &llvmContext, StringRef name) {
   auto *dialect = m->getContext()->getRegisteredDialect<LLVM::LLVMDialect>();
   assert(dialect && "LLVM dialect must be registered");
-  // Lock the LLVM context as we might create new types here.
-  llvm::sys::SmartScopedLock<true> scopedLock(dialect->getLLVMContextMutex());
 
-  auto llvmModule = llvm::CloneModule(dialect->getLLVMModule());
-  if (!llvmModule)
-    return nullptr;
-
-  llvm::LLVMContext &llvmContext = llvmModule->getContext();
-  llvm::IRBuilder<> builder(llvmContext);
+  auto llvmModule = std::make_unique<llvm::Module>(name, llvmContext);
+  llvmModule->setDataLayout(dialect->getDataLayout());
 
   // Inject declarations for `malloc` and `free` functions that can be used in
   // memref allocation/deallocation coming from standard ops lowering.
+  llvm::IRBuilder<> builder(llvmContext);
   llvmModule->getOrInsertFunction("malloc", builder.getInt8PtrTy(),
                                   builder.getInt64Ty());
   llvmModule->getOrInsertFunction("free", builder.getVoidTy(),
                                   builder.getInt8PtrTy());
 
   return llvmModule;
 }

mlir/tools/mlir-rocm-runner/mlir-rocm-runner.cpp

if (!hsacoFile) {
     return failure();
   }
   hsacoBlob.assign(hsacoFile->getBuffer().begin(),
                    hsacoFile->getBuffer().end());
 
   return success();
 }
 
-static std::unique_ptr<llvm::Module> compileModuleToROCDLIR(Operation *m) {
-  auto llvmModule = translateModuleToROCDLIR(m);
+static std::unique_ptr<llvm::Module>
+compileModuleToROCDLIR(Operation *m, llvm::LLVMContext &llvmContext,
+                       StringRef name) {
+  auto llvmModule = translateModuleToROCDLIR(m, llvmContext, name);
   // TODO: Link with ROCm-Device-Libs in case needed (ex: the Module
   // depends on math functions).
   return llvmModule;
 }
 
 static OwnedBlob compileISAToHsaco(const std::string isa, Location loc,
                                    StringRef name) {
   // ISA -> ISA in binary form via MC.