Add Inference Feature to Skip Pinned Memory Creation (#4924)

Summary:
Pull Request resolved: #4924

X-link: facebookresearch/FBGEMM#1948

Reviewed By: q10

Differential Revision: D83100663

fbshipit-source-id: d30d2742dd69834c7f0b7684d55033461fc3024a

Author: Jason Chen

fbgemm_gpu/codegen/inference/embedding_forward_quantized_cpu_template.cpp

@@ -18,6 +18,7 @@
 #include "fbgemm_gpu/embedding_common.h"
 #include "fbgemm/FbgemmEmbedding.h"
 #include "fbgemm_gpu/utils/tensor_utils.h"
+#include "fbgemm_gpu/config/feature_gates.h"
 
 #if defined(__x86_64__) || defined(__i386__) || (defined(_MSC_VER) && (defined(_M_X64) || defined(_M_IX86)))
 #include <immintrin.h>
@@ -190,8 +191,9 @@ Tensor int_nbit_split_embedding{{ "_nobag" if nobag else "" }}_codegen_forward_{
     {% else %}
     TORCH_CHECK(D > 0);
     {% endif %}
+    const static bool disablePinnedMemory = fbgemm_gpu::config::is_feature_enabled_from_env(fbgemm_gpu::config::FeatureGateName::TBE_CPU_OUTPUT_DISABLE_PINNED_MEMORY);
     bool pinned_memory = false;
-    if (at::Context::hasCUDA() && at::getNumGPUs() > 0) {
+    if (!disablePinnedMemory && at::Context::hasCUDA() && at::getNumGPUs() > 0) {
       pinned_memory = true;
     }
 

fbgemm_gpu/include/fbgemm_gpu/config/feature_gates.h

@@ -62,7 +62,8 @@ namespace fbgemm_gpu::config {
   X(TBE_ROCM_INFERENCE_PACKED_BAGS) \
   X(TBE_ROCM_HIP_BACKWARD_KERNEL)   \
   X(BOUNDS_CHECK_INDICES_V2)        \
-  X(TBE_REPORT_INPUT_PARAMS)
+  X(TBE_REPORT_INPUT_PARAMS)        \
+  X(TBE_CPU_OUTPUT_DISABLE_PINNED_MEMORY)
 // X(EXAMPLE_FEATURE_FLAG)
 
 /// @ingroup fbgemm-gpu-config
@@ -91,6 +92,13 @@ bool check_feature_gate_key(const std::string& key);
 /// is enabled.
 bool is_feature_enabled(const FeatureGateName& feature);
 
+/// @ingroup fbgemm-gpu-config
+///
+/// @brief For the given `FeatureGateName`, check if the corresponding inference
+/// feature is enabled in the env vars only. Only applicable for inference
+/// features suitable for env var rollouts
+bool is_feature_enabled_from_env(const FeatureGateName& feature);
+
 #ifdef FBGEMM_FBCODE
 bool is_feature_enabled(const FbFeatureGateName& feature);
 #endif

fbgemm_gpu/src/config/feature_gates.cpp

@@ -45,39 +45,50 @@ bool ev_check_key(const std::string& key) {
   }
 }
 
-DLL_PUBLIC bool check_feature_gate_key(const std::string& key) {
+static bool check_feature_gate_key_impl(
+    const std::string& key,
+    bool check_env_vars_only) {
   // Cache feature flags to avoid repeated JK and env var checks
   static std::map<std::string, bool> feature_flags_cache;
-#ifdef FBGEMM_FBCODE
-  static const auto no_jk = ev_check_key("NO_JK");
-#endif
-
   if (const auto search = feature_flags_cache.find(key);
       search != feature_flags_cache.end()) {
     return search->second;
+  }
+#ifdef FBGEMM_FBCODE
+  const auto value =
+      check_env_vars_only ? ev_check_key(key) : jk_check_key(key);
+#else
+  const auto value = ev_check_key(key);
+#endif
 
-  } else {
-    const auto value =
+  feature_flags_cache.insert({key, value});
+  return value;
+}
+
+DLL_PUBLIC bool check_feature_gate_key(const std::string& key) {
 #ifdef FBGEMM_FBCODE
-        (no_jk) ? ev_check_key(key) : jk_check_key(key);
+  static const auto no_jk = ev_check_key("NO_JK");
 #else
-        ev_check_key(key);
+  static const auto no_jk = false;
 #endif
 
-    feature_flags_cache.insert({key, value});
-    return value;
-  }
+  return check_feature_gate_key_impl(key, no_jk);
 }
 
 DLL_PUBLIC bool is_feature_enabled(const FeatureGateName& feature) {
   return check_feature_gate_key(to_string(feature));
 }
 
+DLL_PUBLIC bool is_feature_enabled_from_env(const FeatureGateName& feature) {
+  return check_feature_gate_key_impl(
+      to_string(feature), /* check_env_vars_only */ true);
+}
+
 #ifdef FBGEMM_FBCODE
 DLL_PUBLIC bool is_feature_enabled(const FbFeatureGateName& feature) {
   return check_feature_gate_key(to_string(feature));
 }
-#endif
+#endif // FBGEMM_FBCODE
 
 } // namespace fbgemm_gpu::config
 

fbgemm_gpu/test/tbe/inference/nbit_forward_test.py

@@ -8,6 +8,7 @@
 # pyre-strict
 # pyre-ignore-all-errors[56]
 
+import os
 import random
 import unittest
 from typing import Any, Callable, Optional, Union
@@ -123,6 +124,12 @@
 
 @optests.generate_opcheck_tests(fast=True, additional_decorators=additional_decorators)
 class NBitFowardTest(NBitFowardTestCommon):
+    def _is_cpu_output_on_pinned_memory(self) -> bool:
+        return (
+            os.getenv("FBGEMM_TBE_CPU_OUTPUT_DISABLE_PINNED_MEMORY") != "1"
+            and torch.cuda.is_available()
+        )
+
     def execute_nbit_forward_fused_pooled_emb_quant_(
         self,
         T: int,
@@ -898,6 +905,9 @@ def test_nbit_forward_cpu_seq_int8(
         lengths = torch.cat(lengths_list, 0)
         offsets = torch.ops.fbgemm.asynchronous_complete_cumsum(lengths)
         quant_cc_output = quant_cc(indices.int(), offsets.int())
+        self.assertEqual(
+            quant_cc_output.is_pinned(), self._is_cpu_output_on_pinned_memory()
+        )
         tables_rows = [
             T for T, _, _ in quant_cc.split_embedding_weights_with_scale_bias(0)
         ]