[iree-benchmark-module] transfer outputs to host in case of enable_output_processing (#24821)

The enable_output_processing feature added to iree-benchmark_module in
commit a70389bb8b [iree-benchmark-module] support processing outputs
after benchmarking (#24789) did not call
iree_tooling_transfer_variants() to make sure the output buffers are
host accessible. So depending on the device used for benchmarking, the
outputs could be accessed on the host to dump those or to check those.
This PR fixes this by calling iree_tooling_transfer_variants() in case
output processing is enabled. This is the same handling as in
iree-run-module.

Signed-off-by: Stefan Schuermans <schuermans@roofline.ai>

tools/BUILD.bazel

diff --git a/tools/BUILD.bazel b/tools/BUILD.bazel
index 3277371..7fe724e 100644
--- a/tools/BUILD.bazel
+++ b/tools/BUILD.bazel
@@ -70,6 +70,7 @@
         "//runtime/src/iree/tooling:context_util",
         "//runtime/src/iree/tooling:device_util",
         "//runtime/src/iree/tooling:function_io",
+        "//runtime/src/iree/tooling:function_util",
         "//runtime/src/iree/tooling:process_results",
         "//runtime/src/iree/vm",
         "@com_google_benchmark//:benchmark",

tools/CMakeLists.txt

diff --git a/tools/CMakeLists.txt b/tools/CMakeLists.txt
index e1598ba..662ba56 100644
--- a/tools/CMakeLists.txt
+++ b/tools/CMakeLists.txt
@@ -101,6 +101,7 @@
       iree::tooling::context_util
       iree::tooling::device_util
       iree::tooling::function_io
+      iree::tooling::function_util
       iree::tooling::process_results
       iree::vm
     COVERAGE ${IREE_ENABLE_RUNTIME_COVERAGE}

tools/iree-benchmark-module-main.cc

diff --git a/tools/iree-benchmark-module-main.cc b/tools/iree-benchmark-module-main.cc
index 6c75474..fbd3293 100644
--- a/tools/iree-benchmark-module-main.cc
+++ b/tools/iree-benchmark-module-main.cc
@@ -69,6 +69,7 @@
 #include "iree/tooling/context_util.h"
 #include "iree/tooling/device_util.h"
 #include "iree/tooling/function_io.h"
+#include "iree/tooling/function_util.h"
 #include "iree/tooling/process_results.h"
 #include "iree/vm/api.h"
 
@@ -612,6 +613,22 @@
                               "no output list to process");
     }
 
+    // Transfer outputs to the host so they can be processed. Only required when
+    // using full HAL device-based execution.
+    if (device_ != NULL) {
+      iree_hal_buffer_params_t target_params = {};
+      target_params.usage =
+          IREE_HAL_BUFFER_USAGE_TRANSFER | IREE_HAL_BUFFER_USAGE_MAPPING;
+      target_params.access = IREE_HAL_MEMORY_ACCESS_ALL;
+      target_params.type =
+          IREE_HAL_MEMORY_TYPE_HOST_LOCAL | IREE_HAL_MEMORY_TYPE_DEVICE_VISIBLE;
+      target_params.queue_affinity = IREE_HAL_QUEUE_AFFINITY_ANY;
+      target_params.min_alignment = 0;
+      IREE_RETURN_IF_ERROR(iree_tooling_transfer_variants(
+          outputs_.get(), device_, device_allocator_.get(), target_params,
+          /*wait_fence=*/NULL, /*signal_fence=*/NULL));
+    }
+
     IREE_RETURN_IF_ERROR(iree_tooling_process_results_and_print(
         device_,
         iree_make_string_view(results_cconv_.data(), results_cconv_.size()),