Decouple ATEN from Feature Gates Deps (#4965)

Summary:
Pull Request resolved: #4965

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

Reviewed By: q10

Differential Revision: D83725615

fbshipit-source-id: 2a548e77425b1c9f2a66e0cd6d5abbaa9b37a312

Author: Jason Chen

fbgemm_gpu/cmake/TbeTraining.cmake

@@ -75,6 +75,7 @@ gpu_cpp_library(
     ${fbgemm_sources_include_directories}
   CPU_SRCS
     src/config/feature_gates.cpp
+    src/config/feature_gates_torch_op.cpp
   DESTINATION
     fbgemm_gpu)
 

fbgemm_gpu/fbgemm_gpu/config/feature_list.py

@@ -11,7 +11,7 @@
 import torch
 
 try:
-    torch.ops.load_library("//deeplearning/fbgemm/fbgemm_gpu:config_cpp")
+    torch.ops.load_library("//deeplearning/fbgemm/fbgemm_gpu:config_cpp_torch_op")
 except Exception:
     import fbgemm_gpu  # noqa F401
 

fbgemm_gpu/include/fbgemm_gpu/utils/function_types.h

@@ -0,0 +1,19 @@
+/*
+ * 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.
+ */
+
+#pragma once
+
+/*
+ * We annotate the public FBGEMM functions and hide the rest. Those
+ * public symbols can be called via fbgemm_gpu::func() or pytorch
+ * operator dispatcher. We'll hide other symbols, especially CUB APIs,
+ * because different .so may include the same CUB CUDA kernels, which
+ * results in confusion and libA may end up calling libB's CUB kernel,
+ * causing failures when we static link libcudart_static.a
+ */
+#define DLL_PUBLIC __attribute__((visibility("default")))

fbgemm_gpu/include/fbgemm_gpu/utils/ops_utils.h

@@ -29,15 +29,8 @@ __builtin_ia32_serialize(void) {
 #include <ATen/core/op_registration/op_registration.h>
 #include <torch/library.h>
 
-/*
- * We annotate the public FBGEMM functions and hide the rest. Those
- * public symbols can be called via fbgemm_gpu::func() or pytorch
- * operator dispatcher. We'll hide other symbols, especially CUB APIs,
- * because different .so may include the same CUB CUDA kernels, which
- * results in confusion and libA may end up calling libB's CUB kernel,
- * causing failures when we static link libcudart_static.a
- */
-#define DLL_PUBLIC __attribute__((visibility("default")))
+// NOLINTNEXTLINE(misc-unused-using-decls)
+#include "fbgemm_gpu/utils/function_types.h"
 
 ////////////////////////////////////////////////////////////////////////////////
 /// Op Dispatch Macros

fbgemm_gpu/src/config/feature_gates.cpp

@@ -7,14 +7,15 @@
  */
 
 #include "fbgemm_gpu/config/feature_gates.h"
-#include "fbgemm_gpu/utils/ops_utils.h"
+#include "fbgemm_gpu/utils/function_types.h"
 
 #ifdef FBGEMM_FBCODE
 #include "fbgemm_gpu/config/feature_gates_fb.h"
 #endif
 
 #include <cstdlib>
 #include <map>
+#include <stdexcept>
 #include <string>
 
 namespace fbgemm_gpu::config {
@@ -80,9 +81,3 @@ DLL_PUBLIC bool is_feature_enabled(const FbFeatureGateName& feature) {
 #endif
 
 } // namespace fbgemm_gpu::config
-
-TORCH_LIBRARY_FRAGMENT(fbgemm, m) {
-  m.def(
-      "check_feature_gate_key(str key) -> bool",
-      fbgemm_gpu::config::check_feature_gate_key);
-}

fbgemm_gpu/src/config/feature_gates_torch_op.cpp

@@ -0,0 +1,16 @@
+/*
+ * 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.
+ */
+
+#include <torch/library.h>
+#include "fbgemm_gpu/config/feature_gates.h"
+
+TORCH_LIBRARY_FRAGMENT(fbgemm, m) {
+  m.def(
+      "check_feature_gate_key(str key) -> bool",
+      fbgemm_gpu::config::check_feature_gate_key);
+}