Matching variadic lists of tensors?

[iksnagreb]

If I understand correctly it is currently not possible to match a pattern with a variadic list of output tensors, such as Split, and similarly matching variadic inputs, such as Concat, seems only possible via some workarounds (e.g., #2373)? Currently it seems possible to match a Split only if you know the number of outputs in advance and it is not possible to match against the _outputs parameter itself. This would actually allow for some workaround instantiating a set of rules sweeping the _outputs argument with some loop, i.e., explicitly instantiating every variadic version of the operator, but this does not seem sustainable as the ONNX spec allows for up to 2147483647 such inputs or outputs blowing up the number of rewrite rules quite a lot...

So, is there any plan to add proper support for variadic inputs/outputs to the rewriter/pattern matcher in the near future or any pointers on how I could approach implementing this?

[justinchuby]

@gramalingam

[justinchuby]

I think you can use _allow_other_input for variations inputs. https://microsoft.github.io/onnxscript/tutorial/rewriter/rewrite_patterns.html#pattern-options

[iksnagreb]

Hm, would it be possible/useful to create something similar for the outputs as well?

[justinchuby]

Could you share an example when you want to match variadic outputs?

[iksnagreb]

The concrete example I am working on right now involves Split operators preceded by an element-wise operation, such as Add. If this element-wise operation involves a constant tensor (that is not strictly necessary, but that is the situation where it makes the most sense), the order of Split and the element-wise operation can be switched by splitting this constant tensor and distributing the operation into the output branches of the Split. Writing this down for a specific instance of this pattern, i.e., when the number and constantness of the inputs/outputs is known, is kind off easy, generalizing to arbitrary element-wise operators is also doable (either by enumerating or checking against some list of known/supported op-types, as there are few enough so this does not blow up the number of rewrite rules). But, I cannot figure out how to extend this to an unknown/arbitrary number of outputs (enumerating all rules is not really feasible in this case, unless you cut it off at some reasonably small number of allowed outputs). To be clear: It is not a problem to get the number of outputs for the match condition (check) or replacement pattern (rewrite), when you already have the match, the problem is getting a match in the first place without knowing how many outputs to expect. If you for example return op.Split(...) from the pattern without specifying the _outputs, it will try to match exactly one output, no matter how the split is actually configured.

By the way, there seems to be an issue here https://github.com/microsoft/onnxscript/blob/main/onnxscript/rewriter/_matcher.py#L181 (that is actually my fix from a few weeks ago, but I think there could be similar instance where this could go wrong): Returning False does not really fail the match but simply breaks out of the matching, sometimes allowing a partial match to be accepted. Funnily, this currently allows to abuse the _outputs on the match pattern to specify matching up to _outputs, however the rewrite pattern then expects this manually specified maximum number of outputs instead of the actually matched number causing a mismatch from which the rewrite cannot recover so it simply errors out... The following lines trying to bind the output probably suffer from the same problem where failing to bind an output could result in accepting a pattern with only partially bound values. Though, I am not sure if there is some post-processing to weed out partial matches later on, so maybe it is not really an issue. But probably matches should always be failed via self.fail to be sure.

[justinchuby]

Thanks for your explanation! I will look into this more closely. Do you have a suggestion on how the api/user experience should be?

[justinchuby]

@gramalingam