Allow reusing input data in TBE benchmark (#3594)

Summary:
X-link: facebookresearch/FBGEMM#674

Pull Request resolved: #3594

Add `--num-requests` in TBE's `device` benchmark to allow for input
batches reuse.  By default, `--num-requests` is set to -1.  In this
case, the benchmark will generate `iters` batches.  If it is set, the
benchmark will generate `num_requests` batches.  If this value is
smaller than `iters`, input batches will be reused (i.e., iter `i`
uses batch `i % num_requests`).

Reviewed By: gajjanag

Differential Revision: D68340968

fbshipit-source-id: fdae703ec499f3ba2656cba3a3b4967c684058f5

Author: sryap

fbgemm_gpu/bench/bench_utils.py

@@ -182,6 +182,7 @@ def benchmark_requests(
     callback_after_warmup: Optional[Callable[[], None]] = None,
     periodic_logs: bool = False,
     warmup_ms: Optional[int] = None,
+    iters: int = -1,
 ) -> float:
     times = []
     # Run at least one warmup iteration to avoid the long cudaLaunchKernel time
@@ -209,17 +210,20 @@ def benchmark_requests(
     if callback_after_warmup is not None:
         callback_after_warmup()
 
-    num_iters = len(requests)
+    num_reqs = len(requests)
+    iters = num_reqs if iters == -1 else iters
 
     if torch.cuda.is_available():
         torch.cuda.synchronize()
-        start_events = [torch.cuda.Event(enable_timing=True) for _ in range(num_iters)]
-        end_events = [torch.cuda.Event(enable_timing=True) for _ in range(num_iters)]
+        start_events = [torch.cuda.Event(enable_timing=True) for _ in range(iters)]
+        end_events = [torch.cuda.Event(enable_timing=True) for _ in range(iters)]
     else:
         start_events = []
         end_events = []
 
-    for it, req in enumerate(requests):
+    for it in range(iters):
+        req = requests[it % num_reqs]
+
         indices, offsets, weights = req.unpack_3()
         if bwd_only:
             # Run forward before profiling if does backward only
@@ -259,15 +263,15 @@ def benchmark_requests(
         ]
 
     if periodic_logs:
-        for it in range(100, num_iters + 1, 100):
+        for it in range(100, iters + 1, 100):
             times_ = times[0:it]
             avg_time = sum(times_) / len(times_) * 1.0e6
             last_100_avg = sum(times_[-100:]) / 100 * 1.0e6
             logging.info(
                 f"Iteration [{it}/{len(requests)}]: Last 100: {last_100_avg:.2f} us, Running avg: {avg_time:.2f} us"
             )
 
-    avg_time = sum(times) / len(requests)
+    avg_time = sum(times) / iters
     median_time = statistics.median(times)
     return median_time if check_median else avg_time
 

fbgemm_gpu/bench/split_table_batched_embeddings_benchmark.py

@@ -161,6 +161,12 @@ def cli() -> None:
     "--ssd-prefix", type=str, default="/tmp/ssd_benchmark", help="SSD directory prefix"
 )
 @click.option("--cache-load-factor", default=0.2)
+@click.option(
+    "--num-requests",
+    default=-1,
+    help="Number of input batches to generate. If the value is smaller than "
+    "iters, the benchmark will reuse the input batches",
+)
 def device(  # noqa C901
     alpha: float,
     bag_size: int,
@@ -191,8 +197,10 @@ def device(  # noqa C901
     ssd: bool,
     ssd_prefix: str,
     cache_load_factor: float,
+    num_requests: int,
 ) -> None:
     assert not ssd or not dense, "--ssd cannot be used together with --dense"
+    num_requests = iters if num_requests == -1 else num_requests
     np.random.seed(42)
     torch.manual_seed(42)
     B = batch_size
@@ -341,7 +349,7 @@ def device(  # noqa C901
         f"Accessed weights per batch: {B * sum(Ds) * L * param_size_multiplier / 1.0e9: .2f} GB"
     )
     requests = generate_requests(
-        iters,
+        num_requests,
         B,
         T,
         L,
@@ -375,6 +383,7 @@ def context_factory(on_trace_ready: Callable[[profile], None]):
             ),
             flush_gpu_cache_size_mb=flush_gpu_cache_size_mb,
             num_warmups=warmup_runs,
+            iters=iters,
         )
 
     logging.info(
@@ -409,6 +418,7 @@ def context_factory(on_trace_ready: Callable[[profile], None]):
             bwd_only=True,
             grad=grad_output,
             num_warmups=warmup_runs,
+            iters=iters,
         )
 
     logging.info(