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 11:22:44 UTC
(tvm) branch main updated: [FEATURE][IR] Add typed OpDef signature registration (#20478)
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 b98f2584d7 [FEATURE][IR] Add typed OpDef signature registration (#20478)
b98f2584d7 is described below
commit b98f2584d7dbbead14a9602d0f1ad54a45c1080e
Author: Tianqi Chen <tq...@users.noreply.github.com>
AuthorDate: Mon Sep 28 07:22:35 2026 -0400
[FEATURE][IR] Add typed OpDef signature registration (#20478)
---
3rdparty/tvm-ffi | 2 +-
include/tvm/ir/base_expr.h | 18 +-
include/tvm/ir/expr.h | 8 +
include/tvm/ir/ffi_structural_compat.h | 118 +++++
include/tvm/ir/op.h | 410 +++++++++++++++++-
include/tvm/relax/op_attr_types.h | 27 +-
python/tvm/backend/cuda/ptx/engine.py | 2 +-
python/tvm/ir/expr.py | 53 ++-
python/tvm/ir/op.py | 52 +++
python/tvm/relax/script/ir_builder/op.py | 4 +-
src/backend/cuda/codegen/intrin_rule_cuda.cc | 32 +-
src/backend/cuda/op/iket.cc | 19 +-
src/backend/cuda/op/target_builtin.cc | 474 +++++++++++----------
.../hexagon/codegen/llvm/intrin_rule_hexagon.cc | 2 +-
src/backend/metal/codegen/intrin_rule_metal.cc | 12 +-
src/backend/metal/op/target_builtin.cc | 42 +-
src/backend/trn/op/target_builtin.cc | 99 +++--
src/backend/webgpu/codegen/intrin_rule_webgpu.cc | 12 +-
src/ir/expr.cc | 20 +-
src/ir/op.cc | 92 +++-
src/ir/prim/builtin.cc | 14 +-
src/relax/analysis/well_formed.cc | 8 +-
.../backend/adreno/annotate_custom_storage.cc | 9 +-
.../backend/adreno/fold_vdevice_scope_change.cc | 3 +-
src/relax/backend/vm/lower_runtime_builtin.cc | 34 +-
src/relax/backend/vm/vm_shape_lower.cc | 16 +-
.../transform/lower_global_view_to_local_view.cc | 4 +-
src/relax/ir/block_builder.cc | 7 +-
src/relax/op/ccl/ccl.cc | 25 +-
src/relax/op/distributed/distributed.cc | 22 +-
src/relax/op/image/resize.cc | 30 +-
src/relax/op/memory/view.cc | 33 +-
src/relax/op/nn/attention.cc | 48 +--
src/relax/op/nn/convolution.cc | 54 +--
src/relax/op/nn/convolution.h | 2 +-
src/relax/op/nn/nn.cc | 178 ++++----
src/relax/op/nn/pooling.cc | 43 +-
src/relax/op/op.cc | 322 ++++++++------
src/relax/op/op_common.cc | 10 +-
src/relax/op/op_common.h | 8 +-
src/relax/op/tensor/binary.h | 31 +-
src/relax/op/tensor/create.cc | 100 +++--
src/relax/op/tensor/datatype.cc | 10 +-
src/relax/op/tensor/grad.cc | 56 +--
src/relax/op/tensor/index.cc | 31 +-
src/relax/op/tensor/inspect.cc | 46 +-
src/relax/op/tensor/linear_algebra.cc | 17 +-
src/relax/op/tensor/manipulate.cc | 164 ++++---
src/relax/op/tensor/qdq.cc | 24 +-
src/relax/op/tensor/sampling.cc | 14 +-
src/relax/op/tensor/search.cc | 33 +-
src/relax/op/tensor/set.cc | 41 +-
src/relax/op/tensor/sorting.cc | 15 +-
src/relax/op/tensor/statistical.cc | 42 +-
src/relax/op/tensor/ternary.cc | 10 +-
src/relax/op/tensor/unary.cc | 225 +++++-----
src/relax/op/vision/multibox_transform_loc.cc | 14 +-
src/relax/op/vision/nms.cc | 43 +-
src/relax/op/vision/roi_align.cc | 11 +-
src/relax/op/vision/roi_pool.cc | 11 +-
src/relax/script/ir_builder/distributed.cc | 2 +-
src/relax/transform/allocate_workspace.cc | 8 +-
src/relax/transform/alter_op_impl.cc | 13 +-
src/relax/transform/call_tir_rewrite.cc | 14 +-
src/relax/transform/convert_layout.cc | 4 +-
src/relax/transform/decompose_ops.cc | 6 +-
src/relax/transform/fuse_ops.cc | 17 +-
src/relax/transform/fuse_tir.cc | 2 +-
src/relax/transform/gradient.cc | 4 +-
src/relax/transform/kill_after_last_use.cc | 12 +-
src/relax/transform/lambda_lift.cc | 17 +-
src/relax/transform/lazy_transform_params.cc | 14 +-
src/relax/transform/legalize_ops.cc | 7 +-
src/relax/transform/lift_transform_params.cc | 8 +-
src/relax/transform/lower_alloc_tensor.cc | 4 +-
src/relax/transform/merge_composite_functions.cc | 4 +-
src/relax/transform/realize_vdevice.cc | 2 +-
src/relax/transform/remove_purity_checking.cc | 15 +-
src/relax/transform/rewrite_cuda_graph.cc | 11 +-
src/relax/transform/run_codegen.cc | 2 +-
.../transform/split_layout_rewrite_preproc.cc | 8 +-
src/relax/transform/static_plan_block_memory.cc | 18 +-
src/target/intrin_rule.cc | 4 +-
src/te/tensor.cc | 3 +-
src/tirx/op/builtin.cc | 317 ++++++--------
src/tirx/op/op.cc | 95 ++---
src/tirx/op/tirx.cc | 178 +++-----
tests/cpp/ir_op_signature_test.cc | 193 +++++++++
tests/python/ir/test_ir_attrs.py | 87 ++++
.../script/test_relax_script_dynamic_shape.py | 2 +-
tests/python/relax/test_analysis_well_formed.py | 34 +-
91 files changed, 2695 insertions(+), 1716 deletions(-)
diff --git a/3rdparty/tvm-ffi b/3rdparty/tvm-ffi
index e251c6abee..0af62893f8 160000
--- a/3rdparty/tvm-ffi
+++ b/3rdparty/tvm-ffi
@@ -1 +1 @@
-Subproject commit e251c6abee7a111c6467f84d854be512e33e2db0
+Subproject commit 0af62893f8033b65c3f424a0dbfd3138b5cbe5d4
diff --git a/include/tvm/ir/base_expr.h b/include/tvm/ir/base_expr.h
index 3b4b86bc42..3433550676 100644
--- a/include/tvm/ir/base_expr.h
+++ b/include/tvm/ir/base_expr.h
@@ -28,6 +28,8 @@
#include <tvm/ffi/dtype.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>
@@ -530,12 +532,24 @@ struct TypeTraits<PrimExpr>
using Base::CheckAnyStrict;
using Base::CopyFromAnyViewAfterCheck;
using Base::CopyToAnyView;
- using Base::GetMismatchTypeInfo;
using Base::MoveFromAnyAfterCheck;
using Base::MoveToAny;
using Base::TryCastFromAnyView;
using Base::TypeSchema;
- using Base::TypeStr;
+
+ TVM_FFI_INLINE static std::string TypeStr() { return "ir.PrimExpr"; }
+
+ TVM_FFI_INLINE static std::string GetMismatchTypeInfo(const TVMFFIAny* source) {
+ if (source->type_index >= TypeIndex::kTVMFFIStaticObjectBegin && source->v_obj != nullptr &&
+ details::IsObjectInstance<ExprNode>(source->type_index)) {
+ const auto* expr = details::ObjectUnsafe::RawObjectPtrFromUnowned<ExprNode>(source->v_obj);
+ if (expr->ty.defined() &&
+ !details::AnyUnsafe::CheckAnyViewStrict<PrimType>(AnyView(expr->ty))) {
+ return TypeIndexToTypeKey(source->type_index) + "[ty=" + expr->ty.GetTypeKey() + "]";
+ }
+ }
+ return Base::GetMismatchTypeInfo(source);
+ }
TVM_DLL static PrimExpr ConvertFallbackValue(StrictBool value);
TVM_DLL static PrimExpr ConvertFallbackValue(int64_t value);
diff --git a/include/tvm/ir/expr.h b/include/tvm/ir/expr.h
index 7fa6dbf532..7bf93aa25d 100644
--- a/include/tvm/ir/expr.h
+++ b/include/tvm/ir/expr.h
@@ -510,8 +510,16 @@ class Call : public Expr {
TVM_DLL Call(Type ret_ty, Expr op, ffi::Array<Expr> args, Attrs attrs = Attrs(),
ffi::Array<Type> ty_args = ffi::Array<Type>(), Span span = Span());
+ /*! \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());
+
TVM_FFI_DEFINE_OBJECT_REF_METHODS_NULLABLE(Call, Expr, CallNode);
TVM_DEFINE_OBJECT_REF_COW_METHOD(CallNode);
+
+ private:
+ Call(Type ret_ty, Expr op, ffi::Array<Expr> args, Attrs attrs, ffi::Array<Type> ty_args,
+ Span span, bool validate);
};
/*! \brief Base node for literal constants. */
diff --git a/include/tvm/ir/ffi_structural_compat.h b/include/tvm/ir/ffi_structural_compat.h
new file mode 100644
index 0000000000..d892cbf8c8
--- /dev/null
+++ b/include/tvm/ir/ffi_structural_compat.h
@@ -0,0 +1,118 @@
+/*
+ * 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 9a6a530ef9..49fb8223e1 100644
--- a/include/tvm/ir/op.h
+++ b/include/tvm/ir/op.h
@@ -25,12 +25,16 @@
#define TVM_IR_OP_H_
#include <tvm/ffi/container/list.h>
+#include <tvm/ffi/expected.h>
+#include <tvm/ffi/optional.h>
+#include <tvm/ffi/reflection/native_function.h>
#include <tvm/ffi/reflection/registry.h>
#include <tvm/ir/attrs.h>
#include <tvm/ir/env_func.h>
#include <tvm/ir/expr.h>
#include <tvm/ir/type.h>
+#include <string>
#include <type_traits>
#include <utility>
@@ -77,15 +81,20 @@ class OpNode : public ExprNode {
ffi::String doc;
/*! \brief Descriptors for the ordered required value-argument prefix. */
ffi::Array<ArgumentInfo> args_info;
- /*! \brief Accept at least args_info.size() value arguments when true, exactly that count
- * otherwise. */
- bool allow_extra_args{false};
- /*! \brief Descriptors for Call.ty_args; inference owns optional/variable type-argument checks. */
+ /*! \brief Descriptors for the ordered required type-argument prefix. */
ffi::Array<ArgumentInfo> ty_args_info;
- /*! \brief Attribute object type key, or empty when unrestricted. */
+ /*! \brief Named variadic value tail, excluded from the required args_info prefix. */
+ ffi::Optional<ArgumentInfo> var_args_info;
+ /*! \brief Named variadic type-argument tail, excluded from ty_args_info. */
+ ffi::Optional<ArgumentInfo> var_ty_args_info;
+ /*! \brief Attribute object type key, or empty when no attrs type is declared. */
ffi::String attrs_type_key;
/*! \brief Runtime index corresponding to attrs_type_key; not serialized. */
uint32_t attrs_type_index{0};
+ /*! \brief Process-local validation callback, omitted from reflection. */
+ ffi::Any validator;
+ /*! \brief Whether the validator was explicitly registered rather than generated. */
+ bool validator_is_custom{false};
TVM_DLL static void RegisterReflection();
@@ -106,13 +115,14 @@ class OpNode : public ExprNode {
* OpDef temporarily builds metadata on the same Op returned by Get. Independent
* registrations may attach different attributes; replacing an existing attribute
* requires explicit override. Cached attribute maps observe subsequent changes.
- * args_info describes required value operands; allow_extra_args permits a suffix.
- * ty_args_info describes type arguments without imposing a runtime count.
+ * args_info describes required value operands; var_args_info permits a suffix.
+ * ty_args_info describes type arguments; a typed signature installs a
+ * validator that checks their count and classes.
*
* \code
* TVM_FFI_STATIC_INIT_BLOCK() {
* OpDef("example.identity", "Return the input expression.")
- * .add_arg("value", "The input expression.")
+ * .signature(sig::arg("value", "The input expression."))
* .set_attr<bool>("FPurity", true);
* }
* // Copies the handle; the canonical node is shared.
@@ -151,12 +161,264 @@ class Op : public Expr {
TVM_DLL static Op Get(const ffi::String& op_name);
/*! \brief List registered operator names. \return Names in unspecified order. */
TVM_DLL static ffi::Array<ffi::String> ListNames();
+ /*! \brief Validate a Call with this Op's callback, when one is registered. */
+ TVM_FFI_INLINE void Validate(const CallNode* call) const {
+ const auto& validator = get()->validator;
+ if (validator != nullptr) {
+ if (TVM_FFI_PREDICT_FALSE(!call)) {
+ ThrowInvalidCall(get());
+ }
+ using View = ffi::reflection::NativeFunctionView<ffi::Expected<void>(const CallNode*)>;
+ // set_validator stores only an owning NativeFunction with this signature.
+ ffi::details::AnyUnsafe::CopyFromAnyViewAfterCheck<View>(validator)(call).value();
+ }
+ }
TVM_FFI_DEFINE_OBJECT_REF_METHODS_NOTNULLABLE(Op, Expr, OpNode);
private:
+ TVM_DLL static void ThrowInvalidCall(const OpNode* op);
TVM_DLL static ffi::List<ffi::Any> GetAttrColumn(const ffi::String& attr_name);
};
+namespace sig {
+
+/*! \brief Base counts for descriptors accepted by OpDef::signature. */
+struct SignatureTrait {
+ /*! \brief Number of required value arguments contributed by this descriptor. */
+ static constexpr size_t kArgsCount = 0;
+ /*! \brief Number of required type arguments contributed by this descriptor. */
+ static constexpr size_t kTyArgsCount = 0;
+ /*! \brief Number of variadic value tails contributed by this descriptor. */
+ static constexpr size_t kVarArgsCount = 0;
+ /*! \brief Number of variadic type tails contributed by this descriptor. */
+ static constexpr size_t kVarTyArgsCount = 0;
+ /*! \brief Number of Call.attrs descriptors contributed by this descriptor. */
+ static constexpr size_t kCallAttrsCount = 0;
+};
+
+/*! \brief Shared inputs and indexes for a typed signature validation fold. */
+struct ValidateState {
+ /*! \brief Call being validated. */
+ const CallNode* call;
+ /*! \brief Canonical operator and its current argument metadata. */
+ const OpNode* op;
+ /*! \brief Index of the next value argument after the upfront arity check. */
+ size_t args_index = 0;
+ /*! \brief Index of the next type argument after the upfront arity check. */
+ size_t ty_args_index = 0;
+};
+
+namespace details {
+
+template <typename T>
+inline bool ReportTypeMismatch(ffi::Expected<void>* out, const ValidateState* state, bool type_arg,
+ size_t index, ffi::AnyView actual) {
+ const auto& fixed = type_arg ? state->op->ty_args_info : state->op->args_info;
+ const auto& tail = type_arg ? state->op->var_ty_args_info : state->op->var_args_info;
+ std::string name;
+ if (index < fixed.size()) {
+ name = fixed[index]->name;
+ } else if (tail.has_value()) {
+ name = tail.value()->name;
+ }
+ std::string name_clause = name.empty() ? "" : " (`" + name + "`)";
+ TVMFFIAny raw = actual.CopyToTVMFFIAny();
+ *out = TVM_FFI_UNEXPECTED(TypeError)
+ << "Op `" << state->op->name << "`: `" << (type_arg ? "Call.ty_args[" : "Call.args[")
+ << index << "]`" << name_clause << " expected `" << ffi::TypeTraits<T>::TypeStr()
+ << "`, got `" << ffi::TypeTraits<T>::GetMismatchTypeInfo(&raw) << "`.";
+ return false;
+}
+
+} // namespace details
+
+/*! \brief One value argument of type T, for example `sig::arg<PrimExpr>("index")`
+ * for a required primitive index expression. */
+template <typename T = Expr>
+struct arg : SignatureTrait {
+ static_assert(std::is_base_of_v<Expr, T>, "arg<T> expects an expression view");
+ /*! \brief Contributes one required value argument. */
+ static constexpr size_t kArgsCount = 1;
+ /*! \brief Argument name recorded in operator metadata. */
+ ffi::String name;
+ /*! \brief Argument documentation recorded in operator metadata. */
+ ffi::String doc;
+
+ explicit arg(ffi::String name, ffi::String doc = "")
+ : name(std::move(name)), doc(std::move(doc)) {}
+
+ void CollectMetadata(OpNode* op) const {
+ auto info = ffi::make_object<ArgumentInfoNode>();
+ info->name = name;
+ info->doc = doc;
+ op->args_info.push_back(ArgumentInfo(std::move(info)));
+ }
+
+ static bool Validate(ValidateState* state, ffi::Expected<void>* out) {
+ if constexpr (!std::is_same_v<T, Expr>) {
+ ffi::AnyView value(state->call->args.GetArrayObj()->begin()[state->args_index]);
+ if (TVM_FFI_PREDICT_FALSE(value == nullptr ||
+ !ffi::details::AnyUnsafe::CheckAnyViewStrict<T>(value))) {
+ return details::ReportTypeMismatch<T>(out, state, false, state->args_index, value);
+ }
+ }
+ ++state->args_index;
+ return true;
+ }
+};
+
+/*! \brief Zero or more trailing value arguments of type T, for example
+ * `sig::var_args<PrimExpr>("indices")` for primitive index operands. */
+template <typename T = Expr>
+struct var_args : SignatureTrait {
+ static_assert(std::is_base_of_v<Expr, T>, "var_args<T> expects an expression view");
+ /*! \brief Contributes one variadic value tail. */
+ static constexpr size_t kVarArgsCount = 1;
+ /*! \brief Tail name recorded in operator metadata. */
+ ffi::String name;
+ /*! \brief Tail documentation recorded in operator metadata. */
+ ffi::String doc;
+
+ explicit var_args(ffi::String name, ffi::String doc = "")
+ : name(std::move(name)), doc(std::move(doc)) {}
+
+ void CollectMetadata(OpNode* op) const {
+ auto info = ffi::make_object<ArgumentInfoNode>();
+ info->name = name;
+ info->doc = doc;
+ op->var_args_info = ArgumentInfo(std::move(info));
+ }
+
+ static bool Validate(ValidateState* state, ffi::Expected<void>* out) {
+ if constexpr (!std::is_same_v<T, Expr>) {
+ const ffi::Any* values = state->call->args.GetArrayObj()->begin();
+ for (; state->args_index != state->call->args.size(); ++state->args_index) {
+ ffi::AnyView value(values[state->args_index]);
+ if (TVM_FFI_PREDICT_FALSE(value == nullptr ||
+ !ffi::details::AnyUnsafe::CheckAnyViewStrict<T>(value))) {
+ return details::ReportTypeMismatch<T>(out, state, false, state->args_index, value);
+ }
+ }
+ } else {
+ state->args_index = state->call->args.size();
+ }
+ return true;
+ }
+};
+
+/*! \brief One type argument of type T, for example `sig::ty_arg<PrimType>("dtype")`
+ * for a required primitive type. */
+template <typename T = Type>
+struct ty_arg : SignatureTrait {
+ static_assert(std::is_base_of_v<Type, T>, "ty_arg<T> expects a type view");
+ /*! \brief Contributes one required type argument. */
+ static constexpr size_t kTyArgsCount = 1;
+ /*! \brief Type-argument name recorded in operator metadata. */
+ ffi::String name;
+ /*! \brief Type-argument documentation recorded in operator metadata. */
+ ffi::String doc;
+
+ explicit ty_arg(ffi::String name, ffi::String doc = "")
+ : name(std::move(name)), doc(std::move(doc)) {}
+
+ void CollectMetadata(OpNode* op) const {
+ auto info = ffi::make_object<ArgumentInfoNode>();
+ info->name = name;
+ info->doc = doc;
+ op->ty_args_info.push_back(ArgumentInfo(std::move(info)));
+ }
+
+ static bool Validate(ValidateState* state, ffi::Expected<void>* out) {
+ if constexpr (!std::is_same_v<T, Type>) {
+ ffi::AnyView value(state->call->ty_args.GetArrayObj()->begin()[state->ty_args_index]);
+ if (TVM_FFI_PREDICT_FALSE(value == nullptr ||
+ !ffi::details::AnyUnsafe::CheckAnyViewStrict<T>(value))) {
+ return details::ReportTypeMismatch<T>(out, state, true, state->ty_args_index, value);
+ }
+ }
+ ++state->ty_args_index;
+ return true;
+ }
+};
+
+/*! \brief Zero or more trailing type arguments of type T, for example
+ * `sig::var_ty_args<PrimType>("dtypes")` for optional primitive types. */
+template <typename T = Type>
+struct var_ty_args : SignatureTrait {
+ static_assert(std::is_base_of_v<Type, T>, "var_ty_args<T> expects a type view");
+ /*! \brief Contributes one variadic type tail. */
+ static constexpr size_t kVarTyArgsCount = 1;
+ /*! \brief Tail name recorded in operator metadata. */
+ ffi::String name;
+ /*! \brief Tail documentation recorded in operator metadata. */
+ ffi::String doc;
+
+ explicit var_ty_args(ffi::String name, ffi::String doc = "")
+ : name(std::move(name)), doc(std::move(doc)) {}
+
+ void CollectMetadata(OpNode* op) const {
+ auto info = ffi::make_object<ArgumentInfoNode>();
+ info->name = name;
+ info->doc = doc;
+ op->var_ty_args_info = ArgumentInfo(std::move(info));
+ }
+
+ static bool Validate(ValidateState* state, ffi::Expected<void>* out) {
+ if constexpr (!std::is_same_v<T, Type>) {
+ const ffi::Any* values = state->call->ty_args.GetArrayObj()->begin();
+ for (; state->ty_args_index != state->call->ty_args.size(); ++state->ty_args_index) {
+ ffi::AnyView value(values[state->ty_args_index]);
+ if (TVM_FFI_PREDICT_FALSE(value == nullptr ||
+ !ffi::details::AnyUnsafe::CheckAnyViewStrict<T>(value))) {
+ return details::ReportTypeMismatch<T>(out, state, true, state->ty_args_index, value);
+ }
+ }
+ } else {
+ state->ty_args_index = state->call->ty_args.size();
+ }
+ return true;
+ }
+};
+
+/*!
+ * \brief Require Call.attrs to hold node type T and record its metadata.
+ *
+ * \code
+ * OpDef("example.with_attrs").signature(sig::call_attrs<DictAttrsNode>());
+ * // Call.attrs must contain a DictAttrsNode for this operator.
+ * \endcode
+ */
+template <typename T>
+struct call_attrs : SignatureTrait {
+ static_assert(std::is_base_of_v<AttrsNode, T>, "call_attrs<T> expects an attribute node type");
+ /*! \brief Contributes one required Call.attrs descriptor. */
+ static constexpr size_t kCallAttrsCount = 1;
+
+ void CollectMetadata(OpNode* op) const {
+ uint32_t index = T::RuntimeTypeIndex();
+ op->attrs_type_key = T::_type_key;
+ op->attrs_type_index = index;
+ }
+
+ static bool Validate(ValidateState* state, ffi::Expected<void>* out) {
+ ffi::AnyView attrs(state->call->attrs);
+ if (TVM_FFI_PREDICT_FALSE(!state->call->attrs.defined() ||
+ !ffi::details::AnyUnsafe::CheckAnyViewStrict<const T*>(attrs))) {
+ TVMFFIAny raw = attrs.CopyToTVMFFIAny();
+ std::string actual = state->call->attrs.defined()
+ ? "`" + ffi::TypeTraits<const T*>::GetMismatchTypeInfo(&raw) + "`"
+ : "None";
+ *out = TVM_FFI_UNEXPECTED(TypeError)
+ << "Op `" << state->op->name << "`: Call.attrs expected `"
+ << ffi::TypeTraits<const T*>::TypeStr() << "`, got " << actual;
+ return false;
+ }
+ return true;
+ }
+};
+
+} // namespace sig
+
/*! \brief Noncopyable temporary builder for a canonical Op; see Op for the example. */
class OpDef {
public:
@@ -181,7 +443,7 @@ class OpDef {
* \brief Append a required value-argument descriptor.
* \param name Argument name.
* \param doc Argument documentation.
- * \return This builder. Without allow_extra_args(), this is part of the complete list.
+ * \return This builder. Without a variadic descriptor, this is part of the complete list.
*/
TVM_DLL OpDef& add_arg(const ffi::String& name, const ffi::String& doc);
/*!
@@ -202,14 +464,6 @@ class OpDef {
* \return This builder. An unknown key raises an error before modifying metadata.
*/
TVM_DLL OpDef& set_attrs_type_key(const ffi::String& key);
- /*!
- * \brief Allow value arguments after the required prefix described by args_info.
- * \return This builder. An empty prefix accepts any number of value arguments.
- */
- OpDef& allow_extra_args() {
- get()->allow_extra_args = true;
- return *this;
- }
/*!
* \brief Append a descriptor for a type argument in Call.ty_args.
* \param name Type-argument name.
@@ -217,6 +471,47 @@ class OpDef {
* \return This builder. Inference retains responsibility for type-argument counts.
*/
TVM_DLL OpDef& add_ty_arg(const ffi::String& name, const ffi::String& doc);
+ /*!
+ * \brief Register a complete typed Call signature and its validator.
+ *
+ * Fixed descriptors require exactly one value/type argument each. A variadic
+ * tail accepts zero or more trailing arguments of its category. Validation
+ * throws on arity or class mismatch; it does not infer types or mutate the
+ * Call. Call construction invokes the validator unless Call::Unchecked is
+ * used. Without call_attrs<T>, Call.attrs is unconstrained. Repeated
+ * registration replaces metadata but retains any existing validator,
+ * including a generated one. Replace the callback explicitly if its
+ * executable checks must change.
+ *
+ * \code
+ * OpDef("example.op")
+ * .signature(
+ * sig::arg("value"),
+ * sig::arg<PrimExpr>("index"),
+ * sig::var_args<PrimExpr>("rest"),
+ * sig::ty_arg("T"),
+ * sig::var_ty_args("Ts"),
+ * sig::call_attrs<DictAttrsNode>());
+ * \endcode
+ */
+ template <typename... Specs>
+ OpDef& signature(const Specs&... specs) {
+ auto updated = ffi::make_object<OpNode>();
+ (ApplySignatureTrait(updated.get(), specs), ...);
+ if (op_->validator == nullptr) {
+ using View = ffi::reflection::NativeFunctionView<ffi::Expected<void>(const CallNode*)>;
+ set_validator(View::FromNative<&ValidateSignature<Specs...>>());
+ get()->validator_is_custom = false;
+ }
+ OpNode* target = get();
+ target->args_info = std::move(updated->args_info);
+ target->ty_args_info = std::move(updated->ty_args_info);
+ target->var_args_info = std::move(updated->var_args_info);
+ target->var_ty_args_info = std::move(updated->var_ty_args_info);
+ target->attrs_type_key = std::move(updated->attrs_type_key);
+ target->attrs_type_index = updated->attrs_type_index;
+ return *this;
+ }
/*!
* \brief Register an extensible attribute, rejecting duplicates unless overridden.
* \tparam ValueType The attribute value type.
@@ -230,6 +525,35 @@ class OpDef {
UpdateAttr(attr_name, ffi::Any(value), override);
return *this;
}
+ /*!
+ * \brief Register a validator that checks a Call with this operator.
+ *
+ * The callback accepts a `const CallNode*` and returns Expected<void>, with
+ * an error for invalid input. A native function pointer can be bound with
+ * NativeFunctionView::FromNative; a borrowed packed function may also be
+ * passed while it remains alive for this call. The setter retains an owning
+ * copy, so the original packed function may then be destroyed.
+ * Ordinary Call construction invokes it; Call::Unchecked skips that initial
+ * check, while Relax normalization and well-formedness may validate later.
+ * A signature installs its generated validator only if none is registered,
+ * so a custom validator registered first takes precedence. Replacing a
+ * generated or custom validator requires override; duplicate registration
+ * without override raises ValueError. NativeFunctionView cannot be null.
+ *
+ * \param validator Callback to install.
+ * \param override Whether to replace the current validator.
+ * \return This builder.
+ */
+ OpDef& set_validator(
+ ffi::reflection::NativeFunctionView<ffi::Expected<void>(const CallNode*)> validator,
+ bool override = false) {
+ TVM_FFI_CHECK(override || op_->validator == nullptr, ValueError)
+ << "Validator of " << op_->name << " is already registered";
+ get()->validator =
+ ffi::reflection::NativeFunction<ffi::Expected<void>(const CallNode*)>::From(validator);
+ get()->validator_is_custom = true;
+ return *this;
+ }
/*!
* \brief Remove the current attribute value, with no priority fallback.
* \param attr_name Attribute column name; missing columns and values are ignored.
@@ -238,6 +562,58 @@ class OpDef {
TVM_DLL OpDef& reset_attr(const ffi::String& attr_name);
private:
+ template <typename Spec>
+ static void ApplySignatureTrait(OpNode* op, const Spec& spec) {
+ static_assert(std::is_base_of_v<sig::SignatureTrait, Spec>, "Unknown signature descriptor");
+ spec.CollectMetadata(op);
+ }
+
+ template <typename... Specs>
+ TVM_FFI_INLINE static ffi::Expected<void> ValidateSignature(const CallNode* call) noexcept {
+ static_assert((std::is_base_of_v<sig::SignatureTrait, Specs> && ...),
+ "Unknown signature descriptor");
+ constexpr size_t n_args = (size_t{0} + ... + Specs::kArgsCount);
+ constexpr size_t n_ty_args = (size_t{0} + ... + Specs::kTyArgsCount);
+ constexpr size_t n_var_args = (size_t{0} + ... + Specs::kVarArgsCount);
+ constexpr size_t n_var_ty_args = (size_t{0} + ... + Specs::kVarTyArgsCount);
+ constexpr size_t n_attrs = (size_t{0} + ... + Specs::kCallAttrsCount);
+ static_assert(n_var_args <= 1 && n_var_ty_args <= 1 && n_attrs <= 1,
+ "Duplicate signature tail or call_attrs");
+ const OpNode* op = call ? call->op.as<OpNode>() : nullptr;
+ if (TVM_FFI_PREDICT_FALSE(!op)) {
+ return TVM_FFI_UNEXPECTED(TypeError) << "Invalid Op Call";
+ }
+ if (TVM_FFI_PREDICT_FALSE(!call->args.defined() || !call->ty_args.defined())) {
+ return TVM_FFI_UNEXPECTED(TypeError) << "Op `" << op->name << "`: invalid Call";
+ }
+ ffi::Expected<void> result;
+ auto* out = &result;
+ if (TVM_FFI_PREDICT_FALSE(n_var_args ? call->args.size() < n_args
+ : call->args.size() != n_args)) {
+ *out = TVM_FFI_UNEXPECTED(TypeError)
+ << "Op `" << op->name << "`: Call.args expected " << (n_var_args ? "at least " : "")
+ << n_args << (n_args == 1 ? " argument" : " arguments") << ", got "
+ << call->args.size();
+ return result;
+ }
+ if (TVM_FFI_PREDICT_FALSE(n_var_ty_args ? call->ty_args.size() < n_ty_args
+ : call->ty_args.size() != n_ty_args)) {
+ *out = TVM_FFI_UNEXPECTED(TypeError)
+ << "Op `" << op->name << "`: Call.ty_args expected "
+ << (n_var_ty_args ? "at least " : "") << n_ty_args
+ << (n_ty_args == 1 ? " type argument" : " type arguments") << ", got "
+ << call->ty_args.size();
+ return result;
+ }
+ if constexpr (sizeof...(Specs) != 0) {
+ sig::ValidateState state{call, op};
+ if (TVM_FFI_PREDICT_FALSE(!(Specs::Validate(&state, out) && ...))) {
+ return result;
+ }
+ }
+ return result;
+ }
+
OpNode* get() { return const_cast<OpNode*>(op_.operator->()); }
TVM_DLL void UpdateAttr(const ffi::String& attr_name, ffi::Any value, bool override);
Op op_;
diff --git a/include/tvm/relax/op_attr_types.h b/include/tvm/relax/op_attr_types.h
index 034b0dd2fe..51d6735534 100644
--- a/include/tvm/relax/op_attr_types.h
+++ b/include/tvm/relax/op_attr_types.h
@@ -76,12 +76,8 @@ using FInferType = ffi::TypedFunction<Type(const Call& call, const BlockBuilder&
* equivalent forms, to normalize to a single representation.
*
* Note: `FNormalize` is applied for each expression as part of the
- * `relax::BlockBuilder`. While operator-specific validation may
- * be performed within the `FNormalize` implementation, ensuring
- * that errors are caught as early as possible, this should only be
- * used when validation is fast to apply. If the validation logic
- * may be slow, it should instead be implemented in `FValidate`,
- * which is only run as part of the well-formed checker.
+ * `relax::BlockBuilder`. Operator-specific validation belongs in
+ * `OpDef::set_validator` and runs after normalization.
*
* \param bb The BlockBuilder context.
*
@@ -90,25 +86,6 @@ using FInferType = ffi::TypedFunction<Type(const Call& call, const BlockBuilder&
*/
using FNormalize = ffi::TypedFunction<Expr(const BlockBuilder& bb, Call call)>;
-/*!
- * \brief The function type of a validation function.
- *
- * A validation function is used to define constraints that should be
- * verified for an operator as part of the well-formed checker.
- *
- * Note: `FValidate` is only applied as part of the well-formed
- * checker. While this minimizes overhead while compiling Relax,
- * this delay between generating an ill-formed `relax::Call` and
- * identifying the ill-formed call may complicate debugging. If
- * the validation logic is very fast to check, and doing so would
- * not introduce a significant overhead, consider validating as part
- * of `FNormalize`, which is applied by the block builder for each
- * `relax::Call`.
- *
- * \param call The call to be validated.
- */
-using FValidate = ffi::TypedFunction<void(const Call& call)>;
-
/*! \brief The function type of a legalization function.
*
* A legalization function is used to replace a `relax::Call` with
diff --git a/python/tvm/backend/cuda/ptx/engine.py b/python/tvm/backend/cuda/ptx/engine.py
index 242ed6b759..3a6ee65bb0 100644
--- a/python/tvm/backend/cuda/ptx/engine.py
+++ b/python/tvm/backend/cuda/ptx/engine.py
@@ -94,7 +94,7 @@ def register_table(table: dict[str, InstructionEntry]) -> None:
register_op_attr(entry.op_name, "TCallEffectKind", _EFFECT_OPAQUE)
# Encoded variants have no shared named operand prefix; codegen checks
# each variant's operands, predicate, modifier tokens, and marker.
- Op.get(entry.op_name).set_allow_extra_args()
+ Op.get(entry.op_name).set_signature(var_args="operands")
# The printer name is the *surface* path a user can type, which is the
# mnemonic (several `mov_*` entries all answer to `T.ptx.mov`), not the
# table key. Reparsing re-dispatches on the operand shape.
diff --git a/python/tvm/ir/expr.py b/python/tvm/ir/expr.py
index 42b8f258da..9cfd4722d3 100644
--- a/python/tvm/ir/expr.py
+++ b/python/tvm/ir/expr.py
@@ -505,6 +505,8 @@ class Call(_CallableExprWithOp):
When ``ret_ty`` is omitted, use a missing type for subsequent normalization.
Builders may supply a known result type explicitly.
+ Operator validation runs during construction. Use :meth:`unchecked` for a
+ provisional call that will be checked after normalization.
"""
op: Expr
@@ -522,6 +524,12 @@ class Call(_CallableExprWithOp):
span: Span | None = None,
ret_ty: "tvm.ir.Type | str | None" = None,
) -> None:
+ self.__init_handle_by_constructor__(
+ _ffi_api.Call, *self._normalize_constructor_args(op, args, attrs, ty_args, span, ret_ty)
+ )
+
+ @staticmethod
+ def _normalize_constructor_args(op, args, attrs, ty_args, span, ret_ty):
# pylint: disable=import-outside-toplevel
from .attrs import DictAttrs
from .op import Op
@@ -539,7 +547,50 @@ class Call(_CallableExprWithOp):
ret_ty = PrimType(ret_ty)
if ty_args is None:
ty_args = []
- self.__init_handle_by_constructor__(_ffi_api.Call, ret_ty, op, args, attrs, ty_args, span)
+ return ret_ty, op, args, attrs, ty_args, span
+
+ @staticmethod
+ def unchecked(
+ op: Expr | str,
+ args: list[Expr] | tuple[Expr, ...],
+ attrs: "tvm.ir.Attrs | dict | None" = None,
+ ty_args: list["tvm.ir.Type"] | tuple["tvm.ir.Type", ...] | None = None,
+ span: Span | None = None,
+ ret_ty: "tvm.ir.Type | str | None" = None,
+ ) -> "Call":
+ """Construct a provisional Call without invoking its Op validator.
+
+ Use this for inputs that cannot be checked until type deduction or
+ normalization. Op validation is skipped at construction, so callers
+ must arrange a later check before relying on the Call. For Relax Calls,
+ operator-specific normalization and :func:`tvm.relax.analysis.well_formed`
+ provide validation paths. Ordinary ``Call(...)`` validates immediately.
+
+ Parameters
+ ----------
+ op : Expr or str
+ Callee expression or the name of a registered Op.
+ args : list[Expr] or tuple[Expr, ...]
+ Positional value arguments.
+ attrs : tvm.ir.Attrs, dict, or None
+ Call attributes. A dict is converted to ``DictAttrs``.
+ ty_args : list[tvm.ir.Type], tuple[tvm.ir.Type, ...], or None
+ Explicit type arguments; ``None`` gives an empty list.
+ span : Span or None
+ Source location of the Call.
+ ret_ty : tvm.ir.Type, str, or None
+ Result type. ``None`` uses ``Type.missing()`` for later inference;
+ ``"handle"`` becomes a void pointer type and other strings become
+ primitive types.
+
+ Returns
+ -------
+ Call
+ The provisional, unvalidated Call.
+ """
+ return _ffi_api.CallUnchecked(
+ *Call._normalize_constructor_args(op, args, attrs, ty_args, span, ret_ty)
+ )
@tvm_ffi.register_object("ir.TensorRegion")
diff --git a/python/tvm/ir/op.py b/python/tvm/ir/op.py
index ae5d500e13..6dc519d9b9 100644
--- a/python/tvm/ir/op.py
+++ b/python/tvm/ir/op.py
@@ -17,6 +17,8 @@
# pylint: disable=invalid-name
"""Primitive operators in the TVM IR."""
+from collections.abc import Sequence
+
import tvm_ffi
from . import _ffi_api
@@ -76,6 +78,56 @@ class Op(Expr):
"""
self._set_attr(attr_name, value, override)
+ def set_signature(self, args=(), *, ty_args=(), var_args=None, var_ty_args=None) -> None:
+ """Replace this Op's argument signature and validate Call arity.
+
+ Each entry is a name string or a ``(name, doc)`` tuple of strings.
+ Fixed entries form required prefixes; a variadic entry allows zero or
+ more additional arguments. Calls must have exactly the required count
+ without a tail, or at least that count with one. This method does not
+ check argument types or Call attrs. It replaces a generated typed
+ validator with a count-only one, while an existing custom validator
+ keeps precedence.
+
+ Parameters
+ ----------
+ args : sequence[str or tuple[str, str]], optional
+ Fixed value arguments, in order.
+ ty_args : sequence[str or tuple[str, str]], optional
+ Fixed type arguments, in order.
+ var_args : str or tuple[str, str], optional
+ Variadic value-argument tail.
+ var_ty_args : str or tuple[str, str], optional
+ Variadic type-argument tail.
+
+ Returns
+ -------
+ None
+ """
+
+ def parse_entry(entry, label):
+ if isinstance(entry, str):
+ return entry, ""
+ if (
+ isinstance(entry, tuple)
+ and len(entry) == 2
+ and all(isinstance(part, str) for part in entry)
+ ):
+ return entry
+ raise TypeError(f"{label} must be a name string or a (name, doc) tuple")
+
+ def parse_fixed(entries, label):
+ if not isinstance(entries, Sequence) or isinstance(entries, str | bytes):
+ raise TypeError(f"{label} must be a sequence of signature entries")
+ parsed = [parse_entry(entry, label) for entry in entries]
+ return [name for name, _ in parsed], [doc for _, doc in parsed]
+
+ arg_names, arg_docs = parse_fixed(args, "args")
+ ty_arg_names, ty_arg_docs = parse_fixed(ty_args, "ty_args")
+ value_tail = [] if var_args is None else list(parse_entry(var_args, "var_args"))
+ type_tail = [] if var_ty_args is None else list(parse_entry(var_ty_args, "var_ty_args"))
+ self._set_signature(arg_names, arg_docs, ty_arg_names, ty_arg_docs, value_tail, type_tail)
+
def register_op_attr(op_name, attr_key, value=None, override=False):
"""Register an operator property of an operator by name.
diff --git a/python/tvm/relax/script/ir_builder/op.py b/python/tvm/relax/script/ir_builder/op.py
index cc9d868651..925599d078 100644
--- a/python/tvm/relax/script/ir_builder/op.py
+++ b/python/tvm/relax/script/ir_builder/op.py
@@ -278,7 +278,7 @@ def call_packed(
if kwargs or not is_default:
attrs = tvm.ir.attrs.make_node(attrs_type_key, **kwargs)
- return Call(op, args, attrs=attrs, ty_args=ty_args)
+ return Call.unchecked(op, args, attrs=attrs, ty_args=ty_args)
def call_py_func(
@@ -369,7 +369,7 @@ def emit_with_type(
The created Relax Call
"""
builtin_call = tvm.ir.Op.get(op)
- return Call(builtin_call, args, attrs=None, ty_args=ty_args)
+ return Call.unchecked(builtin_call, args, attrs=None, ty_args=ty_args)
def _logical_pair(lhs, rhs, operation, primitive, python_operation):
diff --git a/src/backend/cuda/codegen/intrin_rule_cuda.cc b/src/backend/cuda/codegen/intrin_rule_cuda.cc
index f9a3f2b5ff..a8dad1d14d 100644
--- a/src/backend/cuda/codegen/intrin_rule_cuda.cc
+++ b/src/backend/cuda/codegen/intrin_rule_cuda.cc
@@ -274,10 +274,10 @@ void RegisterCudaIntrinRules() {
// Register low-level CUDA device intrinsics.
// TODO(tvm-team): consider make CUDA its own subfolder and create a file for low-level builtins.
OpDef("tirx.cuda.__shfl_sync")
- .add_arg("mask", "The thread mask.")
- .add_arg("var", "The variable to sync.")
- .add_arg("lane", "The source thread id.")
- .add_arg("width", "The warp thread width, must be a power of 2.")
+ .signature(sig::arg("mask", "The thread mask."),
+ sig::arg("var", "The variable to sync."),
+ sig::arg("lane", "The source thread id."),
+ sig::arg("width", "The warp thread width, must be a power of 2."))
.set_attr<tirx::TIRxOpCategory>("TIRxOpCategory", ffi::String("device_intrin"))
.set_attr<tirx::TDeviceIntrinsicNamespace>("TDeviceIntrinsicNamespace", ffi::String("cuda"))
.set_attr<tirx::TScriptPrinterName>("TScriptPrinterName", ffi::String("cuda.__shfl_sync"))
@@ -286,10 +286,10 @@ void RegisterCudaIntrinRules() {
.set_attr<bool>("cuda.need_warp_shuffle", true);
OpDef("tirx.cuda.__shfl_up_sync")
- .add_arg("mask", "The thread mask.")
- .add_arg("var", "The variable to sync.")
- .add_arg("delta", "The source lane id offset to be added.")
- .add_arg("width", "The warp thread width, must be a power of 2.")
+ .signature(sig::arg("mask", "The thread mask."),
+ sig::arg("var", "The variable to sync."),
+ sig::arg("delta", "The source lane id offset to be added."),
+ sig::arg("width", "The warp thread width, must be a power of 2."))
.set_attr<tirx::TIRxOpCategory>("TIRxOpCategory", ffi::String("device_intrin"))
.set_attr<tirx::TDeviceIntrinsicNamespace>("TDeviceIntrinsicNamespace", ffi::String("cuda"))
.set_attr<tirx::TScriptPrinterName>("TScriptPrinterName", ffi::String("cuda.__shfl_up_sync"))
@@ -298,10 +298,10 @@ void RegisterCudaIntrinRules() {
.set_attr<bool>("cuda.need_warp_shuffle", true);
OpDef("tirx.cuda.__shfl_down_sync")
- .add_arg("mask", "The thread mask.")
- .add_arg("var", "The variable to sync.")
- .add_arg("delta", "The source lane id offset to be subtracted.")
- .add_arg("width", "The warp thread width, must be a power of 2.")
+ .signature(sig::arg("mask", "The thread mask."),
+ sig::arg("var", "The variable to sync."),
+ sig::arg("delta", "The source lane id offset to be subtracted."),
+ sig::arg("width", "The warp thread width, must be a power of 2."))
.set_attr<tirx::TIRxOpCategory>("TIRxOpCategory", ffi::String("device_intrin"))
.set_attr<tirx::TDeviceIntrinsicNamespace>("TDeviceIntrinsicNamespace", ffi::String("cuda"))
.set_attr<tirx::TScriptPrinterName>("TScriptPrinterName",
@@ -311,10 +311,10 @@ void RegisterCudaIntrinRules() {
.set_attr<bool>("cuda.need_warp_shuffle", true);
OpDef("tirx.cuda.__shfl_xor_sync")
- .add_arg("mask", "The thread mask.")
- .add_arg("var", "The variable to sync.")
- .add_arg("lane_mask", "The lane mask.")
- .add_arg("width", "The warp thread width, must be a power of 2.")
+ .signature(sig::arg("mask", "The thread mask."),
+ sig::arg("var", "The variable to sync."),
+ sig::arg("lane_mask", "The lane mask."),
+ sig::arg("width", "The warp thread width, must be a power of 2."))
.set_attr<tirx::TIRxOpCategory>("TIRxOpCategory", ffi::String("device_intrin"))
.set_attr<tirx::TDeviceIntrinsicNamespace>("TDeviceIntrinsicNamespace", ffi::String("cuda"))
.set_attr<tirx::TScriptPrinterName>("TScriptPrinterName", ffi::String("cuda.__shfl_xor_sync"))
diff --git a/src/backend/cuda/op/iket.cc b/src/backend/cuda/op/iket.cc
index 2ffb0fab6c..aaec8a8003 100644
--- a/src/backend/cuda/op/iket.cc
+++ b/src/backend/cuda/op/iket.cc
@@ -29,8 +29,7 @@ namespace tirx {
TVM_FFI_STATIC_INIT_BLOCK() {
OpDef("tirx.cuda.iket_mark")
- .add_arg("name", "The name.")
- .allow_extra_args()
+ .signature(sig::arg("name", "The name."), sig::var_args("args"))
.set_attr<TIRxOpCategory>("TIRxOpCategory", ffi::String("device_intrin"))
.set_attr<TDeviceIntrinsicNamespace>("TDeviceIntrinsicNamespace", ffi::String("cuda"))
.set_attr<TScriptPrinterName>("TScriptPrinterName", ffi::String("cuda.iket."
@@ -38,8 +37,7 @@ TVM_FFI_STATIC_INIT_BLOCK() {
.set_attr<TCallEffectKind>("TCallEffectKind", static_cast<int64_t>(CallEffectKind::kOpaque));
OpDef("tirx.cuda.iket_range_start")
- .add_arg("name", "The name.")
- .allow_extra_args()
+ .signature(sig::arg("name", "The name."), sig::var_args("args"))
.set_attr<TIRxOpCategory>("TIRxOpCategory", ffi::String("device_intrin"))
.set_attr<TDeviceIntrinsicNamespace>("TDeviceIntrinsicNamespace", ffi::String("cuda"))
.set_attr<TScriptPrinterName>("TScriptPrinterName", ffi::String("cuda.iket."
@@ -47,8 +45,7 @@ TVM_FFI_STATIC_INIT_BLOCK() {
.set_attr<TCallEffectKind>("TCallEffectKind", static_cast<int64_t>(CallEffectKind::kOpaque));
OpDef("tirx.cuda.iket_range_end")
- .add_arg("token", "The token.")
- .allow_extra_args()
+ .signature(sig::arg("token", "The token."), sig::var_args("args"))
.set_attr<TIRxOpCategory>("TIRxOpCategory", ffi::String("device_intrin"))
.set_attr<TDeviceIntrinsicNamespace>("TDeviceIntrinsicNamespace", ffi::String("cuda"))
.set_attr<TScriptPrinterName>("TScriptPrinterName", ffi::String("cuda.iket."
@@ -56,8 +53,7 @@ TVM_FFI_STATIC_INIT_BLOCK() {
.set_attr<TCallEffectKind>("TCallEffectKind", static_cast<int64_t>(CallEffectKind::kOpaque));
OpDef("tirx.cuda.iket_range_push")
- .add_arg("name", "The name.")
- .allow_extra_args()
+ .signature(sig::arg("name", "The name."), sig::var_args("args"))
.set_attr<TIRxOpCategory>("TIRxOpCategory", ffi::String("device_intrin"))
.set_attr<TDeviceIntrinsicNamespace>("TDeviceIntrinsicNamespace", ffi::String("cuda"))
.set_attr<TScriptPrinterName>("TScriptPrinterName", ffi::String("cuda.iket."
@@ -72,7 +68,7 @@ TVM_FFI_STATIC_INIT_BLOCK() {
.set_attr<TCallEffectKind>("TCallEffectKind", static_cast<int64_t>(CallEffectKind::kOpaque));
OpDef("tirx.cuda.iket_sentinel_token")
- .add_arg("name", "The name.")
+ .signature(sig::arg("name", "The name."))
.set_attr<TIRxOpCategory>("TIRxOpCategory", ffi::String("device_intrin"))
.set_attr<TDeviceIntrinsicNamespace>("TDeviceIntrinsicNamespace", ffi::String("cuda"))
.set_attr<TScriptPrinterName>("TScriptPrinterName", ffi::String("cuda.iket."
@@ -80,9 +76,8 @@ TVM_FFI_STATIC_INIT_BLOCK() {
.set_attr<TCallEffectKind>("TCallEffectKind", static_cast<int64_t>(CallEffectKind::kOpaque));
OpDef("tirx.cuda.iket_official_event")
- .add_arg("event_id", "The event identifier.")
- .add_arg("source_code", "The source code.")
- .allow_extra_args()
+ .signature(sig::arg("event_id", "The event identifier."),
+ sig::arg("source_code", "The source code."), sig::var_args("args"))
.set_attr<TIRxOpCategory>("TIRxOpCategory", ffi::String("device_intrin"))
.set_attr<TDeviceIntrinsicNamespace>("TDeviceIntrinsicNamespace", ffi::String("cuda"))
.set_attr<TScriptPrinterName>("TScriptPrinterName", ffi::String("cuda.iket."
diff --git a/src/backend/cuda/op/target_builtin.cc b/src/backend/cuda/op/target_builtin.cc
index f30749370c..2011ab9481 100644
--- a/src/backend/cuda/op/target_builtin.cc
+++ b/src/backend/cuda/op/target_builtin.cc
@@ -31,6 +31,7 @@
#include <initializer_list>
#include <string>
+#include <utility>
namespace tvm {
namespace tirx {
@@ -46,65 +47,48 @@ void RegisterCudaTargetBuiltins() {
registered = true;
OpDef("tirx.tvm_load_matrix_sync")
- .add_arg("fragment", "The matrix fragment.")
- .add_arg("m", "The M dimension.")
- .add_arg("n", "The N dimension.")
- .add_arg("k", "The K dimension.")
- .add_arg("index", "The index.")
- .add_arg("buffer_ptr", "The buffer pointer.")
- .add_arg("stride", "The stride.")
- .add_arg("layout", "The layout.")
+ .signature(sig::arg("fragment", "The matrix fragment."), sig::arg("m", "The M dimension."),
+ sig::arg("n", "The N dimension."), sig::arg("k", "The K dimension."),
+ sig::arg("index", "The index."), sig::arg("buffer_ptr", "The buffer pointer."),
+ sig::arg("stride", "The stride."), sig::arg("layout", "The layout."))
.set_attr<TScriptPrinterName>("TScriptPrinterName", ffi::String("tvm_load_matrix_sync"))
.set_attr<TIRxOpCategory>("TIRxOpCategory", ffi::String("builtin"))
.set_attr<TCallEffectKind>("TCallEffectKind",
static_cast<int64_t>(CallEffectKind::kReadState));
OpDef("tirx.tvm_mma_sync")
- .add_arg("fragment_d", "The D fragment.")
- .add_arg("index_d", "The D fragment index.")
- .add_arg("fragment_a", "The A fragment.")
- .add_arg("index_a", "The A fragment index.")
- .add_arg("fragment_b", "The B fragment.")
- .add_arg("index_b", "The B fragment index.")
- .add_arg("fragment_c", "The C fragment.")
- .add_arg("index_c", "The C fragment index.")
+ .signature(
+ sig::arg("fragment_d", "The D fragment."), sig::arg("index_d", "The D fragment index."),
+ sig::arg("fragment_a", "The A fragment."), sig::arg("index_a", "The A fragment index."),
+ sig::arg("fragment_b", "The B fragment."), sig::arg("index_b", "The B fragment index."),
+ sig::arg("fragment_c", "The C fragment."), sig::arg("index_c", "The C fragment index."))
.set_attr<TScriptPrinterName>("TScriptPrinterName", ffi::String("tvm_mma_sync"))
.set_attr<TIRxOpCategory>("TIRxOpCategory", ffi::String("builtin"))
.set_attr<TCallEffectKind>("TCallEffectKind", static_cast<int64_t>(CallEffectKind::kOpaque));
OpDef("tirx.tvm_bmma_sync")
- .add_arg("fragment_d", "The D fragment.")
- .add_arg("index_d", "The D fragment index.")
- .add_arg("fragment_a", "The A fragment.")
- .add_arg("index_a", "The A fragment index.")
- .add_arg("fragment_b", "The B fragment.")
- .add_arg("index_b", "The B fragment index.")
- .add_arg("fragment_c", "The C fragment.")
- .add_arg("index_c", "The C fragment index.")
+ .signature(
+ sig::arg("fragment_d", "The D fragment."), sig::arg("index_d", "The D fragment index."),
+ sig::arg("fragment_a", "The A fragment."), sig::arg("index_a", "The A fragment index."),
+ sig::arg("fragment_b", "The B fragment."), sig::arg("index_b", "The B fragment index."),
+ sig::arg("fragment_c", "The C fragment."), sig::arg("index_c", "The C fragment index."))
.set_attr<TScriptPrinterName>("TScriptPrinterName", ffi::String("tvm_bmma_sync"))
.set_attr<TIRxOpCategory>("TIRxOpCategory", ffi::String("builtin"))
.set_attr<TCallEffectKind>("TCallEffectKind", static_cast<int64_t>(CallEffectKind::kOpaque));
OpDef("tirx.tvm_fill_fragment")
- .add_arg("fragment", "The matrix fragment.")
- .add_arg("m", "The M dimension.")
- .add_arg("n", "The N dimension.")
- .add_arg("k", "The K dimension.")
- .add_arg("index", "The index.")
- .add_arg("value", "The value to use.")
+ .signature(sig::arg("fragment", "The matrix fragment."), sig::arg("m", "The M dimension."),
+ sig::arg("n", "The N dimension."), sig::arg("k", "The K dimension."),
+ sig::arg("index", "The index."), sig::arg("value", "The value to use."))
.set_attr<TScriptPrinterName>("TScriptPrinterName", ffi::String("tvm_fill_fragment"))
.set_attr<TIRxOpCategory>("TIRxOpCategory", ffi::String("builtin"))
.set_attr<TCallEffectKind>("TCallEffectKind", static_cast<int64_t>(CallEffectKind::kOpaque));
OpDef("tirx.tvm_store_matrix_sync")
- .add_arg("fragment", "The matrix fragment.")
- .add_arg("m", "The M dimension.")
- .add_arg("n", "The N dimension.")
- .add_arg("k", "The K dimension.")
- .add_arg("index", "The index.")
- .add_arg("buffer_ptr", "The buffer pointer.")
- .add_arg("stride", "The stride.")
- .add_arg("layout", "The layout.")
+ .signature(sig::arg("fragment", "The matrix fragment."), sig::arg("m", "The M dimension."),
+ sig::arg("n", "The N dimension."), sig::arg("k", "The K dimension."),
+ sig::arg("index", "The index."), sig::arg("buffer_ptr", "The buffer pointer."),
+ sig::arg("stride", "The stride."), sig::arg("layout", "The layout."))
.set_attr<TScriptPrinterName>("TScriptPrinterName", ffi::String("tvm_store_matrix_sync"))
.set_attr<TIRxOpCategory>("TIRxOpCategory", ffi::String("builtin"))
.set_attr<TCallEffectKind>("TCallEffectKind", static_cast<int64_t>(CallEffectKind::kOpaque));
@@ -115,29 +99,25 @@ void RegisterCudaTargetBuiltins() {
// to its thread-local index. Used by the s_tir tensor_intrin tensorize
// path so per-thread fragment offsets stay element-accurate.
OpDef("tirx.mma_store_legacy")
- .add_arg("m", "The M dimension.")
- .add_arg("n", "The N dimension.")
- .add_arg("dst_ptr", "The destination pointer.")
- .add_arg("src_ptr", "The source pointer.")
- .add_arg("src_offset", "The source offset.")
- .add_arg("dst_stride", "The destination stride.")
+ .signature(sig::arg("m", "The M dimension."), sig::arg("n", "The N dimension."),
+ sig::arg("dst_ptr", "The destination pointer."),
+ sig::arg("src_ptr", "The source pointer."),
+ sig::arg("src_offset", "The source offset."),
+ sig::arg("dst_stride", "The destination stride."))
.set_attr<TScriptPrinterName>("TScriptPrinterName", ffi::String("cuda.mma_store_legacy"))
.set_attr<TIRxOpCategory>("TIRxOpCategory", ffi::String("builtin"))
.set_attr<TCallEffectKind>("TCallEffectKind", static_cast<int64_t>(CallEffectKind::kOpaque));
OpDef("tirx.mma_fill_legacy")
- .add_arg("local_size", "The local allocation size.")
- .add_arg("local_ptr", "The local pointer.")
- .add_arg("offset", "The offset.")
+ .signature(sig::arg("local_size", "The local allocation size."),
+ sig::arg("local_ptr", "The local pointer."), sig::arg("offset", "The offset."))
.set_attr<TScriptPrinterName>("TScriptPrinterName", ffi::String("cuda.mma_fill_legacy"))
.set_attr<TIRxOpCategory>("TIRxOpCategory", ffi::String("builtin"))
.set_attr<TCallEffectKind>("TCallEffectKind", static_cast<int64_t>(CallEffectKind::kOpaque));
OpDef("tirx.s_tir.ldg32")
- .add_arg("reg", "The register.")
- .add_arg("guard", "The guard predicate.")
- .add_arg("addr", "The address.")
- .add_arg("local_addr", "The local address.")
+ .signature(sig::arg("reg", "The register."), sig::arg("guard", "The guard predicate."),
+ sig::arg("addr", "The address."), sig::arg("local_addr", "The local address."))
.set_attr<TCallEffectKind>("TCallEffectKind", static_cast<int64_t>(CallEffectKind::kPure))
.set_attr<TScriptPrinterName>("TScriptPrinterName", ffi::String("s_tir.ldg32"))
.set_attr<TIRxOpCategory>("TIRxOpCategory", ffi::String("device_intrin"))
@@ -147,12 +127,11 @@ void RegisterCudaTargetBuiltins() {
// the T.s_tir.cp_async_raw.legacy 6-arg surface). It carries the element dtype in Call.dtype
// and prints it dtype-first; user-issued copies go through T.ptx instead.
OpDef("tirx.s_tir.cp_async_raw")
- .add_arg("dst_ptr", "The destination pointer.")
- .add_arg("dst_offset", "The destination offset.")
- .add_arg("src_ptr", "The source pointer.")
- .add_arg("src_offset", "The source offset.")
- .add_arg("cp_size", "The copy size.")
- .allow_extra_args()
+ .signature(sig::arg("dst_ptr", "The destination pointer."),
+ sig::arg("dst_offset", "The destination offset."),
+ sig::arg("src_ptr", "The source pointer."),
+ sig::arg("src_offset", "The source offset."),
+ sig::arg("cp_size", "The copy size."), sig::var_args("args"))
.set_attr<TCallEffectKind>("TCallEffectKind", static_cast<int64_t>(CallEffectKind::kOpaque))
.set_attr<TIRxOpCategory>("TIRxOpCategory", ffi::String("device_intrin"))
.set_attr<TDeviceIntrinsicNamespace>("TDeviceIntrinsicNamespace", ffi::String("s_tir"))
@@ -161,12 +140,11 @@ void RegisterCudaTargetBuiltins() {
static_cast<int64_t>(ScriptDtypePrintLocation::kFirst));
OpDef("tirx.mma_store")
- .add_arg("m", "The M dimension.")
- .add_arg("n", "The N dimension.")
- .add_arg("dst_ptr", "The destination pointer.")
- .add_arg("src_ptr", "The source pointer.")
- .add_arg("src_offset", "The source offset.")
- .add_arg("dst_stride", "The destination stride.")
+ .signature(sig::arg("m", "The M dimension."), sig::arg("n", "The N dimension."),
+ sig::arg("dst_ptr", "The destination pointer."),
+ sig::arg("src_ptr", "The source pointer."),
+ sig::arg("src_offset", "The source offset."),
+ sig::arg("dst_stride", "The destination stride."))
.set_attr<TScriptPrinterName>("TScriptPrinterName", ffi::String("cuda.mma_store"))
.set_attr<TIRxOpCategory>("TIRxOpCategory", ffi::String("builtin"))
.set_attr<TCallEffectKind>("TCallEffectKind", static_cast<int64_t>(CallEffectKind::kOpaque))
@@ -174,9 +152,8 @@ void RegisterCudaTargetBuiltins() {
static_cast<int64_t>(ScriptDtypePrintLocation::kFirst));
OpDef("tirx.mma_fill")
- .add_arg("local_size", "The local allocation size.")
- .add_arg("local_ptr", "The local pointer.")
- .add_arg("offset", "The offset.")
+ .signature(sig::arg("local_size", "The local allocation size."),
+ sig::arg("local_ptr", "The local pointer."), sig::arg("offset", "The offset."))
.set_attr<TScriptPrinterName>("TScriptPrinterName", ffi::String("cuda.mma_fill"))
.set_attr<TIRxOpCategory>("TIRxOpCategory", ffi::String("builtin"))
.set_attr<TCallEffectKind>("TCallEffectKind", static_cast<int64_t>(CallEffectKind::kOpaque))
@@ -184,43 +161,43 @@ void RegisterCudaTargetBuiltins() {
static_cast<int64_t>(ScriptDtypePrintLocation::kFirst));
OpDef("tirx.timer_init_cuda")
- .add_arg("profiler_buffer", "The profiler buffer.")
- .add_arg("profiler_tag", "The profiler tag.")
- .add_arg("profiler_write_offset", "The profiler write offset.")
- .add_arg("num_groups", "The number of groups.")
- .add_arg("group_id", "The group identifier.")
+ .signature(sig::arg("profiler_buffer", "The profiler buffer."),
+ sig::arg("profiler_tag", "The profiler tag."),
+ sig::arg("profiler_write_offset", "The profiler write offset."),
+ sig::arg("num_groups", "The number of groups."),
+ sig::arg("group_id", "The group identifier."))
.set_attr<TScriptPrinterName>("TScriptPrinterName", ffi::String("cuda.timer_init"))
.set_attr<TIRxOpCategory>("TIRxOpCategory", ffi::String("builtin"))
.set_attr<TCallEffectKind>("TCallEffectKind", static_cast<int64_t>(CallEffectKind::kOpaque));
OpDef("tirx.timer_start_cuda")
- .add_arg("event_type", "The event type.")
- .add_arg("profiler_buffer", "The profiler buffer.")
- .add_arg("profiler_tag", "The profiler tag.")
- .add_arg("profiler_write_offset", "The profiler write offset.")
- .add_arg("profiler_write_stride", "The profiler write stride.")
- .add_arg("leader_cond", "The leader condition.")
+ .signature(sig::arg("event_type", "The event type."),
+ sig::arg("profiler_buffer", "The profiler buffer."),
+ sig::arg("profiler_tag", "The profiler tag."),
+ sig::arg("profiler_write_offset", "The profiler write offset."),
+ sig::arg("profiler_write_stride", "The profiler write stride."),
+ sig::arg("leader_cond", "The leader condition."))
.set_attr<TScriptPrinterName>("TScriptPrinterName", ffi::String("cuda.timer_start"))
.set_attr<TIRxOpCategory>("TIRxOpCategory", ffi::String("builtin"))
.set_attr<TCallEffectKind>("TCallEffectKind", static_cast<int64_t>(CallEffectKind::kOpaque));
OpDef("tirx.timer_end_cuda")
- .add_arg("event_type", "The event type.")
- .add_arg("profiler_buffer", "The profiler buffer.")
- .add_arg("profiler_tag", "The profiler tag.")
- .add_arg("profiler_write_offset", "The profiler write offset.")
- .add_arg("profiler_write_stride", "The profiler write stride.")
- .add_arg("leader_cond", "The leader condition.")
+ .signature(sig::arg("event_type", "The event type."),
+ sig::arg("profiler_buffer", "The profiler buffer."),
+ sig::arg("profiler_tag", "The profiler tag."),
+ sig::arg("profiler_write_offset", "The profiler write offset."),
+ sig::arg("profiler_write_stride", "The profiler write stride."),
+ sig::arg("leader_cond", "The leader condition."))
.set_attr<TScriptPrinterName>("TScriptPrinterName", ffi::String("cuda.timer_end"))
.set_attr<TIRxOpCategory>("TIRxOpCategory", ffi::String("builtin"))
.set_attr<TCallEffectKind>("TCallEffectKind", static_cast<int64_t>(CallEffectKind::kOpaque));
OpDef("tirx.timer_finalize_cuda")
- .add_arg("profiler_buffer", "The profiler buffer.")
- .add_arg("profiler_tag", "The profiler tag.")
- .add_arg("profiler_write_offset", "The profiler write offset.")
- .add_arg("profiler_write_stride", "The profiler write stride.")
- .add_arg("leader_cond", "The leader condition.")
+ .signature(sig::arg("profiler_buffer", "The profiler buffer."),
+ sig::arg("profiler_tag", "The profiler tag."),
+ sig::arg("profiler_write_offset", "The profiler write offset."),
+ sig::arg("profiler_write_stride", "The profiler write stride."),
+ sig::arg("leader_cond", "The leader condition."))
.set_attr<TScriptPrinterName>("TScriptPrinterName", ffi::String("cuda.timer_finalize"))
.set_attr<TIRxOpCategory>("TIRxOpCategory", ffi::String("builtin"))
.set_attr<TCallEffectKind>("TCallEffectKind", static_cast<int64_t>(CallEffectKind::kOpaque));
@@ -230,25 +207,22 @@ void RegisterCudaTargetBuiltins() {
namespace {
-struct DeviceIntrinsicRegistration {
- const char* name;
- const char* namespace_name;
- CallEffectKind effect_kind;
- std::initializer_list<const char*> args;
- bool allow_extra_args;
+struct DeviceIntrinsicNames {
+ std::string canonical;
+ std::string printer;
};
-void RegisterDeviceIntrinsic(const DeviceIntrinsicRegistration& reg) {
- std::string name(reg.name);
- std::string namespace_name(reg.namespace_name);
+TVM_FFI_NO_INLINE DeviceIntrinsicNames MakeDeviceIntrinsicNames(const char* op_name,
+ const char* op_namespace) {
+ std::string name(op_name);
+ std::string namespace_name(op_namespace);
std::string prefix = namespace_name + "_";
std::string suffix = name;
if (suffix.rfind(prefix, 0) == 0) {
suffix = suffix.substr(prefix.size());
}
- std::string canonical_op_name = "tirx." + namespace_name + "." + suffix;
- ffi::String namespace_attr(namespace_name);
+ std::string canonical = "tirx." + namespace_name + "." + suffix;
// Match the nested construction namespaces at the canonical registration site.
if (namespace_name == "cuda" &&
(suffix.rfind("tcgen05_", 0) == 0 || suffix.rfind("wgmma_", 0) == 0)) {
@@ -258,144 +232,182 @@ void RegisterDeviceIntrinsic(const DeviceIntrinsicRegistration& reg) {
(suffix.size() >= 5 && suffix.compare(suffix.size() - 5, 5, "_warp") == 0))) {
suffix[suffix.rfind('_')] = '.';
}
- ffi::String printer_name(namespace_name + "." + suffix);
- int64_t effect = static_cast<int64_t>(reg.effect_kind);
+ return {std::move(canonical), namespace_name + "." + suffix};
+}
- OpDef def(canonical_op_name);
- for (const char* name : reg.args) {
- def.add_arg(name, "");
- }
- if (reg.allow_extra_args) {
- def.allow_extra_args();
- }
+TVM_FFI_NO_INLINE void RegisterDeviceIntrinsicAttrs(OpDef& def, const char* op_namespace,
+ CallEffectKind effect_kind,
+ const std::string& printer_name) {
def.set_attr<TIRxOpCategory>("TIRxOpCategory", ffi::String("device_intrin"))
- .set_attr<TDeviceIntrinsicNamespace>("TDeviceIntrinsicNamespace", namespace_attr)
- .set_attr<TCallEffectKind>("TCallEffectKind", effect)
- .set_attr<TScriptPrinterName>("TScriptPrinterName", printer_name);
+ .set_attr<TDeviceIntrinsicNamespace>("TDeviceIntrinsicNamespace", ffi::String(op_namespace))
+ .set_attr<TCallEffectKind>("TCallEffectKind", static_cast<int64_t>(effect_kind))
+ .set_attr<TScriptPrinterName>("TScriptPrinterName", ffi::String(printer_name));
}
-#define TIRX_DEVICE_INTRIN_ALIAS(OpName, Namespace, EffectKind, ...) \
- {#OpName, #Namespace, CallEffectKind::EffectKind, __VA_ARGS__}
-
-const DeviceIntrinsicRegistration kDeviceIntrinsics[] = {
- TIRX_DEVICE_INTRIN_ALIAS(cuda_any_sync, cuda, kPure, {"mask", "pred"}, false),
- TIRX_DEVICE_INTRIN_ALIAS(cuda_atomic_add, cuda, kOpaque, {"res_addr", "value"}, false),
- TIRX_DEVICE_INTRIN_ALIAS(cuda_atomic_cas, cuda, kOpaque, {"ptr", "old_val", "new_val"}, false),
- TIRX_DEVICE_INTRIN_ALIAS(
- cuda_wait_until, cuda, kOpaque,
- {"dst", "ptr", "condition", "scope", "space", "ptx_type", "backoff_ns"}, false),
- TIRX_DEVICE_INTRIN_ALIAS(cuda_ballot_sync, cuda, kOpaque, {"mask", "pred"}, false),
- TIRX_DEVICE_INTRIN_ALIAS(cuda_bfloat1622float2, cuda, kOpaque, {"packed"}, false),
- TIRX_DEVICE_INTRIN_ALIAS(cuda_bfloat162float, cuda, kOpaque, {"src"}, false),
- TIRX_DEVICE_INTRIN_ALIAS(cuda_clock64, cuda, kOpaque, {}, false),
- TIRX_DEVICE_INTRIN_ALIAS(cuda_cluster_sync, cuda, kOpaque, {}, false),
- TIRX_DEVICE_INTRIN_ALIAS(cuda_cta_reduce, cuda, kOpaque,
- {"value", "op", "num_warps", "scratch"}, false),
- TIRX_DEVICE_INTRIN_ALIAS(cuda_cta_sync, cuda, kOpaque, {}, false),
- TIRX_DEVICE_INTRIN_ALIAS(cuda_cvta_generic_to_shared, cuda, kOpaque, {"ptr"}, false),
- TIRX_DEVICE_INTRIN_ALIAS(cuda_elect_sync, cuda, kOpaque, {}, false),
- TIRX_DEVICE_INTRIN_ALIAS(cuda_fadd2_rn, cuda, kOpaque, {"a", "b"}, false),
- TIRX_DEVICE_INTRIN_ALIAS(cuda_fdividef, cuda, kPure, {"x", "y"}, false),
- TIRX_DEVICE_INTRIN_ALIAS(cuda_ffs_u32, cuda, kOpaque, {"value"}, false),
- TIRX_DEVICE_INTRIN_ALIAS(cuda_float22bfloat162_rn, cuda, kOpaque, {"v0", "v1"}, false),
- TIRX_DEVICE_INTRIN_ALIAS(cuda_float22bfloat162_rn_from_float2, cuda, kOpaque, {"packed"},
- false),
- TIRX_DEVICE_INTRIN_ALIAS(cuda_float22half2, cuda, kOpaque, {"dst", "src"}, false),
- TIRX_DEVICE_INTRIN_ALIAS(cuda_float2_x, cuda, kOpaque, {"packed"}, false),
- TIRX_DEVICE_INTRIN_ALIAS(cuda_float2_y, cuda, kOpaque, {"packed"}, false),
- TIRX_DEVICE_INTRIN_ALIAS(cuda_float8tohalf8, cuda, kOpaque, {"src_addr", "dst_addr"}, false),
- TIRX_DEVICE_INTRIN_ALIAS(cuda_float_as_uint, cuda, kOpaque, {"x"}, false),
- TIRX_DEVICE_INTRIN_ALIAS(cuda_fmul2_rn, cuda, kOpaque, {"a", "b"}, false),
- TIRX_DEVICE_INTRIN_ALIAS(cuda_fp8x4_e4m3_from_float4, cuda, kOpaque, {"x", "y", "z", "w"},
- false),
- TIRX_DEVICE_INTRIN_ALIAS(cuda_func_call, cuda, kOpaque, {"func_name"}, true),
- TIRX_DEVICE_INTRIN_ALIAS(cuda_get_tmem_addr, cuda, kOpaque,
- {"addr", "row_offset", "col_offset"}, false),
- TIRX_DEVICE_INTRIN_ALIAS(cuda_grid_sync, cuda, kOpaque, {}, false),
- TIRX_DEVICE_INTRIN_ALIAS(cuda_half2float, cuda, kOpaque, {"src"}, false),
- TIRX_DEVICE_INTRIN_ALIAS(cuda_half8tofloat8, cuda, kOpaque, {"src_addr", "dst_addr"}, false),
- TIRX_DEVICE_INTRIN_ALIAS(cuda_hmax2, cuda, kOpaque, {"a", "b"}, false),
- TIRX_DEVICE_INTRIN_ALIAS(cuda_hmin2, cuda, kOpaque, {"a", "b"}, false),
- TIRX_DEVICE_INTRIN_ALIAS(cuda_ldg, cuda, kOpaque, {}, true),
- TIRX_DEVICE_INTRIN_ALIAS(cuda_make_float2, cuda, kOpaque, {"x", "y"}, false),
- TIRX_DEVICE_INTRIN_ALIAS(cuda_mbarrier_wait, cuda, kOpaque, {"bar", "phase"}, false),
- TIRX_DEVICE_INTRIN_ALIAS(cuda_mbarrier_wait_acquire_cluster, cuda, kOpaque, {"bar", "phase"},
- false),
- TIRX_DEVICE_INTRIN_ALIAS(cuda_mov_sreg, cuda, kPure, {"bits", "reg_name"}, false),
- TIRX_DEVICE_INTRIN_ALIAS(cuda_nano_sleep, cuda, kOpaque, {"time"}, false),
- TIRX_DEVICE_INTRIN_ALIAS(cuda_printf, cuda, kOpaque, {"fmt"}, true),
- TIRX_DEVICE_INTRIN_ALIAS(cuda_reduce_add_sync_u32, cuda, kOpaque, {"mask", "value"}, false),
- TIRX_DEVICE_INTRIN_ALIAS(cuda_reduce_min_sync_u32, cuda, kOpaque, {"mask", "value"}, false),
- TIRX_DEVICE_INTRIN_ALIAS(cuda_runtime_instr_desc, cuda, kOpaque, {"desc", "sf_id"}, false),
- TIRX_DEVICE_INTRIN_ALIAS(cuda_sm100_2sm_leader_smem_addr, cuda, kOpaque, {"ptr"}, false),
- TIRX_DEVICE_INTRIN_ALIAS(cuda_smem_addr_from_uint64, cuda, kOpaque, {"cluster_addr"}, false),
- TIRX_DEVICE_INTRIN_ALIAS(cuda_syncthreads_and, cuda, kOpaque, {"cond"}, false),
- TIRX_DEVICE_INTRIN_ALIAS(cuda_syncthreads_or, cuda, kOpaque, {"cond"}, false),
- TIRX_DEVICE_INTRIN_ALIAS(cuda_tcgen05_encode_instr_descriptor, cuda, kOpaque,
- {"desc", "d_dtype", "a_dtype", "b_dtype", "M", "N", "K", "trans_a",
- "trans_b", "n_cta_groups", "neg_a", "neg_b", "sat_d", "is_sparse"},
- false),
- TIRX_DEVICE_INTRIN_ALIAS(cuda_tcgen05_encode_instr_descriptor_block_scaled, cuda, kOpaque,
- {"desc", "d_dtype", "a_dtype", "b_dtype", "sfa_dtype", "sfb_dtype",
- "sfa_tmem_addr", "sfb_tmem_addr", "M", "N", "K", "trans_a", "trans_b",
- "n_cta_groups", "neg_a", "neg_b", "is_sparse"},
- false),
- TIRX_DEVICE_INTRIN_ALIAS(cuda_tcgen05_encode_matrix_descriptor, cuda, kOpaque,
- {"desc", "addr", "ldo", "sdo", "swizzle"}, false),
- TIRX_DEVICE_INTRIN_ALIAS(cuda_thread_fence, cuda, kOpaque, {}, false),
- TIRX_DEVICE_INTRIN_ALIAS(cuda_thread_rank, cuda, kPure, {}, false),
- TIRX_DEVICE_INTRIN_ALIAS(cuda_trap_when_assert_failed, cuda, kOpaque, {"cond"}, false),
- TIRX_DEVICE_INTRIN_ALIAS(cuda_uint_as_float, cuda, kOpaque, {"bits"}, false),
- TIRX_DEVICE_INTRIN_ALIAS(cuda_warp_reduce, cuda, kOpaque, {"value", "op", "width"}, false),
- TIRX_DEVICE_INTRIN_ALIAS(cuda_warp_sync, cuda, kOpaque, {}, false),
- TIRX_DEVICE_INTRIN_ALIAS(cuda_warpgroup_sync, cuda, kOpaque, {"bar_no"}, false),
- TIRX_DEVICE_INTRIN_ALIAS(cuda_wgmma_encode_matrix_descriptor, cuda, kOpaque,
- {"desc", "addr", "ldo", "sdo", "swizzle"}, false),
- TIRX_DEVICE_INTRIN_ALIAS(cuda_wgmma_noop_barrier, cuda, kOpaque, {"reg"}, false),
- TIRX_DEVICE_INTRIN_ALIAS(nvshmem_barrier_all, nvshmem, kOpaque, {}, false),
- TIRX_DEVICE_INTRIN_ALIAS(nvshmem_fence, nvshmem, kOpaque, {}, false),
- TIRX_DEVICE_INTRIN_ALIAS(nvshmem_getmem_nbi, nvshmem, kOpaque, {"dst", "src", "nelems", "pe"},
- false),
- TIRX_DEVICE_INTRIN_ALIAS(nvshmem_getmem_nbi_block, nvshmem, kOpaque,
- {"dst", "src", "nelems", "pe"}, false),
- TIRX_DEVICE_INTRIN_ALIAS(nvshmem_getmem_nbi_warp, nvshmem, kOpaque,
- {"dst", "src", "nelems", "pe"}, false),
- TIRX_DEVICE_INTRIN_ALIAS(nvshmem_my_pe, nvshmem, kOpaque, {}, false),
- TIRX_DEVICE_INTRIN_ALIAS(nvshmem_n_pes, nvshmem, kOpaque, {}, false),
- TIRX_DEVICE_INTRIN_ALIAS(nvshmem_putmem_nbi, nvshmem, kOpaque, {"dst", "src", "nelems", "pe"},
- false),
- TIRX_DEVICE_INTRIN_ALIAS(nvshmem_putmem_nbi_block, nvshmem, kOpaque,
- {"dst", "src", "nelems", "pe"}, false),
- TIRX_DEVICE_INTRIN_ALIAS(nvshmem_putmem_nbi_warp, nvshmem, kOpaque,
- {"dst", "src", "nelems", "pe"}, false),
- TIRX_DEVICE_INTRIN_ALIAS(nvshmem_putmem_signal_nbi, nvshmem, kOpaque,
- {"dst", "src", "nelems", "sig_addr", "signal", "sig_op", "pe"}, false),
- TIRX_DEVICE_INTRIN_ALIAS(nvshmem_putmem_signal_nbi_block, nvshmem, kOpaque,
- {"dst", "src", "nelems", "sig_addr", "signal", "sig_op", "pe"}, false),
- TIRX_DEVICE_INTRIN_ALIAS(nvshmem_putmem_signal_nbi_warp, nvshmem, kOpaque,
- {"dst", "src", "nelems", "sig_addr", "signal", "sig_op", "pe"}, false),
- TIRX_DEVICE_INTRIN_ALIAS(nvshmem_quiet, nvshmem, kOpaque, {}, false),
- TIRX_DEVICE_INTRIN_ALIAS(nvshmem_signal_op, nvshmem, kOpaque,
- {"sig_addr", "signal", "sig_op", "pe"}, false),
- TIRX_DEVICE_INTRIN_ALIAS(nvshmem_wait_until, nvshmem, kOpaque,
- {"ivar", "cmp", "cmp_value", "type"}, false),
- TIRX_DEVICE_INTRIN_ALIAS(
- ptx_legacy_ldmatrix,
- ptx_legacy, kOpaque,
- {"trans", "num", "dtype", "local_ptr", "local_offset", "smem_ptr", "smem_offset"}, false),
- TIRX_DEVICE_INTRIN_ALIAS(
- ptx_legacy_mma, ptx_legacy, kOpaque,
- {"shape", "a_layout", "b_layout", "a_dtype", "b_dtype", "c_dtype", "a_ptr", "a_offset",
- "b_ptr", "b_offset", "acc_ptr", "c_offset", "saturate"},
- true),
-};
-
-#undef TIRX_DEVICE_INTRIN_ALIAS
+template <typename... Specs>
+void RegisterDeviceIntrinsic(const char* op_name, const char* op_namespace,
+ CallEffectKind effect_kind, const Specs&... specs) {
+ DeviceIntrinsicNames names = MakeDeviceIntrinsicNames(op_name, op_namespace);
+ OpDef def(names.canonical);
+ def.signature(specs...);
+ RegisterDeviceIntrinsicAttrs(def, op_namespace, effect_kind, names.printer);
+}
void RegisterDeviceIntrinsicAliases() {
- for (const auto& reg : kDeviceIntrinsics) {
- RegisterDeviceIntrinsic(reg);
- }
+ RegisterDeviceIntrinsic("cuda_any_sync", "cuda", CallEffectKind::kPure, sig::arg("mask"),
+ sig::arg("pred"));
+ RegisterDeviceIntrinsic("cuda_atomic_add", "cuda", CallEffectKind::kOpaque, sig::arg("res_addr"),
+ sig::arg("value"));
+ RegisterDeviceIntrinsic("cuda_atomic_cas", "cuda", CallEffectKind::kOpaque, sig::arg("ptr"),
+ sig::arg("old_val"), sig::arg("new_val"));
+ RegisterDeviceIntrinsic("cuda_wait_until", "cuda", CallEffectKind::kOpaque, sig::arg("dst"),
+ sig::arg("ptr"), sig::arg("condition"), sig::arg("scope"),
+ sig::arg("space"), sig::arg("ptx_type"), sig::arg("backoff_ns"));
+ RegisterDeviceIntrinsic("cuda_ballot_sync", "cuda", CallEffectKind::kOpaque, sig::arg("mask"),
+ sig::arg("pred"));
+ RegisterDeviceIntrinsic("cuda_bfloat1622float2", "cuda", CallEffectKind::kOpaque,
+ sig::arg("packed"));
+ RegisterDeviceIntrinsic("cuda_bfloat162float", "cuda", CallEffectKind::kOpaque, sig::arg("src"));
+ RegisterDeviceIntrinsic("cuda_clock64", "cuda", CallEffectKind::kOpaque);
+ RegisterDeviceIntrinsic("cuda_cluster_sync", "cuda", CallEffectKind::kOpaque);
+ RegisterDeviceIntrinsic("cuda_cta_reduce", "cuda", CallEffectKind::kOpaque, sig::arg("value"),
+ sig::arg("op"), sig::arg("num_warps"), sig::arg("scratch"));
+ RegisterDeviceIntrinsic("cuda_cta_sync", "cuda", CallEffectKind::kOpaque);
+ RegisterDeviceIntrinsic("cuda_cvta_generic_to_shared", "cuda", CallEffectKind::kOpaque,
+ sig::arg("ptr"));
+ RegisterDeviceIntrinsic("cuda_elect_sync", "cuda", CallEffectKind::kOpaque);
+ RegisterDeviceIntrinsic("cuda_fadd2_rn", "cuda", CallEffectKind::kOpaque, sig::arg("a"),
+ sig::arg("b"));
+ RegisterDeviceIntrinsic("cuda_fdividef", "cuda", CallEffectKind::kPure, sig::arg("x"),
+ sig::arg("y"));
+ RegisterDeviceIntrinsic("cuda_ffs_u32", "cuda", CallEffectKind::kOpaque, sig::arg("value"));
+ RegisterDeviceIntrinsic("cuda_float22bfloat162_rn", "cuda", CallEffectKind::kOpaque,
+ sig::arg("v0"), sig::arg("v1"));
+ RegisterDeviceIntrinsic("cuda_float22bfloat162_rn_from_float2", "cuda", CallEffectKind::kOpaque,
+ sig::arg("packed"));
+ RegisterDeviceIntrinsic("cuda_float22half2", "cuda", CallEffectKind::kOpaque, sig::arg("dst"),
+ sig::arg("src"));
+ RegisterDeviceIntrinsic("cuda_float2_x", "cuda", CallEffectKind::kOpaque, sig::arg("packed"));
+ RegisterDeviceIntrinsic("cuda_float2_y", "cuda", CallEffectKind::kOpaque, sig::arg("packed"));
+ RegisterDeviceIntrinsic("cuda_float8tohalf8", "cuda", CallEffectKind::kOpaque,
+ sig::arg("src_addr"), sig::arg("dst_addr"));
+ RegisterDeviceIntrinsic("cuda_float_as_uint", "cuda", CallEffectKind::kOpaque, sig::arg("x"));
+ RegisterDeviceIntrinsic("cuda_fmul2_rn", "cuda", CallEffectKind::kOpaque, sig::arg("a"),
+ sig::arg("b"));
+ RegisterDeviceIntrinsic("cuda_fp8x4_e4m3_from_float4", "cuda", CallEffectKind::kOpaque,
+ sig::arg("x"), sig::arg("y"), sig::arg("z"), sig::arg("w"));
+ RegisterDeviceIntrinsic("cuda_func_call", "cuda", CallEffectKind::kOpaque, sig::arg("func_name"),
+ sig::var_args("args"));
+ RegisterDeviceIntrinsic("cuda_get_tmem_addr", "cuda", CallEffectKind::kOpaque, sig::arg("addr"),
+ sig::arg("row_offset"), sig::arg("col_offset"));
+ RegisterDeviceIntrinsic("cuda_grid_sync", "cuda", CallEffectKind::kOpaque);
+ RegisterDeviceIntrinsic("cuda_half2float", "cuda", CallEffectKind::kOpaque, sig::arg("src"));
+ RegisterDeviceIntrinsic("cuda_half8tofloat8", "cuda", CallEffectKind::kOpaque,
+ sig::arg("src_addr"), sig::arg("dst_addr"));
+ RegisterDeviceIntrinsic("cuda_hmax2", "cuda", CallEffectKind::kOpaque, sig::arg("a"),
+ sig::arg("b"));
+ RegisterDeviceIntrinsic("cuda_hmin2", "cuda", CallEffectKind::kOpaque, sig::arg("a"),
+ sig::arg("b"));
+ RegisterDeviceIntrinsic("cuda_ldg", "cuda", CallEffectKind::kOpaque, sig::var_args("args"));
+ RegisterDeviceIntrinsic("cuda_make_float2", "cuda", CallEffectKind::kOpaque, sig::arg("x"),
+ sig::arg("y"));
+ RegisterDeviceIntrinsic("cuda_mbarrier_wait", "cuda", CallEffectKind::kOpaque, sig::arg("bar"),
+ sig::arg("phase"));
+ RegisterDeviceIntrinsic("cuda_mbarrier_wait_acquire_cluster", "cuda", CallEffectKind::kOpaque,
+ sig::arg("bar"), sig::arg("phase"));
+ RegisterDeviceIntrinsic("cuda_mov_sreg", "cuda", CallEffectKind::kPure, sig::arg("bits"),
+ sig::arg("reg_name"));
+ RegisterDeviceIntrinsic("cuda_nano_sleep", "cuda", CallEffectKind::kOpaque, sig::arg("time"));
+ RegisterDeviceIntrinsic("cuda_printf", "cuda", CallEffectKind::kOpaque, sig::arg("fmt"),
+ sig::var_args("args"));
+ RegisterDeviceIntrinsic("cuda_reduce_add_sync_u32", "cuda", CallEffectKind::kOpaque,
+ sig::arg("mask"), sig::arg("value"));
+ RegisterDeviceIntrinsic("cuda_reduce_min_sync_u32", "cuda", CallEffectKind::kOpaque,
+ sig::arg("mask"), sig::arg("value"));
+ RegisterDeviceIntrinsic("cuda_runtime_instr_desc", "cuda", CallEffectKind::kOpaque,
+ sig::arg("desc"), sig::arg("sf_id"));
+ RegisterDeviceIntrinsic("cuda_sm100_2sm_leader_smem_addr", "cuda", CallEffectKind::kOpaque,
+ sig::arg("ptr"));
+ RegisterDeviceIntrinsic("cuda_smem_addr_from_uint64", "cuda", CallEffectKind::kOpaque,
+ sig::arg("cluster_addr"));
+ RegisterDeviceIntrinsic("cuda_syncthreads_and", "cuda", CallEffectKind::kOpaque,
+ sig::arg("cond"));
+ RegisterDeviceIntrinsic("cuda_syncthreads_or", "cuda", CallEffectKind::kOpaque, sig::arg("cond"));
+ RegisterDeviceIntrinsic(
+ "cuda_tcgen05_encode_instr_descriptor", "cuda", CallEffectKind::kOpaque, sig::arg("desc"),
+ sig::arg("d_dtype"), sig::arg("a_dtype"), sig::arg("b_dtype"), sig::arg("M"), sig::arg("N"),
+ sig::arg("K"), sig::arg("trans_a"), sig::arg("trans_b"), sig::arg("n_cta_groups"),
+ sig::arg("neg_a"), sig::arg("neg_b"), sig::arg("sat_d"), sig::arg("is_sparse"));
+ RegisterDeviceIntrinsic("cuda_tcgen05_encode_instr_descriptor_block_scaled", "cuda",
+ CallEffectKind::kOpaque, sig::arg("desc"), sig::arg("d_dtype"),
+ sig::arg("a_dtype"), sig::arg("b_dtype"), sig::arg("sfa_dtype"),
+ sig::arg("sfb_dtype"), sig::arg("sfa_tmem_addr"),
+ sig::arg("sfb_tmem_addr"), sig::arg("M"), sig::arg("N"), sig::arg("K"),
+ sig::arg("trans_a"), sig::arg("trans_b"), sig::arg("n_cta_groups"),
+ sig::arg("neg_a"), sig::arg("neg_b"), sig::arg("is_sparse"));
+ RegisterDeviceIntrinsic("cuda_tcgen05_encode_matrix_descriptor", "cuda", CallEffectKind::kOpaque,
+ sig::arg("desc"), sig::arg("addr"), sig::arg("ldo"), sig::arg("sdo"),
+ sig::arg("swizzle"));
+ RegisterDeviceIntrinsic("cuda_thread_fence", "cuda", CallEffectKind::kOpaque);
+ RegisterDeviceIntrinsic("cuda_thread_rank", "cuda", CallEffectKind::kPure);
+ RegisterDeviceIntrinsic("cuda_trap_when_assert_failed", "cuda", CallEffectKind::kOpaque,
+ sig::arg("cond"));
+ RegisterDeviceIntrinsic("cuda_uint_as_float", "cuda", CallEffectKind::kOpaque, sig::arg("bits"));
+ RegisterDeviceIntrinsic("cuda_warp_reduce", "cuda", CallEffectKind::kOpaque, sig::arg("value"),
+ sig::arg("op"), sig::arg("width"));
+ RegisterDeviceIntrinsic("cuda_warp_sync", "cuda", CallEffectKind::kOpaque);
+ RegisterDeviceIntrinsic("cuda_warpgroup_sync", "cuda", CallEffectKind::kOpaque,
+ sig::arg("bar_no"));
+ RegisterDeviceIntrinsic("cuda_wgmma_encode_matrix_descriptor", "cuda", CallEffectKind::kOpaque,
+ sig::arg("desc"), sig::arg("addr"), sig::arg("ldo"), sig::arg("sdo"),
+ sig::arg("swizzle"));
+ RegisterDeviceIntrinsic("cuda_wgmma_noop_barrier", "cuda", CallEffectKind::kOpaque,
+ sig::arg("reg"));
+ RegisterDeviceIntrinsic("nvshmem_barrier_all", "nvshmem", CallEffectKind::kOpaque);
+ RegisterDeviceIntrinsic("nvshmem_fence", "nvshmem", CallEffectKind::kOpaque);
+ RegisterDeviceIntrinsic("nvshmem_getmem_nbi", "nvshmem", CallEffectKind::kOpaque, sig::arg("dst"),
+ sig::arg("src"), sig::arg("nelems"), sig::arg("pe"));
+ RegisterDeviceIntrinsic("nvshmem_getmem_nbi_block", "nvshmem", CallEffectKind::kOpaque,
+ sig::arg("dst"), sig::arg("src"), sig::arg("nelems"), sig::arg("pe"));
+ RegisterDeviceIntrinsic("nvshmem_getmem_nbi_warp", "nvshmem", CallEffectKind::kOpaque,
+ sig::arg("dst"), sig::arg("src"), sig::arg("nelems"), sig::arg("pe"));
+ RegisterDeviceIntrinsic("nvshmem_my_pe", "nvshmem", CallEffectKind::kOpaque);
+ RegisterDeviceIntrinsic("nvshmem_n_pes", "nvshmem", CallEffectKind::kOpaque);
+ RegisterDeviceIntrinsic("nvshmem_putmem_nbi", "nvshmem", CallEffectKind::kOpaque, sig::arg("dst"),
+ sig::arg("src"), sig::arg("nelems"), sig::arg("pe"));
+ RegisterDeviceIntrinsic("nvshmem_putmem_nbi_block", "nvshmem", CallEffectKind::kOpaque,
+ sig::arg("dst"), sig::arg("src"), sig::arg("nelems"), sig::arg("pe"));
+ RegisterDeviceIntrinsic("nvshmem_putmem_nbi_warp", "nvshmem", CallEffectKind::kOpaque,
+ sig::arg("dst"), sig::arg("src"), sig::arg("nelems"), sig::arg("pe"));
+ RegisterDeviceIntrinsic("nvshmem_putmem_signal_nbi", "nvshmem", CallEffectKind::kOpaque,
+ sig::arg("dst"), sig::arg("src"), sig::arg("nelems"),
+ sig::arg("sig_addr"), sig::arg("signal"), sig::arg("sig_op"),
+ sig::arg("pe"));
+ RegisterDeviceIntrinsic("nvshmem_putmem_signal_nbi_block", "nvshmem", CallEffectKind::kOpaque,
+ sig::arg("dst"), sig::arg("src"), sig::arg("nelems"),
+ sig::arg("sig_addr"), sig::arg("signal"), sig::arg("sig_op"),
+ sig::arg("pe"));
+ RegisterDeviceIntrinsic("nvshmem_putmem_signal_nbi_warp", "nvshmem", CallEffectKind::kOpaque,
+ sig::arg("dst"), sig::arg("src"), sig::arg("nelems"),
+ sig::arg("sig_addr"), sig::arg("signal"), sig::arg("sig_op"),
+ sig::arg("pe"));
+ RegisterDeviceIntrinsic("nvshmem_quiet", "nvshmem", CallEffectKind::kOpaque);
+ RegisterDeviceIntrinsic("nvshmem_signal_op", "nvshmem", CallEffectKind::kOpaque,
+ sig::arg("sig_addr"), sig::arg("signal"), sig::arg("sig_op"),
+ sig::arg("pe"));
+ RegisterDeviceIntrinsic("nvshmem_wait_until", "nvshmem", CallEffectKind::kOpaque,
+ sig::arg("ivar"), sig::arg("cmp"), sig::arg("cmp_value"),
+ sig::arg("type"));
+ RegisterDeviceIntrinsic("ptx_legacy_ldmatrix", "ptx_legacy", CallEffectKind::kOpaque,
+ sig::arg("trans"), sig::arg("num"), sig::arg("dtype"),
+ sig::arg("local_ptr"), sig::arg("local_offset"), sig::arg("smem_ptr"),
+ sig::arg("smem_offset"));
+ RegisterDeviceIntrinsic("ptx_legacy_mma", "ptx_legacy", CallEffectKind::kOpaque,
+ sig::arg("shape"), sig::arg("a_layout"), sig::arg("b_layout"),
+ sig::arg("a_dtype"), sig::arg("b_dtype"), sig::arg("c_dtype"),
+ sig::arg("a_ptr"), sig::arg("a_offset"), sig::arg("b_ptr"),
+ sig::arg("b_offset"), sig::arg("acc_ptr"), sig::arg("c_offset"),
+ sig::arg("saturate"), sig::var_args("args"));
}
} // namespace
diff --git a/src/backend/hexagon/codegen/llvm/intrin_rule_hexagon.cc b/src/backend/hexagon/codegen/llvm/intrin_rule_hexagon.cc
index cacbd924d6..7f2bf8de2b 100644
--- a/src/backend/hexagon/codegen/llvm/intrin_rule_hexagon.cc
+++ b/src/backend/hexagon/codegen/llvm/intrin_rule_hexagon.cc
@@ -109,7 +109,7 @@ void RegisterHexagonIntrinRules() {
DispatchLLVMPureIntrin<::llvm::Intrinsic::nearbyint, 1>);
OpDef("tirx.ctpop")
- .add_arg("x", "The input value.")
+ .signature(sig::arg("x", "The input value."))
.set_attr<tirx::TIRxOpCategory>("TIRxOpCategory", ffi::String("builtin"))
.set_attr<FLowerIntrinsic>("hexagon.FLowerIntrinsic",
DispatchLLVMPureIntrin<::llvm::Intrinsic::ctpop, 1>);
diff --git a/src/backend/metal/codegen/intrin_rule_metal.cc b/src/backend/metal/codegen/intrin_rule_metal.cc
index b5eb1f86ae..306e81a55f 100644
--- a/src/backend/metal/codegen/intrin_rule_metal.cc
+++ b/src/backend/metal/codegen/intrin_rule_metal.cc
@@ -147,8 +147,8 @@ void RegisterMetalIntrinRules() {
// Register low-level Metal device intrinsics.
OpDef("tirx.metal.simd_shuffle")
- .add_arg("var", "The variable to sync.")
- .add_arg("lane", "The source thread id.")
+ .signature(sig::arg("var", "The variable to sync."),
+ sig::arg("lane", "The source thread id."))
.set_attr<tirx::TIRxOpCategory>("TIRxOpCategory", ffi::String("device_intrin"))
.set_attr<tirx::TDeviceIntrinsicNamespace>("TDeviceIntrinsicNamespace", ffi::String("metal"))
.set_attr<tirx::TScriptPrinterName>("TScriptPrinterName", ffi::String("metal.simd_shuffle"))
@@ -156,8 +156,8 @@ void RegisterMetalIntrinRules() {
.set_attr<TCallEffectKind>("TCallEffectKind", static_cast<int64_t>(CallEffectKind::kOpaque));
OpDef("tirx.metal.simd_shuffle_up")
- .add_arg("var", "The variable to sync.")
- .add_arg("delta", "The source lane id offset to be added.")
+ .signature(sig::arg("var", "The variable to sync."),
+ sig::arg("delta", "The source lane id offset to be added."))
.set_attr<tirx::TIRxOpCategory>("TIRxOpCategory", ffi::String("device_intrin"))
.set_attr<tirx::TDeviceIntrinsicNamespace>("TDeviceIntrinsicNamespace", ffi::String("metal"))
.set_attr<tirx::TScriptPrinterName>("TScriptPrinterName",
@@ -166,8 +166,8 @@ void RegisterMetalIntrinRules() {
.set_attr<TCallEffectKind>("TCallEffectKind", static_cast<int64_t>(CallEffectKind::kOpaque));
OpDef("tirx.metal.simd_shuffle_down")
- .add_arg("var", "The variable to sync.")
- .add_arg("delta", "The source lane id offset to be subtracted.")
+ .signature(sig::arg("var", "The variable to sync."),
+ sig::arg("delta", "The source lane id offset to be subtracted."))
.set_attr<tirx::TIRxOpCategory>("TIRxOpCategory", ffi::String("device_intrin"))
.set_attr<tirx::TDeviceIntrinsicNamespace>("TDeviceIntrinsicNamespace", ffi::String("metal"))
.set_attr<tirx::TScriptPrinterName>("TScriptPrinterName",
diff --git a/src/backend/metal/op/target_builtin.cc b/src/backend/metal/op/target_builtin.cc
index e6fdcdbcd4..7f6c8bf60e 100644
--- a/src/backend/metal/op/target_builtin.cc
+++ b/src/backend/metal/op/target_builtin.cc
@@ -37,49 +37,37 @@ void RegisterMetalTargetBuiltins() {
registered = true;
OpDef("tirx.make_filled_simdgroup_matrix")
- .add_arg("d", "The D operand.")
- .add_arg("index", "The index.")
- .add_arg("value", "The value to use.")
- .add_arg("col", "The column index.")
- .add_arg("row", "The row index.")
+ .signature(sig::arg("d", "The D operand."), sig::arg("index", "The index."),
+ sig::arg("value", "The value to use."), sig::arg("col", "The column index."),
+ sig::arg("row", "The row index."))
.set_attr<TScriptPrinterName>("TScriptPrinterName",
ffi::String("metal.make_filled_simdgroup_matrix"))
.set_attr<TIRxOpCategory>("TIRxOpCategory", ffi::String("builtin"))
.set_attr<TCallEffectKind>("TCallEffectKind", static_cast<int64_t>(CallEffectKind::kOpaque));
OpDef("tirx.simdgroup_load")
- .add_arg("d", "The D operand.")
- .add_arg("index", "The index.")
- .add_arg("ptr", "The pointer.")
- .add_arg("stride", "The stride.")
- .add_arg("col", "The column index.")
- .add_arg("row", "The row index.")
- .add_arg("transpose_matrix", "Whether to transpose the matrix.")
+ .signature(sig::arg("d", "The D operand."), sig::arg("index", "The index."),
+ sig::arg("ptr", "The pointer."), sig::arg("stride", "The stride."),
+ sig::arg("col", "The column index."), sig::arg("row", "The row index."),
+ sig::arg("transpose_matrix", "Whether to transpose the matrix."))
.set_attr<TScriptPrinterName>("TScriptPrinterName", ffi::String("metal.simdgroup_load"))
.set_attr<TIRxOpCategory>("TIRxOpCategory", ffi::String("builtin"))
.set_attr<TCallEffectKind>("TCallEffectKind", static_cast<int64_t>(CallEffectKind::kOpaque));
OpDef("tirx.simdgroup_store")
- .add_arg("d", "The D operand.")
- .add_arg("index", "The index.")
- .add_arg("ptr", "The pointer.")
- .add_arg("stride", "The stride.")
- .add_arg("col", "The column index.")
- .add_arg("row", "The row index.")
- .add_arg("transpose_matrix", "Whether to transpose the matrix.")
+ .signature(sig::arg("d", "The D operand."), sig::arg("index", "The index."),
+ sig::arg("ptr", "The pointer."), sig::arg("stride", "The stride."),
+ sig::arg("col", "The column index."), sig::arg("row", "The row index."),
+ sig::arg("transpose_matrix", "Whether to transpose the matrix."))
.set_attr<TScriptPrinterName>("TScriptPrinterName", ffi::String("metal.simdgroup_store"))
.set_attr<TIRxOpCategory>("TIRxOpCategory", ffi::String("builtin"))
.set_attr<TCallEffectKind>("TCallEffectKind", static_cast<int64_t>(CallEffectKind::kOpaque));
OpDef("tirx.simdgroup_multiply_accumulate")
- .add_arg("d", "The D operand.")
- .add_arg("index_d", "The D fragment index.")
- .add_arg("a", "The A operand.")
- .add_arg("index_a", "The A fragment index.")
- .add_arg("b", "The B operand.")
- .add_arg("index_b", "The B fragment index.")
- .add_arg("c", "The C operand.")
- .add_arg("index_c", "The C fragment index.")
+ .signature(sig::arg("d", "The D operand."), sig::arg("index_d", "The D fragment index."),
+ sig::arg("a", "The A operand."), sig::arg("index_a", "The A fragment index."),
+ sig::arg("b", "The B operand."), sig::arg("index_b", "The B fragment index."),
+ sig::arg("c", "The C operand."), sig::arg("index_c", "The C fragment index."))
.set_attr<TScriptPrinterName>("TScriptPrinterName",
ffi::String("metal.simdgroup_multiply_accumulate"))
.set_attr<TIRxOpCategory>("TIRxOpCategory", ffi::String("builtin"))
diff --git a/src/backend/trn/op/target_builtin.cc b/src/backend/trn/op/target_builtin.cc
index 64d06b00e0..2beb4dc45e 100644
--- a/src/backend/trn/op/target_builtin.cc
+++ b/src/backend/trn/op/target_builtin.cc
@@ -51,69 +51,66 @@ void RegisterTRNTargetBuiltins() {
namespace {
-struct NKIIntrinsicRegistration {
- const char* name;
- std::initializer_list<const char*> args;
- bool allow_extra_args;
+struct NKIIntrinsicNames {
+ std::string canonical;
+ std::string printer;
};
-void RegisterNKIIntrinsic(const NKIIntrinsicRegistration& reg) {
+TVM_FFI_NO_INLINE NKIIntrinsicNames MakeNKIIntrinsicNames(const char* op_name) {
std::string prefix = "nki_";
- std::string suffix(reg.name);
+ std::string suffix(op_name);
if (suffix.rfind(prefix, 0) == 0) {
suffix = suffix.substr(prefix.size());
}
+ return {"tirx.nki." + suffix, "nki." + suffix};
+}
- std::string canonical_op_name = "tirx.nki." + suffix;
- ffi::String namespace_attr("nki");
- ffi::String printer_name("nki." + suffix);
- int64_t effect = static_cast<int64_t>(CallEffectKind::kOpaque);
-
- OpDef def(canonical_op_name);
- for (const char* name : reg.args) {
- def.add_arg(name, "");
- }
- if (reg.allow_extra_args) {
- def.allow_extra_args();
- }
+TVM_FFI_NO_INLINE void RegisterNKIIntrinsicAttrs(OpDef& def, const std::string& printer_name) {
def.set_attr<TIRxOpCategory>("TIRxOpCategory", ffi::String("device_intrin"))
- .set_attr<TDeviceIntrinsicNamespace>("TDeviceIntrinsicNamespace", namespace_attr)
- .set_attr<TCallEffectKind>("TCallEffectKind", effect)
- .set_attr<TScriptPrinterName>("TScriptPrinterName", printer_name);
+ .set_attr<TDeviceIntrinsicNamespace>("TDeviceIntrinsicNamespace", ffi::String("nki"))
+ .set_attr<TCallEffectKind>("TCallEffectKind", static_cast<int64_t>(CallEffectKind::kOpaque))
+ .set_attr<TScriptPrinterName>("TScriptPrinterName", ffi::String(printer_name));
}
-const NKIIntrinsicRegistration kNKIIntrinsics[] = {
- {"nki_activation", {"result", "data", "opcode", "bias", "scale"}, false},
- {"nki_activation_reduce",
- {"reduce_res", "act_res", "data", "opcode", "reduce_opcode", "bias", "scale"},
- false},
- {"nki_affine_select", {"result", "pred", "true_value", "false_value"}, false},
- {"nki_identity", {"result", "size"}, false},
- {"nki_load", {"res", "data"}, false},
- {"nki_matmul", {"res", "lhs", "rhs", "accum"}, false},
- {"nki_memset", {"result", "value"}, false},
- {"nki_reciprocal", {"result", "data"}, false},
- {"nki_scalar_tensor_scalar",
- {"result", "data", "operand0", "operand1", "opcode0", "opcode1", "reverse0", "reverse1"},
- false},
- {"nki_scalar_tensor_tensor",
- {"result", "data", "operand0", "operand1", "opcode0", "opcode1", "reverse0", "reverse1"},
- false},
- {"nki_store", {"res", "data"}, false},
- {"nki_tensor_copy", {"res", "data"}, false},
- {"nki_tensorreduce", {"result", "data", "opcode", "negate"}, true},
- {"nki_tensorscalar", {"result", "operand0", "operand1", "opcode", "reverse"}, false},
- {"nki_tensorscalar_reduce",
- {"reduce_res", "tensorscalar_res", "operand0", "operand1", "opcode", "reduce_opcode",
- "reverse"},
- false},
- {"nki_tensortensor", {"result", "operand0", "operand1", "opcode"}, false},
-};
+template <typename... Specs>
+void RegisterNKIIntrinsic(const char* op_name, const Specs&... specs) {
+ NKIIntrinsicNames names = MakeNKIIntrinsicNames(op_name);
+ OpDef def(names.canonical);
+ def.signature(specs...);
+ RegisterNKIIntrinsicAttrs(def, names.printer);
+}
void RegisterNKIIntrinsicAliases() {
- for (const auto& reg : kNKIIntrinsics) {
- RegisterNKIIntrinsic(reg);
- }
+ RegisterNKIIntrinsic("nki_activation", sig::arg("result"), sig::arg("data"), sig::arg("opcode"),
+ sig::arg("bias"), sig::arg("scale"));
+ RegisterNKIIntrinsic("nki_activation_reduce", sig::arg("reduce_res"), sig::arg("act_res"),
+ sig::arg("data"), sig::arg("opcode"), sig::arg("reduce_opcode"),
+ sig::arg("bias"), sig::arg("scale"));
+ RegisterNKIIntrinsic("nki_affine_select", sig::arg("result"), sig::arg("pred"),
+ sig::arg("true_value"), sig::arg("false_value"));
+ RegisterNKIIntrinsic("nki_identity", sig::arg("result"), sig::arg("size"));
+ RegisterNKIIntrinsic("nki_load", sig::arg("res"), sig::arg("data"));
+ RegisterNKIIntrinsic("nki_matmul", sig::arg("res"), sig::arg("lhs"), sig::arg("rhs"),
+ sig::arg("accum"));
+ RegisterNKIIntrinsic("nki_memset", sig::arg("result"), sig::arg("value"));
+ RegisterNKIIntrinsic("nki_reciprocal", sig::arg("result"), sig::arg("data"));
+ RegisterNKIIntrinsic("nki_scalar_tensor_scalar", sig::arg("result"), sig::arg("data"),
+ sig::arg("operand0"), sig::arg("operand1"), sig::arg("opcode0"),
+ sig::arg("opcode1"), sig::arg("reverse0"), sig::arg("reverse1"));
+ RegisterNKIIntrinsic("nki_scalar_tensor_tensor", sig::arg("result"), sig::arg("data"),
+ sig::arg("operand0"), sig::arg("operand1"), sig::arg("opcode0"),
+ sig::arg("opcode1"), sig::arg("reverse0"), sig::arg("reverse1"));
+ RegisterNKIIntrinsic("nki_store", sig::arg("res"), sig::arg("data"));
+ RegisterNKIIntrinsic("nki_tensor_copy", sig::arg("res"), sig::arg("data"));
+ RegisterNKIIntrinsic("nki_tensorreduce", sig::arg("result"), sig::arg("data"), sig::arg("opcode"),
+ sig::arg("negate"), sig::var_args("args"));
+ RegisterNKIIntrinsic("nki_tensorscalar", sig::arg("result"), sig::arg("operand0"),
+ sig::arg("operand1"), sig::arg("opcode"), sig::arg("reverse"));
+ RegisterNKIIntrinsic("nki_tensorscalar_reduce", sig::arg("reduce_res"),
+ sig::arg("tensorscalar_res"), sig::arg("operand0"), sig::arg("operand1"),
+ sig::arg("opcode"), sig::arg("reduce_opcode"), sig::arg("reverse"));
+ RegisterNKIIntrinsic("nki_tensortensor", sig::arg("result"), sig::arg("operand0"),
+ sig::arg("operand1"), sig::arg("opcode"));
}
} // namespace
diff --git a/src/backend/webgpu/codegen/intrin_rule_webgpu.cc b/src/backend/webgpu/codegen/intrin_rule_webgpu.cc
index d030a78417..12d0ad8c0b 100644
--- a/src/backend/webgpu/codegen/intrin_rule_webgpu.cc
+++ b/src/backend/webgpu/codegen/intrin_rule_webgpu.cc
@@ -174,8 +174,8 @@ void RegisterWebGPUIntrinRules() {
// Register low-level WebGPU device intrinsics.
OpDef("tirx.webgpu.subgroup_shuffle")
- .add_arg("var", "The variable to sync.")
- .add_arg("lane", "The source thread id.")
+ .signature(sig::arg("var", "The variable to sync."),
+ sig::arg("lane", "The source thread id."))
.set_attr<tirx::TIRxOpCategory>("TIRxOpCategory", ffi::String("device_intrin"))
.set_attr<tirx::TDeviceIntrinsicNamespace>("TDeviceIntrinsicNamespace", ffi::String("webgpu"))
.set_attr<tirx::TScriptPrinterName>("TScriptPrinterName",
@@ -184,8 +184,8 @@ void RegisterWebGPUIntrinRules() {
.set_attr<TCallEffectKind>("TCallEffectKind", static_cast<int64_t>(CallEffectKind::kOpaque));
OpDef("tirx.webgpu.subgroup_shuffle_up")
- .add_arg("var", "The variable to sync.")
- .add_arg("delta", "The source lane id offset to be added.")
+ .signature(sig::arg("var", "The variable to sync."),
+ sig::arg("delta", "The source lane id offset to be added."))
.set_attr<tirx::TIRxOpCategory>("TIRxOpCategory", ffi::String("device_intrin"))
.set_attr<tirx::TDeviceIntrinsicNamespace>("TDeviceIntrinsicNamespace", ffi::String("webgpu"))
.set_attr<tirx::TScriptPrinterName>("TScriptPrinterName",
@@ -194,8 +194,8 @@ void RegisterWebGPUIntrinRules() {
.set_attr<TCallEffectKind>("TCallEffectKind", static_cast<int64_t>(CallEffectKind::kOpaque));
OpDef("tirx.webgpu.subgroup_shuffle_down")
- .add_arg("var", "The variable to sync.")
- .add_arg("delta", "The source lane id offset to be subtracted.")
+ .signature(sig::arg("var", "The variable to sync."),
+ sig::arg("delta", "The source lane id offset to be subtracted."))
.set_attr<tirx::TIRxOpCategory>("TIRxOpCategory", ffi::String("device_intrin"))
.set_attr<tirx::TDeviceIntrinsicNamespace>("TDeviceIntrinsicNamespace", ffi::String("webgpu"))
.set_attr<tirx::TScriptPrinterName>("TScriptPrinterName",
diff --git a/src/ir/expr.cc b/src/ir/expr.cc
index 189e039d9a..02ba0b5768 100644
--- a/src/ir/expr.cc
+++ b/src/ir/expr.cc
@@ -1130,7 +1130,18 @@ TVM_FFI_STATIC_INIT_BLOCK() {
// Call
Call::Call(Type ret_ty, Expr op, ffi::Array<Expr> args, Attrs attrs, ffi::Array<Type> ty_args,
- Span span) {
+ Span span)
+ : Call(std::move(ret_ty), std::move(op), std::move(args), std::move(attrs), std::move(ty_args),
+ std::move(span), true) {}
+
+Call Call::Unchecked(Type ret_ty, Expr op, ffi::Array<Expr> args, Attrs attrs,
+ ffi::Array<Type> ty_args, Span span) {
+ return Call(std::move(ret_ty), std::move(op), std::move(args), std::move(attrs),
+ std::move(ty_args), std::move(span), false);
+}
+
+Call::Call(Type ret_ty, Expr op, ffi::Array<Expr> args, Attrs attrs, ffi::Array<Type> ty_args,
+ Span span, bool validate) {
TVM_FFI_CHECK(op.defined(), ValueError) << "Call expects a defined operator";
ffi::ObjectPtr<CallNode> n = ffi::make_object<CallNode>();
@@ -1140,6 +1151,9 @@ Call::Call(Type ret_ty, Expr op, ffi::Array<Expr> args, Attrs attrs, ffi::Array<
n->attrs = std::move(attrs);
n->ty_args = std::move(ty_args);
n->span = std::move(span);
+ if (validate) {
+ if (auto opt_op = n->op.as<Op>()) opt_op.value().Validate(n.get());
+ }
data_ = std::move(n);
}
@@ -1156,6 +1170,10 @@ TVM_FFI_STATIC_INIT_BLOCK() {
ffi::Array<Type> ty_args, Span span) {
return Call(ret_ty, op, args, attrs, ty_args, span);
});
+ refl::GlobalDef().def("ir.CallUnchecked", [](Type ret_ty, Expr op, ffi::Array<Expr> args,
+ Attrs attrs, ffi::Array<Type> ty_args, Span span) {
+ return Call::Unchecked(ret_ty, op, args, attrs, ty_args, span);
+ });
}
TVM_FFI_STATIC_INIT_BLOCK() {
diff --git a/src/ir/op.cc b/src/ir/op.cc
index 7f63eb7ca2..5499ec21dc 100644
--- a/src/ir/op.cc
+++ b/src/ir/op.cc
@@ -98,6 +98,10 @@ Op Op::Get(const ffi::String& name) { return OpRegistry::Global()->Get(name); }
ffi::Array<ffi::String> Op::ListNames() { return OpRegistry::Global()->ListNames(); }
+void Op::ThrowInvalidCall(const OpNode* op) {
+ TVM_FFI_THROW(TypeError) << "Op `" << op->name << "`: invalid Call";
+}
+
bool Op::HasAttrMap(const ffi::String& name) { return OpRegistry::Global()->HasAttrMap(name); }
ffi::List<ffi::Any> Op::GetAttrColumn(const ffi::String& name) {
@@ -140,6 +144,84 @@ OpDef& OpDef::reset_attr(const ffi::String& name) {
return *this;
}
+namespace {
+
+ffi::Array<ArgumentInfo> MakeArgumentInfos(const ffi::Array<ffi::String>& names,
+ const ffi::Array<ffi::String>& docs) {
+ TVM_FFI_CHECK_EQ(names.size(), docs.size(), ValueError);
+ ffi::Array<ArgumentInfo> infos;
+ for (size_t i = 0; i < names.size(); ++i) {
+ auto info = ffi::make_object<ArgumentInfoNode>();
+ info->name = names[i];
+ info->doc = docs[i];
+ infos.push_back(ArgumentInfo(std::move(info)));
+ }
+ return infos;
+}
+
+ffi::Optional<ArgumentInfo> MakeTailInfo(const ffi::Array<ffi::String>& tail) {
+ TVM_FFI_CHECK(tail.empty() || tail.size() == 2, ValueError)
+ << "A variadic signature entry must have a name and documentation";
+ if (tail.empty()) return std::nullopt;
+ auto info = ffi::make_object<ArgumentInfoNode>();
+ info->name = tail[0];
+ info->doc = tail[1];
+ return ArgumentInfo(std::move(info));
+}
+
+TVM_FFI_INLINE ffi::Expected<void> ValidateCountSignature(const CallNode* call) noexcept {
+ const OpNode* op = call ? call->op.as<OpNode>() : nullptr;
+ if (TVM_FFI_PREDICT_FALSE(!op)) {
+ return TVM_FFI_UNEXPECTED(TypeError) << "Invalid Op Call";
+ }
+ if (TVM_FFI_PREDICT_FALSE(!call->args.defined() || !call->ty_args.defined())) {
+ return TVM_FFI_UNEXPECTED(TypeError) << "Op `" << op->name << "`: invalid Call";
+ }
+ const size_t n_args = op->args_info.size();
+ const bool var_args = op->var_args_info.has_value();
+ if (TVM_FFI_PREDICT_FALSE(var_args ? call->args.size() < n_args : call->args.size() != n_args)) {
+ return TVM_FFI_UNEXPECTED(TypeError)
+ << "Op `" << op->name << "`: Call.args expected " << (var_args ? "at least " : "")
+ << n_args << (n_args == 1 ? " argument" : " arguments") << ", got " << call->args.size();
+ }
+ const size_t n_ty_args = op->ty_args_info.size();
+ const bool var_ty_args = op->var_ty_args_info.has_value();
+ if (TVM_FFI_PREDICT_FALSE(var_ty_args ? call->ty_args.size() < n_ty_args
+ : call->ty_args.size() != n_ty_args)) {
+ return TVM_FFI_UNEXPECTED(TypeError)
+ << "Op `" << op->name << "`: Call.ty_args expected " << (var_ty_args ? "at least " : "")
+ << n_ty_args << (n_ty_args == 1 ? " type argument" : " type arguments") << ", got "
+ << call->ty_args.size();
+ }
+ return {};
+}
+
+void SetOpSignature(Op op, const ffi::Array<ffi::String>& arg_names,
+ const ffi::Array<ffi::String>& arg_docs,
+ const ffi::Array<ffi::String>& ty_arg_names,
+ const ffi::Array<ffi::String>& ty_arg_docs,
+ const ffi::Array<ffi::String>& var_args,
+ const ffi::Array<ffi::String>& var_ty_args) {
+ auto args_info = MakeArgumentInfos(arg_names, arg_docs);
+ auto ty_args_info = MakeArgumentInfos(ty_arg_names, ty_arg_docs);
+ auto var_args_info = MakeTailInfo(var_args);
+ auto var_ty_args_info = MakeTailInfo(var_ty_args);
+ auto* node = const_cast<OpNode*>(op.operator->());
+ if (!node->validator_is_custom) {
+ using View = ffi::reflection::NativeFunctionView<ffi::Expected<void>(const CallNode*)>;
+ node->validator = ffi::reflection::NativeFunction<ffi::Expected<void>(const CallNode*)>::From(
+ View::FromNative<&ValidateCountSignature>());
+ }
+ node->args_info = std::move(args_info);
+ node->ty_args_info = std::move(ty_args_info);
+ node->var_args_info = std::move(var_args_info);
+ node->var_ty_args_info = std::move(var_ty_args_info);
+ node->attrs_type_key = "";
+ node->attrs_type_index = 0;
+}
+
+} // namespace
+
void OpNode::RegisterReflection() {
namespace refl = ffi::reflection;
// clang-format off
@@ -149,8 +231,9 @@ void OpNode::RegisterReflection() {
.def_ro("doc", &OpNode::doc, refl::AttachFieldFlag::SEqHashIgnore())
.def_ro("args_info", &OpNode::args_info, refl::AttachFieldFlag::SEqHashIgnore())
.def_ro("attrs_type_key", &OpNode::attrs_type_key, refl::AttachFieldFlag::SEqHashIgnore())
- .def_ro("allow_extra_args", &OpNode::allow_extra_args, refl::AttachFieldFlag::SEqHashIgnore())
.def_ro("ty_args_info", &OpNode::ty_args_info, refl::AttachFieldFlag::SEqHashIgnore())
+ .def_ro("var_args_info", &OpNode::var_args_info, refl::AttachFieldFlag::SEqHashIgnore())
+ .def_ro("var_ty_args_info", &OpNode::var_ty_args_info, refl::AttachFieldFlag::SEqHashIgnore())
.def_static("get", &Op::Get,
"get(op_name: str) -> Op\n\nReturn the canonical named Op; raise AttributeError "
"if unregistered.")
@@ -172,6 +255,7 @@ void OpNode::RegisterReflection() {
OpDef(op->name)
.set_attr(name, value, override);
})
+ .def("_set_signature", &SetOpSignature)
.def(
"reset_attr", [](Op op, ffi::String name) {
OpDef(op->name)
@@ -187,12 +271,6 @@ void OpNode::RegisterReflection() {
},
"add_arg(name: str, doc: str) -> None\n\nAppend an argument's "
"name and documentation.")
- .def(
- "set_allow_extra_args", [](Op op) {
- OpDef(op->name)
- .allow_extra_args();
- },
- "set_allow_extra_args() -> None\n\nAllow value arguments after the required prefix.")
.def(
"add_ty_arg",
[](Op op, ffi::String name, ffi::String doc) {
diff --git a/src/ir/prim/builtin.cc b/src/ir/prim/builtin.cc
index bb0f6d7d99..76a1eff6f5 100644
--- a/src/ir/prim/builtin.cc
+++ b/src/ir/prim/builtin.cc
@@ -40,32 +40,32 @@ TVM_DEFINE_CACHED_OP_GETTER(clz, "prim.clz")
TVM_FFI_STATIC_INIT_BLOCK() {
OpDef("prim.likely")
- .add_arg("x", "The input value.")
+ .signature(sig::arg("x", "The input value."))
.set_attr<TCallEffectKind>("TCallEffectKind",
static_cast<int64_t>(CallEffectKind::kExprAnnotation))
.set_attr<bool>("TVectorizable", true);
OpDef("prim.if_then_else")
- .add_arg("condition", "The condition.")
- .add_arg("true_value", "The value when the condition is true.")
- .add_arg("false_value", "The value when the condition is false.")
+ .signature(sig::arg("condition", "The condition."),
+ sig::arg("true_value", "The value when the condition is true."),
+ sig::arg("false_value", "The value when the condition is false."))
.set_attr<TCallEffectKind>("TCallEffectKind", static_cast<int64_t>(CallEffectKind::kPure));
OpDef("prim.vscale")
.set_attr<TCallEffectKind>("TCallEffectKind", static_cast<int64_t>(CallEffectKind::kPure));
OpDef("prim.ceil")
- .add_arg("x", "The input value.")
+ .signature(sig::arg("x", "The input value."))
.set_attr<TCallEffectKind>("TCallEffectKind", static_cast<int64_t>(CallEffectKind::kPure))
.set_attr<bool>("TVectorizable", true);
OpDef("prim.log2")
- .add_arg("x", "The input value.")
+ .signature(sig::arg("x", "The input value."))
.set_attr<TCallEffectKind>("TCallEffectKind", static_cast<int64_t>(CallEffectKind::kPure))
.set_attr<bool>("TVectorizable", true);
OpDef("prim.clz")
- .add_arg("x", "The input value.")
+ .signature(sig::arg("x", "The input value."))
.set_attr<TCallEffectKind>("TCallEffectKind", static_cast<int64_t>(CallEffectKind::kPure));
}
diff --git a/src/relax/analysis/well_formed.cc b/src/relax/analysis/well_formed.cc
index 6146707eb8..ee81e4bafb 100644
--- a/src/relax/analysis/well_formed.cc
+++ b/src/relax/analysis/well_formed.cc
@@ -413,13 +413,12 @@ class WellFormedChecker : public relax::ExprVisitor, public relax::TypeVisitor {
}
}
- if (auto func_validate = op_map_validate_.get(call->op, nullptr); func_validate != nullptr) {
+ if (auto op = call->op.as<Op>()) {
try {
- func_validate(ffi::GetRef<Call>(call));
+ op.value().Validate(call);
} catch (std::exception& err) {
TVM_FFI_VISIT_THROW(ValueError, ffi::GetRef<Call>(call))
- << "Operator-specific validation (FValidate) for " << call->op
- << " identified error: \n"
+ << "Operator-specific validation for " << call->op << " identified error: \n"
<< err.what();
}
}
@@ -715,7 +714,6 @@ class WellFormedChecker : public relax::ExprVisitor, public relax::TypeVisitor {
std::unordered_map<Var, const FunctionNode*> param_var_func_map_;
tvm::OpAttrMap<FNormalize> op_map_normalize_ = Op::GetAttrMap<FNormalize>("FNormalize");
- tvm::OpAttrMap<FValidate> op_map_validate_ = Op::GetAttrMap<FValidate>("FValidate");
tvm::OpAttrMap<FInferType> op_map_infer_type_ = Op::GetAttrMap<FInferType>("FInferType");
};
diff --git a/src/relax/backend/adreno/annotate_custom_storage.cc b/src/relax/backend/adreno/annotate_custom_storage.cc
index 1748149407..1e7d9b6966 100644
--- a/src/relax/backend/adreno/annotate_custom_storage.cc
+++ b/src/relax/backend/adreno/annotate_custom_storage.cc
@@ -693,11 +693,11 @@ class DefineVDevice : ExprMutator {
}
if (call->op.same_as(call_tir_op)) {
- return builder_->Normalize(
- Call(Type::Missing(), call_tir_op, {gv, Tuple(new_args)}, call->attrs, {updated_ret_ty}));
+ return builder_->Normalize(Call::Unchecked(
+ Type::Missing(), call_tir_op, {gv, Tuple(new_args)}, call->attrs, {updated_ret_ty}));
} else {
return builder_->Normalize(
- Call(Type::Missing(), call->op, new_args, call->attrs, {updated_ret_ty}));
+ Call::Unchecked(Type::Missing(), call->op, new_args, call->attrs, {updated_ret_ty}));
}
}
@@ -733,7 +733,8 @@ class DefineVDevice : ExprMutator {
attrs->index = vdev->vdevice_id;
attrs->memory_scope = vdev->memory_scope;
- Expr new_arg = Call(Type::Missing(), hint_on_device_op_, {arg}, Attrs{std::move(attrs)}, {});
+ Expr new_arg =
+ Call::Unchecked(Type::Missing(), hint_on_device_op_, {arg}, Attrs{std::move(attrs)}, {});
return new_arg;
}
diff --git a/src/relax/backend/adreno/fold_vdevice_scope_change.cc b/src/relax/backend/adreno/fold_vdevice_scope_change.cc
index de39763fab..1b3f7c394a 100644
--- a/src/relax/backend/adreno/fold_vdevice_scope_change.cc
+++ b/src/relax/backend/adreno/fold_vdevice_scope_change.cc
@@ -88,7 +88,8 @@ std::tuple<DFPattern, ffi::TypedFunction<Expr(Expr, ffi::Map<DFPattern, Expr>)>>
auto shape_arr = tir_out_ty->GetShape().value();
auto new_ty = TensorType(ShapeExpr(shape_arr), tir_out_ty->dtype, vdev_attrs->dst_vdevice);
- return Call(Type::Missing(), call_tir->op, call_tir->args, call_tir->attrs, {new_ty});
+ return Call::Unchecked(Type::Missing(), call_tir->op, call_tir->args, call_tir->attrs,
+ {new_ty});
}
return expr;
};
diff --git a/src/relax/backend/vm/lower_runtime_builtin.cc b/src/relax/backend/vm/lower_runtime_builtin.cc
index 7cde67fc5e..9872cb81f8 100644
--- a/src/relax/backend/vm/lower_runtime_builtin.cc
+++ b/src/relax/backend/vm/lower_runtime_builtin.cc
@@ -87,8 +87,9 @@ class LowerRuntimeBuiltinMutator : public ExprMutator {
PrimExpr runtime_device_index = call->args[1].as_or_throw<PrimExpr>();
StringImm storage_scope = call->args[2].as_or_throw<StringImm>();
DataTypeImm output_dtype = DataTypeImm((DLDataType{kDLUInt, 8, 1}));
- return Call(Type::Missing(), vm_alloc_storage_op_,
- {call->args[0], runtime_device_index, output_dtype, storage_scope}, Attrs());
+ return Call::Unchecked(Type::Missing(), vm_alloc_storage_op_,
+ {call->args[0], runtime_device_index, output_dtype, storage_scope},
+ Attrs());
}
Expr MakeMemAllocTensor(const Call& call) {
@@ -100,12 +101,12 @@ class LowerRuntimeBuiltinMutator : public ExprMutator {
call_args.push_back(call->args[4]);
}
- return Call(Type::Missing(), vm_alloc_tensor_op_, call_args, Attrs());
+ return Call::Unchecked(Type::Missing(), vm_alloc_tensor_op_, call_args, Attrs());
}
Expr MakeMemKillObject(const Call& call) {
TVM_FFI_ICHECK_EQ(call->args.size(), 1);
- return Call(Type::Missing(), vm_kill_object_op_, {call->args[0]}, Attrs());
+ return Call::Unchecked(Type::Missing(), vm_kill_object_op_, {call->args[0]}, Attrs());
}
Expr CallTIRDyn(const Call& call_node) {
@@ -119,7 +120,7 @@ class LowerRuntimeBuiltinMutator : public ExprMutator {
for (Expr arg : tir_args->fields) {
args.push_back(arg);
}
- return Call(Type::Missing(), builtin_call_tir_dyn_, args, Attrs(), {void_ty_});
+ return Call::Unchecked(Type::Missing(), builtin_call_tir_dyn_, args, Attrs(), {void_ty_});
}
Expr Reshape(const Call& call_node) {
@@ -133,21 +134,23 @@ class LowerRuntimeBuiltinMutator : public ExprMutator {
<< "However, in expression " << call_node << ", the shape argument " << arg << " has type "
<< arg->ty;
- return Call(Type::Missing(), builtin_reshape_, call_node->args, Attrs(), {GetType(call_node)});
+ return Call::Unchecked(Type::Missing(), builtin_reshape_, call_node->args, Attrs(),
+ {GetType(call_node)});
}
Expr ShapeOf(const Call& call_node) {
TVM_FFI_ICHECK(call_node->args.size() == 1);
TVM_FFI_ICHECK(!call_node->ty.IsMissing());
- return Call(Type::Missing(), builtin_shape_of_, call_node->args, Attrs(), {GetType(call_node)});
+ return Call::Unchecked(Type::Missing(), builtin_shape_of_, call_node->args, Attrs(),
+ {GetType(call_node)});
}
Expr TensorToShape(const Call& call_node) {
TVM_FFI_ICHECK(call_node->args.size() == 1);
TVM_FFI_ICHECK(!call_node->ty.IsMissing());
- return Call(Type::Missing(), builtin_tensor_to_shape_, call_node->args, Attrs(),
- {GetType(call_node)});
+ return Call::Unchecked(Type::Missing(), builtin_tensor_to_shape_, call_node->args, Attrs(),
+ {GetType(call_node)});
}
Expr CallPyFunc(const Call& call_node) {
@@ -161,8 +164,8 @@ class LowerRuntimeBuiltinMutator : public ExprMutator {
auto combined_tuple = Tuple(tuple_fields);
// Direct call to vm.builtin.call_py_func
- return Call(Type::Missing(), builtin_call_py_func_, {combined_tuple}, call_node->attrs,
- call_node->ty_args, call_node->span);
+ return Call::Unchecked(Type::Missing(), builtin_call_py_func_, {combined_tuple},
+ call_node->attrs, call_node->ty_args, call_node->span);
}
Expr ToDevice(const Call& call_node) {
@@ -180,7 +183,8 @@ class LowerRuntimeBuiltinMutator : public ExprMutator {
args.push_back(IntImm::Int64(dev_type));
args.push_back(IntImm::Int64(dev_id));
args.push_back(storage_scope);
- return Call(Type::Missing(), builtin_to_device_, args, call_node->attrs, {GetType(call_node)});
+ return Call::Unchecked(Type::Missing(), builtin_to_device_, args, call_node->attrs,
+ {GetType(call_node)});
}
Expr MakeClosure(const Call& call_node) {
@@ -197,7 +201,7 @@ class LowerRuntimeBuiltinMutator : public ExprMutator {
args.push_back(arg);
}
- return Call(Type::Missing(), builtin_make_closure_, args, Attrs(), {object_ty_});
+ return Call::Unchecked(Type::Missing(), builtin_make_closure_, args, Attrs(), {object_ty_});
}
Expr InvokeClosure(const Call& call_node) {
@@ -214,8 +218,8 @@ class LowerRuntimeBuiltinMutator : public ExprMutator {
for (Expr arg : invoke_closure_args->fields) {
args.push_back(arg);
}
- return Call(Type::Missing(), call_builtin_with_ctx_op_, {builtin_invoke_closure_, Tuple(args)},
- Attrs(), {object_ty_});
+ return Call::Unchecked(Type::Missing(), call_builtin_with_ctx_op_,
+ {builtin_invoke_closure_, Tuple(args)}, Attrs(), {object_ty_});
}
const Op call_builtin_with_ctx_op_ = Op::Get("relax.call_builtin_with_ctx");
diff --git a/src/relax/backend/vm/vm_shape_lower.cc b/src/relax/backend/vm/vm_shape_lower.cc
index 8017dc4f49..8f1efc3b7f 100644
--- a/src/relax/backend/vm/vm_shape_lower.cc
+++ b/src/relax/backend/vm/vm_shape_lower.cc
@@ -437,8 +437,9 @@ class VMShapeLowerMutator
auto [code, rvalue] = MakeMatchArgs(slot->expr, false);
ffi::Array<Expr> args = {runtime_var, shape_heap_, IntImm::Int64(static_cast<int>(code)),
rvalue, GetErrContext(err_ctx)};
- builder_->Emit(Call(Type::Missing(), builtin_match_prim_value_, args, Attrs(), {void_ty_}),
- "_");
+ builder_->Emit(
+ Call::Unchecked(Type::Missing(), builtin_match_prim_value_, args, Attrs(), {void_ty_}),
+ "_");
this->EmitOutstandingPrimExprCompute();
}
@@ -751,7 +752,7 @@ class VMShapeLowerMutator
WithAttr<tirx::PrimFunc>(std::move(shape_func), tvm::tirx::attr::kIsHostFunc, true);
}
GlobalVar shape_func_var = builder_->AddFunction(shape_func, "shape_func");
- builder_->Emit(Call(Type::Missing(), shape_func_var, {shape_heap_}), "_");
+ builder_->Emit(Call::Unchecked(Type::Missing(), shape_func_var, {shape_heap_}), "_");
return to_compute.size();
}
//-------------------------------------------------------
@@ -792,8 +793,8 @@ class VMShapeLowerMutator
void VisitType_(const StringTypeNode* op, Expr value, bool always_check, bool dynamic_only,
const ffi::String& err_ctx, std::vector<MatchShapeTodoItem>* match_todos) final {
if (always_check || !IsBaseOf(StringType(), GetType(value))) {
- builder_->Emit(Call(Type::Missing(), ExternFunc("vm.builtin.check_string_info"),
- {value, GetErrContext(err_ctx)}, Attrs(), {void_ty_}),
+ builder_->Emit(Call::Unchecked(Type::Missing(), ExternFunc("vm.builtin.check_string_info"),
+ {value, GetErrContext(err_ctx)}, Attrs(), {void_ty_}),
"_");
}
}
@@ -839,8 +840,9 @@ class VMShapeLowerMutator
}
if (always_check || !IsBaseOf(TensorType(op->dtype, op->ndim), GetType(value))) {
// check_tensor_info(value, ndim, dtype, err_ctx)
- Expr dtype_arg = op->IsUnknownDtype() ? Expr(Call(Type::Missing(), null_value_op_, {}))
- : Expr(DataTypeImm(op->dtype.value()->dtype));
+ Expr dtype_arg = op->IsUnknownDtype()
+ ? Expr(Call::Unchecked(Type::Missing(), null_value_op_, {}))
+ : Expr(DataTypeImm(op->dtype.value()->dtype));
Call call(Type::Missing(), builtin_check_tensor_info_,
{value, IntImm::Int64(op->ndim), dtype_arg, GetErrContext(err_ctx)}, Attrs(),
{void_ty_});
diff --git a/src/relax/distributed/transform/lower_global_view_to_local_view.cc b/src/relax/distributed/transform/lower_global_view_to_local_view.cc
index 55a5f7f211..dfc45b0d95 100644
--- a/src/relax/distributed/transform/lower_global_view_to_local_view.cc
+++ b/src/relax/distributed/transform/lower_global_view_to_local_view.cc
@@ -427,8 +427,8 @@ class LowerTIRToLocalView : public ExprMutator {
if (allreduce_kind != "") {
ffi::ObjectPtr<AllReduceAttrs> attrs = ffi::make_object<AllReduceAttrs>();
attrs->op_type = allreduce_kind;
- new_call =
- Call(Type::Missing(), Op::Get("relax.ccl.allreduce"), {new_call}, Attrs(attrs), {});
+ new_call = Call::Unchecked(Type::Missing(), Op::Get("relax.ccl.allreduce"), {new_call},
+ Attrs(attrs), {});
}
ReEmitBinding(binding, this->builder_->Normalize(new_call));
}
diff --git a/src/relax/ir/block_builder.cc b/src/relax/ir/block_builder.cc
index 14f92b7f8d..ed219a36e8 100644
--- a/src/relax/ir/block_builder.cc
+++ b/src/relax/ir/block_builder.cc
@@ -659,7 +659,7 @@ class Normalizer : public BlockBuilderImpl, private ExprFunctor<Expr(const Expr&
if (new_op.same_as(op->op) && new_args.same_as(op->args)) {
call = ffi::GetRef<Call>(op);
} else {
- call = Call(Type::Missing(), new_op, new_args, op->attrs, op->ty_args);
+ call = Call::Unchecked(Type::Missing(), new_op, new_args, op->attrs, op->ty_args);
}
if (call->ty.IsMissing()) {
@@ -680,6 +680,11 @@ class Normalizer : public BlockBuilderImpl, private ExprFunctor<Expr(const Expr&
}
}
+ // TVMScript uses a builder with FNormalize disabled so it can construct
+ // intentionally ill-formed IR for a later check_well_formed call.
+ if (apply_f_normalize_) {
+ if (auto op = call->op.as<Op>()) op.value().Validate(call.get());
+ }
return call;
}
diff --git a/src/relax/op/ccl/ccl.cc b/src/relax/op/ccl/ccl.cc
index f26578afc1..77705b87a0 100644
--- a/src/relax/op/ccl/ccl.cc
+++ b/src/relax/op/ccl/ccl.cc
@@ -43,7 +43,7 @@ Expr allreduce(Expr x, ffi::String op_type, bool in_group) {
attrs->in_group = std::move(in_group);
static const Op op = Op::Get("relax.ccl.allreduce");
- return Call(Type::Missing(), op, {std::move(x)}, Attrs{attrs}, {});
+ return Call::Unchecked(Type::Missing(), op, {std::move(x)}, Attrs{attrs}, {});
}
TVM_FFI_STATIC_INIT_BLOCK() {
@@ -58,8 +58,8 @@ Type InferTypeAllReduce(const Call& call, const BlockBuilder& ctx) {
TVM_FFI_STATIC_INIT_BLOCK() {
OpDef("relax.ccl.allreduce")
- .add_arg("x", "Input to which allreduce will be applied.")
- .attrs_type<AllReduceAttrs>()
+ .signature(sig::arg("x", "Input to which allreduce will be applied."),
+ sig::call_attrs<AllReduceAttrs>())
.set_attr<FInferType>("FInferType", InferTypeAllReduce)
.set_attr<FRelaxInferLayout>("FRelaxInferLayout", InferLayoutUnaryEwise)
.set_attr<bool>("FPurity", true);
@@ -73,7 +73,7 @@ Expr allgather(Expr x, int num_workers, bool in_group) {
attrs->in_group = std::move(in_group);
static const Op op = Op::Get("relax.ccl.allgather");
- return Call(Type::Missing(), op, {std::move(x)}, Attrs{attrs}, {});
+ return Call::Unchecked(Type::Missing(), op, {std::move(x)}, Attrs{attrs}, {});
}
TVM_FFI_STATIC_INIT_BLOCK() {
@@ -99,7 +99,8 @@ Type InferTypeAllGather(const Call& call, const BlockBuilder& ctx) {
TVM_FFI_STATIC_INIT_BLOCK() {
OpDef("relax.ccl.allgather")
- .add_arg("x", "Input to which allgather will be applied.")
+ .signature(sig::arg("x", "Input to which allgather will be applied."),
+ sig::call_attrs<AllGatherAttrs>())
.set_attr<FInferType>("FInferType", InferTypeAllGather)
.set_attr<FRelaxInferLayout>("FRelaxInferLayout", InferLayoutUnaryEwise)
.set_attr<bool>("FPurity", true);
@@ -108,7 +109,7 @@ TVM_FFI_STATIC_INIT_BLOCK() {
/* relax.ccl.broadcast_from_worker0 */
Expr broadcast_from_worker0(Expr x) {
static const Op op = Op::Get("relax.ccl.broadcast_from_worker0");
- return Call(Type::Missing(), op, {std::move(x)}, {}, {});
+ return Call::Unchecked(Type::Missing(), op, {std::move(x)}, {}, {});
}
TVM_FFI_STATIC_INIT_BLOCK() {
@@ -123,7 +124,7 @@ Type InferTypeBroadcastFromZero(const Call& call, const BlockBuilder& ctx) {
TVM_FFI_STATIC_INIT_BLOCK() {
OpDef("relax.ccl.broadcast_from_worker0")
- .add_arg("x", "Input to be broadcast.")
+ .signature(sig::arg("x", "Input to be broadcast."))
.set_attr<FInferType>("FInferType", InferTypeBroadcastFromZero)
.set_attr<FRelaxInferLayout>("FRelaxInferLayout", InferLayoutUnaryEwise)
.set_attr<bool>("FPurity", true);
@@ -137,7 +138,7 @@ Expr scatter_from_worker0(Expr data, int num_workers, int axis) {
attrs->axis = std::move(axis);
static const Op op = Op::Get("relax.ccl.scatter_from_worker0");
- return Call(Type::Missing(), op, {std::move(data)}, Attrs{attrs}, {});
+ return Call::Unchecked(Type::Missing(), op, {std::move(data)}, Attrs{attrs}, {});
}
TVM_FFI_STATIC_INIT_BLOCK() {
@@ -171,9 +172,11 @@ Type InferTypeScatter(const Call& call, const BlockBuilder& ctx) {
TVM_FFI_STATIC_INIT_BLOCK() {
OpDef("relax.ccl.scatter_from_worker0")
- .add_arg("x",
- "The buffer to be divided into equal parts and sent to each worker accordingly.")
- .attrs_type<ScatterCollectiveAttrs>()
+ .signature(
+ sig::arg(
+ "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<bool>("FPurity", true);
}
diff --git a/src/relax/op/distributed/distributed.cc b/src/relax/op/distributed/distributed.cc
index 5075f8d354..591982204d 100644
--- a/src/relax/op/distributed/distributed.cc
+++ b/src/relax/op/distributed/distributed.cc
@@ -50,7 +50,7 @@ Expr annotate_sharding(Expr input, distributed::DeviceMesh device_mesh,
attrs->placement = placement;
static const Op op = Op::Get("relax.dist.annotate_sharding");
- return Call(Type::Missing(), op, {std::move(input)}, Attrs(attrs), {});
+ return Call::Unchecked(Type::Missing(), op, {std::move(input)}, Attrs(attrs), {});
}
TVM_FFI_STATIC_INIT_BLOCK() {
@@ -64,7 +64,7 @@ Type InferTypeAnnotateSharding(const Call& call, const BlockBuilder& ctx) {
TVM_FFI_STATIC_INIT_BLOCK() {
OpDef("relax.dist.annotate_sharding")
- .add_arg("input", "The input tensor.")
+ .signature(sig::arg("input", "The input tensor."), sig::call_attrs<DistributionAttrs>())
.set_attr<FInferType>("FInferType", InferTypeAnnotateSharding)
.set_attr<FInferType>("dist.FInferType", InferTypeAnnotateSharding)
.set_attr<bool>("FPurity", true);
@@ -79,7 +79,7 @@ Expr redistribute(Expr input, distributed::DeviceMesh device_mesh,
attrs->placement = placement;
static const Op op = Op::Get("relax.dist.redistribute");
- return Call(Type::Missing(), op, {std::move(input)}, Attrs(attrs), {});
+ return Call::Unchecked(Type::Missing(), op, {std::move(input)}, Attrs(attrs), {});
}
TVM_FFI_STATIC_INIT_BLOCK() {
@@ -96,7 +96,7 @@ Type InferDistTypeRedistribute(const Call& call, const BlockBuilder& ctx) {
TVM_FFI_STATIC_INIT_BLOCK() {
OpDef("relax.dist.redistribute")
- .add_arg("input", "The input tensor.")
+ .signature(sig::arg("input", "The input tensor."), sig::call_attrs<DistributionAttrs>())
.set_attr<FInferType>("dist.FInferType", InferDistTypeRedistribute)
.set_attr<bool>("FPurity", true);
}
@@ -114,9 +114,9 @@ Type InferTypeCallTIRLocalView(const Call& call, const BlockBuilder& ctx) {
TVM_FFI_STATIC_INIT_BLOCK() {
OpDef("relax.dist.call_tir_local_view")
- .add_arg("func", "The destination-passing-style function.")
- .add_arg("args", "The input arguments.")
- .add_ty_arg("out_type", "The output type.")
+ .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<bool>("FPurity", true);
}
@@ -138,7 +138,7 @@ Expr MakeCallTIRLocalView(Expr func, Tuple args, ffi::Array<distributed::DTensor
}
static const Op op = Op::Get("relax.dist.call_tir_local_view");
- return Call(Type::Missing(), op, {func, args}, {}, {out_ty});
+ return Call::Unchecked(Type::Missing(), op, {func, args}, {}, {out_ty});
}
TVM_FFI_STATIC_INIT_BLOCK() {
@@ -210,7 +210,7 @@ Expr redistribute_replica_to_shard(Expr input, int num_workers, int axis) {
attrs->axis = std::move(axis);
static const Op op = Op::Get("relax.dist.redistribute_replica_to_shard");
- return Call(Type::Missing(), op, {std::move(input)}, Attrs{attrs}, {});
+ return Call::Unchecked(Type::Missing(), op, {std::move(input)}, Attrs{attrs}, {});
}
TVM_FFI_STATIC_INIT_BLOCK() {
@@ -219,8 +219,8 @@ TVM_FFI_STATIC_INIT_BLOCK() {
redistribute_replica_to_shard);
OpDef("relax.dist.redistribute_replica_to_shard")
- .add_arg("input", "The buffer to be sliced.")
- .attrs_type<ScatterCollectiveAttrs>()
+ .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<bool>("FPurity", true);
diff --git a/src/relax/op/image/resize.cc b/src/relax/op/image/resize.cc
index 1476ddd74b..3a23426bd6 100644
--- a/src/relax/op/image/resize.cc
+++ b/src/relax/op/image/resize.cc
@@ -56,7 +56,7 @@ Expr resize2d(Expr data, Expr size, ffi::Array<FloatImm> roi, ffi::String layout
attrs->out_dtype = out_dtype;
static const Op op = Op::Get("relax.image.resize2d");
- return Call(Type::Missing(), op, {std::move(data), std::move(size)}, Attrs(attrs), {});
+ return Call::Unchecked(Type::Missing(), op, {std::move(data), std::move(size)}, Attrs(attrs), {});
}
TVM_FFI_STATIC_INIT_BLOCK() {
@@ -146,9 +146,8 @@ InferLayoutOutput InferLayoutResize2d(
TVM_FFI_STATIC_INIT_BLOCK() {
OpDef("relax.image.resize2d")
- .add_arg("data", "The input tensor.")
- .add_arg("size", "The output image shape.")
- .attrs_type<Resize2DAttrs>()
+ .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<FRelaxInferLayout>("FRelaxInferLayout", InferLayoutResize2d)
.set_attr<TMixedPrecisionPolicy>("TMixedPrecisionPolicy", MixedPrecisionPolicyKind::kFollow)
@@ -173,7 +172,7 @@ Expr resize3d(Expr data, Expr size, ffi::Array<FloatImm> roi, ffi::String layout
attrs->out_dtype = out_dtype;
static const Op op = Op::Get("relax.image.resize3d");
- return Call(Type::Missing(), op, {std::move(data), std::move(size)}, Attrs(attrs), {});
+ return Call::Unchecked(Type::Missing(), op, {std::move(data), std::move(size)}, Attrs(attrs), {});
}
TVM_FFI_STATIC_INIT_BLOCK() {
@@ -261,9 +260,8 @@ InferLayoutOutput InferLayoutResize3d(
TVM_FFI_STATIC_INIT_BLOCK() {
OpDef("relax.image.resize3d")
- .add_arg("data", "The input tensor.")
- .add_arg("size", "The output image shape.")
- .attrs_type<Resize3DAttrs>()
+ .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<FRelaxInferLayout>("FRelaxInferLayout", InferLayoutResize3d)
.set_attr<TMixedPrecisionPolicy>("TMixedPrecisionPolicy", MixedPrecisionPolicyKind::kFollow)
@@ -283,7 +281,7 @@ Expr grid_sample(Expr data, Expr grid, ffi::String method, ffi::String layout,
attrs->align_corners = align_corners;
static const Op op = Op::Get("relax.image.grid_sample");
- return Call(Type::Missing(), op, {std::move(data), std::move(grid)}, Attrs(attrs), {});
+ return Call::Unchecked(Type::Missing(), op, {std::move(data), std::move(grid)}, Attrs(attrs), {});
}
TVM_FFI_STATIC_INIT_BLOCK() {
@@ -353,9 +351,9 @@ Type InferTypeGridSample(const Call& call, const BlockBuilder& ctx) {
TVM_FFI_STATIC_INIT_BLOCK() {
OpDef("relax.image.grid_sample")
- .add_arg("data", "The input tensor.")
- .add_arg("grid", "The grid tensor for sampling.")
- .attrs_type<GridSampleAttrs>()
+ .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<TMixedPrecisionPolicy>("TMixedPrecisionPolicy", MixedPrecisionPolicyKind::kFollow)
.set_attr<bool>("FPurity", true);
@@ -367,7 +365,7 @@ Expr affine_grid(Expr data, Expr size, bool align_corners) {
ffi::ObjectPtr<AffineGridAttrs> attrs = ffi::make_object<AffineGridAttrs>();
attrs->align_corners = align_corners;
static const Op op = Op::Get("relax.image.affine_grid");
- return Call(Type::Missing(), op, {std::move(data), std::move(size)}, Attrs(attrs), {});
+ return Call::Unchecked(Type::Missing(), op, {std::move(data), std::move(size)}, Attrs(attrs), {});
}
TVM_FFI_STATIC_INIT_BLOCK() {
@@ -453,9 +451,9 @@ Type InferTypeAffineGrid(const Call& call, const BlockBuilder& ctx) {
TVM_FFI_STATIC_INIT_BLOCK() {
OpDef("relax.image.affine_grid")
- .add_arg("data", "The input affine matrix tensor.")
- .add_arg("size", "The target output shape (H, W).")
- .attrs_type<AffineGridAttrs>()
+ .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<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 5d25cf3219..af12fbd4b1 100644
--- a/src/relax/op/memory/view.cc
+++ b/src/relax/op/memory/view.cc
@@ -37,13 +37,13 @@ Expr view(Expr x, ffi::Optional<Expr> shape, ffi::Optional<Expr> dtype,
Tuple void_expr(ffi::Array<Expr>{});
static const Op op = Op::Get("relax.memory.view");
- return Call(Type::Missing(), op,
- {
- x,
- shape.value_or(void_expr),
- dtype.value_or(void_expr),
- relative_byte_offset.value_or(void_expr),
- });
+ return Call::Unchecked(Type::Missing(), op,
+ {
+ x,
+ shape.value_or(void_expr),
+ dtype.value_or(void_expr),
+ relative_byte_offset.value_or(void_expr),
+ });
}
TVM_FFI_STATIC_INIT_BLOCK() {
@@ -389,16 +389,16 @@ Expr LowerBuiltinView(const BlockBuilder& bb, const Call& call) {
ExternFunc runtime_view_func("runtime.TVMTensorCreateView", runtime_view_ty);
- return Call(Type::Missing(), runtime_view_func, {data, shape, dtype, relative_byte_offset});
+ return Call::Unchecked(Type::Missing(), runtime_view_func,
+ {data, shape, dtype, relative_byte_offset});
}
TVM_FFI_STATIC_INIT_BLOCK() {
OpDef("relax.memory.view")
- .add_arg("x", "The input tensor.")
- .add_arg("shape", "The view's shape.")
- .add_arg("dtype", "The view's data type.")
- .add_arg("relative_byte_offset",
- "The view's byte offset, relative to the input tensor's byte offset.")
+ .signature(sig::arg("x", "The input tensor."), sig::arg("shape", "The view's shape."),
+ sig::arg("dtype", "The view's data type."),
+ 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<bool>("FPurity", true)
@@ -407,7 +407,7 @@ TVM_FFI_STATIC_INIT_BLOCK() {
Expr ensure_zero_offset(const Expr& x) {
static const Op op = Op::Get("relax.memory.ensure_zero_offset");
- return Call(Type::Missing(), op, {x});
+ return Call::Unchecked(Type::Missing(), op, {x});
}
TVM_FFI_STATIC_INIT_BLOCK() {
@@ -426,12 +426,13 @@ Type InferTypeEnsureZeroOffset(const Call& call, const BlockBuilder& ctx) {
Expr LowerBuiltinEnsureZeroOffset(const BlockBuilder& bb, const Call& call) {
const ExternFunc builtin_ensure_zero_offset_{"vm.builtin.ensure_zero_offset"};
- return Call(Type::Missing(), builtin_ensure_zero_offset_, call->args, Attrs(), {GetType(call)});
+ return Call::Unchecked(Type::Missing(), builtin_ensure_zero_offset_, call->args, Attrs(),
+ {GetType(call)});
}
TVM_FFI_STATIC_INIT_BLOCK() {
OpDef("relax.memory.ensure_zero_offset")
- .add_arg("x", "The input tensor.")
+ .signature(sig::arg("x", "The input tensor."))
.set_attr<bool>("RequiresArgumentShapes", false)
.set_attr<FInferType>("FInferType", InferTypeEnsureZeroOffset)
.set_attr<bool>("FPurity", true)
diff --git a/src/relax/op/nn/attention.cc b/src/relax/op/nn/attention.cc
index b271db5659..bdab2755ca 100644
--- a/src/relax/op/nn/attention.cc
+++ b/src/relax/op/nn/attention.cc
@@ -38,12 +38,12 @@ Expr attention(Expr query, Expr key, Expr value, ffi::Optional<Expr> bias,
attrs->window_size = window_size;
if (bias) {
- return Call(Type::Missing(), Op::Get("relax.nn.attention_bias"),
- {std::move(query), std::move(key), std::move(value), bias.value()}, Attrs(attrs),
- {});
+ return Call::Unchecked(Type::Missing(), Op::Get("relax.nn.attention_bias"),
+ {std::move(query), std::move(key), std::move(value), bias.value()},
+ Attrs(attrs), {});
}
- return Call(Type::Missing(), Op::Get("relax.nn.attention"),
- {std::move(query), std::move(key), std::move(value)}, Attrs(attrs), {});
+ return Call::Unchecked(Type::Missing(), Op::Get("relax.nn.attention"),
+ {std::move(query), std::move(key), std::move(value)}, Attrs(attrs), {});
}
Expr attention_var_len(Expr query, Expr key, Expr value, Expr seqstart_q, Expr seqstart_k,
@@ -54,9 +54,9 @@ Expr attention_var_len(Expr query, Expr key, Expr value, Expr seqstart_q, Expr s
attrs->causal_mask = causal_mask;
attrs->window_size = window_size;
- return Call(Type::Missing(), Op::Get("relax.nn.attention_var_len"),
- {query, key, value, seqstart_q, seqstart_k, max_seqlen_q, max_seqlen_k}, Attrs(attrs),
- {});
+ return Call::Unchecked(Type::Missing(), Op::Get("relax.nn.attention_var_len"),
+ {query, key, value, seqstart_q, seqstart_k, max_seqlen_q, max_seqlen_k},
+ Attrs(attrs), {});
}
TVM_FFI_STATIC_INIT_BLOCK() {
@@ -151,35 +151,33 @@ Call InferMixedPrecisionAttention(const Call& call, DLDataType out_dtype) {
TVM_FFI_STATIC_INIT_BLOCK() {
OpDef("relax.nn.attention")
- .add_arg("query", "The input queries tensor.")
- .add_arg("key", "The input keys tensor.")
- .add_arg("value", "The input values tensor.")
- .attrs_type<AttentionAttrs>()
+ .signature(sig::arg("query", "The input queries tensor."),
+ sig::arg("key", "The input keys tensor."),
+ 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<bool>("FPurity", true);
OpDef("relax.nn.attention_bias")
- .add_arg("query", "The input queries tensor.")
- .add_arg("key", "The input keys tensor.")
- .add_arg("value", "The input values tensor.")
- .add_arg("bias", "The input bias tensor.")
- .attrs_type<AttentionAttrs>()
+ .signature(sig::arg("query", "The input queries tensor."),
+ sig::arg("key", "The input keys tensor."),
+ sig::arg("value", "The input values tensor."),
+ 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<bool>("FPurity", true);
OpDef("relax.nn.attention_var_len")
- .add_arg("query", "The input queries tensor.")
- .add_arg("key", "The input keys tensor.")
- .add_arg("value", "The input values tensor.")
- .add_arg("seqstart_q", "The cumsum of query sequence lengths, prepended with 0.")
- .add_arg("seqstart_k", "The cumsum of key sequence lengths, prepended with 0.")
- .add_arg("max_seqlen_q", "The maximum query sequence length in the batch.")
- .add_arg("max_seqlen_k", "The maximum key sequence length in the batch.")
- .attrs_type<AttentionAttrs>()
+ .signature(sig::arg("query", "The input queries tensor."),
+ sig::arg("key", "The input keys tensor."),
+ sig::arg("value", "The input values tensor."),
+ sig::arg("seqstart_q", "The cumsum of query sequence lengths, prepended with 0."),
+ sig::arg("seqstart_k", "The cumsum of key sequence lengths, prepended with 0."),
+ sig::arg("max_seqlen_q", "The maximum query sequence length in the batch."),
+ sig::arg("max_seqlen_k", "The maximum key sequence length in the batch."),
+ sig::call_attrs<AttentionAttrs>())
.set_attr<TMixedPrecisionPolicy>("TMixedPrecisionPolicy", MixedPrecisionPolicyKind::kAlways)
.set_attr<FInferMixedPrecision>("FInferMixedPrecision", InferMixedPrecisionAttention)
.set_attr<FInferType>("FInferType", InferTypeAttention)
diff --git a/src/relax/op/nn/convolution.cc b/src/relax/op/nn/convolution.cc
index e13541dc10..2c5f0373fd 100644
--- a/src/relax/op/nn/convolution.cc
+++ b/src/relax/op/nn/convolution.cc
@@ -198,10 +198,10 @@ Call InferMixedPrecisionConv1d(const Call& call, DLDataType out_dtype) {
TVM_FFI_STATIC_INIT_BLOCK() {
OpDef("relax.nn.conv1d")
- .add_arg("data", "The input tensor.")
- .add_arg("weight", "The weight tensor.")
- .add_ty_arg("out_type", "Optional output tensor type carrying the virtual device.")
- .attrs_type<Conv1DAttrs>()
+ .signature(
+ 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<FRelaxInferLayout>("FRelaxInferLayout", InferLayoutConv1d)
.set_attr<TMixedPrecisionPolicy>("TMixedPrecisionPolicy", MixedPrecisionPolicyKind::kAlways)
@@ -412,10 +412,10 @@ Call InferMixedPrecisionConv2d(const Call& call, DLDataType out_dtype) {
TVM_FFI_STATIC_INIT_BLOCK() {
OpDef("relax.nn.conv2d")
- .add_arg("data", "The input tensor.")
- .add_arg("weight", "The weight tensor.")
- .add_ty_arg("out_type", "Optional output tensor type carrying the virtual device.")
- .attrs_type<Conv2DAttrs>()
+ .signature(
+ 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<FRelaxInferLayout>("FRelaxInferLayout", InferLayoutConv2d)
.set_attr<TMixedPrecisionPolicy>("TMixedPrecisionPolicy", MixedPrecisionPolicyKind::kAlways)
@@ -600,10 +600,10 @@ Call InferMixedPrecisionConv3d(const Call& call, DLDataType out_dtype) {
TVM_FFI_STATIC_INIT_BLOCK() {
OpDef("relax.nn.conv3d")
- .add_arg("data", "The input tensor.")
- .add_arg("weight", "The weight tensor.")
- .add_ty_arg("out_type", "Optional output tensor type carrying the virtual device.")
- .attrs_type<Conv3DAttrs>()
+ .signature(
+ 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<FRelaxInferLayout>("FRelaxInferLayout", InferLayoutConv3d)
.set_attr<TMixedPrecisionPolicy>("TMixedPrecisionPolicy", MixedPrecisionPolicyKind::kAlways)
@@ -643,7 +643,7 @@ Expr conv1d_transpose(Expr data, Expr weight, ffi::Array<int64_t> strides,
attrs->out_layout = out_layout.value_or(data_layout);
attrs->out_dtype = out_dtype;
const Op op = Op::Get("relax.nn.conv1d_transpose");
- return Call(Type::Missing(), op, {data, weight}, Attrs(attrs), {});
+ return Call::Unchecked(Type::Missing(), op, {data, weight}, Attrs(attrs), {});
}
TVM_FFI_STATIC_INIT_BLOCK() {
@@ -782,10 +782,10 @@ Call InferMixedPrecisionConv1dTranspose(const Call& call, DLDataType out_dtype)
TVM_FFI_STATIC_INIT_BLOCK() {
OpDef("relax.nn.conv1d_transpose")
- .add_arg("data", "The input tensor.")
- .add_arg("weight", "The weight tensor.")
- .add_ty_arg("out_type", "Optional output tensor type carrying the virtual device.")
- .attrs_type<Conv1DTransposeAttrs>()
+ .signature(
+ 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<FRelaxInferLayout>("FRelaxInferLayout", InferLayoutConv1dTranspose)
.set_attr<TMixedPrecisionPolicy>("TMixedPrecisionPolicy", MixedPrecisionPolicyKind::kAlways)
@@ -836,7 +836,7 @@ Expr conv2d_transpose(Expr data, Expr weight, ffi::Array<int64_t> strides,
attrs->out_layout = out_layout.value_or(data_layout);
attrs->out_dtype = out_dtype;
const Op op = Op::Get("relax.nn.conv2d_transpose");
- return Call(Type::Missing(), op, {data, weight}, Attrs(attrs), {});
+ return Call::Unchecked(Type::Missing(), op, {data, weight}, Attrs(attrs), {});
}
TVM_FFI_STATIC_INIT_BLOCK() {
@@ -1013,10 +1013,10 @@ Call InferMixedPrecisionConv2dTranspose(const Call& call, DLDataType out_dtype)
TVM_FFI_STATIC_INIT_BLOCK() {
OpDef("relax.nn.conv2d_transpose")
- .add_arg("data", "The input tensor.")
- .add_arg("weight", "The weight tensor.")
- .add_ty_arg("out_type", "Optional output tensor type carrying the virtual device.")
- .attrs_type<Conv2DTransposeAttrs>()
+ .signature(
+ 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<FRelaxInferLayout>("FRelaxInferLayout", InferLayoutConv2dTranspose)
.set_attr<TMixedPrecisionPolicy>("TMixedPrecisionPolicy", MixedPrecisionPolicyKind::kAlways)
@@ -1070,7 +1070,7 @@ Expr conv3d_transpose(Expr data, Expr weight, ffi::Array<int64_t> strides,
attrs->out_layout = out_layout.value_or(data_layout);
attrs->out_dtype = out_dtype;
const Op op = Op::Get("relax.nn.conv3d_transpose");
- return Call(Type::Missing(), op, {data, weight}, Attrs(attrs), {});
+ return Call::Unchecked(Type::Missing(), op, {data, weight}, Attrs(attrs), {});
}
TVM_FFI_STATIC_INIT_BLOCK() {
@@ -1255,10 +1255,10 @@ Call InferMixedPrecisionConv3dTranspose(const Call& call, DLDataType out_dtype)
TVM_FFI_STATIC_INIT_BLOCK() {
OpDef("relax.nn.conv3d_transpose")
- .add_arg("data", "The input tensor.")
- .add_arg("weight", "The weight tensor.")
- .add_ty_arg("out_type", "Optional output tensor type carrying the virtual device.")
- .attrs_type<Conv3DTransposeAttrs>()
+ .signature(
+ 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<FRelaxInferLayout>("FRelaxInferLayout", InferLayoutConv3dTranspose)
.set_attr<TMixedPrecisionPolicy>("TMixedPrecisionPolicy", MixedPrecisionPolicyKind::kAlways)
diff --git a/src/relax/op/nn/convolution.h b/src/relax/op/nn/convolution.h
index 5b655b797a..d8312dcb32 100644
--- a/src/relax/op/nn/convolution.h
+++ b/src/relax/op/nn/convolution.h
@@ -50,7 +50,7 @@ inline Expr MakeConv(Expr data, Expr weight, ffi::Array<int64_t> strides,
attrs->out_layout = std::move(out_layout);
attrs->out_dtype = out_dtype;
const Op op = Op::Get(op_name);
- return Call(Type::Missing(), op, {data, weight}, Attrs(attrs), {});
+ return Call::Unchecked(Type::Missing(), op, {data, weight}, Attrs(attrs), {});
}
/*! \brief 1D convolution */
diff --git a/src/relax/op/nn/nn.cc b/src/relax/op/nn/nn.cc
index 38d93b4200..f814815033 100644
--- a/src/relax/op/nn/nn.cc
+++ b/src/relax/op/nn/nn.cc
@@ -48,35 +48,36 @@ TVM_FFI_STATIC_INIT_BLOCK() {
/* relax.nn.relu */
Expr relu(Expr x) {
static const Op op = Op::Get("relax.nn.relu");
- return Call(Type::Missing(), op, {std::move(x)}, Attrs(), {});
+ return Call::Unchecked(Type::Missing(), op, {std::move(x)}, Attrs(), {});
}
Expr gelu(Expr x) {
static const Op op = Op::Get("relax.nn.gelu");
- return Call(Type::Missing(), op, {std::move(x)}, Attrs(), {});
+ return Call::Unchecked(Type::Missing(), op, {std::move(x)}, Attrs(), {});
}
Expr gelu_tanh(Expr x) {
static const Op op = Op::Get("relax.nn.gelu_tanh");
- return Call(Type::Missing(), op, {std::move(x)}, Attrs(), {});
+ return Call::Unchecked(Type::Missing(), op, {std::move(x)}, Attrs(), {});
}
Expr selu(Expr x) {
static const Op op = Op::Get("relax.nn.selu");
- return Call(Type::Missing(), op, {std::move(x)}, Attrs(), {});
+ return Call::Unchecked(Type::Missing(), op, {std::move(x)}, Attrs(), {});
}
Expr silu(Expr x) {
static const Op op = Op::Get("relax.nn.silu");
- return Call(Type::Missing(), op, {std::move(x)}, Attrs(), {});
+ return Call::Unchecked(Type::Missing(), op, {std::move(x)}, Attrs(), {});
}
TVM_FFI_STATIC_INIT_BLOCK() {
tvm::ffi::reflection::GlobalDef().def("relax.op.nn.relu", relu);
OpDef("relax.nn.relu")
- .add_arg("x", "The input tensor.")
- .add_ty_arg("out_type", "Optional output tensor type carrying the virtual device.")
+ .signature(
+ sig::arg("x", "The input tensor."),
+ sig::var_ty_args("out_type", "Optional output tensor type carrying the virtual device."))
.set_attr<FRelaxInferLayout>("FRelaxInferLayout", InferLayoutUnaryEwise)
.set_attr<TMixedPrecisionPolicy>("TMixedPrecisionPolicy", MixedPrecisionPolicyKind::kFollow)
.set_attr<bool>("FPurity", true)
@@ -87,8 +88,9 @@ TVM_FFI_STATIC_INIT_BLOCK() {
tvm::ffi::reflection::GlobalDef().def("relax.op.nn.gelu", gelu);
OpDef("relax.nn.gelu")
- .add_arg("x", "The input tensor.")
- .add_ty_arg("out_type", "Optional output tensor type carrying the virtual device.")
+ .signature(
+ sig::arg("x", "The input tensor."),
+ sig::var_ty_args("out_type", "Optional output tensor type carrying the virtual device."))
.set_attr<FRelaxInferLayout>("FRelaxInferLayout", InferLayoutUnaryEwise)
.set_attr<TMixedPrecisionPolicy>("TMixedPrecisionPolicy", MixedPrecisionPolicyKind::kFollow)
.set_attr<bool>("FPurity", true)
@@ -99,8 +101,9 @@ TVM_FFI_STATIC_INIT_BLOCK() {
tvm::ffi::reflection::GlobalDef().def("relax.op.nn.gelu_tanh", gelu_tanh);
OpDef("relax.nn.gelu_tanh")
- .add_arg("x", "The input tensor.")
- .add_ty_arg("out_type", "Optional output tensor type carrying the virtual device.")
+ .signature(
+ sig::arg("x", "The input tensor."),
+ sig::var_ty_args("out_type", "Optional output tensor type carrying the virtual device."))
.set_attr<FRelaxInferLayout>("FRelaxInferLayout", InferLayoutUnaryEwise)
.set_attr<TMixedPrecisionPolicy>("TMixedPrecisionPolicy", MixedPrecisionPolicyKind::kFollow)
.set_attr<bool>("FPurity", true)
@@ -111,8 +114,9 @@ TVM_FFI_STATIC_INIT_BLOCK() {
tvm::ffi::reflection::GlobalDef().def("relax.op.nn.selu", selu);
OpDef("relax.nn.selu")
- .add_arg("x", "The input tensor.")
- .add_ty_arg("out_type", "Optional output tensor type carrying the virtual device.")
+ .signature(
+ sig::arg("x", "The input tensor."),
+ sig::var_ty_args("out_type", "Optional output tensor type carrying the virtual device."))
.set_attr<FRelaxInferLayout>("FRelaxInferLayout", InferLayoutUnaryEwise)
.set_attr<TMixedPrecisionPolicy>("TMixedPrecisionPolicy", MixedPrecisionPolicyKind::kFollow)
.set_attr<bool>("FPurity", true)
@@ -123,8 +127,9 @@ TVM_FFI_STATIC_INIT_BLOCK() {
tvm::ffi::reflection::GlobalDef().def("relax.op.nn.silu", silu);
OpDef("relax.nn.silu")
- .add_arg("x", "The input tensor.")
- .add_ty_arg("out_type", "Optional output tensor type carrying the virtual device.")
+ .signature(
+ sig::arg("x", "The input tensor."),
+ sig::var_ty_args("out_type", "Optional output tensor type carrying the virtual device."))
.set_attr<FRelaxInferLayout>("FRelaxInferLayout", InferLayoutUnaryEwise)
.set_attr<TMixedPrecisionPolicy>("TMixedPrecisionPolicy", MixedPrecisionPolicyKind::kFollow)
.set_attr<bool>("FPurity", true)
@@ -137,7 +142,7 @@ Expr leakyrelu(Expr data, double alpha) {
auto attrs = ffi::make_object<LeakyReluAttrs>();
attrs->alpha = alpha;
static const Op op = Op::Get("relax.nn.leakyrelu");
- return Call(Type::Missing(), op, {data}, Attrs(attrs), {});
+ return Call::Unchecked(Type::Missing(), op, {data}, Attrs(attrs), {});
}
TVM_FFI_STATIC_INIT_BLOCK() {
@@ -145,9 +150,10 @@ TVM_FFI_STATIC_INIT_BLOCK() {
refl::GlobalDef().def("relax.op.nn.leakyrelu", leakyrelu);
OpDef("relax.nn.leakyrelu")
- .add_arg("data", "The input tensor.")
- .add_ty_arg("out_type", "Optional output tensor type carrying the virtual device.")
- .attrs_type<LeakyReluAttrs>()
+ .signature(
+ 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<bool>("FPurity", true);
}
@@ -159,7 +165,7 @@ Expr softplus(Expr data, double beta, double threshold) {
attrs->beta = beta;
attrs->threshold = threshold;
static const Op op = Op::Get("relax.nn.softplus");
- return Call(Type::Missing(), op, {data}, Attrs(attrs), {});
+ return Call::Unchecked(Type::Missing(), op, {data}, Attrs(attrs), {});
}
TVM_FFI_STATIC_INIT_BLOCK() {
@@ -167,9 +173,10 @@ TVM_FFI_STATIC_INIT_BLOCK() {
refl::GlobalDef().def("relax.op.nn.softplus", softplus);
OpDef("relax.nn.softplus")
- .add_arg("data", "The input tensor.")
- .add_ty_arg("out_type", "Optional output tensor type carrying the virtual device.")
- .attrs_type<SoftplusAttrs>()
+ .signature(
+ 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<bool>("FPurity", true);
}
@@ -180,7 +187,7 @@ Expr prelu(Expr data, Expr alpha, int axis = 1) {
auto attrs = ffi::make_object<PReluAttrs>();
attrs->axis = axis;
static const Op op = Op::Get("relax.nn.prelu");
- return Call(Type::Missing(), op, {data, alpha}, Attrs(attrs), {});
+ return Call::Unchecked(Type::Missing(), op, {data, alpha}, Attrs(attrs), {});
}
TVM_FFI_STATIC_INIT_BLOCK() {
@@ -232,9 +239,9 @@ InferLayoutOutput InferLayoutPRelu(
TVM_FFI_STATIC_INIT_BLOCK() {
OpDef("relax.nn.prelu")
- .add_arg("data", "The input tensor.")
- .add_arg("alpha", "The channel-wise learnable slope.")
- .attrs_type<PReluAttrs>()
+ .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<FRelaxInferLayout>("FRelaxInferLayout", InferLayoutPRelu)
.set_attr<bool>("FPurity", true);
@@ -246,7 +253,7 @@ Expr softmax(Expr data, int axis) {
auto attrs = ffi::make_object<SoftmaxAttrs>();
attrs->axis = axis;
static const Op op = Op::Get("relax.nn.softmax");
- return Call(Type::Missing(), op, {data}, Attrs(attrs), {});
+ return Call::Unchecked(Type::Missing(), op, {data}, Attrs(attrs), {});
}
TVM_FFI_STATIC_INIT_BLOCK() {
@@ -299,8 +306,7 @@ InferLayoutOutput InferLayoutSoftmax(
TVM_FFI_STATIC_INIT_BLOCK() {
OpDef("relax.nn.softmax")
- .add_arg("data", "The input tensor.")
- .attrs_type<SoftmaxAttrs>()
+ .signature(sig::arg("data", "The input tensor."), sig::call_attrs<SoftmaxAttrs>())
.set_attr<FInferType>("FInferType", InferTypeSoftmax)
.set_attr<FRelaxInferLayout>("FRelaxInferLayout", InferLayoutSoftmax)
.set_attr<bool>("FPurity", true);
@@ -311,7 +317,7 @@ Expr log_softmax(Expr data, int axis) {
auto attrs = ffi::make_object<SoftmaxAttrs>();
attrs->axis = axis;
static const Op op = Op::Get("relax.nn.log_softmax");
- return Call(Type::Missing(), op, {data}, Attrs(attrs), {});
+ return Call::Unchecked(Type::Missing(), op, {data}, Attrs(attrs), {});
}
TVM_FFI_STATIC_INIT_BLOCK() {
@@ -319,8 +325,7 @@ TVM_FFI_STATIC_INIT_BLOCK() {
refl::GlobalDef().def("relax.op.nn.log_softmax", log_softmax);
OpDef("relax.nn.log_softmax")
- .add_arg("data", "The input tensor.")
- .attrs_type<SoftmaxAttrs>()
+ .signature(sig::arg("data", "The input tensor."), sig::call_attrs<SoftmaxAttrs>())
.set_attr<FInferType>("FInferType", InferTypeSoftmax)
.set_attr<bool>("FPurity", true);
}
@@ -333,7 +338,7 @@ Expr pad(Expr data, ffi::Array<int64_t> pad_width, ffi::String pad_mode, double
attrs->pad_mode = std::move(pad_mode);
attrs->pad_value = pad_value;
static const Op op = Op::Get("relax.nn.pad");
- return Call(Type::Missing(), op, {data}, Attrs(attrs), {});
+ return Call::Unchecked(Type::Missing(), op, {data}, Attrs(attrs), {});
}
TVM_FFI_STATIC_INIT_BLOCK() {
@@ -367,8 +372,7 @@ Type InferTypePad(const Call& call, const BlockBuilder& ctx) {
TVM_FFI_STATIC_INIT_BLOCK() {
OpDef("relax.nn.pad")
- .add_arg("data", "The input tensor.")
- .attrs_type<PadAttrs>()
+ .signature(sig::arg("data", "The input tensor."), sig::call_attrs<PadAttrs>())
.set_attr<FInferType>("FInferType", InferTypePad)
.set_attr<bool>("FPurity", true);
}
@@ -379,7 +383,7 @@ Expr pixel_shuffle(Expr data, int upscale_factor) {
auto attrs = ffi::make_object<PixelShuffleAttrs>();
attrs->upscale_factor = upscale_factor;
static const Op op = Op::Get("relax.nn.pixel_shuffle");
- return Call(Type::Missing(), op, {data}, Attrs(attrs), {});
+ return Call::Unchecked(Type::Missing(), op, {data}, Attrs(attrs), {});
}
TVM_FFI_STATIC_INIT_BLOCK() {
@@ -440,8 +444,7 @@ Type InferTypePixelShuffle(const Call& call, const BlockBuilder& ctx) {
TVM_FFI_STATIC_INIT_BLOCK() {
OpDef("relax.nn.pixel_shuffle")
- .add_arg("data", "The input tensor.")
- .attrs_type<PixelShuffleAttrs>()
+ .signature(sig::arg("data", "The input tensor."), sig::call_attrs<PixelShuffleAttrs>())
.set_attr<FInferType>("FInferType", InferTypePixelShuffle)
.set_attr<bool>("FPurity", true);
}
@@ -529,10 +532,10 @@ Expr batch_norm(Expr data, Expr gamma, Expr beta, Expr moving_mean, Expr moving_
attrs->training = training;
static const Op op = Op::Get("relax.nn.batch_norm");
- return Call(Type::Missing(), op,
- {std::move(data), std::move(gamma), std::move(beta), std::move(moving_mean),
- std::move(moving_var)},
- Attrs{attrs}, {});
+ return Call::Unchecked(Type::Missing(), op,
+ {std::move(data), std::move(gamma), std::move(beta),
+ std::move(moving_mean), std::move(moving_var)},
+ Attrs{attrs}, {});
}
TVM_FFI_STATIC_INIT_BLOCK() {
namespace refl = tvm::ffi::reflection;
@@ -588,12 +591,11 @@ InferLayoutOutput InferLayoutBatchNorm(
TVM_FFI_STATIC_INIT_BLOCK() {
OpDef("relax.nn.batch_norm")
- .add_arg("data", "Input to which batch_norm will be applied.")
- .add_arg("gamma", "The gamma scale factor.")
- .add_arg("beta", "The beta offset factor.")
- .add_arg("moving_mean", "Running mean of input.")
- .add_arg("moving_var", "Running variance of input.")
- .attrs_type<BatchNormAttrs>()
+ .signature(
+ sig::arg("data", "Input to which batch_norm will be applied."),
+ 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<FRelaxInferLayout>("FRelaxInferLayout", InferLayoutBatchNorm)
.set_attr<bool>("FPurity", true);
@@ -610,8 +612,8 @@ Expr layer_norm(Expr data, Expr gamma, Expr beta, ffi::Array<int64_t> axes, doub
attrs->scale = scale;
static const Op op = Op::Get("relax.nn.layer_norm");
- return Call(Type::Missing(), op, {std::move(data), std::move(gamma), std::move(beta)},
- Attrs{attrs}, {});
+ return Call::Unchecked(Type::Missing(), op, {std::move(data), std::move(gamma), std::move(beta)},
+ Attrs{attrs}, {});
}
TVM_FFI_STATIC_INIT_BLOCK() {
@@ -658,10 +660,9 @@ InferLayoutOutput InferLayoutLayerNorm(
TVM_FFI_STATIC_INIT_BLOCK() {
OpDef("relax.nn.layer_norm")
- .add_arg("data", "Input to which layer_norm will be applied.")
- .add_arg("gamma", "The gamma scale factor.")
- .add_arg("beta", "The beta offset factor.")
- .attrs_type<LayerNormAttrs>()
+ .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<FRelaxInferLayout>("FRelaxInferLayout", InferLayoutLayerNorm)
.set_attr<TMixedPrecisionPolicy>("TMixedPrecisionPolicy", MixedPrecisionPolicyKind::kFollow)
@@ -681,8 +682,8 @@ Expr group_norm(Expr data, Expr gamma, Expr beta, int num_groups, int channel_ax
attrs->scale = scale;
static const Op op = Op::Get("relax.nn.group_norm");
- return Call(Type::Missing(), op, {std::move(data), std::move(gamma), std::move(beta)},
- Attrs{attrs}, {});
+ return Call::Unchecked(Type::Missing(), op, {std::move(data), std::move(gamma), std::move(beta)},
+ Attrs{attrs}, {});
}
TVM_FFI_STATIC_INIT_BLOCK() {
@@ -773,10 +774,9 @@ InferLayoutOutput InferLayoutGroupNorm(
TVM_FFI_STATIC_INIT_BLOCK() {
OpDef("relax.nn.group_norm")
- .add_arg("data", "Input to which group_norm will be applied.")
- .add_arg("gamma", "The gamma scale factor.")
- .add_arg("beta", "The beta offset factor.")
- .attrs_type<GroupNormAttrs>()
+ .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<FRelaxInferLayout>("FRelaxInferLayout", InferLayoutGroupNorm)
.set_attr<TMixedPrecisionPolicy>("TMixedPrecisionPolicy", MixedPrecisionPolicyKind::kFollow)
@@ -795,8 +795,8 @@ Expr instance_norm(Expr data, Expr gamma, Expr beta, int channel_axis, ffi::Arra
attrs->scale = scale;
static const Op op = Op::Get("relax.nn.instance_norm");
- return Call(Type::Missing(), op, {std::move(data), std::move(gamma), std::move(beta)},
- Attrs{attrs}, {});
+ return Call::Unchecked(Type::Missing(), op, {std::move(data), std::move(gamma), std::move(beta)},
+ Attrs{attrs}, {});
}
TVM_FFI_STATIC_INIT_BLOCK() {
@@ -876,10 +876,9 @@ InferLayoutOutput InferLayoutInstanceNorm(
TVM_FFI_STATIC_INIT_BLOCK() {
OpDef("relax.nn.instance_norm")
- .add_arg("data", "Input to which instance_norm will be applied.")
- .add_arg("gamma", "The gamma scale factor.")
- .add_arg("beta", "The beta offset factor.")
- .attrs_type<InstanceNormAttrs>()
+ .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<FRelaxInferLayout>("FRelaxInferLayout", InferLayoutInstanceNorm)
.set_attr<TMixedPrecisionPolicy>("TMixedPrecisionPolicy", MixedPrecisionPolicyKind::kFollow)
@@ -893,7 +892,8 @@ Expr rms_norm(Expr data, Expr weight, ffi::Array<int64_t> axes, double epsilon)
attrs->epsilon = epsilon;
static const Op op = Op::Get("relax.nn.rms_norm");
- return Call(Type::Missing(), op, {std::move(data), std::move(weight)}, Attrs{attrs}, {});
+ return Call::Unchecked(Type::Missing(), op, {std::move(data), std::move(weight)}, Attrs{attrs},
+ {});
}
TVM_FFI_STATIC_INIT_BLOCK() {
@@ -937,9 +937,8 @@ InferLayoutOutput InferLayoutRMSNorm(
TVM_FFI_STATIC_INIT_BLOCK() {
OpDef("relax.nn.rms_norm")
- .add_arg("data", "Input to which rms_norm will be applied.")
- .add_arg("weight", "The scale factor.")
- .attrs_type<RMSNormAttrs>()
+ .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<FRelaxInferLayout>("FRelaxInferLayout", InferLayoutRMSNorm)
.set_attr<TMixedPrecisionPolicy>("TMixedPrecisionPolicy", MixedPrecisionPolicyKind::kFollow)
@@ -953,7 +952,7 @@ Expr dropout(Expr data, double rate) {
attrs->rate = rate;
static const Op op = Op::Get("relax.nn.dropout");
- return Call(Type::Missing(), op, {std::move(data)}, Attrs{attrs}, {});
+ return Call::Unchecked(Type::Missing(), op, {std::move(data)}, Attrs{attrs}, {});
}
TVM_FFI_STATIC_INIT_BLOCK() {
@@ -968,8 +967,8 @@ Type InferTypeDropout(const Call& call, const BlockBuilder& ctx) {
TVM_FFI_STATIC_INIT_BLOCK() {
OpDef("relax.nn.dropout")
- .add_arg("data", "Input to which dropout will be applied.")
- .attrs_type<DropoutAttrs>()
+ .signature(sig::arg("data", "Input to which dropout will be applied."),
+ sig::call_attrs<DropoutAttrs>())
.set_attr<FInferType>("FInferType", InferTypeDropout)
.set_attr<FRelaxInferLayout>("FRelaxInferLayout", InferLayoutUnaryEwise)
.set_attr<TMixedPrecisionPolicy>("TMixedPrecisionPolicy", MixedPrecisionPolicyKind::kFollow)
@@ -1023,7 +1022,7 @@ Type InferTypeCrossEntropy(const Call& call, const BlockBuilder& ctx) {
Expr cross_entropy_with_logits(Expr predictions, Expr labels) {
static const Op op = Op::Get("relax.nn.cross_entropy_with_logits");
- return Call(Type::Missing(), op, {std::move(predictions), std::move(labels)}, {}, {});
+ return Call::Unchecked(Type::Missing(), op, {std::move(predictions), std::move(labels)}, {}, {});
}
TVM_FFI_STATIC_INIT_BLOCK() {
@@ -1031,9 +1030,9 @@ TVM_FFI_STATIC_INIT_BLOCK() {
refl::GlobalDef().def("relax.op.nn.cross_entropy_with_logits", cross_entropy_with_logits);
OpDef("relax.nn.cross_entropy_with_logits")
- .add_arg("predictions", "The predictions.")
- .add_arg("labels", "The labels.")
- .add_ty_arg("out_type", "Optional output tensor type carrying the virtual device.")
+ .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<bool>("FPurity", true);
}
@@ -1054,11 +1053,12 @@ Expr nll_loss(Expr predictions, Expr targets, ffi::Optional<Expr> weights, ffi::
static const Op op = Op::Get("relax.nn.nll_loss");
if (weights.has_value()) {
- return Call(Type::Missing(), op, {std::move(predictions), std::move(targets), weights.value()},
- Attrs{attrs}, {});
+ return Call::Unchecked(Type::Missing(), op,
+ {std::move(predictions), std::move(targets), weights.value()},
+ Attrs{attrs}, {});
} else {
- return Call(Type::Missing(), op, {std::move(predictions), std::move(targets)}, Attrs{attrs},
- {});
+ return Call::Unchecked(Type::Missing(), op, {std::move(predictions), std::move(targets)},
+ Attrs{attrs}, {});
}
}
@@ -1259,11 +1259,11 @@ Type InferTypeNLLLoss(const Call& call, const BlockBuilder& ctx) {
TVM_FFI_STATIC_INIT_BLOCK() {
OpDef("relax.nn.nll_loss", "Optional weights: The weight of each target values.")
- .add_arg("predictions", "The prediction tensor.")
- .add_arg("targets", "The target tensor.")
- .allow_extra_args()
- .add_ty_arg("out_type", "Optional output tensor type carrying the virtual device.")
- .attrs_type<NLLLossAttrs>()
+ .signature(
+ sig::arg("predictions", "The prediction tensor."),
+ 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<bool>("FPurity", true);
}
@@ -1272,7 +1272,7 @@ TVM_FFI_STATIC_INIT_BLOCK() {
Expr batch_flatten(Expr data) {
static const Op op = Op::Get("relax.nn.batch_flatten");
- return Call(Type::Missing(), op, {std::move(data)}, {}, {});
+ return Call::Unchecked(Type::Missing(), op, {std::move(data)}, {}, {});
}
TVM_FFI_STATIC_INIT_BLOCK() {
@@ -1313,7 +1313,7 @@ Type InferTypeBatchFlatten(const Call& call, const BlockBuilder& ctx) {
TVM_FFI_STATIC_INIT_BLOCK() {
OpDef("relax.nn.batch_flatten")
- .add_arg("data", "The input tensor.")
+ .signature(sig::arg("data", "The input tensor."))
.set_attr<FInferType>("FInferType", 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 45f5be6b06..2d2d1a78ab 100644
--- a/src/relax/op/nn/pooling.cc
+++ b/src/relax/op/nn/pooling.cc
@@ -65,7 +65,7 @@ Expr MakePool1d(ffi::String op_name, Expr data, ffi::Array<int64_t> pool_size,
attrs->layout = layout;
attrs->out_layout = out_layout.value_or(layout);
const Op op = Op::Get(op_name);
- return Call(Type::Missing(), op, {std::move(data)}, Attrs(attrs), {});
+ return Call::Unchecked(Type::Missing(), op, {std::move(data)}, Attrs(attrs), {});
}
Expr max_pool1d(Expr data, ffi::Array<int64_t> pool_size, ffi::Array<int64_t> strides,
@@ -146,8 +146,7 @@ InferLayoutOutput InferLayoutPool1d(
TVM_FFI_STATIC_INIT_BLOCK() {
OpDef("relax.nn.max_pool1d")
- .add_arg("data", "The input tensor")
- .attrs_type<Pool1DAttrs>()
+ .signature(sig::arg("data", "The input tensor"), sig::call_attrs<Pool1DAttrs>())
.set_attr<FInferType>("FInferType", InferTypePool1D)
.set_attr<FRelaxInferLayout>("FRelaxInferLayout", InferLayoutPool1d)
.set_attr<TMixedPrecisionPolicy>("TMixedPrecisionPolicy", MixedPrecisionPolicyKind::kFollow)
@@ -190,7 +189,7 @@ Expr MakePool2d(ffi::String op_name, Expr data, ffi::Array<int64_t> pool_size,
attrs->layout = layout;
attrs->out_layout = out_layout.value_or(layout);
const Op op = Op::Get(op_name);
- return Call(Type::Missing(), op, {std::move(data)}, Attrs(attrs), {});
+ return Call::Unchecked(Type::Missing(), op, {std::move(data)}, Attrs(attrs), {});
}
Expr max_pool2d(Expr data, ffi::Array<int64_t> pool_size, ffi::Array<int64_t> strides,
@@ -301,8 +300,9 @@ InferLayoutOutput InferLayoutPool2d(
TVM_FFI_STATIC_INIT_BLOCK() {
OpDef("relax.nn.max_pool2d")
- .add_arg("data", "The input tensor")
- .attrs_type<Pool2DAttrs>()
+ .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<FRelaxInferLayout>("FRelaxInferLayout", InferLayoutPool2d)
.set_attr<TMixedPrecisionPolicy>("TMixedPrecisionPolicy", MixedPrecisionPolicyKind::kFollow)
@@ -348,7 +348,7 @@ Expr MakePool3d(ffi::String op_name, Expr data, ffi::Array<int64_t> pool_size,
attrs->layout = layout;
attrs->out_layout = out_layout.value_or(layout);
const Op op = Op::Get(op_name);
- return Call(Type::Missing(), op, {std::move(data)}, Attrs(attrs), {});
+ return Call::Unchecked(Type::Missing(), op, {std::move(data)}, Attrs(attrs), {});
}
Expr max_pool3d(Expr data, ffi::Array<int64_t> pool_size, ffi::Array<int64_t> strides,
@@ -451,8 +451,7 @@ InferLayoutOutput InferLayoutPool3d(
TVM_FFI_STATIC_INIT_BLOCK() {
OpDef("relax.nn.max_pool3d")
- .add_arg("data", "The input tensor")
- .attrs_type<Pool3DAttrs>()
+ .signature(sig::arg("data", "The input tensor"), sig::call_attrs<Pool3DAttrs>())
.set_attr<FInferType>("FInferType", InferTypePool3D)
.set_attr<FRelaxInferLayout>("FRelaxInferLayout", InferLayoutPool3d)
.set_attr<TMixedPrecisionPolicy>("TMixedPrecisionPolicy", MixedPrecisionPolicyKind::kFollow)
@@ -472,8 +471,7 @@ TVM_FFI_STATIC_INIT_BLOCK() {
refl::GlobalDef().def("relax.op.nn.avg_pool1d", avg_pool1d);
OpDef("relax.nn.avg_pool1d")
- .add_arg("data", "The input tensor")
- .attrs_type<Pool1DAttrs>()
+ .signature(sig::arg("data", "The input tensor"), sig::call_attrs<Pool1DAttrs>())
.set_attr<FInferType>("FInferType", InferTypePool1D)
.set_attr<FRelaxInferLayout>("FRelaxInferLayout", InferLayoutPool1d)
.set_attr<TMixedPrecisionPolicy>("TMixedPrecisionPolicy", MixedPrecisionPolicyKind::kFollow)
@@ -493,8 +491,7 @@ TVM_FFI_STATIC_INIT_BLOCK() {
refl::GlobalDef().def("relax.op.nn.avg_pool2d", avg_pool2d);
OpDef("relax.nn.avg_pool2d")
- .add_arg("data", "The input tensor")
- .attrs_type<Pool2DAttrs>()
+ .signature(sig::arg("data", "The input tensor"), sig::call_attrs<Pool2DAttrs>())
.set_attr<FInferType>("FInferType", InferTypePool2D)
.set_attr<FRelaxInferLayout>("FRelaxInferLayout", InferLayoutPool2d)
.set_attr<TMixedPrecisionPolicy>("TMixedPrecisionPolicy", MixedPrecisionPolicyKind::kFollow)
@@ -514,8 +511,7 @@ TVM_FFI_STATIC_INIT_BLOCK() {
refl::GlobalDef().def("relax.op.nn.avg_pool3d", avg_pool3d);
OpDef("relax.nn.avg_pool3d")
- .add_arg("data", "The input tensor")
- .attrs_type<Pool3DAttrs>()
+ .signature(sig::arg("data", "The input tensor"), sig::call_attrs<Pool3DAttrs>())
.set_attr<FInferType>("FInferType", InferTypePool3D)
.set_attr<FRelaxInferLayout>("FRelaxInferLayout", InferLayoutPool3d)
.set_attr<TMixedPrecisionPolicy>("TMixedPrecisionPolicy", MixedPrecisionPolicyKind::kFollow)
@@ -538,7 +534,7 @@ Expr adaptive_avg_pool1d(Expr data, ffi::Optional<ffi::Array<int64_t>> output_si
}
static const Op op = Op::Get("relax.nn.adaptive_avg_pool1d");
- return Call(Type::Missing(), op, {std::move(data)}, Attrs(attrs), {});
+ return Call::Unchecked(Type::Missing(), op, {std::move(data)}, Attrs(attrs), {});
}
TVM_FFI_STATIC_INIT_BLOCK() {
@@ -597,8 +593,7 @@ InferLayoutOutput InferLayoutAdaptiveAvgPool1D(
TVM_FFI_STATIC_INIT_BLOCK() {
OpDef("relax.nn.adaptive_avg_pool1d")
- .add_arg("data", "The input tensor")
- .attrs_type<AdaptivePool1DAttrs>()
+ .signature(sig::arg("data", "The input tensor"), sig::call_attrs<AdaptivePool1DAttrs>())
.set_attr<FInferType>("FInferType", InferTypeAdaptiveAvgPool1D)
.set_attr<FRelaxInferLayout>("FRelaxInferLayout", InferLayoutAdaptiveAvgPool1D)
.set_attr<TMixedPrecisionPolicy>("TMixedPrecisionPolicy", MixedPrecisionPolicyKind::kFollow)
@@ -624,7 +619,7 @@ Expr adaptive_avg_pool2d(Expr data, ffi::Optional<ffi::Array<int64_t>> output_si
}
static const Op op = Op::Get("relax.nn.adaptive_avg_pool2d");
- return Call(Type::Missing(), op, {std::move(data)}, Attrs(attrs), {});
+ return Call::Unchecked(Type::Missing(), op, {std::move(data)}, Attrs(attrs), {});
}
TVM_FFI_STATIC_INIT_BLOCK() {
@@ -701,8 +696,9 @@ InferLayoutOutput InferLayoutAdaptiveAvgPool2D(
TVM_FFI_STATIC_INIT_BLOCK() {
OpDef("relax.nn.adaptive_avg_pool2d")
- .add_arg("data", "The input tensor")
- .attrs_type<AdaptivePool2DAttrs>()
+ .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<FRelaxInferLayout>("FRelaxInferLayout", InferLayoutAdaptiveAvgPool2D)
.set_attr<TMixedPrecisionPolicy>("TMixedPrecisionPolicy", MixedPrecisionPolicyKind::kFollow)
@@ -728,7 +724,7 @@ Expr adaptive_avg_pool3d(Expr data, ffi::Optional<ffi::Array<int64_t>> output_si
}
static const Op op = Op::Get("relax.nn.adaptive_avg_pool3d");
- return Call(Type::Missing(), op, {std::move(data)}, Attrs(attrs), {});
+ return Call::Unchecked(Type::Missing(), op, {std::move(data)}, Attrs(attrs), {});
}
TVM_FFI_STATIC_INIT_BLOCK() {
@@ -789,8 +785,7 @@ InferLayoutOutput InferLayoutAdaptiveAvgPool3D(
TVM_FFI_STATIC_INIT_BLOCK() {
OpDef("relax.nn.adaptive_avg_pool3d")
- .add_arg("data", "The input tensor")
- .attrs_type<AdaptivePool3DAttrs>()
+ .signature(sig::arg("data", "The input tensor"), sig::call_attrs<AdaptivePool3DAttrs>())
.set_attr<FInferType>("FInferType", InferTypeAdaptiveAvgPool3D)
.set_attr<FRelaxInferLayout>("FRelaxInferLayout", InferLayoutAdaptiveAvgPool3D)
.set_attr<TMixedPrecisionPolicy>("TMixedPrecisionPolicy", MixedPrecisionPolicyKind::kFollow)
diff --git a/src/relax/op/op.cc b/src/relax/op/op.cc
index ed913eee58..9dcf647cbc 100644
--- a/src/relax/op/op.cc
+++ b/src/relax/op/op.cc
@@ -25,6 +25,8 @@
#include <tvm/relax/expr.h>
#include <tvm/relax/utils.h>
+#include <exception>
+
#include "../transform/utils.h"
#include "op_common.h"
@@ -113,13 +115,47 @@ Type InferTypeCallPurePacked(const Call& call, const BlockBuilder& ctx) {
}
}
+template <void (*Fn)(const CallNode*)>
+TVM_FFI_INLINE ffi::Expected<void> ValidateAsExpected(const CallNode* call) noexcept {
+ try {
+ Fn(call);
+ return {};
+ } catch (const ffi::Error& err) {
+ return ffi::Unexpected(err);
+ } catch (const std::exception& err) {
+ return ffi::Unexpected(ffi::Error("InternalError", err.what(), ""));
+ }
+}
+
+void ValidateCallPurePacked(const CallNode* call) {
+ TVM_FFI_CHECK(call->args.size() >= 1, TypeError)
+ << "call_pure_packed expects a function argument";
+ for (const ffi::Any* arg = call->args.GetArrayObj()->begin();
+ arg != call->args.GetArrayObj()->end(); ++arg) {
+ TVM_FFI_CHECK(
+ *arg != nullptr && ffi::details::AnyUnsafe::CheckAnyViewStrict<Expr>(ffi::AnyView(*arg)),
+ TypeError)
+ << "call_pure_packed has an invalid value argument";
+ }
+ for (const ffi::Any* ty_arg = call->ty_args.GetArrayObj()->begin();
+ ty_arg != call->ty_args.GetArrayObj()->end(); ++ty_arg) {
+ TVM_FFI_CHECK(*ty_arg != nullptr &&
+ ffi::details::AnyUnsafe::CheckAnyViewStrict<Type>(ffi::AnyView(*ty_arg)),
+ TypeError)
+ << "call_pure_packed has an invalid type argument";
+ }
+}
+
TVM_FFI_STATIC_INIT_BLOCK() {
OpDef("relax.call_pure_packed")
- .add_arg("func",
- "The first argument is the function being called. The rest are the "
- "arguments to that function.")
- .allow_extra_args()
- .add_ty_arg("type_args", "Optional type arguments forwarded to the callee.")
+ .set_validator(
+ ffi::reflection::NativeFunctionView<ffi::Expected<void>(
+ const CallNode*)>::FromNative<&ValidateAsExpected<ValidateCallPurePacked>>())
+ .signature(sig::arg("func",
+ "The first argument is the function being called. The rest are the "
+ "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<bool>("FPurity", true);
}
@@ -131,7 +167,7 @@ Expr MakeCallPurePacked(const Expr& callee, ffi::Array<Expr> args, const Attrs&
for (auto arg : args) {
call_args.push_back(arg);
}
- return Call(Type::Missing(), op, call_args, attrs, ty_args);
+ return Call::Unchecked(Type::Missing(), op, call_args, attrs, ty_args);
}
TVM_FFI_STATIC_INIT_BLOCK() {
@@ -229,12 +265,12 @@ Type InferTypeCallInplacePacked(const Call& call, const BlockBuilder& ctx) {
TVM_FFI_STATIC_INIT_BLOCK() {
OpDef("relax.call_inplace_packed")
- .add_arg("func",
- "The first argument is the function being called. The rest are the "
- "arguments to that function.")
- .allow_extra_args()
- .add_ty_arg("type_args", "Optional type arguments forwarded to the callee.")
- .attrs_type<CallInplacePackedAttrs>()
+ .signature(sig::arg("func",
+ "The first argument is the function being called. The rest are the "
+ "arguments to that function."),
+ 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)
// 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
@@ -251,7 +287,7 @@ Expr MakeCallInplacePacked(Expr func, ffi::Array<Expr> args, ffi::Array<int64_t>
static const Op op = Op::Get("relax.call_inplace_packed");
ffi::Array<Expr> call_args = {func};
call_args.insert(call_args.end(), args.begin(), args.end());
- return Call(Type::Missing(), op, call_args, Attrs(attrs), ty_args);
+ return Call::Unchecked(Type::Missing(), op, call_args, Attrs(attrs), ty_args);
}
TVM_FFI_STATIC_INIT_BLOCK() {
@@ -386,7 +422,8 @@ static ffi::Optional<Type> InferCallTIROutputTypeFromArguments(
args->fields.Map([](const Type& ty) -> Expr { return Var("dummy_arg", ty); });
Type derived_ret_ty = DeriveCallRetType(
- dummy_callee_ty, Call(Type::Missing(), Var("dummy_callee", dummy_callee_ty), dummy_args),
+ dummy_callee_ty,
+ Call::Unchecked(Type::Missing(), Var("dummy_callee", dummy_callee_ty), dummy_args),
BlockBuilder::Create(std::nullopt));
if (derived_ret_ty.IsMissing()) {
return std::nullopt;
@@ -488,7 +525,7 @@ Expr NormalizeCallTIR(const BlockBuilder& ctx, Call call) {
return call;
}
-void ValidateCallTIR(Call call) {
+void ValidateCallTIR(const CallNode* call) {
// This function is used for validation of `relax.call_tir`,
// along with the variants `relax.call_tir_with_grad` and
// `relax.call_tir_inplace`. Therefore, all error messages should
@@ -519,12 +556,13 @@ void ValidateCallTIR(Call call) {
TVM_FFI_STATIC_INIT_BLOCK() {
OpDef("relax.call_tir")
- .add_arg("func", "The destination-passing-style function.")
- .add_arg("args", "The input arguments.")
- .add_ty_arg("out_type", "The output type.")
+ .set_validator(ffi::reflection::NativeFunctionView<ffi::Expected<void>(
+ const CallNode*)>::FromNative<&ValidateAsExpected<ValidateCallTIR>>())
+ .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<FNormalize>("FNormalize", NormalizeCallTIR)
- .set_attr<FValidate>("FValidate", ValidateCallTIR)
.set_attr<bool>("FPurity", true);
}
@@ -545,7 +583,7 @@ Expr MakeCallTIR(Expr func, Tuple args, ffi::Array<TensorType> out_ty_list) {
}
static const Op op = Op::Get("relax.call_tir");
- return Call(Type::Missing(), op, {func, args}, {}, {out_ty});
+ return Call::Unchecked(Type::Missing(), op, {func, args}, {}, {out_ty});
}
TVM_FFI_STATIC_INIT_BLOCK() {
@@ -555,13 +593,14 @@ TVM_FFI_STATIC_INIT_BLOCK() {
// call_tir_with_grad
OpDef("relax.call_tir_with_grad")
- .add_arg("func", "The destination-passing-style function.")
- .add_arg("args", "The input arguments.")
- .add_ty_arg("out_type", "The output type.")
- .attrs_type<CallTIRWithGradAttrs>()
+ .set_validator(ffi::reflection::NativeFunctionView<ffi::Expected<void>(
+ const CallNode*)>::FromNative<&ValidateAsExpected<ValidateCallTIR>>())
+ .signature(sig::arg("func", "The destination-passing-style function."),
+ 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<FNormalize>("FNormalize", NormalizeCallTIR)
- .set_attr<FValidate>("FValidate", ValidateCallTIR)
.set_attr<bool>("FPurity", true);
}
@@ -587,7 +626,7 @@ Expr MakeCallTIRWithGrad(Expr func, Tuple args, ffi::Array<TensorType> out_ty_li
attrs->te_grad_kwargs = te_grad_kwargs;
static const Op op = Op::Get("relax.call_tir_with_grad");
- return Call(Type::Missing(), op, {func, args}, Attrs(attrs), {out_ty});
+ return Call::Unchecked(Type::Missing(), op, {func, args}, Attrs(attrs), {out_ty});
}
TVM_FFI_STATIC_INIT_BLOCK() {
@@ -686,13 +725,14 @@ Expr NormalizeCallTIRInPlace(const BlockBuilder& ctx, Call call) {
TVM_FFI_STATIC_INIT_BLOCK() {
OpDef("relax.call_tir_inplace")
- .add_arg("func", "The destination-passing-style function.")
- .add_arg("args", "The input arguments.")
- .add_ty_arg("out_type", "The output type.")
- .attrs_type<CallTIRInplaceAttrs>()
+ .set_validator(ffi::reflection::NativeFunctionView<ffi::Expected<void>(
+ const CallNode*)>::FromNative<&ValidateAsExpected<ValidateCallTIR>>())
+ .signature(sig::arg("func", "The destination-passing-style function."),
+ 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<FNormalize>("FNormalize", NormalizeCallTIRInPlace)
- .set_attr<FValidate>("FValidate", ValidateCallTIR)
// 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)
@@ -720,7 +760,7 @@ Expr MakeCallTIRInplace(Expr func, Tuple args, ffi::Array<int64_t> inplace_indic
}
static const Op op = Op::Get("relax.call_tir_inplace");
- return Call(Type::Missing(), op, {func, args}, Attrs(attrs), {out_ty});
+ return Call::Unchecked(Type::Missing(), op, {func, args}, Attrs(attrs), {out_ty});
}
TVM_FFI_STATIC_INIT_BLOCK() {
@@ -739,9 +779,9 @@ Type InferTypeCallDPSPacked(const Call& call, const BlockBuilder& ctx) {
TVM_FFI_STATIC_INIT_BLOCK() {
OpDef("relax.call_dps_packed")
- .add_arg("func", "The destination-passing-style function.")
- .add_arg("args", "The input arguments.")
- .add_ty_arg("out_type", "The output type.")
+ .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)
// technically, an impure op could be used with this, but there is
// little reason to use DPS with an impure op
@@ -765,7 +805,7 @@ Expr MakeCallDPSPacked(Expr func, Tuple args, ffi::Array<TensorType> out_ty_list
}
static const Op op = Op::Get("relax.call_dps_packed");
- return Call(Type::Missing(), op, {func, args}, {}, {out_ty});
+ return Call::Unchecked(Type::Missing(), op, {func, args}, {}, {out_ty});
}
TVM_FFI_STATIC_INIT_BLOCK() {
@@ -782,7 +822,7 @@ Type InferTypeCallPyFunc(const Call& call, const BlockBuilder& ctx) {
return call->ty_args[0];
}
-void ValidateCallPyFunc(Call call) {
+void ValidateCallPyFunc(const CallNode* call) {
// Validate that the function name is a string literal
auto func_name = call->args[0];
TVM_FFI_ICHECK(func_name->IsInstance<StringImmNode>())
@@ -798,18 +838,19 @@ void ValidateCallPyFunc(Call call) {
TVM_FFI_ICHECK(arg_tuple.as<TupleNode>() || arg_tuple.as<VarNode>())
<< "Operation " << call->op << " must hold its arguments as an in-line tuple. "
- << "However, " << call << " has arguments " << arg_tuple
+ << "However, " << ffi::GetRef<Call>(call) << " has arguments " << arg_tuple
<< ", which is neither an in-line tuple, "
<< "nor a variable binding that may be normalized to an in-line tuple.";
}
TVM_FFI_STATIC_INIT_BLOCK() {
OpDef("relax.call_py_func")
- .add_arg("func_name", "The name of the Python function to call.")
- .add_arg("args", "The input arguments.")
- .add_ty_arg("out_type", "The output type.")
+ .set_validator(ffi::reflection::NativeFunctionView<ffi::Expected<void>(
+ const CallNode*)>::FromNative<&ValidateAsExpected<ValidateCallPyFunc>>())
+ .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<FValidate>("FValidate", ValidateCallPyFunc)
.set_attr<bool>("FPurity", true);
}
@@ -830,7 +871,7 @@ Expr MakeCallPyFunc(StringImm func_name, Tuple args, ffi::Array<TensorType> out_
}
static const Op op = Op::Get("relax.call_py_func");
- return Call(Type::Missing(), op, {func_name, args}, {}, {out_ty});
+ return Call::Unchecked(Type::Missing(), op, {func_name, args}, {}, {out_ty});
}
TVM_FFI_STATIC_INIT_BLOCK() {
@@ -851,9 +892,9 @@ Type InferTypeCallBuiltinWithCtx(const Call& call, const BlockBuilder& ctx) {
TVM_FFI_STATIC_INIT_BLOCK() {
OpDef("relax.call_builtin_with_ctx")
- .add_arg("func", "The builtin packed func.")
- .add_arg("args", "The input arguments.")
- .add_ty_arg("out_type", "Optional output type; omitted for void.")
+ .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)
// Most builtins are pure, but some are not, like `vm.builtin.attention_kv_cache_append`
.set_attr<bool>("FPurity", false);
@@ -861,7 +902,7 @@ TVM_FFI_STATIC_INIT_BLOCK() {
Expr MakeCallBuiltinWithCtx(Expr func, Tuple args, ffi::Array<Type> ty_args) {
static const Op op = Op::Get("relax.call_builtin_with_ctx");
- return Call(Type::Missing(), op, {func, args}, Attrs(), ty_args);
+ return Call::Unchecked(Type::Missing(), op, {func, args}, Attrs(), ty_args);
}
TVM_FFI_STATIC_INIT_BLOCK() {
@@ -875,7 +916,7 @@ TVM_FFI_STATIC_INIT_BLOCK() {
Expr MakeCallNullValue() {
static const Op op = Op::Get("relax.null_value");
- return Call(Type::Missing(), op, {}, {}, {});
+ return Call::Unchecked(Type::Missing(), op, {}, {}, {});
}
TVM_FFI_STATIC_INIT_BLOCK() {
@@ -885,10 +926,11 @@ TVM_FFI_STATIC_INIT_BLOCK() {
// print
OpDef("relax.print")
- .add_arg("format",
- "The first value is Python-style format string to use to print. The others "
- "are values to print")
- .allow_extra_args()
+ .signature(
+ sig::arg("format",
+ "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<FCallPacked>("FCallPacked", "relax.run.print")
.set_attr<bool>("FPurity", false);
@@ -901,7 +943,7 @@ Expr MakePrint(ffi::Array<Expr> vals, StringImm format) {
params.push_back(val);
}
static const Op op = Op::Get("relax.print");
- return Call(Type::Missing(), op, params);
+ return Call::Unchecked(Type::Missing(), op, params);
}
TVM_FFI_STATIC_INIT_BLOCK() {
@@ -931,11 +973,12 @@ Type InferAssertType(const Call& call, const BlockBuilder& ctx) {
TVM_FFI_STATIC_INIT_BLOCK() {
OpDef("relax.assert_op")
- .add_arg("condition",
- "The first value is used as the assertion condition. The second value is "
- "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.")
- .allow_extra_args()
+ .signature(
+ sig::arg("condition",
+ "The first value is used as the assertion condition. The second value is "
+ "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<FCallPacked>("FCallPacked", "relax.run.assert_op")
.set_attr<bool>("FPurity", false);
@@ -948,7 +991,7 @@ Expr MakeAssertOp(Expr condition, ffi::Array<Expr> vals, StringImm format) {
for (auto val : vals) {
args.push_back(val);
}
- return Call(Type::Missing(), op, args);
+ return Call::Unchecked(Type::Missing(), op, args);
}
TVM_FFI_STATIC_INIT_BLOCK() {
@@ -958,15 +1001,14 @@ TVM_FFI_STATIC_INIT_BLOCK() {
// make_closure
OpDef("relax.make_closure")
- .add_arg("func", "The closure.")
- .add_arg("args", "The captured variables.")
+ .signature(sig::arg("func", "The closure."), sig::arg("args", "The captured variables."))
.set_attr<FInferType>("FInferType", ReturnAnyType)
.set_attr<bool>("FPurity", true);
}
Expr MakeClosure(Expr func, Tuple args) {
static const Op op = Op::Get("relax.make_closure");
- return Call(Type::Missing(), op, {func, args}, {}, {});
+ return Call::Unchecked(Type::Missing(), op, {func, args}, {}, {});
}
TVM_FFI_STATIC_INIT_BLOCK() {
@@ -988,9 +1030,9 @@ Type InferTypeInvokeClosure(const Call& call, const BlockBuilder& ctx) {
TVM_FFI_STATIC_INIT_BLOCK() {
OpDef("relax.invoke_closure")
- .add_arg("closure", "The VMClosure.")
- .add_arg("args", "The captured variables.")
- .add_ty_arg("out_types", "Zero or more output types; multiple entries form a tuple.")
+ .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)
// Not all closures are pure. Use invoke_pure_closure for specifying purity
.set_attr<bool>("FPurity", false);
@@ -998,7 +1040,7 @@ TVM_FFI_STATIC_INIT_BLOCK() {
Expr InvokeClosure(Expr closure, Tuple args, ffi::Array<Type> ty_args) {
static const Op op = Op::Get("relax.invoke_closure");
- return Call(Type::Missing(), op, {closure, args}, {}, ty_args);
+ return Call::Unchecked(Type::Missing(), op, {closure, args}, {}, ty_args);
}
TVM_FFI_STATIC_INIT_BLOCK() {
@@ -1008,16 +1050,16 @@ TVM_FFI_STATIC_INIT_BLOCK() {
// invoke_pure_closure
OpDef("relax.invoke_pure_closure")
- .add_arg("closure", "The VMClosure.")
- .add_arg("args", "The captured variables.")
- .add_ty_arg("out_types", "Zero or more output types; multiple entries form a tuple.")
+ .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<bool>("FPurity", true);
}
Expr InvokePureClosure(Expr closure, Tuple args, ffi::Array<Type> ty_args) {
static const Op op = Op::Get("relax.invoke_pure_closure");
- return Call(Type::Missing(), op, {closure, args}, {}, ty_args);
+ return Call::Unchecked(Type::Missing(), op, {closure, args}, {}, ty_args);
}
TVM_FFI_STATIC_INIT_BLOCK() {
@@ -1027,14 +1069,14 @@ TVM_FFI_STATIC_INIT_BLOCK() {
// shape_of
OpDef("relax.shape_of")
- .add_arg("input", "The input expression")
+ .signature(sig::arg("input", "The input expression"))
.set_attr<FInferType>("FInferType", InferTypeShapeOf)
.set_attr<bool>("FPurity", true);
}
Expr MakeShapeOf(Expr expr) {
static const Op op = Op::Get("relax.shape_of");
- return Call(Type::Missing(), op, {expr}, {}, {});
+ return Call::Unchecked(Type::Missing(), op, {expr}, {}, {});
}
TVM_FFI_STATIC_INIT_BLOCK() {
@@ -1054,14 +1096,14 @@ Type InferTypeSize(const Call& call, const BlockBuilder& ctx) {
TVM_FFI_STATIC_INIT_BLOCK() {
OpDef("relax.size")
- .add_arg("input", "The input tensor")
+ .signature(sig::arg("input", "The input tensor"))
.set_attr<FInferType>("FInferType", InferTypeSize)
.set_attr<bool>("FPurity", true);
}
Expr MakeSize(Expr expr) {
static const Op op = Op::Get("relax.size");
- return Call(Type::Missing(), op, {expr}, {}, {});
+ return Call::Unchecked(Type::Missing(), op, {expr}, {}, {});
}
TVM_FFI_STATIC_INIT_BLOCK() {
@@ -1092,14 +1134,14 @@ Type ReturnTensorToShapeType(const Call& call, const BlockBuilder& ctx) {
TVM_FFI_STATIC_INIT_BLOCK() {
OpDef("relax.tensor_to_shape")
- .add_arg("input", "The input expression")
+ .signature(sig::arg("input", "The input expression"))
.set_attr<FInferType>("FInferType", ReturnTensorToShapeType)
.set_attr<bool>("FPurity", true);
}
Expr MakeTensorToShape(Expr expr) {
static const Op op = Op::Get("relax.tensor_to_shape");
- return Call(Type::Missing(), op, {expr}, {}, {});
+ return Call::Unchecked(Type::Missing(), op, {expr}, {}, {});
}
TVM_FFI_STATIC_INIT_BLOCK() {
@@ -1119,7 +1161,7 @@ Type ReturnShapeToTensorType(const Call& call, const BlockBuilder& ctx) {
TVM_FFI_STATIC_INIT_BLOCK() {
OpDef("relax.shape_to_tensor")
- .add_arg("input", "The input expression")
+ .signature(sig::arg("input", "The input expression"))
.set_attr<FInferType>("FInferType", ReturnShapeToTensorType)
.set_attr<FCallPacked>("FCallPacked", "relax.run.shape_to_tensor")
.set_attr<bool>("FPurity", true);
@@ -1127,7 +1169,7 @@ TVM_FFI_STATIC_INIT_BLOCK() {
Expr MakeShapeToTensor(Expr expr) {
static const Op op = Op::Get("relax.shape_to_tensor");
- return Call(Type::Missing(), op, {expr}, {}, {});
+ return Call::Unchecked(Type::Missing(), op, {expr}, {}, {});
}
TVM_FFI_STATIC_INIT_BLOCK() {
@@ -1161,12 +1203,14 @@ Type InferTypeAllocateTensor(const Call& call, const BlockBuilder& ctx) {
TVM_FFI_STATIC_INIT_BLOCK() {
OpDef("relax.builtin.alloc_tensor")
- .add_arg("shape", "The shape of the tensor to allocate.")
- .add_arg("dtype", "The dtype of the tensor to allocate.")
- .add_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.")
- .add_arg("storage_scope", "The storage scope of the storage to allocate. Default is global.")
+ .signature(sig::arg("shape", "The shape of the tensor to allocate."),
+ sig::arg("dtype", "The dtype of the tensor to allocate."),
+ 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."),
+ 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)
// memory allocation isn't considered a "visible effect" as far as purity is concerned
.set_attr<bool>("FPurity", true)
@@ -1176,8 +1220,8 @@ TVM_FFI_STATIC_INIT_BLOCK() {
Expr MakeAllocTensor(Expr shape, DataTypeImm dtype, PrimExpr runtime_device_index,
StringImm storage_scope) {
static const Op op = Op::Get("relax.builtin.alloc_tensor");
- return Call(Type::Missing(), op, {shape, dtype, runtime_device_index, storage_scope}, Attrs(),
- {});
+ return Call::Unchecked(Type::Missing(), op, {shape, dtype, runtime_device_index, storage_scope},
+ Attrs(), {});
}
TVM_FFI_STATIC_INIT_BLOCK() {
@@ -1187,13 +1231,15 @@ TVM_FFI_STATIC_INIT_BLOCK() {
// memory planning alloc_storage
OpDef("relax.memory.alloc_storage")
- .add_arg("total_space", "The total space of the storage to allocate.")
- .add_arg(
- "virtual_device_index",
- "The virtual device index indicating on which device the storage is to be allocated, "
- "Index -1 is reserved for the host device.")
- .add_arg("storage_scope", "The storage scope of the storage to allocate. Default is global.")
- .add_arg("dtype", "The dtype of the tensor to allocate.")
+ .signature(
+ sig::arg("total_space", "The total space of the storage to allocate."),
+ sig::arg(
+ "virtual_device_index",
+ "The virtual device index indicating on which device the storage is to be allocated, "
+ "Index -1 is reserved for the host device."),
+ 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)
// memory allocation isn't considered a "visible effect" as far as purity is concerned
.set_attr<bool>("FPurity", true)
@@ -1203,7 +1249,8 @@ TVM_FFI_STATIC_INIT_BLOCK() {
Expr MakeAllocStorage(Expr size, PrimExpr virtual_device_index, StringImm storage_scope,
DataTypeImm dtype) {
static const Op op = Op::Get("relax.memory.alloc_storage");
- return Call(Type::Missing(), op, {size, virtual_device_index, storage_scope, dtype}, Attrs(), {});
+ return Call::Unchecked(Type::Missing(), op, {size, virtual_device_index, storage_scope, dtype},
+ Attrs(), {});
}
TVM_FFI_STATIC_INIT_BLOCK() {
@@ -1238,13 +1285,13 @@ Type InferTypeMemAllocTensor(const Call& call, const BlockBuilder& ctx) {
TVM_FFI_STATIC_INIT_BLOCK() {
OpDef("relax.memory.alloc_tensor")
- .add_arg("storage", "The storage to allocate the tensor to.")
- .add_arg("offset", "Storage offset to allocate the tensor.")
- .add_arg("shape", "The shape of the tensor to allocate.")
- .add_arg("dtype", "The dtype of the tensor to allocate.")
- .add_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.")
+ .signature(sig::arg("storage", "The storage to allocate the tensor to."),
+ sig::arg("offset", "Storage offset to allocate the tensor."),
+ sig::arg("shape", "The shape of the tensor to allocate."),
+ sig::arg("dtype", "The dtype of the tensor to allocate."),
+ 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)
// memory allocation isn't considered a "visible effect" as far as purity is concerned
.set_attr<bool>("FPurity", true)
@@ -1254,8 +1301,8 @@ TVM_FFI_STATIC_INIT_BLOCK() {
Expr MakeMemAllocTensor(Expr storage, PrimExpr offset, Expr shape, DataTypeImm dtype,
PrimExpr virtual_device_index) {
static const Op op = Op::Get("relax.memory.alloc_tensor");
- return Call(Type::Missing(), op, {storage, offset, shape, dtype, virtual_device_index}, Attrs(),
- {});
+ return Call::Unchecked(Type::Missing(), op, {storage, offset, shape, dtype, virtual_device_index},
+ Attrs(), {});
}
TVM_FFI_STATIC_INIT_BLOCK() {
@@ -1276,7 +1323,7 @@ TVM_FFI_STATIC_INIT_BLOCK() {
// memory planning kill_storage
OpDef("relax.memory.kill_storage")
- .add_arg("storage", "The storage to be killed.")
+ .signature(sig::arg("storage", "The storage to be killed."))
.set_attr<FInferType>("FInferType", ReturnVoidType)
// We mark this as impure so it wouldn't be removed by "remove_all_unused"
.set_attr<bool>("FPurity", false);
@@ -1284,7 +1331,7 @@ TVM_FFI_STATIC_INIT_BLOCK() {
Expr MakeMemKillStorage(Expr storage) {
static const Op op = Op::Get("relax.memory.kill_storage");
- return Call(Type::Missing(), op, {storage}, {}, {});
+ return Call::Unchecked(Type::Missing(), op, {storage}, {}, {});
}
TVM_FFI_STATIC_INIT_BLOCK() {
@@ -1294,7 +1341,7 @@ TVM_FFI_STATIC_INIT_BLOCK() {
// memory planning kill_tensor
OpDef("relax.memory.kill_tensor")
- .add_arg("tensor", "The tensor to be killed.")
+ .signature(sig::arg("tensor", "The tensor to be killed."))
.set_attr<FInferType>("FInferType", ReturnVoidType)
// We mark this as impure so it wouldn't be removed by "remove_all_unused"
.set_attr<bool>("FPurity", false);
@@ -1302,7 +1349,7 @@ TVM_FFI_STATIC_INIT_BLOCK() {
Expr MakeMemKillTensor(Expr tensor) {
static const Op op = Op::Get("relax.memory.kill_tensor");
- return Call(Type::Missing(), op, {tensor}, {}, {});
+ return Call::Unchecked(Type::Missing(), op, {tensor}, {}, {});
}
TVM_FFI_STATIC_INIT_BLOCK() {
@@ -1312,12 +1359,13 @@ TVM_FFI_STATIC_INIT_BLOCK() {
// vm alloc_storage
OpDef("relax.vm.alloc_storage")
- .add_arg("size", "The size of the storage to allocate.")
- .add_arg("runtime_device_index",
- "The device index indicating on which device the tensor is "
- "to be allocated at runtime.")
- .add_arg("dtype", "The dtype of the tensor to allocate.")
- .add_arg("storage_scope", "The storage scope of the storage to allocate. Default is global.")
+ .signature(sig::arg("size", "The size of the storage to allocate."),
+ sig::arg("runtime_device_index",
+ "The device index indicating on which device the tensor is "
+ "to be allocated at runtime."),
+ 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)
// memory allocation isn't considered a "visible effect" as far as purity is concerned
.set_attr<bool>("FPurity", true)
@@ -1327,7 +1375,8 @@ TVM_FFI_STATIC_INIT_BLOCK() {
Expr MakeVMAllocStorage(Expr size, PrimExpr runtime_device_index, DataTypeImm dtype,
StringImm storage_scope) {
static const Op op = Op::Get("relax.vm.alloc_storage");
- return Call(Type::Missing(), op, {size, runtime_device_index, dtype, storage_scope}, Attrs(), {});
+ return Call::Unchecked(Type::Missing(), op, {size, runtime_device_index, dtype, storage_scope},
+ Attrs(), {});
}
TVM_FFI_STATIC_INIT_BLOCK() {
@@ -1363,13 +1412,13 @@ Type InferTypeVMAllocTensor(const Call& call, const BlockBuilder& ctx) {
TVM_FFI_STATIC_INIT_BLOCK() {
OpDef("relax.vm.alloc_tensor")
- .add_arg("storage", "The storage to allocate the tensor to.")
- .add_arg("offset", "Storage offset to allocate the tensor.")
- .add_arg("shape", "The shape of the tensor to allocate.")
- .add_arg("dtype", "The dtype of the tensor to allocate.")
- .add_arg("runtime_device_index",
- "The device index indicating on which device the tensor is "
- "to be allocated at runtime.")
+ .signature(sig::arg("storage", "The storage to allocate the tensor to."),
+ sig::arg("offset", "Storage offset to allocate the tensor."),
+ sig::arg("shape", "The shape of the tensor to allocate."),
+ sig::arg("dtype", "The dtype of the tensor to allocate."),
+ 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)
// memory allocation isn't considered a "visible effect" as far as purity is concerned
.set_attr<bool>("FPurity", true)
@@ -1379,8 +1428,8 @@ TVM_FFI_STATIC_INIT_BLOCK() {
Expr MakeVMAllocTensor(Expr storage, PrimExpr offset, Expr shape, DataTypeImm dtype,
PrimExpr runtime_device_index) {
static const Op op = Op::Get("relax.vm.alloc_tensor");
- return Call(Type::Missing(), op, {storage, offset, shape, dtype, runtime_device_index}, Attrs(),
- {});
+ return Call::Unchecked(Type::Missing(), op, {storage, offset, shape, dtype, runtime_device_index},
+ Attrs(), {});
}
TVM_FFI_STATIC_INIT_BLOCK() {
@@ -1398,7 +1447,7 @@ TVM_FFI_STATIC_INIT_BLOCK() {
// vm kill_object
OpDef("relax.vm.kill_object")
- .add_arg("obj", "The object to be killed.")
+ .signature(sig::arg("obj", "The object to be killed."))
.set_attr<FInferType>("FInferType", ReturnVoidType)
// We mark this as impure so it wouldn't be removed by "remove_all_unused"
.set_attr<bool>("FPurity", false);
@@ -1406,7 +1455,7 @@ TVM_FFI_STATIC_INIT_BLOCK() {
Expr MakeVMKillObject(Expr obj) {
static const Op op = Op::Get("relax.vm.kill_object");
- return Call(Type::Missing(), op, {std::move(obj)}, Attrs(), {});
+ return Call::Unchecked(Type::Missing(), op, {std::move(obj)}, Attrs(), {});
}
TVM_FFI_STATIC_INIT_BLOCK() {
@@ -1416,8 +1465,9 @@ TVM_FFI_STATIC_INIT_BLOCK() {
// vm call_tir_dyn
OpDef("relax.vm.call_tir_dyn")
- .add_arg("func", "The destination-passing-style function.")
- .add_arg("args", "The input arguments (list of tensors and last argument is ShapeExpr)")
+ .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)
// "relax.vm.call_tir_dyn" works in an in-place way, which is impure.
.set_attr<bool>("FPurity", false);
@@ -1425,7 +1475,7 @@ TVM_FFI_STATIC_INIT_BLOCK() {
Expr MakeCallTIRDyn(Expr func, Tuple args) {
static const Op op = Op::Get("relax.vm.call_tir_dyn");
- return Call(Type::Missing(), op, {func, args}, Attrs(), {});
+ return Call::Unchecked(Type::Missing(), op, {func, args}, Attrs(), {});
}
TVM_FFI_STATIC_INIT_BLOCK() {
@@ -1440,14 +1490,15 @@ Type InferTypeStopLiftParams(const Call& call, const BlockBuilder& ctx) {
TVM_FFI_STATIC_INIT_BLOCK() {
OpDef("relax.builtin.stop_lift_params")
- .add_arg("x", "The input data")
+ .signature(sig::arg("x", "The input data"),
+ sig::var_ty_args("out_type", "Optional output type."))
.set_attr<FInferType>("FInferType", InferTypeStopLiftParams)
.set_attr<bool>("FPurity", true);
}
Expr MakeStopLiftParams(Expr x) {
static const Op op = Op::Get("relax.builtin.stop_lift_params");
- return Call(Type::Missing(), op, {x}, Attrs(), {});
+ return Call::Unchecked(Type::Missing(), op, {x}, Attrs(), {});
}
TVM_FFI_STATIC_INIT_BLOCK() {
@@ -1471,8 +1522,8 @@ Type InferToVDeviceType(const Call& call, const BlockBuilder& ctx) {
TVM_FFI_STATIC_INIT_BLOCK() {
OpDef("relax.to_vdevice")
- .add_arg("data", "The input expression to be copied")
- .attrs_type<ToVDeviceAttrs>()
+ .signature(sig::arg("data", "The input expression to be copied"),
+ sig::call_attrs<ToVDeviceAttrs>())
.set_attr<FInferType>("FInferType", InferToVDeviceType)
.set_attr<bool>("FPurity", true);
}
@@ -1481,7 +1532,7 @@ Expr MakeToVDevice(Expr data, VDevice dst_vdev) {
static const Op op = Op::Get("relax.to_vdevice");
ffi::ObjectPtr<ToVDeviceAttrs> attrs = ffi::make_object<ToVDeviceAttrs>();
attrs->dst_vdevice = dst_vdev;
- return Call(Type::Missing(), op, {data}, Attrs(attrs), {});
+ return Call::Unchecked(Type::Missing(), op, {data}, Attrs(attrs), {});
}
TVM_FFI_STATIC_INIT_BLOCK() {
@@ -1500,8 +1551,7 @@ Type InferHintOnDeviceType(const Call& call, const BlockBuilder& ctx) {
TVM_FFI_STATIC_INIT_BLOCK() {
OpDef("relax.hint_on_device")
- .add_arg("data", "The input expression")
- .attrs_type<HintOnDeviceAttrs>()
+ .signature(sig::arg("data", "The input expression"), sig::call_attrs<HintOnDeviceAttrs>())
.set_attr<FInferType>("FInferType", InferHintOnDeviceType)
.set_attr<bool>("FPurity", true);
}
@@ -1512,7 +1562,7 @@ Expr MakeHintOnDevice(Expr data, Device device, ffi::String memory_scope = "glob
attrs->device_type = static_cast<int32_t>(device.device_type);
attrs->index = device.device_id;
attrs->memory_scope = memory_scope;
- return Call(Type::Missing(), op, {data}, Attrs(attrs), {});
+ return Call::Unchecked(Type::Missing(), op, {data}, Attrs(attrs), {});
}
TVM_FFI_STATIC_INIT_BLOCK() {
diff --git a/src/relax/op/op_common.cc b/src/relax/op/op_common.cc
index 67e23768f7..fd2d39b953 100644
--- a/src/relax/op/op_common.cc
+++ b/src/relax/op/op_common.cc
@@ -43,10 +43,10 @@ ffi::Array<Expr> GetCallArgs(const Call& call) {
void CheckNumArguments(const Call& call, const BlockBuilder& ctx) {
Op op = call->op.as_or_throw<Op>();
int expected_input = op->args_info.size();
- if (op->allow_extra_args ? call->args.size() < op->args_info.size()
- : call->args.size() != op->args_info.size()) {
+ if (op->var_args_info.has_value() ? call->args.size() < op->args_info.size()
+ : call->args.size() != op->args_info.size()) {
TVM_FFI_VISIT_THROW(ValueError, call)
- << "Operator " << op << " expects " << (op->allow_extra_args ? "at least " : "")
+ << "Operator " << op << " expects " << (op->var_args_info.has_value() ? "at least " : "")
<< expected_input << " arguments"
<< ", but was called with " << call->args.size() << " arguments";
}
@@ -55,8 +55,8 @@ void CheckNumArguments(const Call& call, const BlockBuilder& ctx) {
TensorType GetInputTensorType(const Call& call, size_t i_arg, const BlockBuilder& ctx) {
Op op = call->op.as_or_throw<Op>();
- TVM_FFI_ICHECK(op->allow_extra_args ? call->args.size() >= op->args_info.size()
- : call->args.size() == op->args_info.size())
+ TVM_FFI_ICHECK(op->var_args_info.has_value() ? call->args.size() >= op->args_info.size()
+ : call->args.size() == op->args_info.size())
<< "Failure caught by this check "
<< "should have previously been caught by `CheckNumArguments`";
TVM_FFI_ICHECK_LT(i_arg, call->args.size());
diff --git a/src/relax/op/op_common.h b/src/relax/op/op_common.h
index b4a9ac2817..ccfdeb1d09 100644
--- a/src/relax/op/op_common.h
+++ b/src/relax/op/op_common.h
@@ -144,11 +144,11 @@ std::tuple<ArgTypes...> GetArgType(const Call& call, const BlockBuilder& ctx) {
Op op = call->op.as_or_throw<Op>();
size_t n_input = op->args_info.size();
- // Unfortunately, because the `.add_arg()` calls in
- // OpDef occur during initialization of globals and are
+ // Unfortunately, because signature registrations in OpDef
+ // occur during initialization of globals and are
// not available at compile-time, this cannot be a static_assert.
- TVM_FFI_ICHECK(op->allow_extra_args ? sizeof...(ArgTypes) >= n_input
- : sizeof...(ArgTypes) == n_input)
+ TVM_FFI_ICHECK(op->var_args_info.has_value() ? sizeof...(ArgTypes) >= n_input
+ : sizeof...(ArgTypes) == n_input)
<< "Internal error: " << op << " op defines " << n_input << " arguments in its OpDef() call, "
<< "but GetArgType was given " << sizeof...(ArgTypes) << " template arguments.";
diff --git a/src/relax/op/tensor/binary.h b/src/relax/op/tensor/binary.h
index 46ec595b26..1700789465 100644
--- a/src/relax/op/tensor/binary.h
+++ b/src/relax/op/tensor/binary.h
@@ -30,21 +30,22 @@ namespace tvm {
namespace relax {
/* Register a binary constructor and its Op metadata together. */
-#define RELAX_REGISTER_BINARY_OP_AND_IMPL(OpName, ...) \
- Expr OpName(Expr x1, Expr x2) { \
- static const Op op = Op::Get("relax." #OpName); \
- return Call(Type::Missing(), op, {x1, x2}, Attrs(), {}); \
- } \
- TVM_FFI_STATIC_INIT_BLOCK() { \
- tvm::ffi::reflection::GlobalDef().def("relax.op." #OpName, OpName); \
- OpDef("relax." #OpName) \
- .add_arg("x1", "The first input tensor.") \
- .add_arg("x2", "The second input tensor.") \
- .add_ty_arg("out_type", "Optional output tensor type carrying the virtual device.") \
- .set_attr<FRelaxInferLayout>("FRelaxInferLayout", InferLayoutBinaryEwise) \
- .set_attr<TMixedPrecisionPolicy>("TMixedPrecisionPolicy", \
- MixedPrecisionPolicyKind::kFollow) \
- .set_attr<bool>("FPurity", true) __VA_ARGS__; \
+#define RELAX_REGISTER_BINARY_OP_AND_IMPL(OpName, ...) \
+ Expr OpName(Expr x1, Expr x2) { \
+ static const Op op = Op::Get("relax." #OpName); \
+ return Call::Unchecked(Type::Missing(), op, {x1, x2}, Attrs(), {}); \
+ } \
+ TVM_FFI_STATIC_INIT_BLOCK() { \
+ tvm::ffi::reflection::GlobalDef().def("relax.op." #OpName, OpName); \
+ OpDef("relax." #OpName) \
+ .signature(sig::arg("x1", "The first input tensor."), \
+ sig::arg("x2", "The second input tensor."), \
+ sig::var_ty_args("out_type", \
+ "Optional output tensor type carrying the virtual device.")) \
+ .set_attr<FRelaxInferLayout>("FRelaxInferLayout", InferLayoutBinaryEwise) \
+ .set_attr<TMixedPrecisionPolicy>("TMixedPrecisionPolicy", \
+ MixedPrecisionPolicyKind::kFollow) \
+ .set_attr<bool>("FPurity", true) __VA_ARGS__; \
}
#define RELAX_REGISTER_BINARY_BROADCAST_OP_AND_IMPL(OpName) \
diff --git a/src/relax/op/tensor/create.cc b/src/relax/op/tensor/create.cc
index 72504a2d0d..f7dfe9e20a 100644
--- a/src/relax/op/tensor/create.cc
+++ b/src/relax/op/tensor/create.cc
@@ -62,8 +62,8 @@ Expr full(ffi::Variant<Expr, ffi::Array<PrimExpr>> shape, Expr fill_value,
attrs->dtype = dtype;
static const Op op = Op::Get("relax.full");
- return Call(Type::Missing(), op, {std::move(shape_in_expr), std::move(fill_value)}, Attrs(attrs),
- {});
+ return Call::Unchecked(Type::Missing(), op, {std::move(shape_in_expr), std::move(fill_value)},
+ Attrs(attrs), {});
}
TVM_FFI_STATIC_INIT_BLOCK() {
@@ -97,9 +97,9 @@ Type InferTypeFull(const Call& call, const BlockBuilder& ctx) {
TVM_FFI_STATIC_INIT_BLOCK() {
OpDef("relax.full")
- .add_arg("shape", "The shape of the created tensor.")
- .add_arg("fill_value", "The scalar tensor, denoting the value to fill.")
- .attrs_type<InitAttrs>()
+ .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<bool>("RequiresArgumentShapes", false)
.set_attr<bool>("FDataDependent", true)
@@ -112,7 +112,8 @@ Expr full_like(Expr x, Expr fill_value, ffi::Optional<DLDataType> dtype) {
ffi::ObjectPtr<InitAttrs> attrs = ffi::make_object<InitAttrs>();
attrs->dtype = dtype;
static const Op op = Op::Get("relax.full_like");
- return Call(Type::Missing(), op, {std::move(x), std::move(fill_value)}, Attrs(attrs), {});
+ return Call::Unchecked(Type::Missing(), op, {std::move(x), std::move(fill_value)}, Attrs(attrs),
+ {});
}
TVM_FFI_STATIC_INIT_BLOCK() {
@@ -142,9 +143,8 @@ Type InferTypeFullLike(const Call& call, const BlockBuilder& ctx) {
TVM_FFI_STATIC_INIT_BLOCK() {
OpDef("relax.full_like")
- .add_arg("x", "The input tensor.")
- .add_arg("fill_value", "The scalar value to fill.")
- .attrs_type<InitAttrs>()
+ .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<TMixedPrecisionPolicy>("TMixedPrecisionPolicy", MixedPrecisionPolicyKind::kFollow)
.set_attr<bool>("FPurity", true);
@@ -186,14 +186,14 @@ Expr ones(Expr shape, DLDataType dtype) {
attrs->dtype = dtype;
static const Op op = Op::Get("relax.ones");
- return Call(Type::Missing(), op, {std::move(shape)}, Attrs(attrs), {});
+ return Call::Unchecked(Type::Missing(), op, {std::move(shape)}, Attrs(attrs), {});
}
Expr ones_like(Expr x, ffi::Optional<DLDataType> dtype) {
ffi::ObjectPtr<InitAttrs> attrs = ffi::make_object<InitAttrs>();
attrs->dtype = dtype;
static const Op op = Op::Get("relax.ones_like");
- return Call(Type::Missing(), op, {std::move(x)}, Attrs(attrs), {});
+ return Call::Unchecked(Type::Missing(), op, {std::move(x)}, Attrs(attrs), {});
}
TVM_FFI_STATIC_INIT_BLOCK() {
@@ -201,15 +201,14 @@ TVM_FFI_STATIC_INIT_BLOCK() {
refl::GlobalDef().def("relax.op.ones", ones).def("relax.op.ones_like", ones_like);
OpDef("relax.ones")
- .add_arg("shape", "The shape of the created tensor.")
- .attrs_type<InitAttrs>()
+ .signature(sig::arg("shape", "The shape of the created tensor."),
+ sig::call_attrs<InitAttrs>())
.set_attr<FInferType>("FInferType", InferTypeOnesZeros)
.set_attr<TMixedPrecisionPolicy>("TMixedPrecisionPolicy", MixedPrecisionPolicyKind::kFollow)
.set_attr<bool>("FPurity", true);
OpDef("relax.ones_like")
- .add_arg("x", "The input tensor.")
- .attrs_type<InitAttrs>()
+ .signature(sig::arg("x", "The input tensor."), sig::call_attrs<InitAttrs>())
.set_attr<FInferType>("FInferType", InferTypeOnesLikeZerosLike)
.set_attr<bool>("FPurity", true);
}
@@ -220,14 +219,14 @@ Expr zeros(Expr shape, DLDataType dtype) {
attrs->dtype = dtype;
static const Op op = Op::Get("relax.zeros");
- return Call(Type::Missing(), op, {std::move(shape)}, Attrs(attrs), {});
+ return Call::Unchecked(Type::Missing(), op, {std::move(shape)}, Attrs(attrs), {});
}
Expr zeros_like(Expr x, ffi::Optional<DLDataType> dtype) {
ffi::ObjectPtr<InitAttrs> attrs = ffi::make_object<InitAttrs>();
attrs->dtype = dtype;
static const Op op = Op::Get("relax.zeros_like");
- return Call(Type::Missing(), op, {std::move(x)}, Attrs(attrs), {});
+ return Call::Unchecked(Type::Missing(), op, {std::move(x)}, Attrs(attrs), {});
}
TVM_FFI_STATIC_INIT_BLOCK() {
@@ -235,15 +234,14 @@ TVM_FFI_STATIC_INIT_BLOCK() {
refl::GlobalDef().def("relax.op.zeros", zeros).def("relax.op.zeros_like", zeros_like);
OpDef("relax.zeros")
- .add_arg("shape", "The shape of the created tensor.")
- .attrs_type<InitAttrs>()
+ .signature(sig::arg("shape", "The shape of the created tensor."),
+ sig::call_attrs<InitAttrs>())
.set_attr<FInferType>("FInferType", InferTypeOnesZeros)
.set_attr<TMixedPrecisionPolicy>("TMixedPrecisionPolicy", MixedPrecisionPolicyKind::kFollow)
.set_attr<bool>("FPurity", true);
OpDef("relax.zeros_like")
- .add_arg("x", "The input tensor.")
- .attrs_type<InitAttrs>()
+ .signature(sig::arg("x", "The input tensor."), sig::call_attrs<InitAttrs>())
.set_attr<FInferType>("FInferType", InferTypeOnesLikeZerosLike)
.set_attr<bool>("FPurity", true);
}
@@ -253,14 +251,15 @@ Expr eye(PrimExpr n, PrimExpr m, PrimExpr k, DLDataType dtype) {
ffi::ObjectPtr<InitAttrs> attrs = ffi::make_object<InitAttrs>();
attrs->dtype = dtype;
static const Op op = Op::Get("relax.eye");
- return Call(Type::Missing(), op, {std::move(n), std::move(m), std::move(k)}, Attrs(attrs), {});
+ return Call::Unchecked(Type::Missing(), op, {std::move(n), std::move(m), std::move(k)},
+ Attrs(attrs), {});
}
Expr eye_like(Expr x, PrimExpr k, ffi::Optional<DLDataType> dtype) {
ffi::ObjectPtr<InitAttrs> attrs = ffi::make_object<InitAttrs>();
attrs->dtype = dtype;
static const Op op = Op::Get("relax.eye_like");
- return Call(Type::Missing(), op, {std::move(x), std::move(k)}, Attrs(attrs), {});
+ return Call::Unchecked(Type::Missing(), op, {std::move(x), std::move(k)}, Attrs(attrs), {});
}
TVM_FFI_STATIC_INIT_BLOCK() {
@@ -321,18 +320,16 @@ Type InferTypeEyeLike(const Call& call, const BlockBuilder& ctx) {
TVM_FFI_STATIC_INIT_BLOCK() {
OpDef("relax.eye")
- .add_arg("n", "Number of rows in the output.")
- .add_arg("m", "Number of columns in the output.")
- .add_arg("k", "Index of the diagonal.")
- .attrs_type<InitAttrs>()
+ .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<TMixedPrecisionPolicy>("TMixedPrecisionPolicy", MixedPrecisionPolicyKind::kFollow)
.set_attr<bool>("FPurity", true);
OpDef("relax.eye_like")
- .add_arg("x", "The input tensor.")
- .add_arg("k", "Index of the diagonal.")
- .attrs_type<InitAttrs>()
+ .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<bool>("FPurity", true);
}
@@ -342,8 +339,8 @@ Expr arange(PrimExpr start, PrimExpr stop, PrimExpr step, DLDataType dtype) {
ffi::ObjectPtr<InitAttrs> attrs = ffi::make_object<InitAttrs>();
attrs->dtype = dtype;
static const Op op = Op::Get("relax.arange");
- return Call(Type::Missing(), op, {std::move(start), std::move(stop), std::move(step)},
- Attrs(attrs), {});
+ return Call::Unchecked(Type::Missing(), op, {std::move(start), std::move(stop), std::move(step)},
+ Attrs(attrs), {});
}
TVM_FFI_STATIC_INIT_BLOCK() {
@@ -388,10 +385,10 @@ Type InferTypeArange(const Call& call, const BlockBuilder& ctx) {
TVM_FFI_STATIC_INIT_BLOCK() {
OpDef("relax.arange")
- .add_arg("start", "The starting value for the set of points.")
- .add_arg("end", "The ending value for the set of points.")
- .add_arg("step", "The gap between each pair of adjacent points.")
- .attrs_type<InitAttrs>()
+ .signature(sig::arg("start", "The starting value for the set of points."),
+ 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<TMixedPrecisionPolicy>("TMixedPrecisionPolicy", MixedPrecisionPolicyKind::kFollow)
.set_attr<bool>("FPurity", true);
@@ -403,9 +400,10 @@ Expr hamming_window(PrimExpr window_size, PrimExpr periodic, PrimExpr alpha, Pri
ffi::ObjectPtr<InitAttrs> attrs = ffi::make_object<InitAttrs>();
attrs->dtype = dtype;
static const Op op = Op::Get("relax.hamming_window");
- return Call(Type::Missing(), op,
- {std::move(window_size), std::move(periodic), std::move(alpha), std::move(beta)},
- Attrs(attrs), {});
+ return Call::Unchecked(
+ Type::Missing(), op,
+ {std::move(window_size), std::move(periodic), std::move(alpha), std::move(beta)},
+ Attrs(attrs), {});
}
TVM_FFI_STATIC_INIT_BLOCK() {
@@ -443,13 +441,13 @@ Type InferTypeHammingWindow(const Call& call, const BlockBuilder& ctx) {
TVM_FFI_STATIC_INIT_BLOCK() {
OpDef("relax.hamming_window")
- .add_arg("window_size", "The size of the window")
- .add_arg("periodic",
- "If True, returns a window to be used as periodic function. If False, return a "
- "symmetric window")
- .add_arg("alpha", "The coefficient alpha")
- .add_arg("beta", "The coefficient beta")
- .attrs_type<InitAttrs>()
+ .signature(
+ sig::arg("window_size", "The size of the window"),
+ sig::arg("periodic",
+ "If True, returns a window to be used as periodic function. If False, return a "
+ "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<TMixedPrecisionPolicy>("TMixedPrecisionPolicy", MixedPrecisionPolicyKind::kFollow)
.set_attr<bool>("FPurity", true);
@@ -459,12 +457,12 @@ TVM_FFI_STATIC_INIT_BLOCK() {
Expr tril(Expr x, Expr k) {
static const Op op = Op::Get("relax.tril");
- return Call(Type::Missing(), op, {x, k});
+ return Call::Unchecked(Type::Missing(), op, {x, k});
}
Expr triu(Expr x, Expr k) {
static const Op op = Op::Get("relax.triu");
- return Call(Type::Missing(), op, {x, k});
+ return Call::Unchecked(Type::Missing(), op, {x, k});
}
TVM_FFI_STATIC_INIT_BLOCK() {
@@ -488,14 +486,12 @@ Type InferTypeTrilTriu(const Call& call, const BlockBuilder& ctx) {
TVM_FFI_STATIC_INIT_BLOCK() {
OpDef("relax.tril")
- .add_arg("x", "The input tensor.")
- .add_arg("k", "The offset of the diagonal.")
+ .signature(sig::arg("x", "The input tensor."), sig::arg("k", "The offset of the diagonal."))
.set_attr<FInferType>("FInferType", InferTypeTrilTriu)
.set_attr<bool>("FPurity", true);
OpDef("relax.triu")
- .add_arg("x", "The input tensor.")
- .add_arg("k", "The offset of the diagonal.")
+ .signature(sig::arg("x", "The input tensor."), sig::arg("k", "The offset of the diagonal."))
.set_attr<FInferType>("FInferType", InferTypeTrilTriu)
.set_attr<bool>("FPurity", true);
}
diff --git a/src/relax/op/tensor/datatype.cc b/src/relax/op/tensor/datatype.cc
index 87efbb5d91..58b78e27bc 100644
--- a/src/relax/op/tensor/datatype.cc
+++ b/src/relax/op/tensor/datatype.cc
@@ -43,7 +43,7 @@ Expr astype(Expr x, DLDataType dtype) {
attrs->dtype = dtype;
static const Op op = Op::Get("relax.astype");
- return Call(Type::Missing(), op, {std::move(x)}, Attrs(attrs), {});
+ return Call::Unchecked(Type::Missing(), op, {std::move(x)}, Attrs(attrs), {});
}
TVM_FFI_STATIC_INIT_BLOCK() {
@@ -61,8 +61,7 @@ Type InferTypeAstype(const Call& call, const BlockBuilder& ctx) {
TVM_FFI_STATIC_INIT_BLOCK() {
OpDef("relax.astype")
- .add_arg("x", "The input tensor")
- .attrs_type<AstypeAttrs>()
+ .signature(sig::arg("x", "The input tensor"), sig::call_attrs<AstypeAttrs>())
.set_attr<FInferType>("FInferType", InferTypeAstype)
.set_attr<FRelaxInferLayout>("FRelaxInferLayout", InferLayoutUnaryEwise)
.set_attr<TMixedPrecisionPolicy>("TMixedPrecisionPolicy", MixedPrecisionPolicyKind::kFollow)
@@ -76,7 +75,7 @@ Expr MakeWrapParam(Expr data, DLDataType dtype) {
attrs->dtype = dtype;
static const Op op = Op::Get("relax.wrap_param");
- return Call(Type::Missing(), op, {std::move(data)}, Attrs(attrs), {});
+ return Call::Unchecked(Type::Missing(), op, {std::move(data)}, Attrs(attrs), {});
}
TVM_FFI_STATIC_INIT_BLOCK() {
@@ -94,8 +93,7 @@ Type InferTypeWrapParam(const Call& call, const BlockBuilder& ctx) {
TVM_FFI_STATIC_INIT_BLOCK() {
OpDef("relax.wrap_param")
- .add_arg("data", "The input tensor")
- .attrs_type<WrapParamAttrs>()
+ .signature(sig::arg("data", "The input tensor"), sig::call_attrs<WrapParamAttrs>())
.set_attr<FInferType>("FInferType", InferTypeWrapParam)
.set_attr<bool>("FPurity", true);
}
diff --git a/src/relax/op/tensor/grad.cc b/src/relax/op/tensor/grad.cc
index 8b3e222a47..a113512e6c 100644
--- a/src/relax/op/tensor/grad.cc
+++ b/src/relax/op/tensor/grad.cc
@@ -35,7 +35,7 @@ namespace relax {
/* relax.grad.no_grad */
Expr no_grad(Expr input) {
static const Op op = Op::Get("relax.grad.no_grad");
- return Call(Type::Missing(), op, {std::move(input)}, {}, {});
+ return Call::Unchecked(Type::Missing(), op, {std::move(input)}, {}, {});
}
TVM_FFI_STATIC_INIT_BLOCK() {
@@ -47,7 +47,7 @@ Type InferTypeNoGrad(const Call& call, const BlockBuilder& ctx) { return GetType
TVM_FFI_STATIC_INIT_BLOCK() {
OpDef("relax.grad.no_grad")
- .add_arg("x", "The corresponding input tensor.")
+ .signature(sig::arg("x", "The corresponding input tensor."))
.set_attr<FInferType>("FInferType", InferTypeNoGrad)
.set_attr<bool>("FPurity", true);
}
@@ -55,7 +55,7 @@ TVM_FFI_STATIC_INIT_BLOCK() {
/* relax.grad.start_checkpoint */
Expr start_checkpoint(Expr input) {
static const Op op = Op::Get("relax.grad.start_checkpoint");
- return Call(Type::Missing(), op, {std::move(input)}, {}, {});
+ return Call::Unchecked(Type::Missing(), op, {std::move(input)}, {}, {});
}
TVM_FFI_STATIC_INIT_BLOCK() {
@@ -73,7 +73,7 @@ Type InferTypeStartCheckpoint(const Call& call, const BlockBuilder& ctx) {
TVM_FFI_STATIC_INIT_BLOCK() {
OpDef("relax.grad.start_checkpoint")
- .add_arg("x", "The tensor marking the input of the checkpoint stage.")
+ .signature(sig::arg("x", "The tensor marking the input of the checkpoint stage."))
.set_attr<FInferType>("FInferType", InferTypeStartCheckpoint)
.set_attr<bool>("FPurity", true);
}
@@ -81,7 +81,7 @@ TVM_FFI_STATIC_INIT_BLOCK() {
/* relax.grad.end_checkpoint */
Expr end_checkpoint(Expr input) {
static const Op op = Op::Get("relax.grad.end_checkpoint");
- return Call(Type::Missing(), op, {std::move(input)}, {}, {});
+ return Call::Unchecked(Type::Missing(), op, {std::move(input)}, {}, {});
}
TVM_FFI_STATIC_INIT_BLOCK() {
@@ -99,7 +99,7 @@ Type InferTypeEndCheckpoint(const Call& call, const BlockBuilder& ctx) {
TVM_FFI_STATIC_INIT_BLOCK() {
OpDef("relax.grad.end_checkpoint")
- .add_arg("x", "The output of the checkpoint stage.")
+ .signature(sig::arg("x", "The output of the checkpoint stage."))
.set_attr<FInferType>("FInferType", InferTypeEndCheckpoint)
.set_attr<bool>("FPurity", true);
}
@@ -119,9 +119,9 @@ Expr nll_loss_backward(Expr output_grad, Expr predictions, Expr targets,
{std::move(output_grad), std::move(predictions), std::move(targets), weights.value()},
Attrs{attrs}, {});
} else {
- return Call(Type::Missing(), op,
- {std::move(output_grad), std::move(predictions), std::move(targets)}, Attrs{attrs},
- {});
+ return Call::Unchecked(Type::Missing(), op,
+ {std::move(output_grad), std::move(predictions), std::move(targets)},
+ Attrs{attrs}, {});
}
}
@@ -136,11 +136,10 @@ Type InferTypeNLLLossBackward(const Call& call, const BlockBuilder& ctx) {
TVM_FFI_STATIC_INIT_BLOCK() {
OpDef("relax.grad.nll_loss_backward", "Optional weights: The weight of each target values.")
- .add_arg("output_grad", "The output gradient.")
- .add_arg("predictions", "The prediction tensor.")
- .add_arg("targets", "The target tensor.")
- .allow_extra_args()
- .attrs_type<NLLLossAttrs>()
+ .signature(sig::arg("output_grad", "The output gradient."),
+ 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<bool>("FPurity", true);
}
@@ -160,7 +159,8 @@ Expr max_pool2d_backward(Expr output_grad, Expr data, ffi::Array<int64_t> pool_s
attrs->layout = layout;
attrs->out_layout = out_layout.value_or(layout);
static const Op op = Op::Get("relax.grad.max_pool2d_backward");
- return Call(Type::Missing(), op, {std::move(output_grad), std::move(data)}, Attrs(attrs), {});
+ return Call::Unchecked(Type::Missing(), op, {std::move(output_grad), std::move(data)},
+ Attrs(attrs), {});
}
TVM_FFI_STATIC_INIT_BLOCK() {
@@ -174,9 +174,8 @@ Type InferTypeMaxPool2DBackward(const Call& call, const BlockBuilder& ctx) {
TVM_FFI_STATIC_INIT_BLOCK() {
OpDef("relax.grad.max_pool2d_backward")
- .add_arg("output_grad", "The output gradient.")
- .add_arg("data", "The input tensor")
- .attrs_type<Pool2DAttrs>()
+ .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<bool>("FPurity", true);
}
@@ -196,7 +195,8 @@ Expr avg_pool2d_backward(Expr output_grad, Expr data, ffi::Array<int64_t> pool_s
attrs->layout = layout;
attrs->out_layout = out_layout.value_or(layout);
static const Op op = Op::Get("relax.grad.avg_pool2d_backward");
- return Call(Type::Missing(), op, {std::move(output_grad), std::move(data)}, Attrs(attrs), {});
+ return Call::Unchecked(Type::Missing(), op, {std::move(output_grad), std::move(data)},
+ Attrs(attrs), {});
}
TVM_FFI_STATIC_INIT_BLOCK() {
@@ -210,9 +210,8 @@ Type InferTypeAvgPool2DBackward(const Call& call, const BlockBuilder& ctx) {
TVM_FFI_STATIC_INIT_BLOCK() {
OpDef("relax.grad.avg_pool2d_backward")
- .add_arg("output_grad", "The output gradient.")
- .add_arg("data", "The input tensor")
- .attrs_type<Pool2DAttrs>()
+ .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<bool>("FPurity", true);
}
@@ -224,8 +223,9 @@ Expr take_backward(Expr output_grad, Expr x, Expr indices, ffi::Optional<int64_t
attrs->axis = std::move(axis);
static const Op op = Op::Get("relax.grad.take_backward");
- return Call(Type::Missing(), op, {std::move(output_grad), std::move(x), std::move(indices)},
- Attrs(attrs), {});
+ return Call::Unchecked(Type::Missing(), op,
+ {std::move(output_grad), std::move(x), std::move(indices)}, Attrs(attrs),
+ {});
}
TVM_FFI_STATIC_INIT_BLOCK() {
@@ -239,10 +239,10 @@ Type InferTypeTakeBackward(const Call& call, const BlockBuilder& ctx) {
TVM_FFI_STATIC_INIT_BLOCK() {
OpDef("relax.grad.take_backward")
- .add_arg("output_grad", "The output gradient.")
- .add_arg("x", "The source tensor.")
- .add_arg("indices", "The indices of the values to extract.")
- .attrs_type<TakeAttrs>()
+ .signature(sig::arg("output_grad", "The output gradient."),
+ 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<bool>("FPurity", true);
}
diff --git a/src/relax/op/tensor/index.cc b/src/relax/op/tensor/index.cc
index 00ea940555..e1b73c3d15 100644
--- a/src/relax/op/tensor/index.cc
+++ b/src/relax/op/tensor/index.cc
@@ -53,7 +53,7 @@ Expr take(Expr x, Expr indices, ffi::Optional<int64_t> axis, ffi::String mode) {
attrs->mode = std::move(mode);
static const Op op = Op::Get("relax.take");
- return Call(Type::Missing(), op, {std::move(x), std::move(indices)}, Attrs(attrs), {});
+ return Call::Unchecked(Type::Missing(), op, {std::move(x), std::move(indices)}, Attrs(attrs), {});
}
TVM_FFI_STATIC_INIT_BLOCK() {
@@ -129,9 +129,9 @@ Type InferTypeTake(const Call& call, const BlockBuilder& ctx) {
TVM_FFI_STATIC_INIT_BLOCK() {
OpDef("relax.take")
- .add_arg("x", "The source tensor.")
- .add_arg("indices", "The indices of the values to extract.")
- .attrs_type<TakeAttrs>()
+ .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<bool>("FPurity", true);
}
@@ -175,7 +175,7 @@ Expr strided_slice(Expr x, Expr axes, Expr begin, Expr end, ffi::Optional<Expr>
}
static const Op op = Op::Get("relax.strided_slice");
- auto call = Call(Type::Missing(), op, args, Attrs(attrs));
+ auto call = Call::Unchecked(Type::Missing(), op, args, Attrs(attrs));
return call;
}
@@ -489,12 +489,9 @@ InferLayoutOutput InferLayoutStridedSlice(
TVM_FFI_STATIC_INIT_BLOCK() {
OpDef("relax.strided_slice")
- .add_arg("x", "The source tensor to be sliced.")
- .add_arg("axes", "The axes.")
- .add_arg("begin", "The start index.")
- .add_arg("end", "The end index.")
- .allow_extra_args()
- .attrs_type<StridedSliceAttrs>()
+ .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<FRelaxInferLayout>("FRelaxInferLayout", InferLayoutStridedSlice)
.set_attr<TMixedPrecisionPolicy>("TMixedPrecisionPolicy", MixedPrecisionPolicyKind::kFollow)
@@ -507,8 +504,8 @@ Expr dynamic_strided_slice(Expr x, //
Expr end, //
Expr strides) {
static const Op op = Op::Get("relax.dynamic_strided_slice");
- return Call(Type::Missing(), op,
- {std::move(x), std::move(begin), std::move(end), std::move(strides)}, {});
+ return Call::Unchecked(Type::Missing(), op,
+ {std::move(x), std::move(begin), std::move(end), std::move(strides)}, {});
}
TVM_FFI_STATIC_INIT_BLOCK() {
@@ -591,10 +588,10 @@ InferLayoutOutput InferLayoutDynStridedSlice(
TVM_FFI_STATIC_INIT_BLOCK() {
OpDef("relax.dynamic_strided_slice")
- .add_arg("x", "The source tensor to be sliced.")
- .add_arg("begin", "The indices to begin with in the slicing.")
- .add_arg("end", "Indices indicating end of the slice.")
- .add_arg("strides", "The stride values.")
+ .signature(sig::arg("x", "The source tensor to be sliced."),
+ 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<FRelaxInferLayout>("FRelaxInferLayout", InferLayoutDynStridedSlice)
.set_attr<TMixedPrecisionPolicy>("TMixedPrecisionPolicy", MixedPrecisionPolicyKind::kFollow)
diff --git a/src/relax/op/tensor/inspect.cc b/src/relax/op/tensor/inspect.cc
index 49bf821dec..5ef6a9dec7 100644
--- a/src/relax/op/tensor/inspect.cc
+++ b/src/relax/op/tensor/inspect.cc
@@ -118,7 +118,7 @@ Expr NormalizeToKnownPrimExpr(const BlockBuilder&, Call call) { return call; }
Expr tensor_dtype_code(Expr expr) {
static const Op op = Op::Get("relax.inspect.tensor_dtype_code");
- return Call(Type::Missing(), op, {expr});
+ return Call::Unchecked(Type::Missing(), op, {expr});
}
Type InferTypeTensorDtypeCode(const Call& call, const BlockBuilder&) { return PrimType::UInt(8); }
@@ -131,12 +131,12 @@ Expr LegalizeTensorDtypeCode(const BlockBuilder& bb, const Call& call) {
GetDLTensorField(tirx::builtin::TVMStructFieldKind::kDLTensorTypeCode, field_ty);
GlobalVar gvar_getter = bb->AddFunction(getter, "_get_tensor_dtype_code");
- return Call(Type::Missing(), gvar_getter, {arg});
+ return Call::Unchecked(Type::Missing(), gvar_getter, {arg});
}
TVM_FFI_STATIC_INIT_BLOCK() {
OpDef("relax.inspect.tensor_dtype_code")
- .add_arg("tensor", "The tensor to be inspected")
+ .signature(sig::arg("tensor", "The tensor to be inspected"))
.set_attr<FInferType>("FInferType", InferTypeTensorDtypeCode)
.set_attr<FLegalize>("FLegalize", LegalizeTensorDtypeCode)
.set_attr<bool>("RequiresArgumentShapes", false)
@@ -148,7 +148,7 @@ TVM_FFI_STATIC_INIT_BLOCK() {
Expr tensor_dtype_bits(Expr expr) {
static const Op op = Op::Get("relax.inspect.tensor_dtype_bits");
- return Call(Type::Missing(), op, {expr});
+ return Call::Unchecked(Type::Missing(), op, {expr});
}
Type InferTypeTensorDtypeBits(const Call& call, const BlockBuilder&) { return PrimType::UInt(8); }
@@ -161,12 +161,12 @@ Expr LegalizeTensorDtypeBits(const BlockBuilder& bb, const Call& call) {
GetDLTensorField(tirx::builtin::TVMStructFieldKind::kDLTensorTypeBits, field_ty);
GlobalVar gvar_getter = bb->AddFunction(getter, "_get_tensor_dtype_bits");
- return Call(Type::Missing(), gvar_getter, {arg});
+ return Call::Unchecked(Type::Missing(), gvar_getter, {arg});
}
TVM_FFI_STATIC_INIT_BLOCK() {
OpDef("relax.inspect.tensor_dtype_bits")
- .add_arg("tensor", "The tensor to be inspected")
+ .signature(sig::arg("tensor", "The tensor to be inspected"))
.set_attr<FInferType>("FInferType", InferTypeTensorDtypeBits)
.set_attr<FLegalize>("FLegalize", LegalizeTensorDtypeBits)
.set_attr<bool>("RequiresArgumentShapes", false)
@@ -178,7 +178,7 @@ TVM_FFI_STATIC_INIT_BLOCK() {
Expr tensor_dtype_lanes(Expr expr) {
static const Op op = Op::Get("relax.inspect.tensor_dtype_lanes");
- return Call(Type::Missing(), op, {expr});
+ return Call::Unchecked(Type::Missing(), op, {expr});
}
Type InferTypeTensorDtypeLanes(const Call& call, const BlockBuilder&) { return PrimType::UInt(16); }
@@ -191,12 +191,12 @@ Expr LegalizeTensorDtypeLanes(const BlockBuilder& bb, const Call& call) {
GetDLTensorField(tirx::builtin::TVMStructFieldKind::kDLTensorTypeLanes, field_ty);
GlobalVar gvar_getter = bb->AddFunction(getter, "_get_tensor_dtype_lanes");
- return Call(Type::Missing(), gvar_getter, {arg});
+ return Call::Unchecked(Type::Missing(), gvar_getter, {arg});
}
TVM_FFI_STATIC_INIT_BLOCK() {
OpDef("relax.inspect.tensor_dtype_lanes")
- .add_arg("tensor", "The tensor to be inspected")
+ .signature(sig::arg("tensor", "The tensor to be inspected"))
.set_attr<FInferType>("FInferType", InferTypeTensorDtypeLanes)
.set_attr<FLegalize>("FLegalize", LegalizeTensorDtypeLanes)
.set_attr<bool>("RequiresArgumentShapes", false)
@@ -208,7 +208,7 @@ TVM_FFI_STATIC_INIT_BLOCK() {
Expr tensor_ndim(Expr expr) {
static const Op op = Op::Get("relax.inspect.tensor_ndim");
- return Call(Type::Missing(), op, {expr});
+ return Call::Unchecked(Type::Missing(), op, {expr});
}
Type InferTypeTensorNDim(const Call& call, const BlockBuilder&) { return PrimType::Int(32); }
@@ -221,12 +221,12 @@ Expr LegalizeTensorNDim(const BlockBuilder& bb, const Call& call) {
GetDLTensorField(tirx::builtin::TVMStructFieldKind::kDLTensorNDim, field_ty);
GlobalVar gvar_getter = bb->AddFunction(getter, "_get_tensor_ndim");
- return Call(Type::Missing(), gvar_getter, {arg});
+ return Call::Unchecked(Type::Missing(), gvar_getter, {arg});
}
TVM_FFI_STATIC_INIT_BLOCK() {
OpDef("relax.inspect.tensor_ndim")
- .add_arg("tensor", "The tensor to be inspected")
+ .signature(sig::arg("tensor", "The tensor to be inspected"))
.set_attr<FInferType>("FInferType", InferTypeTensorNDim)
.set_attr<FLegalize>("FLegalize", LegalizeTensorNDim)
.set_attr<bool>("RequiresArgumentShapes", false)
@@ -238,7 +238,7 @@ TVM_FFI_STATIC_INIT_BLOCK() {
Expr tensor_shape_i(Expr expr) {
static const Op op = Op::Get("relax.inspect.tensor_shape_i");
- return Call(Type::Missing(), op, {expr});
+ return Call::Unchecked(Type::Missing(), op, {expr});
}
Type InferTypeTensorShape(const Call& call, const BlockBuilder&) {
@@ -301,13 +301,13 @@ Expr LegalizeTensorShape(const BlockBuilder& bb, const Call& call) {
}();
GlobalVar gvar_getter = bb->AddFunction(getter, "_get_tensor_shape_i");
- return Call(Type::Missing(), gvar_getter, call->args);
+ return Call::Unchecked(Type::Missing(), gvar_getter, call->args);
}
TVM_FFI_STATIC_INIT_BLOCK() {
OpDef("relax.inspect.tensor_shape_i")
- .add_arg("tensor", "The tensor to be inspected")
- .add_arg("axis", "The axis whose extent should be returned")
+ .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<FLegalize>("FLegalize", LegalizeTensorShape)
.set_attr<bool>("RequiresArgumentShapes", false)
@@ -319,7 +319,7 @@ TVM_FFI_STATIC_INIT_BLOCK() {
Expr tensor_stride_i(Expr expr) {
static const Op op = Op::Get("relax.inspect.tensor_stride_i");
- return Call(Type::Missing(), op, {expr});
+ return Call::Unchecked(Type::Missing(), op, {expr});
}
Type InferTypeTensorStride(const Call& call, const BlockBuilder&) {
@@ -355,8 +355,8 @@ Type InferTypeTensorStride(const Call& call, const BlockBuilder&) {
TVM_FFI_STATIC_INIT_BLOCK() {
OpDef("relax.inspect.tensor_stride_i")
- .add_arg("tensor", "The tensor to be inspected")
- .add_arg("axis", "The axis whose extent should be returned")
+ .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<bool>("RequiresArgumentShapes", false)
.set_attr<FNormalize>("FNormalize", NormalizeToKnownPrimExpr)
@@ -367,7 +367,7 @@ TVM_FFI_STATIC_INIT_BLOCK() {
Expr tensor_byte_offset(Expr expr) {
static const Op op = Op::Get("relax.inspect.tensor_byte_offset");
- return Call(Type::Missing(), op, {expr});
+ return Call::Unchecked(Type::Missing(), op, {expr});
}
Type InferTypeTensorByteOffset(const Call& call, const BlockBuilder&) {
@@ -388,7 +388,7 @@ Type InferTypeTensorByteOffset(const Call& call, const BlockBuilder&) {
TVM_FFI_STATIC_INIT_BLOCK() {
OpDef("relax.inspect.tensor_byte_offset")
- .add_arg("tensor", "The tensor to be inspected")
+ .signature(sig::arg("tensor", "The tensor to be inspected"))
.set_attr<FInferType>("FInferType", InferTypeTensorByteOffset)
.set_attr<bool>("RequiresArgumentShapes", false)
.set_attr<FNormalize>("FNormalize", NormalizeToKnownPrimExpr)
@@ -399,7 +399,7 @@ TVM_FFI_STATIC_INIT_BLOCK() {
Expr tensor_elem_offset(Expr expr) {
static const Op op = Op::Get("relax.inspect.tensor_elem_offset");
- return Call(Type::Missing(), op, {expr});
+ return Call::Unchecked(Type::Missing(), op, {expr});
}
Type InferTypeTensorElemOffset(const Call& call, const BlockBuilder&) {
@@ -420,7 +420,7 @@ Type InferTypeTensorElemOffset(const Call& call, const BlockBuilder&) {
TVM_FFI_STATIC_INIT_BLOCK() {
OpDef("relax.inspect.tensor_elem_offset")
- .add_arg("tensor", "The tensor to be inspected")
+ .signature(sig::arg("tensor", "The tensor to be inspected"))
.set_attr<FInferType>("FInferType", InferTypeTensorElemOffset)
.set_attr<bool>("RequiresArgumentShapes", false)
.set_attr<FNormalize>("FNormalize", NormalizeToKnownPrimExpr)
diff --git a/src/relax/op/tensor/linear_algebra.cc b/src/relax/op/tensor/linear_algebra.cc
index e9bf495c1b..5ab4566909 100644
--- a/src/relax/op/tensor/linear_algebra.cc
+++ b/src/relax/op/tensor/linear_algebra.cc
@@ -47,7 +47,7 @@ Expr matmul(Expr x1, Expr x2, ffi::Optional<DLDataType> out_dtype) {
attrs->out_dtype = out_dtype;
static const Op op = Op::Get("relax.matmul");
- return Call(Type::Missing(), op, {std::move(x1), std::move(x2)}, Attrs(attrs), {});
+ return Call::Unchecked(Type::Missing(), op, {std::move(x1), std::move(x2)}, Attrs(attrs), {});
}
TVM_FFI_STATIC_INIT_BLOCK() {
@@ -164,8 +164,8 @@ Call InferMixedPrecisionMatmul(const Call& call, DLDataType out_dtype) {
TVM_FFI_STATIC_INIT_BLOCK() {
OpDef("relax.matmul")
- .add_arg("x1", "The first input tensor.")
- .add_arg("x2", "The second input tensor.")
+ .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<TMixedPrecisionPolicy>("TMixedPrecisionPolicy", MixedPrecisionPolicyKind::kAlways)
.set_attr<FInferMixedPrecision>("FInferMixedPrecision", InferMixedPrecisionMatmul)
@@ -179,7 +179,7 @@ Expr einsum(Expr operands, ffi::String subscripts) {
attrs->subscripts = std::move(subscripts);
static const Op op = Op::Get("relax.einsum");
- return Call(Type::Missing(), op, {std::move(operands)}, Attrs{attrs}, {});
+ return Call::Unchecked(Type::Missing(), op, {std::move(operands)}, Attrs{attrs}, {});
}
TVM_FFI_STATIC_INIT_BLOCK() {
@@ -254,8 +254,7 @@ Type InferTypeEinsum(const Call& call, const BlockBuilder& ctx) {
TVM_FFI_STATIC_INIT_BLOCK() {
OpDef("relax.einsum")
- .add_arg("operands", "The input tensors.")
- .attrs_type<EinsumAttrs>()
+ .signature(sig::arg("operands", "The input tensors."), sig::call_attrs<EinsumAttrs>())
.set_attr<FInferType>("FInferType", InferTypeEinsum)
.set_attr<bool>("FPurity", true);
}
@@ -264,7 +263,7 @@ TVM_FFI_STATIC_INIT_BLOCK() {
Expr outer(Expr x1, Expr x2) {
static const Op op = Op::Get("relax.outer");
- return Call(Type::Missing(), op, {std::move(x1), std::move(x2)}, {});
+ return Call::Unchecked(Type::Missing(), op, {std::move(x1), std::move(x2)}, {});
}
TVM_FFI_STATIC_INIT_BLOCK() {
@@ -294,8 +293,8 @@ Type InferTypeOuter(const Call& call, const BlockBuilder& ctx) {
TVM_FFI_STATIC_INIT_BLOCK() {
OpDef("relax.outer")
- .add_arg("x1", "The first input tensor.")
- .add_arg("x2", "The second input tensor.")
+ .signature(sig::arg("x1", "The first input tensor."),
+ sig::arg("x2", "The second input tensor."))
.set_attr<FInferType>("FInferType", 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 64d8d62c75..347582a67c 100644
--- a/src/relax/op/tensor/manipulate.cc
+++ b/src/relax/op/tensor/manipulate.cc
@@ -65,7 +65,7 @@ TVM_FFI_STATIC_INIT_BLOCK() {
/* relax.broadcast_to */
Expr broadcast_to(Expr x, Expr shape) {
static const Op op = Op::Get("relax.broadcast_to");
- return Call(Type::Missing(), op, {std::move(x), std::move(shape)}, Attrs(), {});
+ return Call::Unchecked(Type::Missing(), op, {std::move(x), std::move(shape)}, Attrs(), {});
}
TVM_FFI_STATIC_INIT_BLOCK() {
@@ -133,8 +133,7 @@ Type InferTypeBroadcastTo(const Call& call, const BlockBuilder& ctx) {
TVM_FFI_STATIC_INIT_BLOCK() {
OpDef("relax.broadcast_to")
- .add_arg("x", "The input tensor.")
- .add_arg("shape", "The target shape.")
+ .signature(sig::arg("x", "The input tensor."), sig::arg("shape", "The target shape."))
.set_attr<FInferType>("FInferType", InferTypeBroadcastTo)
.set_attr<TMixedPrecisionPolicy>("TMixedPrecisionPolicy", MixedPrecisionPolicyKind::kFollow)
.set_attr<bool>("FPurity", true);
@@ -147,7 +146,7 @@ Expr concat(Expr tensors, ffi::Optional<int64_t> axis) {
attrs->axis = std::move(axis);
static const Op op = Op::Get("relax.concat");
- return Call(Type::Missing(), op, {std::move(tensors)}, Attrs(attrs), {});
+ return Call::Unchecked(Type::Missing(), op, {std::move(tensors)}, Attrs(attrs), {});
}
TVM_FFI_STATIC_INIT_BLOCK() {
@@ -393,8 +392,7 @@ InferLayoutOutput InferLayoutConcat(
TVM_FFI_STATIC_INIT_BLOCK() {
OpDef("relax.concat")
- .add_arg("tensors", "The input list of tensors.")
- .attrs_type<ConcatAttrs>()
+ .signature(sig::arg("tensors", "The input list of tensors."), sig::call_attrs<ConcatAttrs>())
.set_attr<FInferType>("FInferType", InferTypeConcat)
.set_attr<FRelaxInferLayout>("FRelaxInferLayout", InferLayoutConcat)
.set_attr<TMixedPrecisionPolicy>("TMixedPrecisionPolicy", MixedPrecisionPolicyKind::kFollow)
@@ -408,7 +406,7 @@ Expr expand_dims(Expr x, ffi::Array<int64_t> axis) {
attrs->axis = std::move(axis);
static const Op op = Op::Get("relax.expand_dims");
- return Call(Type::Missing(), op, {std::move(x)}, Attrs{attrs}, {});
+ return Call::Unchecked(Type::Missing(), op, {std::move(x)}, Attrs{attrs}, {});
}
TVM_FFI_STATIC_INIT_BLOCK() {
@@ -498,8 +496,7 @@ InferLayoutOutput InferLayoutExpandDims(
TVM_FFI_STATIC_INIT_BLOCK() {
OpDef("relax.expand_dims")
- .add_arg("x", "The input tensor.")
- .attrs_type<ExpandDimsAttrs>()
+ .signature(sig::arg("x", "The input tensor."), sig::call_attrs<ExpandDimsAttrs>())
.set_attr<FInferType>("FInferType", InferTypeExpandDims)
.set_attr<FRelaxInferLayout>("FRelaxInferLayout", InferLayoutExpandDims)
.set_attr<TMixedPrecisionPolicy>("TMixedPrecisionPolicy", MixedPrecisionPolicyKind::kFollow)
@@ -518,7 +515,7 @@ PrimExpr ComputeShapeProduct(const ffi::Array<PrimExpr>& shape_values) {
/* relax.flatten */
Expr flatten(Expr x) {
static const Op op = Op::Get("relax.flatten");
- return Call(Type::Missing(), op, {std::move(x)}, {}, {});
+ return Call::Unchecked(Type::Missing(), op, {std::move(x)}, {}, {});
}
TVM_FFI_STATIC_INIT_BLOCK() {
@@ -546,7 +543,7 @@ Type InferTypeFlatten(const Call& call, const BlockBuilder& ctx) {
TVM_FFI_STATIC_INIT_BLOCK() {
OpDef("relax.flatten")
- .add_arg("x", "The input tensor.")
+ .signature(sig::arg("x", "The input tensor."))
.set_attr<FInferType>("FInferType", InferTypeFlatten)
.set_attr<TMixedPrecisionPolicy>("TMixedPrecisionPolicy", MixedPrecisionPolicyKind::kFollow)
.set_attr<bool>("FPurity", true);
@@ -556,7 +553,7 @@ TVM_FFI_STATIC_INIT_BLOCK() {
Expr index_tensor(Expr first, Expr tensors) {
static const Op op = Op::Get("relax.index_tensor");
- return Call(Type::Missing(), op, {std::move(first), std::move(tensors)}, Attrs(), {});
+ return Call::Unchecked(Type::Missing(), op, {std::move(first), std::move(tensors)}, Attrs(), {});
}
TVM_FFI_STATIC_INIT_BLOCK() {
@@ -696,8 +693,8 @@ Type InferTypeIndexTensor(const Call& call, const BlockBuilder& ctx) {
TVM_FFI_STATIC_INIT_BLOCK() {
OpDef("relax.index_tensor")
- .add_arg("data", "The input data.")
- .add_arg("indices", "The indices used to index.")
+ .signature(sig::arg("data", "The input data."),
+ sig::arg("indices", "The indices used to index."))
.set_attr<FInferType>("FInferType", InferTypeIndexTensor)
.set_attr<bool>("FPurity", true);
}
@@ -710,7 +707,7 @@ Expr layout_transform(Expr x, tirx::IndexMap index_map, ffi::Optional<PrimExpr>
attrs->pad_value = std::move(pad_value);
static const Op op = Op::Get("relax.layout_transform");
- return Call(Type::Missing(), op, {std::move(x)}, Attrs{attrs}, {});
+ return Call::Unchecked(Type::Missing(), op, {std::move(x)}, Attrs{attrs}, {});
}
TVM_FFI_STATIC_INIT_BLOCK() {
@@ -764,8 +761,7 @@ Type InferTypeLayoutTransform(const Call& call, const BlockBuilder& ctx) {
TVM_FFI_STATIC_INIT_BLOCK() {
OpDef("relax.layout_transform")
- .add_arg("x", "The input tensor.")
- .attrs_type<LayoutTransformAttrs>()
+ .signature(sig::arg("x", "The input tensor."), sig::call_attrs<LayoutTransformAttrs>())
.set_attr<FInferType>("FInferType", InferTypeLayoutTransform)
.set_attr<TMixedPrecisionPolicy>("TMixedPrecisionPolicy", MixedPrecisionPolicyKind::kFollow)
.set_attr<bool>("FPurity", true);
@@ -778,7 +774,7 @@ Expr permute_dims(Expr x, ffi::Optional<ffi::Array<int64_t>> axes) {
attrs->axes = std::move(axes);
static const Op op = Op::Get("relax.permute_dims");
- return Call(Type::Missing(), op, {std::move(x)}, Attrs{attrs}, {});
+ return Call::Unchecked(Type::Missing(), op, {std::move(x)}, Attrs{attrs}, {});
}
TVM_FFI_STATIC_INIT_BLOCK() {
@@ -885,8 +881,7 @@ InferLayoutOutput InferLayoutPermuteDims(
TVM_FFI_STATIC_INIT_BLOCK() {
OpDef("relax.permute_dims")
- .add_arg("x", "The input tensor.")
- .attrs_type<PermuteDimsAttrs>()
+ .signature(sig::arg("x", "The input tensor."), sig::call_attrs<PermuteDimsAttrs>())
.set_attr<FInferType>("FInferType", InferTypePermuteDims)
.set_attr<FRelaxInferLayout>("FRelaxInferLayout", InferLayoutPermuteDims)
.set_attr<TMixedPrecisionPolicy>("TMixedPrecisionPolicy", MixedPrecisionPolicyKind::kFollow)
@@ -998,7 +993,8 @@ Expr ConvertNewShapeToExpr(const Expr& data,
Expr reshape(Expr x, ffi::Variant<Expr, ffi::Array<PrimExpr>> shape) {
Expr shape_in_expr = ConvertNewShapeToExpr(x, shape);
static const Op op = Op::Get("relax.reshape");
- return Call(Type::Missing(), op, {std::move(x), std::move(shape_in_expr)}, Attrs(), {});
+ return Call::Unchecked(Type::Missing(), op, {std::move(x), std::move(shape_in_expr)}, Attrs(),
+ {});
}
TVM_FFI_STATIC_INIT_BLOCK() {
@@ -1051,8 +1047,7 @@ Type InferTypeReshape(const Call& call, const BlockBuilder& ctx) {
TVM_FFI_STATIC_INIT_BLOCK() {
OpDef("relax.reshape")
- .add_arg("x", "The input tensor.")
- .add_arg("shape", "The input new shape.")
+ .signature(sig::arg("x", "The input tensor."), sig::arg("shape", "The input new shape."))
.set_attr<FInferType>("FInferType", InferTypeReshape)
.set_attr<TMixedPrecisionPolicy>("TMixedPrecisionPolicy", MixedPrecisionPolicyKind::kFollow)
.set_attr<bool>("FPurity", true);
@@ -1089,7 +1084,7 @@ Expr split(Expr x, ffi::Variant<IntImm, ffi::Array<IntImm>> indices_or_sections,
attrs->axis = axis;
static const Op op = Op::Get("relax.split");
- return Call(Type::Missing(), op, {std::move(x)}, Attrs(attrs), {});
+ return Call::Unchecked(Type::Missing(), op, {std::move(x)}, Attrs(attrs), {});
}
TVM_FFI_STATIC_INIT_BLOCK() {
@@ -1225,8 +1220,7 @@ InferLayoutOutput InferLayoutSplit(
TVM_FFI_STATIC_INIT_BLOCK() {
OpDef("relax.split")
- .add_arg("x", "The input tensor.")
- .attrs_type<SplitAttrs>()
+ .signature(sig::arg("x", "The input tensor."), sig::call_attrs<SplitAttrs>())
.set_attr<FInferType>("FInferType", InferTypeSplit)
.set_attr<FRelaxInferLayout>("FRelaxInferLayout", InferLayoutSplit)
.set_attr<TMixedPrecisionPolicy>("TMixedPrecisionPolicy", MixedPrecisionPolicyKind::kFollow)
@@ -1240,7 +1234,7 @@ Expr squeeze(Expr x, ffi::Optional<ffi::Array<int64_t>> axis) {
attrs->axis = std::move(axis);
static const Op op = Op::Get("relax.squeeze");
- return Call(Type::Missing(), op, {std::move(x)}, Attrs{attrs}, {});
+ return Call::Unchecked(Type::Missing(), op, {std::move(x)}, Attrs{attrs}, {});
}
TVM_FFI_STATIC_INIT_BLOCK() {
@@ -1385,8 +1379,7 @@ InferLayoutOutput InferLayoutSqueeze(
TVM_FFI_STATIC_INIT_BLOCK() {
OpDef("relax.squeeze")
- .add_arg("x", "The input tensor.")
- .attrs_type<SqueezeAttrs>()
+ .signature(sig::arg("x", "The input tensor."), sig::call_attrs<SqueezeAttrs>())
.set_attr<FInferType>("FInferType", InferTypeSqueeze)
.set_attr<FRelaxInferLayout>("FRelaxInferLayout", InferLayoutSqueeze)
.set_attr<TMixedPrecisionPolicy>("TMixedPrecisionPolicy", MixedPrecisionPolicyKind::kFollow)
@@ -1439,7 +1432,7 @@ Expr stack(Expr tensors, ffi::Optional<int64_t> axis) {
attrs->axis = std::move(axis);
static const Op op = Op::Get("relax.stack");
- return Call(Type::Missing(), op, {std::move(tensors)}, Attrs(attrs), {});
+ return Call::Unchecked(Type::Missing(), op, {std::move(tensors)}, Attrs(attrs), {});
}
TVM_FFI_STATIC_INIT_BLOCK() {
@@ -1638,8 +1631,8 @@ InferLayoutOutput InferLayoutStack(
TVM_FFI_STATIC_INIT_BLOCK() {
OpDef("relax.stack")
- .add_arg("tensors", "The input list of tensors to stack")
- .attrs_type<StackAttrs>()
+ .signature(sig::arg("tensors", "The input list of tensors to stack"),
+ sig::call_attrs<StackAttrs>())
.set_attr<FInferType>("FInferType", InferTypeStack)
.set_attr<FRelaxInferLayout>("FRelaxInferLayout", InferLayoutStack)
.set_attr<TMixedPrecisionPolicy>("TMixedPrecisionPolicy", MixedPrecisionPolicyKind::kFollow)
@@ -1649,7 +1642,8 @@ TVM_FFI_STATIC_INIT_BLOCK() {
/* relax.collapse_sum_like */
Expr collapse_sum_like(Expr data, Expr collapse_target) {
static const Op op = Op::Get("relax.collapse_sum_like");
- return Call(Type::Missing(), op, {std::move(data), std::move(collapse_target)}, Attrs(), {});
+ return Call::Unchecked(Type::Missing(), op, {std::move(data), std::move(collapse_target)},
+ Attrs(), {});
}
TVM_FFI_STATIC_INIT_BLOCK() {
@@ -1687,8 +1681,8 @@ Type InferTypeCollapseSumLike(const Call& call, const BlockBuilder& ctx) {
TVM_FFI_STATIC_INIT_BLOCK() {
OpDef("relax.collapse_sum_like")
- .add_arg("data", "The input tensor.")
- .add_arg("collapse_target", "The tensor whose shape is the shape to collapse to.")
+ .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<bool>("FPurity", true);
}
@@ -1696,7 +1690,7 @@ TVM_FFI_STATIC_INIT_BLOCK() {
/* relax.collapse_sum_to */
Expr collapse_sum_to(Expr data, Expr shape) {
static const Op op = Op::Get("relax.collapse_sum_to");
- return Call(Type::Missing(), op, {std::move(data), std::move(shape)}, Attrs(), {});
+ return Call::Unchecked(Type::Missing(), op, {std::move(data), std::move(shape)}, Attrs(), {});
}
TVM_FFI_STATIC_INIT_BLOCK() {
@@ -1738,8 +1732,8 @@ Type InferTypeCollapseSumTo(const Call& call, const BlockBuilder& ctx) {
TVM_FFI_STATIC_INIT_BLOCK() {
OpDef("relax.collapse_sum_to")
- .add_arg("data", "The input tensor.")
- .add_arg("shape", "The shape to collapse to.")
+ .signature(sig::arg("data", "The input tensor."),
+ sig::arg("shape", "The shape to collapse to."))
.set_attr<FInferType>("FInferType", InferTypeCollapseSumTo)
.set_attr<bool>("FPurity", true);
}
@@ -1752,7 +1746,7 @@ Expr repeat(Expr data, int repeats, ffi::Optional<int64_t> axis) {
attrs->axis = std::move(axis);
static const Op op = Op::Get("relax.repeat");
- return Call(Type::Missing(), op, {std::move(data)}, Attrs{attrs}, {});
+ return Call::Unchecked(Type::Missing(), op, {std::move(data)}, Attrs{attrs}, {});
}
TVM_FFI_STATIC_INIT_BLOCK() {
@@ -1863,8 +1857,7 @@ InferLayoutOutput InferLayoutRepeat(
TVM_FFI_STATIC_INIT_BLOCK() {
OpDef("relax.repeat")
- .add_arg("data", "The input tensor.")
- .attrs_type<RepeatAttrs>()
+ .signature(sig::arg("data", "The input tensor."), sig::call_attrs<RepeatAttrs>())
.set_attr<FInferType>("FInferType", InferTypeRepeat)
.set_attr<FRelaxInferLayout>("FRelaxInferLayout", InferLayoutRepeat)
.set_attr<bool>("FPurity", true);
@@ -1877,7 +1870,7 @@ Expr tile(Expr data, ffi::Array<int64_t> repeats) {
attrs->repeats = std::move(repeats);
static const Op op = Op::Get("relax.tile");
- return Call(Type::Missing(), op, {std::move(data)}, Attrs{attrs}, {});
+ return Call::Unchecked(Type::Missing(), op, {std::move(data)}, Attrs{attrs}, {});
}
TVM_FFI_STATIC_INIT_BLOCK() {
@@ -2008,8 +2001,7 @@ InferLayoutOutput InferLayoutTile(
TVM_FFI_STATIC_INIT_BLOCK() {
OpDef("relax.tile")
- .add_arg("data", "The input tensor.")
- .attrs_type<TileAttrs>()
+ .signature(sig::arg("data", "The input tensor."), sig::call_attrs<TileAttrs>())
.set_attr<FInferType>("FInferType", InferTypeTile)
.set_attr<FRelaxInferLayout>("FRelaxInferLayout", InferLayoutTile)
.set_attr<bool>("FPurity", true);
@@ -2021,7 +2013,7 @@ Expr flip(Expr data, int64_t axis) {
auto attrs = ffi::make_object<FlipAttrs>();
attrs->axis = axis;
static const Op op = Op::Get("relax.flip");
- return Call(Type::Missing(), op, {std::move(data)}, Attrs{attrs}, {});
+ return Call::Unchecked(Type::Missing(), op, {std::move(data)}, Attrs{attrs}, {});
}
TVM_FFI_STATIC_INIT_BLOCK() {
@@ -2081,8 +2073,7 @@ InferLayoutOutput InferLayoutFlip(
TVM_FFI_STATIC_INIT_BLOCK() {
OpDef("relax.flip")
- .add_arg("data", "The input tensor.")
- .attrs_type<FlipAttrs>()
+ .signature(sig::arg("data", "The input tensor."), sig::call_attrs<FlipAttrs>())
.set_attr<FInferType>("FInferType", InferTypeFlip)
.set_attr<FRelaxInferLayout>("FRelaxInferLayout", InferLayoutFlip)
.set_attr<bool>("FPurity", true);
@@ -2095,7 +2086,8 @@ Expr reverse_sequence(Expr data, Expr seq_lengths, int64_t seq_axis, int64_t bat
attrs->seq_axis = seq_axis;
attrs->batch_axis = batch_axis;
static const Op op = Op::Get("relax.reverse_sequence");
- return Call(Type::Missing(), op, {std::move(data), std::move(seq_lengths)}, Attrs{attrs}, {});
+ return Call::Unchecked(Type::Missing(), op, {std::move(data), std::move(seq_lengths)},
+ Attrs{attrs}, {});
}
TVM_FFI_STATIC_INIT_BLOCK() {
@@ -2175,9 +2167,9 @@ Type InferTypeReverseSequence(const Call& call, const BlockBuilder& ctx) {
TVM_FFI_STATIC_INIT_BLOCK() {
OpDef("relax.reverse_sequence")
- .add_arg("data", "The input tensor.")
- .add_arg("seq_lengths", "The sequence length tensor.")
- .attrs_type<ReverseSequenceAttrs>()
+ .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<bool>("FPurity", true);
}
@@ -2188,7 +2180,7 @@ Expr gather_elements(Expr data, Expr indices, int axis) {
auto attrs = ffi::make_object<GatherElementsAttrs>();
attrs->axis = axis;
static const Op op = Op::Get("relax.gather_elements");
- return Call(Type::Missing(), op, {data, indices}, Attrs(attrs), {});
+ return Call::Unchecked(Type::Missing(), op, {data, indices}, Attrs(attrs), {});
}
TVM_FFI_STATIC_INIT_BLOCK() {
@@ -2277,9 +2269,8 @@ InferLayoutOutput InferLayoutGatherElements(
TVM_FFI_STATIC_INIT_BLOCK() {
OpDef("relax.gather_elements")
- .add_arg("data", "The input tensor.")
- .add_arg("indices", "The indices tensor.")
- .attrs_type<GatherElementsAttrs>()
+ .signature(sig::arg("data", "The input tensor."), sig::arg("indices", "The indices tensor."),
+ sig::call_attrs<GatherElementsAttrs>())
.set_attr<FInferType>("FInferType", InferTypeGatherElements)
.set_attr<FRelaxInferLayout>("FRelaxInferLayout", InferLayoutGatherElements)
.set_attr<bool>("FPurity", true);
@@ -2291,7 +2282,7 @@ Expr gather_nd(Expr data, Expr indices, int batch_dims) {
auto attrs = ffi::make_object<GatherNDAttrs>();
attrs->batch_dims = batch_dims;
static const Op op = Op::Get("relax.gather_nd");
- return Call(Type::Missing(), op, {data, indices}, Attrs(attrs), {});
+ return Call::Unchecked(Type::Missing(), op, {data, indices}, Attrs(attrs), {});
}
TVM_FFI_STATIC_INIT_BLOCK() {
@@ -2372,9 +2363,8 @@ Type InferTypeGatherND(const Call& call, const BlockBuilder& ctx) {
TVM_FFI_STATIC_INIT_BLOCK() {
OpDef("relax.gather_nd")
- .add_arg("data", "The input tensor.")
- .add_arg("indices", "The indices tensor.")
- .attrs_type<GatherNDAttrs>()
+ .signature(sig::arg("data", "The input tensor."), sig::arg("indices", "The indices tensor."),
+ sig::call_attrs<GatherNDAttrs>())
.set_attr<FInferType>("FInferType", InferTypeGatherND)
.set_attr<bool>("FPurity", true);
}
@@ -2385,7 +2375,7 @@ Expr index_put(Expr data, Expr indices, Expr values, bool accumulate) {
auto attrs = ffi::make_object<IndexPutAttrs>();
attrs->accumulate = std::move(accumulate);
static const Op op = Op::Get("relax.index_put");
- return Call(Type::Missing(), op, {data, indices, values}, Attrs(attrs), {});
+ return Call::Unchecked(Type::Missing(), op, {data, indices, values}, Attrs(attrs), {});
}
TVM_FFI_STATIC_INIT_BLOCK() {
@@ -2523,10 +2513,9 @@ Type InferTypeIndexPut(const Call& call, const BlockBuilder& ctx) {
TVM_FFI_STATIC_INIT_BLOCK() {
OpDef("relax.index_put")
- .add_arg("data", "The input tensor.")
- .add_arg("indices", "The indices tensor(s).")
- .add_arg("values", "The values to put.")
- .attrs_type<IndexPutAttrs>()
+ .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<bool>("FPurity", true);
}
@@ -2537,7 +2526,7 @@ Expr meshgrid(Expr tensors, ffi::Optional<ffi::String> indexing) {
ffi::ObjectPtr<MeshgridAttrs> attrs = ffi::make_object<MeshgridAttrs>();
attrs->indexing = indexing;
static const Op op = Op::Get("relax.meshgrid");
- return Call(Type::Missing(), op, {std::move(tensors)}, Attrs(attrs), {});
+ return Call::Unchecked(Type::Missing(), op, {std::move(tensors)}, Attrs(attrs), {});
}
TVM_FFI_STATIC_INIT_BLOCK() {
@@ -2630,8 +2619,8 @@ Type InferTypeMeshgrid(const Call& call, const BlockBuilder& ctx) {
TVM_FFI_STATIC_INIT_BLOCK() {
OpDef("relax.meshgrid")
- .add_arg("tensors", "The input list of tensors.")
- .attrs_type<MeshgridAttrs>()
+ .signature(sig::arg("tensors", "The input list of tensors."),
+ sig::call_attrs<MeshgridAttrs>())
.set_attr<FInferType>("FInferType", InferTypeMeshgrid)
.set_attr<TMixedPrecisionPolicy>("TMixedPrecisionPolicy", MixedPrecisionPolicyKind::kFollow)
.set_attr<bool>("FPurity", true);
@@ -2644,7 +2633,7 @@ Expr scatter_elements(Expr data, Expr indices, Expr updates, int axis, ffi::Stri
attrs->axis = std::move(axis);
attrs->reduction = std::move(reduction);
static const Op op = Op::Get("relax.scatter_elements");
- return Call(Type::Missing(), op, {data, indices, updates}, Attrs(attrs), {});
+ return Call::Unchecked(Type::Missing(), op, {data, indices, updates}, Attrs(attrs), {});
}
TVM_FFI_STATIC_INIT_BLOCK() {
@@ -2774,10 +2763,9 @@ InferLayoutOutput InferLayoutScatterElements(
TVM_FFI_STATIC_INIT_BLOCK() {
OpDef("relax.scatter_elements")
- .add_arg("data", "The input tensor.")
- .add_arg("indices", "The indices tensor.")
- .add_arg("updates", "The input tensor of updates.")
- .attrs_type<ScatterElementsAttrs>()
+ .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<FRelaxInferLayout>("FRelaxInferLayout", InferLayoutScatterElements)
.set_attr<bool>("FPurity", true);
@@ -2789,7 +2777,7 @@ Expr scatter_nd(Expr data, Expr indices, Expr updates, ffi::String reduction) {
auto attrs = ffi::make_object<ScatterNDAttrs>();
attrs->reduction = std::move(reduction);
static const Op op = Op::Get("relax.scatter_nd");
- return Call(Type::Missing(), op, {data, indices, updates}, Attrs(attrs), {});
+ return Call::Unchecked(Type::Missing(), op, {data, indices, updates}, Attrs(attrs), {});
}
TVM_FFI_STATIC_INIT_BLOCK() {
@@ -2953,10 +2941,9 @@ InferLayoutOutput InferLayoutScatterND(
TVM_FFI_STATIC_INIT_BLOCK() {
OpDef("relax.scatter_nd")
- .add_arg("data", "The input tensor.")
- .add_arg("indices", "The indices tensor.")
- .add_arg("updates", "The input tensor of updates.")
- .attrs_type<ScatterNDAttrs>()
+ .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<FRelaxInferLayout>("FRelaxInferLayout", InferLayoutScatterND)
.set_attr<bool>("FPurity", true);
@@ -2968,7 +2955,7 @@ Expr slice_scatter(Expr input, Expr src, int axis, PrimExpr start, PrimExpr end,
auto attrs = ffi::make_object<SliceScatterAttrs>();
attrs->axis = std::move(axis);
static const Op op = Op::Get("relax.slice_scatter");
- return Call(Type::Missing(), op, {input, src, start, end, step}, Attrs(attrs), {});
+ return Call::Unchecked(Type::Missing(), op, {input, src, start, end, step}, Attrs(attrs), {});
}
TVM_FFI_STATIC_INIT_BLOCK() {
@@ -3108,12 +3095,11 @@ Type InferTypeSliceScatter(const Call& call, const BlockBuilder& ctx) {
TVM_FFI_STATIC_INIT_BLOCK() {
OpDef("relax.slice_scatter")
- .add_arg("input", "The input tensor.")
- .add_arg("src", "The source tensor to scatter.")
- .add_arg("start", "The starting index of the slice (inclusive).")
- .add_arg("end", "The ending index of the slice (exclusive).")
- .add_arg("step", "The step of the slice.")
- .attrs_type<SliceScatterAttrs>()
+ .signature(sig::arg("input", "The input tensor."),
+ sig::arg("src", "The source tensor to scatter."),
+ 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<bool>("FPurity", true);
}
@@ -3135,7 +3121,7 @@ Expr one_hot(Expr indices, PrimExpr on_value, PrimExpr off_value, int depth, int
TVM_FFI_ICHECK(depth > 0) << "one_hot: depth must be positive, but got " << depth;
static const Op op = Op::Get("relax.one_hot");
- return Call(Type::Missing(), op, {indices, on_value, off_value}, Attrs(attrs), {});
+ return Call::Unchecked(Type::Missing(), op, {indices, on_value, off_value}, Attrs(attrs), {});
} // namespace relax
TVM_FFI_STATIC_INIT_BLOCK() {
@@ -3194,10 +3180,10 @@ Type InferTypeOneHot(const Call& call, const BlockBuilder& ctx) {
TVM_FFI_STATIC_INIT_BLOCK() {
OpDef("relax.one_hot")
- .add_arg("indices", "The indices tensor.")
- .add_arg("on_value", "The value to fill at specified indices.")
- .add_arg("off_value", "The value to fill at other indices.")
- .attrs_type<OneHotAttrs>()
+ .signature(sig::arg("indices", "The indices tensor."),
+ 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<bool>("FPurity", true);
}
diff --git a/src/relax/op/tensor/qdq.cc b/src/relax/op/tensor/qdq.cc
index b262a0edc0..78a6164386 100644
--- a/src/relax/op/tensor/qdq.cc
+++ b/src/relax/op/tensor/qdq.cc
@@ -44,8 +44,8 @@ Expr quantize(Expr data, Expr scale, Expr zero_point, int axis, DLDataType out_d
attrs->axis = axis;
attrs->out_dtype = out_dtype;
static const Op op = Op::Get("relax.quantize");
- return Call(Type::Missing(), op, {std::move(data), std::move(scale), std::move(zero_point)},
- Attrs(attrs));
+ return Call::Unchecked(Type::Missing(), op,
+ {std::move(data), std::move(scale), std::move(zero_point)}, Attrs(attrs));
}
TVM_FFI_STATIC_INIT_BLOCK() {
@@ -157,10 +157,10 @@ Type InferTypeQuantize(const Call& call, const BlockBuilder& ctx) {
TVM_FFI_STATIC_INIT_BLOCK() {
OpDef("relax.quantize")
- .add_arg("data", "The input tensor.")
- .add_arg("scale", "The quantization scale of the output tensor.")
- .add_arg("zero_point", "The quantization zero_point of the output tensor.")
- .attrs_type<QuantizeAttrs>()
+ .signature(sig::arg("data", "The input tensor."),
+ 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<bool>("FPurity", true);
}
@@ -172,8 +172,8 @@ Expr dequantize(Expr data, Expr scale, Expr zero_point, int axis, DLDataType out
attrs->axis = axis;
attrs->out_dtype = out_dtype;
static const Op op = Op::Get("relax.dequantize");
- return Call(Type::Missing(), op, {std::move(data), std::move(scale), std::move(zero_point)},
- Attrs(attrs));
+ return Call::Unchecked(Type::Missing(), op,
+ {std::move(data), std::move(scale), std::move(zero_point)}, Attrs(attrs));
}
TVM_FFI_STATIC_INIT_BLOCK() {
@@ -280,10 +280,10 @@ Type InferTypeDequantize(const Call& call, const BlockBuilder& ctx) {
TVM_FFI_STATIC_INIT_BLOCK() {
OpDef("relax.dequantize")
- .add_arg("data", "The input tensor.")
- .add_arg("scale", "The quantization scale of the input tensor.")
- .add_arg("zero_point", "The quantization zero_point of the input tensor.")
- .attrs_type<QuantizeAttrs>()
+ .signature(sig::arg("data", "The input tensor."),
+ 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<bool>("FPurity", true);
}
diff --git a/src/relax/op/tensor/sampling.cc b/src/relax/op/tensor/sampling.cc
index f560f42bd3..6c3fb3be06 100644
--- a/src/relax/op/tensor/sampling.cc
+++ b/src/relax/op/tensor/sampling.cc
@@ -44,9 +44,9 @@ Expr multinomial_from_uniform(Expr prob, Expr uniform_sample, Expr sample_indice
attrs->dtype = dtype;
static const Op op = Op::Get("relax.multinomial_from_uniform");
- return Call(Type::Missing(), op,
- {std::move(prob), std::move(uniform_sample), std::move(sample_indices)}, Attrs(attrs),
- {});
+ return Call::Unchecked(Type::Missing(), op,
+ {std::move(prob), std::move(uniform_sample), std::move(sample_indices)},
+ Attrs(attrs), {});
}
TVM_FFI_STATIC_INIT_BLOCK() {
@@ -146,10 +146,10 @@ Type InferTypeMultinomialFromUniform(const Call& call, const BlockBuilder& ctx)
TVM_FFI_STATIC_INIT_BLOCK() {
OpDef("relax.multinomial_from_uniform")
- .add_arg("prob", "The probability tensor.")
- .add_arg("uniform_sample", "The uniform sample tensor.")
- .add_arg("sample_indices", "The sample indices tensor.")
- .attrs_type<MultinomialFromUniformAttrs>()
+ .signature(sig::arg("prob", "The probability tensor."),
+ 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<bool>("FPurity", true);
}
diff --git a/src/relax/op/tensor/search.cc b/src/relax/op/tensor/search.cc
index 68d230c630..40b170c9d4 100644
--- a/src/relax/op/tensor/search.cc
+++ b/src/relax/op/tensor/search.cc
@@ -45,8 +45,8 @@ Expr bucketize(Expr input_tensor, Expr boundaries, bool out_int32, bool right) {
attrs->out_int32 = std::move(out_int32);
attrs->right = std::move(right);
static const Op op = Op::Get("relax.bucketize");
- return Call(Type::Missing(), op, {std::move(input_tensor), std::move(boundaries)}, Attrs(attrs),
- {});
+ return Call::Unchecked(Type::Missing(), op, {std::move(input_tensor), std::move(boundaries)},
+ Attrs(attrs), {});
}
TVM_FFI_STATIC_INIT_BLOCK() {
@@ -79,10 +79,13 @@ Type InferTypeBucketize(const Call& call, const BlockBuilder& ctx) {
TVM_FFI_STATIC_INIT_BLOCK() {
OpDef("relax.bucketize")
- .add_arg("input_tensor", " N-D tensor or a Scalar containing the search value(s).")
- .add_arg("boundaries",
- "1-D tensor, must contain a strictly increasing sequence, or the return value is "
- "undefined.")
+ .signature(
+ sig::arg("input_tensor", " N-D tensor or a Scalar containing the search value(s)."),
+ sig::arg(
+ "boundaries",
+ "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<bool>("FPurity", true);
}
@@ -90,8 +93,8 @@ TVM_FFI_STATIC_INIT_BLOCK() {
/* relax.where */
Expr where(Expr condition, Expr x1, Expr x2) {
static const Op op = Op::Get("relax.where");
- return Call(Type::Missing(), op, {std::move(condition), std::move(x1), std::move(x2)}, Attrs(),
- {});
+ return Call::Unchecked(Type::Missing(), op, {std::move(condition), std::move(x1), std::move(x2)},
+ Attrs(), {});
}
TVM_FFI_STATIC_INIT_BLOCK() {
@@ -182,9 +185,9 @@ Type InferTypeWhere(const Call& call, const BlockBuilder& ctx) {
TVM_FFI_STATIC_INIT_BLOCK() {
OpDef("relax.where")
- .add_arg("condition", "When True, yield `x1`; otherwise, yield `x2`.")
- .add_arg("x1", "The first input tensor.")
- .add_arg("x2", "The second input tensor.")
+ .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<bool>("FPurity", true);
}
@@ -254,7 +257,7 @@ Expr argmax(Expr x, ffi::Optional<int64_t> axis, bool keepdims) {
attrs->axis = std::move(axis);
attrs->keepdims = std::move(keepdims);
static const Op op = Op::Get("relax.argmax");
- return Call(Type::Missing(), op, {std::move(x)}, Attrs(attrs));
+ return Call::Unchecked(Type::Missing(), op, {std::move(x)}, Attrs(attrs));
}
Expr argmin(Expr x, ffi::Optional<int64_t> axis, bool keepdims) {
@@ -262,14 +265,14 @@ Expr argmin(Expr x, ffi::Optional<int64_t> axis, bool keepdims) {
attrs->axis = std::move(axis);
attrs->keepdims = std::move(keepdims);
static const Op op = Op::Get("relax.argmin");
- return Call(Type::Missing(), op, {std::move(x)}, Attrs(attrs));
+ return Call::Unchecked(Type::Missing(), op, {std::move(x)}, Attrs(attrs));
}
TVM_FFI_STATIC_INIT_BLOCK() {
tvm::ffi::reflection::GlobalDef().def("relax.op.argmax", argmax);
OpDef("relax.argmax")
- .add_arg("x", "The input data tensor")
+ .signature(sig::arg("x", "The input data tensor"), sig::call_attrs<ArgmaxArgminAttrs>())
.set_attr<FInferType>("FInferType", InferTypeArgmaxArgmin)
.set_attr<bool>("FPurity", true);
};
@@ -278,7 +281,7 @@ TVM_FFI_STATIC_INIT_BLOCK() {
tvm::ffi::reflection::GlobalDef().def("relax.op.argmin", argmin);
OpDef("relax.argmin")
- .add_arg("x", "The input data tensor")
+ .signature(sig::arg("x", "The input data tensor"), sig::call_attrs<ArgmaxArgminAttrs>())
.set_attr<FInferType>("FInferType", InferTypeArgmaxArgmin)
.set_attr<bool>("FPurity", true);
};
diff --git a/src/relax/op/tensor/set.cc b/src/relax/op/tensor/set.cc
index d89bfc7d2d..b711d31558 100644
--- a/src/relax/op/tensor/set.cc
+++ b/src/relax/op/tensor/set.cc
@@ -41,12 +41,13 @@ Expr unique(Expr x, PrimExpr sorted, PrimExpr return_index, PrimExpr return_inve
static const Op op = Op::Get("relax.unique");
Call call;
if (!axis) {
- call = Call(Type::Missing(), op,
- {std::move(x), sorted, return_index, return_inverse, return_counts});
+ call = Call::Unchecked(Type::Missing(), op,
+ {std::move(x), sorted, return_index, return_inverse, return_counts});
} else {
PrimExpr pv_axis = axis.value();
- call = Call(Type::Missing(), op,
- {std::move(x), sorted, return_index, return_inverse, return_counts, pv_axis});
+ call = Call::Unchecked(
+ Type::Missing(), op,
+ {std::move(x), sorted, return_index, return_inverse, return_counts, pv_axis});
}
return call;
}
@@ -144,19 +145,21 @@ TVM_FFI_STATIC_INIT_BLOCK() {
OpDef("relax.unique",
"Optional axis: The dimension to apply unique. If it is std::nullopt, the unique values of "
"the flattened input are returned.")
- .add_arg("x", "The input tensor")
- .add_arg("sorted",
- "Whether to sort the unique elements in ascending order before returning as output.")
- .add_arg(
- "return_index",
- "Whether to return an additional tensor with indices for where elements in the unique "
- "tensor come from the original input.")
- .add_arg("return_inverse",
- "Whether to return an additional tensor with indices for where elements in the "
- "original input ended up in the returned unique list.")
- .add_arg("return_counts",
- "Whether to return an additional tensor with counts of each unique elements")
- .allow_extra_args()
+ .signature(
+ sig::arg("x", "The input tensor"),
+ sig::arg(
+ "sorted",
+ "Whether to sort the unique elements in ascending order before returning as output."),
+ sig::arg("return_index",
+ "Whether to return an additional tensor with indices for where elements in the "
+ "unique "
+ "tensor come from the original input."),
+ sig::arg("return_inverse",
+ "Whether to return an additional tensor with indices for where elements in the "
+ "original input ended up in the returned unique list."),
+ 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<FCallPacked>("FCallPacked", "relax.run.unique")
.set_attr<bool>("FPurity", true);
@@ -165,7 +168,7 @@ TVM_FFI_STATIC_INIT_BLOCK() {
/* relax.nonzero */
Expr nonzero(Expr x) {
static const Op op = Op::Get("relax.nonzero");
- return Call(Type::Missing(), op, {std::move(x)});
+ return Call::Unchecked(Type::Missing(), op, {std::move(x)});
}
TVM_FFI_STATIC_INIT_BLOCK() {
@@ -180,7 +183,7 @@ Type InferTypeNonzero(const Call& call, const BlockBuilder& ctx) {
TVM_FFI_STATIC_INIT_BLOCK() {
OpDef("relax.nonzero")
- .add_arg("x", "The input tensor")
+ .signature(sig::arg("x", "The input tensor"))
.set_attr<FInferType>("FInferType", 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 43bf7ae78f..840315d2fb 100644
--- a/src/relax/op/tensor/sorting.cc
+++ b/src/relax/op/tensor/sorting.cc
@@ -45,7 +45,7 @@ Expr sort(Expr data, int axis, bool descending) {
attrs->descending = std::move(descending);
static const Op op = Op::Get("relax.sort");
- return Call(Type::Missing(), op, {std::move(data)}, Attrs{attrs}, {});
+ return Call::Unchecked(Type::Missing(), op, {std::move(data)}, Attrs{attrs}, {});
}
TVM_FFI_STATIC_INIT_BLOCK() {
@@ -59,8 +59,7 @@ Type InferTypeSort(const Call& call, const BlockBuilder& ctx) {
TVM_FFI_STATIC_INIT_BLOCK() {
OpDef("relax.sort")
- .add_arg("data", "The input tensor.")
- .attrs_type<SortAttrs>()
+ .signature(sig::arg("data", "The input tensor."), sig::call_attrs<SortAttrs>())
.set_attr<FInferType>("FInferType", InferTypeSort)
.set_attr<bool>("FPurity", true);
}
@@ -74,7 +73,7 @@ Expr argsort(Expr data, int axis, bool descending, ffi::Optional<DLDataType> dty
attrs->dtype = std::move(dtype);
static const Op op = Op::Get("relax.argsort");
- return Call(Type::Missing(), op, {std::move(data)}, Attrs{attrs}, {});
+ return Call::Unchecked(Type::Missing(), op, {std::move(data)}, Attrs{attrs}, {});
}
TVM_FFI_STATIC_INIT_BLOCK() {
@@ -96,8 +95,7 @@ Type InferTypeArgsort(const Call& call, const BlockBuilder& ctx) {
TVM_FFI_STATIC_INIT_BLOCK() {
OpDef("relax.argsort")
- .add_arg("data", "The input tensor.")
- .attrs_type<ArgsortAttrs>()
+ .signature(sig::arg("data", "The input tensor."), sig::call_attrs<ArgsortAttrs>())
.set_attr<FInferType>("FInferType", InferTypeArgsort)
.set_attr<bool>("FPurity", true);
}
@@ -114,7 +112,7 @@ Expr topk(Expr data, int k, int axis, ffi::String ret_type, bool largest,
attrs->dtype = std::move(dtype);
static const Op op = Op::Get("relax.topk");
- return Call(Type::Missing(), op, {std::move(data)}, Attrs{attrs}, {});
+ return Call::Unchecked(Type::Missing(), op, {std::move(data)}, Attrs{attrs}, {});
}
TVM_FFI_STATIC_INIT_BLOCK() {
@@ -165,8 +163,7 @@ Type InferTypeTopK(const Call& call, const BlockBuilder& ctx) {
TVM_FFI_STATIC_INIT_BLOCK() {
OpDef("relax.topk")
- .add_arg("data", "The input tensor.")
- .attrs_type<TopKAttrs>()
+ .signature(sig::arg("data", "The input tensor."), sig::call_attrs<TopKAttrs>())
.set_attr<FInferType>("FInferType", InferTypeTopK)
.set_attr<bool>("FPurity", true);
}
diff --git a/src/relax/op/tensor/statistical.cc b/src/relax/op/tensor/statistical.cc
index 3493357144..ec70316fef 100644
--- a/src/relax/op/tensor/statistical.cc
+++ b/src/relax/op/tensor/statistical.cc
@@ -248,7 +248,7 @@ Expr cumprod(Expr data, ffi::Optional<int64_t> axis, ffi::Optional<DLDataType> d
attrs->exclusive = exclusive;
static const Op op = Op::Get("relax.cumprod");
- return Call(Type::Missing(), op, {std::move(data)}, Attrs{attrs}, {});
+ return Call::Unchecked(Type::Missing(), op, {std::move(data)}, Attrs{attrs}, {});
}
TVM_FFI_STATIC_INIT_BLOCK() {
@@ -256,8 +256,7 @@ TVM_FFI_STATIC_INIT_BLOCK() {
refl::GlobalDef().def("relax.op.cumprod", cumprod);
OpDef("relax.cumprod")
- .add_arg("data", "The input tensor.")
- .attrs_type<ScanopAttrs>()
+ .signature(sig::arg("data", "The input tensor."), sig::call_attrs<ScanopAttrs>())
.set_attr<FInferType>("FInferType", InferTypeScan)
.set_attr<bool>("FPurity", true);
}
@@ -271,7 +270,7 @@ Expr cumsum(Expr data, ffi::Optional<int64_t> axis, ffi::Optional<DLDataType> dt
attrs->exclusive = exclusive;
static const Op op = Op::Get("relax.cumsum");
- return Call(Type::Missing(), op, {std::move(data)}, Attrs{attrs}, {});
+ return Call::Unchecked(Type::Missing(), op, {std::move(data)}, Attrs{attrs}, {});
}
TVM_FFI_STATIC_INIT_BLOCK() {
@@ -279,8 +278,7 @@ TVM_FFI_STATIC_INIT_BLOCK() {
refl::GlobalDef().def("relax.op.cumsum", cumsum);
OpDef("relax.cumsum")
- .add_arg("data", "The input tensor.")
- .attrs_type<ScanopAttrs>()
+ .signature(sig::arg("data", "The input tensor."), sig::call_attrs<ScanopAttrs>())
.set_attr<FInferType>("FInferType", InferTypeScan)
.set_attr<bool>("FPurity", true);
}
@@ -291,7 +289,7 @@ Expr median(Expr data, ffi::Optional<ffi::Array<int64_t>> axis, bool keepdims) {
attrs->axis = std::move(axis);
attrs->keepdims = keepdims;
static const Op op = Op::Get("relax.median");
- return Call(Type::Missing(), op, {std::move(data)}, Attrs{attrs}, {});
+ return Call::Unchecked(Type::Missing(), op, {std::move(data)}, Attrs{attrs}, {});
}
TVM_FFI_STATIC_INIT_BLOCK() {
@@ -299,7 +297,7 @@ TVM_FFI_STATIC_INIT_BLOCK() {
refl::GlobalDef().def("relax.op.median", median);
OpDef("relax.median")
- .add_arg("data", "The input tensor.")
+ .signature(sig::arg("data", "The input tensor."), sig::call_attrs<StatisticalAttrs>())
.set_attr<FInferType>("FInferType", InferTypeStatisticalExtension)
.set_attr<bool>("FPurity", true);
}
@@ -309,7 +307,7 @@ Expr max(Expr x, ffi::Optional<ffi::Array<int64_t>> axis, bool keepdims) {
attrs->axis = std::move(axis);
attrs->keepdims = keepdims;
static const Op op = Op::Get("relax.max");
- return Call(Type::Missing(), op, {std::move(x)}, Attrs{attrs}, {});
+ return Call::Unchecked(Type::Missing(), op, {std::move(x)}, Attrs{attrs}, {});
}
Expr mean(Expr x, ffi::Optional<ffi::Array<int64_t>> axis, bool keepdims) {
@@ -317,7 +315,7 @@ Expr mean(Expr x, ffi::Optional<ffi::Array<int64_t>> axis, bool keepdims) {
attrs->axis = std::move(axis);
attrs->keepdims = keepdims;
static const Op op = Op::Get("relax.mean");
- return Call(Type::Missing(), op, {std::move(x)}, Attrs{attrs}, {});
+ return Call::Unchecked(Type::Missing(), op, {std::move(x)}, Attrs{attrs}, {});
}
Expr min(Expr x, ffi::Optional<ffi::Array<int64_t>> axis, bool keepdims) {
@@ -325,7 +323,7 @@ Expr min(Expr x, ffi::Optional<ffi::Array<int64_t>> axis, bool keepdims) {
attrs->axis = std::move(axis);
attrs->keepdims = keepdims;
static const Op op = Op::Get("relax.min");
- return Call(Type::Missing(), op, {std::move(x)}, Attrs{attrs}, {});
+ return Call::Unchecked(Type::Missing(), op, {std::move(x)}, Attrs{attrs}, {});
}
Expr prod(Expr x, ffi::Optional<ffi::Array<int64_t>> axis, bool keepdims) {
@@ -333,7 +331,7 @@ Expr prod(Expr x, ffi::Optional<ffi::Array<int64_t>> axis, bool keepdims) {
attrs->axis = std::move(axis);
attrs->keepdims = keepdims;
static const Op op = Op::Get("relax.prod");
- return Call(Type::Missing(), op, {std::move(x)}, Attrs{attrs}, {});
+ return Call::Unchecked(Type::Missing(), op, {std::move(x)}, Attrs{attrs}, {});
}
Expr std(Expr x, ffi::Optional<ffi::Array<int64_t>> axis, bool keepdims) {
@@ -341,7 +339,7 @@ Expr std(Expr x, ffi::Optional<ffi::Array<int64_t>> axis, bool keepdims) {
attrs->axis = std::move(axis);
attrs->keepdims = keepdims;
static const Op op = Op::Get("relax.std");
- return Call(Type::Missing(), op, {std::move(x)}, Attrs{attrs}, {});
+ return Call::Unchecked(Type::Missing(), op, {std::move(x)}, Attrs{attrs}, {});
}
Expr sum(Expr x, ffi::Optional<ffi::Array<int64_t>> axis, bool keepdims) {
@@ -349,7 +347,7 @@ Expr sum(Expr x, ffi::Optional<ffi::Array<int64_t>> axis, bool keepdims) {
attrs->axis = std::move(axis);
attrs->keepdims = keepdims;
static const Op op = Op::Get("relax.sum");
- return Call(Type::Missing(), op, {std::move(x)}, Attrs{attrs}, {});
+ return Call::Unchecked(Type::Missing(), op, {std::move(x)}, Attrs{attrs}, {});
}
Expr variance(Expr x, ffi::Optional<ffi::Array<int64_t>> axis, bool keepdims) {
@@ -357,14 +355,14 @@ Expr variance(Expr x, ffi::Optional<ffi::Array<int64_t>> axis, bool keepdims) {
attrs->axis = std::move(axis);
attrs->keepdims = keepdims;
static const Op op = Op::Get("relax.variance");
- return Call(Type::Missing(), op, {std::move(x)}, Attrs{attrs}, {});
+ return Call::Unchecked(Type::Missing(), op, {std::move(x)}, Attrs{attrs}, {});
}
TVM_FFI_STATIC_INIT_BLOCK() {
tvm::ffi::reflection::GlobalDef().def("relax.op.max", max);
OpDef("relax.max")
- .add_arg("x", "The input data tensor")
+ .signature(sig::arg("x", "The input data tensor"), sig::call_attrs<StatisticalAttrs>())
.set_attr<FInferType>("FInferType", InferTypeStatistical)
.set_attr<FRelaxInferLayout>("FRelaxInferLayout", InferLayoutStatistical)
.set_attr<bool>("FPurity", true);
@@ -372,7 +370,7 @@ TVM_FFI_STATIC_INIT_BLOCK() {
tvm::ffi::reflection::GlobalDef().def("relax.op.mean", mean);
OpDef("relax.mean")
- .add_arg("x", "The input data tensor")
+ .signature(sig::arg("x", "The input data tensor"), sig::call_attrs<StatisticalAttrs>())
.set_attr<FInferType>("FInferType", InferTypeStatistical)
.set_attr<FRelaxInferLayout>("FRelaxInferLayout", InferLayoutStatistical)
.set_attr<bool>("FPurity", true);
@@ -380,7 +378,7 @@ TVM_FFI_STATIC_INIT_BLOCK() {
tvm::ffi::reflection::GlobalDef().def("relax.op.min", min);
OpDef("relax.min")
- .add_arg("x", "The input data tensor")
+ .signature(sig::arg("x", "The input data tensor"), sig::call_attrs<StatisticalAttrs>())
.set_attr<FInferType>("FInferType", InferTypeStatistical)
.set_attr<FRelaxInferLayout>("FRelaxInferLayout", InferLayoutStatistical)
.set_attr<bool>("FPurity", true);
@@ -388,7 +386,7 @@ TVM_FFI_STATIC_INIT_BLOCK() {
tvm::ffi::reflection::GlobalDef().def("relax.op.prod", prod);
OpDef("relax.prod")
- .add_arg("x", "The input data tensor")
+ .signature(sig::arg("x", "The input data tensor"), sig::call_attrs<StatisticalAttrs>())
.set_attr<FInferType>("FInferType", InferTypeStatistical)
.set_attr<FRelaxInferLayout>("FRelaxInferLayout", InferLayoutStatistical)
.set_attr<bool>("FPurity", true);
@@ -396,7 +394,7 @@ TVM_FFI_STATIC_INIT_BLOCK() {
tvm::ffi::reflection::GlobalDef().def("relax.op.std", std);
OpDef("relax.std")
- .add_arg("x", "The input data tensor")
+ .signature(sig::arg("x", "The input data tensor"), sig::call_attrs<StatisticalAttrs>())
.set_attr<FInferType>("FInferType", InferTypeStatistical)
.set_attr<FRelaxInferLayout>("FRelaxInferLayout", InferLayoutStatistical)
.set_attr<bool>("FPurity", true);
@@ -404,7 +402,7 @@ TVM_FFI_STATIC_INIT_BLOCK() {
tvm::ffi::reflection::GlobalDef().def("relax.op.sum", sum);
OpDef("relax.sum")
- .add_arg("x", "The input data tensor")
+ .signature(sig::arg("x", "The input data tensor"), sig::call_attrs<StatisticalAttrs>())
.set_attr<FInferType>("FInferType", InferTypeStatistical)
.set_attr<FRelaxInferLayout>("FRelaxInferLayout", InferLayoutStatistical)
.set_attr<bool>("FPurity", true);
@@ -412,7 +410,7 @@ TVM_FFI_STATIC_INIT_BLOCK() {
tvm::ffi::reflection::GlobalDef().def("relax.op.variance", variance);
OpDef("relax.variance")
- .add_arg("x", "The input data tensor")
+ .signature(sig::arg("x", "The input data tensor"), sig::call_attrs<StatisticalAttrs>())
.set_attr<FInferType>("FInferType", 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 d0e3cada9d..eb8250d9fa 100644
--- a/src/relax/op/tensor/ternary.cc
+++ b/src/relax/op/tensor/ternary.cc
@@ -133,9 +133,11 @@ InferLayoutOutput InferLayoutEwiseFMA(
TVM_FFI_STATIC_INIT_BLOCK() {
OpDef("relax.ewise_fma")
- .add_arg("x1", "The left hand operand of the multiplication")
- .add_arg("x2", "The right hand operand of the multiplication")
- .add_arg("x3", "The operand of the addition")
+ .signature(
+ sig::arg("x1", "The left hand operand of the multiplication"),
+ 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<FRelaxInferLayout>("FRelaxInferLayout", InferLayoutEwiseFMA)
.set_attr<TMixedPrecisionPolicy>("TMixedPrecisionPolicy", MixedPrecisionPolicyKind::kFollow)
@@ -144,7 +146,7 @@ TVM_FFI_STATIC_INIT_BLOCK() {
Expr ewise_fma(Expr x1, Expr x2, Expr x3) {
static const Op op = Op::Get("relax.ewise_fma");
- return Call(Type::Missing(), op, {x1, x2, x3}, Attrs(), {});
+ return Call::Unchecked(Type::Missing(), op, {x1, x2, x3}, Attrs(), {});
}
TVM_FFI_STATIC_INIT_BLOCK() {
diff --git a/src/relax/op/tensor/unary.cc b/src/relax/op/tensor/unary.cc
index 94dc6cd649..f67acb07c3 100644
--- a/src/relax/op/tensor/unary.cc
+++ b/src/relax/op/tensor/unary.cc
@@ -40,165 +40,166 @@ Type InferTypeUnaryCheck(const Call& call, const BlockBuilder& ctx) {
Expr abs(Expr x) {
static const Op op = Op::Get("relax.abs");
- return Call(Type::Missing(), op, {std::move(x)}, Attrs(), {});
+ return Call::Unchecked(Type::Missing(), op, {std::move(x)}, Attrs(), {});
}
Expr acos(Expr x) {
static const Op op = Op::Get("relax.acos");
- return Call(Type::Missing(), op, {std::move(x)}, Attrs(), {});
+ return Call::Unchecked(Type::Missing(), op, {std::move(x)}, Attrs(), {});
}
Expr acosh(Expr x) {
static const Op op = Op::Get("relax.acosh");
- return Call(Type::Missing(), op, {std::move(x)}, Attrs(), {});
+ return Call::Unchecked(Type::Missing(), op, {std::move(x)}, Attrs(), {});
}
Expr asin(Expr x) {
static const Op op = Op::Get("relax.asin");
- return Call(Type::Missing(), op, {std::move(x)}, Attrs(), {});
+ return Call::Unchecked(Type::Missing(), op, {std::move(x)}, Attrs(), {});
}
Expr asinh(Expr x) {
static const Op op = Op::Get("relax.asinh");
- return Call(Type::Missing(), op, {std::move(x)}, Attrs(), {});
+ return Call::Unchecked(Type::Missing(), op, {std::move(x)}, Attrs(), {});
}
Expr atan(Expr x) {
static const Op op = Op::Get("relax.atan");
- return Call(Type::Missing(), op, {std::move(x)}, Attrs(), {});
+ return Call::Unchecked(Type::Missing(), op, {std::move(x)}, Attrs(), {});
}
Expr atanh(Expr x) {
static const Op op = Op::Get("relax.atanh");
- return Call(Type::Missing(), op, {std::move(x)}, Attrs(), {});
+ return Call::Unchecked(Type::Missing(), op, {std::move(x)}, Attrs(), {});
}
Expr bitwise_not(Expr x) {
static const Op op = Op::Get("relax.bitwise_not");
- return Call(Type::Missing(), op, {std::move(x)}, Attrs(), {});
+ return Call::Unchecked(Type::Missing(), op, {std::move(x)}, Attrs(), {});
}
Expr ceil(Expr x) {
static const Op op = Op::Get("relax.ceil");
- return Call(Type::Missing(), op, {std::move(x)}, Attrs(), {});
+ return Call::Unchecked(Type::Missing(), op, {std::move(x)}, Attrs(), {});
}
Expr cos(Expr x) {
static const Op op = Op::Get("relax.cos");
- return Call(Type::Missing(), op, {std::move(x)}, Attrs(), {});
+ return Call::Unchecked(Type::Missing(), op, {std::move(x)}, Attrs(), {});
}
Expr cosh(Expr x) {
static const Op op = Op::Get("relax.cosh");
- return Call(Type::Missing(), op, {std::move(x)}, Attrs(), {});
+ return Call::Unchecked(Type::Missing(), op, {std::move(x)}, Attrs(), {});
}
Expr exp(Expr x) {
static const Op op = Op::Get("relax.exp");
- return Call(Type::Missing(), op, {std::move(x)}, Attrs(), {});
+ return Call::Unchecked(Type::Missing(), op, {std::move(x)}, Attrs(), {});
}
Expr floor(Expr x) {
static const Op op = Op::Get("relax.floor");
- return Call(Type::Missing(), op, {std::move(x)}, Attrs(), {});
+ return Call::Unchecked(Type::Missing(), op, {std::move(x)}, Attrs(), {});
}
Expr log(Expr x) {
static const Op op = Op::Get("relax.log");
- return Call(Type::Missing(), op, {std::move(x)}, Attrs(), {});
+ return Call::Unchecked(Type::Missing(), op, {std::move(x)}, Attrs(), {});
}
Expr logical_not(Expr x) {
static const Op op = Op::Get("relax.logical_not");
- return Call(Type::Missing(), op, {std::move(x)}, Attrs(), {});
+ return Call::Unchecked(Type::Missing(), op, {std::move(x)}, Attrs(), {});
}
Expr negative(Expr x) {
static const Op op = Op::Get("relax.negative");
- return Call(Type::Missing(), op, {std::move(x)}, Attrs(), {});
+ return Call::Unchecked(Type::Missing(), op, {std::move(x)}, Attrs(), {});
}
Expr round(Expr x) {
static const Op op = Op::Get("relax.round");
- return Call(Type::Missing(), op, {std::move(x)}, Attrs(), {});
+ return Call::Unchecked(Type::Missing(), op, {std::move(x)}, Attrs(), {});
}
Expr rsqrt(Expr x) {
static const Op op = Op::Get("relax.rsqrt");
- return Call(Type::Missing(), op, {std::move(x)}, Attrs(), {});
+ return Call::Unchecked(Type::Missing(), op, {std::move(x)}, Attrs(), {});
}
Expr sigmoid(Expr x) {
static const Op op = Op::Get("relax.sigmoid");
- return Call(Type::Missing(), op, {std::move(x)}, Attrs(), {});
+ return Call::Unchecked(Type::Missing(), op, {std::move(x)}, Attrs(), {});
}
Expr sign(Expr x) {
static const Op op = Op::Get("relax.sign");
- return Call(Type::Missing(), op, {std::move(x)}, Attrs(), {});
+ return Call::Unchecked(Type::Missing(), op, {std::move(x)}, Attrs(), {});
}
Expr sin(Expr x) {
static const Op op = Op::Get("relax.sin");
- return Call(Type::Missing(), op, {std::move(x)}, Attrs(), {});
+ return Call::Unchecked(Type::Missing(), op, {std::move(x)}, Attrs(), {});
}
Expr sinh(Expr x) {
static const Op op = Op::Get("relax.sinh");
- return Call(Type::Missing(), op, {std::move(x)}, Attrs(), {});
+ return Call::Unchecked(Type::Missing(), op, {std::move(x)}, Attrs(), {});
}
Expr square(Expr x) {
static const Op op = Op::Get("relax.square");
- return Call(Type::Missing(), op, {std::move(x)}, Attrs(), {});
+ return Call::Unchecked(Type::Missing(), op, {std::move(x)}, Attrs(), {});
}
Expr sqrt(Expr x) {
static const Op op = Op::Get("relax.sqrt");
- return Call(Type::Missing(), op, {std::move(x)}, Attrs(), {});
+ return Call::Unchecked(Type::Missing(), op, {std::move(x)}, Attrs(), {});
}
Expr tan(Expr x) {
static const Op op = Op::Get("relax.tan");
- return Call(Type::Missing(), op, {std::move(x)}, Attrs(), {});
+ return Call::Unchecked(Type::Missing(), op, {std::move(x)}, Attrs(), {});
}
Expr tanh(Expr x) {
static const Op op = Op::Get("relax.tanh");
- return Call(Type::Missing(), op, {std::move(x)}, Attrs(), {});
+ return Call::Unchecked(Type::Missing(), op, {std::move(x)}, Attrs(), {});
}
Expr trunc(Expr x) {
static const Op op = Op::Get("relax.trunc");
- return Call(Type::Missing(), op, {std::move(x)}, Attrs(), {});
+ return Call::Unchecked(Type::Missing(), op, {std::move(x)}, Attrs(), {});
}
Expr erf(Expr x) {
static const Op op = Op::Get("relax.erf");
- return Call(Type::Missing(), op, {std::move(x)}, Attrs(), {});
+ return Call::Unchecked(Type::Missing(), op, {std::move(x)}, Attrs(), {});
}
Expr isfinite(Expr x) {
static const Op op = Op::Get("relax.isfinite");
- return Call(Type::Missing(), op, {std::move(x)}, Attrs(), {});
+ return Call::Unchecked(Type::Missing(), op, {std::move(x)}, Attrs(), {});
}
Expr isinf(Expr x) {
static const Op op = Op::Get("relax.isinf");
- return Call(Type::Missing(), op, {std::move(x)}, Attrs(), {});
+ return Call::Unchecked(Type::Missing(), op, {std::move(x)}, Attrs(), {});
}
Expr isnan(Expr x) {
static const Op op = Op::Get("relax.isnan");
- return Call(Type::Missing(), op, {std::move(x)}, Attrs(), {});
+ return Call::Unchecked(Type::Missing(), op, {std::move(x)}, Attrs(), {});
}
TVM_FFI_STATIC_INIT_BLOCK() {
tvm::ffi::reflection::GlobalDef().def("relax.op.abs", abs);
OpDef("relax.abs")
- .add_arg("x", "The input tensor.")
- .add_ty_arg("out_type", "Optional output tensor type carrying the virtual device.")
+ .signature(
+ sig::arg("x", "The input tensor."),
+ sig::var_ty_args("out_type", "Optional output tensor type carrying the virtual device."))
.set_attr<FRelaxInferLayout>("FRelaxInferLayout", InferLayoutUnaryEwise)
.set_attr<TMixedPrecisionPolicy>("TMixedPrecisionPolicy", MixedPrecisionPolicyKind::kFollow)
.set_attr<bool>("FPurity", true)
@@ -207,8 +208,9 @@ TVM_FFI_STATIC_INIT_BLOCK() {
tvm::ffi::reflection::GlobalDef().def("relax.op.acos", acos);
OpDef("relax.acos")
- .add_arg("x", "The input tensor.")
- .add_ty_arg("out_type", "Optional output tensor type carrying the virtual device.")
+ .signature(
+ sig::arg("x", "The input tensor."),
+ sig::var_ty_args("out_type", "Optional output tensor type carrying the virtual device."))
.set_attr<FRelaxInferLayout>("FRelaxInferLayout", InferLayoutUnaryEwise)
.set_attr<TMixedPrecisionPolicy>("TMixedPrecisionPolicy", MixedPrecisionPolicyKind::kFollow)
.set_attr<bool>("FPurity", true)
@@ -217,8 +219,9 @@ TVM_FFI_STATIC_INIT_BLOCK() {
tvm::ffi::reflection::GlobalDef().def("relax.op.acosh", acosh);
OpDef("relax.acosh")
- .add_arg("x", "The input tensor.")
- .add_ty_arg("out_type", "Optional output tensor type carrying the virtual device.")
+ .signature(
+ sig::arg("x", "The input tensor."),
+ sig::var_ty_args("out_type", "Optional output tensor type carrying the virtual device."))
.set_attr<FRelaxInferLayout>("FRelaxInferLayout", InferLayoutUnaryEwise)
.set_attr<TMixedPrecisionPolicy>("TMixedPrecisionPolicy", MixedPrecisionPolicyKind::kFollow)
.set_attr<bool>("FPurity", true)
@@ -227,8 +230,9 @@ TVM_FFI_STATIC_INIT_BLOCK() {
tvm::ffi::reflection::GlobalDef().def("relax.op.asin", asin);
OpDef("relax.asin")
- .add_arg("x", "The input tensor.")
- .add_ty_arg("out_type", "Optional output tensor type carrying the virtual device.")
+ .signature(
+ sig::arg("x", "The input tensor."),
+ sig::var_ty_args("out_type", "Optional output tensor type carrying the virtual device."))
.set_attr<FRelaxInferLayout>("FRelaxInferLayout", InferLayoutUnaryEwise)
.set_attr<TMixedPrecisionPolicy>("TMixedPrecisionPolicy", MixedPrecisionPolicyKind::kFollow)
.set_attr<bool>("FPurity", true)
@@ -237,8 +241,9 @@ TVM_FFI_STATIC_INIT_BLOCK() {
tvm::ffi::reflection::GlobalDef().def("relax.op.asinh", asinh);
OpDef("relax.asinh")
- .add_arg("x", "The input tensor.")
- .add_ty_arg("out_type", "Optional output tensor type carrying the virtual device.")
+ .signature(
+ sig::arg("x", "The input tensor."),
+ sig::var_ty_args("out_type", "Optional output tensor type carrying the virtual device."))
.set_attr<FRelaxInferLayout>("FRelaxInferLayout", InferLayoutUnaryEwise)
.set_attr<TMixedPrecisionPolicy>("TMixedPrecisionPolicy", MixedPrecisionPolicyKind::kFollow)
.set_attr<bool>("FPurity", true)
@@ -247,8 +252,9 @@ TVM_FFI_STATIC_INIT_BLOCK() {
tvm::ffi::reflection::GlobalDef().def("relax.op.atan", atan);
OpDef("relax.atan")
- .add_arg("x", "The input tensor.")
- .add_ty_arg("out_type", "Optional output tensor type carrying the virtual device.")
+ .signature(
+ sig::arg("x", "The input tensor."),
+ sig::var_ty_args("out_type", "Optional output tensor type carrying the virtual device."))
.set_attr<FRelaxInferLayout>("FRelaxInferLayout", InferLayoutUnaryEwise)
.set_attr<TMixedPrecisionPolicy>("TMixedPrecisionPolicy", MixedPrecisionPolicyKind::kFollow)
.set_attr<bool>("FPurity", true)
@@ -257,8 +263,9 @@ TVM_FFI_STATIC_INIT_BLOCK() {
tvm::ffi::reflection::GlobalDef().def("relax.op.atanh", atanh);
OpDef("relax.atanh")
- .add_arg("x", "The input tensor.")
- .add_ty_arg("out_type", "Optional output tensor type carrying the virtual device.")
+ .signature(
+ sig::arg("x", "The input tensor."),
+ sig::var_ty_args("out_type", "Optional output tensor type carrying the virtual device."))
.set_attr<FRelaxInferLayout>("FRelaxInferLayout", InferLayoutUnaryEwise)
.set_attr<TMixedPrecisionPolicy>("TMixedPrecisionPolicy", MixedPrecisionPolicyKind::kFollow)
.set_attr<bool>("FPurity", true)
@@ -267,8 +274,9 @@ TVM_FFI_STATIC_INIT_BLOCK() {
tvm::ffi::reflection::GlobalDef().def("relax.op.bitwise_not", bitwise_not);
OpDef("relax.bitwise_not")
- .add_arg("x", "The input tensor.")
- .add_ty_arg("out_type", "Optional output tensor type carrying the virtual device.")
+ .signature(
+ sig::arg("x", "The input tensor."),
+ sig::var_ty_args("out_type", "Optional output tensor type carrying the virtual device."))
.set_attr<FRelaxInferLayout>("FRelaxInferLayout", InferLayoutUnaryEwise)
.set_attr<TMixedPrecisionPolicy>("TMixedPrecisionPolicy", MixedPrecisionPolicyKind::kFollow)
.set_attr<bool>("FPurity", true)
@@ -277,8 +285,9 @@ TVM_FFI_STATIC_INIT_BLOCK() {
tvm::ffi::reflection::GlobalDef().def("relax.op.ceil", ceil);
OpDef("relax.ceil")
- .add_arg("x", "The input tensor.")
- .add_ty_arg("out_type", "Optional output tensor type carrying the virtual device.")
+ .signature(
+ sig::arg("x", "The input tensor."),
+ sig::var_ty_args("out_type", "Optional output tensor type carrying the virtual device."))
.set_attr<FRelaxInferLayout>("FRelaxInferLayout", InferLayoutUnaryEwise)
.set_attr<TMixedPrecisionPolicy>("TMixedPrecisionPolicy", MixedPrecisionPolicyKind::kFollow)
.set_attr<bool>("FPurity", true)
@@ -287,8 +296,9 @@ TVM_FFI_STATIC_INIT_BLOCK() {
tvm::ffi::reflection::GlobalDef().def("relax.op.cos", cos);
OpDef("relax.cos")
- .add_arg("x", "The input tensor.")
- .add_ty_arg("out_type", "Optional output tensor type carrying the virtual device.")
+ .signature(
+ sig::arg("x", "The input tensor."),
+ sig::var_ty_args("out_type", "Optional output tensor type carrying the virtual device."))
.set_attr<FRelaxInferLayout>("FRelaxInferLayout", InferLayoutUnaryEwise)
.set_attr<TMixedPrecisionPolicy>("TMixedPrecisionPolicy", MixedPrecisionPolicyKind::kFollow)
.set_attr<bool>("FPurity", true)
@@ -297,8 +307,9 @@ TVM_FFI_STATIC_INIT_BLOCK() {
tvm::ffi::reflection::GlobalDef().def("relax.op.cosh", cosh);
OpDef("relax.cosh")
- .add_arg("x", "The input tensor.")
- .add_ty_arg("out_type", "Optional output tensor type carrying the virtual device.")
+ .signature(
+ sig::arg("x", "The input tensor."),
+ sig::var_ty_args("out_type", "Optional output tensor type carrying the virtual device."))
.set_attr<FRelaxInferLayout>("FRelaxInferLayout", InferLayoutUnaryEwise)
.set_attr<TMixedPrecisionPolicy>("TMixedPrecisionPolicy", MixedPrecisionPolicyKind::kFollow)
.set_attr<bool>("FPurity", true)
@@ -307,8 +318,9 @@ TVM_FFI_STATIC_INIT_BLOCK() {
tvm::ffi::reflection::GlobalDef().def("relax.op.exp", exp);
OpDef("relax.exp")
- .add_arg("x", "The input tensor.")
- .add_ty_arg("out_type", "Optional output tensor type carrying the virtual device.")
+ .signature(
+ sig::arg("x", "The input tensor."),
+ sig::var_ty_args("out_type", "Optional output tensor type carrying the virtual device."))
.set_attr<FRelaxInferLayout>("FRelaxInferLayout", InferLayoutUnaryEwise)
.set_attr<TMixedPrecisionPolicy>("TMixedPrecisionPolicy", MixedPrecisionPolicyKind::kFollow)
.set_attr<bool>("FPurity", true)
@@ -317,8 +329,9 @@ TVM_FFI_STATIC_INIT_BLOCK() {
tvm::ffi::reflection::GlobalDef().def("relax.op.floor", floor);
OpDef("relax.floor")
- .add_arg("x", "The input tensor.")
- .add_ty_arg("out_type", "Optional output tensor type carrying the virtual device.")
+ .signature(
+ sig::arg("x", "The input tensor."),
+ sig::var_ty_args("out_type", "Optional output tensor type carrying the virtual device."))
.set_attr<FRelaxInferLayout>("FRelaxInferLayout", InferLayoutUnaryEwise)
.set_attr<TMixedPrecisionPolicy>("TMixedPrecisionPolicy", MixedPrecisionPolicyKind::kFollow)
.set_attr<bool>("FPurity", true)
@@ -327,8 +340,9 @@ TVM_FFI_STATIC_INIT_BLOCK() {
tvm::ffi::reflection::GlobalDef().def("relax.op.log", log);
OpDef("relax.log")
- .add_arg("x", "The input tensor.")
- .add_ty_arg("out_type", "Optional output tensor type carrying the virtual device.")
+ .signature(
+ sig::arg("x", "The input tensor."),
+ sig::var_ty_args("out_type", "Optional output tensor type carrying the virtual device."))
.set_attr<FRelaxInferLayout>("FRelaxInferLayout", InferLayoutUnaryEwise)
.set_attr<TMixedPrecisionPolicy>("TMixedPrecisionPolicy", MixedPrecisionPolicyKind::kFollow)
.set_attr<bool>("FPurity", true)
@@ -337,8 +351,9 @@ TVM_FFI_STATIC_INIT_BLOCK() {
tvm::ffi::reflection::GlobalDef().def("relax.op.logical_not", logical_not);
OpDef("relax.logical_not")
- .add_arg("x", "The input tensor.")
- .add_ty_arg("out_type", "Optional output tensor type carrying the virtual device.")
+ .signature(
+ sig::arg("x", "The input tensor."),
+ sig::var_ty_args("out_type", "Optional output tensor type carrying the virtual device."))
.set_attr<FRelaxInferLayout>("FRelaxInferLayout", InferLayoutUnaryEwise)
.set_attr<TMixedPrecisionPolicy>("TMixedPrecisionPolicy", MixedPrecisionPolicyKind::kFollow)
.set_attr<bool>("FPurity", true)
@@ -347,8 +362,9 @@ TVM_FFI_STATIC_INIT_BLOCK() {
tvm::ffi::reflection::GlobalDef().def("relax.op.negative", negative);
OpDef("relax.negative")
- .add_arg("x", "The input tensor.")
- .add_ty_arg("out_type", "Optional output tensor type carrying the virtual device.")
+ .signature(
+ sig::arg("x", "The input tensor."),
+ sig::var_ty_args("out_type", "Optional output tensor type carrying the virtual device."))
.set_attr<FRelaxInferLayout>("FRelaxInferLayout", InferLayoutUnaryEwise)
.set_attr<TMixedPrecisionPolicy>("TMixedPrecisionPolicy", MixedPrecisionPolicyKind::kFollow)
.set_attr<bool>("FPurity", true)
@@ -357,8 +373,9 @@ TVM_FFI_STATIC_INIT_BLOCK() {
tvm::ffi::reflection::GlobalDef().def("relax.op.round", round);
OpDef("relax.round")
- .add_arg("x", "The input tensor.")
- .add_ty_arg("out_type", "Optional output tensor type carrying the virtual device.")
+ .signature(
+ sig::arg("x", "The input tensor."),
+ sig::var_ty_args("out_type", "Optional output tensor type carrying the virtual device."))
.set_attr<FRelaxInferLayout>("FRelaxInferLayout", InferLayoutUnaryEwise)
.set_attr<TMixedPrecisionPolicy>("TMixedPrecisionPolicy", MixedPrecisionPolicyKind::kFollow)
.set_attr<bool>("FPurity", true)
@@ -367,8 +384,9 @@ TVM_FFI_STATIC_INIT_BLOCK() {
tvm::ffi::reflection::GlobalDef().def("relax.op.rsqrt", rsqrt);
OpDef("relax.rsqrt")
- .add_arg("x", "The input tensor.")
- .add_ty_arg("out_type", "Optional output tensor type carrying the virtual device.")
+ .signature(
+ sig::arg("x", "The input tensor."),
+ sig::var_ty_args("out_type", "Optional output tensor type carrying the virtual device."))
.set_attr<FRelaxInferLayout>("FRelaxInferLayout", InferLayoutUnaryEwise)
.set_attr<TMixedPrecisionPolicy>("TMixedPrecisionPolicy", MixedPrecisionPolicyKind::kFollow)
.set_attr<bool>("FPurity", true)
@@ -377,8 +395,9 @@ TVM_FFI_STATIC_INIT_BLOCK() {
tvm::ffi::reflection::GlobalDef().def("relax.op.sigmoid", sigmoid);
OpDef("relax.sigmoid")
- .add_arg("x", "The input tensor.")
- .add_ty_arg("out_type", "Optional output tensor type carrying the virtual device.")
+ .signature(
+ sig::arg("x", "The input tensor."),
+ sig::var_ty_args("out_type", "Optional output tensor type carrying the virtual device."))
.set_attr<FRelaxInferLayout>("FRelaxInferLayout", InferLayoutUnaryEwise)
.set_attr<TMixedPrecisionPolicy>("TMixedPrecisionPolicy", MixedPrecisionPolicyKind::kFollow)
.set_attr<bool>("FPurity", true)
@@ -387,8 +406,9 @@ TVM_FFI_STATIC_INIT_BLOCK() {
tvm::ffi::reflection::GlobalDef().def("relax.op.sign", sign);
OpDef("relax.sign")
- .add_arg("x", "The input tensor.")
- .add_ty_arg("out_type", "Optional output tensor type carrying the virtual device.")
+ .signature(
+ sig::arg("x", "The input tensor."),
+ sig::var_ty_args("out_type", "Optional output tensor type carrying the virtual device."))
.set_attr<FRelaxInferLayout>("FRelaxInferLayout", InferLayoutUnaryEwise)
.set_attr<TMixedPrecisionPolicy>("TMixedPrecisionPolicy", MixedPrecisionPolicyKind::kFollow)
.set_attr<bool>("FPurity", true)
@@ -397,8 +417,9 @@ TVM_FFI_STATIC_INIT_BLOCK() {
tvm::ffi::reflection::GlobalDef().def("relax.op.sin", sin);
OpDef("relax.sin")
- .add_arg("x", "The input tensor.")
- .add_ty_arg("out_type", "Optional output tensor type carrying the virtual device.")
+ .signature(
+ sig::arg("x", "The input tensor."),
+ sig::var_ty_args("out_type", "Optional output tensor type carrying the virtual device."))
.set_attr<FRelaxInferLayout>("FRelaxInferLayout", InferLayoutUnaryEwise)
.set_attr<TMixedPrecisionPolicy>("TMixedPrecisionPolicy", MixedPrecisionPolicyKind::kFollow)
.set_attr<bool>("FPurity", true)
@@ -407,8 +428,9 @@ TVM_FFI_STATIC_INIT_BLOCK() {
tvm::ffi::reflection::GlobalDef().def("relax.op.sinh", sinh);
OpDef("relax.sinh")
- .add_arg("x", "The input tensor.")
- .add_ty_arg("out_type", "Optional output tensor type carrying the virtual device.")
+ .signature(
+ sig::arg("x", "The input tensor."),
+ sig::var_ty_args("out_type", "Optional output tensor type carrying the virtual device."))
.set_attr<FRelaxInferLayout>("FRelaxInferLayout", InferLayoutUnaryEwise)
.set_attr<TMixedPrecisionPolicy>("TMixedPrecisionPolicy", MixedPrecisionPolicyKind::kFollow)
.set_attr<bool>("FPurity", true)
@@ -417,8 +439,9 @@ TVM_FFI_STATIC_INIT_BLOCK() {
tvm::ffi::reflection::GlobalDef().def("relax.op.square", square);
OpDef("relax.square")
- .add_arg("x", "The input tensor.")
- .add_ty_arg("out_type", "Optional output tensor type carrying the virtual device.")
+ .signature(
+ sig::arg("x", "The input tensor."),
+ sig::var_ty_args("out_type", "Optional output tensor type carrying the virtual device."))
.set_attr<FRelaxInferLayout>("FRelaxInferLayout", InferLayoutUnaryEwise)
.set_attr<TMixedPrecisionPolicy>("TMixedPrecisionPolicy", MixedPrecisionPolicyKind::kFollow)
.set_attr<bool>("FPurity", true)
@@ -427,8 +450,9 @@ TVM_FFI_STATIC_INIT_BLOCK() {
tvm::ffi::reflection::GlobalDef().def("relax.op.sqrt", sqrt);
OpDef("relax.sqrt")
- .add_arg("x", "The input tensor.")
- .add_ty_arg("out_type", "Optional output tensor type carrying the virtual device.")
+ .signature(
+ sig::arg("x", "The input tensor."),
+ sig::var_ty_args("out_type", "Optional output tensor type carrying the virtual device."))
.set_attr<FRelaxInferLayout>("FRelaxInferLayout", InferLayoutUnaryEwise)
.set_attr<TMixedPrecisionPolicy>("TMixedPrecisionPolicy", MixedPrecisionPolicyKind::kFollow)
.set_attr<bool>("FPurity", true)
@@ -437,8 +461,9 @@ TVM_FFI_STATIC_INIT_BLOCK() {
tvm::ffi::reflection::GlobalDef().def("relax.op.tan", tan);
OpDef("relax.tan")
- .add_arg("x", "The input tensor.")
- .add_ty_arg("out_type", "Optional output tensor type carrying the virtual device.")
+ .signature(
+ sig::arg("x", "The input tensor."),
+ sig::var_ty_args("out_type", "Optional output tensor type carrying the virtual device."))
.set_attr<FRelaxInferLayout>("FRelaxInferLayout", InferLayoutUnaryEwise)
.set_attr<TMixedPrecisionPolicy>("TMixedPrecisionPolicy", MixedPrecisionPolicyKind::kFollow)
.set_attr<bool>("FPurity", true)
@@ -447,8 +472,9 @@ TVM_FFI_STATIC_INIT_BLOCK() {
tvm::ffi::reflection::GlobalDef().def("relax.op.tanh", tanh);
OpDef("relax.tanh")
- .add_arg("x", "The input tensor.")
- .add_ty_arg("out_type", "Optional output tensor type carrying the virtual device.")
+ .signature(
+ sig::arg("x", "The input tensor."),
+ sig::var_ty_args("out_type", "Optional output tensor type carrying the virtual device."))
.set_attr<FRelaxInferLayout>("FRelaxInferLayout", InferLayoutUnaryEwise)
.set_attr<TMixedPrecisionPolicy>("TMixedPrecisionPolicy", MixedPrecisionPolicyKind::kFollow)
.set_attr<bool>("FPurity", true)
@@ -457,8 +483,9 @@ TVM_FFI_STATIC_INIT_BLOCK() {
tvm::ffi::reflection::GlobalDef().def("relax.op.trunc", trunc);
OpDef("relax.trunc")
- .add_arg("x", "The input tensor.")
- .add_ty_arg("out_type", "Optional output tensor type carrying the virtual device.")
+ .signature(
+ sig::arg("x", "The input tensor."),
+ sig::var_ty_args("out_type", "Optional output tensor type carrying the virtual device."))
.set_attr<FRelaxInferLayout>("FRelaxInferLayout", InferLayoutUnaryEwise)
.set_attr<TMixedPrecisionPolicy>("TMixedPrecisionPolicy", MixedPrecisionPolicyKind::kFollow)
.set_attr<bool>("FPurity", true)
@@ -467,8 +494,9 @@ TVM_FFI_STATIC_INIT_BLOCK() {
tvm::ffi::reflection::GlobalDef().def("relax.op.erf", erf);
OpDef("relax.erf")
- .add_arg("x", "The input tensor.")
- .add_ty_arg("out_type", "Optional output tensor type carrying the virtual device.")
+ .signature(
+ sig::arg("x", "The input tensor."),
+ sig::var_ty_args("out_type", "Optional output tensor type carrying the virtual device."))
.set_attr<FRelaxInferLayout>("FRelaxInferLayout", InferLayoutUnaryEwise)
.set_attr<TMixedPrecisionPolicy>("TMixedPrecisionPolicy", MixedPrecisionPolicyKind::kFollow)
.set_attr<bool>("FPurity", true)
@@ -477,9 +505,9 @@ TVM_FFI_STATIC_INIT_BLOCK() {
// relax.clip
OpDef("relax.clip")
- .add_arg("x", "The input tensor.")
- .add_arg("min", "The lower-bound of the range to be clipped to")
- .add_arg("max", "The upper-bound of the range to be clipped to")
+ .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<bool>("FPurity", true);
}
@@ -492,7 +520,7 @@ Expr clip(Expr x, Expr min, Expr max) {
<< "The argument `max` of relax.clip is expected to be a PrimExpr, but got "
<< max->GetTypeKey();
static const Op op = Op::Get("relax.clip");
- return Call(Type::Missing(), op, {std::move(x), std::move(min), std::move(max)});
+ return Call::Unchecked(Type::Missing(), op, {std::move(x), std::move(min), std::move(max)});
}
TVM_FFI_STATIC_INIT_BLOCK() {
@@ -504,8 +532,9 @@ TVM_FFI_STATIC_INIT_BLOCK() {
tvm::ffi::reflection::GlobalDef().def("relax.op.isfinite", isfinite);
OpDef("relax.isfinite")
- .add_arg("x", "The input tensor.")
- .add_ty_arg("out_type", "Optional output tensor type carrying the virtual device.")
+ .signature(
+ sig::arg("x", "The input tensor."),
+ sig::var_ty_args("out_type", "Optional output tensor type carrying the virtual device."))
.set_attr<FRelaxInferLayout>("FRelaxInferLayout", InferLayoutUnaryEwise)
.set_attr<TMixedPrecisionPolicy>("TMixedPrecisionPolicy", MixedPrecisionPolicyKind::kFollow)
.set_attr<bool>("FPurity", true)
@@ -515,8 +544,9 @@ TVM_FFI_STATIC_INIT_BLOCK() {
tvm::ffi::reflection::GlobalDef().def("relax.op.isinf", isinf);
OpDef("relax.isinf")
- .add_arg("x", "The input tensor.")
- .add_ty_arg("out_type", "Optional output tensor type carrying the virtual device.")
+ .signature(
+ sig::arg("x", "The input tensor."),
+ sig::var_ty_args("out_type", "Optional output tensor type carrying the virtual device."))
.set_attr<FRelaxInferLayout>("FRelaxInferLayout", InferLayoutUnaryEwise)
.set_attr<TMixedPrecisionPolicy>("TMixedPrecisionPolicy", MixedPrecisionPolicyKind::kFollow)
.set_attr<bool>("FPurity", true)
@@ -526,8 +556,9 @@ TVM_FFI_STATIC_INIT_BLOCK() {
tvm::ffi::reflection::GlobalDef().def("relax.op.isnan", isnan);
OpDef("relax.isnan")
- .add_arg("x", "The input tensor.")
- .add_ty_arg("out_type", "Optional output tensor type carrying the virtual device.")
+ .signature(
+ sig::arg("x", "The input tensor."),
+ sig::var_ty_args("out_type", "Optional output tensor type carrying the virtual device."))
.set_attr<FRelaxInferLayout>("FRelaxInferLayout", InferLayoutUnaryEwise)
.set_attr<TMixedPrecisionPolicy>("TMixedPrecisionPolicy", MixedPrecisionPolicyKind::kFollow)
.set_attr<bool>("FPurity", true)
diff --git a/src/relax/op/vision/multibox_transform_loc.cc b/src/relax/op/vision/multibox_transform_loc.cc
index fec105de3a..15daf0543f 100644
--- a/src/relax/op/vision/multibox_transform_loc.cc
+++ b/src/relax/op/vision/multibox_transform_loc.cc
@@ -49,8 +49,9 @@ Expr multibox_transform_loc(Expr cls_pred, Expr loc_pred, Expr anchor, bool clip
attrs->apply_softmax = apply_softmax;
static const Op op = Op::Get("relax.vision.multibox_transform_loc");
- return Call(Type::Missing(), op, {std::move(cls_pred), std::move(loc_pred), std::move(anchor)},
- Attrs(attrs), {});
+ return Call::Unchecked(Type::Missing(), op,
+ {std::move(cls_pred), std::move(loc_pred), std::move(anchor)},
+ Attrs(attrs), {});
}
TVM_FFI_STATIC_INIT_BLOCK() {
@@ -191,10 +192,11 @@ TVM_FFI_STATIC_INIT_BLOCK() {
"Decode SSD/TFLite-style priors and offsets into boxes and class scores. If "
"cls_pred shape is unknown, N-based loc/anchor shape checks are skipped in "
"inference. Very large variances (w,h) can overflow exp in half box sizes.")
- .add_arg("cls_pred", "[B,C,N] class logits or scores.")
- .add_arg("loc_pred", "[B,4*N] box encodings (x,y,w,h); TFLite yxhw order remapped to xywh.")
- .add_arg("anchor", "[1,N,4] priors as ltrb (left,top,right,bottom).")
- .attrs_type<MultiboxTransformLocAttrs>()
+ .signature(sig::arg("cls_pred", "[B,C,N] class logits or scores."),
+ sig::arg("loc_pred",
+ "[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<bool>("FPurity", true);
}
diff --git a/src/relax/op/vision/nms.cc b/src/relax/op/vision/nms.cc
index 6c945918ae..d3a0e74ba6 100644
--- a/src/relax/op/vision/nms.cc
+++ b/src/relax/op/vision/nms.cc
@@ -50,10 +50,11 @@ Expr all_class_non_max_suppression(Expr boxes, Expr scores, Expr max_output_boxe
attrs->output_format = output_format;
static const Op op = Op::Get("relax.vision.all_class_non_max_suppression");
- return Call(Type::Missing(), op,
- {std::move(boxes), std::move(scores), std::move(max_output_boxes_per_class),
- std::move(iou_threshold), std::move(score_threshold)},
- Attrs(attrs), {});
+ return Call::Unchecked(
+ Type::Missing(), op,
+ {std::move(boxes), std::move(scores), std::move(max_output_boxes_per_class),
+ std::move(iou_threshold), std::move(score_threshold)},
+ Attrs(attrs), {});
}
TVM_FFI_STATIC_INIT_BLOCK() {
@@ -104,13 +105,14 @@ Type InferTypeAllClassNMS(const Call& call, const BlockBuilder& ctx) {
TVM_FFI_STATIC_INIT_BLOCK() {
OpDef("relax.vision.all_class_non_max_suppression")
- .add_arg("boxes", "The input boxes in the format [batch, num_boxes, 4].")
- .add_arg("scores",
- "Scores for each box and class in the format [batch, num_classes, num_boxes].")
- .add_arg("max_output_boxes_per_class", "The maximum number of output boxes per class.")
- .add_arg("iou_threshold", "The IoU threshold for box the overlap test.")
- .add_arg("score_threshold", "The score threshold to filter out low score boxes early.")
- .attrs_type<AllClassNonMaximumSuppressionAttrs>()
+ .signature(
+ sig::arg("boxes", "The input boxes in the format [batch, num_boxes, 4]."),
+ sig::arg("scores",
+ "Scores for each box and class in the format [batch, num_classes, num_boxes]."),
+ sig::arg("max_output_boxes_per_class", "The maximum number of output boxes per class."),
+ 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<bool>("FPurity", true);
}
@@ -124,7 +126,7 @@ Expr get_valid_counts(Expr data, double score_threshold, int id_index, int score
attrs->score_index = score_index;
static const Op op = Op::Get("relax.vision.get_valid_counts");
- return Call(Type::Missing(), op, {std::move(data)}, Attrs(attrs), {});
+ return Call::Unchecked(Type::Missing(), op, {std::move(data)}, Attrs(attrs), {});
}
TVM_FFI_STATIC_INIT_BLOCK() {
@@ -184,8 +186,8 @@ Type InferTypeGetValidCounts(const Call& call, const BlockBuilder& ctx) {
TVM_FFI_STATIC_INIT_BLOCK() {
OpDef("relax.vision.get_valid_counts")
- .add_arg("data", "Input data, 3-D tensor [batch_size, num_anchors, elem_length].")
- .attrs_type<GetValidCountsAttrs>()
+ .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<bool>("FPurity", true);
}
@@ -210,8 +212,9 @@ Expr non_max_suppression(Expr data, Expr valid_count, Expr indices, int max_outp
attrs->score_threshold = score_threshold;
static const Op op = Op::Get("relax.vision.non_max_suppression");
- return Call(Type::Missing(), op, {std::move(data), std::move(valid_count), std::move(indices)},
- Attrs(attrs), {});
+ return Call::Unchecked(Type::Missing(), op,
+ {std::move(data), std::move(valid_count), std::move(indices)},
+ Attrs(attrs), {});
}
TVM_FFI_STATIC_INIT_BLOCK() {
@@ -355,10 +358,10 @@ Type InferTypeNMS(const Call& call, const BlockBuilder& ctx) {
TVM_FFI_STATIC_INIT_BLOCK() {
OpDef("relax.vision.non_max_suppression")
- .add_arg("data", "Input data, 3-D tensor [batch_size, num_anchors, elem_length].")
- .add_arg("valid_count", "1-D tensor for valid number of boxes.")
- .add_arg("indices", "2-D tensor with shape [batch_size, num_anchors].")
- .attrs_type<NonMaximumSuppressionAttrs>()
+ .signature(sig::arg("data", "Input data, 3-D tensor [batch_size, num_anchors, elem_length]."),
+ 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<bool>("FPurity", true);
}
diff --git a/src/relax/op/vision/roi_align.cc b/src/relax/op/vision/roi_align.cc
index c8643590db..ea08faad6c 100644
--- a/src/relax/op/vision/roi_align.cc
+++ b/src/relax/op/vision/roi_align.cc
@@ -52,7 +52,7 @@ Expr roi_align(Expr data, Expr rois, ffi::Array<int64_t> pooled_size, double spa
attrs->mode = mode;
static const Op op = Op::Get("relax.vision.roi_align");
- return Call(Type::Missing(), op, {std::move(data), std::move(rois)}, Attrs(attrs), {});
+ return Call::Unchecked(Type::Missing(), op, {std::move(data), std::move(rois)}, Attrs(attrs), {});
}
TVM_FFI_STATIC_INIT_BLOCK() {
@@ -130,10 +130,11 @@ Type InferTypeROIAlign(const Call& call, const BlockBuilder& ctx) {
TVM_FFI_STATIC_INIT_BLOCK() {
OpDef("relax.vision.roi_align")
- .add_arg("data", "The input tensor.")
- .add_arg("rois",
- "The input rois with shape (num_roi, 5) in [batch_idx, x1, y1, x2, y2] format.")
- .attrs_type<ROIAlignAttrs>()
+ .signature(
+ sig::arg("data", "The input tensor."),
+ 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<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 f2c768894f..fdd8f534d2 100644
--- a/src/relax/op/vision/roi_pool.cc
+++ b/src/relax/op/vision/roi_pool.cc
@@ -49,7 +49,7 @@ Expr roi_pool(Expr data, Expr rois, ffi::Array<int64_t> pooled_size, double spat
attrs->layout = layout;
static const Op op = Op::Get("relax.vision.roi_pool");
- return Call(Type::Missing(), op, {std::move(data), std::move(rois)}, Attrs(attrs), {});
+ return Call::Unchecked(Type::Missing(), op, {std::move(data), std::move(rois)}, Attrs(attrs), {});
}
TVM_FFI_STATIC_INIT_BLOCK() {
@@ -118,10 +118,11 @@ Type InferTypeROIPool(const Call& call, const BlockBuilder& ctx) {
TVM_FFI_STATIC_INIT_BLOCK() {
OpDef("relax.vision.roi_pool")
- .add_arg("data", "The input tensor.")
- .add_arg("rois",
- "The input rois with shape (num_roi, 5) in [batch_idx, x1, y1, x2, y2] format.")
- .attrs_type<ROIPoolAttrs>()
+ .signature(
+ sig::arg("data", "The input tensor."),
+ 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<TMixedPrecisionPolicy>("TMixedPrecisionPolicy", MixedPrecisionPolicyKind::kFollow)
.set_attr<bool>("FPurity", true);
diff --git a/src/relax/script/ir_builder/distributed.cc b/src/relax/script/ir_builder/distributed.cc
index fd80fb5adb..24a6746f7c 100644
--- a/src/relax/script/ir_builder/distributed.cc
+++ b/src/relax/script/ir_builder/distributed.cc
@@ -45,7 +45,7 @@ Expr MakeCallTIRDist(Expr func, Tuple args, ffi::Array<distributed::DTensorType>
}
static const Op op = Op::Get("relax.call_tir");
- return Call(Type::Missing(), op, {func, args}, {}, {out_ty});
+ return Call::Unchecked(Type::Missing(), op, {func, args}, {}, {out_ty});
}
TVM_FFI_STATIC_INIT_BLOCK() {
diff --git a/src/relax/transform/allocate_workspace.cc b/src/relax/transform/allocate_workspace.cc
index ab80be5119..41051b557c 100644
--- a/src/relax/transform/allocate_workspace.cc
+++ b/src/relax/transform/allocate_workspace.cc
@@ -88,8 +88,8 @@ class ExternFunctionRewriter : ExprMutator {
auto new_args = call_node->args;
TVM_FFI_ICHECK(workspace_var_param_.defined());
new_args.push_back(workspace_var_param_);
- return Call(Type::Missing(), new_op, new_args, call_node->attrs, call_node->ty_args,
- call_node->span);
+ return Call::Unchecked(Type::Missing(), new_op, new_args, call_node->attrs,
+ call_node->ty_args, call_node->span);
}
}
return ExprMutator::VisitExpr_(call_node);
@@ -175,8 +175,8 @@ class WorkspaceProvider : ExprMutator {
auto new_args = call_node->args;
TVM_FFI_ICHECK(workspace_var_main_.defined());
new_args.push_back(workspace_var_main_);
- return Call(Type::Missing(), new_op, new_args, call_node->attrs, call_node->ty_args,
- call_node->span);
+ return Call::Unchecked(Type::Missing(), new_op, new_args, call_node->attrs,
+ call_node->ty_args, call_node->span);
}
}
diff --git a/src/relax/transform/alter_op_impl.cc b/src/relax/transform/alter_op_impl.cc
index bdcbfd02bc..1d5dbe7a59 100644
--- a/src/relax/transform/alter_op_impl.cc
+++ b/src/relax/transform/alter_op_impl.cc
@@ -142,9 +142,9 @@ class AlterOpImplMutator : public ExprMutator {
TVM_FFI_ICHECK_EQ(call->ty_args.size(), 1) << "call_tir ty_args.size() is expected to be 1";
Type updated_ret_ty = UpdateOutputType(call->ty_args[0], buffer_transforms);
- auto updated_call =
- builder_->Normalize(Call(Type::Missing(), call_tir_op_, {replacement_gv, updated_inputs},
- call->attrs, {updated_ret_ty}));
+ auto updated_call = builder_->Normalize(Call::Unchecked(Type::Missing(), call_tir_op_,
+ {replacement_gv, updated_inputs},
+ call->attrs, {updated_ret_ty}));
// Now transform each of the outputs to previous layout.
return TransformOutputs(updated_call, buffer_transforms, call->ty_args[0]);
@@ -184,7 +184,8 @@ class AlterOpImplMutator : public ExprMutator {
// identical. The scope of vars used in index map initial indices is local to the op. Not doing
// so would confuse the structural equality check.
attrs->index_map = DeepCopyIndexMap(index_map);
- return Call(Type::Missing(), layout_transform_op_, {expr}, Attrs{std::move(attrs)}, {});
+ return Call::Unchecked(Type::Missing(), layout_transform_op_, {expr}, Attrs{std::move(attrs)},
+ {});
}
/*!
@@ -247,8 +248,8 @@ class AlterOpImplMutator : public ExprMutator {
const auto& tensor_ty = padded_expr->ty.as_or_throw<TensorType>();
GlobalVar gv_remove_pad = GetOrCreateRemovePadOp(old_shape, tensor_ty->dtype.value()->dtype);
- return Call(Type::Missing(), call_tir_op_, {gv_remove_pad, Tuple({padded_expr})}, {},
- {old_tensor_ty});
+ return Call::Unchecked(Type::Missing(), call_tir_op_, {gv_remove_pad, Tuple({padded_expr})},
+ {}, {old_tensor_ty});
}
}
diff --git a/src/relax/transform/call_tir_rewrite.cc b/src/relax/transform/call_tir_rewrite.cc
index 3448311fdd..d81134fecd 100644
--- a/src/relax/transform/call_tir_rewrite.cc
+++ b/src/relax/transform/call_tir_rewrite.cc
@@ -137,7 +137,7 @@ class CallTIRMutator : public ExprMutator {
}
}
}
- builder_->Emit(Call(Type::Missing(), call->args[0], args), "_");
+ builder_->Emit(Call::Unchecked(Type::Missing(), call->args[0], args), "_");
} else {
if (!is_inplace) {
args = outs;
@@ -145,7 +145,7 @@ class CallTIRMutator : public ExprMutator {
} else {
args.push_back(call->args[1]);
}
- builder_->Emit(Call(Type::Missing(), call->args[0], args), "_");
+ builder_->Emit(Call::Unchecked(Type::Missing(), call->args[0], args), "_");
}
if (tuple_output_type.has_value()) {
@@ -174,11 +174,11 @@ class CallTIRMutator : public ExprMutator {
dev_index = GetDeviceIndexByScope(mod_, scope);
}
- return builder_->Emit(Call(Type::Missing(), alloc_tensor_op,
- {tensor_ty->shape.value().as_or_throw<ShapeExpr>(),
- DataTypeImm(tensor_ty->dtype.value()->dtype),
- IntImm::Int64(dev_index), StringImm(scope)},
- Attrs(), {tensor_ty}),
+ return builder_->Emit(Call::Unchecked(Type::Missing(), alloc_tensor_op,
+ {tensor_ty->shape.value().as_or_throw<ShapeExpr>(),
+ DataTypeImm(tensor_ty->dtype.value()->dtype),
+ IntImm::Int64(dev_index), StringImm(scope)},
+ Attrs(), {tensor_ty}),
"alloc");
}
diff --git a/src/relax/transform/convert_layout.cc b/src/relax/transform/convert_layout.cc
index 2ab89bbfa8..0401d5f77a 100644
--- a/src/relax/transform/convert_layout.cc
+++ b/src/relax/transform/convert_layout.cc
@@ -133,8 +133,8 @@ class LayoutConvertMutator : public ExprMutator {
ffi::ObjectPtr<LayoutTransformAttrs> attrs = ffi::make_object<LayoutTransformAttrs>();
attrs->index_map = ffi::FromJSONGraph(ffi::ToJSONGraph(index_map)).as_or_throw<IndexMap>();
const Op layout_transform_op_ = Op::Get("relax.layout_transform");
- auto ret_expr =
- Call(Type::Missing(), layout_transform_op_, {expr}, Attrs{std::move(attrs)}, {});
+ auto ret_expr = Call::Unchecked(Type::Missing(), layout_transform_op_, {expr},
+ Attrs{std::move(attrs)}, {});
return ret_expr;
}
};
diff --git a/src/relax/transform/decompose_ops.cc b/src/relax/transform/decompose_ops.cc
index 4e2fe68b96..f0ee758e44 100644
--- a/src/relax/transform/decompose_ops.cc
+++ b/src/relax/transform/decompose_ops.cc
@@ -151,9 +151,9 @@ Expr TensorToShape(const Call& call_node, const BlockBuilder& builder) {
// call builtin function that converts tensor to shape tuple
// TODO(@sunggg): Register operator for "vm.builtin.tensor_to_shape"
static const Op call_pure_packed_op = Op::Get("relax.call_pure_packed");
- Var call = builder->Emit(Call(Type::Missing(), call_pure_packed_op,
- {ExternFunc("vm.builtin.tensor_to_shape"), expr}, {},
- {ffi::GetRef<ShapeType>(ty)}));
+ Var call = builder->Emit(Call::Unchecked(Type::Missing(), call_pure_packed_op,
+ {ExternFunc("vm.builtin.tensor_to_shape"), expr}, {},
+ {ffi::GetRef<ShapeType>(ty)}));
// Operators like reshape take the output of `TensorToShape` as their output shape.
// Because TOPI expects to have such output shape in symbolic shape at least (i.e.,
diff --git a/src/relax/transform/fuse_ops.cc b/src/relax/transform/fuse_ops.cc
index 17a737264c..a96a545f81 100644
--- a/src/relax/transform/fuse_ops.cc
+++ b/src/relax/transform/fuse_ops.cc
@@ -938,7 +938,7 @@ class OperatorFusor : public ExprMutator {
// insertion point and may itself be a dead internal binding.
TVM_FFI_ICHECK(!output_vars.empty());
Var new_var;
- Call call_to_emit = Call(Type::Missing(), gv, UpdateArgs(func_info.arguments_));
+ Call call_to_emit = Call::Unchecked(Type::Missing(), gv, UpdateArgs(func_info.arguments_));
if (output_vars.size() == 1 && !output_vars[0]->IsInstance<DataflowVarNode>()) {
new_var = builder_->EmitOutput(call_to_emit);
@@ -1387,7 +1387,7 @@ class CompositeFunctionAnnotator : public ExprMutator {
Expr VisitExpr_(const CallNode* call_node) final {
if (auto const* gvar = call_node->op.as<GlobalVarNode>()) {
if (auto it = gvar_map_.find(gvar); it != gvar_map_.end()) {
- return Call(Type::Missing(), it->second, call_node->args);
+ return Call::Unchecked(Type::Missing(), it->second, call_node->args);
}
auto func = builder_->GetContextIRModule()->Lookup(ffi::GetRef<GlobalVar>(gvar));
if (auto composite_name = func->GetAttr<ffi::String>(attr::kComposite)) {
@@ -1400,7 +1400,7 @@ class CompositeFunctionAnnotator : public ExprMutator {
builder_->GetContextIRModule()->Remove(ffi::GetRef<GlobalVar>(gvar));
auto new_gvar = builder_->AddFunction(new_func, gsymbol);
gvar_map_[gvar] = new_gvar;
- return Call(Type::Missing(), new_gvar, call_node->args);
+ return Call::Unchecked(Type::Missing(), new_gvar, call_node->args);
}
}
return ExprMutator::VisitExpr_(call_node);
@@ -1431,11 +1431,12 @@ class CompositeFunctionAnnotator : public ExprMutator {
// well-formed Relax IR. As a result, we need to build the SeqExpr ourselves.
Var local_func_var("local_func", GetType(f_inner));
Var output_var("output", f_inner->ret_ty);
- SeqExpr new_body({BindingBlock({
- VarBinding(local_func_var, f_inner),
- VarBinding(output_var, Call(Type::Missing(), local_func_var, params)),
- })},
- output_var);
+ SeqExpr new_body(
+ {BindingBlock({
+ VarBinding(local_func_var, f_inner),
+ VarBinding(output_var, Call::Unchecked(Type::Missing(), local_func_var, params)),
+ })},
+ output_var);
// pure if the inner func is pure (no need to force purity if it's forced for the inner func)
return Function(param_vars, new_body, func_node->ret_ty, f_inner->is_pure);
diff --git a/src/relax/transform/fuse_tir.cc b/src/relax/transform/fuse_tir.cc
index e18a01b04a..72eb56004d 100644
--- a/src/relax/transform/fuse_tir.cc
+++ b/src/relax/transform/fuse_tir.cc
@@ -1145,7 +1145,7 @@ class TIRFuseMutator : public ExprMutator {
inplace_attrs->inplace_indices = replacement.inplace_indices;
call_attrs = Attrs(inplace_attrs);
}
- return Call(Type::Missing(), call_op, call_args, call_attrs, {GetType(call)});
+ return Call::Unchecked(Type::Missing(), call_op, call_args, call_attrs, {GetType(call)});
}
private:
diff --git a/src/relax/transform/gradient.cc b/src/relax/transform/gradient.cc
index 5421192e96..1795859b02 100644
--- a/src/relax/transform/gradient.cc
+++ b/src/relax/transform/gradient.cc
@@ -76,8 +76,8 @@ class CallTIRWithGradEliminator : private ExprMutator {
if (!call_node->op.same_as(Op::Get("relax.call_tir_with_grad"))) {
return ExprMutator::VisitExpr_(call_node);
}
- return Call(Type::Missing(), Op::Get("relax.call_tir"), call_node->args, {}, call_node->ty_args,
- call_node->span);
+ return Call::Unchecked(Type::Missing(), Op::Get("relax.call_tir"), call_node->args, {},
+ call_node->ty_args, call_node->span);
}
};
diff --git a/src/relax/transform/kill_after_last_use.cc b/src/relax/transform/kill_after_last_use.cc
index d4ec4665bf..29c102ee6c 100644
--- a/src/relax/transform/kill_after_last_use.cc
+++ b/src/relax/transform/kill_after_last_use.cc
@@ -245,19 +245,21 @@ class KillInserter : public ExprMutator {
if (auto it = last_usage_.find(binding->var.get()); it != last_usage_.end()) {
static const Op mem_kill_tensor = Op::Get("relax.memory.kill_tensor");
for (const auto& tensor_obj : it->second.tensors) {
- builder_->Emit(Call(Type::Missing(), mem_kill_tensor, {ffi::GetRef<Expr>(tensor_obj)}),
- /*name_hint=*/"_");
+ builder_->Emit(
+ Call::Unchecked(Type::Missing(), mem_kill_tensor, {ffi::GetRef<Expr>(tensor_obj)}),
+ /*name_hint=*/"_");
}
static const Op mem_kill_storage = Op::Get("relax.memory.kill_storage");
for (const VarNode* storage_obj : it->second.storage) {
- builder_->Emit(Call(Type::Missing(), mem_kill_storage, {ffi::GetRef<Expr>(storage_obj)}),
- /*name_hint=*/"_");
+ builder_->Emit(
+ Call::Unchecked(Type::Missing(), mem_kill_storage, {ffi::GetRef<Expr>(storage_obj)}),
+ /*name_hint=*/"_");
}
static const Op vm_kill_object = Op::Get("relax.vm.kill_object");
for (const VarNode* obj : it->second.objects) {
- builder_->Emit(Call(Type::Missing(), vm_kill_object, {ffi::GetRef<Expr>(obj)}),
+ builder_->Emit(Call::Unchecked(Type::Missing(), vm_kill_object, {ffi::GetRef<Expr>(obj)}),
/*name_hint=*/"_");
}
}
diff --git a/src/relax/transform/lambda_lift.cc b/src/relax/transform/lambda_lift.cc
index 410cba0a7a..d3391eb0ae 100644
--- a/src/relax/transform/lambda_lift.cc
+++ b/src/relax/transform/lambda_lift.cc
@@ -310,9 +310,10 @@ class LambdaLifter : public ExprMutator {
// Defining the rewrite rule prior to visiting the body, so that
// recursive closures can be updated.
if (is_recursive && is_closure) {
- nested_closure_map_.emplace(current_lambda_var_.value(),
- Call(Type::Missing(), gvar_lifted_func,
- captured_vars.Map([](Var var) -> Expr { return var; })));
+ nested_closure_map_.emplace(
+ current_lambda_var_.value(),
+ Call::Unchecked(Type::Missing(), gvar_lifted_func,
+ captured_vars.Map([](Var var) -> Expr { return var; })));
}
if (!is_closure) {
@@ -351,7 +352,7 @@ class LambdaLifter : public ExprMutator {
Tuple arg_tuple(captured_vars.Map([](Var var) -> Expr { return var; }));
// Call make_closure intrinsic
callable_value =
- Call(Type::Missing(), make_closure_op_, {gvar_lifted_func, arg_tuple}, {}, {});
+ Call::Unchecked(Type::Missing(), make_closure_op_, {gvar_lifted_func, arg_tuple}, {}, {});
}
return callable_value;
@@ -387,8 +388,9 @@ class LambdaLifter : public ExprMutator {
}();
auto prev = call;
- call = Call(Type::Missing(), is_pure ? invoke_pure_closure_op_ : invoke_closure_op_,
- {var, Tuple(call->args)}, {}, {orig_ty});
+ call =
+ Call::Unchecked(Type::Missing(), is_pure ? invoke_pure_closure_op_ : invoke_closure_op_,
+ {var, Tuple(call->args)}, {}, {orig_ty});
}
}
@@ -403,7 +405,8 @@ class LambdaLifter : public ExprMutator {
}
auto prev = call;
- call = Call(Type::Missing(), nested_call->op, new_args, call->attrs, call->ty_args);
+ call =
+ Call::Unchecked(Type::Missing(), nested_call->op, new_args, call->attrs, call->ty_args);
}
}
diff --git a/src/relax/transform/lazy_transform_params.cc b/src/relax/transform/lazy_transform_params.cc
index a4d80b9cd0..c6065bc8fa 100644
--- a/src/relax/transform/lazy_transform_params.cc
+++ b/src/relax/transform/lazy_transform_params.cc
@@ -97,11 +97,11 @@ class LazyInputMutator : public ExprMutator {
if (plan_) {
Var var = ffi::GetRef<Var>(op);
if (auto it = plan_->param_lookup.find(var); it != plan_->param_lookup.end()) {
- auto untyped = builder_->Emit(Call(Type::Missing(), plan_->fget_param,
- {
- PrimExpr(IntImm::Int64(it->second)),
- StringImm(var->name),
- }),
+ auto untyped = builder_->Emit(Call::Unchecked(Type::Missing(), plan_->fget_param,
+ {
+ PrimExpr(IntImm::Int64(it->second)),
+ StringImm(var->name),
+ }),
var->name + "_untyped");
return builder_->EmitMatchCast(untyped, GetType(var), var->name);
}
@@ -209,8 +209,8 @@ class LazyOutputMutator : public ExprMutator {
if (plan_.has_value()) {
if (auto it = plan_->output_lookup.find(var); it != plan_->output_lookup.end()) {
for (auto output_index : it->second) {
- callback(Call(Type::Missing(), plan_->fset_output,
- {PrimExpr(IntImm::Int64(output_index)), var}));
+ callback(Call::Unchecked(Type::Missing(), plan_->fset_output,
+ {PrimExpr(IntImm::Int64(output_index)), var}));
}
}
}
diff --git a/src/relax/transform/legalize_ops.cc b/src/relax/transform/legalize_ops.cc
index 0972949073..b1b6343d7a 100644
--- a/src/relax/transform/legalize_ops.cc
+++ b/src/relax/transform/legalize_ops.cc
@@ -149,7 +149,8 @@ class LegalizeMutator : public ExprMutator {
for (auto arg : ret->args) {
ret_args.push_back(arg);
}
- return Call(Type::Missing(), call_pure_packed_op, ret_args, ret->attrs, ret->ty_args);
+ return Call::Unchecked(Type::Missing(), call_pure_packed_op, ret_args, ret->attrs,
+ ret->ty_args);
}
ffi::Optional<Target> GetTarget(const ffi::Array<Type>& types) {
@@ -336,8 +337,8 @@ class LegalizeMutator : public ExprMutator {
// Third choice, use an explicit ffi::String replacement. This does not require the shape
ffi::String packed_func_name = call_packed_map[op];
legalization_func = [packed_func_name](const BlockBuilder& bb, const Call& call) -> Expr {
- return Call(Type::Missing(), ExternFunc(packed_func_name), call->args, Attrs(),
- {GetType(call)});
+ return Call::Unchecked(Type::Missing(), ExternFunc(packed_func_name), call->args, Attrs(),
+ {GetType(call)});
};
} else {
// No legalization.
diff --git a/src/relax/transform/lift_transform_params.cc b/src/relax/transform/lift_transform_params.cc
index a4c80ca5af..6b50bf6f2c 100644
--- a/src/relax/transform/lift_transform_params.cc
+++ b/src/relax/transform/lift_transform_params.cc
@@ -646,10 +646,10 @@ class ConsumeBundledParams : public ExprMutator {
ExprMutator::VisitBinding_(binding, tuple_get_item);
auto new_var = VisitExpr(binding->var);
param_remap_[tuple_get_item->index] = new_var;
- builder_->Emit(
- Call(Type::Missing(), call_pure_packed,
- {builtin_tuple_reset_item, tuple_get_item->tuple, PrimExpr(tuple_get_item->index)},
- tvm::Attrs(), {TupleType(ffi::Array<Type>{})}));
+ builder_->Emit(Call::Unchecked(
+ Type::Missing(), call_pure_packed,
+ {builtin_tuple_reset_item, tuple_get_item->tuple, PrimExpr(tuple_get_item->index)},
+ tvm::Attrs(), {TupleType(ffi::Array<Type>{})}));
} else {
ExprMutator::VisitBinding_(binding, tuple_get_item);
}
diff --git a/src/relax/transform/lower_alloc_tensor.cc b/src/relax/transform/lower_alloc_tensor.cc
index 57a697b5c7..728b71d6ce 100644
--- a/src/relax/transform/lower_alloc_tensor.cc
+++ b/src/relax/transform/lower_alloc_tensor.cc
@@ -119,8 +119,8 @@ class Mutator : public ExprMutator {
Type::Missing(), mem_alloc_storage_op,
{size, runtime_device_index, storage_scope, DataTypeImm((DLDataType{kDLUInt, 8, 1}))});
storage = builder_->Emit(storage, "storage");
- Expr tensor = Call(Type::Missing(), mem_alloc_tensor_op,
- {storage, offset, shape_arg, dtype, op->args[2]});
+ Expr tensor = Call::Unchecked(Type::Missing(), mem_alloc_tensor_op,
+ {storage, offset, shape_arg, dtype, op->args[2]});
return tensor;
} else {
return ExprMutator::VisitExpr_(op);
diff --git a/src/relax/transform/merge_composite_functions.cc b/src/relax/transform/merge_composite_functions.cc
index a16139c204..0c75d7e0fe 100644
--- a/src/relax/transform/merge_composite_functions.cc
+++ b/src/relax/transform/merge_composite_functions.cc
@@ -469,7 +469,7 @@ class CompositeInliner : public ExprMutator {
new_func = WithoutAttr(new_func, tvm::relax::attr::kPrimitive);
inlined_functions_.Set(func, new_func);
}
- return Call(Type::Missing(), inlined_functions_[func], call->args);
+ return Call::Unchecked(Type::Missing(), inlined_functions_[func], call->args);
}
}
@@ -529,7 +529,7 @@ class CompositeFunctionAnnotator : public ExprMutator {
// we call new var instead of the old one.
// we don't have to update args since we are just updating the function to call,
// without any change in the arguments.
- return Call(Type::Missing(), new_var, call->args);
+ return Call::Unchecked(Type::Missing(), new_var, call->args);
}
}
return ffi::GetRef<Call>(call);
diff --git a/src/relax/transform/realize_vdevice.cc b/src/relax/transform/realize_vdevice.cc
index ef3734ddd6..1109d8ce46 100644
--- a/src/relax/transform/realize_vdevice.cc
+++ b/src/relax/transform/realize_vdevice.cc
@@ -397,7 +397,7 @@ class VDeviceTypeUpdater : ExprMutator {
} else {
ffi::ObjectPtr<ToVDeviceAttrs> attrs = ffi::make_object<ToVDeviceAttrs>();
attrs->dst_vdevice = output_vdevice;
- return Call(Type::Missing(), to_vdevice_op_, {arg}, Attrs(attrs), {});
+ return Call::Unchecked(Type::Missing(), to_vdevice_op_, {arg}, Attrs(attrs), {});
}
}
diff --git a/src/relax/transform/remove_purity_checking.cc b/src/relax/transform/remove_purity_checking.cc
index 2cecca0e11..68e58d84ee 100644
--- a/src/relax/transform/remove_purity_checking.cc
+++ b/src/relax/transform/remove_purity_checking.cc
@@ -49,20 +49,21 @@ class PurityRemover : public ExprMutator {
Expr VisitExpr_(const CallNode* call) override {
if (call->op.same_as(call_pure_packed_op_)) {
- auto ret = Call(Type::Missing(), call->args[0],
- ffi::Array<Expr>(call->args.begin() + 1, call->args.end()), call->attrs,
- call->ty_args);
+ auto ret = Call::Unchecked(Type::Missing(), call->args[0],
+ ffi::Array<Expr>(call->args.begin() + 1, call->args.end()),
+ call->attrs, call->ty_args);
return VisitExpr(ret);
}
if (call->op.same_as(call_inplace_packed_op_)) {
// call_inplace_packed has its own attrs so we don't pass those down
- auto ret = Call(Type::Missing(), call->args[0],
- ffi::Array<Expr>(call->args.begin() + 1, call->args.end()), tvm::Attrs(),
- call->ty_args);
+ auto ret = Call::Unchecked(Type::Missing(), call->args[0],
+ ffi::Array<Expr>(call->args.begin() + 1, call->args.end()),
+ tvm::Attrs(), call->ty_args);
return VisitExpr(ret);
}
if (call->op.same_as(invoke_pure_closure_op_)) {
- auto ret = Call(Type::Missing(), invoke_closure_op_, call->args, call->attrs, call->ty_args);
+ auto ret = Call::Unchecked(Type::Missing(), invoke_closure_op_, call->args, call->attrs,
+ call->ty_args);
return VisitExpr(ret);
}
return ExprMutator::VisitExpr_(call);
diff --git a/src/relax/transform/rewrite_cuda_graph.cc b/src/relax/transform/rewrite_cuda_graph.cc
index f5c886472d..ab654772e4 100644
--- a/src/relax/transform/rewrite_cuda_graph.cc
+++ b/src/relax/transform/rewrite_cuda_graph.cc
@@ -797,9 +797,9 @@ class CUDAGraphRewriter : public ExprMutator {
auto gv_alloc = gv_global_alloc_.value();
auto ret_ty = gv_alloc->ty.as_or_throw<FuncType>()->ret;
launch_subgraph =
- Call(Type::Missing(), call_builtin_with_ctx_op,
- {builtin_get_cached_alloc, Tuple({gv_alloc, PrimExpr(IntImm::Int64(0))})}, Attrs(),
- {ret_ty});
+ Call::Unchecked(Type::Missing(), call_builtin_with_ctx_op,
+ {builtin_get_cached_alloc, Tuple({gv_alloc, PrimExpr(IntImm::Int64(0))})},
+ Attrs(), {ret_ty});
} else {
auto gv_func = builder_->AddFunction(
plan->func, current_func_.value()->name_hint + "_cuda_graph_capture");
@@ -832,8 +832,9 @@ class CUDAGraphRewriter : public ExprMutator {
// passing it twice simplifies the handling during the capture phase.
tuple_arg_fields.push_back(plan->propogated_tir_vars.value());
}
- launch_subgraph = Call(Type::Missing(), call_builtin_with_ctx_op,
- {builtin_run_or_capture, Tuple(tuple_arg_fields)}, Attrs(), {call_ty});
+ launch_subgraph =
+ Call::Unchecked(Type::Missing(), call_builtin_with_ctx_op,
+ {builtin_run_or_capture, Tuple(tuple_arg_fields)}, Attrs(), {call_ty});
}
Expr ret_value = builder_->Emit(launch_subgraph);
for (const auto& [var, tuple_index] : plan->outputs) {
diff --git a/src/relax/transform/run_codegen.cc b/src/relax/transform/run_codegen.cc
index 20e7e4737f..656aa58120 100644
--- a/src/relax/transform/run_codegen.cc
+++ b/src/relax/transform/run_codegen.cc
@@ -116,7 +116,7 @@ class CodeGenRunner : ExprMutator {
static const Op call_op = Op::Get("relax.call_dps_packed");
- return Call(Type::Missing(), call_op, new_args, tvm::Attrs(), {ret_ty});
+ return Call::Unchecked(Type::Missing(), call_op, new_args, tvm::Attrs(), {ret_ty});
};
auto ret_ty = GetType(call);
diff --git a/src/relax/transform/split_layout_rewrite_preproc.cc b/src/relax/transform/split_layout_rewrite_preproc.cc
index edb041d1e3..f449b71479 100644
--- a/src/relax/transform/split_layout_rewrite_preproc.cc
+++ b/src/relax/transform/split_layout_rewrite_preproc.cc
@@ -325,8 +325,8 @@ class SplitLayoutRewritePreproc : public ExprMutator {
: preproc_ty_list[0];
// Step 6: Call the preproc function
- Expr preproc_call = builder_->Emit(
- Call(Type::Missing(), call_tir_op, {preproc_gv, Tuple(preproc_args)}, {}, {preproc_ty}));
+ Expr preproc_call = builder_->Emit(Call::Unchecked(
+ Type::Missing(), call_tir_op, {preproc_gv, Tuple(preproc_args)}, {}, {preproc_ty}));
if (rewrite_infos.size() == 1) {
call_tir_args.Set(rewrite_infos[0].buffer_index, preproc_call);
} else {
@@ -334,8 +334,8 @@ class SplitLayoutRewritePreproc : public ExprMutator {
call_tir_args.Set(rewrite_infos[i].buffer_index, TupleGetItem(preproc_call, i));
}
}
- Expr main_call = builder_->Emit(
- Call(Type::Missing(), call_tir_op, {compute_gv, Tuple(call_tir_args)}, {}, call->ty_args));
+ Expr main_call = builder_->Emit(Call::Unchecked(
+ Type::Missing(), call_tir_op, {compute_gv, Tuple(call_tir_args)}, {}, call->ty_args));
return main_call;
}
diff --git a/src/relax/transform/static_plan_block_memory.cc b/src/relax/transform/static_plan_block_memory.cc
index b0b2c0b439..340caaf5d3 100644
--- a/src/relax/transform/static_plan_block_memory.cc
+++ b/src/relax/transform/static_plan_block_memory.cc
@@ -970,9 +970,9 @@ class StorageAllocationRewriter : public ExprMutator {
// And always create a `memory.alloc_tensor` for the old `builtin.alloc_tensor`.
PrimExpr offset = IntImm::Int64(0);
DLDataType dtype = ty->dtype.value()->dtype;
- return Call(Type::Missing(), mem_alloc_tensor,
- {storage_var, offset, ty->shape.value(), DataTypeImm(dtype), call->args[2]},
- Attrs());
+ return Call::Unchecked(
+ Type::Missing(), mem_alloc_tensor,
+ {storage_var, offset, ty->shape.value(), DataTypeImm(dtype), call->args[2]}, Attrs());
} else if (plan_dynamic_output_ && call->op.same_as(alloc_tensor_op)) {
// Case 2. For a `alloc_tensor` that is not planned for memory reuse,
// we would still like to allocate **static** memory for the tensor.
@@ -1005,12 +1005,12 @@ class StorageAllocationRewriter : public ExprMutator {
/*storage_scope=*/call->args[3].as_or_throw<StringImm>(), //
/*dtype=*/DataTypeImm(dtype)});
Var storage = builder_->Emit(alloc_storage, "storage");
- return Call(Type::Missing(), mem_alloc_tensor,
- {storage, //
- /*offset=*/IntImm::Int64(0),
- /*shape=*/ffi::GetRef<ShapeExpr>(shape), //
- /*dtype=*/DataTypeImm(dtype),
- /*vdevice_index=*/call->args[2]});
+ return Call::Unchecked(Type::Missing(), mem_alloc_tensor,
+ {storage, //
+ /*offset=*/IntImm::Int64(0),
+ /*shape=*/ffi::GetRef<ShapeExpr>(shape), //
+ /*dtype=*/DataTypeImm(dtype),
+ /*vdevice_index=*/call->args[2]});
}
}
diff --git a/src/target/intrin_rule.cc b/src/target/intrin_rule.cc
index 4721575fa9..6014fb17d3 100644
--- a/src/target/intrin_rule.cc
+++ b/src/target/intrin_rule.cc
@@ -212,7 +212,7 @@ TVM_FFI_STATIC_INIT_BLOCK() {
});
OpDef("tirx.isfinite")
- .add_arg("x", "The input value.")
+ .signature(sig::arg("x", "The input value."))
.set_attr<tirx::TIRxOpCategory>("TIRxOpCategory", ffi::String("builtin"))
.set_attr<FLegalize>("default.FLegalize", [](const PrimExpr& e) -> PrimExpr {
const CallNode* call = e.as<CallNode>();
@@ -221,7 +221,7 @@ TVM_FFI_STATIC_INIT_BLOCK() {
});
OpDef("tirx.isinf")
- .add_arg("x", "The input value.")
+ .signature(sig::arg("x", "The input value."))
.set_attr<tirx::TIRxOpCategory>("TIRxOpCategory", ffi::String("builtin"))
.set_attr<FLegalize>("default.FLegalize", [](const PrimExpr& e) -> PrimExpr {
const CallNode* call = e.as<CallNode>();
diff --git a/src/te/tensor.cc b/src/te/tensor.cc
index 6c990820f7..228fb2b7d4 100644
--- a/src/te/tensor.cc
+++ b/src/te/tensor.cc
@@ -99,8 +99,7 @@ PrimVar var(std::string name_hint, PrimType t) { return PrimVar(name_hint, t); }
TVM_FFI_STATIC_INIT_BLOCK() {
OpDef("te.tensor_load")
- .add_arg("tensor", "The input tensor.")
- .allow_extra_args()
+ .signature(sig::arg("tensor", "The input tensor."), sig::var_args("args"))
.set_attr<TCallEffectKind>("TCallEffectKind",
static_cast<int64_t>(CallEffectKind::kReadState));
}
diff --git a/src/tirx/op/builtin.cc b/src/tirx/op/builtin.cc
index 6817345788..8f04560055 100644
--- a/src/tirx/op/builtin.cc
+++ b/src/tirx/op/builtin.cc
@@ -159,7 +159,7 @@ TVM_DEFINE_CACHED_OP_GETTER(print_buffer, "tirx.print_buffer")
TVM_FFI_STATIC_INIT_BLOCK() {
OpDef("tirx.reinterpret")
- .add_arg("x", "The input value.")
+ .signature(sig::arg("x", "The input value."))
.set_attr<TScriptPrinterName>("TScriptPrinterName", ffi::String("reinterpret"))
.set_attr<TIRxOpCategory>("TIRxOpCategory", ffi::String("builtin"))
.set_attr<TCallEffectKind>("TCallEffectKind", static_cast<int64_t>(CallEffectKind::kPure))
@@ -180,79 +180,73 @@ TVM_FFI_STATIC_INIT_BLOCK() {
// without this wrapper.
OpDef("tirx.filter")
- .add_arg("var", "The thread-axis variable.")
- .add_arg("pred", "The predicate.")
+ .signature(sig::arg("var", "The thread-axis variable."), sig::arg("pred", "The predicate."))
.set_attr<TScriptPrinterName>("TScriptPrinterName", ffi::String("filter"))
.set_attr<TIRxOpCategory>("TIRxOpCategory", ffi::String("builtin"))
.set_attr<TCallEffectKind>("TCallEffectKind", static_cast<int64_t>(CallEffectKind::kPure));
OpDef("tirx.selector")
- .add_arg("var", "The thread-axis variable.")
- .add_arg("pred", "The predicate.")
+ .signature(sig::arg("var", "The thread-axis variable."), sig::arg("pred", "The predicate."))
.set_attr<TScriptPrinterName>("TScriptPrinterName", ffi::String("selector"))
.set_attr<TIRxOpCategory>("TIRxOpCategory", ffi::String("builtin"))
.set_attr<TCallEffectKind>("TCallEffectKind", static_cast<int64_t>(CallEffectKind::kOpaque));
OpDef("tirx.address_of")
- .add_arg("obj", "The referenced object.")
+ .signature(sig::arg("obj", "The referenced object."))
.set_attr<TScriptPrinterName>("TScriptPrinterName", ffi::String("address_of"))
.set_attr<TIRxOpCategory>("TIRxOpCategory", ffi::String("builtin"))
.set_attr<TCallEffectKind>("TCallEffectKind", static_cast<int64_t>(CallEffectKind::kPure));
OpDef("tirx.q_multiply_shift")
- .add_arg("x", "The input value.")
- .add_arg("y", "The second input value.")
- .add_arg("q", "The number of fractional bits.")
- .add_arg("s", "The right shift amount.")
+ .signature(sig::arg("x", "The input value."), sig::arg("y", "The second input value."),
+ sig::arg("q", "The number of fractional bits."),
+ sig::arg("s", "The right shift amount."))
.set_attr<TScriptPrinterName>("TScriptPrinterName", ffi::String("q_multiply_shift"))
.set_attr<TIRxOpCategory>("TIRxOpCategory", ffi::String("builtin"))
.set_attr<TCallEffectKind>("TCallEffectKind", static_cast<int64_t>(CallEffectKind::kPure))
.set_attr<TVectorizable>("TVectorizable", true);
OpDef("tirx.q_multiply_shift_per_axis")
- .add_arg("x", "The input value.")
- .add_arg("y", "The second input value.")
- .add_arg("ls", "The left shift amount.")
- .add_arg("rs", "The right shift amount.")
- .add_arg("q", "The number of fractional bits.")
- .add_arg("is_lshift_required", "Whether the left shift is required.")
- .add_arg("is_rshift_required", "Whether the right shift is required.")
+ .signature(sig::arg("x", "The input value."), sig::arg("y", "The second input value."),
+ sig::arg("ls", "The left shift amount."),
+ sig::arg("rs", "The right shift amount."),
+ sig::arg("q", "The number of fractional bits."),
+ sig::arg("is_lshift_required", "Whether the left shift is required."),
+ sig::arg("is_rshift_required", "Whether the right shift is required."))
.set_attr<TScriptPrinterName>("TScriptPrinterName", ffi::String("q_multiply_shift_per_axis"))
.set_attr<TIRxOpCategory>("TIRxOpCategory", ffi::String("builtin"))
.set_attr<TCallEffectKind>("TCallEffectKind", static_cast<int64_t>(CallEffectKind::kPure))
.set_attr<TVectorizable>("TVectorizable", true);
OpDef("tirx.isnullptr")
- .add_arg("x", "The input value.")
+ .signature(sig::arg("x", "The input value."))
.set_attr<TScriptPrinterName>("TScriptPrinterName", ffi::String("isnullptr"))
.set_attr<TIRxOpCategory>("TIRxOpCategory", ffi::String("builtin"))
.set_attr<TCallEffectKind>("TCallEffectKind", static_cast<int64_t>(CallEffectKind::kPure));
OpDef("tirx.isnan")
- .add_arg("x", "The input value.")
+ .signature(sig::arg("x", "The input value."))
.set_attr<TScriptPrinterName>("TScriptPrinterName", ffi::String("isnan"))
.set_attr<TIRxOpCategory>("TIRxOpCategory", ffi::String("builtin"))
.set_attr<TCallEffectKind>("TCallEffectKind", static_cast<int64_t>(CallEffectKind::kPure));
OpDef("tirx.popcount")
- .add_arg("x", "The input value.")
+ .signature(sig::arg("x", "The input value."))
.set_attr<TScriptPrinterName>("TScriptPrinterName", ffi::String("popcount"))
.set_attr<TIRxOpCategory>("TIRxOpCategory", ffi::String("builtin"))
.set_attr<TCallEffectKind>("TCallEffectKind", static_cast<int64_t>(CallEffectKind::kPure))
.set_attr<TVectorizable>("TVectorizable", true);
OpDef("tirx.fma")
- .add_arg("x", "The input value.")
- .add_arg("y", "The second input value.")
- .add_arg("z", "The third input value.")
+ .signature(sig::arg("x", "The input value."), sig::arg("y", "The second input value."),
+ sig::arg("z", "The third input value."))
.set_attr<TScriptPrinterName>("TScriptPrinterName", ffi::String("fma"))
.set_attr<TIRxOpCategory>("TIRxOpCategory", ffi::String("builtin"))
.set_attr<TCallEffectKind>("TCallEffectKind", static_cast<int64_t>(CallEffectKind::kPure))
.set_attr<TVectorizable>("TVectorizable", true);
OpDef("tirx.call_extern")
- .add_arg("func_name", "The function name.")
- .allow_extra_args()
+ .signature(sig::arg("func_name", "The function name."), sig::var_args("args"))
.set_attr<TScriptPrinterName>("TScriptPrinterName", ffi::String("call_extern"))
.set_attr<TIRxOpCategory>("TIRxOpCategory", ffi::String("builtin"))
.set_attr<TCallEffectKind>("TCallEffectKind", static_cast<int64_t>(CallEffectKind::kOpaque))
@@ -260,8 +254,7 @@ TVM_FFI_STATIC_INIT_BLOCK() {
static_cast<int64_t>(ScriptDtypePrintLocation::kFirst));
OpDef("tirx.call_pure_extern")
- .add_arg("func_name", "The function name.")
- .allow_extra_args()
+ .signature(sig::arg("func_name", "The function name."), sig::var_args("args"))
.set_attr<TScriptPrinterName>("TScriptPrinterName", ffi::String("call_pure_extern"))
.set_attr<TIRxOpCategory>("TIRxOpCategory", ffi::String("builtin"))
.set_attr<TCallEffectKind>("TCallEffectKind", static_cast<int64_t>(CallEffectKind::kPure))
@@ -269,8 +262,7 @@ TVM_FFI_STATIC_INIT_BLOCK() {
static_cast<int64_t>(ScriptDtypePrintLocation::kFirst));
OpDef("tirx.call_llvm_intrin")
- .add_arg("intrin_id", "The intrinsic identifier.")
- .allow_extra_args()
+ .signature(sig::arg("intrin_id", "The intrinsic identifier."), sig::var_args("args"))
.set_attr<TScriptPrinterName>("TScriptPrinterName", ffi::String("call_llvm_intrin"))
.set_attr<TIRxOpCategory>("TIRxOpCategory", ffi::String("builtin"))
.set_attr<TCallEffectKind>("TCallEffectKind", static_cast<int64_t>(CallEffectKind::kOpaque))
@@ -278,8 +270,7 @@ TVM_FFI_STATIC_INIT_BLOCK() {
static_cast<int64_t>(ScriptDtypePrintLocation::kFirst));
OpDef("tirx.call_llvm_pure_intrin")
- .add_arg("intrin_id", "The intrinsic identifier.")
- .allow_extra_args()
+ .signature(sig::arg("intrin_id", "The intrinsic identifier."), sig::var_args("args"))
.set_attr<TScriptPrinterName>("TScriptPrinterName", ffi::String("call_llvm_pure_intrin"))
.set_attr<TIRxOpCategory>("TIRxOpCategory", ffi::String("builtin"))
.set_attr<TCallEffectKind>("TCallEffectKind", static_cast<int64_t>(CallEffectKind::kPure))
@@ -288,36 +279,31 @@ TVM_FFI_STATIC_INIT_BLOCK() {
.set_attr<TVectorizable>("TVectorizable", true);
OpDef("tirx.call_spirv_pure_glsl450")
- .add_arg("intrin_id", "The intrinsic identifier.")
- .allow_extra_args()
+ .signature(sig::arg("intrin_id", "The intrinsic identifier."), sig::var_args("args"))
.set_attr<TScriptPrinterName>("TScriptPrinterName", ffi::String("call_spirv_pure_glsl450"))
.set_attr<TIRxOpCategory>("TIRxOpCategory", ffi::String("builtin"))
.set_attr<TCallEffectKind>("TCallEffectKind", static_cast<int64_t>(CallEffectKind::kPure));
OpDef("tirx.prefetch")
- .add_arg("ptr", "The pointer.")
- .add_arg("rw", "The read/write mode.")
- .add_arg("locality", "The locality hint.")
- .add_arg("cache_type", "The cache policy.")
+ .signature(sig::arg("ptr", "The pointer."), sig::arg("rw", "The read/write mode."),
+ sig::arg("locality", "The locality hint."),
+ sig::arg("cache_type", "The cache policy."))
.set_attr<TScriptPrinterName>("TScriptPrinterName", ffi::String("prefetch"))
.set_attr<TIRxOpCategory>("TIRxOpCategory", ffi::String("builtin"))
.set_attr<TCallEffectKind>("TCallEffectKind", static_cast<int64_t>(CallEffectKind::kOpaque));
OpDef("tirx.tvm_access_ptr")
- .add_arg("ptype", "The pointer type.")
- .add_arg("data", "The input data.")
- .add_arg("offset", "The offset.")
- .add_arg("extent", "The extent.")
- .add_arg("rw_mask", "The read/write mask.")
+ .signature(sig::arg("ptype", "The pointer type."), sig::arg("data", "The input data."),
+ sig::arg("offset", "The offset."), sig::arg("extent", "The extent."),
+ sig::arg("rw_mask", "The read/write mask."))
.set_attr<TScriptPrinterName>("TScriptPrinterName", ffi::String("tvm_access_ptr"))
.set_attr<TIRxOpCategory>("TIRxOpCategory", ffi::String("builtin"))
.set_attr<TCallEffectKind>("TCallEffectKind",
static_cast<int64_t>(CallEffectKind::kSpecialCallArg));
OpDef("tirx.ptr_byte_offset")
- .add_arg("data", "Base pointer.")
- .add_arg("byte_offset", "Offset in bytes.")
- .add_arg("dtype", "Type annotation for pointed-to elements.")
+ .signature(sig::arg("data", "Base pointer."), sig::arg("byte_offset", "Offset in bytes."),
+ sig::arg("dtype", "Type annotation for pointed-to elements."))
.set_attr<TScriptPrinterName>("TScriptPrinterName", ffi::String("ptr_byte_offset"))
.set_attr<TIRxOpCategory>("TIRxOpCategory", ffi::String("builtin"))
.set_attr<TCallEffectKind>("TCallEffectKind", static_cast<int64_t>(CallEffectKind::kPure));
@@ -329,16 +315,14 @@ TVM_FFI_STATIC_INIT_BLOCK() {
static_cast<int64_t>(CallEffectKind::kSpecialCallArg));
OpDef("tirx.handle_add_byte_offset")
- .add_arg("handle", "The handle.")
- .add_arg("offset", "The offset.")
+ .signature(sig::arg("handle", "The handle."), sig::arg("offset", "The offset."))
.set_attr<TScriptPrinterName>("TScriptPrinterName", ffi::String("handle_add_byte_offset"))
.set_attr<TIRxOpCategory>("TIRxOpCategory", ffi::String("builtin"))
.set_attr<TCallEffectKind>("TCallEffectKind", static_cast<int64_t>(CallEffectKind::kPure));
OpDef("tirx.tvm_struct_get")
- .add_arg("arr", "The array.")
- .add_arg("index", "The index.")
- .add_arg("field", "The field index.")
+ .signature(sig::arg("arr", "The array."), sig::arg("index", "The index."),
+ sig::arg("field", "The field index."))
.set_attr<TScriptPrinterName>("TScriptPrinterName", ffi::String("tvm_struct_get"))
.set_attr<TIRxOpCategory>("TIRxOpCategory", ffi::String("builtin"))
.set_attr<TCallEffectKind>("TCallEffectKind",
@@ -347,10 +331,8 @@ TVM_FFI_STATIC_INIT_BLOCK() {
static_cast<int64_t>(ScriptDtypePrintLocation::kLast));
OpDef("tirx.tvm_struct_set")
- .add_arg("arr", "The array.")
- .add_arg("index", "The index.")
- .add_arg("field", "The field index.")
- .add_arg("value", "The value to use.")
+ .signature(sig::arg("arr", "The array."), sig::arg("index", "The index."),
+ sig::arg("field", "The field index."), sig::arg("value", "The value to use."))
.set_attr<TScriptPrinterName>("TScriptPrinterName", ffi::String("tvm_struct_set"))
.set_attr<TIRxOpCategory>("TIRxOpCategory", ffi::String("builtin"))
.set_attr<TCallEffectKind>("TCallEffectKind",
@@ -362,78 +344,70 @@ TVM_FFI_STATIC_INIT_BLOCK() {
.set_attr<TCallEffectKind>("TCallEffectKind", static_cast<int64_t>(CallEffectKind::kOpaque));
OpDef("tirx.tvm_stack_alloca")
- .add_arg("dtype_str", "The data type name.")
- .add_arg("num", "The number of entries.")
+ .signature(sig::arg("dtype_str", "The data type name."),
+ sig::arg("num", "The number of entries."))
.set_attr<TScriptPrinterName>("TScriptPrinterName", ffi::String("tvm_stack_alloca"))
.set_attr<TIRxOpCategory>("TIRxOpCategory", ffi::String("builtin"))
.set_attr<TCallEffectKind>("TCallEffectKind", static_cast<int64_t>(CallEffectKind::kOpaque));
OpDef("tirx.tvm_stack_make_shape")
- .allow_extra_args()
+ .signature(sig::var_args("args"))
.set_attr<TScriptPrinterName>("TScriptPrinterName", ffi::String("tvm_stack_make_shape"))
.set_attr<TIRxOpCategory>("TIRxOpCategory", ffi::String("builtin"))
.set_attr<TCallEffectKind>("TCallEffectKind", static_cast<int64_t>(CallEffectKind::kOpaque));
OpDef("tirx.tvm_stack_make_array")
- .add_arg("data", "The input data.")
- .add_arg("shape", "The shape.")
- .add_arg("strides", "The strides.")
- .add_arg("ndim", "The number of dimensions.")
- .add_arg("arr_dtype", "The array data type.")
- .add_arg("elem_offset", "The element offset.")
+ .signature(sig::arg("data", "The input data."), sig::arg("shape", "The shape."),
+ sig::arg("strides", "The strides."), sig::arg("ndim", "The number of dimensions."),
+ sig::arg("arr_dtype", "The array data type."),
+ sig::arg("elem_offset", "The element offset."))
.set_attr<TScriptPrinterName>("TScriptPrinterName", ffi::String("tvm_stack_make_array"))
.set_attr<TIRxOpCategory>("TIRxOpCategory", ffi::String("builtin"))
.set_attr<TCallEffectKind>("TCallEffectKind", static_cast<int64_t>(CallEffectKind::kOpaque));
OpDef("tirx.tvm_call_packed")
- .add_arg("func_name", "The function name.")
- .allow_extra_args()
+ .signature(sig::arg("func_name", "The function name."), sig::var_args("args"))
.set_attr<TIRxOpCategory>("TIRxOpCategory", ffi::String("builtin"))
.set_attr<TCallEffectKind>("TCallEffectKind", static_cast<int64_t>(CallEffectKind::kOpaque))
.set_attr<TScriptPrinterName>("TScriptPrinterName", ffi::String("call_packed"));
OpDef("tirx.tensormap_encode_tiled")
- .add_arg("descriptor", "The descriptor.")
- .add_arg("data", "The input data.")
- .allow_extra_args()
+ .signature(sig::arg("descriptor", "The descriptor."), sig::arg("data", "The input data."),
+ sig::var_args("args"))
.set_attr<TScriptPrinterName>("TScriptPrinterName", ffi::String("tensormap_encode_tiled"))
.set_attr<TIRxOpCategory>("TIRxOpCategory", ffi::String("builtin"))
.set_attr<TCallEffectKind>("TCallEffectKind", static_cast<int64_t>(CallEffectKind::kOpaque));
OpDef("tirx.call_ffi_kernel")
- .add_arg("kernel", "The kernel.")
- .allow_extra_args()
+ .signature(sig::arg("kernel", "The kernel."), sig::var_args("args"))
.set_attr<TScriptPrinterName>("TScriptPrinterName", ffi::String("call_ffi_kernel"))
.set_attr<TIRxOpCategory>("TIRxOpCategory", ffi::String("builtin"))
.set_attr<TCallEffectKind>("TCallEffectKind", static_cast<int64_t>(CallEffectKind::kOpaque));
OpDef("tirx.tvm_call_cpacked")
- .add_arg("func_name", "The function name.")
- .allow_extra_args()
+ .signature(sig::arg("func_name", "The function name."), sig::var_args("args"))
.set_attr<TIRxOpCategory>("TIRxOpCategory", ffi::String("builtin"))
.set_attr<TCallEffectKind>("TCallEffectKind", static_cast<int64_t>(CallEffectKind::kOpaque))
.set_attr<TScriptPrinterName>("TScriptPrinterName", ffi::String("call_cpacked"));
OpDef("tirx.tvm_thread_invariant")
- .add_arg("cond", "The condition.")
+ .signature(sig::arg("cond", "The condition."))
.set_attr<TScriptPrinterName>("TScriptPrinterName", ffi::String("tvm_thread_invariant"))
.set_attr<TIRxOpCategory>("TIRxOpCategory", ffi::String("builtin"))
.set_attr<TCallEffectKind>("TCallEffectKind", static_cast<int64_t>(CallEffectKind::kPure));
OpDef("tirx.tvm_call_packed_lowered")
- .add_arg("func_name", "The function name.")
- .add_arg("args_stack", "The argument stack.")
- .add_arg("begin", "The start index.")
- .add_arg("end", "The end index.")
+ .signature(sig::arg("func_name", "The function name."),
+ sig::arg("args_stack", "The argument stack."),
+ sig::arg("begin", "The start index."), sig::arg("end", "The end index."))
.set_attr<TIRxOpCategory>("TIRxOpCategory", ffi::String("builtin"))
.set_attr<TCallEffectKind>("TCallEffectKind", static_cast<int64_t>(CallEffectKind::kOpaque))
.set_attr<TScriptPrinterName>("TScriptPrinterName", ffi::String("call_packed_lowered"));
OpDef("tirx.tvm_call_cpacked_lowered")
- .add_arg("func_name", "The function name.")
- .add_arg("args_stack", "The argument stack.")
- .add_arg("begin", "The start index.")
- .add_arg("end", "The end index.")
+ .signature(sig::arg("func_name", "The function name."),
+ sig::arg("args_stack", "The argument stack."),
+ sig::arg("begin", "The start index."), sig::arg("end", "The end index."))
.set_attr<TIRxOpCategory>("TIRxOpCategory", ffi::String("builtin"))
.set_attr<TCallEffectKind>("TCallEffectKind", static_cast<int64_t>(CallEffectKind::kOpaque))
.set_attr<TScriptPrinterName>("TScriptPrinterName", ffi::String("call_cpacked_lowered"));
@@ -441,8 +415,7 @@ TVM_FFI_STATIC_INIT_BLOCK() {
// TODO(tvm-team) revisit storage sync once we have a good memory hierachy structure.
OpDef("tirx.tvm_storage_sync")
- .add_arg("storage_scope", "The storage scope.")
- .allow_extra_args()
+ .signature(sig::arg("storage_scope", "The storage scope."), sig::var_args("args"))
.set_attr<TScriptPrinterName>("TScriptPrinterName", ffi::String("tvm_storage_sync"))
.set_attr<TIRxOpCategory>("TIRxOpCategory", ffi::String("builtin"))
.set_attr<TCallEffectKind>("TCallEffectKind", static_cast<int64_t>(CallEffectKind::kOpaque));
@@ -453,41 +426,33 @@ TVM_FFI_STATIC_INIT_BLOCK() {
.set_attr<TCallEffectKind>("TCallEffectKind", static_cast<int64_t>(CallEffectKind::kOpaque));
OpDef("tirx.tvm_warp_shuffle")
- .add_arg("mask", "The mask.")
- .add_arg("value", "The value to use.")
- .add_arg("warp_id", "The warp identifier.")
- .add_arg("width", "The width.")
- .add_arg("warp_size", "The number of threads per warp.")
+ .signature(sig::arg("mask", "The mask."), sig::arg("value", "The value to use."),
+ sig::arg("warp_id", "The warp identifier."), sig::arg("width", "The width."),
+ sig::arg("warp_size", "The number of threads per warp."))
.set_attr<TScriptPrinterName>("TScriptPrinterName", ffi::String("tvm_warp_shuffle"))
.set_attr<TIRxOpCategory>("TIRxOpCategory", ffi::String("builtin"))
.set_attr<TCallEffectKind>("TCallEffectKind", static_cast<int64_t>(CallEffectKind::kOpaque));
OpDef("tirx.tvm_warp_shuffle_up")
- .add_arg("mask", "The mask.")
- .add_arg("value", "The value to use.")
- .add_arg("offset", "The offset.")
- .add_arg("width", "The width.")
- .add_arg("warp_size", "The number of threads per warp.")
+ .signature(sig::arg("mask", "The mask."), sig::arg("value", "The value to use."),
+ sig::arg("offset", "The offset."), sig::arg("width", "The width."),
+ sig::arg("warp_size", "The number of threads per warp."))
.set_attr<TScriptPrinterName>("TScriptPrinterName", ffi::String("tvm_warp_shuffle_up"))
.set_attr<TIRxOpCategory>("TIRxOpCategory", ffi::String("builtin"))
.set_attr<TCallEffectKind>("TCallEffectKind", static_cast<int64_t>(CallEffectKind::kOpaque));
OpDef("tirx.tvm_warp_shuffle_down")
- .add_arg("mask", "The mask.")
- .add_arg("value", "The value to use.")
- .add_arg("offset", "The offset.")
- .add_arg("width", "The width.")
- .add_arg("warp_size", "The number of threads per warp.")
+ .signature(sig::arg("mask", "The mask."), sig::arg("value", "The value to use."),
+ sig::arg("offset", "The offset."), sig::arg("width", "The width."),
+ sig::arg("warp_size", "The number of threads per warp."))
.set_attr<TScriptPrinterName>("TScriptPrinterName", ffi::String("tvm_warp_shuffle_down"))
.set_attr<TIRxOpCategory>("TIRxOpCategory", ffi::String("builtin"))
.set_attr<TCallEffectKind>("TCallEffectKind", static_cast<int64_t>(CallEffectKind::kOpaque));
OpDef("tirx.tvm_warp_shuffle_xor")
- .add_arg("mask", "The mask.")
- .add_arg("value", "The value to use.")
- .add_arg("lane_mask", "The lane mask.")
- .add_arg("width", "The width.")
- .add_arg("warp_size", "The number of threads per warp.")
+ .signature(sig::arg("mask", "The mask."), sig::arg("value", "The value to use."),
+ sig::arg("lane_mask", "The lane mask."), sig::arg("width", "The width."),
+ sig::arg("warp_size", "The number of threads per warp."))
.set_attr<TScriptPrinterName>("TScriptPrinterName", ffi::String("tvm_warp_shuffle_xor"))
.set_attr<TIRxOpCategory>("TIRxOpCategory", ffi::String("builtin"))
.set_attr<TCallEffectKind>("TCallEffectKind", static_cast<int64_t>(CallEffectKind::kOpaque));
@@ -498,75 +463,63 @@ TVM_FFI_STATIC_INIT_BLOCK() {
.set_attr<TCallEffectKind>("TCallEffectKind", static_cast<int64_t>(CallEffectKind::kOpaque));
OpDef("tirx.tvm_thread_allreduce")
- .add_arg("size", "The size.")
- .allow_extra_args()
+ .signature(sig::arg("size", "The size."), sig::var_args("args"))
.set_attr<TScriptPrinterName>("TScriptPrinterName", ffi::String("tvm_thread_allreduce"))
.set_attr<TIRxOpCategory>("TIRxOpCategory", ffi::String("builtin"))
.set_attr<TCallEffectKind>("TCallEffectKind", static_cast<int64_t>(CallEffectKind::kOpaque));
OpDef("tirx.cooperative_tensor_fill")
- .add_arg("d", "The D operand.")
- .add_arg("index", "The index.")
- .add_arg("value", "The value to use.")
- .add_arg("rows", "The number of rows.")
- .add_arg("cols", "The number of columns.")
+ .signature(sig::arg("d", "The D operand."), sig::arg("index", "The index."),
+ sig::arg("value", "The value to use."), sig::arg("rows", "The number of rows."),
+ sig::arg("cols", "The number of columns."))
.set_attr<TScriptPrinterName>("TScriptPrinterName", ffi::String("cooperative_tensor_fill"))
.set_attr<TIRxOpCategory>("TIRxOpCategory", ffi::String("builtin"))
.set_attr<TCallEffectKind>("TCallEffectKind", static_cast<int64_t>(CallEffectKind::kOpaque));
OpDef("tirx.cooperative_tensor_load")
- .add_arg("d", "The D operand.")
- .add_arg("index", "The index.")
- .add_arg("ptr", "The pointer.")
- .add_arg("stride", "The stride.")
- .add_arg("rows", "The number of rows.")
- .add_arg("cols", "The number of columns.")
- .add_arg("transpose_matrix", "Whether to transpose the matrix.")
- .add_arg("mma_M", "The M dimension of the matrix operation.")
- .add_arg("mma_N", "The N dimension of the matrix operation.")
- .add_arg("mma_K", "The K dimension of the matrix operation.")
- .add_arg("operand_role", "The matrix operand role.")
+ .signature(sig::arg("d", "The D operand."), sig::arg("index", "The index."),
+ sig::arg("ptr", "The pointer."), sig::arg("stride", "The stride."),
+ sig::arg("rows", "The number of rows."),
+ sig::arg("cols", "The number of columns."),
+ sig::arg("transpose_matrix", "Whether to transpose the matrix."),
+ sig::arg("mma_M", "The M dimension of the matrix operation."),
+ sig::arg("mma_N", "The N dimension of the matrix operation."),
+ sig::arg("mma_K", "The K dimension of the matrix operation."),
+ sig::arg("operand_role", "The matrix operand role."))
.set_attr<TScriptPrinterName>("TScriptPrinterName", ffi::String("cooperative_tensor_load"))
.set_attr<TIRxOpCategory>("TIRxOpCategory", ffi::String("builtin"))
.set_attr<TCallEffectKind>("TCallEffectKind", static_cast<int64_t>(CallEffectKind::kOpaque));
OpDef("tirx.cooperative_tensor_store")
- .add_arg("d", "The D operand.")
- .add_arg("index", "The index.")
- .add_arg("ptr", "The pointer.")
- .add_arg("stride", "The stride.")
- .add_arg("rows", "The number of rows.")
- .add_arg("cols", "The number of columns.")
- .add_arg("transpose_matrix", "Whether to transpose the matrix.")
- .add_arg("mma_M", "The M dimension of the matrix operation.")
- .add_arg("mma_N", "The N dimension of the matrix operation.")
- .add_arg("mma_K", "The K dimension of the matrix operation.")
- .add_arg("operand_role", "The matrix operand role.")
+ .signature(sig::arg("d", "The D operand."), sig::arg("index", "The index."),
+ sig::arg("ptr", "The pointer."), sig::arg("stride", "The stride."),
+ sig::arg("rows", "The number of rows."),
+ sig::arg("cols", "The number of columns."),
+ sig::arg("transpose_matrix", "Whether to transpose the matrix."),
+ sig::arg("mma_M", "The M dimension of the matrix operation."),
+ sig::arg("mma_N", "The N dimension of the matrix operation."),
+ sig::arg("mma_K", "The K dimension of the matrix operation."),
+ sig::arg("operand_role", "The matrix operand role."))
.set_attr<TScriptPrinterName>("TScriptPrinterName", ffi::String("cooperative_tensor_store"))
.set_attr<TIRxOpCategory>("TIRxOpCategory", ffi::String("builtin"))
.set_attr<TCallEffectKind>("TCallEffectKind", static_cast<int64_t>(CallEffectKind::kOpaque));
OpDef("tirx.cooperative_tensor_multiply_accumulate")
- .add_arg("d", "The D operand.")
- .add_arg("index_d", "The D fragment index.")
- .add_arg("a", "The A operand.")
- .add_arg("index_a", "The A fragment index.")
- .add_arg("b", "The B operand.")
- .add_arg("index_b", "The B fragment index.")
- .add_arg("c", "The C operand.")
- .add_arg("index_c", "The C fragment index.")
- .add_arg("M", "The M dimension.")
- .add_arg("N", "The N dimension.")
- .add_arg("K", "The K dimension.")
- .add_arg("transpose_a", "Whether to transpose A.")
- .add_arg("transpose_b", "Whether to transpose B.")
+ .signature(sig::arg("d", "The D operand."), sig::arg("index_d", "The D fragment index."),
+ sig::arg("a", "The A operand."), sig::arg("index_a", "The A fragment index."),
+ sig::arg("b", "The B operand."), sig::arg("index_b", "The B fragment index."),
+ sig::arg("c", "The C operand."), sig::arg("index_c", "The C fragment index."),
+ sig::arg("M", "The M dimension."), sig::arg("N", "The N dimension."),
+ sig::arg("K", "The K dimension."),
+ sig::arg("transpose_a", "Whether to transpose A."),
+ sig::arg("transpose_b", "Whether to transpose B."))
.set_attr<TScriptPrinterName>("TScriptPrinterName",
ffi::String("cooperative_tensor_multiply_accumulate"))
.set_attr<TIRxOpCategory>("TIRxOpCategory", ffi::String("builtin"))
.set_attr<TCallEffectKind>("TCallEffectKind", static_cast<int64_t>(CallEffectKind::kOpaque));
OpDef("tirx.vectorhigh")
- .add_arg("vec", "The input vector.")
+ .signature(sig::arg("vec", "The input vector."))
.set_attr<TScriptPrinterName>("TScriptPrinterName", ffi::String("vectorhigh"))
.set_attr<TIRxOpCategory>("TIRxOpCategory", ffi::String("builtin"))
.set_attr<TCallEffectKind>("TCallEffectKind", static_cast<int64_t>(CallEffectKind::kPure))
@@ -574,7 +527,7 @@ TVM_FFI_STATIC_INIT_BLOCK() {
static_cast<int64_t>(ScriptDtypePrintLocation::kFirst));
OpDef("tirx.vectorlow")
- .add_arg("vec", "The input vector.")
+ .signature(sig::arg("vec", "The input vector."))
.set_attr<TScriptPrinterName>("TScriptPrinterName", ffi::String("vectorlow"))
.set_attr<TIRxOpCategory>("TIRxOpCategory", ffi::String("builtin"))
.set_attr<TCallEffectKind>("TCallEffectKind", static_cast<int64_t>(CallEffectKind::kPure))
@@ -582,8 +535,8 @@ TVM_FFI_STATIC_INIT_BLOCK() {
static_cast<int64_t>(ScriptDtypePrintLocation::kFirst));
OpDef("tirx.vectorcombine")
- .add_arg("vec1", "The first input vector.")
- .add_arg("vec2", "The second input vector.")
+ .signature(sig::arg("vec1", "The first input vector."),
+ sig::arg("vec2", "The second input vector."))
.set_attr<TScriptPrinterName>("TScriptPrinterName", ffi::String("vectorcombine"))
.set_attr<TIRxOpCategory>("TIRxOpCategory", ffi::String("builtin"))
.set_attr<TCallEffectKind>("TCallEffectKind", static_cast<int64_t>(CallEffectKind::kPure))
@@ -591,9 +544,8 @@ TVM_FFI_STATIC_INIT_BLOCK() {
static_cast<int64_t>(ScriptDtypePrintLocation::kFirst));
OpDef("tirx.dp4a")
- .add_arg("vec1", "The first input vector.")
- .add_arg("vec2", "The second input vector.")
- .add_arg("acc", "The accumulator.")
+ .signature(sig::arg("vec1", "The first input vector."),
+ sig::arg("vec2", "The second input vector."), sig::arg("acc", "The accumulator."))
.set_attr<TScriptPrinterName>("TScriptPrinterName", ffi::String("dp4a"))
.set_attr<TIRxOpCategory>("TIRxOpCategory", ffi::String("builtin"))
.set_attr<TCallEffectKind>("TCallEffectKind", static_cast<int64_t>(CallEffectKind::kPure))
@@ -601,46 +553,41 @@ TVM_FFI_STATIC_INIT_BLOCK() {
static_cast<int64_t>(ScriptDtypePrintLocation::kFirst));
OpDef("tirx.atomic_add")
- .add_arg("ptr", "The pointer.")
- .add_arg("value", "The value to use.")
+ .signature(sig::arg("ptr", "The pointer."), sig::arg("value", "The value to use."))
.set_attr<TScriptPrinterName>("TScriptPrinterName", ffi::String("atomic_add"))
.set_attr<TIRxOpCategory>("TIRxOpCategory", ffi::String("builtin"))
.set_attr<TCallEffectKind>("TCallEffectKind", static_cast<int64_t>(CallEffectKind::kOpaque));
OpDef("tirx.nd_mem_alloc_with_scope")
- .add_arg("storage_scope", "The storage scope.")
- .add_arg("ndim", "The number of dimensions.")
- .add_arg("shape", "The shape.")
- .allow_extra_args()
+ .signature(sig::arg("storage_scope", "The storage scope."),
+ sig::arg("ndim", "The number of dimensions."), sig::arg("shape", "The shape."),
+ sig::var_args("args"))
.set_attr<TScriptPrinterName>("TScriptPrinterName", ffi::String("nd_mem_alloc_with_scope"))
.set_attr<TIRxOpCategory>("TIRxOpCategory", ffi::String("builtin"))
.set_attr<TCallEffectKind>("TCallEffectKind", static_cast<int64_t>(CallEffectKind::kOpaque));
OpDef("tirx.texture2d_store")
- .add_arg("texture", "The texture.")
- .add_arg("x", "The input value.")
- .add_arg("y", "The second input value.")
- .add_arg("z", "The third input value.")
- .add_arg("channel_size", "The number of channels.")
- .add_arg("value", "The value to use.")
+ .signature(sig::arg("texture", "The texture."), sig::arg("x", "The input value."),
+ sig::arg("y", "The second input value."), sig::arg("z", "The third input value."),
+ sig::arg("channel_size", "The number of channels."),
+ sig::arg("value", "The value to use."))
.set_attr<TScriptPrinterName>("TScriptPrinterName", ffi::String("texture2d_store"))
.set_attr<TIRxOpCategory>("TIRxOpCategory", ffi::String("builtin"))
.set_attr<TVectorizable>("TVectorizable", true)
.set_attr<TCallEffectKind>("TCallEffectKind", static_cast<int64_t>(CallEffectKind::kOpaque));
OpDef("tirx.texture2d_load")
- .add_arg("texture", "The texture.")
- .add_arg("x", "The input value.")
- .add_arg("y", "The second input value.")
- .add_arg("z", "The third input value.")
- .add_arg("channel_size", "The number of channels.")
+ .signature(sig::arg("texture", "The texture."), sig::arg("x", "The input value."),
+ sig::arg("y", "The second input value."), sig::arg("z", "The third input value."),
+ sig::arg("channel_size", "The number of channels."),
+ sig::arg("element_index", "The element index within a texture channel."))
.set_attr<TScriptPrinterName>("TScriptPrinterName", ffi::String("texture2d_load"))
.set_attr<TIRxOpCategory>("TIRxOpCategory", ffi::String("builtin"))
.set_attr<TVectorizable>("TVectorizable", true)
.set_attr<TCallEffectKind>("TCallEffectKind", static_cast<int64_t>(CallEffectKind::kOpaque));
OpDef("tirx.assume")
- .add_arg("cond", "The condition.")
+ .signature(sig::arg("cond", "The condition."))
.set_attr<TScriptPrinterName>("TScriptPrinterName", ffi::String("assume"))
.set_attr<TIRxOpCategory>("TIRxOpCategory", ffi::String("builtin"))
.set_attr<TCallEffectKind>("TCallEffectKind",
@@ -653,8 +600,7 @@ TVM_FFI_STATIC_INIT_BLOCK() {
static_cast<int64_t>(CallEffectKind::kReadState));
OpDef("tirx.get_active_lane_mask")
- .add_arg("base", "The base value.")
- .add_arg("limit", "The limit value.")
+ .signature(sig::arg("base", "The base value."), sig::arg("limit", "The limit value."))
.set_attr<TScriptPrinterName>("TScriptPrinterName", ffi::String("get_active_lane_mask"))
.set_attr<TIRxOpCategory>("TIRxOpCategory", ffi::String("builtin"))
.set_attr<TCallEffectKind>("TCallEffectKind", static_cast<int64_t>(CallEffectKind::kPure))
@@ -662,9 +608,8 @@ TVM_FFI_STATIC_INIT_BLOCK() {
static_cast<int64_t>(ScriptDtypePrintLocation::kFirst));
OpDef("tirx.masked_load")
- .add_arg("buffer", "The buffer.")
- .add_arg("index", "The index.")
- .allow_extra_args()
+ .signature(sig::arg("buffer", "The buffer."), sig::arg("index", "The index."),
+ sig::var_args("args"))
.set_attr<TScriptPrinterName>("TScriptPrinterName", ffi::String("masked_load"))
.set_attr<TIRxOpCategory>("TIRxOpCategory", ffi::String("builtin"))
.set_attr<TCallEffectKind>("TCallEffectKind",
@@ -673,17 +618,15 @@ TVM_FFI_STATIC_INIT_BLOCK() {
static_cast<int64_t>(ScriptDtypePrintLocation::kFirst));
OpDef("tirx.masked_store")
- .add_arg("buffer", "The buffer.")
- .add_arg("value", "The value to use.")
- .add_arg("index", "The index.")
- .allow_extra_args()
+ .signature(sig::arg("buffer", "The buffer."), sig::arg("value", "The value to use."),
+ sig::arg("index", "The index."), sig::var_args("args"))
.set_attr<TScriptPrinterName>("TScriptPrinterName", ffi::String("masked_store"))
.set_attr<TIRxOpCategory>("TIRxOpCategory", ffi::String("builtin"))
.set_attr<TCallEffectKind>("TCallEffectKind",
static_cast<int64_t>(CallEffectKind::kUpdateState));
OpDef("tirx.ignore_loop_partition")
- .add_arg("predicate", "The predicate.")
+ .signature(sig::arg("predicate", "The predicate."))
.set_attr<TScriptPrinterName>("TScriptPrinterName", ffi::String("ignore_loop_partition"))
.set_attr<TIRxOpCategory>("TIRxOpCategory", ffi::String("builtin"))
.set_attr<TCallEffectKind>("TCallEffectKind", static_cast<int64_t>(CallEffectKind::kPure))
@@ -691,24 +634,22 @@ TVM_FFI_STATIC_INIT_BLOCK() {
static_cast<int64_t>(ScriptDtypePrintLocation::kNone));
OpDef("tirx.buffer_offset")
- .add_arg("load", "The buffer load whose offset is returned.")
+ .signature(sig::arg("load", "The buffer load whose offset is returned."))
.set_attr<TScriptPrinterName>("TScriptPrinterName", ffi::String("buffer_offset"))
.set_attr<TIRxOpCategory>("TIRxOpCategory", ffi::String("builtin"))
.set_attr<TCallEffectKind>("TCallEffectKind", static_cast<int64_t>(CallEffectKind::kPure));
OpDef("tirx.buffer_data")
- .add_arg("buffer", "The buffer.")
+ .signature(sig::arg("buffer", "The buffer."))
.set_attr<TScriptPrinterName>("TScriptPrinterName", ffi::String("buffer_data"))
.set_attr<TIRxOpCategory>("TIRxOpCategory", ffi::String("builtin"))
.set_attr<TCallEffectKind>("TCallEffectKind", static_cast<int64_t>(CallEffectKind::kPure));
OpDef("tirx.print_buffer")
- .add_arg("data", "The input data.")
- .add_arg("dtype", "The data type.")
- .add_arg("is_string", "Whether to print as a string.")
- .add_arg("is_scalar", "Whether to print as a scalar.")
- .add_arg("ndim", "The number of dimensions.")
- .allow_extra_args()
+ .signature(sig::arg("data", "The input data."), sig::arg("dtype", "The data type."),
+ sig::arg("is_string", "Whether to print as a string."),
+ sig::arg("is_scalar", "Whether to print as a scalar."),
+ sig::arg("ndim", "The number of dimensions."), sig::var_args("args"))
.set_attr<TScriptPrinterName>("TScriptPrinterName", ffi::String("print_buffer"))
.set_attr<TIRxOpCategory>("TIRxOpCategory", ffi::String("builtin"))
.set_attr<TCallEffectKind>("TCallEffectKind", static_cast<int64_t>(CallEffectKind::kOpaque));
diff --git a/src/tirx/op/op.cc b/src/tirx/op/op.cc
index 27391b9f9a..65a1a48a04 100644
--- a/src/tirx/op/op.cc
+++ b/src/tirx/op/op.cc
@@ -258,8 +258,7 @@ PrimExpr pow(PrimExpr x, PrimExpr y, Span span) {
TVM_FFI_STATIC_INIT_BLOCK() {
OpDef("tirx.pow")
- .add_arg("x", "The input value.")
- .add_arg("y", "The second input value.")
+ .signature(sig::arg("x", "The input value."), sig::arg("y", "The second input value."))
.set_attr<TScriptPrinterName>("TScriptPrinterName", ffi::String("pow"))
.set_attr<TIRxOpCategory>("TIRxOpCategory", ffi::String("builtin"))
.set_attr<TCallEffectKind>("TCallEffectKind", static_cast<int64_t>(CallEffectKind::kPure))
@@ -288,7 +287,7 @@ PrimExpr abs(PrimExpr x, Span span) {
TVM_FFI_STATIC_INIT_BLOCK() {
OpDef("tirx.fabs")
- .add_arg("x", "The input value.")
+ .signature(sig::arg("x", "The input value."))
.set_attr<TScriptPrinterName>("TScriptPrinterName", ffi::String("fabs"))
.set_attr<TIRxOpCategory>("TIRxOpCategory", ffi::String("builtin"))
.set_attr<TCallEffectKind>("TCallEffectKind", static_cast<int64_t>(CallEffectKind::kPure))
@@ -413,8 +412,7 @@ PrimExpr fmod(PrimExpr x, PrimExpr y, Span span) {
TVM_FFI_STATIC_INIT_BLOCK() {
OpDef("tirx.fmod")
- .add_arg("x", "The input value.")
- .add_arg("y", "The second input value.")
+ .signature(sig::arg("x", "The input value."), sig::arg("y", "The second input value."))
.set_attr<TScriptPrinterName>("TScriptPrinterName", ffi::String("fmod"))
.set_attr<TIRxOpCategory>("TIRxOpCategory", ffi::String("builtin"))
.set_attr<TCallEffectKind>("TCallEffectKind", static_cast<int64_t>(CallEffectKind::kPure));
@@ -434,7 +432,7 @@ PrimExpr floor(PrimExpr x, Span span) {
TVM_FFI_STATIC_INIT_BLOCK() {
OpDef("tirx.floor")
- .add_arg("x", "The input value.")
+ .signature(sig::arg("x", "The input value."))
.set_attr<TScriptPrinterName>("TScriptPrinterName", ffi::String("floor"))
.set_attr<TIRxOpCategory>("TIRxOpCategory", ffi::String("builtin"))
.set_attr<TCallEffectKind>("TCallEffectKind", static_cast<int64_t>(CallEffectKind::kPure))
@@ -455,7 +453,7 @@ PrimExpr round(PrimExpr x, Span span) {
TVM_FFI_STATIC_INIT_BLOCK() {
OpDef("tirx.round")
- .add_arg("x", "The input value.")
+ .signature(sig::arg("x", "The input value."))
.set_attr<TScriptPrinterName>("TScriptPrinterName", ffi::String("round"))
.set_attr<TIRxOpCategory>("TIRxOpCategory", ffi::String("builtin"))
.set_attr<TCallEffectKind>("TCallEffectKind", static_cast<int64_t>(CallEffectKind::kPure))
@@ -476,7 +474,7 @@ PrimExpr nearbyint(PrimExpr x, Span span) {
TVM_FFI_STATIC_INIT_BLOCK() {
OpDef("tirx.nearbyint")
- .add_arg("x", "The input value.")
+ .signature(sig::arg("x", "The input value."))
.set_attr<TScriptPrinterName>("TScriptPrinterName", ffi::String("nearbyint"))
.set_attr<TIRxOpCategory>("TIRxOpCategory", ffi::String("builtin"))
.set_attr<TCallEffectKind>("TCallEffectKind", static_cast<int64_t>(CallEffectKind::kPure));
@@ -499,7 +497,7 @@ PrimExpr trunc(PrimExpr x, Span span) {
TVM_FFI_STATIC_INIT_BLOCK() {
OpDef("tirx.trunc")
- .add_arg("x", "The input value.")
+ .signature(sig::arg("x", "The input value."))
.set_attr<TScriptPrinterName>("TScriptPrinterName", ffi::String("trunc"))
.set_attr<TIRxOpCategory>("TIRxOpCategory", ffi::String("builtin"))
.set_attr<TCallEffectKind>("TCallEffectKind", static_cast<int64_t>(CallEffectKind::kPure))
@@ -508,146 +506,146 @@ TVM_FFI_STATIC_INIT_BLOCK() {
// unary op registration.
OpDef("tirx.exp")
- .add_arg("x", "The input value.")
+ .signature(sig::arg("x", "The input value."))
.set_attr<TScriptPrinterName>("TScriptPrinterName", ffi::String("exp"))
.set_attr<TIRxOpCategory>("TIRxOpCategory", ffi::String("builtin"))
.set_attr<TCallEffectKind>("TCallEffectKind", static_cast<int64_t>(CallEffectKind::kPure))
.set_attr<TVectorizable>("TVectorizable", true);
OpDef("tirx.exp2")
- .add_arg("x", "The input value.")
+ .signature(sig::arg("x", "The input value."))
.set_attr<TScriptPrinterName>("TScriptPrinterName", ffi::String("exp2"))
.set_attr<TIRxOpCategory>("TIRxOpCategory", ffi::String("builtin"))
.set_attr<TCallEffectKind>("TCallEffectKind", static_cast<int64_t>(CallEffectKind::kPure))
.set_attr<TVectorizable>("TVectorizable", true);
OpDef("tirx.exp10")
- .add_arg("x", "The input value.")
+ .signature(sig::arg("x", "The input value."))
.set_attr<TScriptPrinterName>("TScriptPrinterName", ffi::String("exp10"))
.set_attr<TIRxOpCategory>("TIRxOpCategory", ffi::String("builtin"))
.set_attr<TCallEffectKind>("TCallEffectKind", static_cast<int64_t>(CallEffectKind::kPure))
.set_attr<TVectorizable>("TVectorizable", true);
OpDef("tirx.erf")
- .add_arg("x", "The input value.")
+ .signature(sig::arg("x", "The input value."))
.set_attr<TScriptPrinterName>("TScriptPrinterName", ffi::String("erf"))
.set_attr<TIRxOpCategory>("TIRxOpCategory", ffi::String("builtin"))
.set_attr<TCallEffectKind>("TCallEffectKind", static_cast<int64_t>(CallEffectKind::kPure));
OpDef("tirx.tanh")
- .add_arg("x", "The input value.")
+ .signature(sig::arg("x", "The input value."))
.set_attr<TScriptPrinterName>("TScriptPrinterName", ffi::String("tanh"))
.set_attr<TIRxOpCategory>("TIRxOpCategory", ffi::String("builtin"))
.set_attr<TCallEffectKind>("TCallEffectKind", static_cast<int64_t>(CallEffectKind::kPure))
.set_attr<TVectorizable>("TVectorizable", true);
OpDef("tirx.sigmoid")
- .add_arg("x", "The input value.")
+ .signature(sig::arg("x", "The input value."))
.set_attr<TScriptPrinterName>("TScriptPrinterName", ffi::String("sigmoid"))
.set_attr<TIRxOpCategory>("TIRxOpCategory", ffi::String("builtin"))
.set_attr<TCallEffectKind>("TCallEffectKind", static_cast<int64_t>(CallEffectKind::kPure))
.set_attr<TVectorizable>("TVectorizable", true);
OpDef("tirx.sqrt")
- .add_arg("x", "The input value.")
+ .signature(sig::arg("x", "The input value."))
.set_attr<TScriptPrinterName>("TScriptPrinterName", ffi::String("sqrt"))
.set_attr<TIRxOpCategory>("TIRxOpCategory", ffi::String("builtin"))
.set_attr<TCallEffectKind>("TCallEffectKind", static_cast<int64_t>(CallEffectKind::kPure))
.set_attr<TVectorizable>("TVectorizable", true);
OpDef("tirx.rsqrt")
- .add_arg("x", "The input value.")
+ .signature(sig::arg("x", "The input value."))
.set_attr<TScriptPrinterName>("TScriptPrinterName", ffi::String("rsqrt"))
.set_attr<TIRxOpCategory>("TIRxOpCategory", ffi::String("builtin"))
.set_attr<TCallEffectKind>("TCallEffectKind", static_cast<int64_t>(CallEffectKind::kPure));
OpDef("tirx.log")
- .add_arg("x", "The input value.")
+ .signature(sig::arg("x", "The input value."))
.set_attr<TScriptPrinterName>("TScriptPrinterName", ffi::String("log"))
.set_attr<TIRxOpCategory>("TIRxOpCategory", ffi::String("builtin"))
.set_attr<TCallEffectKind>("TCallEffectKind", static_cast<int64_t>(CallEffectKind::kPure))
.set_attr<TVectorizable>("TVectorizable", true);
OpDef("tirx.log1p")
- .add_arg("x", "The input value.")
+ .signature(sig::arg("x", "The input value."))
.set_attr<TScriptPrinterName>("TScriptPrinterName", ffi::String("log1p"))
.set_attr<TIRxOpCategory>("TIRxOpCategory", ffi::String("builtin"))
.set_attr<TCallEffectKind>("TCallEffectKind", static_cast<int64_t>(CallEffectKind::kPure));
OpDef("tirx.log10")
- .add_arg("x", "The input value.")
+ .signature(sig::arg("x", "The input value."))
.set_attr<TScriptPrinterName>("TScriptPrinterName", ffi::String("log10"))
.set_attr<TIRxOpCategory>("TIRxOpCategory", ffi::String("builtin"))
.set_attr<TCallEffectKind>("TCallEffectKind", static_cast<int64_t>(CallEffectKind::kPure))
.set_attr<TVectorizable>("TVectorizable", true);
OpDef("tirx.tan")
- .add_arg("x", "The input value.")
+ .signature(sig::arg("x", "The input value."))
.set_attr<TScriptPrinterName>("TScriptPrinterName", ffi::String("tan"))
.set_attr<TIRxOpCategory>("TIRxOpCategory", ffi::String("builtin"))
.set_attr<TCallEffectKind>("TCallEffectKind", static_cast<int64_t>(CallEffectKind::kPure))
.set_attr<TVectorizable>("TVectorizable", true);
OpDef("tirx.cos")
- .add_arg("x", "The input value.")
+ .signature(sig::arg("x", "The input value."))
.set_attr<TScriptPrinterName>("TScriptPrinterName", ffi::String("cos"))
.set_attr<TIRxOpCategory>("TIRxOpCategory", ffi::String("builtin"))
.set_attr<TCallEffectKind>("TCallEffectKind", static_cast<int64_t>(CallEffectKind::kPure))
.set_attr<TVectorizable>("TVectorizable", true);
OpDef("tirx.cosh")
- .add_arg("x", "The input value.")
+ .signature(sig::arg("x", "The input value."))
.set_attr<TScriptPrinterName>("TScriptPrinterName", ffi::String("cosh"))
.set_attr<TIRxOpCategory>("TIRxOpCategory", ffi::String("builtin"))
.set_attr<TCallEffectKind>("TCallEffectKind", static_cast<int64_t>(CallEffectKind::kPure))
.set_attr<TVectorizable>("TVectorizable", true);
OpDef("tirx.sin")
- .add_arg("x", "The input value.")
+ .signature(sig::arg("x", "The input value."))
.set_attr<TScriptPrinterName>("TScriptPrinterName", ffi::String("sin"))
.set_attr<TIRxOpCategory>("TIRxOpCategory", ffi::String("builtin"))
.set_attr<TCallEffectKind>("TCallEffectKind", static_cast<int64_t>(CallEffectKind::kPure))
.set_attr<TVectorizable>("TVectorizable", true);
OpDef("tirx.sinh")
- .add_arg("x", "The input value.")
+ .signature(sig::arg("x", "The input value."))
.set_attr<TScriptPrinterName>("TScriptPrinterName", ffi::String("sinh"))
.set_attr<TIRxOpCategory>("TIRxOpCategory", ffi::String("builtin"))
.set_attr<TCallEffectKind>("TCallEffectKind", static_cast<int64_t>(CallEffectKind::kPure))
.set_attr<TVectorizable>("TVectorizable", true);
OpDef("tirx.asin")
- .add_arg("x", "The input value.")
+ .signature(sig::arg("x", "The input value."))
.set_attr<TScriptPrinterName>("TScriptPrinterName", ffi::String("asin"))
.set_attr<TIRxOpCategory>("TIRxOpCategory", ffi::String("builtin"))
.set_attr<TCallEffectKind>("TCallEffectKind", static_cast<int64_t>(CallEffectKind::kPure));
OpDef("tirx.acos")
- .add_arg("x", "The input value.")
+ .signature(sig::arg("x", "The input value."))
.set_attr<TScriptPrinterName>("TScriptPrinterName", ffi::String("acos"))
.set_attr<TIRxOpCategory>("TIRxOpCategory", ffi::String("builtin"))
.set_attr<TCallEffectKind>("TCallEffectKind", static_cast<int64_t>(CallEffectKind::kPure));
OpDef("tirx.atan")
- .add_arg("x", "The input value.")
+ .signature(sig::arg("x", "The input value."))
.set_attr<TScriptPrinterName>("TScriptPrinterName", ffi::String("atan"))
.set_attr<TIRxOpCategory>("TIRxOpCategory", ffi::String("builtin"))
.set_attr<TCallEffectKind>("TCallEffectKind", static_cast<int64_t>(CallEffectKind::kPure));
OpDef("tirx.acosh")
- .add_arg("x", "The input value.")
+ .signature(sig::arg("x", "The input value."))
.set_attr<TScriptPrinterName>("TScriptPrinterName", ffi::String("acosh"))
.set_attr<TIRxOpCategory>("TIRxOpCategory", ffi::String("builtin"))
.set_attr<TCallEffectKind>("TCallEffectKind", static_cast<int64_t>(CallEffectKind::kPure));
OpDef("tirx.asinh")
- .add_arg("x", "The input value.")
+ .signature(sig::arg("x", "The input value."))
.set_attr<TScriptPrinterName>("TScriptPrinterName", ffi::String("asinh"))
.set_attr<TIRxOpCategory>("TIRxOpCategory", ffi::String("builtin"))
.set_attr<TCallEffectKind>("TCallEffectKind", static_cast<int64_t>(CallEffectKind::kPure));
OpDef("tirx.atanh")
- .add_arg("x", "The input value.")
+ .signature(sig::arg("x", "The input value."))
.set_attr<TScriptPrinterName>("TScriptPrinterName", ffi::String("atanh"))
.set_attr<TIRxOpCategory>("TIRxOpCategory", ffi::String("builtin"))
.set_attr<TCallEffectKind>("TCallEffectKind", static_cast<int64_t>(CallEffectKind::kPure));
@@ -655,55 +653,54 @@ TVM_FFI_STATIC_INIT_BLOCK() {
// binary intrinsics
OpDef("tirx.atan2")
- .add_arg("x1", "The first input value.")
- .add_arg("x2", "The second input value.")
+ .signature(sig::arg("x1", "The first input value."),
+ sig::arg("x2", "The second input value."))
.set_attr<TScriptPrinterName>("TScriptPrinterName", ffi::String("atan2"))
.set_attr<TIRxOpCategory>("TIRxOpCategory", ffi::String("builtin"))
.set_attr<TCallEffectKind>("TCallEffectKind", static_cast<int64_t>(CallEffectKind::kPure));
OpDef("tirx.nextafter")
- .add_arg("x1", "The first input value.")
- .add_arg("x2", "The second input value.")
+ .signature(sig::arg("x1", "The first input value."),
+ sig::arg("x2", "The second input value."))
.set_attr<TScriptPrinterName>("TScriptPrinterName", ffi::String("nextafter"))
.set_attr<TIRxOpCategory>("TIRxOpCategory", ffi::String("builtin"))
.set_attr<TCallEffectKind>("TCallEffectKind", static_cast<int64_t>(CallEffectKind::kPure));
OpDef("tirx.hypot")
- .add_arg("x1", "The first input value.")
- .add_arg("x2", "The second input value.")
+ .signature(sig::arg("x1", "The first input value."),
+ sig::arg("x2", "The second input value."))
.set_attr<TScriptPrinterName>("TScriptPrinterName", ffi::String("hypot"))
.set_attr<TIRxOpCategory>("TIRxOpCategory", ffi::String("builtin"))
.set_attr<TCallEffectKind>("TCallEffectKind", static_cast<int64_t>(CallEffectKind::kPure));
OpDef("tirx.copysign")
- .add_arg("x1", "The first input value.")
- .add_arg("x2", "The second input value.")
+ .signature(sig::arg("x1", "The first input value."),
+ sig::arg("x2", "The second input value."))
.set_attr<TScriptPrinterName>("TScriptPrinterName", ffi::String("copysign"))
.set_attr<TIRxOpCategory>("TIRxOpCategory", ffi::String("builtin"))
.set_attr<TCallEffectKind>("TCallEffectKind", static_cast<int64_t>(CallEffectKind::kPure));
OpDef("tirx.ldexp")
- .add_arg("x1", "The first input value.")
- .add_arg("x2", "The second input value.")
+ .signature(sig::arg("x1", "The first input value."),
+ sig::arg("x2", "The second input value."))
.set_attr<TScriptPrinterName>("TScriptPrinterName", ffi::String("ldexp"))
.set_attr<TIRxOpCategory>("TIRxOpCategory", ffi::String("builtin"))
.set_attr<TCallEffectKind>("TCallEffectKind", static_cast<int64_t>(CallEffectKind::kPure));
OpDef("tirx.TVMBackendAllocWorkspace")
- .add_arg("device_type", "The device type.")
- .add_arg("device_id", "The device index.")
- .add_arg("nbytes", "The number of bytes.")
- .add_arg("dtype_code_hint", "The data type code hint.")
- .add_arg("dtype_bits_hint", "The data type bit-width hint.")
+ .signature(sig::arg("device_type", "The device type."),
+ sig::arg("device_id", "The device index."),
+ sig::arg("nbytes", "The number of bytes."),
+ sig::arg("dtype_code_hint", "The data type code hint."),
+ sig::arg("dtype_bits_hint", "The data type bit-width hint."))
.set_attr<TScriptPrinterName>("TScriptPrinterName", ffi::String("TVMBackendAllocWorkspace"))
.set_attr<TIRxOpCategory>("TIRxOpCategory", ffi::String("builtin"))
.set_attr<TGlobalSymbol>("TGlobalSymbol", "TVMBackendAllocWorkspace")
.set_attr<TCallEffectKind>("TCallEffectKind", static_cast<int64_t>(CallEffectKind::kOpaque));
OpDef("tirx.TVMBackendFreeWorkspace")
- .add_arg("device_type", "The device type.")
- .add_arg("device_id", "The device index.")
- .add_arg("ptr", "The pointer.")
+ .signature(sig::arg("device_type", "The device type."),
+ sig::arg("device_id", "The device index."), sig::arg("ptr", "The pointer."))
.set_attr<TScriptPrinterName>("TScriptPrinterName", ffi::String("TVMBackendFreeWorkspace"))
.set_attr<TIRxOpCategory>("TIRxOpCategory", ffi::String("builtin"))
.set_attr<TGlobalSymbol>("TGlobalSymbol", "TVMBackendFreeWorkspace")
diff --git a/src/tirx/op/tirx.cc b/src/tirx/op/tirx.cc
index bc47e32907..6ce5cf088c 100644
--- a/src/tirx/op/tirx.cc
+++ b/src/tirx/op/tirx.cc
@@ -171,229 +171,193 @@ TVM_DEFINE_CACHED_OP_GETTER(gemm_async, "tirx.tile.gemm_async")
TVM_FFI_STATIC_INIT_BLOCK() {
OpDef("tirx.tile.zero")
- .add_arg("dst", "The destination.")
- .add_arg("src", "The source.")
+ .signature(sig::arg("dst", "The destination."), sig::arg("src", "The source."))
.set_attr<TScriptPrinterName>("TScriptPrinterName", ffi::String("zero"))
.set_attr<TIRxOpCategory>("TIRxOpCategory", ffi::String("tile_primitive"));
OpDef("tirx.tile.sqrt")
- .add_arg("dst", "The destination.")
- .add_arg("src", "The source.")
- .allow_extra_args()
+ .signature(sig::arg("dst", "The destination."), sig::arg("src", "The source."),
+ sig::var_args("args"))
.set_attr<TScriptPrinterName>("TScriptPrinterName", ffi::String("sqrt"))
.set_attr<TIRxOpCategory>("TIRxOpCategory", ffi::String("tile_primitive"));
OpDef("tirx.tile.exp")
- .add_arg("dst", "The destination.")
- .add_arg("src", "The source.")
- .allow_extra_args()
+ .signature(sig::arg("dst", "The destination."), sig::arg("src", "The source."),
+ sig::var_args("args"))
.set_attr<TScriptPrinterName>("TScriptPrinterName", ffi::String("exp"))
.set_attr<TIRxOpCategory>("TIRxOpCategory", ffi::String("tile_primitive"));
OpDef("tirx.tile.exp2")
- .add_arg("dst", "The destination.")
- .add_arg("src", "The source.")
- .allow_extra_args()
+ .signature(sig::arg("dst", "The destination."), sig::arg("src", "The source."),
+ sig::var_args("args"))
.set_attr<TScriptPrinterName>("TScriptPrinterName", ffi::String("exp2"))
.set_attr<TIRxOpCategory>("TIRxOpCategory", ffi::String("tile_primitive"));
OpDef("tirx.tile.log2")
- .add_arg("dst", "The destination.")
- .add_arg("src", "The source.")
- .allow_extra_args()
+ .signature(sig::arg("dst", "The destination."), sig::arg("src", "The source."),
+ sig::var_args("args"))
.set_attr<TScriptPrinterName>("TScriptPrinterName", ffi::String("log2"))
.set_attr<TIRxOpCategory>("TIRxOpCategory", ffi::String("tile_primitive"));
OpDef("tirx.tile.add")
- .add_arg("dst", "The destination.")
- .add_arg("src1", "The first source.")
- .add_arg("src2", "The second source.")
+ .signature(sig::arg("dst", "The destination."), sig::arg("src1", "The first source."),
+ sig::arg("src2", "The second source."))
.set_attr<TScriptPrinterName>("TScriptPrinterName", ffi::String("add"))
.set_attr<TIRxOpCategory>("TIRxOpCategory", ffi::String("tile_primitive"));
OpDef("tirx.tile.sub")
- .add_arg("dst", "The destination.")
- .add_arg("src1", "The first source.")
- .add_arg("src2", "The second source.")
+ .signature(sig::arg("dst", "The destination."), sig::arg("src1", "The first source."),
+ sig::arg("src2", "The second source."))
.set_attr<TScriptPrinterName>("TScriptPrinterName", ffi::String("sub"))
.set_attr<TIRxOpCategory>("TIRxOpCategory", ffi::String("tile_primitive"));
OpDef("tirx.tile.mul")
- .add_arg("dst", "The destination.")
- .add_arg("src1", "The first source.")
- .add_arg("src2", "The second source.")
+ .signature(sig::arg("dst", "The destination."), sig::arg("src1", "The first source."),
+ sig::arg("src2", "The second source."))
.set_attr<TScriptPrinterName>("TScriptPrinterName", ffi::String("mul"))
.set_attr<TIRxOpCategory>("TIRxOpCategory", ffi::String("tile_primitive"));
OpDef("tirx.tile.fdiv")
- .add_arg("dst", "The destination.")
- .add_arg("src1", "The first source.")
- .add_arg("src2", "The second source.")
+ .signature(sig::arg("dst", "The destination."), sig::arg("src1", "The first source."),
+ sig::arg("src2", "The second source."))
.set_attr<TScriptPrinterName>("TScriptPrinterName", ffi::String("fdiv"))
.set_attr<TIRxOpCategory>("TIRxOpCategory", ffi::String("tile_primitive"));
OpDef("tirx.tile.minimum")
- .add_arg("dst", "The destination.")
- .add_arg("src1", "The first source.")
- .add_arg("src2", "The second source.")
+ .signature(sig::arg("dst", "The destination."), sig::arg("src1", "The first source."),
+ sig::arg("src2", "The second source."))
.set_attr<TScriptPrinterName>("TScriptPrinterName", ffi::String("minimum"))
.set_attr<TIRxOpCategory>("TIRxOpCategory", ffi::String("tile_primitive"));
OpDef("tirx.tile.maximum")
- .add_arg("dst", "The destination.")
- .add_arg("src1", "The first source.")
- .add_arg("src2", "The second source.")
+ .signature(sig::arg("dst", "The destination."), sig::arg("src1", "The first source."),
+ sig::arg("src2", "The second source."))
.set_attr<TScriptPrinterName>("TScriptPrinterName", ffi::String("maximum"))
.set_attr<TIRxOpCategory>("TIRxOpCategory", ffi::String("tile_primitive"));
OpDef("tirx.tile.copy")
- .add_arg("dst", "The destination.")
- .add_arg("src", "The source.")
+ .signature(sig::arg("dst", "The destination."), sig::arg("src", "The source."))
.set_attr<TScriptPrinterName>("TScriptPrinterName", ffi::String("copy"))
.set_attr<TIRxOpCategory>("TIRxOpCategory", ffi::String("tile_primitive"));
OpDef("tirx.tile.fill")
- .add_arg("dst", "The destination.")
- .add_arg("value", "The value to use.")
+ .signature(sig::arg("dst", "The destination."), sig::arg("value", "The value to use."))
.set_attr<TScriptPrinterName>("TScriptPrinterName", ffi::String("fill"))
.set_attr<TIRxOpCategory>("TIRxOpCategory", ffi::String("tile_primitive"));
OpDef("tirx.tile.gemm")
- .add_arg("D", "The D tile.")
- .add_arg("A", "The A tile.")
- .add_arg("B", "The B tile.")
- .add_arg("C", "The C tile.")
- .add_arg("transpose_A", "Whether to transpose A.")
- .add_arg("transpose_B", "Whether to transpose B.")
- .add_arg("alpha", "The alpha scale factor.")
- .add_arg("beta", "The beta scale factor.")
+ .signature(
+ sig::arg("D", "The D tile."), sig::arg("A", "The A tile."), sig::arg("B", "The B tile."),
+ sig::arg("C", "The C tile."), sig::arg("transpose_A", "Whether to transpose A."),
+ sig::arg("transpose_B", "Whether to transpose B."),
+ sig::arg("alpha", "The alpha scale factor."), sig::arg("beta", "The beta scale factor."))
.set_attr<TScriptPrinterName>("TScriptPrinterName", ffi::String("gemm"))
.set_attr<TIRxOpCategory>("TIRxOpCategory", ffi::String("tile_primitive"));
OpDef("tirx.tile.reciprocal")
- .add_arg("dst", "The destination.")
- .add_arg("src", "The source.")
+ .signature(sig::arg("dst", "The destination."), sig::arg("src", "The source."))
.set_attr<TScriptPrinterName>("TScriptPrinterName", ffi::String("reciprocal"))
.set_attr<TIRxOpCategory>("TIRxOpCategory", ffi::String("tile_primitive"));
OpDef("tirx.tile.sum")
- .add_arg("dst", "The destination.")
- .add_arg("src", "The source.")
- .add_arg("axes", "The axes.")
- .add_arg("accum", "The accumulator.")
+ .signature(sig::arg("dst", "The destination."), sig::arg("src", "The source."),
+ sig::arg("axes", "The axes."), sig::arg("accum", "The accumulator."))
.set_attr<TScriptPrinterName>("TScriptPrinterName", ffi::String("sum"))
.set_attr<TIRxOpCategory>("TIRxOpCategory", ffi::String("tile_primitive"));
OpDef("tirx.tile.max")
- .add_arg("dst", "The destination.")
- .add_arg("src", "The source.")
- .add_arg("axes", "The axes.")
- .add_arg("accum", "The accumulator.")
+ .signature(sig::arg("dst", "The destination."), sig::arg("src", "The source."),
+ sig::arg("axes", "The axes."), sig::arg("accum", "The accumulator."))
.set_attr<TScriptPrinterName>("TScriptPrinterName", ffi::String("max"))
.set_attr<TIRxOpCategory>("TIRxOpCategory", ffi::String("tile_primitive"));
OpDef("tirx.tile.min")
- .add_arg("dst", "The destination.")
- .add_arg("src", "The source.")
- .add_arg("axes", "The axes.")
- .add_arg("accum", "The accumulator.")
+ .signature(sig::arg("dst", "The destination."), sig::arg("src", "The source."),
+ sig::arg("axes", "The axes."), sig::arg("accum", "The accumulator."))
.set_attr<TScriptPrinterName>("TScriptPrinterName", ffi::String("min"))
.set_attr<TIRxOpCategory>("TIRxOpCategory", ffi::String("tile_primitive"));
OpDef("tirx.tile.memset")
- .add_arg("dst", "The destination.")
- .add_arg("value", "The value to use.")
+ .signature(sig::arg("dst", "The destination."), sig::arg("value", "The value to use."))
.set_attr<TScriptPrinterName>("TScriptPrinterName", ffi::String("memset"))
.set_attr<TIRxOpCategory>("TIRxOpCategory", ffi::String("tile_primitive"));
OpDef("tirx.tile.reduce_negate")
- .add_arg("output", "The output.")
- .add_arg("input", "The input.")
- .add_arg("reduce_axes", "The reduction axes.")
- .add_arg("accum", "The accumulator.")
- .add_arg("reduce_op", "The reduction operation.")
+ .signature(sig::arg("output", "The output."), sig::arg("input", "The input."),
+ sig::arg("reduce_axes", "The reduction axes."),
+ sig::arg("accum", "The accumulator."),
+ sig::arg("reduce_op", "The reduction operation."))
.set_attr<TScriptPrinterName>("TScriptPrinterName", ffi::String("reduce_negate"))
.set_attr<TIRxOpCategory>("TIRxOpCategory", ffi::String("tile_primitive"));
OpDef("tirx.tile.binary_reduce")
- .add_arg("binary_output", "The binary operation output.")
- .add_arg("reduce_output", "The reduction output.")
- .add_arg("binary_input1", "The first binary input.")
- .add_arg("binary_input2", "The second binary input.")
- .add_arg("binary_op", "The binary operation.")
- .add_arg("reduce_op", "The reduction operation.")
- .add_arg("reduce_axes", "The reduction axes.")
+ .signature(sig::arg("binary_output", "The binary operation output."),
+ sig::arg("reduce_output", "The reduction output."),
+ sig::arg("binary_input1", "The first binary input."),
+ sig::arg("binary_input2", "The second binary input."),
+ sig::arg("binary_op", "The binary operation."),
+ sig::arg("reduce_op", "The reduction operation."),
+ sig::arg("reduce_axes", "The reduction axes."))
.set_attr<TScriptPrinterName>("TScriptPrinterName", ffi::String("binary_reduce"))
.set_attr<TIRxOpCategory>("TIRxOpCategory", ffi::String("tile_primitive"));
OpDef("tirx.tile.unary_reduce")
- .add_arg("unary_output", "The unary operation output.")
- .add_arg("reduce_output", "The reduction output.")
- .add_arg("unary_input", "The unary operation input.")
- .add_arg("unary_op", "The unary operation.")
- .add_arg("reduce_op", "The reduction operation.")
- .add_arg("bias", "The bias.")
- .add_arg("scale", "The scale factor.")
- .add_arg("reduce_axes", "The reduction axes.")
+ .signature(sig::arg("unary_output", "The unary operation output."),
+ sig::arg("reduce_output", "The reduction output."),
+ sig::arg("unary_input", "The unary operation input."),
+ sig::arg("unary_op", "The unary operation."),
+ sig::arg("reduce_op", "The reduction operation."), sig::arg("bias", "The bias."),
+ sig::arg("scale", "The scale factor."),
+ sig::arg("reduce_axes", "The reduction axes."))
.set_attr<TScriptPrinterName>("TScriptPrinterName", ffi::String("unary_reduce"))
.set_attr<TIRxOpCategory>("TIRxOpCategory", ffi::String("tile_primitive"));
OpDef("tirx.tile.binary_chain")
- .add_arg("output", "The output.")
- .add_arg("data", "The input data.")
- .add_arg("operand0", "The first operand.")
- .add_arg("operand1", "The second operand.")
- .add_arg("op0", "The first operation.")
- .add_arg("op1", "The second operation.")
- .add_arg("reverse1", "Whether to reverse the second operation.")
+ .signature(sig::arg("output", "The output."), sig::arg("data", "The input data."),
+ sig::arg("operand0", "The first operand."),
+ sig::arg("operand1", "The second operand."),
+ sig::arg("op0", "The first operation."), sig::arg("op1", "The second operation."),
+ sig::arg("reverse1", "Whether to reverse the second operation."))
.set_attr<TScriptPrinterName>("TScriptPrinterName", ffi::String("binary_chain"))
.set_attr<TIRxOpCategory>("TIRxOpCategory", ffi::String("tile_primitive"));
OpDef("tirx.tile.select")
- .add_arg("dst", "The destination.")
- .add_arg("true_value", "The value when the condition is true.")
- .add_arg("false_value", "The value when the condition is false.")
- .add_arg("pred", "The predicate.")
+ .signature(sig::arg("dst", "The destination."),
+ sig::arg("true_value", "The value when the condition is true."),
+ sig::arg("false_value", "The value when the condition is false."),
+ sig::arg("pred", "The predicate."))
.set_attr<TScriptPrinterName>("TScriptPrinterName", ffi::String("select"))
.set_attr<TIRxOpCategory>("TIRxOpCategory", ffi::String("tile_primitive"));
OpDef("tirx.tile.cast")
- .add_arg("dst", "The destination.")
- .add_arg("src", "The source.")
+ .signature(sig::arg("dst", "The destination."), sig::arg("src", "The source."))
.set_attr<TScriptPrinterName>("TScriptPrinterName", ffi::String("cast"))
.set_attr<TIRxOpCategory>("TIRxOpCategory", ffi::String("tile_primitive"));
OpDef("tirx.tile.fma")
- .add_arg("dst", "The destination.")
- .add_arg("src", "The source.")
- .add_arg("scale", "The scale factor.")
- .add_arg("bias", "The bias.")
+ .signature(sig::arg("dst", "The destination."), sig::arg("src", "The source."),
+ sig::arg("scale", "The scale factor."), sig::arg("bias", "The bias."))
.set_attr<TScriptPrinterName>("TScriptPrinterName", ffi::String("fma"))
.set_attr<TIRxOpCategory>("TIRxOpCategory", ffi::String("tile_primitive"));
OpDef("tirx.tile.silu")
- .add_arg("dst", "The destination.")
- .add_arg("src", "The source.")
+ .signature(sig::arg("dst", "The destination."), sig::arg("src", "The source."))
.set_attr<TScriptPrinterName>("TScriptPrinterName", ffi::String("silu"))
.set_attr<TIRxOpCategory>("TIRxOpCategory", ffi::String("tile_primitive"));
OpDef("tirx.tile.permute_layout")
- .add_arg("dst", "The destination.")
- .add_arg("src", "The source.")
+ .signature(sig::arg("dst", "The destination."), sig::arg("src", "The source."))
.set_attr<TScriptPrinterName>("TScriptPrinterName", ffi::String("permute_layout"))
.set_attr<TIRxOpCategory>("TIRxOpCategory", ffi::String("tile_primitive"));
OpDef("tirx.tile.copy_async")
- .add_arg("dst", "The destination.")
- .add_arg("src", "The source.")
+ .signature(sig::arg("dst", "The destination."), sig::arg("src", "The source."))
.set_attr<TScriptPrinterName>("TScriptPrinterName", ffi::String("copy_async"))
.set_attr<TIRxOpCategory>("TIRxOpCategory", ffi::String("tile_primitive"));
OpDef("tirx.tile.gemm_async")
- .add_arg("C", "The C tile.")
- .add_arg("A", "The A tile.")
- .add_arg("B", "The B tile.")
- .allow_extra_args()
+ .signature(sig::arg("C", "The C tile."), sig::arg("A", "The A tile."),
+ sig::arg("B", "The B tile."), sig::var_args("args"))
.set_attr<TScriptPrinterName>("TScriptPrinterName", ffi::String("gemm_async"))
.set_attr<TIRxOpCategory>("TIRxOpCategory", ffi::String("tile_primitive"));
}
diff --git a/tests/cpp/ir_op_signature_test.cc b/tests/cpp/ir_op_signature_test.cc
new file mode 100644
index 0000000000..ccdbed0e57
--- /dev/null
+++ b/tests/cpp/ir_op_signature_test.cc
@@ -0,0 +1,193 @@
+/*
+ * 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.
+ */
+
+#include <gtest/gtest.h>
+#include <tvm/ir/op.h>
+
+#include <string>
+
+namespace tvm {
+
+static std::string ValidationError(const Op& op, const Call& call) {
+ try {
+ op.Validate(call.get());
+ } catch (const ffi::Error& error) {
+ return error.what();
+ }
+ return "";
+}
+
+TEST(OpSignature, MetadataAndValidator) {
+ OpDef def("test.op_signature.typed", "A synthetic signature.");
+ def.signature(sig::arg("value", "The input."), sig::arg<PrimExpr>("index"),
+ sig::var_args<PrimExpr>("rest", "Remaining indices."), sig::ty_arg<PrimType>("T"),
+ sig::var_ty_args<PrimType>("Ts", "More types."), sig::call_attrs<DictAttrsNode>());
+ Op op = def.op();
+ EXPECT_EQ(op->args_info.size(), 2);
+ EXPECT_EQ(op->args_info[0]->name, "value");
+ EXPECT_EQ(op->args_info[0]->doc, "The input.");
+ EXPECT_EQ(op->args_info[1]->name, "index");
+ ASSERT_TRUE(op->var_args_info.has_value());
+ EXPECT_EQ(op->var_args_info.value()->name, "rest");
+ EXPECT_EQ(op->var_args_info.value()->doc, "Remaining indices.");
+ EXPECT_EQ(op->ty_args_info.size(), 1);
+ EXPECT_EQ(op->ty_args_info[0]->name, "T");
+ ASSERT_TRUE(op->var_ty_args_info.has_value());
+ EXPECT_EQ(op->var_ty_args_info.value()->name, "Ts");
+ EXPECT_EQ(op->var_ty_args_info.value()->doc, "More types.");
+ EXPECT_EQ(op->attrs_type_key, DictAttrsNode::_type_key);
+ EXPECT_EQ(op->attrs_type_index, DictAttrsNode::RuntimeTypeIndex());
+
+ Expr ordinary = Var("value", AnyType());
+ PrimExpr index = PrimVar("index");
+ Type scalar = PrimType::Int(32);
+ Call minimum(AnyType(), op, {ordinary, index}, DictAttrs(), {scalar});
+ Call extended(AnyType(), op, {ordinary, index, PrimVar("extra")}, DictAttrs(),
+ {scalar, PrimType::Float(32)});
+ EXPECT_TRUE(minimum.defined());
+ EXPECT_TRUE(extended.defined());
+
+ EXPECT_THROW(Call(AnyType(), op, {ordinary}, DictAttrs(), {scalar}), ffi::Error);
+ Call too_few = Call::Unchecked(AnyType(), op, {ordinary}, DictAttrs(), {scalar});
+ EXPECT_TRUE(too_few.defined());
+
+ std::string fixed_error = ValidationError(
+ op, Call::Unchecked(AnyType(), op, {ordinary, ordinary}, DictAttrs(), {scalar}));
+ EXPECT_EQ(fixed_error,
+ "Op `test.op_signature.typed`: `Call.args[1]` (`index`) expected `ir.PrimExpr`, got "
+ "`ir.Var[ty=ir.AnyType]`.");
+ std::string tail_error = ValidationError(
+ op, Call::Unchecked(AnyType(), op, {ordinary, index, ordinary}, DictAttrs(), {scalar}));
+ EXPECT_EQ(tail_error,
+ "Op `test.op_signature.typed`: `Call.args[2]` (`rest`) expected `ir.PrimExpr`, got "
+ "`ir.Var[ty=ir.AnyType]`.");
+ EXPECT_EQ(ValidationError(op, Call::Unchecked(AnyType(), op, {ordinary, index}, DictAttrs(),
+ {scalar, StringType()})),
+ "Op `test.op_signature.typed`: `Call.ty_args[1]` (`Ts`) expected `ir.PrimType`, got "
+ "`ir.StringType`.");
+ EXPECT_EQ(ValidationError(op, too_few),
+ "Op `test.op_signature.typed`: Call.args expected at least 2 arguments, got 1");
+ EXPECT_EQ(ValidationError(op, Call::Unchecked(AnyType(), op, {ordinary, index}, DictAttrs(), {})),
+ "Op `test.op_signature.typed`: Call.ty_args expected at least 1 type argument, got "
+ "0");
+ EXPECT_EQ(
+ ValidationError(op, Call::Unchecked(AnyType(), op, {ordinary, index}, Attrs(), {scalar})),
+ "Op `test.op_signature.typed`: Call.attrs expected `ir.DictAttrs`, got None");
+ EXPECT_EQ(ValidationError(op, Call::Unchecked(AnyType(), op, {ordinary, index},
+ Attrs(ffi::make_object<AttrsNode>()), {scalar})),
+ "Op `test.op_signature.typed`: Call.attrs expected `ir.DictAttrs`, got `ir.Attrs`");
+
+ ffi::Any non_expr(ffi::String("text"));
+ TVMFFIAny raw = ffi::AnyView(non_expr).CopyToTVMFFIAny();
+ EXPECT_EQ(ffi::TypeTraits<PrimExpr>::GetMismatchTypeInfo(&raw),
+ ffi::TypeTraitsBase::GetMismatchTypeInfo(&raw));
+
+ OpDef(op->name).set_attr<bool>("FSignatureTestMarker", true);
+ EXPECT_EQ(op->args_info.size(), 2);
+ EXPECT_TRUE(op->var_args_info.has_value());
+ EXPECT_NO_THROW(OpDef(op->name).signature(sig::arg("replacement")));
+ ASSERT_EQ(op->args_info.size(), 1);
+ EXPECT_EQ(op->args_info[0]->name, "replacement");
+ EXPECT_FALSE(op->var_args_info.has_value());
+ EXPECT_TRUE(op->ty_args_info.empty());
+ EXPECT_FALSE(op->var_ty_args_info.has_value());
+ EXPECT_TRUE(op->attrs_type_key.empty());
+ EXPECT_NO_THROW(op.Validate(minimum.get()));
+ EXPECT_NO_THROW(OpDef(op->name).add_arg("legacy", ""));
+ EXPECT_NO_THROW(OpDef(op->name).add_ty_arg("legacy_type", ""));
+ EXPECT_NO_THROW(OpDef(op->name).attrs_type<AttrsNode>());
+}
+
+TEST(OpSignature, DefaultsAndRepeatedRegistration) {
+ OpDef plain("test.op_signature.defaults");
+ plain.signature(sig::arg("value"), sig::ty_arg("T"));
+ Op op = plain.op();
+ EXPECT_FALSE(op->var_args_info.has_value());
+ EXPECT_FALSE(op->var_ty_args_info.has_value());
+ EXPECT_NO_THROW(Call(AnyType(), op, {Var("x", AnyType())}, Attrs(), {AnyType()}));
+ EXPECT_NO_THROW(Call(AnyType(), op, {Var("x", AnyType())}, DictAttrs(), {AnyType()}));
+ EXPECT_EQ(
+ ValidationError(op, Call::Unchecked(AnyType(), op, {Var("x", AnyType()), Var("y", AnyType())},
+ Attrs(), {AnyType()})),
+ "Op `test.op_signature.defaults`: Call.args expected 1 argument, got 2");
+
+ OpDef unnamed("test.op_signature.unnamed");
+ unnamed.signature(sig::arg<PrimExpr>(""));
+ EXPECT_EQ(ValidationError(unnamed.op(),
+ Call::Unchecked(AnyType(), unnamed.op(), {Var("x", AnyType())})),
+ "Op `test.op_signature.unnamed`: `Call.args[0]` expected `ir.PrimExpr`, got "
+ "`ir.Var[ty=ir.AnyType]`.");
+
+ OpDef tails("test.op_signature.default_tails");
+ tails.signature(sig::var_args("values"), sig::var_ty_args("types"));
+ EXPECT_TRUE(tails.op()->args_info.empty());
+ EXPECT_TRUE(tails.op()->ty_args_info.empty());
+ EXPECT_EQ(tails.op()->var_args_info.value()->name, "values");
+ EXPECT_EQ(tails.op()->var_ty_args_info.value()->name, "types");
+ EXPECT_NO_THROW(Call(AnyType(), tails.op(), {}, Attrs(), {}));
+ EXPECT_NO_THROW(Call(AnyType(), tails.op(), {Var("x", AnyType())}, Attrs(), {AnyType()}));
+
+ OpDef repeated("test.op_signature.repeated");
+ EXPECT_NO_THROW(repeated.signature(sig::arg("", "value doc"), sig::var_args("", "tail doc"),
+ sig::ty_arg("", "type doc"),
+ sig::var_ty_args("", "type tail doc")));
+ EXPECT_EQ(repeated.op()->args_info[0]->name, "");
+ EXPECT_EQ(repeated.op()->var_args_info.value()->name, "");
+ EXPECT_EQ(repeated.op()->ty_args_info[0]->name, "");
+ EXPECT_EQ(repeated.op()->var_ty_args_info.value()->name, "");
+ EXPECT_TRUE(repeated.op()->validator != nullptr);
+ EXPECT_NO_THROW(repeated.op().Validate(
+ Call(AnyType(), repeated.op(), {Var("x", AnyType())}, Attrs(), {AnyType()}).get()));
+ EXPECT_NO_THROW(repeated.signature(sig::arg("changed"), sig::arg("")));
+ EXPECT_EQ(repeated.op()->args_info.size(), 2);
+ EXPECT_EQ(repeated.op()->args_info[0]->name, "changed");
+ EXPECT_EQ(repeated.op()->args_info[1]->name, "");
+
+ OpDef legacy("test.op_signature.legacy");
+ legacy.add_arg("old", "Existing metadata.");
+ EXPECT_NO_THROW(legacy.signature(sig::arg("new")));
+ EXPECT_EQ(legacy.op()->args_info.size(), 1);
+ EXPECT_EQ(legacy.op()->args_info[0]->name, "new");
+ EXPECT_FALSE(legacy.op()->var_args_info.has_value());
+ EXPECT_TRUE(legacy.op()->validator != nullptr);
+
+ OpDef manual("test.op_signature.manual_hook");
+ int manual_calls = 0;
+ ffi::TypedFunction<ffi::Expected<void>(const CallNode*)> manual_hook(
+ [&](const CallNode*) -> ffi::Expected<void> {
+ ++manual_calls;
+ return {};
+ });
+ manual.set_validator(manual_hook);
+ EXPECT_THROW(manual.set_validator(manual_hook), ffi::Error);
+ EXPECT_NO_THROW(manual.set_validator(manual_hook, true));
+ EXPECT_NO_THROW(manual.signature(sig::arg("value")));
+ EXPECT_EQ(manual.op()->args_info[0]->name, "value");
+ Call manual_call(AnyType(), manual.op(), {Var("x", AnyType())});
+ manual.op().Validate(manual_call.get());
+ EXPECT_EQ(manual_calls, 2);
+ ffi::TypedFunction<ffi::Expected<void>(const CallNode*)> failing_hook(
+ [](const CallNode*) -> ffi::Expected<void> {
+ return TVM_FFI_UNEXPECTED(TypeError) << "packed validator rejected Call";
+ });
+ manual.set_validator(failing_hook, true);
+ EXPECT_THROW(manual.op().Validate(manual_call.get()), ffi::Error);
+}
+
+} // namespace tvm
diff --git a/tests/python/ir/test_ir_attrs.py b/tests/python/ir/test_ir_attrs.py
index 25480f7265..27448ccc47 100644
--- a/tests/python/ir/test_ir_attrs.py
+++ b/tests/python/ir/test_ir_attrs.py
@@ -57,6 +57,93 @@ def test_assert_structural_equal_reports_mismatch():
assert "and rhs at" in message
+def test_op_set_signature_metadata_and_counts():
+ name = "test.python_set_signature"
+ tvm.ir._ffi_api.RegisterOp(name, "Python signature test")
+ op = tvm.ir.Op.get(name)
+ x = tvm.ir.Var("x", tvm.ir.AnyType())
+ ty = tvm.ir.AnyType()
+
+ op.set_signature(
+ ["lhs", ("rhs", "Second input")],
+ ty_args=[("T", "Result type")],
+ var_args=("rest", "Other inputs"),
+ var_ty_args="more_types",
+ )
+ assert [(info.name, info.doc) for info in op.args_info] == [
+ ("lhs", ""),
+ ("rhs", "Second input"),
+ ]
+ assert [(info.name, info.doc) for info in op.ty_args_info] == [("T", "Result type")]
+ assert (op.var_args_info.name, op.var_args_info.doc) == ("rest", "Other inputs")
+ assert op.var_ty_args_info.name == "more_types"
+ tvm.ir.Call(op, [x, x], ty_args=[ty])
+ tvm.ir.Call(op, [x, x, x], ty_args=[ty, ty])
+ with pytest.raises(TypeError, match=r"Call.args expected at least 2 arguments, got 1"):
+ tvm.ir.Call(op, [x], ty_args=[ty])
+ with pytest.raises(TypeError, match=r"Call.ty_args expected at least 1 type argument, got 0"):
+ tvm.ir.Call(op, [x, x])
+
+ op.set_signature(["only"])
+ assert [info.name for info in op.args_info] == ["only"]
+ assert op.var_args_info is None and op.var_ty_args_info is None
+ tvm.ir.Call(op, [x])
+ with pytest.raises(TypeError, match=r"Call.args expected 1 argument, got 2"):
+ tvm.ir.Call(op, [x, x])
+ with pytest.raises(TypeError, match=r"Call.ty_args expected 0 type arguments, got 1"):
+ tvm.ir.Call(op, [x], ty_args=[ty])
+ with pytest.raises(TypeError, match=r"args must be a name string or a \(name, doc\) tuple"):
+ op.set_signature(["new", ("bad", 7)])
+ with pytest.raises(TypeError, match=r"var_ty_args must be a name string"):
+ op.set_signature(["new"], var_ty_args=("bad", 7))
+ assert [info.name for info in op.args_info] == ["only"]
+ tvm.ir.Call(op, [x])
+
+
+def test_op_set_signature_preserves_custom_validator():
+ op = tvm.ir.Op.get("relax.call_pure_packed")
+ original_args = [(info.name, info.doc) for info in op.args_info]
+ original_ty_args = [(info.name, info.doc) for info in op.ty_args_info]
+ original_var_args = (op.var_args_info.name, op.var_args_info.doc)
+ original_var_ty_args = (op.var_ty_args_info.name, op.var_ty_args_info.doc)
+ try:
+ op.set_signature()
+ with pytest.raises(TypeError, match="call_pure_packed expects a function argument"):
+ tvm.ir.Call(op, [])
+ finally:
+ op.set_signature(
+ original_args,
+ ty_args=original_ty_args,
+ var_args=original_var_args,
+ var_ty_args=original_var_ty_args,
+ )
+
+
+def test_op_set_signature_checks_counts_without_types_or_attrs():
+ name = "test.python_set_signature.count_only"
+ tvm.ir._ffi_api.RegisterOp(name, "Count-only Python signature")
+ op = tvm.ir.Op.get(name)
+ op.set_signature(["value"], ty_args=["T"])
+
+ call = tvm.ir.Call(
+ op,
+ [tvm.ir.Var("x", tvm.ir.AnyType())],
+ attrs={"note": "permitted"},
+ ty_args=[tvm.ir.StringType()],
+ )
+ assert call.attrs is not None
+
+
+def test_ptx_ops_have_variadic_call_signature():
+ from tvm.backend.cuda.ptx import TABLE
+
+ op = tvm.ir.Op.get(next(iter(TABLE.values())).op_name)
+ assert not op.args_info
+ assert op.var_args_info.name == "operands"
+ tvm.ir.Call(op, [])
+ tvm.ir.Call(op, [tvm.ir.Var("x", tvm.ir.AnyType())])
+
+
if __name__ == "__main__":
test_dict_attrs()
test_attrs_equal()
diff --git a/tests/python/relax/script/test_relax_script_dynamic_shape.py b/tests/python/relax/script/test_relax_script_dynamic_shape.py
index 41efcd191b..2ee370a9aa 100644
--- a/tests/python/relax/script/test_relax_script_dynamic_shape.py
+++ b/tests/python/relax/script/test_relax_script_dynamic_shape.py
@@ -497,7 +497,7 @@ def test_symbolic_vars_in_shape():
def test_string_shape_expression_is_not_resolved():
"""A quoted expression is ordinary data, not a symbol declaration."""
- with pytest.raises(TypeError, match="Array<ir.Expr>"):
+ with pytest.raises(TypeError, match=r"Array<ir\.PrimExpr>"):
@R.function
def foo(x: R.Tensor(("m + 1", "m * 2"), "float32")):
diff --git a/tests/python/relax/test_analysis_well_formed.py b/tests/python/relax/test_analysis_well_formed.py
index 2e51ce5568..daf834999a 100644
--- a/tests/python/relax/test_analysis_well_formed.py
+++ b/tests/python/relax/test_analysis_well_formed.py
@@ -152,7 +152,7 @@ def test_global_var():
# Error: GlobalVar GlobalVar0 is not defined
gv0 = rx.Var("gv0", R.Tensor([m, n], "float32"))
globalvar = rx.GlobalVar("GlobalVar0")
- call_node = rx.Call(
+ call_node = rx.Call.unchecked(
op=tvm.ir.Op.get("relax.call_tir"),
args=[globalvar, rx.Tuple([x]), rx.ShapeExpr([m, n])],
)
@@ -163,6 +163,24 @@ def test_global_var():
assert not rx.analysis.check_well_formed(mod, check_ty=False)
+def test_unchecked_call_constructor():
+ op = tvm.ir.Op.get("relax.add")
+ with pytest.raises(Exception, match="Call.args expected 2 arguments, got 1"):
+ tvm.ir.Call(op, [x])
+
+ span = tvm.ir.Span(tvm.ir.SourceName("unchecked.py"), 1, 1, 0, 1)
+ call = tvm.ir.Call.unchecked("relax.add", [x], attrs={"key": 1}, span=span)
+ assert isinstance(call, tvm.ir.Call)
+ assert call.op.same_as(op)
+ assert call.ty.is_missing()
+ assert call.span.same_as(span)
+ assert isinstance(call.attrs, tvm.ir.DictAttrs)
+ assert len(call.ty_args) == 0
+ assert isinstance(tvm.ir.Call.unchecked(op, [x], ret_ty="handle").ty, tvm.ir.PointerType)
+ with pytest.raises(TypeError, match="skip_validate"):
+ tvm.ir.Call(op, [x], skip_validate=True)
+
+
def test_symbolic_var():
# Error: Symbolic Var new_s is not defined
new_s = tirx.Var("new_s", "int64")
@@ -444,7 +462,7 @@ def test_inline_prim_func():
),
rx.VarBinding(
var=y,
- value=rx.Call(
+ value=rx.Call.unchecked(
op=tvm.ir.Op.get("relax.call_tir"),
args=[
rx.GlobalVar("GlobalVar0"),
@@ -799,6 +817,8 @@ def test_call_tir_with_incorrect_primitive_argument_dtype():
B[i] = A[i] * T.Cast("float16", scale)
assert not rx.analysis.check_well_formed(Module)
+ with pytest.raises(ValueError, match="Argument 1 type mismatch"):
+ rx.transform.Normalize()(Module)
def test_call_tir_shape_expr_is_not_a_primitive_argument():
@@ -1470,5 +1490,15 @@ def test_incomplete_ty_must_be_consistent():
assert not rx.analysis.check_well_formed(Module)
+def test_stop_lift_params_optional_type_arg():
+ x = rx.Var("x", rx.TensorType((4,), "float32"))
+ out_ty = rx.TensorType((4,), "float32", rx.VDevice("llvm"))
+ for ty_args in ([], [out_ty]):
+ call = tvm.ir.Call("relax.builtin.stop_lift_params", [x], ty_args=ty_args)
+ func = rx.Function([x], call, ty_args[0] if ty_args else x.ty)
+ normalized = rx.transform.Normalize()(tvm.IRModule.from_expr(func))
+ rx.analysis.well_formed(normalized)
+
+
if __name__ == "__main__":
tvm.testing.main()