Function LaunchRowReduction

tensorflow/core/kernels/reduction_gpu_kernels.cu.h:719–760

template <typename T, typename Op, typename OUT_T, typename IN_T>
void LaunchRowReduction(OpKernelContext* ctx, OUT_T out, IN_T in, int num_rows,
                        int num_cols, Op op, T init,
                        const gpuStream_t& cu_stream) {
  if (num_cols < 1024) {
    const int threads_per_block = 128;
    const int warps_per_block = threads_per_block / TF_RED_WARPSIZE;
    int num_blocks = (num_rows + warps_per_block - 1) / warps_per_block;

    TF_CHECK_OK(GpuLaunchKernel(RowReduceKernel<IN_T, OUT_T, Op>, num_blocks,
                                threads_per_block, 0, cu_stream, in, out,
                                num_rows, num_cols, op, init));
    return;
  }

  // setup segment offsets with counting and transform iterator
  RowOffset row_offset_op(num_cols);
  gpuprim::CountingInputIterator<int> counting_iter(0);
  gpuprim::TransformInputIterator<int, RowOffset,
                                  gpuprim::CountingInputIterator<int>>
      transform_iter(counting_iter, row_offset_op);

  size_t temp_storage_bytes = 0;
  auto reduce = [&](void* temp_storage_ptr) {
    auto success = gpuprim::DeviceSegmentedReduce::Reduce(
        temp_storage_ptr, temp_storage_bytes, in, out, num_rows, transform_iter,
        transform_iter + 1, op, init, cu_stream);

    OP_REQUIRES(ctx, success == 0,
                errors::Internal("CUB segmented reduce error",
                                 GpuGetErrorString(success)));
  };

  reduce(nullptr);  // Get required amount of temp storage.

  Tensor temp_storage;
  OP_REQUIRES_OK(
      ctx, ctx->allocate_temp(
               DT_INT8, TensorShape({static_cast<int64>(temp_storage_bytes)}),
               &temp_storage));

  reduce(temp_storage.flat<int8_t>().data());  // Do reduction.
}

template <typename T, typename Op, typename OUT_T, typename IN_T>
void LaunchColumnReduction_LTE16Cols(OpKernelContext* ctx, OUT_T out, IN_T in,