How to add new scheduler in auto-tvm

Hi guys: I have a question. If I have finished writing register_topi_schedule register_topi_compute for a specific op in topi/cuda repository, my goal is that I want to let auto_tvm automatically extract tasks with respect to that new registered scheduler, how I am going to do? I know that I need to go to op/strategy/ to add strategy.add_implementation something like that. What else do I need to do? Is there some discussion related to it? I mean to add new scheduler and compute for auto-tvm? Thank you guys!

1 Like is a tutorial that covers a lot of the strategy framework. If this is an existing operator, you should be fine just adding strategy.add_implementation to op/strategy/

In addition to registering the compute and schedule to Relay op strategy, you also need to register them as an AutoTVM task so that they can be extracted via extract_from_program and tuned. Specifically, you need to add the following decorators to your compute and schedule functions. Here I use conv2d_nchw.cuda as an example.

def conv2d_nchw(cfg, data, kernel, strides, padding, dilation, out_dtype="float32"):
    # Compute function.

def schedule_conv2d_nchw(cfg, outs):
    # Schedule function.

In this example, we registered an AutoTVM task conv2d_nchw.cuda. Since we also have the corresponding op strategy at, this task will be extracted by extract_from_program.

1 Like

Thank you! I will try and test it !