CMakeLists.txt

cmake_minimum_required(VERSION 3.10)
project(musa_plugin)

# Force GCC compiler to align with TF core library ABI
set(CMAKE_C_COMPILER gcc)
set(CMAKE_CXX_COMPILER g++)

set(CMAKE_CXX_STANDARD 14)
set(CMAKE_CXX_STANDARD_REQUIRED ON)

# Kernel debug/timing instrumentation switch.
# OFF by default to keep production performance unchanged.
option(MUSA_KERNEL_DEBUG "Enable MUSA kernel compute timing instrumentation" OFF)

# Set MUSA path
set(MUSA_PATH "/usr/local/musa" CACHE STRING "Path to MUSA installation")
set(MCC_EXECUTABLE "${MUSA_PATH}/bin/mcc")

# Validate MUSA path exists
if(NOT EXISTS "${MUSA_PATH}")
    message(FATAL_ERROR "MUSA_PATH '${MUSA_PATH}' does not exist. Please set correct MUSA installation path.")
endif()

find_library(MUDNNCXX_LIBRARY NAMES mudnncxx PATHS "${MUSA_PATH}/lib" "${MUSA_PATH}/lib64" NO_DEFAULT_PATH)
find_library(MUDNN_LIBRARY NAMES mudnn PATHS "${MUSA_PATH}/lib" "${MUSA_PATH}/lib64" NO_DEFAULT_PATH)

if(MUDNNCXX_LIBRARY)
    set(MUSA_DNN_LIBRARY "${MUDNNCXX_LIBRARY}")
elseif(MUDNN_LIBRARY)
    set(MUSA_DNN_LIBRARY "${MUDNN_LIBRARY}")
else()
    message(FATAL_ERROR "Failed to find mudnncxx or mudnn under '${MUSA_PATH}/lib' or '${MUSA_PATH}/lib64'.")
endif()

# Get TensorFlow parameters
execute_process(COMMAND python3 -c "import tensorflow as tf; print(tf.sysconfig.get_include())"
    OUTPUT_VARIABLE TF_INC OUTPUT_STRIP_TRAILING_WHITESPACE)
execute_process(COMMAND python3 -c "import tensorflow as tf; print(' '.join(tf.sysconfig.get_compile_flags()))"
    OUTPUT_VARIABLE TF_COMPILE_FLAGS OUTPUT_STRIP_TRAILING_WHITESPACE)
execute_process(COMMAND python3 -c "import tensorflow as tf; print(' '.join(tf.sysconfig.get_link_flags()))"
    OUTPUT_VARIABLE TF_LINK_FLAGS OUTPUT_STRIP_TRAILING_WHITESPACE)
execute_process(COMMAND python3 -c "import tensorflow as tf; import os; print(os.path.dirname(tf.__file__))"
    OUTPUT_VARIABLE TF_LIB_DIR OUTPUT_STRIP_TRAILING_WHITESPACE)

# Validate TensorFlow is available
if(NOT TF_INC)
    message(FATAL_ERROR "Failed to get TensorFlow include path. Please ensure TensorFlow is installed.")
endif()

# Clean ABI conflicts in TF_COMPILE_FLAGS
string(REPLACE "-D_GLIBCXX_USE_CXX11_ABI=0" "" TF_COMPILE_FLAGS "${TF_COMPILE_FLAGS}")
string(REPLACE "-D_GLIBCXX_USE_CXX11_ABI=1" "" TF_COMPILE_FLAGS "${TF_COMPILE_FLAGS}")
# Re-add the safest definition
set(TF_COMPILE_FLAGS "${TF_COMPILE_FLAGS} -D_GLIBCXX_USE_CXX11_ABI=0")

# Configure debug/logging-related compile definitions.
#
# Important: this plugin links against the pip TensorFlow wheel, which is built
# in release mode. Some TensorFlow headers (for example refcount.h) contain
# inline DCHECK-dependent side effects. If the plugin is compiled without
# NDEBUG while libtensorflow_framework is compiled with NDEBUG, the plugin and
# framework can disagree on RefCounted teardown semantics and trigger false
# refcount aborts during session/resource destruction.
#
# Keep NDEBUG enabled in every build to match the TensorFlow wheel, and use
# MUSA_KERNEL_DEBUG only for plugin-local timing/debug instrumentation.
if(MUSA_KERNEL_DEBUG)
    set(TF_COMPILE_FLAGS "${TF_COMPILE_FLAGS} -DMUSA_KERNEL_DEBUG -DNDEBUG")
else()
    set(TF_COMPILE_FLAGS "${TF_COMPILE_FLAGS} -DMUSA_DISABLE_DEBUG_LOGGING -DNDEBUG -DMUSA_DISABLE_TRACE_LOGGING -DMUSA_DISABLE_AUTO_TRACE")
endif()

# Include paths
include_directories("${TF_INC}" "${MUSA_PATH}/include")
if(EXISTS "${MUSA_PATH}/include/mudnncxx")
    include_directories("${MUSA_PATH}/include/mudnncxx")
endif()

# Include TensorFlow C++ core headers for common_runtime
# These are required for accessing BFCAllocator and device-specific headers
include_directories("${TF_INC}/tensorflow/core/common_runtime")
include_directories("${TF_INC}/tensorflow/core/common_runtime/device")

# TensorFlow pip headers may place farmhash in different external include layouts.
# fingerprint.h includes <farmhash.h>, so add whichever location exists.
set(TF_FARMHASH_INC_SRC "${TF_INC}/external/farmhash_archive/src")
set(TF_FARMHASH_INC_VIRTUAL "${TF_INC}/external/farmhash_archive/_virtual_includes/farmhash")
if(EXISTS "${TF_FARMHASH_INC_SRC}/farmhash.h")
    include_directories("${TF_FARMHASH_INC_SRC}")
endif()
if(EXISTS "${TF_FARMHASH_INC_VIRTUAL}/farmhash.h")
    include_directories("${TF_FARMHASH_INC_VIRTUAL}")
endif()

# Define source root directory
set(SRC_ROOT "${CMAKE_CURRENT_SOURCE_DIR}/musa_ext")
include_directories("${SRC_ROOT}" "${SRC_ROOT}/mu" "${SRC_ROOT}/kernels")

# Validate source directory exists
if(NOT EXISTS "${SRC_ROOT}")
    message(FATAL_ERROR "Source directory '${SRC_ROOT}' does not exist.")
endif()

# Link paths
link_directories("${MUSA_PATH}/lib" "${MUSA_PATH}/lib64" "${TF_LIB_DIR}")

# Automatically collect source files
file(GLOB_RECURSE CPP_SOURCES
    "${SRC_ROOT}/*.cc"
    "${SRC_ROOT}/mu/*.cc"
    "${SRC_ROOT}/kernels/*.cc"
)

# filter out disabled kernels (if any)
# list(FILTER CPP_SOURCES EXCLUDE REGEX ".*test.*|.*disabled.*")
# Automatically collect MUSA Kernel files
file(GLOB_RECURSE MU_SOURCES "${SRC_ROOT}/kernels/*.mu")

# TensorFlow headers pulled into `.mu` compilation require C++17, and mcc's
# aarch64 host predefines can incorrectly trigger NEON-only host headers.
set(MCC_COMMON_FLAGS
    -O3
    -fPIC
    -std=c++17
    -DEIGEN_DONT_VECTORIZE
    -U__ARM_NEON
    -U__ARM_NEON__
    -I${MUSA_PATH}/include
    -I${TF_INC}
)

# MCC compile MUSA Kernel
set(MU_OBJECTS "")
foreach(mu_file ${MU_SOURCES})
    get_filename_component(filename ${mu_file} NAME)
    if(NOT filename MATCHES "^disabled_")
        file(RELATIVE_PATH mu_rel_path "${SRC_ROOT}/kernels" "${mu_file}")
        string(REPLACE "/" "_" mu_rel_obj "${mu_rel_path}")
        string(REPLACE ".mu" "" mu_rel_obj "${mu_rel_obj}")
        set(obj_file "${CMAKE_CURRENT_BINARY_DIR}/${mu_rel_obj}.mu.o")
        set(mu_full_path "${mu_file}")

        add_custom_command(
            OUTPUT ${obj_file}
            COMMAND ${MCC_EXECUTABLE} -c ${mu_full_path} -o ${obj_file} ${MCC_COMMON_FLAGS}
            DEPENDS ${mu_full_path}
            COMMENT "Compiling MUSA kernel: ${filename}"
        )
        list(APPEND MU_OBJECTS ${obj_file})
    endif()
endforeach()

# Generate plugin library
add_library(musa_plugin SHARED ${CPP_SOURCES} ${MU_OBJECTS})

# Apply TF compile flags (ensure strict ABI consistency)
set_target_properties(musa_plugin PROPERTIES COMPILE_FLAGS "${TF_COMPILE_FLAGS}")

# Set RPATH to ensure runtime library search
set_target_properties(musa_plugin PROPERTIES
    SKIP_BUILD_RPATH FALSE
    BUILD_WITH_INSTALL_RPATH TRUE
    INSTALL_RPATH "${TF_LIB_DIR};${MUSA_PATH}/lib;${MUSA_PATH}/lib64"
)

# Link
target_link_libraries(musa_plugin PRIVATE ${TF_LINK_FLAGS} musa mublas ${MUSA_DNN_LIBRARY} rt)

# Output configuration
set_target_properties(musa_plugin PROPERTIES
    PREFIX "lib"
    SUFFIX ".so"
    OUTPUT_NAME "musa_plugin"
)

# Build completion message
add_custom_command(TARGET musa_plugin POST_BUILD
    COMMAND ${CMAKE_COMMAND} -E echo "Build completed successfully!"
    COMMENT "Plugin is located at: $<TARGET_FILE:musa_plugin>"
)

# Print build summary
list(LENGTH CPP_SOURCES CPP_SOURCES_COUNT)
list(LENGTH MU_SOURCES MU_SOURCES_COUNT)

message(STATUS "TensorFlow MUSA Plugin Configuration:")
message(STATUS "  MUSA Path: ${MUSA_PATH}")
message(STATUS "  MuDNN Library: ${MUSA_DNN_LIBRARY}")
message(STATUS "  TensorFlow Include: ${TF_INC}")
message(STATUS "  MUSA Kernel Debug: ${MUSA_KERNEL_DEBUG}")
message(STATUS "  Source Files: ${CPP_SOURCES_COUNT} C++ files, ${MU_SOURCES_COUNT} MUSA kernels")
message(STATUS "  Output: ${CMAKE_CURRENT_BINARY_DIR}/libmusa_plugin.so")