[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
This commit is contained in:
Nikita Shulga
2025-06-17 17:50:57 -07:00
committed by PyTorch MergeBot
parent 3dabc351bb
commit 541297daae
2 changed files with 3 additions and 2 deletions

View File

@ -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}"

View File

@ -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: