add dense TBE to template and enable VBE support (pytorch#2620)

Summary:

- make the dense TBE headers into a template
- add VBE options to the dense TBE header

Differential Revision: D57017981

fbgemm_gpu/codegen/genscript/generate_backward_split.py

@@ -55,7 +55,7 @@ def render_backward_templates(
         ):
             if nobag and (weighted or vbe):
                 continue
-            if kwargs.get("dense") and (vbe or ssd):
+            if kwargs.get("dense") and ssd:
                 continue
             if ssd and (vbe or is_gwd):
                 continue
@@ -152,7 +152,7 @@ def generate_backward_split_gpu(**kwargs: Any) -> None:
                 **kwargs,
             )
 
-        # Generate the backward splits (non-dense)
+        # Generate the backward splits
         # We generate only the API to preserve the backward compatibility if
         # has_gpu_support=True
         if not kwargs.get("dense"):
@@ -197,6 +197,17 @@ def generate_backward_split_gpu(**kwargs: Any) -> None:
                             filename, is_fbcode=args.is_fbcode, ssd=ssd, **kwargs
                         )
 
+        else:
+            template_filepath = (
+                "training/backward/embedding_backward_split_host_template.cpp"
+            )
+            filename = "gen_embedding_backward_split_dense.cpp"
+            CodeTemplate.load(template_filepath).write(
+                filename,
+                is_forward=False,
+                **kwargs,
+            )
+
     @staticmethod
     def generate_backward_split_cpu(**kwargs: Any) -> None:
         """

fbgemm_gpu/codegen/genscript/generate_forward_split.py

@@ -39,7 +39,7 @@ def render_forward_templates(
         ):
             if nobag and (weighted or vbe):
                 continue
-            if dense and (vbe or ssd):
+            if dense and ssd:
                 continue
             if ssd and (vbe or is_gwd):
                 continue
@@ -64,6 +64,7 @@ def render_forward_templates(
                 is_gwd=is_gwd,
             )
 
+    @staticmethod
     def generate_pt2_wrappers() -> None:
         # Generate PT2 forward wrapper (CUDA)
         CodeTemplate.load(

fbgemm_gpu/codegen/genscript/optimizers.py

@@ -40,7 +40,7 @@ def dense() -> Dict[str, Any]:
         ),
         "has_cpu_support": True,
         "has_gpu_support": True,
-        "has_vbe_support": False,
+        "has_vbe_support": True,
         "has_global_weight_decay_support": False,
         "has_ssd_support": False,
     }