ProtoTask.h

Go to the documentation of this file.

 
#pragma once
 
#include <tvm/api.h>
 
#include <tvm/utils/AffineExpr.h>
#include <tvm/utils/internal/ProtoTaskDetails.h>
 
namespace tvm
{
 
namespace utils
{
 
template<constraint::Type T, typename FunT>
class ProtoTaskCommon;
 
template<constraint::Type T>
using ProtoTask = ProtoTaskCommon<T, FunctionPtr>;
 
template<constraint::Type T>
using LinearProtoTask = ProtoTaskCommon<T, LinearFunctionPtr>;
 
template<constraint::Type T, typename FunT>
class ProtoTaskCommon
{
public:
  ProtoTaskCommon(FunT f, const internal::RHS & rhs) : f_(f), rhs_(rhs)
  {
    if(rhs.type_ == internal::RHSType::Vector && f->rSize() != rhs.v_.size())
    {
      throw std::runtime_error("The vector you provided has not the correct size.");
    }
  }
 
  template<typename FunU>
  ProtoTaskCommon(const ProtoTaskCommon<T, FunU> & pt) : f_(pt.f_), rhs_(pt.rhs_)
  {}
 
  FunT f_;
  internal::RHS rhs_;
};
 
template<typename FunT>
class ProtoTaskCommon<constraint::Type::DOUBLE_SIDED, FunT>
{
public:
  ProtoTaskCommon(FunT f, const internal::RHS & l, const internal::RHS & u) : f_(f), l_(l), u_(u)
  {
    if(l.type_ == internal::RHSType::Vector && f->rSize() != l.v_.size())
    {
      throw std::runtime_error("The lower bound vector you provided has not the correct size.");
    }
    if(u.type_ == internal::RHSType::Vector && f->rSize() != u.v_.size())
    {
      throw std::runtime_error("The upper bound vector you provided has not the correct size.");
    }
  }
 
  template<typename FunU>
  ProtoTaskCommon(const ProtoTaskCommon<constraint::Type::DOUBLE_SIDED, FunU> & pt) : f_(pt.f_), l_(pt.l_), u_(pt.u_)
  {}
 
  FunctionPtr f_;
  internal::RHS l_;
  internal::RHS u_;
};
 
template<typename T, typename IfNotLinearFunction, typename IfLinearFunction>
using ProtoChoice =
    typename std::enable_if<std::is_base_of<tvm::function::abstract::Function, T>::value,
                            typename std::conditional<std::is_base_of<tvm::function::abstract::LinearFunction, T>::value,
                                                      IfLinearFunction,
                                                      IfNotLinearFunction>::type>::type;
 
using ProtoTaskEQ = ProtoTask<constraint::Type::EQUAL>;
using LinearProtoTaskEQ = LinearProtoTask<constraint::Type::EQUAL>;
template<typename T>
using ProtoTaskEQRet = ProtoChoice<T, tvm::utils::ProtoTaskEQ, tvm::utils::LinearProtoTaskEQ>;
 
using ProtoTaskLT = ProtoTask<constraint::Type::LOWER_THAN>;
using LinearProtoTaskLT = LinearProtoTask<constraint::Type::LOWER_THAN>;
template<typename T>
using ProtoTaskLTRet = ProtoChoice<T, tvm::utils::ProtoTaskLT, tvm::utils::LinearProtoTaskLT>;
 
using ProtoTaskGT = ProtoTask<constraint::Type::GREATER_THAN>;
using LinearProtoTaskGT = LinearProtoTask<constraint::Type::GREATER_THAN>;
template<typename T>
using ProtoTaskGTRet = ProtoChoice<T, tvm::utils::ProtoTaskGT, tvm::utils::LinearProtoTaskGT>;
 
using ProtoTaskDS = ProtoTask<constraint::Type::DOUBLE_SIDED>;
using LinearProtoTaskDS = LinearProtoTask<constraint::Type::DOUBLE_SIDED>;
template<typename T>
using ProtoTaskDSRet = ProtoChoice<T, tvm::utils::ProtoTaskDS, tvm::utils::LinearProtoTaskDS>;
 
} // namespace utils
 
} // namespace tvm
 
// clang-format off
template<typename F>
inline tvm::utils::ProtoTaskEQRet<F> operator==(std::shared_ptr<F> f, const tvm::utils::internal::RHS& rhs) { return { f, rhs }; }
template<typename F>
inline tvm::utils::ProtoTaskEQRet<F> operator==(const tvm::utils::internal::RHS& rhs, std::shared_ptr<F> f) { return { f, rhs }; }
template<typename F>
inline tvm::utils::ProtoTaskGTRet<F> operator>=(std::shared_ptr<F> f, const tvm::utils::internal::RHS& rhs) { return { f, rhs }; }
template<typename F>
inline tvm::utils::ProtoTaskLTRet<F> operator>=(const tvm::utils::internal::RHS& rhs, std::shared_ptr<F> f) { return { f, rhs }; }
template<typename F>
inline tvm::utils::ProtoTaskLTRet<F> operator<=(std::shared_ptr<F> f, const tvm::utils::internal::RHS& rhs) { return { f, rhs }; }
template<typename F>
inline tvm::utils::ProtoTaskGTRet<F> operator<=(const tvm::utils::internal::RHS& rhs, std::shared_ptr<F> f) { return { f, rhs }; }
 
inline tvm::utils::ProtoTaskDS operator>=(const tvm::utils::ProtoTaskLT& ptl, const tvm::utils::internal::RHS& rhs) { return { ptl.f_, rhs, ptl.rhs_ }; }
inline tvm::utils::ProtoTaskDS operator<=(const tvm::utils::ProtoTaskGT& ptg, const tvm::utils::internal::RHS& rhs) { return { ptg.f_, ptg.rhs_, rhs }; }
inline tvm::utils::LinearProtoTaskDS operator>=(const tvm::utils::LinearProtoTaskLT& ptl, const tvm::utils::internal::RHS& rhs) { return { ptl.f_, rhs, ptl.rhs_ }; }
inline tvm::utils::LinearProtoTaskDS operator<=(const tvm::utils::LinearProtoTaskGT& ptg, const tvm::utils::internal::RHS& rhs) { return { ptg.f_, ptg.rhs_, rhs }; }
// clang-format on
#define TVM_ID(x) std::make_shared<tvm::function::IdentityFunction>(x)
// clang-format off
inline tvm::utils::LinearProtoTaskEQ operator==(tvm::VariablePtr x, const tvm::utils::internal::RHS& rhs) { return TVM_ID(x) == rhs; }
inline tvm::utils::LinearProtoTaskEQ operator==(const tvm::utils::internal::RHS& rhs, tvm::VariablePtr x) { return TVM_ID(x) == rhs; }
inline tvm::utils::LinearProtoTaskGT operator>=(tvm::VariablePtr x, const tvm::utils::internal::RHS& rhs) { return TVM_ID(x) >= rhs; }
inline tvm::utils::LinearProtoTaskLT operator>=(const tvm::utils::internal::RHS& rhs, tvm::VariablePtr x) { return TVM_ID(x) <= rhs; }
inline tvm::utils::LinearProtoTaskLT operator<=(tvm::VariablePtr x, const tvm::utils::internal::RHS& rhs) { return TVM_ID(x) <= rhs; }
inline tvm::utils::LinearProtoTaskGT operator<=(const tvm::utils::internal::RHS& rhs, tvm::VariablePtr x) { return TVM_ID(x) >= rhs; }
// clang-format on
#undef TVM_ID
 
#define TVM_LIN(x) std::make_shared<tvm::function::BasicLinearFunction>(x)
 
// clang-format off
template<typename Derived>
inline tvm::utils::LinearProtoTaskEQ operator==(const tvm::utils::LinearExpr<Derived>& lin, const tvm::utils::internal::RHS& rhs) { return TVM_LIN(lin) == rhs; }
template<typename Derived>
inline tvm::utils::LinearProtoTaskEQ operator==(const tvm::utils::internal::RHS& rhs, const tvm::utils::LinearExpr<Derived>& lin) { return TVM_LIN(lin) == rhs; }
template<typename Derived>
inline tvm::utils::LinearProtoTaskGT operator>=(const tvm::utils::LinearExpr<Derived>& lin, const tvm::utils::internal::RHS& rhs) { return TVM_LIN(lin) >= rhs; }
template<typename Derived>
inline tvm::utils::LinearProtoTaskLT operator>=(const tvm::utils::internal::RHS& rhs, const tvm::utils::LinearExpr<Derived>& lin) { return TVM_LIN(lin) <= rhs; }
template<typename Derived>
inline tvm::utils::LinearProtoTaskLT operator<=(const tvm::utils::LinearExpr<Derived>& lin, const tvm::utils::internal::RHS& rhs) { return TVM_LIN(lin) <= rhs; }
template<typename Derived>
inline tvm::utils::LinearProtoTaskGT operator<=(const tvm::utils::internal::RHS& rhs, const tvm::utils::LinearExpr<Derived>& lin) { return TVM_LIN(lin) >= rhs; }
 
template<typename CstDerived, typename... Derived>
inline tvm::utils::LinearProtoTaskEQ operator==(const tvm::utils::AffineExpr<CstDerived, Derived...>& aff, const tvm::utils::internal::RHS& rhs) { return TVM_LIN(aff) == rhs; }
template<typename CstDerived, typename... Derived>
inline tvm::utils::LinearProtoTaskEQ operator==(const tvm::utils::internal::RHS& rhs, const tvm::utils::AffineExpr<CstDerived, Derived...>& aff) { return TVM_LIN(aff) == rhs; }
template<typename CstDerived, typename... Derived>
inline tvm::utils::LinearProtoTaskGT operator>=(const tvm::utils::AffineExpr<CstDerived, Derived...>& aff, const tvm::utils::internal::RHS& rhs) { return TVM_LIN(aff) >= rhs; }
template<typename CstDerived, typename... Derived>
inline tvm::utils::LinearProtoTaskLT operator>=(const tvm::utils::internal::RHS& rhs, const tvm::utils::AffineExpr<CstDerived, Derived...>& aff) { return TVM_LIN(aff) <= rhs; }
template<typename CstDerived, typename... Derived>
inline tvm::utils::LinearProtoTaskLT operator<=(const tvm::utils::AffineExpr<CstDerived, Derived...>& aff, const tvm::utils::internal::RHS& rhs) { return TVM_LIN(aff) <= rhs; }
template<typename CstDerived, typename... Derived>
inline tvm::utils::LinearProtoTaskGT operator<=(const tvm::utils::internal::RHS& rhs, const tvm::utils::AffineExpr<CstDerived, Derived...>& aff) { return TVM_LIN(aff) >= rhs; }
// clang-format on
#undef TVM_LIN