tuple.h

Go to the documentation of this file.

/*
 * Licensed to the Apache Software Foundation (ASF) under one
 * or more contributor license agreements.  See the NOTICE file
 * distributed with this work for additional information
 * regarding copyright ownership.  The ASF licenses this file
 * to you under the Apache License, Version 2.0 (the
 * "License"); you may not use this file except in compliance
 * with the License.  You may obtain a copy of the License at
 *
 *   http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing,
 * software distributed under the License is distributed on an
 * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
 * KIND, either express or implied.  See the License for the
 * specific language governing permissions and limitations
 * under the License.
 */
 
#ifndef NNVM_TUPLE_H_
#define NNVM_TUPLE_H_
 
#include <algorithm>
#include <iostream>
#include <string>
#include <type_traits>
#include <utility>
#include <vector>
 
#include "base.h"
 
namespace nnvm {
 
typedef int64_t dim_t;
 
template <typename ValueType>
class Tuple {
 public:
  Tuple() = default;
  inline ~Tuple() { delete[] data_heap_; }
  inline Tuple(const Tuple<ValueType>& s) { this->assign(s.begin(), s.end()); }
  inline Tuple(std::initializer_list<ValueType> init) { this->assign(init.begin(), init.end()); }
  inline Tuple(std::vector<ValueType> init) {  // NOLINT(runtime/explicit)
    this->assign(init.begin(), init.end());
  }
  inline Tuple(Tuple<ValueType>&& src) {  // NOLINT(runtime/explicit)
    this->swap(src);
  }
  template <typename RandomAccessIterator>
  inline Tuple(RandomAccessIterator begin, RandomAccessIterator end) {
    this->assign(begin, end);
  }
  template <typename RandomAccessIterator>
  inline void assign(RandomAccessIterator begin, RandomAccessIterator end) {
    this->SetDim(end - begin);
    std::copy(begin, end, this->begin());
  }
  inline void swap(Tuple<ValueType>& other) {  // NOLINT(*)
    std::swap(ndim_, other.ndim_);
    std::swap(num_heap_allocated_, other.num_heap_allocated_);
    std::swap(data_stack_, other.data_stack_);
    std::swap(data_heap_, other.data_heap_);
  }
  inline Tuple<ValueType>& operator=(const Tuple<ValueType>& src) {
    this->assign(src.begin(), src.end());
    return *this;
  }
  inline Tuple<ValueType>& operator=(Tuple<ValueType>&& src) {
    Tuple<ValueType>(std::move(src)).swap(*this);
    return *this;
  }
  inline Tuple<ValueType>& operator=(std::initializer_list<ValueType> init) {
    this->assign(init.begin(), init.end());
    return *this;
  }
  inline bool operator==(const Tuple<ValueType>& s) const {
    if (ndim_ != s.ndim_) return false;
    return std::equal(begin(), end(), s.begin());
  }
  inline bool operator!=(const Tuple<ValueType>& s) const { return !(*this == s); }
  inline const ValueType* begin() const { return ndim_ <= kStackCache ? data_stack_ : data_heap_; }
  inline ValueType* begin() { return ndim_ <= kStackCache ? data_stack_ : data_heap_; }
  inline const ValueType* end() const {
    return ndim_ <= kStackCache ? (data_stack_ + ndim_) : (data_heap_ + ndim_);
  }
  inline ValueType* end() {
    return ndim_ <= kStackCache ? (data_stack_ + ndim_) : (data_heap_ + ndim_);
  }
  inline uint32_t ndim() const { return ndim_; }
  inline ValueType& operator[](size_t i) { return begin()[i]; }
  inline const ValueType& operator[](size_t i) const { return begin()[i]; }
  inline void Save(dmlc::JSONWriter* writer) const {
    std::vector<ValueType> tmp(begin(), end());
    writer->Write(tmp);
  }
  inline void Load(dmlc::JSONReader* reader) {
    std::vector<ValueType> tmp;
    reader->Read(&tmp);
    this->assign(tmp.begin(), tmp.end());
  }
  friend std::ostream& operator<<(std::ostream& os, const Tuple<ValueType>& t) {
    os << '[';
    const ValueType* begin = t.begin();
    const ValueType* end = t.end();
    for (const ValueType* it = begin; it != end; ++it) {
      if (it != begin) os << ',';
      os << *it;
    }
    os << ']';
    return os;
  }
  friend std::istream& operator>>(std::istream& is, Tuple<ValueType>& t) {
    // get (
    while (true) {
      char ch = is.peek();
      if (isdigit(ch) || ch == '-') {
        ValueType idx;
        if (is >> idx) {
          t.assign(&idx, &idx + 1);
        }
        return is;
      }
      is.get();
      if (ch == '(' || ch == '[') break;
      if (!isspace(ch)) {
        is.setstate(std::ios::failbit);
        return is;
      }
    }
    // Handle empty tuple
    while (isspace(is.peek())) {
      is.get();
    }
    if (is.peek() == ')' || is.peek() == ']') {
      is.get();
      return is;
    }
    // Handle non-empty tuple
    ValueType idx;
    std::vector<ValueType> tmp;
    while (is >> idx) {
      tmp.push_back(idx);
      char ch;
      do {
        ch = is.get();
      } while (isspace(ch));
      if (std::is_integral<ValueType>::value && ch == 'L') {
        ch = is.get();
      }
      if (ch == ',') {
        while (true) {
          ch = is.peek();
          if (isspace(ch)) {
            is.get();
            continue;
          }
          if (ch == ')' || ch == ']') {
            is.get();
            break;
          }
          break;
        }
        if (ch == ')' || ch == ']') break;
      } else if (ch == ')' || ch == ']') {
        break;
      } else {
        is.setstate(std::ios::failbit);
        return is;
      }
    }
    t.assign(tmp.begin(), tmp.end());
    return is;
  }
  template <typename DType = ValueType, typename TStream>
  inline void Save(TStream* strm) const;
  template <typename DType = ValueType, typename TStream>
  inline bool Load(TStream* strm);
 
 protected:
  // stack cache size
  static const uint32_t kStackCache = 4;
  uint32_t ndim_{0};
  uint32_t num_heap_allocated_{0};
  ValueType data_stack_[kStackCache];
  ValueType* data_heap_{nullptr};
  // internal function to change the dimension
  inline void SetDim(uint32_t ndim) {
    if (ndim > kStackCache && ndim > num_heap_allocated_) {
      delete[] data_heap_;
      data_heap_ = new ValueType[ndim];
      num_heap_allocated_ = ndim;
    }
    ndim_ = ndim;
  }
};
 
class TShape : public Tuple<dim_t> {
 public:
  TShape() = default;
  inline TShape(uint32_t ndim) {  // NOLINT(*)
    this->SetDim(ndim);
    std::fill_n(begin(), ndim, 1);
  }
  inline TShape(const Tuple<dim_t>& s) {  // NOLINT(*)
    this->assign(s.begin(), s.end());
  }
  inline TShape(std::initializer_list<dim_t> init) { this->assign(init.begin(), init.end()); }
  inline TShape(Tuple<dim_t>&& s) {  // NOLINT(*)
    this->swap(s);
  }
  template <typename RandomAccessIterator>
  inline TShape(RandomAccessIterator begin, RandomAccessIterator end) {
    this->assign(begin, end);
  }
  inline TShape& operator=(const Tuple<dim_t>& src) {
    this->assign(src.begin(), src.end());
    return *this;
  }
  inline TShape& operator=(Tuple<dim_t>&& src) {  // NOLINT(*)
    TShape(std::move(src)).swap(*this);           // NOLINT(*)
    return *this;
  }
  inline size_t Size() const {
    dim_t size = 1;
    const dim_t *start = begin(), *fin = end();
    for (const dim_t* it = start; it != fin; ++it) {
      size *= *it;
    }
    return size;
  }
  inline size_t ProdShape(int dimstart, int dimend) const {
    dim_t num = 1;
    const dim_t* d = this->data();
    for (int i = dimstart; i < dimend; ++i) {
      num *= d[i];
    }
    return num;
  }
  inline const dim_t* data() const { return begin(); }
  inline dim_t* data() { return begin(); }
#ifdef MSHADOW_XINLINE
  template <int dim>
  inline TShape(const mshadow::Shape<dim>& s) {  // NOLINT(*)
    this->assign(s.shape_, s.shape_ + dim);
  }
 
  template <int dim>
  inline TShape(mshadow::Shape<dim>&& s) {  // NOLINT(*)
    this->assign(s.shape_, s.shape_ + dim);
  }
  template <int dim>
  inline TShape& operator=(const mshadow::Shape<dim>& shape) {
    this->assign(shape.shape_, shape.shape_ + dim);
    return *this;
  }
  template <int dim>
  inline mshadow::Shape<dim> get() const {
    CHECK_EQ(dim, static_cast<int>(ndim()))
        << "dimension do not match target dimension " << dim << " vs " << ndim();
    const dim_t* d = this->data();
    mshadow::Shape<dim> s;
    for (int i = 0; i < dim; ++i) {
      s[i] = d[i];
    }
    return s;
  }
  inline mshadow::Shape<2> FlatTo2D(void) const {
    mshadow::Shape<2> s;
    if (ndim() == 0) return mshadow::Shape2(0, 0);
    const dim_t* d = this->data();
    s.shape_[1] = d[ndim() - 1];
    dim_t ymax = 1;
    for (size_t i = 1; i < ndim(); ++i) {
      ymax *= d[i - 1];
    }
    s.shape_[0] = ymax;
    return s;
  }
  inline mshadow::Shape<3> FlatTo3D(size_t axis_begin, size_t axis_end) const {
    CHECK(axis_end >= axis_begin);
    mshadow::Shape<3> s;
    if (ndim() == 0) return mshadow::Shape3(0, 0, 0);
    const dim_t* d = this->data();
    s.shape_[0] = 1;
    s.shape_[1] = 1;
    s.shape_[2] = 1;
 
    for (size_t i = 0; i < axis_begin; ++i) {
      s.shape_[0] *= d[i];
    }
    for (size_t i = axis_begin; i <= axis_end; ++i) {
      s.shape_[1] *= d[i];
    }
    for (size_t i = axis_end + 1; i < ndim(); ++i) {
      s.shape_[2] *= d[i];
    }
    return s;
  }
  inline mshadow::Shape<3> FlatTo3D(size_t axis) const { return FlatTo3D(axis, axis); }
  inline bool operator==(const TShape& s) const {
    if (ndim() != s.ndim()) return false;
    return std::equal(begin(), end(), s.begin());
  }
  inline bool operator!=(const TShape& s) const { return !(*this == s); }
  template <int dim>
  inline bool operator==(const mshadow::Shape<dim>& s) const {
    if (ndim_ != dim) return false;
    const dim_t* d = dim <= kStackCache ? data_stack_ : data_heap_;
    for (size_t i = 0; i < dim; ++i) {
      if (d[i] != s.shape_[i]) return false;
    }
    return true;
  }
  template <int dim>
  inline bool operator!=(const mshadow::Shape<dim>& s) const {
    return !(*this == s);
  }
#endif
};
 
template <typename SrcIter, typename DstIter>
inline DstIter ShapeTypeCast(const SrcIter begin, const SrcIter end, DstIter dst_begin) {
  typedef typename std::iterator_traits<SrcIter>::value_type SrcDType;
  typedef typename std::iterator_traits<DstIter>::value_type DstDType;
  auto cast = [](const SrcDType& dim) { return static_cast<DstDType>(dim); };
  return std::transform(begin, end, dst_begin, cast);
}
 
template <typename SrcIter>
inline TShape ShapeTypeCast(const SrcIter begin, const SrcIter end) {
  size_t ndim = std::distance(begin, end);
  TShape res(ndim);
  ShapeTypeCast(begin, end, res.begin());
  return res;
}
 
template <typename ValueType>
template <typename DType, typename TStream>
inline void Tuple<ValueType>::Save(TStream* strm) const {
  strm->Write(&ndim_, sizeof(ndim_));
  if (typeid(DType) == typeid(ValueType)) {
    strm->Write(begin(), sizeof(ValueType) * ndim_);
  } else {
    std::vector<DType> buffer(ndim_);
    ShapeTypeCast(begin(), end(), buffer.data());
    strm->Write(buffer.data(), sizeof(DType) * ndim_);
  }
}
 
template <typename ValueType>
template <typename DType, typename TStream>
inline bool Tuple<ValueType>::Load(TStream* strm) {
  if (strm->Read(&ndim_, sizeof(ndim_)) != sizeof(ndim_)) return false;
  this->SetDim(ndim_);
  size_t nread = sizeof(DType) * ndim_;
  if (typeid(DType) == typeid(ValueType)) {
    if (strm->Read(begin(), nread) != nread) return false;
  } else {
    std::vector<DType> buffer(ndim_);
    if (strm->Read(buffer.data(), nread) != nread) return false;
    ShapeTypeCast(buffer.begin(), buffer.end(), begin());
  }
  return true;
}
 
}  // namespace nnvm
 
namespace std {
template <typename T>
struct hash<nnvm::Tuple<T> > {
  size_t operator()(const nnvm::Tuple<T>& val) const {
    std::hash<uint32_t> hash_uint;
    size_t res = hash_uint(val.ndim());
    for (uint32_t i = 0; i < val.ndim(); ++i) {
      res = dmlc::HashCombine(res, val[i]);
    }
    return res;
  }
};
 
template <>
struct hash<nnvm::TShape> {
  size_t operator()(const nnvm::TShape& val) const {
    std::hash<uint32_t> hash_uint;
    size_t res = hash_uint(val.ndim());
    for (uint32_t i = 0; i < val.ndim(); ++i) {
      res = dmlc::HashCombine(res, val[i]);
    }
    return res;
  }
};
}  // namespace std
 
namespace dmlc {
DMLC_DECLARE_TYPE_NAME(optional<nnvm::TShape>, "Shape or None");
// avoid low version of MSVC
#if !defined(_MSC_VER)
template <typename T>
struct type_name_helper<nnvm::Tuple<T> > {
  static inline std::string value() { return "tuple of <" + type_name<T>() + ">"; }
};
#endif
}  // namespace dmlc
#endif  // NNVM_TUPLE_H_