File size: 3,665 Bytes
be3cca2 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 | #ifndef GAIA_SIMD_NEIGHBOR_HPP
#define GAIA_SIMD_NEIGHBOR_HPP
#include <vector>
#include <cmath>
#include <algorithm>
#include <iostream>
#ifdef __AVX2__
#include <immintrin.h>
#endif
namespace gaia {
namespace simd {
struct Vec3 {
float x, y, z;
Vec3(float x=0, float y=0, float z=0) : x(x), y(y), z(z) {}
};
class SIMDNeighborList {
public:
SIMDNeighborList(float cutoff = 10.0f) : cutoff(cutoff) {
simd_width = 8;
}
void build(const std::vector<Vec3>& positions) {
neighbors.clear();
neighbors.resize(positions.size());
#ifdef __AVX2__
for (int i = 0; i < (int)positions.size(); i++) {
float xi = positions[i].x;
float yi = positions[i].y;
float zi = positions[i].z;
__m256 xi_vec = _mm256_set1_ps(xi);
__m256 yi_vec = _mm256_set1_ps(yi);
__m256 zi_vec = _mm256_set1_ps(zi);
__m256 cutoff_sq = _mm256_set1_ps(cutoff * cutoff);
for (int j = i + 1; j < (int)positions.size(); j += simd_width) {
int remaining = std::min(simd_width, (int)positions.size() - j);
float xj_arr[8] = {0}, yj_arr[8] = {0}, zj_arr[8] = {0};
for (int k = 0; k < remaining; k++) {
xj_arr[k] = positions[j + k].x;
yj_arr[k] = positions[j + k].y;
zj_arr[k] = positions[j + k].z;
}
__m256 xj_vec = _mm256_loadu_ps(xj_arr);
__m256 yj_vec = _mm256_loadu_ps(yj_arr);
__m256 zj_vec = _mm256_loadu_ps(zj_arr);
__m256 dx = _mm256_sub_ps(xi_vec, xj_vec);
__m256 dy = _mm256_sub_ps(yi_vec, yj_vec);
__m256 dz = _mm256_sub_ps(zi_vec, zj_vec);
__m256 dx2 = _mm256_mul_ps(dx, dx);
__m256 dy2 = _mm256_mul_ps(dy, dy);
__m256 dz2 = _mm256_mul_ps(dz, dz);
__m256 r2 = _mm256_add_ps(_mm256_add_ps(dx2, dy2), dz2);
__m256 mask = _mm256_cmp_ps(r2, cutoff_sq, _CMP_LT_OQ);
int mask_int = _mm256_movemask_ps(mask);
for (int k = 0; k < remaining; k++) {
if (mask_int & (1 << k)) {
int idx = j + k;
if (idx < (int)positions.size()) {
neighbors[i].push_back(idx);
neighbors[idx].push_back(i);
}
}
}
}
}
#else
for (int i = 0; i < (int)positions.size(); i++) {
for (int j = i + 1; j < (int)positions.size(); j++) {
float dx = positions[i].x - positions[j].x;
float dy = positions[i].y - positions[j].y;
float dz = positions[i].z - positions[j].z;
float r2 = dx*dx + dy*dy + dz*dz;
if (r2 < cutoff * cutoff) {
neighbors[i].push_back(j);
neighbors[j].push_back(i);
}
}
}
#endif
}
const std::vector<int>& get_neighbors(int i) const {
return neighbors[i];
}
size_t total_pairs() const {
size_t sum = 0;
for (const auto& n : neighbors) sum += n.size();
return sum / 2;
}
private:
float cutoff;
int simd_width;
std::vector<std::vector<int>> neighbors;
};
} // namespace simd
} // namespace gaia
#endif
|