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