# 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
    mobilenet_v1
  SRC
    "mobilenet_v1_0.25_224_float.tflite"
  C_IDENTIFIER
    "float_models_mobilenet_v1"
  FLAGS
    "-iree-input-type=tosa"
  RVV_OFF
)

sparrow_modules(
  NAME
    mnist
  SRC
    "$ENV{ROOTDIR}/toolchain/iree/samples/models/mnist.mlir"
  C_IDENTIFIER
    "float_models_mnist"
  FLAGS
    "-iree-input-type=stablehlo"
  RVV_OFF
)

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

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

iree_model_input(
  NAME
    mnist_input
  SHAPE
    "1, 28, 28, 1"
  SRC
    "https://github.com/google/iree/raw/ \
    1e8d1fa96057c47d1fda918f3a2f71f14b073937/samples/iree_vision_inference/mnist_test.png"
  RANGE
    "0, 1"
)

#-------------------------------------------------------------------------------
# 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.

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

iree_cc_binary(
  NAME
    mobilenet_v1_bytecode_static
  SRCS
    "iree_exec/mobilenet_v1.c"
  DEPS
    ::mobilenet_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__=200k"
)

sparrow_test(
  NAME
    mobilenet_v1_emitc_static
  SRCS
    "iree_exec/mobilenet_v1.c"
  DEPS
    ::mobilenet_input_c
    ::mobilenet_v1_c_module_static_emitc
    ::mobilenet_v1_c_module_static_lib
    model_util::util_static
  LINKOPTS
    "LINKER:--defsym=__stack_size__=200k"
  COPTS
    "-DBUILD_EMITC"
  TESTFILES
    "iree_exec/mobilenet_v1_test.run"
    "iree_exec/mobilenet_v1_test.filecheck"
)

iree_cc_binary(
  NAME
    mnist_bytecode_static
  SRCS
    "iree_exec/mnist.c"
  DEPS
    ::mnist_bytecode_module_static_c
    ::mnist_bytecode_module_static_lib
    ::mnist_input_c
    iree::vm::bytecode::module
    model_util::util_static
  LINKOPTS
    "LINKER:--defsym=__stack_size__=100k"
    "LINKER:--defsym=__heap_size__=512k"
)

sparrow_test(
  NAME
    mnist_emitc_static
  SRCS
    "iree_exec/mnist.c"
  DEPS
    ::mnist_c_module_static_emitc
    ::mnist_c_module_static_lib
    ::mnist_input_c
    model_util::util_static
  LINKOPTS
    "LINKER:--defsym=__stack_size__=100k"
    "LINKER:--defsym=__heap_size__=512k"
  COPTS
    "-DBUILD_EMITC"
  TESTFILES
    "iree_exec/mnist_test.run"
    "iree_exec/mnist_test.filecheck"
)