How to add my custom Relay node to pattern matcher? [EthosU example]

Questions
#1

Hello,

I was looking at the Arm EthosU integration in TVM and noticed that there was a new conv2d Relay operator defined. Obviously this operator is only legal/valid for offloading onto the EthosU, so they decide to replace standard Relay operator patterns inside of a Composite with this single operator. That makes it easier to define their own Relay->Tir pipeline since they only need to parse one node and all its parameters/attributes.

I was wondering though, how could this (or any) new Relay operator be added into the pattern matching infrastructure?

I guess what I would like is the following

#Original Relay module using standard Relay operators
%1 = standard_relay_op_1(...)
%2 = standard_relay_op_2(...)

#Assume I have some custom Relay operators that each match to one of those from above
%1 = my_custom_relay_op_1(...)
%2 = my_custom_relay_op_2(...)

#Now, I further have a third Relay operator which matches exactly to the pattern above
%1 = my_custom_rela_op_3(....)

I know in this simple case there is a solution using the already available pattern matcher infrastructure (i.e. defining the pattern to be match of my_custom_relay_op_3 before the others in the pattern list). But let’s say for more complex situations I couldn’t just solve it via position in the pattern list.

How do I add my_custom_relay_op_{1,2} to the pattern matching infrastructure such that I can define the pattern matching for my_custom_relay_op_3 to have statements/expression like is_op('my_custom_relay_op_1') ?

Thank you all

Friendly mention @comaniac :slight_smile:

#2

@manupa-arm @mbaret

So maybe I can ask more directly.

Is the ordering of first pattern matching your offloadable and then replacing, within the extracted composite, the native relay operators with your new ethosu.conv2d relay operator a solution to not being able to do what I said before?

The following would feel like a simpler pipeline

  1. Find and replace in the main Relay module the native Relay op pattern which describes your ethosu.<operator> (just like you did it in your legalization process)
  2. PartitionGraph given that ethosu.<operator> is offloadable
  3. Relay-TIR now works like you already have it
#3

Hi,

Apologies for the delayed reply. I don’t entirely follow your question on pattern matching - it seems like what you’re talking about there is already possible.

Regarding the query about the Ethos-U partitioning/legalizing flow, I think the two routes are essentially equivalent. We partition first (via composite functions) and then legalize, but I agree it would be equally valid to legalize first and then partition. In terms of simplicity, I think both are just as simple. There are fundamentally two pieces of logic that need implementing - something to check whether a given pattern is supported by your offload target, and something to mutate the Relay into your target’s Relay ‘dialect’.

If you legalize first, then you must do both of these things in the same legalization pass. If you partition first, then you test ‘supportedness’ in the partitioner and subsequently do the mutation in the partitioned module. The code that needs writing is about the same, just positioned differently.

#4

@cron I think you’re asking about how to express a situation where a subgraph of two adjacent Relay operators may match two patterns in the graph partitioning logic. could you elaborate on the exact problem you’re having? I think the ordering logic should be sufficient to handle this, but it’d be good to understand your use case a bit further.

It’s worth mentioning that in the Ethos-U case, there are additional reasons to legalize into Ethos-U-specific Relay operators. Specifically, information is needed further down the compilation pipeline from that entire subgraph, and it’s only valid to drive the Ethos-U compilation pipeline with certain patterns of Relay operators and with certain arguments to those operators. So if I’m understanding your original question correctly, you could view the Ethos-U legalization as more of a sanity-check or input validation step rather than a workaround for a pattern-matcher limitation.

also cc @mbs-octoml @mbrookhart

#5

cc @mikepapadim, who’s reworking FuseOps to be more Composite+patterns like.

#6

I’m also a bit confused about your question, can you clarify what custom_relay_op_3 is, and its relationship to custom_relay_op_1 and custom_relay_op_2?