From a3f30d03c1d0bd8bac5da64eec18beaf992da951 Mon Sep 17 00:00:00 2001 From: Filip Leonarski Date: Sat, 17 Oct 2026 09:24:16 +0200 Subject: [PATCH] Rigid-body engines leave the calling thread on its own card; the MR search runs on its transform's card Every RigidBodyGPUEngine call did cudaSetDevice(engine's card) and never set the caller back. Molecular replacement's group threads refine copies on the pool's engines - on any card where every engine on this one is taken - and then run the next translation function on the same thread. MolecularReplacementGPU ran its calls on whatever card the thread was on, so that function's buffers and cuFFT plan were made on the engine's card, beside the transform and the stream on the search's card: an illegal memory access, 9s02 with four cards visible, every run (the further-copies loop: refine a trial on another card's engine, then translation_function on the group thread). - Each engine call (and the constructor) now holds a DeviceGuard: on the engine's card for the call, back on the caller's after it. - MolecularReplacementGPU remembers the card its transform was uploaded to and runs every call (and its destructor) there, whichever card the calling thread is on. - RigidBodyGPU_CallsLeaveTheCallerOnItsCard: one engine per card, all taken from one thread; the thread's card is unchanged after each call (meaningful with two or more cards). No number changes: where every call runs is the same card as before for each object. Co-Authored-By: Claude Opus 5.5 Claude-Session: https://claude.ai/code/session_01PXc1zD1yxHqaR1DGdCTHL4 --- .../MolecularReplacementGPU.cu | 14 ++++++- .../structure_refinement/RigidBodyGPU.cu | 41 +++++++++++++------ tests/ModelValidationTest.cpp | 30 ++++++++++++++ 3 files changed, 72 insertions(+), 13 deletions(-) diff --git a/image_analysis/structure_refinement/MolecularReplacementGPU.cu b/image_analysis/structure_refinement/MolecularReplacementGPU.cu index 23cf234f0..fdb397c31 100644 --- a/image_analysis/structure_refinement/MolecularReplacementGPU.cu +++ b/image_analysis/structure_refinement/MolecularReplacementGPU.cu @@ -349,7 +349,10 @@ gemmi::FPhiGrid MapToFPhiGPU(const gemmi::Grid &map, int reach) { return hkl; } +// Every call runs on the card the transform was uploaded to, whichever card the calling thread is on: its buffers, +// its plans and its kernels beside the transform and the stream. struct MolecularReplacementGPUState { + int device = 0; CudaStream stream; CudaDevicePtr grid; Transform transform{}; @@ -357,6 +360,7 @@ struct MolecularReplacementGPUState { MolecularReplacementGPU::MolecularReplacementGPU(const MolecularTransform &ft) : state_(std::make_unique()) { + cuda_err(cudaGetDevice(&state_->device)); const gemmi::FPhiGrid &g = ft.Grid(); std::vector host(g.data.size()); for (size_t i = 0; i < host.size(); i++) @@ -367,11 +371,18 @@ MolecularReplacementGPU::MolecularReplacementGPU(const MolecularTransform &ft) cuda_err(cudaStreamSynchronize(state_->stream)); } -MolecularReplacementGPU::~MolecularReplacementGPU() = default; +MolecularReplacementGPU::~MolecularReplacementGPU() { + int current = 0; + cudaGetDevice(¤t); + cudaSetDevice(state_->device); + state_.reset(); + cudaSetDevice(current); +} std::vector MolecularReplacementGPU::RotationFunction(const SearchReflections &refl, const gemmi::UnitCell &cell, const gemmi::SpaceGroup &sg, const std::vector &rotations) { + const ScopedGpuPin pin(state_->device); cudaStream_t stream = state_->stream; const int n = static_cast(refl.hkl.size()); const std::vector s = CopyVectors(refl, cell, sg); @@ -553,6 +564,7 @@ std::vector MolecularReplacementGPU::TranslationFunction( const int npairs = static_cast(setup.pair_j.size()); if (terms == 0 || rotations.empty() || (nops < 2 && known == nullptr)) return out; + const ScopedGpuPin pin(state_->device); cudaStream_t stream = state_->stream; const int nu = setup.nu, nv = setup.nv, nw = setup.nw, nwh = nw / 2 + 1; const size_t cells = static_cast(nu) * nv * nwh, points = static_cast(nu) * nv * nw; diff --git a/image_analysis/structure_refinement/RigidBodyGPU.cu b/image_analysis/structure_refinement/RigidBodyGPU.cu index 3f406a766..2b73e5b3d 100644 --- a/image_analysis/structure_refinement/RigidBodyGPU.cu +++ b/image_analysis/structure_refinement/RigidBodyGPU.cu @@ -27,6 +27,23 @@ void cuda_err(cudaError_t val) { throw JFJochException(JFJochExceptionCategory::GPUCUDAError, cudaGetErrorString(val)); } +// The calling thread on the engine's card for one call, and back on its own card after: a caller's later work +// - a function of the search, or memory it allocates - runs on the card the caller chose, not on whichever card +// the last engine it borrowed is on. +class DeviceGuard { +public: + explicit DeviceGuard(int device) { + cuda_err(cudaGetDevice(&previous_)); + cuda_err(cudaSetDevice(device)); + } + ~DeviceGuard() { cudaSetDevice(previous_); } + DeviceGuard(const DeviceGuard &) = delete; + DeviceGuard &operator=(const DeviceGuard &) = delete; + +private: + int previous_ = 0; +}; + void cufft_err(cufftResult val) { if (val != CUFFT_SUCCESS) throw JFJochException(JFJochExceptionCategory::GPUCUDAError, @@ -494,7 +511,7 @@ size_t RigidBodyGPUEngine::DeviceBytes(const RigidBodyGPUCapacity &c) { } RigidBodyGPUEngine::RigidBodyGPUEngine(const RigidBodyGPUCapacity &capacity, int device) { - cuda_err(cudaSetDevice(device)); + const DeviceGuard guard(device); impl_ = std::make_unique(capacity, device); } @@ -513,7 +530,7 @@ int RigidBodyGPUEngine::Device() const { void RigidBodyGPUEngine::SetBody(const std::vector> &relative) { RigidBodyGPUEngineImpl &e = *impl_; - cuda_err(cudaSetDevice(e.device)); + const DeviceGuard guard(e.device); if (relative.size() > e.cap.atoms) throw JFJochException(JFJochExceptionCategory::GPUCUDAError, "rigid body: more atoms than the reserve"); e.n_atoms = static_cast(relative.size()); @@ -524,7 +541,7 @@ void RigidBodyGPUEngine::SetBody(const std::vector> &relat void RigidBodyGPUEngine::SetZone(const RigidBodyGPUZone &zone) { RigidBodyGPUEngineImpl &e = *impl_; - cuda_err(cudaSetDevice(e.device)); + const DeviceGuard guard(e.device); const size_t np = static_cast(zone.nu) * zone.nv * zone.nw; const size_t nc = static_cast(zone.nu / 2 + 1) * zone.nv * zone.nw; if (static_cast(zone.atoms.size()) != e.n_atoms || np > e.cap.grid_points || nc > e.cap.complex_points || @@ -591,7 +608,7 @@ void RigidBodyGPUEngine::SetZone(const RigidBodyGPUZone &zone) { void RigidBodyGPUEngine::Fcalc(const double rotation[9], const double translation[3]) { RigidBodyGPUEngineImpl &e = *impl_; - cuda_err(cudaSetDevice(e.device)); + const DeviceGuard guard(e.device); e.Density(rotation, translation, 0); cufft_err(cufftExecR2C(e.Plan(1), e.grid, reinterpret_cast(e.spectrum.get()))); e.Compose(0, e.f, e.df_dt); @@ -600,7 +617,7 @@ void RigidBodyGPUEngine::Fcalc(const double rotation[9], const double translatio void RigidBodyGPUEngine::SetFixedCopies(const std::vector> &fixed, const std::array &moving) { RigidBodyGPUEngineImpl &e = *impl_; - cuda_err(cudaSetDevice(e.device)); + const DeviceGuard guard(e.device); const size_t copies = fixed.size() + 1; if (e.copies_mask_capacity < copies * e.n_mask_atoms) { e.copies_mask_capacity = copies * e.n_mask_atoms; @@ -641,7 +658,7 @@ void RigidBodyGPUEngine::SetFixedCopies(const std::vector void RigidBodyGPUEngine::Fmask() { RigidBodyGPUEngineImpl &e = *impl_; - cuda_err(cudaSetDevice(e.device)); + const DeviceGuard guard(e.device); float *mask_grid = e.grid.get() + e.GridPoints(); cufftComplex *mask_spectrum = reinterpret_cast(e.spectrum.get() + e.ComplexPoints()); // The mask atoms Fcalc() placed, four doubles each, as ModelMaskAtom lays them out. @@ -657,7 +674,7 @@ void RigidBodyGPUEngine::Fmask() { void RigidBodyGPUEngine::DownloadPoints(std::vector> &fcmol, std::vector> &fmask) { RigidBodyGPUEngineImpl &e = *impl_; - cuda_err(cudaSetDevice(e.device)); + const DeviceGuard guard(e.device); fcmol.resize(e.n_points); fmask.resize(e.n_points); if (e.n_points == 0) @@ -670,7 +687,7 @@ void RigidBodyGPUEngine::DownloadPoints(std::vector> &fcmol void RigidBodyGPUEngine::FitSolvent(double b_hi, double &k_sol, double &b_sol) { RigidBodyGPUEngineImpl &e = *impl_; - cuda_err(cudaSetDevice(e.device)); + const DeviceGuard guard(e.device); e.PackPoints(); const ModelSolventFit fit = e.scale->FitSolvent(e.point_fcmol, e.point_fmask, b_hi); k_sol = fit.k_sol; @@ -679,7 +696,7 @@ void RigidBodyGPUEngine::FitSolvent(double b_hi, double &k_sol, double &b_sol) { void RigidBodyGPUEngine::FitScale(double k_sol, double b_sol, double &k_overall, double b_star[6]) { RigidBodyGPUEngineImpl &e = *impl_; - cuda_err(cudaSetDevice(e.device)); + const DeviceGuard guard(e.device); e.PackPoints(); const ModelScaleParams p = e.scale->Fit(e.point_fcmol, e.point_fmask, k_sol, b_sol); k_overall = p.k_overall; @@ -689,7 +706,7 @@ void RigidBodyGPUEngine::FitScale(double k_sol, double b_sol, double &k_overall, void RigidBodyGPUEngine::Residuals(double k_overall, const double b_star[6], double k_sol, double b_sol, double *residuals) { RigidBodyGPUEngineImpl &e = *impl_; - cuda_err(cudaSetDevice(e.device)); + const DeviceGuard guard(e.device); const Scale sc = e.MakeScale(k_overall, b_star, k_sol, b_sol); residual_kernel<<>>(e.obs_row, e.obs_fobs, static_cast(e.n_obs), e.f, e.fmask, e.row_hkl, e.row_stol2, sc, e.f_mean, @@ -703,7 +720,7 @@ void RigidBodyGPUEngine::Jacobian(const double rotation[3][9], const double tran double k_overall, const double b_star[6], double k_sol, double b_sol, std::vector &jtj, std::vector &jtq) { RigidBodyGPUEngineImpl &e = *impl_; - cuda_err(cudaSetDevice(e.device)); + const DeviceGuard guard(e.device); for (int j = 0; j < 3; j++) e.Density(rotation[j], translation[j], j); cufft_err(cufftExecR2C(e.Plan(3), e.grid, reinterpret_cast(e.spectrum.get()))); @@ -730,7 +747,7 @@ void RigidBodyGPUEngine::Jacobian(const double rotation[3][9], const double tran void RigidBodyGPUEngine::ProjectJacobian(const std::vector &x, double *jacobian) { RigidBodyGPUEngineImpl &e = *impl_; - cuda_err(cudaSetDevice(e.device)); + const DeviceGuard guard(e.device); const int p = 1 + static_cast(e.n_constraints); cuda_err(cudaMemcpyAsync(e.x, x.data(), x.size() * sizeof(double), cudaMemcpyHostToDevice, e.stream)); project_kernel<<>>(e.jacobian, e.jk, p, e.x, static_cast(e.n_obs)); diff --git a/tests/ModelValidationTest.cpp b/tests/ModelValidationTest.cpp index cec04d31d..97698c0d6 100644 --- a/tests/ModelValidationTest.cpp +++ b/tests/ModelValidationTest.cpp @@ -1655,6 +1655,36 @@ TEST_CASE("RigidBodyGPU_Deterministic", "[ModelValidation][gpu]") { } } +// An engine's calls run on the engine's card and leave the calling thread on its own: every engine of a pool with +// one on each card, all taken at once from one thread, evaluated in turn. A thread left on the last engine's card +// would run what it does next - a search's function, memory it allocates - on that card, beside data on another. +TEST_CASE("RigidBodyGPU_CallsLeaveTheCallerOnItsCard", "[ModelValidation][gpu]") { + if (get_gpu_count() == 0) + SKIP("No GPU"); + Logger logger("RigidBodyGPU_CallsLeaveTheCallerOnItsCard"); + const char *cryst = kRigidBodyCrysts[4]; + const gemmi::Structure st = AnisoCluster(cryst); + const gemmi::SpaceGroup *sg = st.find_spacegroup(); + const auto fobs = OwnAmplitudes(cryst, st, 3.5, logger); + const size_t cards = get_gpu_count(); + std::vector models(cards, st.models[0]); + auto pool = TestPool(models[0], st.cell, *sg, 3.5, fobs.v.size(), cards, logger); + const int caller = RigidBodyGPUEngine::CurrentDevice(); + std::vector> targets; + for (size_t k = 0; k < cards; k++) + targets.push_back(std::make_unique(*pool, models[k], st.cell, *sg, 4)); + CHECK(RigidBodyGPUEngine::CurrentDevice() == caller); + const double q0[6] = {0.12, -0.09, 0.07, 0.15, -0.10, 0.05}; + for (auto &t : targets) { + t->SetZone(fobs, 3.5); + CHECK(RigidBodyGPUEngine::CurrentDevice() == caller); + std::vector r(t->NumObservations()), j(6 * t->NumObservations()); + REQUIRE(t->Residuals(q0, r.data())); + REQUIRE(t->Jacobian(q0, j.data())); + CHECK(RigidBodyGPUEngine::CurrentDevice() == caller); + } +} + // The model's structure factors on the GPU against the CPU path model validation takes - gemmi's density and // mask on the symmetrized grid, transformed by FFTW and read off by prepare_asu_data() - in the five groups, // at a resolution where the mask's shrink step does nothing and at one where it does. The same reflections