Implement a stat library for fbgemm embedding (#4339)

Summary:
Pull Request resolved: #4339

X-link: facebookresearch/FBGEMM#1408

This diff implements a framework to collect statistics within the FBGEMM kernel named `EmbeddingStatsTracker`. We have implemented a class method called `recordPattern` to track embedding access patterns inside `EmbeddingSpMDMAutovec.cc`, which generates a log file that counts the frequency of different access patterns. This feature is controlled by the following environment variables:
- `FBGEMM_STATS_ENABLE` (required) enables or disables the feature. Set to 1 to enable, or leave unset to disable.
- `FBGEMM_STATS_FREQ` (optional) controls how many samples are collected for each log file update and has a default value of `1000000`.
- `FBGEMM_STATS_LOGPATH` (optional) specifies where the log file is written and has a default value of `"/tmp/fbgemm_embedding_stats.txt"`.

We made the following changes:

- Implemented a statistics framework for FBGEMM embedding in `EmbeddingStatsTracker.h` and `EmbeddingStatsTracker.cc`.
- Added the corresponding tracking logic inside `EmbeddingSpMDMAutovec.cc`.
- Added environment variable loading logic in `Util.h` and `Util.cc`.

Reviewed By: excelle08

Differential Revision: D76060846

fbshipit-source-id: ad3445a15b7d0b6f0f8976914ca57def04fc68fd

Author: Kaiwei Tu

defs.bzl

@@ -202,6 +202,7 @@ def get_fbgemm_inline_neon_srcs(msvc = False, buck = False):
 def get_fbgemm_autovec_srcs():
     return [
         "src/EmbeddingSpMDMAutovec.cc",
+        "src/EmbeddingStatsTracker.cc",
     ]
 
 def get_fbgemm_tests(skip_tests = ["test/FP32Test.cc"]):

fbgemm_gpu/cmake/Fbgemm.cmake

@@ -18,7 +18,7 @@ set(fbgemm_sources_normal
   "${FBGEMM}/src/Utils.cc")
 
 if(NOT DISABLE_FBGEMM_AUTOVEC)
-  list(APPEND fbgemm_sources_normal "${FBGEMM}/src/EmbeddingSpMDMAutovec.cc")
+  list(APPEND fbgemm_sources_normal "${FBGEMM}/src/EmbeddingSpMDMAutovec.cc" "${FBGEMM}/src/EmbeddingStatsTracker.cc")
 endif()
 
 set(fbgemm_sources_avx2

include/fbgemm/Utils.h

@@ -431,6 +431,7 @@ FBGEMM_API bool is_radix_sort_accelerated_with_openmp();
 FBGEMM_API bool is_autovec_disabled();
 FBGEMM_API bool is_autovec_forced();
 FBGEMM_API bool is_asmjit_disabled();
+FBGEMM_API bool is_stats_enabled();
 
 /**
  * @brief A function to check if the input parameter in the nbit CPU TBE kernel

src/EmbeddingSpMDMAutovec.cc

@@ -10,6 +10,7 @@
 
 #define FBGEMM_EXPORTS
 #include "./EmbeddingSpMDMAutovec.h" // @manual
+#include "./EmbeddingStatsTracker.h"
 #include "./RefImplementations.h" // @manual
 #include "fbgemm/FbgemmBuild.h"
 #include "fbgemm/FloatConversion.h"
@@ -69,6 +70,18 @@ static inline void fill_output(
   }
 }
 
+template <typename OutType>
+static inline EmbeddingStatsTracker::DataType get_output_type(
+    const bool is_bf16_out) {
+  if (std::is_same<OutType, float>::value) {
+    return EmbeddingStatsTracker::DataType::FP32;
+  } else if (std::is_same<OutType, uint16_t>::value && is_bf16_out) {
+    return EmbeddingStatsTracker::DataType::BF16;
+  } else {
+    return EmbeddingStatsTracker::DataType::FP16;
+  }
+}
+
 template <typename IndexType, typename OffsetType, typename OutType>
 static bool ALWAYS_INLINE EmbeddingSpMDM8Bit_autovec(
     const int64_t block_size,
@@ -171,6 +184,15 @@ static bool ALWAYS_INLINE EmbeddingSpMDM8Bit_autovec(
       }
       out += output_stride;
     } // m
+    // Track every forward pass in the no_bag case
+    EmbeddingStatsTracker::getInstance().recordPattern(
+        data_size,
+        block_size,
+        EmbeddingStatsTracker::DataType::INT8,
+        isOutput8bit ? EmbeddingStatsTracker::DataType::INT8
+                     : get_output_type<OutType>(is_bf16_out),
+        output_size,
+        1);
     return true;
   } // no_bag
 
@@ -185,6 +207,15 @@ static bool ALWAYS_INLINE EmbeddingSpMDM8Bit_autovec(
       return false;
     }
 
+    // Track every forward inference with the actual bag size (len)
+    EmbeddingStatsTracker::getInstance().recordPattern(
+        data_size,
+        block_size,
+        EmbeddingStatsTracker::DataType::INT8,
+        get_output_type<OutType>(is_bf16_out),
+        output_size,
+        len);
+
     const float* weights_addr = weights != nullptr
         ? (is_weight_positional ? weights : weights + current)
         : nullptr;
@@ -316,6 +347,7 @@ static bool ALWAYS_INLINE EmbeddingSpMDMNBit_autovec(
       WARN_ONCE("no_bag is only supported for int4 to int4");
       return false;
     }
+
     for (int64_t i = 0; i < output_size; ++i) {
       const auto idx = indices[i];
       if (idx < 0 || idx > data_size) {
@@ -325,6 +357,15 @@ static bool ALWAYS_INLINE EmbeddingSpMDMNBit_autovec(
       memcpy(out, input_row, sizeof(uint8_t) * input_stride);
       out += input_stride;
     }
+
+    // Track every forward pass with the actual bag size (len)
+    EmbeddingStatsTracker::getInstance().recordPattern(
+        data_size,
+        block_size,
+        EmbeddingStatsTracker::DataType::INT4,
+        EmbeddingStatsTracker::DataType::INT4,
+        output_size,
+        1);
     return true;
   }
 
@@ -348,6 +389,16 @@ static bool ALWAYS_INLINE EmbeddingSpMDMNBit_autovec(
     if (end > index_size) {
       return false;
     }
+
+    // Track every forward pass with the actual bag size (len)
+    EmbeddingStatsTracker::getInstance().recordPattern(
+        data_size,
+        block_size,
+        input_bit_rate == 4 ? EmbeddingStatsTracker::DataType::INT4
+                            : EmbeddingStatsTracker::DataType::INT2,
+        get_output_type<OutType>(is_bf16_out),
+        output_size,
+        len);
     memset(buf, 0, sizeof(float) * rounded_block_size);
 
     const float* weights_addr = weights != nullptr
@@ -558,6 +609,16 @@ static bool ALWAYS_INLINE EmbeddingSpMDM_autovec(
       fill_output(out, buf, block_size, is_bf16_out);
       out += output_stride;
     } // m
+
+    EmbeddingStatsTracker::getInstance().recordPattern(
+        data_size,
+        block_size,
+        is_bf16_in ? EmbeddingStatsTracker::DataType::BF16
+                   : EmbeddingStatsTracker::DataType::FP32,
+        get_output_type<OutType>(is_bf16_out),
+        output_size,
+        1);
+
     return true;
   } // no_bag
 
@@ -592,6 +653,15 @@ static bool ALWAYS_INLINE EmbeddingSpMDM_autovec(
     if (current + len > index_size) {
       return false;
     }
+    // Track every inference for actual bag size (len)
+    EmbeddingStatsTracker::getInstance().recordPattern(
+        data_size,
+        block_size,
+        is_bf16_in ? EmbeddingStatsTracker::DataType::BF16
+                   : EmbeddingStatsTracker::DataType::FP32,
+        get_output_type<OutType>(is_bf16_out),
+        output_size,
+        len);
 
     for (int i = 0; i < len; ++i) {
       int64_t idx = indices[current];
@@ -683,6 +753,13 @@ static bool ALWAYS_INLINE EmbeddingSpMDMRowWiseSparse_autovec(
       if (end > index_size) {
         return false;
       }
+      EmbeddingStatsTracker::getInstance().recordPattern(
+          uncompressed_data_size,
+          block_size,
+          EmbeddingStatsTracker::DataType::SPARSE_INT8,
+          EmbeddingStatsTracker::DataType::FP32,
+          output_size,
+          len);
       const float* weights_addr = weights != nullptr
           ? (is_weight_positional ? weights : weights + current)
           : nullptr;
@@ -749,6 +826,14 @@ static bool ALWAYS_INLINE EmbeddingSpMDMRowWiseSparse_autovec(
         return false;
       }
 
+      EmbeddingStatsTracker::getInstance().recordPattern(
+          uncompressed_data_size,
+          block_size,
+          EmbeddingStatsTracker::DataType::SPARSE_FP32,
+          EmbeddingStatsTracker::DataType::FP32,
+          output_size,
+          len);
+
       const float* weights_addr = weights != nullptr
           ? (is_weight_positional ? weights : weights + current)
           : nullptr;
@@ -926,6 +1011,14 @@ static bool ALWAYS_INLINE EmbeddingSpMDMFP8_autovec(
       return false;
     }
 
+    EmbeddingStatsTracker::getInstance().recordPattern(
+        data_size,
+        block_size,
+        EmbeddingStatsTracker::DataType::FP8,
+        get_output_type<OutType>(is_bf16_out),
+        output_size,
+        len);
+
     // Adjust these as necessary to reflect actual batch size
     const int batch_size = block_size; // Assuming the entire block is
                                        // processed at once; adjust if needed

src/EmbeddingStatsTracker.cc

@@ -0,0 +1,64 @@
+/*
+ * 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.
+ */
+
+#define FBGEMM_EXPORTS
+#include "./EmbeddingStatsTracker.h"
+#include <iostream>
+#include "fbgemm/Utils.h"
+
+namespace fbgemm {
+
+EmbeddingStatsTracker& EmbeddingStatsTracker::getInstance() {
+  static EmbeddingStatsTracker instance;
+  return instance;
+}
+
+void EmbeddingStatsTracker::recordPattern(
+    int64_t rows,
+    int64_t dims,
+    DataType input_type,
+    DataType output_type,
+    int64_t batch_size,
+    int64_t bag_size) {
+  if (!is_stats_enabled() || bag_size == 0) {
+    return;
+  }
+  std::lock_guard<std::mutex> lock(mutex_);
+
+  // Create the entry and ensure the pattern exists
+  AccessPatternEntry key(
+      rows, dims, batch_size, bag_size, input_type, output_type);
+  auto result = tables_.find(key);
+  if (result == tables_.end()) {
+    tables_[key] = 1;
+  } else {
+    result->second += 1;
+  }
+
+  sampleCount_ += 1;
+
+  if (sampleCount_ % config_.getLogFreq() == 0) {
+    // Log the table statistics - only try to open the file if it's not
+    logFile_.open(config_.getLogFilePath(), std::ios::out | std::ios::trunc);
+
+    if (!logFile_) {
+      std::cerr << "Failed to open log file: " << config_.getLogFilePath()
+                << '\n';
+      return;
+    }
+    for (const auto& pair : tables_) {
+      const auto& pattern = pair.first;
+      logFile_ << pattern.toString() << "freq=" << pair.second << ";"
+               << std::endl;
+    }
+    logFile_.flush();
+    logFile_.close();
+  }
+}
+
+} // namespace fbgemm