Method Compute

oneflow/user/kernels/eager_p_to_s_kernel.cpp:150–195

 private:
  void Compute(user_op::KernelComputeContext* ctx, user_op::OpKernelState*,
               const user_op::OpKernelCache* cache) const override {
    auto* kernel_cache = dynamic_cast<const EagerPToSOpKernelCache*>(cache);
    CHECK(kernel_cache != nullptr);
    const user_op::Tensor* in = ctx->Tensor4ArgNameAndIndex("in", 0);
    user_op::Tensor* out = ctx->Tensor4ArgNameAndIndex("out", 0);
    user_op::Tensor* tmp_buffer = ctx->Tensor4ArgNameAndIndex("tmp_buffer", 0);
    const void* in_ptr = in->dptr();
    void* tmp_buffer_ptr = tmp_buffer->mut_dptr();

    int64_t elem_cnt_of_this_chunk = kernel_cache->elem_cnt_of_this_chunk();
    const auto& sorted_elem_cnt2_in_tensor_slice_copier =
        kernel_cache->sorted_elem_cnt2_in_tensor_slice_copier();
    const auto& sorted_p2p_pair = kernel_cache->sorted_p2p_pair();
    CHECK_EQ(sorted_elem_cnt2_in_tensor_slice_copier.size(), sorted_p2p_pair.size());

    DeviceType device_type = ctx->device_type();

    std::unique_ptr<ep::primitive::Memset> memset_primitive =
        ep::primitive::NewPrimitive<ep::primitive::MemsetFactory>(device_type);
    CHECK(memset_primitive) << "Can not create Memset primitive for device type " << device_type;
    memset_primitive->Launch(ctx->stream(), out->mut_dptr(), 0,
                             elem_cnt_of_this_chunk * GetSizeOfDataType(out->data_type()));

    std::unique_ptr<ep::primitive::Add> add_primitive =
        ep::primitive::NewPrimitive<ep::primitive::AddFactory>(ctx->device_type(), in->data_type());
    CHECK(add_primitive);
    for (int64_t i = 0; i < sorted_p2p_pair.size(); ++i) {
      const auto& p2p_pair = sorted_p2p_pair.at(i);
      int64_t src = p2p_pair.first;
      int64_t dst = p2p_pair.second;
      if (GlobalProcessCtx::Rank() == src) {
        const auto& tensor_slice_copier = sorted_elem_cnt2_in_tensor_slice_copier.at(i).second;
        int64_t send_elem_cnt = sorted_elem_cnt2_in_tensor_slice_copier.at(i).first;
        tensor_slice_copier->Copy(ctx->stream(), tmp_buffer_ptr, in_ptr);
        CHECK_JUST(Send(reinterpret_cast<const void*>(tmp_buffer_ptr), send_elem_cnt,
                        in->data_type(), dst, device_type, ctx->stream()));
      }
      if (GlobalProcessCtx::Rank() == dst) {
        CHECK_JUST(Recv(tmp_buffer_ptr, elem_cnt_of_this_chunk, out->data_type(), src, device_type,
                        ctx->stream()));
        add_primitive->Launch(ctx->stream(), out->dptr(), tmp_buffer_ptr, out->mut_dptr(),
                              elem_cnt_of_this_chunk);
      }
    }
  }
  bool AlwaysComputeWhenAllOutputsEmpty() const override { return false; }
};