mirror of
https://github.com/pytorch/pytorch.git
synced 2025-10-20 12:54:11 +08:00
ebcacd5e87
Summary:
Make sleef dependency public so that `ATen_CPU_{capability}` libs can depend on it
Pull Request resolved: https://github.com/pytorch/pytorch/pull/37381
Test Plan: CI
Differential Revision: D21273443
Pulled By: malfet
fbshipit-source-id: 7f756c7f3c605e51cf0c27ea37f687913cd48708
45 lines
1.3 KiB
Python
45 lines
1.3 KiB
Python
load("@rules_cc//cc:defs.bzl", "cc_library")
|
|
|
|
CPU_CAPABILITY_NAMES = ["DEFAULT", "AVX", "AVX2"]
|
|
CAPABILITY_COMPILER_FLAGS = {
|
|
"AVX2": ["-mavx2", "-mfma"],
|
|
"AVX": ["-mavx"],
|
|
"DEFAULT": [],
|
|
}
|
|
|
|
PREFIX = "aten/src/ATen/native/"
|
|
|
|
def intern_build_aten_ops(copts, deps):
|
|
for cpu_capability in CPU_CAPABILITY_NAMES:
|
|
srcs = []
|
|
for impl in native.glob(
|
|
[
|
|
PREFIX + "cpu/*.cpp",
|
|
PREFIX + "quantized/cpu/kernels/*.cpp",
|
|
]):
|
|
name = impl.replace(PREFIX, "")
|
|
out = PREFIX + name + "." + cpu_capability + ".cpp"
|
|
native.genrule(
|
|
name = name + "_" + cpu_capability + "_cp",
|
|
srcs = [impl],
|
|
outs = [out],
|
|
cmd = "cp $< $@",
|
|
)
|
|
srcs.append(out)
|
|
|
|
cc_library(
|
|
name = "ATen_CPU_" + cpu_capability,
|
|
srcs = srcs,
|
|
copts = copts + [
|
|
"-DCPU_CAPABILITY=" + cpu_capability,
|
|
"-DCPU_CAPABILITY_" + cpu_capability,
|
|
] + CAPABILITY_COMPILER_FLAGS[cpu_capability],
|
|
deps = deps,
|
|
linkstatic = 1,
|
|
)
|
|
cc_library(
|
|
name = "ATen_CPU",
|
|
deps = [":ATen_CPU_" + cpu_capability for cpu_capability in CPU_CAPABILITY_NAMES],
|
|
linkstatic = 1,
|
|
)
|