1 /* Copyright 2019 The TensorFlow Authors. All Rights Reserved.
2
3 Licensed under the Apache License, Version 2.0 (the "License");
4 you may not use this file except in compliance with the License.
5 You may obtain a copy of the License at
6
7 http://www.apache.org/licenses/LICENSE-2.0
8
9 Unless required by applicable law or agreed to in writing, software
10 distributed under the License is distributed on an "AS IS" BASIS,
11 WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12 See the License for the specific language governing permissions and
13 limitations under the License.
14 ==============================================================================*/
15
16 #include "tensorflow/compiler/xla/service/gpu/nccl_all_gather_thunk.h"
17
18 #include <chrono> // NOLINT (required by TF interfaces)
19 #include <cstdlib>
20 #include <memory>
21 #include <string>
22 #include <utility>
23 #include <vector>
24
25 #include "absl/strings/str_format.h"
26 #include "tensorflow/compiler/xla/layout_util.h"
27 #include "tensorflow/compiler/xla/service/gpu/ir_emission_utils.h"
28 #include "tensorflow/compiler/xla/service/hlo_casting_utils.h"
29 #include "tensorflow/compiler/xla/service/hlo_instructions.h"
30 #include "tensorflow/compiler/xla/util.h"
31
32 #if XLA_ENABLE_XCCL
33 #include "tensorflow/stream_executor/gpu/gpu_stream.h"
34 #endif
35
36 namespace xla {
37 namespace gpu {
38
GetNcclAllGatherConfig(mlir::lmhlo::AllGatherOp op)39 /*static*/ NcclAllGatherConfig NcclAllGatherThunk::GetNcclAllGatherConfig(
40 mlir::lmhlo::AllGatherOp op) {
41 NcclAllGatherConfig config;
42 config.config =
43 GetNcclCollectiveConfigForMlir(op, op.getUseGlobalDeviceIds());
44 return config;
45 }
46
CanImplement(mlir::lmhlo::AllGatherOp op)47 /*static*/ bool NcclAllGatherThunk::CanImplement(mlir::lmhlo::AllGatherOp op) {
48 return absl::c_all_of(op.getInputs(), [&](mlir::Value operand) {
49 Shape shape = GetShape(operand);
50 return LayoutUtil::IsDenseArray(shape) &&
51 IsTypeSupportedByNccl(shape.element_type()) &&
52 LayoutUtil::MinorToMajor(shape).back() == op.getAllGatherDimension();
53 });
54 }
55
NcclAllGatherThunk(ThunkInfo thunk_info,mlir::lmhlo::AllGatherOp op,std::vector<NcclAllGatherThunk::Buffer> buffers)56 NcclAllGatherThunk::NcclAllGatherThunk(
57 ThunkInfo thunk_info, mlir::lmhlo::AllGatherOp op,
58 std::vector<NcclAllGatherThunk::Buffer> buffers)
59 : NcclCollectiveThunk(Thunk::kNcclAllGather, thunk_info),
60 config_(GetNcclAllGatherConfig(op)),
61 buffers_(std::move(buffers)) {
62 CHECK_EQ(config_.config.operand_count, buffers_.size());
63 }
64
RunAllGather(std::vector<DeviceBufferPair> & buffers,se::Stream & stream,ncclComm_t comm)65 Status RunAllGather(std::vector<DeviceBufferPair>& buffers, se::Stream& stream,
66 ncclComm_t comm) {
67 #if XLA_ENABLE_XCCL
68 int device_ordinal = stream.parent()->device_ordinal();
69 VLOG(3) << "Performing all-gather from device ordinal: " << device_ordinal;
70
71 se::gpu::GpuStreamHandle gpu_stream = se::gpu::AsGpuStreamValue(&stream);
72
73 XLA_CUDA_RETURN_IF_ERROR(ncclGroupStart());
74 for (size_t i = 0; i < buffers.size(); ++i) {
75 DeviceBufferPair& buffer = buffers[i];
76 const void* send_buffer = buffer.source_buffer.opaque();
77 void* recv_buffer = buffer.destination_buffer.opaque();
78
79 PrimitiveType element_type = buffer.element_type;
80 TF_ASSIGN_OR_RETURN(auto dtype_and_multiplier,
81 ToNcclDataTypeAndCountMultiplier(element_type));
82 ncclDataType_t dtype = dtype_and_multiplier.first;
83 int element_count = buffer.element_count * dtype_and_multiplier.second;
84
85 VLOG(3) << absl::StreamFormat(
86 "Calling ncclAllGather(send_buffer=%p, recv_buffer=%p, sendcount=%d, "
87 "comm=%p, stream=%p)",
88 send_buffer, recv_buffer, element_count, static_cast<const void*>(comm),
89 gpu_stream);
90
91 XLA_CUDA_RETURN_IF_ERROR(ncclAllGather(
92 send_buffer, recv_buffer, element_count, dtype, comm, gpu_stream));
93 }
94 XLA_CUDA_RETURN_IF_ERROR(ncclGroupEnd());
95
96 VLOG(3) << "Done performing all-gather for ordinal: " << device_ordinal;
97 return OkStatus();
98 #else // XLA_ENABLE_XCCL
99 return Unimplemented(
100 "NCCL support is not available: this binary was not built with a CUDA "
101 "compiler, which is necessary to build the NCCL source library.");
102 #endif // XLA_ENABLE_XCCL
103 }
104
RunNcclCollective(const ExecuteParams & params,ncclComm_t comm)105 Status NcclAllGatherThunk::RunNcclCollective(const ExecuteParams& params,
106 ncclComm_t comm) {
107 se::Stream& stream = *params.stream;
108 TF_ASSIGN_OR_RETURN(
109 std::vector<DeviceBufferPair> device_buffers,
110 ConvertToDeviceBuffers(params, buffers_,
111 config_.config.operand_element_type));
112 return RunAllGather(device_buffers, stream, comm);
113 }
114
115 } // namespace gpu
116 } // namespace xla
117