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

[krzysz00]

Summary

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.

[tschuett]

mlir/lib/Dialect/GPU/Transforms/SerializeToCubin.cpp

83

AMD or NVIDIA?

[fmorac]

mlir/lib/Dialect/GPU/Transforms/SerializeToCubin.cpp

79

Shouldn't this be llvm::once_flag?, they aren't always alias:
https://github.com/llvm/llvm-project/blob/main/llvm/include/llvm/Support/Threading.h#L56C1-L72C7

mlir/lib/Dialect/GPU/Transforms/SerializeToHsaco.cpp

127

same here.

[krzysz00]

Fix typos, migration issues

[fmorac]

Looks good to me.

[fmorac]

Please fix the git-clang-format fail on the build.

[krzysz00]

Fixed clang-format issue caused by using old clang-format version

mlir/lib/Dialect/GPU/Transforms/SerializeToCubin.cpp

 #include "mlir/Dialect/GPU/Transforms/Passes.h"
 #include "llvm/Support/Debug.h"
 
 #if MLIR_GPU_TO_CUBIN_PASS_ENABLE
 #include "mlir/Pass/Pass.h"
 #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>
 
 using namespace mlir;
 
 static void emitCudaError(const llvm::Twine &expr, const char *buffer,
                           CUresult result, Location loc) {
   const char *error;
if (auto status = (expr)) {                                                \
       emitCudaError(#expr, jitErrorBuffer, status, loc);                       \
       return {};                                                               \
     }                                                                          \
   } while (false)
 
 namespace {
 class SerializeToCubinPass
     : public PassWrapper<SerializeToCubinPass, gpu::SerializeToBlobPass> {
+  static llvm::once_flag initializeBackendOnce;
+
 public:
   MLIR_DEFINE_EXPLICIT_INTERNAL_INLINE_TYPE_ID(SerializeToCubinPass)
 
   SerializeToCubinPass(StringRef triple = "nvptx64-nvidia-cuda",
                        StringRef chip = "sm_35", StringRef features = "+ptx60",
                        int optLevel = 2, bool dumpPtx = false);
 
   StringRef getArgument() const override { return "gpu-to-cubin"; }

 } // namespace
 
 // Sets the 'option' to 'value' unless it already has a value.
 static void maybeSetOption(Pass::Option<std::string> &option, StringRef value) {
   if (!option.hasValue())
     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);
   this->dumpPtx = dumpPtx;
   if (this->optLevel.getNumOccurrences() == 0)
     this->optLevel.setValue(optLevel);
 }
 
SerializeToCubinPass::serializeISA(const std::string &isa) {
   RETURN_ON_CUDA_ERROR(cuLinkDestroy(linkState));
   RETURN_ON_CUDA_ERROR(cuCtxDestroy(context));
 
   return result;
 }
 
 // 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,
                                                           StringRef arch,
                                                           StringRef features,
                                                           int optLevel,
                                                           bool dumpPtx) {
   return std::make_unique<SerializeToCubinPass>(triple, arch, features,
                                                 optLevel, dumpPtx);
 }
 
 #else  // MLIR_GPU_TO_CUBIN_PASS_ENABLE
 void mlir::registerGpuSerializeToCubinPass() {}
 #endif // MLIR_GPU_TO_CUBIN_PASS_ENABLE

mlir/lib/Dialect/GPU/Transforms/SerializeToHsaco.cpp

 #include "llvm/MC/TargetRegistry.h"
 
 #include "llvm/Support/CommandLine.h"
 #include "llvm/Support/FileUtilities.h"
 #include "llvm/Support/Path.h"
 #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"
 #include "llvm/Target/TargetOptions.h"
 
 #include "llvm/Transforms/IPO/Internalize.h"
 
-#include <mutex>
 #include <optional>
 
 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)
 
   SerializeToHsacoPass(StringRef triple, StringRef arch, StringRef features,
                        int optLevel);
   SerializeToHsacoPass(const SerializeToHsacoPass &other);
   StringRef getArgument() const override { return "gpu-to-hsaco"; }
   StringRef getDescription() const override {

 
 // Sets the 'option' to 'value' unless it already has a value.
 static void maybeSetOption(Pass::Option<std::string> &option,
                            function_ref<std::string()> getValue) {
   if (!option.hasValue())
     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(); });
   if (this->optLevel.getNumOccurrences() == 0)
     this->optLevel.setValue(optLevel);
 }
 
 void SerializeToHsacoPass::getDependentDialects(
SerializeToHsacoPass::serializeISA(const std::string &isa) {
   if (!isaBinary)
     return {};
   return createHsaco(*isaBinary);
 }
 
 // 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);
   });
 }
 
 /// Create an instance of the GPU kernel function to HSAco binary serialization
 /// pass.
 std::unique_ptr<Pass> mlir::createGpuSerializeToHsacoPass(StringRef triple,