# Copyright 2023 Google LLC
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

include_directories(BEFORE SYSTEM ${CMAKE_CURRENT_BINARY_DIR})

# -------------------------------------------------------------------------------
# Build the mlir bytecode modules with iree-compile. Note the last two flags
# are for RVV support.
# -------------------------------------------------------------------------------

sparrow_modules(
  NAME
    hps
  SRC
    "hps_quant.tflite"
  C_IDENTIFIER
    "quant_models_hps"
  FLAGS
    "-iree-input-type=tosa"
  RVV_OFF
)

sparrow_modules(
  NAME
    mobilenet_v1
  SRC
    "mobilenet_v1_0.25_224_quant.tflite"
  C_IDENTIFIER
    "quant_models_mobilenet_v1"
  FLAGS
    "-iree-input-type=tosa"
    "-riscv-v-fixed-length-vector-lmul-max=8"
)

sparrow_modules(
  NAME
    mobilenet_v2
  SRC
    "mobilenet_v2_1.0_224_quant.tflite"
  C_IDENTIFIER
    "quant_models_mobilenet_v2"
  FLAGS
    "-iree-input-type=tosa"
    "-riscv-v-fixed-length-vector-lmul-max=8"
)

sparrow_modules(
  NAME
    person_detection
  SRC
    "person_detection.tflite"
  C_IDENTIFIER
    "quant_models_person_detection"
  FLAGS
    "-iree-input-type=tosa"
    "-riscv-v-fixed-length-vector-lmul-max=8"
  VMVX
)

# ------------------------------------------------------------------------------
# Binaries to execute the IREE model input
# ------------------------------------------------------------------------------

foreach(N RANGE 0 6)
  iree_model_input(
    NAME
      hps_${N}_quant_input
    SHAPE
      "1, 240, 320, 1"
    SRC
      "$ENV{ROOTDIR}/ml/ml-models-public/test_data/hps_${N}.jpg"
    QUANT
    SIGNED
  )
endforeach()

iree_model_input(
  NAME
    mobilenet_quant_input
  SHAPE
    "1, 224, 224, 3"
  SRC
    "https://storage.googleapis.com/download.tensorflow.org/ \
    example_images/YellowLabradorLooking_new.jpg"
  QUANT
)

iree_model_input(
  NAME
    person_detection_quant_input
  SHAPE
    "1, 96, 96, 1"
  SRC
    "https://github.com/tensorflow/tflite-micro/raw/aeac6f39e5c7475cea20c54e86d41e3a38312546/ \
    tensorflow/lite/micro/examples/person_detection/testdata/person.bmp"
  QUANT
)

# ------------------------------------------------------------------------------
# Binaries to execute the MLIR bytecode modules
# ------------------------------------------------------------------------------

# If the program requires a larger stack size, add
#
# LINKOPTS
# "LINKER:--defsym=__stack_size__=<new stack size>"
#
# to increase it.

iree_cc_binary(
  NAME
    mobilenet_v1_bytecode_static
  SRCS
    "iree_exec/mobilenet_v1.c"
  DEPS
    ::mobilenet_quant_input_c
    ::mobilenet_v1_bytecode_module_static_c
    ::mobilenet_v1_bytecode_module_static_lib
    iree::vm::bytecode::module
    model_util::util_static
  LINKOPTS
    "LINKER:--defsym=__stack_size__=300k"
)

sparrow_test(
  NAME
    mobilenet_v1_emitc_static
  SRCS
    "iree_exec/mobilenet_v1.c"
  DEPS
    ::mobilenet_quant_input_c
    ::mobilenet_v1_c_module_static_emitc
    ::mobilenet_v1_c_module_static_lib
    model_util::util_static
  LINKOPTS
    "LINKER:--defsym=__stack_size__=300k"
    "LINKER:--defsym=__heap_size__=1500k"
    "LINKER:--defsym=__tcm_length__=4M"
  COPTS
    "-DBUILD_EMITC"
  TESTFILES
    "iree_exec/mobilenet_v1_test.run"
    "iree_exec/mobilenet_v1_test.filecheck"
)

iree_cc_binary(
  NAME
    person_detection_bytecode_static
  SRCS
    "iree_exec/person_detection.c"
  DEPS
    ::person_detection_bytecode_module_static_c
    ::person_detection_bytecode_module_static_lib
    ::person_detection_quant_input_c
    iree::vm::bytecode::module
    model_util::util_static
  LINKOPTS
    "LINKER:--defsym=__stack_size__=128k"
    "LINKER:--defsym=__heap_size__=200k"
)

sparrow_test(
  NAME
    person_detection_emitc_static
  SRCS
    "iree_exec/person_detection.c"
  DEPS
    ::person_detection_c_module_static_emitc
    ::person_detection_c_module_static_lib
    ::person_detection_quant_input_c
    model_util::util_static
  LINKOPTS
    "LINKER:--defsym=__stack_size__=128k"
    "LINKER:--defsym=__heap_size__=200k"
  COPTS
    "-DBUILD_EMITC"
  TESTFILES
    "iree_exec/person_detection_test.run"
    "iree_exec/person_detection_test.filecheck"
)

# Do not compile following models for Kelvin
if (${BUILD_WITH_KELVIN})
  return()
endif()

foreach(N RANGE 0 6)
  set(HPS_COPTS "-DHPS_QUANT_INPUT=hps_${N}_quant_input")
  list(APPEND HPS_COPTS "-DHPS_QUANT_INPUT_C_H=\"hps_${N}_quant_input_c.h\"")
  if (NOT BUILD_NO_WMMU)
    list(APPEND HPS_COPTS "-DNO_STATIC_INPUT")
  endif()

  iree_cc_binary(
    NAME
      hps_${N}_bytecode_static
    SRCS
      "iree_exec/hps.c"
    DEPS
      ::hps_${N}_quant_input_c
      ::hps_bytecode_module_static_c
      ::hps_bytecode_module_static_lib
      iree::vm::bytecode::module
      model_util::util_static
    COPTS
      ${HPS_COPTS}
    LINKOPTS
      "LINKER:--defsym=__stack_size__=200k"
      # TODO(b/241297921): determine HPS heap size
      "LINKER:--defsym=__heap_size__=5M"
  )

  sparrow_test(
    NAME
      hps_${N}_emitc_static
    SRCS
      "iree_exec/hps.c"
    DEPS
      ::hps_${N}_quant_input_c
      ::hps_c_module_static_emitc
      ::hps_c_module_static_lib
      model_util::util_static
    LINKOPTS
      "LINKER:--defsym=__stack_size__=200k"
      # TODO(b/241297921): determine HPS heap size
      "LINKER:--defsym=__heap_size__=5M"
    COPTS
      ${HPS_COPTS}
      "-DBUILD_EMITC"
    TESTFILES
      "iree_exec/hps_test.run"
      "iree_exec/hps_${N}_test.filecheck"
  )
endforeach()

iree_cc_binary(
  NAME
    mobilenet_v2_bytecode_static
  SRCS
    "iree_exec/mobilenet_v2.c"
  DEPS
    ::mobilenet_quant_input_c
    ::mobilenet_v2_bytecode_module_static_c
    ::mobilenet_v2_bytecode_module_static_lib
    iree::vm::bytecode::module
    model_util::util_static
  LINKOPTS
    "LINKER:--defsym=__stack_size__=300k"
    "LINKER:--defsym=__heap_size__=2m"
)

sparrow_test(
  NAME
    mobilenet_v2_emitc_static
  SRCS
    "iree_exec/mobilenet_v2.c"
  DEPS
    ::mobilenet_quant_input_c
    ::mobilenet_v2_c_module_static_emitc
    ::mobilenet_v2_c_module_static_lib
    model_util::util_static
  LINKOPTS
    "LINKER:--defsym=__stack_size__=300k"
    "LINKER:--defsym=__heap_size__=2m"
  COPTS
    "-DBUILD_EMITC"
  TESTFILES
    "iree_exec/mobilenet_v2_test.run"
    "iree_exec/mobilenet_v2_test.filecheck"
)

iree_cc_binary(
  NAME
    person_detection_bytecode_vmvx
  SRCS
    "iree_exec/person_detection.c"
  DEPS
    ::person_detection_bytecode_module_vmvx_c
    ::person_detection_quant_input_c
    model_util::util_vmvx
  LINKOPTS
    "LINKER:--defsym=__stack_size__=128k"
    "LINKER:--defsym=__heap_size__=200k"
  COPTS
    "-DBUILD_VMVX"
)

sparrow_test(
  NAME
    person_detection_emitc_vmvx
  SRCS
    "iree_exec/person_detection.c"
  DEPS
    ::person_detection_c_module_vmvx_emitc
    ::person_detection_quant_input_c
    model_util::util_vmvx
  LINKOPTS
    "LINKER:--defsym=__stack_size__=128k"
    "LINKER:--defsym=__heap_size__=200k"
  COPTS
    "-DBUILD_EMITC"
    "-DBUILD_VMVX"
  TESTFILES
    "iree_exec/person_detection_vmvx_test.run"
    "iree_exec/person_detection_test.filecheck"
)