[RFC] Op based annotation for external codegen

I think this method should work for what you want, although maybe we could make it a bit clearer.

This annotation method

@reg.register("nn.batch_norm", "FTVMExternalCompiler")
def batch_norm(attrs, args, compiler):
    if compiler != "dnnl":
        return False
    return # check with attrs and args

allows you to not just say True of False based on the compiler, but also based on the attrs and args. You could query these to see the tensor shapes and add in a check to say it’s not supported if the input tensors are too large (or other things like if the kernel size isn’t supported).

2 Likes

Would be great to consider the context of [RFC][Relay] Program Matching for Relay Pt 1: A pattern language

With the new pattern language, we could hugely simplify MergeComposite. There’s potentially an argument we wouldn’t need the pass at all if we considered everything as patterns in the annotation pass (with single op patterns where necessary). I don’t think we should block progress waiting for it, but we’ll definitely have to consider where it can simplify the implementation when it arrives.