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/10/08 11:14:38 UTC
(tvm) branch main updated: [REFACTOR][IR] Unify Call construction and TVMScript printing (#20578)
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 3c2c6eecee [REFACTOR][IR] Unify Call construction and TVMScript printing (#20578)
3c2c6eecee is described below
commit 3c2c6eecee6a1b16343e1e049b2303a0789df1d7
Author: Tianqi Chen <tq...@users.noreply.github.com>
AuthorDate: Thu Oct 8 07:14:30 2026 -0400
[REFACTOR][IR] Unify Call construction and TVMScript printing (#20578)
Move best-effort result inference into the core Call constructor.
Omitted or None `ty` uses available inference, explicit types are
preserved, and `Call.validate()` checks the operator contract
separately. `I.Call` directly exposes this constructor and preserves
provisional IR without a separate unchecked path.
Use one shared printer for canonical named operations with operands,
semantic attributes, type arguments and result types. Remove
standard-call markers and dialect-wide registration helpers. Retain
concrete hooks for representation-changing syntax, including LLVM names,
PTX operands and Relax call_tir sugar; use exact raw `I.Call`
reconstruction when named syntax cannot preserve the stored fields.
Preserve existing Python operator definitions, documentation and operand
normalization, keeping helper parameters specific to their semantics and
generating only missing exposures. Give assertions the stored
condition/format/variadic-values contract and take-backward dedicated
axis attributes, removing their special printer paths. Preserve
contextual, native-call and pointer inference, reflect literal attribute
defaults, and give the five legacy CUDA helpers explicit result-type
parameters.
---
CMakeLists.txt | 1 +
docs/arch/tvmscript.rst | 12 +-
docs/reference/api/python/script/ir_builder.rst | 2 +-
include/tvm/ir/expr.h | 16 +-
include/tvm/ir/op.h | 19 +-
include/tvm/relax/attrs/ccl.h | 12 +-
include/tvm/relax/attrs/create.h | 3 +-
include/tvm/relax/attrs/datatype.h | 3 +-
include/tvm/relax/attrs/image.h | 59 +-
include/tvm/relax/attrs/index.h | 16 +-
include/tvm/relax/attrs/linear_algebra.h | 3 +-
include/tvm/relax/attrs/manipulate.h | 38 +-
include/tvm/relax/attrs/nn.h | 279 ++++---
include/tvm/relax/attrs/search.h | 13 +-
include/tvm/relax/attrs/sorting.h | 8 +-
include/tvm/relax/attrs/statistical.h | 12 +-
include/tvm/relax/attrs/vision.h | 62 +-
include/tvm/tirx/op_attr_types.h | 35 -
python/tvm/backend/cuda/__init__.py | 2 +-
python/tvm/backend/cuda/lang/pipeline.py | 2 +-
python/tvm/backend/cuda/op.py | 478 ++++++-----
python/tvm/backend/cuda/ptx/engine.py | 33 +-
python/tvm/backend/cuda/script.py | 165 ++--
python/tvm/backend/cuda/tile_primitive/common.py | 6 +-
.../cuda/tile_primitive/copy_async/tcgen05_cp.py | 2 +-
python/tvm/backend/metal/op.py | 81 +-
python/tvm/backend/metal/script.py | 28 +-
python/tvm/backend/trn/__init__.py | 8 +-
python/tvm/backend/trn/op.py | 196 +++--
python/tvm/backend/trn/script.py | 50 +-
python/tvm/ir/expr.py | 56 +-
python/tvm/ir/op.py | 50 +-
.../tvm/relax/frontend/tflite/tflite_frontend.py | 4 +-
.../frontend/torch/exported_program_translator.py | 2 +-
python/tvm/relax/op/base.py | 166 ++--
python/tvm/relax/op/binary.py | 107 +--
python/tvm/relax/op/builtin/builtin.py | 24 +-
python/tvm/relax/op/ccl/ccl.py | 43 +-
python/tvm/relax/op/create.py | 134 ++-
python/tvm/relax/op/datatype.py | 29 +-
python/tvm/relax/op/distributed/distributed.py | 88 +-
python/tvm/relax/op/grad/grad.py | 133 ++-
python/tvm/relax/op/image/image.py | 104 ++-
python/tvm/relax/op/index.py | 52 +-
python/tvm/relax/op/linear_algebra.py | 38 +-
python/tvm/relax/op/manipulate.py | 326 ++++++--
python/tvm/relax/op/memory/memory.py | 40 +-
python/tvm/relax/op/memory/view.py | 28 +-
python/tvm/relax/op/nn/nn.py | 905 ++++++++++++++++-----
python/tvm/relax/op/qdq.py | 49 +-
python/tvm/relax/op/sampling.py | 18 +-
python/tvm/relax/op/search.py | 68 +-
python/tvm/relax/op/set.py | 16 +-
python/tvm/relax/op/sorting.py | 53 +-
python/tvm/relax/op/statistical.py | 193 ++++-
python/tvm/relax/op/ternary.py | 7 +-
python/tvm/relax/op/unary.py | 131 +--
.../tvm/relax/op/vision/multibox_transform_loc.py | 28 +-
python/tvm/relax/op/vision/nms.py | 76 +-
python/tvm/relax/op/vision/roi_align.py | 30 +-
python/tvm/relax/op/vision/roi_pool.py | 20 +-
python/tvm/relax/op/vm/vm.py | 35 +-
python/tvm/relax/script/ir_builder/__init__.py | 46 --
.../tvm/relax/script/ir_builder/distributed/op.py | 18 +-
python/tvm/relax/script/ir_builder/op.py | 73 +-
.../tvm/relax/script/ir_builder/parser_protocol.py | 2 +-
.../tvm/relax/transform/legalize_ops/inspect_op.py | 14 +-
python/tvm/s_tir/schedule/schedule.py | 4 +-
python/tvm/s_tir/tensor_intrin/arm_cpu.py | 46 +-
python/tvm/s_tir/tensor_intrin/cuda.py | 14 +-
.../tvm/s_tir/tensor_intrin/dot_product_common.py | 2 +-
python/tvm/s_tir/tensor_intrin/hexagon.py | 18 +-
python/tvm/s_tir/tensor_intrin/riscv_cpu.py | 24 +-
python/tvm/s_tir/tensor_intrin/rocm.py | 14 +-
python/tvm/s_tir/tensor_intrin/x86.py | 14 +-
python/tvm/script/ir_builder/__init__.py | 4 +
python/tvm/tirx/op.py | 582 ++++++++-----
python/tvm/tirx/script/ir_builder/op.py | 190 ++---
src/backend/cuda/codegen/intrin_rule_cuda.cc | 10 +-
src/backend/cuda/op/target_builtin.cc | 46 +-
src/backend/metal/codegen/intrin_rule_metal.cc | 6 +-
src/backend/metal/op/target_builtin.cc | 4 +
src/backend/webgpu/codegen/intrin_rule_webgpu.cc | 6 +-
src/ir/expr.cc | 67 +-
src/ir/prim/builtin.cc | 11 +
.../backend/adreno/annotate_custom_storage.cc | 10 +-
.../backend/adreno/fold_vdevice_scope_change.cc | 3 +-
src/relax/backend/vm/codegen_vm.cc | 7 +-
src/relax/backend/vm/lower_runtime_builtin.cc | 34 +-
src/relax/backend/vm/vm_shape_lower.cc | 18 +-
.../transform/lower_global_view_to_local_view.cc | 4 +-
src/relax/ir/block_builder.cc | 2 +-
src/relax/op/ccl/ccl.cc | 9 +-
src/relax/op/distributed/distributed.cc | 8 +-
src/relax/op/image/resize.cc | 8 +-
src/relax/op/memory/view.cc | 22 +-
src/relax/op/nn/attention.cc | 16 +-
src/relax/op/nn/convolution.cc | 6 +-
src/relax/op/nn/convolution.h | 2 +-
src/relax/op/nn/nn.cc | 64 +-
src/relax/op/nn/pooling.cc | 12 +-
src/relax/op/op.cc | 116 ++-
src/relax/op/op_common.h | 14 +-
src/relax/op/tensor/binary.h | 2 +-
src/relax/op/tensor/create.cc | 35 +-
src/relax/op/tensor/datatype.cc | 5 +-
src/relax/op/tensor/grad.cc | 38 +-
src/relax/op/tensor/index.cc | 10 +-
src/relax/op/tensor/inspect.cc | 31 +-
src/relax/op/tensor/linear_algebra.cc | 6 +-
src/relax/op/tensor/manipulate.cc | 62 +-
src/relax/op/tensor/qdq.cc | 8 +-
src/relax/op/tensor/sampling.cc | 7 +-
src/relax/op/tensor/search.cc | 12 +-
src/relax/op/tensor/set.cc | 11 +-
src/relax/op/tensor/sorting.cc | 7 +-
src/relax/op/tensor/statistical.cc | 21 +-
src/relax/op/tensor/ternary.cc | 2 +-
src/relax/op/tensor/unary.cc | 64 +-
src/relax/op/vision/multibox_transform_loc.cc | 6 +-
src/relax/op/vision/nms.cc | 17 +-
src/relax/op/vision/roi_align.cc | 3 +-
src/relax/op/vision/roi_pool.cc | 2 +-
src/relax/script/ir_builder/distributed.cc | 2 +-
src/relax/script/printer/call.cc | 236 +-----
src/relax/transform/allocate_workspace.cc | 8 +-
src/relax/transform/alter_op_impl.cc | 13 +-
src/relax/transform/call_tir_rewrite.cc | 28 +-
src/relax/transform/compute_prim_value.cc | 5 +-
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 | 12 +-
src/relax/transform/run_codegen.cc | 2 +-
.../transform/split_layout_rewrite_preproc.cc | 8 +-
src/relax/transform/static_plan_block_memory.cc | 18 +-
src/s_tir/script/printer/expr.cc | 19 -
src/script/printer/ir/call.cc | 141 ++++
src/script/printer/ir/utils.cc | 32 +-
src/tirx/ir/stmt.cc | 6 +-
src/tirx/op/builtin.cc | 133 ++-
src/tirx/op/op.cc | 36 +-
src/tirx/script/printer/buffer.cc | 9 +
src/tirx/script/printer/expr.cc | 346 +++-----
src/tirx/script/printer/stmt.cc | 4 +-
src/tirx/script/printer/utils.h | 7 -
tests/cpp/ir_op_signature_test.cc | 39 +-
tests/cpp/tir_scalable_datatype.cc | 4 +-
tests/python/codegen/test_target_codegen.py | 8 +-
tests/python/codegen/test_target_codegen_c_host.py | 2 +-
tests/python/codegen/test_target_codegen_cuda.py | 4 +-
.../python/codegen/test_target_codegen_cuda_fp4.py | 4 +-
.../python/codegen/test_target_codegen_cuda_fp8.py | 8 +-
tests/python/codegen/test_target_codegen_llvm.py | 22 +-
.../python/codegen/test_target_codegen_llvm_vla.py | 2 +-
tests/python/codegen/test_target_codegen_riscv.py | 10 +-
.../codegen/test_target_codegen_static_init.py | 2 +-
tests/python/codegen/test_target_codegen_vulkan.py | 6 +-
tests/python/ir/test_ir_attrs.py | 22 +-
.../relax/script/test_relax_script_basic_usage.py | 14 +-
.../relax/script/test_relax_script_printer.py | 25 +-
tests/python/relax/test_analysis_well_formed.py | 24 +-
tests/python/relax/test_expr.py | 21 +-
tests/python/relax/test_op_distributed.py | 5 +-
tests/python/relax/test_op_index.py | 6 +-
tests/python/relax/test_op_manipulate.py | 4 +-
tests/python/relax/test_op_nn_pooling.py | 8 +-
tests/python/relax/test_relax_operators.py | 8 +-
tests/python/relax/test_transform_cse.py | 8 +-
.../relax/test_transform_error_enrichment.py | 7 +-
.../test_s_tir_analysis_is_pure_function.py | 4 +-
.../s_tir/analysis/test_sblock_access_region.py | 4 +-
tests/python/s_tir/dlight/test_gpu_matmul.py | 6 +-
.../s_tir/dlight/test_gpu_matmul_tensorize.py | 36 +-
tests/python/s_tir/dlight/test_gpu_transpose.py | 4 +-
...est_meta_schedule_postproc_rewrite_tensorize.py | 8 +-
.../test_meta_schedule_postproc_verify_gpu_code.py | 6 +-
.../test_meta_schedule_trace_apply.py | 12 +-
.../schedule/test_tir_schedule_cache_read_write.py | 20 +-
.../s_tir/schedule/test_tir_schedule_tensorize.py | 16 +-
.../s_tir/script/test_s_tir_script_basic_usage.py | 14 +-
.../script/test_s_tir_script_error_handling.py | 4 +-
.../s_tir/script/test_s_tir_script_printer.py | 99 +--
.../test_s_tir_transform_compact_buffer_region.py | 34 +-
.../transform/test_s_tir_transform_hoist_if.py | 14 +-
.../test_s_tir_transform_inject_permuted_layout.py | 24 +-
.../test_s_tir_transform_inject_ptx_async_copy.py | 10 +-
.../test_s_tir_transform_inject_virtual_thread.py | 17 +-
.../test_s_tir_transform_loop_partition.py | 6 +-
...s_tir_transform_lower_cross_thread_reduction.py | 10 +-
.../test_s_tir_transform_lower_match_buffer.py | 6 +-
...st_s_tir_transform_memhammer_lower_auto_copy.py | 8 +-
...form_merge_dynamic_shared_memory_allocations.py | 12 +-
...sform_plan_update_buffer_allocation_location.py | 10 +-
.../transform/test_s_tir_transform_thread_sync.py | 4 +-
tests/python/script/test_constructor_contracts.py | 14 +-
tests/python/te/test_te_create_function.py | 2 +-
tests/python/tirx-base/test_tir_op_types.py | 4 +-
tests/python/tirx-base/test_tir_ptx_cp_async.py | 8 +-
tests/python/tirx-base/test_tir_ptx_ldmatrix.py | 2 +-
tests/python/tirx-base/test_tir_ptx_mma.py | 34 +-
.../python/tirx-base/test_tir_scalable_datatype.py | 6 +-
tests/python/tirx-base/test_tir_texture_scope.py | 6 +-
.../test_tir_transform_bf16_legalize.py | 28 +-
.../test_tir_transform_common_subexpr_elim.py | 8 +-
...test_tir_transform_force_narrow_index_to_i32.py | 4 +-
.../test_tir_transform_fp8_legalize.py | 22 +-
.../test_tir_transform_lower_tvm_builtin.py | 14 +-
.../test_tir_transform_make_packed_api.py | 14 +-
.../test_tir_transform_remove_no_op.py | 4 +-
.../test_tir_transform_split_host_device.py | 6 +-
.../test_tir_transform_storage_rewrite.py | 26 +-
.../tirx-transform/test_tir_transform_vectorize.py | 40 +-
tests/python/tirx/codegen/test_codegen_cuda.py | 6 +-
tests/python/tirx/codegen/test_codegen_dsmem.py | 2 +-
.../cuda/gemm_async/test_gemm_async.py | 12 +-
.../tirx/script/test_tirx_script_basic_usage.py | 8 +-
tests/python/tirx/script/test_tirx_script_gpu.py | 2 +-
.../tirx/script/test_tirx_script_ir_builder.py | 6 +-
.../python/tirx/script/test_tirx_script_printer.py | 8 +-
.../tirx/transform/test_transform_lower_tirx.py | 4 +-
232 files changed, 5207 insertions(+), 3531 deletions(-)
diff --git a/CMakeLists.txt b/CMakeLists.txt
index 58a296f469..7f8279684c 100644
--- a/CMakeLists.txt
+++ b/CMakeLists.txt
@@ -312,6 +312,7 @@ tvm_file_glob(GLOB_RECURSE COMPILER_SRCS
src/script/printer/config.cc
src/script/printer/doc.cc
src/script/printer/doc_translator.cc
+ src/script/printer/ir/call.cc
src/script/printer/ir/utils.cc
src/script/printer/ir/ir.cc
src/script/printer/ir/misc.cc
diff --git a/docs/arch/tvmscript.rst b/docs/arch/tvmscript.rst
index 24ea1900f2..c7cce1ecf4 100644
--- a/docs/arch/tvmscript.rst
+++ b/docs/arch/tvmscript.rst
@@ -139,10 +139,14 @@ protocol.
Ordinary IR constructors can be used directly in parsed source. Shared exports
such as ``I.Call`` and ``T.Range`` use the same constructor contracts as
``tvm.ir.Call`` and ``tvm.ir.Range``, including keyword arguments, source spans
-and validation. ``Call(..., ty=...)`` supplies an explicit result type; omission
-leaves ``Type.missing()`` for subsequent normalization. Use ``Call.unchecked``
-explicitly for provisional calls that require later validation. Raw printed calls
-use this form with their stored result type to preserve all fields.
+and validation. ``Call(..., ty=...)`` preserves an explicit result type, including
+``Type.missing()``. Omitted or ``None`` results use available inference and retain
+``Type.missing()`` when no deduction is available; inference errors propagate.
+Construction permits provisional IR, and ``call.validate()`` checks the registered
+operator contract separately. Raw printed calls use ``I.Call`` with an explicit
+stored result type to preserve every field. Canonical named operations use the same
+``args``, ``attrs``, ``ty_args`` and ``ty`` construction contract. Concrete printer
+hooks handle representation-changing syntax; ordinary calls share one formatter.
Operations likewise retain their normal argument contracts. A dtype inferred from
operands is not an extra ``dtype`` keyword; operations with an explicit dtype
diff --git a/docs/reference/api/python/script/ir_builder.rst b/docs/reference/api/python/script/ir_builder.rst
index 414ac31cd1..805f69c695 100644
--- a/docs/reference/api/python/script/ir_builder.rst
+++ b/docs/reference/api/python/script/ir_builder.rst
@@ -23,7 +23,7 @@ tvm.script.ir_builder
.. automodule:: tvm.script.ir_builder
:members:
:imported-members:
- :exclude-members: Call, DataTypeImm, FuncType, GenericConst, MissingType, Op, PrimType, Range, StringImm, StringType, Type
+ :exclude-members: Call, DataTypeImm, FuncType, GenericConst, MissingType, Op, PrimType, Range, StringImm, StringType, TensorRegion, Tuple, Type
tvm.relax.script.ir_builder.distributed
***************************************
diff --git a/include/tvm/ir/expr.h b/include/tvm/ir/expr.h
index 1119a7ddc6..fd5a09d8f1 100644
--- a/include/tvm/ir/expr.h
+++ b/include/tvm/ir/expr.h
@@ -969,16 +969,18 @@ class CallNode : public ExprNode {
*/
class Call : public Expr {
public:
- TVM_DLL Call(Type ret_ty, Expr op, ffi::Array<Expr> args, Attrs attrs = Attrs(),
+ /*! \brief Construct a Call, inferring only when ret_ty is absent.
+ * Construction preserves provisional IR; invoke Validate separately when needed.
+ */
+ TVM_DLL Call(ffi::Optional<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());
+ /*! \brief Check this Call against its registered operator contract. */
+ TVM_DLL void Validate() const;
/*! \brief Recompute a result type from the Call's current explicit inputs.
*
* This ignores the Call's stored result type and does not mutate the Call.
- * The operator must register TFixedReturnType or a context-free FInferType hook.
+ * Returns Missing when no deduction is available; inference errors propagate.
* This does not invoke the operator's validator.
*/
TVM_DLL static Type ReinferType(const CallNode* call);
@@ -987,10 +989,6 @@ class Call : public Expr {
TVM_FFI_DEFINE_OBJECT_REF_METHODS_NOTNULLABLE(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 range over one dimension */
diff --git a/include/tvm/ir/op.h b/include/tvm/ir/op.h
index dcb985b882..4b87707175 100644
--- a/include/tvm/ir/op.h
+++ b/include/tvm/ir/op.h
@@ -52,6 +52,11 @@ using TFixedReturnType = Type;
/*! \brief Infer a Call's result type from its explicit inputs without builder state. */
using FInferType = ffi::reflection::NativeFunctionView<Type(const CallNode* call)>;
+/*!
+ * \brief The fully qualified TVMScript name, including its dialect namespace.
+ */
+using TScriptPrinterName = ffi::String;
+
/*! \brief An operator argument's name and documentation. */
class ArgumentInfoNode : public ffi::Object {
public:
@@ -170,6 +175,7 @@ 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;
@@ -488,11 +494,10 @@ class OpDef {
* 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.
+ * Call. Call::Validate invokes the validator separately from construction. 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")
@@ -544,8 +549,8 @@ class OpDef {
* 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.
+ * Call::Validate invokes it explicitly; Relax normalization and
+ * well-formedness also validate when inputs are ready.
* 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
diff --git a/include/tvm/relax/attrs/ccl.h b/include/tvm/relax/attrs/ccl.h
index 031a1de493..2d6a1d183d 100644
--- a/include/tvm/relax/attrs/ccl.h
+++ b/include/tvm/relax/attrs/ccl.h
@@ -40,10 +40,12 @@ struct AllReduceAttrs : public tvm::AttrsNode {
refl::ObjectDef<AllReduceAttrs>()
.def_ro("op_type", &AllReduceAttrs::op_type,
"The type of reduction operation to be applied to the input data. Now only sum is "
- "supported.")
+ "supported.",
+ refl::DefaultValue(ffi::String("sum")))
.def_ro("in_group", &AllReduceAttrs::in_group,
"Whether the reduction operation performs in group or globally or in group as "
- "default.");
+ "default.",
+ refl::DefaultValue(true));
}
TVM_FFI_DECLARE_OBJECT_INFO_FINAL("relax.attrs.AllReduceAttrs", AllReduceAttrs, AttrsNode);
}; // struct AllReduceAttrs
@@ -61,7 +63,8 @@ struct AllGatherAttrs : public tvm::AttrsNode {
"chunked into.")
.def_ro("in_group", &AllGatherAttrs::in_group,
"Whether the allgather operation performs in group or globally or in group as "
- "default.");
+ "default.",
+ refl::DefaultValue(true));
}
TVM_FFI_DECLARE_OBJECT_INFO_FINAL("relax.attrs.AllGatherAttrs", AllGatherAttrs, AttrsNode);
}; // struct AllGatherAttrs
@@ -79,7 +82,8 @@ struct ScatterCollectiveAttrs : public tvm::AttrsNode {
"chunked into.")
.def_ro("axis", &ScatterCollectiveAttrs::axis,
"The axis of the tensor to be scattered. The tensor will be chunked along "
- "this axis.");
+ "this axis.",
+ refl::DefaultValue(0));
}
TVM_FFI_DECLARE_OBJECT_INFO_FINAL("relax.attrs.ScatterCollectiveAttrs", ScatterCollectiveAttrs,
AttrsNode);
diff --git a/include/tvm/relax/attrs/create.h b/include/tvm/relax/attrs/create.h
index 614485eb9e..f2ab653443 100644
--- a/include/tvm/relax/attrs/create.h
+++ b/include/tvm/relax/attrs/create.h
@@ -36,7 +36,8 @@ struct InitAttrs : public AttrsNode {
static void RegisterReflection() {
namespace refl = tvm::ffi::reflection;
refl::ObjectDef<InitAttrs>().def_ro("dtype", &InitAttrs::dtype,
- "The data type of the created tensor.");
+ "The data type of the created tensor.",
+ refl::DefaultValue(ffi::Optional<DLDataType>{}));
}
TVM_FFI_DECLARE_OBJECT_INFO_FINAL("relax.attrs.InitAttrs", InitAttrs, AttrsNode);
}; // struct InitAttrs
diff --git a/include/tvm/relax/attrs/datatype.h b/include/tvm/relax/attrs/datatype.h
index aeac65e644..5ee6e8af33 100644
--- a/include/tvm/relax/attrs/datatype.h
+++ b/include/tvm/relax/attrs/datatype.h
@@ -46,7 +46,8 @@ struct WrapParamAttrs : public AttrsNode {
static void RegisterReflection() {
namespace refl = tvm::ffi::reflection;
- refl::ObjectDef<WrapParamAttrs>().def_ro("dtype", &WrapParamAttrs::dtype, "Target data type");
+ refl::ObjectDef<WrapParamAttrs>().def_ro("dtype", &WrapParamAttrs::dtype, "Target data type",
+ refl::DefaultValue((DLDataType{kDLFloat, 32, 1})));
}
TVM_FFI_DECLARE_OBJECT_INFO_FINAL("relax.attrs.WrapParamAttrs", WrapParamAttrs, AttrsNode);
}; // struct WrapParamAttrs.
diff --git a/include/tvm/relax/attrs/image.h b/include/tvm/relax/attrs/image.h
index 721acf87a9..7e3f8c27d1 100644
--- a/include/tvm/relax/attrs/image.h
+++ b/include/tvm/relax/attrs/image.h
@@ -50,30 +50,36 @@ struct Resize2DAttrs : public AttrsNode {
"Dimension ordering of input data. Can be 'NCHW', 'NHWC', etc."
"'N', 'C', 'H', 'W' stands for batch, channel, height, and width"
"dimensions respectively. Resize is applied on the 'H' and"
- "'W' dimensions.")
+ "'W' dimensions.",
+ refl::DefaultValue(ffi::String("NCHW")))
.def_ro("method", &Resize2DAttrs::method,
"Specify the mode to use for scaling."
"nearest_neighbor - Nearest Neighbor"
"linear - Bilinear Interpolation"
- "cubic - Bicubic Interpolation")
+ "cubic - Bicubic Interpolation",
+ refl::DefaultValue(ffi::String("linear")))
.def_ro("coordinate_transformation_mode", &Resize2DAttrs::coordinate_transformation_mode,
"Describes how to transform the coordinate in the resized tensor"
"to the coordinate in the original tensor."
"Refer to the ONNX Resize operator specification for details"
- "Available options are half_pixel, align_corners and asymmetric")
+ "Available options are half_pixel, align_corners and asymmetric",
+ refl::DefaultValue(ffi::String("half_pixel")))
.def_ro("rounding_method", &Resize2DAttrs::rounding_method,
"indicates how to find the \"nearest\" pixel in nearest_neighbor method"
- "Available options are round, floor, and ceil.")
+ "Available options are round, floor, and ceil.",
+ refl::DefaultValue(ffi::String("round")))
.def_ro("cubic_alpha", &Resize2DAttrs::cubic_alpha,
- "Spline Coefficient for Bicubic Interpolation")
+ "Spline Coefficient for Bicubic Interpolation", refl::DefaultValue(-0.75))
.def_ro("cubic_exclude", &Resize2DAttrs::cubic_exclude,
- "Flag to exclude exterior of the image during bicubic interpolation")
+ "Flag to exclude exterior of the image during bicubic interpolation",
+ refl::DefaultValue(0))
.def_ro("extrapolation_value", &Resize2DAttrs::extrapolation_value,
- "Value to return when roi is outside of the image")
+ "Value to return when roi is outside of the image", refl::DefaultValue(0.0))
.def_ro(
"out_dtype", &Resize2DAttrs::out_dtype,
"The dtype of the output tensor. It it is not specified, the output will have the same "
- "dtype as input if not specified.");
+ "dtype as input if not specified.",
+ refl::DefaultValue(ffi::Optional<DLDataType>{}));
}
TVM_FFI_DECLARE_OBJECT_INFO_FINAL("relax.attrs.Resize2DAttrs", Resize2DAttrs, AttrsNode);
}; // struct Resize2dAttrs
@@ -99,30 +105,36 @@ struct Resize3DAttrs : public AttrsNode {
"Dimension ordering of input data. Can be 'NCDHW', 'NDHWC', etc."
"'N', 'C', 'D', 'H', 'W' stands for batch, channel, depth, height, and width"
"dimensions respectively. Resize is applied on the 'D', 'H' and"
- "'W' dimensions.")
+ "'W' dimensions.",
+ refl::DefaultValue(ffi::String("NCDHW")))
.def_ro("method", &Resize3DAttrs::method,
"Specify the mode to use for scaling."
"nearest_neighbor - Nearest Neighbor"
"linear - Trilinear Interpolation"
- "cubic - Tricubic Interpolation")
+ "cubic - Tricubic Interpolation",
+ refl::DefaultValue(ffi::String("linear")))
.def_ro("coordinate_transformation_mode", &Resize3DAttrs::coordinate_transformation_mode,
"Describes how to transform the coordinate in the resized tensor"
"to the coordinate in the original tensor."
"Refer to the ONNX Resize operator specification for details"
- "Available options are half_pixel, align_corners and asymmetric")
+ "Available options are half_pixel, align_corners and asymmetric",
+ refl::DefaultValue(ffi::String("half_pixel")))
.def_ro("rounding_method", &Resize3DAttrs::rounding_method,
"indicates how to find the \"nearest\" pixel in nearest_neighbor method"
- "Available options are round, floor, and ceil.")
+ "Available options are round, floor, and ceil.",
+ refl::DefaultValue(ffi::String("")))
.def_ro("cubic_alpha", &Resize3DAttrs::cubic_alpha,
- "Spline Coefficient for Tricubic Interpolation")
+ "Spline Coefficient for Tricubic Interpolation", refl::DefaultValue(-0.75))
.def_ro("cubic_exclude", &Resize3DAttrs::cubic_exclude,
- "Flag to exclude exterior of the image during tricubic interpolation")
+ "Flag to exclude exterior of the image during tricubic interpolation",
+ refl::DefaultValue(0))
.def_ro("extrapolation_value", &Resize3DAttrs::extrapolation_value,
- "Value to return when roi is outside of the image")
+ "Value to return when roi is outside of the image", refl::DefaultValue(0.0))
.def_ro(
"out_dtype", &Resize3DAttrs::out_dtype,
"The dtype of the output tensor. It it is not specified, the output will have the same "
- "dtype as input if not specified.");
+ "dtype as input if not specified.",
+ refl::DefaultValue(ffi::Optional<DLDataType>{}));
}
TVM_FFI_DECLARE_OBJECT_INFO_FINAL("relax.attrs.Resize3DAttrs", Resize3DAttrs, AttrsNode);
}; // struct Resize3DAttrs
@@ -138,13 +150,17 @@ struct GridSampleAttrs : public AttrsNode {
namespace refl = tvm::ffi::reflection;
refl::ObjectDef<GridSampleAttrs>()
.def_ro("method", &GridSampleAttrs::method,
- "Interpolation method. Can be 'nearest', 'bilinear', or 'bicubic'.")
+ "Interpolation method. Can be 'nearest', 'bilinear', or 'bicubic'.",
+ refl::DefaultValue(ffi::String("bilinear")))
.def_ro("layout", &GridSampleAttrs::layout,
- "Dimension ordering of input data. Can be 'NCHW', 'NHWC', etc.")
+ "Dimension ordering of input data. Can be 'NCHW', 'NHWC', etc.",
+ refl::DefaultValue(ffi::String("NCHW")))
.def_ro("padding_mode", &GridSampleAttrs::padding_mode,
- "Padding mode for outside grid values. Can be 'zeros', 'border', or 'reflection'.")
+ "Padding mode for outside grid values. Can be 'zeros', 'border', or 'reflection'.",
+ refl::DefaultValue(ffi::String("zeros")))
.def_ro("align_corners", &GridSampleAttrs::align_corners,
- "If True, the corner pixels of the input and output tensors are aligned.");
+ "If True, the corner pixels of the input and output tensors are aligned.",
+ refl::DefaultValue(false));
}
TVM_FFI_DECLARE_OBJECT_INFO_FINAL("relax.attrs.GridSampleAttrs", GridSampleAttrs, AttrsNode);
}; // struct GridSampleAttrs
@@ -157,7 +173,8 @@ struct AffineGridAttrs : public AttrsNode {
namespace refl = tvm::ffi::reflection;
refl::ObjectDef<AffineGridAttrs>().def_ro(
"align_corners", &AffineGridAttrs::align_corners,
- "If True, normalized grid coordinates map to corner pixels; otherwise to pixel centers.");
+ "If True, normalized grid coordinates map to corner pixels; otherwise to pixel centers.",
+ refl::DefaultValue(true));
}
TVM_FFI_DECLARE_OBJECT_INFO_FINAL("relax.attrs.AffineGridAttrs", AffineGridAttrs, AttrsNode);
}; // struct AffineGridAttrs
diff --git a/include/tvm/relax/attrs/index.h b/include/tvm/relax/attrs/index.h
index 6133a6f580..f275f67bb2 100644
--- a/include/tvm/relax/attrs/index.h
+++ b/include/tvm/relax/attrs/index.h
@@ -37,13 +37,27 @@ struct TakeAttrs : public AttrsNode {
static void RegisterReflection() {
namespace refl = tvm::ffi::reflection;
refl::ObjectDef<TakeAttrs>()
- .def_ro("axis", &TakeAttrs::axis, "The axis over which to select values.")
+ .def_ro("axis", &TakeAttrs::axis, "The axis over which to select values.",
+ refl::DefaultValue(ffi::Optional<int64_t>{}))
.def_ro("mode", &TakeAttrs::mode, "The mode for handling out-of-bounds indices.",
refl::DefaultValue("fast"));
}
TVM_FFI_DECLARE_OBJECT_INFO_FINAL("relax.attrs.TakeAttrs", TakeAttrs, AttrsNode);
}; // struct TakeAttrs
+/*! \brief Attributes used in take_backward operator */
+struct TakeBackwardAttrs : public AttrsNode {
+ ffi::Optional<int64_t> axis;
+
+ static void RegisterReflection() {
+ namespace refl = tvm::ffi::reflection;
+ refl::ObjectDef<TakeBackwardAttrs>().def_ro("axis", &TakeBackwardAttrs::axis,
+ "The axis over which values were selected.",
+ refl::DefaultValue(ffi::Optional<int64_t>{}));
+ }
+ TVM_FFI_DECLARE_OBJECT_INFO_FINAL("relax.attrs.TakeBackwardAttrs", TakeBackwardAttrs, AttrsNode);
+};
+
/*! \brief Attributes used in strided_slice operator */
struct StridedSliceAttrs : public AttrsNode {
bool assume_inbound;
diff --git a/include/tvm/relax/attrs/linear_algebra.h b/include/tvm/relax/attrs/linear_algebra.h
index 34605bdabf..856bee34e7 100644
--- a/include/tvm/relax/attrs/linear_algebra.h
+++ b/include/tvm/relax/attrs/linear_algebra.h
@@ -36,7 +36,8 @@ struct MatmulAttrs : public AttrsNode {
static void RegisterReflection() {
namespace refl = tvm::ffi::reflection;
refl::ObjectDef<MatmulAttrs>().def_ro("out_dtype", &MatmulAttrs::out_dtype,
- "The data type of the output tensor");
+ "The data type of the output tensor",
+ refl::DefaultValue(ffi::Optional<DLDataType>{}));
}
TVM_FFI_DECLARE_OBJECT_INFO_FINAL("relax.attrs.MatmulAttrs", MatmulAttrs, AttrsNode);
}; // struct MatmulAttrs
diff --git a/include/tvm/relax/attrs/manipulate.h b/include/tvm/relax/attrs/manipulate.h
index 6a3061b112..1affa55b98 100644
--- a/include/tvm/relax/attrs/manipulate.h
+++ b/include/tvm/relax/attrs/manipulate.h
@@ -38,7 +38,8 @@ struct ConcatAttrs : public AttrsNode {
namespace refl = tvm::ffi::reflection;
refl::ObjectDef<ConcatAttrs>().def_ro("axis", &ConcatAttrs::axis,
"The axis at which the input arrays are concatenated."
- "Should lie in range `[-ndim, ndim)`.");
+ "Should lie in range `[-ndim, ndim)`.",
+ refl::DefaultValue(ffi::Optional<int64_t>{0}));
}
TVM_FFI_DECLARE_OBJECT_INFO_FINAL("relax.attrs.ConcatAttrs", ConcatAttrs, AttrsNode);
}; // struct ConcatAttrs
@@ -73,7 +74,8 @@ struct LayoutTransformAttrs : public AttrsNode {
.def_ro(
"pad_value", &LayoutTransformAttrs::pad_value,
"The specific value to be used to pad if the layout transform would result in implicit "
- "padding. If not specified, the compiler is free to choose any value.");
+ "padding. If not specified, the compiler is free to choose any value.",
+ refl::DefaultValue(ffi::Optional<PrimExpr>{}));
}
TVM_FFI_DECLARE_OBJECT_INFO_FINAL("relax.attrs.LayoutTransformAttrs", LayoutTransformAttrs,
AttrsNode);
@@ -86,7 +88,8 @@ struct PermuteDimsAttrs : public AttrsNode {
static void RegisterReflection() {
namespace refl = tvm::ffi::reflection;
refl::ObjectDef<PermuteDimsAttrs>().def_ro(
- "axes", &PermuteDimsAttrs::axes, "The target axes order, reverse order if not specified.");
+ "axes", &PermuteDimsAttrs::axes, "The target axes order, reverse order if not specified.",
+ refl::DefaultValue(ffi::Optional<ffi::Array<int64_t>>{}));
}
TVM_FFI_DECLARE_OBJECT_INFO_FINAL("relax.attrs.PermuteDimsAttrs", PermuteDimsAttrs, AttrsNode);
}; // struct PermuteDimsAttrs
@@ -101,7 +104,7 @@ struct SplitAttrs : public AttrsNode {
refl::ObjectDef<SplitAttrs>()
.def_ro("indices_or_sections", &SplitAttrs::indices_or_sections,
"The input array of indices or the number of split sections.")
- .def_ro("axis", &SplitAttrs::axis, "The axis to be splitted");
+ .def_ro("axis", &SplitAttrs::axis, "The axis to be splitted", refl::DefaultValue(0));
}
TVM_FFI_DECLARE_OBJECT_INFO_FINAL("relax.attrs.SplitAttrs", SplitAttrs, AttrsNode);
}; // struct SplitAttrs
@@ -112,11 +115,13 @@ struct SqueezeAttrs : public AttrsNode {
static void RegisterReflection() {
namespace refl = tvm::ffi::reflection;
- refl::ObjectDef<SqueezeAttrs>().def_ro("axis", &SqueezeAttrs::axis,
- "The axis to squeeze in the input tensor."
- "If `axis = None`, all axis of dimension 1 get squeezed;"
- "Else, the dimension in axes get squeezed."
- "It is an error if an axis does not has dimension 1.");
+ refl::ObjectDef<SqueezeAttrs>().def_ro(
+ "axis", &SqueezeAttrs::axis,
+ "The axis to squeeze in the input tensor."
+ "If `axis = None`, all axis of dimension 1 get squeezed;"
+ "Else, the dimension in axes get squeezed."
+ "It is an error if an axis does not has dimension 1.",
+ refl::DefaultValue(ffi::Optional<ffi::Array<int64_t>>{}));
}
TVM_FFI_DECLARE_OBJECT_INFO_FINAL("relax.attrs.SqueezeAttrs", SqueezeAttrs, AttrsNode);
}; // struct SqueezeAttrs
@@ -132,7 +137,8 @@ struct StackAttrs : public AttrsNode {
"The axis along which to stack the input tensors. "
"The axis will be inserted at this position in the output, "
"so it must be in range [-ndim-1, ndim] where ndim is the "
- "number of dimensions of the input tensors.");
+ "number of dimensions of the input tensors.",
+ refl::DefaultValue(ffi::Optional<int64_t>{0}));
}
TVM_FFI_DECLARE_OBJECT_INFO_FINAL("relax.attrs.StackAttrs", StackAttrs, AttrsNode);
}; // struct StackAttrs
@@ -149,7 +155,8 @@ struct RepeatAttrs : public AttrsNode {
.def_ro("axis", &RepeatAttrs::axis,
"The axis along which to repeat values. The negative numbers are interpreted "
"counting from the backward. By default, use the flattened input array, and "
- "return a flat output array.");
+ "return a flat output array.",
+ refl::DefaultValue(ffi::Optional<int64_t>{}));
}
TVM_FFI_DECLARE_OBJECT_INFO_FINAL("relax.attrs.RepeatAttrs", RepeatAttrs, AttrsNode);
}; // struct RepeatAttrs
@@ -187,9 +194,9 @@ struct ReverseSequenceAttrs : public AttrsNode {
namespace refl = tvm::ffi::reflection;
refl::ObjectDef<ReverseSequenceAttrs>()
.def_ro("seq_axis", &ReverseSequenceAttrs::seq_axis,
- "The axis along which to reverse variable length slices.")
- .def_ro("batch_axis", &ReverseSequenceAttrs::batch_axis,
- "The axis that indexes the batch.");
+ "The axis along which to reverse variable length slices.", refl::DefaultValue(1))
+ .def_ro("batch_axis", &ReverseSequenceAttrs::batch_axis, "The axis that indexes the batch.",
+ refl::DefaultValue(0));
}
TVM_FFI_DECLARE_OBJECT_INFO_FINAL("relax.attrs.ReverseSequenceAttrs", ReverseSequenceAttrs,
AttrsNode);
@@ -244,7 +251,8 @@ struct MeshgridAttrs : public AttrsNode {
static void RegisterReflection() {
namespace refl = tvm::ffi::reflection;
refl::ObjectDef<MeshgridAttrs>().def_ro("indexing", &MeshgridAttrs::indexing,
- "Specifies how the grid dimensions are ordered.");
+ "Specifies how the grid dimensions are ordered.",
+ refl::DefaultValue(ffi::String("ij")));
}
TVM_FFI_DECLARE_OBJECT_INFO_FINAL("relax.attrs.MeshgridAttrs", MeshgridAttrs, AttrsNode);
};
diff --git a/include/tvm/relax/attrs/nn.h b/include/tvm/relax/attrs/nn.h
index 1e97c8e5f0..e28a3c80d8 100644
--- a/include/tvm/relax/attrs/nn.h
+++ b/include/tvm/relax/attrs/nn.h
@@ -43,32 +43,39 @@ struct Conv1DAttrs : public AttrsNode {
static void RegisterReflection() {
namespace refl = tvm::ffi::reflection;
refl::ObjectDef<Conv1DAttrs>()
- .def_ro("strides", &Conv1DAttrs::strides, "Specifies the strides of the convolution.")
+ .def_ro("strides", &Conv1DAttrs::strides, "Specifies the strides of the convolution.",
+ refl::DefaultValue(ffi::Array<int64_t>{1}))
.def_ro("padding", &Conv1DAttrs::padding,
"If padding is non-zero, then the input is implicitly zero-padded"
"Padding support both symmetric and asymmetric as"
"one int : same padding used on both sides"
- "two int : padding width in the order of (left, right)")
+ "two int : padding width in the order of (left, right)",
+ refl::DefaultValue(ffi::Array<int64_t>{0, 0}))
.def_ro("dilation", &Conv1DAttrs::dilation,
- "Specifies the dilation rate to use for dilated convolution.")
+ "Specifies the dilation rate to use for dilated convolution.",
+ refl::DefaultValue(ffi::Array<int64_t>{1}))
.def_ro("groups", &Conv1DAttrs::groups,
"Number of groups to split the input into for grouped convolution. The number of "
"input and "
- "output channels should be divisible by the number of groups.")
+ "output channels should be divisible by the number of groups.",
+ refl::DefaultValue(1))
.def_ro("data_layout", &Conv1DAttrs::data_layout,
"Dimension ordering of input data. Can be 'NCW', 'NWC', etc."
"'N', 'C', 'W' stands for batch, channel, width"
- "dimensions respectively. Convolution is applied on the 'W' dimensions.")
+ "dimensions respectively. Convolution is applied on the 'W' dimensions.",
+ refl::DefaultValue(ffi::String("NCW")))
.def_ro("kernel_layout", &Conv1DAttrs::kernel_layout,
"Dimension ordering of weight. Can be 'OIW', 'IOW', etc."
"'O', 'I', 'W' stands for num_filter, input_channel, and width"
- "dimensions respectively.")
+ "dimensions respectively.",
+ refl::DefaultValue(ffi::String("OIW")))
.def_ro("out_layout", &Conv1DAttrs::out_layout,
"Dimension ordering of output. Can be 'NCW', 'NWC', etc."
"'N', 'C', 'W' stands for batch, channel, and width"
"dimensions respectively. Default to be same as input layout.")
.def_ro("out_dtype", &Conv1DAttrs::out_dtype,
- "Output data type, set to explicit type under mixed precision setting");
+ "Output data type, set to explicit type under mixed precision setting",
+ refl::DefaultValue(ffi::Optional<DLDataType>{}));
}
TVM_FFI_DECLARE_OBJECT_INFO_FINAL("relax.attrs.Conv1DAttrs", Conv1DAttrs, AttrsNode);
}; // struct Conv1dAttrs
@@ -87,34 +94,41 @@ struct Conv2DAttrs : public AttrsNode {
static void RegisterReflection() {
namespace refl = tvm::ffi::reflection;
refl::ObjectDef<Conv2DAttrs>()
- .def_ro("strides", &Conv2DAttrs::strides, "Specifies the strides of the convolution.")
+ .def_ro("strides", &Conv2DAttrs::strides, "Specifies the strides of the convolution.",
+ refl::DefaultValue(ffi::Array<int64_t>{1, 1}))
.def_ro("padding", &Conv2DAttrs::padding,
"If padding is non-zero, then the input is implicitly zero-padded"
"Padding support both symmetric and asymmetric as"
"one int : same padding used on all sides"
"two int : bottom, right will use same padding as top, left"
- "four int : padding width in the order of (top, left, bottom, right)")
+ "four int : padding width in the order of (top, left, bottom, right)",
+ refl::DefaultValue(ffi::Array<int64_t>{0, 0, 0, 0}))
.def_ro("dilation", &Conv2DAttrs::dilation,
- "Specifies the dilation rate to use for dilated convolution.")
+ "Specifies the dilation rate to use for dilated convolution.",
+ refl::DefaultValue(ffi::Array<int64_t>{1, 1}))
.def_ro("groups", &Conv2DAttrs::groups,
"Number of groups to split the input into for grouped convolution. The number of "
"input and "
- "output channels should be divisible by the number of groups.")
+ "output channels should be divisible by the number of groups.",
+ refl::DefaultValue(1))
.def_ro("data_layout", &Conv2DAttrs::data_layout,
"Dimension ordering of input data. Can be 'NCHW', 'NHWC', etc."
"'N', 'C', 'H', 'W' stands for batch, channel, height, and width"
"dimensions respectively. Convolution is applied on the 'H' and"
- "'W' dimensions.")
+ "'W' dimensions.",
+ refl::DefaultValue(ffi::String("NCHW")))
.def_ro("kernel_layout", &Conv2DAttrs::kernel_layout,
"Dimension ordering of weight. Can be 'OIHW', 'OIHW16o16i', etc."
"'O', 'I', 'H', 'W' stands for num_filter, input_channel, height, and width"
- "dimensions respectively.")
+ "dimensions respectively.",
+ refl::DefaultValue(ffi::String("OIHW")))
.def_ro("out_layout", &Conv2DAttrs::out_layout,
"Dimension ordering of output. Can be 'NCHW', 'NHWC', etc."
"'N', 'C', 'H', 'W' stands for batch, channel, height, and width"
"dimensions respectively. Default to be same as input layout.")
.def_ro("out_dtype", &Conv2DAttrs::out_dtype,
- "Output data type, set to explicit type under mixed precision setting");
+ "Output data type, set to explicit type under mixed precision setting",
+ refl::DefaultValue(ffi::Optional<DLDataType>{}));
}
TVM_FFI_DECLARE_OBJECT_INFO_FINAL("relax.attrs.Conv2DAttrs", Conv2DAttrs, AttrsNode);
}; // struct Conv2dAttrs
@@ -133,36 +147,43 @@ struct Conv3DAttrs : public AttrsNode {
static void RegisterReflection() {
namespace refl = tvm::ffi::reflection;
refl::ObjectDef<Conv3DAttrs>()
- .def_ro("strides", &Conv3DAttrs::strides, "Specifies the strides of the convolution.")
+ .def_ro("strides", &Conv3DAttrs::strides, "Specifies the strides of the convolution.",
+ refl::DefaultValue(ffi::Array<int64_t>{1, 1, 1}))
.def_ro(
"padding", &Conv3DAttrs::padding,
"If padding is non-zero, then the input is implicitly zero-padded"
"Padding support both symmetric and asymmetric as"
"one int : same padding used on all sides"
"two int : bottom, right will use same padding as top, left"
- "four int : padding width in the order of (forward, back, top, left, bottom, right)")
+ "four int : padding width in the order of (forward, back, top, left, bottom, right)",
+ refl::DefaultValue(ffi::Array<int64_t>{0, 0, 0, 0, 0, 0}))
.def_ro("dilation", &Conv3DAttrs::dilation,
- "Specifies the dilation rate to use for dilated convolution.")
+ "Specifies the dilation rate to use for dilated convolution.",
+ refl::DefaultValue(ffi::Array<int64_t>{1, 1, 1}))
.def_ro("groups", &Conv3DAttrs::groups,
"Number of groups to split the input into for grouped convolution. The number of "
"input and "
- "output channels should be divisible by the number of groups.")
+ "output channels should be divisible by the number of groups.",
+ refl::DefaultValue(1))
.def_ro("data_layout", &Conv3DAttrs::data_layout,
"Dimension ordering of input data. Can be 'NCDHW', 'NDHWC', etc."
"'N', 'C', 'D', 'H', 'W' stands for batch, channel, depth, height, and width"
"dimensions respectively. Convolution is applied on the 'D', 'H', and"
- "'W' dimensions.")
+ "'W' dimensions.",
+ refl::DefaultValue(ffi::String("NCDHW")))
.def_ro(
"kernel_layout", &Conv3DAttrs::kernel_layout,
"Dimension ordering of weight. Can be 'OIDHW', 'OIDHW16o16i', etc."
"'O', 'I', 'D', 'H', 'W' stands for num_filter, input_channel, depth, height, and width"
- "dimensions respectively.")
+ "dimensions respectively.",
+ refl::DefaultValue(ffi::String("OIDHW")))
.def_ro("out_layout", &Conv3DAttrs::out_layout,
"Dimension ordering of output. Can be 'NCDHW', 'NDHWC', etc."
"'N', 'C', 'D', 'H', 'W' stands for batch, channel, depth, height, and width"
"dimensions respectively. Default to be same as input layout.")
.def_ro("out_dtype", &Conv3DAttrs::out_dtype,
- "Output data type, set to explicit type under mixed precision setting");
+ "Output data type, set to explicit type under mixed precision setting",
+ refl::DefaultValue(ffi::Optional<DLDataType>{}));
}
TVM_FFI_DECLARE_OBJECT_INFO_FINAL("relax.attrs.Conv3DAttrs", Conv3DAttrs, AttrsNode);
}; // struct Conv3dAttrs
@@ -183,34 +204,42 @@ struct Conv1DTransposeAttrs : public AttrsNode {
namespace refl = tvm::ffi::reflection;
refl::ObjectDef<Conv1DTransposeAttrs>()
.def_ro("strides", &Conv1DTransposeAttrs::strides,
- "Specifies the strides of the convolution.")
+ "Specifies the strides of the convolution.",
+ refl::DefaultValue(ffi::Array<int64_t>{1}))
.def_ro("padding", &Conv1DTransposeAttrs::padding,
"If padding is non-zero, then the input is implicitly zero-padded"
"Padding support both symmetric and asymmetric as"
"one int : same padding used on both sides"
- "two int : padding width in the order of (left, right)")
+ "two int : padding width in the order of (left, right)",
+ refl::DefaultValue(ffi::Array<int64_t>{0, 0}))
.def_ro("output_padding", &Conv1DTransposeAttrs::output_padding,
- "Used to disambiguate the output shape.")
+ "Used to disambiguate the output shape.",
+ refl::DefaultValue(ffi::Array<int64_t>{0}))
.def_ro("dilation", &Conv1DTransposeAttrs::dilation,
- "Specifies the dilation rate to use for dilated convolution.")
+ "Specifies the dilation rate to use for dilated convolution.",
+ refl::DefaultValue(ffi::Array<int64_t>{1}))
.def_ro("groups", &Conv1DTransposeAttrs::groups,
"Number of groups to split the input into for grouped convolution. The number of "
"input and "
- "output channels should be divisible by the number of groups.")
+ "output channels should be divisible by the number of groups.",
+ refl::DefaultValue(1))
.def_ro("data_layout", &Conv1DTransposeAttrs::data_layout,
"Dimension ordering of input data. Can be 'NCW', 'NWC', etc."
"'N', 'C', 'W' stands for batch, channel, width"
- "dimensions respectively. Convolution is applied on the 'W' dimensions.")
+ "dimensions respectively. Convolution is applied on the 'W' dimensions.",
+ refl::DefaultValue(ffi::String("NCW")))
.def_ro("kernel_layout", &Conv1DTransposeAttrs::kernel_layout,
"Dimension ordering of weight. Can be 'OIW', 'IOW', etc."
"'O', 'I', 'W' stands for num_filter, input_channel, and width"
- "dimensions respectively.")
+ "dimensions respectively.",
+ refl::DefaultValue(ffi::String("IOW")))
.def_ro("out_layout", &Conv1DTransposeAttrs::out_layout,
"Dimension ordering of output. Can be 'NCW', 'NWC', etc."
"'N', 'C', 'W' stands for batch, channel, and width"
"dimensions respectively. Default to be same as input layout.")
.def_ro("out_dtype", &Conv1DTransposeAttrs::out_dtype,
- "Output data type, set to explicit type under mixed precision setting");
+ "Output data type, set to explicit type under mixed precision setting",
+ refl::DefaultValue(ffi::Optional<DLDataType>{}));
}
TVM_FFI_DECLARE_OBJECT_INFO_FINAL("relax.attrs.Conv1DTransposeAttrs", Conv1DTransposeAttrs,
AttrsNode);
@@ -232,36 +261,44 @@ struct Conv2DTransposeAttrs : public AttrsNode {
namespace refl = tvm::ffi::reflection;
refl::ObjectDef<Conv2DTransposeAttrs>()
.def_ro("strides", &Conv2DTransposeAttrs::strides,
- "Specifies the strides of the convolution.")
+ "Specifies the strides of the convolution.",
+ refl::DefaultValue(ffi::Array<int64_t>{1, 1}))
.def_ro("padding", &Conv2DTransposeAttrs::padding,
"If padding is non-zero, then the input is implicitly zero-padded"
"Padding support both symmetric and asymmetric as"
"one int : same padding used on all sides"
"two int : bottom, right will use same padding as top, left"
- "four int : padding width in the order of (top, left, bottom, right)")
+ "four int : padding width in the order of (top, left, bottom, right)",
+ refl::DefaultValue(ffi::Array<int64_t>{0, 0, 0, 0}))
.def_ro("output_padding", &Conv2DTransposeAttrs::output_padding,
- "Used to disambiguate the output shape.")
+ "Used to disambiguate the output shape.",
+ refl::DefaultValue(ffi::Array<int64_t>{0, 0}))
.def_ro("dilation", &Conv2DTransposeAttrs::dilation,
- "Specifies the dilation rate to use for dilated convolution.")
+ "Specifies the dilation rate to use for dilated convolution.",
+ refl::DefaultValue(ffi::Array<int64_t>{1, 1}))
.def_ro("groups", &Conv2DTransposeAttrs::groups,
"Number of groups to split the input into for grouped convolution. The number of "
"input and "
- "output channels should be divisible by the number of groups.")
+ "output channels should be divisible by the number of groups.",
+ refl::DefaultValue(1))
.def_ro("data_layout", &Conv2DTransposeAttrs::data_layout,
"Dimension ordering of input data. Can be 'NCHW', 'NHWC', etc."
"'N', 'C', 'H', 'W' stands for batch, channel, height, and width"
"dimensions respectively. Convolution is applied on the 'H' and"
- "'W' dimensions.")
+ "'W' dimensions.",
+ refl::DefaultValue(ffi::String("NCHW")))
.def_ro("kernel_layout", &Conv2DTransposeAttrs::kernel_layout,
"Dimension ordering of weight. Can be 'OIHW', 'OIHW16o16i', etc."
"'O', 'I', 'H', 'W' stands for num_filter, input_channel, height, and width"
- "dimensions respectively.")
+ "dimensions respectively.",
+ refl::DefaultValue(ffi::String("IOHW")))
.def_ro("out_layout", &Conv2DTransposeAttrs::out_layout,
"Dimension ordering of output. Can be 'NCHW', 'NHWC', etc."
"'N', 'C', 'H', 'W' stands for batch, channel, height, and width"
"dimensions respectively. Default to be same as input layout.")
.def_ro("out_dtype", &Conv2DTransposeAttrs::out_dtype,
- "Output data type, set to explicit type under mixed precision setting");
+ "Output data type, set to explicit type under mixed precision setting",
+ refl::DefaultValue(ffi::Optional<DLDataType>{}));
}
TVM_FFI_DECLARE_OBJECT_INFO_FINAL("relax.attrs.Conv2DTransposeAttrs", Conv2DTransposeAttrs,
AttrsNode);
@@ -283,38 +320,46 @@ struct Conv3DTransposeAttrs : public AttrsNode {
namespace refl = tvm::ffi::reflection;
refl::ObjectDef<Conv3DTransposeAttrs>()
.def_ro("strides", &Conv3DTransposeAttrs::strides,
- "Specifies the strides of the convolution.")
+ "Specifies the strides of the convolution.",
+ refl::DefaultValue(ffi::Array<int64_t>{1, 1, 1}))
.def_ro("padding", &Conv3DTransposeAttrs::padding,
"If padding is non-zero, then the input is implicitly zero-padded"
"Padding support both symmetric and asymmetric as"
"one int : same padding used on all sides"
"three int : back/bottom/right will use same padding as front/top/left"
- "six int : padding width in the order of (front, top, left, back, bottom, right)")
+ "six int : padding width in the order of (front, top, left, back, bottom, right)",
+ refl::DefaultValue(ffi::Array<int64_t>{0, 0, 0, 0, 0, 0}))
.def_ro("output_padding", &Conv3DTransposeAttrs::output_padding,
- "Used to disambiguate the output shape.")
+ "Used to disambiguate the output shape.",
+ refl::DefaultValue(ffi::Array<int64_t>{0, 0, 0}))
.def_ro("dilation", &Conv3DTransposeAttrs::dilation,
- "Specifies the dilation rate to use for dilated convolution.")
+ "Specifies the dilation rate to use for dilated convolution.",
+ refl::DefaultValue(ffi::Array<int64_t>{1, 1, 1}))
.def_ro("groups", &Conv3DTransposeAttrs::groups,
"Number of groups to split the input into for grouped convolution. The number of "
"input and "
- "output channels should be divisible by the number of groups.")
+ "output channels should be divisible by the number of groups.",
+ refl::DefaultValue(1))
.def_ro("data_layout", &Conv3DTransposeAttrs::data_layout,
"Dimension ordering of input data. Can be 'NCDHW', 'NDHWC', etc."
"'N', 'C', 'D', 'H', 'W' stands for batch, channel, depth, height, and width"
"dimensions respectively. Convolution is applied on the 'D', 'H', and"
- "'W' dimensions.")
+ "'W' dimensions.",
+ refl::DefaultValue(ffi::String("NCDHW")))
.def_ro(
"kernel_layout", &Conv3DTransposeAttrs::kernel_layout,
"Dimension ordering of weight. Can be 'IODHW', etc."
"'I', 'O', 'D', 'H', 'W' stands for input_channel, output_channel, depth, height, and "
"width"
- "dimensions respectively.")
+ "dimensions respectively.",
+ refl::DefaultValue(ffi::String("IODHW")))
.def_ro("out_layout", &Conv3DTransposeAttrs::out_layout,
"Dimension ordering of output. Can be 'NCDHW', 'NDHWC', etc."
"'N', 'C', 'D', 'H', 'W' stands for batch, channel, depth, height, and width"
"dimensions respectively. Default to be same as input layout.")
.def_ro("out_dtype", &Conv3DTransposeAttrs::out_dtype,
- "Output data type, set to explicit type under mixed precision setting");
+ "Output data type, set to explicit type under mixed precision setting",
+ refl::DefaultValue(ffi::Optional<DLDataType>{}));
}
TVM_FFI_DECLARE_OBJECT_INFO_FINAL("relax.attrs.Conv3DTransposeAttrs", Conv3DTransposeAttrs,
AttrsNode);
@@ -334,18 +379,23 @@ struct Pool1DAttrs : public AttrsNode {
static void RegisterReflection() {
namespace refl = tvm::ffi::reflection;
refl::ObjectDef<Pool1DAttrs>()
- .def_ro("pool_size", &Pool1DAttrs::pool_size, "Size of the pooling windows.")
- .def_ro("strides", &Pool1DAttrs::strides, "Specifies the strides of the convolution.")
- .def_ro("dilation", &Pool1DAttrs::dilation, "Specifies the dilation of the convolution.")
+ .def_ro("pool_size", &Pool1DAttrs::pool_size, "Size of the pooling windows.",
+ refl::DefaultValue(ffi::Array<int64_t>{1}))
+ .def_ro("strides", &Pool1DAttrs::strides, "Specifies the strides of the convolution.",
+ refl::DefaultValue(ffi::Array<int64_t>{1}))
+ .def_ro("dilation", &Pool1DAttrs::dilation, "Specifies the dilation of the convolution.",
+ refl::DefaultValue(ffi::Array<int64_t>{1}))
.def_ro("padding", &Pool1DAttrs::padding,
"If padding is non-zero, then the input is implicitly zero-padded"
"Padding support both symmetric and asymmetric as"
"one int : same padding used on all sides"
- "two int : padding width in the order of (left, right)")
+ "two int : padding width in the order of (left, right)",
+ refl::DefaultValue(ffi::Array<int64_t>{0, 0}))
.def_ro(
"ceil_mode", &Pool1DAttrs::ceil_mode,
"A boolean indicating if use ceil or floor to compute the output shape. By using ceil, "
- "every element in the input tensor will be covered by a sliding window.")
+ "every element in the input tensor will be covered by a sliding window.",
+ refl::DefaultValue(false))
.def_ro("count_include_pad", &Pool1DAttrs::count_include_pad,
"When true, will include padding to compute the average")
.def_ro("layout", &Pool1DAttrs::layout,
@@ -375,26 +425,32 @@ struct Pool2DAttrs : public AttrsNode {
static void RegisterReflection() {
namespace refl = tvm::ffi::reflection;
refl::ObjectDef<Pool2DAttrs>()
- .def_ro("pool_size", &Pool2DAttrs::pool_size, "Size of the pooling windows.")
- .def_ro("strides", &Pool2DAttrs::strides, "Specifies the strides of the convolution.")
- .def_ro("dilation", &Pool2DAttrs::dilation, "Specifies the dilation of the convolution.")
+ .def_ro("pool_size", &Pool2DAttrs::pool_size, "Size of the pooling windows.",
+ refl::DefaultValue(ffi::Array<int64_t>{1, 1}))
+ .def_ro("strides", &Pool2DAttrs::strides, "Specifies the strides of the convolution.",
+ refl::DefaultValue(ffi::Array<int64_t>{1, 1}))
+ .def_ro("dilation", &Pool2DAttrs::dilation, "Specifies the dilation of the convolution.",
+ refl::DefaultValue(ffi::Array<int64_t>{1, 1}))
.def_ro("padding", &Pool2DAttrs::padding,
"If padding is non-zero, then the input is implicitly zero-padded"
"Padding support both symmetric and asymmetric as"
"one int : same padding used on all sides"
"two int : bottom, right will use same padding as top, left"
- "four int : padding width in the order of (top, left, bottom, right)")
+ "four int : padding width in the order of (top, left, bottom, right)",
+ refl::DefaultValue(ffi::Array<int64_t>{0, 0, 0, 0}))
.def_ro(
"ceil_mode", &Pool2DAttrs::ceil_mode,
"A boolean indicating if use ceil or floor to compute the output shape. By using ceil, "
- "every element in the input tensor will be covered by a sliding window.")
+ "every element in the input tensor will be covered by a sliding window.",
+ refl::DefaultValue(false))
.def_ro("count_include_pad", &Pool2DAttrs::count_include_pad,
- "When true, will include padding to compute the average")
+ "When true, will include padding to compute the average", refl::DefaultValue(false))
.def_ro("layout", &Pool2DAttrs::layout,
"Dimension ordering of input data. Can be 'NCHW', 'NHWC', etc."
"'N', 'C', 'H', 'W' stands for batch, channel, height, and width"
"dimensions respectively. Pooling is applied on the 'H' and"
- "'W' dimensions.")
+ "'W' dimensions.",
+ refl::DefaultValue(ffi::String("NCHW")))
.def_ro("out_layout", &Pool2DAttrs::out_layout,
"Dimension ordering of output data. Can be 'NCHW', 'NHWC', etc."
"'N', 'C', 'H', 'W' stands for batch, channel, height, and width"
@@ -418,26 +474,32 @@ struct Pool3DAttrs : public AttrsNode {
static void RegisterReflection() {
namespace refl = tvm::ffi::reflection;
refl::ObjectDef<Pool3DAttrs>()
- .def_ro("pool_size", &Pool3DAttrs::pool_size, "Size of the pooling windows.")
- .def_ro("strides", &Pool3DAttrs::strides, "Specifies the strides of the convolution.")
- .def_ro("dilation", &Pool3DAttrs::dilation, "Specifies the dilation of the convolution.")
+ .def_ro("pool_size", &Pool3DAttrs::pool_size, "Size of the pooling windows.",
+ refl::DefaultValue(ffi::Array<int64_t>{1, 1, 1}))
+ .def_ro("strides", &Pool3DAttrs::strides, "Specifies the strides of the convolution.",
+ refl::DefaultValue(ffi::Array<int64_t>{1, 1, 1}))
+ .def_ro("dilation", &Pool3DAttrs::dilation, "Specifies the dilation of the convolution.",
+ refl::DefaultValue(ffi::Array<int64_t>{1, 1, 1}))
.def_ro("padding", &Pool3DAttrs::padding,
"If padding is non-zero, then the input is implicitly zero-padded"
"Padding support both symmetric and asymmetric as"
"one int : same padding used on all sides"
"three int : back, bottom, right will use same padding as front, top, left"
- "four int : padding width in the order of (front, top, left, back, bottom, right)")
+ "four int : padding width in the order of (front, top, left, back, bottom, right)",
+ refl::DefaultValue(ffi::Array<int64_t>{0, 0, 0, 0, 0, 0}))
.def_ro(
"ceil_mode", &Pool3DAttrs::ceil_mode,
"A boolean indicating if use ceil or floor to compute the output shape. By using ceil, "
- "every element in the input tensor will be covered by a sliding window.")
+ "every element in the input tensor will be covered by a sliding window.",
+ refl::DefaultValue(false))
.def_ro("count_include_pad", &Pool3DAttrs::count_include_pad,
"When true, will include padding to compute the average")
.def_ro("layout", &Pool3DAttrs::layout,
"Dimension ordering of input data. Can be 'NCDHW', 'NDHWC', etc."
"'N', 'C', 'D', 'H', 'W' stands for batch, channel, depth, height, and width"
"dimensions respectively. Pooling is applied on the 'D', 'H' and"
- "'W' dimensions.")
+ "'W' dimensions.",
+ refl::DefaultValue(ffi::String("NCDHW")))
.def_ro("out_layout", &Pool3DAttrs::out_layout,
"Dimension ordering of output data. Can be 'NCDHW', 'NDHWC', etc."
"'N', 'C', 'D', 'H', 'W' stands for batch, channel, depth, height, and width"
@@ -456,12 +518,14 @@ struct AdaptivePool1DAttrs : public AttrsNode {
static void RegisterReflection() {
namespace refl = tvm::ffi::reflection;
refl::ObjectDef<AdaptivePool1DAttrs>()
- .def_ro("output_size", &AdaptivePool1DAttrs::output_size, "Output width.")
+ .def_ro("output_size", &AdaptivePool1DAttrs::output_size, "Output width.",
+ refl::DefaultValue(ffi::Optional<ffi::Array<int64_t>>{}))
.def_ro("layout", &AdaptivePool1DAttrs::layout,
"Dimension ordering of input data. Can be 'NCW', 'NWC', etc."
"'N', 'C', 'W' stands for batch, channel and width"
"dimensions respectively. Pooling is applied on the"
- "'W' dimensions.")
+ "'W' dimensions.",
+ refl::DefaultValue(ffi::String("NCW")))
.def_ro("out_layout", &AdaptivePool1DAttrs::out_layout,
"Dimension ordering of output data. Can be 'NCW', 'NWC', etc."
"'N', 'C', 'W' stands for batch, channel and width"
@@ -481,12 +545,14 @@ struct AdaptivePool2DAttrs : public AttrsNode {
static void RegisterReflection() {
namespace refl = tvm::ffi::reflection;
refl::ObjectDef<AdaptivePool2DAttrs>()
- .def_ro("output_size", &AdaptivePool2DAttrs::output_size, "Output height and width.")
+ .def_ro("output_size", &AdaptivePool2DAttrs::output_size, "Output height and width.",
+ refl::DefaultValue(ffi::Optional<ffi::Array<int64_t>>{}))
.def_ro("layout", &AdaptivePool2DAttrs::layout,
"Dimension ordering of input data. Can be 'NCHW', 'NHWC', etc."
"'N', 'C', 'H', 'W' stands for batch, channel, height, and width"
"dimensions respectively. Pooling is applied on the 'H' and"
- "'W' dimensions.")
+ "'W' dimensions.",
+ refl::DefaultValue(ffi::String("NCHW")))
.def_ro("out_layout", &AdaptivePool2DAttrs::out_layout,
"Dimension ordering of output data. Can be 'NCHW', 'NHWC', etc."
"'N', 'C', 'H', 'W' stands for batch, channel, height, and width"
@@ -506,12 +572,14 @@ struct AdaptivePool3DAttrs : public AttrsNode {
static void RegisterReflection() {
namespace refl = tvm::ffi::reflection;
refl::ObjectDef<AdaptivePool3DAttrs>()
- .def_ro("output_size", &AdaptivePool3DAttrs::output_size, "Output depth, height and width.")
+ .def_ro("output_size", &AdaptivePool3DAttrs::output_size, "Output depth, height and width.",
+ refl::DefaultValue(ffi::Optional<ffi::Array<int64_t>>{}))
.def_ro("layout", &AdaptivePool3DAttrs::layout,
"Dimension ordering of input data. Can be 'NCDHW', 'NDHWC', etc."
"'N', 'C', 'D', 'H', 'W' stands for batch, channel, depth, height, and width"
"dimensions respectively. Pooling is applied on 'D', 'H' and"
- "'W' dimensions.")
+ "'W' dimensions.",
+ refl::DefaultValue(ffi::String("NCDHW")))
.def_ro("out_layout", &AdaptivePool3DAttrs::out_layout,
"Dimension ordering of output data. Can be 'NCDHW', 'NDHWC', etc."
"'N', 'C', 'D', 'H', 'W' stands for batch, channel, depth, height, and width"
@@ -529,7 +597,8 @@ struct SoftmaxAttrs : public AttrsNode {
static void RegisterReflection() {
namespace refl = tvm::ffi::reflection;
refl::ObjectDef<SoftmaxAttrs>().def_ro("axis", &SoftmaxAttrs::axis,
- "The axis to sum over when computing softmax.");
+ "The axis to sum over when computing softmax.",
+ refl::DefaultValue(-1));
}
TVM_FFI_DECLARE_OBJECT_INFO_FINAL("relax.attrs.SoftmaxAttrs", SoftmaxAttrs, AttrsNode);
};
@@ -541,7 +610,8 @@ struct LeakyReluAttrs : public AttrsNode {
static void RegisterReflection() {
namespace refl = tvm::ffi::reflection;
refl::ObjectDef<LeakyReluAttrs>().def_ro("alpha", &LeakyReluAttrs::alpha,
- "The slope of the negative part.");
+ "The slope of the negative part.",
+ refl::DefaultValue(0.01));
}
TVM_FFI_DECLARE_OBJECT_INFO_FINAL("relax.attrs.LeakyReluAttrs", LeakyReluAttrs, AttrsNode);
};
@@ -555,9 +625,11 @@ struct SoftplusAttrs : public AttrsNode {
namespace refl = tvm::ffi::reflection;
refl::ObjectDef<SoftplusAttrs>()
.def_ro("beta", &SoftplusAttrs::beta,
- "Scaling factor controlling the sharpness of the Softplus transition.")
+ "Scaling factor controlling the sharpness of the Softplus transition.",
+ refl::DefaultValue(1.0))
.def_ro("threshold", &SoftplusAttrs::threshold,
- "Value determining when to use linear approximation for numerical stability.");
+ "Value determining when to use linear approximation for numerical stability.",
+ refl::DefaultValue(20.0));
}
TVM_FFI_DECLARE_OBJECT_INFO_FINAL("relax.attrs.SoftplusAttrs", SoftplusAttrs, AttrsNode);
};
@@ -569,7 +641,8 @@ struct PReluAttrs : public AttrsNode {
static void RegisterReflection() {
namespace refl = tvm::ffi::reflection;
refl::ObjectDef<PReluAttrs>().def_ro("axis", &PReluAttrs::axis,
- "The axis along which the alpha values are applied.");
+ "The axis along which the alpha values are applied.",
+ refl::DefaultValue(1));
}
TVM_FFI_DECLARE_OBJECT_INFO_FINAL("relax.attrs.PReluAttrs", PReluAttrs, AttrsNode);
};
@@ -588,15 +661,18 @@ struct BatchNormAttrs : public AttrsNode {
refl::ObjectDef<BatchNormAttrs>()
.def_ro("axis", &BatchNormAttrs::axis, "The axis along which the normalization is applied.")
.def_ro("epsilon", &BatchNormAttrs::epsilon,
- "Small float added to variance to avoid dividing by zero")
+ "Small float added to variance to avoid dividing by zero",
+ refl::DefaultValue(1e-05))
.def_ro("center", &BatchNormAttrs::center,
- "Indicating if the beta offset will be added to the normalized tensor.")
+ "Indicating if the beta offset will be added to the normalized tensor.",
+ refl::DefaultValue(true))
.def_ro("scale", &BatchNormAttrs::scale,
- "Indicating if the gamma scale will be multiplied.")
+ "Indicating if the gamma scale will be multiplied.", refl::DefaultValue(true))
.def_ro("momentum", &BatchNormAttrs::momentum,
- "The value used for the moving_mean and moving_var update.")
+ "The value used for the moving_mean and moving_var update.",
+ refl::DefaultValue(0.1))
.def_ro("training", &BatchNormAttrs::training,
- "Whether we are training (i.e., not in eval mode).");
+ "Whether we are training (i.e., not in eval mode).", refl::DefaultValue(true));
}
TVM_FFI_DECLARE_OBJECT_INFO_FINAL("relax.attrs.BatchNormAttrs", BatchNormAttrs, AttrsNode);
}; // struct BatchNormAttrs
@@ -614,11 +690,13 @@ struct LayerNormAttrs : public AttrsNode {
.def_ro("axes", &LayerNormAttrs::axes,
"The axes that along which the normalization is applied.")
.def_ro("epsilon", &LayerNormAttrs::epsilon,
- "Small float added to variance to avoid dividing by zero")
+ "Small float added to variance to avoid dividing by zero",
+ refl::DefaultValue(1e-05))
.def_ro("center", &LayerNormAttrs::center,
- "Indicating if the beta offset will be added to the normalized tensor.")
+ "Indicating if the beta offset will be added to the normalized tensor.",
+ refl::DefaultValue(true))
.def_ro("scale", &LayerNormAttrs::scale,
- "Indicating if the gamma scale will be multiplied.");
+ "Indicating if the gamma scale will be multiplied.", refl::DefaultValue(true));
}
TVM_FFI_DECLARE_OBJECT_INFO_FINAL("relax.attrs.LayerNormAttrs", LayerNormAttrs, AttrsNode);
}; // struct LayerNormAttrs
@@ -643,11 +721,13 @@ struct GroupNormAttrs : public AttrsNode {
"axes", &GroupNormAttrs::axes,
"The axes that along which the normalization is applied (excluding the channel axis).")
.def_ro("epsilon", &GroupNormAttrs::epsilon,
- "Small float added to variance to avoid dividing by zero")
+ "Small float added to variance to avoid dividing by zero",
+ refl::DefaultValue(1e-05))
.def_ro("center", &GroupNormAttrs::center,
- "Indicating if the beta offset will be added to the normalized tensor.")
+ "Indicating if the beta offset will be added to the normalized tensor.",
+ refl::DefaultValue(true))
.def_ro("scale", &GroupNormAttrs::scale,
- "Indicating if the gamma scale will be multiplied.");
+ "Indicating if the gamma scale will be multiplied.", refl::DefaultValue(true));
}
TVM_FFI_DECLARE_OBJECT_INFO_FINAL("relax.attrs.GroupNormAttrs", GroupNormAttrs, AttrsNode);
}; // struct GroupNormAttrs
@@ -668,11 +748,13 @@ struct InstanceNormAttrs : public AttrsNode {
.def_ro("axes", &InstanceNormAttrs::axes,
"The axes that along which the normalization is applied.")
.def_ro("epsilon", &InstanceNormAttrs::epsilon,
- "Small float added to variance to avoid dividing by zero")
+ "Small float added to variance to avoid dividing by zero",
+ refl::DefaultValue(1e-05))
.def_ro("center", &InstanceNormAttrs::center,
- "Indicating if the beta offset will be added to the normalized tensor.")
+ "Indicating if the beta offset will be added to the normalized tensor.",
+ refl::DefaultValue(true))
.def_ro("scale", &InstanceNormAttrs::scale,
- "Indicating if the gamma scale will be multiplied.");
+ "Indicating if the gamma scale will be multiplied.", refl::DefaultValue(true));
}
TVM_FFI_DECLARE_OBJECT_INFO_FINAL("relax.attrs.InstanceNormAttrs", InstanceNormAttrs, AttrsNode);
}; // struct InstanceNormAttrs
@@ -686,9 +768,11 @@ struct RMSNormAttrs : public AttrsNode {
namespace refl = tvm::ffi::reflection;
refl::ObjectDef<RMSNormAttrs>()
.def_ro("axes", &RMSNormAttrs::axes,
- "The axes that along which the normalization is applied.")
+ "The axes that along which the normalization is applied.",
+ refl::DefaultValue(ffi::Array<int64_t>{-1}))
.def_ro("epsilon", &RMSNormAttrs::epsilon,
- "Small float added to variance to avoid dividing by zero");
+ "Small float added to variance to avoid dividing by zero",
+ refl::DefaultValue(1e-05));
}
TVM_FFI_DECLARE_OBJECT_INFO_FINAL("relax.attrs.RMSNormAttrs", RMSNormAttrs, AttrsNode);
}; // struct RMSNormAttrs
@@ -705,7 +789,8 @@ struct NLLLossAttrs : public AttrsNode {
"The reduction method to apply to the output. Can be"
"'none', 'mean' or 'sum'.",
refl::DefaultValue("mean"))
- .def_ro("ignore_index", &NLLLossAttrs::ignore_index, "The target value to ignore.");
+ .def_ro("ignore_index", &NLLLossAttrs::ignore_index, "The target value to ignore.",
+ refl::DefaultValue(-100));
}
TVM_FFI_DECLARE_OBJECT_INFO_FINAL("relax.attrs.NLLLossAttrs", NLLLossAttrs, AttrsNode);
}; // struct NLLLossAttrs
@@ -718,7 +803,8 @@ struct DropoutAttrs : public AttrsNode {
namespace refl = tvm::ffi::reflection;
refl::ObjectDef<DropoutAttrs>().def_ro(
"rate", &DropoutAttrs::rate,
- "Fraction of the input that gets dropped out during training time");
+ "Fraction of the input that gets dropped out during training time",
+ refl::DefaultValue(0.5));
}
TVM_FFI_DECLARE_OBJECT_INFO_FINAL("relax.attrs.DropoutAttrs", DropoutAttrs, AttrsNode);
}; // struct DropoutAttrs
@@ -734,11 +820,14 @@ struct AttentionAttrs : public AttrsNode {
refl::ObjectDef<AttentionAttrs>()
.def_ro(
"scale", &AttentionAttrs::scale,
- "The custom scale applied before the softmax. The default value is 1 / sqrt(head_dim).")
+ "The custom scale applied before the softmax. The default value is 1 / sqrt(head_dim).",
+ refl::DefaultValue(ffi::Optional<FloatImm>{}))
.def_ro("causal_mask", &AttentionAttrs::causal_mask,
- "The type of the causal mask, i.e. 'TopLeft' and 'BottomRight'.")
+ "The type of the causal mask, i.e. 'TopLeft' and 'BottomRight'.",
+ refl::DefaultValue(ffi::Optional<ffi::String>{}))
.def_ro("window_size", &AttentionAttrs::window_size,
- "The size of the window for sliding-window attention.");
+ "The size of the window for sliding-window attention.",
+ refl::DefaultValue(ffi::Optional<IntImm>{}));
}
TVM_FFI_DECLARE_OBJECT_INFO_FINAL("relax.attrs.AttentionAttrs", AttentionAttrs, AttrsNode);
}; // struct AttentionAttrs
diff --git a/include/tvm/relax/attrs/search.h b/include/tvm/relax/attrs/search.h
index 6b3ee4860a..c750093368 100644
--- a/include/tvm/relax/attrs/search.h
+++ b/include/tvm/relax/attrs/search.h
@@ -38,11 +38,13 @@ struct ArgmaxArgminAttrs : public AttrsNode {
namespace refl = tvm::ffi::reflection;
refl::ObjectDef<ArgmaxArgminAttrs>()
.def_ro("axis", &ArgmaxArgminAttrs::axis,
- "The axis along which to perform the argmin/argmax.")
+ "The axis along which to perform the argmin/argmax.",
+ refl::DefaultValue(ffi::Optional<int64_t>{}))
.def_ro("keepdims", &ArgmaxArgminAttrs::keepdims,
"If this is set to `True`, the reduced axis is left in the result as dimension "
"with size "
- "one.");
+ "one.",
+ refl::DefaultValue(false));
}
TVM_FFI_DECLARE_OBJECT_INFO_FINAL("relax.attrs.ArgmaxArgminAttrs", ArgmaxArgminAttrs, AttrsNode);
}; // struct ArgmaxArgminAttrs
@@ -56,9 +58,10 @@ struct BucketizeAttrs : public tvm::AttrsNode {
namespace refl = tvm::ffi::reflection;
refl::ObjectDef<BucketizeAttrs>()
.def_ro("out_int32", &BucketizeAttrs::out_int32,
- "Indicate the output datatype, int32 if True, int64 otherwise.")
- .def_ro("right", &BucketizeAttrs::right,
- "Determines the behavior for values in boundaries");
+ "Indicate the output datatype, int32 if True, int64 otherwise.",
+ refl::DefaultValue(false))
+ .def_ro("right", &BucketizeAttrs::right, "Determines the behavior for values in boundaries",
+ refl::DefaultValue(false));
}
TVM_FFI_DECLARE_OBJECT_INFO_FINAL("relax.attrs.BucketizeAttrs", BucketizeAttrs, AttrsNode);
}; // struct BucketizeAttrs
diff --git a/include/tvm/relax/attrs/sorting.h b/include/tvm/relax/attrs/sorting.h
index eb0993721f..606a5e3c85 100644
--- a/include/tvm/relax/attrs/sorting.h
+++ b/include/tvm/relax/attrs/sorting.h
@@ -67,7 +67,8 @@ struct ArgsortAttrs : public AttrsNode {
"Whether to argsort in descending order."
"If it is not specified, it defaults to the ascending order.",
refl::DefaultValue(false))
- .def_ro("dtype", &ArgsortAttrs::dtype, "DType of the output indices.");
+ .def_ro("dtype", &ArgsortAttrs::dtype, "DType of the output indices.",
+ refl::DefaultValue((DLDataType{kDLInt, 32, 1})));
}
TVM_FFI_DECLARE_OBJECT_INFO_FINAL("relax.attrs.ArgsortAttrs", ArgsortAttrs, AttrsNode);
}; // struct ArgsortAttrs
@@ -83,7 +84,7 @@ struct TopKAttrs : public AttrsNode {
static void RegisterReflection() {
namespace refl = tvm::ffi::reflection;
refl::ObjectDef<TopKAttrs>()
- .def_ro("k", &TopKAttrs::k, "Number of top elements to select")
+ .def_ro("k", &TopKAttrs::k, "Number of top elements to select", refl::DefaultValue(1))
.def_ro("axis", &TopKAttrs::axis, "Axis along which to sort the input tensor.",
refl::DefaultValue(-1))
.def_ro("ret_type", &TopKAttrs::ret_type,
@@ -96,7 +97,8 @@ struct TopKAttrs : public AttrsNode {
"Whether to return largest or smallest elements."
"By default, return the largest k elements.",
refl::DefaultValue(true))
- .def_ro("dtype", &TopKAttrs::dtype, "Data type of the output indices.");
+ .def_ro("dtype", &TopKAttrs::dtype, "Data type of the output indices.",
+ refl::DefaultValue((DLDataType{kDLInt, 32, 1})));
}
TVM_FFI_DECLARE_OBJECT_INFO_FINAL("relax.attrs.TopKAttrs", TopKAttrs, AttrsNode);
}; // struct TopKAttrs
diff --git a/include/tvm/relax/attrs/statistical.h b/include/tvm/relax/attrs/statistical.h
index fa0ea8f47e..53e4481557 100644
--- a/include/tvm/relax/attrs/statistical.h
+++ b/include/tvm/relax/attrs/statistical.h
@@ -38,11 +38,13 @@ struct StatisticalAttrs : public AttrsNode {
namespace refl = tvm::ffi::reflection;
refl::ObjectDef<StatisticalAttrs>()
.def_ro("axis", &StatisticalAttrs::axis,
- "The axis or axes along which to perform the reduction.")
+ "The axis or axes along which to perform the reduction.",
+ refl::DefaultValue(ffi::Optional<ffi::Array<int64_t>>{}))
.def_ro("keepdims", &StatisticalAttrs::keepdims,
"If this is set to `True`, the reduced axes are left in the result as dimension "
"with size "
- "one.");
+ "one.",
+ refl::DefaultValue(false));
}
TVM_FFI_DECLARE_OBJECT_INFO_FINAL("relax.attrs.StatisticalAttrs", StatisticalAttrs, AttrsNode);
}; // struct StatisticalAttrs
@@ -58,10 +60,12 @@ struct ScanopAttrs : public AttrsNode {
refl::ObjectDef<ScanopAttrs>()
.def_ro("axis", &ScanopAttrs::axis,
"The axis along which to perform the scan computation."
- "The default (None) is to compute over the flattened array.")
+ "The default (None) is to compute over the flattened array.",
+ refl::DefaultValue(ffi::Optional<int64_t>{}))
.def_ro("dtype", &ScanopAttrs::dtype,
"The output data type."
- "If dtype is not specified, it defaults to the dtype of input data.")
+ "If dtype is not specified, it defaults to the dtype of input data.",
+ refl::DefaultValue(ffi::Optional<DLDataType>{}))
.def_ro("exclusive", &ScanopAttrs::exclusive, "The first element is not included",
refl::DefaultValue(false));
}
diff --git a/include/tvm/relax/attrs/vision.h b/include/tvm/relax/attrs/vision.h
index 9a2384a494..8091ad1c32 100644
--- a/include/tvm/relax/attrs/vision.h
+++ b/include/tvm/relax/attrs/vision.h
@@ -40,7 +40,8 @@ struct AllClassNonMaximumSuppressionAttrs : public AttrsNode {
refl::ObjectDef<AllClassNonMaximumSuppressionAttrs>().def_ro(
"output_format", &AllClassNonMaximumSuppressionAttrs::output_format,
"Output format, onnx or tensorflow. Returns outputs in a way that can be easily "
- "consumed by each frontend.");
+ "consumed by each frontend.",
+ refl::DefaultValue(ffi::String("onnx")));
}
TVM_FFI_DECLARE_OBJECT_INFO_FINAL("relax.attrs.AllClassNonMaximumSuppressionAttrs",
AllClassNonMaximumSuppressionAttrs, AttrsNode);
@@ -62,11 +63,15 @@ struct ROIAlignAttrs : public AttrsNode {
.def_ro("spatial_scale", &ROIAlignAttrs::spatial_scale,
"Ratio of input feature map height (or width) to raw image height (or width).")
.def_ro("sample_ratio", &ROIAlignAttrs::sample_ratio,
- "Optional sampling ratio of ROI align, using adaptive size by default.")
+ "Optional sampling ratio of ROI align, using adaptive size by default.",
+ refl::DefaultValue(-1))
.def_ro("aligned", &ROIAlignAttrs::aligned,
- "Whether to use the aligned ROIAlign semantics without the legacy 1-pixel clamp.")
- .def_ro("layout", &ROIAlignAttrs::layout, "Dimension ordering of the input data.")
- .def_ro("mode", &ROIAlignAttrs::mode, "Mode for ROI Align. Can be 'avg' or 'max'.");
+ "Whether to use the aligned ROIAlign semantics without the legacy 1-pixel clamp.",
+ refl::DefaultValue(false))
+ .def_ro("layout", &ROIAlignAttrs::layout, "Dimension ordering of the input data.",
+ refl::DefaultValue(ffi::String("NCHW")))
+ .def_ro("mode", &ROIAlignAttrs::mode, "Mode for ROI Align. Can be 'avg' or 'max'.",
+ refl::DefaultValue(ffi::String("avg")));
}
TVM_FFI_DECLARE_OBJECT_INFO_FINAL("relax.attrs.ROIAlignAttrs", ROIAlignAttrs, AttrsNode);
}; // struct ROIAlignAttrs
@@ -83,7 +88,8 @@ struct ROIPoolAttrs : public AttrsNode {
.def_ro("pooled_size", &ROIPoolAttrs::pooled_size, "Output size of roi pool.")
.def_ro("spatial_scale", &ROIPoolAttrs::spatial_scale,
"Ratio of input feature map height (or width) to raw image height (or width).")
- .def_ro("layout", &ROIPoolAttrs::layout, "Dimension ordering of the input data.");
+ .def_ro("layout", &ROIPoolAttrs::layout, "Dimension ordering of the input data.",
+ refl::DefaultValue(ffi::String("NCHW")));
}
TVM_FFI_DECLARE_OBJECT_INFO_FINAL("relax.attrs.ROIPoolAttrs", ROIPoolAttrs, AttrsNode);
}; // struct ROIPoolAttrs
@@ -98,11 +104,11 @@ struct GetValidCountsAttrs : public AttrsNode {
namespace refl = tvm::ffi::reflection;
refl::ObjectDef<GetValidCountsAttrs>()
.def_ro("score_threshold", &GetValidCountsAttrs::score_threshold,
- "Lower limit of score for valid bounding boxes.")
+ "Lower limit of score for valid bounding boxes.", refl::DefaultValue(0))
.def_ro("id_index", &GetValidCountsAttrs::id_index,
- "Index of the class categories, -1 to disable.")
+ "Index of the class categories, -1 to disable.", refl::DefaultValue(0))
.def_ro("score_index", &GetValidCountsAttrs::score_index,
- "Index of the scores/confidence of boxes.");
+ "Index of the scores/confidence of boxes.", refl::DefaultValue(1));
}
TVM_FFI_DECLARE_OBJECT_INFO_FINAL("relax.attrs.GetValidCountsAttrs", GetValidCountsAttrs,
AttrsNode);
@@ -126,27 +132,30 @@ struct NonMaximumSuppressionAttrs : public AttrsNode {
namespace refl = tvm::ffi::reflection;
refl::ObjectDef<NonMaximumSuppressionAttrs>()
.def_ro("max_output_size", &NonMaximumSuppressionAttrs::max_output_size,
- "Max number of output valid boxes, -1 for no limit.")
+ "Max number of output valid boxes, -1 for no limit.", refl::DefaultValue(-1))
.def_ro("iou_threshold", &NonMaximumSuppressionAttrs::iou_threshold,
- "Non-maximum suppression IoU threshold.")
+ "Non-maximum suppression IoU threshold.", refl::DefaultValue(0.5))
.def_ro("force_suppress", &NonMaximumSuppressionAttrs::force_suppress,
- "Whether to suppress all detections regardless of class_id.")
+ "Whether to suppress all detections regardless of class_id.",
+ refl::DefaultValue(false))
.def_ro("top_k", &NonMaximumSuppressionAttrs::top_k,
- "Keep maximum top k detections before nms, -1 for no limit.")
+ "Keep maximum top k detections before nms, -1 for no limit.",
+ refl::DefaultValue(-1))
.def_ro("coord_start", &NonMaximumSuppressionAttrs::coord_start,
- "Start index of the consecutive 4 coordinates.")
+ "Start index of the consecutive 4 coordinates.", refl::DefaultValue(2))
.def_ro("score_index", &NonMaximumSuppressionAttrs::score_index,
- "Index of the scores/confidence of boxes.")
+ "Index of the scores/confidence of boxes.", refl::DefaultValue(1))
.def_ro("id_index", &NonMaximumSuppressionAttrs::id_index,
- "Index of the class categories, -1 to disable.")
+ "Index of the class categories, -1 to disable.", refl::DefaultValue(0))
.def_ro("return_indices", &NonMaximumSuppressionAttrs::return_indices,
- "Whether to return box indices in input data.")
+ "Whether to return box indices in input data.", refl::DefaultValue(true))
.def_ro("invalid_to_bottom", &NonMaximumSuppressionAttrs::invalid_to_bottom,
- "Whether to move all valid bounding boxes to the top.")
+ "Whether to move all valid bounding boxes to the top.", refl::DefaultValue(false))
.def_ro("soft_nms_sigma", &NonMaximumSuppressionAttrs::soft_nms_sigma,
- "Sigma for soft-NMS; 0.0 means standard hard NMS.")
+ "Sigma for soft-NMS; 0.0 means standard hard NMS.", refl::DefaultValue(0.0))
.def_ro("score_threshold", &NonMaximumSuppressionAttrs::score_threshold,
- "Score threshold for soft-NMS validity check; 0.0 when unused.");
+ "Score threshold for soft-NMS validity check; 0.0 when unused.",
+ refl::DefaultValue(0.0));
}
TVM_FFI_DECLARE_OBJECT_INFO_FINAL("relax.attrs.NonMaximumSuppressionAttrs",
NonMaximumSuppressionAttrs, AttrsNode);
@@ -164,16 +173,19 @@ struct MultiboxTransformLocAttrs : public AttrsNode {
namespace refl = tvm::ffi::reflection;
refl::ObjectDef<MultiboxTransformLocAttrs>()
.def_ro("clip", &MultiboxTransformLocAttrs::clip,
- "Clip decoded ymin,xmin,ymax,xmax to [0,1].")
+ "Clip decoded ymin,xmin,ymax,xmax to [0,1].", refl::DefaultValue(false))
.def_ro("threshold", &MultiboxTransformLocAttrs::threshold,
- "After softmax, zero scores strictly below this value.")
+ "After softmax, zero scores strictly below this value.", refl::DefaultValue(0.0))
.def_ro("variances", &MultiboxTransformLocAttrs::variances,
"(x,y,w,h) scales = TFLite 1/x_scale,1/y_scale,1/w_scale,1/h_scale on "
- "encodings. Very large w/h scales can overflow exp in decode.")
+ "encodings. Very large w/h scales can overflow exp in decode.",
+ refl::DefaultValue(ffi::Array<double>{1.0, 1.0, 1.0, 1.0}))
.def_ro("keep_background", &MultiboxTransformLocAttrs::keep_background,
- "If false, force output scores[:,0,:] to 0 (background class).")
+ "If false, force output scores[:,0,:] to 0 (background class).",
+ refl::DefaultValue(true))
.def_ro("apply_softmax", &MultiboxTransformLocAttrs::apply_softmax,
- "Whether to apply softmax to class predictions before thresholding.");
+ "Whether to apply softmax to class predictions before thresholding.",
+ refl::DefaultValue(true));
}
TVM_FFI_DECLARE_OBJECT_INFO_FINAL("relax.attrs.MultiboxTransformLocAttrs",
MultiboxTransformLocAttrs, AttrsNode);
diff --git a/include/tvm/tirx/op_attr_types.h b/include/tvm/tirx/op_attr_types.h
index aab39d4ca1..b101f861b0 100644
--- a/include/tvm/tirx/op_attr_types.h
+++ b/include/tvm/tirx/op_attr_types.h
@@ -70,41 +70,6 @@ using FLowerIntrinsic = ffi::TypedFunction<PrimExpr(PrimExpr)>;
*/
using FLegalize = ffi::TypedFunction<PrimExpr(PrimExpr)>;
-/*!
- * \brief The fully qualified TVMScript name, including its dialect namespace.
- */
-using TScriptPrinterName = ffi::String;
-
-/*!
- * \brief The published script callable accepts the shared Op construction contract.
- *
- * Published by Python construction namespaces, not by Op definitions: a printer
- * name alone does not establish a legacy wrapper's positional/keyword contract.
- * Semantic syntax hooks retain precedence over ordinary named construction.
- */
-using TScriptStandardCall = bool;
-
-/*!
- * \brief Specifies that TVMScript printer prints the dtype as the first/last argument.
- If not specified, dtype will not be printed.
- */
-enum class ScriptDtypePrintLocation : int {
- /*!
- * \brief Do not print dtype as an argument.
- */
- kNone = 0,
- /*!
- * \brief Print dtype as the first argument.
- */
- kFirst = 1,
- /*!
- * \brief FPrint dtype as the last argument.
- */
- kLast = 2,
-};
-
-using TScriptDtypePrintLocation = int64_t;
-
/*!
* \brief Broad TIRx op category.
*
diff --git a/python/tvm/backend/cuda/__init__.py b/python/tvm/backend/cuda/__init__.py
index df0b40dc79..1929619d60 100644
--- a/python/tvm/backend/cuda/__init__.py
+++ b/python/tvm/backend/cuda/__init__.py
@@ -73,7 +73,7 @@ def register_backend():
pass
register_device_target_detector("cuda", _detect_target_from_device)
for name, namespace in script_namespaces().items():
- builder_op.register_script_namespace(name, namespace, canonical_op_names=name == "cuda")
+ builder_op.register_script_namespace(name, namespace)
# script_namespaces() above pulls in ptx, which only imports the shared
# codegen layer -- not the device-helper modules. This import is the sole
diff --git a/python/tvm/backend/cuda/lang/pipeline.py b/python/tvm/backend/cuda/lang/pipeline.py
index 645019b417..5d21e30484 100644
--- a/python/tvm/backend/cuda/lang/pipeline.py
+++ b/python/tvm/backend/cuda/lang/pipeline.py
@@ -116,7 +116,7 @@ def _map_buffer_into_cta(ptr, rank, depth):
mapped = T.alloc_local([1], "uint64")
T.evaluate(T.ptx.mapa.u64(mapped[0], ptr, T.uint32(rank)))
remote_ptr = TIRVar("remote_mbar_ptr", ptr_ty)
- T.bind(T.reinterpret(ptr_ty, mapped[0]), var=remote_ptr)
+ T.bind(T.reinterpret(mapped[0], ty=ptr_ty), var=remote_ptr)
return T.decl_tensor([depth], "uint64", data=remote_ptr, scope="shared")
diff --git a/python/tvm/backend/cuda/op.py b/python/tvm/backend/cuda/op.py
index 4218095fc8..c437d04e07 100644
--- a/python/tvm/backend/cuda/op.py
+++ b/python/tvm/backend/cuda/op.py
@@ -19,6 +19,8 @@
from __future__ import annotations
+from enum import Enum
+
from tvm import tirx
from tvm.ir import Call, Op, StringImm
from tvm.ir.op import _init_op_api, _make_op_api
@@ -99,7 +101,7 @@ def cuda_func_call(func_name, *args, source_code, return_type="void"):
return call_intrin(return_type, "tirx.cuda.func_call", func_name, *args, source_code)
-def cuda_warp_reduce(value, op, width=32):
+def cuda_warp_reduce(value, op, width=32, *, ty=None, span=None):
"""Warp-level butterfly shuffle-XOR reduction.
Reduces ``value`` across ``width`` adjacent lanes using the specified
@@ -123,7 +125,12 @@ def cuda_warp_reduce(value, op, width=32):
call : Expr
The reduced value (same dtype as *value*).
"""
- return call_intrin(value.ty, "tirx.cuda.warp_reduce", value, op, width)
+ return Call(
+ "tirx.cuda.warp_reduce",
+ [value, op, width],
+ ty=ty,
+ span=span,
+ )
def cuda_warp_sum(value, width=32):
@@ -141,7 +148,7 @@ def cuda_warp_min(value, width=32):
return cuda_warp_reduce(value, "min", width)
-def cuda_cta_reduce(value, op, num_warps, scratch):
+def cuda_cta_reduce(value, op, num_warps, scratch, *, ty=None, span=None):
"""CTA-wide reduction via warp shuffle + shared memory.
Two-step reduction: (1) intra-warp shuffle reduction, (2) warp-0
@@ -167,7 +174,12 @@ def cuda_cta_reduce(value, op, num_warps, scratch):
call : Expr
The reduced value broadcast to all threads (same dtype as *value*).
"""
- return call_intrin(value.ty, "tirx.cuda.cta_reduce", value, op, num_warps, scratch)
+ return Call(
+ "tirx.cuda.cta_reduce",
+ [value, op, num_warps, scratch],
+ ty=ty,
+ span=span,
+ )
def cuda_cta_sum(value, num_warps, scratch):
@@ -254,6 +266,9 @@ def cuda_wait_until(
space="global",
ptx_type=None,
backoff_ns=None,
+ *,
+ ty=None,
+ span=None,
):
"""Read the global word at ``ptr`` into ``dst`` until ``predicate`` holds,
and leave the exit value there.
@@ -306,21 +321,27 @@ def cuda_wait_until(
checker concludes. A timeout that has to print and trap is not that, and
belongs to a loop the kernel writes itself.
"""
+ scope = _static_str(scope)
+ space = _static_str(space)
+ ptx_type = _static_str(ptx_type) if ptx_type is not None else None
_validate_wait_until_attrs(scope, space, ptx_type)
_reject_wide_word_for_predicate(ptx_type, "wait_until")
if tirx.is_tensor_var(dst):
dst = dst[0]
condition = tirx.convert(predicate(dst) if callable(predicate) else predicate)
- return call_intrin(
- "",
+ return Call(
"tirx.cuda.wait_until",
- dst,
- ptr,
- condition,
- scope,
- space,
- ptx_type or "",
- tirx.convert(0 if backoff_ns is None else backoff_ns),
+ [
+ dst,
+ ptr,
+ condition,
+ scope,
+ space,
+ ptx_type or "",
+ tirx.convert(0 if backoff_ns is None else backoff_ns),
+ ],
+ ty=ty,
+ span=span,
)
@@ -337,35 +358,27 @@ def _validate_mbarrier_arrive_attrs(sem, scope, space, remote):
raise ValueError("remote mbarrier.arrive requires space='shared::cluster'")
-def ptx_cp_async_legacy(*all_args):
- """Legacy ``ptx_cp_async`` API taking explicit src/dst offsets.
+_cp_async_raw = _make_op_api(Op.get("tirx.s_tir.cp_async_raw"), __name__)
- Signature: ``(dst_ptr, dst_offset, src_ptr, src_offset, cp_size)``.
- Offsets are folded into the pointers via ``tvm_access_ptr`` and the call
- lowers through the raw ``tirx.s_tir.cp_async_raw`` op.
- ``T.s_tir.cp_async_raw.legacy`` runs through ``_dtype_forward`` which
- prepends a ``dtype=`` kwarg as a leading positional. The dtype names
- the *element* type of the buffer (offsets are in elements of that
- dtype, not bytes), so this function accepts either 5 or 6 positional
- args.
+def ptx_cp_async_legacy(
+ dst_ptr,
+ dst_offset,
+ src_ptr,
+ src_offset,
+ cp_size,
+ *,
+ elem_dtype="int8",
+ ty=None,
+ span=None,
+):
+ """Fold element offsets into the raw cp.async pointer operands.
+
+ ``elem_dtype`` scales element offsets independently of the call's result ``ty``.
"""
- args = list(all_args)
- elem_dtype = "int8"
- if len(args) == 6:
- # Leading positional is the buffer element dtype, used to scale
- # offsets correctly when folding via ``tvm_access_ptr``.
- elem_dtype = args.pop(0)
- if len(args) != 5:
- raise ValueError(
- f"ptx_cp_async_legacy expects 5 args (or 6 with dtype= kwarg "
- f"prepended); got {len(all_args)}"
- )
- dst_ptr, dst_offset, src_ptr, src_offset, cp_size = args
dst_ptr = _wrap_or_fold_access_ptr(dst_ptr, dst_offset, elem_dtype)
src_ptr = _wrap_or_fold_access_ptr(src_ptr, src_offset, elem_dtype)
- # The raw 5-arg Call InjectPTXAsyncCopy emits; offsets are already folded.
- return call_intrin(elem_dtype, "tirx.s_tir.cp_async_raw", dst_ptr, 0, src_ptr, 0, cp_size)
+ return _cp_async_raw(dst_ptr, 0, src_ptr, 0, cp_size, ty=ty, span=span)
def _is_static_unicast_cta_mask(cta_mask):
@@ -377,7 +390,7 @@ def _is_static_unicast_cta_mask(cta_mask):
return False
-def cuda_mov_sreg(bits, reg_name):
+def cuda_mov_sreg(bits, reg_name, *, ty=None, span=None):
"""TVM intrinsic to tvm instrinsics to fetch PTX pre-defined registers
Parameters
@@ -393,63 +406,38 @@ def cuda_mov_sreg(bits, reg_name):
call : Expr
The call expression.
"""
- return call_intrin("int" + str(bits), "tirx.cuda.mov_sreg", bits, reg_name)
-
+ return Call(
+ "tirx.cuda.mov_sreg",
+ [bits, reg_name],
+ ty=ty,
+ span=span,
+ )
-def ptx_legacy_mma(*all_args, operator=None):
- """Legacy ``ptx_mma`` API.
- Signature: ``(shape, A_layout, B_layout, A_dtype, B_dtype, C_dtype,
- multiplicand_a, a_index, multiplicand_b, b_index, accumulator,
- c_index, saturate, operator=None)``. The accumulator is reused as
- both input and output (no separate ``d``/``c`` slot). Translation:
+_legacy_mma = _make_op_api(Op.get("tirx.ptx_legacy.mma"), __name__)
- * ``a_dtype, b_dtype, c_dtype`` → fork ``a_type, b_type, c_type``
- (and reuse ``c_dtype`` as fork ``d_type`` since the accumulator
- dtype is the output dtype here).
- * ``(a_ptr, a_offset)`` and ``(b_ptr, b_offset)`` → folded via
- :func:`tvm_access_ptr`.
- * ``(accumulator, c_index)`` → folded; passed for both ``d_ptr`` and
- ``c_ptr`` since the accumulator is reused as the output.
- ``T.ptx_legacy.mma`` runs through ``_dtype_forward`` which prepends a
- ``dtype=`` kwarg as a leading positional, so this function accepts
- either 13 or 14 positional args.
- """
- args = list(all_args)
- # ``T.ptx_legacy.mma(..., dtype="...")`` has the dtype prepended by
- # ``_dtype_forward``; strip it here.
- if len(args) in (14, 15):
- _ = args.pop(0)
- if len(args) == 14:
- # operator passed positionally as the trailing arg.
- operator = args.pop()
- if len(args) != 13:
- raise ValueError(
- f"ptx_legacy_mma expects 13-15 positional args (with optional "
- f"leading ``call_dtype`` from dtype= kwarg and optional trailing "
- f"``operator``); got {len(all_args)}"
- )
- (
- shape,
- a_layout,
- b_layout,
- a_dtype,
- b_dtype,
- c_dtype,
- a_ptr,
- a_offset,
- b_ptr,
- b_offset,
- acc_ptr,
- c_offset,
- saturate,
- ) = args
- # Emit tirx.ptx_legacy.mma directly with separate (ptr_var, offset)
- # pairs. codegen_cuda.cc uses C pointer arithmetic ``ptr + offset``
- # so element offsets stay element-accurate, and lower_warp_memory
- # rewrites the offset's group component to a thread-local index.
- call_args = [
+def ptx_legacy_mma(
+ shape,
+ a_layout,
+ b_layout,
+ a_dtype,
+ b_dtype,
+ c_dtype,
+ a_ptr,
+ a_offset,
+ b_ptr,
+ b_offset,
+ acc_ptr,
+ c_offset,
+ saturate,
+ operator=None,
+ *,
+ ty=None,
+ span=None,
+):
+ """Construct legacy MMA with element offsets and an optional bit operator."""
+ args = [
shape,
a_layout,
b_layout,
@@ -465,8 +453,8 @@ def ptx_legacy_mma(*all_args, operator=None):
saturate,
]
if operator is not None:
- call_args.append(operator)
- return call_intrin("", "tirx.ptx_legacy.mma", *call_args)
+ args.append(operator)
+ return _legacy_mma(*args, ty=ty, span=span)
def mma_store(dtype, m, n, dst_ptr, src_ptr, src_offset, dst_stride):
@@ -475,18 +463,20 @@ def mma_store(dtype, m, n, dst_ptr, src_ptr, src_offset, dst_stride):
return call_intrin(dtype, "tirx.mma_store", m, n, dst_ptr, src_ptr, src_offset, dst_stride)
-def mma_store_legacy(dtype, m, n, dst_ptr, src_ptr, src_offset, dst_stride):
- """mma_store with apache-style pointer/offset semantics."""
+_mma_store_legacy = _make_op_api(Op.get("tirx.mma_store_legacy"), __name__)
- return call_intrin(
- dtype,
- "tirx.mma_store_legacy",
+
+def mma_store_legacy(m, n, dst_ptr, src_ptr, src_offset, dst_stride, *, ty=None, span=None):
+ """Store MMA registers using explicit pointer and element-offset operands."""
+ return _mma_store_legacy(
m,
n,
dst_ptr,
src_ptr,
src_offset,
dst_stride,
+ ty=ty,
+ span=span,
)
@@ -496,10 +486,12 @@ def mma_fill(dtype, local_size, local_ptr, offset):
return call_intrin(dtype, "tirx.mma_fill", local_size, local_ptr, offset)
-def mma_fill_legacy(dtype, local_size, local_ptr, offset):
- """mma_fill with apache-style pointer/offset semantics."""
+_mma_fill_legacy = _make_op_api(Op.get("tirx.mma_fill_legacy"), __name__)
+
- return call_intrin(dtype, "tirx.mma_fill_legacy", local_size, local_ptr, offset)
+def mma_fill_legacy(local_size, local_ptr, offset, *, ty=None, span=None):
+ """Initialize MMA registers using an explicit element offset."""
+ return _mma_fill_legacy(local_size, local_ptr, offset, ty=ty, span=span)
_PTX_TO_NUMPY_DTYPE = {
@@ -566,33 +558,27 @@ def _wrap_or_fold_access_ptr(ptr, offset, elem_dtype):
return tvm_access_ptr(elem_dtype, ptr, offset, 1, 1)
-def ptx_legacy_ldmatrix(*all_args):
- """Legacy ``ptx_ldmatrix`` API taking explicit offsets.
+_legacy_ldmatrix = _make_op_api(Op.get("tirx.ptx_legacy.ldmatrix"), __name__)
- Signature: ``(trans, num, dtype, local_ptr, local_offset, smem_ptr,
- smem_offset)``. Offsets are folded into the pointers via
- ``tvm_access_ptr``.
- ``T.ptx_legacy.ldmatrix`` runs through ``_dtype_forward`` which
- prepends a ``dtype=`` kwarg as a leading positional naming the buffer
- element type — offsets are in elements of that dtype, not bytes, so
- we forward it to ``tvm_access_ptr`` for correct scaling.
+def ptx_legacy_ldmatrix(
+ trans,
+ num,
+ dtype,
+ local_ptr,
+ local_offset,
+ smem_ptr,
+ smem_offset,
+ *,
+ ty=None,
+ span=None,
+):
+ """Load a matrix with explicit pointer and element-offset operands.
+
+ The legacy lowering uses the result ``ty`` for its element width, including
+ the transposed int8 gather form. ``dtype`` is the PTX instruction type token.
"""
- if len(all_args) == 8:
- elem_dtype, trans, num, dtype, local_ptr, local_offset, smem_ptr, smem_offset = all_args
- elif len(all_args) == 7:
- trans, num, dtype, local_ptr, local_offset, smem_ptr, smem_offset = all_args
- elem_dtype = "int8"
- else:
- raise ValueError(
- f"ptx_legacy_ldmatrix expects 7 args (or 8 with dtype= kwarg "
- f"prepended); got {len(all_args)}"
- )
- # Call.dtype carries the buffer element type so codegen can pick the
- # int8+trans manual-loop fallback (ldmatrix can't transpose int8).
- return call_intrin(
- elem_dtype,
- "tirx.ptx_legacy.ldmatrix",
+ return _legacy_ldmatrix(
trans,
num,
dtype,
@@ -600,6 +586,8 @@ def ptx_legacy_ldmatrix(*all_args):
local_offset,
smem_ptr,
smem_offset,
+ ty=ty,
+ span=span,
)
@@ -818,7 +806,16 @@ def _static_str(value):
# shape tokens in backend/cuda/ptx/table.py.
-def timer_init_cuda(profiler_buffer, profiler_tag, profiler_write_offset, num_groups, group_id):
+def timer_init_cuda(
+ profiler_buffer,
+ profiler_tag,
+ profiler_write_offset,
+ num_groups,
+ group_id,
+ *,
+ ty=None,
+ span=None,
+):
"""TVM intrinsic for initializing the CUDA profiler, and store profiling result in a buffer.
Parameters
@@ -845,14 +842,11 @@ def timer_init_cuda(profiler_buffer, profiler_tag, profiler_write_offset, num_gr
The call expression.
"""
- return call_intrin(
- "void",
+ return Call(
"tirx.timer_init_cuda",
- profiler_buffer,
- profiler_tag,
- profiler_write_offset,
- num_groups,
- group_id,
+ [profiler_buffer, profiler_tag, profiler_write_offset, num_groups, group_id],
+ ty=ty,
+ span=span,
)
@@ -863,6 +857,9 @@ def timer_start_cuda(
profiler_write_offset,
profiler_write_stride,
leader_cond,
+ *,
+ ty=None,
+ span=None,
):
"""TVM intrinsic for starting the timer for profiling a specific event, and storing profiling result in a buffer.
@@ -893,15 +890,18 @@ def timer_start_cuda(
The call expression.
""" # noqa: E501
- return call_intrin(
- "void",
+ return Call(
"tirx.timer_start_cuda",
- event_type.value,
- profiler_buffer,
- profiler_tag,
- profiler_write_offset,
- profiler_write_stride,
- leader_cond,
+ [
+ event_type.value if isinstance(event_type, Enum) else event_type,
+ profiler_buffer,
+ profiler_tag,
+ profiler_write_offset,
+ profiler_write_stride,
+ leader_cond,
+ ],
+ ty=ty,
+ span=span,
)
@@ -912,6 +912,9 @@ def timer_end_cuda(
profiler_write_offset,
profiler_write_stride,
leader_cond,
+ *,
+ ty=None,
+ span=None,
):
"""TVM intrinsic for ending the timer for profiling a specific event, and storing profiling result in a buffer.
@@ -942,20 +945,30 @@ def timer_end_cuda(
The call expression.
""" # noqa: E501
- return call_intrin(
- "void",
+ return Call(
"tirx.timer_end_cuda",
- event_type.value,
- profiler_buffer,
- profiler_tag,
- profiler_write_offset,
- profiler_write_stride,
- leader_cond,
+ [
+ event_type.value if isinstance(event_type, Enum) else event_type,
+ profiler_buffer,
+ profiler_tag,
+ profiler_write_offset,
+ profiler_write_stride,
+ leader_cond,
+ ],
+ ty=ty,
+ span=span,
)
def timer_finalize_cuda(
- profiler_buffer, profiler_tag, profiler_write_offset, profiler_write_stride, leader_cond
+ profiler_buffer,
+ profiler_tag,
+ profiler_write_offset,
+ profiler_write_stride,
+ leader_cond,
+ *,
+ ty=None,
+ span=None,
):
"""TVM intrinsic for finalizing the CUDA profiler, and store profiling result in a buffer.
@@ -983,18 +996,15 @@ def timer_finalize_cuda(
The call expression.
"""
- return call_intrin(
- "void",
+ return Call(
"tirx.timer_finalize_cuda",
- profiler_buffer,
- profiler_tag,
- profiler_write_offset,
- profiler_write_stride,
- leader_cond,
+ [profiler_buffer, profiler_tag, profiler_write_offset, profiler_write_stride, leader_cond],
+ ty=ty,
+ span=span,
)
-def cuda_atomic_add(res_addr, value):
+def cuda_atomic_add(res_addr, value, *, ty=None, span=None):
"""TVM intrinsic to call cuda atomic add instruction
Parameters
@@ -1011,10 +1021,15 @@ def cuda_atomic_add(res_addr, value):
The call expression.
"""
value = tir.convert(value)
- return call_intrin(value.ty, "tirx.cuda.atomic_add", res_addr, value)
+ return Call(
+ "tirx.cuda.atomic_add",
+ [res_addr, value],
+ ty=ty,
+ span=span,
+ )
-def cuda_ldg(addr, dtype, *, dst=None, vec=""):
+def cuda_ldg(addr, dtype, *, dst=None, vec="", ty=None, span=None):
"""TVM intrinsic to call CUDA C++ ``__ldg()``.
Parameters
@@ -1037,7 +1052,7 @@ def cuda_ldg(addr, dtype, *, dst=None, vec=""):
if dst is None:
if vec:
raise ValueError("vector cuda.ldg requires dst")
- return call_intrin(dtype, "tirx.cuda.ldg", addr, dtype)
+ return Call("tirx.cuda.ldg", [addr, dtype], ty=ty, span=span)
if vec not in ("v2", "v4"):
raise ValueError(f"vector cuda.ldg expects vec in {{'v2', 'v4'}}, got {vec!r}")
if not isinstance(dst, list | tuple):
@@ -1045,7 +1060,12 @@ def cuda_ldg(addr, dtype, *, dst=None, vec=""):
vec_len = int(vec[1:])
if len(dst) != vec_len:
raise ValueError(f"cuda.ldg dst length must match {vec}: got {len(dst)}")
- return call_intrin("", "tirx.cuda.ldg", *dst, addr, dtype, vec, vec_len)
+ return Call(
+ "tirx.cuda.ldg",
+ [*dst, addr, dtype, vec, vec_len],
+ ty=ty,
+ span=span,
+ )
def cuda_sm100_2sm_leader_smem_addr(ptr):
@@ -1174,7 +1194,7 @@ def _validate_ptx_address(addr, space, op_name):
)
-def cuda_atomic_cas(ptr, old_val, new_val):
+def cuda_atomic_cas(ptr, old_val, new_val, *, ty=None, span=None):
"""TVM intrinsic to call cuda atomic cas instruction
Parameters
@@ -1194,7 +1214,12 @@ def cuda_atomic_cas(ptr, old_val, new_val):
The call expression.
"""
old_val = tir.convert(old_val)
- return call_intrin(old_val.ty, "tirx.cuda.atomic_cas", ptr, old_val, new_val)
+ return Call(
+ "tirx.cuda.atomic_cas",
+ [ptr, old_val, new_val],
+ ty=ty,
+ span=span,
+ )
########################################################
@@ -1202,7 +1227,7 @@ def cuda_atomic_cas(ptr, old_val, new_val):
########################################################
-def nvshmem_my_pe():
+def nvshmem_my_pe(*, ty=None, span=None):
"""TVM intrinsic to call nvshmem_my_pe()
Returns
@@ -1211,10 +1236,10 @@ def nvshmem_my_pe():
The call expression.
"""
- return call_intrin("int32", "tirx.nvshmem.my_pe")
+ return Call("tirx.nvshmem.my_pe", [], ty=ty, span=span)
-def nvshmem_n_pes():
+def nvshmem_n_pes(*, ty=None, span=None):
"""TVM intrinsic to call nvshmem_n_pes()
Returns
@@ -1223,10 +1248,10 @@ def nvshmem_n_pes():
The call expression.
"""
- return call_intrin("int32", "tirx.nvshmem.n_pes")
+ return Call("tirx.nvshmem.n_pes", [], ty=ty, span=span)
-def nvshmem_getmem_nbi(dst, src, nelems, pe):
+def nvshmem_getmem_nbi(dst, src, nelems, pe, *, ty=None, span=None):
"""TVM intrinsic to call nvshmem_getmem_nbi()
Parameters
@@ -1249,10 +1274,15 @@ def nvshmem_getmem_nbi(dst, src, nelems, pe):
The call expression.
""" # noqa: E501
- return call_intrin("", "tirx.nvshmem.getmem_nbi", dst, src, nelems, pe)
+ return Call(
+ "tirx.nvshmem.getmem_nbi",
+ [dst, src, nelems, pe],
+ ty=ty,
+ span=span,
+ )
-def nvshmem_putmem_nbi(dst, src, nelems, pe):
+def nvshmem_putmem_nbi(dst, src, nelems, pe, *, ty=None, span=None):
"""TVM intrinsic to call nvshmem_putmem_nbi()
Parameters
@@ -1275,10 +1305,15 @@ def nvshmem_putmem_nbi(dst, src, nelems, pe):
The call expression.
"""
- return call_intrin("", "tirx.nvshmem.putmem_nbi", dst, src, nelems, pe)
+ return Call(
+ "tirx.nvshmem.putmem_nbi",
+ [dst, src, nelems, pe],
+ ty=ty,
+ span=span,
+ )
-def nvshmem_getmem_nbi_warp(dst, src, nelems, pe):
+def nvshmem_getmem_nbi_warp(dst, src, nelems, pe, *, ty=None, span=None):
"""TVM intrinsic to call nvshmem_getmem_nbi_warp()
Parameters
@@ -1301,10 +1336,15 @@ def nvshmem_getmem_nbi_warp(dst, src, nelems, pe):
The call expression.
""" # noqa: E501
- return call_intrin("", "tirx.nvshmem.getmem_nbi_warp", dst, src, nelems, pe)
+ return Call(
+ "tirx.nvshmem.getmem_nbi_warp",
+ [dst, src, nelems, pe],
+ ty=ty,
+ span=span,
+ )
-def nvshmem_putmem_nbi_warp(dst, src, nelems, pe):
+def nvshmem_putmem_nbi_warp(dst, src, nelems, pe, *, ty=None, span=None):
"""TVM intrinsic to call nvshmem_putmem_nbi_warp()
Parameters
@@ -1327,10 +1367,15 @@ def nvshmem_putmem_nbi_warp(dst, src, nelems, pe):
The call expression.
"""
- return call_intrin("", "tirx.nvshmem.putmem_nbi_warp", dst, src, nelems, pe)
+ return Call(
+ "tirx.nvshmem.putmem_nbi_warp",
+ [dst, src, nelems, pe],
+ ty=ty,
+ span=span,
+ )
-def nvshmem_getmem_nbi_block(dst, src, nelems, pe):
+def nvshmem_getmem_nbi_block(dst, src, nelems, pe, *, ty=None, span=None):
"""TVM intrinsic to call nvshmem_getmem_nbi_block()
Parameters
@@ -1353,10 +1398,15 @@ def nvshmem_getmem_nbi_block(dst, src, nelems, pe):
The call expression.
""" # noqa: E501
- return call_intrin("", "tirx.nvshmem.getmem_nbi_block", dst, src, nelems, pe)
+ return Call(
+ "tirx.nvshmem.getmem_nbi_block",
+ [dst, src, nelems, pe],
+ ty=ty,
+ span=span,
+ )
-def nvshmem_putmem_nbi_block(dst, src, nelems, pe):
+def nvshmem_putmem_nbi_block(dst, src, nelems, pe, *, ty=None, span=None):
"""TVM intrinsic to call nvshmem_putmem_nbi_block()
Parameters
@@ -1379,10 +1429,15 @@ def nvshmem_putmem_nbi_block(dst, src, nelems, pe):
The call expression.
"""
- return call_intrin("", "tirx.nvshmem.putmem_nbi_block", dst, src, nelems, pe)
+ return Call(
+ "tirx.nvshmem.putmem_nbi_block",
+ [dst, src, nelems, pe],
+ ty=ty,
+ span=span,
+ )
-def nvshmem_signal_op(sig_addr, signal, sig_op, pe):
+def nvshmem_signal_op(sig_addr, signal, sig_op, pe, *, ty=None, span=None):
"""TVM intrinsic to call nvshmem_signal_op()
Parameters
@@ -1406,10 +1461,15 @@ def nvshmem_signal_op(sig_addr, signal, sig_op, pe):
"""
_choice("sig_op", sig_op, _NVSHMEM_SIG_OP)
- return call_intrin("", "tirx.nvshmem.signal_op", sig_addr, signal, sig_op, pe)
+ return Call(
+ "tirx.nvshmem.signal_op",
+ [sig_addr, signal, sig_op, pe],
+ ty=ty,
+ span=span,
+ )
-def nvshmem_wait_until(ivar, cmp, cmp_value, type="uint64_t"):
+def nvshmem_wait_until(ivar, cmp, cmp_value, type="uint64_t", *, ty=None, span=None):
"""TVM intrinsic to call nvshmem_wait_until()
Parameters
@@ -1433,10 +1493,15 @@ def nvshmem_wait_until(ivar, cmp, cmp_value, type="uint64_t"):
"""
_choice("cmp", cmp, _NVSHMEM_CMP)
- return call_intrin("", "tirx.nvshmem.wait_until", ivar, cmp, cmp_value, type)
+ return Call(
+ "tirx.nvshmem.wait_until",
+ [ivar, cmp, cmp_value, type],
+ ty=ty,
+ span=span,
+ )
-def nvshmem_quiet():
+def nvshmem_quiet(*, ty=None, span=None):
"""TVM intrinsic to call nvshmem_quiet()
Returns
@@ -1445,10 +1510,12 @@ def nvshmem_quiet():
The call expression.
"""
- return call_intrin("", "tirx.nvshmem.quiet")
+ return Call("tirx.nvshmem.quiet", [], ty=ty, span=span)
-def nvshmem_putmem_signal_nbi(dst, src, nelems, sig_addr, signal, sig_op, pe):
+def nvshmem_putmem_signal_nbi(
+ dst, src, nelems, sig_addr, signal, sig_op, pe, *, ty=None, span=None
+):
"""TVM intrinsic to call nvshmem_putmem_signal_nbi()
Parameters
@@ -1480,12 +1547,17 @@ def nvshmem_putmem_signal_nbi(dst, src, nelems, sig_addr, signal, sig_op, pe):
The call expression.
""" # noqa: E501
- return call_intrin(
- "", "tirx.nvshmem.putmem_signal_nbi", dst, src, nelems, sig_addr, signal, sig_op, pe
+ return Call(
+ "tirx.nvshmem.putmem_signal_nbi",
+ [dst, src, nelems, sig_addr, signal, sig_op, pe],
+ ty=ty,
+ span=span,
)
-def nvshmem_putmem_signal_nbi_warp(dst, src, nelems, sig_addr, signal, sig_op, pe):
+def nvshmem_putmem_signal_nbi_warp(
+ dst, src, nelems, sig_addr, signal, sig_op, pe, *, ty=None, span=None
+):
"""TVM intrinsic to call nvshmem_putmem_signal_nbi_warp()
Parameters
@@ -1517,12 +1589,17 @@ def nvshmem_putmem_signal_nbi_warp(dst, src, nelems, sig_addr, signal, sig_op, p
The call expression.
""" # noqa: E501
- return call_intrin(
- "", "tirx.nvshmem.putmem_signal_nbi_warp", dst, src, nelems, sig_addr, signal, sig_op, pe
+ return Call(
+ "tirx.nvshmem.putmem_signal_nbi_warp",
+ [dst, src, nelems, sig_addr, signal, sig_op, pe],
+ ty=ty,
+ span=span,
)
-def nvshmem_putmem_signal_nbi_block(dst, src, nelems, sig_addr, signal, sig_op, pe):
+def nvshmem_putmem_signal_nbi_block(
+ dst, src, nelems, sig_addr, signal, sig_op, pe, *, ty=None, span=None
+):
"""TVM intrinsic to call nvshmem_putmem_signal_nbi_block()
Parameters
@@ -1554,12 +1631,15 @@ def nvshmem_putmem_signal_nbi_block(dst, src, nelems, sig_addr, signal, sig_op,
The call expression.
""" # noqa: E501
- return call_intrin(
- "", "tirx.nvshmem.putmem_signal_nbi_block", dst, src, nelems, sig_addr, signal, sig_op, pe
+ return Call(
+ "tirx.nvshmem.putmem_signal_nbi_block",
+ [dst, src, nelems, sig_addr, signal, sig_op, pe],
+ ty=ty,
+ span=span,
)
-def nvshmem_fence():
+def nvshmem_fence(*, ty=None, span=None):
"""TVM intrinsic to call nvshmem_fence()
Returns
@@ -1568,10 +1648,10 @@ def nvshmem_fence():
The call expression.
"""
- return call_intrin("", "tirx.nvshmem.fence")
+ return Call("tirx.nvshmem.fence", [], ty=ty, span=span)
-def nvshmem_barrier_all():
+def nvshmem_barrier_all(*, ty=None, span=None):
"""TVM intrinsic to call nvshmem_barrier_all()
Returns
@@ -1580,7 +1660,7 @@ def nvshmem_barrier_all():
The call expression.
"""
- return call_intrin("", "tirx.nvshmem.barrier_all")
+ return Call("tirx.nvshmem.barrier_all", [], ty=ty, span=span)
# Canonical Op builders also supply the historical direct-import aliases.
diff --git a/python/tvm/backend/cuda/ptx/engine.py b/python/tvm/backend/cuda/ptx/engine.py
index 1d15e5c294..473b5bd20b 100644
--- a/python/tvm/backend/cuda/ptx/engine.py
+++ b/python/tvm/backend/cuda/ptx/engine.py
@@ -37,10 +37,12 @@ per-instruction generated or hand-written code:
``T.ptx.ld.acquire.gpu.global_.b32(val, ptr)``.
"""
+from tvm_ffi import get_global_func, register_global_func, structural_equal
+
from tvm.backend.cuda.codegen.registry import register_codegen
from tvm.backend.cuda.codegen.utils import parse_str
from tvm.backend.cuda.op import cuda_cvta_generic_to_shared, cuda_func_call
-from tvm.ir import Call, Op, TensorLoad
+from tvm.ir import Call, Op, StringImm, TensorLoad
from tvm.ir.op import register_op_attr
from tvm.ir.type import PointerType, PrimType
from tvm.runtime import const
@@ -109,6 +111,12 @@ def register_table(table: dict[str, InstructionEntry]) -> None:
register_op_attr(entry.op_name, "TIRxOpCategory", "device_intrin")
register_op_attr(entry.op_name, "TDeviceIntrinsicNamespace", "ptx")
register_op_attr(entry.op_name, "TFixedReturnType", PrimType("void"))
+ if Op.get(entry.op_name).get_attr("__tvm_doc_translate_op_call__") is None:
+ register_op_attr(
+ entry.op_name,
+ "__tvm_doc_translate_op_call__",
+ get_global_func("script.printer.PTXCallDocTranslate")(),
+ )
register_codegen(f"ptx.{entry.name}")(_make_codegen(entry))
@@ -122,6 +130,12 @@ def register_addr() -> None:
register_op_attr(_ADDR_OP_NAME, "TIRxOpCategory", "device_intrin")
register_op_attr(_ADDR_OP_NAME, "TDeviceIntrinsicNamespace", "ptx")
register_op_attr(_ADDR_OP_NAME, "FInferType", _infer_addr_type)
+ if op.get_attr("__tvm_doc_translate_op_call__") is None:
+ register_op_attr(
+ _ADDR_OP_NAME,
+ "__tvm_doc_translate_op_call__",
+ get_global_func("script.printer.PTXCallDocTranslate")(),
+ )
register_codegen("ptx.addr")(_unconsumed_addr_codegen)
@@ -981,3 +995,20 @@ class PTXNamespace:
def __repr__(self):
return f"<T.ptx: {len(self._family_names())} instruction families>"
+
+
+@register_global_func("script.printer.PTXCallCanRoundtrip")
+def _ptx_call_can_roundtrip(call):
+ """Check that printed operands dispatch to the same PTX instruction."""
+ name = call.op.get_attr("TScriptPrinterName").rsplit(".", 1)[-1]
+ args = []
+ try:
+ for arg in call.args:
+ if isinstance(arg, StringImm):
+ arg = arg.value
+ elif isinstance(arg, Call) and arg.op.same_as(Op.get(_ADDR_OP_NAME)):
+ arg = AddrArg(*arg.args)
+ args.append(arg)
+ return structural_equal(call, getattr(PTXNamespace(), name)(*args))
+ except (ValueError, TypeError, IndexError, AttributeError):
+ return False
diff --git a/python/tvm/backend/cuda/script.py b/python/tvm/backend/cuda/script.py
index 20db026759..cc95702241 100644
--- a/python/tvm/backend/cuda/script.py
+++ b/python/tvm/backend/cuda/script.py
@@ -18,14 +18,8 @@
from __future__ import annotations
-from collections.abc import Callable
-from typing import Any
-
from tvm import ir as _ir
from tvm.backend.cuda import op as _cuda_op
-from tvm.tirx import is_tensor_var
-from tvm.tirx import op as _tir_op
-from tvm.tirx.script.ir_builder.op import _dtype_forward, _op_wrapper
# pylint: disable=protected-access
@@ -37,27 +31,12 @@ class _CpAsyncRaw:
# Legacy variant: takes (dst_ptr, dst_offset, src_ptr, src_offset,
# cp_size). Offsets are folded into the pointers; lowers through the
# raw op below.
- self.legacy = _dtype_forward(_cuda_op.ptx_cp_async_legacy)
+ self.legacy = _cuda_op.ptx_cp_async_legacy
+ op = _ir.Op.get("tirx.s_tir.cp_async_raw")
+ self._construct = _ir.op._make_op_api(op, __name__)
def __call__(self, *args, **kwds):
- # The 6-arg form ``(elem_dtype, dst, dst_off, src, src_off, cp_size)``
- # the printer round-trips for the raw ``tirx.s_tir.cp_async_raw`` Call
- # emitted by ``tvm.backend.cuda.transform.InjectPTXAsyncCopy``. The
- # pass-emitted Call has 5 args (no ``tvm_access_ptr`` fold) and a
- # per-element-dtype Call.dtype, so build it directly.
- if len(args) == 6 and isinstance(args[0], str) and "dtype" not in kwds:
- import tvm
-
- elem_dtype, dst, dst_off, src, src_off, cp_size = args
- return tvm.ir.Call(
- tvm.ir.Op.get("tirx.s_tir.cp_async_raw"),
- [dst, dst_off, src, src_off, cp_size],
- ty=tvm.ir.PrimType(elem_dtype),
- )
- raise TypeError(
- "T.s_tir.cp_async_raw only accepts the printed 6-arg raw form; "
- 'issue new copies through T.ptx["cp.async..."]'
- )
+ return self._construct(*args, **kwds)
class STIRNamespace:
@@ -84,27 +63,29 @@ class PTXLegacyNamespace:
def __init__(self):
# Same lowered asm as T.ptx.mma, but the accumulator doubles as the
# destination and the offsets are explicit.
- self.mma = _dtype_forward(_cuda_op.ptx_legacy_mma)
+ self.mma = _cuda_op.ptx_legacy_mma
# (trans, num, dtype, local_ptr, local_offset, smem_ptr, smem_offset)
- self.ldmatrix = _dtype_forward(_cuda_op.ptx_legacy_ldmatrix)
+ self.ldmatrix = _cuda_op.ptx_legacy_ldmatrix
class CudaWgmmaNamespace:
"""WGMMA companions that are not PTX instructions (pure-C / empty-asm helpers)."""
def __init__(self):
- self.noop_barrier = _op_wrapper(_cuda_op.cuda_wgmma_noop_barrier)
- self.encode_matrix_descriptor = _op_wrapper(_cuda_op.cuda_wgmma_encode_matrix_descriptor)
+ self.noop_barrier = _cuda_op.cuda_wgmma_noop_barrier
+ self.encode_matrix_descriptor = _cuda_op.cuda_wgmma_encode_matrix_descriptor
class CudaTcgen05Namespace:
"""tcgen05 companions that are not PTX instructions (pure-C descriptor packers)."""
def __init__(self):
- self.encode_matrix_descriptor = _op_wrapper(_cuda_op.cuda_tcgen05_encode_matrix_descriptor)
- self.encode_instr_descriptor = _op_wrapper(_cuda_op.cuda_tcgen05_encode_instr_descriptor)
- self.encode_instr_descriptor_block_scaled = _op_wrapper(
- _cuda_op.cuda_tcgen05_encode_instr_descriptor_block_scaled
+ self.encode_matrix_descriptor = _cuda_op.cuda_tcgen05_encode_matrix_descriptor
+ self.encode_instr_descriptor = _ir.op._make_op_api(
+ _ir.Op.get("tirx.cuda.tcgen05_encode_instr_descriptor"), __name__
+ )
+ self.encode_instr_descriptor_block_scaled = _ir.op._make_op_api(
+ _ir.Op.get("tirx.cuda.tcgen05_encode_instr_descriptor_block_scaled"), __name__
)
@@ -112,88 +93,41 @@ class IketNamespace:
"""Frontend-only NVIDIA IKET annotations."""
def __init__(self):
- self.mark = _op_wrapper(_cuda_op.cuda_iket_mark)
- self.range_start = _op_wrapper(_cuda_op.cuda_iket_range_start)
- self.range_end = _op_wrapper(_cuda_op.cuda_iket_range_end)
- self.range_push = _op_wrapper(_cuda_op.cuda_iket_range_push)
- self.range_pop = _op_wrapper(_cuda_op.cuda_iket_range_pop)
- self.sentinel_token = _op_wrapper(_cuda_op.cuda_iket_sentinel_token)
- self.official_event = _op_wrapper(_cuda_op.cuda_iket_official_event)
-
-
-def _shfl_sync(mask, var, lane, width):
- if is_tensor_var(var):
- var = var[0]
- return _tir_op.call_intrin(var.ty, "tirx.cuda.__shfl_sync", mask, var, lane, width)
-
-
-def _shfl_up_sync(mask, var, delta, width):
- if is_tensor_var(var):
- var = var[0]
- return _tir_op.call_intrin(var.ty, "tirx.cuda.__shfl_up_sync", mask, var, delta, width)
-
-
-def _shfl_down_sync(mask, var, delta, width):
- if is_tensor_var(var):
- var = var[0]
- return _tir_op.call_intrin(var.ty, "tirx.cuda.__shfl_down_sync", mask, var, delta, width)
-
-
-def _shfl_xor_sync(mask, var, lane_mask, width):
- if is_tensor_var(var):
- var = var[0]
- return _tir_op.call_intrin(var.ty, "tirx.cuda.__shfl_xor_sync", mask, var, lane_mask, width)
-
-
-def _activemask():
- return _tir_op.call_intrin("uint32", "tirx.cuda.__activemask")
+ self.mark = _cuda_op.cuda_iket_mark
+ self.range_start = _cuda_op.cuda_iket_range_start
+ self.range_end = _cuda_op.cuda_iket_range_end
+ self.range_push = _cuda_op.cuda_iket_range_push
+ self.range_pop = _cuda_op.cuda_iket_range_pop
+ self.sentinel_token = _cuda_op.cuda_iket_sentinel_token
+ self.official_event = _cuda_op.cuda_iket_official_event
iket = IketNamespace()
wgmma = CudaWgmmaNamespace()
tcgen05 = CudaTcgen05Namespace()
-mov_sreg: Callable[..., Any] = _cuda_op.cuda_mov_sreg
-wait_until = _cuda_op.cuda_wait_until
-atomic_add = _cuda_op.cuda_atomic_add
-warp_reduce = _cuda_op.cuda_warp_reduce
warp_sum = _cuda_op.cuda_warp_sum
warp_max = _cuda_op.cuda_warp_max
warp_min = _cuda_op.cuda_warp_min
-cta_reduce = _cuda_op.cuda_cta_reduce
cta_sum = _cuda_op.cuda_cta_sum
cta_max = _cuda_op.cuda_cta_max
cta_min = _cuda_op.cuda_cta_min
-atomic_cas = _cuda_op.cuda_atomic_cas
func_call = _cuda_op.cuda_func_call
-ldg = _cuda_op.cuda_ldg
sm100_2sm_leader_smem_addr = _cuda_op.cuda_sm100_2sm_leader_smem_addr
timer_init = _cuda_op.timer_init_cuda
timer_start = _cuda_op.timer_start_cuda
timer_end = _cuda_op.timer_end_cuda
timer_finalize = _cuda_op.timer_finalize_cuda
-mma_store = _dtype_forward(_cuda_op.mma_store)
-mma_fill = _dtype_forward(_cuda_op.mma_fill)
-mma_store_legacy = _dtype_forward(_cuda_op.mma_store_legacy)
-mma_fill_legacy = _dtype_forward(_cuda_op.mma_fill_legacy)
-mov_sreg.__tvm_op__ = _ir.Op.get("tirx.cuda.mov_sreg")
-wait_until.__tvm_op__ = _ir.Op.get("tirx.cuda.wait_until")
-atomic_add.__tvm_op__ = _ir.Op.get("tirx.cuda.atomic_add")
-warp_reduce.__tvm_op__ = _ir.Op.get("tirx.cuda.warp_reduce")
-cta_reduce.__tvm_op__ = _ir.Op.get("tirx.cuda.cta_reduce")
-atomic_cas.__tvm_op__ = _ir.Op.get("tirx.cuda.atomic_cas")
-func_call.__tvm_op__ = _ir.Op.get("tirx.cuda.func_call")
-ldg.__tvm_op__ = _ir.Op.get("tirx.cuda.ldg")
-__shfl_sync = _shfl_sync
-__shfl_sync.__tvm_op__ = _ir.Op.get("tirx.cuda.__shfl_sync")
-__shfl_up_sync = _shfl_up_sync
-__shfl_up_sync.__tvm_op__ = _ir.Op.get("tirx.cuda.__shfl_up_sync")
-__shfl_down_sync = _shfl_down_sync
-__shfl_down_sync.__tvm_op__ = _ir.Op.get("tirx.cuda.__shfl_down_sync")
-__shfl_xor_sync = _shfl_xor_sync
-__shfl_xor_sync.__tvm_op__ = _ir.Op.get("tirx.cuda.__shfl_xor_sync")
-__activemask = _activemask
-__activemask.__tvm_op__ = _ir.Op.get("tirx.cuda.__activemask")
-
+mma_store = _ir.op._make_op_api(_ir.Op.get("tirx.mma_store"), __name__)
+mma_fill = _ir.op._make_op_api(_ir.Op.get("tirx.mma_fill"), __name__)
+mma_store_legacy = _cuda_op.mma_store_legacy
+mma_fill_legacy = _cuda_op.mma_fill_legacy
+atomic_add = _cuda_op.cuda_atomic_add
+atomic_cas = _cuda_op.cuda_atomic_cas
+cta_reduce = _cuda_op.cuda_cta_reduce
+ldg = _cuda_op.cuda_ldg
+mov_sreg = _cuda_op.cuda_mov_sreg
+wait_until = _cuda_op.cuda_wait_until
+warp_reduce = _cuda_op.cuda_warp_reduce
_ir.op._init_op_api("tirx.cuda", __name__)
@@ -201,13 +135,13 @@ class NVSHMEMNamespace:
"""The NVSHMEM intrinsics submodule."""
def __init__(self):
- self.my_pe = _op_wrapper(_cuda_op.nvshmem_my_pe)
- self.n_pes = _op_wrapper(_cuda_op.nvshmem_n_pes)
- self.signal_op = _op_wrapper(_cuda_op.nvshmem_signal_op)
- self.wait_until = _op_wrapper(_cuda_op.nvshmem_wait_until)
- self.quiet = _op_wrapper(_cuda_op.nvshmem_quiet)
- self.fence = _op_wrapper(_cuda_op.nvshmem_fence)
- self.barrier_all = _op_wrapper(_cuda_op.nvshmem_barrier_all)
+ self.my_pe = _cuda_op.nvshmem_my_pe
+ self.n_pes = _cuda_op.nvshmem_n_pes
+ self.signal_op = _cuda_op.nvshmem_signal_op
+ self.wait_until = _cuda_op.nvshmem_wait_until
+ self.quiet = _cuda_op.nvshmem_quiet
+ self.fence = _cuda_op.nvshmem_fence
+ self.barrier_all = _cuda_op.nvshmem_barrier_all
self.getmem_nbi = NVSHMEMGetMemNBINamespace()
self.putmem_nbi = NVSHMEMPutMemNBINamespace()
self.putmem_signal_nbi = NVSHMEMPutMemSignalNBINamespace()
@@ -217,39 +151,36 @@ class NVSHMEMGetMemNBINamespace:
"""The NVSHMEM GetMemNBI intrinsics submodule."""
def __init__(self):
- self.warp = _op_wrapper(_cuda_op.nvshmem_getmem_nbi_warp)
- self.block = _op_wrapper(_cuda_op.nvshmem_getmem_nbi_block)
+ self.warp = _cuda_op.nvshmem_getmem_nbi_warp
+ self.block = _cuda_op.nvshmem_getmem_nbi_block
def __call__(self, *args, **kwds):
- return _op_wrapper(_cuda_op.nvshmem_getmem_nbi)(*args, **kwds)
+ return _cuda_op.nvshmem_getmem_nbi(*args, **kwds)
# __call__ corresponds to nvshmem_getmem_nbi
- __tir_call_op_name__ = "nvshmem_getmem_nbi"
class NVSHMEMPutMemNBINamespace:
"""The NVSHMEM PutMemNBI intrinsics submodule."""
def __init__(self):
- self.warp = _op_wrapper(_cuda_op.nvshmem_putmem_nbi_warp)
- self.block = _op_wrapper(_cuda_op.nvshmem_putmem_nbi_block)
+ self.warp = _cuda_op.nvshmem_putmem_nbi_warp
+ self.block = _cuda_op.nvshmem_putmem_nbi_block
def __call__(self, *args, **kwds):
- return _op_wrapper(_cuda_op.nvshmem_putmem_nbi)(*args, **kwds)
+ return _cuda_op.nvshmem_putmem_nbi(*args, **kwds)
# __call__ corresponds to nvshmem_putmem_nbi
- __tir_call_op_name__ = "nvshmem_putmem_nbi"
class NVSHMEMPutMemSignalNBINamespace:
"""The NVSHMEM PutMemSignalNBI intrinsics submodule."""
def __init__(self):
- self.warp = _op_wrapper(_cuda_op.nvshmem_putmem_signal_nbi_warp)
- self.block = _op_wrapper(_cuda_op.nvshmem_putmem_signal_nbi_block)
+ self.warp = _cuda_op.nvshmem_putmem_signal_nbi_warp
+ self.block = _cuda_op.nvshmem_putmem_signal_nbi_block
def __call__(self, *args, **kwds):
- return _op_wrapper(_cuda_op.nvshmem_putmem_signal_nbi)(*args, **kwds)
+ return _cuda_op.nvshmem_putmem_signal_nbi(*args, **kwds)
# __call__ corresponds to nvshmem_putmem_signal_nbi
- __tir_call_op_name__ = "nvshmem_putmem_signal_nbi"
diff --git a/python/tvm/backend/cuda/tile_primitive/common.py b/python/tvm/backend/cuda/tile_primitive/common.py
index 977b4abc54..00372b8a26 100644
--- a/python/tvm/backend/cuda/tile_primitive/common.py
+++ b/python/tvm/backend/cuda/tile_primitive/common.py
@@ -58,9 +58,9 @@ def get_indices(nth, start, extent):
def smem_desc_replace_lo(desc_val, desc_lo):
"""Replace the lower address lane of a 64-bit SMEM descriptor."""
- desc_halves = T.reinterpret("uint32x2", desc_val)
+ desc_halves = T.reinterpret(desc_val, ty="uint32x2")
desc_hi = T.Shuffle([desc_halves], [1])
- return T.reinterpret("uint64", T.Shuffle([T.cast(desc_lo, "uint32"), desc_hi], [0, 1]))
+ return T.reinterpret(T.Shuffle([T.cast(desc_lo, "uint32"), desc_hi], [0, 1]), ty="uint64")
def smem_desc_add_16B_offset(desc_val, offset):
@@ -69,7 +69,7 @@ def smem_desc_add_16B_offset(desc_val, offset):
The address lane wraps as uint32 without carrying into the descriptor's
upper control bits.
"""
- desc_halves = T.reinterpret("uint32x2", desc_val)
+ desc_halves = T.reinterpret(desc_val, ty="uint32x2")
desc_lo = T.Shuffle([desc_halves], [0]) + T.cast(offset, "uint32")
return smem_desc_replace_lo(desc_val, desc_lo)
diff --git a/python/tvm/backend/cuda/tile_primitive/copy_async/tcgen05_cp.py b/python/tvm/backend/cuda/tile_primitive/copy_async/tcgen05_cp.py
index c8f8ad4510..2b2ee35d96 100644
--- a/python/tvm/backend/cuda/tile_primitive/copy_async/tcgen05_cp.py
+++ b/python/tvm/backend/cuda/tile_primitive/copy_async/tcgen05_cp.py
@@ -682,7 +682,7 @@ def _get_or_create_desc(sctx, s_buf, ldo, sdo, swizzle):
desc_buf = tvm.tirx.decl_tensor((1,), "uint64", name="cp_desc", scope="local")
encode_call = T.cuda.tcgen05.encode_matrix_descriptor(
- desc_buf.data, T.reinterpret("handle", T.uint64(0)), ldo, sdo, swizzle
+ desc_buf.data, T.reinterpret(T.uint64(0), ty="handle"), ldo, sdo, swizzle
)
wrap = SeqStmt(
[
diff --git a/python/tvm/backend/metal/op.py b/python/tvm/backend/metal/op.py
index 50ef4f39e2..d47f7cec02 100644
--- a/python/tvm/backend/metal/op.py
+++ b/python/tvm/backend/metal/op.py
@@ -18,61 +18,74 @@
from __future__ import annotations
-from tvm.tirx.op import call_intrin
+from tvm.ir import Call
-def make_filled_simdgroup_matrix(d, index, value, col=8, row=8):
+def make_filled_simdgroup_matrix(d, index, value, col=8, row=8, *, ty=None, span=None):
"""Create a filled SIMDGroup matrix."""
- return call_intrin("void", "tirx.make_filled_simdgroup_matrix", d, index, value, col, row)
+ return Call(
+ "tirx.make_filled_simdgroup_matrix",
+ [d, index, value, col, row],
+ ty=ty,
+ span=span,
+ )
-def simdgroup_load(d, index, ptr, stride, col=8, row=8, transpose_matrix=False):
+def simdgroup_load(
+ d,
+ index,
+ ptr,
+ stride,
+ col=8,
+ row=8,
+ transpose_matrix=False,
+ *,
+ ty=None,
+ span=None,
+):
"""Load data from device or threadgroup memory to simdgroup."""
- return call_intrin(
- "void",
+ return Call(
"tirx.simdgroup_load",
- d,
- index,
- ptr,
- stride,
- col,
- row,
- transpose_matrix,
+ [d, index, ptr, stride, col, row, transpose_matrix],
+ ty=ty,
+ span=span,
)
-def simdgroup_store(d, index, ptr, stride, col=8, row=8, transpose_matrix=False):
+def simdgroup_store(
+ d,
+ index,
+ ptr,
+ stride,
+ col=8,
+ row=8,
+ transpose_matrix=False,
+ *,
+ ty=None,
+ span=None,
+):
"""Store data from simdgroup to device or threadgroup memory."""
- return call_intrin(
- "void",
+ return Call(
"tirx.simdgroup_store",
- d,
- index,
- ptr,
- stride,
- col,
- row,
- transpose_matrix,
+ [d, index, ptr, stride, col, row, transpose_matrix],
+ ty=ty,
+ span=span,
)
-def simdgroup_multiply_accumulate(d, index_d, a, index_a, b, index_b, c, index_c):
+def simdgroup_multiply_accumulate(
+ d, index_d, a, index_a, b, index_b, c, index_c, *, ty=None, span=None
+):
"""Multiply and accumulate two matrices in simdgroup."""
- return call_intrin(
- "void",
+ return Call(
"tirx.simdgroup_multiply_accumulate",
- d,
- index_d,
- a,
- index_a,
- b,
- index_b,
- c,
- index_c,
+ [d, index_d, a, index_a, b, index_b, c, index_c],
+ ty=ty,
+ span=span,
)
diff --git a/python/tvm/backend/metal/script.py b/python/tvm/backend/metal/script.py
index 263f309f1e..dd220af853 100644
--- a/python/tvm/backend/metal/script.py
+++ b/python/tvm/backend/metal/script.py
@@ -19,37 +19,41 @@
from __future__ import annotations
from tvm.backend.metal import op as _metal_op
+from tvm.ir import Op
+from tvm.ir.op import _make_op_api
from tvm.tirx import is_tensor_var
-from tvm.tirx import op as _tir_op
-from tvm.tirx.script.ir_builder.op import _op_wrapper
+
+_simd_shuffle = _make_op_api(Op.get("tirx.metal.simd_shuffle"), __name__)
+_simd_shuffle_up = _make_op_api(Op.get("tirx.metal.simd_shuffle_up"), __name__)
+_simd_shuffle_down = _make_op_api(Op.get("tirx.metal.simd_shuffle_down"), __name__)
class MetalNamespace:
"""The Metal intrinsics submodule."""
def __init__(self):
- self.make_filled_simdgroup_matrix = _op_wrapper(_metal_op.make_filled_simdgroup_matrix)
- self.simdgroup_load = _op_wrapper(_metal_op.simdgroup_load)
- self.simdgroup_store = _op_wrapper(_metal_op.simdgroup_store)
- self.simdgroup_multiply_accumulate = _op_wrapper(_metal_op.simdgroup_multiply_accumulate)
+ self.make_filled_simdgroup_matrix = _metal_op.make_filled_simdgroup_matrix
+ self.simdgroup_load = _metal_op.simdgroup_load
+ self.simdgroup_store = _metal_op.simdgroup_store
+ self.simdgroup_multiply_accumulate = _metal_op.simdgroup_multiply_accumulate
@staticmethod
- def simd_shuffle(var, lane):
+ def simd_shuffle(var, lane, *, ty=None, span=None):
if is_tensor_var(var):
var = var[0]
- return _tir_op.call_intrin(var.ty, "tirx.metal.simd_shuffle", var, lane)
+ return _simd_shuffle(var, lane, ty=ty, span=span)
@staticmethod
- def simd_shuffle_up(var, delta):
+ def simd_shuffle_up(var, delta, *, ty=None, span=None):
if is_tensor_var(var):
var = var[0]
- return _tir_op.call_intrin(var.ty, "tirx.metal.simd_shuffle_up", var, delta)
+ return _simd_shuffle_up(var, delta, ty=ty, span=span)
@staticmethod
- def simd_shuffle_down(var, delta):
+ def simd_shuffle_down(var, delta, *, ty=None, span=None):
if is_tensor_var(var):
var = var[0]
- return _tir_op.call_intrin(var.ty, "tirx.metal.simd_shuffle_down", var, delta)
+ return _simd_shuffle_down(var, delta, ty=ty, span=span)
__all__ = ["MetalNamespace"]
diff --git a/python/tvm/backend/trn/__init__.py b/python/tvm/backend/trn/__init__.py
index de66492578..e9f1a80bc5 100644
--- a/python/tvm/backend/trn/__init__.py
+++ b/python/tvm/backend/trn/__init__.py
@@ -46,16 +46,16 @@ def register_backend():
compilation_pipeline.register_tir_pipeline("trn", trn_pipeline.trn_pipeline)
-def script_namespace(op_wrapper=None):
+def script_namespace():
"""Return the Trainium TVMScript namespace object."""
from .script import NKINamespace # pylint: disable=import-outside-toplevel
- return NKINamespace(op_wrapper)
+ return NKINamespace()
-def script_namespaces(op_wrapper=None, **_):
+def script_namespaces(**_):
"""Return Trainium-owned TVMScript namespaces."""
- return {"nki": script_namespace(op_wrapper)}
+ return {"nki": script_namespace()}
def __getattr__(name: str):
diff --git a/python/tvm/backend/trn/op.py b/python/tvm/backend/trn/op.py
index 88cdd048fa..077d0fae6a 100644
--- a/python/tvm/backend/trn/op.py
+++ b/python/tvm/backend/trn/op.py
@@ -18,119 +18,181 @@
from __future__ import annotations
-from tvm.tirx.op import call_intrin
+from tvm.ir import Call
-def nki_load(res, data):
- return call_intrin("", "tirx.nki.load", res, data)
+def nki_load(res, data, *, ty=None, span=None):
+ return Call("tirx.nki.load", [res, data], ty=ty, span=span)
-def nki_store(res, data):
- return call_intrin("", "tirx.nki.store", res, data)
+def nki_store(res, data, *, ty=None, span=None):
+ return Call("tirx.nki.store", [res, data], ty=ty, span=span)
-def nki_tensor_copy(res, data):
- return call_intrin("", "tirx.nki.tensor_copy", res, data)
+def nki_tensor_copy(res, data, *, ty=None, span=None):
+ return Call("tirx.nki.tensor_copy", [res, data], ty=ty, span=span)
-def nki_matmul(res, lhs, rhs, accum=True):
- return call_intrin("", "tirx.nki.matmul", res, lhs, rhs, accum)
+def nki_matmul(res, lhs, rhs, accum=True, *, ty=None, span=None):
+ return Call(
+ "tirx.nki.matmul",
+ [res, lhs, rhs, accum],
+ ty=ty,
+ span=span,
+ )
-def nki_activation(result, data, opcode, bias=0.0, scale=1.0):
- return call_intrin("", "tirx.nki.activation", result, data, opcode, bias, scale)
+def nki_activation(result, data, opcode, bias=0.0, scale=1.0, *, ty=None, span=None):
+ return Call(
+ "tirx.nki.activation",
+ [result, data, opcode, bias, scale],
+ ty=ty,
+ span=span,
+ )
-def nki_reciprocal(result, data):
- return call_intrin("", "tirx.nki.reciprocal", result, data)
+def nki_reciprocal(result, data, *, ty=None, span=None):
+ return Call(
+ "tirx.nki.reciprocal",
+ [result, data],
+ ty=ty,
+ span=span,
+ )
-def nki_tensorreduce(result, data, opcode, negate, *axes):
- return call_intrin("", "tirx.nki.tensorreduce", result, data, opcode, negate, *axes)
+def nki_tensorreduce(result, data, opcode, negate, *axes, ty=None, span=None):
+ return Call(
+ "tirx.nki.tensorreduce",
+ [result, data, opcode, negate, *axes],
+ ty=ty,
+ span=span,
+ )
-def nki_tensortensor(result, operand0, operand1, opcode):
- return call_intrin("", "tirx.nki.tensortensor", result, operand0, operand1, opcode)
+def nki_tensortensor(result, operand0, operand1, opcode, *, ty=None, span=None):
+ return Call(
+ "tirx.nki.tensortensor",
+ [result, operand0, operand1, opcode],
+ ty=ty,
+ span=span,
+ )
-def nki_tensorscalar(result, operand0, operand1, opcode, reverse=False):
- return call_intrin("", "tirx.nki.tensorscalar", result, operand0, operand1, opcode, reverse)
+def nki_tensorscalar(
+ result,
+ operand0,
+ operand1,
+ opcode,
+ reverse=False,
+ *,
+ ty=None,
+ span=None,
+):
+ return Call(
+ "tirx.nki.tensorscalar",
+ [result, operand0, operand1, opcode, reverse],
+ ty=ty,
+ span=span,
+ )
-def nki_memset(result, value):
- return call_intrin("", "tirx.nki.memset", result, value)
+def nki_memset(result, value, *, ty=None, span=None):
+ return Call("tirx.nki.memset", [result, value], ty=ty, span=span)
-def nki_activation_reduce(reduce_res, act_res, data, opcode, reduce_opcode, bias=0.0, scale=1.0):
- return call_intrin(
- "",
+def nki_activation_reduce(
+ reduce_res,
+ act_res,
+ data,
+ opcode,
+ reduce_opcode,
+ bias=0.0,
+ scale=1.0,
+ *,
+ ty=None,
+ span=None,
+):
+ return Call(
"tirx.nki.activation_reduce",
- reduce_res,
- act_res,
- data,
- opcode,
- reduce_opcode,
- bias,
- scale,
+ [reduce_res, act_res, data, opcode, reduce_opcode, bias, scale],
+ ty=ty,
+ span=span,
)
def nki_tensorscalar_reduce(
- reduce_res, tensorscalar_res, operand0, operand1, opcode, reduce_opcode, reverse=False
+ reduce_res,
+ tensorscalar_res,
+ operand0,
+ operand1,
+ opcode,
+ reduce_opcode,
+ reverse=False,
+ *,
+ ty=None,
+ span=None,
):
- return call_intrin(
- "",
+ return Call(
"tirx.nki.tensorscalar_reduce",
- reduce_res,
- tensorscalar_res,
- operand0,
- operand1,
- opcode,
- reduce_opcode,
- reverse,
+ [reduce_res, tensorscalar_res, operand0, operand1, opcode, reduce_opcode, reverse],
+ ty=ty,
+ span=span,
)
-def nki_identity(result, size):
- return call_intrin("", "tirx.nki.identity", result, size)
+def nki_identity(result, size, *, ty=None, span=None):
+ return Call("tirx.nki.identity", [result, size], ty=ty, span=span)
def nki_scalar_tensor_tensor(
- result, data, operand0, operand1, opcode0, opcode1, reverse0=False, reverse1=False
+ result,
+ data,
+ operand0,
+ operand1,
+ opcode0,
+ opcode1,
+ reverse0=False,
+ reverse1=False,
+ *,
+ ty=None,
+ span=None,
):
- return call_intrin(
- "",
+ return Call(
"tirx.nki.scalar_tensor_tensor",
- result,
- data,
- operand0,
- operand1,
- opcode0,
- opcode1,
- reverse0,
- reverse1,
+ [result, data, operand0, operand1, opcode0, opcode1, reverse0, reverse1],
+ ty=ty,
+ span=span,
)
def nki_scalar_tensor_scalar(
- result, data, operand0, operand1, opcode0, opcode1, reverse0=False, reverse1=False
+ result,
+ data,
+ operand0,
+ operand1,
+ opcode0,
+ opcode1,
+ reverse0=False,
+ reverse1=False,
+ *,
+ ty=None,
+ span=None,
):
- return call_intrin(
- "",
+ return Call(
"tirx.nki.scalar_tensor_scalar",
- result,
- data,
- operand0,
- operand1,
- opcode0,
- opcode1,
- reverse0,
- reverse1,
+ [result, data, operand0, operand1, opcode0, opcode1, reverse0, reverse1],
+ ty=ty,
+ span=span,
)
-def nki_affine_select(result, pred, true_value, false_value):
- return call_intrin("", "tirx.nki.affine_select", result, pred, true_value, false_value)
+def nki_affine_select(result, pred, true_value, false_value, *, ty=None, span=None):
+ return Call(
+ "tirx.nki.affine_select",
+ [result, pred, true_value, false_value],
+ ty=ty,
+ span=span,
+ )
__all__ = [
diff --git a/python/tvm/backend/trn/script.py b/python/tvm/backend/trn/script.py
index 1bd1baee14..11d335d23b 100644
--- a/python/tvm/backend/trn/script.py
+++ b/python/tvm/backend/trn/script.py
@@ -18,43 +18,29 @@
from __future__ import annotations
-from collections.abc import Callable
-from typing import Any
-
from . import op as _trn_op
-OpWrapper = Callable[[Callable[..., Any]], Callable[..., Any]]
-
-
-def _default_op_wrapper(func: Callable[..., Any]) -> Callable[..., Any]:
- """Create the default facade for a returned-node NKI intrinsic."""
- from tvm.tirx.script.ir_builder.op import _op_wrapper # pylint: disable=import-outside-toplevel
-
- # NKI producers return a single call_intrin node and never emit separately.
- return _op_wrapper(func)
-
class NKINamespace:
"""The NKI instructions submodule."""
- def __init__(self, op_wrapper: OpWrapper | None = None):
- wrap = _default_op_wrapper if op_wrapper is None else op_wrapper
- self.load = wrap(_trn_op.nki_load)
- self.store = wrap(_trn_op.nki_store)
- self.tensor_copy = wrap(_trn_op.nki_tensor_copy)
- self.matmul = wrap(_trn_op.nki_matmul)
- self.activation = wrap(_trn_op.nki_activation)
- self.activation_reduce = wrap(_trn_op.nki_activation_reduce)
- self.reciprocal = wrap(_trn_op.nki_reciprocal)
- self.tensorreduce = wrap(_trn_op.nki_tensorreduce)
- self.tensortensor = wrap(_trn_op.nki_tensortensor)
- self.tensorscalar = wrap(_trn_op.nki_tensorscalar)
- self.tensorscalar_reduce = wrap(_trn_op.nki_tensorscalar_reduce)
- self.scalar_tensor_tensor = wrap(_trn_op.nki_scalar_tensor_tensor)
- self.scalar_tensor_scalar = wrap(_trn_op.nki_scalar_tensor_scalar)
- self.memset = wrap(_trn_op.nki_memset)
- self.identity = wrap(_trn_op.nki_identity)
- self.affine_select = wrap(_trn_op.nki_affine_select)
+ def __init__(self):
+ self.load = _trn_op.nki_load
+ self.store = _trn_op.nki_store
+ self.tensor_copy = _trn_op.nki_tensor_copy
+ self.matmul = _trn_op.nki_matmul
+ self.activation = _trn_op.nki_activation
+ self.activation_reduce = _trn_op.nki_activation_reduce
+ self.reciprocal = _trn_op.nki_reciprocal
+ self.tensorreduce = _trn_op.nki_tensorreduce
+ self.tensortensor = _trn_op.nki_tensortensor
+ self.tensorscalar = _trn_op.nki_tensorscalar
+ self.tensorscalar_reduce = _trn_op.nki_tensorscalar_reduce
+ self.scalar_tensor_tensor = _trn_op.nki_scalar_tensor_tensor
+ self.scalar_tensor_scalar = _trn_op.nki_scalar_tensor_scalar
+ self.memset = _trn_op.nki_memset
+ self.identity = _trn_op.nki_identity
+ self.affine_select = _trn_op.nki_affine_select
@staticmethod
def tensorized_instruction():
@@ -66,4 +52,4 @@ class NKINamespace:
return region("tirx.nki.tensorized_instruction", [])
-__all__ = ["NKINamespace", "OpWrapper"]
+__all__ = ["NKINamespace"]
diff --git a/python/tvm/ir/expr.py b/python/tvm/ir/expr.py
index 9d67db2683..3d8f368c2d 100644
--- a/python/tvm/ir/expr.py
+++ b/python/tvm/ir/expr.py
@@ -511,10 +511,11 @@ class StringImm(Constant):
class Call(_CallableExprWithOp):
"""Core function call node.
- When ``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.
+ Omitted or ``None`` ``ty`` uses available result inference, or a missing
+ type when no deduction is available. Explicit types, including
+ ``Type.missing()``, are preserved exactly. Inference errors propagate.
+ Construction permits provisional IR; :meth:`validate` checks the operator
+ contract explicitly after inputs are ready.
"""
op: Expr
@@ -547,8 +548,6 @@ class Call(_CallableExprWithOp):
op = Op.get(op)
if attrs is not None and isinstance(attrs, dict):
attrs = DictAttrs(attrs)
- if ty is None:
- ty = Type.missing()
if isinstance(ty, str) and ty == "handle":
ty = PointerType(PrimType("void"))
elif ty is not None and not isinstance(ty, Type):
@@ -557,48 +556,9 @@ class Call(_CallableExprWithOp):
ty_args = []
return 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,
- 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.
- 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, ty)
- )
+ def validate(self) -> None:
+ """Check the registered operator contract without changing this Call."""
+ _ffi_api.CallValidate(self)
def reinfer_type(call: Call) -> "tvm.ir.Type":
diff --git a/python/tvm/ir/op.py b/python/tvm/ir/op.py
index 0b406b556f..5c7941ed75 100644
--- a/python/tvm/ir/op.py
+++ b/python/tvm/ir/op.py
@@ -27,21 +27,18 @@ import tvm_ffi
from . import _ffi_api
from .expr import Expr
-_RESULT_TYPE_UNSET = object()
-
def _make_op_api(op, module_name):
"""Build a callable whose operands and result are governed by its Op."""
from .attrs import make_node # pylint: disable=import-outside-toplevel
- from .expr import Call, reinfer_type # pylint: disable=import-outside-toplevel
+ from .expr import Call # pylint: disable=import-outside-toplevel
def call(
*args,
attrs=None,
ty_args=None,
span=None,
- ty=_RESULT_TYPE_UNSET,
- ret_ty=_RESULT_TYPE_UNSET,
+ ty=None,
**kwargs,
):
# Bind named operands using the registered signature, without inventing
@@ -55,30 +52,19 @@ def _make_op_api(op, module_name):
operands.append(kwargs.pop(info.name))
else:
raise TypeError(f"{op.name}: missing operand {info.name!r}")
- if kwargs:
+ if kwargs or (attrs is None and op.attrs_type_key):
if not op.attrs_type_key:
raise TypeError(f"{op.name}: unexpected keyword operands {tuple(kwargs)}")
if attrs is not None:
raise TypeError(f"{op.name}: cannot mix attrs with attribute keywords")
attrs = make_node(op.attrs_type_key, **kwargs)
- if ty is not _RESULT_TYPE_UNSET:
- if ret_ty is not _RESULT_TYPE_UNSET:
- raise TypeError(f"{op.name}: ty and ret_ty cannot both be specified")
- ret_ty = ty
- elif (ret_ty is _RESULT_TYPE_UNSET or ret_ty is None) and (
- op.get_attr("TFixedReturnType") is not None or op.get_attr("FInferType") is not None
- ):
- provisional = Call.unchecked(op, operands, attrs=attrs, ty_args=ty_args, span=span)
- ret_ty = reinfer_type(provisional)
- if ret_ty is _RESULT_TYPE_UNSET:
- ret_ty = None
- return Call(op, operands, attrs=attrs, ty_args=ty_args, span=span, ty=ret_ty)
+ return Call(op, operands, attrs=attrs, ty_args=ty_args, span=span, ty=ty)
call.__name__ = op.name.rsplit(".", 1)[-1]
call.__module__ = module_name
call.__doc__ = op.doc or f"Construct a call to {op.name}."
- call.__tvm_op__ = op
- call.__tvm_standard_call__ = True
+ if op.get_attr("TScriptPrinterName") is None:
+ op.set_attr("TScriptPrinterName", op.name)
return call
@@ -91,18 +77,18 @@ def _init_op_api(namespace, target_module_name=None):
module or SimpleNamespace containers; underscores are ordinary name parts.
Generated functions accept registered positional/named operands plus
- ``attrs``, ``ty_args``, ``span`` and ``ty`` (also spelled ``ret_ty``).
+ ``attrs``, ``ty_args``, ``span`` and ``ty``.
Omitting the result invokes an available Op inference hook; without one,
Call retains a missing type.
- Explicit results and inference/validation errors are preserved. Attribute
+ Explicit results and inference errors are preserved; Call.validate checks
+ the operator contract separately. Attribute
keywords construct the registered attrs schema, when one is declared.
- Existing generated functions are reused. A deliberate wrapper may declare
- ``__tvm_op__ = Op.get(name)`` to retain ownership of that name. This declares
- identity, not semantic equivalence: the wrapper must accept printed calls
- or have an appropriate exceptional printer hook. Other collisions fail
- before any functions are installed. Script namespaces separately publish
- their callable names to the printer after initialization. Returns None.
+ Existing Python callables retain ownership of their names; only missing
+ exposures are generated. Existing functions must accept printed calls or
+ have an appropriate exceptional printer hook. Non-callable collisions fail
+ before any functions are installed. Generated constructors publish a canonical
+ printer name when no explicit alias is registered. Returns None.
"""
if not namespace or any(
not part.isidentifier() or keyword.iskeyword(part) for part in namespace.split(".")
@@ -110,6 +96,7 @@ def _init_op_api(namespace, target_module_name=None):
raise ValueError(f"Invalid Op namespace {namespace!r}")
target = sys.modules[target_module_name or namespace]
pending = []
+ existing = []
destinations = set()
for name in sorted(Op.list_op_names()):
if not name.startswith(namespace + "."):
@@ -129,13 +116,16 @@ def _init_op_api(namespace, target_module_name=None):
op = Op.get(name)
if parts[-1] in vars(container):
current = vars(container)[parts[-1]]
- identity = getattr(current, "__tvm_op__", None)
- if not callable(current) or not isinstance(identity, Op) or not identity.same_as(op):
+ if not callable(current):
raise ValueError(f"Op {name!r} conflicts with an existing Python attribute")
+ existing.append(op)
else:
pending.append((container, parts[-1], _make_op_api(op, target.__name__)))
for container, name, call in pending:
setattr(container, name, call)
+ for op in existing:
+ if op.get_attr("TScriptPrinterName") is None:
+ op.set_attr("TScriptPrinterName", op.name)
@tvm_ffi.register_object("ir.Op")
diff --git a/python/tvm/relax/frontend/tflite/tflite_frontend.py b/python/tvm/relax/frontend/tflite/tflite_frontend.py
index f22d013c4e..fad6ff548c 100644
--- a/python/tvm/relax/frontend/tflite/tflite_frontend.py
+++ b/python/tvm/relax/frontend/tflite/tflite_frontend.py
@@ -8608,8 +8608,8 @@ def _build_stablehlo_rng_bit_generator_function(algorithm, state_len, out_dtype,
if is_64bit:
low = _u64(words[2 * write_index])
high = _u64(words[2 * write_index + 1])
- return T.reinterpret(out_dtype, low | (high << T.uint64(32)))
- return T.reinterpret(out_dtype, words[write_index])
+ return T.reinterpret(low | (high << T.uint64(32)), ty=out_dtype)
+ return T.reinterpret(words[write_index], ty=out_dtype)
if algorithm == "threefry":
diff --git a/python/tvm/relax/frontend/torch/exported_program_translator.py b/python/tvm/relax/frontend/torch/exported_program_translator.py
index 44804e4d2a..e42cf012a0 100644
--- a/python/tvm/relax/frontend/torch/exported_program_translator.py
+++ b/python/tvm/relax/frontend/torch/exported_program_translator.py
@@ -1735,7 +1735,7 @@ class ExportedProgramImporter(BaseFXGraphImporter):
else node.kwargs.get("assert_msg", "Assertion Failed")
)
return self.block_builder.emit(
- relax.op.assert_op(condition, [tvm.ir.StringImm(message)], format="{}")
+ relax.op.assert_op(condition, "{}", tvm.ir.StringImm(message))
)
def create_convert_map(
diff --git a/python/tvm/relax/op/base.py b/python/tvm/relax/op/base.py
index 1c7a6a7f51..f67ef1ef03 100644
--- a/python/tvm/relax/op/base.py
+++ b/python/tvm/relax/op/base.py
@@ -24,6 +24,7 @@ import tvm_ffi
import tvm
import tvm.runtime
from tvm.ir import Call, Op, StringImm
+from tvm.ir.attrs import make_node as _make_attrs
from tvm.runtime import Object, ObjectConvertible
from ..expr import Expr, ExternFunc, GlobalVar, Var
@@ -62,7 +63,7 @@ def register_gradient(
return tvm.ir.register_op_attr(op_name, "FPrimalGradient", fgradient, override)
-def null_value() -> Call:
+def null_value(*, ty=None, span=None) -> Call:
"""Create a call node that represents a null value object.
Returns
@@ -70,7 +71,7 @@ def null_value() -> Call:
ret: Call
The created call node.
"""
- return _ffi_api.null_value() # type: ignore
+ return Call("relax.null_value", [], ty=ty, span=span) # type: ignore
def _wrap_inline_arg_tuple(args) -> Expr:
@@ -130,7 +131,7 @@ def call_tir(
return _ffi_api.call_tir(gvar, args, out_ty) # type: ignore
-def call_tir_packed(gvar: GlobalVar, args: Expr) -> Call:
+def call_tir_packed(gvar: GlobalVar, args: Expr, *, ty=None, span=None) -> Call:
"""Call a TIRx Function through its native packed-call contract.
Every native parameter is supplied explicitly, in order. Unlike
@@ -203,7 +204,7 @@ def call_tir_packed(gvar: GlobalVar, args: Expr) -> Call:
relax.call_tir_packed(copy, (source, destination))
"""
args = _wrap_inline_arg_tuple(args)
- return Call.unchecked("relax.call_tir_packed", [gvar, args])
+ return Call("relax.call_tir_packed", [gvar, args], ty=ty, span=span)
def call_tir_with_grad(
@@ -307,7 +308,7 @@ def call_tir_inplace(
"""
args = _wrap_inline_arg_tuple(args)
- if not isinstance(inplace_indices, list):
+ if inplace_indices is not None and not isinstance(inplace_indices, list):
inplace_indices = [inplace_indices]
if not isinstance(out_ty, list):
@@ -402,6 +403,8 @@ def call_builtin_with_ctx(
args: Expr,
*,
ty_args: Type | list[Type] | None = None,
+ ty=None,
+ span=None,
) -> Call:
"""Call a builtin function func.
@@ -429,17 +432,25 @@ def call_builtin_with_ctx(
if ty_args is not None and not isinstance(ty_args, list | tuple):
ty_args = [ty_args]
- return _ffi_api.call_builtin_with_ctx( # type: ignore
- func,
- args,
- ty_args, # type: ignore
+ if ty_args is not None:
+ ty_args = [
+ value()
+ if callable(value)
+ else value.asobject()
+ if isinstance(value, ObjectConvertible)
+ else value
+ for value in ty_args
+ ]
+ return Call(
+ "relax.call_builtin_with_ctx",
+ [func, args],
+ ty_args=ty_args,
+ ty=ty,
+ span=span,
)
-def make_closure(
- func: Expr,
- args: Expr,
-) -> Object:
+def make_closure(func: Expr, args: Expr, *, ty=None, span=None) -> Object:
"""
Create a closure with free variables and return the closure.
@@ -460,13 +471,11 @@ def make_closure(
args = _wrap_inline_arg_tuple(args)
- return _ffi_api.make_closure(func, args) # type: ignore
+ return Call("relax.make_closure", [func, args], ty=ty, span=span) # type: ignore
def invoke_closure(
- closure: Expr,
- args: Expr,
- ty_args: list[Type] | Type,
+ closure: Expr, args: Expr, ty_args: list[Type] | Type, *, ty=None, span=None
) -> Call:
"""
Invoke a closure.
@@ -492,7 +501,13 @@ def invoke_closure(
if not isinstance(ty_args, list | tuple):
ty_args = [ty_args]
- return _ffi_api.invoke_closure(closure, args, ty_args) # type: ignore
+ return Call(
+ "relax.invoke_closure",
+ [closure, args],
+ ty_args=ty_args,
+ ty=ty,
+ span=span,
+ ) # type: ignore
def render_object(val: tvm.Object) -> str:
@@ -585,8 +600,8 @@ def relax_assert_op(condition: tvm.Object, format_str: str, *format_args: tvm.Ob
If the condition is true, then the operator does nothing.
If the condition is false, then the operator raises an assertion error.
- Arguments after the first value serve as format arguments for the error message;
- the last argument must be a format string for the error message (empty by default).
+ The second argument is the format string for the error message, followed by
+ its format arguments.
If the format string is the empty string, then the error message will simply include
a comma-separated list of the format arguments.
The condition argument is not included in the format string.
@@ -597,7 +612,7 @@ def relax_assert_op(condition: tvm.Object, format_str: str, *format_args: tvm.Ob
The assertion condition. Must be a boolean scalar.
format_str: str
- The last argument is a Python-style format string for printing the value
+ A Python-style format string for the error message.
format_args: List[tvm.Object]
Values used for formatting the string.
@@ -640,8 +655,10 @@ def relax_assert_op(condition: tvm.Object, format_str: str, *format_args: tvm.Ob
def assert_op(
condition: Expr,
- format_args: Expr | list[Expr] | None = None,
format: str | Expr = "",
+ *values: Expr,
+ ty=None,
+ span=None,
) -> Expr:
"""
Create a call to Relax's assert_op operation (`assert` is reserved in Python,
@@ -652,32 +669,23 @@ def assert_op(
condition: Expr
The assertion condition.
- format_args: Optional[Union[Expr, List[Expr]]]
- Format arguments for the error message if the condition fails.
-
format: Union[str, Expr]
- The format string or StringImm for the error message.
+ The format string or StringImm for the error message. If empty, the
+ values are rendered as a comma-separated list.
+
+ values: Expr
+ Values used to format the error message if the condition fails.
+ A tuple-valued expression is one value unless explicitly expanded.
Returns
-------
result : Expr
A Call to the Relax assert operation.
"""
- if not isinstance(condition, Expr):
- condition = tvm.relax.prim_value(condition)
-
- if format_args is None:
- format_args = []
- elif isinstance(format_args, Expr):
- format_args = [format_args]
-
- if isinstance(format, str):
- format = StringImm(format)
+ return Call("relax.assert_op", [condition, format, *values], ty=ty, span=span)
- return _ffi_api.assert_op(condition, format_args, format) # type: ignore
-
-def shape_of(expr: Expr) -> Expr:
+def shape_of(expr: Expr, *, ty=None, span=None) -> Expr:
"""Get shape of a tensor.
Parameters
@@ -690,10 +698,10 @@ def shape_of(expr: Expr) -> Expr:
result : Expr
A relax Call, which gets the shape of the input
"""
- return _ffi_api.shape_of(expr) # type: ignore # pylint: disable=no-member
+ return Call("relax.shape_of", [expr], ty=ty, span=span) # type: ignore # pylint: disable=no-member
-def size(expr: Expr) -> Expr:
+def size(expr: Expr, *, ty=None, span=None) -> Expr:
"""Get the total number of elements in a tensor.
Parameters
@@ -706,10 +714,10 @@ def size(expr: Expr) -> Expr:
result : Expr
A scalar tensor of dtype int64 containing the total number of elements.
"""
- return _ffi_api.size(expr) # type: ignore # pylint: disable=no-member
+ return Call("relax.size", [expr], ty=ty, span=span) # type: ignore # pylint: disable=no-member
-def tensor_to_shape(expr: Expr) -> Expr:
+def tensor_to_shape(expr: Expr, *, ty=None, span=None) -> Expr:
"""Convert tensor to shape expr.
Parameters
----------
@@ -720,10 +728,10 @@ def tensor_to_shape(expr: Expr) -> Expr:
result : Expr
A relax Call, which transforms the tensor values to the shape
"""
- return _ffi_api.tensor_to_shape(expr) # type: ignore # pylint: disable=no-member
+ return Call("relax.tensor_to_shape", [expr], ty=ty, span=span) # type: ignore # pylint: disable=no-member
-def shape_to_tensor(expr: Expr) -> Expr:
+def shape_to_tensor(expr: Expr, *, ty=None, span=None) -> Expr:
"""Convert shape to tensor expr.
Parameters
----------
@@ -734,14 +742,16 @@ def shape_to_tensor(expr: Expr) -> Expr:
result : Expr
A relax Call, which transforms the shape values to the tensor
"""
- return _ffi_api.shape_to_tensor(expr) # type: ignore # pylint: disable=no-member
+ return Call("relax.shape_to_tensor", [expr], ty=ty, span=span) # type: ignore # pylint: disable=no-member
def call_inplace_packed(
func: str | ExternFunc | GlobalVar,
*args: Expr,
- inplace_indices: int | list[int],
- ty_args: Type | list[Type],
+ inplace_indices: int | list[int] | None = None,
+ ty_args: Type | list[Type] | None = None,
+ ty=None,
+ span=None,
) -> Expr:
"""
Construct a call to a packed function that consumes some of its arguments "in-place"
@@ -786,27 +796,33 @@ def call_inplace_packed(
A Relax call, corresponding to
`call_pure_packed(ExternFunc(func), args, DictAttrs(kwargs), ty_args)`
"""
- if isinstance(func, ExternFunc):
- func = func.global_symbol
-
- op = ExternFunc(func)
+ op = ExternFunc(func) if isinstance(func, str) else func
args = tuple(convert_to_expr(a) for a in args)
if ty_args is None:
- raise ValueError("R.call_pure_packed is required to have type_args")
+ ty_args = []
if isinstance(ty_args, tuple): # type: ignore
ty_args = list(ty_args)
elif not isinstance(ty_args, list):
ty_args = [ty_args]
- if not isinstance(inplace_indices, list):
+ if inplace_indices is not None and not isinstance(inplace_indices, list):
inplace_indices = [inplace_indices]
- return _ffi_api.call_inplace_packed(op, args, inplace_indices, ty_args) # type: ignore # pylint: disable=no-member
+ return Call(
+ "relax.call_inplace_packed",
+ [op, *args],
+ attrs=_make_attrs("relax.attrs.CallInplacePackedAttrs", inplace_indices=inplace_indices),
+ ty_args=ty_args,
+ ty=ty,
+ span=span,
+ ) # type: ignore # pylint: disable=no-member
def call_pure_packed(
func: str | ExternFunc | GlobalVar | Op,
*args: Expr,
ty_args: Type | list[Type] | None = None,
+ ty=None,
+ span=None,
) -> Expr:
"""
Construct a call to a packed function that should be treated as pure,
@@ -840,17 +856,11 @@ def call_pure_packed(
A Relax call, corresponding to
`call_pure_packed(ExternFunc(func), args, DictAttrs(kwargs), ty_args)`
"""
- if isinstance(func, ExternFunc):
- func = func.global_symbol
-
- op = func if isinstance(func, Op) else ExternFunc(func)
+ op = ExternFunc(func) if isinstance(func, str) else func
args = tuple(convert_to_expr(a) for a in args)
if ty_args is None:
- if isinstance(op, Op) and op.same_as(Op.get("relax.call_tir_packed")):
- ty_args = []
- else:
- raise ValueError("R.call_pure_packed is required to have type_args")
+ ty_args = []
if isinstance(ty_args, tuple): # type: ignore
ty_args = list(ty_args)
@@ -864,13 +874,17 @@ def call_pure_packed(
# note: if we need attributes, we can also take them here
- return _ffi_api.call_pure_packed(op, args, None, ty_args) # type: ignore # pylint: disable=no-member
+ return Call(
+ "relax.call_pure_packed",
+ [op, *args],
+ ty_args=ty_args,
+ ty=ty,
+ span=span,
+ ) # type: ignore # pylint: disable=no-member
def invoke_pure_closure(
- closure: Expr,
- args: Expr,
- ty_args: list[Type] | Type,
+ closure: Expr, args: Expr, ty_args: list[Type] | Type, *, ty=None, span=None
) -> Call:
"""
Invoke a closure and indicate to the compiler that it is pure.
@@ -902,10 +916,16 @@ def invoke_pure_closure(
if not isinstance(ty_args, list | tuple):
ty_args = [ty_args]
- return _ffi_api.invoke_pure_closure(closure, args, ty_args) # type: ignore
+ return Call(
+ "relax.invoke_pure_closure",
+ [closure, args],
+ ty_args=ty_args,
+ ty=ty,
+ span=span,
+ ) # type: ignore
-def to_vdevice(data, dst_vdevice) -> Expr:
+def to_vdevice(data, dst_vdevice, *, ty=None, span=None) -> Expr:
"""Copy data to the destination device. This
operator helps data transferring between difference devices for
heterogeneous execution.
@@ -923,7 +943,13 @@ def to_vdevice(data, dst_vdevice) -> Expr:
result : Expr
The copied result.
"""
- return _ffi_api.to_vdevice(data, dst_vdevice) # type: ignore
+ return Call(
+ "relax.to_vdevice",
+ [data],
+ attrs=_make_attrs("relax.attrs.ToVDeviceAttrs", dst_vdevice=dst_vdevice),
+ ty=ty,
+ span=span,
+ ) # type: ignore
def hint_on_device(data, dst_vdevice, memory_scope="global") -> Expr:
diff --git a/python/tvm/relax/op/binary.py b/python/tvm/relax/op/binary.py
index 4b26226276..6a57161a20 100644
--- a/python/tvm/relax/op/binary.py
+++ b/python/tvm/relax/op/binary.py
@@ -17,13 +17,14 @@
# pylint: disable=redefined-builtin, invalid-name
"""Relax binary arithmetic and comparison operators."""
+from tvm.ir import Call as _Call
+
from ..expr import Expr
-from . import _ffi_api
###################### Arithmetic operators ######################
-def add(x1: Expr, x2: Expr) -> Expr:
+def add(x1: Expr, x2: Expr, *, ty_args=None, ty=None, span=None) -> Expr:
"""Addition with numpy-style broadcasting.
Parameters
@@ -47,10 +48,10 @@ def add(x1: Expr, x2: Expr) -> Expr:
b = relax.Var("b", relax.TensorType(shape=(2, 1), dtype="float32"))
c = bb.normalize(relax.op.add(a, b)) # c has TensorType(shape=(2, 3), dtype="float32")
"""
- return _ffi_api.add(x1, x2) # type: ignore
+ return _Call("relax.add", [x1, x2], ty_args=ty_args, ty=ty, span=span) # type: ignore
-def divide(x1: Expr, x2: Expr) -> Expr:
+def divide(x1: Expr, x2: Expr, *, ty_args=None, ty=None, span=None) -> Expr:
"""Division with numpy-style broadcasting.
Parameters
@@ -65,10 +66,10 @@ def divide(x1: Expr, x2: Expr) -> Expr:
result : relax.Expr
The computed result.
"""
- return _ffi_api.divide(x1, x2) # type: ignore
+ return _Call("relax.divide", [x1, x2], ty_args=ty_args, ty=ty, span=span) # type: ignore
-def floor_divide(x1: Expr, x2: Expr) -> Expr:
+def floor_divide(x1: Expr, x2: Expr, *, ty_args=None, ty=None, span=None) -> Expr:
"""Floor division with numpy-style broadcasting.
Parameters
@@ -83,10 +84,10 @@ def floor_divide(x1: Expr, x2: Expr) -> Expr:
result : relax.Expr
The computed result.
"""
- return _ffi_api.floor_divide(x1, x2) # type: ignore
+ return _Call("relax.floor_divide", [x1, x2], ty_args=ty_args, ty=ty, span=span) # type: ignore
-def log_add_exp(x1: Expr, x2: Expr) -> Expr:
+def log_add_exp(x1: Expr, x2: Expr, *, ty_args=None, ty=None, span=None) -> Expr:
"""
Compute the log of the sum of exponentials of the inputs, element-wise.
@@ -102,10 +103,10 @@ def log_add_exp(x1: Expr, x2: Expr) -> Expr:
Expr
The element-wise log-sum-exp of `x1` and `x2`.
"""
- return _ffi_api.log_add_exp(x1, x2)
+ return _Call("relax.log_add_exp", [x1, x2], ty_args=ty_args, ty=ty, span=span)
-def multiply(x1: Expr, x2: Expr) -> Expr:
+def multiply(x1: Expr, x2: Expr, *, ty_args=None, ty=None, span=None) -> Expr:
"""Multiplication with numpy-style broadcasting.
Parameters
@@ -120,10 +121,10 @@ def multiply(x1: Expr, x2: Expr) -> Expr:
result : Expr
The computed result.
"""
- return _ffi_api.multiply(x1, x2) # type: ignore
+ return _Call("relax.multiply", [x1, x2], ty_args=ty_args, ty=ty, span=span) # type: ignore
-def power(x1: Expr, x2: Expr):
+def power(x1: Expr, x2: Expr, *, ty_args=None, ty=None, span=None):
"""Power with numpy-style broadcasting.
Parameters
@@ -138,10 +139,10 @@ def power(x1: Expr, x2: Expr):
result : relax.Expr
The computed result.
"""
- return _ffi_api.power(x1, x2) # type: ignore
+ return _Call("relax.power", [x1, x2], ty_args=ty_args, ty=ty, span=span) # type: ignore
-def atan2(x1: Expr, x2: Expr) -> Expr:
+def atan2(x1: Expr, x2: Expr, *, ty_args=None, ty=None, span=None) -> Expr:
"""Atan2 with numpy-style broadcasting.
Parameters
@@ -156,10 +157,10 @@ def atan2(x1: Expr, x2: Expr) -> Expr:
result : relax.Expr
The computed result.
"""
- return _ffi_api.atan2(x1, x2) # type: ignore
+ return _Call("relax.atan2", [x1, x2], ty_args=ty_args, ty=ty, span=span) # type: ignore
-def subtract(x1: Expr, x2: Expr) -> Expr:
+def subtract(x1: Expr, x2: Expr, *, ty_args=None, ty=None, span=None) -> Expr:
"""Subtraction with numpy-style broadcasting.
Parameters
@@ -174,10 +175,10 @@ def subtract(x1: Expr, x2: Expr) -> Expr:
result : relax.Expr
The computed result.
"""
- return _ffi_api.subtract(x1, x2) # type: ignore
+ return _Call("relax.subtract", [x1, x2], ty_args=ty_args, ty=ty, span=span) # type: ignore
-def mod(x1: Expr, x2: Expr) -> Expr:
+def mod(x1: Expr, x2: Expr, *, ty_args=None, ty=None, span=None) -> Expr:
"""Modulo with numpy-style broadcasting.
Parameters
@@ -187,10 +188,10 @@ def mod(x1: Expr, x2: Expr) -> Expr:
x2 : Expr
The second input tensor.
"""
- return _ffi_api.mod(x1, x2) # type: ignore
+ return _Call("relax.mod", [x1, x2], ty_args=ty_args, ty=ty, span=span) # type: ignore
-def floor_mod(x1: Expr, x2: Expr) -> Expr:
+def floor_mod(x1: Expr, x2: Expr, *, ty_args=None, ty=None, span=None) -> Expr:
"""Floor modulo with numpy-style broadcasting.
Parameters
@@ -200,13 +201,13 @@ def floor_mod(x1: Expr, x2: Expr) -> Expr:
x2 : Expr
The second input tensor.
"""
- return _ffi_api.floor_mod(x1, x2) # type: ignore
+ return _Call("relax.floor_mod", [x1, x2], ty_args=ty_args, ty=ty, span=span) # type: ignore
###################### Comparison operators ######################
-def equal(x1: Expr, x2: Expr) -> Expr:
+def equal(x1: Expr, x2: Expr, *, ty_args=None, ty=None, span=None) -> Expr:
"""Broadcasted element-wise test for (lhs == rhs).
Parameters
@@ -221,10 +222,10 @@ def equal(x1: Expr, x2: Expr) -> Expr:
result : relax.Expr
The computed result.
"""
- return _ffi_api.equal(x1, x2) # type: ignore
+ return _Call("relax.equal", [x1, x2], ty_args=ty_args, ty=ty, span=span) # type: ignore
-def greater(x1: Expr, x2: Expr) -> Expr:
+def greater(x1: Expr, x2: Expr, *, ty_args=None, ty=None, span=None) -> Expr:
"""Broadcasted element-wise test for (lhs > rhs).
Parameters
@@ -239,10 +240,10 @@ def greater(x1: Expr, x2: Expr) -> Expr:
result : relax.Expr
The computed result.
"""
- return _ffi_api.greater(x1, x2) # type: ignore
+ return _Call("relax.greater", [x1, x2], ty_args=ty_args, ty=ty, span=span) # type: ignore
-def greater_equal(x1: Expr, x2: Expr) -> Expr:
+def greater_equal(x1: Expr, x2: Expr, *, ty_args=None, ty=None, span=None) -> Expr:
"""Broadcasted element-wise test for (lhs >= rhs).
Parameters
@@ -257,10 +258,10 @@ def greater_equal(x1: Expr, x2: Expr) -> Expr:
result : relax.Expr
The computed result.
"""
- return _ffi_api.greater_equal(x1, x2) # type: ignore
+ return _Call("relax.greater_equal", [x1, x2], ty_args=ty_args, ty=ty, span=span) # type: ignore
-def less(x1: Expr, x2: Expr) -> Expr:
+def less(x1: Expr, x2: Expr, *, ty_args=None, ty=None, span=None) -> Expr:
"""Broadcasted element-wise test for (lhs < rhs).
Parameters
@@ -275,10 +276,10 @@ def less(x1: Expr, x2: Expr) -> Expr:
result : relax.Expr
The computed result.
"""
- return _ffi_api.less(x1, x2) # type: ignore
+ return _Call("relax.less", [x1, x2], ty_args=ty_args, ty=ty, span=span) # type: ignore
-def less_equal(x1: Expr, x2: Expr) -> Expr:
+def less_equal(x1: Expr, x2: Expr, *, ty_args=None, ty=None, span=None) -> Expr:
"""Broadcasted element-wise test for (lhs <= rhs).
Parameters
@@ -293,10 +294,10 @@ def less_equal(x1: Expr, x2: Expr) -> Expr:
result : relax.Expr
The computed result.
"""
- return _ffi_api.less_equal(x1, x2) # type: ignore
+ return _Call("relax.less_equal", [x1, x2], ty_args=ty_args, ty=ty, span=span) # type: ignore
-def not_equal(x1: Expr, x2: Expr) -> Expr:
+def not_equal(x1: Expr, x2: Expr, *, ty_args=None, ty=None, span=None) -> Expr:
"""Broadcasted element-wise test for (lhs != rhs).
Parameters
@@ -311,10 +312,10 @@ def not_equal(x1: Expr, x2: Expr) -> Expr:
result : relax.Expr
The computed result.
"""
- return _ffi_api.not_equal(x1, x2) # type: ignore
+ return _Call("relax.not_equal", [x1, x2], ty_args=ty_args, ty=ty, span=span) # type: ignore
-def maximum(x1: Expr, x2: Expr) -> Expr:
+def maximum(x1: Expr, x2: Expr, *, ty_args=None, ty=None, span=None) -> Expr:
"""Element-wise maximum
Parameters
@@ -329,10 +330,10 @@ def maximum(x1: Expr, x2: Expr) -> Expr:
result : relax.Expr
The computed result.
"""
- return _ffi_api.maximum(x1, x2)
+ return _Call("relax.maximum", [x1, x2], ty_args=ty_args, ty=ty, span=span)
-def minimum(x1: Expr, x2: Expr) -> Expr:
+def minimum(x1: Expr, x2: Expr, *, ty_args=None, ty=None, span=None) -> Expr:
"""Element-wise minimum
Parameters
@@ -347,13 +348,13 @@ def minimum(x1: Expr, x2: Expr) -> Expr:
result : relax.Expr
The computed result.
"""
- return _ffi_api.minimum(x1, x2)
+ return _Call("relax.minimum", [x1, x2], ty_args=ty_args, ty=ty, span=span)
###################### Logical operators ######################
-def logical_and(x1: Expr, x2: Expr) -> Expr:
+def logical_and(x1: Expr, x2: Expr, *, ty_args=None, ty=None, span=None) -> Expr:
"""Logical AND
Parameters
----------
@@ -366,10 +367,10 @@ def logical_and(x1: Expr, x2: Expr) -> Expr:
result : relax.Expr
The computed result.
"""
- return _ffi_api.logical_and(x1, x2)
+ return _Call("relax.logical_and", [x1, x2], ty_args=ty_args, ty=ty, span=span)
-def logical_or(x1: Expr, x2: Expr) -> Expr:
+def logical_or(x1: Expr, x2: Expr, *, ty_args=None, ty=None, span=None) -> Expr:
"""Logical OR
Parameters
----------
@@ -382,10 +383,10 @@ def logical_or(x1: Expr, x2: Expr) -> Expr:
result : relax.Expr
The computed result.
"""
- return _ffi_api.logical_or(x1, x2)
+ return _Call("relax.logical_or", [x1, x2], ty_args=ty_args, ty=ty, span=span)
-def logical_xor(x1: Expr, x2: Expr) -> Expr:
+def logical_xor(x1: Expr, x2: Expr, *, ty_args=None, ty=None, span=None) -> Expr:
"""Logical XOR
Parameters
----------
@@ -398,13 +399,13 @@ def logical_xor(x1: Expr, x2: Expr) -> Expr:
result : relax.Expr
The computed result.
"""
- return _ffi_api.logical_xor(x1, x2)
+ return _Call("relax.logical_xor", [x1, x2], ty_args=ty_args, ty=ty, span=span)
###################### Bitwise operators ######################
-def bitwise_and(x1: Expr, x2: Expr) -> Expr:
+def bitwise_and(x1: Expr, x2: Expr, *, ty_args=None, ty=None, span=None) -> Expr:
"""Bitwise AND
Parameters
----------
@@ -417,10 +418,10 @@ def bitwise_and(x1: Expr, x2: Expr) -> Expr:
result : relax.Expr
The computed result.
"""
- return _ffi_api.bitwise_and(x1, x2)
+ return _Call("relax.bitwise_and", [x1, x2], ty_args=ty_args, ty=ty, span=span)
-def bitwise_or(x1: Expr, x2: Expr) -> Expr:
+def bitwise_or(x1: Expr, x2: Expr, *, ty_args=None, ty=None, span=None) -> Expr:
"""Bitwise OR
Parameters
----------
@@ -433,10 +434,10 @@ def bitwise_or(x1: Expr, x2: Expr) -> Expr:
result : relax.Expr
The computed result.
"""
- return _ffi_api.bitwise_or(x1, x2)
+ return _Call("relax.bitwise_or", [x1, x2], ty_args=ty_args, ty=ty, span=span)
-def bitwise_xor(x1: Expr, x2: Expr) -> Expr:
+def bitwise_xor(x1: Expr, x2: Expr, *, ty_args=None, ty=None, span=None) -> Expr:
"""Bitwise XOR
Parameters
----------
@@ -449,10 +450,10 @@ def bitwise_xor(x1: Expr, x2: Expr) -> Expr:
result : relax.Expr
The computed result.
"""
- return _ffi_api.bitwise_xor(x1, x2)
+ return _Call("relax.bitwise_xor", [x1, x2], ty_args=ty_args, ty=ty, span=span)
-def left_shift(x1: Expr, x2: Expr) -> Expr:
+def left_shift(x1: Expr, x2: Expr, *, ty_args=None, ty=None, span=None) -> Expr:
"""Bitwise Shift Left
Parameters
----------
@@ -465,10 +466,10 @@ def left_shift(x1: Expr, x2: Expr) -> Expr:
result : relax.Expr
The computed result.
"""
- return _ffi_api.left_shift(x1, x2)
+ return _Call("relax.left_shift", [x1, x2], ty_args=ty_args, ty=ty, span=span)
-def right_shift(x1: Expr, x2: Expr) -> Expr:
+def right_shift(x1: Expr, x2: Expr, *, ty_args=None, ty=None, span=None) -> Expr:
"""Bitwise Shift Right
Parameters
----------
@@ -481,4 +482,4 @@ def right_shift(x1: Expr, x2: Expr) -> Expr:
result : relax.Expr
The computed result.
"""
- return _ffi_api.right_shift(x1, x2)
+ return _Call("relax.right_shift", [x1, x2], ty_args=ty_args, ty=ty, span=span)
diff --git a/python/tvm/relax/op/builtin/builtin.py b/python/tvm/relax/op/builtin/builtin.py
index 5592e0d402..a905b87999 100644
--- a/python/tvm/relax/op/builtin/builtin.py
+++ b/python/tvm/relax/op/builtin/builtin.py
@@ -16,10 +16,10 @@
"""The builtin Relax operators."""
from tvm.ir import Call, DataTypeImm, GenericConst, StringImm
+from tvm.ir import Call as _Call
from ...expr import Expr, prim_value
from ...utils import convert_to_expr
-from . import _ffi_api
def alloc_tensor(
@@ -27,6 +27,10 @@ def alloc_tensor(
dtype: str | Expr,
runtime_device_index: int | Expr,
storage_scope: str | Expr = "global",
+ *,
+ ty_args=None,
+ ty=None,
+ span=None,
) -> Call:
"""Construct a Call to allocate a tensor with specific shape, dtype, runtime_device_index.
@@ -64,10 +68,16 @@ def alloc_tensor(
f"but {storage_scope} is got."
)
- return _ffi_api.alloc_tensor(shape, dtype, runtime_device_index, storage_scope) # type: ignore
+ return _Call(
+ "relax.builtin.alloc_tensor",
+ [shape, dtype, runtime_device_index, storage_scope],
+ ty_args=ty_args,
+ ty=ty,
+ span=span,
+ ) # type: ignore
-def stop_lift_params(x: Expr) -> Expr:
+def stop_lift_params(x: Expr, *, ty_args=None, ty=None, span=None) -> Expr:
"""
An indicator that the consumers of input tensor should not be
lifted to transform_params function
@@ -82,4 +92,10 @@ def stop_lift_params(x: Expr) -> Expr:
result : relax.Expr
The result tensor that is the same as input tensor
"""
- return _ffi_api.stop_lift_params(x) # type: ignore
+ return _Call(
+ "relax.builtin.stop_lift_params",
+ [x],
+ ty_args=ty_args,
+ ty=ty,
+ span=span,
+ ) # type: ignore
diff --git a/python/tvm/relax/op/ccl/ccl.py b/python/tvm/relax/op/ccl/ccl.py
index e49614f4de..da7b3fddee 100644
--- a/python/tvm/relax/op/ccl/ccl.py
+++ b/python/tvm/relax/op/ccl/ccl.py
@@ -16,11 +16,13 @@
# under the License.
"""Relax Collective Communications Library (CCL) operators"""
+from tvm.ir import Call as _Call
+from tvm.ir.attrs import make_node as _make_attrs
+
from ...expr import Expr
-from . import _ffi_api
-def allreduce(x, op_type: str = "sum", in_group: bool = True): # pylint: disable=invalid-name
+def allreduce(x, op_type: str = "sum", in_group: bool = True, *, ty=None, span=None): # pylint: disable=invalid-name
"""Allreduce operator
Parameters
@@ -45,10 +47,16 @@ def allreduce(x, op_type: str = "sum", in_group: bool = True): # pylint: disabl
"Allreduce only supports limited reduction operations, "
f"including {supported_op_types}, but got {op_type}."
)
- return _ffi_api.allreduce(x, op_type, in_group) # type: ignore # pylint: disable=no-member
+ return _Call(
+ "relax.ccl.allreduce",
+ [x],
+ attrs=_make_attrs("relax.attrs.AllReduceAttrs", op_type=op_type, in_group=in_group),
+ ty=ty,
+ span=span,
+ ) # type: ignore # pylint: disable=no-member
-def allgather(x, num_workers: int, in_group: bool = True): # pylint: disable=invalid-name
+def allgather(x, num_workers: int, in_group: bool = True, *, ty=None, span=None): # pylint: disable=invalid-name
"""AllGather operator
Parameters
@@ -67,10 +75,16 @@ def allgather(x, num_workers: int, in_group: bool = True): # pylint: disable=in
result : relax.Expr
The result of allgather.
"""
- return _ffi_api.allgather(x, num_workers, in_group) # type: ignore # pylint: disable=no-member
+ return _Call(
+ "relax.ccl.allgather",
+ [x],
+ attrs=_make_attrs("relax.attrs.AllGatherAttrs", num_workers=num_workers, in_group=in_group),
+ ty=ty,
+ span=span,
+ ) # type: ignore # pylint: disable=no-member
-def broadcast_from_worker0(x: Expr) -> Expr:
+def broadcast_from_worker0(x: Expr, *, ty=None, span=None) -> Expr:
"""Broadcast data from worker-0 to all other workers.
Parameters
@@ -83,10 +97,15 @@ def broadcast_from_worker0(x: Expr) -> Expr:
result : relax.Expr
The same tensor, which has been broadcast to all other workers.
"""
- return _ffi_api.broadcast_from_worker0(x)
+ return _Call(
+ "relax.ccl.broadcast_from_worker0",
+ [x],
+ ty=ty,
+ span=span,
+ )
-def scatter_from_worker0(x: Expr, num_workers: int, axis: int = 0) -> Expr:
+def scatter_from_worker0(x: Expr, num_workers: int, axis: int = 0, *, ty=None, span=None) -> Expr:
"""Perform a scatter operation from worker-0, chunking the given buffer into equal parts.
Parameters
@@ -105,4 +124,10 @@ def scatter_from_worker0(x: Expr, num_workers: int, axis: int = 0) -> Expr:
result : relax.Expr
Chunked Tensor received by different workers.
"""
- return _ffi_api.scatter_from_worker0(x, num_workers, axis)
+ return _Call(
+ "relax.ccl.scatter_from_worker0",
+ [x],
+ attrs=_make_attrs("relax.attrs.ScatterCollectiveAttrs", num_workers=num_workers, axis=axis),
+ ty=ty,
+ span=span,
+ )
diff --git a/python/tvm/relax/op/create.py b/python/tvm/relax/op/create.py
index 5188140222..1e602c5a36 100644
--- a/python/tvm/relax/op/create.py
+++ b/python/tvm/relax/op/create.py
@@ -17,10 +17,11 @@
"""Creation operators."""
from tvm import DataType, DataTypeCode
+from tvm.ir import Call as _Call
from tvm.ir import PrimType, is_prim_expr
+from tvm.ir.attrs import make_node as _make_attrs
from ..expr import Expr, ShapeExpr, prim_value
-from . import _ffi_api
PrimExprLike = int | Expr
@@ -41,6 +42,9 @@ def full(
shape: tuple[PrimExprLike] | Expr,
fill_value: Expr,
dtype: str | DataType | None = None,
+ *,
+ ty=None,
+ span=None,
) -> Expr:
"""Fill array with scalar value.
@@ -62,10 +66,23 @@ def full(
The result tensor.
"""
shape = _normalize_shape(shape)
- return _ffi_api.full(shape, fill_value, _raw_dtype(dtype)) # type: ignore
+ return _Call(
+ "relax.full",
+ [shape, fill_value],
+ attrs=_make_attrs("relax.attrs.InitAttrs", dtype=_raw_dtype(dtype)),
+ ty=ty,
+ span=span,
+ ) # type: ignore
-def full_like(x: Expr, fill_value: Expr, dtype: str | DataType | None = None) -> Expr:
+def full_like(
+ x: Expr,
+ fill_value: Expr,
+ dtype: str | DataType | None = None,
+ *,
+ ty=None,
+ span=None,
+) -> Expr:
"""Construct a tensor such that
- its shape is the same as the input data tensor's shape,
- its value is filled with the input scalar fill value.
@@ -88,10 +105,22 @@ def full_like(x: Expr, fill_value: Expr, dtype: str | DataType | None = None) ->
result : relax.Expr
The result tensor.
"""
- return _ffi_api.full_like(x, fill_value, _raw_dtype(dtype)) # type: ignore
+ return _Call(
+ "relax.full_like",
+ [x, fill_value],
+ attrs=_make_attrs("relax.attrs.InitAttrs", dtype=_raw_dtype(dtype)),
+ ty=ty,
+ span=span,
+ ) # type: ignore
-def ones(shape: tuple[PrimExprLike] | Expr, dtype: str | DataType) -> Expr:
+def ones(
+ shape: tuple[PrimExprLike] | Expr,
+ dtype: str | DataType,
+ *,
+ ty=None,
+ span=None,
+) -> Expr:
"""Construct a tensor of all ones, with the input shape and dtype.
Parameters
@@ -108,10 +137,16 @@ def ones(shape: tuple[PrimExprLike] | Expr, dtype: str | DataType) -> Expr:
The result tensor.
"""
shape = _normalize_shape(shape)
- return _ffi_api.ones(shape, _raw_dtype(dtype)) # type: ignore
+ return _Call(
+ "relax.ones",
+ [shape],
+ attrs=_make_attrs("relax.attrs.InitAttrs", dtype=_raw_dtype(dtype)),
+ ty=ty,
+ span=span,
+ ) # type: ignore
-def ones_like(x: Expr, dtype: str | DataType | None = None) -> Expr:
+def ones_like(x: Expr, dtype: str | DataType | None = None, *, ty=None, span=None) -> Expr:
"""Construct a tensor with all ones, with shape of the input tensor shape.
Parameters
@@ -129,10 +164,22 @@ def ones_like(x: Expr, dtype: str | DataType | None = None) -> Expr:
result : relax.Expr
The result tensor.
"""
- return _ffi_api.ones_like(x, _raw_dtype(dtype)) # type: ignore
+ return _Call(
+ "relax.ones_like",
+ [x],
+ attrs=_make_attrs("relax.attrs.InitAttrs", dtype=_raw_dtype(dtype)),
+ ty=ty,
+ span=span,
+ ) # type: ignore
-def zeros(shape: tuple[PrimExprLike] | Expr, dtype: str | DataType) -> Expr:
+def zeros(
+ shape: tuple[PrimExprLike] | Expr,
+ dtype: str | DataType,
+ *,
+ ty=None,
+ span=None,
+) -> Expr:
"""Construct a tensor of all zeros, with the input shape and dtype.
Parameters
@@ -149,10 +196,16 @@ def zeros(shape: tuple[PrimExprLike] | Expr, dtype: str | DataType) -> Expr:
The result tensor.
"""
shape = _normalize_shape(shape)
- return _ffi_api.zeros(shape, _raw_dtype(dtype)) # type: ignore
+ return _Call(
+ "relax.zeros",
+ [shape],
+ attrs=_make_attrs("relax.attrs.InitAttrs", dtype=_raw_dtype(dtype)),
+ ty=ty,
+ span=span,
+ ) # type: ignore
-def zeros_like(x: Expr, dtype: str | DataType | None = None) -> Expr:
+def zeros_like(x: Expr, dtype: str | DataType | None = None, *, ty=None, span=None) -> Expr:
"""Construct a tensor with all zeros, with shape of the input tensor shape.
Parameters
@@ -170,7 +223,13 @@ def zeros_like(x: Expr, dtype: str | DataType | None = None) -> Expr:
result : relax.Expr
The result tensor.
"""
- return _ffi_api.zeros_like(x, _raw_dtype(dtype)) # type: ignore
+ return _Call(
+ "relax.zeros_like",
+ [x],
+ attrs=_make_attrs("relax.attrs.InitAttrs", dtype=_raw_dtype(dtype)),
+ ty=ty,
+ span=span,
+ ) # type: ignore
def eye(
@@ -178,6 +237,9 @@ def eye(
m: PrimExprLike | None = None,
k: PrimExprLike = 0,
dtype: str | DataType = "float32",
+ *,
+ ty=None,
+ span=None,
) -> Expr:
"""Construct a 2-D tensor with ones on the diagonal and zeros elsewhere.
@@ -206,13 +268,22 @@ def eye(
n = prim_value(n)
m = prim_value(m)
k = prim_value(k)
- return _ffi_api.eye(n, m, k, _raw_dtype(dtype)) # type: ignore
+ return _Call(
+ "relax.eye",
+ [n, m, k],
+ attrs=_make_attrs("relax.attrs.InitAttrs", dtype=_raw_dtype(dtype)),
+ ty=ty,
+ span=span,
+ ) # type: ignore
def eye_like(
x: Expr,
k: PrimExprLike = 0,
dtype: str | DataType | None = None,
+ *,
+ ty=None,
+ span=None,
) -> Expr:
"""Return a 2-D tensor with ones on the diagonal and zeros elsewhere,
with the same shape as the input tensor.
@@ -238,7 +309,13 @@ def eye_like(
The result tensor.
"""
k = prim_value(k)
- return _ffi_api.eye_like(x, k, _raw_dtype(dtype)) # type: ignore
+ return _Call(
+ "relax.eye_like",
+ [x, k],
+ attrs=_make_attrs("relax.attrs.InitAttrs", dtype=_raw_dtype(dtype)),
+ ty=ty,
+ span=span,
+ ) # type: ignore
def arange(
@@ -246,6 +323,9 @@ def arange(
end: PrimExprLike | None = None,
step: PrimExprLike = 1,
dtype: str | DataType | PrimType | None = None,
+ *,
+ ty=None,
+ span=None,
) -> Expr:
"""Construct a tensor with evenly spaced elements.
@@ -288,10 +368,16 @@ def arange(
start = prim_value(start)
end = prim_value(end)
step = prim_value(step)
- return _ffi_api.arange(start, end, step, _raw_dtype(dtype)) # type: ignore
+ return _Call(
+ "relax.arange",
+ [start, end, step],
+ attrs=_make_attrs("relax.attrs.InitAttrs", dtype=_raw_dtype(dtype)),
+ ty=ty,
+ span=span,
+ ) # type: ignore
-def hamming_window(window_size, periodic, alpha, beta, dtype):
+def hamming_window(window_size, periodic, alpha, beta, dtype, *, ty=None, span=None):
"""Hamming window function.
Parameters
@@ -323,10 +409,16 @@ def hamming_window(window_size, periodic, alpha, beta, dtype):
if not is_prim_expr(beta):
beta = prim_value(beta)
- return _ffi_api.hamming_window(window_size, periodic, alpha, beta, dtype)
+ return _Call(
+ "relax.hamming_window",
+ [window_size, periodic, alpha, beta],
+ attrs=_make_attrs("relax.attrs.InitAttrs", dtype=dtype),
+ ty=ty,
+ span=span,
+ )
-def tril(x: Expr, k: int | Expr = 0) -> Expr:
+def tril(x: Expr, k: int | Expr = 0, *, ty=None, span=None) -> Expr:
"""Return the lower triangular part of a matrix or a batch of matrices.
Parameters
@@ -349,10 +441,10 @@ def tril(x: Expr, k: int | Expr = 0) -> Expr:
if not is_prim_expr(k):
k = prim_value(k)
- return _ffi_api.tril(x, k) # type: ignore
+ return _Call("relax.tril", [x, k], ty=ty, span=span) # type: ignore
-def triu(x: Expr, k: int | Expr = 0) -> Expr:
+def triu(x: Expr, k: int | Expr = 0, *, ty=None, span=None) -> Expr:
"""Return the upper triangular part of a matrix or a batch of matrices.
Parameters
@@ -375,4 +467,4 @@ def triu(x: Expr, k: int | Expr = 0) -> Expr:
if not is_prim_expr(k):
k = prim_value(k)
- return _ffi_api.triu(x, k) # type: ignore
+ return _Call("relax.triu", [x, k], ty=ty, span=span) # type: ignore
diff --git a/python/tvm/relax/op/datatype.py b/python/tvm/relax/op/datatype.py
index bcba5b4dd3..9e1c7186ca 100644
--- a/python/tvm/relax/op/datatype.py
+++ b/python/tvm/relax/op/datatype.py
@@ -17,17 +17,18 @@
"""Datatype operators."""
from tvm import DataType
+from tvm.ir import Call as _Call
from tvm.ir import PrimType
+from tvm.ir.attrs import make_node as _make_attrs
from ..expr import Expr
-from . import _ffi_api
def _raw_dtype(dtype):
return dtype.dtype if isinstance(dtype, PrimType) else dtype
-def astype(x: Expr, dtype: str | DataType | PrimType) -> Expr:
+def astype(x: Expr, dtype: str | DataType | PrimType, *, ty=None, span=None) -> Expr:
"""Cast input tensor to the given data type.
Parameters
@@ -43,10 +44,22 @@ def astype(x: Expr, dtype: str | DataType | PrimType) -> Expr:
result : relax.Expr
The casted result.
"""
- return _ffi_api.astype(x, _raw_dtype(dtype)) # type: ignore
+ return _Call(
+ "relax.astype",
+ [x],
+ attrs=_make_attrs("relax.attrs.AstypeAttrs", dtype=_raw_dtype(dtype)),
+ ty=ty,
+ span=span,
+ ) # type: ignore
-def wrap_param(data: Expr, dtype: str | DataType | PrimType = "float32") -> Expr:
+def wrap_param(
+ data: Expr,
+ dtype: str | DataType | PrimType = "float32",
+ *,
+ ty=None,
+ span=None,
+) -> Expr:
"""Cast input tensor which is model param to data type if the dtype of the input data is not
the same as the given dtype.
Parameters
@@ -60,4 +73,10 @@ def wrap_param(data: Expr, dtype: str | DataType | PrimType = "float32") -> Expr
result : relax.Expr
The casted result.
"""
- return _ffi_api.wrap_param(data, _raw_dtype(dtype)) # type: ignore
+ return _Call(
+ "relax.wrap_param",
+ [data],
+ attrs=_make_attrs("relax.attrs.WrapParamAttrs", dtype=_raw_dtype(dtype)),
+ ty=ty,
+ span=span,
+ ) # type: ignore
diff --git a/python/tvm/relax/op/distributed/distributed.py b/python/tvm/relax/op/distributed/distributed.py
index f2d0595ad3..faebc735c6 100644
--- a/python/tvm/relax/op/distributed/distributed.py
+++ b/python/tvm/relax/op/distributed/distributed.py
@@ -17,16 +17,25 @@
# pylint: disable=redefined-builtin
"""Operators for distributed Relax."""
+from tvm.error import InternalError as _InternalError
from tvm.ir import Call
+from tvm.ir.attrs import make_node
+from tvm.relax import ShapeExpr, TupleType
from tvm.relax.distributed import DeviceMesh, DTensorType, Placement
from ...expr import Expr, GlobalVar
from ...expr import Tuple as RxTuple
from ...utils import convert_to_expr
-from . import _ffi_api
-def annotate_sharding(input: Expr, device_mesh: DeviceMesh, placement: Placement) -> Expr:
+def annotate_sharding(
+ input: Expr,
+ device_mesh: DeviceMesh,
+ placement: Placement,
+ *,
+ ty=None,
+ span=None,
+) -> Expr:
"""Annotate sharding plan for tensor
Parameters
@@ -43,10 +52,25 @@ def annotate_sharding(input: Expr, device_mesh: DeviceMesh, placement: Placement
result : relax.Expr
The tensor unmodified.
"""
- return _ffi_api.annotate_sharding(input, device_mesh, placement) # type: ignore
-
-
-def redistribute(input: Expr, device_mesh: DeviceMesh, placement: Placement) -> Expr:
+ return Call(
+ "relax.dist.annotate_sharding",
+ [input],
+ attrs=make_node(
+ "relax.attrs.DistributionAttrs", device_mesh=device_mesh, placement=placement
+ ),
+ ty=ty,
+ span=span,
+ ) # type: ignore
+
+
+def redistribute(
+ input: Expr,
+ device_mesh: DeviceMesh,
+ placement: Placement,
+ *,
+ ty=None,
+ span=None,
+) -> Expr:
"""Redistribute tensor
Parameters
@@ -62,13 +86,25 @@ def redistribute(input: Expr, device_mesh: DeviceMesh, placement: Placement) ->
result : relax.Expr
The tensor after redistribution.
"""
- return _ffi_api.redistribute(input, device_mesh, placement) # type: ignore
+ return Call(
+ "relax.dist.redistribute",
+ [input],
+ attrs=make_node(
+ "relax.attrs.DistributionAttrs", device_mesh=device_mesh, placement=placement
+ ),
+ ty=ty,
+ span=span,
+ ) # type: ignore
def call_tir_local_view(
gvar: GlobalVar,
args: Expr,
- out_ty: DTensorType | list[DTensorType],
+ out_ty: DTensorType | list[DTensorType] | None = None,
+ *,
+ ty_args=None,
+ ty=None,
+ span=None,
) -> Call:
"""
Call a tirx.function and return the output. The function should be a worker-local function
@@ -99,13 +135,27 @@ def call_tir_local_view(
elif isinstance(args, Expr) and not isinstance(args, RxTuple): # type: ignore
args = RxTuple((args,))
- if not isinstance(out_ty, list):
- out_ty = [out_ty]
-
- return _ffi_api.call_tir_local_view(gvar, args, out_ty) # type: ignore
-
-
-def redistribute_replica_to_shard(input: Expr, num_workers: int, axis: int) -> Expr:
+ if out_ty is not None:
+ if ty_args is not None:
+ raise TypeError("Specify either out_ty or ty_args")
+ if not isinstance(out_ty, list):
+ out_ty = [out_ty]
+ for output in out_ty:
+ if not isinstance(output.tensor_ty.shape, ShapeExpr):
+ raise _InternalError("out_ty must have a defined ShapeExpr shape")
+ ty_args = [out_ty[0] if len(out_ty) == 1 else TupleType(out_ty)]
+ return Call(
+ "relax.dist.call_tir_local_view",
+ [gvar, args],
+ ty_args=ty_args,
+ ty=ty,
+ span=span,
+ )
+
+
+def redistribute_replica_to_shard(
+ input: Expr, num_workers: int, axis: int, *, ty=None, span=None
+) -> Expr:
"""Slice tensor into several parts along one axis,
and each worker takes one part.
input.ty.shape[axis] % num_workers == 0 is required.
@@ -128,4 +178,10 @@ def redistribute_replica_to_shard(input: Expr, num_workers: int, axis: int) -> E
result : relax.Expr
Sliced Tensor kept by each device.
"""
- return _ffi_api.redistribute_replica_to_shard(input, num_workers, axis)
+ return Call(
+ "relax.dist.redistribute_replica_to_shard",
+ [input],
+ attrs=make_node("relax.attrs.ScatterCollectiveAttrs", num_workers=num_workers, axis=axis),
+ ty=ty,
+ span=span,
+ )
diff --git a/python/tvm/relax/op/grad/grad.py b/python/tvm/relax/op/grad/grad.py
index 86c0792b54..90cebab85e 100644
--- a/python/tvm/relax/op/grad/grad.py
+++ b/python/tvm/relax/op/grad/grad.py
@@ -22,11 +22,13 @@ gradients for forward operators. The ty inference for grad operators just return
ty of the input.
"""
+from tvm.ir import Call as _Call
+from tvm.ir.attrs import make_node as _make_attrs
+
from ...expr import Expr
-from . import _ffi_api
-def no_grad(input: Expr) -> Expr:
+def no_grad(input: Expr, *, ty=None, span=None) -> Expr:
"""No gradient dummy operator w.r.t. the input.
Parameters
@@ -39,10 +41,10 @@ def no_grad(input: Expr) -> Expr:
result : relax.Expr
The no-gradient representation w.r.t. input.
"""
- return _ffi_api.no_grad(input) # type: ignore
+ return _Call("relax.grad.no_grad", [input], ty=ty, span=span) # type: ignore
-def start_checkpoint(input: Expr) -> Expr:
+def start_checkpoint(input: Expr, *, ty=None, span=None) -> Expr:
"""Mark the start of the checkpoint stage. The computation between start_checkpoint and
end_checkpoint will be marked as the checkpoint stage.
@@ -76,10 +78,15 @@ def start_checkpoint(input: Expr) -> Expr:
result : relax.Expr
The same tensor as the input.
"""
- return _ffi_api.start_checkpoint(input) # type: ignore
+ return _Call(
+ "relax.grad.start_checkpoint",
+ [input],
+ ty=ty,
+ span=span,
+ ) # type: ignore
-def end_checkpoint(input: Expr) -> Expr:
+def end_checkpoint(input: Expr, *, ty=None, span=None) -> Expr:
"""Mark the end of checkpoint stage. See tvm.relax.op.grad.start_checkpoint.
Parameters
@@ -92,7 +99,12 @@ def end_checkpoint(input: Expr) -> Expr:
result : relax.Expr
The same tensor as the input.
"""
- return _ffi_api.end_checkpoint(input) # type: ignore
+ return _Call(
+ "relax.grad.end_checkpoint",
+ [input],
+ ty=ty,
+ span=span,
+ ) # type: ignore
def nll_loss_backward(
@@ -102,6 +114,9 @@ def nll_loss_backward(
weights: Expr | None = None,
reduction: str = "mean",
ignore_index: int = -100,
+ *,
+ ty=None,
+ span=None,
) -> Expr:
"""Backward operator of relax.nn.nll_loss. All parameters except output_grad is the same as
relax.nn.nll_loss. Returns the gradient w.r.t. predictions.
@@ -116,8 +131,14 @@ def nll_loss_backward(
result : relax.Expr
The gradient w.r.t. predictions.
"""
- return _ffi_api.nll_loss_backward( # type: ignore
- output_grad, predictions, targets, weights, reduction, ignore_index
+ return _Call(
+ "relax.grad.nll_loss_backward",
+ [output_grad, predictions, targets, *([] if weights is None else [weights])],
+ attrs=_make_attrs(
+ "relax.attrs.NLLLossAttrs", reduction=reduction, ignore_index=ignore_index
+ ),
+ ty=ty,
+ span=span,
)
@@ -132,6 +153,9 @@ def max_pool2d_backward(
count_include_pad: bool = False,
layout: str = "NCHW",
out_layout: str | None = None,
+ *,
+ ty=None,
+ span=None,
) -> Expr:
"""Backward operator of relax.nn.max_pool2d. All parameters except output_grad is the same as
relax.nn.max_pool2d. Returns the gradient w.r.t. data.
@@ -146,17 +170,30 @@ def max_pool2d_backward(
result : relax.Expr
The gradient w.r.t. data.
"""
- return _ffi_api.max_pool2d_backward( # type: ignore
- output_grad,
- data,
- pool_size,
- strides,
- padding,
- dilation,
- ceil_mode,
- count_include_pad,
- layout,
- out_layout,
+ if len(padding) == 2:
+ padding = tuple(padding) + tuple(padding)
+ if len(strides) == 1:
+ strides = tuple(strides) * 2
+ if len(dilation) == 1:
+ dilation = tuple(dilation) * 2
+ if len(pool_size) == 1:
+ pool_size = tuple(pool_size) * 2
+ return _Call(
+ "relax.grad.max_pool2d_backward",
+ [output_grad, data],
+ attrs=_make_attrs(
+ "relax.attrs.Pool2DAttrs",
+ pool_size=pool_size,
+ strides=strides,
+ padding=padding,
+ dilation=dilation,
+ ceil_mode=ceil_mode,
+ count_include_pad=count_include_pad,
+ layout=layout,
+ out_layout=(layout if out_layout is None else out_layout),
+ ),
+ ty=ty,
+ span=span,
)
@@ -171,6 +208,9 @@ def avg_pool2d_backward(
count_include_pad: bool = False,
layout: str = "NCHW",
out_layout: str | None = None,
+ *,
+ ty=None,
+ span=None,
) -> Expr:
"""Backward operator of relax.nn.avg_pool2d. All parameters except output_grad is the same as
relax.nn.avg_pool2d. Returns the gradient w.r.t. data.
@@ -185,21 +225,42 @@ def avg_pool2d_backward(
result : relax.Expr
The gradient w.r.t. data.
"""
- return _ffi_api.avg_pool2d_backward( # type: ignore
- output_grad,
- data,
- pool_size,
- strides,
- padding,
- dilation,
- ceil_mode,
- count_include_pad,
- layout,
- out_layout,
+ if len(padding) == 2:
+ padding = tuple(padding) + tuple(padding)
+ if len(strides) == 1:
+ strides = tuple(strides) * 2
+ if len(dilation) == 1:
+ dilation = tuple(dilation) * 2
+ if len(pool_size) == 1:
+ pool_size = tuple(pool_size) * 2
+ return _Call(
+ "relax.grad.avg_pool2d_backward",
+ [output_grad, data],
+ attrs=_make_attrs(
+ "relax.attrs.Pool2DAttrs",
+ pool_size=pool_size,
+ strides=strides,
+ padding=padding,
+ dilation=dilation,
+ ceil_mode=ceil_mode,
+ count_include_pad=count_include_pad,
+ layout=layout,
+ out_layout=(layout if out_layout is None else out_layout),
+ ),
+ ty=ty,
+ span=span,
)
-def take_backward(output_grad: Expr, x: Expr, indices: Expr, axis: int | None = None) -> Expr:
+def take_backward(
+ output_grad: Expr,
+ x: Expr,
+ indices: Expr,
+ axis: int | None = None,
+ *,
+ ty=None,
+ span=None,
+) -> Expr:
"""Backward operator of relax.take. All parameters except output_grad is the same as
relax.take. Returns the gradient w.r.t. x.
@@ -213,4 +274,10 @@ def take_backward(output_grad: Expr, x: Expr, indices: Expr, axis: int | None =
result : relax.Expr
The gradient w.r.t. x.
"""
- return _ffi_api.take_backward(output_grad, x, indices, axis) # type: ignore
+ return _Call(
+ "relax.grad.take_backward",
+ [output_grad, x, indices],
+ attrs=_make_attrs("relax.attrs.TakeBackwardAttrs", axis=axis),
+ ty=ty,
+ span=span,
+ ) # type: ignore
diff --git a/python/tvm/relax/op/image/image.py b/python/tvm/relax/op/image/image.py
index cb71b85a28..7f8ad29da6 100644
--- a/python/tvm/relax/op/image/image.py
+++ b/python/tvm/relax/op/image/image.py
@@ -19,10 +19,12 @@
from typing import cast
from tvm import DataType
+from tvm.ir import Call as _Call
+from tvm.ir import PrimType as _PrimType
from tvm.ir import is_prim_expr
+from tvm.ir.attrs import make_node as _make_attrs
from ...expr import Expr, ShapeExpr
-from . import _ffi_api
PrimExprLike = int | Expr
SizeLike = PrimExprLike | tuple[PrimExprLike, ...]
@@ -40,6 +42,9 @@ def resize2d(
cubic_exclude: int = 0,
extrapolation_value: float = 0.0,
out_dtype: str | DataType | None = None,
+ *,
+ ty=None,
+ span=None,
) -> Expr:
"""Image resize2d operator.
@@ -118,18 +123,23 @@ def resize2d(
else:
size = ShapeExpr(size)
- return _ffi_api.resize2d( # type: ignore
- data,
- size,
- roi,
- layout,
- method,
- coordinate_transformation_mode,
- rounding_method,
- cubic_alpha,
- cubic_exclude,
- extrapolation_value,
- out_dtype,
+ return _Call(
+ "relax.image.resize2d",
+ [data, size],
+ attrs=_make_attrs(
+ "relax.attrs.Resize2DAttrs",
+ roi=roi,
+ layout=layout,
+ method=method,
+ coordinate_transformation_mode=coordinate_transformation_mode,
+ rounding_method=rounding_method,
+ cubic_alpha=cubic_alpha,
+ cubic_exclude=cubic_exclude,
+ extrapolation_value=extrapolation_value,
+ out_dtype=(out_dtype.dtype if isinstance(out_dtype, _PrimType) else out_dtype),
+ ),
+ ty=ty,
+ span=span,
)
@@ -145,6 +155,9 @@ def resize3d(
cubic_exclude: int = 0,
extrapolation_value: float = 0.0,
out_dtype: str | DataType | None = None,
+ *,
+ ty=None,
+ span=None,
) -> Expr:
"""Image resize3d operator.
@@ -170,18 +183,23 @@ def resize3d(
else:
size = ShapeExpr(size)
- return _ffi_api.resize3d( # type: ignore
- data,
- size,
- roi,
- layout,
- method,
- coordinate_transformation_mode,
- rounding_method,
- cubic_alpha,
- cubic_exclude,
- extrapolation_value,
- out_dtype,
+ return _Call(
+ "relax.image.resize3d",
+ [data, size],
+ attrs=_make_attrs(
+ "relax.attrs.Resize3DAttrs",
+ roi=roi,
+ layout=layout,
+ method=method,
+ coordinate_transformation_mode=coordinate_transformation_mode,
+ rounding_method=rounding_method,
+ cubic_alpha=cubic_alpha,
+ cubic_exclude=cubic_exclude,
+ extrapolation_value=extrapolation_value,
+ out_dtype=(out_dtype.dtype if isinstance(out_dtype, _PrimType) else out_dtype),
+ ),
+ ty=ty,
+ span=span,
)
@@ -192,6 +210,9 @@ def grid_sample(
layout: str = "NCHW",
padding_mode: str = "zeros",
align_corners: bool = False,
+ *,
+ ty=None,
+ span=None,
) -> Expr:
"""Applies grid sampling to input feature map.
@@ -224,13 +245,18 @@ def grid_sample(
result : relax.Expr
The sampled output tensor with shape [N, C, H_out, W_out].
"""
- return _ffi_api.grid_sample( # type: ignore
- data,
- grid,
- method,
- layout,
- padding_mode,
- align_corners,
+ return _Call(
+ "relax.image.grid_sample",
+ [data, grid],
+ attrs=_make_attrs(
+ "relax.attrs.GridSampleAttrs",
+ method=method,
+ layout=layout,
+ padding_mode=padding_mode,
+ align_corners=align_corners,
+ ),
+ ty=ty,
+ span=span,
)
@@ -238,6 +264,9 @@ def affine_grid(
data: Expr,
size: SizeLike,
align_corners: bool = True,
+ *,
+ ty=None,
+ span=None,
) -> Expr:
"""Generate a 2D or 3D sampling grid using an affine transformation matrix.
@@ -271,4 +300,13 @@ def affine_grid(
if isinstance(size, tuple | list):
size = ShapeExpr(size)
- return cast(Expr, _ffi_api.affine_grid(data, size, align_corners))
+ return cast(
+ Expr,
+ _Call(
+ "relax.image.affine_grid",
+ [data, size],
+ attrs=_make_attrs("relax.attrs.AffineGridAttrs", align_corners=align_corners),
+ ty=ty,
+ span=span,
+ ),
+ )
diff --git a/python/tvm/relax/op/index.py b/python/tvm/relax/op/index.py
index dd872f643e..a7fda22b45 100644
--- a/python/tvm/relax/op/index.py
+++ b/python/tvm/relax/op/index.py
@@ -16,14 +16,26 @@
# under the License.
"""Indexing operators."""
+from tvm.error import InternalError as _InternalError
+from tvm.ir import Call as _Call
+from tvm.ir.attrs import make_node as _make_attrs
+from tvm.relax.expr import Tuple as _Tuple
+
from ..expr import Expr
from ..utils import convert_to_expr
-from . import _ffi_api
PrimExprLike = int | Expr
-def take(x: Expr, indices: Expr, axis: int | None = None, mode: str = "fast") -> Expr:
+def take(
+ x: Expr,
+ indices: Expr,
+ axis: int | None = None,
+ mode: str = "fast",
+ *,
+ ty=None,
+ span=None,
+) -> Expr:
"""Take elements from a tensor along an axis.
Its semantic is mostly similar to `numpy.take`
(https://numpy.org/doc/stable/reference/generated/numpy.take.html),
@@ -53,7 +65,13 @@ def take(x: Expr, indices: Expr, axis: int | None = None, mode: str = "fast") ->
ret : relax.Expr
The taken result.
"""
- return _ffi_api.take(x, indices, axis, mode) # type: ignore
+ return _Call(
+ "relax.take",
+ [x, indices],
+ attrs=_make_attrs("relax.attrs.TakeAttrs", axis=axis, mode=mode),
+ ty=ty,
+ span=span,
+ ) # type: ignore
def strided_slice(
@@ -63,6 +81,9 @@ def strided_slice(
end: Expr,
strides: Expr | None = None,
assume_inbound: bool = False,
+ *,
+ ty=None,
+ span=None,
) -> Expr:
"""Strided slice of a tensor.
@@ -103,14 +124,22 @@ def strided_slice(
end = convert_to_expr(end)
if strides is not None:
strides = convert_to_expr(strides)
- return _ffi_api.strided_slice(x, axes, begin, end, strides, assume_inbound) # type: ignore
+ lengths = [
+ len(value.fields) for value in (axes, begin, end, strides) if isinstance(value, _Tuple)
+ ]
+ if lengths and any(length != lengths[0] for length in lengths):
+ raise _InternalError("axes, begin, end, and strides must have the same length")
+ return _Call(
+ "relax.strided_slice",
+ [x, axes, begin, end, *([] if strides is None else [strides])],
+ attrs=_make_attrs("relax.attrs.StridedSliceAttrs", assume_inbound=assume_inbound),
+ ty=ty,
+ span=span,
+ ) # type: ignore
def dynamic_strided_slice(
- x: Expr,
- begin: Expr,
- end: Expr,
- strides: Expr,
+ x: Expr, begin: Expr, end: Expr, strides: Expr, *, ty=None, span=None
) -> Expr:
"""Dynamic strided slice of a tensor. `begin`, `end`, `strides` can be computed at runtime.
@@ -140,4 +169,9 @@ def dynamic_strided_slice(
dyn_strided_slice require the input `begin`, `end` and `strides` to have the
same length as rank of `data` tensor.
"""
- return _ffi_api.dynamic_strided_slice(x, begin, end, strides) # type: ignore
+ return _Call(
+ "relax.dynamic_strided_slice",
+ [x, begin, end, strides],
+ ty=ty,
+ span=span,
+ ) # type: ignore
diff --git a/python/tvm/relax/op/linear_algebra.py b/python/tvm/relax/op/linear_algebra.py
index 8cc9e4def8..314470cabd 100644
--- a/python/tvm/relax/op/linear_algebra.py
+++ b/python/tvm/relax/op/linear_algebra.py
@@ -18,14 +18,23 @@
"""Relax linear algebra operators"""
from tvm import DataType
+from tvm.ir import Call as _Call
+from tvm.ir import PrimType as _PrimType
+from tvm.ir.attrs import make_node as _make_attrs
from ..expr import Expr
from ..expr import Tuple as RxTuple
-from . import _ffi_api
from .manipulate import permute_dims
-def matmul(x1: Expr, x2: Expr, out_dtype: str | DataType | None = None) -> Expr:
+def matmul(
+ x1: Expr,
+ x2: Expr,
+ out_dtype: str | DataType | None = None,
+ *,
+ ty=None,
+ span=None,
+) -> Expr:
"""General matrix multiplication of two tensors, with broadcasting on batched dimensions.
The semantics and output shape deduction rule is specified as
@@ -48,7 +57,16 @@ def matmul(x1: Expr, x2: Expr, out_dtype: str | DataType | None = None) -> Expr:
result : relax.Expr
The computed result.
"""
- return _ffi_api.matmul(x1, x2, out_dtype) # type: ignore
+ return _Call(
+ "relax.matmul",
+ [x1, x2],
+ attrs=_make_attrs(
+ "relax.attrs.MatmulAttrs",
+ out_dtype=(out_dtype.dtype if isinstance(out_dtype, _PrimType) else out_dtype),
+ ),
+ ty=ty,
+ span=span,
+ ) # type: ignore
def linear(
@@ -90,7 +108,7 @@ def linear(
return x + bias if bias is not None else x
-def einsum(operands, subscripts):
+def einsum(operands, subscripts, *, ty=None, span=None):
"""Evaluates the Einstein summation convention on data
Parameters
@@ -109,10 +127,16 @@ def einsum(operands, subscripts):
if isinstance(operands, list | tuple):
operands = RxTuple(operands)
- return _ffi_api.einsum(operands, subscripts) # type: ignore
+ return _Call(
+ "relax.einsum",
+ [operands],
+ attrs=_make_attrs("relax.attrs.EinsumAttrs", subscripts=subscripts),
+ ty=ty,
+ span=span,
+ ) # type: ignore
-def outer(x1: Expr, x2: Expr) -> Expr:
+def outer(x1: Expr, x2: Expr, *, ty=None, span=None) -> Expr:
"""
Computes the outer product of two input expressions.
@@ -136,4 +160,4 @@ def outer(x1: Expr, x2: Expr) -> Expr:
result : relax.Expr
The resulting expression representing the outer product.
"""
- return _ffi_api.outer(x1, x2)
+ return _Call("relax.outer", [x1, x2], ty=ty, span=span)
diff --git a/python/tvm/relax/op/manipulate.py b/python/tvm/relax/op/manipulate.py
index 9c7ad20197..58c58c7a63 100644
--- a/python/tvm/relax/op/manipulate.py
+++ b/python/tvm/relax/op/manipulate.py
@@ -18,7 +18,10 @@
from collections.abc import Callable
+from tvm.error import InternalError as _InternalError
+from tvm.ir import Call as _Call
from tvm.ir import is_prim_expr
+from tvm.ir.attrs import make_node as _make_attrs
from tvm.runtime import DataTypeCode
from tvm.tirx import FloatImm, IndexMap, IntImm
@@ -29,7 +32,7 @@ from . import _ffi_api
PrimExprLike = int | Expr
-def broadcast_to(x: Expr, shape: tuple[PrimExprLike] | Expr) -> Expr:
+def broadcast_to(x: Expr, shape: tuple[PrimExprLike] | Expr, *, ty=None, span=None) -> Expr:
"""Broadcasts a tensor to a specified shape.
Parameters
@@ -47,10 +50,16 @@ def broadcast_to(x: Expr, shape: tuple[PrimExprLike] | Expr) -> Expr:
"""
if isinstance(shape, tuple | list):
shape = ShapeExpr(shape)
- return _ffi_api.broadcast_to(x, shape) # type: ignore
+ return _Call("relax.broadcast_to", [x, shape], ty=ty, span=span) # type: ignore
-def concat(tensors: Expr | list[Expr], axis: int | None = 0) -> Expr:
+def concat(
+ tensors: Expr | list[Expr],
+ axis: int | None = 0,
+ *,
+ ty=None,
+ span=None,
+) -> Expr:
"""Concatenate the input tensors along the given axis.
Parameters
@@ -70,10 +79,16 @@ def concat(tensors: Expr | list[Expr], axis: int | None = 0) -> Expr:
"""
if isinstance(tensors, list | tuple):
tensors = RxTuple(tensors)
- return _ffi_api.concat(tensors, axis) # type: ignore
+ return _Call(
+ "relax.concat",
+ [tensors],
+ attrs=_make_attrs("relax.attrs.ConcatAttrs", axis=axis),
+ ty=ty,
+ span=span,
+ ) # type: ignore
-def expand_dims(x: Expr, axis: int | list[int]) -> Expr:
+def expand_dims(x: Expr, axis: int | list[int], *, ty=None, span=None) -> Expr:
"""Insert new axes at the positions given by `axis`.
Parameters
@@ -93,10 +108,16 @@ def expand_dims(x: Expr, axis: int | list[int]) -> Expr:
"""
if isinstance(axis, int):
axis = [axis]
- return _ffi_api.expand_dims(x, axis) # type: ignore
+ return _Call(
+ "relax.expand_dims",
+ [x],
+ attrs=_make_attrs("relax.attrs.ExpandDimsAttrs", axis=axis),
+ ty=ty,
+ span=span,
+ ) # type: ignore
-def flatten(x: Expr) -> Expr:
+def flatten(x: Expr, *, ty=None, span=None) -> Expr:
"""Flatten all the tensor dimensions into one.
Parameters
@@ -109,13 +130,16 @@ def flatten(x: Expr) -> Expr:
result : relax.Expr
The flattened result.
"""
- return _ffi_api.flatten(x) # type: ignore
+ return _Call("relax.flatten", [x], ty=ty, span=span) # type: ignore
def layout_transform(
x: Expr,
index_map: Callable | IndexMap,
pad_value: int | float | Expr | None = None,
+ *,
+ ty=None,
+ span=None,
):
"""Modifies the layout of a tensor.
@@ -155,10 +179,18 @@ def layout_transform(
pad_value = FloatImm(x_dtype.dtype, float(pad_value))
pad_value = prim_value(pad_value)
- return _ffi_api.layout_transform(x, index_map, pad_value)
+ return _Call(
+ "relax.layout_transform",
+ [x],
+ attrs=_make_attrs(
+ "relax.attrs.LayoutTransformAttrs", index_map=index_map, pad_value=pad_value
+ ),
+ ty=ty,
+ span=span,
+ )
-def permute_dims(x: Expr, axes: list[int] | None = None) -> Expr:
+def permute_dims(x: Expr, axes: list[int] | None = None, *, ty=None, span=None) -> Expr:
"""Permutes the dimensions of an array.
Parameters
@@ -174,10 +206,16 @@ def permute_dims(x: Expr, axes: list[int] | None = None) -> Expr:
result : relax.Expr
The transposed result.
"""
- return _ffi_api.permute_dims(x, axes) # type: ignore
+ return _Call(
+ "relax.permute_dims",
+ [x],
+ attrs=_make_attrs("relax.attrs.PermuteDimsAttrs", axes=axes),
+ ty=ty,
+ span=span,
+ ) # type: ignore
-def reshape(x: Expr, shape: tuple[PrimExprLike] | Expr) -> Expr:
+def reshape(x: Expr, shape: tuple[PrimExprLike] | Expr, *, ty=None, span=None) -> Expr:
"""Reshape the input array.
``-1`` infers the dimension of the output shape by using the remainder of
@@ -211,13 +249,17 @@ def reshape(x: Expr, shape: tuple[PrimExprLike] | Expr) -> Expr:
"""
if not isinstance(shape, tuple | list | Expr) or is_prim_expr(shape):
raise TypeError("shape must be a tuple/list or a Relax shape expression")
- return _ffi_api.reshape(x, shape) # type: ignore
+ shape = _ffi_api._NormalizeReshapeShape(x, shape)
+ return _Call("relax.reshape", [x, shape], ty=ty, span=span) # type: ignore
def split(
x: Expr,
indices_or_sections: int | list[PrimExprLike],
axis: int = 0,
+ *,
+ ty=None,
+ span=None,
) -> Expr:
"""Split input tensor along axis by sections or indices.
@@ -246,10 +288,25 @@ def split(
"""
if isinstance(indices_or_sections, int):
indices_or_sections = IntImm("int64", indices_or_sections)
- return _ffi_api.split(x, indices_or_sections, axis) # type: ignore
-
-
-def squeeze(x: Expr, axis: int | list[int] | None = None) -> Expr:
+ elif isinstance(indices_or_sections, list | tuple):
+ indices_or_sections = [IntImm("int64", int(index)) for index in indices_or_sections]
+ if isinstance(indices_or_sections, IntImm):
+ if int(indices_or_sections) <= 0:
+ raise _InternalError("The number of split sections must be positive")
+ elif not isinstance(indices_or_sections, list | tuple):
+ raise TypeError("indices_or_sections must be an integer or a list of integers")
+ return _Call(
+ "relax.split",
+ [x],
+ attrs=_make_attrs(
+ "relax.attrs.SplitAttrs", indices_or_sections=indices_or_sections, axis=axis
+ ),
+ ty=ty,
+ span=span,
+ ) # type: ignore
+
+
+def squeeze(x: Expr, axis: int | list[int] | None = None, *, ty=None, span=None) -> Expr:
"""Squeeze axes in the array.
Parameters
@@ -269,10 +326,16 @@ def squeeze(x: Expr, axis: int | list[int] | None = None) -> Expr:
"""
if isinstance(axis, int):
axis = [axis]
- return _ffi_api.squeeze(x, axis) # type: ignore
+ return _Call(
+ "relax.squeeze",
+ [x],
+ attrs=_make_attrs("relax.attrs.SqueezeAttrs", axis=axis),
+ ty=ty,
+ span=span,
+ ) # type: ignore
-def stack(tensors: Expr | list[Expr], axis: int = 0) -> Expr:
+def stack(tensors: Expr | list[Expr], axis: int = 0, *, ty=None, span=None) -> Expr:
"""Stack the input tensors along a new axis.
Parameters
@@ -293,10 +356,16 @@ def stack(tensors: Expr | list[Expr], axis: int = 0) -> Expr:
"""
if isinstance(tensors, list | tuple):
tensors = RxTuple(tensors)
- return _ffi_api.stack(tensors, axis) # type: ignore
+ return _Call(
+ "relax.stack",
+ [tensors],
+ attrs=_make_attrs("relax.attrs.StackAttrs", axis=axis),
+ ty=ty,
+ span=span,
+ ) # type: ignore
-def collapse_sum_like(data: Expr, collapse_target: Expr) -> Expr:
+def collapse_sum_like(data: Expr, collapse_target: Expr, *, ty=None, span=None) -> Expr:
"""Return a summation of data to the shape of collapse_target.
For details, please see relax.op.collapse_sum_to.
@@ -314,10 +383,15 @@ def collapse_sum_like(data: Expr, collapse_target: Expr) -> Expr:
result : relax.Expr
The result tensor after summation.
"""
- return _ffi_api.collapse_sum_like(data, collapse_target) # type: ignore
+ return _Call(
+ "relax.collapse_sum_like",
+ [data, collapse_target],
+ ty=ty,
+ span=span,
+ ) # type: ignore
-def collapse_sum_to(data: Expr, shape: tuple[PrimExprLike] | Expr) -> Expr:
+def collapse_sum_to(data: Expr, shape: tuple[PrimExprLike] | Expr, *, ty=None, span=None) -> Expr:
"""Return a summation of data to the given shape.
collapse_sum_to is intended as the backward operator of tvm.relax.op.broadcast_to and
@@ -346,10 +420,22 @@ def collapse_sum_to(data: Expr, shape: tuple[PrimExprLike] | Expr) -> Expr:
"""
if isinstance(shape, tuple | list):
shape = ShapeExpr(shape)
- return _ffi_api.collapse_sum_to(data, shape) # type: ignore
+ return _Call(
+ "relax.collapse_sum_to",
+ [data, shape],
+ ty=ty,
+ span=span,
+ ) # type: ignore
-def repeat(data: Expr, repeats: int, axis: int | None = None) -> Expr:
+def repeat(
+ data: Expr,
+ repeats: int,
+ axis: int | None = None,
+ *,
+ ty=None,
+ span=None,
+) -> Expr:
"""Repeats elements of an array.
Parameters
@@ -379,10 +465,22 @@ def repeat(data: Expr, repeats: int, axis: int | None = None) -> Expr:
lv2 = R.repeat(x, repeats=2, axis=1) # lv2 == [[1., 1., 2., 2.],
# [3., 3., 4., 4.]]
"""
- return _ffi_api.repeat(data, repeats, axis) # type: ignore
+ return _Call(
+ "relax.repeat",
+ [data],
+ attrs=_make_attrs("relax.attrs.RepeatAttrs", repeats=repeats, axis=axis),
+ ty=ty,
+ span=span,
+ ) # type: ignore
-def tile(data: Expr, repeats: int | tuple[int] | list[int]) -> Expr:
+def tile(
+ data: Expr,
+ repeats: int | tuple[int] | list[int],
+ *,
+ ty=None,
+ span=None,
+) -> Expr:
"""Construct an array by repeating data the number of times given by repeats.
If repeats has length l, and data has dimension d, the result will have dimension of max(l, d).
@@ -421,10 +519,16 @@ def tile(data: Expr, repeats: int | tuple[int] | list[int]) -> Expr:
"""
if isinstance(repeats, int):
repeats = [repeats]
- return _ffi_api.tile(data, repeats) # type: ignore
+ return _Call(
+ "relax.tile",
+ [data],
+ attrs=_make_attrs("relax.attrs.TileAttrs", repeats=repeats),
+ ty=ty,
+ span=span,
+ ) # type: ignore
-def flip(data, axis):
+def flip(data, axis, *, ty=None, span=None):
"""Reverses the order of elements along given axis while preserving array shape.
Parameters
@@ -449,10 +553,24 @@ def flip(data, axis):
relax.flip(x, axis=1) = [[2., 1.], [4., 3.]]
"""
- return _ffi_api.flip(data, axis) # type: ignore
+ return _Call(
+ "relax.flip",
+ [data],
+ attrs=_make_attrs("relax.attrs.FlipAttrs", axis=axis),
+ ty=ty,
+ span=span,
+ ) # type: ignore
-def reverse_sequence(data: Expr, seq_lengths: Expr, seq_axis: int = 1, batch_axis: int = 0) -> Expr:
+def reverse_sequence(
+ data: Expr,
+ seq_lengths: Expr,
+ seq_axis: int = 1,
+ batch_axis: int = 0,
+ *,
+ ty=None,
+ span=None,
+) -> Expr:
"""Reverses variable length slices.
Parameters
@@ -474,10 +592,18 @@ def reverse_sequence(data: Expr, seq_lengths: Expr, seq_axis: int = 1, batch_axi
ret : relax.Expr
The computed result.
"""
- return _ffi_api.reverse_sequence(data, seq_lengths, seq_axis, batch_axis) # type: ignore
-
-
-def gather_elements(data: Expr, indices: Expr, axis: int = 0) -> Expr:
+ return _Call(
+ "relax.reverse_sequence",
+ [data, seq_lengths],
+ attrs=_make_attrs(
+ "relax.attrs.ReverseSequenceAttrs", seq_axis=seq_axis, batch_axis=batch_axis
+ ),
+ ty=ty,
+ span=span,
+ ) # type: ignore
+
+
+def gather_elements(data: Expr, indices: Expr, axis: int = 0, *, ty=None, span=None) -> Expr:
"""Gather elements from data according to indices along the specified axis.
Parameters
@@ -510,10 +636,16 @@ def gather_elements(data: Expr, indices: Expr, axis: int = 0) -> Expr:
axis = 0
output = [[4, 5, 6]]
"""
- return _ffi_api.gather_elements(data, indices, axis) # type: ignore
+ return _Call(
+ "relax.gather_elements",
+ [data, indices],
+ attrs=_make_attrs("relax.attrs.GatherElementsAttrs", axis=axis),
+ ty=ty,
+ span=span,
+ ) # type: ignore
-def gather_nd(data: Expr, indices: Expr, batch_dims: int = 0) -> Expr:
+def gather_nd(data: Expr, indices: Expr, batch_dims: int = 0, *, ty=None, span=None) -> Expr:
"""Update data at positions defined by indices with values in updates.
Parameters
@@ -547,10 +679,16 @@ def gather_nd(data: Expr, indices: Expr, batch_dims: int = 0) -> Expr:
output = [[2,3],[4,5]] # output_shape = [2, 2]
"""
- return _ffi_api.gather_nd(data, indices, batch_dims) # type: ignore
+ return _Call(
+ "relax.gather_nd",
+ [data, indices],
+ attrs=_make_attrs("relax.attrs.GatherNDAttrs", batch_dims=batch_dims),
+ ty=ty,
+ span=span,
+ ) # type: ignore
-def index_tensor(data: Expr, indices: Expr | list[Expr]) -> Expr:
+def index_tensor(data: Expr, indices: Expr | list[Expr], *, ty=None, span=None) -> Expr:
"""Advanced-tensor indexing (NumPy/PyTorch-style).
Given k index tensors ``indices = (I0, I1, …, Ik-1)`` this
@@ -610,7 +748,12 @@ def index_tensor(data: Expr, indices: Expr | list[Expr]) -> Expr:
"""
if isinstance(indices, list | tuple):
indices = RxTuple(indices)
- return _ffi_api.index_tensor(data, indices) # type: ignore
+ return _Call(
+ "relax.index_tensor",
+ [data, indices],
+ ty=ty,
+ span=span,
+ ) # type: ignore
def index_put(
@@ -618,6 +761,9 @@ def index_put(
indices: Expr | tuple[Expr],
values: Expr,
accumulate: bool = False,
+ *,
+ ty=None,
+ span=None,
) -> Expr:
"""This operation updates values in `data` at positions
specified by `indices` with corresponding values from `values`. The `indices` is a tuple
@@ -662,10 +808,22 @@ def index_put(
"""
if isinstance(indices, list | tuple):
indices = RxTuple(indices)
- return _ffi_api.index_put(data, indices, values, accumulate) # type: ignore
-
-
-def meshgrid(tensors: Expr | list[Expr], indexing: str | None = "ij") -> Expr:
+ return _Call(
+ "relax.index_put",
+ [data, indices, values],
+ attrs=_make_attrs("relax.attrs.IndexPutAttrs", accumulate=accumulate),
+ ty=ty,
+ span=span,
+ ) # type: ignore
+
+
+def meshgrid(
+ tensors: Expr | list[Expr],
+ indexing: str | None = "ij",
+ *,
+ ty=None,
+ span=None,
+) -> Expr:
"""Generate coordinate grids from input tensors.
Parameters
@@ -685,11 +843,24 @@ def meshgrid(tensors: Expr | list[Expr], indexing: str | None = "ij") -> Expr:
"""
if isinstance(tensors, list | tuple):
tensors = RxTuple(tensors)
- return _ffi_api.meshgrid(tensors, indexing)
+ return _Call(
+ "relax.meshgrid",
+ [tensors],
+ attrs=_make_attrs("relax.attrs.MeshgridAttrs", indexing=indexing),
+ ty=ty,
+ span=span,
+ )
def scatter_elements(
- data: Expr, indices: Expr, updates: Expr, axis: int = 0, reduction: str = "update"
+ data: Expr,
+ indices: Expr,
+ updates: Expr,
+ axis: int = 0,
+ reduction: str = "update",
+ *,
+ ty=None,
+ span=None,
):
"""ONNX style scatter elements. This operation updates its value in `data` to values
specified by `updates` at specific index positions specified by `indices`.
@@ -763,10 +934,24 @@ def scatter_elements(
]
"""
- return _ffi_api.scatter_elements(data, indices, updates, axis, reduction) # type: ignore
+ return _Call(
+ "relax.scatter_elements",
+ [data, indices, updates],
+ attrs=_make_attrs("relax.attrs.ScatterElementsAttrs", axis=axis, reduction=reduction),
+ ty=ty,
+ span=span,
+ ) # type: ignore
-def scatter_nd(data: Expr, indices: Expr, updates: Expr, reduction: str = "update") -> Expr:
+def scatter_nd(
+ data: Expr,
+ indices: Expr,
+ updates: Expr,
+ reduction: str = "update",
+ *,
+ ty=None,
+ span=None,
+) -> Expr:
"""Scatter updates into an array according to indices.
Parameters
@@ -802,10 +987,26 @@ def scatter_nd(data: Expr, indices: Expr, updates: Expr, reduction: str = "updat
output = [1, 11, 3, 10, 9, 6, 7, 12]
"""
- return _ffi_api.scatter_nd(data, indices, updates, reduction) # type: ignore
-
-
-def slice_scatter(input_tensor: Expr, src: Expr, start, end, step, axis=0):
+ return _Call(
+ "relax.scatter_nd",
+ [data, indices, updates],
+ attrs=_make_attrs("relax.attrs.ScatterNDAttrs", reduction=reduction),
+ ty=ty,
+ span=span,
+ ) # type: ignore
+
+
+def slice_scatter(
+ input_tensor: Expr,
+ src: Expr,
+ start,
+ end,
+ step,
+ axis=0,
+ *,
+ ty=None,
+ span=None,
+):
"""Embeds the values of the src tensor into input at the given dimension.
Parameters
@@ -840,7 +1041,13 @@ def slice_scatter(input_tensor: Expr, src: Expr, start, end, step, axis=0):
end = prim_value(end)
if not is_prim_expr(step):
step = prim_value(step)
- return _ffi_api.slice_scatter(input_tensor, src, axis, start, end, step)
+ return _Call(
+ "relax.slice_scatter",
+ [input_tensor, src, start, end, step],
+ attrs=_make_attrs("relax.attrs.SliceScatterAttrs", axis=axis),
+ ty=ty,
+ span=span,
+ )
def one_hot(
@@ -849,6 +1056,9 @@ def one_hot(
off_value: int | float | Expr,
depth: int,
axis: int = -1,
+ *,
+ ty=None,
+ span=None,
) -> Expr:
"""Returns a one-hot tensor.
@@ -890,4 +1100,12 @@ def one_hot(
"""
on_value = prim_value(on_value)
off_value = prim_value(off_value)
- return _ffi_api.one_hot(indices, on_value, off_value, depth, axis) # type: ignore
+ if depth <= 0:
+ raise _InternalError("one_hot: depth must be positive")
+ return _Call(
+ "relax.one_hot",
+ [indices, on_value, off_value],
+ attrs=_make_attrs("relax.attrs.OneHotAttrs", depth=depth, axis=axis),
+ ty=ty,
+ span=span,
+ ) # type: ignore
diff --git a/python/tvm/relax/op/memory/memory.py b/python/tvm/relax/op/memory/memory.py
index 2cdcbdfe7b..1e0b9c30bd 100644
--- a/python/tvm/relax/op/memory/memory.py
+++ b/python/tvm/relax/op/memory/memory.py
@@ -16,10 +16,10 @@
"""Relax memory primitives."""
from tvm.ir import Call, DataTypeImm, GenericConst, StringImm
+from tvm.ir import Call as _Call
from ...expr import Expr, prim_value
from ...utils import convert_to_expr
-from . import _ffi_api
def alloc_storage(
@@ -27,6 +27,9 @@ def alloc_storage(
virtual_device_index: int | Expr,
storage_scope: str | Expr,
dtype: str | Expr,
+ *,
+ ty=None,
+ span=None,
) -> Call:
"""Construct a Call to allocate a storage with specific size, virtual_device_index,
storage_scope and dtype.
@@ -58,7 +61,12 @@ def alloc_storage(
storage_scope = StringImm(storage_scope)
if isinstance(virtual_device_index, int):
virtual_device_index = prim_value(virtual_device_index)
- return _ffi_api.alloc_storage(size, virtual_device_index, storage_scope, dtype) # type: ignore
+ return _Call(
+ "relax.memory.alloc_storage",
+ [size, virtual_device_index, storage_scope, dtype],
+ ty=ty,
+ span=span,
+ ) # type: ignore
def alloc_tensor(
@@ -67,6 +75,9 @@ def alloc_tensor(
shape: Expr,
dtype: str | Expr,
runtime_device_ind: int | Expr = prim_value(0),
+ *,
+ ty=None,
+ span=None,
) -> Call:
"""Construct a Call to allocate a tensor on a certain storage starting from the given offset.
@@ -100,10 +111,15 @@ def alloc_tensor(
dtype = DataTypeImm(dtype.value if isinstance(dtype, GenericConst) else dtype)
if isinstance(runtime_device_ind, int):
runtime_device_ind = prim_value(runtime_device_ind)
- return _ffi_api.alloc_tensor(storage, offset, shape, dtype, runtime_device_ind) # type: ignore
+ return _Call(
+ "relax.memory.alloc_tensor",
+ [storage, offset, shape, dtype, runtime_device_ind],
+ ty=ty,
+ span=span,
+ ) # type: ignore
-def kill_storage(storage: Expr) -> Call:
+def kill_storage(storage: Expr, *, ty=None, span=None) -> Call:
"""Construct a Call to kill a storage.
Parameters
@@ -116,10 +132,15 @@ def kill_storage(storage: Expr) -> Call:
result : Call
A relax Call to kill a storage.
"""
- return _ffi_api.kill_storage(storage) # type: ignore
+ return _Call(
+ "relax.memory.kill_storage",
+ [storage],
+ ty=ty,
+ span=span,
+ ) # type: ignore
-def kill_tensor(tensor: Expr) -> Call:
+def kill_tensor(tensor: Expr, *, ty=None, span=None) -> Call:
"""Construct a Call to kill a tensor.
Parameters
@@ -132,4 +153,9 @@ def kill_tensor(tensor: Expr) -> Call:
result : Call
A relax Call to kill a tensor.
"""
- return _ffi_api.kill_tensor(tensor) # type: ignore
+ return _Call(
+ "relax.memory.kill_tensor",
+ [tensor],
+ ty=ty,
+ span=span,
+ ) # type: ignore
diff --git a/python/tvm/relax/op/memory/view.py b/python/tvm/relax/op/memory/view.py
index b3ecc57842..df9daea4c4 100644
--- a/python/tvm/relax/op/memory/view.py
+++ b/python/tvm/relax/op/memory/view.py
@@ -28,12 +28,12 @@ while keeping the same underlying data.
from collections.abc import Sequence
+from tvm.ir import Call as _Call
from tvm.ir import DataTypeImm, GenericConst
-from tvm.relax import Expr, ShapeExpr
+from tvm.relax import Expr, ShapeExpr, Tuple
from tvm.relax.expr import prim_value
from ..base import null_value
-from . import _ffi_api
PrimExprLike = int | Expr
@@ -43,6 +43,9 @@ def view(
shape: Sequence[PrimExprLike] | Expr | None = None,
dtype: Expr | None = None,
relative_byte_offset: Expr | None = None,
+ *,
+ ty=None,
+ span=None,
) -> Expr:
"""Provide a view into an existing tensor
@@ -99,10 +102,20 @@ def view(
else prim_value(relative_byte_offset)
)
- return _ffi_api.view(data, shape, dtype, relative_byte_offset) # type: ignore
+ return _Call(
+ "relax.memory.view",
+ [
+ data,
+ Tuple([]) if shape is None else shape,
+ dtype,
+ Tuple([]) if relative_byte_offset is None else relative_byte_offset,
+ ],
+ ty=ty,
+ span=span,
+ ) # type: ignore
-def ensure_zero_offset(data: Expr) -> Expr:
+def ensure_zero_offset(data: Expr, *, ty=None, span=None) -> Expr:
"""
Ensure the tensor has elem_offset == 0. A copy will be made if necessary.
@@ -116,4 +129,9 @@ def ensure_zero_offset(data: Expr) -> Expr:
result : relax.Expr
The tensor with elem_offset == 0
"""
- return _ffi_api.ensure_zero_offset(data) # type: ignore
+ return _Call(
+ "relax.memory.ensure_zero_offset",
+ [data],
+ ty=ty,
+ span=span,
+ ) # type: ignore
diff --git a/python/tvm/relax/op/nn/nn.py b/python/tvm/relax/op/nn/nn.py
index 6755782fda..22b4d48d88 100644
--- a/python/tvm/relax/op/nn/nn.py
+++ b/python/tvm/relax/op/nn/nn.py
@@ -17,10 +17,13 @@
"""Relax Neural Network (NN) operators"""
from tvm import DataType, relax
+from tvm.error import InternalError as _InternalError
+from tvm.ir import Call as _Call
+from tvm.ir import PrimType as _PrimType
+from tvm.ir.attrs import make_node as _make_attrs
from tvm.tirx import FloatImm
from ...expr import Expr
-from . import _ffi_api
def conv1d(
@@ -34,6 +37,10 @@ def conv1d(
kernel_layout: str = "OIW",
out_layout: str | None = None,
out_dtype: str | DataType | None = None,
+ *,
+ ty_args=None,
+ ty=None,
+ span=None,
) -> Expr:
r"""1D convolution.
@@ -107,17 +114,33 @@ def conv1d(
if isinstance(padding, int):
padding = (padding, padding)
- return _ffi_api.conv1d( # type: ignore
- data,
- weight,
- strides,
- padding,
- dilation,
- groups,
- data_layout,
- kernel_layout,
- out_layout,
- out_dtype,
+ if len(padding) == 1:
+ padding = tuple(padding) + tuple(padding)
+ if groups <= 0:
+ raise _InternalError("The number of groups must be positive")
+ if len(strides) != 1:
+ raise _InternalError("strides must contain 1 values")
+ if len(padding) != 2:
+ raise _InternalError("padding must contain 2 values")
+ if len(dilation) != 1:
+ raise _InternalError("dilation must contain 1 values")
+ return _Call(
+ "relax.nn.conv1d",
+ [data, weight],
+ attrs=_make_attrs(
+ "relax.attrs.Conv1DAttrs",
+ strides=strides,
+ padding=padding,
+ dilation=dilation,
+ groups=groups,
+ data_layout=data_layout,
+ kernel_layout=kernel_layout,
+ out_layout=(data_layout if out_layout is None else out_layout),
+ out_dtype=(out_dtype.dtype if isinstance(out_dtype, _PrimType) else out_dtype),
+ ),
+ ty_args=ty_args,
+ ty=ty,
+ span=span,
)
@@ -132,6 +155,10 @@ def conv2d(
kernel_layout: str = "OIHW",
out_layout: str | None = None,
out_dtype: str | DataType | None = None,
+ *,
+ ty_args=None,
+ ty=None,
+ span=None,
) -> Expr:
r"""2D convolution.
@@ -205,17 +232,37 @@ def conv2d(
if isinstance(padding, int):
padding = (padding, padding, padding, padding)
- return _ffi_api.conv2d( # type: ignore
- data,
- weight,
- strides,
- padding,
- dilation,
- groups,
- data_layout,
- kernel_layout,
- out_layout,
- out_dtype,
+ if len(padding) == 2:
+ padding = tuple(padding) + tuple(padding)
+ if len(strides) == 1:
+ strides = tuple(strides) * 2
+ if len(dilation) == 1:
+ dilation = tuple(dilation) * 2
+ if groups <= 0:
+ raise _InternalError("The number of groups must be positive")
+ if len(strides) != 2:
+ raise _InternalError("strides must contain 2 values")
+ if len(padding) != 4:
+ raise _InternalError("padding must contain 4 values")
+ if len(dilation) != 2:
+ raise _InternalError("dilation must contain 2 values")
+ return _Call(
+ "relax.nn.conv2d",
+ [data, weight],
+ attrs=_make_attrs(
+ "relax.attrs.Conv2DAttrs",
+ strides=strides,
+ padding=padding,
+ dilation=dilation,
+ groups=groups,
+ data_layout=data_layout,
+ kernel_layout=kernel_layout,
+ out_layout=(data_layout if out_layout is None else out_layout),
+ out_dtype=(out_dtype.dtype if isinstance(out_dtype, _PrimType) else out_dtype),
+ ),
+ ty_args=ty_args,
+ ty=ty,
+ span=span,
)
@@ -230,6 +277,10 @@ def conv3d(
kernel_layout: str = "OIDHW",
out_layout: str | None = None,
out_dtype: str | DataType | None = None,
+ *,
+ ty_args=None,
+ ty=None,
+ span=None,
) -> Expr:
r"""3D convolution.
@@ -309,17 +360,37 @@ def conv3d(
if isinstance(padding, int):
padding = (padding, padding, padding, padding, padding, padding)
- return _ffi_api.conv3d( # type: ignore
- data,
- weight,
- strides,
- padding,
- dilation,
- groups,
- data_layout,
- kernel_layout,
- out_layout,
- out_dtype,
+ if len(padding) == 3:
+ padding = tuple(padding) + tuple(padding)
+ if len(strides) == 1:
+ strides = tuple(strides) * 3
+ if len(dilation) == 1:
+ dilation = tuple(dilation) * 3
+ if groups <= 0:
+ raise _InternalError("The number of groups must be positive")
+ if len(strides) != 3:
+ raise _InternalError("strides must contain 3 values")
+ if len(padding) != 6:
+ raise _InternalError("padding must contain 6 values")
+ if len(dilation) != 3:
+ raise _InternalError("dilation must contain 3 values")
+ return _Call(
+ "relax.nn.conv3d",
+ [data, weight],
+ attrs=_make_attrs(
+ "relax.attrs.Conv3DAttrs",
+ strides=strides,
+ padding=padding,
+ dilation=dilation,
+ groups=groups,
+ data_layout=data_layout,
+ kernel_layout=kernel_layout,
+ out_layout=(data_layout if out_layout is None else out_layout),
+ out_dtype=(out_dtype.dtype if isinstance(out_dtype, _PrimType) else out_dtype),
+ ),
+ ty_args=ty_args,
+ ty=ty,
+ span=span,
)
@@ -335,6 +406,10 @@ def conv1d_transpose(
kernel_layout: str = "IOW",
out_layout: str | None = None,
out_dtype: str | DataType | None = None,
+ *,
+ ty_args=None,
+ ty=None,
+ span=None,
) -> Expr:
r"""1D transposed convolution operator.
@@ -399,18 +474,36 @@ def conv1d_transpose(
if isinstance(output_padding, int):
output_padding = (output_padding,)
- return _ffi_api.conv1d_transpose( # type: ignore
- data,
- weight,
- strides,
- padding,
- output_padding,
- dilation,
- groups,
- data_layout,
- kernel_layout,
- out_layout,
- out_dtype,
+ if len(padding) == 1:
+ padding = tuple(padding) + tuple(padding)
+ if groups <= 0:
+ raise _InternalError("The number of groups must be positive")
+ if len(strides) != 1:
+ raise _InternalError("strides must contain 1 values")
+ if len(padding) != 2:
+ raise _InternalError("padding must contain 2 values")
+ if len(dilation) != 1:
+ raise _InternalError("dilation must contain 1 values")
+ if len(output_padding) != 1:
+ raise _InternalError("output_padding must contain 1 values")
+ return _Call(
+ "relax.nn.conv1d_transpose",
+ [data, weight],
+ attrs=_make_attrs(
+ "relax.attrs.Conv1DTransposeAttrs",
+ strides=strides,
+ padding=padding,
+ output_padding=output_padding,
+ dilation=dilation,
+ groups=groups,
+ data_layout=data_layout,
+ kernel_layout=kernel_layout,
+ out_layout=(data_layout if out_layout is None else out_layout),
+ out_dtype=(out_dtype.dtype if isinstance(out_dtype, _PrimType) else out_dtype),
+ ),
+ ty_args=ty_args,
+ ty=ty,
+ span=span,
)
@@ -426,6 +519,10 @@ def conv2d_transpose(
kernel_layout: str = "IOHW",
out_layout: str | None = None,
out_dtype: str | DataType | None = None,
+ *,
+ ty_args=None,
+ ty=None,
+ span=None,
) -> Expr:
r"""Two dimensional transposed convolution operator.
@@ -501,18 +598,42 @@ def conv2d_transpose(
if isinstance(output_padding, int):
output_padding = (output_padding, output_padding)
- return _ffi_api.conv2d_transpose( # type: ignore
- data,
- weight,
- strides,
- padding,
- output_padding,
- dilation,
- groups,
- data_layout,
- kernel_layout,
- out_layout,
- out_dtype,
+ if len(padding) == 2:
+ padding = tuple(padding) + tuple(padding)
+ if len(strides) == 1:
+ strides = tuple(strides) * 2
+ if len(dilation) == 1:
+ dilation = tuple(dilation) * 2
+ if len(output_padding) == 1:
+ output_padding = tuple(output_padding) * 2
+ if groups <= 0:
+ raise _InternalError("The number of groups must be positive")
+ if len(strides) != 2:
+ raise _InternalError("strides must contain 2 values")
+ if len(padding) != 4:
+ raise _InternalError("padding must contain 4 values")
+ if len(dilation) != 2:
+ raise _InternalError("dilation must contain 2 values")
+ if len(output_padding) != 2:
+ raise _InternalError("output_padding must contain 2 values")
+ return _Call(
+ "relax.nn.conv2d_transpose",
+ [data, weight],
+ attrs=_make_attrs(
+ "relax.attrs.Conv2DTransposeAttrs",
+ strides=strides,
+ padding=padding,
+ output_padding=output_padding,
+ dilation=dilation,
+ groups=groups,
+ data_layout=data_layout,
+ kernel_layout=kernel_layout,
+ out_layout=(data_layout if out_layout is None else out_layout),
+ out_dtype=(out_dtype.dtype if isinstance(out_dtype, _PrimType) else out_dtype),
+ ),
+ ty_args=ty_args,
+ ty=ty,
+ span=span,
)
@@ -528,6 +649,10 @@ def conv3d_transpose(
kernel_layout: str = "IODHW",
out_layout: str | None = None,
out_dtype: str | DataType | None = None,
+ *,
+ ty_args=None,
+ ty=None,
+ span=None,
) -> Expr:
r"""Three dimensional transposed convolution operator.
@@ -604,18 +729,42 @@ def conv3d_transpose(
if isinstance(output_padding, int):
output_padding = (output_padding, output_padding, output_padding)
- return _ffi_api.conv3d_transpose( # type: ignore
- data,
- weight,
- strides,
- padding,
- output_padding,
- dilation,
- groups,
- data_layout,
- kernel_layout,
- out_layout,
- out_dtype,
+ if len(padding) == 3:
+ padding = tuple(padding) + tuple(padding)
+ if len(strides) == 1:
+ strides = tuple(strides) * 3
+ if len(dilation) == 1:
+ dilation = tuple(dilation) * 3
+ if len(output_padding) == 1:
+ output_padding = tuple(output_padding) * 3
+ if groups <= 0:
+ raise _InternalError("The number of groups must be positive")
+ if len(strides) != 3:
+ raise _InternalError("strides must contain 3 values")
+ if len(padding) != 6:
+ raise _InternalError("padding must contain 6 values")
+ if len(dilation) != 3:
+ raise _InternalError("dilation must contain 3 values")
+ if len(output_padding) != 3:
+ raise _InternalError("output_padding must contain 3 values")
+ return _Call(
+ "relax.nn.conv3d_transpose",
+ [data, weight],
+ attrs=_make_attrs(
+ "relax.attrs.Conv3DTransposeAttrs",
+ strides=strides,
+ padding=padding,
+ output_padding=output_padding,
+ dilation=dilation,
+ groups=groups,
+ data_layout=data_layout,
+ kernel_layout=kernel_layout,
+ out_layout=(data_layout if out_layout is None else out_layout),
+ out_dtype=(out_dtype.dtype if isinstance(out_dtype, _PrimType) else out_dtype),
+ ),
+ ty_args=ty_args,
+ ty=ty,
+ span=span,
)
@@ -624,6 +773,9 @@ def pad(
pad_width: list[int] | tuple[int, ...],
pad_mode: str | None = "constant",
pad_value: float | None = 0.0,
+ *,
+ ty=None,
+ span=None,
):
r"""Padding
@@ -652,10 +804,18 @@ def pad(
result : relax.Expr
The computed result.
"""
- return _ffi_api.pad(data, pad_width, pad_mode, pad_value)
+ return _Call(
+ "relax.nn.pad",
+ [data],
+ attrs=_make_attrs(
+ "relax.attrs.PadAttrs", pad_width=pad_width, pad_mode=pad_mode, pad_value=pad_value
+ ),
+ ty=ty,
+ span=span,
+ )
-def pixel_shuffle(data: Expr, upscale_factor: int):
+def pixel_shuffle(data: Expr, upscale_factor: int, *, ty=None, span=None):
r"""
Pixel Shuffle Operator
@@ -685,7 +845,13 @@ def pixel_shuffle(data: Expr, upscale_factor: int):
If the input tensor has shape (1, 8, 10, 15) and `upscale_factor` is 2,
the resulting tensor will have shape (1, 2, 20, 30).
"""
- return _ffi_api.pixel_shuffle(data, upscale_factor)
+ return _Call(
+ "relax.nn.pixel_shuffle",
+ [data],
+ attrs=_make_attrs("relax.attrs.PixelShuffleAttrs", upscale_factor=upscale_factor),
+ ty=ty,
+ span=span,
+ )
def max_pool1d(
@@ -698,6 +864,9 @@ def max_pool1d(
count_include_pad: bool = False,
layout: str = "NCW",
out_layout: str | None = None,
+ *,
+ ty=None,
+ span=None,
) -> Expr:
r"""1D maximum pooling operator.
@@ -756,16 +925,32 @@ def max_pool1d(
if isinstance(padding, int):
padding = (padding, padding)
- return _ffi_api.max_pool1d( # type: ignore
- data,
- pool_size,
- strides,
- padding,
- dilation,
- ceil_mode,
- count_include_pad,
- layout,
- out_layout,
+ if len(padding) == 1:
+ padding = tuple(padding) + tuple(padding)
+ if len(strides) != 1:
+ raise _InternalError("strides must contain 1 values")
+ if len(padding) != 2:
+ raise _InternalError("padding must contain 2 values")
+ if len(dilation) != 1:
+ raise _InternalError("dilation must contain 1 values")
+ if len(pool_size) != 1:
+ raise _InternalError("pool_size must contain 1 values")
+ return _Call(
+ "relax.nn.max_pool1d",
+ [data],
+ attrs=_make_attrs(
+ "relax.attrs.Pool1DAttrs",
+ pool_size=pool_size,
+ strides=strides,
+ padding=padding,
+ dilation=dilation,
+ ceil_mode=ceil_mode,
+ count_include_pad=count_include_pad,
+ layout=layout,
+ out_layout=(layout if out_layout is None else out_layout),
+ ),
+ ty=ty,
+ span=span,
)
@@ -779,6 +964,10 @@ def max_pool2d(
count_include_pad: bool = False,
layout: str = "NCHW",
out_layout: str | None = None,
+ *,
+ ty_args=None,
+ ty=None,
+ span=None,
) -> Expr:
r"""2D maximum pooling operator.
@@ -844,16 +1033,39 @@ def max_pool2d(
if isinstance(padding, int):
padding = (padding, padding, padding, padding)
- return _ffi_api.max_pool2d( # type: ignore
- data,
- pool_size,
- strides,
- padding,
- dilation,
- ceil_mode,
- count_include_pad,
- layout,
- out_layout,
+ if len(padding) == 2:
+ padding = tuple(padding) + tuple(padding)
+ if len(strides) == 1:
+ strides = tuple(strides) * 2
+ if len(dilation) == 1:
+ dilation = tuple(dilation) * 2
+ if len(pool_size) == 1:
+ pool_size = tuple(pool_size) * 2
+ if len(strides) != 2:
+ raise _InternalError("strides must contain 2 values")
+ if len(padding) != 4:
+ raise _InternalError("padding must contain 4 values")
+ if len(dilation) != 2:
+ raise _InternalError("dilation must contain 2 values")
+ if len(pool_size) != 2:
+ raise _InternalError("pool_size must contain 2 values")
+ return _Call(
+ "relax.nn.max_pool2d",
+ [data],
+ attrs=_make_attrs(
+ "relax.attrs.Pool2DAttrs",
+ pool_size=pool_size,
+ strides=strides,
+ padding=padding,
+ dilation=dilation,
+ ceil_mode=ceil_mode,
+ count_include_pad=count_include_pad,
+ layout=layout,
+ out_layout=(layout if out_layout is None else out_layout),
+ ),
+ ty_args=ty_args,
+ ty=ty,
+ span=span,
)
@@ -867,6 +1079,9 @@ def max_pool3d(
count_include_pad: bool = False,
layout: str = "NCDHW",
out_layout: str | None = None,
+ *,
+ ty=None,
+ span=None,
) -> Expr:
r"""3D maximum pooling operator.
@@ -926,16 +1141,38 @@ def max_pool3d(
if isinstance(padding, int):
padding = (padding, padding, padding, padding, padding, padding)
- return _ffi_api.max_pool3d( # type: ignore
- data,
- pool_size,
- strides,
- padding,
- dilation,
- ceil_mode,
- count_include_pad,
- layout,
- out_layout,
+ if len(padding) == 3:
+ padding = tuple(padding) + tuple(padding)
+ if len(strides) == 1:
+ strides = tuple(strides) * 3
+ if len(dilation) == 1:
+ dilation = tuple(dilation) * 3
+ if len(pool_size) == 1:
+ pool_size = tuple(pool_size) * 3
+ if len(strides) != 3:
+ raise _InternalError("strides must contain 3 values")
+ if len(padding) != 6:
+ raise _InternalError("padding must contain 6 values")
+ if len(dilation) != 3:
+ raise _InternalError("dilation must contain 3 values")
+ if len(pool_size) != 3:
+ raise _InternalError("pool_size must contain 3 values")
+ return _Call(
+ "relax.nn.max_pool3d",
+ [data],
+ attrs=_make_attrs(
+ "relax.attrs.Pool3DAttrs",
+ pool_size=pool_size,
+ strides=strides,
+ padding=padding,
+ dilation=dilation,
+ ceil_mode=ceil_mode,
+ count_include_pad=count_include_pad,
+ layout=layout,
+ out_layout=(layout if out_layout is None else out_layout),
+ ),
+ ty=ty,
+ span=span,
)
@@ -949,6 +1186,9 @@ def avg_pool1d(
count_include_pad: bool = True,
layout: str = "NCW",
out_layout: str | None = None,
+ *,
+ ty=None,
+ span=None,
) -> Expr:
r"""1D average pooling operator.
@@ -1006,16 +1246,32 @@ def avg_pool1d(
dilation = (dilation,)
if isinstance(padding, int):
padding = (padding, padding)
- return _ffi_api.avg_pool1d( # type: ignore
- data,
- pool_size,
- strides,
- padding,
- dilation,
- ceil_mode,
- count_include_pad,
- layout,
- out_layout,
+ if len(padding) == 1:
+ padding = tuple(padding) + tuple(padding)
+ if len(strides) != 1:
+ raise _InternalError("strides must contain 1 values")
+ if len(padding) != 2:
+ raise _InternalError("padding must contain 2 values")
+ if len(dilation) != 1:
+ raise _InternalError("dilation must contain 1 values")
+ if len(pool_size) != 1:
+ raise _InternalError("pool_size must contain 1 values")
+ return _Call(
+ "relax.nn.avg_pool1d",
+ [data],
+ attrs=_make_attrs(
+ "relax.attrs.Pool1DAttrs",
+ pool_size=pool_size,
+ strides=strides,
+ padding=padding,
+ dilation=dilation,
+ ceil_mode=ceil_mode,
+ count_include_pad=count_include_pad,
+ layout=layout,
+ out_layout=(layout if out_layout is None else out_layout),
+ ),
+ ty=ty,
+ span=span,
)
@@ -1029,6 +1285,9 @@ def avg_pool2d(
count_include_pad: bool = False,
layout: str = "NCHW",
out_layout: str | None = None,
+ *,
+ ty=None,
+ span=None,
) -> Expr:
r"""2D average pooling operator.
@@ -1094,16 +1353,38 @@ def avg_pool2d(
dilation = (dilation, dilation)
if isinstance(padding, int):
padding = (padding, padding, padding, padding)
- return _ffi_api.avg_pool2d( # type: ignore
- data,
- pool_size,
- strides,
- padding,
- dilation,
- ceil_mode,
- count_include_pad,
- layout,
- out_layout,
+ if len(padding) == 2:
+ padding = tuple(padding) + tuple(padding)
+ if len(strides) == 1:
+ strides = tuple(strides) * 2
+ if len(dilation) == 1:
+ dilation = tuple(dilation) * 2
+ if len(pool_size) == 1:
+ pool_size = tuple(pool_size) * 2
+ if len(strides) != 2:
+ raise _InternalError("strides must contain 2 values")
+ if len(padding) != 4:
+ raise _InternalError("padding must contain 4 values")
+ if len(dilation) != 2:
+ raise _InternalError("dilation must contain 2 values")
+ if len(pool_size) != 2:
+ raise _InternalError("pool_size must contain 2 values")
+ return _Call(
+ "relax.nn.avg_pool2d",
+ [data],
+ attrs=_make_attrs(
+ "relax.attrs.Pool2DAttrs",
+ pool_size=pool_size,
+ strides=strides,
+ padding=padding,
+ dilation=dilation,
+ ceil_mode=ceil_mode,
+ count_include_pad=count_include_pad,
+ layout=layout,
+ out_layout=(layout if out_layout is None else out_layout),
+ ),
+ ty=ty,
+ span=span,
)
@@ -1117,6 +1398,9 @@ def avg_pool3d(
count_include_pad: bool = True,
layout: str = "NCDHW",
out_layout: str | None = None,
+ *,
+ ty=None,
+ span=None,
) -> Expr:
r"""2D average pooling operator.
@@ -1176,16 +1460,38 @@ def avg_pool3d(
if isinstance(padding, int):
padding = (padding, padding, padding, padding, padding, padding)
- return _ffi_api.avg_pool3d( # type: ignore
- data,
- pool_size,
- strides,
- padding,
- dilation,
- ceil_mode,
- count_include_pad,
- layout,
- out_layout,
+ if len(padding) == 3:
+ padding = tuple(padding) + tuple(padding)
+ if len(strides) == 1:
+ strides = tuple(strides) * 3
+ if len(dilation) == 1:
+ dilation = tuple(dilation) * 3
+ if len(pool_size) == 1:
+ pool_size = tuple(pool_size) * 3
+ if len(strides) != 3:
+ raise _InternalError("strides must contain 3 values")
+ if len(padding) != 6:
+ raise _InternalError("padding must contain 6 values")
+ if len(dilation) != 3:
+ raise _InternalError("dilation must contain 3 values")
+ if len(pool_size) != 3:
+ raise _InternalError("pool_size must contain 3 values")
+ return _Call(
+ "relax.nn.avg_pool3d",
+ [data],
+ attrs=_make_attrs(
+ "relax.attrs.Pool3DAttrs",
+ pool_size=pool_size,
+ strides=strides,
+ padding=padding,
+ dilation=dilation,
+ ceil_mode=ceil_mode,
+ count_include_pad=count_include_pad,
+ layout=layout,
+ out_layout=(layout if out_layout is None else out_layout),
+ ),
+ ty=ty,
+ span=span,
)
@@ -1194,6 +1500,9 @@ def adaptive_avg_pool1d(
output_size: int | tuple[int] | None = None,
layout: str = "NCW",
out_layout: str | None = None,
+ *,
+ ty=None,
+ span=None,
) -> Expr:
r"""1D adaptive average pooling operator. This operator is experimental.
@@ -1239,7 +1548,20 @@ def adaptive_avg_pool1d(
"""
if isinstance(output_size, int):
output_size = (output_size,)
- return _ffi_api.adaptive_avg_pool1d(data, output_size, layout, out_layout) # type: ignore
+ if output_size is not None and len(output_size) != 1:
+ raise _InternalError("output_size must contain 1 values")
+ return _Call(
+ "relax.nn.adaptive_avg_pool1d",
+ [data],
+ attrs=_make_attrs(
+ "relax.attrs.AdaptivePool1DAttrs",
+ output_size=output_size,
+ layout=layout,
+ out_layout=(layout if out_layout is None else out_layout),
+ ),
+ ty=ty,
+ span=span,
+ ) # type: ignore
def adaptive_avg_pool2d(
@@ -1247,6 +1569,10 @@ def adaptive_avg_pool2d(
output_size: int | tuple[int, int] | None = None,
layout: str = "NCHW",
out_layout: str | None = None,
+ *,
+ ty_args=None,
+ ty=None,
+ span=None,
) -> Expr:
r"""2D adaptive average pooling operator. This operator is experimental.
@@ -1295,7 +1621,21 @@ def adaptive_avg_pool2d(
"""
if isinstance(output_size, int):
output_size = (output_size, output_size)
- return _ffi_api.adaptive_avg_pool2d(data, output_size, layout, out_layout) # type: ignore
+ if output_size is not None and len(output_size) != 2:
+ raise _InternalError("output_size must contain 2 values")
+ return _Call(
+ "relax.nn.adaptive_avg_pool2d",
+ [data],
+ attrs=_make_attrs(
+ "relax.attrs.AdaptivePool2DAttrs",
+ output_size=output_size,
+ layout=layout,
+ out_layout=(layout if out_layout is None else out_layout),
+ ),
+ ty_args=ty_args,
+ ty=ty,
+ span=span,
+ ) # type: ignore
def adaptive_avg_pool3d(
@@ -1303,6 +1643,9 @@ def adaptive_avg_pool3d(
output_size: int | tuple[int, int] | None = None,
layout: str = "NCDHW",
out_layout: str | None = None,
+ *,
+ ty=None,
+ span=None,
) -> Expr:
r"""3D adaptive average pooling operator. This operator is experimental.
@@ -1351,10 +1694,23 @@ def adaptive_avg_pool3d(
"""
if isinstance(output_size, int):
output_size = (output_size, output_size, output_size)
- return _ffi_api.adaptive_avg_pool3d(data, output_size, layout, out_layout) # type: ignore
+ if output_size is not None and len(output_size) != 3:
+ raise _InternalError("output_size must contain 3 values")
+ return _Call(
+ "relax.nn.adaptive_avg_pool3d",
+ [data],
+ attrs=_make_attrs(
+ "relax.attrs.AdaptivePool3DAttrs",
+ output_size=output_size,
+ layout=layout,
+ out_layout=(layout if out_layout is None else out_layout),
+ ),
+ ty=ty,
+ span=span,
+ ) # type: ignore
-def relu(data: Expr) -> Expr:
+def relu(data: Expr, *, ty_args=None, ty=None, span=None) -> Expr:
r"""Rectified linear unit.
.. math::
@@ -1370,7 +1726,7 @@ def relu(data: Expr) -> Expr:
result : relax.Expr
The computed result.
"""
- return _ffi_api.relu(data) # type: ignore
+ return _Call("relax.nn.relu", [data], ty_args=ty_args, ty=ty, span=span) # type: ignore
def relu6(data: Expr) -> Expr:
@@ -1392,7 +1748,7 @@ def relu6(data: Expr) -> Expr:
return relax.op.clip(data, 0, 6)
-def leakyrelu(data: Expr, alpha: float = 0.01) -> Expr:
+def leakyrelu(data: Expr, alpha: float = 0.01, *, ty_args=None, ty=None, span=None) -> Expr:
"""Rectified linear unit.
.. math::
@@ -1412,10 +1768,17 @@ def leakyrelu(data: Expr, alpha: float = 0.01) -> Expr:
result : relax.Expr
The computed result.
"""
- return _ffi_api.leakyrelu(data, alpha) # type: ignore
+ return _Call(
+ "relax.nn.leakyrelu",
+ [data],
+ attrs=_make_attrs("relax.attrs.LeakyReluAttrs", alpha=alpha),
+ ty_args=ty_args,
+ ty=ty,
+ span=span,
+ ) # type: ignore
-def gelu(data: Expr) -> Expr:
+def gelu(data: Expr, *, ty_args=None, ty=None, span=None) -> Expr:
r"""Gaussian Error Linear Units function
.. math::
@@ -1437,10 +1800,10 @@ def gelu(data: Expr) -> Expr:
----
The input tensor is required to have float dtype
"""
- return _ffi_api.gelu(data) # type: ignore
+ return _Call("relax.nn.gelu", [data], ty_args=ty_args, ty=ty, span=span) # type: ignore
-def gelu_tanh(data: Expr) -> Expr:
+def gelu_tanh(data: Expr, *, ty_args=None, ty=None, span=None) -> Expr:
r"""Gaussian Error Linear Units function with tanh approximation
.. math::
@@ -1460,10 +1823,10 @@ def gelu_tanh(data: Expr) -> Expr:
----
The input tensor is required to have float dtype
"""
- return _ffi_api.gelu_tanh(data) # type: ignore
+ return _Call("relax.nn.gelu_tanh", [data], ty_args=ty_args, ty=ty, span=span) # type: ignore
-def selu(data: Expr) -> Expr:
+def selu(data: Expr, *, ty_args=None, ty=None, span=None) -> Expr:
r"""Scaled Exponential Linear Unit (SELU).
.. math::
@@ -1484,10 +1847,10 @@ def selu(data: Expr) -> Expr:
result : relax.Expr
The computed result.
"""
- return _ffi_api.selu(data)
+ return _Call("relax.nn.selu", [data], ty_args=ty_args, ty=ty, span=span)
-def silu(data: Expr) -> Expr:
+def silu(data: Expr, *, ty_args=None, ty=None, span=None) -> Expr:
r"""Sigmoid Linear Unit function
.. math::
@@ -1507,10 +1870,10 @@ def silu(data: Expr) -> Expr:
----
The input tensor is required to have float dtype
"""
- return _ffi_api.silu(data) # type: ignore
+ return _Call("relax.nn.silu", [data], ty_args=ty_args, ty=ty, span=span) # type: ignore
-def softmax(data: Expr, axis: int = -1) -> Expr:
+def softmax(data: Expr, axis: int = -1, *, ty=None, span=None) -> Expr:
r"""Computes softmax.
.. math:: \text{softmax}(x)_i = \frac{\exp(x_i)}{\sum_j \exp(x_j)}
@@ -1534,10 +1897,24 @@ def softmax(data: Expr, axis: int = -1) -> Expr:
----
The input tensor is required to have float dtype
"""
- return _ffi_api.softmax(data, axis) # type: ignore
+ return _Call(
+ "relax.nn.softmax",
+ [data],
+ attrs=_make_attrs("relax.attrs.SoftmaxAttrs", axis=axis),
+ ty=ty,
+ span=span,
+ ) # type: ignore
-def softplus(data: Expr, beta: float = 1.0, threshold: float = 20.0) -> Expr:
+def softplus(
+ data: Expr,
+ beta: float = 1.0,
+ threshold: float = 20.0,
+ *,
+ ty_args=None,
+ ty=None,
+ span=None,
+) -> Expr:
r"""Softplus activation function.
.. math:: \text{Softplus}(x) = \frac{1}{\beta} \log(1 + e^{\beta x})
@@ -1559,10 +1936,17 @@ def softplus(data: Expr, beta: float = 1.0, threshold: float = 20.0) -> Expr:
result : relax.Expr
The computed result.
"""
- return _ffi_api.softplus(data, beta, threshold)
+ return _Call(
+ "relax.nn.softplus",
+ [data],
+ attrs=_make_attrs("relax.attrs.SoftplusAttrs", beta=beta, threshold=threshold),
+ ty_args=ty_args,
+ ty=ty,
+ span=span,
+ )
-def log_softmax(data: Expr, axis: int = -1) -> Expr:
+def log_softmax(data: Expr, axis: int = -1, *, ty=None, span=None) -> Expr:
r"""Computes log softmax.
.. math::
@@ -1587,10 +1971,16 @@ def log_softmax(data: Expr, axis: int = -1) -> Expr:
result : relax.Expr
The computed result.
"""
- return _ffi_api.log_softmax(data, axis) # type: ignore
+ return _Call(
+ "relax.nn.log_softmax",
+ [data],
+ attrs=_make_attrs("relax.attrs.SoftmaxAttrs", axis=axis),
+ ty=ty,
+ span=span,
+ ) # type: ignore
-def prelu(data: Expr, alpha: Expr, axis: int = 1) -> Expr:
+def prelu(data: Expr, alpha: Expr, axis: int = 1, *, ty=None, span=None) -> Expr:
r"""Parametric Rectified Linear Unit (PReLU).
.. math::
@@ -1613,7 +2003,13 @@ def prelu(data: Expr, alpha: Expr, axis: int = 1) -> Expr:
result : relax.Expr
The computed result.
"""
- return _ffi_api.prelu(data, alpha, axis)
+ return _Call(
+ "relax.nn.prelu",
+ [data, alpha],
+ attrs=_make_attrs("relax.attrs.PReluAttrs", axis=axis),
+ ty=ty,
+ span=span,
+ )
def batch_norm(
@@ -1628,6 +2024,9 @@ def batch_norm(
scale: bool = True,
momentum: float = 0.1,
training: bool = True,
+ *,
+ ty=None,
+ span=None,
) -> Expr:
r"""
Batch normalization layer (Ioffe and Szegedy, 2014).
@@ -1726,18 +2125,20 @@ def batch_norm(
result : relax.Expr
The computed result.
"""
- return _ffi_api.batch_norm( # type: ignore
- data,
- gamma,
- beta,
- moving_mean,
- moving_var,
- axis,
- epsilon,
- center,
- scale,
- momentum,
- training,
+ return _Call(
+ "relax.nn.batch_norm",
+ [data, gamma, beta, moving_mean, moving_var],
+ attrs=_make_attrs(
+ "relax.attrs.BatchNormAttrs",
+ axis=axis,
+ epsilon=epsilon,
+ center=center,
+ scale=scale,
+ momentum=momentum,
+ training=training,
+ ),
+ ty=ty,
+ span=span,
)
@@ -1749,6 +2150,9 @@ def layer_norm(
epsilon: float = 1e-5,
center: bool = True,
scale: bool = True,
+ *,
+ ty=None,
+ span=None,
) -> Expr:
r"""
Layer normalization (Lei Ba and et al., 2016).
@@ -1799,7 +2203,15 @@ def layer_norm(
"""
if isinstance(axes, int):
axes = [axes]
- return _ffi_api.layer_norm(data, gamma, beta, axes, epsilon, center, scale) # type: ignore
+ return _Call(
+ "relax.nn.layer_norm",
+ [data, gamma, beta],
+ attrs=_make_attrs(
+ "relax.attrs.LayerNormAttrs", axes=axes, epsilon=epsilon, center=center, scale=scale
+ ),
+ ty=ty,
+ span=span,
+ ) # type: ignore
def group_norm(
@@ -1812,6 +2224,9 @@ def group_norm(
epsilon: float = 1e-5,
center: bool = True,
scale: bool = True,
+ *,
+ ty=None,
+ span=None,
) -> Expr:
r"""
Group normalization (Yuxin Wu and et al., 2016).
@@ -1855,8 +2270,20 @@ def group_norm(
"""
if isinstance(axes, int):
axes = [axes]
- return _ffi_api.group_norm( # type: ignore
- data, gamma, beta, num_groups, channel_axis, axes, epsilon, center, scale
+ return _Call(
+ "relax.nn.group_norm",
+ [data, gamma, beta],
+ attrs=_make_attrs(
+ "relax.attrs.GroupNormAttrs",
+ num_groups=num_groups,
+ channel_axis=channel_axis,
+ axes=axes,
+ epsilon=epsilon,
+ center=center,
+ scale=scale,
+ ),
+ ty=ty,
+ span=span,
)
@@ -1869,6 +2296,9 @@ def instance_norm(
epsilon: float = 1e-5,
center: bool = True,
scale: bool = True,
+ *,
+ ty=None,
+ span=None,
) -> Expr:
r"""
Instance normalization
@@ -1903,15 +2333,19 @@ def instance_norm(
"""
if isinstance(axes, int):
axes = [axes]
- return _ffi_api.instance_norm( # type: ignore
- data,
- gamma,
- beta,
- channel_axis,
- axes,
- epsilon,
- center,
- scale,
+ return _Call(
+ "relax.nn.instance_norm",
+ [data, gamma, beta],
+ attrs=_make_attrs(
+ "relax.attrs.InstanceNormAttrs",
+ channel_axis=channel_axis,
+ axes=axes,
+ epsilon=epsilon,
+ center=center,
+ scale=scale,
+ ),
+ ty=ty,
+ span=span,
)
@@ -1920,6 +2354,9 @@ def rms_norm(
weight: Expr,
axes: int | list[int] = -1,
epsilon: float = 1e-5,
+ *,
+ ty=None,
+ span=None,
) -> Expr:
r"""
Root mean square normalization (Biao Zhang and et al., 2019).
@@ -1952,10 +2389,16 @@ def rms_norm(
"""
if isinstance(axes, int):
axes = [axes]
- return _ffi_api.rms_norm(data, weight, axes, epsilon) # type: ignore
+ return _Call(
+ "relax.nn.rms_norm",
+ [data, weight],
+ attrs=_make_attrs("relax.attrs.RMSNormAttrs", axes=axes, epsilon=epsilon),
+ ty=ty,
+ span=span,
+ ) # type: ignore
-def dropout(data: Expr, rate: float = 0.5) -> Expr:
+def dropout(data: Expr, rate: float = 0.5, *, ty=None, span=None) -> Expr:
"""Applies the dropout operation to the input tensor.
During training, each element of the input is set to zero with
@@ -1977,10 +2420,18 @@ def dropout(data: Expr, rate: float = 0.5) -> Expr:
The first one is the original tensor and the second one is a
mask tensor (1.0 where element not dropped, 0.0 where dropped)
"""
- return _ffi_api.dropout(data, rate) # type: ignore
+ return _Call(
+ "relax.nn.dropout",
+ [data],
+ attrs=_make_attrs("relax.attrs.DropoutAttrs", rate=rate),
+ ty=ty,
+ span=span,
+ ) # type: ignore
-def cross_entropy_with_logits(predictions: Expr, labels: Expr) -> Expr:
+def cross_entropy_with_logits(
+ predictions: Expr, labels: Expr, *, ty_args=None, ty=None, span=None
+) -> Expr:
r"""CrossEntropy with logits between the predictions and labels.
The shape of predictions and labels must be the same. And when ndim >= 2,
@@ -2004,7 +2455,13 @@ def cross_entropy_with_logits(predictions: Expr, labels: Expr) -> Expr:
result : relax.Expr
The computed result.
"""
- return _ffi_api.cross_entropy_with_logits(predictions, labels) # type: ignore
+ return _Call(
+ "relax.nn.cross_entropy_with_logits",
+ [predictions, labels],
+ ty_args=ty_args,
+ ty=ty,
+ span=span,
+ ) # type: ignore
def nll_loss(
@@ -2013,6 +2470,10 @@ def nll_loss(
weights: Expr | None = None,
reduction: str = "mean",
ignore_index: int = -100,
+ *,
+ ty_args=None,
+ ty=None,
+ span=None,
) -> Expr:
"""Negative log likelihood loss.
@@ -2049,7 +2510,18 @@ def nll_loss(
result : relax.Expr
The computed result.
"""
- return _ffi_api.nll_loss(predictions, targets, weights, reduction, ignore_index) # type: ignore
+ if reduction not in ("none", "sum", "mean"):
+ raise _InternalError("reduction must be none, sum, or mean")
+ return _Call(
+ "relax.nn.nll_loss",
+ [predictions, targets, *([] if weights is None else [weights])],
+ attrs=_make_attrs(
+ "relax.attrs.NLLLossAttrs", reduction=reduction, ignore_index=ignore_index
+ ),
+ ty_args=ty_args,
+ ty=ty,
+ span=span,
+ ) # type: ignore
def attention(
@@ -2060,6 +2532,9 @@ def attention(
scale: FloatImm | None = None,
causal_mask: str | None = None,
window_size: int | None = None,
+ *,
+ ty=None,
+ span=None,
) -> Expr:
r"""Computes fused multi head attention.
@@ -2140,7 +2615,18 @@ def attention(
The computed result. The layout of the output should be
(batch_size, seq_len, num_head, head_dim_v).
"""
- return _ffi_api.attention(query, key, value, bias, scale, causal_mask, window_size) # type: ignore
+ return _Call(
+ "relax.nn.attention_bias" if bias is not None else "relax.nn.attention",
+ [query, key, value, *([] if bias is None else [bias])],
+ attrs=_make_attrs(
+ "relax.attrs.AttentionAttrs",
+ scale=scale,
+ causal_mask=causal_mask,
+ window_size=window_size,
+ ),
+ ty=ty,
+ span=span,
+ ) # type: ignore
def attention_bias(
@@ -2151,6 +2637,9 @@ def attention_bias(
scale: FloatImm | None = None,
causal_mask: str | None = None,
window_size: int | None = None,
+ *,
+ ty=None,
+ span=None,
) -> Expr:
r"""Computes fused multi head attention.
@@ -2235,7 +2724,18 @@ def attention_bias(
The computed result. The layout of the output should be
(batch_size, seq_len, num_head, head_dim_v).
"""
- return _ffi_api.attention(query, key, value, bias, scale, causal_mask, window_size) # type: ignore
+ return _Call(
+ "relax.nn.attention_bias",
+ [query, key, value, bias],
+ attrs=_make_attrs(
+ "relax.attrs.AttentionAttrs",
+ scale=scale,
+ causal_mask=causal_mask,
+ window_size=window_size,
+ ),
+ ty=ty,
+ span=span,
+ ) # type: ignore
def attention_var_len(
@@ -2249,6 +2749,9 @@ def attention_var_len(
scale: FloatImm | None = None,
causal_mask: str | None = None,
window_size: int | None = None,
+ *,
+ ty=None,
+ span=None,
) -> Expr:
"""Computes fused multi head attention over batched sequences of variable lengths.
@@ -2339,21 +2842,21 @@ def attention_var_len(
seqstart_k = seqstart_q
if max_seqlen_k is None:
max_seqlen_k = max_seqlen_q
- return _ffi_api.attention_var_len(
- queries,
- keys,
- values,
- seqstart_q,
- seqstart_k,
- max_seqlen_q,
- max_seqlen_k,
- scale,
- causal_mask,
- window_size,
+ return _Call(
+ "relax.nn.attention_var_len",
+ [queries, keys, values, seqstart_q, seqstart_k, max_seqlen_q, max_seqlen_k],
+ attrs=_make_attrs(
+ "relax.attrs.AttentionAttrs",
+ scale=scale,
+ causal_mask=causal_mask,
+ window_size=window_size,
+ ),
+ ty=ty,
+ span=span,
) # type: ignore
-def batch_flatten(data: Expr) -> Expr:
+def batch_flatten(data: Expr, *, ty=None, span=None) -> Expr:
"""Flatten all dimensions except the first (batch) dimension.
This operation flattens a tensor of shape `(N, C, H, W, ...)` into
@@ -2369,4 +2872,4 @@ def batch_flatten(data: Expr) -> Expr:
result : relax.Expr
The flattened result with shape `(batch_size, flattened_features)`.
"""
- return _ffi_api.batch_flatten(data) # type: ignore
+ return _Call("relax.nn.batch_flatten", [data], ty=ty, span=span) # type: ignore
diff --git a/python/tvm/relax/op/qdq.py b/python/tvm/relax/op/qdq.py
index 6231b5f9b9..a883ec52d7 100644
--- a/python/tvm/relax/op/qdq.py
+++ b/python/tvm/relax/op/qdq.py
@@ -16,11 +16,23 @@
# under the License.
"""Relax quantize/dequantize operators"""
+from tvm.ir import Call as _Call
+from tvm.ir import PrimType as _PrimType
+from tvm.ir.attrs import make_node as _make_attrs
+
from ..expr import Expr
-from . import _ffi_api
-def quantize(data: Expr, scale: Expr, zero_point: Expr, axis: int = -1, out_dtype: str = "int8"):
+def quantize(
+ data: Expr,
+ scale: Expr,
+ zero_point: Expr,
+ axis: int = -1,
+ out_dtype: str = "int8",
+ *,
+ ty=None,
+ span=None,
+):
r"""Quantize op
This operator takes input and produces quantized output. The input tensor can be of any shape.
The output shape is the same as input shape.
@@ -50,11 +62,28 @@ def quantize(data: Expr, scale: Expr, zero_point: Expr, axis: int = -1, out_dtyp
The computed result.
"""
- return _ffi_api.quantize(data, scale, zero_point, axis, out_dtype)
+ return _Call(
+ "relax.quantize",
+ [data, scale, zero_point],
+ attrs=_make_attrs(
+ "relax.attrs.QuantizeAttrs",
+ axis=axis,
+ out_dtype=(out_dtype.dtype if isinstance(out_dtype, _PrimType) else out_dtype),
+ ),
+ ty=ty,
+ span=span,
+ )
def dequantize(
- data: Expr, scale: Expr, zero_point: Expr, axis: int = -1, out_dtype: str = "float32"
+ data: Expr,
+ scale: Expr,
+ zero_point: Expr,
+ axis: int = -1,
+ out_dtype: str = "float32",
+ *,
+ ty=None,
+ span=None,
):
r"""Dequantize op
This operator takes input and produces dequantized output. The input tensor can be of any shape.
@@ -85,4 +114,14 @@ def dequantize(
The computed result.
"""
- return _ffi_api.dequantize(data, scale, zero_point, axis, out_dtype)
+ return _Call(
+ "relax.dequantize",
+ [data, scale, zero_point],
+ attrs=_make_attrs(
+ "relax.attrs.QuantizeAttrs",
+ axis=axis,
+ out_dtype=(out_dtype.dtype if isinstance(out_dtype, _PrimType) else out_dtype),
+ ),
+ ty=ty,
+ span=span,
+ )
diff --git a/python/tvm/relax/op/sampling.py b/python/tvm/relax/op/sampling.py
index cd4cad9253..4f973676ae 100644
--- a/python/tvm/relax/op/sampling.py
+++ b/python/tvm/relax/op/sampling.py
@@ -16,8 +16,10 @@
# under the License.
"""Sampling operators."""
+from tvm.ir import Call as _Call
+from tvm.ir.attrs import make_node as _make_attrs
+
from ..expr import Expr
-from . import _ffi_api
def multinomial_from_uniform(
@@ -25,6 +27,9 @@ def multinomial_from_uniform(
uniform_sample: Expr,
sample_indices: Expr,
dtype: str = "int64",
+ *,
+ ty=None,
+ span=None,
) -> Expr:
"""Returns a tensor where each row contains the index sampled from the multinomial
probability distribution located in the corresponding row of tensor prob.
@@ -77,9 +82,10 @@ def multinomial_from_uniform(
"""
- return _ffi_api.multinomial_from_uniform( # type: ignore
- prob,
- uniform_sample,
- sample_indices,
- dtype,
+ return _Call(
+ "relax.multinomial_from_uniform",
+ [prob, uniform_sample, sample_indices],
+ attrs=_make_attrs("relax.attrs.MultinomialFromUniformAttrs", dtype=dtype),
+ ty=ty,
+ span=span,
)
diff --git a/python/tvm/relax/op/search.py b/python/tvm/relax/op/search.py
index 8428400a45..85ce91db2d 100644
--- a/python/tvm/relax/op/search.py
+++ b/python/tvm/relax/op/search.py
@@ -17,11 +17,13 @@
# pylint: disable=invalid-name
"""Search operators."""
+from tvm.ir import Call as _Call
+from tvm.ir.attrs import make_node as _make_attrs
+
from ..expr import Expr
-from . import _ffi_api
-def where(condition: Expr, x1: Expr, x2: Expr) -> Expr:
+def where(condition: Expr, x1: Expr, x2: Expr, *, ty=None, span=None) -> Expr:
"""Selecting elements from either the input tensors depending on the value of the
condition.
@@ -48,10 +50,17 @@ def where(condition: Expr, x1: Expr, x2: Expr) -> Expr:
result : relax.Expr
The result tensor.
"""
- return _ffi_api.where(condition, x1, x2) # type: ignore
+ return _Call("relax.where", [condition, x1, x2], ty=ty, span=span) # type: ignore
-def argmax(x: Expr, axis: int | None = None, keepdims: bool = False) -> Expr:
+def argmax(
+ x: Expr,
+ axis: int | None = None,
+ keepdims: bool = False,
+ *,
+ ty=None,
+ span=None,
+) -> Expr:
"""Computes the argmax of tensor elements over given axis.
Parameters
@@ -74,10 +83,23 @@ def argmax(x: Expr, axis: int | None = None, keepdims: bool = False) -> Expr:
result : relax.Expr
The computed result.
"""
- return _ffi_api.argmax(x, axis, keepdims) # type: ignore
-
-
-def argmin(x: Expr, axis: int | None = None, keepdims: bool = False) -> Expr:
+ return _Call(
+ "relax.argmax",
+ [x],
+ attrs=_make_attrs("relax.attrs.ArgmaxArgminAttrs", axis=axis, keepdims=keepdims),
+ ty=ty,
+ span=span,
+ ) # type: ignore
+
+
+def argmin(
+ x: Expr,
+ axis: int | None = None,
+ keepdims: bool = False,
+ *,
+ ty=None,
+ span=None,
+) -> Expr:
"""Computes the argmin of tensor elements over given axis.
Parameters
@@ -100,10 +122,24 @@ def argmin(x: Expr, axis: int | None = None, keepdims: bool = False) -> Expr:
result : relax.Expr
The computed result.
"""
- return _ffi_api.argmin(x, axis, keepdims) # type: ignore
-
-
-def bucketize(input_tensor, boundaries, out_int32=False, right=False):
+ return _Call(
+ "relax.argmin",
+ [x],
+ attrs=_make_attrs("relax.attrs.ArgmaxArgminAttrs", axis=axis, keepdims=keepdims),
+ ty=ty,
+ span=span,
+ ) # type: ignore
+
+
+def bucketize(
+ input_tensor,
+ boundaries,
+ out_int32=False,
+ right=False,
+ *,
+ ty=None,
+ span=None,
+):
"""Returns the indices of the buckets to which each value in the input belongs.
Parameters
@@ -125,4 +161,10 @@ def bucketize(input_tensor, boundaries, out_int32=False, right=False):
result : relax.Expr
The computed result with same shape as input_tensor.
"""
- return _ffi_api.bucketize(input_tensor, boundaries, out_int32, right)
+ return _Call(
+ "relax.bucketize",
+ [input_tensor, boundaries],
+ attrs=_make_attrs("relax.attrs.BucketizeAttrs", out_int32=out_int32, right=right),
+ ty=ty,
+ span=span,
+ )
diff --git a/python/tvm/relax/op/set.py b/python/tvm/relax/op/set.py
index 28979d58d2..59e3ead1fe 100644
--- a/python/tvm/relax/op/set.py
+++ b/python/tvm/relax/op/set.py
@@ -20,9 +20,9 @@
import numpy as np # type: ignore
import tvm
+from tvm.ir import Call as _Call
from ..expr import Expr, prim_value
-from . import _ffi_api
def unique(
@@ -32,6 +32,9 @@ def unique(
return_inverse: bool | Expr = False,
return_counts: bool | Expr = False,
axis: int | Expr | None = None,
+ *,
+ ty=None,
+ span=None,
) -> Expr:
"""Find the unique elements in a given tensor.
In addition, it optionally returns
@@ -79,8 +82,11 @@ def unique(
return_counts = prim_value(return_counts)
if axis is not None and isinstance(axis, int):
axis = prim_value(axis)
- return _ffi_api.unique( # type: ignore
- x, sorted, return_index, return_inverse, return_counts, axis
+ return _Call(
+ "relax.unique",
+ [x, sorted, return_index, return_inverse, return_counts, *([] if axis is None else [axis])],
+ ty=ty,
+ span=span,
)
@@ -179,7 +185,7 @@ def numpy_unique(
return tuple(output_list)
-def nonzero(x: Expr) -> Expr:
+def nonzero(x: Expr, *, ty=None, span=None) -> Expr:
"""Find the indices of elements of a tensor that are non-zero.
Parameters
@@ -207,7 +213,7 @@ def nonzero(x: Expr) -> Expr:
[1, 0]]
"""
- return _ffi_api.nonzero(x) # type: ignore
+ return _Call("relax.nonzero", [x], ty=ty, span=span) # type: ignore
@tvm.register_global_func("relax.run.nonzero")
diff --git a/python/tvm/relax/op/sorting.py b/python/tvm/relax/op/sorting.py
index 36919540e5..f892a43947 100644
--- a/python/tvm/relax/op/sorting.py
+++ b/python/tvm/relax/op/sorting.py
@@ -16,13 +16,21 @@
# under the License.
"""Sortings operators."""
+from tvm.ir import Call as _Call
from tvm.ir import GenericConst
+from tvm.ir.attrs import make_node as _make_attrs
from ..expr import Expr
-from . import _ffi_api
-def sort(x: Expr, axis: int = -1, descending: bool = False):
+def sort(
+ x: Expr,
+ axis: int = -1,
+ descending: bool = False,
+ *,
+ ty=None,
+ span=None,
+):
"""Performs sorting along the given axis and returns an array
in sorted order.
@@ -44,10 +52,24 @@ def sort(x: Expr, axis: int = -1, descending: bool = False):
Sorted tensor.
"""
- return _ffi_api.sort(x, axis, descending) # type: ignore
+ return _Call(
+ "relax.sort",
+ [x],
+ attrs=_make_attrs("relax.attrs.SortAttrs", axis=axis, descending=descending),
+ ty=ty,
+ span=span,
+ ) # type: ignore
-def argsort(data: Expr, axis: int = -1, descending: bool = False, dtype: str = "int32"):
+def argsort(
+ data: Expr,
+ axis: int = -1,
+ descending: bool = False,
+ dtype: str = "int32",
+ *,
+ ty=None,
+ span=None,
+):
"""Performs sorting along the given axis and returns an array of indices
having same shape as an input array that index data in sorted order.
@@ -70,7 +92,15 @@ def argsort(data: Expr, axis: int = -1, descending: bool = False, dtype: str = "
out : relax.Expr
Tensor with same shape as data.
"""
- return _ffi_api.argsort(data, axis, descending, dtype) # type: ignore
+ return _Call(
+ "relax.argsort",
+ [data],
+ attrs=_make_attrs(
+ "relax.attrs.ArgsortAttrs", axis=axis, descending=descending, dtype=dtype
+ ),
+ ty=ty,
+ span=span,
+ ) # type: ignore
def topk(
@@ -80,6 +110,9 @@ def topk(
ret_type: str = "both",
largest: bool = True,
dtype: str = "int32",
+ *,
+ ty=None,
+ span=None,
):
"""Get the top k elements in an input tensor along the given axis.
@@ -116,4 +149,12 @@ def topk(
"""
if isinstance(k, GenericConst):
k = k.value.numpy().item()
- return _ffi_api.topk(data, k, axis, ret_type, largest, dtype) # type: ignore
+ return _Call(
+ "relax.topk",
+ [data],
+ attrs=_make_attrs(
+ "relax.attrs.TopKAttrs", k=k, axis=axis, ret_type=ret_type, largest=largest, dtype=dtype
+ ),
+ ty=ty,
+ span=span,
+ ) # type: ignore
diff --git a/python/tvm/relax/op/statistical.py b/python/tvm/relax/op/statistical.py
index 051de2ce8c..80fa0ea115 100644
--- a/python/tvm/relax/op/statistical.py
+++ b/python/tvm/relax/op/statistical.py
@@ -18,17 +18,25 @@
"""Statistical operators."""
from tvm import DataType
+from tvm.ir import Call as _Call
from tvm.ir import PrimType
+from tvm.ir.attrs import make_node as _make_attrs
from ..expr import Expr
-from . import _ffi_api
def _raw_dtype(dtype):
return dtype.dtype if isinstance(dtype, PrimType) else dtype
-def max(x: Expr, axis: int | list[int] | None = None, keepdims: bool = False) -> Expr:
+def max(
+ x: Expr,
+ axis: int | list[int] | None = None,
+ keepdims: bool = False,
+ *,
+ ty=None,
+ span=None,
+) -> Expr:
"""Computes the max of tensor elements over given axes.
Parameters
@@ -53,10 +61,23 @@ def max(x: Expr, axis: int | list[int] | None = None, keepdims: bool = False) ->
"""
if isinstance(axis, int):
axis = [axis]
- return _ffi_api.max(x, axis, keepdims) # type: ignore
-
-
-def mean(x: Expr, axis: int | list[int] | None = None, keepdims: bool = False) -> Expr:
+ return _Call(
+ "relax.max",
+ [x],
+ attrs=_make_attrs("relax.attrs.StatisticalAttrs", axis=axis, keepdims=keepdims),
+ ty=ty,
+ span=span,
+ ) # type: ignore
+
+
+def mean(
+ x: Expr,
+ axis: int | list[int] | None = None,
+ keepdims: bool = False,
+ *,
+ ty=None,
+ span=None,
+) -> Expr:
"""Computes the mean of tensor elements over given axes.
Parameters
@@ -81,10 +102,23 @@ def mean(x: Expr, axis: int | list[int] | None = None, keepdims: bool = False) -
"""
if isinstance(axis, int):
axis = [axis]
- return _ffi_api.mean(x, axis, keepdims) # type: ignore
-
-
-def min(x: Expr, axis: int | list[int] | None = None, keepdims: bool = False) -> Expr:
+ return _Call(
+ "relax.mean",
+ [x],
+ attrs=_make_attrs("relax.attrs.StatisticalAttrs", axis=axis, keepdims=keepdims),
+ ty=ty,
+ span=span,
+ ) # type: ignore
+
+
+def min(
+ x: Expr,
+ axis: int | list[int] | None = None,
+ keepdims: bool = False,
+ *,
+ ty=None,
+ span=None,
+) -> Expr:
"""Computes the min of tensor elements over given axes.
Parameters
@@ -109,10 +143,23 @@ def min(x: Expr, axis: int | list[int] | None = None, keepdims: bool = False) ->
"""
if isinstance(axis, int):
axis = [axis]
- return _ffi_api.min(x, axis, keepdims) # type: ignore
-
-
-def prod(x: Expr, axis: int | list[int] | None = None, keepdims: bool = False) -> Expr:
+ return _Call(
+ "relax.min",
+ [x],
+ attrs=_make_attrs("relax.attrs.StatisticalAttrs", axis=axis, keepdims=keepdims),
+ ty=ty,
+ span=span,
+ ) # type: ignore
+
+
+def prod(
+ x: Expr,
+ axis: int | list[int] | None = None,
+ keepdims: bool = False,
+ *,
+ ty=None,
+ span=None,
+) -> Expr:
"""Computes the product of tensor elements over given axes.
Parameters
@@ -137,10 +184,23 @@ def prod(x: Expr, axis: int | list[int] | None = None, keepdims: bool = False) -
"""
if isinstance(axis, int):
axis = [axis]
- return _ffi_api.prod(x, axis, keepdims) # type: ignore
-
-
-def std(x: Expr, axis: int | list[int] | None = None, keepdims: bool = False) -> Expr:
+ return _Call(
+ "relax.prod",
+ [x],
+ attrs=_make_attrs("relax.attrs.StatisticalAttrs", axis=axis, keepdims=keepdims),
+ ty=ty,
+ span=span,
+ ) # type: ignore
+
+
+def std(
+ x: Expr,
+ axis: int | list[int] | None = None,
+ keepdims: bool = False,
+ *,
+ ty=None,
+ span=None,
+) -> Expr:
"""Computes the standard deviation of tensor elements over given axes.
Parameters
@@ -165,10 +225,23 @@ def std(x: Expr, axis: int | list[int] | None = None, keepdims: bool = False) ->
"""
if isinstance(axis, int):
axis = [axis]
- return _ffi_api.std(x, axis, keepdims) # type: ignore
-
-
-def sum(x: Expr, axis: int | list[int] | None = None, keepdims: bool = False) -> Expr:
+ return _Call(
+ "relax.std",
+ [x],
+ attrs=_make_attrs("relax.attrs.StatisticalAttrs", axis=axis, keepdims=keepdims),
+ ty=ty,
+ span=span,
+ ) # type: ignore
+
+
+def sum(
+ x: Expr,
+ axis: int | list[int] | None = None,
+ keepdims: bool = False,
+ *,
+ ty=None,
+ span=None,
+) -> Expr:
"""Computes the sum of tensor elements over given axes.
Parameters
@@ -193,7 +266,13 @@ def sum(x: Expr, axis: int | list[int] | None = None, keepdims: bool = False) ->
"""
if isinstance(axis, int):
axis = [axis]
- return _ffi_api.sum(x, axis, keepdims) # type: ignore
+ return _Call(
+ "relax.sum",
+ [x],
+ attrs=_make_attrs("relax.attrs.StatisticalAttrs", axis=axis, keepdims=keepdims),
+ ty=ty,
+ span=span,
+ ) # type: ignore
def cumprod(
@@ -201,6 +280,9 @@ def cumprod(
axis: int | None = None,
dtype: str | DataType | None = None,
exclusive: bool = False,
+ *,
+ ty=None,
+ span=None,
):
"""Numpy style cumprod op. Return the cumulative product of the elements along
a given axis.
@@ -255,7 +337,15 @@ def cumprod(
if exclusive is None:
exclusive = False
- return _ffi_api.cumprod(data, axis, _raw_dtype(dtype), exclusive) # type: ignore
+ return _Call(
+ "relax.cumprod",
+ [data],
+ attrs=_make_attrs(
+ "relax.attrs.ScanopAttrs", axis=axis, dtype=_raw_dtype(dtype), exclusive=exclusive
+ ),
+ ty=ty,
+ span=span,
+ ) # type: ignore
def cumsum(
@@ -263,6 +353,9 @@ def cumsum(
axis: int | None = None,
dtype: str | DataType | None = None,
exclusive: bool = False,
+ *,
+ ty=None,
+ span=None,
):
"""Numpy style cumsum op. Return the cumulative inclusive sum of the elements along
a given axis.
@@ -317,10 +410,25 @@ def cumsum(
if exclusive is None:
exclusive = False
- return _ffi_api.cumsum(data, axis, _raw_dtype(dtype), exclusive) # type: ignore
-
-
-def variance(x: Expr, axis: int | list[int] | None = None, keepdims: bool = False) -> Expr:
+ return _Call(
+ "relax.cumsum",
+ [data],
+ attrs=_make_attrs(
+ "relax.attrs.ScanopAttrs", axis=axis, dtype=_raw_dtype(dtype), exclusive=exclusive
+ ),
+ ty=ty,
+ span=span,
+ ) # type: ignore
+
+
+def variance(
+ x: Expr,
+ axis: int | list[int] | None = None,
+ keepdims: bool = False,
+ *,
+ ty=None,
+ span=None,
+) -> Expr:
"""Computes the variance of tensor elements over given axes.
Parameters
@@ -345,10 +453,23 @@ def variance(x: Expr, axis: int | list[int] | None = None, keepdims: bool = Fals
"""
if isinstance(axis, int):
axis = [axis]
- return _ffi_api.variance(x, axis, keepdims) # type: ignore
-
-
-def median(x: Expr, axis: int | list[int] | None = None, keepdims: bool = False) -> Expr:
+ return _Call(
+ "relax.variance",
+ [x],
+ attrs=_make_attrs("relax.attrs.StatisticalAttrs", axis=axis, keepdims=keepdims),
+ ty=ty,
+ span=span,
+ ) # type: ignore
+
+
+def median(
+ x: Expr,
+ axis: int | list[int] | None = None,
+ keepdims: bool = False,
+ *,
+ ty=None,
+ span=None,
+) -> Expr:
"""Computes the median of tensor elements over given axes.
Parameters
@@ -372,4 +493,10 @@ def median(x: Expr, axis: int | list[int] | None = None, keepdims: bool = False)
"""
if isinstance(axis, int):
axis = [axis]
- return _ffi_api.median(x, axis, keepdims) # type: ignore
+ return _Call(
+ "relax.median",
+ [x],
+ attrs=_make_attrs("relax.attrs.StatisticalAttrs", axis=axis, keepdims=keepdims),
+ ty=ty,
+ span=span,
+ ) # type: ignore
diff --git a/python/tvm/relax/op/ternary.py b/python/tvm/relax/op/ternary.py
index c576b3123d..5a5c258d8d 100644
--- a/python/tvm/relax/op/ternary.py
+++ b/python/tvm/relax/op/ternary.py
@@ -17,11 +17,12 @@
# pylint: disable=redefined-builtin, invalid-name
"""Relax ternary arithmetic operators."""
+from tvm.ir import Call as _Call
+
from ..expr import Expr
-from . import _ffi_api
-def ewise_fma(x1: Expr, x2: Expr, x3: Expr) -> Expr:
+def ewise_fma(x1: Expr, x2: Expr, x3: Expr, *, ty_args=None, ty=None, span=None) -> Expr:
"""Elementwise fused multiply-add operator
Returns elementwise result of :math:`x1 * x2 + x3`
@@ -41,4 +42,4 @@ def ewise_fma(x1: Expr, x2: Expr, x3: Expr) -> Expr:
result : relax.Expr
The computed result.
"""
- return _ffi_api.ewise_fma(x1, x2, x3) # type: ignore
+ return _Call("relax.ewise_fma", [x1, x2, x3], ty_args=ty_args, ty=ty, span=span) # type: ignore
diff --git a/python/tvm/relax/op/unary.py b/python/tvm/relax/op/unary.py
index e7b800b5c9..ed9cb0f543 100644
--- a/python/tvm/relax/op/unary.py
+++ b/python/tvm/relax/op/unary.py
@@ -17,14 +17,15 @@
# pylint: disable=redefined-builtin, invalid-name
"""Relax unary arithmetic operators."""
+from tvm.ir import Call as _Call
+
from ..expr import Expr
from ..utils import convert_to_expr
-from . import _ffi_api
###################### Arithmetic operators ######################
-def abs(x: Expr) -> Expr:
+def abs(x: Expr, *, ty_args=None, ty=None, span=None) -> Expr:
"""Compute element-wise absolute value of the input data.
Parameters
@@ -37,10 +38,10 @@ def abs(x: Expr) -> Expr:
result : relax.Expr
The computed result.
"""
- return _ffi_api.abs(x) # type: ignore
+ return _Call("relax.abs", [x], ty_args=ty_args, ty=ty, span=span) # type: ignore
-def acos(x: Expr) -> Expr:
+def acos(x: Expr, *, ty_args=None, ty=None, span=None) -> Expr:
"""Compute element-wise arc cos of the input data.
Parameters
@@ -57,10 +58,10 @@ def acos(x: Expr) -> Expr:
----
The input tensor is required to have float dtype
"""
- return _ffi_api.acos(x) # type: ignore
+ return _Call("relax.acos", [x], ty_args=ty_args, ty=ty, span=span) # type: ignore
-def acosh(x: Expr) -> Expr:
+def acosh(x: Expr, *, ty_args=None, ty=None, span=None) -> Expr:
"""Compute element-wise arc cosh of the input data.
Parameters
@@ -77,10 +78,10 @@ def acosh(x: Expr) -> Expr:
----
The input tensor is required to have float dtype
"""
- return _ffi_api.acosh(x) # type: ignore
+ return _Call("relax.acosh", [x], ty_args=ty_args, ty=ty, span=span) # type: ignore
-def asin(x: Expr) -> Expr:
+def asin(x: Expr, *, ty_args=None, ty=None, span=None) -> Expr:
"""Compute element-wise arc sin of the input data.
Parameters
@@ -97,10 +98,10 @@ def asin(x: Expr) -> Expr:
----
The input tensor is required to have float dtype
"""
- return _ffi_api.asin(x) # type: ignore
+ return _Call("relax.asin", [x], ty_args=ty_args, ty=ty, span=span) # type: ignore
-def asinh(x: Expr) -> Expr:
+def asinh(x: Expr, *, ty_args=None, ty=None, span=None) -> Expr:
"""Compute element-wise arc sinh of the input data.
Parameters
@@ -117,10 +118,10 @@ def asinh(x: Expr) -> Expr:
----
The input tensor is required to have float dtype
"""
- return _ffi_api.asinh(x) # type: ignore
+ return _Call("relax.asinh", [x], ty_args=ty_args, ty=ty, span=span) # type: ignore
-def atan(x: Expr) -> Expr:
+def atan(x: Expr, *, ty_args=None, ty=None, span=None) -> Expr:
"""Compute element-wise arc tan of the input data.
Parameters
@@ -137,10 +138,10 @@ def atan(x: Expr) -> Expr:
----
The input tensor is required to have float dtype
"""
- return _ffi_api.atan(x) # type: ignore
+ return _Call("relax.atan", [x], ty_args=ty_args, ty=ty, span=span) # type: ignore
-def atanh(x: Expr) -> Expr:
+def atanh(x: Expr, *, ty_args=None, ty=None, span=None) -> Expr:
"""Compute element-wise arc tanh of the input data.
Parameters
@@ -157,10 +158,10 @@ def atanh(x: Expr) -> Expr:
----
The input tensor is required to have float dtype
"""
- return _ffi_api.atanh(x) # type: ignore
+ return _Call("relax.atanh", [x], ty_args=ty_args, ty=ty, span=span) # type: ignore
-def bitwise_not(x: Expr) -> Expr:
+def bitwise_not(x: Expr, *, ty_args=None, ty=None, span=None) -> Expr:
"""Compute bitwise NOT of the input data.
Parameters
@@ -173,10 +174,10 @@ def bitwise_not(x: Expr) -> Expr:
result : relax.Expr
The computed result.
"""
- return _ffi_api.bitwise_not(x) # type: ignore
+ return _Call("relax.bitwise_not", [x], ty_args=ty_args, ty=ty, span=span) # type: ignore
-def ceil(x: Expr) -> Expr:
+def ceil(x: Expr, *, ty_args=None, ty=None, span=None) -> Expr:
"""Take ceil of input data.
Parameters
@@ -189,10 +190,10 @@ def ceil(x: Expr) -> Expr:
result : relax.Expr
The computed result.
"""
- return _ffi_api.ceil(x) # type: ignore
+ return _Call("relax.ceil", [x], ty_args=ty_args, ty=ty, span=span) # type: ignore
-def cos(x: Expr) -> Expr:
+def cos(x: Expr, *, ty_args=None, ty=None, span=None) -> Expr:
"""Compute element-wise cos of the input data.
Parameters
@@ -209,10 +210,10 @@ def cos(x: Expr) -> Expr:
----
The input tensor is required to have float dtype
"""
- return _ffi_api.cos(x) # type: ignore
+ return _Call("relax.cos", [x], ty_args=ty_args, ty=ty, span=span) # type: ignore
-def cosh(x: Expr) -> Expr:
+def cosh(x: Expr, *, ty_args=None, ty=None, span=None) -> Expr:
"""Compute element-wise cosh of the input data.
Parameters
@@ -229,10 +230,10 @@ def cosh(x: Expr) -> Expr:
----
The input tensor is required to have float dtype
"""
- return _ffi_api.cosh(x) # type: ignore
+ return _Call("relax.cosh", [x], ty_args=ty_args, ty=ty, span=span) # type: ignore
-def exp(x: Expr) -> Expr:
+def exp(x: Expr, *, ty_args=None, ty=None, span=None) -> Expr:
"""Compute element-wise exp of data.
Parameters
@@ -249,10 +250,10 @@ def exp(x: Expr) -> Expr:
----
The input tensor is required to have float dtype
"""
- return _ffi_api.exp(x) # type: ignore
+ return _Call("relax.exp", [x], ty_args=ty_args, ty=ty, span=span) # type: ignore
-def floor(x: Expr) -> Expr:
+def floor(x: Expr, *, ty_args=None, ty=None, span=None) -> Expr:
"""Take floor of input data.
Parameters
@@ -265,10 +266,10 @@ def floor(x: Expr) -> Expr:
result : relax.Expr
The computed result.
"""
- return _ffi_api.floor(x) # type: ignore
+ return _Call("relax.floor", [x], ty_args=ty_args, ty=ty, span=span) # type: ignore
-def log(x: Expr) -> Expr:
+def log(x: Expr, *, ty_args=None, ty=None, span=None) -> Expr:
"""Compute element-wise natural logarithm of the input data.
Parameters
@@ -285,10 +286,10 @@ def log(x: Expr) -> Expr:
----
The input tensor is required to have float dtype
"""
- return _ffi_api.log(x) # type: ignore
+ return _Call("relax.log", [x], ty_args=ty_args, ty=ty, span=span) # type: ignore
-def logical_not(x: Expr) -> Expr:
+def logical_not(x: Expr, *, ty_args=None, ty=None, span=None) -> Expr:
"""Compute logical NOT of the input data.
Parameters
@@ -301,10 +302,10 @@ def logical_not(x: Expr) -> Expr:
result : relax.Expr
The computed result.
"""
- return _ffi_api.logical_not(x) # type: ignore
+ return _Call("relax.logical_not", [x], ty_args=ty_args, ty=ty, span=span) # type: ignore
-def negative(x: Expr) -> Expr:
+def negative(x: Expr, *, ty_args=None, ty=None, span=None) -> Expr:
"""Compute element-wise negative of the input data.
Parameters
@@ -317,10 +318,10 @@ def negative(x: Expr) -> Expr:
result : relax.Expr
The computed result
"""
- return _ffi_api.negative(x) # type: ignore
+ return _Call("relax.negative", [x], ty_args=ty_args, ty=ty, span=span) # type: ignore
-def round(x: Expr) -> Expr:
+def round(x: Expr, *, ty_args=None, ty=None, span=None) -> Expr:
"""Rounds each element of the input data to nearest integer.
Parameters
@@ -333,10 +334,10 @@ def round(x: Expr) -> Expr:
result : relax.Expr
The computed result.
"""
- return _ffi_api.round(x) # type: ignore
+ return _Call("relax.round", [x], ty_args=ty_args, ty=ty, span=span) # type: ignore
-def rsqrt(x: Expr) -> Expr:
+def rsqrt(x: Expr, *, ty_args=None, ty=None, span=None) -> Expr:
"""Compute element-wise reciprocal square root of the input data.
.. math::
@@ -357,10 +358,10 @@ def rsqrt(x: Expr) -> Expr:
----
The input tensor is required to have float dtype
"""
- return _ffi_api.rsqrt(x) # type: ignore
+ return _Call("relax.rsqrt", [x], ty_args=ty_args, ty=ty, span=span) # type: ignore
-def sigmoid(x: Expr) -> Expr:
+def sigmoid(x: Expr, *, ty_args=None, ty=None, span=None) -> Expr:
"""Compute element-wise sigmoid of the input data.
Parameters
@@ -377,10 +378,10 @@ def sigmoid(x: Expr) -> Expr:
----
The input tensor is required to have float dtype
"""
- return _ffi_api.sigmoid(x) # type: ignore
+ return _Call("relax.sigmoid", [x], ty_args=ty_args, ty=ty, span=span) # type: ignore
-def sign(x: Expr) -> Expr:
+def sign(x: Expr, *, ty_args=None, ty=None, span=None) -> Expr:
"""Returns an indication of the sign of a number for each element of the input data.
Parameters
@@ -393,10 +394,10 @@ def sign(x: Expr) -> Expr:
result : relax.Expr
The computed result.
"""
- return _ffi_api.sign(x) # type: ignore
+ return _Call("relax.sign", [x], ty_args=ty_args, ty=ty, span=span) # type: ignore
-def sin(x: Expr) -> Expr:
+def sin(x: Expr, *, ty_args=None, ty=None, span=None) -> Expr:
"""Compute element-wise sin of the input data.
Parameters
@@ -413,10 +414,10 @@ def sin(x: Expr) -> Expr:
----
The input tensor is required to have float dtype
"""
- return _ffi_api.sin(x) # type: ignore
+ return _Call("relax.sin", [x], ty_args=ty_args, ty=ty, span=span) # type: ignore
-def sinh(x: Expr) -> Expr:
+def sinh(x: Expr, *, ty_args=None, ty=None, span=None) -> Expr:
"""Compute element-wise sinh of the input data.
Parameters
@@ -433,10 +434,10 @@ def sinh(x: Expr) -> Expr:
----
The input tensor is required to have float dtype
"""
- return _ffi_api.sinh(x) # type: ignore
+ return _Call("relax.sinh", [x], ty_args=ty_args, ty=ty, span=span) # type: ignore
-def square(x: Expr) -> Expr:
+def square(x: Expr, *, ty_args=None, ty=None, span=None) -> Expr:
"""Squares each element of the input data.
Parameters
@@ -449,10 +450,10 @@ def square(x: Expr) -> Expr:
result : relax.Expr
The computed result.
"""
- return _ffi_api.square(x) # type: ignore
+ return _Call("relax.square", [x], ty_args=ty_args, ty=ty, span=span) # type: ignore
-def sqrt(x: Expr) -> Expr:
+def sqrt(x: Expr, *, ty_args=None, ty=None, span=None) -> Expr:
"""Compute element-wise square root of the input data.
Parameters
@@ -469,10 +470,10 @@ def sqrt(x: Expr) -> Expr:
----
The input tensor is required to have float dtype
"""
- return _ffi_api.sqrt(x) # type: ignore
+ return _Call("relax.sqrt", [x], ty_args=ty_args, ty=ty, span=span) # type: ignore
-def tan(x: Expr) -> Expr:
+def tan(x: Expr, *, ty_args=None, ty=None, span=None) -> Expr:
"""Compute element-wise tan of the input data.
Parameters
@@ -489,10 +490,10 @@ def tan(x: Expr) -> Expr:
----
The input tensor is required to have float dtype
"""
- return _ffi_api.tan(x) # type: ignore
+ return _Call("relax.tan", [x], ty_args=ty_args, ty=ty, span=span) # type: ignore
-def tanh(x: Expr) -> Expr:
+def tanh(x: Expr, *, ty_args=None, ty=None, span=None) -> Expr:
"""Compute element-wise tanh of the input data.
Parameters
@@ -509,10 +510,10 @@ def tanh(x: Expr) -> Expr:
----
The input tensor is required to have float dtype
"""
- return _ffi_api.tanh(x) # type: ignore
+ return _Call("relax.tanh", [x], ty_args=ty_args, ty=ty, span=span) # type: ignore
-def trunc(x: Expr) -> Expr:
+def trunc(x: Expr, *, ty_args=None, ty=None, span=None) -> Expr:
"""Take trunc of input data.
Parameters
----------
@@ -523,10 +524,10 @@ def trunc(x: Expr) -> Expr:
result : relax.Expr
The computed result.
"""
- return _ffi_api.trunc(x) # type: ignore
+ return _Call("relax.trunc", [x], ty_args=ty_args, ty=ty, span=span) # type: ignore
-def clip(x: Expr, min: Expr, max: Expr) -> Expr:
+def clip(x: Expr, min: Expr, max: Expr, *, ty=None, span=None) -> Expr:
"""Clips tensor values to a specified min and max.
Parameters
@@ -547,10 +548,10 @@ def clip(x: Expr, min: Expr, max: Expr) -> Expr:
"""
min = convert_to_expr(min)
max = convert_to_expr(max)
- return _ffi_api.clip(x, min, max) # type: ignore
+ return _Call("relax.clip", [x, min, max], ty=ty, span=span) # type: ignore
-def erf(x: Expr) -> Expr:
+def erf(x: Expr, *, ty_args=None, ty=None, span=None) -> Expr:
"""Computes the error function of the input.
Parameters
@@ -563,13 +564,13 @@ def erf(x: Expr) -> Expr:
result : relax.Expr
Computed error function for each element.
"""
- return _ffi_api.erf(x) # type: ignore
+ return _Call("relax.erf", [x], ty_args=ty_args, ty=ty, span=span) # type: ignore
###################### Check operators ######################
-def isfinite(x: Expr) -> Expr:
+def isfinite(x: Expr, *, ty_args=None, ty=None, span=None) -> Expr:
"""Check if input value is finite.
Parameters
@@ -582,10 +583,10 @@ def isfinite(x: Expr) -> Expr:
result : relax.Expr
The computed result.
"""
- return _ffi_api.isfinite(x) # type: ignore
+ return _Call("relax.isfinite", [x], ty_args=ty_args, ty=ty, span=span) # type: ignore
-def isinf(x: Expr) -> Expr:
+def isinf(x: Expr, *, ty_args=None, ty=None, span=None) -> Expr:
"""Check if input value is infinite.
Parameters
@@ -598,10 +599,10 @@ def isinf(x: Expr) -> Expr:
result : relax.Expr
The computed result.
"""
- return _ffi_api.isinf(x) # type: ignore
+ return _Call("relax.isinf", [x], ty_args=ty_args, ty=ty, span=span) # type: ignore
-def isnan(x: Expr) -> Expr:
+def isnan(x: Expr, *, ty_args=None, ty=None, span=None) -> Expr:
"""Check if input value is Nan.
Parameters
@@ -614,4 +615,4 @@ def isnan(x: Expr) -> Expr:
result : relax.Expr
The computed result.
"""
- return _ffi_api.isnan(x) # type: ignore
+ return _Call("relax.isnan", [x], ty_args=ty_args, ty=ty, span=span) # type: ignore
diff --git a/python/tvm/relax/op/vision/multibox_transform_loc.py b/python/tvm/relax/op/vision/multibox_transform_loc.py
index 6f5118407a..dcbc37d900 100644
--- a/python/tvm/relax/op/vision/multibox_transform_loc.py
+++ b/python/tvm/relax/op/vision/multibox_transform_loc.py
@@ -16,7 +16,8 @@
# under the License.
"""Multibox location transform for object detection."""
-from . import _ffi_api
+from tvm.ir import Call as _Call
+from tvm.ir.attrs import make_node as _make_attrs
def multibox_transform_loc(
@@ -28,6 +29,9 @@ def multibox_transform_loc(
variances=(1.0, 1.0, 1.0, 1.0),
keep_background=True,
apply_softmax=True,
+ *,
+ ty=None,
+ span=None,
):
"""SSD / TFLite-style decode: priors + offsets → boxes; prepare class scores.
@@ -77,13 +81,17 @@ def multibox_transform_loc(
(ranks, dtypes, ``loc_pred.shape[1] % 4 == 0`` when known, batch match when both batch
axes are known, etc.) still run where applicable.
"""
- return _ffi_api.multibox_transform_loc(
- cls_pred,
- loc_pred,
- anchor,
- clip,
- threshold,
- variances,
- keep_background,
- apply_softmax,
+ return _Call(
+ "relax.vision.multibox_transform_loc",
+ [cls_pred, loc_pred, anchor],
+ attrs=_make_attrs(
+ "relax.attrs.MultiboxTransformLocAttrs",
+ clip=clip,
+ threshold=threshold,
+ variances=variances,
+ keep_background=keep_background,
+ apply_softmax=apply_softmax,
+ ),
+ ty=ty,
+ span=span,
)
diff --git a/python/tvm/relax/op/vision/nms.py b/python/tvm/relax/op/vision/nms.py
index 427d68d113..2715d20bc7 100644
--- a/python/tvm/relax/op/vision/nms.py
+++ b/python/tvm/relax/op/vision/nms.py
@@ -16,7 +16,8 @@
# under the License.
"""Non-maximum suppression operators."""
-from . import _ffi_api
+from tvm.ir import Call as _Call
+from tvm.ir.attrs import make_node as _make_attrs
def all_class_non_max_suppression(
@@ -26,6 +27,9 @@ def all_class_non_max_suppression(
iou_threshold,
score_threshold,
output_format="onnx",
+ *,
+ ty=None,
+ span=None,
):
"""Non-maximum suppression operator for object detection, corresponding to ONNX
NonMaxSuppression and TensorFlow combined_non_max_suppression.
@@ -68,12 +72,26 @@ def all_class_non_max_suppression(
So the box indices and scores for the class 0 come first in a sorted order, followed by
the class 1 etc.
"""
- return _ffi_api.all_class_non_max_suppression(
- boxes, scores, max_output_boxes_per_class, iou_threshold, score_threshold, output_format
+ return _Call(
+ "relax.vision.all_class_non_max_suppression",
+ [boxes, scores, max_output_boxes_per_class, iou_threshold, score_threshold],
+ attrs=_make_attrs(
+ "relax.attrs.AllClassNonMaximumSuppressionAttrs", output_format=output_format
+ ),
+ ty=ty,
+ span=span,
)
-def get_valid_counts(data, score_threshold=0, id_index=0, score_index=1):
+def get_valid_counts(
+ data,
+ score_threshold=0,
+ id_index=0,
+ score_index=1,
+ *,
+ ty=None,
+ span=None,
+):
"""Get valid count of bounding boxes given a score threshold.
Also moves valid boxes to the top of input data.
@@ -99,7 +117,18 @@ def get_valid_counts(data, score_threshold=0, id_index=0, score_index=1):
``[batch_size, num_anchors, elem_length]``, and ``out_indices`` has shape
``[batch_size, num_anchors]``.
"""
- return _ffi_api.get_valid_counts(data, score_threshold, id_index, score_index)
+ return _Call(
+ "relax.vision.get_valid_counts",
+ [data],
+ attrs=_make_attrs(
+ "relax.attrs.GetValidCountsAttrs",
+ score_threshold=score_threshold,
+ id_index=id_index,
+ score_index=score_index,
+ ),
+ ty=ty,
+ span=span,
+ )
def non_max_suppression(
@@ -117,6 +146,9 @@ def non_max_suppression(
invalid_to_bottom=False,
soft_nms_sigma=0.0,
score_threshold=0.0,
+ *,
+ ty=None,
+ span=None,
):
"""Non-maximum suppression operator for object detection.
@@ -186,19 +218,23 @@ def non_max_suppression(
data.
Otherwise returns the modified data tensor.
"""
- return _ffi_api.non_max_suppression(
- data,
- valid_count,
- indices,
- max_output_size,
- iou_threshold,
- force_suppress,
- top_k,
- coord_start,
- score_index,
- id_index,
- return_indices,
- invalid_to_bottom,
- soft_nms_sigma,
- score_threshold,
+ return _Call(
+ "relax.vision.non_max_suppression",
+ [data, valid_count, indices],
+ attrs=_make_attrs(
+ "relax.attrs.NonMaximumSuppressionAttrs",
+ max_output_size=max_output_size,
+ iou_threshold=iou_threshold,
+ force_suppress=force_suppress,
+ top_k=top_k,
+ coord_start=coord_start,
+ score_index=score_index,
+ id_index=id_index,
+ return_indices=return_indices,
+ invalid_to_bottom=invalid_to_bottom,
+ soft_nms_sigma=soft_nms_sigma,
+ score_threshold=score_threshold,
+ ),
+ ty=ty,
+ span=span,
)
diff --git a/python/tvm/relax/op/vision/roi_align.py b/python/tvm/relax/op/vision/roi_align.py
index 8db694c7f2..36075342d8 100644
--- a/python/tvm/relax/op/vision/roi_align.py
+++ b/python/tvm/relax/op/vision/roi_align.py
@@ -16,8 +16,10 @@
# under the License.
"""ROI Align operator"""
+from tvm.ir import Call as _Call
+from tvm.ir.attrs import make_node as _make_attrs
+
from ..base import Expr
-from . import _ffi_api
def roi_align(
@@ -29,6 +31,9 @@ def roi_align(
aligned: bool = False,
layout: str = "NCHW",
mode: str = "avg",
+ *,
+ ty=None,
+ span=None,
):
"""ROI Align operator.
@@ -66,13 +71,18 @@ def roi_align(
"""
if isinstance(pooled_size, int):
pooled_size = (pooled_size, pooled_size)
- return _ffi_api.roi_align(
- data,
- rois,
- pooled_size,
- spatial_scale,
- sample_ratio,
- aligned,
- layout,
- mode,
+ return _Call(
+ "relax.vision.roi_align",
+ [data, rois],
+ attrs=_make_attrs(
+ "relax.attrs.ROIAlignAttrs",
+ pooled_size=pooled_size,
+ spatial_scale=spatial_scale,
+ sample_ratio=sample_ratio,
+ aligned=aligned,
+ layout=layout,
+ mode=mode,
+ ),
+ ty=ty,
+ span=span,
)
diff --git a/python/tvm/relax/op/vision/roi_pool.py b/python/tvm/relax/op/vision/roi_pool.py
index f8b7f11463..50f97412dc 100644
--- a/python/tvm/relax/op/vision/roi_pool.py
+++ b/python/tvm/relax/op/vision/roi_pool.py
@@ -16,8 +16,10 @@
# under the License.
"""ROI Pool operator"""
+from tvm.ir import Call as _Call
+from tvm.ir.attrs import make_node as _make_attrs
+
from ..base import Expr
-from . import _ffi_api
def roi_pool(
@@ -26,6 +28,9 @@ def roi_pool(
pooled_size: int | tuple[int, int] | list[int],
spatial_scale: float,
layout: str = "NCHW",
+ *,
+ ty=None,
+ span=None,
):
"""ROI Pool operator.
@@ -54,4 +59,15 @@ def roi_pool(
"""
if isinstance(pooled_size, int):
pooled_size = (pooled_size, pooled_size)
- return _ffi_api.roi_pool(data, rois, pooled_size, spatial_scale, layout)
+ return _Call(
+ "relax.vision.roi_pool",
+ [data, rois],
+ attrs=_make_attrs(
+ "relax.attrs.ROIPoolAttrs",
+ pooled_size=pooled_size,
+ spatial_scale=spatial_scale,
+ layout=layout,
+ ),
+ ty=ty,
+ span=span,
+ )
diff --git a/python/tvm/relax/op/vm/vm.py b/python/tvm/relax/op/vm/vm.py
index 922b9836d7..f0222e4d74 100644
--- a/python/tvm/relax/op/vm/vm.py
+++ b/python/tvm/relax/op/vm/vm.py
@@ -16,10 +16,10 @@
"""Relax vm primitives."""
from tvm.ir import Call, DataTypeImm, GenericConst, StringImm
+from tvm.ir import Call as _Call
from ...expr import Expr, Tuple, prim_value
from ...utils import convert_to_expr
-from . import _ffi_api
def alloc_storage(
@@ -27,6 +27,9 @@ def alloc_storage(
runtime_device_index: int | Expr,
dtype: str | Expr,
storage_scope: str | StringImm = "global",
+ *,
+ ty=None,
+ span=None,
) -> Call:
"""Construct a Call to allocate a storage with specific size,
runtime_device_index, and dtype.
@@ -58,7 +61,12 @@ def alloc_storage(
storage_scope = StringImm(storage_scope)
if isinstance(runtime_device_index, int):
runtime_device_index = prim_value(runtime_device_index)
- return _ffi_api.alloc_storage(shape, runtime_device_index, dtype, storage_scope) # type: ignore
+ return _Call(
+ "relax.vm.alloc_storage",
+ [shape, runtime_device_index, dtype, storage_scope],
+ ty=ty,
+ span=span,
+ ) # type: ignore
def alloc_tensor(
@@ -67,6 +75,9 @@ def alloc_tensor(
shape: Expr,
dtype: str | Expr,
runtime_device_ind: int | Expr = prim_value(0),
+ *,
+ ty=None,
+ span=None,
) -> Call:
"""Construct a Call to allocate a tensor on a certain storage starting from the given offset.
@@ -100,10 +111,15 @@ def alloc_tensor(
dtype = DataTypeImm(dtype.value if isinstance(dtype, GenericConst) else dtype)
if isinstance(runtime_device_ind, int):
runtime_device_ind = prim_value(runtime_device_ind)
- return _ffi_api.alloc_tensor(storage, offset, shape, dtype, runtime_device_ind) # type: ignore
+ return _Call(
+ "relax.vm.alloc_tensor",
+ [storage, offset, shape, dtype, runtime_device_ind],
+ ty=ty,
+ span=span,
+ ) # type: ignore
-def kill_object(obj: Expr) -> Call:
+def kill_object(obj: Expr, *, ty=None, span=None) -> Call:
"""Construct a Call to set the register corresponding to the input object to
null at runtime, in order to kill the input object.
@@ -117,10 +133,10 @@ def kill_object(obj: Expr) -> Call:
result : Call
CallNode that kills the input object.
"""
- return _ffi_api.kill_object(obj) # type: ignore
+ return _Call("relax.vm.kill_object", [obj], ty=ty, span=span) # type: ignore
-def call_tir_dyn(func: Expr, args: Tuple) -> Call:
+def call_tir_dyn(func: Expr, args: Tuple, *, ty=None, span=None) -> Call:
"""Construct a Call to call_tir_dyn (invoke the given TIR Function)
consisting of the input tensors and the shape of the result.
@@ -141,4 +157,9 @@ def call_tir_dyn(func: Expr, args: Tuple) -> Call:
if isinstance(args, list | tuple):
args = Tuple(args)
- return _ffi_api.call_tir_dyn(func, args) # type: ignore
+ return _Call(
+ "relax.vm.call_tir_dyn",
+ [func, args],
+ ty=ty,
+ span=span,
+ ) # type: ignore
diff --git a/python/tvm/relax/script/ir_builder/__init__.py b/python/tvm/relax/script/ir_builder/__init__.py
index 2289459b5b..91c87aca82 100644
--- a/python/tvm/relax/script/ir_builder/__init__.py
+++ b/python/tvm/relax/script/ir_builder/__init__.py
@@ -33,49 +33,3 @@ from .parser_protocol import (
)
__all__ = [*_native.__all__, *_op.__all__, *_protocol.__all__, "constexpr", "dist"]
-
-
-def _register_printer_names():
- import types # pylint: disable=import-outside-toplevel
-
- from tvm.ir import Op, register_op_attr # pylint: disable=import-outside-toplevel
-
- active = set()
-
- def register(value, path):
- # These constructors use specialized out_ty/argument conventions.
- name = path.rsplit(".", 1)[-1]
- if not callable(value) or (name.startswith("call_tir") and name != "call_tir_packed"):
- return
- try:
- op = Op.get(path)
- except AttributeError:
- return
- if op.has_attr("TScriptPrinterName") and op.get_attr("TScriptPrinterName") == path:
- return
- register_op_attr(path, "TScriptPrinterName", path)
-
- def visit(value, path):
- if id(value) in active:
- return
- active.add(id(value))
- register(value, path)
- for name in getattr(value, "__all__", dir(value)):
- if name.startswith("_"):
- continue
- try:
- member = getattr(value, name)
- except AttributeError:
- continue
- member_path = f"{path}.{name}"
- if isinstance(member, types.ModuleType) and member.__name__.startswith("tvm.relax.op."):
- visit(member, member_path)
- elif callable(member):
- register(member, member_path)
- active.remove(id(value))
-
- visit(_op, "relax")
- visit(dist, "relax.dist")
-
-
-_register_printer_names()
diff --git a/python/tvm/relax/script/ir_builder/distributed/op.py b/python/tvm/relax/script/ir_builder/distributed/op.py
index c4f6ae1825..edc31f33d9 100644
--- a/python/tvm/relax/script/ir_builder/distributed/op.py
+++ b/python/tvm/relax/script/ir_builder/distributed/op.py
@@ -73,23 +73,33 @@ def call_tir(
def annotate_sharding(
- value: Expr, device_mesh: py_str | DeviceMesh, placement: py_str | Placement
+ value: Expr,
+ device_mesh: py_str | DeviceMesh,
+ placement: py_str | Placement,
+ *,
+ ty=None,
+ span=None,
) -> Expr:
if isinstance(device_mesh, py_str):
device_mesh = _lookup_device_mesh(device_mesh)
if isinstance(placement, py_str):
placement = Placement.from_text(placement)
- return _annotate_sharding(value, device_mesh, placement)
+ return _annotate_sharding(value, device_mesh, placement, ty=ty, span=span)
def redistribute(
- value: Expr, device_mesh: py_str | DeviceMesh, placement: py_str | Placement
+ value: Expr,
+ device_mesh: py_str | DeviceMesh,
+ placement: py_str | Placement,
+ *,
+ ty=None,
+ span=None,
) -> Expr:
if isinstance(device_mesh, py_str):
device_mesh = _lookup_device_mesh(device_mesh)
if isinstance(placement, py_str):
placement = Placement.from_text(placement)
- return _redistribute(value, device_mesh, placement)
+ return _redistribute(value, device_mesh, placement, ty=ty, span=span)
__all__ = [
diff --git a/python/tvm/relax/script/ir_builder/op.py b/python/tvm/relax/script/ir_builder/op.py
index 8920ea70f4..8813d25933 100644
--- a/python/tvm/relax/script/ir_builder/op.py
+++ b/python/tvm/relax/script/ir_builder/op.py
@@ -20,9 +20,8 @@
from __future__ import annotations
import builtins
-import functools
-import inspect
import numbers as _numbers
+from types import SimpleNamespace as _Namespace
from typing import Any
import tvm_ffi as _ffi
@@ -32,6 +31,7 @@ from tvm import ir as _ir
from tvm import relax as _relax
from tvm import tirx as _tir
from tvm.ir import StringImm
+from tvm.ir.op import _init_op_api
from tvm.ir.prim import _ffi_api as _prim_ffi
from tvm.relax import Call, Expr, ExternFunc
from tvm.relax.global_info import VDevice
@@ -188,15 +188,18 @@ from tvm.relax.op import (
zeros,
zeros_like,
)
-from tvm.relax.op import (
- call_py_func as _call_py_func,
-)
+from tvm.relax.op import call_py_func as _call_py_func
+from tvm.relax.op import logical_and as _logical_and
+from tvm.relax.op import logical_not as _logical_not
+from tvm.relax.op import logical_or as _logical_or
+from tvm.relax.op import to_vdevice as _to_vdevice
from tvm.relax.op.builtin import stop_lift_params
from tvm.relax.type import Type
from tvm.relax.utils import convert_to_expr
from tvm.runtime import ObjectConvertible
from tvm.script.ir_builder import base as _base
+from . import distributed as dist # noqa: F401
from .ir import _value, lookup_vdevice
py_print = builtins.print
@@ -205,7 +208,13 @@ py_str = str
_Span = _base.SpanEntry | _ir.Span | None
-def to_vdevice(data: Expr, dst_vdevice: py_str | VDevice) -> Expr:
+def to_vdevice(
+ data: Expr,
+ dst_vdevice: py_str | VDevice = None,
+ *,
+ ty=None,
+ span=None,
+) -> Expr:
"""Copy data to the destination device.
Parameters
@@ -228,7 +237,7 @@ def to_vdevice(data: Expr, dst_vdevice: py_str | VDevice) -> Expr:
else:
dst_vdevice = lookup_vdevice(dst_vdevice, 0)
- return tvm.relax.op.to_vdevice(data, dst_vdevice)
+ return _to_vdevice(data, dst_vdevice, ty=ty, span=span)
def call_packed(
@@ -279,7 +288,7 @@ def call_packed(
if kwargs or not is_default:
attrs = tvm.ir.attrs.make_node(attrs_type_key, **kwargs)
- return Call.unchecked(op, args, attrs=attrs, ty_args=ty_args)
+ return Call(op, args, attrs=attrs, ty_args=ty_args)
def call_py_func(
@@ -327,28 +336,6 @@ def call_py_func(
return _call_py_func(func_name_imm, args, out_ty)
-def _ty_arg_wrapper(func):
- """Normalize callable and object-convertible type arguments for operators."""
-
- def _convert_tensor_type(args):
- if isinstance(args, list | py_tuple): # type: ignore
- new_args = [_convert_tensor_type(x) for x in args]
- return type(args)(new_args)
- if isinstance(args, dict):
- return {_convert_tensor_type(k): _convert_tensor_type(v) for k, v in args.items()}
- if inspect.isfunction(args):
- args = args()
- if isinstance(args, ObjectConvertible):
- args = args.asobject()
- return args
-
- @functools.wraps(func)
- def wrapped(*args, **kwargs):
- return func(*_convert_tensor_type(args), **_convert_tensor_type(kwargs))
-
- return wrapped # type: ignore
-
-
def emit_with_type(
op: str,
args: Expr,
@@ -369,8 +356,7 @@ def emit_with_type(
call: Call
The created Relax Call
"""
- builtin_call = tvm.ir.Op.get(op)
- return Call.unchecked(builtin_call, args, attrs=None, ty_args=ty_args)
+ return Call(op, args, ty_args=ty_args)
def _logical_pair(lhs, rhs, operation, primitive, python_operation):
@@ -381,7 +367,7 @@ def _logical_pair(lhs, rhs, operation, primitive, python_operation):
return operation(_value(lhs), _value(rhs))
-def logical_and(*values):
+def logical_and(*values, ty_args=None, ty=None, span=None):
"""Construct conjunction of host, primitive, or tensor values.
Parameters
@@ -401,15 +387,17 @@ def logical_and(*values):
All arguments are evaluated before this call; it does not provide Python
short-circuit evaluation of the argument expressions.
"""
+ if any(field is not None for field in (ty_args, ty, span)):
+ return _logical_and(*values, ty_args=ty_args, ty=ty, span=span)
if not values:
raise TypeError("logical_and requires at least one operand")
result = values[0]
for value in values[1:]:
- result = _logical_pair(result, value, _relax.op.logical_and, _tir.And, lambda a, b: a and b)
+ result = _logical_pair(result, value, _logical_and, _tir.And, lambda a, b: a and b)
return result
-def logical_or(*values):
+def logical_or(*values, ty_args=None, ty=None, span=None):
"""Construct disjunction of host, primitive, or tensor values.
Parameters
@@ -429,15 +417,17 @@ def logical_or(*values):
All arguments are evaluated before this call; it does not provide Python
short-circuit evaluation of the argument expressions.
"""
+ if any(field is not None for field in (ty_args, ty, span)):
+ return _logical_or(*values, ty_args=ty_args, ty=ty, span=span)
if not values:
raise TypeError("logical_or requires at least one operand")
result = values[0]
for value in values[1:]:
- result = _logical_pair(result, value, _relax.op.logical_or, _tir.Or, lambda a, b: a or b)
+ result = _logical_pair(result, value, _logical_or, _tir.Or, lambda a, b: a or b)
return result
-def logical_not(value):
+def logical_not(value, *, ty_args=None, ty=None, span=None):
"""Negate a host, primitive, or tensor condition.
Parameters
@@ -451,10 +441,12 @@ def logical_not(value):
result : Expr or bool
The logical negation without testing an IR expression as a Python bool.
"""
+ if any(field is not None for field in (ty_args, ty, span)):
+ return _logical_not(value, ty_args=ty_args, ty=ty, span=span)
if _ir.is_prim_expr(value):
return _tir.Not(value)
if isinstance(value, _ir.Expr):
- return _relax.op.logical_not(value)
+ return _logical_not(value)
return not value
@@ -578,8 +570,8 @@ def ne_(lhs: Any, rhs: Any, *, span: _Span = None) -> _ir.Expr:
return _prim_ffi._OpNE(lhs, rhs, span.span if isinstance(span, _base.SpanEntry) else span)
-invoke_closure = _ty_arg_wrapper(invoke_closure)
-call_builtin_with_ctx = _ty_arg_wrapper(call_builtin_with_ctx)
+inspect = _Namespace()
+_init_op_api("relax", __name__)
__all__ = [
"abs",
@@ -658,6 +650,7 @@ __all__ = [
"image",
"index_put",
"index_tensor",
+ "inspect",
"invoke_closure",
"invoke_pure_closure",
"isfinite",
diff --git a/python/tvm/relax/script/ir_builder/parser_protocol.py b/python/tvm/relax/script/ir_builder/parser_protocol.py
index 3d09966ed7..f5df2b67e1 100644
--- a/python/tvm/relax/script/ir_builder/parser_protocol.py
+++ b/python/tvm/relax/script/ir_builder/parser_protocol.py
@@ -533,7 +533,7 @@ def assert_(
"""Implements :func:`tvm.script.ir_builder.parser_protocol.assert_`."""
if not isinstance(message, _python.str):
raise TypeError("An assertion message must be construction-time text")
- emit_(_base.at_(span, _op.assert_op(condition, format=message)), span=span)
+ emit_(_base.at_(span, _op.assert_op(condition, message)), span=span)
# --------------------------------------
diff --git a/python/tvm/relax/transform/legalize_ops/inspect_op.py b/python/tvm/relax/transform/legalize_ops/inspect_op.py
index 638c7a4dce..78fe31d1e2 100644
--- a/python/tvm/relax/transform/legalize_ops/inspect_op.py
+++ b/python/tvm/relax/transform/legalize_ops/inspect_op.py
@@ -60,18 +60,18 @@ def _tensor_stride_i(bb: BlockBuilder, call: Call) -> Expr:
T.func_attr({"tirx.is_host_func": True, "tirx.is_scheduled": True})
assert T.int64(0) <= axis, "Specified axis may not be negative"
ndim: T.let[T.int32] = T.tvm_struct_get(
- dlpack_handle, 0, int(TVMStructFieldKind.kDLTensorNDim), "int32"
+ dlpack_handle, 0, int(TVMStructFieldKind.kDLTensorNDim), ty="int32"
)
assert axis < T.Cast("int64", ndim), (
"Specified axis may not be larger than the tensor's dimensionality"
)
stride_ptr: T.let[T.handle("int64")] = T.tvm_struct_get(
- dlpack_handle, 0, int(TVMStructFieldKind.kDLTensorStrides), T.handle("int64").ty
+ dlpack_handle, 0, int(TVMStructFieldKind.kDLTensorStrides), ty=T.handle("int64").ty
)
if T.isnullptr(stride_ptr):
shape_ptr: T.let[T.handle("int64")] = T.tvm_struct_get(
- dlpack_handle, 0, int(TVMStructFieldKind.kDLTensorShape), T.handle("int64").ty
+ dlpack_handle, 0, int(TVMStructFieldKind.kDLTensorShape), ty=T.handle("int64").ty
)
shape = T.decl_tensor(ndim, "int64", data=shape_ptr)
@@ -101,7 +101,7 @@ def _tensor_byte_offset(bb: BlockBuilder, call: Call) -> Expr:
def _get_tensor_byte_offset(dlpack_handle: T.handle) -> T.uint64:
T.func_attr({"tirx.is_host_func": True, "tirx.is_scheduled": True})
byte_offset: T.let[T.uint64] = T.tvm_struct_get(
- dlpack_handle, 0, int(TVMStructFieldKind.kDLTensorByteOffset), "uint64"
+ dlpack_handle, 0, int(TVMStructFieldKind.kDLTensorByteOffset), ty="uint64"
)
return byte_offset
@@ -115,13 +115,13 @@ def _tensor_elem_offset(bb: BlockBuilder, call: Call) -> Expr:
def _get_tensor_elem_offset(dlpack_handle: T.handle) -> T.uint64:
T.func_attr({"tirx.is_host_func": True, "tirx.is_scheduled": True})
byte_offset: T.let[T.uint64] = T.tvm_struct_get(
- dlpack_handle, 0, int(TVMStructFieldKind.kDLTensorByteOffset), "uint64"
+ dlpack_handle, 0, int(TVMStructFieldKind.kDLTensorByteOffset), ty="uint64"
)
scalar_bits: T.let[T.uint8] = T.tvm_struct_get(
- dlpack_handle, 0, int(TVMStructFieldKind.kDLTensorTypeBits), "uint8"
+ dlpack_handle, 0, int(TVMStructFieldKind.kDLTensorTypeBits), ty="uint8"
)
lanes: T.let[T.uint16] = T.tvm_struct_get(
- dlpack_handle, 0, int(TVMStructFieldKind.kDLTensorTypeLanes), "uint16"
+ dlpack_handle, 0, int(TVMStructFieldKind.kDLTensorTypeLanes), ty="uint16"
)
bytes_per_element: T.let[T.uint64] = T.ceildiv(
scalar_bits.astype("uint64") * lanes.astype("uint64"), 8
diff --git a/python/tvm/s_tir/schedule/schedule.py b/python/tvm/s_tir/schedule/schedule.py
index 05c714fbbd..b8917ea8ab 100644
--- a/python/tvm/s_tir/schedule/schedule.py
+++ b/python/tvm/s_tir/schedule/schedule.py
@@ -1786,7 +1786,7 @@ class Schedule(Object):
with Ts.sblock("A"):
Ts.reads(data_io[:64])
Ts.writes(data_io[:64])
- T.evaluate(T.call_extern("call_impl", data_io.data, dtype=""))
+ T.evaluate(T.call_extern("call_impl", data_io.data, ty="void"))
Create the schedule and cache_inplace:
@@ -1814,7 +1814,7 @@ class Schedule(Object):
with Ts.sblock("A"):
Ts.reads(data_io_local[0 : 64])
Ts.writes(data_io_local[0 : 64])
- T.evaluate(T.call_extern("call_impl", data_io_local.data, dtype=""))
+ T.evaluate(T.call_extern("call_impl", data_io_local.data, ty="void"))
for ax0 in T.serial(64):
with Ts.sblock("data_io_local"):
v0 = Ts.axis.spatial(64, ax0)
diff --git a/python/tvm/s_tir/tensor_intrin/arm_cpu.py b/python/tvm/s_tir/tensor_intrin/arm_cpu.py
index 501e95b694..86ba0dc7e8 100644
--- a/python/tvm/s_tir/tensor_intrin/arm_cpu.py
+++ b/python/tvm/s_tir/tensor_intrin/arm_cpu.py
@@ -64,48 +64,48 @@ def neon_4x4_i8i8i32_impl(
Ts.writes(C[0:4])
A_int8 = A.vload([0], "int8x4")
- re_int32 = T.reinterpret(A_int8, dtype="int32")
+ re_int32 = T.reinterpret(A_int8, ty="int32")
vec_ai32 = T.broadcast(re_int32, 2)
- vec_a = T.reinterpret(vec_ai32, dtype="int8x8")
+ vec_a = T.reinterpret(vec_ai32, ty="int8x8")
vec_b = B.vload([0, 0], dtype="int8x16")
# TODO(masahi): Remove duplication when inlined function call is supported
- vec_b_low = T.vectorlow(vec_b, dtype="int8x8")
+ vec_b_low = T.vectorlow(vec_b, ty="int8x8")
multiply_low = T.call_llvm_pure_intrin(
T.llvm_lookup_intrinsic_id("llvm.aarch64.neon.smull.v8i16"),
vec_a,
vec_b_low,
- dtype="int16x8",
+ ty="int16x8",
)
pairwise_reduction_low = T.call_llvm_pure_intrin(
T.llvm_lookup_intrinsic_id("llvm.aarch64.neon.saddlp.v4i32.v8i16"),
multiply_low,
- dtype="int32x4",
+ ty="int32x4",
)
- vec_b_high = T.vectorhigh(vec_b, dtype="int8x8")
+ vec_b_high = T.vectorhigh(vec_b, ty="int8x8")
multiply_high = T.call_llvm_pure_intrin(
T.llvm_lookup_intrinsic_id("llvm.aarch64.neon.smull.v8i16"),
vec_a,
vec_b_high,
- dtype="int16x8",
+ ty="int16x8",
)
pairwise_reduction_high = T.call_llvm_pure_intrin(
T.llvm_lookup_intrinsic_id("llvm.aarch64.neon.saddlp.v4i32.v8i16"),
multiply_high,
- dtype="int32x4",
+ ty="int32x4",
)
C[T.ramp(T.int32(0), 1, 4)] += T.call_llvm_pure_intrin(
T.llvm_lookup_intrinsic_id("llvm.aarch64.neon.addp.v4i32"),
pairwise_reduction_low,
pairwise_reduction_high,
- dtype="int32x4",
+ ty="int32x4",
)
@@ -147,9 +147,9 @@ def get_dotprod_intrin(in_dtype, out_dtype):
Ts.writes(C[0:4])
A_i8x4 = A.vload([0], in_dtype_x4)
- A_i32 = T.reinterpret(A_i8x4, dtype=out_dtype)
+ A_i32 = T.reinterpret(A_i8x4, ty=out_dtype)
vec_ai32 = T.broadcast(A_i32, 4)
- vec_a = T.reinterpret(vec_ai32, dtype=in_dtype_x16)
+ vec_a = T.reinterpret(vec_ai32, ty=in_dtype_x16)
vec_b = B.vload([0, 0], dtype=in_dtype_x16)
@@ -160,7 +160,7 @@ def get_dotprod_intrin(in_dtype, out_dtype):
vec_c,
vec_a,
vec_b,
- dtype=out_dtype_x4,
+ ty=out_dtype_x4,
)
return dot_prod_desc, dot_prod_impl
@@ -211,9 +211,7 @@ def _create_active_lane_mask(tensor, relative_offsets, vertical_limit):
limit = T.Min(limit, T.Cast("int32", vertical_limit) * stride)
return T.get_active_lane_mask(
- "uint1xvscalex4",
- T.Cast("int32", base),
- T.Cast("int32", limit),
+ T.Cast("int32", base), T.Cast("int32", limit), ty="uint1xvscalex4"
)
@@ -303,12 +301,12 @@ def get_sme_transpose_interleave_2svlx2svl_fp32_intrin(cols, rows):
)
T.evaluate(
T.call_llvm_intrin(
- "void",
"llvm.aarch64.sme.ld1w.horiz",
predicate,
input_ptr,
sub_tile,
slice_idx,
+ ty="void",
)
)
@@ -326,12 +324,12 @@ def get_sme_transpose_interleave_2svlx2svl_fp32_intrin(cols, rows):
)
T.evaluate(
T.call_llvm_intrin(
- "void",
"llvm.aarch64.sme.st1w.vert",
predicate,
output_ptr,
sub_tile,
slice_idx,
+ ty="void",
)
)
@@ -426,23 +424,23 @@ def get_sme_transpose_interleave_block2_2svl_fp16_intrin():
input_ptr = A.access_ptr("r", offset=offset)
T.evaluate(
T.call_llvm_intrin(
- "void",
"llvm.aarch64.sme.ld1h.horiz",
ptrue_fp16,
input_ptr,
sub_tile_idx,
slice_idx * 2,
+ ty="void",
)
)
input_ptr = A.access_ptr("r", offset=offset + (SVF // 2) * A.strides[0])
T.evaluate(
T.call_llvm_intrin(
- "void",
"llvm.aarch64.sme.ld1h.horiz",
ptrue_fp16,
input_ptr,
sub_tile_idx,
slice_idx * 2 + 1,
+ ty="void",
)
)
@@ -453,23 +451,23 @@ def get_sme_transpose_interleave_block2_2svl_fp16_intrin():
output_ptr = A_t.access_ptr("w", offset=offset)
T.evaluate(
T.call_llvm_intrin(
- "void",
"llvm.aarch64.sme.st1w.vert",
ptrue_fp32,
output_ptr,
sub_tile_idx,
slice_idx,
+ ty="void",
)
)
output_ptr = A_t.access_ptr("w", offset=offset + A_t.strides[0])
T.evaluate(
T.call_llvm_intrin(
- "void",
"llvm.aarch64.sme.st1w.vert",
ptrue_fp32,
output_ptr,
sub_tile_idx + 2,
slice_idx,
+ ty="void",
)
)
@@ -679,13 +677,13 @@ def get_sme_gemm_interleaved_mopa_2svlx2svl_intrin(M, K, in_dtype):
T.evaluate(
T.call_llvm_intrin(
- "void",
fmopa_intrin,
sub_tile,
input_1[1],
input_2[1],
input_1[0],
input_2[0],
+ ty="void",
)
)
@@ -699,7 +697,6 @@ def get_sme_gemm_interleaved_mopa_2svlx2svl_intrin(M, K, in_dtype):
T.evaluate(
T.call_llvm_intrin(
- "void",
"llvm.aarch64.sme.st1w.horiz",
_create_active_lane_mask(
C, (vert_offset + slice_idx, horiz_offset), M
@@ -707,6 +704,7 @@ def get_sme_gemm_interleaved_mopa_2svlx2svl_intrin(M, K, in_dtype):
output_ptr,
T.int32(sub_tile_idx),
T.int32(slice_idx),
+ ty="void",
)
)
@@ -737,7 +735,7 @@ def get_sme_init_intrin():
Ts.reads()
Ts.writes(C[0:SVF2, 0:SVF2])
clear_all_tiles = T.int32(255)
- T.evaluate(T.call_llvm_intrin("void", "llvm.aarch64.sme.zero", clear_all_tiles))
+ T.evaluate(T.call_llvm_intrin("llvm.aarch64.sme.zero", clear_all_tiles, ty="void"))
return desc, impl
diff --git a/python/tvm/s_tir/tensor_intrin/cuda.py b/python/tvm/s_tir/tensor_intrin/cuda.py
index 76de76eaa4..bff7b6723e 100644
--- a/python/tvm/s_tir/tensor_intrin/cuda.py
+++ b/python/tvm/s_tir/tensor_intrin/cuda.py
@@ -206,7 +206,7 @@ def get_ldmatrix_intrin(
warp.elem_offset + lift(local_size) * tx,
shared.access_ptr("r"),
smem_offset(tx, s0),
- dtype=dtype,
+ ty=dtype,
)
)
@@ -413,7 +413,7 @@ def get_mma_intrin(
C.data,
C.elem_offset + tx * lift(local_size_out),
False,
- dtype=out_dtype,
+ ty=out_dtype,
)
)
@@ -432,7 +432,7 @@ def get_mma_intrin(
C.data,
C.elem_offset + tx * lift(local_size_out) + lift(local_size_out) // 2,
False,
- dtype=out_dtype,
+ ty=out_dtype,
)
)
@@ -545,7 +545,7 @@ def get_mma_fill_intrin(dtype, local_size):
for tx in T.thread_binding(0, WARP_SIZE, "threadIdx.x"):
T.evaluate(
- T.cuda.mma_fill_legacy(local_size, C_warp.data, C_warp.elem_offset, dtype=dtype)
+ T.cuda.mma_fill_legacy(local_size, C_warp.data, C_warp.elem_offset, ty=dtype)
)
return mma_fill_desc, mma_fill_impl
@@ -606,7 +606,7 @@ def get_mma_store_intrin(dtype, local_size, scope="global", use_mma_store_intrin
C_warp.data,
C_warp.elem_offset,
s0,
- dtype=dtype,
+ ty=dtype,
)
)
@@ -1513,7 +1513,7 @@ def get_mma_load_intrin(
get_index(dst.elem_offset, d0),
src.access_ptr("r"),
get_tx_index(tx, s0),
- dtype=dtype,
+ ty=dtype,
)
)
@@ -1608,7 +1608,7 @@ def get_mma_sync_intrin(
C.data,
get_index_C(C.elem_offset, c0),
False,
- dtype=out_dtype,
+ ty=out_dtype,
)
)
diff --git a/python/tvm/s_tir/tensor_intrin/dot_product_common.py b/python/tvm/s_tir/tensor_intrin/dot_product_common.py
index 4d1f6a8302..25e8c53eec 100644
--- a/python/tvm/s_tir/tensor_intrin/dot_product_common.py
+++ b/python/tvm/s_tir/tensor_intrin/dot_product_common.py
@@ -58,7 +58,7 @@ def get_dp4a_intrin(dtype_a, dtype_b, dtype_c):
A.vload([0], vec_type_a),
B.vload([0], vec_type_b),
T.uint32(0) if T.constexpr(dtype_c == "uint32") else T.int32(0),
- dtype=dtype_c,
+ ty=dtype_c,
)
return dp4a_desc, dp4a_impl
diff --git a/python/tvm/s_tir/tensor_intrin/hexagon.py b/python/tvm/s_tir/tensor_intrin/hexagon.py
index d7b5e218cb..ff7d8eb26a 100644
--- a/python/tvm/s_tir/tensor_intrin/hexagon.py
+++ b/python/tvm/s_tir/tensor_intrin/hexagon.py
@@ -104,17 +104,17 @@ def generate_dot_product_32x4_u8u8i32(mem_scope="global"):
Ts.writes(C[0:32])
A_u8x4 = A.vload([0], "uint8x4")
- A_i32 = T.reinterpret(A_u8x4, dtype="int32")
+ A_i32 = T.reinterpret(A_u8x4, ty="int32")
B_i8x128 = B.vload([0, 0], dtype="uint8x128")
- B_i32x32 = T.reinterpret(B_i8x128, dtype="int32x32")
+ B_i32x32 = T.reinterpret(B_i8x128, ty="int32x32")
C[T.ramp(T.int32(0), 1, 32)] = T.call_llvm_pure_intrin(
T.llvm_lookup_intrinsic_id("llvm.hexagon.V6.vrmpyub.acc.128B"),
C[T.ramp(T.int32(0), 1, 32)],
B_i32x32,
A_i32,
- dtype="int32x32",
+ ty="int32x32",
)
return dot_product_32x4_u8u8i32_desc, dot_product_32x4_u8u8i32_vrmpy
@@ -147,17 +147,17 @@ def generate_dot_product_32x4_u8i8i32(mem_scope="global"):
Ts.writes(C[0:32])
A_u8x4 = A.vload([0], "uint8x4")
- A_i32 = T.reinterpret(A_u8x4, dtype="int32")
+ A_i32 = T.reinterpret(A_u8x4, ty="int32")
B_i8x128 = B.vload([0, 0], dtype="int8x128")
- B_i32x32 = T.reinterpret(B_i8x128, dtype="int32x32")
+ B_i32x32 = T.reinterpret(B_i8x128, ty="int32x32")
C[T.ramp(T.int32(0), 1, 32)] = T.call_llvm_pure_intrin(
T.llvm_lookup_intrinsic_id("llvm.hexagon.V6.vrmpybusv.acc.128B"),
C[T.ramp(T.int32(0), 1, 32)],
T.broadcast(A_i32, 32),
B_i32x32,
- dtype="int32x32",
+ ty="int32x32",
)
return dot_product_32x4_u8i8i32_desc, dot_product_32x4_u8i8i32_vrmpy
@@ -190,17 +190,17 @@ def generate_dot_product_32x2_i16i16i32(mem_scope="global"):
Ts.writes(C[0:32])
A_i16x2 = A.vload([0], "int16x2")
- A_i32 = T.reinterpret(A_i16x2, dtype="int32")
+ A_i32 = T.reinterpret(A_i16x2, ty="int32")
B_i16x64 = B.vload([0, 0], dtype="int16x64")
- B_i32x32 = T.reinterpret(B_i16x64, dtype="int32x32")
+ B_i32x32 = T.reinterpret(B_i16x64, ty="int32x32")
C[T.ramp(T.int32(0), 1, 32)] = T.call_llvm_pure_intrin(
T.llvm_lookup_intrinsic_id("llvm.hexagon.V6.vdmpyhvsat.acc.128B"),
C[T.ramp(T.int32(0), 1, 32)],
T.Broadcast(A_i32, 32),
B_i32x32,
- dtype="int32x32",
+ ty="int32x32",
)
return dot_product_32x2_i16i16i32_desc, dot_product_32x2_i16i16i32_vdmpy
diff --git a/python/tvm/s_tir/tensor_intrin/riscv_cpu.py b/python/tvm/s_tir/tensor_intrin/riscv_cpu.py
index 918c2c97d4..14fbe1fadd 100644
--- a/python/tvm/s_tir/tensor_intrin/riscv_cpu.py
+++ b/python/tvm/s_tir/tensor_intrin/riscv_cpu.py
@@ -117,11 +117,10 @@ def rvv_vec_dot_product_kernels(
Ts.writes(C[0:n_lanes])
vec_A = T.call_llvm_intrin(
- f"{data_dtype}xvscalex{d_dtype_lanes}",
"llvm.riscv.vle",
T.broadcast(T.Cast(data_dtype, 0), T.vscale() * d_dtype_lanes),
- T.tvm_access_ptr(T.type_annotation(data_dtype), A.data, 0, n_elems, 1),
- T.int64(n_elems))
+ T.tvm_access_ptr(T.type_annotation(ty=data_dtype), A.data, 0, n_elems, 1),
+ T.int64(n_elems), ty=f"{data_dtype}xvscalex{d_dtype_lanes}")
for i in range(n_lanes):
with Ts.sblock("update"):
@@ -129,14 +128,12 @@ def rvv_vec_dot_product_kernels(
Ts.writes(C[i])
vec_B_row = T.call_llvm_intrin(
- f"{weight_dtype}xvscalex{w_dtype_lanes}",
"llvm.riscv.vle",
T.broadcast(T.Cast(data_dtype, 0), T.vscale() * w_dtype_lanes),
- T.tvm_access_ptr(T.type_annotation(weight_dtype), B.data, i * n_elems, n_elems, 1),
- T.int64(n_elems))
+ T.tvm_access_ptr(T.type_annotation(ty=weight_dtype), B.data, i * n_elems, n_elems, 1),
+ T.int64(n_elems), ty=f"{weight_dtype}xvscalex{w_dtype_lanes}")
product = T.call_llvm_intrin(
... 10953 lines suppressed ...