From 4bf509889bae68ef20668ff6edadbb293c7e431f Mon Sep 17 00:00:00 2001
From: Keren Zhou <kerenzhou@openai.com>
Date: Fri, 1 Jul 2022 12:17:22 -0700
Subject: [PATCH] [BUILD] Change the default build type to Release (#571)

---
 CMakeLists.txt  |  5 -----
 python/setup.py | 23 +++++++++++++++--------
 2 files changed, 15 insertions(+), 13 deletions(-)

diff --git a/CMakeLists.txt b/CMakeLists.txt
index c3aadf9c7..e2bc7f309 100644
--- a/CMakeLists.txt
+++ b/CMakeLists.txt
@@ -3,11 +3,6 @@ include(ExternalProject)
 
 set(CMAKE_CXX_STANDARD 17)
 
-if(NOT TRITON_LLVM_BUILD_DIR)
-    set(TRITON_LLVM_BUILD_DIR ${CMAKE_BINARY_DIR})
-endif()
-
-
 project(triton)
 include(CTest)
 if(NOT WIN32)
diff --git a/python/setup.py b/python/setup.py
index af1fa3068..7ed6ab444 100644
--- a/python/setup.py
+++ b/python/setup.py
@@ -7,7 +7,6 @@ import shutil
 import subprocess
 import sys
 import tarfile
-import tempfile
 import urllib.request
 from distutils.version import LooseVersion
 
@@ -15,6 +14,20 @@ from setuptools import Extension, setup
 from setuptools.command.build_ext import build_ext
 
 
+# Taken from https://github.com/pytorch/pytorch/blob/master/tools/setup_helpers/env.py
+def check_env_flag(name: str, default: str = "") -> bool:
+    return os.getenv(name, default).upper() in ["ON", "1", "YES", "TRUE", "Y"]
+
+
+def get_build_type():
+    if check_env_flag("DEBUG"):
+        return "Debug"
+    elif check_env_flag("REL_WITH_DEB_INFO"):
+        return "RelWithDebInfo"
+    else:
+        return "Release"
+
+
 def get_llvm():
     # tries to find system LLVM
     versions = ['-11.0', '-11', '-11-64']
@@ -80,15 +93,10 @@ class CMakeBuild(build_ext):
 
     def build_extension(self, ext):
         llvm_include_dir, llvm_library_dir = get_llvm()
-        self.debug = True
         extdir = os.path.abspath(os.path.dirname(self.get_ext_fullpath(ext.path)))
         # create build directories
-        build_suffix = 'debug' if self.debug else 'release'
-        llvm_build_dir = os.path.join(tempfile.gettempdir(), "llvm-" + build_suffix)
         if not os.path.exists(self.build_temp):
             os.makedirs(self.build_temp)
-        if not os.path.exists(llvm_build_dir):
-            os.makedirs(llvm_build_dir)
         # python directories
         python_include_dirs = [distutils.sysconfig.get_python_inc()] + ['/usr/local/cuda/include']
         cmake_args = [
@@ -99,11 +107,10 @@ class CMakeBuild(build_ext):
             "-DLLVM_LIBRARY_DIR=" + llvm_library_dir,
             # '-DPYTHON_EXECUTABLE=' + sys.executable,
             # '-DCMAKE_VERBOSE_MAKEFILE:BOOL=ON',
-            "-DTRITON_LLVM_BUILD_DIR=" + llvm_build_dir,
             "-DPYTHON_INCLUDE_DIRS=" + ";".join(python_include_dirs)
         ]
         # configuration
-        cfg = "Debug" if self.debug else "Release"
+        cfg = get_build_type()
         build_args = ["--config", cfg]
 
         if platform.system() == "Windows":