[AOTInductor] Add wrapper and kernel code to debug code logging (#153181)

benjaminglass1 · pytorchmergebot · commit 01cbf5a30a29 · 2025-05-10T15:31:18.000Z
This is a simple PR to make the AOTInductor wrapper and kernel code get output by `TORCH_COMPILE_DEBUG=1`. Pull Request resolved: #153181 Approved by: https://github.com/desertfire
diff --git a/torch/_inductor/codecache.py b/torch/_inductor/codecache.py
@@ -98,6 +98,7 @@
 from .runtime import autotune_cache
 from .runtime.autotune_cache import AutotuneCacheBundler
 from .triton_bundler import TritonBundler
+from .virtualized import V
 
 
 if config.is_fbcode():
@@ -1604,6 +1605,12 @@ def compile(
             specified_dir=specified_output_path,
         )
 
+        # Log the AOTInductor wrapper and kernel code, if needed.
+        with tempfile.NamedTemporaryFile("w+") as t:
+            t.writelines((wrapper_code, "\n", kernel_code, "\n"))
+            t.flush()
+            V.debug.output_code(t.name, extension="cpp")
+
         if config.aot_inductor.package:
             generated_files.append(wrapper_path)
             if not config.aot_inductor.package_cpp_only:
diff --git a/torch/_inductor/debug.py b/torch/_inductor/debug.py
@@ -559,8 +559,8 @@ def draw_orig_fx_graph(
             dot_graph_shape=config.trace.dot_graph_shape,
         )
 
-    def output_code(self, filename: str) -> None:
-        shutil.copy(filename, self.filename("output_code.py"))
+    def output_code(self, filename: str, extension: str = "py") -> None:
+        shutil.copy(filename, self.filename(f"output_code.{extension}"))
 
     def log_inductor_triton_kernel_to_post_grad_node_info(
         self, filename: str = "inductor_triton_kernel_to_post_grad_nodes.json"

Original file line number	Diff line number	Diff line change
`@@ -559,8 +559,8 @@ def draw_orig_fx_graph(`
`559`	`559`	`dot_graph_shape=config.trace.dot_graph_shape,`
`560`	`560`	`)`
`561`	`561`
`562`		`- def output_code(self, filename: str) -> None:`
`563`		`- shutil.copy(filename, self.filename("output_code.py"))`
	`562`	`+ def output_code(self, filename: str, extension: str = "py") -> None:`
	`563`	`+ shutil.copy(filename, self.filename(f"output_code.{extension}"))`
`564`	`564`
`565`	`565`	`def log_inductor_triton_kernel_to_post_grad_node_info(`
`566`	`566`	`self, filename: str = "inductor_triton_kernel_to_post_grad_nodes.json"`