// Copyright 2020 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

//===----------------------------------------------------------------------===//
// iree-benchmark-module: benchmarks public functions in an IREE VM module
//===----------------------------------------------------------------------===//
//
// This runs exported functions using flags specified on the command line.
// Each function is measured independently and the numbers reported will be for
// the full end-to-end CPU and wall times.
//
// From an ML perspective this is an integration benchmark for measuring total
// user-visible latency of model entry points. It is *not* a microbenchmarking
// tool for individual device-side dispatch functions (aka ops aka kernels).
// If interested in the precise time of a particular dispatch then tracy,
// executable_library_benchmark, and platform/vendor tooling (nsight, perf, etc)
// are to be used instead and attaching them to this tool is often useful in
// order to get a large sample set.
//
// By default all functions taking no inputs will be benchmarked. If a function
// takes inputs then the user will need to specify them using --input=
// flags. Depending on the input program the -iree-flow-export-benchmark-funcs
// flag can be passed to the compiler to attempt to wrap each function with
// dummy inputs however this will fail in programs with dynamically shaped
// inputs. The workaround for avoiding the need for flags is to provide the
// input program in a form with no inputs from the start.
//
// It's important to remember that IREE is not a BLAS library and is meant to
// run entire programs. It's not generally appropriate to benchmark a model with
// a single matmul, for example, as that's just treating IREE as a BLAS library.
// Note also that user-level ops in a frontend environment don't map to the
// dispatches that IREE executes: IREE is a compiler like any other and does not
// guarantee a source line of code translates into an atomically divisible and
// independently measurable execution command. In other words don't expect to be
// able to benchmark the cost of a broadcasting elementwise tf.add op within a
// model: by the time we are running the program that's fused itself into a
// single machine instruction operating as part of some other ops.
//
// For coarse dispatch testing and triaging it can still be useful to remove
// some of the overheads introduced by whole-program execution and the compiler
// flag --iree-hal-benchmark-dispatch-repeat-count=N is provided to enable
// batching. Whatever N is chosen must then be passed to this tool via
// --batch_size=N so that the benchmark reporting properly reflects the
// batching. As an example --iree-hal-benchmark-dispatch-repeat-count=32 +
// --batch_size=32 will reduce the overheads by 32x. Think of this as a way to
// control the p value in Amdahl's law representing the amount of time spent in
// dispatches relative to the rest of the program. This isn't representative of
// how the full program will run, though, and YMMV. Always verify timings with
// an appropriate device-specific tool before trusting the more generic and
// higher-level numbers from this tool.

#include <array>
#include <cstdio>
#include <iterator>
#include <string>
#include <type_traits>
#include <utility>
#include <vector>

#include "benchmark/benchmark.h"
#include "iree/base/api.h"
#include "iree/base/tooling/flags.h"
#include "iree/hal/api.h"
#include "iree/modules/hal/types.h"
#include "iree/tooling/context_util.h"
#include "iree/tooling/device_util.h"
#include "iree/tooling/function_io.h"
#include "iree/vm/api.h"

constexpr char kNanosecondsUnitString[] = "ns";
constexpr char kMicrosecondsUnitString[] = "us";
constexpr char kMillisecondsUnitString[] = "ms";

// TODO(hanchung): Extract the batch size using
// iree_vm_function_lookup_attr_by_name.
IREE_FLAG(int32_t, batch_size, 1,
          "Number of invocations per iteration, which for dispatch benchmarks "
          "must match the --iree-hal-benchmark-dispatch-repeat-count value "
          "used during compilation.");
IREE_FLAG(int32_t, batch_concurrency, 1,
          "Number of invocations within a batch that should run concurrently.");

IREE_FLAG(string, function, "",
          "Name of a function contained in the module specified by --module= "
          "to run. If this is not set, all the exported functions will be "
          "benchmarked and they are expected to not have input arguments.");

IREE_FLAG(bool, print_statistics, false,
          "Prints runtime statistics to stderr on exit.");

IREE_FLAG_LIST(
    string, input,
    "An input value or buffer of the format:\n"
    "  [shape]xtype=[value]\n"
    "  --input=\"2x2xi32=1 2 3 4\"\n"
    "Optionally, brackets may be used to separate the element values:\n"
    "  --input=\"2x2xi32=[[1 2][3 4]]\"\n"
    "Raw binary files can be read to provide buffer contents:\n"
    "  --input=2x2xi32=@some/file.bin\n"
    "numpy npy files (from numpy.save) can be read to provide 1+ values:\n"
    "  --input=@some.npy\n"
    "Each occurrence of the flag indicates an input in the order they were\n"
    "specified on the command line.");

static iree_status_t parse_time_unit(iree_string_view_t flag_name,
                                     void* storage, iree_string_view_t value) {
  auto* unit = (std::pair<bool, benchmark::TimeUnit>*)storage;
  auto unit_string = std::string(value.data, value.size);
  if (unit_string == kMillisecondsUnitString) {
    *unit = {true, benchmark::kMillisecond};
    return iree_ok_status();
  } else if (unit_string == kMicrosecondsUnitString) {
    *unit = {true, benchmark::kMicrosecond};
    return iree_ok_status();
  } else if (unit_string == kNanosecondsUnitString) {
    *unit = {true, benchmark::kNanosecond};
    return iree_ok_status();
  }
  return iree_make_status(IREE_STATUS_INVALID_ARGUMENT,
                          "unsupported time unit");
}
static void print_time_unit(iree_string_view_t flag_name, void* storage,
                            FILE* file) {
  auto* unit = (std::pair<bool, benchmark::TimeUnit>*)storage;
  if (!unit->first) {
    return;
  }
  std::string unit_string;
  switch (unit->second) {
    case benchmark::kMillisecond:
      unit_string = kMillisecondsUnitString;
      break;
    case benchmark::kMicrosecond:
      unit_string = kMicrosecondsUnitString;
      break;
    case benchmark::kNanosecond:
      unit_string = kNanosecondsUnitString;
      break;
    default:
      assert(false && "Unexpected time unit.");
  }
  fprintf(file, "--%.*s=\"%s\"\n", (int)flag_name.size, flag_name.data,
          unit_string.c_str());
}
// Time unit to be printed. If the first field is false, each place will use its
// default time unit.
static std::pair<bool, benchmark::TimeUnit> FLAG_time_unit = {
    false, benchmark::kNanosecond};
IREE_FLAG_CALLBACK(
    parse_time_unit, print_time_unit, &FLAG_time_unit, time_unit,
    "The time unit to be printed in the results. Can be 'ms', 'us', or 'ns'.");

namespace iree {
namespace {

static void BenchmarkGenericFunction(const std::string& benchmark_name,
                                     int32_t batch_size,
                                     iree_hal_device_t* device,
                                     iree_vm_context_t* context,
                                     iree_vm_function_t function,
                                     iree_vm_list_t* inputs,
                                     benchmark::State& state) {
  IREE_TRACE_ZONE_BEGIN_NAMED_DYNAMIC(z0, benchmark_name.data(),
                                      benchmark_name.size());
  IREE_TRACE_FRAME_MARK();

  vm::ref<iree_vm_list_t> outputs;
  IREE_CHECK_OK(iree_vm_list_create(iree_vm_make_undefined_type_def(), 16,
                                    iree_allocator_system(), &outputs));

  // Benchmarking loop.
  while (state.KeepRunningBatch(batch_size)) {
    IREE_TRACE_ZONE_BEGIN_NAMED(z1, "BenchmarkIteration");
    IREE_TRACE_FRAME_MARK_NAMED("Iteration");
    IREE_CHECK_OK(iree_vm_invoke(
        context, function, IREE_VM_INVOCATION_FLAG_NONE, /*policy=*/nullptr,
        inputs, outputs.get(), iree_allocator_system()));
    IREE_CHECK_OK(iree_vm_list_resize(outputs.get(), 0));
    IREE_TRACE_ZONE_END(z1);
    if (device) {
      state.PauseTiming();
      IREE_CHECK_OK(iree_hal_device_profiling_flush(device));
      state.ResumeTiming();
    }
  }
  state.SetItemsProcessed(state.iterations());

  IREE_TRACE_ZONE_END(z0);
}

void RegisterGenericBenchmark(const std::string& function_name,
                              iree_hal_device_t* device,
                              iree_vm_context_t* context,
                              iree_vm_function_t function,
                              iree_vm_list_t* inputs) {
  auto benchmark_name = "BM_" + function_name;
  int32_t batch_size = FLAG_batch_size;
  benchmark::RegisterBenchmark(benchmark_name.c_str(),
                               [=](benchmark::State& state) -> void {
                                 BenchmarkGenericFunction(
                                     benchmark_name, batch_size, device,
                                     context, function, inputs, state);
                               })
      // By default only the main thread is included in CPU time. Include all
      // the threads instead.
      ->MeasureProcessCPUTime()
      // To make single and multi-threaded benchmarks more comparable, use the
      // wall time to determine how many iterations to run. See
      // https://github.com/google/benchmark#cpu-timers,
      ->UseRealTime()
      ->Unit(FLAG_time_unit.first ? FLAG_time_unit.second
                                  : benchmark::kMillisecond);
}

// Runs up to |batch_size| pipelined invocations in sequence along with
// concurrency. Example:
//   batch_size=1, concurrency=1:
//     [invocation 0]
//   batch_size=2, concurrency=1:
//     [invocation 0] -> [invocation 1]
//   batch_size=2, concurrency=2:
//     [invocation 0]
//     [invocation 1]
//   batch_size=4, concurrency=2:
//     [invocation 0] -> [invocation 2]
//     [invocation 1] -> [invocation 3]
static void BenchmarkAsyncFunction(
    const std::string& benchmark_name, int32_t batch_size,
    int32_t batch_concurrency, iree_hal_device_t* device,
    iree_vm_context_t* context, iree_vm_function_t function,
    iree_vm_list_t* common_inputs, benchmark::State& state) {
  IREE_TRACE_ZONE_BEGIN_NAMED_DYNAMIC(z0, benchmark_name.data(),
                                      benchmark_name.size());
  IREE_TRACE_FRAME_MARK();
  iree_allocator_t host_allocator = iree_allocator_system();

  // Round up batch size to some multiple of concurrency.
  batch_size = (int32_t)iree_host_align(batch_size, batch_concurrency);

  // Benchmarking loop.
  while (state.KeepRunningBatch(batch_size)) {
    state.PauseTiming();
    IREE_TRACE_ZONE_BEGIN_NAMED(z1, "BenchmarkIteration");
    IREE_TRACE_FRAME_MARK_NAMED("Iteration");

    IREE_TRACE_ZONE_BEGIN_NAMED(z_begin, "PrepareBatch");

    // Each concurrent track of execution gets its own semaphore.
    std::vector<vm::ref<iree_hal_semaphore_t>> timeline_semaphores;
    for (int32_t i = 0; i < batch_concurrency; ++i) {
      vm::ref<iree_hal_semaphore_t> timeline_semaphore;
      IREE_CHECK_OK(iree_hal_semaphore_create(
          device, IREE_HAL_QUEUE_AFFINITY_ANY, 0ull,
          IREE_HAL_SEMAPHORE_FLAG_DEFAULT, &timeline_semaphore));
      timeline_semaphores.push_back(std::move(timeline_semaphore));
    }

    // Preallocate fences and I/O for each invocation.
    // The same inputs are used for each but we need a unique list to hold the
    // unique fences. Each fence represents when the invocation has completed.
    std::vector<vm::ref<iree_hal_fence_t>> invocation_fences;
    std::vector<vm::ref<iree_vm_list_t>> invocation_inputs;
    std::vector<vm::ref<iree_vm_list_t>> invocation_outputs;
    vm::ref<iree_hal_fence_t> completion_fence;
    IREE_CHECK_OK(iree_hal_fence_create(batch_concurrency, host_allocator,
                                        &completion_fence));
    for (int32_t i = 0; i < batch_size / batch_concurrency; ++i) {
      for (int32_t j = 0; j < batch_concurrency; ++j) {
        // Chain each concurrent minibatch to the previous. Note that to start
        // we wait on nothing and begin executing immediately.
        vm::ref<iree_hal_fence_t> wait_fence;
        if (i > 0) {
          wait_fence = vm::retain_ref(
              invocation_fences[(i - 1) * batch_concurrency + j]);
        }
        uint64_t signal_value = i + 1;
        vm::ref<iree_hal_fence_t> signal_fence;
        IREE_CHECK_OK(iree_hal_fence_create_at(timeline_semaphores[j].get(),
                                               signal_value, host_allocator,
                                               &signal_fence));
        invocation_fences.push_back(vm::retain_ref(signal_fence));

        // Join the final minibatch on the completion fence.
        if (i == batch_size / batch_concurrency - 1) {
          IREE_CHECK_OK(iree_hal_fence_insert(completion_fence.get(),
                                              timeline_semaphores[j].get(),
                                              signal_value));
        }

        // Clone common inputs and add the invocation-specific fences.
        vm::ref<iree_vm_list_t> inputs;
        IREE_CHECK_OK(
            iree_vm_list_clone(common_inputs, host_allocator, &inputs));
        IREE_CHECK_OK(iree_vm_list_push_ref_move(inputs.get(), wait_fence));
        IREE_CHECK_OK(iree_vm_list_push_ref_move(inputs.get(), signal_fence));
        invocation_inputs.push_back(std::move(inputs));

        // Setup empty outputs.
        vm::ref<iree_vm_list_t> outputs;
        IREE_CHECK_OK(iree_vm_list_create(iree_vm_make_undefined_type_def(), 16,
                                          host_allocator, &outputs));
        invocation_outputs.push_back(std::move(outputs));
      }
    }

    IREE_TRACE_ZONE_END(z_begin);

    state.ResumeTiming();
    {
      // TODO(benvanik): replace with async invocations. Today if the invocation
      // performs any waits this will block on the initial invoke instead of
      // actually overlapping things.
      for (int32_t i = 0; i < batch_size; ++i) {
        IREE_CHECK_OK(
            iree_vm_invoke(context, function, IREE_VM_INVOCATION_FLAG_NONE,
                           /*policy=*/nullptr, invocation_inputs[i].get(),
                           invocation_outputs[i].get(), host_allocator));
      }
      IREE_CHECK_OK(iree_hal_fence_wait(completion_fence.get(),
                                        iree_infinite_timeout(),
                                        IREE_HAL_WAIT_FLAG_DEFAULT));
    }
    state.PauseTiming();

    IREE_TRACE_ZONE_BEGIN_NAMED(z_end, "CleanupBatch");
    for (int32_t i = 0; i < batch_size; ++i) {
      iree_vm_list_clear(invocation_outputs[i].get());
    }
    invocation_fences.clear();
    invocation_inputs.clear();
    invocation_outputs.clear();
    completion_fence.reset();
    timeline_semaphores.clear();
    IREE_TRACE_ZONE_END(z_end);

    IREE_TRACE_ZONE_END(z1);
    if (device) {
      IREE_CHECK_OK(iree_hal_device_profiling_flush(device));
    }
    state.ResumeTiming();
  }
  state.SetItemsProcessed(state.iterations());

  IREE_TRACE_ZONE_END(z0);
}

void RegisterAsyncBenchmark(const std::string& function_name,
                            iree_hal_device_t* device,
                            iree_vm_context_t* context,
                            iree_vm_function_t function,
                            iree_vm_list_t* inputs) {
  auto benchmark_name = "BM_" + function_name;
  int32_t batch_size = FLAG_batch_size;
  int32_t batch_concurrency = FLAG_batch_concurrency;
  benchmark::RegisterBenchmark(
      benchmark_name.c_str(),
      [=](benchmark::State& state) -> void {
        BenchmarkAsyncFunction(benchmark_name, batch_size, batch_concurrency,
                               device, context, function, inputs, state);
      })
      // By default only the main thread is included in CPU time. Include all
      // the threads instead.
      ->MeasureProcessCPUTime()
      // To make single and multi-threaded benchmarks more comparable, use the
      // wall time to determine how many iterations to run. See
      // https://github.com/google/benchmark#cpu-timers,
      ->UseRealTime()
      ->Unit(FLAG_time_unit.first ? FLAG_time_unit.second
                                  : benchmark::kMillisecond);
}

static void BenchmarkDispatchFunction(const std::string& benchmark_name,
                                      iree_vm_context_t* context,
                                      iree_vm_function_t function,
                                      benchmark::State& state) {
  IREE_TRACE_ZONE_BEGIN_NAMED_DYNAMIC(z0, benchmark_name.data(),
                                      benchmark_name.size());
  IREE_TRACE_FRAME_MARK();

  vm::ref<iree_vm_list_t> inputs;
  IREE_CHECK_OK(iree_vm_list_create(iree_vm_make_undefined_type_def(), 16,
                                    iree_allocator_system(), &inputs));
  iree_vm_value_t batch_size = iree_vm_value_make_i32(FLAG_batch_size);
  IREE_CHECK_OK(iree_vm_list_push_value(inputs.get(), &batch_size));

  vm::ref<iree_vm_list_t> outputs;
  IREE_CHECK_OK(iree_vm_list_create(iree_vm_make_undefined_type_def(), 16,
                                    iree_allocator_system(), &outputs));

  // Benchmarking loop.
  while (state.KeepRunningBatch(FLAG_batch_size)) {
    IREE_TRACE_ZONE_BEGIN_NAMED(z1, "BenchmarkIteration");
    IREE_TRACE_FRAME_MARK_NAMED("Iteration");
    IREE_CHECK_OK(iree_vm_invoke(
        context, function, IREE_VM_INVOCATION_FLAG_NONE, /*policy=*/nullptr,
        inputs.get(), outputs.get(), iree_allocator_system()));
    IREE_CHECK_OK(iree_vm_list_resize(outputs.get(), 0));
    IREE_TRACE_ZONE_END(z1);
  }
  state.SetItemsProcessed(state.iterations());

  IREE_TRACE_ZONE_END(z0);
}

void RegisterDispatchBenchmark(const std::string& function_name,
                               iree_vm_context_t* context,
                               iree_vm_function_t function) {
  auto benchmark_name = "BM_" + function_name;
  benchmark::RegisterBenchmark(
      benchmark_name.c_str(),
      [benchmark_name, context, function](benchmark::State& state) -> void {
        BenchmarkDispatchFunction(benchmark_name, context, function, state);
      })
      // By default only the main thread is included in CPU time. Include all
      // the threads instead.
      ->MeasureProcessCPUTime()
      // To make single and multi-threaded benchmarks more comparable, use the
      // wall time to determine how many iterations to run. See
      // https://github.com/google/benchmark#cpu-timers,
      ->UseRealTime()
      ->Unit(FLAG_time_unit.first ? FLAG_time_unit.second
                                  : benchmark::kMicrosecond);
}

// The lifetime of IREEBenchmark should be as long as
// ::benchmark::RunSpecifiedBenchmarks() where the resources are used during
// benchmarking.
class IREEBenchmark {
 public:
  IREEBenchmark() { iree_tooling_module_list_initialize(&module_list_); }

  ~IREEBenchmark() {
    IREE_TRACE_SCOPE_NAMED("IREEBenchmark::dtor");

    // Order matters. Tear down modules first to release resources.
    inputs_.reset();
    context_.reset();
    iree_tooling_module_list_reset(&module_list_);
    instance_.reset();

    // Tear down device last in order to get accurate statistics.
    if (device_allocator_ && FLAG_print_statistics) {
      IREE_IGNORE_ERROR(iree_hal_allocator_statistics_fprint(
          stderr, device_allocator_.get()));
    }
    device_allocator_.reset();
    device_.reset();
  };

  iree_hal_device_t* device() const { return device_.get(); }

  iree_status_t Register() {
    IREE_TRACE_SCOPE_NAMED("IREEBenchmark::Register");

    if (!instance_ || !device_allocator_ || !context_ || !module_list_.count) {
      IREE_RETURN_IF_ERROR(Init());
    }

    auto function_name = std::string(FLAG_function);
    if (!function_name.empty()) {
      IREE_RETURN_IF_ERROR(RegisterSpecificFunction(function_name));
    } else {
      IREE_RETURN_IF_ERROR(RegisterAllExportedFunctions());
    }
    return iree_ok_status();
  }

 private:
  iree_status_t Init() {
    IREE_TRACE_SCOPE_NAMED("IREEBenchmark::Init");
    IREE_TRACE_FRAME_MARK_BEGIN_NAMED("init");

    iree_allocator_t host_allocator = iree_allocator_system();
    IREE_RETURN_IF_ERROR(
        iree_tooling_create_instance(host_allocator, &instance_));

    IREE_RETURN_IF_ERROR(iree_tooling_load_modules_from_flags(
        instance_.get(), host_allocator, &module_list_));

    IREE_RETURN_IF_ERROR(iree_tooling_create_context_from_flags(
        instance_.get(), module_list_.count, module_list_.values,
        /*default_device_uri=*/iree_string_view_empty(), host_allocator,
        &context_, &device_, &device_allocator_));

    IREE_TRACE_FRAME_MARK_END_NAMED("init");
    return iree_ok_status();
  }

  iree_status_t RegisterSpecificFunction(const std::string& function_name) {
    IREE_TRACE_SCOPE_NAMED("IREEBenchmark::RegisterSpecificFunction");

    iree_vm_module_t* main_module =
        iree_tooling_module_list_back(&module_list_);
    iree_vm_function_t function;
    IREE_RETURN_IF_ERROR(iree_vm_module_lookup_function_by_name(
        main_module, IREE_VM_FUNCTION_LINKAGE_EXPORT,
        iree_string_view_t{function_name.data(),
                           (iree_host_size_t)function_name.size()},
        &function));
    iree_vm_function_signature_t signature =
        iree_vm_function_signature(&function);
    iree_string_view_t arguments_cconv, results_cconv;
    IREE_RETURN_IF_ERROR(iree_vm_function_call_get_cconv_fragments(
        &signature, &arguments_cconv, &results_cconv));

    IREE_CHECK_OK(iree_tooling_parse_variants(
        arguments_cconv, FLAG_input_list(), device_.get(),
        device_allocator_.get(), iree_vm_instance_allocator(instance_.get()),
        &inputs_));

    iree_string_view_t invocation_model = iree_vm_function_lookup_attr_by_name(
        &function, IREE_SV("iree.abi.model"));
    if (iree_string_view_equal(invocation_model, IREE_SV("coarse-fences"))) {
      // Asynchronous invocation.
      iree::RegisterAsyncBenchmark(function_name, device_.get(), context_.get(),
                                   function, inputs_.get());
    } else {
      // Synchronous invocation.
      iree::RegisterGenericBenchmark(function_name, device_.get(),
                                     context_.get(), function, inputs_.get());
    }
    return iree_ok_status();
  }

  iree_status_t RegisterAllExportedFunctions() {
    IREE_TRACE_SCOPE_NAMED("IREEBenchmark::RegisterAllExportedFunctions");
    iree_vm_module_t* main_module =
        iree_tooling_module_list_back(&module_list_);
    iree_vm_module_signature_t signature =
        iree_vm_module_signature(main_module);
    for (iree_host_size_t i = 0; i < signature.export_function_count; ++i) {
      iree_vm_function_t function;
      IREE_RETURN_IF_ERROR(iree_vm_module_lookup_function_by_ordinal(
          main_module, IREE_VM_FUNCTION_LINKAGE_EXPORT, i, &function));
      iree_string_view_t function_name = iree_vm_function_name(&function);

      // We run anything with the 'benchmark' attribute.
      // If the attribute is not present we'll run anything that looks runnable.
      iree_string_view_t benchmark_type = iree_vm_function_lookup_attr_by_name(
          &function, IREE_SV("iree.benchmark"));
      if (iree_string_view_equal(benchmark_type, IREE_SV("dispatch"))) {
        iree::RegisterDispatchBenchmark(
            std::string(function_name.data, function_name.size), context_.get(),
            function);
      } else if (iree_string_view_equal(benchmark_type, IREE_SV("entry"))) {
        iree::RegisterGenericBenchmark(
            std::string(function_name.data, function_name.size), device_.get(),
            context_.get(), function,
            /*inputs=*/nullptr);
      } else {
        // Pick up generic () -> () functions.
        if (iree_string_view_starts_with(function_name,
                                         iree_make_cstring_view("__")) ||
            iree_string_view_find_char(function_name, '$', 0) !=
                IREE_STRING_VIEW_NPOS) {
          // Skip internal or special functions.
          continue;
        }

        // Query function information to determine how to run it.
        iree_vm_function_signature_t signature =
            iree_vm_function_signature(&function);
        iree_host_size_t argument_count = 0;
        iree_host_size_t result_count = 0;
        IREE_RETURN_IF_ERROR(iree_vm_function_call_count_arguments_and_results(
            &signature, &argument_count, &result_count));
        iree_string_view_t invocation_model =
            iree_vm_function_lookup_attr_by_name(&function,
                                                 IREE_SV("iree.abi.model"));
        if (iree_string_view_equal(invocation_model,
                                   IREE_SV("coarse-fences"))) {
          // Asynchronous invocation with coarse fences. Expect just those.
          if (argument_count == 2) {
            // Only functions taking a (wait, signal) fence pair are run.
            iree::RegisterAsyncBenchmark(
                std::string(function_name.data, function_name.size),
                device_.get(), context_.get(), function,
                /*inputs=*/nullptr);
          }
        } else {
          // Basic synchronous invocation.
          if (argument_count == 0) {
            // Only functions with no inputs are run (because we can't pass
            // anything).
            iree::RegisterGenericBenchmark(
                std::string(function_name.data, function_name.size),
                device_.get(), context_.get(), function,
                /*inputs=*/nullptr);
          }
        }
      }
    }
    return iree_ok_status();
  }

  iree::vm::ref<iree_vm_instance_t> instance_;
  iree::vm::ref<iree_vm_context_t> context_;
  iree::vm::ref<iree_hal_device_t> device_;
  iree::vm::ref<iree_hal_allocator_t> device_allocator_;
  iree_tooling_module_list_t module_list_;
  iree::vm::ref<iree_vm_list_t> inputs_;
};
}  // namespace
}  // namespace iree

static int runMain(int argc, char** argv) {
  IREE_TRACE_ZONE_BEGIN_NAMED(z0, "iree-benchmark-module");

  // Pass through flags to benchmark (allowing --help to fall through).
  iree_flags_set_usage(
      "iree-benchmark-module",
      "Benchmarks a function within a compiled IREE module and handles I/O\n"
      "parsing. Modules can be provided by file path (`--module=file.vmfb`)\n"
      "or read from stdin (`--module=-`) and the function to execute\n"
      "matches the original name provided to the compiler\n"
      "(`--function=foo` for `func.func @foo`).\n");
  iree_flags_parse_checked(IREE_FLAGS_PARSE_MODE_UNDEFINED_OK |
                               IREE_FLAGS_PARSE_MODE_CONTINUE_AFTER_HELP,
                           &argc, &argv);
  ::benchmark::Initialize(&argc, argv);

  iree::IREEBenchmark iree_benchmark;
  iree_status_t status = iree_benchmark.Register();
  if (!iree_status_is_ok(status)) {
    int exit_code = static_cast<int>(iree_status_code(status));
    printf("%s\n", iree::Status(std::move(status)).ToString().c_str());
    IREE_TRACE_ZONE_END(z0);
    return exit_code;
  }
  IREE_CHECK_OK(iree_hal_begin_profiling_from_flags(iree_benchmark.device()));
  ::benchmark::RunSpecifiedBenchmarks();
  IREE_CHECK_OK(iree_hal_end_profiling_from_flags(iree_benchmark.device()));

  IREE_TRACE_ZONE_END(z0);
  return EXIT_SUCCESS;
}

int main(int argc, char** argv) {
  IREE_TRACE_APP_ENTER();
  int exit_code = runMain(argc, argv);
  IREE_TRACE_APP_EXIT(exit_code);
  return exit_code;
}