include(CMakeParseArguments)

# springbok_bytecode_module()
#
# A wrapper for the iree_bytecode_module to apply common iree-translate flags
# Parameters:
# NAME: Name of target.
# SRC: Source file to compile into a bytecode module. Support relative path.
# FLAGS: Flags to pass to the translation tool (list of strings).
# C_IDENTIFIER: Identifier to use for generate c embed code.
#     If omitted then no C embed code will be generated.
#
# Examples:
# springbok_bytecode_module(
#   NAME
#     dare_devel_bytecode_module_dylib
#   SRC
#     "daredevil_quant.tflite"
#   C_IDENTIFIER
#     "daredevil_bytecode_module_dylib"
#   FLAGS
#     "-iree-input-type=tosa"
# )
#
# springbok_bytecode_module(
#   NAME
#     simple_float_mul_bytecode_module_dylib
#   SRC
#     "simple_float_mul.mlir"
#   C_IDENTIFIER
#     "simple_float_mul_bytecode_module_dylib"
#   FLAGS
#     "-iree-input-type=mhlo"
#   PUBLIC
# )
#
function(springbok_bytecode_module)
  cmake_parse_arguments(
    _RULE
    ""
    "NAME;SRC;C_IDENTIFIER"
    "FLAGS"
    ${ARGN}
  )

  set(_MLIR_SRC "${_RULE_SRC}")
  string(FIND "${_RULE_SRC}" ".tflite" _IS_TFLITE REVERSE)
  if(${_IS_TFLITE} GREATER 0)
    find_program(IREE_IMPORT_TFLITE_TOOL "iree-import-tflite" REQUIRED)
    set(_MLIR_SRC "${CMAKE_CURRENT_BINARY_DIR}/${_RULE_NAME}.mlir")
    get_filename_component(_SRC_PATH "${_RULE_SRC}" REALPATH)
    set(_ARGS "${_SRC_PATH}")
    list(APPEND _ARGS "-o")
    list(APPEND _ARGS "${_RULE_NAME}.mlir")
    # Only add the custom_command here. The output is passed to
    # iree_bytecode_module as the source.
    add_custom_command(
      OUTPUT
        "${_RULE_NAME}.mlir"
      COMMAND
        ${IREE_IMPORT_TFLITE_TOOL}
        ${_ARGS}
      DEPENDS
        ${IREE_IMPORT_TFLITE_TOOL}
    )
  endif()

  iree_bytecode_module(
    NAME
      "${_RULE_NAME}"
    SRC
      "${_MLIR_SRC}"
    C_IDENTIFIER
      "${_RULE_C_IDENTIFIER}"
    FLAGS
      "-iree-mlir-to-vm-bytecode-module"
      "-iree-hal-target-backends=dylib-llvm-aot"
      "-iree-llvm-target-triple=riscv32-pc-linux-elf"
      "-iree-llvm-target-cpu=generic-rv32"
      "-iree-llvm-target-cpu-features=+m,+f,+experimental-v"
      "-iree-llvm-target-abi=ilp32"
      "-iree-llvm-link-embedded=true"
      "-iree-llvm-debug-symbols=false"
      "${_RULE_FLAGS}"
    PUBLIC
  )
endfunction()