Fix lldMain access to the LLVM command line

Author: draganmladjenovic

tensorflow/compiler/mlir/tools/kernel_gen/hlo_to_kernel.cc

@@ -127,17 +127,17 @@ absl::StatusOr<std::string> EmitToBinary(llvm::StringRef host_triple,
   return ostream.str().str();
 }
 
-absl::Status Run(llvm::StringRef input_file, llvm::StringRef output_file,
-                 llvm::StringRef host_triple,
-                 llvm::ArrayRef<std::string> architectures,
-                 llvm::ArrayRef<int64_t> tile_sizes,
-                 llvm::ArrayRef<int64_t> unroll_factors, bool print_ptx,
+absl::Status Run(std::string input_file, std::string output_file,
+                 std::string host_triple,
+                 std::vector<std::string> architectures,
+                 std::vector<int64_t> tile_sizes,
+                 std::vector<int64_t> unroll_factors, bool print_ptx,
                  bool print_llvmir, bool enable_ftz, bool index_64bit,
                  bool jit_compile, bool jit_i64_indexed_for_large_tensors) {
   // Read TF code.
   std::string hlo_code;
   TF_RETURN_IF_ERROR(
-      ReadFileToString(Env::Default(), input_file.str(), &hlo_code));
+      ReadFileToString(Env::Default(), input_file, &hlo_code));
 
   // Compile.
   mlir::DialectRegistry registry;
@@ -160,7 +160,7 @@ absl::Status Run(llvm::StringRef input_file, llvm::StringRef output_file,
 
   // Write .a file.
   TF_RETURN_IF_ERROR(
-      WriteStringToFile(Env::Default(), output_file.str(), binary));
+      WriteStringToFile(Env::Default(), output_file, binary));
   return absl::OkStatus();
 }
 

third_party/xla/xla/service/gpu/llvm_gpu_backend/amdgpu_backend.cc

@@ -194,7 +194,8 @@ void HsacoCache::Add(const std::string& ir, uint64_t hash,
 // TargetMachine for the AMDGPU target.
 absl::StatusOr<std::vector<uint8_t>> EmitModuleToHsaco(
     llvm::Module* module, llvm::TargetMachine* target_machine,
-    const DebugOptions& debug_options) {
+    const DebugOptions& debug_options,
+    llvm_ir::LLVMCommandLineOptionsLock& llvm_lock) {
   auto* env = tsl::Env::Default();
   std::vector<std::string> tempdir_vector;
   env->GetLocalTempDirectories(&tempdir_vector);
@@ -258,8 +259,6 @@ absl::StatusOr<std::vector<uint8_t>> EmitModuleToHsaco(
 
   if (debug_options.xla_gpu_use_inprocess_lld()) {
 #ifdef HAS_SUPPORT_FOR_LLD_AS_A_LIBRARY
-    static absl::Mutex lld_mu(absl::kConstInit);
-
     std::array<const char*, 7> args{
         "ld.lld",           "--threads=1",       "-shared",
         "--no-undefined",   isabin_path.c_str(), "-o",
@@ -270,7 +269,7 @@ absl::StatusOr<std::vector<uint8_t>> EmitModuleToHsaco(
     llvm::raw_string_ostream os(error_message);
     lld::Result result;
     {
-      absl::MutexLock lock(&lld_mu);
+      llvm_lock.UpgradeToExclusiveAccessToRawLLVMCommandLine();
       result =
           lld::lldMain(args, llvm::nulls(), os, {{lld::Gnu, &lld::elf::link}});
     }
@@ -640,7 +639,7 @@ absl::StatusOr<std::vector<uint8_t>> CompileToHsaco(
 
     // Lower optimized LLVM module to HSA code object.
     TF_ASSIGN_OR_RETURN(
-        hsaco, EmitModuleToHsaco(module, target_machine.get(), debug_options));
+        hsaco, EmitModuleToHsaco(module, target_machine.get(), debug_options, llvm_lock));
     HsacoCache::Add(str, hash, gcn_arch_name, hsaco);
   }
   return hsaco;

third_party/xla/xla/service/llvm_ir/llvm_command_line_options.cc

@@ -20,6 +20,7 @@ limitations under the License.
 #include <vector>
 
 #include "absl/algorithm/container.h"
+#include "absl/base/macros.h"
 #include "absl/base/thread_annotations.h"
 #include "absl/hash/hash.h"
 #include "absl/log/check.h"
@@ -89,5 +90,26 @@ LLVMCommandLineOptionsLock::~LLVMCommandLineOptionsLock() {
   num_active_clients_ -= 1;
 }
 
+void LLVMCommandLineOptionsLock::
+    UpgradeToExclusiveAccessToRawLLVMCommandLine() {
+  {
+    absl::MutexLock lock(lock_);
+    CHECK_GT(num_active_clients_, 0);
+    if (ABSL_PREDICT_TRUE(num_active_clients_ == 1)) {
+      active_client_signature_ = 0;
+      return;
+    }
+    num_active_clients_ -= 1;
+  }
+  // Slow path
+  auto no_other_clients = []() ABSL_EXCLUSIVE_LOCKS_REQUIRED(lock_) {
+    return num_active_clients_ == 0;
+  };
+  lock_.LockWhen(absl::Condition(&no_other_clients));
+  active_client_signature_ = 0;
+  num_active_clients_ = 1;
+  lock_.unlock();
+}
+
 }  // namespace llvm_ir
 }  // namespace xla

third_party/xla/xla/service/llvm_ir/llvm_command_line_options.h

@@ -76,6 +76,8 @@ class ABSL_SCOPED_LOCKABLE LLVMCommandLineOptionsLock {
 
   ~LLVMCommandLineOptionsLock() ABSL_UNLOCK_FUNCTION();
 
+  void UpgradeToExclusiveAccessToRawLLVMCommandLine();
+
   static std::vector<std::string>& GetGlobalOptions() {
     // absl::NoDestructor is not available in OSS XLA.
     static std::vector<std::string>* global_options =