• Home
  • Line#
  • Scopes#
  • Navigate#
  • Raw
  • Download
1# Macros for building CUDA code.
2def if_cuda(if_true, if_false = []):
3    """Shorthand for select()'ing on whether we're building with CUDA.
4
5    Returns a select statement which evaluates to if_true if we're building
6    with CUDA enabled.  Otherwise, the select statement evaluates to if_false.
7
8    """
9    return select({
10        "@local_config_cuda//cuda:using_nvcc": if_true,
11        "@local_config_cuda//cuda:using_clang": if_true,
12        "//conditions:default": if_false,
13    })
14
15def if_cuda_clang(if_true, if_false = []):
16   """Shorthand for select()'ing on wheteher we're building with cuda-clang.
17
18    Returns a select statement which evaluates to if_true if we're building
19    with cuda-clang.  Otherwise, the select statement evaluates to if_false.
20
21   """
22   return select({
23       "@local_config_cuda//cuda:using_clang": if_true,
24       "//conditions:default": if_false
25   })
26
27def if_cuda_clang_opt(if_true, if_false = []):
28   """Shorthand for select()'ing on wheteher we're building with cuda-clang
29   in opt mode.
30
31    Returns a select statement which evaluates to if_true if we're building
32    with cuda-clang in opt mode. Otherwise, the select statement evaluates to
33    if_false.
34
35   """
36   return select({
37       "@local_config_cuda//cuda:using_clang_opt": if_true,
38       "//conditions:default": if_false
39   })
40
41def cuda_default_copts():
42    """Default options for all CUDA compilations."""
43    return if_cuda([
44        "-x", "cuda",
45        "-DGOOGLE_CUDA=1",
46        "-Xcuda-fatbinary=--compress-all",
47        "--no-cuda-include-ptx=all"
48    ] + %{cuda_extra_copts}) + if_cuda_clang_opt(
49        # Some important CUDA optimizations are only enabled at O3.
50        ["-O3"]
51    )
52
53def cuda_gpu_architectures():
54    """Returns a list of supported GPU architectures."""
55    return %{cuda_gpu_architectures}
56
57def if_cuda_is_configured(x):
58    """Tests if the CUDA was enabled during the configure process.
59
60    Unlike if_cuda(), this does not require that we are building with
61    --config=cuda. Used to allow non-CUDA code to depend on CUDA libraries.
62    """
63    if %{cuda_is_configured}:
64      return select({"//conditions:default": x})
65    return select({"//conditions:default": []})
66
67def cuda_header_library(
68        name,
69        hdrs,
70        include_prefix = None,
71        strip_include_prefix = None,
72        deps = [],
73        **kwargs):
74    """Generates a cc_library containing both virtual and system include paths.
75
76    Generates both a header-only target with virtual includes plus the full
77    target without virtual includes. This works around the fact that bazel can't
78    mix 'includes' and 'include_prefix' in the same target."""
79
80    native.cc_library(
81        name = name + "_virtual",
82        hdrs = hdrs,
83        include_prefix = include_prefix,
84        strip_include_prefix = strip_include_prefix,
85        deps = deps,
86        visibility = ["//visibility:private"],
87    )
88
89    native.cc_library(
90        name = name,
91        textual_hdrs = hdrs,
92        deps = deps + [":%s_virtual" % name],
93        **kwargs
94    )
95
96def cuda_library(copts = [], **kwargs):
97    """Wrapper over cc_library which adds default CUDA options."""
98    native.cc_library(copts = cuda_default_copts() + copts, **kwargs)
99