Shamrock 2025.10.0
Astrophysical Code
Loading...
Searching...
No Matches
pyCommonUtils.cpp
Go to the documentation of this file.
1// -------------------------------------------------------//
2//
3// SHAMROCK code for hydrodynamics
4// Copyright (c) 2021-2026 Timothée David--Cléris <tim.shamrock@proton.me>
5// SPDX-License-Identifier: CeCILL Free Software License Agreement v2.1
6// Shamrock is licensed under the CeCILL 2.1 License, see LICENSE for more information
7//
8// -------------------------------------------------------//
9
15
16#include "shambase/DistributedData.hpp"
17#include "shamalgs/primitives/compute_histogram.hpp"
18#include "shamalgs/primitives/upper_bound.hpp"
19#include "shambackends/DeviceBuffer.hpp"
20#include "shambackends/MemPerfInfos.hpp"
21#include "shambackends/kernel_call.hpp"
22#include "shambindings/pybind11_stl.hpp"
23#include "shambindings/pybindaliases.hpp"
24#include "shambindings/pytypealias.hpp"
25#include "shamcomm/logs.hpp"
26#include "shammodels/common/config/enum_NeighCacheStrategy.hpp"
27#include "shamrock/solvergraph/Field.hpp"
28#include "shamsys/NodeInstance.hpp"
29#include <pybind11/cast.h>
30#include <pybind11/complex.h>
31#include <shambackends/sycl.hpp>
32#include <vector>
33
34// Define the operator += for sham::DeviceBuffer, implementation to be done later.
35namespace sham {
36 template<typename T>
37 DeviceBuffer<T> &operator+=(DeviceBuffer<T> &lhs, const DeviceBuffer<T> &rhs) {
38 sham::kernel_call(
39 rhs.get_queue(),
40 sham::MultiRef{rhs},
41 sham::MultiRef{lhs},
42 lhs.get_size(),
43 [](u32 n, const T *rhs, T *lhs) {
44 lhs[n] += rhs[n];
45 });
46 return lhs;
47 }
48
49 template<typename T>
50 DeviceBuffer<T> &operator/=(DeviceBuffer<T> &lhs, const DeviceBuffer<T> &rhs) {
51 sham::kernel_call(
52 rhs.get_queue(),
53 sham::MultiRef{rhs},
54 sham::MultiRef{lhs},
55 lhs.get_size(),
56 [](u32 n, const T *rhs, T *lhs) {
57 auto r = rhs[n];
58 if (r != 0) {
59 lhs[n] /= r;
60 } else {
61 lhs[n] = std::numeric_limits<f64>::quiet_NaN();
62 }
63 });
64 return lhs;
65 }
66
67} // namespace sham
68
69ON_PYTHON_INIT {
70 auto &m = root_module;
71
72 py::enum_<shammodels::NeighCacheStrategy>(
73 m,
74 "NeighCacheStrategy",
75 R"==(
76 Strategy used to build the neighbours cache out of the tree traversal.
77
78 Usage
79 -----
80 >>> from shamrock import NeighCacheStrategy
81 >>> cfg.set_neigh_cache_strategy(NeighCacheStrategy.SingleStage)
82)==")
83 .value(
84 "SingleStage",
85 shammodels::NeighCacheStrategy::SingleStage,
86 R"==(
87 Single tree traversal per particle.
88
89 Each particle walks the tree itself and writes its neighbours straight to the
90 cache. Prefer this one when the tree ends up with giant leaves, as on a chaotic
91 disc: there the leaf bounding boxes grow so large that the two stage search makes
92 each particle scan far more candidates than it keeps.
93)==")
94 .value(
95 "TwoStage",
96 shammodels::NeighCacheStrategy::TwoStage,
97 R"==(
98 Two stage neighbours search (see the shamrock paper). This is the default.
99
100 A first pass walks the tree once per leaf to build a leaf to leaf neighbour map,
101 then each particle only scans the particles held by its own leaf's neighbour
102 leaves. This is usually the faster of the two, since the tree traversal is paid
103 once per leaf instead of once per particle.
104)==");
105
106 m.def(
107 "compute_histogram",
108 [](std::vector<f64> bin_edges,
109 shamrock::solvergraph::Field<f64> &x_field,
110 shamrock::solvergraph::Field<f64> &y_field,
111 bool do_average) -> std::vector<f64> {
112 auto dev_sched = shamsys::instance::get_compute_scheduler_ptr();
113
114 u32 nx = bin_edges.size() - 1;
115 std::vector<f64> bin_edge_inf(nx);
116 std::vector<f64> bin_edge_sup(nx);
117
118 for (size_t i = 0; i < nx; i++) {
119 bin_edge_inf[i] = bin_edges[i];
120 bin_edge_sup[i] = bin_edges[i + 1];
121 }
122
123 sham::DeviceBuffer<f64> ret(bin_edge_inf.size(), dev_sched);
124 ret.fill(0);
125
126 sham::DeviceBuffer<f64> bin_inf(bin_edge_inf.size(), dev_sched);
127 sham::DeviceBuffer<f64> bin_sup(bin_edge_inf.size(), dev_sched);
128 bin_inf.copy_from_stdvec(bin_edge_inf);
129 bin_sup.copy_from_stdvec(bin_edge_sup);
130
131 shambase::DistributedData<u32> obj_cnts = x_field.get_obj_cnts();
132
133 obj_cnts.for_each([&](u64 id_patch, const unsigned int &obj_cnt) {
134 ret += shamalgs::primitives::compute_histogram<f64>(
135 dev_sched,
136 bin_inf,
137 bin_sup,
138 obj_cnt,
139 [](const f64 &bin_edge_inf,
140 const f64 &bin_edge_sup,
141 const f64 &x_val,
142 const f64 &y_val,
143 bool &has_value) {
144 has_value = x_val >= bin_edge_inf && x_val < bin_edge_sup;
145 return has_value ? y_val : 0;
146 },
147 x_field.get_buf(id_patch),
148 y_field.get_buf(id_patch));
149 });
150
151 shamalgs::collective::reduce_buffer_in_place_sum(ret, MPI_COMM_WORLD);
152
153 if (do_average) {
154
155 sham::DeviceBuffer<f64> norm(bin_edge_inf.size(), dev_sched);
156 norm.fill(0);
157
158 obj_cnts.for_each([&](u64 id_patch, const unsigned int &obj_cnt) {
159 sham::DeviceBuffer<f64> unit_buf(obj_cnt, dev_sched);
160 unit_buf.fill(1);
161
162 norm += shamalgs::primitives::compute_histogram<f64>(
163 dev_sched,
164 bin_inf,
165 bin_sup,
166 obj_cnt,
167 [](const f64 &bin_edge_inf,
168 const f64 &bin_edge_sup,
169 const f64 &x_val,
170 const f64 &y_val,
171 bool &has_value) {
172 has_value = x_val >= bin_edge_inf && x_val < bin_edge_sup;
173 return has_value ? y_val : 0;
174 },
175 x_field.get_buf(id_patch),
176 unit_buf);
177 });
178
179 shamalgs::collective::reduce_buffer_in_place_sum(norm, MPI_COMM_WORLD);
180
181 ret /= norm;
182 }
183
184 return ret.copy_to_stdvec();
185 },
186 py::kw_only{},
187 py::arg("bin_edges"),
188 py::arg("x_field"),
189 py::arg("y_field"),
190 py::arg("do_average") = false);
191
192 m.def(
193 "compute_histogram_convolve_x",
194 [](std::vector<f64> bin_edges,
195 shamrock::solvergraph::Field<f64> &x_field,
196 shamrock::solvergraph::Field<f64> &y_field,
197 shamrock::solvergraph::Field<f64> &size_field,
198 bool do_average) -> std::vector<f64> {
199 auto dev_sched = shamsys::instance::get_compute_scheduler_ptr();
200
201 u32 nx = bin_edges.size() - 1;
202 std::vector<f64> bin_edge_inf(nx);
203 std::vector<f64> bin_edge_sup(nx);
204
205 for (size_t i = 0; i < nx; i++) {
206 bin_edge_inf[i] = bin_edges[i];
207 bin_edge_sup[i] = bin_edges[i + 1];
208 }
209
210 sham::DeviceBuffer<f64> ret(bin_edge_inf.size(), dev_sched);
211 ret.fill(0);
212
213 sham::DeviceBuffer<f64> bin_inf(bin_edge_inf.size(), dev_sched);
214 sham::DeviceBuffer<f64> bin_sup(bin_edge_inf.size(), dev_sched);
215 bin_inf.copy_from_stdvec(bin_edge_inf);
216 bin_sup.copy_from_stdvec(bin_edge_sup);
217
218 shambase::DistributedData<u32> obj_cnts = x_field.get_obj_cnts();
219
220 obj_cnts.for_each([&](u64 id_patch, const unsigned int &obj_cnt) {
221 ret += shamalgs::primitives::compute_histogram<f64>(
222 dev_sched,
223 bin_inf,
224 bin_sup,
225 obj_cnt,
226 [](const f64 &bin_edge_inf,
227 const f64 &bin_edge_sup,
228 const f64 &x_val,
229 const f64 &y_val,
230 const f64 &size_val,
231 bool &has_value) {
232 has_value
233 = x_val >= bin_edge_inf - size_val && x_val < bin_edge_sup + size_val;
234 return has_value ? y_val : 0;
235 },
236 x_field.get_buf(id_patch),
237 y_field.get_buf(id_patch),
238 size_field.get_buf(id_patch));
239 });
240
241 shamalgs::collective::reduce_buffer_in_place_sum(ret, MPI_COMM_WORLD);
242
243 if (do_average) {
244
245 sham::DeviceBuffer<f64> norm(bin_edge_inf.size(), dev_sched);
246 norm.fill(0);
247
248 obj_cnts.for_each([&](u64 id_patch, const unsigned int &obj_cnt) {
249 sham::DeviceBuffer<f64> unit_buf(obj_cnt, dev_sched);
250 unit_buf.fill(1);
251
252 norm += shamalgs::primitives::compute_histogram<f64>(
253 dev_sched,
254 bin_inf,
255 bin_sup,
256 obj_cnt,
257 [](const f64 &bin_edge_inf,
258 const f64 &bin_edge_sup,
259 const f64 &x_val,
260 const f64 &y_val,
261 const f64 &size_val,
262 bool &has_value) {
263 has_value = x_val >= bin_edge_inf - size_val
264 && x_val < bin_edge_sup + size_val;
265 return has_value ? y_val : 0;
266 },
267 x_field.get_buf(id_patch),
268 unit_buf,
269 size_field.get_buf(id_patch));
270 });
271
272 shamalgs::collective::reduce_buffer_in_place_sum(norm, MPI_COMM_WORLD);
273
274 ret /= norm;
275 }
276
277 return ret.copy_to_stdvec();
278 },
279 py::kw_only{},
280 py::arg("bin_edges"),
281 py::arg("x_field"),
282 py::arg("y_field"),
283 py::arg("size_field"),
284 py::arg("do_average") = false);
285
286 m.def(
287 "compute_histogram_2d",
288 [](std::vector<f64> bin_edges_x,
289 std::vector<f64> bin_edges_y,
290 shamrock::solvergraph::Field<f64> &x_field,
291 shamrock::solvergraph::Field<f64> &y_field) -> std::vector<u64> {
292 auto dev_sched = shamsys::instance::get_compute_scheduler_ptr();
293
294 u32 nx = bin_edges_x.size() - 1;
295 u32 ny = bin_edges_y.size() - 1;
296
297 sham::DeviceBuffer<u64> ret(nx * ny, dev_sched);
298 ret.fill(0);
299
300 shambase::DistributedData<u32> obj_cnts = x_field.get_obj_cnts();
301
302 sham::DeviceBuffer<f64> binsx(bin_edges_x.size(), dev_sched);
303 sham::DeviceBuffer<f64> binsy(bin_edges_y.size(), dev_sched);
304 binsx.copy_from_stdvec(bin_edges_x);
305 binsy.copy_from_stdvec(bin_edges_y);
306
307 obj_cnts.for_each([&](u64 id_patch, const unsigned int &obj_cnt) {
308 sham::kernel_call(
309 dev_sched->get_queue(),
310 sham::MultiRef{
311 binsx, binsy, x_field.get_buf(id_patch), y_field.get_buf(id_patch)},
312 sham::MultiRef{ret},
313 obj_cnt,
314 [nx, ny](
315 u32 id,
316 const f64 *__restrict x_bins,
317 const f64 *__restrict y_bins,
318 const f64 *__restrict x_field,
319 const f64 *__restrict y_field,
320 u64 *__restrict pic) {
321 auto get_pic_coord = [&](u32 ix, u32 iy) {
322 return ix + iy * nx;
323 };
324
325 f64 x_val = x_field[id];
326 f64 y_val = y_field[id];
327
328 bool is_in_x_range = x_bins[0] <= x_val && x_val <= x_bins[nx];
329 bool is_in_y_range = y_bins[0] <= y_val && y_val <= y_bins[ny];
330
331 if (!(is_in_x_range && is_in_y_range)) {
332 return;
333 }
334
335 u32 ix = shamalgs::primitives::binary_search_upper_bound(
336 x_bins, 0, nx + 1, x_val);
337 u32 iy = shamalgs::primitives::binary_search_upper_bound(
338 y_bins, 0, ny + 1, y_val);
339
340 if (ix >= nx || iy >= ny) {
341 return;
342 }
343
344 using atomic_ref_T = sycl::atomic_ref<
345 u64,
346 sycl::memory_order_relaxed,
347 sycl::memory_scope_device,
348 sycl::access::address_space::global_space>;
349
350 atomic_ref_T pic_ref(pic[get_pic_coord(ix, iy)]);
351 pic_ref++;
352 });
353 });
354
355 shamalgs::collective::reduce_buffer_in_place_sum(ret, MPI_COMM_WORLD);
356
357 return ret.copy_to_stdvec();
358 },
359 py::kw_only{},
360 py::arg("bin_edges_x"),
361 py::arg("bin_edges_y"),
362 py::arg("x_field"),
363 py::arg("y_field"));
364}
NodeInstance.hpp
Header file describing a Node Instance.
f64
double f64
Alias for double.
Definition aliases_float.hpp:20
u32
std::uint32_t u32
32 bit unsigned integer
Definition aliases_int.hpp:27
u64
std::uint64_t u64
64 bit unsigned integer
Definition aliases_int.hpp:26
sham::DeviceBuffer
A buffer allocated in USM (Unified Shared Memory).
Definition DeviceBuffer.hpp:106
shambase::DistributedData
Represents a collection of objects distributed across patches identified by a u64 id.
Definition DistributedData.hpp:43
shambase::DistributedData::for_each
void for_each(std::function< void(u64, T &)> &&f)
Applies a function to each object in the collection.
Definition DistributedData.hpp:93
compute_histogram.hpp
enum_NeighCacheStrategy.hpp
Neighbour cache build strategy enum + json serialization/deserialization.
kernel_call.hpp
sham
namespace for backends this one is named only sham since shambackends is too long to write
Definition fma_chains.hpp:24
sham::kernel_call
void kernel_call(sham::DeviceQueue &q, RefIn in, RefOut in_out, u32 n, Functor &&func, SourceLocation &&callsite=SourceLocation{})
Submit a kernel to a SYCL queue.
Definition kernel_call.hpp:344
shamalgs::primitives::binary_search_upper_bound
constexpr u32 binary_search_upper_bound(const Tkey *__restrict__ key, u32 first, u32 last, const Tkey &value)
GPU compatible implementation of std::upper_bound.
Definition upper_bound.hpp:25
shammodels::SingleStage
@ SingleStage
Single tree traversal per particle.
Definition enum_NeighCacheStrategy.hpp:32
shammodels::TwoStage
@ TwoStage
Two stage neighbours search (see shamrock paper).
Definition enum_NeighCacheStrategy.hpp:33
pybind11_stl.hpp
pybindaliases.hpp
Pybind11 include and definitions.
ON_PYTHON_INIT
#define ON_PYTHON_INIT
Register a Python module init function using static initialization.
Definition pybindaliases.hpp:69
sham::MultiRef
A class that references multiple buffers or similar objects.
Definition MultiRef.hpp:33
upper_bound.hpp
GPU compatible implementation of std::upper_bound.
Generated on for Shamrock by doxygen 1.15.0