GPU Support for rascaline

[Luthaf]

From @clecust in #270, split into it's own issue

GPU Support for rascaline.torch.LodeSphericalExpansion

I encountered an issue when attempting to use rascaline.torch.LodeSphericalExpansion() with GPU input data. It appears that the computation primarily occurs on the CPU, and the results (x) are displayed on the CPU as well. Below is a snippet of my main code:

# lode
LR_HYPERS = {
    "cutoff": 4.0,
    "max_radial": 1,
    "max_angular": 1,
    "atomic_gaussian_width": 2.0,
    "center_atom_weight": 0.0,
    "radial_basis": {"Gto": {}},
    "potential_exponent": 1,
}
lr = rascaline.torch.LodeSphericalExpansion(**LR_HYPERS).to("cuda")
# data
sys = rascaline.torch.System(
    torch.tensor(atom0.numbers, dtype=torch.int64).to("cuda"),
    pos.to("cuda"),
    cell.to("cuda"),
)
# compute
lrv = lr.forward([sys, sys], gradients=["positions"])
# get block
lrvp = lr.keys_to_samples("species_center")
lrvp = lrvp.components_to_properties(["spherical_harmonics_m"])
lrvp = lrvp.keys_to_properties(["species_neighbor", "spherical_harmonics_l"])
x = lrvp.block().values

[Luthaf]

Yes, rascaline currently only supports calculations on CPU. All the CPU-intensif code is written in Rust, and then exported with a compatibilty shim to PyTorch, but the main calculation happens outside of PyTorch itself, and without using PyTorch primitives. Adding support for GPU will require us to write custom CUDA/HIP/... kernels for everything, which think we will want to do at some point. This will however require a significant refactoring and coding effort, unlikely to happen before ~1year.

@kvhuguenin and @PicoCentauri are working on a pure PyTorch version of LODE that should support GPU, but it's currently private.

[clecust]

Yes, rascaline currently only supports calculations on CPU. All the CPU-intensif code is written in Rust, and then exported with a compatibilty shim to PyTorch, but the main calculation happens outside of PyTorch itself, and without using PyTorch primitives. Adding support for GPU will require us to write custom CUDA/HIP/... kernels for everything, which think we will want to do at some point. This will however require a significant refactoring and coding effort, unlikely to happen before ~1year.

@kvhuguenin and @PicoCentauri are working on a pure PyTorch version of LODE that should support GPU, but it's currently private.

I would like to run a simulation of a charged system with about 10,000 particles. So using rascaline and pytroch together may not be practical at this stage.

Thank you very much for your excellent work, I'll keep an eye on it.
I look forward to that day sooner rather than later!