Method Compute

oneflow/user/kernels/argmax_kernel.cpp:29–58

 private:
  void Compute(user_op::KernelComputeContext* ctx) const override {
    const user_op::Tensor* in = ctx->Tensor4ArgNameAndIndex("in", 0);
    user_op::Tensor* out = ctx->Tensor4ArgNameAndIndex("out", 0);

    const int32_t elem_cnt = in->shape_view().elem_cnt();
    CHECK_GE(elem_cnt, 0);
    if (elem_cnt == 0) { return; }

    const T* in_ptr = in->dptr<T>();
    int64_t* out_ptr = out->mut_dptr<int64_t>();

    const int64_t instance_size = in->shape_view().At(in->shape_view().NumAxes() - 1);
    const int64_t instance_num = elem_cnt / instance_size;
    const int64_t num_thread =
        std::min(instance_num, (int64_t)Singleton<ThreadPool>::Get()->thread_num());
    const BalancedSplitter bs(instance_num, num_thread);
    BlockingCounter bc(num_thread);
    FOR_RANGE(int64_t, thread_id, 0, num_thread) {
      const Range range = bs.At(thread_id);
      Singleton<ThreadPool>::Get()->AddWork([=, &bc]() {
        FOR_RANGE(int64_t, i, range.begin(), range.end()) {
          const T* in_ptr_i = in_ptr + i * instance_size;
          out_ptr[i] =
              std::distance(in_ptr_i, std::max_element(in_ptr_i, in_ptr_i + instance_size));
        }
        bc.Decrease();
      });
    }
    bc.WaitForeverUntilCntEqualZero();
  }
  bool AlwaysComputeWhenAllOutputsEmpty() const override { return false; }
};