From 6fc4d254aac75e3f4ced126f305bdd555f6d0196 Mon Sep 17 00:00:00 2001
From: Joe Wang <wangj@meta.com>
Date: Thu, 29 Aug 2024 14:26:55 -0700
Subject: [PATCH] add ods logging for l2 cache perf (#2335)

Summary:
Pull Request resolved: https://github.com/pytorch/torchrec/pull/2335

X-link: https://github.com/pytorch/FBGEMM/pull/3031

X-link: https://github.com/facebookresearch/FBGEMM/pull/129

collect performance related metrics from KV store and export them to ODS

Reviewed By: q10

Differential Revision: D61417980
---
 torchrec/distributed/batched_embedding_kernel.py | 4 ++++
 1 file changed, 4 insertions(+)

diff --git a/torchrec/distributed/batched_embedding_kernel.py b/torchrec/distributed/batched_embedding_kernel.py
index dc4033b4e..090544a48 100644
--- a/torchrec/distributed/batched_embedding_kernel.py
+++ b/torchrec/distributed/batched_embedding_kernel.py
@@ -1156,6 +1156,10 @@ def __init__(
             **ssd_tbe_params,
         ).to(device)
 
+        logger.info(
+            f"tbe_unique_id:{self._emb_module.tbe_unique_id} => table name to count dict:{self.table_name_to_count}"
+        )
+
         self._optim: KeyValueEmbeddingFusedOptimizer = KeyValueEmbeddingFusedOptimizer(
             config,
             self._emb_module,