diff --git a/CMakeLists.txt b/CMakeLists.txt index 69d318d4e0..1a547547cf 100644 --- a/CMakeLists.txt +++ b/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) diff --git a/python/setup.py b/python/setup.py index d06128bd17..6921800f48 100644 --- a/python/setup.py +++ b/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]) ] diff --git a/test/lit.site.cfg.py.in b/test/lit.site.cfg.py.in index fd1fb486dd..4053a8e7df 100644 --- a/test/lit.site.cfg.py.in +++ b/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") diff --git a/third_party/amd/CMakeLists.txt b/third_party/amd/CMakeLists.txt index a09bab8e1c..8228c3d391 100644 --- a/third_party/amd/CMakeLists.txt +++ b/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) diff --git a/third_party/nvidia/CMakeLists.txt b/third_party/nvidia/CMakeLists.txt index bab189bcbd..75f98fa8f7 100644 --- a/third_party/nvidia/CMakeLists.txt +++ b/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) diff --git a/third_party/proton/CMakeLists.txt b/third_party/proton/CMakeLists.txt index ba967c5695..e2d9152c96 100644 --- a/third_party/proton/CMakeLists.txt +++ b/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}) diff --git a/third_party/proton/dialect/CMakeLists.txt b/third_party/proton/dialect/CMakeLists.txt index cfa5938873..c7b5413a0e 100644 --- a/third_party/proton/dialect/CMakeLists.txt +++ b/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()