Method operator()

tensorflow/core/kernels/segment_reduction_ops_gpu.cu.cc:223–297

template <typename T, typename Index, typename InitialValueF,
          typename EmptySegmentValueF, typename ReductionF>
void SegmentReductionFunctor<
    T, Index, InitialValueF, EmptySegmentValueF,
    ReductionF>::operator()(OpKernelContext* ctx, const GPUDevice& d,
                            const Index output_rows,
                            const TensorShape& segment_ids_shape, bool is_mean,
                            typename TTypes<Index>::ConstFlat segment_ids,
                            const Index data_size, const T* data,
                            typename TTypes<T, 2>::Tensor output) {
  if (output.size() == 0) {
    return;
  }

  // Launch kernel(s) to compute sorted segment reduction.
  // Notes:
  // *) 'input_total_size' is the total number of elements to process.
  // *) 'segment_ids.shape' is a prefix of data's shape.
  // *) 'input_outer_dim_size' is the total number of segments to process.
  const Index input_total_size = data_size;
  const Index input_outer_dim_size = segment_ids.dimension(0);
  const Index input_inner_dim_size = input_total_size / input_outer_dim_size;
  const Index num_segments = output.size() / input_inner_dim_size;

  // Set 'output' to initial value.
  GpuLaunchConfig config = GetGpuLaunchConfig(output.size(), d);
  const T initial_value = InitialValueF()();
  TF_CHECK_OK(GpuLaunchKernel(SetToValue<T>, config.block_count,
                              config.thread_per_block, 0, d.stream(),
                              output.size(), output.data(), initial_value));
  if (data_size == 0 || segment_ids_shape.num_elements() == 0) {
    return;
  }

  const int OuterDimTileSize = 8;

  const Index input_outer_dim_num_stripe =
      Eigen::divup(input_outer_dim_size, Index(OuterDimTileSize));

  const Index total_stripe_count =
      input_inner_dim_size * input_outer_dim_num_stripe;

  config = GetGpuLaunchConfig(total_stripe_count, d);
  TF_CHECK_OK(GpuLaunchKernel(
      SortedSegmentReductionCustomKernel<
          T, Index, OuterDimTileSize,
          typename ReduceUpdateOpFor<ReductionF>::nonatomic_op,
          typename ReduceUpdateOpFor<ReductionF>::atomic_op>,
      config.block_count, config.thread_per_block, 0, d.stream(),
      input_outer_dim_size, input_inner_dim_size, output_rows,
      segment_ids.data(), data, output.data(), total_stripe_count,
      initial_value));

  const T empty_value = EmptySegmentValueF()();
  if (is_mean || initial_value != empty_value) {
    Tensor segment_offsets;
    OP_REQUIRES_OK(ctx, ctx->allocate_temp(DataTypeToEnum<Index>::value,
                                           TensorShape({num_segments + 1}),
                                           &segment_offsets));
    Index* segment_offsets_ptr = segment_offsets.flat<Index>().data();