broadcast.h

Go to the documentation of this file.

#ifndef MSHADOW_EXTENSION_BROADCAST_H_
#define MSHADOW_EXTENSION_BROADCAST_H_
#include "../extension.h"
namespace mshadow {
namespace expr {
template<typename SrcExp, typename DType, int dimdst, int dimdst_m_cast>
struct Broadcast1DExp:
      public MakeTensorExp<Broadcast1DExp<SrcExp, DType, dimdst, dimdst_m_cast>,
                           SrcExp, dimdst, DType> {
  const SrcExp &src_;
  Broadcast1DExp(const SrcExp &src, Shape<dimdst> shape)
      : src_(src) {
    this->shape_ = shape;
  }
};

template<typename SrcExp, typename DType, int dimdst>
struct BroadcastScalarExp:
      public MakeTensorExp<BroadcastScalarExp<SrcExp, DType, dimdst>,
                           SrcExp, dimdst, DType> {
  const SrcExp &src_;
  BroadcastScalarExp(const SrcExp &src, Shape<dimdst> shape)
      : src_(src) {
    this->shape_ = shape;
  }
};

template<int dimcast, typename SrcExp, typename DType,
         int etype, int dimdst>
inline Broadcast1DExp<SrcExp, DType, dimdst, dimdst - dimcast>
broadcast(const expr::Exp<SrcExp, DType, etype> &src, Shape<dimdst> shape) {
  TypeCheckPass<dimcast < dimdst && ExpInfo<SrcExp>::kDim == 1>
                ::Error_Expression_Does_Not_Meet_Dimension_Req();
  typedef ShapeCheck<1, SrcExp> ShapeCheckDim1SrcExp;
  CHECK_EQ(ShapeCheckDim1SrcExp::Check(src.self())[0], shape[dimcast])
    << "broadcast, shape mismatch";
  return Broadcast1DExp<SrcExp, DType, dimdst,
                        dimdst - dimcast>(src.self(), shape);
}

template<typename SrcExp, typename DType, int etype, int dimdst>
inline BroadcastScalarExp<SrcExp, DType, dimdst>
broadcast_scalar(const expr::Exp<SrcExp, DType, etype> &src, Shape<dimdst> shape) {
  TypeCheckPass<ExpInfo<SrcExp>::kDim == 1>
                ::Error_Expression_Does_Not_Meet_Dimension_Req();
  typedef ShapeCheck<1, SrcExp> ShapeCheckDim1SrcExp;
  CHECK_EQ(ShapeCheckDim1SrcExp::Check(src.self())[0], 1U)
      << "broadcast_scalar, source need to be scalar expression";
  return BroadcastScalarExp<SrcExp, DType, dimdst>(src.self(), shape);
}
// short cut functions
template<typename SrcExp, typename DType, int etype>
inline Broadcast1DExp<SrcExp, DType, 2, 1>
repmat(const expr::Exp<SrcExp, DType, etype> &src, index_t nrow) {
  return broadcast<1>
      (src, Shape2(nrow, ShapeCheck<1, SrcExp>::Check(src.self())[0]));
}
//----------------------
// Execution plan
//----------------------
template<typename SrcExp, typename DType, int dimdst, int dimdst_m_cast>
struct Plan<Broadcast1DExp<SrcExp, DType, dimdst, dimdst_m_cast>, DType> {
 public:
  static const int dimcast = dimdst - dimdst_m_cast;
  explicit Plan(const Broadcast1DExp<SrcExp, DType, dimdst, dimdst_m_cast> &e)
      : src_(MakePlan(e.src_)),
        ystride_(e.shape_.ProdShape(dimcast + 1, dimdst - 1)),
        length_(e.shape_[dimcast]) {
    TypeCheckPass<dimcast != dimdst - 1>
        ::Error_Expression_Does_Not_Meet_Dimension_Req();
  }
  MSHADOW_XINLINE DType Eval(index_t y, index_t x) const {
    return src_.Eval(0, (y / ystride_) % length_);
  }

 private:
  expr::Plan<SrcExp, DType> src_;
  const index_t  ystride_, length_;
};

template<typename SrcExp, typename DType, int dimdst>
struct Plan<Broadcast1DExp<SrcExp, DType, dimdst, 1>, DType>{
 public:
  explicit Plan(const Broadcast1DExp<SrcExp, DType, dimdst, 1> &e)
      : src_(MakePlan(e.src_)) {}
  MSHADOW_XINLINE DType Eval(index_t y, index_t x) const {
    return src_.Eval(0, x);
  }

 private:
  expr::Plan<SrcExp, DType> src_;
};

template<typename SrcExp, typename DType, int dimdst>
struct Plan<BroadcastScalarExp<SrcExp, DType, dimdst>, DType>{
 public:
  explicit Plan(const BroadcastScalarExp<SrcExp, DType, dimdst> &e)
      : src_(MakePlan(e.src_)) {}
  MSHADOW_XINLINE DType Eval(index_t y, index_t x) const {
    return src_.Eval(0, 0);
  }

 private:
  expr::Plan<SrcExp, DType> src_;
};
}  // namespace expr
}  // namespace mshadow
#endif  // MSHADOW_EXTENSION_BROADCAST_H_