From 2a7b3ab54341705dc554ee39e8f3aa19be896134 Mon Sep 17 00:00:00 2001 From: Benson Ma Date: Tue, 30 Jan 2024 17:07:20 -0800 Subject: [PATCH] Fix sparse ops tests (#2296) Summary: Pull Request resolved: https://github.com/pytorch/FBGEMM/pull/2296 - Fix sparse ops tests to pass again Reviewed By: spcyppt Differential Revision: D53244522 fbshipit-source-id: aa6d5365dd34ba0a61b3d65cad1bbd3e717b187e --- fbgemm_gpu/test/sparse/block_bucketize_test.py | 7 +------ fbgemm_gpu/test/sparse/common.py | 11 +++++++++++ fbgemm_gpu/test/sparse/cumsum_test.py | 7 +------ .../sparse/histogram_binning_calibration_test.py | 7 +------ fbgemm_gpu/test/sparse/index_select_test.py | 8 +------- fbgemm_gpu/test/sparse/misc_ops_test.py | 8 +------- fbgemm_gpu/test/sparse/pack_segments_test.py | 6 +----- fbgemm_gpu/test/sparse/permute_embeddings_test.py | 7 +------ fbgemm_gpu/test/sparse/permute_indices_test.py | 12 ++++++------ .../test/sparse/permute_sparse_features_test.py | 7 +------ fbgemm_gpu/test/sparse/reorder_batched_test.py | 6 +----- fbgemm_gpu/test/uvm/copy_test.py | 4 +--- 12 files changed, 27 insertions(+), 63 deletions(-) diff --git a/fbgemm_gpu/test/sparse/block_bucketize_test.py b/fbgemm_gpu/test/sparse/block_bucketize_test.py index 7a451ba5a8..bacd731666 100644 --- a/fbgemm_gpu/test/sparse/block_bucketize_test.py +++ b/fbgemm_gpu/test/sparse/block_bucketize_test.py @@ -10,16 +10,11 @@ import unittest from typing import Optional, Type -import fbgemm_gpu import hypothesis.strategies as st import torch from hypothesis import given, settings, Verbosity -from .common import extend_test_class - - -# pyre-fixme[16]: Module `fbgemm_gpu` has no attribute `open_source`. -open_source: bool = getattr(fbgemm_gpu, "open_source", False) +from .common import extend_test_class, open_source if open_source: # pyre-ignore[21] diff --git a/fbgemm_gpu/test/sparse/common.py b/fbgemm_gpu/test/sparse/common.py index cd43d3c8ae..2b245b99cc 100644 --- a/fbgemm_gpu/test/sparse/common.py +++ b/fbgemm_gpu/test/sparse/common.py @@ -18,6 +18,17 @@ from torch._utils_internal import get_file_path_2 from torch.testing._internal.optests import generate_opcheck_tests +# pyre-fixme[16]: Module `fbgemm_gpu` has no attribute `open_source`. +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/codegen:index_select_ops") suppressed_list: List[HealthCheck] = ( # pyre-fixme[16]: Module `HealthCheck` has no attribute `differing_executors`. diff --git a/fbgemm_gpu/test/sparse/cumsum_test.py b/fbgemm_gpu/test/sparse/cumsum_test.py index c80c669cea..c0c351f5e6 100644 --- a/fbgemm_gpu/test/sparse/cumsum_test.py +++ b/fbgemm_gpu/test/sparse/cumsum_test.py @@ -9,17 +9,12 @@ import unittest -import fbgemm_gpu import hypothesis.strategies as st import numpy as np import torch from hypothesis import given, settings, Verbosity -from .common import extend_test_class - - -# pyre-fixme[16]: Module `fbgemm_gpu` has no attribute `open_source`. -open_source: bool = getattr(fbgemm_gpu, "open_source", False) +from .common import extend_test_class, open_source if open_source: # pyre-ignore[21] diff --git a/fbgemm_gpu/test/sparse/histogram_binning_calibration_test.py b/fbgemm_gpu/test/sparse/histogram_binning_calibration_test.py index 930c6c2932..531210f291 100644 --- a/fbgemm_gpu/test/sparse/histogram_binning_calibration_test.py +++ b/fbgemm_gpu/test/sparse/histogram_binning_calibration_test.py @@ -10,16 +10,11 @@ import random import unittest -import fbgemm_gpu import hypothesis.strategies as st import torch from hypothesis import given, settings, Verbosity -from .common import extend_test_class - - -# pyre-fixme[16]: Module `fbgemm_gpu` has no attribute `open_source`. -open_source: bool = getattr(fbgemm_gpu, "open_source", False) +from .common import extend_test_class, open_source if open_source: # pyre-ignore[21] diff --git a/fbgemm_gpu/test/sparse/index_select_test.py b/fbgemm_gpu/test/sparse/index_select_test.py index 6b174135e6..872b255d72 100644 --- a/fbgemm_gpu/test/sparse/index_select_test.py +++ b/fbgemm_gpu/test/sparse/index_select_test.py @@ -14,18 +14,12 @@ import unittest from typing import Callable, Dict, List -import fbgemm_gpu - import hypothesis.strategies as st import numpy as np import torch from hypothesis import given, settings, Verbosity -from .common import extend_test_class - - -# pyre-fixme[16]: Module `fbgemm_gpu` has no attribute `open_source`. -open_source: bool = getattr(fbgemm_gpu, "open_source", False) +from .common import extend_test_class, open_source if open_source: # pyre-ignore[21] diff --git a/fbgemm_gpu/test/sparse/misc_ops_test.py b/fbgemm_gpu/test/sparse/misc_ops_test.py index a8536e4971..a209e7f0a0 100644 --- a/fbgemm_gpu/test/sparse/misc_ops_test.py +++ b/fbgemm_gpu/test/sparse/misc_ops_test.py @@ -12,18 +12,12 @@ import unittest from typing import Type, Union -import fbgemm_gpu - import hypothesis.strategies as st import numpy as np import torch from hypothesis import given, settings, Verbosity -from .common import extend_test_class - - -# pyre-fixme[16]: Module `fbgemm_gpu` has no attribute `open_source`. -open_source: bool = getattr(fbgemm_gpu, "open_source", False) +from .common import extend_test_class, open_source if open_source: # pyre-ignore[21] diff --git a/fbgemm_gpu/test/sparse/pack_segments_test.py b/fbgemm_gpu/test/sparse/pack_segments_test.py index e1fe73915c..fc6a39f659 100644 --- a/fbgemm_gpu/test/sparse/pack_segments_test.py +++ b/fbgemm_gpu/test/sparse/pack_segments_test.py @@ -11,16 +11,12 @@ import unittest from typing import Callable, Optional -import fbgemm_gpu import hypothesis.strategies as st import numpy as np import torch from hypothesis import given, settings -from .common import extend_test_class - -# pyre-fixme[16]: Module `fbgemm_gpu` has no attribute `open_source`. -open_source: bool = getattr(fbgemm_gpu, "open_source", False) +from .common import extend_test_class, open_source if open_source: # pyre-ignore[21] diff --git a/fbgemm_gpu/test/sparse/permute_embeddings_test.py b/fbgemm_gpu/test/sparse/permute_embeddings_test.py index d99dbe4202..12e23c2059 100644 --- a/fbgemm_gpu/test/sparse/permute_embeddings_test.py +++ b/fbgemm_gpu/test/sparse/permute_embeddings_test.py @@ -11,16 +11,11 @@ import unittest from typing import Any, Callable, Tuple -import fbgemm_gpu import hypothesis.strategies as st import torch from hypothesis import given, settings, Verbosity -from .common import extend_test_class, permute_indices_ref_ - - -# pyre-fixme[16]: Module `fbgemm_gpu` has no attribute `open_source`. -open_source: bool = getattr(fbgemm_gpu, "open_source", False) +from .common import extend_test_class, open_source, permute_indices_ref_ if open_source: # pyre-ignore[21] diff --git a/fbgemm_gpu/test/sparse/permute_indices_test.py b/fbgemm_gpu/test/sparse/permute_indices_test.py index 70a96611e2..4b41c23b4e 100644 --- a/fbgemm_gpu/test/sparse/permute_indices_test.py +++ b/fbgemm_gpu/test/sparse/permute_indices_test.py @@ -12,16 +12,16 @@ from itertools import accumulate from typing import List, Optional -import fbgemm_gpu import hypothesis.strategies as st import torch from hypothesis import given, settings, Verbosity -from .common import extend_test_class, permute_indices_ref_, permute_scripted - - -# pyre-fixme[16]: Module `fbgemm_gpu` has no attribute `open_source`. -open_source: bool = getattr(fbgemm_gpu, "open_source", False) +from .common import ( + extend_test_class, + open_source, + permute_indices_ref_, + permute_scripted, +) if open_source: # pyre-ignore[21] diff --git a/fbgemm_gpu/test/sparse/permute_sparse_features_test.py b/fbgemm_gpu/test/sparse/permute_sparse_features_test.py index 0fc2c65469..72e3eee9b5 100644 --- a/fbgemm_gpu/test/sparse/permute_sparse_features_test.py +++ b/fbgemm_gpu/test/sparse/permute_sparse_features_test.py @@ -11,16 +11,11 @@ import unittest from typing import cast, Optional, Tuple -import fbgemm_gpu import hypothesis.strategies as st import torch from hypothesis import given, settings -from .common import extend_test_class, permute_indices_ref_ - - -# pyre-fixme[16]: Module `fbgemm_gpu` has no attribute `open_source`. -open_source: bool = getattr(fbgemm_gpu, "open_source", False) +from .common import extend_test_class, open_source, permute_indices_ref_ if open_source: # pyre-ignore[21] diff --git a/fbgemm_gpu/test/sparse/reorder_batched_test.py b/fbgemm_gpu/test/sparse/reorder_batched_test.py index 5d032d763d..13a21a262b 100644 --- a/fbgemm_gpu/test/sparse/reorder_batched_test.py +++ b/fbgemm_gpu/test/sparse/reorder_batched_test.py @@ -10,15 +10,11 @@ import random import unittest -import fbgemm_gpu import hypothesis.strategies as st import torch from hypothesis import given, settings, Verbosity -from .common import extend_test_class - -# pyre-fixme[16]: Module `fbgemm_gpu` has no attribute `open_source`. -open_source: bool = getattr(fbgemm_gpu, "open_source", False) +from .common import extend_test_class, open_source if open_source: # pyre-ignore[21] diff --git a/fbgemm_gpu/test/uvm/copy_test.py b/fbgemm_gpu/test/uvm/copy_test.py index 5c719cd96e..f82e4e03f7 100644 --- a/fbgemm_gpu/test/uvm/copy_test.py +++ b/fbgemm_gpu/test/uvm/copy_test.py @@ -25,9 +25,7 @@ from fbgemm_gpu.test.test_utils import gpu_available, gpu_unavailable, skipIfRocm if gpu_available: - # pyre-ignore[21] - from fbgemm_gpu.uvm import cudaMemAdvise, cudaMemoryAdvise, cudaMemPrefetchAsync - + import fbgemm_gpu.uvm MAX_EXAMPLES = 40