chore: import upstream snapshot with attribution

paddle/phi/backends/gpu/CMakeLists.txt

@@ -0,0 +1 @@
+add_subdirectory(cuda)

paddle/phi/backends/gpu/cuda/CMakeLists.txt

@@ -0,0 +1,15 @@
+collect_srcs(backends_srcs SRCS cudnn_workspace_helper.cc)
+
+if(WITH_GPU)
+  if(WIN32)
+    nv_library(
+      gpu_event_timer
+      SRCS gpu_event_timer.cc
+      DEPS glog phi)
+  else()
+    nv_library(
+      gpu_event_timer
+      SRCS gpu_event_timer.cc
+      DEPS phi_core glog)
+  endif()
+endif()

paddle/phi/backends/gpu/cuda/cuda_device_function.h

@@ -0,0 +1,182 @@
+/* Copyright (c) 2022 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. */
+
+#pragma once
+
+#include "paddle/phi/common/bfloat16.h"
+#include "paddle/phi/common/complex.h"
+#include "paddle/phi/common/float16.h"
+#include "paddle/phi/core/enforce.h"
+
+namespace phi {
+namespace backends {
+namespace gpu {
+
+#define FULL_WARP_MASK 0xFFFFFFFF
+#define CREATE_SHFL_MASK(mask, predicate) \
+  mask = __ballot_sync(FULL_WARP_MASK, (predicate))
+
+#define CUDA_LAUNCH_KERNEL_BASE(dim, ...)  \
+  case (dim): {                            \
+    constexpr auto kPowerOfTwoDim = (dim); \
+    __VA_ARGS__;                           \
+  } break
+
+#define CUDA_LAUNCH_KERNEL_HELPER(...)          \
+  CUDA_LAUNCH_KERNEL_BASE(1024, ##__VA_ARGS__); \
+  CUDA_LAUNCH_KERNEL_BASE(512, ##__VA_ARGS__);  \
+  CUDA_LAUNCH_KERNEL_BASE(256, ##__VA_ARGS__);  \
+  CUDA_LAUNCH_KERNEL_BASE(128, ##__VA_ARGS__);  \
+  CUDA_LAUNCH_KERNEL_BASE(64, ##__VA_ARGS__);   \
+  CUDA_LAUNCH_KERNEL_BASE(32, ##__VA_ARGS__);
+
+template <typename T>
+__forceinline__ __device__ T
+CudaShuffleDownSync(unsigned mask, T val, int delta, int width = warpSize) {
+  return __shfl_down_sync(mask, val, static_cast<unsigned>(delta), width);
+}
+
+template <typename T>
+__forceinline__ __device__ T CudaShuffleXorSync(unsigned mask,
+                                                T val,
+                                                int width = warpSize) {
+  return __shfl_xor_sync(mask, val, width);
+}
+
+template <>
+__forceinline__ __device__ phi::dtype::float16 CudaShuffleDownSync(
+    unsigned mask, phi::dtype::float16 val, int delta, int width) {
+  return phi::dtype::float16(__shfl_down_sync(
+      mask, val.to_half(), static_cast<unsigned>(delta), width));
+}
+
+template <>
+__forceinline__ __device__ phi::dtype::bfloat16 CudaShuffleDownSync(
+    unsigned mask, phi::dtype::bfloat16 val, int delta, int width) {
+#if defined(PADDLE_CUDA_BF16)
+  return phi::dtype::bfloat16(__shfl_down_sync(
+      mask, val.to_nv_bfloat16(), static_cast<unsigned>(delta), width));
+#else
+  PADDLE_ENFORCE(
+      false, "__shfl_down_sync with bfloat16 is not supported on cuda <= 11.");
+#endif
+}
+
+template <>
+__forceinline__ __device__ phi::dtype::complex<float> CudaShuffleDownSync(
+    unsigned mask, phi::dtype::complex<float> val, int delta, int width) {
+  float real = static_cast<float>(__shfl_down_sync(
+      mask, static_cast<float>(val.real), static_cast<unsigned>(delta), width));
+  float imag = static_cast<float>(__shfl_down_sync(
+      mask, static_cast<float>(val.imag), static_cast<unsigned>(delta), width));
+  return phi::dtype::complex<float>(real, imag);
+}
+
+template <>
+__forceinline__ __device__ phi::dtype::complex<double> CudaShuffleDownSync(
+    unsigned mask, phi::dtype::complex<double> val, int delta, int width) {
+  double real =
+      static_cast<double>(__shfl_down_sync(mask,
+                                           static_cast<double>(val.real),
+                                           static_cast<unsigned>(delta),
+                                           width));
+  double imag =
+      static_cast<double>(__shfl_down_sync(mask,
+                                           static_cast<double>(val.imag),
+                                           static_cast<unsigned>(delta),
+                                           width));
+  return phi::dtype::complex<double>(real, imag);
+}
+
+template <>
+__forceinline__ __device__ phi::dtype::float16 CudaShuffleXorSync(
+    unsigned mask, phi::dtype::float16 val, int width) {
+  return phi::dtype::float16(__shfl_xor_sync(mask, val.to_half(), width));
+}
+
+template <>
+__forceinline__ __device__ phi::dtype::bfloat16 CudaShuffleXorSync(
+    unsigned mask, phi::dtype::bfloat16 val, int width) {
+  return phi::dtype::bfloat16(
+      __shfl_xor_sync(mask, val.to_nv_bfloat16(), width));
+}
+
+template <>
+__forceinline__ __device__ phi::dtype::complex<float> CudaShuffleXorSync(
+    unsigned mask, phi::dtype::complex<float> val, int width) {
+  float real = static_cast<float>(
+      __shfl_xor_sync(mask, static_cast<float>(val.real), width));
+  float imag = static_cast<float>(
+      __shfl_xor_sync(mask, static_cast<float>(val.imag), width));
+  return phi::dtype::complex<float>(real, imag);
+}
+
+template <>
+__forceinline__ __device__ phi::dtype::complex<double> CudaShuffleXorSync(
+    unsigned mask, phi::dtype::complex<double> val, int width) {
+  double real = static_cast<double>(
+      __shfl_xor_sync(mask, static_cast<double>(val.real), width));
+  double imag = static_cast<double>(
+      __shfl_xor_sync(mask, static_cast<double>(val.imag), width));
+  return phi::dtype::complex<double>(real, imag);
+}
+
+template <typename T>
+__forceinline__ __device__ T
+CudaShuffleSync(unsigned mask, T val, int src_line, int width = 32) {
+  return __shfl_sync(mask, val, src_line, width);
+}
+
+template <typename T>
+HOSTDEVICE T Infinity() {
+  return INFINITY;
+}
+
+template <typename T>
+__device__ T reduceSum(T val, int tid, int len) {
+  // NOTE(zcd): The warp size should be taken from the
+  // parameters of the GPU but not specified as 32 simply.
+  // To make the reduceSum more efficiently,
+  // I use Warp-Level Parallelism and assume the Warp size
+  // is 32 which may be different for different GPU,
+  // but most card's warp size is 32.
+  const int warpSize = 32;
+  __shared__ T shm[warpSize];
+  unsigned mask = 0u;
+  CREATE_SHFL_MASK(mask, tid < len);
+
+  for (int offset = warpSize / 2; offset > 0; offset /= 2)
+    val += phi::backends::gpu::CudaShuffleDownSync(mask, val, offset);
+
+  if (tid < warpSize) shm[tid] = 0;
+  __syncthreads();
+
+  if (tid % warpSize == 0) {
+    shm[tid / warpSize] = val;
+  }
+  __syncthreads();
+
+  CREATE_SHFL_MASK(mask, tid < warpSize);
+
+  if (tid < warpSize) {
+    val = shm[tid];
+    for (int offset = warpSize / 2; offset > 0; offset /= 2)
+      val += phi::backends::gpu::CudaShuffleDownSync(mask, val, offset);
+  }
+  return val;
+}
+
+}  // namespace gpu
+}  // namespace backends
+}  // namespace phi

paddle/phi/backends/gpu/cuda/cuda_graph.cc

@@ -0,0 +1,477 @@
+// Copyright (c) 2022 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.
+
+#include "paddle/phi/backends/gpu/cuda/cuda_graph.h"
+#include "glog/logging.h"
+#include "paddle/common/flags.h"
+
+#ifdef PADDLE_WITH_CUDA
+
+COMMON_DECLARE_bool(use_cuda_malloc_async_allocator);
+COMMON_DECLARE_bool(auto_free_cudagraph_allocations_on_launch);
+
+namespace phi::backends::gpu {
+
+std::unique_ptr<CUDAGraph> CUDAGraph::capturing_graph_{nullptr};
+paddle::optional<std::thread::id> CUDAGraph::capturing_thread_id_{paddle::none};
+std::vector<std::function<void()>> CUDAGraph::cudagraph_pre_capture_callbacks_;
+
+static std::vector<cudaGraphNode_t> ToposortCUDAGraph(cudaGraph_t graph) {
+  size_t num_nodes;
+  PADDLE_ENFORCE_GPU_SUCCESS(cudaGraphGetNodes(graph, nullptr, &num_nodes));
+  std::vector<cudaGraphNode_t> nodes(num_nodes);
+  PADDLE_ENFORCE_GPU_SUCCESS(
+      cudaGraphGetNodes(graph, nodes.data(), &num_nodes));
+
+  size_t num_edges;
+#if CUDA_VERSION < 13000
+  PADDLE_ENFORCE_GPU_SUCCESS(
+      cudaGraphGetEdges(graph, nullptr, nullptr, &num_edges));
+  std::vector<cudaGraphNode_t> from(num_edges), to(num_edges);
+  PADDLE_ENFORCE_GPU_SUCCESS(
+      cudaGraphGetEdges(graph, from.data(), to.data(), &num_edges));
+#else
+  PADDLE_ENFORCE_GPU_SUCCESS(
+      cudaGraphGetEdges(graph, nullptr, nullptr, nullptr, &num_edges));
+  std::vector<cudaGraphNode_t> from(num_edges), to(num_edges);
+  PADDLE_ENFORCE_GPU_SUCCESS(
+      cudaGraphGetEdges(graph, from.data(), to.data(), nullptr, &num_edges));
+#endif
+
+  std::unordered_map<cudaGraphNode_t, std::unordered_set<cudaGraphNode_t>>
+      in_edges, out_edges;
+  for (auto node : nodes) {
+    in_edges[node];
+    out_edges[node];
+  }
+
+  for (size_t i = 0; i < num_edges; ++i) {
+    in_edges[to[i]].insert(from[i]);
+    out_edges[from[i]].insert(to[i]);
+  }
+
+  std::queue<cudaGraphNode_t> q;
+  for (const auto &pair : in_edges) {
+    if (pair.second.empty()) {
+      q.push(pair.first);
+    }
+  }
+
+  nodes.clear();
+  while (!q.empty()) {
+    auto cur = q.front();
+    q.pop();
+    nodes.push_back(cur);
+
+    for (auto out_node : out_edges.at(cur)) {
+      auto &in_nodes = in_edges.at(out_node);
+      in_nodes.erase(cur);
+      if (in_nodes.empty()) {
+        q.push(out_node);
+      }
+    }
+  }
+  PADDLE_ENFORCE_EQ(
+      nodes.size(),
+      num_nodes,
+      common::errors::InvalidArgument("Toposort error, this may be a bug."));
+  return nodes;
+}
+
+CUDAGraphID CUDAGraph::UniqueID() {
+  static std::atomic<CUDAGraphID> id;
+  return id.fetch_add(1);
+}
+
+int64_t CUDAGraph::UniqueMemoryPoolID() {
+  static std::atomic<int64_t> id(CUDAGraph::kDefaultPoolID + 1);
+  return id.fetch_add(1);
+}
+
+void CUDAGraph::Reset() {
+  if (is_reset_) return;
+  for (auto graph : graphs_) {
+    PADDLE_ENFORCE_GPU_SUCCESS(cudaGraphDestroy(graph));
+  }
+  graphs_.clear();
+  for (auto exec_graph : exec_graphs_) {
+    PADDLE_ENFORCE_GPU_SUCCESS(cudaGraphExecDestroy(exec_graph));
+  }
+  exec_graphs_.clear();
+  // callback should be called in reverse order because the latter added
+  // callback may rely on the former added callback.
+  for (auto iter = cudagraph_post_reset_callbacks_.rbegin();
+       iter != cudagraph_post_reset_callbacks_.rend();
+       ++iter) {
+    (*iter)(*this);
+  }
+  cudagraph_post_reset_callbacks_.clear();
+  is_reset_ = true;
+}
+
+void CUDAGraph::Replay() {
+  is_replayed_ = true;
+  PADDLE_ENFORCE_EQ(is_reset_,
+                    false,
+                    common::errors::PermissionDenied(
+                        "Cannot replay the CUDA Graph after reset is called."));
+  size_t n = exec_graphs_.size();
+  for (size_t i = 0; i < n; ++i) {
+    if (!is_first_run_) {
+      for (auto &hook : cudagraph_pre_replay_callbacks_[i]) {
+        hook(exec_graphs_[i]);
+      }
+    }
+    PADDLE_ENFORCE_GPU_SUCCESS(cudaGraphLaunch(exec_graphs_[i], stream_));
+  }
+  is_first_run_ = false;
+}
+
+void CUDAGraph::BeginSegmentCapture() {
+  ThrowErrorIfNotSupportCUDAGraph();
+  PADDLE_ENFORCE_EQ(IsCapturing(),
+                    true,
+                    common::errors::PermissionDenied(
+                        "BeginSegmentCapture should be called when CUDA "
+                        "Graph is capturing."));
+  if (IsThreadLocalCapturing()) {
+    PADDLE_ENFORCE_EQ(IsThisThreadCapturing(),
+                      true,
+                      common::errors::PermissionDenied(
+                          "When capturing CUDA Graph in the thread local mode, "
+                          "you cannot begin segmented capturing in the thread "
+                          "which is not the one that starts the capturing."));
+  }
+
+  for (auto &hook : cudagraph_pre_capture_callbacks_) {
+    hook();
+  }
+
+  PADDLE_ENFORCE_GPU_SUCCESS(cudaStreamBeginCapture(
+      capturing_graph_->stream_, capturing_graph_->capture_mode_));
+  PADDLE_ENFORCE_EQ(IsValidCapturing(),
+                    true,
+                    common::errors::PermissionDenied(
+                        "CUDA Graph should not be invalidated."));
+  VLOG(10) << "Begin to capture CUDA Graph with ID " << capturing_graph_->id_
+           << ", segment id " << capturing_graph_->graphs_.size()
+           << ", memory pool id " << capturing_graph_->pool_id_;
+}
+
+void CUDAGraph::BeginCapture(phi::GPUPlace place,
+                             cudaStream_t stream,
+                             cudaStreamCaptureMode mode,
+                             bool enable_replace) {
+  ThrowErrorIfNotSupportCUDAGraph();
+  PADDLE_ENFORCE_EQ(IsCapturing(),
+                    false,
+                    common::errors::PermissionDenied(
+                        "CUDA Graph can only captured one by one."));
+  PADDLE_ENFORCE_NOT_NULL(
+      stream,
+      common::errors::PermissionDenied(
+          "CUDA Graph cannot be captured in default CUDA stream 0."));
+  capturing_graph_.reset(new CUDAGraph(enable_replace));
+  capturing_graph_->place_ = place;
+  capturing_graph_->stream_ = stream;
+  capturing_graph_->capture_mode_ = mode;
+  if (mode == cudaStreamCaptureModeThreadLocal) {
+    capturing_thread_id_ = std::this_thread::get_id();
+    VLOG(10) << "Capturing CUDA Graph in thread local mode, thread id: "
+             << capturing_thread_id_;
+  }
+  BeginSegmentCapture();
+}
+
+inline void sync_streams(gpuStream_t to_record, gpuStream_t to_wait) {
+  if (to_record == to_wait) return;
+  cudaEvent_t event = nullptr;
+  PADDLE_ENFORCE_GPU_SUCCESS(
+      cudaEventCreateWithFlags(&event, cudaEventDisableTiming));
+  PADDLE_ENFORCE_GPU_SUCCESS(cudaEventRecord(event, to_record));
+  PADDLE_ENFORCE_GPU_SUCCESS(cudaStreamWaitEvent(to_wait, event));
+  PADDLE_ENFORCE_GPU_SUCCESS(cudaEventDestroy(event));
+}
+
+void CUDAGraph::EndSegmentCapture() {
+  ThrowErrorIfNotSupportCUDAGraph();
+  PADDLE_ENFORCE_EQ(
+      IsCapturing(),
+      true,
+      common::errors::PermissionDenied("No CUDA Graph is capturing."));
+
+  for (const auto &stream : capturing_graph_->streams_to_join_) {
+    VLOG(10) << "Joining steam when the capture is going to end stream ="
+             << stream;
+    sync_streams(stream, capturing_graph_->stream_);
+  }
+  capturing_graph_->streams_to_join_.clear();
+
+  cudaGraph_t graph;
+  PADDLE_ENFORCE_GPU_SUCCESS(
+      cudaStreamEndCapture(capturing_graph_->stream_, &graph));
+  auto num_nodes = static_cast<size_t>(-1);
+  PADDLE_ENFORCE_GPU_SUCCESS(cudaGraphGetNodes(graph, nullptr, &num_nodes));
+  if (num_nodes == 0) {
+    PADDLE_ENFORCE_GPU_SUCCESS(cudaGraphDestroy(graph));
+    VLOG(10) << "Skip empty CUDA Graph with ID " << capturing_graph_->id_
+             << ", segment id " << capturing_graph_->graphs_.size()
+             << ", memory pool id " << capturing_graph_->pool_id_;
+    return;
+  }
+
+  for (auto &cudagraph_post_capture_callback :
+       capturing_graph_->cudagraph_post_capture_callbacks_) {
+    cudagraph_post_capture_callback();
+  }
+  capturing_graph_->cudagraph_post_capture_callbacks_.clear();
+
+  capturing_graph_->cudagraph_pre_replay_callbacks_.emplace_back(
+      CUDAGraphNodeLauncher::Instance().GetParameterSettersForExecGraph(graph));
+
+  cudaGraphExec_t exec_graph;
+  if (FLAGS_use_cuda_malloc_async_allocator &&
+      FLAGS_auto_free_cudagraph_allocations_on_launch) {
+#if CUDA_VERSION >= 11040
+    VLOG(1) << "cudaGraphInstantiateFlagAutoFreeOnLaunch is enabled!";
+    PADDLE_ENFORCE_GPU_SUCCESS(cudaGraphInstantiateWithFlags(
+        &exec_graph, graph, cudaGraphInstantiateFlagAutoFreeOnLaunch));
+#else
+    PADDLE_THROW(common::errors::Unimplemented(
+        "The cudaGraphInstantiateFlagAutoFreeOnLaunch is only supported when "
+        "CUDA version >= 11.4.0"));
+#endif
+  } else {
+    PADDLE_ENFORCE_GPU_SUCCESS(
+        cudaGraphInstantiate(&exec_graph, graph, nullptr, nullptr, 0));
+  }
+  VLOG(10) << "End to capture CUDA Graph with ID " << capturing_graph_->id_
+           << ", segment id " << capturing_graph_->graphs_.size()
+           << ", memory pool id " << capturing_graph_->pool_id_;
+  capturing_graph_->graphs_.emplace_back(graph);
+  capturing_graph_->exec_graphs_.emplace_back(exec_graph);
+  if (capturing_graph_->enable_replace_) {
+    capturing_graph_->CacheKernelNodeInfos(capturing_graph_->graphs_.size() -
+                                           1);
+  }
+}
+
+std::unique_ptr<CUDAGraph> CUDAGraph::EndCapture() {
+  EndSegmentCapture();
+  capturing_thread_id_ = paddle::none;
+  return std::move(capturing_graph_);
+}
+
+bool CUDAGraph::IsValidCapturing() {
+  if (!IsCapturing()) return false;
+  cudaStreamCaptureStatus status;
+  CUDAGraphID id;
+  PADDLE_ENFORCE_GPU_SUCCESS(
+      cudaStreamGetCaptureInfo(capturing_graph_->stream_, &status, &id));
+  return status == cudaStreamCaptureStatusActive;
+}
+
+static std::string ConcatPath(const std::string &dirname,
+                              const std::string &filename) {
+#ifdef _WIN32
+  const std::array<char, 3> kFileSep = {"\\"};
+#else
+  const std::array<char, 2> kFileSep = {"/"};
+#endif
+  if (!dirname.empty() && dirname.back() == kFileSep[0]) {
+    return dirname + filename;
+  } else {
+    return dirname + kFileSep.data() + filename;
+  }
+}
+
+void CUDAGraph::PrintToDotFiles(const std::string &dirname,
+                                unsigned int flags) {
+  ThrowErrorIfNotSupportCUDAGraph();
+#if CUDA_VERSION >= 11030
+  for (size_t i = 0; i < graphs_.size(); ++i) {
+    auto filename =
+        ConcatPath(dirname, "segment_" + std::to_string(i) + ".dot");
+    VLOG(10) << "Save the " << i << "-th segment of graph " << id_ << " to "
+             << filename;
+    PADDLE_ENFORCE_GPU_SUCCESS(
+        cudaGraphDebugDotPrint(graphs_[i], filename.c_str(), flags));
+  }
+#else
+  PADDLE_THROW(common::errors::Unimplemented(
+      "The print_to_dot_files() method is only supported when CUDA version >= "
+      "11.3."));
+#endif
+}
+
+void CUDAGraph::CacheKernelNodeInfos(size_t segment_idx) {
+  auto &graph = graphs_[segment_idx];
+  size_t numNodes = 0;
+  cudaGraphGetNodes(graph, nullptr, &numNodes);
+  std::vector<cudaGraphNode_t> nodes(numNodes);
+  cudaGraphGetNodes(graph, nodes.data(), &numNodes);
+
+  std::vector<KernelNodeInfo> kernel_nodes;
+  for (auto &node : nodes) {
+    cudaGraphNodeType type;
+    cudaGraphNodeGetType(node, &type);
+    if (type == cudaGraphNodeTypeKernel) {
+      KernelNodeInfo info;
+      info.node = node;
+      memset(&info.params, 0, sizeof(info.params));
+      // Use Driver API to get params (works for both Runtime and Driver API
+      // kernels, unlike Runtime API which may return invalid kernelParams
+      // for JIT kernels such as DeepGEMM)
+      dynload::cuGraphKernelNodeGetParams(static_cast<CUgraphNode>(node),
+                                          &info.params);
+      info.param_infos =
+          GetKernelParamInfos(static_cast<CUfunction>(info.params.func));
+      kernel_nodes.emplace_back(std::move(info));
+    }
+  }
+  VLOG(4) << "Cached " << kernel_nodes.size() << " kernel nodes for segment "
+          << segment_idx;
+  cached_kernel_nodes_.emplace_back(std::move(kernel_nodes));
+}
+
+void CUDAGraph::ReplaceInputPtrs(const std::vector<void *> &old_ptrs,
+                                 const std::vector<void *> &new_ptrs) {
+  PADDLE_ENFORCE_EQ(
+      enable_replace_,
+      true,
+      common::errors::PermissionDenied(
+          "ReplaceInputPtrs requires enable_replace to be set to true "
+          "when creating CUDAGraph."));
+#if CUDA_VERSION >= 12040
+  for (size_t i = 0; i < cached_kernel_nodes_.size(); ++i) {
+    for (auto &kernel_info : cached_kernel_nodes_[i]) {
+      auto &params = kernel_info.params;
+      bool modified = false;
+
+      for (size_t k = 0; k < kernel_info.param_infos.size(); k++) {
+        size_t param_size = kernel_info.param_infos[k].size;
+        char *param_base = reinterpret_cast<char *>(params.kernelParams[k]);
+
+        for (size_t offset = 0; offset + sizeof(void *) <= param_size;
+             offset += sizeof(void *)) {
+          void *actual_val = *(reinterpret_cast<void **>(param_base + offset));
+          for (size_t j = 0; j < old_ptrs.size(); j++) {
+            if (old_ptrs[j] == actual_val) {
+              VLOG(4) << "cuda func " << params.func << " match old ptr "
+                      << actual_val << " at param " << k << " offset " << offset
+                      << ", replace with " << new_ptrs[j];
+              *(reinterpret_cast<void **>(param_base + offset)) = new_ptrs[j];
+              modified = true;
+              break;
+            }
+          }
+        }
+      }
+      if (modified) {
+        dynload::cuGraphExecKernelNodeSetParams(
+            static_cast<CUgraphExec>(exec_graphs_[i]),
+            static_cast<CUgraphNode>(kernel_info.node),
+            &params);
+      }
+    }
+  }
+#endif
+}
+
+std::vector<CUDAGraph::KernelParamInfo> CUDAGraph::GetKernelParamInfos(
+    CUfunction func) {
+  std::vector<KernelParamInfo> infos;
+#if CUDA_VERSION >= 12040
+  size_t paramOffset, paramSize;
+  int k = 0;
+
+  while (dynload::cuFuncGetParamInfo(func, k, &paramOffset, &paramSize) ==
+         CUDA_SUCCESS) {
+    infos.push_back({paramOffset, paramSize});
+    VLOG(4) << "[GetKernelParamInfos] func " << func << " param[" << k
+            << "] offset=" << paramOffset << " size=" << paramSize;
+    k++;
+  }
+#endif
+  return infos;
+}
+
+void CUDAGraphNodeLauncher::KernelNodeLaunch(
+    parameterSetter_t parameterSetter, gpuKernelCallback_t cudaKernelCallback) {
+  if (UNLIKELY(phi::backends::gpu::CUDAGraph::IsThisThreadCapturing())) {
+    unsigned int id = GenerateIdentifier();
+    auto cudaFunc = cudaKernelCallback(id);
+
+    parameterSetters[cudaFunc][id] = parameterSetter;
+    VLOG(10) << "[KernelNodeLaunch] Launch kernel with cudaFunc = " << cudaFunc
+             << " id = " << id;
+  } else {
+    cudaKernelCallback(0);
+  }
+}
+
+std::vector<cudaGraphExecuterSetter_t>
+CUDAGraphNodeLauncher::GetParameterSettersForExecGraph(cudaGraph_t graph) {
+  size_t num_nodes;
+  PADDLE_ENFORCE_GPU_SUCCESS(cudaGraphGetNodes(graph, nullptr, &num_nodes));
+  std::vector<cudaGraphNode_t> nodes(num_nodes);
+  PADDLE_ENFORCE_GPU_SUCCESS(
+      cudaGraphGetNodes(graph, nodes.data(), &num_nodes));
+
+  std::vector<std::function<void(cudaGraphExec_t)>> hooks;
+  for (auto node : nodes) {
+    CUgraphNode cuNode = node;
+    CUgraphNodeType pType;
+    PADDLE_ENFORCE_GPU_SUCCESS(dynload::cuGraphNodeGetType(cuNode, &pType));
+    if (pType == CU_GRAPH_NODE_TYPE_KERNEL) {
+      CUDA_KERNEL_NODE_PARAMS cuParams;
+
+      PADDLE_ENFORCE_GPU_SUCCESS(
+          dynload::cuGraphKernelNodeGetParams(cuNode, &cuParams));
+      gpuKernelParams kernel_params(cuParams.kernelParams);
+      auto kernel =
+          parameterSetters.find(static_cast<cudaFunction_t>(cuParams.func));
+      VLOG(10) << "[GetParameterSettersForExecGraph] cuParams.func = "
+               << cuParams.func;
+      // There exists a parameter setter
+      if (kernel != parameterSetters.end()) {
+        auto launchSequence = kernel->second;
+        unsigned int id = kernel_params.As<int>(0);
+
+        VLOG(10) << "[GetParameterSettersForExecGraph] Find launch kernel id = "
+                 << id;
+        auto parameterSetter = launchSequence.find(id);
+        if (parameterSetter != launchSequence.end()) {
+          auto setter = parameterSetter->second;
+          hooks.emplace_back([setter, cuNode, cuParams](
+                                 cudaGraphExec_t exec_graph) {
+            gpuKernelParams kernel_params(cuParams.kernelParams);
+            setter(kernel_params);
+            PADDLE_ENFORCE_GPU_SUCCESS(dynload::cuGraphExecKernelNodeSetParams(
+                static_cast<CUgraphExec>(exec_graph), cuNode, &cuParams));
+          });
+        } else {
+          PADDLE_THROW(common::errors::InvalidArgument(
+              "Error: does not find launch id"));
+        }
+      }
+    }
+  }
+
+  return hooks;
+}
+
+}  // namespace phi::backends::gpu
+
+#endif

paddle/phi/backends/gpu/cuda/cuda_graph.h

@@ -0,0 +1,408 @@
+// Copyright (c) 2022 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.
+
+#pragma once
+
+#include <array>
+#include <atomic>
+#include <functional>
+#include <future>
+#include <memory>
+#include <mutex>
+#include <queue>
+#include <set>
+#include <thread>
+#include <unordered_map>
+#include <unordered_set>
+#include <vector>
+
+#include "paddle/common/errors.h"
+#include "paddle/common/macros.h"
+#include "paddle/phi/backends/context_pool.h"
+#include "paddle/phi/backends/device_code.h"
+#include "paddle/phi/backends/gpu/gpu_context.h"
+#include "paddle/phi/common/memory_utils.h"
+#include "paddle/phi/common/place.h"
+#include "paddle/phi/core/enforce.h"
+#include "paddle/utils/optional.h"
+
+#ifdef PADDLE_WITH_CUDA
+namespace phi {
+namespace backends {
+namespace gpu {
+
+class CUDAGraphContextManager {
+ public:
+  using DeviceContextMap =
+      std::map<Place, std::shared_future<std::unique_ptr<DeviceContext>>>;
+
+  static CUDAGraphContextManager &Instance() {
+    static CUDAGraphContextManager *cuda_graph_ctx_manager =
+        new CUDAGraphContextManager;
+    return *cuda_graph_ctx_manager;
+  }
+
+  DeviceContext *Get(int64_t pool_id, const Place &place, int stream_priority) {
+    std::lock_guard<std::mutex> lk(ctx_mtx_);
+
+    DeviceContextMap &ctxs = cuda_graph_ctx_pool_[pool_id];
+    if (ctxs.find(place) == ctxs.end()) {
+      phi::memory_utils::EmplaceDeviceContexts(
+          &ctxs,
+          {place},
+          /*disable_setting_default_stream_for_allocator=*/true,
+          stream_priority);
+    }
+    return ctxs[place].get().get();
+  }
+
+  void RecordCapturingDeviceContext(DeviceContext *dev_ctx) {
+    capturing_ctxs_.insert(dev_ctx);
+  }
+
+  std::set<DeviceContext *> GetAllCapturingDeviceContexts() const {
+    return capturing_ctxs_;
+  }
+
+  void ClearDeviceContextsRecords() { capturing_ctxs_.clear(); }
+
+ private:
+  CUDAGraphContextManager() {}
+  DISABLE_COPY_AND_ASSIGN(CUDAGraphContextManager);
+
+  std::mutex ctx_mtx_;
+  std::unordered_map<int64_t, DeviceContextMap> cuda_graph_ctx_pool_;
+  std::set<DeviceContext *> capturing_ctxs_;
+};
+
+class gpuKernelParams {
+ public:
+  explicit gpuKernelParams(void **params) : kernelParams(params) {}
+
+  template <typename T>
+  T &As(size_t idx) const {
+    return *reinterpret_cast<T *>(kernelParams[idx]);
+  }
+
+  void **getParams() const { return kernelParams; }
+
+ private:
+  void **kernelParams;
+};
+
+using cudaGraphExecuterSetter_t = std::function<void(cudaGraphExec_t)>;
+
+//  ** class CUDAGraphNodeLauncher
+//
+//  This class offers a interface for launching CUDA kernels in CUDA Graph, we
+//  utilize the `cudaGraphExecKernelNodeSetParams` function for parameter setup.
+//  Launching kernels via this class ensures proper management.
+//
+//  NOTE: It's essential that the first parameter for any kernel launched
+//  through this class is an `unsigned int` identifier. This identifier plays a
+//  crucial role in linking the CUDA kernel to its corresponding CUDA graph
+//  node. We tag each kernel launch with a unique identifier to maintain
+//  structured linkage with its CUDA graph node.
+//
+//  NOTE: This class use a singleton design pattern ensures there's only a
+//  single global instance accessible via the `Instance()` method.
+class CUDAGraphNodeLauncher {
+ public:
+  //  [Parameter Setter Callback]
+  //  Sets the kernel's parameters BEFORE activating the CUDA graph. It enables
+  //  dynamic determination and setup of kernel arguments.
+  //
+  //  parameterSetter_t parameterSetter = [saved_state](gpuKernelParams
+  //  &param){
+  //      // Code to compute and the parameter values from the saved_state
+  //      // ...
+  //      param.As<type>(idx) = calculated_value;
+  //  };
+  using parameterSetter_t = std::function<void(gpuKernelParams &)>;
+
+  //  [CUDA Kernel Callback]
+  //  Acts as the launcher for the kernel. It accepts an `unsigned int`
+  //  identifier and uses it for the kernel launch.
+  //  The `cudaGetFuncBySymbol` method can be used to fetch the `cudaFunction_t`
+  //  reference of the kernel from the kernel pointer.
+  //  gpuKernelCallback_t cudaKernelCallback = [=](unsigned int id) {
+  //      // cudaFunction_t is REQUIRED to get here
+  //      cudaFunction_t cudaFunc;
+  //      PADDLE_ENFORCE_GPU_SUCCESS(cudaGetFuncBySymbol(&cudaFunc, &kernel));
+  //
+  //      kernel<<<>>>(id, ...);  // Launching the kernel with id
+  //      return cudaFunc;
+  //  };
+  using gpuKernelCallback_t = std::function<cudaFunction_t(unsigned int)>;
+
+  //  [Kernel Launch]
+  //  With the callbacks defined and the CUDA function obtained, the kernel can
+  //  be launched using the `KernelNodeLaunch` method.
+  void KernelNodeLaunch(parameterSetter_t parameterSetter,
+                        gpuKernelCallback_t cudaKernelCallback);
+
+  std::vector<cudaGraphExecuterSetter_t> GetParameterSettersForExecGraph(
+      cudaGraph_t graph);
+
+  parameterSetter_t GetParameterSetter(const gpuKernelParams &params);
+
+  static CUDAGraphNodeLauncher &Instance() {
+    static CUDAGraphNodeLauncher *launcher = new CUDAGraphNodeLauncher;
+    return *launcher;
+  }
+
+ private:
+  CUDAGraphNodeLauncher() : id(0) {}
+  DISABLE_COPY_AND_ASSIGN(CUDAGraphNodeLauncher);
+
+  unsigned int GenerateIdentifier() { return id++; }
+
+  unsigned int id;
+  std::unordered_map<cudaFunction_t, std::map<unsigned int, parameterSetter_t>>
+      parameterSetters;
+};
+
+static void ThrowErrorIfNotSupportCUDAGraph() {}
+
+using CUDAGraphID = unsigned long long;  // NOLINT
+
+// NOTE: Currently, we do not support to capture CUDA graph in parallel
+// NOTE: Do not use this class directly because it should be used with
+//       the memory pool.
+class CUDAGraph {
+  DISABLE_COPY_AND_ASSIGN(CUDAGraph);
+
+  // Since the constructor would throw error is CUDA_VERSION < 10010.
+  // The non-static method of CUDAGraph need not check CUDA_VERSION
+  // again.
+  explicit CUDAGraph(bool enable_replace = false)
+      : enable_replace_(enable_replace) {
+    ThrowErrorIfNotSupportCUDAGraph();
+    id_ = UniqueID();
+  }
+
+ public:
+  static constexpr int64_t kDefaultPoolID = 0;
+  static constexpr int64_t kInvalidPoolID = -1;
+
+  ~CUDAGraph() { Reset(); }
+
+  CUDAGraphID ID() const { return id_; }
+
+  static int64_t SetMemoryPoolID(int64_t pool_id) {
+    auto &pool_id_ = capturing_graph_->pool_id_;
+    PADDLE_ENFORCE_EQ(pool_id_,
+                      kInvalidPoolID,
+                      common::errors::InvalidArgument(
+                          "Cannot reset memory pool id twice, the "
+                          "former memory pool id is %d.",
+                          pool_id_));
+    if (pool_id <= kInvalidPoolID) {
+      pool_id_ = UniqueMemoryPoolID();
+    } else {
+      PADDLE_ENFORCE_GE(pool_id,
+                        kDefaultPoolID,
+                        common::errors::InvalidArgument(
+                            "Invalid memory pool id %d.", pool_id));
+      pool_id_ = pool_id;
+    }
+    return pool_id_;
+  }
+
+  int64_t PoolID() const { return pool_id_; }
+
+  static int64_t CapturingPoolID() { return capturing_graph_->pool_id_; }
+
+  void Replay();
+
+  void Reset();
+
+  void AddPostResetCallback(
+      std::function<void(paddle::optional<const CUDAGraph &>)> callback) {
+    std::lock_guard<std::mutex> guard(mtx_);
+    cudagraph_post_reset_callbacks_.push_back(std::move(callback));
+  }
+
+  static void AddPreCaptureCallback(std::function<void()> callback) {
+    cudagraph_pre_capture_callbacks_.push_back(std::move(callback));
+  }
+
+  void AddPostCaptureCallback(std::function<void()> callback) {
+    std::lock_guard<std::mutex> guard(mtx_);
+    cudagraph_post_capture_callbacks_.push_back(std::move(callback));
+  }
+
+  void AddJoiningStream(cudaStream_t stream) {
+    streams_to_join_.insert(stream);
+  }
+
+  void PrintToDotFiles(const std::string &dirname, unsigned int flags);
+
+  bool IsReplayed() const { return is_replayed_; }
+
+  static void BeginCapture(phi::GPUPlace place,
+                           cudaStream_t stream,
+                           gpuStreamCaptureMode mode,
+                           bool enable_replace = false);
+  static std::unique_ptr<CUDAGraph> EndCapture();
+
+  static void BeginSegmentCapture();
+  static void EndSegmentCapture();
+
+  static void AddJoiningStreamDuringCapturing(cudaStream_t stream) {
+    capturing_graph_->AddJoiningStream(stream);
+  }
+
+  static void AddPostResetCallbackDuringCapturing(
+      std::function<void(paddle::optional<const CUDAGraph &>)> callback) {
+    capturing_graph_->AddPostResetCallback(std::move(callback));
+  }
+
+  static void AddPostCaptureCallbackDuringCapturing(
+      std::function<void()> callback) {
+    capturing_graph_->AddPostCaptureCallback(std::move(callback));
+  }
+
+  // No need to add CUDA_VERSION macro because capturing_graph_ would
+  // always be nullptr (constructor throws error)
+  static bool IsCapturing() { return capturing_graph_ != nullptr; }
+
+  static CUDAGraphID CapturingID() { return capturing_graph_->id_; }
+
+  static phi::GPUPlace CapturingPlace() { return capturing_graph_->place_; }
+
+  // This API can be used to debug which GPU operation is not
+  // supported during capturing CUDA Graph.
+  static bool IsValidCapturing();
+
+  static bool IsThreadLocalCapturing() {
+    return IsCapturing() &&
+           capturing_graph_->capture_mode_ == cudaStreamCaptureModeThreadLocal;
+  }
+
+  static bool IsThisThreadCapturing() {
+    if (UNLIKELY(IsCapturing())) {
+      return IsThreadLocalCapturing()
+                 ? capturing_thread_id_.get() == std::this_thread::get_id()
+                 : true;
+    } else {
+      return false;
+    }
+  }
+
+  using SetSeedFunc = std::function<bool(gpuKernelParams *, bool)>;
+  static void RecordRandomKernelInfo(SetSeedFunc set_seed_func) {
+    std::lock_guard<std::mutex> guard(capturing_graph_->func_mtx_);
+    capturing_graph_->set_seed_funcs_.emplace_back(std::move(set_seed_func));
+  }
+
+  static int64_t UniqueMemoryPoolID();
+
+  void ReplaceInputPtrs(const std::vector<void *> &old_ptrs,
+                        const std::vector<void *> &new_ptrs);
+
+  struct KernelParamInfo {
+    size_t offset;
+    size_t size;
+  };
+
+  struct KernelNodeInfo {
+    cudaGraphNode_t node;
+    CUDA_KERNEL_NODE_PARAMS params;
+    std::vector<KernelParamInfo> param_infos;
+  };
+
+ private:
+  static CUDAGraphID UniqueID();
+  std::vector<KernelParamInfo> GetKernelParamInfos(CUfunction func);
+  void CacheKernelNodeInfos(size_t segment_idx);
+
+ private:
+  std::vector<cudaGraph_t> graphs_;
+  std::vector<cudaGraphExec_t> exec_graphs_;
+  std::vector<std::vector<KernelNodeInfo>> cached_kernel_nodes_;
+  gpuStreamCaptureMode capture_mode_;
+  cudaStream_t stream_{nullptr};
+  phi::GPUPlace place_;
+  CUDAGraphID id_;
+  int64_t pool_id_{kInvalidPoolID};
+  bool is_reset_{false};
+  bool is_replayed_{false};
+  bool enable_replace_{false};
+  std::mutex mtx_;
+
+  std::vector<SetSeedFunc> set_seed_funcs_;
+
+  std::unordered_set<cudaStream_t> streams_to_join_;
+
+  // Holds callbacks that are triggered after the CUDA graph is reset. These
+  // callbacks are used for operations that need to be performed following the
+  // reset of a CUDA graph.
+  std::vector<std::function<void(paddle::optional<const CUDAGraph &>)>>
+      cudagraph_post_reset_callbacks_;
+
+  static std::vector<std::function<void()>> cudagraph_pre_capture_callbacks_;
+
+  // Contains callbacks that are invoked after the CUDA graph has been captured.
+  // These callbacks are crucial for managing memory allocations related to the
+  // CUDA graph. They ensure that memory blocks not associated with a graph (as
+  // detailed in cuda_malloc_async_allocator) are not erroneously released
+  // during the graph's lifecycle.
+  std::vector<std::function<void()>> cudagraph_post_capture_callbacks_;
+
+  // Maintains a collection of 'pre-hooks' - functions that are executed before
+  // the CUDA graph is replayed. These pre-hooks are essential for setting up
+  // the necessary conditions or states required for the correct execution of
+  // the CUDA graph.
+  std::vector<std::vector<cudaGraphExecuterSetter_t>>
+      cudagraph_pre_replay_callbacks_;
+
+  std::mutex func_mtx_;
+
+  bool is_first_run_{true};
+
+  static paddle::optional<std::thread::id> capturing_thread_id_;
+  static std::unique_ptr<CUDAGraph> capturing_graph_;
+};
+
+class CUDAGraphCaptureModeGuard {
+  DISABLE_COPY_AND_ASSIGN(CUDAGraphCaptureModeGuard);
+
+ public:
+  explicit CUDAGraphCaptureModeGuard(
+      gpuStreamCaptureMode mode = cudaStreamCaptureModeRelaxed) {
+    if (UNLIKELY(CUDAGraph::IsCapturing())) {
+      PADDLE_ENFORCE_GPU_SUCCESS(cudaThreadExchangeStreamCaptureMode(&mode));
+      // After cudaThreadExchangeStreamCaptureMode is called,
+      // the variable "mode" would be set to the old capturing mode.
+      old_mode_ = mode;
+    }
+  }
+
+  ~CUDAGraphCaptureModeGuard() PADDLE_MAY_THROW {
+    if (UNLIKELY(CUDAGraph::IsCapturing())) {
+      PADDLE_ENFORCE_GPU_SUCCESS(
+          cudaThreadExchangeStreamCaptureMode(&old_mode_));
+    }
+  }
+
+ private:
+  gpuStreamCaptureMode old_mode_;
+};
+
+}  // namespace gpu
+}  // namespace backends
+}  // namespace phi
+
+#endif

paddle/phi/backends/gpu/cuda/cuda_graph_with_memory_pool.h

@@ -0,0 +1,88 @@
+// Copyright (c) 2021 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.
+
+#pragma once
+
+#include <cstddef>
+#include <utility>
+
+#if defined(PADDLE_WITH_CUDA) || defined(PADDLE_WITH_HIP) || \
+    defined(PADDLE_WITH_CUSTOM_DEVICE)
+#include "paddle/phi/backends/context_pool.h"
+#if defined(PADDLE_WITH_CUDA)
+#include "paddle/phi/backends/gpu/cuda/cuda_graph.h"
+#elif defined(PADDLE_WITH_HIP)
+#include "paddle/phi/backends/gpu/rocm/hip_graph.h"
+#elif defined(PADDLE_WITH_CUSTOM_DEVICE)
+#include "paddle/phi/backends/custom/cuda_graph.h"
+#endif
+#endif
+
+namespace phi {
+namespace backends {
+namespace gpu {
+
+inline bool IsCUDAGraphCapturing() {
+#if defined(PADDLE_WITH_CUDA) || defined(PADDLE_WITH_HIP) || \
+    defined(PADDLE_WITH_CUSTOM_DEVICE)
+  return CUDAGraph::IsCapturing();
+#else
+  return false;
+#endif
+}
+
+// Add reset callback if CUDA Graph is capturing.
+// Otherwise, invoke callback directly.
+template <typename Callback>
+inline void AddPostResetCallbackIfCapturingCUDAGraph(Callback &&callback) {
+#if defined(PADDLE_WITH_CUDA) || defined(PADDLE_WITH_HIP) || \
+    defined(PADDLE_WITH_CUSTOM_DEVICE)
+  if (UNLIKELY(IsCUDAGraphCapturing())) {
+    return CUDAGraph::AddPostResetCallbackDuringCapturing(
+        std::forward<Callback>(callback));
+  }
+#endif
+  callback({});
+}
+
+template <typename T>
+inline T *RestoreHostMemIfCapturingCUDAGraph(T *host_mem, size_t size) {
+  static_assert(std::is_trivial<T>::value, "T must be trivial type");
+  static_assert(!std::is_same<T, void>::value, "T cannot be void");
+#if defined(PADDLE_WITH_CUDA) || defined(PADDLE_WITH_HIP) || \
+    defined(PADDLE_WITH_CUSTOM_DEVICE)
+  if (UNLIKELY(IsCUDAGraphCapturing())) {
+    size_t nbytes = size * sizeof(T);
+    // NOTE: Use new[]/delete[] (plain heap) instead of cudaMallocHost /
+    // hipMallocHost here. cudaMallocHost and hipMallocHost are prohibited
+    // operations during CUDA/HIP Graph stream capture on CUDA 12.x / HIP,
+    // returning cudaErrorStreamCaptureUnsupported (error 900). Plain heap
+    // memory is safe to allocate at any time, and the captured
+    // cudaMemcpyAsync node records the host pointer address; the buffer
+    // lifetime is guaranteed by the post-reset callback below.
+    void *new_host_mem = new uint8_t[nbytes];
+    std::memcpy(new_host_mem, host_mem, nbytes);
+    AddPostResetCallbackIfCapturingCUDAGraph(
+        [=](paddle::optional<const CUDAGraph &> graph) {
+          delete[] reinterpret_cast<uint8_t *>(new_host_mem);
+        });
+    return reinterpret_cast<T *>(new_host_mem);
+  }
+#endif
+  return host_mem;
+}
+
+}  // namespace gpu
+}  // namespace backends
+}  // namespace phi

paddle/phi/backends/gpu/cuda/cuda_helper.h

@@ -0,0 +1,105 @@
+// Copyright (c) 2019 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.
+
+#pragma once
+
+#ifdef PADDLE_WITH_CUDA
+#include <cuda_runtime.h>  // NOLINT
+#include "paddle/phi/common/data_type.h"
+#include "paddle/phi/core/enforce.h"
+
+namespace phi {
+namespace backends {
+namespace gpu {
+
+/*
+ * Summary: Grid stride looping macro in CUDA kernel
+ *
+ *  [ Why need this macro? ]
+ *
+ *    The original looping in CUDA kernel is:
+ *
+ *    `for (int i = blockIdx.x * blockDim.x + threadIdx.x; i < (n); \
+ *        i += blockDim.x * gridDim.x)`
+ *
+ *    This for condition is risky. The value of `blockIdx.x * blockDim.x`
+ *    may be large, such as over 1GB, the first iteration is no problem here,
+ *    but when `i += blockDim.x * gridDim.x` is executed, the value of i
+ *    will greater than INT_MAX and overflow becomes negative value, at
+ *    this time, the cycle condition `i < (n)` is still satisfied, so it
+ *    will cause illegal access to cuda memory.
+ *
+ *    Here is a real example in ERNIE, it will trigger above error.
+ *    The related data are:
+ *      - blockIdx.x = 2172938
+ *      - blockDim.x = 512
+ *      - blockIdx.x * blockDim.x = 1112543864
+ *      - INT_MAX = 2147483647
+ *
+ *    So we polish the for condition as follow, the int64_t __index__ will
+ *    prevent overflow in the loop increment.
+ *
+ * Parameters:
+ *    - i: loop index
+ *    - num: total element numbers
+ *
+ * Examples:
+ *    template <typename T>
+ *    __global__ void Scale(T* logit_grad, const T* loss_grad, const int num,
+ *                      const int d, const int remain) {
+ *    CUDA_KERNEL_LOOP(index, num) {
+ *      int idx_n = index / d;
+ *      int idx_remain = index % remain;
+ *      logit_grad[index] *= loss_grad[idx_n * remain + idx_remain];
+ *      }
+ *    }
+ *
+ */
+
+#define CUDA_KERNEL_LOOP_TYPE(i, num, index_type)                    \
+  int64_t __index__ =                                                \
+      static_cast<int64_t>(blockIdx.x) * blockDim.x + threadIdx.x;   \
+  int64_t __stride__ = static_cast<int64_t>(blockDim.x) * gridDim.x; \
+  for (index_type i = __index__; __index__ < (num);                  \
+       __index__ += __stride__, i = __index__)
+
+template <typename T>
+cudaDataType_t ToCudaDataType() {
+  if (std::is_same<T, float>::value) {
+    return CUDA_R_32F;
+  } else if (std::is_same<T, double>::value) {
+    return CUDA_R_64F;
+  } else if (std::is_same<T, phi::dtype::float16>::value) {
+    return CUDA_R_16F;
+#if CUDA_VERSION >= 11000
+  } else if (std::is_same<T, phi::dtype::bfloat16>::value) {
+    return CUDA_R_16BF;
+#endif
+#if CUDA_VERSION >= 11060
+  } else if (std::is_same<T, int8_t>::value) {
+    return CUDA_R_8I;
+  } else if (std::is_same<T, int32_t>::value) {
+    return CUDA_R_32I;
+#endif
+  } else {
+    PADDLE_THROW(common::errors::InvalidArgument(
+        "DataType %s is unsupported for CUDA.",
+        DataTypeToString(phi::CppTypeToDataType<T>::Type())));
+  }
+}
+
+}  // namespace gpu
+}  // namespace backends
+}  // namespace phi
+#endif

paddle/phi/backends/gpu/cuda/cuda_info.cc

@@ -0,0 +1,365 @@
+// Copyright (c) 2022 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.
+#include <mutex>
+
+#include "paddle/phi/backends/gpu/gpu_info.h"
+
+#include "glog/logging.h"
+
+#include "paddle/phi/core/enforce.h"
+
+static std::once_flag g_device_props_size_init_flag;
+static std::vector<std::unique_ptr<std::once_flag>> g_device_props_init_flags;
+static std::vector<phi::gpuDeviceProp> g_device_props;
+
+namespace phi::backends::gpu {
+
+#ifndef PADDLE_WITH_CUSTOM_DEVICE
+int DnnVersion() {
+  if (!dynload::HasCUDNN()) return -1;
+  return dynload::cudnnGetVersion();  // NOLINT
+}
+#endif
+
+static int GetGPUDeviceCountImpl() {
+  int driverVersion = 0;
+  cudaError_t status = cudaDriverGetVersion(&driverVersion);
+
+  if (!(status == gpuSuccess && driverVersion != 0)) {
+    // No GPU driver
+    VLOG(2) << "GPU Driver Version can't be detected. No GPU driver!";
+    return 0;
+  }
+
+  const auto *cuda_visible_devices = std::getenv("CUDA_VISIBLE_DEVICES");
+
+  if (cuda_visible_devices != nullptr) {
+    std::string cuda_visible_devices_str(cuda_visible_devices);
+    if (!cuda_visible_devices_str.empty()) {
+      cuda_visible_devices_str.erase(
+          0, cuda_visible_devices_str.find_first_not_of('\''));
+      cuda_visible_devices_str.erase(
+          cuda_visible_devices_str.find_last_not_of('\'') + 1);
+      cuda_visible_devices_str.erase(
+          0, cuda_visible_devices_str.find_first_not_of('\"'));
+      cuda_visible_devices_str.erase(
+          cuda_visible_devices_str.find_last_not_of('\"') + 1);
+    }
+    if (std::all_of(cuda_visible_devices_str.begin(),
+                    cuda_visible_devices_str.end(),
+                    [](char ch) { return ch == ' '; })) {
+      VLOG(2) << "CUDA_VISIBLE_DEVICES is set to be "
+                 "empty. No GPU detected.";
+      return 0;
+    }
+  }
+  int count;
+  status = cudaGetDeviceCount(&count);
+  if (status != cudaSuccess) {
+    VLOG(2) << "You have gpu driver and cuda installed, but the machine not "
+               "has any gpu card.";
+    count = 0;
+  }
+  return count;
+}
+
+int GetGPUDeviceCount() {
+  // cache the count
+  static auto dev_cnt = GetGPUDeviceCountImpl();
+  return dev_cnt;
+}
+
+int GetGPUComputeCapability(int id) {
+  PADDLE_ENFORCE_LT(id,
+                    GetGPUDeviceCount(),
+                    common::errors::InvalidArgument(
+                        "Device id must be less than GPU count, "
+                        "but received id is: %d. GPU count is: %d.",
+                        id,
+                        GetGPUDeviceCount()));
+  int major, minor;
+  auto major_error_code =
+      cudaDeviceGetAttribute(&major, cudaDevAttrComputeCapabilityMajor, id);
+  auto minor_error_code =
+      cudaDeviceGetAttribute(&minor, cudaDevAttrComputeCapabilityMinor, id);
+
+  PADDLE_ENFORCE_GPU_SUCCESS(major_error_code);
+  PADDLE_ENFORCE_GPU_SUCCESS(minor_error_code);
+  return major * 10 + minor;
+}
+
+int GetGPURuntimeVersion(int id) {
+  PADDLE_ENFORCE_LT(id,
+                    GetGPUDeviceCount(),
+                    common::errors::InvalidArgument(
+                        "Device id must be less than GPU count, "
+                        "but received id is: %d. GPU count is: %d.",
+                        id,
+                        GetGPUDeviceCount()));
+  int runtime_version = 0;
+  PADDLE_ENFORCE_GPU_SUCCESS(cudaRuntimeGetVersion(&runtime_version));
+  return runtime_version;
+}
+
+int GetGPUDriverVersion(int id) {
+  PADDLE_ENFORCE_LT(id,
+                    GetGPUDeviceCount(),
+                    common::errors::InvalidArgument(
+                        "Device id must be less than GPU count, "
+                        "but received id is: %d. GPU count is: %d.",
+                        id,
+                        GetGPUDeviceCount()));
+  int driver_version = 0;
+  PADDLE_ENFORCE_GPU_SUCCESS(cudaDriverGetVersion(&driver_version));
+  return driver_version;
+}
+
+bool TensorCoreAvailable() {
+  int device = GetCurrentDeviceId();
+  int driver_version = GetGPUComputeCapability(device);
+  return driver_version >= 70;
+}
+
+int GetGPUMultiProcessors(int id) {
+  PADDLE_ENFORCE_LT(id,
+                    GetGPUDeviceCount(),
+                    common::errors::InvalidArgument(
+                        "Device id must be less than GPU count, "
+                        "but received id is: %d. GPU count is: %d.",
+                        id,
+                        GetGPUDeviceCount()));
+  int count;
+  PADDLE_ENFORCE_GPU_SUCCESS(
+      cudaDeviceGetAttribute(&count, cudaDevAttrMultiProcessorCount, id));
+  return count;
+}
+
+int GetGPUMaxThreadsPerMultiProcessor(int id) {
+  PADDLE_ENFORCE_LT(id,
+                    GetGPUDeviceCount(),
+                    common::errors::InvalidArgument(
+                        "Device id must be less than GPU count, "
+                        "but received id is: %d. GPU count is: %d.",
+                        id,
+                        GetGPUDeviceCount()));
+  int count;
+  PADDLE_ENFORCE_GPU_SUCCESS(cudaDeviceGetAttribute(
+      &count, cudaDevAttrMaxThreadsPerMultiProcessor, id));
+
+  return count;
+}
+
+int GetGPUMaxThreadsPerBlock(int id) {
+  PADDLE_ENFORCE_LT(id,
+                    GetGPUDeviceCount(),
+                    common::errors::InvalidArgument(
+                        "Device id must be less than GPU count, "
+                        "but received id is: %d. GPU count is: %d.",
+                        id,
+                        GetGPUDeviceCount()));
+  int count;
+  PADDLE_ENFORCE_GPU_SUCCESS(
+      cudaDeviceGetAttribute(&count, cudaDevAttrMaxThreadsPerBlock, id));
+  return count;
+}
+
+int GetCurrentDeviceId() {
+  int device_id;
+  PADDLE_ENFORCE_GPU_SUCCESS(cudaGetDevice(&device_id));
+  return device_id;
+}
+
+std::array<unsigned int, 3> GetGpuMaxGridDimSize(int id) {
+  PADDLE_ENFORCE_LT(id,
+                    GetGPUDeviceCount(),
+                    common::errors::InvalidArgument(
+                        "Device id must be less than GPU count, "
+                        "but received id is: %d. GPU count is: %d.",
+                        id,
+                        GetGPUDeviceCount()));
+  std::array<unsigned int, 3> ret = {};
+  int size;
+  auto error_code_x = cudaDeviceGetAttribute(&size, cudaDevAttrMaxGridDimX, id);
+  PADDLE_ENFORCE_GPU_SUCCESS(error_code_x);
+  ret[0] = size;
+
+  auto error_code_y = cudaDeviceGetAttribute(&size, cudaDevAttrMaxGridDimY, id);
+  PADDLE_ENFORCE_GPU_SUCCESS(error_code_y);
+  ret[1] = size;
+
+  auto error_code_z = cudaDeviceGetAttribute(&size, cudaDevAttrMaxGridDimZ, id);
+  PADDLE_ENFORCE_GPU_SUCCESS(error_code_z);
+  ret[2] = size;
+  return ret;
+}
+
+std::pair<int, int> GetGpuStreamPriorityRange() {
+  int least_priority, greatest_priority;
+  PADDLE_ENFORCE_GPU_SUCCESS(
+      cudaDeviceGetStreamPriorityRange(&least_priority, &greatest_priority));
+  return std::make_pair(least_priority, greatest_priority);
+}
+
+const gpuDeviceProp &GetDeviceProperties(int id) {
+  std::call_once(g_device_props_size_init_flag, [&] {
+    int gpu_num = 0;
+    gpu_num = GetGPUDeviceCount();
+    g_device_props_init_flags.resize(gpu_num);
+    g_device_props.resize(gpu_num);
+    for (int i = 0; i < gpu_num; ++i) {
+      g_device_props_init_flags[i] = std::make_unique<std::once_flag>();
+    }
+  });
+
+  if (id == -1) {
+    id = GetCurrentDeviceId();
+  }
+
+  if (id < 0 || id >= static_cast<int>(g_device_props.size())) {
+    PADDLE_THROW(common::errors::OutOfRange(
+        "The device id %d is out of range [0, %d), where %d is the number of "
+        "devices on this machine. Because the device id should be greater than "
+        "or equal to zero and smaller than the number of gpus. Please input "
+        "appropriate device again!",
+        id,
+        static_cast<int>(g_device_props.size()),
+        static_cast<int>(g_device_props.size())));
+  }
+
+  std::call_once(*(g_device_props_init_flags[id]), [&] {
+    PADDLE_ENFORCE_GPU_SUCCESS(
+        cudaGetDeviceProperties(&g_device_props[id], id));
+  });
+
+  return g_device_props[id];
+}
+
+void SetDeviceId(int id) {
+  static thread_local bool first_call = true;
+  if (first_call) {
+    PADDLE_ENFORCE_LT(id,
+                      GetGPUDeviceCount(),
+                      common::errors::InvalidArgument(
+                          "Device id must be less than GPU count, "
+                          "but received id is: %d. GPU count is: %d.",
+                          id,
+                          GetGPUDeviceCount()));
+
+    PADDLE_RETRY_CUDA_SUCCESS(cudaSetDevice(id));
+    VLOG(4) << "SetDeviceId " << id;
+    first_call = false;
+    return;
+  }
+
+  int prev_id;
+  PADDLE_ENFORCE_GPU_SUCCESS(cudaGetDevice(&prev_id));
+  if (prev_id != id) {
+    PADDLE_ENFORCE_LT(id,
+                      GetGPUDeviceCount(),
+                      common::errors::InvalidArgument(
+                          "Device id must be less than GPU count, "
+                          "but received id is: %d. GPU count is: %d.",
+                          id,
+                          GetGPUDeviceCount()));
+
+    PADDLE_RETRY_CUDA_SUCCESS(cudaSetDevice(id));
+    VLOG(4) << "SetDeviceId " << id;
+  }
+}
+
+void GpuMemcpyAsync(void *dst,
+                    const void *src,
+                    size_t count,
+                    gpuMemcpyKind kind,
+                    gpuStream_t stream) {
+  PADDLE_ENFORCE_GPU_SUCCESS(cudaMemcpyAsync(dst, src, count, kind, stream));
+}
+
+void GpuMemcpySync(void *dst,
+                   const void *src,
+                   size_t count,
+                   gpuMemcpyKind kind) {
+  PADDLE_ENFORCE_GPU_SUCCESS(cudaMemcpy(dst, src, count, kind));
+}
+
+void GpuMemcpyPeerAsync(void *dst,
+                        int dst_device,
+                        const void *src,
+                        int src_device,
+                        size_t count,
+                        gpuStream_t stream) {
+  PADDLE_ENFORCE_GPU_SUCCESS(
+      cudaMemcpyPeerAsync(dst, dst_device, src, src_device, count, stream));
+}
+
+void GpuMemcpyPeerSync(
+    void *dst, int dst_device, const void *src, int src_device, size_t count) {
+  PADDLE_ENFORCE_GPU_SUCCESS(
+      cudaMemcpyPeer(dst, dst_device, src, src_device, count));
+}
+
+void GpuMemsetAsync(void *dst, int value, size_t count, gpuStream_t stream) {
+  PADDLE_ENFORCE_GPU_SUCCESS(cudaMemsetAsync(dst, value, count, stream));
+}
+
+void GpuStreamSync(gpuStream_t stream) {
+  PADDLE_ENFORCE_GPU_SUCCESS(cudaStreamSynchronize(stream));
+}
+
+void GpuDestroyStream(gpuStream_t stream) {
+  PADDLE_ENFORCE_GPU_SUCCESS(cudaStreamDestroy(stream));
+}
+
+void GpuDeviceSync() { PADDLE_ENFORCE_GPU_SUCCESS(cudaDeviceSynchronize()); }
+
+gpuError_t GpuGetLastError() { return cudaGetLastError(); }
+
+// See
+// https://docs.nvidia.com/cuda/cuda-c-programming-guide/index.html#um-requirements
+// for more detail about managed memory requirements
+bool IsGPUManagedMemorySupported(int dev_id) {
+  PADDLE_ENFORCE_LT(dev_id,
+                    GetGPUDeviceCount(),
+                    common::errors::InvalidArgument(
+                        "Device id must be less than GPU count, "
+                        "but received id is: %d. GPU count is: %d.",
+                        dev_id,
+                        GetGPUDeviceCount()));
+#if defined(__linux__) || defined(_WIN32)
+  int ManagedMemoryAttr;
+  PADDLE_ENFORCE_GPU_SUCCESS(cudaDeviceGetAttribute(
+      &ManagedMemoryAttr, cudaDevAttrManagedMemory, dev_id));
+  return ManagedMemoryAttr != 0;
+#else
+  return false;
+#endif
+}
+
+bool IsGPUManagedMemoryOversubscriptionSupported(int dev_id) {
+  PADDLE_ENFORCE_LT(dev_id,
+                    GetGPUDeviceCount(),
+                    common::errors::InvalidArgument(
+                        "Device id must be less than GPU count, "
+                        "but received id is: %d. GPU count is: %d.",
+                        dev_id,
+                        GetGPUDeviceCount()));
+#ifdef __linux__
+  return IsGPUManagedMemorySupported(dev_id) &&
+         GetGPUComputeCapability(dev_id) >= 60;
+#else
+  return false;
+#endif
+}
+
+}  // namespace phi::backends::gpu

paddle/phi/backends/gpu/cuda/cudnn_desc.h

@@ -0,0 +1,313 @@
+// Copyright (c) 2022 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.
+
+#pragma once
+
+#include <algorithm>
+#include <functional>
+#include <iostream>
+#include <iterator>
+#include <memory>
+#include <numeric>
+#include <string>
+#include <vector>
+
+#include "paddle/phi/backends/gpu/cuda/cudnn_helper.h"
+#include "paddle/phi/core/utils/data_type.h"
+
+namespace phi {
+namespace backends {
+namespace gpu {
+
+template <typename T>
+inline std::vector<T> TransformDimOrder(const std::vector<T>& dims) {
+  std::vector<T> transformed_dims(dims.begin(), dims.end());
+  if (dims.size() < 4) {
+    return transformed_dims;
+  }
+  T H, W, D, C;
+  if (dims.size() == 4) {
+    H = dims[1];
+    W = dims[2];
+    C = dims[3];
+    transformed_dims[1] = C;
+    transformed_dims[2] = H;
+    transformed_dims[3] = W;
+  } else {
+    D = dims[1];
+    H = dims[2];
+    W = dims[3];
+    C = dims[4];
+    transformed_dims[1] = C;
+    transformed_dims[2] = D;
+    transformed_dims[3] = H;
+    transformed_dims[4] = W;
+  }
+  return transformed_dims;
+}
+
+inline cudnnDataType_t ToCudnnDataType(const DataType& t) {
+  cudnnDataType_t type = CUDNN_DATA_FLOAT;
+  switch (t) {
+    case DataType::FLOAT16:
+      type = CUDNN_DATA_HALF;
+      break;
+    case DataType::FLOAT32:
+      type = CUDNN_DATA_FLOAT;
+      break;
+    case DataType::FLOAT64:
+      type = CUDNN_DATA_DOUBLE;
+      break;
+#if CUDNN_VERSION_MIN(8, 6, 0) && CUDA_VERSION >= 11080
+    case DataType::FLOAT8_E4M3FN:
+      type = CUDNN_DATA_FP8_E4M3;
+      break;
+    case DataType::FLOAT8_E5M2:
+      type = CUDNN_DATA_FP8_E5M2;
+      break;
+#endif
+#if CUDNN_VERSION_MIN(8, 1, 0)
+    case DataType::BFLOAT16:
+      type = CUDNN_DATA_BFLOAT16;
+      break;
+#endif
+    default:
+      break;
+  }
+  return type;
+}
+
+class ActivationDescriptor {
+ public:
+  using T = cudnnActivationStruct;
+  struct Deleter {
+    void operator()(T* t) {
+      if (t != nullptr) {
+        PADDLE_ENFORCE_GPU_SUCCESS(
+            phi::dynload::cudnnDestroyActivationDescriptor(t));
+        t = nullptr;
+      }
+    }
+  };
+  ActivationDescriptor() {
+    T* raw_ptr;
+    PADDLE_ENFORCE_GPU_SUCCESS(
+        phi::dynload::cudnnCreateActivationDescriptor(&raw_ptr));
+    desc_.reset(raw_ptr);
+  }
+  template <typename T>
+  void set(cudnnActivationMode_t mode, const T& coef) {
+    PADDLE_ENFORCE_GPU_SUCCESS(phi::dynload::cudnnSetActivationDescriptor(
+        desc_.get(), mode, CUDNN_NOT_PROPAGATE_NAN, static_cast<double>(coef)));
+  }
+
+  T* desc() { return desc_.get(); }
+  T* desc() const { return desc_.get(); }
+
+ private:
+  std::unique_ptr<T, Deleter> desc_;
+};
+
+class TensorDescriptor {
+ public:
+  using T = cudnnTensorStruct;
+  struct Deleter {
+    void operator()(T* t) {
+      if (t != nullptr) {
+        PADDLE_ENFORCE_GPU_SUCCESS(
+            phi::dynload::cudnnDestroyTensorDescriptor(t));
+        t = nullptr;
+      }
+    }
+  };
+  TensorDescriptor() {
+    T* raw_ptr;
+    PADDLE_ENFORCE_GPU_SUCCESS(
+        phi::dynload::cudnnCreateTensorDescriptor(&raw_ptr));
+    desc_.reset(raw_ptr);
+  }
+  T* desc() { return desc_.get(); }
+  T* desc() const { return desc_.get(); }
+  void set(const DenseTensor& tensor, const int groups = 1) {
+    auto dims = common::vectorize<int>(tensor.dims());
+    std::vector<int> strides(dims.size());
+    strides[dims.size() - 1] = 1;
+    for (int i = dims.size() - 2; i >= 0; i--) {
+      strides[i] = dims[i + 1] * strides[i + 1];
+    }
+    std::vector<int> dims_with_group(dims.begin(), dims.end());
+    if (groups > 1) {
+      dims_with_group[1] = dims_with_group[1] / groups;
+    }
+    PADDLE_ENFORCE_GPU_SUCCESS(phi::dynload::cudnnSetTensorNdDescriptor(
+        desc_.get(),
+        ToCudnnDataType(tensor.dtype()),
+        dims_with_group.size(),
+        dims_with_group.data(),
+        strides.data()));
+  }
+
+  void set(const std::vector<int>& dims,
+           const cudnnTensorFormat_t format,
+           const cudnnDataType_t dtype) {
+    std::vector<int> transformed_dims;
+    if (format == CUDNN_TENSOR_NHWC) {
+      transformed_dims = TransformDimOrder(dims);
+    } else {
+      transformed_dims = dims;
+    }
+    PADDLE_ENFORCE_GPU_SUCCESS(
+        phi::dynload::cudnnSetTensorNdDescriptorEx(desc_.get(),
+                                                   format,
+                                                   dtype,
+                                                   transformed_dims.size(),
+                                                   transformed_dims.data()));
+  }
+
+  void set(const DenseTensor& tensor, const cudnnTensorFormat_t format) {
+    auto dims = common::vectorize<int>(tensor.dims());
+    auto dtype = ToCudnnDataType(tensor.dtype());
+    set(dims, format, dtype);
+  }
+
+ private:
+  std::unique_ptr<T, Deleter> desc_;
+};
+
+class FilterDescriptor {
+ public:
+  using T = cudnnFilterStruct;
+  struct Deleter {
+    void operator()(T* t) {
+      if (t != nullptr) {
+        PADDLE_ENFORCE_GPU_SUCCESS(
+            phi::dynload::cudnnDestroyFilterDescriptor(t));
+        t = nullptr;
+      }
+    }
+  };
+  FilterDescriptor() {
+    T* raw_ptr;
+    PADDLE_ENFORCE_GPU_SUCCESS(
+        phi::dynload::cudnnCreateFilterDescriptor(&raw_ptr));
+    desc_.reset(raw_ptr);
+  }
+  T* desc() { return desc_.get(); }
+  T* desc() const { return desc_.get(); }
+
+  void set(const std::vector<int>& dims,
+           const cudnnTensorFormat_t format,
+           const cudnnDataType_t dtype,
+           const int groups = 1) {
+    std::vector<int> transformed_dims;
+    if (format == CUDNN_TENSOR_NHWC) {
+      transformed_dims = TransformDimOrder(dims);
+    } else {
+      transformed_dims = dims;
+    }
+    if (groups > 1) {
+      transformed_dims[1] = transformed_dims[1] / groups;
+    }
+    PADDLE_ENFORCE_GPU_SUCCESS(
+        phi::dynload::cudnnSetFilterNdDescriptor(desc_.get(),
+                                                 dtype,
+                                                 format,
+                                                 transformed_dims.size(),
+                                                 transformed_dims.data()));
+  }
+
+  void set(const DenseTensor& tensor,
+           const cudnnTensorFormat_t format,
+           const int groups = 1) {
+    auto dims = common::vectorize<int>(tensor.dims());
+    auto dtype = ToCudnnDataType(tensor.dtype());
+    set(dims, format, dtype, groups);
+  }
+
+ private:
+  std::unique_ptr<T, Deleter> desc_;
+};
+
+class ConvolutionDescriptor {
+ public:
+  using T = cudnnConvolutionStruct;
+  struct Deleter {
+    void operator()(T* t) {
+      if (t != nullptr) {
+        PADDLE_ENFORCE_GPU_SUCCESS(
+            phi::dynload::cudnnDestroyConvolutionDescriptor(t));
+        t = nullptr;
+      }
+    }
+  };
+  ConvolutionDescriptor() {
+    T* raw_ptr;
+    PADDLE_ENFORCE_GPU_SUCCESS(
+        phi::dynload::cudnnCreateConvolutionDescriptor(&raw_ptr));
+    desc_.reset(raw_ptr);
+  }
+  T* desc() { return desc_.get(); }
+  T* desc() const { return desc_.get(); }
+
+  void set(cudnnDataType_t dtype,
+           const std::vector<int>& pads,
+           const std::vector<int>& strides,
+           const std::vector<int>& dilations,
+           bool allow_tf32,
+           const int groups = 1) {
+    allow_tf32_ = allow_tf32;
+    cudnnDataType_t compute_type =
+        (dtype == CUDNN_DATA_DOUBLE) ? CUDNN_DATA_DOUBLE : CUDNN_DATA_FLOAT;
+    T* desc = desc_.get();
+    PADDLE_ENFORCE_GPU_SUCCESS(
+        phi::dynload::cudnnSetConvolutionNdDescriptor(desc,
+                                                      pads.size(),
+                                                      pads.data(),
+                                                      strides.data(),
+                                                      dilations.data(),
+                                                      CUDNN_CROSS_CORRELATION,
+                                                      compute_type));
+#if CUDNN_VERSION_MIN(7, 0, 1)
+    PADDLE_ENFORCE_GPU_SUCCESS(
+        phi::dynload::cudnnSetConvolutionGroupCount(desc, groups));
+#if CUDA_VERSION >= 9000 && CUDNN_VERSION_MIN(7, 0, 1)
+    PADDLE_ENFORCE_GPU_SUCCESS(
+        phi::dynload::cudnnSetConvolutionMathType(desc, CUDNN_DEFAULT_MATH));
+    if (dtype == CUDNN_DATA_HALF) {
+      PADDLE_ENFORCE_GPU_SUCCESS(phi::dynload::cudnnSetConvolutionMathType(
+          desc, CUDNN_TENSOR_OP_MATH));
+#if CUDA_VERSION >= 11000
+#if CUDNN_VERSION_MIN(8, 1, 0)
+    } else if (dtype == CUDNN_DATA_BFLOAT16) {
+      PADDLE_ENFORCE_GPU_SUCCESS(phi::dynload::cudnnSetConvolutionMathType(
+          desc, CUDNN_TENSOR_OP_MATH));
+#endif  // CUDNN_VERSION_MIN(8,1,0)
+    } else if (dtype == CUDNN_DATA_FLOAT && !allow_tf32) {
+      PADDLE_ENFORCE_GPU_SUCCESS(
+          phi::dynload::cudnnSetConvolutionMathType(desc, CUDNN_FMA_MATH));
+#endif  // CUDA_VERSION >= 11000
+    }
+#endif
+#endif
+  }
+
+  bool allow_tf32_;
+
+ private:
+  std::unique_ptr<T, Deleter> desc_;
+};
+
+}  // namespace gpu
+}  // namespace backends
+}  // namespace phi

paddle/phi/backends/gpu/cuda/cudnn_helper.h

@@ -0,0 +1,676 @@
+/* Copyright (c) 2022 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. */
+
+#pragma once
+
+#include <string>
+#include <vector>
+
+#include "paddle/common/errors.h"
+#include "paddle/common/flags.h"
+#include "paddle/common/macros.h"
+#include "paddle/phi/backends/dynload/cudnn.h"
+#include "paddle/phi/common/bfloat16.h"
+#include "paddle/phi/common/float16.h"
+#include "paddle/phi/common/place.h"
+#include "paddle/phi/core/dense_tensor.h"
+#include "paddle/phi/core/enforce.h"
+
+COMMON_DECLARE_bool(cudnn_deterministic);
+
+namespace phi {
+namespace backends {
+namespace gpu {
+
+#define CUDNN_VERSION_COMPUTE(major, minor, patch)     \
+  ((major) <= 8 ? (major)*1000 + (minor)*100 + (patch) \
+                : (major)*10000 + (minor)*100 + (patch))
+
+#define CUDNN_VERSION_MIN(major, minor, patch) \
+  (CUDNN_VERSION >= CUDNN_VERSION_COMPUTE(major, minor, patch))
+
+enum class PoolingMode {
+  kMaximum,
+  kMaximumDeterministic,
+  kAverageExclusive,
+  kAverageInclusive,
+};
+
+enum class ActivationMode {
+  kNone,  // activation identity
+  kSigmoid,
+  kRelu,
+  kRelu6,
+  kReluX,
+  kTanh,
+  kBandPass,
+};
+
+inline cudnnPoolingMode_t GetPoolingMode(const PoolingMode& mode) {
+  switch (mode) {
+    case PoolingMode::kMaximumDeterministic:
+      return CUDNN_POOLING_MAX_DETERMINISTIC;
+    case PoolingMode::kAverageExclusive:
+      return CUDNN_POOLING_AVERAGE_COUNT_EXCLUDE_PADDING;
+    case PoolingMode::kAverageInclusive:
+      return CUDNN_POOLING_AVERAGE_COUNT_INCLUDE_PADDING;
+    case PoolingMode::kMaximum:
+      return CUDNN_POOLING_MAX;
+    default:
+      PADDLE_THROW(
+          common::errors::Unimplemented("Unexpected CUDNN pooling mode."));
+  }
+}
+
+inline ActivationMode StringToActivationMode(const std::string& str) {
+  if (str == "identity") {
+    return ActivationMode::kNone;
+  } else if (str == "sigmoid") {
+    return ActivationMode::kSigmoid;
+  } else if (str == "relu") {
+    return ActivationMode::kRelu;
+  } else if (str == "relu6") {
+    return ActivationMode::kRelu6;
+  } else if (str == "relux") {
+    return ActivationMode::kReluX;
+  } else if (str == "tanh") {
+    return ActivationMode::kTanh;
+  } else if (str == "bandpass") {
+    return ActivationMode::kBandPass;
+  } else {
+    PADDLE_THROW(common::errors::Unimplemented(
+        "Unknown CUDNN activation string: %s.", str));
+  }
+}
+
+template <typename T>
+class CudnnDataType;
+
+// CUDNN_DATA_FLOAT8 is not valid before cudnn8.6
+#if CUDNN_VERSION_MIN(8, 6, 0) && CUDA_VERSION >= 11080
+template <>
+class CudnnDataType<phi::dtype::float8_e4m3fn> {
+ public:
+  static const cudnnDataType_t type = CUDNN_DATA_FP8_E4M3;
+  using ScalingParamType = const float;
+  using BatchNormParamType = float;
+  static ScalingParamType* kOne() {
+    static ScalingParamType v = 1.0;
+    return &v;
+  }
+  static ScalingParamType* kZero() {
+    static ScalingParamType v = 0.0;
+    return &v;
+  }
+};
+#endif
+
+// CUDNN_DATA_BFLOAT16 is not valid before cudnn8.1
+#if CUDNN_VERSION_MIN(8, 1, 0)
+template <>
+class CudnnDataType<phi::dtype::bfloat16> {
+ public:
+  static const cudnnDataType_t type = CUDNN_DATA_BFLOAT16;
+  using ScalingParamType = const float;
+  using BatchNormParamType = float;
+  static ScalingParamType* kOne() {
+    static ScalingParamType v = 1.0;
+    return &v;
+  }
+  static ScalingParamType* kZero() {
+    static ScalingParamType v = 0.0;
+    return &v;
+  }
+};
+#endif
+
+template <>
+class CudnnDataType<phi::dtype::float16> {
+ public:
+  static const cudnnDataType_t type = CUDNN_DATA_HALF;
+  // The scaling param type is float for HALF and FLOAT tensors
+  using ScalingParamType = const float;
+  using BatchNormParamType = float;
+  static ScalingParamType* kOne() {
+    static ScalingParamType v = 1.0;
+    return &v;
+  }
+  static ScalingParamType* kZero() {
+    static ScalingParamType v = 0.0;
+    return &v;
+  }
+};
+
+template <>
+class CudnnDataType<float> {
+ public:
+  static const cudnnDataType_t type = CUDNN_DATA_FLOAT;
+  using ScalingParamType = const float;
+  using BatchNormParamType = float;
+  static ScalingParamType* kOne() {
+    static ScalingParamType v = 1.0;
+    return &v;
+  }
+  static ScalingParamType* kZero() {
+    static ScalingParamType v = 0.0;
+    return &v;
+  }
+};
+
+template <>
+class CudnnDataType<double> {
+ public:
+  static const cudnnDataType_t type = CUDNN_DATA_DOUBLE;
+  using ScalingParamType = const double;
+  using BatchNormParamType = double;
+  static ScalingParamType* kOne() {
+    static ScalingParamType v = 1.0;
+    return &v;
+  }
+  static ScalingParamType* kZero() {
+    static ScalingParamType v = 0.0;
+    return &v;
+  }
+};
+
+inline cudnnTensorFormat_t GetCudnnTensorFormat(
+    const DataLayout& order) {  // Not use
+  switch (order) {
+    case DataLayout::NHWC:
+      return CUDNN_TENSOR_NHWC;
+    case DataLayout::NCHW:
+      return CUDNN_TENSOR_NCHW;
+    case DataLayout::NCDHW:
+      return CUDNN_TENSOR_NCHW;  // NOTE: cudnn treat NdTensor as the same
+    case DataLayout::NDHWC:
+      return CUDNN_TENSOR_NHWC;  // add, liyamei
+    default:
+      PADDLE_THROW(common::errors::Unimplemented(
+          "CUDNN has no equivalent dataLayout for input order."));
+  }
+  return CUDNN_TENSOR_NCHW;
+}
+
+class ScopedTensorDescriptor {
+ public:
+  ScopedTensorDescriptor() {
+    PADDLE_ENFORCE_GPU_SUCCESS(
+        phi::dynload::cudnnCreateTensorDescriptor(&desc_));
+  }
+  ~ScopedTensorDescriptor() PADDLE_MAY_THROW {
+    PADDLE_ENFORCE_GPU_SUCCESS(
+        phi::dynload::cudnnDestroyTensorDescriptor(desc_));
+  }
+
+  inline cudnnTensorDescriptor_t descriptor(const cudnnTensorFormat_t format,
+                                            const cudnnDataType_t type,
+                                            const std::vector<int>& dims,
+                                            const int groups = 1) {
+    // the format is not used now, will add later
+    std::vector<int> strides(dims.size());
+    strides[dims.size() - 1] = 1;
+    for (int i = dims.size() - 2; i >= 0; i--) {
+      strides[i] = dims[i + 1] * strides[i + 1];
+    }
+    // Update tensor descriptor dims setting if groups > 1
+    // NOTE: Here, Assume using NCHW or NCDHW order
+    std::vector<int> dims_with_group(dims.begin(), dims.end());
+    if (groups > 1) {
+      dims_with_group[1] = dims_with_group[1] / groups;
+    }
+
+    if (dims.size() == 4) {
+      if (format == CUDNN_TENSOR_NCHW) {
+        PADDLE_ENFORCE_GPU_SUCCESS(
+            phi::dynload::cudnnSetTensorNdDescriptor(desc_,
+                                                     type,
+                                                     dims_with_group.size(),
+                                                     dims_with_group.data(),
+                                                     strides.data()));
+      } else {  // CUDNN_TENSOR_NHWC
+        PADDLE_ENFORCE_GPU_SUCCESS(phi::dynload::cudnnSetTensor4dDescriptor(
+            desc_, format, type, dims[0], dims[3], dims[1], dims[2]));
+      }
+    } else if (dims.size() == 5) {
+      if (format == CUDNN_TENSOR_NCHW) {
+        PADDLE_ENFORCE_GPU_SUCCESS(
+            phi::dynload::cudnnSetTensorNdDescriptor(desc_,
+                                                     type,
+                                                     dims_with_group.size(),
+                                                     dims_with_group.data(),
+                                                     strides.data()));
+      } else {  // CUDNN_TENSOR_NHWC
+        PADDLE_ENFORCE_GPU_SUCCESS(phi::dynload::cudnnSetTensorNdDescriptorEx(
+            desc_, format, type, dims.size(), dims.data()));
+      }
+    }
+    return desc_;
+  }
+
+  template <typename T>
+  inline cudnnTensorDescriptor_t descriptor(const DataLayout& order,
+                                            const std::vector<int>& dims,
+                                            const int groups = 1) {
+    return descriptor(
+        GetCudnnTensorFormat(order), CudnnDataType<T>::type, dims, groups);
+  }
+
+  inline cudnnTensorDescriptor_t descriptor(const cudnnDataType_t cudnn_type,
+                                            const std::vector<int>& dim,
+                                            const std::vector<int>& stride) {
+    PADDLE_ENFORCE_GPU_SUCCESS(phi::dynload::cudnnSetTensorNdDescriptor(
+        desc_, cudnn_type, dim.size(), dim.data(), stride.data()));
+    return desc_;
+  }
+
+  template <typename T>
+  inline cudnnTensorDescriptor_t descriptor(const std::vector<int>& dim,
+                                            const std::vector<int>& stride) {
+    return descriptor(CudnnDataType<T>::type, dim, stride);
+  }
+
+  inline cudnnTensorDescriptor_t desc() { return desc_; }
+
+ private:
+  cudnnTensorDescriptor_t desc_;
+  DISABLE_COPY_AND_ASSIGN(ScopedTensorDescriptor);
+};
+
+class ScopedRNNTensorDescriptor {
+ public:
+  ScopedRNNTensorDescriptor() {
+    PADDLE_ENFORCE_GPU_SUCCESS(
+        phi::dynload::cudnnCreateRNNDataDescriptor(&desc_));
+  }
+
+  ~ScopedRNNTensorDescriptor() PADDLE_MAY_THROW {
+    PADDLE_ENFORCE_GPU_SUCCESS(
+        phi::dynload::cudnnDestroyRNNDataDescriptor(desc_));
+  }
+
+  inline cudnnRNNDataDescriptor_t descriptor(
+      const cudnnDataType_t cudnn_type,
+      int max_seq_length,
+      int batch_size,
+      int input_size,
+      bool time_major,
+      const std::vector<int>& seq_length) {
+    static double padding_fill = 0.0f;
+    cudnnRNNDataLayout_t layout;
+
+    if (time_major) {
+      layout = CUDNN_RNN_DATA_LAYOUT_SEQ_MAJOR_UNPACKED;
+    } else {
+      layout = CUDNN_RNN_DATA_LAYOUT_BATCH_MAJOR_UNPACKED;
+    }
+
+    PADDLE_ENFORCE_GPU_SUCCESS(phi::dynload::cudnnSetRNNDataDescriptor(
+        desc_,
+        cudnn_type,
+        layout,
+        max_seq_length,
+        batch_size,
+        input_size,
+        seq_length.data(),
+        static_cast<void*>(&padding_fill)));
+
+    return desc_;
+  }
+
+  template <typename T>
+  inline cudnnRNNDataDescriptor_t descriptor(
+      int max_length,
+      int batch_size,
+      int input_size,
+      bool time_major,
+      const std::vector<int>& seq_length) {
+    return descriptor(CudnnDataType<T>::type,
+                      max_length,
+                      batch_size,
+                      input_size,
+                      time_major,
+                      seq_length);
+  }
+
+  inline cudnnRNNDataDescriptor_t desc() { return desc_; }
+
+ private:
+  cudnnRNNDataDescriptor_t desc_;
+  DISABLE_COPY_AND_ASSIGN(ScopedRNNTensorDescriptor);
+};
+
+class ScopedDropoutDescriptor {
+ public:
+  ScopedDropoutDescriptor() {
+    PADDLE_ENFORCE_GPU_SUCCESS(
+        phi::dynload::cudnnCreateDropoutDescriptor(&desc_));
+  }
+  ~ScopedDropoutDescriptor() PADDLE_MAY_THROW {
+    PADDLE_ENFORCE_GPU_SUCCESS(
+        phi::dynload::cudnnDestroyDropoutDescriptor(desc_));
+  }
+
+  inline cudnnDropoutDescriptor_t descriptor(const cudnnHandle_t& handle,
+                                             const Place& place UNUSED,
+                                             bool initialized,
+                                             float dropout_prob_,
+                                             DenseTensor* dropout_state_,
+                                             int seed,
+                                             size_t state_size) {
+    if (dropout_state_ == nullptr) {  // for no dropout or test
+      PADDLE_ENFORCE_GPU_SUCCESS(
+          phi::dynload::cudnnSetDropoutDescriptor(desc_,
+                                                  handle,
+                                                  0 /* dropout */,
+                                                  nullptr,
+                                                  0 /* state_size */,
+                                                  0 /* seed */));
+      return desc_;
+    }
+    auto* dropout_state_data = dropout_state_->data<uint8_t>();
+    if (!initialized) {
+      PADDLE_ENFORCE_GPU_SUCCESS(phi::dynload::cudnnSetDropoutDescriptor(
+          desc_, handle, dropout_prob_, dropout_state_data, state_size, seed));
+    } else {
+      auto dropout_state_dims =
+          common::vectorize<int64_t>(dropout_state_->dims());
+      state_size = dropout_state_dims[0];
+      PADDLE_ENFORCE_GPU_SUCCESS(phi::dynload::cudnnRestoreDropoutDescriptor(
+          desc_, handle, dropout_prob_, dropout_state_data, state_size, 0));
+    }
+    return desc_;
+  }
+  inline cudnnDropoutDescriptor_t desc() { return desc_; }
+
+ private:
+  cudnnDropoutDescriptor_t desc_;
+  DISABLE_COPY_AND_ASSIGN(ScopedDropoutDescriptor);
+};
+
+class ScopedRNNDescriptor {
+ public:
+  ScopedRNNDescriptor() {
+    PADDLE_ENFORCE_GPU_SUCCESS(phi::dynload::cudnnCreateRNNDescriptor(&desc_));
+  }
+  ~ScopedRNNDescriptor() PADDLE_MAY_THROW {
+    PADDLE_ENFORCE_GPU_SUCCESS(phi::dynload::cudnnDestroyRNNDescriptor(desc_));
+  }
+
+  inline cudnnRNNDescriptor_t desc() { return desc_; }
+
+ private:
+  cudnnRNNDescriptor_t desc_;
+  DISABLE_COPY_AND_ASSIGN(ScopedRNNDescriptor);
+};
+
+class ScopedFilterDescriptor {
+ public:
+  ScopedFilterDescriptor() {
+    PADDLE_ENFORCE_GPU_SUCCESS(
+        phi::dynload::cudnnCreateFilterDescriptor(&desc_));
+  }
+  ~ScopedFilterDescriptor() PADDLE_MAY_THROW {
+    PADDLE_ENFORCE_GPU_SUCCESS(
+        phi::dynload::cudnnDestroyFilterDescriptor(desc_));
+  }
+
+  inline cudnnFilterDescriptor_t descriptor(const cudnnTensorFormat_t format,
+                                            const cudnnDataType_t type,
+                                            const std::vector<int>& kernel,
+                                            const int groups = 1) {
+    // filter layout: MCHW(MCDHW), where M is the number of
+    // output image channels, C is the number of input image channels,
+    // D is the depth of the filter, H is the height of the filter, and W is the
+    // width of the filter.
+    std::vector<int> kernel_with_group(kernel.begin(), kernel.end());
+    if (groups > 1) {
+      kernel_with_group[0] /= groups;
+      // NOTE: input filter(C) of the filter is already asserted to be C/groups.
+    }
+    PADDLE_ENFORCE_GPU_SUCCESS(
+        phi::dynload::cudnnSetFilterNdDescriptor(desc_,
+                                                 type,
+                                                 format,
+                                                 kernel_with_group.size(),
+                                                 kernel_with_group.data()));
+    return desc_;
+  }
+
+  template <typename T>
+  inline cudnnFilterDescriptor_t descriptor(const DataLayout& order,
+                                            const std::vector<int>& kernel,
+                                            const int groups = 1) {
+    return descriptor(
+        GetCudnnTensorFormat(order), CudnnDataType<T>::type, kernel, groups);
+  }
+
+  inline cudnnFilterDescriptor_t desc() { return desc_; }
+
+ private:
+  cudnnFilterDescriptor_t desc_;
+  DISABLE_COPY_AND_ASSIGN(ScopedFilterDescriptor);
+};
+
+class ScopedConvolutionDescriptor {
+ public:
+  ScopedConvolutionDescriptor() {
+    PADDLE_ENFORCE_GPU_SUCCESS(
+        phi::dynload::cudnnCreateConvolutionDescriptor(&desc_));
+  }
+  ~ScopedConvolutionDescriptor() PADDLE_MAY_THROW {
+    PADDLE_ENFORCE_GPU_SUCCESS(
+        phi::dynload::cudnnDestroyConvolutionDescriptor(desc_));
+  }
+
+  inline cudnnConvolutionDescriptor_t descriptor(
+      cudnnDataType_t type,
+      const std::vector<int>& pads,
+      const std::vector<int>& strides,
+      const std::vector<int>& dilations) {
+    PADDLE_ENFORCE_EQ(pads.size(),
+                      strides.size(),
+                      common::errors::InvalidArgument(
+                          "The size of pads and strides should be equal. But "
+                          "received size of pads is %d, size of strides is %d.",
+                          pads.size(),
+                          strides.size()));
+    PADDLE_ENFORCE_EQ(
+        pads.size(),
+        dilations.size(),
+        common::errors::InvalidArgument(
+            "The size of pads and dilations should be equal. But received size "
+            "of pads is %d, size of dilations is %d.",
+            pads.size(),
+            dilations.size()));
+
+    cudnnDataType_t compute_type =
+        (type == CUDNN_DATA_DOUBLE) ? CUDNN_DATA_DOUBLE : CUDNN_DATA_FLOAT;
+    PADDLE_ENFORCE_GPU_SUCCESS(
+        phi::dynload::cudnnSetConvolutionNdDescriptor(desc_,
+                                                      pads.size(),
+                                                      pads.data(),
+                                                      strides.data(),
+                                                      dilations.data(),
+                                                      CUDNN_CROSS_CORRELATION,
+                                                      compute_type));
+    return desc_;
+  }
+
+  template <typename T>
+  inline cudnnConvolutionDescriptor_t descriptor(
+      const std::vector<int>& pads,
+      const std::vector<int>& strides,
+      const std::vector<int>& dilations) {
+    return descriptor(CudnnDataType<T>::type, pads, strides, dilations);
+  }
+
+ private:
+  cudnnConvolutionDescriptor_t desc_;
+  DISABLE_COPY_AND_ASSIGN(ScopedConvolutionDescriptor);
+};
+
+class ScopedPoolingDescriptor {
+ public:
+  ScopedPoolingDescriptor() {
+    PADDLE_ENFORCE_GPU_SUCCESS(
+        phi::dynload::cudnnCreatePoolingDescriptor(&desc_));
+  }
+  ~ScopedPoolingDescriptor() PADDLE_MAY_THROW {
+    PADDLE_ENFORCE_GPU_SUCCESS(
+        phi::dynload::cudnnDestroyPoolingDescriptor(desc_));
+  }
+
+  inline cudnnPoolingDescriptor_t descriptor(const PoolingMode& mode,
+                                             const std::vector<int>& kernel,
+                                             const std::vector<int>& pads,
+                                             const std::vector<int>& strides) {
+    PADDLE_ENFORCE_EQ(kernel.size(),
+                      pads.size(),
+                      common::errors::InvalidArgument(
+                          "The size of kernel and pads should be equal. But "
+                          "received size of kernel is %d, size of pads is %d.",
+                          kernel.size(),
+                          pads.size()));
+    PADDLE_ENFORCE_EQ(
+        kernel.size(),
+        strides.size(),
+        common::errors::InvalidArgument(
+            "The size of kernel and strides should be equal. But "
+            "received size of kernel is %d, size of strides is %d.",
+            kernel.size(),
+            strides.size()));
+    PADDLE_ENFORCE_GPU_SUCCESS(phi::dynload::cudnnSetPoolingNdDescriptor(
+        desc_,
+        (GetPoolingMode(mode)),
+        CUDNN_PROPAGATE_NAN,  // Always propagate nans.
+        kernel.size(),
+        kernel.data(),
+        pads.data(),
+        strides.data()));
+    return desc_;
+  }
+
+ private:
+  cudnnPoolingDescriptor_t desc_;
+  DISABLE_COPY_AND_ASSIGN(ScopedPoolingDescriptor);
+};
+
+class ScopedSpatialTransformerDescriptor {
+ public:
+  ScopedSpatialTransformerDescriptor() {
+    PADDLE_ENFORCE_GPU_SUCCESS(
+        phi::dynload::cudnnCreateSpatialTransformerDescriptor(&desc_));
+  }
+  ~ScopedSpatialTransformerDescriptor() PADDLE_MAY_THROW {
+    PADDLE_ENFORCE_GPU_SUCCESS(
+        phi::dynload::cudnnDestroySpatialTransformerDescriptor(desc_));
+  }
+
+  template <typename T>
+  inline cudnnSpatialTransformerDescriptor_t descriptor(const int nbDims,
+                                                        const int dimA[]) {
+    PADDLE_ENFORCE_GPU_SUCCESS(
+        phi::dynload::cudnnSetSpatialTransformerNdDescriptor(
+            desc_,
+            CUDNN_SAMPLER_BILINEAR,
+            CudnnDataType<T>::type,
+            nbDims,
+            dimA));
+    return desc_;
+  }
+
+ private:
+  cudnnSpatialTransformerDescriptor_t desc_;
+  DISABLE_COPY_AND_ASSIGN(ScopedSpatialTransformerDescriptor);
+};
+
+class ScopedActivationDescriptor {
+ public:
+  ScopedActivationDescriptor() {
+    PADDLE_ENFORCE_GPU_SUCCESS(
+        phi::dynload::cudnnCreateActivationDescriptor(&desc_));
+  }
+  ~ScopedActivationDescriptor() PADDLE_MAY_THROW {
+    PADDLE_ENFORCE_GPU_SUCCESS(
+        phi::dynload::cudnnDestroyActivationDescriptor(desc_));
+  }
+
+  template <typename T>
+  inline cudnnActivationDescriptor_t descriptor(
+      const std::string& act, double value_max = static_cast<double>(0.)) {
+    double relu_ceiling = 0.0;
+    ActivationMode activation_mode = StringToActivationMode(act);
+    cudnnActivationMode_t mode;
+    switch (activation_mode) {
+      case ActivationMode::kNone:
+        mode = CUDNN_ACTIVATION_IDENTITY;
+        break;
+      case ActivationMode::kRelu6:
+        relu_ceiling = 6.0;
+        mode = CUDNN_ACTIVATION_CLIPPED_RELU;
+        break;
+      case ActivationMode::kReluX:
+        relu_ceiling = value_max;
+        mode = CUDNN_ACTIVATION_CLIPPED_RELU;
+        break;
+      case ActivationMode::kRelu:
+        mode = CUDNN_ACTIVATION_RELU;
+        break;
+      case ActivationMode::kSigmoid:
+        mode = CUDNN_ACTIVATION_SIGMOID;
+        break;
+      case ActivationMode::kTanh:
+        mode = CUDNN_ACTIVATION_TANH;
+        break;
+      default:
+        PADDLE_THROW(common::errors::Unimplemented(
+            "Unrecognized CUDNN activation mode: %d.",
+            static_cast<int>(activation_mode)));
+    }
+    PADDLE_ENFORCE_GPU_SUCCESS(phi::dynload::cudnnSetActivationDescriptor(
+        desc_, mode, CUDNN_NOT_PROPAGATE_NAN, relu_ceiling));
+    return desc_;
+  }
+
+ private:
+  cudnnActivationDescriptor_t desc_;
+  DISABLE_COPY_AND_ASSIGN(ScopedActivationDescriptor);
+};
+
+class ScopedCTCLossDescriptor {
+ public:
+  ScopedCTCLossDescriptor() {
+    PADDLE_ENFORCE_GPU_SUCCESS(
+        phi::dynload::cudnnCreateCTCLossDescriptor(&desc_));
+  }
+  ~ScopedCTCLossDescriptor() PADDLE_MAY_THROW {
+    PADDLE_ENFORCE_GPU_SUCCESS(
+        phi::dynload::cudnnDestroyCTCLossDescriptor(desc_));
+  }
+
+  template <typename T>
+  inline cudnnCTCLossDescriptor_t descriptor() {
+    PADDLE_ENFORCE_GPU_SUCCESS(
+        phi::dynload::cudnnSetCTCLossDescriptor(desc_, CudnnDataType<T>::type));
+    return desc_;
+  }
+
+ private:
+  cudnnCTCLossDescriptor_t desc_;
+  DISABLE_COPY_AND_ASSIGN(ScopedCTCLossDescriptor);
+};
+
+}  // namespace gpu
+}  // namespace backends
+}  // namespace phi

paddle/phi/backends/gpu/cuda/cudnn_workspace_helper.cc

@@ -0,0 +1,32 @@
+// Copyright (c) 2022 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.
+
+#include "paddle/phi/backends/gpu/cuda/cudnn_workspace_helper.h"
+
+#include <cstdlib>
+#include <string>
+
+namespace phi::backends::gpu {
+
+static int GetDefaultConvWorkspaceSizeLimitMBImpl() {
+  const char *env_str = std::getenv("FLAGS_conv_workspace_size_limit");
+  return env_str ? std::stoi(std::string(env_str))
+                 : kDefaultConvWorkspaceSizeLimitMB;
+}
+
+int GetDefaultConvWorkspaceSizeLimitMB() {
+  static auto workspace_size = GetDefaultConvWorkspaceSizeLimitMBImpl();
+  return workspace_size;
+}
+}  // namespace phi::backends::gpu

paddle/phi/backends/gpu/cuda/cudnn_workspace_helper.h

@@ -0,0 +1,38 @@
+// Copyright (c) 2022 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.
+
+#pragma once
+#include "paddle/common/macros.h"
+namespace phi {
+namespace backends {
+namespace gpu {
+
+#define CUDNN_ENFORCE_TENSOR_SIZE_SUPPORTED(tensor)                            \
+  do {                                                                         \
+    int64_t largest = (1LL << 31) - 1;                                         \
+    PADDLE_ENFORCE_LE(                                                         \
+        tensor.numel(),                                                        \
+        largest,                                                               \
+        ::common::errors::PreconditionNotMet(                                  \
+            "The element size of " #tensor " should be <= INT_MAX(2147483647)" \
+            ", but got %lld",                                                  \
+            tensor.numel()));                                                  \
+  } while (0)
+
+static constexpr int kDefaultConvWorkspaceSizeLimitMB = 512;
+
+PADDLE_API int GetDefaultConvWorkspaceSizeLimitMB();
+}  // namespace gpu
+}  // namespace backends
+}  // namespace phi

paddle/phi/backends/gpu/cuda/gpu_event_timer.cc

@@ -0,0 +1,145 @@
+// Copyright (c) 2024 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.
+
+#include "paddle/phi/backends/gpu/cuda/gpu_event_timer.h"
+#include <numeric>
+#include "glog/logging.h"
+#include "paddle/common/errors.h"
+#include "paddle/phi/backends/context_pool.h"
+#include "paddle/phi/backends/gpu/gpu_context.h"
+#include "paddle/phi/core/enforce.h"
+
+namespace phi {
+
+EventPair::EventPair() {
+  PADDLE_ENFORCE_GPU_SUCCESS(cudaEventCreate(&start_));
+  PADDLE_ENFORCE_GPU_SUCCESS(cudaEventCreate(&end_));
+  started_ = false;
+}
+
+EventPair::~EventPair() PADDLE_MAY_THROW {
+  PADDLE_ENFORCE_GPU_SUCCESS(cudaEventDestroy(start_));
+  PADDLE_ENFORCE_GPU_SUCCESS(cudaEventDestroy(end_));
+}
+
+void EventPair::Start(cudaStream_t stream) {
+  PADDLE_ENFORCE_EQ(
+      started_,
+      false,
+      common::errors::InvalidArgument("GPUEventTimer has been started."));
+  PADDLE_ENFORCE_GPU_SUCCESS(cudaEventRecord(start_, stream));
+  started_ = true;
+}
+
+void EventPair::Stop(cudaStream_t stream) {
+  PADDLE_ENFORCE_EQ(
+      started_,
+      true,
+      common::errors::InvalidArgument("GPUEventTimer has not been started."));
+  PADDLE_ENFORCE_GPU_SUCCESS(cudaEventRecord(end_, stream));
+  started_ = false;
+}
+
+void EventPair::Reset() { started_ = false; }
+
+double EventPair::Elapsed() {
+  PADDLE_ENFORCE_EQ(
+      started_,
+      false,
+      common::errors::InvalidArgument("GPUEventTimer has not been stopped."));
+  PADDLE_ENFORCE_GPU_SUCCESS(cudaEventSynchronize(start_));
+  PADDLE_ENFORCE_GPU_SUCCESS(cudaEventSynchronize(end_));
+  float ms;
+  PADDLE_ENFORCE_GPU_SUCCESS(cudaEventElapsedTime(&ms, start_, end_));
+  return ms / 1000.0;
+}
+
+GPUEventTimer::GPUEventTimer(phi::GPUPlace place) : length_(0) {
+  auto *dev_ctx = phi::DeviceContextPool::Instance().GetByPlace(place);
+  default_stream_ = dev_ctx->stream();
+}
+
+EventPair *GPUEventTimer::GetLatest() {
+  PADDLE_ENFORCE_GT(
+      length_,
+      0,
+      common::errors::InvalidArgument("GPUEventTimer has not been started."));
+  auto &back = events_[length_ - 1];
+  if (back == nullptr) {
+    back.reset(new EventPair());
+  }
+  return back.get();
+}
+
+void GPUEventTimer::Start(cudaStream_t stream) {
+  if (length_ == events_.size()) {
+    VLOG(10) << "Expand when length = " << length_;
+    events_.emplace_back();
+  }
+  ++length_;
+  GetLatest()->Start(stream);
+}
+
+void GPUEventTimer::Stop(cudaStream_t stream) { GetLatest()->Stop(stream); }
+
+void GPUEventTimer::Start() { Start(default_stream_); }
+
+void GPUEventTimer::Stop() { Stop(default_stream_); }
+
+void GPUEventTimer::Reset() {
+  for (size_t i = 0; i < length_; ++i) {
+    events_[i]->Reset();
+  }
+  length_ = 0;
+}
+
+double GPUEventTimer::Elapsed(bool reset) {
+  double ret = 0;
+  for (size_t i = 0; i < length_; ++i) {
+    ret += events_[i]->Elapsed();
+  }
+  if (reset) {
+    Reset();
+  }
+  return ret;
+}
+
+std::vector<double> GPUEventTimer::ElapsedList(bool reset) {
+  std::vector<double> values(length_);
+  for (size_t i = 0; i < length_; ++i) {
+    values[i] = events_[i]->Elapsed();
+  }
+  if (reset) {
+    Reset();
+  }
+  return values;
+}
+
+void GPUEventTimer::PreAlloc(size_t n) {
+  if (events_.size() >= n) return;
+  events_.resize(n);
+  for (auto &pair : events_) {
+    if (pair == nullptr) {
+      pair.reset(new EventPair());
+    }
+  }
+}
+
+void GPUEventTimer::ShrinkToFit() { events_.resize(length_); }
+
+size_t GPUEventTimer::Size() const { return length_; }
+
+size_t GPUEventTimer::Capacity() const { return events_.size(); }
+
+}  // namespace phi

paddle/phi/backends/gpu/cuda/gpu_event_timer.h

@@ -0,0 +1,85 @@
+// Copyright (c) 2024 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.
+
+#pragma once
+
+#include <memory>
+#include <vector>
+#include "cuda_runtime.h"  // NOLINT
+#include "paddle/common/enforce.h"
+#include "paddle/common/macros.h"
+#include "paddle/phi/common/place.h"
+
+namespace phi {
+
+class EventPair {
+  DISABLE_COPY_AND_ASSIGN(EventPair);
+
+ public:
+  EventPair();
+
+  ~EventPair() PADDLE_MAY_THROW;
+
+  void Start(cudaStream_t stream);
+
+  void Stop(cudaStream_t stream);
+
+  void Reset();
+
+  double Elapsed();
+
+ private:
+  cudaEvent_t start_;
+  cudaEvent_t end_;
+  bool started_;
+};
+
+class GPUEventTimer {
+  DISABLE_COPY_AND_ASSIGN(GPUEventTimer);
+
+ public:
+  explicit GPUEventTimer(phi::GPUPlace place);
+
+  void Start(cudaStream_t stream);
+
+  void Stop(cudaStream_t stream);
+
+  void Start();
+
+  void Stop();
+
+  void Reset();
+
+  double Elapsed(bool reset);
+
+  std::vector<double> ElapsedList(bool reset);
+
+  void PreAlloc(size_t n);
+
+  void ShrinkToFit();
+
+  size_t Size() const;
+
+  size_t Capacity() const;
+
+ private:
+  EventPair *GetLatest();
+
+ private:
+  std::vector<std::unique_ptr<EventPair>> events_;
+  size_t length_;
+  cudaStream_t default_stream_;
+};
+
+}  // namespace phi

paddle/phi/backends/gpu/forwards.h

@@ -0,0 +1,115 @@
+/* Copyright (c) 2022 PaddlePaddle Authors. All Rights Reserved.
+Copyright (c) 2022 NVIDIA Corporation. 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. */
+
+// Forward-declares CUDA API types used in platform-agnostic wrapper headers.
+#pragma once
+
+/// Forward declaration of Eigen types.
+namespace Eigen {
+struct GpuDevice;
+}  // namespace Eigen
+
+/// Forward declaration of CUDA types.
+
+// Forward declaration of CUDA runtime types.
+using cudaStream_t = struct CUstream_st *;
+using cudaEvent_t = struct CUevent_st *;
+
+// Forward declaration of cuBLAS types.
+using cublasHandle_t = struct cublasContext *;
+
+// Forward declaration of cuBLASLt types.
+using cublasLtHandle_t = struct cublasLtContext *;
+
+#ifndef PADDLE_WITH_CUSTOM_DEVICE
+// Forward declaration of cuDNN types.
+using cudnnHandle_t = struct cudnnContext *;
+using cudnnTensorDescriptor_t = struct cudnnTensorStruct *;
+using cudnnConvolutionDescriptor_t = struct cudnnConvolutionStruct *;
+using cudnnPoolingDescriptor_t = struct cudnnPoolingStruct *;
+using cudnnFilterDescriptor_t = struct cudnnFilterStruct *;
+using cudnnLRNDescriptor_t = struct cudnnLRNStruct *;
+using cudnnActivationDescriptor_t = struct cudnnActivationStruct *;
+using cudnnSpatialTransformerDescriptor_t =
+    struct cudnnSpatialTransformerStruct *;
+using cudnnOpTensorDescriptor_t = struct cudnnOpTensorStruct *;
+using cudnnReduceTensorDescriptor_t = struct cudnnReduceTensorStruct *;
+using cudnnCTCLossDescriptor_t = struct cudnnCTCLossStruct *;
+using cudnnTensorTransformDescriptor_t = struct cudnnTensorTransformStruct *;
+using cudnnDropoutDescriptor_t = struct cudnnDropoutStruct *;
+using cudnnRNNDescriptor_t = struct cudnnRNNStruct *;
+using cudnnPersistentRNNPlan_t = struct cudnnPersistentRNNPlan *;
+using cudnnRNNDataDescriptor_t = struct cudnnRNNDataStruct *;
+using cudnnAlgorithmDescriptor_t = struct cudnnAlgorithmStruct *;
+using cudnnAlgorithmPerformance_t = struct cudnnAlgorithmPerformanceStruct *;
+using cudnnSeqDataDescriptor_t = struct cudnnSeqDataStruct *;
+using cudnnAttnDescriptor_t = struct cudnnAttnStruct *;
+using cudnnFusedOpsConstParamPack_t = struct cudnnFusedOpsConstParamStruct *;
+using cudnnFusedOpsVariantParamPack_t =
+    struct cudnnFusedOpsVariantParamStruct *;
+using cudnnFusedOpsPlan_t = struct cudnnFusedOpsPlanStruct *;
+
+// Forward declaration of cuSOLVER types.
+using cusolverDnHandle_t = struct cusolverDnContext *;
+
+// Forward declaration of cuSparse types.
+using cusparseHandle_t = struct cusparseContext *;
+
+// Forward declaration of NCCL types.
+using ncclComm_t = struct ncclComm *;
+
+/// Forward declaration of ROCM types.
+#include <cstddef>
+
+using hipDevice_t = int;
+using hipCtx_t = struct ihipCtx_t *;
+using hipStream_t = struct ihipStream_t *;
+using hipEvent_t = struct ihipEvent_t *;
+
+// Forward declaration of MIOpen types.
+using miopenHandle_t = struct miopenHandle *;
+using miopenAcceleratorQueue_t = hipStream_t;
+using miopenFusionOpDescriptor_t = struct miopenFusionOpDescriptor *;
+using miopenTensorDescriptor_t = struct miopenTensorDescriptor *;
+using miopenConvolutionDescriptor_t = struct miopenConvolutionDescriptor *;
+using miopenPoolingDescriptor_t = struct miopenPoolingDescriptor *;
+using miopenLRNDescriptor_t = struct miopenLRNDescriptor *;
+using miopenActivationDescriptor_t = struct miopenActivationDescriptor *;
+using miopenRNNDescriptor_t = struct miopenRNNDescriptor *;
+using miopenCTCLossDescriptor_t = struct miopenCTCLossDescriptor *;
+using miopenDropoutDescriptor_t = struct miopenDropoutDescriptor *;
+using miopenFusionPlanDescriptor_t = struct miopenFusionPlanDescriptor *;
+using miopenOperatorDescriptor_t = struct miopenOperatorDescriptor *;
+using miopenOperatorArgs_t = struct miopenOperatorArgs *;
+using miopenAllocatorFunction = void *(*)(void *context, size_t sizeBytes);
+// using miopenDeallocatorFunction = void *(*)(void *context, void *memory);
+// struct miopenConvAlgoPerf_t;
+// struct miopenConvSolution_t;
+
+// Forward declaration of rocBLAS types.
+using rocblas_handle = struct _rocblas_handle *;
+
+// Forward declaration of hipblaslt types.
+using hipblasLtHandle_t = void *;
+
+// Forward declaration of hipfft types.
+using hipfftHandle = struct hipfftHandle_t *;
+
+// Forward declaration of rocSOLVER types.
+using rocsolver_handle = rocblas_handle;
+
+// Forward declaration of rocSparse types.
+using rocsparse_handle = struct _rocsparse_handle *;
+#endif

paddle/phi/backends/gpu/gpu_context.h

@@ -0,0 +1,329 @@
+/* Copyright (c) 2022 PaddlePaddle Authors. All Rights Reserved.
+Copyright (c) 2022 NVIDIA Corporation. 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. */
+
+#pragma once
+
+#ifdef PADDLE_WITH_CUSTOM_DEVICE
+#include "paddle/phi/backends/custom/custom_context.h"
+#include "paddle/phi/backends/gpu/gpu_helper.h"
+#include "paddle/phi/backends/gpu/gpu_info.h"
+#else
+#if defined(PADDLE_WITH_CUDA) || defined(PADDLE_WITH_HIP) || \
+    defined(PADDLE_WITH_XPU_KP)
+
+#include <array>
+#include <functional>
+#include <mutex>
+#include <utility>
+
+#include "paddle/common/enforce.h"
+#include "paddle/phi/backends/gpu/forwards.h"
+#include "paddle/phi/backends/gpu/gpu_decls.h"
+#include "paddle/phi/backends/gpu/gpu_helper.h"
+#include "paddle/phi/backends/gpu/gpu_info.h"
+#include "paddle/phi/common/place.h"
+#include "paddle/phi/core/attribute.h"
+#include "paddle/phi/core/device_context.h"
+
+namespace phi {
+
+class CUDAStream;
+
+class DnnWorkspaceHandle {
+ public:
+  inline DnnWorkspaceHandle(Allocator* allocator, gpuStream_t stream)
+      : allocator_(allocator), stream_(stream) {
+    mtx_ = std::make_unique<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 */
+  PADDLE_API void RunFuncSync(const std::function<void(void*)>& cudnn_func,
+                              size_t required_workspace_bytes,
+                              bool use_cached_allocation = true);
+
+  inline size_t WorkspaceSize() {
+    if (allocation_ == nullptr) {
+      return 0;
+    }
+    return allocation_->size();
+  }
+
+  PADDLE_API void ResetWorkspace();
+
+  TEST_API void ReallocWorkspace(size_t required_workspace_bytes);
+
+  DnnWorkspaceHandle(DnnWorkspaceHandle&&) = default;
+  DnnWorkspaceHandle& operator=(DnnWorkspaceHandle&&) = delete;
+
+ private:
+  Allocator::AllocationPtr allocation_{nullptr};
+  Allocator* allocator_{nullptr};  // Not owned
+  gpuStream_t stream_{nullptr};    // Not owned
+  std::unique_ptr<std::mutex> mtx_;
+};
+
+class PADDLE_API GPUContext : public DeviceContext,
+                              public TypeInfoTraits<DeviceContext, GPUContext> {
+ public:
+  explicit GPUContext(const GPUPlace& place,
+                      bool init = true,
+                      int stream_priority = 0);
+
+  GPUContext(GPUContext&&);
+  GPUContext& operator=(GPUContext&&);
+
+  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 CUDAStream in the device context. */
+  CUDAStream* cuda_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;
+
+  /*! \brief  Return cublasLt handle in the device context. */
+  blasLtHandle_t cublaslt_handle() const;
+
+  /*! \brief  Return persistent cublasLt workspace (grow-only, multi-stream
+   * safe). */
+  std::pair<void*, size_t> cublaslt_workspace(size_t required_size) const;
+
+  /*! \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<unsigned 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.
+   */
+  // TODO(wilber): The return type is a pointer, to be modified later.
+  DnnWorkspaceHandle cudnn_workspace_handle() const;
+
+ 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;
+
+  // Several methods for adapting Dnn-specific attributes
+  bool HasDnnAttr(const std::string& attr_name) const;
+  const Attribute& GetDnnAttr(const std::string& attr_name) const;
+  void SetDnnAttr(const std::string& attr_name, Attribute attr);
+  void ClearDnnAttr();
+
+  static const char* name() { return "GPUContext"; }
+
+ public:
+  /*! \brief  Return nccl communicators. */
+  ncclComm_t nccl_comm() const;
+
+  /*! \brief  Set nccl communicators. */
+  void set_nccl_comm(ncclComm_t comm);
+
+  // 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);
+
+ 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(int stream_priority = 0);
+  // 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();
+
+  // Note that this function is a trick implementation since all 'set' methods
+  // are protected by default.
+  // clear: whether clear the original CUDAStream or not
+  void SetCUDAStream(CUDAStream*, bool clear = true);
+
+ protected:
+  void SetEigenDevice(Eigen::GpuDevice*);
+  void SetEigenDevice(std::function<Eigen::GpuDevice*()>&&);
+
+  void SetBlasHandle(blasHandle_t);
+  void SetBlasHandle(std::function<blasHandle_t()>&&);
+
+  void SetBlasTensorCoreHandle(blasHandle_t);
+  void SetBlasTensorCoreHandle(std::function<blasHandle_t()>&&);
+
+  void SetBlasTF32Handle(blasHandle_t);
+  void SetBlasTF32Handle(std::function<blasHandle_t()>&&);
+
+  void SetBlasLtHandle(blasLtHandle_t);
+  void SetBlasLtHandle(std::function<blasLtHandle_t()>&&);
+
+  void SetDnnHandle(dnnHandle_t);
+  void SetDnnHandle(std::function<dnnHandle_t()>&&);
+
+  void SetSolverHandle(solverHandle_t);
+  void SetSolverHandle(std::function<solverHandle_t()>&&);
+
+  void SetSparseHandle(sparseHandle_t);
+  void SetSparseHandle(std::function<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<unsigned int, 3>& val);
+
+  void SetDriverVersion(int val);
+
+  void SetRuntimeVersion(int val);
+
+ private:
+  struct Impl;
+  std::unique_ptr<Impl> impl_;
+};
+
+// Note: In order to register the kernel of CUDNN, DnnContext 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 Dnn kernel function, so if we using DnnContext = GPUContext, we
+// must use different function name for cudnn kernel
+using GPUDNNContext = GPUContext;
+
+// 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
+
+}  // namespace phi
+
+namespace Eigen {
+struct DefaultDevice;
+}  // namespace Eigen
+
+namespace phi {
+#if defined(PADDLE_WITH_CUDA) || defined(PADDLE_WITH_HIP)
+// Currently, GPUPinnedContext is only used to data copying.
+class GPUPinnedContext
+    : public DeviceContext,
+      public phi::TypeInfoTraits<DeviceContext, GPUPinnedContext> {
+ public:
+  PADDLE_API GPUPinnedContext();
+  PADDLE_API explicit GPUPinnedContext(GPUPinnedPlace place);
+
+  const Place& GetPlace() const override;
+
+  Eigen::DefaultDevice* eigen_device() const;
+
+  dnnHandle_t cudnn_handle() const override {
+    PADDLE_THROW(common::errors::Unavailable(
+        "GPUPinnedContext does not support cudnn_handle()."));
+  }
+
+  static const char* name() { return "GPUPinnedContext"; }
+
+ private:
+  GPUPinnedPlace place_;
+  std::unique_ptr<Eigen::DefaultDevice> eigen_device_;
+};
+#endif
+}  // namespace phi
+
+#endif
+#endif

paddle/phi/backends/gpu/gpu_decls.h

@@ -0,0 +1,74 @@
+// Copyright (c) 2022 PaddlePaddle Authors. All Rights Reserved.
+// Copyright (c) 2022 NVIDIA Corporation. 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.
+
+#pragma once
+
+#include "paddle/phi/backends/gpu/forwards.h"
+
+namespace phi {
+
+#ifdef PADDLE_WITH_HIP
+#define DECLARE_TYPE_FOR_GPU(GPU_TYPE, CUDA_TYPE, ROCM_TYPE) \
+  using GPU_TYPE = ROCM_TYPE;
+
+#else  // PADDLE_WITH_CUDA
+
+#define DECLARE_TYPE_FOR_GPU(GPU_TYPE, CUDA_TYPE, ROCM_TYPE) \
+  using GPU_TYPE = CUDA_TYPE;
+#endif
+
+DECLARE_TYPE_FOR_GPU(gpuStream_t, cudaStream_t, hipStream_t);
+DECLARE_TYPE_FOR_GPU(gpuEvent_t, cudaEvent_t, hipEvent_t);
+
+#ifndef PADDLE_WITH_CUSTOM_DEVICE
+DECLARE_TYPE_FOR_GPU(dnnActivationDescriptor,
+                     cudnnActivationStruct,
+                     miopenActivationDescriptor);
+DECLARE_TYPE_FOR_GPU(dnnTensorDescriptor,
+                     cudnnTensorStruct,
+                     miopenTensorDescriptor);
+DECLARE_TYPE_FOR_GPU(dnnFilterDescriptor,
+                     cudnnFilterStruct,
+                     miopenTensorDescriptor);
+DECLARE_TYPE_FOR_GPU(dnnFilterDescriptor_t,
+                     cudnnFilterDescriptor_t,
+                     miopenTensorDescriptor_t);
+DECLARE_TYPE_FOR_GPU(dnnConvolutionDescriptor,
+                     cudnnConvolutionStruct,
+                     miopenConvolutionDescriptor);
+DECLARE_TYPE_FOR_GPU(dnnConvolutionDescriptor_t,
+                     cudnnConvolutionDescriptor_t,
+                     miopenConvolutionDescriptor_t);
+DECLARE_TYPE_FOR_GPU(dnnPoolingDescriptor_t,
+                     cudnnPoolingDescriptor_t,
+                     miopenPoolingDescriptor_t);
+DECLARE_TYPE_FOR_GPU(dnnDropoutDescriptor_t,
+                     cudnnDropoutDescriptor_t,
+                     miopenDropoutDescriptor_t);
+
+DECLARE_TYPE_FOR_GPU(blasHandle_t, cublasHandle_t, rocblas_handle);
+
+DECLARE_TYPE_FOR_GPU(blasLtHandle_t, cublasLtHandle_t, hipblasLtHandle_t);
+
+DECLARE_TYPE_FOR_GPU(solverHandle_t, cusolverDnHandle_t, rocsolver_handle);
+
+DECLARE_TYPE_FOR_GPU(sparseHandle_t, cusparseHandle_t, rocsparse_handle);
+
+#undef DECLARE_TYPE_FOR_GPU
+
+using CUDAGraphID = unsigned long long;  // NOLINT
+#endif
+
+}  // namespace phi

paddle/phi/backends/gpu/gpu_device_function.h

@@ -0,0 +1,24 @@
+/* Copyright (c) 2022 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. */
+
+#pragma once
+#if defined(PADDLE_WITH_CUDA) || defined(PADDLE_WITH_HIP)
+
+#ifdef PADDLE_WITH_HIP
+#include "paddle/phi/backends/gpu/rocm/rocm_device_function.h"
+#else
+#include "paddle/phi/backends/gpu/cuda/cuda_device_function.h"
+#endif
+
+#endif

paddle/phi/backends/gpu/gpu_dnn.h

@@ -0,0 +1,27 @@
+// Copyright (c) 2022 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.
+
+#pragma once
+
+#if defined(PADDLE_WITH_CUDA) || defined(PADDLE_WITH_HIP)
+
+#ifdef PADDLE_WITH_HIP
+#include "paddle/phi/backends/gpu/rocm/miopen_desc.h"
+#include "paddle/phi/backends/gpu/rocm/miopen_helper.h"
+#else  // CUDA
+#include "paddle/phi/backends/gpu/cuda/cudnn_desc.h"
+#include "paddle/phi/backends/gpu/cuda/cudnn_helper.h"
+#endif
+
+#endif

paddle/phi/backends/gpu/gpu_helper.h

@@ -0,0 +1,26 @@
+// Copyright (c) 2021 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.
+
+#pragma once
+#if defined(PADDLE_WITH_CUDA) || defined(PADDLE_WITH_HIP)
+
+#ifdef PADDLE_WITH_HIP
+#include "paddle/phi/backends/gpu/rocm/rocm_helper.h"
+#else
+#include "paddle/phi/backends/gpu/cuda/cuda_helper.h"
+#endif
+
+#define CUDA_KERNEL_LOOP(i, num) CUDA_KERNEL_LOOP_TYPE(i, num, int)
+
+#endif

paddle/phi/backends/gpu/gpu_info.cc

@@ -0,0 +1,84 @@
+/* Copyright (c) 2016 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. */
+
+#include "paddle/phi/backends/gpu/gpu_info.h"
+
+#include <sstream>
+#include <vector>
+
+#include "glog/logging.h"
+#include "paddle/common/flags.h"
+
+#include "paddle/phi/common/memory_utils.h"
+
+COMMON_DECLARE_string(selected_gpus);
+
+namespace phi::backends::gpu {
+
+static inline std::vector<std::string> Split(std::string const& original,
+                                             char separator) {
+  std::vector<std::string> results;
+  std::string token;
+  std::istringstream is(original);
+  while (std::getline(is, token, separator)) {
+    if (!token.empty()) {
+      results.push_back(token);
+    }
+  }
+  return results;
+}
+
+//! Get a list of device ids from environment variable or use all.
+std::vector<int> GetSelectedDevices() {
+  // use user specified GPUs in single-node multi-process mode.
+  std::vector<int> devices;
+  if (!FLAGS_selected_gpus.empty()) {
+    auto devices_str = Split(FLAGS_selected_gpus, ',');
+    for (auto const& id : devices_str) {
+      devices.push_back(atoi(id.c_str()));
+    }
+  } else {
+    int count = GetGPUDeviceCount();
+    for (int i = 0; i < count; ++i) {
+      devices.push_back(i);
+    }
+  }
+  return devices;
+}
+
+constexpr static float fraction_reserve_gpu_memory = 0.05f;
+
+size_t GpuAvailableMemToAlloc() {
+  size_t total = 0;
+  size_t available = 0;
+  memory_utils::GpuMemoryUsage(&available, &total);
+  size_t reserving =
+      static_cast<size_t>(fraction_reserve_gpu_memory * available);  // NOLINT
+  // If available size is less than minimum chunk size, no usable memory exists
+  size_t available_to_alloc = available - reserving;
+  size_t min_chunk_size = GpuMinChunkSize();
+  if (available_to_alloc < min_chunk_size) {
+    available_to_alloc = 0;
+  }
+  VLOG(10) << "GPU usage " << (available >> 20) << "M/" << (total >> 20)
+           << "M, " << (available_to_alloc >> 20) << "M available to allocate";
+  return available_to_alloc;
+}
+
+size_t GpuMinChunkSize() {
+  // Allow to allocate the minimum chunk size is 256 bytes.
+  return 1 << 8;
+}
+
+}  // namespace phi::backends::gpu

paddle/phi/backends/gpu/gpu_info.h

@@ -0,0 +1,148 @@
+/* Copyright (c) 2016 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. */
+
+#pragma once
+
+#if defined(PADDLE_WITH_CUDA) || defined(PADDLE_WITH_HIP)
+
+#include <stddef.h>
+
+#include <array>
+#include <string>
+#include <utility>
+#include <vector>
+
+#include "paddle/common/macros.h"
+#include "paddle/phi/backends/gpu/gpu_types.h"
+
+namespace phi {
+namespace backends {
+namespace gpu {
+
+//! Get the version of dnn
+int DnnVersion();
+
+//! Get the total number of GPU devices in system.
+PADDLE_API int GetGPUDeviceCount();
+
+//! Get the compute capability of the ith GPU (format: major * 10 + minor)
+int GetGPUComputeCapability(int id);
+
+//! Get the runtime version of the ith GPU
+int GetGPURuntimeVersion(int id);
+
+//! Get the driver version of the ith GPU
+int GetGPUDriverVersion(int id);
+
+//! Whether the current device support TensorCore
+PADDLE_API bool TensorCoreAvailable();
+
+//! Get the MultiProcessors of the ith GPU.
+int GetGPUMultiProcessors(int id);
+
+//! Get the MaxThreads of each MultiProcessor of the ith GPU.
+int GetGPUMaxThreadsPerMultiProcessor(int id);
+
+//! Get the MaxThreads of each block of the ith GPU.
+int GetGPUMaxThreadsPerBlock(int id);
+
+//! Get the current GPU device id in system.
+PADDLE_API int GetCurrentDeviceId();
+
+//! Get the maximum GridDim size for GPU buddy allocator.
+std::array<unsigned int, 3> GetGpuMaxGridDimSize(int);
+
+std::pair<int, int> GetGpuStreamPriorityRange();
+
+//! Get a list of device ids from environment variable or use all.
+std::vector<int> GetSelectedDevices();
+
+//! Get the properties of the ith GPU device.
+PADDLE_API const gpuDeviceProp &GetDeviceProperties(int id);
+
+//! Set the GPU device id for next execution.
+PADDLE_API void SetDeviceId(int device_id);
+
+//! Get the available memory to allocate, which is the size of available gpu
+//! minus reserving.
+size_t GpuAvailableMemToAlloc();
+
+//! Get the minimum chunk size for GPU buddy allocator.
+size_t GpuMinChunkSize();
+
+//! Copy memory from address src to dst asynchronously.
+void GpuMemcpyAsync(void *dst,
+                    const void *src,
+                    size_t count,
+                    gpuMemcpyKind kind,
+                    gpuStream_t stream);
+
+//! Copy memory from address src to dst synchronously.
+void GpuMemcpySync(void *dst,
+                   const void *src,
+                   size_t count,
+                   gpuMemcpyKind kind);
+
+//! Copy memory from one device to another device asynchronously.
+void GpuMemcpyPeerAsync(void *dst,
+                        int dst_device,
+                        const void *src,
+                        int src_device,
+                        size_t count,
+                        gpuStream_t stream);
+
+//! Copy memory from one device to another device synchronously.
+void GpuMemcpyPeerSync(
+    void *dst, int dst_device, const void *src, int src_device, size_t count);
+
+//! Set memory dst with value count size asynchronously
+void GpuMemsetAsync(void *dst, int value, size_t count, gpuStream_t stream);
+
+//! Blocks until stream has completed all operations.
+PADDLE_API void GpuStreamSync(gpuStream_t stream);
+
+void GpuDestroyStream(gpuStream_t stream);
+
+// ! Blocks until device has completed all operations.
+void GpuDeviceSync();
+
+gpuError_t GpuGetLastError();
+
+bool IsGPUManagedMemorySupported(int dev_id);
+
+bool IsGPUManagedMemoryOversubscriptionSupported(int dev_id);
+
+class GPUDeviceGuard {
+ public:
+  explicit inline GPUDeviceGuard(int dev_id) {
+    int prev_id = GetCurrentDeviceId();
+    if (prev_id != dev_id) {
+      prev_id_ = prev_id;
+      SetDeviceId(dev_id);
+    }
+  }
+  inline ~GPUDeviceGuard() {
+    if (prev_id_ != -1) {
+      SetDeviceId(prev_id_);
+    }
+  }
+  GPUDeviceGuard(const GPUDeviceGuard &o) = delete;
+  GPUDeviceGuard &operator=(const GPUDeviceGuard &o) = delete;
+
+ private:
+  int prev_id_{-1};
+};
+
+}  // namespace gpu
+}  // namespace backends
+}  // namespace phi
+
+#endif

paddle/phi/backends/gpu/gpu_launch_config.h

@@ -0,0 +1,247 @@
+// Copyright (c) 2019 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.
+
+// Used for compute gpu launch parameter config
+
+#pragma once
+
+#if defined(PADDLE_WITH_CUDA) || defined(PADDLE_WITH_HIP)
+
+#ifdef PADDLE_WITH_CUDA
+#include <cuda_runtime.h>
+#else
+#include <hip/hip_runtime.h>
+#endif
+
+#include <stddef.h>
+
+#include <algorithm>
+#include <string>
+#include <vector>
+
+#include "glog/logging.h"
+#include "paddle/phi/backends/gpu/gpu_context.h"
+#include "paddle/phi/core/enforce.h"
+
+// CUDA performs better when thread_per_block is between [64, 512]
+#define PREDEFINED_BLOCK_SIZE 512
+
+namespace phi {
+namespace backends {
+namespace gpu {
+
+// Limitation of the setting in one dimension of cuda grid.
+constexpr int kMultiDimslimit = 65536;
+
+template <typename T = int64_t>
+inline T DivUp(T a, T b) {
+  return (a + b - 1) / b;
+}
+
+// https://graphics.stanford.edu/~seander/bithacks.html#RoundUpPowerOf2
+//   for round integer value into next highest power of 2.
+inline int64_t RoundToNextHighPowOfTwo(int64_t n, int64_t min_val = 1) {
+  n--;
+  n |= (n >> 1);
+  n |= (n >> 2);
+  n |= (n >> 4);
+  n |= (n >> 8);
+  n |= (n >> 16);
+  return std::max(min_val, (n + 1));
+}
+
+inline int64_t RoundToPowerOfTwo(int64_t n) {
+  constexpr int64_t min_val = 32;
+  int64_t num = RoundToNextHighPowOfTwo(n, min_val);
+  int64_t max_val = 1024;
+  return std::min(max_val, num);
+}
+
+#ifdef WITH_NV_JETSON
+// The number of threads cannot be assigned 1024 in some cases when the device
+// is nano or tx2 .
+inline void ChangeThreadNum(const GPUContext& dev_ctx,
+                            int* num_thread,
+                            int alternative_num_thread = 512) {
+  if (dev_ctx.GetComputeCapability() == 53 ||
+      dev_ctx.GetComputeCapability() == 62) {
+    *num_thread = alternative_num_thread;
+  }
+}
+#endif
+
+struct GpuLaunchConfig {
+ public:
+  GpuLaunchConfig() {}
+
+  size_t GetThreadNum() const { return GetBlockSize() * GetGridSize(); }
+
+  size_t GetGridSize() const {
+    return block_per_grid.x * block_per_grid.y * block_per_grid.z;
+  }
+
+  size_t GetBlockSize() const {
+    return thread_per_block.x * thread_per_block.y * thread_per_block.z;
+  }
+
+  int compute_capability = 0;
+  dim3 thread_per_block = dim3(1, 1, 1);
+  dim3 block_per_grid = dim3(1, 1, 1);
+};
+
+/* According to NVIDIA, if number of threads per block is 64/128/256/512,
+ * cuda performs better. And number of blocks should be greater (at least
+ * 2x~4x) than number of SMs. Hence, SM count is took into account within
+ * this function to determine the right number of threads per block. */
+inline GpuLaunchConfig GetGpuLaunchConfig1D(const GPUContext& dev_ctx,
+                                            int64_t numel,
+                                            int vec_size = 1) {
+  PADDLE_ENFORCE_GE(numel,
+                    0,
+                    common::errors::InvalidArgument(
+                        "numel is expected to be greater than or equal 0,"
+                        " but received %d.",
+                        numel));
+  PADDLE_ENFORCE_GE(
+      vec_size,
+      1,
+      common::errors::InvalidArgument(
+          "vec_size is expected greater than 0, but received %d.", vec_size));
+  // Get compute_capability
+  const int capability = dev_ctx.GetComputeCapability();
+  // If thread number per block is 64/128/256/512, cuda performs better.
+  int limit_threads =
+      std::min(PREDEFINED_BLOCK_SIZE, dev_ctx.GetMaxThreadsPerBlock());
+#ifdef WITH_NV_JETSON
+  if (capability == 53 || capability == 62) {
+    limit_threads = 512;
+  }
+#endif
+  int threads = limit_threads;
+  int sm_count = dev_ctx.GetSMCount();
+  int64_t active_threads_num = numel / vec_size;
+  if (active_threads_num / (sm_count << 1) < limit_threads) {
+    // Round up threads number into an exponential multiple of 2, while number
+    // of active blocks is about twice of SM, to acquire better performance.
+    threads = RoundToPowerOfTwo(active_threads_num / (sm_count << 1));
+  } else if (active_threads_num / (sm_count << 2) < limit_threads) {
+    // Round up threads number into an exponential multiple of 2, while number
+    // of active blocks is about 4 times of SM, to acquire better performance.
+    threads = RoundToPowerOfTwo(active_threads_num / (sm_count << 2));
+  }
+  // Number of threads per block shall be larger than 64.
+  threads = std::max(64, threads);
+  int64_t blocks = DivUp<int64_t>(DivUp<int64_t>(numel, vec_size), threads);
+  int64_t limit_blocks = dev_ctx.GetCUDAMaxGridDimSize()[0];
+  if (blocks > limit_blocks) {
+    blocks = limit_blocks;
+  }
+
+  GpuLaunchConfig config;
+  config.thread_per_block.x = threads;
+  config.block_per_grid.x = static_cast<uint32_t>(blocks);
+  config.compute_capability = capability;
+
+  VLOG(7) << "Get 1-D launch config: numel=" << numel
+          << ", vec_size=" << vec_size << ", block_size=" << threads
+          << ", grid_size=" << blocks << ", limit_blocks=" << limit_blocks
+          << ", limit_threads=" << limit_threads;
+  return config;
+}
+
+inline GpuLaunchConfig GetGpuLaunchConfig2D(const GPUContext& dev_ctx,
+                                            int64_t x_dim,
+                                            int64_t y_dim) {
+  PADDLE_ENFORCE_GT(
+      x_dim,
+      0,
+      common::errors::InvalidArgument("x dim number should greater than 0,"
+                                      " but received value is: %d",
+                                      x_dim));
+  PADDLE_ENFORCE_GT(
+      y_dim,
+      0,
+      common::errors::InvalidArgument("y dim number should greater than 0,"
+                                      " but received value is: %d",
+                                      y_dim));
+
+  const int kThreadsPerBlock = 256;
+  int block_cols = std::min<int64_t>(x_dim, kThreadsPerBlock);
+  int block_rows = std::max(kThreadsPerBlock / block_cols, 1);
+
+  int max_physical_threads = dev_ctx.GetMaxPhysicalThreadCount();
+  const int max_blocks = std::max(max_physical_threads / kThreadsPerBlock, 1);
+
+  GpuLaunchConfig config;
+  // Noticed, block size is not align to 32, if needed do it yourself.
+  config.thread_per_block = dim3(block_cols, block_rows, 1);
+
+  int grid_x = std::min<int64_t>(DivUp<int64_t>(x_dim, block_cols), max_blocks);
+  int grid_y = std::min<int64_t>(max_blocks / grid_x,
+                                 std::max<int64_t>(y_dim / block_rows, 1));
+
+  config.block_per_grid = dim3(grid_x, grid_y, 1);
+  return config;
+}
+
+static inline int GetLastPow2(int n) {
+  n |= (n >> 1);
+  n |= (n >> 2);
+  n |= (n >> 4);
+  n |= (n >> 8);
+  n |= (n >> 16);
+  return std::max(1, n - (n >> 1));
+}
+
+inline GpuLaunchConfig GetGpuLaunchConfig3D(const GPUContext& dev_ctx,
+                                            int num_img,
+                                            int height,
+                                            int width) {
+  const int kThreadsPerBlock = 256;
+  int max_threads_per_block = dev_ctx.GetMaxThreadsPerBlock();  // 1024
+  int max_threads = std::min(kThreadsPerBlock, max_threads_per_block);
+
+  int block_x = std::min(GetLastPow2(width), max_threads);
+  int block_y = std::min(GetLastPow2(height), max_threads / block_x);
+  int block_z = std::min(num_img, max_threads / block_x / block_y);
+
+  std::array<unsigned int, 3> max_grid_dim = dev_ctx.GetCUDAMaxGridDimSize();
+  unsigned int grid_x =
+      std::min(max_grid_dim[0], DivUp<unsigned int>(width, block_x));
+  unsigned int grid_y =
+      std::min(max_grid_dim[1], DivUp<unsigned int>(height, block_y));
+  unsigned int grid_z =
+      std::min(max_grid_dim[2], DivUp<unsigned int>(num_img, block_z * 4));
+
+  const int capability = dev_ctx.GetComputeCapability();
+  GpuLaunchConfig config;
+  config.compute_capability = capability;
+  config.thread_per_block = dim3(block_x, block_y, block_z);
+  config.block_per_grid = dim3(grid_x, grid_y, grid_z);
+  return config;
+}
+
+template <typename Context>
+void LimitGridDim(const Context& dev_ctx, dim3* grid_dim) {
+  auto max_grid_dim =
+      reinterpret_cast<const GPUContext&>(dev_ctx).GetCUDAMaxGridDimSize();
+  grid_dim->x = grid_dim->x < max_grid_dim[0] ? grid_dim->x : max_grid_dim[0];
+  grid_dim->y = grid_dim->y < max_grid_dim[1] ? grid_dim->y : max_grid_dim[1];
+  grid_dim->z = grid_dim->z < max_grid_dim[2] ? grid_dim->z : max_grid_dim[2];
+}
+}  // namespace gpu
+}  // namespace backends
+}  // namespace phi
+
+#endif

paddle/phi/backends/gpu/gpu_resources.cc

@@ -0,0 +1,324 @@
+// Copyright (c) 2022 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.
+
+#include "paddle/phi/backends/gpu/gpu_resources.h"
+
+#include <set>
+
+#include <map>
+#include "paddle/phi/api/include/tensor.h"
+#include "paddle/phi/backends/gpu/gpu_decls.h"
+#include "paddle/phi/backends/gpu/gpu_info.h"
+#include "paddle/phi/common/place.h"
+#include "paddle/phi/core/allocator.h"
+#ifdef PADDLE_WITH_CUDA
+#include "paddle/phi/backends/dynload/cublas.h"
+#include "paddle/phi/backends/dynload/cublasLt.h"
+#include "paddle/phi/backends/dynload/cudnn.h"
+#include "paddle/phi/backends/dynload/cusolver.h"
+#include "paddle/phi/backends/dynload/cusparse.h"
+#if !defined(__APPLE__) && defined(PADDLE_WITH_NCCL)
+#include "paddle/phi/backends/dynload/nccl.h"
+#endif  // !defined(__APPLE__) && defined(PADDLE_WITH_NCCL)
+#endif  // PADDLE_WITH_CUDA
+#ifdef PADDLE_WITH_HIP
+#include "paddle/phi/backends/dynload/hipblasLt.h"
+#include "paddle/phi/backends/dynload/rocsparse.h"
+#endif
+
+#include "glog/logging.h"
+#include "unsupported/Eigen/CXX11/Tensor"
+
+#include "paddle/phi/core/enforce.h"
+
+namespace phi {
+
+void InitGpuProperties(Place place,
+                       int* compute_capability,
+                       int* runtime_version,
+                       int* driver_version,
+                       int* multi_process,
+                       int* max_threads_per_mp,
+                       int* max_threads_per_block,
+                       std::array<unsigned int, 3>* max_grid_dim_size) {
+  backends::gpu::GPUDeviceGuard guard(place.GetDeviceId());
+  *compute_capability =
+      backends::gpu::GetGPUComputeCapability(place.GetDeviceId());
+  *multi_process = backends::gpu::GetGPUMultiProcessors(place.GetDeviceId());
+  *max_threads_per_mp =
+      backends::gpu::GetGPUMaxThreadsPerMultiProcessor(place.GetDeviceId());
+  *max_grid_dim_size = backends::gpu::GetGpuMaxGridDimSize(place.GetDeviceId());
+  *max_threads_per_block =
+      backends::gpu::GetGPUMaxThreadsPerBlock(place.GetDeviceId());
+  *driver_version = backends::gpu::GetGPUDriverVersion(place.GetDeviceId());
+  *runtime_version = backends::gpu::GetGPURuntimeVersion(place.GetDeviceId());
+
+  const gpuDeviceProp& prop =
+      backends::gpu::GetDeviceProperties(place.GetDeviceId());
+
+#ifdef PADDLE_WITH_CUDA
+  static const std::set<int> compiled_archs{CUDA_REAL_ARCHS};
+  // Make sure compiled cuda arch is as same as runtime cuda arch.
+  if (compiled_archs.find(*compute_capability) == compiled_archs.cend() &&
+      compiled_archs.find(prop.major * 10) == compiled_archs.cend()) {
+    static std::atomic<bool> once_flag(false);
+    if (!once_flag.exchange(true)) {
+      std::string compile_arch_str = "";
+      for (const int32_t& arch : compiled_archs) {
+        compile_arch_str += std::to_string(arch) + " ";
+      }
+      std::map<int, std::string> arch_computing_mapping_table = {
+          {20, "Fermi"},
+          {30, "Kepler"},
+          {35, "Kapler"},
+          {37, "Kepler"},
+          {50, "Maxwell"},
+          {52, "Maxwell"},
+          {60, "Pascal"},
+          {61, "Pascal"},
+          {70, "Volta"},
+          {75, "Turing"},
+          {80, "Ampere"},
+          {86, "Ampere"},
+          {89, "Ada Lovelace"},
+          {90, "Hopper"},
+          {100, "Blackwell"},
+          {120, "Blackwell"},
+      };
+      if (!arch_computing_mapping_table.count(*compute_capability)) {
+        LOG(ERROR) << "Mismatched GPU Architecture: "
+                   << "The installed PaddlePaddle package was compiled for "
+                   << compile_arch_str << ",but your current GPU is "
+                   << *compute_capability
+                   << " Solution: Install the correct wheel package built for "
+                      "your GPU "
+                   << "from the official PaddlePaddle website: "
+                   << "https://www.paddlepaddle.org.cn/";
+
+        throw std::runtime_error("Unsupported GPU architecture");
+      }
+    }
+  }
+#endif
+
+  // TODO(wilber): glog may be replaced in the future?
+  LOG_FIRST_N(WARNING, 1) << "Please NOTE: device: "
+                          << static_cast<int>(place.device)
+                          << ", GPU Compute Capability: "
+                          << *compute_capability / 10 << "."
+                          << *compute_capability % 10
+                          << ", Driver API Version: " << *driver_version / 1000
+                          << "." << (*driver_version % 100) / 10
+                          << ", Runtime API Version: "
+                          << *runtime_version / 1000 << "."
+                          << (*runtime_version % 100) / 10;
+#ifdef PADDLE_WITH_HIP
+  size_t miopen_major, miopen_minor, miopen_patch;
+  PADDLE_ENFORCE_GPU_SUCCESS(
+      dynload::miopenGetVersion(&miopen_major, &miopen_minor, &miopen_patch));
+  auto cudnn_dso_ver =
+      (miopen_major * 1000 + miopen_minor * 10 + miopen_patch) / 10;
+  auto compile_miopen_version = MIOPEN_VERSION / 10;
+  if (cudnn_dso_ver < static_cast<size_t>(compile_miopen_version)) {
+    LOG_FIRST_N(WARNING, 1)
+        << "WARNING: device: " << static_cast<int>(place.device)
+        << ". The installed Paddle is compiled with MIOPEN "
+        << compile_miopen_version / 100 << "." << compile_miopen_version % 100
+        << ", but MIOPEN version in your machine is " << cudnn_dso_ver / 100
+        << "." << cudnn_dso_ver % 100
+        << ", which may cause serious incompatible bug. "
+        << "Please recompile or reinstall Paddle with compatible MIOPEN "
+           "version.";
+  }
+#endif
+}
+
+void InitStream(gpuStream_t* stream) {
+#ifdef PADDLE_WITH_HIP
+  PADDLE_ENFORCE_GPU_SUCCESS(
+      hipStreamCreateWithPriority(stream, hipStreamDefault, 0));
+#else
+  PADDLE_ENFORCE_GPU_SUCCESS(
+      cudaStreamCreateWithPriority(stream, cudaStreamDefault, 0));
+#endif
+}
+
+void DestroyStream(gpuStream_t stream) {
+  if (stream != nullptr) {
+#ifdef PADDLE_WITH_HIP
+    PADDLE_ENFORCE_GPU_SUCCESS(hipStreamDestroy(stream));
+#else
+    PADDLE_ENFORCE_GPU_SUCCESS(cudaStreamDestroy(stream));
+#endif
+  }
+  stream = nullptr;
+}
+
+#ifndef PADDLE_WITH_CUSTOM_DEVICE
+void InitBlasHandle(blasHandle_t* blas_handle, gpuStream_t stream) {
+#ifdef PADDLE_WITH_HIP
+  phi::dynload::rocblas_create_handle(blas_handle);
+  phi::dynload::rocblas_set_stream(*blas_handle, stream);
+#else   // PADDLE_WITH_CUDA
+  PADDLE_RETRY_CUDA_SUCCESS(phi::dynload::cublasCreate(blas_handle));
+  PADDLE_RETRY_CUDA_SUCCESS(
+      phi::dynload::cublasSetStream(*blas_handle, stream));
+#endif  // PADDLE_WITH_HIP
+}
+
+void DestroyBlasHandle(blasHandle_t handle) {
+#ifdef PADDLE_WITH_HIP
+  if (handle != nullptr) {
+    phi::dynload::rocblas_destroy_handle(handle);
+    handle = nullptr;
+  }
+#else
+  if (handle != nullptr) {
+    phi::dynload::cublasDestroy(handle);
+    handle = nullptr;
+  }
+#endif  // PADDLE_WITH_HIP
+}
+
+void InitBlasLtHandle(blasLtHandle_t* blaslt_handle) {
+#if defined(PADDLE_WITH_CUDA) && CUDA_VERSION >= 11060
+  phi::dynload::cublasLtCreate(blaslt_handle);
+#elif defined(PADDLE_WITH_HIP)
+  phi::dynload::hipblasLtCreate(blaslt_handle);
+#endif
+}
+
+void DestroyBlasLtHandle(blasLtHandle_t handle) {
+#if defined(PADDLE_WITH_CUDA) && CUDA_VERSION >= 11060
+  if (handle != nullptr) {
+    phi::dynload::cublasLtDestroy(handle);
+    handle = nullptr;
+  }
+#elif defined(PADDLE_WITH_HIP)
+  if (handle != nullptr) {
+    phi::dynload::hipblasLtDestroy(handle);
+    handle = nullptr;
+  }
+#endif
+}
+
+void InitDnnHandle(dnnHandle_t* handle, gpuStream_t stream, Place place) {
+  if (phi::dynload::HasCUDNN()) {
+#ifdef PADDLE_WITH_HIP
+    size_t miopen_major, miopen_minor, miopen_patch;
+    PADDLE_ENFORCE_GPU_SUCCESS(
+        dynload::miopenGetVersion(&miopen_major, &miopen_minor, &miopen_patch));
+    auto local_miopen_version =
+        (miopen_major * 1000 + miopen_minor * 10 + miopen_patch) / 10;
+    auto compile_miopen_version = MIOPEN_VERSION / 10;
+    if (local_miopen_version < static_cast<size_t>(compile_miopen_version)) {
+      LOG_FIRST_N(WARNING, 1)
+          << "WARNING: device: " << static_cast<int>(place.device)
+          << ". The installed Paddle is compiled with MIOPEN "
+          << compile_miopen_version / 100 << "." << compile_miopen_version % 100
+          << ", but MIOPEN version in your machine is "
+          << local_miopen_version / 100 << "." << local_miopen_version % 100
+          << ", which may cause serious incompatible bug. "
+          << "Please recompile or reinstall Paddle with compatible MIOPEN "
+             "version.";
+    }
+    PADDLE_ENFORCE_GPU_SUCCESS(dynload::miopenCreate(handle));
+    PADDLE_ENFORCE_GPU_SUCCESS(dynload::miopenSetStream(*handle, stream));
+#else
+    auto version = phi::dynload::cudnnGetVersion();
+    auto local_cudnn_major =
+        (version < 9000) ? version / 1000 : version / 10000;
+    auto local_cudnn_minor =
+        (version < 9000) ? (version % 1000) / 100 : (version % 10000) / 100;
+    if (version < static_cast<size_t>(CUDNN_VERSION)) {
+      LOG_FIRST_N(WARNING, 1)
+          << "WARNING: device: " << static_cast<int>(place.device)
+          << ". The installed Paddle is compiled with CUDNN " << CUDNN_MAJOR
+          << "." << CUDNN_MINOR << ", but CUDNN version in your machine is "
+          << local_cudnn_major << "." << local_cudnn_minor
+          << ", which may cause serious incompatible bug. "
+          << "Please recompile or reinstall Paddle with compatible CUDNN "
+             "version.";
+    }
+    PADDLE_RETRY_CUDA_SUCCESS(phi::dynload::cudnnCreate(handle));
+    PADDLE_RETRY_CUDA_SUCCESS(phi::dynload::cudnnSetStream(*handle, stream));
+#endif
+  } else {
+    *handle = nullptr;
+  }
+}
+
+void DestroyDnnHandle(dnnHandle_t handle) {
+#ifdef PADDLE_WITH_HIP
+  if (handle != nullptr) {
+    PADDLE_ENFORCE_GPU_SUCCESS(phi::dynload::miopenDestroy(handle));
+    handle = nullptr;
+  }
+#else
+  if (handle != nullptr) {
+    PADDLE_ENFORCE_GPU_SUCCESS(phi::dynload::cudnnDestroy(handle));
+    handle = nullptr;
+  }
+#endif  // PADDLE_WITH_HIP
+}
+
+void InitSolverHandle(solverHandle_t* handle, gpuStream_t stream) {
+#if defined(PADDLE_WITH_CUDA)
+  PADDLE_RETRY_CUDA_SUCCESS(phi::dynload::cusolverDnCreate(handle));
+  PADDLE_RETRY_CUDA_SUCCESS(phi::dynload::cusolverDnSetStream(*handle, stream));
+#elif defined(PADDLE_WITH_HIP)
+  phi::dynload::rocblas_create_handle(handle);
+  phi::dynload::rocblas_set_stream(*handle, stream);
+#endif
+}
+
+void DestroySolverHandle(solverHandle_t solver_handle) {
+#if defined(PADDLE_WITH_CUDA)
+  if (solver_handle != nullptr) {
+    PADDLE_ENFORCE_GPU_SUCCESS(phi::dynload::cusolverDnDestroy(solver_handle));
+    solver_handle = nullptr;
+  }
+#elif defined(PADDLE_WITH_HIP)
+  if (solver_handle != nullptr) {
+    phi::dynload::rocblas_destroy_handle(solver_handle);
+    solver_handle = nullptr;
+  }
+#endif
+}
+
+void InitSparseHandle(sparseHandle_t* handle, gpuStream_t stream) {
+#if defined(PADDLE_WITH_CUDA)
+  PADDLE_RETRY_CUDA_SUCCESS(dynload::cusparseCreate(handle));
+  PADDLE_RETRY_CUDA_SUCCESS(dynload::cusparseSetStream(*handle, stream));
+#elif defined(PADDLE_WITH_HIP)
+  phi::dynload::rocsparse_create_handle(handle);
+  phi::dynload::rocsparse_set_stream(*handle, stream);
+#endif
+}
+
+void DestroySparseHandle(sparseHandle_t handle) {
+#ifdef PADDLE_WITH_CUDA
+  if (handle != nullptr) {
+    PADDLE_RETRY_CUDA_SUCCESS(dynload::cusparseDestroy(handle));
+    handle = nullptr;
+  }
+#elif defined(PADDLE_WITH_HIP)
+  if (handle != nullptr) {
+    phi::dynload::rocsparse_destroy_handle(handle);
+    handle = nullptr;
+  }
+#endif
+}
+#endif
+}  // namespace phi

paddle/phi/backends/gpu/gpu_resources.h

@@ -0,0 +1,57 @@
+// Copyright (c) 2022 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.
+#pragma once
+
+#include <array>
+
+#include "paddle/phi/backends/gpu/gpu_decls.h"
+#include "paddle/phi/common/place.h"
+#include "paddle/phi/core/device_context.h"
+
+namespace phi {
+
+PADDLE_API void InitGpuProperties(
+    Place place,
+    int* compute_capability,
+    int* runtime_version,
+    int* driver_version,
+    int* multi_process,
+    int* max_threads_per_mp,
+    int* max_threads_per_block,
+    std::array<unsigned int, 3>* max_grid_dim_size);
+
+PADDLE_API void InitStream(gpuStream_t* stream);
+PADDLE_API void DestroyStream(gpuStream_t stream);
+
+#ifndef PADDLE_WITH_CUSTOM_DEVICE
+PADDLE_API void InitBlasHandle(blasHandle_t* blas_handle, gpuStream_t stream);
+PADDLE_API void DestroyBlasHandle(blasHandle_t handle);
+
+PADDLE_API void InitBlasLtHandle(blasLtHandle_t* blaslt_handle);
+PADDLE_API void DestroyBlasLtHandle(blasLtHandle_t handle);
+
+PADDLE_API void InitDnnHandle(dnnHandle_t* handle,
+                              gpuStream_t stream,
+                              Place place);
+PADDLE_API void DestroyDnnHandle(dnnHandle_t handle);
+
+PADDLE_API void InitSolverHandle(solverHandle_t* handle, gpuStream_t stream);
+PADDLE_API void DestroySolverHandle(solverHandle_t solver_handle);
+
+PADDLE_API void InitSparseHandle(sparseHandle_t* handle, gpuStream_t stream);
+PADDLE_API void DestroySparseHandle(sparseHandle_t handle);
+#endif
+// void InitDnnWorkspace();
+
+}  // namespace phi

paddle/phi/backends/gpu/gpu_types.h

@@ -0,0 +1,173 @@
+// Copyright (c) 2021 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.
+
+#pragma once
+
+#include "paddle/phi/backends/gpu/forwards.h"
+#include "paddle/phi/backends/gpu/gpu_decls.h"
+
+#if defined(PADDLE_WITH_CUDA) || defined(PADDLE_WITH_HIP)
+
+#ifdef PADDLE_WITH_HIP
+#include "paddle/phi/backends/dynload/miopen.h"
+#include "paddle/phi/backends/dynload/rocblas.h"
+#else  // PADDLE_WITH_CUDA
+#ifndef PADDLE_WITH_CUSTOM_DEVICE
+#include "paddle/phi/backends/dynload/cublas.h"
+#include "paddle/phi/backends/dynload/cudnn.h"
+#else
+#include <cuda_runtime.h>
+#endif
+#endif
+
+namespace phi {
+
+// Note(qili93): CUDA Runtime API supported by HIP
+// https://github.com/ROCm/HIPIFY/blob/master/doc/markdown/CUDA_Runtime_API_functions_supported_by_HIP.md
+
+#ifdef PADDLE_WITH_HIP
+#define DECLARE_TYPE_FOR_GPU(GPU_TYPE, CUDA_TYPE, ROCM_TYPE) \
+  using GPU_TYPE = ROCM_TYPE;
+
+#else  // PADDLE_WITH_CUDA
+
+#define DECLARE_TYPE_FOR_GPU(GPU_TYPE, CUDA_TYPE, ROCM_TYPE) \
+  using GPU_TYPE = CUDA_TYPE;
+#endif
+
+DECLARE_TYPE_FOR_GPU(gpuError_t, cudaError_t, hipError_t);
+DECLARE_TYPE_FOR_GPU(gpuMemcpyKind, cudaMemcpyKind, hipMemcpyKind);
+DECLARE_TYPE_FOR_GPU(gpuDeviceProp, cudaDeviceProp, hipDeviceProp_t);
+#ifndef PADDLE_WITH_CUSTOM_DEVICE
+DECLARE_TYPE_FOR_GPU(dnnDataType_t, cudnnDataType_t, miopenDataType_t);
+DECLARE_TYPE_FOR_GPU(dnnPoolingMode_t, cudnnPoolingMode_t, miopenPoolingMode_t);
+DECLARE_TYPE_FOR_GPU(dnnTensorFormat_t,
+                     cudnnTensorFormat_t,
+                     miopenTensorFormat_t);
+DECLARE_TYPE_FOR_GPU(dnnActivationMode_t,
+                     cudnnActivationMode_t,
+                     miopenActivationMode_t);
+#endif
+DECLARE_TYPE_FOR_GPU(gpuGraph_t, cudaGraph_t, hipGraph_t);
+DECLARE_TYPE_FOR_GPU(gpuFunction_t, cudaFunction_t, hipFunction_t);
+DECLARE_TYPE_FOR_GPU(gpuGraphExec_t, cudaGraphExec_t, hipGraphExec_t);
+DECLARE_TYPE_FOR_GPU(gpuGraphNode_t, cudaGraphNode_t, hipGraphNode_t);
+DECLARE_TYPE_FOR_GPU(gpuGraphNodeType, cudaGraphNodeType, hipGraphNodeType);
+DECLARE_TYPE_FOR_GPU(gpuKernelNodeParams,
+                     cudaKernelNodeParams,
+                     hipKernelNodeParams);
+DECLARE_TYPE_FOR_GPU(gpuStreamCaptureMode,
+                     cudaStreamCaptureMode,
+                     hipStreamCaptureMode);
+DECLARE_TYPE_FOR_GPU(gpuStreamCaptureStatus,
+                     cudaStreamCaptureStatus,
+                     hipStreamCaptureStatus);
+
+#undef DECLARE_TYPE_FOR_GPU
+
+#ifdef PADDLE_WITH_HIP
+#define DECLARE_CONSTANT_FOR_GPU(GPU_CV, CUDA_CV, ROCM_CV) \
+  constexpr auto GPU_CV = ROCM_CV;
+#else  // PADDLE_WITH_CUDA
+#define DECLARE_CONSTANT_FOR_GPU(GPU_CV, CUDA_CV, ROCM_CV) \
+  constexpr auto GPU_CV = CUDA_CV;
+#endif
+
+DECLARE_CONSTANT_FOR_GPU(gpuErrorOutOfMemory,
+                         cudaErrorMemoryAllocation,
+                         hipErrorOutOfMemory);
+DECLARE_CONSTANT_FOR_GPU(gpuErrorNotReady, cudaErrorNotReady, hipErrorNotReady);
+DECLARE_CONSTANT_FOR_GPU(gpuSuccess, cudaSuccess, hipSuccess);
+
+DECLARE_CONSTANT_FOR_GPU(gpuMemcpyHostToDevice,
+                         cudaMemcpyKind::cudaMemcpyHostToDevice,
+                         hipMemcpyKind::hipMemcpyHostToDevice);
+DECLARE_CONSTANT_FOR_GPU(gpuMemcpyDeviceToHost,
+                         cudaMemcpyKind::cudaMemcpyDeviceToHost,
+                         hipMemcpyKind::hipMemcpyDeviceToHost);
+DECLARE_CONSTANT_FOR_GPU(gpuMemcpyDeviceToDevice,
+                         cudaMemcpyKind::cudaMemcpyDeviceToDevice,
+                         hipMemcpyKind::hipMemcpyDeviceToDevice);
+DECLARE_CONSTANT_FOR_GPU(gpuEventDisableTiming,
+                         cudaEventDisableTiming,
+                         hipEventDisableTiming);
+DECLARE_CONSTANT_FOR_GPU(gpuStreamNonBlocking,
+                         cudaStreamNonBlocking,
+                         hipStreamNonBlocking);
+DECLARE_CONSTANT_FOR_GPU(gpuStreamCaptureModeThreadLocal,
+                         cudaStreamCaptureModeThreadLocal,
+                         hipStreamCaptureModeThreadLocal);
+DECLARE_CONSTANT_FOR_GPU(gpuStreamCaptureModeRelaxed,
+                         cudaStreamCaptureModeRelaxed,
+                         hipStreamCaptureModeRelaxed);
+DECLARE_CONSTANT_FOR_GPU(gpuStreamCaptureStatusActive,
+                         cudaStreamCaptureStatusActive,
+                         hipStreamCaptureStatusActive);
+DECLARE_CONSTANT_FOR_GPU(gpuGraphNodeTypeKernel,
+                         cudaGraphNodeTypeKernel,
+                         hipGraphNodeTypeKernel);
+
+#undef DECLARE_CONSTANT_FOR_GPU
+
+#ifdef PADDLE_WITH_HIP
+#define DECLARE_FUNCTION_FOR_GPU(GPU_FUNC, CUDA_FUNC, ROCM_FUNC) \
+  const auto GPU_FUNC = ROCM_FUNC;
+#elif defined(PADDLE_WITH_CUDA)  // PADDLE_WITH_CUDA
+#define DECLARE_FUNCTION_FOR_GPU(GPU_FUNC, CUDA_FUNC, ROCM_FUNC) \
+  const auto GPU_FUNC = CUDA_FUNC;
+#endif
+
+DECLARE_FUNCTION_FOR_GPU(gpuGraphGetNodes, cudaGraphGetNodes, hipGraphGetNodes);
+DECLARE_FUNCTION_FOR_GPU(gpuGraphGetEdges, cudaGraphGetEdges, hipGraphGetEdges);
+DECLARE_FUNCTION_FOR_GPU(gpuGraphLaunch, cudaGraphLaunch, hipGraphLaunch);
+DECLARE_FUNCTION_FOR_GPU(gpuGraphDestroy, cudaGraphDestroy, hipGraphDestroy);
+DECLARE_FUNCTION_FOR_GPU(gpuGraphExecDestroy,
+                         cudaGraphExecDestroy,
+                         hipGraphExecDestroy);
+DECLARE_FUNCTION_FOR_GPU(gpuGraphNodeGetType,
+                         cudaGraphNodeGetType,
+                         hipGraphNodeGetType);
+DECLARE_FUNCTION_FOR_GPU(gpuGraphExecKernelNodeSetParams,
+                         cudaGraphExecKernelNodeSetParams,
+                         hipGraphExecKernelNodeSetParams);
+DECLARE_FUNCTION_FOR_GPU(gpuGraphKernelNodeGetParams,
+                         cudaGraphKernelNodeGetParams,
+                         hipGraphKernelNodeGetParams);
+DECLARE_FUNCTION_FOR_GPU(gpuStreamCreateWithPriority,
+                         cudaStreamCreateWithPriority,
+                         hipStreamCreateWithPriority);
+DECLARE_FUNCTION_FOR_GPU(gpuStreamBeginCapture,
+                         cudaStreamBeginCapture,
+                         hipStreamBeginCapture);
+DECLARE_FUNCTION_FOR_GPU(gpuStreamEndCapture,
+                         cudaStreamEndCapture,
+                         hipStreamEndCapture);
+DECLARE_FUNCTION_FOR_GPU(gpuStreamGetCaptureInfo,
+                         cudaStreamGetCaptureInfo,
+                         hipStreamGetCaptureInfo);
+DECLARE_FUNCTION_FOR_GPU(gpuEventCreateWithFlags,
+                         cudaEventCreateWithFlags,
+                         hipEventCreateWithFlags);
+DECLARE_FUNCTION_FOR_GPU(gpuEventRecord, cudaEventRecord, hipEventRecord);
+DECLARE_FUNCTION_FOR_GPU(gpuEventDestroy, cudaEventDestroy, hipEventDestroy);
+DECLARE_FUNCTION_FOR_GPU(gpuEventQuery, cudaEventQuery, hipEventQuery);
+DECLARE_FUNCTION_FOR_GPU(gpuEventSynchronize,
+                         cudaEventSynchronize,
+                         hipEventSynchronize);
+
+#undef DECLARE_FUNCTION_FOR_GPU
+
+}  // namespace phi
+
+#endif  // defined(PADDLE_WITH_CUDA) || defined(PADDLE_WITH_HIP)

paddle/phi/backends/gpu/gpu_utils.h

@@ -0,0 +1,104 @@
+// 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.
+
+#pragma once
+
+#define EIGEN_USE_GPU
+
+#include <array>
+
+#include "paddle/phi/core/enforce.h"
+#include "unsupported/Eigen/CXX11/Tensor"
+
+namespace phi {
+namespace funcs {
+
+template <typename T, int Size, T DefaultValue>
+struct DeviceArray {
+  EIGEN_DEVICE_FUNC EIGEN_STRONG_INLINE const T& operator[](int index) const {
+    return data[index];
+  }
+  EIGEN_DEVICE_FUNC EIGEN_STRONG_INLINE T& operator[](int index) {
+    return data[index];
+  }
+  EIGEN_DEVICE_FUNC EIGEN_STRONG_INLINE DeviceArray() {
+    for (int i = 0; i < Size; i++) {
+      data[i] = DefaultValue;
+    }
+  }
+  EIGEN_DEVICE_FUNC EIGEN_STRONG_INLINE DeviceArray(T a0) {
+    data[0] = a0;
+    for (int i = 1; i < Size; i++) {
+      data[i] = DefaultValue;
+    }
+  }
+  EIGEN_DEVICE_FUNC EIGEN_STRONG_INLINE DeviceArray(T a0, T a1) {
+    data[0] = a0;
+    data[1] = a1;
+    for (int i = 2; i < Size; i++) {
+      data[i] = DefaultValue;
+    }
+  }
+  EIGEN_DEVICE_FUNC EIGEN_STRONG_INLINE DeviceArray(T a0, T a1, T a2) {
+    data[0] = a0;
+    data[1] = a1;
+    data[2] = a2;
+    for (int i = 3; i < Size; i++) {
+      data[i] = DefaultValue;
+    }
+  }
+  EIGEN_STRONG_INLINE DeviceArray(const std::array<T, Size>& sa) {
+    for (int i = 0; i < Size; i++) {
+      data[i] = sa[i];
+    }
+  }
+  T data[Size];
+};
+
+template <typename T>
+struct Dim3 : DeviceArray<T, 3, 1> {
+  typedef DeviceArray<T, 3, 1> Base;
+  EIGEN_DEVICE_FUNC EIGEN_STRONG_INLINE Dim3() : Base() {}
+  EIGEN_DEVICE_FUNC EIGEN_STRONG_INLINE Dim3(T a0, T a1, T a2)
+      : Base(a0, a1, a2) {}
+};
+
+// Flat index with real dimension
+template <typename IndexType = int>
+EIGEN_DEVICE_FUNC EIGEN_STRONG_INLINE IndexType
+FlatTensorIndex(const Dim3<IndexType>& index, const Dim3<IndexType>& dims) {
+  IndexType flat_index = index[0];
+#pragma unroll
+  for (int i = 1; i < 3; ++i) {
+    flat_index = flat_index * dims[i] + index[i];
+  }
+  return flat_index;
+}
+
+// Convert index to tensor index with dimension.
+template <typename IndexType = int>
+EIGEN_DEVICE_FUNC EIGEN_STRONG_INLINE Dim3<IndexType> ConvertTensorIndex(
+    IndexType index, const Dim3<IndexType>& dims) {
+  Dim3<IndexType> tensor_index;
+#pragma unroll
+  for (int i = 2; i >= 0; --i) {
+    IndexType new_index = index / dims[i];
+    tensor_index[i] = static_cast<int>(index - dims[i] * new_index);
+    index = new_index;
+  }
+  return tensor_index;
+}
+
+}  // namespace funcs
+}  // namespace phi

paddle/phi/backends/gpu/rocm/hip_graph.cc

@@ -0,0 +1,333 @@
+// Copyright (c) 2022 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.
+
+#include "paddle/phi/backends/gpu/rocm/hip_graph.h"
+#include "glog/logging.h"
+#include "paddle/common/flags.h"
+
+COMMON_DECLARE_bool(use_cuda_malloc_async_allocator);
+COMMON_DECLARE_bool(auto_free_cudagraph_allocations_on_launch);
+
+#ifdef PADDLE_WITH_HIP
+
+namespace phi {
+namespace backends {
+namespace gpu {
+
+std::unique_ptr<CUDAGraph> CUDAGraph::capturing_graph_{nullptr};
+paddle::optional<std::thread::id> CUDAGraph::capturing_thread_id_{paddle::none};
+
+static std::vector<hipGraphNode_t> ToposortCUDAGraph(hipGraph_t graph) {
+  size_t num_nodes;
+  PADDLE_ENFORCE_GPU_SUCCESS(hipGraphGetNodes(graph, nullptr, &num_nodes));
+  std::vector<hipGraphNode_t> nodes(num_nodes);
+  PADDLE_ENFORCE_GPU_SUCCESS(hipGraphGetNodes(graph, nodes.data(), &num_nodes));
+
+  size_t num_edges;
+  PADDLE_ENFORCE_GPU_SUCCESS(
+      hipGraphGetEdges(graph, nullptr, nullptr, &num_edges));
+  std::vector<hipGraphNode_t> from(num_edges), to(num_edges);
+  PADDLE_ENFORCE_GPU_SUCCESS(
+      hipGraphGetEdges(graph, from.data(), to.data(), &num_edges));
+
+  std::unordered_map<hipGraphNode_t, std::unordered_set<hipGraphNode_t>>
+      in_edges, out_edges;
+  for (auto node : nodes) {
+    in_edges[node];
+    out_edges[node];
+  }
+
+  for (size_t i = 0; i < num_edges; ++i) {
+    in_edges[to[i]].insert(from[i]);
+    out_edges[from[i]].insert(to[i]);
+  }
+
+  std::queue<hipGraphNode_t> q;
+  for (const auto &pair : in_edges) {
+    if (pair.second.empty()) {
+      q.push(pair.first);
+    }
+  }
+
+  nodes.clear();
+  while (!q.empty()) {
+    auto cur = q.front();
+    q.pop();
+    nodes.push_back(cur);
+
+    for (auto out_node : out_edges.at(cur)) {
+      auto &in_nodes = in_edges.at(out_node);
+      in_nodes.erase(cur);
+      if (in_nodes.empty()) {
+        q.push(out_node);
+      }
+    }
+  }
+  PADDLE_ENFORCE_EQ(
+      nodes.size(),
+      num_nodes,
+      common::errors::InvalidArgument("Toposort error, this may be a bug."));
+  return nodes;
+}
+
+CUDAGraphID CUDAGraph::UniqueID() {
+  static std::atomic<CUDAGraphID> id;
+  return id.fetch_add(1);
+}
+
+int64_t CUDAGraph::UniqueMemoryPoolID() {
+  static std::atomic<int64_t> id(CUDAGraph::kDefaultPoolID + 1);
+  return id.fetch_add(1);
+}
+
+void CUDAGraph::Reset() {
+  if (is_reset_) return;
+  for (auto graph : graphs_) {
+    PADDLE_ENFORCE_GPU_SUCCESS(hipGraphDestroy(graph));
+  }
+  graphs_.clear();
+  for (auto exec_graph : exec_graphs_) {
+    PADDLE_ENFORCE_GPU_SUCCESS(hipGraphExecDestroy(exec_graph));
+  }
+  exec_graphs_.clear();
+  // callback should be called in reverse order because the latter added
+  // callback may rely on the former added callback.
+  for (auto iter = cudagraph_post_reset_callbacks_.rbegin();
+       iter != cudagraph_post_reset_callbacks_.rend();
+       ++iter) {
+    (*iter)(*this);
+  }
+  cudagraph_post_reset_callbacks_.clear();
+  is_reset_ = true;
+}
+
+void CUDAGraph::Replay() {
+  PADDLE_ENFORCE_EQ(is_reset_,
+                    false,
+                    common::errors::PermissionDenied(
+                        "Cannot replay the CUDA Graph after reset is called."));
+  size_t n = exec_graphs_.size();
+  for (size_t i = 0; i < n; ++i) {
+    if (!is_first_run_) {
+      for (auto &hook : cudagraph_pre_replay_callbacks_[i]) {
+        hook(exec_graphs_[i]);
+      }
+    }
+    PADDLE_ENFORCE_GPU_SUCCESS(hipGraphLaunch(exec_graphs_[i], stream_));
+  }
+  is_first_run_ = false;
+}
+
+void CUDAGraph::BeginSegmentCapture() {
+  ThrowErrorIfNotSupportCUDAGraph();
+  PADDLE_ENFORCE_EQ(IsCapturing(),
+                    true,
+                    common::errors::PermissionDenied(
+                        "BeginSegmentCapture should be called when CUDA "
+                        "Graph is capturing."));
+  if (IsThreadLocalCapturing()) {
+    PADDLE_ENFORCE_EQ(IsThisThreadCapturing(),
+                      true,
+                      common::errors::PermissionDenied(
+                          "When capturing CUDA Graph in the thread local mode, "
+                          "you cannot begin segmented capturing in the thread "
+                          "which is not the one that starts the capturing."));
+  }
+  PADDLE_ENFORCE_GPU_SUCCESS(hipStreamBeginCapture(
+      capturing_graph_->stream_, capturing_graph_->capture_mode_));
+  PADDLE_ENFORCE_EQ(IsValidCapturing(),
+                    true,
+                    common::errors::PermissionDenied(
+                        "CUDA Graph should not be invalidated."));
+  VLOG(10) << "Begin to capture CUDA Graph with ID " << capturing_graph_->id_
+           << ", segment id " << capturing_graph_->graphs_.size()
+           << ", memory pool id " << capturing_graph_->pool_id_;
+}
+
+void CUDAGraph::BeginCapture(phi::GPUPlace place,
+                             gpuStream_t stream,
+                             hipStreamCaptureMode mode,
+                             bool enable_replace) {
+  ThrowErrorIfNotSupportCUDAGraph();
+  PADDLE_ENFORCE_EQ(IsCapturing(),
+                    false,
+                    common::errors::PermissionDenied(
+                        "CUDA Graph can only captured one by one."));
+  PADDLE_ENFORCE_NOT_NULL(
+      stream,
+      common::errors::PermissionDenied(
+          "CUDA Graph cannot be captured in default CUDA stream 0."));
+  capturing_graph_.reset(new CUDAGraph(enable_replace));
+  capturing_graph_->place_ = place;
+  capturing_graph_->stream_ = stream;
+  capturing_graph_->capture_mode_ = mode;
+  if (mode == hipStreamCaptureModeThreadLocal) {
+    capturing_thread_id_ = std::this_thread::get_id();
+    VLOG(10) << "Capturing CUDA Graph in thread local mode, thread id: "
+             << capturing_thread_id_;
+  }
+  BeginSegmentCapture();
+}
+
+void CUDAGraph::EndSegmentCapture() {
+  ThrowErrorIfNotSupportCUDAGraph();
+  PADDLE_ENFORCE_EQ(
+      IsCapturing(),
+      true,
+      common::errors::PermissionDenied("No CUDA Graph is capturing."));
+  hipGraph_t graph;
+  PADDLE_ENFORCE_GPU_SUCCESS(
+      hipStreamEndCapture(capturing_graph_->stream_, &graph));
+  auto num_nodes = static_cast<size_t>(-1);
+  PADDLE_ENFORCE_GPU_SUCCESS(hipGraphGetNodes(graph, nullptr, &num_nodes));
+  if (num_nodes == 0) {
+    PADDLE_ENFORCE_GPU_SUCCESS(hipGraphDestroy(graph));
+    VLOG(10) << "Skip empty CUDA Graph with ID " << capturing_graph_->id_
+             << ", segment id " << capturing_graph_->graphs_.size()
+             << ", memory pool id " << capturing_graph_->pool_id_;
+    return;
+  }
+
+  for (auto &cudagraph_post_capture_callback :
+       capturing_graph_->cudagraph_post_capture_callbacks_) {
+    cudagraph_post_capture_callback();
+  }
+  capturing_graph_->cudagraph_post_capture_callbacks_.clear();
+
+  capturing_graph_->cudagraph_pre_replay_callbacks_.emplace_back(
+      CUDAGraphNodeLauncher::Instance().GetParameterSettersForExecGraph(graph));
+
+  // if forward graph is registered, this graph is a backward graph
+  // we check whether there is remain blocks that is unreleased by this
+  hipGraphExec_t exec_graph;
+  if (FLAGS_use_cuda_malloc_async_allocator &&
+      FLAGS_auto_free_cudagraph_allocations_on_launch) {
+    VLOG(1) << "hipGraphInstantiateFlagAutoFreeOnLaunch is enabled!";
+    PADDLE_ENFORCE_GPU_SUCCESS(hipGraphInstantiateWithFlags(
+        &exec_graph, graph, hipGraphInstantiateFlagAutoFreeOnLaunch));
+  } else {
+    PADDLE_ENFORCE_GPU_SUCCESS(
+        hipGraphInstantiate(&exec_graph, graph, nullptr, nullptr, 0));
+  }
+  VLOG(10) << "End to capture CUDA Graph with ID " << capturing_graph_->id_
+           << ", segment id " << capturing_graph_->graphs_.size()
+           << ", memory pool id " << capturing_graph_->pool_id_;
+  capturing_graph_->graphs_.emplace_back(graph);
+  capturing_graph_->exec_graphs_.emplace_back(exec_graph);
+}
+
+std::unique_ptr<CUDAGraph> CUDAGraph::EndCapture() {
+  EndSegmentCapture();
+  capturing_thread_id_ = paddle::none;
+  return std::move(capturing_graph_);
+}
+
+bool CUDAGraph::IsValidCapturing() {
+  if (!IsCapturing()) return false;
+  hipStreamCaptureStatus status;
+  CUDAGraphID id;
+  PADDLE_ENFORCE_GPU_SUCCESS(
+      hipStreamGetCaptureInfo(capturing_graph_->stream_, &status, &id));
+  return status == hipStreamCaptureStatusActive;
+}
+
+static std::string ConcatPath(const std::string &dirname,
+                              const std::string &filename) {
+#ifdef _WIN32
+  const std::array<char, 3> kFileSep = {"\\"};
+#else
+  const std::array<char, 2> kFileSep = {"/"};
+#endif
+  if (!dirname.empty() && dirname.back() == kFileSep[0]) {
+    return dirname + filename;
+  } else {
+    return dirname + kFileSep.data() + filename;
+  }
+}
+
+void CUDAGraph::PrintToDotFiles(const std::string &dirname,
+                                unsigned int flags) {
+  ThrowErrorIfNotSupportCUDAGraph();
+  PADDLE_THROW(common::errors::Unimplemented(
+      "The print_to_dot_files() method is not supported on ROCm/HIP"));
+}
+
+void CUDAGraphNodeLauncher::KernelNodeLaunch(
+    parameterSetter_t parameterSetter, gpuKernelCallback_t cudaKernelCallback) {
+  if (UNLIKELY(phi::backends::gpu::CUDAGraph::IsThisThreadCapturing())) {
+    unsigned int id = GenerateIdentifier();
+    auto cudaFunc = cudaKernelCallback(id);
+
+    parameterSetters[cudaFunc][id] = parameterSetter;
+    VLOG(10) << "[KernelNodeLaunch] Launch kernel with cudaFunc = " << cudaFunc
+             << " id = " << id;
+  } else {
+    cudaKernelCallback(0);
+  }
+}
+
+std::vector<cudaGraphExecuterSetter_t>
+CUDAGraphNodeLauncher::GetParameterSettersForExecGraph(hipGraph_t graph) {
+  size_t num_nodes;
+  PADDLE_ENFORCE_GPU_SUCCESS(hipGraphGetNodes(graph, nullptr, &num_nodes));
+  std::vector<hipGraphNode_t> nodes(num_nodes);
+  PADDLE_ENFORCE_GPU_SUCCESS(hipGraphGetNodes(graph, nodes.data(), &num_nodes));
+
+  std::vector<std::function<void(hipGraphExec_t)>> hooks;
+  for (auto node : nodes) {
+    hipGraphNode_t gpuNode = node;
+    hipGraphNodeType pType;
+    PADDLE_ENFORCE_GPU_SUCCESS(hipGraphNodeGetType(gpuNode, &pType));
+    if (pType == hipGraphNodeTypeKernel) {
+      hipKernelNodeParams gpuParams;
+      PADDLE_ENFORCE_GPU_SUCCESS(
+          gpuGraphKernelNodeGetParams(gpuNode, &gpuParams));
+      gpuKernelParams kernel_params(gpuParams.kernelParams);
+      auto kernel =
+          parameterSetters.find(static_cast<gpuFunction_t>(gpuParams.func));
+      VLOG(10) << "[GetParameterSettersForExecGraph] gpuParams.func = "
+               << gpuParams.func;
+      // There exists a parameter setter
+      if (kernel != parameterSetters.end()) {
+        auto launchSequence = kernel->second;
+        unsigned int id = kernel_params.As<int>(0);
+
+        VLOG(10) << "[GetParameterSettersForExecGraph] Find launch kernel id = "
+                 << id;
+        auto parameterSetter = launchSequence.find(id);
+        if (parameterSetter != launchSequence.end()) {
+          auto setter = parameterSetter->second;
+          hooks.emplace_back(
+              [setter, gpuNode, gpuParams](hipGraphExec_t exec_graph) {
+                gpuKernelParams kernel_params(gpuParams.kernelParams);
+                setter(kernel_params);
+                PADDLE_ENFORCE_GPU_SUCCESS(hipGraphExecKernelNodeSetParams(
+                    exec_graph, gpuNode, &gpuParams));
+              });
+        } else {
+          PADDLE_THROW(common::errors::InvalidArgument(
+              "Error: does not find launch id"));
+        }
+      }
+    }
+  }
+
+  return hooks;
+}
+
+}  // namespace gpu
+}  // namespace backends
+}  // namespace phi
+
+#endif

paddle/phi/backends/gpu/rocm/hip_graph.h

@@ -0,0 +1,374 @@
+// Copyright (c) 2022 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.
+
+#pragma once
+
+#include <array>
+#include <atomic>
+#include <functional>
+#include <future>
+#include <memory>
+#include <mutex>
+#include <queue>
+#include <set>
+#include <thread>
+#include <unordered_map>
+#include <unordered_set>
+#include <vector>
+
+#include "paddle/common/errors.h"
+#include "paddle/common/macros.h"
+#include "paddle/phi/backends/context_pool.h"
+#include "paddle/phi/backends/device_code.h"
+#include "paddle/phi/backends/gpu/gpu_context.h"
+#include "paddle/phi/common/memory_utils.h"
+#include "paddle/phi/common/place.h"
+#include "paddle/phi/core/enforce.h"
+#include "paddle/utils/optional.h"
+
+#ifdef PADDLE_WITH_HIP
+
+namespace phi {
+namespace backends {
+namespace gpu {
+
+class CUDAGraphContextManager {
+ public:
+  using DeviceContextMap =
+      std::map<Place, std::shared_future<std::unique_ptr<DeviceContext>>>;
+
+  static CUDAGraphContextManager &Instance() {
+    static CUDAGraphContextManager *cuda_graph_ctx_manager =
+        new CUDAGraphContextManager;
+    return *cuda_graph_ctx_manager;
+  }
+
+  DeviceContext *Get(int64_t pool_id, const Place &place, int stream_priority) {
+    std::lock_guard<std::mutex> lk(ctx_mtx_);
+    VLOG(6) << "Get cuda graph device context for " << place;
+
+    DeviceContextMap &ctxs = cuda_graph_ctx_pool_[pool_id];
+    if (ctxs.find(place) == ctxs.end()) {
+      phi::memory_utils::EmplaceDeviceContexts(
+          &ctxs,
+          {place},
+          /*disable_setting_default_stream_for_allocator=*/true,
+          stream_priority);
+    }
+    return ctxs[place].get().get();
+  }
+
+  void RecordCapturingDeviceContext(DeviceContext *dev_ctx) {
+    capturing_ctxs_.insert(dev_ctx);
+  }
+
+  std::set<DeviceContext *> GetAllCapturingDeviceContexts() const {
+    return capturing_ctxs_;
+  }
+
+  void ClearDeviceContextsRecords() { capturing_ctxs_.clear(); }
+
+ private:
+  CUDAGraphContextManager() {}
+  DISABLE_COPY_AND_ASSIGN(CUDAGraphContextManager);
+
+  std::mutex ctx_mtx_;
+  std::unordered_map<int64_t, DeviceContextMap> cuda_graph_ctx_pool_;
+  std::set<DeviceContext *> capturing_ctxs_;
+};
+
+class gpuKernelParams {
+ public:
+  explicit gpuKernelParams(void **params) : kernelParams(params) {}
+
+  template <typename T>
+  T &As(size_t idx) const {
+    return *reinterpret_cast<T *>(kernelParams[idx]);
+  }
+
+  void **getParams() const { return kernelParams; }
+
+ private:
+  void **kernelParams;
+};
+
+using cudaGraphExecuterSetter_t = std::function<void(hipGraphExec_t)>;
+
+//  ** class CUDAGraphNodeLauncher
+//
+//  This class offers a interface for launching CUDA kernels in CUDA Graph, we
+//  utilize the `cudaGraphExecKernelNodeSetParams` function for parameter setup.
+//  Launching kernels via this class ensures proper management.
+//
+//  NOTE: It's essential that the first parameter for any kernel launched
+//  through this class is an `unsigned int` identifier. This identifier plays a
+//  crucial role in linking the CUDA kernel to its corresponding CUDA graph
+//  node. We tag each kernel launch with a unique identifier to maintain
+//  structured linkage with its CUDA graph node.
+//
+//  NOTE: This class use a singleton design pattern ensures there's only a
+//  single global instance accessible via the `Instance()` method.
+class CUDAGraphNodeLauncher {
+ public:
+  //  [Parameter Setter Callback]
+  //  Sets the kernel's parameters BEFORE activating the CUDA graph. It enables
+  //  dynamic determination and setup of kernel arguments.
+  //
+  //  parameterSetter_t parameterSetter = [saved_state](gpuKernelParams
+  //  &param){
+  //      // Code to compute and the parameter values from the saved_state
+  //      // ...
+  //      param.As<type>(idx) = calculated_value;
+  //  };
+  using parameterSetter_t = std::function<void(gpuKernelParams &)>;
+
+  //  [CUDA Kernel Callback]
+  //  Acts as the launcher for the kernel. It accepts an `unsigned int`
+  //  identifier and uses it for the kernel launch.
+  //  The `cudaGetFuncBySymbol` method can be used to fetch the `cudaFunction_t`
+  //  reference of the kernel from the kernel pointer.
+  //  gpuKernelCallback_t cudaKernelCallback = [=](unsigned int id) {
+  //      // cudaFunction_t is REQUIRED to get here
+  //      cudaFunction_t cudaFunc;
+  //      PADDLE_ENFORCE_GPU_SUCCESS(cudaGetFuncBySymbol(&cudaFunc, &kernel));
+  //
+  //      kernel<<<>>>(id, ...);  // Launching the kernel with id
+  //      return cudaFunc;
+  //  };
+  using gpuKernelCallback_t = std::function<hipFunction_t(unsigned int)>;
+
+  //  [Kernel Launch]
+  //  With the callbacks defined and the CUDA function obtained, the kernel can
+  //  be launched using the `KernelNodeLaunch` method.
+  void KernelNodeLaunch(parameterSetter_t parameterSetter,
+                        gpuKernelCallback_t cudaKernelCallback);
+
+  std::vector<cudaGraphExecuterSetter_t> GetParameterSettersForExecGraph(
+      hipGraph_t graph);
+
+  parameterSetter_t GetParameterSetter(const gpuKernelParams &params);
+
+  static CUDAGraphNodeLauncher &Instance() {
+    static CUDAGraphNodeLauncher *launcher = new CUDAGraphNodeLauncher;
+    return *launcher;
+  }
+
+ private:
+  CUDAGraphNodeLauncher() : id(0) {}
+  DISABLE_COPY_AND_ASSIGN(CUDAGraphNodeLauncher);
+
+  unsigned int GenerateIdentifier() { return id++; }
+
+  unsigned int id;
+  std::unordered_map<hipFunction_t, std::map<unsigned int, parameterSetter_t>>
+      parameterSetters;
+};
+
+static void ThrowErrorIfNotSupportCUDAGraph() {}
+
+using CUDAGraphID = unsigned long long;  // NOLINT
+
+// NOTE: Currently, we do not support to capture CUDA graph in parallel
+// NOTE: Do not use this class directly because it should be used with
+//       the memory pool.
+class CUDAGraph {
+  DISABLE_COPY_AND_ASSIGN(CUDAGraph);
+
+  // Since the constructor would throw error is CUDA_VERSION < 10010.
+  // The non-static method of CUDAGraph need not check CUDA_VERSION
+  // again.
+  explicit CUDAGraph(bool enable_replace = false)
+      : enable_replace_(enable_replace) {
+    ThrowErrorIfNotSupportCUDAGraph();
+    id_ = UniqueID();
+  }
+
+ public:
+  static constexpr int64_t kDefaultPoolID = 0;
+  static constexpr int64_t kInvalidPoolID = -1;
+
+  ~CUDAGraph() { Reset(); }
+
+  CUDAGraphID ID() const { return id_; }
+
+  static int64_t SetMemoryPoolID(int64_t pool_id) {
+    auto &pool_id_ = capturing_graph_->pool_id_;
+    PADDLE_ENFORCE_EQ(pool_id_,
+                      kInvalidPoolID,
+                      common::errors::InvalidArgument(
+                          "Cannot reset memory pool id twice, the "
+                          "former memory pool id is %d.",
+                          pool_id_));
+    if (pool_id <= kInvalidPoolID) {
+      pool_id_ = UniqueMemoryPoolID();
+    } else {
+      PADDLE_ENFORCE_GE(pool_id,
+                        kDefaultPoolID,
+                        common::errors::InvalidArgument(
+                            "Invalid memory pool id %d.", pool_id));
+      pool_id_ = pool_id;
+    }
+    return pool_id_;
+  }
+
+  int64_t PoolID() const { return pool_id_; }
+
+  static int64_t CapturingPoolID() { return capturing_graph_->pool_id_; }
+
+  void Replay();
+
+  void Reset();
+
+  void AddPostResetCallback(
+      std::function<void(paddle::optional<const CUDAGraph &>)> callback) {
+    std::lock_guard<std::mutex> guard(mtx_);
+    cudagraph_post_reset_callbacks_.push_back(std::move(callback));
+  }
+
+  void AddPostCaptureCallback(std::function<void()> callback) {
+    std::lock_guard<std::mutex> guard(mtx_);
+    cudagraph_post_capture_callbacks_.push_back(std::move(callback));
+  }
+
+  void PrintToDotFiles(const std::string &dirname, unsigned int flags);
+
+  static void BeginCapture(phi::GPUPlace place,
+                           gpuStream_t stream,
+                           gpuStreamCaptureMode mode,
+                           bool enable_replace = false);
+  static std::unique_ptr<CUDAGraph> EndCapture();
+
+  static void BeginSegmentCapture();
+  static void EndSegmentCapture();
+
+  static void AddPostResetCallbackDuringCapturing(
+      std::function<void(paddle::optional<const CUDAGraph &>)> callback) {
+    capturing_graph_->AddPostResetCallback(std::move(callback));
+  }
+
+  static void AddPostCaptureCallbackDuringCapturing(
+      std::function<void()> callback) {
+    capturing_graph_->AddPostCaptureCallback(std::move(callback));
+  }
+
+  // No need to add CUDA_VERSION macro because capturing_graph_ would
+  // always be nullptr (constructor throws error)
+  static bool IsCapturing() { return capturing_graph_ != nullptr; }
+
+  static CUDAGraphID CapturingID() { return capturing_graph_->id_; }
+
+  static phi::GPUPlace CapturingPlace() { return capturing_graph_->place_; }
+
+  // This API can be used to debug which GPU operation is not
+  // supported during capturing CUDA Graph.
+  static bool IsValidCapturing();
+
+  static bool IsThreadLocalCapturing() {
+    return IsCapturing() &&
+           capturing_graph_->capture_mode_ == hipStreamCaptureModeThreadLocal;
+  }
+
+  static bool IsThisThreadCapturing() {
+    if (UNLIKELY(IsCapturing())) {
+      return IsThreadLocalCapturing()
+                 ? capturing_thread_id_.get() == std::this_thread::get_id()
+                 : true;
+    } else {
+      return false;
+    }
+  }
+
+  using SetSeedFunc = std::function<bool(gpuKernelParams *, bool)>;
+  static void RecordRandomKernelInfo(SetSeedFunc set_seed_func) {
+    std::lock_guard<std::mutex> guard(capturing_graph_->func_mtx_);
+    capturing_graph_->set_seed_funcs_.emplace_back(std::move(set_seed_func));
+  }
+
+  static int64_t UniqueMemoryPoolID();
+
+ private:
+  static CUDAGraphID UniqueID();
+
+ private:
+  std::vector<hipGraph_t> graphs_;
+  std::vector<hipGraphExec_t> exec_graphs_;
+  gpuStreamCaptureMode capture_mode_;
+  gpuStream_t stream_{nullptr};
+  phi::GPUPlace place_;
+  CUDAGraphID id_;
+  int64_t pool_id_{kInvalidPoolID};
+  bool is_reset_{false};
+  bool enable_replace_{false};
+  std::mutex mtx_;
+
+  std::vector<SetSeedFunc> set_seed_funcs_;
+
+  // Holds callbacks that are triggered after the CUDA graph is reset. These
+  // callbacks are used for operations that need to be performed following the
+  // reset of a CUDA graph.
+  std::vector<std::function<void(paddle::optional<const CUDAGraph &>)>>
+      cudagraph_post_reset_callbacks_;
+
+  // Contains callbacks that are invoked after the CUDA graph has been captured.
+  // These callbacks are crucial for managing memory allocations related to the
+  // CUDA graph. They ensure that memory blocks not associated with a graph (as
+  // detailed in cuda_malloc_async_allocator) are not erroneously released
+  // during the graph's lifecycle.
+  std::vector<std::function<void()>> cudagraph_post_capture_callbacks_;
+
+  // Maintains a collection of 'pre-hooks' - functions that are executed before
+  // the CUDA graph is replayed. These pre-hooks are essential for setting up
+  // the necessary conditions or states required for the correct execution of
+  // the CUDA graph.
+  std::vector<std::vector<cudaGraphExecuterSetter_t>>
+      cudagraph_pre_replay_callbacks_;
+
+  std::mutex func_mtx_;
+
+  bool is_first_run_{true};
+
+  static paddle::optional<std::thread::id> capturing_thread_id_;
+  static std::unique_ptr<CUDAGraph> capturing_graph_;
+};
+
+class CUDAGraphCaptureModeGuard {
+  DISABLE_COPY_AND_ASSIGN(CUDAGraphCaptureModeGuard);
+
+ public:
+  explicit CUDAGraphCaptureModeGuard(
+      gpuStreamCaptureMode mode = hipStreamCaptureModeRelaxed) {
+    if (UNLIKELY(CUDAGraph::IsCapturing())) {
+      PADDLE_ENFORCE_GPU_SUCCESS(hipThreadExchangeStreamCaptureMode(&mode));
+      // After cudaThreadExchangeStreamCaptureMode is called,
+      // the variable "mode" would be set to the old capturing mode.
+      old_mode_ = mode;
+    }
+  }
+
+  ~CUDAGraphCaptureModeGuard() PADDLE_MAY_THROW {
+    if (UNLIKELY(CUDAGraph::IsCapturing())) {
+      PADDLE_ENFORCE_GPU_SUCCESS(
+          hipThreadExchangeStreamCaptureMode(&old_mode_));
+    }
+  }
+
+ private:
+  gpuStreamCaptureMode old_mode_;
+};
+
+}  // namespace gpu
+}  // namespace backends
+}  // namespace phi
+
+#endif

paddle/phi/backends/gpu/rocm/miopen_desc.h

@@ -0,0 +1,267 @@
+// Copyright (c) 2022 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.
+
+#pragma once
+
+#include <algorithm>
+#include <functional>
+#include <iostream>
+#include <iterator>
+#include <memory>
+#include <numeric>
+#include <string>
+#include <vector>
+
+#include "paddle/phi/backends/gpu/rocm/miopen_helper.h"
+#include "paddle/phi/core/utils/data_type.h"
+
+namespace phi {
+namespace backends {
+namespace gpu {
+
+inline std::vector<int> TransformDimOrder(const std::vector<int>& dims) {
+  std::vector<int> transformed_dims(dims.begin(), dims.end());
+  int H, W, D, C;
+  if (dims.size() == 4) {
+    H = dims[1];
+    W = dims[2];
+    C = dims[3];
+    transformed_dims[1] = C;
+    transformed_dims[2] = H;
+    transformed_dims[3] = W;
+  } else {
+    D = dims[1];
+    H = dims[2];
+    W = dims[3];
+    C = dims[4];
+    transformed_dims[1] = C;
+    transformed_dims[2] = D;
+    transformed_dims[3] = H;
+    transformed_dims[4] = W;
+  }
+  return transformed_dims;
+}
+
+inline miopenDataType_t ToCudnnDataType(const DataType& t) {
+  miopenDataType_t type = miopenFloat;
+  switch (t) {
+    case DataType::FLOAT16:
+      type = miopenHalf;
+      break;
+    case DataType::FLOAT32:
+      type = miopenFloat;
+      break;
+    case DataType::BFLOAT16:
+      type = miopenBFloat16;
+      break;
+    default:
+      break;
+  }
+  return type;
+}
+
+class ActivationDescriptor {
+ public:
+  using T = miopenActivationDescriptor;
+  struct Deleter {
+    void operator()(T* t) {
+      if (t != nullptr) {
+        PADDLE_ENFORCE_GPU_SUCCESS(
+            phi::dynload::miopenDestroyActivationDescriptor(t));
+        t = nullptr;
+      }
+    }
+  };
+  ActivationDescriptor() {
+    T* raw_ptr;
+    PADDLE_ENFORCE_GPU_SUCCESS(
+        phi::dynload::miopenCreateActivationDescriptor(&raw_ptr));
+    desc_.reset(raw_ptr);
+  }
+  template <typename T>
+  void set(miopenActivationMode_t mode, const T& coef) {
+    PADDLE_ENFORCE_GPU_SUCCESS(phi::dynload::miopenSetActivationDescriptor(
+        desc_.get(), mode, static_cast<double>(coef), 0.0, 0.0));
+  }
+
+  T* desc() { return desc_.get(); }
+  T* desc() const { return desc_.get(); }
+
+ private:
+  std::unique_ptr<T, Deleter> desc_;
+};
+
+class TensorDescriptor {
+ public:
+  using T = miopenTensorDescriptor;
+  struct Deleter {
+    void operator()(T* t) {
+      if (t != nullptr) {
+        PADDLE_ENFORCE_GPU_SUCCESS(
+            phi::dynload::miopenDestroyTensorDescriptor(t));
+        t = nullptr;
+      }
+    }
+  };
+  TensorDescriptor() {
+    T* raw_ptr;
+    PADDLE_ENFORCE_GPU_SUCCESS(
+        phi::dynload::miopenCreateTensorDescriptor(&raw_ptr));
+    desc_.reset(raw_ptr);
+  }
+  T* desc() { return desc_.get(); }
+  T* desc() const { return desc_.get(); }
+
+  void set(const DenseTensor& tensor, const int groups = 1) {
+    auto dims = common::vectorize<int>(tensor.dims());
+    std::vector<int> strides(dims.size());
+    strides[dims.size() - 1] = 1;
+    for (int i = dims.size() - 2; i >= 0; i--) {
+      strides[i] = dims[i + 1] * strides[i + 1];
+    }
+    std::vector<int> dims_with_group(dims.begin(), dims.end());
+    if (groups > 1) {
+      dims_with_group[1] = dims_with_group[1] / groups;
+    }
+    PADDLE_ENFORCE_GPU_SUCCESS(phi::dynload::miopenSetTensorDescriptor(
+        (miopenTensorDescriptor_t)(desc_.get()),
+        ToCudnnDataType(tensor.dtype()),
+        static_cast<int>(dims_with_group.size()),
+        const_cast<int*>(dims_with_group.data()),
+        const_cast<int*>(strides.data())));
+  }
+
+  void set(const DenseTensor& tensor, const miopenTensorFormat_t format) {
+    const int groups = 1;
+    PADDLE_ENFORCE_EQ(format,
+                      MIOPEN_TENSOR_NCHW,
+                      common::errors::InvalidArgument(
+                          "format should ONLY be NCHW in MIOPEN."));
+    auto dims = common::vectorize<int>(tensor.dims());
+    std::vector<int> strides(dims.size());
+    strides[dims.size() - 1] = 1;
+    for (int i = dims.size() - 2; i >= 0; i--) {
+      strides[i] = dims[i + 1] * strides[i + 1];
+    }
+    std::vector<int> dims_with_group(dims.begin(), dims.end());
+    if (groups > 1) {
+      dims_with_group[1] = dims_with_group[1] / groups;
+    }
+    PADDLE_ENFORCE_GPU_SUCCESS(phi::dynload::miopenSetTensorDescriptor(
+        (miopenTensorDescriptor_t)(desc_.get()),
+        ToCudnnDataType(tensor.dtype()),
+        static_cast<int>(dims_with_group.size()),
+        const_cast<int*>(dims_with_group.data()),
+        const_cast<int*>(strides.data())));
+  }
+
+ private:
+  std::unique_ptr<T, Deleter> desc_;
+};
+
+class FilterDescriptor {
+ public:
+  using T = miopenTensorDescriptor;
+  struct Deleter {
+    void operator()(T* t) {
+      if (t != nullptr) {
+        PADDLE_ENFORCE_GPU_SUCCESS(
+            phi::dynload::miopenDestroyTensorDescriptor(t));
+        t = nullptr;
+      }
+    }
+  };
+  FilterDescriptor() {
+    T* raw_ptr;
+    PADDLE_ENFORCE_GPU_SUCCESS(
+        phi::dynload::miopenCreateTensorDescriptor(&raw_ptr));
+    desc_.reset(raw_ptr);
+  }
+  T* desc() { return desc_.get(); }
+  T* desc() const { return desc_.get(); }
+
+  void set(const DenseTensor& tensor,
+           const miopenTensorFormat_t format,
+           const int groups = 1) {
+    PADDLE_ENFORCE_EQ(format,
+                      MIOPEN_TENSOR_NCHW,
+                      common::errors::InvalidArgument(
+                          "format should ONLY be NCHW in MIOPEN."));
+    auto dims = common::vectorize<int>(tensor.dims());
+    std::vector<int> strides(dims.size());
+    strides[dims.size() - 1] = 1;
+    for (int i = dims.size() - 2; i >= 0; i--) {
+      strides[i] = dims[i + 1] * strides[i + 1];
+    }
+    std::vector<int> dims_with_group(dims.begin(), dims.end());
+    if (groups > 1) {
+      dims_with_group[1] = dims_with_group[1] / groups;
+    }
+    PADDLE_ENFORCE_GPU_SUCCESS(phi::dynload::miopenSetTensorDescriptor(
+        (miopenTensorDescriptor_t)(desc_.get()),
+        ToCudnnDataType(tensor.dtype()),
+        static_cast<int>(dims_with_group.size()),
+        const_cast<int*>(dims_with_group.data()),
+        const_cast<int*>(strides.data())));
+  }
+
+ private:
+  std::unique_ptr<T, Deleter> desc_;
+};
+
+class ConvolutionDescriptor {
+ public:
+  using T = miopenConvolutionDescriptor;
+  struct Deleter {
+    void operator()(T* t) {
+      if (t != nullptr) {
+        PADDLE_ENFORCE_GPU_SUCCESS(
+            phi::dynload::miopenDestroyConvolutionDescriptor(t));
+        t = nullptr;
+      }
+    }
+  };
+  ConvolutionDescriptor() {
+    T* raw_ptr;
+    PADDLE_ENFORCE_GPU_SUCCESS(
+        phi::dynload::miopenCreateConvolutionDescriptor(&raw_ptr));
+    desc_.reset(raw_ptr);
+  }
+  T* desc() { return desc_.get(); }
+  T* desc() const { return desc_.get(); }
+
+  void set(miopenDataType_t dtype,
+           const std::vector<int>& pads,
+           const std::vector<int>& strides,
+           const std::vector<int>& dilations,
+           bool allow_tf32,
+           const int groups = 1) {
+    PADDLE_ENFORCE_GPU_SUCCESS(phi::dynload::miopenInitConvolutionNdDescriptor(
+        (miopenConvolutionDescriptor_t)desc_.get(),
+        static_cast<int>(pads.size()),
+        const_cast<int*>(pads.data()),
+        const_cast<int*>(strides.data()),
+        const_cast<int*>(dilations.data()),
+        miopenConvolution));
+    PADDLE_ENFORCE_GPU_SUCCESS(phi::dynload::miopenSetConvolutionGroupCount(
+        (miopenConvolutionDescriptor_t)desc_.get(), groups));
+  }
+
+ private:
+  std::unique_ptr<T, Deleter> desc_;
+};
+
+}  // namespace gpu
+}  // namespace backends
+}  // namespace phi

paddle/phi/backends/gpu/rocm/miopen_helper.h

@@ -0,0 +1,590 @@
+/* Copyright (c) 2022 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. */
+
+#pragma once
+
+#include <string>
+#include <vector>
+
+#include "paddle/common/flags.h"
+
+#include "paddle/common/errors.h"
+#include "paddle/common/macros.h"
+#include "paddle/phi/backends/dynload/miopen.h"
+#include "paddle/phi/common/bfloat16.h"
+#include "paddle/phi/common/float16.h"
+#include "paddle/phi/common/place.h"
+#include "paddle/phi/core/dense_tensor.h"
+#include "paddle/phi/core/enforce.h"
+
+// MIOPEN do not have epsilon definition
+#define CUDNN_BN_MIN_EPSILON 1e-05
+
+COMMON_DECLARE_bool(cudnn_deterministic);
+
+namespace phi {
+namespace backends {
+namespace gpu {
+
+inline const char* miopenGetErrorString(miopenStatus_t status) {
+  switch (status) {
+    case miopenStatusSuccess:
+      return "miopenStatusSuccess";
+    case miopenStatusNotInitialized:
+      return "miopenStatusNotInitialized";
+    case miopenStatusAllocFailed:
+      return "miopenStatusAllocFailed";
+    case miopenStatusBadParm:        // typos: disable-line
+      return "miopenStatusBadParm";  // typos: disable-line
+    case miopenStatusInternalError:
+      return "miopenStatusInternalError";
+    case miopenStatusInvalidValue:
+      return "miopenStatusInvalidValue";
+    case miopenStatusUnknownError:
+      return "miopenStatusUnknownError";
+    case miopenStatusNotImplemented:
+      return "miopenStatusNotImplemented";
+    default:
+      return "Unknown miopen error number";
+  }
+}
+
+// no use, but will have compiling error if not defined
+#define CUDNN_VERSION_COMPUTE(major, minor, patch)     \
+  ((major) <= 8 ? (major)*1000 + (minor)*100 + (patch) \
+                : (major)*10000 + (minor)*100 + (patch))
+
+#define CUDNN_VERSION_MIN(major, minor, patch) \
+  (CUDNN_VERSION >= CUDNN_VERSION_COMPUTE(major, minor, patch))
+
+enum class PoolingMode {
+  kMaximum,
+  kMaximumDeterministic,
+  kAverageExclusive,
+  kAverageInclusive,
+};
+
+enum class ActivationMode {
+  kNone,  // activation identity
+  kSigmoid,
+  kRelu,
+  kRelu6,
+  kReluX,
+  kTanh,
+  kBandPass,
+};
+
+inline miopenPoolingMode_t GetPoolingMode(const PoolingMode& mode) {
+  switch (mode) {
+    case PoolingMode::kMaximumDeterministic:
+      return miopenPoolingMax;
+    case PoolingMode::kAverageExclusive:
+      return miopenPoolingAverage;
+    case PoolingMode::kAverageInclusive:
+      return miopenPoolingAverageInclusive;
+    case PoolingMode::kMaximum:
+      return miopenPoolingMax;
+    default:
+      PADDLE_THROW(
+          common::errors::Unimplemented("Unexpected MIOPEN pooling mode."));
+  }
+}
+
+inline ActivationMode StringToActivationMode(const std::string& str) {
+  if (str == "identity") {
+    return ActivationMode::kNone;
+  } else if (str == "sigmoid") {
+    return ActivationMode::kSigmoid;
+  } else if (str == "relu") {
+    return ActivationMode::kRelu;
+  } else if (str == "relu6") {
+    return ActivationMode::kRelu6;
+  } else if (str == "relux") {
+    return ActivationMode::kReluX;
+  } else if (str == "tanh") {
+    return ActivationMode::kTanh;
+  } else if (str == "bandpass") {
+    return ActivationMode::kBandPass;
+  } else {
+    PADDLE_THROW(common::errors::Unimplemented(
+        "Unknown MIOPEN activation string: %s.", str));
+  }
+}
+
+template <typename T>
+class CudnnDataType;
+
+template <>
+class CudnnDataType<phi::dtype::float16> {
+ public:
+  static const miopenDataType_t type = miopenHalf;
+  // The scaling param type is float for HALF and FLOAT tensors
+  using ScalingParamType = const float;
+  using BatchNormParamType = float;
+  static ScalingParamType* kOne() {
+    static ScalingParamType v = 1.0;
+    return &v;
+  }
+  static ScalingParamType* kZero() {
+    static ScalingParamType v = 0.0;
+    return &v;
+  }
+};
+
+template <>
+class CudnnDataType<phi::dtype::bfloat16> {
+ public:
+  static const miopenDataType_t type = miopenBFloat16;
+  // The scaling param type is float for HALF and FLOAT tensors
+  using ScalingParamType = const float;
+  using BatchNormParamType = float;
+  static ScalingParamType* kOne() {
+    static ScalingParamType v = 1.0;
+    return &v;
+  }
+  static ScalingParamType* kZero() {
+    static ScalingParamType v = 0.0;
+    return &v;
+  }
+};
+
+template <>
+class CudnnDataType<float> {
+ public:
+  static const miopenDataType_t type = miopenFloat;
+  using ScalingParamType = const float;
+  using BatchNormParamType = float;
+  static ScalingParamType* kOne() {
+    static ScalingParamType v = 1.0;
+    return &v;
+  }
+  static ScalingParamType* kZero() {
+    static ScalingParamType v = 0.0;
+    return &v;
+  }
+};
+
+inline miopenTensorFormat_t GetCudnnTensorFormat(const DataLayout& order) {
+  switch (order) {
+    case DataLayout::NHWC:
+      return MIOPEN_TENSOR_NHWC;
+    case DataLayout::NCHW:
+      return MIOPEN_TENSOR_NCHW;
+    case DataLayout::NCDHW:
+      return MIOPEN_TENSOR_NCHW;
+    case DataLayout::NDHWC:
+      return MIOPEN_TENSOR_NHWC;
+    default:
+      PADDLE_THROW(common::errors::Unimplemented(
+          "MIOPEN has no equivalent dataLayout for input order."));
+  }
+  return MIOPEN_TENSOR_NCHW;
+}
+class ScopedTensorDescriptor {
+ public:
+  ScopedTensorDescriptor() {
+    PADDLE_ENFORCE_GPU_SUCCESS(
+        phi::dynload::miopenCreateTensorDescriptor(&desc_));
+  }
+  ~ScopedTensorDescriptor() PADDLE_MAY_THROW {
+    PADDLE_ENFORCE_GPU_SUCCESS(
+        phi::dynload::miopenDestroyTensorDescriptor(desc_));
+  }
+
+  inline miopenTensorDescriptor_t descriptor(const miopenTensorFormat_t format,
+                                             const miopenDataType_t type,
+                                             const std::vector<int>& dims,
+                                             const int groups = 1) {
+    // the format is not used now, will add later
+    std::vector<int> strides(dims.size());
+    strides[dims.size() - 1] = 1;
+    for (int i = dims.size() - 2; i >= 0; i--) {
+      strides[i] = dims[i + 1] * strides[i + 1];
+    }
+    // Update tensor descriptor dims setting if groups > 1
+    // NOTE: Here, Assume using NCHW or NCDHW order
+    std::vector<int> dims_with_group(dims.begin(), dims.end());
+    if (groups > 1) {
+      dims_with_group[1] = dims_with_group[1] / groups;
+    }
+
+    // MIOPEN ONLY support data layout of NCHW
+    PADDLE_ENFORCE_EQ(format,
+                      MIOPEN_TENSOR_NCHW,
+                      common::errors::InvalidArgument(
+                          "format should ONLY be NCHW in MIOPEN."));
+    if (dims.size() == 4) {
+      PADDLE_ENFORCE_GPU_SUCCESS(phi::dynload::miopenSetTensorDescriptor(
+          desc_,
+          type,
+          dims_with_group.size(),
+          const_cast<int*>(dims_with_group.data()),
+          const_cast<int*>(strides.data())));
+    } else if (dims.size() == 5) {
+      PADDLE_ENFORCE_GPU_SUCCESS(phi::dynload::miopenSetTensorDescriptor(
+          desc_,
+          type,
+          dims_with_group.size(),
+          const_cast<int*>(dims_with_group.data()),
+          const_cast<int*>(strides.data())));
+    }
+    return desc_;
+  }
+
+  template <typename T>
+  inline miopenTensorDescriptor_t descriptor(const DataLayout& order,
+                                             const std::vector<int>& dims,
+                                             const int groups = 1) {
+    return descriptor(
+        GetCudnnTensorFormat(order), CudnnDataType<T>::type, dims, groups);
+  }
+
+  inline miopenTensorDescriptor_t descriptor(const miopenDataType_t miopen_type,
+                                             const std::vector<int>& dim,
+                                             const std::vector<int>& stride) {
+    PADDLE_ENFORCE_GPU_SUCCESS(phi::dynload::miopenSetTensorDescriptor(
+        desc_,
+        miopen_type,
+        dim.size(),
+        const_cast<int*>(dim.data()),
+        const_cast<int*>(stride.data())));
+    return desc_;
+  }
+
+  template <typename T>
+  inline miopenTensorDescriptor_t descriptor(const std::vector<int>& dim,
+                                             const std::vector<int>& stride) {
+    return descriptor(CudnnDataType<T>::type, dim, stride);
+  }
+
+  inline miopenTensorDescriptor_t desc() { return desc_; }
+
+ private:
+  miopenTensorDescriptor_t desc_;
+  DISABLE_COPY_AND_ASSIGN(ScopedTensorDescriptor);
+};
+
+class ScopedDropoutDescriptor {
+ public:
+  ScopedDropoutDescriptor() {
+    PADDLE_ENFORCE_GPU_SUCCESS(
+        phi::dynload::miopenCreateDropoutDescriptor(&desc_));
+  }
+  ~ScopedDropoutDescriptor() PADDLE_MAY_THROW {
+    PADDLE_ENFORCE_GPU_SUCCESS(
+        phi::dynload::miopenDestroyDropoutDescriptor(desc_));
+  }
+
+  inline miopenDropoutDescriptor_t descriptor(const miopenHandle_t& handle,
+                                              const Place& place,
+                                              bool initialized,
+                                              float dropout_prob_,
+                                              DenseTensor* dropout_state_,
+                                              int seed,
+                                              size_t state_size) {
+    if (dropout_state_ == nullptr) {  // for no dropout or test
+      PADDLE_ENFORCE_GPU_SUCCESS(
+          phi::dynload::miopenSetDropoutDescriptor(desc_,
+                                                   handle,
+                                                   0 /* dropout */,
+                                                   nullptr,
+                                                   0 /* state_size */,
+                                                   0 /* seed */,
+                                                   false,
+                                                   false,
+                                                   MIOPEN_RNG_PSEUDO_XORWOW));
+      return desc_;
+    }
+    auto* dropout_state_data = dropout_state_->data<uint8_t>();
+    if (!initialized) {
+      PADDLE_ENFORCE_GPU_SUCCESS(
+          phi::dynload::miopenSetDropoutDescriptor(desc_,
+                                                   handle,
+                                                   dropout_prob_,
+                                                   dropout_state_data,
+                                                   state_size,
+                                                   seed,
+                                                   false,
+                                                   false,
+                                                   MIOPEN_RNG_PSEUDO_XORWOW));
+    } else {
+      auto dropout_state_dims = dropout_state_->dims();
+      state_size = dropout_state_dims[0];
+      PADDLE_ENFORCE_GPU_SUCCESS(phi::dynload::miopenRestoreDropoutDescriptor(
+          desc_,
+          handle,
+          dropout_prob_,
+          dropout_state_data,
+          state_size,
+          0,
+          false,
+          false,
+          MIOPEN_RNG_PSEUDO_XORWOW));
+    }
+    return desc_;
+  }
+  inline miopenDropoutDescriptor_t desc() { return desc_; }
+
+ private:
+  miopenDropoutDescriptor_t desc_;
+  DISABLE_COPY_AND_ASSIGN(ScopedDropoutDescriptor);
+};
+
+class ScopedRNNDescriptor {
+ public:
+  ScopedRNNDescriptor() {
+    PADDLE_ENFORCE_GPU_SUCCESS(phi::dynload::miopenCreateRNNDescriptor(&desc_));
+  }
+  ~ScopedRNNDescriptor() PADDLE_MAY_THROW {
+    PADDLE_ENFORCE_GPU_SUCCESS(phi::dynload::miopenDestroyRNNDescriptor(desc_));
+  }
+
+  inline miopenRNNDescriptor_t desc() { return desc_; }
+
+ private:
+  miopenRNNDescriptor_t desc_;
+  DISABLE_COPY_AND_ASSIGN(ScopedRNNDescriptor);
+};
+
+class ScopedFilterDescriptor {
+ public:
+  ScopedFilterDescriptor() {
+    PADDLE_ENFORCE_GPU_SUCCESS(
+        phi::dynload::miopenCreateTensorDescriptor(&desc_));
+  }
+  ~ScopedFilterDescriptor() PADDLE_MAY_THROW {
+    PADDLE_ENFORCE_GPU_SUCCESS(
+        phi::dynload::miopenDestroyTensorDescriptor(desc_));
+  }
+
+  inline miopenTensorDescriptor_t descriptor(const miopenTensorFormat_t format,
+                                             const miopenDataType_t type,
+                                             const std::vector<int>& kernel,
+                                             const int groups = 1) {
+    // filter layout: MCHW(MCDHW), where M is the number of
+    // output image channels, C is the number of input image channels,
+    // D is the depth of the filter, H is the height of the filter, and W is the
+    // width of the filter.
+    std::vector<int> kernel_with_group(kernel.begin(), kernel.end());
+    if (groups > 1) {
+      kernel_with_group[0] /= groups;
+      // NOTE: input filter(C) of the filter is already asserted to be C/groups.
+    }
+    std::vector<int> stride_dim(kernel_with_group.size());
+    stride_dim[kernel_with_group.size() - 1] = 1;
+    for (int k = kernel_with_group.size() - 2; k >= 0; k--) {
+      stride_dim[k] = stride_dim[k + 1] * kernel_with_group[k + 1];
+    }
+    PADDLE_ENFORCE_GPU_SUCCESS(phi::dynload::miopenSetTensorDescriptor(
+        desc_,
+        type,
+        kernel_with_group.size(),
+        const_cast<int*>(kernel_with_group.data()),
+        const_cast<int*>(stride_dim.data())));
+    return desc_;
+  }
+
+  template <typename T>
+  inline miopenTensorDescriptor_t descriptor(const DataLayout& order,
+                                             const std::vector<int>& kernel,
+                                             const int groups = 1) {
+    return descriptor(
+        GetCudnnTensorFormat(order), CudnnDataType<T>::type, kernel, groups);
+  }
+
+  inline miopenTensorDescriptor_t desc() { return desc_; }
+
+ private:
+  miopenTensorDescriptor_t desc_;
+  DISABLE_COPY_AND_ASSIGN(ScopedFilterDescriptor);
+};
+
+class ScopedConvolutionDescriptor {
+ public:
+  ScopedConvolutionDescriptor() {
+    PADDLE_ENFORCE_GPU_SUCCESS(
+        phi::dynload::miopenCreateConvolutionDescriptor(&desc_));
+  }
+  ~ScopedConvolutionDescriptor() PADDLE_MAY_THROW {
+    PADDLE_ENFORCE_GPU_SUCCESS(
+        phi::dynload::miopenDestroyConvolutionDescriptor(desc_));
+  }
+
+  inline miopenConvolutionDescriptor_t descriptor(
+      miopenDataType_t type,
+      const std::vector<int>& pads,
+      const std::vector<int>& strides,
+      const std::vector<int>& dilations) {
+    PADDLE_ENFORCE_EQ(pads.size(),
+                      strides.size(),
+                      common::errors::InvalidArgument(
+                          "The size of pads and strides should be equal. But "
+                          "received size of pads is %d, size of strides is %d.",
+                          pads.size(),
+                          strides.size()));
+    PADDLE_ENFORCE_EQ(
+        pads.size(),
+        dilations.size(),
+        common::errors::InvalidArgument(
+            "The size of pads and dilations should be equal. But received size "
+            "of pads is %d, size of dilations is %d.",
+            pads.size(),
+            dilations.size()));
+    PADDLE_ENFORCE_GPU_SUCCESS(phi::dynload::miopenInitConvolutionNdDescriptor(
+        desc_,
+        pads.size(),
+        const_cast<int*>(pads.data()),
+        const_cast<int*>(strides.data()),
+        const_cast<int*>(dilations.data()),
+        miopenConvolution));
+    return desc_;
+  }
+
+  template <typename T>
+  inline miopenConvolutionDescriptor_t descriptor(
+      const std::vector<int>& pads,
+      const std::vector<int>& strides,
+      const std::vector<int>& dilations) {
+    return descriptor(CudnnDataType<T>::type, pads, strides, dilations);
+  }
+
+ private:
+  miopenConvolutionDescriptor_t desc_;
+  DISABLE_COPY_AND_ASSIGN(ScopedConvolutionDescriptor);
+};
+
+class ScopedPoolingDescriptor {
+ public:
+  ScopedPoolingDescriptor() {
+    PADDLE_ENFORCE_GPU_SUCCESS(
+        phi::dynload::miopenCreatePoolingDescriptor(&desc_));
+  }
+  ~ScopedPoolingDescriptor() PADDLE_MAY_THROW {
+    PADDLE_ENFORCE_GPU_SUCCESS(
+        phi::dynload::miopenDestroyPoolingDescriptor(desc_));
+  }
+
+  inline miopenPoolingDescriptor_t descriptor(const PoolingMode& mode,
+                                              const std::vector<int>& kernel,
+                                              const std::vector<int>& pads,
+                                              const std::vector<int>& strides) {
+    PADDLE_ENFORCE_EQ(kernel.size(),
+                      pads.size(),
+                      common::errors::InvalidArgument(
+                          "The size of kernel and pads should be equal. But "
+                          "received size of kernel is %d, size of pads is %d.",
+                          kernel.size(),
+                          pads.size()));
+    PADDLE_ENFORCE_EQ(
+        kernel.size(),
+        strides.size(),
+        common::errors::InvalidArgument(
+            "The size of kernel and strides should be equal. But "
+            "received size of kernel is %d, size of strides is %d.",
+            kernel.size(),
+            strides.size()));
+    PADDLE_ENFORCE_GPU_SUCCESS(phi::dynload::miopenSetNdPoolingDescriptor(
+        desc_,
+        GetPoolingMode(mode),
+        kernel.size(),
+        const_cast<int*>(kernel.data()),
+        const_cast<int*>(pads.data()),
+        const_cast<int*>(strides.data())));
+    return desc_;
+  }
+
+ private:
+  miopenPoolingDescriptor_t desc_;
+  DISABLE_COPY_AND_ASSIGN(ScopedPoolingDescriptor);
+};
+
+class ScopedActivationDescriptor {
+ public:
+  ScopedActivationDescriptor() {
+    PADDLE_ENFORCE_GPU_SUCCESS(
+        phi::dynload::miopenCreateActivationDescriptor(&desc_));
+  }
+  ~ScopedActivationDescriptor() PADDLE_MAY_THROW {
+    PADDLE_ENFORCE_GPU_SUCCESS(
+        phi::dynload::miopenDestroyActivationDescriptor(desc_));
+  }
+
+  template <typename T>
+  inline miopenActivationDescriptor_t descriptor(
+      const std::string& act, double value_max = static_cast<double>(0.)) {
+    double relu_ceiling = 0.0;
+    ActivationMode activation_mode = StringToActivationMode(act);
+    miopenActivationMode_t mode;
+    switch (activation_mode) {
+      case ActivationMode::kNone:
+        mode = miopenActivationPASTHRU;
+        break;
+      case ActivationMode::kRelu6:
+        relu_ceiling = 6.0;
+        mode = miopenActivationCLIPPEDRELU;
+        break;
+      case ActivationMode::kReluX:
+        relu_ceiling = value_max;
+        mode = miopenActivationCLIPPEDRELU;
+        break;
+      case ActivationMode::kRelu:
+        mode = miopenActivationRELU;
+        break;
+      case ActivationMode::kSigmoid:
+        mode = miopenActivationLOGISTIC;
+        break;
+      case ActivationMode::kTanh:
+        mode = miopenActivationTANH;
+        break;
+      default:
+        PADDLE_THROW(common::errors::Unimplemented(
+            "Unrecognized MIOPEN activation mode: %d.",
+            static_cast<int>(activation_mode)));
+    }
+    PADDLE_ENFORCE_GPU_SUCCESS(phi::dynload::miopenSetActivationDescriptor(
+        desc_, mode, relu_ceiling, 0.0, 0.0));
+    return desc_;
+  }
+
+ private:
+  miopenActivationDescriptor_t desc_;
+  DISABLE_COPY_AND_ASSIGN(ScopedActivationDescriptor);
+};
+
+class ScopedCTCLossDescriptor {
+ public:
+  ScopedCTCLossDescriptor() {
+    PADDLE_ENFORCE_GPU_SUCCESS(
+        phi::dynload::miopenCreateCTCLossDescriptor(&desc_));
+  }
+  ~ScopedCTCLossDescriptor() PADDLE_MAY_THROW {
+    PADDLE_ENFORCE_GPU_SUCCESS(
+        phi::dynload::miopenDestroyCTCLossDescriptor(desc_));
+  }
+
+  template <typename T>
+  inline miopenCTCLossDescriptor_t descriptor() {
+    PADDLE_ENFORCE_GPU_SUCCESS(phi::dynload::miopenSetCTCLossDescriptor(
+        desc_, CudnnDataType<T>::type, 0, false));
+    return desc_;
+  }
+
+ private:
+  miopenCTCLossDescriptor_t desc_;
+  DISABLE_COPY_AND_ASSIGN(ScopedCTCLossDescriptor);
+};
+
+}  // namespace gpu
+}  // namespace backends
+}  // namespace phi

paddle/phi/backends/gpu/rocm/rocm_device_function.h

@@ -0,0 +1,163 @@
+/* Copyright (c) 2022 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. */
+
+#pragma once
+
+#include "paddle/phi/common/bfloat16.h"
+#include "paddle/phi/common/complex.h"
+#include "paddle/phi/common/float16.h"
+
+namespace phi {
+namespace backends {
+namespace gpu {
+
+#define CREATE_SHFL_MASK(mask, predicate) mask = __ballot((predicate))
+
+#define CUDA_LAUNCH_KERNEL_BASE(dim, ...)  \
+  case (dim): {                            \
+    constexpr auto kPowerOfTwoDim = (dim); \
+    __VA_ARGS__;                           \
+  } break
+
+#define CUDA_LAUNCH_KERNEL_HELPER(...)          \
+  CUDA_LAUNCH_KERNEL_BASE(1024, ##__VA_ARGS__); \
+  CUDA_LAUNCH_KERNEL_BASE(512, ##__VA_ARGS__);  \
+  CUDA_LAUNCH_KERNEL_BASE(256, ##__VA_ARGS__);  \
+  CUDA_LAUNCH_KERNEL_BASE(128, ##__VA_ARGS__);  \
+  CUDA_LAUNCH_KERNEL_BASE(64, ##__VA_ARGS__);   \
+  CUDA_LAUNCH_KERNEL_BASE(32, ##__VA_ARGS__);
+
+template <typename T>
+__forceinline__ __device__ T
+CudaShuffleDownSync(unsigned mask, T val, int delta, int width = warpSize) {
+  return __shfl_down(val, delta, width);
+}
+
+template <typename T>
+__forceinline__ __device__ T CudaShuffleXorSync(unsigned mask,
+                                                T val,
+                                                int width = warpSize) {
+  return __shfl_xor(val, width);
+}
+
+template <>
+__forceinline__ __device__ phi::dtype::float16 CudaShuffleDownSync(
+    unsigned mask, phi::dtype::float16 val, int delta, int width) {
+  return phi::dtype::float16(__shfl_down(
+      static_cast<float>(val), static_cast<unsigned>(delta), width));
+}
+
+template <>
+__forceinline__ __device__ phi::dtype::bfloat16 CudaShuffleDownSync(
+    unsigned mask, phi::dtype::bfloat16 val, int delta, int width) {
+  return phi::dtype::bfloat16(__shfl_down(
+      static_cast<float>(val), static_cast<unsigned>(delta), width));
+}
+
+template <>
+__forceinline__ __device__ phi::dtype::complex<float> CudaShuffleDownSync(
+    unsigned mask, phi::dtype::complex<float> val, int delta, int width) {
+  float real = __shfl_down(val.real, delta, width);
+  float imag = __shfl_down(val.imag, delta, width);
+  return phi::dtype::complex<float>(real, imag);
+}
+
+template <>
+__forceinline__ __device__ phi::dtype::complex<double> CudaShuffleDownSync(
+    unsigned mask, phi::dtype::complex<double> val, int delta, int width) {
+  double real = __shfl_down(val.real, delta, width);
+  double imag = __shfl_down(val.imag, delta, width);
+  return phi::dtype::complex<double>(real, imag);
+}
+
+template <>
+__forceinline__ __device__ phi::dtype::float16 CudaShuffleXorSync(
+    unsigned mask, phi::dtype::float16 val, int width) {
+  return phi::dtype::float16(__shfl_xor(static_cast<float>(val), width));
+}
+
+template <>
+__forceinline__ __device__ phi::dtype::bfloat16 CudaShuffleXorSync(
+    unsigned mask, phi::dtype::bfloat16 val, int width) {
+  return phi::dtype::bfloat16(__shfl_xor(static_cast<float>(val), width));
+}
+
+template <>
+__forceinline__ __device__ phi::dtype::complex<float> CudaShuffleXorSync(
+    unsigned mask, phi::dtype::complex<float> val, int width) {
+  float real = __shfl_xor(val.real, width);
+  float imag = __shfl_xor(val.imag, width);
+  return phi::dtype::complex<float>(real, imag);
+}
+
+template <>
+__forceinline__ __device__ phi::dtype::complex<double> CudaShuffleXorSync(
+    unsigned mask, phi::dtype::complex<double> val, int width) {
+  double real = __shfl_xor(val.real, width);
+  double imag = __shfl_xor(val.imag, width);
+  return phi::dtype::complex<double>(real, imag);
+}
+
+template <typename T>
+__forceinline__ __device__ T
+CudaShuffleSync(unsigned mask, T val, int src_line, int width = 32) {
+  return __shfl(val, src_line, width);
+}
+
+template <typename T>
+HOSTDEVICE T Infinity() {
+  return INFINITY;
+}
+
+template <typename T>
+__device__ T reduceSum(T val, int tid, int len) {
+  // NOTE(zcd): The warp size should be taken from the
+  // parameters of the GPU but not specified as 32 simply.
+  // To make the reduceSum more efficiently,
+  // I use Warp-Level Parallelism and assume the Warp size
+  // is 32 which may be different for different GPU,
+  // but most card's warp size is 32.
+#ifdef PADDLE_WITH_HIP
+  const int warpSize = 64;
+#else
+  const int warpSize = 32;
+#endif
+  __shared__ T shm[warpSize];
+  unsigned mask = 0u;
+  CREATE_SHFL_MASK(mask, tid < len);
+
+  for (int offset = warpSize / 2; offset > 0; offset /= 2)
+    val += phi::backends::gpu::CudaShuffleDownSync(mask, val, offset);
+
+  if (tid < warpSize) shm[tid] = 0;
+  __syncthreads();
+
+  if (tid % warpSize == 0) {
+    shm[tid / warpSize] = val;
+  }
+  __syncthreads();
+
+  CREATE_SHFL_MASK(mask, tid < warpSize);
+
+  if (tid < warpSize) {
+    val = shm[tid];
+    for (int offset = warpSize / 2; offset > 0; offset /= 2)
+      val += phi::backends::gpu::CudaShuffleDownSync(mask, val, offset);
+  }
+  return val;
+}
+
+}  // namespace gpu
+}  // namespace backends
+}  // namespace phi

paddle/phi/backends/gpu/rocm/rocm_helper.h

@@ -0,0 +1,138 @@
+// Copyright (c) 2019 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.
+
+#pragma once
+
+#ifdef PADDLE_WITH_HIP
+#include <hip/hip_runtime.h>
+#include <hip/library_types.h>
+#include <hipblaslt/hipblaslt.h>
+
+#include "paddle/phi/common/bfloat16.h"
+#include "paddle/phi/common/data_type.h"
+#include "paddle/phi/common/float16.h"
+#include "paddle/phi/core/enforce.h"
+
+#if HIP_VERSION >= 60100000
+using hipDataType_t = hipDataType;
+constexpr hipDataType_t HIP_DATATYPE_R_32F = hipDataType::HIP_R_32F;
+constexpr hipDataType_t HIP_DATATYPE_R_64F = hipDataType::HIP_R_64F;
+constexpr hipDataType_t HIP_DATATYPE_R_16F = hipDataType::HIP_R_16F;
+constexpr hipDataType_t HIP_DATATYPE_R_8I = hipDataType::HIP_R_8I;
+constexpr hipDataType_t HIP_DATATYPE_R_16BF = hipDataType::HIP_R_16BF;
+constexpr hipDataType_t HIP_DATATYPE_R_32I = hipDataType::HIP_R_32I;
+#endif
+
+namespace phi {
+namespace backends {
+namespace gpu {
+
+/*
+ * Summary: Grid stride looping macro in CUDA kernel
+ *
+ *  [ Why need this macro? ]
+ *
+ *    The original looping in CUDA kernel is:
+ *
+ *    `for (int i = blockIdx.x * blockDim.x + threadIdx.x; i < (n); \
+ *        i += blockDim.x * gridDim.x)`
+ *
+ *    This for condition is risky. The value of `blockIdx.x * blockDim.x`
+ *    may be large, such as over 1GB, the first iteration is no problem here,
+ *    but when `i += blockDim.x * gridDim.x` is executed, the value of i
+ *    will greater than INT_MAX and overflow becomes negative value, at
+ *    this time, the cycle condition `i < (n)` is still satisfied, so it
+ *    will cause illegal access to cuda memory.
+ *
+ *    Here is a real example in ERINE, it will trigger above error.
+ *    The related data are:
+ *      - blockIdx.x = 2172938
+ *      - blockDim.x = 512
+ *      - blockIdx.x * blockDim.x = 1112543864
+ *      - INT_MAX = 2147483647
+ *
+ *    So we polish the for condition as follow, the int64_t __index__ will
+ *    prevent overflow in the loop increment.
+ *
+ * Parameters:
+ *    - i: loop index
+ *    - num: total element numbers
+ *
+ * Examples:
+ *    template <typename T>
+ *    __global__ void Scale(T* logit_grad, const T* loss_grad, const int num,
+ *                      const int d, const int remain) {
+ *    CUDA_KERNEL_LOOP(index, num) {
+ *      int idx_n = index / d;
+ *      int idx_remain = index % remain;
+ *      logit_grad[index] *= loss_grad[idx_n * remain + idx_remain];
+ *      }
+ *    }
+ *
+ */
+
+#define CUDA_KERNEL_LOOP_TYPE(i, num, index_type)                           \
+  int64_t __index__ =                                                       \
+      static_cast<int64_t>(hipBlockIdx_x) * hipBlockDim_x + hipThreadIdx_x; \
+  int64_t __stride__ = static_cast<int64_t>(hipBlockDim_x) * hipGridDim_x;  \
+  for (index_type i = __index__; __index__ < (num);                         \
+       __index__ += __stride__, i = __index__)
+
+template <typename T>
+hipDataType ToHipDataType() {
+  if (std::is_same<T, float>::value) {
+    return HIP_R_32F;
+  } else if (std::is_same<T, double>::value) {
+    return HIP_R_64F;
+  } else if (std::is_same<T, phi::dtype::float16>::value) {
+    return HIP_R_16F;
+  } else if (std::is_same<T, phi::dtype::bfloat16>::value) {
+    return HIP_R_16BF;
+  } else if (std::is_same<T, int8_t>::value) {
+    return HIP_R_8I;
+  } else if (std::is_same<T, int32_t>::value) {
+    return HIP_R_32I;
+  } else {
+    PADDLE_THROW(common::errors::InvalidArgument(
+        "DataType %s is unsupported for ROCm.",
+        DataTypeToString(phi::CppTypeToDataType<T>::Type())));
+  }
+}
+
+template <typename T>
+hipDataType_t ToHipBlasLtDataType() {
+  if (std::is_same<T, float>::value) {
+    return HIP_DATATYPE_R_32F;
+  } else if (std::is_same<T, double>::value) {
+    return HIP_DATATYPE_R_64F;
+  } else if (std::is_same<T, phi::dtype::float16>::value) {
+    return HIP_DATATYPE_R_16F;
+  } else if (std::is_same<T, phi::dtype::bfloat16>::value) {
+    return HIP_DATATYPE_R_16BF;
+  } else if (std::is_same<T, int8_t>::value) {
+    return HIP_DATATYPE_R_8I;
+  } else if (std::is_same<T, int32_t>::value) {
+    return HIP_DATATYPE_R_32I;
+  } else {
+    PADDLE_THROW(common::errors::InvalidArgument(
+        "DataType %s is unsupported for ROCm.",
+        DataTypeToString(phi::CppTypeToDataType<T>::Type())));
+  }
+}
+
+}  // namespace gpu
+}  // namespace backends
+}  // namespace phi
+
+#endif

paddle/phi/backends/gpu/rocm/rocm_info.cc

@@ -0,0 +1,360 @@
+// Copyright (c) 2021 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.
+
+#include <array>
+
+#include "paddle/phi/backends/gpu/gpu_info.h"
+
+#include "paddle/phi/core/enforce.h"
+
+static std::once_flag g_device_props_size_init_flag;
+static std::vector<std::unique_ptr<std::once_flag>> g_device_props_init_flags;
+static std::vector<phi::gpuDeviceProp> g_device_props;
+
+namespace phi {
+namespace backends {
+namespace gpu {
+
+int DnnVersion() {
+  if (!dynload::HasCUDNN()) return -1;
+  size_t version_major, version_minor, version_patch;
+  PADDLE_ENFORCE_GPU_SUCCESS(dynload::miopenGetVersion(
+      &version_major, &version_minor, &version_patch));
+  return version_major * 100 + version_minor * 10 + version_patch;
+}
+
+static int GetGPUDeviceCountImpl() {
+  int driverVersion = 0;
+  hipError_t status = hipDriverGetVersion(&driverVersion);
+
+  if (!(status == gpuSuccess && driverVersion != 0)) {
+    // No GPU driver
+    VLOG(2) << "GPU Driver Version can't be detected. No GPU driver!";
+    return 0;
+  }
+
+  const auto *cuda_visible_devices = std::getenv("HIP_VISIBLE_DEVICES");
+
+  if (cuda_visible_devices != nullptr) {
+    std::string cuda_visible_devices_str(cuda_visible_devices);
+    if (!cuda_visible_devices_str.empty()) {
+      cuda_visible_devices_str.erase(
+          0, cuda_visible_devices_str.find_first_not_of('\''));
+      cuda_visible_devices_str.erase(
+          cuda_visible_devices_str.find_last_not_of('\'') + 1);
+      cuda_visible_devices_str.erase(
+          0, cuda_visible_devices_str.find_first_not_of('\"'));
+      cuda_visible_devices_str.erase(
+          cuda_visible_devices_str.find_last_not_of('\"') + 1);
+    }
+    if (std::all_of(cuda_visible_devices_str.begin(),
+                    cuda_visible_devices_str.end(),
+                    [](char ch) { return ch == ' '; })) {
+      VLOG(2) << "HIP_VISIBLE_DEVICES is set to be "
+                 "empty. No GPU detected.";
+      return 0;
+    }
+  }
+  int count;
+  status = hipGetDeviceCount(&count);
+  if (status != hipSuccess) {
+    VLOG(2) << "You have gpu driver and rocm installed, but the machine does "
+               "not have any visible gpu card.";
+    count = 0;
+  }
+  return count;
+}
+
+int GetGPUDeviceCount() {
+  // cache the count
+  static auto dev_cnt = GetGPUDeviceCountImpl();
+  return dev_cnt;
+}
+
+int GetGPUComputeCapability(int id) {
+  PADDLE_ENFORCE_LT(id,
+                    GetGPUDeviceCount(),
+                    common::errors::InvalidArgument(
+                        "Device id must be less than GPU count, "
+                        "but received id is: %d. GPU count is: %d.",
+                        id,
+                        GetGPUDeviceCount()));
+  int major, minor;
+  auto major_error_code = hipDeviceGetAttribute(
+      &major, hipDeviceAttributeComputeCapabilityMajor, id);
+  auto minor_error_code = hipDeviceGetAttribute(
+      &minor, hipDeviceAttributeComputeCapabilityMinor, id);
+
+  PADDLE_ENFORCE_GPU_SUCCESS(major_error_code);
+  PADDLE_ENFORCE_GPU_SUCCESS(minor_error_code);
+  return major * 100 + minor;
+}
+
+int GetGPURuntimeVersion(int id) {
+  PADDLE_ENFORCE_LT(id,
+                    GetGPUDeviceCount(),
+                    common::errors::InvalidArgument(
+                        "Device id must be less than GPU count, "
+                        "but received id is: %d. GPU count is: %d.",
+                        id,
+                        GetGPUDeviceCount()));
+  int runtime_version = 0;
+  PADDLE_ENFORCE_GPU_SUCCESS(hipRuntimeGetVersion(&runtime_version));
+  return runtime_version;
+}
+
+int GetGPUDriverVersion(int id) {
+  PADDLE_ENFORCE_LT(id,
+                    GetGPUDeviceCount(),
+                    common::errors::InvalidArgument(
+                        "Device id must be less than GPU count, "
+                        "but received id is: %d. GPU count is: %d.",
+                        id,
+                        GetGPUDeviceCount()));
+  int driver_version = 0;
+  PADDLE_ENFORCE_GPU_SUCCESS(hipDriverGetVersion(&driver_version));
+  return driver_version;
+}
+
+bool TensorCoreAvailable() { return false; }
+
+int GetGPUMultiProcessors(int id) {
+  PADDLE_ENFORCE_LT(id,
+                    GetGPUDeviceCount(),
+                    common::errors::InvalidArgument(
+                        "Device id must be less than GPU count, "
+                        "but received id is: %d. GPU count is: %d.",
+                        id,
+                        GetGPUDeviceCount()));
+  int count;
+  PADDLE_ENFORCE_GPU_SUCCESS(
+      hipDeviceGetAttribute(&count, hipDeviceAttributeMultiprocessorCount, id));
+  return count;
+}
+
+int GetGPUMaxThreadsPerMultiProcessor(int id) {
+  PADDLE_ENFORCE_LT(id,
+                    GetGPUDeviceCount(),
+                    common::errors::InvalidArgument(
+                        "Device id must be less than GPU count, "
+                        "but received id is: %d. GPU count is: %d.",
+                        id,
+                        GetGPUDeviceCount()));
+  int count;
+  PADDLE_ENFORCE_GPU_SUCCESS(hipDeviceGetAttribute(
+      &count, hipDeviceAttributeMaxThreadsPerMultiProcessor, id));
+
+  return count;
+}
+
+int GetGPUMaxThreadsPerBlock(int id) {
+  PADDLE_ENFORCE_LT(id,
+                    GetGPUDeviceCount(),
+                    common::errors::InvalidArgument(
+                        "Device id must be less than GPU count, "
+                        "but received id is: %d. GPU count is: %d.",
+                        id,
+                        GetGPUDeviceCount()));
+  int count;
+  PADDLE_ENFORCE_GPU_SUCCESS(
+      hipDeviceGetAttribute(&count, hipDeviceAttributeMaxThreadsPerBlock, id));
+  return count;
+}
+
+int GetCurrentDeviceId() {
+  int device_id;
+  PADDLE_ENFORCE_GPU_SUCCESS(hipGetDevice(&device_id));
+  return device_id;
+}
+
+std::array<unsigned int, 3> GetGpuMaxGridDimSize(int id) {
+  PADDLE_ENFORCE_LT(id,
+                    GetGPUDeviceCount(),
+                    common::errors::InvalidArgument(
+                        "Device id must be less than GPU count, "
+                        "but received id is: %d. GPU count is: %d.",
+                        id,
+                        GetGPUDeviceCount()));
+  std::array<unsigned int, 3> ret;
+  int size;
+  auto error_code_x =
+      hipDeviceGetAttribute(&size, hipDeviceAttributeMaxGridDimX, id);
+  PADDLE_ENFORCE_GPU_SUCCESS(error_code_x);
+  ret[0] = size;
+
+  auto error_code_y =
+      hipDeviceGetAttribute(&size, hipDeviceAttributeMaxGridDimY, id);
+  PADDLE_ENFORCE_GPU_SUCCESS(error_code_y);
+  ret[1] = size;
+
+  auto error_code_z =
+      hipDeviceGetAttribute(&size, hipDeviceAttributeMaxGridDimZ, id);
+  PADDLE_ENFORCE_GPU_SUCCESS(error_code_z);
+  ret[2] = size;
+  return ret;
+}
+
+std::pair<int, int> GetGpuStreamPriorityRange() {
+  int least_priority, greatest_priority;
+  PADDLE_ENFORCE_GPU_SUCCESS(
+      hipDeviceGetStreamPriorityRange(&least_priority, &greatest_priority));
+  return std::make_pair(least_priority, greatest_priority);
+}
+
+const gpuDeviceProp &GetDeviceProperties(int id) {
+  std::call_once(g_device_props_size_init_flag, [&] {
+    int gpu_num = 0;
+    gpu_num = GetGPUDeviceCount();
+    g_device_props_init_flags.resize(gpu_num);
+    g_device_props.resize(gpu_num);
+    for (int i = 0; i < gpu_num; ++i) {
+      g_device_props_init_flags[i] = std::make_unique<std::once_flag>();
+    }
+  });
+
+  if (id == -1) {
+    id = GetCurrentDeviceId();
+  }
+
+  if (id < 0 || id >= static_cast<int>(g_device_props.size())) {
+    PADDLE_THROW(common::errors::OutOfRange(
+        "The device id %d is out of range [0, %d), where %d is the number of "
+        "devices on this machine. Because the device id should be greater than "
+        "or equal to zero and smaller than the number of gpus. Please input "
+        "appropriate device again!",
+        id,
+        static_cast<int>(g_device_props.size()),
+        static_cast<int>(g_device_props.size())));
+  }
+
+  std::call_once(*(g_device_props_init_flags[id]), [&] {
+    PADDLE_ENFORCE_GPU_SUCCESS(hipGetDeviceProperties(&g_device_props[id], id));
+  });
+
+  return g_device_props[id];
+}
+
+void SetDeviceId(int id) {
+  static thread_local bool first_call = true;
+  if (first_call) {
+    PADDLE_ENFORCE_LT(id,
+                      GetGPUDeviceCount(),
+                      common::errors::InvalidArgument(
+                          "Device id must be less than GPU count, "
+                          "but received id is: %d. GPU count is: %d.",
+                          id,
+                          GetGPUDeviceCount()));
+
+    PADDLE_RETRY_CUDA_SUCCESS(hipSetDevice(id));
+    VLOG(4) << "SetDeviceId " << id;
+    first_call = false;
+    return;
+  }
+
+  int prev_id;
+  PADDLE_ENFORCE_GPU_SUCCESS(hipGetDevice(&prev_id));
+  if (prev_id != id) {
+    PADDLE_ENFORCE_LT(id,
+                      GetGPUDeviceCount(),
+                      common::errors::InvalidArgument(
+                          "Device id must be less than GPU count, "
+                          "but received id is: %d. GPU count is: %d.",
+                          id,
+                          GetGPUDeviceCount()));
+
+    PADDLE_RETRY_CUDA_SUCCESS(hipSetDevice(id));
+    VLOG(4) << "SetDeviceId " << id;
+  }
+}
+
+void GpuMemcpyAsync(void *dst,
+                    const void *src,
+                    size_t count,
+                    gpuMemcpyKind kind,
+                    gpuStream_t stream) {
+  PADDLE_ENFORCE_GPU_SUCCESS(hipMemcpyAsync(dst, src, count, kind, stream));
+}
+
+void GpuMemcpySync(void *dst,
+                   const void *src,
+                   size_t count,
+                   gpuMemcpyKind kind) {
+  PADDLE_ENFORCE_GPU_SUCCESS(hipMemcpy(dst, src, count, kind));
+}
+
+void GpuMemcpyPeerAsync(void *dst,
+                        int dst_device,
+                        const void *src,
+                        int src_device,
+                        size_t count,
+                        gpuStream_t stream) {
+  PADDLE_ENFORCE_GPU_SUCCESS(
+      hipMemcpyPeerAsync(dst, dst_device, src, src_device, count, stream));
+}
+
+void GpuMemcpyPeerSync(
+    void *dst, int dst_device, const void *src, int src_device, size_t count) {
+  PADDLE_ENFORCE_GPU_SUCCESS(
+      hipMemcpyPeer(dst, dst_device, src, src_device, count));
+}
+
+void GpuMemsetAsync(void *dst, int value, size_t count, gpuStream_t stream) {
+  PADDLE_ENFORCE_GPU_SUCCESS(hipMemsetAsync(dst, value, count, stream));
+}
+
+void GpuStreamSync(gpuStream_t stream) {
+  PADDLE_ENFORCE_GPU_SUCCESS(hipStreamSynchronize(stream));
+}
+
+void GpuDestroyStream(gpuStream_t stream) {
+  PADDLE_ENFORCE_GPU_SUCCESS(hipStreamDestroy(stream));
+}
+
+void GpuDeviceSync() { PADDLE_ENFORCE_GPU_SUCCESS(hipDeviceSynchronize()); }
+
+gpuError_t GpuGetLastError() { return hipGetLastError(); }
+
+bool IsGPUManagedMemorySupported(int dev_id) {
+  PADDLE_ENFORCE_LT(dev_id,
+                    GetGPUDeviceCount(),
+                    common::errors::InvalidArgument(
+                        "Device id must be less than GPU count, "
+                        "but received id is: %d. GPU count is: %d.",
+                        dev_id,
+                        GetGPUDeviceCount()));
+  // TODO(qili93): Hygon DTK (21.04 and 22.04) not support
+  // hipDeviceAttributeManagedMemory, temporary disable by default, to be
+  // verified in next DTK release
+  return false;
+}
+
+bool IsGPUManagedMemoryOversubscriptionSupported(int dev_id) {
+  PADDLE_ENFORCE_LT(dev_id,
+                    GetGPUDeviceCount(),
+                    common::errors::InvalidArgument(
+                        "Device id must be less than GPU count, "
+                        "but received id is: %d. GPU count is: %d.",
+                        dev_id,
+                        GetGPUDeviceCount()));
+#ifdef __linux__
+  return IsGPUManagedMemorySupported(dev_id) &&
+         GetGPUComputeCapability(dev_id) >= 60;
+#else
+  return false;
+#endif
+}
+
+}  // namespace gpu
+}  // namespace backends
+}  // namespace phi