Fix segfault in MatMulNBits prepacking when weights/scales are parent-graph initializers in subgraph

onnxruntime/core/framework/kernel_registry_manager.cc

@@ -21,7 +21,7 @@ Status KernelRegistryManager::CreateKernel(const Node& node,
                                            const KernelCreateInfo& kernel_create_info,
                                            std::unique_ptr<OpKernel>& out) const {
   OpKernelInfo kernel_info(node, *kernel_create_info.kernel_def, execution_provider,
-                           session_state.GetConstantInitializedTensors(),
+                           session_state.GetConstantInitializedTensorsForKernelCreation(),
                            session_state.GetOrtValueNameIdxMap(),
                            session_state.GetDataTransferMgr(),
                            session_state.GetAllocators(),

onnxruntime/core/framework/session_state.cc

@@ -376,6 +376,11 @@ const std::unordered_map<int, OrtValue>& SessionState::GetConstantInitializedTen
   return constant_initialized_tensors_;
 }
 
+const std::unordered_map<int, OrtValue>& SessionState::GetConstantInitializedTensorsForKernelCreation() const {
+  return outer_scope_augmented_map_built_ ? outer_scope_augmented_constant_tensors_
+                                          : constant_initialized_tensors_;
+}
+
 const PrepackedWeightsForGraph& onnxruntime::SessionState::GetPrepackedIniitializersForGraph() const {
   return graph_.GetPrepacked();
 }
@@ -1753,6 +1758,61 @@ Status SessionState::FinalizeSessionStateImpl(const std::basic_string<PATH_CHAR_
     CleanInitializedTensorsFromGraph();
   }
 
+  // For subgraph session states: build a merged constant-tensor map (outer_scope_augmented_constant_tensors_)
+  // that extends this subgraph's own constant_initialized_tensors_ with any outer-scope constant initializers
+  // from parent graphs, re-indexed into this subgraph's ort_value_name_idx_map_. This allows
+  // TryGetConstantInput (called from kernel constructors and PrePack) to resolve constants that live in
+  // parent-graph scope, such as the scales/zero_points paired with a MatMulNBits B tensor when all three
+  // initializers belong to the outer graph but the MatMulNBits node is inside a subgraph (e.g. If branch).
+  // The augmented map is only used for kernel creation; PrepackConstantInitializedTensors keeps using the
+  // unaugmented constant_initialized_tensors_ to avoid double-prepacking outer-scope tensors.
+  // After the subgraph finalization loop below, the parent-scope OrtValue copies in this map are erased
+  // (via outer_scope_parent_only_indices_) to allow the parent's constant_initialized_tensors_ to release
+  // memory once all prepack use counts reach zero.
+  if (parent_node != nullptr && parent_ != nullptr) {
+    outer_scope_augmented_constant_tensors_ = constant_initialized_tensors_;
+
+    for (const NodeArg* outer_arg : parent_node->ImplicitInputDefs()) {
+      const std::string& name = outer_arg->Name();
+
+      // Get the index for this name in the current (subgraph) scope.
+      int current_idx = -1;
+      if (!ort_value_name_idx_map_.GetIdx(name, current_idx).IsOK()) {
+        continue;
+      }
+
+      // Skip if this name is already covered by the subgraph's own constants.
+      if (constant_initialized_tensors_.count(current_idx) > 0) {
+        continue;
+      }
+
+      // Walk the parent session state chain (handles deeply nested subgraphs) until we either
+      // find the constant or exhaust all ancestors.
+      const SessionState* p = parent_;
+      while (p != nullptr) {
+        int parent_idx = -1;
+        if (!p->ort_value_name_idx_map_.GetIdx(name, parent_idx).IsOK()) {
+          p = p->parent_;
+          continue;
+        }
+
+        // Use the parent's already-augmented map so that constants from grandparent scopes
+        // are also visible (the parent's FinalizeSessionStateImpl has already run).
+        const auto& parent_const = p->GetConstantInitializedTensorsForKernelCreation();
+        auto it = parent_const.find(parent_idx);
+        if (it != parent_const.end()) {
+          outer_scope_augmented_constant_tensors_.emplace(current_idx, it->second);
+          outer_scope_parent_only_indices_.insert(current_idx);
+          break;
+        }
+
+        p = p->parent_;
+      }
+    }
+
+    outer_scope_augmented_map_built_ = true;
+  }
+
   ORT_RETURN_IF_ERROR(CreateKernels(kernel_registry_manager));
 
   if (!disable_prepacking) {
@@ -1763,6 +1823,7 @@ Status SessionState::FinalizeSessionStateImpl(const std::basic_string<PATH_CHAR_
   ORT_RETURN_IF_ERROR(
       session_state_utils::SaveInputOutputNamesToNodeMapping(*graph_viewer_, *this, valid_outer_scope_node_args));
 
+
   // Need to recurse into subgraph session state instances to finalize them and add the execution info
 
   // Currently all subgraphs need to be executed using the sequential EP due to potential deadlock with the current
@@ -1815,6 +1876,19 @@ Status SessionState::FinalizeSessionStateImpl(const std::basic_string<PATH_CHAR_
     // locations for these would be the locations they are explicitly consumed on in nested subgraphs.
   }
 
+  // Release the extra OrtValue copies from parent scope that were added to the augmented constant
+  // map for this subgraph session state. These copies were needed to let TryGetConstantInput resolve
+  // parent-scope constants during CreateKernels and PrepackConstantInitializedTensors above, and to
+  // let any nested (grandchild) subgraph session states read them via GetConstantInitializedTensorsForKernelCreation().
+  // Now that all subgraphs are fully finalized, the copies are no longer needed. Releasing them lets
+  // the parent's constant_initialized_tensors_ free the underlying tensor buffers once all prepack
+  // use counts reach zero, avoiding a memory regression for large initializers such as quantized B
+  // matrices and scales in MatMulNBits subgraph patterns.
+  for (int idx : outer_scope_parent_only_indices_) {
+    outer_scope_augmented_constant_tensors_.erase(idx);
+  }
+  outer_scope_parent_only_indices_.clear();
+
   return Status::OK();
 }
 

onnxruntime/core/framework/session_state.h

@@ -6,6 +6,7 @@
 #include <memory>
 #include <map>
 #include <unordered_map>
+#include <unordered_set>
 #include <string>
 #include <vector>
 
@@ -166,6 +167,16 @@ class SessionState {
    */
   const std::unordered_map<int, OrtValue>& GetConstantInitializedTensors() const;
 
+  /**
+   * Gets the constant initialized tensors for use during kernel creation (constructor and PrePack).
+   * For subgraph session states this extends GetConstantInitializedTensors() with outer-scope
+   * constant initializers from parent graphs, re-indexed to the current subgraph's value map.
+   * This allows TryGetConstantInput to resolve parent-scope constants inside kernel constructors
+   * and PrePack (e.g. MatMulNBits packing scales alongside B when both live in the parent graph).
+   * For top-level graphs this is identical to GetConstantInitializedTensors().
+   */
+  const std::unordered_map<int, OrtValue>& GetConstantInitializedTensorsForKernelCreation() const;
+
   const PrepackedWeightsForGraph& GetPrepackedIniitializersForGraph() const;
 
 #if !defined(DISABLE_SPARSE_TENSORS)
@@ -494,6 +505,20 @@ class SessionState {
   std::unordered_map<int, OrtValue> initialized_tensors_;  // key is ort_value_index
   // subset of initialized_tensors_ that are constant and cannot be overridden at runtime
   std::unordered_map<int, OrtValue> constant_initialized_tensors_;
+  // For subgraph session states: constant_initialized_tensors_ extended with outer-scope constant
+  // initializers (from parent graphs) re-indexed to this subgraph's ort_value_name_idx_map_.
+  // Populated in FinalizeSessionStateImpl before CreateKernels; empty for top-level graphs.
+  // Returned by GetConstantInitializedTensorsForKernelCreation() so kernel constructors and
+  // PrePack can resolve parent-graph constants via TryGetConstantInput.
+  // After the subgraph finalization loop, the parent-scope copies in this map are erased
+  // (tracked via outer_scope_parent_only_indices_) to allow parent memory to be freed once
+  // prepack use counts reach zero.
+  std::unordered_map<int, OrtValue> outer_scope_augmented_constant_tensors_;
+  // Indices in outer_scope_augmented_constant_tensors_ that were added from parent scope
+  // (not present in constant_initialized_tensors_).  Erased from the augmented map after
+  // all subgraphs are finalized to release the extra OrtValue refcounts.
+  std::unordered_set<int> outer_scope_parent_only_indices_;
+  bool outer_scope_augmented_map_built_{false};
 
 #if !defined(DISABLE_SPARSE_TENSORS)
   // This is an auxiliary lookup to check if the OrtValue was actually a sparse tensor

onnxruntime/test/contrib_ops/matmul_4bits_test.cc

@@ -4,13 +4,15 @@
 #ifndef ORT_MINIMAL_BUILD
 
 #include <optional>
+#include <sstream>
 
 #include "gtest/gtest.h"
 #include "gmock/gmock.h"
 
 #include "core/common/narrow.h"
 #include "core/common/span_utils.h"
 #include "core/framework/tensor.h"
+#include "core/graph/onnx_protobuf.h"
 #include "core/mlas/inc/mlas_qnbit.h"
 #include "core/mlas/inc/mlas_q4.h"
 #include "core/mlas/inc/mlas.h"
@@ -1754,6 +1756,181 @@ TEST(MatMulNBits, PrePack_LegacyFlattenedShapes_Accepted) {
            {}, nullptr, &execution_providers);
 }
 
+// Regression test for https://github.com/microsoft/onnxruntime/issues/31137
+// MatMulNBits inside an If subgraph where B and scales are parent-graph initializers.
+// On ARM64 with accuracy_level=4 (KleidiAI path) this triggered a segfault during
+// session initialization because TryGetConstantInput could not find parent-scope
+// constants when building the kernel's OpKernelInfo.
+TEST(MatMulNBits, SubgraphParentScopeInitializers) {
+#if !defined(MLAS_TARGET_ARM64)
+  GTEST_SKIP() << "This test targets the Arm64 KleidiAI path.";
+#else
+  if (!MlasQNBitGemmScalesPacked(64, QBits, 32, SQNBIT_CompInt8, true, nullptr)) {
+    GTEST_SKIP() << "KleidiAI Q4 packed-scales path is not active.";
+  }
+#endif
+
+  constexpr int64_t M = 4, K = 64, N = 64, BLK = 32, BITS = 4;
+  constexpr int64_t nblk = K / BLK;
+  constexpr int64_t blob = BLK * BITS / 8;
+
+  // Generate quantized B and scales.
+  std::vector<float> b_f(static_cast<size_t>(N * K), 0.1f);
+  std::vector<uint8_t> b_quant(static_cast<size_t>(N * nblk * blob));
+  std::vector<float> scales(static_cast<size_t>(N * nblk));
+  QuantizeDequantize(b_f, b_quant, scales, nullptr,
+                     static_cast<int32_t>(N), static_cast<int32_t>(K), static_cast<int32_t>(BLK));
+
+  // Build the If-branch subgraph containing a MatMulNBits node.
+  // The branch consumes A, Bq, Bs from outer scope and outputs Y.
+  auto build_branch = [&](const std::string& branch_name) -> ONNX_NAMESPACE::GraphProto {
+    ONNX_NAMESPACE::GraphProto g;
+    g.set_name(branch_name);
+
+    auto* out_vi = g.add_output();
+    out_vi->set_name("Y");
+    auto* out_type = out_vi->mutable_type()->mutable_tensor_type();
+    out_type->set_elem_type(ONNX_NAMESPACE::TensorProto_DataType_FLOAT);
+    auto* out_shape = out_type->mutable_shape();
+    out_shape->add_dim()->set_dim_value(M);
+    out_shape->add_dim()->set_dim_value(N);
+
+    auto* node = g.add_node();
+    node->set_op_type("MatMulNBits");
+    node->set_domain("com.microsoft");
+    node->set_name(branch_name + "_mmn");
+    node->add_input("A");
+    node->add_input("Bq");
+    node->add_input("Bs");
+    node->add_output("Y");
+
+    auto add_int_attr = [&](const std::string& attr_name, int64_t val) {
+      auto* attr = node->add_attribute();
+      attr->set_name(attr_name);
+      attr->set_type(ONNX_NAMESPACE::AttributeProto_AttributeType_INT);
+      attr->set_i(val);
+    };
+    add_int_attr("K", K);
+    add_int_attr("N", N);
+    add_int_attr("bits", BITS);
+    add_int_attr("block_size", BLK);
+    // accuracy_level=4 triggers the KleidiAI / high-accuracy prepack path on ARM64
+    // that asserts scales != nullptr; this is the crashing case in the bug report.
+    add_int_attr("accuracy_level", 4);
+
+    return g;
+  };
+
+  // Build the parent model: cond + A as inputs, Bq + Bs as parent-graph initializers,
+  // an If node that dispatches to MatMulNBits in both branches.
+  ONNX_NAMESPACE::ModelProto model;
+  model.set_ir_version(9);
+
+  auto* opset_default = model.add_opset_import();
+  opset_default->set_version(17);
+  auto* opset_ms = model.add_opset_import();
+  opset_ms->set_domain("com.microsoft");
+  opset_ms->set_version(1);
+
+  auto* graph = model.mutable_graph();
+  graph->set_name("main");
+
+  // Graph inputs.
+  {
+    auto* cond_vi = graph->add_input();
+    cond_vi->set_name("cond");
+    cond_vi->mutable_type()->mutable_tensor_type()->set_elem_type(
+        ONNX_NAMESPACE::TensorProto_DataType_BOOL);
+
+    auto* a_vi = graph->add_input();
+    a_vi->set_name("A");
+    auto* a_type = a_vi->mutable_type()->mutable_tensor_type();
+    a_type->set_elem_type(ONNX_NAMESPACE::TensorProto_DataType_FLOAT);
+    auto* a_shape = a_type->mutable_shape();
+    a_shape->add_dim()->set_dim_value(M);
+    a_shape->add_dim()->set_dim_value(K);
+  }
+
+  // Parent-graph initializers (Bq and Bs live here, NOT inside the subgraph).
+  {
+    auto* bq_init = graph->add_initializer();
+    bq_init->set_name("Bq");
+    bq_init->set_data_type(ONNX_NAMESPACE::TensorProto_DataType_UINT8);
+    bq_init->add_dims(N);
+    bq_init->add_dims(nblk);
+    bq_init->add_dims(blob);
+    bq_init->set_raw_data(b_quant.data(), b_quant.size() * sizeof(uint8_t));
+
+    auto* bs_init = graph->add_initializer();
+    bs_init->set_name("Bs");
+    bs_init->set_data_type(ONNX_NAMESPACE::TensorProto_DataType_FLOAT);
+    bs_init->add_dims(N * nblk);
+    bs_init->set_raw_data(scales.data(), scales.size() * sizeof(float));
+  }
+
+  // Graph output.
+  {
+    auto* out_vi = graph->add_output();
+    out_vi->set_name("Yout");
+    auto* out_type = out_vi->mutable_type()->mutable_tensor_type();
+    out_type->set_elem_type(ONNX_NAMESPACE::TensorProto_DataType_FLOAT);
+    auto* out_shape = out_type->mutable_shape();
+    out_shape->add_dim()->set_dim_value(M);
+    out_shape->add_dim()->set_dim_value(N);
+  }
+
+  // If node referencing Bq/Bs from parent scope inside both branches.
+  {
+    auto* if_node = graph->add_node();
+    if_node->set_op_type("If");
+    if_node->set_name("if0");
+    if_node->add_input("cond");
+    if_node->add_output("Yout");
+
+    auto* then_attr = if_node->add_attribute();
+    then_attr->set_name("then_branch");
+    then_attr->set_type(ONNX_NAMESPACE::AttributeProto_AttributeType_GRAPH);
+    *then_attr->mutable_g() = build_branch("then");
+
+    auto* else_attr = if_node->add_attribute();
+    else_attr->set_name("else_branch");
+    else_attr->set_type(ONNX_NAMESPACE::AttributeProto_AttributeType_GRAPH);
+    *else_attr->mutable_g() = build_branch("else");
+  }
+
+  std::string model_str;
+  ASSERT_TRUE(model.SerializeToString(&model_str));
+
+  // Session initialization must not crash/segfault.
+  SessionOptions so;
+  so.session_logid = "MatMulNBitsSubgraphTest";
+  InferenceSession session{so, GetEnvironment()};
+  std::istringstream model_stream(model_str);
+  ASSERT_STATUS_OK(session.Load(model_stream));
+  ASSERT_STATUS_OK(session.Initialize());
+
+  // Run inference and verify output shape is correct.
+  std::vector<float> a_vals(static_cast<size_t>(M * K), 1.0f);
+  OrtValue a_val;
+  CreateMLValue<float>(TestCPUExecutionProvider()->CreatePreferredAllocators()[0],
+                       {M, K}, a_vals, &a_val);
+
+  OrtValue cond_val;
+  CreateMLValue<bool>(TestCPUExecutionProvider()->CreatePreferredAllocators()[0],
+                      {}, {true}, &cond_val);
+
+  NameMLValMap feeds = {{"A", a_val}, {"cond", cond_val}};
+  std::vector<std::string> output_names = {"Yout"};
+  std::vector<OrtValue> fetches;
+  RunOptions run_options;
+  ASSERT_STATUS_OK(session.Run(run_options, feeds, output_names, &fetches));
+
+  ASSERT_EQ(fetches.size(), 1u);
+  const auto& result = fetches[0].Get<Tensor>();
+  EXPECT_EQ(result.Shape()[0], M);
+  EXPECT_EQ(result.Shape()[1], N);
+}
+
 }  // namespace test
 }  // namespace onnxruntime