[PR] Shape Inference for Fused Nodes

[pranav-prakash]

Is your feature request related to a problem? Please describe.
Currently ORT does not set any shape inference function for fused nodes. That is, we see in the createSchema function:

static std::unique_ptr<ONNX_NAMESPACE::OpSchema> CreateSchema(const Graph& graph,
                                                              const IndexedSubGraph& nodes_to_fuse) {
  const auto* meta_def = nodes_to_fuse.GetMetaDef();
  auto op_schema = onnxruntime::make_unique<ONNX_NAMESPACE::OpSchema>();
  op_schema->SetName(meta_def->name);
  op_schema->SetDomain(meta_def->domain);
  op_schema->SetDoc(meta_def->doc_string);
  op_schema->SinceVersion(meta_def->since_version);
  int i = 0;

  for (auto& input : meta_def->inputs) {
    auto input_arg = graph.GetNodeArg(input);
    // inputs must have a type. can be inferred for outputs.
    ORT_ENFORCE(input_arg->Type() != nullptr);
    op_schema->Input(i, input, "", *input_arg->Type());
    ++i;
  }
  i = 0;
  for (auto& output : meta_def->outputs) {
    auto output_arg = graph.GetNodeArg(output);
    op_schema->Output(i, output, "", *output_arg->Type());
    ++i;
  }
  op_schema->Finalize();

  return op_schema;
}

that because TypeAndShapeInferenceFunction is not explicitly set, it falls back to the default. While this doesn't cause an issue for inference, it prevents the allocator from reusing buffers since it has less information to work with. Instead for these fused nodes I think it would be better to define a shape inference function that is the composition of the inference functions for the nodes of the subgraph to be fused. For instance, consider a simple example of fusing a Conv + Relu node. The shape inference for this fused subgraph would run shape inference for Conv, take the resulting output TensorShape, and feed that as input to shape inference for Relu.

Implementing this might be a little bit involved because ONNX seems to require all shape inference to be mediated by an InferenceContext. That is, one cannot directly pass in input TensorShapeProtos to the shape inference function and get an output TensorShapeProto. Thus to implement, I suppose you'll have to subclass InferenceContext so that we have direct access to the TensorShapeProtos. Then the fused shape inference function can be defined by simulating the inference via our subclassed InferenceContext.

At the very least, an option should be exposed to the user to manually define a shape inference function in the meta_def which will be applied to the op_schema created for the fused node.

EDIT: I submitted a PR (#7007) to at least allow a custom shape inference function to be set. And now that I think about it more, letting the fused inference function be the composition of the individual nodes' inference functions only works if all attributes from each node are present for the fused node – so we'd want to allow the inference function to be manually set anyway.

[skottmckay]

Do you have a model that reproduces the issue?

AFAIK there are two paths to a Function based node.

1. An ONNX function based operator is expanded to its component nodes as we don't have an implementation of the function.
2. Nodes are fused by an execution provider into a single function based node during graph partitioning.

For #1, the code that expands the ONNX FunctionProto sets the type/shape inferencing function here:

onnxruntime/onnxruntime/core/graph/function.cc

Lines 322 to 330 in 368e4a3

	if (!cached_op_schema || !cached_op_schema->has_type_and_shape_inference_function()) {
	op_schema_->TypeAndShapeInferenceFunction(
	[this](ONNX_NAMESPACE::InferenceContext& ctx) {
	auto schema_registry = ONNX_NAMESPACE::OpSchemaRegistry::Instance();
	ONNX_NAMESPACE::shape_inference::InferShapeForFunctionNode(&onnx_func_proto_, body_.MainGraph().DomainToVersionMap(), schema_registry, ctx);
	});
	} else {
	op_schema_->TypeAndShapeInferenceFunction(cached_op_schema->GetTypeAndShapeInferenceFunction());
	}

For #2, the ONNX type/shape inferencing would have already happened prior to fusion (shape inferencing is run during model loading, which completes before partitioning), so the output values from the fusion should already have type/shape info.

[pranav-prakash]

Scott McKay (@skottmckay) We make use of #2 in our custom EP (link to the node-fusion we define in our EP)

ONNX type/shape inferencing would have already happened prior to fusion

Is this still true if Level 3 graph transforms are performed? I observed that these are performed after partitioning, and this was the case motivating the creation of the PR – I recall that the outputs of fused ops had no shape info after the shape-inference pass following level 3 transforms (but it's been a few months so I might need to double check)

[pranav-prakash]

I looked again and the sequence of steps that reproduces this issue is this

• After partitioning, we create a fused node for e.g. QLinearAdd + Relu in our EP
• A level 3 graph transformer (e.g. NHWC conversion) is registered that replaces the inputs and outputs of the fused node.
• Since these are newly created NodeArgs, the previous shape inference info is no longer present. Hence when graph.Resolve() is called and shape-inference is again performed, since the fused node lacks any shape inference function we end up with no shape on the output (which prevents shape info from being computed for any downstream nodes as well).

[skottmckay]

That makes more sense. I added some debug output and ran one of the unit tests that does some node fusion (InternalTestingEP.TestSaveAndLoadOrtModel) and the shape info was maintained as expected in that case.

Tracy Sharpe (@tracysh) Can the NHWC transformer add shape info to newly created NodeArgs when it runs?

[pranav-prakash]

Scott McKay (@skottmckay)
(We use our own nhwc transformer – not the upstream one – but last I checked the upstream one doesn't set shape info for nodeargs either).

So just to clarify, is the suggested solution to set shape info in the transformer itself when processing a fused node? That works for me (but it feels more like a workaround – I still feel allowing the inference_function to be manually set when creating a fused node is a cleaner solution; albeit either achieves the same end result).

[skottmckay]

I think the expected result of fusing nodes is that the same output names are still produced, and hence no need for special handling of the shape info.

I assume if you're creating new NodeArg instances you're also having to update downstream nodes to change input names as well, and that's getting beyond simply fusing nodes.

[skottmckay]

FWIW I'm not against adding a shape inferencing function to the onnxruntime::Function based node if there's a clear need for it. Just because it doesn't fit the normal node fusion scenarios doesn't mean it wouldn't be valuable, however there should be a unit test and clear explanation of when and how it should be used in a PR that adds it.

[pranav-prakash]

Scott McKay (@skottmckay)
I'd be happy to update the PR to add the unit test.

As for applicable use-cases: even though strictly speaking performing shape-inference for fused nodes by calling into a custom function during graph.Resolve() will have the same end-result as manually setting shape info during graph transforms, doing the latter leads to a weird asymmetry between processing "normal" nodes and "fused ones" – graph transforms for the former can be agnostic to the input/output shapes, whereas transforms for the latter must explicitly set shape info.

Given that it requires < 4 lines of code to add support for this, I think it's worth adding it so graph transforms can be kept as simple as possible.

Edit: I've updated the PR if you'd like to take a look.