[RFC] Calculate atom groups on GPU

[HanatoK]

Hi @giacomofiorin and @jhenin ! This WIP PR is a partial implementation of the plan discussed in #816. This PR only implements the calculation in cvm::atom_group in CUDA when the proxy sets the flag support_gpu to true. The test code is in tests/functional_gpu/run_colvars_test_cuda.cpp. Needless to say, there are many still questions and issues:

colvarproxy

• What's the best practice to allow a proxy to declare GPU support? Currently this PR uses the support_gpu flag in colvarproxy_gpu for it but sometimes people may prefer CPU even if the proxy supports GPU. Should we also check smp, or provide a way to set the support_gpu flag from the config file? (Done, smp gpu is supported)

colvarmodule and colvar

• The logic of executing GPU kernels in this PR is running all kernels in a single stream but with possibly different graphs. To improve the parallelization, this PR unpacks all atom groups from CVCs in colvar {...} and computes them in a single CUDA graph, and then copies the data to the host memory if necessary. The CVCs and colvar {...} blocks are computed after that. In other words, the order of computation looks like reverse level order traversal of a tree, which is different from the CPU code. However, the current implementation relies partially on C++ exceptions, and I think it would be better to convert colvardeps* into an AST (as disscussed in [RFC] Work towards reusable CVCs #700), and move the feature-level dependency resolution into a separate class/functor that traverses the AST.

colvaratoms

• I still think colvaratoms is too heavy and the optimal alignment should be implemented as a separate CVC, and Colvars should allow referencing it in other CVCs;
• I would like to incorporate the reading of total forces into the read_data graph, but cvm::atom_group::read_total_forces can be called anywhere from the CVC class, so it is difficult to know if the total forces are required beforehand. Is it possible to know the parents of an atom group requesting total forces in read_data of the atom group?

tests

• How can I add tests/functional_gpu/run_colvars_test_cuda.cpp to the cmake tests?
• How can I compare the results of tests/functional_gpu/run_colvars_test_cuda.cpp with run_colvars_test.cpp?
• The tests can be run by the following commands

cd tests/functional_gpu
mkdir build
cd build/
cmake ../
make -j8
make test

[HanatoK]

In my benchmarks, when using the GPU code for atom groups, the TRPV1 RMSD simulation is about 5% slower than the CPU code. There are multiple reasons:

• Although the optimal alignment is now performed on GPU, the RMSD CVC is still CPU only, which requires copying the atoms from the GPU memory to the host memory after rotation. Also, the gradients of RMSD is CPU only, and the GPU needs to wait for the CPU to complete the gradient calculation;
• The CPU code is much faster after switching to SOA, and it can be run in parallel with the NAMD NB kernel (the most computational expensive kernel). The Colvars GPU atom group kernels have to wait for the NB kernel. The situation might be changed if the CVCs in Colvars are more expensive than the NAMD NB kernel.

[jhenin]

Thanks @HanatoK ! About user control, I think a behavior similar to the current smp flag would be good, that is, enabled by default if the hardware supports it, and can be disabled explicitly. If the value is a string, it opens the possibility for finer control, as you recently did with smp to allow for selecting the parallelization scheme.

[jhenin]

I see your point about dependency resolution, and I agree that moving to an explicit AST is better.

[jhenin]

At this point, there is only a moderate amount of code dedicated to the optimal alignment in the colvaratoms class. And the orientation CVC is already mostly equivalent to what you describe. So maybe this goal could be achieved by adapting the orientation CVC?

[jhenin]

I would like to incorporate the reading of total forces into the read_data graph, but cvm::atom_group::read_total_forces can be called anywhere from the CVC class, so it is difficult to know if the total forces are required beforehand. Is it possible to know the parents of an atom group requesting total forces in read_data of the atom group?

We could certainly add constraints on when read_total_forces gets called, if that helps. For now, it is only called by calc_force_invgrads. We can make that a constraint.

[HanatoK]

Thanks @HanatoK ! About user control, I think a behavior similar to the current smp flag would be good, that is, enabled by default if the hardware supports it, and can be disabled explicitly. If the value is a string, it opens the possibility for finer control, as you recently did with smp to allow for selecting the parallelization scheme.

Done in 7f36ea2.

[HanatoK]

At this point, there is only a moderate amount of code dedicated to the optimal alignment in the colvaratoms class. And the orientation CVC is already mostly equivalent to what you describe. So maybe this goal could be achieved by adapting the orientation CVC?

I am not sure. The orientation CVC only gives a quaternion, but the cvm::atom_group calculates a rotated frame. If we want to adapt the orientation CVC, then we may have to move a lot of code from cvm::atom_group to colvar::orientation. But I think before changing code layout of the orientation CVC and the atom group, we have to allow referencing a CVC in another CVCs.

[giacomofiorin]

At this point, there is only a moderate amount of code dedicated to the optimal alignment in the colvaratoms class. And the orientation CVC is already mostly equivalent to what you describe. So maybe this goal could be achieved by adapting the orientation CVC?

I am not sure. The orientation CVC only gives a quaternion, but the cvm::atom_group calculates a rotated frame. If we want to adapt the orientation CVC, then we may have to move a lot of code from cvm::atom_group to colvar::orientation. But I think before changing code layout of the orientation CVC and the atom group, we have to allow referencing a CVC in another CVCs.

That's exactly right! Just stated looking at this new PR yesterday, but I wanted to reply on this point, which is that it would require refactoring of the CVC code that we haven't done yet. I think it's okay if this PR focuses on the basic properties like centers of mass for now.

[HanatoK]

At this point, there is only a moderate amount of code dedicated to the optimal alignment in the colvaratoms class. And the orientation CVC is already mostly equivalent to what you describe. So maybe this goal could be achieved by adapting the orientation CVC?

I am not sure. The orientation CVC only gives a quaternion, but the cvm::atom_group calculates a rotated frame. If we want to adapt the orientation CVC, then we may have to move a lot of code from cvm::atom_group to colvar::orientation. But I think before changing code layout of the orientation CVC and the atom group, we have to allow referencing a CVC in another CVCs.

That's exactly right! Just stated looking at this new PR yesterday, but I wanted to reply on this point, which is that it would require refactoring of the CVC code that we haven't done yet. I think it's okay if this PR focuses on the basic properties like centers of mass for now.

This PR implements almost all operations in the atom group class in CUDA, including the calculations of COM and COG, optimal alignment, rotation, and the rotation derivative. If you would like to restructure the CVC code and move the optimal alignment out of the atom group class, I can adapt the GPU code path accordingly.

[HanatoK]

The tests/functional_gpu/run_colvars_test_cuda.cpp can be compiled with -DCOLVARS_DEBUG and outputs the CUDA graphs. You can run

./run_colvars_test_cuda  ../../input_files/orientationangle-fitgroup_harmonic-fixed/test.in ../../input_files/trajectory.xyz xxx

with enabling fit gradients (commenting out enableFitGradients no) and get the graphs as follows:

Read data:
xxx_read_data.pdf

Fit gradients:
xxx_fit_gradients.pdf

Apply forces:
xxx_apply_forces.pdf

[HanatoK]

There is a race condition and I am still debugging, so I mark this PR as draft.

[HanatoK]

The tests/functional_gpu/run_colvars_test_cuda.cpp can be compiled with -DCOLVARS_DEBUG and outputs the CUDA graphs. You can run

./run_colvars_test_cuda  ../../input_files/orientationangle-fitgroup_harmonic-fixed/test.in ../../input_files/trajectory.xyz xxx

with enabling fit gradients (commenting out enableFitGradients no) and get the graphs as follows:

Read data: xxx_read_data.pdf

Fit gradients: xxx_fit_gradients.pdf

Apply forces: xxx_apply_forces.pdf

The code PR has been significantly revised to merge many small memset and memcpy nodes into kernels. An updated CUDA graph of reading data is attached:
xxx_read_data11.pdf

[HanatoK]

The failed tests of the GPU code are:

          6 - coordnum-aniso_harmonic-fixed_spiff (Failed)
         10 - coordnum-pairlist_harmonic-fixed_spiff (Failed)
         12 - coordnum_harmonic-fixed_spiff (Failed)
         13 - customfunction_harmonic-fixed (Subprocess aborted)
         14 - customfunction_harmonic-fixed_spiff (Failed)
        120 - distanceinv_harmonic-fixed_spiff (Failed)
        138 - distancez-axis-fitgroup_harmonic-fixed_spiff (Failed)
        146 - eigenvector-difference_harmonic-fixed_spiff (Failed)
        148 - eigenvector-normalized_harmonic-fixed_spiff (Failed)
        150 - eigenvector_harmonic-fixed_spiff (Failed)
        182 - polar_harmonic-fixed_spiff (Failed)
        186 - rmsd-fitgroup_harmonic-fixed_spiff (Failed)

I am still looking into the issues.

[HanatoK]

The failed tests of the GPU code are:

          6 - coordnum-aniso_harmonic-fixed_spiff (Failed)
         10 - coordnum-pairlist_harmonic-fixed_spiff (Failed)
         12 - coordnum_harmonic-fixed_spiff (Failed)
         13 - customfunction_harmonic-fixed (Subprocess aborted)
         14 - customfunction_harmonic-fixed_spiff (Failed)
        120 - distanceinv_harmonic-fixed_spiff (Failed)
        138 - distancez-axis-fitgroup_harmonic-fixed_spiff (Failed)
        146 - eigenvector-difference_harmonic-fixed_spiff (Failed)
        148 - eigenvector-normalized_harmonic-fixed_spiff (Failed)
        150 - eigenvector_harmonic-fixed_spiff (Failed)
        182 - polar_harmonic-fixed_spiff (Failed)
        186 - rmsd-fitgroup_harmonic-fixed_spiff (Failed)

I am still looking into the issues.

Commit 3703bb1 with also #868 should fix most of the problems except:

 13 - customfunction_harmonic-fixed (Subprocess aborted)
 14 - customfunction_harmonic-fixed_spiff (Failed)

The custom function does not work because I have not figured out how to add lepton to cmake yet.

[jhenin]

Looking good in general!

How does this interact with scripting? I expected it to be incompatible, but there are changes to colvarscript.

[HanatoK]

Looking good in general!

How does this interact with scripting? I expected it to be incompatible, but there are changes to colvarscript.

This PR should not have any conflicts with scripting, but there's no way to test scripting because the CudaGlobalMaster itself does not support the traditional way of scripting.

[jhenin]

I see!

Aside from the comment I made, I see no significant obstacles to merging.

[jhenin]

I've pushed a fix that's more satisfactory to me: make log() and error() check before they try to dereference the proxy pointer! Hopefully this avoids dangerous behaviour in other situations.

[giacomofiorin]

Looks awesome! Approved following conversation offline.