shape.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 MXNET_CPP_SHAPE_H_
#define MXNET_CPP_SHAPE_H_
 
#include <istream>
#include <ostream>
#include <algorithm>
#include <vector>
#include "mxnet-cpp/base.h"
 
namespace mxnet {
namespace cpp {
 
struct Shape {
 public:
  Shape() : ndim_(0), num_heap_allocated_(0), data_heap_(nullptr) {}
  explicit Shape(const std::vector<index_t>& v) : ndim_(v.size()) {
    if (ndim_ <= kStackCache) {
      data_heap_          = nullptr;
      num_heap_allocated_ = 0;
      std::copy(v.begin(), v.end(), data_stack_);
    } else {
      data_heap_          = new index_t[ndim_];
      num_heap_allocated_ = ndim_;
      std::copy(v.begin(), v.end(), data_heap_);
    }
  }
  explicit Shape(index_t s1) : ndim_(1) {
    if (ndim_ <= kStackCache) {
      data_heap_          = nullptr;
      num_heap_allocated_ = 0;
      data_stack_[0]      = s1;
    } else {
      data_heap_          = new index_t[ndim_];
      num_heap_allocated_ = ndim_;
      data_heap_[0]       = s1;
    }
  }
  Shape(index_t s1, index_t s2) : ndim_(2) {
    if (ndim_ <= kStackCache) {
      data_heap_          = nullptr;
      num_heap_allocated_ = 0;
      data_stack_[0]      = s1;
      data_stack_[1]      = s2;
    } else {
      data_heap_          = new index_t[ndim_];
      num_heap_allocated_ = ndim_;
      data_heap_[0]       = s1;
      data_heap_[1]       = s2;
    }
  }
  Shape(index_t s1, index_t s2, index_t s3) : ndim_(3) {
    if (ndim_ <= kStackCache) {
      data_heap_          = nullptr;
      num_heap_allocated_ = 0;
      data_stack_[0]      = s1;
      data_stack_[1]      = s2;
      data_stack_[2]      = s3;
    } else {
      data_heap_          = new index_t[ndim_];
      num_heap_allocated_ = ndim_;
      data_heap_[0]       = s1;
      data_heap_[1]       = s2;
      data_heap_[2]       = s3;
    }
  }
  Shape(index_t s1, index_t s2, index_t s3, index_t s4) : ndim_(4) {
    if (ndim_ <= kStackCache) {
      data_heap_          = nullptr;
      num_heap_allocated_ = 0;
      data_stack_[0]      = s1;
      data_stack_[1]      = s2;
      data_stack_[2]      = s3;
      data_stack_[3]      = s4;
    } else {
      data_heap_          = new index_t[ndim_];
      num_heap_allocated_ = ndim_;
      data_heap_[0]       = s1;
      data_heap_[1]       = s2;
      data_heap_[2]       = s3;
      data_heap_[3]       = s4;
    }
  }
  Shape(index_t s1, index_t s2, index_t s3, index_t s4, index_t s5) : ndim_(5) {
    if (ndim_ <= kStackCache) {
      data_heap_          = nullptr;
      num_heap_allocated_ = 0;
      data_stack_[0]      = s1;
      data_stack_[1]      = s2;
      data_stack_[2]      = s3;
      data_stack_[3]      = s4;
      data_stack_[4]      = s5;
    } else {
      data_heap_          = new index_t[ndim_];
      num_heap_allocated_ = ndim_;
      data_heap_[0]       = s1;
      data_heap_[1]       = s2;
      data_heap_[2]       = s3;
      data_heap_[3]       = s4;
      data_heap_[4]       = s5;
    }
  }
  Shape(const Shape& s) : ndim_(s.ndim_) {
    if (ndim_ <= kStackCache) {
      data_heap_          = nullptr;
      num_heap_allocated_ = 0;
      std::copy(s.data_stack_, s.data_stack_ + ndim_, data_stack_);
    } else {
      data_heap_          = new index_t[ndim_];
      num_heap_allocated_ = ndim_;
      std::copy(s.data_heap_, s.data_heap_ + ndim_, data_heap_);
    }
  }
#if MSHADOW_IN_CXX11
 
  Shape(Shape&& s)
      : ndim_(s.ndim_), num_heap_allocated_(s.num_heap_allocated_), data_heap_(s.data_heap_) {
    if (ndim_ <= kStackCache) {
      std::copy(s.data_stack_, s.data_stack_ + ndim_, data_stack_);
    }
    // remove data heap space from s
    s.data_heap_ = nullptr;
  }
#endif
 
  ~Shape() {
    // data_heap_ can be nullptr
    delete[] data_heap_;
  }
  template <typename RandomAccessIterator>
  inline void CopyFrom(RandomAccessIterator begin, RandomAccessIterator end) {
    this->SetDim(end - begin);
    std::copy(begin, end, data());
  }
  inline Shape& operator=(const Shape& shape) {
    this->SetDim(shape.ndim_);
    const index_t* src = shape.data();
    std::copy(src, src + ndim_, data());
    return *this;
  }
  inline Shape& operator=(const std::vector<index_t>& shape) {
    this->CopyFrom(shape.begin(), shape.end());
    return *this;
  }
  inline const index_t* data() const {
    return ndim_ <= kStackCache ? data_stack_ : data_heap_;
  }
  inline index_t* data() {
    return ndim_ <= kStackCache ? data_stack_ : data_heap_;
  }
  inline index_t ndim(void) const {
    return ndim_;
  }
  inline index_t& operator[](index_t i) {
    return data()[i];
  }
  inline const index_t& operator[](index_t i) const {
    return data()[i];
  }
  inline size_t Size(void) const {
    size_t size      = 1;
    const index_t* d = this->data();
    for (index_t i = 0; i < ndim_; ++i) {
      size *= d[i];
    }
    return size;
  }
  inline bool operator==(const Shape& s) const {
    if (ndim_ != s.ndim_)
      return false;
    if (ndim_ <= kStackCache) {
      for (index_t i = 0; i < ndim_; ++i) {
        if (data_stack_[i] != s.data_stack_[i])
          return false;
      }
    } else {
      for (index_t i = 0; i < ndim_; ++i) {
        if (data_heap_[i] != s.data_heap_[i])
          return false;
      }
    }
    return true;
  }
  inline bool operator!=(const Shape& s) const {
    return !(*this == s);
  }
 
  friend std::ostream& operator<<(std::ostream& os, const Shape& shape);
  friend std::istream& operator>>(std::istream& is, Shape& shape);
 
 private:
  // the shape will be stored in data_stack_
  // when dimension is smaller than kStackCache
  // when it is bigger, it will be stored in data_heap_;
  static const index_t kStackCache = 5;
  index_t ndim_;
  index_t num_heap_allocated_;
  index_t data_stack_[kStackCache];
  index_t* data_heap_;
  inline void SetDim(index_t dim) {
    if (dim > kStackCache && dim > num_heap_allocated_) {
      // data_heap_ can be nullptr
      delete[] data_heap_;
      data_heap_          = new index_t[dim];
      num_heap_allocated_ = dim;
    }
    ndim_ = dim;
  }
};
 
inline std::ostream& operator<<(std::ostream& os, const Shape& shape) {
  os << '(';
  for (index_t i = 0; i < shape.ndim(); ++i) {
    if (i != 0)
      os << ',';
    os << static_cast<int>(shape[i]);  // Supports negative Shape 'special codes' for inferring
  }
  // python style tuple
  if (shape.ndim() == 1)
    os << ',';
  os << ')';
  return os;
}
 
inline std::istream& operator>>(std::istream& is, Shape& shape) {
  // get (
  while (true) {
    char ch = is.get();
    if (ch == '(')
      break;
    if (!isspace(ch)) {
      is.setstate(std::ios::failbit);
      return is;
    }
  }
  index_t idx;
  std::vector<index_t> tmp;
  while (is >> idx) {
    tmp.push_back(idx);
    char ch;
    do {
      ch = is.get();
    } while (isspace(ch));
    if (ch == ',') {
      while (true) {
        ch = is.peek();
        if (isspace(ch)) {
          is.get();
          continue;
        }
        if (ch == ')') {
          is.get();
          break;
        }
        break;
      }
      if (ch == ')')
        break;
    } else if (ch == ')') {
      break;
    } else {
      is.setstate(std::ios::failbit);
      return is;
    }
  }
  shape.CopyFrom(tmp.begin(), tmp.end());
  return is;
}
 
}  // namespace cpp
}  // namespace mxnet
 
#endif  // MXNET_CPP_SHAPE_H_