include fb only ops in private fbgemm build (pytorch#2972)

Summary:
Pull Request resolved: pytorch#2972

X-link: facebookresearch/FBGEMM#69

On fb only source files, we have CMake files based on USE_FB_ONLY compiler flag:

```
# Set the source file for FB only CPP
if(USE_FB_ONLY)
  file(GLOB fb_only_ops_sources
      fb/src/*/*.cu
      fb/src/*/*.cpp)
  list(APPEND experimental_gen_ai_cpp_source_files ${fb_only_ops_sources})
endif()
```

Reviewed By: GD06

Differential Revision: D61100617

fbshipit-source-id: 8ac9ce0c81f3ad2ecc205f255df826b3c7a4ac0d

fbgemm_gpu/CMakeLists.txt

@@ -22,9 +22,9 @@ set(CMAKE_VERBOSE_MAKEFILE on)
 # FBGEMM_GPU Build Options
 ################################################################################
 
-option(FBGEMM_CPU_ONLY  "Build FBGEMM_GPU without GPU support" OFF)
-option(USE_ROCM         "Build FBGEMM_GPU for ROCm" OFF)
-option(FBGEMM_GENAI_ONLY  "Build FBGEMM_GPU with GEN AI only support" OFF)
+option(FBGEMM_CPU_ONLY   "Build FBGEMM_GPU without GPU support" OFF)
+option(USE_ROCM          "Build FBGEMM_GPU for ROCm" OFF)
+option(FBGEMM_GENAI_ONLY "Build FBGEMM_GPU with GEN AI only support" OFF)
 option(USE_FB_ONLY       "Build FBGEMM_GPU FB only operators" OFF)
 
 if((NOT FBGEMM_CPU_ONLY) AND

fbgemm_gpu/setup.py

@@ -76,6 +76,11 @@ def from_args(cls, argv: List[str]):
             default=None,
             help="NCCL (libnccl.so.2) filepath. This is required for building certain targets.",
         )
+        parser.add_argument(
+            "--use_fb_only",
+            action="store_true",
+            help="Build FB only operators.",
+        )
         parser.add_argument(
             "--cxxprefix",
             type=str,
@@ -276,6 +281,10 @@ def _get_cxx11_abi():
                 ]
             )
 
+        if self.args.use_fb_only:
+            print("[SETUP.PY] Building the FB ONLY operators of FBGEMM_GPU ...")
+            cmake_args.append("-DUSE_FB_ONLY=ON")
+
         if self.args.cxxprefix:
             print("[SETUP.PY] Setting CMake flags ...")
             path = self.args.cxxprefix