Function MultiMapCompute

tensorflow/core/kernels/unique_ali_op_util.h:485–620

template<typename TIndex, class HashMap>
void MultiMapCompute(OpKernelContext* context, const Tensor& input,
                     Tensor* idx, int64 axis, int64* uniq_size_out,
                     int32 num_buckets, int64 unique_ratio_hint, int num_sparse,
                     google::dense_hash_map<int, TIndex>* counter_map, Tensor* output) {
  auto Tin = input.vec<int64>();
  const int64 N = input.NumElements();

  auto idx_vec = idx->template vec<TIndex>();

  auto thread_pool =
    context->device()->tensorflow_cpu_worker_threads()->workers;

  // Parallel Step 0: Partition.
  int32 num_partitions = num_buckets;
  std::unique_ptr<int64[]> partitions{new int64[num_partitions * num_buckets]};

  static IdHash hasher;
  Partitioner map_parter(N, num_partitions);
  auto PartitionTask = [N, num_buckets, &Tin, &partitions, &map_parter, &idx_vec]
      (int32 task_id, int32 num_tasks) {
    auto st = Status::OK();
    int64* partition = partitions.get() + task_id * num_buckets;
    for (int64 i = 0; i < num_buckets; ++i) {
      partition[i] = -1;
    }
    const Range* range = map_parter.GetRange(task_id);

    for (int64 i = range->Start(); i < range->End(); ++i) {
      auto& id = Tin(i);
      if (unlikely(id == kPreseverdEmptyKey)) {
        st = errors::InvalidArgument(
            "Input id is preserved key of dense_hash_map, "
            "not supported: ", id);
        break;
      }
      int64 bucket = (hasher(id) >> 54) % num_buckets;
      idx_vec(i) = partition[bucket];
      partition[bucket] = i;
    }
    return st;
  };

  SummaryTaskRunner<Status, StatusSummaryUpdater> t0_runner(
      PartitionTask, Status::OK(), thread_pool, num_partitions);
  t0_runner.Run();
  OP_REQUIRES_OK(context, t0_runner.summary());

  // Parallel Step 1: Build hash maps.
  HashMap uniq_maps_ptr[num_buckets];
  HashMap* uniq_maps = uniq_maps_ptr;
  auto MapBuildTask = [N, num_partitions, &Tin, &partitions,
    uniq_maps, &idx_vec, unique_ratio_hint] (int32 task_id, int32 num_tasks) {
    auto& uniq = uniq_maps[task_id];
    HashMapInitializer<HashMap>::Reserve(&uniq, N / num_tasks / unique_ratio_hint);

    for (int64 k = 0; k < num_partitions; ++k) {
      int64* partition = partitions.get() + k * num_tasks + task_id;
      int64 next_idx = *partition;