# Author: Simon-Pierre Boucher — contact@spboucher.ai
cmake_minimum_required(VERSION 3.24)
project(forge LANGUAGES CXX)

set(CMAKE_CXX_STANDARD 20)
set(CMAKE_CXX_STANDARD_REQUIRED ON)
set(CMAKE_CXX_EXTENSIONS OFF)

if(NOT CMAKE_BUILD_TYPE)
  set(CMAKE_BUILD_TYPE Release)
endif()

add_compile_options(-Wall -Wextra)

# ---------------------------------------------------------------------------
# Metal shader library: .metal -> .air -> forge.metallib
# Compiled with -std=metal3.2 (macOS 15+ baseline: simdgroup_matrix, bfloat).
# Debug builds embed sources for the Xcode GPU debugger.
# ---------------------------------------------------------------------------
file(GLOB METAL_SOURCES CONFIGURE_DEPENDS ${CMAKE_SOURCE_DIR}/src/kernels/*.metal)
set(AIR_FILES "")
foreach(shader ${METAL_SOURCES})
  get_filename_component(shader_name ${shader} NAME_WE)
  set(air ${CMAKE_BINARY_DIR}/${shader_name}.air)
  # Metal Performance Primitives (cooperative tensors) need the Metal 4
  # language; everything else targets 3.2 so it runs on macOS 15+.
  if(shader_name MATCHES "_mpp$")
    set(metal_std "-std=metal4.0")
  else()
    set(metal_std "-std=metal3.2")
  endif()
  add_custom_command(
    OUTPUT ${air}
    # -frecord-sources is unconditional, not Debug-only: without shader sources
    # embedded in the metallib, the Metal debugger and gpudebug report
    # "no source" and every per-shader/per-line profiling view comes back
    # empty. It costs file size, nothing else.
    COMMAND xcrun -sdk macosx metal ${metal_std} -O2
            -gline-tables-only -frecord-sources
            -c ${shader} -o ${air}
    DEPENDS ${shader}
    COMMAND_EXPAND_LISTS
    COMMENT "Compiling Metal shader ${shader_name}.metal"
    VERBATIM)
  list(APPEND AIR_FILES ${air})
endforeach()

add_custom_command(
  OUTPUT ${CMAKE_BINARY_DIR}/forge.metallib
  # Link with `metal`, not `metallib`: the sources only survive into the
  # library if -frecord-sources is passed at link time too, and `metallib`
  # rejects that flag ("unknown argument").
  COMMAND xcrun -sdk macosx metal -frecord-sources ${AIR_FILES}
          -o ${CMAKE_BINARY_DIR}/forge.metallib
  DEPENDS ${AIR_FILES}
  COMMENT "Linking forge.metallib"
  VERBATIM)
add_custom_target(forge_kernels ALL DEPENDS ${CMAKE_BINARY_DIR}/forge.metallib)

# ---------------------------------------------------------------------------
# Core library
# ---------------------------------------------------------------------------
add_library(forge_core STATIC
  src/core/metal_impl.cpp
  src/core/tensor.cpp
  src/core/allocator.cpp
  src/core/device.cpp
  src/core/autograd.cpp
  src/core/fmodel.cpp
  src/ops/cpu/cpu_ops.cpp
  src/ops/metal/metal_ops.cpp
  src/ops/ops.cpp
  src/nn/linear.cpp
  src/train/optimizer.cpp
  src/train/dataloader.cpp
  src/train/checkpoint.cpp
  src/train/trainer.cpp
  src/tokenizer/bpe.cpp)

target_include_directories(forge_core PUBLIC
  ${CMAKE_SOURCE_DIR}/src
  ${CMAKE_SOURCE_DIR}/third_party
  ${CMAKE_SOURCE_DIR}/third_party/metal-cpp)

target_link_libraries(forge_core PUBLIC
  "-framework Metal" "-framework Foundation" "-framework QuartzCore")

add_dependencies(forge_core forge_kernels)

# ---------------------------------------------------------------------------
# CLI
# ---------------------------------------------------------------------------
add_executable(forge src/main.cpp)
target_link_libraries(forge PRIVATE forge_core)

# ---------------------------------------------------------------------------
# Tests
# ---------------------------------------------------------------------------
enable_testing()
foreach(t test_ops test_gradcheck test_overfit test_tokenizer)
  add_executable(${t} tests/${t}.cpp)
  target_link_libraries(${t} PRIVATE forge_core)
  add_test(NAME ${t} COMMAND ${t} WORKING_DIRECTORY ${CMAKE_BINARY_DIR})
endforeach()

# Benchmarks: built, not run by ctest.
foreach(b bench_matmul bench_attention bench_precision mppcheck)
  add_executable(${b} tests/${b}.cpp)
  target_link_libraries(${b} PRIVATE forge_core)
endforeach()