Create PythonTorchForce

[peastman]

This adds PythonTorchForce. It's very similar to PythonForce, except that the positions are passed as a Tensor instead of a State, and the forces are returned as a Tensor instead of a NumPy array. When using the CUDA platform, this lets us skip some data copies between host and device, reducing overhead.

In principle we could do the same with HIP, but I haven't managed to compile anything against the HIP build of PyTorch. See pytorch/pytorch#181939.

So how much does it help? I tried using AceFF 1.1 (a pretty fast potential as ML models go) to simulate a 2 nm water box with 774 atoms on a RTX 4080. With the current PythonForce implementation, 100 time steps take 1.92 seconds. Switching to PythonTorchForce reduces it to 1.88 seconds. Oh well. Looking at the profiler I can see that each invocation of the model takes about 18 ms. In between there's a gap where it does data transfers and integration. With PythonForce, that gap is about 1 ms. PythonTorchForce reduces it to 0.5 ms.

So this is only going to be signficant for really fast models. Still, it's something.

[peastman]

There, I finally got all the tests to pass! The CI workflow was very out of date.

[epretti]

Should this be CalcPythonTorchForceKernel::Name()?

[epretti]

The test failure looks strange...

Fatal Python error: Segmentation fault

Thread 0x00007f6cf76c7740 [pytest] (most recent call first):
  File "/home/runner/miniconda3/envs/build/lib/python3.14/site-packages/openmm/openmm.py", line 16692 in minimize
  File "/home/runner/work/openmm-torch/openmm-torch/python/tests/TestPythonTorchForce.py", line 190 in testMinimization
  ...

Current thread's C stack trace (most recent call first):
  ...
  Binary file "/home/runner/.cache/pocl/kcache/AH/MDDDDLEPEJELHLNMABPIBDIIJPEBOHCNCGCAN/getScale/64-1-1-goffs0-smallgrid/getScale.so", at _pocl_kernel_getScale_workgroup+0x1132 [0x7f6c884d8132]

...
/home/runner/work/_temp/e48d8ed4-04b5-4621-addf-7b9a2dc460e1.sh: line 3: 24673 Segmentation fault      (core dumped) pytest --verbose Test*
TestPythonTorchForce.py::TestPythonTorchForce::testMinimization 

It's crashing in a kernel inside the GPU minimizer. Is it running with some CPU OpenCL implementation that's getting installed in the CI environment? I can't reproduce it locally with either NVIDIA (GPU) or Intel CPU OpenCL; everything passes on my machine.

[peastman]

I've been trying to reproduce it, also without success. It's using pocl, which is a very buggy OpenCL. That might or might not be the problem.

[peastman]

After switching from pocl to Intel OpenCL (which is still buggy, but not quite as buggy), the test passes. And I can see it running OpenCL tests, so it's not just that the platform gets skipped.

TestTorchForce.py::testForce[OpenCL-True-../../tests/central.pt-False-False] PASSED [ 73%]
TestTorchForce.py::testForce[OpenCL-True-../../tests/forces.pt-True-False] PASSED [ 75%]
TestTorchForce.py::testForce[OpenCL-True-../../tests/forces.pt-True-True] PASSED [ 77%]

[epretti]

Looks good. I tried installing pocl on my machine and could reproduce the failure from before; the minimizer test and several others from the main OpenMM test suite failed with it, so I think it is just a broken implementation. In any case, it seems not to be a problem with PythonTorchForce or the minimizer.