// Copyright 2022 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 <stdint.h>
#include <stdio.h>
#include <string.h>

// Must be first.
#include "experimental/webgpu/platform/webgpu.h"

// NOTE: include order matters.
#include "experimental/webgpu/buffer.h"
#include "experimental/webgpu/webgpu_device.h"
#include "iree/base/api.h"
#include "iree/hal/api.h"
#include "iree/modules/hal/module.h"
#include "iree/runtime/api.h"
#include "iree/vm/bytecode/module.h"

//===----------------------------------------------------------------------===//
// Public API
//===----------------------------------------------------------------------===//

// Opaque state for the sample, shared between multiple loaded programs.
typedef struct iree_sample_state_t iree_sample_state_t;

// Initializes the sample and returns its state.
iree_sample_state_t* setup_sample();

// Shuts down the sample and frees its state.
// Requires that all programs first be unloaded with |unload_program|.
void cleanup_sample(iree_sample_state_t* sample_state);

// Opaque state for an individual loaded program.
typedef struct iree_program_state_t iree_program_state_t;

// Loads a program into the sample from the provided data.
// Note: this takes ownership of |vmfb_data|.
iree_program_state_t* load_program(iree_sample_state_t* sample_state,
                                   uint8_t* vmfb_data, size_t length);

// Inspects metadata about a loaded program, printing to stdout.
void inspect_program(iree_program_state_t* program_state);

// Unloads a program and frees its state.
void unload_program(iree_program_state_t* program_state);

// Calls a function synchronously.
//
// Returns a semicolon-delimited list of formatted outputs on success or the
// empty string on failure. Note: This is in need of some real API bindings
// that marshal structured data between C <-> JS.
//
// * |function_name| is the fully qualified function name, like 'module.abs'.
// * |inputs| is a semicolon delimited list of VM scalars and buffers, as
//   described in iree/tooling/vm_util and used in IREE's CLI tools.
//   For example, the CLI `--function_input=f32=1 --function_input=f32=2`
//   should be passed here as `f32=1;f32=2`.
// * |iterations| is the number of times to call the function, for benchmarking
const char* call_function(iree_program_state_t* program_state,
                          const char* function_name, const char* inputs,
                          int iterations);

//===----------------------------------------------------------------------===//
// Implementation
//===----------------------------------------------------------------------===//

typedef struct iree_sample_state_t {
  iree_runtime_instance_t* instance;
  iree_hal_device_t* device;
} iree_sample_state_t;

typedef struct iree_program_state_t {
  iree_runtime_session_t* session;
  iree_vm_module_t* module;
} iree_program_state_t;

extern iree_status_t create_device(iree_allocator_t host_allocator,
                                   iree_hal_device_t** out_device);

iree_sample_state_t* setup_sample() {
  iree_sample_state_t* sample_state = NULL;
  iree_status_t status =
      iree_allocator_malloc(iree_allocator_system(),
                            sizeof(iree_sample_state_t), (void**)&sample_state);

  iree_runtime_instance_options_t instance_options;
  iree_runtime_instance_options_initialize(&instance_options);
  // Note: no call to iree_runtime_instance_options_use_all_available_drivers().

  if (iree_status_is_ok(status)) {
    status = iree_runtime_instance_create(
        &instance_options, iree_allocator_system(), &sample_state->instance);
  }

  if (iree_status_is_ok(status)) {
    status = create_device(iree_allocator_system(), &sample_state->device);
  }

  if (!iree_status_is_ok(status)) {
    iree_status_fprint(stderr, status);
    iree_status_free(status);
    cleanup_sample(sample_state);
    return NULL;
  }

  return sample_state;
}

void cleanup_sample(iree_sample_state_t* sample_state) {
  iree_hal_device_release(sample_state->device);
  iree_runtime_instance_release(sample_state->instance);
  free(sample_state);
}

iree_program_state_t* load_program(iree_sample_state_t* sample_state,
                                   uint8_t* vmfb_data, size_t length) {
  iree_program_state_t* program_state = NULL;
  iree_status_t status = iree_allocator_malloc(iree_allocator_system(),
                                               sizeof(iree_program_state_t),
                                               (void**)&program_state);

  iree_runtime_session_options_t session_options;
  iree_runtime_session_options_initialize(&session_options);
  if (iree_status_is_ok(status)) {
    status = iree_runtime_session_create_with_device(
        sample_state->instance, &session_options, sample_state->device,
        iree_runtime_instance_host_allocator(sample_state->instance),
        &program_state->session);
  }

  if (iree_status_is_ok(status)) {
    // Take ownership of the FlatBuffer data so JavaScript doesn't need to
    // explicitly call `Module._free()`.
    status = iree_vm_bytecode_module_create(
        iree_runtime_instance_vm_instance(sample_state->instance),
        iree_make_const_byte_span(vmfb_data, length),
        /*flatbuffer_allocator=*/iree_allocator_system(),
        iree_allocator_system(), &program_state->module);
  } else {
    // Must clean up the FlatBuffer data directly.
    iree_allocator_free(iree_allocator_system(), (void*)vmfb_data);
  }

  if (iree_status_is_ok(status)) {
    status = iree_runtime_session_append_module(program_state->session,
                                                program_state->module);
  }

  if (!iree_status_is_ok(status)) {
    iree_status_fprint(stderr, status);
    iree_status_free(status);
    unload_program(program_state);
    return NULL;
  }

  return program_state;
}

void inspect_program(iree_program_state_t* program_state) {
  fprintf(stdout, "=== program properties ===\n");

  iree_vm_module_t* module = program_state->module;
  iree_string_view_t module_name = iree_vm_module_name(module);
  fprintf(stdout, "  module name: '%.*s'\n", (int)module_name.size,
          module_name.data);

  iree_vm_module_signature_t module_signature =
      iree_vm_module_signature(module);
  fprintf(stdout, "  module signature:\n");
  fprintf(stdout, "    %" PRIhsz " imported functions\n",
          module_signature.import_function_count);
  fprintf(stdout, "    %" PRIhsz " exported functions\n",
          module_signature.export_function_count);
  fprintf(stdout, "    %" PRIhsz " internal functions\n",
          module_signature.internal_function_count);

  fprintf(stdout, "  exported functions:\n");
  for (iree_host_size_t i = 0; i < module_signature.export_function_count;
       ++i) {
    iree_vm_function_t function;
    iree_status_t status = iree_vm_module_lookup_function_by_ordinal(
        module, IREE_VM_FUNCTION_LINKAGE_EXPORT, i, &function);
    if (!iree_status_is_ok(status)) {
      iree_status_fprint(stderr, status);
      iree_status_free(status);
      continue;
    }

    iree_string_view_t function_name = iree_vm_function_name(&function);
    iree_vm_function_signature_t function_signature =
        iree_vm_function_signature(&function);
    iree_string_view_t calling_convention =
        function_signature.calling_convention;
    fprintf(stdout, "    function name: '%.*s', calling convention: %.*s'\n",
            (int)function_name.size, function_name.data,
            (int)calling_convention.size, calling_convention.data);
  }
}

void unload_program(iree_program_state_t* program_state) {
  iree_vm_module_release(program_state->module);
  iree_runtime_session_release(program_state->session);
  free(program_state);
}

static iree_status_t parse_input_into_call(
    iree_runtime_call_t* call, iree_hal_allocator_t* device_allocator,
    iree_string_view_t input) {
  bool has_equal =
      iree_string_view_find_char(input, '=', 0) != IREE_STRING_VIEW_NPOS;
  bool has_x =
      iree_string_view_find_char(input, 'x', 0) != IREE_STRING_VIEW_NPOS;
  if (has_equal || has_x) {
    // Buffer view (either just a shape or a shape=value) or buffer.
    bool is_storage_reference =
        iree_string_view_consume_prefix(&input, iree_make_cstring_view("&"));
    iree_hal_buffer_view_t* buffer_view = NULL;
    IREE_RETURN_IF_ERROR(
        iree_hal_buffer_view_parse(input, device_allocator, &buffer_view),
        "parsing value '%.*s'", (int)input.size, input.data);
    if (is_storage_reference) {
      // Storage buffer reference; just take the storage for the buffer view -
      // it'll still have whatever contents were specified (or 0) but we'll
      // discard the metadata.
      iree_vm_ref_t buffer_ref =
          iree_hal_buffer_retain_ref(iree_hal_buffer_view_buffer(buffer_view));
      iree_hal_buffer_view_release(buffer_view);
      return iree_vm_list_push_ref_move(call->inputs, &buffer_ref);
    } else {
      iree_vm_ref_t buffer_view_ref =
          iree_hal_buffer_view_move_ref(buffer_view);
      return iree_vm_list_push_ref_move(call->inputs, &buffer_view_ref);
    }
  } else {
    // Scalar.
    bool has_dot =
        iree_string_view_find_char(input, '.', 0) != IREE_STRING_VIEW_NPOS;
    iree_vm_value_t val;
    if (has_dot) {
      // Float.
      val = iree_vm_value_make_f32(0.0f);
      if (!iree_string_view_atof(input, &val.f32)) {
        return iree_make_status(IREE_STATUS_INVALID_ARGUMENT,
                                "parsing value '%.*s' as f32", (int)input.size,
                                input.data);
      }
    } else {
      // Integer.
      val = iree_vm_value_make_i32(0);
      if (!iree_string_view_atoi_int32(input, &val.i32)) {
        return iree_make_status(IREE_STATUS_INVALID_ARGUMENT,
                                "parsing value '%.*s' as i32", (int)input.size,
                                input.data);
      }
    }
    return iree_vm_list_push_value(call->inputs, &val);
  }

  return iree_make_status(IREE_STATUS_INVALID_ARGUMENT,
                          "Unhandled function input (unreachable?)");
}

static iree_status_t parse_inputs_into_call(
    iree_runtime_call_t* call, iree_hal_allocator_t* device_allocator,
    iree_string_view_t inputs) {
  if (inputs.size == 0) return iree_ok_status();

  // Inputs are provided in a semicolon-delimited list.
  // Split inputs from the list until no semicolons are left.
  iree_string_view_t remaining_inputs = inputs;
  intptr_t split_index = 0;
  do {
    iree_string_view_t next_input;
    split_index = iree_string_view_split(remaining_inputs, ';', &next_input,
                                         &remaining_inputs);
    IREE_RETURN_IF_ERROR(
        parse_input_into_call(call, device_allocator, next_input));
  } while (split_index != -1);

  return iree_ok_status();
}

typedef struct iree_buffer_map_userdata_t {
  iree_hal_buffer_view_t* source_buffer_view;
  iree_hal_buffer_t* readback_buffer;
} iree_buffer_map_userdata_t;

static void iree_webgpu_mapped_buffer_release(void* user_data,
                                              iree_hal_buffer_t* buffer) {
  WGPUBuffer buffer_handle = (WGPUBuffer)user_data;
  wgpuBufferUnmap(buffer_handle);
}

// TODO(scotttodd): move async mapping into webgpu/buffer.h/.c?
static void buffer_map_sync_callback(WGPUBufferMapAsyncStatus map_status,
                                     void* userdata_ptr) {
  iree_buffer_map_userdata_t* userdata =
      (iree_buffer_map_userdata_t*)userdata_ptr;
  switch (map_status) {
    case WGPUBufferMapAsyncStatus_Success:
      break;
    case WGPUBufferMapAsyncStatus_Error:
      fprintf(stderr, "  buffer_map_sync_callback status: Error\n");
      break;
    case WGPUBufferMapAsyncStatus_DeviceLost:
      fprintf(stderr, "  buffer_map_sync_callback status: DeviceLost\n");
      break;
    case WGPUBufferMapAsyncStatus_Unknown:
    default:
      fprintf(stderr, "  buffer_map_sync_callback status: Unknown\n");
      break;
  }

  if (map_status != WGPUBufferMapAsyncStatus_Success) {
    iree_hal_buffer_view_release(userdata->source_buffer_view);
    iree_hal_buffer_release(userdata->readback_buffer);
    iree_allocator_free(iree_allocator_system(), userdata);
    return;
  }

  iree_status_t status = iree_ok_status();

  // TODO(scotttodd): bubble result(s) up to the caller (async + callback API)

  iree_device_size_t data_offset = iree_hal_buffer_byte_offset(
      iree_hal_buffer_view_buffer(userdata->source_buffer_view));
  iree_device_size_t data_length =
      iree_hal_buffer_view_byte_length(userdata->source_buffer_view);
  WGPUBuffer buffer_handle =
      iree_hal_webgpu_buffer_handle(userdata->readback_buffer);

  // For this sample we want to print arbitrary buffers, which is easiest
  // using the |iree_hal_buffer_view_format| function. Internally, that
  // function requires synchronous buffer mapping, so we'll first wrap the
  // already (async) mapped GPU memory into a heap buffer. In a less general
  // application (or one not requiring pretty logging like this), we could
  // skip a few buffer copies and other data transformations here.

  const void* data_ptr =
      wgpuBufferGetConstMappedRange(buffer_handle, data_offset, data_length);

  iree_hal_buffer_t* heap_buffer = NULL;
  if (iree_status_is_ok(status)) {
    // The buffer we get from WebGPU may not be aligned to 64.
    iree_hal_memory_access_t memory_access =
        IREE_HAL_MEMORY_ACCESS_READ | IREE_HAL_MEMORY_ACCESS_UNALIGNED;
    status = iree_hal_heap_buffer_wrap(
        userdata->readback_buffer->device_allocator,
        IREE_HAL_MEMORY_TYPE_HOST_LOCAL, memory_access,
        IREE_HAL_BUFFER_USAGE_MAPPING, data_length,
        iree_make_byte_span((void*)data_ptr, data_length),
        (iree_hal_buffer_release_callback_t){
            .fn = iree_webgpu_mapped_buffer_release,
            .user_data = buffer_handle,
        },
        &heap_buffer);
  }

  // Copy the original buffer_view, backed by the mapped heap buffer instead.
  iree_hal_buffer_view_t* heap_buffer_view = NULL;
  if (iree_status_is_ok(status)) {
    status = iree_hal_buffer_view_create_like(
        heap_buffer, userdata->source_buffer_view, iree_allocator_system(),
        &heap_buffer_view);
  }

  if (iree_status_is_ok(status)) {
    fprintf(stdout, "Call output:\n");
    status = iree_hal_buffer_view_fprint(stdout, heap_buffer_view,
                                         /*max_element_count=*/4096,
                                         iree_allocator_system());
    fprintf(stdout, "\n");
  }
  iree_hal_buffer_view_release(heap_buffer_view);
  iree_hal_buffer_release(heap_buffer);

  if (!iree_status_is_ok(status)) {
    fprintf(stderr, "buffer_map_sync_callback error:\n");
    iree_status_fprint(stderr, status);
    iree_status_free(status);
  }

  iree_hal_buffer_view_release(userdata->source_buffer_view);
  iree_hal_buffer_release(userdata->readback_buffer);
  iree_allocator_free(iree_allocator_system(), userdata);
}

static iree_status_t print_buffer_view(iree_hal_device_t* device,
                                       iree_hal_buffer_view_t* buffer_view) {
  iree_status_t status = iree_ok_status();

  iree_hal_buffer_t* buffer = iree_hal_buffer_view_buffer(buffer_view);
  iree_device_size_t data_offset = iree_hal_buffer_byte_offset(buffer);
  iree_device_size_t data_length =
      iree_hal_buffer_view_byte_length(buffer_view);

  // ----------------------------------------------
  // Allocate mappable host memory.
  // Note: iree_hal_webgpu_simple_allocator_allocate_buffer only supports
  // CopySrc today, so we'll create the buffer directly with
  // wgpuDeviceCreateBuffer and then wrap it using iree_hal_webgpu_buffer_wrap.
  WGPUBufferDescriptor descriptor = {
      .nextInChain = NULL,
      .label = "IREE_readback",
      .usage = WGPUBufferUsage_MapRead | WGPUBufferUsage_CopyDst,
      .size = data_length,
      .mappedAtCreation = false,
  };
  WGPUBuffer readback_buffer_handle = NULL;
  if (iree_status_is_ok(status)) {
    readback_buffer_handle = wgpuDeviceCreateBuffer(
        iree_hal_webgpu_device_handle(device), &descriptor);
    if (!readback_buffer_handle) {
      status = iree_make_status(IREE_STATUS_RESOURCE_EXHAUSTED,
                                "unable to allocate buffer of size %" PRIdsz,
                                data_length);
    }
  }
  iree_device_size_t target_offset = 0;
  const iree_hal_buffer_params_t target_params = {
      .usage = IREE_HAL_BUFFER_USAGE_TRANSFER | IREE_HAL_BUFFER_USAGE_MAPPING,
      .type =
          IREE_HAL_MEMORY_TYPE_HOST_LOCAL | IREE_HAL_MEMORY_TYPE_DEVICE_VISIBLE,
      .access = IREE_HAL_MEMORY_ACCESS_ALL,
  };
  iree_hal_buffer_t* readback_buffer = NULL;
  if (iree_status_is_ok(status)) {
    status = iree_hal_webgpu_buffer_wrap(
        device, iree_hal_device_allocator(device), target_params.type,
        target_params.access, target_params.usage, data_length,
        /*byte_offset=*/0,
        /*byte_length=*/data_length, readback_buffer_handle,
        iree_allocator_system(), &readback_buffer);
  }
  // ----------------------------------------------

  // ----------------------------------------------
  // Transfer from device memory to mappable host memory.
  const iree_hal_transfer_command_t transfer_command = {
      .type = IREE_HAL_TRANSFER_COMMAND_TYPE_COPY,
      .copy =
          {
              .source_buffer = buffer,
              .source_offset = data_offset,
              .target_buffer = readback_buffer,
              .target_offset = target_offset,
              .length = data_length,
          },
  };
  iree_hal_command_buffer_t* command_buffer = NULL;
  if (iree_status_is_ok(status)) {
    status = iree_hal_create_transfer_command_buffer(
        device, IREE_HAL_COMMAND_BUFFER_MODE_ONE_SHOT,
        IREE_HAL_QUEUE_AFFINITY_ANY, /*transfer_count=*/1, &transfer_command,
        &command_buffer);
  }
  iree_hal_semaphore_t* fence_semaphore = NULL;
  if (iree_status_is_ok(status)) {
    status = iree_hal_semaphore_create(device, 0ull, &fence_semaphore);
  }
  uint64_t signal_value = 1ull;
  if (iree_status_is_ok(status)) {
    iree_hal_semaphore_list_t signal_semaphores = {
        .count = 1,
        .semaphores = &fence_semaphore,
        .payload_values = &signal_value,
    };
    status = iree_hal_device_queue_execute(
        device, IREE_HAL_QUEUE_AFFINITY_ANY, iree_hal_semaphore_list_empty(),
        signal_semaphores, 1, &command_buffer);
  }
  // TODO(scotttodd): Make this async - pass a wait source to iree_loop_wait_one
  if (iree_status_is_ok(status)) {
    status = iree_hal_semaphore_wait(fence_semaphore, signal_value,
                                     iree_infinite_timeout());
  }
  iree_hal_command_buffer_release(command_buffer);
  iree_hal_semaphore_release(fence_semaphore);
  // ----------------------------------------------

  iree_buffer_map_userdata_t* userdata = NULL;
  if (iree_status_is_ok(status)) {
    status = iree_allocator_malloc(iree_allocator_system(),
                                   sizeof(iree_buffer_map_userdata_t),
                                   (void**)&userdata);
    iree_hal_buffer_view_retain(buffer_view);  // Released in the callback.
    userdata->source_buffer_view = buffer_view;
    userdata->readback_buffer = readback_buffer;
  }

  if (iree_status_is_ok(status)) {
    wgpuBufferMapAsync(readback_buffer_handle, WGPUMapMode_Read, /*offset=*/0,
                       /*size=*/data_length, buffer_map_sync_callback,
                       /*userdata=*/userdata);
  }

  return status;
}

static iree_status_t print_outputs_from_call(
    iree_runtime_call_t* call, iree_string_builder_t* outputs_builder) {
  iree_vm_list_t* variants_list = iree_runtime_call_outputs(call);
  for (iree_host_size_t i = 0; i < iree_vm_list_size(variants_list); ++i) {
    iree_vm_variant_t variant = iree_vm_variant_empty();
    IREE_RETURN_IF_ERROR(
        iree_vm_list_get_variant_assign(variants_list, i, &variant),
        "variant %" PRIhsz " not present", i);

    if (iree_vm_variant_is_value(variant)) {
      switch (iree_vm_type_def_as_value(variant.type)) {
        case IREE_VM_VALUE_TYPE_I8: {
          IREE_RETURN_IF_ERROR(iree_string_builder_append_format(
              outputs_builder, "i8=%" PRIi8, variant.i8));
          break;
        }
        case IREE_VM_VALUE_TYPE_I16: {
          IREE_RETURN_IF_ERROR(iree_string_builder_append_format(
              outputs_builder, "i16=%" PRIi16, variant.i16));
          break;
        }
        case IREE_VM_VALUE_TYPE_I32: {
          IREE_RETURN_IF_ERROR(iree_string_builder_append_format(
              outputs_builder, "i32=%" PRIi32, variant.i32));
          break;
        }
        case IREE_VM_VALUE_TYPE_I64: {
          IREE_RETURN_IF_ERROR(iree_string_builder_append_format(
              outputs_builder, "i64=%" PRIi64, variant.i64));
          break;
        }
        case IREE_VM_VALUE_TYPE_F32: {
          IREE_RETURN_IF_ERROR(iree_string_builder_append_format(
              outputs_builder, "f32=%f", variant.f32));
          break;
        }
        case IREE_VM_VALUE_TYPE_F64: {
          IREE_RETURN_IF_ERROR(iree_string_builder_append_format(
              outputs_builder, "f64=%lf", variant.f64));
          break;
        }
        default: {
          IREE_RETURN_IF_ERROR(
              iree_string_builder_append_cstring(outputs_builder, "?"));
          break;
        }
      }
    } else if (iree_vm_variant_is_ref(variant)) {
      if (iree_hal_buffer_view_isa(variant.ref)) {
        iree_hal_buffer_view_t* buffer_view =
            iree_hal_buffer_view_deref(variant.ref);
        // TODO(scotttodd): join async outputs together and return to caller
        iree_hal_device_t* device = iree_runtime_session_device(call->session);
        IREE_RETURN_IF_ERROR(print_buffer_view(device, buffer_view));
      } else {
        IREE_RETURN_IF_ERROR(iree_string_builder_append_cstring(
            outputs_builder, "(no printer)"));
      }
    } else {
      IREE_RETURN_IF_ERROR(
          iree_string_builder_append_cstring(outputs_builder, "(null)"));
    }

    if (i < iree_vm_list_size(variants_list) - 1) {
      IREE_RETURN_IF_ERROR(
          iree_string_builder_append_cstring(outputs_builder, ";"));
    }
  }

  iree_vm_list_resize(variants_list, 0);

  return iree_ok_status();
}

iree_status_t invoke_callback(void* user_data, iree_loop_t loop,
                              iree_status_t status, iree_vm_list_t* outputs) {
  iree_vm_async_invoke_state_t* invoke_state =
      (iree_vm_async_invoke_state_t*)user_data;

  if (!iree_status_is_ok(status)) {
    fprintf(stderr, "iree_vm_async_invoke_callback_fn_t error:\n");
    iree_status_fprint(stderr, status);
    iree_status_free(status);
  }

  iree_vm_list_release(outputs);

  iree_allocator_free(iree_allocator_system(), (void*)invoke_state);
  return iree_ok_status();
}

const char* call_function(iree_program_state_t* program_state,
                          const char* function_name, const char* inputs,
                          int iterations) {
  iree_status_t status = iree_ok_status();

  // Fully qualify the function name. This sample only supports loading one
  // module (i.e. 'program') per session, so we can do this.
  iree_string_builder_t name_builder;
  iree_string_builder_initialize(iree_allocator_system(), &name_builder);
  if (iree_status_is_ok(status)) {
    iree_string_view_t module_name = iree_vm_module_name(program_state->module);
    status = iree_string_builder_append_format(&name_builder, "%.*s.%s",
                                               (int)module_name.size,
                                               module_name.data, function_name);
  }

  iree_runtime_call_t call;
  if (iree_status_is_ok(status)) {
    status = iree_runtime_call_initialize_by_name(
        program_state->session, iree_string_builder_view(&name_builder), &call);
  }
  iree_string_builder_deinitialize(&name_builder);

  if (iree_status_is_ok(status)) {
    status = parse_inputs_into_call(
        &call, iree_runtime_session_device_allocator(program_state->session),
        iree_make_cstring_view(inputs));
  }

  // Note: Timing has ~millisecond precision on the web to mitigate timing /
  // side-channel security threats.
  // https://developer.mozilla.org/en-US/docs/Web/API/Performance/now#reduced_time_precision
  iree_time_t start_time = iree_time_now();

  // TODO(scotttodd): benchmark iterations (somehow with async)

  iree_vm_async_invoke_state_t* invoke_state = NULL;
  if (iree_status_is_ok(status)) {
    status = iree_allocator_malloc(iree_allocator_system(),
                                   sizeof(iree_vm_async_invoke_state_t),
                                   (void**)&invoke_state);
  }
  // TODO(scotttodd): emscripten / browser loop here
  iree_status_t loop_status = iree_ok_status();
  iree_loop_t loop = iree_loop_inline(&loop_status);
  if (iree_status_is_ok(status)) {
    iree_vm_context_t* vm_context = iree_runtime_session_context(call.session);
    iree_vm_function_t vm_function = call.function;
    iree_vm_list_t* inputs = call.inputs;
    iree_vm_list_t* outputs = call.outputs;

    status = iree_vm_async_invoke(loop, invoke_state, vm_context, vm_function,
                                  IREE_VM_INVOCATION_FLAG_NONE, /*policy=*/NULL,
                                  inputs, outputs, iree_allocator_system(),
                                  invoke_callback,
                                  /*user_data=*/invoke_state);
  }

  // TODO(scotttodd): record end time in async callback instead of here
  // TODO(scotttodd): print outputs in async callback instead of here

  iree_time_t end_time = iree_time_now();
  iree_time_t time_elapsed = end_time - start_time;

  iree_string_builder_t outputs_builder;
  iree_string_builder_initialize(iree_allocator_system(), &outputs_builder);

  // Output a JSON object as a string:
  // {
  //   "total_invoke_time_ms": [number],
  //   "outputs": [semicolon delimited list of formatted outputs]
  // }
  if (iree_status_is_ok(status)) {
    status = iree_string_builder_append_format(
        &outputs_builder,
        "{ \"total_invoke_time_ms\": %" PRId64 ", \"outputs\": \"",
        time_elapsed / 1000000);
  }
  if (iree_status_is_ok(status)) {
    status = print_outputs_from_call(&call, &outputs_builder);
  }
  if (iree_status_is_ok(status)) {
    status = iree_string_builder_append_cstring(&outputs_builder, "\"}");
  }

  if (!iree_status_is_ok(status)) {
    iree_string_builder_deinitialize(&outputs_builder);
    iree_status_fprint(stderr, status);
    iree_status_free(status);
    return "";
  }

  // Note: this leaks the buffer. It's up to the caller to free it after use.
  char* outputs_string = strdup(iree_string_builder_buffer(&outputs_builder));
  iree_string_builder_deinitialize(&outputs_builder);
  return outputs_string;
}