cudaCAC/kernels/forces.cuh
2025-09-12 06:16:37 -04:00

74 lines
2.2 KiB
Text

#ifndef FORCES_CUH
#define FORCES_CUH
#include "potentials/pair_potentials.cuh"
#include "precision.hpp"
#include <cstdio>
#include <type_traits>
#include <variant>
#include <vector>
namespace CAC {
inline void reset_forces_and_energies(int n_particles, real *forces,
real *energies) {
cudaMemset(forces, 0, n_particles * sizeof(real) * 3);
cudaMemset(energies, 0, n_particles * sizeof(real));
}
template <typename PotentialType>
__global__ void calc_forces_and_energies(real *xs, real *forces, real *energies,
int n_particles, real *box_len,
PotentialType potential) {
int i = blockIdx.x * blockDim.x + threadIdx.x;
if (i < n_particles) {
real xi = xs[3 * i];
real yi = xs[3 * i + 1];
real zi = xs[3 * i + 2];
for (int j = 0; j < n_particles; j++) {
if (i != j) {
real xj = xs[3 * j];
real yj = xs[3 * j + 1];
real zj = xs[3 * j + 2];
real dx = xi - xj;
real dy = yi - yj;
real dz = zi - zj;
// Apply periodic boundary conditions
dx -= box_len[0] * round(dx / box_len[0]);
dy -= box_len[1] * round(dy / box_len[1]);
dz -= box_len[2] * round(dz / box_len[2]);
ForceAndEnergy sol = potential.calc_force_and_energy({dx, dy, dz});
forces[3 * i] += sol.force.x;
forces[3 * i + 1] += sol.force.y;
forces[3 * i + 2] += sol.force.z;
energies[i] += sol.energy;
}
}
}
}
inline void launch_force_kernels(real *xs, real *forces, real *energies,
int n_particles, real *box_len,
std::vector<PairPotentials> potentials,
int grid_size, int block_size) {
reset_forces_and_energies(n_particles, forces, energies);
for (const auto &potential : potentials) {
std::visit(
[&](const auto &potential) {
using PotentialType = std::decay_t<decltype(potential)>;
calc_forces_and_energies<PotentialType><<<grid_size, block_size>>>(
xs, forces, energies, n_particles, box_len, potential);
},
potential);
cudaDeviceSynchronize();
}
}
} // namespace CAC
#endif