Revert "[BUILD] Some CMake cleanup/modernisation (#5271)"

This reverts commit 2003685.

Author: whitneywhtsang

CMakeLists.txt

@@ -168,6 +168,10 @@ include_directories(${PROJECT_BINARY_DIR}/third_party) # Tablegen'd files
 add_subdirectory(include)
 add_subdirectory(lib)
 
+# find_package(PythonLibs REQUIRED)
+set(TRITON_SOURCE_DIR "${CMAKE_CURRENT_SOURCE_DIR}")
+set(TRITON_BINARY_DIR "${CMAKE_CURRENT_BINARY_DIR}")
+
 # TODO: Figure out which target is sufficient to fix errors; triton is
 # apparently not enough. Currently set linking libstdc++fs for all targets
 # to support some old version GCC compilers like 8.3.0.
@@ -184,12 +188,24 @@ if(TRITON_BUILD_PYTHON_MODULE)
   set(PYTHON_SRC_PATH ${CMAKE_CURRENT_SOURCE_DIR}/python/src)
   include_directories(${PYTHON_SRC_PATH})
 
-  # Python Interpreter is used to run lit tests
-  find_package(Python3 REQUIRED COMPONENTS Development Interpreter)
-  find_package(pybind11 CONFIG REQUIRED HINTS "${Python3_SITELIB}")
   if(PYTHON_INCLUDE_DIRS)
+    # We have PYTHON_INCLUDE_DIRS set--this is what we expect when building
+    # using pip install.
+    include_directories(${PYTHON_INCLUDE_DIRS})
+    include_directories(${PYBIND11_INCLUDE_DIR})
     message(STATUS "PYTHON_LIB_DIRS ${PYTHON_LIB_DIRS}")
     link_directories(${PYTHON_LIB_DIRS})
+  else()
+    # Otherwise, we might be building from top CMakeLists.txt directly.
+    # Try to find Python and pybind11 packages.
+    find_package(Python3 REQUIRED COMPONENTS Development Interpreter)
+    find_package(pybind11 CONFIG REQUIRED HINTS "${Python3_SITELIB}")
+    include_directories(${Python3_INCLUDE_DIRS})
+    include_directories(${pybind11_INCLUDE_DIR})
+    link_directories(${Python3_LIBRARY_DIRS})
+    link_libraries(${Python3_LIBRARIES})
+    add_link_options(${Python3_LINK_OPTIONS})
+  endif()
 
   if (DEFINED TRITON_PLUGIN_DIRS)
     foreach(PLUGIN_DIR ${TRITON_PLUGIN_DIRS})
@@ -256,9 +272,6 @@ if(TRITON_BUILD_PYTHON_MODULE)
     LLVMAMDGPUCodeGen
     LLVMAMDGPUAsmParser
 
-    Python3::Module
-    pybind11::headers
-
   )
   if(CMAKE_SYSTEM_PROCESSOR MATCHES "aarch64" OR # Linux arm64
      CMAKE_SYSTEM_PROCESSOR MATCHES "arm64" OR # macOS arm64
@@ -299,7 +312,7 @@ if(TRITON_BUILD_PYTHON_MODULE)
                   ${PYTHON_SRC_PATH}/llvm.cc)
 
   # Link triton with its dependencies
-  target_link_libraries(triton PRIVATE ${TRITON_LIBRARIES})
+  target_link_libraries(triton PUBLIC ${TRITON_LIBRARIES})
   if(WIN32)
     target_link_libraries(triton PRIVATE ${CMAKE_DL_LIBS})
     set_target_properties(triton PROPERTIES SUFFIX ".pyd")
@@ -322,7 +335,7 @@ if(TRITON_BUILD_PYTHON_MODULE AND NOT WIN32)
     set(PYTHON_LDFLAGS "-undefined dynamic_lookup")
   endif()
 
-  target_link_options(triton PRIVATE ${PYTHON_LDFLAGS})
+  target_link_libraries(triton PRIVATE ${PYTHON_LDFLAGS})
 endif()
 
 if(NOT TRITON_BUILD_PYTHON_MODULE)

python/setup.py

@@ -453,7 +453,7 @@ def get_pybind11_cmake_args(self):
             pybind11_include_dir = os.path.join(pybind11_sys_path, "include")
         else:
             pybind11_include_dir = pybind11.get_include()
-        return [f"-Dpybind11_INCLUDE_DIR='{pybind11_include_dir}'", f"-Dpybind11_DIR='{pybind11.get_cmake_dir()}'"]
+        return [f"-DPYBIND11_INCLUDE_DIR={pybind11_include_dir}"]
 
     def get_proton_cmake_args(self):
         cmake_args = get_thirdparty_packages([get_json_package_info()])
@@ -493,7 +493,7 @@ def build_extension(self, ext):
             "-DCMAKE_EXPORT_COMPILE_COMMANDS=ON", "-DLLVM_ENABLE_WERROR=ON",
             "-DCMAKE_LIBRARY_OUTPUT_DIRECTORY=" + extdir, "-DTRITON_BUILD_TUTORIALS=OFF",
             "-DTRITON_BUILD_PYTHON_MODULE=ON", "-DPython3_EXECUTABLE:FILEPATH=" + sys.executable,
-            "-DCMAKE_VERBOSE_MAKEFILE:BOOL=ON", "-DPython3_INCLUDE_DIR=" + python_include_dir,
+            "-DCMAKE_VERBOSE_MAKEFILE:BOOL=ON", "-DPYTHON_INCLUDE_DIRS=" + python_include_dir,
             "-DTRITON_CODEGEN_BACKENDS=" + ';'.join([b.name for b in backends if not b.is_external]),
             "-DTRITON_PLUGIN_DIRS=" + ';'.join([b.src_dir for b in backends if b.is_external])
         ]

test/lit.site.cfg.py.in

@@ -2,7 +2,7 @@
 
 import sys
 
-config.triton_obj_root = "@triton_BINARY_DIR@"
+config.triton_obj_root = "@TRITON_BINARY_DIR@"
 config.llvm_src_root = "@LLVM_SOURCE_DIR@"
 config.llvm_obj_root = "@LLVM_BINARY_DIR@"
 config.llvm_tools_dir = "@LLVM_TOOLS_DIR@"
@@ -20,4 +20,4 @@ import lit.llvm
 lit.llvm.initialize(lit_config, config)
 
 # Let the main config do the real work
-lit_config.load_config(config, "@triton_SOURCE_DIR@/test/lit.cfg.py")
+lit_config.load_config(config, "@TRITON_SOURCE_DIR@/test/lit.cfg.py")

third_party/amd/CMakeLists.txt

@@ -4,7 +4,6 @@ add_subdirectory(include)
 add_subdirectory(lib)
 if(TRITON_BUILD_PYTHON_MODULE)
   add_triton_plugin(TritonAMD ${CMAKE_CURRENT_SOURCE_DIR}/python/triton_amd.cc LINK_LIBS TritonAMDGPUToLLVM TritonAMDGPUTransforms TritonAMDGPUDialectToLLVM)
-  target_link_libraries(TritonAMD PRIVATE Python3::Module pybind11::headers)
 endif()
 if(TRITON_BUILD_UT)
   add_subdirectory(unittest)

third_party/nvidia/CMakeLists.txt

@@ -4,7 +4,6 @@ add_subdirectory(include)
 add_subdirectory(lib)
 if(TRITON_BUILD_PYTHON_MODULE)
   add_triton_plugin(TritonNVIDIA ${CMAKE_CURRENT_SOURCE_DIR}/triton_nvidia.cc LINK_LIBS TritonNVIDIAGPUToLLVM NVGPUToLLVM)
-  target_link_libraries(TritonNVIDIA PRIVATE Python3::Module pybind11::headers)
 endif()
 if(TRITON_BUILD_UT)
   add_subdirectory(unittest)

third_party/proton/CMakeLists.txt

@@ -19,8 +19,19 @@ include_directories(${JSON_INCLUDE_DIR})
 include_directories(${PROTON_SRC_DIR}/include)
 include_directories(${PROTON_EXTERN_DIR})
 
-find_package(Python3 REQUIRED Interpreter Development)
-find_package(pybind11 CONFIG REQUIRED HINTS "${Python3_SITELIB}")
+if(PYTHON_INCLUDE_DIRS)
+  # We have PYTHON_INCLUDE_DIRS set--this is what we expect when building
+  # using pip install.
+  include_directories(${PYTHON_INCLUDE_DIRS})
+  include_directories(${PYBIND11_INCLUDE_DIR})
+else()
+  # Otherwise, we might be building from top CMakeLists.txt directly.
+  # Try to find Python and pybind11 packages.
+  find_package(Python3 REQUIRED Interpreter Development)
+  find_package(pybind11 CONFIG REQUIRED HINTS "${Python3_SITELIB}")
+  include_directories(${Python3_INCLUDE_DIRS})
+  include_directories(${pybind11_INCLUDE_DIR})
+endif()
 
 # Check if the platform is MacOS
 if(APPLE)
@@ -38,5 +49,4 @@ include_directories(${CUPTI_INCLUDE_DIR})
 include_directories(SYSTEM ${ROCTRACER_INCLUDE_DIR})
 target_compile_definitions(proton PRIVATE __HIP_PLATFORM_AMD__)
 
-target_link_libraries(proton PRIVATE Python3::Module pybind11::headers)
-target_link_options(proton PRIVATE ${PROTON_PYTHON_LDFLAGS})
+target_link_libraries(proton PRIVATE ${Python_LIBRARIES} ${PROTON_PYTHON_LDFLAGS})

third_party/proton/dialect/CMakeLists.txt

@@ -3,6 +3,5 @@ include_directories(${CMAKE_CURRENT_BINARY_DIR}/include)
 add_subdirectory(include)
 add_subdirectory(lib)
 if(TRITON_BUILD_PYTHON_MODULE)
-  add_triton_plugin(TritonProton ${CMAKE_CURRENT_SOURCE_DIR}/triton_proton.cc)
-  target_link_libraries(TritonProton PRIVATE ProtonIR Python3::Module pybind11::headers)
+  add_triton_plugin(TritonProton ${CMAKE_CURRENT_SOURCE_DIR}/triton_proton.cc LINK_LIBS ProtonIR)
 endif()