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_S_TIR_DATA_LAYOUT_H_
#define TVM_S_TIR_DATA_LAYOUT_H_
 
#include <tvm/ffi/reflection/registry.h>
#include <tvm/tirx/expr.h>
#include <tvm/tirx/op.h>
 
#include <algorithm>
#include <sstream>
#include <string>
#include <utility>
#include <vector>
 
#include "tvm/tirx/var.h"
 
namespace tvm {
namespace tirx {
 
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 tirx::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:
  ffi::String name;
  ffi::Array<tirx::IterVar> axes;
 
  static void RegisterReflection() {
    namespace refl = tvm::ffi::reflection;
    refl::ObjectDef<LayoutNode>()
        .def_ro("name", &LayoutNode::name)
        .def_ro("axes", &LayoutNode::axes);
  }
  TVM_FFI_DECLARE_OBJECT_INFO_FINAL("s_tir.Layout", LayoutNode, Object);
};
 
class Layout : public ObjectRef {
 public:
  explicit Layout(const ffi::Array<tirx::IterVar>& axes);
 
  Layout(const tvm::ffi::String& name) : Layout(name.operator std::string()) {}  // NOLINT(*)
 
  Layout(const char* name) : Layout(std::string(name)) {}  // NOLINT(*)
 
  TVM_DLL Layout(const std::string& name, DataType dtype = DataType::Int(32));  // NOLINT(*)
 
  LayoutNode* operator->() { return static_cast<LayoutNode*>(get_mutable()); }
 
  static const Layout& Undef() {
    static Layout undef;
    return undef;
  }
 
  static IterVar PackIterVar(ffi::Array<IterVar> iters);
 
  static ffi::Array<IterVar> UnpackIterVar(IterVar packed_iter);
 
  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 px : operator->()->axes) {
      auto iter_vars = UnpackIterVar(px);
      for (auto x : iter_vars) {
        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 packed_axis : dst_layout->axes) {
      auto iter_vars = UnpackIterVar(packed_axis);
      for (auto dst_axis : iter_vars) {
        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 std::string& 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) return static_cast<int32_t>(i);
    }
    return -1;
  }
 
  inline int32_t IndexOf(const LayoutAxis& axis) const { return IndexOf(axis.name()); }
 
  inline int32_t IndexOf(const tirx::IterVar& iter) const { return IndexOf(iter->var->name_hint); }
 
  int32_t FactorOf(const LayoutAxis& axis) const;
 
  bool Contains(const LayoutAxis& axis) const {
    if (!defined()) return false;
    for (const tirx::IterVar packed_var : operator->()->axes) {
      auto iter_vars = UnpackIterVar(packed_var);
      for (auto var : iter_vars) {
        if (var->var->name_hint == axis.name()) {
          return true;
        }
      }
    }
    return false;
  }
 
  const LayoutAxis& operator[](int32_t i) const {
    TVM_FFI_ICHECK(defined()) << "Try to access axis from an undefined layout.";
    int32_t index = i < 0 ? static_cast<int32_t>(ndim() + i) : i;
    TVM_FFI_ICHECK(index >= 0 && static_cast<size_t>(index) < ndim()) << "Invalid index " << i;
    const tirx::IterVar axis = operator->()->axes[index];
    return LayoutAxis::Get(axis);
  }
 
  IterVar PackedAxisAt(int32_t i) const {
    TVM_FFI_ICHECK(defined()) << "Try to access axis from an undefined layout.";
    int32_t index = i < 0 ? static_cast<int32_t>(ndim() + i) : i;
    TVM_FFI_ICHECK(index >= 0 && static_cast<size_t>(index) < ndim()) << "Invalid index " << i;
    const tirx::IterVar axis = operator->()->axes[index];
    return 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_FFI_DEFINE_OBJECT_REF_METHODS_NULLABLE(Layout, ObjectRef, LayoutNode);
};
 
// Internal node container BijectiveLayout
class BijectiveLayoutNode : public Object {
 public:
  ffi::Array<PrimExpr> index_forward_rule;
  ffi::Array<PrimExpr> index_backward_rule;
  ffi::Array<PrimExpr> shape_forward_rule;
  ffi::Array<PrimExpr> shape_backward_rule;
 
  Layout src_layout;
  Layout dst_layout;
 
  static void RegisterReflection() {
    namespace refl = tvm::ffi::reflection;
    refl::ObjectDef<BijectiveLayoutNode>()
        .def_ro("src_layout", &BijectiveLayoutNode::src_layout)
        .def_ro("dst_layout", &BijectiveLayoutNode::dst_layout)
        .def_ro("index_forward_rule", &BijectiveLayoutNode::index_forward_rule)
        .def_ro("index_backward_rule", &BijectiveLayoutNode::index_backward_rule)
        .def_ro("shape_forward_rule", &BijectiveLayoutNode::shape_forward_rule)
        .def_ro("shape_backward_rule", &BijectiveLayoutNode::shape_backward_rule);
  }
  TVM_FFI_DECLARE_OBJECT_INFO_FINAL("s_tir.BijectiveLayout", 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 ffi::Array<PrimExpr> ForwardShape(const ffi::Array<PrimExpr>& shape) const;
  // Given the destination shape, recover the source shape.
  TVM_DLL ffi::Array<PrimExpr> BackwardShape(const ffi::Array<PrimExpr>& dst_shape) const;
  // Given the destination indices, infer the destination indices.
  TVM_DLL ffi::Array<PrimExpr> ForwardIndex(const ffi::Array<PrimExpr>& index) const;
  // Given the destination indices, recover the source indices.
  TVM_DLL ffi::Array<PrimExpr> BackwardIndex(const ffi::Array<PrimExpr>& dst_index) const;
 
  TVM_FFI_DEFINE_OBJECT_REF_METHODS_NULLABLE(BijectiveLayout, ObjectRef, BijectiveLayoutNode);
};
 
}  // namespace tirx
}  // namespace tvm
 
#endif  // TVM_S_TIR_DATA_LAYOUT_H_