session.h

Go to the documentation of this file.

/*
 * Licensed to the Apache Software Foundation (ASF) under one
 * or more contributor license agreements.  See the NOTICE file
 * distributed with this work for additional information
 * regarding copyright ownership.  The ASF licenses this file
 * to you under the Apache License, Version 2.0 (the
 * "License"); you may not use this file except in compliance
 * with the License.  You may obtain a copy of the License at
 *
 *   http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing,
 * software distributed under the License is distributed on an
 * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
 * KIND, either express or implied.  See the License for the
 * specific language governing permissions and limitations
 * under the License.
 */
#ifndef TVM_RUNTIME_DISCO_SESSION_H_
#define TVM_RUNTIME_DISCO_SESSION_H_
 
#include <tvm/ffi/function.h>
#include <tvm/runtime/int_tuple.h>
#include <tvm/runtime/ndarray.h>
#include <tvm/runtime/object.h>
 
#include <queue>
#include <string>
#include <utility>
 
namespace tvm {
namespace runtime {
 
enum class DiscoAction : int32_t {
  kShutDown = 0,
  kKillReg = 1,
  kGetGlobalFunc = 2,
  kCallPacked = 3,
  kSyncWorker = 4,
  kCopyFromWorker0 = 5,
  kCopyToWorker0 = 6,
  kDebugGetFromRemote = 7,
  kDebugSetRegister = 8,
};
 
inline std::string DiscoAction2String(DiscoAction action) {
  switch (action) {
    case DiscoAction::kShutDown:
      return "kShutDown";
    case DiscoAction::kKillReg:
      return "kKillReg";
    case DiscoAction::kGetGlobalFunc:
      return "kGetGlobalFunc";
    case DiscoAction::kCallPacked:
      return "kCallPacked";
    case DiscoAction::kSyncWorker:
      return "kSyncWorker";
    case DiscoAction::kCopyFromWorker0:
      return "kCopyFromWorker0";
    case DiscoAction::kCopyToWorker0:
      return "kCopyToWorker0";
    case DiscoAction::kDebugGetFromRemote:
      return "kDebugGetFromRemote";
    case DiscoAction::kDebugSetRegister:
      return "kDebugSetRegister";
  }
  LOG(FATAL) << "ValueError: Unknown DiscoAction: " << static_cast<int>(action);
}
 
class SessionObj;
 
class DRefObj : public Object {
 public:
  inline ~DRefObj();
  inline ffi::Any DebugGetFromRemote(int worker_id);
  inline void DebugCopyFrom(int worker_id, ffi::AnyView source);
 
  static constexpr const char* _type_key = "runtime.disco.DRef";
  static constexpr const uint32_t _type_index = TypeIndex::kRuntimeDiscoDRef;
  static const constexpr bool _type_final = true;
  TVM_FFI_DECLARE_STATIC_OBJECT_INFO(DRefObj, Object);
 
  int64_t reg_id;
  ObjectRef session{nullptr};
 
 private:
  inline SessionObj* GetSession();
};
 
class DRef : public ObjectRef {
 public:
  TVM_DEFINE_MUTABLE_NOTNULLABLE_OBJECT_REF_METHODS(DRef, ObjectRef, DRefObj);
};
 
class SessionObj : public Object {
 public:
  virtual ~SessionObj() = default;
  template <typename... Args>
  DRef TVM_ALWAYS_INLINE CallPacked(const DRef& func, Args&&... args);
  TVM_DLL virtual DRef CallWithPacked(const ffi::PackedArgs& args) = 0;
  TVM_DLL virtual int64_t GetNumWorkers() = 0;
  TVM_DLL virtual DRef GetGlobalFunc(const std::string& name) = 0;
  TVM_DLL virtual void CopyFromWorker0(const NDArray& host_array, const DRef& remote_array) = 0;
  TVM_DLL virtual void CopyToWorker0(const NDArray& host_array, const DRef& remote_array) = 0;
  TVM_DLL virtual void SyncWorker(int worker_id) = 0;
  TVM_DLL virtual void Shutdown() = 0;
  TVM_DLL virtual void InitCCL(String ccl, IntTuple device_ids) = 0;
  TVM_DLL virtual ffi::Any DebugGetFromRemote(int64_t reg_id, int worker_id) = 0;
  TVM_DLL virtual void DebugSetRegister(int64_t reg_id, ffi::AnyView value, int worker_id) = 0;
 
  struct FFI;
  friend struct SessionObj::FFI;
  friend class DRefObj;
  static constexpr const char* _type_key = "runtime.disco.Session";
  TVM_DECLARE_BASE_OBJECT_INFO(SessionObj, Object);
 
 protected:
  virtual void DeallocReg(int reg_id) = 0;
};
 
class Session : public ObjectRef {
 public:
  TVM_DLL static Session ThreadedSession(int num_workers, int num_groups);
  TVM_DLL static Session ProcessSession(int num_workers, int num_groups,
                                        String process_pool_creator, String entrypoint);
 
  TVM_FFI_DEFINE_MUTABLE_OBJECT_REF_METHODS(Session, ObjectRef, SessionObj);
};
 
class DiscoChannel {
 public:
  virtual ~DiscoChannel() = default;
  virtual void Send(const ffi::PackedArgs& args) = 0;
  virtual ffi::PackedArgs Recv() = 0;
  virtual void Reply(const ffi::PackedArgs& args) = 0;
  virtual ffi::PackedArgs RecvReply() = 0;
};
 
class WorkerZeroData {
 public:
  std::queue<NDArray> host_arrays;
  std::mutex queue_mutex_;
};
 
// Implementation details
 
inline SessionObj* DRefObj::GetSession() {
  return const_cast<SessionObj*>(static_cast<const SessionObj*>(session.get()));
}
 
DRefObj::~DRefObj() {
  if (this->session.defined()) {
    GetSession()->DeallocReg(reg_id);
  }
}
 
ffi::Any DRefObj::DebugGetFromRemote(int worker_id) {
  return GetSession()->DebugGetFromRemote(this->reg_id, worker_id);
}
 
void DRefObj::DebugCopyFrom(int worker_id, ffi::AnyView value) {
  return GetSession()->DebugSetRegister(this->reg_id, value, worker_id);
}
 
template <typename... Args>
DRef SessionObj::CallPacked(const DRef& func, Args&&... args) {
  constexpr int offset = 3;
  constexpr int kNumArgs = offset + sizeof...(Args);
  ffi::AnyView packed_args[kNumArgs];
  ffi::PackedArgs::Fill(packed_args,
                        /*.0=*/static_cast<int>(DiscoAction::kCallPacked),  // action
                        /*.1=*/0,     // reg_id, which will be updated by this->CallWithPacked
                        /*.2=*/func,  // the function to be called
                        std::forward<Args>(args)...);
  return this->CallWithPacked(ffi::PackedArgs(packed_args, kNumArgs));
}
 
}  // namespace runtime
}  // namespace tvm
#endif  // TVM_RUNTIME_DISCO_SESSION_H_