Method Compute

oneflow/user/kernels/eager_p_to_b_kernel.cpp:81–119

 private:
  void Compute(user_op::KernelComputeContext* ctx, user_op::OpKernelState*,
               const user_op::OpKernelCache* cache) const override {
    auto* kernel_cache = dynamic_cast<const EagerPToBOpKernelCache*>(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();

    const int64_t total_elem_cnt = ctx->Attr<Shape>("shape").elem_cnt();
    const auto& p2p_pair = kernel_cache->p2p_pair();

    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,
                             total_elem_cnt * 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 (const auto& pair : p2p_pair) {
      int64_t src = pair.first;
      int64_t dst = pair.second;

      if (GlobalProcessCtx::Rank() == src) {
        CHECK_JUST(Send(in_ptr, total_elem_cnt, in->data_type(), dst, device_type, ctx->stream()));
      }
      if (GlobalProcessCtx::Rank() == dst) {
        CHECK_JUST(Recv(tmp_buffer_ptr, total_elem_cnt, out->data_type(), src, device_type,
                        ctx->stream()));
        add_primitive->Launch(ctx->stream(), out->dptr(), tmp_buffer_ptr, out->mut_dptr(),
                              total_elem_cnt);
      }
    }
  };
  bool AlwaysComputeWhenAllOutputsEmpty() const override { return false; }
};