Migrate TBE benchmark utilities over to TBE, pt 3 (pytorch#3785)

Summary:
X-link: facebookresearch/FBGEMM#870

Pull Request resolved: pytorch#3785

- Migrate TBE benchmark utilities over to TBE, pt 3

Reviewed By: jianyuh

Differential Revision: D70745336

Author: q10

fbgemm_gpu/bench/split_table_batched_embeddings_benchmark.py

@@ -49,6 +49,7 @@
 )
 from fbgemm_gpu.tbe.bench import (
     bench_warmup,
+    benchmark_cpu_requests,
     benchmark_eval_compression,
     benchmark_pipelined_requests,
     benchmark_requests,
@@ -978,24 +979,6 @@ def cache(  # noqa C901
     )
 
 
-def benchmark_cpu_requests(
-    requests: List[TBERequest],
-    func: Callable[[Tensor, Tensor, Optional[Tensor]], Tensor],
-    num_warmups: int = 0,
-) -> float:
-    import time
-
-    if num_warmups > 0:
-        for _ in range(num_warmups):
-            func(*(requests[0].unpack_3()))
-
-    start_time = time.perf_counter()
-    for req in requests:
-        func(*(req.unpack_3()))
-    end_time = time.perf_counter()
-    return (end_time - start_time) / len(requests)
-
-
 @cli.command()
 @click.option("--alpha", default=1.0)
 @click.option("--bag-size", default=20)

fbgemm_gpu/bench/split_embeddings_cache_benchmark.py renamed to fbgemm_gpu/bench/tbe/tbe_cache_benchmark.py

@@ -13,6 +13,7 @@
 )
 from .bench_runs import (  # noqa F401
     bench_warmup,
+    benchmark_cpu_requests,
     benchmark_pipelined_requests,
     benchmark_requests,
     benchmark_requests_refer,

fbgemm_gpu/bench/ssd_table_batched_embeddings_benchmark.py renamed to fbgemm_gpu/bench/tbe/tbe_ssd_benchmark.py

@@ -12,7 +12,9 @@
 from typing import Callable, List, Optional, Tuple
 
 import torch
+
 from fbgemm_gpu.tbe.utils import b_indices, TBERequest  # noqa: F401
+from torch import Tensor
 
 logging.basicConfig(level=logging.DEBUG)
 
@@ -39,6 +41,24 @@ def bench_warmup(
                 out.backward(grad)
 
 
+def benchmark_cpu_requests(
+    requests: List[TBERequest],
+    func: Callable[[Tensor, Tensor, Optional[Tensor]], Tensor],
+    num_warmups: int = 0,
+) -> float:
+    import time
+
+    if num_warmups > 0:
+        for _ in range(num_warmups):
+            func(*(requests[0].unpack_3()))
+
+    start_time = time.perf_counter()
+    for req in requests:
+        func(*(req.unpack_3()))
+    end_time = time.perf_counter()
+    return (end_time - start_time) / len(requests)
+
+
 def benchmark_requests(  # noqa: C901
     requests: List[TBERequest],
     func: Callable[[torch.Tensor, torch.Tensor, Optional[torch.Tensor]], torch.Tensor],