add_library(modelstudio_native STATIC
  backends/cpu/cpu_backend.cpp
  backends/cpu/kernels/add.cpp
  backends/cpu/kernels/matmul.cpp
  backends/cpu/kernels/mul.cpp
  backends/cpu/kernels/relu.cpp
)

target_include_directories(modelstudio_native PUBLIC ${CMAKE_CURRENT_SOURCE_DIR})

if(MODELSTUDIO_ENABLE_CUDA)
  target_sources(modelstudio_native PRIVATE
    backends/cuda/cuda_backend.cu
    backends/cuda/cuda_context.cu
    backends/cuda/cuda_memory.cu
    backends/cuda/cuda_stream.cu
    backends/cuda/kernels/elementwise.cu
    backends/cuda/kernels/reductions.cu
    backends/cuda/kernels/matmul.cu
  )
  target_compile_definitions(modelstudio_native PUBLIC MODELSTUDIO_ENABLE_CUDA=1)
  target_link_libraries(modelstudio_native PUBLIC CUDA::cudart CUDA::cublas)
endif()

if(MODELSTUDIO_ENABLE_ROCM)
  target_sources(modelstudio_native PRIVATE backends/rocm/rocm_backend.cpp)
  target_compile_definitions(modelstudio_native PUBLIC MODELSTUDIO_ENABLE_ROCM=1)
endif()

if(MODELSTUDIO_ENABLE_ONEAPI)
  target_sources(modelstudio_native PRIVATE backends/oneapi/oneapi_backend.cpp)
  target_compile_definitions(modelstudio_native PUBLIC MODELSTUDIO_ENABLE_ONEAPI=1)
endif()
