Merge mr-multigpu-crash-fix into rc176: rigid-body engine calls restore the caller's card and the MR GPU functions run on their own card, fixing an illegal memory access in MR with engines on other cards; output unchanged on one card

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_01PXc1zD1yxHqaR1DGdCTHL4

image_analysis/structure_refinement/MolecularReplacementGPU.cu

@@ -349,7 +349,10 @@ gemmi::FPhiGrid<float> MapToFPhiGPU(const gemmi::Grid<float> &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<float2> grid;
     Transform transform{};
@@ -357,6 +360,7 @@ struct MolecularReplacementGPUState {
 
 MolecularReplacementGPU::MolecularReplacementGPU(const MolecularTransform &ft)
     : state_(std::make_unique<MolecularReplacementGPUState>()) {
+    cuda_err(cudaGetDevice(&state_->device));
     const gemmi::FPhiGrid<float> &g = ft.Grid();
     std::vector<float2> 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(&current);
+    cudaSetDevice(state_->device);
+    state_.reset();
+    cudaSetDevice(current);
+}
 
 std::vector<double> MolecularReplacementGPU::RotationFunction(const SearchReflections &refl,
                                                               const gemmi::UnitCell &cell, const gemmi::SpaceGroup &sg,
                                                               const std::vector<gemmi::Mat33> &rotations) {
+    const ScopedGpuPin pin(state_->device);
     cudaStream_t stream = state_->stream;
     const int n = static_cast<int>(refl.hkl.size());
     const std::vector<gemmi::Vec3> s = CopyVectors(refl, cell, sg);
@@ -553,6 +564,7 @@ std::vector<TranslationPeak> MolecularReplacementGPU::TranslationFunction(
     const int npairs = static_cast<int>(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<size_t>(nu) * nv * nwh, points = static_cast<size_t>(nu) * nv * nw;

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<RigidBodyGPUEngineImpl>(capacity, device);
 }
 
@@ -513,7 +530,7 @@ int RigidBodyGPUEngine::Device() const {
 
 void RigidBodyGPUEngine::SetBody(const std::vector<std::array<double, 3>> &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<int>(relative.size());
@@ -524,7 +541,7 @@ void RigidBodyGPUEngine::SetBody(const std::vector<std::array<double, 3>> &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<size_t>(zone.nu) * zone.nv * zone.nw;
     const size_t nc = static_cast<size_t>(zone.nu / 2 + 1) * zone.nv * zone.nw;
     if (static_cast<int>(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<cufftComplex *>(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<std::array<double, 12>> &fixed,
                                         const std::array<double, 12> &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<std::array<double, 12>
 
 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<cufftComplex *>(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<std::array<float, 2>> &fcmol,
                                         std::vector<std::array<float, 2>> &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<std::array<float, 2>> &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<<<blocks(e.n_obs), THREADS, 0, e.stream>>>(e.obs_row, e.obs_fobs, static_cast<int>(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<double> &jtj, std::vector<double> &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<cufftComplex *>(e.spectrum.get())));
@@ -730,7 +747,7 @@ void RigidBodyGPUEngine::Jacobian(const double rotation[3][9], const double tran
 
 void RigidBodyGPUEngine::ProjectJacobian(const std::vector<double> &x, double *jacobian) {
     RigidBodyGPUEngineImpl &e = *impl_;
-    cuda_err(cudaSetDevice(e.device));
+    const DeviceGuard guard(e.device);
     const int p = 1 + static_cast<int>(e.n_constraints);
     cuda_err(cudaMemcpyAsync(e.x, x.data(), x.size() * sizeof(double), cudaMemcpyHostToDevice, e.stream));
     project_kernel<<<blocks(e.n_obs), THREADS, 0, e.stream>>>(e.jacobian, e.jk, p, e.x, static_cast<int>(e.n_obs));

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<gemmi::Model> 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<std::unique_ptr<RigidBodyTargetGPU>> targets;
+    for (size_t k = 0; k < cards; k++)
+        targets.push_back(std::make_unique<RigidBodyTargetGPU>(*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<double> 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