-
Notifications
You must be signed in to change notification settings - Fork 74k
/
build_defs.bzl
34 lines (27 loc) · 1.18 KB
/
build_defs.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
"""Starlark macros for oneDNN.
if_mkldnn_openmp checks if we are building x86 backend with OpenMP.
if_mkldnn_aarch64_acl checks if we are building with Arm Compute Library.
if_mkldnn_aarch64_acl_openmp checks if we are building ACL with OpenMP.
"""
def if_mkldnn_openmp(if_true, if_false = []):
"""Returns `if_true` if OpenMP is used with oneDNN.
Shorthand for select()'ing on whether we're building with
oneDNN open source library only with openmp
Returns a select statement which evaluates to if_true if we're building
with oneDNN open source library only with OpenMP. Otherwise, the
select statement evaluates to if_false.
"""
return select({
"@local_tsl//third_party/mkl_dnn:build_with_mkldnn_openmp": if_true,
"//conditions:default": if_false,
})
def if_mkldnn_aarch64_acl(if_true, if_false = []):
return select({
"@local_xla//xla/tsl/mkl:build_with_mkl_aarch64": if_true,
"//conditions:default": if_false,
})
def if_mkldnn_aarch64_acl_openmp(if_true, if_false = []):
return select({
"@local_tsl//third_party/mkl_dnn:build_with_mkl_aarch64_openmp": if_true,
"//conditions:default": if_false,
})