Method ParallelSplit

tensorflow/core/kernels/string_split_op.cc:202–324

  }

  void ParallelSplit(OpKernelContext* ctx,
                      const Eigen::TensorMap<
                        Eigen::Tensor<const string, 1, 1, long>,
                        16, Eigen::MakePointer> input_vec,
                      const int64 batch_size, const string& delimiter) {
    ThreadPool* thread_pool =
    ctx->device()->tensorflow_cpu_worker_threads()->workers;
    const int64 num_threads = thread_pool->NumThreads() + 1;

    std::vector<int64> num_indices(batch_size);
    num_indices[0] = 0;

    std::vector<WorkerInfo> w_array;
    for (int i = 0; i < num_threads; i++) {
      WorkerInfo w(num_threads, batch_size);
      w_array.emplace_back(w);
    }
    std::vector<std::vector<int64>> id_to_worker(batch_size);

    thread_pool->ParallelForWithWorkerId(
      batch_size,
      element_cost_,
      [&w_array, &id_to_worker, &input_vec,
      &delimiter, ctx, this, &num_indices]
      (int64 start, int64 end, int64 worker_id){
        int64 position_in_worker = 0;
        for (int64 i = start; i < end; ++i) {
          std::vector<StringPiece> parts =
              skip_empty_ ?
               Split(input_vec(i), delimiter, str_util::SkipEmpty())
               : Split(input_vec(i), delimiter, str_util::AllowEmpty());
          int64 n_entries = parts.size();
          id_to_worker[i].emplace_back(worker_id);
          id_to_worker[i].emplace_back(w_array[worker_id].counter_for_thread);
          id_to_worker[i].emplace_back(w_array[worker_id].output_size);
          num_indices[i] = n_entries;
          position_in_worker += n_entries;
          w_array[worker_id].num_indices_buffer[
              w_array[worker_id].counter_for_thread] = n_entries;
          w_array[worker_id].output_size += n_entries;
          w_array[worker_id].max_num_entries =
              std::max(w_array[worker_id].max_num_entries, n_entries);
          w_array[worker_id].tokens_buffer.insert(
              w_array[worker_id].tokens_buffer.end(),
              std::make_move_iterator(parts.begin()),
              std::make_move_iterator(parts.end()));
          w_array[worker_id].counter_for_thread++;
       }
     });

    int64 output_size = 0;
    int64 max_num_entries = 0;
    for (int i = 0; i < num_threads; i++) {
      output_size += w_array[i].output_size;
      max_num_entries = std::max(w_array[i].max_num_entries, max_num_entries);
    }

    std::vector<int64> id_to_index(batch_size);