alchemiops.py: Flag as unavailable on non-NVIDIA torch builds

[yantar92]

Summary

On ROCm builds, nvalchemiops correctly imports but does not work during runtime, yielding

File "<...>/lib/python3.12/site-packages/torch_sim/neighbors/init.py", line 86, in torchsim_nl
return alchemiops_nl_n2(
^^^^^^^^^^^^^^^^^
File "<...>/lib/python3.12/site-packages/torch_sim/neighbors/alchemiops.py", line 62, in alchemiops_nl_n2
res = _batch_naive_neighbor_list(
^^^^^^^^^^^^^^^^^^^^^^^^^^^
File "<...>/lib/python3.12/site-packages/nvalchemiops/torch/neighbors/batch_naive.py", line 1513, in batch_naive_neighbor_list
compute_naive_num_shifts(cell, cutoff, pbc)
File "<...>/lib/python3.12/site-packages/nvalchemiops/torch/neighbors/neighbor_utils.py", line 326, in compute_naive_num_shifts
wp_device = wp.device_from_torch(device)
^^^^^^^^^^^^^^^^^^^^^^^^^^^^
File "<...>/lib/python3.12/site-packages/warp/_src/torch.py", line 39, in device_from_torch
return warp._src.context.runtime.cuda_devices[torch_device.index]
~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~^^^^^^^^^^^^^^^^^^^^
IndexError: list index out of range

Checklist

Before a pull request can be merged, the following items must be checked:

• Doc strings have been added in the Google docstring format, but not in _import_nvalchemiops_batch_neighbors, which I did not touch.
• Run ruff on your code. (no new failures)
• Tests have been added for any new functionality or bug fixes.

[CompRhys]

Thanks for the contribution, I totally understand the issue but I am not sure this is the right fix because on CPU achemiops should be the default so the check needs to check if there is a non-cuda accelerator not just if there is not a cuda accelerator.

[yantar92]

Oops. Did not realize that alchemiops works on CPUs. I have updated the commit.

[lil-lon]

Btw, if you want to detect ROCm, you can use torch.version.hip variable.

[yantar92]

Btw, if you want to detect ROCm, you can use torch.version.hip variable.

I have considered it, but AFAIU alchemiops will fail on any GPU that is not CUDA.

[CompRhys]

I think that perhaps it's hard to do in logic for everyone's preferences so perhaps set 'TS_USE_ALCHEMIOPS=false' and skip on that is most direct approach and easy to document. If not set default to true.

[yantar92]

I am ok with that approach, but that would require more significant changes.

1. ALCHEMIOPS_AVAILABLE is currently used for load-time checks. They will likely need to be moved to runtime.
2. Failure to run alcheimops-related functions should better provide some hints about existance of this user toggle. At least, that will greatly simplify debugging the issue I encountered.

[CompRhys]

1. ALCHEMIOPS_AVAILABLE is currently used for load-time checks. They will likely need to be moved to runtime.

I don't think that adding another check for setting ALCHEMIOPS_AVAILABLE is a significant change it is less lines of code that this solution, indeed it could be a single line change without needing explicit tests imo. My opposition to this solution is that relying on 'torch.cuda.is_available()' to be true for a non-cuda device to get the logic tree to work for cpus is just very indirect and impossible to parse without knowing that amd just hides behind the cuda api and I would hope amd would avoid that obvious misdirection in future api updates.

2. Failure to run alcheimops-related functions should better provide some hints about existance of this user toggle. At least, that will greatly simplify debugging the issue I encountered.

Any brief note about the env variable would be surfaced instantly by an agent in debugging so I think that the env variable covers that documentation side. However if you're using TS on ROCm a lot it could also be cool to create a dedicated docs page called something like 'TorchSim on AMD' to capture all the learnings you have.

[yantar92]

I don't think that adding another check for setting ALCHEMIOPS_AVAILABLE is a significant change it is less lines of code that this solution, indeed it could be a single line change without needing explicit tests imo.

Do I understand correctly that you are suggesting something like

if not TS_USE_ALCHEMIOPS:
    return None

inside _import_nvalchemiops_batch_neighbors?

My concern here is that _import_nvalchemiops_batch_neighbors is called early, during torch_sim.neighbors module init. It means that the users will have to carefully arrange setting the proposed new flag to be set before loading anything that triggers this import. That's why I though that it will require something more complex that a one-liner change. (Just in case, I have no problem with the proposed approach with a flag, just discussing the implementation).

My opposition to this solution is that relying on 'torch.cuda.is_available()' to be true for a non-cuda device to get the logic tree to work for cpus is just very indirect and impossible to parse without knowing that amd just hides behind the cuda api and I would hope amd would avoid that obvious misdirection in future api updates.

Fair. What about the alternative proposed by @lil-lon with torch.version.hip? I am thinking about both disabling alchemiops for ROCm and providing the user-level flag.

Any brief note about the env variable would be surfaced instantly by an agent in debugging so I think that the env variable covers that documentation side.

I personally prefer to more direct approach, but you are the one to decide.

However if you're using TS on ROCm a lot it could also be cool to create a dedicated docs page called something like 'TorchSim on AMD' to capture all the learnings you have.

This is the first critical failure I encountered, really. Usually, on ROCm some NVIDIA-only optimizations are simply disabled (with a warning), but that does not crash the code.

[yantar92]

Do I understand correctly that you are suggesting something like

Oh, you meant environment variable. Then, my concern is indeed moot and the code change will be simple.