Move TBE SSD into tbe/ directory (#2749)

Summary: Pull Request resolved: #2749 - Move TBE SSD into tbe/ directory Reviewed By: sryap Differential Revision: D58696893 fbshipit-source-id: 7404b9bdeef877d5738362c5112164618d717889
pytorch · Jun 24, 2024 · dbe678d · dbe678d
1 parent 61e1058
commit dbe678d
Show file tree

Hide file tree

Showing 7 changed files with 1,394 additions and 1,280 deletions.
diff --git a/fbgemm_gpu/bench/ssd_table_batched_embeddings_benchmark.py b/fbgemm_gpu/bench/ssd_table_batched_embeddings_benchmark.py
@@ -25,20 +25,20 @@
     round_up,
     TBERequest,
 )
+from fbgemm_gpu.split_table_batched_embeddings_ops_common import (
+    CacheAlgorithm,
+    EmbeddingLocation,
+    PoolingMode,
+)
 from fbgemm_gpu.split_table_batched_embeddings_ops_inference import (
     IntNBitTableBatchedEmbeddingBagsCodegen,
 )
 from fbgemm_gpu.split_table_batched_embeddings_ops_training import (
     ComputeDevice,
     SplitTableBatchedEmbeddingBagsCodegen,
 )
-from fbgemm_gpu.ssd_split_table_batched_embeddings_ops import (
-    CacheAlgorithm,
-    EmbeddingLocation,
-    PoolingMode,
-    # Inference
+from fbgemm_gpu.tbe.ssd import (
     SSDIntNBitTableBatchedEmbeddingBags,
-    # Training
     SSDTableBatchedEmbeddingBags,
 )
 from torch.profiler import profile