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/optional.h>
#include <tvm/ffi/reflection/registry.h>
#include <tvm/ir/prim/expr.h>
#include <tvm/ir/prim/op.h>
#include <tvm/s_tir/iter_var.h>
 
#include <algorithm>
#include <sstream>
#include <string>
#include <utility>
#include <vector>
 
#include "tvm/ir/expr.h"
 
namespace tvm {
namespace tirx {
 
class SLayout;
 
class SLayoutAxis {
 public:
  static const SLayoutAxis& Get(const char name);
 
  // Get the singleton SLayoutAxis using itvar->var->name
  static const SLayoutAxis& Get(const s_tir::IterVar& itvar);
 
  // Get the singleton SLayoutAxis using name[0] (size of name must be 1).
  static const SLayoutAxis& 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 SLayoutAxis& ToDual() const {
    if (name_ >= 'A' && name_ <= 'Z') {
      return SLayoutAxis::Get(name_ - 'A' + 'a');
    } else {
      return SLayoutAxis::Get(name_ - 'a' + 'A');
    }
  }
 
  // return the primal axis. If it is already primal, return itself.
  const SLayoutAxis& ToPrimal() const { return IsPrimal() ? *this : ToDual(); }
 
  // return the subordinate axis. If it is already subordinate, return itself.
  const SLayoutAxis& ToSubordinate() const { return IsPrimal() ? ToDual() : *this; }
 
  inline bool operator==(const SLayoutAxis& rhs) const { return name_ == rhs.name_; }
 
  friend std::ostream& operator<<(std::ostream& os, const SLayoutAxis& l) {
    os << l.name();
    return os;
  }
 
 private:
  static const SLayoutAxis UPPER_CASE[];
  static const SLayoutAxis LOWER_CASE[];
  SLayoutAxis(const SLayoutAxis&);
  SLayoutAxis& operator=(const SLayoutAxis&);
  explicit SLayoutAxis(const char name) : name_(name) {}
 
  const char name_;
};
 
class SLayoutNode : public ffi::Object {
 public:
  ffi::String name;
  ffi::Array<s_tir::IterVar> axes;
 
  static void RegisterReflection() {
    namespace refl = tvm::ffi::reflection;
    refl::ObjectDef<SLayoutNode>()
        .def_ro("name", &SLayoutNode::name)
        .def_ro("axes", &SLayoutNode::axes);
  }
  TVM_FFI_DECLARE_OBJECT_INFO_FINAL("s_tir.SLayout", SLayoutNode, ffi::Object);
};
 
class SLayout : public ffi::ObjectRef {
 public:
  explicit SLayout(const ffi::Array<s_tir::IterVar>& axes);
 
  SLayout(const tvm::ffi::String& name) : SLayout(name.operator std::string()) {}  // NOLINT(*)
 
  SLayout(const char* name) : SLayout(std::string(name)) {}  // NOLINT(*)
 
  TVM_DLL SLayout(const std::string& name, PrimType index_ty = PrimType::Int(32));  // NOLINT(*)
  TVM_DLL static ffi::Optional<SLayout> Create(const std::string& name,
                                               PrimType index_ty = PrimType::Int(32));
 
  SLayoutNode* operator->() { return static_cast<SLayoutNode*>(get_mutable()); }
 
  static ffi::Optional<SLayout> Undef() { return std::nullopt; }
 
  static s_tir::IterVar PackIterVar(ffi::Array<s_tir::IterVar> iters);
 
  static ffi::Array<s_tir::IterVar> UnpackIterVar(s_tir::IterVar packed_iter);
 
  ffi::Optional<SLayout> SubLayout(size_t pos, size_t len) const;
 
  SLayout Split(const SLayoutAxis& 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 (SLayoutAxis::Get(x).IsPrimal()) {
          ct++;
        }
      }
    }
    return ct;
  }
 
  inline ffi::Optional<SLayout> ExpandPrimal(const SLayout& dst_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 (SLayoutAxis::Get(dst_axis).IsPrimal()) {
          if (!this->Contains(SLayoutAxis::Get(dst_axis))) {
            new_src_layout_str += dst_axis->var->name;
          }
        }
      }
    }
    // 2) Now, add the primal axis of the current layout.
    new_src_layout_str += this->name();
    return SLayout::Create(new_src_layout_str);
  }
 
  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 == axis) return static_cast<int32_t>(i);
    }
    return -1;
  }
 
  inline int32_t IndexOf(const SLayoutAxis& axis) const { return IndexOf(axis.name()); }
 
  inline int32_t IndexOf(const s_tir::IterVar& iter) const { return IndexOf(iter->var->name); }
 
  int32_t FactorOf(const SLayoutAxis& axis) const;
 
  bool Contains(const SLayoutAxis& axis) const {
    if (!defined()) return false;
    for (const s_tir::IterVar packed_var : operator->()->axes) {
      auto iter_vars = UnpackIterVar(packed_var);
      for (auto var : iter_vars) {
        if (var->var->name == axis.name()) {
          return true;
        }
      }
    }
    return false;
  }
 
  const SLayoutAxis& 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 s_tir::IterVar axis = operator->()->axes[index];
    return SLayoutAxis::Get(axis);
  }
 
  s_tir::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 s_tir::IterVar axis = operator->()->axes[index];
    return axis;
  }
 
  inline std::string name() const {
    if (!defined()) return "__undef__";
    return operator->()->name;
  }
 
  inline bool Equals(const SLayout& rhs) const { return name() == rhs.name(); }
 
  friend std::ostream& operator<<(std::ostream& os, const SLayout& l) {
    os << l.name();
    return os;
  }
 
  TVM_FFI_DEFINE_OBJECT_REF_METHODS_NOTNULLABLE(SLayout, ffi::ObjectRef, SLayoutNode);
};
 
// Internal node container SBijectiveLayout
class SBijectiveLayoutNode : public ffi::Object {
 public:
  explicit SBijectiveLayoutNode(ffi::UnsafeInit tag) : src_layout(tag), dst_layout(tag) {}
 
  SBijectiveLayoutNode(SLayout src, SLayout dst)
      : src_layout(std::move(src)), dst_layout(std::move(dst)) {}
  ffi::Array<PrimExpr> index_forward_rule;
  ffi::Array<PrimExpr> index_backward_rule;
  ffi::Array<PrimExpr> shape_forward_rule;
  ffi::Array<PrimExpr> shape_backward_rule;
 
  SLayout src_layout;
  SLayout dst_layout;
 
  static void RegisterReflection() {
    namespace refl = tvm::ffi::reflection;
    refl::ObjectDef<SBijectiveLayoutNode>()
        .def_ro("src_layout", &SBijectiveLayoutNode::src_layout)
        .def_ro("dst_layout", &SBijectiveLayoutNode::dst_layout)
        .def_ro("index_forward_rule", &SBijectiveLayoutNode::index_forward_rule)
        .def_ro("index_backward_rule", &SBijectiveLayoutNode::index_backward_rule)
        .def_ro("shape_forward_rule", &SBijectiveLayoutNode::shape_forward_rule)
        .def_ro("shape_backward_rule", &SBijectiveLayoutNode::shape_backward_rule);
  }
  TVM_FFI_DECLARE_OBJECT_INFO_FINAL("s_tir.SBijectiveLayout", SBijectiveLayoutNode, ffi::Object);
};
 
class SBijectiveLayout : public ffi::ObjectRef {
 public:
  explicit SBijectiveLayout(ffi::UnsafeInit tag, ffi::ObjectPtr<SBijectiveLayoutNode> data)
      : ffi::ObjectRef(tag) {
    data_ = std::move(data);
  }
 
  TVM_DLL SBijectiveLayout(SLayout src_layout, SLayout dst_layout);
  TVM_DLL static ffi::Optional<SBijectiveLayout> Create(SLayout src_layout, SLayout 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_NOTNULLABLE(SBijectiveLayout, ffi::ObjectRef,
                                                SBijectiveLayoutNode);
};
 
}  // namespace tirx
}  // namespace tvm
 
#endif  // TVM_S_TIR_DATA_LAYOUT_H_