OpenMP + distopia together

[PardhavMaradani]

In the current MDAnalysis support for distance computations, the choice between an OpenMP or distopia accelerated backends is mutually exclusive.

Both provide great speedups on their own, but they don't have to be mutually exclusive and when combined, they could provide the best of both worlds (OpenMP for outer loops along with explicit SIMD by distopia within). Given that I have a dev setup ready from working on PR #188, I gave this a quick try to explore. I took the two most expensive methods (based on the benchmark results) - DihedralsTriclinicOutBox and AnglesTriclinicOutBox to add OpenMP parallelization.

Here are the benchmark results:

Just distopia:

$ ./bench --benchmark_filter="DihedralsTriclinicOutBox" 
2026-03-13T15:28:46+05:30
Running ./bench
Run on (8 X 4500 MHz CPU s)
CPU Caches:
  L1 Data 32 KiB (x4)
  L1 Instruction 32 KiB (x4)
  L2 Unified 256 KiB (x4)
  L3 Unified 8192 KiB (x1)
Load Average: 2.24, 1.10, 0.82
***WARNING*** CPU scaling is enabled, the benchmark real time measurements may be noisy and will incur extra overhead.
-------------------------------------------------------------------------------------------------------------------------------
Benchmark                                            Time             CPU   Iterations UserCounters...
-------------------------------------------------------------------------------------------------------------------------------
DihedralsTriclinicOutBoxFloat/10/0/0               428 ns          428 ns      1395434 Per Result=42.7919ns items_per_second=23.3689M/s
DihedralsTriclinicOutBoxFloat/100/0/0             2613 ns         2613 ns       264777 Per Result=26.1312ns items_per_second=38.2684M/s
DihedralsTriclinicOutBoxFloat/1000/0/0           25284 ns        25284 ns        27574 Per Result=25.2836ns items_per_second=39.5513M/s
DihedralsTriclinicOutBoxFloat/10000/0/0         254897 ns       254891 ns         2765 Per Result=25.4891ns items_per_second=39.2324M/s
DihedralsTriclinicOutBoxFloat/100000/0/0       2536143 ns      2536111 ns          276 Per Result=25.3611ns items_per_second=39.4305M/s
DihedralsTriclinicOutBoxFloat/1000000/0/0     25977657 ns     25977458 ns           27 Per Result=25.9775ns items_per_second=38.4949M/s
DihedralsTriclinicOutBoxFloat/10000000/0/0   263232647 ns    263219733 ns            3 Per Result=26.322ns items_per_second=37.9911M/s
DihedralsTriclinicOutBoxDouble/10/0/0              657 ns          656 ns      1006522 Per Result=65.6171ns items_per_second=15.2399M/s
DihedralsTriclinicOutBoxDouble/100/0/0            5356 ns         5355 ns       130071 Per Result=53.5544ns items_per_second=18.6726M/s
DihedralsTriclinicOutBoxDouble/1000/0/0          53403 ns        53392 ns        12983 Per Result=53.3921ns items_per_second=18.7294M/s
DihedralsTriclinicOutBoxDouble/10000/0/0        539185 ns       539110 ns         1301 Per Result=53.911ns items_per_second=18.5491M/s
DihedralsTriclinicOutBoxDouble/100000/0/0      5466642 ns      5466520 ns          128 Per Result=54.6652ns items_per_second=18.2932M/s
DihedralsTriclinicOutBoxDouble/1000000/0/0    55280199 ns     55277639 ns           13 Per Result=55.2776ns items_per_second=18.0905M/s
DihedralsTriclinicOutBoxDouble/10000000/0/0  567188054 ns    566542366 ns            1 Per Result=56.6542ns items_per_second=17.6509M/s
$ 
$ ./bench --benchmark_filter="AnglesTriclinicOutBox" 
2026-03-13T15:29:03+05:30
Running ./bench
Run on (8 X 4500 MHz CPU s)
CPU Caches:
  L1 Data 32 KiB (x4)
  L1 Instruction 32 KiB (x4)
  L2 Unified 256 KiB (x4)
  L3 Unified 8192 KiB (x1)
Load Average: 1.95, 1.09, 0.82
***WARNING*** CPU scaling is enabled, the benchmark real time measurements may be noisy and will incur extra overhead.
----------------------------------------------------------------------------------------------------------------------------
Benchmark                                         Time             CPU   Iterations UserCounters...
----------------------------------------------------------------------------------------------------------------------------
AnglesTriclinicOutBoxFloat/10/0/0               294 ns          294 ns      2376366 Per Result=29.3817ns items_per_second=34.0348M/s
AnglesTriclinicOutBoxFloat/100/0/0             1843 ns         1843 ns       379912 Per Result=18.4323ns items_per_second=54.2525M/s
AnglesTriclinicOutBoxFloat/1000/0/0           17738 ns        17737 ns        39537 Per Result=17.7371ns items_per_second=56.3791M/s
AnglesTriclinicOutBoxFloat/10000/0/0         178189 ns       178183 ns         3918 Per Result=17.8183ns items_per_second=56.1221M/s
AnglesTriclinicOutBoxFloat/100000/0/0       1777268 ns      1777223 ns          395 Per Result=17.7722ns items_per_second=56.2676M/s
AnglesTriclinicOutBoxFloat/1000000/0/0     18273611 ns     18272463 ns           38 Per Result=18.2725ns items_per_second=54.7272M/s
AnglesTriclinicOutBoxFloat/10000000/0/0   183798089 ns    183789962 ns            4 Per Result=18.379ns items_per_second=54.4099M/s
AnglesTriclinicOutBoxDouble/10/0/0              466 ns          466 ns      1371526 Per Result=46.591ns items_per_second=21.4634M/s
AnglesTriclinicOutBoxDouble/100/0/0            3785 ns         3784 ns       183485 Per Result=37.8406ns items_per_second=26.4266M/s
AnglesTriclinicOutBoxDouble/1000/0/0          38297 ns        38290 ns        18318 Per Result=38.2905ns items_per_second=26.1162M/s
AnglesTriclinicOutBoxDouble/10000/0/0        383619 ns       383567 ns         1826 Per Result=38.3567ns items_per_second=26.0711M/s
AnglesTriclinicOutBoxDouble/100000/0/0      3873396 ns      3873199 ns          181 Per Result=38.732ns items_per_second=25.8185M/s
AnglesTriclinicOutBoxDouble/1000000/0/0    39662545 ns     39661628 ns           18 Per Result=39.6616ns items_per_second=25.2133M/s
AnglesTriclinicOutBoxDouble/10000000/0/0  401398198 ns    401369460 ns            2 Per Result=40.1369ns items_per_second=24.9147M/s
$ 

distopia + OpenMP:

$ ./bench --benchmark_filter="DihedralsTriclinicOutBox" 
2026-03-13T15:28:00+05:30
Running ./bench
Run on (8 X 4500 MHz CPU s)
CPU Caches:
  L1 Data 32 KiB (x4)
  L1 Instruction 32 KiB (x4)
  L2 Unified 256 KiB (x4)
  L3 Unified 8192 KiB (x1)
Load Average: 0.20, 0.55, 0.63
***WARNING*** CPU scaling is enabled, the benchmark real time measurements may be noisy and will incur extra overhead.
-------------------------------------------------------------------------------------------------------------------------------
Benchmark                                            Time             CPU   Iterations UserCounters...
-------------------------------------------------------------------------------------------------------------------------------
DihedralsTriclinicOutBoxFloat/10/0/0              1828 ns         1828 ns       380229 Per Result=182.786ns items_per_second=5.47089M/s
DihedralsTriclinicOutBoxFloat/100/0/0             4094 ns         4077 ns       127738 Per Result=40.7715ns items_per_second=24.527M/s
DihedralsTriclinicOutBoxFloat/1000/0/0           29645 ns        29634 ns        23498 Per Result=29.634ns items_per_second=33.745M/s
DihedralsTriclinicOutBoxFloat/10000/0/0          85524 ns        85165 ns         7521 Per Result=8.51646ns items_per_second=117.42M/s
DihedralsTriclinicOutBoxFloat/100000/0/0        553676 ns       551436 ns         1196 Per Result=5.51436ns items_per_second=181.345M/s
DihedralsTriclinicOutBoxFloat/1000000/0/0      5358949 ns      5358476 ns          138 Per Result=5.35848ns items_per_second=186.62M/s
DihedralsTriclinicOutBoxFloat/10000000/0/0    53266276 ns     53263145 ns           14 Per Result=5.32631ns items_per_second=187.747M/s
DihedralsTriclinicOutBoxDouble/10/0/0             1992 ns         1992 ns       346537 Per Result=199.16ns items_per_second=5.02109M/s
DihedralsTriclinicOutBoxDouble/100/0/0            7703 ns         7696 ns        81124 Per Result=76.9567ns items_per_second=12.9943M/s
DihedralsTriclinicOutBoxDouble/1000/0/0          61942 ns        61703 ns         9016 Per Result=61.7034ns items_per_second=16.2066M/s
DihedralsTriclinicOutBoxDouble/10000/0/0        153414 ns       153410 ns         4539 Per Result=15.341ns items_per_second=65.1849M/s
DihedralsTriclinicOutBoxDouble/100000/0/0      1147630 ns      1146558 ns          628 Per Result=11.4656ns items_per_second=87.2175M/s
DihedralsTriclinicOutBoxDouble/1000000/0/0    11257817 ns     11250386 ns           65 Per Result=11.2504ns items_per_second=88.8858M/s
DihedralsTriclinicOutBoxDouble/10000000/0/0  112400040 ns    112397671 ns            6 Per Result=11.2398ns items_per_second=88.9698M/s
$ 
$ ./bench --benchmark_filter="AnglesTriclinicOutBox" 
2026-03-13T15:28:19+05:30
Running ./bench
Run on (8 X 4500 MHz CPU s)
CPU Caches:
  L1 Data 32 KiB (x4)
  L1 Instruction 32 KiB (x4)
  L2 Unified 256 KiB (x4)
  L3 Unified 8192 KiB (x1)
Load Average: 1.30, 0.78, 0.71
***WARNING*** CPU scaling is enabled, the benchmark real time measurements may be noisy and will incur extra overhead.
----------------------------------------------------------------------------------------------------------------------------
Benchmark                                         Time             CPU   Iterations UserCounters...
----------------------------------------------------------------------------------------------------------------------------
AnglesTriclinicOutBoxFloat/10/0/0              1705 ns         1705 ns       404350 Per Result=170.456ns items_per_second=5.86661M/s
AnglesTriclinicOutBoxFloat/100/0/0             3272 ns         3260 ns       214320 Per Result=32.6017ns items_per_second=30.6732M/s
AnglesTriclinicOutBoxFloat/1000/0/0           21271 ns        21266 ns        31986 Per Result=21.2657ns items_per_second=47.0242M/s
AnglesTriclinicOutBoxFloat/10000/0/0          54648 ns        54496 ns        12885 Per Result=5.44956ns items_per_second=183.501M/s
AnglesTriclinicOutBoxFloat/100000/0/0        416300 ns       416145 ns         1759 Per Result=4.16145ns items_per_second=240.301M/s
AnglesTriclinicOutBoxFloat/1000000/0/0      3895758 ns      3894971 ns          181 Per Result=3.89497ns items_per_second=256.741M/s
AnglesTriclinicOutBoxFloat/10000000/0/0    39117019 ns     39113263 ns           18 Per Result=3.91133ns items_per_second=255.668M/s
AnglesTriclinicOutBoxDouble/10/0/0             1740 ns         1739 ns       343928 Per Result=173.949ns items_per_second=5.74881M/s
AnglesTriclinicOutBoxDouble/100/0/0            5671 ns         5670 ns       122964 Per Result=56.6983ns items_per_second=17.6372M/s
AnglesTriclinicOutBoxDouble/1000/0/0          46262 ns        46123 ns        15110 Per Result=46.1235ns items_per_second=21.6809M/s
AnglesTriclinicOutBoxDouble/10000/0/0        114818 ns       114798 ns         5795 Per Result=11.4798ns items_per_second=87.1092M/s
AnglesTriclinicOutBoxDouble/100000/0/0       844467 ns       841197 ns          851 Per Result=8.41197ns items_per_second=118.878M/s
AnglesTriclinicOutBoxDouble/1000000/0/0     8254578 ns      8254282 ns           88 Per Result=8.25428ns items_per_second=121.149M/s
AnglesTriclinicOutBoxDouble/10000000/0/0   84946337 ns     83777935 ns            9 Per Result=8.37779ns items_per_second=119.363M/s
$ 

Tables above reformatted to make them less wider here

Unsurprisingly, there is a penalty for small sizes (setup cost - we can probably choose to skip parallelization for such cases), but the gains becomes considerable for large sizes (factor of 4 to 5).

Here are the example code changes for the DihedralsTriclinic<float> case, which shows the OpenMP parallel for during the dispatch:

Before

    template <> void DihedralsTriclinic(const float *a, const float *b, const float *c, const float *d, int n, const float *box, float *out) {
        HWY_EXPORT_AND_DYNAMIC_DISPATCH_T(DihedralsTriclinicHwy<float>)(a, b, c, d, n, box, out);
    }

After

    template <> void DihedralsTriclinic(const float *a, const float *b, const float *c, const float *d, int n, const float *box, float *out) {
        hwy::GetChosenTarget().Update(hwy::SupportedTargets());
        size_t chunk_size = 1024;
        #pragma omp parallel for
        for (size_t i = 0; i < n; i += chunk_size) {
            size_t chunk = std::min(chunk_size, n - i);
            HWY_EXPORT_AND_DYNAMIC_DISPATCH_T(DihedralsTriclinicHwy<float>)(a + i, b + i, c + i, d + i, chunk, box, out + i);
        }
    }

There are probably a lot more things to consider, but sharing this data here for any future exploration along these lines as this seems to be one potential way to speed things up even further. Thanks

[hmacdope]

Hey @PardhavMaradani, sorry for the no / slow reply on your PRs and also here. Short on time at the moment but will try and put a coherent response together ASAP. Thanks for getting involved, appreciate the effort, youll hear from me soon.

[PardhavMaradani]

PR #196 adds distance array benchmarks.

Here is a comparison for DistanceArrayTriclinic with a benchmark range that starts with slightly higher values:

distopia:

$ ./bench --benchmark_filter="DistanceArrayTriclinic" 
2026-04-04T18:22:49+05:30
Running ./bench
Run on (8 X 4375.89 MHz CPU s)
CPU Caches:
  L1 Data 32 KiB (x4)
  L1 Instruction 32 KiB (x4)
  L2 Unified 256 KiB (x4)
  L3 Unified 8192 KiB (x1)
Load Average: 0.82, 1.11, 1.22
--------------------------------------------------------------------------------------------------------------------------------------------------
Benchmark                                                  Time             CPU   Iterations UserCounters...
--------------------------------------------------------------------------------------------------------------------------------------------------
DistanceArrayTriclinicFloat/5000/5000/0            202511087 ns    202672306 ns            3 Per Result=8.10689ns items_per_second=123.352M/s
DistanceArrayTriclinicFloat/15625/5000/0           672905449 ns    673446346 ns            1 Per Result=8.62011ns items_per_second=116.008M/s
DistanceArrayTriclinicFloat/35000/5000/0          1522639911 ns   1523748323 ns            1 Per Result=8.70713ns items_per_second=114.848M/s
DistanceArrayTriclinicFloat/5000/15625/0           687715929 ns    688204391 ns            1 Per Result=8.80902ns items_per_second=113.52M/s
DistanceArrayTriclinicFloat/15625/15625/0         2106007128 ns   2107428104 ns            1 Per Result=8.63203ns items_per_second=115.848M/s
DistanceArrayTriclinicFloat/35000/15625/0         4776978010 ns   4832494347 ns            1 Per Result=8.83656ns items_per_second=113.166M/s
DistanceArrayTriclinicFloat/5000/35000/0          1532434902 ns   1544443738 ns            1 Per Result=8.82539ns items_per_second=113.309M/s
DistanceArrayTriclinicFloat/15625/35000/0         4783292047 ns   4808875015 ns            1 Per Result=8.79337ns items_per_second=113.722M/s
DistanceArrayTriclinicFloat/35000/35000/0         1.1071e+10 ns   1.1096e+10 ns            1 Per Result=9.05819ns items_per_second=110.397M/s
DistanceArrayTriclinicDouble/5000/5000/0           436234036 ns    436714913 ns            2 Per Result=17.4686ns items_per_second=57.2456M/s
DistanceArrayTriclinicDouble/15625/5000/0         1420126403 ns   1421480655 ns            1 Per Result=18.195ns items_per_second=54.9603M/s
DistanceArrayTriclinicDouble/35000/5000/0         3232769744 ns   3235554330 ns            1 Per Result=18.4889ns items_per_second=54.0866M/s
DistanceArrayTriclinicDouble/5000/15625/0         1464918910 ns   1465992300 ns            1 Per Result=18.7647ns items_per_second=53.2915M/s
DistanceArrayTriclinicDouble/15625/15625/0        4467543618 ns   4469928068 ns            1 Per Result=18.3088ns items_per_second=54.6185M/s
DistanceArrayTriclinicDouble/35000/15625/0        1.0239e+10 ns   1.0243e+10 ns            1 Per Result=18.7303ns items_per_second=53.3894M/s
DistanceArrayTriclinicDouble/5000/35000/0         3305918878 ns   3307255671 ns            1 Per Result=18.8986ns items_per_second=52.914M/s
DistanceArrayTriclinicDouble/15625/35000/0        1.0179e+10 ns   1.0184e+10 ns            1 Per Result=18.6222ns items_per_second=53.6992M/s
Killed

distopia+OpenMP:

$ ./bench --benchmark_filter="/DistanceArrayTriclinic" 
2026-04-04T18:15:02+05:30
Running ./bench
Run on (8 X 4500 MHz CPU s)
CPU Caches:
  L1 Data 32 KiB (x4)
  L1 Instruction 32 KiB (x4)
  L2 Unified 256 KiB (x4)
  L3 Unified 8192 KiB (x1)
Load Average: 0.66, 0.97, 1.18
----------------------------------------------------------------------------------------------------------------------------------------------
Benchmark                                              Time             CPU   Iterations UserCounters...
----------------------------------------------------------------------------------------------------------------------------------------------
DistanceArrayTriclinicFloat/5000/5000/0         63046989 ns     59004510 ns           11 Per Result=2.36018ns items_per_second=423.696M/s
DistanceArrayTriclinicFloat/15625/5000/0       139520874 ns    136323563 ns            5 Per Result=1.74494ns items_per_second=573.085M/s
DistanceArrayTriclinicFloat/35000/5000/0       325629320 ns    318447489 ns            2 Per Result=1.8197ns items_per_second=549.541M/s
DistanceArrayTriclinicFloat/5000/15625/0       140148646 ns    137168983 ns            5 Per Result=1.75576ns items_per_second=569.553M/s
DistanceArrayTriclinicFloat/15625/15625/0      445444676 ns    403984235 ns            2 Per Result=1.65472ns items_per_second=604.332M/s
DistanceArrayTriclinicFloat/35000/15625/0     1018628400 ns   1004013831 ns            1 Per Result=1.83591ns items_per_second=544.689M/s
DistanceArrayTriclinicFloat/5000/35000/0       327345441 ns    327169904 ns            2 Per Result=1.86954ns items_per_second=534.89M/s
DistanceArrayTriclinicFloat/15625/35000/0     1031270989 ns   1030706243 ns            1 Per Result=1.88472ns items_per_second=530.583M/s
DistanceArrayTriclinicFloat/35000/35000/0     2309469447 ns   2308271684 ns            1 Per Result=1.8843ns items_per_second=530.7M/s
DistanceArrayTriclinicDouble/5000/5000/0       135291948 ns    110524314 ns            5 Per Result=4.42097ns items_per_second=226.195M/s
DistanceArrayTriclinicDouble/15625/5000/0      292543395 ns    271714998 ns            3 Per Result=3.47795ns items_per_second=287.526M/s
DistanceArrayTriclinicDouble/35000/5000/0      680464155 ns    660608874 ns            1 Per Result=3.77491ns items_per_second=264.907M/s
DistanceArrayTriclinicDouble/5000/15625/0      294929853 ns    294764165 ns            3 Per Result=3.77298ns items_per_second=265.042M/s
DistanceArrayTriclinicDouble/15625/15625/0     950249978 ns    930115088 ns            1 Per Result=3.80975ns items_per_second=262.484M/s
DistanceArrayTriclinicDouble/35000/15625/0    2114889210 ns   2102703250 ns            1 Per Result=3.84494ns items_per_second=260.082M/s
DistanceArrayTriclinicDouble/5000/35000/0      690552102 ns    690050211 ns            1 Per Result=3.94314ns items_per_second=253.605M/s
DistanceArrayTriclinicDouble/15625/35000/0    2170352400 ns   2168608105 ns            1 Per Result=3.96545ns items_per_second=252.178M/s
DistanceArrayTriclinicDouble/35000/35000/0    4817893478 ns   4812911058 ns            1 Per Result=3.92891ns items_per_second=254.524M/s