From 076dad385d3b4fcbe9c39c24c603496117b3e583 Mon Sep 17 00:00:00 2001 From: Benson Ma Date: Wed, 28 Aug 2024 22:35:26 -0700 Subject: [PATCH] Redefine FBGEMM targets with split_embeddings_utils [1/N] (#3049) Summary: Pull Request resolved: https://github.com/pytorch/FBGEMM/pull/3049 X-link: https://github.com/facebookresearch/FBGEMM/pull/146 - Redefine `split_embeddings_utils`, `embedding_ops_common`, and `index_select_ops` targets with `gpu_cpp_library()` Reviewed By: duduyi2013 Differential Revision: D61889098 fbshipit-source-id: 7c6fae12a5e2bccf382d7746bf510b6b812a3ba5 --- .../training/backward/embedding_backward_split_template.cu | 2 +- .../forward/embedding_forward_split_kernel_template.cu | 1 - .../training/forward/embedding_forward_split_template.cu | 3 ++- fbgemm_gpu/fbgemm_gpu/sparse_ops.py | 4 +--- 4 files changed, 4 insertions(+), 6 deletions(-) diff --git a/fbgemm_gpu/codegen/training/backward/embedding_backward_split_template.cu b/fbgemm_gpu/codegen/training/backward/embedding_backward_split_template.cu index 62758d949f..8c52cadb6f 100644 --- a/fbgemm_gpu/codegen/training/backward/embedding_backward_split_template.cu +++ b/fbgemm_gpu/codegen/training/backward/embedding_backward_split_template.cu @@ -22,8 +22,8 @@ #include "fbgemm_gpu/embedding_backward_template_helpers.cuh" #include "fbgemm_gpu/utils/tensor_accessor.h" -#include "fbgemm_gpu/split_embeddings_utils.cuh" #include "fbgemm_gpu/sparse_ops.h" +#include "fbgemm_gpu/split_embeddings_utils.cuh" #include "fbgemm_gpu/utils/ops_utils.h" using Tensor = at::Tensor; diff --git a/fbgemm_gpu/codegen/training/forward/embedding_forward_split_kernel_template.cu b/fbgemm_gpu/codegen/training/forward/embedding_forward_split_kernel_template.cu index 1f4b21ef3c..7b07d1d1da 100644 --- a/fbgemm_gpu/codegen/training/forward/embedding_forward_split_kernel_template.cu +++ b/fbgemm_gpu/codegen/training/forward/embedding_forward_split_kernel_template.cu @@ -34,7 +34,6 @@ {%- set locs_or_addrs_idx = "row_idx" if ssd else "cache_idx" %} #include "fbgemm_gpu/embedding_forward_template_helpers.cuh" -#include "fbgemm_gpu/split_embeddings_cache_cuda.cuh" using Tensor = at::Tensor; using namespace fbgemm_gpu; diff --git a/fbgemm_gpu/codegen/training/forward/embedding_forward_split_template.cu b/fbgemm_gpu/codegen/training/forward/embedding_forward_split_template.cu index 19bbdfc4e5..947231e8e7 100644 --- a/fbgemm_gpu/codegen/training/forward/embedding_forward_split_template.cu +++ b/fbgemm_gpu/codegen/training/forward/embedding_forward_split_template.cu @@ -34,14 +34,15 @@ {%- if not is_index_select %} //////////////////////////////////////////////////////////////////////////////// // Required for op registrations -#include "fbgemm_gpu/utils/ops_utils.h" //////////////////////////////////////////////////////////////////////////////// +#include "fbgemm_gpu/utils/ops_utils.h" {%- endif %} #include "fbgemm_gpu/embedding_forward_template_helpers.cuh" #include "fbgemm_gpu/split_embeddings_cache_cuda.cuh" using Tensor = at::Tensor; using namespace fbgemm_gpu; + //////////////////////////////////////////////////////////////////////////////// // External Function Declarations //////////////////////////////////////////////////////////////////////////////// diff --git a/fbgemm_gpu/fbgemm_gpu/sparse_ops.py b/fbgemm_gpu/fbgemm_gpu/sparse_ops.py index 29a1236db8..024170a68b 100644 --- a/fbgemm_gpu/fbgemm_gpu/sparse_ops.py +++ b/fbgemm_gpu/fbgemm_gpu/sparse_ops.py @@ -26,9 +26,7 @@ torch.ops.load_library( "//deeplearning/fbgemm/fbgemm_gpu/codegen:embedding_ops_hip" ) - torch.ops.load_library( - "//deeplearning/fbgemm/fbgemm_gpu/codegen:index_select_ops_hip" - ) + else: torch.ops.load_library("//deeplearning/fbgemm/fbgemm_gpu:sparse_ops") torch.ops.load_library(