Skip to content

Commit c195bc4

Browse files
committed
enable cpp kernl building
1 parent 6a8887f commit c195bc4

File tree

1 file changed

+14
-13
lines changed

1 file changed

+14
-13
lines changed

setup.py

Lines changed: 14 additions & 13 deletions
Original file line numberDiff line numberDiff line change
@@ -385,19 +385,20 @@ def get_extensions():
385385
extra_compile_args["cxx"].extend(
386386
["-O3" if not debug_mode else "-O0", "-fdiagnostics-color=always"]
387387
)
388-
389-
# TODO(future PR): make this work without using `TORCH_VERSION_AT_LEAST_2_7`,
390-
# because we should not be using anything from `torchao` to build `torchao`.
391-
# if use_cpu_kernels and is_linux and TORCH_VERSION_AT_LEAST_2_7:
392-
# if torch._C._cpu._is_avx512_supported():
393-
# extra_compile_args["cxx"].extend(
394-
# [
395-
# "-DCPU_CAPABILITY_AVX512",
396-
# "-march=native",
397-
# "-mfma",
398-
# "-fopenmp",
399-
# ]
400-
# )
388+
if (
389+
use_cpu_kernels
390+
and is_linux
391+
and hasattr(torch._C._cpu, "_is_avx512_supported")
392+
and torch._C._cpu._is_avx512_supported()
393+
):
394+
extra_compile_args["cxx"].extend(
395+
[
396+
"-DCPU_CAPABILITY_AVX512",
397+
"-march=native",
398+
"-mfma",
399+
"-fopenmp",
400+
]
401+
)
401402

402403
if debug_mode:
403404
extra_compile_args["cxx"].append("-g")

0 commit comments

Comments
 (0)