From 5b412a0b229b4cea7ea0e5c5241ca6318627f0c2 Mon Sep 17 00:00:00 2001
From: Benson Ma <bensonma415@meta.com>
Date: Fri, 22 Nov 2024 17:38:20 -0800
Subject: [PATCH] Patch D66310520 to make it build in OSS (#3409)

Summary:
Pull Request resolved: https://github.com/pytorch/FBGEMM/pull/3409

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

- Patch D66310520 to make the code build in OSS

Differential Revision: D66399304
---
 fbgemm_gpu/FbgemmGpu.cmake                    | 23 ++++++++++++++++++-
 .../genscript/generate_backward_split.py      |  2 +-
 .../fbgemm_gpu/rocm/split_embeddings_common.h |  1 +
 3 files changed, 24 insertions(+), 2 deletions(-)

diff --git a/fbgemm_gpu/FbgemmGpu.cmake b/fbgemm_gpu/FbgemmGpu.cmake
index 74172ae3da..f20f64473a 100644
--- a/fbgemm_gpu/FbgemmGpu.cmake
+++ b/fbgemm_gpu/FbgemmGpu.cmake
@@ -295,19 +295,30 @@ foreach(optimizer ${SSD_OPTIMIZERS})
       "gen_embedding_backward_${optimizer}_ssd_${wdesc}_kernel_cta.cu"
       "gen_embedding_backward_${optimizer}_ssd_${wdesc}_kernel_warp.cu")
   endforeach()
+
   foreach(wdesc weighted unweighted)
     list(APPEND gen_gpu_kernel_source_files
       "gen_embedding_backward_${optimizer}_ssd_${wdesc}_vbe_cuda.cu"
       "gen_embedding_backward_${optimizer}_ssd_${wdesc}_vbe_kernel_cta.cu"
       "gen_embedding_backward_${optimizer}_ssd_${wdesc}_vbe_kernel_warp.cu")
   endforeach()
-
 endforeach()
 
 list(APPEND gen_defused_optim_py_files
     ${CMAKE_BINARY_DIR}/optimizer_args.py)
 
 
+################################################################################
+# FBGEMM_GPU Generated HIP-Specific Sources
+################################################################################
+
+set(gen_hip_kernel_source_files)
+foreach(wdesc weighted unweighted unweighted_nobag)
+  list(APPEND gen_hip_kernel_source_files
+    "gen_embedding_backward_split_${wdesc}_device_kernel_hip.hip")
+endforeach()
+
+
 ################################################################################
 # FBGEMM_GPU Generated Sources
 ################################################################################
@@ -563,6 +574,9 @@ set(fbgemm_gpu_sources_gpu_gen
   ${gen_gpu_host_source_files}
   ${gen_defused_optim_source_files})
 
+set(fbgemm_gpu_sources_hip_gen
+  ${gen_hip_kernel_source_files})
+
 if(USE_ROCM)
   prepend_filepaths(
     PREFIX ${CMAKE_BINARY_DIR}
@@ -573,6 +587,11 @@ if(USE_ROCM)
     PREFIX ${CMAKE_BINARY_DIR}
     INPUT ${fbgemm_gpu_sources_gpu_gen}
     OUTPUT fbgemm_gpu_sources_gpu_gen)
+
+  prepend_filepaths(
+    PREFIX ${CMAKE_BINARY_DIR}
+    INPUT ${fbgemm_gpu_sources_hip_gen}
+    OUTPUT fbgemm_gpu_sources_hip_gen)
 endif()
 
 
@@ -591,6 +610,8 @@ gpu_cpp_library(
   GPU_SRCS
     ${fbgemm_gpu_sources_gpu_static}
     ${fbgemm_gpu_sources_gpu_gen}
+  HIP_SPECIFIC_SRCS
+    ${fbgemm_gpu_sources_hip_gen}
   OTHER_SRCS
     ${asmjit_sources}
     ${fbgemm_sources}
diff --git a/fbgemm_gpu/codegen/genscript/generate_backward_split.py b/fbgemm_gpu/codegen/genscript/generate_backward_split.py
index 5e01defc83..c977148578 100644
--- a/fbgemm_gpu/codegen/genscript/generate_backward_split.py
+++ b/fbgemm_gpu/codegen/genscript/generate_backward_split.py
@@ -390,7 +390,7 @@ def generate() -> None:
             BackwardSplitGenerator.generate_backward_split(
                 ssd_tensors=ssd_tensors, **optimizer
             )
-        BackwardSplitGenerator.generate_rocm_backward_split(**optimizer)
+        BackwardSplitGenerator.generate_rocm_backward_split()
 
         # Generate common device kernels for backwards
         BackwardSplitGenerator.generate_backward_device()
diff --git a/fbgemm_gpu/include/fbgemm_gpu/rocm/split_embeddings_common.h b/fbgemm_gpu/include/fbgemm_gpu/rocm/split_embeddings_common.h
index 0058548e22..b3a56c4b52 100644
--- a/fbgemm_gpu/include/fbgemm_gpu/rocm/split_embeddings_common.h
+++ b/fbgemm_gpu/include/fbgemm_gpu/rocm/split_embeddings_common.h
@@ -21,6 +21,7 @@
  *
  ******************************************************************************/
 #pragma once
+#include <c10/util/Half.h>
 #include <hip/hip_fp16.h>
 #include <hip/hip_runtime.h>