data_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 TVM_TIR_DATA_LAYOUT_H_
#define TVM_TIR_DATA_LAYOUT_H_

#include <tvm/tir/expr.h>
#include <tvm/tir/op.h>

#include <algorithm>
#include <sstream>
#include <string>
#include <utility>
#include <vector>

namespace tvm {
namespace tir {

class Layout;

class LayoutAxis {
 public:
  static const LayoutAxis& Get(const char name);

  // Get the singleton LayoutAxis using itvar->var->name_hint
  static const LayoutAxis& Get(const tir::IterVar& itvar);

  // Get the singleton LayoutAxis using name[0] (size of name must be 1).
  static const LayoutAxis& Get(const std::string& name);

  inline bool IsPrimal() const { return name_ >= 'A' && name_ <= 'Z'; }
  inline std::string name() const { return std::string(1, name_); }

  // if current axis is primal, switch the axis to its subordinate one,
  // else switch to the primal.
  inline const LayoutAxis& ToDual() const {
    if (name_ >= 'A' && name_ <= 'Z') {
      return LayoutAxis::Get(name_ - 'A' + 'a');
    } else {
      return LayoutAxis::Get(name_ - 'a' + 'A');
    }
  }

  // return the primal axis. If it is already primal, return itself.
  const LayoutAxis& ToPrimal() const { return IsPrimal() ? *this : ToDual(); }

  // return the subordinate axis. If it is already subordinate, return itself.
  const LayoutAxis& ToSubordinate() const { return IsPrimal() ? ToDual() : *this; }

  inline bool operator==(const LayoutAxis& rhs) const { return name_ == rhs.name_; }

  friend std::ostream& operator<<(std::ostream& os, const LayoutAxis& l) {
    os << l.name();
    return os;
  }

 private:
  static const LayoutAxis UPPER_CASE[];
  static const LayoutAxis LOWER_CASE[];
  LayoutAxis(const LayoutAxis&);
  LayoutAxis& operator=(const LayoutAxis&);
  explicit LayoutAxis(const char name) : name_(name) {}

  const char name_;
};

class LayoutNode : public Object {
 public:
  String name;
  Array<tir::IterVar> axes;

  void VisitAttrs(AttrVisitor* v) {
    v->Visit("name", &name);
    v->Visit("axes", &axes);
  }

  static constexpr const char* _type_key = "tir.Layout";
  TVM_DECLARE_FINAL_OBJECT_INFO(LayoutNode, Object);
};

class Layout : public ObjectRef {
 public:
  explicit Layout(const Array<tir::IterVar>& axes);

  Layout(const tvm::String& name) : Layout(name.operator std::string()) {}  // NOLINT(*)

  Layout(const char* name) : Layout(std::string(name)) {}  // NOLINT(*)

  TVM_DLL Layout(const std::string& name);  // NOLINT(*)

  LayoutNode* operator->() { return static_cast<LayoutNode*>(get_mutable()); }

  static const Layout& Undef() {
    static Layout undef;
    return undef;
  }

  Layout SubLayout(size_t pos, size_t len) const;

  Layout Split(const LayoutAxis& axis, size_t target_pos, int32_t factor) const;

  inline size_t ndim() const {
    if (!defined()) return 0;
    return operator->()->axes.size();
  }

  inline size_t ndim_primal() const {
    if (!defined()) return 0;
    size_t ct = 0;
    for (auto x : operator->()->axes) {
      if (LayoutAxis::Get(x).IsPrimal()) {
        ct++;
      }
    }
    return ct;
  }

  inline Layout ExpandPrimal(const Layout& dst_layout) {
    Layout new_src_layout;
    // 1) Find the axis which are missing in the current layout. Make them the prefix.
    std::string new_src_layout_str = "";
    for (auto dst_axis : dst_layout->axes) {
      if (LayoutAxis::Get(dst_axis).IsPrimal()) {
        if (!this->Contains(LayoutAxis::Get(dst_axis))) {
          new_src_layout_str += dst_axis->var->name_hint;
        }
      }
    }
    // 2) Now, add the primal axis of the current layout.
    new_src_layout_str += this->name();
    new_src_layout = Layout(new_src_layout_str);
    return new_src_layout;
  }

  inline int32_t IndexOf(const LayoutAxis& axis) const {
    if (!this->defined()) return -1;
    const auto axes = operator->()->axes;
    for (size_t i = 0; i < axes.size(); ++i) {
      if (axes[i]->var->name_hint == axis.name()) return static_cast<int32_t>(i);
    }
    return -1;
  }

  int32_t FactorOf(const LayoutAxis& axis) const;

  bool Contains(const LayoutAxis& axis) const {
    if (!defined()) return false;
    for (const tir::IterVar var : operator->()->axes) {
      if (var->var->name_hint == axis.name()) {
        return true;
      }
    }
    return false;
  }

  const LayoutAxis& operator[](int32_t i) const {
    ICHECK(defined()) << "Try to access axis from an undefined layout.";
    int32_t index = i < 0 ? static_cast<int32_t>(ndim() + i) : i;
    ICHECK(index >= 0 && static_cast<size_t>(index) < ndim()) << "Invalid index " << i;
    const tir::IterVar axis = operator->()->axes[index];
    return LayoutAxis::Get(axis);
  }

  inline std::string name() const {
    if (!defined()) return "__undef__";
    return operator->()->name;
  }

  inline bool Equals(const Layout& rhs) const { return name() == rhs.name(); }

  friend std::ostream& operator<<(std::ostream& os, const Layout& l) {
    os << l.name();
    return os;
  }

  TVM_DEFINE_OBJECT_REF_METHODS(Layout, ObjectRef, LayoutNode);
};

// Internal node container BijectiveLayout
class BijectiveLayoutNode : public Object {
 public:
  Array<PrimExpr> index_forward_rule;
  Array<PrimExpr> index_backward_rule;
  Array<PrimExpr> shape_forward_rule;
  Array<PrimExpr> shape_backward_rule;

  Layout src_layout;
  Layout dst_layout;

  void VisitAttrs(AttrVisitor* v) {
    v->Visit("src_layout", &src_layout);
    v->Visit("dst_layout", &dst_layout);
    v->Visit("index_forward_rule", &index_forward_rule);
    v->Visit("index_backward_rule", &index_backward_rule);
    v->Visit("shape_forward_rule", &shape_forward_rule);
    v->Visit("shape_backward_rule", &shape_backward_rule);
  }

  static constexpr const char* _type_key = "tir.BijectiveLayout";
  TVM_DECLARE_FINAL_OBJECT_INFO(BijectiveLayoutNode, Object);
};

class BijectiveLayout : public ObjectRef {
 public:
  TVM_DLL BijectiveLayout(Layout src_layout, Layout dst_layout);

  // Given the source shape, infer the destination shape.
  TVM_DLL Array<PrimExpr> ForwardShape(const Array<PrimExpr>& shape) const;
  // Given the destination shape, recover the source shape.
  TVM_DLL Array<PrimExpr> BackwardShape(const Array<PrimExpr>& dst_shape) const;
  // Given the destination indices, infer the destination indices.
  TVM_DLL Array<PrimExpr> ForwardIndex(const Array<PrimExpr>& index) const;
  // Given the destination indices, recover the source indices.
  TVM_DLL Array<PrimExpr> BackwardIndex(const Array<PrimExpr>& dst_index) const;

  TVM_DEFINE_OBJECT_REF_METHODS(BijectiveLayout, ObjectRef, BijectiveLayoutNode);
};

}  // namespace tir
}  // namespace tvm

#endif  // TVM_TIR_DATA_LAYOUT_H_