diff --git a/demo/binary.c b/demo/binary.c index 9fca83a..70202bb 100644 --- a/demo/binary.c +++ b/demo/binary.c @@ -7,11 +7,11 @@ #define POLESITTER_IMPLEMENTATION #include "../src/polesitter.h" -#define PARTICLE_COUNT 5000 +#define PARTICLE_COUNT 10000 #define BLACKHOLE_MASS 1000.0F #define ARENA_SIZE (1024ULL * 1024ULL * 256ULL) #define MAX_EXPECTED_SPEED 20 -#define THETA 3.4641F // sqrt of 12 +#define THETA 2.0F // sqrt of 12 float vx[PARTICLE_COUNT] = {0}; float vy[PARTICLE_COUNT] = {0}; @@ -51,7 +51,7 @@ void* init_polesitter(void) { return NULL; } - ps_config_t cfg = {buffer, ARENA_SIZE, THETA}; + ps_config_t cfg = {buffer, ARENA_SIZE, PARTICLE_COUNT, THETA, 1}; ps_init(&ctx, &cfg); arrs.x = px; @@ -155,7 +155,7 @@ void update(void) { UpdateCamera(&camera, CAMERA_ORBITAL); - ps_arena_clear(&ctx->arena); + ps_impl_arena_clear(&ctx->arena); for (int i = 0; i < PARTICLE_COUNT; ++i) { id[i] = i; @@ -253,7 +253,7 @@ void draw(void) { (unsigned char)(180 - (f * 100)), (unsigned char)(255 - (f * 255)), 255}; } - DrawPoint3D((Vector3){px[i], py[i], pz[i]}, p_color); + DrawSphereEx((Vector3){px[i], py[i], pz[i]}, 0.1F, 3, 3, p_color); } EndBlendMode(); diff --git a/demo/unary.c b/demo/unary.c index cd06d52..eaf4bda 100644 --- a/demo/unary.c +++ b/demo/unary.c @@ -4,14 +4,15 @@ #include #include +#define PS_MULTITHREADING #define POLESITTER_IMPLEMENTATION #include "../src/polesitter.h" -#define PARTICLE_COUNT 5000 +#define PARTICLE_COUNT 40000 #define BLACKHOLE_MASS 1000.0F #define ARENA_SIZE (1024ULL * 1024ULL * 256ULL) #define MAX_EXPECTED_SPEED 20 -#define THETA 3.4641F // sqrt of 12 +#define THETA 2.00F // sqrt of 12 float vx[PARTICLE_COUNT] = {0}; float vy[PARTICLE_COUNT] = {0}; @@ -51,7 +52,7 @@ void* init_polesitter(void) { return NULL; } - ps_config_t cfg = {buffer, ARENA_SIZE, THETA}; + ps_config_t cfg = {buffer, ARENA_SIZE, PARTICLE_COUNT, THETA, 4}; ps_init(&ctx, &cfg); arrs.x = px; @@ -116,7 +117,7 @@ int main(void) { void update(void) { UpdateCamera(&camera, CAMERA_ORBITAL); - ps_arena_clear(&ctx->arena); + ps_impl_arena_clear(&ctx->arena); for (int i = 0; i < PARTICLE_COUNT; ++i) { id[i] = i; diff --git a/docs/performance_graph.svg b/docs/performance_graph.svg index 426e5e3..3a035fb 100644 --- a/docs/performance_graph.svg +++ b/docs/performance_graph.svg @@ -6,7 +6,7 @@ - 2026-08-02T23:08:26.949210 + 2026-08-06T00:13:40.672698 image/svg+xml @@ -44,16 +44,16 @@ z +" clip-path="url(#p93e5e5b620)" style="fill: none; stroke-dasharray: 2.96,1.28; stroke-dashoffset: 0; stroke: #8b949e; stroke-opacity: 0.3; stroke-width: 0.8"/> - - + @@ -112,11 +112,11 @@ z +" clip-path="url(#p93e5e5b620)" style="fill: none; stroke-dasharray: 2.96,1.28; stroke-dashoffset: 0; stroke: #8b949e; stroke-opacity: 0.3; stroke-width: 0.8"/> - + @@ -198,11 +198,11 @@ z +" clip-path="url(#p93e5e5b620)" style="fill: none; stroke-dasharray: 2.96,1.28; stroke-dashoffset: 0; stroke: #8b949e; stroke-opacity: 0.3; stroke-width: 0.8"/> - + @@ -220,11 +220,11 @@ L 170.227492 37.466667 +" clip-path="url(#p93e5e5b620)" style="fill: none; stroke-dasharray: 2.96,1.28; stroke-dashoffset: 0; stroke: #8b949e; stroke-opacity: 0.3; stroke-width: 0.8"/> - + @@ -267,11 +267,11 @@ z +" clip-path="url(#p93e5e5b620)" style="fill: none; stroke-dasharray: 2.96,1.28; stroke-dashoffset: 0; stroke: #8b949e; stroke-opacity: 0.3; stroke-width: 0.8"/> - + @@ -309,11 +309,11 @@ z +" clip-path="url(#p93e5e5b620)" style="fill: none; stroke-dasharray: 2.96,1.28; stroke-dashoffset: 0; stroke: #8b949e; stroke-opacity: 0.3; stroke-width: 0.8"/> - + @@ -332,11 +332,11 @@ L 355.428731 37.466667 +" clip-path="url(#p93e5e5b620)" style="fill: none; stroke-dasharray: 2.96,1.28; stroke-dashoffset: 0; stroke: #8b949e; stroke-opacity: 0.3; stroke-width: 0.8"/> - + @@ -355,11 +355,11 @@ L 417.162477 37.466667 +" clip-path="url(#p93e5e5b620)" style="fill: none; stroke-dasharray: 2.96,1.28; stroke-dashoffset: 0; stroke: #8b949e; stroke-opacity: 0.3; stroke-width: 0.8"/> - + @@ -378,11 +378,11 @@ L 478.896223 37.466667 +" clip-path="url(#p93e5e5b620)" style="fill: none; stroke-dasharray: 2.96,1.28; stroke-dashoffset: 0; stroke: #8b949e; stroke-opacity: 0.3; stroke-width: 0.8"/> - + @@ -806,16 +806,16 @@ z +" clip-path="url(#p93e5e5b620)" style="fill: none; stroke-dasharray: 2.96,1.28; stroke-dashoffset: 0; stroke: #8b949e; stroke-opacity: 0.3; stroke-width: 0.8"/> - - + @@ -827,36 +827,36 @@ L -3.5 0 - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + @@ -1021,43 +1021,24 @@ L 565.2 137.954479 - + - + - + - - - - - - - - - - - - - - - - - - + @@ -1141,22 +1122,22 @@ z - - + + - - - - - - - - - - - - - - + + + + + + + + + + + + + - - + + - - - - - - - - - - - - - - + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + @@ -1245,7 +1268,7 @@ L 565.2 316.04 L 565.2 37.466667 " style="fill: none; stroke: #8b949e; stroke-width: 0.8; stroke-linejoin: miter; stroke-linecap: square"/> - + @@ -1476,28 +1499,28 @@ z - - + - + - + @@ -1614,18 +1637,56 @@ z - + - + - - + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + @@ -1634,13 +1695,20 @@ L 78.66 70.356667 + + + + + + + - + diff --git a/docs/performance_graph_log.svg b/docs/performance_graph_log.svg index 881c2a6..9ca98e5 100644 --- a/docs/performance_graph_log.svg +++ b/docs/performance_graph_log.svg @@ -6,7 +6,7 @@ - 2026-08-02T23:08:27.151108 + 2026-08-06T00:13:40.888825 image/svg+xml @@ -44,16 +44,16 @@ z +" clip-path="url(#p93610dacca)" style="fill: none; stroke-dasharray: 2.96,1.28; stroke-dashoffset: 0; stroke: #8b949e; stroke-opacity: 0.3; stroke-width: 0.8"/> - - + @@ -162,11 +162,11 @@ z +" clip-path="url(#p93610dacca)" style="fill: none; stroke-dasharray: 2.96,1.28; stroke-dashoffset: 0; stroke: #8b949e; stroke-opacity: 0.3; stroke-width: 0.8"/> - + @@ -217,11 +217,11 @@ z +" clip-path="url(#p93610dacca)" style="fill: none; stroke-dasharray: 2.96,1.28; stroke-dashoffset: 0; stroke: #8b949e; stroke-opacity: 0.3; stroke-width: 0.8"/> - + @@ -262,16 +262,16 @@ z +" clip-path="url(#p93610dacca)" style="fill: none; stroke-dasharray: 2.96,1.28; stroke-dashoffset: 0; stroke: #8b949e; stroke-opacity: 0.3; stroke-width: 0.8"/> - - + @@ -279,11 +279,11 @@ L 0 2 +" clip-path="url(#p93610dacca)" style="fill: none; stroke-dasharray: 2.96,1.28; stroke-dashoffset: 0; stroke: #8b949e; stroke-opacity: 0.3; stroke-width: 0.8"/> - + @@ -291,11 +291,11 @@ L 70.056918 37.466667 +" clip-path="url(#p93610dacca)" style="fill: none; stroke-dasharray: 2.96,1.28; stroke-dashoffset: 0; stroke: #8b949e; stroke-opacity: 0.3; stroke-width: 0.8"/> - + @@ -303,11 +303,11 @@ L 139.805357 37.466667 +" clip-path="url(#p93610dacca)" style="fill: none; stroke-dasharray: 2.96,1.28; stroke-dashoffset: 0; stroke: #8b949e; stroke-opacity: 0.3; stroke-width: 0.8"/> - + @@ -315,11 +315,11 @@ L 175.22212 37.466667 +" clip-path="url(#p93610dacca)" style="fill: none; stroke-dasharray: 2.96,1.28; stroke-dashoffset: 0; stroke: #8b949e; stroke-opacity: 0.3; stroke-width: 0.8"/> - + @@ -327,11 +327,11 @@ L 200.350714 37.466667 +" clip-path="url(#p93610dacca)" style="fill: none; stroke-dasharray: 2.96,1.28; stroke-dashoffset: 0; stroke: #8b949e; stroke-opacity: 0.3; stroke-width: 0.8"/> - + @@ -339,11 +339,11 @@ L 219.841965 37.466667 +" clip-path="url(#p93610dacca)" style="fill: none; stroke-dasharray: 2.96,1.28; stroke-dashoffset: 0; stroke: #8b949e; stroke-opacity: 0.3; stroke-width: 0.8"/> - + @@ -351,11 +351,11 @@ L 235.767477 37.466667 +" clip-path="url(#p93610dacca)" style="fill: none; stroke-dasharray: 2.96,1.28; stroke-dashoffset: 0; stroke: #8b949e; stroke-opacity: 0.3; stroke-width: 0.8"/> - + @@ -363,11 +363,11 @@ L 249.232305 37.466667 +" clip-path="url(#p93610dacca)" style="fill: none; stroke-dasharray: 2.96,1.28; stroke-dashoffset: 0; stroke: #8b949e; stroke-opacity: 0.3; stroke-width: 0.8"/> - + @@ -375,11 +375,11 @@ L 260.89607 37.466667 +" clip-path="url(#p93610dacca)" style="fill: none; stroke-dasharray: 2.96,1.28; stroke-dashoffset: 0; stroke: #8b949e; stroke-opacity: 0.3; stroke-width: 0.8"/> - + @@ -387,11 +387,11 @@ L 271.18424 37.466667 +" clip-path="url(#p93610dacca)" style="fill: none; stroke-dasharray: 2.96,1.28; stroke-dashoffset: 0; stroke: #8b949e; stroke-opacity: 0.3; stroke-width: 0.8"/> - + @@ -399,11 +399,11 @@ L 340.932678 37.466667 +" clip-path="url(#p93610dacca)" style="fill: none; stroke-dasharray: 2.96,1.28; stroke-dashoffset: 0; stroke: #8b949e; stroke-opacity: 0.3; stroke-width: 0.8"/> - + @@ -411,11 +411,11 @@ L 376.349442 37.466667 +" clip-path="url(#p93610dacca)" style="fill: none; stroke-dasharray: 2.96,1.28; stroke-dashoffset: 0; stroke: #8b949e; stroke-opacity: 0.3; stroke-width: 0.8"/> - + @@ -423,11 +423,11 @@ L 401.478035 37.466667 +" clip-path="url(#p93610dacca)" style="fill: none; stroke-dasharray: 2.96,1.28; stroke-dashoffset: 0; stroke: #8b949e; stroke-opacity: 0.3; stroke-width: 0.8"/> - + @@ -435,11 +435,11 @@ L 420.969286 37.466667 +" clip-path="url(#p93610dacca)" style="fill: none; stroke-dasharray: 2.96,1.28; stroke-dashoffset: 0; stroke: #8b949e; stroke-opacity: 0.3; stroke-width: 0.8"/> - + @@ -447,11 +447,11 @@ L 436.894798 37.466667 +" clip-path="url(#p93610dacca)" style="fill: none; stroke-dasharray: 2.96,1.28; stroke-dashoffset: 0; stroke: #8b949e; stroke-opacity: 0.3; stroke-width: 0.8"/> - + @@ -459,11 +459,11 @@ L 450.359627 37.466667 +" clip-path="url(#p93610dacca)" style="fill: none; stroke-dasharray: 2.96,1.28; stroke-dashoffset: 0; stroke: #8b949e; stroke-opacity: 0.3; stroke-width: 0.8"/> - + @@ -471,11 +471,11 @@ L 462.023392 37.466667 +" clip-path="url(#p93610dacca)" style="fill: none; stroke-dasharray: 2.96,1.28; stroke-dashoffset: 0; stroke: #8b949e; stroke-opacity: 0.3; stroke-width: 0.8"/> - + @@ -483,11 +483,11 @@ L 472.311562 37.466667 +" clip-path="url(#p93610dacca)" style="fill: none; stroke-dasharray: 2.96,1.28; stroke-dashoffset: 0; stroke: #8b949e; stroke-opacity: 0.3; stroke-width: 0.8"/> - + @@ -898,23 +898,23 @@ z - + - - + - + - + - + - + - + - + - + @@ -1042,18 +1042,18 @@ L 565.2 183.962425 - + - + - + @@ -1062,18 +1062,18 @@ L 565.2 122.340185 - + - + - + @@ -1082,426 +1082,498 @@ L 565.2 60.717945 - + - - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + @@ -1589,22 +1661,22 @@ z - + +" clip-path="url(#p93610dacca)" style="fill: none; stroke: #d29922; stroke-width: 2.5; stroke-linecap: square"/> - - - - - - - - - - - - - - + + + + + + + + + + + + + - - + + - - - - - - - - - - - - - - + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + @@ -1815,25 +1929,25 @@ z - - + - + @@ -1972,17 +2086,17 @@ z - + - + - + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + - + diff --git a/docs/plot.py b/docs/plot.py index 1f5f7e5..582fd6c 100644 --- a/docs/plot.py +++ b/docs/plot.py @@ -4,7 +4,8 @@ COL_TEXT = '#C9D1D9' COL_BORDER = '#8B949E' COL_NAIVE = '#D29922' -COL_FMM = '#1F6FEB' +COL_FMM_ST = '#1F6FEB' +COL_FMM_MT = '#8957E5' COL_BG = '#21262D' plt.rcParams.update({ @@ -26,12 +27,14 @@ data = pd.read_csv('results.csv') N = data['N'] t_direct = data['Direct_Time'] -t_fmm = data['FMM_Time'] +t_fmm_st = data['FMM_Time_ST'] +t_fmm_mt = data['FMM_Time_MT'] fig, ax = plt.subplots(figsize=(8, 5)) ax.plot(N, t_direct, marker='o', linewidth=2.5, color=COL_NAIVE, label='naive $\mathcal{O}(N^2)$') -ax.plot(N, t_fmm, marker='o', linewidth=2.5, color=COL_FMM, label='FMM $\mathcal{O}(N)$') +ax.plot(N, t_fmm_st, marker='o', linewidth=2.5, color=COL_FMM_ST, label='FMM $\mathcal{O}(N)$ 1-core') +ax.plot(N, t_fmm_mt, marker='o', linewidth=2.5, color=COL_FMM_MT, label='FMM $\mathcal{O}(N)$ 8-core') ax.set_title('perf scaling: direct evaluation vs FMM', fontsize=14, pad=15, fontweight='bold') ax.set_xlabel('# of particles (N)', fontsize=12) @@ -48,16 +51,14 @@ fig_log, ax_log = plt.subplots(figsize=(8, 5)) -ax_log.plot(N, t_direct, marker='o', linewidth=2.5, color=COL_NAIVE, -label='naive $\mathcal{O}(N^2)$') -ax_log.plot(N, t_fmm, marker='o', linewidth=2.5, color=COL_FMM, -label='FMM $\mathcal{O}(N)$') +ax_log.plot(N, t_direct, marker='o', linewidth=2.5, color=COL_NAIVE, label='naive $\mathcal{O}(N^2)$') +ax_log.plot(N, t_fmm_st, marker='o', linewidth=2.5, color=COL_FMM_ST, label='FMM $\mathcal{O}(N)$ 1-core') +ax_log.plot(N, t_fmm_mt, marker='o', linewidth=2.5, color=COL_FMM_MT, label='FMM $\mathcal{O}(N)$ 8-core') ax_log.set_xscale('log') ax_log.set_yscale('log') -ax_log.set_title('algorithmic scaling (log-log)', fontsize=14, pad=15, -fontweight='bold') +ax_log.set_title('algorithmic scaling (log-log)', fontsize=14, pad=15, fontweight='bold') ax_log.set_xlabel('# of particles (N)', fontsize=12) ax_log.set_ylabel('exec time (secs)', fontsize=12) diff --git a/src/polesitter.h b/src/polesitter.h index 25bfe12..d9ade04 100644 --- a/src/polesitter.h +++ b/src/polesitter.h @@ -1,16 +1,19 @@ // clang-format off /*- - polesitter - v1.0 - Fast Multipole Method (FFM) N-body solver written in C99 + polesitter - v2.0 - Fast Multipole Method (FFM) N-body solver written in C99 Do this: #define POLESITTER_IMPLEMENTATION before you include this file in *one* C or C++ file to create the implementation. + Use #define PS_MULTITHREADING to enable multithreading capabilities. + // i.e. it should look like this: #include ... #include ... #include ... + #define PS_MULTITHREADING #define POLESITTER_IMPLEMENTATION #include "polesitter.h" @@ -32,6 +35,7 @@ QUICKSTART ``` + #define PS_MULTITHREADING #define POLESITTER_IMPLEMENTATION #include "polesitter.h" #include @@ -47,8 +51,13 @@ void* memory_block = malloc(MEMORY_SIZE); // initialize the context - - ps_config_t cfg = { memory_block, MEMORY_SIZE }; + ps_config_t cfg = { + .buff = memory_block, + .buff_size = MEMORY_SIZE, + .max_particles = PARTICLE_CNT, + .theta = 2.0F, + .thrd_cnt = 4 + }; ps_context_t* ctx = NULL; ps_init(&ctx, &cfg); @@ -153,6 +162,7 @@ } // cleanup + ps_destroy(ctx); free(memory_block); return 0; } @@ -164,8 +174,10 @@ #define POLESITTER_H #include +#include #include #include +#include #if defined(__AVX2__) @@ -192,6 +204,85 @@ #endif // PS_RESTRICT +#ifdef PS_MULTITHREADING + +#ifdef _WIN32 + +#define WIN32_LEAN_AND_MEAN +#include + +typedef HANDLE ps_thrd_t; +typedef CRITICAL_SECTION ps_mtx_t; +typedef CONDITION_VARIABLE ps_cond_t; +typedef DWORD(WINAPI* ps_thrd_func_t)(void*); + +#define PS_THRD_RET_TYPE DWORD WINAPI +#define PS_THRD_RET_VAL 0 + +#if defined(_WIN64) + +// clang-format off + +#define ps_atomic_fetch_add_size_t(ptr, val) \ + (size_t)InterlockedExchangeAdd64((volatile LONG64*)(ptr), (LONG64)(val)) +#else +#define ps_atomic_fetch_add_size_t(ptr, val) \ + (size_t)InterlockedExchangeAdd((volatile LONG*)(ptr), (LONG)(val)) +#endif + +// clang-format on + +#else // POSIX + +#include + +typedef pthread_t ps_thrd_t; +typedef pthread_mutex_t ps_mtx_t; +typedef pthread_cond_t ps_cond_t; +typedef void* (*ps_thrd_func_t)(void*); + +#define PS_THRD_RET_TYPE void* +#define PS_THRD_RET_VAL NULL + +#define ps_atomic_fetch_add_size_t(ptr, val) __sync_fetch_and_add((ptr), (val)) + +#endif // POSIX + +#if defined(__x86_64__) || defined(__i386__) || defined(_M_X64) || \ + defined(_M_IX86) + +#define PS_YIELD() _mm_pause() +#elif defined(__aarch64__) || defined(_M_ARM64) || defined(__arm__) +#define PS_YIELD() __asm__ volatile("yield" ::: "memory") +#else +#define PS_YIELD() +#endif + +typedef volatile int ps_spinlock_t; + +#ifdef _WIN32 + +#define ps_spin_lock(lock) \ + while (InterlockedExchange((volatile LONG*)(lock), 1)) { \ + while (*(lock)) \ + PS_YIELD(); \ + } +#define ps_spin_unlock(lock) InterlockedExchange((volatile LONG*)(lock), 0) + +#else + +#define ps_spin_lock(lock) \ + while (__sync_lock_test_and_set((lock), 1)) { \ + while (*(lock)) \ + PS_YIELD(); \ + } + +#define ps_spin_unlock(lock) __sync_lock_release(lock) + +#endif + +#endif // PS_MULTITHREADING + // ===================================================================== // public api // ===================================================================== @@ -200,14 +291,18 @@ typedef enum { PS_OK = 0, PS_EOOM = -1, // out of memory PS_EINVAL = -2, // invalid argument + PS_ETHRD = -3, // thread creation failed + PS_EALLOC = -4, // allocation failed } ps_result_t; typedef struct ps_context ps_context_t; typedef struct { - void* buff; // ram provided by host - size_t buff_size; // total size in B - float theta; // MAC threshold. 0.0F defaults to 1.0F + void* buff; // ram provided by host + size_t buff_size; // total size in B + size_t max_particles; // maximum capacity for the solver + float theta; // MAC threshold. 0.0F defaults to 1.0F + size_t thrd_cnt; // 0 defaults to 1 (single-threaded impl) } ps_config_t; typedef struct { @@ -223,7 +318,7 @@ typedef struct { } ps_particle_arrs_t; // init pipeline with a pre-allocated buffer. -ps_result_t ps_init(ps_context_t** out_ctx, const ps_config_t* conf); +ps_result_t ps_init(ps_context_t** out_ctx, const ps_config_t* cfg); // compute forces on particles using FMM ps_result_t ps_calc_forces(ps_context_t* ctx, const ps_particle_arrs_t* arrs, @@ -236,8 +331,12 @@ ps_result_t ps_prepare_particles(ps_particle_arrs_t* arrs, uint32_t* out_morton_codes, float* out_min_b, float* out_max_b, float* out_range); +// shuts down bg threads safely before freeing memory +ps_result_t ps_destroy(ps_context_t* ctx); + #endif // POLESITTER_H +#define POLESITTER_IMPLEMENTATION #ifdef POLESITTER_IMPLEMENTATION // ===================================================================== @@ -248,13 +347,15 @@ ps_result_t ps_prepare_particles(ps_particle_arrs_t* arrs, // internal structs // --------------------------------------------------------------------- +// represents one threads memory slice typedef struct { - uint8_t* mem; - size_t cap; - size_t off; + uint8_t* mem; + size_t cap; + volatile size_t off; } ps_arena_t; // octree node +// exactly 64 bytes (1 cache line) typedef struct ps_node { // physics, 16B float x; @@ -264,28 +365,404 @@ typedef struct ps_node { // FMM payload, 32B (monopole + dipole) float multipole[4]; - float local[4]; + union { + uint32_t children_offs[8]; + + struct { + // leaf metadata, 16B + uint32_t is_leaf; + uint32_t particle_cnt; + uint32_t first_particle_idx; + uint32_t _pad[5]; + } leaf; + } data; +} ps_node_t; - // structure, 64B - struct ps_node* children[8]; +#ifndef PS_MAX_THRDS +#define PS_MAX_THRDS 32 +#endif // PS_MAX_THRDS - // leaf metadata, 16B - uint32_t is_leaf; - uint32_t particle_cnt; - uint32_t first_particle_idx; - uint32_t _pad; -} ps_node_t; +#ifndef PS_MAX_JOBS +#define PS_MAX_JOBS 512 +#endif // PS_MAX_JOBS + +typedef enum { + PS_JOB_EXIT = 0, + PS_JOB_UPWARD, + PS_JOB_INTERACTION, + PS_JOB_DOWNWARD, + PS_JOB_P2P, + PS_JOB_RADIX_MAP, + PS_JOB_RADIX_SCATTER, + PS_JOB_TREE_BUILD, +} ps_job_type_t; + +typedef struct { + ps_job_type_t type; + uint32_t thrd_id; // which worker is executing this + + union { + // for fmm passes + struct { + ps_node_t* target; + ps_node_t* src; + const ps_particle_arrs_t* arrs; + } fmm; + + // for radix + struct { + uint32_t start_idx; + uint32_t end_idx; + uint32_t chunk_id; + } array; + + // for tree + struct { + ps_node_t* node; + uint32_t start_idx; + uint32_t end_idx; + uint32_t depth; + } tree; + } data; +} ps_job_t; + +#ifdef PS_MULTITHREADING + +typedef struct { + ps_thrd_t thrds[PS_MAX_THRDS]; + uint32_t thrd_cnt; + + ps_job_t queue[PS_MAX_JOBS]; + volatile int hd; + volatile int tl; + volatile int cnt; + volatile int active_jobs; + volatile int shutdown_flag; + + ps_spinlock_t lock; +} ps_thrd_pool_t; + +// payload passed to each worker thread when it spawns +typedef struct { + ps_thrd_pool_t* pool; + ps_context_t* ctx; + uint32_t thrd_id; // 0 is main thrd +} ps_worker_arg_t; + +#endif // PS_MULTITHREADING + +#define PS_RADIX_BUF_SIZE 16 + +typedef struct { + uint32_t m[256][PS_RADIX_BUF_SIZE]; + uint32_t id[256][PS_RADIX_BUF_SIZE]; + float x[256][PS_RADIX_BUF_SIZE]; + float y[256][PS_RADIX_BUF_SIZE]; + float z[256][PS_RADIX_BUF_SIZE]; + float mass[256][PS_RADIX_BUF_SIZE]; + uint8_t cnt[256]; +} ps_scatter_buf_t; + +// holds the shared state for the radix pipeline +typedef struct { + uint32_t* m_src; + uint32_t* m_dst; + uint32_t* id_src; + uint32_t* id_dst; + float* x_src; + float* x_dst; + float* y_src; + float* y_dst; + float* z_src; + float* z_dst; + float* mass_src; + float* mass_dst; + + ps_scatter_buf_t* bufs; + uint8_t pass; // 0-3 + + // each thrd gets its own 256-bin histogram and off array + uint32_t hists[PS_MAX_THRDS + 1][256]; + uint32_t offs[PS_MAX_THRDS + 1][256]; +} ps_radix_state_t; struct ps_context { - ps_arena_t arena; ps_node_t* root; float theta; + ps_arena_t arena; + float (*local_exp)[4]; + ps_radix_state_t radix_state; + +#ifdef PS_MULTITHREADING + + ps_thrd_pool_t pool; + ps_worker_arg_t worker_args[PS_MAX_THRDS]; + +#endif // PS_MULTITHREADING }; // --------------------------------------------------------------------- // internal functions // --------------------------------------------------------------------- +// separated from the above for clarity purposes +#ifdef PS_MULTITHREADING + +#ifdef _WIN32 + +static inline void ps_mtx_init(ps_mtx_t* m) { + InitializeCriticalSection(m); +} +static inline void ps_mtx_destroy(ps_mtx_t* m) { + DeleteCriticalSection(m); +} +static inline void ps_mtx_lock(ps_mtx_t* m) { + EnterCriticalSection(m); +} +static inline void ps_mtx_unlock(ps_mtx_t* m) { + LeaveCriticalSection(m); +} + +static inline void ps_cond_init(ps_cond_t* c) { + InitializeConditionVariable(c); +} +static inline void ps_cond_destroy(ps_cond_t* c) { + (void)c; /* no-op */ +} +static inline void ps_cond_wait(ps_cond_t* c, ps_mtx_t* m) { + SleepConditionVariableCS(c, m, INFINITE); +} +static inline void ps_cond_signal(ps_cond_t* c) { + WakeConditionVariable(c); +} +static inline void ps_cond_bcast(ps_cond_t* c) { + WakeAllConditionVariable(c); +} + +static inline int ps_thrd_create(ps_thrd_t* t, ps_thrd_func_t f, void* a) { + *t = CreateThread(NULL, 0, f, a, 0, NULL); + return (*t != NULL) ? 0 : -1; +} +static inline void ps_thrd_join(ps_thrd_t t) { + WaitForSingleObject(t, INFINITE); + CloseHandle(t); +} + +#else // POSIX + +static inline void ps_mtx_init(ps_mtx_t* m) { + pthread_mutex_init(m, NULL); +} +static inline void ps_mtx_destroy(ps_mtx_t* m) { + pthread_mutex_destroy(m); +} +static inline void ps_mtx_lock(ps_mtx_t* m) { + pthread_mutex_lock(m); +} +static inline void ps_mtx_unlock(ps_mtx_t* m) { + pthread_mutex_unlock(m); +} + +static inline void ps_cond_init(ps_cond_t* c) { + pthread_cond_init(c, NULL); +} +static inline void ps_cond_destroy(ps_cond_t* c) { + pthread_cond_destroy(c); +} +static inline void ps_cond_wait(ps_cond_t* c, ps_mtx_t* m) { + pthread_cond_wait(c, m); +} +static inline void ps_cond_signal(ps_cond_t* c) { + pthread_cond_signal(c); +} +static inline void ps_cond_bcast(ps_cond_t* c) { + pthread_cond_broadcast(c); +} + +static inline int ps_thrd_create(ps_thrd_t* t, ps_thrd_func_t f, void* a) { + return pthread_create(t, NULL, f, a); +} +static inline void ps_thrd_join(ps_thrd_t t) { + pthread_join(t, NULL); +} + +#endif + +#endif // PS_MULTITHREADING + +// forward declare passes for router +static void ps_impl_fmm_upward_pass(ps_context_t* ctx, ps_node_t* node, + const ps_particle_arrs_t* arrs); +static void ps_impl_fmm_interaction_pass(ps_context_t* ctx, ps_node_t* target, + ps_node_t* src, float theta); +static void ps_impl_fmm_downward_pass(ps_context_t* ctx, ps_node_t* node, + const ps_particle_arrs_t* arrs); +static void ps_impl_fmm_p2p_pass(ps_context_t* ctx, ps_node_t* target, + ps_node_t* src, const ps_particle_arrs_t* arrs, + float theta); +static void ps_impl_radix_map(ps_context_t* ctx, uint32_t chunk_id, + size_t start_idx, size_t end_idx); +static void ps_impl_radix_scatter(ps_context_t* ctx, uint32_t chunk_id, + size_t start_idx, size_t end_idx); +static void ps_impl_build_tree(ps_context_t* ctx, uint32_t thrd_id, + ps_node_t* node, uint32_t depth, + size_t start_idx, size_t end_idx, + const uint32_t* morton_codes); + +static inline void ps_impl_exec_job(ps_context_t* ctx, ps_job_t* job) { + if (job->type == PS_JOB_UPWARD) { + ps_impl_fmm_upward_pass(ctx, job->data.fmm.src, job->data.fmm.arrs); + } else if (job->type == PS_JOB_INTERACTION) { + ps_impl_fmm_interaction_pass(ctx, job->data.fmm.target, + job->data.fmm.src, ctx->theta); + } else if (job->type == PS_JOB_DOWNWARD) { + ps_impl_fmm_downward_pass(ctx, job->data.fmm.src, job->data.fmm.arrs); + } else if (job->type == PS_JOB_P2P) { + ps_impl_fmm_p2p_pass(ctx, job->data.fmm.target, job->data.fmm.src, + job->data.fmm.arrs, ctx->theta); + } else if (job->type == PS_JOB_RADIX_MAP) { + ps_impl_radix_map(ctx, job->data.array.chunk_id, + job->data.array.start_idx, job->data.array.end_idx); + } else if (job->type == PS_JOB_RADIX_SCATTER) { + ps_impl_radix_scatter(ctx, job->data.array.chunk_id, + job->data.array.start_idx, + job->data.array.end_idx); + } else if (job->type == PS_JOB_TREE_BUILD) { + ps_impl_build_tree(ctx, job->thrd_id, job->data.tree.node, + job->data.tree.depth, job->data.tree.start_idx, + job->data.tree.end_idx, ctx->radix_state.m_src); + } +} + +static void ps_impl_pool_submit(ps_context_t* ctx, ps_job_t job) { +#ifdef PS_MULTITHREADING + + if (ctx->pool.thrd_cnt > 0) { + ps_thrd_pool_t* pool = &ctx->pool; + + // block if buffer is full + while (pool->cnt == PS_MAX_JOBS && !pool->shutdown_flag) { + PS_YIELD(); + } + + // enqueue + ps_spin_lock(&pool->lock); + pool->queue[pool->tl] = job; + pool->tl = (pool->tl + 1) % PS_MAX_JOBS; + pool->cnt++; + ps_spin_unlock(&pool->lock); + return; + } + +#endif + // single-threaded + + job.thrd_id = 0; + ps_impl_exec_job(ctx, &job); +} + +// no-op if single-threaded +static void ps_impl_pool_wait(ps_context_t* ctx) { + (void)ctx; +#ifdef PS_MULTITHREADING + ps_thrd_pool_t* pool = &ctx->pool; + + while (1) { + if (pool->cnt == 0 && pool->active_jobs == 0) { + + ps_spin_lock(&pool->lock); + int done = (pool->cnt == 0 && pool->active_jobs == 0); + ps_spin_unlock(&pool->lock); + + if (done) + return; + } + PS_YIELD(); + } +#endif +} + +#ifdef PS_MULTITHREADING + +static PS_THRD_RET_TYPE ps_impl_worker_loop(void* arg) { + ps_worker_arg_t* w_arg = (ps_worker_arg_t*)arg; + ps_thrd_pool_t* pool = w_arg->pool; + + while (1) { + while (pool->cnt == 0 && !pool->shutdown_flag) { + PS_YIELD(); + } + + ps_spin_lock(&pool->lock); + + if (pool->cnt == 0) { + ps_spin_unlock(&pool->lock); + if (pool->shutdown_flag) { + break; + } + + continue; + } + + // dequeue + ps_job_t job = pool->queue[pool->hd]; + pool->hd = (pool->hd + 1) % PS_MAX_JOBS; + + pool->active_jobs++; + pool->cnt--; + ps_spin_unlock(&pool->lock); + + // execute the job + job.thrd_id = w_arg->thrd_id; + ps_impl_exec_job(w_arg->ctx, &job); + + // mark done + ps_spin_lock(&pool->lock); + pool->active_jobs--; + ps_spin_unlock(&pool->lock); + } + + return PS_THRD_RET_VAL; +} + +static int ps_impl_pool_init(ps_context_t* ctx, uint32_t num_thrds) { + ps_thrd_pool_t* pool = &ctx->pool; + ps_worker_arg_t* worker_args = ctx->worker_args; + + pool->hd = 0; + pool->tl = 0; + pool->cnt = 0; + pool->active_jobs = 0; + pool->shutdown_flag = 0; + pool->thrd_cnt = num_thrds; + + if (num_thrds == 0) { + return PS_OK; + } + + if (pool->thrd_cnt > PS_MAX_THRDS) { + pool->thrd_cnt = PS_MAX_THRDS; + } + + pool->lock = 0; + + for (uint32_t i = 0; i < pool->thrd_cnt; ++i) { + worker_args[i].pool = pool; + worker_args[i].ctx = ctx; + worker_args[i].thrd_id = i + 1; + + if (ps_thrd_create(&pool->thrds[i], ps_impl_worker_loop, + &worker_args[i]) != 0) { + + return PS_ETHRD; + } + } + + return PS_OK; +} + +#endif // PS_MULTITHREADING + // take a 10b num and expand it to 30b by inserting 2 0s between each b. static uint32_t ps_impl_expand_bits(uint32_t v) { v &= 0x000003FF; // only look at 10 ls bits @@ -311,142 +788,190 @@ static uint32_t ps_impl_morton_encode(uint32_t x, uint32_t y, uint32_t z) { } // rounds 'v' up to the nearest multiple of 'align', which must be a power of 2 -static inline size_t ps_align_forward(size_t v, size_t align) { +static inline size_t ps_impl_align_forward(size_t v, size_t align) { return (v + (align - 1)) & ~(align - 1); } -// init arena with pre-allocated/externally provided buffer -static void ps_arena_init(ps_arena_t* arena, void* buffer, size_t cap) { - arena->mem = (uint8_t*)buffer; - arena->cap = cap; - arena->off = 0; -} +// alloc raw bytes +static void* ps_impl_arena_alloc(ps_context_t* ctx, size_t size) { + if (!ctx) { + return NULL; + } -// alloc raw bytes with a specified align (16 for SIMD) -static void* ps_arena_alloc(ps_arena_t* arena, size_t size, size_t align) { - // current aligned address offset - size_t curr_addr = (size_t)(arena->mem + arena->off); - size_t align_addr = ps_align_forward(curr_addr, align); + size = (size + 63) & ~63; - // how much extra padding was added to align the address - size_t pad = align_addr - curr_addr; +#ifdef PS_MULTITHREADING + size_t old_off = ps_atomic_fetch_add_size_t(&ctx->arena.off, size); +#else + size_t old_off = ctx->arena.off; + ctx->arena.off += size; +#endif // check if we have enough space in the arena - size_t total_req = size + pad; - if (arena->off + total_req > arena->cap) { + if (old_off + size > ctx->arena.cap) { return NULL; } // advance bump pointer and return aligned address - void* ptr = (void*)align_addr; - arena->off += total_req; - - return ptr; -} - -// clear for next frame -static void ps_arena_clear(ps_arena_t* arena) { - arena->off = 0; + return ctx->arena.mem + old_off; } -#define PS_MAX_DEPTH \ - 10 // max depth of octree, 10 levels = 1024 (2^10) leaf nodes +// helper to grab a new node +static inline ps_node_t* ps_impl_alloc_node(ps_context_t* ctx) { + if (!ctx) { + return NULL; + } -// zero out a newly allocated node -static void ps_impl_node_init(ps_node_t* node) { - node->x = 0.0F; - node->y = 0.0F; - node->z = 0.0F; - node->half_width = 0.0F; + size_t size = 64; - for (int i = 0; i < 4; ++i) { - node->multipole[i] = 0.0F; - node->local[i] = 0.0F; - } +#ifdef PS_MULTITHREADING + size_t old_off = ps_atomic_fetch_add_size_t(&ctx->arena.off, size); +#else + size_t old_off = ctx->arena.off; + ctx->arena.off += size; +#endif - for (int i = 0; i < 8; ++i) { - node->children[i] = NULL; + // check if we have enough space in the arena + if (old_off + size > ctx->arena.cap) { + return NULL; } - node->is_leaf = 0; - node->particle_cnt = 0; - node->first_particle_idx = 0; -} + ps_node_t* node = (ps_node_t*)(ctx->arena.mem + old_off); -// walk the morton code and build the tree branches -static ps_result_t ps_impl_tree_insert(ps_arena_t* arena, ps_node_t* root, - uint32_t morton_code, - uint32_t particle_idx) { - if (!root) { - return PS_EINVAL; - } + memset(node, 0, 64); - ps_node_t* curr = root; + uint32_t idx = (uint32_t)(old_off / 64); + ctx->local_exp[idx][0] = 0.0F; + ctx->local_exp[idx][1] = 0.0F; + ctx->local_exp[idx][2] = 0.0F; + ctx->local_exp[idx][3] = 0.0F; - float curr_x = root->x; - float curr_y = root->y; - float curr_z = root->z; - float curr_hw = root->half_width; + return node; +} - for (int depth = 0; depth < PS_MAX_DEPTH; ++depth) { - // shift starts at 27, decreases by 3 each lvl - int shift = 27 - (depth * 3); - uint32_t octant = (morton_code >> shift) & 0x7; // 3 bits for octant +static inline ps_node_t* ps_impl_get_node(ps_context_t* ctx, uint32_t off) { + if (off == 0) { + return NULL; + } - curr_hw *= 0.5F; + return (ps_node_t*)(ctx->arena.mem + off); +} - curr_x += (octant & 0x1 /* bit 0 */) ? curr_hw : -curr_hw; - curr_y += (octant & 0x2 /* bit 1 */) ? curr_hw : -curr_hw; - curr_z += (octant & 0x4 /* bit 2 */) ? curr_hw : -curr_hw; +static inline float* ps_impl_get_local(ps_context_t* ctx, ps_node_t* node) { + uint32_t idx = (uint32_t)(((uint8_t*)node - ctx->arena.mem) / 64); - if (!curr->children[octant]) { - ps_node_t* new_node = - (ps_node_t*)ps_arena_alloc(arena, sizeof(ps_node_t), 32); - if (!new_node) { - return PS_EOOM; - } + return ctx->local_exp[idx]; +} - ps_impl_node_init(new_node); +// clear for next frame +static void ps_impl_arena_clear(ps_arena_t* arena) { + arena->off = 0; +} - new_node->x = curr_x; - new_node->y = curr_y; - new_node->z = curr_z; - new_node->half_width = curr_hw; +#define PS_MAX_DEPTH \ + 10 // max depth of octree, 10 levels = 1024 (2^10) leaf nodes - curr->children[octant] = new_node; +// bin search to find the idx where target oct begins +static inline size_t ps_impl_find_split(const uint32_t* codes, size_t start, + size_t end, int shift, + uint32_t target_oct) { + size_t left = start; + size_t right = end; + while (left < right) { + size_t mid = left + ((right - left) / 2); + uint32_t oct = (codes[mid] >> shift) & 0x7; + + if (oct < target_oct) { + left = mid + 1; + } else { + right = mid; } - - curr = curr->children[octant]; } - curr->is_leaf = 1; + return left; +} - if (curr->particle_cnt == 0) { - curr->first_particle_idx = particle_idx; +// walk the morton code and build the tree branches +static void ps_impl_build_tree(ps_context_t* ctx, uint32_t thrd_id, + ps_node_t* node, uint32_t depth, + size_t start_idx, size_t end_idx, + const uint32_t* morton_codes) { + // max depth reached or only 1 particle left + if (depth == PS_MAX_DEPTH || end_idx - start_idx <= 1) { + node->data.leaf.is_leaf = 1; + node->data.leaf.first_particle_idx = (uint32_t)start_idx; + node->data.leaf.particle_cnt = (uint32_t)(end_idx - start_idx); + return; } - curr->particle_cnt++; + // shift starts at 27, decreases by 3 each lvl + uint8_t shift = (uint8_t)(27 - (depth * 3)); + size_t curr_start = start_idx; + + // subdivide into 8 octants + for (uint32_t oct = 0; oct < 8; ++oct) { + // find where this octnat ends in the sorted arr + size_t oct_end = ps_impl_find_split(morton_codes, curr_start, end_idx, + shift, oct + 1); + + // if this octant has particles build a branch + if (oct_end > curr_start) { + ps_node_t* child = ps_impl_alloc_node(ctx); + if (!child) { + return; + } + + float hw = node->half_width * 0.5F; + child->x = node->x + ((oct & 1) ? hw : -hw); + child->y = node->y + ((oct & 2) ? hw : -hw); + child->z = node->z + ((oct & 4) ? hw : -hw); + child->half_width = hw; + + node->data.children_offs[oct] = + (uint32_t)((uint8_t*)child - ctx->arena.mem); + +// spawn job for top layers +#ifdef PS_MULTITHREADING + if (depth == 0) { + ps_job_t job; + job.type = PS_JOB_TREE_BUILD; + job.data.tree.node = child; + job.data.tree.start_idx = (uint32_t)curr_start; + job.data.tree.end_idx = (uint32_t)oct_end; + job.data.tree.depth = depth + 1; + ps_impl_pool_submit(ctx, job); + } else +#endif // PS_MULTITHREADING + { + ps_impl_build_tree(ctx, thrd_id, child, depth + 1, curr_start, + oct_end, morton_codes); + } + } - return PS_OK; + curr_start = oct_end; + if (curr_start == end_idx) { + break; // all particles handled + } + } } // pass 1 // p2m (leaves) and m2m (parents) in post-order traversal -static void ps_impl_fmm_upward_pass(ps_node_t* node, +static void ps_impl_fmm_upward_pass(ps_context_t* ctx, ps_node_t* node, const ps_particle_arrs_t* arrs) { if (!node) { return; } // p2m - if (node->is_leaf) { + if (node->data.leaf.is_leaf & 1) { float m0 = 0.0F; // monopole (total mass) float mx = 0.0F; // dipole x (mass moment) float my = 0.0F; // dipole y float mz = 0.0F; // dipole z - for (uint32_t i = 0; i < node->particle_cnt; ++i) { - uint32_t idx = node->first_particle_idx + i; + for (uint32_t i = 0; i < node->data.leaf.particle_cnt; ++i) { + uint32_t idx = node->data.leaf.first_particle_idx + i; float mass = arrs->mass[idx]; // dist from particle pos to center of this voxel @@ -475,29 +1000,31 @@ static void ps_impl_fmm_upward_pass(ps_node_t* node, float p_mz = 0.0F; for (int i = 0; i < 8; ++i) { - ps_node_t* child = node->children[i]; - if (child) { - // calc the children first - ps_impl_fmm_upward_pass(child, arrs); - - // dist vector from the child's center - // up to the parent's center - float dx = child->x - node->x; - float dy = child->y - node->y; - float dz = child->z - node->z; - - float c_m0 = child->multipole[0]; - float c_mx = child->multipole[1]; - float c_my = child->multipole[2]; - float c_mz = child->multipole[3]; - - // shift the childs expansion to the parents center and accumulate - // dipole requires the monopole * dist - p_m0 += c_m0; - p_mx += c_mx + (c_m0 * dx); - p_my += c_my + (c_m0 * dy); - p_mz += c_mz + (c_m0 * dz); + ps_node_t* child = ps_impl_get_node(ctx, node->data.children_offs[i]); + if (!child) { + continue; } + + // calc the children first + ps_impl_fmm_upward_pass(ctx, child, arrs); + + // dist vector from the child's center + // up to the parent's center + float dx = child->x - node->x; + float dy = child->y - node->y; + float dz = child->z - node->z; + + float c_m0 = child->multipole[0]; + float c_mx = child->multipole[1]; + float c_my = child->multipole[2]; + float c_mz = child->multipole[3]; + + // shift the childs expansion to the parents center and accumulate + // dipole requires the monopole * dist + p_m0 += c_m0; + p_mx += c_mx + (c_m0 * dx); + p_my += c_my + (c_m0 * dy); + p_mz += c_mz + (c_m0 * dz); } // store aggregated expansion in parent @@ -509,8 +1036,8 @@ static void ps_impl_fmm_upward_pass(ps_node_t* node, // pass 2 // m2l dual-tree traversal to find well-separated nodes and translate expansions -static void ps_impl_fmm_interaction_pass(ps_node_t* target, ps_node_t* src, - float theta) { +static void ps_impl_fmm_interaction_pass(ps_context_t* ctx, ps_node_t* target, + ps_node_t* src, float theta) { if (!target || !src) { return; } @@ -548,82 +1075,80 @@ static void ps_impl_fmm_interaction_pass(ps_node_t* target, ps_node_t* src, float force_dip_z = (mz * inv_r3) - (dz * dipole_coeff); // accumulate into target's local expansion - target->local[1] += force_m0_x + force_dip_x; - target->local[2] += force_m0_y + force_dip_y; - target->local[3] += force_m0_z + force_dip_z; + float* t_local = ps_impl_get_local(ctx, target); + t_local[1] += force_m0_x + force_dip_x; + t_local[2] += force_m0_y + force_dip_y; + t_local[3] += force_m0_z + force_dip_z; return; } // not well-separated - if (target->is_leaf && src->is_leaf) { + if ((target->data.leaf.is_leaf & 1) && (src->data.leaf.is_leaf & 1)) { // both are leaves, near-field p2p pass will handle exact dists return; } - if (target->is_leaf) { + if (target->data.leaf.is_leaf & 1) { // target is as small as possible, open the src for (int i = 0; i < 8; ++i) { - ps_impl_fmm_interaction_pass(target, src->children[i], theta); + ps_node_t* src_child = + ps_impl_get_node(ctx, src->data.children_offs[i]); + if (!src_child) { + continue; + } + + ps_impl_fmm_interaction_pass(ctx, target, src_child, theta); } - } else if (src->is_leaf) { + } else if (src->data.leaf.is_leaf & 1) { // src is as small as possible, open the target for (int i = 0; i < 8; ++i) { - ps_impl_fmm_interaction_pass(target->children[i], src, theta); + ps_node_t* target_child = + ps_impl_get_node(ctx, target->data.children_offs[i]); + if (!target_child) { + continue; + } + + ps_impl_fmm_interaction_pass(ctx, target_child, src, theta); } } else { // subdivide both and pair all 64 permutations for (int i = 0; i < 8; ++i) { - if (!target->children[i]) { + ps_node_t* target_child = + ps_impl_get_node(ctx, target->data.children_offs[i]); + if (!target_child) { continue; } for (int j = 0; j < 8; ++j) { - if (!src->children[j]) { + ps_node_t* src_child = + ps_impl_get_node(ctx, src->data.children_offs[j]); + if (!src_child) { continue; } - ps_impl_fmm_interaction_pass(target->children[i], - src->children[j], theta); + ps_impl_fmm_interaction_pass(ctx, target_child, src_child, + theta); } } } } -// pass 3 +// pass 3/4 // l2l, pushes the accumulated background field from parents down to their // children -static void ps_impl_fmm_downward_pass(ps_node_t* node) { - if (!node || node->is_leaf) { - return; - } - - for (int i = 0; i < 8; ++i) { - ps_node_t* child = node->children[i]; - if (child) { - // 1-st order local expansion, - // shifting it is adding the parent's field to the child's field - child->local[1] += node->local[1]; - child->local[2] += node->local[2]; - child->local[3] += node->local[3]; - - ps_impl_fmm_downward_pass(child); - } - } -} - -// pass 4 // l2p, applies the accumulated bg field to the particles inside the leaf -static void ps_impl_fmm_l2p_pass(ps_node_t* node, - const ps_particle_arrs_t* arrs) { +static void ps_impl_fmm_downward_pass(ps_context_t* ctx, ps_node_t* node, + const ps_particle_arrs_t* arrs) { if (!node) { return; } - if (node->is_leaf) { - float field_x = node->local[1]; - float field_y = node->local[2]; - float field_z = node->local[3]; + if (node->data.leaf.is_leaf & 1) { + float* t_local = ps_impl_get_local(ctx, node); + float field_x = t_local[1]; + float field_y = t_local[2]; + float field_z = t_local[3]; uint32_t i = 0; @@ -634,8 +1159,8 @@ static void ps_impl_fmm_l2p_pass(ps_node_t* node, __m256 f_z_vec = _mm256_set1_ps(field_z); // process in chunks of 8 - for (; i + 7 < node->particle_cnt; i += 8) { - uint32_t idx = node->first_particle_idx + i; + for (; i + 7 < node->data.leaf.particle_cnt; i += 8) { + uint32_t idx = node->data.leaf.first_particle_idx + i; // load 8 masses __m256 m_vec = _mm256_loadu_ps(&arrs->mass[idx]); @@ -662,8 +1187,8 @@ static void ps_impl_fmm_l2p_pass(ps_node_t* node, float32x4_t f_z_vec = vdupq_n_f32(field_z); // process in chunks of 4 - for (; i + 3 < node->particle_cnt; i += 4) { - uint32_t idx = node->first_particle_idx + i; + for (; i + 3 < node->data.leaf.particle_cnt; i += 4) { + uint32_t idx = node->data.leaf.first_particle_idx + i; // load 4 masses float32x4_t m_vec = vld1q_f32(&arrs->mass[idx]); @@ -686,8 +1211,8 @@ static void ps_impl_fmm_l2p_pass(ps_node_t* node, #endif // fallback for the remainder - for (; i < node->particle_cnt; ++i) { - uint32_t idx = node->first_particle_idx + i; + for (; i < node->data.leaf.particle_cnt; ++i) { + uint32_t idx = node->data.leaf.first_particle_idx + i; float mass = arrs->mass[idx]; // F = m * a @@ -699,16 +1224,29 @@ static void ps_impl_fmm_l2p_pass(ps_node_t* node, return; } - // cascade for (int i = 0; i < 8; ++i) { - ps_impl_fmm_l2p_pass(node->children[i], arrs); + ps_node_t* child = ps_impl_get_node(ctx, node->data.children_offs[i]); + if (!child) { + continue; + } + + // 1-st order local expansion, + // shifting it is adding the parent's field to the child's field + float* child_local = ps_impl_get_local(ctx, child); + const float* node_local = ps_impl_get_local(ctx, node); + child_local[1] += node_local[1]; + child_local[2] += node_local[2]; + child_local[3] += node_local[3]; + + ps_impl_fmm_downward_pass(ctx, child, arrs); } } // pass 5 // p2p, dual-tree traversal to calc N-body forces for near-field neighbors -static void ps_impl_fmm_p2p_pass(ps_node_t* target, ps_node_t* src, - const ps_particle_arrs_t* arrs) { +static void ps_impl_fmm_p2p_pass(ps_context_t* ctx, ps_node_t* target, + ps_node_t* src, const ps_particle_arrs_t* arrs, + float theta) { if (!target || !src) { return; } @@ -718,19 +1256,18 @@ static void ps_impl_fmm_p2p_pass(ps_node_t* target, ps_node_t* src, float dz = src->z - target->z; float dist_sq = (dx * dx) + (dy * dy) + (dz * dz); - float theta = 1.0F; float hw_sum = target->half_width + src->half_width; // if well-separated M2L already handled it - if (dist_sq > 12.0F * (theta * theta) * (hw_sum * hw_sum)) { + if (dist_sq > (theta * theta) * (hw_sum * hw_sum)) { return; } // if both are leaves and too close // do direct N-body force - if (target->is_leaf && src->is_leaf) { - for (uint32_t i = 0; i < target->particle_cnt; ++i) { - uint32_t t_idx = target->first_particle_idx + i; + if ((target->data.leaf.is_leaf & 1) && (src->data.leaf.is_leaf & 1)) { + for (uint32_t i = 0; i < target->data.leaf.particle_cnt; ++i) { + uint32_t t_idx = target->data.leaf.first_particle_idx + i; float t_x = arrs->x[t_idx]; float t_y = arrs->y[t_idx]; float t_z = arrs->z[t_idx]; @@ -759,8 +1296,8 @@ static void ps_impl_fmm_p2p_pass(ps_node_t* target, ps_node_t* src, __m256 f_z_vec = _mm256_setzero_ps(); // process in chunks of 8 - for (; j + 7 < src->particle_cnt; j += 8) { - uint32_t s_idx = src->first_particle_idx + j; + for (; j + 7 < src->data.leaf.particle_cnt; j += 8) { + uint32_t s_idx = src->data.leaf.first_particle_idx + j; // load 8 source coordinates and masses __m256 s_x_vec = _mm256_loadu_ps(&sx[s_idx]); @@ -820,8 +1357,8 @@ static void ps_impl_fmm_p2p_pass(ps_node_t* target, ps_node_t* src, float32x4_t f_z_vec = vdupq_n_f32(0.0F); // process in chunks of 4 - for (; j + 3 < src->particle_cnt; j += 4) { - uint32_t s_idx = src->first_particle_idx + j; + for (; j + 3 < src->data.leaf.particle_cnt; j += 4) { + uint32_t s_idx = src->data.leaf.first_particle_idx + j; // load 4 source coordinates and masses float32x4_t s_x_vec = vld1q_f32(&sx[s_idx]); @@ -873,8 +1410,8 @@ static void ps_impl_fmm_p2p_pass(ps_node_t* target, ps_node_t* src, #endif // fallback for the remainder (previous operations left n in mod 8 // particles) - for (; j < src->particle_cnt; ++j) { - uint32_t s_idx = src->first_particle_idx + j; + for (; j < src->data.leaf.particle_cnt; ++j) { + uint32_t s_idx = src->data.leaf.first_particle_idx + j; float p_dx = sx[s_idx] - t_x; float p_dy = sy[s_idx] - t_y; @@ -904,32 +1441,132 @@ static void ps_impl_fmm_p2p_pass(ps_node_t* target, ps_node_t* src, } // otherwise subdivide and recurse (like in m2l) - if (target->is_leaf) { + if (target->data.leaf.is_leaf & 1) { for (int i = 0; i < 8; ++i) { - ps_impl_fmm_p2p_pass(target, src->children[i], arrs); + ps_node_t* src_child = + ps_impl_get_node(ctx, src->data.children_offs[i]); + if (!src_child) { + continue; + } + + ps_impl_fmm_p2p_pass(ctx, target, src_child, arrs, theta); } - } else if (src->is_leaf) { + } else if (src->data.leaf.is_leaf & 1) { for (int i = 0; i < 8; ++i) { - ps_impl_fmm_p2p_pass(target->children[i], src, arrs); + ps_node_t* target_child = + ps_impl_get_node(ctx, target->data.children_offs[i]); + if (!target_child) { + continue; + } + + ps_impl_fmm_p2p_pass(ctx, target_child, src, arrs, theta); } } else { for (int i = 0; i < 8; ++i) { - if (!target->children[i]) { + ps_node_t* target_child = + ps_impl_get_node(ctx, target->data.children_offs[i]); + if (!target_child) { continue; } for (int j = 0; j < 8; ++j) { - if (!src->children[j]) { + ps_node_t* src_child = + ps_impl_get_node(ctx, src->data.children_offs[j]); + if (!src_child) { continue; } - ps_impl_fmm_p2p_pass(target->children[i], src->children[j], - arrs); + ps_impl_fmm_p2p_pass(ctx, target_child, src_child, arrs, theta); } } } } +// generates a histogram for a threads chunk of sorted memory +static void ps_impl_radix_map(ps_context_t* ctx, uint32_t chunk_id, + size_t start_idx, size_t end_idx) { + int shift = ctx->radix_state.pass * 8; + const uint32_t* m_src = ctx->radix_state.m_src; + uint32_t* hist = ctx->radix_state.hists[chunk_id]; + + // count frequencies for this threads chunk + for (size_t i = start_idx; i < end_idx; ++i) { + uint8_t bucket = (uint8_t)((m_src[i] >> shift) & 0xFF); + hist[bucket]++; + } +} + +static void ps_impl_radix_scatter(ps_context_t* ctx, uint32_t chunk_id, + size_t start_idx, size_t end_idx) { + int shift = ctx->radix_state.pass * 8; + ps_radix_state_t* rs = &ctx->radix_state; + ps_scatter_buf_t* buf = &rs->bufs[chunk_id]; + + // reset buffer counts for this pass + for (int i = 0; i < 256; ++i) { + buf->cnt[i] = 0; + } + + // make a local copy of this threads start offsets + // so we can increment them safely in regs + uint32_t local_offs[256]; + for (int i = 0; i < 256; ++i) { + local_offs[i] = rs->offs[chunk_id][i]; + } + + for (size_t i = start_idx; i < end_idx; ++i) { + uint8_t bucket = (uint8_t)((rs->m_src[i] >> shift) & 0xFF); + uint8_t c = buf->cnt[bucket]; + + buf->m[bucket][c] = rs->m_src[i]; + buf->id[bucket][c] = rs->id_src[i]; + buf->x[bucket][c] = rs->x_src[i]; + buf->y[bucket][c] = rs->y_src[i]; + buf->z[bucket][c] = rs->z_src[i]; + buf->mass[bucket][c] = rs->mass_src[i]; + + c++; + + // if buffer is full, flush to main memory in one + if (c == PS_RADIX_BUF_SIZE) { + uint32_t dst_idx = local_offs[bucket]; + + for (int j = 0; j < PS_RADIX_BUF_SIZE; ++j) { + rs->m_dst[dst_idx + j] = buf->m[bucket][j]; + rs->id_dst[dst_idx + j] = buf->id[bucket][j]; + rs->x_dst[dst_idx + j] = buf->x[bucket][j]; + rs->y_dst[dst_idx + j] = buf->y[bucket][j]; + rs->z_dst[dst_idx + j] = buf->z[bucket][j]; + rs->mass_dst[dst_idx + j] = buf->mass[bucket][j]; + } + + local_offs[bucket] += PS_RADIX_BUF_SIZE; + buf->cnt[bucket] = 0; + } else { + buf->cnt[bucket] = c; + } + } + + for (int bucket = 0; bucket < 256; ++bucket) { + uint8_t rem = buf->cnt[bucket]; + + if (rem > 0) { + uint32_t dst_idx = local_offs[bucket]; + + for (int j = 0; j < rem; ++j) { + rs->m_dst[dst_idx + j] = buf->m[bucket][j]; + rs->id_dst[dst_idx + j] = buf->id[bucket][j]; + rs->x_dst[dst_idx + j] = buf->x[bucket][j]; + rs->y_dst[dst_idx + j] = buf->y[bucket][j]; + rs->z_dst[dst_idx + j] = buf->z[bucket][j]; + rs->mass_dst[dst_idx + j] = buf->mass[bucket][j]; + } + + local_offs[bucket] += rem; + } + } +} + #define PS_IMPL_SWAP_PTR(type, a, b) \ do { \ type tmp = a; \ @@ -937,86 +1574,105 @@ static void ps_impl_fmm_p2p_pass(ps_node_t* target, ps_node_t* src, (b) = tmp; \ } while (0) -static ps_result_t ps_impl_sort_particles(ps_arena_t* arena, - uint32_t* morton_codes, +static ps_result_t ps_impl_sort_particles(ps_context_t* ctx, + uint32_t* morton_codes, const ps_particle_arrs_t* arrs) { size_t cnt = arrs->cnt; if (cnt == 0) { return PS_OK; } - // used later to reclaim temp memory - size_t tmp_off = arena->off; + ps_radix_state_t* rs = &ctx->radix_state; + rs->m_src = morton_codes; + rs->id_src = arrs->id; + rs->x_src = arrs->x; + rs->y_src = arrs->y; + rs->z_src = arrs->z; + rs->mass_src = arrs->mass; + + // how many chunks split into +#ifdef PS_MULTITHREADING + uint32_t num_chunks = (ctx->pool.thrd_cnt > 0) ? ctx->pool.thrd_cnt : 1; +#else + uint32_t num_chunks = 1; +#endif + +#define DEPLOY_JOB(job_type) \ + for (uint32_t i = 0; i < num_chunks; ++i) { \ + size_t start = i * chunk_size; \ + size_t end = start + chunk_size; \ + if (end > cnt) { \ + end = cnt; \ + } \ + if (start >= cnt) { \ + break; \ + } \ + ps_job_t job; \ + job.type = job_type; \ + job.data.array.start_idx = (uint32_t)start; \ + job.data.array.end_idx = (uint32_t)end; \ + job.data.array.chunk_id = i; \ + ps_impl_pool_submit(ctx, job); \ + } - // borrow temp SoA buffers from the arena - uint32_t* m_tmp = - (uint32_t*)ps_arena_alloc(arena, cnt * sizeof(uint32_t), 32); - uint32_t* id_tmp = - (uint32_t*)ps_arena_alloc(arena, cnt * sizeof(uint32_t), 32); - float* x_tmp = (float*)ps_arena_alloc(arena, cnt * sizeof(float), 32); - float* y_tmp = (float*)ps_arena_alloc(arena, cnt * sizeof(float), 32); - float* z_tmp = (float*)ps_arena_alloc(arena, cnt * sizeof(float), 32); - float* mass_tmp = (float*)ps_arena_alloc(arena, cnt * sizeof(float), 32); + // ceiling div so last chunk takes remainder + size_t chunk_size = (cnt + num_chunks - 1) / num_chunks; - if (!m_tmp || !x_tmp || !y_tmp || !z_tmp || !mass_tmp) { + size_t pre_sort_off = ctx->arena.off; + rs->bufs = (ps_scatter_buf_t*)ps_impl_arena_alloc( + ctx, num_chunks * sizeof(ps_scatter_buf_t)); + + if (!rs->bufs) { + ctx->arena.off = pre_sort_off; return PS_EOOM; } - uint32_t* m_src = morton_codes; - uint32_t* m_dst = m_tmp; - uint32_t* id_src = arrs->id; - uint32_t* id_dst = id_tmp; - float* x_src = arrs->x; - float* x_dst = x_tmp; - float* y_src = arrs->y; - float* y_dst = y_tmp; - float* z_src = arrs->z; - float* z_dst = z_tmp; - float* mass_src = arrs->mass; - float* mass_dst = mass_tmp; - // 4 passes of 8b radix sort for (int pass = 0; pass < 4; ++pass) { - uint32_t counts[256] = {0}; - uint32_t offsets[256] = {0}; - int shift = pass * 8; - - // count frequencies of current 8b chunk - for (size_t i = 0; i < cnt; ++i) { - uint8_t bucket = (m_src[i] >> shift) & 0xFF; - counts[bucket]++; - } + rs->pass = (uint8_t)pass; - // calc prefix sums to find the start off per bucket - offsets[0] = 0; - for (int i = 1; i < 256; ++i) { - offsets[i] = offsets[i - 1] + counts[i - 1]; + // clear histograms on main thread + for (int bucket = 0; bucket < 256; ++bucket) { + for (uint32_t chunk = 0; chunk < num_chunks; ++chunk) { + rs->hists[chunk][bucket] = 0; + } } - // distribute the data into dest buffers - for (size_t i = 0; i < cnt; ++i) { - uint8_t bucket = (m_src[i] >> shift) & 0xFF; - uint32_t dst_idx = offsets[bucket]++; - - m_dst[dst_idx] = m_src[i]; - id_dst[dst_idx] = id_src[i]; - x_dst[dst_idx] = x_src[i]; - y_dst[dst_idx] = y_src[i]; - z_dst[dst_idx] = z_src[i]; - mass_dst[dst_idx] = mass_src[i]; + // local histograms/map + DEPLOY_JOB(PS_JOB_RADIX_MAP); + + // wait for all threads to finish mapping + ps_impl_pool_wait(ctx); + + // reduce + // calculates exactly where each thread should start writing for each + // bucket + uint32_t run_off = 0; + for (int bucket = 0; bucket < 256; ++bucket) { + for (uint32_t chunk = 0; chunk < num_chunks; ++chunk) { + rs->offs[chunk][bucket] = run_off; + run_off += rs->hists[chunk][bucket]; + } } - // pingpong pointers for next pass - PS_IMPL_SWAP_PTR(uint32_t*, m_src, m_dst); - PS_IMPL_SWAP_PTR(uint32_t*, id_src, id_dst); - PS_IMPL_SWAP_PTR(float*, x_src, x_dst); - PS_IMPL_SWAP_PTR(float*, y_src, y_dst); - PS_IMPL_SWAP_PTR(float*, z_src, z_dst); - PS_IMPL_SWAP_PTR(float*, mass_src, mass_dst); + // scatter + DEPLOY_JOB(PS_JOB_RADIX_SCATTER); + + // wait for all data to be scattered + ps_impl_pool_wait(ctx); + + // ping-pong pointers + PS_IMPL_SWAP_PTR(uint32_t*, rs->m_src, rs->m_dst); + PS_IMPL_SWAP_PTR(uint32_t*, rs->id_src, rs->id_dst); + PS_IMPL_SWAP_PTR(float*, rs->x_src, rs->x_dst); + PS_IMPL_SWAP_PTR(float*, rs->y_src, rs->y_dst); + PS_IMPL_SWAP_PTR(float*, rs->z_src, rs->z_dst); + PS_IMPL_SWAP_PTR(float*, rs->mass_src, rs->mass_dst); } - // reclaim memory - arena->off = tmp_off; +#undef DEPLOY_JOB + + ctx->arena.off = pre_sort_off; // 4 is an even number, so m_src is guaranteed to be pointing back to the // orig morton_codes arr, and x_src back to arrs->x. the sorted data is @@ -1029,25 +1685,147 @@ static ps_result_t ps_impl_sort_particles(ps_arena_t* arena, // public api implementation // ===================================================================== -ps_result_t ps_init(ps_context_t** out_ctx, const ps_config_t* conf) { - if (!conf || !conf->buff || conf->buff_size < sizeof(ps_context_t)) { +ps_result_t ps_init(ps_context_t** out_ctx, const ps_config_t* cfg) { + if (!cfg || !cfg->buff || cfg->buff_size < sizeof(ps_context_t)) { return PS_EINVAL; } // place ctx at beginning of buffer - ps_context_t* ctx = (ps_context_t*)conf->buff; + ps_context_t* ctx = (ps_context_t*)cfg->buff; ctx->root = NULL; - ctx->theta = (conf->theta > 0.0F) ? conf->theta : 1.0F; + ctx->theta = (cfg->theta > 0.0F) ? cfg->theta : 1.0F; + + uint32_t total_thrds = (uint32_t)cfg->thrd_cnt; + if (total_thrds > PS_MAX_THRDS + 1) { + total_thrds = PS_MAX_THRDS + 1; + } else if (total_thrds < 1) { + total_thrds = 1; + } + + // calculate space needed for temp radix arrs + size_t radix_tmp_size = 6 * cfg->max_particles * sizeof(float); + + // ensure we have enough mem + size_t arena_start = ps_impl_align_forward(sizeof(ps_context_t), 64); + if (arena_start + radix_tmp_size > cfg->buff_size) { + return PS_EOOM; + } + + // temp arrs go at end of the buffer + uint8_t* radix_mem = (uint8_t*)cfg->buff + cfg->buff_size - radix_tmp_size; + + // wire permanent ptrs to radix_state + ctx->radix_state.m_dst = + (uint32_t*)(radix_mem + (0 * cfg->max_particles * sizeof(float))); + ctx->radix_state.id_dst = + (uint32_t*)(radix_mem + (1 * cfg->max_particles * sizeof(float))); + ctx->radix_state.x_dst = + (float*)(radix_mem + (2 * cfg->max_particles * sizeof(float))); + ctx->radix_state.y_dst = + (float*)(radix_mem + (3 * cfg->max_particles * sizeof(float))); + ctx->radix_state.z_dst = + (float*)(radix_mem + (4 * cfg->max_particles * sizeof(float))); + ctx->radix_state.mass_dst = + (float*)(radix_mem + (5 * cfg->max_particles * sizeof(float))); + + size_t usable = cfg->buff_size - arena_start - radix_tmp_size; + size_t max_nodes = usable / (sizeof(ps_node_t) + 16); + + // global arena gets remaining space + ctx->arena.mem = (uint8_t*)cfg->buff + arena_start; + ctx->arena.cap = max_nodes * sizeof(ps_node_t); + ctx->arena.off = 0; + + // clang-format off - // arena takes the rest - size_t arena_start = ps_align_forward(sizeof(ps_context_t), 16); - ps_arena_init(&ctx->arena, (uint8_t*)conf->buff + arena_start, - conf->buff_size - arena_start); + ctx->local_exp = (float (*)[4])(ctx->arena.mem + ctx->arena.cap); + + // clang-format on + +#ifdef PS_MULTITHREADING + + // initialize the pool + if (ps_impl_pool_init(ctx, total_thrds - 1) != 0) { + return PS_EALLOC; + } + +#endif // PS_MULTITHREADING *out_ctx = ctx; return PS_OK; } +#ifdef PS_MULTITHREADING + +// dispatches target nodes for M2L/P2P passes +static void ps_impl_dispatch_fmm(ps_context_t* ctx, ps_job_type_t job_type, + ps_node_t* target, ps_node_t* src, + const ps_particle_arrs_t* arrs, int curr_depth, + int target_depth) { + if (!target) { + return; + } + + if (curr_depth == target_depth || (target->data.leaf.is_leaf & 1)) { + ps_job_t job; + job.type = job_type; + job.data.fmm.target = target; + job.data.fmm.src = src; + job.data.fmm.arrs = arrs; + ps_impl_pool_submit(ctx, job); + return; + } + + // otherwise keep traversing down + for (int i = 0; i < 8; ++i) { + ps_node_t* target_child = + ps_impl_get_node(ctx, target->data.children_offs[i]); + if (!target_child) { + continue; + } + + ps_impl_dispatch_fmm(ctx, job_type, target_child, src, arrs, + curr_depth + 1, target_depth); + } +} + +// dispatches nodes for downward pass +static void ps_impl_dispatch_downward(ps_context_t* ctx, ps_node_t* target, + const ps_particle_arrs_t* arrs, + int curr_depth, int target_depth) { + if (!target) { + return; + } + + if (curr_depth == target_depth || (target->data.leaf.is_leaf & 1)) { + ps_job_t job; + job.type = PS_JOB_DOWNWARD; + job.data.fmm.src = target; + job.data.fmm.arrs = arrs; + ps_impl_pool_submit(ctx, job); + return; + } + + // above target depth, manually push local field down before dispatch + for (int i = 0; i < 8; ++i) { + ps_node_t* child = ps_impl_get_node(ctx, target->data.children_offs[i]); + if (!child) { + continue; + } + + float* child_local = ps_impl_get_local(ctx, child); + float* target_local = ps_impl_get_local(ctx, target); + + child_local[1] += target_local[1]; + child_local[2] += target_local[2]; + child_local[3] += target_local[3]; + ps_impl_dispatch_downward(ctx, child, arrs, curr_depth + 1, + target_depth); + } +} + +#endif // PS_MULTITHREADING + ps_result_t ps_calc_forces(ps_context_t* ctx, const ps_particle_arrs_t* arrs, uint32_t* morton_codes, float root_cx, float root_cy, float root_cz, float root_hw) { @@ -1056,18 +1834,16 @@ ps_result_t ps_calc_forces(ps_context_t* ctx, const ps_particle_arrs_t* arrs, } // clear arena for new frame - ps_arena_clear(&ctx->arena); + ps_impl_arena_clear(&ctx->arena); // sort the arrays based on morton_codes as keys - ps_impl_sort_particles(&ctx->arena, morton_codes, arrs); - - // allocate the root node - ctx->root = (ps_node_t*)ps_arena_alloc(&ctx->arena, sizeof(ps_node_t), 32); - if (!ctx->root) { - return PS_EOOM; + ps_result_t res = ps_impl_sort_particles(ctx, morton_codes, arrs); + if (res != PS_OK) { + return res; } - ps_impl_node_init(ctx->root); + // allocate the root node + ctx->root = ps_impl_alloc_node(ctx); // seed geometry ctx->root->x = root_cx; @@ -1075,28 +1851,48 @@ ps_result_t ps_calc_forces(ps_context_t* ctx, const ps_particle_arrs_t* arrs, ctx->root->z = root_cz; ctx->root->half_width = root_hw; - // build the octree - for (size_t i = 0; i < arrs->cnt; ++i) { - ps_impl_tree_insert(&ctx->arena, ctx->root, morton_codes[i], - (uint32_t)i); + // seed the radix state so the workers can read the morton codes + ctx->radix_state.m_src = morton_codes; + for (size_t i = 0; i < arrs->cnt; ++i) { arrs->fx[i] = 0.0F; arrs->fy[i] = 0.0F; arrs->fz[i] = 0.0F; } - // upward pass (p2m -> m2m) - ps_impl_fmm_upward_pass(ctx->root, arrs); + // build the octree + ps_impl_build_tree(ctx, 0, ctx->root, 0, 0, arrs->cnt, morton_codes); + ps_impl_pool_wait(ctx); - // interaction pass (m2l) - ps_impl_fmm_interaction_pass(ctx->root, ctx->root, ctx->theta); + // upward pass (p2m -> m2m) + ps_impl_fmm_upward_pass(ctx, ctx->root, arrs); + +// interaction pass (m2l) +#ifdef PS_MULTITHREADING + int dispatch_depth = 2; + ps_impl_dispatch_fmm(ctx, PS_JOB_INTERACTION, ctx->root, ctx->root, arrs, 0, + dispatch_depth); + ps_impl_pool_wait(ctx); +#else + ps_impl_fmm_interaction_pass(ctx, ctx->root, ctx->root, ctx->theta); +#endif - // downward pass (l2l) - ps_impl_fmm_downward_pass(ctx->root); + // downward pass (l2l -> l2p) +#ifdef PS_MULTITHREADING + ps_impl_dispatch_downward(ctx, ctx->root, arrs, 0, dispatch_depth); + ps_impl_pool_wait(ctx); +#else + ps_impl_fmm_downward_pass(ctx, ctx->root, arrs); +#endif - // evaluation pass (l2p->p2p) - ps_impl_fmm_l2p_pass(ctx->root, arrs); - ps_impl_fmm_p2p_pass(ctx->root, ctx->root, arrs); +// evaluation pass (p2p) +#ifdef PS_MULTITHREADING + ps_impl_dispatch_fmm(ctx, PS_JOB_P2P, ctx->root, ctx->root, arrs, 0, + dispatch_depth); + ps_impl_pool_wait(ctx); +#else + ps_impl_fmm_p2p_pass(ctx, ctx->root, ctx->root, arrs, ctx->theta); +#endif return PS_OK; } @@ -1362,6 +2158,28 @@ ps_result_t ps_prepare_particles(ps_particle_arrs_t* arrs, return PS_OK; } +ps_result_t ps_destroy(ps_context_t* ctx) { + if (!ctx) { + return PS_EINVAL; + } + +#ifdef PS_MULTITHREADING + + ps_thrd_pool_t* pool = &ctx->pool; + + ps_spin_lock(&pool->lock); + pool->shutdown_flag = 1; + ps_spin_unlock(&pool->lock); + + for (uint32_t i = 0; i < pool->thrd_cnt; ++i) { + ps_thrd_join(pool->thrds[i]); + } + +#endif // PS_MULTITHREADING + + return PS_OK; +} + #endif // POLESITTER_IMPLEMENTATION /* diff --git a/tests/accuracy.c b/tests/accuracy.c index c724eee..4ef0726 100644 --- a/tests/accuracy.c +++ b/tests/accuracy.c @@ -3,12 +3,13 @@ #include #include +#define PS_MULTITHREADING #define POLESITTER_IMPLEMENTATION #include "../src/polesitter.h" #define ARENA_SIZE (1024ULL * 1024ULL * 32ULL) #define SOFTENING_SQ 0.1F -#define THETA 3.4641F // sqrt of 12 +#define THETA 2.0F static int g_failures = 0; @@ -39,14 +40,14 @@ static float rand_unit(uint32_t* state) { } static void direct_reference(const float* x, const float* y, const float* z, - const float* mass, size_t count, double* fx, + const float* mass, size_t cnt, double* fx, double* fy, double* fz) { - for (size_t i = 0; i < count; ++i) { + for (size_t i = 0; i < cnt; ++i) { double ax = 0.0; double ay = 0.0; double az = 0.0; - for (size_t j = 0; j < count; ++j) { + for (size_t j = 0; j < cnt; ++j) { if (i == j) { continue; } @@ -71,39 +72,46 @@ static void direct_reference(const float* x, const float* y, const float* z, } static void run_dipole_sign_test(void) { - ps_node_t target; - ps_node_t source; - ps_impl_node_init(&target); - ps_impl_node_init(&source); + void* arena = malloc(ARENA_SIZE); + ps_context_t* ctx = NULL; + ps_config_t cfg = {arena, ARENA_SIZE, 2, THETA, 1}; + ps_init(&ctx, &cfg); - target.x = 0.0F; - target.y = 0.0F; - target.z = 0.0F; - target.half_width = 1.0F; + ps_node_t* target = ps_impl_alloc_node(ctx); + ps_node_t* source = ps_impl_alloc_node(ctx); - source.x = 100.0F; - source.y = 0.0F; - source.z = 0.0F; - source.half_width = 1.0F; + target->x = 0.0F; + target->y = 0.0F; + target->z = 0.0F; + target->half_width = 1.0F; + + source->x = 100.0F; + source->y = 0.0F; + source->z = 0.0F; + source->half_width = 1.0F; /* One unit mass displaced +1 from the source-cell center. */ - source.multipole[0] = 1.0F; - source.multipole[1] = 1.0F; + source->multipole[0] = 1.0F; + source->multipole[1] = 1.0F; - ps_impl_fmm_interaction_pass(&target, &source, THETA); + ps_impl_fmm_interaction_pass(ctx, target, source, THETA); /* * First-order softened expansion for R=100 and delta=+1 is * approximately 9.799856e-05. The large separation keeps this * analytic check independent of the production MAC threshold. */ - check_close("M2L dipole orientation", target.local[1], 0.00009799856, 5e-9); + float* target_local = ps_impl_get_local(ctx, target); + check_close("M2L dipole orientation", target_local[1], 0.00009799856, 5e-9); + + ps_destroy(ctx); + free(arena); } static void fill_uniform(float* x, float* y, float* z, float* mass, - uint32_t* id, size_t count) { + uint32_t* id, size_t cnt) { uint32_t state = 7331U; - for (size_t i = 0; i < count; ++i) { + for (size_t i = 0; i < cnt; ++i) { x[i] = (rand_unit(&state) * 100.0F) - 50.0F; y[i] = (rand_unit(&state) * 100.0F) - 50.0F; z[i] = (rand_unit(&state) * 100.0F) - 50.0F; @@ -113,11 +121,11 @@ static void fill_uniform(float* x, float* y, float* z, float* mass, } static void fill_two_clusters(float* x, float* y, float* z, float* mass, - uint32_t* id, size_t count) { + uint32_t* id, size_t cnt) { uint32_t state = 20260803U; - size_t half = count / 2; + size_t half = cnt / 2; - for (size_t i = 0; i < count; ++i) { + for (size_t i = 0; i < cnt; ++i) { float center = i < half ? -12.0F : 12.0F; float skew = i < half ? -0.75F : 0.75F; x[i] = center + skew + ((rand_unit(&state) - 0.5F) * 2.0F); @@ -128,21 +136,308 @@ static void fill_two_clusters(float* x, float* y, float* z, float* mass, } } -static double run_accuracy_case(const char* name, size_t count, +static int comp_arrs(const char* phase, size_t cnt, ps_particle_arrs_t* st, + ps_particle_arrs_t* mt, const uint32_t* morton_st, + const uint32_t* morton_mt) { + int mismatches = 0; + for (size_t i = 0; i < cnt; ++i) { + if (st->id[i] != mt->id[i] || morton_st[i] != morton_mt[i]) { + printf( + "[%s] Order mismatch at array index %zu: ST_ID=%u MT_ID=%u\n", + phase, i, st->id[i], mt->id[i]); + mismatches++; + } + if (fabs(st->fx[i] - mt->fx[i]) > 1e-4F || + fabs(st->fy[i] - mt->fy[i]) > 1e-4F || + fabs(st->fz[i] - mt->fz[i]) > 1e-4F) { + printf("[%s] Force mismatch at array index %zu (ID %u): " + "ST(%.2f,%.2f,%.2f) MT(%.2f,%.2f,%.2f)", + phase, i, st->id[i], st->fx[i], st->fy[i], st->fz[i], + mt->fx[i], mt->fy[i], mt->fz[i]); + + if (mismatches >= 5) { + break; + } + } + } + + if (mismatches > 0) { + return 0; + } + + printf("[PASS] %s\n", phase); + + return 1; +} + +static int comp_trees_rec(ps_context_t* ctx_st, ps_context_t* ctx_mt, + ps_node_t* a, ps_node_t* b, int depth, + const char* phase) { + if (!a && !b) { + return 1; + } + + if (!a || !b) { + printf("[%s] Topology mismatch at depth %d: one node is NULL\n", phase, + depth); + return 0; + } + + if (a->data.leaf.is_leaf & 1) { + if (a->data.leaf.particle_cnt != b->data.leaf.particle_cnt || + a->data.leaf.first_particle_idx != + b->data.leaf.first_particle_idx) { + printf("[%s] Leaf data mismatch at depth %d. ST(cnt=%u, idx=%u) " + "MT(cnt=%u, idx=%u)\n", + phase, depth, a->data.leaf.particle_cnt, + a->data.leaf.first_particle_idx, b->data.leaf.particle_cnt, + b->data.leaf.first_particle_idx); + return 0; + } + } + + for (int i = 0; i < 4; ++i) { + if (fabs(a->multipole[i] - b->multipole[i]) > 1e-4F) { + printf( + "[%s] Multipole mismatch at depth %d, idx %d: ST=%f, MT=%f\n", + phase, depth, i, a->multipole[i], b->multipole[i]); + return 0; + } + + float* a_local = ps_impl_get_local(ctx_st, a); + float* b_local = ps_impl_get_local(ctx_mt, b); + if (fabs(a_local[i] - b_local[i]) > 1e-4F) { + printf( + "[%s] Local field mismatch at depth %d, idx %d: ST=%f, MT=%f\n", + phase, depth, i, a_local[i], b_local[i]); + return 0; + } + } + + if (!(a->data.leaf.is_leaf & 1)) { + for (int i = 0; i < 8; ++i) { + ps_node_t* a_child = + ps_impl_get_node(ctx_st, a->data.children_offs[i]); + ps_node_t* b_child = + ps_impl_get_node(ctx_mt, b->data.children_offs[i]); + if (!comp_trees_rec(ctx_st, ctx_mt, a_child, b_child, depth + 1, + phase)) { + return 0; + } + } + } + + return 1; +} + +static void run_st_vs_mt_test(void) { + size_t cnt = 512; + + float* x_st = (float*)malloc(cnt * sizeof(float)); + float* y_st = (float*)malloc(cnt * sizeof(float)); + float* z_st = (float*)malloc(cnt * sizeof(float)); + float* m_st = (float*)malloc(cnt * sizeof(float)); + + float* x_mt = (float*)malloc(cnt * sizeof(float)); + float* y_mt = (float*)malloc(cnt * sizeof(float)); + float* z_mt = (float*)malloc(cnt * sizeof(float)); + float* m_mt = (float*)malloc(cnt * sizeof(float)); + + float* fx_st = (float*)calloc(cnt, sizeof(float)); + float* fy_st = (float*)calloc(cnt, sizeof(float)); + float* fz_st = (float*)calloc(cnt, sizeof(float)); + uint32_t* id_st = (uint32_t*)malloc(cnt * sizeof(uint32_t)); + uint32_t* mo_st = (uint32_t*)malloc(cnt * sizeof(uint32_t)); + + float* fx_mt = (float*)calloc(cnt, sizeof(float)); + float* fy_mt = (float*)calloc(cnt, sizeof(float)); + float* fz_mt = (float*)calloc(cnt, sizeof(float)); + uint32_t* id_mt = (uint32_t*)malloc(cnt * sizeof(uint32_t)); + uint32_t* mo_mt = (uint32_t*)malloc(cnt * sizeof(uint32_t)); + + fill_two_clusters(x_st, y_st, z_st, m_st, id_st, cnt); + for (size_t i = 0; i < cnt; ++i) { + x_mt[i] = x_st[i]; + y_mt[i] = y_st[i]; + z_mt[i] = z_st[i]; + m_mt[i] = m_st[i]; + id_mt[i] = id_st[i]; + } + + ps_particle_arrs_t arrs_st = {x_st, y_st, z_st, m_st, fx_st, + fy_st, fz_st, id_st, cnt}; + ps_particle_arrs_t arrs_mt = {x_mt, y_mt, z_mt, m_mt, fx_mt, + fy_mt, fz_mt, id_mt, cnt}; + + void* arena_st = malloc(ARENA_SIZE); + void* arena_mt = malloc(ARENA_SIZE); + + ps_context_t* ctx_st = NULL; + ps_context_t* ctx_mt = NULL; + ps_config_t cfg_st = {arena_st, ARENA_SIZE, cnt, THETA, 1}; + ps_config_t cfg_mt = {arena_mt, ARENA_SIZE, cnt, THETA, 4}; + + ps_init(&ctx_st, &cfg_st); + ps_init(&ctx_mt, &cfg_mt); + + float min_b, max_b, range; + ps_prepare_particles(&arrs_st, mo_st, &min_b, &max_b, &range); + ps_prepare_particles(&arrs_mt, mo_mt, &min_b, &max_b, &range); + float rc = min_b + (range * 0.5F); + float hw = range * 0.5F; + + ps_impl_arena_clear(&ctx_st->arena); + ps_impl_arena_clear(&ctx_mt->arena); + + ps_impl_sort_particles(ctx_st, mo_st, &arrs_st); + ps_impl_sort_particles(ctx_mt, mo_mt, &arrs_mt); + if (!comp_arrs("RADIX SORT", cnt, &arrs_st, &arrs_mt, mo_st, mo_mt)) { + return; + } + + // allocate the root node + ctx_st->root = ps_impl_alloc_node(ctx_st); + ctx_mt->root = ps_impl_alloc_node(ctx_mt); + + // seed geometry + ctx_st->root->x = rc; + ctx_st->root->y = rc; + ctx_st->root->z = rc; + ctx_st->root->half_width = hw; + + ctx_mt->root->x = rc; + ctx_mt->root->y = rc; + ctx_mt->root->z = rc; + ctx_mt->root->half_width = hw; + + // seed the radix state so the workers can read the morton codes + ctx_st->radix_state.m_src = mo_st; + ctx_mt->radix_state.m_src = mo_mt; + + for (size_t i = 0; i < cnt; ++i) { + arrs_st.fx[i] = 0.0F; + arrs_st.fy[i] = 0.0F; + arrs_st.fz[i] = 0.0F; + } + for (size_t i = 0; i < cnt; ++i) { + arrs_mt.fx[i] = 0.0F; + arrs_mt.fy[i] = 0.0F; + arrs_mt.fz[i] = 0.0F; + } + + // build the octree + ps_impl_build_tree(ctx_st, 0, ctx_st->root, 0, 0, cnt, mo_st); + ps_impl_pool_wait(ctx_st); + ps_impl_build_tree(ctx_mt, 0, ctx_mt->root, 0, 0, cnt, mo_mt); + ps_impl_pool_wait(ctx_mt); + + if (!comp_trees_rec(ctx_st, ctx_mt, ctx_st->root, ctx_mt->root, 0, + "TREE BUILD")) { + return; + } + + // upward pass (p2m -> m2m) + ps_impl_fmm_upward_pass(ctx_st, ctx_st->root, &arrs_st); + ps_impl_fmm_upward_pass(ctx_mt, ctx_mt->root, &arrs_mt); + if (!comp_trees_rec(ctx_st, ctx_mt, ctx_st->root, ctx_mt->root, 0, + "UPWARD PASS")) { + return; + } + + ps_impl_fmm_interaction_pass(ctx_st, ctx_st->root, ctx_st->root, THETA); + for (int i = 0; i < 8; ++i) { + ps_node_t* root_child = + ps_impl_get_node(ctx_mt, ctx_mt->root->data.children_offs[i]); + if (!root_child) { + continue; + } + + ps_job_t job; + job.type = PS_JOB_INTERACTION; + job.data.fmm.target = root_child; + job.data.fmm.src = ctx_mt->root; + ps_impl_pool_submit(ctx_mt, job); + } + ps_impl_pool_wait(ctx_mt); + if (!comp_trees_rec(ctx_st, ctx_mt, ctx_st->root, ctx_mt->root, 0, + "INTERACTION PASS")) { + return; + } + + // downward pass (l2l -> l2p) + ps_impl_fmm_downward_pass(ctx_st, ctx_st->root, &arrs_st); + for (int i = 0; i < 8; ++i) { + ps_node_t* root_child = + ps_impl_get_node(ctx_mt, ctx_mt->root->data.children_offs[i]); + if (!root_child) { + continue; + } + + // shift the roots bg field into the child sequentially + float* root_child_local = ps_impl_get_local(ctx_mt, root_child); + float* root_local = ps_impl_get_local(ctx_mt, ctx_mt->root); + root_child_local[1] += root_local[1]; + root_child_local[2] += root_local[2]; + root_child_local[3] += root_local[3]; + + ps_job_t job; + job.type = PS_JOB_DOWNWARD; + job.data.fmm.src = root_child; + job.data.fmm.arrs = &arrs_mt; + ps_impl_pool_submit(ctx_mt, job); + } + ps_impl_pool_wait(ctx_mt); + if (!comp_arrs("DOWNWARD PASS ARRAYS", cnt, &arrs_st, &arrs_mt, mo_st, + mo_mt)) { + return; + } + if (!comp_trees_rec(ctx_st, ctx_mt, ctx_st->root, ctx_mt->root, 0, + "DOWNWARD PASS TREE")) { + return; + } + + // evaluation pass (p2p) + ps_impl_fmm_p2p_pass(ctx_st, ctx_st->root, ctx_st->root, &arrs_st, + ctx_st->theta); + for (int i = 0; i < 8; ++i) { + ps_node_t* root_child = + ps_impl_get_node(ctx_mt, ctx_mt->root->data.children_offs[i]); + if (!root_child) { + continue; + } + + ps_job_t job; + job.type = PS_JOB_P2P; + job.data.fmm.target = root_child; + job.data.fmm.src = ctx_mt->root; + job.data.fmm.arrs = &arrs_mt; + ps_impl_pool_submit(ctx_mt, job); + } + ps_impl_pool_wait(ctx_mt); + if (!comp_arrs("P2P PASS", cnt, &arrs_st, &arrs_mt, mo_st, mo_mt)) { + return; + } + + ps_destroy(ctx_st); + ps_destroy(ctx_mt); + printf("[SUCCESS] ST/MT pipeline match.\n"); +} + +static double run_accuracy_case(const char* name, size_t cnt, void (*fill)(float*, float*, float*, float*, uint32_t*, size_t)) { - float* x = (float*)malloc(count * sizeof(float)); - float* y = (float*)malloc(count * sizeof(float)); - float* z = (float*)malloc(count * sizeof(float)); - float* mass = (float*)malloc(count * sizeof(float)); - float* fx = (float*)calloc(count, sizeof(float)); - float* fy = (float*)calloc(count, sizeof(float)); - float* fz = (float*)calloc(count, sizeof(float)); - uint32_t* id = (uint32_t*)malloc(count * sizeof(uint32_t)); - uint32_t* morton = (uint32_t*)malloc(count * sizeof(uint32_t)); - double* ref_fx = (double*)malloc(count * sizeof(double)); - double* ref_fy = (double*)malloc(count * sizeof(double)); - double* ref_fz = (double*)malloc(count * sizeof(double)); + float* x = (float*)malloc(cnt * sizeof(float)); + float* y = (float*)malloc(cnt * sizeof(float)); + float* z = (float*)malloc(cnt * sizeof(float)); + float* mass = (float*)malloc(cnt * sizeof(float)); + float* fx = (float*)calloc(cnt, sizeof(float)); + float* fy = (float*)calloc(cnt, sizeof(float)); + float* fz = (float*)calloc(cnt, sizeof(float)); + uint32_t* id = (uint32_t*)malloc(cnt * sizeof(uint32_t)); + uint32_t* morton = (uint32_t*)malloc(cnt * sizeof(uint32_t)); + double* ref_fx = (double*)malloc(cnt * sizeof(double)); + double* ref_fy = (double*)malloc(cnt * sizeof(double)); + double* ref_fz = (double*)malloc(cnt * sizeof(double)); void* arena = malloc(ARENA_SIZE); if (!x || !y || !z || !mass || !fx || !fy || !fz || !id || !morton || @@ -165,12 +460,12 @@ static double run_accuracy_case(const char* name, size_t count, return -1.0; } - fill(x, y, z, mass, id, count); - direct_reference(x, y, z, mass, count, ref_fx, ref_fy, ref_fz); + fill(x, y, z, mass, id, cnt); + direct_reference(x, y, z, mass, cnt, ref_fx, ref_fy, ref_fz); - ps_particle_arrs_t arrs = {x, y, z, mass, fx, fy, fz, id, count}; + ps_particle_arrs_t arrs = {x, y, z, mass, fx, fy, fz, id, cnt}; ps_context_t* ctx = NULL; - ps_config_t cfg = {arena, ARENA_SIZE, THETA}; + ps_config_t cfg = {arena, ARENA_SIZE, cnt, THETA, 4}; if (ps_init(&ctx, &cfg) != PS_OK) { (void)fprintf(stderr, "[FAIL] %s: ps_init failed\n", name); @@ -210,7 +505,7 @@ static double run_accuracy_case(const char* name, size_t count, double ref2 = 0.0; double max_rel = 0.0; - for (size_t i = 0; i < count; ++i) { + for (size_t i = 0; i < cnt; ++i) { uint32_t original = id[i]; double dx = (double)fx[i] - ref_fx[original]; double dy = (double)fy[i] - ref_fy[original]; @@ -234,6 +529,8 @@ static double run_accuracy_case(const char* name, size_t count, "[METRIC] %s: rms_relative_error=%.6f max_relative_error=%.6f\n", name, rms_rel, max_rel); + ps_destroy(ctx); + free(x); free(y); free(z); @@ -253,12 +550,14 @@ static double run_accuracy_case(const char* name, size_t count, int main(void) { run_dipole_sign_test(); + run_st_vs_mt_test(); double uniform = run_accuracy_case("uniform-512", 512, fill_uniform); double clusters = run_accuracy_case("two-clusters-512", 512, fill_two_clusters); - /* Intentionally generous initial gates: these are validation smoke tests, + /* Intentionally generous initial gates: these are validation smoke + tests, not production accuracy targets. */ if (uniform > 0.10) { failf("uniform RMS relative error", uniform, 0.10, 0.0); diff --git a/tests/benchmark.c b/tests/benchmark.c index 7a91112..33d3239 100644 --- a/tests/benchmark.c +++ b/tests/benchmark.c @@ -1,9 +1,10 @@ +#define _POSIX_C_SOURCE 199309L #include #include #include #include -#include +#define PS_MULTITHREADING #define POLESITTER_IMPLEMENTATION #include "../src/polesitter.h" @@ -21,127 +22,190 @@ #endif -#define ARENA_SIZE (1024ULL * 1024ULL * 128ULL) +#define ARENA_SIZE (1024ULL * 1024ULL * 512ULL) +#define THETA 1.0F + +#ifdef _WIN32 + +#define WIN32_LEAN_AND_MEAN +#include + +double get_wall_time(void) { + static LARGE_INTEGER freq = {0}; + if (freq.QuadPart == 0) { + QueryPerformanceFrequency(&freq); + } + + LARGE_INTEGER time; + QueryPerformanceCounter(&time); + + return (double)time.QuadPart / (double)freq.QuadPart; +} + +#else + +#include + +double get_wall_time(void) { + struct timespec ts; + clock_gettime(CLOCK_MONOTONIC, &ts); + return (double)ts.tv_sec + ((double)ts.tv_nsec / 1e9); +} + +#endif + +double get_baseline_time(int count) { + float* px = ALIGNED_MALLOC(count * sizeof(float), 32); + float* py = ALIGNED_MALLOC(count * sizeof(float), 32); + float* pz = ALIGNED_MALLOC(count * sizeof(float), 32); + float* mass = ALIGNED_MALLOC(count * sizeof(float), 32); + float* fx = ALIGNED_MALLOC(count * sizeof(float), 32); + float* fy = ALIGNED_MALLOC(count * sizeof(float), 32); + float* fz = ALIGNED_MALLOC(count * sizeof(float), 32); + uint32_t* id = ALIGNED_MALLOC(count * sizeof(uint32_t), 32); + uint32_t* morton_codes = ALIGNED_MALLOC(count * sizeof(uint32_t), 32); + memset(fx, 0, count * sizeof(float)); + memset(fy, 0, count * sizeof(float)); + memset(fz, 0, count * sizeof(float)); + + srand(7331); + for (int i = 0; i < count; i++) { + px[i] = ((float)rand() / RAND_MAX) * 100.0F; + py[i] = ((float)rand() / RAND_MAX) * 100.0F; + pz[i] = ((float)rand() / RAND_MAX) * 100.0F; + mass[i] = 1.0F; + id[i] = i; + } -void direct_nbody_baseline(const ps_particle_arrs_t* arrs) { - size_t cnt = arrs->cnt; + ps_particle_arrs_t arrs = {px, py, pz, mass, fx, fy, fz, id, count}; - for (size_t i = 0; i < cnt; ++i) { - arrs->fx[i] = 0.0F; - arrs->fy[i] = 0.0F; - arrs->fz[i] = 0.0F; + for (int i = 0; i < count; ++i) { + arrs.fx[i] = 0.0F; + arrs.fy[i] = 0.0F; + arrs.fz[i] = 0.0F; } - for (size_t i = 0; i < cnt; ++i) { - float t_x = arrs->x[i]; - float t_y = arrs->y[i]; - float t_z = arrs->z[i]; - float t_mass = arrs->mass[i]; + double start = get_wall_time(); + for (int i = 0; i < count; ++i) { + float t_x = arrs.x[i]; + float t_y = arrs.y[i]; + float t_z = arrs.z[i]; + float t_mass = arrs.mass[i]; float f_x = 0.0F; float f_y = 0.0F; float f_z = 0.0F; - for (size_t j = 0; j < cnt; ++j) { + for (int j = 0; j < count; ++j) { float mask = (float)(i != j); - float p_dx = arrs->x[j] - t_x; - float p_dy = arrs->y[j] - t_y; - float p_dz = arrs->z[j] - t_z; + float p_dx = arrs.x[j] - t_x; + float p_dy = arrs.y[j] - t_y; + float p_dz = arrs.z[j] - t_z; float p_dist_sq = (p_dx * p_dx) + (p_dy * p_dy) + (p_dz * p_dz) + 0.1F; float inv_dist = 1.0F / sqrtf(p_dist_sq); float inv_dist3 = inv_dist * inv_dist * inv_dist; - float force = arrs->mass[j] * inv_dist3 * mask; + float force = arrs.mass[j] * inv_dist3 * mask; f_x += p_dx * force; f_y += p_dy * force; f_z += p_dz * force; } - arrs->fx[i] += t_mass * f_x; - arrs->fy[i] += t_mass * f_y; - arrs->fz[i] += t_mass * f_z; + arrs.fx[i] += t_mass * f_x; + arrs.fy[i] += t_mass * f_y; + arrs.fz[i] += t_mass * f_z; } + double end = get_wall_time(); + + ALIGNED_FREE(px); + ALIGNED_FREE(py); + ALIGNED_FREE(pz); + ALIGNED_FREE(mass); + ALIGNED_FREE(fx); + ALIGNED_FREE(fy); + ALIGNED_FREE(fz); + ALIGNED_FREE(id); + ALIGNED_FREE(morton_codes); + + return end - start; } -int main(void) { - int test_sizes[] = {1000, 2500, 5000, 10000, 15000, 20000, - 25000, 50000, 75000, 100000, 150000, 200000}; - int num_tests = sizeof(test_sizes) / sizeof(test_sizes[0]); - - printf("N,Direct_Time,FMM_Time\n"); - - for (int t = 0; t < num_tests; ++t) { - int count = test_sizes[t]; - - float* px = ALIGNED_MALLOC(count * sizeof(float), 32); - float* py = ALIGNED_MALLOC(count * sizeof(float), 32); - float* pz = ALIGNED_MALLOC(count * sizeof(float), 32); - float* mass = ALIGNED_MALLOC(count * sizeof(float), 32); - float* fx = ALIGNED_MALLOC(count * sizeof(float), 32); - float* fy = ALIGNED_MALLOC(count * sizeof(float), 32); - float* fz = ALIGNED_MALLOC(count * sizeof(float), 32); - uint32_t* id = ALIGNED_MALLOC(count * sizeof(uint32_t), 32); - uint32_t* morton_codes = ALIGNED_MALLOC(count * sizeof(uint32_t), 32); - memset(fx, 0, count * sizeof(float)); - memset(fy, 0, count * sizeof(float)); - memset(fz, 0, count * sizeof(float)); - - srand(7331); - for (int i = 0; i < count; i++) { - px[i] = ((float)rand() / RAND_MAX) * 100.0F; - py[i] = ((float)rand() / RAND_MAX) * 100.0F; - pz[i] = ((float)rand() / RAND_MAX) * 100.0F; - mass[i] = 1.0F; - id[i] = i; - } +const int test_sizes[] = {1000, 2500, 5000, 10000, 15000, 20000, + 25000, 50000, 75000, 100000, 150000, 200000}; + +double get_test_time(int count, size_t thrd_cnt) { + void* buffer = malloc(ARENA_SIZE); + ps_context_t* ctx = NULL; + ps_config_t cfg = {buffer, ARENA_SIZE, count, THETA, thrd_cnt}; + ps_init(&ctx, &cfg); + + float* px = ALIGNED_MALLOC(count * sizeof(float), 32); + float* py = ALIGNED_MALLOC(count * sizeof(float), 32); + float* pz = ALIGNED_MALLOC(count * sizeof(float), 32); + float* mass = ALIGNED_MALLOC(count * sizeof(float), 32); + float* fx = ALIGNED_MALLOC(count * sizeof(float), 32); + float* fy = ALIGNED_MALLOC(count * sizeof(float), 32); + float* fz = ALIGNED_MALLOC(count * sizeof(float), 32); + uint32_t* id = ALIGNED_MALLOC(count * sizeof(uint32_t), 32); + uint32_t* morton_codes = ALIGNED_MALLOC(count * sizeof(uint32_t), 32); + memset(fx, 0, count * sizeof(float)); + memset(fy, 0, count * sizeof(float)); + memset(fz, 0, count * sizeof(float)); + + srand(7331); + for (int i = 0; i < count; i++) { + px[i] = ((float)rand() / RAND_MAX) * 100.0F; + py[i] = ((float)rand() / RAND_MAX) * 100.0F; + pz[i] = ((float)rand() / RAND_MAX) * 100.0F; + mass[i] = 1.0F; + id[i] = i; + } - ps_particle_arrs_t arrs = {px, py, pz, mass, fx, fy, fz, id, count}; + ps_particle_arrs_t arrs = {px, py, pz, mass, fx, fy, fz, id, count}; - clock_t start_direct = clock(); - direct_nbody_baseline(&arrs); - clock_t end_direct = clock(); - double time_direct = - (double)(end_direct - start_direct) / CLOCKS_PER_SEC; + for (int i = 0; i < count; ++i) { + fx[i] = 0; + fy[i] = 0; + fz[i] = 0; + } - for (int i = 0; i < count; ++i) { - fx[i] = 0; - fy[i] = 0; - fz[i] = 0; - } + double start = get_wall_time(); - void* buffer = malloc(ARENA_SIZE); - ps_context_t* ctx = NULL; - ps_config_t cfg = {buffer, ARENA_SIZE, 1.0F}; - ps_init(&ctx, &cfg); + float min_b, max_b, range; + ps_prepare_particles(&arrs, morton_codes, &min_b, &max_b, &range); - clock_t start_fmm = clock(); + float root_c = min_b + (range / 2.0F); + float root_hw = range / 2.0F; + ps_calc_forces(ctx, &arrs, morton_codes, root_c, root_c, root_c, root_hw); - float min_b, max_b, range; - ps_prepare_particles(&arrs, morton_codes, &min_b, &max_b, &range); + double end = get_wall_time(); - float root_c = min_b + (range / 2.0F); - float root_hw = range / 2.0F; - ps_calc_forces(ctx, &arrs, morton_codes, root_c, root_c, root_c, - root_hw); + ps_destroy(ctx); - clock_t end_fmm = clock(); - double time_fmm = (double)(end_fmm - start_fmm) / CLOCKS_PER_SEC; + ALIGNED_FREE(px); + ALIGNED_FREE(py); + ALIGNED_FREE(pz); + ALIGNED_FREE(mass); + ALIGNED_FREE(fx); + ALIGNED_FREE(fy); + ALIGNED_FREE(fz); + ALIGNED_FREE(id); + ALIGNED_FREE(morton_codes); - printf("%d,%f,%f\n", count, time_direct, time_fmm); + return end - start; +} + +int main(void) { + int num_tests = sizeof(test_sizes) / sizeof(test_sizes[0]); - ALIGNED_FREE(buffer); - ALIGNED_FREE(px); - ALIGNED_FREE(py); - ALIGNED_FREE(pz); - ALIGNED_FREE(mass); - ALIGNED_FREE(fx); - ALIGNED_FREE(fy); - ALIGNED_FREE(fz); - ALIGNED_FREE(id); - ALIGNED_FREE(morton_codes); + printf("N,Direct_Time,FMM_Time_ST,FMM_Time_MT\n"); + for (int i = 0; i < num_tests; ++i) { + int count = test_sizes[i]; + printf("%d,%.6f,%.6f,%.6f\n", count, get_baseline_time(count), + get_test_time(count, 1), get_test_time(count, 8)); } return 0; diff --git a/tests/main.c b/tests/main.c index 92433fe..c418d50 100644 --- a/tests/main.c +++ b/tests/main.c @@ -4,11 +4,26 @@ #include #include +#define PS_MULTITHREADING #define POLESITTER_IMPLEMENTATION #include "../src/polesitter.h" #define THETA 3.4641F // sqrt of 12 +#if defined(_MSC_VER) + +#include +#define ALIGNED_MALLOC(size, align) _aligned_malloc((size), (align)) +#define ALIGNED_FREE(ptr) _aligned_free((ptr)) + +#else + +#include +#define ALIGNED_MALLOC(size, align) _mm_malloc((size), (align)) +#define ALIGNED_FREE(ptr) _mm_free((ptr)) + +#endif + static void ps_assert(int cond, const char* expr_str, const char* msg, const char* file, int line) { if (!cond) { @@ -57,105 +72,136 @@ void test_morton_encoding(void) { } void test_arena_allocator(void) { - void* buffer = malloc(64); + size_t ctx_size = ps_impl_align_forward(sizeof(ps_context_t), 64); + size_t total_mem = ctx_size + 256; + + void* buffer = ALIGNED_MALLOC(total_mem, 64); TEST_ASSERT(buffer != NULL, "Test buffer allocation failed"); - ps_arena_t arena; - ps_arena_init(&arena, buffer, 64); + ps_context_t* ctx = NULL; + ps_config_t cfg = {buffer, total_mem, 4, THETA, 4}; + + ps_result_t res = ps_init(&ctx, &cfg); + TEST_ASSERT(res == PS_OK, "Failed to initialize context"); + TEST_ASSERT(ctx != NULL, "Context pointer is null"); - TEST_ASSERT(arena.cap == 64, "Arena capacity mismatch"); - TEST_ASSERT(arena.off == 0, "Arena offset should start at 0"); + TEST_ASSERT(ctx->arena.cap == total_mem - ctx_size - 128, + "Arena capacity mismatch"); + TEST_ASSERT(ctx->arena.off == 0, "Arena offset should start at 0"); // basic alloc with SIMD align - void* ptr1 = ps_arena_alloc(&arena, 10, 16); + void* ptr1 = ps_impl_arena_alloc(ctx, 10); TEST_ASSERT(ptr1 != NULL, "1st allocation failed"); - TEST_ASSERT((uintptr_t)ptr1 % 16 == 0, "1st allocation not aligned to 16B"); + TEST_ASSERT((uintptr_t)ptr1 % 64 == 0, "1st allocation not aligned to 64B"); // padding test - void* ptr2 = ps_arena_alloc(&arena, 10, 16); + void* ptr2 = ps_impl_arena_alloc(ctx, 10); TEST_ASSERT(ptr2 != NULL, "2nd allocation failed"); - TEST_ASSERT((uintptr_t)ptr2 % 16 == 0, "2nd allocation not aligned to 16B"); + TEST_ASSERT((uintptr_t)ptr2 % 64 == 0, "2nd allocation not aligned to 64B"); - // ptr2 should be 16B after ptr1 - TEST_ASSERT((uintptr_t)ptr2 == (uintptr_t)ptr1 + 16, + // ptr2 should be 64B after ptr1 + TEST_ASSERT((uintptr_t)ptr2 == (uintptr_t)ptr1 + 64, "2nd allocation not 16B after 1st allocation"); // oom test - void* ptr3 = ps_arena_alloc(&arena, 64, 16); + void* ptr3 = ps_impl_arena_alloc(ctx, 64); TEST_ASSERT(ptr3 == NULL, "3rd allocation should have failed due to OOM"); // clear test - ps_arena_clear(&arena); - TEST_ASSERT(arena.off == 0, "Arena offset should be reset to 0"); + ps_impl_arena_clear(&ctx->arena); + TEST_ASSERT(ctx->arena.off == 0, "Arena offset should be reset to 0"); // reusability test - void* ptr4 = ps_arena_alloc(&arena, 32, 16); + void* ptr4 = ps_impl_arena_alloc(ctx, 32); TEST_ASSERT(ptr4 != NULL, "4th allocation failed after clear"); TEST_ASSERT( ptr4 == ptr1, "4th allocation should reuse the same memory as 1st allocation"); - free(buffer); + ps_destroy(ctx); + ALIGNED_FREE(buffer); } void test_octree_insertion(void) { - void* buffer = malloc(1024ULL * 4); + size_t ctx_size = ps_impl_align_forward(sizeof(ps_context_t), 64); + size_t total_mem = ctx_size + (1024ULL * 16); + + void* buffer = malloc(total_mem); TEST_ASSERT(buffer != NULL, "Test buffer allocation failed"); ps_context_t* ctx = NULL; - ps_config_t cfg = {buffer, 1024ULL * 4, THETA}; + ps_config_t cfg = {buffer, total_mem, 2, THETA, 4}; ps_result_t res = ps_init(&ctx, &cfg); TEST_ASSERT(res == PS_OK, "Failed to initialize context"); TEST_ASSERT(ctx != NULL, "Context pointer is null"); // alloc the root node directly from arena - ctx->root = (ps_node_t*)ps_arena_alloc(&ctx->arena, sizeof(ps_node_t), 16); + ctx->root = ps_impl_alloc_node(ctx); TEST_ASSERT(ctx->root != NULL, "Failed to allocate root node"); - ps_impl_node_init(ctx->root); - // morton code is 0, tree should traverse down children[0] at every level - uint32_t morton_zero = ps_impl_morton_encode(0, 0, 0); - res = ps_impl_tree_insert(&ctx->arena, ctx->root, morton_zero, 42); - TEST_ASSERT(res == PS_OK, "Failed to insert origin particle"); + // seed geometry for hw + ctx->root->x = 0.0F; + ctx->root->y = 0.0F; + ctx->root->z = 0.0F; + ctx->root->half_width = 10.0F; + + // create a sorted array containing both particles + uint32_t morton_zero = ps_impl_morton_encode(0, 0, 0); + uint32_t morton_max = ps_impl_morton_encode(1023, 1023, 1023); + uint32_t morton_codes[2] = {morton_zero, morton_max}; + + // bind to radix_state so worker threads can read it + ctx->radix_state.m_src = morton_codes; + + ps_impl_build_tree(ctx, 0, ctx->root, 0, 0, 2, morton_codes); + ps_impl_pool_wait(ctx); ps_node_t* curr = ctx->root; - for (int i = 0; i < 10 /* PS_MAX_DEPTH */; ++i) { - TEST_ASSERT(curr->children[0] != NULL, "Missing child in origin path"); - curr = curr->children[0]; + while (!(curr->data.leaf.is_leaf & 1)) { + ps_node_t* curr_child = + ps_impl_get_node(ctx, curr->data.children_offs[0]); + TEST_ASSERT(curr_child != NULL, "Missing child in origin path"); + curr = curr_child; } - TEST_ASSERT(curr->is_leaf, "Bottom node is not marked as leaf"); - TEST_ASSERT(curr->particle_cnt == 1, "Leaf particle count is wrong"); - TEST_ASSERT(curr->first_particle_idx == 42, - "Leaf stored wrong particle idx"); - // maximum bounds test, tree should traverse down children[7] at every level - uint32_t morton_max = ps_impl_morton_encode(1023, 1023, 1023); - res = ps_impl_tree_insert(&ctx->arena, ctx->root, morton_max, 99); - TEST_ASSERT(res == PS_OK, "Failed to insert max bounds particle"); + TEST_ASSERT(curr->data.leaf.is_leaf, "Bottom node is not marked as leaf"); + TEST_ASSERT(curr->data.leaf.particle_cnt == 1, + "Leaf particle count is wrong"); + TEST_ASSERT(curr->data.leaf.first_particle_idx == 0, + "Leaf stored wrong particle idx"); curr = ctx->root; - for (int i = 0; i < 10 /* PS_MAX_DEPTH */; ++i) { - TEST_ASSERT(curr->children[7] != NULL, - "Missing child in max bounds path"); - curr = curr->children[7]; + while (!(curr->data.leaf.is_leaf & 1)) { + ps_node_t* curr_child = + ps_impl_get_node(ctx, curr->data.children_offs[7]); + TEST_ASSERT(curr_child != NULL, "Missing child in max bounds path"); + curr = curr_child; } - TEST_ASSERT(curr->is_leaf, "Bottom node is not marked as leaf"); - TEST_ASSERT(curr->particle_cnt == 1, "Leaf particle count is wrong"); - TEST_ASSERT(curr->first_particle_idx == 99, + TEST_ASSERT(curr->data.leaf.is_leaf, "Bottom node is not marked as leaf"); + TEST_ASSERT(curr->data.leaf.particle_cnt == 1, + "Leaf particle count is wrong"); + TEST_ASSERT(curr->data.leaf.first_particle_idx == 1, "Leaf stored wrong particle idx"); + ps_destroy(ctx); free(buffer); } void test_radix_sort(void) { - void* buffer = malloc(1024ULL * 4); + size_t ctx_size = ps_impl_align_forward(sizeof(ps_context_t), 64); + size_t total_mem = ctx_size + (1024ULL * 1024); + + void* buffer = ALIGNED_MALLOC(total_mem, 64); TEST_ASSERT(buffer != NULL, "Test buffer allocation failed"); ps_context_t* ctx = NULL; - ps_config_t cfg = {buffer, 1024ULL * 4, THETA}; - ps_init(&ctx, &cfg); + ps_config_t cfg = {buffer, total_mem, 4, THETA, 4}; + + ps_result_t res = ps_init(&ctx, &cfg); + TEST_ASSERT(res == PS_OK, "Failed to initialize context"); + TEST_ASSERT(ctx != NULL, "Context pointer is null"); uint32_t morton_codes[4] = {999, 10, 500, 42}; uint32_t ids[4] = {0, 1, 2, 3}; @@ -167,7 +213,8 @@ void test_radix_sort(void) { ps_particle_arrs_t arrs = {x, y, z, mass, fx, fy, fz, ids, 4}; - ps_impl_sort_particles(&ctx->arena, morton_codes, &arrs); + res = ps_impl_sort_particles(ctx, morton_codes, &arrs); + TEST_ASSERT(res == PS_OK, "Sort function aborted early"); // verify morton codes are strictly ascending TEST_ASSERT(morton_codes[0] == 10, "Sort failed at idx 0"); @@ -181,15 +228,19 @@ void test_radix_sort(void) { TEST_ASSERT_FLOAT_EQ(5.0F, arrs.mass[2], 1e-6F); TEST_ASSERT_FLOAT_EQ(9.0F, arrs.mass[3], 1e-6F); - free(buffer); + ps_destroy(ctx); + ALIGNED_FREE(buffer); } void test_fmm_upward_pass(void) { - void* buffer = malloc(1024ULL * 64); + size_t ctx_size = ps_impl_align_forward(sizeof(ps_context_t), 64); + size_t total_mem = ctx_size + (1024ULL * 32); + + void* buffer = malloc(total_mem); TEST_ASSERT(buffer != NULL, "Test buffer allocation failed"); ps_context_t* ctx = NULL; - ps_config_t cfg = {buffer, 1024ULL * 64, THETA}; + ps_config_t cfg = {buffer, total_mem, 2, THETA, 4}; ps_init(&ctx, &cfg); float x[2] = {2.0F, -1.0F}; @@ -201,16 +252,15 @@ void test_fmm_upward_pass(void) { ps_particle_arrs_t arrs = {x, y, z, mass, fx, fy, fz, id, 2}; - ctx->root = (ps_node_t*)ps_arena_alloc(&ctx->arena, sizeof(ps_node_t), 16); - ps_impl_node_init(ctx->root); - ctx->root->is_leaf = 1; - ctx->root->particle_cnt = 2; - ctx->root->first_particle_idx = 0; - ctx->root->x = 0.0F; - ctx->root->y = 0.0F; - ctx->root->z = 0.0F; + ctx->root = ps_impl_alloc_node(ctx); + ctx->root->data.leaf.is_leaf = 1; + ctx->root->data.leaf.particle_cnt = 2; + ctx->root->data.leaf.first_particle_idx = 0; + ctx->root->x = 0.0F; + ctx->root->y = 0.0F; + ctx->root->z = 0.0F; - ps_impl_fmm_upward_pass(ctx->root, &arrs); + ps_impl_fmm_upward_pass(ctx, ctx->root, &arrs); // M0 = sum(mass) = 2.0 + 3.0 = 5.0 // MX = sum(mass * (x - root_x)) = 2.0*2.0 + 3.0*-1.0 = 1.0 @@ -225,23 +275,23 @@ void test_fmm_upward_pass(void) { TEST_ASSERT_FLOAT_EQ(0.0F, root->multipole[2], 1e-4F); // MY TEST_ASSERT_FLOAT_EQ(-1.0F, root->multipole[3], 1e-4F); // MZ + ps_destroy(ctx); free(buffer); } void test_fmm_interaction_pass(void) { - void* buffer = malloc(1024ULL * 64); + size_t ctx_size = ps_impl_align_forward(sizeof(ps_context_t), 64); + size_t total_mem = ctx_size + (1024ULL * 32); + + void* buffer = malloc(total_mem); TEST_ASSERT(buffer != NULL, "Test buffer allocation failed"); ps_context_t* ctx = NULL; - ps_config_t cfg = {buffer, 1024ULL * 64, THETA}; + ps_config_t cfg = {buffer, total_mem, 2, THETA, 4}; ps_init(&ctx, &cfg); - ps_node_t* node_a = - (ps_node_t*)ps_arena_alloc(&ctx->arena, sizeof(ps_node_t), 16); - ps_node_t* node_b = - (ps_node_t*)ps_arena_alloc(&ctx->arena, sizeof(ps_node_t), 16); - ps_impl_node_init(node_a); - ps_impl_node_init(node_b); + ps_node_t* node_a = ps_impl_alloc_node(ctx); + ps_node_t* node_b = ps_impl_alloc_node(ctx); node_a->x = -5.0F; node_a->y = -5.0F; @@ -255,7 +305,7 @@ void test_fmm_interaction_pass(void) { node_b->multipole[0] = 1.0F; - ps_impl_fmm_interaction_pass(node_a, node_b, THETA); + ps_impl_fmm_interaction_pass(ctx, node_a, node_b, THETA); // vector from A to B: dx=10, dy=10, dz=10 // softened dist_sq = 300.1, dist ~= 17.3234 @@ -263,60 +313,88 @@ void test_fmm_interaction_pass(void) { float expected_field = 0.00192354F; // verify node a local expansion - TEST_ASSERT_FLOAT_EQ(expected_field, node_a->local[1], 1e-6F); // F_x - TEST_ASSERT_FLOAT_EQ(expected_field, node_a->local[2], 1e-6F); // F_y - TEST_ASSERT_FLOAT_EQ(expected_field, node_a->local[3], 1e-6F); // F_z + float* a_local = ps_impl_get_local(ctx, node_a); + TEST_ASSERT_FLOAT_EQ(expected_field, a_local[1], 1e-6F); // F_x + TEST_ASSERT_FLOAT_EQ(expected_field, a_local[2], 1e-6F); // F_y + TEST_ASSERT_FLOAT_EQ(expected_field, a_local[3], 1e-6F); // F_z + ps_destroy(ctx); free(buffer); } void test_fmm_downward_pass(void) { - void* buffer = malloc(1024ULL * 64); + size_t ctx_size = ps_impl_align_forward(sizeof(ps_context_t), 64); + size_t total_mem = ctx_size + (1024ULL * 16); + + void* buffer = malloc(total_mem); TEST_ASSERT(buffer != NULL, "Test buffer allocation failed"); ps_context_t* ctx = NULL; - ps_config_t cfg = {buffer, 1024ULL * 64, THETA}; - ps_init(&ctx, &cfg); + ps_config_t cfg = {buffer, total_mem, 1, THETA, 4}; + ps_result_t res = ps_init(&ctx, &cfg); + TEST_ASSERT(res == PS_OK, "Init failed"); + + uint32_t morton_max = ps_impl_morton_encode(1023, 1023, 1023); + uint32_t morton_codes[1] = {morton_max}; + + float x[1] = {0.0F}, y[1] = {0.0F}, z[1] = {0.0F}, mass[1] = {2.5F}; + float fx[1] = {0.0F}, fy[1] = {0.0F}, fz[1] = {0.0F}; + uint32_t ids[1] = {7331}; + ps_particle_arrs_t arrs = {x, y, z, mass, fx, fy, fz, ids, 1}; // manually allocate and send the root - ctx->root = (ps_node_t*)ps_arena_alloc(&ctx->arena, sizeof(ps_node_t), 16); - ps_impl_node_init(ctx->root); + ctx->root = ps_impl_alloc_node(ctx); + ctx->root->x = 0.0F; + ctx->root->y = 0.0F; + ctx->root->z = 0.0F; ctx->root->half_width = 10.0F; - // insert a single particle to create a deep branch - uint32_t morton_max = ps_impl_morton_encode(1023, 1023, 1023); - ps_impl_tree_insert(&ctx->arena, ctx->root, morton_max, 42); + ps_impl_build_tree(ctx, 0, ctx->root, 0, 0, 1, morton_codes); + ps_impl_pool_wait(ctx); // inject a bg field at the root node - ctx->root->local[1] = 5.0F; // F_x - ctx->root->local[2] = -3.5F; // F_y - ctx->root->local[3] = 42.0F; // F_z + float* root_local = ps_impl_get_local(ctx, ctx->root); + root_local[1] = 5.0F; // F_x + root_local[2] = -3.5F; // F_y + root_local[3] = 42.0F; // F_z - ps_impl_fmm_downward_pass(ctx->root); + ps_impl_fmm_downward_pass(ctx, ctx->root, &arrs); - // traverse to the bottom leaf node + // traverse down until we hit the adaptive leaf node ps_node_t* curr = ctx->root; - for (int i = 0; i < 10 /* PS_MAX_DEPTH */; ++i) { - TEST_ASSERT(curr->children[7] != NULL, "Missing child in path"); - curr = curr->children[7]; + while (!curr->data.leaf.is_leaf) { + ps_node_t* curr_child = + ps_impl_get_node(ctx, curr->data.children_offs[7]); + TEST_ASSERT(curr_child != NULL, "Missing child in path"); + curr = curr_child; } - TEST_ASSERT(curr->is_leaf, "Bottom nod eis not marked as leaf"); + TEST_ASSERT(curr->data.leaf.is_leaf, "Bottom node is not marked as leaf"); + + // verify the field cascaded down properly (L2L) + float* curr_local = ps_impl_get_local(ctx, curr); + TEST_ASSERT_FLOAT_EQ(5.0F, curr_local[1], 1e-6F); // F_x + TEST_ASSERT_FLOAT_EQ(-3.5F, curr_local[2], 1e-6F); // F_y + TEST_ASSERT_FLOAT_EQ(42.0F, curr_local[3], 1e-6F); // F_z - // verify the field cascaded down properly - TEST_ASSERT_FLOAT_EQ(5.0F, curr->local[1], 1e-6F); // F_x - TEST_ASSERT_FLOAT_EQ(-3.5F, curr->local[2], 1e-6F); // F_y - TEST_ASSERT_FLOAT_EQ(42.0F, curr->local[3], 1e-6F); // F_z + // verify the L2P pass + TEST_ASSERT_FLOAT_EQ(12.5F, arrs.fx[0], 1e-6F); // F_x + TEST_ASSERT_FLOAT_EQ(-8.75F, arrs.fy[0], 1e-6F); // F_y + TEST_ASSERT_FLOAT_EQ(105.0F, arrs.fz[0], 1e-6F); // F_z + ps_destroy(ctx); free(buffer); } void test_fmm_l2p_pass(void) { - void* buffer = malloc(1024ULL * 4); + size_t ctx_size = ps_impl_align_forward(sizeof(ps_context_t), 64); + size_t total_mem = ctx_size + (1024ULL * 8); + + void* buffer = malloc(total_mem); TEST_ASSERT(buffer != NULL, "Test buffer allocation failed"); ps_context_t* ctx = NULL; - ps_config_t cfg = {buffer, 1024ULL * 4, THETA}; + ps_config_t cfg = {buffer, total_mem, 1, THETA, 4}; ps_init(&ctx, &cfg); float x[1] = {0.0F}, y[1] = {0.0F}, z[1] = {0.0F}, mass[1] = {2.5F}; @@ -324,33 +402,37 @@ void test_fmm_l2p_pass(void) { ps_particle_arrs_t arrs = {x, y, z, mass, fx, fy, fz, 0, 1}; // manually create a leaf node - ctx->root = (ps_node_t*)ps_arena_alloc(&ctx->arena, sizeof(ps_node_t), 16); - ps_impl_node_init(ctx->root); - ctx->root->is_leaf = 1; - ctx->root->particle_cnt = 1; - ctx->root->first_particle_idx = 0; + ctx->root = ps_impl_alloc_node(ctx); + ctx->root->data.leaf.is_leaf = 1; + ctx->root->data.leaf.particle_cnt = 1; + ctx->root->data.leaf.first_particle_idx = 0; // bg accel field - ctx->root->local[1] = 2.0F; - ctx->root->local[2] = -1.0F; - ctx->root->local[3] = 4.0F; + float* root_local = ps_impl_get_local(ctx, ctx->root); + root_local[1] = 2.0F; + root_local[2] = -1.0F; + root_local[3] = 4.0F; - ps_impl_fmm_l2p_pass(ctx->root, &arrs); + ps_impl_fmm_downward_pass(ctx, ctx->root, &arrs); // check F = m * a TEST_ASSERT_FLOAT_EQ(5.0F, arrs.fx[0], 1e-6F); // 2.5 * 2.0 TEST_ASSERT_FLOAT_EQ(-2.5F, arrs.fy[0], 1e-6F); // 2.5 * -1.0 TEST_ASSERT_FLOAT_EQ(10.0F, arrs.fz[0], 1e-6F); // 2.5 * 4.0 + ps_destroy(ctx); free(buffer); } void test_fmm_p2p_pass(void) { - void* buffer = malloc(1024ULL * 4); + size_t ctx_size = ps_impl_align_forward(sizeof(ps_context_t), 64); + size_t total_mem = ctx_size + (1024ULL * 8); + + void* buffer = malloc(total_mem); TEST_ASSERT(buffer != NULL, "Test buffer allocation failed"); ps_context_t* ctx = NULL; - ps_config_t cfg = {buffer, 1024ULL * 4, THETA}; + ps_config_t cfg = {buffer, total_mem, 2, THETA, 4}; ps_init(&ctx, &cfg); // 2 particles offset by 1.0 unit on the x axis @@ -366,14 +448,13 @@ void test_fmm_p2p_pass(void) { const ps_particle_arrs_t arrs = {x, y, z, mass, fx, fy, fz, id, 2}; // manually create a leaf node containing both - ctx->root = (ps_node_t*)ps_arena_alloc(&ctx->arena, sizeof(ps_node_t), 16); - ps_impl_node_init(ctx->root); - ctx->root->is_leaf = 1; - ctx->root->particle_cnt = 2; - ctx->root->first_particle_idx = 0; - ctx->root->half_width = 10.0F; // not well-separated + ctx->root = ps_impl_alloc_node(ctx); + ctx->root->data.leaf.is_leaf = 1; + ctx->root->data.leaf.particle_cnt = 2; + ctx->root->data.leaf.first_particle_idx = 0; + ctx->root->half_width = 10.0F; // not well-separated - ps_impl_fmm_p2p_pass(ctx->root, ctx->root, &arrs); + ps_impl_fmm_p2p_pass(ctx, ctx->root, ctx->root, &arrs, THETA); // dist = 1.0 // dist_sq = 1.0^2 + 0.1 = 1.1 @@ -390,6 +471,7 @@ void test_fmm_p2p_pass(void) { TEST_ASSERT_FLOAT_EQ(0.0F, arrs.fy[1], 1e-6F); TEST_ASSERT_FLOAT_EQ(0.0F, arrs.fz[1], 1e-6F); + ps_destroy(ctx); free(buffer); }