# Copyright (c) Advanced Micro Devices, Inc., or its affiliates.
# SPDX-License-Identifier: MIT
#
# Tile Engine build for batched_contraction_multi_abd.
#
# This op uses the dispatcher bridge pattern:
#   - Kernel headers are generated by dispatcher/codegen/unified_contraction_multi_abd_codegen.py
#   - The C++ benchmark entry points (benchmark_single.cpp) are compiled with
#     -include <kernel.hpp> -DCK_TILE_SINGLE_KERNEL_INCLUDE (same as the ctypes libs)
#   - No shared tile_engine utils are modified; all needed structs are local.

set(CONTRACTION_MULTI_ABD_DATATYPE "fp16" CACHE STRING
    "List of datatypes for contraction_multi_abd (semicolon-separated)")
set(CONTRACTION_MULTI_ABD_LAYOUT "rcr" CACHE STRING
    "List of layouts for contraction_multi_abd (semicolon-separated, 3 chars: a+b+e)")
set(CONTRACTION_MULTI_ABD_CONFIG_FILE "" CACHE STRING
    "Custom config file name (without path, must be in configs/ folder)")
set(CONTRACTION_MULTI_ABD_NUM_A_TENSORS "1" CACHE STRING "Number of A tensors")
set(CONTRACTION_MULTI_ABD_NUM_B_TENSORS "1" CACHE STRING "Number of B tensors")
set(CONTRACTION_MULTI_ABD_NUM_D_TENSORS "1" CACHE STRING "Number of D tensors")
set(CONTRACTION_MULTI_ABD_NUM_DIM_G "1" CACHE STRING "Number of G (batch) dimensions")
set(CONTRACTION_MULTI_ABD_NUM_DIM_M "2" CACHE STRING "Number of M dimensions")
set(CONTRACTION_MULTI_ABD_NUM_DIM_N "2" CACHE STRING "Number of N dimensions")
set(CONTRACTION_MULTI_ABD_NUM_DIM_K "1" CACHE STRING "Number of K dimensions")
set(CONTRACTION_MULTI_ABD_A_ELEMENTWISE "PassThrough" CACHE STRING "A elementwise function")
set(CONTRACTION_MULTI_ABD_B_ELEMENTWISE "PassThrough" CACHE STRING "B elementwise function")
set(CONTRACTION_MULTI_ABD_CDE_ELEMENTWISE "MultiDAdd" CACHE STRING "CDE elementwise function")
set(CONTRACTION_MULTI_ABD_MAX_INSTANCES "" CACHE STRING
    "Max kernel instances (empty = no cap)")

option(ENABLE_CCACHE_CONTRACTION_MULTI_ABD
    "Enable ccache for contraction_multi_abd compilation" OFF)

set(CONTRACTION_MULTI_ABD_SOURCE_DIR ${CMAKE_CURRENT_LIST_DIR})

# Path to the dispatcher codegen script
set(_CMA_CODEGEN_SCRIPT
    "${CMAKE_CURRENT_LIST_DIR}/../../../../dispatcher/codegen/unified_contraction_multi_abd_codegen.py")

# Path to the dispatcher ctypes C++ source (force-included compile)
set(_CMA_CTYPES_SRC
    "${CMAKE_CURRENT_LIST_DIR}/../../../../dispatcher/bindings/ctypes/batched_contraction_multi_abd_ctypes_lib.cpp")

# ---------------------------------------------------------------------------
# GPU target filter (same set as batched_contraction: gfx90a, gfx942, gfx950)
# ---------------------------------------------------------------------------
set(CONTRACTION_MULTI_ABD_GPU_TARGETS "")
set(_CMA_DESIRED_TARGETS "gfx90a;gfx942;gfx950")
foreach(target IN LISTS SUPPORTED_GPU_TARGETS)
    if(target IN_LIST _CMA_DESIRED_TARGETS)
        list(APPEND CONTRACTION_MULTI_ABD_GPU_TARGETS ${target})
    endif()
endforeach()

if(NOT CONTRACTION_MULTI_ABD_GPU_TARGETS)
    message(WARNING
        "Skipping contraction_multi_abd build: no supported GPU targets "
        "(gfx90a, gfx942, gfx950) found in SUPPORTED_GPU_TARGETS: ${SUPPORTED_GPU_TARGETS}")
    return()
endif()

message(STATUS "contraction_multi_abd: building for GPU targets: ${CONTRACTION_MULTI_ABD_GPU_TARGETS}")

# ---------------------------------------------------------------------------
# Compiler cache
# ---------------------------------------------------------------------------
if(ENABLE_CCACHE_CONTRACTION_MULTI_ABD)
    find_program(CCACHE_PROGRAM ccache)
    if(CCACHE_PROGRAM)
        set(_PREV_CXX_LAUNCHER "${CMAKE_CXX_COMPILER_LAUNCHER}")
        set(CMAKE_CXX_COMPILER_LAUNCHER ${CCACHE_PROGRAM})
        set(_CCACHE_APPLIED_CMA TRUE)
    else()
        message(WARNING "ccache requested but not found for contraction_multi_abd")
    endif()
endif()

# ---------------------------------------------------------------------------
# Helper function: create one benchmark target for (datatype, layout, kernel_name)
# ---------------------------------------------------------------------------
function(create_contraction_multi_abd_target datatype layout kernel_name instance_header)
    if(NOT CONTRACTION_MULTI_ABD_GPU_TARGETS)
        return()
    endif()

    set(target "benchmark_contraction_multi_abd_${kernel_name}")

    add_executable(${target}
        EXCLUDE_FROM_ALL
        ${CONTRACTION_MULTI_ABD_SOURCE_DIR}/contraction_multi_abd_benchmark_single.cpp
        ${instance_header}
    )

    set_property(TARGET ${target} PROPERTY HIP_ARCHITECTURES ${CONTRACTION_MULTI_ABD_GPU_TARGETS})

    target_include_directories(${target} PRIVATE
        ${CONTRACTION_MULTI_ABD_SOURCE_DIR}
    )

    target_compile_options(${target} PRIVATE
        -Wno-undefined-func-template
        -Wno-float-equal
        -DCK_TILE_SINGLE_KERNEL_INCLUDE
        -include ${instance_header}
    )

    add_dependencies(benchmark_contraction_multi_abd_all ${target})
    add_dependencies(benchmark_contraction_multi_abd_${datatype} ${target})
    add_dependencies(benchmark_contraction_multi_abd_${layout} ${target})
    add_dependencies(benchmark_contraction_multi_abd_${datatype}_${layout} ${target})
endfunction()

# ---------------------------------------------------------------------------
# Collection targets
# ---------------------------------------------------------------------------
add_custom_target(benchmark_contraction_multi_abd_all)

foreach(dt IN LISTS CONTRACTION_MULTI_ABD_DATATYPE)
    add_custom_target(benchmark_contraction_multi_abd_${dt})
endforeach()

foreach(l IN LISTS CONTRACTION_MULTI_ABD_LAYOUT)
    add_custom_target(benchmark_contraction_multi_abd_${l})
endforeach()

foreach(dt IN LISTS CONTRACTION_MULTI_ABD_DATATYPE)
    foreach(l IN LISTS CONTRACTION_MULTI_ABD_LAYOUT)
        add_custom_target(benchmark_contraction_multi_abd_${dt}_${l})
    endforeach()
endforeach()

# ---------------------------------------------------------------------------
# For each (datatype, layout): run codegen to list kernels, then create targets
# ---------------------------------------------------------------------------
foreach(datatype IN LISTS CONTRACTION_MULTI_ABD_DATATYPE)
    foreach(layout IN LISTS CONTRACTION_MULTI_ABD_LAYOUT)
        set(working_path "${CMAKE_CURRENT_BINARY_DIR}/${datatype}/${layout}")
        file(MAKE_DIRECTORY ${working_path})

        # Select config file
        if(DEFINED ENV{CONTRACTION_MULTI_ABD_CONFIG_FILE}
           AND NOT "$ENV{CONTRACTION_MULTI_ABD_CONFIG_FILE}" STREQUAL "")
            set(config_file "$ENV{CONTRACTION_MULTI_ABD_CONFIG_FILE}")
            set(json_blob "${CMAKE_CURRENT_LIST_DIR}/configs/${config_file}")
        elseif(NOT "${CONTRACTION_MULTI_ABD_CONFIG_FILE}" STREQUAL "")
            set(json_blob "${CMAKE_CURRENT_LIST_DIR}/configs/${CONTRACTION_MULTI_ABD_CONFIG_FILE}")
        else()
            set(json_blob "${CMAKE_CURRENT_LIST_DIR}/configs/default_config.json")
        endif()

        if(NOT EXISTS ${json_blob})
            message(FATAL_ERROR "Config file not found: ${json_blob}")
        endif()

        # Build the per-operation overrides (dtype, layout, tensor counts, dim counts,
        # elementwise ops) as a JSON fragment that will be MERGED on top of the tile/trait
        # ranges read from json_blob. Tile/trait ranges come from the selected config file;
        # per-operation scalars override only the fields listed here.
        # Include max_instances only when non-empty so codegen can cap kernel count.
        if(CONTRACTION_MULTI_ABD_MAX_INSTANCES STREQUAL "")
            set(_cma_max_inst_json "")
        else()
            set(_cma_max_inst_json ",\"max_instances\":${CONTRACTION_MULTI_ABD_MAX_INSTANCES}")
        endif()

        set(_cma_overrides_json
            "{\"dtypes\":[\"${datatype}\"],\"layouts\":[\"${layout}\"],\
\"num_a_tensors\":[${CONTRACTION_MULTI_ABD_NUM_A_TENSORS}],\
\"num_b_tensors\":[${CONTRACTION_MULTI_ABD_NUM_B_TENSORS}],\
\"num_d_tensors\":[${CONTRACTION_MULTI_ABD_NUM_D_TENSORS}],\
\"dim_combos\":[{\"num_dim_g\":${CONTRACTION_MULTI_ABD_NUM_DIM_G},\
\"num_dim_m\":${CONTRACTION_MULTI_ABD_NUM_DIM_M},\
\"num_dim_n\":${CONTRACTION_MULTI_ABD_NUM_DIM_N},\
\"num_dim_k\":${CONTRACTION_MULTI_ABD_NUM_DIM_K}}],\
\"a_elementwise\":\"${CONTRACTION_MULTI_ABD_A_ELEMENTWISE}\",\
\"b_elementwise\":\"${CONTRACTION_MULTI_ABD_B_ELEMENTWISE}\",\
\"cde_elementwise\":\"${CONTRACTION_MULTI_ABD_CDE_ELEMENTWISE}\"${_cma_max_inst_json}}"
        )

        # Merge: start from json_blob (tile/trait ranges), then layer overrides on top.
        # Use a Python one-liner so we do not need to parse JSON in CMake.
        set(_cma_config_file "${working_path}/codegen_config.json")
        execute_process(
            COMMAND ${Python3_EXECUTABLE} -c
                "import json,sys; \
base=json.load(open('${json_blob}')); \
ovr=json.loads(sys.argv[1]); \
base.update(ovr); \
open('${_cma_config_file}','w').write(json.dumps(base))"
                "${_cma_overrides_json}"
            RESULT_VARIABLE _cma_merge_ret
            ERROR_VARIABLE  _cma_merge_err
        )
        if(NOT _cma_merge_ret EQUAL 0)
            message(FATAL_ERROR
                "Failed to merge contraction_multi_abd config JSON: ${_cma_merge_err}")
        endif()

        # Run codegen to list kernel names (--list-name)
        execute_process(
            COMMAND ${Python3_EXECUTABLE} ${_CMA_CODEGEN_SCRIPT}
                    --output-dir ${working_path}
                    --config ${_cma_config_file}
                    --list-name
            WORKING_DIRECTORY ${CMAKE_CURRENT_LIST_DIR}
            RESULT_VARIABLE _cma_ret
            OUTPUT_VARIABLE _cma_names
            ERROR_VARIABLE  _cma_err
        )
        if(NOT _cma_ret EQUAL 0)
            message(FATAL_ERROR
                "contraction_multi_abd codegen (--list-name) failed for "
                "${datatype} ${layout}:\n${_cma_err}")
        endif()

        # Generate all headers
        execute_process(
            COMMAND ${Python3_EXECUTABLE} ${_CMA_CODEGEN_SCRIPT}
                    --output-dir ${working_path}
                    --config ${_cma_config_file}
            WORKING_DIRECTORY ${CMAKE_CURRENT_LIST_DIR}
            RESULT_VARIABLE _cma_ret2
            OUTPUT_VARIABLE _cma_out2
            ERROR_VARIABLE  _cma_err2
        )
        if(NOT _cma_ret2 EQUAL 0)
            message(FATAL_ERROR
                "contraction_multi_abd codegen failed for ${datatype} ${layout}:\n${_cma_err2}")
        endif()

        # Parse kernel names and create per-kernel targets
        string(REPLACE "\n" ";" _cma_name_list "${_cma_names}")
        foreach(kernel_name IN LISTS _cma_name_list)
            string(STRIP "${kernel_name}" kernel_name)
            if(NOT kernel_name STREQUAL "")
                set(_hpp "${working_path}/${kernel_name}.hpp")
                if(EXISTS ${_hpp})
                    create_contraction_multi_abd_target(
                        "${datatype}" "${layout}" "${kernel_name}" "${_hpp}")
                else()
                    message(WARNING
                        "contraction_multi_abd: expected header not found: ${_hpp}")
                endif()
            endif()
        endforeach()
    endforeach()
endforeach()

# ---------------------------------------------------------------------------
# Restore CXX launcher if ccache was applied
# ---------------------------------------------------------------------------
if(_CCACHE_APPLIED_CMA)
    set(CMAKE_CXX_COMPILER_LAUNCHER "${_PREV_CXX_LAUNCHER}")
    unset(_CCACHE_APPLIED_CMA)
endif()
