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/cuda.py 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
https://tvm.apache.org/docs/dev/relay_op_strategy.html 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/cuda.py.
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.
@autotvm.register_topi_compute("conv2d_nchw.cuda")
def conv2d_nchw(cfg, data, kernel, strides, padding, dilation, out_dtype="float32"):
# Compute function.
@autotvm.register_topi_schedule("conv2d_nchw.cuda")
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 https://github.com/apache/incubator-tvm/blob/main/python/tvm/relay/op/strategy/cuda.py#L128, this task will be extracted by extract_from_program.
1 Like
Thank you! I will try and test it !