AffineExpr.h

Go to the documentation of this file.

 
#pragma once
 
#include <tvm/api.h>
#include <tvm/defs.h>
 
#include <tvm/Variable.h>
#include <tvm/internal/meta.h>
#include <tvm/utils/internal/AffineExprDetail.h>
 
#include <Eigen/Core>
 
#include <tuple>
 
namespace tvm
{
namespace utils
{
template<typename Derived>
class LinearExpr
{
public:
  LinearExpr(const Eigen::MatrixBase<Derived> & matrix, const VariablePtr & v) : matrix_(matrix.derived()), var_(v)
  { assert(matrix.cols() == v->size()); }
 
  template<class T = Derived, typename std::enable_if_t<std::is_same_v<T, internal::IdentityType>, int> = 0>
  LinearExpr(const tvm::VariablePtr & v) : matrix_(Eigen::MatrixXd::Identity(v->size(), v->size())), var_(v)
  {}
 
  template<class T = Derived, typename std::enable_if_t<std::is_same_v<T, internal::MultIdentityType>, int> = 0>
  LinearExpr(double a, const tvm::VariablePtr & v)
  : matrix_(a * Eigen::MatrixXd::Identity(v->size(), v->size())), var_(v)
  {}
 
  template<class T = Derived, typename std::enable_if_t<std::is_same_v<T, internal::MinusIdentityType>, int> = 0>
  LinearExpr(const tvm::VariablePtr & v) : matrix_(-Eigen::MatrixXd::Identity(v->size(), v->size())), var_(v)
  {}
 
  const Derived & matrix() const { return matrix_; };
  const VariablePtr & variable() const { return var_; }
 
private:
  typename internal::RefSelector_t<Derived> matrix_;
  const VariablePtr & var_;
};
 
template<typename CstDerived, typename... Derived>
class AffineExpr
{
public:
  template<class T = CstDerived, typename std::enable_if_t<std::is_same_v<T, internal::NoConstant>, int> = 0>
  AffineExpr(const LinearExpr<Derived> &... linear)
  : linear_(std::forward_as_tuple(linear...)), constant_(internal::NoConstant())
  {}
 
  template<class T = CstDerived, typename std::enable_if_t<!std::is_same_v<T, internal::NoConstant>, int> = 0>
  AffineExpr(const Eigen::MatrixBase<CstDerived> & constant, const LinearExpr<Derived> &... linear)
  : linear_(std::forward_as_tuple(linear...)), constant_(constant.derived())
  { EIGEN_STATIC_ASSERT_VECTOR_ONLY(CstDerived) }
 
  template<class T = CstDerived, typename std::enable_if_t<std::is_same_v<T, internal::NoConstant>, int> = 0>
  AffineExpr(const internal::NoConstant & constant, const std::tuple<LinearExpr<Derived>...> & linear)
  : linear_(linear), constant_(constant)
  {}
 
  template<class T = CstDerived, typename std::enable_if_t<!std::is_same_v<T, internal::NoConstant>, int> = 0>
  AffineExpr(const Eigen::MatrixBase<CstDerived> & constant, const std::tuple<LinearExpr<Derived>...> & linear)
  : linear_(linear), constant_(constant.derived())
  { EIGEN_STATIC_ASSERT_VECTOR_ONLY(CstDerived) }
 
  const std::tuple<LinearExpr<Derived>...> & linear() const { return linear_; }
  const CstDerived & constant() const { return constant_; }
 
private:
  // ConstantType::Type is NoConstant in case CstType==NoConstant and ref_selector<CstDerived> otherwise.
  // Use of ref_selector<CstDerived>: for a vector, we need to keep a const ref, while for a vector
  // expression we need to take a copy of the expression (same use as in e.g. CWiseBinaryOp).
  template<typename T>
  struct ConstantType
  {
    using Type = typename internal::RefSelector_t<CstDerived>;
  };
 
  std::tuple<LinearExpr<Derived>...> linear_;
  typename ConstantType<CstDerived>::Type constant_;
};
 
template<typename CstDerived, typename... Derived>
AffineExpr<CstDerived, Derived...> make_AffineExpr(const CstDerived & constant,
                                                   const std::tuple<LinearExpr<Derived>...> & linear)
{ return {constant, linear}; }
} // namespace utils
} // namespace tvm
 
#include <tvm/utils/internal/AffineExprOperators.h>