Function BiasGradNCHW_SharedAtomics

tensorflow/core/kernels/bias_op_gpu.cu.cc:146–205

template <typename T>
__global__ void BiasGradNCHW_SharedAtomics(const T* output_backprop,
                                           T* bias_backprop, int32 batch,
                                           int32 bias_size, int32 image_size,
                                           int group_size) {
  // Initialize the shared memory.
  typedef typename AccumulatorType<T>::type AccT;
  const int32 kSDataSize = 32;
  __shared__ AccT s_data[kSDataSize];
  for (int32 index = threadIdx.x; index < kSDataSize; index += blockDim.x) {
    s_data[index] = AccT(0);
  }
  __syncthreads();

  // Accumulate all the values within this thread. They all have the same bias
  // index.
  int32 bias_index = blockIdx.x % bias_size;
  int32 group_index = blockIdx.x / bias_size;
  int32 total_count = batch * image_size;
  AccT sum(0);
  for (int32 index = group_index * blockDim.x + threadIdx.x;
       index < total_count; index += blockDim.x * group_size) {
    int32 image_offset = index % image_size;
    int32 batch = index / image_size;
    T val = ldg(output_backprop +
                (batch * bias_size + bias_index) * image_size + image_offset);
    sum += AccT(val);
  }

  // Write the accumulated sum in this thread to the shared memory. Each thread
  // shifts their write location to avoid bank conflict.
  int bias_offset = threadIdx.x % 32;
  GpuAtomicAdd(s_data + bias_offset, sum);
  __syncthreads();

  // Accumulate the results in the shared memory into the first element.
  // No syncthreads is needed since this is only in the same warp.
  int32 thread_index = threadIdx.x;
#if GOOGLE_CUDA
  if (thread_index < 32) {
    AccT data = s_data[thread_index];
    for (int32 delta = warpSize / 2; delta > 0; delta /= 2) {
      data += GpuShuffleXorSync(kCudaWarpAll, data, delta);
    }
    if (thread_index == 0) {
      GpuAtomicAdd(bias_backprop + bias_index, T(data));
    }
  }
#elif TENSORFLOW_USE_ROCM
  if (thread_index < 16) s_data[thread_index] += s_data[thread_index + 16];
  if (thread_index < 8) s_data[thread_index] += s_data[thread_index + 8];
  if (thread_index < 4) s_data[thread_index] += s_data[thread_index + 4];
  if (thread_index < 2) s_data[thread_index] += s_data[thread_index + 2];
  if (thread_index < 1) s_data[thread_index] += s_data[thread_index + 1];

  // The first thread writes out the accumulated result to the global location.
  if (thread_index == 0) {
    GpuAtomicAdd(bias_backprop + bias_index, T(s_data[0]));
  }