From 06862bb306c9e174c10c8db1e2e7496746e4a794 Mon Sep 17 00:00:00 2001 From: Fujun Han Date: Fri, 3 Apr 2026 13:59:03 +0800 Subject: [PATCH 1/4] [Multi-Vendor Support]: add FLYDSL_TARGET_STACK and move FlyJitRuntime to lib/Runtime - Introduce cmake/FlyDSLTargetStack.cmake (cache + rocdl descriptor fields). - Build FlyJitRuntime from lib/Runtime; Python CMake sets output dir under _mlir_libs. Made-with: Cursor Signed-off-by: Fujun Han --- CMakeLists.txt | 3 +++ cmake/FlyDSLTargetStack.cmake | 23 ++++++++++++++++++ lib/CMakeLists.txt | 1 + lib/Runtime/CMakeLists.txt | 15 ++++++++++++ python/mlir_flydsl/CMakeLists.txt | 39 +++++++------------------------ 5 files changed, 50 insertions(+), 31 deletions(-) create mode 100644 cmake/FlyDSLTargetStack.cmake create mode 100644 lib/Runtime/CMakeLists.txt diff --git a/CMakeLists.txt b/CMakeLists.txt index 8b374a1a..86ecf895 100644 --- a/CMakeLists.txt +++ b/CMakeLists.txt @@ -8,6 +8,9 @@ set(CMAKE_EXPORT_COMPILE_COMMANDS ON) option(FLY_UNITTEST "Build Fly Unit Tests" OFF) +include("${CMAKE_CURRENT_SOURCE_DIR}/cmake/FlyDSLTargetStack.cmake") +message(STATUS "FLYDSL_TARGET_STACK: ${FLYDSL_TARGET_STACK}") + find_package(MLIR REQUIRED CONFIG) message(STATUS "Found MLIR: ${MLIR_DIR}") diff --git a/cmake/FlyDSLTargetStack.cmake b/cmake/FlyDSLTargetStack.cmake new file mode 100644 index 00000000..093de6ab --- /dev/null +++ b/cmake/FlyDSLTargetStack.cmake @@ -0,0 +1,23 @@ +# SPDX-License-Identifier: Apache-2.0 +# Copyright (c) 2025 FlyDSL Project Contributors +# +# Active Fly target lowering stack (one per build). Sets FLYDSL_TARGET_STACK and +# stack-specific variables for other CMakeLists. + +set(FLYDSL_TARGET_STACK "rocdl" + CACHE STRING "Active Fly target lowering stack for this build (one stack per artifact).") +set_property(CACHE FLYDSL_TARGET_STACK PROPERTY STRINGS rocdl) + +# Extend the list and the if-branch when supporting additional target stacks. +set(_FLYDSL_TARGET_STACK_ALLOWED rocdl) +if(NOT FLYDSL_TARGET_STACK IN_LIST _FLYDSL_TARGET_STACK_ALLOWED) + message(FATAL_ERROR + "FLYDSL_TARGET_STACK='${FLYDSL_TARGET_STACK}' is not supported. " + "Allowed values: ${_FLYDSL_TARGET_STACK_ALLOWED}") +endif() + +if(FLYDSL_TARGET_STACK STREQUAL "rocdl") + set(FLYDSL_STACK_MLIR_TARGET_DIALECT_NAME "fly_rocdl") + set(FLYDSL_STACK_CPP_TARGET_TOKEN "ROCDL") + set(FLYDSL_STACK_UPSTREAM_MLIR_PY_DIALECT "rocdl") +endif() diff --git a/lib/CMakeLists.txt b/lib/CMakeLists.txt index 8426d229..8dad5a27 100644 --- a/lib/CMakeLists.txt +++ b/lib/CMakeLists.txt @@ -1,3 +1,4 @@ add_subdirectory(Conversion) add_subdirectory(Dialect) add_subdirectory(CAPI) +add_subdirectory(Runtime) diff --git a/lib/Runtime/CMakeLists.txt b/lib/Runtime/CMakeLists.txt new file mode 100644 index 00000000..e6326668 --- /dev/null +++ b/lib/Runtime/CMakeLists.txt @@ -0,0 +1,15 @@ +# SPDX-License-Identifier: Apache-2.0 +# Copyright (c) 2025 FlyDSL Project Contributors + +file(GLOB _flydsl_rocm_paths LIST_DIRECTORIES true "/opt/rocm*") +list(SORT _flydsl_rocm_paths ORDER DESCENDING) +find_package(hip REQUIRED CONFIG PATHS ${_flydsl_rocm_paths}) + +add_library(FlyJitRuntime SHARED FlyRocmRuntimeWrappers.cpp) +target_include_directories(FlyJitRuntime PRIVATE + ${LLVM_INCLUDE_DIRS} + ${MLIR_INCLUDE_DIRS} +) +target_compile_features(FlyJitRuntime PRIVATE cxx_std_17) +target_link_libraries(FlyJitRuntime PRIVATE hip::host hip::amdhip64) +set_target_properties(FlyJitRuntime PROPERTIES OUTPUT_NAME "fly_jit_runtime") diff --git a/python/mlir_flydsl/CMakeLists.txt b/python/mlir_flydsl/CMakeLists.txt index 07a51ac0..5f67fae3 100644 --- a/python/mlir_flydsl/CMakeLists.txt +++ b/python/mlir_flydsl/CMakeLists.txt @@ -56,10 +56,10 @@ declare_mlir_python_extension(FlyPythonSources.Core.fly_rocdl ) # NOTE: Do NOT link MLIRFlyToROCDL or other C++ libs via PRIVATE_LINK_LIBS. -# Doing so statically embeds a copy of MLIRPass (and its pass registry) into +# Doing so statically embeds a copy of MLIRPass (and the pass registry) into # _mlirRegisterEverything.so. Passes registered via registerFlyPasses() then # go into that LOCAL registry, while PassManager.parse() queries the GLOBAL -# registry in FlyPythonCAPI.so → "does not refer to a registered pass". +# registry in FlyPythonCAPI.so -> "does not refer to a registered pass". # # Instead, pass registration is done through CAPI functions # (mlirRegisterFlyPasses, mlirRegisterFlyToROCDLConversionPass) defined in @@ -121,8 +121,8 @@ add_mlir_python_common_capi_library(FlyPythonCAPI # (libamd_comgr.so). # # Strategy: -# 1. -fvisibility=hidden — hide symbols compiled by FlyDSL itself -# 2. Version script — controls the FINAL exported symbol set: +# 1. -fvisibility=hidden - hide symbols compiled by FlyDSL itself +# 2. Version script - controls the FINAL exported symbol set: # exports mlir C API + mlir:: C++ symbols # (needed by _mlirDialectsFly.so, _mlir.so, etc.), # hides everything else (llvm::, LLVM C API). @@ -130,7 +130,7 @@ add_mlir_python_common_capi_library(FlyPythonCAPI # NOTE: Triton uses --exclude-libs,ALL because it ships a single # libtriton.so with no downstream consumers. FlyDSL cannot use # --exclude-libs,ALL because GNU ld's --exclude-libs forces symbols -# to LOCAL *before* the version script is evaluated — wildcard patterns +# to LOCAL *before* the version script is evaluated - wildcard patterns # in the version script cannot re-export them. The version script # alone is sufficient: its `local: *;` catch-all hides LLVM internals. # --------------------------------------------------------------------------- @@ -168,6 +168,9 @@ add_mlir_python_modules(FlyPythonModules set(_FLYDSL_PYTHON_PACKAGES_DIR "${MLIR_BINARY_DIR}/python_packages") set(_MLIR_LIBS_DIR "${FlyPythonModules_ROOT_PREFIX}/_mlir_libs") +set_target_properties(FlyJitRuntime PROPERTIES + LIBRARY_OUTPUT_DIRECTORY "${_MLIR_LIBS_DIR}") +add_dependencies(FlyPythonCAPI FlyJitRuntime) set(_STUB_MARKER_FILE "${_MLIR_LIBS_DIR}/.stubs_generated") add_custom_command( @@ -223,29 +226,3 @@ add_custom_target(CopyFlyPythonSources ALL ) add_dependencies(CopyFlyPythonSources FlyPythonModules) - - -################################################################################ -# FlyJitRuntime — thin ROCm runtime with GPU module caching -################################################################################ - -file(GLOB _ROCM_SEARCH_PATHS LIST_DIRECTORIES true "/opt/rocm*") -list(SORT _ROCM_SEARCH_PATHS ORDER DESCENDING) -find_package(hip REQUIRED CONFIG PATHS ${_ROCM_SEARCH_PATHS}) - -add_library(FlyJitRuntime SHARED - ${PROJECT_SOURCE_DIR}/lib/Runtime/FlyRocmRuntimeWrappers.cpp -) -target_include_directories(FlyJitRuntime PRIVATE - ${LLVM_INCLUDE_DIRS} - ${MLIR_INCLUDE_DIRS} -) -target_compile_features(FlyJitRuntime PRIVATE cxx_std_17) -target_link_libraries(FlyJitRuntime PRIVATE hip::host hip::amdhip64) -set_target_properties(FlyJitRuntime PROPERTIES - OUTPUT_NAME "fly_jit_runtime" - LIBRARY_OUTPUT_DIRECTORY "${_MLIR_LIBS_DIR}" -) -add_dependencies(FlyPythonCAPI FlyJitRuntime) - - From 5027a4bb0e4752e321145d28349c0bc54a4f0e0e Mon Sep 17 00:00:00 2001 From: Fujun Han Date: Fri, 3 Apr 2026 14:17:33 +0800 Subject: [PATCH 2/4] [Multi-Vendor Support]: generate FlyRegisterEverything.cpp from stack descriptor - Add FlyRegisterEverything.cpp.in and configure_file into the build tree. - Drive EMBED_CAPI_LINK_LIBS from FLYDSL_REGISTER_EMBED_CAPI_LINK_LIBS in cmake/FlyDSLTargetStack.cmake (rocdl row). Made-with: Cursor Signed-off-by: Fujun Han --- cmake/FlyDSLTargetStack.cmake | 21 +++++++++++++++++++ python/mlir_flydsl/CMakeLists.txt | 20 +++++++----------- ...thing.cpp => FlyRegisterEverything.cpp.in} | 8 +++---- 3 files changed, 31 insertions(+), 18 deletions(-) rename python/mlir_flydsl/{FlyRegisterEverything.cpp => FlyRegisterEverything.cpp.in} (75%) diff --git a/cmake/FlyDSLTargetStack.cmake b/cmake/FlyDSLTargetStack.cmake index 093de6ab..e5226243 100644 --- a/cmake/FlyDSLTargetStack.cmake +++ b/cmake/FlyDSLTargetStack.cmake @@ -20,4 +20,25 @@ if(FLYDSL_TARGET_STACK STREQUAL "rocdl") set(FLYDSL_STACK_MLIR_TARGET_DIALECT_NAME "fly_rocdl") set(FLYDSL_STACK_CPP_TARGET_TOKEN "ROCDL") set(FLYDSL_STACK_UPSTREAM_MLIR_PY_DIALECT "rocdl") + + set(FLYDSL_REGISTER_EXTRA_INCLUDES "#include \"flydsl-c/FlyROCDLDialect.h\"") + set(FLYDSL_REGISTER_TARGET_DIALECT_HANDLES [[ MlirDialectHandle flyROCDLHandle = mlirGetDialectHandle__fly_rocdl__(); + mlirDialectHandleInsertDialect(flyROCDLHandle, registry);]]) + set(FLYDSL_REGISTER_TARGET_PASSES [[ mlirRegisterFlyToROCDLConversionPass(); + mlirRegisterFlyROCDLClusterAttrPass();]]) + + set(FLYDSL_REGISTER_EMBED_CAPI_LINK_LIBS + MLIRCAPIIR + MLIRCPIFly + MLIRCPIFlyROCDL + MLIRCPILLVMOption + MLIRCAPIArith + MLIRCAPIGPU + MLIRCAPILLVM + MLIRCAPIMath + MLIRCAPIVector + MLIRCAPIConversion + MLIRCAPITransforms + MLIRCAPIRegisterEverything + ) endif() diff --git a/python/mlir_flydsl/CMakeLists.txt b/python/mlir_flydsl/CMakeLists.txt index 5f67fae3..3cfc5986 100644 --- a/python/mlir_flydsl/CMakeLists.txt +++ b/python/mlir_flydsl/CMakeLists.txt @@ -1,5 +1,10 @@ include(AddMLIRPython) +configure_file( + "${CMAKE_CURRENT_SOURCE_DIR}/FlyRegisterEverything.cpp.in" + "${CMAKE_CURRENT_BINARY_DIR}/FlyRegisterEverything.cpp" + @ONLY) + add_compile_definitions("MLIR_PYTHON_PACKAGE_PREFIX=flydsl.${MLIR_PYTHON_PACKAGE_PREFIX}.") declare_mlir_python_sources(FlyPythonSources) @@ -71,22 +76,11 @@ declare_mlir_python_extension(FlyPythonSources.RegisterEverything ADD_TO_PARENT FlyPythonSources PYTHON_BINDINGS_LIBRARY nanobind SOURCES - FlyRegisterEverything.cpp + "${CMAKE_CURRENT_BINARY_DIR}/FlyRegisterEverything.cpp" PRIVATE_LINK_LIBS LLVMSupport EMBED_CAPI_LINK_LIBS - MLIRCAPIIR - MLIRCPIFly - MLIRCPIFlyROCDL - MLIRCPILLVMOption - MLIRCAPIArith - MLIRCAPIGPU - MLIRCAPILLVM - MLIRCAPIMath - MLIRCAPIVector - MLIRCAPIConversion - MLIRCAPITransforms - MLIRCAPIRegisterEverything + ${FLYDSL_REGISTER_EMBED_CAPI_LINK_LIBS} ) diff --git a/python/mlir_flydsl/FlyRegisterEverything.cpp b/python/mlir_flydsl/FlyRegisterEverything.cpp.in similarity index 75% rename from python/mlir_flydsl/FlyRegisterEverything.cpp rename to python/mlir_flydsl/FlyRegisterEverything.cpp.in index b25b6b4f..7c4f04a2 100644 --- a/python/mlir_flydsl/FlyRegisterEverything.cpp +++ b/python/mlir_flydsl/FlyRegisterEverything.cpp.in @@ -6,7 +6,7 @@ #include "mlir/Bindings/Python/NanobindAdaptors.h" #include "flydsl-c/FlyDialect.h" -#include "flydsl-c/FlyROCDLDialect.h" +@FLYDSL_REGISTER_EXTRA_INCLUDES@ NB_MODULE(_mlirRegisterEverything, m) { m.doc() = "MLIR All Upstream Dialects, Translations and Passes Registration"; @@ -16,14 +16,12 @@ NB_MODULE(_mlirRegisterEverything, m) { MlirDialectHandle flyHandle = mlirGetDialectHandle__fly__(); mlirDialectHandleInsertDialect(flyHandle, registry); - MlirDialectHandle flyROCDLHandle = mlirGetDialectHandle__fly_rocdl__(); - mlirDialectHandleInsertDialect(flyROCDLHandle, registry); +@FLYDSL_REGISTER_TARGET_DIALECT_HANDLES@ }); m.def("register_llvm_translations", [](MlirContext context) { mlirRegisterAllLLVMTranslations(context); }); mlirRegisterAllPasses(); mlirRegisterFlyPasses(); - mlirRegisterFlyToROCDLConversionPass(); - mlirRegisterFlyROCDLClusterAttrPass(); +@FLYDSL_REGISTER_TARGET_PASSES@ } From 1c96b0824cbf14d0b94f84615d3206b04b6982f4 Mon Sep 17 00:00:00 2001 From: Fujun Han Date: Fri, 3 Apr 2026 14:19:17 +0800 Subject: [PATCH 3/4] [Multi-Vendor Support]: gate ROCDL stack behind FLYDSL_BUILD_ROCDL_STACK - Conditionally add FlyROCDL / FlyToROCDL / Runtime subtrees in include/ and lib/. - fly-opt: conditional links and FLYDSL_HAS_ROCDL_TARGET_STACK; guard Passes.h. - MLIRCPIFlyROCDL: define FLYDSL_HAS_ROCDL_TARGET_STACK for Conversion/Passes.h. - python/mlir_flydsl: gate FlyROCDL bindings, stubgen, tablegen copies, FlyJitRuntime hookup. Made-with: Cursor Signed-off-by: Fujun Han --- cmake/FlyDSLTargetStack.cmake | 2 + include/flydsl/Conversion/CMakeLists.txt | 4 +- include/flydsl/Conversion/Passes.h | 4 ++ include/flydsl/Dialect/CMakeLists.txt | 4 +- lib/CAPI/Dialect/CMakeLists.txt | 4 +- lib/CAPI/Dialect/FlyROCDL/CMakeLists.txt | 2 + lib/CMakeLists.txt | 4 +- lib/Conversion/CMakeLists.txt | 4 +- lib/Dialect/CMakeLists.txt | 4 +- python/mlir_flydsl/CMakeLists.txt | 61 +++++++++++++++++------- tools/fly-opt/CMakeLists.txt | 13 ++++- tools/fly-opt/fly-opt.cpp | 6 +++ 12 files changed, 88 insertions(+), 24 deletions(-) diff --git a/cmake/FlyDSLTargetStack.cmake b/cmake/FlyDSLTargetStack.cmake index e5226243..5f28070f 100644 --- a/cmake/FlyDSLTargetStack.cmake +++ b/cmake/FlyDSLTargetStack.cmake @@ -16,7 +16,9 @@ if(NOT FLYDSL_TARGET_STACK IN_LIST _FLYDSL_TARGET_STACK_ALLOWED) "Allowed values: ${_FLYDSL_TARGET_STACK_ALLOWED}") endif() +set(FLYDSL_BUILD_ROCDL_STACK OFF) if(FLYDSL_TARGET_STACK STREQUAL "rocdl") + set(FLYDSL_BUILD_ROCDL_STACK ON) set(FLYDSL_STACK_MLIR_TARGET_DIALECT_NAME "fly_rocdl") set(FLYDSL_STACK_CPP_TARGET_TOKEN "ROCDL") set(FLYDSL_STACK_UPSTREAM_MLIR_PY_DIALECT "rocdl") diff --git a/include/flydsl/Conversion/CMakeLists.txt b/include/flydsl/Conversion/CMakeLists.txt index be704b62..08dd9fca 100644 --- a/include/flydsl/Conversion/CMakeLists.txt +++ b/include/flydsl/Conversion/CMakeLists.txt @@ -1 +1,3 @@ -add_subdirectory(FlyToROCDL) +if(FLYDSL_BUILD_ROCDL_STACK) + add_subdirectory(FlyToROCDL) +endif() diff --git a/include/flydsl/Conversion/Passes.h b/include/flydsl/Conversion/Passes.h index 1475c25f..d9bb462d 100644 --- a/include/flydsl/Conversion/Passes.h +++ b/include/flydsl/Conversion/Passes.h @@ -5,6 +5,8 @@ #ifndef FLYDSL_CONVERSION_PASSES_H #define FLYDSL_CONVERSION_PASSES_H +#ifdef FLYDSL_HAS_ROCDL_TARGET_STACK + #include "flydsl/Conversion/FlyToROCDL/FlyToROCDL.h" namespace mlir { @@ -14,4 +16,6 @@ namespace mlir { } // namespace mlir +#endif // FLYDSL_HAS_ROCDL_TARGET_STACK + #endif // FLYDSL_CONVERSION_PASSES_H diff --git a/include/flydsl/Dialect/CMakeLists.txt b/include/flydsl/Dialect/CMakeLists.txt index 0b152044..b1ff28b6 100644 --- a/include/flydsl/Dialect/CMakeLists.txt +++ b/include/flydsl/Dialect/CMakeLists.txt @@ -1,2 +1,4 @@ add_subdirectory(Fly) -add_subdirectory(FlyROCDL) +if(FLYDSL_BUILD_ROCDL_STACK) + add_subdirectory(FlyROCDL) +endif() diff --git a/lib/CAPI/Dialect/CMakeLists.txt b/lib/CAPI/Dialect/CMakeLists.txt index 0b152044..b1ff28b6 100644 --- a/lib/CAPI/Dialect/CMakeLists.txt +++ b/lib/CAPI/Dialect/CMakeLists.txt @@ -1,2 +1,4 @@ add_subdirectory(Fly) -add_subdirectory(FlyROCDL) +if(FLYDSL_BUILD_ROCDL_STACK) + add_subdirectory(FlyROCDL) +endif() diff --git a/lib/CAPI/Dialect/FlyROCDL/CMakeLists.txt b/lib/CAPI/Dialect/FlyROCDL/CMakeLists.txt index c30d361d..2ba6ac2d 100644 --- a/lib/CAPI/Dialect/FlyROCDL/CMakeLists.txt +++ b/lib/CAPI/Dialect/FlyROCDL/CMakeLists.txt @@ -4,3 +4,5 @@ add_mlir_public_c_api_library(MLIRCPIFlyROCDL MLIRFlyROCDLDialect MLIRFlyToROCDL ) + +target_compile_definitions(MLIRCPIFlyROCDL PRIVATE FLYDSL_HAS_ROCDL_TARGET_STACK) diff --git a/lib/CMakeLists.txt b/lib/CMakeLists.txt index 8dad5a27..d7a2c610 100644 --- a/lib/CMakeLists.txt +++ b/lib/CMakeLists.txt @@ -1,4 +1,6 @@ add_subdirectory(Conversion) add_subdirectory(Dialect) add_subdirectory(CAPI) -add_subdirectory(Runtime) +if(FLYDSL_BUILD_ROCDL_STACK) + add_subdirectory(Runtime) +endif() diff --git a/lib/Conversion/CMakeLists.txt b/lib/Conversion/CMakeLists.txt index be704b62..08dd9fca 100644 --- a/lib/Conversion/CMakeLists.txt +++ b/lib/Conversion/CMakeLists.txt @@ -1 +1,3 @@ -add_subdirectory(FlyToROCDL) +if(FLYDSL_BUILD_ROCDL_STACK) + add_subdirectory(FlyToROCDL) +endif() diff --git a/lib/Dialect/CMakeLists.txt b/lib/Dialect/CMakeLists.txt index 0b152044..b1ff28b6 100644 --- a/lib/Dialect/CMakeLists.txt +++ b/lib/Dialect/CMakeLists.txt @@ -1,2 +1,4 @@ add_subdirectory(Fly) -add_subdirectory(FlyROCDL) +if(FLYDSL_BUILD_ROCDL_STACK) + add_subdirectory(FlyROCDL) +endif() diff --git a/python/mlir_flydsl/CMakeLists.txt b/python/mlir_flydsl/CMakeLists.txt index 3cfc5986..51daeb1f 100644 --- a/python/mlir_flydsl/CMakeLists.txt +++ b/python/mlir_flydsl/CMakeLists.txt @@ -20,6 +20,7 @@ declare_mlir_dialect_python_bindings( GEN_ENUM_BINDINGS ) +if(FLYDSL_BUILD_ROCDL_STACK) declare_mlir_dialect_python_bindings( ADD_TO_PARENT FlyPythonSources ROOT_DIR "${CMAKE_CURRENT_SOURCE_DIR}/" @@ -29,6 +30,7 @@ declare_mlir_dialect_python_bindings( DIALECT_NAME fly_rocdl GEN_ENUM_BINDINGS ) +endif() # NOTE: Do NOT link MLIRFlyDialect/MLIRFlyROCDLDialect here via PRIVATE_LINK_LIBS. # These symbols are already provided by FlyPythonCAPI.so (via MLIRCPIFly's transitive @@ -49,6 +51,7 @@ declare_mlir_python_extension(FlyPythonSources.Core.fly LLVMSupport ) +if(FLYDSL_BUILD_ROCDL_STACK) declare_mlir_python_extension(FlyPythonSources.Core.fly_rocdl MODULE_NAME _mlirDialectsFlyROCDL ADD_TO_PARENT FlyPythonSources @@ -59,6 +62,7 @@ declare_mlir_python_extension(FlyPythonSources.Core.fly_rocdl PRIVATE_LINK_LIBS LLVMSupport ) +endif() # NOTE: Do NOT link MLIRFlyToROCDL or other C++ libs via PRIVATE_LINK_LIBS. # Doing so statically embeds a copy of MLIRPass (and the pass registry) into @@ -95,7 +99,11 @@ set(MLIRFlyDSLSources MLIRPythonSources.Dialects.func MLIRPythonSources.Dialects.cf MLIRPythonSources.Dialects.scf - MLIRPythonSources.Dialects.rocdl +) +if(FLYDSL_BUILD_ROCDL_STACK) + list(APPEND MLIRFlyDSLSources MLIRPythonSources.Dialects.rocdl) +endif() +list(APPEND MLIRFlyDSLSources MLIRPythonSources.Dialects.vector MLIRPythonSources.Dialects.llvm MLIRPythonSources.ExecutionEngine @@ -162,22 +170,35 @@ add_mlir_python_modules(FlyPythonModules set(_FLYDSL_PYTHON_PACKAGES_DIR "${MLIR_BINARY_DIR}/python_packages") set(_MLIR_LIBS_DIR "${FlyPythonModules_ROOT_PREFIX}/_mlir_libs") -set_target_properties(FlyJitRuntime PROPERTIES - LIBRARY_OUTPUT_DIRECTORY "${_MLIR_LIBS_DIR}") -add_dependencies(FlyPythonCAPI FlyJitRuntime) +if(TARGET FlyJitRuntime) + set_target_properties(FlyJitRuntime PROPERTIES + LIBRARY_OUTPUT_DIRECTORY "${_MLIR_LIBS_DIR}") + add_dependencies(FlyPythonCAPI FlyJitRuntime) +endif() set(_STUB_MARKER_FILE "${_MLIR_LIBS_DIR}/.stubs_generated") +set(_flydsl_stubgen_modules + flydsl._mlir._mlir_libs._mlir + flydsl._mlir._mlir_libs._mlirDialectsFly +) +if(FLYDSL_BUILD_ROCDL_STACK) + list(APPEND _flydsl_stubgen_modules flydsl._mlir._mlir_libs._mlirDialectsFlyROCDL) +endif() +list(APPEND _flydsl_stubgen_modules + flydsl._mlir._mlir_libs._mlirDialectsGPU + flydsl._mlir._mlir_libs._mlirDialectsLLVM +) +set(_flydsl_stubgen_m_flags "") +foreach(_m IN LISTS _flydsl_stubgen_modules) + string(APPEND _flydsl_stubgen_m_flags " -m ${_m}") +endforeach() + add_custom_command( OUTPUT "${_STUB_MARKER_FILE}" COMMAND /bin/bash -c "\ PYTHONPATH='${_FLYDSL_PYTHON_PACKAGES_DIR}' \ '${Python3_EXECUTABLE}' -m nanobind.stubgen \ - -q -r \ - -m flydsl._mlir._mlir_libs._mlir \ - -m flydsl._mlir._mlir_libs._mlirDialectsFly \ - -m flydsl._mlir._mlir_libs._mlirDialectsFlyROCDL \ - -m flydsl._mlir._mlir_libs._mlirDialectsGPU \ - -m flydsl._mlir._mlir_libs._mlirDialectsLLVM \ + -q -r${_flydsl_stubgen_m_flags} \ -O '${_MLIR_LIBS_DIR}' \ || echo 'Warning: nanobind.stubgen not available -- skipping stub generation'" COMMAND ${CMAKE_COMMAND} -E touch "${_STUB_MARKER_FILE}" @@ -190,6 +211,19 @@ add_custom_target(FlyPythonStubs ALL DEPENDS "${_STUB_MARKER_FILE}" ) +if(FLYDSL_BUILD_ROCDL_STACK) + set(_FLY_COPY_ROCDL_TABLEGEN + COMMAND ${CMAKE_COMMAND} -E copy_if_different + "${MLIR_BINARY_DIR}/python/mlir_flydsl/dialects/_fly_rocdl_ops_gen.py" + "${MLIR_BINARY_DIR}/python_packages/flydsl/_mlir/dialects/_fly_rocdl_ops_gen.py" + COMMAND ${CMAKE_COMMAND} -E copy_if_different + "${MLIR_BINARY_DIR}/python/mlir_flydsl/dialects/_fly_rocdl_enum_gen.py" + "${MLIR_BINARY_DIR}/python_packages/flydsl/_mlir/dialects/_fly_rocdl_enum_gen.py" + ) +else() + set(_FLY_COPY_ROCDL_TABLEGEN "") +endif() + add_custom_target(CopyFlyPythonSources ALL COMMAND ${CMAKE_COMMAND} -E copy_directory "${PROJECT_SOURCE_DIR}/python/flydsl" @@ -201,12 +235,7 @@ add_custom_target(CopyFlyPythonSources ALL COMMAND ${CMAKE_COMMAND} -E copy_if_different "${MLIR_BINARY_DIR}/python/mlir_flydsl/dialects/_fly_enum_gen.py" "${MLIR_BINARY_DIR}/python_packages/flydsl/_mlir/dialects/_fly_enum_gen.py" - COMMAND ${CMAKE_COMMAND} -E copy_if_different - "${MLIR_BINARY_DIR}/python/mlir_flydsl/dialects/_fly_rocdl_ops_gen.py" - "${MLIR_BINARY_DIR}/python_packages/flydsl/_mlir/dialects/_fly_rocdl_ops_gen.py" - COMMAND ${CMAKE_COMMAND} -E copy_if_different - "${MLIR_BINARY_DIR}/python/mlir_flydsl/dialects/_fly_rocdl_enum_gen.py" - "${MLIR_BINARY_DIR}/python_packages/flydsl/_mlir/dialects/_fly_rocdl_enum_gen.py" + ${_FLY_COPY_ROCDL_TABLEGEN} COMMAND ${CMAKE_COMMAND} -E copy_if_different "$" "${_MLIR_LIBS_DIR}/libmlir_c_runner_utils.so" diff --git a/tools/fly-opt/CMakeLists.txt b/tools/fly-opt/CMakeLists.txt index 9e0e8558..eb78b61a 100644 --- a/tools/fly-opt/CMakeLists.txt +++ b/tools/fly-opt/CMakeLists.txt @@ -1,12 +1,17 @@ set(LIBS MLIRFlyDialect - MLIRFlyROCDLDialect - MLIRFlyToROCDL MLIRRegisterAllExtensions MLIRMlirOptMain ) +if(FLYDSL_BUILD_ROCDL_STACK) + list(APPEND LIBS + MLIRFlyROCDLDialect + MLIRFlyToROCDL + ) +endif() + add_llvm_tool(fly-opt fly-opt.cpp ) @@ -14,3 +19,7 @@ add_llvm_tool(fly-opt target_link_libraries(fly-opt PRIVATE ${LIBS}) llvm_map_components_to_libnames(llvm_libs Support) target_link_libraries(fly-opt PRIVATE ${llvm_libs}) + +if(FLYDSL_BUILD_ROCDL_STACK) + target_compile_definitions(fly-opt PRIVATE FLYDSL_HAS_ROCDL_TARGET_STACK) +endif() diff --git a/tools/fly-opt/fly-opt.cpp b/tools/fly-opt/fly-opt.cpp index 2feefec9..be85e17c 100644 --- a/tools/fly-opt/fly-opt.cpp +++ b/tools/fly-opt/fly-opt.cpp @@ -14,19 +14,25 @@ #include "flydsl/Conversion/Passes.h" #include "flydsl/Dialect/Fly/IR/FlyDialect.h" #include "flydsl/Dialect/Fly/Transforms/Passes.h" +#ifdef FLYDSL_HAS_ROCDL_TARGET_STACK #include "flydsl/Dialect/FlyROCDL/IR/Dialect.h" +#endif int main(int argc, char **argv) { mlir::registerAllPasses(); mlir::fly::registerFlyPasses(); +#ifdef FLYDSL_HAS_ROCDL_TARGET_STACK mlir::registerFlyToROCDLConversionPass(); mlir::registerFlyROCDLClusterAttrPass(); +#endif mlir::DialectRegistry registry; mlir::registerAllDialects(registry); mlir::registerAllExtensions(registry); registry.insert(); +#ifdef FLYDSL_HAS_ROCDL_TARGET_STACK registry.insert(); +#endif return mlir::asMainReturnCode( mlir::MlirOptMain(argc, argv, "Fly Optimizer Driver\n", registry)); From d0152f93145fc2d5580a5cf4206e6ff10e0bdcdb Mon Sep 17 00:00:00 2001 From: Fujun Han Date: Fri, 3 Apr 2026 17:40:30 +0800 Subject: [PATCH 4/4] [Multi-Vendor Supprot]: generated _build_config and Python install caps - Generate flydsl/_build_config.py from cmake/FlyDSLBuildConfig.py.in; copy into python_packages during CopyFlyPythonSources. - FLYDSL_BUILD_CONFIG_* in FlyDSLTargetStack (rocdl row) drive allow-lists. - _install_limits + checks in get_backend and ensure_compile_runtime_pairing_from_env. - Document stack bring-up in ADDING_TARGET_STACK.md (pointer in FlyDSLTargetStack). Made-with: Cursor Signed-off-by: Fujun Han --- ADDING_TARGET_STACK.md | 130 ++++++++++++++++++ README.md | 5 + cmake/FlyDSLBuildConfig.py.in | 12 ++ cmake/FlyDSLTargetStack.cmake | 5 + python/flydsl/_install_limits.py | 43 ++++++ python/flydsl/compiler/backends/__init__.py | 6 + .../flydsl/runtime/device_runtime/__init__.py | 6 + python/mlir_flydsl/CMakeLists.txt | 9 ++ tests/README.md | 6 +- tests/unit/test_device_runtime.py | 8 +- 10 files changed, 225 insertions(+), 5 deletions(-) create mode 100644 ADDING_TARGET_STACK.md create mode 100644 cmake/FlyDSLBuildConfig.py.in create mode 100644 python/flydsl/_install_limits.py diff --git a/ADDING_TARGET_STACK.md b/ADDING_TARGET_STACK.md new file mode 100644 index 00000000..512c9995 --- /dev/null +++ b/ADDING_TARGET_STACK.md @@ -0,0 +1,130 @@ +# Adding a New Fly Target Stack + +This document is the maintainer guide for extending FlyDSL with an additional value of **`FLYDSL_TARGET_STACK`**. Today the repository allows only **`rocdl`** (ROCm / HIP / FlyROCDL). Use the checklists below when introducing a second stack (for example, another GPU vendor or a host-only toolchain). + +## Scope and assumptions + +- **One build, one stack.** The produced binaries and Python package must contain only the dialects, conversions, CAPI libraries, Python bindings, JIT runtime, and registration logic for the selected stack. +- **Single CMake selector.** End users choose the stack with **`FLYDSL_TARGET_STACK`**. Do not add parallel per-stack **`ENABLE_*`** cache options; encode stack-specific behavior in the descriptor row for that stack id. +- **Consistency.** The active row in **`cmake/FlyDSLTargetStack.cmake`**, the strings substituted into **`python/mlir_flydsl/FlyRegisterEverything.cpp.in`**, and the allow-lists emitted into **`flydsl._build_config`** (**`FLYDSL_BUILD_CONFIG_*`**) must describe the **same** compile backends and device runtime kinds. + +## Design constraints + +### Registry safety (`_mlirRegisterEverything`) + +Pass registration for the Python extension **`_mlirRegisterEverything`** must go through **CAPI entry points** linked via **`EMBED_CAPI_LINK_LIBS`** only. + +Do **not** link C++ pass libraries (for example **`MLIRFlyToROCDL`**) into **`PRIVATE_LINK_LIBS`** for that extension. Doing so embeds a separate **`MLIRPass`** registry instance; passes then register in the wrong registry and **`PassManager.parse()`** fails at runtime. + +### Generated sources + +Do not hand-edit build-tree outputs: + +| Generated file | Source | +|----------------|--------| +| **`FlyRegisterEverything.cpp`** | **`configure_file`** from **`python/mlir_flydsl/FlyRegisterEverything.cpp.in`**, driven by **`FLYDSL_REGISTER_*`** variables in **`cmake/FlyDSLTargetStack.cmake`**. | +| **`flydsl/_build_config.py`** | **`configure_file`** from **`cmake/FlyDSLBuildConfig.py.in`**, then copied into **`python_packages/flydsl/`** during the **`CopyFlyPythonSources`** custom target. | + +--- + +## 1. Central descriptor: `cmake/FlyDSLTargetStack.cmake` + +| Step | Action | +|------|--------| +| 1 | Append the new stack id to **`_FLYDSL_TARGET_STACK_ALLOWED`** and to **`set_property(CACHE FLYDSL_TARGET_STACK PROPERTY STRINGS ...)`** so the value appears in CMake GUI tools. | +| 2 | Replace the single **`if(FLYDSL_TARGET_STACK STREQUAL "rocdl")`** block with **`if` / `elseif` / `endif()`** (or factor shared variables) and add a branch for the new stack. | +| 3 | In that branch, set **descriptor variables** used elsewhere: **`FLYDSL_STACK_MLIR_TARGET_DIALECT_NAME`** (segment in **`mlirGetDialectHandle____()`**), **`FLYDSL_STACK_CPP_TARGET_TOKEN`** (naming pattern for **`MLIRFly${TOKEN}Dialect`**, CAPI targets, and so on), **`FLYDSL_STACK_UPSTREAM_MLIR_PY_DIALECT`** (suffix for **`MLIRPythonSources.Dialects.`**). | +| 4 | Set **`FLYDSL_REGISTER_EXTRA_INCLUDES`**, **`FLYDSL_REGISTER_TARGET_DIALECT_HANDLES`**, and **`FLYDSL_REGISTER_TARGET_PASSES`** so the generated **`FlyRegisterEverything.cpp`** registers the correct dialect handles and CAPI pass registration functions. Preserve C++ indentation inside bracket strings. | +| 5 | Set **`FLYDSL_REGISTER_EMBED_CAPI_LINK_LIBS`** to the **CAPI** targets for the Fly dialect, the stack-specific CAPI library, and the upstream MLIR CAPI set. Omit non-CAPI pass libraries. | +| 6 | Set **`FLYDSL_BUILD_CONFIG_BACKEND_IDS`** and **`FLYDSL_BUILD_CONFIG_RUNTIME_KINDS`** to fragments that expand to valid Python tuple elements inside **`cmake/FlyDSLBuildConfig.py.in`** (for example **`"\"rocm\", "`** for a single id). | +| 7 | **Stack-wide CMake flag.** If the new stack shares the same HIP-based JIT layout as ROCm, you may keep using **`FLYDSL_BUILD_ROCDL_STACK`**. If not (no HIP, different runtime), introduce **`FLYDSL_BUILD__STACK`** (or a unified flag) and update **every** **`if(FLYDSL_BUILD_ROCDL_STACK)`** in CMake and every related preprocessor macro in C++ consistently. | + +The header comment in **`cmake/FlyDSLTargetStack.cmake`** may point reviewers to this file (**`ADDING_TARGET_STACK.md`** at the repository root). + +--- + +## 2. C++ and TableGen layout: `include/flydsl` and `lib` + +### Parent `CMakeLists.txt` files to gate + +Wrap stack-specific **`add_subdirectory(...)`** calls in the **same** condition as in **`FlyDSLTargetStack.cmake`** (today **`if(FLYDSL_BUILD_ROCDL_STACK)`**). + +| File | Purpose | +|------|---------| +| **`include/flydsl/Dialect/CMakeLists.txt`** | Target dialect TableGen and headers (for example **`FlyROCDL`**). | +| **`include/flydsl/Conversion/CMakeLists.txt`** | Target lowering TableGen (for example **`FlyToROCDL`**). | +| **`lib/Dialect/CMakeLists.txt`** | Target dialect implementation libraries. | +| **`lib/Conversion/CMakeLists.txt`** | Target conversion pass libraries. | +| **`lib/CAPI/Dialect/CMakeLists.txt`** | Target stack CAPI wrapper subdirectory. | +| **`lib/CMakeLists.txt`** | **`lib/Runtime`** (JIT shared library). Omit **`add_subdirectory(Runtime)`** when the stack does not ship that target. | + +### New directories (typical layout) + +Mirror the existing ROCDL layout with stack-specific names, for example: + +- **`include/flydsl/Dialect//`** +- **`include/flydsl/Conversion/FlyTo/`** +- **`lib/Dialect//`**, **`lib/Conversion/FlyTo/`** +- **`lib/CAPI/Dialect//`** with **`add_mlir_public_c_api_library`** + +Any translation unit that includes **`flydsl/Conversion/Passes.h`** and expects generated **`Passes.h.inc`** content must receive **`target_compile_definitions`** for the same preprocessor macro used in **`Passes.h`** (today **`FLYDSL_HAS_ROCDL_TARGET_STACK`** for the ROCDL conversion block). + +--- + +## 3. `include/flydsl/Conversion/Passes.h` and `tools/fly-opt` + +| Artifact | Requirement | +|----------|-------------| +| **`include/flydsl/Conversion/Passes.h`** | Guard stack-specific includes and **`GEN_PASS_REGISTRATION`** / **`Passes.h.inc`** inclusion with a preprocessor macro so builds without that conversion still configure. | +| **`tools/fly-opt/fly-opt.cpp`** | Match the same macro for extra includes, **`register*Pass`**, and **`registry.insert`** for stack dialects. | +| **`tools/fly-opt/CMakeLists.txt`** | Link only libraries that exist for the active stack; apply **`target_compile_definitions`** for the macro used above. | +| **`lib/CAPI/Dialect//CMakeLists.txt`** | If sources include **`Passes.h`**, define the macro on that target so pass registration declarations are visible. | + +--- + +## 4. JIT runtime: `lib/Runtime/` + +| Item | Notes | +|------|--------| +| **`lib/Runtime/CMakeLists.txt`** | Defines **`FlyJitRuntime`** (or a renamed / additional target). Vendor **`find_package`**, link libraries, and sources differ per stack. | +| **Output name** | The default **`OUTPUT_NAME`** is **`fly_jit_runtime`**. If you produce multiple runtime shared objects per product line, update **`jit_runtime_lib_basenames()`** in the relevant **`python/flydsl/compiler/backends/*.py`** and any install or copy rules. | +| **`python/mlir_flydsl/CMakeLists.txt`** | After **`_MLIR_LIBS_DIR`** is set, assign **`LIBRARY_OUTPUT_DIRECTORY`** for **`FlyJitRuntime`** and **`add_dependencies(FlyPythonCAPI FlyJitRuntime)`**. If a configuration omits **`FlyJitRuntime`**, wrap those lines in **`if(TARGET FlyJitRuntime)`**. | + +--- + +## 5. Python MLIR bindings: `python/mlir_flydsl/` + +| Area | Location / variable | Notes | +|------|---------------------|--------| +| Registration C++ template | **`FlyRegisterEverything.cpp.in`** | Usually one template for all stacks; fill with **`FLYDSL_REGISTER_*`** from CMake. Empty fragments must still yield valid C++. | +| Configure steps | **`CMakeLists.txt`** (top) | **`configure_file`** for **`FlyRegisterEverything.cpp.in`**; **`configure_file`** for **`cmake/FlyDSLBuildConfig.py.in`** into a generated path under **`CMAKE_CURRENT_BINARY_DIR`**. | +| Dialect bindings | **`declare_mlir_dialect_python_bindings`** | Declare stack-specific **`.td`** and Python stubs only when the stack is enabled. | +| Extension modules | **`declare_mlir_python_extension`** | Same gating for stack-specific nanobind modules (for example **`_mlirDialectsFlyROCDL`**). | +| Embed list | **`EMBED_CAPI_LINK_LIBS`** | Use **`${FLYDSL_REGISTER_EMBED_CAPI_LINK_LIBS}`**; avoid duplicating the list. | +| Upstream dialect bundle | **`MLIRFlyDSLSources`** | **`list(APPEND ... MLIRPythonSources.Dialects.)`** when the stack is enabled; align **``** with **`FLYDSL_STACK_UPSTREAM_MLIR_PY_DIALECT`**. | +| Stub generation | **`_flydsl_stubgen_modules`** | Append **`-m`** entries only for extension modules built in that configuration. | +| TableGen Python copies | **`CopyFlyPythonSources`** | **`copy_if_different`** only for generated **`_fly_*`** files that exist for the stack; use a conditional command list (similar to **`_FLY_COPY_ROCDL_TABLEGEN`**) when a stack is disabled. | +| Install tree | **`CopyFlyPythonSources`** | After copying **`python/flydsl`**, copy **`_build_config.py`** into **`python_packages/flydsl/`** so imports work from the build/install layout. | + +--- + +## 6. Python language surface: compilers, runtimes, install caps + +| File | Change | +|------|--------| +| **`python/flydsl/compiler/backends/__init__.py`** | **`register_backend("", ...)`**. The id must appear in **`flydsl._build_config.ENABLED_COMPILE_BACKEND_IDS`** (generated from CMake). | +| **`python/flydsl/compiler/backends/.py`** (or equivalent) | Implement **`BaseBackend`**: **`pipeline_fragments`**, **`jit_runtime_lib_basenames`**, target detection, and any stack-specific options. | +| **`python/flydsl/runtime/device_runtime/__init__.py`** | Extend **`COMPILE_BACKEND_TO_RUNTIME_KIND`** and **`_builtin_runtimes`**; document extension hooks (**`register_compile_runtime_mapping`**) if third parties map custom backend names to an existing runtime kind. Runtime kinds must appear in **`ENABLED_RUNTIME_KINDS`**. | +| **`python/flydsl/_install_limits.py`** | **`_FALLBACK`** is used when **`flydsl._build_config`** is missing (for example **`PYTHONPATH`** pointed only at a bare source tree). Update it if the default development story assumes a different stack or allow-list. | + +**`get_backend`** and **`ensure_compile_runtime_pairing_from_env`** enforce install caps against **`_build_config`**. CMake-generated allow-lists must stay aligned with **`register_backend`** and device runtime registration. + +--- + +## 7. Testing and CI + +| Layer | Recommendation | +|-------|------------------| +| CMake | Add a configure job with **`FLYDSL_TARGET_STACK=`**; optionally build **`fly-opt`** and selected Python extension targets. | +| Python | Use **`pytest.importorskip`** or markers for modules that exist only on some stacks. When tests monkeypatch **`install_caps`** or **`install_limits`**, keep allow-lists consistent with the scenario under test (see **`tests/unit/test_device_runtime.py`**). | +| Integration | Run a minimal kernel or a subset of **`tests/kernels`** on representative hardware or container images for that stack. | diff --git a/README.md b/README.md index 2b3ef3d7..801590dd 100644 --- a/README.md +++ b/README.md @@ -40,9 +40,14 @@ FlyDSL/ ├── kernels/ # Python GPU kernels (importable as `kernels.*`) ├── tests/ # pytest-based tests ├── CMakeLists.txt # top-level CMake +├── ADDING_TARGET_STACK.md # maintainer guide: adding a Fly target stack (`FLYDSL_TARGET_STACK`) └── setup.py # Python packaging ``` +### Target stacks (contributors) + +Each build selects a single lowering stack via CMake cache variable **`FLYDSL_TARGET_STACK`** (today **`rocdl`** only). To add another stack or wire new dialects, passes, Python bindings, and JIT runtime pieces consistently, follow [**`ADDING_TARGET_STACK.md`**](ADDING_TARGET_STACK.md). It documents the descriptor in **`cmake/FlyDSLTargetStack.cmake`**, which CMake targets to gate, safe **`_mlirRegisterEverything`** linking, generated **`flydsl._build_config`**, and related tests. + ## Getting started ### Prerequisites diff --git a/cmake/FlyDSLBuildConfig.py.in b/cmake/FlyDSLBuildConfig.py.in new file mode 100644 index 00000000..8f03cfba --- /dev/null +++ b/cmake/FlyDSLBuildConfig.py.in @@ -0,0 +1,12 @@ +# SPDX-License-Identifier: Apache-2.0 +# Copyright (c) 2025 FlyDSL Project Contributors +# +# Generated by CMake - do not edit. Values reflect the Fly target stack for this build. + +TARGET_STACK = "@FLYDSL_TARGET_STACK@" + +# Allowed ``FLYDSL_COMPILE_BACKEND`` values for this install (lowercase ids). +ENABLED_COMPILE_BACKEND_IDS = (@FLYDSL_BUILD_CONFIG_BACKEND_IDS@) + +# Allowed ``FLYDSL_RUNTIME_KIND`` values for this install (lowercase ids). +ENABLED_RUNTIME_KINDS = (@FLYDSL_BUILD_CONFIG_RUNTIME_KINDS@) diff --git a/cmake/FlyDSLTargetStack.cmake b/cmake/FlyDSLTargetStack.cmake index 5f28070f..df9957c9 100644 --- a/cmake/FlyDSLTargetStack.cmake +++ b/cmake/FlyDSLTargetStack.cmake @@ -3,6 +3,8 @@ # # Active Fly target lowering stack (one per build). Sets FLYDSL_TARGET_STACK and # stack-specific variables for other CMakeLists. +# +# How to add a stack: ADDING_TARGET_STACK.md (repository root) set(FLYDSL_TARGET_STACK "rocdl" CACHE STRING "Active Fly target lowering stack for this build (one stack per artifact).") @@ -43,4 +45,7 @@ if(FLYDSL_TARGET_STACK STREQUAL "rocdl") MLIRCAPITransforms MLIRCAPIRegisterEverything ) + + set(FLYDSL_BUILD_CONFIG_BACKEND_IDS "\"rocm\", ") + set(FLYDSL_BUILD_CONFIG_RUNTIME_KINDS "\"rocm\", ") endif() diff --git a/python/flydsl/_install_limits.py b/python/flydsl/_install_limits.py new file mode 100644 index 00000000..f5c6d3ef --- /dev/null +++ b/python/flydsl/_install_limits.py @@ -0,0 +1,43 @@ +# SPDX-License-Identifier: Apache-2.0 +# Copyright (c) 2025 FlyDSL Project Contributors + +"""Install-time compile/runtime capability (from CMake-generated ``_build_config``).""" + +from __future__ import annotations + +from typing import Tuple + +# Used when ``flydsl._build_config`` is missing (e.g. ``PYTHONPATH`` to a bare source tree). +_FALLBACK: Tuple[str, tuple[str, ...], tuple[str, ...]] = ("rocdl", ("rocm",), ("rocm",)) + + +def install_caps() -> Tuple[str, tuple[str, ...], tuple[str, ...]]: + """Return ``(TARGET_STACK, compile_backend_ids, runtime_kinds)`` for this install.""" + try: + from flydsl import _build_config as bc + except ImportError: + return _FALLBACK + stack = str(getattr(bc, "TARGET_STACK", _FALLBACK[0])) + backends = tuple(getattr(bc, "ENABLED_COMPILE_BACKEND_IDS", ()) or ()) + runtimes = tuple(getattr(bc, "ENABLED_RUNTIME_KINDS", ()) or ()) + return (stack, backends, runtimes) + + +def ensure_compile_backend_in_build(name: str) -> None: + name = name.strip().lower() + stack, backends, _ = install_caps() + if name not in backends: + raise RuntimeError( + f"This FlyDSL install targets stack {stack!r} and does not support compile backend {name!r}. " + f"Allowed FLYDSL_COMPILE_BACKEND values: {sorted(backends)!r}." + ) + + +def ensure_runtime_kind_in_build(kind: str) -> None: + kind = kind.strip().lower() + stack, _, kinds = install_caps() + if kind not in kinds: + raise RuntimeError( + f"This FlyDSL install targets stack {stack!r} and does not support device runtime kind {kind!r}. " + f"Allowed FLYDSL_RUNTIME_KIND values: {sorted(kinds)!r}." + ) diff --git a/python/flydsl/compiler/backends/__init__.py b/python/flydsl/compiler/backends/__init__.py index 7393d03f..ddf6440c 100644 --- a/python/flydsl/compiler/backends/__init__.py +++ b/python/flydsl/compiler/backends/__init__.py @@ -18,6 +18,7 @@ from functools import lru_cache from typing import Dict, Optional, Type +from ... import _install_limits from ...utils import env from .base import BaseBackend, GPUTarget @@ -74,10 +75,15 @@ def get_backend(name: Optional[str] = None, *, arch: str = "") -> BaseBackend: (via :func:`flydsl.runtime.device_runtime.ensure_compile_runtime_pairing_from_env`) and again on first :func:`flydsl.runtime.device_runtime.get_device_runtime`, not here. + + Raises ``RuntimeError`` if the requested compile backend is not supported by this + install (CMake ``FLYDSL_TARGET_STACK`` / generated ``flydsl._build_config``). """ + if name is None: name = compile_backend_name() name = name.lower() + _install_limits.ensure_compile_backend_in_build(name) if not arch: backend_cls = _registry.get(name) if backend_cls is None: diff --git a/python/flydsl/runtime/device_runtime/__init__.py b/python/flydsl/runtime/device_runtime/__init__.py index d0aedb99..2b611c03 100644 --- a/python/flydsl/runtime/device_runtime/__init__.py +++ b/python/flydsl/runtime/device_runtime/__init__.py @@ -18,6 +18,7 @@ from typing import Dict, Optional, Type +from ... import _install_limits from ...utils import env from .base import DeviceRuntime from .rocm import RocmDeviceRuntime @@ -115,6 +116,11 @@ def ensure_compile_runtime_pairing_from_env(compile_backend_id: str) -> None: :class:`DeviceRuntime`. Suitable for compiler paths (e.g. ``COMPILE_ONLY``) where initializing the runtime is unnecessary. """ + _install_limits.ensure_compile_backend_in_build(compile_backend_id) + if _runtime_cls_override is None: + _install_limits.ensure_runtime_kind_in_build( + (env.runtime.kind or "rocm").strip().lower() + ) expected = _expected_runtime_kind_for_compile_backend(compile_backend_id) actual = _selected_runtime_kind_from_env() if actual != expected: diff --git a/python/mlir_flydsl/CMakeLists.txt b/python/mlir_flydsl/CMakeLists.txt index 51daeb1f..5b6757f2 100644 --- a/python/mlir_flydsl/CMakeLists.txt +++ b/python/mlir_flydsl/CMakeLists.txt @@ -5,6 +5,12 @@ configure_file( "${CMAKE_CURRENT_BINARY_DIR}/FlyRegisterEverything.cpp" @ONLY) +set(_FLYDSL_GENERATED_BUILD_CONFIG "${CMAKE_CURRENT_BINARY_DIR}/generated_flydsl_build_config/_build_config.py") +configure_file( + "${PROJECT_SOURCE_DIR}/cmake/FlyDSLBuildConfig.py.in" + "${_FLYDSL_GENERATED_BUILD_CONFIG}" + @ONLY) + add_compile_definitions("MLIR_PYTHON_PACKAGE_PREFIX=flydsl.${MLIR_PYTHON_PACKAGE_PREFIX}.") declare_mlir_python_sources(FlyPythonSources) @@ -228,6 +234,9 @@ add_custom_target(CopyFlyPythonSources ALL COMMAND ${CMAKE_COMMAND} -E copy_directory "${PROJECT_SOURCE_DIR}/python/flydsl" "${MLIR_BINARY_DIR}/python_packages/flydsl" + COMMAND ${CMAKE_COMMAND} -E copy_if_different + "${_FLYDSL_GENERATED_BUILD_CONFIG}" + "${MLIR_BINARY_DIR}/python_packages/flydsl/_build_config.py" # Copy auto-generated Python bindings from mlir-tblgen to the correct location COMMAND ${CMAKE_COMMAND} -E copy_if_different "${MLIR_BINARY_DIR}/python/mlir_flydsl/dialects/_fly_ops_gen.py" diff --git a/tests/README.md b/tests/README.md index fdeaa474..2cc3ffdf 100644 --- a/tests/README.md +++ b/tests/README.md @@ -2,9 +2,9 @@ Pytest configuration lives in [`pytest.ini`](pytest.ini) in this directory. Run pytest from the **repository root** (or pass `-c tests/pytest.ini`) so this file is picked up. -## Test tiering (RFC 0003) +## Test tiering -The project uses a **layered model** so CI and contributors can select tests by dependency (CPU-only vs MLIR with ROCDL vs real GPU). The full specification is [**RFC 0003: Test tiering and multi-backend CI matrix**](https://github.com/ROCm/FlyDSL/issues/275). +The project uses a **layered model** so CI and contributors can select tests by dependency (CPU-only vs MLIR with ROCDL vs real GPU). Background and design discussion: [GitHub issue #275 — test tiering and multi-backend CI matrix](https://github.com/ROCm/FlyDSL/issues/275). | Tier | Meaning | |------|---------| @@ -89,5 +89,5 @@ export FLYDSL_RUNTIME_ENABLE_CACHE=0 ## MLIR FileCheck tests -`tests/mlir/**/*.mlir` checks are driven by **`scripts/run_tests.sh`** (FileCheck + `fly-opt`), not by pytest. Tiering for those may be documented in parallel in this README as the RFC rollout continues; see RFC open questions. +`tests/mlir/**/*.mlir` checks are driven by **`scripts/run_tests.sh`** (FileCheck + `fly-opt`), not by pytest. Tiering for those may be documented in parallel in this README as the FileCheck runner and marker story evolve. diff --git a/tests/unit/test_device_runtime.py b/tests/unit/test_device_runtime.py index 47798c73..6231ecdf 100644 --- a/tests/unit/test_device_runtime.py +++ b/tests/unit/test_device_runtime.py @@ -47,7 +47,7 @@ def test_unknown_runtime_kind_env(monkeypatch): """Invalid FLYDSL_RUNTIME_KIND fails at compile/runtime pairing first.""" monkeypatch.setenv("FLYDSL_RUNTIME_KIND", "not_a_real_kind") monkeypatch.setenv("FLYDSL_COMPILE_BACKEND", "rocm") - with pytest.raises(RuntimeError, match="requires device runtime kind"): + with pytest.raises(RuntimeError, match="does not support device runtime kind"): dr.get_device_runtime() @@ -56,6 +56,10 @@ def test_unknown_runtime_kind_after_pairing_passes(monkeypatch): dr.register_compile_runtime_mapping("custom", "weird_kind") monkeypatch.setenv("FLYDSL_COMPILE_BACKEND", "custom") monkeypatch.setenv("FLYDSL_RUNTIME_KIND", "weird_kind") + monkeypatch.setattr( + "flydsl._install_limits.install_caps", + lambda: ("rocdl", ("rocm", "custom"), ("rocm", "weird_kind")), + ) try: with pytest.raises(ValueError, match="Unknown FLYDSL_RUNTIME_KIND"): dr.get_device_runtime() @@ -82,5 +86,5 @@ def test_pairing_from_env_no_singleton(monkeypatch): def test_pairing_from_env_mismatch_raises(monkeypatch): monkeypatch.setenv("FLYDSL_COMPILE_BACKEND", "rocm") monkeypatch.setenv("FLYDSL_RUNTIME_KIND", "not_a_registered_kind") - with pytest.raises(RuntimeError, match="requires device runtime kind"): + with pytest.raises(RuntimeError, match="does not support device runtime kind"): dr.ensure_compile_runtime_pairing_from_env("rocm")