[BYOC] Question about registering operator

Questions
#1

Hi, I’m a beginner with TVM. I’m successful to integrate our accelerator library into TVM using JSON Runtime and run several deep learning networks. Thank you for kind guide.

For additional testing, I want some operators to need CPU fallback in certain input shapes. But I could not find a guide for this, and I only saw a guide to offload a specific operator using “@tvm.ir.register_op_attr”. Is there a way to use the CPU device, not our accelerator library, only in a specific input shape?

#2

You can use the annotation to achieve this.

@tvm.ir.register_op_attr("nn.conv2d", "target.accel")
def conv2d(expr):
    args = expr.args
    data_shape = args[0].checked_type.shape
    if not check_shape(data_shape):
        return False
    return True
1 Like
#3

Thank you very much!