GPU timing and basic reporting framework (rebase of D52716004) (#2314)

Summary:
Implements the reporting framework for internal state per TBE for better visibility.

Pull Request resolved: #2314

Reviewed By: sryap

Differential Revision: D53028585

fbshipit-source-id: fa8564068aef533b03aac93abf82c95969b2f9c2

Author: levythu

fbgemm_gpu/fbgemm_gpu/runtime_monitor.py

@@ -0,0 +1,61 @@
+#!/usr/bin/env python3
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+# All rights reserved.
+#
+# This source code is licensed under the BSD-style license found in the
+# LICENSE file in the root directory of this source tree.
+
+import abc
+
+from dataclasses import dataclass
+from typing import Optional
+
+
+class TBEStatsReporter(abc.ABC):
+    """
+    Interface for TBE runtime stats reporting. Actual implementation may do
+    custome aggregation (on intended group-key) and reporting destination.
+
+    All the report_XXX functions should be light weighted and fail-safe.
+    """
+
+    @abc.abstractmethod
+    def should_report(self, iteration_step: int) -> bool:
+        """
+        Return whether we should report metrics during this step.
+        This function should be cheap, side-effect free and return immediately.
+        """
+        ...
+
+    @abc.abstractmethod
+    def report_duration(
+        self,
+        iteration_step: int,
+        event_name: str,
+        duration_ms: float,
+        embedding_id: str = "",
+        tbe_id: str = "",
+    ) -> None:
+        """
+        Report the duration of a timed event.
+        """
+        ...
+
+
+@dataclass
+class TBEStatsReporterConfig:
+    """
+    Configuration for TBEStatsReporter. It eventually instantiates the actual
+    reporter, so it can be deep-copied without incurring the actual reporter
+    getting copied.
+    """
+
+    # Collect required batches every given batches. Non-positive stands for
+    # no collection or reporting
+    interval: int = -1
+
+    def create_reporter(self) -> Optional[TBEStatsReporter]:
+        assert (
+            self.interval <= 0
+        ), "Cannot specify interval without an actual implementation of reporter"
+        return None

fbgemm_gpu/fbgemm_gpu/split_table_batched_embeddings_ops_training.py

@@ -20,6 +20,7 @@
 from torch import nn, Tensor  # usort:skip
 
 import fbgemm_gpu.split_embedding_codegen_lookup_invokers as invokers
+from fbgemm_gpu.runtime_monitor import TBEStatsReporter, TBEStatsReporterConfig
 from fbgemm_gpu.split_embedding_configs import EmbOptimType as OptimType, SparseType
 from fbgemm_gpu.split_table_batched_embeddings_ops_common import (
     BoundsCheckMode,
@@ -349,6 +350,7 @@ def __init__(  # noqa C901
         # If a separate stream is used for prefetch, the optional forward_stream arg of prefetch function
         # should be set.
         prefetch_pipeline: bool = False,
+        stats_reporter_config: Optional[TBEStatsReporterConfig] = None,
     ) -> None:
         super(SplitTableBatchedEmbeddingBagsCodegen, self).__init__()
 
@@ -442,6 +444,13 @@ def __init__(  # noqa C901
         # 0: N_calls, 1: N_requested_indices, 2: N_unique_indices, 3: N_unique_misses,
         # 4: N_conflict_unique_misses, 5: N_conflict_misses
 
+        # Reporter to collect runtime performance stats bottom-up. Reporter may
+        # do aggregation across TBEs and publish results per training batch.
+        # Example of stats include UVM cache hit rate, table I/O size, etc.
+        self.stats_reporter: Optional[TBEStatsReporter] = (
+            stats_reporter_config.create_reporter() if stats_reporter_config else None
+        )
+
         self.int8_emb_row_dim_offset: int = INT8_EMB_ROW_DIM_OFFSET
 
         self.feature_table_map: List[int] = (

fbgemm_gpu/test/tbe/cache/cache_common.py

@@ -7,10 +7,12 @@
 
 # pyre-ignore-all-errors[56]
 
-from typing import Tuple
+from dataclasses import dataclass
+from typing import List, Optional, Tuple, Union
 
 import numpy as np
 import torch
+from fbgemm_gpu.runtime_monitor import TBEStatsReporter, TBEStatsReporterConfig
 from fbgemm_gpu.split_embedding_configs import SparseType
 
 from fbgemm_gpu.split_embedding_utils import round_up
@@ -37,6 +39,33 @@
 VERBOSITY: Verbosity = Verbosity.verbose
 
 
+class TestingStatsReporter(TBEStatsReporter):
+    def __init__(self, reporting_interval: int = 1) -> None:
+        self.reported_data: List[List[Union[int, str, float]]] = []
+        self.reporting_interval = reporting_interval
+
+    def should_report(self, iteration_step: int) -> bool:
+        return (iteration_step - 1) % self.reporting_interval == 0
+
+    def report_duration(
+        self,
+        iteration_step: int,
+        event_name: str,
+        duration_ms: float,
+        embedding_id: str = "",
+        tbe_id: str = "",
+    ) -> None:
+        self.reported_data.append(
+            [iteration_step, event_name, duration_ms, embedding_id, tbe_id]
+        )
+
+
+@dataclass
+class TestingStatsReporterConfig(TBEStatsReporterConfig):
+    def create_reporter(self) -> Optional[TBEStatsReporter]:
+        return TestingStatsReporter(reporting_interval=self.interval)
+
+
 def generate_cache_tbes(
     T: int,
     D: int,
@@ -49,6 +78,7 @@ def generate_cache_tbes(
     weights_cache_precision: SparseType = SparseType.FP32,
     stochastic_rounding: bool = False,
     gather_uvm_cache_stats: bool = False,
+    reporter_config: Optional[TestingStatsReporterConfig] = None,
 ) -> Tuple[
     SplitTableBatchedEmbeddingBagsCodegen,
     SplitTableBatchedEmbeddingBagsCodegen,
@@ -105,6 +135,7 @@ def generate_cache_tbes(
         weights_precision=weights_cache_precision,
         cache_precision=weights_cache_precision,
         gather_uvm_cache_stats=gather_uvm_cache_stats,
+        stats_reporter_config=reporter_config,
     )
 
     if use_int_weight:

fbgemm_gpu/test/tbe/cache/cache_test.py

@@ -38,6 +38,7 @@
     generate_cache_tbes,
     gpu_unavailable,
     optests,
+    TestingStatsReporterConfig,
     VERBOSITY,
 )
 
@@ -126,6 +127,7 @@ def _test_cache_prefetch_pipeline(  # noqa C901
         """
 
         assert prefetch_location in ["before_fwd", "between_fwd_bwd"]
+        reporter = TestingStatsReporterConfig(interval=2)
         cc, cc_ref, min_Es, sum_Ds = generate_cache_tbes(
             T,
             D,
@@ -137,6 +139,7 @@ def _test_cache_prefetch_pipeline(  # noqa C901
             weights_cache_precision=weights_cache_precision,
             stochastic_rounding=stochastic_rounding,
             gather_uvm_cache_stats=gather_uvm_cache_stats,
+            reporter_config=reporter,
         )
         iters = 5
         requests = generate_requests(iters, B, T, L, min_Es, reuse=0.1)