Fix help output for iree-benchmark-module (#3261)
Benchmark and Absl fight. Right now, benchmark eats the `--help` flag
and so we lose all the help information for iree-benchmark-module.
This is not elegant, but it works and includes tests. A number of better
options are not available because both Absl and Benchmark bake behavior
into the only functions callable through their public API.
Fixes https://github.com/google/iree/issues/3247
diff --git a/iree/tools/BUILD b/iree/tools/BUILD
index 530e835..9ccc7bf 100644
--- a/iree/tools/BUILD
+++ b/iree/tools/BUILD
@@ -40,6 +40,8 @@
deps = [
":vm_util",
"@com_google_absl//absl/flags:flag",
+ "@com_google_absl//absl/flags:parse",
+ "@com_google_absl//absl/flags:usage",
"@com_google_absl//absl/strings",
"@com_google_benchmark//:benchmark",
"//iree/base:init",
diff --git a/iree/tools/CMakeLists.txt b/iree/tools/CMakeLists.txt
index 9a073c6..9d73449 100644
--- a/iree/tools/CMakeLists.txt
+++ b/iree/tools/CMakeLists.txt
@@ -59,6 +59,8 @@
DEPS
::vm_util
absl::flags
+ absl::flags_parse
+ absl::flags_usage
absl::strings
benchmark
iree::base::init
diff --git a/iree/tools/iree-benchmark-module-main.cc b/iree/tools/iree-benchmark-module-main.cc
index 188a170..0c260be 100644
--- a/iree/tools/iree-benchmark-module-main.cc
+++ b/iree/tools/iree-benchmark-module-main.cc
@@ -13,6 +13,8 @@
// limitations under the License.
#include "absl/flags/flag.h"
+#include "absl/flags/internal/parse.h"
+#include "absl/flags/usage.h"
#include "absl/strings/string_view.h"
#include "benchmark/benchmark.h"
#include "iree/base/file_io.h"
@@ -164,8 +166,12 @@
} // namespace
-void RegisterModuleBenchmarks() {
+Status RegisterModuleBenchmarks() {
auto function_name = absl::GetFlag(FLAGS_entry_function);
+ if (function_name.empty()) {
+ return InvalidArgumentErrorBuilder(IREE_LOC)
+ << "Must specify an entry_function";
+ }
auto benchmark_name = "BM_" + function_name;
benchmark::RegisterBenchmark(benchmark_name.c_str(),
[function_name](benchmark::State& state) {
@@ -184,13 +190,47 @@
// significant digits. If we end up wanting precision beyond microseconds,
// we can make this setting configurable with a custom command line flag.
->Unit(benchmark::kMillisecond);
+ return OkStatus();
}
} // namespace iree
int main(int argc, char** argv) {
+ // We have to contend with two flag parsing libraries here: absl's and
+ // benchmark's. To make matters worse, both define the `--help` flag. To
+ // ensure that each is able to parse its own flags, we use an absl "internal"
+ // function (still with public visibility) to parse while ignoring undefined
+ // flags. If it sees `--help` it will exit here, so we include the benchmark
+ // library usage information in the manually-set help output. Then we let
+ // benchmark parse its flags. Finally we call the normal initialization
+ // function to do other IREE initialization including flag parsing with
+ // normal options. Any remaining flags will be unknown and result in an error.
+ absl::SetProgramUsageMessage(
+ "iree-benchmark-module \n"
+ " --input_file=module.vmfb\n"
+ " --entry_function=exported_function_to_benchmark\n"
+ " [--inputs=2xi32=1 2,1x2xf32=2 1 | --inputs_file=file_with_inputs]\n"
+ " [--driver=vmla]\n"
+ "\n\n"
+ " Optional flags from third_party/benchmark/src/benchmark.cc:\n"
+ " [--benchmark_list_tests={true|false}]\n"
+ " [--benchmark_filter=<regex>]\n"
+ " [--benchmark_min_time=<min_time>]\n"
+ " [--benchmark_repetitions=<num_repetitions>]\n"
+ " [--benchmark_report_aggregates_only={true|false}]\n"
+ " [--benchmark_display_aggregates_only={true|false}]\n"
+ " [--benchmark_format=<console|json|csv>]\n"
+ " [--benchmark_out=<filename>]\n"
+ " [--benchmark_out_format=<json|console|csv>]\n"
+ " [--benchmark_color={auto|true|false}]\n"
+ " [--benchmark_counters_tabular={true|false}]\n"
+ " [--v=<verbosity>]\n");
+ absl::flags_internal::ParseCommandLineImpl(
+ argc, argv, absl::flags_internal::ArgvListAction::kRemoveParsedArgs,
+ absl::flags_internal::UsageFlagsAction::kHandleUsage,
+ absl::flags_internal::OnUndefinedFlag::kIgnoreUndefined);
::benchmark::Initialize(&argc, argv);
iree::InitializeEnvironment(&argc, &argv);
- iree::RegisterModuleBenchmarks();
+ IREE_CHECK_OK(iree::RegisterModuleBenchmarks());
::benchmark::RunSpecifiedBenchmarks();
return 0;
}
diff --git a/iree/tools/test/BUILD b/iree/tools/test/BUILD
index 2446837..5e4c337 100644
--- a/iree/tools/test/BUILD
+++ b/iree/tools/test/BUILD
@@ -34,3 +34,13 @@
],
tags = ["hostonly"],
)
+
+iree_lit_test_suite(
+ name = "benchmark_flags",
+ srcs = ["benchmark_flags.txt"],
+ data = [
+ "//iree/tools:IreeFileCheck",
+ "//iree/tools:iree-benchmark-module",
+ ],
+ tags = ["hostonly"],
+)
diff --git a/iree/tools/test/CMakeLists.txt b/iree/tools/test/CMakeLists.txt
index 658aa8e..26bad30 100644
--- a/iree/tools/test/CMakeLists.txt
+++ b/iree/tools/test/CMakeLists.txt
@@ -29,3 +29,15 @@
LABELS
"hostonly"
)
+
+iree_lit_test_suite(
+ NAME
+ benchmark_flags
+ SRCS
+ "benchmark_flags.txt"
+ DATA
+ iree::tools::IreeFileCheck
+ iree::tools::iree-benchmark-module
+ LABELS
+ "hostonly"
+)
diff --git a/iree/tools/test/benchmark_flags.txt b/iree/tools/test/benchmark_flags.txt
new file mode 100644
index 0000000..70f279c
--- /dev/null
+++ b/iree/tools/test/benchmark_flags.txt
@@ -0,0 +1,12 @@
+// HELP: iree-benchmark-module
+// HELP: --input_file
+// HELP: --benchmark_list_tests
+// RUN: ( iree-benchmark-module --help || [[ $? == 1 ]] ) | IreeFileCheck --check-prefix=HELP %s
+// RUN: ( iree-benchmark-module --helpshort || [[ $? == 1 ]] ) | IreeFileCheck --check-prefix=HELP %s
+
+// UNKNOWN: unknown-flag
+// RUN: ( iree-benchmark-module --unknown-flag 2>&1 || [[ $? == 1 ]] ) | IreeFileCheck --check-prefix=UNKNOWN %s
+// RUN: ( iree-benchmark-module --driver=vmla --unknown-flag --benchmark_list_tests 2>&1 || [[ $? == 1 ]] ) | IreeFileCheck --check-prefix=UNKNOWN %s
+
+// LIST-BENCHMARKS: BM_some_function
+// RUN: iree-benchmark-module --benchmark_list_tests --entry_function=some_function | IreeFileCheck --check-prefix=LIST-BENCHMARKS %s
