hip_torch_nl

PyTorch C++ Extension for GPU-accelerated neighbor list computation on AMD hardware.

Requirements

• ROCm 5.0+
• PyTorch with ROCm support

Installation

cd hip_torch_nl

# Set environment for your GPU architecture (example: gfx1102)
export HSA_OVERRIDE_GFX_VERSION=11.0.0

pip install -e .

Usage

import torch
from hip_torch_nl import hip_torch_nl, HIP_TORCH_NL_AVAILABLE

if HIP_TORCH_NL_AVAILABLE:
    positions = torch.rand(1000, 3, device='cuda') * 10.0
    cell = torch.eye(3, device='cuda') * 10.0
    pbc = torch.tensor([True, True, True], device='cuda')
    cutoff = torch.tensor(3.0)

    mapping, shifts = hip_torch_nl(positions, cell, pbc, cutoff)

Architecture

hip_torch_nl/
├── __init__.py                     # Python interface
├── setup.py                        # Build configuration
└── csrc/
    ├── hip_neighborlist.cpp        # C++ bindings
    └── hip_neighborlist_kernel.cu  # CUDA/HIP kernel

The extension uses PyTorch's memory allocator and stream management for compatibility with other PyTorch GPU operations.

License

BSD-3-Clause