wehub-resource-sync
463 lines
19 KiB
C++
463 lines
19 KiB
C++
// Copyright (c) 2022 PaddlePaddle Authors. All Rights Reserved.
|
|
// Copyright 2022 The DGL team for some useful functions.
|
|
|
|
// 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
|
|
// thrust headers require nvcc/hipcc
|
|
// (rocThrust 7.0+ pulls in rocprim)
|
|
#if defined(__NVCC__) || defined(__HIPCC__)
|
|
#include <thrust/device_vector.h>
|
|
#include <thrust/fill.h>
|
|
#endif
|
|
|
|
#include "paddle/common/hostdevice.h"
|
|
#include "paddle/phi/backends/gpu/gpu_context.h"
|
|
#include "paddle/phi/backends/gpu/gpu_primitives.h"
|
|
#include "paddle/phi/kernels/impl/graph_message_passing_impl.h"
|
|
|
|
namespace phi {
|
|
|
|
#define CUDA_MAX_NUM_BLOCKS_X 0x7FFFFFFF
|
|
#define CUDA_MAX_NUM_BLOCKS_Y 0xFFFF
|
|
#define CUDA_MAX_NUM_BLOCKS_Z 0xFFFF
|
|
|
|
inline void CopyBCastOff(const BroadCastInfo& bcast_info,
|
|
thrust::device_vector<int64_t>* l_bcastoff,
|
|
thrust::device_vector<int64_t>* r_bcastoff) {
|
|
l_bcastoff->resize(bcast_info.out_len);
|
|
r_bcastoff->resize(bcast_info.out_len);
|
|
#ifdef PADDLE_WITH_HIP
|
|
hipMemcpy(thrust::raw_pointer_cast(l_bcastoff->data()),
|
|
bcast_info.l_offset.data(),
|
|
sizeof(int64_t) * bcast_info.out_len,
|
|
hipMemcpyHostToDevice);
|
|
hipMemcpy(thrust::raw_pointer_cast(r_bcastoff->data()),
|
|
bcast_info.r_offset.data(),
|
|
sizeof(int64_t) * bcast_info.out_len,
|
|
hipMemcpyHostToDevice);
|
|
#else
|
|
cudaMemcpy(thrust::raw_pointer_cast(l_bcastoff->data()),
|
|
bcast_info.l_offset.data(),
|
|
sizeof(int64_t) * bcast_info.out_len,
|
|
cudaMemcpyHostToDevice);
|
|
cudaMemcpy(thrust::raw_pointer_cast(r_bcastoff->data()),
|
|
bcast_info.r_offset.data(),
|
|
sizeof(int64_t) * bcast_info.out_len,
|
|
cudaMemcpyHostToDevice);
|
|
#endif
|
|
}
|
|
|
|
inline int FindNumThreads(int64_t dim, int max_num_threads) {
|
|
PADDLE_ENFORCE_GE(dim,
|
|
0,
|
|
common::errors::PreconditionNotMet(
|
|
"Required dim >= 0, but received dim = %d", dim));
|
|
int res = max_num_threads;
|
|
if (dim == 0) res = 1;
|
|
while (res > dim) {
|
|
res = res >> 1;
|
|
}
|
|
res = res <= 32 ? 32 : res;
|
|
return res;
|
|
}
|
|
|
|
inline int FindNumBlocks(char axis, int64_t nblocks, int max_num_blocks = -1) {
|
|
int default_max_num_blocks = -1;
|
|
switch (axis) {
|
|
case 'x':
|
|
default_max_num_blocks = CUDA_MAX_NUM_BLOCKS_X;
|
|
break;
|
|
case 'y':
|
|
default_max_num_blocks = CUDA_MAX_NUM_BLOCKS_Y;
|
|
break;
|
|
case 'z':
|
|
default_max_num_blocks = CUDA_MAX_NUM_BLOCKS_Z;
|
|
break;
|
|
default:
|
|
PADDLE_THROW(
|
|
common::errors::InvalidArgument("%c axis is not recognized", axis));
|
|
}
|
|
if (max_num_blocks == -1) {
|
|
max_num_blocks = default_max_num_blocks;
|
|
}
|
|
PADDLE_ENFORCE_GE(nblocks,
|
|
0,
|
|
common::errors::InvalidArgument(
|
|
"The number of CUDA blocks must be non-negative. "
|
|
"Expected nblocks >= 0, but received nblocks = %ld.",
|
|
nblocks));
|
|
PADDLE_ENFORCE_GT(
|
|
max_num_blocks,
|
|
0,
|
|
common::errors::InvalidArgument("max_num_blocks should be larger than 0, "
|
|
"but received %d",
|
|
max_num_blocks));
|
|
if (nblocks < max_num_blocks) {
|
|
return static_cast<int>(nblocks);
|
|
}
|
|
return max_num_blocks;
|
|
}
|
|
|
|
template <typename T>
|
|
struct GraphSendUERecvSumCUDAFunctor {
|
|
DEVICE inline void operator()(T* output, T val) {
|
|
CudaAtomicAdd(output, val);
|
|
}
|
|
};
|
|
|
|
template <typename T>
|
|
struct GraphSendUERecvMaxCUDAFunctor {
|
|
DEVICE inline void operator()(T* output, T val) {
|
|
CudaAtomicMax(output, val);
|
|
}
|
|
};
|
|
|
|
template <typename T>
|
|
struct GraphSendUERecvMinCUDAFunctor {
|
|
DEVICE inline void operator()(T* output, T val) {
|
|
CudaAtomicMin(output, val);
|
|
}
|
|
};
|
|
|
|
template <typename T,
|
|
typename IndexT,
|
|
typename ReduceFunctor,
|
|
typename ComputeFunctor>
|
|
__global__ void GraphSendUERecvCUDAKernel(const T* x_data,
|
|
const T* e_data,
|
|
const IndexT* src_indices,
|
|
const IndexT* dst_indices,
|
|
const int64_t* xbcast_off,
|
|
const int64_t* ebcast_off,
|
|
T* output,
|
|
int64_t index_size,
|
|
int64_t x_len,
|
|
int64_t e_len,
|
|
int64_t out_len,
|
|
bool use_bcast,
|
|
ComputeFunctor cfunctor,
|
|
ReduceFunctor rfunctor) {
|
|
IndexT ty = static_cast<IndexT>(blockIdx.y) * blockDim.y + threadIdx.y;
|
|
const IndexT stride_y = static_cast<IndexT>(blockDim.y) * gridDim.y;
|
|
|
|
while (ty < index_size) {
|
|
IndexT src = src_indices[ty];
|
|
IndexT dst = dst_indices[ty];
|
|
int64_t tx = static_cast<int64_t>(blockIdx.x) * blockDim.x + threadIdx.x;
|
|
int64_t stride_x = blockDim.x * static_cast<int64_t>(gridDim.x);
|
|
|
|
const T* x_off = x_data + src * x_len;
|
|
const T* e_off = e_data + ty * e_len;
|
|
T* out_off = output + dst * out_len;
|
|
while (tx < out_len) {
|
|
int64_t x_add = use_bcast ? xbcast_off[tx] : tx;
|
|
int64_t e_add = use_bcast ? ebcast_off[tx] : tx;
|
|
T val = cfunctor(x_off[x_add], e_off[e_add]);
|
|
rfunctor(out_off + tx, val);
|
|
tx += stride_x;
|
|
}
|
|
ty += stride_y;
|
|
}
|
|
}
|
|
|
|
// x_grad: for backward mean with mul.
|
|
template <typename T, typename IndexT>
|
|
__global__ void ManipulateMeanGradCUDAKernelForMulX(const T* out_grad_data,
|
|
const T* e_data,
|
|
const IndexT* src_indices,
|
|
const IndexT* dst_indices,
|
|
const int* dst_count,
|
|
const int64_t* l_bcastoff,
|
|
const int64_t* r_bcastoff,
|
|
T* x_grad,
|
|
int64_t index_size,
|
|
int64_t l_len,
|
|
int64_t r_len,
|
|
int64_t out_len,
|
|
bool use_bcast) {
|
|
IndexT ty = blockIdx.y * blockDim.y + threadIdx.y;
|
|
const IndexT stride_y = blockDim.y * gridDim.y;
|
|
|
|
while (ty < index_size) {
|
|
IndexT src = src_indices[ty];
|
|
IndexT dst = dst_indices[ty];
|
|
int64_t tx =
|
|
static_cast<int64_t>(blockIdx.x) * static_cast<int64_t>(blockDim.x) +
|
|
static_cast<int64_t>(threadIdx.x);
|
|
int64_t stride_x =
|
|
static_cast<int64_t>(blockDim.x) * static_cast<int64_t>(gridDim.x);
|
|
|
|
const T* out_grad_off = out_grad_data + src * l_len;
|
|
const T* e_off = e_data + ty * r_len;
|
|
T* x_grad_off = x_grad + dst * out_len;
|
|
while (tx < out_len) {
|
|
int64_t o_add = use_bcast ? l_bcastoff[tx] : tx;
|
|
int64_t e_add = use_bcast ? r_bcastoff[tx] : tx;
|
|
T val = out_grad_off[o_add] * e_off[e_add];
|
|
CudaAtomicAdd(x_grad_off + tx, val / static_cast<T>(dst_count[src]));
|
|
tx += stride_x;
|
|
}
|
|
ty += stride_y;
|
|
}
|
|
}
|
|
|
|
// e_grad: backward sum for add.
|
|
template <typename T, typename IndexT>
|
|
__global__ void ManipulateSumGradCUDAKernelForAddE(const T* out_grad_data,
|
|
const IndexT* dst_indices,
|
|
const int64_t* r_bcastoff,
|
|
T* e_grad,
|
|
int64_t index_size,
|
|
int64_t r_len,
|
|
int64_t out_len,
|
|
bool use_bcast) {
|
|
IndexT ty = blockIdx.y * blockDim.y + threadIdx.y;
|
|
const IndexT stride_y = blockDim.y * gridDim.y;
|
|
|
|
while (ty < index_size) {
|
|
IndexT dst = dst_indices[ty];
|
|
int64_t tx =
|
|
static_cast<int64_t>(blockIdx.x) * static_cast<int64_t>(blockDim.x) +
|
|
static_cast<int64_t>(threadIdx.x);
|
|
int64_t stride_x =
|
|
static_cast<int64_t>(blockDim.x) * static_cast<int64_t>(gridDim.x);
|
|
|
|
T* e_grad_off = e_grad + ty * r_len;
|
|
const T* out_grad_off = out_grad_data + dst * out_len;
|
|
while (tx < out_len) {
|
|
int64_t e_add = use_bcast ? r_bcastoff[tx] : tx;
|
|
CudaAtomicAdd(e_grad_off + e_add, out_grad_off[tx]);
|
|
tx += stride_x;
|
|
}
|
|
ty += stride_y;
|
|
}
|
|
}
|
|
|
|
// e_grad: backward sum for mul.
|
|
template <typename T, typename IndexT>
|
|
__global__ void ManipulateSumGradCUDAKernelForMulE(const T* x_data,
|
|
const T* out_grad_data,
|
|
const IndexT* src_indices,
|
|
const IndexT* dst_indices,
|
|
const int64_t* l_bcastoff,
|
|
const int64_t* r_bcastoff,
|
|
T* e_grad,
|
|
int64_t index_size,
|
|
int64_t l_len,
|
|
int64_t r_len,
|
|
int64_t out_len,
|
|
bool use_bcast) {
|
|
IndexT ty = blockIdx.y * blockDim.y + threadIdx.y;
|
|
const IndexT stride_y = blockDim.y * gridDim.y;
|
|
|
|
while (ty < index_size) {
|
|
IndexT src = src_indices[ty];
|
|
IndexT dst = dst_indices[ty];
|
|
int64_t tx =
|
|
static_cast<int64_t>(blockIdx.x) * static_cast<int64_t>(blockDim.x) +
|
|
static_cast<int64_t>(threadIdx.x);
|
|
int64_t stride_x =
|
|
static_cast<int64_t>(blockDim.x) * static_cast<int64_t>(gridDim.x);
|
|
|
|
const T* x_off = x_data + src * l_len;
|
|
T* e_grad_off = e_grad + ty * r_len;
|
|
const T* out_grad_off = out_grad_data + dst * out_len;
|
|
while (tx < out_len) {
|
|
int64_t x_add = use_bcast ? l_bcastoff[tx] : tx;
|
|
int64_t e_add = use_bcast ? r_bcastoff[tx] : tx;
|
|
CudaAtomicAdd(e_grad_off + e_add, out_grad_off[tx] * x_off[x_add]);
|
|
tx += stride_x;
|
|
}
|
|
ty += stride_y;
|
|
}
|
|
}
|
|
|
|
// e_grad: backward mean for add
|
|
template <typename T, typename IndexT>
|
|
__global__ void ManipulateMeanGradCUDAKernelForAddE(const T* out_grad_data,
|
|
const IndexT* dst_indices,
|
|
const int* dst_count,
|
|
const int64_t* r_bcastoff,
|
|
T* e_grad,
|
|
int64_t index_size,
|
|
int64_t r_len,
|
|
int64_t out_len,
|
|
bool use_bcast) {
|
|
IndexT ty = blockIdx.y * blockDim.y + threadIdx.y;
|
|
const IndexT stride_y = blockDim.y * gridDim.y;
|
|
|
|
while (ty < index_size) {
|
|
IndexT dst = dst_indices[ty];
|
|
int64_t tx =
|
|
static_cast<int64_t>(blockIdx.x) * static_cast<int64_t>(blockDim.x) +
|
|
static_cast<int64_t>(threadIdx.x);
|
|
int64_t stride_x =
|
|
static_cast<int64_t>(blockDim.x) * static_cast<int64_t>(gridDim.x);
|
|
|
|
T* e_grad_off = e_grad + ty * r_len;
|
|
const T* out_grad_off = out_grad_data + dst * out_len;
|
|
while (tx < out_len) {
|
|
int64_t e_add = use_bcast ? r_bcastoff[tx] : tx;
|
|
CudaAtomicAdd(e_grad_off + e_add,
|
|
out_grad_off[tx] / static_cast<T>(dst_count[dst]));
|
|
tx += stride_x;
|
|
}
|
|
ty += stride_y;
|
|
}
|
|
}
|
|
|
|
// e_grad: backward mean for mul.
|
|
template <typename T, typename IndexT>
|
|
__global__ void ManipulateMeanGradCUDAKernelForMulE(const T* x_data,
|
|
const T* out_grad_data,
|
|
const IndexT* src_indices,
|
|
const IndexT* dst_indices,
|
|
const int* dst_count,
|
|
const int64_t* l_bcastoff,
|
|
const int64_t* r_bcastoff,
|
|
T* e_grad,
|
|
int64_t index_size,
|
|
int64_t l_len,
|
|
int64_t r_len,
|
|
int64_t out_len,
|
|
bool use_bcast) {
|
|
IndexT ty = blockIdx.y * blockDim.y + threadIdx.y;
|
|
const IndexT stride_y = blockDim.y * gridDim.y;
|
|
|
|
while (ty < index_size) {
|
|
IndexT src = src_indices[ty];
|
|
IndexT dst = dst_indices[ty];
|
|
int64_t tx =
|
|
static_cast<int64_t>(blockIdx.x) * static_cast<int64_t>(blockDim.x) +
|
|
static_cast<int64_t>(threadIdx.x);
|
|
int64_t stride_x =
|
|
static_cast<int64_t>(blockDim.x) * static_cast<int64_t>(gridDim.x);
|
|
|
|
const T* x_off = x_data + src * l_len;
|
|
T* e_grad_off = e_grad + ty * r_len;
|
|
const T* out_grad_off = out_grad_data + dst * out_len;
|
|
while (tx < out_len) {
|
|
int64_t x_add = use_bcast ? l_bcastoff[tx] : tx;
|
|
int64_t e_add = use_bcast ? r_bcastoff[tx] : tx;
|
|
CudaAtomicAdd(
|
|
e_grad_off + e_add,
|
|
out_grad_off[tx] * x_off[x_add] / static_cast<T>(dst_count[dst]));
|
|
tx += stride_x;
|
|
}
|
|
ty += stride_y;
|
|
}
|
|
}
|
|
|
|
// x_grad, e_grad: backward min and max for add.
|
|
template <typename T, typename IndexT>
|
|
__global__ void ManipulateMinMaxGradCUDAKernelForAdd(const T* x_data,
|
|
const T* e_data,
|
|
const T* out,
|
|
const T* out_grad,
|
|
const IndexT* src_indices,
|
|
const IndexT* dst_indices,
|
|
const int64_t* xbcast_off,
|
|
const int64_t* ebcast_off,
|
|
T* x_grad,
|
|
T* e_grad,
|
|
int64_t index_size,
|
|
int64_t x_len,
|
|
int64_t e_len,
|
|
int64_t out_len,
|
|
bool use_bcast) {
|
|
IndexT ty = blockIdx.y * blockDim.y + threadIdx.y;
|
|
const IndexT stride_y = blockDim.y * gridDim.y;
|
|
|
|
while (ty < index_size) {
|
|
IndexT src = src_indices[ty];
|
|
IndexT dst = dst_indices[ty];
|
|
int64_t tx =
|
|
static_cast<int64_t>(blockIdx.x) * static_cast<int64_t>(blockDim.x) +
|
|
static_cast<int64_t>(threadIdx.x);
|
|
int64_t stride_x =
|
|
static_cast<int64_t>(blockDim.x) * static_cast<int64_t>(gridDim.x);
|
|
|
|
const T* x_off = x_data + dst * x_len;
|
|
const T* e_off = e_data + ty * e_len;
|
|
const T* out_off = out + src * out_len;
|
|
const T* out_grad_off = out_grad + src * out_len;
|
|
T* x_grad_off = x_grad + dst * x_len;
|
|
T* e_grad_off = e_grad + ty * e_len;
|
|
while (tx < out_len) {
|
|
int64_t x_add = use_bcast ? xbcast_off[tx] : tx;
|
|
int64_t e_add = use_bcast ? ebcast_off[tx] : tx;
|
|
T val = x_off[x_add] + e_off[e_add];
|
|
CudaAtomicAdd(x_grad_off + x_add,
|
|
out_grad_off[tx] * static_cast<T>(val == out_off[tx]));
|
|
CudaAtomicAdd(e_grad_off + e_add,
|
|
out_grad_off[tx] * static_cast<T>(val == out_off[tx]));
|
|
tx += stride_x;
|
|
}
|
|
ty += stride_y;
|
|
}
|
|
}
|
|
|
|
// x_grad, e_grad: backward min and max for mul.
|
|
template <typename T, typename IndexT>
|
|
__global__ void ManipulateMinMaxGradCUDAKernelForMul(const T* x_data,
|
|
const T* e_data,
|
|
const T* out,
|
|
const T* out_grad,
|
|
const IndexT* src_indices,
|
|
const IndexT* dst_indices,
|
|
const int64_t* xbcast_off,
|
|
const int64_t* ebcast_off,
|
|
T* x_grad,
|
|
T* e_grad,
|
|
int64_t index_size,
|
|
int64_t x_len,
|
|
int64_t e_len,
|
|
int64_t out_len,
|
|
bool use_bcast) {
|
|
IndexT ty = blockIdx.y * blockDim.y + threadIdx.y;
|
|
const IndexT stride_y = blockDim.y * gridDim.y;
|
|
|
|
while (ty < index_size) {
|
|
IndexT src = src_indices[ty];
|
|
IndexT dst = dst_indices[ty];
|
|
int64_t tx =
|
|
static_cast<int64_t>(blockIdx.x) * static_cast<int64_t>(blockDim.x) +
|
|
static_cast<int64_t>(threadIdx.x);
|
|
int64_t stride_x =
|
|
static_cast<int64_t>(blockDim.x) * static_cast<int64_t>(gridDim.x);
|
|
|
|
const T* x_off = x_data + dst * x_len;
|
|
const T* e_off = e_data + ty * e_len;
|
|
const T* out_off = out + src * out_len;
|
|
const T* out_grad_off = out_grad + src * out_len;
|
|
T* x_grad_off = x_grad + dst * x_len;
|
|
T* e_grad_off = e_grad + ty * e_len;
|
|
while (tx < out_len) {
|
|
int64_t x_add = use_bcast ? xbcast_off[tx] : tx;
|
|
int64_t e_add = use_bcast ? ebcast_off[tx] : tx;
|
|
T val = x_off[x_add] * e_off[e_add];
|
|
CudaAtomicAdd(
|
|
x_grad_off + x_add,
|
|
out_grad_off[tx] * static_cast<T>(val == out_off[tx]) * e_off[e_add]);
|
|
CudaAtomicAdd(
|
|
e_grad_off + e_add,
|
|
out_grad_off[tx] * static_cast<T>(val == out_off[tx]) * x_off[x_add]);
|
|
tx += stride_x;
|
|
}
|
|
ty += stride_y;
|
|
}
|
|
}
|
|
|
|
} // namespace phi
|