[Mlir-commits] [mlir] db647f5 - [mlir][GPU] Initialize LLVM exactly once during GPU compiles

[Krzysztof Drewniak]

Author: Krzysztof Drewniak
Date: 2023-07-14T19:10:52Z
New Revision: db647f5bd8fd337260bf636897f45ecfac8eb4cb

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

LOG: [mlir][GPU] Initialize LLVM exactly once during GPU compiles

No matter how one constructs their SerializeTo* pass, we want to
ensure that the LLVM initialization code runs once and only once. This
commit adds a static once_flag to ensure that.

I've run into mysterious segfaults when calling MLIR GPU compiles from
multiple threads, and this commit is a potential fix for the issue.

Reviewed By: fmorac

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

Added: 
    

Modified: 
    mlir/lib/Dialect/GPU/Transforms/SerializeToCubin.cpp
    mlir/lib/Dialect/GPU/Transforms/SerializeToHsaco.cpp

Removed: 
    


################################################################################
diff  --git a/mlir/lib/Dialect/GPU/Transforms/SerializeToCubin.cpp b/mlir/lib/Dialect/GPU/Transforms/SerializeToCubin.cpp
index 172f006c7b6fb1..0077debe26a0b2 100644
--- a/mlir/lib/Dialect/GPU/Transforms/SerializeToCubin.cpp
+++ b/mlir/lib/Dialect/GPU/Transforms/SerializeToCubin.cpp
@@ -19,6 +19,7 @@
 #include "mlir/Target/LLVMIR/Dialect/NVVM/NVVMToLLVMIRTranslation.h"
 #include "mlir/Target/LLVMIR/Export.h"
 #include "llvm/Support/TargetSelect.h"
+#include "llvm/Support/Threading.h"
 
 #include <cuda.h>
 
@@ -46,6 +47,8 @@ static void emitCudaError(const llvm::Twine &expr, const char *buffer,
 namespace {
 class SerializeToCubinPass
     : public PassWrapper<SerializeToCubinPass, gpu::SerializeToBlobPass> {
+  static llvm::once_flag initializeBackendOnce;
+
 public:
   MLIR_DEFINE_EXPLICIT_INTERNAL_INLINE_TYPE_ID(SerializeToCubinPass)
 
@@ -73,9 +76,21 @@ static void maybeSetOption(Pass::Option<std::string> &option, StringRef value) {
     option = value.str();
 }
 
+llvm::once_flag SerializeToCubinPass::initializeBackendOnce;
+
 SerializeToCubinPass::SerializeToCubinPass(StringRef triple, StringRef chip,
                                            StringRef features, int optLevel,
                                            bool dumpPtx) {
+  // No matter how this pass is constructed, ensure that the NVPTX backend
+  // is initialized exactly once.
+  llvm::call_once(initializeBackendOnce, []() {
+    // Initialize LLVM NVPTX backend.
+    LLVMInitializeNVPTXTarget();
+    LLVMInitializeNVPTXTargetInfo();
+    LLVMInitializeNVPTXTargetMC();
+    LLVMInitializeNVPTXAsmPrinter();
+  });
+
   maybeSetOption(this->triple, triple);
   maybeSetOption(this->chip, chip);
   maybeSetOption(this->features, features);
@@ -144,15 +159,8 @@ SerializeToCubinPass::serializeISA(const std::string &isa) {
 
 // Register pass to serialize GPU kernel functions to a CUBIN binary annotation.
 void mlir::registerGpuSerializeToCubinPass() {
-  PassRegistration<SerializeToCubinPass> registerSerializeToCubin([] {
-    // Initialize LLVM NVPTX backend.
-    LLVMInitializeNVPTXTarget();
-    LLVMInitializeNVPTXTargetInfo();
-    LLVMInitializeNVPTXTargetMC();
-    LLVMInitializeNVPTXAsmPrinter();
-
-    return std::make_unique<SerializeToCubinPass>();
-  });
+  PassRegistration<SerializeToCubinPass> registerSerializeToCubin(
+      [] { return std::make_unique<SerializeToCubinPass>(); });
 }
 
 std::unique_ptr<Pass> mlir::createGpuSerializeToCubinPass(StringRef triple,

diff  --git a/mlir/lib/Dialect/GPU/Transforms/SerializeToHsaco.cpp b/mlir/lib/Dialect/GPU/Transforms/SerializeToHsaco.cpp
index 108b8ab945facb..e204a232ed0abe 100644
--- a/mlir/lib/Dialect/GPU/Transforms/SerializeToHsaco.cpp
+++ b/mlir/lib/Dialect/GPU/Transforms/SerializeToHsaco.cpp
@@ -47,6 +47,7 @@
 #include "llvm/Support/Program.h"
 #include "llvm/Support/SourceMgr.h"
 #include "llvm/Support/TargetSelect.h"
+#include "llvm/Support/Threading.h"
 #include "llvm/Support/WithColor.h"
 
 #include "llvm/Target/TargetMachine.h"
@@ -54,7 +55,6 @@
 
 #include "llvm/Transforms/IPO/Internalize.h"
 
-#include <mutex>
 #include <optional>
 
 using namespace mlir;
@@ -62,6 +62,8 @@ using namespace mlir;
 namespace {
 class SerializeToHsacoPass
     : public PassWrapper<SerializeToHsacoPass, gpu::SerializeToBlobPass> {
+  static llvm::once_flag initializeBackendOnce;
+
 public:
   MLIR_DEFINE_EXPLICIT_INTERNAL_INLINE_TYPE_ID(SerializeToHsacoPass)
 
@@ -122,8 +124,20 @@ static void maybeSetOption(Pass::Option<std::string> &option,
     option = getValue();
 }
 
+llvm::once_flag SerializeToHsacoPass::initializeBackendOnce;
+
 SerializeToHsacoPass::SerializeToHsacoPass(StringRef triple, StringRef arch,
                                            StringRef features, int optLevel) {
+  // No matter how this pass is constructed, ensure that the AMDGPU backend
+  // is initialized exactly once.
+  llvm::call_once(initializeBackendOnce, []() {
+    // Initialize LLVM AMDGPU backend.
+    LLVMInitializeAMDGPUAsmParser();
+    LLVMInitializeAMDGPUAsmPrinter();
+    LLVMInitializeAMDGPUTarget();
+    LLVMInitializeAMDGPUTargetInfo();
+    LLVMInitializeAMDGPUTargetMC();
+  });
   maybeSetOption(this->triple, [&triple] { return triple.str(); });
   maybeSetOption(this->chip, [&arch] { return arch.str(); });
   maybeSetOption(this->features, [&features] { return features.str(); });
@@ -437,13 +451,6 @@ SerializeToHsacoPass::serializeISA(const std::string &isa) {
 // Register pass to serialize GPU kernel functions to a HSACO binary annotation.
 void mlir::registerGpuSerializeToHsacoPass() {
   PassRegistration<SerializeToHsacoPass> registerSerializeToHSACO([] {
-    // Initialize LLVM AMDGPU backend.
-    LLVMInitializeAMDGPUAsmParser();
-    LLVMInitializeAMDGPUAsmPrinter();
-    LLVMInitializeAMDGPUTarget();
-    LLVMInitializeAMDGPUTargetInfo();
-    LLVMInitializeAMDGPUTargetMC();
-
     return std::make_unique<SerializeToHsacoPass>("amdgcn-amd-amdhsa", "", "",
                                                   2);
   });