Re-organize TBE tests, pt 5 (#2305)

Summary: Pull Request resolved: #2305 - Re-organize TBE tests, pt 5 Reviewed By: sryap Differential Revision: D53331779 fbshipit-source-id: c65c8565fc3bec883c70b5797549dd532e8d1add
pytorch · Feb 2, 2024 · c218dc3 · c218dc3
1 parent 79f4862
commit c218dc3
Show file tree

Hide file tree

Showing 26 changed files with 1,007 additions and 34 deletions.
diff --git a/fbgemm_gpu/test/tbe/cache/__init__.py b/fbgemm_gpu/test/tbe/cache/__init__.py
@@ -0,0 +1,6 @@
+#!/usr/bin/env python3
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+# All rights reserved.
+#
+# This source code is licensed under the BSD-style license found in the
+# LICENSE file in the root directory of this source tree.
diff --git a/fbgemm_gpu/test/tbe/cache_common.py → fbgemm_gpu/test/tbe/cache/cache_common.py b/fbgemm_gpu/test/tbe/cache_common.py → fbgemm_gpu/test/tbe/cache/cache_common.py
@@ -25,7 +25,7 @@
 
 from hypothesis import Verbosity
 
-from .common import open_source
+from ..common import open_source
 
 if open_source:
     # pyre-ignore[21]

diff --git a/fbgemm_gpu/test/tbe/cache_overflow_test.py → ...gpu/test/tbe/cache/cache_overflow_test.py b/fbgemm_gpu/test/tbe/cache_overflow_test.py → ...gpu/test/tbe/cache/cache_overflow_test.py
@@ -15,10 +15,9 @@
 from fbgemm_gpu.split_embedding_utils import to_device
 from fbgemm_gpu.split_table_batched_embeddings_ops_training import DEFAULT_ASSOC
 
+from ..common import assert_torch_equal
 from .cache_common import generate_cache_tbes, gpu_unavailable
 
-from .common import assert_torch_equal
-
 
 class CacheOverflowTest(unittest.TestCase):
     @unittest.skipIf(*gpu_unavailable)

diff --git a/fbgemm_gpu/test/tbe/cache_test.py → fbgemm_gpu/test/tbe/cache/cache_test.py b/fbgemm_gpu/test/tbe/cache_test.py → fbgemm_gpu/test/tbe/cache/cache_test.py
@@ -30,9 +30,9 @@
 )
 from hypothesis import given, settings
 
-from .cache_common import generate_cache_tbes, gpu_unavailable, optests, VERBOSITY
+from ..common import assert_torch_equal, MAX_EXAMPLES  # noqa E402
 
-from .common import assert_torch_equal, MAX_EXAMPLES  # noqa E402
+from .cache_common import generate_cache_tbes, gpu_unavailable, optests, VERBOSITY
 
 
 @optests.generate_opcheck_tests(fast=True)

diff --git a/fbgemm_gpu/test/tbe/failures_dict_fast.json → ...pu/test/tbe/cache/failures_dict_fast.json b/fbgemm_gpu/test/tbe/failures_dict_fast.json → ...pu/test/tbe/cache/failures_dict_fast.json
diff --git a/.../test/tbe/linearize_cache_indices_test.py → ...tbe/cache/linearize_cache_indices_test.py b/.../test/tbe/linearize_cache_indices_test.py → ...tbe/cache/linearize_cache_indices_test.py
@@ -12,8 +12,8 @@
 import torch
 from hypothesis import Verbosity
 
-from . import common  # noqa E402
-from .common import open_source
+from .. import common  # noqa E402
+from ..common import open_source
 
 if open_source:
     # pyre-ignore[21]

diff --git a/fbgemm_gpu/test/tbe/lxu_cache_test.py → fbgemm_gpu/test/tbe/cache/lxu_cache_test.py b/fbgemm_gpu/test/tbe/lxu_cache_test.py → fbgemm_gpu/test/tbe/cache/lxu_cache_test.py
@@ -20,8 +20,8 @@
 from hypothesis import given, settings, Verbosity
 from torch import Tensor
 
-from . import common  # noqa E402
-from .common import MAX_EXAMPLES, open_source
+from .. import common  # noqa E402
+from ..common import MAX_EXAMPLES, open_source
 
 if open_source:
     # pyre-ignore[21]

diff --git a/fbgemm_gpu/test/tbe/inference/__init__.py b/fbgemm_gpu/test/tbe/inference/__init__.py
@@ -0,0 +1,6 @@
+#!/usr/bin/env python3
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+# All rights reserved.
+#
+# This source code is licensed under the BSD-style license found in the
+# LICENSE file in the root directory of this source tree.
diff --git a/fbgemm_gpu/test/tbe/inference/failures_dict_fast.json b/fbgemm_gpu/test/tbe/inference/failures_dict_fast.json
@@ -0,0 +1,316 @@
+{
+  "_description": "This is a dict containing failures for tests autogenerated by generate_opcheck_tests. For more details, please see https://docs.google.com/document/d/1Pj5HRZvdOq3xpFpbEjUZp2hBovhy7Wnxw14m6lF2154/edit",
+  "_version": 1,
+  "data": {
+    "fbgemm::asynchronous_complete_cumsum": {},
+    "fbgemm::bounds_check_indices": {},
+    "fbgemm::dense_embedding_codegen_lookup_function": {
+      "BackwardDenseTest.test_autograd_registration__test_backward_dense": {
+        "comment": "",
+        "status": "xfail"
+      },
+      "BackwardDenseTest.test_faketensor__test_backward_dense": {
+        "comment": "",
+        "status": "xfail"
+      }
+    },
+    "fbgemm::direct_mapped_lru_cache_populate_byte": {},
+    "fbgemm::direct_mapped_lxu_cache_lookup": {
+      "NBitForwardTest.test_faketensor__test_nbit_forward_uvm_cache": {
+        "comment": "",
+        "status": "xfail"
+      }
+    },
+    "fbgemm::emb_inplace_update": {},
+    "fbgemm::get_unique_indices": {
+      "LXUCacheTest.test_faketensor__test_unique_lxu_cache_lookup": {
+        "comment": "",
+        "status": "xfail"
+      }
+    },
+    "fbgemm::int_nbit_split_embedding_codegen_lookup_function": {
+      "NBitCacheTest.test_faketensor__test_nbit_cache_miss_counter": {
+        "comment": "",
+        "status": "xsuccess"
+      },
+      "NBitForwardTest.test_faketensor__test_nbit_forward_fused_pooled_emb_quant": {
+        "comment": "",
+        "status": "xsuccess"
+      },
+      "NBitForwardTest.test_faketensor__test_nbit_forward_uvm_cache": {
+        "comment": "",
+        "status": "xsuccess"
+      },
+      "NBitSplitEmbeddingsTest.test_faketensor__test_int_nbit_split_embedding_uvm_caching_codegen_lookup_function": {
+        "comment": "",
+        "status": "xsuccess"
+      }
+    },
+    "fbgemm::int_nbit_split_embedding_uvm_caching_codegen_lookup_function": {
+      "NBitSplitEmbeddingsTest.test_faketensor__test_int_nbit_split_embedding_uvm_caching_codegen_lookup_function": {
+        "comment": "",
+        "status": "xfail"
+      },
+      "NBitSplitEmbeddingsTest.test_schema__test_int_nbit_split_embedding_uvm_caching_codegen_lookup_function": {
+        "comment": "",
+        "status": "xfail"
+      }
+    },
+    "fbgemm::lfu_cache_populate": {
+      "BackwardSGDTest.test_faketensor__test_backward_sgd": {
+        "comment": "",
+        "status": "skip"
+      },
+      "BackwardSGDTest.test_faketensor__test_backward_sgd_really_long_segments": {
+        "comment": "",
+        "status": "skip"
+      }
+    },
+    "fbgemm::lfu_cache_populate_byte": {
+      "NBitForwardTest.test_faketensor__test_nbit_forward_uvm_cache": {
+        "comment": "",
+        "status": "xfail"
+      }
+    },
+    "fbgemm::linearize_cache_indices": {
+      "BackwardAdagradTest.test_faketensor__test_backward_adagrad_fp16_pmMEAN": {
+        "comment": "",
+        "status": "skip"
+      },
+      "BackwardAdagradTest.test_faketensor__test_backward_adagrad_fp16_pmNONE": {
+        "comment": "",
+        "status": "skip"
+      },
+      "BackwardAdagradTest.test_faketensor__test_backward_adagrad_fp16_pmSUM": {
+        "comment": "",
+        "status": "skip"
+      },
+      "BackwardAdagradTest.test_faketensor__test_backward_adagrad_fp32_pmMEAN": {
+        "comment": "",
+        "status": "skip"
+      },
+      "BackwardAdagradTest.test_faketensor__test_backward_adagrad_fp32_pmNONE": {
+        "comment": "",
+        "status": "skip"
+      },
+      "BackwardAdagradTest.test_faketensor__test_backward_adagrad_fp32_pmSUM": {
+        "comment": "",
+        "status": "skip"
+      },
+      "BackwardSGDTest.test_faketensor__test_backward_sgd": {
+        "comment": "",
+        "status": "xfail"
+      },
+      "BackwardSGDTest.test_faketensor__test_backward_sgd_really_long_segments": {
+        "comment": "",
+        "status": "xfail"
+      },
+      "CacheTest.test_faketensor__test_cache_miss_counter": {
+        "comment": "",
+        "status": "xfail"
+      },
+      "LXUCacheTest.test_faketensor__test_unique_lxu_cache_lookup": {
+        "comment": "",
+        "status": "xfail"
+      },
+      "LinearizeCacheIndicesTest.test_faketensor__test_linearize_cache_indices": {
+        "comment": "",
+        "status": "xfail"
+      },
+      "NBitCacheTest.test_faketensor__test_nbit_cache_miss_counter": {
+        "comment": "",
+        "status": "xfail"
+      },
+      "NBitForwardTest.test_faketensor__test_nbit_forward_uvm_cache": {
+        "comment": "",
+        "status": "xfail"
+      },
+      "NBitSplitEmbeddingsTest.test_faketensor__test_int_nbit_split_embedding_uvm_caching_codegen_lookup_function": {
+        "comment": "",
+        "status": "xfail"
+      },
+      "SplitTableBatchedEmbeddingsTest.test_faketensor__test_stb_uvm_cache_stats": {
+        "comment": "",
+        "status": "xfail"
+      }
+    },
+    "fbgemm::linearize_cache_indices_from_row_idx": {
+      "LinearizeCacheIndicesTest.test_faketensor__test_linearize_cache_indices_from_row_idx": {
+        "comment": "",
+        "status": "xfail"
+      }
+    },
+    "fbgemm::lru_cache_populate": {
+      "BackwardAdagradTest.test_faketensor__test_backward_adagrad_fp16_pmMEAN": {
+        "comment": "",
+        "status": "skip"
+      },
+      "BackwardAdagradTest.test_faketensor__test_backward_adagrad_fp16_pmNONE": {
+        "comment": "",
+        "status": "skip"
+      },
+      "BackwardAdagradTest.test_faketensor__test_backward_adagrad_fp16_pmSUM": {
+        "comment": "",
+        "status": "skip"
+      },
+      "BackwardAdagradTest.test_faketensor__test_backward_adagrad_fp32_pmMEAN": {
+        "comment": "",
+        "status": "skip"
+      },
+      "BackwardAdagradTest.test_faketensor__test_backward_adagrad_fp32_pmNONE": {
+        "comment": "",
+        "status": "skip"
+      },
+      "BackwardAdagradTest.test_faketensor__test_backward_adagrad_fp32_pmSUM": {
+        "comment": "",
+        "status": "skip"
+      },
+      "BackwardSGDTest.test_faketensor__test_backward_sgd": {
+        "comment": "",
+        "status": "skip"
+      },
+      "BackwardSGDTest.test_faketensor__test_backward_sgd_really_long_segments": {
+        "comment": "",
+        "status": "skip"
+      },
+      "CacheTest.test_faketensor__test_cache_miss_counter": {
+        "comment": "",
+        "status": "skip"
+      },
+      "SplitTableBatchedEmbeddingsTest.test_faketensor__test_stb_uvm_cache_stats": {
+        "comment": "",
+        "status": "skip"
+      }
+    },
+    "fbgemm::lru_cache_populate_byte": {},
+    "fbgemm::lxu_cache_flush": {},
+    "fbgemm::lxu_cache_locking_counter_decrement": {},
+    "fbgemm::lxu_cache_lookup": {
+      "BackwardAdagradTest.test_faketensor__test_backward_adagrad_fp16_pmMEAN": {
+        "comment": "",
+        "status": "xfail"
+      },
+      "BackwardAdagradTest.test_faketensor__test_backward_adagrad_fp16_pmNONE": {
+        "comment": "",
+        "status": "xfail"
+      },
+      "BackwardAdagradTest.test_faketensor__test_backward_adagrad_fp16_pmSUM": {
+        "comment": "",
+        "status": "xfail"
+      },
+      "BackwardAdagradTest.test_faketensor__test_backward_adagrad_fp32_pmMEAN": {
+        "comment": "",
+        "status": "xfail"
+      },
+      "BackwardAdagradTest.test_faketensor__test_backward_adagrad_fp32_pmNONE": {
+        "comment": "",
+        "status": "xfail"
+      },
+      "BackwardAdagradTest.test_faketensor__test_backward_adagrad_fp32_pmSUM": {
+        "comment": "",
+        "status": "xfail"
+      },
+      "BackwardSGDTest.test_faketensor__test_backward_sgd": {
+        "comment": "",
+        "status": "xfail"
+      },
+      "BackwardSGDTest.test_faketensor__test_backward_sgd_really_long_segments": {
+        "comment": "",
+        "status": "xfail"
+      },
+      "CacheTest.test_faketensor__test_cache_miss_counter": {
+        "comment": "",
+        "status": "xfail"
+      },
+      "LXUCacheTest.test_faketensor__test_lxu_cache_lookup": {
+        "comment": "",
+        "status": "xfail"
+      },
+      "LXUCacheTest.test_faketensor__test_unique_lxu_cache_lookup": {
+        "comment": "",
+        "status": "xfail"
+      },
+      "NBitCacheTest.test_faketensor__test_nbit_cache_miss_counter": {
+        "comment": "",
+        "status": "xfail"
+      },
+      "NBitForwardTest.test_faketensor__test_nbit_forward_uvm_cache": {
+        "comment": "",
+        "status": "xfail"
+      },
+      "NBitSplitEmbeddingsTest.test_faketensor__test_int_nbit_split_embedding_uvm_caching_codegen_lookup_function": {
+        "comment": "",
+        "status": "xfail"
+      },
+      "SplitTableBatchedEmbeddingsTest.test_faketensor__test_stb_uvm_cache_stats": {
+        "comment": "",
+        "status": "xfail"
+      }
+    },
+    "fbgemm::new_managed_tensor": {},
+    "fbgemm::new_unified_tensor": {
+      "NBitCacheTest.test_faketensor__test_nbit_cache_miss_counter": {
+        "comment": "",
+        "status": "xfail"
+      },
+      "NBitCacheTest.test_faketensor__test_nbit_cache_update_function": {
+        "comment": "",
+        "status": "xfail"
+      },
+      "NBitForwardTest.test_faketensor__test_nbit_forward_gpu_no_cache": {
+        "comment": "",
+        "status": "xfail"
+      },
+      "NBitForwardTest.test_faketensor__test_nbit_forward_gpu_no_cache_fp8_2048": {
+        "comment": "",
+        "status": "xfail"
+      },
+      "NBitForwardTest.test_faketensor__test_nbit_forward_uvm_cache": {
+        "comment": "",
+        "status": "xfail"
+      },
+      "NBitSplitEmbeddingsTest.test_faketensor__test_int_nbit_split_embedding_uvm_caching_codegen_lookup_function": {
+        "comment": "",
+        "status": "xfail"
+      }
+    },
+    "fbgemm::pruned_array_lookup": {
+      "NBitForwardTest.test_faketensor__test_nbit_forward_uvm_cache": {
+        "comment": "",
+        "status": "xfail"
+      }
+    },
+    "fbgemm::pruned_hashmap_insert": {},
+    "fbgemm::pruned_hashmap_lookup": {
+      "NBitForwardTest.test_faketensor__test_nbit_forward_uvm_cache": {
+        "comment": "",
+        "status": "xfail"
+      }
+    },
+    "fbgemm::reset_weight_momentum": {
+      "SplitTableBatchedEmbeddingsTest.test_faketensor__test_reset_embedding_weight_momentum": {
+        "comment": "",
+        "status": "skip"
+      },
+      "SplitTableBatchedEmbeddingsTest.test_schema__test_reset_embedding_weight_momentum": {
+        "comment": "",
+        "status": "skip"
+      }
+    },
+    "fbgemm::split_embedding_codegen_forward_unweighted_cuda": {},
+    "fbgemm::split_embedding_codegen_forward_weighted_cuda": {},
+    "fbgemm::split_embedding_codegen_lookup_adagrad_function": {},
+    "fbgemm::split_embedding_codegen_lookup_adagrad_function_cpu": {},
+    "fbgemm::split_embedding_codegen_lookup_adam_function": {},
+    "fbgemm::split_embedding_codegen_lookup_lamb_function": {},
+    "fbgemm::split_embedding_codegen_lookup_lars_sgd_function": {},
+    "fbgemm::split_embedding_codegen_lookup_none_function": {},
+    "fbgemm::split_embedding_codegen_lookup_partial_rowwise_adam_function": {},
+    "fbgemm::split_embedding_codegen_lookup_partial_rowwise_lamb_function": {},
+    "fbgemm::split_embedding_codegen_lookup_rowwise_adagrad_function": {},
+    "fbgemm::split_embedding_codegen_lookup_rowwise_adagrad_function_cpu": {},
+    "fbgemm::split_embedding_codegen_lookup_rowwise_weighted_adagrad_function": {},
+    "fbgemm::split_embedding_codegen_lookup_sgd_function": {},
+    "fbgemm::split_embedding_codegen_lookup_sgd_function_cpu": {}
+  }
+}
diff --git a/..._gpu/test/tbe/inference_converter_test.py → ...tbe/inference/inference_converter_test.py b/..._gpu/test/tbe/inference_converter_test.py → ...tbe/inference/inference_converter_test.py
@@ -36,8 +36,8 @@
 from hypothesis import given, settings, Verbosity
 from torch import nn
 
-from . import common  # noqa E402
-from .common import open_source
+from .. import common  # noqa E402
+from ..common import open_source
 
 
 if open_source: