layout.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_LAYOUT_H_
#define NNVM_LAYOUT_H_
 
#include <dmlc/parameter.h>
 
#include <algorithm>
#include <sstream>
#include <string>
#include <utility>
#include <vector>
 
namespace nnvm {
 
class Layout {
 public:
  using LayoutDim = char;
 
  Layout() : name_("__undef__") {}  // NOLINT(*)
 
  inline Layout(const std::string& layout) {  // NOLINT(*)
    parse(layout);
  }
  inline Layout(const Layout& s) {  // NOLINT(*)
    this->parse(s.name_);
  }
  inline Layout(Layout&& src) {  // NOLINT(*)
    this->swap(src);
  }
  inline Layout& operator=(const Layout& src) {
    this->parse(src.name_);
    return *this;
  }
  inline Layout& operator=(Layout&& src) {
    Layout(std::move(src)).swap(*this);  // NOLINT(*)
    return *this;
  }
  inline Layout& operator=(const std::string& src) {
    this->parse(src);
    return *this;
  }
  inline bool operator==(const Layout& s) const { return name_ == s.name_; }
  inline bool operator!=(const Layout& s) const { return !(*this == s); }
 
  inline Layout operator+(const Layout& other) const {
    if (!this->defined() && !other.defined()) {
      return Layout::Undef();
    } else if (!this->defined()) {
      return other;
    } else if (!other.defined()) {
      return *this;
    }
    return Layout(this->name_ + other.name_);
  }
 
  static inline bool is_superdim(LayoutDim dim) { return dim >= 'A' && dim <= 'Z'; }
 
  static inline bool is_subdim(LayoutDim dim) { return dim >= 'a' && dim <= 'z'; }
 
  static inline LayoutDim to_superdim(LayoutDim dim) {
    if (is_subdim(dim)) {
      return dim - 'a' + 'A';
    }
    return dim;
  }
 
  static inline LayoutDim to_subdim(LayoutDim dim) {
    if (is_superdim(dim)) {
      return dim - 'A' + 'a';
    }
    return dim;
  }
 
  static inline const Layout& Undef() {
    static Layout undef;
    return undef;
  }
 
  inline void swap(Layout& other) {  // NOLINT(*)
    std::swap(name_, other.name_);
    std::swap(superdim_pos_, other.superdim_pos_);
    std::swap(subdim_pos_, other.subdim_pos_);
    std::swap(subdim_size_, other.subdim_size_);
    std::swap(layout_simplified_, other.layout_simplified_);
  }
 
  inline bool convertible(const Layout& dst) const {
    if (!this->defined() || !dst.defined()) return false;
    for (size_t i = 0; i < kUniqueDim; ++i) {
      if ((superdim_pos_[i] >= 0 && dst.superdim_pos_[i] < 0) ||
          (superdim_pos_[i] < 0 && dst.superdim_pos_[i] >= 0)) {
        return false;
      }
    }
    return true;
  }
 
  inline Layout sublayout(size_t pos, size_t len) const {
    if (pos > ndim()) return Layout::Undef();
    if (pos + len > ndim()) len = ndim() - pos;
    if (len == 0) return Layout::Undef();
    std::ostringstream new_layout;
    for (size_t i = pos; i < pos + len; ++i) {
      if (is_subdim(layout_simplified_[i])) {
        auto block_size = this->subsizeof(layout_simplified_[i]);
        CHECK_GT(block_size, 0);
        new_layout << block_size;
      }
      new_layout << layout_simplified_[i];
    }
    return Layout(new_layout.str());
  }
 
  inline Layout reverse() const {
    if (!this->defined()) return Layout::Undef();
    std::ostringstream new_layout;
    for (int64_t i = this->ndim() - 1; i >= 0; --i) {
      if (is_subdim(layout_simplified_[i])) {
        auto block_size = this->subsizeof(layout_simplified_[i]);
        CHECK_GT(block_size, 0);
        new_layout << block_size;
      }
      new_layout << layout_simplified_[i];
    }
    return Layout(new_layout.str());
  }
 
  inline Layout split(LayoutDim dim, size_t target_pos, uint32_t size) const {
    CHECK(target_pos <= this->ndim())
        << "Invalid split position " << target_pos << " for layout " << name_;
    CHECK(is_superdim(dim)) << "Cannot split a sub-dimension " << dim;
    CHECK(this->contains(dim)) << "Axis " << dim << " does not exist in " << name_;
    CHECK(!this->contains(to_subdim(dim)))
        << "Dimension " << dim << " has already been split in " << name_;
    CHECK(size > 0) << "Invalid split size " << size;
    std::ostringstream new_layout;
    for (size_t i = 0; i <= this->ndim(); ++i) {
      if (i == target_pos) {
        new_layout << size << Layout::to_subdim(dim);
      }
      if (i == this->ndim()) break;
      new_layout << this->at(i);
    }
    Layout x(new_layout.str());
    return x;
  }
 
  using iterator = std::vector<LayoutDim>::const_iterator;
  using reverse_iterator = std::vector<LayoutDim>::const_reverse_iterator;
 
  inline iterator begin() const { return layout_simplified_.begin(); }
  inline iterator end() const { return layout_simplified_.end(); }
  inline reverse_iterator rbegin() const { return layout_simplified_.rbegin(); }
  inline reverse_iterator rend() const { return layout_simplified_.rend(); }
 
  inline size_t ndim() const { return layout_simplified_.size(); }
 
  inline std::string at(size_t i) const {
    CHECK_LT(i, this->ndim()) << "position " << i << " exceeds ndim=" << this->ndim();
    std::ostringstream repr;
    if (is_subdim(layout_simplified_[i])) {
      auto factor = subsizeof(layout_simplified_[i]);
      CHECK_GT(factor, 0);
      repr << factor;
    }
    repr << layout_simplified_[i];
    return repr.str();
  }
 
  inline int32_t indexof(LayoutDim dim) const {
    if (!this->defined())
      return -1;
    else if (is_superdim(dim))
      return superdim_pos_[dim - 'A'];
    else if (is_subdim(dim))
      return subdim_pos_[dim - 'a'];
    return -1;
  }
 
  inline int64_t subsizeof(LayoutDim dim) const {
    CHECK(is_superdim(dim) || is_subdim(dim)) << "Invalid dim " << dim;
    if (!this->defined() || !this->contains(to_subdim(dim))) {
      return -1;
    }
    int idx = to_subdim(dim) - 'a';
    return subdim_size_[idx];
  }
 
  inline bool contains(LayoutDim dim) const {
    if (is_superdim(dim)) {
      return superdim_pos_[dim - 'A'] >= 0;
    } else if (is_subdim(dim)) {
      return subdim_pos_[dim - 'a'] >= 0;
    }
    return false;
  }
 
  inline LayoutDim operator[](size_t i) const { return layout_simplified_[i]; }
 
  inline bool defined() const { return name_ != "__undef__"; }
 
  inline const std::string& name() const { return name_; }
 
  inline void Save(dmlc::JSONWriter* writer) const { writer->Write(name_); }
 
  inline void Load(dmlc::JSONReader* reader) {
    std::string tmp;
    reader->Read(&tmp);
    this->parse(tmp);
  }
 
  friend std::ostream& operator<<(std::ostream& os, const Layout& l) {
    os << l.name_;
    return os;
  }
 
 private:
  static const uint32_t kUniqueDim = 26;
 
  std::string name_;
  int32_t superdim_pos_[kUniqueDim];
  int32_t subdim_pos_[kUniqueDim];
  int64_t subdim_size_[kUniqueDim];
  std::vector<LayoutDim> layout_simplified_;
 
  void parse(const std::string& layout) {
    name_ = layout;
    std::fill_n(superdim_pos_, kUniqueDim, -1);
    std::fill_n(subdim_pos_, kUniqueDim, -1);
    std::fill_n(subdim_size_, kUniqueDim, -1);
    layout_simplified_.clear();
 
    if (layout == "__undef__") return;
 
    int32_t factor = 0;
    uint32_t curr = 0;
    for (size_t i = 0; i < layout.size(); ++i) {
      const LayoutDim c = layout.at(i);
      if (is_superdim(c)) {
        int pos = c - 'A';
        CHECK_EQ(factor, 0) << "Invalid layout " << layout << ": invalid factor size " << factor
                            << " before dimension " << c;
        CHECK_EQ(superdim_pos_[pos], -1)
            << "Invalid layout " << layout << ": duplicate dimension " << c;
        superdim_pos_[pos] = curr++;
        layout_simplified_.push_back(c);
      } else if (is_subdim(c)) {
        int pos = c - 'a';
        CHECK_GT(factor, 0) << "Invalid layout " << layout << ": invalid factor size " << factor
                            << " for dimension " << c;
        CHECK_EQ(subdim_pos_[pos], -1)
            << "Invalid layout " << layout << ": duplicate dimension " << c;
        CHECK_EQ(subdim_size_[pos], -1)
            << "Invalid layout " << layout << ": duplicate dimension " << c;
        subdim_pos_[pos] = curr++;
        subdim_size_[pos] = factor;
        layout_simplified_.push_back(c);
        factor = 0;
      } else if (c >= '0' && c <= '9') {
        CHECK(factor >= 0) << "Invalid layout " << layout << ": _ is adjacent to a number.";
        factor = factor * 10 + c - '0';
      } else {
        LOG(FATAL) << "Invalid layout " << layout;
      }
    }
    CHECK(!layout_simplified_.empty()) << "Invalid layout " << layout;
    for (LayoutDim dim : layout_simplified_) {
      CHECK(is_superdim(dim) || superdim_pos_[dim - 'a'] >= 0)
          << "Invalid layout " << layout << ": missing axis " << static_cast<char>(dim - 'a' + 'A');
    }
  }
};
 
}  // namespace nnvm
 
#endif  // NNVM_LAYOUT_H_