device_context.cc 15.0 KB
Newer Older
1
/* Copyright (c) 2016 PaddlePaddle Authors. All Rights Reserved.
2 3
Copyright (c) 2022 NVIDIA Corporation. All rights reserved.

Q
qijun 已提交
4 5 6 7
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
8

Q
qijun 已提交
9 10 11 12 13
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. */
14

Y
Yi Wang 已提交
15
#include "paddle/fluid/platform/device_context.h"
16

W
Wilber 已提交
17
#include <functional>
18
#include <memory>
19
#include <set>
20

21 22 23 24
#include "glog/logging.h"
#include "paddle/fluid/framework/generator.h"
#include "paddle/fluid/memory/allocation/allocator_facade.h"
#include "paddle/fluid/platform/device/device_wrapper.h"
W
Wilber 已提交
25
#include "paddle/fluid/platform/place.h"
26 27
#include "paddle/fluid/platform/profiler.h"
#include "paddle/fluid/platform/profiler/event_tracing.h"
28 29
#include "paddle/phi/backends/gpu/gpu_context.h"
#include "paddle/phi/core/allocator.h"
30
#include "paddle/phi/core/expect.h"
31

32
#if defined(PADDLE_WITH_CUDA) || defined(PADDLE_WITH_HIP)
33
#include "paddle/fluid/memory/allocation/cuda_device_context_allocator.h"
S
sneaxiy 已提交
34
#include "paddle/fluid/platform/cuda_device_guard.h"
35
#endif
36

F
fwenguang 已提交
37 38 39 40
#ifdef PADDLE_WITH_MLU
#include "paddle/fluid/platform/device/mlu/device_context.h"
#include "paddle/fluid/platform/device/mlu/device_context_allocator.h"
#endif
41

Q
qijun 已提交
42 43 44
namespace paddle {
namespace platform {

45
#if defined(PADDLE_WITH_CUDA) || defined(PADDLE_WITH_HIP)
46 47 48
bool allow_tf32_cublas = true;
void SetAllowTF32Cublas(bool active) { allow_tf32_cublas = active; }
bool AllowTF32Cublas() { return allow_tf32_cublas; }
A
AshburnLee 已提交
49 50 51 52

bool allow_tf32_cudnn = true;
void SetAllowTF32Cudnn(bool active) { allow_tf32_cudnn = active; }
bool AllowTF32Cudnn() { return allow_tf32_cudnn; }
53 54
#endif  // PADDLE_WITH_CUDA

55 56 57 58 59 60 61
DeviceType Place2DeviceType(const platform::Place& place) {
  if (platform::is_cpu_place(place)) {
    return platform::DeviceType::CPU;
  } else if (platform::is_gpu_place(place)) {
    return platform::DeviceType::CUDA;
  } else if (platform::is_xpu_place(place)) {
    return platform::DeviceType::XPU;
62 63
  } else if (platform::is_ipu_place(place)) {
    return platform::DeviceType::IPU;
64 65
  } else if (platform::is_npu_place(place)) {
    return platform::DeviceType::NPU;
F
fwenguang 已提交
66 67
  } else if (platform::is_mlu_place(place)) {
    return platform::DeviceType::MLU;
68 69
  } else if (platform::is_custom_place(place)) {
    return platform::DeviceType::CUSTOM_DEVICE;
70 71 72 73 74 75
  } else {
    PADDLE_THROW(platform::errors::Unavailable(
        "Unsupported place %s to convert into platform::DeviceType.", place));
  }
}

76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99
static DeviceContextPool* pool = nullptr;

DeviceContextPool& DeviceContextPool::Instance() {
  PADDLE_ENFORCE_NOT_NULL(pool,
                          phi::errors::PreconditionNotMet(
                              "Need to Create DeviceContextPool firstly!"));
  return *pool;
}

/*! \brief  Create should only called by Init function */
DeviceContextPool& DeviceContextPool::Init(
    const std::vector<platform::Place>& places) {
  if (pool == nullptr) {
    pool = new DeviceContextPool(places);
  }
  return *pool;
}

bool DeviceContextPool::IsInitialized() { return pool != nullptr; }

void DeviceContextPool::SetPool(DeviceContextPool* dev_pool) {
  pool = dev_pool;
}

100 101 102
thread_local const std::map<Place,
                            std::shared_future<std::unique_ptr<DeviceContext>>>*
    DeviceContextPool::external_device_contexts_ = nullptr;
D
dzhwinter 已提交
103

Y
Yu Yang 已提交
104
platform::DeviceContext* DeviceContextPool::Get(const platform::Place& place) {
105
  VLOG(6) << "DeviceContextPool Get: " << place;
106 107 108 109 110 111 112 113 114 115
  const std::map<Place, std::shared_future<std::unique_ptr<DeviceContext>>>*
      ptr;
  if (external_device_contexts_ && external_device_contexts_->count(place)) {
    ptr = external_device_contexts_;
  } else {
    ptr = &device_contexts_;
  }

  auto it = ptr->find(place);
  if (it == ptr->end()) {
G
GaoWei8 已提交
116 117
    PADDLE_THROW(platform::errors::Unimplemented(
        "Place %s is not supported. Please check that your paddle compiles "
F
fwenguang 已提交
118 119
        "with WITH_GPU, WITH_XPU, WITH_IPU, WITH_MLU or WITH_ASCEND_CL option "
        "or check "
J
jianghaicheng 已提交
120 121
        "that your train process set the correct device id if you use "
        "Executor.",
G
GaoWei8 已提交
122
        place));
D
dzhwinter 已提交
123
  }
124
  return it->second.get().get();
D
dzhwinter 已提交
125 126
}

127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147
size_t DeviceContextPool::size() const {
  if (external_device_contexts_) {
    return external_device_contexts_->size();
  }
  return device_contexts_.size();
}

const std::map<Place, std::shared_future<std::unique_ptr<DeviceContext>>>&
DeviceContextPool::device_contexts() const {
  if (external_device_contexts_) {
    return *external_device_contexts_;
  }
  return device_contexts_;
}

void DeviceContextPool::SetDeviceContexts(
    const std::map<Place, std::shared_future<std::unique_ptr<DeviceContext>>>*
        dev_ctxs) {
  external_device_contexts_ = dev_ctxs;
}

W
Wilber 已提交
148
template <typename DevCtx>
149 150 151
std::unique_ptr<DeviceContext> CreateDeviceContext(
    const platform::Place& p,
    bool disable_setting_default_stream_for_allocator = false) {
152
  using PtrType = std::unique_ptr<DeviceContext>;
153 154
  auto* dev_ctx = new DevCtx(p);
  if (is_gpu_place(p)) {
155
#if defined(PADDLE_WITH_CUDA) || defined(PADDLE_WITH_HIP)
L
Leo Chen 已提交
156
    auto* cuda_ctx = dynamic_cast<phi::GPUContext*>(dev_ctx);
157 158 159
    PADDLE_ENFORCE_NOT_NULL(
        cuda_ctx,
        platform::errors::InvalidArgument(
L
Leo Chen 已提交
160
            "Failed to dynamic_cast dev_ctx into phi::GPUContext."));
161 162 163 164 165

    auto& instance = memory::allocation::AllocatorFacade::Instance();
    if (!disable_setting_default_stream_for_allocator) {
      instance.SetDefaultStream(CUDAPlace(p.GetDeviceId()), cuda_ctx->stream());
    }
166
    dev_ctx->SetAllocator(instance.GetAllocator(p, cuda_ctx->stream()).get());
167 168 169 170 171 172
    dev_ctx->SetPinnedAllocator(
        instance.GetAllocator(paddle::platform::CUDAPinnedPlace()).get());

    cuda_ctx->PartialInitWithAllocator();
    dev_ctx->SetGenerator(
        framework::DefaultCUDAGenerator(p.GetDeviceId()).get());
Q
QingshuChen 已提交
173 174 175 176 177 178 179
#endif
  } else if (is_xpu_place(p)) {
#if defined(PADDLE_WITH_XPU)
    dev_ctx->SetAllocator(
        memory::allocation::AllocatorFacade::Instance().GetAllocator(p).get());
    dev_ctx->SetGenerator(
        framework::DefaultXPUGenerator(p.GetDeviceId()).get());
180
#endif
181 182 183 184 185 186 187 188 189 190 191 192
  } else {
    dev_ctx->SetAllocator(
        memory::allocation::AllocatorFacade::Instance().GetAllocator(p).get());
    dev_ctx->SetGenerator(framework::DefaultCPUGenerator().get());
  }
  dev_ctx->SetHostGenerator(framework::DefaultCPUGenerator().get());
  dev_ctx->SetHostAllocator(memory::allocation::AllocatorFacade::Instance()
                                .GetAllocator(platform::CPUPlace())
                                .get());
  dev_ctx->SetZeroAllocator(memory::allocation::AllocatorFacade::Instance()
                                .GetZeroAllocator(p)
                                .get());
193 194 195
  dev_ctx->SetHostZeroAllocator(memory::allocation::AllocatorFacade::Instance()
                                    .GetZeroAllocator(platform::CPUPlace())
                                    .get());
196
  return PtrType(dev_ctx);
C
chengduozh 已提交
197 198
}

199 200 201 202
template <typename DevCtx>
inline void EmplaceDeviceContext(
    std::map<Place, std::shared_future<std::unique_ptr<DeviceContext>>>*
        place_to_device_context,
203 204
    platform::Place place,
    bool disable_setting_default_stream_for_allocator) {
205 206
  // lazy evaluation. i.e., only create device context at first `Get`
  place_to_device_context->emplace(
207 208 209 210 211
      place,
      std::async(std::launch::deferred,
                 CreateDeviceContext<DevCtx>,
                 place,
                 disable_setting_default_stream_for_allocator));
212 213 214 215 216 217 218
}

void EmplaceDeviceContexts(
    std::map<Place, std::shared_future<std::unique_ptr<DeviceContext>>>*
        place_to_device_context,
    const std::vector<platform::Place>& places,
    bool disable_setting_default_stream_for_allocator) {
G
GaoWei8 已提交
219
  PADDLE_ENFORCE_GT(
220 221
      places.size(),
      0,
G
GaoWei8 已提交
222 223 224
      platform::errors::InvalidArgument("The number of platform places should "
                                        "be larger than 0. But received %d.",
                                        places.size()));
225

226
  std::set<Place> set;
Y
Yu Yang 已提交
227 228 229
  for (auto& p : places) {
    set.insert(p);
  }
230

Y
Yu Yang 已提交
231 232
  for (auto& p : set) {
    if (platform::is_cpu_place(p)) {
233
#ifdef PADDLE_WITH_MKLDNN
234
      EmplaceDeviceContext<phi::OneDNNContext>(
235 236
          place_to_device_context,
          p,
237
          disable_setting_default_stream_for_allocator);
238
#else
L
Leo Chen 已提交
239
      EmplaceDeviceContext<phi::CPUContext>(
240 241
          place_to_device_context,
          p,
242
          disable_setting_default_stream_for_allocator);
243
#endif
Y
Yu Yang 已提交
244
    } else if (platform::is_gpu_place(p)) {
245
#if defined(PADDLE_WITH_CUDA) || defined(PADDLE_WITH_HIP)
L
Leo Chen 已提交
246
      EmplaceDeviceContext<phi::GPUContext>(
247 248
          place_to_device_context,
          p,
249
          disable_setting_default_stream_for_allocator);
D
dzhwinter 已提交
250
#else
G
GaoWei8 已提交
251 252 253
      PADDLE_THROW(
          platform::errors::Unimplemented("CUDAPlace is not supported. Please "
                                          "re-compile with WITH_GPU option."));
C
chengduoZH 已提交
254 255
#endif
    } else if (platform::is_cuda_pinned_place(p)) {
256
#if defined(PADDLE_WITH_CUDA) || defined(PADDLE_WITH_HIP)
257
      EmplaceDeviceContext<CUDAPinnedDeviceContext>(
258 259
          place_to_device_context,
          p,
260
          disable_setting_default_stream_for_allocator);
C
chengduoZH 已提交
261
#else
G
GaoWei8 已提交
262
      PADDLE_THROW(platform::errors::Unimplemented(
G
GaoWei8 已提交
263 264
          "CUDAPlace is not supported. Please re-compile with WITH_GPU "
          "option."));
265 266 267
#endif
    } else if (platform::is_xpu_place(p)) {
#ifdef PADDLE_WITH_XPU
268
      EmplaceDeviceContext<XPUDeviceContext>(
269 270
          place_to_device_context,
          p,
271
          disable_setting_default_stream_for_allocator);
272 273 274 275
#else
      PADDLE_THROW(
          platform::errors::Unimplemented("XPUPlace is not supported. Please "
                                          "re-compile with WITH_XPU option."));
F
fwenguang 已提交
276 277 278
#endif
    } else if (platform::is_mlu_place(p)) {
#ifdef PADDLE_WITH_MLU
279
      EmplaceDeviceContext<MLUDeviceContext>(
280 281
          place_to_device_context,
          p,
282
          disable_setting_default_stream_for_allocator);
F
fwenguang 已提交
283 284 285 286
#else
      PADDLE_THROW(
          platform::errors::Unimplemented("MLUPlace is not supported. Please "
                                          "re-compile with WITH_MLU option."));
J
jianghaicheng 已提交
287 288 289
#endif
    } else if (platform::is_ipu_place(p)) {
#ifdef PADDLE_WITH_IPU
290
      EmplaceDeviceContext<IPUDeviceContext>(
291 292
          place_to_device_context,
          p,
293
          disable_setting_default_stream_for_allocator);
J
jianghaicheng 已提交
294 295 296 297
#else
      PADDLE_THROW(
          platform::errors::Unimplemented("IPUPlace is not supported. Please "
                                          "re-compile with WITH_IPU option."));
298 299 300
#endif
    } else if (platform::is_npu_place(p)) {
#ifdef PADDLE_WITH_ASCEND_CL
301
      EmplaceDeviceContext<NPUDeviceContext>(
302 303
          place_to_device_context,
          p,
304
          disable_setting_default_stream_for_allocator);
305 306 307 308
#else
      PADDLE_THROW(platform::errors::Unimplemented(
          "NPUPlace is not supported. Please "
          "re-compile with WITH_ASCEND_CL option."));
309 310 311
#endif
    } else if (platform::is_npu_pinned_place(p)) {
#ifdef PADDLE_WITH_ASCEND_CL
312
      EmplaceDeviceContext<NPUPinnedDeviceContext>(
313 314
          place_to_device_context,
          p,
315
          disable_setting_default_stream_for_allocator);
316 317 318 319 320
#else
      PADDLE_THROW(platform::errors::Unimplemented(
          "NPUPinnedPlace is not supported. Please re-compile with "
          "WITH_ASCEND_CL "
          "option."));
321 322 323
#endif
    } else if (platform::is_custom_place(p)) {
#ifdef PADDLE_WITH_CUSTOM_DEVICE
324
      EmplaceDeviceContext<CustomDeviceContext>(
325 326
          place_to_device_context,
          p,
327
          disable_setting_default_stream_for_allocator);
328 329 330 331 332
#else
      PADDLE_THROW(platform::errors::Unimplemented(
          "CustomPlace is not supported. Please re-compile with "
          "WITH_CUSTOM_DEVICE "
          "option."));
D
dzhwinter 已提交
333 334 335 336 337
#endif
    }
  }
}

338 339
DeviceContextPool::DeviceContextPool(
    const std::vector<platform::Place>& places) {
340 341
  EmplaceDeviceContexts(&device_contexts_,
                        places,
342 343 344
                        /*disable_setting_default_stream_for_allocator=*/false);
}

J
jianghaicheng 已提交
345
#ifdef PADDLE_WITH_IPU
A
Allen Guo 已提交
346
IPUDeviceContext::IPUDeviceContext(IPUPlace place) : place_(place) {}
J
jianghaicheng 已提交
347

W
Wilber 已提交
348
const Place& IPUDeviceContext::GetPlace() const { return place_; }
A
Allen Guo 已提交
349

J
jianghaicheng 已提交
350 351 352 353 354 355 356
void IPUDeviceContext::Wait() const {
  /*! \brief  Wait for all operations completion in the stream. */
}

IPUDeviceContext::~IPUDeviceContext() {}

#endif
357
#ifdef PADDLE_WITH_XPU
358 359
XPUDeviceContext::XPUDeviceContext() : phi::XPUContext() {
  phi::XPUContext::Init();
W
Wilber 已提交
360
}
361

362
XPUDeviceContext::~XPUDeviceContext() {}
363

364 365
XPUDeviceContext::XPUDeviceContext(XPUPlace place) : phi::XPUContext(place) {
  phi::XPUContext::Init();
366
  LOG_FIRST_N(WARNING, 1) << "Please NOTE: xpu device: "
W
Wilber 已提交
367
                          << static_cast<int>(place.device);
368 369 370
}
#endif

371 372 373 374 375 376 377
#ifdef PADDLE_WITH_ASCEND_CL
NPUDeviceContext::NPUDeviceContext(NPUPlace place) : place_(place) {
  NPUDeviceGuard guard(place_.device);
  // PADDLE_ENFORCE_NPU_SUCCESS(aclrtCreateContext(&context_, place_.device));
  // NOTE(zhiqiu): Usually, no need to create context explicitly,
  // ACL creates a default context which contains 1 default stream
  // and 1 sync strean after aclrtSetDevice.
378
  platform::GetCurrentNPUContext(&context_);
379 380 381 382 383 384 385
  stream_.reset(new stream::NPUStream(place));
}

NPUDeviceContext::~NPUDeviceContext() {
  // NPUDeviceGuard guard(place_.device);
  // PADDLE_ENFORCE_NPU_SUCCESS(aclrtDestroyContext(context_));
}
386

387
void NPUDeviceContext::Wait() const {
388 389
  platform::RecordEvent record_event(
      "NPUDeviceContext/wait", platform::TracerEventType::UserDefined, 2);
390 391
  VLOG(4) << "NPU context(" << this << ")  Wait";
  stream_->Wait();
392 393 394 395
}

aclrtStream NPUDeviceContext::stream() const { return stream_->raw_stream(); }

W
Wilber 已提交
396
const Place& NPUDeviceContext::GetPlace() const { return place_; }
397 398

aclrtContext NPUDeviceContext::context() const { return context_; }
399 400 401 402 403 404 405 406 407 408 409 410 411 412

NPUPinnedDeviceContext::NPUPinnedDeviceContext() {
  eigen_device_.reset(new Eigen::DefaultDevice());
}

NPUPinnedDeviceContext::NPUPinnedDeviceContext(NPUPinnedPlace place)
    : place_(place) {
  eigen_device_.reset(new Eigen::DefaultDevice());
}

Eigen::DefaultDevice* NPUPinnedDeviceContext::eigen_device() const {
  return eigen_device_.get();
}

W
Wilber 已提交
413
const Place& NPUPinnedDeviceContext::GetPlace() const { return place_; }
414

415 416 417
#endif

#if defined(PADDLE_WITH_CUDA) || defined(PADDLE_WITH_HIP)
418

C
chengduoZH 已提交
419 420 421 422 423 424 425 426 427 428 429 430 431
CUDAPinnedDeviceContext::CUDAPinnedDeviceContext() {
  eigen_device_.reset(new Eigen::DefaultDevice());
}

CUDAPinnedDeviceContext::CUDAPinnedDeviceContext(CUDAPinnedPlace place)
    : place_(place) {
  eigen_device_.reset(new Eigen::DefaultDevice());
}

Eigen::DefaultDevice* CUDAPinnedDeviceContext::eigen_device() const {
  return eigen_device_.get();
}

W
Wilber 已提交
432
const Place& CUDAPinnedDeviceContext::GetPlace() const { return place_; }
L
Luo Tao 已提交
433
#endif
Q
qijun 已提交
434

435
#ifdef PADDLE_WITH_CUSTOM_DEVICE
436 437 438
CustomDeviceContext::CustomDeviceContext(CustomPlace place)
    : phi::CustomContext(place) {
  Init();
439
  stream_.reset(new phi::stream::Stream(place, stream()));
440 441 442
}

CustomDeviceContext::~CustomDeviceContext() {}
T
tensor-tang 已提交
443
#endif
Q
qijun 已提交
444
}  // namespace platform
Q
qijun 已提交
445
}  // namespace paddle