Note
Go to the end to download the full example code.
reduction performance benchmarks#
This example benchmarks the reduction performance for the different algorithms available in Shamrock
9 import json
10 import random
11 import time
12
13 import matplotlib.pyplot as plt
14 import numpy as np
15 from matplotlib import colors
16 from shamrock.utils.plot import make_std_bench_plot
17
18 import shamrock
19
20 # If we use the shamrock executable to run this script instead of the python interpreter,
21 # we should not initialize the system as the shamrock executable needs to handle specific MPI logic
22 if not shamrock.sys.is_initialized():
23 shamrock.change_loglevel(1)
24 shamrock.sys.init("0:0")
Use shamrock documentation style for matplotlib
30 shamrock.matplotlib.set_shamrock_mpl_style()
Recover microbenchmark results
34 microbench_results = shamrock.sys.get_microbench_results(allow_run=True)
35 if len(microbench_results) == 0:
36 print("no microbench results, please run with --benchmark-mpi")
37 raise ValueError("no microbench results")
Main benchmark functions
42 def benchmark_f32(N, nb_repeat=10):
43 times = []
44 for i in range(nb_repeat):
45 buf = shamrock.backends.DeviceBuffer_f32()
46 buf.resize(N)
47 buf.fill(0)
48 times.append(shamrock.algs.benchmark_reduction_sum(buf, N))
49 return min(times), max(times), sum(times) / nb_repeat
50
51
52 def benchmark_f64(N, nb_repeat=10):
53 times = []
54 for i in range(nb_repeat):
55 buf = shamrock.backends.DeviceBuffer_f64()
56 buf.resize(N)
57 buf.fill(0)
58 times.append(shamrock.algs.benchmark_reduction_sum(buf, N))
59 return min(times), max(times), sum(times) / nb_repeat
Run the performance test for all parameters
64 def run_performance_sweep():
65 # Define parameter ranges
66 # logspace as array
67 particle_counts = np.logspace(2, 7, 20).astype(int).tolist()
68
69 # Initialize results matrix
70 results_f32 = []
71 results_f64 = []
72
73 print(f"Particle counts: {particle_counts}")
74
75 total_runs = len(particle_counts)
76 current_run = 0
77
78 for i, N in enumerate(particle_counts):
79 current_run += 1
80
81 print(
82 f"[{current_run:2d}/{total_runs}] Running N={N:5d}...",
83 end=" ",
84 )
85
86 start_time = time.time()
87 min_time, max_time, mean_time = benchmark_f32(N)
88 results_f32.append(min_time)
89 min_time, max_time, mean_time = benchmark_f64(N)
90 results_f64.append(min_time)
91 elapsed = time.time() - start_time
92
93 print(f"mean={mean_time:.3f}s (took {elapsed:.1f}s)")
94
95 return particle_counts, results_f32, results_f64
List current implementation
100 if not shamrock.algs.is_impl_set_reduction():
101 shamrock.algs.autoselect_impl_reduction()
102 current_impl = shamrock.algs.get_current_impl_reduction()
103
104 print(current_impl)
{"implementation":"group_reduction","parameters":{"group_size":128}}
List all implementations available
108 all_default_impls = shamrock.algs.get_default_impl_list_reduction()
109
110 print(all_default_impls)
['{"implementation":"fallback","parameters":{}}', '{"implementation":"group_reduction","parameters":{"group_size":16}}', '{"implementation":"group_reduction","parameters":{"group_size":128}}', '{"implementation":"group_reduction","parameters":{"group_size":256}}']
Run the performance benchmarks for all implementations
115 dic_bench = {}
116 for impl in all_default_impls:
117 shamrock.algs.set_impl_reduction(impl)
118
119 impl_json = json.loads(impl)
120 impl_name = impl_json["implementation"]
121 impl_params = impl_json.get("parameters", {})
122 if impl_params:
123 # Disambiguate implementations that expose multiple default parameter sets
124 # (e.g. several group sizes) under the same "implementation" name
125 params_str = ", ".join(f"{k}={v}" for k, v in impl_params.items())
126 impl_name = f"{impl_name} ({params_str})"
127
128 print(f"Running reduction performance benchmarks for {impl}...")
129
130 # Run the performance sweep
131 particle_counts, results_f32, results_f64 = run_performance_sweep()
132
133 dic_bench[impl_name] = {
134 "particle_counts": particle_counts,
135 "results_f32": results_f32,
136 "results_f64": results_f64,
137 }
Info: setting reduction implementation to impl : {"implementation":"fallback","parameters":{}} [algs][rank=0]
Running reduction performance benchmarks for {"implementation":"fallback","parameters":{}}...
Particle counts: [100, 183, 335, 615, 1128, 2069, 3792, 6951, 12742, 23357, 42813, 78475, 143844, 263665, 483293, 885866, 1623776, 2976351, 5455594, 10000000]
[ 1/20] Running N= 100... mean=0.000s (took 0.0s)
[ 2/20] Running N= 183... mean=0.000s (took 0.0s)
[ 3/20] Running N= 335... mean=0.000s (took 0.0s)
[ 4/20] Running N= 615... mean=0.000s (took 0.0s)
[ 5/20] Running N= 1128... mean=0.000s (took 0.0s)
[ 6/20] Running N= 2069... mean=0.000s (took 0.0s)
[ 7/20] Running N= 3792... mean=0.000s (took 0.0s)
[ 8/20] Running N= 6951... mean=0.000s (took 0.0s)
[ 9/20] Running N=12742... mean=0.000s (took 0.0s)
[10/20] Running N=23357... mean=0.000s (took 0.0s)
[11/20] Running N=42813... mean=0.000s (took 0.0s)
[12/20] Running N=78475... mean=0.000s (took 0.0s)
[13/20] Running N=143844... mean=0.000s (took 0.0s)
[14/20] Running N=263665... mean=0.001s (took 0.0s)
[15/20] Running N=483293... mean=0.002s (took 0.0s)
[16/20] Running N=885866... mean=0.003s (took 0.1s)
[17/20] Running N=1623776... mean=0.005s (took 0.1s)
[18/20] Running N=2976351... mean=0.005s (took 0.1s)
[19/20] Running N=5455594... mean=0.011s (took 0.2s)
[20/20] Running N=10000000... mean=0.022s (took 0.4s)
Info: setting reduction implementation to impl : {"implementation":"group_reduction","parameters":{"group_size":16}} [algs][rank=0]
Running reduction performance benchmarks for {"implementation":"group_reduction","parameters":{"group_size":16}}...
Particle counts: [100, 183, 335, 615, 1128, 2069, 3792, 6951, 12742, 23357, 42813, 78475, 143844, 263665, 483293, 885866, 1623776, 2976351, 5455594, 10000000]
[ 1/20] Running N= 100... mean=0.000s (took 0.0s)
[ 2/20] Running N= 183... mean=0.000s (took 0.0s)
[ 3/20] Running N= 335... mean=0.000s (took 0.0s)
[ 4/20] Running N= 615... mean=0.000s (took 0.0s)
[ 5/20] Running N= 1128... mean=0.000s (took 0.0s)
[ 6/20] Running N= 2069... mean=0.000s (took 0.0s)
[ 7/20] Running N= 3792... mean=0.000s (took 0.0s)
[ 8/20] Running N= 6951... mean=0.000s (took 0.0s)
[ 9/20] Running N=12742... mean=0.000s (took 0.0s)
[10/20] Running N=23357... mean=0.000s (took 0.0s)
[11/20] Running N=42813... mean=0.000s (took 0.0s)
[12/20] Running N=78475... mean=0.000s (took 0.0s)
[13/20] Running N=143844... mean=0.000s (took 0.0s)
[14/20] Running N=263665... mean=0.001s (took 0.0s)
[15/20] Running N=483293... mean=0.002s (took 0.0s)
[16/20] Running N=885866... mean=0.001s (took 0.0s)
[17/20] Running N=1623776... mean=0.003s (took 0.1s)
[18/20] Running N=2976351... mean=0.005s (took 0.1s)
[19/20] Running N=5455594... mean=0.009s (took 0.2s)
[20/20] Running N=10000000... mean=0.023s (took 0.4s)
Info: setting reduction implementation to impl : {"implementation":"group_reduction","parameters":{"group_size":128}} [algs][rank=0]
Running reduction performance benchmarks for {"implementation":"group_reduction","parameters":{"group_size":128}}...
Particle counts: [100, 183, 335, 615, 1128, 2069, 3792, 6951, 12742, 23357, 42813, 78475, 143844, 263665, 483293, 885866, 1623776, 2976351, 5455594, 10000000]
[ 1/20] Running N= 100... mean=0.000s (took 0.0s)
[ 2/20] Running N= 183... mean=0.000s (took 0.0s)
[ 3/20] Running N= 335... mean=0.000s (took 0.0s)
[ 4/20] Running N= 615... mean=0.000s (took 0.0s)
[ 5/20] Running N= 1128... mean=0.000s (took 0.0s)
[ 6/20] Running N= 2069... mean=0.000s (took 0.0s)
[ 7/20] Running N= 3792... mean=0.000s (took 0.0s)
[ 8/20] Running N= 6951... mean=0.000s (took 0.0s)
[ 9/20] Running N=12742... mean=0.000s (took 0.0s)
[10/20] Running N=23357... mean=0.000s (took 0.0s)
[11/20] Running N=42813... mean=0.000s (took 0.0s)
[12/20] Running N=78475... mean=0.000s (took 0.0s)
[13/20] Running N=143844... mean=0.000s (took 0.0s)
[14/20] Running N=263665... mean=0.001s (took 0.0s)
[15/20] Running N=483293... mean=0.001s (took 0.0s)
[16/20] Running N=885866... mean=0.001s (took 0.0s)
[17/20] Running N=1623776... mean=0.003s (took 0.1s)
[18/20] Running N=2976351... mean=0.006s (took 0.1s)
[19/20] Running N=5455594... mean=0.009s (took 0.2s)
[20/20] Running N=10000000... mean=0.022s (took 0.4s)
Info: setting reduction implementation to impl : {"implementation":"group_reduction","parameters":{"group_size":256}} [algs][rank=0]
Running reduction performance benchmarks for {"implementation":"group_reduction","parameters":{"group_size":256}}...
Particle counts: [100, 183, 335, 615, 1128, 2069, 3792, 6951, 12742, 23357, 42813, 78475, 143844, 263665, 483293, 885866, 1623776, 2976351, 5455594, 10000000]
[ 1/20] Running N= 100... mean=0.000s (took 0.0s)
[ 2/20] Running N= 183... mean=0.000s (took 0.0s)
[ 3/20] Running N= 335... mean=0.000s (took 0.0s)
[ 4/20] Running N= 615... mean=0.000s (took 0.0s)
[ 5/20] Running N= 1128... mean=0.000s (took 0.0s)
[ 6/20] Running N= 2069... mean=0.000s (took 0.0s)
[ 7/20] Running N= 3792... mean=0.000s (took 0.0s)
[ 8/20] Running N= 6951... mean=0.000s (took 0.0s)
[ 9/20] Running N=12742... mean=0.000s (took 0.0s)
[10/20] Running N=23357... mean=0.000s (took 0.0s)
[11/20] Running N=42813... mean=0.000s (took 0.0s)
[12/20] Running N=78475... mean=0.000s (took 0.0s)
[13/20] Running N=143844... mean=0.000s (took 0.0s)
[14/20] Running N=263665... mean=0.001s (took 0.0s)
[15/20] Running N=483293... mean=0.001s (took 0.0s)
[16/20] Running N=885866... mean=0.001s (took 0.0s)
[17/20] Running N=1623776... mean=0.003s (took 0.1s)
[18/20] Running N=2976351... mean=0.005s (took 0.1s)
[19/20] Running N=5455594... mean=0.009s (took 0.2s)
[20/20] Running N=10000000... mean=0.022s (took 0.4s)
Plot results (time)
143 color_cycle = plt.rcParams["axes.prop_cycle"].by_key()["color"]
144
145 plot_data = {}
146 for i, (label, item) in enumerate(dic_bench.items()):
147 color = color_cycle[i % len(color_cycle)]
148 plot_data[label + " (f64)"] = {
149 "x": item["particle_counts"],
150 "y": item["results_f64"],
151 "color": color,
152 "label": label + " (f64)",
153 "linestyle": "--",
154 "marker": ".",
155 }
156 plot_data[label + " (f32)"] = {
157 "x": item["particle_counts"],
158 "y": item["results_f32"],
159 "color": color,
160 "label": label + " (f32)",
161 "linestyle": ":",
162 "marker": None,
163 }
164
165
166 def before_plot(ax_plot):
167 particle_counts = next(iter(dic_bench.values()))["particle_counts"]
168 Nobj = np.array(particle_counts)
169 Time100M = Nobj / 1e8
170 ax_plot.plot(
171 particle_counts,
172 Time100M,
173 color="grey",
174 linestyle="-",
175 alpha=0.7,
176 label="100M obj/sec",
177 )
178
179
180 make_std_bench_plot(
181 plot_data,
182 xlabel="Number of elements",
183 ylabel="Time (s)",
184 title="reduction performance benchmarks",
185 end_label_fmt=lambda y: f"{y:.2e} s",
186 before_plot_func=before_plot,
187 )
188 plt.show()

Plot results (bandwidth)
194 peak_bw_f32 = microbench_results["saxpy_f32"]
195 peak_bw_f64 = microbench_results["saxpy_f64"]
196
197 color_cycle = plt.rcParams["axes.prop_cycle"].by_key()["color"]
198
199 plot_data = {}
200 for i, (label, item) in enumerate(dic_bench.items()):
201 color = color_cycle[i % len(color_cycle)]
202 Nobj = np.array(item["particle_counts"])
203
204 Bytes_f64 = 8 * Nobj # 1 read f64 (sizeof = 8)
205 BW_f64 = Bytes_f64 / np.array(item["results_f64"])
206 plot_data[label + " (f64)"] = {
207 "x": item["particle_counts"],
208 "y": BW_f64,
209 "color": color,
210 "label": label + " (f64)",
211 "linestyle": "-",
212 "marker": "x",
213 }
214
215 Bytes_f32 = 4 * Nobj # 1 read f32 (sizeof = 4)
216 BW_f32 = Bytes_f32 / np.array(item["results_f32"])
217 plot_data[label + " (f32)"] = {
218 "x": item["particle_counts"],
219 "y": BW_f32,
220 "color": color,
221 "label": label + " (f32)",
222 "linestyle": ":",
223 "marker": "x",
224 }
225
226
227 def before_plot(ax_plot):
228 ax_plot.axhline(
229 y=peak_bw_f64,
230 color="black",
231 linestyle=":",
232 label="microbenchmark peak BW f64",
233 )
234 ax_plot.axhline(
235 y=peak_bw_f32,
236 color="black",
237 linestyle="--",
238 label="microbenchmark peak BW f32",
239 )
240
241
242 make_std_bench_plot(
243 plot_data,
244 xlabel="Number of elements",
245 ylabel="Bandwidth (B.s^-1)",
246 title="reduction performance benchmarks",
247 end_label_fmt=lambda y: f"{y / 1e9:.2f} GB.s^-1",
248 before_plot_func=before_plot,
249 )
250 plt.show()

Total running time of the script: (0 minutes 6.149 seconds)
Estimated memory usage: 322 MB