From 0113cf00b6f43ebfa902447e6ee23f9b37021a16 Mon Sep 17 00:00:00 2001
From: Jez Ng <jezng@fb.com>
Date: Thu, 25 Mar 2021 14:39:45 -0400
Subject: [PATCH] [lld-macho] Add support for --threads

Code and test are largely identical to the LLD-ELF equivalents.

Reviewed By: MaskRay

Differential Revision: https://reviews.llvm.org/D99312
---
 lld/MachO/Driver.cpp     | 11 +++++++++++
 lld/MachO/Options.td     |  3 +++
 lld/test/MachO/threads.s | 16 ++++++++++++++++
 3 files changed, 30 insertions(+)
 create mode 100644 lld/test/MachO/threads.s

diff --git a/lld/MachO/Driver.cpp b/lld/MachO/Driver.cpp
index 57ff369956c2..35f3ddf8b2a7 100644
--- a/lld/MachO/Driver.cpp
+++ b/lld/MachO/Driver.cpp
@@ -38,6 +38,7 @@
 #include "llvm/Support/FileSystem.h"
 #include "llvm/Support/Host.h"
 #include "llvm/Support/MemoryBuffer.h"
+#include "llvm/Support/Parallel.h"
 #include "llvm/Support/Path.h"
 #include "llvm/Support/TarWriter.h"
 #include "llvm/Support/TargetSelect.h"
@@ -861,6 +862,16 @@ bool macho::link(ArrayRef<const char *> argsArr, bool canExitEarly,
   depTracker =
       make<DependencyTracker>(args.getLastArgValue(OPT_dependency_info, ""));
 
+  if (auto *arg = args.getLastArg(OPT_threads_eq)) {
+    StringRef v(arg->getValue());
+    unsigned threads = 0;
+    if (!llvm::to_integer(v, threads, 0) || threads == 0)
+      error(arg->getSpelling() + ": expected a positive integer, but got '" +
+            arg->getValue() + "'");
+    parallel::strategy = hardware_concurrency(threads);
+    // FIXME: use this to configure ThinLTO concurrency too
+  }
+
   config->entry = symtab->addUndefined(args.getLastArgValue(OPT_e, "_main"),
                                        /*file=*/nullptr,
                                        /*isWeakRef=*/false);
diff --git a/lld/MachO/Options.td b/lld/MachO/Options.td
index 073cb5b11621..cdb8a4676145 100644
--- a/lld/MachO/Options.td
+++ b/lld/MachO/Options.td
@@ -20,6 +20,9 @@ def color_diagnostics_eq: Joined<["--"], "color-diagnostics=">,
     HelpText<"Use colors in diagnostics (default: auto)">,
     MetaVarName<"[auto,always,never]">,
     Group<grp_lld>;
+def threads_eq : Joined<["--"], "threads=">,
+    HelpText<"Number of threads. '1' disables multi-threading. By default all available hardware threads are used">,
+    Group<grp_lld>;
 def reproduce: Separate<["--"], "reproduce">,
     Group<grp_lld>;
 def reproduce_eq: Joined<["--"], "reproduce=">,
diff --git a/lld/test/MachO/threads.s b/lld/test/MachO/threads.s
new file mode 100644
index 000000000000..b3176388de1e
--- /dev/null
+++ b/lld/test/MachO/threads.s
@@ -0,0 +1,16 @@
+# REQUIRES: x86
+# RUN: llvm-mc -filetype=obj -triple=x86_64-apple-darwin %s -o %t.o
+
+## A positive integer is allowed.
+# RUN: %lld --threads=1 %t.o -o /dev/null
+# RUN: %lld --threads=2 %t.o -o /dev/null
+
+# RUN: not %lld --threads=all %t.o -o /dev/null 2>&1 | FileCheck %s -DN=all
+# RUN: not %lld --threads=0 %t.o -o /dev/null 2>&1 | FileCheck %s -DN=0
+# RUN: not %lld --threads=-1 %t.o -o /dev/null 2>&1 | FileCheck %s -DN=-1
+
+# CHECK: error: --threads=: expected a positive integer, but got '[[N]]'
+
+.globl _main
+_main:
+  ret