This pr modifies the topi implementation (which is also the legalizer's backend of Relax) of LayerNorm and GroupNorm operators to allow them to accept fp16 inputs, cast to fp32 internally, and produce fp16 outputs.
This can help eliminate unnecessary casts caused by AMP.
As more and more ML models nowadays contain the group normalization
computation, we find it beneficial to introduce this op to TOPI level.
It will enable us to optimize the group normalization operation as a
whole in a more convenient way.
This PR introduces the group normalization op to TOPI. The group norm
operation was introduced in https://arxiv.org/abs/1803.08494. The
implementation uses tuple reduction, same as the implementation of layer
norm. Implemented with tuple reduction, the corresponding generated TIR
function can be optimized by cross-thread reduction or rfactor through
MetaSchedule.
Co-authored-by: Bohan Hou <spectrometerh@gmail.com>
Tianqi Chen