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/ffi/container/array.h>
#include <tvm/ffi/optional.h>
#include <tvm/relax/expr.h>
#include <tvm/relax/struct_info.h>
 
#include <string>
#include <utility>
#include <vector>
 
namespace tvm {
namespace relax {
 
template <typename T>
class NestedMsg {
 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;
  NestedMsg(std::nullopt_t) {}  // 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(*)
      : data_(std::move(other)) {}
  NestedMsg<T>& operator=(T other) {
    data_ = std::move(other);
    return *this;
  }
  // ffi::Array<NestedMsg<T>> handling
  NestedMsg(ffi::Array<NestedMsg<T>, void> other)  // NOLINT(*)
      : data_(other) {}
 
  NestedMsg<T>& operator=(ffi::Array<NestedMsg<T>, void> other) {
    data_ = std::move(other);
    return *this;
  }
 
  // initializer list handling
  NestedMsg(std::initializer_list<NestedMsg<T>> other)  // NOLINT(*)
      : NestedMsg(ffi::Array<NestedMsg<T>, void>(other)) {}
  NestedMsg<T>& operator=(std::initializer_list<NestedMsg<T>> other) {
    return operator=(ffi::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_.type_index() != ffi::TypeIndex::kTVMFFINone &&
           data_.type_index() != ffi::TypeIndex::kTVMFFIArray;
  }
 
  bool IsNull() const { return data_.type_index() == ffi::TypeIndex::kTVMFFINone; }
 
  bool IsNested() const { return data_.type_index() == ffi::TypeIndex::kTVMFFIArray; }
 
  T LeafValue() const {
    ICHECK(IsLeaf());
    return ffi::details::AnyUnsafe::CopyFromAnyViewAfterCheck<T>(data_);
  }
 
  ffi::Array<NestedMsg<T>, void> NestedArray() const {
    return ffi::details::AnyUnsafe::CopyFromAnyViewAfterCheck<ffi::Array<NestedMsg<T>, void>>(
        data_);
  }
 
 private:
  ffi::Any data_;
  // private constructor
  explicit NestedMsg(ffi::Any data) : data_(data) {}
  template <typename, typename>
  friend struct ffi::TypeTraits;
};
 
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;
    ffi::Array<NestedMsg<T>> arr_lhs = lhs.NestedArray();
    ffi::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>()) {
    ffi::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>()) {
    ffi::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>()) {
    ffi::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(std::nullopt);
  } else if (msg.IsLeaf()) {
    return fmapleaf(msg.LeafValue());
  } else {
    ICHECK(msg.IsNested());
    ffi::Array<NestedMsg<T>> arr = msg.NestedArray();
    ffi::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, [](ffi::Array<Expr> arr) {
    ffi::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";
    ffi::Array<NestedMsg<T>> arr_lhs = lhs.NestedArray();
    ffi::Array<NestedMsg<T>> arr_rhs = rhs.NestedArray();
    ICHECK_EQ(arr_lhs.size(), arr_rhs.size())
        << "Cannot combine two nested array with different sizes";
    ffi::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());
    ffi::Array<NestedMsg<T>> arr = msg.NestedArray();
    ffi::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";
    ffi::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<ffi::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;
    ffi::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<ffi::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;
    ffi::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 ffi {
 
template <typename T>
inline constexpr bool use_default_type_traits_v<relax::NestedMsg<T>> = false;
 
template <typename T>
struct TypeTraits<relax::NestedMsg<T>> : public TypeTraitsBase {
  TVM_FFI_INLINE static void CopyToAnyView(const relax::NestedMsg<T>& src, TVMFFIAny* result) {
    *result = ffi::AnyView(src.data_).CopyToTVMFFIAny();
  }
 
  TVM_FFI_INLINE static void MoveToAny(relax::NestedMsg<T> src, TVMFFIAny* result) {
    *result = details::AnyUnsafe::MoveAnyToTVMFFIAny(std::move(src.data_));
  }
 
  TVM_FFI_INLINE static std::string GetMismatchTypeInfo(const TVMFFIAny* src) {
    return TypeTraitsBase::GetMismatchTypeInfo(src);
  }
 
  static bool CheckAnyStrict(const TVMFFIAny* src) {
    if (src->type_index == TypeIndex::kTVMFFINone) {
      return true;
    }
    if (TypeTraits<T>::CheckAnyStrict(src)) {
      return true;
    }
    if (src->type_index == TypeIndex::kTVMFFIArray) {
      const ffi::ArrayObj* array = reinterpret_cast<const ffi::ArrayObj*>(src->v_obj);
      for (size_t i = 0; i < array->size(); ++i) {
        const Any& any_v = (*array)[i];
        if (!details::AnyUnsafe::CheckAnyStrict<relax::NestedMsg<T>>(any_v)) return false;
      }
    }
    return true;
  }
 
  TVM_FFI_INLINE static relax::NestedMsg<T> CopyFromAnyViewAfterCheck(const TVMFFIAny* src) {
    return relax::NestedMsg<T>(Any(AnyView::CopyFromTVMFFIAny(*src)));
  }
 
  TVM_FFI_INLINE static relax::NestedMsg<T> MoveFromAnyAfterCheck(TVMFFIAny* src) {
    return relax::NestedMsg<T>(details::AnyUnsafe::MoveTVMFFIAnyToAny(std::move(*src)));
  }
 
  static std::optional<relax::NestedMsg<T>> TryCastFromAnyView(const TVMFFIAny* src) {
    if (CheckAnyStrict(src)) {
      return CopyFromAnyViewAfterCheck(src);
    }
    // slow path run conversion
    if (src->type_index == TypeIndex::kTVMFFINone) {
      return relax::NestedMsg<T>(std::nullopt);
    }
    if (auto opt_value = TypeTraits<T>::TryCastFromAnyView(src)) {
      return relax::NestedMsg<T>(*std::move(opt_value));
    }
    if (src->type_index == TypeIndex::kTVMFFIArray) {
      const ArrayObj* n = reinterpret_cast<const ArrayObj*>(src->v_obj);
      ffi::Array<relax::NestedMsg<T>> result;
      result.reserve(n->size());
      for (size_t i = 0; i < n->size(); i++) {
        const Any& any_v = (*n)[i];
        if (auto opt_v = any_v.try_cast<relax::NestedMsg<T>>()) {
          result.push_back(*std::move(opt_v));
        } else {
          return std::nullopt;
        }
      }
      return relax::NestedMsg<T>(result);
    }
    return std::nullopt;
  }
 
  TVM_FFI_INLINE static std::string TypeStr() {
    return "NestedMsg<" + details::Type2Str<T>::v() + ">";
  }
};
}  // namespace ffi
}  // namespace tvm
#endif  // TVM_RELAX_NESTED_MSG_H_