• Home
  • Line#
  • Scopes#
  • Navigate#
  • Raw
  • Download
1# Macros for building ROCm code.
2def if_rocm(if_true, if_false = []):
3    """Shorthand for select()'ing on whether we're building with ROCm.
4
5    Returns a select statement which evaluates to if_true if we're building
6    with ROCm enabled.  Otherwise, the select statement evaluates to if_false.
7
8    """
9    return select({
10        "@local_config_rocm//rocm:using_hipcc": if_true,
11        "//conditions:default": if_false
12    })
13
14
15def rocm_default_copts():
16    """Default options for all ROCm compilations."""
17    return if_rocm(["-x", "rocm"] + %{rocm_extra_copts})
18
19def rocm_copts(opts = []):
20    """Gets the appropriate set of copts for (maybe) ROCm compilation.
21
22      If we're doing ROCm compilation, returns copts for our particular ROCm
23      compiler.  If we're not doing ROCm compilation, returns an empty list.
24
25      """
26    return rocm_default_copts() + select({
27        "//conditions:default": [],
28        "@local_config_rocm//rocm:using_hipcc": ([
29            "",
30        ]),
31    }) + if_rocm_is_configured(opts)
32
33def rocm_gpu_architectures():
34    """Returns a list of supported GPU architectures."""
35    return %{rocm_gpu_architectures}
36
37def rocm_version_number():
38    """Returns a list of supported GPU architectures."""
39    return %{rocm_version_number}
40
41def if_rocm_is_configured(x):
42    """Tests if the ROCm was enabled during the configure process.
43
44    Unlike if_rocm(), this does not require that we are building with
45    --config=rocm. Used to allow non-ROCm code to depend on ROCm libraries.
46    """
47    if %{rocm_is_configured}:
48      return select({"//conditions:default": x})
49    return select({"//conditions:default": []})
50
51def rocm_library(copts = [], **kwargs):
52    """Wrapper over cc_library which adds default ROCm options."""
53    native.cc_library(copts = rocm_default_copts() + copts, **kwargs)
54