| #include "caffe2/operators/pool_op.h" |
| |
| #include <array> |
| #include <functional> |
| #include <limits> |
| #include <numeric> |
| |
| #include "caffe2/core/context_gpu.h" |
| #include "caffe2/utils/math.h" |
| |
| namespace caffe2 { |
| |
| namespace { |
| |
| template <typename T> |
| __global__ void AveragePool1DForwardNCHWCUDAKernel( |
| const int X_size, |
| const int Y_size, |
| const int kernel, |
| const int stride, |
| const int pad, |
| const bool count_include_pad, |
| const T* X, |
| T* Y) { |
| const int nc = blockIdx.x; |
| const T* X_ptr = X + nc * X_size; |
| T* Y_ptr = Y + nc * Y_size; |
| for (int y = threadIdx.x; y < Y_size; y += blockDim.x) { |
| const int x = y * stride - pad; |
| const int l = max(x, 0); |
| const int r = min(x + kernel, X_size); |
| const T scale = T(1) / static_cast<T>(count_include_pad ? kernel : r - l); |
| T sum = 0; |
| for (int i = l; i < r; ++i) { |
| #if __CUDA_ARCH__ >= 350 |
| sum += __ldg(X_ptr + i); |
| #else |
| sum += X_ptr[i]; |
| #endif |
| } |
| Y_ptr[y] = sum * scale; |
| } |
| } |
| |
| template <typename T> |
| __global__ void AveragePool1DForwardNHWCCUDAKernel( |
| const int C, |
| const int X_size, |
| const int Y_size, |
| const int kernel, |
| const int stride, |
| const int pad, |
| const bool count_include_pad, |
| const T* X, |
| T* Y) { |
| const int n = blockIdx.x / Y_size; |
| const int y = blockIdx.x % Y_size; |
| const int x = y * stride - pad; |
| const int l = max(x, 0); |
| const int r = min(x + kernel, X_size); |
| const T scale = T(1) / static_cast<T>(count_include_pad ? kernel : r - l); |
| const T* X_ptr = X + n * X_size * C; |
| T* Y_ptr = Y + n * Y_size * C; |
| for (int c = threadIdx.x; c < C; c += blockDim.x) { |
| T sum = 0; |
| for (int i = l; i < r; ++i) { |
| #if __CUDA_ARCH__ >= 350 |
| sum += __ldg(X_ptr + i * C + c); |
| #else |
| sum += X_ptr[i * C + c]; |
| #endif |
| } |
| Y_ptr[y * C + c] = sum * scale; |
| } |
| } |
| |
| template <typename T> |
| __global__ void AveragePool2DForwardNCHWCUDAKernel( |
| const int X_H, |
| const int X_W, |
| const int Y_H, |
| const int Y_W, |
| const int kernel_h, |
| const int kernel_w, |
| const int stride_h, |
| const int stride_w, |
| const int pad_t, |
| const int pad_l, |
| const bool count_include_pad, |
| const T* X, |
| T* Y) { |
| const int X_HxW = X_H * X_W; |
| const int Y_HxW = Y_H * Y_W; |
| const int nc = blockIdx.x / Y_H; |
| const int yh = blockIdx.x % Y_H; |
| const T* X_ptr = X + nc * X_HxW; |
| T* Y_ptr = Y + nc * Y_HxW; |
| const int xh = yh * stride_h - pad_t; |
| const int t = max(xh, 0); |
| const int b = min(xh + kernel_h, X_H); |
| for (int yw = threadIdx.x; yw < Y_W; yw += blockDim.x) { |
| const int xw = yw * stride_w - pad_l; |
| const int l = max(xw, 0); |
| const int r = min(xw + kernel_w, X_W); |
| const T scale = T(1) / |
| static_cast<T>(count_include_pad ? kernel_h * kernel_w |
| : (b - t) * (r - l)); |
| T sum = 0; |
| for (int i = t; i < b; ++i) { |
| for (int j = l; j < r; ++j) { |
| #if __CUDA_ARCH__ >= 350 |
| sum += __ldg(X_ptr + i * X_W + j); |
| #else |
| sum += X_ptr[i * X_W + j]; |
| #endif |
| } |
| } |
| Y_ptr[yh * Y_W + yw] = sum * scale; |
| } |
| } |
| |
| template <typename T> |
| __global__ void AveragePool2DForwardNHWCCUDAKernel( |
| const int C, |
| const int X_H, |
| const int X_W, |
| const int Y_H, |
| const int Y_W, |
| const int kernel_h, |
| const int kernel_w, |
| const int stride_h, |
| const int stride_w, |
| const int pad_t, |
| const int pad_l, |
| const bool count_include_pad, |
| const T* X, |
| T* Y) { |
| const int X_HxW = X_H * X_W; |
| const int Y_HxW = Y_H * Y_W; |
| const int n = blockIdx.x / Y_HxW; |
| const int y = blockIdx.x % Y_HxW; |
| const int yh = y / Y_W; |
| const int yw = y % Y_W; |
| const int xh = yh * stride_h - pad_t; |
| const int xw = yw * stride_w - pad_l; |
| const int t = max(xh, 0); |
| const int b = min(xh + kernel_h, X_H); |
| const int l = max(xw, 0); |
| const int r = min(xw + kernel_w, X_W); |
| const T scale = T(1) / |
| static_cast<T>(count_include_pad ? kernel_h * kernel_w |
| : (b - t) * (r - l)); |
| const T* X_ptr = X + n * X_HxW * C; |
| T* Y_ptr = Y + n * Y_HxW * C; |
| for (int c = threadIdx.x; c < C; c += blockDim.x) { |
| T sum = 0; |
| for (int i = t; i < b; ++i) { |
| for (int j = l; j < r; ++j) { |
| #if __CUDA_ARCH__ >= 350 |
| sum += __ldg(X_ptr + (i * X_W + j) * C + c); |
| #else |
| sum += X_ptr[(i * X_W + j) * C + c]; |
| #endif |
| } |
| } |
| Y_ptr[y * C + c] = sum * scale; |
| } |
| } |
| |
| template <typename T> |
| __global__ void AveragePool3DForwardNCHWCUDAKernel( |
| const int X_D, |
| const int X_H, |
| const int X_W, |
| const int Y_D, |
| const int Y_H, |
| const int Y_W, |
| const int kernel_d, |
| const int kernel_h, |
| const int kernel_w, |
| const int stride_d, |
| const int stride_h, |
| const int stride_w, |
| const int pad_p, |
| const int pad_t, |
| const int pad_l, |
| const bool count_include_pad, |
| const T* X, |
| T* Y) { |
| const int X_HxW = X_D * X_H * X_W; |
| const int Y_HxW = Y_D * Y_H * Y_W; |
| const int yy = blockIdx.x / Y_H; |
| const int nc = yy / Y_D; |
| const int yd = yy % Y_D; |
| const int yh = blockIdx.x % Y_H; |
| const T* X_ptr = X + nc * X_HxW; |
| T* Y_ptr = Y + nc * Y_HxW; |
| const int xd = yd * stride_d - pad_p; |
| const int xh = yh * stride_h - pad_t; |
| const int p = max(xd, 0); |
| const int a = min(xd + kernel_d, X_D); |
| const int t = max(xh, 0); |
| const int b = min(xh + kernel_h, X_H); |
| for (int yw = threadIdx.x; yw < Y_W; yw += blockDim.x) { |
| const int xw = yw * stride_w - pad_l; |
| const int l = max(xw, 0); |
| const int r = min(xw + kernel_w, X_W); |
| const T scale = T(1) / |
| static_cast<T>(count_include_pad ? kernel_d * kernel_h * kernel_w |
| : (a - p) * (b - t) * (r - l)); |
| T sum = 0; |
| for (int i = p; i < a; ++i) { |
| for (int j = t; j < b; ++j) { |
| for (int k = l; k < r; ++k) { |
| #if __CUDA_ARCH__ >= 350 |
| sum += __ldg(X_ptr + (i * X_H + j) * X_W + k); |
| #else |
| sum += X_ptr[(i * X_H + j) * X_W + k]; |
| #endif |
| } |
| } |
| } |
| Y_ptr[(yd * Y_H + yh) * Y_W + yw] = sum * scale; |
| } |
| } |
| |
| template <typename T> |
| __global__ void AveragePool3DForwardNHWCCUDAKernel( |
| const int C, |
| const int X_D, |
| const int X_H, |
| const int X_W, |
| const int Y_D, |
| const int Y_H, |
| const int Y_W, |
| const int kernel_d, |
| const int kernel_h, |
| const int kernel_w, |
| const int stride_d, |
| const int stride_h, |
| const int stride_w, |
| const int pad_p, |
| const int pad_t, |
| const int pad_l, |
| const bool count_include_pad, |
| const T* X, |
| T* Y) { |
| const int X_HxW = X_D * X_H * X_W; |
| const int Y_HxW = Y_D * Y_H * Y_W; |
| const int n = blockIdx.x / Y_HxW; |
| const int y = blockIdx.x % Y_HxW; |
| const int yy = y / Y_W; |
| const int yd = yy / Y_H; |
| const int yh = yy % Y_H; |
| const int yw = y % Y_W; |
| const int xd = yd * stride_d - pad_p; |
| const int xh = yh * stride_h - pad_t; |
| const int xw = yw * stride_w - pad_l; |
| const int p = max(xd, 0); |
| const int a = min(xd + kernel_d, X_D); |
| const int t = max(xh, 0); |
| const int b = min(xh + kernel_h, X_H); |
| const int l = max(xw, 0); |
| const int r = min(xw + kernel_w, X_W); |
| const T scale = T(1) / |
| static_cast<T>(count_include_pad ? kernel_d * kernel_h * kernel_w |
| : (a - p) * (b - t) * (r - l)); |
| const T* X_ptr = X + n * X_HxW * C; |
| T* Y_ptr = Y + n * Y_HxW * C; |
| for (int c = threadIdx.x; c < C; c += blockDim.x) { |
| T sum = 0; |
| for (int i = p; i < a; ++i) { |
| for (int j = t; j < b; ++j) { |
| for (int k = l; k < r; ++k) { |
| #if __CUDA_ARCH__ >= 350 |
| sum += __ldg(X_ptr + ((i * X_H + j) * X_W + k) * C + c); |
| #else |
| sum += X_ptr[((i * X_H + j) * X_W + k) * C + c]; |
| #endif |
| } |
| } |
| } |
| Y_ptr[y * C + c] = sum * scale; |
| } |
| } |
| |
| template <typename T> |
| __global__ void GlobalAveragePoolingBackwardNCHWCUDAKernel( |
| const int K, |
| const int HxW, |
| const T scale, |
| const T* dY, |
| T* dX) { |
| const int nc = blockIdx.x / K; |
| const int block = blockIdx.x % K; |
| const int x = threadIdx.x + block * CAFFE_CUDA_NUM_THREADS; |
| if (x < HxW) { |
| #if __CUDA_ARCH__ >= 350 |
| dX[nc * HxW + x] = __ldg(dY + nc) * scale; |
| #else |
| dX[nc * HxW + x] = dY[nc] * scale; |
| #endif |
| } |
| } |
| |
| template <typename T> |
| __global__ void GlobalAveragePoolingBackwardNHWCCUDAKernel( |
| const int C, |
| const int HxW, |
| const T scale, |
| const T* dY, |
| T* dX) { |
| const int n = blockIdx.x / HxW; |
| for (int c = threadIdx.x; c < C; c += blockDim.x) { |
| #if __CUDA_ARCH__ >= 350 |
| dX[blockIdx.x * C + c] = __ldg(dY + n * C + c) * scale; |
| #else |
| dX[blockIdx.x * C + c] = dY[n * C + c] * scale; |
| #endif |
| } |
| } |
| |
| template <typename T, bool kCountIncludePad> |
| __global__ void AveragePool1DBackwardNCHWCUDAKernel( |
| const int X_size, |
| const int Y_size, |
| const int kernel, |
| const int stride, |
| const int pad, |
| const T* dY, |
| T* dX) { |
| const int nc = blockIdx.x; |
| const T* dY_ptr = dY + nc * Y_size; |
| T* dX_ptr = dX + nc * X_size; |
| for (int x = threadIdx.x; x < X_size; x += blockDim.x) { |
| const int w = x + pad; |
| const int l = w < kernel ? 0 : (w - kernel) / stride + 1; |
| const int r = min(w / stride + 1, Y_size); |
| T sum = 0; |
| for (int i = l; i < r; ++i) { |
| if (kCountIncludePad) { |
| #if __CUDA_ARCH__ >= 350 |
| sum += __ldg(dY_ptr + i); |
| #else |
| sum += dY_ptr[i]; |
| #endif |
| } else { |
| const int xx = i * stride - pad; |
| const int xl = max(xx, 0); |
| const int xr = min(xx + kernel, X_size); |
| #if __CUDA_ARCH__ >= 350 |
| sum += __ldg(dY_ptr + i) / static_cast<T>(xr - xl); |
| #else |
| sum += dY_ptr[i] / static_cast<T>(xr - xl); |
| #endif |
| } |
| } |
| if (kCountIncludePad) { |
| dX_ptr[x] = sum / static_cast<T>(kernel); |
| } else { |
| dX_ptr[x] = sum; |
| } |
| } |
| } |
| |
| template <typename T, bool kCountIncludePad> |
| __global__ void AveragePool1DBackwardNHWCCUDAKernel( |
| const int C, |
| const int X_size, |
| const int Y_size, |
| const int kernel, |
| const int stride, |
| const int pad, |
| const T* dY, |
| T* dX) { |
| const int n = blockIdx.x / X_size; |
| const int x = blockIdx.x % X_size; |
| const int w = x + pad; |
| const int l = w < kernel ? 0 : (w - kernel) / stride + 1; |
| const int r = min(w / stride + 1, Y_size); |
| const T scale = T(1) / static_cast<T>(kernel); |
| const T* dY_ptr = dY + n * Y_size * C; |
| T* dX_ptr = dX + n * X_size * C; |
| for (int c = threadIdx.x; c < C; c += blockDim.x) { |
| T sum = 0; |
| for (int i = l; i < r; ++i) { |
| if (kCountIncludePad) { |
| #if __CUDA_ARCH__ >= 350 |
| sum += __ldg(dY_ptr + i * C + c); |
| #else |
| sum += dY_ptr[i * C + c]; |
| #endif |
| } else { |
| const int xx = i * stride - pad; |
| const int xl = max(xx, 0); |
| const int xr = min(xx + kernel, X_size); |
| #if __CUDA_ARCH__ >= 350 |
| sum += __ldg(dY_ptr + i * C + c) / static_cast<T>(xr - xl); |
| #else |
| sum += dY_ptr[i * C + c] / static_cast<T>(xr - xl); |
| #endif |
| } |
| } |
| if (kCountIncludePad) { |
| dX_ptr[x * C + c] = sum * scale; |
| } else { |
| dX_ptr[x * C + c] = sum; |
| } |
| } |
| } |
| |
| template <typename T, bool kCountIncludePad> |
| __global__ void AveragePool2DBackwardNCHWCUDAKernel( |
| const int X_H, |
| const int X_W, |
| const int Y_H, |
| const int Y_W, |
| const int kernel_h, |
| const int kernel_w, |
| const int stride_h, |
| const int stride_w, |
| const int pad_t, |
| const int pad_l, |
| const T* dY, |
| T* dX) { |
| const int X_HxW = X_H * X_W; |
| const int Y_HxW = Y_H * Y_W; |
| const int nc = blockIdx.x / X_H; |
| const int hh = blockIdx.x % X_H; |
| const T* dY_ptr = dY + nc * Y_HxW; |
| T* dX_ptr = dX + nc * X_HxW; |
| const int h = hh + pad_t; |
| const int t = h < kernel_h ? 0 : (h - kernel_h) / stride_h + 1; |
| const int b = min(h / stride_h + 1, Y_H); |
| for (int ww = threadIdx.x; ww < X_W; ww += blockDim.x) { |
| const int w = ww + pad_l; |
| const int l = w < kernel_w ? 0 : (w - kernel_w) / stride_w + 1; |
| const int r = min(w / stride_w + 1, Y_W); |
| T sum = 0; |
| for (int i = t; i < b; ++i) { |
| for (int j = l; j < r; ++j) { |
| if (kCountIncludePad) { |
| #if __CUDA_ARCH__ >= 350 |
| sum += __ldg(dY_ptr + i * Y_W + j); |
| #else |
| sum += dY_ptr[i * Y_W + j]; |
| #endif |
| } else { |
| const int xh = i * stride_h - pad_t; |
| const int xw = j * stride_w - pad_l; |
| const int xt = max(xh, 0); |
| const int xb = min(xh + kernel_h, X_H); |
| const int xl = max(xw, 0); |
| const int xr = min(xw + kernel_w, X_W); |
| #if __CUDA_ARCH__ >= 350 |
| sum += __ldg(dY_ptr + i * Y_W + j) / |
| static_cast<T>((xb - xt) * (xr - xl)); |
| #else |
| sum += dY_ptr[i * Y_W + j] / static_cast<T>((xb - xt) * (xr - xl)); |
| #endif |
| } |
| } |
| } |
| if (kCountIncludePad) { |
| dX_ptr[hh * X_W + ww] = sum / static_cast<T>(kernel_h * kernel_w); |
| } else { |
| dX_ptr[hh * X_W + ww] = sum; |
| } |
| } |
| } |
| |
| template <typename T, bool kCountIncludePad> |
| __global__ void AveragePool2DBackwardNHWCCUDAKernel( |
| const int C, |
| const int X_H, |
| const int X_W, |
| const int Y_H, |
| const int Y_W, |
| const int kernel_h, |
| const int kernel_w, |
| const int stride_h, |
| const int stride_w, |
| const int pad_t, |
| const int pad_l, |
| const T* dY, |
| T* dX) { |
| const int X_HxW = X_H * X_W; |
| const int Y_HxW = Y_H * Y_W; |
| const int n = blockIdx.x / X_HxW; |
| const int x = blockIdx.x % X_HxW; |
| const int h = x / X_W + pad_t; |
| const int w = x % X_W + pad_l; |
| const int t = h < kernel_h ? 0 : (h - kernel_h) / stride_h + 1; |
| const int b = min(h / stride_h + 1, Y_H); |
| const int l = w < kernel_w ? 0 : (w - kernel_w) / stride_w + 1; |
| const int r = min(w / stride_w + 1, Y_W); |
| const T scale = T(1) / static_cast<T>(kernel_h * kernel_w); |
| const T* dY_ptr = dY + n * Y_HxW * C; |
| T* dX_ptr = dX + n * X_HxW * C; |
| for (int c = threadIdx.x; c < C; c += blockDim.x) { |
| T sum = 0; |
| for (int i = t; i < b; ++i) { |
| for (int j = l; j < r; ++j) { |
| if (kCountIncludePad) { |
| #if __CUDA_ARCH__ >= 350 |
| sum += __ldg(dY_ptr + (i * Y_W + j) * C + c); |
| #else |
| sum += dY_ptr[(i * Y_W + j) * C + c]; |
| #endif |
| } else { |
| const int xh = i * stride_h - pad_t; |
| const int xw = j * stride_w - pad_l; |
| const int xt = max(xh, 0); |
| const int xb = min(xh + kernel_h, X_H); |
| const int xl = max(xw, 0); |
| const int xr = min(xw + kernel_w, X_W); |
| #if __CUDA_ARCH__ >= 350 |
| sum += __ldg(dY_ptr + (i * Y_W + j) * C + c) / |
| static_cast<T>((xb - xt) * (xr - xl)); |
| #else |
| sum += dY_ptr[(i * Y_W + j) * C + c] / |
| static_cast<T>((xb - xt) * (xr - xl)); |
| #endif |
| } |
| } |
| } |
| if (kCountIncludePad) { |
| dX_ptr[x * C + c] = sum * scale; |
| } else { |
| dX_ptr[x * C + c] = sum; |
| } |
| } |
| } |
| |
| template <typename T, bool kCountIncludePad> |
| __global__ void AveragePool3DBackwardNCHWCUDAKernel( |
| const int X_D, |
| const int X_H, |
| const int X_W, |
| const int Y_D, |
| const int Y_H, |
| const int Y_W, |
| const int kernel_d, |
| const int kernel_h, |
| const int kernel_w, |
| const int stride_d, |
| const int stride_h, |
| const int stride_w, |
| const int pad_p, |
| const int pad_t, |
| const int pad_l, |
| const T* dY, |
| T* dX) { |
| const int X_HxW = X_D * X_H * X_W; |
| const int Y_HxW = Y_D * Y_H * Y_W; |
| const int xx = blockIdx.x / X_H; |
| const int nc = xx / X_D; |
| const int dd = xx % X_D; |
| const int hh = blockIdx.x % X_H; |
| const T* dY_ptr = dY + nc * Y_HxW; |
| T* dX_ptr = dX + nc * X_HxW; |
| const int d = dd + pad_p; |
| const int h = hh + pad_t; |
| const int p = d < kernel_d ? 0 : (d - kernel_d) / stride_d + 1; |
| const int a = min(d / stride_d + 1, Y_D); |
| const int t = h < kernel_h ? 0 : (h - kernel_h) / stride_h + 1; |
| const int b = min(h / stride_h + 1, Y_H); |
| for (int ww = threadIdx.x; ww < X_W; ww += blockDim.x) { |
| const int w = ww + pad_l; |
| const int l = w < kernel_w ? 0 : (w - kernel_w) / stride_w + 1; |
| const int r = min(w / stride_w + 1, Y_W); |
| T sum = 0; |
| for (int i = p; i < a; ++i) { |
| for (int j = t; j < b; ++j) { |
| for (int k = l; k < r; ++k) { |
| if (kCountIncludePad) { |
| #if __CUDA_ARCH__ >= 350 |
| sum += __ldg(dY_ptr + (i * Y_H + j) * Y_W + k); |
| #else |
| sum += dY_ptr[(i * Y_H + j) * Y_W + k]; |
| #endif |
| } else { |
| const int xd = i * stride_d - pad_p; |
| const int xh = j * stride_h - pad_t; |
| const int xw = k * stride_w - pad_l; |
| const int xp = max(xd, 0); |
| const int xa = min(xd + kernel_d, X_D); |
| const int xt = max(xh, 0); |
| const int xb = min(xh + kernel_h, X_H); |
| const int xl = max(xw, 0); |
| const int xr = min(xw + kernel_w, X_W); |
| #if __CUDA_ARCH__ >= 350 |
| sum += __ldg(dY_ptr + (i * Y_H + j) * Y_W + k) / |
| static_cast<T>((xa - xp) * (xb - xt) * (xr - xl)); |
| #else |
| sum += dY_ptr[(i * Y_H + j) * Y_W + k] / |
| static_cast<T>((xa - xp) * (xb - xt) * (xr - xl)); |
| #endif |
| } |
| } |
| } |
| } |
| if (kCountIncludePad) { |
| dX_ptr[(dd * X_H + hh) * X_W + ww] = |
| sum / static_cast<T>(kernel_d * kernel_h * kernel_w); |
| } else { |
| dX_ptr[(dd * X_H + hh) * X_W + ww] = sum; |
| } |
| } |
| } |
| |
| template <typename T, bool kCountIncludePad> |
| __global__ void AveragePool3DBackwardNHWCCUDAKernel( |
| const int C, |
| const int X_D, |
| const int X_H, |
| const int X_W, |
| const int Y_D, |
| const int Y_H, |
| const int Y_W, |
| const int kernel_d, |
| const int kernel_h, |
| const int kernel_w, |
| const int stride_d, |
| const int stride_h, |
| const int stride_w, |
| const int pad_p, |
| const int pad_t, |
| const int pad_l, |
| const T* dY, |
| T* dX) { |
| const int X_HxW = X_D * X_H * X_W; |
| const int Y_HxW = Y_D * Y_H * Y_W; |
| const int n = blockIdx.x / X_HxW; |
| const int x = blockIdx.x % X_HxW; |
| const int xx = x / X_W; |
| const int d = xx / X_H + pad_p; |
| const int h = xx % X_H + pad_t; |
| const int w = x % X_W + pad_l; |
| const int p = d < kernel_d ? 0 : (d - kernel_d) / stride_d + 1; |
| const int a = min(d / stride_d + 1, Y_D); |
| const int t = h < kernel_h ? 0 : (h - kernel_h) / stride_h + 1; |
| const int b = min(h / stride_h + 1, Y_H); |
| const int l = w < kernel_w ? 0 : (w - kernel_w) / stride_w + 1; |
| const int r = min(w / stride_w + 1, Y_W); |
| const T scale = T(1) / static_cast<T>(kernel_d * kernel_h * kernel_w); |
| const T* dY_ptr = dY + n * Y_HxW * C; |
| T* dX_ptr = dX + n * X_HxW * C; |
| for (int c = threadIdx.x; c < C; c += blockDim.x) { |
| T sum = 0; |
| for (int i = p; i < a; ++i) { |
| for (int j = t; j < b; ++j) { |
| for (int k = l; k < r; ++k) { |
| if (kCountIncludePad) { |
| #if __CUDA_ARCH__ >= 350 |
| sum += __ldg(dY_ptr + ((i * Y_H + j) * Y_W + k) * C + c); |
| #else |
| sum += dY_ptr[((i * Y_H + j) * Y_W + k) * C + c]; |
| #endif |
| } else { |
| const int xd = i * stride_d - pad_p; |
| const int xh = j * stride_h - pad_t; |
| const int xw = k * stride_w - pad_l; |
| const int xp = max(xd, 0); |
| const int xa = min(xd + kernel_d, X_D); |
| const int xt = max(xh, 0); |
| const int xb = min(xh + kernel_h, X_H); |
| const int xl = max(xw, 0); |
| const int xr = min(xw + kernel_w, X_W); |
| #if __CUDA_ARCH__ >= 350 |
| sum += __ldg(dY_ptr + ((i * Y_H + j) * Y_W + k) * C + c) / |
| static_cast<T>((xa - xp) * (xb - xt) * (xr - xl)); |
| #else |
| sum += dY_ptr[((i * Y_H + j) * Y_W + k) * C + c] / |
| static_cast<T>((xa - xp) * (xb - xt) * (xr - xl)); |
| #endif |
| } |
| } |
| } |
| } |
| if (kCountIncludePad) { |
| dX_ptr[x * C + c] = sum * scale; |
| } else { |
| dX_ptr[x * C + c] = sum; |
| } |
| } |
| } |
| |
| } // namespace |
| |
| template <> |
| template <> |
| bool AveragePoolFunctor<CUDAContext>:: |
| GlobalPoolingForward<float, StorageOrder::NCHW>( |
| const int N, |
| const int C, |
| const int HxW, |
| const float* X, |
| float* Y, |
| CUDAContext* context) const { |
| const std::array<int, 2> X_dims = {N * C, HxW}; |
| const std::array<int, 2> Y_dims = {N * C, 1}; |
| math::ReduceMean<float, CUDAContext>( |
| 2, X_dims.data(), Y_dims.data(), 1.0f, X, Y, context); |
| return true; |
| } |
| |
| template <> |
| template <> |
| bool AveragePoolFunctor<CUDAContext>:: |
| GlobalPoolingForward<float, StorageOrder::NHWC>( |
| const int N, |
| const int C, |
| const int HxW, |
| const float* X, |
| float* Y, |
| CUDAContext* context) const { |
| if (ones.numel() != HxW) { |
| ones.Resize(HxW); |
| math::Set<float, CUDAContext>( |
| HxW, 1.0f, ones.mutable_data<float>(), context); |
| } |
| math::GemmStridedBatched<float, CUDAContext>( |
| CblasTrans, |
| CblasNoTrans, |
| N, |
| C, |
| 1, |
| HxW, |
| 1.0f / static_cast<float>(HxW), |
| X, |
| HxW * C, |
| ones.data<float>(), |
| 0, |
| 0.0f, |
| Y, |
| C, |
| context); |
| return true; |
| } |
| |
| template <> |
| template <> |
| bool AveragePoolFunctor<CUDAContext>::Forward<float, StorageOrder::NCHW>( |
| const int N, |
| const int C, |
| const std::vector<int>& X_dims, |
| const std::vector<int>& Y_dims, |
| const std::vector<int>& kernel, |
| const std::vector<int>& /* dilation */, |
| const std::vector<int>& stride, |
| const std::vector<int>& pads, |
| const float* X, |
| float* Y, |
| CUDAContext* context) const { |
| const int ndim = X_dims.size(); |
| switch (ndim) { |
| case 1: { |
| const int num_blocks = N * C; |
| AveragePool1DForwardNCHWCUDAKernel<float> |
| <<<num_blocks, CAFFE_CUDA_NUM_THREADS, 0, context->cuda_stream()>>>( |
| X_dims[0], |
| Y_dims[0], |
| kernel[0], |
| stride[0], |
| pads[0], |
| count_include_pad, |
| X, |
| Y); |
| return true; |
| } |
| case 2: { |
| const int num_blocks = N * C * Y_dims[0]; |
| AveragePool2DForwardNCHWCUDAKernel<float> |
| <<<num_blocks, CAFFE_CUDA_NUM_THREADS, 0, context->cuda_stream()>>>( |
| X_dims[0], |
| X_dims[1], |
| Y_dims[0], |
| Y_dims[1], |
| kernel[0], |
| kernel[1], |
| stride[0], |
| stride[1], |
| pads[0], |
| pads[1], |
| count_include_pad, |
| X, |
| Y); |
| return true; |
| } |
| case 3: { |
| const int num_blocks = N * C * Y_dims[0] * Y_dims[1]; |
| AveragePool3DForwardNCHWCUDAKernel<float> |
| <<<num_blocks, CAFFE_CUDA_NUM_THREADS, 0, context->cuda_stream()>>>( |
| X_dims[0], |
| X_dims[1], |
| X_dims[2], |
| Y_dims[0], |
| Y_dims[1], |
| Y_dims[2], |
| kernel[0], |
| kernel[1], |
| kernel[2], |
| stride[0], |
| stride[1], |
| stride[2], |
| pads[0], |
| pads[1], |
| pads[2], |
| count_include_pad, |
| X, |
| Y); |
| return true; |
| } |
| default: { |
| CAFFE_THROW("Unsupported pooling dim: ", ndim); |
| return false; |
| } |
| } |
| } |
| |
| template <> |
| template <> |
| bool AveragePoolFunctor<CUDAContext>::Forward<float, StorageOrder::NHWC>( |
| const int N, |
| const int C, |
| const std::vector<int>& X_dims, |
| const std::vector<int>& Y_dims, |
| const std::vector<int>& kernel, |
| const std::vector<int>& /* dilation */, |
| const std::vector<int>& stride, |
| const std::vector<int>& pads, |
| const float* X, |
| float* Y, |
| CUDAContext* context) const { |
| // Each CUDA block handles one point, one thread per channel. |
| const int ndim = X_dims.size(); |
| const int Y_HxW = std::accumulate( |
| Y_dims.cbegin(), Y_dims.cend(), 1, std::multiplies<int>()); |
| switch (ndim) { |
| case 1: { |
| AveragePool1DForwardNHWCCUDAKernel<float> |
| <<<N * Y_HxW, CAFFE_CUDA_NUM_THREADS, 0, context->cuda_stream()>>>( |
| C, |
| X_dims[0], |
| Y_dims[0], |
| kernel[0], |
| stride[0], |
| pads[0], |
| count_include_pad, |
| X, |
| Y); |
| return true; |
| } |
| case 2: { |
| AveragePool2DForwardNHWCCUDAKernel<float> |
| <<<N * Y_HxW, CAFFE_CUDA_NUM_THREADS, 0, context->cuda_stream()>>>( |
| C, |
| X_dims[0], |
| X_dims[1], |
| Y_dims[0], |
| Y_dims[1], |
| kernel[0], |
| kernel[1], |
| stride[0], |
| stride[1], |
| pads[0], |
| pads[1], |
| count_include_pad, |
| X, |
| Y); |
| return true; |
| } |
| case 3: { |
| AveragePool3DForwardNHWCCUDAKernel<float> |
| <<<N * Y_HxW, CAFFE_CUDA_NUM_THREADS, 0, context->cuda_stream()>>>( |
| C, |
| X_dims[0], |
| X_dims[1], |
| X_dims[2], |
| Y_dims[0], |
| Y_dims[1], |
| Y_dims[2], |
| kernel[0], |
| kernel[1], |
| kernel[2], |
| stride[0], |
| stride[1], |
| stride[2], |
| pads[0], |
| pads[1], |
| pads[2], |
| count_include_pad, |
| X, |
| Y); |
| return true; |
| } |
| default: { |
| CAFFE_THROW("Unsupported pooling dim: ", ndim); |
| return false; |
| } |
| } |
| } |
| |
| template <> |
| template <> |
| bool AveragePoolFunctor<CUDAContext>:: |
| GlobalPoolingBackward<float, StorageOrder::NCHW>( |
| const int N, |
| const int C, |
| const int HxW, |
| const float* dY, |
| const float* /* X */, |
| const float* /* Y */, |
| float* dX, |
| CUDAContext* context) const { |
| const float scale = 1.0f / static_cast<float>(HxW); |
| const int K = (HxW + CAFFE_CUDA_NUM_THREADS - 1) / CAFFE_CUDA_NUM_THREADS; |
| GlobalAveragePoolingBackwardNCHWCUDAKernel<float> |
| <<<N * C * K, CAFFE_CUDA_NUM_THREADS, 0, context->cuda_stream()>>>( |
| K, HxW, scale, dY, dX); |
| return true; |
| } |
| |
| template <> |
| template <> |
| bool AveragePoolFunctor<CUDAContext>:: |
| GlobalPoolingBackward<float, StorageOrder::NHWC>( |
| const int N, |
| const int C, |
| const int HxW, |
| const float* dY, |
| const float* /* X */, |
| const float* /* Y */, |
| float* dX, |
| CUDAContext* context) const { |
| const float scale = 1.0f / static_cast<float>(HxW); |
| GlobalAveragePoolingBackwardNHWCCUDAKernel<float> |
| <<<N * HxW, CAFFE_CUDA_NUM_THREADS, 0, context->cuda_stream()>>>( |
| C, HxW, scale, dY, dX); |
| return true; |
| } |
| |
| #define DISPATCH_KERNEL_FUNCTION_BY_BOOL_WITH_TYPE_1( \ |
| cond, Func, T, num_blocks, threads_per_block, cuda_stream, ...) \ |
| do { \ |
| if (cond) { \ |
| Func<T, true> \ |
| <<<num_blocks, threads_per_block, 0, cuda_stream>>>(__VA_ARGS__); \ |
| } else { \ |
| Func<T, false> \ |
| <<<num_blocks, threads_per_block, 0, cuda_stream>>>(__VA_ARGS__); \ |
| } \ |
| } while (false) |
| |
| template <> |
| template <> |
| bool AveragePoolFunctor<CUDAContext>::Backward<float, StorageOrder::NCHW>( |
| const int N, |
| const int C, |
| const std::vector<int>& X_dims, |
| const std::vector<int>& Y_dims, |
| const std::vector<int>& kernel, |
| const std::vector<int>& /* dilation */, |
| const std::vector<int>& stride, |
| const std::vector<int>& pads, |
| const float* dY, |
| const float* /* X */, |
| const float* /* Y */, |
| float* dX, |
| CUDAContext* context) const { |
| const int ndim = X_dims.size(); |
| switch (ndim) { |
| case 1: { |
| const int num_blocks = N * C; |
| DISPATCH_KERNEL_FUNCTION_BY_BOOL_WITH_TYPE_1( |
| count_include_pad, |
| AveragePool1DBackwardNCHWCUDAKernel, |
| float, |
| num_blocks, |
| CAFFE_CUDA_NUM_THREADS, |
| context->cuda_stream(), |
| X_dims[0], |
| Y_dims[0], |
| kernel[0], |
| stride[0], |
| pads[0], |
| dY, |
| dX); |
| return true; |
| } |
| case 2: { |
| const int num_blocks = N * C * X_dims[0]; |
| DISPATCH_KERNEL_FUNCTION_BY_BOOL_WITH_TYPE_1( |
| count_include_pad, |
| AveragePool2DBackwardNCHWCUDAKernel, |
| float, |
| num_blocks, |
| CAFFE_CUDA_NUM_THREADS, |
| context->cuda_stream(), |
| X_dims[0], |
| X_dims[1], |
| Y_dims[0], |
| Y_dims[1], |
| kernel[0], |
| kernel[1], |
| stride[0], |
| stride[1], |
| pads[0], |
| pads[1], |
| dY, |
| dX); |
| return true; |
| } |
| case 3: { |
| const int num_blocks = N * C * X_dims[0] * X_dims[1]; |
| DISPATCH_KERNEL_FUNCTION_BY_BOOL_WITH_TYPE_1( |
| count_include_pad, |
| AveragePool3DBackwardNCHWCUDAKernel, |
| float, |
| num_blocks, |
| CAFFE_CUDA_NUM_THREADS, |
| context->cuda_stream(), |
| X_dims[0], |
| X_dims[1], |
| X_dims[2], |
| Y_dims[0], |
| Y_dims[1], |
| Y_dims[2], |
| kernel[0], |
| kernel[1], |
| kernel[2], |
| stride[0], |
| stride[1], |
| stride[2], |
| pads[0], |
| pads[1], |
| pads[2], |
| dY, |
| dX); |
| return true; |
| } |
| default: { |
| CAFFE_THROW("Unsupported pooling dim: ", ndim); |
| return false; |
| } |
| } |
| } |
| |
| template <> |
| template <> |
| bool AveragePoolFunctor<CUDAContext>::Backward<float, StorageOrder::NHWC>( |
| const int N, |
| const int C, |
| const std::vector<int>& X_dims, |
| const std::vector<int>& Y_dims, |
| const std::vector<int>& kernel, |
| const std::vector<int>& /* dilation */, |
| const std::vector<int>& stride, |
| const std::vector<int>& pads, |
| const float* dY, |
| const float* /* X */, |
| const float* /* Y */, |
| float* dX, |
| CUDAContext* context) const { |
| const int ndim = X_dims.size(); |
| const int X_HxW = std::accumulate( |
| X_dims.cbegin(), X_dims.cend(), 1, std::multiplies<int>()); |
| const int num_blocks = N * X_HxW; |
| switch (ndim) { |
| case 1: { |
| DISPATCH_KERNEL_FUNCTION_BY_BOOL_WITH_TYPE_1( |
| count_include_pad, |
| AveragePool1DBackwardNHWCCUDAKernel, |
| float, |
| num_blocks, |
| CAFFE_CUDA_NUM_THREADS, |
| context->cuda_stream(), |
| C, |
| X_dims[0], |
| Y_dims[0], |
| kernel[0], |
| stride[0], |
| pads[0], |
| dY, |
| dX); |
| return true; |
| } |
| case 2: { |
| DISPATCH_KERNEL_FUNCTION_BY_BOOL_WITH_TYPE_1( |
| count_include_pad, |
| AveragePool2DBackwardNHWCCUDAKernel, |
| float, |
| num_blocks, |
| CAFFE_CUDA_NUM_THREADS, |
| context->cuda_stream(), |
| C, |
| X_dims[0], |
| X_dims[1], |
| Y_dims[0], |
| Y_dims[1], |
| kernel[0], |
| kernel[1], |
| stride[0], |
| stride[1], |
| pads[0], |
| pads[1], |
| dY, |
| dX); |
| return true; |
| } |
| case 3: { |
| DISPATCH_KERNEL_FUNCTION_BY_BOOL_WITH_TYPE_1( |
| count_include_pad, |
| AveragePool3DBackwardNHWCCUDAKernel, |
| float, |
| num_blocks, |
| CAFFE_CUDA_NUM_THREADS, |
| context->cuda_stream(), |
| C, |
| X_dims[0], |
| X_dims[1], |
| X_dims[2], |
| Y_dims[0], |
| Y_dims[1], |
| Y_dims[2], |
| kernel[0], |
| kernel[1], |
| kernel[2], |
| stride[0], |
| stride[1], |
| stride[2], |
| pads[0], |
| pads[1], |
| pads[2], |
| dY, |
| dX); |
| return true; |
| } |
| default: { |
| CAFFE_THROW("Unsupported pooling dim: ", ndim); |
| return false; |
| } |
| } |
| } |
| |
| #undef DISPATCH_KERNEL_FUNCTION_BY_BOOL_WITH_TYPE_1 |
| |
| namespace { |
| |
| template <typename T> |
| __global__ void MaxPool1DForwardNCHWCUDAKernel( |
| const int X_size, |
| const int Y_size, |
| const int kernel, |
| const int stride, |
| const int pad, |
| const T* X, |
| T* Y) { |
| const int nc = blockIdx.x; |
| const T* X_ptr = X + nc * X_size; |
| T* Y_ptr = Y + nc * Y_size; |
| for (int y = threadIdx.x; y < Y_size; y += blockDim.x) { |
| const int x = y * stride; |
| const int l = max(x - pad, 0); |
| const int r = min(x - pad + kernel, X_size); |
| T val = std::numeric_limits<T>::lowest(); |
| for (int i = l; i < r; ++i) { |
| #if __CUDA_ARCH__ >= 350 |
| val = max(val, __ldg(X_ptr + i)); |
| #else |
| val = max(val, X_ptr[i]); |
| #endif |
| } |
| Y_ptr[y] = val; |
| } |
| } |
| |
| template <typename T> |
| __global__ void MaxPool1DForwardNHWCCUDAKernel( |
| const int C, |
| const int X_size, |
| const int Y_size, |
| const int kernel, |
| const int stride, |
| const int pad, |
| const T* X, |
| T* Y) { |
| const int n = blockIdx.x / Y_size; |
| const int y = blockIdx.x % Y_size; |
| const int x = y * stride; |
| const int l = max(x - pad, 0); |
| const int r = min(x - pad + kernel, X_size); |
| const T* X_ptr = X + n * X_size * C; |
| T* Y_ptr = Y + n * Y_size * C; |
| for (int c = threadIdx.x; c < C; c += blockDim.x) { |
| T val = std::numeric_limits<T>::lowest(); |
| for (int i = l; i < r; ++i) { |
| #if __CUDA_ARCH__ >= 350 |
| val = max(val, __ldg(X_ptr + i * C + c)); |
| #else |
| val = max(val, X_ptr[i * C + c]); |
| #endif |
| } |
| Y_ptr[y * C + c] = val; |
| } |
| } |
| |
| template <typename T> |
| __global__ void MaxPool2DForwardNCHWCUDAKernel( |
| const int X_H, |
| const int X_W, |
| const int Y_H, |
| const int Y_W, |
| const int kernel_h, |
| const int kernel_w, |
| const int stride_h, |
| const int stride_w, |
| const int pad_t, |
| const int pad_l, |
| const T* X, |
| T* Y) { |
| const int X_HxW = X_H * X_W; |
| const int Y_HxW = Y_H * Y_W; |
| const int nc = blockIdx.x / Y_H; |
| const int yh = blockIdx.x % Y_H; |
| const T* X_ptr = X + nc * X_HxW; |
| T* Y_ptr = Y + nc * Y_HxW; |
| const int xh = yh * stride_h; |
| const int t = max(xh - pad_t, 0); |
| const int b = min(xh - pad_t + kernel_h, X_H); |
| for (int yw = threadIdx.x; yw < Y_W; yw += blockDim.x) { |
| const int xw = yw * stride_w; |
| const int l = max(xw - pad_l, 0); |
| const int r = min(xw - pad_l + kernel_w, X_W); |
| T val = std::numeric_limits<T>::lowest(); |
| for (int i = t; i < b; ++i) { |
| for (int j = l; j < r; ++j) { |
| #if __CUDA_ARCH__ >= 350 |
| val = max(val, __ldg(X_ptr + i * X_W + j)); |
| #else |
| val = max(val, X_ptr[i * X_W + j]); |
| #endif |
| } |
| } |
| Y_ptr[yh * Y_W + yw] = val; |
| } |
| } |
| |
| template <typename T> |
| __global__ void MaxPool2DForwardNHWCCUDAKernel( |
| const int C, |
| const int X_H, |
| const int X_W, |
| const int Y_H, |
| const int Y_W, |
| const int kernel_h, |
| const int kernel_w, |
| const int stride_h, |
| const int stride_w, |
| const int pad_t, |
| const int pad_l, |
| const T* X, |
| T* Y) { |
| const int X_HxW = X_H * X_W; |
| const int Y_HxW = Y_H * Y_W; |
| const int n = blockIdx.x / Y_HxW; |
| const int y = blockIdx.x % Y_HxW; |
| const int yh = y / Y_W; |
| const int yw = y % Y_W; |
| const int xh = yh * stride_h; |
| const int xw = yw * stride_w; |
| const int t = max(xh - pad_t, 0); |
| const int b = min(xh - pad_t + kernel_h, X_H); |
| const int l = max(xw - pad_l, 0); |
| const int r = min(xw - pad_l + kernel_w, X_W); |
| const T* X_ptr = X + n * X_HxW * C; |
| T* Y_ptr = Y + n * Y_HxW * C; |
| for (int c = threadIdx.x; c < C; c += blockDim.x) { |
| T val = std::numeric_limits<T>::lowest(); |
| for (int i = t; i < b; ++i) { |
| for (int j = l; j < r; ++j) { |
| #if __CUDA_ARCH__ >= 350 |
| val = max(val, __ldg(X_ptr + (i * X_W + j) * C + c)); |
| #else |
| val = max(val, X_ptr[(i * X_W + j) * C + c]); |
| #endif |
| } |
| } |
| Y_ptr[y * C + c] = val; |
| } |
| } |
| |
| template <typename T> |
| __global__ void MaxPool3DForwardNCHWCUDAKernel( |
| const int X_D, |
| const int X_H, |
| const int X_W, |
| const int Y_D, |
| const int Y_H, |
| const int Y_W, |
| const int kernel_d, |
| const int kernel_h, |
| const int kernel_w, |
| const int stride_d, |
| const int stride_h, |
| const int stride_w, |
| const int pad_p, |
| const int pad_t, |
| const int pad_l, |
| const T* X, |
| T* Y) { |
| const int X_HxW = X_D * X_H * X_W; |
| const int Y_HxW = Y_D * Y_H * Y_W; |
| const int yy = blockIdx.x / Y_H; |
| const int nc = yy / Y_D; |
| const int yd = yy % Y_D; |
| const int yh = blockIdx.x % Y_H; |
| const T* X_ptr = X + nc * X_HxW; |
| T* Y_ptr = Y + nc * Y_HxW; |
| const int xd = yd * stride_d; |
| const int xh = yh * stride_h; |
| const int p = max(xd - pad_p, 0); |
| const int a = min(xd - pad_p + kernel_d, X_D); |
| const int t = max(xh - pad_t, 0); |
| const int b = min(xh - pad_t + kernel_h, X_H); |
| for (int yw = threadIdx.x; yw < Y_W; yw += blockDim.x) { |
| const int xw = yw * stride_w; |
| const int l = max(xw - pad_l, 0); |
| const int r = min(xw - pad_l + kernel_w, X_W); |
| T val = std::numeric_limits<T>::lowest(); |
| for (int i = p; i < a; ++i) { |
| for (int j = t; j < b; ++j) { |
| for (int k = l; k < r; ++k) { |
| #if __CUDA_ARCH__ >= 350 |
| val = max(val, __ldg(X_ptr + (i * X_H + j) * X_W + k)); |
| #else |
| val = max(val, X_ptr[(i * X_H + j) * X_W + k]); |
| #endif |
| } |
| } |
| } |
| Y_ptr[(yd * Y_H + yh) * Y_W + yw] = val; |
| } |
| } |
| |
| template <typename T> |
| __global__ void MaxPool3DForwardNHWCCUDAKernel( |
| const int C, |
| const int X_D, |
| const int X_H, |
| const int X_W, |
| const int Y_D, |
| const int Y_H, |
| const int Y_W, |
| const int kernel_d, |
| const int kernel_h, |
| const int kernel_w, |
| const int stride_d, |
| const int stride_h, |
| const int stride_w, |
| const int pad_p, |
| const int pad_t, |
| const int pad_l, |
| const T* X, |
| T* Y) { |
| const int X_HxW = X_D * X_H * X_W; |
| const int Y_HxW = Y_D * Y_H * Y_W; |
| const int n = blockIdx.x / Y_HxW; |
| const int y = blockIdx.x % Y_HxW; |
| const int yy = y / Y_W; |
| const int yw = y % Y_W; |
| const int yh = yy % Y_H; |
| const int yd = yy / Y_H; |
| const int xd = yd * stride_d; |
| const int xh = yh * stride_h; |
| const int xw = yw * stride_w; |
| const int p = max(xd - pad_p, 0); |
| const int a = min(xd - pad_p + kernel_d, X_D); |
| const int t = max(xh - pad_t, 0); |
| const int b = min(xh - pad_t + kernel_h, X_H); |
| const int l = max(xw - pad_l, 0); |
| const int r = min(xw - pad_l + kernel_w, X_W); |
| const T* X_ptr = X + n * X_HxW * C; |
| T* Y_ptr = Y + n * Y_HxW * C; |
| for (int c = threadIdx.x; c < C; c += blockDim.x) { |
| T val = std::numeric_limits<T>::lowest(); |
| for (int i = p; i < a; ++i) { |
| for (int j = t; j < b; ++j) { |
| for (int k = l; k < r; ++k) { |
| #if __CUDA_ARCH__ >= 350 |
| val = max(val, __ldg(X_ptr + ((i * X_H + j) * X_W + k) * C + c)); |
| #else |
| val = max(val, X_ptr[((i * X_H + j) * X_W + k) * C + c]); |
| #endif |
| } |
| } |
| } |
| Y_ptr[y * C + c] = val; |
| } |
| } |
| |
| template <typename T> |
| __global__ void GlobalMaxPoolingBackwardNCHWCUDAKernel( |
| const int K, |
| const int HxW, |
| const T* dY, |
| const T* X, |
| const T* Y, |
| T* dX) { |
| const int nc = blockIdx.x / K; |
| const int block = blockIdx.x % K; |
| const int x = threadIdx.x + block * CAFFE_CUDA_NUM_THREADS; |
| if (x < HxW) { |
| #if __CUDA_ARCH__ >= 350 |
| dX[nc * HxW + x] = |
| (__ldg(X + nc * HxW + x) == __ldg(Y + nc)) ? __ldg(dY + nc) : T(0); |
| #else |
| dX[nc * HxW + x] = (X[nc * HxW + x] == Y[nc]) ? dY[nc] : T(0); |
| #endif |
| } |
| } |
| |
| template <typename T> |
| __global__ void GlobalMaxPoolingBackwardNHWCCUDAKernel( |
| const int C, |
| const int HxW, |
| const T* dY, |
| const T* X, |
| const T* Y, |
| T* dX) { |
| const int n = blockIdx.x / HxW; |
| for (int c = threadIdx.x; c < C; c += blockDim.x) { |
| #if __CUDA_ARCH__ >= 350 |
| dX[blockIdx.x * C + c] = |
| (__ldg(X + blockIdx.x * C + c) == __ldg(Y + n * C + c)) |
| ? __ldg(dY + n * C + c) |
| : T(0); |
| #else |
| dX[blockIdx.x * C + c] = |
| (X[blockIdx.x * C + c] == Y[n * C + c]) ? dY[n * C + c] : T(0); |
| #endif |
| } |
| } |
| |
| template <typename T> |
| __global__ void MaxPool1DBackwardNCHWCUDAKernel( |
| const int X_size, |
| const int Y_size, |
| const int kernel, |
| const int stride, |
| const int pad, |
| const T* dY, |
| const T* X, |
| const T* Y, |
| T* dX) { |
| const int nc = blockIdx.x; |
| const T* dY_ptr = dY + nc * Y_size; |
| const T* X_ptr = X + nc * X_size; |
| const T* Y_ptr = Y + nc * Y_size; |
| T* dX_ptr = dX + nc * X_size; |
| for (int x = threadIdx.x; x < X_size; x += blockDim.x) { |
| const int w = x + pad; |
| const int l = w < kernel ? 0 : (w - kernel) / stride + 1; |
| const int r = min(w / stride + 1, Y_size); |
| T sum = 0; |
| for (int i = l; i < r; ++i) { |
| #if __CUDA_ARCH__ >= 350 |
| if (__ldg(X_ptr + x) == __ldg(Y_ptr + i)) { |
| sum += __ldg(dY_ptr + i); |
| } |
| #else |
| if (X_ptr[x] == Y_ptr[i]) { |
| sum += dY_ptr[i]; |
| } |
| #endif |
| } |
| dX_ptr[x] = sum; |
| } |
| } |
| |
| template <typename T> |
| __global__ void MaxPool1DBackwardNHWCCUDAKernel( |
| const int C, |
| const int X_size, |
| const int Y_size, |
| const int kernel, |
| const int stride, |
| const int pad, |
| const T* dY, |
| const T* X, |
| const T* Y, |
| T* dX) { |
| const int n = blockIdx.x / X_size; |
| const int x = blockIdx.x % X_size; |
| const int w = x + pad; |
| const int l = w < kernel ? 0 : (w - kernel) / stride + 1; |
| const int r = min(w / stride + 1, Y_size); |
| const T* dY_ptr = dY + n * Y_size * C; |
| const T* X_ptr = X + n * X_size * C; |
| const T* Y_ptr = Y + n * Y_size * C; |
| T* dX_ptr = dX + n * X_size * C; |
| for (int c = threadIdx.x; c < C; c += blockDim.x) { |
| T sum = 0; |
| for (int i = l; i < r; ++i) { |
| #if __CUDA_ARCH__ >= 350 |
| if (__ldg(X_ptr + x * C + c) == __ldg(Y_ptr + i * C + c)) { |
| sum += __ldg(dY_ptr + i * C + c); |
| } |
| #else |
| if (X_ptr[x * C + c] == Y_ptr[i * C + c]) { |
| sum += dY_ptr[i * C + c]; |
| } |
| #endif |
| } |
| dX_ptr[x * C + c] = sum; |
| } |
| } |
| |
| template <typename T> |
| __global__ void MaxPool2DBackwardNCHWCUDAKernel( |
| const int X_H, |
| const int X_W, |
| const int Y_H, |
| const int Y_W, |
| const int kernel_h, |
| const int kernel_w, |
| const int stride_h, |
| const int stride_w, |
| const int pad_t, |
| const int pad_l, |
| const T* dY, |
| const T* X, |
| const T* Y, |
| T* dX) { |
| const int X_HxW = X_H * X_W; |
| const int Y_HxW = Y_H * Y_W; |
| const int nc = blockIdx.x / X_H; |
| const int xh = blockIdx.x % X_H; |
| const T* dY_ptr = dY + nc * Y_HxW; |
| const T* X_ptr = X + nc * X_HxW; |
| const T* Y_ptr = Y + nc * Y_HxW; |
| T* dX_ptr = dX + nc * X_HxW; |
| const int h = xh + pad_t; |
| const int t = h < kernel_h ? 0 : (h - kernel_h) / stride_h + 1; |
| const int b = min(h / stride_h + 1, Y_H); |
| for (int xw = threadIdx.x; xw < X_W; xw += blockDim.x) { |
| const int w = xw + pad_l; |
| const int l = w < kernel_w ? 0 : (w - kernel_w) / stride_w + 1; |
| const int r = min(w / stride_w + 1, Y_W); |
| const int x = xh * X_W + xw; |
| T sum = 0; |
| for (int i = t; i < b; ++i) { |
| for (int j = l; j < r; ++j) { |
| const int y = i * Y_W + j; |
| #if __CUDA_ARCH__ >= 350 |
| if (__ldg(X_ptr + x) == __ldg(Y_ptr + y)) { |
| sum += __ldg(dY_ptr + y); |
| } |
| #else |
| if (X_ptr[x] == Y_ptr[y]) { |
| sum += dY_ptr[y]; |
| } |
| #endif |
| } |
| } |
| dX_ptr[x] = sum; |
| } |
| } |
| |
| template <typename T> |
| __global__ void MaxPool2DBackwardNHWCCUDAKernel( |
| const int C, |
| const int X_H, |
| const int X_W, |
| const int Y_H, |
| const int Y_W, |
| const int kernel_h, |
| const int kernel_w, |
| const int stride_h, |
| const int stride_w, |
| const int pad_t, |
| const int pad_l, |
| const T* dY, |
| const T* X, |
| const T* Y, |
| T* dX) { |
| const int X_HxW = X_H * X_W; |
| const int Y_HxW = Y_H * Y_W; |
| const int n = blockIdx.x / X_HxW; |
| const int x = blockIdx.x % X_HxW; |
| const int h = x / X_W + pad_t; |
| const int w = x % X_W + pad_l; |
| const int t = h < kernel_h ? 0 : (h - kernel_h) / stride_h + 1; |
| const int b = min(h / stride_h + 1, Y_H); |
| const int l = w < kernel_w ? 0 : (w - kernel_w) / stride_w + 1; |
| const int r = min(w / stride_w + 1, Y_W); |
| const T* dY_ptr = dY + n * Y_HxW * C; |
| const T* X_ptr = X + n * X_HxW * C; |
| const T* Y_ptr = Y + n * Y_HxW * C; |
| T* dX_ptr = dX + n * X_HxW * C; |
| for (int c = threadIdx.x; c < C; c += blockDim.x) { |
| T sum = 0; |
| for (int i = t; i < b; ++i) { |
| for (int j = l; j < r; ++j) { |
| const int y = i * Y_W + j; |
| #if __CUDA_ARCH__ >= 350 |
| if (__ldg(X_ptr + x * C + c) == __ldg(Y_ptr + y * C + c)) { |
| sum += __ldg(dY_ptr + y * C + c); |
| } |
| #else |
| if (X_ptr[x * C + c] == Y_ptr[y * C + c]) { |
| sum += dY_ptr[y * C + c]; |
| } |
| #endif |
| } |
| } |
| dX_ptr[x * C + c] = sum; |
| } |
| } |
| |
| template <typename T> |
| __global__ void MaxPool3DBackwardNCHWCUDAKernel( |
| const int X_D, |
| const int X_H, |
| const int X_W, |
| const int Y_D, |
| const int Y_H, |
| const int Y_W, |
| const int kernel_d, |
| const int kernel_h, |
| const int kernel_w, |
| const int stride_d, |
| const int stride_h, |
| const int stride_w, |
| const int pad_p, |
| const int pad_t, |
| const int pad_l, |
| const T* dY, |
| const T* X, |
| const T* Y, |
| T* dX) { |
| const int X_HxW = X_D * X_H * X_W; |
| const int Y_HxW = Y_D * Y_H * Y_W; |
| const int xx = blockIdx.x / X_H; |
| const int nc = xx / X_D; |
| const int xd = xx % X_D; |
| const int xh = blockIdx.x % X_H; |
| const T* dY_ptr = dY + nc * Y_HxW; |
| const T* X_ptr = X + nc * X_HxW; |
| const T* Y_ptr = Y + nc * Y_HxW; |
| T* dX_ptr = dX + nc * X_HxW; |
| const int d = xd + pad_p; |
| const int h = xh + pad_t; |
| const int p = d < kernel_d ? 0 : (d - kernel_d) / stride_d + 1; |
| const int a = min(d / stride_d + 1, Y_D); |
| const int t = h < kernel_h ? 0 : (h - kernel_h) / stride_h + 1; |
| const int b = min(h / stride_h + 1, Y_H); |
| for (int xw = threadIdx.x; xw < X_W; xw += blockDim.x) { |
| const int w = xw + pad_l; |
| const int l = w < kernel_w ? 0 : (w - kernel_w) / stride_w + 1; |
| const int r = min(w / stride_w + 1, Y_W); |
| const int x = (xd * X_H + xh) * X_W + xw; |
| T sum = 0; |
| for (int i = p; i < a; ++i) { |
| for (int j = t; j < b; ++j) { |
| for (int k = l; k < r; ++k) { |
| const int y = (i * Y_H + j) * Y_W + k; |
| #if __CUDA_ARCH__ >= 350 |
| if (__ldg(X_ptr + x) == __ldg(Y_ptr + y)) { |
| sum += __ldg(dY_ptr + y); |
| } |
| #else |
| if (X_ptr[x] == Y_ptr[y]) { |
| sum += dY_ptr[y]; |
| } |
| #endif |
| } |
| } |
| } |
| dX_ptr[x] = sum; |
| } |
| } |
| |
| template <typename T> |
| __global__ void MaxPool3DBackwardNHWCCUDAKernel( |
| const int C, |
| const int X_D, |
| const int X_H, |
| const int X_W, |
| const int Y_D, |
| const int Y_H, |
| const int Y_W, |
| const int kernel_d, |
| const int kernel_h, |
| const int kernel_w, |
| const int stride_d, |
| const int stride_h, |
| const int stride_w, |
| const int pad_p, |
| const int pad_t, |
| const int pad_l, |
| const T* dY, |
| const T* X, |
| const T* Y, |
| T* dX) { |
| const int X_HxW = X_D * X_H * X_W; |
| const int Y_HxW = Y_D * Y_H * Y_W; |
| const int n = blockIdx.x / X_HxW; |
| const int x = blockIdx.x % X_HxW; |
| const int xx = x / X_W; |
| const int d = xx / X_H + pad_p; |
| const int h = xx % X_H + pad_t; |
| const int w = x % X_W + pad_l; |
| const int p = d < kernel_d ? 0 : (d - kernel_d) / stride_d + 1; |
| const int a = min(d / stride_d + 1, Y_D); |
| const int t = h < kernel_h ? 0 : (h - kernel_h) / stride_h + 1; |
| const int b = min(h / stride_h + 1, Y_H); |
| const int l = w < kernel_w ? 0 : (w - kernel_w) / stride_w + 1; |
| const int r = min(w / stride_w + 1, Y_W); |
| const T* dY_ptr = dY + n * Y_HxW * C; |
| const T* X_ptr = X + n * X_HxW * C; |
| const T* Y_ptr = Y + n * Y_HxW * C; |
| T* dX_ptr = dX + n * X_HxW * C; |
| for (int c = threadIdx.x; c < C; c += blockDim.x) { |
| T sum = 0; |
| for (int i = p; i < a; ++i) { |
| for (int j = t; j < b; ++j) { |
| for (int k = l; k < r; ++k) { |
| const int y = (i * Y_H + j) * Y_W + k; |
| #if __CUDA_ARCH__ >= 350 |
| if (__ldg(X_ptr + x * C + c) == __ldg(Y_ptr + y * C + c)) { |
| sum += __ldg(dY_ptr + y * C + c); |
| } |
| #else |
| if (X_ptr[x * C + c] == Y_ptr[y * C + c]) { |
| sum += dY_ptr[y * C + c]; |
| } |
| #endif |
| } |
| } |
| } |
| dX_ptr[x * C + c] = sum; |
| } |
| } |
| |
| } // namespace |
| |
| template <> |
| template <> |
| bool MaxPoolFunctor<CUDAContext>:: |
| GlobalPoolingForward<float, StorageOrder::NCHW>( |
| const int N, |
| const int C, |
| const int HxW, |
| const float* X, |
| float* Y, |
| CUDAContext* context) const { |
| const std::array<int, 2> X_dims = {N * C, HxW}; |
| const std::array<int, 2> Y_dims = {N * C, 1}; |
| math::ReduceMax<float, CUDAContext>( |
| 2, X_dims.data(), Y_dims.data(), 1.0f, X, Y, context); |
| return true; |
| } |
| |
| template <> |
| template <> |
| bool MaxPoolFunctor<CUDAContext>:: |
| GlobalPoolingForward<float, StorageOrder::NHWC>( |
| const int N, |
| const int C, |
| const int HxW, |
| const float* X, |
| float* Y, |
| CUDAContext* context) const { |
| const std::array<int, 3> X_dims = {N, HxW, C}; |
| const std::array<int, 3> Y_dims = {N, 1, C}; |
| math::ReduceMax<float, CUDAContext>( |
| 3, X_dims.data(), Y_dims.data(), 1.0f, X, Y, context); |
| return true; |
| } |
| |
| template <> |
| template <> |
| bool MaxPoolFunctor<CUDAContext>::Forward<float, StorageOrder::NCHW>( |
| const int N, |
| const int C, |
| const std::vector<int>& X_dims, |
| const std::vector<int>& Y_dims, |
| const std::vector<int>& kernel, |
| const std::vector<int>& /* dilation */, |
| const std::vector<int>& stride, |
| const std::vector<int>& pads, |
| const float* X, |
| float* Y, |
| CUDAContext* context) const { |
| const int ndim = X_dims.size(); |
| switch (ndim) { |
| case 1: { |
| const int num_blocks = N * C; |
| MaxPool1DForwardNCHWCUDAKernel<float> |
| <<<num_blocks, CAFFE_CUDA_NUM_THREADS, 0, context->cuda_stream()>>>( |
| X_dims[0], Y_dims[0], kernel[0], stride[0], pads[0], X, Y); |
| return true; |
| } |
| case 2: { |
| const int num_blocks = N * C * Y_dims[0]; |
| MaxPool2DForwardNCHWCUDAKernel<float> |
| <<<num_blocks, CAFFE_CUDA_NUM_THREADS, 0, context->cuda_stream()>>>( |
| X_dims[0], |
| X_dims[1], |
| Y_dims[0], |
| Y_dims[1], |
| kernel[0], |
| kernel[1], |
| stride[0], |
| stride[1], |
| pads[0], |
| pads[1], |
| X, |
| Y); |
| return true; |
| } |
| case 3: { |
| const int num_blocks = N * C * Y_dims[0] * Y_dims[1]; |
| MaxPool3DForwardNCHWCUDAKernel<float> |
| <<<num_blocks, CAFFE_CUDA_NUM_THREADS, 0, context->cuda_stream()>>>( |
| X_dims[0], |
| X_dims[1], |
| X_dims[2], |
| Y_dims[0], |
| Y_dims[1], |
| Y_dims[2], |
| kernel[0], |
| kernel[1], |
| kernel[2], |
| stride[0], |
| stride[1], |
| stride[2], |
| pads[0], |
| pads[1], |
| pads[2], |
| X, |
| Y); |
| return true; |
| } |
| default: { |
| CAFFE_THROW("Unsupported pooling dim: ", ndim); |
| return false; |
| } |
| } |
| } |
| |
| template <> |
| template <> |
| bool MaxPoolFunctor<CUDAContext>::Forward<float, StorageOrder::NHWC>( |
| const int N, |
| const int C, |
| const std::vector<int>& X_dims, |
| const std::vector<int>& Y_dims, |
| const std::vector<int>& kernel, |
| const std::vector<int>& /* dilation */, |
| const std::vector<int>& stride, |
| const std::vector<int>& pads, |
| const float* X, |
| float* Y, |
| CUDAContext* context) const { |
| // Each CUDA block handles one point, one thread per channel. |
| const int ndim = X_dims.size(); |
| const int Y_HxW = std::accumulate( |
| Y_dims.cbegin(), Y_dims.cend(), 1, std::multiplies<int>()); |
| switch (ndim) { |
| case 1: { |
| MaxPool1DForwardNHWCCUDAKernel<float> |
| <<<N * Y_HxW, CAFFE_CUDA_NUM_THREADS, 0, context->cuda_stream()>>>( |
| C, X_dims[0], Y_dims[0], kernel[0], stride[0], pads[0], X, Y); |
| return true; |
| } |
| case 2: { |
| MaxPool2DForwardNHWCCUDAKernel<float> |
| <<<N * Y_HxW, CAFFE_CUDA_NUM_THREADS, 0, context->cuda_stream()>>>( |
| C, |
| X_dims[0], |
| X_dims[1], |
| Y_dims[0], |
| Y_dims[1], |
| kernel[0], |
| kernel[1], |
| stride[0], |
| stride[1], |
| pads[0], |
| pads[1], |
| X, |
| Y); |
| return true; |
| } |
| case 3: { |
| MaxPool3DForwardNHWCCUDAKernel<float> |
| <<<N * Y_HxW, CAFFE_CUDA_NUM_THREADS, 0, context->cuda_stream()>>>( |
| C, |
| X_dims[0], |
| X_dims[1], |
| X_dims[2], |
| Y_dims[0], |
| Y_dims[1], |
| Y_dims[2], |
| kernel[0], |
| kernel[1], |
| kernel[2], |
| stride[0], |
| stride[1], |
| stride[2], |
| pads[0], |
| pads[1], |
| pads[2], |
| X, |
| Y); |
| return true; |
| } |
| default: { |
| CAFFE_THROW("Unsupported pooling dim: ", ndim); |
| return false; |
| } |
| } |
| } |
| |
| template <> |
| template <> |
| bool MaxPoolFunctor<CUDAContext>:: |
| GlobalPoolingBackward<float, StorageOrder::NCHW>( |
| const int N, |
| const int C, |
| const int HxW, |
| const float* dY, |
| const float* X, |
| const float* Y, |
| float* dX, |
| CUDAContext* context) const { |
| const int K = (HxW + CAFFE_CUDA_NUM_THREADS - 1) / CAFFE_CUDA_NUM_THREADS; |
| GlobalMaxPoolingBackwardNCHWCUDAKernel<float> |
| <<<N * C * K, CAFFE_CUDA_NUM_THREADS, 0, context->cuda_stream()>>>( |
| K, HxW, dY, X, Y, dX); |
| return true; |
| } |
| |
| template <> |
| template <> |
| bool MaxPoolFunctor<CUDAContext>:: |
| GlobalPoolingBackward<float, StorageOrder::NHWC>( |
| const int N, |
| const int C, |
| const int HxW, |
| const float* dY, |
| const float* X, |
| const float* Y, |
| float* dX, |
| CUDAContext* context) const { |
| GlobalMaxPoolingBackwardNHWCCUDAKernel<float> |
| <<<N * HxW, CAFFE_CUDA_NUM_THREADS, 0, context->cuda_stream()>>>( |
| C, HxW, dY, X, Y, dX); |
| return true; |
| } |
| |
| template <> |
| template <> |
| bool MaxPoolFunctor<CUDAContext>::Backward<float, StorageOrder::NCHW>( |
| const int N, |
| const int C, |
| const std::vector<int>& X_dims, |
| const std::vector<int>& Y_dims, |
| const std::vector<int>& kernel, |
| const std::vector<int>& /* dilation */, |
| const std::vector<int>& stride, |
| const std::vector<int>& pads, |
| const float* dY, |
| const float* X, |
| const float* Y, |
| float* dX, |
| CUDAContext* context) const { |
| const int ndim = X_dims.size(); |
| switch (ndim) { |
| case 1: { |
| const int num_blocks = N * C; |
| MaxPool1DBackwardNCHWCUDAKernel<float> |
| <<<num_blocks, CAFFE_CUDA_NUM_THREADS, 0, context->cuda_stream()>>>( |
| X_dims[0], |
| Y_dims[0], |
| kernel[0], |
| stride[0], |
| pads[0], |
| dY, |
| X, |
| Y, |
| dX); |
| return true; |
| } |
| case 2: { |
| const int num_blocks = N * C * X_dims[0]; |
| MaxPool2DBackwardNCHWCUDAKernel<float> |
| <<<num_blocks, CAFFE_CUDA_NUM_THREADS, 0, context->cuda_stream()>>>( |
| X_dims[0], |
| X_dims[1], |
| Y_dims[0], |
| Y_dims[1], |
| kernel[0], |
| kernel[1], |
| stride[0], |
| stride[1], |
| pads[0], |
| pads[1], |
| dY, |
| X, |
| Y, |
| dX); |
| return true; |
| } |
| case 3: { |
| const int num_blocks = N * C * X_dims[0] * X_dims[1]; |
| MaxPool3DBackwardNCHWCUDAKernel<float> |
| <<<num_blocks, CAFFE_CUDA_NUM_THREADS, 0, context->cuda_stream()>>>( |
| X_dims[0], |
| X_dims[1], |
| X_dims[2], |
| Y_dims[0], |
| Y_dims[1], |
| Y_dims[2], |
| kernel[0], |
| kernel[1], |
| kernel[2], |
| stride[0], |
| stride[1], |
| stride[2], |
| pads[0], |
| pads[1], |
| pads[2], |
| dY, |
| X, |
| Y, |
| dX); |
| return true; |
| } |
| default: { |
| CAFFE_THROW("Unsupported pooling dim: ", ndim); |
| return false; |
| } |
| } |
| } |
| |
| template <> |
| template <> |
| bool MaxPoolFunctor<CUDAContext>::Backward<float, StorageOrder::NHWC>( |
| const int N, |
| const int C, |
| const std::vector<int>& X_dims, |
| const std::vector<int>& Y_dims, |
| const std::vector<int>& kernel, |
| const std::vector<int>& /* dilation */, |
| const std::vector<int>& stride, |
| const std::vector<int>& pads, |
| const float* dY, |
| const float* X, |
| const float* Y, |
| float* dX, |
| CUDAContext* context) const { |
| const int ndim = X_dims.size(); |
| const int X_HxW = std::accumulate( |
| X_dims.cbegin(), X_dims.cend(), 1, std::multiplies<int>()); |
| switch (ndim) { |
| case 1: { |
| MaxPool1DBackwardNHWCCUDAKernel<float> |
| <<<N * X_HxW, CAFFE_CUDA_NUM_THREADS, 0, context->cuda_stream()>>>( |
| C, |
| X_dims[0], |
| Y_dims[0], |
| kernel[0], |
| stride[0], |
| pads[0], |
| dY, |
| X, |
| Y, |
| dX); |
| return true; |
| } |
| case 2: { |
| MaxPool2DBackwardNHWCCUDAKernel<float> |
| <<<N * X_HxW, CAFFE_CUDA_NUM_THREADS, 0, context->cuda_stream()>>>( |
| C, |
| X_dims[0], |
| X_dims[1], |
| Y_dims[0], |
| Y_dims[1], |
| kernel[0], |
| kernel[1], |
| stride[0], |
| stride[1], |
| pads[0], |
| pads[1], |
| dY, |
| X, |
| Y, |
| dX); |
| return true; |
| } |
| case 3: { |
| MaxPool3DBackwardNHWCCUDAKernel<float> |
| <<<N * X_HxW, CAFFE_CUDA_NUM_THREADS, 0, context->cuda_stream()>>>( |
| C, |
| X_dims[0], |
| X_dims[1], |
| X_dims[2], |
| Y_dims[0], |
| Y_dims[1], |
| Y_dims[2], |
| kernel[0], |
| kernel[1], |
| kernel[2], |
| stride[0], |
| stride[1], |
| stride[2], |
| pads[0], |
| pads[1], |
| pads[2], |
| dY, |
| X, |
| Y, |
| dX); |
| return true; |
| } |
| default: { |
| CAFFE_THROW("Unsupported pooling dim: ", ndim); |
| return false; |
| } |
| } |
| } |
| |
| REGISTER_CUDA_OPERATOR( |
| AveragePool, |
| PoolOp<float, CUDAContext, AveragePoolFunctor<CUDAContext>>); |
| REGISTER_CUDA_OPERATOR( |
| AveragePoolGradient, |
| PoolGradientOp<float, CUDAContext, AveragePoolFunctor<CUDAContext>>); |
| |
| REGISTER_CUDA_OPERATOR( |
| AveragePool1D, |
| PoolOp<float, CUDAContext, AveragePoolFunctor<CUDAContext>>); |
| REGISTER_CUDA_OPERATOR( |
| AveragePool1DGradient, |
| PoolGradientOp<float, CUDAContext, AveragePoolFunctor<CUDAContext>>); |
| |
| REGISTER_CUDA_OPERATOR( |
| AveragePool2D, |
| PoolOp<float, CUDAContext, AveragePoolFunctor<CUDAContext>>); |
| REGISTER_CUDA_OPERATOR( |
| AveragePool2DGradient, |
| PoolGradientOp<float, CUDAContext, AveragePoolFunctor<CUDAContext>>); |
| |
| REGISTER_CUDA_OPERATOR( |
| AveragePool3D, |
| PoolOp<float, CUDAContext, AveragePoolFunctor<CUDAContext>>); |
| REGISTER_CUDA_OPERATOR( |
| AveragePool3DGradient, |
| PoolGradientOp<float, CUDAContext, AveragePoolFunctor<CUDAContext>>); |
| |
| REGISTER_CUDA_OPERATOR( |
| MaxPool, |
| PoolOp<float, CUDAContext, MaxPoolFunctor<CUDAContext>>); |
| REGISTER_CUDA_OPERATOR( |
| MaxPoolGradient, |
| PoolGradientOp<float, CUDAContext, MaxPoolFunctor<CUDAContext>>); |
| |
| REGISTER_CUDA_OPERATOR( |
| MaxPool1D, |
| PoolOp<float, CUDAContext, MaxPoolFunctor<CUDAContext>>); |
| REGISTER_CUDA_OPERATOR( |
| MaxPool1DGradient, |
| PoolGradientOp<float, CUDAContext, MaxPoolFunctor<CUDAContext>>); |
| |
| REGISTER_CUDA_OPERATOR( |
| MaxPool2D, |
| PoolOp<float, CUDAContext, MaxPoolFunctor<CUDAContext>>); |
| REGISTER_CUDA_OPERATOR( |
| MaxPool2DGradient, |
| PoolGradientOp<float, CUDAContext, MaxPoolFunctor<CUDAContext>>); |
| |
| REGISTER_CUDA_OPERATOR( |
| MaxPool3D, |
| PoolOp<float, CUDAContext, MaxPoolFunctor<CUDAContext>>); |
| REGISTER_CUDA_OPERATOR( |
| MaxPool3DGradient, |
| PoolGradientOp<float, CUDAContext, MaxPoolFunctor<CUDAContext>>); |
| |
| } // namespace caffe2 |
