gpu_context.h 8.2 KB
Newer Older
W
Wilber 已提交
1
/* Copyright (c) 2022 PaddlePaddle Authors. All Rights Reserved.
2
Copyright (c) 2022 NVIDIA Corporation. All rights reserved.
3 4 5 6 7 8 9 10 11 12 13 14 15 16 17

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. */

#pragma once

W
Wilber 已提交
18 19
#include <array>
#include <functional>
W
Wilber 已提交
20
#include <mutex>
21 22 23
#include "paddle/phi/backends/gpu/forwards.h"
#include "paddle/phi/backends/gpu/gpu_decls.h"
#include "paddle/phi/backends/gpu/gpu_helper.h"
24
#include "paddle/phi/backends/gpu/gpu_info.h"
25 26
#include "paddle/phi/common/place.h"
#include "paddle/phi/core/device_context.h"
27

28
namespace phi {
29

W
Wilber 已提交
30 31
class DnnWorkspaceHandle {
 public:
32 33
  inline DnnWorkspaceHandle(Allocator* allocator, gpuStream_t stream)
      : allocator_(allocator), stream_(stream) {
W
Wilber 已提交
34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51
    mtx_.reset(new std::mutex());
  }

  inline void RunFunc(const std::function<void(void*)>& cudnn_func,
                      size_t required_workspace_bytes) {
    if (required_workspace_bytes > WorkspaceSize()) {
      ReallocWorkspace(required_workspace_bytes);
    }
    {
      std::lock_guard<std::mutex> guard(*mtx_);
      cudnn_func(allocation_ ? allocation_->ptr() : nullptr);
    }
  }

  /*! \brief Thread which call RunFuncSync() would release gpu memory after
   *  running the function. Currently this function is only used when cudnn
   *  exhaustive searching and callers have to guarantee that the input function
   *  is host blocking */
52 53 54
  void RunFuncSync(const std::function<void(void*)>& cudnn_func,
                   size_t required_workspace_bytes,
                   bool use_cached_allocation = true);
W
Wilber 已提交
55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71

  inline size_t WorkspaceSize() {
    if (allocation_ == nullptr) {
      return 0;
    }
    return allocation_->size();
  }

  void ResetWorkspace();

  void ReallocWorkspace(size_t required_workspace_bytes);

  DnnWorkspaceHandle(DnnWorkspaceHandle&&) = default;
  DnnWorkspaceHandle& operator=(DnnWorkspaceHandle&&) = delete;

 private:
  Allocator::AllocationPtr allocation_{nullptr};
72 73
  Allocator* allocator_{nullptr};  // Not owned
  gpuStream_t stream_{nullptr};    // Not owned
W
Wilber 已提交
74 75
  std::unique_ptr<std::mutex> mtx_;
};
W
Wilber 已提交
76

77
class PADDLE_API GPUContext : public DeviceContext {
W
Wilber 已提交
78 79
 public:
  GPUContext();
W
Wilber 已提交
80 81
  GPUContext(GPUContext&&);
  GPUContext& operator=(GPUContext&&);
W
Wilber 已提交
82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98

  explicit GPUContext(const GPUPlace& place);

  virtual ~GPUContext();

  /*! \brief  Return place in the device context. */
  const Place& GetPlace() const override;

  /*! \brief  Return gpu stream in the device context. */
  gpuStream_t stream() const;

  /*! \brief  Return cudnn  handle in the device context. */
  dnnHandle_t cudnn_handle() const;

  /*! \brief  Return cublas handle in the device context. */
  blasHandle_t cublas_handle() const;

99 100 101
  /*! \brief  Return cublasLt handle in the device context. */
  blasLtHandle_t cublaslt_handle() const;

W
Wilber 已提交
102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140
  /*! \brief  Return cusolver handle in the device context. */
  solverHandle_t cusolver_dn_handle() const;

  /*! \brief  Return cusparse handle in the device context. */
  sparseHandle_t cusparse_handle() const;

  /*! \brief  Wait for all operations completion in the stream. */
  void Wait() const override;

  /*! \brief  Wait for event in the stream. */
  void WaitEvent(gpuEvent_t ev) const;

  /*! \brief  Check whether tensor core is supported */
  bool tensor_core_available() const;

  /*! \brief  Return compute capability in the device context. */
  int GetComputeCapability() const;

  /*! \brief  Return the max physical thread count in the device context */
  int GetMaxPhysicalThreadCount() const;

  /*! \brief  Return the SM count in the device context */
  int GetSMCount() const;

  /*! \brief  Return the Max thread num of block in the device context */
  int GetMaxThreadsPerBlock() const;

  /*! \brief  Return the max grid dim size in the device context */
  std::array<int, 3> GetCUDAMaxGridDimSize() const;

  /*! \brief  Return eigen device in the device context. */
  Eigen::GpuDevice* eigen_device() const;

  /*! \brief  Return a cudnn workspace handle to call multiple cudnn
   *  functions without interrupting by other threads.
   *  Once the first cudnn function is called by the handle, a lock
   *  would be acquired to prevent other threads from accessing the
   *  workspace. Once the handle is destructed, the lock would be released.
   */
W
Wilber 已提交
141 142
  // TODO(wilber): The return type is a pointer, to be modified later.
  DnnWorkspaceHandle cudnn_workspace_handle() const;
W
Wilber 已提交
143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 184 185 186 187 188 189 190 191 192 193 194 195 196 197 198 199 200 201

 public:
  /*! \brief  Call cublas function safely. */
  void CublasCall(const std::function<void(blasHandle_t)>&) const;

  /*! \brief  Call cublas function with Tensor Core safely. If
      Tensor Core is not available, use DEFAULT_MATH instead. */
  void TensorCoreCublasCallIfAvailable(
      const std::function<void(blasHandle_t)>&) const;

  /*! \brief  Call cusparse function safely. */
  void CusparseCall(const std::function<void(sparseHandle_t)>&) const;

  void RecordEvent(gpuEvent_t ev, const std::function<void()>& callback) const;

  void RecordEvent(gpuEvent_t ev) const;

  void AddStreamCallback(const std::function<void()>& callback) const;

  void WaitStreamCallback() const;

 public:
  /*! \brief  Return nccl communicators. */
  ncclComm_t nccl_comm() const;

  /*! \brief  Set nccl communicators. */
  void set_nccl_comm(ncclComm_t comm);

 public:
  // NOTE: DeviceContext hold resources. Used in training scenarios.
  // The interface used by the training scene, DeviceContext will initialize
  // all resources and delete them when destructing.
  // Note that you must set the Allocator before calling Init function.
  void Init();

  // TODO(wilber): Why does the GetAllocator interface require a stream
  // parameter?
  // The temporary trick method bypasses this problem, and the following
  // interfaces
  // need to be deleted later.

  // Note that this is a trick implementation, which can be used to partially
  // initialize when the SetAllocator interface is not called.
  void PartialInitWithoutAllocator();
  // Note that this is a trick implementation that can be used to initialize
  // resources that require an Allocator when the SetAllocator interface is
  // called.
  void PartialInitWithAllocator();

 protected:
  // NOTE: External users manage resources. Used in inference scenarios.
  // The Set interface is for inference only, DeviceContext will mark the
  // resource as external, and will not delete any resource when destructing.
  void SetStream(gpuStream_t);

  void SetEigenDevice(Eigen::GpuDevice*);

  void SetBlasHandle(blasHandle_t);

W
Wilber 已提交
202 203 204 205
  void SetBlasTensorCoreHandle(blasHandle_t);

  void SetBlasTF32Handle(blasHandle_t);

206 207
  void SetBlasLtHandle(blasLtHandle_t);

W
Wilber 已提交
208 209 210 211 212 213 214 215 216 217 218 219 220 221 222 223 224 225 226 227 228 229 230 231 232 233 234
  void SetDnnHandle(dnnHandle_t);

  void SetSolverHandle(solverHandle_t);

  void SetSparseHandle(sparseHandle_t);

  void SetDnnWorkspaceHandle(DnnWorkspaceHandle*);

  void SetComputeCapability(int val);

  void SetMaxThreadsPerMultiProcessor(int val);

  void SetMultiProcessors(int val);

  void SetMaxThreadsPerBlock(int val);

  void SetMaxGridDimSize(const std::array<int, 3>& val);

  void SetDriverVersion(int val);

  void SetRuntimeVersion(int val);

 private:
  struct Impl;
  std::unique_ptr<Impl> impl_;
};

235 236 237 238 239 240 241
// Note: In order to register the kernel of CUDNN, GPUDNNContext is required.
// Currently, CUDNN kernel directly uses GPUContext. But if the kernel function
// has the same name, this will lead to duplicate instantiations of GPU kernel
// and GPUDNN kernel function, so if we using GPUDNNContext = GPUContext, we
// must use different function name for cudnn kernel
using GPUDNNContext = GPUContext;

242 243 244 245 246 247 248 249
// KPS (Kernel PrimitiveS API) needs to exist as a kind of backend,
// because we want to implement a KPS-based kernel and make it run
// on GPU and XPU at the same time, so we need KPSContext when registering
// KPS Kernel. Note: XPU and GPU cannot be compiled at the same time!
#if defined(PADDLE_WITH_CUDA) || defined(PADDLE_WITH_HIP)
using KPSContext = GPUContext;
#endif

250
}  // namespace phi