Modify gpu_numerics_check to accept multiple models.

LiteRT-PiperOrigin-RevId: 826195884

Author: gcarranza-1

litert/core/filesystem.cc

@@ -73,6 +73,10 @@ std::string Join(const std::vector<absl::string_view>& paths) {
   return std_path.generic_string();
 }
 
+std::string Stem(absl::string_view path) {
+  return MakeStdPath(path).stem().generic_string();
+}
+
 bool Exists(absl::string_view path) { return StdExists(MakeStdPath(path)); }
 
 Expected<size_t> Size(absl::string_view path) {

litert/core/filesystem.h

@@ -32,6 +32,9 @@ namespace litert::internal {
 // Append all given subpaths together (e.g. os.path.join).
 std::string Join(const std::vector<absl::string_view>& paths);
 
+// Returns the stem of the given path.
+std::string Stem(absl::string_view path);
+
 // Make a new empty file at the given path.
 void Touch(absl::string_view path);
 

litert/core/filesystem_test.cc

@@ -24,6 +24,8 @@ namespace {
 static constexpr absl::string_view kPrefix = "a/prefix";
 static constexpr absl::string_view kInfix = "an/infix";
 static constexpr absl::string_view kSuffix = "suffix.ext";
+static constexpr absl::string_view kPath = "a/prefix.ext";
+static constexpr absl::string_view kStem = "prefix";
 
 TEST(FilesystemTest, JoinTwo) {
   const auto path = Join({kPrefix, kSuffix});
@@ -35,5 +37,10 @@ TEST(FilesystemTest, JoinMany) {
   EXPECT_EQ(path, absl::StrFormat("%s/%s/%s", kPrefix, kInfix, kSuffix));
 }
 
+TEST(FilesystemTest, Stem){
+  const auto stem = Stem(kPath);
+  EXPECT_EQ(stem, kStem);
+}
+
 }  // namespace
 }  // namespace litert::internal

litert/tools/BUILD

@@ -511,6 +511,7 @@ cc_binary(
     srcs = ["gpu_numerics_check.cc"],
     deps = NUMERICS_CHECK_DEPS + [
         "//litert/cc/options:litert_gpu_options",
+        "//litert/core:filesystem",
     ] + GPU_ACCELERATOR_DEPS,
 )
 
@@ -519,6 +520,7 @@ cc_binary(
     srcs = ["gpu_numerics_check.cc"],
     deps = NUMERICS_CHECK_DEPS + [
         "//litert/cc/options:litert_gpu_options",
+        "//litert/core:filesystem",
         # copybara:uncomment_begin(google-only)
         # "//litert/runtime/accelerators/gpu:ml_drift_cl_gl_accelerator",  # buildcleaner: keep
         # copybara:uncomment_end
@@ -531,6 +533,7 @@ cc_binary(
     tags = ["requires-gpu-nvidia"],
     deps = NUMERICS_CHECK_DEPS + [
         "//litert/cc/options:litert_gpu_options",
+        "//litert/core:filesystem",
         # copybara:uncomment_begin(google-only)
         # "//litert/runtime/accelerators/gpu:ml_drift_vulkan_accelerator",  # buildcleaner: keep
         # "//third_party/vulkan_loader",
@@ -548,6 +551,7 @@ cc_binary(
     ],
     deps = NUMERICS_CHECK_DEPS + [
         "//litert/cc/options:litert_gpu_options",
+        "//litert/core:filesystem",
         # copybara:uncomment_begin(google-only)
         # "//litert/runtime/accelerators/gpu/google:jet_gpu_accelerator",  # buildcleaner: keep
         # copybara:uncomment_end

litert/tools/gpu_numerics_check.cc

@@ -29,6 +29,7 @@
 #include "absl/flags/parse.h"  // from @com_google_absl
 #include "absl/log/absl_log.h"  // from @com_google_absl
 #include "absl/strings/str_format.h"  // from @com_google_absl
+#include "absl/strings/string_view.h"  // from @com_google_absl
 #include "absl/types/span.h"  // from @com_google_absl
 #include "litert/c/litert_common.h"
 #include "litert/cc/litert_compiled_model.h"
@@ -40,8 +41,14 @@
 #include "litert/cc/litert_options.h"
 #include "litert/cc/litert_tensor_buffer.h"
 #include "litert/cc/options/litert_gpu_options.h"
-
-ABSL_FLAG(std::string, graph, "", "Model filename to use for testing.");
+#include "litert/core/filesystem.h"
+
+ABSL_FLAG(bool, print_diff_stats, false,
+          "Whether to print the diff stats CSV.");
+ABSL_FLAG(std::string, model_dir, "",
+          "Optional base directory to prepend to models provide in --graph.");
+ABSL_FLAG(std::vector<std::string>, graph, {},
+          "Model file(s) to use for testing.");
 ABSL_FLAG(size_t, signature_index, 0, "Index of the signature to run.");
 ABSL_FLAG(float, epsilon, 1e-4f,
           "Threshold value for gpu / cpu inference comparison");
@@ -56,8 +63,33 @@ ABSL_FLAG(bool, enable_constant_tensors_sharing, false,
           "Whether to enable constant tensors sharing.");
 
 namespace litert {
+
 namespace {
 
+struct BufferDiffStats {
+  // Index of output buffer.
+  size_t buffer_idx;
+  // Total number of elements in the buffer.
+  size_t total_elements;
+  // Number of elements with absolute difference greater than epsilon.
+  size_t diff_elements;
+  // Epsilon value used for comparison.
+  double epsilon;
+  // Maximum absolute difference between CPU and GPU values.
+  double max_diff;
+  // Minimum absolute difference between CPU and GPU values.
+  double min_diff;
+  // Mean absolute difference between CPU and GPU values.
+  double mean_diff;
+  // Mean squared error between CPU and GPU values.
+  double mse;
+};
+
+struct ModelRunResult {
+  std::string model_name;
+  std::vector<BufferDiffStats> diff_stats;
+};
+
 Expected<Environment> GetEnvironment() {
   std::vector<litert::Environment::Option> environment_options = {};
 
@@ -154,13 +186,14 @@ Expected<std::vector<TensorBuffer>> CreateOutputBuffers(
 }
 
 // Compares a single pair of output buffers and prints the results.
-Expected<void> CompareSingleOutputBuffer(TensorBuffer& cpu_buffer,
-                                         TensorBuffer& gpu_buffer,
-                                         size_t buffer_index, float epsilon) {
+Expected<BufferDiffStats> CompareSingleOutputBuffer(TensorBuffer& cpu_buffer,
+                                                    TensorBuffer& gpu_buffer,
+                                                    size_t buffer_index,
+                                                    float epsilon) {
   std::vector<std::pair<float, int>> all_diffs;
   const int kMaxPrint = 20;
   int printed = 0;
-  int total_different = 0;
+  size_t total_different = 0;
   double mean_squared_error = 0;
   float mean_diff = 0;
 
@@ -299,10 +332,19 @@ Expected<void> CompareSingleOutputBuffer(TensorBuffer& cpu_buffer,
   std::cout << "Total " << total_different << " out of " << total_elements
             << " are different elements, for output #" << buffer_index
             << ", threshold - " << epsilon << std::endl;
-  return {};
+  return BufferDiffStats{
+      .buffer_idx = buffer_index,
+      .total_elements = total_elements,
+      .diff_elements = total_different,
+      .epsilon = epsilon,
+      .max_diff = all_diffs.back().first,
+      .min_diff = all_diffs.front().first,
+      .mean_diff = mean_diff / all_diffs.size(),
+      .mse = mean_squared_error / total_elements,
+  };
 }
 
-Expected<void> CompareOutputBuffers(
+Expected<std::vector<BufferDiffStats>> CompareOutputBuffers(
     std::vector<TensorBuffer>& cpu_output_buffers,
     std::vector<TensorBuffer>& gpu_output_buffers) {
   if (cpu_output_buffers.size() != gpu_output_buffers.size()) {
@@ -312,26 +354,24 @@ Expected<void> CompareOutputBuffers(
 
   float epsilon = absl::GetFlag(FLAGS_epsilon);
   size_t num_output_buffers = cpu_output_buffers.size();
+  std::vector<BufferDiffStats> diff_stats;
   for (size_t i = 0; i < num_output_buffers; ++i) {
     auto& cpu_buffer = cpu_output_buffers[i];
     auto& gpu_buffer = gpu_output_buffers[i];
-    LITERT_RETURN_IF_ERROR(
+    LITERT_ASSIGN_OR_RETURN(
+        auto diff_stat,
         CompareSingleOutputBuffer(cpu_buffer, gpu_buffer, i, epsilon));
+    diff_stats.push_back(std::move(diff_stat));
   }
-  return {};
+  return diff_stats;
 }
 
-Expected<void> RunModel() {
-  if (absl::GetFlag(FLAGS_graph).empty()) {
-    return Error(kLiteRtStatusErrorInvalidArgument,
-                 "model filename is empty. Use --graph to provide it.");
-  }
-
-  ABSL_LOG(INFO) << "Model: " << absl::GetFlag(FLAGS_graph);
+Expected<std::vector<BufferDiffStats>> RunModel(absl::string_view model_path) {
+  ABSL_LOG(INFO) << "Model: " << model_path;
   LITERT_ASSIGN_OR_RETURN(auto cpu_model,
-                          Model::CreateFromFile(absl::GetFlag(FLAGS_graph)));
+                          Model::CreateFromFile(std::string(model_path)));
   LITERT_ASSIGN_OR_RETURN(auto gpu_model,
-                          Model::CreateFromFile(absl::GetFlag(FLAGS_graph)));
+                          Model::CreateFromFile(std::string(model_path)));
 
   LITERT_ASSIGN_OR_RETURN(auto env, GetEnvironment());
 
@@ -373,10 +413,53 @@ Expected<void> RunModel() {
       signature_index, gpu_input_buffers, gpu_output_buffers));
 
   // Compare output buffers
-  LITERT_RETURN_IF_ERROR(
+  LITERT_ASSIGN_OR_RETURN(
+      auto diff_stats,
       CompareOutputBuffers(cpu_output_buffers, gpu_output_buffers));
+  return diff_stats;
+}
+
+Expected<std::vector<ModelRunResult>> RunModels() {
+  std::vector<std::string> relative_model_paths = absl::GetFlag(FLAGS_graph);
+  if (relative_model_paths.empty()) {
+    return Error(kLiteRtStatusErrorInvalidArgument,
+                 "No model provided. Use --graph to provide it.");
+  }
 
-  return {};
+  std::string model_dir = absl::GetFlag(FLAGS_model_dir);
+  std::vector<std::string> full_model_paths;
+  full_model_paths.reserve(relative_model_paths.size());
+  for (auto& model_path : relative_model_paths) {
+    full_model_paths.push_back(internal::Join({model_dir, model_path}));
+  }
+
+  std::vector<ModelRunResult> results;
+  for (const auto& model_path : full_model_paths) {
+    LITERT_ASSIGN_OR_RETURN(std::vector<BufferDiffStats> diff_stats,
+                            RunModel(model_path));
+    results.push_back(ModelRunResult{
+        .model_name = internal::Stem(model_path),
+        .diff_stats = std::move(diff_stats),
+    });
+  }
+
+  return results;
+}
+
+void PrintDiffStats(const std::vector<litert::ModelRunResult>& results) {
+  // Print CSV header
+  std::cout << "model_name, buffer_idx, total_elements, diff_elements, "
+               "epsilon, max_diff, min_diff, mean_diff, mse"
+            << std::endl;
+  for (const auto& result : results) {
+    for (const auto& diff_stat : result.diff_stats) {
+      std::cout << result.model_name << ", " << diff_stat.buffer_idx << ", "
+                << diff_stat.total_elements << ", " << diff_stat.diff_elements
+                << ", " << diff_stat.epsilon << ", " << diff_stat.max_diff
+                << ", " << diff_stat.min_diff << ", " << diff_stat.mean_diff
+                << ", " << diff_stat.mse << std::endl;
+    }
+  }
 }
 
 }  // namespace
@@ -385,10 +468,15 @@ Expected<void> RunModel() {
 int main(int argc, char** argv) {
   absl::ParseCommandLine(argc, argv);
 
-  auto res = litert::RunModel();
+  auto res = litert::RunModels();
   if (!res) {
     ABSL_LOG(ERROR) << res.Error().Message();
     return EXIT_FAILURE;
   }
+
+  if (absl::GetFlag(FLAGS_print_diff_stats)) {
+    litert::PrintDiffStats(*res);
+  }
+
   return EXIT_SUCCESS;
 }