From ba7df90792bfe3bcf46759536b4f0ff4edb2503d Mon Sep 17 00:00:00 2001 From: Benson Ma Date: Sat, 28 Sep 2024 17:19:47 -0700 Subject: [PATCH] Redefine FBGEMM targets with gpu_cpp_library [20/N] (#3193) Summary: Pull Request resolved: https://github.com/pytorch/FBGEMM/pull/3193 X-link: https://github.com/facebookresearch/FBGEMM/pull/289 - Combine `gpu_cpp_library` targets into `fbgemm_gpu:sparse_ops` Differential Revision: D63591546 --- fbgemm_gpu/bench/batched_unary_embeddings_benchmark.py | 6 +----- .../bench/histogram_binning_calibration_benchmark.py | 6 +----- fbgemm_gpu/bench/jagged_tensor_benchmark.py | 6 +----- fbgemm_gpu/bench/quantize_ops_benchmark.py | 6 +----- fbgemm_gpu/bench/stride_gemm_benchmark.py | 6 +----- fbgemm_gpu/fbgemm_gpu/batched_unary_embeddings_ops.py | 9 +++------ fbgemm_gpu/fbgemm_gpu/sparse_ops.py | 4 +--- fbgemm_gpu/test/batched_unary_embeddings_test.py | 8 ++------ fbgemm_gpu/test/jagged/common.py | 8 +------- fbgemm_gpu/test/layout_transform_ops_test.py | 8 ++------ fbgemm_gpu/test/quantize/common.py | 7 +------ fbgemm_gpu/test/sparse/common.py | 7 +------ 12 files changed, 16 insertions(+), 65 deletions(-) diff --git a/fbgemm_gpu/bench/batched_unary_embeddings_benchmark.py b/fbgemm_gpu/bench/batched_unary_embeddings_benchmark.py index 6091bcc8e8..6d354900a1 100644 --- a/fbgemm_gpu/bench/batched_unary_embeddings_benchmark.py +++ b/fbgemm_gpu/bench/batched_unary_embeddings_benchmark.py @@ -25,11 +25,7 @@ else: from fbgemm_gpu.bench.bench_utils import benchmark_torch_function - if torch.version.hip: - torch.ops.load_library("//deeplearning/fbgemm/fbgemm_gpu:sparse_ops_hip") - else: - torch.ops.load_library("//deeplearning/fbgemm/fbgemm_gpu:sparse_ops") - torch.ops.load_library("//deeplearning/fbgemm/fbgemm_gpu:sparse_ops_cpu") + torch.ops.load_library("//deeplearning/fbgemm/fbgemm_gpu:sparse_ops") def generate_unary_feature( diff --git a/fbgemm_gpu/bench/histogram_binning_calibration_benchmark.py b/fbgemm_gpu/bench/histogram_binning_calibration_benchmark.py index c919199ee2..e43106b8c9 100644 --- a/fbgemm_gpu/bench/histogram_binning_calibration_benchmark.py +++ b/fbgemm_gpu/bench/histogram_binning_calibration_benchmark.py @@ -21,11 +21,7 @@ # pyre-ignore[21] from fbgemm_gpu import open_source # noqa: F401 except Exception: - if torch.version.hip: - torch.ops.load_library("//deeplearning/fbgemm/fbgemm_gpu:sparse_ops_hip") - else: - torch.ops.load_library("//deeplearning/fbgemm/fbgemm_gpu:sparse_ops") - torch.ops.load_library("//deeplearning/fbgemm/fbgemm_gpu:sparse_ops_cpu") + torch.ops.load_library("//deeplearning/fbgemm/fbgemm_gpu:sparse_ops") def benchmark_hbc_function( diff --git a/fbgemm_gpu/bench/jagged_tensor_benchmark.py b/fbgemm_gpu/bench/jagged_tensor_benchmark.py index 46337701e3..814f950b07 100644 --- a/fbgemm_gpu/bench/jagged_tensor_benchmark.py +++ b/fbgemm_gpu/bench/jagged_tensor_benchmark.py @@ -31,10 +31,7 @@ else: from fbgemm_gpu.bench.bench_utils import benchmark_torch_function - if torch.version.hip: - torch.ops.load_library("//deeplearning/fbgemm/fbgemm_gpu:sparse_ops_hip") - else: - torch.ops.load_library("//deeplearning/fbgemm/fbgemm_gpu:sparse_ops") + torch.ops.load_library("//deeplearning/fbgemm/fbgemm_gpu:sparse_ops") torch.ops.load_library( "//deeplearning/fbgemm/fbgemm_gpu:permute_pooled_embedding_ops_cpu" ) @@ -47,7 +44,6 @@ torch.ops.load_library( "//deeplearning/fbgemm/fbgemm_gpu:permute_multi_embedding_ops_gpu" ) - torch.ops.load_library("//deeplearning/fbgemm/fbgemm_gpu:sparse_ops_cpu") @click.group() diff --git a/fbgemm_gpu/bench/quantize_ops_benchmark.py b/fbgemm_gpu/bench/quantize_ops_benchmark.py index 81eb07bead..54755fff67 100644 --- a/fbgemm_gpu/bench/quantize_ops_benchmark.py +++ b/fbgemm_gpu/bench/quantize_ops_benchmark.py @@ -34,11 +34,7 @@ else: from fbgemm_gpu.bench.bench_utils import benchmark_torch_function - if torch.version.hip: - torch.ops.load_library("//deeplearning/fbgemm/fbgemm_gpu:sparse_ops_hip") - else: - torch.ops.load_library("//deeplearning/fbgemm/fbgemm_gpu:sparse_ops") - torch.ops.load_library("//deeplearning/fbgemm/fbgemm_gpu:sparse_ops_cpu") + torch.ops.load_library("//deeplearning/fbgemm/fbgemm_gpu:sparse_ops") @click.group() diff --git a/fbgemm_gpu/bench/stride_gemm_benchmark.py b/fbgemm_gpu/bench/stride_gemm_benchmark.py index 2609f7fbfa..bca34b8a92 100644 --- a/fbgemm_gpu/bench/stride_gemm_benchmark.py +++ b/fbgemm_gpu/bench/stride_gemm_benchmark.py @@ -20,11 +20,7 @@ # pyre-ignore[21] from fbgemm_gpu import open_source # noqa: F401 except Exception: - if torch.version.hip: - torch.ops.load_library("//deeplearning/fbgemm/fbgemm_gpu:sparse_ops_hip") - else: - torch.ops.load_library("//deeplearning/fbgemm/fbgemm_gpu:sparse_ops") - torch.ops.load_library("//deeplearning/fbgemm/fbgemm_gpu:sparse_ops_cpu") + torch.ops.load_library("//deeplearning/fbgemm/fbgemm_gpu:sparse_ops") @click.group() diff --git a/fbgemm_gpu/fbgemm_gpu/batched_unary_embeddings_ops.py b/fbgemm_gpu/fbgemm_gpu/batched_unary_embeddings_ops.py index 5ef0f1c329..db2260df4d 100644 --- a/fbgemm_gpu/fbgemm_gpu/batched_unary_embeddings_ops.py +++ b/fbgemm_gpu/fbgemm_gpu/batched_unary_embeddings_ops.py @@ -13,16 +13,13 @@ import torch +from fbgemm_gpu.utils.loader import load_torch_module + try: # pyre-ignore[21] from fbgemm_gpu import open_source # noqa: F401 except Exception: - if torch.version.hip: - torch.ops.load_library("//deeplearning/fbgemm/fbgemm_gpu:sparse_ops_hip") - else: - torch.ops.load_library("//deeplearning/fbgemm/fbgemm_gpu:sparse_ops") - - torch.ops.load_library("//deeplearning/fbgemm/fbgemm_gpu:sparse_ops_cpu") + load_torch_module("//deeplearning/fbgemm/fbgemm_gpu:sparse_ops") def wrap_weight_to_parameter(weights: List[torch.Tensor]) -> List[torch.Tensor]: diff --git a/fbgemm_gpu/fbgemm_gpu/sparse_ops.py b/fbgemm_gpu/fbgemm_gpu/sparse_ops.py index de9a21ef97..71e0e2ccc6 100644 --- a/fbgemm_gpu/fbgemm_gpu/sparse_ops.py +++ b/fbgemm_gpu/fbgemm_gpu/sparse_ops.py @@ -20,20 +20,18 @@ from fbgemm_gpu import open_source # noqa: F401 except Exception: load_torch_module("//deeplearning/fbgemm/fbgemm_gpu:merge_pooled_embeddings") + load_torch_module("//deeplearning/fbgemm/fbgemm_gpu:sparse_ops") if torch.version.hip: - torch.ops.load_library("//deeplearning/fbgemm/fbgemm_gpu:sparse_ops_hip") torch.ops.load_library( "//deeplearning/fbgemm/fbgemm_gpu/codegen:embedding_ops_hip" ) else: - torch.ops.load_library("//deeplearning/fbgemm/fbgemm_gpu:sparse_ops") torch.ops.load_library("//deeplearning/fbgemm/fbgemm_gpu/codegen:embedding_ops") torch.ops.load_library("//deeplearning/fbgemm/fbgemm_gpu:input_combine") - torch.ops.load_library("//deeplearning/fbgemm/fbgemm_gpu:sparse_ops_cpu") torch.ops.load_library( "//deeplearning/fbgemm/fbgemm_gpu:merge_pooled_embeddings_cpu" ) diff --git a/fbgemm_gpu/test/batched_unary_embeddings_test.py b/fbgemm_gpu/test/batched_unary_embeddings_test.py index 3f1124effa..37d0aaf7b5 100644 --- a/fbgemm_gpu/test/batched_unary_embeddings_test.py +++ b/fbgemm_gpu/test/batched_unary_embeddings_test.py @@ -26,14 +26,10 @@ from test_utils import gpu_unavailable except Exception: - if torch.version.hip: - torch.ops.load_library("//deeplearning/fbgemm/fbgemm_gpu:sparse_ops_hip") - else: - torch.ops.load_library("//deeplearning/fbgemm/fbgemm_gpu:sparse_ops") - - torch.ops.load_library("//deeplearning/fbgemm/fbgemm_gpu:sparse_ops_cpu") from fbgemm_gpu.test.test_utils import gpu_unavailable + torch.ops.load_library("//deeplearning/fbgemm/fbgemm_gpu:sparse_ops") + # Relative tolerances # pyre-fixme[5]: Global expression must be annotated. diff --git a/fbgemm_gpu/test/jagged/common.py b/fbgemm_gpu/test/jagged/common.py index 3bdcacf980..491176f769 100644 --- a/fbgemm_gpu/test/jagged/common.py +++ b/fbgemm_gpu/test/jagged/common.py @@ -24,13 +24,7 @@ open_source: bool = getattr(fbgemm_gpu, "open_source", False) if not open_source: - if torch.version.hip: - torch.ops.load_library("//deeplearning/fbgemm/fbgemm_gpu:sparse_ops_hip") - else: - torch.ops.load_library("//deeplearning/fbgemm/fbgemm_gpu:sparse_ops") - - torch.ops.load_library("//deeplearning/fbgemm/fbgemm_gpu:sparse_ops_cpu") - + torch.ops.load_library("//deeplearning/fbgemm/fbgemm_gpu:sparse_ops") suppressed_list: List[HealthCheck] = ( [HealthCheck.differing_executors] diff --git a/fbgemm_gpu/test/layout_transform_ops_test.py b/fbgemm_gpu/test/layout_transform_ops_test.py index 658d773f3a..57a7b02638 100644 --- a/fbgemm_gpu/test/layout_transform_ops_test.py +++ b/fbgemm_gpu/test/layout_transform_ops_test.py @@ -22,14 +22,10 @@ from test_utils import gpu_unavailable except Exception: - if torch.version.hip: - torch.ops.load_library("//deeplearning/fbgemm/fbgemm_gpu:sparse_ops_hip") - else: - torch.ops.load_library("//deeplearning/fbgemm/fbgemm_gpu:sparse_ops") - - torch.ops.load_library("//deeplearning/fbgemm/fbgemm_gpu:sparse_ops_cpu") from fbgemm_gpu.test.test_utils import gpu_unavailable + torch.ops.load_library("//deeplearning/fbgemm/fbgemm_gpu:sparse_ops") + MAX_EXAMPLES = 20 diff --git a/fbgemm_gpu/test/quantize/common.py b/fbgemm_gpu/test/quantize/common.py index 5333cc8933..6a720a1748 100644 --- a/fbgemm_gpu/test/quantize/common.py +++ b/fbgemm_gpu/test/quantize/common.py @@ -23,12 +23,7 @@ from fbgemm_gpu import open_source # noqa: F401 except Exception: - if torch.version.hip: - torch.ops.load_library("//deeplearning/fbgemm/fbgemm_gpu:sparse_ops_hip") - else: - torch.ops.load_library("//deeplearning/fbgemm/fbgemm_gpu:sparse_ops") - - torch.ops.load_library("//deeplearning/fbgemm/fbgemm_gpu:sparse_ops_cpu") + torch.ops.load_library("//deeplearning/fbgemm/fbgemm_gpu:sparse_ops") # Eigen/Python round 0.5 away from 0, Numpy rounds to even round_to_nearest: Callable[[npt.NDArray], npt.NDArray] = np.vectorize(round) diff --git a/fbgemm_gpu/test/sparse/common.py b/fbgemm_gpu/test/sparse/common.py index 69b6e34772..8abddca752 100644 --- a/fbgemm_gpu/test/sparse/common.py +++ b/fbgemm_gpu/test/sparse/common.py @@ -23,12 +23,7 @@ open_source: bool = getattr(fbgemm_gpu, "open_source", False) if not open_source: - if torch.version.hip: - torch.ops.load_library("//deeplearning/fbgemm/fbgemm_gpu:sparse_ops_hip") - else: - torch.ops.load_library("//deeplearning/fbgemm/fbgemm_gpu:sparse_ops") - - torch.ops.load_library("//deeplearning/fbgemm/fbgemm_gpu:sparse_ops_cpu") + torch.ops.load_library("//deeplearning/fbgemm/fbgemm_gpu:sparse_ops") torch.ops.load_library("//deeplearning/fbgemm/fbgemm_gpu/codegen:index_select_ops") suppressed_list: List[HealthCheck] = (