frozenleaves/pytorch
1
0
Fork 0
mirror of https://github.com/pytorch/pytorch.git synced 2025-10-20 21:14:14 +08:00
Code Issues Packages Projects Releases Wiki Activity

[Build] Allow metal shaders to include ATen headers (#156256)

Browse Source 
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
Download Patch File Download Diff File Expand all files Collapse all files

2
cmake/Metal.cmake
View File

@ -8,7 +8,7 @@ if(WERROR)
endif() endif()
function(metal_to_air SRC TARGET FLAGS) 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} DEPENDS ${SRC}
OUTPUT ${TARGET} OUTPUT ${TARGET}
COMMENT "Compiling ${SRC} to ${TARGET}" COMMENT "Compiling ${SRC} to ${TARGET}"

3
torch/utils/_cpp_embed_headers.py
View File

@ -39,7 +39,8 @@ def embed_headers(
fname: str, include_dirs: Optional[Union[Sequence[str], Sequence[Path], str]] = None fname: str, include_dirs: Optional[Union[Sequence[str], Sequence[Path], str]] = None
) -> str: ) -> str:
if include_dirs is None: 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): elif isinstance(include_dirs, str):
include_dirs = [Path(include_dirs)] include_dirs = [Path(include_dirs)]
else: else: