From 541297daae63f74a90bac52c5db3540bf84bc971 Mon Sep 17 00:00:00 2001
From: Nikita Shulga <nshulga@meta.com>
Date: Tue, 17 Jun 2025 17:50:57 -0700
Subject: [PATCH] [Build] Allow metal shaders to include ATen headers (#156256)

No-op change that will be used later to share structs between CPU and Metal
Pull Request resolved: https://github.com/pytorch/pytorch/pull/156256
Approved by: https://github.com/dcci
---
 cmake/Metal.cmake                 | 2 +-
 torch/utils/_cpp_embed_headers.py | 3 ++-
 2 files changed, 3 insertions(+), 2 deletions(-)

diff --git a/cmake/Metal.cmake b/cmake/Metal.cmake
index cc8e1932f1a1..c9565e2fc0e9 100644
--- a/cmake/Metal.cmake
+++ b/cmake/Metal.cmake
@@ -8,7 +8,7 @@ if(WERROR)
 endif()
 
 function(metal_to_air SRC TARGET FLAGS)
-    add_custom_command(COMMAND xcrun metal -c ${SRC} -I ${CMAKE_SOURCE_DIR} -o ${TARGET} ${FLAGS} ${METAL_CFLAGS}
+    add_custom_command(COMMAND xcrun metal -c ${SRC} -I ${CMAKE_SOURCE_DIR} -I ${CMAKE_SOURCE_DIR}/aten/src -o ${TARGET} ${FLAGS} ${METAL_CFLAGS}
                        DEPENDS ${SRC}
                        OUTPUT ${TARGET}
                        COMMENT "Compiling ${SRC} to ${TARGET}"
diff --git a/torch/utils/_cpp_embed_headers.py b/torch/utils/_cpp_embed_headers.py
index 9cb0fee3a3f8..6bcf8d583f0c 100644
--- a/torch/utils/_cpp_embed_headers.py
+++ b/torch/utils/_cpp_embed_headers.py
@@ -39,7 +39,8 @@ def embed_headers(
     fname: str, include_dirs: Optional[Union[Sequence[str], Sequence[Path], str]] = None
 ) -> str:
     if include_dirs is None:
-        include_dirs = [Path(__file__).parent.parent.parent]
+        base_dir = Path(__file__).parent.parent.parent
+        include_dirs = [base_dir, base_dir / "aten" / "src"]
     elif isinstance(include_dirs, str):
         include_dirs = [Path(include_dirs)]
     else: