1""" 2 Macros for selecting with / without various GPU libraries. Most of these are meant to be used 3 directly by tensorflow in place of their build's own configure.py + bazel-gen system. 4""" 5 6load("@bazel_skylib//lib:selects.bzl", "selects") 7 8def if_cuda(if_true, if_false = []): 9 """Helper for selecting based on the whether CUDA is configured. """ 10 return selects.with_or({ 11 "@//tools/config:cuda_enabled_and_capable": if_true, 12 "//conditions:default": if_false, 13 }) 14 15def if_tensorrt(if_true, if_false = []): 16 """Helper for selecting based on the whether TensorRT is configured. """ 17 return select({ 18 "//conditions:default": if_false, 19 }) 20 21def if_rocm(if_true, if_false = []): 22 """Helper for selecting based on the whether ROCM is configured. """ 23 return select({ 24 "//conditions:default": if_false, 25 }) 26 27def if_sycl(if_true, if_false = []): 28 """Helper for selecting based on the whether SYCL/ComputeCPP is configured.""" 29 30 # NOTE: Tensorflow expects some stange behavior (see their if_sycl) if we 31 # actually plan on supporting this at some point. 32 return select({ 33 "//conditions:default": if_false, 34 }) 35 36def if_ccpp(if_true, if_false = []): 37 """Helper for selecting based on the whether ComputeCPP is configured. """ 38 return select({ 39 "//conditions:default": if_false, 40 }) 41 42def cuda_default_copts(): 43 return if_cuda(["-DGOOGLE_CUDA=1"]) 44 45def cuda_default_features(): 46 return if_cuda(["-per_object_debug_info", "-use_header_modules", "cuda_clang"]) 47 48def rocm_default_copts(): 49 return if_rocm(["-x", "rocm"]) 50 51def rocm_copts(opts = []): 52 return rocm_default_copts() + if_rocm(opts) 53 54def cuda_is_configured(): 55 # FIXME(dcollins): currently only used by tensorflow's xla stuff, which we aren't building. However bazel 56 # query hits it so this needs to be defined. Because bazel doesn't actually resolve config at macro expansion 57 # time, `select` can't be used here (since xla expects lists of strings and not lists of select objects). 58 # Instead, the xla build rules must be rewritten to use `if_cuda_is_configured` 59 return False 60 61def if_cuda_is_configured(x): 62 return if_cuda(x, []) 63 64def if_rocm_is_configured(x): 65 return if_rocm(x, []) 66