(1/2) Implementing Compiler Pass for AutoTP (#8204)
Working on #8104
- Added support for AutoTP. Added two primitives: copy_to_tp and
reduce_from_tp - the f and g nodes.
- Identifies column/row-parallel matmuls by the injected layer type in
nn_module_stack, reading back.
- Wrote a test to verify correctness of module injection and compiler
pass.
cc @tohtana
---------
Signed-off-by: Naveenraj Kamalakannan <therealnaveenkamal@gmail.com>
Signed-off-by: Masahiro Tanaka <mtanaka@anyscale.com>
Co-authored-by: Masahiro Tanaka <mtanaka@anyscale.com>
Co-authored-by: Masahiro Tanaka <81312776+tohtana@users.noreply.github.com>
