diff --git a/onnxruntime/core/framework/kernel_registry_manager.cc b/onnxruntime/core/framework/kernel_registry_manager.cc index 721353854a474..f9c82c0495916 100644 --- a/onnxruntime/core/framework/kernel_registry_manager.cc +++ b/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& 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(), diff --git a/onnxruntime/core/framework/session_state.cc b/onnxruntime/core/framework/session_state.cc index 5d71011197378..ebd6592a37ebb 100644 --- a/onnxruntime/core/framework/session_state.cc +++ b/onnxruntime/core/framework/session_state.cc @@ -376,6 +376,11 @@ const std::unordered_map& SessionState::GetConstantInitializedTen return constant_initialized_tensors_; } +const std::unordered_map& 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_stringImplicitInputDefs()) { + 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) { @@ -1815,6 +1875,19 @@ Status SessionState::FinalizeSessionStateImpl(const std::basic_string #include #include +#include #include #include @@ -166,6 +167,16 @@ class SessionState { */ const std::unordered_map& 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& GetConstantInitializedTensorsForKernelCreation() const; + const PrepackedWeightsForGraph& GetPrepackedIniitializersForGraph() const; #if !defined(DISABLE_SPARSE_TENSORS) @@ -494,6 +505,20 @@ class SessionState { std::unordered_map initialized_tensors_; // key is ort_value_index // subset of initialized_tensors_ that are constant and cannot be overridden at runtime std::unordered_map 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 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 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 diff --git a/onnxruntime/test/contrib_ops/matmul_4bits_test.cc b/onnxruntime/test/contrib_ops/matmul_4bits_test.cc index 2bad749cdd528..b69e9241b101d 100644 --- a/onnxruntime/test/contrib_ops/matmul_4bits_test.cc +++ b/onnxruntime/test/contrib_ops/matmul_4bits_test.cc @@ -4,6 +4,7 @@ #ifndef ORT_MINIMAL_BUILD #include +#include #include "gtest/gtest.h" #include "gmock/gmock.h" @@ -11,6 +12,7 @@ #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 b_f(static_cast(N * K), 0.1f); + std::vector b_quant(static_cast(N * nblk * blob)); + std::vector scales(static_cast(N * nblk)); + QuantizeDequantize(b_f, b_quant, scales, nullptr, + static_cast(N), static_cast(K), static_cast(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 a_vals(static_cast(M * K), 1.0f); + OrtValue a_val; + CreateMLValue(TestCPUExecutionProvider()->CreatePreferredAllocators()[0], + {M, K}, a_vals, &a_val); + + OrtValue cond_val; + CreateMLValue(TestCPUExecutionProvider()->CreatePreferredAllocators()[0], + {}, {true}, &cond_val); + + NameMLValMap feeds = {{"A", a_val}, {"cond", cond_val}}; + std::vector output_names = {"Yout"}; + std::vector 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(); + EXPECT_EQ(result.Shape()[0], M); + EXPECT_EQ(result.Shape()[1], N); +} + } // namespace test } // namespace onnxruntime