github.com/johnnyeven/libtools@v0.0.0-20191126065708-61829c1adf46/third_party/nccl/archive.BUILD (about)

     1  # NVIDIA NCCL 2
     2  # A package of optimized primitives for collective multi-GPU communication.
     3  
     4  licenses(["notice"])
     5  
     6  exports_files(["LICENSE.txt"])
     7  
     8  load(
     9      "@local_config_nccl//:build_defs.bzl",
    10      "cuda_rdc_library",
    11      "gen_device_srcs",
    12  )
    13  
    14  cc_library(
    15      name = "src_hdrs",
    16      hdrs = [
    17          "src/collectives.h",
    18          "src/collectives/collectives.h",
    19          "src/nccl.h",
    20      ],
    21      strip_include_prefix = "src",
    22  )
    23  
    24  cc_library(
    25      name = "include_hdrs",
    26      hdrs = glob(["src/include/*.h"]),
    27      strip_include_prefix = "src/include",
    28      deps = ["@local_config_cuda//cuda:cuda_headers"],
    29  )
    30  
    31  cc_library(
    32      name = "device_hdrs",
    33      hdrs = glob(["src/collectives/device/*.h"]),
    34      strip_include_prefix = "src/collectives/device",
    35  )
    36  
    37  # NCCL compiles the same source files with different NCCL_OP/NCCL_TYPE defines.
    38  # RDC compilation requires that each compiled module has a unique ID. Clang
    39  # derives the module ID from the path only so we need to copy the files to get
    40  # different IDs for different parts of compilation. NVCC does not have that
    41  # problem because it generates IDs based on preprocessed content.
    42  gen_device_srcs(
    43      name = "device_srcs",
    44      srcs = [
    45          "src/collectives/device/all_gather.cu.cc",
    46          "src/collectives/device/all_reduce.cu.cc",
    47          "src/collectives/device/broadcast.cu.cc",
    48          "src/collectives/device/reduce.cu.cc",
    49          "src/collectives/device/reduce_scatter.cu.cc",
    50      ],
    51  )
    52  
    53  cuda_rdc_library(
    54      name = "device",
    55      srcs = [
    56          "src/collectives/device/functions.cu.cc",
    57          ":device_srcs",
    58      ] + glob([
    59          # Required for header inclusion checking, see below for details.
    60          "src/collectives/device/*.h",
    61          "src/nccl.h",
    62      ]),
    63      deps = [
    64          ":device_hdrs",
    65          ":include_hdrs",
    66          ":src_hdrs",
    67          "@local_config_cuda//cuda:cuda_headers",
    68      ],
    69  )
    70  
    71  # Primary NCCL target.
    72  cc_library(
    73      name = "nccl",
    74      srcs = glob(
    75          include = ["src/**/*.cc"],
    76          # Exclude device-library code.
    77          exclude = ["src/collectives/device/**"],
    78      ) + [
    79          # Required for header inclusion checking (see
    80          # http://docs.bazel.build/versions/master/be/c-cpp.html#hdrs).
    81          # Files in src/ which #include "nccl.h" load it from there rather than
    82          # from the virtual includes directory.
    83          "src/collectives.h",
    84          "src/collectives/collectives.h",
    85          "src/nccl.h",
    86      ],
    87      hdrs = ["src/nccl.h"],
    88      include_prefix = "third_party/nccl",
    89      strip_include_prefix = "src",
    90      visibility = ["//visibility:public"],
    91      deps = [
    92          ":device",
    93          ":include_hdrs",
    94          ":src_hdrs",
    95          "@local_config_cuda//cuda:cudart_static",
    96      ],
    97  )