[ORT] onnxruntime_perf_test: Add multi-shape profiling support to perftest tool (--data_shape)

onnxruntime/test/perftest/command_args_parser.cc

@@ -214,6 +214,15 @@ ABSL_FLAG(bool, compile_ep_context, DefaultPerformanceTestConfig().run_config.co
 ABSL_FLAG(std::string, compile_model_path, "model_ctx.onnx", "The compiled model path for saving EP context model. Overwrites if already exists");
 ABSL_FLAG(bool, compile_binary_embed, DefaultPerformanceTestConfig().run_config.compile_binary_embed, "Embed binary blob within EP context node");
 ABSL_FLAG(bool, compile_only, DefaultPerformanceTestConfig().run_config.compile_only, "Only compile EP context model without running it");
+ABSL_FLAG(std::string, data_shape, "",
+          "Specifies input shapes for multi-shape profiling within a single session.\n"
+          "The model is compiled once and run with each shape group in round-robin order.\n"
+          "[Usage]: --data_shape \"input_name:[d0,d1,...][d0,d1,...] ...\"\n"
+          "[Example]: --data_shape \"input:[1,3,224,224][1,3,448,448][1,3,112,112]\"\n"
+          "With -I: generates random input of the specified shapes.\n"
+          "Without -I: selects test data folders whose tensor shapes match.\n"
+          "All inputs must have the same number of shape groups.\n"
+          "All dimension values must be positive integers.");
 ABSL_FLAG(bool, h, false, "Print program usage.");
 
 namespace onnxruntime {
@@ -481,6 +490,16 @@ bool CommandLineParser::ParseArguments(PerformanceTestConfig& test_config, int a
   // -I
   test_config.run_config.generate_model_input_binding = absl::GetFlag(FLAGS_I);
 
+  // --data_shape
+  {
+    const auto& data_shape_str = absl::GetFlag(FLAGS_data_shape);
+    if (!data_shape_str.empty()) {
+      if (!ParseDataShapeGroups(data_shape_str, test_config.run_config.data_shape_groups)) {
+        return false;
+      }
+    }
+  }
+
   // -d
   if (absl::GetFlag(FLAGS_d) < 0) return false;
   test_config.run_config.cudnn_conv_algo = absl::GetFlag(FLAGS_d);

onnxruntime/test/perftest/ort_test_session.cc

@@ -42,9 +42,14 @@ namespace onnxruntime {
 namespace perftest {
 
 RunTiming OnnxRuntimeTestSession::Run() {
-  // Randomly pick one OrtValueArray from test_inputs_. (NOT ThreadSafe)
-  const std::uniform_int_distribution<int>::param_type p(0, static_cast<int>(test_inputs_.size() - 1));
-  const size_t id = static_cast<size_t>(dist_(rand_engine_, p));
+  // Select input set: round-robin for multi-shape mode, random otherwise.
+  size_t id;
+  if (use_round_robin_ && test_inputs_.size() > 1) {
+    id = shape_group_counter_.fetch_add(1, std::memory_order_relaxed) % test_inputs_.size();
+  } else {
+    const std::uniform_int_distribution<int>::param_type p(0, static_cast<int>(test_inputs_.size() - 1));
+    id = static_cast<size_t>(dist_(rand_engine_, p));
+  }
 
   auto& input = test_inputs_.at(id);
   auto start = std::chrono::high_resolution_clock::now();
@@ -88,6 +93,7 @@ RunTiming OnnxRuntimeTestSession::Run() {
     timing.submit_timing = std::chrono::high_resolution_clock::now() - start;
     timing.total_timing = timing.submit_timing;
   }
+  timing.shape_group_index = id;
   return timing;
 }
 
@@ -1131,6 +1137,99 @@ bool OnnxRuntimeTestSession::PopulateGeneratedInputTestData(int32_t seed) {
   }
   return true;
 }
+
+bool OnnxRuntimeTestSession::PopulateMultiShapeInputTestData(
+    int32_t seed,
+    const std::map<std::string, std::vector<std::vector<int64_t>>>& data_shape_groups) {
+  // Validate that all input names in data_shape_groups exist in the model
+  for (const auto& [name, groups] : data_shape_groups) {
+    bool found = false;
+    for (int i = 0; i < input_length_; i++) {
+      if (input_names_str_[i] == name) {
+        found = true;
+        break;
+      }
+    }
+    if (!found) {
+      std::cerr << "Error: --data_shape specifies unknown input '" << name << "'." << std::endl;
+      return false;
+    }
+  }
+
+  const size_t num_groups = data_shape_groups.begin()->second.size();
+  Ort::AllocatorWithDefaultOptions default_allocator;
+
+  for (size_t g = 0; g < num_groups; g++) {
+    for (size_t i = 0; i < static_cast<size_t>(input_length_); i++) {
+      Ort::TypeInfo type_info = session_.GetInputTypeInfo(i);
+      if (type_info.GetONNXType() != ONNX_TYPE_TENSOR) {
+        continue;
+      }
+
+      auto tensor_info = type_info.GetTensorTypeAndShapeInfo();
+      std::vector<int64_t> input_node_dim;
+
+      // Use user-specified shape if available, otherwise fall back to model metadata
+      auto it = data_shape_groups.find(input_names_str_[i]);
+      if (it != data_shape_groups.end()) {
+        input_node_dim = it->second[g];
+        const auto model_shape = tensor_info.GetShape();
+        if (!model_shape.empty() && input_node_dim.size() != model_shape.size()) {
+          std::cerr << "Error: --data_shape rank mismatch for input '" << input_names_str_[i]
+                    << "': expected " << model_shape.size() << " dims but got " << input_node_dim.size() << "." << std::endl;
+          return false;
+        }
+      } else {
+        input_node_dim = tensor_info.GetShape();
+        auto transform_fcn = [](int64_t input) { return (input == -1) ? -input : input; };
+        std::transform(input_node_dim.begin(), input_node_dim.end(), input_node_dim.begin(), transform_fcn);
+      }
+
+      if (device_memory_name_ != CUDA) {
+        Ort::Value input_tensor = Ort::Value::CreateTensor(allocator_, (const int64_t*)input_node_dim.data(),
+                                                           input_node_dim.size(), tensor_info.GetElementType());
+        InitializeTensorWithSeed(seed, input_tensor);
+        PreLoadTestData(g, i, std::move(input_tensor));
+      }
+#if defined(USE_CUDA) || defined(USE_TENSORRT) || defined(USE_NV)
+      else {
+        Ort::Value default_tensor = Ort::Value::CreateTensor(default_allocator, (const int64_t*)input_node_dim.data(),
+                                                             input_node_dim.size(), tensor_info.GetElementType());
+        InitializeTensorWithSeed(seed, default_tensor);
+
+        const void* default_ptr = default_tensor.GetTensorRawData();
+        size_t total_bytes = default_tensor.GetTensorSizeInBytes();
+
+        Ort::Value cuda_tensor = Ort::Value::CreateTensor(allocator_, input_node_dim.data(),
+                                                          input_node_dim.size(), tensor_info.GetElementType());
+        void* cuda_ptr = cuda_tensor.GetTensorMutableData<void>();
+
+        cudaError_t cuda_err = cudaMemcpy(cuda_ptr, default_ptr, total_bytes, cudaMemcpyHostToDevice);
+        if (cuda_err != cudaSuccess) {
+          ORT_THROW("Failed to copy tensor data from CPU to CUDA device. CUDA Error: ", cudaGetErrorString(cuda_err));
+        }
+        PreLoadTestData(g, i, std::move(cuda_tensor));
+      }
+#endif
+    }
+  }
+  use_round_robin_ = true;
+  return true;
+}
+
+std::vector<int64_t> OnnxRuntimeTestSession::GetLoadedInputShape(size_t test_data_id, size_t input_id) const {
+  return test_inputs_.at(test_data_id).at(input_id).GetTensorTypeAndShapeInfo().GetShape();
+}
+
+void OnnxRuntimeTestSession::SelectTestDataSets(const std::vector<size_t>& selected_ids) {
+  std::vector<std::vector<Ort::Value>> filtered;
+  filtered.reserve(selected_ids.size());
+  for (size_t id : selected_ids) {
+    filtered.push_back(std::move(test_inputs_.at(id)));
+  }
+  test_inputs_ = std::move(filtered);
+}
+
 OnnxRuntimeTestSession::~OnnxRuntimeTestSession() {
 #ifdef USE_CUDA
   if (device_memory_name_ == CUDA && stream_ != nullptr) {

onnxruntime/test/perftest/ort_test_session.h

@@ -2,6 +2,7 @@
 // Licensed under the MIT License.
 
 #pragma once
+#include <atomic>
 #include <core/session/onnxruntime_cxx_api.h>
 #include <random>
 #include "test_configuration.h"
@@ -31,6 +32,13 @@ class OnnxRuntimeTestSession : public TestSession {
   }
 
   bool PopulateGeneratedInputTestData(int32_t seed);
+  bool PopulateMultiShapeInputTestData(
+      int32_t seed,
+      const std::map<std::string, std::vector<std::vector<int64_t>>>& data_shape_groups);
+
+  std::vector<int64_t> GetLoadedInputShape(size_t test_data_id, size_t input_id) const;
+  void SelectTestDataSets(const std::vector<size_t>& selected_ids);
+  void SetUseRoundRobin(bool v) { use_round_robin_ = v; }
 
   ~OnnxRuntimeTestSession();
 
@@ -60,6 +68,8 @@ class OnnxRuntimeTestSession : public TestSession {
   std::string device_memory_name_;  // Device memory type name to use from the list in allocator.h
   const std::unordered_map<std::string, std::string>& run_config_entries_;
   bool has_dynamic_output_shapes_ = false;
+  std::atomic<size_t> shape_group_counter_{0};
+  bool use_round_robin_{false};
 #if defined(USE_CUDA) || defined(USE_TENSORRT) || defined(USE_NV)
   cudaStream_t stream_;  // Device stream if required by IO bindings
 #endif