Error when not all operations can be lowered to xnnpack delegate

[mylifeasazucchini]

Similar to the issue #1169

I have a custom pytorch model that I am trying to lower to xnnpack delegate and I do as follows:

edge = export_to_edge( model, sample_inputs, edge_compile_config=EdgeCompileConfig( _check_ir_validity=True, ), )

which works as expected however when I run:

edge = edge.to_backend(XnnpackPartitioner)

I get the error:

_RuntimeError: For aten__native_batch_norm_legit_no_training_default, call_function:aten.native_batch_norm_legit_no_training.default is not supported in XNNPACK Delegate

I understood from the argument comments here that I can overcome this unsupported delegation by passing a dict of "supported" method and it's corresponding "supported" Partitioner.

I was wondering if there was any example of doing this and what the dict would look like ? to be precise what exactly should I pass in the "method" key of the dict.

Also, for the non lowerable part of the program how or what deleagte should I pass in the dict (if any!)

Thanks (sorry if there were any naive questions here, I am just learning this as we go!)

[mylifeasazucchini]

cc-ing for faster reach @angelayi since I have been following your tutorial on this!

[kimishpatel]

cc @digantdesai @mcr229 . For non lowerable part, it should be handled by partitioner by not partitionering it. I do see that batch_norm is not supported in xnnpack delegate, however it expects to fuse batch_norm with conv. I suspect your model does not have conv + batch_norm and instead batch norm has some other preceding op. In that case partitioner should take care of it.

[mcr229]

yea, right now the only batch_norm supported by the XNNPACK delegate is when it succeeds convolution. In which we perform conv + bn fusion.

@digantdesai @GregoryComer We should look into either adding support to not partition batch norms which we do not support, or supporting other bn variants through transformation to dw conv.

[cbilgin]

@mcr229 Ideal state is having these other variants as well so we can probably start by adding this particular one?

[mylifeasazucchini]

@kimishpatel thanks for your prompt inputs, based on your inputs I looked at my conv-bn layers and to the best of my knowledge bn layers follow conv layers in my own model but regardless it seems there is some place where the order seems to be conv-relu-bn (that I cannot see in my code atm!) but either way I would still think the partitioner would just handle it but seems to be not the case:

I have a minimal reproducible example code where the partitioner deviates from its intended way of working, can you test this on your end too ?

import torch
import torch.nn as nn
from executorch.backends.xnnpack.partition.xnnpack_partitioner import XnnpackPartitioner
from executorch.exir import EdgeCompileConfig
import torch._export as export
from executorch.examples.portable.utils import export_to_edge, save_pte_program

class MySmallModel(torch.nn.Module):
    def __init__(self):
        super().__init__() 
        self.convl = nn.Conv2d(1, 32, 3,stride=1, padding=1)
        self.relu = nn.ReLU()
        self.bn1 = nn.BatchNorm2d(32)
        self.t_convl = nn.ConvTranspose2d(32, 3, kernel_size=2, stride=2)

    def forward(self, x):
        x = self.convl(x) #torch.Size([1, 32, 256, 256])
        x = self.relu(x) #if the order of relu and bn1 was switched it works fine!
        x = self.bn1(x)
        x = self.t_convl(x)
        return x
    
model=MySmallModel()
model.eval()

sample_inputs = (torch.randn(1, 1, 256, 256),)

model = export.capture_pre_autograd_graph(model, sample_inputs)

edge = export_to_edge(
        model,
        sample_inputs,
        edge_compile_config=EdgeCompileConfig(
            _check_ir_validity=True,
        ),
    )

print("exporting graph!")

print(edge.exported_program().graph)

edge = edge.to_backend(XnnpackPartitioner)
exec_prog = edge.to_executorch()
save_pte_program(exec_prog.buffer, "model_xnnpack")

[kimishpatel]

@GregoryComer @mcr229 can you figure out how much work this would be and eta

[mcr229]

@mylifeasazucchini Yup this does seem like a reproducible error on ourside. Instead of changing the partitioner to not partition BN in this case, we will try to support BN after ReLU among other cases.

cc. @GregoryComer who will be working on this, please feel free to chime in when you might be able to take this on. Let's try to add a test case with @mylifeasazucchini's example above as well

[GregoryComer]

After taking a quick look, I anticipate that I can implement the standalone batchnorm with a few days of effort. I'll prioritize this and try to land by end of next week or earlier.

[GregoryComer]

I apologize for the delay. I didn't quite manage to finish this before the holidays, but I am back and have a pull request out to resolve this issue: #1501. Once merged, this should resolve the bug.

[GregoryComer]

PR was merged, so it should hopefully resolve the issue. Feel free to let me know if you still see any problems.

30732fe