Pattern matching has_attr using any value

I need to create a pattern matching to match all qnn.conv2d operators which are depthwise convolution operators. I saw that the EthosU integration does this by checking the kernel_layout attribute using the has_attr function. My problem is, I did a layout conversion before so both conv2d and depthwise conv2d have the same kernel_layout attribute.

I noticed that only the depthwise convolution operators have the groups attribute, so I was trying to match them, but I couldn’t find a way of saying “match any qnn.conv2d with attribute groups, no matter the value of the attribute”. Using something like this:

is_op("qnn.conv2d")(...).has_attr({"groups": 512})

Works only for the qnn.conv2d operators which has the groups attribute set to 512. I was hoping I could do something like:

is_op("qnn.conv2d")(...).has_attr({"groups": wildcard()})

Or:

is_op("qnn.conv2d")(...).has_attr({"groups": is_constant()})

But sadly, none of them work.

Does anyone know a way to do this? Or perhaps another workaround?

So, I solved this in a ratter dirty way, so the question still remains open, but I will post my solution in case someone needs something similar.

Motivation

I need to do pattern matching to insert my custom operators for 2d convolutions and 2d depthwise convolutions.

Problem

I also need to convert the kernel layout of the 2d convolution in order to correctly schedule the convolution, so I used the ConvertLayout relay transformation. But I only need to do this for the 2d convolution, not for the depthwise convolution. But because both are represented with qnn.conv2d, both get converted with the layout transformation.

Current (dirty) solution

I have two pattern matching tables:

  • one that matches the 2d depthwise convolution operators using the trick used by the EthosU (looking for qnn.conv2d with kernel_layour equal to HWOI)
  • one that matches de 2d convolution and all other necesary operators.

First, I call the relay transformation MergeComposite using the depthwise convolution pattern matching table. This replaces the qnn.conv2d that are depthwise convolutions with my custom operator.

Then, I call the ConvertLayout transformation, which modifies the kernel_layout of the remaining qnn.conv2d operators.

Finally, I call the relay transformation MergeComposite with my second pattern matching table, replacing all other operators with my custom ones.