Skip to content

Commit 2560695

Browse files
zhzhcookieflagtree-bot
zhzhcookie
and
flagtree-bot
authored
[SPEC] Spec enflame cpp and cleanup enflame cpp patch (#940)
--------- Co-authored-by: flagtree-bot <flagtree_ai@163.com>
1 parent bf12939 commit 2560695

20 files changed

Lines changed: 3751 additions & 1346 deletions

File tree

‎cmake/FlagTreeOptions.cmake‎

Lines changed: 17 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -454,6 +454,14 @@ function(flagtree_add_tle_generated_header_dependencies)
454454
list(APPEND _flagtree_tle_codegen_deps TritonTLETransformsIncGen)
455455
endif()
456456
457+
set(_flagtree_enflame_tle_header_targets
458+
MLIRTritonToGCU_gcu300
459+
MLIRTritonToGCU_gcu400
460+
MLIRGCUTritonToTritonGPU_gcu400
461+
MLIRTritonGCUTransforms_gcu400
462+
triton_gcu300_core
463+
triton_gcu400_core)
464+
457465
# Native compiler targets include TLE generated headers under __TLE__ guards.
458466
# The TLE dialect is added after the core libraries, so the dependency must be
459467
# attached explicitly once the TLE tablegen targets exist; otherwise a clean
@@ -466,16 +474,21 @@ function(flagtree_add_tle_generated_header_dependencies)
466474
TritonNVIDIAGPUToLLVM
467475
TritonGPUToLLVM
468476
NVHopperTransforms
477+
${_flagtree_enflame_tle_header_targets}
469478
triton
470479
triton-opt
471480
triton-reduce
472481
triton-lsp
473482
triton-llvm-opt
474483
triton-tensor-layout)
475-
if(TARGET ${_flagtree_tle_header_target})
476-
add_dependencies(${_flagtree_tle_header_target}
477-
${_flagtree_tle_codegen_deps})
478-
endif()
484+
foreach(_flagtree_tle_dependency_target IN ITEMS
485+
${_flagtree_tle_header_target}
486+
obj.${_flagtree_tle_header_target})
487+
if(TARGET ${_flagtree_tle_dependency_target})
488+
add_dependencies(${_flagtree_tle_dependency_target}
489+
${_flagtree_tle_codegen_deps})
490+
endif()
491+
endforeach()
479492
endforeach()
480493
endfunction()
481494

‎include/triton/Dialect/TritonGPU/Transforms/CMakeLists.txt‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -1,3 +1,3 @@
1-
set(LLVM_TARGET_DEFINITIONS Passes.td)
1+
flagtree_spec_td_set(LLVM_TARGET_DEFINITIONS Passes.td)
22
mlir_tablegen(Passes.h.inc -gen-pass-decls -name TritonGPU)
33
add_public_tablegen_target(TritonGPUTransformsIncGen)

‎python/setup_tools/utils/enflame.py‎

Lines changed: 0 additions & 13 deletions
Original file line numberDiff line numberDiff line change
@@ -85,19 +85,6 @@ def install_extension(*args, **kargs):
8585

8686
project_root_dir = cmake_dir.parent.parent
8787

88-
# Modify nvidia driver's is_active() to return False for enflame backend
89-
drvfile = project_root_dir / 'third_party' / 'nvidia' / 'backend' / 'driver.py'
90-
if drvfile.exists():
91-
with open(drvfile, 'r') as f:
92-
lines = f.readlines()
93-
for i, line in enumerate(lines):
94-
if 'def is_active():' in line:
95-
if i + 1 < len(lines) and 'return False' not in lines[i + 1]:
96-
lines.insert(i + 1, ' return False\n')
97-
break
98-
with open(drvfile, 'w') as f:
99-
f.writelines(lines)
100-
10188
dst_dir = project_root_dir / "third_party" / "enflame" / "backend"
10289
dst_dir.mkdir(parents=True, exist_ok=True)
10390

‎third_party/enflame/cmake/triton_gcu.cmake‎

Lines changed: 106 additions & 130 deletions
Original file line numberDiff line numberDiff line change
@@ -1,33 +1,12 @@
1-
# For Triton
2-
#
3-
# Staged pipeline matches kurama/cmake/triton_gcu.cmake:
4-
# same macro names and call order; implementations differ (FetchContent vs in-tree).
1+
# Flagtree builds Triton once in the top-level CMake graph. Enflame links those
2+
# targets directly, so backend source overrides under spec_cpp are applied by
3+
# FlagTreeBackendSpec.cmake to the same targets consumed here.
4+
include_guard(GLOBAL)
55

6-
include(triton_gcu_common)
7-
include(triton_gcu_llvm)
8-
9-
macro(triton_gcu_local_triton_source)
10-
set(third_party_triton_${arch}_fetch_src "${CMAKE_SOURCE_DIR}")
11-
file(GLOB_RECURSE third_party_triton_${arch}_src "${CMAKE_SOURCE_DIR}/include/*" "${CMAKE_SOURCE_DIR}/lib/*" "${CMAKE_SOURCE_DIR}/third_party/f2reduce/*" "${CMAKE_SOURCE_DIR}/third_party/proton/*")
12-
endmacro()
13-
14-
# --- Stages (keep names aligned with kurama for side-by-side review) ---
15-
16-
macro(triton_gcu_stage_init)
17-
set(CMAKE_CXX_FLAGS "${CMAKE_CXX_FLAGS} -Wall -Wextra -Werror -Wno-unused-parameter -Wno-unused-but-set-parameter -Wno-attributes")
18-
set(_triton_build_target triton_build_in_flagtree)
19-
set(third_party_triton_${arch}_fetch_bin "${CMAKE_BINARY_DIR}/triton_bin")
20-
triton_gcu_local_triton_source()
21-
22-
triton_gcu_add_llvm_triton_enflame_base_include_directories(${arch} "${MLIR_INCLUDE_DIRS}" "${LLVM_INCLUDE_DIRS}")
23-
24-
# In flagtree Triton is built in-tree: CMAKE_SOURCE_DIR IS the Triton root.
25-
# All variables that setup_triton_fetch would set via PARENT_SCOPE are defined
26-
# here instead (macros share the caller's scope, avoiding PARENT_SCOPE issues).
27-
set(TRITON_SOURCE_DIR ${CMAKE_SOURCE_DIR})
28-
set(TRITON_BINARY_DIR ${CMAKE_BINARY_DIR})
29-
set(TRITON_BUILD_TARGET "triton_build_in_flagtree")
30-
set(TRITON_OUTPUT_FILES "")
6+
function(setup_triton_in_tree)
7+
set(TRITON_SOURCE_DIR ${CMAKE_SOURCE_DIR} PARENT_SCOPE)
8+
set(TRITON_BINARY_DIR ${CMAKE_BINARY_DIR} PARENT_SCOPE)
9+
set(TRITON_ORIG_VERSION "${TRITON_VERSION}" PARENT_SCOPE)
3110
set(TRITON_CORE_LIBS
3211
TritonIR
3312
TritonGPUIR
@@ -40,111 +19,108 @@ macro(triton_gcu_stage_init)
4019
TritonTools
4120
GluonIR
4221
GluonTransforms
22+
PARENT_SCOPE
4323
)
44-
include_directories(SYSTEM ${TRITON_SOURCE_DIR}/include)
45-
include_directories(SYSTEM ${TRITON_BINARY_DIR}/include)
46-
endmacro()
47-
48-
macro(triton_gcu_stage_nested_upstream_triton_build)
49-
if(TARGET ${_triton_build_target})
50-
message(STATUS "[triton-${arch}] reuse ${_triton_build_target}")
51-
else()
52-
file(MAKE_DIRECTORY ${third_party_triton_${arch}_fetch_bin})
53-
54-
triton_gcu_append_nested_triton_cmake_args(triton_cmake_args "${MLIR_DIR}" "${LLVM_LIBRARY_DIR}")
55-
triton_gcu_apply_common_triton_upstream_patches("${third_party_triton_${arch}_fetch_src}")
24+
set(TRITON_CORE_TABLEGEN_TARGETS
25+
TritonTableGen
26+
TritonGPUTableGen
27+
TritonGPUAttrDefsIncGen
28+
TritonGPUCTAAttrIncGen
29+
TritonGPUTypeInterfacesIncGen
30+
TritonGPUOpInterfacesIncGen
31+
TritonGPUTransformsIncGen
32+
TritonConversionPassIncGen
33+
TritonTransformsIncGen
34+
TritonNvidiaGPUTableGen
35+
TritonNvidiaGPUAttrDefsIncGen
36+
TritonNvidiaGPUOpInterfacesIncGen
37+
TritonNvidiaGPUTransformsIncGen
38+
PARENT_SCOPE
39+
)
40+
endfunction()
5641

57-
execute_process(
58-
COMMAND sed -i "/auto dTy = cast<ShapedType>(\\$_op.getD().getType());/d"
59-
${third_party_triton_${arch}_fetch_src}/include/triton/Dialect/Triton/IR/TritonOpInterfaces.td
60-
ERROR_QUIET
42+
function(build_triton_python_bindings ARCH_NAME SOURCE_DIR)
43+
string(TOUPPER "${ARCH_NAME}" _ARCH_UPPER)
44+
set(_ARCH_TAG "${ARCH_NAME}")
45+
46+
if(NOT DEFINED ${_ARCH_UPPER}_PYTHON_VERSIONS)
47+
set(_CANDIDATE_VERSIONS "3.9;3.10;3.11;3.12")
48+
if(DEFINED ENV{PYTHON_VERSION} AND NOT "$ENV{PYTHON_VERSION}" STREQUAL "")
49+
list(APPEND _CANDIDATE_VERSIONS "$ENV{PYTHON_VERSION}")
50+
list(REMOVE_DUPLICATES _CANDIDATE_VERSIONS)
51+
endif()
52+
set(${_ARCH_UPPER}_PYTHON_VERSIONS "")
53+
foreach(_cv IN LISTS _CANDIDATE_VERSIONS)
54+
unset(_cv_exe)
55+
unset(_cv_exe CACHE)
56+
find_program(_cv_exe "python${_cv}" NO_CACHE)
57+
if(_cv_exe)
58+
list(APPEND ${_ARCH_UPPER}_PYTHON_VERSIONS "${_cv}")
59+
endif()
60+
endforeach()
61+
if(NOT ${_ARCH_UPPER}_PYTHON_VERSIONS)
62+
set(${_ARCH_UPPER}_PYTHON_VERSIONS "3.10")
63+
endif()
64+
endif()
65+
message(STATUS "[${_ARCH_TAG}-python] Building bindings for: ${${_ARCH_UPPER}_PYTHON_VERSIONS}")
66+
67+
set(_BINDING_TARGETS "")
68+
foreach(_pyver IN LISTS ${_ARCH_UPPER}_PYTHON_VERSIONS)
69+
unset(_PY_EXE)
70+
unset(_PY_EXE CACHE)
71+
find_program(_PY_EXE "python${_pyver}" NO_CACHE)
72+
if(NOT _PY_EXE)
73+
message(WARNING "[${_ARCH_TAG}-python] python${_pyver} not found -- skipping")
74+
continue()
75+
endif()
76+
77+
execute_process(COMMAND "${_PY_EXE}" -c
78+
"import pybind11; print(pybind11.get_include())"
79+
OUTPUT_VARIABLE _pb_inc OUTPUT_STRIP_TRAILING_WHITESPACE
80+
ERROR_QUIET RESULT_VARIABLE _pb_rc)
81+
if(NOT _pb_rc EQUAL 0)
82+
message(WARNING "[${_ARCH_TAG}-python] pybind11 not available for python${_pyver} -- installing")
83+
execute_process(COMMAND "${_PY_EXE}" -m pip install --user pybind11)
84+
execute_process(COMMAND "${_PY_EXE}" -c
85+
"import pybind11; print(pybind11.get_include())"
86+
OUTPUT_VARIABLE _pb_inc OUTPUT_STRIP_TRAILING_WHITESPACE)
87+
endif()
88+
89+
execute_process(COMMAND "${_PY_EXE}" -c
90+
"import sysconfig; print(sysconfig.get_path('include'))"
91+
OUTPUT_VARIABLE _py_inc OUTPUT_STRIP_TRAILING_WHITESPACE)
92+
execute_process(COMMAND "${_PY_EXE}" -c
93+
"import sysconfig; print(sysconfig.get_config_var('EXT_SUFFIX'))"
94+
OUTPUT_VARIABLE _py_ext_suffix OUTPUT_STRIP_TRAILING_WHITESPACE)
95+
string(REPLACE "." "" _py_tag "${_pyver}")
96+
97+
set(_tgt "_triton_${_ARCH_TAG}_py${_py_tag}")
98+
add_library(${_tgt} MODULE "${SOURCE_DIR}/triton_${_ARCH_TAG}_module.cpp")
99+
target_compile_features(${_tgt} PRIVATE cxx_std_17)
100+
target_include_directories(${_tgt} PRIVATE
101+
"${SOURCE_DIR}"
102+
"${CMAKE_CURRENT_SOURCE_DIR}/lib"
61103
)
62-
63-
add_custom_command(
64-
OUTPUT ${triton_${arch}_objs}
65-
COMMAND sed -i "s/-Wno-covered-switch-default//g" ${third_party_triton_${arch}_fetch_src}/CMakeLists.txt
66-
COMMAND find ${third_party_triton_${arch}_fetch_src} -name "CMakeLists.txt" -exec sed -i "s/-Wno-covered-switch-default//g" {} +
67-
COMMAND cmake -S ${third_party_triton_${arch}_fetch_src} -B ${third_party_triton_${arch}_fetch_bin} ${triton_cmake_args} -DTRITON_CODEGEN_BACKENDS='nvidia\;amd' -DCMAKE_CXX_FLAGS='-Wno-reorder -Wno-error=comment -Wno-unknown-warning-option' -G Ninja
68-
COMMAND cmake --build ${third_party_triton_${arch}_fetch_bin} --target all ${JOB_SETTING}
69-
DEPENDS ${third_party_triton_${arch}_src}
104+
target_include_directories(${_tgt} SYSTEM PUBLIC
105+
"${_pb_inc}" "${_py_inc}"
106+
)
107+
target_link_libraries(${_tgt} PRIVATE triton_${_ARCH_TAG}_core)
108+
109+
set_target_properties(${_tgt} PROPERTIES
110+
OUTPUT_NAME "_triton_${_ARCH_TAG}"
111+
PREFIX ""
112+
SUFFIX "${_py_ext_suffix}"
113+
LIBRARY_OUTPUT_DIRECTORY "${CMAKE_BINARY_DIR}/lib"
114+
BUILD_RPATH "$ORIGIN"
115+
INSTALL_RPATH "$ORIGIN"
116+
LINK_FLAGS "-Wl,--version-script=${SOURCE_DIR}/triton_${_ARCH_TAG}_py.map"
70117
)
71118

72-
add_custom_target(${_triton_build_target} ALL DEPENDS ${triton_${arch}_objs})
73-
message(STATUS "[triton-${arch}] created ${_triton_build_target}")
74-
endif()
75-
76-
add_custom_target(third_party_triton_${arch}_fetch_build ALL)
77-
add_dependencies(third_party_triton_${arch}_fetch_build ${_triton_build_target})
78-
endmacro()
79-
80-
macro(triton_gcu_stage_enflame_subdirectory_and_extra_deps)
81-
triton_gcu_add_triton_enflame_subdirectory_bundle(${arch} "${third_party_triton_${arch}_fetch_src}" "${third_party_triton_${arch}_fetch_bin}")
82-
endmacro()
83-
84-
macro(triton_gcu_stage_triton_opt_link_and_compile)
85-
triton_gcu_add_triton_opt_toolchain(${arch})
86-
87-
target_link_options(triton-${arch}-opt PRIVATE -Wl,--allow-multiple-definition)
88-
89-
target_link_libraries(triton-${arch}-opt PRIVATE
90-
TleIR
91-
TleToLLVM
92-
TritonTLETransforms
93-
TritonGPUTransforms
94-
)
95-
endmacro()
96-
97-
macro(triton_gcu_stage_unittests)
98-
# Kurama registers gtests here; flagtree has no triton_gcu unit tests in-tree.
99-
endmacro()
100-
119+
list(APPEND _BINDING_TARGETS "${_tgt}")
120+
message(STATUS "[${_ARCH_TAG}-python] ${_pyver} -> ${_tgt} (${_PY_EXE})")
121+
endforeach()
101122

102-
# -----------------------------------------------------------------------------
103-
# Function: setup_triton_fetch
104-
# -----------------------------------------------------------------------------
105-
# Sets up Triton source fetch and build configuration (organized by commit).
106-
# If the same commit was already fetched and built, this function skips the work.
107-
#
108-
# Parameters:
109-
# TRITON_COMMIT : Triton commit hash/tag
110-
# MLIR_DIR : MLIR cmake directory
111-
# LLVM_LIBRARY_DIR : LLVM library directory
112-
# OBJECT_FILES : List of Triton object files to build
113-
# GIT_URL_DIR : Base git URL directory for cloning triton.git
114-
# MLIR_TABLEGEN_EXE : Path to mlir-tblgen executable
115-
#
116-
# Exports to PARENT_SCOPE:
117-
# TRITON_BUILD_TARGET : Name of the build target
118-
# TRITON_SOURCE_DIR : Triton source directory
119-
# TRITON_BINARY_DIR : Triton build directory
120-
# TRITON_OUTPUT_FILES : Triton object files (absolute paths)
121-
# TRITON_OBJECT_LIB : CMake OBJECT library target name (nested Triton objs; triton_upstream_objs_<arch>)
122-
# TRITON_ORIG_VERSION : Extracted Triton version
123-
#
124-
function(setup_triton_fetch)
125-
cmake_parse_arguments(ARG "" "TRITON_COMMIT;MLIR_DIR;LLVM_LIBRARY_DIR;GIT_URL_DIR;MLIR_TABLEGEN_EXE" "OBJECT_FILES" ${ARGN})
126-
set(TRITON_SOURCE_DIR ${CMAKE_SOURCE_DIR} PARENT_SCOPE)
127-
set(TRITON_BINARY_DIR ${CMAKE_BINARY_DIR} PARENT_SCOPE)
128-
set(TRITON_BUILD_TARGET "triton_build_in_flagtree" PARENT_SCOPE)
129-
# In flagtree, Triton is built in-tree as OBJECT libraries rather than
130-
# via nested build .o files. TRITON_OUTPUT_FILES is intentionally empty;
131-
# TRITON_CORE_LIBS carries the in-tree OBJECT library targets that the
132-
# core shared library must link against.
133-
set(TRITON_OUTPUT_FILES "" PARENT_SCOPE)
134-
set(TRITON_CORE_LIBS
135-
TritonIR
136-
TritonGPUIR
137-
TritonGPUTransforms
138-
TritonTransforms
139-
TritonToTritonGPU
140-
TritonAnalysis
141-
TritonGPUToLLVM
142-
TritonLLVMIR
143-
TritonTools
144-
GluonIR
145-
GluonTransforms
146-
PARENT_SCOPE
147-
)
123+
if(_BINDING_TARGETS)
124+
add_custom_target(_triton_${_ARCH_TAG} ALL DEPENDS ${_BINDING_TARGETS})
125+
endif()
148126
endfunction()
149-
150-
triton_gcu_pipeline(${arch} 0 "${project_git_url_dir}" "${MLIR_DIR}" "${LLVM_LIBRARY_DIR}" "${MLIR_INCLUDE_DIRS}" "${LLVM_INCLUDE_DIRS}")

0 commit comments

Comments
 (0)