tools/iree-benchmark-trace-main.c - 3p/openxla/iree - Git at Google

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

 // DISCLAIMER: this is leaky under error conditions as it's a benchmark tool and
 // not a correctness test.

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

 #include "iree/base/api.h"
 #include "iree/base/internal/flags.h"
 #include "iree/base/internal/path.h"
 #include "iree/hal/api.h"
 #include "iree/testing/benchmark.h"
 #include "iree/tooling/device_util.h"
 #include "iree/tooling/trace_replay.h"
 #include "iree/tooling/yaml_util.h"
 #include "iree/vm/api.h"

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

 IREE_FLAG(int32_t, call_iterations, 1,
           "Number of times to invoke each call in the trace. May break usage "
           "with stateful models.");

 // A benchmark registration for each file to run.
 typedef struct iree_replay_benchmark_registration_t {
   iree_benchmark_def_t benchmark_def;  // Must be first.
   // Used for relative file path lookup when referencing files in the trace.
   iree_string_view_t root_path;
   // Trace file path.
   iree_string_view_t file_path;
   // Shared VM instance. Unowned; callers must retain for the valid lifetime of
   // the registration.
   iree_vm_instance_t* instance;
 } iree_replay_benchmark_registration_t;

 // A parsed call event from the trace file.
 typedef struct iree_replay_benchmark_call_t {
   iree_vm_function_t function;
   iree_vm_list_t* input_list;
   iree_vm_list_t* output_list;
 } iree_replay_benchmark_call_t;

 // A growable list of calls.
 typedef struct iree_replay_benchmark_call_list_t {
   size_t count;
   size_t capacity;
   iree_replay_benchmark_call_t* items;
 } iree_replay_benchmark_call_list_t;

 // Initializes |out_list| with an initial allocation.
 static void iree_replay_benchmark_call_list_initialize(
     iree_replay_benchmark_call_list_t* out_list) {
   out_list->count = 0;
   out_list->capacity = 32;
   out_list->items = (iree_replay_benchmark_call_t*)malloc(
       out_list->capacity * sizeof(*out_list->items));
 }

 // Deinitializes |list| and frees all call resources.
 static void iree_replay_benchmark_call_list_deinitialize(
     iree_replay_benchmark_call_list_t* list) {
   if (!list->items) return;
   for (size_t i = 0; i < list->count; ++i) {
     iree_vm_list_release(list->items[i].input_list);
     iree_vm_list_release(list->items[i].output_list);
   }
   free(list->items);
   memset(list, 0, sizeof(*list));
 }

 // Acquires a call from the back of the call list, growing as needed.
 // The returned pointer is only valid until the next acquire.
 static iree_replay_benchmark_call_t*
 iree_replay_benchmark_call_list_acquire_back(
     iree_replay_benchmark_call_list_t* list) {
   if (list->count >= list->capacity) {
     list->capacity *= 2;
     list->items = (iree_replay_benchmark_call_t*)realloc(
         list->items, list->capacity * sizeof(*list->items));
   }
   return &list->items[list->count++];
 }

 // Processes a call trace event by preparing the inputs and appending it to the
 // provided |call_list|.
 static iree_status_t iree_replay_benchmark_prepare_call(
     iree_trace_replay_t* replay, yaml_document_t* document,
     yaml_node_t* event_node, iree_replay_benchmark_call_list_t* call_list) {
   // Acquire a call in the list.
   iree_replay_benchmark_call_t* call =
       iree_replay_benchmark_call_list_acquire_back(call_list);
   memset(call, 0, sizeof(*call));

   // Prepare the call event state.
   IREE_RETURN_IF_ERROR(iree_trace_replay_event_call_prepare(
       replay, document, event_node, &call->function, &call->input_list));

   // To avoid allocations in the inner benchmark loop we preallocate the outputs
   // here. To avoid a memory leak we'll need to trim it as soon as the call
   // returns as otherwise the list is retained for the lifetime of the
   // registration.
   IREE_RETURN_IF_ERROR(
       iree_vm_list_create(/*element_type=*/NULL, /*initial_capacity=*/8,
                           replay->host_allocator, &call->output_list));

   return iree_ok_status();
 }

 // Processes a trace event by either setting up the |replay| or appending a call
 // to the |call_list|.
 static iree_status_t iree_replay_benchmark_process_event(
     iree_trace_replay_t* replay, yaml_document_t* document,
     yaml_node_t* event_node, iree_replay_benchmark_call_list_t* call_list) {
   if (event_node->type != YAML_MAPPING_NODE) {
     return iree_make_status(IREE_STATUS_INVALID_ARGUMENT,
                             "(%zu): expected mapping node",
                             event_node->start_mark.line);
   }
   yaml_node_t* type_node = NULL;
   IREE_RETURN_IF_ERROR(iree_yaml_mapping_find(
       document, event_node, iree_make_cstring_view("type"), &type_node));
   if (iree_yaml_string_equal(type_node,
                              iree_make_cstring_view("context_load"))) {
     return iree_trace_replay_event_context_load(replay, document, event_node);
   } else if (iree_yaml_string_equal(type_node,
                                     iree_make_cstring_view("module_load"))) {
     return iree_trace_replay_event_module_load(replay, document, event_node);
   } else if (iree_yaml_string_equal(type_node,
                                     iree_make_cstring_view("call"))) {
     return iree_replay_benchmark_prepare_call(replay, document, event_node,
                                               call_list);
   }
   return iree_make_status(
       IREE_STATUS_UNIMPLEMENTED, "(%zu): unhandled type '%.*s'",
       event_node->start_mark.line, (int)type_node->data.scalar.length,
       type_node->data.scalar.value);
 }

 // Loads the trace at |file_path| and configures |replay| for making calls in
 // the |call_list|.
 static iree_status_t iree_replay_benchmark_load_trace(
     iree_string_view_t file_path, iree_trace_replay_t* replay,
     iree_replay_benchmark_call_list_t* call_list) {
   // Open trace YAML file from the given file_path.
   FILE* file = fopen(file_path.data, "rb");
   if (!file) {
     return iree_make_status(iree_status_code_from_errno(errno),
                             "failed to open trace file '%.*s'",
                             (int)file_path.size, file_path.data);
   }

   // One-pass parsing through the file.
   yaml_parser_t parser;
   if (!yaml_parser_initialize(&parser)) {
     return iree_make_status(IREE_STATUS_INTERNAL,
                             "yaml_parser_initialize failed");
   }
   yaml_parser_set_input_file(&parser, file);

   // Read each event document in the file until EOF.
   iree_status_t status = iree_ok_status();
   for (bool document_eof = false; !document_eof;) {
     yaml_document_t document;
     if (!yaml_parser_load(&parser, &document)) {
       status = iree_status_from_yaml_parser_error(&parser);
       break;
     }
     yaml_node_t* event_node = yaml_document_get_root_node(&document);
     if (event_node) {
       status = iree_replay_benchmark_process_event(replay, &document,
                                                    event_node, call_list);
     } else {
       document_eof = true;
     }
     yaml_document_delete(&document);
     if (!iree_status_is_ok(status)) break;
   }

   yaml_parser_delete(&parser);
   fclose(file);
   return status;
 }

 // Benchmark function that runs a trace file.
 static iree_status_t iree_replay_benchmark_run_file(
     const iree_benchmark_def_t* benchmark_def,
     iree_benchmark_state_t* benchmark_state) {
   const iree_replay_benchmark_registration_t* registration =
       (const iree_replay_benchmark_registration_t*)benchmark_def->user_data;

   // Setup replay state used for this benchmark.
   iree_trace_replay_t replay;
   IREE_RETURN_IF_ERROR(iree_trace_replay_initialize(
       registration->root_path, registration->instance,
       IREE_VM_CONTEXT_FLAG_NONE, iree_hal_available_driver_registry(),
       iree_allocator_system(), &replay));

   // Query device overrides, if any. When omitted the devices from the trace
   // file will be used.
   // TODO(#5724): remove this and instead provide a device set on initialize.
   iree_host_size_t device_uri_count = 0;
   iree_string_view_t* device_uris = NULL;
   iree_hal_get_devices_flag_list(&device_uri_count, &device_uris);
   iree_trace_replay_set_hal_devices_override(&replay, device_uri_count,
                                              device_uris);

   // Load YAML file and setup replay state with all modules loaded and ready.
   iree_replay_benchmark_call_list_t call_list;
   iree_replay_benchmark_call_list_initialize(&call_list);
   IREE_RETURN_IF_ERROR(iree_replay_benchmark_load_trace(registration->file_path,
                                                         &replay, &call_list));

   // Call the functions within the trace in order.
   while (iree_benchmark_keep_running(benchmark_state,
                                      /*batch_count=*/FLAG_call_iterations)) {
     for (size_t i = 0; i < call_list.count; ++i) {
       iree_replay_benchmark_call_t* call = &call_list.items[i];
       for (int32_t j = 0; j < FLAG_call_iterations; ++j) {
         IREE_RETURN_IF_ERROR(iree_vm_invoke(
             replay.context, call->function, IREE_VM_INVOCATION_FLAG_NONE,
             /*policy=*/NULL, call->input_list, call->output_list,
             replay.host_allocator));
         IREE_RETURN_IF_ERROR(iree_vm_list_resize(call->output_list, 0));
       }
     }
   }

   iree_replay_benchmark_call_list_deinitialize(&call_list);
   iree_trace_replay_deinitialize(
       &replay, FLAG_print_statistics
                    ? IREE_TRACE_REPLAY_SHUTDOWN_PRINT_STATISTICS
                    : IREE_TRACE_REPLAY_SHUTDOWN_QUIET);
   return iree_ok_status();
 }

 // Registers benchmarks for each trace file.
 static void iree_replay_benchmark_register_trace_files(
     int file_count, char** file_paths, iree_vm_instance_t* instance) {
   static iree_replay_benchmark_registration_t* registrations = NULL;
   if (!registrations) free(registrations);
   registrations = (iree_replay_benchmark_registration_t*)malloc(
       file_count * sizeof(iree_replay_benchmark_registration_t));

   for (int i = 0; i < file_count; ++i) {
     iree_string_view_t file_path = iree_make_cstring_view(file_paths[i]);
     registrations[i].root_path = iree_file_path_dirname(file_path);
     registrations[i].file_path = file_path;
     registrations[i].instance = instance;
     registrations[i].benchmark_def = (iree_benchmark_def_t){
         .flags = IREE_BENCHMARK_FLAG_MEASURE_PROCESS_CPU_TIME |
                  IREE_BENCHMARK_FLAG_USE_REAL_TIME,
         .time_unit = IREE_BENCHMARK_UNIT_MILLISECOND,
         .minimum_duration_ns = 0,
         .iteration_count = 0,
         .run = iree_replay_benchmark_run_file,
         .user_data = &registrations[i],
     };
     iree_benchmark_register(iree_file_path_stem(file_path),
                             &registrations[i].benchmark_def);
   }
 }

 int main(int argc, char** argv) {
   // Pass through flags to benchmark (allowing --help to fall through).
   iree_flags_parse_checked(IREE_FLAGS_PARSE_MODE_UNDEFINED_OK |
                                IREE_FLAGS_PARSE_MODE_CONTINUE_AFTER_HELP,
                            &argc, &argv);
   iree_benchmark_initialize(&argc, argv);
   if (argc <= 1) {
     fprintf(stderr,
             "no trace files provided; pass one or more yaml file paths");
     return 1;
   }

   // Setup shared instance used for each benchmark.
   iree_vm_instance_t* instance = NULL;
   IREE_CHECK_OK(iree_vm_instance_create(iree_allocator_system(), &instance));

   // Register a benchmark per file provided and run them.
   iree_replay_benchmark_register_trace_files(argc - 1, argv + 1, instance);
   iree_benchmark_run_specified();

   iree_vm_instance_release(instance);
   return 0;
 }
	// Copyright 2021 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

	// DISCLAIMER: this is leaky under error conditions as it's a benchmark tool and
	// not a correctness test.

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

	#include "iree/base/api.h"
	#include "iree/base/internal/flags.h"
	#include "iree/base/internal/path.h"
	#include "iree/hal/api.h"
	#include "iree/testing/benchmark.h"
	#include "iree/tooling/device_util.h"
	#include "iree/tooling/trace_replay.h"
	#include "iree/tooling/yaml_util.h"
	#include "iree/vm/api.h"

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

	IREE_FLAG(int32_t, call_iterations, 1,
	"Number of times to invoke each call in the trace. May break usage "
	"with stateful models.");

	// A benchmark registration for each file to run.
	typedef struct iree_replay_benchmark_registration_t {
	iree_benchmark_def_t benchmark_def; // Must be first.
	// Used for relative file path lookup when referencing files in the trace.
	iree_string_view_t root_path;
	// Trace file path.
	iree_string_view_t file_path;
	// Shared VM instance. Unowned; callers must retain for the valid lifetime of
	// the registration.
	iree_vm_instance_t* instance;
	} iree_replay_benchmark_registration_t;

	// A parsed call event from the trace file.
	typedef struct iree_replay_benchmark_call_t {
	iree_vm_function_t function;
	iree_vm_list_t* input_list;
	iree_vm_list_t* output_list;
	} iree_replay_benchmark_call_t;

	// A growable list of calls.
	typedef struct iree_replay_benchmark_call_list_t {
	size_t count;
	size_t capacity;
	iree_replay_benchmark_call_t* items;
	} iree_replay_benchmark_call_list_t;

	// Initializes \|out_list\| with an initial allocation.
	static void iree_replay_benchmark_call_list_initialize(
	iree_replay_benchmark_call_list_t* out_list) {
	out_list->count = 0;
	out_list->capacity = 32;
	out_list->items = (iree_replay_benchmark_call_t*)malloc(
	out_list->capacity * sizeof(*out_list->items));
	}

	// Deinitializes \|list\| and frees all call resources.
	static void iree_replay_benchmark_call_list_deinitialize(
	iree_replay_benchmark_call_list_t* list) {
	if (!list->items) return;
	for (size_t i = 0; i < list->count; ++i) {
	iree_vm_list_release(list->items[i].input_list);
	iree_vm_list_release(list->items[i].output_list);
	}
	free(list->items);
	memset(list, 0, sizeof(*list));
	}

	// Acquires a call from the back of the call list, growing as needed.
	// The returned pointer is only valid until the next acquire.
	static iree_replay_benchmark_call_t*
	iree_replay_benchmark_call_list_acquire_back(
	iree_replay_benchmark_call_list_t* list) {
	if (list->count >= list->capacity) {
	list->capacity *= 2;
	list->items = (iree_replay_benchmark_call_t*)realloc(
	list->items, list->capacity * sizeof(*list->items));
	}
	return &list->items[list->count++];
	}

	// Processes a call trace event by preparing the inputs and appending it to the
	// provided \|call_list\|.
	static iree_status_t iree_replay_benchmark_prepare_call(
	iree_trace_replay_t* replay, yaml_document_t* document,
	yaml_node_t* event_node, iree_replay_benchmark_call_list_t* call_list) {
	// Acquire a call in the list.
	iree_replay_benchmark_call_t* call =
	iree_replay_benchmark_call_list_acquire_back(call_list);
	memset(call, 0, sizeof(*call));

	// Prepare the call event state.
	IREE_RETURN_IF_ERROR(iree_trace_replay_event_call_prepare(
	replay, document, event_node, &call->function, &call->input_list));

	// To avoid allocations in the inner benchmark loop we preallocate the outputs
	// here. To avoid a memory leak we'll need to trim it as soon as the call
	// returns as otherwise the list is retained for the lifetime of the
	// registration.
	IREE_RETURN_IF_ERROR(
	iree_vm_list_create(/element_type=/NULL, /initial_capacity=/8,
	replay->host_allocator, &call->output_list));

	return iree_ok_status();
	}

	// Processes a trace event by either setting up the \|replay\| or appending a call
	// to the \|call_list\|.
	static iree_status_t iree_replay_benchmark_process_event(
	iree_trace_replay_t* replay, yaml_document_t* document,
	yaml_node_t* event_node, iree_replay_benchmark_call_list_t* call_list) {
	if (event_node->type != YAML_MAPPING_NODE) {
	return iree_make_status(IREE_STATUS_INVALID_ARGUMENT,
	"(%zu): expected mapping node",
	event_node->start_mark.line);
	}
	yaml_node_t* type_node = NULL;
	IREE_RETURN_IF_ERROR(iree_yaml_mapping_find(
	document, event_node, iree_make_cstring_view("type"), &type_node));
	if (iree_yaml_string_equal(type_node,
	iree_make_cstring_view("context_load"))) {
	return iree_trace_replay_event_context_load(replay, document, event_node);
	} else if (iree_yaml_string_equal(type_node,
	iree_make_cstring_view("module_load"))) {
	return iree_trace_replay_event_module_load(replay, document, event_node);
	} else if (iree_yaml_string_equal(type_node,
	iree_make_cstring_view("call"))) {
	return iree_replay_benchmark_prepare_call(replay, document, event_node,
	call_list);
	}
	return iree_make_status(
	IREE_STATUS_UNIMPLEMENTED, "(%zu): unhandled type '%.*s'",
	event_node->start_mark.line, (int)type_node->data.scalar.length,
	type_node->data.scalar.value);
	}

	// Loads the trace at \|file_path\| and configures \|replay\| for making calls in
	// the \|call_list\|.
	static iree_status_t iree_replay_benchmark_load_trace(
	iree_string_view_t file_path, iree_trace_replay_t* replay,
	iree_replay_benchmark_call_list_t* call_list) {
	// Open trace YAML file from the given file_path.
	FILE* file = fopen(file_path.data, "rb");
	if (!file) {
	return iree_make_status(iree_status_code_from_errno(errno),
	"failed to open trace file '%.*s'",
	(int)file_path.size, file_path.data);
	}

	// One-pass parsing through the file.
	yaml_parser_t parser;
	if (!yaml_parser_initialize(&parser)) {
	return iree_make_status(IREE_STATUS_INTERNAL,
	"yaml_parser_initialize failed");
	}
	yaml_parser_set_input_file(&parser, file);

	// Read each event document in the file until EOF.
	iree_status_t status = iree_ok_status();
	for (bool document_eof = false; !document_eof;) {
	yaml_document_t document;
	if (!yaml_parser_load(&parser, &document)) {
	status = iree_status_from_yaml_parser_error(&parser);
	break;
	}
	yaml_node_t* event_node = yaml_document_get_root_node(&document);
	if (event_node) {
	status = iree_replay_benchmark_process_event(replay, &document,
	event_node, call_list);
	} else {
	document_eof = true;
	}
	yaml_document_delete(&document);
	if (!iree_status_is_ok(status)) break;
	}

	yaml_parser_delete(&parser);
	fclose(file);
	return status;
	}

	// Benchmark function that runs a trace file.
	static iree_status_t iree_replay_benchmark_run_file(
	const iree_benchmark_def_t* benchmark_def,
	iree_benchmark_state_t* benchmark_state) {
	const iree_replay_benchmark_registration_t* registration =
	(const iree_replay_benchmark_registration_t*)benchmark_def->user_data;

	// Setup replay state used for this benchmark.
	iree_trace_replay_t replay;
	IREE_RETURN_IF_ERROR(iree_trace_replay_initialize(
	registration->root_path, registration->instance,
	IREE_VM_CONTEXT_FLAG_NONE, iree_hal_available_driver_registry(),
	iree_allocator_system(), &replay));

	// Query device overrides, if any. When omitted the devices from the trace
	// file will be used.
	// TODO(#5724): remove this and instead provide a device set on initialize.
	iree_host_size_t device_uri_count = 0;
	iree_string_view_t* device_uris = NULL;
	iree_hal_get_devices_flag_list(&device_uri_count, &device_uris);
	iree_trace_replay_set_hal_devices_override(&replay, device_uri_count,
	device_uris);

	// Load YAML file and setup replay state with all modules loaded and ready.
	iree_replay_benchmark_call_list_t call_list;
	iree_replay_benchmark_call_list_initialize(&call_list);
	IREE_RETURN_IF_ERROR(iree_replay_benchmark_load_trace(registration->file_path,
	&replay, &call_list));

	// Call the functions within the trace in order.
	while (iree_benchmark_keep_running(benchmark_state,
	/batch_count=/FLAG_call_iterations)) {
	for (size_t i = 0; i < call_list.count; ++i) {
	iree_replay_benchmark_call_t* call = &call_list.items[i];
	for (int32_t j = 0; j < FLAG_call_iterations; ++j) {
	IREE_RETURN_IF_ERROR(iree_vm_invoke(
	replay.context, call->function, IREE_VM_INVOCATION_FLAG_NONE,
	/policy=/NULL, call->input_list, call->output_list,
	replay.host_allocator));
	IREE_RETURN_IF_ERROR(iree_vm_list_resize(call->output_list, 0));
	}
	}
	}

	iree_replay_benchmark_call_list_deinitialize(&call_list);
	iree_trace_replay_deinitialize(
	&replay, FLAG_print_statistics
	? IREE_TRACE_REPLAY_SHUTDOWN_PRINT_STATISTICS
	: IREE_TRACE_REPLAY_SHUTDOWN_QUIET);
	return iree_ok_status();
	}

	// Registers benchmarks for each trace file.
	static void iree_replay_benchmark_register_trace_files(
	int file_count, char** file_paths, iree_vm_instance_t* instance) {
	static iree_replay_benchmark_registration_t* registrations = NULL;
	if (!registrations) free(registrations);
	registrations = (iree_replay_benchmark_registration_t*)malloc(
	file_count * sizeof(iree_replay_benchmark_registration_t));

	for (int i = 0; i < file_count; ++i) {
	iree_string_view_t file_path = iree_make_cstring_view(file_paths[i]);
	registrations[i].root_path = iree_file_path_dirname(file_path);
	registrations[i].file_path = file_path;
	registrations[i].instance = instance;
	registrations[i].benchmark_def = (iree_benchmark_def_t){
	.flags = IREE_BENCHMARK_FLAG_MEASURE_PROCESS_CPU_TIME \|
	IREE_BENCHMARK_FLAG_USE_REAL_TIME,
	.time_unit = IREE_BENCHMARK_UNIT_MILLISECOND,
	.minimum_duration_ns = 0,
	.iteration_count = 0,
	.run = iree_replay_benchmark_run_file,
	.user_data = &registrations[i],
	};
	iree_benchmark_register(iree_file_path_stem(file_path),
	&registrations[i].benchmark_def);
	}
	}

	int main(int argc, char** argv) {
	// Pass through flags to benchmark (allowing --help to fall through).
	iree_flags_parse_checked(IREE_FLAGS_PARSE_MODE_UNDEFINED_OK \|
	IREE_FLAGS_PARSE_MODE_CONTINUE_AFTER_HELP,
	&argc, &argv);
	iree_benchmark_initialize(&argc, argv);
	if (argc <= 1) {
	fprintf(stderr,
	"no trace files provided; pass one or more yaml file paths");
	return 1;
	}

	// Setup shared instance used for each benchmark.
	iree_vm_instance_t* instance = NULL;
	IREE_CHECK_OK(iree_vm_instance_create(iree_allocator_system(), &instance));

	// Register a benchmark per file provided and run them.
	iree_replay_benchmark_register_trace_files(argc - 1, argv + 1, instance);
	iree_benchmark_run_specified();

	iree_vm_instance_release(instance);
	return 0;
	}