Fix sparse ops tests (#2296)

Summary:

- Fix sparse ops tests to pass again

Differential Revision: D53244522

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]

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`.

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]

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]

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]

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]

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]

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]

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]

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]

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]

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