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/container/shape.h>
#include <tvm/ffi/function.h>
#include <tvm/runtime/tensor.h>
 
#include <mutex>
#include <queue>
#include <string>
#include <utility>
 
namespace tvm {
namespace runtime {
 
constexpr int32_t kRuntimeDiscoDRef = TVMFFITypeIndex::kTVMFFIDynObjectBegin - 14;
 
static_assert(kRuntimeDiscoDRef >= TVMFFITypeIndex::kTVMFFIStaticObjectEnd &&
                  kRuntimeDiscoDRef < TVMFFITypeIndex::kTVMFFIDynObjectBegin,
              "kRuntimeDiscoDRef must live in the static custom-index slot range");
 
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";
  }
  TVM_FFI_THROW(ValueError) << "Unknown DiscoAction: " << static_cast<int>(action);
}
 
class SessionObj;
 
class DRefObj : public ffi::Object {
 public:
  inline ~DRefObj();
  inline ffi::Any DebugGetFromRemote(int worker_id);
  inline void DebugCopyFrom(int worker_id, ffi::AnyView source);
 
  static constexpr const uint32_t _type_index = kRuntimeDiscoDRef;
  static const constexpr bool _type_final = true;
  static constexpr const bool _type_mutable = true;
  TVM_FFI_DECLARE_OBJECT_INFO_STATIC("runtime.disco.DRef", DRefObj, ffi::Object);
 
  int64_t reg_id;
  ffi::ObjectRef session{nullptr};
 
 private:
  inline SessionObj* GetSession();
};
 
class DRef : public ffi::ObjectRef {
 public:
  explicit DRef(ffi::ObjectPtr<DRefObj> data) : ffi::ObjectRef(data) {
    TVM_FFI_ICHECK(data != nullptr);
  }
  TVM_FFI_DEFINE_OBJECT_REF_METHODS_NOTNULLABLE(DRef, ffi::ObjectRef, DRefObj);
};
 
class SessionObj : public ffi::Object {
 public:
  virtual ~SessionObj() = default;
  template <typename... Args>
  TVM_FFI_INLINE DRef CallPacked(const DRef& func, Args&&... args);
  TVM_RUNTIME_DLL virtual DRef CallWithPacked(const ffi::PackedArgs& args) = 0;
  TVM_RUNTIME_DLL virtual int64_t GetNumWorkers() = 0;
  TVM_RUNTIME_DLL virtual DRef GetGlobalFunc(const std::string& name) = 0;
  TVM_RUNTIME_DLL virtual void CopyFromWorker0(const ffi::Tensor& host_array,
                                               const DRef& remote_array) = 0;
  TVM_RUNTIME_DLL virtual void CopyToWorker0(const ffi::Tensor& host_array,
                                             const DRef& remote_array) = 0;
  TVM_RUNTIME_DLL virtual void SyncWorker(int worker_id) = 0;
  TVM_RUNTIME_DLL virtual void Shutdown() = 0;
  TVM_RUNTIME_DLL virtual void InitCCL(ffi::String ccl, ffi::Shape device_ids) = 0;
  TVM_RUNTIME_DLL virtual ffi::Any DebugGetFromRemote(int64_t reg_id, int worker_id) = 0;
  TVM_RUNTIME_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 bool _type_mutable = true;
  TVM_FFI_DECLARE_OBJECT_INFO("runtime.disco.Session", SessionObj, ffi::Object);
 
 protected:
  virtual void DeallocReg(int reg_id) = 0;
};
 
class Session : public ffi::ObjectRef {
 public:
  TVM_RUNTIME_DLL static Session ThreadedSession(int num_workers, int num_groups);
  TVM_RUNTIME_DLL static Session ProcessSession(int num_workers, int num_groups,
                                                ffi::String process_pool_creator,
                                                ffi::String entrypoint);
 
  explicit Session(ffi::UnsafeInit tag, ffi::ObjectPtr<SessionObj> n) : ffi::ObjectRef(tag) {
    data_ = std::move(n);
  }
 
  TVM_FFI_DEFINE_OBJECT_REF_METHODS_NOTNULLABLE(Session, ffi::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<ffi::Tensor> 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_