Decouple embedding_ssd_{}_pt2_autograd from CUDA files (#4389)

Summary:
Pull Request resolved: #4389

X-link: facebookresearch/FBGEMM#1460

This file contains `{{ bwd_mdesc }}_embedding_codegen_lookup_{{ optimizer }}_function_pt2`, which is not a CUDA-only function (i.e. it can also be called just fine from the CPU). Therefore, we extract it from the CUDA codegen to put it inside the CPU codegen instead.

Reviewed By: q10

Differential Revision: D76866940

fbshipit-source-id: 2c09ae92cc0d0474b4cf524a38145595d1f42ba7

Author: PatriceVignola

fbgemm_gpu/fbgemm_gpu/utils/loader.py

@@ -14,7 +14,10 @@
 
 
 def load_torch_module(
-    unified_path: str, cuda_path: Optional[str] = None, hip_path: Optional[str] = None
+    unified_path: str,
+    cuda_path: Optional[str] = None,
+    hip_path: Optional[str] = None,
+    mtia_path: Optional[str] = None,
 ) -> None:
     try:
         torch.ops.load_library(unified_path)
@@ -24,9 +27,16 @@ def load_torch_module(
                 hip_path = f"{unified_path}_hip"
             torch.ops.load_library(hip_path)
         else:
-            if not cuda_path:
-                cuda_path = f"{unified_path}_cuda"
-            torch.ops.load_library(cuda_path)
+            try:
+                # pyre-ignore-next-line[21]
+                import mtia.host_runtime.torch_mtia.dynamic_library  # noqa
+
+                if mtia_path is not None:
+                    torch.ops.load_library(mtia_path)
+            except OSError:
+                if not cuda_path:
+                    cuda_path = f"{unified_path}_cuda"
+                torch.ops.load_library(cuda_path)
 
 
 def load_torch_module_bc(new_path: str, old_path: str) -> None: