Skip to content

D3 implementation as Pytorch cpp extension #317

Description

@YutackPark

It solves two problems.

  1. Redundant CPU <=> GPU memory transfer.
  2. 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

Metadata

Metadata

Assignees

No one assigned

    Labels

    No labels
    No labels

    Type

    No type

    Projects

    No projects

      Milestone

      No milestone

      Relationships

      None yet

      Development

      No branches or pull requests

      Issue actions