OSSForks/pytorch
1
0
Fork 0
mirror of https://github.com/zebrajr/pytorch.git synced 2025-12-06 12:20:52 +01:00
Code Issues Actions 38 Packages Projects Releases Wiki Activity

[inductor] optimize cpp builder configuration code (#129577)

Browse Source 
Changes:
1. Combine choose isa condition dispatch code.
2. Unificate MacOS openmp configuration code.
3. Clean up useless code.

Co-authored-by: Jason Ansel <jansel@jansel.net>
Pull Request resolved: https://github.com/pytorch/pytorch/pull/129577
Approved by: https://github.com/jgong5, https://github.com/jansel
This commit is contained in:
Xu Han 2024-06-28 15:08:53 +00:00 committed by PyTorch MergeBot
parent eabe6574c0
commit 2e3ff394bf
1 changed files with 14 additions and 24 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

38
torch/_inductor/cpp_builder.py
View File

@ -320,26 +320,14 @@ def _get_optimization_cflags() -> List[str]:
if not config.cpp.enable_floating_point_contract_flag:
cflags.append("ffp-contract=off")
if config.is_fbcode():
# FIXME: passing `-fopenmp` adds libgomp.so to the generated shared library's dependencies.
# This causes `ldopen` to fail in fbcode, because libgomp does not exist in the default paths.
# We will fix it later by exposing the lib path.
return cflags
if sys.platform == "darwin":
# Per https://mac.r-project.org/openmp/ right way to pass `openmp` flags to MacOS is via `-Xclang`
# Also, `-march=native` is unrecognized option on M1
cflags.append("Xclang")
else:
if sys.platform != "darwin":
# https://stackoverflow.com/questions/65966969/why-does-march-native-not-work-on-apple-m1
# `-march=native` is unrecognized option on M1
if platform.machine() == "ppc64le":
cflags.append("mcpu=native")
else:
cflags.append("march=native")
# Internal cannot find libgomp.so
if not config.is_fbcode():
cflags.append("fopenmp")
return cflags
@ -547,13 +535,13 @@ def _get_build_args_of_chosen_isa(vec_isa: VecISA):
build_flags = [vec_isa.build_arch_flags()]
if config.is_fbcode() and vec_isa != invalid_vec_isa:
cap = str(vec_isa).upper()
macros = [
f"CPU_CAPABILITY={cap}",
f"CPU_CAPABILITY_{cap}",
f"HAVE_{cap}_CPU_DEFINITION",
]
if config.is_fbcode():
cap = str(vec_isa).upper()
macros = [
f"CPU_CAPABILITY={cap}",
f"CPU_CAPABILITY_{cap}",
f"HAVE_{cap}_CPU_DEFINITION",
]
return macros, build_flags
@ -615,6 +603,10 @@ def _get_openmp_args(cpp_compiler):
libs: List[str] = []
passthough_args: List[str] = []
if _IS_MACOS:
# Per https://mac.r-project.org/openmp/ right way to pass `openmp` flags to MacOS is via `-Xclang`
cflags.append("Xclang")
cflags.append("fopenmp")
from torch._inductor.codecache import (
homebrew_libomp,
is_conda_llvm_openmp_installed,
@ -665,10 +657,8 @@ def _get_openmp_args(cpp_compiler):
# /openmp, /openmp:llvm
# llvm on Windows, new openmp: https://devblogs.microsoft.com/cppblog/msvc-openmp-update/
# msvc openmp: https://learn.microsoft.com/zh-cn/cpp/build/reference/openmp-enable-openmp-2-0-support?view=msvc-170
cflags.append("openmp")
cflags.append("openmp:experimental") # MSVC CL
libs = []
else:
if config.is_fbcode():
include_dir_paths.append(build_paths.openmp())