broadcast_op_handle.cc 6.0 KB
Newer Older
C
chengduoZH 已提交
1 2 3 4 5 6 7 8 9 10 11 12 13 14
//   Copyright (c) 2018 PaddlePaddle Authors. All Rights Reserved.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

C
chengduoZH 已提交
15
#include "paddle/fluid/framework/details/broadcast_op_handle.h"
Y
Yu Yang 已提交
16 17
#include "paddle/fluid/framework/details/container_cast.h"
#include "paddle/fluid/framework/details/variable_visitor.h"
Y
Yancey1989 已提交
18
#include "paddle/fluid/platform/profiler.h"
C
chengduoZH 已提交
19 20 21 22 23

namespace paddle {
namespace framework {
namespace details {

C
chengduoZH 已提交
24
void BroadcastOpHandle::RunImpl() {
25
  platform::RecordEvent record_event(Name());
Y
Yancey1989 已提交
26

C
chengduoZH 已提交
27
  if (places_.size() == 1) return;
Y
Yu Yang 已提交
28

C
chengduoZH 已提交
29
  // The input and output may have dummy vars.
C
chengduo 已提交
30
  auto in_var_handles = DynamicCast<VarHandle>(inputs_);
Y
Yu Yang 已提交
31
  auto out_var_handles = DynamicCast<VarHandle>(outputs_);
C
chengduoZH 已提交
32

C
chengduo 已提交
33 34
  PADDLE_ENFORCE_EQ(in_var_handles.size(), 1UL,
                    "The number of input should be one.");
C
chengduoZH 已提交
35
  PADDLE_ENFORCE_EQ(
36
      out_var_handles.size(), places_.size(),
C
chengduoZH 已提交
37
      "The number of output should equal to the number of places.");
C
chengduoZH 已提交
38

C
chengduo 已提交
39 40
  VarHandle *in_var_handle = in_var_handles[0];

C
chengduoZH 已提交
41
  WaitInputVarGenerated();
C
chengduoZH 已提交
42

43
  BroadcastOneVar(*in_var_handle, out_var_handles, local_exec_scopes_);
44 45 46 47 48
}

void BroadcastOpHandle::BroadcastOneVar(
    const VarHandle &in_var_handle,
    const std::vector<VarHandle *> &out_var_handles,
49
    const std::vector<Scope *> &var_scopes) {
C
chengduoZH 已提交
50
  auto *in_var =
G
gongweibao 已提交
51
      var_scopes.at(in_var_handle.scope_idx())->FindVar(in_var_handle.name());
Y
Yu Yang 已提交
52 53
  PADDLE_ENFORCE_NOT_NULL(in_var);
  Tensor &in_tensor = VariableVisitor::GetMutableTensor(in_var);
54
  if (UNLIKELY(!in_tensor.IsInitialized())) {
G
gongweibao 已提交
55
    VLOG(3) << "in var " << in_var_handle.name() << "not inited, return!";
56 57
    return;
  }
C
chengduoZH 已提交
58

59
  InitOutputValue(in_var_handle, out_var_handles);
C
chengduoZH 已提交
60

C
chengduoZH 已提交
61
  if (platform::is_cpu_place(in_tensor.place())) {
C
chengduoZH 已提交
62
    for (auto *out_var_handle : out_var_handles) {
63
      if (out_var_handle->IsTheSameVar(in_var_handle)) {
C
chengduoZH 已提交
64 65
        continue;
      }
G
gongweibao 已提交
66 67 68
      auto &out_p = out_var_handle->place();
      auto *out_var = var_scopes.at(out_var_handle->scope_idx())
                          ->FindVar(out_var_handle->name());
C
chengduoZH 已提交
69

C
chengduoZH 已提交
70
      RunAndRecordEvent(out_p, [in_tensor, out_var] {
C
chengduoZH 已提交
71
        paddle::framework::TensorCopy(
C
chengduoZH 已提交
72
            in_tensor, platform::CPUPlace(),
C
chengduoZH 已提交
73 74 75 76
            &VariableVisitor::GetMutableTensor(out_var));
      });
    }
  } else {
P
peizhilin 已提交
77
#if defined(PADDLE_WITH_CUDA) && !defined(_WIN32)
C
chengduoZH 已提交
78 79
    VarHandle *out_handle = nullptr;
    int root_id = boost::get<platform::CUDAPlace>(in_tensor.place()).device;
C
chengduoZH 已提交
80 81
    std::vector<std::function<void()>> broadcast_calls;

C
chengduoZH 已提交
82 83 84
    int type = platform::ToNCCLDataType(in_tensor.type());
    size_t numel = static_cast<size_t>(in_tensor.numel());

C
chengduoZH 已提交
85
    for (auto out_var_handle : out_var_handles) {
G
gongweibao 已提交
86 87
      Variable *out_var = var_scopes.at(out_var_handle->scope_idx())
                              ->FindVar(out_var_handle->name());
C
chengduoZH 已提交
88

C
chengduoZH 已提交
89
      int dst_id =
G
gongweibao 已提交
90
          boost::get<platform::CUDAPlace>(out_var_handle->place()).device;
C
chengduoZH 已提交
91

C
chengduoZH 已提交
92
      auto &nccl_ctx = nccl_ctxs_->at(dst_id);
C
chengduoZH 已提交
93 94

      void *send_recv_buffer = nullptr;
C
chengduoZH 已提交
95
      if (root_id == dst_id) {
C
chengduoZH 已提交
96 97 98
        send_recv_buffer = const_cast<void *>(in_tensor.data<void>());
        out_handle = out_var_handle;
      } else {
C
chengduoZH 已提交
99 100
        send_recv_buffer = VariableVisitor::GetMutableTensor(out_var)
                               .Resize(in_tensor.dims())
G
gongweibao 已提交
101
                               .mutable_data(out_var_handle->place());
C
chengduoZH 已提交
102 103
      }

C
chengduoZH 已提交
104 105 106 107 108 109
      broadcast_calls.emplace_back(
          [send_recv_buffer, numel, type, root_id, &nccl_ctx] {
            PADDLE_ENFORCE(platform::dynload::ncclBcast(
                send_recv_buffer, numel, static_cast<ncclDataType_t>(type),
                root_id, nccl_ctx.comm_, nccl_ctx.stream()));
          });
Y
Yu Yang 已提交
110 111
    }

112 113 114 115 116
    this->RunAndRecordEvent([&] {
      {
        platform::NCCLGroupGuard guard;
        for (auto &call : broadcast_calls) {
          call();
C
chengduoZH 已提交
117
        }
118
      }
C
chengduoZH 已提交
119

120
      if (!out_handle->IsTheSameVar(in_var_handle)) {
G
gongweibao 已提交
121 122
        auto out_var = var_scopes.at(in_var_handle.scope_idx())
                           ->FindVar(out_var_handles[0]->name());
123
        paddle::framework::TensorCopy(
G
gongweibao 已提交
124 125
            in_tensor, in_var_handle.place(),
            *(dev_ctxes_.at(in_var_handle.place())),
126 127 128
            &VariableVisitor::GetMutableTensor(out_var));
      }
    });
C
chengduo 已提交
129 130 131
    for (auto &p : places_) {
      nccl_ctxs_->DevCtx(p)->Wait();
    }
C
chengduoZH 已提交
132
#else
C
chengduoZH 已提交
133
    PADDLE_THROW("CUDA is not enabled.");
C
chengduoZH 已提交
134
#endif
C
chengduoZH 已提交
135 136 137
  }
}

138 139 140
void BroadcastOpHandle::InitOutputValue(
    const VarHandle &in_var_handle,
    const std::vector<VarHandle *> &out_var_handles) const {
141
  auto &var_scopes = local_exec_scopes_;
142
  auto *in_var =
G
gongweibao 已提交
143
      var_scopes.at(in_var_handle.scope_idx())->FindVar(in_var_handle.name());
144 145 146 147 148 149 150 151 152

  Tensor &in_tensor = VariableVisitor::GetMutableTensor(in_var);

  // NOTE: The tensors' Place of input and output must be all on GPU or all on
  // CPU.
  for (auto *out_var_handle : out_var_handles) {
    if (out_var_handle->IsTheSameVar(in_var_handle)) {
      continue;
    }
G
gongweibao 已提交
153 154 155
    auto t_out_p = out_var_handle->place();
    auto *out_var = var_scopes.at(out_var_handle->scope_idx())
                        ->FindVar(out_var_handle->name());
156 157 158 159 160 161 162 163 164 165 166 167 168
    PADDLE_ENFORCE_NOT_NULL(out_var);
    if (is_gpu_place(in_tensor.place())) {
      PADDLE_ENFORCE(platform::is_gpu_place(t_out_p),
                     "Places of input and output must be all on GPU.");
    } else {
      t_out_p = platform::CPUPlace();
    }
    VariableVisitor::ShareDimsAndLoD(*in_var, out_var);
    VariableVisitor::GetMutableTensor(out_var).mutable_data(t_out_p,
                                                            in_tensor.type());
  }
}

C
chengduoZH 已提交
169
std::string BroadcastOpHandle::Name() const { return "broadcast"; }
C
chengduoZH 已提交
170 171 172
}  // namespace details
}  // namespace framework
}  // namespace paddle