Hey everyone,
We have been working on a way to add a post-search exploitation phase to the MetaSchedule. We have a PR ready (#17104). I would like to discuss the idea in the mailing list, in case you have ideas to improve it. The idea is as follows:
i. Run MetaSchedule over a given end-to-end model that must be optimized;
ii. Select the best implementation of that model that MetaSchedule finds;
iii. Use Droplet Search to exploit that candidate.
By combining MetaSchedule with Droplet Search’s coordinate descent, we can reduce the number of trials that MetaSchedule explores, while still obtaining higher quality kernels. We have been able to demonstrate that we can get faster kernels with less search time in three architectures: Nvidia A100, AMD x86, and ARM A64FX. A summary of these results is available in this manuscript: report
The proposed extension does not change the way MetaSchedule is used. Its original implementation can still be invoked via the same commands, without any modification. If the user wants to apply Droplet Search on the best model found by MetaSchedule after a number of trials, all that she needs to do is to run a second command. Section 2.4 of the PDF explains how to do it. In terms of implementation, the new patch modifies 498 lines of code in the TVM code base. I’ve also prepared a docker container with the experiments here: bennu, in case someone wants to reproduce our experiments.
Regards,
Michael Canesche, Gaurav Verma and Fernando Pereira