Display results as we progress during benchmark (facebookresearch#357)

Co-authored-by: danthe3rd <danthe3rd>

xformers/benchmarks/benchmark_mem_eff_attention.py

@@ -84,7 +84,10 @@ def benchmark_forward(shape, num_threads: int, use_attn_bias: bool, dtype):
     if True:
         r = xformers.ops.memory_efficient_attention(q, q, q, attn_bias, op=op).float()
         rr = ref_attention(
-            q.float(), q.float(), q.float(), attn_bias.float() if attn_bias else None
+            q.float(),
+            q.float(),
+            q.float(),
+            attn_bias.float() if attn_bias is not None else None,
         )
         assert (r - rr).abs().max() < 2e-4, (r - rr).abs().max()
         del r, rr

xformers/benchmarks/utils.py

@@ -212,6 +212,7 @@ def benchmark_main_helper(
             # pbar.write(f"Skipped (NotImplementedError)")
             continue
 
+        name = None
         for benchmark_object, is_optimized in zip(benchmarks_generator, [True, False]):
             if benchmark_object is None:
                 continue
@@ -236,6 +237,23 @@ def benchmark_main_helper(
             memory = torch.cuda.max_memory_allocated() / 2**20
             mem_use[name][measurement.task_spec.sub_label] = memory
             pbar.write(f"{name}: memory used: {memory} MB")
+        # Display results for benchmarks we just calculated
+        if name is not None:
+
+            def matches_current(r):
+                return (
+                    r.task_spec.sub_label == results[-1].task_spec.sub_label
+                    and r.task_spec.label == results[-1].task_spec.label
+                )
+
+            pbar.write(
+                str(
+                    benchmark.Compare(
+                        list(filter(matches_current, results))
+                        + list(filter(matches_current, results_compare_to))
+                    )
+                )
+            )
 
     pprint.pprint(mem_use)
     benchmark.Compare(results + results_compare_to).print()