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_(NULL) {}
  explicit Shape(const std::vector<index_t> &v)
    : ndim_(v.size()) {
    if (ndim_ <= kStackCache) {
      data_heap_ = NULL;
      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_ = NULL;
      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_ = NULL;
      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_ = NULL;
      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_ = NULL;
      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_ = NULL;
      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_ = NULL;
      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_ = NULL;
  }
#endif

  ~Shape() {
    // data_heap_ can be NULL
    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 NULL
      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_