It solves two problems.
- Redundant CPU <=> GPU memory transfer.
- Remove duplicates (LAMMPS D3, ASE D3, Batched D3).
FlashTP is a good example that achieves the same and is already integrated into 7net.
Related FlashTP module:
https://github.com/SNU-ARC/flashTP/blob/main/flashTP_e3nn/sptp_exp_opt/fused_e3nn_exp_opt.py
Related torch docs:
https://docs.pytorch.org/docs/2.12/cpp_extension.html
It solves two problems.
FlashTP is a good example that achieves the same and is already integrated into 7net.
Related FlashTP module:
https://github.com/SNU-ARC/flashTP/blob/main/flashTP_e3nn/sptp_exp_opt/fused_e3nn_exp_opt.py
Related torch docs:
https://docs.pytorch.org/docs/2.12/cpp_extension.html