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

Alex Zinenko llvmlistbot at llvm.org
Fri Aug 7 05:22:39 PDT 2020


Author: Alex Zinenko
Date: 2020-08-07T14:22:30+02:00
New Revision: db1c197bf8247d8dced41ae2f579168c7b54d9ef

URL: https://github.com/llvm/llvm-project/commit/db1c197bf8247d8dced41ae2f579168c7b54d9ef
DIFF: https://github.com/llvm/llvm-project/commit/db1c197bf8247d8dced41ae2f579168c7b54d9ef.diff

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

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.

Reviewed By: rriddle, mehdi_amini

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

Added: 
    

Modified: 
    mlir/docs/Tutorials/Toy/Ch-6.md
    mlir/examples/toy/Ch6/toyc.cpp
    mlir/examples/toy/Ch7/toyc.cpp
    mlir/include/mlir/Conversion/GPUCommon/GPUCommonPass.h
    mlir/include/mlir/Dialect/LLVMIR/LLVMDialect.h
    mlir/include/mlir/Target/LLVMIR.h
    mlir/include/mlir/Target/LLVMIR/ModuleTranslation.h
    mlir/include/mlir/Target/NVVMIR.h
    mlir/include/mlir/Target/ROCDLIR.h
    mlir/lib/Conversion/GPUCommon/ConvertKernelFuncToBlob.cpp
    mlir/lib/Dialect/LLVMIR/IR/LLVMDialect.cpp
    mlir/lib/ExecutionEngine/ExecutionEngine.cpp
    mlir/lib/Target/LLVMIR/ConvertToLLVMIR.cpp
    mlir/lib/Target/LLVMIR/ConvertToNVVMIR.cpp
    mlir/lib/Target/LLVMIR/ConvertToROCDLIR.cpp
    mlir/lib/Target/LLVMIR/LLVMAVX512Intr.cpp
    mlir/lib/Target/LLVMIR/ModuleTranslation.cpp
    mlir/tools/mlir-rocm-runner/mlir-rocm-runner.cpp

Removed: 
    


################################################################################
diff  --git a/mlir/docs/Tutorials/Toy/Ch-6.md b/mlir/docs/Tutorials/Toy/Ch-6.md
index 06f5cd60dca5..86805c10831a 100644
--- a/mlir/docs/Tutorials/Toy/Ch-6.md
+++ b/mlir/docs/Tutorials/Toy/Ch-6.md
@@ -238,14 +238,17 @@ define void @main()
 }
 ```
 
-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;

diff  --git a/mlir/examples/toy/Ch6/toyc.cpp b/mlir/examples/toy/Ch6/toyc.cpp
index 2f984cbe5666..bdcdf1af7ea8 100644
--- a/mlir/examples/toy/Ch6/toyc.cpp
+++ b/mlir/examples/toy/Ch6/toyc.cpp
@@ -189,7 +189,9 @@ int dumpAST() {
 }
 
 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;

diff  --git a/mlir/examples/toy/Ch7/toyc.cpp b/mlir/examples/toy/Ch7/toyc.cpp
index ef2bf8316421..c1cc207a406c 100644
--- a/mlir/examples/toy/Ch7/toyc.cpp
+++ b/mlir/examples/toy/Ch7/toyc.cpp
@@ -190,7 +190,9 @@ int dumpAST() {
 }
 
 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;

diff  --git a/mlir/include/mlir/Conversion/GPUCommon/GPUCommonPass.h b/mlir/include/mlir/Conversion/GPUCommon/GPUCommonPass.h
index 0935786af27a..5cb987081bb5 100644
--- a/mlir/include/mlir/Conversion/GPUCommon/GPUCommonPass.h
+++ b/mlir/include/mlir/Conversion/GPUCommon/GPUCommonPass.h
@@ -35,8 +35,8 @@ class LLVMDialect;
 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.

diff  --git a/mlir/include/mlir/Dialect/LLVMIR/LLVMDialect.h b/mlir/include/mlir/Dialect/LLVMIR/LLVMDialect.h
index 1614a244070f..04700f0aa17d 100644
--- a/mlir/include/mlir/Dialect/LLVMIR/LLVMDialect.h
+++ b/mlir/include/mlir/Dialect/LLVMIR/LLVMDialect.h
@@ -65,12 +65,6 @@ Value createGlobalString(Location loc, OpBuilder &builder, StringRef name,
 /// 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
 

diff  --git a/mlir/include/mlir/Target/LLVMIR.h b/mlir/include/mlir/Target/LLVMIR.h
index d18ed0f643f9..ffd1a4c2d180 100644
--- a/mlir/include/mlir/Target/LLVMIR.h
+++ b/mlir/include/mlir/Target/LLVMIR.h
@@ -13,6 +13,8 @@
 #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.
@@ -31,7 +33,9 @@ class ModuleOp;
 /// 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

diff  --git a/mlir/include/mlir/Target/LLVMIR/ModuleTranslation.h b/mlir/include/mlir/Target/LLVMIR/ModuleTranslation.h
index 8b67c249c099..b87dde307867 100644
--- a/mlir/include/mlir/Target/LLVMIR/ModuleTranslation.h
+++ b/mlir/include/mlir/Target/LLVMIR/ModuleTranslation.h
@@ -50,14 +50,15 @@ class LLVMFuncOp;
 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);
 
@@ -94,7 +95,9 @@ class ModuleTranslation {
   /// 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);
@@ -122,8 +125,6 @@ class ModuleTranslation {
   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;

diff  --git a/mlir/include/mlir/Target/NVVMIR.h b/mlir/include/mlir/Target/NVVMIR.h
index ab41fb841a95..0cd7688e275b 100644
--- a/mlir/include/mlir/Target/NVVMIR.h
+++ b/mlir/include/mlir/Target/NVVMIR.h
@@ -13,10 +13,12 @@
 #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
 
@@ -28,7 +30,9 @@ class Operation;
 /// 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
 

diff  --git a/mlir/include/mlir/Target/ROCDLIR.h b/mlir/include/mlir/Target/ROCDLIR.h
index 45dc5a15c7ac..e2cb812a173d 100644
--- a/mlir/include/mlir/Target/ROCDLIR.h
+++ b/mlir/include/mlir/Target/ROCDLIR.h
@@ -14,10 +14,12 @@
 #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
 
@@ -29,7 +31,9 @@ class Operation;
 /// 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
 

diff  --git a/mlir/lib/Conversion/GPUCommon/ConvertKernelFuncToBlob.cpp b/mlir/lib/Conversion/GPUCommon/ConvertKernelFuncToBlob.cpp
index cf41523d3b29..a42df092b156 100644
--- a/mlir/lib/Conversion/GPUCommon/ConvertKernelFuncToBlob.cpp
+++ b/mlir/lib/Conversion/GPUCommon/ConvertKernelFuncToBlob.cpp
@@ -59,14 +59,11 @@ class GpuKernelToBlobPass
   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();
 
@@ -109,17 +106,12 @@ 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;

diff  --git a/mlir/lib/Dialect/LLVMIR/IR/LLVMDialect.cpp b/mlir/lib/Dialect/LLVMIR/IR/LLVMDialect.cpp
index 876fd05aa1b5..96f2ecd049bb 100644
--- a/mlir/lib/Dialect/LLVMIR/IR/LLVMDialect.cpp
+++ b/mlir/lib/Dialect/LLVMIR/IR/LLVMDialect.cpp
@@ -1794,16 +1794,3 @@ 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));
-}

diff  --git a/mlir/lib/ExecutionEngine/ExecutionEngine.cpp b/mlir/lib/ExecutionEngine/ExecutionEngine.cpp
index 130ca3c02ad4..83db795f8ce8 100644
--- a/mlir/lib/ExecutionEngine/ExecutionEngine.cpp
+++ b/mlir/lib/ExecutionEngine/ExecutionEngine.cpp
@@ -223,7 +223,7 @@ Expected<std::unique_ptr<ExecutionEngine>> ExecutionEngine::create(
       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
@@ -232,12 +232,7 @@ Expected<std::unique_ptr<ExecutionEngine>> ExecutionEngine::create(
   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.
@@ -295,7 +290,7 @@ Expected<std::unique_ptr<ExecutionEngine>> ExecutionEngine::create(
                    .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); }));

diff  --git a/mlir/lib/Target/LLVMIR/ConvertToLLVMIR.cpp b/mlir/lib/Target/LLVMIR/ConvertToLLVMIR.cpp
index 75bcf38aae04..6027ab344a56 100644
--- a/mlir/lib/Target/LLVMIR/ConvertToLLVMIR.cpp
+++ b/mlir/lib/Target/LLVMIR/ConvertToLLVMIR.cpp
@@ -21,15 +21,19 @@
 
 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();
 

diff  --git a/mlir/lib/Target/LLVMIR/ConvertToNVVMIR.cpp b/mlir/lib/Target/LLVMIR/ConvertToNVVMIR.cpp
index 51686ad73666..fc2f650840c1 100644
--- a/mlir/lib/Target/LLVMIR/ConvertToNVVMIR.cpp
+++ b/mlir/lib/Target/LLVMIR/ConvertToNVVMIR.cpp
@@ -65,9 +65,11 @@ class ModuleTranslation : public 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;
 
@@ -98,7 +100,8 @@ 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();
 

diff  --git a/mlir/lib/Target/LLVMIR/ConvertToROCDLIR.cpp b/mlir/lib/Target/LLVMIR/ConvertToROCDLIR.cpp
index 943ceea7f63d..5bd04e4e3ef8 100644
--- a/mlir/lib/Target/LLVMIR/ConvertToROCDLIR.cpp
+++ b/mlir/lib/Target/LLVMIR/ConvertToROCDLIR.cpp
@@ -75,10 +75,12 @@ class ModuleTranslation : public 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.
@@ -102,7 +104,8 @@ 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();
 

diff  --git a/mlir/lib/Target/LLVMIR/LLVMAVX512Intr.cpp b/mlir/lib/Target/LLVMIR/LLVMAVX512Intr.cpp
index 7335fab84bb9..64f9fef91847 100644
--- a/mlir/lib/Target/LLVMIR/LLVMAVX512Intr.cpp
+++ b/mlir/lib/Target/LLVMIR/LLVMAVX512Intr.cpp
@@ -34,9 +34,11 @@ class LLVMAVX512ModuleTranslation : public LLVM::ModuleTranslation {
   }
 };
 
-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
 
@@ -44,7 +46,9 @@ 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();
 

diff  --git a/mlir/lib/Target/LLVMIR/ModuleTranslation.cpp b/mlir/lib/Target/LLVMIR/ModuleTranslation.cpp
index 9275875e2560..f58d02845ccd 100644
--- a/mlir/lib/Target/LLVMIR/ModuleTranslation.cpp
+++ b/mlir/lib/Target/LLVMIR/ModuleTranslation.cpp
@@ -304,7 +304,6 @@ ModuleTranslation::ModuleTranslation(Operation *module,
           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.");
@@ -688,9 +687,6 @@ LogicalResult ModuleTranslation::convertBlock(Block &bb, bool ignoreArguments) {
 /// 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);
@@ -892,10 +888,6 @@ LogicalResult ModuleTranslation::checkSupportedModuleOps(Operation *m) {
 }
 
 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>()) {
@@ -916,10 +908,6 @@ LogicalResult ModuleTranslation::convertFunctionSignatures() {
 }
 
 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.
@@ -934,8 +922,6 @@ LogicalResult ModuleTranslation::convertFunctions() {
 }
 
 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);
 }
 
@@ -951,22 +937,17 @@ ModuleTranslation::lookupValues(ValueRange values) {
   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(),

diff  --git a/mlir/tools/mlir-rocm-runner/mlir-rocm-runner.cpp b/mlir/tools/mlir-rocm-runner/mlir-rocm-runner.cpp
index 5b3d23adcb89..de07a2370857 100644
--- a/mlir/tools/mlir-rocm-runner/mlir-rocm-runner.cpp
+++ b/mlir/tools/mlir-rocm-runner/mlir-rocm-runner.cpp
@@ -196,8 +196,10 @@ static LogicalResult createHsaco(const Blob &isaBlob, StringRef name,
   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;


        


More information about the Mlir-commits mailing list