[atenlib] `trace_only=True` and onnx local function

[BowenBao]

Due to various limitations, many more aten lib operators are authored with trace_only=True, e.g. #419.

Consider constructing those as local function, to preserve exact signature as original aten call.

[justinchuby]

Does it mean we will have many variants of the same function? How should these functions be named and managed? Would an alternative of just preserving the pre-transform arguments be possible?

[BowenBao]

Does it mean we will have many variants of the same function?

Yes.

How should these functions be named and managed?

I think they are similar to overloads (edit: atenlib overloads, not overloads of aten op itself), differs in they are created on the fly and are scoped locally under the model. It should contain the full original aten qualified name, with means (e.g., postfix) to be uniquely identifiable in the scope of the model. I'm not worried about duplication. As long as the signature is accurate, backend can choose to replace them all with the same kernel. It should help with fusion too.

Would an alternative of just preserving the pre-transform arguments be possible?

The motivation is to preserve exact same signature with aten. The internal local function subgraph may be using entirely different attribute values, which are computed during export from the trace_only=True parts. And that is ok.
If the aten signature, both arguments and attributes and their order, can be preserved by alternative methods, then no we don't need local function.

[justinchuby]

I've been thinking about the kind of guarantees we are offering with these functions. Is the goal of preserving the signature for human readability, or machine recognition?

If the former, I wonder if there are ways to represent that information in the text area of a node, or just store them as unused attributes, so that people can see. We generally strive to contain an aten op within a function even when trace_only is true. So that plain onnx nodes in the main graph is hopefully rare.

If it means the latter, I am personally cautious about the potential for any implication to optimizers/runtimes that a function is defined by its signature alone and they can substitute the function solely by looking at its signature, because that may require the functions to be versioned/stable. (1) We know that aten is not versioned or stable and (2) this may imply that we now maintain an opset (aten), which may not be a sustainable thing to do, imo.

If we do not make those guarantees however, runtime cannot use the signature as a robust mechanism for matching, and should look at the underlying graph instead. The benefit of the function here, I think, is that we can provide a clear boundary for optimization, and offer a relatively stable pattern for machines to match.

Curious on your thoughts?

[BowenBao]

human readability, or machine recognition?

Both. Even for human readability I feel it's better when nodes are all in the same style.

I am personally cautious about the potential for any implication to optimizers/runtimes that a function is defined by its signature alone and they can substitute the function solely by looking at its signature, because that may require the functions to be versioned/stable. (1) We know that aten is not versioned or stable and (2) this may imply that we now maintain an opset (aten), which may not be a sustainable thing to do, imo.

Hmm I think the motivation is only to elevate the trace_only=True ones to have the same minimal surface area as those trace_only=False ones, where the signature is very similar to aten, if not identical. The function subgraph is always there to serve as reference. But yea there is this potential implication that backends substitutes solely based on signatures, but this argument applies to all atenlib ops in general.

Agree with rest of your comment on the benefit of function.

[justinchuby]

Good point. Lets talk more

[gramalingam]

We had a similar discussion in an earlier weekly meeting, focusing primarily on machine recognition.

I agree that, as Justin says, a robust matching mechanism would be better. There are some interesting possibilities we can explore: (i) a runtime graph-equivalence matcher, or (ii) some kind of sophisticated hashing algorithm that is invariant under such a graph match, etc.

However, realistically, many backend-implementations want to take the simpler route of name-based matching. Of course, it is up to a library-designer to decide whether the names they use come with a guarantee/contract (about its specification). It is part of the library's contract. ONNX does allow the library-designer to use versioning if they want.

Anyway: one suggestion, from an ONNX perspective, is to add a simple "string overload" field to both NodeProto and FunctionProto to allow multiple overloaded variants of the same function to exist in a model/library.

[justinchuby]

Anyway: one suggestion, from an ONNX perspective, is to add a simple "string overload" field to both NodeProto and FunctionProto to allow multiple overloaded variants of the same function to exist in a model/library.

Thinking about this again, I agree this is a nice solution for us to not have to manage function names. E.g. "What do you call a aten_transpose function with rank=4? aten_transpose_rank_4?" and more complex cases where we have more attributes or sequences as attributes.

However, realistically, many backend-implementations want to take the simpler route of name-based matching.

I wonder if there is a way for onnx to make the pattern matching robust and simple enough so name matching isn't as much appealing. I just think that providing a contract on names from the library's perspective places a lot of constraints on how we can evolve the implementation. My current feeling is that it's not a wise path to take, unless we are willing to formally offer an aten opset

I re-thought about this as I wrote - maybe this isn't as bad as I thought, because we will (1) need to match pytorch anyways (2) pytorch isn't updated frequently. So we might as well version the ops

[titaiwangms]

cc @abock