[DON'T MERGE]PoC for Conversion from FP32 to mixed precision model

[anirudh2290]

Description

This is a PoC for #14584
prereq is the AMP PR: #14173
I have pulled some of the operator code for amp_cast and amp_multicast etc. into my branch for now.

Approach

• Currently, the pass uses the fp16_op_names, fp32_op_names and widest_op_names to make decisions on operators to insert in the computation graph.
• The logic for casting currently follows the PR: MXNet AMP (automatic mixed precision) #14173.
• For fp16_ops an amp_cast with target_dtype of fp16 is inserted before the op, fp32 is inserted before op for fp32 ops, for widest_type_casts, amp_multicast is used.
• The code is currently specific to fp16 and fp32 ops but i plan to change it to use target_dtype, low_precision_ops, full_precision_ops.

Testing

Test script for the PoC:

import mxnet as mx

data = mx.sym.var("data")
data2 = mx.sym.var("data2")
data3 = mx.sym.var("data3")
x = mx.sym.exp(data)
x2 = mx.sym.sin(data)
x3 = mx.sym.cos(data)
sym = x + x2 + x3
result = mx.sym.add_n(sym, data2, data3)
x = mx.viz.plot_network(result)
casted_result = mx.contrib.amp._convert_symbol(result, fp32_op_names=["elemwise_add"], fp16_op_names=["sin", "cos", "exp"], widest_type_op_names=["add_n"])
y = mx.viz.plot_network(casted_result)
x.render('test-output/round-table.gv', view=False)
y.render('test-output/round-table.gv', view=False)

Model before:

Model after:

@ptrendx @DickJC123 @Caenorst @ZhennanQin @pengzhao-intel

[anirudh2290]

Closing this Pull Request. Final PR: #15118