nested_msg.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_RELAX_NESTED_MSG_H_
#define TVM_RELAX_NESTED_MSG_H_
 
#include <tvm/relax/expr.h>
#include <tvm/relax/struct_info.h>
#include <tvm/runtime/container/array.h>
#include <tvm/runtime/container/optional.h>
 
#include <utility>
#include <vector>
 
namespace tvm {
namespace relax {
 
template <typename T>
class NestedMsg : public ObjectRef {
 public:
  // default constructors.
  NestedMsg() = default;
  NestedMsg(const NestedMsg<T>&) = default;
  NestedMsg(NestedMsg<T>&&) = default;
  NestedMsg<T>& operator=(const NestedMsg<T>&) = default;
  NestedMsg<T>& operator=(NestedMsg<T>&&) = default;
  explicit NestedMsg(ObjectPtr<Object> ptr) : ObjectRef(ptr) {}
  NestedMsg(runtime::NullOptType) {}  // NOLINT(*)
  // nullptr handling.
  // disallow implicit conversion as 0 can be implicitly converted to nullptr_t
  explicit NestedMsg(std::nullptr_t) {}
  NestedMsg<T>& operator=(std::nullptr_t) {
    data_ = nullptr;
    return *this;
  }
  // normal value handling.
  NestedMsg(T other)  // NOLINT(*)
      : ObjectRef(std::move(other)) {}
  NestedMsg<T>& operator=(T other) {
    ObjectRef::operator=(std::move(other));
    return *this;
  }
  // Array<NestedMsg<T>> handling
  NestedMsg(Array<NestedMsg<T>, void> other)  // NOLINT(*)
      : ObjectRef(std::move(other)) {}
  NestedMsg<T>& operator=(Array<NestedMsg<T>, void> other) {
    ObjectRef::operator=(std::move(other));
    return *this;
  }
 
  // initializer list handling
  NestedMsg(std::initializer_list<NestedMsg<T>> other)  // NOLINT(*)
      : NestedMsg(Array<NestedMsg<T>, void>(other)) {}
  NestedMsg<T>& operator=(std::initializer_list<NestedMsg<T>> other) {
    return operator=(Array<NestedMsg<T>, void>(other));
  }
 
  // delete the int constructor
  // since NestedMsg<Integer>(0) is ambiguous
  // 0 can be implicitly casted to nullptr_t
  explicit NestedMsg(int val) = delete;
  NestedMsg<T>& operator=(int val) = delete;
  // operator overloadings
  bool operator==(std::nullptr_t) const { return data_ == nullptr; }
  bool operator!=(std::nullptr_t) const { return data_ != nullptr; }
 
  bool IsLeaf() const { return data_ != nullptr && data_->IsInstance<LeafContainerType>(); }
 
  bool IsNull() const { return data_ == nullptr; }
 
  bool IsNested() const { return data_ != nullptr && data_->IsInstance<ArrayNode>(); }
 
  T LeafValue() const {
    ICHECK(IsLeaf());
    return T(data_);
  }
 
  Array<NestedMsg<T>, void> NestedArray() const {
    ICHECK(IsNested());
    return Array<NestedMsg<T>, void>(data_);
  }
 
  using ContainerType = Object;
  using LeafContainerType = typename T::ContainerType;
 
  static_assert(std::is_base_of<ObjectRef, T>::value, "NestedMsg is only defined for ObjectRef.");
 
  static constexpr bool _type_is_nullable = true;
};
 
template <typename T, typename FType>
void ForEachLeaf(const NestedMsg<T>& msg, FType fvisit) {
  if (msg == nullptr) return;
  if (msg.IsLeaf()) {
    fvisit(msg.LeafValue());
  } else {
    for (NestedMsg<T> x : msg.NestedArray()) {
      ForEachLeaf(x, fvisit);
    }
  }
}
 
template <typename T, typename FType>
bool Equal(const NestedMsg<T>& lhs, const NestedMsg<T>& rhs, FType fequal) {
  if (lhs.IsNull()) return rhs.IsNull();
  if (rhs.IsNull()) return lhs.IsNull();
  if (lhs.IsLeaf()) {
    return rhs.IsLeaf() && fequal(lhs.LeafValue(), rhs.LeafValue());
  } else {
    if (!rhs.IsNested()) return false;
    Array<NestedMsg<T>> arr_lhs = lhs.NestedArray();
    Array<NestedMsg<T>> arr_rhs = rhs.NestedArray();
    if (arr_lhs.size() != arr_rhs.size()) return false;
    for (size_t i = 0; i < arr_lhs.size(); ++i) {
      if (!Equal(arr_lhs[i], arr_rhs[i], fequal)) return false;
    }
    return true;
  }
}
 
template <typename T, typename FType>
NestedMsg<T> MapToNestedMsg(Expr expr, FType fmapleaf) {
  if (auto* tuple = expr.as<TupleNode>()) {
    Array<NestedMsg<T>> res;
    res.reserve(tuple->fields.size());
    for (Expr x : tuple->fields) {
      res.push_back(MapToNestedMsg<T, FType>(x, fmapleaf));
    }
    return res;
  } else {
    return fmapleaf(expr);
  }
}
 
template <typename T, typename FType>
NestedMsg<T> MapToNestedMsg(StructInfo sinfo, FType fmapleaf) {
  if (auto* tuple = sinfo.as<TupleStructInfoNode>()) {
    Array<NestedMsg<T>> res;
    res.reserve(tuple->fields.size());
    for (StructInfo x : tuple->fields) {
      res.push_back(MapToNestedMsg<T, FType>(x, fmapleaf));
    }
    return res;
  } else {
    return fmapleaf(sinfo);
  }
}
 
template <typename T, typename FType>
NestedMsg<T> MapToNestedMsgBySInfo(Expr expr, FType fmapleaf) {
  auto sinfo = GetStructInfo(expr);
  if (auto* tuple = sinfo.as<TupleStructInfoNode>()) {
    Array<NestedMsg<T>> res;
    res.reserve(tuple->fields.size());
    for (size_t i = 0; i < tuple->fields.size(); ++i) {
      Expr field;
      if (const auto* expr_tuple = expr.as<TupleNode>()) {
        field = expr_tuple->fields[i];
      } else {
        field = TupleGetItem(expr, i);
      }
      res.push_back(MapToNestedMsgBySInfo<T, FType>(field, fmapleaf));
    }
    return res;
  } else {
    return fmapleaf(expr);
  }
}
 
template <typename TargetType, typename T, typename FMapLeaf, typename FCombine>
TargetType NestedMsgTo(NestedMsg<T> msg, FMapLeaf fmapleaf, FCombine fcombine) {
  if (msg.IsNull()) {
    return fmapleaf(NullOpt);
  } else if (msg.IsLeaf()) {
    return fmapleaf(msg.LeafValue());
  } else {
    ICHECK(msg.IsNested());
    Array<NestedMsg<T>> arr = msg.NestedArray();
    Array<TargetType> subexpr;
    subexpr.reserve(arr.size());
    for (size_t i = 0; i < arr.size(); ++i) {
      subexpr.push_back(NestedMsgTo<TargetType>(arr[i], fmapleaf, fcombine));
    }
    return fcombine(subexpr);
  }
}
 
template <typename T, typename FType>
Expr NestedMsgToExpr(NestedMsg<T> msg, FType fmapleaf) {
  return NestedMsgTo<Expr>(msg, fmapleaf, [](Array<Expr> arr) {
    Optional<Expr> simplified_tuple;
    bool simplified_flag = false;
    if (arr.size() >= 1) {
      simplified_flag = true;
      for (size_t i = 0; i < arr.size() && simplified_flag; ++i) {
        auto* node = arr[i].as<TupleGetItemNode>();
        if (node == nullptr || node->index != static_cast<int>(i)) {
          simplified_flag = false;
        } else {
          if (simplified_tuple.defined()) {
            simplified_flag &= (simplified_tuple == node->tuple);
          } else {
            simplified_tuple = node->tuple;
            ICHECK(simplified_tuple.defined());
          }
        }
      }
    }
    return simplified_flag ? simplified_tuple.value() : Tuple(arr);
  });
}
 
template <typename T, typename FType>
NestedMsg<T> CombineNestedMsg(NestedMsg<T> lhs, NestedMsg<T> rhs, FType fcombine) {
  if (lhs.IsNull()) return rhs;
  if (rhs.IsNull()) return lhs;
 
  if (lhs.IsLeaf()) {
    ICHECK(rhs.IsLeaf()) << "Cannot combine leaf with nested";
    return NestedMsg<T>(fcombine(lhs.LeafValue(), rhs.LeafValue()));
  } else {
    ICHECK(lhs.IsNested());
    ICHECK(rhs.IsNested()) << "Cannot combine leaf with nested";
    Array<NestedMsg<T>> arr_lhs = lhs.NestedArray();
    Array<NestedMsg<T>> arr_rhs = rhs.NestedArray();
    ICHECK_EQ(arr_lhs.size(), arr_rhs.size())
        << "Cannot combine two nested array with different sizes";
    Array<NestedMsg<T>> res;
    res.reserve(arr_lhs.size());
    for (size_t i = 0; i < arr_lhs.size(); ++i) {
      res.push_back(CombineNestedMsg<T, FType>(arr_lhs[i], arr_rhs[i], fcombine));
    }
    return NestedMsg<T>(res);
  }
}
 
template <typename T, typename FType>
NestedMsg<T> MapNestedMsg(NestedMsg<T> msg, FType fmapleaf) {
  if (msg.IsNull()) {
    return msg;
  } else if (msg.IsLeaf()) {
    return fmapleaf(msg.LeafValue());
  } else {
    ICHECK(msg.IsNested());
    Array<NestedMsg<T>> arr = msg.NestedArray();
    Array<NestedMsg<T>> res;
    res.reserve(arr.size());
    for (int i = 0; i < static_cast<int>(arr.size()); ++i) {
      res.push_back(MapNestedMsg(arr[i], fmapleaf));
    }
    return NestedMsg<T>(res);
  }
}
 
template <typename T, typename FType>
void DecomposeNestedMsg(Expr expr, NestedMsg<T> msg, FType fvisitleaf) {
  if (auto* tuple = expr.as<TupleNode>()) {
    ICHECK(msg.IsNested()) << "Expected nested to match tuple";
    Array<NestedMsg<T>> arr = msg.NestedArray();
    ICHECK_EQ(arr.size(), tuple->fields.size()) << "Expected nested array size to match tuple size";
    for (size_t i = 0; i < arr.size(); ++i) {
      DecomposeNestedMsg(tuple->fields[i], arr[i], fvisitleaf);
    }
  } else {
    fvisitleaf(expr, msg);
  }
}
 
template <typename T, std::size_t N, typename FType>
Expr TransformTupleLeaf(Expr expr, std::array<NestedMsg<T>, N> msgs, FType ftransleaf) {
  StructInfo sinfo = GetStructInfo(expr);
  if (const auto* tuple = sinfo.as<TupleStructInfoNode>()) {
    std::array<Array<NestedMsg<T>>, N> msg_arrays;
    for (size_t i = 0; i < N; ++i) {
      ICHECK(msgs[i].IsNested()) << "Expected nested to match tuple";
      msg_arrays[i] = msgs[i].NestedArray();
    }
    bool same = true;
    Array<Expr> fields;
    fields.reserve(tuple->fields.size());
    for (size_t i = 0; i < tuple->fields.size(); ++i) {
      Expr field;
      if (const auto* expr_tuple = expr.as<TupleNode>()) {
        field = expr_tuple->fields[i];
      } else {
        field = TupleGetItem(expr, i);
      }
      std::array<NestedMsg<T>, N> sub_msgs;
      for (size_t j = 0; j < N; ++j) {
        sub_msgs[j] = msg_arrays[j][i];
      }
      fields.push_back(TransformTupleLeaf(field, std::move(sub_msgs), ftransleaf));
      same &= (fields.back().same_as(field));
    }
    return same ? expr : Tuple(fields);
  } else {
    for (const auto& msg : msgs) {
      ICHECK(msg.IsLeaf()) << "Expected leaf to match non-tuple";
    }
    return ftransleaf(expr, msgs);
  }
}
 
template <typename T, std::size_t N, typename FType>
StructInfo TransformTupleLeaf(StructInfo sinfo, std::array<NestedMsg<T>, N> msgs,
                              FType ftransleaf) {
  if (const auto* tuple = sinfo.as<TupleStructInfoNode>()) {
    std::array<Array<NestedMsg<T>>, N> msg_arrays;
    for (size_t i = 0; i < N; ++i) {
      ICHECK(msgs[i].IsNested()) << "Expected nested to match tuple";
      msg_arrays[i] = msgs[i].NestedArray();
    }
    bool same = true;
    Array<StructInfo> fields;
    fields.reserve(tuple->fields.size());
    for (size_t i = 0; i < tuple->fields.size(); ++i) {
      StructInfo field = tuple->fields[i];
      std::array<NestedMsg<T>, N> sub_msgs;
      for (size_t j = 0; j < N; ++j) {
        sub_msgs[j] = msg_arrays[j][i];
      }
      fields.push_back(TransformTupleLeaf(field, std::move(sub_msgs), ftransleaf));
      same &= (fields.back().same_as(field));
    }
    return same ? sinfo : TupleStructInfo(fields);
  } else {
    for (const auto& msg : msgs) {
      ICHECK(msg.IsLeaf()) << "Expected leaf to match non-tuple";
    }
    return ftransleaf(sinfo, msgs);
  }
}
 
}  // namespace relax
}  // namespace tvm
#endif  // TVM_RELAX_NESTED_MSG_H_