From 648dff7723cfd0f272ffa4ab2ecfa3ae264d75a6 Mon Sep 17 00:00:00 2001
From: Peter Bell <peterbell10@openai.com>
Date: Thu, 21 Nov 2024 15:00:08 +0000
Subject: [PATCH] [Build] Add option to limit number of parallel link jobs

git-pr-chain: ci_testing_limit_number_of_parallel_link_a412
---
 CMakeLists.txt  | 7 +++++++
 python/setup.py | 1 +
 2 files changed, 8 insertions(+)

diff --git a/CMakeLists.txt b/CMakeLists.txt
index 6ee47748602f..56564c38964c 100644
--- a/CMakeLists.txt
+++ b/CMakeLists.txt
@@ -44,6 +44,13 @@ if(TRITON_BUILD_WITH_CCACHE)
   endif()
 endif()
 
+set(TRITON_PARALLEL_LINK_JOBS "" CACHE STRING
+  "Define the maximum number of concurrent link jobs (Ninja only).")
+if (TRITON_PARALLEL_LINK_JOBS)
+    set_property(GLOBAL APPEND PROPERTY JOB_POOLS link_job_pool=${TRITON_PARALLEL_LINK_JOBS})
+    set(CMAKE_JOB_POOL_LINK link_job_pool)
+endif()
+
 
 # Ensure Python3 vars are set correctly
 # used conditionally in this file and by lit tests
diff --git a/python/setup.py b/python/setup.py
index 86c4013b0718..607670fc5e17 100644
--- a/python/setup.py
+++ b/python/setup.py
@@ -465,6 +465,7 @@ def build_extension(self, ext):
             "TRITON_BUILD_PROTON",
             "TRITON_BUILD_TUTORIALS",
             "TRITON_BUILD_WITH_CCACHE",
+            "TRITON_PARALLEL_LINK_JOBS",
         ]
         cmake_args += [f"-D{option}={os.getenv(option)}" for option in passthrough_args if option in os.environ]