Skip to content

Fix EP assignment reporting after NHWC layout transformation #33184

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
mirounga merged 1 commit into from
Oct 8, 2026
Merged

Fix EP assignment reporting after NHWC layout transformation #33184

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
26 changes: 21 additions & 5 deletions onnxruntime/core/framework/graph_partitioner.cc
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -204,6 +204,7 @@ struct GetCapabilityForEPParams {
std::reference_wrapper<const CheckLoadCancellationFn> check_load_cancellation_fn;
LayeringIndex* layering_index; // Added member
uint32_t registered_ep_context_data_callbacks;
InlinedHashSet<NodeIndex>* layout_assigned_nodes = nullptr;
};

auto get_capabilities = [](const IExecutionProvider& ep,
Expand Down Expand Up @@ -384,6 +385,10 @@ static Status GetCapabilityForEP(const GetCapabilityForEPParams& params, const l
for (auto& capability : capabilities) {
TryAssignNodes(graph, *capability->sub_graph, ep_type, &nodes_temporarily_assigned_to_ep);
}
if (params.layout_assigned_nodes != nullptr) {
params.layout_assigned_nodes->insert(nodes_temporarily_assigned_to_ep.begin(),
nodes_temporarily_assigned_to_ep.end());
}

const NodeIndex first_new_node = graph.MaxNodeIndex();

Expand Down Expand Up @@ -846,6 +851,10 @@ static Status GetCapabilityForEP(const GetCapabilityForEPParams& params, const l

for (NodeIndex idx = first_new_node; idx < end_node; ++idx) {
const Node* node = graph.GetNode(idx);
if (params.layout_assigned_nodes != nullptr && node != nullptr &&
node->GetExecutionProviderType() == ep_type) {
params.layout_assigned_nodes->insert(idx);
}
if (node != nullptr && node->Domain() == kMSInternalNHWCDomain) {
if (new_nodes_in_capabilities.count(node->Index()) == 0) {
return ORT_MAKE_STATUS(
Expand Down Expand Up @@ -898,12 +907,16 @@ static Status GetCapabilityForEPForAotInlining(const GraphViewer& graph_viewer,
static bool IsIndexedSubGraphAvailableForAssignment(Graph& graph,
const IndexedSubGraph& capability,
GraphPartitioner::Mode mode,
const std::string& provider_type) {
const std::string& provider_type,
const InlinedHashSet<NodeIndex>* layout_assigned_nodes = nullptr) {
// The provider can run a single node in the <graph> if not using meta-defs.
if (capability.GetMetaDef() == nullptr && capability.nodes.size() == 1) {
auto* node = graph.GetNode(capability.nodes[0]);
if (nullptr != node && node->GetExecutionProviderType().empty()) {
// The node was not fused or assigned.
if (nullptr != node &&
(node->GetExecutionProviderType().empty() ||
(node->GetExecutionProviderType() == provider_type && layout_assigned_nodes != nullptr &&
layout_assigned_nodes->contains(node->Index())))) {
// Accept provisional NHWC tags, but do not reprocess assignments from an earlier partitioning round.
return true;
}
return false;
Expand Down Expand Up @@ -1091,6 +1104,7 @@ static Status PartitionOnnxFormatModelImpl(Graph& graph, FuncManager& func_mgr,
// run the function by SessionOption, we should create a function kernel for it and
// delegate the compute to the functions inside the dlls.
std::vector<std::unique_ptr<ComputeCapability>> capabilities;
InlinedHashSet<NodeIndex> layout_assigned_nodes;
const auto get_capability_params = GetCapabilityForEPParams{
std::ref(graph),
std::cref(kernel_registry_mgr),
Expand All @@ -1103,7 +1117,8 @@ static Status PartitionOnnxFormatModelImpl(Graph& graph, FuncManager& func_mgr,
std::ref(graph_optimizer_registry),
std::cref(check_load_cancellation_fn),
layering_index,
registered_ep_context_data_callbacks};
registered_ep_context_data_callbacks,
&layout_assigned_nodes};

ORT_RETURN_IF_ERROR(GetCapabilityForEP(get_capability_params, logger));

Expand Down Expand Up @@ -1131,7 +1146,8 @@ static Status PartitionOnnxFormatModelImpl(Graph& graph, FuncManager& func_mgr,
// The <provider> can run a fused <sub_graph> in the <graph>.
// Check whether any node in the <sub_graph> was already assigned. If so it cannot be stolen as assignment is done
// in order of EP priority
bool sub_graph_available_for_assignment = IsIndexedSubGraphAvailableForAssignment(graph, *capability->sub_graph, mode, type);
bool sub_graph_available_for_assignment =
IsIndexedSubGraphAvailableForAssignment(graph, *capability->sub_graph, mode, type, &layout_assigned_nodes);

// If the <sub_graph> is available to be assigned to the EP and the ComputeCapability has nodes_to_optimize,
// run EP related optimizations and update ComputeCapability.
Expand Down
123 changes: 123 additions & 0 deletions onnxruntime/test/internal_testing_ep/internal_testing_partitioning_tests.cc
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -802,6 +802,129 @@ Status PartitionDirectAssignmentExternalEpContext(bool produces_ep_context_nodes

#define ORT_MODEL_FOLDER ORT_TSTR("testdata/")

#if !defined(DISABLE_CONTRIB_OPS)
TEST(InternalTestingEP, StaticKernelEpGraphAssignmentInfoAfterFunctionInlining) {
class OnnxOnlyStaticExecutionProvider : public InternalTestingExecutionProvider {
public:
explicit OnnxOnlyStaticExecutionProvider(DataLayout layout)
: InternalTestingExecutionProvider({}, {}, layout) {
EnableStaticKernels().TakeAllNodes();
}

std::vector<std::unique_ptr<ComputeCapability>> GetCapability(
const GraphViewer& graph_viewer, const IKernelLookup& kernel_lookup,
const GraphOptimizerRegistry& optimizer_registry, IResourceAccountant* accountant) const override {
auto capabilities = InternalTestingExecutionProvider::GetCapability(
graph_viewer, kernel_lookup, optimizer_registry, accountant);
capabilities.erase(std::remove_if(capabilities.begin(), capabilities.end(),
[&](const auto& capability) {
return graph_viewer.GetNode(capability->sub_graph->nodes[0])->Domain() != kOnnxDomain;
}),
capabilities.end());
return capabilities;
}
};

FunctionProto function;
function.set_domain("local");
function.set_name("IdentityFunction");
function.add_input("X");
function.add_output("Y");
function.add_opset_import()->set_version(13);
auto* body = function.add_node();
body->set_op_type("Identity");
body->add_input("X");
body->add_output("Y");

// AOT inlining covers one partitioning round; fallback inlining must not re-report the assigned Add.
for (DataLayout layout : {DataLayout::NCHW, DataLayout::NHWC}) {
for (bool disable_aot : {false, true}) {
SCOPED_TRACE(layout == DataLayout::NHWC ? "NHWC" : "NCHW");
SCOPED_TRACE(disable_aot ? "Fallback inlining" : "AOT inlining");
Model model("StaticKernelEpGraphAssignmentInfoAfterFunctionInlining", false,
ModelMetaData(), PathString(), IOnnxRuntimeOpSchemaRegistryList(),
{{kOnnxDomain, 13}, {"local", 1}}, {function},
DefaultLoggingManager().DefaultLogger());
Graph& graph = model.MainGraph();
ModelTestBuilder builder(graph);
auto* input_a = builder.MakeInput<float>(std::vector<int64_t>{3});
auto* input_b = builder.MakeInput<float>(std::vector<int64_t>{3});
auto* intermediate = builder.MakeIntermediate<float>(std::vector<int64_t>{3});
auto* output = builder.MakeOutput<float>(std::vector<int64_t>{3});
const std::string add_name = builder.AddNode("Add", {input_a, input_b}, {intermediate}).Name();
builder.AddNode("IdentityFunction", {intermediate}, {output}, "local");
builder.SetGraphOutputs();
ASSERT_STATUS_OK(graph.Resolve());

SessionOptions options;
options.graph_optimization_level = TransformerLevel::Default;
ASSERT_STATUS_OK(options.config_options.AddConfigEntry(kOrtSessionOptionsRecordEpGraphAssignmentInfo, "1"));
ASSERT_STATUS_OK(options.config_options.AddConfigEntry(kOrtSessionOptionsDisableCPUEPFallback, "1"));
ASSERT_STATUS_OK(options.config_options.AddConfigEntry(
kOrtSessionOptionsDisableAheadOfTimeFunctionInlining, disable_aot ? "1" : "0"));
InferenceSessionWrapper session(options, GetEnvironment());
ASSERT_STATUS_OK(session.RegisterExecutionProvider(std::make_unique<OnnxOnlyStaticExecutionProvider>(layout)));
const std::string model_bytes = model.ToProto().SerializeAsString();
ASSERT_STATUS_OK(session.Load(model_bytes.data(), static_cast<int>(model_bytes.size())));
ASSERT_STATUS_OK(session.Initialize());

const auto& subgraphs = session.GetEpGraphAssignmentInfo();
ASSERT_EQ(subgraphs.size(), 2u);
InlinedVector<std::string> assigned_ops;
for (const auto* subgraph : subgraphs) {
EXPECT_EQ(subgraph->ep_name, kInternalTestingExecutionProvider);
ASSERT_EQ(subgraph->nodes.size(), 1u);
const auto* node = subgraph->nodes[0];
EXPECT_EQ(node->domain, kOnnxDomain);
assigned_ops.push_back(node->op_type);
if (node->op_type == "Add") {
EXPECT_EQ(node->name, add_name);
}
}
EXPECT_THAT(assigned_ops, testing::UnorderedElementsAre("Add", "Identity"));
}
}
}

TEST(InternalTestingEP, NhwcTransformedStaticKernelEpGraphAssignmentInfo) {
Model model("NhwcTransformedStaticKernelEpGraphAssignmentInfo", false,
ModelMetaData(), PathString(), IOnnxRuntimeOpSchemaRegistryList(),
{{kOnnxDomain, 13}}, {}, DefaultLoggingManager().DefaultLogger());
Graph& graph = model.MainGraph();
ModelTestBuilder builder(graph);
auto* input = builder.MakeInput<float>(std::vector<int64_t>{1, 1, 3, 3});
auto* weights = builder.MakeInitializer<float>(std::vector<int64_t>{1, 1, 1, 1}, std::vector<float>{1.0f});
auto* output = builder.MakeOutput<float>(std::vector<int64_t>{1, 1, 3, 3});
builder.AddConvNode(input, weights, output);
builder.SetGraphOutputs();
ASSERT_STATUS_OK(graph.Resolve());

SessionOptions options;
options.graph_optimization_level = TransformerLevel::Default;
ASSERT_STATUS_OK(options.config_options.AddConfigEntry(kOrtSessionOptionsRecordEpGraphAssignmentInfo, "1"));
InferenceSessionWrapper session(options, GetEnvironment());
auto ep = std::make_unique<InternalTestingExecutionProvider>(
std::unordered_set<std::string>{"Conv"}, std::unordered_set<std::string>{}, DataLayout::NHWC);
ep->EnableStaticKernels();
ASSERT_STATUS_OK(session.RegisterExecutionProvider(std::move(ep)));
const std::string model_bytes = model.ToProto().SerializeAsString();
ASSERT_STATUS_OK(session.Load(model_bytes.data(), static_cast<int>(model_bytes.size())));
ASSERT_STATUS_OK(session.Initialize());

size_t conv_count = 0;
for (const auto* subgraph : session.GetEpGraphAssignmentInfo()) {
for (const auto* node : subgraph->nodes) {
if (node->op_type == "Conv") {
++conv_count;
EXPECT_EQ(subgraph->ep_name, kInternalTestingExecutionProvider);
EXPECT_EQ(node->domain, kMSInternalNHWCDomain);
}
}
}
EXPECT_EQ(conv_count, 1u);
}
#endif // !defined(DISABLE_CONTRIB_OPS)

auto RunTest(const std::string& op, const ORTCHAR_T* model_path) {
SessionOptions so;
auto session = std::make_unique<InferenceSessionWrapper>(so, GetEnvironment());
Expand Down