// Float simple_mul bytecode loading and input/output processes

#include "iree/base/api.h"
#include "iree/hal/api.h"
#include "samples/util/util.h"

// Compiled module embedded here to avoid file IO:
#include "samples/simple_vec_mul/simple_float_mul_bytecode_module_dylib_c.h"

const MlModel kModel = {
    .num_input = 2,
    .num_input_dim = {1, 1},
    .input_shape = {{1024}, {1024}},
    .input_length = {1024, 1024},
    .input_size_bytes = {sizeof(float), sizeof(float)},
    .num_output = 1,
    .output_length = {1024},
    .output_size_bytes = sizeof(float),
    .hal_element_type = IREE_HAL_ELEMENT_TYPE_FLOAT_32,
    .entry_func = "module.simple_mul",
    .model_name = "simple_float_vec_mul",
};

const iree_const_byte_span_t load_bytecode_module_data() {
  const struct iree_file_toc_t *module_file_toc =
      samples_simple_vec_mul_simple_float_mul_bytecode_module_dylib_create();
  return iree_make_const_byte_span(module_file_toc->data,
                                   module_file_toc->size);
}

iree_status_t load_input_data(const MlModel *model, void **buffer) {
  iree_status_t result = alloc_input_buffer(model, buffer);
  // Populate initial values
  // arg0 = 0, 1/4, 1/2, 3/4... 1023/4
  // arg1 = 0, 1/2, 1, 3/2... 1023/2
  if (iree_status_is_ok(result)) {
    for (int i = 0; i < model->input_length[0]; ++i) {
      ((float *)buffer[0])[i] = i / 4.0f;
      ((float *)buffer[1])[i] = i / 2.0f;
    }
  }
  return result;
}

iree_status_t check_output_data(const MlModel *model,
                                iree_hal_buffer_mapping_t *mapped_memory,
                                int index_output) {
  iree_status_t result = iree_ok_status();
  for (int i = 0; i < mapped_memory->contents.data_length / sizeof(float);
       ++i) {
    if (((const float *)mapped_memory->contents.data)[i] != i * i / 8.0f) {
      result = iree_make_status(IREE_STATUS_UNKNOWN, "result mismatches");
      break;
    }
  }
  return result;
}