You are viewing a plain text version of this content. The canonical link for it is here.
Posted to commits@tvm.apache.org by tq...@apache.org on 2026/09/28 16:03:22 UTC
(tvm) branch main updated: Add context-free Call reinference and binding type propagation (#20480)
This is an automated email from the ASF dual-hosted git repository.
tqchen pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/tvm.git
The following commit(s) were added to refs/heads/main by this push:
new a7fffaa63a Add context-free Call reinference and binding type propagation (#20480)
a7fffaa63a is described below
commit a7fffaa63a5e91c811a536100aa51921b7937d9d
Author: Tianqi Chen <tq...@users.noreply.github.com>
AuthorDate: Mon Sep 28 12:03:16 2026 -0400
Add context-free Call reinference and binding type propagation (#20480)
- Separate native context-free Call type inference from
builder-dependent Relax and distributed inference, with strict read-only
reinference from current Call inputs.
- Infer `tirx.buffer_data` from its BufferVar argument and reinfer after
pointer storage-scope rewrites.
- Migrate TVM structural visit and mutation hooks to native tvm-ffi
callbacks and remove the structural compatibility header.
- Propagate Bind and Let value type changes through their binder
definitions and later uses while preserving unchanged node identity.
- Reuse existing Call result Type objects when inference establishes an
equivalent result cheaply.
---
3rdparty/tvm-ffi | 2 +-
include/tvm/ir/base_expr.h | 4 +-
include/tvm/ir/expr.h | 10 +
include/tvm/ir/expr_functor.h | 13 +
include/tvm/ir/ffi_structural_compat.h | 118 ------
include/tvm/ir/op.h | 3 +
include/tvm/relax/op_attr_types.h | 2 +-
python/tvm/ir/__init__.py | 1 +
python/tvm/ir/expr.py | 9 +
.../relax/transform/legalize_ops/linear_algebra.py | 2 +-
.../tvm/relax/transform/legalize_ops/manipulate.py | 4 +-
src/ir/expr.cc | 364 ++++++++++-------
src/ir/expr_functor.cc | 49 ++-
src/ir/op.cc | 21 +-
src/ir/prim/expr.cc | 382 +++++++++++-------
src/ir/prim/vector_expr.cc | 79 ++--
src/ir/type.cc | 154 ++++---
src/relax/analysis/well_formed.cc | 5 +-
.../backend/adreno/annotate_custom_storage.cc | 23 +-
src/relax/backend/vm/vm_shape_lower.cc | 5 +-
src/relax/distributed/type.cc | 26 +-
src/relax/ir/block_builder.cc | 25 +-
src/relax/ir/dependent_type.cc | 76 ++--
src/relax/ir/expr.cc | 158 ++++----
src/relax/ir/type.cc | 23 +-
src/relax/op/ccl/ccl.cc | 35 +-
src/relax/op/distributed/binary.cc | 53 +--
src/relax/op/distributed/ccl.cc | 2 +-
src/relax/op/distributed/distributed.cc | 14 +-
src/relax/op/distributed/linear_algebra.cc | 2 +-
src/relax/op/distributed/manipulate.cc | 4 +-
src/relax/op/distributed/nn.cc | 2 +-
src/relax/op/distributed/op.cc | 5 +-
src/relax/op/distributed/statistical.h | 2 +-
src/relax/op/distributed/unary.h | 4 +-
src/relax/op/image/resize.cc | 52 ++-
src/relax/op/memory/view.cc | 11 +-
src/relax/op/nn/attention.cc | 6 +-
src/relax/op/nn/convolution.cc | 12 +-
src/relax/op/nn/nn.cc | 109 +++--
src/relax/op/nn/pooling.cc | 66 +--
src/relax/op/op.cc | 172 +++++---
src/relax/op/op_common.cc | 16 +-
src/relax/op/op_common.h | 174 ++++++--
src/relax/op/tensor/binary.cc | 5 +-
src/relax/op/tensor/binary.h | 10 +-
src/relax/op/tensor/create.cc | 93 +++--
src/relax/op/tensor/datatype.cc | 22 +-
src/relax/op/tensor/grad.cc | 65 ++-
src/relax/op/tensor/index.cc | 26 +-
src/relax/op/tensor/inspect.cc | 16 +-
src/relax/op/tensor/linear_algebra.cc | 24 +-
src/relax/op/tensor/manipulate.cc | 149 ++++---
src/relax/op/tensor/qdq.cc | 4 +-
src/relax/op/tensor/sampling.cc | 18 +-
src/relax/op/tensor/search.cc | 28 +-
src/relax/op/tensor/set.cc | 22 +-
src/relax/op/tensor/sorting.cc | 33 +-
src/relax/op/tensor/statistical.cc | 51 ++-
src/relax/op/tensor/ternary.cc | 2 +-
src/relax/op/tensor/unary.cc | 115 ++++--
src/relax/op/vision/multibox_transform_loc.cc | 11 +-
src/relax/op/vision/nms.cc | 22 +-
src/relax/op/vision/roi_align.cc | 9 +-
src/relax/op/vision/roi_pool.cc | 9 +-
src/relax/script/printer/utils.h | 15 +-
src/relax/transform/compute_prim_value.cc | 5 +-
src/relax/transform/fold_constant.cc | 5 +-
src/relax/transform/gradient.cc | 4 +-
src/relax/transform/run_codegen.cc | 4 +-
src/s_tir/stmt.cc | 82 ++--
src/te/tensor.cc | 21 +-
src/tirx/ir/function.cc | 25 +-
src/tirx/ir/iter_var.cc | 27 +-
src/tirx/ir/layout/tile_core.cc | 53 +--
src/tirx/ir/stmt.cc | 443 ++++++++++++---------
src/tirx/ir/stmt_functor.cc | 25 +-
src/tirx/ir/tirx_stmt.cc | 28 +-
src/tirx/ir/type.cc | 74 ++--
src/tirx/op/builtin.cc | 13 +
src/tirx/transform/update_pointer_storage_scope.cc | 16 +-
tests/python/relax/test_expr.py | 49 +++
tests/python/relax/test_op_distributed.py | 6 +
tests/python/relax/test_transform_legalize_ops.py | 10 +-
...st_transform_operator_specific_normalization.py | 2 +-
.../tirx-base/test_binding_type_propagation.py | 166 ++++++++
tests/python/tirx/test_buffer_data_reinfer_type.py | 39 ++
87 files changed, 2661 insertions(+), 1484 deletions(-)
diff --git a/3rdparty/tvm-ffi b/3rdparty/tvm-ffi
index 0af62893f8..4ee5a8621a 160000
--- a/3rdparty/tvm-ffi
+++ b/3rdparty/tvm-ffi
@@ -1 +1 @@
-Subproject commit 0af62893f8033b65c3f424a0dbfd3138b5cbe5d4
+Subproject commit 4ee5a8621a075150ba1f2356893397d9d08d073e
diff --git a/include/tvm/ir/base_expr.h b/include/tvm/ir/base_expr.h
index 3433550676..87855537b4 100644
--- a/include/tvm/ir/base_expr.h
+++ b/include/tvm/ir/base_expr.h
@@ -26,10 +26,10 @@
#include <tvm/ffi/cast.h>
#include <tvm/ffi/dtype.h>
+#include <tvm/ffi/extra/structural_mutate.h>
+#include <tvm/ffi/extra/structural_visit.h>
#include <tvm/ffi/reflection/registry.h>
#include <tvm/ffi/string.h>
-// Keep raw and typed structural hook return paths available to all TVM IR nodes.
-#include <tvm/ir/ffi_structural_compat.h>
#include <tvm/ir/source_map.h>
#include <cstddef>
diff --git a/include/tvm/ir/expr.h b/include/tvm/ir/expr.h
index 7bf93aa25d..d30833e8a6 100644
--- a/include/tvm/ir/expr.h
+++ b/include/tvm/ir/expr.h
@@ -385,6 +385,9 @@ class Var : public Expr {
/*! \brief Return a fresh ordinary Var with a new primitive type. */
TVM_DLL Var CopyWithDType(PrimType dtype) const;
+ /*! \brief Return a fresh ordinary Var with a new type, retaining its metadata. */
+ TVM_DLL Var CopyWithType(Type type) const;
+
TVM_FFI_DEFINE_OBJECT_REF_METHODS_NULLABLE(Var, Expr, VarNode);
};
@@ -513,6 +516,13 @@ class Call : public Expr {
/*! \brief Construct a provisional Call without invoking its Op validator. */
TVM_DLL static Call Unchecked(Type ret_ty, Expr op, ffi::Array<Expr> args, Attrs attrs = Attrs(),
ffi::Array<Type> ty_args = ffi::Array<Type>(), Span span = Span());
+ /*! \brief Recompute a result type from the Call's current explicit inputs.
+ *
+ * This ignores the Call's stored result type and does not mutate the Call.
+ * A context-free inference hook must be registered for the operator. Missing
+ * or invalid inputs are reported by the hook rather than screened here.
+ */
+ TVM_DLL static Type ReinferType(const CallNode* call);
TVM_FFI_DEFINE_OBJECT_REF_METHODS_NULLABLE(Call, Expr, CallNode);
TVM_DEFINE_OBJECT_REF_COW_METHOD(CallNode);
diff --git a/include/tvm/ir/expr_functor.h b/include/tvm/ir/expr_functor.h
index cc735b4db2..95bd226a77 100644
--- a/include/tvm/ir/expr_functor.h
+++ b/include/tvm/ir/expr_functor.h
@@ -510,6 +510,19 @@ class TVM_DLL ExprMutator : public ObjectMutator {
virtual UnchangedOr<PrimExpr> Mutate_(const prim::ShuffleNode* node, InplaceMode inplace_mode);
protected:
+ /*!
+ * \brief Reinfer the result type of a Call after ordinary mutation.
+ * \param mutated The ordinary mutation result, or Unchanged for the original Call.
+ * \param original The borrowed Call before ordinary mutation.
+ * \param inplace_mode Permission to modify the original Call in place.
+ * \return The combined mutation result, preserving Unchanged when possible.
+ * \note A unique replacement Call may be updated in place even when the original
+ * could not be. Inference errors propagate to the caller.
+ */
+ static UnchangedOr<Expr> ReinferMutatedCallType(UnchangedOr<Expr> mutated,
+ const CallNode* original,
+ InplaceMode inplace_mode);
+
/*!
* \brief Construct a mutator with an extended native dispatch table.
* \param vtable The finalized table, which must outlive this mutator.
diff --git a/include/tvm/ir/ffi_structural_compat.h b/include/tvm/ir/ffi_structural_compat.h
deleted file mode 100644
index d892cbf8c8..0000000000
--- a/include/tvm/ir/ffi_structural_compat.h
+++ /dev/null
@@ -1,118 +0,0 @@
-/*
- * 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 ffi_structural_compat.h
- * \brief Keep TVM's raw structural hooks compatible with typed tvm-ffi traversal.
- */
-#ifndef TVM_IR_FFI_STRUCTURAL_COMPAT_H_
-#define TVM_IR_FFI_STRUCTURAL_COMPAT_H_
-
-#include <tvm/ffi/extra/structural_mutate.h>
-#include <tvm/ffi/extra/structural_visit.h>
-
-#include <type_traits>
-#include <utility>
-
-namespace tvm {
-namespace structural_compat {
-
-template <typename T>
-class ExpectedReturn {
- public:
- explicit ExpectedReturn(ffi::Expected<T>&& value) : value_(std::move(value)) {}
-
- operator TVMFFIAny() && noexcept {
- return ffi::details::ExpectedUnsafe::MoveToTVMFFIAny(std::move(value_));
- }
-
- operator ffi::Expected<T>() && noexcept { return std::move(value_); }
-
- private:
- ffi::Expected<T> value_;
-};
-
-class ErrorReturn {
- public:
- explicit ErrorReturn(ffi::Unexpected<ffi::Error>&& value) : value_(std::move(value)) {}
-
- operator TVMFFIAny() && noexcept {
- return ffi::details::ExpectedUnsafe::MoveToTVMFFIAny(
- ffi::Expected<ffi::Any>(std::move(value_)));
- }
-
- template <typename T>
- operator ffi::Expected<T>() && noexcept {
- return std::move(value_);
- }
-
- private:
- ffi::Unexpected<ffi::Error> value_;
-};
-
-template <typename T>
-auto VisitReturn(T&& result) {
- if constexpr (std::is_same_v<std::remove_cv_t<std::remove_reference_t<T>>,
- ffi::Optional<ffi::VisitInterrupt>>) {
- return std::forward<T>(result);
- } else {
- return ExpectedReturn(std::forward<T>(result));
- }
-}
-
-} // namespace structural_compat
-} // namespace tvm
-
-// tvm-ffi's macros now return Expected directly. TVM still has raw ABI hooks,
-// so select the representation from the enclosing hook's declared return type.
-#undef TVM_FFI_S_VISIT_MAYBE_EARLY_RETURN
-#define TVM_FFI_S_VISIT_MAYBE_EARLY_RETURN(Result) \
- do { \
- auto&& tvm_ffi_res_ = (Result); \
- if (TVM_FFI_PREDICT_FALSE( \
- ::tvm::ffi::details::StructuralVisitNeedEarlyReturn(tvm_ffi_res_))) { \
- return ::tvm::structural_compat::VisitReturn(::std::move(tvm_ffi_res_)); \
- } \
- } while (0)
-
-#undef TVM_FFI_S_MUTATE_MAYBE_EARLY_RETURN
-#define TVM_FFI_S_MUTATE_MAYBE_EARLY_RETURN(Result) \
- do { \
- auto&& tvm_ffi_res_ = (Result); \
- if (TVM_FFI_PREDICT_FALSE(tvm_ffi_res_.is_err())) { \
- return ::tvm::structural_compat::ErrorReturn( \
- ::tvm::ffi::Unexpected(::std::move(tvm_ffi_res_).error())); \
- } \
- } while (0)
-
-#undef TVM_FFI_S_MUTATE_ASSIGN_OR_RETURN_IMPL_
-#define TVM_FFI_S_MUTATE_ASSIGN_OR_RETURN_IMPL_(Result, Type, Name, ResultExpr) \
- auto Result = (ResultExpr); /* NOLINT(bugprone-macro-parentheses) */ \
- TVM_FFI_S_MUTATE_MAYBE_EARLY_RETURN(Result); \
- if constexpr (!::tvm::ffi::type_subsumes_v<::tvm::ffi::Expected<Type>, decltype(Result)>) { \
- if (TVM_FFI_PREDICT_FALSE(!::tvm::ffi::details::AnyUnsafe::CheckAnyStrict<Type>( \
- ::tvm::ffi::details::ExpectedUnsafe::GetData(Result)))) { \
- return ::tvm::structural_compat::ErrorReturn( \
- ::tvm::ffi::details::SMutateDeclaredTypeError()); \
- } \
- } \
- Type Name = /* NOLINT(bugprone-macro-parentheses) */ \
- ::tvm::ffi::details::AnyUnsafe::MoveFromAnyAfterCheck<Type>( \
- ::std::move(::tvm::ffi::details::ExpectedUnsafe::GetData(Result)))
-
-#endif // TVM_IR_FFI_STRUCTURAL_COMPAT_H_
diff --git a/include/tvm/ir/op.h b/include/tvm/ir/op.h
index 49fb8223e1..c1c704db13 100644
--- a/include/tvm/ir/op.h
+++ b/include/tvm/ir/op.h
@@ -43,6 +43,9 @@ namespace tvm {
template <typename>
class OpAttrMap;
+/*! \brief Infer a Call's result type from its explicit inputs without builder state. */
+using FInferType = ffi::reflection::NativeFunctionView<ffi::Expected<Type>(const CallNode* call)>;
+
/*! \brief An operator argument's name and documentation. */
class ArgumentInfoNode : public ffi::Object {
public:
diff --git a/include/tvm/relax/op_attr_types.h b/include/tvm/relax/op_attr_types.h
index 51d6735534..1ff6749116 100644
--- a/include/tvm/relax/op_attr_types.h
+++ b/include/tvm/relax/op_attr_types.h
@@ -66,7 +66,7 @@ using FCallPacked = ffi::String;
* \param call The call expression to be derived.
* \param ctx The builder context.
*/
-using FInferType = ffi::TypedFunction<Type(const Call& call, const BlockBuilder& ctx)>;
+using FInferTypeWithBuilder = ffi::TypedFunction<Type(const Call& call, const BlockBuilder& ctx)>;
/*!
* \brief The function type of a normalization function.
diff --git a/python/tvm/ir/__init__.py b/python/tvm/ir/__init__.py
index de84cecf66..fe264b16ba 100644
--- a/python/tvm/ir/__init__.py
+++ b/python/tvm/ir/__init__.py
@@ -53,6 +53,7 @@ from .expr import (
Var,
is_prim_expr,
is_prim_var,
+ reinfer_type,
)
from . import prim
from .function import BaseFunc, CallingConv
diff --git a/python/tvm/ir/expr.py b/python/tvm/ir/expr.py
index 9cfd4722d3..6129ea8c50 100644
--- a/python/tvm/ir/expr.py
+++ b/python/tvm/ir/expr.py
@@ -593,6 +593,15 @@ class Call(_CallableExprWithOp):
)
+def reinfer_type(call: Call) -> "tvm.ir.Type":
+ """Derive a Call's result type from its current inputs without changing the Call.
+
+ The operator must have a context-free inference rule. Invalid inputs and
+ missing rules raise errors.
+ """
+ return _ffi_api.reinfer_type(call)
+
+
@tvm_ffi.register_object("ir.TensorRegion")
class TensorRegion(Expr, Scriptable):
"""A region of an arbitrary tensor expression.
diff --git a/python/tvm/relax/transform/legalize_ops/linear_algebra.py b/python/tvm/relax/transform/legalize_ops/linear_algebra.py
index fc000cf40e..7e0846e05e 100644
--- a/python/tvm/relax/transform/legalize_ops/linear_algebra.py
+++ b/python/tvm/relax/transform/legalize_ops/linear_algebra.py
@@ -52,7 +52,7 @@ def _matmul(bb: BlockBuilder, call: Call) -> Expr:
a_relax = relax.Var("a", relax.TensorType(a.shape))
b_relax = relax.Var("b", relax.TensorType(b.shape))
- f_infer_ty = call.op.get_attr("FInferType")
+ f_infer_ty = call.op.get_attr("relax.FInferTypeWithBuilder")
output_shape = f_infer_ty(relax.op.matmul(a_relax, b_relax), bb).shape
if isinstance(a_shape[-1], tirx.IntImm) and a_shape[-1] == 0:
return te.compute(
diff --git a/python/tvm/relax/transform/legalize_ops/manipulate.py b/python/tvm/relax/transform/legalize_ops/manipulate.py
index 4301641b05..ecc904b175 100644
--- a/python/tvm/relax/transform/legalize_ops/manipulate.py
+++ b/python/tvm/relax/transform/legalize_ops/manipulate.py
@@ -78,8 +78,8 @@ def _concat(bb: BlockBuilder, call: Call) -> Expr:
def _expand_dims(bb: BlockBuilder, call: Call) -> Expr:
def te_expand_dims(data, axis):
data_relax = relax.Var("data", relax.TensorType(data.shape))
- f_infer_ty = call.op.get_attr("FInferType")
- output_shape = f_infer_ty(relax.op.expand_dims(data_relax, axis), bb).shape
+ output_ty = tvm.ir.reinfer_type(relax.op.expand_dims(data_relax, axis))
+ output_shape = output_ty.shape
output_ndim = len(output_shape)
data_dims = []
diff --git a/src/ir/expr.cc b/src/ir/expr.cc
index 02ba0b5768..1bfb833d88 100644
--- a/src/ir/expr.cc
+++ b/src/ir/expr.cc
@@ -42,22 +42,25 @@ namespace tvm {
namespace {
template <typename TNode>
-TVMFFIAny ConstantVisit(ffi::StructuralVisitorObj* visitor, ffi::AnyView value) noexcept {
+TVM_FFI_INLINE ffi::Expected<ffi::Optional<ffi::VisitInterrupt>> ConstantVisit(
+ ffi::StructuralVisitorObj* visitor, ffi::AnyView value) noexcept {
const TNode* self =
ffi::details::AnyUnsafe::RawObjectPtrFromAnyViewAfterCheck<const TNode>(value);
TVM_FFI_S_VISIT_MAYBE_EARLY_RETURN(visitor->VisitExpected(self->ty));
- return ffi::AnyView(nullptr).CopyToTVMFFIAny();
+ return std::nullopt;
}
-TVMFFIAny DataTypeImmVisit(ffi::StructuralVisitorObj* visitor, ffi::AnyView value) noexcept {
+TVM_FFI_INLINE ffi::Expected<ffi::Optional<ffi::VisitInterrupt>> DataTypeImmVisit(
+ ffi::StructuralVisitorObj* visitor, ffi::AnyView value) noexcept {
const DataTypeImmNode* self =
ffi::details::AnyUnsafe::RawObjectPtrFromAnyViewAfterCheck<const DataTypeImmNode>(value);
TVM_FFI_S_VISIT_MAYBE_EARLY_RETURN(visitor->VisitExpected(self->ty));
TVM_FFI_S_VISIT_MAYBE_EARLY_RETURN(visitor->VisitExpected(self->value));
- return ffi::AnyView(nullptr).CopyToTVMFFIAny();
+ return std::nullopt;
}
-TVMFFIAny DataTypeImmMutate(ffi::StructuralMutatorObj* mutator, ffi::AnyView value) noexcept {
+TVM_FFI_INLINE ffi::Expected<ffi::UnchangedOr<ffi::Any>> DataTypeImmMutate(
+ ffi::StructuralMutatorObj* mutator, ffi::AnyView value) noexcept {
const DataTypeImmNode* self =
ffi::details::AnyUnsafe::RawObjectPtrFromAnyViewAfterCheck<const DataTypeImmNode>(value);
TVM_FFI_S_MUTATE_ASSIGN_OR_RETURN(ffi::UnchangedOr<Type>, mapped_ty,
@@ -65,16 +68,16 @@ TVMFFIAny DataTypeImmMutate(ffi::StructuralMutatorObj* mutator, ffi::AnyView val
TVM_FFI_S_MUTATE_ASSIGN_OR_RETURN(ffi::UnchangedOr<DLDataType>, mapped_value,
mutator->MutateExpected(self->value));
if (mapped_ty.UnchangedOrSameAs(self->ty) && mapped_value.UnchangedOrSameAs(self->value)) {
- return ffi::Unchanged().CopyToTVMFFIAny();
+ return ffi::Unchanged();
}
ffi::ObjectPtr<DataTypeImmNode> copy = ffi::make_object<DataTypeImmNode>(*self);
copy->ty = std::move(mapped_ty).ValueOrUnchanged(std::move(copy->ty));
copy->value = std::move(mapped_value).ValueOrUnchanged(copy->value);
- return ffi::details::AnyUnsafe::MoveAnyToTVMFFIAny(ffi::Any(std::move(copy)));
+ return ffi::Any(std::move(copy));
}
-TVMFFIAny DataTypeImmMaybeInplaceMutate(ffi::StructuralMutatorObj* mutator,
- ffi::AnyView value) noexcept {
+TVM_FFI_INLINE ffi::Expected<ffi::UnchangedOr<ffi::Any>> DataTypeImmMaybeInplaceMutate(
+ ffi::StructuralMutatorObj* mutator, ffi::AnyView value) noexcept {
DataTypeImmNode* self = const_cast<DataTypeImmNode*>(
ffi::details::AnyUnsafe::RawObjectPtrFromAnyViewAfterCheck<const DataTypeImmNode>(value));
TVM_FFI_S_MUTATE_ASSIGN_OR_RETURN(ffi::UnchangedOr<Type>, mapped_ty,
@@ -83,24 +86,25 @@ TVMFFIAny DataTypeImmMaybeInplaceMutate(ffi::StructuralMutatorObj* mutator,
mutator->MutateExpected(self->value, ffi::InplaceMode::kAllow));
if (!mapped_ty.IsUnchanged()) self->ty = std::move(mapped_ty).ValueUnchecked();
if (!mapped_value.IsUnchanged()) self->value = std::move(mapped_value).ValueUnchecked();
- return ffi::Unchanged().CopyToTVMFFIAny();
+ return ffi::Unchanged();
}
template <typename TNode>
-TVMFFIAny ConstantMutate(ffi::StructuralMutatorObj* mutator, ffi::AnyView value) noexcept {
+TVM_FFI_INLINE ffi::Expected<ffi::UnchangedOr<ffi::Any>> ConstantMutate(
+ ffi::StructuralMutatorObj* mutator, ffi::AnyView value) noexcept {
const TNode* self =
ffi::details::AnyUnsafe::RawObjectPtrFromAnyViewAfterCheck<const TNode>(value);
TVM_FFI_S_MUTATE_ASSIGN_OR_RETURN(ffi::UnchangedOr<Type>, mapped_ty,
mutator->MutateExpected(self->ty));
- if (mapped_ty.UnchangedOrSameAs(self->ty)) return ffi::Unchanged().CopyToTVMFFIAny();
+ if (mapped_ty.UnchangedOrSameAs(self->ty)) return ffi::Unchanged();
ffi::ObjectPtr<TNode> copy = ffi::make_object<TNode>(*self);
copy->ty = std::move(mapped_ty).ValueOrUnchanged(std::move(copy->ty));
- return ffi::details::AnyUnsafe::MoveAnyToTVMFFIAny(ffi::Any(std::move(copy)));
+ return ffi::Any(std::move(copy));
}
template <typename TNode>
-TVMFFIAny ConstantMaybeInplaceMutate(ffi::StructuralMutatorObj* mutator,
- ffi::AnyView value) noexcept {
+TVM_FFI_INLINE ffi::Expected<ffi::UnchangedOr<ffi::Any>> ConstantMaybeInplaceMutate(
+ ffi::StructuralMutatorObj* mutator, ffi::AnyView value) noexcept {
TNode* self = const_cast<TNode*>(
ffi::details::AnyUnsafe::RawObjectPtrFromAnyViewAfterCheck<const TNode>(value));
TVM_FFI_S_MUTATE_ASSIGN_OR_RETURN(ffi::UnchangedOr<Type>, mapped_ty,
@@ -108,53 +112,57 @@ TVMFFIAny ConstantMaybeInplaceMutate(ffi::StructuralMutatorObj* mutator,
if (!mapped_ty.UnchangedOrSameAs(self->ty)) {
self->ty = std::move(mapped_ty).ValueUnchecked();
}
- return ffi::Unchanged().CopyToTVMFFIAny();
+ return ffi::Unchanged();
}
-TVMFFIAny OpaqueExprVisit(ffi::StructuralVisitorObj* visitor, ffi::AnyView value) noexcept {
+TVM_FFI_INLINE ffi::Expected<ffi::Optional<ffi::VisitInterrupt>> OpaqueExprVisit(
+ ffi::StructuralVisitorObj* visitor, ffi::AnyView value) noexcept {
const OpaqueExprNode* self =
ffi::details::AnyUnsafe::RawObjectPtrFromAnyViewAfterCheck<const OpaqueExprNode>(value);
TVM_FFI_S_VISIT_MAYBE_EARLY_RETURN(visitor->VisitExpected(self->ty));
// Any None is Expected<Optional<VisitInterrupt>>'s successful empty value.
- return ffi::AnyView(nullptr).CopyToTVMFFIAny();
+ return std::nullopt;
}
-TVMFFIAny OpaqueExprMutate(ffi::StructuralMutatorObj* mutator, ffi::AnyView value) noexcept {
+TVM_FFI_INLINE ffi::Expected<ffi::UnchangedOr<ffi::Any>> OpaqueExprMutate(
+ ffi::StructuralMutatorObj* mutator, ffi::AnyView value) noexcept {
const OpaqueExprNode* self =
ffi::details::AnyUnsafe::RawObjectPtrFromAnyViewAfterCheck<const OpaqueExprNode>(value);
TVM_FFI_S_MUTATE_ASSIGN_OR_RETURN(ffi::UnchangedOr<Type>, mapped_ty_u,
mutator->MutateExpected(self->ty));
if (mapped_ty_u.UnchangedOrSameAs(self->ty)) {
- return ffi::Unchanged().CopyToTVMFFIAny();
+ return ffi::Unchanged();
}
ffi::ObjectPtr<OpaqueExprNode> copy = ffi::make_object<OpaqueExprNode>(*self);
if (!mapped_ty_u.IsUnchanged()) copy->ty = std::move(mapped_ty_u).ValueUnchecked();
- return ffi::details::AnyUnsafe::MoveAnyToTVMFFIAny(ffi::Any(std::move(copy)));
+ return ffi::Any(std::move(copy));
}
-TVMFFIAny OpaqueExprMaybeInplaceMutate(ffi::StructuralMutatorObj* mutator,
- ffi::AnyView value) noexcept {
+TVM_FFI_INLINE ffi::Expected<ffi::UnchangedOr<ffi::Any>> OpaqueExprMaybeInplaceMutate(
+ ffi::StructuralMutatorObj* mutator, ffi::AnyView value) noexcept {
OpaqueExprNode* self = const_cast<OpaqueExprNode*>(
ffi::details::AnyUnsafe::RawObjectPtrFromAnyViewAfterCheck<const OpaqueExprNode>(value));
TVM_FFI_S_MUTATE_ASSIGN_OR_RETURN(ffi::UnchangedOr<Type>, mapped_ty_u,
mutator->MutateExpected(self->ty, ffi::InplaceMode::kAllow));
if (mapped_ty_u.UnchangedOrSameAs(self->ty)) {
- return ffi::Unchanged().CopyToTVMFFIAny();
+ return ffi::Unchanged();
}
if (!mapped_ty_u.IsUnchanged()) self->ty = std::move(mapped_ty_u).ValueUnchecked();
- return ffi::Unchanged().CopyToTVMFFIAny();
+ return ffi::Unchanged();
}
-TVMFFIAny TensorLoadVisit(ffi::StructuralVisitorObj* visitor, ffi::AnyView value) noexcept {
+TVM_FFI_INLINE ffi::Expected<ffi::Optional<ffi::VisitInterrupt>> TensorLoadVisit(
+ ffi::StructuralVisitorObj* visitor, ffi::AnyView value) noexcept {
const TensorLoadNode* self =
ffi::details::AnyUnsafe::RawObjectPtrFromAnyViewAfterCheck<const TensorLoadNode>(value);
TVM_FFI_S_VISIT_MAYBE_EARLY_RETURN(visitor->VisitExpected(self->ty));
TVM_FFI_S_VISIT_MAYBE_EARLY_RETURN(visitor->VisitExpected(self->source));
TVM_FFI_S_VISIT_MAYBE_EARLY_RETURN(visitor->VisitExpected(self->indices));
- return ffi::AnyView(nullptr).CopyToTVMFFIAny();
+ return std::nullopt;
}
-TVMFFIAny TensorLoadMutate(ffi::StructuralMutatorObj* mutator, ffi::AnyView value) noexcept {
+TVM_FFI_INLINE ffi::Expected<ffi::UnchangedOr<ffi::Any>> TensorLoadMutate(
+ ffi::StructuralMutatorObj* mutator, ffi::AnyView value) noexcept {
const TensorLoadNode* self =
ffi::details::AnyUnsafe::RawObjectPtrFromAnyViewAfterCheck<const TensorLoadNode>(value);
TVM_FFI_S_MUTATE_ASSIGN_OR_RETURN(ffi::UnchangedOr<Type>, mapped_ty_u,
@@ -165,17 +173,17 @@ TVMFFIAny TensorLoadMutate(ffi::StructuralMutatorObj* mutator, ffi::AnyView valu
mutator->MutateExpected(self->indices));
if (mapped_ty_u.UnchangedOrSameAs(self->ty) && mapped_source_u.UnchangedOrSameAs(self->source) &&
mapped_indices_u.UnchangedOrSameAs(self->indices)) {
- return ffi::Unchanged().CopyToTVMFFIAny();
+ return ffi::Unchanged();
}
ffi::ObjectPtr<TensorLoadNode> copy = ffi::make_object<TensorLoadNode>(*self);
if (!mapped_ty_u.IsUnchanged()) copy->ty = std::move(mapped_ty_u).ValueUnchecked();
if (!mapped_source_u.IsUnchanged()) copy->source = std::move(mapped_source_u).ValueUnchecked();
if (!mapped_indices_u.IsUnchanged()) copy->indices = std::move(mapped_indices_u).ValueUnchecked();
- return ffi::details::AnyUnsafe::MoveAnyToTVMFFIAny(ffi::Any(std::move(copy)));
+ return ffi::Any(std::move(copy));
}
-TVMFFIAny TensorLoadMaybeInplaceMutate(ffi::StructuralMutatorObj* mutator,
- ffi::AnyView value) noexcept {
+TVM_FFI_INLINE ffi::Expected<ffi::UnchangedOr<ffi::Any>> TensorLoadMaybeInplaceMutate(
+ ffi::StructuralMutatorObj* mutator, ffi::AnyView value) noexcept {
TensorLoadNode* self = const_cast<TensorLoadNode*>(
ffi::details::AnyUnsafe::RawObjectPtrFromAnyViewAfterCheck<const TensorLoadNode>(value));
TVM_FFI_S_MUTATE_ASSIGN_OR_RETURN(ffi::UnchangedOr<Type>, mapped_ty_u,
@@ -188,24 +196,26 @@ TVMFFIAny TensorLoadMaybeInplaceMutate(ffi::StructuralMutatorObj* mutator,
mutator->MutateExpected(self->indices, ffi::InplaceMode::kAllow));
if (mapped_ty_u.UnchangedOrSameAs(self->ty) && mapped_source_u.UnchangedOrSameAs(self->source) &&
mapped_indices_u.UnchangedOrSameAs(self->indices)) {
- return ffi::Unchanged().CopyToTVMFFIAny();
+ return ffi::Unchanged();
}
if (!mapped_ty_u.IsUnchanged()) self->ty = std::move(mapped_ty_u).ValueUnchecked();
if (!mapped_source_u.IsUnchanged()) self->source = std::move(mapped_source_u).ValueUnchecked();
if (!mapped_indices_u.IsUnchanged()) self->indices = std::move(mapped_indices_u).ValueUnchecked();
- return ffi::Unchanged().CopyToTVMFFIAny();
+ return ffi::Unchanged();
}
-TVMFFIAny TensorRegionVisit(ffi::StructuralVisitorObj* visitor, ffi::AnyView value) noexcept {
+TVM_FFI_INLINE ffi::Expected<ffi::Optional<ffi::VisitInterrupt>> TensorRegionVisit(
+ ffi::StructuralVisitorObj* visitor, ffi::AnyView value) noexcept {
const TensorRegionNode* self =
ffi::details::AnyUnsafe::RawObjectPtrFromAnyViewAfterCheck<const TensorRegionNode>(value);
TVM_FFI_S_VISIT_MAYBE_EARLY_RETURN(visitor->VisitExpected(self->ty));
TVM_FFI_S_VISIT_MAYBE_EARLY_RETURN(visitor->VisitExpected(self->source));
TVM_FFI_S_VISIT_MAYBE_EARLY_RETURN(visitor->VisitExpected(self->region));
- return ffi::AnyView(nullptr).CopyToTVMFFIAny();
+ return std::nullopt;
}
-TVMFFIAny TensorRegionMutate(ffi::StructuralMutatorObj* mutator, ffi::AnyView value) noexcept {
+TVM_FFI_INLINE ffi::Expected<ffi::UnchangedOr<ffi::Any>> TensorRegionMutate(
+ ffi::StructuralMutatorObj* mutator, ffi::AnyView value) noexcept {
const TensorRegionNode* self =
ffi::details::AnyUnsafe::RawObjectPtrFromAnyViewAfterCheck<const TensorRegionNode>(value);
TVM_FFI_S_MUTATE_ASSIGN_OR_RETURN(ffi::UnchangedOr<Type>, mapped_ty_u,
@@ -216,17 +226,17 @@ TVMFFIAny TensorRegionMutate(ffi::StructuralMutatorObj* mutator, ffi::AnyView va
mutator->MutateExpected(self->region));
if (mapped_ty_u.UnchangedOrSameAs(self->ty) && mapped_source_u.UnchangedOrSameAs(self->source) &&
mapped_region_u.UnchangedOrSameAs(self->region)) {
- return ffi::Unchanged().CopyToTVMFFIAny();
+ return ffi::Unchanged();
}
ffi::ObjectPtr<TensorRegionNode> copy = ffi::make_object<TensorRegionNode>(*self);
if (!mapped_ty_u.IsUnchanged()) copy->ty = std::move(mapped_ty_u).ValueUnchecked();
if (!mapped_source_u.IsUnchanged()) copy->source = std::move(mapped_source_u).ValueUnchecked();
if (!mapped_region_u.IsUnchanged()) copy->region = std::move(mapped_region_u).ValueUnchecked();
- return ffi::details::AnyUnsafe::MoveAnyToTVMFFIAny(ffi::Any(std::move(copy)));
+ return ffi::Any(std::move(copy));
}
-TVMFFIAny TensorRegionMaybeInplaceMutate(ffi::StructuralMutatorObj* mutator,
- ffi::AnyView value) noexcept {
+TVM_FFI_INLINE ffi::Expected<ffi::UnchangedOr<ffi::Any>> TensorRegionMaybeInplaceMutate(
+ ffi::StructuralMutatorObj* mutator, ffi::AnyView value) noexcept {
TensorRegionNode* self = const_cast<TensorRegionNode*>(
ffi::details::AnyUnsafe::RawObjectPtrFromAnyViewAfterCheck<const TensorRegionNode>(value));
TVM_FFI_S_MUTATE_ASSIGN_OR_RETURN(ffi::UnchangedOr<Type>, mapped_ty_u,
@@ -239,23 +249,25 @@ TVMFFIAny TensorRegionMaybeInplaceMutate(ffi::StructuralMutatorObj* mutator,
mutator->MutateExpected(self->region, ffi::InplaceMode::kAllow));
if (mapped_ty_u.UnchangedOrSameAs(self->ty) && mapped_source_u.UnchangedOrSameAs(self->source) &&
mapped_region_u.UnchangedOrSameAs(self->region)) {
- return ffi::Unchanged().CopyToTVMFFIAny();
+ return ffi::Unchanged();
}
if (!mapped_ty_u.IsUnchanged()) self->ty = std::move(mapped_ty_u).ValueUnchecked();
if (!mapped_source_u.IsUnchanged()) self->source = std::move(mapped_source_u).ValueUnchecked();
if (!mapped_region_u.IsUnchanged()) self->region = std::move(mapped_region_u).ValueUnchecked();
- return ffi::Unchanged().CopyToTVMFFIAny();
+ return ffi::Unchanged();
}
-TVMFFIAny TupleVisit(ffi::StructuralVisitorObj* visitor, ffi::AnyView value) noexcept {
+TVM_FFI_INLINE ffi::Expected<ffi::Optional<ffi::VisitInterrupt>> TupleVisit(
+ ffi::StructuralVisitorObj* visitor, ffi::AnyView value) noexcept {
const TupleNode* self =
ffi::details::AnyUnsafe::RawObjectPtrFromAnyViewAfterCheck<const TupleNode>(value);
TVM_FFI_S_VISIT_MAYBE_EARLY_RETURN(visitor->VisitExpected(self->ty));
TVM_FFI_S_VISIT_MAYBE_EARLY_RETURN(visitor->VisitExpected(self->fields));
- return ffi::AnyView(nullptr).CopyToTVMFFIAny();
+ return std::nullopt;
}
-TVMFFIAny TupleMutate(ffi::StructuralMutatorObj* mutator, ffi::AnyView value) noexcept {
+TVM_FFI_INLINE ffi::Expected<ffi::UnchangedOr<ffi::Any>> TupleMutate(
+ ffi::StructuralMutatorObj* mutator, ffi::AnyView value) noexcept {
const TupleNode* self =
ffi::details::AnyUnsafe::RawObjectPtrFromAnyViewAfterCheck<const TupleNode>(value);
TVM_FFI_S_MUTATE_ASSIGN_OR_RETURN(ffi::UnchangedOr<Type>, mapped_ty_u,
@@ -263,15 +275,16 @@ TVMFFIAny TupleMutate(ffi::StructuralMutatorObj* mutator, ffi::AnyView value) no
TVM_FFI_S_MUTATE_ASSIGN_OR_RETURN(ffi::UnchangedOr<ffi::Array<Expr>>, mapped_fields_u,
mutator->MutateExpected(self->fields));
if (mapped_ty_u.UnchangedOrSameAs(self->ty) && mapped_fields_u.UnchangedOrSameAs(self->fields)) {
- return ffi::Unchanged().CopyToTVMFFIAny();
+ return ffi::Unchanged();
}
ffi::ObjectPtr<TupleNode> copy = ffi::make_object<TupleNode>(*self);
if (!mapped_ty_u.IsUnchanged()) copy->ty = std::move(mapped_ty_u).ValueUnchecked();
if (!mapped_fields_u.IsUnchanged()) copy->fields = std::move(mapped_fields_u).ValueUnchecked();
- return ffi::details::AnyUnsafe::MoveAnyToTVMFFIAny(ffi::Any(std::move(copy)));
+ return ffi::Any(std::move(copy));
}
-TVMFFIAny TupleMaybeInplaceMutate(ffi::StructuralMutatorObj* mutator, ffi::AnyView value) noexcept {
+TVM_FFI_INLINE ffi::Expected<ffi::UnchangedOr<ffi::Any>> TupleMaybeInplaceMutate(
+ ffi::StructuralMutatorObj* mutator, ffi::AnyView value) noexcept {
TupleNode* self = const_cast<TupleNode*>(
ffi::details::AnyUnsafe::RawObjectPtrFromAnyViewAfterCheck<const TupleNode>(value));
TVM_FFI_S_MUTATE_ASSIGN_OR_RETURN(ffi::UnchangedOr<Type>, mapped_ty_u,
@@ -280,23 +293,25 @@ TVMFFIAny TupleMaybeInplaceMutate(ffi::StructuralMutatorObj* mutator, ffi::AnyVi
ffi::UnchangedOr<ffi::Array<Expr>>, mapped_fields_u,
mutator->MutateExpected(self->fields, ffi::InplaceMode::kAllow));
if (mapped_ty_u.UnchangedOrSameAs(self->ty) && mapped_fields_u.UnchangedOrSameAs(self->fields)) {
- return ffi::Unchanged().CopyToTVMFFIAny();
+ return ffi::Unchanged();
}
if (!mapped_ty_u.IsUnchanged()) self->ty = std::move(mapped_ty_u).ValueUnchecked();
if (!mapped_fields_u.IsUnchanged()) self->fields = std::move(mapped_fields_u).ValueUnchecked();
- return ffi::Unchanged().CopyToTVMFFIAny();
+ return ffi::Unchanged();
}
-TVMFFIAny TupleGetItemVisit(ffi::StructuralVisitorObj* visitor, ffi::AnyView value) noexcept {
+TVM_FFI_INLINE ffi::Expected<ffi::Optional<ffi::VisitInterrupt>> TupleGetItemVisit(
+ ffi::StructuralVisitorObj* visitor, ffi::AnyView value) noexcept {
// skips: index
const TupleGetItemNode* self =
ffi::details::AnyUnsafe::RawObjectPtrFromAnyViewAfterCheck<const TupleGetItemNode>(value);
TVM_FFI_S_VISIT_MAYBE_EARLY_RETURN(visitor->VisitExpected(self->ty));
TVM_FFI_S_VISIT_MAYBE_EARLY_RETURN(visitor->VisitExpected(self->tuple));
- return ffi::AnyView(nullptr).CopyToTVMFFIAny();
+ return std::nullopt;
}
-TVMFFIAny TupleGetItemMutate(ffi::StructuralMutatorObj* mutator, ffi::AnyView value) noexcept {
+TVM_FFI_INLINE ffi::Expected<ffi::UnchangedOr<ffi::Any>> TupleGetItemMutate(
+ ffi::StructuralMutatorObj* mutator, ffi::AnyView value) noexcept {
// skips: index
const TupleGetItemNode* self =
ffi::details::AnyUnsafe::RawObjectPtrFromAnyViewAfterCheck<const TupleGetItemNode>(value);
@@ -305,16 +320,16 @@ TVMFFIAny TupleGetItemMutate(ffi::StructuralMutatorObj* mutator, ffi::AnyView va
TVM_FFI_S_MUTATE_ASSIGN_OR_RETURN(ffi::UnchangedOr<Expr>, mapped_tuple_u,
mutator->MutateExpected(self->tuple));
if (mapped_ty_u.UnchangedOrSameAs(self->ty) && mapped_tuple_u.UnchangedOrSameAs(self->tuple)) {
- return ffi::Unchanged().CopyToTVMFFIAny();
+ return ffi::Unchanged();
}
ffi::ObjectPtr<TupleGetItemNode> copy = ffi::make_object<TupleGetItemNode>(*self);
if (!mapped_ty_u.IsUnchanged()) copy->ty = std::move(mapped_ty_u).ValueUnchecked();
if (!mapped_tuple_u.IsUnchanged()) copy->tuple = std::move(mapped_tuple_u).ValueUnchecked();
- return ffi::details::AnyUnsafe::MoveAnyToTVMFFIAny(ffi::Any(std::move(copy)));
+ return ffi::Any(std::move(copy));
}
-TVMFFIAny TupleGetItemMaybeInplaceMutate(ffi::StructuralMutatorObj* mutator,
- ffi::AnyView value) noexcept {
+TVM_FFI_INLINE ffi::Expected<ffi::UnchangedOr<ffi::Any>> TupleGetItemMaybeInplaceMutate(
+ ffi::StructuralMutatorObj* mutator, ffi::AnyView value) noexcept {
// skips: index
TupleGetItemNode* self = const_cast<TupleGetItemNode*>(
ffi::details::AnyUnsafe::RawObjectPtrFromAnyViewAfterCheck<const TupleGetItemNode>(value));
@@ -323,52 +338,60 @@ TVMFFIAny TupleGetItemMaybeInplaceMutate(ffi::StructuralMutatorObj* mutator,
TVM_FFI_S_MUTATE_ASSIGN_OR_RETURN(ffi::UnchangedOr<Expr>, mapped_tuple_u,
mutator->MutateExpected(self->tuple, ffi::InplaceMode::kAllow));
if (mapped_ty_u.UnchangedOrSameAs(self->ty) && mapped_tuple_u.UnchangedOrSameAs(self->tuple)) {
- return ffi::Unchanged().CopyToTVMFFIAny();
+ return ffi::Unchanged();
}
if (!mapped_ty_u.IsUnchanged()) self->ty = std::move(mapped_ty_u).ValueUnchecked();
if (!mapped_tuple_u.IsUnchanged()) self->tuple = std::move(mapped_tuple_u).ValueUnchecked();
- return ffi::Unchanged().CopyToTVMFFIAny();
+ return ffi::Unchanged();
}
-TVMFFIAny IntImmVisit(ffi::StructuralVisitorObj*, ffi::AnyView) noexcept {
+TVM_FFI_INLINE ffi::Expected<ffi::Optional<ffi::VisitInterrupt>> IntImmVisit(
+ ffi::StructuralVisitorObj*, ffi::AnyView) noexcept {
// skips: value
- return ffi::AnyView(nullptr).CopyToTVMFFIAny();
+ return std::nullopt;
}
-TVMFFIAny IntImmMutate(ffi::StructuralMutatorObj*, ffi::AnyView) noexcept {
+TVM_FFI_INLINE ffi::Expected<ffi::UnchangedOr<ffi::Any>> IntImmMutate(ffi::StructuralMutatorObj*,
+ ffi::AnyView) noexcept {
// skips: value
- return ffi::Unchanged().CopyToTVMFFIAny();
+ return ffi::Unchanged();
}
-TVMFFIAny IntImmMaybeInplaceMutate(ffi::StructuralMutatorObj*, ffi::AnyView) noexcept {
+TVM_FFI_INLINE ffi::Expected<ffi::UnchangedOr<ffi::Any>> IntImmMaybeInplaceMutate(
+ ffi::StructuralMutatorObj*, ffi::AnyView) noexcept {
// skips: value
- return ffi::Unchanged().CopyToTVMFFIAny();
+ return ffi::Unchanged();
}
-TVMFFIAny FloatImmVisit(ffi::StructuralVisitorObj*, ffi::AnyView) noexcept {
+TVM_FFI_INLINE ffi::Expected<ffi::Optional<ffi::VisitInterrupt>> FloatImmVisit(
+ ffi::StructuralVisitorObj*, ffi::AnyView) noexcept {
// skips: value
- return ffi::AnyView(nullptr).CopyToTVMFFIAny();
+ return std::nullopt;
}
-TVMFFIAny FloatImmMutate(ffi::StructuralMutatorObj*, ffi::AnyView) noexcept {
+TVM_FFI_INLINE ffi::Expected<ffi::UnchangedOr<ffi::Any>> FloatImmMutate(ffi::StructuralMutatorObj*,
+ ffi::AnyView) noexcept {
// skips: value
- return ffi::Unchanged().CopyToTVMFFIAny();
+ return ffi::Unchanged();
}
-TVMFFIAny FloatImmMaybeInplaceMutate(ffi::StructuralMutatorObj*, ffi::AnyView) noexcept {
+TVM_FFI_INLINE ffi::Expected<ffi::UnchangedOr<ffi::Any>> FloatImmMaybeInplaceMutate(
+ ffi::StructuralMutatorObj*, ffi::AnyView) noexcept {
// skips: value
- return ffi::Unchanged().CopyToTVMFFIAny();
+ return ffi::Unchanged();
}
-TVMFFIAny RangeVisit(ffi::StructuralVisitorObj* visitor, ffi::AnyView value) noexcept {
+TVM_FFI_INLINE ffi::Expected<ffi::Optional<ffi::VisitInterrupt>> RangeVisit(
+ ffi::StructuralVisitorObj* visitor, ffi::AnyView value) noexcept {
const RangeNode* self =
ffi::details::AnyUnsafe::RawObjectPtrFromAnyViewAfterCheck<const RangeNode>(value);
TVM_FFI_S_VISIT_MAYBE_EARLY_RETURN(visitor->VisitExpected(self->min));
TVM_FFI_S_VISIT_MAYBE_EARLY_RETURN(visitor->VisitExpected(self->extent));
- return ffi::AnyView(nullptr).CopyToTVMFFIAny();
+ return std::nullopt;
}
-TVMFFIAny RangeMutate(ffi::StructuralMutatorObj* mutator, ffi::AnyView value) noexcept {
+TVM_FFI_INLINE ffi::Expected<ffi::UnchangedOr<ffi::Any>> RangeMutate(
+ ffi::StructuralMutatorObj* mutator, ffi::AnyView value) noexcept {
const RangeNode* self =
ffi::details::AnyUnsafe::RawObjectPtrFromAnyViewAfterCheck<const RangeNode>(value);
TVM_FFI_S_MUTATE_ASSIGN_OR_RETURN(ffi::UnchangedOr<PrimExpr>, mapped_min_u,
@@ -377,15 +400,16 @@ TVMFFIAny RangeMutate(ffi::StructuralMutatorObj* mutator, ffi::AnyView value) no
mutator->MutateExpected(self->extent));
if (mapped_min_u.UnchangedOrSameAs(self->min) &&
mapped_extent_u.UnchangedOrSameAs(self->extent)) {
- return ffi::Unchanged().CopyToTVMFFIAny();
+ return ffi::Unchanged();
}
ffi::ObjectPtr<RangeNode> copy = ffi::make_object<RangeNode>(*self);
if (!mapped_min_u.IsUnchanged()) copy->min = std::move(mapped_min_u).ValueUnchecked();
if (!mapped_extent_u.IsUnchanged()) copy->extent = std::move(mapped_extent_u).ValueUnchecked();
- return ffi::details::AnyUnsafe::MoveAnyToTVMFFIAny(ffi::Any(std::move(copy)));
+ return ffi::Any(std::move(copy));
}
-TVMFFIAny RangeMaybeInplaceMutate(ffi::StructuralMutatorObj* mutator, ffi::AnyView value) noexcept {
+TVM_FFI_INLINE ffi::Expected<ffi::UnchangedOr<ffi::Any>> RangeMaybeInplaceMutate(
+ ffi::StructuralMutatorObj* mutator, ffi::AnyView value) noexcept {
RangeNode* self = const_cast<RangeNode*>(
ffi::details::AnyUnsafe::RawObjectPtrFromAnyViewAfterCheck<const RangeNode>(value));
TVM_FFI_S_MUTATE_ASSIGN_OR_RETURN(ffi::UnchangedOr<PrimExpr>, mapped_min_u,
@@ -395,16 +419,17 @@ TVMFFIAny RangeMaybeInplaceMutate(ffi::StructuralMutatorObj* mutator, ffi::AnyVi
mutator->MutateExpected(self->extent, ffi::InplaceMode::kAllow));
if (mapped_min_u.UnchangedOrSameAs(self->min) &&
mapped_extent_u.UnchangedOrSameAs(self->extent)) {
- return ffi::Unchanged().CopyToTVMFFIAny();
+ return ffi::Unchanged();
}
if (!mapped_min_u.IsUnchanged()) self->min = std::move(mapped_min_u).ValueUnchecked();
if (!mapped_extent_u.IsUnchanged()) self->extent = std::move(mapped_extent_u).ValueUnchecked();
- return ffi::Unchanged().CopyToTVMFFIAny();
+ return ffi::Unchanged();
}
// DataflowVarNode duplicates this protocol because structural hooks do not inherit. Keep the two
// hook triples in lockstep when changing remap, PrimType-skip, or definition-region behavior.
-TVMFFIAny VarVisit(ffi::StructuralVisitorObj* visitor, ffi::AnyView value) noexcept {
+TVM_FFI_INLINE ffi::Expected<ffi::Optional<ffi::VisitInterrupt>> VarVisit(
+ ffi::StructuralVisitorObj* visitor, ffi::AnyView value) noexcept {
// skips: name
const VarNode* self =
ffi::details::AnyUnsafe::RawObjectPtrFromAnyViewAfterCheck<const VarNode>(value);
@@ -420,10 +445,11 @@ TVMFFIAny VarVisit(ffi::StructuralVisitorObj* visitor, ffi::AnyView value) noexc
TVM_FFI_S_VISIT_MAYBE_EARLY_RETURN(visitor->VisitExpected(self->ty));
}
}
- return ffi::AnyView(nullptr).CopyToTVMFFIAny();
+ return std::nullopt;
}
-TVMFFIAny VarMutate(ffi::StructuralMutatorObj* mutator, ffi::AnyView value) noexcept {
+TVM_FFI_INLINE ffi::Expected<ffi::UnchangedOr<ffi::Any>> VarMutate(
+ ffi::StructuralMutatorObj* mutator, ffi::AnyView value) noexcept {
// skips: name
const VarNode* self =
ffi::details::AnyUnsafe::RawObjectPtrFromAnyViewAfterCheck<const VarNode>(value);
@@ -431,10 +457,10 @@ TVMFFIAny VarMutate(ffi::StructuralMutatorObj* mutator, ffi::AnyView value) noex
TVM_FFI_S_MUTATE_MAYBE_EARLY_RETURN(remap_result);
if (ffi::details::ExpectedUnsafe::GetData(remap_result).type_index() !=
ffi::TypeIndex::kTVMFFINone) {
- return ffi::details::ExpectedUnsafe::MoveToTVMFFIAny(std::move(remap_result));
+ return std::move(remap_result);
}
if (mutator->def_region_kind() == kTVMFFIDefRegionKindNone) {
- return ffi::Unchanged().CopyToTVMFFIAny();
+ return ffi::Unchanged();
}
ffi::UnchangedOr<ffi::Any> result_u = ffi::Unchanged();
ffi::Any mapped_value = ffi::Unchanged();
@@ -459,13 +485,14 @@ TVMFFIAny VarMutate(ffi::StructuralMutatorObj* mutator, ffi::AnyView value) noex
if (!result_u.IsUnchanged() || mutator->def_region_kind() == kTVMFFIDefRegionKindPattern) {
auto set_result = mutator->VarRemapSetExpected(value, mapped_value);
if (TVM_FFI_PREDICT_FALSE(set_result.is_err())) {
- return ffi::details::AnyUnsafe::MoveAnyToTVMFFIAny(ffi::Any(std::move(set_result).error()));
+ return ffi::Unexpected(std::move(set_result).error());
}
}
- return ffi::details::UnchangedOrUnsafe::MoveToTVMFFIAny(std::move(result_u));
+ return std::move(result_u);
}
-TVMFFIAny VarMaybeInplaceMutate(ffi::StructuralMutatorObj* mutator, ffi::AnyView value) noexcept {
+TVM_FFI_INLINE ffi::Expected<ffi::UnchangedOr<ffi::Any>> VarMaybeInplaceMutate(
+ ffi::StructuralMutatorObj* mutator, ffi::AnyView value) noexcept {
// skips: name
VarNode* self = const_cast<VarNode*>(
ffi::details::AnyUnsafe::RawObjectPtrFromAnyViewAfterCheck<const VarNode>(value));
@@ -473,10 +500,10 @@ TVMFFIAny VarMaybeInplaceMutate(ffi::StructuralMutatorObj* mutator, ffi::AnyView
TVM_FFI_S_MUTATE_MAYBE_EARLY_RETURN(remap_result);
if (ffi::details::ExpectedUnsafe::GetData(remap_result).type_index() !=
ffi::TypeIndex::kTVMFFINone) {
- return ffi::details::ExpectedUnsafe::MoveToTVMFFIAny(std::move(remap_result));
+ return std::move(remap_result);
}
if (mutator->def_region_kind() == kTVMFFIDefRegionKindNone) {
- return ffi::Unchanged().CopyToTVMFFIAny();
+ return ffi::Unchanged();
}
ffi::UnchangedOr<ffi::Any> result_u = ffi::Unchanged();
ffi::Any mapped_value = ffi::Unchanged();
@@ -501,29 +528,33 @@ TVMFFIAny VarMaybeInplaceMutate(ffi::StructuralMutatorObj* mutator, ffi::AnyView
if (!result_u.IsUnchanged() || mutator->def_region_kind() == kTVMFFIDefRegionKindPattern) {
auto set_result = mutator->VarRemapSetExpected(value, mapped_value);
if (TVM_FFI_PREDICT_FALSE(set_result.is_err())) {
- return ffi::details::AnyUnsafe::MoveAnyToTVMFFIAny(ffi::Any(std::move(set_result).error()));
+ return ffi::Unexpected(std::move(set_result).error());
}
}
- return ffi::details::UnchangedOrUnsafe::MoveToTVMFFIAny(std::move(result_u));
+ return std::move(result_u);
}
-TVMFFIAny GlobalVarVisit(ffi::StructuralVisitorObj*, ffi::AnyView) noexcept {
+TVM_FFI_INLINE ffi::Expected<ffi::Optional<ffi::VisitInterrupt>> GlobalVarVisit(
+ ffi::StructuralVisitorObj*, ffi::AnyView) noexcept {
// GlobalVar is a module-level symbol. name_hint is scalar identity and ty is derived from the
// referenced function, matching GlobalVarNode's custom structural equality/hash definition.
// It has no definition site where this hook could establish a VarRemap. A callback that renames
// GlobalVars is therefore responsible for returning one stable replacement per module symbol.
- return ffi::AnyView(nullptr).CopyToTVMFFIAny();
+ return std::nullopt;
}
-TVMFFIAny GlobalVarMutate(ffi::StructuralMutatorObj*, ffi::AnyView) noexcept {
- return ffi::Unchanged().CopyToTVMFFIAny();
+TVM_FFI_INLINE ffi::Expected<ffi::UnchangedOr<ffi::Any>> GlobalVarMutate(ffi::StructuralMutatorObj*,
+ ffi::AnyView) noexcept {
+ return ffi::Unchanged();
}
-TVMFFIAny GlobalVarMaybeInplaceMutate(ffi::StructuralMutatorObj*, ffi::AnyView) noexcept {
- return ffi::Unchanged().CopyToTVMFFIAny();
+TVM_FFI_INLINE ffi::Expected<ffi::UnchangedOr<ffi::Any>> GlobalVarMaybeInplaceMutate(
+ ffi::StructuralMutatorObj*, ffi::AnyView) noexcept {
+ return ffi::Unchanged();
}
-TVMFFIAny CallVisit(ffi::StructuralVisitorObj* visitor, ffi::AnyView value) noexcept {
+TVM_FFI_INLINE ffi::Expected<ffi::Optional<ffi::VisitInterrupt>> CallVisit(
+ ffi::StructuralVisitorObj* visitor, ffi::AnyView value) noexcept {
// skips: attrs, constant metadata left untouched like the classic Expr functors.
const CallNode* self =
ffi::details::AnyUnsafe::RawObjectPtrFromAnyViewAfterCheck<const CallNode>(value);
@@ -543,10 +574,11 @@ TVMFFIAny CallVisit(ffi::StructuralVisitorObj* visitor, ffi::AnyView value) noex
if (!self->ty_args.empty()) {
TVM_FFI_S_VISIT_MAYBE_EARLY_RETURN(visitor->VisitExpected(self->ty_args));
}
- return ffi::AnyView(nullptr).CopyToTVMFFIAny();
+ return std::nullopt;
}
-TVMFFIAny CallMutate(ffi::StructuralMutatorObj* mutator, ffi::AnyView value) noexcept {
+TVM_FFI_INLINE ffi::Expected<ffi::UnchangedOr<ffi::Any>> CallMutate(
+ ffi::StructuralMutatorObj* mutator, ffi::AnyView value) noexcept {
// skips: attrs, constant metadata left untouched like the classic Expr functors.
const CallNode* self =
ffi::details::AnyUnsafe::RawObjectPtrFromAnyViewAfterCheck<const CallNode>(value);
@@ -579,17 +611,18 @@ TVMFFIAny CallMutate(ffi::StructuralMutatorObj* mutator, ffi::AnyView value) noe
if (mapped_ty_u.UnchangedOrSameAs(self->ty) && mapped_op_u.UnchangedOrSameAs(self->op) &&
mapped_args_u.UnchangedOrSameAs(self->args) &&
mapped_ty_args_u.UnchangedOrSameAs(self->ty_args)) {
- return ffi::Unchanged().CopyToTVMFFIAny();
+ return ffi::Unchanged();
}
ffi::ObjectPtr<CallNode> copy = ffi::make_object<CallNode>(*self);
if (!mapped_ty_u.IsUnchanged()) copy->ty = std::move(mapped_ty_u).ValueUnchecked();
if (!mapped_op_u.IsUnchanged()) copy->op = std::move(mapped_op_u).ValueUnchecked();
if (!mapped_args_u.IsUnchanged()) copy->args = std::move(mapped_args_u).ValueUnchecked();
if (!mapped_ty_args_u.IsUnchanged()) copy->ty_args = std::move(mapped_ty_args_u).ValueUnchecked();
- return ffi::details::AnyUnsafe::MoveAnyToTVMFFIAny(ffi::Any(std::move(copy)));
+ return ffi::Any(std::move(copy));
}
-TVMFFIAny CallMaybeInplaceMutate(ffi::StructuralMutatorObj* mutator, ffi::AnyView value) noexcept {
+TVM_FFI_INLINE ffi::Expected<ffi::UnchangedOr<ffi::Any>> CallMaybeInplaceMutate(
+ ffi::StructuralMutatorObj* mutator, ffi::AnyView value) noexcept {
// skips: attrs, constant metadata left untouched like the classic Expr functors.
CallNode* self = const_cast<CallNode*>(
ffi::details::AnyUnsafe::RawObjectPtrFromAnyViewAfterCheck<const CallNode>(value));
@@ -623,13 +656,13 @@ TVMFFIAny CallMaybeInplaceMutate(ffi::StructuralMutatorObj* mutator, ffi::AnyVie
if (mapped_ty_u.UnchangedOrSameAs(self->ty) && mapped_op_u.UnchangedOrSameAs(self->op) &&
mapped_args_u.UnchangedOrSameAs(self->args) &&
mapped_ty_args_u.UnchangedOrSameAs(self->ty_args)) {
- return ffi::Unchanged().CopyToTVMFFIAny();
+ return ffi::Unchanged();
}
if (!mapped_ty_u.IsUnchanged()) self->ty = std::move(mapped_ty_u).ValueUnchecked();
if (!mapped_op_u.IsUnchanged()) self->op = std::move(mapped_op_u).ValueUnchecked();
if (!mapped_args_u.IsUnchanged()) self->args = std::move(mapped_args_u).ValueUnchecked();
if (!mapped_ty_args_u.IsUnchanged()) self->ty_args = std::move(mapped_ty_args_u).ValueUnchecked();
- return ffi::Unchanged().CopyToTVMFFIAny();
+ return ffi::Unchanged();
}
} // namespace
@@ -642,20 +675,24 @@ TVM_FFI_STATIC_INIT_BLOCK() {
namespace refl = tvm::ffi::reflection;
OpaqueExprNode::RegisterReflection();
refl::TypeAttrDef<OpaqueExprNode>()
- .attr(refl::type_attr::kStructuralVisit, reinterpret_cast<void*>(&OpaqueExprVisit))
- .attr(refl::type_attr::kStructuralMutate, reinterpret_cast<void*>(&OpaqueExprMutate))
+ .attr(refl::type_attr::kStructuralVisit,
+ ffi::FStructuralVisit::FromNative<&OpaqueExprVisit>())
+ .attr(refl::type_attr::kStructuralMutate,
+ ffi::FStructuralMutate::FromNative<&OpaqueExprMutate>())
.attr(refl::type_attr::kStructuralMaybeInplaceMutate,
- reinterpret_cast<void*>(&OpaqueExprMaybeInplaceMutate));
+ ffi::FStructuralMutate::FromNative<&OpaqueExprMaybeInplaceMutate>());
}
TVM_FFI_STATIC_INIT_BLOCK() {
namespace refl = tvm::ffi::reflection;
TensorLoadNode::RegisterReflection();
refl::TypeAttrDef<TensorLoadNode>()
- .attr(refl::type_attr::kStructuralVisit, reinterpret_cast<void*>(&TensorLoadVisit))
- .attr(refl::type_attr::kStructuralMutate, reinterpret_cast<void*>(&TensorLoadMutate))
+ .attr(refl::type_attr::kStructuralVisit,
+ ffi::FStructuralVisit::FromNative<&TensorLoadVisit>())
+ .attr(refl::type_attr::kStructuralMutate,
+ ffi::FStructuralMutate::FromNative<&TensorLoadMutate>())
.attr(refl::type_attr::kStructuralMaybeInplaceMutate,
- reinterpret_cast<void*>(&TensorLoadMaybeInplaceMutate));
+ ffi::FStructuralMutate::FromNative<&TensorLoadMaybeInplaceMutate>());
}
TensorRegion::TensorRegion(Expr source, ffi::Array<Range> region, Type ty, Span span) {
@@ -671,10 +708,12 @@ TVM_FFI_STATIC_INIT_BLOCK() {
namespace refl = tvm::ffi::reflection;
TensorRegionNode::RegisterReflection();
refl::TypeAttrDef<TensorRegionNode>()
- .attr(refl::type_attr::kStructuralVisit, reinterpret_cast<void*>(&TensorRegionVisit))
- .attr(refl::type_attr::kStructuralMutate, reinterpret_cast<void*>(&TensorRegionMutate))
+ .attr(refl::type_attr::kStructuralVisit,
+ ffi::FStructuralVisit::FromNative<&TensorRegionVisit>())
+ .attr(refl::type_attr::kStructuralMutate,
+ ffi::FStructuralMutate::FromNative<&TensorRegionMutate>())
.attr(refl::type_attr::kStructuralMaybeInplaceMutate,
- reinterpret_cast<void*>(&TensorRegionMaybeInplaceMutate));
+ ffi::FStructuralMutate::FromNative<&TensorRegionMaybeInplaceMutate>());
refl::GlobalDef().def("ir.TensorRegion",
[](Expr source, ffi::Array<Range> region, Type ty, Span span) {
return TensorRegion(source, region, ty, span);
@@ -707,10 +746,10 @@ TVM_FFI_STATIC_INIT_BLOCK() {
namespace refl = tvm::ffi::reflection;
TupleNode::RegisterReflection();
refl::TypeAttrDef<TupleNode>()
- .attr(refl::type_attr::kStructuralVisit, reinterpret_cast<void*>(&TupleVisit))
- .attr(refl::type_attr::kStructuralMutate, reinterpret_cast<void*>(&TupleMutate))
+ .attr(refl::type_attr::kStructuralVisit, ffi::FStructuralVisit::FromNative<&TupleVisit>())
+ .attr(refl::type_attr::kStructuralMutate, ffi::FStructuralMutate::FromNative<&TupleMutate>())
.attr(refl::type_attr::kStructuralMaybeInplaceMutate,
- reinterpret_cast<void*>(&TupleMaybeInplaceMutate));
+ ffi::FStructuralMutate::FromNative<&TupleMaybeInplaceMutate>());
refl::GlobalDef().def("ir.Tuple",
[](ffi::Array<Expr> fields, Span span) { return Tuple(fields, span); });
@@ -737,10 +776,12 @@ TVM_FFI_STATIC_INIT_BLOCK() {
namespace refl = tvm::ffi::reflection;
TupleGetItemNode::RegisterReflection();
refl::TypeAttrDef<TupleGetItemNode>()
- .attr(refl::type_attr::kStructuralVisit, reinterpret_cast<void*>(&TupleGetItemVisit))
- .attr(refl::type_attr::kStructuralMutate, reinterpret_cast<void*>(&TupleGetItemMutate))
+ .attr(refl::type_attr::kStructuralVisit,
+ ffi::FStructuralVisit::FromNative<&TupleGetItemVisit>())
+ .attr(refl::type_attr::kStructuralMutate,
+ ffi::FStructuralMutate::FromNative<&TupleGetItemMutate>())
.attr(refl::type_attr::kStructuralMaybeInplaceMutate,
- reinterpret_cast<void*>(&TupleGetItemMaybeInplaceMutate));
+ ffi::FStructuralMutate::FromNative<&TupleGetItemMaybeInplaceMutate>());
refl::GlobalDef().def("ir.TupleGetItem", [](Expr tuple, int index, Span span) {
return TupleGetItem(tuple, index, span);
@@ -815,25 +856,27 @@ TVM_FFI_STATIC_INIT_BLOCK() {
GenericConstNode::RegisterReflection();
refl::TypeAttrDef<GenericConstNode>()
.attr(refl::type_attr::kStructuralVisit,
- reinterpret_cast<void*>(&ConstantVisit<GenericConstNode>))
+ ffi::FStructuralVisit::FromNative<&ConstantVisit<GenericConstNode>>())
.attr(refl::type_attr::kStructuralMutate,
- reinterpret_cast<void*>(&ConstantMutate<GenericConstNode>))
+ ffi::FStructuralMutate::FromNative<&ConstantMutate<GenericConstNode>>())
.attr(refl::type_attr::kStructuralMaybeInplaceMutate,
- reinterpret_cast<void*>(&ConstantMaybeInplaceMutate<GenericConstNode>));
+ ffi::FStructuralMutate::FromNative<&ConstantMaybeInplaceMutate<GenericConstNode>>());
StringImmNode::RegisterReflection();
refl::TypeAttrDef<StringImmNode>()
.attr(refl::type_attr::kStructuralVisit,
- reinterpret_cast<void*>(&ConstantVisit<StringImmNode>))
+ ffi::FStructuralVisit::FromNative<&ConstantVisit<StringImmNode>>())
.attr(refl::type_attr::kStructuralMutate,
- reinterpret_cast<void*>(&ConstantMutate<StringImmNode>))
+ ffi::FStructuralMutate::FromNative<&ConstantMutate<StringImmNode>>())
.attr(refl::type_attr::kStructuralMaybeInplaceMutate,
- reinterpret_cast<void*>(&ConstantMaybeInplaceMutate<StringImmNode>));
+ ffi::FStructuralMutate::FromNative<&ConstantMaybeInplaceMutate<StringImmNode>>());
DataTypeImmNode::RegisterReflection();
refl::TypeAttrDef<DataTypeImmNode>()
- .attr(refl::type_attr::kStructuralVisit, reinterpret_cast<void*>(&DataTypeImmVisit))
- .attr(refl::type_attr::kStructuralMutate, reinterpret_cast<void*>(&DataTypeImmMutate))
+ .attr(refl::type_attr::kStructuralVisit,
+ ffi::FStructuralVisit::FromNative<&DataTypeImmVisit>())
+ .attr(refl::type_attr::kStructuralMutate,
+ ffi::FStructuralMutate::FromNative<&DataTypeImmMutate>())
.attr(refl::type_attr::kStructuralMaybeInplaceMutate,
- reinterpret_cast<void*>(&DataTypeImmMaybeInplaceMutate));
+ ffi::FStructuralMutate::FromNative<&DataTypeImmMaybeInplaceMutate>());
refl::GlobalDef()
.def("ir.GenericConst",
[](ffi::Any value, Type ty, Span span) {
@@ -894,10 +937,10 @@ TVM_FFI_STATIC_INIT_BLOCK() {
namespace refl = tvm::ffi::reflection;
IntImmNode::RegisterReflection();
refl::TypeAttrDef<IntImmNode>()
- .attr(refl::type_attr::kStructuralVisit, reinterpret_cast<void*>(&IntImmVisit))
- .attr(refl::type_attr::kStructuralMutate, reinterpret_cast<void*>(&IntImmMutate))
+ .attr(refl::type_attr::kStructuralVisit, ffi::FStructuralVisit::FromNative<&IntImmVisit>())
+ .attr(refl::type_attr::kStructuralMutate, ffi::FStructuralMutate::FromNative<&IntImmMutate>())
.attr(refl::type_attr::kStructuralMaybeInplaceMutate,
- reinterpret_cast<void*>(&IntImmMaybeInplaceMutate));
+ ffi::FStructuralMutate::FromNative<&IntImmMaybeInplaceMutate>());
refl::GlobalDef().def("ir.IntImm", [](DLDataType dtype, ffi::BigInt value, Span span) {
return IntImm(PrimType(dtype), value, span);
@@ -1021,10 +1064,11 @@ TVM_FFI_STATIC_INIT_BLOCK() {
namespace refl = tvm::ffi::reflection;
FloatImmNode::RegisterReflection();
refl::TypeAttrDef<FloatImmNode>()
- .attr(refl::type_attr::kStructuralVisit, reinterpret_cast<void*>(&FloatImmVisit))
- .attr(refl::type_attr::kStructuralMutate, reinterpret_cast<void*>(&FloatImmMutate))
+ .attr(refl::type_attr::kStructuralVisit, ffi::FStructuralVisit::FromNative<&FloatImmVisit>())
+ .attr(refl::type_attr::kStructuralMutate,
+ ffi::FStructuralMutate::FromNative<&FloatImmMutate>())
.attr(refl::type_attr::kStructuralMaybeInplaceMutate,
- reinterpret_cast<void*>(&FloatImmMaybeInplaceMutate));
+ ffi::FStructuralMutate::FromNative<&FloatImmMaybeInplaceMutate>());
refl::GlobalDef().def("ir.FloatImm", [](DLDataType dtype, double value, Span span) {
return FloatImm(PrimType(dtype), value, span);
@@ -1040,10 +1084,10 @@ TVM_FFI_STATIC_INIT_BLOCK() {
namespace refl = tvm::ffi::reflection;
RangeNode::RegisterReflection();
refl::TypeAttrDef<RangeNode>()
- .attr(refl::type_attr::kStructuralVisit, reinterpret_cast<void*>(&RangeVisit))
- .attr(refl::type_attr::kStructuralMutate, reinterpret_cast<void*>(&RangeMutate))
+ .attr(refl::type_attr::kStructuralVisit, ffi::FStructuralVisit::FromNative<&RangeVisit>())
+ .attr(refl::type_attr::kStructuralMutate, ffi::FStructuralMutate::FromNative<&RangeMutate>())
.attr(refl::type_attr::kStructuralMaybeInplaceMutate,
- reinterpret_cast<void*>(&RangeMaybeInplaceMutate));
+ ffi::FStructuralMutate::FromNative<&RangeMaybeInplaceMutate>());
refl::GlobalDef().def("ir.Range",
[](PrimExpr begin, ffi::Optional<PrimExpr> end, Span span) -> Range {
@@ -1079,10 +1123,10 @@ TVM_FFI_STATIC_INIT_BLOCK() {
namespace refl = tvm::ffi::reflection;
VarNode::RegisterReflection();
refl::TypeAttrDef<VarNode>()
- .attr(refl::type_attr::kStructuralVisit, reinterpret_cast<void*>(&VarVisit))
- .attr(refl::type_attr::kStructuralMutate, reinterpret_cast<void*>(&VarMutate))
+ .attr(refl::type_attr::kStructuralVisit, ffi::FStructuralVisit::FromNative<&VarVisit>())
+ .attr(refl::type_attr::kStructuralMutate, ffi::FStructuralMutate::FromNative<&VarMutate>())
.attr(refl::type_attr::kStructuralMaybeInplaceMutate,
- reinterpret_cast<void*>(&VarMaybeInplaceMutate));
+ ffi::FStructuralMutate::FromNative<&VarMaybeInplaceMutate>());
refl::GlobalDef().def("ir.Var", [](ffi::String name, ffi::Optional<Type> ty_annotation,
Span span) { return Var(name, ty_annotation, span); });
@@ -1100,11 +1144,13 @@ Var Var::CopyWithSuffix(const ffi::String& suffix) const {
return CopyWithName(get()->name + suffix);
}
-Var Var::CopyWithDType(PrimType dtype) const {
+Var Var::CopyWithDType(PrimType dtype) const { return CopyWithType(std::move(dtype)); }
+
+Var Var::CopyWithType(Type type) const {
TVM_FFI_CHECK_EQ(type_index(), VarNode::RuntimeTypeIndex(), TypeError)
<< "Cannot copy a Var runtime subtype as an ordinary Var";
ffi::ObjectPtr<VarNode> copy = ffi::make_object<VarNode>(*get());
- copy->ExprNode::ty = std::move(dtype);
+ copy->ExprNode::ty = std::move(type);
return Var(std::move(copy));
}
@@ -1120,10 +1166,11 @@ TVM_FFI_STATIC_INIT_BLOCK() {
namespace refl = tvm::ffi::reflection;
GlobalVarNode::RegisterReflection();
refl::TypeAttrDef<GlobalVarNode>()
- .attr(refl::type_attr::kStructuralVisit, reinterpret_cast<void*>(&GlobalVarVisit))
- .attr(refl::type_attr::kStructuralMutate, reinterpret_cast<void*>(&GlobalVarMutate))
+ .attr(refl::type_attr::kStructuralVisit, ffi::FStructuralVisit::FromNative<&GlobalVarVisit>())
+ .attr(refl::type_attr::kStructuralMutate,
+ ffi::FStructuralMutate::FromNative<&GlobalVarMutate>())
.attr(refl::type_attr::kStructuralMaybeInplaceMutate,
- reinterpret_cast<void*>(&GlobalVarMaybeInplaceMutate));
+ ffi::FStructuralMutate::FromNative<&GlobalVarMaybeInplaceMutate>());
refl::GlobalDef().def("ir.GlobalVar", [](ffi::String name) { return GlobalVar(name); });
}
@@ -1157,14 +1204,29 @@ Call::Call(Type ret_ty, Expr op, ffi::Array<Expr> args, Attrs attrs, ffi::Array<
data_ = std::move(n);
}
+Type Call::ReinferType(const CallNode* call) {
+ TVM_FFI_CHECK(call != nullptr, ValueError) << "Call::ReinferType expects a defined Call";
+ auto op = call->op.as<Op>();
+ TVM_FFI_CHECK(op.has_value(), ValueError) << "Call::ReinferType requires an Op callee";
+ TVM_FFI_CHECK(Op::HasAttrMap("FInferType"), ValueError)
+ << "No context-free FInferType hook is registered for " << op.value();
+ static auto infer_type = Op::GetAttrMap<FInferType>("FInferType");
+ TVM_FFI_CHECK(infer_type.count(op.value()), ValueError)
+ << "No context-free FInferType hook is registered for " << op.value();
+ Type result = infer_type[op.value()](call).value();
+ TVM_FFI_CHECK(!result.IsMissing(), InternalError)
+ << "FInferType for " << op.value() << " returned Type::Missing()";
+ return result;
+}
+
TVM_FFI_STATIC_INIT_BLOCK() {
namespace refl = tvm::ffi::reflection;
CallNode::RegisterReflection();
refl::TypeAttrDef<CallNode>()
- .attr(refl::type_attr::kStructuralVisit, reinterpret_cast<void*>(&CallVisit))
- .attr(refl::type_attr::kStructuralMutate, reinterpret_cast<void*>(&CallMutate))
+ .attr(refl::type_attr::kStructuralVisit, ffi::FStructuralVisit::FromNative<&CallVisit>())
+ .attr(refl::type_attr::kStructuralMutate, ffi::FStructuralMutate::FromNative<&CallMutate>())
.attr(refl::type_attr::kStructuralMaybeInplaceMutate,
- reinterpret_cast<void*>(&CallMaybeInplaceMutate));
+ ffi::FStructuralMutate::FromNative<&CallMaybeInplaceMutate>());
refl::GlobalDef().def("ir.Call", [](Type ret_ty, Expr op, ffi::Array<Expr> args, Attrs attrs,
ffi::Array<Type> ty_args, Span span) {
@@ -1174,6 +1236,8 @@ TVM_FFI_STATIC_INIT_BLOCK() {
Attrs attrs, ffi::Array<Type> ty_args, Span span) {
return Call::Unchecked(ret_ty, op, args, attrs, ty_args, span);
});
+ refl::GlobalDef().def("ir.reinfer_type",
+ [](const Call& call) { return Call::ReinferType(call.get()); });
}
TVM_FFI_STATIC_INIT_BLOCK() {
diff --git a/src/ir/expr_functor.cc b/src/ir/expr_functor.cc
index 046508e987..2bcb127527 100644
--- a/src/ir/expr_functor.cc
+++ b/src/ir/expr_functor.cc
@@ -509,6 +509,22 @@ UnchangedOr<Expr> ExprMutator::Mutate_(const CallNode* node, InplaceMode inplace
return Expr(std::move(copy));
}
+UnchangedOr<Expr> ExprMutator::ReinferMutatedCallType(UnchangedOr<Expr> mutated,
+ const CallNode* original,
+ InplaceMode inplace_mode) {
+ bool is_replacement = !mutated.IsUnchanged();
+ const CallNode* call = is_replacement ? ffi::AnyView(mutated).as<CallNode>() : original;
+ Type type = Call::ReinferType(call);
+ if (call->ty.same_as(type)) return mutated;
+ if (call->unique() && (is_replacement || inplace_mode == InplaceMode::kAllow)) {
+ const_cast<CallNode*>(call)->ty = std::move(type);
+ return mutated;
+ }
+ auto copy = ffi::make_object<CallNode>(*call);
+ copy->ty = std::move(type);
+ return Expr(std::move(copy));
+}
+
UnchangedOr<Expr> ExprMutator::Mutate_(const GenericConstNode* node, InplaceMode inplace_mode) {
auto ty = Mutate(node->ty, inplace_mode).as_or_throw<UnchangedOr<Type>>();
if (ty.UnchangedOrSameAs(node->ty)) return ffi::Unchanged();
@@ -651,25 +667,40 @@ UnchangedOr<PrimExpr> ExprMutator::Mutate_(const prim::SelectNode* node, Inplace
}
UnchangedOr<PrimExpr> ExprMutator::Mutate_(const prim::LetNode* node, InplaceMode inplace_mode) {
+ Type old_value_type = node->value->ty;
+ auto value_u = Mutate(node->value, inplace_mode);
+ bool value_changed = !value_u.UnchangedOrSameAs(node->value);
+ PrimExpr value = std::move(value_u).ValueOrUnchanged(node->value);
+ value_changed |= !old_value_type.same_as(value->ty);
+ Type old_var_type = node->var->ty;
auto var_u = WithDefRegionKind(kTVMFFIDefRegionKindSimple, [&] {
return Mutate(node->var, inplace_mode);
}).as_or_throw<UnchangedOr<Var>>();
- auto value_u = Mutate(node->value, inplace_mode);
+ bool var_changed = !var_u.UnchangedOrSameAs(node->var);
+ Var var = std::move(var_u).ValueOrUnchanged(node->var);
+ var_changed |= !old_var_type.same_as(var->ty);
+ if ((value_changed || var_changed) && !var->ty.same_as(value->ty)) {
+ var = var.CopyWithType(value->ty);
+ }
+ if (!var.same_as(node->var)) VarRemapSet(node->var, var);
auto body_u = Mutate(node->body, inplace_mode);
- if (var_u.UnchangedOrSameAs(node->var) && value_u.UnchangedOrSameAs(node->value) &&
- body_u.UnchangedOrSameAs(node->body))
+ PrimExpr body = std::move(body_u).ValueOrUnchanged(node->body);
+ if (var.same_as(node->var) && value.same_as(node->value) && body.same_as(node->body) &&
+ body->ty.same_as(node->ty))
return ffi::Unchanged();
if (inplace_mode == InplaceMode::kAllow) {
auto* writable = const_cast<prim::LetNode*>(node);
- if (!var_u.IsUnchanged()) writable->var = std::move(var_u).ValueUnchecked();
- if (!value_u.IsUnchanged()) writable->value = std::move(value_u).ValueUnchecked();
- if (!body_u.IsUnchanged()) writable->body = std::move(body_u).ValueUnchecked();
+ writable->var = std::move(var);
+ writable->value = std::move(value);
+ writable->body = std::move(body);
+ writable->ty = writable->body->ty;
return ffi::Unchanged();
}
auto copy = ffi::make_object<prim::LetNode>(*node);
- if (!var_u.IsUnchanged()) copy->var = std::move(var_u).ValueUnchecked();
- if (!value_u.IsUnchanged()) copy->value = std::move(value_u).ValueUnchecked();
- if (!body_u.IsUnchanged()) copy->body = std::move(body_u).ValueUnchecked();
+ copy->var = std::move(var);
+ copy->value = std::move(value);
+ copy->body = std::move(body);
+ copy->ty = copy->body->ty;
return PrimExpr(std::move(copy));
}
diff --git a/src/ir/op.cc b/src/ir/op.cc
index 5499ec21dc..0d518dc555 100644
--- a/src/ir/op.cc
+++ b/src/ir/op.cc
@@ -293,17 +293,20 @@ void OpNode::RegisterReflection() {
namespace {
-TVMFFIAny OpVisit(ffi::StructuralVisitorObj*, ffi::AnyView) noexcept {
+TVM_FFI_INLINE ffi::Expected<ffi::Optional<ffi::VisitInterrupt>> OpVisit(ffi::StructuralVisitorObj*,
+ ffi::AnyView) noexcept {
// Ops are unique registry atoms. Avoid reflecting through their registry metadata.
- return ffi::AnyView(nullptr).CopyToTVMFFIAny();
+ return std::nullopt;
}
-TVMFFIAny OpMutate(ffi::StructuralMutatorObj*, ffi::AnyView) noexcept {
- return ffi::Unchanged().CopyToTVMFFIAny();
+TVM_FFI_INLINE ffi::Expected<ffi::UnchangedOr<ffi::Any>> OpMutate(ffi::StructuralMutatorObj*,
+ ffi::AnyView) noexcept {
+ return ffi::Unchanged();
}
-TVMFFIAny OpMaybeInplaceMutate(ffi::StructuralMutatorObj*, ffi::AnyView) noexcept {
- return ffi::Unchanged().CopyToTVMFFIAny();
+TVM_FFI_INLINE ffi::Expected<ffi::UnchangedOr<ffi::Any>> OpMaybeInplaceMutate(
+ ffi::StructuralMutatorObj*, ffi::AnyView) noexcept {
+ return ffi::Unchanged();
}
} // namespace
@@ -313,10 +316,10 @@ TVM_FFI_STATIC_INIT_BLOCK() {
ArgumentInfoNode::RegisterReflection();
OpNode::RegisterReflection();
refl::TypeAttrDef<OpNode>()
- .attr(refl::type_attr::kStructuralVisit, reinterpret_cast<void*>(&OpVisit))
- .attr(refl::type_attr::kStructuralMutate, reinterpret_cast<void*>(&OpMutate))
+ .attr(refl::type_attr::kStructuralVisit, ffi::FStructuralVisit::FromNative<&OpVisit>())
+ .attr(refl::type_attr::kStructuralMutate, ffi::FStructuralMutate::FromNative<&OpMutate>())
.attr(refl::type_attr::kStructuralMaybeInplaceMutate,
- reinterpret_cast<void*>(&OpMaybeInplaceMutate))
+ ffi::FStructuralMutate::FromNative<&OpMaybeInplaceMutate>())
.def("__data_to_json__", [](const OpNode* node) { return node->name; })
.def("__data_from_json__", &Op::Get);
// clang-format off
diff --git a/src/ir/prim/expr.cc b/src/ir/prim/expr.cc
index 78ca9b032c..71f24e4d4a 100644
--- a/src/ir/prim/expr.cc
+++ b/src/ir/prim/expr.cc
@@ -49,17 +49,19 @@ TVM_FFI_INLINE const PrimTypeNode* GetPrimTypeNode(const PrimExpr& expr) {
// Structural traversal hooks
template <typename TNode>
-TVMFFIAny BinaryVisit(ffi::StructuralVisitorObj* visitor, ffi::AnyView value) noexcept {
+TVM_FFI_INLINE ffi::Expected<ffi::Optional<ffi::VisitInterrupt>> BinaryVisit(
+ ffi::StructuralVisitorObj* visitor, ffi::AnyView value) noexcept {
// skips: PrimExpr types are always PrimType and remain unchanged in normal mutation.
const TNode* self =
ffi::details::AnyUnsafe::RawObjectPtrFromAnyViewAfterCheck<const TNode>(value);
TVM_FFI_S_VISIT_MAYBE_EARLY_RETURN(visitor->VisitExpected(self->a));
TVM_FFI_S_VISIT_MAYBE_EARLY_RETURN(visitor->VisitExpected(self->b));
- return ffi::AnyView(nullptr).CopyToTVMFFIAny();
+ return std::nullopt;
}
template <typename TNode>
-TVMFFIAny BinaryMutate(ffi::StructuralMutatorObj* mutator, ffi::AnyView value) noexcept {
+TVM_FFI_INLINE ffi::Expected<ffi::UnchangedOr<ffi::Any>> BinaryMutate(
+ ffi::StructuralMutatorObj* mutator, ffi::AnyView value) noexcept {
// skips: PrimExpr types are always PrimType and remain unchanged in normal mutation.
const TNode* self =
ffi::details::AnyUnsafe::RawObjectPtrFromAnyViewAfterCheck<const TNode>(value);
@@ -68,7 +70,7 @@ TVMFFIAny BinaryMutate(ffi::StructuralMutatorObj* mutator, ffi::AnyView value) n
TVM_FFI_S_MUTATE_ASSIGN_OR_RETURN(ffi::UnchangedOr<PrimExpr>, b_u,
mutator->MutateExpected(self->b));
if (a_u.UnchangedOrSameAs(self->a) && b_u.UnchangedOrSameAs(self->b)) {
- return ffi::Unchanged().CopyToTVMFFIAny();
+ return ffi::Unchanged();
}
ffi::ObjectPtr<TNode> copy = ffi::make_object<TNode>(*self);
// StructuralMap preserves node types. A rewrite that changes operand dtypes must keep the
@@ -76,12 +78,12 @@ TVMFFIAny BinaryMutate(ffi::StructuralMutatorObj* mutator, ffi::AnyView value) n
// casts that would require.
if (!a_u.IsUnchanged()) copy->a = std::move(a_u).ValueUnchecked();
if (!b_u.IsUnchanged()) copy->b = std::move(b_u).ValueUnchecked();
- return ffi::details::AnyUnsafe::MoveAnyToTVMFFIAny(ffi::Any(std::move(copy)));
+ return ffi::Any(std::move(copy));
}
template <typename TNode>
-TVMFFIAny BinaryMaybeInplaceMutate(ffi::StructuralMutatorObj* mutator,
- ffi::AnyView value) noexcept {
+TVM_FFI_INLINE ffi::Expected<ffi::UnchangedOr<ffi::Any>> BinaryMaybeInplaceMutate(
+ ffi::StructuralMutatorObj* mutator, ffi::AnyView value) noexcept {
// skips: PrimExpr types are always PrimType and remain unchanged in normal mutation.
TNode* self = const_cast<TNode*>(
ffi::details::AnyUnsafe::RawObjectPtrFromAnyViewAfterCheck<const TNode>(value));
@@ -90,130 +92,140 @@ TVMFFIAny BinaryMaybeInplaceMutate(ffi::StructuralMutatorObj* mutator,
TVM_FFI_S_MUTATE_ASSIGN_OR_RETURN(ffi::UnchangedOr<PrimExpr>, b_u,
mutator->MutateExpected(self->b, ffi::InplaceMode::kAllow));
if (a_u.UnchangedOrSameAs(self->a) && b_u.UnchangedOrSameAs(self->b)) {
- return ffi::Unchanged().CopyToTVMFFIAny();
+ return ffi::Unchanged();
}
if (!a_u.IsUnchanged()) self->a = std::move(a_u).ValueUnchecked();
if (!b_u.IsUnchanged()) self->b = std::move(b_u).ValueUnchecked();
- return ffi::Unchanged().CopyToTVMFFIAny();
+ return ffi::Unchanged();
}
-TVMFFIAny CastVisit(ffi::StructuralVisitorObj* visitor, ffi::AnyView value) noexcept {
+TVM_FFI_INLINE ffi::Expected<ffi::Optional<ffi::VisitInterrupt>> CastVisit(
+ ffi::StructuralVisitorObj* visitor, ffi::AnyView value) noexcept {
// skips: PrimExpr types are always PrimType and remain unchanged in normal mutation.
const CastNode* self =
ffi::details::AnyUnsafe::RawObjectPtrFromAnyViewAfterCheck<const CastNode>(value);
TVM_FFI_S_VISIT_MAYBE_EARLY_RETURN(visitor->VisitExpected(self->value));
- return ffi::AnyView(nullptr).CopyToTVMFFIAny();
+ return std::nullopt;
}
-TVMFFIAny CastMutate(ffi::StructuralMutatorObj* mutator, ffi::AnyView value) noexcept {
+TVM_FFI_INLINE ffi::Expected<ffi::UnchangedOr<ffi::Any>> CastMutate(
+ ffi::StructuralMutatorObj* mutator, ffi::AnyView value) noexcept {
// skips: PrimExpr types are always PrimType and remain unchanged in normal mutation.
const CastNode* self =
ffi::details::AnyUnsafe::RawObjectPtrFromAnyViewAfterCheck<const CastNode>(value);
TVM_FFI_S_MUTATE_ASSIGN_OR_RETURN(ffi::UnchangedOr<PrimExpr>, mapped_value_u,
mutator->MutateExpected(self->value));
if (mapped_value_u.UnchangedOrSameAs(self->value)) {
- return ffi::Unchanged().CopyToTVMFFIAny();
+ return ffi::Unchanged();
}
ffi::ObjectPtr<CastNode> copy = ffi::make_object<CastNode>(*self);
if (!mapped_value_u.IsUnchanged()) copy->value = std::move(mapped_value_u).ValueUnchecked();
- return ffi::details::AnyUnsafe::MoveAnyToTVMFFIAny(ffi::Any(std::move(copy)));
+ return ffi::Any(std::move(copy));
}
-TVMFFIAny CastMaybeInplaceMutate(ffi::StructuralMutatorObj* mutator, ffi::AnyView value) noexcept {
+TVM_FFI_INLINE ffi::Expected<ffi::UnchangedOr<ffi::Any>> CastMaybeInplaceMutate(
+ ffi::StructuralMutatorObj* mutator, ffi::AnyView value) noexcept {
// skips: PrimExpr types are always PrimType and remain unchanged in normal mutation.
CastNode* self = const_cast<CastNode*>(
ffi::details::AnyUnsafe::RawObjectPtrFromAnyViewAfterCheck<const CastNode>(value));
TVM_FFI_S_MUTATE_ASSIGN_OR_RETURN(ffi::UnchangedOr<PrimExpr>, mapped_value_u,
mutator->MutateExpected(self->value, ffi::InplaceMode::kAllow));
if (mapped_value_u.UnchangedOrSameAs(self->value)) {
- return ffi::Unchanged().CopyToTVMFFIAny();
+ return ffi::Unchanged();
}
if (!mapped_value_u.IsUnchanged()) self->value = std::move(mapped_value_u).ValueUnchecked();
- return ffi::Unchanged().CopyToTVMFFIAny();
+ return ffi::Unchanged();
}
-TVMFFIAny NotVisit(ffi::StructuralVisitorObj* visitor, ffi::AnyView value) noexcept {
+TVM_FFI_INLINE ffi::Expected<ffi::Optional<ffi::VisitInterrupt>> NotVisit(
+ ffi::StructuralVisitorObj* visitor, ffi::AnyView value) noexcept {
// skips: PrimExpr types are always PrimType and remain unchanged in normal mutation.
const NotNode* self =
ffi::details::AnyUnsafe::RawObjectPtrFromAnyViewAfterCheck<const NotNode>(value);
TVM_FFI_S_VISIT_MAYBE_EARLY_RETURN(visitor->VisitExpected(self->a));
- return ffi::AnyView(nullptr).CopyToTVMFFIAny();
+ return std::nullopt;
}
-TVMFFIAny NotMutate(ffi::StructuralMutatorObj* mutator, ffi::AnyView value) noexcept {
+TVM_FFI_INLINE ffi::Expected<ffi::UnchangedOr<ffi::Any>> NotMutate(
+ ffi::StructuralMutatorObj* mutator, ffi::AnyView value) noexcept {
// skips: PrimExpr types are always PrimType and remain unchanged in normal mutation.
const NotNode* self =
ffi::details::AnyUnsafe::RawObjectPtrFromAnyViewAfterCheck<const NotNode>(value);
TVM_FFI_S_MUTATE_ASSIGN_OR_RETURN(ffi::UnchangedOr<PrimExpr>, mapped_a_u,
mutator->MutateExpected(self->a));
if (mapped_a_u.UnchangedOrSameAs(self->a)) {
- return ffi::Unchanged().CopyToTVMFFIAny();
+ return ffi::Unchanged();
}
ffi::ObjectPtr<NotNode> copy = ffi::make_object<NotNode>(*self);
if (!mapped_a_u.IsUnchanged()) copy->a = std::move(mapped_a_u).ValueUnchecked();
- return ffi::details::AnyUnsafe::MoveAnyToTVMFFIAny(ffi::Any(std::move(copy)));
+ return ffi::Any(std::move(copy));
}
-TVMFFIAny NotMaybeInplaceMutate(ffi::StructuralMutatorObj* mutator, ffi::AnyView value) noexcept {
+TVM_FFI_INLINE ffi::Expected<ffi::UnchangedOr<ffi::Any>> NotMaybeInplaceMutate(
+ ffi::StructuralMutatorObj* mutator, ffi::AnyView value) noexcept {
// skips: PrimExpr types are always PrimType and remain unchanged in normal mutation.
NotNode* self = const_cast<NotNode*>(
ffi::details::AnyUnsafe::RawObjectPtrFromAnyViewAfterCheck<const NotNode>(value));
TVM_FFI_S_MUTATE_ASSIGN_OR_RETURN(ffi::UnchangedOr<PrimExpr>, mapped_a_u,
mutator->MutateExpected(self->a, ffi::InplaceMode::kAllow));
if (mapped_a_u.UnchangedOrSameAs(self->a)) {
- return ffi::Unchanged().CopyToTVMFFIAny();
+ return ffi::Unchanged();
}
if (!mapped_a_u.IsUnchanged()) self->a = std::move(mapped_a_u).ValueUnchecked();
- return ffi::Unchanged().CopyToTVMFFIAny();
+ return ffi::Unchanged();
}
-TVMFFIAny BitwiseNotVisit(ffi::StructuralVisitorObj* visitor, ffi::AnyView value) noexcept {
+TVM_FFI_INLINE ffi::Expected<ffi::Optional<ffi::VisitInterrupt>> BitwiseNotVisit(
+ ffi::StructuralVisitorObj* visitor, ffi::AnyView value) noexcept {
// skips: PrimExpr types are always PrimType and remain unchanged in normal mutation.
const BitwiseNotNode* self =
ffi::details::AnyUnsafe::RawObjectPtrFromAnyViewAfterCheck<const BitwiseNotNode>(value);
TVM_FFI_S_VISIT_MAYBE_EARLY_RETURN(visitor->VisitExpected(self->a));
- return ffi::AnyView(nullptr).CopyToTVMFFIAny();
+ return std::nullopt;
}
-TVMFFIAny BitwiseNotMutate(ffi::StructuralMutatorObj* mutator, ffi::AnyView value) noexcept {
+TVM_FFI_INLINE ffi::Expected<ffi::UnchangedOr<ffi::Any>> BitwiseNotMutate(
+ ffi::StructuralMutatorObj* mutator, ffi::AnyView value) noexcept {
// skips: PrimExpr types are always PrimType and remain unchanged in normal mutation.
const BitwiseNotNode* self =
ffi::details::AnyUnsafe::RawObjectPtrFromAnyViewAfterCheck<const BitwiseNotNode>(value);
TVM_FFI_S_MUTATE_ASSIGN_OR_RETURN(ffi::UnchangedOr<PrimExpr>, mapped_a_u,
mutator->MutateExpected(self->a));
if (mapped_a_u.UnchangedOrSameAs(self->a)) {
- return ffi::Unchanged().CopyToTVMFFIAny();
+ return ffi::Unchanged();
}
ffi::ObjectPtr<BitwiseNotNode> copy = ffi::make_object<BitwiseNotNode>(*self);
if (!mapped_a_u.IsUnchanged()) copy->a = std::move(mapped_a_u).ValueUnchecked();
- return ffi::details::AnyUnsafe::MoveAnyToTVMFFIAny(ffi::Any(std::move(copy)));
+ return ffi::Any(std::move(copy));
}
-TVMFFIAny BitwiseNotMaybeInplaceMutate(ffi::StructuralMutatorObj* mutator,
- ffi::AnyView value) noexcept {
+TVM_FFI_INLINE ffi::Expected<ffi::UnchangedOr<ffi::Any>> BitwiseNotMaybeInplaceMutate(
+ ffi::StructuralMutatorObj* mutator, ffi::AnyView value) noexcept {
// skips: PrimExpr types are always PrimType and remain unchanged in normal mutation.
BitwiseNotNode* self = const_cast<BitwiseNotNode*>(
ffi::details::AnyUnsafe::RawObjectPtrFromAnyViewAfterCheck<const BitwiseNotNode>(value));
TVM_FFI_S_MUTATE_ASSIGN_OR_RETURN(ffi::UnchangedOr<PrimExpr>, mapped_a_u,
mutator->MutateExpected(self->a, ffi::InplaceMode::kAllow));
if (mapped_a_u.UnchangedOrSameAs(self->a)) {
- return ffi::Unchanged().CopyToTVMFFIAny();
+ return ffi::Unchanged();
}
if (!mapped_a_u.IsUnchanged()) self->a = std::move(mapped_a_u).ValueUnchecked();
- return ffi::Unchanged().CopyToTVMFFIAny();
+ return ffi::Unchanged();
}
-TVMFFIAny SelectVisit(ffi::StructuralVisitorObj* visitor, ffi::AnyView value) noexcept {
+TVM_FFI_INLINE ffi::Expected<ffi::Optional<ffi::VisitInterrupt>> SelectVisit(
+ ffi::StructuralVisitorObj* visitor, ffi::AnyView value) noexcept {
// skips: PrimExpr types are always PrimType and remain unchanged in normal mutation.
const SelectNode* self =
ffi::details::AnyUnsafe::RawObjectPtrFromAnyViewAfterCheck<const SelectNode>(value);
TVM_FFI_S_VISIT_MAYBE_EARLY_RETURN(visitor->VisitExpected(self->condition));
TVM_FFI_S_VISIT_MAYBE_EARLY_RETURN(visitor->VisitExpected(self->true_value));
TVM_FFI_S_VISIT_MAYBE_EARLY_RETURN(visitor->VisitExpected(self->false_value));
- return ffi::AnyView(nullptr).CopyToTVMFFIAny();
+ return std::nullopt;
}
-TVMFFIAny SelectMutate(ffi::StructuralMutatorObj* mutator, ffi::AnyView value) noexcept {
+TVM_FFI_INLINE ffi::Expected<ffi::UnchangedOr<ffi::Any>> SelectMutate(
+ ffi::StructuralMutatorObj* mutator, ffi::AnyView value) noexcept {
// skips: PrimExpr types are always PrimType and remain unchanged in normal mutation.
const SelectNode* self =
ffi::details::AnyUnsafe::RawObjectPtrFromAnyViewAfterCheck<const SelectNode>(value);
@@ -226,7 +238,7 @@ TVMFFIAny SelectMutate(ffi::StructuralMutatorObj* mutator, ffi::AnyView value) n
if (mapped_condition_u.UnchangedOrSameAs(self->condition) &&
mapped_true_value_u.UnchangedOrSameAs(self->true_value) &&
mapped_false_value_u.UnchangedOrSameAs(self->false_value)) {
- return ffi::Unchanged().CopyToTVMFFIAny();
+ return ffi::Unchanged();
}
ffi::ObjectPtr<SelectNode> copy = ffi::make_object<SelectNode>(*self);
if (!mapped_condition_u.IsUnchanged())
@@ -235,11 +247,11 @@ TVMFFIAny SelectMutate(ffi::StructuralMutatorObj* mutator, ffi::AnyView value) n
copy->true_value = std::move(mapped_true_value_u).ValueUnchecked();
if (!mapped_false_value_u.IsUnchanged())
copy->false_value = std::move(mapped_false_value_u).ValueUnchecked();
- return ffi::details::AnyUnsafe::MoveAnyToTVMFFIAny(ffi::Any(std::move(copy)));
+ return ffi::Any(std::move(copy));
}
-TVMFFIAny SelectMaybeInplaceMutate(ffi::StructuralMutatorObj* mutator,
- ffi::AnyView value) noexcept {
+TVM_FFI_INLINE ffi::Expected<ffi::UnchangedOr<ffi::Any>> SelectMaybeInplaceMutate(
+ ffi::StructuralMutatorObj* mutator, ffi::AnyView value) noexcept {
// skips: PrimExpr types are always PrimType and remain unchanged in normal mutation.
SelectNode* self = const_cast<SelectNode*>(
ffi::details::AnyUnsafe::RawObjectPtrFromAnyViewAfterCheck<const SelectNode>(value));
@@ -255,7 +267,7 @@ TVMFFIAny SelectMaybeInplaceMutate(ffi::StructuralMutatorObj* mutator,
if (mapped_condition_u.UnchangedOrSameAs(self->condition) &&
mapped_true_value_u.UnchangedOrSameAs(self->true_value) &&
mapped_false_value_u.UnchangedOrSameAs(self->false_value)) {
- return ffi::Unchanged().CopyToTVMFFIAny();
+ return ffi::Unchanged();
}
if (!mapped_condition_u.IsUnchanged())
self->condition = std::move(mapped_condition_u).ValueUnchecked();
@@ -263,10 +275,11 @@ TVMFFIAny SelectMaybeInplaceMutate(ffi::StructuralMutatorObj* mutator,
self->true_value = std::move(mapped_true_value_u).ValueUnchecked();
if (!mapped_false_value_u.IsUnchanged())
self->false_value = std::move(mapped_false_value_u).ValueUnchecked();
- return ffi::Unchanged().CopyToTVMFFIAny();
+ return ffi::Unchanged();
}
-TVMFFIAny LetVisit(ffi::StructuralVisitorObj* visitor, ffi::AnyView value) noexcept {
+TVM_FFI_INLINE ffi::Expected<ffi::Optional<ffi::VisitInterrupt>> LetVisit(
+ ffi::StructuralVisitorObj* visitor, ffi::AnyView value) noexcept {
// skips: PrimExpr types are always PrimType and remain unchanged in normal mutation.
const LetNode* self =
ffi::details::AnyUnsafe::RawObjectPtrFromAnyViewAfterCheck<const LetNode>(value);
@@ -274,53 +287,87 @@ TVMFFIAny LetVisit(ffi::StructuralVisitorObj* visitor, ffi::AnyView value) noexc
kTVMFFIDefRegionKindSimple, [&]() { return visitor->VisitExpected(self->var); }));
TVM_FFI_S_VISIT_MAYBE_EARLY_RETURN(visitor->VisitExpected(self->value));
TVM_FFI_S_VISIT_MAYBE_EARLY_RETURN(visitor->VisitExpected(self->body));
- return ffi::AnyView(nullptr).CopyToTVMFFIAny();
+ return std::nullopt;
}
-TVMFFIAny LetMutate(ffi::StructuralMutatorObj* mutator, ffi::AnyView value) noexcept {
- // skips: PrimExpr types are always PrimType and remain unchanged in normal mutation.
+TVM_FFI_INLINE ffi::Expected<ffi::UnchangedOr<ffi::Any>> LetMutate(
+ ffi::StructuralMutatorObj* mutator, ffi::AnyView value) noexcept {
const LetNode* self =
ffi::details::AnyUnsafe::RawObjectPtrFromAnyViewAfterCheck<const LetNode>(value);
+ Type old_value_type = self->value->ty;
+ TVM_FFI_S_MUTATE_ASSIGN_OR_RETURN(ffi::UnchangedOr<PrimExpr>, mapped_value_u,
+ mutator->MutateExpected(self->value));
+ bool value_changed = !mapped_value_u.UnchangedOrSameAs(self->value);
+ PrimExpr new_value = std::move(mapped_value_u).ValueOrUnchanged(self->value);
+ value_changed |= !old_value_type.same_as(new_value->ty);
+ Type old_var_type = self->var->ty;
TVM_FFI_S_MUTATE_ASSIGN_OR_RETURN(ffi::UnchangedOr<Var>, mapped_var_u,
mutator->WithDefRegionKind(kTVMFFIDefRegionKindSimple, [&]() {
return mutator->MutateExpected(self->var);
}));
- TVM_FFI_S_MUTATE_ASSIGN_OR_RETURN(ffi::UnchangedOr<PrimExpr>, mapped_value_u,
- mutator->MutateExpected(self->value));
+ bool var_changed = !mapped_var_u.UnchangedOrSameAs(self->var);
+ Var new_var = std::move(mapped_var_u).ValueOrUnchanged(self->var);
+ var_changed |= !old_var_type.same_as(new_var->ty);
+ if ((value_changed || var_changed) && !new_var->ty.same_as(new_value->ty)) {
+ new_var = new_var.CopyWithType(new_value->ty);
+ }
+ if (!new_var.same_as(self->var)) {
+ auto remap = mutator->VarRemapSetExpected(self->var, new_var);
+ TVM_FFI_S_MUTATE_MAYBE_EARLY_RETURN(remap);
+ }
TVM_FFI_S_MUTATE_ASSIGN_OR_RETURN(ffi::UnchangedOr<PrimExpr>, mapped_body_u,
mutator->MutateExpected(self->body));
- if (mapped_var_u.UnchangedOrSameAs(self->var) && mapped_value_u.UnchangedOrSameAs(self->value) &&
- mapped_body_u.UnchangedOrSameAs(self->body)) {
- return ffi::Unchanged().CopyToTVMFFIAny();
+ PrimExpr new_body = std::move(mapped_body_u).ValueOrUnchanged(self->body);
+ if (new_var.same_as(self->var) && new_value.same_as(self->value) &&
+ new_body.same_as(self->body) && new_body->ty.same_as(self->ty)) {
+ return ffi::Unchanged();
}
ffi::ObjectPtr<LetNode> copy = ffi::make_object<LetNode>(*self);
- if (!mapped_var_u.IsUnchanged()) copy->var = std::move(mapped_var_u).ValueUnchecked();
- if (!mapped_value_u.IsUnchanged()) copy->value = std::move(mapped_value_u).ValueUnchecked();
- if (!mapped_body_u.IsUnchanged()) copy->body = std::move(mapped_body_u).ValueUnchecked();
- return ffi::details::AnyUnsafe::MoveAnyToTVMFFIAny(ffi::Any(std::move(copy)));
+ copy->var = std::move(new_var);
+ copy->value = std::move(new_value);
+ copy->body = std::move(new_body);
+ copy->ty = copy->body->ty;
+ return ffi::Any(std::move(copy));
}
-TVMFFIAny LetMaybeInplaceMutate(ffi::StructuralMutatorObj* mutator, ffi::AnyView value) noexcept {
- // skips: PrimExpr types are always PrimType and remain unchanged in normal mutation.
+TVM_FFI_INLINE ffi::Expected<ffi::UnchangedOr<ffi::Any>> LetMaybeInplaceMutate(
+ ffi::StructuralMutatorObj* mutator, ffi::AnyView value) noexcept {
LetNode* self = const_cast<LetNode*>(
ffi::details::AnyUnsafe::RawObjectPtrFromAnyViewAfterCheck<const LetNode>(value));
+ Type old_value_type = self->value->ty;
+ TVM_FFI_S_MUTATE_ASSIGN_OR_RETURN(ffi::UnchangedOr<PrimExpr>, mapped_value_u,
+ mutator->MutateExpected(self->value, ffi::InplaceMode::kAllow));
+ bool value_changed = !mapped_value_u.UnchangedOrSameAs(self->value);
+ PrimExpr new_value = std::move(mapped_value_u).ValueOrUnchanged(self->value);
+ value_changed |= !old_value_type.same_as(new_value->ty);
+ Type old_var_type = self->var->ty;
TVM_FFI_S_MUTATE_ASSIGN_OR_RETURN(ffi::UnchangedOr<Var>, mapped_var_u,
mutator->WithDefRegionKind(kTVMFFIDefRegionKindSimple, [&]() {
return mutator->MutateExpected(self->var,
ffi::InplaceMode::kAllow);
}));
- TVM_FFI_S_MUTATE_ASSIGN_OR_RETURN(ffi::UnchangedOr<PrimExpr>, mapped_value_u,
- mutator->MutateExpected(self->value, ffi::InplaceMode::kAllow));
+ bool var_changed = !mapped_var_u.UnchangedOrSameAs(self->var);
+ Var new_var = std::move(mapped_var_u).ValueOrUnchanged(self->var);
+ var_changed |= !old_var_type.same_as(new_var->ty);
+ if ((value_changed || var_changed) && !new_var->ty.same_as(new_value->ty)) {
+ new_var = new_var.CopyWithType(new_value->ty);
+ }
+ if (!new_var.same_as(self->var)) {
+ auto remap = mutator->VarRemapSetExpected(self->var, new_var);
+ TVM_FFI_S_MUTATE_MAYBE_EARLY_RETURN(remap);
+ }
TVM_FFI_S_MUTATE_ASSIGN_OR_RETURN(ffi::UnchangedOr<PrimExpr>, mapped_body_u,
mutator->MutateExpected(self->body, ffi::InplaceMode::kAllow));
- if (mapped_var_u.UnchangedOrSameAs(self->var) && mapped_value_u.UnchangedOrSameAs(self->value) &&
- mapped_body_u.UnchangedOrSameAs(self->body)) {
- return ffi::Unchanged().CopyToTVMFFIAny();
+ PrimExpr new_body = std::move(mapped_body_u).ValueOrUnchanged(self->body);
+ if (new_var.same_as(self->var) && new_value.same_as(self->value) &&
+ new_body.same_as(self->body) && new_body->ty.same_as(self->ty)) {
+ return ffi::Unchanged();
}
- if (!mapped_var_u.IsUnchanged()) self->var = std::move(mapped_var_u).ValueUnchecked();
- if (!mapped_value_u.IsUnchanged()) self->value = std::move(mapped_value_u).ValueUnchecked();
- if (!mapped_body_u.IsUnchanged()) self->body = std::move(mapped_body_u).ValueUnchecked();
- return ffi::Unchanged().CopyToTVMFFIAny();
+ self->var = std::move(new_var);
+ self->value = std::move(new_value);
+ self->body = std::move(new_body);
+ self->ty = self->body->ty;
+ return ffi::Unchanged();
}
} // namespace
@@ -412,10 +459,10 @@ TVM_FFI_STATIC_INIT_BLOCK() {
namespace refl = tvm::ffi::reflection;
CastNode::RegisterReflection();
refl::TypeAttrDef<CastNode>()
- .attr(refl::type_attr::kStructuralVisit, reinterpret_cast<void*>(&CastVisit))
- .attr(refl::type_attr::kStructuralMutate, reinterpret_cast<void*>(&CastMutate))
+ .attr(refl::type_attr::kStructuralVisit, ffi::FStructuralVisit::FromNative<&CastVisit>())
+ .attr(refl::type_attr::kStructuralMutate, ffi::FStructuralMutate::FromNative<&CastMutate>())
.attr(refl::type_attr::kStructuralMaybeInplaceMutate,
- reinterpret_cast<void*>(&CastMaybeInplaceMutate));
+ ffi::FStructuralMutate::FromNative<&CastMaybeInplaceMutate>());
refl::GlobalDef().def("prim.Cast", [](PrimType dtype, PrimExpr value, Span span) {
return Cast(dtype, value, span);
@@ -431,10 +478,12 @@ TVM_FFI_STATIC_INIT_BLOCK() {
refl::GlobalDef().def("prim.Add",
[](PrimExpr a, PrimExpr b, Span span) { return Add(a, b, span); });
refl::TypeAttrDef<AddNode>()
- .attr(refl::type_attr::kStructuralVisit, reinterpret_cast<void*>(&BinaryVisit<AddNode>))
- .attr(refl::type_attr::kStructuralMutate, reinterpret_cast<void*>(&BinaryMutate<AddNode>))
+ .attr(refl::type_attr::kStructuralVisit,
+ ffi::FStructuralVisit::FromNative<&BinaryVisit<AddNode>>())
+ .attr(refl::type_attr::kStructuralMutate,
+ ffi::FStructuralMutate::FromNative<&BinaryMutate<AddNode>>())
.attr(refl::type_attr::kStructuralMaybeInplaceMutate,
- reinterpret_cast<void*>(&BinaryMaybeInplaceMutate<AddNode>));
+ ffi::FStructuralMutate::FromNative<&BinaryMaybeInplaceMutate<AddNode>>());
}
// LShift
@@ -446,10 +495,12 @@ TVM_FFI_STATIC_INIT_BLOCK() {
refl::GlobalDef().def("prim.LShift",
[](PrimExpr a, PrimExpr b, Span span) { return LShift(a, b, span); });
refl::TypeAttrDef<LShiftNode>()
- .attr(refl::type_attr::kStructuralVisit, reinterpret_cast<void*>(&BinaryVisit<LShiftNode>))
- .attr(refl::type_attr::kStructuralMutate, reinterpret_cast<void*>(&BinaryMutate<LShiftNode>))
+ .attr(refl::type_attr::kStructuralVisit,
+ ffi::FStructuralVisit::FromNative<&BinaryVisit<LShiftNode>>())
+ .attr(refl::type_attr::kStructuralMutate,
+ ffi::FStructuralMutate::FromNative<&BinaryMutate<LShiftNode>>())
.attr(refl::type_attr::kStructuralMaybeInplaceMutate,
- reinterpret_cast<void*>(&BinaryMaybeInplaceMutate<LShiftNode>));
+ ffi::FStructuralMutate::FromNative<&BinaryMaybeInplaceMutate<LShiftNode>>());
}
// RShift
@@ -461,10 +512,12 @@ TVM_FFI_STATIC_INIT_BLOCK() {
refl::GlobalDef().def("prim.RShift",
[](PrimExpr a, PrimExpr b, Span span) { return RShift(a, b, span); });
refl::TypeAttrDef<RShiftNode>()
- .attr(refl::type_attr::kStructuralVisit, reinterpret_cast<void*>(&BinaryVisit<RShiftNode>))
- .attr(refl::type_attr::kStructuralMutate, reinterpret_cast<void*>(&BinaryMutate<RShiftNode>))
+ .attr(refl::type_attr::kStructuralVisit,
+ ffi::FStructuralVisit::FromNative<&BinaryVisit<RShiftNode>>())
+ .attr(refl::type_attr::kStructuralMutate,
+ ffi::FStructuralMutate::FromNative<&BinaryMutate<RShiftNode>>())
.attr(refl::type_attr::kStructuralMaybeInplaceMutate,
- reinterpret_cast<void*>(&BinaryMaybeInplaceMutate<RShiftNode>));
+ ffi::FStructuralMutate::FromNative<&BinaryMaybeInplaceMutate<RShiftNode>>());
}
// BitwiseAnd
@@ -477,11 +530,11 @@ TVM_FFI_STATIC_INIT_BLOCK() {
[](PrimExpr a, PrimExpr b, Span span) { return BitwiseAnd(a, b, span); });
refl::TypeAttrDef<BitwiseAndNode>()
.attr(refl::type_attr::kStructuralVisit,
- reinterpret_cast<void*>(&BinaryVisit<BitwiseAndNode>))
+ ffi::FStructuralVisit::FromNative<&BinaryVisit<BitwiseAndNode>>())
.attr(refl::type_attr::kStructuralMutate,
- reinterpret_cast<void*>(&BinaryMutate<BitwiseAndNode>))
+ ffi::FStructuralMutate::FromNative<&BinaryMutate<BitwiseAndNode>>())
.attr(refl::type_attr::kStructuralMaybeInplaceMutate,
- reinterpret_cast<void*>(&BinaryMaybeInplaceMutate<BitwiseAndNode>));
+ ffi::FStructuralMutate::FromNative<&BinaryMaybeInplaceMutate<BitwiseAndNode>>());
}
// BitwiseOr
@@ -493,11 +546,12 @@ TVM_FFI_STATIC_INIT_BLOCK() {
refl::GlobalDef().def("prim.BitwiseOr",
[](PrimExpr a, PrimExpr b, Span span) { return BitwiseOr(a, b, span); });
refl::TypeAttrDef<BitwiseOrNode>()
- .attr(refl::type_attr::kStructuralVisit, reinterpret_cast<void*>(&BinaryVisit<BitwiseOrNode>))
+ .attr(refl::type_attr::kStructuralVisit,
+ ffi::FStructuralVisit::FromNative<&BinaryVisit<BitwiseOrNode>>())
.attr(refl::type_attr::kStructuralMutate,
- reinterpret_cast<void*>(&BinaryMutate<BitwiseOrNode>))
+ ffi::FStructuralMutate::FromNative<&BinaryMutate<BitwiseOrNode>>())
.attr(refl::type_attr::kStructuralMaybeInplaceMutate,
- reinterpret_cast<void*>(&BinaryMaybeInplaceMutate<BitwiseOrNode>));
+ ffi::FStructuralMutate::FromNative<&BinaryMaybeInplaceMutate<BitwiseOrNode>>());
}
// BitwiseXor
@@ -510,11 +564,11 @@ TVM_FFI_STATIC_INIT_BLOCK() {
[](PrimExpr a, PrimExpr b, Span span) { return BitwiseXor(a, b, span); });
refl::TypeAttrDef<BitwiseXorNode>()
.attr(refl::type_attr::kStructuralVisit,
- reinterpret_cast<void*>(&BinaryVisit<BitwiseXorNode>))
+ ffi::FStructuralVisit::FromNative<&BinaryVisit<BitwiseXorNode>>())
.attr(refl::type_attr::kStructuralMutate,
- reinterpret_cast<void*>(&BinaryMutate<BitwiseXorNode>))
+ ffi::FStructuralMutate::FromNative<&BinaryMutate<BitwiseXorNode>>())
.attr(refl::type_attr::kStructuralMaybeInplaceMutate,
- reinterpret_cast<void*>(&BinaryMaybeInplaceMutate<BitwiseXorNode>));
+ ffi::FStructuralMutate::FromNative<&BinaryMaybeInplaceMutate<BitwiseXorNode>>());
}
// Sub
@@ -526,10 +580,12 @@ TVM_FFI_STATIC_INIT_BLOCK() {
refl::GlobalDef().def("prim.Sub",
[](PrimExpr a, PrimExpr b, Span span) { return Sub(a, b, span); });
refl::TypeAttrDef<SubNode>()
- .attr(refl::type_attr::kStructuralVisit, reinterpret_cast<void*>(&BinaryVisit<SubNode>))
- .attr(refl::type_attr::kStructuralMutate, reinterpret_cast<void*>(&BinaryMutate<SubNode>))
+ .attr(refl::type_attr::kStructuralVisit,
+ ffi::FStructuralVisit::FromNative<&BinaryVisit<SubNode>>())
+ .attr(refl::type_attr::kStructuralMutate,
+ ffi::FStructuralMutate::FromNative<&BinaryMutate<SubNode>>())
.attr(refl::type_attr::kStructuralMaybeInplaceMutate,
- reinterpret_cast<void*>(&BinaryMaybeInplaceMutate<SubNode>));
+ ffi::FStructuralMutate::FromNative<&BinaryMaybeInplaceMutate<SubNode>>());
}
// Mul
@@ -541,10 +597,12 @@ TVM_FFI_STATIC_INIT_BLOCK() {
refl::GlobalDef().def("prim.Mul",
[](PrimExpr a, PrimExpr b, Span span) { return Mul(a, b, span); });
refl::TypeAttrDef<MulNode>()
- .attr(refl::type_attr::kStructuralVisit, reinterpret_cast<void*>(&BinaryVisit<MulNode>))
- .attr(refl::type_attr::kStructuralMutate, reinterpret_cast<void*>(&BinaryMutate<MulNode>))
+ .attr(refl::type_attr::kStructuralVisit,
+ ffi::FStructuralVisit::FromNative<&BinaryVisit<MulNode>>())
+ .attr(refl::type_attr::kStructuralMutate,
+ ffi::FStructuralMutate::FromNative<&BinaryMutate<MulNode>>())
.attr(refl::type_attr::kStructuralMaybeInplaceMutate,
- reinterpret_cast<void*>(&BinaryMaybeInplaceMutate<MulNode>));
+ ffi::FStructuralMutate::FromNative<&BinaryMaybeInplaceMutate<MulNode>>());
}
// Div
@@ -556,10 +614,12 @@ TVM_FFI_STATIC_INIT_BLOCK() {
refl::GlobalDef().def("prim.Div",
[](PrimExpr a, PrimExpr b, Span span) { return Div(a, b, span); });
refl::TypeAttrDef<DivNode>()
- .attr(refl::type_attr::kStructuralVisit, reinterpret_cast<void*>(&BinaryVisit<DivNode>))
- .attr(refl::type_attr::kStructuralMutate, reinterpret_cast<void*>(&BinaryMutate<DivNode>))
+ .attr(refl::type_attr::kStructuralVisit,
+ ffi::FStructuralVisit::FromNative<&BinaryVisit<DivNode>>())
+ .attr(refl::type_attr::kStructuralMutate,
+ ffi::FStructuralMutate::FromNative<&BinaryMutate<DivNode>>())
.attr(refl::type_attr::kStructuralMaybeInplaceMutate,
- reinterpret_cast<void*>(&BinaryMaybeInplaceMutate<DivNode>));
+ ffi::FStructuralMutate::FromNative<&BinaryMaybeInplaceMutate<DivNode>>());
}
// Mod
@@ -571,10 +631,12 @@ TVM_FFI_STATIC_INIT_BLOCK() {
refl::GlobalDef().def("prim.Mod",
[](PrimExpr a, PrimExpr b, Span span) { return Mod(a, b, span); });
refl::TypeAttrDef<ModNode>()
- .attr(refl::type_attr::kStructuralVisit, reinterpret_cast<void*>(&BinaryVisit<ModNode>))
- .attr(refl::type_attr::kStructuralMutate, reinterpret_cast<void*>(&BinaryMutate<ModNode>))
+ .attr(refl::type_attr::kStructuralVisit,
+ ffi::FStructuralVisit::FromNative<&BinaryVisit<ModNode>>())
+ .attr(refl::type_attr::kStructuralMutate,
+ ffi::FStructuralMutate::FromNative<&BinaryMutate<ModNode>>())
.attr(refl::type_attr::kStructuralMaybeInplaceMutate,
- reinterpret_cast<void*>(&BinaryMaybeInplaceMutate<ModNode>));
+ ffi::FStructuralMutate::FromNative<&BinaryMaybeInplaceMutate<ModNode>>());
}
// FloorDiv
@@ -586,11 +648,12 @@ TVM_FFI_STATIC_INIT_BLOCK() {
refl::GlobalDef().def("prim.FloorDiv",
[](PrimExpr a, PrimExpr b, Span span) { return FloorDiv(a, b, span); });
refl::TypeAttrDef<FloorDivNode>()
- .attr(refl::type_attr::kStructuralVisit, reinterpret_cast<void*>(&BinaryVisit<FloorDivNode>))
+ .attr(refl::type_attr::kStructuralVisit,
+ ffi::FStructuralVisit::FromNative<&BinaryVisit<FloorDivNode>>())
.attr(refl::type_attr::kStructuralMutate,
- reinterpret_cast<void*>(&BinaryMutate<FloorDivNode>))
+ ffi::FStructuralMutate::FromNative<&BinaryMutate<FloorDivNode>>())
.attr(refl::type_attr::kStructuralMaybeInplaceMutate,
- reinterpret_cast<void*>(&BinaryMaybeInplaceMutate<FloorDivNode>));
+ ffi::FStructuralMutate::FromNative<&BinaryMaybeInplaceMutate<FloorDivNode>>());
}
// FloorMod
@@ -602,11 +665,12 @@ TVM_FFI_STATIC_INIT_BLOCK() {
refl::GlobalDef().def("prim.FloorMod",
[](PrimExpr a, PrimExpr b, Span span) { return FloorMod(a, b, span); });
refl::TypeAttrDef<FloorModNode>()
- .attr(refl::type_attr::kStructuralVisit, reinterpret_cast<void*>(&BinaryVisit<FloorModNode>))
+ .attr(refl::type_attr::kStructuralVisit,
+ ffi::FStructuralVisit::FromNative<&BinaryVisit<FloorModNode>>())
.attr(refl::type_attr::kStructuralMutate,
- reinterpret_cast<void*>(&BinaryMutate<FloorModNode>))
+ ffi::FStructuralMutate::FromNative<&BinaryMutate<FloorModNode>>())
.attr(refl::type_attr::kStructuralMaybeInplaceMutate,
- reinterpret_cast<void*>(&BinaryMaybeInplaceMutate<FloorModNode>));
+ ffi::FStructuralMutate::FromNative<&BinaryMaybeInplaceMutate<FloorModNode>>());
}
// Min
@@ -618,10 +682,12 @@ TVM_FFI_STATIC_INIT_BLOCK() {
refl::GlobalDef().def("prim.Min",
[](PrimExpr a, PrimExpr b, Span span) { return Min(a, b, span); });
refl::TypeAttrDef<MinNode>()
- .attr(refl::type_attr::kStructuralVisit, reinterpret_cast<void*>(&BinaryVisit<MinNode>))
- .attr(refl::type_attr::kStructuralMutate, reinterpret_cast<void*>(&BinaryMutate<MinNode>))
+ .attr(refl::type_attr::kStructuralVisit,
+ ffi::FStructuralVisit::FromNative<&BinaryVisit<MinNode>>())
+ .attr(refl::type_attr::kStructuralMutate,
+ ffi::FStructuralMutate::FromNative<&BinaryMutate<MinNode>>())
.attr(refl::type_attr::kStructuralMaybeInplaceMutate,
- reinterpret_cast<void*>(&BinaryMaybeInplaceMutate<MinNode>));
+ ffi::FStructuralMutate::FromNative<&BinaryMaybeInplaceMutate<MinNode>>());
}
// Max
@@ -633,10 +699,12 @@ TVM_FFI_STATIC_INIT_BLOCK() {
refl::GlobalDef().def("prim.Max",
[](PrimExpr a, PrimExpr b, Span span) { return Max(a, b, span); });
refl::TypeAttrDef<MaxNode>()
- .attr(refl::type_attr::kStructuralVisit, reinterpret_cast<void*>(&BinaryVisit<MaxNode>))
- .attr(refl::type_attr::kStructuralMutate, reinterpret_cast<void*>(&BinaryMutate<MaxNode>))
+ .attr(refl::type_attr::kStructuralVisit,
+ ffi::FStructuralVisit::FromNative<&BinaryVisit<MaxNode>>())
+ .attr(refl::type_attr::kStructuralMutate,
+ ffi::FStructuralMutate::FromNative<&BinaryMutate<MaxNode>>())
.attr(refl::type_attr::kStructuralMaybeInplaceMutate,
- reinterpret_cast<void*>(&BinaryMaybeInplaceMutate<MaxNode>));
+ ffi::FStructuralMutate::FromNative<&BinaryMaybeInplaceMutate<MaxNode>>());
}
// EQ
@@ -648,10 +716,12 @@ TVM_FFI_STATIC_INIT_BLOCK() {
refl::GlobalDef().def("prim.EQ",
[](PrimExpr a, PrimExpr b, Span span) { return EQ(a, b, span); });
refl::TypeAttrDef<EQNode>()
- .attr(refl::type_attr::kStructuralVisit, reinterpret_cast<void*>(&BinaryVisit<EQNode>))
- .attr(refl::type_attr::kStructuralMutate, reinterpret_cast<void*>(&BinaryMutate<EQNode>))
+ .attr(refl::type_attr::kStructuralVisit,
+ ffi::FStructuralVisit::FromNative<&BinaryVisit<EQNode>>())
+ .attr(refl::type_attr::kStructuralMutate,
+ ffi::FStructuralMutate::FromNative<&BinaryMutate<EQNode>>())
.attr(refl::type_attr::kStructuralMaybeInplaceMutate,
- reinterpret_cast<void*>(&BinaryMaybeInplaceMutate<EQNode>));
+ ffi::FStructuralMutate::FromNative<&BinaryMaybeInplaceMutate<EQNode>>());
}
// NE
@@ -663,10 +733,12 @@ TVM_FFI_STATIC_INIT_BLOCK() {
refl::GlobalDef().def("prim.NE",
[](PrimExpr a, PrimExpr b, Span span) { return NE(a, b, span); });
refl::TypeAttrDef<NENode>()
- .attr(refl::type_attr::kStructuralVisit, reinterpret_cast<void*>(&BinaryVisit<NENode>))
- .attr(refl::type_attr::kStructuralMutate, reinterpret_cast<void*>(&BinaryMutate<NENode>))
+ .attr(refl::type_attr::kStructuralVisit,
+ ffi::FStructuralVisit::FromNative<&BinaryVisit<NENode>>())
+ .attr(refl::type_attr::kStructuralMutate,
+ ffi::FStructuralMutate::FromNative<&BinaryMutate<NENode>>())
.attr(refl::type_attr::kStructuralMaybeInplaceMutate,
- reinterpret_cast<void*>(&BinaryMaybeInplaceMutate<NENode>));
+ ffi::FStructuralMutate::FromNative<&BinaryMaybeInplaceMutate<NENode>>());
}
// LT
@@ -678,10 +750,12 @@ TVM_FFI_STATIC_INIT_BLOCK() {
refl::GlobalDef().def("prim.LT",
[](PrimExpr a, PrimExpr b, Span span) { return LT(a, b, span); });
refl::TypeAttrDef<LTNode>()
- .attr(refl::type_attr::kStructuralVisit, reinterpret_cast<void*>(&BinaryVisit<LTNode>))
- .attr(refl::type_attr::kStructuralMutate, reinterpret_cast<void*>(&BinaryMutate<LTNode>))
+ .attr(refl::type_attr::kStructuralVisit,
+ ffi::FStructuralVisit::FromNative<&BinaryVisit<LTNode>>())
+ .attr(refl::type_attr::kStructuralMutate,
+ ffi::FStructuralMutate::FromNative<&BinaryMutate<LTNode>>())
.attr(refl::type_attr::kStructuralMaybeInplaceMutate,
- reinterpret_cast<void*>(&BinaryMaybeInplaceMutate<LTNode>));
+ ffi::FStructuralMutate::FromNative<&BinaryMaybeInplaceMutate<LTNode>>());
}
// LE
@@ -693,10 +767,12 @@ TVM_FFI_STATIC_INIT_BLOCK() {
refl::GlobalDef().def("prim.LE",
[](PrimExpr a, PrimExpr b, Span span) { return LE(a, b, span); });
refl::TypeAttrDef<LENode>()
- .attr(refl::type_attr::kStructuralVisit, reinterpret_cast<void*>(&BinaryVisit<LENode>))
- .attr(refl::type_attr::kStructuralMutate, reinterpret_cast<void*>(&BinaryMutate<LENode>))
+ .attr(refl::type_attr::kStructuralVisit,
+ ffi::FStructuralVisit::FromNative<&BinaryVisit<LENode>>())
+ .attr(refl::type_attr::kStructuralMutate,
+ ffi::FStructuralMutate::FromNative<&BinaryMutate<LENode>>())
.attr(refl::type_attr::kStructuralMaybeInplaceMutate,
- reinterpret_cast<void*>(&BinaryMaybeInplaceMutate<LENode>));
+ ffi::FStructuralMutate::FromNative<&BinaryMaybeInplaceMutate<LENode>>());
}
// GT
@@ -708,10 +784,12 @@ TVM_FFI_STATIC_INIT_BLOCK() {
refl::GlobalDef().def("prim.GT",
[](PrimExpr a, PrimExpr b, Span span) { return GT(a, b, span); });
refl::TypeAttrDef<GTNode>()
- .attr(refl::type_attr::kStructuralVisit, reinterpret_cast<void*>(&BinaryVisit<GTNode>))
- .attr(refl::type_attr::kStructuralMutate, reinterpret_cast<void*>(&BinaryMutate<GTNode>))
+ .attr(refl::type_attr::kStructuralVisit,
+ ffi::FStructuralVisit::FromNative<&BinaryVisit<GTNode>>())
+ .attr(refl::type_attr::kStructuralMutate,
+ ffi::FStructuralMutate::FromNative<&BinaryMutate<GTNode>>())
.attr(refl::type_attr::kStructuralMaybeInplaceMutate,
- reinterpret_cast<void*>(&BinaryMaybeInplaceMutate<GTNode>));
+ ffi::FStructuralMutate::FromNative<&BinaryMaybeInplaceMutate<GTNode>>());
}
// GE
@@ -723,10 +801,12 @@ TVM_FFI_STATIC_INIT_BLOCK() {
refl::GlobalDef().def("prim.GE",
[](PrimExpr a, PrimExpr b, Span span) { return GE(a, b, span); });
refl::TypeAttrDef<GENode>()
- .attr(refl::type_attr::kStructuralVisit, reinterpret_cast<void*>(&BinaryVisit<GENode>))
- .attr(refl::type_attr::kStructuralMutate, reinterpret_cast<void*>(&BinaryMutate<GENode>))
+ .attr(refl::type_attr::kStructuralVisit,
+ ffi::FStructuralVisit::FromNative<&BinaryVisit<GENode>>())
+ .attr(refl::type_attr::kStructuralMutate,
+ ffi::FStructuralMutate::FromNative<&BinaryMutate<GENode>>())
.attr(refl::type_attr::kStructuralMaybeInplaceMutate,
- reinterpret_cast<void*>(&BinaryMaybeInplaceMutate<GENode>));
+ ffi::FStructuralMutate::FromNative<&BinaryMaybeInplaceMutate<GENode>>());
}
// And
@@ -753,10 +833,12 @@ TVM_FFI_STATIC_INIT_BLOCK() {
refl::GlobalDef().def("prim.And",
[](PrimExpr a, PrimExpr b, Span span) { return And(a, b, span); });
refl::TypeAttrDef<AndNode>()
- .attr(refl::type_attr::kStructuralVisit, reinterpret_cast<void*>(&BinaryVisit<AndNode>))
- .attr(refl::type_attr::kStructuralMutate, reinterpret_cast<void*>(&BinaryMutate<AndNode>))
+ .attr(refl::type_attr::kStructuralVisit,
+ ffi::FStructuralVisit::FromNative<&BinaryVisit<AndNode>>())
+ .attr(refl::type_attr::kStructuralMutate,
+ ffi::FStructuralMutate::FromNative<&BinaryMutate<AndNode>>())
.attr(refl::type_attr::kStructuralMaybeInplaceMutate,
- reinterpret_cast<void*>(&BinaryMaybeInplaceMutate<AndNode>));
+ ffi::FStructuralMutate::FromNative<&BinaryMaybeInplaceMutate<AndNode>>());
}
// Or
@@ -783,10 +865,12 @@ TVM_FFI_STATIC_INIT_BLOCK() {
refl::GlobalDef().def("prim.Or",
[](PrimExpr a, PrimExpr b, Span span) { return Or(a, b, span); });
refl::TypeAttrDef<OrNode>()
- .attr(refl::type_attr::kStructuralVisit, reinterpret_cast<void*>(&BinaryVisit<OrNode>))
- .attr(refl::type_attr::kStructuralMutate, reinterpret_cast<void*>(&BinaryMutate<OrNode>))
+ .attr(refl::type_attr::kStructuralVisit,
+ ffi::FStructuralVisit::FromNative<&BinaryVisit<OrNode>>())
+ .attr(refl::type_attr::kStructuralMutate,
+ ffi::FStructuralMutate::FromNative<&BinaryMutate<OrNode>>())
.attr(refl::type_attr::kStructuralMaybeInplaceMutate,
- reinterpret_cast<void*>(&BinaryMaybeInplaceMutate<OrNode>));
+ ffi::FStructuralMutate::FromNative<&BinaryMaybeInplaceMutate<OrNode>>());
}
// Not
@@ -806,10 +890,10 @@ TVM_FFI_STATIC_INIT_BLOCK() {
namespace refl = tvm::ffi::reflection;
NotNode::RegisterReflection();
refl::TypeAttrDef<NotNode>()
- .attr(refl::type_attr::kStructuralVisit, reinterpret_cast<void*>(&NotVisit))
- .attr(refl::type_attr::kStructuralMutate, reinterpret_cast<void*>(&NotMutate))
+ .attr(refl::type_attr::kStructuralVisit, ffi::FStructuralVisit::FromNative<&NotVisit>())
+ .attr(refl::type_attr::kStructuralMutate, ffi::FStructuralMutate::FromNative<&NotMutate>())
.attr(refl::type_attr::kStructuralMaybeInplaceMutate,
- reinterpret_cast<void*>(&NotMaybeInplaceMutate));
+ ffi::FStructuralMutate::FromNative<&NotMaybeInplaceMutate>());
refl::GlobalDef().def("prim.Not", [](PrimExpr a, Span span) { return Not(a, span); });
}
@@ -834,10 +918,12 @@ TVM_FFI_STATIC_INIT_BLOCK() {
namespace refl = tvm::ffi::reflection;
BitwiseNotNode::RegisterReflection();
refl::TypeAttrDef<BitwiseNotNode>()
- .attr(refl::type_attr::kStructuralVisit, reinterpret_cast<void*>(&BitwiseNotVisit))
- .attr(refl::type_attr::kStructuralMutate, reinterpret_cast<void*>(&BitwiseNotMutate))
+ .attr(refl::type_attr::kStructuralVisit,
+ ffi::FStructuralVisit::FromNative<&BitwiseNotVisit>())
+ .attr(refl::type_attr::kStructuralMutate,
+ ffi::FStructuralMutate::FromNative<&BitwiseNotMutate>())
.attr(refl::type_attr::kStructuralMaybeInplaceMutate,
- reinterpret_cast<void*>(&BitwiseNotMaybeInplaceMutate));
+ ffi::FStructuralMutate::FromNative<&BitwiseNotMaybeInplaceMutate>());
refl::GlobalDef().def("prim.BitwiseNot",
[](PrimExpr a, Span span) { return BitwiseNot(a, span); });
@@ -871,10 +957,10 @@ TVM_FFI_STATIC_INIT_BLOCK() {
namespace refl = tvm::ffi::reflection;
SelectNode::RegisterReflection();
refl::TypeAttrDef<SelectNode>()
- .attr(refl::type_attr::kStructuralVisit, reinterpret_cast<void*>(&SelectVisit))
- .attr(refl::type_attr::kStructuralMutate, reinterpret_cast<void*>(&SelectMutate))
+ .attr(refl::type_attr::kStructuralVisit, ffi::FStructuralVisit::FromNative<&SelectVisit>())
+ .attr(refl::type_attr::kStructuralMutate, ffi::FStructuralMutate::FromNative<&SelectMutate>())
.attr(refl::type_attr::kStructuralMaybeInplaceMutate,
- reinterpret_cast<void*>(&SelectMaybeInplaceMutate));
+ ffi::FStructuralMutate::FromNative<&SelectMaybeInplaceMutate>());
refl::GlobalDef().def(
"prim.Select", [](PrimExpr condition, PrimExpr true_value, PrimExpr false_value, Span span) {
@@ -901,10 +987,10 @@ TVM_FFI_STATIC_INIT_BLOCK() {
namespace refl = tvm::ffi::reflection;
LetNode::RegisterReflection();
refl::TypeAttrDef<LetNode>()
- .attr(refl::type_attr::kStructuralVisit, reinterpret_cast<void*>(&LetVisit))
- .attr(refl::type_attr::kStructuralMutate, reinterpret_cast<void*>(&LetMutate))
+ .attr(refl::type_attr::kStructuralVisit, ffi::FStructuralVisit::FromNative<&LetVisit>())
+ .attr(refl::type_attr::kStructuralMutate, ffi::FStructuralMutate::FromNative<&LetMutate>())
.attr(refl::type_attr::kStructuralMaybeInplaceMutate,
- reinterpret_cast<void*>(&LetMaybeInplaceMutate));
+ ffi::FStructuralMutate::FromNative<&LetMaybeInplaceMutate>());
refl::GlobalDef().def("prim.Let", [](Var var, PrimExpr value, PrimExpr body, Span span) {
return Let(var, value, body, span);
diff --git a/src/ir/prim/vector_expr.cc b/src/ir/prim/vector_expr.cc
index d91602b332..8a70948b63 100644
--- a/src/ir/prim/vector_expr.cc
+++ b/src/ir/prim/vector_expr.cc
@@ -53,17 +53,19 @@ std::optional<int> ExtractVscaleFactor(const PrimExpr& lanes) {
return std::nullopt;
}
-TVMFFIAny RampVisit(ffi::StructuralVisitorObj* visitor, ffi::AnyView value) noexcept {
+TVM_FFI_INLINE ffi::Expected<ffi::Optional<ffi::VisitInterrupt>> RampVisit(
+ ffi::StructuralVisitorObj* visitor, ffi::AnyView value) noexcept {
// skips: PrimExpr types are always PrimType and remain unchanged in normal mutation.
const RampNode* self =
ffi::details::AnyUnsafe::RawObjectPtrFromAnyViewAfterCheck<const RampNode>(value);
TVM_FFI_S_VISIT_MAYBE_EARLY_RETURN(visitor->VisitExpected(self->base));
TVM_FFI_S_VISIT_MAYBE_EARLY_RETURN(visitor->VisitExpected(self->stride));
TVM_FFI_S_VISIT_MAYBE_EARLY_RETURN(visitor->VisitExpected(self->lanes));
- return ffi::AnyView(nullptr).CopyToTVMFFIAny();
+ return std::nullopt;
}
-TVMFFIAny RampMutate(ffi::StructuralMutatorObj* mutator, ffi::AnyView value) noexcept {
+TVM_FFI_INLINE ffi::Expected<ffi::UnchangedOr<ffi::Any>> RampMutate(
+ ffi::StructuralMutatorObj* mutator, ffi::AnyView value) noexcept {
// skips: PrimExpr types are always PrimType and remain unchanged in normal mutation.
const RampNode* self =
ffi::details::AnyUnsafe::RawObjectPtrFromAnyViewAfterCheck<const RampNode>(value);
@@ -76,16 +78,17 @@ TVMFFIAny RampMutate(ffi::StructuralMutatorObj* mutator, ffi::AnyView value) noe
if (mapped_base_u.UnchangedOrSameAs(self->base) &&
mapped_stride_u.UnchangedOrSameAs(self->stride) &&
mapped_lanes_u.UnchangedOrSameAs(self->lanes)) {
- return ffi::Unchanged().CopyToTVMFFIAny();
+ return ffi::Unchanged();
}
ffi::ObjectPtr<RampNode> copy = ffi::make_object<RampNode>(*self);
if (!mapped_base_u.IsUnchanged()) copy->base = std::move(mapped_base_u).ValueUnchecked();
if (!mapped_stride_u.IsUnchanged()) copy->stride = std::move(mapped_stride_u).ValueUnchecked();
if (!mapped_lanes_u.IsUnchanged()) copy->lanes = std::move(mapped_lanes_u).ValueUnchecked();
- return ffi::details::AnyUnsafe::MoveAnyToTVMFFIAny(ffi::Any(std::move(copy)));
+ return ffi::Any(std::move(copy));
}
-TVMFFIAny RampMaybeInplaceMutate(ffi::StructuralMutatorObj* mutator, ffi::AnyView value) noexcept {
+TVM_FFI_INLINE ffi::Expected<ffi::UnchangedOr<ffi::Any>> RampMaybeInplaceMutate(
+ ffi::StructuralMutatorObj* mutator, ffi::AnyView value) noexcept {
// skips: PrimExpr types are always PrimType and remain unchanged in normal mutation.
RampNode* self = const_cast<RampNode*>(
ffi::details::AnyUnsafe::RawObjectPtrFromAnyViewAfterCheck<const RampNode>(value));
@@ -99,24 +102,26 @@ TVMFFIAny RampMaybeInplaceMutate(ffi::StructuralMutatorObj* mutator, ffi::AnyVie
if (mapped_base_u.UnchangedOrSameAs(self->base) &&
mapped_stride_u.UnchangedOrSameAs(self->stride) &&
mapped_lanes_u.UnchangedOrSameAs(self->lanes)) {
- return ffi::Unchanged().CopyToTVMFFIAny();
+ return ffi::Unchanged();
}
if (!mapped_base_u.IsUnchanged()) self->base = std::move(mapped_base_u).ValueUnchecked();
if (!mapped_stride_u.IsUnchanged()) self->stride = std::move(mapped_stride_u).ValueUnchecked();
if (!mapped_lanes_u.IsUnchanged()) self->lanes = std::move(mapped_lanes_u).ValueUnchecked();
- return ffi::Unchanged().CopyToTVMFFIAny();
+ return ffi::Unchanged();
}
-TVMFFIAny BroadcastVisit(ffi::StructuralVisitorObj* visitor, ffi::AnyView value) noexcept {
+TVM_FFI_INLINE ffi::Expected<ffi::Optional<ffi::VisitInterrupt>> BroadcastVisit(
+ ffi::StructuralVisitorObj* visitor, ffi::AnyView value) noexcept {
// skips: PrimExpr types are always PrimType and remain unchanged in normal mutation.
const BroadcastNode* self =
ffi::details::AnyUnsafe::RawObjectPtrFromAnyViewAfterCheck<const BroadcastNode>(value);
TVM_FFI_S_VISIT_MAYBE_EARLY_RETURN(visitor->VisitExpected(self->value));
TVM_FFI_S_VISIT_MAYBE_EARLY_RETURN(visitor->VisitExpected(self->lanes));
- return ffi::AnyView(nullptr).CopyToTVMFFIAny();
+ return std::nullopt;
}
-TVMFFIAny BroadcastMutate(ffi::StructuralMutatorObj* mutator, ffi::AnyView value) noexcept {
+TVM_FFI_INLINE ffi::Expected<ffi::UnchangedOr<ffi::Any>> BroadcastMutate(
+ ffi::StructuralMutatorObj* mutator, ffi::AnyView value) noexcept {
// skips: PrimExpr types are always PrimType and remain unchanged in normal mutation.
const BroadcastNode* self =
ffi::details::AnyUnsafe::RawObjectPtrFromAnyViewAfterCheck<const BroadcastNode>(value);
@@ -126,16 +131,16 @@ TVMFFIAny BroadcastMutate(ffi::StructuralMutatorObj* mutator, ffi::AnyView value
mutator->MutateExpected(self->lanes));
if (mapped_value_u.UnchangedOrSameAs(self->value) &&
mapped_lanes_u.UnchangedOrSameAs(self->lanes)) {
- return ffi::Unchanged().CopyToTVMFFIAny();
+ return ffi::Unchanged();
}
ffi::ObjectPtr<BroadcastNode> copy = ffi::make_object<BroadcastNode>(*self);
if (!mapped_value_u.IsUnchanged()) copy->value = std::move(mapped_value_u).ValueUnchecked();
if (!mapped_lanes_u.IsUnchanged()) copy->lanes = std::move(mapped_lanes_u).ValueUnchecked();
- return ffi::details::AnyUnsafe::MoveAnyToTVMFFIAny(ffi::Any(std::move(copy)));
+ return ffi::Any(std::move(copy));
}
-TVMFFIAny BroadcastMaybeInplaceMutate(ffi::StructuralMutatorObj* mutator,
- ffi::AnyView value) noexcept {
+TVM_FFI_INLINE ffi::Expected<ffi::UnchangedOr<ffi::Any>> BroadcastMaybeInplaceMutate(
+ ffi::StructuralMutatorObj* mutator, ffi::AnyView value) noexcept {
// skips: PrimExpr types are always PrimType and remain unchanged in normal mutation.
BroadcastNode* self = const_cast<BroadcastNode*>(
ffi::details::AnyUnsafe::RawObjectPtrFromAnyViewAfterCheck<const BroadcastNode>(value));
@@ -145,23 +150,25 @@ TVMFFIAny BroadcastMaybeInplaceMutate(ffi::StructuralMutatorObj* mutator,
mutator->MutateExpected(self->lanes, ffi::InplaceMode::kAllow));
if (mapped_value_u.UnchangedOrSameAs(self->value) &&
mapped_lanes_u.UnchangedOrSameAs(self->lanes)) {
- return ffi::Unchanged().CopyToTVMFFIAny();
+ return ffi::Unchanged();
}
if (!mapped_value_u.IsUnchanged()) self->value = std::move(mapped_value_u).ValueUnchecked();
if (!mapped_lanes_u.IsUnchanged()) self->lanes = std::move(mapped_lanes_u).ValueUnchecked();
- return ffi::Unchanged().CopyToTVMFFIAny();
+ return ffi::Unchanged();
}
-TVMFFIAny ShuffleVisit(ffi::StructuralVisitorObj* visitor, ffi::AnyView value) noexcept {
+TVM_FFI_INLINE ffi::Expected<ffi::Optional<ffi::VisitInterrupt>> ShuffleVisit(
+ ffi::StructuralVisitorObj* visitor, ffi::AnyView value) noexcept {
// skips: PrimExpr types are always PrimType and remain unchanged in normal mutation.
const ShuffleNode* self =
ffi::details::AnyUnsafe::RawObjectPtrFromAnyViewAfterCheck<const ShuffleNode>(value);
TVM_FFI_S_VISIT_MAYBE_EARLY_RETURN(visitor->VisitExpected(self->vectors));
TVM_FFI_S_VISIT_MAYBE_EARLY_RETURN(visitor->VisitExpected(self->indices));
- return ffi::AnyView(nullptr).CopyToTVMFFIAny();
+ return std::nullopt;
}
-TVMFFIAny ShuffleMutate(ffi::StructuralMutatorObj* mutator, ffi::AnyView value) noexcept {
+TVM_FFI_INLINE ffi::Expected<ffi::UnchangedOr<ffi::Any>> ShuffleMutate(
+ ffi::StructuralMutatorObj* mutator, ffi::AnyView value) noexcept {
// skips: PrimExpr types are always PrimType and remain unchanged in normal mutation.
const ShuffleNode* self =
ffi::details::AnyUnsafe::RawObjectPtrFromAnyViewAfterCheck<const ShuffleNode>(value);
@@ -171,16 +178,16 @@ TVMFFIAny ShuffleMutate(ffi::StructuralMutatorObj* mutator, ffi::AnyView value)
mutator->MutateExpected(self->indices));
if (mapped_vectors_u.UnchangedOrSameAs(self->vectors) &&
mapped_indices_u.UnchangedOrSameAs(self->indices)) {
- return ffi::Unchanged().CopyToTVMFFIAny();
+ return ffi::Unchanged();
}
ffi::ObjectPtr<ShuffleNode> copy = ffi::make_object<ShuffleNode>(*self);
if (!mapped_vectors_u.IsUnchanged()) copy->vectors = std::move(mapped_vectors_u).ValueUnchecked();
if (!mapped_indices_u.IsUnchanged()) copy->indices = std::move(mapped_indices_u).ValueUnchecked();
- return ffi::details::AnyUnsafe::MoveAnyToTVMFFIAny(ffi::Any(std::move(copy)));
+ return ffi::Any(std::move(copy));
}
-TVMFFIAny ShuffleMaybeInplaceMutate(ffi::StructuralMutatorObj* mutator,
- ffi::AnyView value) noexcept {
+TVM_FFI_INLINE ffi::Expected<ffi::UnchangedOr<ffi::Any>> ShuffleMaybeInplaceMutate(
+ ffi::StructuralMutatorObj* mutator, ffi::AnyView value) noexcept {
// skips: PrimExpr types are always PrimType and remain unchanged in normal mutation.
ShuffleNode* self = const_cast<ShuffleNode*>(
ffi::details::AnyUnsafe::RawObjectPtrFromAnyViewAfterCheck<const ShuffleNode>(value));
@@ -192,11 +199,11 @@ TVMFFIAny ShuffleMaybeInplaceMutate(ffi::StructuralMutatorObj* mutator,
mutator->MutateExpected(self->indices, ffi::InplaceMode::kAllow));
if (mapped_vectors_u.UnchangedOrSameAs(self->vectors) &&
mapped_indices_u.UnchangedOrSameAs(self->indices)) {
- return ffi::Unchanged().CopyToTVMFFIAny();
+ return ffi::Unchanged();
}
if (!mapped_vectors_u.IsUnchanged()) self->vectors = std::move(mapped_vectors_u).ValueUnchecked();
if (!mapped_indices_u.IsUnchanged()) self->indices = std::move(mapped_indices_u).ValueUnchecked();
- return ffi::Unchanged().CopyToTVMFFIAny();
+ return ffi::Unchanged();
}
} // namespace
@@ -240,10 +247,10 @@ TVM_FFI_STATIC_INIT_BLOCK() {
namespace refl = tvm::ffi::reflection;
RampNode::RegisterReflection();
refl::TypeAttrDef<RampNode>()
- .attr(refl::type_attr::kStructuralVisit, reinterpret_cast<void*>(&RampVisit))
- .attr(refl::type_attr::kStructuralMutate, reinterpret_cast<void*>(&RampMutate))
+ .attr(refl::type_attr::kStructuralVisit, ffi::FStructuralVisit::FromNative<&RampVisit>())
+ .attr(refl::type_attr::kStructuralMutate, ffi::FStructuralMutate::FromNative<&RampMutate>())
.attr(refl::type_attr::kStructuralMaybeInplaceMutate,
- reinterpret_cast<void*>(&RampMaybeInplaceMutate));
+ ffi::FStructuralMutate::FromNative<&RampMaybeInplaceMutate>());
refl::GlobalDef().def("prim.Ramp", [](PrimExpr base, PrimExpr stride, PrimExpr lanes, Span span) {
return Ramp(base, stride, lanes, span);
@@ -283,10 +290,11 @@ TVM_FFI_STATIC_INIT_BLOCK() {
namespace refl = tvm::ffi::reflection;
BroadcastNode::RegisterReflection();
refl::TypeAttrDef<BroadcastNode>()
- .attr(refl::type_attr::kStructuralVisit, reinterpret_cast<void*>(&BroadcastVisit))
- .attr(refl::type_attr::kStructuralMutate, reinterpret_cast<void*>(&BroadcastMutate))
+ .attr(refl::type_attr::kStructuralVisit, ffi::FStructuralVisit::FromNative<&BroadcastVisit>())
+ .attr(refl::type_attr::kStructuralMutate,
+ ffi::FStructuralMutate::FromNative<&BroadcastMutate>())
.attr(refl::type_attr::kStructuralMaybeInplaceMutate,
- reinterpret_cast<void*>(&BroadcastMaybeInplaceMutate));
+ ffi::FStructuralMutate::FromNative<&BroadcastMaybeInplaceMutate>());
refl::GlobalDef().def("prim.Broadcast", [](PrimExpr value, PrimExpr lanes, Span span) {
return Broadcast(value, lanes, span);
@@ -320,10 +328,11 @@ TVM_FFI_STATIC_INIT_BLOCK() {
namespace refl = tvm::ffi::reflection;
ShuffleNode::RegisterReflection();
refl::TypeAttrDef<ShuffleNode>()
- .attr(refl::type_attr::kStructuralVisit, reinterpret_cast<void*>(&ShuffleVisit))
- .attr(refl::type_attr::kStructuralMutate, reinterpret_cast<void*>(&ShuffleMutate))
+ .attr(refl::type_attr::kStructuralVisit, ffi::FStructuralVisit::FromNative<&ShuffleVisit>())
+ .attr(refl::type_attr::kStructuralMutate,
+ ffi::FStructuralMutate::FromNative<&ShuffleMutate>())
.attr(refl::type_attr::kStructuralMaybeInplaceMutate,
- reinterpret_cast<void*>(&ShuffleMaybeInplaceMutate));
+ ffi::FStructuralMutate::FromNative<&ShuffleMaybeInplaceMutate>());
refl::GlobalDef().def("prim.Shuffle",
[](ffi::Array<PrimExpr> vectors, ffi::Array<PrimExpr> indices, Span span) {
diff --git a/src/ir/type.cc b/src/ir/type.cc
index ba5996858d..e824b4a380 100644
--- a/src/ir/type.cc
+++ b/src/ir/type.cc
@@ -61,30 +61,36 @@ ffi::ObjectPtr<PrimTypeNode> GetCachedPrimTypeNode(DLDataType dtype) {
// Structural traversal hooks
-TVMFFIAny TypeVisit(ffi::StructuralVisitorObj*, ffi::AnyView) noexcept {
+TVM_FFI_INLINE ffi::Expected<ffi::Optional<ffi::VisitInterrupt>> TypeVisit(
+ ffi::StructuralVisitorObj*, ffi::AnyView) noexcept {
// Type::Missing() is the only concrete TypeNode value; span is ignored debug metadata.
- return ffi::AnyView(nullptr).CopyToTVMFFIAny();
+ return std::nullopt;
}
-TVMFFIAny TypeMutate(ffi::StructuralMutatorObj*, ffi::AnyView) noexcept {
- return ffi::Unchanged().CopyToTVMFFIAny();
+TVM_FFI_INLINE ffi::Expected<ffi::UnchangedOr<ffi::Any>> TypeMutate(ffi::StructuralMutatorObj*,
+ ffi::AnyView) noexcept {
+ return ffi::Unchanged();
}
-TVMFFIAny TypeMaybeInplaceMutate(ffi::StructuralMutatorObj*, ffi::AnyView) noexcept {
- return ffi::Unchanged().CopyToTVMFFIAny();
+TVM_FFI_INLINE ffi::Expected<ffi::UnchangedOr<ffi::Any>> TypeMaybeInplaceMutate(
+ ffi::StructuralMutatorObj*, ffi::AnyView) noexcept {
+ return ffi::Unchanged();
}
-TVMFFIAny OpaqueTypeVisit(ffi::StructuralVisitorObj*, ffi::AnyView) noexcept {
+TVM_FFI_INLINE ffi::Expected<ffi::Optional<ffi::VisitInterrupt>> OpaqueTypeVisit(
+ ffi::StructuralVisitorObj*, ffi::AnyView) noexcept {
// OpaqueType is a field-less construction-time marker; span is ignored debug metadata.
- return ffi::AnyView(nullptr).CopyToTVMFFIAny();
+ return std::nullopt;
}
-TVMFFIAny OpaqueTypeMutate(ffi::StructuralMutatorObj*, ffi::AnyView) noexcept {
- return ffi::Unchanged().CopyToTVMFFIAny();
+TVM_FFI_INLINE ffi::Expected<ffi::UnchangedOr<ffi::Any>> OpaqueTypeMutate(
+ ffi::StructuralMutatorObj*, ffi::AnyView) noexcept {
+ return ffi::Unchanged();
}
-TVMFFIAny OpaqueTypeMaybeInplaceMutate(ffi::StructuralMutatorObj*, ffi::AnyView) noexcept {
- return ffi::Unchanged().CopyToTVMFFIAny();
+TVM_FFI_INLINE ffi::Expected<ffi::UnchangedOr<ffi::Any>> OpaqueTypeMaybeInplaceMutate(
+ ffi::StructuralMutatorObj*, ffi::AnyView) noexcept {
+ return ffi::Unchanged();
}
int64_t PrimTypeAnyHash(const ffi::Any& src) {
@@ -95,49 +101,54 @@ bool PrimTypeAnyEqual(const ffi::Any& lhs, const ffi::Any& rhs) {
return lhs.cast<PrimType>()->dtype == rhs.cast<PrimType>()->dtype;
}
-TVMFFIAny PrimTypeVisit(ffi::StructuralVisitorObj*, ffi::AnyView) noexcept {
+TVM_FFI_INLINE ffi::Expected<ffi::Optional<ffi::VisitInterrupt>> PrimTypeVisit(
+ ffi::StructuralVisitorObj*, ffi::AnyView) noexcept {
// dtype is a constant: reflected for StructuralEqual/Hash,
// not traversed by the visitor/mutator contract.
- return ffi::AnyView(nullptr).CopyToTVMFFIAny();
+ return std::nullopt;
}
-TVMFFIAny PrimTypeMutate(ffi::StructuralMutatorObj*, ffi::AnyView) noexcept {
+TVM_FFI_INLINE ffi::Expected<ffi::UnchangedOr<ffi::Any>> PrimTypeMutate(ffi::StructuralMutatorObj*,
+ ffi::AnyView) noexcept {
// dtype is a constant: reflected for StructuralEqual/Hash,
// not traversed by the visitor/mutator contract.
- return ffi::Unchanged().CopyToTVMFFIAny();
+ return ffi::Unchanged();
}
-TVMFFIAny PrimTypeMaybeInplaceMutate(ffi::StructuralMutatorObj*, ffi::AnyView) noexcept {
+TVM_FFI_INLINE ffi::Expected<ffi::UnchangedOr<ffi::Any>> PrimTypeMaybeInplaceMutate(
+ ffi::StructuralMutatorObj*, ffi::AnyView) noexcept {
// dtype is a constant: reflected for StructuralEqual/Hash,
// not traversed by the visitor/mutator contract.
- return ffi::Unchanged().CopyToTVMFFIAny();
+ return ffi::Unchanged();
}
-TVMFFIAny PointerTypeVisit(ffi::StructuralVisitorObj* visitor, ffi::AnyView value) noexcept {
+TVM_FFI_INLINE ffi::Expected<ffi::Optional<ffi::VisitInterrupt>> PointerTypeVisit(
+ ffi::StructuralVisitorObj* visitor, ffi::AnyView value) noexcept {
// skips: storage_scope (scalar)
const PointerTypeNode* self =
ffi::details::AnyUnsafe::RawObjectPtrFromAnyViewAfterCheck<const PointerTypeNode>(value);
TVM_FFI_S_VISIT_MAYBE_EARLY_RETURN(visitor->VisitExpected(self->element_type));
- return ffi::AnyView(nullptr).CopyToTVMFFIAny();
+ return std::nullopt;
}
-TVMFFIAny PointerTypeMutate(ffi::StructuralMutatorObj* mutator, ffi::AnyView value) noexcept {
+TVM_FFI_INLINE ffi::Expected<ffi::UnchangedOr<ffi::Any>> PointerTypeMutate(
+ ffi::StructuralMutatorObj* mutator, ffi::AnyView value) noexcept {
// skips: storage_scope (scalar)
const PointerTypeNode* self =
ffi::details::AnyUnsafe::RawObjectPtrFromAnyViewAfterCheck<const PointerTypeNode>(value);
TVM_FFI_S_MUTATE_ASSIGN_OR_RETURN(ffi::UnchangedOr<Type>, mapped_element_type_u,
mutator->MutateExpected(self->element_type));
if (mapped_element_type_u.UnchangedOrSameAs(self->element_type)) {
- return ffi::Unchanged().CopyToTVMFFIAny();
+ return ffi::Unchanged();
}
ffi::ObjectPtr<PointerTypeNode> copy = ffi::make_object<PointerTypeNode>(*self);
if (!mapped_element_type_u.IsUnchanged())
copy->element_type = std::move(mapped_element_type_u).ValueUnchecked();
- return ffi::details::AnyUnsafe::MoveAnyToTVMFFIAny(ffi::Any(std::move(copy)));
+ return ffi::Any(std::move(copy));
}
-TVMFFIAny PointerTypeMaybeInplaceMutate(ffi::StructuralMutatorObj* mutator,
- ffi::AnyView value) noexcept {
+TVM_FFI_INLINE ffi::Expected<ffi::UnchangedOr<ffi::Any>> PointerTypeMaybeInplaceMutate(
+ ffi::StructuralMutatorObj* mutator, ffi::AnyView value) noexcept {
// skips: storage_scope (scalar)
PointerTypeNode* self = const_cast<PointerTypeNode*>(
ffi::details::AnyUnsafe::RawObjectPtrFromAnyViewAfterCheck<const PointerTypeNode>(value));
@@ -147,18 +158,20 @@ TVMFFIAny PointerTypeMaybeInplaceMutate(ffi::StructuralMutatorObj* mutator,
if (!mapped_element_type_u.UnchangedOrSameAs(self->element_type)) {
self->element_type = std::move(mapped_element_type_u).ValueUnchecked();
}
- return ffi::Unchanged().CopyToTVMFFIAny();
+ return ffi::Unchanged();
}
-TVMFFIAny FuncTypeVisit(ffi::StructuralVisitorObj* visitor, ffi::AnyView value) noexcept {
+TVM_FFI_INLINE ffi::Expected<ffi::Optional<ffi::VisitInterrupt>> FuncTypeVisit(
+ ffi::StructuralVisitorObj* visitor, ffi::AnyView value) noexcept {
const FuncTypeNode* self =
ffi::details::AnyUnsafe::RawObjectPtrFromAnyViewAfterCheck<const FuncTypeNode>(value);
TVM_FFI_S_VISIT_MAYBE_EARLY_RETURN(visitor->VisitExpected(self->arg_types));
TVM_FFI_S_VISIT_MAYBE_EARLY_RETURN(visitor->VisitExpected(self->ret_type));
- return ffi::AnyView(nullptr).CopyToTVMFFIAny();
+ return std::nullopt;
}
-TVMFFIAny FuncTypeMutate(ffi::StructuralMutatorObj* mutator, ffi::AnyView value) noexcept {
+TVM_FFI_INLINE ffi::Expected<ffi::UnchangedOr<ffi::Any>> FuncTypeMutate(
+ ffi::StructuralMutatorObj* mutator, ffi::AnyView value) noexcept {
const FuncTypeNode* self =
ffi::details::AnyUnsafe::RawObjectPtrFromAnyViewAfterCheck<const FuncTypeNode>(value);
TVM_FFI_S_MUTATE_ASSIGN_OR_RETURN(ffi::UnchangedOr<ffi::Array<Type>>, mapped_arg_types_u,
@@ -167,18 +180,18 @@ TVMFFIAny FuncTypeMutate(ffi::StructuralMutatorObj* mutator, ffi::AnyView value)
mutator->MutateExpected(self->ret_type));
if (mapped_arg_types_u.UnchangedOrSameAs(self->arg_types) &&
mapped_ret_type_u.UnchangedOrSameAs(self->ret_type)) {
- return ffi::Unchanged().CopyToTVMFFIAny();
+ return ffi::Unchanged();
}
ffi::ObjectPtr<FuncTypeNode> copy = ffi::make_object<FuncTypeNode>(*self);
if (!mapped_arg_types_u.IsUnchanged())
copy->arg_types = std::move(mapped_arg_types_u).ValueUnchecked();
if (!mapped_ret_type_u.IsUnchanged())
copy->ret_type = std::move(mapped_ret_type_u).ValueUnchecked();
- return ffi::details::AnyUnsafe::MoveAnyToTVMFFIAny(ffi::Any(std::move(copy)));
+ return ffi::Any(std::move(copy));
}
-TVMFFIAny FuncTypeMaybeInplaceMutate(ffi::StructuralMutatorObj* mutator,
- ffi::AnyView value) noexcept {
+TVM_FFI_INLINE ffi::Expected<ffi::UnchangedOr<ffi::Any>> FuncTypeMaybeInplaceMutate(
+ ffi::StructuralMutatorObj* mutator, ffi::AnyView value) noexcept {
FuncTypeNode* self = const_cast<FuncTypeNode*>(
ffi::details::AnyUnsafe::RawObjectPtrFromAnyViewAfterCheck<const FuncTypeNode>(value));
TVM_FFI_S_MUTATE_ASSIGN_OR_RETURN(
@@ -193,31 +206,33 @@ TVMFFIAny FuncTypeMaybeInplaceMutate(ffi::StructuralMutatorObj* mutator,
if (!mapped_ret_type_u.UnchangedOrSameAs(self->ret_type)) {
self->ret_type = std::move(mapped_ret_type_u).ValueUnchecked();
}
- return ffi::Unchanged().CopyToTVMFFIAny();
+ return ffi::Unchanged();
}
-TVMFFIAny TupleTypeVisit(ffi::StructuralVisitorObj* visitor, ffi::AnyView value) noexcept {
+TVM_FFI_INLINE ffi::Expected<ffi::Optional<ffi::VisitInterrupt>> TupleTypeVisit(
+ ffi::StructuralVisitorObj* visitor, ffi::AnyView value) noexcept {
const TupleTypeNode* self =
ffi::details::AnyUnsafe::RawObjectPtrFromAnyViewAfterCheck<const TupleTypeNode>(value);
TVM_FFI_S_VISIT_MAYBE_EARLY_RETURN(visitor->VisitExpected(self->fields));
- return ffi::AnyView(nullptr).CopyToTVMFFIAny();
+ return std::nullopt;
}
-TVMFFIAny TupleTypeMutate(ffi::StructuralMutatorObj* mutator, ffi::AnyView value) noexcept {
+TVM_FFI_INLINE ffi::Expected<ffi::UnchangedOr<ffi::Any>> TupleTypeMutate(
+ ffi::StructuralMutatorObj* mutator, ffi::AnyView value) noexcept {
const TupleTypeNode* self =
ffi::details::AnyUnsafe::RawObjectPtrFromAnyViewAfterCheck<const TupleTypeNode>(value);
TVM_FFI_S_MUTATE_ASSIGN_OR_RETURN(ffi::UnchangedOr<ffi::Array<Type>>, mapped_fields_u,
mutator->MutateExpected(self->fields));
if (mapped_fields_u.UnchangedOrSameAs(self->fields)) {
- return ffi::Unchanged().CopyToTVMFFIAny();
+ return ffi::Unchanged();
}
ffi::ObjectPtr<TupleTypeNode> copy = ffi::make_object<TupleTypeNode>(*self);
if (!mapped_fields_u.IsUnchanged()) copy->fields = std::move(mapped_fields_u).ValueUnchecked();
- return ffi::details::AnyUnsafe::MoveAnyToTVMFFIAny(ffi::Any(std::move(copy)));
+ return ffi::Any(std::move(copy));
}
-TVMFFIAny TupleTypeMaybeInplaceMutate(ffi::StructuralMutatorObj* mutator,
- ffi::AnyView value) noexcept {
+TVM_FFI_INLINE ffi::Expected<ffi::UnchangedOr<ffi::Any>> TupleTypeMaybeInplaceMutate(
+ ffi::StructuralMutatorObj* mutator, ffi::AnyView value) noexcept {
TupleTypeNode* self = const_cast<TupleTypeNode*>(
ffi::details::AnyUnsafe::RawObjectPtrFromAnyViewAfterCheck<const TupleTypeNode>(value));
TVM_FFI_S_MUTATE_ASSIGN_OR_RETURN(
@@ -226,7 +241,7 @@ TVMFFIAny TupleTypeMaybeInplaceMutate(ffi::StructuralMutatorObj* mutator,
if (!mapped_fields_u.UnchangedOrSameAs(self->fields)) {
self->fields = std::move(mapped_fields_u).ValueUnchecked();
}
- return ffi::Unchanged().CopyToTVMFFIAny();
+ return ffi::Unchanged();
}
} // namespace
@@ -244,10 +259,10 @@ TVM_FFI_STATIC_INIT_BLOCK() {
namespace refl = tvm::ffi::reflection;
TypeNode::RegisterReflection();
refl::TypeAttrDef<TypeNode>()
- .attr(refl::type_attr::kStructuralVisit, reinterpret_cast<void*>(&TypeVisit))
- .attr(refl::type_attr::kStructuralMutate, reinterpret_cast<void*>(&TypeMutate))
+ .attr(refl::type_attr::kStructuralVisit, ffi::FStructuralVisit::FromNative<&TypeVisit>())
+ .attr(refl::type_attr::kStructuralMutate, ffi::FStructuralMutate::FromNative<&TypeMutate>())
.attr(refl::type_attr::kStructuralMaybeInplaceMutate,
- reinterpret_cast<void*>(&TypeMaybeInplaceMutate));
+ ffi::FStructuralMutate::FromNative<&TypeMaybeInplaceMutate>());
refl::GlobalDef()
.def("ir.TypeMissing", []() { return Type::Missing(); })
.def("ir.TypeIsMissing", [](Type type) { return type.IsMissing(); });
@@ -265,10 +280,10 @@ TVM_FFI_STATIC_INIT_BLOCK() {
namespace refl = tvm::ffi::reflection;
AnyTypeNode::RegisterReflection();
refl::TypeAttrDef<AnyTypeNode>()
- .attr(refl::type_attr::kStructuralVisit, reinterpret_cast<void*>(&TypeVisit))
- .attr(refl::type_attr::kStructuralMutate, reinterpret_cast<void*>(&TypeMutate))
+ .attr(refl::type_attr::kStructuralVisit, ffi::FStructuralVisit::FromNative<&TypeVisit>())
+ .attr(refl::type_attr::kStructuralMutate, ffi::FStructuralMutate::FromNative<&TypeMutate>())
.attr(refl::type_attr::kStructuralMaybeInplaceMutate,
- reinterpret_cast<void*>(&TypeMaybeInplaceMutate));
+ ffi::FStructuralMutate::FromNative<&TypeMaybeInplaceMutate>());
refl::GlobalDef().def("ir.AnyType", [](Span span) { return AnyType(span); });
}
@@ -278,10 +293,12 @@ TVM_FFI_STATIC_INIT_BLOCK() {
namespace refl = tvm::ffi::reflection;
OpaqueTypeNode::RegisterReflection();
refl::TypeAttrDef<OpaqueTypeNode>()
- .attr(refl::type_attr::kStructuralVisit, reinterpret_cast<void*>(&OpaqueTypeVisit))
- .attr(refl::type_attr::kStructuralMutate, reinterpret_cast<void*>(&OpaqueTypeMutate))
+ .attr(refl::type_attr::kStructuralVisit,
+ ffi::FStructuralVisit::FromNative<&OpaqueTypeVisit>())
+ .attr(refl::type_attr::kStructuralMutate,
+ ffi::FStructuralMutate::FromNative<&OpaqueTypeMutate>())
.attr(refl::type_attr::kStructuralMaybeInplaceMutate,
- reinterpret_cast<void*>(&OpaqueTypeMaybeInplaceMutate));
+ ffi::FStructuralMutate::FromNative<&OpaqueTypeMaybeInplaceMutate>());
refl::GlobalDef().def("ir.OpaqueType", []() { return OpaqueType(); });
}
@@ -305,10 +322,11 @@ TVM_FFI_STATIC_INIT_BLOCK() {
refl::TypeAttrDef<PrimTypeNode>()
.attr(refl::type_attr::kAnyHash, reinterpret_cast<void*>(&PrimTypeAnyHash))
.attr(refl::type_attr::kAnyEqual, reinterpret_cast<void*>(&PrimTypeAnyEqual))
- .attr(refl::type_attr::kStructuralVisit, reinterpret_cast<void*>(&PrimTypeVisit))
- .attr(refl::type_attr::kStructuralMutate, reinterpret_cast<void*>(&PrimTypeMutate))
+ .attr(refl::type_attr::kStructuralVisit, ffi::FStructuralVisit::FromNative<&PrimTypeVisit>())
+ .attr(refl::type_attr::kStructuralMutate,
+ ffi::FStructuralMutate::FromNative<&PrimTypeMutate>())
.attr(refl::type_attr::kStructuralMaybeInplaceMutate,
- reinterpret_cast<void*>(&PrimTypeMaybeInplaceMutate));
+ ffi::FStructuralMutate::FromNative<&PrimTypeMaybeInplaceMutate>());
refl::GlobalDef().def("ir.PrimType", [](DLDataType dtype) { return PrimType(dtype); });
}
@@ -363,10 +381,10 @@ TVM_FFI_STATIC_INIT_BLOCK() {
namespace refl = tvm::ffi::reflection;
StringTypeNode::RegisterReflection();
refl::TypeAttrDef<StringTypeNode>()
- .attr(refl::type_attr::kStructuralVisit, reinterpret_cast<void*>(&TypeVisit))
- .attr(refl::type_attr::kStructuralMutate, reinterpret_cast<void*>(&TypeMutate))
+ .attr(refl::type_attr::kStructuralVisit, ffi::FStructuralVisit::FromNative<&TypeVisit>())
+ .attr(refl::type_attr::kStructuralMutate, ffi::FStructuralMutate::FromNative<&TypeMutate>())
.attr(refl::type_attr::kStructuralMaybeInplaceMutate,
- reinterpret_cast<void*>(&TypeMaybeInplaceMutate));
+ ffi::FStructuralMutate::FromNative<&TypeMaybeInplaceMutate>());
refl::GlobalDef().def("ir.StringType", []() { return StringType(); });
}
@@ -387,10 +405,12 @@ TVM_FFI_STATIC_INIT_BLOCK() {
namespace refl = tvm::ffi::reflection;
PointerTypeNode::RegisterReflection();
refl::TypeAttrDef<PointerTypeNode>()
- .attr(refl::type_attr::kStructuralVisit, reinterpret_cast<void*>(&PointerTypeVisit))
- .attr(refl::type_attr::kStructuralMutate, reinterpret_cast<void*>(&PointerTypeMutate))
+ .attr(refl::type_attr::kStructuralVisit,
+ ffi::FStructuralVisit::FromNative<&PointerTypeVisit>())
+ .attr(refl::type_attr::kStructuralMutate,
+ ffi::FStructuralMutate::FromNative<&PointerTypeMutate>())
.attr(refl::type_attr::kStructuralMaybeInplaceMutate,
- reinterpret_cast<void*>(&PointerTypeMaybeInplaceMutate));
+ ffi::FStructuralMutate::FromNative<&PointerTypeMaybeInplaceMutate>());
refl::GlobalDef().def("ir.PointerType", [](Type element_type, ffi::String storage_scope = "") {
return PointerType(element_type, storage_scope);
@@ -414,10 +434,11 @@ TVM_FFI_STATIC_INIT_BLOCK() {
namespace refl = tvm::ffi::reflection;
FuncTypeNode::RegisterReflection();
refl::TypeAttrDef<FuncTypeNode>()
- .attr(refl::type_attr::kStructuralVisit, reinterpret_cast<void*>(&FuncTypeVisit))
- .attr(refl::type_attr::kStructuralMutate, reinterpret_cast<void*>(&FuncTypeMutate))
+ .attr(refl::type_attr::kStructuralVisit, ffi::FStructuralVisit::FromNative<&FuncTypeVisit>())
+ .attr(refl::type_attr::kStructuralMutate,
+ ffi::FStructuralMutate::FromNative<&FuncTypeMutate>())
.attr(refl::type_attr::kStructuralMaybeInplaceMutate,
- reinterpret_cast<void*>(&FuncTypeMaybeInplaceMutate));
+ ffi::FStructuralMutate::FromNative<&FuncTypeMaybeInplaceMutate>());
refl::GlobalDef().def("ir.FuncType", [](tvm::ffi::Array<Type> arg_types, Type ret_type) {
return FuncType(arg_types, ret_type);
@@ -435,10 +456,11 @@ TVM_FFI_STATIC_INIT_BLOCK() {
namespace refl = tvm::ffi::reflection;
TupleTypeNode::RegisterReflection();
refl::TypeAttrDef<TupleTypeNode>()
- .attr(refl::type_attr::kStructuralVisit, reinterpret_cast<void*>(&TupleTypeVisit))
- .attr(refl::type_attr::kStructuralMutate, reinterpret_cast<void*>(&TupleTypeMutate))
+ .attr(refl::type_attr::kStructuralVisit, ffi::FStructuralVisit::FromNative<&TupleTypeVisit>())
+ .attr(refl::type_attr::kStructuralMutate,
+ ffi::FStructuralMutate::FromNative<&TupleTypeMutate>())
.attr(refl::type_attr::kStructuralMaybeInplaceMutate,
- reinterpret_cast<void*>(&TupleTypeMaybeInplaceMutate));
+ ffi::FStructuralMutate::FromNative<&TupleTypeMaybeInplaceMutate>());
refl::GlobalDef().def("ir.TupleType",
[](ffi::Array<Type> fields, Span span) { return TupleType(fields, span); });
diff --git a/src/relax/analysis/well_formed.cc b/src/relax/analysis/well_formed.cc
index ee81e4bafb..6e1965b85a 100644
--- a/src/relax/analysis/well_formed.cc
+++ b/src/relax/analysis/well_formed.cc
@@ -425,7 +425,8 @@ class WellFormedChecker : public relax::ExprVisitor, public relax::TypeVisitor {
bool has_infer_type = true;
if (auto op = call->op.as<Op>()) {
- has_infer_type = op_map_infer_type_.count(op.value());
+ has_infer_type =
+ op_map_infer_type_.count(op.value()) || op_map_infer_type_with_builder_.count(op.value());
}
if (check_ty && !call->ty.IsMissing() && (!call->ty.as<PrimTypeNode>() || has_infer_type)) {
// The `InferType` method isn't currently exposed by the
@@ -715,6 +716,8 @@ class WellFormedChecker : public relax::ExprVisitor, public relax::TypeVisitor {
tvm::OpAttrMap<FNormalize> op_map_normalize_ = Op::GetAttrMap<FNormalize>("FNormalize");
tvm::OpAttrMap<FInferType> op_map_infer_type_ = Op::GetAttrMap<FInferType>("FInferType");
+ tvm::OpAttrMap<FInferTypeWithBuilder> op_map_infer_type_with_builder_ =
+ Op::GetAttrMap<FInferTypeWithBuilder>("relax.FInferTypeWithBuilder");
};
void WellFormed(ffi::Variant<IRModule, Function> obj, bool check_ty) {
diff --git a/src/relax/backend/adreno/annotate_custom_storage.cc b/src/relax/backend/adreno/annotate_custom_storage.cc
index 1e7d9b6966..e71e9d38d8 100644
--- a/src/relax/backend/adreno/annotate_custom_storage.cc
+++ b/src/relax/backend/adreno/annotate_custom_storage.cc
@@ -122,12 +122,12 @@
* 3: DefineVDevice: Pass does injects hint_on_device for each argument. It also tries to update
* out Type containing VDevice information. This update for tirx calls is straight forward
* as ty_args in CallNode is meant for this purpose. This ty_args for other calls by
- * design is invalid as we do this by "FInferType".
- * Another issue we have with "FInferType" per op is they can't decide this
+ * design is invalid as we do this by per-op type inference.
+ * Another issue with type inference per op is that it cannot decide this
* memory scope information which is done by this pass based on consumer demand.
* Hence, we are going to use the ty_args to indicate this information.
- * So, this pass attributes ty_args for regumar calls too and FInferType implmentation
- * do take VDevice information fro this hint. This also solves the issue of mixed VDevice
+ * So, this pass attributes ty_args for regular calls too, and type inference implementations
+ * take VDevice information from this hint. This also solves the issue of mixed VDevice
* for arguments of an op.
* After these steps the mod looks like
*
@@ -501,13 +501,18 @@ class CollectProducerScopeInfo : public ExprVisitor {
if (call->op.same_as(call_tir_op)) {
out_ty = call->ty_args[0];
} else {
- tvm::OpAttrMap<FInferType> op_map_infer_ty = Op::GetAttrMap<FInferType>("FInferType");
-
auto* op_ptr = call->op.as<OpNode>();
Op op = ffi::GetRef<Op>(op_ptr);
- TVM_FFI_ICHECK(op_map_infer_ty.count(op))
- << " Cannot find the FInferType attribute registered to op: " << op->name;
- out_ty = op_map_infer_ty[op](ffi::GetRef<Call>(call), builder_);
+ static auto op_map_context_free = Op::GetAttrMap<FInferType>("FInferType");
+ if (op_map_context_free.count(op)) {
+ out_ty = Call::ReinferType(call);
+ } else {
+ static auto op_map_infer_ty =
+ Op::GetAttrMap<FInferTypeWithBuilder>("relax.FInferTypeWithBuilder");
+ TVM_FFI_ICHECK(op_map_infer_ty.count(op))
+ << " Cannot find a type inference attribute registered to op: " << op->name;
+ out_ty = op_map_infer_ty[op](ffi::GetRef<Call>(call), builder_);
+ }
}
std::unordered_map<ffi::String, int> scope_count;
diff --git a/src/relax/backend/vm/vm_shape_lower.cc b/src/relax/backend/vm/vm_shape_lower.cc
index 8f1efc3b7f..47ce4ce39c 100644
--- a/src/relax/backend/vm/vm_shape_lower.cc
+++ b/src/relax/backend/vm/vm_shape_lower.cc
@@ -84,8 +84,11 @@ static bool IsRelaxOwnedCall(const CallNode* call) {
auto op = call->op.as<Op>();
if (!op) return true;
static auto infer_type_map = Op::GetAttrMap<FInferType>("FInferType");
+ static auto infer_type_with_builder_map =
+ Op::GetAttrMap<FInferTypeWithBuilder>("relax.FInferTypeWithBuilder");
static auto legalize_map = Op::GetAttrMap<FLegalize>("FLegalize");
- return infer_type_map.count(op.value()) || legalize_map.count(op.value());
+ return infer_type_map.count(op.value()) || infer_type_with_builder_map.count(op.value()) ||
+ legalize_map.count(op.value());
}
// Collector to collect PrimExprSlotMap
diff --git a/src/relax/distributed/type.cc b/src/relax/distributed/type.cc
index 91a423b886..78b4310fdb 100644
--- a/src/relax/distributed/type.cc
+++ b/src/relax/distributed/type.cc
@@ -32,16 +32,18 @@ namespace distributed {
namespace {
-TVMFFIAny DTensorTypeVisit(ffi::StructuralVisitorObj* visitor, ffi::AnyView value) noexcept {
+TVM_FFI_INLINE ffi::Expected<ffi::Optional<ffi::VisitInterrupt>> DTensorTypeVisit(
+ ffi::StructuralVisitorObj* visitor, ffi::AnyView value) noexcept {
const DTensorTypeNode* self =
ffi::details::AnyUnsafe::RawObjectPtrFromAnyViewAfterCheck<const DTensorTypeNode>(value);
TVM_FFI_S_VISIT_MAYBE_EARLY_RETURN(visitor->VisitExpected(self->device_mesh));
TVM_FFI_S_VISIT_MAYBE_EARLY_RETURN(visitor->VisitExpected(self->placement));
TVM_FFI_S_VISIT_MAYBE_EARLY_RETURN(visitor->VisitExpected(self->tensor_ty));
- return ffi::AnyView(nullptr).CopyToTVMFFIAny();
+ return std::nullopt;
}
-TVMFFIAny DTensorTypeMutate(ffi::StructuralMutatorObj* mutator, ffi::AnyView value) noexcept {
+TVM_FFI_INLINE ffi::Expected<ffi::UnchangedOr<ffi::Any>> DTensorTypeMutate(
+ ffi::StructuralMutatorObj* mutator, ffi::AnyView value) noexcept {
const DTensorTypeNode* self =
ffi::details::AnyUnsafe::RawObjectPtrFromAnyViewAfterCheck<const DTensorTypeNode>(value);
TVM_FFI_S_MUTATE_ASSIGN_OR_RETURN(ffi::UnchangedOr<DeviceMesh>, mapped_device_mesh,
@@ -53,17 +55,17 @@ TVMFFIAny DTensorTypeMutate(ffi::StructuralMutatorObj* mutator, ffi::AnyView val
if (mapped_device_mesh.UnchangedOrSameAs(self->device_mesh) &&
mapped_placement.UnchangedOrSameAs(self->placement) &&
mapped_tensor_ty.UnchangedOrSameAs(self->tensor_ty)) {
- return ffi::Unchanged().CopyToTVMFFIAny();
+ return ffi::Unchanged();
}
ffi::ObjectPtr<DTensorTypeNode> copy = ffi::make_object<DTensorTypeNode>(*self);
copy->device_mesh = std::move(mapped_device_mesh).ValueOrUnchanged(std::move(copy->device_mesh));
copy->placement = std::move(mapped_placement).ValueOrUnchanged(std::move(copy->placement));
copy->tensor_ty = std::move(mapped_tensor_ty).ValueOrUnchanged(std::move(copy->tensor_ty));
- return ffi::details::AnyUnsafe::MoveAnyToTVMFFIAny(ffi::Any(std::move(copy)));
+ return ffi::Any(std::move(copy));
}
-TVMFFIAny DTensorTypeMaybeInplaceMutate(ffi::StructuralMutatorObj* mutator,
- ffi::AnyView value) noexcept {
+TVM_FFI_INLINE ffi::Expected<ffi::UnchangedOr<ffi::Any>> DTensorTypeMaybeInplaceMutate(
+ ffi::StructuralMutatorObj* mutator, ffi::AnyView value) noexcept {
DTensorTypeNode* self = const_cast<DTensorTypeNode*>(
ffi::details::AnyUnsafe::RawObjectPtrFromAnyViewAfterCheck<const DTensorTypeNode>(value));
TVM_FFI_S_MUTATE_ASSIGN_OR_RETURN(
@@ -84,7 +86,7 @@ TVMFFIAny DTensorTypeMaybeInplaceMutate(ffi::StructuralMutatorObj* mutator,
if (!mapped_tensor_ty.UnchangedOrSameAs(self->tensor_ty)) {
self->tensor_ty = std::move(mapped_tensor_ty).ValueUnchecked();
}
- return ffi::Unchanged().CopyToTVMFFIAny();
+ return ffi::Unchanged();
}
} // namespace
@@ -201,10 +203,12 @@ TVM_FFI_STATIC_INIT_BLOCK() {
namespace refl = tvm::ffi::reflection;
DTensorTypeNode::RegisterReflection();
refl::TypeAttrDef<DTensorTypeNode>()
- .attr(refl::type_attr::kStructuralVisit, reinterpret_cast<void*>(&DTensorTypeVisit))
- .attr(refl::type_attr::kStructuralMutate, reinterpret_cast<void*>(&DTensorTypeMutate))
+ .attr(refl::type_attr::kStructuralVisit,
+ ffi::FStructuralVisit::FromNative<&DTensorTypeVisit>())
+ .attr(refl::type_attr::kStructuralMutate,
+ ffi::FStructuralMutate::FromNative<&DTensorTypeMutate>())
.attr(refl::type_attr::kStructuralMaybeInplaceMutate,
- reinterpret_cast<void*>(&DTensorTypeMaybeInplaceMutate));
+ ffi::FStructuralMutate::FromNative<&DTensorTypeMaybeInplaceMutate>());
}
TVM_FFI_STATIC_INIT_BLOCK() {
diff --git a/src/relax/ir/block_builder.cc b/src/relax/ir/block_builder.cc
index ed219a36e8..06efd797a7 100644
--- a/src/relax/ir/block_builder.cc
+++ b/src/relax/ir/block_builder.cc
@@ -830,7 +830,7 @@ class Normalizer : public BlockBuilderImpl, private ExprFunctor<Expr(const Expr&
// Helper function to infer the type of a Call.
Type InferType(const Call& call) {
if (auto* op_ptr = call->op.as<OpNode>()) {
- // Case 1: the op field is a primitive op, look up FInferType attribute
+ // Case 1: the op field is a primitive op, look up FInferTypeWithBuilder attribute
Op op = ffi::GetRef<Op>(op_ptr);
bool is_dist_op = false;
for (const auto& arg : call->args) {
@@ -845,11 +845,20 @@ class Normalizer : public BlockBuilderImpl, private ExprFunctor<Expr(const Expr&
<< "Distributed operator must take DTensor instead of Tensor as input";
}
TVM_FFI_ICHECK(op_map_dist_infer_ty.count(op))
- << " Cannot find the dist.FInferType attribute registered to op: " << op->name;
+ << " Cannot find the dist.FInferTypeWithBuilder attribute registered to op: "
+ << op->name;
return op_map_dist_infer_ty[op](call, ffi::GetRef<BlockBuilder>(this));
}
- TVM_FFI_ICHECK(op_map_infer_ty.count(op))
- << " Cannot find the FInferType attribute registered to op: " << op->name;
+ bool has_context_free = op_map_context_free_infer_ty.count(op);
+ bool has_contextual = op_map_infer_ty.count(op);
+ TVM_FFI_ICHECK(!(has_context_free && has_contextual))
+ << "Operator " << op->name << " registers both FInferType and FInferTypeWithBuilder";
+ if (has_context_free) {
+ return Call::ReinferType(call.get());
+ }
+ TVM_FFI_ICHECK(has_contextual)
+ << " Cannot find the FInferType or FInferTypeWithBuilder attribute registered to op: "
+ << op->name;
return op_map_infer_ty[op](call, ffi::GetRef<BlockBuilder>(this));
} else {
// derive using function parameters
@@ -1034,8 +1043,12 @@ class Normalizer : public BlockBuilderImpl, private ExprFunctor<Expr(const Expr&
}
/*! \brief Operator type inference map. */
- tvm::OpAttrMap<FInferType> op_map_infer_ty = Op::GetAttrMap<FInferType>("FInferType");
- tvm::OpAttrMap<FInferType> op_map_dist_infer_ty = Op::GetAttrMap<FInferType>("dist.FInferType");
+ tvm::OpAttrMap<FInferType> op_map_context_free_infer_ty =
+ Op::GetAttrMap<FInferType>("FInferType");
+ tvm::OpAttrMap<FInferTypeWithBuilder> op_map_infer_ty =
+ Op::GetAttrMap<FInferTypeWithBuilder>("relax.FInferTypeWithBuilder");
+ tvm::OpAttrMap<FInferTypeWithBuilder> op_map_dist_infer_ty =
+ Op::GetAttrMap<FInferTypeWithBuilder>("relax.dist.FInferTypeWithBuilder");
/*! \brief Operator normalization function */
tvm::OpAttrMap<FNormalize> op_map_normalize_ = Op::GetAttrMap<FNormalize>("FNormalize");
diff --git a/src/relax/ir/dependent_type.cc b/src/relax/ir/dependent_type.cc
index 406ca9c654..6cc4070b35 100644
--- a/src/relax/ir/dependent_type.cc
+++ b/src/relax/ir/dependent_type.cc
@@ -35,28 +35,30 @@ using namespace tvm::prim;
namespace {
-TVMFFIAny ShapeTypeVisit(ffi::StructuralVisitorObj* visitor, ffi::AnyView value) noexcept {
+TVM_FFI_INLINE ffi::Expected<ffi::Optional<ffi::VisitInterrupt>> ShapeTypeVisit(
+ ffi::StructuralVisitorObj* visitor, ffi::AnyView value) noexcept {
// skips: ndim (scalar)
const ShapeTypeNode* self =
ffi::details::AnyUnsafe::RawObjectPtrFromAnyViewAfterCheck<const ShapeTypeNode>(value);
TVM_FFI_S_VISIT_MAYBE_EARLY_RETURN(visitor->VisitExpected(self->values));
- return ffi::AnyView(nullptr).CopyToTVMFFIAny();
+ return std::nullopt;
}
-TVMFFIAny ShapeTypeMutate(ffi::StructuralMutatorObj* mutator, ffi::AnyView value) noexcept {
+TVM_FFI_INLINE ffi::Expected<ffi::UnchangedOr<ffi::Any>> ShapeTypeMutate(
+ ffi::StructuralMutatorObj* mutator, ffi::AnyView value) noexcept {
// skips: ndim (scalar)
const ShapeTypeNode* self =
ffi::details::AnyUnsafe::RawObjectPtrFromAnyViewAfterCheck<const ShapeTypeNode>(value);
TVM_FFI_S_MUTATE_ASSIGN_OR_RETURN(ffi::UnchangedOr<ffi::Optional<ffi::Array<PrimExpr>>>,
mapped_values, mutator->MutateExpected(self->values));
- if (mapped_values.UnchangedOrSameAs(self->values)) return ffi::Unchanged().CopyToTVMFFIAny();
+ if (mapped_values.UnchangedOrSameAs(self->values)) return ffi::Unchanged();
ffi::ObjectPtr<ShapeTypeNode> copy = ffi::make_object<ShapeTypeNode>(*self);
copy->values = std::move(mapped_values).ValueOrUnchanged(std::move(copy->values));
- return ffi::details::AnyUnsafe::MoveAnyToTVMFFIAny(ffi::Any(std::move(copy)));
+ return ffi::Any(std::move(copy));
}
-TVMFFIAny ShapeTypeMaybeInplaceMutate(ffi::StructuralMutatorObj* mutator,
- ffi::AnyView value) noexcept {
+TVM_FFI_INLINE ffi::Expected<ffi::UnchangedOr<ffi::Any>> ShapeTypeMaybeInplaceMutate(
+ ffi::StructuralMutatorObj* mutator, ffi::AnyView value) noexcept {
// skips: ndim (scalar)
ShapeTypeNode* self = const_cast<ShapeTypeNode*>(
ffi::details::AnyUnsafe::RawObjectPtrFromAnyViewAfterCheck<const ShapeTypeNode>(value));
@@ -66,20 +68,22 @@ TVMFFIAny ShapeTypeMaybeInplaceMutate(ffi::StructuralMutatorObj* mutator,
if (!mapped_values.UnchangedOrSameAs(self->values)) {
self->values = std::move(mapped_values).ValueUnchecked();
}
- return ffi::Unchanged().CopyToTVMFFIAny();
+ return ffi::Unchanged();
}
-TVMFFIAny TensorTypeVisit(ffi::StructuralVisitorObj* visitor, ffi::AnyView value) noexcept {
+TVM_FFI_INLINE ffi::Expected<ffi::Optional<ffi::VisitInterrupt>> TensorTypeVisit(
+ ffi::StructuralVisitorObj* visitor, ffi::AnyView value) noexcept {
// skips: ndim (scalar)
const TensorTypeNode* self =
ffi::details::AnyUnsafe::RawObjectPtrFromAnyViewAfterCheck<const TensorTypeNode>(value);
TVM_FFI_S_VISIT_MAYBE_EARLY_RETURN(visitor->VisitExpected(self->shape));
TVM_FFI_S_VISIT_MAYBE_EARLY_RETURN(visitor->VisitExpected(self->dtype));
TVM_FFI_S_VISIT_MAYBE_EARLY_RETURN(visitor->VisitExpected(self->vdevice));
- return ffi::AnyView(nullptr).CopyToTVMFFIAny();
+ return std::nullopt;
}
-TVMFFIAny TensorTypeMutate(ffi::StructuralMutatorObj* mutator, ffi::AnyView value) noexcept {
+TVM_FFI_INLINE ffi::Expected<ffi::UnchangedOr<ffi::Any>> TensorTypeMutate(
+ ffi::StructuralMutatorObj* mutator, ffi::AnyView value) noexcept {
// skips: ndim (scalar)
const TensorTypeNode* self =
ffi::details::AnyUnsafe::RawObjectPtrFromAnyViewAfterCheck<const TensorTypeNode>(value);
@@ -91,17 +95,17 @@ TVMFFIAny TensorTypeMutate(ffi::StructuralMutatorObj* mutator, ffi::AnyView valu
mutator->MutateExpected(self->vdevice));
if (mapped_shape.UnchangedOrSameAs(self->shape) && mapped_dtype.UnchangedOrSameAs(self->dtype) &&
mapped_vdevice.UnchangedOrSameAs(self->vdevice)) {
- return ffi::Unchanged().CopyToTVMFFIAny();
+ return ffi::Unchanged();
}
ffi::ObjectPtr<TensorTypeNode> copy = ffi::make_object<TensorTypeNode>(*self);
copy->shape = std::move(mapped_shape).ValueOrUnchanged(std::move(copy->shape));
copy->dtype = std::move(mapped_dtype).ValueOrUnchanged(std::move(copy->dtype));
copy->vdevice = std::move(mapped_vdevice).ValueOrUnchanged(std::move(copy->vdevice));
- return ffi::details::AnyUnsafe::MoveAnyToTVMFFIAny(ffi::Any(std::move(copy)));
+ return ffi::Any(std::move(copy));
}
-TVMFFIAny TensorTypeMaybeInplaceMutate(ffi::StructuralMutatorObj* mutator,
- ffi::AnyView value) noexcept {
+TVM_FFI_INLINE ffi::Expected<ffi::UnchangedOr<ffi::Any>> TensorTypeMaybeInplaceMutate(
+ ffi::StructuralMutatorObj* mutator, ffi::AnyView value) noexcept {
// skips: ndim (scalar)
TensorTypeNode* self = const_cast<TensorTypeNode*>(
ffi::details::AnyUnsafe::RawObjectPtrFromAnyViewAfterCheck<const TensorTypeNode>(value));
@@ -121,20 +125,22 @@ TVMFFIAny TensorTypeMaybeInplaceMutate(ffi::StructuralMutatorObj* mutator,
if (!mapped_vdevice.UnchangedOrSameAs(self->vdevice)) {
self->vdevice = std::move(mapped_vdevice).ValueUnchecked();
}
- return ffi::Unchanged().CopyToTVMFFIAny();
+ return ffi::Unchanged();
}
-TVMFFIAny FuncTypeVisit(ffi::StructuralVisitorObj* visitor, ffi::AnyView value) noexcept {
+TVM_FFI_INLINE ffi::Expected<ffi::Optional<ffi::VisitInterrupt>> FuncTypeVisit(
+ ffi::StructuralVisitorObj* visitor, ffi::AnyView value) noexcept {
// skips: derive_func (environment-backed callable metadata), purity (scalar)
const FuncTypeNode* self =
ffi::details::AnyUnsafe::RawObjectPtrFromAnyViewAfterCheck<const FuncTypeNode>(value);
TVM_FFI_S_VISIT_MAYBE_EARLY_RETURN(visitor->WithDefRegionKind(
kTVMFFIDefRegionKindPattern, [&]() { return visitor->VisitExpected(self->params); }));
TVM_FFI_S_VISIT_MAYBE_EARLY_RETURN(visitor->VisitExpected(self->ret));
- return ffi::AnyView(nullptr).CopyToTVMFFIAny();
+ return std::nullopt;
}
-TVMFFIAny FuncTypeMutate(ffi::StructuralMutatorObj* mutator, ffi::AnyView value) noexcept {
+TVM_FFI_INLINE ffi::Expected<ffi::UnchangedOr<ffi::Any>> FuncTypeMutate(
+ ffi::StructuralMutatorObj* mutator, ffi::AnyView value) noexcept {
// skips: derive_func (environment-backed callable metadata), purity (scalar)
const FuncTypeNode* self =
ffi::details::AnyUnsafe::RawObjectPtrFromAnyViewAfterCheck<const FuncTypeNode>(value);
@@ -145,16 +151,16 @@ TVMFFIAny FuncTypeMutate(ffi::StructuralMutatorObj* mutator, ffi::AnyView value)
TVM_FFI_S_MUTATE_ASSIGN_OR_RETURN(ffi::UnchangedOr<Type>, mapped_ret,
mutator->MutateExpected(self->ret));
if (mapped_params.UnchangedOrSameAs(self->params) && mapped_ret.UnchangedOrSameAs(self->ret)) {
- return ffi::Unchanged().CopyToTVMFFIAny();
+ return ffi::Unchanged();
}
ffi::ObjectPtr<FuncTypeNode> copy = ffi::make_object<FuncTypeNode>(*self);
copy->params = std::move(mapped_params).ValueOrUnchanged(std::move(copy->params));
copy->ret = std::move(mapped_ret).ValueOrUnchanged(std::move(copy->ret));
- return ffi::details::AnyUnsafe::MoveAnyToTVMFFIAny(ffi::Any(std::move(copy)));
+ return ffi::Any(std::move(copy));
}
-TVMFFIAny FuncTypeMaybeInplaceMutate(ffi::StructuralMutatorObj* mutator,
- ffi::AnyView value) noexcept {
+TVM_FFI_INLINE ffi::Expected<ffi::UnchangedOr<ffi::Any>> FuncTypeMaybeInplaceMutate(
+ ffi::StructuralMutatorObj* mutator, ffi::AnyView value) noexcept {
// skips: derive_func (environment-backed callable metadata), purity (scalar)
FuncTypeNode* self = const_cast<FuncTypeNode*>(
ffi::details::AnyUnsafe::RawObjectPtrFromAnyViewAfterCheck<const FuncTypeNode>(value));
@@ -171,7 +177,7 @@ TVMFFIAny FuncTypeMaybeInplaceMutate(ffi::StructuralMutatorObj* mutator,
if (!mapped_ret.UnchangedOrSameAs(self->ret)) {
self->ret = std::move(mapped_ret).ValueUnchecked();
}
- return ffi::Unchanged().CopyToTVMFFIAny();
+ return ffi::Unchanged();
}
} // namespace
@@ -204,10 +210,11 @@ TVM_FFI_STATIC_INIT_BLOCK() {
namespace refl = tvm::ffi::reflection;
ShapeTypeNode::RegisterReflection();
refl::TypeAttrDef<ShapeTypeNode>()
- .attr(refl::type_attr::kStructuralVisit, reinterpret_cast<void*>(&ShapeTypeVisit))
- .attr(refl::type_attr::kStructuralMutate, reinterpret_cast<void*>(&ShapeTypeMutate))
+ .attr(refl::type_attr::kStructuralVisit, ffi::FStructuralVisit::FromNative<&ShapeTypeVisit>())
+ .attr(refl::type_attr::kStructuralMutate,
+ ffi::FStructuralMutate::FromNative<&ShapeTypeMutate>())
.attr(refl::type_attr::kStructuralMaybeInplaceMutate,
- reinterpret_cast<void*>(&ShapeTypeMaybeInplaceMutate));
+ ffi::FStructuralMutate::FromNative<&ShapeTypeMaybeInplaceMutate>());
}
TVM_FFI_STATIC_INIT_BLOCK() {
@@ -259,10 +266,12 @@ TVM_FFI_STATIC_INIT_BLOCK() {
namespace refl = tvm::ffi::reflection;
TensorTypeNode::RegisterReflection();
refl::TypeAttrDef<TensorTypeNode>()
- .attr(refl::type_attr::kStructuralVisit, reinterpret_cast<void*>(&TensorTypeVisit))
- .attr(refl::type_attr::kStructuralMutate, reinterpret_cast<void*>(&TensorTypeMutate))
+ .attr(refl::type_attr::kStructuralVisit,
+ ffi::FStructuralVisit::FromNative<&TensorTypeVisit>())
+ .attr(refl::type_attr::kStructuralMutate,
+ ffi::FStructuralMutate::FromNative<&TensorTypeMutate>())
.attr(refl::type_attr::kStructuralMaybeInplaceMutate,
- reinterpret_cast<void*>(&TensorTypeMaybeInplaceMutate))
+ ffi::FStructuralMutate::FromNative<&TensorTypeMaybeInplaceMutate>())
.def("__subscript_expr_realize__",
[](Expr value,
ffi::Array<ffi::Variant<ffi::Tuple<ffi::Optional<PrimExpr>, ffi::Optional<PrimExpr>,
@@ -311,10 +320,11 @@ TVM_FFI_STATIC_INIT_BLOCK() {
namespace refl = tvm::ffi::reflection;
FuncTypeNode::RegisterReflection();
refl::TypeAttrDef<FuncTypeNode>()
- .attr(refl::type_attr::kStructuralVisit, reinterpret_cast<void*>(&FuncTypeVisit))
- .attr(refl::type_attr::kStructuralMutate, reinterpret_cast<void*>(&FuncTypeMutate))
+ .attr(refl::type_attr::kStructuralVisit, ffi::FStructuralVisit::FromNative<&FuncTypeVisit>())
+ .attr(refl::type_attr::kStructuralMutate,
+ ffi::FStructuralMutate::FromNative<&FuncTypeMutate>())
.attr(refl::type_attr::kStructuralMaybeInplaceMutate,
- reinterpret_cast<void*>(&FuncTypeMaybeInplaceMutate));
+ ffi::FStructuralMutate::FromNative<&FuncTypeMaybeInplaceMutate>());
refl::GlobalDef().def("relax.FuncType",
[](ffi::Array<Type> params, Type ret, bool purity, Span span) {
diff --git a/src/relax/ir/expr.cc b/src/relax/ir/expr.cc
index 472dee24a3..0a48baeec5 100644
--- a/src/relax/ir/expr.cc
+++ b/src/relax/ir/expr.cc
@@ -34,28 +34,30 @@ namespace {
// Traverses only ExprNode::ty. The per-node payload is an intentional constant leaf:
// ExternFuncNode::global_symbol is scalar and BaseFuncNode::attrs is metadata.
template <typename TNode>
-TVMFFIAny TypeOnlyExprVisit(ffi::StructuralVisitorObj* visitor, ffi::AnyView value) noexcept {
+TVM_FFI_INLINE ffi::Expected<ffi::Optional<ffi::VisitInterrupt>> TypeOnlyExprVisit(
+ ffi::StructuralVisitorObj* visitor, ffi::AnyView value) noexcept {
const TNode* self =
ffi::details::AnyUnsafe::RawObjectPtrFromAnyViewAfterCheck<const TNode>(value);
TVM_FFI_S_VISIT_MAYBE_EARLY_RETURN(visitor->VisitExpected(self->ty));
- return ffi::AnyView(nullptr).CopyToTVMFFIAny();
+ return std::nullopt;
}
template <typename TNode>
-TVMFFIAny TypeOnlyExprMutate(ffi::StructuralMutatorObj* mutator, ffi::AnyView value) noexcept {
+TVM_FFI_INLINE ffi::Expected<ffi::UnchangedOr<ffi::Any>> TypeOnlyExprMutate(
+ ffi::StructuralMutatorObj* mutator, ffi::AnyView value) noexcept {
const TNode* self =
ffi::details::AnyUnsafe::RawObjectPtrFromAnyViewAfterCheck<const TNode>(value);
TVM_FFI_S_MUTATE_ASSIGN_OR_RETURN(ffi::UnchangedOr<Type>, mapped_ty,
mutator->MutateExpected(self->ty));
- if (mapped_ty.UnchangedOrSameAs(self->ty)) return ffi::Unchanged().CopyToTVMFFIAny();
+ if (mapped_ty.UnchangedOrSameAs(self->ty)) return ffi::Unchanged();
ffi::ObjectPtr<TNode> copy = ffi::make_object<TNode>(*self);
copy->ty = std::move(mapped_ty).ValueOrUnchanged(std::move(copy->ty));
- return ffi::details::AnyUnsafe::MoveAnyToTVMFFIAny(ffi::Any(std::move(copy)));
+ return ffi::Any(std::move(copy));
}
template <typename TNode>
-TVMFFIAny TypeOnlyExprMaybeInplaceMutate(ffi::StructuralMutatorObj* mutator,
- ffi::AnyView value) noexcept {
+TVM_FFI_INLINE ffi::Expected<ffi::UnchangedOr<ffi::Any>> TypeOnlyExprMaybeInplaceMutate(
+ ffi::StructuralMutatorObj* mutator, ffi::AnyView value) noexcept {
TNode* self = const_cast<TNode*>(
ffi::details::AnyUnsafe::RawObjectPtrFromAnyViewAfterCheck<const TNode>(value));
TVM_FFI_S_MUTATE_ASSIGN_OR_RETURN(ffi::UnchangedOr<Type>, mapped_ty,
@@ -63,20 +65,22 @@ TVMFFIAny TypeOnlyExprMaybeInplaceMutate(ffi::StructuralMutatorObj* mutator,
if (!mapped_ty.UnchangedOrSameAs(self->ty)) {
self->ty = std::move(mapped_ty).ValueUnchecked();
}
- return ffi::Unchanged().CopyToTVMFFIAny();
+ return ffi::Unchanged();
}
-TVMFFIAny IfVisit(ffi::StructuralVisitorObj* visitor, ffi::AnyView value) noexcept {
+TVM_FFI_INLINE ffi::Expected<ffi::Optional<ffi::VisitInterrupt>> IfVisit(
+ ffi::StructuralVisitorObj* visitor, ffi::AnyView value) noexcept {
const IfNode* self =
ffi::details::AnyUnsafe::RawObjectPtrFromAnyViewAfterCheck<const IfNode>(value);
TVM_FFI_S_VISIT_MAYBE_EARLY_RETURN(visitor->VisitExpected(self->ty));
TVM_FFI_S_VISIT_MAYBE_EARLY_RETURN(visitor->VisitExpected(self->cond));
TVM_FFI_S_VISIT_MAYBE_EARLY_RETURN(visitor->VisitExpected(self->true_branch));
TVM_FFI_S_VISIT_MAYBE_EARLY_RETURN(visitor->VisitExpected(self->false_branch));
- return ffi::AnyView(nullptr).CopyToTVMFFIAny();
+ return std::nullopt;
}
-TVMFFIAny IfMutate(ffi::StructuralMutatorObj* mutator, ffi::AnyView value) noexcept {
+TVM_FFI_INLINE ffi::Expected<ffi::UnchangedOr<ffi::Any>> IfMutate(
+ ffi::StructuralMutatorObj* mutator, ffi::AnyView value) noexcept {
const IfNode* self =
ffi::details::AnyUnsafe::RawObjectPtrFromAnyViewAfterCheck<const IfNode>(value);
TVM_FFI_S_MUTATE_ASSIGN_OR_RETURN(ffi::UnchangedOr<Type>, mapped_ty,
@@ -90,7 +94,7 @@ TVMFFIAny IfMutate(ffi::StructuralMutatorObj* mutator, ffi::AnyView value) noexc
if (mapped_ty.UnchangedOrSameAs(self->ty) && mapped_cond.UnchangedOrSameAs(self->cond) &&
mapped_true_branch.UnchangedOrSameAs(self->true_branch) &&
mapped_false_branch.UnchangedOrSameAs(self->false_branch)) {
- return ffi::Unchanged().CopyToTVMFFIAny();
+ return ffi::Unchanged();
}
ffi::ObjectPtr<IfNode> copy = ffi::make_object<IfNode>(*self);
copy->ty = std::move(mapped_ty).ValueOrUnchanged(std::move(copy->ty));
@@ -98,10 +102,11 @@ TVMFFIAny IfMutate(ffi::StructuralMutatorObj* mutator, ffi::AnyView value) noexc
copy->true_branch = std::move(mapped_true_branch).ValueOrUnchanged(std::move(copy->true_branch));
copy->false_branch =
std::move(mapped_false_branch).ValueOrUnchanged(std::move(copy->false_branch));
- return ffi::details::AnyUnsafe::MoveAnyToTVMFFIAny(ffi::Any(std::move(copy)));
+ return ffi::Any(std::move(copy));
}
-TVMFFIAny IfMaybeInplaceMutate(ffi::StructuralMutatorObj* mutator, ffi::AnyView value) noexcept {
+TVM_FFI_INLINE ffi::Expected<ffi::UnchangedOr<ffi::Any>> IfMaybeInplaceMutate(
+ ffi::StructuralMutatorObj* mutator, ffi::AnyView value) noexcept {
IfNode* self = const_cast<IfNode*>(
ffi::details::AnyUnsafe::RawObjectPtrFromAnyViewAfterCheck<const IfNode>(value));
TVM_FFI_S_MUTATE_ASSIGN_OR_RETURN(ffi::UnchangedOr<Type>, mapped_ty,
@@ -124,18 +129,20 @@ TVMFFIAny IfMaybeInplaceMutate(ffi::StructuralMutatorObj* mutator, ffi::AnyView
if (!mapped_false_branch.UnchangedOrSameAs(self->false_branch)) {
self->false_branch = std::move(mapped_false_branch).ValueUnchecked();
}
- return ffi::Unchanged().CopyToTVMFFIAny();
+ return ffi::Unchanged();
}
-TVMFFIAny ShapeExprVisit(ffi::StructuralVisitorObj* visitor, ffi::AnyView value) noexcept {
+TVM_FFI_INLINE ffi::Expected<ffi::Optional<ffi::VisitInterrupt>> ShapeExprVisit(
+ ffi::StructuralVisitorObj* visitor, ffi::AnyView value) noexcept {
const ShapeExprNode* self =
ffi::details::AnyUnsafe::RawObjectPtrFromAnyViewAfterCheck<const ShapeExprNode>(value);
TVM_FFI_S_VISIT_MAYBE_EARLY_RETURN(visitor->VisitExpected(self->ty));
TVM_FFI_S_VISIT_MAYBE_EARLY_RETURN(visitor->VisitExpected(self->values));
- return ffi::AnyView(nullptr).CopyToTVMFFIAny();
+ return std::nullopt;
}
-TVMFFIAny ShapeExprMutate(ffi::StructuralMutatorObj* mutator, ffi::AnyView value) noexcept {
+TVM_FFI_INLINE ffi::Expected<ffi::UnchangedOr<ffi::Any>> ShapeExprMutate(
+ ffi::StructuralMutatorObj* mutator, ffi::AnyView value) noexcept {
const ShapeExprNode* self =
ffi::details::AnyUnsafe::RawObjectPtrFromAnyViewAfterCheck<const ShapeExprNode>(value);
TVM_FFI_S_MUTATE_ASSIGN_OR_RETURN(ffi::UnchangedOr<Type>, mapped_ty,
@@ -143,16 +150,16 @@ TVMFFIAny ShapeExprMutate(ffi::StructuralMutatorObj* mutator, ffi::AnyView value
TVM_FFI_S_MUTATE_ASSIGN_OR_RETURN(ffi::UnchangedOr<ffi::Array<PrimExpr>>, mapped_values,
mutator->MutateExpected(self->values));
if (mapped_ty.UnchangedOrSameAs(self->ty) && mapped_values.UnchangedOrSameAs(self->values)) {
- return ffi::Unchanged().CopyToTVMFFIAny();
+ return ffi::Unchanged();
}
ffi::ObjectPtr<ShapeExprNode> copy = ffi::make_object<ShapeExprNode>(*self);
copy->ty = std::move(mapped_ty).ValueOrUnchanged(std::move(copy->ty));
copy->values = std::move(mapped_values).ValueOrUnchanged(std::move(copy->values));
- return ffi::details::AnyUnsafe::MoveAnyToTVMFFIAny(ffi::Any(std::move(copy)));
+ return ffi::Any(std::move(copy));
}
-TVMFFIAny ShapeExprMaybeInplaceMutate(ffi::StructuralMutatorObj* mutator,
- ffi::AnyView value) noexcept {
+TVM_FFI_INLINE ffi::Expected<ffi::UnchangedOr<ffi::Any>> ShapeExprMaybeInplaceMutate(
+ ffi::StructuralMutatorObj* mutator, ffi::AnyView value) noexcept {
ShapeExprNode* self = const_cast<ShapeExprNode*>(
ffi::details::AnyUnsafe::RawObjectPtrFromAnyViewAfterCheck<const ShapeExprNode>(value));
TVM_FFI_S_MUTATE_ASSIGN_OR_RETURN(ffi::UnchangedOr<Type>, mapped_ty,
@@ -164,12 +171,13 @@ TVMFFIAny ShapeExprMaybeInplaceMutate(ffi::StructuralMutatorObj* mutator,
if (!mapped_values.UnchangedOrSameAs(self->values)) {
self->values = std::move(mapped_values).ValueUnchecked();
}
- return ffi::Unchanged().CopyToTVMFFIAny();
+ return ffi::Unchanged();
}
// Hooks do not inherit, so DataflowVar must mirror the base VarNode remap, PrimType-skip, and
// Simple-to-None definition-region protocol. Keep this hook triple in lockstep with VarNode.
-TVMFFIAny DataflowVarVisit(ffi::StructuralVisitorObj* visitor, ffi::AnyView value) noexcept {
+TVM_FFI_INLINE ffi::Expected<ffi::Optional<ffi::VisitInterrupt>> DataflowVarVisit(
+ ffi::StructuralVisitorObj* visitor, ffi::AnyView value) noexcept {
const DataflowVarNode* self =
ffi::details::AnyUnsafe::RawObjectPtrFromAnyViewAfterCheck<const DataflowVarNode>(value);
if (!self->ty.as<PrimTypeNode>()) {
@@ -180,20 +188,21 @@ TVMFFIAny DataflowVarVisit(ffi::StructuralVisitorObj* visitor, ffi::AnyView valu
TVM_FFI_S_VISIT_MAYBE_EARLY_RETURN(visitor->VisitExpected(self->ty));
}
}
- return ffi::AnyView(nullptr).CopyToTVMFFIAny();
+ return std::nullopt;
}
-TVMFFIAny DataflowVarMutate(ffi::StructuralMutatorObj* mutator, ffi::AnyView value) noexcept {
+TVM_FFI_INLINE ffi::Expected<ffi::UnchangedOr<ffi::Any>> DataflowVarMutate(
+ ffi::StructuralMutatorObj* mutator, ffi::AnyView value) noexcept {
const DataflowVarNode* self =
ffi::details::AnyUnsafe::RawObjectPtrFromAnyViewAfterCheck<const DataflowVarNode>(value);
ffi::Expected<ffi::Any> remap_result = mutator->VarRemapGetExpected(value);
TVM_FFI_S_MUTATE_MAYBE_EARLY_RETURN(remap_result);
if (ffi::details::ExpectedUnsafe::GetData(remap_result).type_index() !=
ffi::TypeIndex::kTVMFFINone) {
- return ffi::details::ExpectedUnsafe::MoveToTVMFFIAny(std::move(remap_result));
+ return std::move(remap_result);
}
if (mutator->def_region_kind() == kTVMFFIDefRegionKindNone) {
- return ffi::Unchanged().CopyToTVMFFIAny();
+ return ffi::Unchanged();
}
ffi::UnchangedOr<ffi::Any> result = ffi::Unchanged();
ffi::Any mapped_value;
@@ -216,24 +225,24 @@ TVMFFIAny DataflowVarMutate(ffi::StructuralMutatorObj* mutator, ffi::AnyView val
ffi::AnyView value_to_store = result.IsUnchanged() ? value : ffi::AnyView(mapped_value);
auto set_result = mutator->VarRemapSetExpected(value, value_to_store);
if (TVM_FFI_PREDICT_FALSE(set_result.is_err())) {
- return ffi::details::AnyUnsafe::MoveAnyToTVMFFIAny(ffi::Any(std::move(set_result).error()));
+ return ffi::Unexpected(std::move(set_result).error());
}
}
- return ffi::details::UnchangedOrUnsafe::MoveToTVMFFIAny(std::move(result));
+ return std::move(result);
}
-TVMFFIAny DataflowVarMaybeInplaceMutate(ffi::StructuralMutatorObj* mutator,
- ffi::AnyView value) noexcept {
+TVM_FFI_INLINE ffi::Expected<ffi::UnchangedOr<ffi::Any>> DataflowVarMaybeInplaceMutate(
+ ffi::StructuralMutatorObj* mutator, ffi::AnyView value) noexcept {
DataflowVarNode* self = const_cast<DataflowVarNode*>(
ffi::details::AnyUnsafe::RawObjectPtrFromAnyViewAfterCheck<const DataflowVarNode>(value));
ffi::Expected<ffi::Any> remap_result = mutator->VarRemapGetExpected(value);
TVM_FFI_S_MUTATE_MAYBE_EARLY_RETURN(remap_result);
if (ffi::details::ExpectedUnsafe::GetData(remap_result).type_index() !=
ffi::TypeIndex::kTVMFFINone) {
- return ffi::details::ExpectedUnsafe::MoveToTVMFFIAny(std::move(remap_result));
+ return std::move(remap_result);
}
if (mutator->def_region_kind() == kTVMFFIDefRegionKindNone) {
- return ffi::Unchanged().CopyToTVMFFIAny();
+ return ffi::Unchanged();
}
ffi::UnchangedOr<ffi::Any> result = ffi::Unchanged();
ffi::Any mapped_value;
@@ -256,22 +265,24 @@ TVMFFIAny DataflowVarMaybeInplaceMutate(ffi::StructuralMutatorObj* mutator,
ffi::AnyView value_to_store = result.IsUnchanged() ? value : ffi::AnyView(mapped_value);
auto set_result = mutator->VarRemapSetExpected(value, value_to_store);
if (TVM_FFI_PREDICT_FALSE(set_result.is_err())) {
- return ffi::details::AnyUnsafe::MoveAnyToTVMFFIAny(ffi::Any(std::move(set_result).error()));
+ return ffi::Unexpected(std::move(set_result).error());
}
}
- return ffi::details::UnchangedOrUnsafe::MoveToTVMFFIAny(std::move(result));
+ return std::move(result);
}
-TVMFFIAny SeqExprVisit(ffi::StructuralVisitorObj* visitor, ffi::AnyView value) noexcept {
+TVM_FFI_INLINE ffi::Expected<ffi::Optional<ffi::VisitInterrupt>> SeqExprVisit(
+ ffi::StructuralVisitorObj* visitor, ffi::AnyView value) noexcept {
const SeqExprNode* self =
ffi::details::AnyUnsafe::RawObjectPtrFromAnyViewAfterCheck<const SeqExprNode>(value);
TVM_FFI_S_VISIT_MAYBE_EARLY_RETURN(visitor->VisitExpected(self->blocks));
TVM_FFI_S_VISIT_MAYBE_EARLY_RETURN(visitor->VisitExpected(self->ty));
TVM_FFI_S_VISIT_MAYBE_EARLY_RETURN(visitor->VisitExpected(self->body));
- return ffi::AnyView(nullptr).CopyToTVMFFIAny();
+ return std::nullopt;
}
-TVMFFIAny SeqExprMutate(ffi::StructuralMutatorObj* mutator, ffi::AnyView value) noexcept {
+TVM_FFI_INLINE ffi::Expected<ffi::UnchangedOr<ffi::Any>> SeqExprMutate(
+ ffi::StructuralMutatorObj* mutator, ffi::AnyView value) noexcept {
const SeqExprNode* self =
ffi::details::AnyUnsafe::RawObjectPtrFromAnyViewAfterCheck<const SeqExprNode>(value);
TVM_FFI_S_MUTATE_ASSIGN_OR_RETURN(ffi::UnchangedOr<ffi::Array<BindingBlock>>, mapped_blocks,
@@ -282,17 +293,17 @@ TVMFFIAny SeqExprMutate(ffi::StructuralMutatorObj* mutator, ffi::AnyView value)
mutator->MutateExpected(self->body));
if (mapped_blocks.UnchangedOrSameAs(self->blocks) && mapped_ty.UnchangedOrSameAs(self->ty) &&
mapped_body.UnchangedOrSameAs(self->body)) {
- return ffi::Unchanged().CopyToTVMFFIAny();
+ return ffi::Unchanged();
}
ffi::ObjectPtr<SeqExprNode> copy = ffi::make_object<SeqExprNode>(*self);
copy->blocks = std::move(mapped_blocks).ValueOrUnchanged(std::move(copy->blocks));
copy->ty = std::move(mapped_ty).ValueOrUnchanged(std::move(copy->ty));
copy->body = std::move(mapped_body).ValueOrUnchanged(std::move(copy->body));
- return ffi::details::AnyUnsafe::MoveAnyToTVMFFIAny(ffi::Any(std::move(copy)));
+ return ffi::Any(std::move(copy));
}
-TVMFFIAny SeqExprMaybeInplaceMutate(ffi::StructuralMutatorObj* mutator,
- ffi::AnyView value) noexcept {
+TVM_FFI_INLINE ffi::Expected<ffi::UnchangedOr<ffi::Any>> SeqExprMaybeInplaceMutate(
+ ffi::StructuralMutatorObj* mutator, ffi::AnyView value) noexcept {
SeqExprNode* self = const_cast<SeqExprNode*>(
ffi::details::AnyUnsafe::RawObjectPtrFromAnyViewAfterCheck<const SeqExprNode>(value));
TVM_FFI_S_MUTATE_ASSIGN_OR_RETURN(
@@ -309,12 +320,13 @@ TVMFFIAny SeqExprMaybeInplaceMutate(ffi::StructuralMutatorObj* mutator,
if (!mapped_body.UnchangedOrSameAs(self->body)) {
self->body = std::move(mapped_body).ValueUnchecked();
}
- return ffi::Unchanged().CopyToTVMFFIAny();
+ return ffi::Unchanged();
}
// Parameters precede the reflected ty field in this hook triple so their pattern region establishes
// the remap before the derived function type can refer to those symbolic definitions.
-TVMFFIAny FunctionVisit(ffi::StructuralVisitorObj* visitor, ffi::AnyView value) noexcept {
+TVM_FFI_INLINE ffi::Expected<ffi::Optional<ffi::VisitInterrupt>> FunctionVisit(
+ ffi::StructuralVisitorObj* visitor, ffi::AnyView value) noexcept {
// skips: attrs (metadata), is_pure (scalar)
const FunctionNode* self =
ffi::details::AnyUnsafe::RawObjectPtrFromAnyViewAfterCheck<const FunctionNode>(value);
@@ -323,10 +335,11 @@ TVMFFIAny FunctionVisit(ffi::StructuralVisitorObj* visitor, ffi::AnyView value)
TVM_FFI_S_VISIT_MAYBE_EARLY_RETURN(visitor->VisitExpected(self->ty));
TVM_FFI_S_VISIT_MAYBE_EARLY_RETURN(visitor->VisitExpected(self->body));
TVM_FFI_S_VISIT_MAYBE_EARLY_RETURN(visitor->VisitExpected(self->ret_ty));
- return ffi::AnyView(nullptr).CopyToTVMFFIAny();
+ return std::nullopt;
}
-TVMFFIAny FunctionMutate(ffi::StructuralMutatorObj* mutator, ffi::AnyView value) noexcept {
+TVM_FFI_INLINE ffi::Expected<ffi::UnchangedOr<ffi::Any>> FunctionMutate(
+ ffi::StructuralMutatorObj* mutator, ffi::AnyView value) noexcept {
// skips: attrs (metadata), is_pure (scalar)
const FunctionNode* self =
ffi::details::AnyUnsafe::RawObjectPtrFromAnyViewAfterCheck<const FunctionNode>(value);
@@ -342,18 +355,18 @@ TVMFFIAny FunctionMutate(ffi::StructuralMutatorObj* mutator, ffi::AnyView value)
mutator->MutateExpected(self->ret_ty));
if (mapped_params.UnchangedOrSameAs(self->params) && mapped_ty.UnchangedOrSameAs(self->ty) &&
mapped_body.UnchangedOrSameAs(self->body) && mapped_ret_ty.UnchangedOrSameAs(self->ret_ty)) {
- return ffi::Unchanged().CopyToTVMFFIAny();
+ return ffi::Unchanged();
}
ffi::ObjectPtr<FunctionNode> copy = ffi::make_object<FunctionNode>(*self);
copy->params = std::move(mapped_params).ValueOrUnchanged(std::move(copy->params));
copy->ty = std::move(mapped_ty).ValueOrUnchanged(std::move(copy->ty));
copy->body = std::move(mapped_body).ValueOrUnchanged(std::move(copy->body));
copy->ret_ty = std::move(mapped_ret_ty).ValueOrUnchanged(std::move(copy->ret_ty));
- return ffi::details::AnyUnsafe::MoveAnyToTVMFFIAny(ffi::Any(std::move(copy)));
+ return ffi::Any(std::move(copy));
}
-TVMFFIAny FunctionMaybeInplaceMutate(ffi::StructuralMutatorObj* mutator,
- ffi::AnyView value) noexcept {
+TVM_FFI_INLINE ffi::Expected<ffi::UnchangedOr<ffi::Any>> FunctionMaybeInplaceMutate(
+ ffi::StructuralMutatorObj* mutator, ffi::AnyView value) noexcept {
// skips: attrs (metadata), is_pure (scalar)
FunctionNode* self = const_cast<FunctionNode*>(
ffi::details::AnyUnsafe::RawObjectPtrFromAnyViewAfterCheck<const FunctionNode>(value));
@@ -379,7 +392,7 @@ TVMFFIAny FunctionMaybeInplaceMutate(ffi::StructuralMutatorObj* mutator,
if (!mapped_ret_ty.UnchangedOrSameAs(self->ret_ty)) {
self->ret_ty = std::move(mapped_ret_ty).ValueUnchecked();
}
- return ffi::Unchanged().CopyToTVMFFIAny();
+ return ffi::Unchanged();
}
} // namespace
@@ -399,10 +412,10 @@ TVM_FFI_STATIC_INIT_BLOCK() {
namespace refl = tvm::ffi::reflection;
IfNode::RegisterReflection();
refl::TypeAttrDef<IfNode>()
- .attr(refl::type_attr::kStructuralVisit, reinterpret_cast<void*>(&IfVisit))
- .attr(refl::type_attr::kStructuralMutate, reinterpret_cast<void*>(&IfMutate))
+ .attr(refl::type_attr::kStructuralVisit, ffi::FStructuralVisit::FromNative<&IfVisit>())
+ .attr(refl::type_attr::kStructuralMutate, ffi::FStructuralMutate::FromNative<&IfMutate>())
.attr(refl::type_attr::kStructuralMaybeInplaceMutate,
- reinterpret_cast<void*>(&IfMaybeInplaceMutate));
+ ffi::FStructuralMutate::FromNative<&IfMaybeInplaceMutate>());
refl::GlobalDef().def("relax.If", [](Expr cond, Expr true_branch, Expr false_branch, Span span) {
return If(cond, true_branch, false_branch, span);
@@ -429,10 +442,11 @@ TVM_FFI_STATIC_INIT_BLOCK() {
namespace refl = tvm::ffi::reflection;
ShapeExprNode::RegisterReflection();
refl::TypeAttrDef<ShapeExprNode>()
- .attr(refl::type_attr::kStructuralVisit, reinterpret_cast<void*>(&ShapeExprVisit))
- .attr(refl::type_attr::kStructuralMutate, reinterpret_cast<void*>(&ShapeExprMutate))
+ .attr(refl::type_attr::kStructuralVisit, ffi::FStructuralVisit::FromNative<&ShapeExprVisit>())
+ .attr(refl::type_attr::kStructuralMutate,
+ ffi::FStructuralMutate::FromNative<&ShapeExprMutate>())
.attr(refl::type_attr::kStructuralMaybeInplaceMutate,
- reinterpret_cast<void*>(&ShapeExprMaybeInplaceMutate));
+ ffi::FStructuralMutate::FromNative<&ShapeExprMaybeInplaceMutate>());
refl::GlobalDef().def("relax.ShapeExpr", [](ffi::Array<PrimExpr> values, Span span) {
return ShapeExpr(values, span);
@@ -453,10 +467,12 @@ TVM_FFI_STATIC_INIT_BLOCK() {
namespace refl = tvm::ffi::reflection;
DataflowVarNode::RegisterReflection();
refl::TypeAttrDef<DataflowVarNode>()
- .attr(refl::type_attr::kStructuralVisit, reinterpret_cast<void*>(&DataflowVarVisit))
- .attr(refl::type_attr::kStructuralMutate, reinterpret_cast<void*>(&DataflowVarMutate))
+ .attr(refl::type_attr::kStructuralVisit,
+ ffi::FStructuralVisit::FromNative<&DataflowVarVisit>())
+ .attr(refl::type_attr::kStructuralMutate,
+ ffi::FStructuralMutate::FromNative<&DataflowVarMutate>())
.attr(refl::type_attr::kStructuralMaybeInplaceMutate,
- reinterpret_cast<void*>(&DataflowVarMaybeInplaceMutate));
+ ffi::FStructuralMutate::FromNative<&DataflowVarMaybeInplaceMutate>());
refl::GlobalDef().def("relax.DataflowVar",
[](ffi::String name, ffi::Optional<Type> ty_annotation, Span span) {
@@ -615,10 +631,11 @@ TVM_FFI_STATIC_INIT_BLOCK() {
namespace refl = tvm::ffi::reflection;
SeqExprNode::RegisterReflection();
refl::TypeAttrDef<SeqExprNode>()
- .attr(refl::type_attr::kStructuralVisit, reinterpret_cast<void*>(&SeqExprVisit))
- .attr(refl::type_attr::kStructuralMutate, reinterpret_cast<void*>(&SeqExprMutate))
+ .attr(refl::type_attr::kStructuralVisit, ffi::FStructuralVisit::FromNative<&SeqExprVisit>())
+ .attr(refl::type_attr::kStructuralMutate,
+ ffi::FStructuralMutate::FromNative<&SeqExprMutate>())
.attr(refl::type_attr::kStructuralMaybeInplaceMutate,
- reinterpret_cast<void*>(&SeqExprMaybeInplaceMutate));
+ ffi::FStructuralMutate::FromNative<&SeqExprMaybeInplaceMutate>());
refl::GlobalDef().def("relax.SeqExpr", [](ffi::Array<BindingBlock> blocks, Expr body, Span span) {
return SeqExpr(blocks, body, span);
@@ -690,10 +707,11 @@ TVM_FFI_STATIC_INIT_BLOCK() {
namespace refl = tvm::ffi::reflection;
FunctionNode::RegisterReflection();
refl::TypeAttrDef<FunctionNode>()
- .attr(refl::type_attr::kStructuralVisit, reinterpret_cast<void*>(&FunctionVisit))
- .attr(refl::type_attr::kStructuralMutate, reinterpret_cast<void*>(&FunctionMutate))
+ .attr(refl::type_attr::kStructuralVisit, ffi::FStructuralVisit::FromNative<&FunctionVisit>())
+ .attr(refl::type_attr::kStructuralMutate,
+ ffi::FStructuralMutate::FromNative<&FunctionMutate>())
.attr(refl::type_attr::kStructuralMaybeInplaceMutate,
- reinterpret_cast<void*>(&FunctionMaybeInplaceMutate));
+ ffi::FStructuralMutate::FromNative<&FunctionMaybeInplaceMutate>());
refl::GlobalDef().def("relax.Function",
[](ffi::Array<Var> params, Expr body, ffi::Optional<Type> ret_ty,
@@ -785,11 +803,11 @@ TVM_FFI_STATIC_INIT_BLOCK() {
ExternFuncNode::RegisterReflection();
refl::TypeAttrDef<ExternFuncNode>()
.attr(refl::type_attr::kStructuralVisit,
- reinterpret_cast<void*>(&TypeOnlyExprVisit<ExternFuncNode>))
+ ffi::FStructuralVisit::FromNative<&TypeOnlyExprVisit<ExternFuncNode>>())
.attr(refl::type_attr::kStructuralMutate,
- reinterpret_cast<void*>(&TypeOnlyExprMutate<ExternFuncNode>))
+ ffi::FStructuralMutate::FromNative<&TypeOnlyExprMutate<ExternFuncNode>>())
.attr(refl::type_attr::kStructuralMaybeInplaceMutate,
- reinterpret_cast<void*>(&TypeOnlyExprMaybeInplaceMutate<ExternFuncNode>));
+ ffi::FStructuralMutate::FromNative<&TypeOnlyExprMaybeInplaceMutate<ExternFuncNode>>());
refl::GlobalDef().def("relax.ExternFunc",
[](ffi::String global_symbol, ffi::Optional<Type> ty, Span span) {
diff --git a/src/relax/ir/type.cc b/src/relax/ir/type.cc
index 6da87e6d2f..7b01a4b137 100644
--- a/src/relax/ir/type.cc
+++ b/src/relax/ir/type.cc
@@ -32,16 +32,19 @@ namespace relax {
namespace {
-TVMFFIAny PackedFuncTypeVisit(ffi::StructuralVisitorObj*, ffi::AnyView) noexcept {
- return ffi::AnyView(nullptr).CopyToTVMFFIAny();
+TVM_FFI_INLINE ffi::Expected<ffi::Optional<ffi::VisitInterrupt>> PackedFuncTypeVisit(
+ ffi::StructuralVisitorObj*, ffi::AnyView) noexcept {
+ return std::nullopt;
}
-TVMFFIAny PackedFuncTypeMutate(ffi::StructuralMutatorObj*, ffi::AnyView) noexcept {
- return ffi::Unchanged().CopyToTVMFFIAny();
+TVM_FFI_INLINE ffi::Expected<ffi::UnchangedOr<ffi::Any>> PackedFuncTypeMutate(
+ ffi::StructuralMutatorObj*, ffi::AnyView) noexcept {
+ return ffi::Unchanged();
}
-TVMFFIAny PackedFuncTypeMaybeInplaceMutate(ffi::StructuralMutatorObj*, ffi::AnyView) noexcept {
- return ffi::Unchanged().CopyToTVMFFIAny();
+TVM_FFI_INLINE ffi::Expected<ffi::UnchangedOr<ffi::Any>> PackedFuncTypeMaybeInplaceMutate(
+ ffi::StructuralMutatorObj*, ffi::AnyView) noexcept {
+ return ffi::Unchanged();
}
} // namespace
@@ -56,10 +59,12 @@ TVM_FFI_STATIC_INIT_BLOCK() {
namespace refl = tvm::ffi::reflection;
PackedFuncTypeNode::RegisterReflection();
refl::TypeAttrDef<PackedFuncTypeNode>()
- .attr(refl::type_attr::kStructuralVisit, reinterpret_cast<void*>(&PackedFuncTypeVisit))
- .attr(refl::type_attr::kStructuralMutate, reinterpret_cast<void*>(&PackedFuncTypeMutate))
+ .attr(refl::type_attr::kStructuralVisit,
+ ffi::FStructuralVisit::FromNative<&PackedFuncTypeVisit>())
+ .attr(refl::type_attr::kStructuralMutate,
+ ffi::FStructuralMutate::FromNative<&PackedFuncTypeMutate>())
.attr(refl::type_attr::kStructuralMaybeInplaceMutate,
- reinterpret_cast<void*>(&PackedFuncTypeMaybeInplaceMutate));
+ ffi::FStructuralMutate::FromNative<&PackedFuncTypeMaybeInplaceMutate>());
refl::GlobalDef().def("relax.PackedFuncType", [](Span span) { return PackedFuncType(span); });
}
diff --git a/src/relax/op/ccl/ccl.cc b/src/relax/op/ccl/ccl.cc
index 77705b87a0..0ac13cf5c1 100644
--- a/src/relax/op/ccl/ccl.cc
+++ b/src/relax/op/ccl/ccl.cc
@@ -51,16 +51,21 @@ TVM_FFI_STATIC_INIT_BLOCK() {
refl::GlobalDef().def("relax.op.ccl.allreduce", allreduce);
}
-Type InferTypeAllReduce(const Call& call, const BlockBuilder& ctx) {
- TensorType input_ty = GetUnaryInputTensorType(call, ctx);
+ffi::Expected<Type> InferTypeAllReduce(const CallNode* call_node) noexcept try {
+ const Call call = ffi::GetRef<Call>(call_node);
+ TensorType input_ty = GetUnaryInputTensorType(call);
return input_ty;
+} catch (const ffi::Error& error) {
+ return ffi::Unexpected(error);
+} catch (const std::exception& error) {
+ return ffi::Unexpected(ffi::Error("InternalError", error.what(), ""));
}
TVM_FFI_STATIC_INIT_BLOCK() {
OpDef("relax.ccl.allreduce")
.signature(sig::arg("x", "Input to which allreduce will be applied."),
sig::call_attrs<AllReduceAttrs>())
- .set_attr<FInferType>("FInferType", InferTypeAllReduce)
+ .set_attr<FInferType>("FInferType", FInferType::FromNative<&InferTypeAllReduce>())
.set_attr<FRelaxInferLayout>("FRelaxInferLayout", InferLayoutUnaryEwise)
.set_attr<bool>("FPurity", true);
}
@@ -81,8 +86,9 @@ TVM_FFI_STATIC_INIT_BLOCK() {
refl::GlobalDef().def("relax.op.ccl.allgather", allgather);
}
-Type InferTypeAllGather(const Call& call, const BlockBuilder& ctx) {
- TensorType input_ty = GetUnaryInputTensorType(call, ctx);
+ffi::Expected<Type> InferTypeAllGather(const CallNode* call_node) noexcept try {
+ const Call call = ffi::GetRef<Call>(call_node);
+ TensorType input_ty = GetUnaryInputTensorType(call);
const auto* attrs = call->attrs.as<AllGatherAttrs>();
int num_workers = attrs->num_workers;
@@ -95,13 +101,17 @@ Type InferTypeAllGather(const Call& call, const BlockBuilder& ctx) {
ffi::Array<PrimExpr> output_shape = input_shape.value();
output_shape.Set(0, floor(output_shape[0] * num_workers));
return TensorType(ShapeExpr(output_shape), output_dtype, input_ty->vdevice);
+} catch (const ffi::Error& error) {
+ return ffi::Unexpected(error);
+} catch (const std::exception& error) {
+ return ffi::Unexpected(ffi::Error("InternalError", error.what(), ""));
}
TVM_FFI_STATIC_INIT_BLOCK() {
OpDef("relax.ccl.allgather")
.signature(sig::arg("x", "Input to which allgather will be applied."),
sig::call_attrs<AllGatherAttrs>())
- .set_attr<FInferType>("FInferType", InferTypeAllGather)
+ .set_attr<FInferType>("FInferType", FInferType::FromNative<&InferTypeAllGather>())
.set_attr<FRelaxInferLayout>("FRelaxInferLayout", InferLayoutUnaryEwise)
.set_attr<bool>("FPurity", true);
}
@@ -117,15 +127,20 @@ TVM_FFI_STATIC_INIT_BLOCK() {
refl::GlobalDef().def("relax.op.ccl.broadcast_from_worker0", broadcast_from_worker0);
}
-Type InferTypeBroadcastFromZero(const Call& call, const BlockBuilder& ctx) {
- TensorType input_ty = GetUnaryInputTensorType(call, ctx);
+ffi::Expected<Type> InferTypeBroadcastFromZero(const CallNode* call_node) noexcept try {
+ const Call call = ffi::GetRef<Call>(call_node);
+ TensorType input_ty = GetUnaryInputTensorType(call);
return input_ty;
+} catch (const ffi::Error& error) {
+ return ffi::Unexpected(error);
+} catch (const std::exception& error) {
+ return ffi::Unexpected(ffi::Error("InternalError", error.what(), ""));
}
TVM_FFI_STATIC_INIT_BLOCK() {
OpDef("relax.ccl.broadcast_from_worker0")
.signature(sig::arg("x", "Input to be broadcast."))
- .set_attr<FInferType>("FInferType", InferTypeBroadcastFromZero)
+ .set_attr<FInferType>("FInferType", FInferType::FromNative<&InferTypeBroadcastFromZero>())
.set_attr<FRelaxInferLayout>("FRelaxInferLayout", InferLayoutUnaryEwise)
.set_attr<bool>("FPurity", true);
}
@@ -177,7 +192,7 @@ TVM_FFI_STATIC_INIT_BLOCK() {
"x",
"The buffer to be divided into equal parts and sent to each worker accordingly."),
sig::call_attrs<ScatterCollectiveAttrs>())
- .set_attr<FInferType>("FInferType", InferTypeScatter)
+ .set_attr<FInferTypeWithBuilder>("relax.FInferTypeWithBuilder", InferTypeScatter)
.set_attr<bool>("FPurity", true);
}
diff --git a/src/relax/op/distributed/binary.cc b/src/relax/op/distributed/binary.cc
index 2ec0550ed4..5d98e6757b 100644
--- a/src/relax/op/distributed/binary.cc
+++ b/src/relax/op/distributed/binary.cc
@@ -24,7 +24,10 @@ namespace relax {
namespace distributed {
Type InferDistTypeBroadcastArith(const Call& call, const BlockBuilder& ctx) {
- return InferDistTypeBroadcast(call, ctx, InferBinaryArithOpOutDtype);
+ return InferDistTypeBroadcast(
+ call, ctx, [](const Call& call, const BlockBuilder& ctx, const Type& lhs, const Type& rhs) {
+ return InferBinaryArithOpOutDtype(call, ctx, lhs, rhs);
+ });
}
Type InferDistTypeBroadcastCMP(const Call& call, const BlockBuilder& ctx) {
@@ -39,84 +42,84 @@ Type InferDistTypeBroadcastCMP(const Call& call, const BlockBuilder& ctx) {
TVM_FFI_STATIC_INIT_BLOCK() {
// clang-format off
OpDef("relax.add")
- .set_attr<FInferType>("dist.FInferType", InferDistTypeBroadcastArith);
+ .set_attr<FInferTypeWithBuilder>("relax.dist.FInferTypeWithBuilder", InferDistTypeBroadcastArith);
OpDef("relax.divide")
- .set_attr<FInferType>("dist.FInferType", InferDistTypeBroadcastArith);
+ .set_attr<FInferTypeWithBuilder>("relax.dist.FInferTypeWithBuilder", InferDistTypeBroadcastArith);
OpDef("relax.floor_divide")
- .set_attr<FInferType>("dist.FInferType", InferDistTypeBroadcastArith);
+ .set_attr<FInferTypeWithBuilder>("relax.dist.FInferTypeWithBuilder", InferDistTypeBroadcastArith);
OpDef("relax.multiply")
- .set_attr<FInferType>("dist.FInferType", InferDistTypeBroadcastArith);
+ .set_attr<FInferTypeWithBuilder>("relax.dist.FInferTypeWithBuilder", InferDistTypeBroadcastArith);
OpDef("relax.power")
- .set_attr<FInferType>("dist.FInferType", InferDistTypeBroadcastArith);
+ .set_attr<FInferTypeWithBuilder>("relax.dist.FInferTypeWithBuilder", InferDistTypeBroadcastArith);
OpDef("relax.subtract")
- .set_attr<FInferType>("dist.FInferType", InferDistTypeBroadcastArith);
+ .set_attr<FInferTypeWithBuilder>("relax.dist.FInferTypeWithBuilder", InferDistTypeBroadcastArith);
OpDef("relax.mod")
- .set_attr<FInferType>("dist.FInferType", InferDistTypeBroadcastArith);
+ .set_attr<FInferTypeWithBuilder>("relax.dist.FInferTypeWithBuilder", InferDistTypeBroadcastArith);
OpDef("relax.floor_mod")
- .set_attr<FInferType>("dist.FInferType", InferDistTypeBroadcastArith);
+ .set_attr<FInferTypeWithBuilder>("relax.dist.FInferTypeWithBuilder", InferDistTypeBroadcastArith);
/***************** Comparison operators *****************/
OpDef("relax.equal")
- .set_attr<FInferType>("dist.FInferType", InferDistTypeBroadcastCMP);
+ .set_attr<FInferTypeWithBuilder>("relax.dist.FInferTypeWithBuilder", InferDistTypeBroadcastCMP);
OpDef("relax.greater")
- .set_attr<FInferType>("dist.FInferType", InferDistTypeBroadcastCMP);
+ .set_attr<FInferTypeWithBuilder>("relax.dist.FInferTypeWithBuilder", InferDistTypeBroadcastCMP);
OpDef("relax.greater_equal")
- .set_attr<FInferType>("dist.FInferType", InferDistTypeBroadcastCMP);
+ .set_attr<FInferTypeWithBuilder>("relax.dist.FInferTypeWithBuilder", InferDistTypeBroadcastCMP);
OpDef("relax.less")
- .set_attr<FInferType>("dist.FInferType", InferDistTypeBroadcastCMP);
+ .set_attr<FInferTypeWithBuilder>("relax.dist.FInferTypeWithBuilder", InferDistTypeBroadcastCMP);
OpDef("relax.less_equal")
- .set_attr<FInferType>("dist.FInferType", InferDistTypeBroadcastCMP);
+ .set_attr<FInferTypeWithBuilder>("relax.dist.FInferTypeWithBuilder", InferDistTypeBroadcastCMP);
OpDef("relax.not_equal")
- .set_attr<FInferType>("dist.FInferType", InferDistTypeBroadcastCMP);
+ .set_attr<FInferTypeWithBuilder>("relax.dist.FInferTypeWithBuilder", InferDistTypeBroadcastCMP);
/***************** Min/Max operators *****************/
OpDef("relax.minimum")
- .set_attr<FInferType>("dist.FInferType", InferDistTypeBroadcastArith);
+ .set_attr<FInferTypeWithBuilder>("relax.dist.FInferTypeWithBuilder", InferDistTypeBroadcastArith);
OpDef("relax.maximum")
- .set_attr<FInferType>("dist.FInferType", InferDistTypeBroadcastArith);
+ .set_attr<FInferTypeWithBuilder>("relax.dist.FInferTypeWithBuilder", InferDistTypeBroadcastArith);
/***************** Logical operators *****************/
OpDef("relax.logical_and")
- .set_attr<FInferType>("dist.FInferType", InferDistTypeBroadcastArith);
+ .set_attr<FInferTypeWithBuilder>("relax.dist.FInferTypeWithBuilder", InferDistTypeBroadcastArith);
OpDef("relax.logical_or")
- .set_attr<FInferType>("dist.FInferType", InferDistTypeBroadcastArith);
+ .set_attr<FInferTypeWithBuilder>("relax.dist.FInferTypeWithBuilder", InferDistTypeBroadcastArith);
OpDef("relax.logical_xor")
- .set_attr<FInferType>("dist.FInferType", InferDistTypeBroadcastArith);
+ .set_attr<FInferTypeWithBuilder>("relax.dist.FInferTypeWithBuilder", InferDistTypeBroadcastArith);
/***************** Bitwise operators *****************/
OpDef("relax.bitwise_and")
- .set_attr<FInferType>("dist.FInferType", InferDistTypeBroadcastArith);
+ .set_attr<FInferTypeWithBuilder>("relax.dist.FInferTypeWithBuilder", InferDistTypeBroadcastArith);
OpDef("relax.bitwise_or")
- .set_attr<FInferType>("dist.FInferType", InferDistTypeBroadcastArith);
+ .set_attr<FInferTypeWithBuilder>("relax.dist.FInferTypeWithBuilder", InferDistTypeBroadcastArith);
OpDef("relax.bitwise_xor")
- .set_attr<FInferType>("dist.FInferType", InferDistTypeBroadcastArith);
+ .set_attr<FInferTypeWithBuilder>("relax.dist.FInferTypeWithBuilder", InferDistTypeBroadcastArith);
OpDef("relax.left_shift")
- .set_attr<FInferType>("dist.FInferType", InferDistTypeBroadcastArith);
+ .set_attr<FInferTypeWithBuilder>("relax.dist.FInferTypeWithBuilder", InferDistTypeBroadcastArith);
OpDef("relax.right_shift")
- .set_attr<FInferType>("dist.FInferType", InferDistTypeBroadcastArith);
+ .set_attr<FInferTypeWithBuilder>("relax.dist.FInferTypeWithBuilder", InferDistTypeBroadcastArith);
// clang-format on
}
diff --git a/src/relax/op/distributed/ccl.cc b/src/relax/op/distributed/ccl.cc
index 8ec459d25d..73a35335e1 100644
--- a/src/relax/op/distributed/ccl.cc
+++ b/src/relax/op/distributed/ccl.cc
@@ -40,7 +40,7 @@ Type InferDistTypeAllReduce(const Call& call, const BlockBuilder& ctx) {
TVM_FFI_STATIC_INIT_BLOCK() {
// clang-format off
OpDef("relax.ccl.allreduce")
- .set_attr<FInferType>("dist.FInferType", InferDistTypeAllReduce);
+ .set_attr<FInferTypeWithBuilder>("relax.dist.FInferTypeWithBuilder", InferDistTypeAllReduce);
// clang-format on
}
diff --git a/src/relax/op/distributed/distributed.cc b/src/relax/op/distributed/distributed.cc
index 591982204d..4aad954c69 100644
--- a/src/relax/op/distributed/distributed.cc
+++ b/src/relax/op/distributed/distributed.cc
@@ -65,8 +65,9 @@ Type InferTypeAnnotateSharding(const Call& call, const BlockBuilder& ctx) {
TVM_FFI_STATIC_INIT_BLOCK() {
OpDef("relax.dist.annotate_sharding")
.signature(sig::arg("input", "The input tensor."), sig::call_attrs<DistributionAttrs>())
- .set_attr<FInferType>("FInferType", InferTypeAnnotateSharding)
- .set_attr<FInferType>("dist.FInferType", InferTypeAnnotateSharding)
+ .set_attr<FInferTypeWithBuilder>("relax.FInferTypeWithBuilder", InferTypeAnnotateSharding)
+ .set_attr<FInferTypeWithBuilder>("relax.dist.FInferTypeWithBuilder",
+ InferTypeAnnotateSharding)
.set_attr<bool>("FPurity", true);
}
@@ -97,7 +98,8 @@ Type InferDistTypeRedistribute(const Call& call, const BlockBuilder& ctx) {
TVM_FFI_STATIC_INIT_BLOCK() {
OpDef("relax.dist.redistribute")
.signature(sig::arg("input", "The input tensor."), sig::call_attrs<DistributionAttrs>())
- .set_attr<FInferType>("dist.FInferType", InferDistTypeRedistribute)
+ .set_attr<FInferTypeWithBuilder>("relax.dist.FInferTypeWithBuilder",
+ InferDistTypeRedistribute)
.set_attr<bool>("FPurity", true);
}
@@ -117,7 +119,7 @@ TVM_FFI_STATIC_INIT_BLOCK() {
.signature(sig::arg("func", "The destination-passing-style function."),
sig::arg("args", "The input arguments."),
sig::ty_arg("out_type", "The output type."))
- .set_attr<FInferType>("FInferType", InferTypeCallTIRLocalView)
+ .set_attr<FInferTypeWithBuilder>("relax.FInferTypeWithBuilder", InferTypeCallTIRLocalView)
.set_attr<bool>("FPurity", true);
}
@@ -221,8 +223,8 @@ TVM_FFI_STATIC_INIT_BLOCK() {
OpDef("relax.dist.redistribute_replica_to_shard")
.signature(sig::arg("input", "The buffer to be sliced."),
sig::call_attrs<ScatterCollectiveAttrs>())
- .set_attr<FInferType>("FInferType", InferTypeRtoS)
- .set_attr<FInferType>("dist.FInferType", InferDistTypeRtoS)
+ .set_attr<FInferTypeWithBuilder>("relax.FInferTypeWithBuilder", InferTypeRtoS)
+ .set_attr<FInferTypeWithBuilder>("relax.dist.FInferTypeWithBuilder", InferDistTypeRtoS)
.set_attr<bool>("FPurity", true);
}
diff --git a/src/relax/op/distributed/linear_algebra.cc b/src/relax/op/distributed/linear_algebra.cc
index e7240c1bc0..0c54338d39 100644
--- a/src/relax/op/distributed/linear_algebra.cc
+++ b/src/relax/op/distributed/linear_algebra.cc
@@ -99,7 +99,7 @@ Type InferDistTypeMatmul(const Call& call, const BlockBuilder& ctx) {
TVM_FFI_STATIC_INIT_BLOCK() {
// clang-format off
OpDef("relax.matmul")
- .set_attr<FInferType>("dist.FInferType", InferDistTypeMatmul);
+ .set_attr<FInferTypeWithBuilder>("relax.dist.FInferTypeWithBuilder", InferDistTypeMatmul);
// clang-format on
}
diff --git a/src/relax/op/distributed/manipulate.cc b/src/relax/op/distributed/manipulate.cc
index 2febd5dd6d..0c90ecab1f 100644
--- a/src/relax/op/distributed/manipulate.cc
+++ b/src/relax/op/distributed/manipulate.cc
@@ -80,7 +80,7 @@ Type InferDistTypePermuteDims(const Call& call, const BlockBuilder& ctx) {
TVM_FFI_STATIC_INIT_BLOCK() {
// clang-format off
OpDef("relax.permute_dims")
- .set_attr<FInferType>("dist.FInferType", InferDistTypePermuteDims);
+ .set_attr<FInferTypeWithBuilder>("relax.dist.FInferTypeWithBuilder", InferDistTypePermuteDims);
// clang-format on
}
@@ -135,7 +135,7 @@ Type InferDistTypeReshape(const Call& call, const BlockBuilder& ctx) {
TVM_FFI_STATIC_INIT_BLOCK() {
// clang-format off
OpDef("relax.reshape")
- .set_attr<FInferType>("dist.FInferType", InferDistTypeReshape);
+ .set_attr<FInferTypeWithBuilder>("relax.dist.FInferTypeWithBuilder", InferDistTypeReshape);
// clang-format on
}
diff --git a/src/relax/op/distributed/nn.cc b/src/relax/op/distributed/nn.cc
index 48fac7c458..9d2324a760 100644
--- a/src/relax/op/distributed/nn.cc
+++ b/src/relax/op/distributed/nn.cc
@@ -49,7 +49,7 @@ Type InferDistTypeSoftmax(const Call& call, const BlockBuilder& ctx) {
TVM_FFI_STATIC_INIT_BLOCK() {
// clang-format off
OpDef("relax.nn.softmax")
- .set_attr<FInferType>("dist.FInferType", InferDistTypeSoftmax);
+ .set_attr<FInferTypeWithBuilder>("relax.dist.FInferTypeWithBuilder", InferDistTypeSoftmax);
RELAX_REGISTER_UNARY_ARITH_DIST_INFER_TYPE(nn.relu, /*require_float_dtype=*/false);
diff --git a/src/relax/op/distributed/op.cc b/src/relax/op/distributed/op.cc
index 697b7fee5d..a7cc19e0cf 100644
--- a/src/relax/op/distributed/op.cc
+++ b/src/relax/op/distributed/op.cc
@@ -38,7 +38,7 @@ Type InferDistTypeCallTIR(const Call& call, const BlockBuilder& ctx) {
TVM_FFI_STATIC_INIT_BLOCK() {
// clang-format off
OpDef("relax.call_tir")
- .set_attr<FInferType>("dist.FInferType", InferDistTypeCallTIR);
+ .set_attr<FInferTypeWithBuilder>("relax.dist.FInferTypeWithBuilder", InferDistTypeCallTIR);
// clang-format on
}
@@ -51,7 +51,8 @@ Type InferDistTypeStopLiftParams(const Call& call, const BlockBuilder& ctx) {
TVM_FFI_STATIC_INIT_BLOCK() {
OpDef("relax.builtin.stop_lift_params")
- .set_attr<FInferType>("dist.FInferType", InferDistTypeStopLiftParams);
+ .set_attr<FInferTypeWithBuilder>("relax.dist.FInferTypeWithBuilder",
+ InferDistTypeStopLiftParams);
}
} // namespace distributed
diff --git a/src/relax/op/distributed/statistical.h b/src/relax/op/distributed/statistical.h
index 7d44e1532e..dc4cfbcabc 100644
--- a/src/relax/op/distributed/statistical.h
+++ b/src/relax/op/distributed/statistical.h
@@ -36,7 +36,7 @@ Type InferDistTypeStatistical(const Call& call, const BlockBuilder& ctx);
// clang-format off
#define RELAX_REGISTER_STATISTICAL_DIST_INFER_TYPE(OpName) \
OpDef("relax." #OpName) \
- .set_attr<FInferType>("dist.FInferType", InferDistTypeStatistical)
+ .set_attr<FInferTypeWithBuilder>("relax.dist.FInferTypeWithBuilder", InferDistTypeStatistical)
// clang-format on
} // namespace distributed
diff --git a/src/relax/op/distributed/unary.h b/src/relax/op/distributed/unary.h
index 4e2bb8f64e..325a5462d2 100644
--- a/src/relax/op/distributed/unary.h
+++ b/src/relax/op/distributed/unary.h
@@ -67,11 +67,11 @@ Type InferDistTypeUnaryCheck(const Call& call, const BlockBuilder& ctx);
// clang-format off
#define RELAX_REGISTER_UNARY_ARITH_DIST_INFER_TYPE(OpName, RequireFloatDtype) \
OpDef("relax." #OpName) \
- .set_attr<FInferType>("dist.FInferType", InferDistTypeUnaryArith<RequireFloatDtype>)
+ .set_attr<FInferTypeWithBuilder>("relax.dist.FInferTypeWithBuilder", InferDistTypeUnaryArith<RequireFloatDtype>)
#define RELAX_REGISTER_UNARY_CHECK_DIST_INFER_TYPE(OpName) \
OpDef("relax." #OpName) \
- .set_attr<FInferType>("dist.FInferType", InferDistTypeUnaryCheck)
+ .set_attr<FInferTypeWithBuilder>("relax.dist.FInferTypeWithBuilder", InferDistTypeUnaryCheck)
// clang-format on
} // namespace distributed
diff --git a/src/relax/op/image/resize.cc b/src/relax/op/image/resize.cc
index 3a23426bd6..8aa43ab971 100644
--- a/src/relax/op/image/resize.cc
+++ b/src/relax/op/image/resize.cc
@@ -64,7 +64,8 @@ TVM_FFI_STATIC_INIT_BLOCK() {
refl::GlobalDef().def("relax.op.image.resize2d", resize2d);
}
-Type InferTypeResize2D(const Call& call, const BlockBuilder& ctx) {
+ffi::Expected<Type> InferTypeResize2D(const CallNode* call_node) noexcept try {
+ const Call call = ffi::GetRef<Call>(call_node);
if (call->args.size() != 2) {
TVM_FFI_VISIT_THROW(ValueError, call)
<< "Resize2D expects 2 arguments, while the given number of arguments is "
@@ -91,8 +92,8 @@ Type InferTypeResize2D(const Call& call, const BlockBuilder& ctx) {
}
const auto* attrs = call->attrs.as<Resize2DAttrs>();
- auto [data_layout, data2NCHW] = CheckTensorLayout(call, ctx, attrs->layout, //
- /*tgt_layout=*/"NCHW", //
+ auto [data_layout, data2NCHW] = CheckTensorLayout(call, attrs->layout, //
+ /*tgt_layout=*/"NCHW", //
/*tensor_name=*/"data");
ffi::Optional<PrimType> out_dtype =
@@ -100,7 +101,7 @@ Type InferTypeResize2D(const Call& call, const BlockBuilder& ctx) {
: data_ty->dtype;
ffi::Optional<ShapeExpr> data_shape =
- CheckNdimPerLayoutAndGetShape(call, ctx, ffi::GetRef<TensorType>(data_ty), data_layout);
+ CheckNdimPerLayoutAndGetShape(call, ffi::GetRef<TensorType>(data_ty), data_layout);
if (!data_shape.has_value() || size_value == nullptr) {
return TensorType(out_dtype, data_layout.ndim(), data_ty->vdevice);
}
@@ -112,6 +113,10 @@ Type InferTypeResize2D(const Call& call, const BlockBuilder& ctx) {
ffi::Array<PrimExpr> out_shape = data2NCHW.BackwardShape(out_NCHW_shape);
return TensorType(ShapeExpr(out_shape), out_dtype, data_ty->vdevice);
+} catch (const ffi::Error& error) {
+ return ffi::Unexpected(error);
+} catch (const std::exception& error) {
+ return ffi::Unexpected(ffi::Error("InternalError", error.what(), ""));
}
InferLayoutOutput InferLayoutResize2d(
@@ -148,7 +153,7 @@ TVM_FFI_STATIC_INIT_BLOCK() {
OpDef("relax.image.resize2d")
.signature(sig::arg("data", "The input tensor."), sig::arg("size", "The output image shape."),
sig::call_attrs<Resize2DAttrs>())
- .set_attr<FInferType>("FInferType", InferTypeResize2D)
+ .set_attr<FInferType>("FInferType", FInferType::FromNative<&InferTypeResize2D>())
.set_attr<FRelaxInferLayout>("FRelaxInferLayout", InferLayoutResize2d)
.set_attr<TMixedPrecisionPolicy>("TMixedPrecisionPolicy", MixedPrecisionPolicyKind::kFollow)
.set_attr<bool>("FPurity", true);
@@ -180,7 +185,8 @@ TVM_FFI_STATIC_INIT_BLOCK() {
refl::GlobalDef().def("relax.op.image.resize3d", resize3d);
}
-Type InferTypeResize3D(const Call& call, const BlockBuilder& ctx) {
+ffi::Expected<Type> InferTypeResize3D(const CallNode* call_node) noexcept try {
+ const Call call = ffi::GetRef<Call>(call_node);
if (call->args.size() != 2) {
TVM_FFI_VISIT_THROW(ValueError, call)
<< "Resize3D expects 2 arguments, while the given number of arguments is "
@@ -207,8 +213,8 @@ Type InferTypeResize3D(const Call& call, const BlockBuilder& ctx) {
}
const auto* attrs = call->attrs.as<Resize3DAttrs>();
- auto [data_layout, data2NCDHW] = CheckTensorLayout(call, ctx, attrs->layout, //
- /*tgt_layout=*/"NCDHW", //
+ auto [data_layout, data2NCDHW] = CheckTensorLayout(call, attrs->layout, //
+ /*tgt_layout=*/"NCDHW", //
/*tensor_name=*/"data");
ffi::Optional<PrimType> out_dtype =
@@ -216,7 +222,7 @@ Type InferTypeResize3D(const Call& call, const BlockBuilder& ctx) {
: data_ty->dtype;
ffi::Optional<ShapeExpr> data_shape =
- CheckNdimPerLayoutAndGetShape(call, ctx, ffi::GetRef<TensorType>(data_ty), data_layout);
+ CheckNdimPerLayoutAndGetShape(call, ffi::GetRef<TensorType>(data_ty), data_layout);
if (!data_shape.has_value() || size_value == nullptr) {
return TensorType(out_dtype, data_layout.ndim(), data_ty->vdevice);
}
@@ -229,6 +235,10 @@ Type InferTypeResize3D(const Call& call, const BlockBuilder& ctx) {
ffi::Array<PrimExpr> out_shape = data2NCDHW.BackwardShape(out_NCDHW_shape);
return TensorType(ShapeExpr(out_shape), out_dtype, data_ty->vdevice);
+} catch (const ffi::Error& error) {
+ return ffi::Unexpected(error);
+} catch (const std::exception& error) {
+ return ffi::Unexpected(ffi::Error("InternalError", error.what(), ""));
}
InferLayoutOutput InferLayoutResize3d(
@@ -262,7 +272,7 @@ TVM_FFI_STATIC_INIT_BLOCK() {
OpDef("relax.image.resize3d")
.signature(sig::arg("data", "The input tensor."), sig::arg("size", "The output image shape."),
sig::call_attrs<Resize3DAttrs>())
- .set_attr<FInferType>("FInferType", InferTypeResize3D)
+ .set_attr<FInferType>("FInferType", FInferType::FromNative<&InferTypeResize3D>())
.set_attr<FRelaxInferLayout>("FRelaxInferLayout", InferLayoutResize3d)
.set_attr<TMixedPrecisionPolicy>("TMixedPrecisionPolicy", MixedPrecisionPolicyKind::kFollow)
.set_attr<bool>("FPurity", true);
@@ -289,7 +299,8 @@ TVM_FFI_STATIC_INIT_BLOCK() {
refl::GlobalDef().def("relax.op.image.grid_sample", grid_sample);
}
-Type InferTypeGridSample(const Call& call, const BlockBuilder& ctx) {
+ffi::Expected<Type> InferTypeGridSample(const CallNode* call_node) noexcept try {
+ const Call call = ffi::GetRef<Call>(call_node);
if (call->args.size() != 2) {
TVM_FFI_VISIT_THROW(ValueError, call)
<< "GridSample expects two arguments, while the given number of arguments is "
@@ -317,14 +328,14 @@ Type InferTypeGridSample(const Call& call, const BlockBuilder& ctx) {
// treated as the 2D NCHW path so existing behavior is preserved.
const bool is_ncdhw = (attrs->layout == "NCDHW");
- auto [data_layout, data2tgt] = CheckTensorLayout(call, ctx, attrs->layout,
+ auto [data_layout, data2tgt] = CheckTensorLayout(call, attrs->layout,
/*tgt_layout=*/is_ncdhw ? "NCDHW" : "NCHW",
/*tensor_name=*/"data");
ffi::Optional<PrimType> out_dtype = data_ty->dtype;
ffi::Optional<ShapeExpr> data_shape =
- CheckNdimPerLayoutAndGetShape(call, ctx, ffi::GetRef<TensorType>(data_ty), data_layout);
+ CheckNdimPerLayoutAndGetShape(call, ffi::GetRef<TensorType>(data_ty), data_layout);
const auto* grid_shape = grid_ty->shape.as<ShapeExprNode>();
if (!data_shape.has_value() || grid_shape == nullptr) {
@@ -347,6 +358,10 @@ Type InferTypeGridSample(const Call& call, const BlockBuilder& ctx) {
ffi::Array<PrimExpr> out_shape = data2tgt.BackwardShape(out_tgt_shape);
return TensorType(ShapeExpr(out_shape), out_dtype, data_ty->vdevice);
+} catch (const ffi::Error& error) {
+ return ffi::Unexpected(error);
+} catch (const std::exception& error) {
+ return ffi::Unexpected(ffi::Error("InternalError", error.what(), ""));
}
TVM_FFI_STATIC_INIT_BLOCK() {
@@ -354,7 +369,7 @@ TVM_FFI_STATIC_INIT_BLOCK() {
.signature(sig::arg("data", "The input tensor."),
sig::arg("grid", "The grid tensor for sampling."),
sig::call_attrs<GridSampleAttrs>())
- .set_attr<FInferType>("FInferType", InferTypeGridSample)
+ .set_attr<FInferType>("FInferType", FInferType::FromNative<&InferTypeGridSample>())
.set_attr<TMixedPrecisionPolicy>("TMixedPrecisionPolicy", MixedPrecisionPolicyKind::kFollow)
.set_attr<bool>("FPurity", true);
}
@@ -375,7 +390,8 @@ TVM_FFI_STATIC_INIT_BLOCK() {
refl::GlobalDef().def("relax.op.image.affine_grid", affine_grid);
}
-Type InferTypeAffineGrid(const Call& call, const BlockBuilder& ctx) {
+ffi::Expected<Type> InferTypeAffineGrid(const CallNode* call_node) noexcept try {
+ const Call call = ffi::GetRef<Call>(call_node);
if (call->args.size() != 2) {
TVM_FFI_VISIT_THROW(ValueError, call)
<< "AffineGrid expects two arguments, while the given number of arguments is "
@@ -447,6 +463,10 @@ Type InferTypeAffineGrid(const Call& call, const BlockBuilder& ctx) {
}
return TensorType(ShapeExpr(out_shape), out_dtype, data_ty->vdevice);
+} catch (const ffi::Error& error) {
+ return ffi::Unexpected(error);
+} catch (const std::exception& error) {
+ return ffi::Unexpected(ffi::Error("InternalError", error.what(), ""));
}
TVM_FFI_STATIC_INIT_BLOCK() {
@@ -454,7 +474,7 @@ TVM_FFI_STATIC_INIT_BLOCK() {
.signature(sig::arg("data", "The input affine matrix tensor."),
sig::arg("size", "The target output shape (H, W)."),
sig::call_attrs<AffineGridAttrs>())
- .set_attr<FInferType>("FInferType", InferTypeAffineGrid)
+ .set_attr<FInferType>("FInferType", FInferType::FromNative<&InferTypeAffineGrid>())
.set_attr<TMixedPrecisionPolicy>("TMixedPrecisionPolicy", MixedPrecisionPolicyKind::kFollow)
.set_attr<bool>("FPurity", true);
}
diff --git a/src/relax/op/memory/view.cc b/src/relax/op/memory/view.cc
index af12fbd4b1..7b481da192 100644
--- a/src/relax/op/memory/view.cc
+++ b/src/relax/op/memory/view.cc
@@ -400,7 +400,7 @@ TVM_FFI_STATIC_INIT_BLOCK() {
sig::arg("relative_byte_offset",
"The view's byte offset, relative to the input tensor's byte offset."))
.set_attr<bool>("RequiresArgumentShapes", false)
- .set_attr<FInferType>("FInferType", InferTypeView)
+ .set_attr<FInferTypeWithBuilder>("relax.FInferTypeWithBuilder", InferTypeView)
.set_attr<bool>("FPurity", true)
.set_attr<FLowerBuiltin>("FLowerBuiltin", LowerBuiltinView);
}
@@ -415,13 +415,18 @@ TVM_FFI_STATIC_INIT_BLOCK() {
refl::GlobalDef().def("relax.op.memory.ensure_zero_offset", ensure_zero_offset);
}
-Type InferTypeEnsureZeroOffset(const Call& call, const BlockBuilder& ctx) {
+ffi::Expected<Type> InferTypeEnsureZeroOffset(const CallNode* call_node) noexcept try {
+ const Call call = ffi::GetRef<Call>(call_node);
if (call->args.size() != 1) {
TVM_FFI_VISIT_THROW(ValueError, call)
<< "Operator " << call->op << " should receive 1 argument, "
<< "but received " << call->args;
}
return GetType(call->args[0]);
+} catch (const ffi::Error& error) {
+ return ffi::Unexpected(error);
+} catch (const std::exception& error) {
+ return ffi::Unexpected(ffi::Error("InternalError", error.what(), ""));
}
Expr LowerBuiltinEnsureZeroOffset(const BlockBuilder& bb, const Call& call) {
@@ -434,7 +439,7 @@ TVM_FFI_STATIC_INIT_BLOCK() {
OpDef("relax.memory.ensure_zero_offset")
.signature(sig::arg("x", "The input tensor."))
.set_attr<bool>("RequiresArgumentShapes", false)
- .set_attr<FInferType>("FInferType", InferTypeEnsureZeroOffset)
+ .set_attr<FInferType>("FInferType", FInferType::FromNative<&InferTypeEnsureZeroOffset>())
.set_attr<bool>("FPurity", true)
.set_attr<FLowerBuiltin>("FLowerBuiltin", LowerBuiltinEnsureZeroOffset);
}
diff --git a/src/relax/op/nn/attention.cc b/src/relax/op/nn/attention.cc
index bdab2755ca..4a6b8796e7 100644
--- a/src/relax/op/nn/attention.cc
+++ b/src/relax/op/nn/attention.cc
@@ -156,7 +156,7 @@ TVM_FFI_STATIC_INIT_BLOCK() {
sig::arg("value", "The input values tensor."), sig::call_attrs<AttentionAttrs>())
.set_attr<TMixedPrecisionPolicy>("TMixedPrecisionPolicy", MixedPrecisionPolicyKind::kAlways)
.set_attr<FInferMixedPrecision>("FInferMixedPrecision", InferMixedPrecisionAttention)
- .set_attr<FInferType>("FInferType", InferTypeAttention)
+ .set_attr<FInferTypeWithBuilder>("relax.FInferTypeWithBuilder", InferTypeAttention)
.set_attr<bool>("FPurity", true);
OpDef("relax.nn.attention_bias")
@@ -166,7 +166,7 @@ TVM_FFI_STATIC_INIT_BLOCK() {
sig::arg("bias", "The input bias tensor."), sig::call_attrs<AttentionAttrs>())
.set_attr<TMixedPrecisionPolicy>("TMixedPrecisionPolicy", MixedPrecisionPolicyKind::kAlways)
.set_attr<FInferMixedPrecision>("FInferMixedPrecision", InferMixedPrecisionAttention)
- .set_attr<FInferType>("FInferType", InferTypeAttention)
+ .set_attr<FInferTypeWithBuilder>("relax.FInferTypeWithBuilder", InferTypeAttention)
.set_attr<bool>("FPurity", true);
OpDef("relax.nn.attention_var_len")
@@ -180,7 +180,7 @@ TVM_FFI_STATIC_INIT_BLOCK() {
sig::call_attrs<AttentionAttrs>())
.set_attr<TMixedPrecisionPolicy>("TMixedPrecisionPolicy", MixedPrecisionPolicyKind::kAlways)
.set_attr<FInferMixedPrecision>("FInferMixedPrecision", InferMixedPrecisionAttention)
- .set_attr<FInferType>("FInferType", InferTypeAttention)
+ .set_attr<FInferTypeWithBuilder>("relax.FInferTypeWithBuilder", InferTypeAttention)
.set_attr<bool>("FPurity", true);
AttentionAttrs::RegisterReflection();
diff --git a/src/relax/op/nn/convolution.cc b/src/relax/op/nn/convolution.cc
index 2c5f0373fd..852452d814 100644
--- a/src/relax/op/nn/convolution.cc
+++ b/src/relax/op/nn/convolution.cc
@@ -202,7 +202,7 @@ TVM_FFI_STATIC_INIT_BLOCK() {
sig::arg("data", "The input tensor."), sig::arg("weight", "The weight tensor."),
sig::var_ty_args("out_type", "Optional output tensor type carrying the virtual device."),
sig::call_attrs<Conv1DAttrs>())
- .set_attr<FInferType>("FInferType", InferTypeConv1d)
+ .set_attr<FInferTypeWithBuilder>("relax.FInferTypeWithBuilder", InferTypeConv1d)
.set_attr<FRelaxInferLayout>("FRelaxInferLayout", InferLayoutConv1d)
.set_attr<TMixedPrecisionPolicy>("TMixedPrecisionPolicy", MixedPrecisionPolicyKind::kAlways)
.set_attr<FInferMixedPrecision>("FInferMixedPrecision", InferMixedPrecisionConv1d)
@@ -416,7 +416,7 @@ TVM_FFI_STATIC_INIT_BLOCK() {
sig::arg("data", "The input tensor."), sig::arg("weight", "The weight tensor."),
sig::var_ty_args("out_type", "Optional output tensor type carrying the virtual device."),
sig::call_attrs<Conv2DAttrs>())
- .set_attr<FInferType>("FInferType", InferTypeConv2d)
+ .set_attr<FInferTypeWithBuilder>("relax.FInferTypeWithBuilder", InferTypeConv2d)
.set_attr<FRelaxInferLayout>("FRelaxInferLayout", InferLayoutConv2d)
.set_attr<TMixedPrecisionPolicy>("TMixedPrecisionPolicy", MixedPrecisionPolicyKind::kAlways)
.set_attr<FInferMixedPrecision>("FInferMixedPrecision", InferMixedPrecisionConv2d)
@@ -604,7 +604,7 @@ TVM_FFI_STATIC_INIT_BLOCK() {
sig::arg("data", "The input tensor."), sig::arg("weight", "The weight tensor."),
sig::var_ty_args("out_type", "Optional output tensor type carrying the virtual device."),
sig::call_attrs<Conv3DAttrs>())
- .set_attr<FInferType>("FInferType", InferTypeConv3d)
+ .set_attr<FInferTypeWithBuilder>("relax.FInferTypeWithBuilder", InferTypeConv3d)
.set_attr<FRelaxInferLayout>("FRelaxInferLayout", InferLayoutConv3d)
.set_attr<TMixedPrecisionPolicy>("TMixedPrecisionPolicy", MixedPrecisionPolicyKind::kAlways)
.set_attr<FInferMixedPrecision>("FInferMixedPrecision", InferMixedPrecisionConv3d)
@@ -786,7 +786,7 @@ TVM_FFI_STATIC_INIT_BLOCK() {
sig::arg("data", "The input tensor."), sig::arg("weight", "The weight tensor."),
sig::var_ty_args("out_type", "Optional output tensor type carrying the virtual device."),
sig::call_attrs<Conv1DTransposeAttrs>())
- .set_attr<FInferType>("FInferType", InferTypeConv1dTranspose)
+ .set_attr<FInferTypeWithBuilder>("relax.FInferTypeWithBuilder", InferTypeConv1dTranspose)
.set_attr<FRelaxInferLayout>("FRelaxInferLayout", InferLayoutConv1dTranspose)
.set_attr<TMixedPrecisionPolicy>("TMixedPrecisionPolicy", MixedPrecisionPolicyKind::kAlways)
.set_attr<FInferMixedPrecision>("FInferMixedPrecision", InferMixedPrecisionConv1dTranspose)
@@ -1017,7 +1017,7 @@ TVM_FFI_STATIC_INIT_BLOCK() {
sig::arg("data", "The input tensor."), sig::arg("weight", "The weight tensor."),
sig::var_ty_args("out_type", "Optional output tensor type carrying the virtual device."),
sig::call_attrs<Conv2DTransposeAttrs>())
- .set_attr<FInferType>("FInferType", InferTypeConv2dTranspose)
+ .set_attr<FInferTypeWithBuilder>("relax.FInferTypeWithBuilder", InferTypeConv2dTranspose)
.set_attr<FRelaxInferLayout>("FRelaxInferLayout", InferLayoutConv2dTranspose)
.set_attr<TMixedPrecisionPolicy>("TMixedPrecisionPolicy", MixedPrecisionPolicyKind::kAlways)
.set_attr<FInferMixedPrecision>("FInferMixedPrecision", InferMixedPrecisionConv2dTranspose)
@@ -1259,7 +1259,7 @@ TVM_FFI_STATIC_INIT_BLOCK() {
sig::arg("data", "The input tensor."), sig::arg("weight", "The weight tensor."),
sig::var_ty_args("out_type", "Optional output tensor type carrying the virtual device."),
sig::call_attrs<Conv3DTransposeAttrs>())
- .set_attr<FInferType>("FInferType", InferTypeConv3dTranspose)
+ .set_attr<FInferTypeWithBuilder>("relax.FInferTypeWithBuilder", InferTypeConv3dTranspose)
.set_attr<FRelaxInferLayout>("FRelaxInferLayout", InferLayoutConv3dTranspose)
.set_attr<TMixedPrecisionPolicy>("TMixedPrecisionPolicy", MixedPrecisionPolicyKind::kAlways)
.set_attr<FInferMixedPrecision>("FInferMixedPrecision", InferMixedPrecisionConv3dTranspose)
diff --git a/src/relax/op/nn/nn.cc b/src/relax/op/nn/nn.cc
index f814815033..0f4b5374bc 100644
--- a/src/relax/op/nn/nn.cc
+++ b/src/relax/op/nn/nn.cc
@@ -81,7 +81,8 @@ TVM_FFI_STATIC_INIT_BLOCK() {
.set_attr<FRelaxInferLayout>("FRelaxInferLayout", InferLayoutUnaryEwise)
.set_attr<TMixedPrecisionPolicy>("TMixedPrecisionPolicy", MixedPrecisionPolicyKind::kFollow)
.set_attr<bool>("FPurity", true)
- .set_attr<FInferType>("FInferType", InferTypeUnaryArith<false>);
+ .set_attr<FInferType>("FInferType",
+ FInferType::FromNative<&InferTypeUnaryArithContextFree<false>>());
/* relax.nn.gelu */
@@ -94,7 +95,8 @@ TVM_FFI_STATIC_INIT_BLOCK() {
.set_attr<FRelaxInferLayout>("FRelaxInferLayout", InferLayoutUnaryEwise)
.set_attr<TMixedPrecisionPolicy>("TMixedPrecisionPolicy", MixedPrecisionPolicyKind::kFollow)
.set_attr<bool>("FPurity", true)
- .set_attr<FInferType>("FInferType", InferTypeUnaryArith<true>);
+ .set_attr<FInferType>("FInferType",
+ FInferType::FromNative<&InferTypeUnaryArithContextFree<true>>());
/* relax.nn.gelu_tanh */
@@ -107,7 +109,8 @@ TVM_FFI_STATIC_INIT_BLOCK() {
.set_attr<FRelaxInferLayout>("FRelaxInferLayout", InferLayoutUnaryEwise)
.set_attr<TMixedPrecisionPolicy>("TMixedPrecisionPolicy", MixedPrecisionPolicyKind::kFollow)
.set_attr<bool>("FPurity", true)
- .set_attr<FInferType>("FInferType", InferTypeUnaryArith<true>);
+ .set_attr<FInferType>("FInferType",
+ FInferType::FromNative<&InferTypeUnaryArithContextFree<true>>());
/* relax.nn.selu */
@@ -120,7 +123,8 @@ TVM_FFI_STATIC_INIT_BLOCK() {
.set_attr<FRelaxInferLayout>("FRelaxInferLayout", InferLayoutUnaryEwise)
.set_attr<TMixedPrecisionPolicy>("TMixedPrecisionPolicy", MixedPrecisionPolicyKind::kFollow)
.set_attr<bool>("FPurity", true)
- .set_attr<FInferType>("FInferType", InferTypeUnaryArith<true>);
+ .set_attr<FInferType>("FInferType",
+ FInferType::FromNative<&InferTypeUnaryArithContextFree<true>>());
/* relax.nn.silu */
@@ -133,7 +137,8 @@ TVM_FFI_STATIC_INIT_BLOCK() {
.set_attr<FRelaxInferLayout>("FRelaxInferLayout", InferLayoutUnaryEwise)
.set_attr<TMixedPrecisionPolicy>("TMixedPrecisionPolicy", MixedPrecisionPolicyKind::kFollow)
.set_attr<bool>("FPurity", true)
- .set_attr<FInferType>("FInferType", InferTypeUnaryArith<true>);
+ .set_attr<FInferType>("FInferType",
+ FInferType::FromNative<&InferTypeUnaryArithContextFree<true>>());
}
/* relax.nn.leakyrelu */
@@ -154,7 +159,9 @@ TVM_FFI_STATIC_INIT_BLOCK() {
sig::arg("data", "The input tensor."),
sig::var_ty_args("out_type", "Optional output tensor type carrying the virtual device."),
sig::call_attrs<LeakyReluAttrs>())
- .set_attr<FInferType>("FInferType", InferTypeUnaryArith</*require_float_dtype=*/true>)
+ .set_attr<FInferType>(
+ "FInferType",
+ FInferType::FromNative<&InferTypeUnaryArithContextFree</*require_float_dtype=*/true>>())
.set_attr<bool>("FPurity", true);
}
@@ -177,7 +184,9 @@ TVM_FFI_STATIC_INIT_BLOCK() {
sig::arg("data", "The input tensor."),
sig::var_ty_args("out_type", "Optional output tensor type carrying the virtual device."),
sig::call_attrs<SoftplusAttrs>())
- .set_attr<FInferType>("FInferType", InferTypeUnaryArith</*require_float_dtype=*/true>)
+ .set_attr<FInferType>(
+ "FInferType",
+ FInferType::FromNative<&InferTypeUnaryArithContextFree</*require_float_dtype=*/true>>())
.set_attr<bool>("FPurity", true);
}
@@ -195,8 +204,9 @@ TVM_FFI_STATIC_INIT_BLOCK() {
refl::GlobalDef().def("relax.op.nn.prelu", prelu);
}
-Type InferTypePRelu(const Call& call, const BlockBuilder& ctx) {
- TensorType data_ty = GetUnaryInputTensorType(call, ctx);
+ffi::Expected<Type> InferTypePRelu(const CallNode* call_node) noexcept try {
+ const Call call = ffi::GetRef<Call>(call_node);
+ TensorType data_ty = GetUnaryInputTensorType(call);
if (data_ty->IsUnknownNdim()) {
return data_ty;
}
@@ -207,9 +217,13 @@ Type InferTypePRelu(const Call& call, const BlockBuilder& ctx) {
<< data_ty->dtype;
}
const auto* attrs = call->attrs.as<PReluAttrs>();
- NormalizeAxis(call, ctx, data_ty->ndim, attrs->axis);
+ NormalizeAxis(call, data_ty->ndim, attrs->axis);
return data_ty;
+} catch (const ffi::Error& error) {
+ return ffi::Unexpected(error);
+} catch (const std::exception& error) {
+ return ffi::Unexpected(ffi::Error("InternalError", error.what(), ""));
}
InferLayoutOutput InferLayoutPRelu(
@@ -242,7 +256,7 @@ TVM_FFI_STATIC_INIT_BLOCK() {
.signature(sig::arg("data", "The input tensor."),
sig::arg("alpha", "The channel-wise learnable slope."),
sig::call_attrs<PReluAttrs>())
- .set_attr<FInferType>("FInferType", InferTypePRelu)
+ .set_attr<FInferType>("FInferType", FInferType::FromNative<&InferTypePRelu>())
.set_attr<FRelaxInferLayout>("FRelaxInferLayout", InferLayoutPRelu)
.set_attr<bool>("FPurity", true);
}
@@ -261,8 +275,9 @@ TVM_FFI_STATIC_INIT_BLOCK() {
refl::GlobalDef().def("relax.op.nn.softmax", softmax);
}
-Type InferTypeSoftmax(const Call& call, const BlockBuilder& ctx) {
- TensorType data_ty = GetUnaryInputTensorType(call, ctx);
+ffi::Expected<Type> InferTypeSoftmax(const CallNode* call_node) noexcept try {
+ const Call call = ffi::GetRef<Call>(call_node);
+ TensorType data_ty = GetUnaryInputTensorType(call);
if (data_ty->IsUnknownNdim()) {
return data_ty;
}
@@ -276,9 +291,13 @@ Type InferTypeSoftmax(const Call& call, const BlockBuilder& ctx) {
}
}
const auto* attrs = call->attrs.as<SoftmaxAttrs>();
- NormalizeAxis(call, ctx, data_ty->ndim, attrs->axis);
+ NormalizeAxis(call, data_ty->ndim, attrs->axis);
return data_ty;
+} catch (const ffi::Error& error) {
+ return ffi::Unexpected(error);
+} catch (const std::exception& error) {
+ return ffi::Unexpected(ffi::Error("InternalError", error.what(), ""));
}
InferLayoutOutput InferLayoutSoftmax(
@@ -307,7 +326,7 @@ InferLayoutOutput InferLayoutSoftmax(
TVM_FFI_STATIC_INIT_BLOCK() {
OpDef("relax.nn.softmax")
.signature(sig::arg("data", "The input tensor."), sig::call_attrs<SoftmaxAttrs>())
- .set_attr<FInferType>("FInferType", InferTypeSoftmax)
+ .set_attr<FInferType>("FInferType", FInferType::FromNative<&InferTypeSoftmax>())
.set_attr<FRelaxInferLayout>("FRelaxInferLayout", InferLayoutSoftmax)
.set_attr<bool>("FPurity", true);
}
@@ -326,7 +345,7 @@ TVM_FFI_STATIC_INIT_BLOCK() {
OpDef("relax.nn.log_softmax")
.signature(sig::arg("data", "The input tensor."), sig::call_attrs<SoftmaxAttrs>())
- .set_attr<FInferType>("FInferType", InferTypeSoftmax)
+ .set_attr<FInferType>("FInferType", FInferType::FromNative<&InferTypeSoftmax>())
.set_attr<bool>("FPurity", true);
}
@@ -346,8 +365,9 @@ TVM_FFI_STATIC_INIT_BLOCK() {
refl::GlobalDef().def("relax.op.nn.pad", pad);
}
-Type InferTypePad(const Call& call, const BlockBuilder& ctx) {
- ffi::Array<TensorType> input_ty = GetInputTensorType(call, ctx);
+ffi::Expected<Type> InferTypePad(const CallNode* call_node) noexcept try {
+ const Call call = ffi::GetRef<Call>(call_node);
+ ffi::Array<TensorType> input_ty = GetInputTensorType(call);
const auto* attrs = call->attrs.as<PadAttrs>();
int ndim = input_ty[0]->ndim;
ffi::Array<int64_t> pad_width = attrs->pad_width;
@@ -368,12 +388,16 @@ Type InferTypePad(const Call& call, const BlockBuilder& ctx) {
return TensorType(input_ty[0]->dtype, ndim);
}
return TensorType(ShapeExpr(out_shape), input_ty[0]->dtype);
+} catch (const ffi::Error& error) {
+ return ffi::Unexpected(error);
+} catch (const std::exception& error) {
+ return ffi::Unexpected(ffi::Error("InternalError", error.what(), ""));
}
TVM_FFI_STATIC_INIT_BLOCK() {
OpDef("relax.nn.pad")
.signature(sig::arg("data", "The input tensor."), sig::call_attrs<PadAttrs>())
- .set_attr<FInferType>("FInferType", InferTypePad)
+ .set_attr<FInferType>("FInferType", FInferType::FromNative<&InferTypePad>())
.set_attr<bool>("FPurity", true);
}
@@ -391,8 +415,9 @@ TVM_FFI_STATIC_INIT_BLOCK() {
refl::GlobalDef().def("relax.op.nn.pixel_shuffle", pixel_shuffle);
}
-Type InferTypePixelShuffle(const Call& call, const BlockBuilder& ctx) {
- ffi::Array<TensorType> input_ty = GetInputTensorType(call, ctx);
+ffi::Expected<Type> InferTypePixelShuffle(const CallNode* call_node) noexcept try {
+ const Call call = ffi::GetRef<Call>(call_node);
+ ffi::Array<TensorType> input_ty = GetInputTensorType(call);
const auto* attrs = call->attrs.as<PixelShuffleAttrs>();
int r = attrs->upscale_factor;
TVM_FFI_ICHECK_GT(r, 0) << "Upscale factor must be positive";
@@ -440,12 +465,16 @@ Type InferTypePixelShuffle(const Call& call, const BlockBuilder& ctx) {
}
return TensorType(ShapeExpr(out_shape), input->dtype);
+} catch (const ffi::Error& error) {
+ return ffi::Unexpected(error);
+} catch (const std::exception& error) {
+ return ffi::Unexpected(ffi::Error("InternalError", error.what(), ""));
}
TVM_FFI_STATIC_INIT_BLOCK() {
OpDef("relax.nn.pixel_shuffle")
.signature(sig::arg("data", "The input tensor."), sig::call_attrs<PixelShuffleAttrs>())
- .set_attr<FInferType>("FInferType", InferTypePixelShuffle)
+ .set_attr<FInferType>("FInferType", FInferType::FromNative<&InferTypePixelShuffle>())
.set_attr<bool>("FPurity", true);
}
@@ -596,7 +625,7 @@ TVM_FFI_STATIC_INIT_BLOCK() {
sig::arg("gamma", "The gamma scale factor."), sig::arg("beta", "The beta offset factor."),
sig::arg("moving_mean", "Running mean of input."),
sig::arg("moving_var", "Running variance of input."), sig::call_attrs<BatchNormAttrs>())
- .set_attr<FInferType>("FInferType", InferTypeBatchNorm)
+ .set_attr<FInferTypeWithBuilder>("relax.FInferTypeWithBuilder", InferTypeBatchNorm)
.set_attr<FRelaxInferLayout>("FRelaxInferLayout", InferLayoutBatchNorm)
.set_attr<bool>("FPurity", true);
}
@@ -663,7 +692,7 @@ TVM_FFI_STATIC_INIT_BLOCK() {
.signature(sig::arg("data", "Input to which layer_norm will be applied."),
sig::arg("gamma", "The gamma scale factor."),
sig::arg("beta", "The beta offset factor."), sig::call_attrs<LayerNormAttrs>())
- .set_attr<FInferType>("FInferType", InferTypeLayerNorm)
+ .set_attr<FInferTypeWithBuilder>("relax.FInferTypeWithBuilder", InferTypeLayerNorm)
.set_attr<FRelaxInferLayout>("FRelaxInferLayout", InferLayoutLayerNorm)
.set_attr<TMixedPrecisionPolicy>("TMixedPrecisionPolicy", MixedPrecisionPolicyKind::kFollow)
.set_attr<bool>("FPurity", true);
@@ -777,7 +806,7 @@ TVM_FFI_STATIC_INIT_BLOCK() {
.signature(sig::arg("data", "Input to which group_norm will be applied."),
sig::arg("gamma", "The gamma scale factor."),
sig::arg("beta", "The beta offset factor."), sig::call_attrs<GroupNormAttrs>())
- .set_attr<FInferType>("FInferType", InferTypeGroupNorm)
+ .set_attr<FInferTypeWithBuilder>("relax.FInferTypeWithBuilder", InferTypeGroupNorm)
.set_attr<FRelaxInferLayout>("FRelaxInferLayout", InferLayoutGroupNorm)
.set_attr<TMixedPrecisionPolicy>("TMixedPrecisionPolicy", MixedPrecisionPolicyKind::kFollow)
.set_attr<bool>("FPurity", true);
@@ -879,7 +908,7 @@ TVM_FFI_STATIC_INIT_BLOCK() {
.signature(sig::arg("data", "Input to which instance_norm will be applied."),
sig::arg("gamma", "The gamma scale factor."),
sig::arg("beta", "The beta offset factor."), sig::call_attrs<InstanceNormAttrs>())
- .set_attr<FInferType>("FInferType", InferTypeInstanceNorm)
+ .set_attr<FInferTypeWithBuilder>("relax.FInferTypeWithBuilder", InferTypeInstanceNorm)
.set_attr<FRelaxInferLayout>("FRelaxInferLayout", InferLayoutInstanceNorm)
.set_attr<TMixedPrecisionPolicy>("TMixedPrecisionPolicy", MixedPrecisionPolicyKind::kFollow)
.set_attr<bool>("FPurity", true);
@@ -939,7 +968,7 @@ TVM_FFI_STATIC_INIT_BLOCK() {
OpDef("relax.nn.rms_norm")
.signature(sig::arg("data", "Input to which rms_norm will be applied."),
sig::arg("weight", "The scale factor."), sig::call_attrs<RMSNormAttrs>())
- .set_attr<FInferType>("FInferType", InferTypeRMSNorm)
+ .set_attr<FInferTypeWithBuilder>("relax.FInferTypeWithBuilder", InferTypeRMSNorm)
.set_attr<FRelaxInferLayout>("FRelaxInferLayout", InferLayoutRMSNorm)
.set_attr<TMixedPrecisionPolicy>("TMixedPrecisionPolicy", MixedPrecisionPolicyKind::kFollow)
.set_attr<bool>("FPurity", true);
@@ -960,16 +989,21 @@ TVM_FFI_STATIC_INIT_BLOCK() {
refl::GlobalDef().def("relax.op.nn.dropout", dropout);
}
-Type InferTypeDropout(const Call& call, const BlockBuilder& ctx) {
- TensorType data_ty = GetUnaryInputTensorType(call, ctx);
+ffi::Expected<Type> InferTypeDropout(const CallNode* call_node) noexcept try {
+ const Call call = ffi::GetRef<Call>(call_node);
+ TensorType data_ty = GetUnaryInputTensorType(call);
return TupleType({data_ty, data_ty});
+} catch (const ffi::Error& error) {
+ return ffi::Unexpected(error);
+} catch (const std::exception& error) {
+ return ffi::Unexpected(ffi::Error("InternalError", error.what(), ""));
}
TVM_FFI_STATIC_INIT_BLOCK() {
OpDef("relax.nn.dropout")
.signature(sig::arg("data", "Input to which dropout will be applied."),
sig::call_attrs<DropoutAttrs>())
- .set_attr<FInferType>("FInferType", InferTypeDropout)
+ .set_attr<FInferType>("FInferType", FInferType::FromNative<&InferTypeDropout>())
.set_attr<FRelaxInferLayout>("FRelaxInferLayout", InferLayoutUnaryEwise)
.set_attr<TMixedPrecisionPolicy>("TMixedPrecisionPolicy", MixedPrecisionPolicyKind::kFollow)
.set_attr<bool>("FPurity", true);
@@ -1033,7 +1067,7 @@ TVM_FFI_STATIC_INIT_BLOCK() {
.signature(
sig::arg("predictions", "The predictions."), sig::arg("labels", "The labels."),
sig::var_ty_args("out_type", "Optional output tensor type carrying the virtual device."))
- .set_attr<FInferType>("FInferType", InferTypeCrossEntropy)
+ .set_attr<FInferTypeWithBuilder>("relax.FInferTypeWithBuilder", InferTypeCrossEntropy)
.set_attr<bool>("FPurity", true);
}
@@ -1264,7 +1298,7 @@ TVM_FFI_STATIC_INIT_BLOCK() {
sig::arg("targets", "The target tensor."), sig::var_args("args"),
sig::var_ty_args("out_type", "Optional output tensor type carrying the virtual device."),
sig::call_attrs<NLLLossAttrs>())
- .set_attr<FInferType>("FInferType", InferTypeNLLLoss)
+ .set_attr<FInferTypeWithBuilder>("relax.FInferTypeWithBuilder", InferTypeNLLLoss)
.set_attr<bool>("FPurity", true);
}
@@ -1280,8 +1314,9 @@ TVM_FFI_STATIC_INIT_BLOCK() {
refl::GlobalDef().def("relax.op.nn.batch_flatten", batch_flatten);
}
-Type InferTypeBatchFlatten(const Call& call, const BlockBuilder& ctx) {
- TensorType data_ty = GetUnaryInputTensorType(call, ctx);
+ffi::Expected<Type> InferTypeBatchFlatten(const CallNode* call_node) noexcept try {
+ const Call call = ffi::GetRef<Call>(call_node);
+ TensorType data_ty = GetUnaryInputTensorType(call);
if (data_ty->IsUnknownNdim()) {
return TensorType(data_ty->dtype, /*ndim=*/2, data_ty->vdevice);
@@ -1309,12 +1344,16 @@ Type InferTypeBatchFlatten(const Call& call, const BlockBuilder& ctx) {
}
return TensorType(ShapeExpr({batch_dim, flat_dim}), data_ty->dtype, data_ty->vdevice);
+} catch (const ffi::Error& error) {
+ return ffi::Unexpected(error);
+} catch (const std::exception& error) {
+ return ffi::Unexpected(ffi::Error("InternalError", error.what(), ""));
}
TVM_FFI_STATIC_INIT_BLOCK() {
OpDef("relax.nn.batch_flatten")
.signature(sig::arg("data", "The input tensor."))
- .set_attr<FInferType>("FInferType", InferTypeBatchFlatten)
+ .set_attr<FInferType>("FInferType", FInferType::FromNative<&InferTypeBatchFlatten>())
.set_attr<TMixedPrecisionPolicy>("TMixedPrecisionPolicy", MixedPrecisionPolicyKind::kFollow)
.set_attr<bool>("FPurity", true);
}
diff --git a/src/relax/op/nn/pooling.cc b/src/relax/op/nn/pooling.cc
index 2d2d1a78ab..4d8acb6fcc 100644
--- a/src/relax/op/nn/pooling.cc
+++ b/src/relax/op/nn/pooling.cc
@@ -147,7 +147,7 @@ InferLayoutOutput InferLayoutPool1d(
TVM_FFI_STATIC_INIT_BLOCK() {
OpDef("relax.nn.max_pool1d")
.signature(sig::arg("data", "The input tensor"), sig::call_attrs<Pool1DAttrs>())
- .set_attr<FInferType>("FInferType", InferTypePool1D)
+ .set_attr<FInferTypeWithBuilder>("relax.FInferTypeWithBuilder", InferTypePool1D)
.set_attr<FRelaxInferLayout>("FRelaxInferLayout", InferLayoutPool1d)
.set_attr<TMixedPrecisionPolicy>("TMixedPrecisionPolicy", MixedPrecisionPolicyKind::kFollow)
.set_attr<bool>("FPurity", true);
@@ -303,7 +303,7 @@ TVM_FFI_STATIC_INIT_BLOCK() {
.signature(
sig::arg("data", "The input tensor"), sig::call_attrs<Pool2DAttrs>(),
sig::var_ty_args("out_type", "Optional output tensor type carrying the virtual device."))
- .set_attr<FInferType>("FInferType", InferTypePool2D)
+ .set_attr<FInferTypeWithBuilder>("relax.FInferTypeWithBuilder", InferTypePool2D)
.set_attr<FRelaxInferLayout>("FRelaxInferLayout", InferLayoutPool2d)
.set_attr<TMixedPrecisionPolicy>("TMixedPrecisionPolicy", MixedPrecisionPolicyKind::kFollow)
.set_attr<bool>("FPurity", true);
@@ -452,7 +452,7 @@ InferLayoutOutput InferLayoutPool3d(
TVM_FFI_STATIC_INIT_BLOCK() {
OpDef("relax.nn.max_pool3d")
.signature(sig::arg("data", "The input tensor"), sig::call_attrs<Pool3DAttrs>())
- .set_attr<FInferType>("FInferType", InferTypePool3D)
+ .set_attr<FInferTypeWithBuilder>("relax.FInferTypeWithBuilder", InferTypePool3D)
.set_attr<FRelaxInferLayout>("FRelaxInferLayout", InferLayoutPool3d)
.set_attr<TMixedPrecisionPolicy>("TMixedPrecisionPolicy", MixedPrecisionPolicyKind::kFollow)
.set_attr<bool>("FPurity", true);
@@ -472,7 +472,7 @@ TVM_FFI_STATIC_INIT_BLOCK() {
OpDef("relax.nn.avg_pool1d")
.signature(sig::arg("data", "The input tensor"), sig::call_attrs<Pool1DAttrs>())
- .set_attr<FInferType>("FInferType", InferTypePool1D)
+ .set_attr<FInferTypeWithBuilder>("relax.FInferTypeWithBuilder", InferTypePool1D)
.set_attr<FRelaxInferLayout>("FRelaxInferLayout", InferLayoutPool1d)
.set_attr<TMixedPrecisionPolicy>("TMixedPrecisionPolicy", MixedPrecisionPolicyKind::kFollow)
.set_attr<bool>("FPurity", true);
@@ -492,7 +492,7 @@ TVM_FFI_STATIC_INIT_BLOCK() {
OpDef("relax.nn.avg_pool2d")
.signature(sig::arg("data", "The input tensor"), sig::call_attrs<Pool2DAttrs>())
- .set_attr<FInferType>("FInferType", InferTypePool2D)
+ .set_attr<FInferTypeWithBuilder>("relax.FInferTypeWithBuilder", InferTypePool2D)
.set_attr<FRelaxInferLayout>("FRelaxInferLayout", InferLayoutPool2d)
.set_attr<TMixedPrecisionPolicy>("TMixedPrecisionPolicy", MixedPrecisionPolicyKind::kFollow)
.set_attr<bool>("FPurity", true);
@@ -512,7 +512,7 @@ TVM_FFI_STATIC_INIT_BLOCK() {
OpDef("relax.nn.avg_pool3d")
.signature(sig::arg("data", "The input tensor"), sig::call_attrs<Pool3DAttrs>())
- .set_attr<FInferType>("FInferType", InferTypePool3D)
+ .set_attr<FInferTypeWithBuilder>("relax.FInferTypeWithBuilder", InferTypePool3D)
.set_attr<FRelaxInferLayout>("FRelaxInferLayout", InferLayoutPool3d)
.set_attr<TMixedPrecisionPolicy>("TMixedPrecisionPolicy", MixedPrecisionPolicyKind::kFollow)
.set_attr<bool>("FPurity", true);
@@ -542,19 +542,19 @@ TVM_FFI_STATIC_INIT_BLOCK() {
refl::GlobalDef().def("relax.op.nn.adaptive_avg_pool1d", adaptive_avg_pool1d);
}
-Type InferTypeAdaptiveAvgPool1D(const Call& call, const BlockBuilder& ctx) {
- TensorType data_ty = GetUnaryInputTensorType(call, ctx);
+ffi::Expected<Type> InferTypeAdaptiveAvgPool1D(const CallNode* call_node) noexcept try {
+ const Call call = ffi::GetRef<Call>(call_node);
+ TensorType data_ty = GetUnaryInputTensorType(call);
const auto* attrs = call->attrs.as<AdaptivePool1DAttrs>();
- auto [data_layout, data2NCW] = CheckTensorLayout(call, ctx, attrs->layout,
+ auto [data_layout, data2NCW] = CheckTensorLayout(call, attrs->layout,
/*tgt_layout=*/"NCW",
/*tensor_name=*/"data");
- auto [out_layout, out2NCW] = CheckTensorLayout(call, ctx, attrs->out_layout,
+ auto [out_layout, out2NCW] = CheckTensorLayout(call, attrs->out_layout,
/*tgt_layout=*/"NCW",
/*tensor_name=*/"output");
- ffi::Optional<ShapeExpr> data_shape =
- CheckNdimPerLayoutAndGetShape(call, ctx, data_ty, data_layout);
+ ffi::Optional<ShapeExpr> data_shape = CheckNdimPerLayoutAndGetShape(call, data_ty, data_layout);
if (!data_shape.has_value()) {
if (data_ty->shape.has_value() && attrs->out_layout == attrs->layout &&
!attrs->output_size.has_value()) {
@@ -572,6 +572,10 @@ Type InferTypeAdaptiveAvgPool1D(const Call& call, const BlockBuilder& ctx) {
ffi::Array<PrimExpr> out_shape = out2NCW.BackwardShape(out_NCW_shape);
return TensorType(ShapeExpr(out_shape), data_ty->dtype, data_ty->vdevice);
+} catch (const ffi::Error& error) {
+ return ffi::Unexpected(error);
+} catch (const std::exception& error) {
+ return ffi::Unexpected(ffi::Error("InternalError", error.what(), ""));
}
InferLayoutOutput InferLayoutAdaptiveAvgPool1D(
@@ -594,7 +598,7 @@ InferLayoutOutput InferLayoutAdaptiveAvgPool1D(
TVM_FFI_STATIC_INIT_BLOCK() {
OpDef("relax.nn.adaptive_avg_pool1d")
.signature(sig::arg("data", "The input tensor"), sig::call_attrs<AdaptivePool1DAttrs>())
- .set_attr<FInferType>("FInferType", InferTypeAdaptiveAvgPool1D)
+ .set_attr<FInferType>("FInferType", FInferType::FromNative<&InferTypeAdaptiveAvgPool1D>())
.set_attr<FRelaxInferLayout>("FRelaxInferLayout", InferLayoutAdaptiveAvgPool1D)
.set_attr<TMixedPrecisionPolicy>("TMixedPrecisionPolicy", MixedPrecisionPolicyKind::kFollow)
.set_attr<bool>("FPurity", true);
@@ -627,19 +631,19 @@ TVM_FFI_STATIC_INIT_BLOCK() {
refl::GlobalDef().def("relax.op.nn.adaptive_avg_pool2d", adaptive_avg_pool2d);
}
-Type InferTypeAdaptiveAvgPool2D(const Call& call, const BlockBuilder& ctx) {
- TensorType data_ty = GetUnaryInputTensorType(call, ctx);
+ffi::Expected<Type> InferTypeAdaptiveAvgPool2D(const CallNode* call_node) noexcept try {
+ const Call call = ffi::GetRef<Call>(call_node);
+ TensorType data_ty = GetUnaryInputTensorType(call);
const auto* attrs = call->attrs.as<AdaptivePool2DAttrs>();
- auto [data_layout, data2NCHW] = CheckTensorLayout(call, ctx, attrs->layout,
+ auto [data_layout, data2NCHW] = CheckTensorLayout(call, attrs->layout,
/*tgt_layout=*/"NCHW",
/*tensor_name=*/"data");
- auto [out_layout, out2NCHW] = CheckTensorLayout(call, ctx, attrs->out_layout,
+ auto [out_layout, out2NCHW] = CheckTensorLayout(call, attrs->out_layout,
/*tgt_layout=*/"NCHW",
/*tensor_name=*/"output");
- ffi::Optional<ShapeExpr> data_shape =
- CheckNdimPerLayoutAndGetShape(call, ctx, data_ty, data_layout);
+ ffi::Optional<ShapeExpr> data_shape = CheckNdimPerLayoutAndGetShape(call, data_ty, data_layout);
if (!data_shape.has_value()) {
if (data_ty->shape.has_value() && attrs->out_layout == attrs->layout &&
!attrs->output_size.has_value()) {
@@ -658,6 +662,10 @@ Type InferTypeAdaptiveAvgPool2D(const Call& call, const BlockBuilder& ctx) {
ffi::Array<PrimExpr> out_shape = out2NCHW.BackwardShape(out_NCHW_shape);
return TensorType(ShapeExpr(out_shape), data_ty->dtype, data_ty->vdevice);
+} catch (const ffi::Error& error) {
+ return ffi::Unexpected(error);
+} catch (const std::exception& error) {
+ return ffi::Unexpected(ffi::Error("InternalError", error.what(), ""));
}
InferLayoutOutput InferLayoutAdaptiveAvgPool2D(
@@ -699,7 +707,7 @@ TVM_FFI_STATIC_INIT_BLOCK() {
.signature(
sig::arg("data", "The input tensor"), sig::call_attrs<AdaptivePool2DAttrs>(),
sig::var_ty_args("out_type", "Optional output tensor type carrying the virtual device."))
- .set_attr<FInferType>("FInferType", InferTypeAdaptiveAvgPool2D)
+ .set_attr<FInferType>("FInferType", FInferType::FromNative<&InferTypeAdaptiveAvgPool2D>())
.set_attr<FRelaxInferLayout>("FRelaxInferLayout", InferLayoutAdaptiveAvgPool2D)
.set_attr<TMixedPrecisionPolicy>("TMixedPrecisionPolicy", MixedPrecisionPolicyKind::kFollow)
.set_attr<bool>("FPurity", true);
@@ -732,19 +740,19 @@ TVM_FFI_STATIC_INIT_BLOCK() {
refl::GlobalDef().def("relax.op.nn.adaptive_avg_pool3d", adaptive_avg_pool3d);
}
-Type InferTypeAdaptiveAvgPool3D(const Call& call, const BlockBuilder& ctx) {
- TensorType data_ty = GetUnaryInputTensorType(call, ctx);
+ffi::Expected<Type> InferTypeAdaptiveAvgPool3D(const CallNode* call_node) noexcept try {
+ const Call call = ffi::GetRef<Call>(call_node);
+ TensorType data_ty = GetUnaryInputTensorType(call);
const auto* attrs = call->attrs.as<AdaptivePool3DAttrs>();
- auto [data_layout, data2NCDHW] = CheckTensorLayout(call, ctx, attrs->layout,
+ auto [data_layout, data2NCDHW] = CheckTensorLayout(call, attrs->layout,
/*tgt_layout=*/"NCDHW",
/*tensor_name=*/"data");
- auto [out_layout, out2NCDHW] = CheckTensorLayout(call, ctx, attrs->out_layout,
+ auto [out_layout, out2NCDHW] = CheckTensorLayout(call, attrs->out_layout,
/*tgt_layout=*/"NCDHW",
/*tensor_name=*/"output");
- ffi::Optional<ShapeExpr> data_shape =
- CheckNdimPerLayoutAndGetShape(call, ctx, data_ty, data_layout);
+ ffi::Optional<ShapeExpr> data_shape = CheckNdimPerLayoutAndGetShape(call, data_ty, data_layout);
if (!data_shape.has_value()) {
if (data_ty->shape.has_value() && attrs->out_layout == attrs->layout &&
!attrs->output_size.has_value()) {
@@ -764,6 +772,10 @@ Type InferTypeAdaptiveAvgPool3D(const Call& call, const BlockBuilder& ctx) {
ffi::Array<PrimExpr> out_shape = out2NCDHW.BackwardShape(out_NCDHW_shape);
return TensorType(ShapeExpr(out_shape), data_ty->dtype, data_ty->vdevice);
+} catch (const ffi::Error& error) {
+ return ffi::Unexpected(error);
+} catch (const std::exception& error) {
+ return ffi::Unexpected(ffi::Error("InternalError", error.what(), ""));
}
InferLayoutOutput InferLayoutAdaptiveAvgPool3D(
@@ -786,7 +798,7 @@ InferLayoutOutput InferLayoutAdaptiveAvgPool3D(
TVM_FFI_STATIC_INIT_BLOCK() {
OpDef("relax.nn.adaptive_avg_pool3d")
.signature(sig::arg("data", "The input tensor"), sig::call_attrs<AdaptivePool3DAttrs>())
- .set_attr<FInferType>("FInferType", InferTypeAdaptiveAvgPool3D)
+ .set_attr<FInferType>("FInferType", FInferType::FromNative<&InferTypeAdaptiveAvgPool3D>())
.set_attr<FRelaxInferLayout>("FRelaxInferLayout", InferLayoutAdaptiveAvgPool3D)
.set_attr<TMixedPrecisionPolicy>("TMixedPrecisionPolicy", MixedPrecisionPolicyKind::kFollow)
.set_attr<bool>("FPurity", true);
diff --git a/src/relax/op/op.cc b/src/relax/op/op.cc
index 9dcf647cbc..02bf168805 100644
--- a/src/relax/op/op.cc
+++ b/src/relax/op/op.cc
@@ -62,13 +62,24 @@ bool EqualCheck(const PrimExpr& lhs, const PrimExpr& rhs) {
return false;
}
-Type ReturnVoidType(const Call& call, const BlockBuilder& ctx) {
+ffi::Expected<Type> ReturnVoidType(const CallNode*) noexcept try {
return TupleType(ffi::Array<Type>());
+} catch (const ffi::Error& error) {
+ return ffi::Unexpected(error);
+} catch (const std::exception& error) {
+ return ffi::Unexpected(ffi::Error("InternalError", error.what(), ""));
}
-Type ReturnAnyType(const Call& call, const BlockBuilder& ctx) { return AnyType(); }
+ffi::Expected<Type> ReturnAnyType(const CallNode*) noexcept try {
+ return AnyType();
+} catch (const ffi::Error& error) {
+ return ffi::Unexpected(error);
+} catch (const std::exception& error) {
+ return ffi::Unexpected(ffi::Error("InternalError", error.what(), ""));
+}
-Type InferTypeShapeOf(const Call& call, const BlockBuilder& ctx) {
+ffi::Expected<Type> InferTypeShapeOf(const CallNode* call_node) noexcept try {
+ const Call call = ffi::GetRef<Call>(call_node);
// use the Type of the argument
auto arg_ty = GetType(call->args[0]);
auto* tensor_ty = GetType(call->args[0]).as<TensorTypeNode>();
@@ -85,6 +96,10 @@ Type InferTypeShapeOf(const Call& call, const BlockBuilder& ctx) {
auto* tensor_shape = tensor_ty->shape.as<ShapeExprNode>();
TVM_FFI_ICHECK(tensor_shape);
return ShapeType(tensor_shape->values);
+} catch (const ffi::Error& error) {
+ return ffi::Unexpected(error);
+} catch (const std::exception& error) {
+ return ffi::Unexpected(ffi::Error("InternalError", error.what(), ""));
}
// call_pure_packed
@@ -156,7 +171,7 @@ TVM_FFI_STATIC_INIT_BLOCK() {
"arguments to that function."),
sig::var_args("args"),
sig::var_ty_args("type_args", "Optional type arguments forwarded to the callee."))
- .set_attr<FInferType>("FInferType", InferTypeCallPurePacked)
+ .set_attr<FInferTypeWithBuilder>("relax.FInferTypeWithBuilder", InferTypeCallPurePacked)
.set_attr<bool>("FPurity", true);
}
@@ -271,7 +286,7 @@ TVM_FFI_STATIC_INIT_BLOCK() {
sig::var_args("args"),
sig::var_ty_args("type_args", "Optional type arguments forwarded to the callee."),
sig::call_attrs<CallInplacePackedAttrs>())
- .set_attr<FInferType>("FInferType", InferTypeCallInplacePacked)
+ .set_attr<FInferTypeWithBuilder>("relax.FInferTypeWithBuilder", InferTypeCallInplacePacked)
// Warning: considered pure, but it has the potential to create visible effects!
// This should only be used if it has been *checked* that it is safe (no aliases, in-place
// arguments will no longer be live) and the user believes the packed func to have no
@@ -432,7 +447,8 @@ static ffi::Optional<Type> InferCallTIROutputTypeFromArguments(
return derived_ret_ty;
}
-Type InferTypeCallTIR(const Call& call, const BlockBuilder& ctx) {
+ffi::Expected<Type> InferTypeCallTIR(const CallNode* call_node) noexcept try {
+ const Call call = ffi::GetRef<Call>(call_node);
if (call->ty_args.size() != 1) {
TVM_FFI_VISIT_THROW(InternalError, call) << "ty_args should have exactly 1 output type.";
}
@@ -444,6 +460,10 @@ Type InferTypeCallTIR(const Call& call, const BlockBuilder& ctx) {
Type explicit_ty = call->ty_args[0];
return explicit_ty;
+} catch (const ffi::Error& error) {
+ return ffi::Unexpected(error);
+} catch (const std::exception& error) {
+ return ffi::Unexpected(ffi::Error("InternalError", error.what(), ""));
}
Expr NormalizeCallTIR(const BlockBuilder& ctx, Call call) {
@@ -561,7 +581,7 @@ TVM_FFI_STATIC_INIT_BLOCK() {
.signature(sig::arg("func", "The destination-passing-style function."),
sig::arg("args", "The input arguments."),
sig::ty_arg("out_type", "The output type."))
- .set_attr<FInferType>("FInferType", InferTypeCallTIR)
+ .set_attr<FInferType>("FInferType", FInferType::FromNative<&InferTypeCallTIR>())
.set_attr<FNormalize>("FNormalize", NormalizeCallTIR)
.set_attr<bool>("FPurity", true);
}
@@ -599,7 +619,7 @@ TVM_FFI_STATIC_INIT_BLOCK() {
sig::arg("args", "The input arguments."),
sig::ty_arg("out_type", "The output type."),
sig::call_attrs<CallTIRWithGradAttrs>())
- .set_attr<FInferType>("FInferType", InferTypeCallTIR)
+ .set_attr<FInferType>("FInferType", FInferType::FromNative<&InferTypeCallTIR>())
.set_attr<FNormalize>("FNormalize", NormalizeCallTIR)
.set_attr<bool>("FPurity", true);
}
@@ -731,7 +751,7 @@ TVM_FFI_STATIC_INIT_BLOCK() {
sig::arg("args", "The input arguments."),
sig::ty_arg("out_type", "The output type."),
sig::call_attrs<CallTIRInplaceAttrs>())
- .set_attr<FInferType>("FInferType", InferTypeCallTIR)
+ .set_attr<FInferType>("FInferType", FInferType::FromNative<&InferTypeCallTIR>())
.set_attr<FNormalize>("FNormalize", NormalizeCallTIRInPlace)
// Warning: considered pure, but it has the potential to create visible effects!
// This should only be used if it has been *checked* that it is safe (no aliases, in-place
@@ -770,11 +790,16 @@ TVM_FFI_STATIC_INIT_BLOCK() {
// call_dps_packed
-Type InferTypeCallDPSPacked(const Call& call, const BlockBuilder& ctx) {
+ffi::Expected<Type> InferTypeCallDPSPacked(const CallNode* call_node) noexcept try {
+ const Call call = ffi::GetRef<Call>(call_node);
if (call->ty_args.size() != 1) {
TVM_FFI_VISIT_THROW(InternalError, call) << "ty_args should have exact 1 output type.";
}
return call->ty_args[0];
+} catch (const ffi::Error& error) {
+ return ffi::Unexpected(error);
+} catch (const std::exception& error) {
+ return ffi::Unexpected(ffi::Error("InternalError", error.what(), ""));
}
TVM_FFI_STATIC_INIT_BLOCK() {
@@ -782,7 +807,7 @@ TVM_FFI_STATIC_INIT_BLOCK() {
.signature(sig::arg("func", "The destination-passing-style function."),
sig::arg("args", "The input arguments."),
sig::ty_arg("out_type", "The output type."))
- .set_attr<FInferType>("FInferType", InferTypeCallDPSPacked)
+ .set_attr<FInferType>("FInferType", FInferType::FromNative<&InferTypeCallDPSPacked>())
// technically, an impure op could be used with this, but there is
// little reason to use DPS with an impure op
.set_attr<bool>("FPurity", true);
@@ -815,11 +840,16 @@ TVM_FFI_STATIC_INIT_BLOCK() {
// call_py_func
-Type InferTypeCallPyFunc(const Call& call, const BlockBuilder& ctx) {
+ffi::Expected<Type> InferTypeCallPyFunc(const CallNode* call_node) noexcept try {
+ const Call call = ffi::GetRef<Call>(call_node);
if (call->ty_args.size() != 1) {
TVM_FFI_VISIT_THROW(InternalError, call) << "ty_args should have exact 1 output type.";
}
return call->ty_args[0];
+} catch (const ffi::Error& error) {
+ return ffi::Unexpected(error);
+} catch (const std::exception& error) {
+ return ffi::Unexpected(ffi::Error("InternalError", error.what(), ""));
}
void ValidateCallPyFunc(const CallNode* call) {
@@ -850,7 +880,7 @@ TVM_FFI_STATIC_INIT_BLOCK() {
.signature(sig::arg("func_name", "The name of the Python function to call."),
sig::arg("args", "The input arguments."),
sig::ty_arg("out_type", "The output type."))
- .set_attr<FInferType>("FInferType", InferTypeCallPyFunc)
+ .set_attr<FInferType>("FInferType", FInferType::FromNative<&InferTypeCallPyFunc>())
.set_attr<bool>("FPurity", true);
}
@@ -880,7 +910,8 @@ TVM_FFI_STATIC_INIT_BLOCK() {
}
// call builtin
-Type InferTypeCallBuiltinWithCtx(const Call& call, const BlockBuilder& ctx) {
+ffi::Expected<Type> InferTypeCallBuiltinWithCtx(const CallNode* call_node) noexcept try {
+ const Call call = ffi::GetRef<Call>(call_node);
if (call->ty_args.size() == 0) {
// by default return void.
return TupleType(ffi::Array<Type>());
@@ -888,6 +919,10 @@ Type InferTypeCallBuiltinWithCtx(const Call& call, const BlockBuilder& ctx) {
TVM_FFI_ICHECK_EQ(call->ty_args.size(), 1);
return call->ty_args[0];
}
+} catch (const ffi::Error& error) {
+ return ffi::Unexpected(error);
+} catch (const std::exception& error) {
+ return ffi::Unexpected(ffi::Error("InternalError", error.what(), ""));
}
TVM_FFI_STATIC_INIT_BLOCK() {
@@ -895,7 +930,7 @@ TVM_FFI_STATIC_INIT_BLOCK() {
.signature(sig::arg("func", "The builtin packed func."),
sig::arg("args", "The input arguments."),
sig::var_ty_args("out_type", "Optional output type; omitted for void."))
- .set_attr<FInferType>("FInferType", InferTypeCallBuiltinWithCtx)
+ .set_attr<FInferType>("FInferType", FInferType::FromNative<&InferTypeCallBuiltinWithCtx>())
// Most builtins are pure, but some are not, like `vm.builtin.attention_kv_cache_append`
.set_attr<bool>("FPurity", false);
}
@@ -910,7 +945,7 @@ TVM_FFI_STATIC_INIT_BLOCK() {
refl::GlobalDef().def("relax.op.call_builtin_with_ctx", MakeCallBuiltinWithCtx);
OpDef("relax.null_value")
- .set_attr<FInferType>("FInferType", ReturnAnyType)
+ .set_attr<FInferType>("FInferType", FInferType::FromNative<&ReturnAnyType>())
.set_attr<bool>("FPurity", true);
}
@@ -931,7 +966,7 @@ TVM_FFI_STATIC_INIT_BLOCK() {
"The first value is Python-style format string to use to print. The others "
"are values to print"),
sig::var_args("args"))
- .set_attr<FInferType>("FInferType", ReturnVoidType)
+ .set_attr<FInferType>("FInferType", FInferType::FromNative<&ReturnVoidType>())
.set_attr<FCallPacked>("FCallPacked", "relax.run.print")
.set_attr<bool>("FPurity", false);
}
@@ -955,7 +990,8 @@ TVM_FFI_STATIC_INIT_BLOCK() {
// can't actually name it assert or else Python will consider it a syntax error
-Type InferAssertType(const Call& call, const BlockBuilder& ctx) {
+ffi::Expected<Type> InferAssertType(const CallNode* call_node) noexcept try {
+ const Call call = ffi::GetRef<Call>(call_node);
// Ensure that the condition argument is a boolean scalar.
// Also permitted is a tensor with unknown shape and unknown dtype
// (checked dynamically in that case). Returns void.
@@ -968,7 +1004,11 @@ Type InferAssertType(const Call& call, const BlockBuilder& ctx) {
TVM_FFI_VISIT_THROW(TypeError, call)
<< "The argument to assert must be a boolean scalar, but received " << arg_ty;
}
- return ReturnVoidType(call, ctx);
+ return ReturnVoidType(call.get());
+} catch (const ffi::Error& error) {
+ return ffi::Unexpected(error);
+} catch (const std::exception& error) {
+ return ffi::Unexpected(ffi::Error("InternalError", error.what(), ""));
}
TVM_FFI_STATIC_INIT_BLOCK() {
@@ -979,7 +1019,7 @@ TVM_FFI_STATIC_INIT_BLOCK() {
"Python-style format string to use for displaying an error message, if the "
"assert fails. The others are used as format arguments if there is an error."),
sig::var_args("args"))
- .set_attr<FInferType>("FInferType", InferAssertType)
+ .set_attr<FInferType>("FInferType", FInferType::FromNative<&InferAssertType>())
.set_attr<FCallPacked>("FCallPacked", "relax.run.assert_op")
.set_attr<bool>("FPurity", false);
}
@@ -1002,7 +1042,7 @@ TVM_FFI_STATIC_INIT_BLOCK() {
OpDef("relax.make_closure")
.signature(sig::arg("func", "The closure."), sig::arg("args", "The captured variables."))
- .set_attr<FInferType>("FInferType", ReturnAnyType)
+ .set_attr<FInferType>("FInferType", FInferType::FromNative<&ReturnAnyType>())
.set_attr<bool>("FPurity", true);
}
@@ -1018,7 +1058,8 @@ TVM_FFI_STATIC_INIT_BLOCK() {
// invoke_closure
-Type InferTypeInvokeClosure(const Call& call, const BlockBuilder& ctx) {
+ffi::Expected<Type> InferTypeInvokeClosure(const CallNode* call_node) noexcept try {
+ const Call call = ffi::GetRef<Call>(call_node);
if (call->ty_args.empty()) {
return AnyType();
} else if (call->ty_args.size() == 1) {
@@ -1026,6 +1067,10 @@ Type InferTypeInvokeClosure(const Call& call, const BlockBuilder& ctx) {
} else {
return TupleType(call->ty_args);
}
+} catch (const ffi::Error& error) {
+ return ffi::Unexpected(error);
+} catch (const std::exception& error) {
+ return ffi::Unexpected(ffi::Error("InternalError", error.what(), ""));
}
TVM_FFI_STATIC_INIT_BLOCK() {
@@ -1033,7 +1078,7 @@ TVM_FFI_STATIC_INIT_BLOCK() {
.signature(sig::arg("closure", "The VMClosure."), sig::arg("args", "The captured variables."),
sig::var_ty_args("out_types",
"Zero or more output types; multiple entries form a tuple."))
- .set_attr<FInferType>("FInferType", InferTypeInvokeClosure)
+ .set_attr<FInferType>("FInferType", FInferType::FromNative<&InferTypeInvokeClosure>())
// Not all closures are pure. Use invoke_pure_closure for specifying purity
.set_attr<bool>("FPurity", false);
}
@@ -1053,7 +1098,7 @@ TVM_FFI_STATIC_INIT_BLOCK() {
.signature(sig::arg("closure", "The VMClosure."), sig::arg("args", "The captured variables."),
sig::var_ty_args("out_types",
"Zero or more output types; multiple entries form a tuple."))
- .set_attr<FInferType>("FInferType", InferTypeInvokeClosure)
+ .set_attr<FInferType>("FInferType", FInferType::FromNative<&InferTypeInvokeClosure>())
.set_attr<bool>("FPurity", true);
}
@@ -1070,7 +1115,7 @@ TVM_FFI_STATIC_INIT_BLOCK() {
OpDef("relax.shape_of")
.signature(sig::arg("input", "The input expression"))
- .set_attr<FInferType>("FInferType", InferTypeShapeOf)
+ .set_attr<FInferType>("FInferType", FInferType::FromNative<&InferTypeShapeOf>())
.set_attr<bool>("FPurity", true);
}
@@ -1086,18 +1131,23 @@ TVM_FFI_STATIC_INIT_BLOCK() {
// size
-Type InferTypeSize(const Call& call, const BlockBuilder& ctx) {
+ffi::Expected<Type> InferTypeSize(const CallNode* call_node) noexcept try {
+ const Call call = ffi::GetRef<Call>(call_node);
auto arg_ty = GetType(call->args[0]);
auto* tensor_ty = GetType(call->args[0]).as<TensorTypeNode>();
TVM_FFI_ICHECK(tensor_ty) << "size expects a tensor input, but received " << arg_ty
<< "; use MatchCast if necessary";
return TensorType(ShapeExpr(ffi::Array<PrimExpr>{}), PrimType::Int(64));
+} catch (const ffi::Error& error) {
+ return ffi::Unexpected(error);
+} catch (const std::exception& error) {
+ return ffi::Unexpected(ffi::Error("InternalError", error.what(), ""));
}
TVM_FFI_STATIC_INIT_BLOCK() {
OpDef("relax.size")
.signature(sig::arg("input", "The input tensor"))
- .set_attr<FInferType>("FInferType", InferTypeSize)
+ .set_attr<FInferType>("FInferType", FInferType::FromNative<&InferTypeSize>())
.set_attr<bool>("FPurity", true);
}
@@ -1113,7 +1163,8 @@ TVM_FFI_STATIC_INIT_BLOCK() {
// tensor_to_shape
-Type ReturnTensorToShapeType(const Call& call, const BlockBuilder& ctx) {
+ffi::Expected<Type> ReturnTensorToShapeType(const CallNode* call_node) noexcept try {
+ const Call call = ffi::GetRef<Call>(call_node);
TVM_FFI_ICHECK(call->args.size() == 1);
TVM_FFI_ICHECK(!call->args[0]->ty.IsMissing());
const auto* tensor_ty = GetTypeAs<TensorTypeNode>(call->args[0]);
@@ -1130,12 +1181,16 @@ Type ReturnTensorToShapeType(const Call& call, const BlockBuilder& ctx) {
}
}
return ShapeType(kUnknownNDim);
+} catch (const ffi::Error& error) {
+ return ffi::Unexpected(error);
+} catch (const std::exception& error) {
+ return ffi::Unexpected(ffi::Error("InternalError", error.what(), ""));
}
TVM_FFI_STATIC_INIT_BLOCK() {
OpDef("relax.tensor_to_shape")
.signature(sig::arg("input", "The input expression"))
- .set_attr<FInferType>("FInferType", ReturnTensorToShapeType)
+ .set_attr<FInferType>("FInferType", FInferType::FromNative<&ReturnTensorToShapeType>())
.set_attr<bool>("FPurity", true);
}
@@ -1150,19 +1205,24 @@ TVM_FFI_STATIC_INIT_BLOCK() {
}
// shape_to_tensor
-Type ReturnShapeToTensorType(const Call& call, const BlockBuilder& ctx) {
+ffi::Expected<Type> ReturnShapeToTensorType(const CallNode* call_node) noexcept try {
+ const Call call = ffi::GetRef<Call>(call_node);
TVM_FFI_ICHECK(call->args.size() == 1);
TVM_FFI_ICHECK(!call->args[0]->ty.IsMissing());
const auto* ty = GetTypeAs<ShapeTypeNode>(call->args[0]);
TVM_FFI_ICHECK(ty);
int32_t ndim = ty->ndim;
return TensorType(ShapeExpr({PrimExpr(ndim)}), PrimType::Int(64));
+} catch (const ffi::Error& error) {
+ return ffi::Unexpected(error);
+} catch (const std::exception& error) {
+ return ffi::Unexpected(ffi::Error("InternalError", error.what(), ""));
}
TVM_FFI_STATIC_INIT_BLOCK() {
OpDef("relax.shape_to_tensor")
.signature(sig::arg("input", "The input expression"))
- .set_attr<FInferType>("FInferType", ReturnShapeToTensorType)
+ .set_attr<FInferType>("FInferType", FInferType::FromNative<&ReturnShapeToTensorType>())
.set_attr<FCallPacked>("FCallPacked", "relax.run.shape_to_tensor")
.set_attr<bool>("FPurity", true);
}
@@ -1211,7 +1271,7 @@ TVM_FFI_STATIC_INIT_BLOCK() {
sig::arg("storage_scope",
"The storage scope of the storage to allocate. Default is global."),
sig::var_ty_args("out_type", "Optional output type used by allocation rewrites."))
- .set_attr<FInferType>("FInferType", InferTypeAllocateTensor)
+ .set_attr<FInferTypeWithBuilder>("relax.FInferTypeWithBuilder", InferTypeAllocateTensor)
// memory allocation isn't considered a "visible effect" as far as purity is concerned
.set_attr<bool>("FPurity", true)
.set_attr<bool>("TAllocator", true);
@@ -1240,7 +1300,7 @@ TVM_FFI_STATIC_INIT_BLOCK() {
sig::arg("storage_scope",
"The storage scope of the storage to allocate. Default is global."),
sig::arg("dtype", "The dtype of the tensor to allocate."))
- .set_attr<FInferType>("FInferType", ReturnAnyType)
+ .set_attr<FInferType>("FInferType", FInferType::FromNative<&ReturnAnyType>())
// memory allocation isn't considered a "visible effect" as far as purity is concerned
.set_attr<bool>("FPurity", true)
.set_attr<bool>("TAllocator", true);
@@ -1292,7 +1352,7 @@ TVM_FFI_STATIC_INIT_BLOCK() {
sig::arg("runtime_device_index",
"The device index indicating on which device the tensor is to be "
"allocated at runtime. Index -1 is reserved for the host device."))
- .set_attr<FInferType>("FInferType", InferTypeMemAllocTensor)
+ .set_attr<FInferTypeWithBuilder>("relax.FInferTypeWithBuilder", InferTypeMemAllocTensor)
// memory allocation isn't considered a "visible effect" as far as purity is concerned
.set_attr<bool>("FPurity", true)
.set_attr<bool>("TAllocator", true);
@@ -1324,7 +1384,7 @@ TVM_FFI_STATIC_INIT_BLOCK() {
OpDef("relax.memory.kill_storage")
.signature(sig::arg("storage", "The storage to be killed."))
- .set_attr<FInferType>("FInferType", ReturnVoidType)
+ .set_attr<FInferType>("FInferType", FInferType::FromNative<&ReturnVoidType>())
// We mark this as impure so it wouldn't be removed by "remove_all_unused"
.set_attr<bool>("FPurity", false);
}
@@ -1342,7 +1402,7 @@ TVM_FFI_STATIC_INIT_BLOCK() {
OpDef("relax.memory.kill_tensor")
.signature(sig::arg("tensor", "The tensor to be killed."))
- .set_attr<FInferType>("FInferType", ReturnVoidType)
+ .set_attr<FInferType>("FInferType", FInferType::FromNative<&ReturnVoidType>())
// We mark this as impure so it wouldn't be removed by "remove_all_unused"
.set_attr<bool>("FPurity", false);
}
@@ -1366,7 +1426,7 @@ TVM_FFI_STATIC_INIT_BLOCK() {
sig::arg("dtype", "The dtype of the tensor to allocate."),
sig::arg("storage_scope",
"The storage scope of the storage to allocate. Default is global."))
- .set_attr<FInferType>("FInferType", ReturnAnyType)
+ .set_attr<FInferType>("FInferType", FInferType::FromNative<&ReturnAnyType>())
// memory allocation isn't considered a "visible effect" as far as purity is concerned
.set_attr<bool>("FPurity", true)
.set_attr<bool>("TAllocator", true);
@@ -1419,7 +1479,7 @@ TVM_FFI_STATIC_INIT_BLOCK() {
sig::arg("runtime_device_index",
"The device index indicating on which device the tensor is "
"to be allocated at runtime."))
- .set_attr<FInferType>("FInferType", InferTypeVMAllocTensor)
+ .set_attr<FInferTypeWithBuilder>("relax.FInferTypeWithBuilder", InferTypeVMAllocTensor)
// memory allocation isn't considered a "visible effect" as far as purity is concerned
.set_attr<bool>("FPurity", true)
.set_attr<bool>("TAllocator", true);
@@ -1448,7 +1508,7 @@ TVM_FFI_STATIC_INIT_BLOCK() {
OpDef("relax.vm.kill_object")
.signature(sig::arg("obj", "The object to be killed."))
- .set_attr<FInferType>("FInferType", ReturnVoidType)
+ .set_attr<FInferType>("FInferType", FInferType::FromNative<&ReturnVoidType>())
// We mark this as impure so it wouldn't be removed by "remove_all_unused"
.set_attr<bool>("FPurity", false);
}
@@ -1468,7 +1528,7 @@ TVM_FFI_STATIC_INIT_BLOCK() {
.signature(
sig::arg("func", "The destination-passing-style function."),
sig::arg("args", "The input arguments (list of tensors and last argument is ShapeExpr)"))
- .set_attr<FInferType>("FInferType", ReturnVoidType)
+ .set_attr<FInferType>("FInferType", FInferType::FromNative<&ReturnVoidType>())
// "relax.vm.call_tir_dyn" works in an in-place way, which is impure.
.set_attr<bool>("FPurity", false);
}
@@ -1484,15 +1544,19 @@ TVM_FFI_STATIC_INIT_BLOCK() {
}
// builtin stop_lift_params
-Type InferTypeStopLiftParams(const Call& call, const BlockBuilder& ctx) {
- return InferTypeUnaryArith<false>(call, ctx);
+ffi::Expected<Type> InferTypeStopLiftParams(const CallNode* call) noexcept try {
+ return InferTypeUnaryArithContextFree<false>(call);
+} catch (const ffi::Error& error) {
+ return ffi::Unexpected(error);
+} catch (const std::exception& error) {
+ return ffi::Unexpected(ffi::Error("InternalError", error.what(), ""));
}
TVM_FFI_STATIC_INIT_BLOCK() {
OpDef("relax.builtin.stop_lift_params")
.signature(sig::arg("x", "The input data"),
sig::var_ty_args("out_type", "Optional output type."))
- .set_attr<FInferType>("FInferType", InferTypeStopLiftParams)
+ .set_attr<FInferType>("FInferType", FInferType::FromNative<&InferTypeStopLiftParams>())
.set_attr<bool>("FPurity", true);
}
@@ -1508,23 +1572,28 @@ TVM_FFI_STATIC_INIT_BLOCK() {
// to_vdevice
-Type InferToVDeviceType(const Call& call, const BlockBuilder& ctx) {
+ffi::Expected<Type> InferToVDeviceType(const CallNode* call_node) noexcept try {
+ const Call call = ffi::GetRef<Call>(call_node);
TVM_FFI_ICHECK(call->args.size() == 1);
TVM_FFI_ICHECK(!call->args[0]->ty.IsMissing());
- TensorType data_ty = GetUnaryInputTensorType(call, ctx);
+ TensorType data_ty = GetUnaryInputTensorType(call);
auto attrs = call->attrs.as<ToVDeviceAttrs>();
VDevice vdev = attrs->dst_vdevice;
if (data_ty->shape.has_value()) {
return TensorType(data_ty->shape.value(), data_ty->dtype, vdev, data_ty->span);
}
return TensorType(data_ty->dtype, data_ty->ndim, vdev, data_ty->span);
+} catch (const ffi::Error& error) {
+ return ffi::Unexpected(error);
+} catch (const std::exception& error) {
+ return ffi::Unexpected(ffi::Error("InternalError", error.what(), ""));
}
TVM_FFI_STATIC_INIT_BLOCK() {
OpDef("relax.to_vdevice")
.signature(sig::arg("data", "The input expression to be copied"),
sig::call_attrs<ToVDeviceAttrs>())
- .set_attr<FInferType>("FInferType", InferToVDeviceType)
+ .set_attr<FInferType>("FInferType", FInferType::FromNative<&InferToVDeviceType>())
.set_attr<bool>("FPurity", true);
}
@@ -1542,17 +1611,22 @@ TVM_FFI_STATIC_INIT_BLOCK() {
// hint_on_device
-Type InferHintOnDeviceType(const Call& call, const BlockBuilder& ctx) {
+ffi::Expected<Type> InferHintOnDeviceType(const CallNode* call_node) noexcept try {
+ const Call call = ffi::GetRef<Call>(call_node);
TVM_FFI_ICHECK(call->args.size() == 1);
TVM_FFI_ICHECK(!call->args[0]->ty.IsMissing());
- TensorType data_ty = GetUnaryInputTensorType(call, ctx);
+ TensorType data_ty = GetUnaryInputTensorType(call);
return data_ty;
+} catch (const ffi::Error& error) {
+ return ffi::Unexpected(error);
+} catch (const std::exception& error) {
+ return ffi::Unexpected(ffi::Error("InternalError", error.what(), ""));
}
TVM_FFI_STATIC_INIT_BLOCK() {
OpDef("relax.hint_on_device")
.signature(sig::arg("data", "The input expression"), sig::call_attrs<HintOnDeviceAttrs>())
- .set_attr<FInferType>("FInferType", InferHintOnDeviceType)
+ .set_attr<FInferType>("FInferType", FInferType::FromNative<&InferHintOnDeviceType>())
.set_attr<bool>("FPurity", true);
}
diff --git a/src/relax/op/op_common.cc b/src/relax/op/op_common.cc
index fd2d39b953..5cbbbc9ebc 100644
--- a/src/relax/op/op_common.cc
+++ b/src/relax/op/op_common.cc
@@ -40,7 +40,7 @@ ffi::Array<Expr> GetCallArgs(const Call& call) {
return args;
}
-void CheckNumArguments(const Call& call, const BlockBuilder& ctx) {
+void CheckNumArguments(const Call& call) {
Op op = call->op.as_or_throw<Op>();
int expected_input = op->args_info.size();
if (op->var_args_info.has_value() ? call->args.size() < op->args_info.size()
@@ -52,7 +52,7 @@ void CheckNumArguments(const Call& call, const BlockBuilder& ctx) {
}
}
-TensorType GetInputTensorType(const Call& call, size_t i_arg, const BlockBuilder& ctx) {
+TensorType GetInputTensorType(const Call& call, size_t i_arg) {
Op op = call->op.as_or_throw<Op>();
TVM_FFI_ICHECK(op->var_args_info.has_value() ? call->args.size() >= op->args_info.size()
@@ -77,19 +77,18 @@ TensorType GetInputTensorType(const Call& call, size_t i_arg, const BlockBuilder
}
}
-ffi::Array<TensorType> GetInputTensorType(const Call& call, const BlockBuilder& ctx) {
- CheckNumArguments(call, ctx);
+ffi::Array<TensorType> GetInputTensorType(const Call& call) {
+ CheckNumArguments(call);
Op op = call->op.as_or_throw<Op>();
ffi::Array<TensorType> input_tensor_ty;
for (size_t i = 0; i < call->args.size(); ++i) {
- input_tensor_ty.push_back(GetInputTensorType(call, i, ctx));
+ input_tensor_ty.push_back(GetInputTensorType(call, i));
}
return input_tensor_ty;
}
-ffi::Array<TensorType> GetTensorTypeFromTuple(const Call& call, const BlockBuilder& ctx,
- const Expr& tup) {
+ffi::Array<TensorType> GetTensorTypeFromTuple(const Call& call, const Expr& tup) {
const auto* tuple_ty = GetTypeAs<TupleTypeNode>(tup);
if (tuple_ty == nullptr) {
TVM_FFI_VISIT_THROW(TypeError, call)
@@ -176,8 +175,7 @@ ffi::Optional<ffi::Array<PrimExpr>> InferBinaryBroadcastShape(
TVM_FFI_UNREACHABLE();
}
-std::vector<int> NormalizeAxes(const Call& call, const BlockBuilder& ctx, int ndim,
- const ffi::Array<int64_t>& axes) {
+std::vector<int> NormalizeAxes(const Call& call, int ndim, const ffi::Array<int64_t>& axes) {
TVM_FFI_ICHECK_NE(ndim, kUnknownNDim) << "The ndim is required to be known for this function.";
std::vector<bool> appeared_dims_set;
std::vector<int> axes_non_neg;
diff --git a/src/relax/op/op_common.h b/src/relax/op/op_common.h
index ccfdeb1d09..e71b166e54 100644
--- a/src/relax/op/op_common.h
+++ b/src/relax/op/op_common.h
@@ -56,7 +56,8 @@ using namespace tvm::prim;
*
* \param ctx The error reporting context.
*/
-void CheckNumArguments(const Call& call, const BlockBuilder& ctx);
+void CheckNumArguments(const Call& call);
+inline void CheckNumArguments(const Call& call, const BlockBuilder&) { CheckNumArguments(call); }
/*!
* \brief Get the tensor type of the operator input.
@@ -65,7 +66,10 @@ void CheckNumArguments(const Call& call, const BlockBuilder& ctx);
* \param ctx The error reporting context.
* \return The tensor type of the argument
*/
-TensorType GetInputTensorType(const Call& call, size_t i_arg, const BlockBuilder& ctx);
+TensorType GetInputTensorType(const Call& call, size_t i_arg);
+inline TensorType GetInputTensorType(const Call& call, size_t i_arg, const BlockBuilder&) {
+ return GetInputTensorType(call, i_arg);
+}
/*!
* \brief Get the tensor type of the operator input.
@@ -75,7 +79,10 @@ TensorType GetInputTensorType(const Call& call, size_t i_arg, const BlockBuilder
* \note This function require every input to be Tensor. The number of call arguments is required
* to match the number of inputs of the op being called.
*/
-ffi::Array<TensorType> GetInputTensorType(const Call& call, const BlockBuilder& ctx);
+ffi::Array<TensorType> GetInputTensorType(const Call& call);
+inline ffi::Array<TensorType> GetInputTensorType(const Call& call, const BlockBuilder&) {
+ return GetInputTensorType(call);
+}
/*!
* \brief Get the tensor type of the unary operator input.
@@ -85,8 +92,9 @@ ffi::Array<TensorType> GetInputTensorType(const Call& call, const BlockBuilder&
* \throw Throw exception if the number of input is not one, or the type of the input is not
* a tensor type.
*/
-inline TensorType GetUnaryInputTensorType(const Call& call, const BlockBuilder& ctx) {
- return GetInputTensorType(call, ctx)[0];
+inline TensorType GetUnaryInputTensorType(const Call& call) { return GetInputTensorType(call)[0]; }
+inline TensorType GetUnaryInputTensorType(const Call& call, const BlockBuilder&) {
+ return GetUnaryInputTensorType(call);
}
/*!
@@ -97,13 +105,16 @@ inline TensorType GetUnaryInputTensorType(const Call& call, const BlockBuilder&
* \return The tensor types of tuple input.
* \throw Throw exception if input expression is not a tuple.
*/
-ffi::Array<TensorType> GetTensorTypeFromTuple(const Call& call, const BlockBuilder& ctx,
- const Expr& tup);
+ffi::Array<TensorType> GetTensorTypeFromTuple(const Call& call, const Expr& tup);
+inline ffi::Array<TensorType> GetTensorTypeFromTuple(const Call& call, const BlockBuilder&,
+ const Expr& tup) {
+ return GetTensorTypeFromTuple(call, tup);
+}
namespace detail {
/*! \brief Implementation helper for GetArgType */
template <typename ArgType>
-ArgType GetArgTypeByIndex(const Call& call, const Op& op, const BlockBuilder& ctx, size_t index) {
+ArgType GetArgTypeByIndex(const Call& call, const Op& op, size_t index) {
if (call->args[index]->ty.IsMissing()) {
TVM_FFI_VISIT_THROW(InternalError, call)
<< op << " op should have arguments with defined Type. "
@@ -124,9 +135,9 @@ ArgType GetArgTypeByIndex(const Call& call, const Op& op, const BlockBuilder& ct
/*! \brief Implementation helper for GetArgType */
template <typename... ArgTypes, size_t... Indices>
-std::tuple<ArgTypes...> GetArgTypeHelper(const Call& call, const Op& op, const BlockBuilder& ctx,
+std::tuple<ArgTypes...> GetArgTypeHelper(const Call& call, const Op& op,
std::index_sequence<Indices...>) {
- return std::tuple<ArgTypes...>{GetArgTypeByIndex<ArgTypes>(call, op, ctx, Indices)...};
+ return std::tuple<ArgTypes...>{GetArgTypeByIndex<ArgTypes>(call, op, Indices)...};
}
} // namespace detail
@@ -140,7 +151,7 @@ std::tuple<ArgTypes...> GetArgTypeHelper(const Call& call, const Op& op, const B
* \throw Throw exception if input expression is not a tuple.
*/
template <typename... ArgTypes>
-std::tuple<ArgTypes...> GetArgType(const Call& call, const BlockBuilder& ctx) {
+std::tuple<ArgTypes...> GetArgType(const Call& call) {
Op op = call->op.as_or_throw<Op>();
size_t n_input = op->args_info.size();
@@ -152,12 +163,17 @@ std::tuple<ArgTypes...> GetArgType(const Call& call, const BlockBuilder& ctx) {
<< "Internal error: " << op << " op defines " << n_input << " arguments in its OpDef() call, "
<< "but GetArgType was given " << sizeof...(ArgTypes) << " template arguments.";
- CheckNumArguments(call, ctx);
+ CheckNumArguments(call);
TVM_FFI_ICHECK_GE(call->args.size(), sizeof...(ArgTypes));
- return detail::GetArgTypeHelper<ArgTypes...>(call, op, ctx,
+ return detail::GetArgTypeHelper<ArgTypes...>(call, op,
std::make_index_sequence<sizeof...(ArgTypes)>());
}
+template <typename... ArgTypes>
+std::tuple<ArgTypes...> GetArgType(const Call& call, const BlockBuilder&) {
+ return GetArgType<ArgTypes...>(call);
+}
+
/************ Op registration macro ************/
/*!
@@ -171,8 +187,8 @@ std::tuple<ArgTypes...> GetArgType(const Call& call, const BlockBuilder& ctx) {
* \return The inferred type.
*/
template <bool require_float_dtype, typename FType>
-inline Type InferTypeUnary(const Call& call, const BlockBuilder& ctx, FType f_compute_out_dtype) {
- TensorType input_ty = GetUnaryInputTensorType(call, ctx);
+inline Type InferTypeUnary(const Call& call, FType f_compute_out_dtype) {
+ TensorType input_ty = GetUnaryInputTensorType(call);
if (require_float_dtype && !input_ty->IsUnknownDtype() &&
!input_ty->dtype.value().MatchesCode(DLDataTypeCode::kDLFloat, DLDataTypeCode::kDLBfloat)) {
TVM_FFI_VISIT_THROW(TypeError, call)
@@ -180,19 +196,46 @@ inline Type InferTypeUnary(const Call& call, const BlockBuilder& ctx, FType f_co
<< " requires the input tensor to have float dtype. However, the given input dtype is "
<< input_ty->dtype.value();
}
- auto output_ty = ffi::make_object<TensorTypeNode>(*input_ty.get());
ffi::Optional<PrimType> computed_dtype = f_compute_out_dtype(input_ty);
- output_ty->dtype = computed_dtype;
- if (call->ty_args.size() > 0) {
- auto defined_ty = call->ty_args[0].as<TensorTypeNode>();
- TVM_FFI_ICHECK(defined_ty);
- auto shape = output_ty->GetShape();
- TVM_FFI_ICHECK(shape.has_value());
- TVM_FFI_ICHECK(defined_ty->vdevice.has_value());
- return TensorType(ShapeExpr(shape.value()), output_ty->dtype, defined_ty->vdevice.value());
- } else {
- return TensorType(output_ty);
+ bool same_dtype = computed_dtype.has_value() == input_ty->dtype.has_value() &&
+ (!computed_dtype.has_value() ||
+ computed_dtype.value()->dtype == input_ty->dtype.value()->dtype);
+ TensorType inferred = input_ty;
+ if (!call->ty_args.empty() || !same_dtype) {
+ auto output_ty = ffi::make_object<TensorTypeNode>(*input_ty.get());
+ output_ty->dtype = computed_dtype;
+ if (call->ty_args.size() > 0) {
+ auto defined_ty = call->ty_args[0].as<TensorTypeNode>();
+ TVM_FFI_ICHECK(defined_ty);
+ auto shape = output_ty->GetShape();
+ TVM_FFI_ICHECK(shape.has_value());
+ TVM_FFI_ICHECK(defined_ty->vdevice.has_value());
+ inferred =
+ TensorType(ShapeExpr(shape.value()), output_ty->dtype, defined_ty->vdevice.value());
+ } else {
+ inferred = TensorType(output_ty);
+ }
+ }
+ if (const auto* old = call->ty.as<TensorTypeNode>()) {
+ auto same_optional_ref = [](const auto& lhs, const auto& rhs) {
+ return lhs.has_value() == rhs.has_value() &&
+ (!lhs.has_value() || lhs.value().same_as(rhs.value()));
+ };
+ bool same_output_dtype =
+ old->dtype.has_value() == inferred->dtype.has_value() &&
+ (!old->dtype.has_value() || old->dtype.value()->dtype == inferred->dtype.value()->dtype);
+ if (old->ndim == inferred->ndim && same_output_dtype &&
+ same_optional_ref(old->shape, inferred->shape) &&
+ same_optional_ref(old->vdevice, inferred->vdevice)) {
+ return call->ty;
+ }
}
+ return inferred;
+}
+
+template <bool require_float_dtype, typename FType>
+inline Type InferTypeUnary(const Call& call, const BlockBuilder&, FType f_compute_out_dtype) {
+ return InferTypeUnary<require_float_dtype>(call, f_compute_out_dtype);
}
/*!
@@ -215,6 +258,24 @@ Type ReturnTypeFromArg(const Call& call, const BlockBuilder& ctx) {
return GetType(call->args[arg_index]);
}
+template <int arg_index>
+ffi::Expected<Type> ReturnTypeFromArgContextFree(const CallNode* call_node) noexcept try {
+ Call call = ffi::GetRef<Call>(call_node);
+ Op op = call->op.as_or_throw<Op>();
+ CheckNumArguments(call);
+ int n_input = call->args.size();
+ if (arg_index >= n_input) {
+ TVM_FFI_VISIT_THROW(IndexError, call)
+ << op << " op has only " << n_input << "arguments, but try to get the arg with index "
+ << arg_index;
+ }
+ return GetType(call->args[arg_index]);
+} catch (const ffi::Error& error) {
+ return ffi::Unexpected(error);
+} catch (const std::exception& error) {
+ return ffi::Unexpected(ffi::Error("InternalError", error.what(), ""));
+}
+
/*!
* \brief Infer the type for unary arithmetic elementwise ops. It's also
* used in some NN operators.
@@ -229,6 +290,17 @@ Type InferTypeUnaryArith(const Call& call, const BlockBuilder& ctx) {
call, ctx, [](const TensorType& input_ty) { return input_ty->dtype; });
}
+template <bool require_float_dtype>
+ffi::Expected<Type> InferTypeUnaryArithContextFree(const CallNode* call_node) noexcept try {
+ Call call = ffi::GetRef<Call>(call_node);
+ return InferTypeUnary<require_float_dtype>(
+ call, [](const TensorType& input_ty) { return input_ty->dtype; });
+} catch (const ffi::Error& error) {
+ return ffi::Unexpected(error);
+} catch (const std::exception& error) {
+ return ffi::Unexpected(ffi::Error("InternalError", error.what(), ""));
+}
+
/*!
* \brief SLayout infer util for unary elementwise ops. It will simply take the layout of the input.
* \param call The context Call to the operator.
@@ -273,8 +345,8 @@ inline std::optional<PrimType> GetElementDType(const Type& ty) {
* \return The inferred output dtype.
* \throw Throw exception if the dtype of two input TensorType don’t match
*/
-inline ffi::Optional<PrimType> InferBinaryArithOpOutDtype(const Call& call, const BlockBuilder& ctx,
- const Type& lhs_ty, const Type& rhs_ty) {
+inline ffi::Optional<PrimType> InferBinaryArithOpOutDtype(const Call& call, const Type& lhs_ty,
+ const Type& rhs_ty) {
auto opt_lhs_dtype = GetElementDType(lhs_ty);
if (!opt_lhs_dtype) {
if (const auto* lhs_tensor = lhs_ty.as<TensorTypeNode>()) {
@@ -311,6 +383,11 @@ inline ffi::Optional<PrimType> InferBinaryArithOpOutDtype(const Call& call, cons
return lhs_dtype;
}
+inline ffi::Optional<PrimType> InferBinaryArithOpOutDtype(const Call& call, const BlockBuilder&,
+ const Type& lhs_ty, const Type& rhs_ty) {
+ return InferBinaryArithOpOutDtype(call, lhs_ty, rhs_ty);
+}
+
/*!
* \brief Infer the output virtual device for binary arithmetic operators.
* \param call The context Call to the operator.
@@ -320,9 +397,8 @@ inline ffi::Optional<PrimType> InferBinaryArithOpOutDtype(const Call& call, cons
* \return The inferred output vdevice.
* \throw Throw exception if the vdevice of two input TensorType don’t match
*/
-inline ffi::Optional<VDevice> InferBinaryArithOpOutVDevice(const Call& call,
- const BlockBuilder& ctx,
- const Type& lhs_ty, const Type& rhs_ty) {
+inline ffi::Optional<VDevice> InferBinaryArithOpOutVDevice(const Call& call, const Type& lhs_ty,
+ const Type& rhs_ty) {
auto get_vdevice = [&](const Type& ty) -> ffi::Optional<VDevice> {
if (const auto* tensor = ty.as<TensorTypeNode>()) {
return tensor->vdevice;
@@ -359,6 +435,11 @@ inline ffi::Optional<VDevice> InferBinaryArithOpOutVDevice(const Call& call,
return lhs_vdevice;
}
+inline ffi::Optional<VDevice> InferBinaryArithOpOutVDevice(const Call& call, const BlockBuilder&,
+ const Type& lhs_ty, const Type& rhs_ty) {
+ return InferBinaryArithOpOutVDevice(call, lhs_ty, rhs_ty);
+}
+
/*! \brief Result of binary broadcast shape inference without diagnostic context. */
struct BinaryBroadcastShapeInferResult {
enum class Status {
@@ -413,8 +494,11 @@ ffi::Optional<ffi::Array<PrimExpr>> InferBinaryBroadcastShape(const Call& call,
* \return The input axes in non-negative indexing.
* \throw Throw exception if there exists out-of-range axis index or repetitive indices.
*/
-std::vector<int> NormalizeAxes(const Call& call, const BlockBuilder& ctx, int ndim,
- const ffi::Array<int64_t>& axes);
+std::vector<int> NormalizeAxes(const Call& call, int ndim, const ffi::Array<int64_t>& axes);
+inline std::vector<int> NormalizeAxes(const Call& call, const BlockBuilder&, int ndim,
+ const ffi::Array<int64_t>& axes) {
+ return NormalizeAxes(call, ndim, axes);
+}
/*!
* \brief Convert the given axis to non-negative index. Meanwhile check if the axis is in range
@@ -426,8 +510,11 @@ std::vector<int> NormalizeAxes(const Call& call, const BlockBuilder& ctx, int nd
* \return The input axis in non-negative indexing.
* \throw Throw exception the given axis is out-of-range.
*/
-inline int NormalizeAxis(const Call& call, const BlockBuilder& ctx, int ndim, int axis) {
- return NormalizeAxes(call, ctx, ndim, {axis})[0];
+inline int NormalizeAxis(const Call& call, int ndim, int axis) {
+ return NormalizeAxes(call, ndim, {axis})[0];
+}
+inline int NormalizeAxis(const Call& call, const BlockBuilder&, int ndim, int axis) {
+ return NormalizeAxis(call, ndim, axis);
}
/*!
@@ -540,8 +627,8 @@ inline ffi::Array<int64_t> GetCompletePadding3D(ffi::Array<int64_t> padding) {
* tirx::SBijectiveLayout accordingly.
*/
inline std::pair<tirx::SLayout, tirx::SBijectiveLayout> CheckTensorLayout(
- const Call& call, const BlockBuilder& ctx, const ffi::String& tensor_layout,
- const ffi::String& tgt_layout, const ffi::String& tensor_name) {
+ const Call& call, const ffi::String& tensor_layout, const ffi::String& tgt_layout,
+ const ffi::String& tensor_name) {
tvm::PrimType i64_ty = tvm::PrimType::Int(64);
tirx::SLayout _tensor_layout(tensor_layout, i64_ty);
tirx::SBijectiveLayout tensor2tgt(_tensor_layout, tirx::SLayout(tgt_layout, i64_ty));
@@ -554,6 +641,12 @@ inline std::pair<tirx::SLayout, tirx::SBijectiveLayout> CheckTensorLayout(
return {_tensor_layout, tensor2tgt};
}
+inline std::pair<tirx::SLayout, tirx::SBijectiveLayout> CheckTensorLayout(
+ const Call& call, const BlockBuilder&, const ffi::String& tensor_layout,
+ const ffi::String& tgt_layout, const ffi::String& tensor_name) {
+ return CheckTensorLayout(call, tensor_layout, tgt_layout, tensor_name);
+}
+
/*!
* \brief Check if the given tensor type has expected ndim per the given layout (or the ndim
* is unknown), and try to cast the shape to ShapeExpr.
@@ -564,7 +657,6 @@ inline std::pair<tirx::SLayout, tirx::SBijectiveLayout> CheckTensorLayout(
* \return The shape of the input tensor in ShapeExpr, or `std::nullopt` if the shape is unknown.
*/
inline ffi::Optional<ShapeExpr> CheckNdimPerLayoutAndGetShape(const Call& call,
- const BlockBuilder& ctx,
const TensorType& ty,
const tirx::SLayout& layout) {
if (!ty->IsUnknownNdim() && ty->ndim != static_cast<int>(layout.ndim())) {
@@ -578,6 +670,12 @@ inline ffi::Optional<ShapeExpr> CheckNdimPerLayoutAndGetShape(const Call& call,
return std::nullopt;
}
+inline ffi::Optional<ShapeExpr> CheckNdimPerLayoutAndGetShape(const Call& call, const BlockBuilder&,
+ const TensorType& ty,
+ const tirx::SLayout& layout) {
+ return CheckNdimPerLayoutAndGetShape(call, ty, layout);
+}
+
Expr MakeVMAllocStorage(Expr size, PrimExpr runtime_device_index, DataTypeImm dtype,
StringImm storage_scope = StringImm("global"));
Expr MakeVMAllocTensor(Expr storage, PrimExpr offset, Expr shape, DataTypeImm dtype,
diff --git a/src/relax/op/tensor/binary.cc b/src/relax/op/tensor/binary.cc
index 1337db324e..2e1f0522e5 100644
--- a/src/relax/op/tensor/binary.cc
+++ b/src/relax/op/tensor/binary.cc
@@ -127,7 +127,10 @@ Type InferTypeBroadcast(const Call& call, const BlockBuilder& ctx, FType f_compu
}
Type InferTypeBroadcastArith(const Call& call, const BlockBuilder& ctx) {
- return InferTypeBroadcast(call, ctx, InferBinaryArithOpOutDtype);
+ return InferTypeBroadcast(
+ call, ctx, [](const Call& call, const BlockBuilder& ctx, const Type& lhs, const Type& rhs) {
+ return InferBinaryArithOpOutDtype(call, ctx, lhs, rhs);
+ });
}
Type InferTypeBroadcastCMP(const Call& call, const BlockBuilder& ctx) {
diff --git a/src/relax/op/tensor/binary.h b/src/relax/op/tensor/binary.h
index 1700789465..7693045210 100644
--- a/src/relax/op/tensor/binary.h
+++ b/src/relax/op/tensor/binary.h
@@ -49,12 +49,14 @@ namespace relax {
}
#define RELAX_REGISTER_BINARY_BROADCAST_OP_AND_IMPL(OpName) \
- RELAX_REGISTER_BINARY_OP_AND_IMPL(OpName, \
- .set_attr<FInferType>("FInferType", InferTypeBroadcastArith))
+ RELAX_REGISTER_BINARY_OP_AND_IMPL( \
+ OpName, \
+ .set_attr<FInferTypeWithBuilder>("relax.FInferTypeWithBuilder", InferTypeBroadcastArith))
#define RELAX_REGISTER_CMP_OP_AND_IMPL(OpName) \
- RELAX_REGISTER_BINARY_OP_AND_IMPL(OpName, \
- .set_attr<FInferType>("FInferType", InferTypeBroadcastCMP))
+ RELAX_REGISTER_BINARY_OP_AND_IMPL( \
+ OpName, \
+ .set_attr<FInferTypeWithBuilder>("relax.FInferTypeWithBuilder", InferTypeBroadcastCMP))
/***************** Arithmetic operators *****************/
diff --git a/src/relax/op/tensor/create.cc b/src/relax/op/tensor/create.cc
index f7dfe9e20a..d23c943be5 100644
--- a/src/relax/op/tensor/create.cc
+++ b/src/relax/op/tensor/create.cc
@@ -71,7 +71,8 @@ TVM_FFI_STATIC_INIT_BLOCK() {
refl::GlobalDef().def("relax.op.full", full);
}
-Type InferTypeFull(const Call& call, const BlockBuilder& ctx) {
+ffi::Expected<Type> InferTypeFull(const CallNode* call_node) noexcept try {
+ const Call call = ffi::GetRef<Call>(call_node);
if (call->args.size() != 2) {
TVM_FFI_VISIT_THROW(ValueError, call) << "Full op should have 2 arguments";
}
@@ -93,6 +94,10 @@ Type InferTypeFull(const Call& call, const BlockBuilder& ctx) {
? ffi::Optional<PrimType>(PrimType(attrs->dtype.value()))
: fill_value_ty->dtype;
return TensorType(/*shape=*/call->args[0], out_dtype, fill_value_ty->vdevice);
+} catch (const ffi::Error& error) {
+ return ffi::Unexpected(error);
+} catch (const std::exception& error) {
+ return ffi::Unexpected(ffi::Error("InternalError", error.what(), ""));
}
TVM_FFI_STATIC_INIT_BLOCK() {
@@ -100,7 +105,7 @@ TVM_FFI_STATIC_INIT_BLOCK() {
.signature(sig::arg("shape", "The shape of the created tensor."),
sig::arg("fill_value", "The scalar tensor, denoting the value to fill."),
sig::call_attrs<InitAttrs>())
- .set_attr<FInferType>("FInferType", InferTypeFull)
+ .set_attr<FInferType>("FInferType", FInferType::FromNative<&InferTypeFull>())
.set_attr<bool>("RequiresArgumentShapes", false)
.set_attr<bool>("FDataDependent", true)
.set_attr<TMixedPrecisionPolicy>("TMixedPrecisionPolicy", MixedPrecisionPolicyKind::kFollow)
@@ -121,8 +126,9 @@ TVM_FFI_STATIC_INIT_BLOCK() {
refl::GlobalDef().def("relax.op.full_like", full_like);
}
-Type InferTypeFullLike(const Call& call, const BlockBuilder& ctx) {
- ffi::Array<TensorType> input_ty = GetInputTensorType(call, ctx);
+ffi::Expected<Type> InferTypeFullLike(const CallNode* call_node) noexcept try {
+ const Call call = ffi::GetRef<Call>(call_node);
+ ffi::Array<TensorType> input_ty = GetInputTensorType(call);
TensorType data_ty = input_ty[0];
TensorType fill_value_ty = input_ty[1];
if (fill_value_ty->ndim != 0) {
@@ -139,19 +145,24 @@ Type InferTypeFullLike(const Call& call, const BlockBuilder& ctx) {
output_ty->dtype = PrimType(attrs->dtype.value());
return TensorType(output_ty);
}
+} catch (const ffi::Error& error) {
+ return ffi::Unexpected(error);
+} catch (const std::exception& error) {
+ return ffi::Unexpected(ffi::Error("InternalError", error.what(), ""));
}
TVM_FFI_STATIC_INIT_BLOCK() {
OpDef("relax.full_like")
.signature(sig::arg("x", "The input tensor."),
sig::arg("fill_value", "The scalar value to fill."), sig::call_attrs<InitAttrs>())
- .set_attr<FInferType>("FInferType", InferTypeFullLike)
+ .set_attr<FInferType>("FInferType", FInferType::FromNative<&InferTypeFullLike>())
.set_attr<TMixedPrecisionPolicy>("TMixedPrecisionPolicy", MixedPrecisionPolicyKind::kFollow)
.set_attr<bool>("FPurity", true);
}
// Structure info inference for ones and zeros
-Type InferTypeOnesZeros(const Call& call, const BlockBuilder& ctx) {
+ffi::Expected<Type> InferTypeOnesZeros(const CallNode* call_node) noexcept try {
+ const Call call = ffi::GetRef<Call>(call_node);
if (call->args.size() != 1) {
TVM_FFI_VISIT_THROW(ValueError, call) << "Ones/Zeros should have 1 argument";
}
@@ -165,11 +176,16 @@ Type InferTypeOnesZeros(const Call& call, const BlockBuilder& ctx) {
const auto* attrs = call->attrs.as<InitAttrs>();
TVM_FFI_ICHECK(attrs->dtype.has_value());
return TensorType(/*shape=*/call->args[0], PrimType(attrs->dtype.value()));
+} catch (const ffi::Error& error) {
+ return ffi::Unexpected(error);
+} catch (const std::exception& error) {
+ return ffi::Unexpected(ffi::Error("InternalError", error.what(), ""));
}
// Structure info inference for ones_like and zeros_like
-Type InferTypeOnesLikeZerosLike(const Call& call, const BlockBuilder& ctx) {
- TensorType data_ty = GetUnaryInputTensorType(call, ctx);
+ffi::Expected<Type> InferTypeOnesLikeZerosLike(const CallNode* call_node) noexcept try {
+ const Call call = ffi::GetRef<Call>(call_node);
+ TensorType data_ty = GetUnaryInputTensorType(call);
const auto* attrs = call->attrs.as<InitAttrs>();
if (!attrs->dtype.has_value()) {
return data_ty;
@@ -178,6 +194,10 @@ Type InferTypeOnesLikeZerosLike(const Call& call, const BlockBuilder& ctx) {
output_ty->dtype = PrimType(attrs->dtype.value());
return TensorType(output_ty);
}
+} catch (const ffi::Error& error) {
+ return ffi::Unexpected(error);
+} catch (const std::exception& error) {
+ return ffi::Unexpected(ffi::Error("InternalError", error.what(), ""));
}
/* relax.ones & relax.ones_like */
@@ -203,13 +223,13 @@ TVM_FFI_STATIC_INIT_BLOCK() {
OpDef("relax.ones")
.signature(sig::arg("shape", "The shape of the created tensor."),
sig::call_attrs<InitAttrs>())
- .set_attr<FInferType>("FInferType", InferTypeOnesZeros)
+ .set_attr<FInferType>("FInferType", FInferType::FromNative<&InferTypeOnesZeros>())
.set_attr<TMixedPrecisionPolicy>("TMixedPrecisionPolicy", MixedPrecisionPolicyKind::kFollow)
.set_attr<bool>("FPurity", true);
OpDef("relax.ones_like")
.signature(sig::arg("x", "The input tensor."), sig::call_attrs<InitAttrs>())
- .set_attr<FInferType>("FInferType", InferTypeOnesLikeZerosLike)
+ .set_attr<FInferType>("FInferType", FInferType::FromNative<&InferTypeOnesLikeZerosLike>())
.set_attr<bool>("FPurity", true);
}
@@ -236,13 +256,13 @@ TVM_FFI_STATIC_INIT_BLOCK() {
OpDef("relax.zeros")
.signature(sig::arg("shape", "The shape of the created tensor."),
sig::call_attrs<InitAttrs>())
- .set_attr<FInferType>("FInferType", InferTypeOnesZeros)
+ .set_attr<FInferType>("FInferType", FInferType::FromNative<&InferTypeOnesZeros>())
.set_attr<TMixedPrecisionPolicy>("TMixedPrecisionPolicy", MixedPrecisionPolicyKind::kFollow)
.set_attr<bool>("FPurity", true);
OpDef("relax.zeros_like")
.signature(sig::arg("x", "The input tensor."), sig::call_attrs<InitAttrs>())
- .set_attr<FInferType>("FInferType", InferTypeOnesLikeZerosLike)
+ .set_attr<FInferType>("FInferType", FInferType::FromNative<&InferTypeOnesLikeZerosLike>())
.set_attr<bool>("FPurity", true);
}
@@ -267,7 +287,8 @@ TVM_FFI_STATIC_INIT_BLOCK() {
refl::GlobalDef().def("relax.op.eye", eye).def("relax.op.eye_like", eye_like);
}
-Type InferTypeEye(const Call& call, const BlockBuilder& ctx) {
+ffi::Expected<Type> InferTypeEye(const CallNode* call_node) noexcept try {
+ const Call call = ffi::GetRef<Call>(call_node);
if (call->args.size() != 3) {
TVM_FFI_VISIT_THROW(ValueError, call) << "Eye op should have 3 arguments: n, m, and k, but got "
<< call->args.size() << " arguments";
@@ -289,9 +310,14 @@ Type InferTypeEye(const Call& call, const BlockBuilder& ctx) {
TVM_FFI_ICHECK(attrs->dtype.has_value());
DLDataType dtype = attrs->dtype.value();
return TensorType(ShapeExpr({n, m}), PrimType(dtype));
+} catch (const ffi::Error& error) {
+ return ffi::Unexpected(error);
+} catch (const std::exception& error) {
+ return ffi::Unexpected(ffi::Error("InternalError", error.what(), ""));
}
-Type InferTypeEyeLike(const Call& call, const BlockBuilder& ctx) {
+ffi::Expected<Type> InferTypeEyeLike(const CallNode* call_node) noexcept try {
+ const Call call = ffi::GetRef<Call>(call_node);
if (call->args.size() != 2) {
TVM_FFI_VISIT_THROW(ValueError, call)
<< "Eye_like op should have 2 arguments: x and k, but got " << call->args.size()
@@ -316,6 +342,10 @@ Type InferTypeEyeLike(const Call& call, const BlockBuilder& ctx) {
: x_ty->dtype;
return TensorType(x_ty->shape.value(), out_dtype, x_ty->vdevice);
+} catch (const ffi::Error& error) {
+ return ffi::Unexpected(error);
+} catch (const std::exception& error) {
+ return ffi::Unexpected(ffi::Error("InternalError", error.what(), ""));
}
TVM_FFI_STATIC_INIT_BLOCK() {
@@ -323,14 +353,14 @@ TVM_FFI_STATIC_INIT_BLOCK() {
.signature(sig::arg("n", "Number of rows in the output."),
sig::arg("m", "Number of columns in the output."),
sig::arg("k", "Index of the diagonal."), sig::call_attrs<InitAttrs>())
- .set_attr<FInferType>("FInferType", InferTypeEye)
+ .set_attr<FInferType>("FInferType", FInferType::FromNative<&InferTypeEye>())
.set_attr<TMixedPrecisionPolicy>("TMixedPrecisionPolicy", MixedPrecisionPolicyKind::kFollow)
.set_attr<bool>("FPurity", true);
OpDef("relax.eye_like")
.signature(sig::arg("x", "The input tensor."), sig::arg("k", "Index of the diagonal."),
sig::call_attrs<InitAttrs>())
- .set_attr<FInferType>("FInferType", InferTypeEyeLike)
+ .set_attr<FInferType>("FInferType", FInferType::FromNative<&InferTypeEyeLike>())
.set_attr<bool>("FPurity", true);
}
@@ -348,7 +378,8 @@ TVM_FFI_STATIC_INIT_BLOCK() {
refl::GlobalDef().def("relax.op.arange", arange);
}
-Type InferTypeArange(const Call& call, const BlockBuilder& ctx) {
+ffi::Expected<Type> InferTypeArange(const CallNode* call_node) noexcept try {
+ const Call call = ffi::GetRef<Call>(call_node);
if (call->args.size() != 3) {
TVM_FFI_VISIT_THROW(ValueError, call)
<< "Arange should have 3 arguments, which are `start`, `end` and `step`, but got "
@@ -381,6 +412,10 @@ Type InferTypeArange(const Call& call, const BlockBuilder& ctx) {
sym::Analyzer analyzer;
num_elem = analyzer->Simplify(num_elem);
return TensorType(ShapeExpr({num_elem}), PrimType(dtype));
+} catch (const ffi::Error& error) {
+ return ffi::Unexpected(error);
+} catch (const std::exception& error) {
+ return ffi::Unexpected(ffi::Error("InternalError", error.what(), ""));
}
TVM_FFI_STATIC_INIT_BLOCK() {
@@ -389,7 +424,7 @@ TVM_FFI_STATIC_INIT_BLOCK() {
sig::arg("end", "The ending value for the set of points."),
sig::arg("step", "The gap between each pair of adjacent points."),
sig::call_attrs<InitAttrs>())
- .set_attr<FInferType>("FInferType", InferTypeArange)
+ .set_attr<FInferType>("FInferType", FInferType::FromNative<&InferTypeArange>())
.set_attr<TMixedPrecisionPolicy>("TMixedPrecisionPolicy", MixedPrecisionPolicyKind::kFollow)
.set_attr<bool>("FPurity", true);
}
@@ -411,7 +446,8 @@ TVM_FFI_STATIC_INIT_BLOCK() {
refl::GlobalDef().def("relax.op.hamming_window", hamming_window);
}
-Type InferTypeHammingWindow(const Call& call, const BlockBuilder& ctx) {
+ffi::Expected<Type> InferTypeHammingWindow(const CallNode* call_node) noexcept try {
+ const Call call = ffi::GetRef<Call>(call_node);
const auto* attrs = call->attrs.as<InitAttrs>();
TVM_FFI_ICHECK(attrs->dtype.has_value());
DLDataType dtype = attrs->dtype.value();
@@ -437,6 +473,10 @@ Type InferTypeHammingWindow(const Call& call, const BlockBuilder& ctx) {
}
window_size = analyzer->Simplify(window_size);
return TensorType(ShapeExpr({window_size}), PrimType(dtype));
+} catch (const ffi::Error& error) {
+ return ffi::Unexpected(error);
+} catch (const std::exception& error) {
+ return ffi::Unexpected(ffi::Error("InternalError", error.what(), ""));
}
TVM_FFI_STATIC_INIT_BLOCK() {
@@ -448,7 +488,7 @@ TVM_FFI_STATIC_INIT_BLOCK() {
"symmetric window"),
sig::arg("alpha", "The coefficient alpha"), sig::arg("beta", "The coefficient beta"),
sig::call_attrs<InitAttrs>())
- .set_attr<FInferType>("FInferType", InferTypeHammingWindow)
+ .set_attr<FInferType>("FInferType", FInferType::FromNative<&InferTypeHammingWindow>())
.set_attr<TMixedPrecisionPolicy>("TMixedPrecisionPolicy", MixedPrecisionPolicyKind::kFollow)
.set_attr<bool>("FPurity", true);
}
@@ -472,8 +512,9 @@ TVM_FFI_STATIC_INIT_BLOCK() {
.def("relax.op.triu", static_cast<Expr (*)(Expr, Expr)>(triu));
}
-Type InferTypeTrilTriu(const Call& call, const BlockBuilder& ctx) {
- auto [data_ty, offset] = GetArgType<TensorType, PrimType>(call, ctx);
+ffi::Expected<Type> InferTypeTrilTriu(const CallNode* call_node) noexcept try {
+ const Call call = ffi::GetRef<Call>(call_node);
+ auto [data_ty, offset] = GetArgType<TensorType, PrimType>(call);
if (!data_ty->IsUnknownNdim() && data_ty->ndim < 2) {
TVM_FFI_VISIT_THROW(ValueError, call) << call->op
@@ -482,17 +523,21 @@ Type InferTypeTrilTriu(const Call& call, const BlockBuilder& ctx) {
<< data_ty->ndim << " dimension(s).";
}
return data_ty;
+} catch (const ffi::Error& error) {
+ return ffi::Unexpected(error);
+} catch (const std::exception& error) {
+ return ffi::Unexpected(ffi::Error("InternalError", error.what(), ""));
}
TVM_FFI_STATIC_INIT_BLOCK() {
OpDef("relax.tril")
.signature(sig::arg("x", "The input tensor."), sig::arg("k", "The offset of the diagonal."))
- .set_attr<FInferType>("FInferType", InferTypeTrilTriu)
+ .set_attr<FInferType>("FInferType", FInferType::FromNative<&InferTypeTrilTriu>())
.set_attr<bool>("FPurity", true);
OpDef("relax.triu")
.signature(sig::arg("x", "The input tensor."), sig::arg("k", "The offset of the diagonal."))
- .set_attr<FInferType>("FInferType", InferTypeTrilTriu)
+ .set_attr<FInferType>("FInferType", FInferType::FromNative<&InferTypeTrilTriu>())
.set_attr<bool>("FPurity", true);
}
diff --git a/src/relax/op/tensor/datatype.cc b/src/relax/op/tensor/datatype.cc
index 58b78e27bc..bd60f13ee2 100644
--- a/src/relax/op/tensor/datatype.cc
+++ b/src/relax/op/tensor/datatype.cc
@@ -51,18 +51,23 @@ TVM_FFI_STATIC_INIT_BLOCK() {
refl::GlobalDef().def("relax.op.astype", astype);
}
-Type InferTypeAstype(const Call& call, const BlockBuilder& ctx) {
- TensorType ty = GetUnaryInputTensorType(call, ctx);
+ffi::Expected<Type> InferTypeAstype(const CallNode* call_node) noexcept try {
+ const Call call = ffi::GetRef<Call>(call_node);
+ TensorType ty = GetUnaryInputTensorType(call);
const auto* attrs = call->attrs.as<AstypeAttrs>();
ffi::ObjectPtr<TensorTypeNode> new_ty = ffi::make_object<TensorTypeNode>(*ty.get());
new_ty->dtype = PrimType(attrs->dtype);
return TensorType(new_ty);
+} catch (const ffi::Error& error) {
+ return ffi::Unexpected(error);
+} catch (const std::exception& error) {
+ return ffi::Unexpected(ffi::Error("InternalError", error.what(), ""));
}
TVM_FFI_STATIC_INIT_BLOCK() {
OpDef("relax.astype")
.signature(sig::arg("x", "The input tensor"), sig::call_attrs<AstypeAttrs>())
- .set_attr<FInferType>("FInferType", InferTypeAstype)
+ .set_attr<FInferType>("FInferType", FInferType::FromNative<&InferTypeAstype>())
.set_attr<FRelaxInferLayout>("FRelaxInferLayout", InferLayoutUnaryEwise)
.set_attr<TMixedPrecisionPolicy>("TMixedPrecisionPolicy", MixedPrecisionPolicyKind::kFollow)
.set_attr<bool>("FPurity", true);
@@ -83,18 +88,23 @@ TVM_FFI_STATIC_INIT_BLOCK() {
refl::GlobalDef().def("relax.op.wrap_param", MakeWrapParam);
}
-Type InferTypeWrapParam(const Call& call, const BlockBuilder& ctx) {
- TensorType ty = GetUnaryInputTensorType(call, ctx);
+ffi::Expected<Type> InferTypeWrapParam(const CallNode* call_node) noexcept try {
+ const Call call = ffi::GetRef<Call>(call_node);
+ TensorType ty = GetUnaryInputTensorType(call);
const auto* attrs = call->attrs.as<WrapParamAttrs>();
ffi::ObjectPtr<TensorTypeNode> new_ty = ffi::make_object<TensorTypeNode>(*ty.get());
new_ty->dtype = PrimType(attrs->dtype);
return TensorType(new_ty);
+} catch (const ffi::Error& error) {
+ return ffi::Unexpected(error);
+} catch (const std::exception& error) {
+ return ffi::Unexpected(ffi::Error("InternalError", error.what(), ""));
}
TVM_FFI_STATIC_INIT_BLOCK() {
OpDef("relax.wrap_param")
.signature(sig::arg("data", "The input tensor"), sig::call_attrs<WrapParamAttrs>())
- .set_attr<FInferType>("FInferType", InferTypeWrapParam)
+ .set_attr<FInferType>("FInferType", FInferType::FromNative<&InferTypeWrapParam>())
.set_attr<bool>("FPurity", true);
}
diff --git a/src/relax/op/tensor/grad.cc b/src/relax/op/tensor/grad.cc
index a113512e6c..e0e109442f 100644
--- a/src/relax/op/tensor/grad.cc
+++ b/src/relax/op/tensor/grad.cc
@@ -43,12 +43,19 @@ TVM_FFI_STATIC_INIT_BLOCK() {
refl::GlobalDef().def("relax.op.grad.no_grad", no_grad);
}
-Type InferTypeNoGrad(const Call& call, const BlockBuilder& ctx) { return GetType(call->args[0]); }
+ffi::Expected<Type> InferTypeNoGrad(const CallNode* call_node) noexcept try {
+ const Call call = ffi::GetRef<Call>(call_node);
+ return GetType(call->args[0]);
+} catch (const ffi::Error& error) {
+ return ffi::Unexpected(error);
+} catch (const std::exception& error) {
+ return ffi::Unexpected(ffi::Error("InternalError", error.what(), ""));
+}
TVM_FFI_STATIC_INIT_BLOCK() {
OpDef("relax.grad.no_grad")
.signature(sig::arg("x", "The corresponding input tensor."))
- .set_attr<FInferType>("FInferType", InferTypeNoGrad)
+ .set_attr<FInferType>("FInferType", FInferType::FromNative<&InferTypeNoGrad>())
.set_attr<bool>("FPurity", true);
}
@@ -63,18 +70,23 @@ TVM_FFI_STATIC_INIT_BLOCK() {
refl::GlobalDef().def("relax.op.grad.start_checkpoint", start_checkpoint);
}
-Type InferTypeStartCheckpoint(const Call& call, const BlockBuilder& ctx) {
+ffi::Expected<Type> InferTypeStartCheckpoint(const CallNode* call_node) noexcept try {
+ const Call call = ffi::GetRef<Call>(call_node);
if (!call->args[0].as<VarNode>()) {
TVM_FFI_VISIT_THROW(TypeError, call)
<< "The argument of relax.op.grad.start_checkpoint should be a Var.";
}
return GetType(call->args[0]);
+} catch (const ffi::Error& error) {
+ return ffi::Unexpected(error);
+} catch (const std::exception& error) {
+ return ffi::Unexpected(ffi::Error("InternalError", error.what(), ""));
}
TVM_FFI_STATIC_INIT_BLOCK() {
OpDef("relax.grad.start_checkpoint")
.signature(sig::arg("x", "The tensor marking the input of the checkpoint stage."))
- .set_attr<FInferType>("FInferType", InferTypeStartCheckpoint)
+ .set_attr<FInferType>("FInferType", FInferType::FromNative<&InferTypeStartCheckpoint>())
.set_attr<bool>("FPurity", true);
}
@@ -89,18 +101,23 @@ TVM_FFI_STATIC_INIT_BLOCK() {
refl::GlobalDef().def("relax.op.grad.end_checkpoint", end_checkpoint);
}
-Type InferTypeEndCheckpoint(const Call& call, const BlockBuilder& ctx) {
+ffi::Expected<Type> InferTypeEndCheckpoint(const CallNode* call_node) noexcept try {
+ const Call call = ffi::GetRef<Call>(call_node);
if (!call->args[0].as<VarNode>()) {
TVM_FFI_VISIT_THROW(TypeError, call)
<< "The argument of relax.op.grad.end_checkpoint should be a Var.";
}
return GetType(call->args[0]);
+} catch (const ffi::Error& error) {
+ return ffi::Unexpected(error);
+} catch (const std::exception& error) {
+ return ffi::Unexpected(ffi::Error("InternalError", error.what(), ""));
}
TVM_FFI_STATIC_INIT_BLOCK() {
OpDef("relax.grad.end_checkpoint")
.signature(sig::arg("x", "The output of the checkpoint stage."))
- .set_attr<FInferType>("FInferType", InferTypeEndCheckpoint)
+ .set_attr<FInferType>("FInferType", FInferType::FromNative<&InferTypeEndCheckpoint>())
.set_attr<bool>("FPurity", true);
}
@@ -130,8 +147,13 @@ TVM_FFI_STATIC_INIT_BLOCK() {
refl::GlobalDef().def("relax.op.grad.nll_loss_backward", nll_loss_backward);
}
-Type InferTypeNLLLossBackward(const Call& call, const BlockBuilder& ctx) {
+ffi::Expected<Type> InferTypeNLLLossBackward(const CallNode* call_node) noexcept try {
+ const Call call = ffi::GetRef<Call>(call_node);
return GetType(call->args[1]);
+} catch (const ffi::Error& error) {
+ return ffi::Unexpected(error);
+} catch (const std::exception& error) {
+ return ffi::Unexpected(ffi::Error("InternalError", error.what(), ""));
}
TVM_FFI_STATIC_INIT_BLOCK() {
@@ -140,7 +162,7 @@ TVM_FFI_STATIC_INIT_BLOCK() {
sig::arg("predictions", "The prediction tensor."),
sig::arg("targets", "The target tensor."), sig::var_args("args"),
sig::call_attrs<NLLLossAttrs>())
- .set_attr<FInferType>("FInferType", InferTypeNLLLossBackward)
+ .set_attr<FInferType>("FInferType", FInferType::FromNative<&InferTypeNLLLossBackward>())
.set_attr<bool>("FPurity", true);
}
@@ -168,15 +190,20 @@ TVM_FFI_STATIC_INIT_BLOCK() {
refl::GlobalDef().def("relax.op.grad.max_pool2d_backward", max_pool2d_backward);
}
-Type InferTypeMaxPool2DBackward(const Call& call, const BlockBuilder& ctx) {
+ffi::Expected<Type> InferTypeMaxPool2DBackward(const CallNode* call_node) noexcept try {
+ const Call call = ffi::GetRef<Call>(call_node);
return GetType(call->args[1]);
+} catch (const ffi::Error& error) {
+ return ffi::Unexpected(error);
+} catch (const std::exception& error) {
+ return ffi::Unexpected(ffi::Error("InternalError", error.what(), ""));
}
TVM_FFI_STATIC_INIT_BLOCK() {
OpDef("relax.grad.max_pool2d_backward")
.signature(sig::arg("output_grad", "The output gradient."),
sig::arg("data", "The input tensor"), sig::call_attrs<Pool2DAttrs>())
- .set_attr<FInferType>("FInferType", InferTypeMaxPool2DBackward)
+ .set_attr<FInferType>("FInferType", FInferType::FromNative<&InferTypeMaxPool2DBackward>())
.set_attr<bool>("FPurity", true);
}
@@ -204,15 +231,20 @@ TVM_FFI_STATIC_INIT_BLOCK() {
refl::GlobalDef().def("relax.op.grad.avg_pool2d_backward", avg_pool2d_backward);
}
-Type InferTypeAvgPool2DBackward(const Call& call, const BlockBuilder& ctx) {
+ffi::Expected<Type> InferTypeAvgPool2DBackward(const CallNode* call_node) noexcept try {
+ const Call call = ffi::GetRef<Call>(call_node);
return GetType(call->args[1]);
+} catch (const ffi::Error& error) {
+ return ffi::Unexpected(error);
+} catch (const std::exception& error) {
+ return ffi::Unexpected(ffi::Error("InternalError", error.what(), ""));
}
TVM_FFI_STATIC_INIT_BLOCK() {
OpDef("relax.grad.avg_pool2d_backward")
.signature(sig::arg("output_grad", "The output gradient."),
sig::arg("data", "The input tensor"), sig::call_attrs<Pool2DAttrs>())
- .set_attr<FInferType>("FInferType", InferTypeAvgPool2DBackward)
+ .set_attr<FInferType>("FInferType", FInferType::FromNative<&InferTypeAvgPool2DBackward>())
.set_attr<bool>("FPurity", true);
}
@@ -233,8 +265,13 @@ TVM_FFI_STATIC_INIT_BLOCK() {
refl::GlobalDef().def("relax.op.grad.take_backward", take_backward);
}
-Type InferTypeTakeBackward(const Call& call, const BlockBuilder& ctx) {
+ffi::Expected<Type> InferTypeTakeBackward(const CallNode* call_node) noexcept try {
+ const Call call = ffi::GetRef<Call>(call_node);
return GetType(call->args[1]);
+} catch (const ffi::Error& error) {
+ return ffi::Unexpected(error);
+} catch (const std::exception& error) {
+ return ffi::Unexpected(ffi::Error("InternalError", error.what(), ""));
}
TVM_FFI_STATIC_INIT_BLOCK() {
@@ -243,7 +280,7 @@ TVM_FFI_STATIC_INIT_BLOCK() {
sig::arg("x", "The source tensor."),
sig::arg("indices", "The indices of the values to extract."),
sig::call_attrs<TakeAttrs>())
- .set_attr<FInferType>("FInferType", InferTypeTakeBackward)
+ .set_attr<FInferType>("FInferType", FInferType::FromNative<&InferTypeTakeBackward>())
.set_attr<bool>("FPurity", true);
}
diff --git a/src/relax/op/tensor/index.cc b/src/relax/op/tensor/index.cc
index e1b73c3d15..67d7cf0181 100644
--- a/src/relax/op/tensor/index.cc
+++ b/src/relax/op/tensor/index.cc
@@ -61,9 +61,10 @@ TVM_FFI_STATIC_INIT_BLOCK() {
refl::GlobalDef().def("relax.op.take", take);
}
-Type InferTypeTake(const Call& call, const BlockBuilder& ctx) {
- CheckNumArguments(call, ctx);
- TensorType data_ty = GetInputTensorType(call, 0, ctx);
+ffi::Expected<Type> InferTypeTake(const CallNode* call_node) noexcept try {
+ const Call call = ffi::GetRef<Call>(call_node);
+ CheckNumArguments(call);
+ TensorType data_ty = GetInputTensorType(call, 0);
// Type inference when the index is a PrimExpr is equivalent
// to that of a scalar (0-d) tensor.
@@ -108,7 +109,7 @@ Type InferTypeTake(const Call& call, const BlockBuilder& ctx) {
int axis = 0;
if (attrs->axis.has_value()) {
- axis = NormalizeAxis(call, ctx, data_ty->ndim, attrs->axis.value());
+ axis = NormalizeAxis(call, data_ty->ndim, attrs->axis.value());
}
const auto* data_shape = data_ty->shape.as<ShapeExprNode>();
const auto* indices_shape = indices_ty->shape.as<ShapeExprNode>();
@@ -125,6 +126,10 @@ Type InferTypeTake(const Call& call, const BlockBuilder& ctx) {
}
}
return TensorType(ShapeExpr(output_shape), data_ty->dtype, data_ty->vdevice);
+} catch (const ffi::Error& error) {
+ return ffi::Unexpected(error);
+} catch (const std::exception& error) {
+ return ffi::Unexpected(ffi::Error("InternalError", error.what(), ""));
}
TVM_FFI_STATIC_INIT_BLOCK() {
@@ -132,7 +137,7 @@ TVM_FFI_STATIC_INIT_BLOCK() {
.signature(sig::arg("x", "The source tensor."),
sig::arg("indices", "The indices of the values to extract."),
sig::call_attrs<TakeAttrs>())
- .set_attr<FInferType>("FInferType", InferTypeTake)
+ .set_attr<FInferType>("FInferType", FInferType::FromNative<&InferTypeTake>())
.set_attr<bool>("FPurity", true);
}
@@ -492,7 +497,7 @@ TVM_FFI_STATIC_INIT_BLOCK() {
.signature(sig::arg("x", "The source tensor to be sliced."), sig::arg("axes", "The axes."),
sig::arg("begin", "The start index."), sig::arg("end", "The end index."),
sig::var_args("args"), sig::call_attrs<StridedSliceAttrs>())
- .set_attr<FInferType>("FInferType", InferTypeStridedSlice)
+ .set_attr<FInferTypeWithBuilder>("relax.FInferTypeWithBuilder", InferTypeStridedSlice)
.set_attr<FRelaxInferLayout>("FRelaxInferLayout", InferLayoutStridedSlice)
.set_attr<TMixedPrecisionPolicy>("TMixedPrecisionPolicy", MixedPrecisionPolicyKind::kFollow)
.set_attr<bool>("FPurity", true);
@@ -513,7 +518,8 @@ TVM_FFI_STATIC_INIT_BLOCK() {
refl::GlobalDef().def("relax.op.dynamic_strided_slice", dynamic_strided_slice);
}
-Type InferTypeDynStridedSlice(const Call& call, const BlockBuilder& ctx) {
+ffi::Expected<Type> InferTypeDynStridedSlice(const CallNode* call_node) noexcept try {
+ const Call call = ffi::GetRef<Call>(call_node);
const auto* data_ty = GetTypeAs<TensorTypeNode>(call->args[0]);
const auto* begin_ty = GetTypeAs<TensorTypeNode>(call->args[1]);
const auto* end_ty = GetTypeAs<TensorTypeNode>(call->args[2]);
@@ -566,6 +572,10 @@ Type InferTypeDynStridedSlice(const Call& call, const BlockBuilder& ctx) {
// TODO(tvm-team): Currently, it is unable to express partially-static shape. Revisit when
// PrimExpr lands.
return TensorType(data_ty->dtype, n_axis, data_ty->vdevice);
+} catch (const ffi::Error& error) {
+ return ffi::Unexpected(error);
+} catch (const std::exception& error) {
+ return ffi::Unexpected(ffi::Error("InternalError", error.what(), ""));
}
InferLayoutOutput InferLayoutDynStridedSlice(
@@ -592,7 +602,7 @@ TVM_FFI_STATIC_INIT_BLOCK() {
sig::arg("begin", "The indices to begin with in the slicing."),
sig::arg("end", "Indices indicating end of the slice."),
sig::arg("strides", "The stride values."))
- .set_attr<FInferType>("FInferType", InferTypeDynStridedSlice)
+ .set_attr<FInferType>("FInferType", FInferType::FromNative<&InferTypeDynStridedSlice>())
.set_attr<FRelaxInferLayout>("FRelaxInferLayout", InferLayoutDynStridedSlice)
.set_attr<TMixedPrecisionPolicy>("TMixedPrecisionPolicy", MixedPrecisionPolicyKind::kFollow)
.set_attr<bool>("FPurity", true)
diff --git a/src/relax/op/tensor/inspect.cc b/src/relax/op/tensor/inspect.cc
index 5ef6a9dec7..cfa264a0c3 100644
--- a/src/relax/op/tensor/inspect.cc
+++ b/src/relax/op/tensor/inspect.cc
@@ -137,7 +137,7 @@ Expr LegalizeTensorDtypeCode(const BlockBuilder& bb, const Call& call) {
TVM_FFI_STATIC_INIT_BLOCK() {
OpDef("relax.inspect.tensor_dtype_code")
.signature(sig::arg("tensor", "The tensor to be inspected"))
- .set_attr<FInferType>("FInferType", InferTypeTensorDtypeCode)
+ .set_attr<FInferTypeWithBuilder>("relax.FInferTypeWithBuilder", InferTypeTensorDtypeCode)
.set_attr<FLegalize>("FLegalize", LegalizeTensorDtypeCode)
.set_attr<bool>("RequiresArgumentShapes", false)
.set_attr<FNormalize>("FNormalize", NormalizeToKnownPrimExpr)
@@ -167,7 +167,7 @@ Expr LegalizeTensorDtypeBits(const BlockBuilder& bb, const Call& call) {
TVM_FFI_STATIC_INIT_BLOCK() {
OpDef("relax.inspect.tensor_dtype_bits")
.signature(sig::arg("tensor", "The tensor to be inspected"))
- .set_attr<FInferType>("FInferType", InferTypeTensorDtypeBits)
+ .set_attr<FInferTypeWithBuilder>("relax.FInferTypeWithBuilder", InferTypeTensorDtypeBits)
.set_attr<FLegalize>("FLegalize", LegalizeTensorDtypeBits)
.set_attr<bool>("RequiresArgumentShapes", false)
.set_attr<FNormalize>("FNormalize", NormalizeToKnownPrimExpr)
@@ -197,7 +197,7 @@ Expr LegalizeTensorDtypeLanes(const BlockBuilder& bb, const Call& call) {
TVM_FFI_STATIC_INIT_BLOCK() {
OpDef("relax.inspect.tensor_dtype_lanes")
.signature(sig::arg("tensor", "The tensor to be inspected"))
- .set_attr<FInferType>("FInferType", InferTypeTensorDtypeLanes)
+ .set_attr<FInferTypeWithBuilder>("relax.FInferTypeWithBuilder", InferTypeTensorDtypeLanes)
.set_attr<FLegalize>("FLegalize", LegalizeTensorDtypeLanes)
.set_attr<bool>("RequiresArgumentShapes", false)
.set_attr<FNormalize>("FNormalize", NormalizeToKnownPrimExpr)
@@ -227,7 +227,7 @@ Expr LegalizeTensorNDim(const BlockBuilder& bb, const Call& call) {
TVM_FFI_STATIC_INIT_BLOCK() {
OpDef("relax.inspect.tensor_ndim")
.signature(sig::arg("tensor", "The tensor to be inspected"))
- .set_attr<FInferType>("FInferType", InferTypeTensorNDim)
+ .set_attr<FInferTypeWithBuilder>("relax.FInferTypeWithBuilder", InferTypeTensorNDim)
.set_attr<FLegalize>("FLegalize", LegalizeTensorNDim)
.set_attr<bool>("RequiresArgumentShapes", false)
.set_attr<FNormalize>("FNormalize", NormalizeToKnownPrimExpr)
@@ -308,7 +308,7 @@ TVM_FFI_STATIC_INIT_BLOCK() {
OpDef("relax.inspect.tensor_shape_i")
.signature(sig::arg("tensor", "The tensor to be inspected"),
sig::arg("axis", "The axis whose extent should be returned"))
- .set_attr<FInferType>("FInferType", InferTypeTensorShape)
+ .set_attr<FInferTypeWithBuilder>("relax.FInferTypeWithBuilder", InferTypeTensorShape)
.set_attr<FLegalize>("FLegalize", LegalizeTensorShape)
.set_attr<bool>("RequiresArgumentShapes", false)
.set_attr<FNormalize>("FNormalize", NormalizeToKnownPrimExpr)
@@ -357,7 +357,7 @@ TVM_FFI_STATIC_INIT_BLOCK() {
OpDef("relax.inspect.tensor_stride_i")
.signature(sig::arg("tensor", "The tensor to be inspected"),
sig::arg("axis", "The axis whose extent should be returned"))
- .set_attr<FInferType>("FInferType", InferTypeTensorStride)
+ .set_attr<FInferTypeWithBuilder>("relax.FInferTypeWithBuilder", InferTypeTensorStride)
.set_attr<bool>("RequiresArgumentShapes", false)
.set_attr<FNormalize>("FNormalize", NormalizeToKnownPrimExpr)
.set_attr<bool>("FPurity", true);
@@ -389,7 +389,7 @@ Type InferTypeTensorByteOffset(const Call& call, const BlockBuilder&) {
TVM_FFI_STATIC_INIT_BLOCK() {
OpDef("relax.inspect.tensor_byte_offset")
.signature(sig::arg("tensor", "The tensor to be inspected"))
- .set_attr<FInferType>("FInferType", InferTypeTensorByteOffset)
+ .set_attr<FInferTypeWithBuilder>("relax.FInferTypeWithBuilder", InferTypeTensorByteOffset)
.set_attr<bool>("RequiresArgumentShapes", false)
.set_attr<FNormalize>("FNormalize", NormalizeToKnownPrimExpr)
.set_attr<bool>("FPurity", true);
@@ -421,7 +421,7 @@ Type InferTypeTensorElemOffset(const Call& call, const BlockBuilder&) {
TVM_FFI_STATIC_INIT_BLOCK() {
OpDef("relax.inspect.tensor_elem_offset")
.signature(sig::arg("tensor", "The tensor to be inspected"))
- .set_attr<FInferType>("FInferType", InferTypeTensorElemOffset)
+ .set_attr<FInferTypeWithBuilder>("relax.FInferTypeWithBuilder", InferTypeTensorElemOffset)
.set_attr<bool>("RequiresArgumentShapes", false)
.set_attr<FNormalize>("FNormalize", NormalizeToKnownPrimExpr)
.set_attr<bool>("FPurity", true);
diff --git a/src/relax/op/tensor/linear_algebra.cc b/src/relax/op/tensor/linear_algebra.cc
index 5ab4566909..e8bf42b3ff 100644
--- a/src/relax/op/tensor/linear_algebra.cc
+++ b/src/relax/op/tensor/linear_algebra.cc
@@ -166,7 +166,7 @@ TVM_FFI_STATIC_INIT_BLOCK() {
OpDef("relax.matmul")
.signature(sig::arg("x1", "The first input tensor."),
sig::arg("x2", "The second input tensor."), sig::call_attrs<MatmulAttrs>())
- .set_attr<FInferType>("FInferType", InferTypeMatmul)
+ .set_attr<FInferTypeWithBuilder>("relax.FInferTypeWithBuilder", InferTypeMatmul)
.set_attr<TMixedPrecisionPolicy>("TMixedPrecisionPolicy", MixedPrecisionPolicyKind::kAlways)
.set_attr<FInferMixedPrecision>("FInferMixedPrecision", InferMixedPrecisionMatmul)
.set_attr<bool>("FPurity", true);
@@ -187,11 +187,12 @@ TVM_FFI_STATIC_INIT_BLOCK() {
refl::GlobalDef().def("relax.op.einsum", einsum);
}
-Type InferTypeEinsum(const Call& call, const BlockBuilder& ctx) {
+ffi::Expected<Type> InferTypeEinsum(const CallNode* call_node) noexcept try {
+ const Call call = ffi::GetRef<Call>(call_node);
if (call->args.size() != 1) {
TVM_FFI_VISIT_THROW(ValueError, call) << "Einsum op should take 1 argument";
}
- ffi::Array<TensorType> operands_tensor_ty = GetTensorTypeFromTuple(call, ctx, call->args[0]);
+ ffi::Array<TensorType> operands_tensor_ty = GetTensorTypeFromTuple(call, call->args[0]);
if (operands_tensor_ty.empty()) {
TVM_FFI_VISIT_THROW(ValueError, call)
<< "Einsum op expects at least one tensor in the input Tuple. However, the "
@@ -250,12 +251,16 @@ Type InferTypeEinsum(const Call& call, const BlockBuilder& ctx) {
return TensorType(ShapeExpr(oshape), operand_ty, vdev);
}
return TensorType(ShapeExpr(oshape), operand_ty);
+} catch (const ffi::Error& error) {
+ return ffi::Unexpected(error);
+} catch (const std::exception& error) {
+ return ffi::Unexpected(ffi::Error("InternalError", error.what(), ""));
}
TVM_FFI_STATIC_INIT_BLOCK() {
OpDef("relax.einsum")
.signature(sig::arg("operands", "The input tensors."), sig::call_attrs<EinsumAttrs>())
- .set_attr<FInferType>("FInferType", InferTypeEinsum)
+ .set_attr<FInferType>("FInferType", FInferType::FromNative<&InferTypeEinsum>())
.set_attr<bool>("FPurity", true);
}
@@ -271,8 +276,9 @@ TVM_FFI_STATIC_INIT_BLOCK() {
refl::GlobalDef().def("relax.op.outer", outer);
}
-Type InferTypeOuter(const Call& call, const BlockBuilder& ctx) {
- auto input_ty = GetInputTensorType(call, ctx);
+ffi::Expected<Type> InferTypeOuter(const CallNode* call_node) noexcept try {
+ const Call call = ffi::GetRef<Call>(call_node);
+ auto input_ty = GetInputTensorType(call);
auto x1_ty = input_ty[0];
auto x2_ty = input_ty[1];
@@ -289,13 +295,17 @@ Type InferTypeOuter(const Call& call, const BlockBuilder& ctx) {
}
ffi::Array<PrimExpr> output_shape = {x1_shape->values[0], x2_shape->values[0]};
return TensorType(ShapeExpr(output_shape), x1_ty->dtype);
+} catch (const ffi::Error& error) {
+ return ffi::Unexpected(error);
+} catch (const std::exception& error) {
+ return ffi::Unexpected(ffi::Error("InternalError", error.what(), ""));
}
TVM_FFI_STATIC_INIT_BLOCK() {
OpDef("relax.outer")
.signature(sig::arg("x1", "The first input tensor."),
sig::arg("x2", "The second input tensor."))
- .set_attr<FInferType>("FInferType", InferTypeOuter)
+ .set_attr<FInferType>("FInferType", FInferType::FromNative<&InferTypeOuter>())
.set_attr<TMixedPrecisionPolicy>("TMixedPrecisionPolicy", MixedPrecisionPolicyKind::kAlways)
.set_attr<bool>("FPurity", true);
}
diff --git a/src/relax/op/tensor/manipulate.cc b/src/relax/op/tensor/manipulate.cc
index 347582a67c..eac58903d8 100644
--- a/src/relax/op/tensor/manipulate.cc
+++ b/src/relax/op/tensor/manipulate.cc
@@ -134,7 +134,7 @@ Type InferTypeBroadcastTo(const Call& call, const BlockBuilder& ctx) {
TVM_FFI_STATIC_INIT_BLOCK() {
OpDef("relax.broadcast_to")
.signature(sig::arg("x", "The input tensor."), sig::arg("shape", "The target shape."))
- .set_attr<FInferType>("FInferType", InferTypeBroadcastTo)
+ .set_attr<FInferTypeWithBuilder>("relax.FInferTypeWithBuilder", InferTypeBroadcastTo)
.set_attr<TMixedPrecisionPolicy>("TMixedPrecisionPolicy", MixedPrecisionPolicyKind::kFollow)
.set_attr<bool>("FPurity", true);
}
@@ -393,7 +393,7 @@ InferLayoutOutput InferLayoutConcat(
TVM_FFI_STATIC_INIT_BLOCK() {
OpDef("relax.concat")
.signature(sig::arg("tensors", "The input list of tensors."), sig::call_attrs<ConcatAttrs>())
- .set_attr<FInferType>("FInferType", InferTypeConcat)
+ .set_attr<FInferTypeWithBuilder>("relax.FInferTypeWithBuilder", InferTypeConcat)
.set_attr<FRelaxInferLayout>("FRelaxInferLayout", InferLayoutConcat)
.set_attr<TMixedPrecisionPolicy>("TMixedPrecisionPolicy", MixedPrecisionPolicyKind::kFollow)
.set_attr<bool>("FPurity", true);
@@ -414,8 +414,9 @@ TVM_FFI_STATIC_INIT_BLOCK() {
refl::GlobalDef().def("relax.op.expand_dims", expand_dims);
}
-Type InferTypeExpandDims(const Call& call, const BlockBuilder& ctx) {
- TensorType data_ty = GetUnaryInputTensorType(call, ctx);
+ffi::Expected<Type> InferTypeExpandDims(const CallNode* call_node) noexcept try {
+ const Call call = ffi::GetRef<Call>(call_node);
+ TensorType data_ty = GetUnaryInputTensorType(call);
const auto* attrs = call->attrs.as<ExpandDimsAttrs>();
if (attrs->axis.empty()) {
return data_ty;
@@ -427,7 +428,7 @@ Type InferTypeExpandDims(const Call& call, const BlockBuilder& ctx) {
int n_new_dim = attrs->axis.size();
int output_ndim = data_ty->ndim + n_new_dim;
- std::vector<int> axes = NormalizeAxes(call, ctx, output_ndim, attrs->axis);
+ std::vector<int> axes = NormalizeAxes(call, output_ndim, attrs->axis);
const auto* data_shape = data_ty->shape.as<ShapeExprNode>();
if (data_shape == nullptr) {
@@ -451,6 +452,10 @@ Type InferTypeExpandDims(const Call& call, const BlockBuilder& ctx) {
}
TVM_FFI_ICHECK_EQ(i_data_shape, data_ty->ndim);
return TensorType(ShapeExpr(output_shape), data_ty->dtype, data_ty->vdevice);
+} catch (const ffi::Error& error) {
+ return ffi::Unexpected(error);
+} catch (const std::exception& error) {
+ return ffi::Unexpected(ffi::Error("InternalError", error.what(), ""));
}
InferLayoutOutput InferLayoutExpandDims(
@@ -497,7 +502,7 @@ InferLayoutOutput InferLayoutExpandDims(
TVM_FFI_STATIC_INIT_BLOCK() {
OpDef("relax.expand_dims")
.signature(sig::arg("x", "The input tensor."), sig::call_attrs<ExpandDimsAttrs>())
- .set_attr<FInferType>("FInferType", InferTypeExpandDims)
+ .set_attr<FInferType>("FInferType", FInferType::FromNative<&InferTypeExpandDims>())
.set_attr<FRelaxInferLayout>("FRelaxInferLayout", InferLayoutExpandDims)
.set_attr<TMixedPrecisionPolicy>("TMixedPrecisionPolicy", MixedPrecisionPolicyKind::kFollow)
.set_attr<bool>("FPurity", true);
@@ -523,8 +528,9 @@ TVM_FFI_STATIC_INIT_BLOCK() {
refl::GlobalDef().def("relax.op.flatten", flatten);
}
-Type InferTypeFlatten(const Call& call, const BlockBuilder& ctx) {
- TensorType data_ty = GetUnaryInputTensorType(call, ctx);
+ffi::Expected<Type> InferTypeFlatten(const CallNode* call_node) noexcept try {
+ const Call call = ffi::GetRef<Call>(call_node);
+ TensorType data_ty = GetUnaryInputTensorType(call);
if (data_ty->IsUnknownNdim()) {
return TensorType(data_ty->dtype, /*ndim=*/1, data_ty->vdevice);
} else if (data_ty->ndim == 0) {
@@ -539,12 +545,16 @@ Type InferTypeFlatten(const Call& call, const BlockBuilder& ctx) {
}
PrimExpr shape_prod = ComputeShapeProduct(data_shape->values);
return TensorType(ShapeExpr({std::move(shape_prod)}), data_ty->dtype, data_ty->vdevice);
+} catch (const ffi::Error& error) {
+ return ffi::Unexpected(error);
+} catch (const std::exception& error) {
+ return ffi::Unexpected(ffi::Error("InternalError", error.what(), ""));
}
TVM_FFI_STATIC_INIT_BLOCK() {
OpDef("relax.flatten")
.signature(sig::arg("x", "The input tensor."))
- .set_attr<FInferType>("FInferType", InferTypeFlatten)
+ .set_attr<FInferType>("FInferType", FInferType::FromNative<&InferTypeFlatten>())
.set_attr<TMixedPrecisionPolicy>("TMixedPrecisionPolicy", MixedPrecisionPolicyKind::kFollow)
.set_attr<bool>("FPurity", true);
}
@@ -695,7 +705,7 @@ TVM_FFI_STATIC_INIT_BLOCK() {
OpDef("relax.index_tensor")
.signature(sig::arg("data", "The input data."),
sig::arg("indices", "The indices used to index."))
- .set_attr<FInferType>("FInferType", InferTypeIndexTensor)
+ .set_attr<FInferTypeWithBuilder>("relax.FInferTypeWithBuilder", InferTypeIndexTensor)
.set_attr<bool>("FPurity", true);
}
@@ -715,8 +725,9 @@ TVM_FFI_STATIC_INIT_BLOCK() {
refl::GlobalDef().def("relax.op.layout_transform", layout_transform);
}
-Type InferTypeLayoutTransform(const Call& call, const BlockBuilder& ctx) {
- TensorType data_ty = GetUnaryInputTensorType(call, ctx);
+ffi::Expected<Type> InferTypeLayoutTransform(const CallNode* call_node) noexcept try {
+ const Call call = ffi::GetRef<Call>(call_node);
+ TensorType data_ty = GetUnaryInputTensorType(call);
const auto* attrs = call->attrs.as<LayoutTransformAttrs>();
tirx::IndexMap index_map = attrs->index_map;
ffi::Optional<PrimExpr> optional_pad_value = attrs->pad_value;
@@ -757,12 +768,16 @@ Type InferTypeLayoutTransform(const Call& call, const BlockBuilder& ctx) {
sym::Analyzer analyzer;
ffi::Array<PrimExpr> output_shape = index_map->MapShape(shape_ty->values.value(), analyzer);
return TensorType(ShapeExpr(output_shape), data_ty->dtype, data_ty->vdevice);
+} catch (const ffi::Error& error) {
+ return ffi::Unexpected(error);
+} catch (const std::exception& error) {
+ return ffi::Unexpected(ffi::Error("InternalError", error.what(), ""));
}
TVM_FFI_STATIC_INIT_BLOCK() {
OpDef("relax.layout_transform")
.signature(sig::arg("x", "The input tensor."), sig::call_attrs<LayoutTransformAttrs>())
- .set_attr<FInferType>("FInferType", InferTypeLayoutTransform)
+ .set_attr<FInferType>("FInferType", FInferType::FromNative<&InferTypeLayoutTransform>())
.set_attr<TMixedPrecisionPolicy>("TMixedPrecisionPolicy", MixedPrecisionPolicyKind::kFollow)
.set_attr<bool>("FPurity", true);
}
@@ -791,8 +806,9 @@ bool IsIdentityPermutation(const std::vector<int>& permutation) {
return true;
}
-Type InferTypePermuteDims(const Call& call, const BlockBuilder& ctx) {
- TensorType data_ty = GetUnaryInputTensorType(call, ctx);
+ffi::Expected<Type> InferTypePermuteDims(const CallNode* call_node) noexcept try {
+ const Call call = ffi::GetRef<Call>(call_node);
+ TensorType data_ty = GetUnaryInputTensorType(call);
const auto* attrs = call->attrs.as<PermuteDimsAttrs>();
@@ -814,7 +830,7 @@ Type InferTypePermuteDims(const Call& call, const BlockBuilder& ctx) {
std::vector<int> axes;
if (attrs->axes.has_value()) {
- axes = NormalizeAxes(call, ctx, data_ty->ndim, attrs->axes.value());
+ axes = NormalizeAxes(call, data_ty->ndim, attrs->axes.value());
} else {
// Construct the reverse permutation via std::iota
axes.resize(data_ty->ndim);
@@ -834,6 +850,10 @@ Type InferTypePermuteDims(const Call& call, const BlockBuilder& ctx) {
new_shape.push_back(data_shape->values[axes[i]]);
}
return TensorType(ShapeExpr(new_shape), data_ty->dtype, data_ty->vdevice);
+} catch (const ffi::Error& error) {
+ return ffi::Unexpected(error);
+} catch (const std::exception& error) {
+ return ffi::Unexpected(ffi::Error("InternalError", error.what(), ""));
}
InferLayoutOutput InferLayoutPermuteDims(
@@ -882,7 +902,7 @@ InferLayoutOutput InferLayoutPermuteDims(
TVM_FFI_STATIC_INIT_BLOCK() {
OpDef("relax.permute_dims")
.signature(sig::arg("x", "The input tensor."), sig::call_attrs<PermuteDimsAttrs>())
- .set_attr<FInferType>("FInferType", InferTypePermuteDims)
+ .set_attr<FInferType>("FInferType", FInferType::FromNative<&InferTypePermuteDims>())
.set_attr<FRelaxInferLayout>("FRelaxInferLayout", InferLayoutPermuteDims)
.set_attr<TMixedPrecisionPolicy>("TMixedPrecisionPolicy", MixedPrecisionPolicyKind::kFollow)
.set_attr<bool>("FPurity", true);
@@ -1048,7 +1068,7 @@ Type InferTypeReshape(const Call& call, const BlockBuilder& ctx) {
TVM_FFI_STATIC_INIT_BLOCK() {
OpDef("relax.reshape")
.signature(sig::arg("x", "The input tensor."), sig::arg("shape", "The input new shape."))
- .set_attr<FInferType>("FInferType", InferTypeReshape)
+ .set_attr<FInferTypeWithBuilder>("relax.FInferTypeWithBuilder", InferTypeReshape)
.set_attr<TMixedPrecisionPolicy>("TMixedPrecisionPolicy", MixedPrecisionPolicyKind::kFollow)
.set_attr<bool>("FPurity", true);
}
@@ -1221,7 +1241,7 @@ InferLayoutOutput InferLayoutSplit(
TVM_FFI_STATIC_INIT_BLOCK() {
OpDef("relax.split")
.signature(sig::arg("x", "The input tensor."), sig::call_attrs<SplitAttrs>())
- .set_attr<FInferType>("FInferType", InferTypeSplit)
+ .set_attr<FInferTypeWithBuilder>("relax.FInferTypeWithBuilder", InferTypeSplit)
.set_attr<FRelaxInferLayout>("FRelaxInferLayout", InferLayoutSplit)
.set_attr<TMixedPrecisionPolicy>("TMixedPrecisionPolicy", MixedPrecisionPolicyKind::kFollow)
.set_attr<bool>("FPurity", true);
@@ -1242,8 +1262,9 @@ TVM_FFI_STATIC_INIT_BLOCK() {
refl::GlobalDef().def("relax.op.squeeze", squeeze);
}
-Type InferTypeSqueeze(const Call& call, const BlockBuilder& ctx) {
- TensorType data_ty = GetUnaryInputTensorType(call, ctx);
+ffi::Expected<Type> InferTypeSqueeze(const CallNode* call_node) noexcept try {
+ const Call call = ffi::GetRef<Call>(call_node);
+ TensorType data_ty = GetUnaryInputTensorType(call);
const auto* attrs = call->attrs.as<SqueezeAttrs>();
if (attrs->axis.has_value() && attrs->axis.value().empty()) {
return data_ty;
@@ -1262,7 +1283,7 @@ Type InferTypeSqueeze(const Call& call, const BlockBuilder& ctx) {
axis_removal_mask.resize(data_ty->ndim, /*value=*/false);
if (attrs->axis.has_value()) {
- std::vector<int> axes = NormalizeAxes(call, ctx, data_ty->ndim, attrs->axis.value());
+ std::vector<int> axes = NormalizeAxes(call, data_ty->ndim, attrs->axis.value());
if (!shape_value.has_value()) {
return TensorType(data_ty->dtype, data_ty->ndim - axes.size(), data_ty->vdevice);
@@ -1315,6 +1336,10 @@ Type InferTypeSqueeze(const Call& call, const BlockBuilder& ctx) {
} else {
return TensorType(ShapeExpr(output_shape), data_ty->dtype, data_ty->vdevice);
}
+} catch (const ffi::Error& error) {
+ return ffi::Unexpected(error);
+} catch (const std::exception& error) {
+ return ffi::Unexpected(ffi::Error("InternalError", error.what(), ""));
}
InferLayoutOutput InferLayoutSqueeze(
@@ -1380,7 +1405,7 @@ InferLayoutOutput InferLayoutSqueeze(
TVM_FFI_STATIC_INIT_BLOCK() {
OpDef("relax.squeeze")
.signature(sig::arg("x", "The input tensor."), sig::call_attrs<SqueezeAttrs>())
- .set_attr<FInferType>("FInferType", InferTypeSqueeze)
+ .set_attr<FInferType>("FInferType", FInferType::FromNative<&InferTypeSqueeze>())
.set_attr<FRelaxInferLayout>("FRelaxInferLayout", InferLayoutSqueeze)
.set_attr<TMixedPrecisionPolicy>("TMixedPrecisionPolicy", MixedPrecisionPolicyKind::kFollow)
.set_attr<bool>("FPurity", true);
@@ -1633,7 +1658,7 @@ TVM_FFI_STATIC_INIT_BLOCK() {
OpDef("relax.stack")
.signature(sig::arg("tensors", "The input list of tensors to stack"),
sig::call_attrs<StackAttrs>())
- .set_attr<FInferType>("FInferType", InferTypeStack)
+ .set_attr<FInferTypeWithBuilder>("relax.FInferTypeWithBuilder", InferTypeStack)
.set_attr<FRelaxInferLayout>("FRelaxInferLayout", InferLayoutStack)
.set_attr<TMixedPrecisionPolicy>("TMixedPrecisionPolicy", MixedPrecisionPolicyKind::kFollow)
.set_attr<bool>("FPurity", true);
@@ -1683,7 +1708,7 @@ TVM_FFI_STATIC_INIT_BLOCK() {
OpDef("relax.collapse_sum_like")
.signature(sig::arg("data", "The input tensor."),
sig::arg("collapse_target", "The tensor whose shape is the shape to collapse to."))
- .set_attr<FInferType>("FInferType", InferTypeCollapseSumLike)
+ .set_attr<FInferTypeWithBuilder>("relax.FInferTypeWithBuilder", InferTypeCollapseSumLike)
.set_attr<bool>("FPurity", true);
}
@@ -1734,7 +1759,7 @@ TVM_FFI_STATIC_INIT_BLOCK() {
OpDef("relax.collapse_sum_to")
.signature(sig::arg("data", "The input tensor."),
sig::arg("shape", "The shape to collapse to."))
- .set_attr<FInferType>("FInferType", InferTypeCollapseSumTo)
+ .set_attr<FInferTypeWithBuilder>("relax.FInferTypeWithBuilder", InferTypeCollapseSumTo)
.set_attr<bool>("FPurity", true);
}
@@ -1858,7 +1883,7 @@ InferLayoutOutput InferLayoutRepeat(
TVM_FFI_STATIC_INIT_BLOCK() {
OpDef("relax.repeat")
.signature(sig::arg("data", "The input tensor."), sig::call_attrs<RepeatAttrs>())
- .set_attr<FInferType>("FInferType", InferTypeRepeat)
+ .set_attr<FInferTypeWithBuilder>("relax.FInferTypeWithBuilder", InferTypeRepeat)
.set_attr<FRelaxInferLayout>("FRelaxInferLayout", InferLayoutRepeat)
.set_attr<bool>("FPurity", true);
}
@@ -2002,7 +2027,7 @@ InferLayoutOutput InferLayoutTile(
TVM_FFI_STATIC_INIT_BLOCK() {
OpDef("relax.tile")
.signature(sig::arg("data", "The input tensor."), sig::call_attrs<TileAttrs>())
- .set_attr<FInferType>("FInferType", InferTypeTile)
+ .set_attr<FInferTypeWithBuilder>("relax.FInferTypeWithBuilder", InferTypeTile)
.set_attr<FRelaxInferLayout>("FRelaxInferLayout", InferLayoutTile)
.set_attr<bool>("FPurity", true);
}
@@ -2021,11 +2046,12 @@ TVM_FFI_STATIC_INIT_BLOCK() {
refl::GlobalDef().def("relax.op.flip", flip);
}
-Type InferTypeFlip(const Call& call, const BlockBuilder& ctx) {
+ffi::Expected<Type> InferTypeFlip(const CallNode* call_node) noexcept try {
+ const Call call = ffi::GetRef<Call>(call_node);
if (call->args.size() != 1) {
TVM_FFI_VISIT_THROW(ValueError, call) << "Flip op should take 1 argument";
}
- TensorType data_ty = GetUnaryInputTensorType(call, ctx);
+ TensorType data_ty = GetUnaryInputTensorType(call);
const auto* attrs = call->attrs.as<FlipAttrs>();
int axis = static_cast<int>(attrs->axis);
if (!data_ty->IsUnknownNdim()) {
@@ -2037,6 +2063,10 @@ Type InferTypeFlip(const Call& call, const BlockBuilder& ctx) {
}
}
return data_ty;
+} catch (const ffi::Error& error) {
+ return ffi::Unexpected(error);
+} catch (const std::exception& error) {
+ return ffi::Unexpected(ffi::Error("InternalError", error.what(), ""));
}
InferLayoutOutput InferLayoutFlip(
@@ -2074,7 +2104,7 @@ InferLayoutOutput InferLayoutFlip(
TVM_FFI_STATIC_INIT_BLOCK() {
OpDef("relax.flip")
.signature(sig::arg("data", "The input tensor."), sig::call_attrs<FlipAttrs>())
- .set_attr<FInferType>("FInferType", InferTypeFlip)
+ .set_attr<FInferType>("FInferType", FInferType::FromNative<&InferTypeFlip>())
.set_attr<FRelaxInferLayout>("FRelaxInferLayout", InferLayoutFlip)
.set_attr<bool>("FPurity", true);
}
@@ -2170,7 +2200,7 @@ TVM_FFI_STATIC_INIT_BLOCK() {
.signature(sig::arg("data", "The input tensor."),
sig::arg("seq_lengths", "The sequence length tensor."),
sig::call_attrs<ReverseSequenceAttrs>())
- .set_attr<FInferType>("FInferType", InferTypeReverseSequence)
+ .set_attr<FInferTypeWithBuilder>("relax.FInferTypeWithBuilder", InferTypeReverseSequence)
.set_attr<bool>("FPurity", true);
}
@@ -2188,7 +2218,8 @@ TVM_FFI_STATIC_INIT_BLOCK() {
refl::GlobalDef().def("relax.op.gather_elements", gather_elements);
}
-Type InferTypeGatherElements(const Call& call, const BlockBuilder& ctx) {
+ffi::Expected<Type> InferTypeGatherElements(const CallNode* call_node) noexcept try {
+ const Call call = ffi::GetRef<Call>(call_node);
const auto* data_ty = GetTypeAs<TensorTypeNode>(call->args[0]);
const auto* indices_ty = GetTypeAs<TensorTypeNode>(call->args[1]);
const auto* attrs = call->attrs.as<GatherElementsAttrs>();
@@ -2232,6 +2263,10 @@ Type InferTypeGatherElements(const Call& call, const BlockBuilder& ctx) {
return TensorType(indices_ty->shape.value(), data_ty->dtype, data_ty->vdevice);
}
return TensorType(data_ty->dtype, indices_ty->ndim, data_ty->vdevice);
+} catch (const ffi::Error& error) {
+ return ffi::Unexpected(error);
+} catch (const std::exception& error) {
+ return ffi::Unexpected(ffi::Error("InternalError", error.what(), ""));
}
InferLayoutOutput InferLayoutGatherElements(
@@ -2271,7 +2306,7 @@ TVM_FFI_STATIC_INIT_BLOCK() {
OpDef("relax.gather_elements")
.signature(sig::arg("data", "The input tensor."), sig::arg("indices", "The indices tensor."),
sig::call_attrs<GatherElementsAttrs>())
- .set_attr<FInferType>("FInferType", InferTypeGatherElements)
+ .set_attr<FInferType>("FInferType", FInferType::FromNative<&InferTypeGatherElements>())
.set_attr<FRelaxInferLayout>("FRelaxInferLayout", InferLayoutGatherElements)
.set_attr<bool>("FPurity", true);
}
@@ -2290,7 +2325,8 @@ TVM_FFI_STATIC_INIT_BLOCK() {
refl::GlobalDef().def("relax.op.gather_nd", gather_nd);
}
-Type InferTypeGatherND(const Call& call, const BlockBuilder& ctx) {
+ffi::Expected<Type> InferTypeGatherND(const CallNode* call_node) noexcept try {
+ const Call call = ffi::GetRef<Call>(call_node);
const auto* data_ty = GetTypeAs<TensorTypeNode>(call->args[0]);
const auto* indices_ty = GetTypeAs<TensorTypeNode>(call->args[1]);
const auto* attrs = call->attrs.as<GatherNDAttrs>();
@@ -2359,13 +2395,17 @@ Type InferTypeGatherND(const Call& call, const BlockBuilder& ctx) {
}
TVM_FFI_ICHECK_EQ(out_shape.size(), output_ndim);
return TensorType(ShapeExpr(out_shape), data_ty->dtype, data_ty->vdevice);
+} catch (const ffi::Error& error) {
+ return ffi::Unexpected(error);
+} catch (const std::exception& error) {
+ return ffi::Unexpected(ffi::Error("InternalError", error.what(), ""));
}
TVM_FFI_STATIC_INIT_BLOCK() {
OpDef("relax.gather_nd")
.signature(sig::arg("data", "The input tensor."), sig::arg("indices", "The indices tensor."),
sig::call_attrs<GatherNDAttrs>())
- .set_attr<FInferType>("FInferType", InferTypeGatherND)
+ .set_attr<FInferType>("FInferType", FInferType::FromNative<&InferTypeGatherND>())
.set_attr<bool>("FPurity", true);
}
@@ -2383,7 +2423,8 @@ TVM_FFI_STATIC_INIT_BLOCK() {
refl::GlobalDef().def("relax.op.index_put", index_put);
}
-Type InferTypeIndexPut(const Call& call, const BlockBuilder& ctx) {
+ffi::Expected<Type> InferTypeIndexPut(const CallNode* call_node) noexcept try {
+ const Call call = ffi::GetRef<Call>(call_node);
const auto* data_ty = GetTypeAs<TensorTypeNode>(call->args[0]);
const auto* values_ty = GetTypeAs<TensorTypeNode>(call->args[2]);
@@ -2509,6 +2550,10 @@ Type InferTypeIndexPut(const Call& call, const BlockBuilder& ctx) {
return TensorType(ShapeExpr(data_shape->values), data_ty->dtype, data_ty->vdevice);
}
return TensorType(data_ty->dtype, data_ty->ndim, data_ty->vdevice);
+} catch (const ffi::Error& error) {
+ return ffi::Unexpected(error);
+} catch (const std::exception& error) {
+ return ffi::Unexpected(ffi::Error("InternalError", error.what(), ""));
}
TVM_FFI_STATIC_INIT_BLOCK() {
@@ -2516,7 +2561,7 @@ TVM_FFI_STATIC_INIT_BLOCK() {
.signature(sig::arg("data", "The input tensor."),
sig::arg("indices", "The indices tensor(s)."),
sig::arg("values", "The values to put."), sig::call_attrs<IndexPutAttrs>())
- .set_attr<FInferType>("FInferType", InferTypeIndexPut)
+ .set_attr<FInferType>("FInferType", FInferType::FromNative<&InferTypeIndexPut>())
.set_attr<bool>("FPurity", true);
}
@@ -2534,11 +2579,12 @@ TVM_FFI_STATIC_INIT_BLOCK() {
refl::GlobalDef().def("relax.op.meshgrid", meshgrid);
}
-Type InferTypeMeshgrid(const Call& call, const BlockBuilder& ctx) {
+ffi::Expected<Type> InferTypeMeshgrid(const CallNode* call_node) noexcept try {
+ const Call call = ffi::GetRef<Call>(call_node);
if (call->args.size() != 1) {
TVM_FFI_VISIT_THROW(ValueError, call) << "meshgrid op expects 1 Tuple input argument.";
}
- ffi::Array<TensorType> input_ty = GetTensorTypeFromTuple(call, ctx, call->args[0]);
+ ffi::Array<TensorType> input_ty = GetTensorTypeFromTuple(call, call->args[0]);
int n_inputs = input_ty.size();
@@ -2615,13 +2661,17 @@ Type InferTypeMeshgrid(const Call& call, const BlockBuilder& ctx) {
}
return TupleType(out_fields);
+} catch (const ffi::Error& error) {
+ return ffi::Unexpected(error);
+} catch (const std::exception& error) {
+ return ffi::Unexpected(ffi::Error("InternalError", error.what(), ""));
}
TVM_FFI_STATIC_INIT_BLOCK() {
OpDef("relax.meshgrid")
.signature(sig::arg("tensors", "The input list of tensors."),
sig::call_attrs<MeshgridAttrs>())
- .set_attr<FInferType>("FInferType", InferTypeMeshgrid)
+ .set_attr<FInferType>("FInferType", FInferType::FromNative<&InferTypeMeshgrid>())
.set_attr<TMixedPrecisionPolicy>("TMixedPrecisionPolicy", MixedPrecisionPolicyKind::kFollow)
.set_attr<bool>("FPurity", true);
}
@@ -2766,7 +2816,7 @@ TVM_FFI_STATIC_INIT_BLOCK() {
.signature(sig::arg("data", "The input tensor."), sig::arg("indices", "The indices tensor."),
sig::arg("updates", "The input tensor of updates."),
sig::call_attrs<ScatterElementsAttrs>())
- .set_attr<FInferType>("FInferType", InferTypeScatterElements)
+ .set_attr<FInferTypeWithBuilder>("relax.FInferTypeWithBuilder", InferTypeScatterElements)
.set_attr<FRelaxInferLayout>("FRelaxInferLayout", InferLayoutScatterElements)
.set_attr<bool>("FPurity", true);
}
@@ -2944,7 +2994,7 @@ TVM_FFI_STATIC_INIT_BLOCK() {
.signature(sig::arg("data", "The input tensor."), sig::arg("indices", "The indices tensor."),
sig::arg("updates", "The input tensor of updates."),
sig::call_attrs<ScatterNDAttrs>())
- .set_attr<FInferType>("FInferType", InferTypeScatterND)
+ .set_attr<FInferTypeWithBuilder>("relax.FInferTypeWithBuilder", InferTypeScatterND)
.set_attr<FRelaxInferLayout>("FRelaxInferLayout", InferLayoutScatterND)
.set_attr<bool>("FPurity", true);
}
@@ -3100,7 +3150,7 @@ TVM_FFI_STATIC_INIT_BLOCK() {
sig::arg("start", "The starting index of the slice (inclusive)."),
sig::arg("end", "The ending index of the slice (exclusive)."),
sig::arg("step", "The step of the slice."), sig::call_attrs<SliceScatterAttrs>())
- .set_attr<FInferType>("FInferType", InferTypeSliceScatter)
+ .set_attr<FInferTypeWithBuilder>("relax.FInferTypeWithBuilder", InferTypeSliceScatter)
.set_attr<bool>("FPurity", true);
}
@@ -3129,8 +3179,9 @@ TVM_FFI_STATIC_INIT_BLOCK() {
refl::GlobalDef().def("relax.op.one_hot", one_hot);
}
-Type InferTypeOneHot(const Call& call, const BlockBuilder& ctx) {
- TensorType indices_ty = GetInputTensorType(call, 0, ctx);
+ffi::Expected<Type> InferTypeOneHot(const CallNode* call_node) noexcept try {
+ const Call call = ffi::GetRef<Call>(call_node);
+ TensorType indices_ty = GetInputTensorType(call, 0);
const auto* attrs = call->attrs.as<OneHotAttrs>();
PrimExpr on_value = call->args[1].as_or_throw<PrimExpr>();
PrimExpr off_value = call->args[2].as_or_throw<PrimExpr>();
@@ -3176,6 +3227,10 @@ Type InferTypeOneHot(const Call& call, const BlockBuilder& ctx) {
output_shape.insert(output_shape.begin() + axis, attrs->depth);
return TensorType(ShapeExpr(output_shape), dtype, indices_ty->vdevice);
+} catch (const ffi::Error& error) {
+ return ffi::Unexpected(error);
+} catch (const std::exception& error) {
+ return ffi::Unexpected(ffi::Error("InternalError", error.what(), ""));
}
TVM_FFI_STATIC_INIT_BLOCK() {
@@ -3184,7 +3239,7 @@ TVM_FFI_STATIC_INIT_BLOCK() {
sig::arg("on_value", "The value to fill at specified indices."),
sig::arg("off_value", "The value to fill at other indices."),
sig::call_attrs<OneHotAttrs>())
- .set_attr<FInferType>("FInferType", InferTypeOneHot)
+ .set_attr<FInferType>("FInferType", FInferType::FromNative<&InferTypeOneHot>())
.set_attr<bool>("FPurity", true);
}
diff --git a/src/relax/op/tensor/qdq.cc b/src/relax/op/tensor/qdq.cc
index 78a6164386..ef4890b76e 100644
--- a/src/relax/op/tensor/qdq.cc
+++ b/src/relax/op/tensor/qdq.cc
@@ -161,7 +161,7 @@ TVM_FFI_STATIC_INIT_BLOCK() {
sig::arg("scale", "The quantization scale of the output tensor."),
sig::arg("zero_point", "The quantization zero_point of the output tensor."),
sig::call_attrs<QuantizeAttrs>())
- .set_attr<FInferType>("FInferType", InferTypeQuantize)
+ .set_attr<FInferTypeWithBuilder>("relax.FInferTypeWithBuilder", InferTypeQuantize)
.set_attr<bool>("FPurity", true);
}
@@ -284,7 +284,7 @@ TVM_FFI_STATIC_INIT_BLOCK() {
sig::arg("scale", "The quantization scale of the input tensor."),
sig::arg("zero_point", "The quantization zero_point of the input tensor."),
sig::call_attrs<QuantizeAttrs>())
- .set_attr<FInferType>("FInferType", InferTypeDequantize)
+ .set_attr<FInferTypeWithBuilder>("relax.FInferTypeWithBuilder", InferTypeDequantize)
.set_attr<bool>("FPurity", true);
}
diff --git a/src/relax/op/tensor/sampling.cc b/src/relax/op/tensor/sampling.cc
index 6c3fb3be06..84a5ca34be 100644
--- a/src/relax/op/tensor/sampling.cc
+++ b/src/relax/op/tensor/sampling.cc
@@ -54,11 +54,12 @@ TVM_FFI_STATIC_INIT_BLOCK() {
refl::GlobalDef().def("relax.op.multinomial_from_uniform", multinomial_from_uniform);
}
-Type InferTypeMultinomialFromUniform(const Call& call, const BlockBuilder& ctx) {
- CheckNumArguments(call, ctx);
- TensorType prob_ty = GetInputTensorType(call, 0, ctx);
- TensorType uniform_sample_ty = GetInputTensorType(call, 1, ctx);
- TensorType sample_indices_ty = GetInputTensorType(call, 2, ctx);
+ffi::Expected<Type> InferTypeMultinomialFromUniform(const CallNode* call_node) noexcept try {
+ const Call call = ffi::GetRef<Call>(call_node);
+ CheckNumArguments(call);
+ TensorType prob_ty = GetInputTensorType(call, 0);
+ TensorType uniform_sample_ty = GetInputTensorType(call, 1);
+ TensorType sample_indices_ty = GetInputTensorType(call, 2);
const auto* attrs = call->attrs.as<MultinomialFromUniformAttrs>();
// Only the element kind matters here; shape inference does not depend on vector lanes.
@@ -142,6 +143,10 @@ Type InferTypeMultinomialFromUniform(const Call& call, const BlockBuilder& ctx)
<< sample_indices_ty->shape;
}
return TensorType(ShapeExpr({n, 1}), PrimType(attrs->dtype), prob_ty->vdevice);
+} catch (const ffi::Error& error) {
+ return ffi::Unexpected(error);
+} catch (const std::exception& error) {
+ return ffi::Unexpected(ffi::Error("InternalError", error.what(), ""));
}
TVM_FFI_STATIC_INIT_BLOCK() {
@@ -150,7 +155,8 @@ TVM_FFI_STATIC_INIT_BLOCK() {
sig::arg("uniform_sample", "The uniform sample tensor."),
sig::arg("sample_indices", "The sample indices tensor."),
sig::call_attrs<MultinomialFromUniformAttrs>())
- .set_attr<FInferType>("FInferType", InferTypeMultinomialFromUniform)
+ .set_attr<FInferType>("FInferType",
+ FInferType::FromNative<&InferTypeMultinomialFromUniform>())
.set_attr<bool>("FPurity", true);
}
diff --git a/src/relax/op/tensor/search.cc b/src/relax/op/tensor/search.cc
index 40b170c9d4..fa46ec670b 100644
--- a/src/relax/op/tensor/search.cc
+++ b/src/relax/op/tensor/search.cc
@@ -54,8 +54,9 @@ TVM_FFI_STATIC_INIT_BLOCK() {
refl::GlobalDef().def("relax.op.bucketize", bucketize);
}
-Type InferTypeBucketize(const Call& call, const BlockBuilder& ctx) {
- ffi::Array<TensorType> input_ty = GetInputTensorType(call, ctx);
+ffi::Expected<Type> InferTypeBucketize(const CallNode* call_node) noexcept try {
+ const Call call = ffi::GetRef<Call>(call_node);
+ ffi::Array<TensorType> input_ty = GetInputTensorType(call);
TensorType input_tensor_info = input_ty[0];
TensorType boundaries_info = input_ty[1];
@@ -75,6 +76,10 @@ Type InferTypeBucketize(const Call& call, const BlockBuilder& ctx) {
return TensorType(ShapeExpr(data_shape->values), out_dtype, input_tensor_info->vdevice);
}
return TensorType(out_dtype, input_tensor_info->ndim, input_tensor_info->vdevice);
+} catch (const ffi::Error& error) {
+ return ffi::Unexpected(error);
+} catch (const std::exception& error) {
+ return ffi::Unexpected(ffi::Error("InternalError", error.what(), ""));
}
TVM_FFI_STATIC_INIT_BLOCK() {
@@ -86,7 +91,7 @@ TVM_FFI_STATIC_INIT_BLOCK() {
"1-D tensor, must contain a strictly increasing sequence, or the return value is "
"undefined."),
sig::call_attrs<BucketizeAttrs>())
- .set_attr<FInferType>("FInferType", InferTypeBucketize)
+ .set_attr<FInferType>("FInferType", FInferType::FromNative<&InferTypeBucketize>())
.set_attr<bool>("FPurity", true);
}
@@ -188,19 +193,20 @@ TVM_FFI_STATIC_INIT_BLOCK() {
.signature(sig::arg("condition", "When True, yield `x1`; otherwise, yield `x2`."),
sig::arg("x1", "The first input tensor."),
sig::arg("x2", "The second input tensor."))
- .set_attr<FInferType>("FInferType", InferTypeWhere)
+ .set_attr<FInferTypeWithBuilder>("relax.FInferTypeWithBuilder", InferTypeWhere)
.set_attr<bool>("FPurity", true);
}
/* relax.argmax & relax.argmin */
-Type InferTypeArgmaxArgmin(const Call& call, const BlockBuilder& ctx) {
- TensorType data_ty = GetUnaryInputTensorType(call, ctx);
+ffi::Expected<Type> InferTypeArgmaxArgmin(const CallNode* call_node) noexcept try {
+ const Call call = ffi::GetRef<Call>(call_node);
+ TensorType data_ty = GetUnaryInputTensorType(call);
const auto* attrs = call->attrs.as<ArgmaxArgminAttrs>();
int axis = -1;
if (!data_ty->IsUnknownNdim() && attrs->axis.has_value()) {
- axis = NormalizeAxis(call, ctx, data_ty->ndim, attrs->axis.value());
+ axis = NormalizeAxis(call, data_ty->ndim, attrs->axis.value());
}
int out_ndim;
@@ -250,6 +256,10 @@ Type InferTypeArgmaxArgmin(const Call& call, const BlockBuilder& ctx) {
}
TVM_FFI_ICHECK_EQ(static_cast<int>(out_shape.size()), out_ndim);
return TensorType(ShapeExpr(out_shape), out_dtype, data_ty->vdevice);
+} catch (const ffi::Error& error) {
+ return ffi::Unexpected(error);
+} catch (const std::exception& error) {
+ return ffi::Unexpected(ffi::Error("InternalError", error.what(), ""));
}
Expr argmax(Expr x, ffi::Optional<int64_t> axis, bool keepdims) {
@@ -273,7 +283,7 @@ TVM_FFI_STATIC_INIT_BLOCK() {
OpDef("relax.argmax")
.signature(sig::arg("x", "The input data tensor"), sig::call_attrs<ArgmaxArgminAttrs>())
- .set_attr<FInferType>("FInferType", InferTypeArgmaxArgmin)
+ .set_attr<FInferType>("FInferType", FInferType::FromNative<&InferTypeArgmaxArgmin>())
.set_attr<bool>("FPurity", true);
};
@@ -282,7 +292,7 @@ TVM_FFI_STATIC_INIT_BLOCK() {
OpDef("relax.argmin")
.signature(sig::arg("x", "The input data tensor"), sig::call_attrs<ArgmaxArgminAttrs>())
- .set_attr<FInferType>("FInferType", InferTypeArgmaxArgmin)
+ .set_attr<FInferType>("FInferType", FInferType::FromNative<&InferTypeArgmaxArgmin>())
.set_attr<bool>("FPurity", true);
};
diff --git a/src/relax/op/tensor/set.cc b/src/relax/op/tensor/set.cc
index b711d31558..e77ecfba1b 100644
--- a/src/relax/op/tensor/set.cc
+++ b/src/relax/op/tensor/set.cc
@@ -57,7 +57,8 @@ TVM_FFI_STATIC_INIT_BLOCK() {
refl::GlobalDef().def("relax.op.unique", unique);
}
-Type InferTypeUnique(const Call& call, const BlockBuilder& ctx) {
+ffi::Expected<Type> InferTypeUnique(const CallNode* call_node) noexcept try {
+ const Call call = ffi::GetRef<Call>(call_node);
TensorType data_ty = call->args[0]->ty.as_or_throw<TensorType>();
PrimExpr axis, return_index, return_inverse, return_counts;
if (call->args.size() == 6) {
@@ -68,7 +69,7 @@ Type InferTypeUnique(const Call& call, const BlockBuilder& ctx) {
if (!data_ty->IsUnknownNdim() && axis.defined()) {
// Normalize the axis for sanity check purpose.
if (const auto* axis_int = axis.as<IntImmNode>()) {
- NormalizeAxis(call, ctx, data_ty->ndim, axis_int->value.as<int>().value());
+ NormalizeAxis(call, data_ty->ndim, axis_int->value.as<int>().value());
}
}
TVM_FFI_ICHECK(call->args[2].as<PrimExpr>());
@@ -139,6 +140,10 @@ Type InferTypeUnique(const Call& call, const BlockBuilder& ctx) {
} else {
return TupleType(output_ty);
}
+} catch (const ffi::Error& error) {
+ return ffi::Unexpected(error);
+} catch (const std::exception& error) {
+ return ffi::Unexpected(ffi::Error("InternalError", error.what(), ""));
}
TVM_FFI_STATIC_INIT_BLOCK() {
@@ -160,7 +165,7 @@ TVM_FFI_STATIC_INIT_BLOCK() {
sig::arg("return_counts",
"Whether to return an additional tensor with counts of each unique elements"),
sig::var_args("args"))
- .set_attr<FInferType>("FInferType", InferTypeUnique)
+ .set_attr<FInferType>("FInferType", FInferType::FromNative<&InferTypeUnique>())
.set_attr<FCallPacked>("FCallPacked", "relax.run.unique")
.set_attr<bool>("FPurity", true);
}
@@ -176,15 +181,20 @@ TVM_FFI_STATIC_INIT_BLOCK() {
refl::GlobalDef().def("relax.op.nonzero", nonzero);
}
-Type InferTypeNonzero(const Call& call, const BlockBuilder& ctx) {
- TensorType data_ty = GetInputTensorType(call, 0, ctx);
+ffi::Expected<Type> InferTypeNonzero(const CallNode* call_node) noexcept try {
+ const Call call = ffi::GetRef<Call>(call_node);
+ TensorType data_ty = GetInputTensorType(call, 0);
return TensorType(PrimType::Int(64), 2, data_ty->vdevice);
+} catch (const ffi::Error& error) {
+ return ffi::Unexpected(error);
+} catch (const std::exception& error) {
+ return ffi::Unexpected(ffi::Error("InternalError", error.what(), ""));
}
TVM_FFI_STATIC_INIT_BLOCK() {
OpDef("relax.nonzero")
.signature(sig::arg("x", "The input tensor"))
- .set_attr<FInferType>("FInferType", InferTypeNonzero)
+ .set_attr<FInferType>("FInferType", FInferType::FromNative<&InferTypeNonzero>())
.set_attr<FCallPacked>("FCallPacked", "relax.run.nonzero")
.set_attr<bool>("FPurity", true);
}
diff --git a/src/relax/op/tensor/sorting.cc b/src/relax/op/tensor/sorting.cc
index 840315d2fb..cef5d9709e 100644
--- a/src/relax/op/tensor/sorting.cc
+++ b/src/relax/op/tensor/sorting.cc
@@ -53,14 +53,19 @@ TVM_FFI_STATIC_INIT_BLOCK() {
refl::GlobalDef().def("relax.op.sort", sort);
}
-Type InferTypeSort(const Call& call, const BlockBuilder& ctx) {
- return GetUnaryInputTensorType(call, ctx);
+ffi::Expected<Type> InferTypeSort(const CallNode* call_node) noexcept try {
+ const Call call = ffi::GetRef<Call>(call_node);
+ return GetUnaryInputTensorType(call);
+} catch (const ffi::Error& error) {
+ return ffi::Unexpected(error);
+} catch (const std::exception& error) {
+ return ffi::Unexpected(ffi::Error("InternalError", error.what(), ""));
}
TVM_FFI_STATIC_INIT_BLOCK() {
OpDef("relax.sort")
.signature(sig::arg("data", "The input tensor."), sig::call_attrs<SortAttrs>())
- .set_attr<FInferType>("FInferType", InferTypeSort)
+ .set_attr<FInferType>("FInferType", FInferType::FromNative<&InferTypeSort>())
.set_attr<bool>("FPurity", true);
}
@@ -81,8 +86,9 @@ TVM_FFI_STATIC_INIT_BLOCK() {
refl::GlobalDef().def("relax.op.argsort", argsort);
}
-Type InferTypeArgsort(const Call& call, const BlockBuilder& ctx) {
- TensorType data_ty = GetUnaryInputTensorType(call, ctx);
+ffi::Expected<Type> InferTypeArgsort(const CallNode* call_node) noexcept try {
+ const Call call = ffi::GetRef<Call>(call_node);
+ TensorType data_ty = GetUnaryInputTensorType(call);
const auto* attrs = call->attrs.as<ArgsortAttrs>();
ffi::Optional<PrimType> out_type = attrs->dtype.has_value()
? ffi::Optional<PrimType>(PrimType(attrs->dtype.value()))
@@ -91,12 +97,16 @@ Type InferTypeArgsort(const Call& call, const BlockBuilder& ctx) {
return TensorType(data_ty->shape.value(), out_type, data_ty->vdevice);
}
return TensorType(out_type, data_ty->ndim, data_ty->vdevice);
+} catch (const ffi::Error& error) {
+ return ffi::Unexpected(error);
+} catch (const std::exception& error) {
+ return ffi::Unexpected(ffi::Error("InternalError", error.what(), ""));
}
TVM_FFI_STATIC_INIT_BLOCK() {
OpDef("relax.argsort")
.signature(sig::arg("data", "The input tensor."), sig::call_attrs<ArgsortAttrs>())
- .set_attr<FInferType>("FInferType", InferTypeArgsort)
+ .set_attr<FInferType>("FInferType", FInferType::FromNative<&InferTypeArgsort>())
.set_attr<bool>("FPurity", true);
}
@@ -120,8 +130,9 @@ TVM_FFI_STATIC_INIT_BLOCK() {
refl::GlobalDef().def("relax.op.topk", topk);
}
-Type InferTypeTopK(const Call& call, const BlockBuilder& ctx) {
- TensorType data_ty = GetUnaryInputTensorType(call, ctx);
+ffi::Expected<Type> InferTypeTopK(const CallNode* call_node) noexcept try {
+ const Call call = ffi::GetRef<Call>(call_node);
+ TensorType data_ty = GetUnaryInputTensorType(call);
const auto* data_shape = data_ty->shape.as<ShapeExprNode>();
const auto* attrs = call->attrs.as<TopKAttrs>();
ffi::Optional<PrimType> indices_type =
@@ -159,12 +170,16 @@ Type InferTypeTopK(const Call& call, const BlockBuilder& ctx) {
}
TVM_FFI_THROW(InternalError) << "Unsupported ret type: " << ret_type;
TVM_FFI_UNREACHABLE();
+} catch (const ffi::Error& error) {
+ return ffi::Unexpected(error);
+} catch (const std::exception& error) {
+ return ffi::Unexpected(ffi::Error("InternalError", error.what(), ""));
}
TVM_FFI_STATIC_INIT_BLOCK() {
OpDef("relax.topk")
.signature(sig::arg("data", "The input tensor."), sig::call_attrs<TopKAttrs>())
- .set_attr<FInferType>("FInferType", InferTypeTopK)
+ .set_attr<FInferType>("FInferType", FInferType::FromNative<&InferTypeTopK>())
.set_attr<bool>("FPurity", true);
}
diff --git a/src/relax/op/tensor/statistical.cc b/src/relax/op/tensor/statistical.cc
index ec70316fef..d1803da6f4 100644
--- a/src/relax/op/tensor/statistical.cc
+++ b/src/relax/op/tensor/statistical.cc
@@ -37,13 +37,14 @@ TVM_FFI_STATIC_INIT_BLOCK() {
ScanopAttrs::RegisterReflection();
}
-Type InferTypeStatistical(const Call& call, const BlockBuilder& ctx) {
- TensorType data_ty = GetUnaryInputTensorType(call, ctx);
+ffi::Expected<Type> InferTypeStatistical(const CallNode* call_node) noexcept try {
+ const Call call = ffi::GetRef<Call>(call_node);
+ TensorType data_ty = GetUnaryInputTensorType(call);
const auto* attrs = call->attrs.as<StatisticalAttrs>();
std::vector<int> axes;
if (!data_ty->IsUnknownNdim() && attrs->axis.has_value()) {
- axes = NormalizeAxes(call, ctx, data_ty->ndim, attrs->axis.value());
+ axes = NormalizeAxes(call, data_ty->ndim, attrs->axis.value());
}
int out_ndim;
@@ -88,6 +89,10 @@ Type InferTypeStatistical(const Call& call, const BlockBuilder& ctx) {
}
TVM_FFI_ICHECK_EQ(static_cast<int>(out_shape.size()), out_ndim);
return TensorType(ShapeExpr(out_shape), data_ty->dtype, data_ty->vdevice);
+} catch (const ffi::Error& error) {
+ return ffi::Unexpected(error);
+} catch (const std::exception& error) {
+ return ffi::Unexpected(ffi::Error("InternalError", error.what(), ""));
}
InferLayoutOutput InferLayoutStatistical(
@@ -151,8 +156,9 @@ InferLayoutOutput InferLayoutStatistical(
Attrs(new_attrs));
}
-Type InferTypeScan(const Call& call, const BlockBuilder& ctx) {
- TensorType data_ty = GetUnaryInputTensorType(call, ctx);
+ffi::Expected<Type> InferTypeScan(const CallNode* call_node) noexcept try {
+ const Call call = ffi::GetRef<Call>(call_node);
+ TensorType data_ty = GetUnaryInputTensorType(call);
const auto* attrs = call->attrs.as<ScanopAttrs>();
ffi::Optional<PrimType> out_type = attrs->dtype.has_value()
@@ -178,15 +184,20 @@ Type InferTypeScan(const Call& call, const BlockBuilder& ctx) {
} else {
return TensorType(out_type, data_ty->ndim, data_ty->vdevice);
}
+} catch (const ffi::Error& error) {
+ return ffi::Unexpected(error);
+} catch (const std::exception& error) {
+ return ffi::Unexpected(ffi::Error("InternalError", error.what(), ""));
}
-Type InferTypeStatisticalExtension(const Call& call, const BlockBuilder& ctx) {
- TensorType data_ty = GetUnaryInputTensorType(call, ctx);
+ffi::Expected<Type> InferTypeStatisticalExtension(const CallNode* call_node) noexcept try {
+ const Call call = ffi::GetRef<Call>(call_node);
+ TensorType data_ty = GetUnaryInputTensorType(call);
const auto* attrs = call->attrs.as<StatisticalAttrs>();
std::vector<int> axes;
if (!data_ty->IsUnknownNdim() && attrs->axis.has_value()) {
- axes = NormalizeAxes(call, ctx, data_ty->ndim, attrs->axis.value());
+ axes = NormalizeAxes(call, data_ty->ndim, attrs->axis.value());
}
int out_ndim;
@@ -237,6 +248,10 @@ Type InferTypeStatisticalExtension(const Call& call, const BlockBuilder& ctx) {
else
return TupleType({TensorType(ShapeExpr(out_shape), data_ty->dtype, data_ty->vdevice),
TensorType(ShapeExpr(out_shape), PrimType::Int(64), data_ty->vdevice)});
+} catch (const ffi::Error& error) {
+ return ffi::Unexpected(error);
+} catch (const std::exception& error) {
+ return ffi::Unexpected(ffi::Error("InternalError", error.what(), ""));
}
/* relax.cumprod */
@@ -257,7 +272,7 @@ TVM_FFI_STATIC_INIT_BLOCK() {
OpDef("relax.cumprod")
.signature(sig::arg("data", "The input tensor."), sig::call_attrs<ScanopAttrs>())
- .set_attr<FInferType>("FInferType", InferTypeScan)
+ .set_attr<FInferType>("FInferType", FInferType::FromNative<&InferTypeScan>())
.set_attr<bool>("FPurity", true);
}
@@ -279,7 +294,7 @@ TVM_FFI_STATIC_INIT_BLOCK() {
OpDef("relax.cumsum")
.signature(sig::arg("data", "The input tensor."), sig::call_attrs<ScanopAttrs>())
- .set_attr<FInferType>("FInferType", InferTypeScan)
+ .set_attr<FInferType>("FInferType", FInferType::FromNative<&InferTypeScan>())
.set_attr<bool>("FPurity", true);
}
@@ -298,7 +313,7 @@ TVM_FFI_STATIC_INIT_BLOCK() {
OpDef("relax.median")
.signature(sig::arg("data", "The input tensor."), sig::call_attrs<StatisticalAttrs>())
- .set_attr<FInferType>("FInferType", InferTypeStatisticalExtension)
+ .set_attr<FInferType>("FInferType", FInferType::FromNative<&InferTypeStatisticalExtension>())
.set_attr<bool>("FPurity", true);
}
@@ -363,7 +378,7 @@ TVM_FFI_STATIC_INIT_BLOCK() {
OpDef("relax.max")
.signature(sig::arg("x", "The input data tensor"), sig::call_attrs<StatisticalAttrs>())
- .set_attr<FInferType>("FInferType", InferTypeStatistical)
+ .set_attr<FInferType>("FInferType", FInferType::FromNative<&InferTypeStatistical>())
.set_attr<FRelaxInferLayout>("FRelaxInferLayout", InferLayoutStatistical)
.set_attr<bool>("FPurity", true);
@@ -371,7 +386,7 @@ TVM_FFI_STATIC_INIT_BLOCK() {
OpDef("relax.mean")
.signature(sig::arg("x", "The input data tensor"), sig::call_attrs<StatisticalAttrs>())
- .set_attr<FInferType>("FInferType", InferTypeStatistical)
+ .set_attr<FInferType>("FInferType", FInferType::FromNative<&InferTypeStatistical>())
.set_attr<FRelaxInferLayout>("FRelaxInferLayout", InferLayoutStatistical)
.set_attr<bool>("FPurity", true);
@@ -379,7 +394,7 @@ TVM_FFI_STATIC_INIT_BLOCK() {
OpDef("relax.min")
.signature(sig::arg("x", "The input data tensor"), sig::call_attrs<StatisticalAttrs>())
- .set_attr<FInferType>("FInferType", InferTypeStatistical)
+ .set_attr<FInferType>("FInferType", FInferType::FromNative<&InferTypeStatistical>())
.set_attr<FRelaxInferLayout>("FRelaxInferLayout", InferLayoutStatistical)
.set_attr<bool>("FPurity", true);
@@ -387,7 +402,7 @@ TVM_FFI_STATIC_INIT_BLOCK() {
OpDef("relax.prod")
.signature(sig::arg("x", "The input data tensor"), sig::call_attrs<StatisticalAttrs>())
- .set_attr<FInferType>("FInferType", InferTypeStatistical)
+ .set_attr<FInferType>("FInferType", FInferType::FromNative<&InferTypeStatistical>())
.set_attr<FRelaxInferLayout>("FRelaxInferLayout", InferLayoutStatistical)
.set_attr<bool>("FPurity", true);
@@ -395,7 +410,7 @@ TVM_FFI_STATIC_INIT_BLOCK() {
OpDef("relax.std")
.signature(sig::arg("x", "The input data tensor"), sig::call_attrs<StatisticalAttrs>())
- .set_attr<FInferType>("FInferType", InferTypeStatistical)
+ .set_attr<FInferType>("FInferType", FInferType::FromNative<&InferTypeStatistical>())
.set_attr<FRelaxInferLayout>("FRelaxInferLayout", InferLayoutStatistical)
.set_attr<bool>("FPurity", true);
@@ -403,7 +418,7 @@ TVM_FFI_STATIC_INIT_BLOCK() {
OpDef("relax.sum")
.signature(sig::arg("x", "The input data tensor"), sig::call_attrs<StatisticalAttrs>())
- .set_attr<FInferType>("FInferType", InferTypeStatistical)
+ .set_attr<FInferType>("FInferType", FInferType::FromNative<&InferTypeStatistical>())
.set_attr<FRelaxInferLayout>("FRelaxInferLayout", InferLayoutStatistical)
.set_attr<bool>("FPurity", true);
@@ -411,7 +426,7 @@ TVM_FFI_STATIC_INIT_BLOCK() {
OpDef("relax.variance")
.signature(sig::arg("x", "The input data tensor"), sig::call_attrs<StatisticalAttrs>())
- .set_attr<FInferType>("FInferType", InferTypeStatistical)
+ .set_attr<FInferType>("FInferType", FInferType::FromNative<&InferTypeStatistical>())
.set_attr<FRelaxInferLayout>("FRelaxInferLayout", InferLayoutStatistical)
.set_attr<bool>("FPurity", true);
}
diff --git a/src/relax/op/tensor/ternary.cc b/src/relax/op/tensor/ternary.cc
index eb8250d9fa..c4357e2551 100644
--- a/src/relax/op/tensor/ternary.cc
+++ b/src/relax/op/tensor/ternary.cc
@@ -138,7 +138,7 @@ TVM_FFI_STATIC_INIT_BLOCK() {
sig::arg("x2", "The right hand operand of the multiplication"),
sig::arg("x3", "The operand of the addition"),
sig::var_ty_args("out_type", "Optional output tensor type carrying the virtual device."))
- .set_attr<FInferType>("FInferType", InferTypeEwiseFMA)
+ .set_attr<FInferTypeWithBuilder>("relax.FInferTypeWithBuilder", InferTypeEwiseFMA)
.set_attr<FRelaxInferLayout>("FRelaxInferLayout", InferLayoutEwiseFMA)
.set_attr<TMixedPrecisionPolicy>("TMixedPrecisionPolicy", MixedPrecisionPolicyKind::kFollow)
.set_attr<bool>("FPurity", true);
diff --git a/src/relax/op/tensor/unary.cc b/src/relax/op/tensor/unary.cc
index f67acb07c3..9ee0a8c561 100644
--- a/src/relax/op/tensor/unary.cc
+++ b/src/relax/op/tensor/unary.cc
@@ -31,9 +31,13 @@
namespace tvm {
namespace relax {
-Type InferTypeUnaryCheck(const Call& call, const BlockBuilder& ctx) {
- return InferTypeUnary<false>(call, ctx,
- [](const TensorType& input_ty) { return PrimType::Bool(); });
+ffi::Expected<Type> InferTypeUnaryCheck(const CallNode* call_node) noexcept try {
+ const Call call = ffi::GetRef<Call>(call_node);
+ return InferTypeUnary<false>(call, [](const TensorType& input_ty) { return PrimType::Bool(); });
+} catch (const ffi::Error& error) {
+ return ffi::Unexpected(error);
+} catch (const std::exception& error) {
+ return ffi::Unexpected(ffi::Error("InternalError", error.what(), ""));
}
/***************** Arithmetic operators *****************/
@@ -203,7 +207,8 @@ TVM_FFI_STATIC_INIT_BLOCK() {
.set_attr<FRelaxInferLayout>("FRelaxInferLayout", InferLayoutUnaryEwise)
.set_attr<TMixedPrecisionPolicy>("TMixedPrecisionPolicy", MixedPrecisionPolicyKind::kFollow)
.set_attr<bool>("FPurity", true)
- .set_attr<FInferType>("FInferType", InferTypeUnaryArith<false>);
+ .set_attr<FInferType>("FInferType",
+ FInferType::FromNative<&InferTypeUnaryArithContextFree<false>>());
tvm::ffi::reflection::GlobalDef().def("relax.op.acos", acos);
@@ -214,7 +219,8 @@ TVM_FFI_STATIC_INIT_BLOCK() {
.set_attr<FRelaxInferLayout>("FRelaxInferLayout", InferLayoutUnaryEwise)
.set_attr<TMixedPrecisionPolicy>("TMixedPrecisionPolicy", MixedPrecisionPolicyKind::kFollow)
.set_attr<bool>("FPurity", true)
- .set_attr<FInferType>("FInferType", InferTypeUnaryArith<true>);
+ .set_attr<FInferType>("FInferType",
+ FInferType::FromNative<&InferTypeUnaryArithContextFree<true>>());
tvm::ffi::reflection::GlobalDef().def("relax.op.acosh", acosh);
@@ -225,7 +231,8 @@ TVM_FFI_STATIC_INIT_BLOCK() {
.set_attr<FRelaxInferLayout>("FRelaxInferLayout", InferLayoutUnaryEwise)
.set_attr<TMixedPrecisionPolicy>("TMixedPrecisionPolicy", MixedPrecisionPolicyKind::kFollow)
.set_attr<bool>("FPurity", true)
- .set_attr<FInferType>("FInferType", InferTypeUnaryArith<true>);
+ .set_attr<FInferType>("FInferType",
+ FInferType::FromNative<&InferTypeUnaryArithContextFree<true>>());
tvm::ffi::reflection::GlobalDef().def("relax.op.asin", asin);
@@ -236,7 +243,8 @@ TVM_FFI_STATIC_INIT_BLOCK() {
.set_attr<FRelaxInferLayout>("FRelaxInferLayout", InferLayoutUnaryEwise)
.set_attr<TMixedPrecisionPolicy>("TMixedPrecisionPolicy", MixedPrecisionPolicyKind::kFollow)
.set_attr<bool>("FPurity", true)
- .set_attr<FInferType>("FInferType", InferTypeUnaryArith<true>);
+ .set_attr<FInferType>("FInferType",
+ FInferType::FromNative<&InferTypeUnaryArithContextFree<true>>());
tvm::ffi::reflection::GlobalDef().def("relax.op.asinh", asinh);
@@ -247,7 +255,8 @@ TVM_FFI_STATIC_INIT_BLOCK() {
.set_attr<FRelaxInferLayout>("FRelaxInferLayout", InferLayoutUnaryEwise)
.set_attr<TMixedPrecisionPolicy>("TMixedPrecisionPolicy", MixedPrecisionPolicyKind::kFollow)
.set_attr<bool>("FPurity", true)
- .set_attr<FInferType>("FInferType", InferTypeUnaryArith<true>);
+ .set_attr<FInferType>("FInferType",
+ FInferType::FromNative<&InferTypeUnaryArithContextFree<true>>());
tvm::ffi::reflection::GlobalDef().def("relax.op.atan", atan);
@@ -258,7 +267,8 @@ TVM_FFI_STATIC_INIT_BLOCK() {
.set_attr<FRelaxInferLayout>("FRelaxInferLayout", InferLayoutUnaryEwise)
.set_attr<TMixedPrecisionPolicy>("TMixedPrecisionPolicy", MixedPrecisionPolicyKind::kFollow)
.set_attr<bool>("FPurity", true)
- .set_attr<FInferType>("FInferType", InferTypeUnaryArith<true>);
+ .set_attr<FInferType>("FInferType",
+ FInferType::FromNative<&InferTypeUnaryArithContextFree<true>>());
tvm::ffi::reflection::GlobalDef().def("relax.op.atanh", atanh);
@@ -269,7 +279,8 @@ TVM_FFI_STATIC_INIT_BLOCK() {
.set_attr<FRelaxInferLayout>("FRelaxInferLayout", InferLayoutUnaryEwise)
.set_attr<TMixedPrecisionPolicy>("TMixedPrecisionPolicy", MixedPrecisionPolicyKind::kFollow)
.set_attr<bool>("FPurity", true)
- .set_attr<FInferType>("FInferType", InferTypeUnaryArith<true>);
+ .set_attr<FInferType>("FInferType",
+ FInferType::FromNative<&InferTypeUnaryArithContextFree<true>>());
tvm::ffi::reflection::GlobalDef().def("relax.op.bitwise_not", bitwise_not);
@@ -280,7 +291,8 @@ TVM_FFI_STATIC_INIT_BLOCK() {
.set_attr<FRelaxInferLayout>("FRelaxInferLayout", InferLayoutUnaryEwise)
.set_attr<TMixedPrecisionPolicy>("TMixedPrecisionPolicy", MixedPrecisionPolicyKind::kFollow)
.set_attr<bool>("FPurity", true)
- .set_attr<FInferType>("FInferType", InferTypeUnaryArith<false>);
+ .set_attr<FInferType>("FInferType",
+ FInferType::FromNative<&InferTypeUnaryArithContextFree<false>>());
tvm::ffi::reflection::GlobalDef().def("relax.op.ceil", ceil);
@@ -291,7 +303,8 @@ TVM_FFI_STATIC_INIT_BLOCK() {
.set_attr<FRelaxInferLayout>("FRelaxInferLayout", InferLayoutUnaryEwise)
.set_attr<TMixedPrecisionPolicy>("TMixedPrecisionPolicy", MixedPrecisionPolicyKind::kFollow)
.set_attr<bool>("FPurity", true)
- .set_attr<FInferType>("FInferType", InferTypeUnaryArith<false>);
+ .set_attr<FInferType>("FInferType",
+ FInferType::FromNative<&InferTypeUnaryArithContextFree<false>>());
tvm::ffi::reflection::GlobalDef().def("relax.op.cos", cos);
@@ -302,7 +315,8 @@ TVM_FFI_STATIC_INIT_BLOCK() {
.set_attr<FRelaxInferLayout>("FRelaxInferLayout", InferLayoutUnaryEwise)
.set_attr<TMixedPrecisionPolicy>("TMixedPrecisionPolicy", MixedPrecisionPolicyKind::kFollow)
.set_attr<bool>("FPurity", true)
- .set_attr<FInferType>("FInferType", InferTypeUnaryArith<true>);
+ .set_attr<FInferType>("FInferType",
+ FInferType::FromNative<&InferTypeUnaryArithContextFree<true>>());
tvm::ffi::reflection::GlobalDef().def("relax.op.cosh", cosh);
@@ -313,7 +327,8 @@ TVM_FFI_STATIC_INIT_BLOCK() {
.set_attr<FRelaxInferLayout>("FRelaxInferLayout", InferLayoutUnaryEwise)
.set_attr<TMixedPrecisionPolicy>("TMixedPrecisionPolicy", MixedPrecisionPolicyKind::kFollow)
.set_attr<bool>("FPurity", true)
- .set_attr<FInferType>("FInferType", InferTypeUnaryArith<true>);
+ .set_attr<FInferType>("FInferType",
+ FInferType::FromNative<&InferTypeUnaryArithContextFree<true>>());
tvm::ffi::reflection::GlobalDef().def("relax.op.exp", exp);
@@ -324,7 +339,8 @@ TVM_FFI_STATIC_INIT_BLOCK() {
.set_attr<FRelaxInferLayout>("FRelaxInferLayout", InferLayoutUnaryEwise)
.set_attr<TMixedPrecisionPolicy>("TMixedPrecisionPolicy", MixedPrecisionPolicyKind::kFollow)
.set_attr<bool>("FPurity", true)
- .set_attr<FInferType>("FInferType", InferTypeUnaryArith<true>);
+ .set_attr<FInferType>("FInferType",
+ FInferType::FromNative<&InferTypeUnaryArithContextFree<true>>());
tvm::ffi::reflection::GlobalDef().def("relax.op.floor", floor);
@@ -335,7 +351,8 @@ TVM_FFI_STATIC_INIT_BLOCK() {
.set_attr<FRelaxInferLayout>("FRelaxInferLayout", InferLayoutUnaryEwise)
.set_attr<TMixedPrecisionPolicy>("TMixedPrecisionPolicy", MixedPrecisionPolicyKind::kFollow)
.set_attr<bool>("FPurity", true)
- .set_attr<FInferType>("FInferType", InferTypeUnaryArith<false>);
+ .set_attr<FInferType>("FInferType",
+ FInferType::FromNative<&InferTypeUnaryArithContextFree<false>>());
tvm::ffi::reflection::GlobalDef().def("relax.op.log", log);
@@ -346,7 +363,8 @@ TVM_FFI_STATIC_INIT_BLOCK() {
.set_attr<FRelaxInferLayout>("FRelaxInferLayout", InferLayoutUnaryEwise)
.set_attr<TMixedPrecisionPolicy>("TMixedPrecisionPolicy", MixedPrecisionPolicyKind::kFollow)
.set_attr<bool>("FPurity", true)
- .set_attr<FInferType>("FInferType", InferTypeUnaryArith<true>);
+ .set_attr<FInferType>("FInferType",
+ FInferType::FromNative<&InferTypeUnaryArithContextFree<true>>());
tvm::ffi::reflection::GlobalDef().def("relax.op.logical_not", logical_not);
@@ -357,7 +375,8 @@ TVM_FFI_STATIC_INIT_BLOCK() {
.set_attr<FRelaxInferLayout>("FRelaxInferLayout", InferLayoutUnaryEwise)
.set_attr<TMixedPrecisionPolicy>("TMixedPrecisionPolicy", MixedPrecisionPolicyKind::kFollow)
.set_attr<bool>("FPurity", true)
- .set_attr<FInferType>("FInferType", InferTypeUnaryArith<false>);
+ .set_attr<FInferType>("FInferType",
+ FInferType::FromNative<&InferTypeUnaryArithContextFree<false>>());
tvm::ffi::reflection::GlobalDef().def("relax.op.negative", negative);
@@ -368,7 +387,8 @@ TVM_FFI_STATIC_INIT_BLOCK() {
.set_attr<FRelaxInferLayout>("FRelaxInferLayout", InferLayoutUnaryEwise)
.set_attr<TMixedPrecisionPolicy>("TMixedPrecisionPolicy", MixedPrecisionPolicyKind::kFollow)
.set_attr<bool>("FPurity", true)
- .set_attr<FInferType>("FInferType", InferTypeUnaryArith<false>);
+ .set_attr<FInferType>("FInferType",
+ FInferType::FromNative<&InferTypeUnaryArithContextFree<false>>());
tvm::ffi::reflection::GlobalDef().def("relax.op.round", round);
@@ -379,7 +399,8 @@ TVM_FFI_STATIC_INIT_BLOCK() {
.set_attr<FRelaxInferLayout>("FRelaxInferLayout", InferLayoutUnaryEwise)
.set_attr<TMixedPrecisionPolicy>("TMixedPrecisionPolicy", MixedPrecisionPolicyKind::kFollow)
.set_attr<bool>("FPurity", true)
- .set_attr<FInferType>("FInferType", InferTypeUnaryArith<false>);
+ .set_attr<FInferType>("FInferType",
+ FInferType::FromNative<&InferTypeUnaryArithContextFree<false>>());
tvm::ffi::reflection::GlobalDef().def("relax.op.rsqrt", rsqrt);
@@ -390,7 +411,8 @@ TVM_FFI_STATIC_INIT_BLOCK() {
.set_attr<FRelaxInferLayout>("FRelaxInferLayout", InferLayoutUnaryEwise)
.set_attr<TMixedPrecisionPolicy>("TMixedPrecisionPolicy", MixedPrecisionPolicyKind::kFollow)
.set_attr<bool>("FPurity", true)
- .set_attr<FInferType>("FInferType", InferTypeUnaryArith<true>);
+ .set_attr<FInferType>("FInferType",
+ FInferType::FromNative<&InferTypeUnaryArithContextFree<true>>());
tvm::ffi::reflection::GlobalDef().def("relax.op.sigmoid", sigmoid);
@@ -401,7 +423,8 @@ TVM_FFI_STATIC_INIT_BLOCK() {
.set_attr<FRelaxInferLayout>("FRelaxInferLayout", InferLayoutUnaryEwise)
.set_attr<TMixedPrecisionPolicy>("TMixedPrecisionPolicy", MixedPrecisionPolicyKind::kFollow)
.set_attr<bool>("FPurity", true)
- .set_attr<FInferType>("FInferType", InferTypeUnaryArith<true>);
+ .set_attr<FInferType>("FInferType",
+ FInferType::FromNative<&InferTypeUnaryArithContextFree<true>>());
tvm::ffi::reflection::GlobalDef().def("relax.op.sign", sign);
@@ -412,7 +435,8 @@ TVM_FFI_STATIC_INIT_BLOCK() {
.set_attr<FRelaxInferLayout>("FRelaxInferLayout", InferLayoutUnaryEwise)
.set_attr<TMixedPrecisionPolicy>("TMixedPrecisionPolicy", MixedPrecisionPolicyKind::kFollow)
.set_attr<bool>("FPurity", true)
- .set_attr<FInferType>("FInferType", InferTypeUnaryArith<false>);
+ .set_attr<FInferType>("FInferType",
+ FInferType::FromNative<&InferTypeUnaryArithContextFree<false>>());
tvm::ffi::reflection::GlobalDef().def("relax.op.sin", sin);
@@ -423,7 +447,8 @@ TVM_FFI_STATIC_INIT_BLOCK() {
.set_attr<FRelaxInferLayout>("FRelaxInferLayout", InferLayoutUnaryEwise)
.set_attr<TMixedPrecisionPolicy>("TMixedPrecisionPolicy", MixedPrecisionPolicyKind::kFollow)
.set_attr<bool>("FPurity", true)
- .set_attr<FInferType>("FInferType", InferTypeUnaryArith<true>);
+ .set_attr<FInferType>("FInferType",
+ FInferType::FromNative<&InferTypeUnaryArithContextFree<true>>());
tvm::ffi::reflection::GlobalDef().def("relax.op.sinh", sinh);
@@ -434,7 +459,8 @@ TVM_FFI_STATIC_INIT_BLOCK() {
.set_attr<FRelaxInferLayout>("FRelaxInferLayout", InferLayoutUnaryEwise)
.set_attr<TMixedPrecisionPolicy>("TMixedPrecisionPolicy", MixedPrecisionPolicyKind::kFollow)
.set_attr<bool>("FPurity", true)
- .set_attr<FInferType>("FInferType", InferTypeUnaryArith<true>);
+ .set_attr<FInferType>("FInferType",
+ FInferType::FromNative<&InferTypeUnaryArithContextFree<true>>());
tvm::ffi::reflection::GlobalDef().def("relax.op.square", square);
@@ -445,7 +471,8 @@ TVM_FFI_STATIC_INIT_BLOCK() {
.set_attr<FRelaxInferLayout>("FRelaxInferLayout", InferLayoutUnaryEwise)
.set_attr<TMixedPrecisionPolicy>("TMixedPrecisionPolicy", MixedPrecisionPolicyKind::kFollow)
.set_attr<bool>("FPurity", true)
- .set_attr<FInferType>("FInferType", InferTypeUnaryArith<false>);
+ .set_attr<FInferType>("FInferType",
+ FInferType::FromNative<&InferTypeUnaryArithContextFree<false>>());
tvm::ffi::reflection::GlobalDef().def("relax.op.sqrt", sqrt);
@@ -456,7 +483,8 @@ TVM_FFI_STATIC_INIT_BLOCK() {
.set_attr<FRelaxInferLayout>("FRelaxInferLayout", InferLayoutUnaryEwise)
.set_attr<TMixedPrecisionPolicy>("TMixedPrecisionPolicy", MixedPrecisionPolicyKind::kFollow)
.set_attr<bool>("FPurity", true)
- .set_attr<FInferType>("FInferType", InferTypeUnaryArith<true>);
+ .set_attr<FInferType>("FInferType",
+ FInferType::FromNative<&InferTypeUnaryArithContextFree<true>>());
tvm::ffi::reflection::GlobalDef().def("relax.op.tan", tan);
@@ -467,7 +495,8 @@ TVM_FFI_STATIC_INIT_BLOCK() {
.set_attr<FRelaxInferLayout>("FRelaxInferLayout", InferLayoutUnaryEwise)
.set_attr<TMixedPrecisionPolicy>("TMixedPrecisionPolicy", MixedPrecisionPolicyKind::kFollow)
.set_attr<bool>("FPurity", true)
- .set_attr<FInferType>("FInferType", InferTypeUnaryArith<true>);
+ .set_attr<FInferType>("FInferType",
+ FInferType::FromNative<&InferTypeUnaryArithContextFree<true>>());
tvm::ffi::reflection::GlobalDef().def("relax.op.tanh", tanh);
@@ -478,7 +507,8 @@ TVM_FFI_STATIC_INIT_BLOCK() {
.set_attr<FRelaxInferLayout>("FRelaxInferLayout", InferLayoutUnaryEwise)
.set_attr<TMixedPrecisionPolicy>("TMixedPrecisionPolicy", MixedPrecisionPolicyKind::kFollow)
.set_attr<bool>("FPurity", true)
- .set_attr<FInferType>("FInferType", InferTypeUnaryArith<true>);
+ .set_attr<FInferType>("FInferType",
+ FInferType::FromNative<&InferTypeUnaryArithContextFree<true>>());
tvm::ffi::reflection::GlobalDef().def("relax.op.trunc", trunc);
@@ -489,7 +519,8 @@ TVM_FFI_STATIC_INIT_BLOCK() {
.set_attr<FRelaxInferLayout>("FRelaxInferLayout", InferLayoutUnaryEwise)
.set_attr<TMixedPrecisionPolicy>("TMixedPrecisionPolicy", MixedPrecisionPolicyKind::kFollow)
.set_attr<bool>("FPurity", true)
- .set_attr<FInferType>("FInferType", InferTypeUnaryArith<false>);
+ .set_attr<FInferType>("FInferType",
+ FInferType::FromNative<&InferTypeUnaryArithContextFree<false>>());
tvm::ffi::reflection::GlobalDef().def("relax.op.erf", erf);
@@ -500,7 +531,8 @@ TVM_FFI_STATIC_INIT_BLOCK() {
.set_attr<FRelaxInferLayout>("FRelaxInferLayout", InferLayoutUnaryEwise)
.set_attr<TMixedPrecisionPolicy>("TMixedPrecisionPolicy", MixedPrecisionPolicyKind::kFollow)
.set_attr<bool>("FPurity", true)
- .set_attr<FInferType>("FInferType", InferTypeUnaryArith<true>);
+ .set_attr<FInferType>("FInferType",
+ FInferType::FromNative<&InferTypeUnaryArithContextFree<true>>());
// relax.clip
@@ -508,7 +540,8 @@ TVM_FFI_STATIC_INIT_BLOCK() {
.signature(sig::arg("x", "The input tensor."),
sig::arg("min", "The lower-bound of the range to be clipped to"),
sig::arg("max", "The upper-bound of the range to be clipped to"))
- .set_attr<FInferType>("FInferType", ReturnTypeFromArg<0>)
+ .set_attr<FInferType>("FInferType",
+ FInferType::FromNative<&ReturnTypeFromArgContextFree<0>>())
.set_attr<bool>("FPurity", true);
}
@@ -538,8 +571,10 @@ TVM_FFI_STATIC_INIT_BLOCK() {
.set_attr<FRelaxInferLayout>("FRelaxInferLayout", InferLayoutUnaryEwise)
.set_attr<TMixedPrecisionPolicy>("TMixedPrecisionPolicy", MixedPrecisionPolicyKind::kFollow)
.set_attr<bool>("FPurity", true)
- .set_attr<FInferType>("FInferType",
- InferTypeUnaryCheck); // require_float_dtype=false for check op
+ .set_attr<FInferType>(
+ "FInferType",
+ FInferType::FromNative<&InferTypeUnaryCheck>()); // require_float_dtype=false
+ // for check op
tvm::ffi::reflection::GlobalDef().def("relax.op.isinf", isinf);
@@ -550,8 +585,10 @@ TVM_FFI_STATIC_INIT_BLOCK() {
.set_attr<FRelaxInferLayout>("FRelaxInferLayout", InferLayoutUnaryEwise)
.set_attr<TMixedPrecisionPolicy>("TMixedPrecisionPolicy", MixedPrecisionPolicyKind::kFollow)
.set_attr<bool>("FPurity", true)
- .set_attr<FInferType>("FInferType",
- InferTypeUnaryCheck); // require_float_dtype=false for check op
+ .set_attr<FInferType>(
+ "FInferType",
+ FInferType::FromNative<&InferTypeUnaryCheck>()); // require_float_dtype=false
+ // for check op
tvm::ffi::reflection::GlobalDef().def("relax.op.isnan", isnan);
@@ -562,8 +599,10 @@ TVM_FFI_STATIC_INIT_BLOCK() {
.set_attr<FRelaxInferLayout>("FRelaxInferLayout", InferLayoutUnaryEwise)
.set_attr<TMixedPrecisionPolicy>("TMixedPrecisionPolicy", MixedPrecisionPolicyKind::kFollow)
.set_attr<bool>("FPurity", true)
- .set_attr<FInferType>("FInferType",
- InferTypeUnaryCheck); // require_float_dtype=false for check op
+ .set_attr<FInferType>(
+ "FInferType",
+ FInferType::FromNative<&InferTypeUnaryCheck>()); // require_float_dtype=false
+ // for check op
}
} // namespace relax
diff --git a/src/relax/op/vision/multibox_transform_loc.cc b/src/relax/op/vision/multibox_transform_loc.cc
index 15daf0543f..58e2712732 100644
--- a/src/relax/op/vision/multibox_transform_loc.cc
+++ b/src/relax/op/vision/multibox_transform_loc.cc
@@ -68,7 +68,8 @@ TVM_FFI_STATIC_INIT_BLOCK() {
* skips those N-based relations; other checks (ndim, dtype, loc dim divisible by 4, etc.)
* still apply when their inputs are known.
*/
-Type InferTypeMultiboxTransformLoc(const Call& call, const BlockBuilder& ctx) {
+ffi::Expected<Type> InferTypeMultiboxTransformLoc(const CallNode* call_node) noexcept try {
+ const Call call = ffi::GetRef<Call>(call_node);
if (call->args.size() != 3) {
TVM_FFI_VISIT_THROW(ValueError, call)
<< "multibox_transform_loc: expected 3 inputs (cls_pred, loc_pred, anchor), "
@@ -76,7 +77,7 @@ Type InferTypeMultiboxTransformLoc(const Call& call, const BlockBuilder& ctx) {
<< call->args.size();
}
- ffi::Array<TensorType> input_ty = GetInputTensorType(call, ctx);
+ ffi::Array<TensorType> input_ty = GetInputTensorType(call);
const auto cls_ty = input_ty[0];
const auto loc_ty = input_ty[1];
const auto anchor_ty = input_ty[2];
@@ -185,6 +186,10 @@ Type InferTypeMultiboxTransformLoc(const Call& call, const BlockBuilder& ctx) {
ffi::Array<Type> fields = {TensorType(ShapeExpr(boxes_shape), cls_ty->dtype, vdev),
TensorType(ShapeExpr(scores_shape), cls_ty->dtype, vdev)};
return TupleType(fields);
+} catch (const ffi::Error& error) {
+ return ffi::Unexpected(error);
+} catch (const std::exception& error) {
+ return ffi::Unexpected(ffi::Error("InternalError", error.what(), ""));
}
TVM_FFI_STATIC_INIT_BLOCK() {
@@ -197,7 +202,7 @@ TVM_FFI_STATIC_INIT_BLOCK() {
"[B,4*N] box encodings (x,y,w,h); TFLite yxhw order remapped to xywh."),
sig::arg("anchor", "[1,N,4] priors as ltrb (left,top,right,bottom)."),
sig::call_attrs<MultiboxTransformLocAttrs>())
- .set_attr<FInferType>("FInferType", InferTypeMultiboxTransformLoc)
+ .set_attr<FInferType>("FInferType", FInferType::FromNative<&InferTypeMultiboxTransformLoc>())
.set_attr<bool>("FPurity", true);
}
diff --git a/src/relax/op/vision/nms.cc b/src/relax/op/vision/nms.cc
index d3a0e74ba6..59a41f37a3 100644
--- a/src/relax/op/vision/nms.cc
+++ b/src/relax/op/vision/nms.cc
@@ -63,8 +63,9 @@ TVM_FFI_STATIC_INIT_BLOCK() {
all_class_non_max_suppression);
}
-Type InferTypeAllClassNMS(const Call& call, const BlockBuilder& ctx) {
- tvm::ffi::Array<TensorType> input_ty = GetInputTensorType(call, ctx);
+ffi::Expected<Type> InferTypeAllClassNMS(const CallNode* call_node) noexcept try {
+ const Call call = ffi::GetRef<Call>(call_node);
+ tvm::ffi::Array<TensorType> input_ty = GetInputTensorType(call);
const auto boxes_ty = input_ty[0];
const auto scores_ty = input_ty[1];
TVM_FFI_ICHECK(!boxes_ty->IsUnknownNdim()) << "Only support known ndim";
@@ -101,6 +102,10 @@ Type InferTypeAllClassNMS(const Call& call, const BlockBuilder& ctx) {
TensorType(scores_shape, PrimType::Float(32), vdev),
TensorType(counts_shape, PrimType::Int(64), vdev)};
return TupleType(fields);
+} catch (const ffi::Error& error) {
+ return ffi::Unexpected(error);
+} catch (const std::exception& error) {
+ return ffi::Unexpected(ffi::Error("InternalError", error.what(), ""));
}
TVM_FFI_STATIC_INIT_BLOCK() {
@@ -113,7 +118,7 @@ TVM_FFI_STATIC_INIT_BLOCK() {
sig::arg("iou_threshold", "The IoU threshold for box the overlap test."),
sig::arg("score_threshold", "The score threshold to filter out low score boxes early."),
sig::call_attrs<AllClassNonMaximumSuppressionAttrs>())
- .set_attr<FInferType>("FInferType", InferTypeAllClassNMS)
+ .set_attr<FInferType>("FInferType", FInferType::FromNative<&InferTypeAllClassNMS>())
.set_attr<bool>("FPurity", true);
}
@@ -134,7 +139,8 @@ TVM_FFI_STATIC_INIT_BLOCK() {
refl::GlobalDef().def("relax.op.vision.get_valid_counts", get_valid_counts);
}
-Type InferTypeGetValidCounts(const Call& call, const BlockBuilder& ctx) {
+ffi::Expected<Type> InferTypeGetValidCounts(const CallNode* call_node) noexcept try {
+ const Call call = ffi::GetRef<Call>(call_node);
if (call->args.size() != 1) {
TVM_FFI_VISIT_THROW(ValueError, call)
<< "get_valid_counts expects 1 argument, got " << call->args.size();
@@ -182,13 +188,17 @@ Type InferTypeGetValidCounts(const Call& call, const BlockBuilder& ctx) {
TensorType(ShapeExpr({batch, num_anchors, elem_length}), data_ty->dtype, vdev),
TensorType(ShapeExpr({batch, num_anchors}), PrimType::Int(32), vdev)};
return TupleType(fields);
+} catch (const ffi::Error& error) {
+ return ffi::Unexpected(error);
+} catch (const std::exception& error) {
+ return ffi::Unexpected(ffi::Error("InternalError", error.what(), ""));
}
TVM_FFI_STATIC_INIT_BLOCK() {
OpDef("relax.vision.get_valid_counts")
.signature(sig::arg("data", "Input data, 3-D tensor [batch_size, num_anchors, elem_length]."),
sig::call_attrs<GetValidCountsAttrs>())
- .set_attr<FInferType>("FInferType", InferTypeGetValidCounts)
+ .set_attr<FInferType>("FInferType", FInferType::FromNative<&InferTypeGetValidCounts>())
.set_attr<bool>("FPurity", true);
}
@@ -362,7 +372,7 @@ TVM_FFI_STATIC_INIT_BLOCK() {
sig::arg("valid_count", "1-D tensor for valid number of boxes."),
sig::arg("indices", "2-D tensor with shape [batch_size, num_anchors]."),
sig::call_attrs<NonMaximumSuppressionAttrs>())
- .set_attr<FInferType>("FInferType", InferTypeNMS)
+ .set_attr<FInferTypeWithBuilder>("relax.FInferTypeWithBuilder", InferTypeNMS)
.set_attr<bool>("FPurity", true);
}
diff --git a/src/relax/op/vision/roi_align.cc b/src/relax/op/vision/roi_align.cc
index ea08faad6c..2ba5820b8e 100644
--- a/src/relax/op/vision/roi_align.cc
+++ b/src/relax/op/vision/roi_align.cc
@@ -60,7 +60,8 @@ TVM_FFI_STATIC_INIT_BLOCK() {
refl::GlobalDef().def("relax.op.vision.roi_align", roi_align);
}
-Type InferTypeROIAlign(const Call& call, const BlockBuilder& ctx) {
+ffi::Expected<Type> InferTypeROIAlign(const CallNode* call_node) noexcept try {
+ const Call call = ffi::GetRef<Call>(call_node);
if (call->args.size() != 2) {
TVM_FFI_VISIT_THROW(ValueError, call)
<< "ROIAlign expects two arguments, while the given number of arguments is "
@@ -126,6 +127,10 @@ Type InferTypeROIAlign(const Call& call, const BlockBuilder& ctx) {
IntImm::Int32(attrs->pooled_size[1]), data_shape[3]};
}
return TensorType(ShapeExpr(out_shape), data_ty->dtype, data_ty->vdevice);
+} catch (const ffi::Error& error) {
+ return ffi::Unexpected(error);
+} catch (const std::exception& error) {
+ return ffi::Unexpected(ffi::Error("InternalError", error.what(), ""));
}
TVM_FFI_STATIC_INIT_BLOCK() {
@@ -135,7 +140,7 @@ TVM_FFI_STATIC_INIT_BLOCK() {
sig::arg("rois",
"The input rois with shape (num_roi, 5) in [batch_idx, x1, y1, x2, y2] format."),
sig::call_attrs<ROIAlignAttrs>())
- .set_attr<FInferType>("FInferType", InferTypeROIAlign)
+ .set_attr<FInferType>("FInferType", FInferType::FromNative<&InferTypeROIAlign>())
.set_attr<TMixedPrecisionPolicy>("TMixedPrecisionPolicy", MixedPrecisionPolicyKind::kFollow)
.set_attr<bool>("FPurity", true);
}
diff --git a/src/relax/op/vision/roi_pool.cc b/src/relax/op/vision/roi_pool.cc
index fdd8f534d2..54af51a697 100644
--- a/src/relax/op/vision/roi_pool.cc
+++ b/src/relax/op/vision/roi_pool.cc
@@ -57,7 +57,8 @@ TVM_FFI_STATIC_INIT_BLOCK() {
refl::GlobalDef().def("relax.op.vision.roi_pool", roi_pool);
}
-Type InferTypeROIPool(const Call& call, const BlockBuilder& ctx) {
+ffi::Expected<Type> InferTypeROIPool(const CallNode* call_node) noexcept try {
+ const Call call = ffi::GetRef<Call>(call_node);
if (call->args.size() != 2) {
TVM_FFI_VISIT_THROW(ValueError, call)
<< "ROIPool expects two arguments, while the given number of arguments is "
@@ -114,6 +115,10 @@ Type InferTypeROIPool(const Call& call, const BlockBuilder& ctx) {
IntImm::Int32(attrs->pooled_size[0]),
IntImm::Int32(attrs->pooled_size[1])};
return TensorType(ShapeExpr(out_shape), data_ty->dtype, data_ty->vdevice);
+} catch (const ffi::Error& error) {
+ return ffi::Unexpected(error);
+} catch (const std::exception& error) {
+ return ffi::Unexpected(ffi::Error("InternalError", error.what(), ""));
}
TVM_FFI_STATIC_INIT_BLOCK() {
@@ -123,7 +128,7 @@ TVM_FFI_STATIC_INIT_BLOCK() {
sig::arg("rois",
"The input rois with shape (num_roi, 5) in [batch_idx, x1, y1, x2, y2] format."),
sig::call_attrs<ROIPoolAttrs>())
- .set_attr<FInferType>("FInferType", InferTypeROIPool)
+ .set_attr<FInferType>("FInferType", FInferType::FromNative<&InferTypeROIPool>())
.set_attr<TMixedPrecisionPolicy>("TMixedPrecisionPolicy", MixedPrecisionPolicyKind::kFollow)
.set_attr<bool>("FPurity", true);
}
diff --git a/src/relax/script/printer/utils.h b/src/relax/script/printer/utils.h
index fade104fdc..9b42a43dc8 100644
--- a/src/relax/script/printer/utils.h
+++ b/src/relax/script/printer/utils.h
@@ -106,10 +106,17 @@ inline ffi::Optional<ExprDoc> TypeAsAnn(const tvm::Var& v, const AccessPath& v_p
if (auto opt = call->op.as<Op>()) {
auto op = opt.value();
- static auto op_map_infer_ty = Op::GetAttrMap<relax::FInferType>("FInferType");
-
- auto temp_builder = relax::BlockBuilder::Create(std::nullopt);
- inferred_ty = op_map_infer_ty[op](call, temp_builder);
+ static auto op_map_context_free = Op::GetAttrMap<FInferType>("FInferType");
+ if (op_map_context_free.count(op)) {
+ inferred_ty = Call::ReinferType(call.get());
+ } else {
+ static auto op_map_infer_ty =
+ Op::GetAttrMap<relax::FInferTypeWithBuilder>("relax.FInferTypeWithBuilder");
+ if (op_map_infer_ty.count(op)) {
+ auto temp_builder = relax::BlockBuilder::Create(std::nullopt);
+ inferred_ty = op_map_infer_ty[op](call, temp_builder);
+ }
+ }
} else if (auto opt = call->op.as<relax::FuncType>()) {
auto temp_builder = relax::BlockBuilder::Create(std::nullopt);
inferred_ty =
diff --git a/src/relax/transform/compute_prim_value.cc b/src/relax/transform/compute_prim_value.cc
index b18aa3506e..fb87d10396 100644
--- a/src/relax/transform/compute_prim_value.cc
+++ b/src/relax/transform/compute_prim_value.cc
@@ -36,8 +36,11 @@ bool HasRelaxCallCapabilities(const CallNode* call) {
auto op = call->op.as<Op>();
if (!op) return true;
static auto infer_type_map = Op::GetAttrMap<FInferType>("FInferType");
+ static auto infer_type_with_builder_map =
+ Op::GetAttrMap<FInferTypeWithBuilder>("relax.FInferTypeWithBuilder");
static auto legalize_map = Op::GetAttrMap<FLegalize>("FLegalize");
- return infer_type_map.count(op.value()) || legalize_map.count(op.value());
+ return infer_type_map.count(op.value()) || infer_type_with_builder_map.count(op.value()) ||
+ legalize_map.count(op.value());
}
class PrimExprComputeInjector : public ExprMutator {
diff --git a/src/relax/transform/fold_constant.cc b/src/relax/transform/fold_constant.cc
index cbcdf0c3eb..142d4cb75b 100644
--- a/src/relax/transform/fold_constant.cc
+++ b/src/relax/transform/fold_constant.cc
@@ -314,6 +314,8 @@ class ConstantFolder : public ExprMutator {
static const Op call_tir_op = Op::Get("relax.call_tir");
static const auto& infer_type_map = Op::GetAttrMap<FInferType>("FInferType");
+ static const auto& infer_type_with_builder_map =
+ Op::GetAttrMap<FInferTypeWithBuilder>("relax.FInferTypeWithBuilder");
static const auto& legalize_map = Op::GetAttrMap<FLegalize>("FLegalize");
auto* op_node = post_call->op.as<OpNode>();
@@ -347,7 +349,8 @@ class ConstantFolder : public ExprMutator {
new_args.push_back(arg);
}
Type ret_ty = Type::Missing();
- if (post_call->ty.as<PrimTypeNode>() && !infer_type_map.count(op)) {
+ if (post_call->ty.as<PrimTypeNode>() && !infer_type_map.count(op) &&
+ !infer_type_with_builder_map.count(op)) {
ret_ty = post_call->ty.as_or_throw<Type>();
}
post_call = Call(ret_ty, post_call->op, new_args, post_call->attrs, post_call->ty_args,
diff --git a/src/relax/transform/gradient.cc b/src/relax/transform/gradient.cc
index 1795859b02..4d82b2f6f7 100644
--- a/src/relax/transform/gradient.cc
+++ b/src/relax/transform/gradient.cc
@@ -269,7 +269,9 @@ class CheckpointGenerator : private ExprMutator {
if (call_node->ty.as<PrimTypeNode>()) {
if (auto op = call_node->op.as<Op>()) {
static auto infer_type_map = Op::GetAttrMap<FInferType>("FInferType");
- if (!infer_type_map.count(op.value())) {
+ static auto infer_type_with_builder_map =
+ Op::GetAttrMap<FInferTypeWithBuilder>("relax.FInferTypeWithBuilder");
+ if (!infer_type_map.count(op.value()) && !infer_type_with_builder_map.count(op.value())) {
ret_ty = call_node->ty.as_or_throw<Type>();
}
}
diff --git a/src/relax/transform/run_codegen.cc b/src/relax/transform/run_codegen.cc
index 656aa58120..90aeb4582f 100644
--- a/src/relax/transform/run_codegen.cc
+++ b/src/relax/transform/run_codegen.cc
@@ -147,7 +147,9 @@ class CodeGenRunner : ExprMutator {
if (call_node->ty.as<PrimTypeNode>()) {
if (auto op = call_node->op.as<Op>()) {
static auto infer_type_map = Op::GetAttrMap<FInferType>("FInferType");
- if (!infer_type_map.count(op.value())) {
+ static auto infer_type_with_builder_map =
+ Op::GetAttrMap<FInferTypeWithBuilder>("relax.FInferTypeWithBuilder");
+ if (!infer_type_map.count(op.value()) && !infer_type_with_builder_map.count(op.value())) {
ret_ty = call_node->ty.as_or_throw<Type>();
}
}
diff --git a/src/s_tir/stmt.cc b/src/s_tir/stmt.cc
index 45cbd47b0b..d119d6cb1d 100644
--- a/src/s_tir/stmt.cc
+++ b/src/s_tir/stmt.cc
@@ -35,17 +35,19 @@ using namespace tvm::prim;
namespace {
-TVMFFIAny MatchBufferRegionVisit(ffi::StructuralVisitorObj* visitor, ffi::AnyView value) noexcept {
+TVM_FFI_INLINE ffi::Expected<ffi::Optional<ffi::VisitInterrupt>> MatchBufferRegionVisit(
+ ffi::StructuralVisitorObj* visitor, ffi::AnyView value) noexcept {
const MatchBufferRegionNode* self =
ffi::details::AnyUnsafe::RawObjectPtrFromAnyViewAfterCheck<const MatchBufferRegionNode>(
value);
TVM_FFI_S_VISIT_MAYBE_EARLY_RETURN(visitor->WithDefRegionKind(
kTVMFFIDefRegionKindSimple, [&]() { return visitor->VisitExpected(self->buffer); }));
TVM_FFI_S_VISIT_MAYBE_EARLY_RETURN(visitor->VisitExpected(self->source));
- return ffi::AnyView(nullptr).CopyToTVMFFIAny();
+ return std::nullopt;
}
-TVMFFIAny MatchBufferRegionMutate(ffi::StructuralMutatorObj* mutator, ffi::AnyView value) noexcept {
+TVM_FFI_INLINE ffi::Expected<ffi::UnchangedOr<ffi::Any>> MatchBufferRegionMutate(
+ ffi::StructuralMutatorObj* mutator, ffi::AnyView value) noexcept {
const MatchBufferRegionNode* self =
ffi::details::AnyUnsafe::RawObjectPtrFromAnyViewAfterCheck<const MatchBufferRegionNode>(
value);
@@ -57,16 +59,16 @@ TVMFFIAny MatchBufferRegionMutate(ffi::StructuralMutatorObj* mutator, ffi::AnyVi
mutator->MutateExpected(self->source));
if (mapped_buffer.UnchangedOrSameAs(self->buffer) &&
mapped_source.UnchangedOrSameAs(self->source)) {
- return ffi::Unchanged().CopyToTVMFFIAny();
+ return ffi::Unchanged();
}
ffi::ObjectPtr<MatchBufferRegionNode> copy = ffi::make_object<MatchBufferRegionNode>(*self);
copy->buffer = std::move(mapped_buffer).ValueOrUnchanged(std::move(copy->buffer));
copy->source = std::move(mapped_source).ValueOrUnchanged(std::move(copy->source));
- return ffi::details::AnyUnsafe::MoveAnyToTVMFFIAny(ffi::Any(std::move(copy)));
+ return ffi::Any(std::move(copy));
}
-TVMFFIAny MatchBufferRegionMaybeInplaceMutate(ffi::StructuralMutatorObj* mutator,
- ffi::AnyView value) noexcept {
+TVM_FFI_INLINE ffi::Expected<ffi::UnchangedOr<ffi::Any>> MatchBufferRegionMaybeInplaceMutate(
+ ffi::StructuralMutatorObj* mutator, ffi::AnyView value) noexcept {
MatchBufferRegionNode* self = const_cast<MatchBufferRegionNode*>(
ffi::details::AnyUnsafe::RawObjectPtrFromAnyViewAfterCheck<const MatchBufferRegionNode>(
value));
@@ -80,14 +82,15 @@ TVMFFIAny MatchBufferRegionMaybeInplaceMutate(ffi::StructuralMutatorObj* mutator
mutator->MutateExpected(self->source, ffi::InplaceMode::kAllow));
if (mapped_buffer.UnchangedOrSameAs(self->buffer) &&
mapped_source.UnchangedOrSameAs(self->source)) {
- return ffi::Unchanged().CopyToTVMFFIAny();
+ return ffi::Unchanged();
}
if (!mapped_buffer.IsUnchanged()) self->buffer = std::move(mapped_buffer).ValueUnchecked();
if (!mapped_source.IsUnchanged()) self->source = std::move(mapped_source).ValueUnchecked();
- return ffi::Unchanged().CopyToTVMFFIAny();
+ return ffi::Unchanged();
}
-TVMFFIAny SBlockVisit(ffi::StructuralVisitorObj* visitor, ffi::AnyView value) noexcept {
+TVM_FFI_INLINE ffi::Expected<ffi::Optional<ffi::VisitInterrupt>> SBlockVisit(
+ ffi::StructuralVisitorObj* visitor, ffi::AnyView value) noexcept {
// Establish allocation and match-buffer definitions before their region uses.
// Whole iterators and annotations remain part of structural traversal.
// skips: name_hint
@@ -102,10 +105,11 @@ TVMFFIAny SBlockVisit(ffi::StructuralVisitorObj* visitor, ffi::AnyView value) no
TVM_FFI_S_VISIT_MAYBE_EARLY_RETURN(visitor->VisitExpected(self->annotations));
TVM_FFI_S_VISIT_MAYBE_EARLY_RETURN(visitor->VisitExpected(self->init));
TVM_FFI_S_VISIT_MAYBE_EARLY_RETURN(visitor->VisitExpected(self->body));
- return ffi::AnyView(nullptr).CopyToTVMFFIAny();
+ return std::nullopt;
}
-TVMFFIAny SBlockMutate(ffi::StructuralMutatorObj* mutator, ffi::AnyView value) noexcept {
+TVM_FFI_INLINE ffi::Expected<ffi::UnchangedOr<ffi::Any>> SBlockMutate(
+ ffi::StructuralMutatorObj* mutator, ffi::AnyView value) noexcept {
// Establish allocation and match-buffer definitions before their region uses.
// Whole iterators and annotations remain part of structural traversal.
// skips: name_hint
@@ -138,7 +142,7 @@ TVMFFIAny SBlockMutate(ffi::StructuralMutatorObj* mutator, ffi::AnyView value) n
mapped_match_buffers.UnchangedOrSameAs(self->match_buffers) &&
mapped_annotations.UnchangedOrSameAs(self->annotations) &&
mapped_init.UnchangedOrSameAs(self->init) && mapped_body.UnchangedOrSameAs(self->body)) {
- return ffi::Unchanged().CopyToTVMFFIAny();
+ return ffi::Unchanged();
}
ffi::ObjectPtr<SBlockNode> copy = ffi::make_object<SBlockNode>(*self);
copy->iter_vars = std::move(mapped_iter_vars).ValueOrUnchanged(std::move(copy->iter_vars));
@@ -151,11 +155,11 @@ TVMFFIAny SBlockMutate(ffi::StructuralMutatorObj* mutator, ffi::AnyView value) n
copy->annotations = std::move(mapped_annotations).ValueOrUnchanged(std::move(copy->annotations));
copy->init = std::move(mapped_init).ValueOrUnchanged(std::move(copy->init));
copy->body = std::move(mapped_body).ValueOrUnchanged(std::move(copy->body));
- return ffi::details::AnyUnsafe::MoveAnyToTVMFFIAny(ffi::Any(std::move(copy)));
+ return ffi::Any(std::move(copy));
}
-TVMFFIAny SBlockMaybeInplaceMutate(ffi::StructuralMutatorObj* mutator,
- ffi::AnyView value) noexcept {
+TVM_FFI_INLINE ffi::Expected<ffi::UnchangedOr<ffi::Any>> SBlockMaybeInplaceMutate(
+ ffi::StructuralMutatorObj* mutator, ffi::AnyView value) noexcept {
// Establish allocation and match-buffer definitions before their region uses.
// Whole iterators and annotations remain part of structural traversal.
// skips: name_hint
@@ -192,7 +196,7 @@ TVMFFIAny SBlockMaybeInplaceMutate(ffi::StructuralMutatorObj* mutator,
mapped_match_buffers.UnchangedOrSameAs(self->match_buffers) &&
mapped_annotations.UnchangedOrSameAs(self->annotations) &&
mapped_init.UnchangedOrSameAs(self->init) && mapped_body.UnchangedOrSameAs(self->body)) {
- return ffi::Unchanged().CopyToTVMFFIAny();
+ return ffi::Unchanged();
}
if (!mapped_iter_vars.IsUnchanged())
self->iter_vars = std::move(mapped_iter_vars).ValueUnchecked();
@@ -208,19 +212,21 @@ TVMFFIAny SBlockMaybeInplaceMutate(ffi::StructuralMutatorObj* mutator,
self->annotations = std::move(mapped_annotations).ValueUnchecked();
if (!mapped_init.IsUnchanged()) self->init = std::move(mapped_init).ValueUnchecked();
if (!mapped_body.IsUnchanged()) self->body = std::move(mapped_body).ValueUnchecked();
- return ffi::Unchanged().CopyToTVMFFIAny();
+ return ffi::Unchanged();
}
-TVMFFIAny SBlockRealizeVisit(ffi::StructuralVisitorObj* visitor, ffi::AnyView value) noexcept {
+TVM_FFI_INLINE ffi::Expected<ffi::Optional<ffi::VisitInterrupt>> SBlockRealizeVisit(
+ ffi::StructuralVisitorObj* visitor, ffi::AnyView value) noexcept {
const SBlockRealizeNode* self =
ffi::details::AnyUnsafe::RawObjectPtrFromAnyViewAfterCheck<const SBlockRealizeNode>(value);
TVM_FFI_S_VISIT_MAYBE_EARLY_RETURN(visitor->VisitExpected(self->iter_values));
TVM_FFI_S_VISIT_MAYBE_EARLY_RETURN(visitor->VisitExpected(self->predicate));
TVM_FFI_S_VISIT_MAYBE_EARLY_RETURN(visitor->VisitExpected(self->block));
- return ffi::AnyView(nullptr).CopyToTVMFFIAny();
+ return std::nullopt;
}
-TVMFFIAny SBlockRealizeMutate(ffi::StructuralMutatorObj* mutator, ffi::AnyView value) noexcept {
+TVM_FFI_INLINE ffi::Expected<ffi::UnchangedOr<ffi::Any>> SBlockRealizeMutate(
+ ffi::StructuralMutatorObj* mutator, ffi::AnyView value) noexcept {
const SBlockRealizeNode* self =
ffi::details::AnyUnsafe::RawObjectPtrFromAnyViewAfterCheck<const SBlockRealizeNode>(value);
TVM_FFI_S_MUTATE_ASSIGN_OR_RETURN(ffi::UnchangedOr<ffi::Array<PrimExpr>>, mapped_iter_values,
@@ -232,17 +238,17 @@ TVMFFIAny SBlockRealizeMutate(ffi::StructuralMutatorObj* mutator, ffi::AnyView v
if (mapped_iter_values.UnchangedOrSameAs(self->iter_values) &&
mapped_predicate.UnchangedOrSameAs(self->predicate) &&
mapped_block.UnchangedOrSameAs(self->block)) {
- return ffi::Unchanged().CopyToTVMFFIAny();
+ return ffi::Unchanged();
}
ffi::ObjectPtr<SBlockRealizeNode> copy = ffi::make_object<SBlockRealizeNode>(*self);
copy->iter_values = std::move(mapped_iter_values).ValueOrUnchanged(std::move(copy->iter_values));
copy->predicate = std::move(mapped_predicate).ValueOrUnchanged(std::move(copy->predicate));
copy->block = std::move(mapped_block).ValueOrUnchanged(std::move(copy->block));
- return ffi::details::AnyUnsafe::MoveAnyToTVMFFIAny(ffi::Any(std::move(copy)));
+ return ffi::Any(std::move(copy));
}
-TVMFFIAny SBlockRealizeMaybeInplaceMutate(ffi::StructuralMutatorObj* mutator,
- ffi::AnyView value) noexcept {
+TVM_FFI_INLINE ffi::Expected<ffi::UnchangedOr<ffi::Any>> SBlockRealizeMaybeInplaceMutate(
+ ffi::StructuralMutatorObj* mutator, ffi::AnyView value) noexcept {
SBlockRealizeNode* self = const_cast<SBlockRealizeNode*>(
ffi::details::AnyUnsafe::RawObjectPtrFromAnyViewAfterCheck<const SBlockRealizeNode>(value));
TVM_FFI_S_MUTATE_ASSIGN_OR_RETURN(
@@ -256,14 +262,14 @@ TVMFFIAny SBlockRealizeMaybeInplaceMutate(ffi::StructuralMutatorObj* mutator,
if (mapped_iter_values.UnchangedOrSameAs(self->iter_values) &&
mapped_predicate.UnchangedOrSameAs(self->predicate) &&
mapped_block.UnchangedOrSameAs(self->block)) {
- return ffi::Unchanged().CopyToTVMFFIAny();
+ return ffi::Unchanged();
}
if (!mapped_iter_values.IsUnchanged())
self->iter_values = std::move(mapped_iter_values).ValueUnchecked();
if (!mapped_predicate.IsUnchanged())
self->predicate = std::move(mapped_predicate).ValueUnchecked();
if (!mapped_block.IsUnchanged()) self->block = std::move(mapped_block).ValueUnchecked();
- return ffi::Unchanged().CopyToTVMFFIAny();
+ return ffi::Unchanged();
}
} // namespace
@@ -319,10 +325,12 @@ TVM_FFI_STATIC_INIT_BLOCK() {
namespace refl = tvm::ffi::reflection;
MatchBufferRegionNode::RegisterReflection();
refl::TypeAttrDef<MatchBufferRegionNode>()
- .attr(refl::type_attr::kStructuralVisit, reinterpret_cast<void*>(&MatchBufferRegionVisit))
- .attr(refl::type_attr::kStructuralMutate, reinterpret_cast<void*>(&MatchBufferRegionMutate))
+ .attr(refl::type_attr::kStructuralVisit,
+ ffi::FStructuralVisit::FromNative<&MatchBufferRegionVisit>())
+ .attr(refl::type_attr::kStructuralMutate,
+ ffi::FStructuralMutate::FromNative<&MatchBufferRegionMutate>())
.attr(refl::type_attr::kStructuralMaybeInplaceMutate,
- reinterpret_cast<void*>(&MatchBufferRegionMaybeInplaceMutate));
+ ffi::FStructuralMutate::FromNative<&MatchBufferRegionMaybeInplaceMutate>());
refl::GlobalDef().def("s_tir.MatchBufferRegion", [](BufferVar buffer, TensorRegion source) {
return MatchBufferRegion(buffer, source);
@@ -375,10 +383,10 @@ TVM_FFI_STATIC_INIT_BLOCK() {
namespace refl = tvm::ffi::reflection;
SBlockNode::RegisterReflection();
refl::TypeAttrDef<SBlockNode>()
- .attr(refl::type_attr::kStructuralVisit, reinterpret_cast<void*>(&SBlockVisit))
- .attr(refl::type_attr::kStructuralMutate, reinterpret_cast<void*>(&SBlockMutate))
+ .attr(refl::type_attr::kStructuralVisit, ffi::FStructuralVisit::FromNative<&SBlockVisit>())
+ .attr(refl::type_attr::kStructuralMutate, ffi::FStructuralMutate::FromNative<&SBlockMutate>())
.attr(refl::type_attr::kStructuralMaybeInplaceMutate,
- reinterpret_cast<void*>(&SBlockMaybeInplaceMutate));
+ ffi::FStructuralMutate::FromNative<&SBlockMaybeInplaceMutate>());
refl::GlobalDef().def("s_tir.SBlock",
[](ffi::Array<IterVar> iter_vars, ffi::Array<TensorRegion> reads,
@@ -411,10 +419,12 @@ TVM_FFI_STATIC_INIT_BLOCK() {
namespace refl = tvm::ffi::reflection;
SBlockRealizeNode::RegisterReflection();
refl::TypeAttrDef<SBlockRealizeNode>()
- .attr(refl::type_attr::kStructuralVisit, reinterpret_cast<void*>(&SBlockRealizeVisit))
- .attr(refl::type_attr::kStructuralMutate, reinterpret_cast<void*>(&SBlockRealizeMutate))
+ .attr(refl::type_attr::kStructuralVisit,
+ ffi::FStructuralVisit::FromNative<&SBlockRealizeVisit>())
+ .attr(refl::type_attr::kStructuralMutate,
+ ffi::FStructuralMutate::FromNative<&SBlockRealizeMutate>())
.attr(refl::type_attr::kStructuralMaybeInplaceMutate,
- reinterpret_cast<void*>(&SBlockRealizeMaybeInplaceMutate));
+ ffi::FStructuralMutate::FromNative<&SBlockRealizeMaybeInplaceMutate>());
refl::GlobalDef().def("s_tir.SBlockRealize", [](ffi::Array<PrimExpr> iter_values,
PrimExpr predicate, SBlock block, Span span) {
diff --git a/src/te/tensor.cc b/src/te/tensor.cc
index 228fb2b7d4..059b519258 100644
--- a/src/te/tensor.cc
+++ b/src/te/tensor.cc
@@ -63,16 +63,19 @@ ffi::Array<PrimExpr> ValidateTensorLoad(const Call& call, Tensor* tensor_out) {
return indices;
}
-TVMFFIAny TensorVisit(ffi::StructuralVisitorObj*, ffi::AnyView) noexcept {
- return ffi::AnyView(nullptr).CopyToTVMFFIAny();
+TVM_FFI_INLINE ffi::Expected<ffi::Optional<ffi::VisitInterrupt>> TensorVisit(
+ ffi::StructuralVisitorObj*, ffi::AnyView) noexcept {
+ return std::nullopt;
}
-TVMFFIAny TensorMutate(ffi::StructuralMutatorObj*, ffi::AnyView) noexcept {
- return ffi::Unchanged().CopyToTVMFFIAny();
+TVM_FFI_INLINE ffi::Expected<ffi::UnchangedOr<ffi::Any>> TensorMutate(ffi::StructuralMutatorObj*,
+ ffi::AnyView) noexcept {
+ return ffi::Unchanged();
}
-TVMFFIAny TensorMaybeInplaceMutate(ffi::StructuralMutatorObj*, ffi::AnyView) noexcept {
- return ffi::Unchanged().CopyToTVMFFIAny();
+TVM_FFI_INLINE ffi::Expected<ffi::UnchangedOr<ffi::Any>> TensorMaybeInplaceMutate(
+ ffi::StructuralMutatorObj*, ffi::AnyView) noexcept {
+ return ffi::Unchanged();
}
} // namespace
@@ -171,10 +174,10 @@ TVM_FFI_STATIC_INIT_BLOCK() {
namespace refl = tvm::ffi::reflection;
TensorNode::RegisterReflection();
refl::TypeAttrDef<TensorNode>()
- .attr(refl::type_attr::kStructuralVisit, reinterpret_cast<void*>(&TensorVisit))
- .attr(refl::type_attr::kStructuralMutate, reinterpret_cast<void*>(&TensorMutate))
+ .attr(refl::type_attr::kStructuralVisit, ffi::FStructuralVisit::FromNative<&TensorVisit>())
+ .attr(refl::type_attr::kStructuralMutate, ffi::FStructuralMutate::FromNative<&TensorMutate>())
.attr(refl::type_attr::kStructuralMaybeInplaceMutate,
- reinterpret_cast<void*>(&TensorMaybeInplaceMutate));
+ ffi::FStructuralMutate::FromNative<&TensorMaybeInplaceMutate>());
}
bool IsTensorLoad(const Expr& expr) {
diff --git a/src/tirx/ir/function.cc b/src/tirx/ir/function.cc
index ccc2b25228..92e69cb246 100644
--- a/src/tirx/ir/function.cc
+++ b/src/tirx/ir/function.cc
@@ -76,7 +76,8 @@ tvm::Type InferType(const PrimFunc& prim_func) {
return relax::FuncType(params, ret, purity);
}
-TVMFFIAny PrimFuncVisit(ffi::StructuralVisitorObj* visitor, ffi::AnyView value) noexcept {
+TVM_FFI_INLINE ffi::Expected<ffi::Optional<ffi::VisitInterrupt>> PrimFuncVisit(
+ ffi::StructuralVisitorObj* visitor, ffi::AnyView value) noexcept {
// skips: attrs (metadata), ty (derived by InferType)
const PrimFuncNode* self =
ffi::details::AnyUnsafe::RawObjectPtrFromAnyViewAfterCheck<const PrimFuncNode>(value);
@@ -84,10 +85,11 @@ TVMFFIAny PrimFuncVisit(ffi::StructuralVisitorObj* visitor, ffi::AnyView value)
kTVMFFIDefRegionKindPattern, [&]() { return visitor->VisitExpected(self->params); }));
TVM_FFI_S_VISIT_MAYBE_EARLY_RETURN(visitor->VisitExpected(self->ret_type));
TVM_FFI_S_VISIT_MAYBE_EARLY_RETURN(visitor->VisitExpected(self->body));
- return ffi::AnyView(nullptr).CopyToTVMFFIAny();
+ return std::nullopt;
}
-TVMFFIAny PrimFuncMutate(ffi::StructuralMutatorObj* mutator, ffi::AnyView value) noexcept {
+TVM_FFI_INLINE ffi::Expected<ffi::UnchangedOr<ffi::Any>> PrimFuncMutate(
+ ffi::StructuralMutatorObj* mutator, ffi::AnyView value) noexcept {
// skips: attrs (metadata), ty (derived by InferType)
const PrimFuncNode* self =
ffi::details::AnyUnsafe::RawObjectPtrFromAnyViewAfterCheck<const PrimFuncNode>(value);
@@ -102,17 +104,17 @@ TVMFFIAny PrimFuncMutate(ffi::StructuralMutatorObj* mutator, ffi::AnyView value)
if (mapped_params.UnchangedOrSameAs(self->params) &&
mapped_ret_type.UnchangedOrSameAs(self->ret_type) &&
mapped_body.UnchangedOrSameAs(self->body)) {
- return ffi::Unchanged().CopyToTVMFFIAny();
+ return ffi::Unchanged();
}
ffi::ObjectPtr<PrimFuncNode> copy = ffi::make_object<PrimFuncNode>(*self);
copy->params = std::move(mapped_params).ValueOrUnchanged(std::move(copy->params));
copy->ret_type = std::move(mapped_ret_type).ValueOrUnchanged(std::move(copy->ret_type));
copy->body = std::move(mapped_body).ValueOrUnchanged(std::move(copy->body));
- return ffi::details::AnyUnsafe::MoveAnyToTVMFFIAny(ffi::Any(std::move(copy)));
+ return ffi::Any(std::move(copy));
}
-TVMFFIAny PrimFuncMaybeInplaceMutate(ffi::StructuralMutatorObj* mutator,
- ffi::AnyView value) noexcept {
+TVM_FFI_INLINE ffi::Expected<ffi::UnchangedOr<ffi::Any>> PrimFuncMaybeInplaceMutate(
+ ffi::StructuralMutatorObj* mutator, ffi::AnyView value) noexcept {
// skips: attrs (metadata), ty (derived by InferType)
PrimFuncNode* self = const_cast<PrimFuncNode*>(
ffi::details::AnyUnsafe::RawObjectPtrFromAnyViewAfterCheck<const PrimFuncNode>(value));
@@ -135,7 +137,7 @@ TVMFFIAny PrimFuncMaybeInplaceMutate(ffi::StructuralMutatorObj* mutator,
if (!mapped_body.UnchangedOrSameAs(self->body)) {
self->body = std::move(mapped_body).ValueUnchecked();
}
- return ffi::Unchanged().CopyToTVMFFIAny();
+ return ffi::Unchanged();
}
} // namespace
@@ -163,10 +165,11 @@ TVM_FFI_STATIC_INIT_BLOCK() {
namespace refl = tvm::ffi::reflection;
PrimFuncNode::RegisterReflection();
refl::TypeAttrDef<PrimFuncNode>()
- .attr(refl::type_attr::kStructuralVisit, reinterpret_cast<void*>(&PrimFuncVisit))
- .attr(refl::type_attr::kStructuralMutate, reinterpret_cast<void*>(&PrimFuncMutate))
+ .attr(refl::type_attr::kStructuralVisit, ffi::FStructuralVisit::FromNative<&PrimFuncVisit>())
+ .attr(refl::type_attr::kStructuralMutate,
+ ffi::FStructuralMutate::FromNative<&PrimFuncMutate>())
.attr(refl::type_attr::kStructuralMaybeInplaceMutate,
- reinterpret_cast<void*>(&PrimFuncMaybeInplaceMutate));
+ ffi::FStructuralMutate::FromNative<&PrimFuncMaybeInplaceMutate>());
refl::GlobalDef().def("tirx.PrimFunc",
[](ffi::Array<tirx::Var> params, Stmt body, Type ret_type, DictAttrs attrs,
diff --git a/src/tirx/ir/iter_var.cc b/src/tirx/ir/iter_var.cc
index 1ab656d3ee..42e2d4cff2 100644
--- a/src/tirx/ir/iter_var.cc
+++ b/src/tirx/ir/iter_var.cc
@@ -33,17 +33,19 @@ namespace tirx {
namespace {
-TVMFFIAny IterVarVisit(ffi::StructuralVisitorObj* visitor, ffi::AnyView value) noexcept {
+TVM_FFI_INLINE ffi::Expected<ffi::Optional<ffi::VisitInterrupt>> IterVarVisit(
+ ffi::StructuralVisitorObj* visitor, ffi::AnyView value) noexcept {
// skips: iter_type, thread_tag
const IterVarNode* self =
ffi::details::AnyUnsafe::RawObjectPtrFromAnyViewAfterCheck<const IterVarNode>(value);
TVM_FFI_S_VISIT_MAYBE_EARLY_RETURN(visitor->VisitExpected(self->dom));
TVM_FFI_S_VISIT_MAYBE_EARLY_RETURN(visitor->WithDefRegionKind(
kTVMFFIDefRegionKindSimple, [&]() { return visitor->VisitExpected(self->var); }));
- return ffi::AnyView(nullptr).CopyToTVMFFIAny();
+ return std::nullopt;
}
-TVMFFIAny IterVarMutate(ffi::StructuralMutatorObj* mutator, ffi::AnyView value) noexcept {
+TVM_FFI_INLINE ffi::Expected<ffi::UnchangedOr<ffi::Any>> IterVarMutate(
+ ffi::StructuralMutatorObj* mutator, ffi::AnyView value) noexcept {
// skips: iter_type, thread_tag
const IterVarNode* self =
ffi::details::AnyUnsafe::RawObjectPtrFromAnyViewAfterCheck<const IterVarNode>(value);
@@ -54,16 +56,16 @@ TVMFFIAny IterVarMutate(ffi::StructuralMutatorObj* mutator, ffi::AnyView value)
return mutator->MutateExpected(self->var);
}));
if (mapped_dom.UnchangedOrSameAs(self->dom) && mapped_var.UnchangedOrSameAs(self->var)) {
- return ffi::Unchanged().CopyToTVMFFIAny();
+ return ffi::Unchanged();
}
ffi::ObjectPtr<IterVarNode> copy = ffi::make_object<IterVarNode>(*self);
copy->dom = std::move(mapped_dom).ValueOrUnchanged(std::move(copy->dom));
copy->var = std::move(mapped_var).ValueOrUnchanged(std::move(copy->var));
- return ffi::details::AnyUnsafe::MoveAnyToTVMFFIAny(ffi::Any(std::move(copy)));
+ return ffi::Any(std::move(copy));
}
-TVMFFIAny IterVarMaybeInplaceMutate(ffi::StructuralMutatorObj* mutator,
- ffi::AnyView value) noexcept {
+TVM_FFI_INLINE ffi::Expected<ffi::UnchangedOr<ffi::Any>> IterVarMaybeInplaceMutate(
+ ffi::StructuralMutatorObj* mutator, ffi::AnyView value) noexcept {
// skips: iter_type, thread_tag
IterVarNode* self = const_cast<IterVarNode*>(
ffi::details::AnyUnsafe::RawObjectPtrFromAnyViewAfterCheck<const IterVarNode>(value));
@@ -75,11 +77,11 @@ TVMFFIAny IterVarMaybeInplaceMutate(ffi::StructuralMutatorObj* mutator,
ffi::InplaceMode::kAllow);
}));
if (mapped_dom.UnchangedOrSameAs(self->dom) && mapped_var.UnchangedOrSameAs(self->var)) {
- return ffi::Unchanged().CopyToTVMFFIAny();
+ return ffi::Unchanged();
}
self->dom = std::move(mapped_dom).ValueOrUnchanged(std::move(self->dom));
self->var = std::move(mapped_var).ValueOrUnchanged(std::move(self->var));
- return ffi::Unchanged().CopyToTVMFFIAny();
+ return ffi::Unchanged();
}
} // namespace
@@ -110,10 +112,11 @@ TVM_FFI_STATIC_INIT_BLOCK() {
namespace refl = tvm::ffi::reflection;
IterVarNode::RegisterReflection();
refl::TypeAttrDef<IterVarNode>()
- .attr(refl::type_attr::kStructuralVisit, reinterpret_cast<void*>(&IterVarVisit))
- .attr(refl::type_attr::kStructuralMutate, reinterpret_cast<void*>(&IterVarMutate))
+ .attr(refl::type_attr::kStructuralVisit, ffi::FStructuralVisit::FromNative<&IterVarVisit>())
+ .attr(refl::type_attr::kStructuralMutate,
+ ffi::FStructuralMutate::FromNative<&IterVarMutate>())
.attr(refl::type_attr::kStructuralMaybeInplaceMutate,
- reinterpret_cast<void*>(&IterVarMaybeInplaceMutate));
+ ffi::FStructuralMutate::FromNative<&IterVarMaybeInplaceMutate>());
}
TVM_FFI_STATIC_INIT_BLOCK() {
diff --git a/src/tirx/ir/layout/tile_core.cc b/src/tirx/ir/layout/tile_core.cc
index e3f4dedff4..7af13da0e1 100644
--- a/src/tirx/ir/layout/tile_core.cc
+++ b/src/tirx/ir/layout/tile_core.cc
@@ -35,16 +35,18 @@ using namespace tvm::prim;
namespace {
-TVMFFIAny IterVisit(ffi::StructuralVisitorObj* visitor, ffi::AnyView value) noexcept {
+TVM_FFI_INLINE ffi::Expected<ffi::Optional<ffi::VisitInterrupt>> IterVisit(
+ ffi::StructuralVisitorObj* visitor, ffi::AnyView value) noexcept {
const IterNode* self =
ffi::details::AnyUnsafe::RawObjectPtrFromAnyViewAfterCheck<const IterNode>(value);
TVM_FFI_S_VISIT_MAYBE_EARLY_RETURN(visitor->VisitExpected(self->extent));
TVM_FFI_S_VISIT_MAYBE_EARLY_RETURN(visitor->VisitExpected(self->stride));
TVM_FFI_S_VISIT_MAYBE_EARLY_RETURN(visitor->VisitExpected(self->axis));
- return ffi::AnyView(nullptr).CopyToTVMFFIAny();
+ return std::nullopt;
}
-TVMFFIAny IterMutate(ffi::StructuralMutatorObj* mutator, ffi::AnyView value) noexcept {
+TVM_FFI_INLINE ffi::Expected<ffi::UnchangedOr<ffi::Any>> IterMutate(
+ ffi::StructuralMutatorObj* mutator, ffi::AnyView value) noexcept {
const IterNode* self =
ffi::details::AnyUnsafe::RawObjectPtrFromAnyViewAfterCheck<const IterNode>(value);
TVM_FFI_S_MUTATE_ASSIGN_OR_RETURN(ffi::UnchangedOr<PrimExpr>, mapped_extent,
@@ -55,16 +57,17 @@ TVMFFIAny IterMutate(ffi::StructuralMutatorObj* mutator, ffi::AnyView value) noe
mutator->MutateExpected(self->axis));
if (mapped_extent.UnchangedOrSameAs(self->extent) &&
mapped_stride.UnchangedOrSameAs(self->stride) && mapped_axis.UnchangedOrSameAs(self->axis)) {
- return ffi::Unchanged().CopyToTVMFFIAny();
+ return ffi::Unchanged();
}
ffi::ObjectPtr<IterNode> copy = ffi::make_object<IterNode>(*self);
copy->extent = std::move(mapped_extent).ValueOrUnchanged(std::move(copy->extent));
copy->stride = std::move(mapped_stride).ValueOrUnchanged(std::move(copy->stride));
copy->axis = std::move(mapped_axis).ValueOrUnchanged(std::move(copy->axis));
- return ffi::details::AnyUnsafe::MoveAnyToTVMFFIAny(ffi::Any(std::move(copy)));
+ return ffi::Any(std::move(copy));
}
-TVMFFIAny IterMaybeInplaceMutate(ffi::StructuralMutatorObj* mutator, ffi::AnyView value) noexcept {
+TVM_FFI_INLINE ffi::Expected<ffi::UnchangedOr<ffi::Any>> IterMaybeInplaceMutate(
+ ffi::StructuralMutatorObj* mutator, ffi::AnyView value) noexcept {
IterNode* self = const_cast<IterNode*>(
ffi::details::AnyUnsafe::RawObjectPtrFromAnyViewAfterCheck<const IterNode>(value));
TVM_FFI_S_MUTATE_ASSIGN_OR_RETURN(
@@ -77,24 +80,26 @@ TVMFFIAny IterMaybeInplaceMutate(ffi::StructuralMutatorObj* mutator, ffi::AnyVie
mutator->MutateExpected(self->axis, ffi::InplaceMode::kAllow));
if (mapped_extent.UnchangedOrSameAs(self->extent) &&
mapped_stride.UnchangedOrSameAs(self->stride) && mapped_axis.UnchangedOrSameAs(self->axis)) {
- return ffi::Unchanged().CopyToTVMFFIAny();
+ return ffi::Unchanged();
}
self->extent = std::move(mapped_extent).ValueOrUnchanged(std::move(self->extent));
self->stride = std::move(mapped_stride).ValueOrUnchanged(std::move(self->stride));
self->axis = std::move(mapped_axis).ValueOrUnchanged(std::move(self->axis));
- return ffi::Unchanged().CopyToTVMFFIAny();
+ return ffi::Unchanged();
}
-TVMFFIAny TileLayoutVisit(ffi::StructuralVisitorObj* visitor, ffi::AnyView value) noexcept {
+TVM_FFI_INLINE ffi::Expected<ffi::Optional<ffi::VisitInterrupt>> TileLayoutVisit(
+ ffi::StructuralVisitorObj* visitor, ffi::AnyView value) noexcept {
const TileLayoutNode* self =
ffi::details::AnyUnsafe::RawObjectPtrFromAnyViewAfterCheck<const TileLayoutNode>(value);
TVM_FFI_S_VISIT_MAYBE_EARLY_RETURN(visitor->VisitExpected(self->shard));
TVM_FFI_S_VISIT_MAYBE_EARLY_RETURN(visitor->VisitExpected(self->replica));
TVM_FFI_S_VISIT_MAYBE_EARLY_RETURN(visitor->VisitExpected(self->offset));
- return ffi::AnyView(nullptr).CopyToTVMFFIAny();
+ return std::nullopt;
}
-TVMFFIAny TileLayoutMutate(ffi::StructuralMutatorObj* mutator, ffi::AnyView value) noexcept {
+TVM_FFI_INLINE ffi::Expected<ffi::UnchangedOr<ffi::Any>> TileLayoutMutate(
+ ffi::StructuralMutatorObj* mutator, ffi::AnyView value) noexcept {
const TileLayoutNode* self =
ffi::details::AnyUnsafe::RawObjectPtrFromAnyViewAfterCheck<const TileLayoutNode>(value);
TVM_FFI_S_MUTATE_ASSIGN_OR_RETURN(ffi::UnchangedOr<ffi::Array<Iter>>, mapped_shard,
@@ -107,17 +112,17 @@ TVMFFIAny TileLayoutMutate(ffi::StructuralMutatorObj* mutator, ffi::AnyView valu
if (mapped_shard.UnchangedOrSameAs(self->shard) &&
mapped_replica.UnchangedOrSameAs(self->replica) &&
mapped_offset.UnchangedOrSameAs(self->offset)) {
- return ffi::Unchanged().CopyToTVMFFIAny();
+ return ffi::Unchanged();
}
ffi::ObjectPtr<TileLayoutNode> copy = ffi::make_object<TileLayoutNode>(*self);
copy->shard = std::move(mapped_shard).ValueOrUnchanged(std::move(copy->shard));
copy->replica = std::move(mapped_replica).ValueOrUnchanged(std::move(copy->replica));
copy->offset = std::move(mapped_offset).ValueOrUnchanged(std::move(copy->offset));
- return ffi::details::AnyUnsafe::MoveAnyToTVMFFIAny(ffi::Any(std::move(copy)));
+ return ffi::Any(std::move(copy));
}
-TVMFFIAny TileLayoutMaybeInplaceMutate(ffi::StructuralMutatorObj* mutator,
- ffi::AnyView value) noexcept {
+TVM_FFI_INLINE ffi::Expected<ffi::UnchangedOr<ffi::Any>> TileLayoutMaybeInplaceMutate(
+ ffi::StructuralMutatorObj* mutator, ffi::AnyView value) noexcept {
TileLayoutNode* self = const_cast<TileLayoutNode*>(
ffi::details::AnyUnsafe::RawObjectPtrFromAnyViewAfterCheck<const TileLayoutNode>(value));
TVM_FFI_S_MUTATE_ASSIGN_OR_RETURN(ffi::UnchangedOr<ffi::Array<Iter>>, mapped_shard,
@@ -132,12 +137,12 @@ TVMFFIAny TileLayoutMaybeInplaceMutate(ffi::StructuralMutatorObj* mutator,
if (mapped_shard.UnchangedOrSameAs(self->shard) &&
mapped_replica.UnchangedOrSameAs(self->replica) &&
mapped_offset.UnchangedOrSameAs(self->offset)) {
- return ffi::Unchanged().CopyToTVMFFIAny();
+ return ffi::Unchanged();
}
self->shard = std::move(mapped_shard).ValueOrUnchanged(std::move(self->shard));
self->replica = std::move(mapped_replica).ValueOrUnchanged(std::move(self->replica));
self->offset = std::move(mapped_offset).ValueOrUnchanged(std::move(self->offset));
- return ffi::Unchanged().CopyToTVMFFIAny();
+ return ffi::Unchanged();
}
} // namespace
@@ -157,10 +162,10 @@ TVM_FFI_STATIC_INIT_BLOCK() {
namespace refl = tvm::ffi::reflection;
IterNode::RegisterReflection();
refl::TypeAttrDef<IterNode>()
- .attr(refl::type_attr::kStructuralVisit, reinterpret_cast<void*>(&IterVisit))
- .attr(refl::type_attr::kStructuralMutate, reinterpret_cast<void*>(&IterMutate))
+ .attr(refl::type_attr::kStructuralVisit, ffi::FStructuralVisit::FromNative<&IterVisit>())
+ .attr(refl::type_attr::kStructuralMutate, ffi::FStructuralMutate::FromNative<&IterMutate>())
.attr(refl::type_attr::kStructuralMaybeInplaceMutate,
- reinterpret_cast<void*>(&IterMaybeInplaceMutate));
+ ffi::FStructuralMutate::FromNative<&IterMaybeInplaceMutate>());
refl::GlobalDef().def("tirx.Iter", [](PrimExpr extent, PrimExpr stride, Axis axis) {
return Iter(extent, stride, axis);
@@ -181,10 +186,12 @@ TVM_FFI_STATIC_INIT_BLOCK() {
namespace refl = tvm::ffi::reflection;
TileLayoutNode::RegisterReflection();
refl::TypeAttrDef<TileLayoutNode>()
- .attr(refl::type_attr::kStructuralVisit, reinterpret_cast<void*>(&TileLayoutVisit))
- .attr(refl::type_attr::kStructuralMutate, reinterpret_cast<void*>(&TileLayoutMutate))
+ .attr(refl::type_attr::kStructuralVisit,
+ ffi::FStructuralVisit::FromNative<&TileLayoutVisit>())
+ .attr(refl::type_attr::kStructuralMutate,
+ ffi::FStructuralMutate::FromNative<&TileLayoutMutate>())
.attr(refl::type_attr::kStructuralMaybeInplaceMutate,
- reinterpret_cast<void*>(&TileLayoutMaybeInplaceMutate));
+ ffi::FStructuralMutate::FromNative<&TileLayoutMaybeInplaceMutate>());
refl::GlobalDef().def("tirx.TileLayout", [](ffi::Array<Iter> shard, ffi::Array<Iter> replica,
ffi::Map<Axis, PrimExpr> offset) {
diff --git a/src/tirx/ir/stmt.cc b/src/tirx/ir/stmt.cc
index 6e9dd8ac3d..0bfbc1be25 100644
--- a/src/tirx/ir/stmt.cc
+++ b/src/tirx/ir/stmt.cc
@@ -69,62 +69,97 @@ bool IntImmValueFits(const ffi::BigInt& value, const PrimType& ty) {
// Structural traversal hooks
-TVMFFIAny BindVisit(ffi::StructuralVisitorObj* visitor, ffi::AnyView value) noexcept {
+TVM_FFI_INLINE ffi::Expected<ffi::Optional<ffi::VisitInterrupt>> BindVisit(
+ ffi::StructuralVisitorObj* visitor, ffi::AnyView value) noexcept {
const BindNode* self =
ffi::details::AnyUnsafe::RawObjectPtrFromAnyViewAfterCheck<const BindNode>(value);
TVM_FFI_S_VISIT_MAYBE_EARLY_RETURN(visitor->WithDefRegionKind(
kTVMFFIDefRegionKindSimple, [&]() { return visitor->VisitExpected(self->var); }));
TVM_FFI_S_VISIT_MAYBE_EARLY_RETURN(visitor->VisitExpected(self->value));
- return ffi::AnyView(nullptr).CopyToTVMFFIAny();
+ return std::nullopt;
}
-TVMFFIAny BindMutate(ffi::StructuralMutatorObj* mutator, ffi::AnyView value) noexcept {
+TVM_FFI_INLINE ffi::Expected<ffi::UnchangedOr<ffi::Any>> BindMutate(
+ ffi::StructuralMutatorObj* mutator, ffi::AnyView value) noexcept {
const BindNode* self =
ffi::details::AnyUnsafe::RawObjectPtrFromAnyViewAfterCheck<const BindNode>(value);
+ Type old_value_type = self->value->ty;
+ TVM_FFI_S_MUTATE_ASSIGN_OR_RETURN(ffi::UnchangedOr<Expr>, mapped_value,
+ mutator->MutateExpected(self->value));
+ bool value_changed = !mapped_value.UnchangedOrSameAs(self->value);
+ Expr new_value = std::move(mapped_value).ValueOrUnchanged(self->value);
+ value_changed |= !old_value_type.same_as(new_value->ty);
+ Type old_var_type = self->var->ty;
TVM_FFI_S_MUTATE_ASSIGN_OR_RETURN(ffi::UnchangedOr<Var>, mapped_var,
mutator->WithDefRegionKind(kTVMFFIDefRegionKindSimple, [&]() {
return mutator->MutateExpected(self->var);
}));
- TVM_FFI_S_MUTATE_ASSIGN_OR_RETURN(ffi::UnchangedOr<Expr>, mapped_value,
- mutator->MutateExpected(self->value));
- if (mapped_var.UnchangedOrSameAs(self->var) && mapped_value.UnchangedOrSameAs(self->value)) {
- return ffi::Unchanged().CopyToTVMFFIAny();
+ bool var_changed = !mapped_var.UnchangedOrSameAs(self->var);
+ Var new_var = std::move(mapped_var).ValueOrUnchanged(self->var);
+ var_changed |= !old_var_type.same_as(new_var->ty);
+ if ((value_changed || var_changed) && !new_var->ty.same_as(new_value->ty)) {
+ new_var = new_var.CopyWithType(new_value->ty);
+ }
+ if (!new_var.same_as(self->var)) {
+ auto remap = mutator->VarRemapSetExpected(self->var, new_var);
+ TVM_FFI_S_MUTATE_MAYBE_EARLY_RETURN(remap);
+ }
+ if (new_var.same_as(self->var) && new_value.same_as(self->value)) {
+ return ffi::Unchanged();
}
ffi::ObjectPtr<BindNode> copy = ffi::make_object<BindNode>(*self);
- copy->var = std::move(mapped_var).ValueOrUnchanged(std::move(copy->var));
- copy->value = std::move(mapped_value).ValueOrUnchanged(std::move(copy->value));
- return ffi::details::AnyUnsafe::MoveAnyToTVMFFIAny(ffi::Any(std::move(copy)));
+ copy->var = std::move(new_var);
+ copy->value = std::move(new_value);
+ return ffi::Any(std::move(copy));
}
-TVMFFIAny BindMaybeInplaceMutate(ffi::StructuralMutatorObj* mutator, ffi::AnyView value) noexcept {
+TVM_FFI_INLINE ffi::Expected<ffi::UnchangedOr<ffi::Any>> BindMaybeInplaceMutate(
+ ffi::StructuralMutatorObj* mutator, ffi::AnyView value) noexcept {
BindNode* self = const_cast<BindNode*>(
ffi::details::AnyUnsafe::RawObjectPtrFromAnyViewAfterCheck<const BindNode>(value));
+ Type old_value_type = self->value->ty;
+ TVM_FFI_S_MUTATE_ASSIGN_OR_RETURN(ffi::UnchangedOr<Expr>, mapped_value,
+ mutator->MutateExpected(self->value, ffi::InplaceMode::kAllow));
+ bool value_changed = !mapped_value.UnchangedOrSameAs(self->value);
+ Expr new_value = std::move(mapped_value).ValueOrUnchanged(self->value);
+ value_changed |= !old_value_type.same_as(new_value->ty);
+ Type old_var_type = self->var->ty;
TVM_FFI_S_MUTATE_ASSIGN_OR_RETURN(ffi::UnchangedOr<Var>, mapped_var,
mutator->WithDefRegionKind(kTVMFFIDefRegionKindSimple, [&]() {
return mutator->MutateExpected(self->var,
ffi::InplaceMode::kAllow);
}));
- TVM_FFI_S_MUTATE_ASSIGN_OR_RETURN(ffi::UnchangedOr<Expr>, mapped_value,
- mutator->MutateExpected(self->value, ffi::InplaceMode::kAllow));
- if (mapped_var.UnchangedOrSameAs(self->var) && mapped_value.UnchangedOrSameAs(self->value)) {
- return ffi::Unchanged().CopyToTVMFFIAny();
+ bool var_changed = !mapped_var.UnchangedOrSameAs(self->var);
+ Var new_var = std::move(mapped_var).ValueOrUnchanged(self->var);
+ var_changed |= !old_var_type.same_as(new_var->ty);
+ if ((value_changed || var_changed) && !new_var->ty.same_as(new_value->ty)) {
+ new_var = new_var.CopyWithType(new_value->ty);
}
- if (!mapped_var.IsUnchanged()) self->var = std::move(mapped_var).ValueUnchecked();
- if (!mapped_value.IsUnchanged()) self->value = std::move(mapped_value).ValueUnchecked();
- return ffi::Unchanged().CopyToTVMFFIAny();
+ if (!new_var.same_as(self->var)) {
+ auto remap = mutator->VarRemapSetExpected(self->var, new_var);
+ TVM_FFI_S_MUTATE_MAYBE_EARLY_RETURN(remap);
+ }
+ if (new_var.same_as(self->var) && new_value.same_as(self->value)) {
+ return ffi::Unchanged();
+ }
+ self->var = std::move(new_var);
+ self->value = std::move(new_value);
+ return ffi::Unchanged();
}
-TVMFFIAny AttrStmtVisit(ffi::StructuralVisitorObj* visitor, ffi::AnyView value) noexcept {
+TVM_FFI_INLINE ffi::Expected<ffi::Optional<ffi::VisitInterrupt>> AttrStmtVisit(
+ ffi::StructuralVisitorObj* visitor, ffi::AnyView value) noexcept {
// skips: attr_key
const AttrStmtNode* self =
ffi::details::AnyUnsafe::RawObjectPtrFromAnyViewAfterCheck<const AttrStmtNode>(value);
TVM_FFI_S_VISIT_MAYBE_EARLY_RETURN(visitor->VisitExpected(self->node));
TVM_FFI_S_VISIT_MAYBE_EARLY_RETURN(visitor->VisitExpected(self->value));
TVM_FFI_S_VISIT_MAYBE_EARLY_RETURN(visitor->VisitExpected(self->body));
- return ffi::AnyView(nullptr).CopyToTVMFFIAny();
+ return std::nullopt;
}
-TVMFFIAny AttrStmtMutate(ffi::StructuralMutatorObj* mutator, ffi::AnyView value) noexcept {
+TVM_FFI_INLINE ffi::Expected<ffi::UnchangedOr<ffi::Any>> AttrStmtMutate(
+ ffi::StructuralMutatorObj* mutator, ffi::AnyView value) noexcept {
// skips: attr_key
const AttrStmtNode* self =
ffi::details::AnyUnsafe::RawObjectPtrFromAnyViewAfterCheck<const AttrStmtNode>(value);
@@ -136,17 +171,17 @@ TVMFFIAny AttrStmtMutate(ffi::StructuralMutatorObj* mutator, ffi::AnyView value)
mutator->MutateExpected(self->body));
if (mapped_node.UnchangedOrSameAs(self->node) && mapped_value.UnchangedOrSameAs(self->value) &&
mapped_body.UnchangedOrSameAs(self->body)) {
- return ffi::Unchanged().CopyToTVMFFIAny();
+ return ffi::Unchanged();
}
ffi::ObjectPtr<AttrStmtNode> copy = ffi::make_object<AttrStmtNode>(*self);
copy->node = std::move(mapped_node).ValueOrUnchanged(std::move(copy->node));
copy->value = std::move(mapped_value).ValueOrUnchanged(std::move(copy->value));
copy->body = std::move(mapped_body).ValueOrUnchanged(std::move(copy->body));
- return ffi::details::AnyUnsafe::MoveAnyToTVMFFIAny(ffi::Any(std::move(copy)));
+ return ffi::Any(std::move(copy));
}
-TVMFFIAny AttrStmtMaybeInplaceMutate(ffi::StructuralMutatorObj* mutator,
- ffi::AnyView value) noexcept {
+TVM_FFI_INLINE ffi::Expected<ffi::UnchangedOr<ffi::Any>> AttrStmtMaybeInplaceMutate(
+ ffi::StructuralMutatorObj* mutator, ffi::AnyView value) noexcept {
// skips: attr_key
AttrStmtNode* self = const_cast<AttrStmtNode*>(
ffi::details::AnyUnsafe::RawObjectPtrFromAnyViewAfterCheck<const AttrStmtNode>(value));
@@ -158,38 +193,40 @@ TVMFFIAny AttrStmtMaybeInplaceMutate(ffi::StructuralMutatorObj* mutator,
mutator->MutateExpected(self->body, ffi::InplaceMode::kAllow));
if (mapped_node.UnchangedOrSameAs(self->node) && mapped_value.UnchangedOrSameAs(self->value) &&
mapped_body.UnchangedOrSameAs(self->body)) {
- return ffi::Unchanged().CopyToTVMFFIAny();
+ return ffi::Unchanged();
}
if (!mapped_node.IsUnchanged()) self->node = std::move(mapped_node).ValueUnchecked();
if (!mapped_value.IsUnchanged()) self->value = std::move(mapped_value).ValueUnchecked();
if (!mapped_body.IsUnchanged()) self->body = std::move(mapped_body).ValueUnchecked();
- return ffi::Unchanged().CopyToTVMFFIAny();
+ return ffi::Unchanged();
}
-TVMFFIAny AssertStmtVisit(ffi::StructuralVisitorObj* visitor, ffi::AnyView value) noexcept {
+TVM_FFI_INLINE ffi::Expected<ffi::Optional<ffi::VisitInterrupt>> AssertStmtVisit(
+ ffi::StructuralVisitorObj* visitor, ffi::AnyView value) noexcept {
// skips: error_kind and message_parts, which are constant assertion metadata.
const AssertStmtNode* self =
ffi::details::AnyUnsafe::RawObjectPtrFromAnyViewAfterCheck<const AssertStmtNode>(value);
TVM_FFI_S_VISIT_MAYBE_EARLY_RETURN(visitor->VisitExpected(self->condition));
- return ffi::AnyView(nullptr).CopyToTVMFFIAny();
+ return std::nullopt;
}
-TVMFFIAny AssertStmtMutate(ffi::StructuralMutatorObj* mutator, ffi::AnyView value) noexcept {
+TVM_FFI_INLINE ffi::Expected<ffi::UnchangedOr<ffi::Any>> AssertStmtMutate(
+ ffi::StructuralMutatorObj* mutator, ffi::AnyView value) noexcept {
// skips: error_kind and message_parts, which are constant assertion metadata.
const AssertStmtNode* self =
ffi::details::AnyUnsafe::RawObjectPtrFromAnyViewAfterCheck<const AssertStmtNode>(value);
TVM_FFI_S_MUTATE_ASSIGN_OR_RETURN(ffi::UnchangedOr<PrimExpr>, mapped_condition,
mutator->MutateExpected(self->condition));
if (mapped_condition.UnchangedOrSameAs(self->condition)) {
- return ffi::Unchanged().CopyToTVMFFIAny();
+ return ffi::Unchanged();
}
ffi::ObjectPtr<AssertStmtNode> copy = ffi::make_object<AssertStmtNode>(*self);
copy->condition = std::move(mapped_condition).ValueOrUnchanged(std::move(copy->condition));
- return ffi::details::AnyUnsafe::MoveAnyToTVMFFIAny(ffi::Any(std::move(copy)));
+ return ffi::Any(std::move(copy));
}
-TVMFFIAny AssertStmtMaybeInplaceMutate(ffi::StructuralMutatorObj* mutator,
- ffi::AnyView value) noexcept {
+TVM_FFI_INLINE ffi::Expected<ffi::UnchangedOr<ffi::Any>> AssertStmtMaybeInplaceMutate(
+ ffi::StructuralMutatorObj* mutator, ffi::AnyView value) noexcept {
// skips: error_kind and message_parts, which are constant assertion metadata.
AssertStmtNode* self = const_cast<AssertStmtNode*>(
ffi::details::AnyUnsafe::RawObjectPtrFromAnyViewAfterCheck<const AssertStmtNode>(value));
@@ -197,14 +234,15 @@ TVMFFIAny AssertStmtMaybeInplaceMutate(ffi::StructuralMutatorObj* mutator,
ffi::UnchangedOr<PrimExpr>, mapped_condition,
mutator->MutateExpected(self->condition, ffi::InplaceMode::kAllow));
if (mapped_condition.UnchangedOrSameAs(self->condition)) {
- return ffi::Unchanged().CopyToTVMFFIAny();
+ return ffi::Unchanged();
}
if (!mapped_condition.IsUnchanged())
self->condition = std::move(mapped_condition).ValueUnchecked();
- return ffi::Unchanged().CopyToTVMFFIAny();
+ return ffi::Unchanged();
}
-TVMFFIAny ForVisit(ffi::StructuralVisitorObj* visitor, ffi::AnyView value) noexcept {
+TVM_FFI_INLINE ffi::Expected<ffi::Optional<ffi::VisitInterrupt>> ForVisit(
+ ffi::StructuralVisitorObj* visitor, ffi::AnyView value) noexcept {
// skips: kind and constant annotations; unlike SBlock annotations, these carry no expressions.
const ForNode* self =
ffi::details::AnyUnsafe::RawObjectPtrFromAnyViewAfterCheck<const ForNode>(value);
@@ -215,10 +253,11 @@ TVMFFIAny ForVisit(ffi::StructuralVisitorObj* visitor, ffi::AnyView value) noexc
TVM_FFI_S_VISIT_MAYBE_EARLY_RETURN(visitor->VisitExpected(self->body));
TVM_FFI_S_VISIT_MAYBE_EARLY_RETURN(visitor->VisitExpected(self->thread_binding));
TVM_FFI_S_VISIT_MAYBE_EARLY_RETURN(visitor->VisitExpected(self->step));
- return ffi::AnyView(nullptr).CopyToTVMFFIAny();
+ return std::nullopt;
}
-TVMFFIAny ForMutate(ffi::StructuralMutatorObj* mutator, ffi::AnyView value) noexcept {
+TVM_FFI_INLINE ffi::Expected<ffi::UnchangedOr<ffi::Any>> ForMutate(
+ ffi::StructuralMutatorObj* mutator, ffi::AnyView value) noexcept {
// skips: kind and constant annotations; unlike SBlock annotations, these carry no expressions.
const ForNode* self =
ffi::details::AnyUnsafe::RawObjectPtrFromAnyViewAfterCheck<const ForNode>(value);
@@ -241,7 +280,7 @@ TVMFFIAny ForMutate(ffi::StructuralMutatorObj* mutator, ffi::AnyView value) noex
mapped_body.UnchangedOrSameAs(self->body) &&
mapped_thread_binding.UnchangedOrSameAs(self->thread_binding) &&
mapped_step.UnchangedOrSameAs(self->step)) {
- return ffi::Unchanged().CopyToTVMFFIAny();
+ return ffi::Unchanged();
}
ffi::ObjectPtr<ForNode> copy = ffi::make_object<ForNode>(*self);
copy->loop_var = std::move(mapped_loop_var).ValueOrUnchanged(std::move(copy->loop_var));
@@ -251,10 +290,11 @@ TVMFFIAny ForMutate(ffi::StructuralMutatorObj* mutator, ffi::AnyView value) noex
copy->thread_binding =
std::move(mapped_thread_binding).ValueOrUnchanged(std::move(copy->thread_binding));
copy->step = std::move(mapped_step).ValueOrUnchanged(std::move(copy->step));
- return ffi::details::AnyUnsafe::MoveAnyToTVMFFIAny(ffi::Any(std::move(copy)));
+ return ffi::Any(std::move(copy));
}
-TVMFFIAny ForMaybeInplaceMutate(ffi::StructuralMutatorObj* mutator, ffi::AnyView value) noexcept {
+TVM_FFI_INLINE ffi::Expected<ffi::UnchangedOr<ffi::Any>> ForMaybeInplaceMutate(
+ ffi::StructuralMutatorObj* mutator, ffi::AnyView value) noexcept {
// skips: kind and constant annotations; unlike SBlock annotations, these carry no expressions.
ForNode* self = const_cast<ForNode*>(
ffi::details::AnyUnsafe::RawObjectPtrFromAnyViewAfterCheck<const ForNode>(value));
@@ -280,7 +320,7 @@ TVMFFIAny ForMaybeInplaceMutate(ffi::StructuralMutatorObj* mutator, ffi::AnyView
mapped_body.UnchangedOrSameAs(self->body) &&
mapped_thread_binding.UnchangedOrSameAs(self->thread_binding) &&
mapped_step.UnchangedOrSameAs(self->step)) {
- return ffi::Unchanged().CopyToTVMFFIAny();
+ return ffi::Unchanged();
}
if (!mapped_loop_var.IsUnchanged()) self->loop_var = std::move(mapped_loop_var).ValueUnchecked();
if (!mapped_min.IsUnchanged()) self->min = std::move(mapped_min).ValueUnchecked();
@@ -290,18 +330,20 @@ TVMFFIAny ForMaybeInplaceMutate(ffi::StructuralMutatorObj* mutator, ffi::AnyView
self->thread_binding = std::move(mapped_thread_binding).ValueUnchecked();
}
if (!mapped_step.IsUnchanged()) self->step = std::move(mapped_step).ValueUnchecked();
- return ffi::Unchanged().CopyToTVMFFIAny();
+ return ffi::Unchanged();
}
-TVMFFIAny WhileVisit(ffi::StructuralVisitorObj* visitor, ffi::AnyView value) noexcept {
+TVM_FFI_INLINE ffi::Expected<ffi::Optional<ffi::VisitInterrupt>> WhileVisit(
+ ffi::StructuralVisitorObj* visitor, ffi::AnyView value) noexcept {
const WhileNode* self =
ffi::details::AnyUnsafe::RawObjectPtrFromAnyViewAfterCheck<const WhileNode>(value);
TVM_FFI_S_VISIT_MAYBE_EARLY_RETURN(visitor->VisitExpected(self->condition));
TVM_FFI_S_VISIT_MAYBE_EARLY_RETURN(visitor->VisitExpected(self->body));
- return ffi::AnyView(nullptr).CopyToTVMFFIAny();
+ return std::nullopt;
}
-TVMFFIAny WhileMutate(ffi::StructuralMutatorObj* mutator, ffi::AnyView value) noexcept {
+TVM_FFI_INLINE ffi::Expected<ffi::UnchangedOr<ffi::Any>> WhileMutate(
+ ffi::StructuralMutatorObj* mutator, ffi::AnyView value) noexcept {
const WhileNode* self =
ffi::details::AnyUnsafe::RawObjectPtrFromAnyViewAfterCheck<const WhileNode>(value);
TVM_FFI_S_MUTATE_ASSIGN_OR_RETURN(ffi::UnchangedOr<PrimExpr>, mapped_condition,
@@ -310,15 +352,16 @@ TVMFFIAny WhileMutate(ffi::StructuralMutatorObj* mutator, ffi::AnyView value) no
mutator->MutateExpected(self->body));
if (mapped_condition.UnchangedOrSameAs(self->condition) &&
mapped_body.UnchangedOrSameAs(self->body)) {
- return ffi::Unchanged().CopyToTVMFFIAny();
+ return ffi::Unchanged();
}
ffi::ObjectPtr<WhileNode> copy = ffi::make_object<WhileNode>(*self);
copy->condition = std::move(mapped_condition).ValueOrUnchanged(std::move(copy->condition));
copy->body = std::move(mapped_body).ValueOrUnchanged(std::move(copy->body));
- return ffi::details::AnyUnsafe::MoveAnyToTVMFFIAny(ffi::Any(std::move(copy)));
+ return ffi::Any(std::move(copy));
}
-TVMFFIAny WhileMaybeInplaceMutate(ffi::StructuralMutatorObj* mutator, ffi::AnyView value) noexcept {
+TVM_FFI_INLINE ffi::Expected<ffi::UnchangedOr<ffi::Any>> WhileMaybeInplaceMutate(
+ ffi::StructuralMutatorObj* mutator, ffi::AnyView value) noexcept {
WhileNode* self = const_cast<WhileNode*>(
ffi::details::AnyUnsafe::RawObjectPtrFromAnyViewAfterCheck<const WhileNode>(value));
TVM_FFI_S_MUTATE_ASSIGN_OR_RETURN(
@@ -328,81 +371,91 @@ TVMFFIAny WhileMaybeInplaceMutate(ffi::StructuralMutatorObj* mutator, ffi::AnyVi
mutator->MutateExpected(self->body, ffi::InplaceMode::kAllow));
if (mapped_condition.UnchangedOrSameAs(self->condition) &&
mapped_body.UnchangedOrSameAs(self->body)) {
- return ffi::Unchanged().CopyToTVMFFIAny();
+ return ffi::Unchanged();
}
if (!mapped_condition.IsUnchanged())
self->condition = std::move(mapped_condition).ValueUnchecked();
if (!mapped_body.IsUnchanged()) self->body = std::move(mapped_body).ValueUnchecked();
- return ffi::Unchanged().CopyToTVMFFIAny();
+ return ffi::Unchanged();
}
-TVMFFIAny ReturnVisit(ffi::StructuralVisitorObj* visitor, ffi::AnyView value) noexcept {
+TVM_FFI_INLINE ffi::Expected<ffi::Optional<ffi::VisitInterrupt>> ReturnVisit(
+ ffi::StructuralVisitorObj* visitor, ffi::AnyView value) noexcept {
const ReturnNode* self =
ffi::details::AnyUnsafe::RawObjectPtrFromAnyViewAfterCheck<const ReturnNode>(value);
TVM_FFI_S_VISIT_MAYBE_EARLY_RETURN(visitor->VisitExpected(self->value));
- return ffi::AnyView(nullptr).CopyToTVMFFIAny();
+ return std::nullopt;
}
-TVMFFIAny ReturnMutate(ffi::StructuralMutatorObj* mutator, ffi::AnyView value) noexcept {
+TVM_FFI_INLINE ffi::Expected<ffi::UnchangedOr<ffi::Any>> ReturnMutate(
+ ffi::StructuralMutatorObj* mutator, ffi::AnyView value) noexcept {
const ReturnNode* self =
ffi::details::AnyUnsafe::RawObjectPtrFromAnyViewAfterCheck<const ReturnNode>(value);
TVM_FFI_S_MUTATE_ASSIGN_OR_RETURN(ffi::UnchangedOr<Expr>, mapped_value,
mutator->MutateExpected(self->value));
if (mapped_value.UnchangedOrSameAs(self->value)) {
- return ffi::Unchanged().CopyToTVMFFIAny();
+ return ffi::Unchanged();
}
ffi::ObjectPtr<ReturnNode> copy = ffi::make_object<ReturnNode>(*self);
copy->value = std::move(mapped_value).ValueOrUnchanged(std::move(copy->value));
- return ffi::details::AnyUnsafe::MoveAnyToTVMFFIAny(ffi::Any(std::move(copy)));
+ return ffi::Any(std::move(copy));
}
-TVMFFIAny ReturnMaybeInplaceMutate(ffi::StructuralMutatorObj* mutator,
- ffi::AnyView value) noexcept {
+TVM_FFI_INLINE ffi::Expected<ffi::UnchangedOr<ffi::Any>> ReturnMaybeInplaceMutate(
+ ffi::StructuralMutatorObj* mutator, ffi::AnyView value) noexcept {
ReturnNode* self = const_cast<ReturnNode*>(
ffi::details::AnyUnsafe::RawObjectPtrFromAnyViewAfterCheck<const ReturnNode>(value));
TVM_FFI_S_MUTATE_ASSIGN_OR_RETURN(ffi::UnchangedOr<Expr>, mapped_value,
mutator->MutateExpected(self->value, ffi::InplaceMode::kAllow));
if (mapped_value.UnchangedOrSameAs(self->value)) {
- return ffi::Unchanged().CopyToTVMFFIAny();
+ return ffi::Unchanged();
}
if (!mapped_value.IsUnchanged()) self->value = std::move(mapped_value).ValueUnchecked();
- return ffi::Unchanged().CopyToTVMFFIAny();
+ return ffi::Unchanged();
}
-TVMFFIAny BreakVisit(ffi::StructuralVisitorObj*, ffi::AnyView) noexcept {
- return ffi::AnyView(nullptr).CopyToTVMFFIAny();
+TVM_FFI_INLINE ffi::Expected<ffi::Optional<ffi::VisitInterrupt>> BreakVisit(
+ ffi::StructuralVisitorObj*, ffi::AnyView) noexcept {
+ return std::nullopt;
}
-TVMFFIAny BreakMutate(ffi::StructuralMutatorObj*, ffi::AnyView) noexcept {
- return ffi::Unchanged().CopyToTVMFFIAny();
+TVM_FFI_INLINE ffi::Expected<ffi::UnchangedOr<ffi::Any>> BreakMutate(ffi::StructuralMutatorObj*,
+ ffi::AnyView) noexcept {
+ return ffi::Unchanged();
}
-TVMFFIAny BreakMaybeInplaceMutate(ffi::StructuralMutatorObj*, ffi::AnyView) noexcept {
- return ffi::Unchanged().CopyToTVMFFIAny();
+TVM_FFI_INLINE ffi::Expected<ffi::UnchangedOr<ffi::Any>> BreakMaybeInplaceMutate(
+ ffi::StructuralMutatorObj*, ffi::AnyView) noexcept {
+ return ffi::Unchanged();
}
-TVMFFIAny ContinueVisit(ffi::StructuralVisitorObj*, ffi::AnyView) noexcept {
- return ffi::AnyView(nullptr).CopyToTVMFFIAny();
+TVM_FFI_INLINE ffi::Expected<ffi::Optional<ffi::VisitInterrupt>> ContinueVisit(
+ ffi::StructuralVisitorObj*, ffi::AnyView) noexcept {
+ return std::nullopt;
}
-TVMFFIAny ContinueMutate(ffi::StructuralMutatorObj*, ffi::AnyView) noexcept {
- return ffi::Unchanged().CopyToTVMFFIAny();
+TVM_FFI_INLINE ffi::Expected<ffi::UnchangedOr<ffi::Any>> ContinueMutate(ffi::StructuralMutatorObj*,
+ ffi::AnyView) noexcept {
+ return ffi::Unchanged();
}
-TVMFFIAny ContinueMaybeInplaceMutate(ffi::StructuralMutatorObj*, ffi::AnyView) noexcept {
- return ffi::Unchanged().CopyToTVMFFIAny();
+TVM_FFI_INLINE ffi::Expected<ffi::UnchangedOr<ffi::Any>> ContinueMaybeInplaceMutate(
+ ffi::StructuralMutatorObj*, ffi::AnyView) noexcept {
+ return ffi::Unchanged();
}
-TVMFFIAny DeclBufferVisit(ffi::StructuralVisitorObj* visitor, ffi::AnyView value) noexcept {
+TVM_FFI_INLINE ffi::Expected<ffi::Optional<ffi::VisitInterrupt>> DeclBufferVisit(
+ ffi::StructuralVisitorObj* visitor, ffi::AnyView value) noexcept {
const DeclBufferNode* self =
ffi::details::AnyUnsafe::RawObjectPtrFromAnyViewAfterCheck<const DeclBufferNode>(value);
TVM_FFI_S_VISIT_MAYBE_EARLY_RETURN(visitor->WithDefRegionKind(
kTVMFFIDefRegionKindSimple, [&]() { return visitor->VisitExpected(self->buffer); }));
TVM_FFI_S_VISIT_MAYBE_EARLY_RETURN(visitor->VisitExpected(self->data));
- return ffi::AnyView(nullptr).CopyToTVMFFIAny();
+ return std::nullopt;
}
-TVMFFIAny DeclBufferMutate(ffi::StructuralMutatorObj* mutator, ffi::AnyView value) noexcept {
+TVM_FFI_INLINE ffi::Expected<ffi::UnchangedOr<ffi::Any>> DeclBufferMutate(
+ ffi::StructuralMutatorObj* mutator, ffi::AnyView value) noexcept {
const DeclBufferNode* self =
ffi::details::AnyUnsafe::RawObjectPtrFromAnyViewAfterCheck<const DeclBufferNode>(value);
TVM_FFI_S_MUTATE_ASSIGN_OR_RETURN(ffi::UnchangedOr<BufferVar>, mapped_buffer,
@@ -412,16 +465,16 @@ TVMFFIAny DeclBufferMutate(ffi::StructuralMutatorObj* mutator, ffi::AnyView valu
TVM_FFI_S_MUTATE_ASSIGN_OR_RETURN(ffi::UnchangedOr<Expr>, mapped_data,
mutator->MutateExpected(self->data));
if (mapped_buffer.UnchangedOrSameAs(self->buffer) && mapped_data.UnchangedOrSameAs(self->data)) {
- return ffi::Unchanged().CopyToTVMFFIAny();
+ return ffi::Unchanged();
}
ffi::ObjectPtr<DeclBufferNode> copy = ffi::make_object<DeclBufferNode>(*self);
copy->buffer = std::move(mapped_buffer).ValueOrUnchanged(std::move(copy->buffer));
copy->data = std::move(mapped_data).ValueOrUnchanged(std::move(copy->data));
- return ffi::details::AnyUnsafe::MoveAnyToTVMFFIAny(ffi::Any(std::move(copy)));
+ return ffi::Any(std::move(copy));
}
-TVMFFIAny DeclBufferMaybeInplaceMutate(ffi::StructuralMutatorObj* mutator,
- ffi::AnyView value) noexcept {
+TVM_FFI_INLINE ffi::Expected<ffi::UnchangedOr<ffi::Any>> DeclBufferMaybeInplaceMutate(
+ ffi::StructuralMutatorObj* mutator, ffi::AnyView value) noexcept {
DeclBufferNode* self = const_cast<DeclBufferNode*>(
ffi::details::AnyUnsafe::RawObjectPtrFromAnyViewAfterCheck<const DeclBufferNode>(value));
TVM_FFI_S_MUTATE_ASSIGN_OR_RETURN(ffi::UnchangedOr<BufferVar>, mapped_buffer,
@@ -432,23 +485,25 @@ TVMFFIAny DeclBufferMaybeInplaceMutate(ffi::StructuralMutatorObj* mutator,
TVM_FFI_S_MUTATE_ASSIGN_OR_RETURN(ffi::UnchangedOr<Expr>, mapped_data,
mutator->MutateExpected(self->data, ffi::InplaceMode::kAllow));
if (mapped_buffer.UnchangedOrSameAs(self->buffer) && mapped_data.UnchangedOrSameAs(self->data)) {
- return ffi::Unchanged().CopyToTVMFFIAny();
+ return ffi::Unchanged();
}
if (!mapped_buffer.IsUnchanged()) self->buffer = std::move(mapped_buffer).ValueUnchecked();
if (!mapped_data.IsUnchanged()) self->data = std::move(mapped_data).ValueUnchecked();
- return ffi::Unchanged().CopyToTVMFFIAny();
+ return ffi::Unchanged();
}
-TVMFFIAny AllocBufferVisit(ffi::StructuralVisitorObj* visitor, ffi::AnyView value) noexcept {
+TVM_FFI_INLINE ffi::Expected<ffi::Optional<ffi::VisitInterrupt>> AllocBufferVisit(
+ ffi::StructuralVisitorObj* visitor, ffi::AnyView value) noexcept {
// skips: constant annotations; unlike SBlock annotations, these carry no expressions.
const AllocBufferNode* self =
ffi::details::AnyUnsafe::RawObjectPtrFromAnyViewAfterCheck<const AllocBufferNode>(value);
TVM_FFI_S_VISIT_MAYBE_EARLY_RETURN(visitor->WithDefRegionKind(
kTVMFFIDefRegionKindSimple, [&]() { return visitor->VisitExpected(self->buffer); }));
- return ffi::AnyView(nullptr).CopyToTVMFFIAny();
+ return std::nullopt;
}
-TVMFFIAny AllocBufferMutate(ffi::StructuralMutatorObj* mutator, ffi::AnyView value) noexcept {
+TVM_FFI_INLINE ffi::Expected<ffi::UnchangedOr<ffi::Any>> AllocBufferMutate(
+ ffi::StructuralMutatorObj* mutator, ffi::AnyView value) noexcept {
// skips: constant annotations; unlike SBlock annotations, these carry no expressions.
const AllocBufferNode* self =
ffi::details::AnyUnsafe::RawObjectPtrFromAnyViewAfterCheck<const AllocBufferNode>(value);
@@ -457,15 +512,15 @@ TVMFFIAny AllocBufferMutate(ffi::StructuralMutatorObj* mutator, ffi::AnyView val
return mutator->MutateExpected(self->buffer);
}));
if (mapped_buffer.UnchangedOrSameAs(self->buffer)) {
- return ffi::Unchanged().CopyToTVMFFIAny();
+ return ffi::Unchanged();
}
ffi::ObjectPtr<AllocBufferNode> copy = ffi::make_object<AllocBufferNode>(*self);
copy->buffer = std::move(mapped_buffer).ValueOrUnchanged(std::move(copy->buffer));
- return ffi::details::AnyUnsafe::MoveAnyToTVMFFIAny(ffi::Any(std::move(copy)));
+ return ffi::Any(std::move(copy));
}
-TVMFFIAny AllocBufferMaybeInplaceMutate(ffi::StructuralMutatorObj* mutator,
- ffi::AnyView value) noexcept {
+TVM_FFI_INLINE ffi::Expected<ffi::UnchangedOr<ffi::Any>> AllocBufferMaybeInplaceMutate(
+ ffi::StructuralMutatorObj* mutator, ffi::AnyView value) noexcept {
// skips: constant annotations; unlike SBlock annotations, these carry no expressions.
AllocBufferNode* self = const_cast<AllocBufferNode*>(
ffi::details::AnyUnsafe::RawObjectPtrFromAnyViewAfterCheck<const AllocBufferNode>(value));
@@ -475,21 +530,23 @@ TVMFFIAny AllocBufferMaybeInplaceMutate(ffi::StructuralMutatorObj* mutator,
ffi::InplaceMode::kAllow);
}));
if (mapped_buffer.UnchangedOrSameAs(self->buffer)) {
- return ffi::Unchanged().CopyToTVMFFIAny();
+ return ffi::Unchanged();
}
if (!mapped_buffer.IsUnchanged()) self->buffer = std::move(mapped_buffer).ValueUnchecked();
- return ffi::Unchanged().CopyToTVMFFIAny();
+ return ffi::Unchanged();
}
-TVMFFIAny SeqStmtVisit(ffi::StructuralVisitorObj* visitor, ffi::AnyView value) noexcept {
+TVM_FFI_INLINE ffi::Expected<ffi::Optional<ffi::VisitInterrupt>> SeqStmtVisit(
+ ffi::StructuralVisitorObj* visitor, ffi::AnyView value) noexcept {
const SeqStmtNode* self =
ffi::details::AnyUnsafe::RawObjectPtrFromAnyViewAfterCheck<const SeqStmtNode>(value);
TVM_FFI_S_VISIT_MAYBE_EARLY_RETURN(visitor->VisitExpected(self->seq));
- return ffi::AnyView(nullptr).CopyToTVMFFIAny();
+ return std::nullopt;
}
-TVMFFIAny MutateSeqStmtStructural(ffi::StructuralMutatorObj* mutator, ffi::AnyView value,
- ffi::InplaceMode inplace_mode) noexcept {
+TVM_FFI_INLINE ffi::Expected<ffi::UnchangedOr<ffi::Any>> MutateSeqStmtStructural(
+ ffi::StructuralMutatorObj* mutator, ffi::AnyView value,
+ ffi::InplaceMode inplace_mode) noexcept {
const auto* self =
ffi::details::AnyUnsafe::RawObjectPtrFromAnyViewAfterCheck<const SeqStmtNode>(value);
auto mutate = [mutator](ffi::AnyView element,
@@ -498,29 +555,31 @@ TVMFFIAny MutateSeqStmtStructural(ffi::StructuralMutatorObj* mutator, ffi::AnyVi
mutator->MutateExpected(element, mode));
return mapped;
};
- return ffi::details::ExpectedUnsafe::MoveToTVMFFIAny(
- detail::MutateSeqStmt(self, inplace_mode, mutate));
+ return detail::MutateSeqStmt(self, inplace_mode, mutate);
}
-TVMFFIAny SeqStmtMutate(ffi::StructuralMutatorObj* mutator, ffi::AnyView value) noexcept {
+TVM_FFI_INLINE ffi::Expected<ffi::UnchangedOr<ffi::Any>> SeqStmtMutate(
+ ffi::StructuralMutatorObj* mutator, ffi::AnyView value) noexcept {
return MutateSeqStmtStructural(mutator, value, ffi::InplaceMode::kDisallow);
}
-TVMFFIAny SeqStmtMaybeInplaceMutate(ffi::StructuralMutatorObj* mutator,
- ffi::AnyView value) noexcept {
+TVM_FFI_INLINE ffi::Expected<ffi::UnchangedOr<ffi::Any>> SeqStmtMaybeInplaceMutate(
+ ffi::StructuralMutatorObj* mutator, ffi::AnyView value) noexcept {
return MutateSeqStmtStructural(mutator, value, ffi::InplaceMode::kAllow);
}
-TVMFFIAny IfThenElseVisit(ffi::StructuralVisitorObj* visitor, ffi::AnyView value) noexcept {
+TVM_FFI_INLINE ffi::Expected<ffi::Optional<ffi::VisitInterrupt>> IfThenElseVisit(
+ ffi::StructuralVisitorObj* visitor, ffi::AnyView value) noexcept {
const IfThenElseNode* self =
ffi::details::AnyUnsafe::RawObjectPtrFromAnyViewAfterCheck<const IfThenElseNode>(value);
TVM_FFI_S_VISIT_MAYBE_EARLY_RETURN(visitor->VisitExpected(self->condition));
TVM_FFI_S_VISIT_MAYBE_EARLY_RETURN(visitor->VisitExpected(self->then_case));
TVM_FFI_S_VISIT_MAYBE_EARLY_RETURN(visitor->VisitExpected(self->else_case));
- return ffi::AnyView(nullptr).CopyToTVMFFIAny();
+ return std::nullopt;
}
-TVMFFIAny IfThenElseMutate(ffi::StructuralMutatorObj* mutator, ffi::AnyView value) noexcept {
+TVM_FFI_INLINE ffi::Expected<ffi::UnchangedOr<ffi::Any>> IfThenElseMutate(
+ ffi::StructuralMutatorObj* mutator, ffi::AnyView value) noexcept {
const IfThenElseNode* self =
ffi::details::AnyUnsafe::RawObjectPtrFromAnyViewAfterCheck<const IfThenElseNode>(value);
TVM_FFI_S_MUTATE_ASSIGN_OR_RETURN(ffi::UnchangedOr<PrimExpr>, mapped_condition,
@@ -532,17 +591,17 @@ TVMFFIAny IfThenElseMutate(ffi::StructuralMutatorObj* mutator, ffi::AnyView valu
if (mapped_condition.UnchangedOrSameAs(self->condition) &&
mapped_then_case.UnchangedOrSameAs(self->then_case) &&
mapped_else_case.UnchangedOrSameAs(self->else_case)) {
- return ffi::Unchanged().CopyToTVMFFIAny();
+ return ffi::Unchanged();
}
ffi::ObjectPtr<IfThenElseNode> copy = ffi::make_object<IfThenElseNode>(*self);
copy->condition = std::move(mapped_condition).ValueOrUnchanged(std::move(copy->condition));
copy->then_case = std::move(mapped_then_case).ValueOrUnchanged(std::move(copy->then_case));
copy->else_case = std::move(mapped_else_case).ValueOrUnchanged(std::move(copy->else_case));
- return ffi::details::AnyUnsafe::MoveAnyToTVMFFIAny(ffi::Any(std::move(copy)));
+ return ffi::Any(std::move(copy));
}
-TVMFFIAny IfThenElseMaybeInplaceMutate(ffi::StructuralMutatorObj* mutator,
- ffi::AnyView value) noexcept {
+TVM_FFI_INLINE ffi::Expected<ffi::UnchangedOr<ffi::Any>> IfThenElseMaybeInplaceMutate(
+ ffi::StructuralMutatorObj* mutator, ffi::AnyView value) noexcept {
IfThenElseNode* self = const_cast<IfThenElseNode*>(
ffi::details::AnyUnsafe::RawObjectPtrFromAnyViewAfterCheck<const IfThenElseNode>(value));
TVM_FFI_S_MUTATE_ASSIGN_OR_RETURN(
@@ -557,7 +616,7 @@ TVMFFIAny IfThenElseMaybeInplaceMutate(ffi::StructuralMutatorObj* mutator,
if (mapped_condition.UnchangedOrSameAs(self->condition) &&
mapped_then_case.UnchangedOrSameAs(self->then_case) &&
mapped_else_case.UnchangedOrSameAs(self->else_case)) {
- return ffi::Unchanged().CopyToTVMFFIAny();
+ return ffi::Unchanged();
}
if (!mapped_condition.IsUnchanged())
self->condition = std::move(mapped_condition).ValueUnchecked();
@@ -565,52 +624,56 @@ TVMFFIAny IfThenElseMaybeInplaceMutate(ffi::StructuralMutatorObj* mutator,
self->then_case = std::move(mapped_then_case).ValueUnchecked();
if (!mapped_else_case.IsUnchanged())
self->else_case = std::move(mapped_else_case).ValueUnchecked();
- return ffi::Unchanged().CopyToTVMFFIAny();
+ return ffi::Unchanged();
}
-TVMFFIAny EvaluateVisit(ffi::StructuralVisitorObj* visitor, ffi::AnyView value) noexcept {
+TVM_FFI_INLINE ffi::Expected<ffi::Optional<ffi::VisitInterrupt>> EvaluateVisit(
+ ffi::StructuralVisitorObj* visitor, ffi::AnyView value) noexcept {
const EvaluateNode* self =
ffi::details::AnyUnsafe::RawObjectPtrFromAnyViewAfterCheck<const EvaluateNode>(value);
TVM_FFI_S_VISIT_MAYBE_EARLY_RETURN(visitor->VisitExpected(self->value));
- return ffi::AnyView(nullptr).CopyToTVMFFIAny();
+ return std::nullopt;
}
-TVMFFIAny EvaluateMutate(ffi::StructuralMutatorObj* mutator, ffi::AnyView value) noexcept {
+TVM_FFI_INLINE ffi::Expected<ffi::UnchangedOr<ffi::Any>> EvaluateMutate(
+ ffi::StructuralMutatorObj* mutator, ffi::AnyView value) noexcept {
const EvaluateNode* self =
ffi::details::AnyUnsafe::RawObjectPtrFromAnyViewAfterCheck<const EvaluateNode>(value);
TVM_FFI_S_MUTATE_ASSIGN_OR_RETURN(ffi::UnchangedOr<Expr>, mapped_value,
mutator->MutateExpected(self->value));
if (mapped_value.UnchangedOrSameAs(self->value)) {
- return ffi::Unchanged().CopyToTVMFFIAny();
+ return ffi::Unchanged();
}
ffi::ObjectPtr<EvaluateNode> copy = ffi::make_object<EvaluateNode>(*self);
copy->value = std::move(mapped_value).ValueOrUnchanged(std::move(copy->value));
- return ffi::details::AnyUnsafe::MoveAnyToTVMFFIAny(ffi::Any(std::move(copy)));
+ return ffi::Any(std::move(copy));
}
-TVMFFIAny EvaluateMaybeInplaceMutate(ffi::StructuralMutatorObj* mutator,
- ffi::AnyView value) noexcept {
+TVM_FFI_INLINE ffi::Expected<ffi::UnchangedOr<ffi::Any>> EvaluateMaybeInplaceMutate(
+ ffi::StructuralMutatorObj* mutator, ffi::AnyView value) noexcept {
EvaluateNode* self = const_cast<EvaluateNode*>(
ffi::details::AnyUnsafe::RawObjectPtrFromAnyViewAfterCheck<const EvaluateNode>(value));
TVM_FFI_S_MUTATE_ASSIGN_OR_RETURN(ffi::UnchangedOr<Expr>, mapped_value,
mutator->MutateExpected(self->value, ffi::InplaceMode::kAllow));
if (mapped_value.UnchangedOrSameAs(self->value)) {
- return ffi::Unchanged().CopyToTVMFFIAny();
+ return ffi::Unchanged();
}
if (!mapped_value.IsUnchanged()) self->value = std::move(mapped_value).ValueUnchecked();
- return ffi::Unchanged().CopyToTVMFFIAny();
+ return ffi::Unchanged();
}
-TVMFFIAny BufferStoreVisit(ffi::StructuralVisitorObj* visitor, ffi::AnyView value) noexcept {
+TVM_FFI_INLINE ffi::Expected<ffi::Optional<ffi::VisitInterrupt>> BufferStoreVisit(
+ ffi::StructuralVisitorObj* visitor, ffi::AnyView value) noexcept {
const BufferStoreNode* self =
ffi::details::AnyUnsafe::RawObjectPtrFromAnyViewAfterCheck<const BufferStoreNode>(value);
TVM_FFI_S_VISIT_MAYBE_EARLY_RETURN(visitor->VisitExpected(self->buffer));
TVM_FFI_S_VISIT_MAYBE_EARLY_RETURN(visitor->VisitExpected(self->value));
TVM_FFI_S_VISIT_MAYBE_EARLY_RETURN(visitor->VisitExpected(self->indices));
- return ffi::AnyView(nullptr).CopyToTVMFFIAny();
+ return std::nullopt;
}
-TVMFFIAny BufferStoreMutate(ffi::StructuralMutatorObj* mutator, ffi::AnyView value) noexcept {
+TVM_FFI_INLINE ffi::Expected<ffi::UnchangedOr<ffi::Any>> BufferStoreMutate(
+ ffi::StructuralMutatorObj* mutator, ffi::AnyView value) noexcept {
const BufferStoreNode* self =
ffi::details::AnyUnsafe::RawObjectPtrFromAnyViewAfterCheck<const BufferStoreNode>(value);
TVM_FFI_S_MUTATE_ASSIGN_OR_RETURN(ffi::UnchangedOr<BufferVar>, mapped_buffer,
@@ -622,17 +685,17 @@ TVMFFIAny BufferStoreMutate(ffi::StructuralMutatorObj* mutator, ffi::AnyView val
if (mapped_buffer.UnchangedOrSameAs(self->buffer) &&
mapped_value.UnchangedOrSameAs(self->value) &&
mapped_indices.UnchangedOrSameAs(self->indices)) {
- return ffi::Unchanged().CopyToTVMFFIAny();
+ return ffi::Unchanged();
}
ffi::ObjectPtr<BufferStoreNode> copy = ffi::make_object<BufferStoreNode>(*self);
copy->buffer = std::move(mapped_buffer).ValueOrUnchanged(std::move(copy->buffer));
copy->value = std::move(mapped_value).ValueOrUnchanged(std::move(copy->value));
copy->indices = std::move(mapped_indices).ValueOrUnchanged(std::move(copy->indices));
- return ffi::details::AnyUnsafe::MoveAnyToTVMFFIAny(ffi::Any(std::move(copy)));
+ return ffi::Any(std::move(copy));
}
-TVMFFIAny BufferStoreMaybeInplaceMutate(ffi::StructuralMutatorObj* mutator,
- ffi::AnyView value) noexcept {
+TVM_FFI_INLINE ffi::Expected<ffi::UnchangedOr<ffi::Any>> BufferStoreMaybeInplaceMutate(
+ ffi::StructuralMutatorObj* mutator, ffi::AnyView value) noexcept {
BufferStoreNode* self = const_cast<BufferStoreNode*>(
ffi::details::AnyUnsafe::RawObjectPtrFromAnyViewAfterCheck<const BufferStoreNode>(value));
TVM_FFI_S_MUTATE_ASSIGN_OR_RETURN(
@@ -646,46 +709,48 @@ TVMFFIAny BufferStoreMaybeInplaceMutate(ffi::StructuralMutatorObj* mutator,
if (mapped_buffer.UnchangedOrSameAs(self->buffer) &&
mapped_value.UnchangedOrSameAs(self->value) &&
mapped_indices.UnchangedOrSameAs(self->indices)) {
- return ffi::Unchanged().CopyToTVMFFIAny();
+ return ffi::Unchanged();
}
if (!mapped_buffer.IsUnchanged()) self->buffer = std::move(mapped_buffer).ValueUnchecked();
if (!mapped_value.IsUnchanged()) self->value = std::move(mapped_value).ValueUnchecked();
if (!mapped_indices.IsUnchanged()) self->indices = std::move(mapped_indices).ValueUnchecked();
- return ffi::Unchanged().CopyToTVMFFIAny();
+ return ffi::Unchanged();
}
-TVMFFIAny ScopeIdDefStmtVisit(ffi::StructuralVisitorObj* visitor, ffi::AnyView value) noexcept {
+TVM_FFI_INLINE ffi::Expected<ffi::Optional<ffi::VisitInterrupt>> ScopeIdDefStmtVisit(
+ ffi::StructuralVisitorObj* visitor, ffi::AnyView value) noexcept {
const ScopeIdDefStmtNode* self =
ffi::details::AnyUnsafe::RawObjectPtrFromAnyViewAfterCheck<const ScopeIdDefStmtNode>(value);
// ScopeIdDef reflection marks only def_ids as definitions; its extent fields remain uses.
TVM_FFI_S_VISIT_MAYBE_EARLY_RETURN(visitor->VisitExpected(self->def));
- return ffi::AnyView(nullptr).CopyToTVMFFIAny();
+ return std::nullopt;
}
-TVMFFIAny ScopeIdDefStmtMutate(ffi::StructuralMutatorObj* mutator, ffi::AnyView value) noexcept {
+TVM_FFI_INLINE ffi::Expected<ffi::UnchangedOr<ffi::Any>> ScopeIdDefStmtMutate(
+ ffi::StructuralMutatorObj* mutator, ffi::AnyView value) noexcept {
const ScopeIdDefStmtNode* self =
ffi::details::AnyUnsafe::RawObjectPtrFromAnyViewAfterCheck<const ScopeIdDefStmtNode>(value);
TVM_FFI_S_MUTATE_ASSIGN_OR_RETURN(ffi::UnchangedOr<ScopeIdDef>, mapped_def,
mutator->MutateExpected(self->def));
if (mapped_def.UnchangedOrSameAs(self->def)) {
- return ffi::Unchanged().CopyToTVMFFIAny();
+ return ffi::Unchanged();
}
ffi::ObjectPtr<ScopeIdDefStmtNode> copy = ffi::make_object<ScopeIdDefStmtNode>(*self);
copy->def = std::move(mapped_def).ValueOrUnchanged(std::move(copy->def));
- return ffi::details::AnyUnsafe::MoveAnyToTVMFFIAny(ffi::Any(std::move(copy)));
+ return ffi::Any(std::move(copy));
}
-TVMFFIAny ScopeIdDefStmtMaybeInplaceMutate(ffi::StructuralMutatorObj* mutator,
- ffi::AnyView value) noexcept {
+TVM_FFI_INLINE ffi::Expected<ffi::UnchangedOr<ffi::Any>> ScopeIdDefStmtMaybeInplaceMutate(
+ ffi::StructuralMutatorObj* mutator, ffi::AnyView value) noexcept {
ScopeIdDefStmtNode* self = const_cast<ScopeIdDefStmtNode*>(
ffi::details::AnyUnsafe::RawObjectPtrFromAnyViewAfterCheck<const ScopeIdDefStmtNode>(value));
TVM_FFI_S_MUTATE_ASSIGN_OR_RETURN(ffi::UnchangedOr<ScopeIdDef>, mapped_def,
mutator->MutateExpected(self->def, ffi::InplaceMode::kAllow));
if (mapped_def.UnchangedOrSameAs(self->def)) {
- return ffi::Unchanged().CopyToTVMFFIAny();
+ return ffi::Unchanged();
}
if (!mapped_def.IsUnchanged()) self->def = std::move(mapped_def).ValueUnchecked();
- return ffi::Unchanged().CopyToTVMFFIAny();
+ return ffi::Unchanged();
}
} // namespace
@@ -708,10 +773,10 @@ TVM_FFI_STATIC_INIT_BLOCK() {
namespace refl = tvm::ffi::reflection;
BindNode::RegisterReflection();
refl::TypeAttrDef<BindNode>()
- .attr(refl::type_attr::kStructuralVisit, reinterpret_cast<void*>(&BindVisit))
- .attr(refl::type_attr::kStructuralMutate, reinterpret_cast<void*>(&BindMutate))
+ .attr(refl::type_attr::kStructuralVisit, ffi::FStructuralVisit::FromNative<&BindVisit>())
+ .attr(refl::type_attr::kStructuralMutate, ffi::FStructuralMutate::FromNative<&BindMutate>())
.attr(refl::type_attr::kStructuralMaybeInplaceMutate,
- reinterpret_cast<void*>(&BindMaybeInplaceMutate));
+ ffi::FStructuralMutate::FromNative<&BindMaybeInplaceMutate>());
refl::GlobalDef().def("tirx.Bind",
[](Var var, Expr value, Span span) { return Bind(var, value, span); });
@@ -732,10 +797,11 @@ TVM_FFI_STATIC_INIT_BLOCK() {
namespace refl = tvm::ffi::reflection;
AttrStmtNode::RegisterReflection();
refl::TypeAttrDef<AttrStmtNode>()
- .attr(refl::type_attr::kStructuralVisit, reinterpret_cast<void*>(&AttrStmtVisit))
- .attr(refl::type_attr::kStructuralMutate, reinterpret_cast<void*>(&AttrStmtMutate))
+ .attr(refl::type_attr::kStructuralVisit, ffi::FStructuralVisit::FromNative<&AttrStmtVisit>())
+ .attr(refl::type_attr::kStructuralMutate,
+ ffi::FStructuralMutate::FromNative<&AttrStmtMutate>())
.attr(refl::type_attr::kStructuralMaybeInplaceMutate,
- reinterpret_cast<void*>(&AttrStmtMaybeInplaceMutate));
+ ffi::FStructuralMutate::FromNative<&AttrStmtMaybeInplaceMutate>());
refl::GlobalDef().def("tirx.AttrStmt",
[](Any node, ffi::String attr_key, Expr value, Stmt body, Span span) {
@@ -765,10 +831,12 @@ TVM_FFI_STATIC_INIT_BLOCK() {
namespace refl = tvm::ffi::reflection;
AssertStmtNode::RegisterReflection();
refl::TypeAttrDef<AssertStmtNode>()
- .attr(refl::type_attr::kStructuralVisit, reinterpret_cast<void*>(&AssertStmtVisit))
- .attr(refl::type_attr::kStructuralMutate, reinterpret_cast<void*>(&AssertStmtMutate))
+ .attr(refl::type_attr::kStructuralVisit,
+ ffi::FStructuralVisit::FromNative<&AssertStmtVisit>())
+ .attr(refl::type_attr::kStructuralMutate,
+ ffi::FStructuralMutate::FromNative<&AssertStmtMutate>())
.attr(refl::type_attr::kStructuralMaybeInplaceMutate,
- reinterpret_cast<void*>(&AssertStmtMaybeInplaceMutate));
+ ffi::FStructuralMutate::FromNative<&AssertStmtMaybeInplaceMutate>());
namespace refl = tvm::ffi::reflection;
refl::GlobalDef().def("tirx.AssertStmt", [](PrimExpr condition, StringImm error_kind,
@@ -846,10 +914,10 @@ TVM_FFI_STATIC_INIT_BLOCK() {
namespace refl = tvm::ffi::reflection;
ForNode::RegisterReflection();
refl::TypeAttrDef<ForNode>()
- .attr(refl::type_attr::kStructuralVisit, reinterpret_cast<void*>(&ForVisit))
- .attr(refl::type_attr::kStructuralMutate, reinterpret_cast<void*>(&ForMutate))
+ .attr(refl::type_attr::kStructuralVisit, ffi::FStructuralVisit::FromNative<&ForVisit>())
+ .attr(refl::type_attr::kStructuralMutate, ffi::FStructuralMutate::FromNative<&ForMutate>())
.attr(refl::type_attr::kStructuralMaybeInplaceMutate,
- reinterpret_cast<void*>(&ForMaybeInplaceMutate));
+ ffi::FStructuralMutate::FromNative<&ForMaybeInplaceMutate>());
refl::GlobalDef().def("tirx.For", [](PrimVar loop_var, PrimExpr min, PrimExpr extent, int kind,
Stmt body, ffi::Optional<IterVar> thread_binding,
@@ -900,10 +968,10 @@ TVM_FFI_STATIC_INIT_BLOCK() {
namespace refl = tvm::ffi::reflection;
WhileNode::RegisterReflection();
refl::TypeAttrDef<WhileNode>()
- .attr(refl::type_attr::kStructuralVisit, reinterpret_cast<void*>(&WhileVisit))
- .attr(refl::type_attr::kStructuralMutate, reinterpret_cast<void*>(&WhileMutate))
+ .attr(refl::type_attr::kStructuralVisit, ffi::FStructuralVisit::FromNative<&WhileVisit>())
+ .attr(refl::type_attr::kStructuralMutate, ffi::FStructuralMutate::FromNative<&WhileMutate>())
.attr(refl::type_attr::kStructuralMaybeInplaceMutate,
- reinterpret_cast<void*>(&WhileMaybeInplaceMutate));
+ ffi::FStructuralMutate::FromNative<&WhileMaybeInplaceMutate>());
refl::GlobalDef().def("tirx.While", [](PrimExpr condition, Stmt body, Span span) {
return While(condition, body, span);
@@ -924,10 +992,10 @@ TVM_FFI_STATIC_INIT_BLOCK() {
namespace refl = tvm::ffi::reflection;
ReturnNode::RegisterReflection();
refl::TypeAttrDef<ReturnNode>()
- .attr(refl::type_attr::kStructuralVisit, reinterpret_cast<void*>(&ReturnVisit))
- .attr(refl::type_attr::kStructuralMutate, reinterpret_cast<void*>(&ReturnMutate))
+ .attr(refl::type_attr::kStructuralVisit, ffi::FStructuralVisit::FromNative<&ReturnVisit>())
+ .attr(refl::type_attr::kStructuralMutate, ffi::FStructuralMutate::FromNative<&ReturnMutate>())
.attr(refl::type_attr::kStructuralMaybeInplaceMutate,
- reinterpret_cast<void*>(&ReturnMaybeInplaceMutate));
+ ffi::FStructuralMutate::FromNative<&ReturnMaybeInplaceMutate>());
refl::GlobalDef().def("tirx.Return", [](Expr value, Span span) { return Return(value, span); });
}
@@ -943,10 +1011,10 @@ TVM_FFI_STATIC_INIT_BLOCK() {
namespace refl = tvm::ffi::reflection;
BreakNode::RegisterReflection();
refl::TypeAttrDef<BreakNode>()
- .attr(refl::type_attr::kStructuralVisit, reinterpret_cast<void*>(&BreakVisit))
- .attr(refl::type_attr::kStructuralMutate, reinterpret_cast<void*>(&BreakMutate))
+ .attr(refl::type_attr::kStructuralVisit, ffi::FStructuralVisit::FromNative<&BreakVisit>())
+ .attr(refl::type_attr::kStructuralMutate, ffi::FStructuralMutate::FromNative<&BreakMutate>())
.attr(refl::type_attr::kStructuralMaybeInplaceMutate,
- reinterpret_cast<void*>(&BreakMaybeInplaceMutate));
+ ffi::FStructuralMutate::FromNative<&BreakMaybeInplaceMutate>());
refl::GlobalDef().def("tirx.Break", [](Span span) { return Break(span); });
}
@@ -962,10 +1030,11 @@ TVM_FFI_STATIC_INIT_BLOCK() {
namespace refl = tvm::ffi::reflection;
ContinueNode::RegisterReflection();
refl::TypeAttrDef<ContinueNode>()
- .attr(refl::type_attr::kStructuralVisit, reinterpret_cast<void*>(&ContinueVisit))
- .attr(refl::type_attr::kStructuralMutate, reinterpret_cast<void*>(&ContinueMutate))
+ .attr(refl::type_attr::kStructuralVisit, ffi::FStructuralVisit::FromNative<&ContinueVisit>())
+ .attr(refl::type_attr::kStructuralMutate,
+ ffi::FStructuralMutate::FromNative<&ContinueMutate>())
.attr(refl::type_attr::kStructuralMaybeInplaceMutate,
- reinterpret_cast<void*>(&ContinueMaybeInplaceMutate));
+ ffi::FStructuralMutate::FromNative<&ContinueMaybeInplaceMutate>());
refl::GlobalDef().def("tirx.Continue", [](Span span) { return Continue(span); });
}
@@ -996,10 +1065,12 @@ TVM_FFI_STATIC_INIT_BLOCK() {
namespace refl = tvm::ffi::reflection;
DeclBufferNode::RegisterReflection();
refl::TypeAttrDef<DeclBufferNode>()
- .attr(refl::type_attr::kStructuralVisit, reinterpret_cast<void*>(&DeclBufferVisit))
- .attr(refl::type_attr::kStructuralMutate, reinterpret_cast<void*>(&DeclBufferMutate))
+ .attr(refl::type_attr::kStructuralVisit,
+ ffi::FStructuralVisit::FromNative<&DeclBufferVisit>())
+ .attr(refl::type_attr::kStructuralMutate,
+ ffi::FStructuralMutate::FromNative<&DeclBufferMutate>())
.attr(refl::type_attr::kStructuralMaybeInplaceMutate,
- reinterpret_cast<void*>(&DeclBufferMaybeInplaceMutate));
+ ffi::FStructuralMutate::FromNative<&DeclBufferMaybeInplaceMutate>());
refl::GlobalDef().def("tirx.DeclBuffer", [](BufferVar buffer, Expr data, Span span) {
return DeclBuffer(buffer, data, span);
@@ -1019,10 +1090,12 @@ TVM_FFI_STATIC_INIT_BLOCK() {
namespace refl = tvm::ffi::reflection;
AllocBufferNode::RegisterReflection();
refl::TypeAttrDef<AllocBufferNode>()
- .attr(refl::type_attr::kStructuralVisit, reinterpret_cast<void*>(&AllocBufferVisit))
- .attr(refl::type_attr::kStructuralMutate, reinterpret_cast<void*>(&AllocBufferMutate))
+ .attr(refl::type_attr::kStructuralVisit,
+ ffi::FStructuralVisit::FromNative<&AllocBufferVisit>())
+ .attr(refl::type_attr::kStructuralMutate,
+ ffi::FStructuralMutate::FromNative<&AllocBufferMutate>())
.attr(refl::type_attr::kStructuralMaybeInplaceMutate,
- reinterpret_cast<void*>(&AllocBufferMaybeInplaceMutate));
+ ffi::FStructuralMutate::FromNative<&AllocBufferMaybeInplaceMutate>());
namespace refl = tvm::ffi::reflection;
refl::GlobalDef().def(
@@ -1063,10 +1136,11 @@ TVM_FFI_STATIC_INIT_BLOCK() {
namespace refl = tvm::ffi::reflection;
SeqStmtNode::RegisterReflection();
refl::TypeAttrDef<SeqStmtNode>()
- .attr(refl::type_attr::kStructuralVisit, reinterpret_cast<void*>(&SeqStmtVisit))
- .attr(refl::type_attr::kStructuralMutate, reinterpret_cast<void*>(&SeqStmtMutate))
+ .attr(refl::type_attr::kStructuralVisit, ffi::FStructuralVisit::FromNative<&SeqStmtVisit>())
+ .attr(refl::type_attr::kStructuralMutate,
+ ffi::FStructuralMutate::FromNative<&SeqStmtMutate>())
.attr(refl::type_attr::kStructuralMaybeInplaceMutate,
- reinterpret_cast<void*>(&SeqStmtMaybeInplaceMutate));
+ ffi::FStructuralMutate::FromNative<&SeqStmtMaybeInplaceMutate>());
refl::GlobalDef().def("tirx.SeqStmt", [](ffi::Array<Stmt> seq, Span span) {
return SeqStmt(std::move(seq), span);
@@ -1091,10 +1165,12 @@ TVM_FFI_STATIC_INIT_BLOCK() {
namespace refl = tvm::ffi::reflection;
IfThenElseNode::RegisterReflection();
refl::TypeAttrDef<IfThenElseNode>()
- .attr(refl::type_attr::kStructuralVisit, reinterpret_cast<void*>(&IfThenElseVisit))
- .attr(refl::type_attr::kStructuralMutate, reinterpret_cast<void*>(&IfThenElseMutate))
+ .attr(refl::type_attr::kStructuralVisit,
+ ffi::FStructuralVisit::FromNative<&IfThenElseVisit>())
+ .attr(refl::type_attr::kStructuralMutate,
+ ffi::FStructuralMutate::FromNative<&IfThenElseMutate>())
.attr(refl::type_attr::kStructuralMaybeInplaceMutate,
- reinterpret_cast<void*>(&IfThenElseMaybeInplaceMutate));
+ ffi::FStructuralMutate::FromNative<&IfThenElseMaybeInplaceMutate>());
refl::GlobalDef().def("tirx.IfThenElse",
[](PrimExpr condition, Stmt then_case, Stmt else_case, Span span) {
@@ -1119,10 +1195,11 @@ TVM_FFI_STATIC_INIT_BLOCK() {
namespace refl = tvm::ffi::reflection;
EvaluateNode::RegisterReflection();
refl::TypeAttrDef<EvaluateNode>()
- .attr(refl::type_attr::kStructuralVisit, reinterpret_cast<void*>(&EvaluateVisit))
- .attr(refl::type_attr::kStructuralMutate, reinterpret_cast<void*>(&EvaluateMutate))
+ .attr(refl::type_attr::kStructuralVisit, ffi::FStructuralVisit::FromNative<&EvaluateVisit>())
+ .attr(refl::type_attr::kStructuralMutate,
+ ffi::FStructuralMutate::FromNative<&EvaluateMutate>())
.attr(refl::type_attr::kStructuralMaybeInplaceMutate,
- reinterpret_cast<void*>(&EvaluateMaybeInplaceMutate));
+ ffi::FStructuralMutate::FromNative<&EvaluateMaybeInplaceMutate>());
refl::GlobalDef().def("tirx.Evaluate",
[](Expr value, Span span) { return Evaluate(value, span); });
@@ -1194,10 +1271,12 @@ TVM_FFI_STATIC_INIT_BLOCK() {
namespace refl = tvm::ffi::reflection;
BufferStoreNode::RegisterReflection();
refl::TypeAttrDef<BufferStoreNode>()
- .attr(refl::type_attr::kStructuralVisit, reinterpret_cast<void*>(&BufferStoreVisit))
- .attr(refl::type_attr::kStructuralMutate, reinterpret_cast<void*>(&BufferStoreMutate))
+ .attr(refl::type_attr::kStructuralVisit,
+ ffi::FStructuralVisit::FromNative<&BufferStoreVisit>())
+ .attr(refl::type_attr::kStructuralMutate,
+ ffi::FStructuralMutate::FromNative<&BufferStoreMutate>())
.attr(refl::type_attr::kStructuralMaybeInplaceMutate,
- reinterpret_cast<void*>(&BufferStoreMaybeInplaceMutate));
+ ffi::FStructuralMutate::FromNative<&BufferStoreMaybeInplaceMutate>());
refl::GlobalDef().def("tirx.BufferStore",
[](BufferVar buffer, PrimExpr value, ffi::Array<PrimExpr> indices,
@@ -1217,10 +1296,12 @@ TVM_FFI_STATIC_INIT_BLOCK() {
namespace refl = tvm::ffi::reflection;
ScopeIdDefStmtNode::RegisterReflection();
refl::TypeAttrDef<ScopeIdDefStmtNode>()
- .attr(refl::type_attr::kStructuralVisit, reinterpret_cast<void*>(&ScopeIdDefStmtVisit))
- .attr(refl::type_attr::kStructuralMutate, reinterpret_cast<void*>(&ScopeIdDefStmtMutate))
+ .attr(refl::type_attr::kStructuralVisit,
+ ffi::FStructuralVisit::FromNative<&ScopeIdDefStmtVisit>())
+ .attr(refl::type_attr::kStructuralMutate,
+ ffi::FStructuralMutate::FromNative<&ScopeIdDefStmtMutate>())
.attr(refl::type_attr::kStructuralMaybeInplaceMutate,
- reinterpret_cast<void*>(&ScopeIdDefStmtMaybeInplaceMutate));
+ ffi::FStructuralMutate::FromNative<&ScopeIdDefStmtMaybeInplaceMutate>());
refl::GlobalDef().def("tirx.ScopeIdDefStmt",
[](ScopeIdDef def, Span span) { return ScopeIdDefStmt(def, span); });
diff --git a/src/tirx/ir/stmt_functor.cc b/src/tirx/ir/stmt_functor.cc
index 35d997d58f..7bc8762aa6 100644
--- a/src/tirx/ir/stmt_functor.cc
+++ b/src/tirx/ir/stmt_functor.cc
@@ -292,15 +292,32 @@ void StmtExprMutator::InitVTable(VTable* vtable) {
}
UnchangedOr<Stmt> StmtExprMutator::Mutate_(const BindNode* op, InplaceMode inplace_mode) {
- auto value = Mutate(op->value, inplace_mode);
- if (value.UnchangedOrSameAs(op->value)) return ffi::Unchanged();
+ Type old_value_type = op->value->ty;
+ auto value_u = Mutate(op->value, inplace_mode);
+ bool value_changed = !value_u.UnchangedOrSameAs(op->value);
+ Expr value = std::move(value_u).ValueOrUnchanged(op->value);
+ value_changed |= !old_value_type.same_as(value->ty);
+ Type old_var_type = op->var->ty;
+ auto var_u = WithDefRegionKind(kTVMFFIDefRegionKindSimple, [&] {
+ return Mutate(op->var, inplace_mode);
+ }).as_or_throw<UnchangedOr<Var>>();
+ bool var_changed = !var_u.UnchangedOrSameAs(op->var);
+ Var var = std::move(var_u).ValueOrUnchanged(op->var);
+ var_changed |= !old_var_type.same_as(var->ty);
+ if ((value_changed || var_changed) && !var->ty.same_as(value->ty)) {
+ var = var.CopyWithType(value->ty);
+ }
+ if (!var.same_as(op->var)) VarRemapSet(op->var, var);
+ if (value.same_as(op->value) && var.same_as(op->var)) return ffi::Unchanged();
if (inplace_mode == InplaceMode::kAllow) {
auto* writable = const_cast<BindNode*>(op);
- if (!value.IsUnchanged()) writable->value = std::move(value).ValueUnchecked();
+ writable->value = std::move(value);
+ writable->var = std::move(var);
return ffi::Unchanged();
}
auto copy = ffi::make_object<BindNode>(*op);
- if (!value.IsUnchanged()) copy->value = std::move(value).ValueUnchecked();
+ copy->value = std::move(value);
+ copy->var = std::move(var);
return Stmt(std::move(copy));
}
diff --git a/src/tirx/ir/tirx_stmt.cc b/src/tirx/ir/tirx_stmt.cc
index be57749739..571593c982 100644
--- a/src/tirx/ir/tirx_stmt.cc
+++ b/src/tirx/ir/tirx_stmt.cc
@@ -59,7 +59,8 @@ TilePrimitiveCall::TilePrimitiveCall(tvm::Op op, ffi::Array<ffi::Any> args,
namespace {
-TVMFFIAny TilePrimitiveCallVisit(ffi::StructuralVisitorObj* visitor, ffi::AnyView value) noexcept {
+TVM_FFI_INLINE ffi::Expected<ffi::Optional<ffi::VisitInterrupt>> TilePrimitiveCallVisit(
+ ffi::StructuralVisitorObj* visitor, ffi::AnyView value) noexcept {
// skips: op, dispatch, scope
const TilePrimitiveCallNode* self =
ffi::details::AnyUnsafe::RawObjectPtrFromAnyViewAfterCheck<const TilePrimitiveCallNode>(
@@ -67,10 +68,11 @@ TVMFFIAny TilePrimitiveCallVisit(ffi::StructuralVisitorObj* visitor, ffi::AnyVie
TVM_FFI_S_VISIT_MAYBE_EARLY_RETURN(visitor->VisitExpected(self->args));
TVM_FFI_S_VISIT_MAYBE_EARLY_RETURN(visitor->VisitExpected(self->workspace));
TVM_FFI_S_VISIT_MAYBE_EARLY_RETURN(visitor->VisitExpected(self->config));
- return ffi::AnyView(nullptr).CopyToTVMFFIAny();
+ return std::nullopt;
}
-TVMFFIAny TilePrimitiveCallMutate(ffi::StructuralMutatorObj* mutator, ffi::AnyView value) noexcept {
+TVM_FFI_INLINE ffi::Expected<ffi::UnchangedOr<ffi::Any>> TilePrimitiveCallMutate(
+ ffi::StructuralMutatorObj* mutator, ffi::AnyView value) noexcept {
// skips: op, dispatch, scope
const TilePrimitiveCallNode* self =
ffi::details::AnyUnsafe::RawObjectPtrFromAnyViewAfterCheck<const TilePrimitiveCallNode>(
@@ -88,17 +90,17 @@ TVMFFIAny TilePrimitiveCallMutate(ffi::StructuralMutatorObj* mutator, ffi::AnyVi
if (mapped_args.UnchangedOrSameAs(self->args) &&
mapped_workspace.UnchangedOrSameAs(self->workspace) &&
mapped_config.UnchangedOrSameAs(self->config)) {
- return ffi::Unchanged().CopyToTVMFFIAny();
+ return ffi::Unchanged();
}
ffi::ObjectPtr<TilePrimitiveCallNode> copy = ffi::make_object<TilePrimitiveCallNode>(*self);
copy->args = std::move(mapped_args).ValueOrUnchanged(std::move(copy->args));
copy->workspace = std::move(mapped_workspace).ValueOrUnchanged(std::move(copy->workspace));
copy->config = std::move(mapped_config).ValueOrUnchanged(std::move(copy->config));
- return ffi::details::AnyUnsafe::MoveAnyToTVMFFIAny(ffi::Any(std::move(copy)));
+ return ffi::Any(std::move(copy));
}
-TVMFFIAny TilePrimitiveCallMaybeInplaceMutate(ffi::StructuralMutatorObj* mutator,
- ffi::AnyView value) noexcept {
+TVM_FFI_INLINE ffi::Expected<ffi::UnchangedOr<ffi::Any>> TilePrimitiveCallMaybeInplaceMutate(
+ ffi::StructuralMutatorObj* mutator, ffi::AnyView value) noexcept {
// skips: op, dispatch, scope
TilePrimitiveCallNode* self = const_cast<TilePrimitiveCallNode*>(
ffi::details::AnyUnsafe::RawObjectPtrFromAnyViewAfterCheck<const TilePrimitiveCallNode>(
@@ -118,12 +120,12 @@ TVMFFIAny TilePrimitiveCallMaybeInplaceMutate(ffi::StructuralMutatorObj* mutator
if (mapped_args.UnchangedOrSameAs(self->args) &&
mapped_workspace.UnchangedOrSameAs(self->workspace) &&
mapped_config.UnchangedOrSameAs(self->config)) {
- return ffi::Unchanged().CopyToTVMFFIAny();
+ return ffi::Unchanged();
}
self->args = std::move(mapped_args).ValueOrUnchanged(std::move(self->args));
self->workspace = std::move(mapped_workspace).ValueOrUnchanged(std::move(self->workspace));
self->config = std::move(mapped_config).ValueOrUnchanged(std::move(self->config));
- return ffi::Unchanged().CopyToTVMFFIAny();
+ return ffi::Unchanged();
}
} // namespace
@@ -139,10 +141,12 @@ TVM_FFI_STATIC_INIT_BLOCK() {
return TilePrimitiveCall(op, args, workspace, config, dispatch, scope);
});
refl::TypeAttrDef<TilePrimitiveCallNode>()
- .attr(refl::type_attr::kStructuralVisit, reinterpret_cast<void*>(&TilePrimitiveCallVisit))
- .attr(refl::type_attr::kStructuralMutate, reinterpret_cast<void*>(&TilePrimitiveCallMutate))
+ .attr(refl::type_attr::kStructuralVisit,
+ ffi::FStructuralVisit::FromNative<&TilePrimitiveCallVisit>())
+ .attr(refl::type_attr::kStructuralMutate,
+ ffi::FStructuralMutate::FromNative<&TilePrimitiveCallMutate>())
.attr(refl::type_attr::kStructuralMaybeInplaceMutate,
- reinterpret_cast<void*>(&TilePrimitiveCallMaybeInplaceMutate));
+ ffi::FStructuralMutate::FromNative<&TilePrimitiveCallMaybeInplaceMutate>());
}
TVM_FFI_STATIC_INIT_BLOCK() {
diff --git a/src/tirx/ir/type.cc b/src/tirx/ir/type.cc
index 2ef4ad53ac..9163b41918 100644
--- a/src/tirx/ir/type.cc
+++ b/src/tirx/ir/type.cc
@@ -32,19 +32,23 @@
namespace tvm::tirx {
namespace {
-TVMFFIAny TensorMapTypeVisit(ffi::StructuralVisitorObj*, ffi::AnyView) noexcept {
- return ffi::AnyView(nullptr).CopyToTVMFFIAny();
+TVM_FFI_INLINE ffi::Expected<ffi::Optional<ffi::VisitInterrupt>> TensorMapTypeVisit(
+ ffi::StructuralVisitorObj*, ffi::AnyView) noexcept {
+ return std::nullopt;
}
-TVMFFIAny TensorMapTypeMutate(ffi::StructuralMutatorObj*, ffi::AnyView) noexcept {
- return ffi::Unchanged().CopyToTVMFFIAny();
+TVM_FFI_INLINE ffi::Expected<ffi::UnchangedOr<ffi::Any>> TensorMapTypeMutate(
+ ffi::StructuralMutatorObj*, ffi::AnyView) noexcept {
+ return ffi::Unchanged();
}
-TVMFFIAny TensorMapTypeMaybeInplaceMutate(ffi::StructuralMutatorObj*, ffi::AnyView) noexcept {
- return ffi::Unchanged().CopyToTVMFFIAny();
+TVM_FFI_INLINE ffi::Expected<ffi::UnchangedOr<ffi::Any>> TensorMapTypeMaybeInplaceMutate(
+ ffi::StructuralMutatorObj*, ffi::AnyView) noexcept {
+ return ffi::Unchanged();
}
-TVMFFIAny BufferTypeVisit(ffi::StructuralVisitorObj* visitor, ffi::AnyView value) noexcept {
+TVM_FFI_INLINE ffi::Expected<ffi::Optional<ffi::VisitInterrupt>> BufferTypeVisit(
+ ffi::StructuralVisitorObj* visitor, ffi::AnyView value) noexcept {
// skips: storage_scope, data_alignment, offset_factor
const BufferTypeNode* self =
ffi::details::AnyUnsafe::RawObjectPtrFromAnyViewAfterCheck<const BufferTypeNode>(value);
@@ -62,10 +66,11 @@ TVMFFIAny BufferTypeVisit(ffi::StructuralVisitorObj* visitor, ffi::AnyView value
if (!self->allocated_addr.empty()) {
TVM_FFI_S_VISIT_MAYBE_EARLY_RETURN(visitor->VisitExpected(self->allocated_addr));
}
- return ffi::AnyView(nullptr).CopyToTVMFFIAny();
+ return std::nullopt;
}
-TVMFFIAny BufferTypeMutate(ffi::StructuralMutatorObj* mutator, ffi::AnyView value) noexcept {
+TVM_FFI_INLINE ffi::Expected<ffi::UnchangedOr<ffi::Any>> BufferTypeMutate(
+ ffi::StructuralMutatorObj* mutator, ffi::AnyView value) noexcept {
// skips: storage_scope, data_alignment, offset_factor
const BufferTypeNode* self =
ffi::details::AnyUnsafe::RawObjectPtrFromAnyViewAfterCheck<const BufferTypeNode>(value);
@@ -99,7 +104,7 @@ TVMFFIAny BufferTypeMutate(ffi::StructuralMutatorObj* mutator, ffi::AnyView valu
mapped_elem_offset.UnchangedOrSameAs(self->elem_offset) &&
mapped_layout.UnchangedOrSameAs(self->layout) &&
mapped_allocated_addr.UnchangedOrSameAs(self->allocated_addr)) {
- return ffi::Unchanged().CopyToTVMFFIAny();
+ return ffi::Unchanged();
}
ffi::ObjectPtr<BufferTypeNode> copy = ffi::make_object<BufferTypeNode>(*self);
copy->dtype = std::move(mapped_dtype).ValueOrUnchanged(std::move(copy->dtype));
@@ -109,11 +114,11 @@ TVMFFIAny BufferTypeMutate(ffi::StructuralMutatorObj* mutator, ffi::AnyView valu
copy->layout = std::move(mapped_layout).ValueOrUnchanged(std::move(copy->layout));
copy->allocated_addr =
std::move(mapped_allocated_addr).ValueOrUnchanged(std::move(copy->allocated_addr));
- return ffi::details::AnyUnsafe::MoveAnyToTVMFFIAny(ffi::Any(std::move(copy)));
+ return ffi::Any(std::move(copy));
}
-TVMFFIAny BufferTypeMaybeInplaceMutate(ffi::StructuralMutatorObj* mutator,
- ffi::AnyView value) noexcept {
+TVM_FFI_INLINE ffi::Expected<ffi::UnchangedOr<ffi::Any>> BufferTypeMaybeInplaceMutate(
+ ffi::StructuralMutatorObj* mutator, ffi::AnyView value) noexcept {
// skips: storage_scope, data_alignment, offset_factor
BufferTypeNode* self = const_cast<BufferTypeNode*>(
ffi::details::AnyUnsafe::RawObjectPtrFromAnyViewAfterCheck<const BufferTypeNode>(value));
@@ -150,7 +155,7 @@ TVMFFIAny BufferTypeMaybeInplaceMutate(ffi::StructuralMutatorObj* mutator,
mapped_elem_offset.UnchangedOrSameAs(self->elem_offset) &&
mapped_layout.UnchangedOrSameAs(self->layout) &&
mapped_allocated_addr.UnchangedOrSameAs(self->allocated_addr)) {
- return ffi::Unchanged().CopyToTVMFFIAny();
+ return ffi::Unchanged();
}
if (!mapped_dtype.IsUnchanged()) self->dtype = std::move(mapped_dtype).ValueUnchecked();
if (!mapped_shape.IsUnchanged()) self->shape = std::move(mapped_shape).ValueUnchecked();
@@ -161,19 +166,22 @@ TVMFFIAny BufferTypeMaybeInplaceMutate(ffi::StructuralMutatorObj* mutator,
if (!mapped_allocated_addr.IsUnchanged()) {
self->allocated_addr = std::move(mapped_allocated_addr).ValueUnchecked();
}
- return ffi::Unchanged().CopyToTVMFFIAny();
+ return ffi::Unchanged();
}
-TVMFFIAny BufferRegionTypeVisit(ffi::StructuralVisitorObj*, ffi::AnyView) noexcept {
- return ffi::AnyView(nullptr).CopyToTVMFFIAny();
+TVM_FFI_INLINE ffi::Expected<ffi::Optional<ffi::VisitInterrupt>> BufferRegionTypeVisit(
+ ffi::StructuralVisitorObj*, ffi::AnyView) noexcept {
+ return std::nullopt;
}
-TVMFFIAny BufferRegionTypeMutate(ffi::StructuralMutatorObj*, ffi::AnyView) noexcept {
- return ffi::Unchanged().CopyToTVMFFIAny();
+TVM_FFI_INLINE ffi::Expected<ffi::UnchangedOr<ffi::Any>> BufferRegionTypeMutate(
+ ffi::StructuralMutatorObj*, ffi::AnyView) noexcept {
+ return ffi::Unchanged();
}
-TVMFFIAny BufferRegionTypeMaybeInplaceMutate(ffi::StructuralMutatorObj*, ffi::AnyView) noexcept {
- return ffi::Unchanged().CopyToTVMFFIAny();
+TVM_FFI_INLINE ffi::Expected<ffi::UnchangedOr<ffi::Any>> BufferRegionTypeMaybeInplaceMutate(
+ ffi::StructuralMutatorObj*, ffi::AnyView) noexcept {
+ return ffi::Unchanged();
}
} // namespace
@@ -205,10 +213,12 @@ TVM_FFI_STATIC_INIT_BLOCK() {
namespace refl = tvm::ffi::reflection;
BufferTypeNode::RegisterReflection();
refl::TypeAttrDef<BufferTypeNode>()
- .attr(refl::type_attr::kStructuralVisit, reinterpret_cast<void*>(&BufferTypeVisit))
- .attr(refl::type_attr::kStructuralMutate, reinterpret_cast<void*>(&BufferTypeMutate))
+ .attr(refl::type_attr::kStructuralVisit,
+ ffi::FStructuralVisit::FromNative<&BufferTypeVisit>())
+ .attr(refl::type_attr::kStructuralMutate,
+ ffi::FStructuralMutate::FromNative<&BufferTypeMutate>())
.attr(refl::type_attr::kStructuralMaybeInplaceMutate,
- reinterpret_cast<void*>(&BufferTypeMaybeInplaceMutate));
+ ffi::FStructuralMutate::FromNative<&BufferTypeMaybeInplaceMutate>());
refl::GlobalDef().def(
"tirx.BufferType",
@@ -231,10 +241,12 @@ TVM_FFI_STATIC_INIT_BLOCK() {
namespace refl = tvm::ffi::reflection;
BufferRegionTypeNode::RegisterReflection();
refl::TypeAttrDef<BufferRegionTypeNode>()
- .attr(refl::type_attr::kStructuralVisit, reinterpret_cast<void*>(&BufferRegionTypeVisit))
- .attr(refl::type_attr::kStructuralMutate, reinterpret_cast<void*>(&BufferRegionTypeMutate))
+ .attr(refl::type_attr::kStructuralVisit,
+ ffi::FStructuralVisit::FromNative<&BufferRegionTypeVisit>())
+ .attr(refl::type_attr::kStructuralMutate,
+ ffi::FStructuralMutate::FromNative<&BufferRegionTypeMutate>())
.attr(refl::type_attr::kStructuralMaybeInplaceMutate,
- reinterpret_cast<void*>(&BufferRegionTypeMaybeInplaceMutate));
+ ffi::FStructuralMutate::FromNative<&BufferRegionTypeMaybeInplaceMutate>());
refl::GlobalDef().def("tirx.BufferRegionType", []() { return BufferRegionType(); });
}
@@ -249,10 +261,12 @@ TVM_FFI_STATIC_INIT_BLOCK() {
namespace refl = tvm::ffi::reflection;
TensorMapTypeNode::RegisterReflection();
refl::TypeAttrDef<TensorMapTypeNode>()
- .attr(refl::type_attr::kStructuralVisit, reinterpret_cast<void*>(&TensorMapTypeVisit))
- .attr(refl::type_attr::kStructuralMutate, reinterpret_cast<void*>(&TensorMapTypeMutate))
+ .attr(refl::type_attr::kStructuralVisit,
+ ffi::FStructuralVisit::FromNative<&TensorMapTypeVisit>())
+ .attr(refl::type_attr::kStructuralMutate,
+ ffi::FStructuralMutate::FromNative<&TensorMapTypeMutate>())
.attr(refl::type_attr::kStructuralMaybeInplaceMutate,
- reinterpret_cast<void*>(&TensorMapTypeMaybeInplaceMutate));
+ ffi::FStructuralMutate::FromNative<&TensorMapTypeMaybeInplaceMutate>());
refl::GlobalDef().def("tirx.TensorMapType", [](Span span) { return TensorMapType(span); });
}
diff --git a/src/tirx/op/builtin.cc b/src/tirx/op/builtin.cc
index 8f04560055..0d2125073e 100644
... 394 lines suppressed ...