Run benchmarks using NVBench.
(args: Args)
| 156 | |
| 157 | |
| 158 | def run_nvbench(args: Args) -> None: |
| 159 | """Run benchmarks using NVBench.""" |
| 160 | |
| 161 | def nvbench_kernel(state: "nvbench.State"): |
| 162 | provider = state.get_string("Provider") |
| 163 | case = as_case(args, provider) |
| 164 | kwargs = case.make_fn_kwargs() |
| 165 | sampler = get_sampler(provider, weights=kwargs["weights"]) |
| 166 | sampler.prepare() |
| 167 | |
| 168 | # Warmup (compile, autotune, etc.) |
| 169 | sampler.sample(**kwargs) |
| 170 | torch.cuda.synchronize() |
| 171 | |
| 172 | def launcher(launch: "nvbench.Launch"): |
| 173 | stream = _as_torch_stream(launch.get_stream()) |
| 174 | with torch.cuda.stream(stream): |
| 175 | sampler.sample(**kwargs) |
| 176 | |
| 177 | state.exec(launcher, batched=False) |
| 178 | |
| 179 | csv_args = [] |
| 180 | if args.tgt_dir is not None: |
| 181 | args.tgt_dir.mkdir(parents=True, exist_ok=True) |
| 182 | csv_path = args.tgt_dir / "nvbench.csv" |
| 183 | csv_args = ["--csv", str(csv_path)] |
| 184 | |
| 185 | b = nvbench.register(nvbench_kernel) |
| 186 | b.add_string_axis("Provider", args.providers()) |
| 187 | b.add_string_axis("Case", [args.case]) |
| 188 | nvbench.run_all_benchmarks(["speed_test"] + csv_args) |
| 189 | |
| 190 | if args.tgt_dir is not None: |
| 191 | df = pd.read_csv(csv_path) |
| 192 | df = assign_col_time_ms(df).sort_values("GPU Time (sec)") |
| 193 | df.to_csv(csv_path, index=False) |
| 194 | print("Saved results to", csv_path) |
| 195 | |
| 196 | |
| 197 | def assign_col_time_ms(df: pd.DataFrame) -> pd.DataFrame: |
no test coverage detected