UMA ConvertLayout relay pass

Still trying to integrate a custom accelerator within TVM using the UMA framework. The next problem we’re hitting is the conversion of the data layout for conv2d operations. I’m a bit puzzled where to put this in the UMA framework.

My first guess would be to register a relay pass in our accelerator backend, but I’m not able to make this work. It seems like this relay pass does not happen because of no PassContext being initialized. On top, it feels a bit strange that this relay pass is happening in the partition call.

I have created a reproducible test case, where I’m able to do the layout conversion manually, but, as said, the same conversion does not happen through the UMA ‘pass’. You can check my test case in this fork. Any guidance on how to proceed? Do these passes belong somewhere else?

Maybe some extra context. I do want this layout conversion to be associated with my accelerator, as it only needs to happen there. Unsupported operations do not need their layout changed.

That is why I chose to have this relay pass happen @ PassPhase.POST_PARTITIONING_1. Not sure whether it makes a real difference with PassPhase.POST_PARTITIONING_0.