Note

Go to the end to download the full example code.

AMR benchmark for homogeneous density box#

This example tests the the performance of the AMR solver for a homogeneous density box, the resolution is automatically adapted to the available memory and number of processes.

  9 import datetime
 10 import json
 11 import math
 12 from statistics import mean, stdev
 13
 14 import shamrock
 15
 16 device_properties = shamrock.sys.get_compute_device_properties()
 17
 18 microbench_results = shamrock.sys.get_microbench_results(allow_run=True)
 19 if len(microbench_results) == 0:
 20     print("no microbench results, please run with --benchmark-mpi")
 21     raise ValueError("no microbench results")
 22
 23 memory_gb = device_properties["global_mem_size"] / (1e9)
 24
 25 sz = 1 << 1
 26
 27 N_target_blocks = 2 ** (math.floor(math.log2(memory_gb * 1e6 / (1.5 * 8))))
 28
 29
 30 print(f"N_target_blocks = {N_target_blocks}")
 31 print(f"memory_gb = {memory_gb}")
 32 print(f"device_properties = {device_properties}")
 33
 34 N_target_blocks = min(N_target_blocks, 2**22)
 35
 36 if device_properties["type"] == "CPU":
 37     N_target_blocks = min(N_target_blocks, 2**20)
 38
 39 # make_base_grid takes a per-axis block count, not a total, so convert the
 40 # memory-based total block budget above into a per-axis size (kept as a power
 41 # of two so the base grid splits evenly across the 3 axes).
 42 N_per_axis = 2 ** (int(math.log2(N_target_blocks)) // 3)
 43
 44 shamrock.backends.reset_mem_info_max()
 45
 46
 47 compute_multiplier = shamrock.sys.world_size()
 48 # compute_multiplier = 12
 49 scheduler_split_val = int(2e7)
 50 scheduler_merge_val = 1
 51
 52
 53 if shamrock.sys.world_rank() == 0:
 54     print("N_target_block", N_target_blocks)
 55     print("N_per_axis", N_per_axis)
 56     print("scheduler_split_val", scheduler_split_val)
 57     print("scheduler_merge_val", scheduler_merge_val)
 58     print("N_target", N_per_axis**3 * 8)
 59
 60
 61 ctx = shamrock.Context()
 62 ctx.pdata_layout_new()
 63
 64 model = shamrock.get_Model_Ramses(context=ctx, vector_type="f64_3", grid_repr="i64_3")
 65
 66 multx = 1
 67 multy = 1
 68 multz = 1
 69
 70
 71 cfg = model.gen_default_config()
 72 scale_fact = 1 / (sz * N_per_axis * multx)
 73 cfg.set_scale_factor(scale_fact)
 74 cfg.set_riemann_solver_hllc()
 75 cfg.set_eos_gamma(1.66667)
 76 cfg.set_slope_lim_vanleer_sym()
 77 cfg.set_face_time_interpolation(True)
 78 model.init_scheduler(scheduler_split_val, scheduler_merge_val)
 79 model.make_base_grid(
 80     (0, 0, 0),
 81     (sz, sz, sz),
 82     (N_per_axis * multx, N_per_axis * multy, N_per_axis * multz),
 83 )
 84
 85
 86 def rho_map(rmin, rmax):
 87     x, y, z = rmin
 88     if x > 0.25 and x < 0.75:
 89         return 2
 90     return 1.0
 91
 92
 93 def rhoe_map(rmin, rmax):
 94     rho = rho_map(rmin, rmax)
 95     return 1.0 * rho
 96
 97
 98 def rhovel_map(rmin, rmax):
 99     rho = rho_map(rmin, rmax)
100     return (1 * rho, 0 * rho, 0 * rho)
101
102
103 model.set_field_value_lambda_f64("rho", rho_map)
104 model.set_field_value_lambda_f64("rhoetot", rhoe_map)
105 model.set_field_value_lambda_f64_3("rhovel", rhovel_map)
106
107
108 # Now run the actual benchmark for 5 steps
109 res_rates = []
110 res_cnts = []
111 res_system_metrics = []
112 res_mpi_timers = []
113
114 """
115 Here we insert callbacks to measure solver MPI usage by fetching the timers twice at the begining and end of the step
116 """
117 before_mpi_timers, after_mpi_timers = None, None
118
119
120 def callback_before_mpi_timer():
121     global before_mpi_timers
122     # print(shamrock.sys.world_rank(), "register before_mpi_timers")
123     before_mpi_timers = shamrock.comm.get_timers()
124
125
126 def callback_after_mpi_timer():
127     global after_mpi_timers
128     # print(shamrock.sys.world_rank(), "register after_mpi_timers")
129     after_mpi_timers = shamrock.comm.get_timers()
130
131
132 model.add_timestep_callback(step_begin=callback_before_mpi_timer, step_end=callback_after_mpi_timer)
133
134 for i in range(10):
135     if shamrock.sys.world_rank() == 0:
136         print("running step ", i + 1, "/", 10, " ...")
137
138     shamrock.sys.mpi_barrier()
139
140     # To replay the same step
141     model.set_next_dt(0.0)
142     model.timestep()
143
144     if shamrock.sys.world_rank() == 0:
145         print("collecting results ...")
146
147     tmp_res_rate, tmp_res_cnt, tmp_system_metrics = (
148         model.solver_logs_last_rate(),
149         model.solver_logs_last_obj_count(),
150         model.solver_logs_last_system_metrics(),
151     )
152     res_rates.append(tmp_res_rate)
153     res_cnts.append(tmp_res_cnt)
154     res_system_metrics.append(tmp_system_metrics)
155     res_mpi_timers.append(shamrock.comm.mpi_timers_delta(before_mpi_timers, after_mpi_timers))
156
157     if shamrock.sys.world_rank() == 0:
158         print("sleeping 1 second ...")
159
160     import time
161
162     time.sleep(1)
163
164     if shamrock.sys.world_rank() == 0:
165         print("done sleeping 1 second ...")
166
167 # result is the best rate of the 5 steps
168 res_rate, res_cnt = max(res_rates), res_cnts[0]
169
170 # index of the max rate
171 max_rate_index = res_rates.index(max(res_rates))
172 max_rate_system_metrics = res_system_metrics[max_rate_index]
173 max_mpi_timers = res_mpi_timers[max_rate_index]
174 step_time = res_cnt / res_rate
175
176 if shamrock.sys.world_rank() == 0:
177     result_text = ""
178     result_text += f"--- final score for N_target_block={N_target_blocks} ---"
179     result_text += f"world size  : {shamrock.sys.world_size()}\n"
180     result_text += f"result rate : {res_rate}\n"
181     result_text += f"result cnt  : {res_cnt}\n"
182     result_text += f"cnt/rank    : {res_cnt / shamrock.sys.world_size()}\n"
183     result_text += f"result rate per rank : {res_rate / shamrock.sys.world_size()}\n"
184     result_text += f"rates infos : max={max(res_rates)}, min={min(res_rates)}, mean={mean(res_rates)}, stddev={stdev(res_rates)}\n"
185     result_text += f"res_rates = {res_rates}\n"
186     result_text += f"res_cnts = {res_cnts}\n"
187     result_text += f"step time = {step_time}\n"
188
189     dic_out = {
190         "device_properties": device_properties,
191         "microbench_results": shamrock.sys.get_microbench_results(),
192         "shamrock_version": shamrock.version_string(),
193         "shamrock_compiler_id_string": shamrock.get_compiler_id_string(),
194         "shamrock_compile_flags": shamrock.get_compile_arg(),
195         "date": datetime.datetime.now().strftime("%Y-%m-%d %H:%M:%S"),
196         "world_size": shamrock.sys.world_size(),
197         "rate": res_rate,
198         "cnt": res_cnt,
199         "step_time": step_time,
200         "mpi_timers": max_mpi_timers,
201     }
202
203     # print the system metrics
204     metrics_duration = max_rate_system_metrics["duration"]
205     result_text += "system metrics:\n"
206     for key, value in max_rate_system_metrics.items():
207         if not key == "duration":
208             result_text += f"{key}: {value} J\n"
209             dic_out[key] = value
210
211     for key, value in max_rate_system_metrics.items():
212         if not key == "duration":
213             result_text += f"avg power {key} / step time : {value / metrics_duration} W\n"
214             dic_out[f"power_{key}"] = value / metrics_duration
215
216     dic_out["system_metric_duration"] = metrics_duration
217
218     result_text += "---------submit this result--------\n"
219     result_text += f"{json.dumps(dic_out, indent=4)}\n"
220     result_text += "-----------------------------------\n"
221
222     print("current results:")
223     print(result_text)

Estimated memory usage: 0 MB

Gallery generated by Sphinx-Gallery