Program Listing for File aggregate_batching_pool.h

↰ Return to documentation for file (src/translator/aggregate_batching_pool.h)

#ifndef SRC_BERGAMOT_AGGREGATE_BATCHING_POOL_H_
#define SRC_BERGAMOT_AGGREGATE_BATCHING_POOL_H_

#include <memory>
#include <queue>

#include "data/types.h"
#include "translation_model.h"

namespace marian {
namespace bergamot {

template <class T>
struct HashPtr {
  size_t operator()(const std::shared_ptr<T>& t) const {
    size_t address = reinterpret_cast<size_t>(t.get());
    return std::hash<size_t>()(address);
  }
};

//
class AggregateBatchingPool {
 public:
  AggregateBatchingPool();

  size_t enqueueRequest(Ptr<TranslationModel> model, Ptr<Request> request);

  size_t generateBatch(Ptr<TranslationModel>& model, Batch& batch);

  void clear();

 private:
  std::unordered_set<std::shared_ptr<TranslationModel>, HashPtr<TranslationModel>> aggregateQueue_;
};

}  // namespace bergamot
}  // namespace marian

#endif  //  SRC_BERGAMOT_AGGREGATE_BATCHING_POOL_H_