I have a small collection of relay passes that I found useful and that might be interesting for more general use:
- Merge consecutive transpose ops (useful for converted models when the other source framework had different conventions, I’m using PyTorch),
- merge equal (shape) constants to enable CSE (also for converted models, but probably also our own),
- specialize
_like
ops for static shapes to_to
(and similar) ops (broadcast, collapse_sum, zeros, ones, reshape) - useful after taking gradients because can reduce the complexity of the graph immensely. Removing broadcast/collapse_sum when not needed, - remove *1 and +0 (with tensor ones and zeros, these appear in gradients),
- spell out training LayerNorm (batch norm etc. could be done, too).
I implemented these (mostly) using the PatternMatching in Python. Moving them to TVM likely involves going to C++.
Are these useful and if so, would we want to automatically include them in optimization (and how)?
Best regards
Thomas