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 <string>
#include <sstream>
#include <vector>
#include <utility>
#include <algorithm>

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_