[Draft] Add internal MatMulNBits core operation-count estimator

include/onnxruntime/core/framework/op_kernel_context.h

@@ -5,6 +5,7 @@ namespace onnxruntime {
 class IExecutionFrame;
 #if !defined(ORT_MINIMAL_BUILD)
 class KernelPilot;
+struct RuntimeWorkloadRecorder;
 #endif
 class Stream;
 namespace concurrency {
@@ -216,6 +217,8 @@ class OpKernelContext {
   // Keep new virtuals after existing declarations to preserve their vtable slots.
   // Session-owned pilot for this kernel; nullptr when unavailable.
   virtual KernelPilot* GetKernelPilot() const { return nullptr; }
+  // Internal typed diagnostic transport. Standalone/constant-folding contexts do not collect.
+  virtual const RuntimeWorkloadRecorder* GetRuntimeWorkloadRecorder() const { return nullptr; }
 #endif
 
  private:

include/onnxruntime/core/framework/run_options.h

@@ -11,6 +11,7 @@
 #include "core/framework/config_options.h"
 
 namespace onnxruntime {
+class RuntimeWorkloadCollector;
 namespace lora {
 class LoraAdapter;
 }
@@ -38,6 +39,12 @@ struct OrtRunOptions {
   // Set to 'true' to enable profiling for this run.
   bool enable_profiling = false;
 
+#if !defined(ORT_MINIMAL_BUILD)
+  // Internal diagnostic hook. Caller owns a fresh collector through completion of this Run.
+  // No C API setter or report/export lifecycle is exposed.
+  onnxruntime::RuntimeWorkloadCollector* runtime_workload_collector = nullptr;
+#endif
+
   // File prefix for profiling result for this run.
   // The actual filename will be: <profile_file_prefix>_<timestamp>.json
   // Only used when enable_profiling is true.

include/onnxruntime/core/session/onnxruntime_session_options_config_keys.h

@@ -577,6 +577,23 @@ static const char* const kOrtSessionOptionsMaxShapeOverride = "session.max_shape
 static const char* const kOrtSessionOptionsStrictWorkspaceVerification =
     "session.strict_workspace_verification";
 
+// Opt-in runtime compute-estimation JSON reports. Value is a UTF-8 output prefix, not a literal filename.
+// Empty/unset disables automatic capture/export. After each canonical Run, append a process token, session ID,
+// run ID and ".json". Parent directories are not created. Use a new prefix for independent application launches.
+// Replaces the legacy initialization-time static export; initialization no longer writes workload reports.
+// Reports contain actual shapes, graph/node identities, CPU MatMulNBits core dense-equivalent and logical bias counts,
+// route provenance, run outcome and collection diagnostics. Unsupported operators/providers are not instrumented;
+// captured totals are not whole-model totals or hardware instruction counts. No additional device synchronization.
+// File failures warn without failing inference. Minimal builds warn that export is unsupported.
+// Internal callers supplying a RunOptions collector use that collector's limits instead of automatic limits.
+static const char* const kOrtSessionOptionsComputeEstimateReportPath = "session.compute_estimate_report_path";
+
+// Nonnegative decimal per-run automatic capture budgets; defaults: 10000 records and 16777216 bytes.
+// Zero retains no records. Bytes are accounted payload, not a hard allocator/peak-memory cap.
+// Invalid values fail initialization. Limits do not bound total report files/disk usage across runs.
+static const char* const kOrtSessionOptionsComputeEstimateMaxRecords = "session.compute_estimate_max_records";
+static const char* const kOrtSessionOptionsComputeEstimateMaxPayloadBytes = "session.compute_estimate_max_payload_bytes";
+
 // Enable EP context feature to dump the partitioned graph which includes the EP context into Onnx file.
 // The dumped Onnx model with EP context can be used for future inference to avoid the EP graph partitioning/compile overhead.
 // "0": disable. (default)

onnxruntime/contrib_ops/cpu/quantization/matmul_nbits.cc

@@ -12,6 +12,7 @@
 #include "core/common/narrow.h"
 #include "core/common/safeint.h"
 #include "core/framework/op_kernel.h"
+#include "core/framework/runtime_workload_collector.h"
 #include "core/mlas/inc/mlas_qnbit.h"
 #include "core/mlas/inc/mlas_q4.h"
 #include "core/providers/cpu/math/matmul_helper.h"
@@ -1535,8 +1536,22 @@ Status MatMulNBits<T1>::Compute(OpKernelContext* ctx) const {
 
   Tensor* y = ctx->Output(0, helper.OutputShape());
 
+#if !defined(ORT_MINIMAL_BUILD)
+  const auto* recorder = ctx->GetRuntimeWorkloadRecorder();
+  auto record = [&](MatMulNBitsRuntimeRoute route, MatMulNBitsRuntimeComputeMode mode) {
+    if (recorder) {
+      recorder->RecordMatMulNBits(a->Shape(), y->Shape(), helper.M(), helper.N(), helper.K(),
+                                  helper.OutputOffsets().size(), bias != nullptr, a->GetElementType(),
+                                  static_cast<uint32_t>(nbits_), static_cast<uint32_t>(block_size_), route, mode);
+    }
+  };
+#endif
+
   // Bail out early if the output is going to be empty
   if (y->Shape().Size() == 0) {
+#if !defined(ORT_MINIMAL_BUILD)
+    record(MatMulNBitsRuntimeRoute::EmptyOutput, MatMulNBitsRuntimeComputeMode::NotExecuted);
+#endif
     return Status::OK();
   }
 
@@ -1556,11 +1571,38 @@ Status MatMulNBits<T1>::Compute(OpKernelContext* ctx) const {
                     // MlasQNBitGemmPackQuantBDataSize() returns 0, we can consider calling MlasQNBitGemmBatch()
                     // with B directly too.
     if (prefer_lut_gemm_) {
-      return ComputeBPackedLUT(a, bias, y, thread_pool, helper);
+      ORT_RETURN_IF_ERROR(ComputeBPackedLUT(a, bias, y, thread_pool, helper));
+#if !defined(ORT_MINIMAL_BUILD)
+      record(MatMulNBitsRuntimeRoute::PackedLut, MatMulNBitsRuntimeComputeMode::Lut);
+#endif
+      return Status::OK();
     }
 
     if (MlasIsQNBitGemmAvailable(nbits_, block_size_, compute_type_)) {
-      return ComputeBPacked(a, scales, zero_points, bias, y, allocator, thread_pool, helper);
+      ORT_RETURN_IF_ERROR(ComputeBPacked(a, scales, zero_points, bias, y, allocator, thread_pool, helper));
+#if !defined(ORT_MINIMAL_BUILD)
+      // HQNBIT_CompInt8 converts half activations and calls the SQNBIT backend on every platform.
+      std::optional<MatMulNBitsRuntimeComputeMode> mode;
+      switch (compute_type_) {
+        case SQNBIT_CompFp32:
+          mode = MatMulNBitsRuntimeComputeMode::SQNBitFp32;
+          break;
+        case HQNBIT_CompFp16:
+          mode = MatMulNBitsRuntimeComputeMode::HQNBitFp16;
+          break;
+        case BHQNBIT_CompBf16:
+          mode = MatMulNBitsRuntimeComputeMode::BHQNBitBf16;
+          break;
+        case SQNBIT_CompInt8:
+        case HQNBIT_CompInt8:
+          mode = MatMulNBitsRuntimeComputeMode::SQNBitInt8;
+          break;
+      }
+      if (mode) {
+        record(MatMulNBitsRuntimeRoute::PackedQNBit, *mode);
+      }
+#endif
+      return Status::OK();
     }
   }
 
@@ -1571,7 +1613,11 @@ Status MatMulNBits<T1>::Compute(OpKernelContext* ctx) const {
                                "This is because MLAS doesn't have an optimized quantized kernel "
                                "for the requested compute configuration.";
 
-  return ComputeBUnpacked(a, b, scales, zero_points, reorder_idx, bias, y, allocator, thread_pool, helper);
+  ORT_RETURN_IF_ERROR(ComputeBUnpacked(a, b, scales, zero_points, reorder_idx, bias, y, allocator, thread_pool, helper));
+#if !defined(ORT_MINIMAL_BUILD)
+  record(MatMulNBitsRuntimeRoute::Unpacked, MatMulNBitsRuntimeComputeMode::Float32Gemm);
+#endif
+  return Status::OK();
 }
 
 #define REGISTER_MatMulNBits(T1)                                         \

onnxruntime/contrib_ops/cpu/transformers/beam_search_impl_gpt.h

@@ -304,7 +304,8 @@ Status BeamSearchGpt<T>::Execute(const FeedsFetchesManager* init_run_feeds_fetch
                                       this->context_.Logger(),
                                       this->ort_stream_,
                                       /*sync_subgraph_fetches*/ false,
-                                      this->context_.GetRunProfiler());
+                                      this->context_.GetRunProfiler(),
+                                      this->context_.GetRuntimeWorkloadCollector());
     } else {
 #ifdef DEBUG_NODE_INPUTS_OUTPUTS
       const_cast<SessionState&>(this->decoder_session_state_).IncrementGraphExecutionCounter();
@@ -319,7 +320,8 @@ Status BeamSearchGpt<T>::Execute(const FeedsFetchesManager* init_run_feeds_fetch
                                       this->context_.Logger(),
                                       this->ort_stream_,
                                       /*sync_subgraph_fetches*/ false,
-                                      this->context_.GetRunProfiler());
+                                      this->context_.GetRunProfiler(),
+                                      this->context_.GetRuntimeWorkloadCollector());
     }
 
     ORT_RETURN_IF_ERROR(status);

onnxruntime/contrib_ops/cpu/transformers/beam_search_impl_t5.h

@@ -182,7 +182,8 @@ Status BeamSearchT5<T>::Execute(const FeedsFetchesManager& encoder_feeds_fetches
                                   this->context_.Logger(),
                                   this->ort_stream_,
                                   /*sync_subgraph_fetches*/ false,
-                                  this->context_.GetRunProfiler());
+                                  this->context_.GetRunProfiler(),
+                                  this->context_.GetRuntimeWorkloadCollector());
   ORT_RETURN_IF_ERROR(status);
 
 #ifdef DEBUG_GENERATION
@@ -362,7 +363,8 @@ Status BeamSearchT5<T>::Execute(const FeedsFetchesManager& encoder_feeds_fetches
                                     this->context_.Logger(),
                                     this->ort_stream_,
                                     /*sync_subgraph_fetches*/ false,
-                                    this->context_.GetRunProfiler());
+                                    this->context_.GetRunProfiler(),
+                                    this->context_.GetRuntimeWorkloadCollector());
 
     ORT_RETURN_IF_ERROR(status);
 

onnxruntime/contrib_ops/cpu/transformers/beam_search_impl_whisper.h

@@ -188,7 +188,8 @@ Status BeamSearchWhisper<T>::Execute(const FeedsFetchesManager& encoder_feeds_fe
                                   this->context_.Logger(),
                                   this->ort_stream_,
                                   /*sync_subgraph_fetches*/ false,
-                                  this->context_.GetRunProfiler());
+                                  this->context_.GetRunProfiler(),
+                                  this->context_.GetRuntimeWorkloadCollector());
   ORT_RETURN_IF_ERROR(status);
 
 #ifdef DEBUG_GENERATION
@@ -388,7 +389,8 @@ Status BeamSearchWhisper<T>::Execute(const FeedsFetchesManager& encoder_feeds_fe
                                     this->context_.Logger(),
                                     this->ort_stream_,
                                     /*sync_subgraph_fetches*/ false,
-                                    this->context_.GetRunProfiler());
+                                    this->context_.GetRunProfiler(),
+                                    this->context_.GetRuntimeWorkloadCollector());
 
     ORT_RETURN_IF_ERROR(status);
 

onnxruntime/contrib_ops/cpu/transformers/greedy_search_impl_gpt.h

@@ -296,7 +296,8 @@ Status GreedySearchGpt<T, ParametersT>::Execute(const FeedsFetchesManager* init_
                                       this->context_.Logger(),
                                       this->ort_stream_,
                                       /*sync_subgraph_fetches*/ false,
-                                      this->context_.GetRunProfiler());
+                                      this->context_.GetRunProfiler(),
+                                      this->context_.GetRuntimeWorkloadCollector());
     } else {
 #ifdef DEBUG_NODE_INPUTS_OUTPUTS
       const_cast<SessionState&>(this->decoder_session_state_).IncrementGraphExecutionCounter();
@@ -311,7 +312,8 @@ Status GreedySearchGpt<T, ParametersT>::Execute(const FeedsFetchesManager* init_
                                       this->context_.Logger(),
                                       this->ort_stream_,
                                       /*sync_subgraph_fetches*/ false,
-                                      this->context_.GetRunProfiler());
+                                      this->context_.GetRunProfiler(),
+                                      this->context_.GetRuntimeWorkloadCollector());
     }
 
     ORT_RETURN_IF_ERROR(status);