// Copyright 2024 The IREE Authors
//
// Licensed under the Apache License v2.0 with LLVM Exceptions.
// See https://llvm.org/LICENSE.txt for license information.
// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception

#include <math.h>
#include <stdio.h>
#include <stdlib.h>
#include <string.h>

#include "iree/base/api.h"
#include "iree/base/internal/math.h"
#include "iree/base/tooling/flags.h"
#include "iree/hal/api.h"
#include "iree/modules/hal/module.h"
#include "iree/tooling/context_util.h"
#include "iree/tooling/device_util.h"
#include "iree/vm/api.h"
#include "iree/vm/native_module_cc.h"
#include "tools/testing/e2e/test_utils.h"

//===----------------------------------------------------------------------===//
// Reference conv2d (NCHW-FCHW)
//===----------------------------------------------------------------------===//

// Conversion from 4D indices in row major order to 1D index.
static int convert_to_1d_index(iree_hal_dim_t channels, iree_hal_dim_t height,
                               iree_hal_dim_t width, iree_hal_dim_t n,
                               iree_hal_dim_t c, iree_hal_dim_t h,
                               iree_hal_dim_t w) {
  return n * (channels * height * width) + c * (height * width) + h * width + w;
}

// [f16 <= f16 * f16 + f16]
static void reference_conv2d_f16_f16_f16_f16(
    iree_hal_dim_t n_size, iree_hal_dim_t c_size, iree_hal_dim_t h_size,
    iree_hal_dim_t w_size, iree_hal_dim_t f_size, iree_hal_dim_t kh_size,
    iree_hal_dim_t kw_size, iree_hal_dim_t sh_size, iree_hal_dim_t sw_size,
    iree_hal_dim_t dh_size, iree_hal_dim_t dw_size, iree_hal_dim_t oh_size,
    iree_hal_dim_t ow_size, const uint16_t* input_data,
    const uint16_t* kernel_data, const uint16_t* acc_data,
    uint16_t* result_data, iree_hal_dim_t n, iree_hal_dim_t oc,
    iree_hal_dim_t oh, iree_hal_dim_t ow) {
  iree_hal_dim_t out_idx =
      convert_to_1d_index(f_size, oh_size, ow_size, n, oc, oh, ow);

  float acc = acc_data ? iree_math_f16_to_f32(acc_data[out_idx]) : 0.f;

  for (iree_hal_dim_t ic = 0; ic < c_size; ++ic) {
    for (iree_hal_dim_t kh = 0; kh < kh_size; ++kh) {
      for (iree_hal_dim_t kw = 0; kw < kw_size; ++kw) {
        iree_hal_dim_t inp_idx = convert_to_1d_index(
            c_size, h_size, w_size, n, ic, (oh * sh_size + kh * dh_size),
            (ow * sw_size + kw * dw_size));
        iree_hal_dim_t krnl_idx =
            convert_to_1d_index(c_size, kh_size, kw_size, oc, ic, kh, kw);

        acc += iree_math_f16_to_f32(input_data[inp_idx]) *
               iree_math_f16_to_f32(kernel_data[krnl_idx]);
      }
    }
    result_data[out_idx] = iree_math_f32_to_f16(acc);
  }
}

static void reference_conv2d_f32_f32_f32_f32(
    iree_hal_dim_t n_size, iree_hal_dim_t c_size, iree_hal_dim_t h_size,
    iree_hal_dim_t w_size, iree_hal_dim_t f_size, iree_hal_dim_t kh_size,
    iree_hal_dim_t kw_size, iree_hal_dim_t sh_size, iree_hal_dim_t sw_size,
    iree_hal_dim_t dh_size, iree_hal_dim_t dw_size, iree_hal_dim_t oh_size,
    iree_hal_dim_t ow_size, const float* input_data, const float* kernel_data,
    const float* acc_data, float* result_data, iree_hal_dim_t n,
    iree_hal_dim_t oc, iree_hal_dim_t oh, iree_hal_dim_t ow) {
  iree_hal_dim_t out_idx =
      convert_to_1d_index(f_size, oh_size, ow_size, n, oc, oh, ow);

  float acc = acc_data ? acc_data[out_idx] : 0;

  for (iree_hal_dim_t ic = 0; ic < c_size; ++ic) {
    for (iree_hal_dim_t kh = 0; kh < kh_size; ++kh) {
      for (iree_hal_dim_t kw = 0; kw < kw_size; ++kw) {
        iree_hal_dim_t inp_idx = convert_to_1d_index(
            c_size, h_size, w_size, n, ic, (oh * sh_size + kh * dh_size),
            (ow * sw_size + kw * dw_size));
        iree_hal_dim_t krnl_idx =
            convert_to_1d_index(c_size, kh_size, kw_size, oc, ic, kh, kw);

        acc += input_data[inp_idx] * kernel_data[krnl_idx];
      }
    }
    result_data[out_idx] = acc;
  }
}

// Helper for reference_conv2d.
static iree_status_t reference_conv2d_element(
    iree_hal_dim_t n_size, iree_hal_dim_t c_size, iree_hal_dim_t h_size,
    iree_hal_dim_t w_size, iree_hal_dim_t f_size, iree_hal_dim_t kh_size,
    iree_hal_dim_t kw_size, iree_hal_dim_t sh_size, iree_hal_dim_t sw_size,
    iree_hal_dim_t dh_size, iree_hal_dim_t dw_size, iree_hal_dim_t oh_size,
    iree_hal_dim_t ow_size, iree_hal_element_type_t input_type,
    iree_hal_element_type_t kernel_type, iree_hal_element_type_t acc_type,
    void* input_data, void* kernel_data, void* acc_data, void* result_data,
    iree_hal_dim_t n, iree_hal_dim_t oc, iree_hal_dim_t oh, iree_hal_dim_t ow) {
  if (input_type == IREE_HAL_ELEMENT_TYPE_FLOAT_32 &&
      kernel_type == IREE_HAL_ELEMENT_TYPE_FLOAT_32 &&
      acc_type == IREE_HAL_ELEMENT_TYPE_FLOAT_32) {
    reference_conv2d_f32_f32_f32_f32(
        n_size, c_size, h_size, w_size, f_size, kh_size, kw_size, sh_size,
        sw_size, dh_size, dw_size, oh_size, ow_size, (const float*)input_data,
        (const float*)kernel_data, (const float*)acc_data, (float*)result_data,
        n, oc, oh, ow);
  } else if (input_type == IREE_HAL_ELEMENT_TYPE_FLOAT_16 &&
             kernel_type == IREE_HAL_ELEMENT_TYPE_FLOAT_16 &&
             acc_type == IREE_HAL_ELEMENT_TYPE_FLOAT_16) {
    reference_conv2d_f16_f16_f16_f16(
        n_size, c_size, h_size, w_size, f_size, kh_size, kw_size, sh_size,
        sw_size, dh_size, dw_size, oh_size, ow_size,
        (const uint16_t*)input_data, (const uint16_t*)kernel_data,
        (const uint16_t*)acc_data, (uint16_t*)result_data, n, oc, oh, ow);
  } else {
    return iree_make_status(IREE_STATUS_INVALID_ARGUMENT,
                            "unhandled combination of element types in conv2d");
  }
  return iree_ok_status();
}

// Calculate the output shape given the dilation and strides.
static iree_hal_dim_t out_shape_calc(iree_hal_dim_t i_shape,
                                     iree_hal_dim_t k_shape,
                                     iree_hal_dim_t stride,
                                     iree_hal_dim_t dilation) {
  iree_hal_dim_t x = (k_shape - 1) * (dilation - 1);
  x = i_shape - k_shape - x;
  return floor(x / stride) + 1;
}

// Reference conv2d-NCHW-FCHW implementation, used to compare conv2d results
// against.
static iree_status_t reference_conv2d(
    iree_hal_dim_t n_size, iree_hal_dim_t c_size, iree_hal_dim_t h_size,
    iree_hal_dim_t w_size, iree_hal_dim_t f_size, iree_hal_dim_t kh_size,
    iree_hal_dim_t kw_size, iree_hal_dim_t sh_size, iree_hal_dim_t sw_size,
    iree_hal_dim_t dh_size, iree_hal_dim_t dw_size,
    iree_hal_element_type_t input_type, iree_hal_element_type_t kernel_type,
    iree_hal_element_type_t acc_type, iree_byte_span_t input_contents,
    iree_byte_span_t kernel_contents, iree_byte_span_t acc_contents,
    iree_byte_span_t result_contents, int compute_every) {
  IREE_TRACE_ZONE_BEGIN(z0);
  IREE_TRACE_ZONE_APPEND_VALUE_I64(z0, n_size);
  IREE_TRACE_ZONE_APPEND_VALUE_I64(z0, c_size);
  IREE_TRACE_ZONE_APPEND_VALUE_I64(z0, h_size);
  IREE_TRACE_ZONE_APPEND_VALUE_I64(z0, w_size);

  iree_hal_dim_t oh_size = out_shape_calc(h_size, kh_size, sh_size, dh_size);
  iree_hal_dim_t ow_size = out_shape_calc(w_size, kw_size, sw_size, dw_size);

  for (iree_hal_dim_t n = 0; n < n_size; ++n) {
    for (iree_hal_dim_t oc = 0; oc < f_size; ++oc) {
      for (iree_hal_dim_t oh = 0; oh < oh_size; ++oh) {
        for (iree_hal_dim_t ow = 0; ow < ow_size; ++ow) {
          IREE_RETURN_AND_END_ZONE_IF_ERROR(
              z0, reference_conv2d_element(
                      n_size, c_size, h_size, w_size, f_size, kh_size, kw_size,
                      sh_size, sw_size, dh_size, dw_size, oh_size, ow_size,
                      input_type, kernel_type, acc_type, input_contents.data,
                      kernel_contents.data, acc_contents.data,
                      result_contents.data, n, oc, oh, ow));
        }
      }
    }
  }

  IREE_TRACE_ZONE_END(z0);
  return iree_ok_status();
}

//===----------------------------------------------------------------------===//
// Conv2d comparison/logging
//===----------------------------------------------------------------------===//

typedef struct {
  iree_allocator_t host_allocator;
  iree_hal_dim_t n;   // batch dim
  iree_hal_dim_t c;   // input channels
  iree_hal_dim_t h;   // input height
  iree_hal_dim_t w;   // input width
  iree_hal_dim_t f;   // output channels
  iree_hal_dim_t kh;  // kernel height
  iree_hal_dim_t kw;  // kernel width
  iree_hal_dim_t sh;  // stride along height dim
  iree_hal_dim_t sw;  // stride along width dim
  iree_hal_dim_t dh;  // dilation along height dim
  iree_hal_dim_t dw;  // dilation along width dim
  iree_hal_element_type_t input_type;
  iree_hal_element_type_t kernel_type;
  iree_hal_element_type_t acc_type;
  iree_hal_element_type_t result_type;
  iree_byte_span_t input_contents;
  iree_byte_span_t kernel_contents;
  iree_byte_span_t acc_contents;
  iree_byte_span_t actual_contents;
  iree_byte_span_t expected_contents;
} conv2d_results_t;

static void conv2d_results_deinitialize(conv2d_results_t* results);

static iree_status_t conv2d_results_initialize(
    iree_hal_device_t* device, iree_hal_dim_t n_size, iree_hal_dim_t c_size,
    iree_hal_dim_t h_size, iree_hal_dim_t w_size, iree_hal_dim_t f_size,
    iree_hal_dim_t kh_size, iree_hal_dim_t kw_size, iree_hal_dim_t sh_size,
    iree_hal_dim_t sw_size, iree_hal_dim_t dh_size, iree_hal_dim_t dw_size,
    iree_hal_buffer_view_t* input, iree_hal_buffer_view_t* kernel,
    iree_hal_buffer_view_t* acc, iree_hal_buffer_view_t* result,
    iree_allocator_t host_allocator, conv2d_results_t* out_results) {
  IREE_TRACE_ZONE_BEGIN(z0);

  memset(out_results, 0, sizeof(*out_results));
  out_results->host_allocator = host_allocator;

  out_results->n = n_size;
  out_results->c = c_size;
  out_results->h = h_size;
  out_results->w = w_size;
  out_results->f = f_size;
  out_results->kh = kh_size;
  out_results->kw = kw_size;
  out_results->sh = sh_size;
  out_results->sw = sw_size;
  out_results->dh = dh_size;
  out_results->dw = dw_size;

  out_results->input_type = iree_hal_buffer_view_element_type(input);
  out_results->kernel_type = iree_hal_buffer_view_element_type(kernel);
  out_results->acc_type = iree_hal_buffer_view_element_type(acc);
  out_results->result_type = iree_hal_buffer_view_element_type(result);

  iree_hal_buffer_t* input_buffer = iree_hal_buffer_view_buffer(input);
  iree_hal_buffer_t* kernel_buffer = iree_hal_buffer_view_buffer(kernel);
  iree_hal_buffer_t* acc_buffer = acc ? iree_hal_buffer_view_buffer(acc) : NULL;
  iree_hal_buffer_t* result_buffer = iree_hal_buffer_view_buffer(result);

  iree_status_t status = iree_ok_status();

  if (iree_status_is_ok(status)) {
    out_results->input_contents.data_length =
        iree_hal_buffer_byte_length(input_buffer);
    status = iree_allocator_malloc(host_allocator,
                                   out_results->input_contents.data_length,
                                   (void**)&out_results->input_contents.data);
  }
  if (iree_status_is_ok(status)) {
    status = iree_hal_device_transfer_d2h(
        device, input_buffer, 0, out_results->input_contents.data,
        out_results->input_contents.data_length,
        IREE_HAL_TRANSFER_BUFFER_FLAG_DEFAULT, iree_infinite_timeout());
  }

  if (iree_status_is_ok(status)) {
    out_results->kernel_contents.data_length =
        iree_hal_buffer_byte_length(kernel_buffer);
    status = iree_allocator_malloc(host_allocator,
                                   out_results->kernel_contents.data_length,
                                   (void**)&out_results->kernel_contents.data);
  }
  if (iree_status_is_ok(status)) {
    status = iree_hal_device_transfer_d2h(
        device, kernel_buffer, 0, out_results->kernel_contents.data,
        out_results->kernel_contents.data_length,
        IREE_HAL_TRANSFER_BUFFER_FLAG_DEFAULT, iree_infinite_timeout());
  }

  if (acc_buffer) {
    if (iree_status_is_ok(status)) {
      out_results->acc_contents.data_length =
          iree_hal_buffer_byte_length(acc_buffer);
      status = iree_allocator_malloc(host_allocator,
                                     out_results->acc_contents.data_length,
                                     (void**)&out_results->acc_contents.data);
    }
    if (iree_status_is_ok(status)) {
      status = iree_hal_device_transfer_d2h(
          device, acc_buffer, 0, out_results->acc_contents.data,
          out_results->acc_contents.data_length,
          IREE_HAL_TRANSFER_BUFFER_FLAG_DEFAULT, iree_infinite_timeout());
    }
  }

  if (iree_status_is_ok(status)) {
    out_results->actual_contents.data_length =
        iree_hal_buffer_byte_length(result_buffer);
    status = iree_allocator_malloc(host_allocator,
                                   out_results->actual_contents.data_length,
                                   (void**)&out_results->actual_contents.data);
  }
  if (iree_status_is_ok(status)) {
    status = iree_hal_device_transfer_d2h(
        device, result_buffer, 0, out_results->actual_contents.data,
        out_results->actual_contents.data_length,
        IREE_HAL_TRANSFER_BUFFER_FLAG_DEFAULT, iree_infinite_timeout());
  }

  if (iree_status_is_ok(status)) {
    out_results->expected_contents.data_length =
        iree_hal_buffer_byte_length(result_buffer);
    status = iree_allocator_malloc(
        host_allocator, out_results->expected_contents.data_length,
        (void**)&out_results->expected_contents.data);
  }

  if (!iree_status_is_ok(status)) {
    conv2d_results_deinitialize(out_results);
  }
  IREE_TRACE_ZONE_END(z0);
  return status;
}

static void conv2d_results_deinitialize(conv2d_results_t* results) {
  IREE_TRACE_ZONE_BEGIN(z0);

  iree_allocator_free(results->host_allocator, results->input_contents.data);
  iree_allocator_free(results->host_allocator, results->kernel_contents.data);
  if (!iree_byte_span_is_empty(results->acc_contents)) {
    iree_allocator_free(results->host_allocator, results->acc_contents.data);
  }
  iree_allocator_free(results->host_allocator, results->actual_contents.data);
  iree_allocator_free(results->host_allocator, results->expected_contents.data);

  IREE_TRACE_ZONE_END(z0);
}

// Helper for check_conv2d: the actual interesting part once we've
// obtained and validated the {n, f, oh, ow}_size values. On error, the first
// index is returned where the actual and expected value doesn't match. TODO:
// Add detailed logging to |file|.
static iree_status_t check_conv2d_results_impl(FILE* file,
                                               const conv2d_results_t* results,
                                               int check_every) {
  IREE_TRACE_ZONE_BEGIN(z0);

  IREE_RETURN_AND_END_ZONE_IF_ERROR(
      z0, reference_conv2d(results->n, results->c, results->h, results->w,
                           results->f, results->kh, results->kw, results->sh,
                           results->sw, results->dh, results->dw,
                           results->input_type, results->acc_type,
                           results->kernel_type, results->input_contents,
                           results->kernel_contents, results->acc_contents,
                           results->expected_contents, check_every));

  int count = 0;

  iree_hal_dim_t oh_size =
      out_shape_calc(results->h, results->kh, results->sh, results->dh);
  iree_hal_dim_t ow_size =
      out_shape_calc(results->w, results->kw, results->sw, results->dw);

  for (iree_hal_dim_t n = 0; n < results->n; ++n) {
    for (iree_hal_dim_t oc = 0; oc < results->f; ++oc) {
      for (iree_hal_dim_t oh = 0; oh < oh_size; ++oh) {
        for (iree_hal_dim_t ow = 0; ow < ow_size; ++ow) {
          if (++count < check_every) continue;
          count = 0;
          iree_hal_dim_t idx =
              convert_to_1d_index(results->f, oh_size, ow_size, n, oc, oh, ow);
          iree_test_utils_e2e_value_t actual_value =
              iree_test_utils_read_buffer_element(
                  idx, results->result_type, results->actual_contents.data);
          iree_test_utils_e2e_value_t expected_value =
              iree_test_utils_read_buffer_element(
                  idx, results->result_type, results->expected_contents.data);
          if (!iree_test_utils_result_elements_agree(actual_value,
                                                     expected_value)) {
            fprintf(
                file,
                "\n\nerror: the actual and expected result tensors disagree "
                "at n %" PRIdim ", oc %" PRIdim ", oh %" PRIdim ", ow %" PRIdim
                ".\n\n",
                n, oc, oh, ow);
            IREE_TRACE_ZONE_END(z0);
            return iree_make_status(IREE_STATUS_ABORTED);
          }
        }
      }
    }
  }

  IREE_TRACE_ZONE_END(z0);
  return iree_ok_status();
}

// Given an actual conv2d's inputs and output (all host-local), uses a
// reference conv2d implementation on the same inputs to check if the output
// is correct. On error, the first index is returned where the actual and
// expected value doesn't match. TODO: Add detailed logging to |file|.
static iree_status_t check_conv2d_results(FILE* file,
                                          const conv2d_results_t* results) {
  IREE_TRACE_ZONE_BEGIN(z0);
  // TODO: Increase the check every param to reduce the number of comparisons.
  int check_every = 1;
  iree_status_t status = check_conv2d_results_impl(file, results, check_every);
  if (!iree_status_is_ok(status) && check_every > 1) {
    // If we got a failure with check_every>1, that didn't log a useful
    // numerical summary, as most of the reference tensor entries hadn't been
    // computed. Rerun now with check_every=1 to get that numerical logging.
    iree_status_ignore(status);
    status = check_conv2d_results_impl(file, results, 1);
  }
  IREE_TRACE_ZONE_END(z0);
  return status;
}

//===----------------------------------------------------------------------===//
// `conv2d_test` custom module
//===----------------------------------------------------------------------===//
// This uses the C++ wrapper to keep things simple. Though easier to use it's
// got additional overhead/code-size bloat that doesn't matter in a test like
// this. Making a C module builder API that removes the boilerplate there is
// TBD so this file is written in C besides this module so that we can swap it
// back to being pure C in the future.

namespace iree {

class Conv2dTestModuleState final {
 public:
  explicit Conv2dTestModuleState(iree_allocator_t host_allocator)
      : host_allocator_(host_allocator) {}
  ~Conv2dTestModuleState() = default;

  // Fills the destination span with pseudorandom values of the given
  // |element_type|. The given |seed| is passed to the pseudorandom generator.
  // The pseudorandom values are reproducible both across runs and across
  // machines.
  StatusOr<vm::ref<iree_hal_buffer_view_t>> GenerateRandom4dTensor(
      iree_hal_device_t* device, int64_t dim0, int64_t dim1, int64_t dim2,
      int64_t dim3, iree_hal_element_type_t element_type, int32_t seed) {
    iree_hal_dim_t dims[4] = {
        (iree_hal_dim_t)dim0,
        (iree_hal_dim_t)dim1,
        (iree_hal_dim_t)dim2,
        (iree_hal_dim_t)dim3,
    };
    iree_hal_buffer_params_t buffer_params = {0};
    buffer_params.usage = IREE_HAL_BUFFER_USAGE_DEFAULT;
    buffer_params.access = IREE_HAL_MEMORY_ACCESS_ALL;
    buffer_params.type = IREE_HAL_MEMORY_TYPE_OPTIMAL_FOR_DEVICE;
    vm::ref<iree_hal_buffer_view_t> result_view;
    struct callback_state_t {
      iree_hal_element_type_t element_type;
      int32_t seed;
    } callback_state = {
        element_type,
        seed,
    };
    IREE_RETURN_IF_ERROR(iree_hal_buffer_view_generate_buffer(
        device, iree_hal_device_allocator(device), IREE_ARRAYSIZE(dims), dims,
        element_type, IREE_HAL_ENCODING_TYPE_DENSE_ROW_MAJOR, buffer_params,
        +[](iree_hal_buffer_mapping_t* mapping, void* user_data) {
          callback_state_t callback_state = *(callback_state_t*)user_data;
          iree_byte_span_t span = mapping->contents;
          // Generate "uniform" integer-valued numbers in the range [min,
          // max].
          int32_t min = 0;
          int32_t max = 0;
          iree_test_utils_get_min_max_for_element_type(
              callback_state.element_type, &min, &max);
          uint32_t range = (max - min + 1);
          iree_host_size_t element_byte_count =
              iree_hal_element_dense_byte_count(callback_state.element_type);
          uint8_t* data_end = span.data + span.data_length;
          uint32_t state = callback_state.seed;
          for (uint8_t* data = span.data; data < data_end;
               data += element_byte_count) {
            int32_t value =
                (int32_t)iree_test_utils_pseudorandom_range(&state, range) +
                min;
            iree_test_utils_write_element(callback_state.element_type, value,
                                          data);
          }
          return iree_ok_status();
        },
        &callback_state, &result_view));
    return std::move(result_view);
  }

  Status CheckConv2dResults(iree_hal_device_t* device, int64_t n, int64_t c,
                            int64_t h, int64_t w, int64_t f, int64_t kh,
                            int64_t kw, int64_t sh, int64_t sw, int64_t dh,
                            int64_t dw, iree_hal_buffer_view_t* input,
                            iree_hal_buffer_view_t* kernel,
                            iree_hal_buffer_view_t* acc,
                            iree_hal_buffer_view_t* actual_result) {
    conv2d_results_t results = {};
    IREE_RETURN_IF_ERROR(conv2d_results_initialize(
        device, (iree_hal_dim_t)n, (iree_hal_dim_t)c, (iree_hal_dim_t)h,
        (iree_hal_dim_t)w, (iree_hal_dim_t)f, (iree_hal_dim_t)kh,
        (iree_hal_dim_t)kw, (iree_hal_dim_t)sh, (iree_hal_dim_t)sw,
        (iree_hal_dim_t)dh, (iree_hal_dim_t)dw, input, kernel, acc,
        actual_result, host_allocator_, &results));
    iree_status_t status = check_conv2d_results(stderr, &results);
    conv2d_results_deinitialize(&results);
    return status;
  }

 private:
  iree_allocator_t host_allocator_;
};

static const vm::NativeFunction<Conv2dTestModuleState>
    kConv2dTestModuleFunctions[] = {
        vm::MakeNativeFunction("generate_random_tensor",
                               &Conv2dTestModuleState::GenerateRandom4dTensor),
        vm::MakeNativeFunction("check_conv2d_results",
                               &Conv2dTestModuleState::CheckConv2dResults),
};

struct Conv2dTestModule final : public vm::NativeModule<Conv2dTestModuleState> {
  using vm::NativeModule<Conv2dTestModuleState>::NativeModule;
  StatusOr<std::unique_ptr<Conv2dTestModuleState>> CreateState(
      iree_allocator_t host_allocator) override {
    return std::make_unique<Conv2dTestModuleState>(host_allocator);
  }
  StatusOr<std::unique_ptr<Conv2dTestModuleState>> ForkState(
      Conv2dTestModuleState* parent_state,
      iree_allocator_t host_allocator) override {
    return CreateState(host_allocator);
  }
};

}  // namespace iree

static iree_status_t conv2d_test_module_create(iree_vm_instance_t* instance,
                                               iree_allocator_t host_allocator,
                                               iree_vm_module_t** out_module) {
  IREE_ASSERT_ARGUMENT(out_module);
  *out_module = NULL;
  auto module = std::make_unique<iree::Conv2dTestModule>(
      "conv2d_test", /*version=*/0, instance, host_allocator,
      iree::span<const iree::vm::NativeFunction<iree::Conv2dTestModuleState>>(
          iree::kConv2dTestModuleFunctions));
  *out_module = module.release()->interface();
  return iree_ok_status();
}

int main(int argc, char** argv) {
  IREE_TRACE_APP_ENTER();

  iree_flags_parse_checked(IREE_FLAGS_PARSE_MODE_DEFAULT, &argc, &argv);
  if (argc != 1) {
    fprintf(stderr, "use --module= flags to specify the modules to run\n");
    IREE_TRACE_APP_EXIT(EXIT_FAILURE);
    return EXIT_FAILURE;
  }

  // Run the tests. Note that some modules may be compiled for other platforms
  // and not have the required architectures for execution within them - to keep
  // the test runner dumber we gracefully fail those cases by returning success.
  iree_status_t status = iree_test_utils_load_and_run_e2e_tests(
      iree_allocator_system(), conv2d_test_module_create);
  int exit_code = EXIT_SUCCESS;
  if (!iree_status_is_ok(status)) {
    iree_status_fprint(stderr, status);
    bool is_device_unavailable = iree_status_is_not_found(status);
    iree_status_free(status);
    exit_code = is_device_unavailable ? EXIT_SUCCESS : EXIT_FAILURE;
  }

  IREE_TRACE_APP_EXIT(exit_code);
  return exit_code;
}