MLModelRunner.h

Go to the documentation of this file.

//===- MLModelRunner.h ---- ML model runner interface -----------*- C++ -*-===//
//
// Part of the MLCompilerBridge Project, under the Apache License v2.0 with LLVM
// Exceptions. See the LICENSE file for license information.
// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
//
// (Preliminary version adopted from MLModelRunner.h of LLVM 17.X)
//
//===----------------------------------------------------------------------===//
//===----------------------------------------------------------------------===//
 
#ifndef ML_MODEL_RUNNER_H
#define ML_MODEL_RUNNER_H
 
#include "SerDes/baseSerDes.h"
#include "SerDes/bitstreamSerDes.h"
#include "SerDes/jsonSerDes.h"
 
#include <cstdlib>
#include <future>
#include <memory>
#include <string>
#include <type_traits>
 
#ifndef C_LIBRARY
#include "SerDes/protobufSerDes.h"
#include "SerDes/tensorflowSerDes.h"
#endif
namespace MLBridge {
 
class MLModelRunner {
public:
  // Disallows copy and assign.
  MLModelRunner(const MLModelRunner &) = delete;
  MLModelRunner &operator=(const MLModelRunner &) = delete;
  virtual ~MLModelRunner() = default;
 
  template <typename T>
  typename std::enable_if<std::is_fundamental<T>::value, T>::type evaluate() {
    return *reinterpret_cast<T *>(evaluateUntyped());
  }
 
  template <typename T>
  typename std::enable_if<
      std::is_fundamental<typename std::remove_pointer<T>::type>::value,
      void>::type
  evaluate(T &data, size_t &dataSize) {
    using BaseType = typename std::remove_pointer<T>::type;
    void *res = evaluateUntyped();
    T ret = static_cast<T>(malloc(SerDes->getMessageLength()));
    memcpy(ret, res, SerDes->getMessageLength());
    dataSize = SerDes->getMessageLength() / sizeof(BaseType);
    data = ret;
  }
 
  enum class Kind : int { Unknown, Pipe, gRPC, ONNX, TFAOT };
 
  Kind getKind() const { return Type; }
  SerDesKind getSerDesKind() const { return SerDesType; }
 
  virtual void requestExit() = 0;
 
 
  template <typename U, typename T, typename... Types>
  void populateFeatures(const std::pair<U, T> &var1,
                        const std::pair<U, Types> &...var2) {
    SerDes->setFeature(var1.first, var1.second);
    populateFeatures(var2...);
  }
 
  template <typename U, typename T, typename... Types>
  void populateFeatures(const std::pair<U, T> &&var1,
                        const std::pair<U, Types> &&...var2) {
    SerDes->setFeature(var1.first, var1.second);
    populateFeatures(var2...);
  }
 
  void populateFeatures() {}
 
  void setRequest(void *request) { SerDes->setRequest(request); }
 
  void setResponse(void *response) { SerDes->setResponse(response); }
 
protected:
  MLModelRunner(Kind Type, SerDesKind SerDesType,
                llvm::LLVMContext *Ctx = nullptr)
      : Ctx(Ctx), Type(Type), SerDesType(SerDesType) {
    assert(Type != Kind::Unknown);
    initSerDes();
  }
 
  MLModelRunner(Kind Type, llvm::LLVMContext *Ctx = nullptr)
      : Ctx(Ctx), Type(Type), SerDesType(SerDesKind::Unknown) {
    SerDes = nullptr;
  };
 
  virtual void *evaluateUntyped() = 0;
 
  llvm::LLVMContext *Ctx;
  const Kind Type;
  const SerDesKind SerDesType;
 
protected:
  std::unique_ptr<BaseSerDes> SerDes;
 
private:
  void initSerDes() {
    switch (SerDesType) {
    case SerDesKind::Json:
      SerDes = std::make_unique<JsonSerDes>();
      break;
    case SerDesKind::Bitstream:
      SerDes = std::make_unique<BitstreamSerDes>();
      break;
#ifndef C_LIBRARY
    case SerDesKind::Protobuf:
      SerDes = std::make_unique<ProtobufSerDes>();
      break;
    case SerDesKind::Tensorflow:
      SerDes = std::make_unique<TensorflowSerDes>();
      break;
#endif
    case SerDesKind::Unknown:
      SerDes = nullptr;
      break;
    }
  }
};
} // namespace MLBridge
 
#endif // LLVM_MLMODELRUNNER_H