Method Compute

oneflow/user/kernels/mode_kernel.cpp:34–87

 private:
  void Compute(user_op::KernelComputeContext* ctx) const override {
    const user_op::Tensor* in = ctx->Tensor4ArgNameAndIndex("input", 0);
    const int64_t num_axes = in->shape_view().NumAxes();
    const int64_t size = in->shape_view().elem_cnt();
    if (size == 0) return;
    const int64_t stride = in->shape_view().At(num_axes - 1);
    const int64_t instance_num = size / stride;
    user_op::Tensor* values = ctx->Tensor4ArgNameAndIndex("values", 0);
    user_op::Tensor* indices = ctx->Tensor4ArgNameAndIndex("indices", 0);
    user_op::Tensor* tmp_buffer = ctx->Tensor4ArgNameAndIndex("tmp_buffer", 0);

    auto memcpy = NewMemcpyPrimitive(ctx);
    CHECK(memcpy);
    memcpy->Launch(ctx->stream(), tmp_buffer->mut_dptr<void>(), in->dptr<void>(), size * sizeof(T));
    const int64_t thread_num =
        std::min(instance_num, (int64_t)Singleton<ThreadPool>::Get()->thread_num());
    const BalancedSplitter bs(instance_num, thread_num);
    BlockingCounter bc(thread_num);
    FOR_RANGE(int64_t, thread_id, 0, thread_num) {
      const Range range = bs.At(thread_id);
      Singleton<ThreadPool>::Get()->AddWork([=, &bc]() {
        FOR_RANGE(int64_t, i, range.begin(), range.end()) {
          T* in_ptr = tmp_buffer->mut_dptr<T>() + i * stride;
          T* val_ptr = values->mut_dptr<T>() + i;
          int64_t* ind_ptr = indices->mut_dptr<int64_t>() + i;
          std::vector<std::pair<T, int64_t>> elements(stride);
          T mode = 0;
          int64_t mode_idx = 0;
          int64_t temp_freq = 0;
          int64_t max_freq = 0;
          FOR_RANGE(int64_t, idx, 0, stride) {
            elements[idx] = std::make_pair(*(in_ptr + idx), idx);
          }
          std::sort(elements.begin(), elements.end(),
                    [=](const auto& i, const auto& j) { return i.first < j.first; });
          FOR_RANGE(int64_t, idx, 0, stride) {
            temp_freq++;
            if ((idx == stride - 1) || (elements[idx].first != elements[idx + 1].first)) {
              if (temp_freq > max_freq) {
                mode = elements[idx].first;
                mode_idx = elements[idx].second;
                max_freq = temp_freq;
              }
              temp_freq = 0;
            }
          }
          *val_ptr = mode;
          *ind_ptr = mode_idx;
        }
        bc.Decrease();
      });
    }
    bc.WaitForeverUntilCntEqualZero();
  }
  bool AlwaysComputeWhenAllOutputsEmpty() const override { return false; }
};