TritonMemResource.h

Go to the documentation of this file.

#ifndef HeterogeneousCore_SonicTriton_TritonMemResource
#define HeterogeneousCore_SonicTriton_TritonMemResource
 
#include <string>
#include <memory>
 
#include "grpc_client.h"
 
//forward declaration
template <typename IO>
class TritonData;
 
//base class for memory operations
template <typename IO>
class TritonMemResource {
public:
  TritonMemResource(TritonData<IO>* data, const std::string& name, size_t size);
  virtual ~TritonMemResource() {}
  uint8_t* addr() { return addr_; }
  size_t size() const { return size_; }
  virtual void close() {}
  //used for input
  virtual void copyInput(const void* values, size_t offset) {}
  //used for output
  virtual const uint8_t* copyOutput() { return nullptr; }
  virtual void set();
 
protected:
  //member variables
  TritonData<IO>* data_;
  std::string name_;
  size_t size_;
  uint8_t* addr_;
  bool closed_;
};
 
template <typename IO>
class TritonHeapResource : public TritonMemResource<IO> {
public:
  TritonHeapResource(TritonData<IO>* data, const std::string& name, size_t size);
  ~TritonHeapResource() override {}
  void copyInput(const void* values, size_t offset) override {}
  const uint8_t* copyOutput() override { return nullptr; }
  void set() override {}
};
 
template <typename IO>
class TritonCpuShmResource : public TritonMemResource<IO> {
public:
  TritonCpuShmResource(TritonData<IO>* data, const std::string& name, size_t size);
  ~TritonCpuShmResource() override;
  void close() override;
  void copyInput(const void* values, size_t offset) override {}
  const uint8_t* copyOutput() override { return nullptr; }
};
 
using TritonInputHeapResource = TritonHeapResource<triton::client::InferInput>;
using TritonInputCpuShmResource = TritonCpuShmResource<triton::client::InferInput>;
using TritonOutputHeapResource = TritonHeapResource<triton::client::InferRequestedOutput>;
using TritonOutputCpuShmResource = TritonCpuShmResource<triton::client::InferRequestedOutput>;
 
//avoid "explicit specialization after instantiation" error
template <>
void TritonInputHeapResource::copyInput(const void* values, size_t offset);
template <>
void TritonInputCpuShmResource::copyInput(const void* values, size_t offset);
template <>
const uint8_t* TritonOutputHeapResource::copyOutput();
template <>
const uint8_t* TritonOutputCpuShmResource::copyOutput();
 
#ifdef TRITON_ENABLE_GPU
#include "cuda_runtime_api.h"
 
template <typename IO>
class TritonGpuShmResource : public TritonMemResource<IO> {
public:
  TritonGpuShmResource(TritonData<IO>* data, const std::string& name, size_t size);
  ~TritonGpuShmResource() override;
  void close() override;
  void copyInput(const void* values, size_t offset) override {}
  const uint8_t* copyOutput() override { return nullptr; }
 
protected:
  int deviceId_;
  std::shared_ptr<cudaIpcMemHandle_t> handle_;
};
 
using TritonInputGpuShmResource = TritonGpuShmResource<triton::client::InferInput>;
using TritonOutputGpuShmResource = TritonGpuShmResource<triton::client::InferRequestedOutput>;
 
//avoid "explicit specialization after instantiation" error
template <>
void TritonInputGpuShmResource::copyInput(const void* values, size_t offset);
template <>
const uint8_t* TritonOutputGpuShmResource::copyOutput();
#endif
 
#endif