wehub-resource-sync
688 lines
24 KiB
C++
688 lines
24 KiB
C++
/*
|
||
* 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.
|
||
*/
|
||
|
||
/*!
|
||
* \file tvm/relax/nested_msg.h
|
||
* \brief Helper container to store nested message for robust tuple-aware analysis.
|
||
*
|
||
* Please see NestedMsg for description of usage.
|
||
*
|
||
* \sa NestedMsg
|
||
*/
|
||
#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/type.h>
|
||
|
||
#include <string>
|
||
#include <utility>
|
||
#include <vector>
|
||
|
||
namespace tvm {
|
||
namespace relax {
|
||
|
||
/*!
|
||
* \brief Container that stores possibly nested message with leaf message type T.
|
||
*
|
||
* NestedMsg is a helper structure to store intermediate
|
||
* message state in pass analysis so we can robustly handle message
|
||
* passing with the presence of nested tuple types.
|
||
*
|
||
* Under the hood, NestedMsg[T] = Union[T, std::nullopt, Array[NestedMsg[T]]].
|
||
* Each nested message corresponds to the same nesting structure as
|
||
* the nested tuple types when we encounter them in analysis.
|
||
*
|
||
* Relax support nested tuple structures in the IR. Nested tuple structure
|
||
* is important to support advanced groupings in cases such as gradient calculation
|
||
* and other scenarios.
|
||
*
|
||
* The possible presence of nested tuple does mean that we need to
|
||
* to robustly handle analysis that contains nested tuple structures
|
||
* in a dataflow graph.
|
||
*
|
||
* \code
|
||
*
|
||
* v1 = relu(v0)
|
||
* v2 = exp(v0)
|
||
* t = ((v0, v1), (v2,), v0)
|
||
* t1 = t[0]
|
||
* v3 = concat(t1)
|
||
* v4 = t[2]
|
||
* v5 = add(v4, v3)
|
||
*
|
||
* \endcode
|
||
*
|
||
* Consider the above code sequence that contains a mixture of tuple
|
||
* nesting and normal operations. A common message-passing-based analysis
|
||
* will track messages attached to each intermediate variable.
|
||
*
|
||
* Because the intermediate value can contain nested-tuples, we need to have
|
||
* abilities to nest messages according to tuple structure and propagate them
|
||
* along the way. In python, this simply corresponds to using a tuple to hold
|
||
* nested messages. This class provides a helper wrapper in C++ to present such
|
||
* possibly nested message for a given leaf message.
|
||
*
|
||
* This design pattern is necessary to handle tuple values regardless of
|
||
* the normal form design of the IR to enable different messages for each
|
||
* tuple component without enforcing all tuple elements to have the same message.
|
||
*
|
||
* Please consider the following patterns in our pass:
|
||
*
|
||
* On a forward propagation message passing analysis:
|
||
* - Create map [leafnode=>NestedMsg<T>], scan forward
|
||
* - input_msg = [MapToNestedMsg<T>(x, lookup_map) for x in call->args]
|
||
* - output_msg = ForwardProp[call->op](input_msg, call)
|
||
* - map[binding->var] = output_msg
|
||
* - Use MapToNestedMsg to remap the remaining body.
|
||
*
|
||
* On a backward propagation message passing analysis:
|
||
* - Create map [leafnode=>NestedMsg<T>], scan backward
|
||
* - output_msg = lookup map(binding->var)
|
||
* - handle case when output_msg is null
|
||
* - input_msg = BackProp[call->op](out_msg, call)
|
||
* - for arg, msg in zip(call->args, input_msg),
|
||
* DecomposeNestedMessage(arg, msg, lambda node, m: update_map(node, m))
|
||
* - update_map(node, m) => CombineNestedMessage(map[node], m)
|
||
*
|
||
* Here leafnode is a node that you would like to propagate messages to
|
||
* such as constant, var and should not include tuple.
|
||
*
|
||
* We also recommend writing unit-test cases that involve nested tuple composition
|
||
* and decomposition.
|
||
*
|
||
* \sa MapToNestedMsg, DecomposeNestedMsg, CombineNestedMsg, ForEachLeaf, Equal
|
||
*
|
||
* \note If you want to write robust message passing-based analysis for
|
||
* programs that can contain nested tuples, you likely need to
|
||
* use this class or logic of a similar kind.
|
||
*/
|
||
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;
|
||
/*! \brief Nullopt handling */
|
||
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<IntImm>(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; }
|
||
|
||
/*! \return Whether the nested message is not-null leaf value */
|
||
bool IsLeaf() const {
|
||
return data_.type_index() != ffi::TypeIndex::kTVMFFINone &&
|
||
data_.type_index() != ffi::TypeIndex::kTVMFFIArray;
|
||
}
|
||
|
||
/*! \return Whether the nested message is null */
|
||
bool IsNull() const { return data_.type_index() == ffi::TypeIndex::kTVMFFINone; }
|
||
|
||
/*! \return Whether the nested message is nested */
|
||
bool IsNested() const { return data_.type_index() == ffi::TypeIndex::kTVMFFIArray; }
|
||
|
||
/*!
|
||
* \return The underlying leaf value.
|
||
* \note This function checks if the msg is leaf.
|
||
*/
|
||
T LeafValue() const {
|
||
TVM_FFI_ICHECK(IsLeaf());
|
||
return ffi::details::AnyUnsafe::CopyFromAnyViewAfterCheck<T>(data_);
|
||
}
|
||
|
||
/*!
|
||
* \return a corresponding nested array.
|
||
* \note This checks if the underlying data type is array.
|
||
*/
|
||
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;
|
||
};
|
||
|
||
/*!
|
||
* \brief Apply fvisit for each leaf elements in the nested message.
|
||
* \param fvisit The visit callback.
|
||
* \param msg The input nested message.
|
||
* \tparam T the content type of nested msg
|
||
* \tparam FType the visitor type with signature void fvisit(T)
|
||
*/
|
||
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);
|
||
}
|
||
}
|
||
}
|
||
|
||
/*!
|
||
* \brief Recursively compare two nested messages.
|
||
*
|
||
* \param lhs The left operand.
|
||
* \param rhs The right operand.
|
||
* \param fequal The equal functor with signature bool fequal(T, T)
|
||
* \tparam T the content type of nested msg
|
||
* \tparam FType the equal comparator type
|
||
*/
|
||
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;
|
||
}
|
||
}
|
||
|
||
/*!
|
||
* \brief Map expr with possible nested-tuple to nested message.
|
||
*
|
||
* This function will unpack recursive tuples and run fmapleaf for each leaf,
|
||
* then recursively combines the results together into a NestedMsg.
|
||
*
|
||
* The nesting structure will corresponds to the tuple structure.
|
||
*
|
||
* \param expr The input expression.
|
||
* \param fmapleaf The mapping function for each leaf with signature `NestedMsg<T> fmap(Expr)`
|
||
* \tparam T the content type of nested msg
|
||
* \tparam FType The mapping function type
|
||
*/
|
||
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);
|
||
}
|
||
}
|
||
|
||
/*!
|
||
* \brief Map structinfo with possible nested-ty to nested message.
|
||
*
|
||
* This function will unpack recursive ty and run fmapleaf for each leaf,
|
||
* then recursively combines the results together into a NestedMsg.
|
||
*
|
||
* The nesting structure will corresponds to the tuple structure.
|
||
*
|
||
* \param ty The input type.
|
||
* \param fmapleaf The mapping function for each leaf with signature `NestedMsg<T> fmap(Type)`
|
||
* \tparam T the content type of nested msg
|
||
* \tparam FType The mapping function type
|
||
*/
|
||
template <typename T, typename FType>
|
||
NestedMsg<T> MapToNestedMsg(Type ty, FType fmapleaf) {
|
||
if (auto* tuple = ty.as<TupleTypeNode>()) {
|
||
ffi::Array<NestedMsg<T>> res;
|
||
res.reserve(tuple->fields.size());
|
||
for (Type x : tuple->fields) {
|
||
res.push_back(MapToNestedMsg<T, FType>(x, fmapleaf));
|
||
}
|
||
return res;
|
||
} else {
|
||
return fmapleaf(ty);
|
||
}
|
||
}
|
||
|
||
/*!
|
||
* \brief Map expr with possible nested-tuple to nested message.
|
||
*
|
||
* This function will unpack recursive expr by its type and
|
||
* run fmapleaf for each leaf, then recursively combines the results
|
||
* together into a NestedMsg.
|
||
*
|
||
* The nesting structure will corresponds to the type of expr.
|
||
*
|
||
* \param expr The input expression which should have type.
|
||
* \param fmapleaf The mapping function for each leaf with signature `NestedMsg<T> fmapleaf(Expr)`
|
||
* \tparam T the content type of nested msg
|
||
* \tparam FType The mapping function type
|
||
*/
|
||
template <typename T, typename FType>
|
||
NestedMsg<T> MapToNestedMsgByType(Expr expr, FType fmapleaf) {
|
||
auto ty = GetType(expr);
|
||
if (auto* tuple = ty.as<TupleTypeNode>()) {
|
||
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(MapToNestedMsgByType<T, FType>(field, fmapleaf));
|
||
}
|
||
return res;
|
||
} else {
|
||
return fmapleaf(expr);
|
||
}
|
||
}
|
||
|
||
/*!
|
||
* \brief Map nested message back to TargetType.
|
||
*
|
||
* This function will decompose the nested message and
|
||
* run fmapleaf for each leaf message and get the leaf value,
|
||
* then recursively combines the results by fcombine.
|
||
*
|
||
* \param msg The input nested message.
|
||
* \param fmapleaf The mapping function for each leaf with signature
|
||
* `TargetType fmapleaf(ffi::Optional<T>)`.
|
||
* \param fcombine The function for combining all childs of a node into TargetType with signature
|
||
* `TargetType fmapleaf(ffi::Array<TargetType>)`.
|
||
* \tparam TargetType the target type to map nested msg to.
|
||
* \tparam T the content type of nested msg.
|
||
* \tparam FMapLeaf The leaf mapping function type.
|
||
* \tparam FCombine The combining function type.
|
||
*/
|
||
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 {
|
||
TVM_FFI_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);
|
||
}
|
||
}
|
||
|
||
/*!
|
||
* \brief Map nested message back to the expr.
|
||
*
|
||
* This function will decompose the nested message and
|
||
* run fmapleaf for each leaf message and get the leaf expr,
|
||
* then recursively combines the results as tuple expr.
|
||
*
|
||
* \param msg The input nested message.
|
||
* \param fmapleaf The mapping function for each leaf with signature `Expr
|
||
* fmapleaf(ffi::Optional<T>)`. \tparam T the content type of nested msg. \tparam FType The mapping
|
||
* function type.
|
||
*/
|
||
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.has_value()) {
|
||
simplified_flag &= simplified_tuple.value().same_as(node->tuple);
|
||
} else {
|
||
simplified_tuple = node->tuple;
|
||
TVM_FFI_ICHECK(simplified_tuple.has_value());
|
||
}
|
||
}
|
||
}
|
||
}
|
||
return simplified_flag ? simplified_tuple.value() : Tuple(arr);
|
||
});
|
||
}
|
||
|
||
/*!
|
||
* \brief Recursively combine two nested message into one.
|
||
*
|
||
* This function requires the two messages to be compatible with each other.
|
||
* The combination rule is as follows:
|
||
* - combine(null, msg) => msg
|
||
* - combine(leaf1, leaf2) => fcombine(leaf1, leaf2)
|
||
* - combine(array1, array2) => [combine(x, y) for x, y in zip(array1, array2)]
|
||
* - This function will throw an error if array have different size
|
||
*
|
||
* \param lhs The left operand.
|
||
* \param rhs The right operand.
|
||
* \param fcombine with signature T fcombine(T lhs, T rhs)
|
||
* \tparam T the content type of nested msg
|
||
* \tparam FType combine function type.
|
||
*/
|
||
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()) {
|
||
TVM_FFI_ICHECK(rhs.IsLeaf()) << "Cannot combine leaf with nested";
|
||
return NestedMsg<T>(fcombine(lhs.LeafValue(), rhs.LeafValue()));
|
||
} else {
|
||
TVM_FFI_ICHECK(lhs.IsNested());
|
||
TVM_FFI_ICHECK(rhs.IsNested()) << "Cannot combine leaf with nested";
|
||
ffi::Array<NestedMsg<T>> arr_lhs = lhs.NestedArray();
|
||
ffi::Array<NestedMsg<T>> arr_rhs = rhs.NestedArray();
|
||
TVM_FFI_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);
|
||
}
|
||
}
|
||
|
||
/*!
|
||
* \brief Recursively map a nested message to another one, with leaf mapped by the input fmapleaf.
|
||
* \param msg The nested message to be mapped.
|
||
* \param fmapleaf The leaf map function, with signature NestedMsg<T> fmapleaf(T msg)
|
||
* \tparam T The content type of nested message.
|
||
* \tparam FType The leaf map function type.
|
||
* \return The new nested message.
|
||
*/
|
||
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 {
|
||
TVM_FFI_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);
|
||
}
|
||
}
|
||
|
||
/*!
|
||
* \brief Recursively decompose the tuple structure in expr and msg along with it.
|
||
*
|
||
* This function will call fvisitleaf for each leaf expression in expr.
|
||
* This function will throw an error if the nesting structure in msg does not
|
||
* match the tuple nesting structure in expr.
|
||
*
|
||
* \param expr The input expression to be decomposed.
|
||
* \param msg The input nested message.
|
||
* \param fvisitleaf with signature fvisitleaf(Expr expr, NestedMsg<T> msg)
|
||
* \tparam T the content type of nested msg
|
||
* \tparam FType The visit function type.
|
||
*/
|
||
template <typename T, typename FType>
|
||
void DecomposeNestedMsg(Expr expr, NestedMsg<T> msg, FType fvisitleaf) {
|
||
if (auto* tuple = expr.as<TupleNode>()) {
|
||
TVM_FFI_ICHECK(msg.IsNested()) << "Expected nested to match tuple";
|
||
ffi::Array<NestedMsg<T>> arr = msg.NestedArray();
|
||
TVM_FFI_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);
|
||
}
|
||
}
|
||
|
||
/*!
|
||
* \brief Recursively transform the tuple structure in expr and msgs along with it.
|
||
*
|
||
* This function will call ftransleaf for each leaf expression in expr.
|
||
* This function will throw an error if the nesting structure in msg does not
|
||
* match the tuple nesting structure in expr.
|
||
*
|
||
* \param expr The input expression to be transform.
|
||
* \param msgs The input messages to guide the transformation.
|
||
* \param ftransleaf with signature ftransleaf(Expr, ffi::Array<NestedMsg<T>>)->Expr
|
||
* \tparam T the content type of nested msg
|
||
* \tparam N the number of messages
|
||
* \tparam FType The visit function type.
|
||
*/
|
||
template <typename T, std::size_t N, typename FType>
|
||
Expr TransformTupleLeaf(Expr expr, std::array<NestedMsg<T>, N> msgs, FType ftransleaf) {
|
||
Type ty = GetType(expr);
|
||
if (const auto* tuple = ty.as<TupleTypeNode>()) {
|
||
std::array<ffi::Array<NestedMsg<T>>, N> msg_arrays;
|
||
for (size_t i = 0; i < N; ++i) {
|
||
TVM_FFI_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) {
|
||
TVM_FFI_ICHECK(msg.IsLeaf()) << "Expected leaf to match non-tuple";
|
||
}
|
||
return ftransleaf(expr, msgs);
|
||
}
|
||
}
|
||
|
||
/*!
|
||
* \brief Recursively transform the tuple structure in ty and msgs along with it.
|
||
*
|
||
* This function will call ftransleaf for each leaf ty in ty.
|
||
* This function will throw an error if the nesting structure in msg does not
|
||
* match the tuple nesting structure in ty.
|
||
*
|
||
* \param ty The input ty to be transform.
|
||
* \param msgs The input messages to guide the transformation.
|
||
* \param ftransleaf with signature ftransleaf(Type, ffi::Array<NestedMsg<T>>)->Type
|
||
* \tparam T the content type of nested msg
|
||
* \tparam N the number of messages
|
||
* \tparam FType The visit function type.
|
||
*/
|
||
template <typename T, std::size_t N, typename FType>
|
||
Type TransformTupleLeaf(Type ty, std::array<NestedMsg<T>, N> msgs, FType ftransleaf) {
|
||
if (const auto* tuple = ty.as<TupleTypeNode>()) {
|
||
std::array<ffi::Array<NestedMsg<T>>, N> msg_arrays;
|
||
for (size_t i = 0; i < N; ++i) {
|
||
TVM_FFI_ICHECK(msgs[i].IsNested()) << "Expected nested to match tuple";
|
||
msg_arrays[i] = msgs[i].NestedArray();
|
||
}
|
||
bool same = true;
|
||
ffi::Array<Type> fields;
|
||
fields.reserve(tuple->fields.size());
|
||
for (size_t i = 0; i < tuple->fields.size(); ++i) {
|
||
Type 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 ? ty : TupleType(fields);
|
||
} else {
|
||
for (const auto& msg : msgs) {
|
||
TVM_FFI_ICHECK(msg.IsLeaf()) << "Expected leaf to match non-tuple";
|
||
}
|
||
return ftransleaf(ty, 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(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() + ">";
|
||
}
|
||
|
||
TVM_FFI_INLINE static std::string TypeSchema() {
|
||
std::ostringstream oss;
|
||
oss << R"({"type":"NestedMsg","args":[)";
|
||
oss << details::TypeSchema<T>::v();
|
||
oss << "]}";
|
||
return oss.str();
|
||
}
|
||
};
|
||
} // namespace ffi
|
||
} // namespace tvm
|
||
#endif // TVM_RELAX_NESTED_MSG_H_
|