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!
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
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
Thank you! I will try and test it !