#include "caffe2/operators/spatial_batch_norm_op.h"
-#include <cub/block/block_reduce.cuh>
-
-#include "caffe2/core/context_gpu.h"
-#include "caffe2/utils/math.h"
+#include "caffe2/operators/spatial_batch_norm_op_gpu_impl.cuh"
namespace caffe2 {
-namespace {
-
-template <typename T>
-using BlockReduce = cub::BlockReduce<T, CAFFE_CUDA_NUM_THREADS>;
-
-template <typename T, int kBlockDimX, int kBlockDimY>
-using BlockReduce2D = cub::
- BlockReduce<T, kBlockDimX, cub::BLOCK_REDUCE_WARP_REDUCTIONS, kBlockDimY>;
-
-template <typename T>
-__global__ void ComputeFusedParamCUDAKernel(
- const int C,
- const T epsilon,
- const T* scale,
- const T* bias,
- const T* mean,
- const T* var,
- T* alpha,
- T* beta);
-
-template <>
-__global__ void ComputeFusedParamCUDAKernel<float>(
- const int C,
- const float epsilon,
- const float* scale,
- const float* bias,
- const float* mean,
- const float* var,
- float* alpha,
- float* beta) {
- const int c = blockIdx.x * CAFFE_CUDA_NUM_THREADS + threadIdx.x;
- if (c < C) {
-#if __CUDA_ARCH__ >= 350 || defined(__HIP_PLATFORM_HCC__)
- const float scale_x_rstd =
- __ldg(scale + c) * rsqrtf(__ldg(var + c) + epsilon);
- alpha[c] = scale_x_rstd;
- beta[c] = fmaf(-scale_x_rstd, __ldg(mean + c), __ldg(bias + c));
-#else
- const float scale_x_rstd = scale[c] * rsqrtf(var[c] + epsilon);
- alpha[c] = scale_x_rstd;
- beta[c] = fmaf(-scale_x_rstd, mean[c], bias[c]);
-#endif
- }
-}
-
-template <typename T>
-__global__ void ComputeBatchMomentsCUDAKernel(
- const int C,
- const T scale,
- const T* batch_mean_sum,
- const T* batch_var_sum,
- T* mean,
- T* var);
-
-template <>
-__global__ void ComputeBatchMomentsCUDAKernel<float>(
- const int C,
- const float scale,
- const float* batch_mean_sum,
- const float* batch_var_sum,
- float* mean,
- float* var) {
- const int c = blockIdx.x * CAFFE_CUDA_NUM_THREADS + threadIdx.x;
- if (c < C) {
-#if __CUDA_ARCH__ >= 350 || defined(__HIP_PLATFORM_HCC__)
- const float mu = __ldg(batch_mean_sum + c) * scale;
- mean[c] = mu;
- var[c] = fmaf(__ldg(batch_var_sum + c), scale, -mu * mu);
-#else
- const float mu = batch_mean_sum[c] * scale;
- mean[c] = mu;
- var[c] = fmaf(batch_var_sum[c], scale, -mu * mu);
-#endif
- }
-}
-
-template <typename T>
-__global__ void ComputeRunningMomentsAndFusedParamCUDAKernel(
- const int C,
- const T momentum,
- const T epsilon,
- const T* scale,
- const T* bias,
- const T* mean,
- const T* var,
- T* running_mean,
- T* running_var,
- T* rstd,
- T* alpha,
- T* beta);
-
-template <>
-__global__ void ComputeRunningMomentsAndFusedParamCUDAKernel<float>(
- const int C,
- const float momentum,
- const float epsilon,
- const float* scale,
- const float* bias,
- const float* mean,
- const float* var,
- float* running_mean,
- float* running_var,
- float* rstd,
- float* alpha,
- float* beta) {
- const float a = 1.0f - momentum;
- const float b = momentum;
- const int c = blockIdx.x * CAFFE_CUDA_NUM_THREADS + threadIdx.x;
- if (c < C) {
-#if __CUDA_ARCH__ >= 350 || defined(__HIP_PLATFORM_HCC__)
- running_mean[c] = fmaf(a, __ldg(mean + c), b * __ldg(running_mean + c));
- running_var[c] = fmaf(a, __ldg(var + c), b * __ldg(running_var + c));
- const float rstd_val = rsqrtf(__ldg(var + c) + epsilon);
- const float scale_x_rstd = __ldg(scale + c) * rstd_val;
- rstd[c] = rstd_val;
- alpha[c] = scale_x_rstd;
- beta[c] = fmaf(-scale_x_rstd, __ldg(mean + c), __ldg(bias + c));
-#else
- running_mean[c] = fmaf(a, mean[c], b * running_mean[c]);
- running_var[c] = fmaf(a, var[c], b * running_var[c]);
- const float rstd_val = rsqrtf(var[c] + epsilon);
- const float scale_x_rstd = scale[c] * rstd_val;
- rstd[c] = rstd_val;
- alpha[c] = scale_x_rstd;
- beta[c] = fmaf(-scale_x_rstd, mean[c], bias[c]);
-#endif
- }
-}
-
-template <typename T>
-__global__ void ComputeMultiBatchScaleBiasGradientsAndFusedParamsCUDAKernel(
- const int C,
- const T inv_num_batches,
- const T inv_nhw,
- const T* scale,
- const T* mean,
- const T* rstd,
- const T* dscale_sum,
- const T* dbias_sum,
- T* dscale,
- T* dbias,
- T* alpha,
- T* beta,
- T* gamma);
-
-template <>
-__global__ void
-ComputeMultiBatchScaleBiasGradientsAndFusedParamsCUDAKernel<float>(
- const int C,
- const float inv_num_batches,
- const float inv_nhw,
- const float* scale,
- const float* mean,
- const float* rstd,
- const float* dscale_sum,
- const float* dbias_sum,
- float* dscale,
- float* dbias,
- float* alpha,
- float* beta,
- float* gamma) {
- const int c = blockIdx.x * CAFFE_CUDA_NUM_THREADS + threadIdx.x;
- if (c < C) {
-#if __CUDA_ARCH__ >= 350 || defined(__HIP_PLATFORM_HCC__)
- const float dscale_val = __ldg(dscale_sum + c) * inv_num_batches;
- const float dbias_val = __ldg(dbias_sum + c) * inv_num_batches;
- const float scale_x_rstd = __ldg(scale + c) * __ldg(rstd + c);
- const float dscale_x_rstd = dscale_val * __ldg(rstd + c);
- dscale[c] = dscale_val;
- dbias[c] = dbias_val;
- alpha[c] = scale_x_rstd;
- beta[c] = -scale_x_rstd * dscale_x_rstd * inv_nhw;
- gamma[c] = scale_x_rstd * fmaf(__ldg(mean + c), dscale_x_rstd, -dbias_val) *
- inv_nhw;
-#else
- const float dscale_val = dscale_sum[c] * inv_num_batches;
- const float dbias_val = dbias_sum[c] * inv_num_batches;
- const float scale_x_rstd = scale[c] * rstd[c];
- const float dscale_x_rstd = dscale_val * rstd[c];
- dscale[c] = dscale_val;
- dbias[c] = dbias_val;
- alpha[c] = scale_x_rstd;
- beta[c] = -scale_x_rstd * dscale_x_rstd * inv_nhw;
- gamma[c] =
- scale_x_rstd * fmaf(mean[c], dscale_x_rstd, -dbias_val) * inv_nhw;
-#endif
- }
-}
-
-template <typename T, int kBlockDimX, int kBlockDimY>
-__global__ void ComputeScaleBiasGradientsAndFusedParamsNCHWCUDAKernel(
- const int N,
- const int C,
- const int HxW,
- const T* dY,
- const T* X,
- const T* scale,
- const T* mean,
- const T* rstd,
- T* dscale,
- T* dbias,
- T* alpha,
- T* beta,
- T* gamma) {
- __shared__
- typename BlockReduce2D<T, kBlockDimX, kBlockDimY>::TempStorage ds_storage;
- __shared__
- typename BlockReduce2D<T, kBlockDimX, kBlockDimY>::TempStorage db_storage;
- const T inv_nhw = T(1) / static_cast<T>(N * HxW);
- const int c = blockIdx.x;
- T ds_val = 0;
- T db_val = 0;
- for (int i = threadIdx.x; i < N; i += blockDim.x) {
- for (int j = threadIdx.y; j < HxW; j += blockDim.y) {
- const int index = (i * C + c) * HxW + j;
-#if __CUDA_ARCH__ >= 350 || defined(__HIP_PLATFORM_HCC__)
- ds_val += __ldg(dY + index) * __ldg(X + index);
- db_val += __ldg(dY + index);
-#else
- ds_val += dY[index] * X[index];
- db_val += dY[index];
-#endif
- }
- }
- ds_val = BlockReduce2D<T, kBlockDimX, kBlockDimY>(ds_storage).Sum(ds_val);
- db_val = BlockReduce2D<T, kBlockDimX, kBlockDimY>(db_storage).Sum(db_val);
- if (threadIdx.x == 0 && threadIdx.y == 0) {
-#if __CUDA_ARCH__ >= 350 || defined(__HIP_PLATFORM_HCC__)
- ds_val = (ds_val - db_val * __ldg(mean + c)) * __ldg(rstd + c);
- const T scale_x_rstd = __ldg(scale + c) * __ldg(rstd + c);
- const T dscale_x_rstd = ds_val * __ldg(rstd + c);
- dscale[c] = ds_val;
- dbias[c] = db_val;
- alpha[c] = scale_x_rstd;
- beta[c] = -scale_x_rstd * dscale_x_rstd * inv_nhw;
- gamma[c] =
- scale_x_rstd * (__ldg(mean + c) * dscale_x_rstd - db_val) * inv_nhw;
-#else
- ds_val = (ds_val - db_val * mean[c]) * rstd[c];
- const T scale_x_rstd = scale[c] * rstd[c];
- const T dscale_x_rstd = ds_val * rstd[c];
- dscale[c] = ds_val;
- dbias[c] = db_val;
- alpha[c] = scale_x_rstd;
- beta[c] = -scale_x_rstd * dscale_x_rstd * inv_nhw;
- gamma[c] = scale_x_rstd * (mean[c] * dscale_x_rstd - db_val) * inv_nhw;
-#endif
- }
-}
-
-template <typename T>
-__global__ void ComputeScaleGradientAndFusedParamsNHWCCUDAKernel(
- const int C,
- const T inv_nhw,
- const T* dYxX,
- const T* dbias,
- const T* scale,
- const T* mean,
- const T* rstd,
- T* dscale,
- T* alpha,
- T* beta,
- T* gamma);
-
-template <>
-__global__ void ComputeScaleGradientAndFusedParamsNHWCCUDAKernel<float>(
- const int C,
- const float inv_nhw,
- const float* dYxX,
- const float* dbias,
- const float* scale,
- const float* mean,
- const float* rstd,
- float* dscale,
- float* alpha,
- float* beta,
- float* gamma) {
- const int c = blockIdx.x * CAFFE_CUDA_NUM_THREADS + threadIdx.x;
- if (c < C) {
-#if __CUDA_ARCH__ >= 350 || defined(__HIP_PLATFORM_HCC__)
- const float ds = fmaf(-__ldg(dbias + c), __ldg(mean + c), __ldg(dYxX + c)) *
- __ldg(rstd + c);
- dscale[c] = ds;
- const float scale_x_rstd = __ldg(scale + c) * __ldg(rstd + c);
- const float dscale_x_rstd = ds * __ldg(rstd + c);
- alpha[c] = scale_x_rstd;
- beta[c] = -scale_x_rstd * dscale_x_rstd * inv_nhw;
- gamma[c] = scale_x_rstd *
- fmaf(__ldg(mean + c), dscale_x_rstd, -__ldg(dbias + c)) * inv_nhw;
-#else
- const float ds = fmaf(-dbias[c], mean[c], dYxX[c]) * rstd[c];
- dscale[c] = ds;
- const float scale_x_rstd = scale[c] * rstd[c];
- const float dscale_x_rstd = ds * rstd[c];
- alpha[c] = scale_x_rstd;
- beta[c] = -scale_x_rstd * dscale_x_rstd * inv_nhw;
- gamma[c] = scale_x_rstd * fmaf(mean[c], dscale_x_rstd, -dbias[c]) * inv_nhw;
-#endif
- }
-}
-
-template <typename T>
-__global__ void ComputeXGradientNCHWCUDAKernel(
- const int C,
- const int HxW,
- const int K,
- const T* dY,
- const T* X,
- const T* alpha,
- const T* beta,
- const T* gamma,
- T* dX);
-
-template <>
-__global__ void ComputeXGradientNCHWCUDAKernel<float>(
- const int C,
- const int HxW,
- const int K,
- const float* dY,
- const float* X,
- const float* alpha,
- const float* beta,
- const float* gamma,
- float* dX) {
- const int nc = blockIdx.x / K;
- const int block = blockIdx.x % K;
- const int c = nc % C;
- const int w = block * CAFFE_CUDA_NUM_THREADS + threadIdx.x;
- if (w < HxW) {
- const int index = nc * HxW + w;
-#if __CUDA_ARCH__ >= 350 || defined(__HIP_PLATFORM_HCC__)
- dX[index] = fmaf(
- __ldg(alpha + c),
- __ldg(dY + index),
- fmaf(__ldg(beta + c), __ldg(X + index), __ldg(gamma + c)));
-#else
- dX[index] = fmaf(alpha[c], dY[index], fma(beta[c], X[index], gamma[c]));
-#endif
- }
-}
-
-template <typename T>
-__global__ void ComputeXGradientNHWCCUDAKernel(
- const int C,
- const T* dY,
- const T* X,
- const T* alpha,
- const T* beta,
- const T* gamma,
- T* dX);
-
-template <>
-__global__ void ComputeXGradientNHWCCUDAKernel<float>(
- const int C,
- const float* dY,
- const float* X,
- const float* alpha,
- const float* beta,
- const float* gamma,
- float* dX) {
- for (int c = threadIdx.x; c < C; c += blockDim.x) {
- const int index = blockIdx.x * C + c;
-#if __CUDA_ARCH__ >= 350 || defined(__HIP_PLATFORM_HCC__)
- dX[index] = fmaf(
- __ldg(alpha + c),
- __ldg(dY + index),
- fmaf(__ldg(beta + c), __ldg(X + index), __ldg(gamma + c)));
-#else
- dX[index] = fmaf(alpha[c], dY[index], fma(beta[c], X[index], gamma[c]));
-#endif
- }
-}
-
-} // namespace
-
-template <>
-template <>
-void SpatialBNOp<CUDAContext>::ComputeFusedParam<float>(
- const int C,
- const float* scale,
- const float* bias,
- const float* mean,
- const float* var,
- float* alpha,
- float* beta) {
- const int K = math::DivUp(C, CAFFE_CUDA_NUM_THREADS);
- ComputeFusedParamCUDAKernel<float>
- <<<K, CAFFE_CUDA_NUM_THREADS, 0, context_.cuda_stream()>>>(
- C, static_cast<float>(epsilon_), scale, bias, mean, var, alpha, beta);
-}
-
-template <>
-template <>
-void SpatialBNOp<CUDAContext>::ComputeBatchMoments<float>(
- const int N,
- const int C,
- const int HxW,
- const float* batch_mean_sum,
- const float* batch_var_sum,
- float* mean,
- float* var) {
- const int K = math::DivUp(C, CAFFE_CUDA_NUM_THREADS);
- const float scale = 1.0f / static_cast<float>(num_batches_ * N * HxW);
- ComputeBatchMomentsCUDAKernel<float>
- <<<K, CAFFE_CUDA_NUM_THREADS, 0, context_.cuda_stream()>>>(
- C, scale, batch_mean_sum, batch_var_sum, mean, var);
-}
-
-template <>
-template <>
-void SpatialBNOp<CUDAContext>::ComputeRunningMomentsAndFusedParam<float>(
- const int C,
- const float* scale,
- const float* bias,
- const float* mean,
- const float* var,
- float* running_mean,
- float* running_var,
- float* rstd,
- float* alpha,
- float* beta) {
- const int K = math::DivUp(C, CAFFE_CUDA_NUM_THREADS);
- ComputeRunningMomentsAndFusedParamCUDAKernel<float>
- <<<K, CAFFE_CUDA_NUM_THREADS, 0, context_.cuda_stream()>>>(
- C,
- static_cast<float>(momentum_),
- static_cast<float>(epsilon_),
- scale,
- bias,
- mean,
- var,
- running_mean,
- running_var,
- rstd,
- alpha,
- beta);
-}
-
-template <>
-template <>
-void SpatialBNGradientOp<CUDAContext>::
- ComputeMultiBatchScaleBiasGradientsAndFusedParams<float>(
- const int N,
- const int C,
- const int HxW,
- const float* scale,
- const float* mean,
- const float* rstd,
- const float* dscale_sum,
- const float* dbias_sum,
- float* dscale,
- float* dbias,
- float* alpha,
- float* beta,
- float* gamma) {
- const float inv_num_batches = 1.0f / static_cast<float>(num_batches_);
- const float inv_nhw = 1.0f / static_cast<float>(N * HxW);
- const int K = math::DivUp(C, CAFFE_CUDA_NUM_THREADS);
- ComputeMultiBatchScaleBiasGradientsAndFusedParamsCUDAKernel<float>
- <<<K, CAFFE_CUDA_NUM_THREADS, 0, context_.cuda_stream()>>>(
- C,
- inv_num_batches,
- inv_nhw,
- scale,
- mean,
- rstd,
- dscale_sum,
- dbias_sum,
- dscale,
- dbias,
- alpha,
- beta,
- gamma);
-}
-
-template <>
-template <>
-void SpatialBNGradientOp<CUDAContext>::ComputeScaleBiasGradientsAndFusedParams<
- float>(
- const int N,
- const int C,
- const int HxW,
- const float* dY,
- const float* X,
- const float* scale,
- const float* mean,
- const float* rstd,
- float* dscale,
- float* dbias,
- float* alpha,
- float* beta,
- float* gamma,
- float* scratch) {
- if (order_ == StorageOrder::NCHW) {
- if (HxW >= 128) {
- ComputeScaleBiasGradientsAndFusedParamsNCHWCUDAKernel<float, 1, 128>
- <<<C, dim3(1, 128), 0, context_.cuda_stream()>>>(
- N,
- C,
- HxW,
- dY,
- X,
- scale,
- mean,
- rstd,
- dscale,
- dbias,
- alpha,
- beta,
- gamma);
- } else if (HxW >= 64) {
- ComputeScaleBiasGradientsAndFusedParamsNCHWCUDAKernel<float, 2, 64>
- <<<C, dim3(2, 64), 0, context_.cuda_stream()>>>(
- N,
- C,
- HxW,
- dY,
- X,
- scale,
- mean,
- rstd,
- dscale,
- dbias,
- alpha,
- beta,
- gamma);
- } else if (HxW >= 32) {
- ComputeScaleBiasGradientsAndFusedParamsNCHWCUDAKernel<float, 4, 32>
- <<<C, dim3(4, 32), 0, context_.cuda_stream()>>>(
- N,
- C,
- HxW,
- dY,
- X,
- scale,
- mean,
- rstd,
- dscale,
- dbias,
- alpha,
- beta,
- gamma);
- } else {
- ComputeScaleBiasGradientsAndFusedParamsNCHWCUDAKernel<float, 8, 16>
- <<<C, dim3(8, 16), 0, context_.cuda_stream()>>>(
- N,
- C,
- HxW,
- dY,
- X,
- scale,
- mean,
- rstd,
- dscale,
- dbias,
- alpha,
- beta,
- gamma);
- }
- } else {
- ReinitializeTensor(&ones_, {N * HxW}, at::dtype<float>().device(CUDA));
- math::Set<float, CUDAContext>(
- N * HxW, 1.0f, ones_.mutable_data<float>(), &context_);
- const float* ones_data = ones_.data<float>();
- math::Mul<float, CUDAContext>(N * C * HxW, dY, X, scratch, &context_);
- math::Gemm<float, CUDAContext>(
- CblasTrans,
- CblasNoTrans,
- C,
- 1,
- N * HxW,
- 1.0f,
- scratch,
- ones_data,
- 0.0f,
- dscale,
- &context_);
- math::Gemm<float, CUDAContext>(
- CblasTrans,
- CblasNoTrans,
- C,
- 1,
- N * HxW,
- 1.0f,
- dY,
- ones_data,
- 0.0f,
- dbias,
- &context_);
- const int K = math::DivUp(C, CAFFE_CUDA_NUM_THREADS);
- ComputeScaleGradientAndFusedParamsNHWCCUDAKernel<float>
- <<<K, CAFFE_CUDA_NUM_THREADS, 0, context_.cuda_stream()>>>(
- C,
- 1.0f / static_cast<float>(N * HxW),
- dscale,
- dbias,
- scale,
- mean,
- rstd,
- dscale,
- alpha,
- beta,
- gamma);
- }
-}
-
-template <>
-template <>
-void SpatialBNGradientOp<CUDAContext>::ComputeXGradient<float>(
- const int N,
- const int C,
- const int HxW,
- const float* dY,
- const float* X,
- const float* alpha,
- const float* beta,
- const float* gamma,
- float* dX) {
- if (order_ == StorageOrder::NCHW) {
- const int K = math::DivUp(HxW, CAFFE_CUDA_NUM_THREADS);
- ComputeXGradientNCHWCUDAKernel<float>
- <<<N * C * K, CAFFE_CUDA_NUM_THREADS, 0, context_.cuda_stream()>>>(
- C, HxW, K, dY, X, alpha, beta, gamma, dX);
- } else {
- ComputeXGradientNHWCCUDAKernel<float>
- <<<N * HxW, CAFFE_CUDA_NUM_THREADS, 0, context_.cuda_stream()>>>(
- C, dY, X, alpha, beta, gamma, dX);
- }
-}
-
REGISTER_CUDA_OPERATOR(SpatialBN, SpatialBNOp<CUDAContext>);
REGISTER_CUDA_OPERATOR(SpatialBNGradient, SpatialBNGradientOp<CUDAContext>);
--- /dev/null
+#include "caffe2/operators/spatial_batch_norm_op.h"
+
+#include <cub/block/block_reduce.cuh>
+
+#include "caffe2/core/context_gpu.h"
+#include "caffe2/utils/math.h"
+
+namespace caffe2 {
+
+namespace {
+
+template <typename T>
+using BlockReduce = cub::BlockReduce<T, CAFFE_CUDA_NUM_THREADS>;
+
+template <typename T>
+__global__ void ComputeFusedParamCUDAKernel(
+ const int C,
+ const T epsilon,
+ const T* scale,
+ const T* bias,
+ const T* mean,
+ const T* var,
+ T* alpha,
+ T* beta);
+
+template <>
+__global__ void ComputeFusedParamCUDAKernel<float>(
+ const int C,
+ const float epsilon,
+ const float* scale,
+ const float* bias,
+ const float* mean,
+ const float* var,
+ float* alpha,
+ float* beta) {
+ CUDA_1D_KERNEL_LOOP(i, C) {
+#if __CUDA_ARCH__ >= 350
+ const float scale_x_rstd =
+ __ldg(scale + i) * rsqrtf(__ldg(var + i) + epsilon);
+ alpha[i] = scale_x_rstd;
+ beta[i] = __ldg(bias + i) - scale_x_rstd * __ldg(mean + i);
+#else
+ const float scale_x_rstd = scale[i] * rsqrtf(var[i] + epsilon);
+ alpha[i] = scale_x_rstd;
+ beta[i] = bias[i] - scale_x_rstd * mean[i];
+#endif
+ }
+}
+
+template <typename T>
+__global__ void ComputeBatchMomentsCUDAKernel(
+ const int C,
+ const T scale,
+ const T* batch_mean_sum,
+ const T* batch_var_sum,
+ T* mean,
+ T* var) {
+ CUDA_1D_KERNEL_LOOP(i, C) {
+#if __CUDA_ARCH__ >= 350
+ const T mu = __ldg(batch_mean_sum + i) * scale;
+ mean[i] = mu;
+ var[i] = __ldg(batch_var_sum + i) * scale - mu * mu;
+#else
+ const T mu = batch_mean_sum[i] * scale;
+ mean[i] = mu;
+ var[i] = batch_var_sum[i] * scale - mu * mu;
+#endif
+ }
+}
+
+template <typename T>
+__global__ void ComputeRunningMomentsAndFusedParamCUDAKernel(
+ const int C,
+ const T momentum,
+ const T epsilon,
+ const T* scale,
+ const T* bias,
+ const T* mean,
+ const T* var,
+ T* running_mean,
+ T* running_var,
+ T* rstd,
+ T* alpha,
+ T* beta);
+
+template <>
+__global__ void ComputeRunningMomentsAndFusedParamCUDAKernel<float>(
+ const int C,
+ const float momentum,
+ const float epsilon,
+ const float* scale,
+ const float* bias,
+ const float* mean,
+ const float* var,
+ float* running_mean,
+ float* running_var,
+ float* rstd,
+ float* alpha,
+ float* beta) {
+ const float a = 1.0f - momentum;
+ const float b = momentum;
+ CUDA_1D_KERNEL_LOOP(i, C) {
+#if __CUDA_ARCH__ >= 350
+ running_mean[i] = a * __ldg(mean + i) + b * __ldg(running_mean + i);
+ running_var[i] = a * __ldg(var + i) + b * __ldg(running_var + i);
+ const float rstd_val = rsqrtf(__ldg(var + i) + epsilon);
+ const float scale_x_rstd = __ldg(scale + i) * rstd_val;
+ rstd[i] = rstd_val;
+ alpha[i] = scale_x_rstd;
+ beta[i] = bias[i] - scale_x_rstd * __ldg(mean + i);
+#else
+ running_mean[i] = a * mean[i] + b * running_mean[i];
+ running_var[i] = a * var[i] + b * running_var[i];
+ const float rstd_val = rsqrtf(var[i] + epsilon);
+ const float scale_x_rstd = scale[i] * rstd_val;
+ rstd[i] = rstd_val;
+ alpha[i] = scale_x_rstd;
+ beta[i] = bias[i] - scale_x_rstd * mean[i];
+#endif
+ }
+}
+
+template <typename T>
+__global__ void ComputeMultiBatchScaleBiasGradientsAndFusedParamsCUDAKernel(
+ const int C,
+ const T inv_num_batches,
+ const T inv_nhw,
+ const T* scale,
+ const T* mean,
+ const T* rstd,
+ const T* dscale_sum,
+ const T* dbias_sum,
+ T* dscale,
+ T* dbias,
+ T* alpha,
+ T* beta,
+ T* gamma) {
+ CUDA_1D_KERNEL_LOOP(i, C) {
+#if __CUDA_ARCH__ >= 350
+ const T dscale_val = __ldg(dscale_sum + i) * inv_num_batches;
+ const T dbias_val = __ldg(dbias_sum + i) * inv_num_batches;
+ const T scale_x_rstd = __ldg(scale + i) * __ldg(rstd + i);
+ const T dscale_x_rstd = dscale_val * __ldg(rstd + i);
+ dscale[i] = dscale_val;
+ dbias[i] = dbias_val;
+ alpha[i] = scale_x_rstd;
+ beta[i] = -scale_x_rstd * dscale_x_rstd * inv_nhw;
+ gamma[i] =
+ scale_x_rstd * (__ldg(mean + i) * dscale_x_rstd - dbias_val) * inv_nhw;
+#else
+ const T dscale_val = dscale_sum[i] * inv_num_batches;
+ const T dbias_val = dbias_sum[i] * inv_num_batches;
+ const T scale_x_rstd = scale[i] * rstd[i];
+ const T dscale_x_rstd = dscale_val * rstd[i];
+ dscale[i] = dscale_val;
+ dbias[i] = dbias_val;
+ alpha[i] = scale_x_rstd;
+ beta[i] = -scale_x_rstd * dscale_x_rstd * inv_nhw;
+ gamma[i] = scale_x_rstd * (mean[i] * dscale_x_rstd - dbias_val) * inv_nhw;
+#endif
+ }
+}
+
+template <typename T, StorageOrder kOrder>
+__global__ void ComputeScaleBiasGradientsAndFusedParamsCUDAKernel(
+ const int N,
+ const int C,
+ const int HxW,
+ const T* dY,
+ const T* X,
+ const T* scale,
+ const T* mean,
+ const T* rstd,
+ T* dscale,
+ T* dbias,
+ T* alpha,
+ T* beta,
+ T* gamma) {
+ const int outer_size = C;
+ const int inner_size = N * HxW;
+ const T inv_nhw = T(1) / static_cast<T>(N * HxW);
+ __shared__ typename BlockReduce<T>::TempStorage ds_storage;
+ __shared__ typename BlockReduce<T>::TempStorage db_storage;
+ for (int i = blockIdx.x; i < outer_size; i += gridDim.x) {
+ T ds_val = 0;
+ T db_val = 0;
+ for (int j = threadIdx.x; j < inner_size; j += blockDim.x) {
+ const int index = kOrder == StorageOrder::NCHW
+ ? (j / HxW * C + i) * HxW + j % HxW
+ : j * C + i;
+ ds_val += dY[index] * (X[index] - mean[i]) * rstd[i];
+ db_val += dY[index];
+ }
+ ds_val = BlockReduce<T>(ds_storage).Sum(ds_val);
+ db_val = BlockReduce<T>(db_storage).Sum(db_val);
+ if (threadIdx.x == 0) {
+#if __CUDA_ARCH__ >= 350
+ const T scale_x_rstd = __ldg(scale + i) * __ldg(rstd + i);
+ const T dscale_x_rstd = ds_val * __ldg(rstd + i);
+ dscale[i] = ds_val;
+ dbias[i] = db_val;
+ alpha[i] = scale_x_rstd;
+ beta[i] = -scale_x_rstd * dscale_x_rstd * inv_nhw;
+ gamma[i] =
+ scale_x_rstd * (__ldg(mean + i) * dscale_x_rstd - db_val) * inv_nhw;
+#else
+ const T scale_x_rstd = scale[i] * rstd[i];
+ const T dscale_x_rstd = ds_val * rstd[i];
+ dscale[i] = ds_val;
+ dbias[i] = db_val;
+ alpha[i] = scale_x_rstd;
+ beta[i] = -scale_x_rstd * dscale_x_rstd * inv_nhw;
+ gamma[i] = scale_x_rstd * (mean[i] * dscale_x_rstd - db_val) * inv_nhw;
+#endif
+ }
+ __syncthreads();
+ }
+}
+
+template <typename T, StorageOrder kOrder>
+__global__ void ComputeXGradientCUDAKernel(
+ const int size,
+ const int C,
+ const int HxW,
+ const T* dY,
+ const T* X,
+ const T* alpha,
+ const T* beta,
+ const T* gamma,
+ T* dX) {
+ CUDA_1D_KERNEL_LOOP(i, size) {
+ const int c = kOrder == StorageOrder::NCHW ? i / HxW % C : i % C;
+#if __CUDA_ARCH__ >= 350
+ dX[i] = __ldg(alpha + c) * __ldg(dY + i) + __ldg(beta + c) * __ldg(X + i) +
+ __ldg(gamma + c);
+#else
+ dX[i] = alpha[c] * dY[i] + beta[c] * X[i] + gamma[c];
+#endif
+ }
+}
+
+} // namespace
+
+template <>
+template <typename T>
+void SpatialBNOp<CUDAContext>::ComputeFusedParam(
+ const int C,
+ const T* scale,
+ const T* bias,
+ const T* mean,
+ const T* var,
+ T* alpha,
+ T* beta) {
+ ComputeFusedParamCUDAKernel<T>
+ <<<CAFFE_GET_BLOCKS(C),
+ CAFFE_CUDA_NUM_THREADS,
+ 0,
+ context_.cuda_stream()>>>(
+ C, static_cast<T>(epsilon_), scale, bias, mean, var, alpha, beta);
+}
+
+template <>
+template <typename T>
+void SpatialBNOp<CUDAContext>::ComputeBatchMoments(
+ const int N,
+ const int C,
+ const int HxW,
+ const T* batch_mean_sum,
+ const T* batch_var_sum,
+ T* mean,
+ T* var) {
+ const T scale = T(1) / static_cast<T>(num_batches_ * N * HxW);
+ ComputeBatchMomentsCUDAKernel<T>
+ <<<CAFFE_GET_BLOCKS(C),
+ CAFFE_CUDA_NUM_THREADS,
+ 0,
+ context_.cuda_stream()>>>(
+ C, scale, batch_mean_sum, batch_var_sum, mean, var);
+}
+
+template <>
+template <typename T>
+void SpatialBNOp<CUDAContext>::ComputeRunningMomentsAndFusedParam(
+ const int C,
+ const T* scale,
+ const T* bias,
+ const T* mean,
+ const T* var,
+ T* running_mean,
+ T* running_var,
+ T* rstd,
+ T* alpha,
+ T* beta) {
+ ComputeRunningMomentsAndFusedParamCUDAKernel<T>
+ <<<CAFFE_GET_BLOCKS(C),
+ CAFFE_CUDA_NUM_THREADS,
+ 0,
+ context_.cuda_stream()>>>(
+ C,
+ static_cast<T>(momentum_),
+ static_cast<T>(epsilon_),
+ scale,
+ bias,
+ mean,
+ var,
+ running_mean,
+ running_var,
+ rstd,
+ alpha,
+ beta);
+}
+
+template <>
+template <typename T>
+void SpatialBNGradientOp<CUDAContext>::
+ ComputeMultiBatchScaleBiasGradientsAndFusedParams(
+ const int N,
+ const int C,
+ const int HxW,
+ const T* scale,
+ const T* mean,
+ const T* rstd,
+ const T* dscale_sum,
+ const T* dbias_sum,
+ T* dscale,
+ T* dbias,
+ T* alpha,
+ T* beta,
+ T* gamma) {
+ const T inv_num_batches = T(1) / static_cast<T>(num_batches_);
+ const T inv_nhw = T(1) / static_cast<T>(N * HxW);
+ ComputeMultiBatchScaleBiasGradientsAndFusedParamsCUDAKernel<T>
+ <<<CAFFE_GET_BLOCKS(C),
+ CAFFE_CUDA_NUM_THREADS,
+ 0,
+ context_.cuda_stream()>>>(
+ C,
+ inv_num_batches,
+ inv_nhw,
+ scale,
+ mean,
+ rstd,
+ dscale_sum,
+ dbias_sum,
+ dscale,
+ dbias,
+ alpha,
+ beta,
+ gamma);
+}
+
+template <>
+template <typename T>
+void SpatialBNGradientOp<CUDAContext>::ComputeScaleBiasGradientsAndFusedParams(
+ const int N,
+ const int C,
+ const int HxW,
+ const T* dY,
+ const T* X,
+ const T* scale,
+ const T* mean,
+ const T* rstd,
+ T* dscale,
+ T* dbias,
+ T* alpha,
+ T* beta,
+ T* gamma) {
+ if (order_ == StorageOrder::NCHW) {
+ ComputeScaleBiasGradientsAndFusedParamsCUDAKernel<T, StorageOrder::NCHW>
+ <<<std::min(C, CAFFE_MAXIMUM_NUM_BLOCKS),
+ CAFFE_CUDA_NUM_THREADS,
+ 0,
+ context_.cuda_stream()>>>(
+ N,
+ C,
+ HxW,
+ dY,
+ X,
+ scale,
+ mean,
+ rstd,
+ dscale,
+ dbias,
+ alpha,
+ beta,
+ gamma);
+ } else {
+ ComputeScaleBiasGradientsAndFusedParamsCUDAKernel<T, StorageOrder::NHWC>
+ <<<std::min(C, CAFFE_MAXIMUM_NUM_BLOCKS),
+ CAFFE_CUDA_NUM_THREADS,
+ 0,
+ context_.cuda_stream()>>>(
+ N,
+ C,
+ HxW,
+ dY,
+ X,
+ scale,
+ mean,
+ rstd,
+ dscale,
+ dbias,
+ alpha,
+ beta,
+ gamma);
+ }
+}
+
+template <>
+template <typename T>
+void SpatialBNGradientOp<CUDAContext>::ComputeXGradient(
+ const int N,
+ const int C,
+ const int HxW,
+ const T* dY,
+ const T* X,
+ const T* alpha,
+ const T* beta,
+ const T* gamma,
+ T* dX) {
+ const int size = N * C * HxW;
+ if (order_ == StorageOrder::NCHW) {
+ ComputeXGradientCUDAKernel<T, StorageOrder::NCHW>
+ <<<CAFFE_GET_BLOCKS(size),
+ CAFFE_CUDA_NUM_THREADS,
+ 0,
+ context_.cuda_stream()>>>(
+ size, C, HxW, dY, X, alpha, beta, gamma, dX);
+ } else {
+ ComputeXGradientCUDAKernel<T, StorageOrder::NHWC>
+ <<<CAFFE_GET_BLOCKS(size),
+ CAFFE_CUDA_NUM_THREADS,
+ 0,
+ context_.cuda_stream()>>>(
+ size, C, HxW, dY, X, alpha, beta, gamma, dX);
+ }
+}
+
+} // namespace caffe2
projects / platform / upstream / pytorch.git / commitdiff