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

set(IREE_RUNTIME_SOURCE_DIR "${IREE_SOURCE_DIR}/runtime/src/iree")

# use IREE builtin reference ukernel implementation if RVV is off
set(PREFIX_UKERNEL "${IREE_RUNTIME_SOURCE_DIR}/builtins/ukernel")
set(PREFIX_VMVX "${IREE_RUNTIME_SOURCE_DIR}/modules/vmvx")
if (${BUILD_WITH_RVV})
  set(PREFIX_UKERNEL "${CMAKE_CURRENT_SOURCE_DIR}")
  set(PREFIX_VMVX "${CMAKE_CURRENT_SOURCE_DIR}")
elseif (${BUILD_WITH_KELVIN})
  set(PREFIX_UKERNEL "${CMAKE_CURRENT_SOURCE_DIR}")
endif()

iree_cc_library(
  NAME
    ukernel
  HDRS
    "${IREE_RUNTIME_SOURCE_DIR}/builtins/ukernel/api.h"
  SRCS
    "${IREE_RUNTIME_SOURCE_DIR}/builtins/ukernel/fallback.c"
    "${IREE_RUNTIME_SOURCE_DIR}/builtins/ukernel/mmt4d.c"
    "${IREE_RUNTIME_SOURCE_DIR}/builtins/ukernel/pack.c"
    "${IREE_RUNTIME_SOURCE_DIR}/builtins/ukernel/pack_tile.c"
    "${IREE_RUNTIME_SOURCE_DIR}/builtins/ukernel/query_tile_sizes.c"
    "${IREE_RUNTIME_SOURCE_DIR}/builtins/ukernel/unpack.c"
    "${IREE_RUNTIME_SOURCE_DIR}/builtins/ukernel/unpack_tile.c"
    "${PREFIX_UKERNEL}/mmt4d_tile.c"
  DEPS
    iree::builtins::ukernel::internal_headers
  PUBLIC
)

if (${BUILD_WITH_KELVIN})
  target_compile_definitions(ukernel_ukernel PUBLIC BUILD_KELVIN)
endif()

iree_cc_library(
  NAME
    vmvx
  HDRS
    "${IREE_RUNTIME_SOURCE_DIR}/modules/vmvx/module.h"
  TEXTUAL_HDRS
    "${IREE_RUNTIME_SOURCE_DIR}/modules/vmvx/exports.inl"
  SRCS
    "${IREE_RUNTIME_SOURCE_DIR}/modules/vmvx/module.c"
    "${PREFIX_VMVX}/elementwise.c"
  DEFINES
    "IREE_HAVE_VMVX_MODULE"
  DEPS
    ::ukernel
  PUBLIC
)

iree_cc_library(
  NAME
    vmvx_module_loader
  HDRS
    "${IREE_RUNTIME_SOURCE_DIR}/hal/local/loaders/vmvx_module_loader.h"
  SRCS
    "${IREE_RUNTIME_SOURCE_DIR}/hal/local/loaders/vmvx_module_loader.c"
  DEPS
    ::vmvx
    iree::hal::local::executable_loader
    iree::vm::bytecode::module
  DEFINES
    "IREE_HAVE_HAL_EXECUTABLE_LOADER_VMVX_MODULE=1"
  PUBLIC
)

iree_cc_library(
  NAME
    ukernel_importer
  SRCS
    "importer.c"
  DEPS
    ::ukernel
)