From 2a3746bac51a78025bb9f516a4765158ceda0e87 Mon Sep 17 00:00:00 2001
From: Justin Lebar <justin.lebar@gmail.com>
Date: Sun, 17 Sep 2023 02:08:04 -0700
Subject: [PATCH] [BUILD] use ninja (#2318)

---
 .github/workflows/integration-tests.yml | 1 +
 README.md                               | 4 ++--
 python/pyproject.toml                   | 2 +-
 python/setup.py                         | 3 +++
 4 files changed, 7 insertions(+), 3 deletions(-)

diff --git a/.github/workflows/integration-tests.yml b/.github/workflows/integration-tests.yml
index c44737a312af..e8c2f41b1eb3 100644
--- a/.github/workflows/integration-tests.yml
+++ b/.github/workflows/integration-tests.yml
@@ -65,6 +65,7 @@ jobs:
           cd python
           python3 -m pip install --upgrade pip
           python3 -m pip install cmake==3.24
+          python3 -m pip install ninja
           python3 -m pip install --no-build-isolation -vvv '.[tests]'
           python3 -m pip install pytest-xdist
 
diff --git a/README.md b/README.md
index 675a41ce4d6e..dbbe18dd795f 100644
--- a/README.md
+++ b/README.md
@@ -41,7 +41,7 @@ git checkout triton-mlir
 # Build
 ```
 cd python
-pip3 install cmake; # build time dependency
+pip3 install ninja cmake; # build time dependencies
 pip3 install -e .
 ```
 # Run tests:
@@ -60,7 +60,7 @@ lit -v test
 ```
 git clone https://github.com/openai/triton.git;
 cd triton/python;
-pip install cmake; # build-time dependency
+pip install ninja cmake; # build-time dependencies
 pip install -e .
 ```
 
diff --git a/python/pyproject.toml b/python/pyproject.toml
index 6430c0c154dc..8bd8093e720d 100644
--- a/python/pyproject.toml
+++ b/python/pyproject.toml
@@ -1,6 +1,6 @@
 
 [build-system]
-requires = ["setuptools>=40.8.0", "wheel", "cmake>=3.18"]
+requires = ["setuptools>=40.8.0", "wheel", "cmake>=3.18", "ninja>=1.11.1"]
 
 [tool.autopep8]
 aggressive = 1
diff --git a/python/setup.py b/python/setup.py
index 69479ceb13bb..c1e33e004f83 100644
--- a/python/setup.py
+++ b/python/setup.py
@@ -202,6 +202,7 @@ def get_cmake_dir(self):
 
     def build_extension(self, ext):
         lit_dir = shutil.which('lit')
+        ninja_dir = shutil.which('ninja')
         user_home = os.getenv("HOME") or os.getenv("USERPROFILE") or \
             os.getenv("HOMEPATH") or None
         if not user_home:
@@ -216,6 +217,8 @@ def build_extension(self, ext):
         # python directories
         python_include_dir = sysconfig.get_path("platinclude")
         cmake_args = [
+            "-G", "Ninja",  # Ninja is much faster than make
+            "-DCMAKE_MAKE_PROGRAM=" + ninja_dir,  # Pass explicit path to ninja otherwise cmake may cache a temporary path
             "-DCMAKE_EXPORT_COMPILE_COMMANDS=ON",
             "-DLLVM_ENABLE_WERROR=ON",
             "-DCMAKE_LIBRARY_OUTPUT_DIRECTORY=" + extdir,