Function BiasGrad2DInternal

tensorflow/core/kernels/bias_grad_ali_op_cpu.h:336–392

template <typename T>
void BiasGrad2DInternal(const CPUDevice& d, typename TTypes<T>::ConstFlat input,
                        Eigen::DSizes<int, 2>& two_dims,
                        typename TTypes<T>::Flat output) {
  const int sum_size = two_dims[0];
  const int channel = two_dims[1];
  // NOTE(zycao): This is only threads number of Eigen threadpool. In MKL
  // handeled cases, threads number for computing will be decided by OpenMP.
  const T* in = input.data();
  T* out = output.data();

  // NOTE(zycao): These conditions are based on modern CPU architechure
  // features and verified by batch of tests on CPU. They are expected to
  // make positive impact on most cases.
#define CPU_CACHE_LINE_SIZE 64
#define HALF_L1_CACHE_SIZE 16384

  const int num_threads = d.numThreads();

  // Small cases would be explicitly done by single thread.
  if (sizeof(T) * sum_size * channel <= HALF_L1_CACHE_SIZE ||
      sum_size / num_threads <= 2 || num_threads == 1) {
    SumIntoOneRow(in, sum_size, channel, channel, out);
    return;
  }
  // Seperate the array by rows and parallel sum into temp array,
  // then sum the temp array in to output vector.
  std::vector<std::vector<T>> sum_vec(num_threads + 1);

  auto work_on_rows = [&in, &sum_vec, d, channel]
      (int64 start, int64 end) {
    // If running in caller thread, currentThreadId would return -1.
    int tid = d.currentThreadId() + 1;
    std::vector<T>& vec = sum_vec[tid];
    if (vec.empty()) vec.resize(channel, static_cast<T>(0));
    SumIntoOneRow(&(in[start * channel]), end - start, channel, channel,
                  vec.data(), false);
  };
  auto cost = Eigen::TensorOpCost(sizeof(T) * channel, // ld bytes
                                  sizeof(T) * channel, // st bytes
                                  channel); // compute cycles
  d.parallelFor(sum_size, cost, work_on_rows);

  // Sum temp array to output vector.
  for (int j = 0; j < channel; ++j) {
    out[j] = static_cast<T>(0);
  }
  for (int i = 0; i < num_threads + 1; ++i) {
    std::vector<T>& vec = sum_vec[i];
    if (!vec.empty()) {
      for (int j = 0; j < channel; ++j) {
        out[j] += vec[j];
      }
    }
  }
#undef CPU_CACHE_LINE_SIZE
#undef HALF_L1_CACHE_SIZE
}