Refine acc actor (#6444)

Co-authored-by: oneflow-ci-bot <[email protected]>

cmake/cfg.cmake

@@ -45,7 +45,6 @@ function(GENERATE_CFG_AND_PYBIND11_CPP SRCS HDRS PYBIND_SRCS ROOT_DIR)
       oneflow/core/job/sbp_parallel.proto
       oneflow/core/graph/boxing/collective_boxing.proto
       oneflow/core/register/blob_desc.proto
-      oneflow/core/register/pod.proto
       oneflow/core/job/scope.proto
       oneflow/core/job/mirrored_parallel.proto
       oneflow/core/operator/op_attribute.proto
@@ -100,7 +99,6 @@ function(GENERATE_CFG_AND_PYBIND11_CPP SRCS HDRS PYBIND_SRCS ROOT_DIR)
       oneflow/core/operator/interface_blob_conf.proto
       oneflow/core/common/shape.proto
       oneflow/core/register/blob_desc.proto
-      oneflow/core/register/pod.proto
       oneflow/core/operator/op_conf.proto
   )
 

oneflow/core/actor/acc_actor.cpp

@@ -23,18 +23,14 @@ class AccActor final : public Actor {
   AccActor() = default;
   ~AccActor() override = default;
 
-  using Actor::Init;
-
  private:
   void Act() override;
   void VirtualAsyncSendNaiveProducedRegstMsgToConsumer() override;
 
   void VirtualActorInit(const TaskProto& proto) override;
-  void Init(const TaskProto&, int32_t max_acc_cnt);
 
-  std::function<void(DeviceCtx*, void* dst, const void* src, size_t)> cpy_func_;
-  int32_t acc_cnt_;
-  int32_t max_acc_cnt_;
+  int32_t acc_cnt_{};
+  int32_t max_acc_cnt_{};
 };
 
 void AccActor::VirtualActorInit(const TaskProto& proto) {
@@ -45,44 +41,21 @@ void AccActor::VirtualActorInit(const TaskProto& proto) {
                                     ->RegstDesc4RegstDescId(Name2SoleRegstDescId("out"))
                                     .data_regst_time_shape();
   CHECK_GE(in_time_shape.elem_cnt(), out_time_shape.elem_cnt());
-  Init(proto, in_time_shape.elem_cnt() / out_time_shape.elem_cnt());
-}
-
-void AccActor::Init(const TaskProto& task_proto, int32_t max_acc_cnt) {
-  using namespace std::placeholders;
-  if (GetDeviceType() == DeviceType::kCPU) {
-    cpy_func_ = std::bind(Memcpy<DeviceType::kCPU>, _1, _2, _3, _4);
-  } else {
-#ifdef WITH_CUDA
-    cpy_func_ = std::bind(Memcpy<DeviceType::kGPU>, _1, _2, _3, _4);
-#else
-    UNIMPLEMENTED();
-#endif
-  }
-  OF_SET_MSG_HANDLER(&AccActor::HandlerNormal);
+  max_acc_cnt_ = in_time_shape.elem_cnt() / out_time_shape.elem_cnt();
   acc_cnt_ = 0;
-  max_acc_cnt_ = max_acc_cnt;
+  OF_SET_MSG_HANDLER(&AccActor::HandlerNormal);
 }
 
 void AccActor::Act() {
-  Regst* out_regst = GetNaiveCurWriteable("out");
-  Regst* in_regst = GetNaiveCurReadable("in");
   if (acc_cnt_ == 0) {
+    Regst* out_regst = GetNaiveCurWriteable("out");
+    Regst* in_regst = GetNaiveCurReadable("in");
     const Blob* in_blob = in_regst->GetMutSoleBlob();
     Blob* out_blob = out_regst->GetMutSoleBlob();
-    if (GetDeviceType() == DeviceType::kCPU) {
-      Memcpy<DeviceType::kCPU>(mut_device_ctx().get(), out_blob->ForceMutDptr(), in_blob->dptr(),
-                               out_blob->ByteSizeOfBlobBody());
-    } else if (GetDeviceType() == DeviceType::kGPU) {
-#ifdef WITH_CUDA
-      Memcpy<DeviceType::kGPU>(mut_device_ctx().get(), out_blob->ForceMutDptr(), in_blob->dptr(),
-                               out_blob->ByteSizeOfBlobBody());
-#else
-      UNIMPLEMENTED();
-#endif
-    } else {
-      UNIMPLEMENTED();
-    }
+    const size_t size = in_blob->ByteSizeOfBlobBody();
+    CHECK_EQ(out_blob->ByteSizeOfBlobBody(), size);
+    AutoMemcpy(mut_device_ctx().get(), out_blob->ForceMutDptr(), in_blob->dptr(), size,
+               out_blob->mem_case(), in_blob->mem_case());
   } else {
     AsyncLaunchKernel();
   }

oneflow/core/actor/actor.cpp

@@ -297,10 +297,6 @@ void Actor::ForEachProducedRegst(const std::function<void(Regst*)>& Handler) con
   }
 }
 
-DeviceType Actor::GetDeviceType() const {
-  return Global<IDMgr>::Get()->GetDeviceTypeFromActorId(actor_id_);
-}
-
 int64_t Actor::Name2SoleRegstDescId(const std::string& name) const {
   auto find_it = name2regst_desc_id_.find(name);
   if (find_it != name2regst_desc_id_.end()) {

oneflow/core/actor/actor.h

@@ -68,7 +68,6 @@ class Actor : public ActorBase {
   const ParallelContext* parallel_ctx() const { return parallel_ctx_.get(); }
   bool ReceiveAllEordMsg() const { return remaining_eord_cnt_ == 0; }
   bool ReceiveEordMsg(int64_t regst_desc_id) const;
-  DeviceType GetDeviceType() const;
   virtual void VirtualActorInit(const TaskProto&) {}
   int64_t Name2SoleRegstDescId(const std::string& name) const;
   const std::vector<int64_t>& Name2RegstDescIds(const std::string& name) const;

oneflow/core/framework/op_arg_util.h

@@ -25,7 +25,6 @@ limitations under the License.
 #include "oneflow/core/common/shape.cfg.h"
 #include "oneflow/core/register/logical_blob_id.cfg.h"
 #include "oneflow/core/operator/interface_blob_conf.cfg.h"
-#include "oneflow/core/register/pod.cfg.h"
 #include "oneflow/core/register/blob_desc.cfg.h"
 #include "oneflow/core/operator/op_node_signature.cfg.h"
 #include "oneflow/core/job/parallel_signature.cfg.h"
@@ -42,6 +41,9 @@ class OpArgBlobAttribute {
                      const std::string& logical_blob_name);
 
   OpArgBlobAttribute(const OpArgBlobAttribute& op_arg_blob_attr) = default;
+  OpArgBlobAttribute(OpArgBlobAttribute&& op_arg_blob_attr) = delete;
+  OpArgBlobAttribute& operator=(const OpArgBlobAttribute&) = delete;
+  OpArgBlobAttribute& operator=(OpArgBlobAttribute&&) = delete;
   virtual ~OpArgBlobAttribute() = default;
 
   std::shared_ptr<cfg::BlobDescProto> blob_desc() const;
@@ -78,6 +80,9 @@ class OpArgParallelAttribute {
                          const std::shared_ptr<cfg::OptMirroredParallel>& opt_mirrored_parallel);
 
   OpArgParallelAttribute(const OpArgParallelAttribute& op_arg_para_attr) = default;
+  OpArgParallelAttribute(OpArgParallelAttribute&& op_arg_blob_attr) = delete;
+  OpArgParallelAttribute& operator=(const OpArgParallelAttribute&) = delete;
+  OpArgParallelAttribute& operator=(OpArgParallelAttribute&&) = delete;
   virtual ~OpArgParallelAttribute() = default;
 
   std::shared_ptr<ParallelDesc> parallel_desc_symbol() const;

oneflow/core/job/id_manager.cpp

@@ -19,10 +19,6 @@ limitations under the License.
 
 namespace oneflow {
 
-DeviceType IDMgr::GetDeviceTypeFromActorId(int64_t actor_id) const {
-  return DeserializeTaskIdFromInt64(actor_id).stream_id().device_id().device_type();
-}
-
 int64_t IDMgr::MachineId4ActorId(int64_t actor_id) const {
   // TODO: change this inferface semantics, rank does not indicate machine_id in multi-client
   return DeserializeTaskIdFromInt64(actor_id).stream_id().device_id().rank();

oneflow/core/job/id_manager.h

@@ -35,7 +35,6 @@ class IDMgr final {
   int64_t NewChunkId() { return chunk_id_count_++; }
 
   // Runtime
-  DeviceType GetDeviceTypeFromActorId(int64_t actor_id) const;
   int64_t MachineId4ActorId(int64_t actor_id) const;
   int64_t ThrdId4ActorId(int64_t actor_id) const;