/*
 * 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.
 */

// conv 1x1 tosa op test.

#include "model_util/util.h"
#if defined(BUILD_VMVX)
#include "samples/microbenchmarks/conv1x1_test_vmvx_emitc.h"
#else
#include "samples/microbenchmarks/conv1x1_test_c.h"
#include "samples/microbenchmarks/conv1x1_test_emitc.h"
#endif

const MlModel kModel = {
    .num_input = 1,
    .num_input_dim = {4},
    .input_shape = {{1, 112, 112, 8}},
    .input_length = {112 * 112 * 8},
    .input_size_bytes = {sizeof(int8_t)},
    .num_output = 1,
    .output_length = {112 * 112 * 16},
    .output_size_bytes = sizeof(int8_t),
    .hal_element_type = IREE_HAL_ELEMENT_TYPE_SINT_8,
    .entry_func = "module.main",
    .model_name = "conv1x1_quant",
};

iree_status_t create_module(iree_vm_instance_t *instance,
                            iree_vm_module_t **module) {
  return module_create(instance, iree_allocator_system(), module);
}

#if !defined(BUILD_VMVX)
iree_hal_executable_library_query_fn_t library_query(void) {
  return conv1x1_test_linked_llvm_cpu_library_query;
}
#endif

iree_status_t load_input_data(const MlModel *model, void **buffer,
                              iree_const_byte_span_t **byte_span) {
  iree_status_t result = alloc_input_buffer(model, buffer);
  // Populate initial value to effectively 0 w.r.t conv1x1, so the output is the
  // same as the bias.
  memset((int8_t *)buffer[0], -128,
         model->input_length[0] * model->input_size_bytes[0]);
  byte_span[0] = malloc(sizeof(iree_const_byte_span_t));
  *byte_span[0] = iree_make_const_byte_span(
      buffer[0], model->input_size_bytes[0] * model->input_length[0]);
  return result;
}

iree_status_t process_output(const MlModel *model,
                             iree_hal_buffer_mapping_t *buffers,
                             uint32_t *output_length, uint32_t *output_ptr) {
  iree_status_t result = iree_ok_status();
  // Output is ((bias + input_zp) * multiplier) >> shift + output_zp after
  // rescale.
  const int8_t kExpectedOutput[] = {-128, -128, -125, -128, -128, 127,
                                    -107, -128, -128, -128, -128, -105,
                                    74,   127,  -128, 69};
  for (int i = 0; i < model->output_length[0] / sizeof(kExpectedOutput); ++i) {
    for (int j = 0; j < sizeof(kExpectedOutput); ++j) {
      if ((((const int8_t *)buffers[0]
                .contents.data)[i * sizeof(kExpectedOutput) + j]) !=
          kExpectedOutput[j]) {
        result = iree_make_status(IREE_STATUS_UNKNOWN, "result mismatches");
        break;
      }
    }
  }
  return result;
}