forked from pytorch/pytorch
-
Notifications
You must be signed in to change notification settings - Fork 0
/
aten.bzl
96 lines (88 loc) · 2.94 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
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
load("@bazel_skylib//lib:paths.bzl", "paths")
load("@rules_cc//cc:defs.bzl", "cc_library")
CPU_CAPABILITY_NAMES = ["DEFAULT", "AVX2"]
CAPABILITY_COMPILER_FLAGS = {
"AVX2": ["-mavx2", "-mfma"],
"DEFAULT": [],
}
PREFIX = "aten/src/ATen/native/"
EXTRA_PREFIX = "aten/src/ATen/"
def intern_build_aten_ops(copts, deps, extra_impls):
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)
for impl in extra_impls:
name = impl.replace(EXTRA_PREFIX, "")
out = EXTRA_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,
)
def generate_aten_impl(ctx):
# Declare the entire ATen/ops/ directory as an output
ops_dir = ctx.actions.declare_directory("aten/src/ATen/ops")
outputs = [ops_dir] + ctx.outputs.outs
install_dir = paths.dirname(ops_dir.path)
tool_inputs, tool_inputs_manifest = ctx.resolve_tools(tools = [ctx.attr.generator])
ctx.actions.run_shell(
outputs = outputs,
inputs = ctx.files.srcs,
command = ctx.executable.generator.path + " $@",
arguments = [
"--source-path",
"aten/src/ATen",
"--per-operator-headers",
"--install_dir",
install_dir,
],
tools = tool_inputs,
input_manifests = tool_inputs_manifest,
use_default_shell_env = True,
mnemonic = "GenerateAten",
)
return [DefaultInfo(files = depset(outputs))]
generate_aten = rule(
implementation = generate_aten_impl,
attrs = {
"generator": attr.label(
executable = True,
allow_files = True,
mandatory = True,
cfg = "exec",
),
"outs": attr.output_list(),
"srcs": attr.label_list(allow_files = True),
},
)