tools/iree-run-module-main.cc - 3p/openxla/iree - Git at Google


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


            #include <array>

            
            #include <cstdio>

            
            #include <cstdlib>

            
            #include <iterator>

            
            #include <string>

            
            #include <utility>

            
            #include <vector>


            #include "iree/base/api.h"

            
            #include "iree/base/internal/flags.h"

            
            #include "iree/base/tracing.h"

            
            #include "iree/hal/api.h"

            
            #include "iree/modules/hal/types.h"

            
            #include "iree/tooling/comparison.h"

            
            #include "iree/tooling/context_util.h"

            
            #include "iree/tooling/device_util.h"

            
            #include "iree/tooling/vm_util.h"

            
            #include "iree/vm/api.h"


            IREE_FLAG(string, function, "",

            
                      "Name of a function contained in the module specified by --module= "

            
                      "to run.");


            IREE_FLAG(bool, print_statistics, false,

            
                      "Prints runtime statistics to stderr on exit.");


            IREE_FLAG_LIST(

            
                string, input,

            
                "An input (a) value or (b) buffer of the format:\n"

            
                "  (a) scalar value\n"

            
                "     value\n"

            
                "     e.g.: --input=\"3.14\"\n"

            
                "  (b) buffer:\n"

            
                "     [shape]xtype=[value]\n"

            
                "     e.g.: --input=\"2x2xi32=1 2 3 4\"\n"

            
                "Optionally, brackets may be used to separate the element values:\n"

            
                "  2x2xi32=[[1 2][3 4]]\n"

            
                "Raw binary files can be read to provide buffer contents:\n"

            
                "  2x2xi32=@some/file.bin\n"

            
                "\n"

            
                "Numpy npy files from numpy.save can be read to provide 1+ values:\n"

            
                "  @some.npy\n"

            
                "\n"

            
                "Each occurrence of the flag indicates an input in the order they were\n"

            
                "specified on the command line.");


            IREE_FLAG_LIST(

            
                string, output,

            
                "Specifies how to handle an output from the invocation:\n"

            
                "  `` (empty): ignore output\n"

            
                "     e.g.: --output=\n"

            
                "  `-`: print textual form to stdout\n"

            
                "     e.g.: --output=-\n"

            
                "  `@file.npy`: create/overwrite a numpy npy file and write buffer view\n"

            
                "     e.g.: --output=@file.npy\n"

            
                "  `+file.npy`: create/append a numpy npy file and write buffer view\n"

            
                "     e.g.: --output=+file.npy\n"

            
                "\n"

            
                "Numpy npy files can be read in Python using numpy.load, for example an\n"

            
                "invocation producing two outputs can be concatenated as:\n"

            
                "    --output=@file.npy --output=+file.npy\n"

            
                "And then loaded in Python by reading from the same file:\n"

            
                "  with open('file.npy', 'rb') as f:\n"

            
                "    print(numpy.load(f))\n"

            
                "    print(numpy.load(f))\n"

            
                "\n"

            
                "Each occurrence of the flag indicates an output in the order they were\n"

            
                "specified on the command line.");


            IREE_FLAG_LIST(string, expected_output,

            
                           "An expected function output following the same format as "

            
                           "--input. When present the results of the "

            
                           "invocation will be compared against these values and the "

            
                           "tool will return non-zero if any differ. If the value of a "

            
                           "particular output is not of interest provide `(ignored)`.");


            IREE_FLAG(int32_t, output_max_element_count, 1024,

            
                      "Prints up to the maximum number of elements of output tensors, "

            
                      "eliding the remainder.");


            namespace iree {

            
            namespace {


            iree_status_t Run(int* out_exit_code) {

            
              IREE_TRACE_SCOPE0("iree-run-module");


              iree_allocator_t host_allocator = iree_allocator_system();

            
              vm::ref<iree_vm_instance_t> instance;

            
              IREE_RETURN_IF_ERROR(iree_tooling_create_instance(host_allocator, &instance),

            
                                   "creating instance");


              vm::ref<iree_vm_module_t> main_module;

            
              IREE_RETURN_IF_ERROR(iree_tooling_load_module_from_flags(

            
                  instance.get(), host_allocator, &main_module));


              vm::ref<iree_vm_context_t> context;

            
              vm::ref<iree_hal_device_t> device;

            
              vm::ref<iree_hal_allocator_t> device_allocator;

            
              IREE_RETURN_IF_ERROR(iree_tooling_create_context_from_flags(

            
                  instance.get(), /*user_module_count=*/1, /*user_modules=*/&main_module,

            
                  /*default_device_uri=*/iree_string_view_empty(), host_allocator, &context,

            
                  &device, &device_allocator));


              std::string function_name = std::string(FLAG_function);

            
              iree_vm_function_t function;

            
              if (function_name.empty()) {

            
                return iree_make_status(IREE_STATUS_INVALID_ARGUMENT,

            
                                        "no --function= specified");

            
              } else {

            
                IREE_RETURN_IF_ERROR(

            
                    iree_vm_module_lookup_function_by_name(

            
                        main_module.get(), IREE_VM_FUNCTION_LINKAGE_EXPORT,

            
                        iree_string_view_t{function_name.data(), function_name.size()},

            
                        &function),

            
                    "looking up function '%s'", function_name.c_str());

            
              }


              IREE_RETURN_IF_ERROR(iree_hal_begin_profiling_from_flags(device.get()));


              vm::ref<iree_vm_list_t> inputs;

            
              IREE_RETURN_IF_ERROR(iree_tooling_parse_to_variant_list(

            
                  device_allocator.get(), FLAG_input_list().values, FLAG_input_list().count,

            
                  host_allocator, &inputs));


              // If the function is async add fences so we can invoke it synchronously.

            
              vm::ref<iree_hal_fence_t> finish_fence;

            
              IREE_RETURN_IF_ERROR(iree_tooling_append_async_fence_inputs(

            
                  inputs.get(), &function, device.get(), /*wait_fence=*/NULL,

            
                  &finish_fence));


              vm::ref<iree_vm_list_t> outputs;

            
              IREE_RETURN_IF_ERROR(iree_vm_list_create(/*element_type=*/nullptr, 16,

            
                                                       host_allocator, &outputs));


              printf("EXEC @%s\n", function_name.c_str());

            
              IREE_RETURN_IF_ERROR(

            
                  iree_vm_invoke(context.get(), function, IREE_VM_INVOCATION_FLAG_NONE,

            
                                 /*policy=*/nullptr, inputs.get(), outputs.get(),

            
                                 host_allocator),

            
                  "invoking function '%s'", function_name.c_str());


              // If the function is async we need to wait for it to complete.

            
              if (finish_fence) {

            
                IREE_RETURN_IF_ERROR(

            
                    iree_hal_fence_wait(finish_fence.get(), iree_infinite_timeout()));

            
              }


              IREE_RETURN_IF_ERROR(iree_hal_end_profiling_from_flags(device.get()));


              if (FLAG_expected_output_list().count == 0) {

            
                if (FLAG_output_list().count == 0) {

            
                  IREE_RETURN_IF_ERROR(

            
                      iree_tooling_variant_list_fprint(

            
                          outputs.get(), (iree_host_size_t)FLAG_output_max_element_count,

            
                          stdout),

            
                      "printing results");

            
                } else {

            
                  IREE_RETURN_IF_ERROR(

            
                      iree_tooling_output_variant_list(

            
                          outputs.get(), FLAG_output_list().values,

            
                          FLAG_output_list().count,

            
                          (iree_host_size_t)FLAG_output_max_element_count, stdout),

            
                      "outputting results");

            
                }

            
              } else {

            
                // Parse expected list into host-local memory that we can easily access.

            
                // Note that we return a status here as this can fail on user inputs.

            
                vm::ref<iree_hal_allocator_t> heap_allocator;

            
                IREE_RETURN_IF_ERROR(iree_hal_allocator_create_heap(

            
                    IREE_SV("heap"), host_allocator, host_allocator, &heap_allocator));

            
                vm::ref<iree_vm_list_t> expected_list;

            
                IREE_RETURN_IF_ERROR(iree_tooling_parse_to_variant_list(

            
                    heap_allocator.get(), FLAG_expected_output_list().values,

            
                    FLAG_expected_output_list().count, host_allocator, &expected_list));


                // Compare expected vs actual lists and output diffs.

            
                bool did_match = iree_tooling_compare_variant_lists(

            
                    expected_list.get(), outputs.get(), host_allocator, stdout);

            
                if (did_match) {

            
                  printf("[SUCCESS] all function outputs matched their expected values.\n");

            
                }

            
                *out_exit_code = did_match ? EXIT_SUCCESS : EXIT_FAILURE;

            
              }


              // Release resources before gathering statistics.

            
              inputs.reset();

            
              outputs.reset();

            
              main_module.reset();

            
              context.reset();


              if (device_allocator && FLAG_print_statistics) {

            
                IREE_IGNORE_ERROR(

            
                    iree_hal_allocator_statistics_fprint(stderr, device_allocator.get()));

            
              }


              device_allocator.reset();

            
              device.reset();

            
              instance.reset();

            
              return iree_ok_status();

            
            }


            }  // namespace


            extern "C" int main(int argc, char** argv) {

            
              iree_flags_parse_checked(IREE_FLAGS_PARSE_MODE_DEFAULT, &argc, &argv);

            
              if (argc > 1) {

            
                // Avoid iree-run-module spinning endlessly on stdin if the user uses single

            
                // dashes for flags.

            
                printf(

            
                    "[ERROR] unexpected positional argument (expected none)."

            
                    " Did you use pass a flag with a single dash ('-')?"

            
                    " Use '--' instead.\n");

            
                return 1;

            
              }


              int exit_code = EXIT_SUCCESS;

            
              iree_status_t status = Run(&exit_code);

            
              if (!iree_status_is_ok(status)) {

            
                iree_status_fprint(stderr, status);

            
                iree_status_free(status);

            
                return EXIT_FAILURE;

            
              }


              return exit_code;

            
            }


            }  // namespace iree
	// 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

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

	#include "iree/base/api.h"
	#include "iree/base/internal/flags.h"
	#include "iree/base/tracing.h"
	#include "iree/hal/api.h"
	#include "iree/modules/hal/types.h"
	#include "iree/tooling/comparison.h"
	#include "iree/tooling/context_util.h"
	#include "iree/tooling/device_util.h"
	#include "iree/tooling/vm_util.h"
	#include "iree/vm/api.h"

	IREE_FLAG(string, function, "",
	"Name of a function contained in the module specified by --module= "
	"to run.");

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

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

	IREE_FLAG_LIST(
	string, output,
	"Specifies how to handle an output from the invocation:\n"
	" `` (empty): ignore output\n"
	" e.g.: --output=\n"
	" `-`: print textual form to stdout\n"
	" e.g.: --output=-\n"
	" `@file.npy`: create/overwrite a numpy npy file and write buffer view\n"
	" e.g.: --output=@file.npy\n"
	" `+file.npy`: create/append a numpy npy file and write buffer view\n"
	" e.g.: --output=+file.npy\n"
	"\n"
	"Numpy npy files can be read in Python using numpy.load, for example an\n"
	"invocation producing two outputs can be concatenated as:\n"
	" --output=@file.npy --output=+file.npy\n"
	"And then loaded in Python by reading from the same file:\n"
	" with open('file.npy', 'rb') as f:\n"
	" print(numpy.load(f))\n"
	" print(numpy.load(f))\n"
	"\n"
	"Each occurrence of the flag indicates an output in the order they were\n"
	"specified on the command line.");

	IREE_FLAG_LIST(string, expected_output,
	"An expected function output following the same format as "
	"--input. When present the results of the "
	"invocation will be compared against these values and the "
	"tool will return non-zero if any differ. If the value of a "
	"particular output is not of interest provide `(ignored)`.");

	IREE_FLAG(int32_t, output_max_element_count, 1024,
	"Prints up to the maximum number of elements of output tensors, "
	"eliding the remainder.");

	namespace iree {
	namespace {

	iree_status_t Run(int* out_exit_code) {
	IREE_TRACE_SCOPE0("iree-run-module");

	iree_allocator_t host_allocator = iree_allocator_system();
	vm::ref<iree_vm_instance_t> instance;
	IREE_RETURN_IF_ERROR(iree_tooling_create_instance(host_allocator, &instance),
	"creating instance");

	vm::ref<iree_vm_module_t> main_module;
	IREE_RETURN_IF_ERROR(iree_tooling_load_module_from_flags(
	instance.get(), host_allocator, &main_module));

	vm::ref<iree_vm_context_t> context;
	vm::ref<iree_hal_device_t> device;
	vm::ref<iree_hal_allocator_t> device_allocator;
	IREE_RETURN_IF_ERROR(iree_tooling_create_context_from_flags(
	instance.get(), /user_module_count=/1, /user_modules=/&main_module,
	/default_device_uri=/iree_string_view_empty(), host_allocator, &context,
	&device, &device_allocator));

	std::string function_name = std::string(FLAG_function);
	iree_vm_function_t function;
	if (function_name.empty()) {
	return iree_make_status(IREE_STATUS_INVALID_ARGUMENT,
	"no --function= specified");
	} else {
	IREE_RETURN_IF_ERROR(
	iree_vm_module_lookup_function_by_name(
	main_module.get(), IREE_VM_FUNCTION_LINKAGE_EXPORT,
	iree_string_view_t{function_name.data(), function_name.size()},
	&function),
	"looking up function '%s'", function_name.c_str());
	}

	IREE_RETURN_IF_ERROR(iree_hal_begin_profiling_from_flags(device.get()));

	vm::ref<iree_vm_list_t> inputs;
	IREE_RETURN_IF_ERROR(iree_tooling_parse_to_variant_list(
	device_allocator.get(), FLAG_input_list().values, FLAG_input_list().count,
	host_allocator, &inputs));

	// If the function is async add fences so we can invoke it synchronously.
	vm::ref<iree_hal_fence_t> finish_fence;
	IREE_RETURN_IF_ERROR(iree_tooling_append_async_fence_inputs(
	inputs.get(), &function, device.get(), /wait_fence=/NULL,
	&finish_fence));

	vm::ref<iree_vm_list_t> outputs;
	IREE_RETURN_IF_ERROR(iree_vm_list_create(/element_type=/nullptr, 16,
	host_allocator, &outputs));

	printf("EXEC @%s\n", function_name.c_str());
	IREE_RETURN_IF_ERROR(
	iree_vm_invoke(context.get(), function, IREE_VM_INVOCATION_FLAG_NONE,
	/policy=/nullptr, inputs.get(), outputs.get(),
	host_allocator),
	"invoking function '%s'", function_name.c_str());

	// If the function is async we need to wait for it to complete.
	if (finish_fence) {
	IREE_RETURN_IF_ERROR(
	iree_hal_fence_wait(finish_fence.get(), iree_infinite_timeout()));
	}

	IREE_RETURN_IF_ERROR(iree_hal_end_profiling_from_flags(device.get()));

	if (FLAG_expected_output_list().count == 0) {
	if (FLAG_output_list().count == 0) {
	IREE_RETURN_IF_ERROR(
	iree_tooling_variant_list_fprint(
	outputs.get(), (iree_host_size_t)FLAG_output_max_element_count,
	stdout),
	"printing results");
	} else {
	IREE_RETURN_IF_ERROR(
	iree_tooling_output_variant_list(
	outputs.get(), FLAG_output_list().values,
	FLAG_output_list().count,
	(iree_host_size_t)FLAG_output_max_element_count, stdout),
	"outputting results");
	}
	} else {
	// Parse expected list into host-local memory that we can easily access.
	// Note that we return a status here as this can fail on user inputs.
	vm::ref<iree_hal_allocator_t> heap_allocator;
	IREE_RETURN_IF_ERROR(iree_hal_allocator_create_heap(
	IREE_SV("heap"), host_allocator, host_allocator, &heap_allocator));
	vm::ref<iree_vm_list_t> expected_list;
	IREE_RETURN_IF_ERROR(iree_tooling_parse_to_variant_list(
	heap_allocator.get(), FLAG_expected_output_list().values,
	FLAG_expected_output_list().count, host_allocator, &expected_list));

	// Compare expected vs actual lists and output diffs.
	bool did_match = iree_tooling_compare_variant_lists(
	expected_list.get(), outputs.get(), host_allocator, stdout);
	if (did_match) {
	printf("[SUCCESS] all function outputs matched their expected values.\n");
	}
	*out_exit_code = did_match ? EXIT_SUCCESS : EXIT_FAILURE;
	}

	// Release resources before gathering statistics.
	inputs.reset();
	outputs.reset();
	main_module.reset();
	context.reset();

	if (device_allocator && FLAG_print_statistics) {
	IREE_IGNORE_ERROR(
	iree_hal_allocator_statistics_fprint(stderr, device_allocator.get()));
	}

	device_allocator.reset();
	device.reset();
	instance.reset();
	return iree_ok_status();
	}

	} // namespace

	extern "C" int main(int argc, char** argv) {
	iree_flags_parse_checked(IREE_FLAGS_PARSE_MODE_DEFAULT, &argc, &argv);
	if (argc > 1) {
	// Avoid iree-run-module spinning endlessly on stdin if the user uses single
	// dashes for flags.
	printf(
	"[ERROR] unexpected positional argument (expected none)."
	" Did you use pass a flag with a single dash ('-')?"
	" Use '--' instead.\n");
	return 1;
	}

	int exit_code = EXIT_SUCCESS;
	iree_status_t status = Run(&exit_code);
	if (!iree_status_is_ok(status)) {
	iree_status_fprint(stderr, status);
	iree_status_free(status);
	return EXIT_FAILURE;
	}

	return exit_code;
	}

	} // namespace iree