heter_resource.cc 3.5 KB
Newer Older
T
Thunderbrook 已提交
1 2 3 4 5 6 7 8 9 10 11 12 13 14
/* Copyright (c) 2020 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. */

T
Thunderbrook 已提交
15
#ifdef PADDLE_WITH_HETERPS
T
Thunderbrook 已提交
16 17 18 19 20 21
#include "heter_resource.h"
#include "paddle/fluid/platform/cuda_device_guard.h"

namespace paddle {
namespace framework {

22
GPUResource::GPUResource(std::vector<int>& dev_ids, int index) {
T
Thunderbrook 已提交
23
  index_ = index;
24 25
  dev_ids_ = dev_ids;
  dev_id_ = dev_ids_[index];
T
Thunderbrook 已提交
26 27

  platform::CUDADeviceGuard guard(dev_id_);
28 29
  local_streams_.resize(dev_ids_.size());
  comm_streams_.resize(dev_ids_.size());
30
  remote_streams_.resize(dev_ids_.size());
31 32

  for (size_t i = 0; i < dev_ids_.size(); ++i) {
33
    PADDLE_ENFORCE_GPU_SUCCESS(
34
        cudaStreamCreateWithFlags(&local_streams_[i], cudaStreamNonBlocking));
35
    PADDLE_ENFORCE_GPU_SUCCESS(
36
        cudaStreamCreateWithFlags(&comm_streams_[i], cudaStreamNonBlocking));
37
    PADDLE_ENFORCE_GPU_SUCCESS(
38
        cudaStreamCreateWithFlags(&remote_streams_[i], cudaStreamNonBlocking));
39
  }
T
Thunderbrook 已提交
40 41 42 43
}

GPUResource::~GPUResource() {
  platform::CUDADeviceGuard guard(dev_id_);
44
  for (size_t i = 0; i < local_streams_.size(); ++i) {
45
    PADDLE_ENFORCE_GPU_SUCCESS(cudaStreamDestroy(local_streams_[i]));
46 47
  }
  for (size_t i = 0; i < comm_streams_.size(); ++i) {
48
    PADDLE_ENFORCE_GPU_SUCCESS(cudaStreamDestroy(comm_streams_[i]));
49
  }
50
  for (size_t i = 0; i < remote_streams_.size(); ++i) {
51
    PADDLE_ENFORCE_GPU_SUCCESS(cudaStreamDestroy(remote_streams_[i]));
52
  }
T
Thunderbrook 已提交
53 54 55 56 57 58 59 60
}

void HeterPsResource::enable_p2p() {
  for (size_t i = 0; i < dev_ids_.size(); ++i) {
    platform::CUDADeviceGuard guard(dev_ids_[i]);
    for (size_t j = 0; j < dev_ids_.size(); ++j) {
      if (i != j) {
        int p2p_flag;
61
        PADDLE_ENFORCE_GPU_SUCCESS(
T
Thunderbrook 已提交
62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80
            cudaDeviceCanAccessPeer(&p2p_flag, dev_ids_[i], dev_ids_[j]));
        if (p2p_flag == 1) {
          cudaError_t ret = cudaDeviceEnablePeerAccess(dev_ids_[j], 0);
          if (ret != cudaSuccess && ret != cudaErrorPeerAccessAlreadyEnabled) {
            VLOG(0) << " Cuda error(" << ret << "), " << cudaGetErrorString(ret)
                    << ".";
          } else {
            cudaGetLastError();
          }
        }
      }
    }
  }
}

HeterPsResource::HeterPsResource(const std::vector<int>& dev_ids) {
  dev_ids_ = dev_ids;
  for (size_t i = 0; i < dev_ids_.size(); ++i) {
    std::shared_ptr<GPUResource> resource =
81
        std::make_shared<GPUResource>(dev_ids_, i);
T
Thunderbrook 已提交
82 83 84 85 86
    resources_.push_back(resource);
    devid_2_index_[dev_ids_[i]] = i;
  }
}

87 88 89 90 91 92
cudaStream_t HeterPsResource::comm_stream(int gpu_num, int stream_num) {
  return resources_[gpu_num]->comm_stream(stream_num);
}

cudaStream_t HeterPsResource::local_stream(int gpu_num, int stream_num) {
  return resources_[gpu_num]->local_stream(stream_num);
T
Thunderbrook 已提交
93 94
}

95 96
cudaStream_t HeterPsResource::remote_stream(int gpu_num, int stream_num) {
  return resources_[gpu_num]->remote_stream(stream_num);
T
Thunderbrook 已提交
97 98 99 100 101 102 103 104 105 106 107 108 109
}

int HeterPsResource::dev_id(int num) { return dev_ids_[num]; }

int HeterPsResource::get_index_by_devid(int devid) {
  return devid_2_index_[devid];
}

int HeterPsResource::total_gpu() { return dev_ids_.size(); }

}  // end namespace framework
}  // end namespace paddle
#endif