Refactor the codegen directory, pt 10 (#2479)

Summary: Pull Request resolved: #2479 - Migrate optimizer templates in `codegen/` over to `codegen/training/optimizer` Reviewed By: spcyppt Differential Revision: D55762970 fbshipit-source-id: 14544d8e928594a453a6f569ac6f3dfcb508a9f1
pytorch · Apr 5, 2024 · aa4127f · aa4127f
1 parent 126cdb1
commit aa4127f
Show file tree

Hide file tree

Showing 7 changed files with 30 additions and 27 deletions.
diff --git a/fbgemm_gpu/codegen/genscript/generate_backward_split.py b/fbgemm_gpu/codegen/genscript/generate_backward_split.py
@@ -92,7 +92,9 @@ def generate_backward_split_gpu(**kwargs: Any) -> None:
             )
 
         # Generate optimizer kernel
-        CodeTemplate.load("embedding_optimizer_split_device_kernel_template.cuh").write(
+        CodeTemplate.load(
+            "training/optimizer/embedding_optimizer_split_device_kernel_template.cuh"
+        ).write(
             f"gen_embedding_optimizer_{optimizer}_split_device_kernel.cuh", **kwargs
         )
 

diff --git a/fbgemm_gpu/codegen/genscript/generate_embedding_optimizer.py b/fbgemm_gpu/codegen/genscript/generate_embedding_optimizer.py
@@ -40,32 +40,33 @@ def generate(**kwargs: Any) -> None:
         )
         kwargs["args"] = kwargs["args"].cuda
 
-        # Generate CUDA host code
-        CodeTemplate.load("embedding_optimizer_split_template.cu").write(
-            f"gen_embedding_optimizer_{optimizer}_split_cuda.cu", **kwargs
-        )
-
-        # Generate CUDA kernel code
-        CodeTemplate.load("embedding_optimizer_split_kernel_template.cu").write(
-            f"gen_embedding_optimizer_{optimizer}_split_kernel.cu", **kwargs
-        )
-
-        # Generate host code
-        CodeTemplate.load("embedding_optimizer_split_host_template.cpp").write(
-            f"gen_embedding_optimizer_{optimizer}_split.cpp", **kwargs
-        )
-
-        # Generates Python invoker for CUDA
-        CodeTemplate.load("split_embedding_optimizer_codegen.template").write(
-            f"split_embedding_optimizer_{optimizer}.py",
-            is_fbcode=args.is_fbcode,
-            **kwargs,
-        )
-
-        # Generate optimizer kernel headers
-        CodeTemplate.load("embedding_optimizer_split_device_kernel_template.cuh").write(
-            f"gen_embedding_optimizer_{optimizer}_split_device_kernel.cuh", **kwargs
-        )
+        PREFIX = "training/optimizer"
+
+        for template_filepath, filename in [
+            (  # CUDA host code
+                f"{PREFIX}/embedding_optimizer_split_template.cu",
+                f"gen_embedding_optimizer_{optimizer}_split_cuda.cu",
+            ),
+            (  # CUDA kernel code
+                f"{PREFIX}/embedding_optimizer_split_kernel_template.cu",
+                f"gen_embedding_optimizer_{optimizer}_split_kernel.cu",
+            ),
+            (  # CPU code
+                f"{PREFIX}/embedding_optimizer_split_host_template.cpp",
+                f"gen_embedding_optimizer_{optimizer}_split.cpp",
+            ),
+            (  # Optimizer kernel headers
+                f"{PREFIX}/embedding_optimizer_split_device_kernel_template.cuh",
+                f"gen_embedding_optimizer_{optimizer}_split_device_kernel.cuh",
+            ),
+            (  # Python kernel invokers
+                "training/python/split_embedding_optimizer_codegen.template",
+                f"split_embedding_optimizer_{optimizer}.py",
+            ),
+        ]:
+            CodeTemplate.load(template_filepath).write(
+                filename, is_fbcode=args.is_fbcode, **kwargs
+            )
 
 
 def main() -> None:

diff --git a/...ptimizer_split_device_kernel_template.cuh → ...ptimizer_split_device_kernel_template.cuh b/...ptimizer_split_device_kernel_template.cuh → ...ptimizer_split_device_kernel_template.cuh
diff --git a/...bedding_optimizer_split_host_template.cpp → ...bedding_optimizer_split_host_template.cpp b/...bedding_optimizer_split_host_template.cpp → ...bedding_optimizer_split_host_template.cpp
diff --git a/...edding_optimizer_split_kernel_template.cu → ...edding_optimizer_split_kernel_template.cu b/...edding_optimizer_split_kernel_template.cu → ...edding_optimizer_split_kernel_template.cu
diff --git a/...gen/embedding_optimizer_split_template.cu → ...zer/embedding_optimizer_split_template.cu b/...gen/embedding_optimizer_split_template.cu → ...zer/embedding_optimizer_split_template.cu
diff --git a/...plit_embedding_optimizer_codegen.template → ...plit_embedding_optimizer_codegen.template b/...plit_embedding_optimizer_codegen.template → ...plit_embedding_optimizer_codegen.template