Method ForwardDataContent

oneflow/core/kernel/collective_boxing_kernels.cpp:65–87

}

void CollectiveBoxingGenericKernel::ForwardDataContent(KernelContext* ctx) const {
  RequestHandle* request_handle =
      CHECK_NOTNULL(dynamic_cast<CollectiveBoxingKernelState*>(ctx->state().get()))
          ->request_handle();
  const void* send_buff = nullptr;
  void* recv_buff = nullptr;
  const RankDesc& rank_desc = this->op_conf().collective_boxing_generic_conf().rank_desc();
  const DataType data_type = rank_desc.op_desc().data_type();
  if (GenericOpHasInput(rank_desc)) {
    const Blob* in = ctx->BnInOp2Blob("in");
    CHECK_EQ(in->data_type(), data_type);
    CHECK(in->shape() == ShapeView(GenericOpGetInputShape(rank_desc)));
    send_buff = in->dptr();
  }
  if (GenericOpHasOutput(rank_desc)) {
    Blob* out = ctx->BnInOp2Blob("out");
    CHECK_EQ(out->data_type(), data_type);
    CHECK(out->shape() == ShapeView(GenericOpGetOutputShape(rank_desc)));
    recv_buff = out->mut_dptr();
  }
  auto* actor_ctx = GetCollectiveBoxingActorContext(ctx);
  actor_ctx->Schedule(request_handle, send_buff, recv_buff);
}

REGISTER_KERNEL(OperatorConf::kCollectiveBoxingGenericConf, CollectiveBoxingGenericKernel);