Adding the `--print_max_element_count=` flag to iree-run-module. (#8105)

This allows overriding the 1024 limit in cases where the user wants to
see the full tensor outputs (hopefully piping to a file).

iree/tools/iree-run-module-main.cc

diff --git a/iree/tools/iree-run-module-main.cc b/iree/tools/iree-run-module-main.cc
index ec24203..fe9b311 100644
--- a/iree/tools/iree-run-module-main.cc
+++ b/iree/tools/iree-run-module-main.cc
@@ -37,6 +37,10 @@
 
 IREE_FLAG(string, driver, "vmvx", "Backend driver to use.");
 
+IREE_FLAG(int32_t, print_max_element_count, 1024,
+          "Prints up to the maximum number of elements of output tensors, "
+          "eliding the remainder.");
+
 IREE_FLAG(bool, print_statistics, false,
           "Prints runtime statistics to stderr on exit.");
 
@@ -152,7 +156,9 @@
                      iree_allocator_system()),
       "invoking function '%s'", function_name.c_str());
 
-  IREE_RETURN_IF_ERROR(PrintVariantList(outputs.get()), "printing results");
+  IREE_RETURN_IF_ERROR(
+      PrintVariantList(outputs.get(), (size_t)FLAG_print_max_element_count),
+      "printing results");
 
   inputs.reset();
   outputs.reset();

iree/tools/utils/vm_util.cc

diff --git a/iree/tools/utils/vm_util.cc b/iree/tools/utils/vm_util.cc
index 39e1331..c3db733 100644
--- a/iree/tools/utils/vm_util.cc
+++ b/iree/tools/utils/vm_util.cc
@@ -134,7 +134,8 @@
   return OkStatus();
 }
 
-Status PrintVariantList(iree_vm_list_t* variant_list, std::ostream* os) {
+Status PrintVariantList(iree_vm_list_t* variant_list, size_t max_element_count,
+                        std::ostream* os) {
   for (iree_host_size_t i = 0; i < iree_vm_list_size(variant_list); ++i) {
     iree_vm_variant_t variant = iree_vm_variant_empty();
     IREE_RETURN_IF_ERROR(iree_vm_list_get_variant(variant_list, i, &variant),
@@ -175,9 +176,9 @@
         iree_status_t status;
         do {
           iree_host_size_t actual_length = 0;
-          status = iree_hal_buffer_view_format(
-              buffer_view, /*max_element_count=*/1024, result_str.size() + 1,
-              &result_str[0], &actual_length);
+          status = iree_hal_buffer_view_format(buffer_view, max_element_count,
+                                               result_str.size() + 1,
+                                               &result_str[0], &actual_length);
           result_str.resize(actual_length);
         } while (iree_status_is_out_of_range(status));
         IREE_RETURN_IF_ERROR(status);

iree/tools/utils/vm_util.h

diff --git a/iree/tools/utils/vm_util.h b/iree/tools/utils/vm_util.h
index 5da92e4..697e5a1 100644
--- a/iree/tools/utils/vm_util.h
+++ b/iree/tools/utils/vm_util.h
@@ -46,8 +46,15 @@
 // described in
 // https://github.com/google/iree/tree/main/iree/hal/api.h
 // Uses descriptors in |descs| for type information and validation.
-Status PrintVariantList(iree_vm_list_t* variant_list,
-                        std::ostream* os = &std::cout);
+Status PrintVariantList(iree_vm_list_t* variant_list, size_t max_element_count,
+                        std::ostream* os);
+inline Status PrintVariantList(iree_vm_list_t* variant_list, std::ostream* os) {
+  return PrintVariantList(variant_list, 1024, os);
+}
+inline Status PrintVariantList(iree_vm_list_t* variant_list,
+                               size_t max_element_count = 1024) {
+  return PrintVariantList(variant_list, max_element_count, &std::cout);
+}
 
 // Creates the default device for |driver| in |out_device|.
 // The returned |out_device| must be released by the caller.