Fix run-level profiling for subgraph operators

onnxruntime/core/framework/op_kernel_context_internal.h

@@ -22,10 +22,12 @@ class OpKernelContextInternal : public OpKernelContext {
                                    const OpKernel& kernel,
                                    const logging::Logger& logger,
                                    const bool& terminate_flag,
-                                   Stream* stream)
+                                   Stream* stream,
+                                   profiling::Profiler* run_profiler = nullptr)
       : OpKernelContext(&frame, &kernel, stream, session_state.GetThreadPool(), logger),
         session_state_(session_state),
-        terminate_flag_(terminate_flag) {
+        terminate_flag_(terminate_flag),
+        run_profiler_(run_profiler) {
     const auto& implicit_inputs = kernel.Node().ImplicitInputDefs();
     int num_implicit_inputs = static_cast<int>(implicit_inputs.size());
     implicit_input_values_.reserve(num_implicit_inputs);
@@ -104,6 +106,8 @@ class OpKernelContextInternal : public OpKernelContext {
 
   const bool& GetTerminateFlag() const noexcept { return terminate_flag_; }
 
+  profiling::Profiler* GetRunProfiler() const noexcept { return run_profiler_; }
+
  private:
 #if !defined(ORT_MINIMAL_BUILD)
   class AccountingAllocator : public IAllocator {
@@ -137,6 +141,7 @@ class OpKernelContextInternal : public OpKernelContext {
 
   const SessionState& session_state_;
   const bool& terminate_flag_;
+  profiling::Profiler* run_profiler_;
   std::vector<const OrtValue*> implicit_input_values_;
 };
 

onnxruntime/core/framework/sequential_executor.cc

@@ -253,6 +253,8 @@ class SessionScope {
     return run_profiler_ && run_profiler_->IsEnabled();
   }
 
+  profiling::Profiler* GetRunProfiler() const { return run_profiler_; }
+
   void StopProfilingIfEnabled(profiling::EventCategory category,
                               const std::string& event_name,
                               const TimePoint& start_time,
@@ -490,7 +492,8 @@ onnxruntime::Status ExecuteKernel(StreamExecutionContext& ctx,
                                      *p_kernel,
                                      ctx.GetLogger(),
                                      terminate_flag,
-                                     ctx.GetDeviceStream(stream_idx));
+                                     ctx.GetDeviceStream(stream_idx),
+                                     session_scope.GetRunProfiler());
   onnxruntime::Status status;
   auto& logger = ctx.GetLogger();
   if (p_kernel->IsAsync()) {

onnxruntime/core/framework/utils.cc

@@ -883,18 +883,21 @@ common::Status ExecuteSubgraph(const SessionState& session_state, const FeedsFet
                                const std::unordered_map<size_t, IExecutor::CustomAllocator>& fetch_allocators,
                                ExecutionMode execution_mode, const bool& terminate_flag, const logging::Logger& logger,
                                Stream* parent_stream,
-                               bool sync_subgraph_fetches) {
+                               bool sync_subgraph_fetches,
+                               profiling::Profiler* run_profiler) {
 #ifdef ORT_ENABLE_STREAM
   DeviceStreamCollectionHolder device_stream_collection_holder(&session_state);
   DeviceStreamCollection* device_stream_collection = device_stream_collection_holder.p_.get();
 
   auto retval = ExecuteGraphImpl(session_state, feeds_fetches_manager, feeds, fetches, fetch_allocators,
-                                 execution_mode, terminate_flag, logger, device_stream_collection, false, parent_stream);
+                                 execution_mode, terminate_flag, logger, device_stream_collection, false, parent_stream,
+                                 run_profiler);
   if (device_stream_collection)
     ORT_CHECK_AND_SET_RETVAL(device_stream_collection->CleanUp(false));
 #else
   auto retval = ExecuteGraphImpl(session_state, feeds_fetches_manager, feeds, fetches, fetch_allocators,
-                                 execution_mode, terminate_flag, logger, false, parent_stream);
+                                 execution_mode, terminate_flag, logger, false, parent_stream,
+                                 run_profiler);
 #endif
   if (retval.IsOK() && sync_subgraph_fetches && parent_stream) {
     parent_stream->Flush();

onnxruntime/core/framework/utils.h

@@ -115,7 +115,8 @@ common::Status ExecuteSubgraph(const SessionState& session_state, const FeedsFet
                                /*when this is enabled, we will sync the parent stream to make sure the subgraph fetches
                                is complete. this is mainly used when the parent kernel depends on the CPU value of the
                                subgraph fetches, i.e. the loop condition*/
-                               bool sync_subgraph_fetches = false);
+                               bool sync_subgraph_fetches = false,
+                               profiling::Profiler* run_profiler = nullptr);
 
 bool IsInputOnCpu(const Node& node, const KernelCreateInfo* p_kci, size_t index);
 bool IsOutputOnCpu(const Node& node, const KernelCreateInfo* p_kci, size_t index);

onnxruntime/core/providers/cpu/controlflow/if.cc

@@ -468,7 +468,9 @@ Status IfImpl::Execute(const FeedsFetchesManager& ffm) {
 
   status = utils::ExecuteSubgraph(session_state_, ffm, feeds, fetches, fetch_allocators,
                                   ExecutionMode::ORT_SEQUENTIAL, context_.GetTerminateFlag(),
-                                  context_.Logger(), context_.GetComputeStream());
+                                  context_.Logger(), context_.GetComputeStream(),
+                                  /*sync_subgraph_fetches*/ false,
+                                  context_.GetRunProfiler());
 
   ORT_RETURN_IF_ERROR(status);
 

onnxruntime/core/providers/cpu/controlflow/loop.cc

@@ -546,7 +546,8 @@ Status LoopImpl::Execute(const FeedsFetchesManager& ffm) {
                                     context_.GetComputeStream(),
                                     // because the fetch[0] is the loop condition which we need to access on CPU,
                                     // have to perofrm a stream sync to make sure the data arrived.
-                                    true);
+                                    true,
+                                    context_.GetRunProfiler());
     ORT_RETURN_IF_ERROR(status);
 
     condition_mlvalue_ = fetches[0];

onnxruntime/core/providers/cpu/controlflow/scan_utils.cc

@@ -272,7 +272,9 @@ Status IterateSequence(OpKernelContextInternal& context, const SessionState& ses
     // Create Executor and run graph.
     status = utils::ExecuteSubgraph(session_state, ffm, feeds, fetches, fetch_allocators,
                                     ExecutionMode::ORT_SEQUENTIAL, context.GetTerminateFlag(), context.Logger(),
-                                    context.GetComputeStream());
+                                    context.GetComputeStream(),
+                                    /*sync_subgraph_fetches*/ false,
+                                    context.GetRunProfiler());
 
     ORT_RETURN_IF_ERROR(status);
 

onnxruntime/test/framework/inference_session_test.cc

@@ -14,6 +14,9 @@
 #include <fstream>
 #include <random>
 
+#include "nlohmann/json.hpp"
+#include "onnxruntime_cxx_api.h"
+
 #include <google/protobuf/io/zero_copy_stream_impl.h>
 #include "core/common/denormal.h"
 #include "core/common/logging/logging.h"
@@ -63,6 +66,8 @@ using namespace ONNX_NAMESPACE;
 using namespace onnxruntime::logging;
 using namespace onnxruntime::concurrency;
 
+extern std::unique_ptr<Ort::Env> ort_env;
+
 namespace {
 struct KernelRegistryAndStatus {
   std::shared_ptr<onnxruntime::KernelRegistry> kernel_registry = std::make_shared<onnxruntime::KernelRegistry>();
@@ -620,7 +625,89 @@ TEST(InferenceSessionTests, CheckRunProfilerWithRunOptions) {
   // Clean up the profile file
   std::remove(profile_file.c_str());
 }
-#endif  // __wasm__
+#endif  // !defined(__wasm__) && !defined(_WIN32)
+
+#ifndef __wasm__
+// Test that run-level profiling captures operators inside subgraphs (e.g., If branches).
+TEST(InferenceSessionTests, CheckRunProfilerWithSubgraph) {
+  Ort::SessionOptions session_options;
+
+  Ort::Session session(*ort_env, ORT_TSTR("testdata/if_mul.onnx"), session_options);
+
+  // Prepare inputs for if_mul.onnx:
+  //   A (bool, [1])    - condition (true -> then branch: C = B * 2)
+  //   B (float, [3,2]) - data
+  Ort::MemoryInfo memory_info = Ort::MemoryInfo::CreateCpu(OrtDeviceAllocator, OrtMemTypeCPU);
+  std::array<int64_t, 1> a_shape = {1};
+  std::array<int64_t, 2> b_shape = {3, 2};
+
+  std::array<bool, 1> a_data = {true};
+  std::array<float, 6> b_data = {2.f, 3.f, 4.f, -5.f, 6.f, 7.f};
+
+  std::vector<Ort::Value> ort_inputs;
+  ort_inputs.emplace_back(
+      Ort::Value::CreateTensor<bool>(memory_info, a_data.data(), a_data.size(), a_shape.data(), a_shape.size()));
+  ort_inputs.emplace_back(
+      Ort::Value::CreateTensor<float>(memory_info, b_data.data(), b_data.size(), b_shape.data(), b_shape.size()));
+
+  std::array ort_input_names{"A", "B"};
+  std::array output_names{"C"};
+
+  // Enable run-level profiling via RunOptions
+  Ort::RunOptions run_options;
+  run_options.EnableProfiling(ORT_TSTR("ort_run_profile_subgraph_test"));
+
+  std::vector<Ort::Value> ort_outputs = session.Run(run_options, ort_input_names.data(), ort_inputs.data(),
+                                                    ort_inputs.size(), output_names.data(), output_names.size());
+
+  // Verify output: condition=true -> then branch -> B * 2
+  const float* output_data = ort_outputs[0].GetTensorData<float>();
+  gsl::span<const float> output_span(output_data, 6);
+  EXPECT_THAT(output_span, ::testing::ElementsAre(4.f, 6.f, 8.f, -10.f, 12.f, 14.f));
+
+  // Find the generated profile JSON file
+  std::string profile_file;
+  for (const auto& entry : std::filesystem::directory_iterator(".")) {
+    std::string filename = entry.path().filename().string();
+    if (filename.find("ort_run_profile_subgraph_test") == 0 && filename.find(".json") != std::string::npos) {
+      profile_file = entry.path().string();
+      break;
+    }
+  }
+
+  ASSERT_FALSE(profile_file.empty()) << "Profile file with prefix 'ort_run_profile_subgraph_test' not found";
+
+  // Ensure the profile file is cleaned up even if an assertion fails.
+  auto cleanup = gsl::finally([&profile_file]() { std::remove(profile_file.c_str()); });
+
+  // Parse the profile JSON
+  std::ifstream profile_stream(profile_file);
+  ASSERT_TRUE(profile_stream.good()) << "Failed to open profile file: " << profile_file;
+
+  nlohmann::json profile_json;
+  profile_stream >> profile_json;
+  profile_stream.close();
+
+  ASSERT_TRUE(profile_json.is_array()) << "Profile JSON should be an array";
+  ASSERT_FALSE(profile_json.empty()) << "Profile JSON should not be empty";
+
+  // Check that the profile contains an entry for the Mul op inside the If's then-branch (mul_0).
+  bool found_subgraph_mul = false;
+  for (const auto& entry : profile_json) {
+    if (entry.contains("name")) {
+      const std::string name = entry["name"].get<std::string>();
+      if (name.find("mul_0") != std::string::npos) {
+        found_subgraph_mul = true;
+        break;
+      }
+    }
+  }
+
+  EXPECT_TRUE(found_subgraph_mul)
+      << "Profile should contain an entry for 'mul_0' (Mul op inside If's then-branch). "
+      << "Profile contents: " << profile_json;
+}
+#endif  // !defined(__wasm__)
 
 TEST(InferenceSessionTests, CheckRunProfilerStartTime) {
   // Test whether the InferenceSession can access the profiler's start time