forked from pytorch/pytorch
-
Notifications
You must be signed in to change notification settings - Fork 0
/
aten.bzl
44 lines (40 loc) · 1.32 KB
/
aten.bzl
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
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,
)