diff --git a/CMakeLists.txt b/CMakeLists.txt index 65dd8e9867..4a3653f761 100644 --- a/CMakeLists.txt +++ b/CMakeLists.txt @@ -272,7 +272,7 @@ endfunction() # Disable warnings that show up in external code (gtest;pybind11) if(NOT MSVC) set(TRITON_DISABLE_EH_RTTI_FLAGS "$<$:-fno-exceptions;-fno-rtti>") - set(CMAKE_CXX_FLAGS "${CMAKE_CXX_FLAGS} -Werror -Wno-covered-switch-default -fvisibility=hidden") + set(CMAKE_CXX_FLAGS "${CMAKE_CXX_FLAGS} -Werror -Wno-covered-switch-default") else() set(CMAKE_CXX_FLAGS "${CMAKE_CXX_FLAGS} /wd4244 /wd4624 /wd4715 /wd4530") endif() @@ -419,8 +419,19 @@ if(TRITON_BUILD_PYTHON_MODULE) ${PYTHON_SRC_PATH}/llvm.cc ${PYTHON_SRC_PATH}/specialize.cc) + # Fixes error: ‘GluonLayouts’ declared with greater visibility than the type of its field + if (NOT MSVC) + target_compile_options(triton PRIVATE -Wno-error=attributes) + endif() + # Link triton with its dependencies target_link_libraries(triton PRIVATE ${TRITON_LIBRARIES}) + + # Do not propagate libraries that libtriton depends on. This ensures that + # targets that link against libtriton do not accidentally link in their own + # copies of core Triton code and LLVM. + set_target_properties(triton PROPERTIES INTERFACE_LINK_LIBRARIES "") + if(WIN32) target_link_libraries(triton PRIVATE ${CMAKE_DL_LIBS}) set_target_properties(triton PROPERTIES SUFFIX ".pyd") @@ -449,10 +460,8 @@ if(TRITON_BUILD_PYTHON_MODULE) "${TRITON_WHEEL_DIR}/FileCheck${CMAKE_EXECUTABLE_SUFFIX}" COPYONLY) -endif() - -if (UNIX AND NOT APPLE) - set(CMAKE_SHARED_LINKER_FLAGS "${CMAKE_SHARED_LINKER_FLAGS} -Wl,--exclude-libs,ALL") + # Only build plugins when building libtriton since they depend on libtriton. + add_subdirectory(examples/plugins) endif() if(TRITON_BUILD_PYTHON_MODULE AND NOT WIN32) @@ -478,7 +487,6 @@ find_package(Threads REQUIRED) add_subdirectory(third_party/f2reduce) add_subdirectory(bin) add_subdirectory(test) -add_subdirectory(examples) if(TRITON_BUILD_UT) add_subdirectory(unittest) diff --git a/examples/CMakeLists.txt b/examples/CMakeLists.txt deleted file mode 100644 index 0e89371e07..0000000000 --- a/examples/CMakeLists.txt +++ /dev/null @@ -1 +0,0 @@ -add_subdirectory(plugins) diff --git a/examples/plugins/CMakeLists.txt b/examples/plugins/CMakeLists.txt index 6b4a14952b..a48298c8bb 100644 --- a/examples/plugins/CMakeLists.txt +++ b/examples/plugins/CMakeLists.txt @@ -24,7 +24,10 @@ foreach( plugin ${TRITON_PLUGIN_PASSES} ) TritonCanonicalizeIncGen TritonPluginsIncGen ) - target_link_libraries(${plugin} PRIVATE MLIRPass) + target_link_libraries(${plugin} PRIVATE + $<$:MLIRPass> + $<$>:triton> + ) # CMAKE_LIBRARY_OUTPUT_DIRECTORY is only set during the Python # build. It is empty if building directly from the root diff --git a/examples/plugins/DialectPlugins/DialectPlugin/lib/DialectPlugin/CMakeLists.txt b/examples/plugins/DialectPlugins/DialectPlugin/lib/DialectPlugin/CMakeLists.txt index 2e02718000..a51fb24edc 100644 --- a/examples/plugins/DialectPlugins/DialectPlugin/lib/DialectPlugin/CMakeLists.txt +++ b/examples/plugins/DialectPlugins/DialectPlugin/lib/DialectPlugin/CMakeLists.txt @@ -20,10 +20,13 @@ add_mlir_dialect_library(MLIRDialectPlugin MLIRDialectPluginPassesIncGen LINK_LIBS PUBLIC - MLIRPass - LLVMSupport - MLIRSupport - TritonNVIDIAGPUToLLVM + $<$: + MLIRPass + LLVMSupport + MLIRSupport + TritonNVIDIAGPUToLLVM + > + $<$>:triton> "$<$:-undefined dynamic_lookup>" ) diff --git a/python/test/unit/plugins/test_dialect_plugin.py b/python/test/unit/plugins/test_dialect_plugin.py index 55ba6c3695..8614615150 100644 --- a/python/test/unit/plugins/test_dialect_plugin.py +++ b/python/test/unit/plugins/test_dialect_plugin.py @@ -2,6 +2,7 @@ import subprocess import pathlib import pytest +import re from triton._internal_testing import is_cuda, is_hip, is_hip_cdna2 @@ -9,8 +10,6 @@ def test_override(tmp_path: pathlib.Path): - if os.environ.get('LLVM_BUILD_SHARED_LIBS', '0') == '0': - return dir_path = os.path.dirname(os.path.realpath(__file__)) # Run once to get the file dumps @@ -51,8 +50,10 @@ def test_override(tmp_path: pathlib.Path): # # Add ttgir instrumentation with open(filename, "w") as outfile: for line in file_str: - if "tt.get_program_id x" in line: - line = ' %pid_base = arith.constant 0 : i32\n %pid = plugin.magic %pid_base : i32\n' + match = re.search(r'(%\w+)\s*=\s*tt\.get_program_id\s+x', line) + if match: + ssa_name = match.group(1) + line = f' %pid_base = arith.constant 0 : i32\n {ssa_name} = plugin.magic %pid_base : i32\n' outfile.write(line) # # # Run again with kernel override diff --git a/python/test/unit/plugins/test_plugin.py b/python/test/unit/plugins/test_plugin.py index 9a895174b1..3687420a4a 100644 --- a/python/test/unit/plugins/test_plugin.py +++ b/python/test/unit/plugins/test_plugin.py @@ -1,11 +1,11 @@ import torch import pytest -import os import triton import triton.language as tl from triton import knobs +from triton._internal_testing import is_hip import custom_stages @@ -21,10 +21,8 @@ def kernel2(BLOCK_SIZE: tl.constexpr): return +@pytest.mark.skipif(is_hip(), reason="plugin not supported/tested on AMD yet") def test_op(capfd, device: str): - if os.environ.get('LLVM_BUILD_SHARED_LIBS', '0') == '0': - return - size = 98432 x = torch.rand(size, device=device) output = torch.empty_like(x)