refactor: use separate CUDA streams for children of colvardeps and use CUDA events to synchronize the calculations - #938
refactor: use separate CUDA streams for children of colvardeps and use CUDA events to synchronize the calculations#938HanatoK wants to merge 15 commits into
Conversation
There was a problem hiding this comment.
Pull request overview
Refactors GPU execution to use per-object streams, CUDA/HIP events, and reusable component-level graphs while sharing the standard CPU orchestration path.
Changes:
- Adds stream/event synchronization across proxies, atom groups, and CVCs.
- Moves RMSD GPU work into an optional graph implementation.
- Integrates GPU calculations, force propagation, buffer reallocation, and NAMD startup handling.
Reviewed changes
Copilot reviewed 35 out of 35 changed files in this pull request and generated 10 comments.
Show a summary per file
| File | Description |
|---|---|
tests/functional_gpu/run_colvars_test_cuda.cpp |
Updates asynchronous GPU test buffers and events. |
src/colvartypes.h |
Adds rotation completion event API. |
src/colvartypes.cpp |
Implements rotation event synchronization. |
src/colvarproxy.cpp |
Initializes proxy GPU events. |
src/colvarproxy_gpu.h |
Adds GPU metadata and proxy events. |
src/colvarproxy_gpu.cpp |
Implements GPU metadata and event lifecycle. |
src/colvarmodule.h |
Exposes buffer-reallocation notification. |
src/colvarmodule.cpp |
Integrates GPU force application and graph resets. |
src/colvardeps.h |
Adds object types and per-object streams. |
src/colvardeps.cpp |
Manages streams and recursive reallocation callbacks. |
src/colvarcomp.h |
Defines asynchronous CVC GPU interfaces. |
src/colvarcomp.cpp |
Integrates GPU data, gradients, and events. |
src/colvarcomp_gpath.cpp |
Uses the GPU-support feature flag. |
src/colvarcomp_distances.cpp |
Refactors RMSD GPU graphs and calculations. |
src/colvarcomp_combination.cpp |
Updates nested-CVC GPU detection. |
src/colvarbias.cpp |
Assigns dependency object types. |
src/colvarbias_restraint.cpp |
Initializes typed virtual dependency bases. |
src/colvarbias_opes.cpp |
Initializes the typed dependency base. |
src/colvarbias_meta.cpp |
Initializes the typed dependency base. |
src/colvarbias_histogram.cpp |
Initializes the typed dependency base. |
src/colvarbias_histogram_reweight_amd.cpp |
Initializes the typed dependency base. |
src/colvarbias_alb.cpp |
Initializes the typed dependency base. |
src/colvarbias_abmd.cpp |
Initializes the typed dependency base. |
src/colvarbias_abf.cpp |
Initializes the typed dependency base. |
src/colvaratoms.h |
Declares GPU graph reset handling. |
src/colvaratoms.cpp |
Adds typed dependencies and graph reset handling. |
src/colvaratoms_gpu.h |
Defines atom-group graphs and events. |
src/colvaratoms_gpu.cpp |
Implements asynchronous atom-group GPU workflows. |
src/colvar.cpp |
Dispatches CVC calculations across CPU and GPU. |
src/colvar_gpu_support.h |
Adds reusable CUDA graph infrastructure. |
src/colvar_gpu_support.cpp |
Implements graph lifecycle management. |
src/colvar_gpu_calc.h |
Defines component-level GPU graph orchestration. |
src/colvar_gpu_calc.cpp |
Implements CVC graph launches and synchronization. |
namd/cudaglobalmaster/colvarproxy_cudaglobalmaster.h |
Supports the newer NAMD step API. |
namd/cudaglobalmaster/colvarproxy_cudaglobalmaster.C |
Adds NAMD events, reallocation, and startup handling. |
💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.
| // Fit gradients of atom groups. | ||
| // TODO: Ideally, this loop should loop over all atom groups from colvarmodule, | ||
| // since it is possible to share the same atom group over multiple different | ||
| // CVCs and different colvar {...}, but this is how Colvars designed... |
There was a problem hiding this comment.
I agree with this in principle, and this has been discussed on and off. I think flexibility in the architecture should be a goal moving forward. The plan was to make re-used atom groups host gradients of several CVCs that depend on them. That does change the structure of the dependency tree, so that takes us back to your proposal of an explicit AST.
There was a problem hiding this comment.
That doesn't even require an explicit AST. It is easier to do that by adding a container in colvarmodule and modifying colvar::cvc::register_atom_group to add the atom group requested by CVCs to that container. Here's an example patch for the idea (based on this PR):
diff --git a/src/colvar.cpp b/src/colvar.cpp
index 8a756b78..c9468c5c 100644
--- a/src/colvar.cpp
+++ b/src/colvar.cpp
@@ -1513,11 +1513,6 @@ int colvar::calc_cvc_values(int first_cvc, size_t num_cvcs)
i++) {
if (!cvcs[i]->is_enabled()) continue;
cvc_count++;
- if (use_gpu) {
- error_code |= (cvcs[i])->read_data_gpu();
- } else {
- (cvcs[i])->read_data();
- }
if (use_gpu && cvcs[i]->is_enabled(f_cvc_support_gpu)) {
error_code |= (cvcs[i])->calc_value_gpu();
} else {
@@ -1651,26 +1646,6 @@ int colvar::calc_cvc_gradients(int first_cvc, size_t num_cvcs)
cvmodule->log("Done calculating gradients of colvar \""+this->name+"\".\n");
}
- // Fit gradients of atom groups.
- // TODO: Ideally, this loop should loop over all atom groups from colvarmodule,
- // since it is possible to share the same atom group over multiple different
- // CVCs and different colvar {...}, but this is how Colvars designed...
- for (i = first_cvc, cvc_count = 0;
- (i < cvcs.size()) && (cvc_count < cvc_max_count);
- i++) {
- if (!cvcs[i]->is_enabled()) continue;
- cvc_count++;
- if ((cvcs[i])->is_enabled(f_cvc_gradient)) {
- if (use_gpu) {
- error_code |= (cvcs[i])->calc_fit_gradients_gpu();
- } else {
- // if requested, propagate (via chain rule) the gradients above
- // to the atoms used to define the roto-translation
- (cvcs[i])->calc_fit_gradients();
- }
- }
- }
-
if (use_gpu) {
for (i = first_cvc, cvc_count = 0;
(i < cvcs.size()) && (cvc_count < cvc_max_count);
@@ -1685,6 +1660,17 @@ int colvar::calc_cvc_gradients(int first_cvc, size_t num_cvcs)
}
}
+ cvmodule->decrease_depth();
+
+ return error_code;
+}
+
+int colvar::debug_cvc_gradients(int first_cvc, size_t num_cvcs) {
+ cvmodule->increase_depth();
+ int error_code = COLVARS_OK;
+ size_t const cvc_max_count = num_cvcs ? num_cvcs : num_active_cvcs();
+ size_t i, cvc_count;
+ const bool use_gpu = cvmodule->proxy->get_smp_mode() == colvarproxy_smp::smp_mode_t::gpu;
// Debug gradients
for (i = first_cvc, cvc_count = 0;
(i < cvcs.size()) && (cvc_count < cvc_max_count);
@@ -1699,9 +1685,7 @@ int colvar::calc_cvc_gradients(int first_cvc, size_t num_cvcs)
}
}
}
-
cvmodule->decrease_depth();
-
return error_code;
}
diff --git a/src/colvar.h b/src/colvar.h
index c8498a7e..9195bac5 100644
--- a/src/colvar.h
+++ b/src/colvar.h
@@ -319,6 +319,8 @@ public:
/// \brief Same as \link colvar::calc_cvc_values \endlink but for Jacobian derivatives/forces
int calc_cvc_Jacobians(int first, size_t num_cvcs);
+ int debug_cvc_gradients(int first = 0, size_t num_cvcs = 0);
+
/// \brief Collect quantities from CVCs and update aggregated data for the colvar
int collect_cvc_data();
diff --git a/src/colvarcomp.cpp b/src/colvarcomp.cpp
index f6e88909..2bcb827a 100644
--- a/src/colvarcomp.cpp
+++ b/src/colvarcomp.cpp
@@ -393,6 +393,9 @@ colvar::cvc::~cvc()
{
free_children_deps();
remove_all_children();
+ for (size_t i = 0; i < atom_groups.size(); i++) {
+ cvmodule->unregister_atom_group_from_cvc(atom_groups[i]);
+ }
for (size_t i = 0; i < atom_groups.size(); i++) {
if (atom_groups[i] != NULL) delete atom_groups[i];
}
@@ -453,6 +456,7 @@ void colvar::cvc::register_atom_group(cvm::atom_group *ag)
atom_groups.push_back(ag);
add_child(ag);
enable(f_cvc_explicit_atom_groups);
+ cvmodule->register_atom_group_from_cvc(ag);
}
@@ -490,48 +494,6 @@ int colvar::cvc::set_param(std::string const ¶m_name,
}
-void colvar::cvc::read_data()
-{
- if (is_enabled(f_cvc_explicit_atom_groups)) {
- for (auto agi = atom_groups.begin(); agi != atom_groups.end(); agi++) {
- auto &atoms = *(*agi);
- atoms.reset_atoms_data();
- atoms.read_positions();
- atoms.calc_required_properties();
- // each atom group will take care of its own fitting_group, if defined
- }
- }
-}
-
-int colvar::cvc::read_data_gpu() {
- int error_code = COLVARS_OK;
-#if defined (COLVARS_CUDA) || defined (COLVARS_HIP)
- colvarproxy* proxy = cvmodule->proxy;
- if ((proxy->get_smp_mode() == colvarproxy_smp::smp_mode_t::gpu) && is_enabled(f_cvc_explicit_atom_groups)) {
- if (is_enabled(colvardeps::f_cvc_require_cpu_buffers)) {
- // We need to reset the atom group data if it requires CPU buffers
- for (auto agi = atom_groups.begin(); agi != atom_groups.end(); agi++) {
- (*agi)->reset_atoms_data();
- }
- }
- for (auto agi = atom_groups.begin(); agi != atom_groups.end(); agi++) {
- error_code |= (*agi)->get_gpu_atom_group()->read_data_gpu(*agi);
- }
- // NOTE: In after_read_data_sync, there are calls to synchronize the
- // CUDA events for checking if the eigendecompositions are failed, so
- // I separate read_data_gpu and after_read_data_sync into two loops for possibly
- // better parallelization.
- for (auto agi = atom_groups.begin(); agi != atom_groups.end(); agi++) {
- // Synchronize the results to CPU if necessary
- error_code |= (*agi)->get_gpu_atom_group()->after_read_data_sync(
- *agi, is_enabled(colvardeps::f_cvc_require_cpu_buffers));
- }
- }
-#endif
- return error_code;
-}
-
-
std::vector<std::vector<int>> colvar::cvc::get_atom_lists()
{
std::vector<std::vector<int>> lists;
@@ -611,16 +573,6 @@ void colvar::cvc::calc_Jacobian_derivative()
}
-void colvar::cvc::calc_fit_gradients()
-{
- if (is_enabled(f_cvc_explicit_gradient)) {
- for (size_t ig = 0; ig < atom_groups.size(); ig++) {
- atom_groups[ig]->calc_fit_gradients();
- }
- }
-}
-
-
void colvar::cvc::apply_force(colvarvalue const &cvforce)
{
if (is_enabled(f_cvc_explicit_atom_groups)) {
@@ -1240,19 +1192,6 @@ void colvar::cvc::wrap(colvarvalue &x_unwrapped) const
}
-int colvar::cvc::calc_fit_gradients_gpu() {
- int error_code = COLVARS_OK;
-#if defined (COLVARS_CUDA) || defined (COLVARS_HIP)
- if (is_enabled(f_cvc_explicit_gradient)) {
- for (auto agi = atom_groups.begin(); agi != atom_groups.end(); agi++) {
- error_code |= (*agi)->get_gpu_atom_group()->calc_fit_gradients_gpu(*agi);
- }
- }
-#endif
- return error_code;
-}
-
-
int colvar::cvc::begin_apply_force_from_cpu_to_gpu() {
int error_code = COLVARS_OK;
#if defined (COLVARS_CUDA) || defined (COLVARS_HIP)
diff --git a/src/colvarcomp.h b/src/colvarcomp.h
index abbd31ef..20c735b8 100644
--- a/src/colvarcomp.h
+++ b/src/colvarcomp.h
@@ -136,10 +136,6 @@ public:
/// \brief Get vector of vectors of atom IDs for all atom groups
virtual std::vector<std::vector<int> > get_atom_lists();
- /// \brief Obtain data needed for the calculation for the backend
- virtual void read_data();
- virtual int read_data_gpu();
-
/// \brief Calculate the variable
virtual void calc_value() = 0;
virtual int calc_value_gpu() { return COLVARS_NOT_IMPLEMENTED; }
@@ -153,10 +149,6 @@ public:
/// \brief CPU-side calculation after the graph in add_calc_gradients_node is done on GPU
virtual int calc_gradients_after_gpu() { return COLVARS_OK; }
- /// \brief Calculate the atomic fit gradients
- void calc_fit_gradients();
- int calc_fit_gradients_gpu();
-
/// \brief Calculate finite-difference gradients alongside the analytical ones, for each Cartesian component
virtual void debug_gradients();
diff --git a/src/colvarmodule.cpp b/src/colvarmodule.cpp
index c1c6140e..95419232 100644
--- a/src/colvarmodule.cpp
+++ b/src/colvarmodule.cpp
@@ -882,6 +882,24 @@ void colvarmodule::unregister_named_atom_group_soa(atom_group *ag) {
}
}
+void colvarmodule::register_atom_group_from_cvc(atom_group *ag) {
+ std::lock_guard<std::mutex> lock(registered_atom_groups_mutex);
+ registered_atom_groups[ag].fetch_add(1, std::memory_order_relaxed);
+}
+
+void colvarmodule::unregister_atom_group_from_cvc(atom_group* ag) {
+ std::lock_guard<std::mutex> lock(registered_atom_groups_mutex);
+ auto it = registered_atom_groups.find(ag);
+ if (it != registered_atom_groups.end()) {
+ const int old_count = it->second.fetch_sub(1);
+ if (old_count == 1) {
+ registered_atom_groups.erase(it);
+ }
+ } else {
+ return;
+ }
+}
+
int colvarmodule::change_configuration(std::string const &bias_name,
std::string const &conf)
{
@@ -943,7 +961,60 @@ int colvarmodule::calc()
this->to_str(this->step_absolute())+"\n");
}
+ const bool use_gpu = proxy->get_smp_mode() == colvarproxy_smp::smp_mode_t::gpu;
+ increase_depth();
+ // Read atom groups
+ if (use_gpu) {
+#if defined (COLVARS_CUDA) || defined (COLVARS_HIP)
+ for (auto it = registered_atom_groups.begin(); it != registered_atom_groups.end(); ++it) {
+ auto parent_cvcs = it->first->get_parents();
+ bool use_cpu_buffers = false;
+ for (auto c = parent_cvcs.begin(); c != parent_cvcs.end(); ++c) {
+ if ((*c)->get_object_type() == colvardeps::object_t::colvarcomp) {
+ auto cvc = dynamic_cast<colvar::cvc*>(*c);
+ use_cpu_buffers |= cvc->is_enabled(colvardeps::f_cvc_require_cpu_buffers);
+ }
+ }
+ if (use_cpu_buffers) {
+ it->first->reset_atoms_data();
+ }
+ error_code |= it->first->get_gpu_atom_group()->read_data_gpu(it->first);
+ error_code |= it->first->get_gpu_atom_group()->after_read_data_sync(it->first, use_cpu_buffers);
+ }
+#endif
+ } else {
+ for (auto it = registered_atom_groups.begin(); it != registered_atom_groups.end(); ++it) {
+ it->first->reset_atoms_data();
+ it->first->read_positions();
+ it->first->calc_required_properties();
+ }
+ }
+ decrease_depth();
error_code |= calc_colvars();
+ increase_depth();
+ // Calculate atom group fit gradients
+ if (use_gpu) {
+#if defined (COLVARS_CUDA) || defined (COLVARS_HIP)
+ for (auto it = registered_atom_groups.begin(); it != registered_atom_groups.end(); ++it) {
+ if (it->first->is_enabled(colvardeps::f_ag_fit_gradients)) {
+ error_code |= (it->first)->get_gpu_atom_group()->calc_fit_gradients_gpu(it->first);
+ }
+ }
+#endif
+ } else {
+ for (auto it = registered_atom_groups.begin(); it != registered_atom_groups.end(); ++it) {
+ if (it->first->is_enabled(colvardeps::f_ag_fit_gradients)) {
+ it->first->calc_fit_gradients();
+ }
+ }
+ }
+ // Debug gradients
+ for (auto cvi = variables_active()->begin(); cvi != variables_active()->end(); cvi++) {
+ if ((*cvi)->is_enabled(colvardeps::f_cv_active)) {
+ error_code |= (*cvi)->debug_cvc_gradients();
+ }
+ }
+ decrease_depth();
error_code |= calc_biases();
error_code |= update_colvar_forces();
@@ -1226,30 +1297,12 @@ int colvarmodule::update_colvar_forces()
// into host-pinned buffers. But before that, I still need to clear these
// buffers of all atom groups that may have forces applied on.
// TODO: How can I avoid constructing forced_atom_groups every step?
- std::vector<atom_group *> forced_atom_groups;
+ // std::vector<atom_group *> forced_atom_groups;
if (proxy->get_smp_mode() == colvarproxy_smp::smp_mode_t::gpu) {
#if defined (COLVARS_CUDA) || defined (COLVARS_HIP)
- for (cvi = variables_active()->begin(); cvi != variables_active()->end(); cvi++) {
- if ((*cvi)->is_enabled(colvardeps::f_cv_apply_force)) {
- auto& cvcs = (*cvi)->get_cvcs();
- for (size_t i = 0; i < cvcs.size(); i++) {
- if (!cvcs[i]->is_enabled()) continue;
- std::vector<atom_group *>& ags = (cvcs[i])->atom_groups;
- for (auto ag = ags.begin(); ag != ags.end(); ++ag) {
- if ((*ag)->size() > 0) {
- forced_atom_groups.push_back(*ag);
- }
- }
- }
- }
- }
- // Since atom groups can be shared, I have to sort and deduplicate the array.
- std::sort(forced_atom_groups.begin(), forced_atom_groups.end());
- auto last = std::unique(forced_atom_groups.begin(), forced_atom_groups.end());
- forced_atom_groups.erase(last, forced_atom_groups.end());
// Clear the host-pinned buffers
- for (auto ag = forced_atom_groups.begin(); ag != forced_atom_groups.end(); ++ag) {
- error_code |= (*ag)->get_gpu_atom_group()->begin_apply_force_gpu();
+ for (auto ag = registered_atom_groups.begin(); ag != registered_atom_groups.end(); ++ag) {
+ error_code |= ag->first->get_gpu_atom_group()->begin_apply_force_gpu();
}
#endif
}
@@ -1269,8 +1322,8 @@ int colvarmodule::update_colvar_forces()
// own stream, the CUDA kernel/Graph for adding forces are always pushed after the
// CUDA graph for fit gradients. As a result, we don't need to synchronize the
// fit gradients event.
- for (auto ag = forced_atom_groups.begin(); ag != forced_atom_groups.end(); ++ag) {
- error_code |= (*ag)->get_gpu_atom_group()->add_force_to_proxy_gpu((*ag));
+ for (auto ag = registered_atom_groups.begin(); ag != registered_atom_groups.end(); ++ag) {
+ error_code |= ag->first->get_gpu_atom_group()->add_force_to_proxy_gpu(ag->first);
}
#endif
}
diff --git a/src/colvarmodule.h b/src/colvarmodule.h
index 3d2b0b0d..30d4c9ae 100644
--- a/src/colvarmodule.h
+++ b/src/colvarmodule.h
@@ -14,6 +14,8 @@
#include <unordered_map>
#include <memory>
#include <cstdio>
+#include <atomic>
+#include <mutex>
#include "colvars_version.h"
@@ -322,11 +324,18 @@ private:
/// Array of named atom groups
std::vector<atom_group *> named_atom_groups_soa;
+ /// Array of all registered atom groups from CVCs
+ std::unordered_map<atom_group *, std::atomic<int>> registered_atom_groups;
+ std::mutex registered_atom_groups_mutex;
+
public:
void register_named_atom_group_soa(atom_group *ag);
void unregister_named_atom_group_soa(atom_group *ag);
+ void register_atom_group_from_cvc(atom_group *ag);
+ void unregister_atom_group_from_cvc(atom_group *ag);
+
/// Array of collective variables
std::vector<colvar *> *variables();
There was a problem hiding this comment.
Looks good! That looks like the beginning of an answer to your first question in the PR.
There was a problem hiding this comment.
Although it works, I suspect that this patch may not be safe enough to include in this PR. The reason is that it doesn't track correctly which CVCs are enabled. In colvar::calc_cvc_gradients, the loop is
for (i = first_cvc, cvc_count = 0;
(i < cvcs.size()) && (cvc_count < cvc_max_count);
i++) {
if (!cvcs[i]->is_enabled()) continue;
cvc_count++;
if ((cvcs[i])->is_enabled(f_cvc_gradient)) {
if (use_gpu) {
error_code |= (cvcs[i])->calc_fit_gradients_gpu();
} else {
// if requested, propagate (via chain rule) the gradients above
// to the atoms used to define the roto-translation
(cvcs[i])->calc_fit_gradients();
}
}
}which skips the disabled CVCs. The patch above, however, assumes all atom groups called in colvar::cvc::register_atom_group are used. I don't know whether there are cases that (i) a CVC is not enabled, and (ii) the children atom groups of the CVC are registered, and if so, how to handle these cases correctly.
There was a problem hiding this comment.
Iterating over all groups and having update functions exit immediately when unnecessary would give the clearest code. This can be deferred to another PR though.
|
Point of detail: the CUDAGlobalMaster OPES example fails now |
Because that register is meant for resolving references to named groups by By the way, this should be renamed to just |
Good design choice! |
Thanks! I see that you helped to fix it by adding |
Should I remove the |
According to the CUDA API documentation, this could improve the performance.
Keep in mind that not every CVC use CUDA graphs for the GPU implementation. Consequently, whether to use CUDA graphs or not is considered to be an implementation detail. This commit strips the CUDA graphs in colvar::cvc to a separate class to colvar_gpu_calc.*. The RMSD GPU code is modified to use it.
Co-authored-by: Copilot Autofix powered by AI <175728472+Copilot@users.noreply.github.com>
Co-authored-by: Copilot Autofix powered by AI <175728472+Copilot@users.noreply.github.com>
Co-authored-by: Copilot Autofix powered by AI <175728472+Copilot@users.noreply.github.com>
Co-authored-by: Copilot Autofix powered by AI <175728472+Copilot@users.noreply.github.com>
Change inconsistent outputFreq parameter
…ault_block_size "constexpr static" is used for muting the compiler warnings.
This PR implements the idea in #816 (comment). Overall, this PR should make the
smp gpucode path simpler and share the code base with the CPU code.Here's an overview of the design of this PR:
colvardepshas its own CUDA/HIP stream (or SYCL queue in the future);colvardeps, it is expected to emit different CUDA/HIP events after the calculations. For example, after reading data and calculating the required properties, an atom group should callcudaEventRecordwithcolvaratoms_gpu::event_type::read_and_calculateand its own stream;has_gpu_implementation() == trueshould implementscalc_value_gpu(). Incalc_value_gpu(), if it requires any data from its children atom groups, then it should iterate over the children, check if a child is ofcvm::atom_grouptype, and if so, callcudaStreamWaitEventwith theread_and_calculateevent from the child and its own stream. In other words, if the parent requires any data to be ready from the children, it should make its own stream wait for the corresponding events from the children;calc_gradients_gpu()(if necessary)calc_Jacobian_derivative_gpu()(if necessary)calc_force_invgrads_gpu()(if necessary)This PR should also supersede #899 once the NAMD MR https://gitlab.com/tcbgUIUC/namd/-/merge_requests/492 is merged.
I still have a few design questions so I add the request for comments tag:
colvarmodule?register_named_atom_group_soaand the C++ vectornamed_atom_groups_soa, and it looks like it could be used for solving the issue above, but I don't know why only named atom groups are registered.I have moved the CUDA graphs code there, so whether a CVC GPU implementation uses CUDA graphs or not keeps as an implementation details. A CVC can now fully optionally use CUDA graphs.colvar_gpu_calc.*are empty now and could be removed. What's the best way to remove the files without breaking the compilation and CI tests?