You are viewing a plain text version of this content. The canonical link for it is here.
Posted to commits@tvm.apache.org by bo...@apache.org on 2026/06/21 14:12:18 UTC
(tvm) branch main updated: [REFACTOR][IR] Unify StructInfo and Type (#19853)
This is an automated email from the ASF dual-hosted git repository.
spectrometerHBH 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 1bb5cf6102 [REFACTOR][IR] Unify StructInfo and Type (#19853)
1bb5cf6102 is described below
commit 1bb5cf6102da318f4dece97c5665047d0ee7717a
Author: Tianqi Chen <tq...@users.noreply.github.com>
AuthorDate: Sun Jun 21 10:12:12 2026 -0400
[REFACTOR][IR] Unify StructInfo and Type (#19853)
## Summary
- unify Relax's former StructInfo surface into the Type vocabulary and
Expr.ty storage path
- remove leftover DependentTypeNode and legacy OpNode::op_type storage
- keep base Type nullable while concrete Relax/DTensor type refs are
non-nullable
- clean stale StructInfo/TensorStructInfo/sinfo vocabulary in
Python/docs and distributed-op macros
- address Gemini follow-ups for parser annotations, BlockBuilder
docstring, and Adreno TensorType cast audit
---
docs/arch/fusion.rst | 6 +-
docs/arch/tvmscript.rst | 24 +-
docs/deep_dive/relax/dpl.rst | 14 +-
docs/deep_dive/relax/learning.rst | 22 +-
docs/deep_dive/relax/tutorials/relax_creation.py | 4 +-
.../tutorials/mix_python_and_tvm_with_pymodule.py | 18 +-
docs/reference/security.rst | 6 +-
include/tvm/ir/expr.h | 72 +-
include/tvm/ir/op.h | 3 -
include/tvm/ir/type.h | 62 +-
include/tvm/relax/analysis.h | 182 +-
include/tvm/relax/attrs/distributed.h | 2 +-
include/tvm/relax/binding_rewrite.h | 4 +-
include/tvm/relax/block_builder.h | 6 +-
include/tvm/relax/dataflow_pattern.h | 35 +-
include/tvm/relax/dataflow_pattern_functor.h | 7 +-
include/tvm/relax/distributed/axis_group_graph.h | 2 +-
.../relax/distributed/{struct_info.h => type.h} | 56 +-
include/tvm/relax/expr.h | 155 +-
include/tvm/relax/expr_functor.h | 102 +-
include/tvm/relax/nested_msg.h | 59 +-
include/tvm/relax/op_attr_types.h | 7 +-
include/tvm/relax/script/builder/frame.h | 14 +-
include/tvm/relax/script/builder/ir.h | 24 +-
include/tvm/relax/struct_info.h | 424 ----
include/tvm/relax/struct_info_functor.h | 160 --
include/tvm/relax/transform.h | 2 +-
include/tvm/relax/type.h | 405 +++-
include/tvm/{ir => relax}/type_functor.h | 84 +-
include/tvm/relax/utils.h | 15 +-
include/tvm/script/printer/config.h | 4 +-
python/tvm/contrib/cutlass/build.py | 32 +-
python/tvm/contrib/hexagon/generate_take_op.py | 6 +-
python/tvm/contrib/hexagon/hexagon_unary_ops.py | 14 +-
python/tvm/ir/expr.py | 13 +-
python/tvm/ir/global_info.py | 2 +-
python/tvm/ir/type.py | 4 +-
python/tvm/relax/__init__.py | 16 +-
python/tvm/relax/analysis/__init__.py | 10 +-
python/tvm/relax/analysis/analysis.py | 138 +-
python/tvm/relax/backend/adreno/clml.py | 40 +-
.../relax/backend/contrib/example_npu/patterns.py | 2 +-
python/tvm/relax/backend/cuda/cublas.py | 14 +-
python/tvm/relax/backend/cuda/cudnn.py | 14 +-
python/tvm/relax/backend/cuda/cutlass.py | 46 +-
python/tvm/relax/backend/dispatch_sampling.py | 6 +-
python/tvm/relax/backend/dispatch_sort_scan.py | 26 +-
python/tvm/relax/backend/metal/coreml.py | 18 +-
python/tvm/relax/backend/patterns.py | 26 +-
python/tvm/relax/backend/rocm/hipblas.py | 12 +-
python/tvm/relax/backend/utils.py | 22 +-
python/tvm/relax/base_py_module.py | 32 +-
python/tvm/relax/block_builder.py | 36 +-
python/tvm/relax/distributed/__init__.py | 2 +-
.../relax/distributed/{struct_info.py => type.py} | 22 +-
python/tvm/relax/dpl/pattern.py | 22 +-
python/tvm/relax/dpl/rewrite.py | 2 +-
python/tvm/relax/expr.py | 178 +-
python/tvm/relax/expr_functor.py | 10 +-
python/tvm/relax/frontend/common.py | 4 +-
python/tvm/relax/frontend/nn/core.py | 66 +-
python/tvm/relax/frontend/nn/exporter.py | 14 +-
python/tvm/relax/frontend/nn/extern.py | 4 +-
python/tvm/relax/frontend/nn/llm/kv_cache.py | 50 +-
python/tvm/relax/frontend/nn/modules.py | 10 +-
python/tvm/relax/frontend/nn/op.py | 24 +-
python/tvm/relax/frontend/nn/subroutine.py | 30 +-
python/tvm/relax/frontend/onnx/onnx_frontend.py | 358 ++--
.../frontend/stablehlo/stablehlo_translator.py | 10 +-
.../tvm/relax/frontend/tflite/tflite_frontend.py | 64 +-
.../frontend/torch/base_fx_graph_translator.py | 178 +-
.../frontend/torch/exported_program_translator.py | 58 +-
python/tvm/relax/frontend/torch/fx_translator.py | 76 +-
python/tvm/relax/ir/instrument.py | 10 +-
python/tvm/relax/op/_op_gradient.py | 36 +-
python/tvm/relax/op/base.py | 179 +-
python/tvm/relax/op/binary.py | 6 +-
python/tvm/relax/op/distributed/distributed.py | 21 +-
python/tvm/relax/op/grad/grad.py | 8 +-
python/tvm/relax/op/manipulate.py | 2 +-
python/tvm/relax/op/mask.py | 2 +-
.../tvm/relax/op/vision/multibox_transform_loc.py | 4 +-
python/tvm/relax/relax_to_pyfunc_converter.py | 12 +-
python/tvm/relax/script/builder/distributed/ir.py | 30 +-
python/tvm/relax/script/builder/ir.py | 150 +-
python/tvm/relax/script/parser/dist.py | 22 +-
python/tvm/relax/script/parser/entry.py | 146 +-
python/tvm/relax/script/parser/parser.py | 78 +-
python/tvm/relax/testing/ast_printer.py | 94 +-
python/tvm/relax/testing/attention.py | 2 +-
python/tvm/relax/testing/nn.py | 6 +-
python/tvm/relax/training/loss.py | 64 +-
python/tvm/relax/training/optimizer.py | 52 +-
python/tvm/relax/training/setup_trainer.py | 6 +-
python/tvm/relax/training/trainer.py | 10 +-
python/tvm/relax/transform/__init__.py | 2 +-
.../fold_batch_norm_to_conv2d_for_inference.py | 6 +-
.../tvm/relax/transform/fuse_transpose_matmul.py | 12 +-
.../tvm/relax/transform/lazy_transform_params.py | 40 +-
.../transform/legalize_ops/adreno/convolution.py | 2 +-
python/tvm/relax/transform/legalize_ops/ccl.py | 38 +-
python/tvm/relax/transform/legalize_ops/common.py | 4 +-
python/tvm/relax/transform/legalize_ops/create.py | 18 +-
.../relax/transform/legalize_ops/distributed.py | 10 +-
python/tvm/relax/transform/legalize_ops/index.py | 12 +-
.../relax/transform/legalize_ops/linear_algebra.py | 22 +-
.../tvm/relax/transform/legalize_ops/manipulate.py | 32 +-
python/tvm/relax/transform/legalize_ops/nn.py | 22 +-
python/tvm/relax/transform/legalize_ops/search.py | 4 +-
.../relax/transform/legalize_ops/statistical.py | 8 +-
python/tvm/relax/transform/legalize_ops/vision.py | 2 +-
.../relax/transform/lower_gpu_ipc_alloc_storage.py | 6 +-
.../relax/transform/optimize_layout_transform.py | 4 +-
.../relax/transform/remove_redundant_reshape.py | 4 +-
python/tvm/relax/transform/transform.py | 24 +-
python/tvm/relax/ty.py | 45 +-
python/tvm/relax/{struct_info.py => type.py} | 112 +-
python/tvm/relax/utils.py | 44 +-
python/tvm/runtime/script_printer.py | 16 +-
python/tvm/s_tir/dlight/benchmark/bench.py | 8 +-
python/tvm/s_tir/dlight/benchmark/extract.py | 22 +-
python/tvm/s_tir/dlight/benchmark/utils.py | 8 +-
python/tvm/tirx/op.py | 8 +-
src/ir/module.cc | 1 -
src/ir/type.cc | 3 +-
src/ir/type_functor.cc | 93 -
src/relax/analysis/analysis.cc | 4 +-
src/relax/analysis/computable_at_compile_time.cc | 2 +-
.../{struct_info_analysis.cc => type_analysis.cc} | 557 +++--
src/relax/analysis/well_formed.cc | 163 +-
.../backend/adreno/annotate_custom_storage.cc | 205 +-
.../backend/adreno/fold_vdevice_scope_change.cc | 19 +-
src/relax/backend/contrib/codegen_c/codegen_c.h | 24 +-
.../backend/contrib/codegen_json/codegen_json.h | 33 +-
src/relax/backend/contrib/cublas/codegen.cc | 6 +-
src/relax/backend/contrib/cudnn/codegen.cc | 9 +-
src/relax/backend/contrib/cutlass/codegen.cc | 18 +-
src/relax/backend/contrib/tensorrt/codegen.cc | 6 +-
src/relax/backend/vm/codegen_vm.cc | 8 +-
src/relax/backend/vm/codegen_vm_tir.cc | 2 +-
src/relax/backend/vm/lower_runtime_builtin.cc | 36 +-
src/relax/backend/vm/vm_shape_lower.cc | 167 +-
src/relax/distributed/axis_group_graph.cc | 72 +-
.../distributed/transform/legalize_redistribute.cc | 10 +-
src/relax/distributed/transform/lower_distir.cc | 95 +-
.../transform/lower_global_view_to_local_view.cc | 28 +-
.../distributed/transform/propagate_sharding.cc | 155 +-
src/relax/distributed/transform/utils.cc | 30 +-
src/relax/distributed/transform/utils.h | 8 +-
src/relax/distributed/{struct_info.cc => type.cc} | 29 +-
src/relax/ir/block_builder.cc | 194 +-
src/relax/ir/dataflow_block_rewriter.cc | 2 +-
src/relax/ir/dataflow_expr_rewriter.cc | 28 +-
src/relax/ir/dataflow_matcher.cc | 22 +-
src/relax/ir/dataflow_matcher.h | 4 +-
src/relax/ir/dataflow_pattern.cc | 27 +-
src/relax/ir/dataflow_pattern_functor.cc | 4 +-
src/relax/ir/dependent_type.cc | 219 ++
src/relax/ir/emit_te.cc | 13 +-
src/relax/ir/emit_te.h | 2 +-
src/relax/ir/expr.cc | 282 ++-
src/relax/ir/expr_functor.cc | 215 +-
src/relax/ir/py_expr_functor.cc | 8 +-
src/relax/ir/struct_info.cc | 244 ---
src/relax/ir/struct_info_functor.cc | 156 --
src/relax/ir/transform.cc | 14 +-
src/relax/ir/type.cc | 54 +-
src/relax/ir/type_functor.cc | 154 ++
src/relax/op/ccl/ccl.cc | 43 +-
src/relax/op/distributed/binary.cc | 60 +-
src/relax/op/distributed/binary.h | 47 +-
src/relax/op/distributed/ccl.cc | 20 +-
src/relax/op/distributed/distributed.cc | 88 +-
src/relax/op/distributed/linear_algebra.cc | 33 +-
src/relax/op/distributed/linear_algebra.h | 4 +-
src/relax/op/distributed/manipulate.cc | 76 +-
src/relax/op/distributed/manipulate.h | 6 +-
src/relax/op/distributed/nn.cc | 30 +-
src/relax/op/distributed/nn.h | 4 +-
src/relax/op/distributed/op.cc | 18 +-
src/relax/op/distributed/statistical.cc | 39 +-
src/relax/op/distributed/statistical.h | 10 +-
src/relax/op/distributed/unary.cc | 66 +-
src/relax/op/distributed/unary.h | 50 +-
src/relax/op/distributed/utils.cc | 71 +-
src/relax/op/distributed/utils.h | 20 +-
src/relax/op/image/resize.cc | 104 +-
src/relax/op/memory/view.cc | 125 +-
src/relax/op/nn/attention.cc | 44 +-
src/relax/op/nn/convolution.cc | 174 +-
src/relax/op/nn/nn.cc | 427 ++--
src/relax/op/nn/nn.h | 4 +-
src/relax/op/nn/pooling.cc | 138 +-
src/relax/op/op.cc | 629 +++---
src/relax/op/op_common.cc | 49 +-
src/relax/op/op_common.h | 209 +-
src/relax/op/tensor/binary.cc | 96 +-
src/relax/op/tensor/binary.h | 12 +-
src/relax/op/tensor/create.cc | 122 +-
src/relax/op/tensor/datatype.cc | 26 +-
src/relax/op/tensor/grad.cc | 42 +-
src/relax/op/tensor/index.cc | 235 +--
src/relax/op/tensor/inspect.cc | 166 +-
src/relax/op/tensor/inspect.h | 32 +-
src/relax/op/tensor/linear_algebra.cc | 127 +-
src/relax/op/tensor/manipulate.cc | 1153 +++++-----
src/relax/op/tensor/qdq.cc | 132 +-
src/relax/op/tensor/sampling.cc | 54 +-
src/relax/op/tensor/search.cc | 115 +-
src/relax/op/tensor/set.cc | 69 +-
src/relax/op/tensor/sorting.cc | 53 +-
src/relax/op/tensor/statistical.cc | 99 +-
src/relax/op/tensor/statistical.h | 2 +-
src/relax/op/tensor/ternary.cc | 32 +-
src/relax/op/tensor/unary.cc | 8 +-
src/relax/op/tensor/unary.h | 10 +-
src/relax/op/vision/multibox_transform_loc.cc | 64 +-
src/relax/op/vision/nms.cc | 167 +-
src/relax/op/vision/roi_align.cc | 30 +-
src/relax/op/vision/roi_pool.cc | 30 +-
src/relax/script/builder/distributed.cc | 29 +-
src/relax/script/builder/frame.cc | 4 +-
src/relax/script/builder/ir.cc | 43 +-
src/relax/script/builder/utils.h | 6 +-
src/relax/script/printer/binding.cc | 18 +-
src/relax/script/printer/call.cc | 36 +-
.../printer/{struct_info.cc => dependent_type.cc} | 72 +-
src/relax/script/printer/distributed.cc | 24 +-
src/relax/script/printer/expr.cc | 8 +-
src/relax/script/printer/function.cc | 20 +-
src/relax/script/printer/type.cc | 25 -
src/relax/script/printer/utils.h | 41 +-
src/relax/training/utils.cc | 18 +-
src/relax/transform/adjust_matmul_order.cc | 8 +-
src/relax/transform/allocate_workspace.cc | 19 +-
src/relax/transform/alter_op_impl.cc | 119 +-
src/relax/transform/attach_global_symbol.cc | 6 +-
src/relax/transform/bind_params.cc | 41 +-
src/relax/transform/bind_symbolic_vars.cc | 1 -
src/relax/transform/bundle_model_params.cc | 6 +-
src/relax/transform/call_tir_rewrite.cc | 55 +-
src/relax/transform/canonicalize_bindings.cc | 34 +-
src/relax/transform/combine_parallel_matmul.cc | 18 +-
src/relax/transform/compute_prim_value.cc | 2 +-
src/relax/transform/convert_layout.cc | 31 +-
src/relax/transform/dataflow_inplace.cc | 72 +-
src/relax/transform/decompose_ops.cc | 50 +-
src/relax/transform/eliminate_common_subexpr.cc | 13 +-
src/relax/transform/expand_tuple_arguments.cc | 25 +-
src/relax/transform/fold_constant.cc | 54 +-
src/relax/transform/fuse_ops.cc | 61 +-
src/relax/transform/fuse_tir.cc | 111 +-
src/relax/transform/gradient.cc | 86 +-
src/relax/transform/gradient_simplifier.cc | 16 +-
src/relax/transform/infer_amp_utils.cc | 12 +-
src/relax/transform/infer_amp_utils.h | 4 +-
src/relax/transform/infer_layout_utils.cc | 12 +-
src/relax/transform/infer_layout_utils.h | 8 +-
src/relax/transform/inline_functions.cc | 2 +-
src/relax/transform/kill_after_last_use.cc | 13 +-
src/relax/transform/lambda_lift.cc | 31 +-
src/relax/transform/lazy_transform_params.cc | 24 +-
src/relax/transform/legalize_ops.cc | 55 +-
src/relax/transform/lift_transform_params.cc | 53 +-
src/relax/transform/lower_alloc_tensor.cc | 6 +-
src/relax/transform/merge_composite_functions.cc | 6 +-
src/relax/transform/normalize.cc | 12 +-
src/relax/transform/realize_vdevice.cc | 38 +-
src/relax/transform/remove_purity_checking.cc | 11 +-
src/relax/transform/remove_unused_outputs.cc | 36 +-
src/relax/transform/remove_unused_parameters.cc | 18 +-
.../transform/reorder_permute_dims_after_concat.cc | 12 +-
src/relax/transform/reorder_take_after_matmul.cc | 33 +-
src/relax/transform/rewrite_cuda_graph.cc | 58 +-
src/relax/transform/rewrite_dataflow_reshape.cc | 20 +-
src/relax/transform/run_codegen.cc | 13 +-
.../specialize_primfunc_based_on_callsite.cc | 58 +-
src/relax/transform/split_call_tir_by_pattern.cc | 4 +-
.../transform/split_layout_rewrite_preproc.cc | 16 +-
src/relax/transform/static_plan_block_memory.cc | 40 +-
src/relax/transform/to_mixed_precision.cc | 48 +-
src/relax/transform/to_non_dataflow.cc | 3 +-
src/relax/transform/topological_sort.cc | 2 +-
...e_param_struct_info.cc => update_param_type.cc} | 24 +-
src/relax/transform/update_vdevice.cc | 11 +-
src/relax/transform/utils.cc | 30 +-
src/relax/transform/utils.h | 50 +-
src/relax/utils.cc | 71 +-
src/runtime/vm/builtin.cc | 2 +-
src/script/ir_builder/ir/ir.cc | 25 +-
src/script/printer/script_printer.cc | 4 +-
src/tirx/ir/function.cc | 31 +-
src/tirx/script/builder/ir.cc | 2 +-
tests/cpp/nested_msg_test.cc | 46 +-
.../python/codegen/test_target_codegen_cuda_fp8.py | 14 +-
.../python/contrib/test_tir_triton_integration.py | 2 +-
tests/python/disco/test_loader.py | 4 +-
tests/python/disco/test_nvshmem.py | 2 +-
tests/python/disco/test_session.py | 6 +-
tests/python/nightly/test_nnapi/infrastructure.py | 10 +-
tests/python/relax/backend/adreno/mod_utils.py | 4 +-
.../adreno/test_transform_annotate_custom_scope.py | 24 +-
.../test_transform_fold_vdevice_scope_change.py | 14 +-
tests/python/relax/backend/adreno/utils.py | 4 +-
...r_sinfo.py => test_distributed_dtensor_type.py} | 44 +-
.../test_distributed_transform_lower_distir.py | 28 +-
...ributed_transform_lower_global_to_local_view.py | 152 +-
...est_distributed_transform_propagate_sharding.py | 200 +-
.../test_distributed_tvmscript_parser.py | 10 +-
.../test_distributed_tvmscript_printer.py | 30 +-
tests/python/relax/test_analysis.py | 28 +-
.../relax/test_analysis_contains_impure_call.py | 4 +-
.../relax/test_analysis_estimate_memory_usage.py | 2 +-
..._analysis.py => test_analysis_type_analysis.py} | 460 ++--
tests/python/relax/test_analysis_well_formed.py | 236 +--
tests/python/relax/test_ast_printer.py | 127 +-
.../python/relax/test_backend_dispatch_sampling.py | 4 +-
.../relax/test_backend_transform_shape_lower.py | 126 +-
tests/python/relax/test_base_py_module_printer.py | 36 +-
.../relax/test_base_py_module_symbolic_shape.py | 20 +-
tests/python/relax/test_bind_symbolic_vars.py | 24 +-
tests/python/relax/test_blockbuilder_core.py | 198 +-
tests/python/relax/test_blockbuilder_emit_te.py | 4 +-
tests/python/relax/test_codegen_coreml.py | 44 +-
tests/python/relax/test_codegen_cutlass.py | 48 +-
tests/python/relax/test_codegen_tensorrt.py | 2 +-
tests/python/relax/test_contrib_vllm.py | 6 +-
tests/python/relax/test_dataflow_inplace.py | 34 +-
tests/python/relax/test_dataflow_pattern.py | 56 +-
tests/python/relax/test_dataflow_rewriter.py | 100 +-
...eliminate_pad_branch_using_buffer_assumption.py | 12 +-
tests/python/relax/test_expr.py | 84 +-
tests/python/relax/test_expr_functor.py | 12 +-
tests/python/relax/test_frontend_common.py | 10 +-
.../relax/test_frontend_from_exported_program.py | 32 +-
tests/python/relax/test_frontend_from_fx.py | 4 +-
.../python/relax/test_frontend_nn_extern_module.py | 4 +-
tests/python/relax/test_frontend_nn_modules.py | 6 +-
tests/python/relax/test_frontend_nn_op.py | 18 +-
tests/python/relax/test_frontend_onnx.py | 6 +-
tests/python/relax/test_frontend_tflite.py | 84 +-
tests/python/relax/test_inline_functions.py | 6 +-
tests/python/relax/test_kill_after_last_use.py | 4 +-
tests/python/relax/test_op_binary.py | 164 +-
tests/python/relax/test_op_ccl.py | 176 +-
tests/python/relax/test_op_create.py | 640 +++---
tests/python/relax/test_op_datatype.py | 72 +-
tests/python/relax/test_op_distributed.py | 4 +-
tests/python/relax/test_op_grad.py | 26 +-
tests/python/relax/test_op_gradient_numeric.py | 26 +-
tests/python/relax/test_op_image.py | 162 +-
tests/python/relax/test_op_index.py | 544 +++--
tests/python/relax/test_op_linear_algebra.py | 190 +-
tests/python/relax/test_op_manipulate.py | 2224 +++++++++-----------
tests/python/relax/test_op_misc.py | 18 +-
tests/python/relax/test_op_nn.py | 824 ++++----
tests/python/relax/test_op_nn_convolution.py | 606 +++---
tests/python/relax/test_op_nn_pooling.py | 774 +++----
tests/python/relax/test_op_qdq.py | 34 +-
tests/python/relax/test_op_sampling.py | 4 +-
tests/python/relax/test_op_search.py | 344 ++-
tests/python/relax/test_op_set.py | 554 +++--
tests/python/relax/test_op_sort.py | 150 +-
tests/python/relax/test_op_statistical.py | 284 ++-
tests/python/relax/test_op_ternary.py | 82 +-
tests/python/relax/test_op_unary.py | 74 +-
tests/python/relax/test_op_view.py | 114 +-
tests/python/relax/test_op_vision.py | 170 +-
.../python/relax/test_optimize_layout_transform.py | 34 +-
tests/python/relax/test_pipeline.py | 6 +-
tests/python/relax/test_pytorch_integration.py | 4 +-
tests/python/relax/test_relax_operators.py | 26 +-
.../python/relax/test_relax_to_pyfunc_converter.py | 14 +-
tests/python/relax/test_tir_call_source_kernel.py | 2 +-
tests/python/relax/test_training_append_loss.py | 2 +-
tests/python/relax/test_training_loss.py | 46 +-
tests/python/relax/test_training_setup_trainer.py | 14 +-
.../python/relax/test_training_trainer_numeric.py | 16 +-
tests/python/relax/test_transform.py | 30 +-
.../relax/test_transform_adjust_matmul_order.py | 10 +-
tests/python/relax/test_transform_alter_op_impl.py | 40 +-
...st_transform_attach_attr_layout_free_buffers.py | 24 +-
tests/python/relax/test_transform_bind_params.py | 22 +-
.../relax/test_transform_bind_symbolic_vars.py | 28 +-
.../relax/test_transform_canonicalize_bindings.py | 22 +-
tests/python/relax/test_transform_codegen_pass.py | 8 +-
.../relax/test_transform_compute_prim_value.py | 8 +-
tests/python/relax/test_transform_cse.py | 8 +-
.../relax/test_transform_dead_code_elimination.py | 6 +-
tests/python/relax/test_transform_decompose_ops.py | 2 +-
.../relax/test_transform_error_enrichment.py | 10 +-
tests/python/relax/test_transform_fold_constant.py | 4 +-
tests/python/relax/test_transform_fuse_ops.py | 68 +-
.../relax/test_transform_fuse_ops_by_pattern.py | 6 +-
tests/python/relax/test_transform_fuse_tir.py | 114 +-
.../relax/test_transform_fuse_transpose_matmul.py | 8 +-
.../relax/test_transform_gradient_checkpoint.py | 22 +-
.../relax/test_transform_gradient_numeric.py | 8 +-
.../relax/test_transform_gradient_te_register.py | 32 +-
.../test_transform_inline_private_functions.py | 2 +-
tests/python/relax/test_transform_lambda_lift.py | 14 +-
.../relax/test_transform_lazy_transform_params.py | 126 +-
tests/python/relax/test_transform_legalize_ops.py | 30 +-
.../relax/test_transform_legalize_ops_binary.py | 8 +-
.../relax/test_transform_legalize_ops_ccl.py | 22 +-
.../test_transform_legalize_ops_create_datatype.py | 2 +-
.../test_transform_legalize_ops_distributed.py | 4 +-
.../relax/test_transform_legalize_ops_grad.py | 14 +-
..._transform_legalize_ops_index_linear_algebra.py | 20 +-
.../test_transform_legalize_ops_manipulate.py | 44 +-
.../python/relax/test_transform_legalize_ops_nn.py | 64 +-
.../relax/test_transform_legalize_ops_qdq.py | 24 +-
...st_transform_legalize_ops_search_statistical.py | 16 +-
.../relax/test_transform_lift_transform_params.py | 24 +-
.../test_transform_lower_gpu_ipc_alloc_storage.py | 4 +-
.../test_transform_merge_composite_functions.py | 4 +-
.../relax/test_transform_meta_schedule_tuning.py | 4 +-
tests/python/relax/test_transform_normalize.py | 28 +-
.../relax/test_transform_normalize_global_var.py | 2 +-
...st_transform_operator_specific_normalization.py | 22 +-
.../relax/test_transform_rewrite_cuda_graph.py | 76 +-
.../test_transform_rewrite_dataflow_reshape.py | 42 +-
...nsform_specialize_primfunc_based_on_callsite.py | 48 +-
.../test_transform_split_layout_rewrite_preproc.py | 22 +-
.../test_transform_static_plan_block_memory.py | 14 +-
.../relax/test_transform_to_mixed_precision.py | 4 +-
...info.py => test_transform_update_param_type.py} | 10 +-
tests/python/relax/test_tvmscript_ir_builder.py | 62 +-
tests/python/relax/test_tvmscript_parser.py | 170 +-
tests/python/relax/test_tvmscript_printer_relax.py | 181 +-
.../relax/{test_struct_info.py => test_type.py} | 108 +-
tests/python/relax/test_utils.py | 2 +-
tests/python/relax/test_vm_build.py | 22 +-
tests/python/relax/test_vm_builtin.py | 2 +-
tests/python/relax/test_vm_builtin_lower.py | 12 +-
tests/python/relax/test_vm_codegen_only.py | 24 +-
tests/python/relax/test_vm_codegen_tir.py | 8 +-
tests/python/relax/test_vm_cuda_graph.py | 6 +-
tests/python/s_tir/dlight/test_benchmark.py | 12 +-
tests/python/tirx-base/test_tir_specialize.py | 20 +-
.../python/tvmscript/test_tvmscript_parser_tir.py | 62 +-
tests/python/tvmscript/test_tvmscript_roundtrip.py | 30 +-
442 files changed, 13578 insertions(+), 15301 deletions(-)
diff --git a/docs/arch/fusion.rst b/docs/arch/fusion.rst
index 5f4a575c08..edb991d93b 100644
--- a/docs/arch/fusion.rst
+++ b/docs/arch/fusion.rst
@@ -195,9 +195,9 @@ via ``cls.func_name``:
@R.function
def main(x: R.Tensor((10, 20), "float32")):
with R.dataflow():
- lv0 = R.call_tir(add, (x, const_1), out_sinfo=R.Tensor((10, 20), "float32"))
- lv1 = R.call_tir(exp, (lv0,), out_sinfo=R.Tensor((10, 20), "float32"))
- gv = R.call_tir(squeeze, (lv1,), out_sinfo=R.Tensor((10, 20), "float32"))
+ lv0 = R.call_tir(add, (x, const_1), out_ty=R.Tensor((10, 20), "float32"))
+ lv1 = R.call_tir(exp, (lv0,), out_ty=R.Tensor((10, 20), "float32"))
+ gv = R.call_tir(squeeze, (lv1,), out_ty=R.Tensor((10, 20), "float32"))
R.output(gv)
return gv
diff --git a/docs/arch/tvmscript.rst b/docs/arch/tvmscript.rst
index 8dbeb18448..1d5efa8de0 100644
--- a/docs/arch/tvmscript.rst
+++ b/docs/arch/tvmscript.rst
@@ -116,7 +116,7 @@ These can be composed:
y: R.Tensor((128,), "float32")) -> R.Tensor((128,), "float32"):
with R.dataflow():
out = R.call_tir(cls.add_kernel, (x, y),
- out_sinfo=R.Tensor((128,), "float32"))
+ out_ty=R.Tensor((128,), "float32"))
R.output(out)
return out
@@ -268,18 +268,18 @@ The Relax builder (``ir_builder/relax/ir.py``) provides:
**Emit**:
- ``R.emit(value)`` → emit a binding, returns a ``Var``
-- ``R.emit_match_cast(value, struct_info)`` → emit with type assertion
+- ``R.emit_match_cast(value, ty)`` → emit with type assertion
**Type annotations**:
-- ``R.Tensor(shape, dtype)`` — tensor struct info
-- ``R.Tuple(*fields)`` — tuple struct info
-- ``R.Shape(values)`` — shape struct info
-- ``R.Object()`` — opaque object struct info
+- ``R.Tensor(shape, dtype)`` — tensor type
+- ``R.Tuple(*fields)`` — tuple type
+- ``R.Shape(values)`` — shape type
+- ``R.Object()`` — opaque object type
**Calling conventions**:
-- ``R.call_tir(func, args, out_sinfo)`` — call a TIR function
+- ``R.call_tir(func, args, out_ty)`` — call a TIR function
- ``R.call_packed(name, *args)`` — call a PackedFunc
- ``R.call_dps_packed(func, *args)`` — call using destination-passing style
@@ -326,7 +326,7 @@ It maintains:
Each IR dialect registers its own converters:
- ``src/script/printer/tirx/`` — converts PrimFunc, Buffer, SBlock, loops, expressions.
-- ``src/script/printer/relax/`` — converts relax.Function, bindings, struct info, operators.
+- ``src/script/printer/relax/`` — converts relax.Function, bindings, types, operators.
- ``src/script/printer/ir/`` — converts IRModule, shared types.
The final step calls ``DocToPythonScript()`` (``src/script/printer/doc_printer/python_doc_printer.cc``)
@@ -491,8 +491,8 @@ Function definition
# function body
return result
-- ``R.Tensor(shape, dtype)`` — tensor type annotation (struct info).
-- ``R.Tuple(...)``, ``R.Shape(...)``, ``R.Object()`` — other struct info types.
+- ``R.Tensor(shape, dtype)`` — tensor type annotation.
+- ``R.Tuple(...)``, ``R.Shape(...)``, ``R.Object()`` — other Relax type annotations.
- ``R.function(private=True)`` — marks the function as module-private.
- ``R.function(pure=False)`` — marks the function as having side effects.
@@ -514,10 +514,10 @@ Calling TIR functions
.. code-block:: python
- out = R.call_tir(cls.my_kernel, (x, y), out_sinfo=R.Tensor((128,), "float32"))
+ out = R.call_tir(cls.my_kernel, (x, y), out_ty=R.Tensor((128,), "float32"))
- ``cls.my_kernel`` — references a TIR ``PrimFunc`` in the same module.
-- ``out_sinfo`` — the struct info (shape and dtype) of the output tensor.
+- ``out_ty`` — the type (shape and dtype) of the output tensor.
Control flow
~~~~~~~~~~~~
diff --git a/docs/deep_dive/relax/dpl.rst b/docs/deep_dive/relax/dpl.rst
index 5389e0acc6..9a4e8369fa 100644
--- a/docs/deep_dive/relax/dpl.rst
+++ b/docs/deep_dive/relax/dpl.rst
@@ -123,7 +123,7 @@ Any pattern can be further narrowed by attaching constraints:
- ``.has_dtype(dtype)`` -- the matched expression must have the given data type.
- ``.has_shape(shape)`` -- the matched expression must have the given shape.
- ``.has_attr(attrs)`` -- the matched call must carry the given attributes.
-- ``.has_struct_info(struct_info)`` -- the matched expression must have the given struct info.
+- ``.has_ty(ty)`` -- the matched expression must have the given type.
.. code:: python
@@ -286,7 +286,7 @@ The callback receives *variables* rather than expressions:
...
- ``matchings[pat]`` returns the **bound variable** (``Var``) whose right-hand
- side matched ``pat``. The ``Var`` itself carries ``struct_info`` and can be
+ side matched ``pat``. The ``Var`` itself carries ``ty`` and can be
used directly in new expressions.
- ``bindings`` maps each ``Var`` to its bound ``Expr`` (the right-hand side),
useful when you need to inspect the original expression.
@@ -311,7 +311,7 @@ The callback receives *variables* rather than expressions:
W1 = matchings[w1]
W2 = matchings[w2]
W3 = matchings[w3]
- width = W1.struct_info.shape[1]
+ width = W1.ty.shape[1]
concat_w = R.concat([W1, W2, W3], axis=1)
merged = R.matmul(inp, concat_w)
@@ -361,7 +361,7 @@ object that can be applied directly.
"my_fast_add",
A,
B,
- sinfo_args=R.Tensor([16], "float32"),
+ ty_args=R.Tensor([16], "float32"),
)
return C
@@ -440,7 +440,7 @@ structurally (dtype restrictions, shape compatibility, attribute values, etc.):
def my_check_fn(ctx: PatternCheckContext) -> bool:
matmul_expr = ctx.annotated_expr["matmul"]
# Only accept float16 output
- if matmul_expr.struct_info.dtype != "float16":
+ if matmul_expr.ty.dtype != "float16":
return False
return True
@@ -464,7 +464,7 @@ sub-function.
def check(ctx):
transpose_call = ctx.annotated_expr["wT"]
- ndim = transpose_call.args[0].struct_info.ndim
+ ndim = transpose_call.args[0].ty.ndim
if ndim == -1:
return False
if ndim == 2 and transpose_call.attrs.axes is None:
@@ -513,7 +513,7 @@ Quick Reference
- Match ``R.call_packed``
* - ``make_fused_bias_activation_pattern(...)``
- Build ``op + bias + activation`` chain
- * - ``.has_dtype()`` / ``.has_shape()`` / ``.has_attr()`` / ``.has_struct_info()``
+ * - ``.has_dtype()`` / ``.has_shape()`` / ``.has_attr()`` / ``.has_ty()``
- Attach constraints
* - ``|`` / ``&`` / ``~``
- Or / And / Not combinators
diff --git a/docs/deep_dive/relax/learning.rst b/docs/deep_dive/relax/learning.rst
index 59b97daaa0..d9c169150e 100644
--- a/docs/deep_dive/relax/learning.rst
+++ b/docs/deep_dive/relax/learning.rst
@@ -170,9 +170,9 @@ for the end-to-end model execution. The code block below shows a TVMScript imple
cls = Module
n = T.int64()
with R.dataflow():
- lv = R.call_tir(cls.linear, (x, w0, b0), out_sinfo=R.Tensor((n, 256), dtype="float32"))
- lv1 = R.call_tir(cls.relu, (lv,), out_sinfo=R.Tensor((n, 256), dtype="float32"))
- lv2 = R.call_tir(cls.linear, (lv1, w1, b1), out_sinfo=R.Tensor((n, 10), dtype="float32"))
+ lv = R.call_tir(cls.linear, (x, w0, b0), out_ty=R.Tensor((n, 256), dtype="float32"))
+ lv1 = R.call_tir(cls.relu, (lv,), out_ty=R.Tensor((n, 256), dtype="float32"))
+ lv2 = R.call_tir(cls.linear, (lv1, w1, b1), out_ty=R.Tensor((n, 10), dtype="float32"))
R.output(lv2)
return lv2
@@ -194,10 +194,10 @@ Key Elements of Relax
This section will introduce the key elements of Relax abstraction and how it enables optimization
in ML compilers.
-Structure Info
-~~~~~~~~~~~~~~
-Structure info is a new concept in Relax that represents the type of relax expressions. It can
-be ``TensorStructInfo``, ``TupleStructInfo``, etc. In the above example, we use ``TensorStructInfo``
+Type
+~~~~
+Type is the Relax representation of expression type information. It can
+be ``TensorType``, ``TupleType``, etc. In the above example, we use ``TensorType``
(short in ``R.Tensor`` in TVMScript) to represent the shape and dtype of the tensor of the inputs,
outputs, and intermediate results.
@@ -210,7 +210,7 @@ Taking one line from the above code as an example:
.. code:: python
- lv = R.call_tir(cls.linear, (x, w0, b0), out_sinfo=R.Tensor((n, 256), dtype="float32"))
+ lv = R.call_tir(cls.linear, (x, w0, b0), out_ty=R.Tensor((n, 256), dtype="float32"))
To explain what does ``R.call_tir`` work, let us review an equivalent low-level numpy
implementation of the operation, as follows:
@@ -238,9 +238,9 @@ Another important element in a relax function is the R.dataflow() scope annotati
.. code:: python
with R.dataflow():
- lv = R.call_tir(cls.linear, (x, w0, b0), out_sinfo=R.Tensor((n, 256), dtype="float32"))
- lv1 = R.call_tir(cls.relu, (lv,), out_sinfo=R.Tensor((n, 256), dtype="float32"))
- lv2 = R.call_tir(cls.linear, (lv1, w1, b1), out_sinfo=R.Tensor((n, 10), dtype="float32"))
+ lv = R.call_tir(cls.linear, (x, w0, b0), out_ty=R.Tensor((n, 256), dtype="float32"))
+ lv1 = R.call_tir(cls.relu, (lv,), out_ty=R.Tensor((n, 256), dtype="float32"))
+ lv2 = R.call_tir(cls.linear, (lv1, w1, b1), out_ty=R.Tensor((n, 10), dtype="float32"))
R.output(lv2)
Before we talk about the dataflow block, let us first introduce the concept of **pure** and
diff --git a/docs/deep_dive/relax/tutorials/relax_creation.py b/docs/deep_dive/relax/tutorials/relax_creation.py
index e0f8e2c613..5ae3ef4b59 100644
--- a/docs/deep_dive/relax/tutorials/relax_creation.py
+++ b/docs/deep_dive/relax/tutorials/relax_creation.py
@@ -254,7 +254,7 @@ with bb.function("forward", [x, fc1_weight, fc1_bias, fc2_weight, fc2_bias]):
relax.call_dps_packed(
"env.linear",
[x, fc1_weight, fc1_bias],
- out_sinfo=relax.TensorStructInfo((n, 128), "float32"),
+ out_ty=relax.TensorType((n, 128), "float32"),
)
)
lv1 = bb.emit_te(topi.nn.relu, lv0)
@@ -263,7 +263,7 @@ with bb.function("forward", [x, fc1_weight, fc1_bias, fc2_weight, fc2_bias]):
relax.call_tir(
tir_gv,
[lv1, fc2_weight, fc2_bias],
- out_sinfo=relax.TensorStructInfo((n, 10), "float32"),
+ out_ty=relax.TensorType((n, 10), "float32"),
)
)
bb.emit_output(gv)
diff --git a/docs/how_to/tutorials/mix_python_and_tvm_with_pymodule.py b/docs/how_to/tutorials/mix_python_and_tvm_with_pymodule.py
index 6a3be7622f..9515f6874d 100644
--- a/docs/how_to/tutorials/mix_python_and_tvm_with_pymodule.py
+++ b/docs/how_to/tutorials/mix_python_and_tvm_with_pymodule.py
@@ -99,9 +99,7 @@ if RUN_EXAMPLE:
"""Takes PyTorch tensors, calls TIR, returns PyTorch tensors."""
x_tvm = self._convert_pytorch_to_tvm(x)
y_tvm = self._convert_pytorch_to_tvm(y)
- result = self.call_tir(
- self.add_tir, [x_tvm, y_tvm], out_sinfo=R.Tensor((4,), "float32")
- )
+ result = self.call_tir(self.add_tir, [x_tvm, y_tvm], out_ty=R.Tensor((4,), "float32"))
return self._convert_tvm_to_pytorch(result)
# TIR functions are JIT-compiled at instantiation
@@ -157,7 +155,7 @@ if RUN_EXAMPLE:
out = self.call_tir(
self.matmul_tir,
[x_tvm, w_tvm],
- out_sinfo=R.Tensor((x.shape[0], 3), "float32"),
+ out_ty=R.Tensor((x.shape[0], 3), "float32"),
)
logits = self._convert_tvm_to_pytorch(out)
@@ -231,7 +229,7 @@ if RUN_EXAMPLE:
h = self.call_tir(
self.matmul_tir,
[x_tvm, w_tvm],
- out_sinfo=R.Tensor((2, 3), "float32"),
+ out_ty=R.Tensor((2, 3), "float32"),
)
h_pt = self._convert_tvm_to_pytorch(h)
@@ -239,7 +237,7 @@ if RUN_EXAMPLE:
h_biased = self.call_dps_packed(
"my_bias_add",
[h_pt, bias],
- out_sinfo=R.Tensor((2, 3), "float32"),
+ out_ty=R.Tensor((2, 3), "float32"),
)
# 3. Python/PyTorch activation
@@ -294,7 +292,7 @@ if RUN_EXAMPLE:
h_bias = R.call_tir(
cls.bias_add_tir,
(h, b),
- out_sinfo=R.Tensor((2, 4), "float32"),
+ out_ty=R.Tensor((2, 4), "float32"),
)
return R.nn.relu(h_bias)
@@ -364,8 +362,8 @@ if RUN_EXAMPLE:
x: R.Tensor((4, 8), "float32"),
) -> R.Tensor((4, 8), "float32"):
# The VM calls back into Python for these two ops
- h = R.call_py_func("layer_norm", (x,), out_sinfo=R.Tensor((4, 8), "float32"))
- out = R.call_py_func("silu", (h,), out_sinfo=R.Tensor((4, 8), "float32"))
+ h = R.call_py_func("layer_norm", (x,), out_ty=R.Tensor((4, 8), "float32"))
+ out = R.call_py_func("silu", (h,), out_ty=R.Tensor((4, 8), "float32"))
return out
mod = HybridVMModule(device=tvm.cpu(0))
@@ -438,7 +436,7 @@ if RUN_EXAMPLE:
# Python → TIR with symbolic output shape
n = T.int64()
x7 = torch.randn(7)
- scaled = mod.call_tir("scale_tir", [x7], relax.TensorStructInfo((n,), "float32"))
+ scaled = mod.call_tir("scale_tir", [x7], relax.TensorType((n,), "float32"))
print("scale_tir(len=7):", scaled)
assert torch.allclose(torch.tensor(scaled.numpy()), x7 * 2.0, atol=1e-5)
diff --git a/docs/reference/security.rst b/docs/reference/security.rst
index f723ded053..08fabb1e90 100644
--- a/docs/reference/security.rst
+++ b/docs/reference/security.rst
@@ -62,7 +62,7 @@ Subroutine Cache Hash Collision
``SubroutineMixin._get_subroutine()`` in ``python/tvm/relax/frontend/nn/subroutine.py``
used ``tvm_ffi.structural_hash`` as the sole cache lookup key without a subsequent
-``structural_equal`` verification. If two different ``arg_sinfo`` values produced the
+``structural_equal`` verification. If two different ``arg_ty`` values produced the
same 64-bit hash, the cache would return a previously compiled function with
mismatched parameter shapes, leading to silently incorrect compiled output.
@@ -73,11 +73,11 @@ The issue is primarily a correctness defect rather than a practically exploitabl
security vulnerability.
**Root Cause**: The subroutine cache (``cls._gvar``) was keyed by
-``(structural_hash(arg_sinfo, map_free_vars=True), is_dataflow)``.
+``(structural_hash(arg_ty, map_free_vars=True), is_dataflow)``.
A hash match was treated as proof of structural equality, skipping the necessary
``structural_equal`` check.
-**Fix**: The cache now stores a list of ``(arg_sinfo, result)`` pairs per hash bucket.
+**Fix**: The cache now stores a list of ``(arg_ty, result)`` pairs per hash bucket.
On lookup, each candidate is verified with ``structural_equal`` before returning.
This follows the standard hash-table pattern: hash for bucket selection, equality
for final verification.
diff --git a/include/tvm/ir/expr.h b/include/tvm/ir/expr.h
index eb924be63d..5be6e7c08f 100644
--- a/include/tvm/ir/expr.h
+++ b/include/tvm/ir/expr.h
@@ -30,7 +30,7 @@
#include <tvm/ir/cast.h>
#include <tvm/ir/cow.h>
#include <tvm/ir/source_map.h>
-#include <tvm/ir/type.h>
+#include <tvm/runtime/data_type.h>
#include <algorithm>
#include <functional>
@@ -43,6 +43,48 @@ namespace tvm {
// Forward-declare VirtualDevice to avoid circular imports.
class VirtualDevice;
+/*!
+ * \brief Type is the base type of all types.
+ *
+ * TVM's type system contains following subclasses:
+ *
+ * - PrimType: type of primitive type values used in the low-level IR.
+ * - FuncType: type of a function.
+ * - TensorType: type of certain Tensor values in the expression.
+ *
+ * There are also advanced types to support generic(polymorphic types).
+ * \sa Type
+ */
+class TypeNode : public ffi::Object {
+ public:
+ /*!
+ * \brief Span that points to the original source code.
+ * Reserved debug information.
+ */
+ mutable Span span;
+
+ static void RegisterReflection() {
+ namespace refl = tvm::ffi::reflection;
+ // span do not participate in structural equal and hash.
+ refl::ObjectDef<TypeNode>().def_ro("span", &TypeNode::span, refl::DefaultValue(Span()),
+ refl::AttachFieldFlag::SEqHashIgnore());
+ }
+
+ static constexpr TVMFFISEqHashKind _type_s_eq_hash_kind = kTVMFFISEqHashKindTreeNode;
+
+ static constexpr const uint32_t _type_child_slots = 14;
+ TVM_FFI_DECLARE_OBJECT_INFO("ir.Type", TypeNode, ffi::Object);
+};
+
+/*!
+ * \brief Managed reference to TypeNode.
+ * \sa TypeNode
+ */
+class Type : public ffi::ObjectRef {
+ public:
+ TVM_FFI_DEFINE_OBJECT_REF_METHODS_NULLABLE(Type, ffi::ObjectRef, TypeNode);
+};
+
/*!
* \brief Base type of all the expressions.
* \sa Expr
@@ -55,11 +97,23 @@ class BaseExprNode : public ffi::Object {
*/
mutable Span span;
+ /*!
+ * \brief The deduced or annotated type of the expression.
+ *
+ * This field is intentionally nullable because type information may
+ * be populated by later analysis passes instead of expression
+ * constructors.
+ */
+ mutable Type ty;
+
static void RegisterReflection() {
namespace refl = tvm::ffi::reflection;
- // span do not participate in structural equal and hash.
- refl::ObjectDef<BaseExprNode>().def_ro("span", &BaseExprNode::span, refl::DefaultValue(Span()),
- refl::AttachFieldFlag::SEqHashIgnore());
+ // span and ty do not participate in structural equal and hash.
+ refl::ObjectDef<BaseExprNode>()
+ .def_ro("span", &BaseExprNode::span, refl::DefaultValue(Span()),
+ refl::AttachFieldFlag::SEqHashIgnore())
+ .def_ro("ty", &BaseExprNode::ty, refl::DefaultValue(Type()),
+ refl::AttachFieldFlag::SEqHashIgnore());
}
static constexpr TVMFFISEqHashKind _type_s_eq_hash_kind = kTVMFFISEqHashKindTreeNode;
@@ -414,17 +468,9 @@ TVM_DLL PrimExpr operator~(PrimExpr a);
*/
class RelaxExprNode : public BaseExprNode {
public:
- /*!
- * \brief Stores the result of structure information of the
- * expression that encapsulate both static shape and
- * runtime information such as shape.
- */
- mutable ffi::Optional<ffi::ObjectRef> struct_info_ = ffi::Optional<ffi::ObjectRef>();
-
static void RegisterReflection() {
namespace refl = tvm::ffi::reflection;
- refl::ObjectDef<RelaxExprNode>().def_ro("struct_info_", &RelaxExprNode::struct_info_,
- refl::AttachFieldFlag::SEqHashIgnore());
+ refl::ObjectDef<RelaxExprNode>();
}
static constexpr const uint32_t _type_child_slots = 22;
diff --git a/include/tvm/ir/op.h b/include/tvm/ir/op.h
index 3fd39c1060..640f1e9f01 100644
--- a/include/tvm/ir/op.h
+++ b/include/tvm/ir/op.h
@@ -95,8 +95,6 @@ class OpNode : public RelaxExprNode {
public:
/*! \brief name of the operator */
ffi::String name;
- /*! \brief the type of the operator */
- mutable FuncType op_type;
/*!
* \brief detailed description of the operator
* This can be used to generate docstring automatically for the operator.
@@ -130,7 +128,6 @@ class OpNode : public RelaxExprNode {
namespace refl = tvm::ffi::reflection;
refl::ObjectDef<OpNode>()
.def_ro("name", &OpNode::name)
- .def_ro("op_type", &OpNode::op_type, refl::AttachFieldFlag::SEqHashIgnore())
.def_ro("description", &OpNode::description, refl::AttachFieldFlag::SEqHashIgnore())
.def_ro("arguments", &OpNode::arguments, refl::AttachFieldFlag::SEqHashIgnore())
.def_ro("attrs_type_key", &OpNode::attrs_type_key, refl::AttachFieldFlag::SEqHashIgnore())
diff --git a/include/tvm/ir/type.h b/include/tvm/ir/type.h
index db6151fda7..9c56d03764 100644
--- a/include/tvm/ir/type.h
+++ b/include/tvm/ir/type.h
@@ -52,6 +52,7 @@
#include <tvm/ffi/cast.h>
#include <tvm/ffi/container/array.h>
#include <tvm/ffi/reflection/registry.h>
+#include <tvm/ir/expr.h>
#include <tvm/ir/source_map.h>
#include <tvm/runtime/data_type.h>
@@ -59,48 +60,6 @@
namespace tvm {
-/*!
- * \brief Type is the base type of all types.
- *
- * TVM's type system contains following subclasses:
- *
- * - PrimType: type of primitive type values used in the low-level IR.
- * - FuncType: type of a function.
- * - TensorType: type of certain Tensor values in the expression.
- *
- * There are also advanced types to support generic(polymorphic types).
- * \sa Type
- */
-class TypeNode : public ffi::Object {
- public:
- /*!
- * \brief Span that points to the original source code.
- * Reserved debug information.
- */
- mutable Span span;
-
- static void RegisterReflection() {
- namespace refl = tvm::ffi::reflection;
- // span do not participate in structural equal and hash.
- refl::ObjectDef<TypeNode>().def_ro("span", &TypeNode::span, refl::DefaultValue(Span()),
- refl::AttachFieldFlag::SEqHashIgnore());
- }
-
- static constexpr TVMFFISEqHashKind _type_s_eq_hash_kind = kTVMFFISEqHashKindTreeNode;
-
- static constexpr const uint32_t _type_child_slots = 14;
- TVM_FFI_DECLARE_OBJECT_INFO("ir.Type", TypeNode, ffi::Object);
-};
-
-/*!
- * \brief Managed reference to TypeNode.
- * \sa TypeNode
- */
-class Type : public ffi::ObjectRef {
- public:
- TVM_FFI_DEFINE_OBJECT_REF_METHODS_NULLABLE(Type, ffi::ObjectRef, TypeNode);
-};
-
/*!
* \brief Primitive data types used in the low-level IR.
*
@@ -136,7 +95,7 @@ class PrimType : public Type {
*/
TVM_DLL explicit PrimType(runtime::DataType dtype, Span span = Span());
- TVM_FFI_DEFINE_OBJECT_REF_METHODS_NULLABLE(PrimType, Type, PrimTypeNode);
+ TVM_FFI_DEFINE_OBJECT_REF_METHODS_NOTNULLABLE(PrimType, Type, PrimTypeNode);
};
/*!
@@ -182,7 +141,7 @@ class PointerType : public Type {
*/
TVM_DLL explicit PointerType(Type element_type, ffi::String storage_scope = "");
- TVM_FFI_DEFINE_OBJECT_REF_METHODS_NULLABLE(PointerType, Type, PointerTypeNode);
+ TVM_FFI_DEFINE_OBJECT_REF_METHODS_NOTNULLABLE(PointerType, Type, PointerTypeNode);
};
/*!
@@ -222,7 +181,7 @@ class TupleType : public Type {
*/
TVM_DLL TupleType static Empty();
- TVM_FFI_DEFINE_OBJECT_REF_METHODS_NULLABLE(TupleType, Type, TupleTypeNode);
+ TVM_FFI_DEFINE_OBJECT_REF_METHODS_NOTNULLABLE(TupleType, Type, TupleTypeNode);
};
/*!
@@ -278,7 +237,7 @@ class FuncType : public Type {
*/
TVM_DLL FuncType(ffi::Array<Type> arg_types, Type ret_type, Span span = Span());
- TVM_FFI_DEFINE_OBJECT_REF_METHODS_NULLABLE(FuncType, Type, FuncTypeNode);
+ TVM_FFI_DEFINE_OBJECT_REF_METHODS_NOTNULLABLE(FuncType, Type, FuncTypeNode);
};
/*!
@@ -302,16 +261,7 @@ class TensorMapType : public Type {
public:
TVM_DLL TensorMapType(Span span = Span());
- explicit TensorMapType(::tvm::ffi::ObjectPtr<TensorMapTypeNode> n) : Type(n) {}
- TensorMapType(const TensorMapType&) = default;
- TensorMapType(TensorMapType&&) = default;
- TensorMapType& operator=(const TensorMapType&) = default;
- TensorMapType& operator=(TensorMapType&&) = default;
- const TensorMapTypeNode* operator->() const {
- return static_cast<const TensorMapTypeNode*>(data_.get());
- }
- const TensorMapTypeNode* get() const { return operator->(); }
- using ContainerType = TensorMapTypeNode;
+ TVM_FFI_DEFINE_OBJECT_REF_METHODS_NOTNULLABLE(TensorMapType, Type, TensorMapTypeNode);
};
} // namespace tvm
diff --git a/include/tvm/relax/analysis.h b/include/tvm/relax/analysis.h
index a34a583b2e..248efd66a4 100644
--- a/include/tvm/relax/analysis.h
+++ b/include/tvm/relax/analysis.h
@@ -28,7 +28,7 @@
#include <tvm/ir/module.h>
#include <tvm/relax/expr.h>
#include <tvm/relax/op_attr_types.h>
-#include <tvm/relax/struct_info.h>
+#include <tvm/relax/type.h>
#include <tvm/tirx/function.h>
#include <tvm/tirx/index_map.h>
@@ -70,54 +70,53 @@ TVM_DLL bool CanProveShapeEqual(const ffi::Array<PrimExpr>& lhs, const ffi::Arra
TVM_DLL bool CanProveShapeEqual(const Expr& lhs, const Expr& rhs, const arith::Analyzer& ana);
//-----------------------------------
-// Foundational StructInfo analysis
+// Foundational Type analysis
//-----------------------------------
/*!
- * \brief Get the corresponding static type from a given struct info.
- * \param info The struct info.
+ * \brief Get the corresponding static type from a given type.
+ * \param info The type.
* \return the corresponding static type.
*/
-TVM_DLL Type GetStaticType(const StructInfo& info);
+TVM_DLL Type GetStaticType(const Type& info);
/*!
- * \brief Get the corresponding struct info from static type.
+ * \brief Get the corresponding type from static type.
* \param type The input type
- * \return the corresponding struct info.
+ * \return the corresponding type.
*/
-TVM_DLL StructInfo StructInfoFromType(const Type& type);
+TVM_DLL Type TypeFromStaticType(const Type& type);
/*!
- * \return Derive the call's ret value struct info from inputs.
- * \param finfo The function struct info.
+ * \return Derive the call's ret value type from inputs.
+ * \param finfo The function type.
* \param call The call expression to be derived.
* \param ctx The builder context.
- * \return The derived struct info of the call.
+ * \return The derived type of the call.
* \note call->op field is ignored during derivation and we only rely on information
- * presented by func_sinfo.
+ * presented by func_ty.
*/
-TVM_DLL StructInfo DeriveCallRetStructInfo(const FuncStructInfo& finfo, const Call& call,
- const BlockBuilder& ctx);
+TVM_DLL Type DeriveCallRetType(const FuncType& finfo, const Call& call, const BlockBuilder& ctx);
/*!
- * \brief Derive the call's ret value struct info using a caller-provided analyzer.
- * \param finfo The function struct info.
+ * \brief Derive the call's ret value type using a caller-provided analyzer.
+ * \param finfo The function type.
* \param call The call expression to be derived.
* \param ctx The builder context.
* \param ana Context analyzer to prove symbolic expression equality.
- * \return The derived struct info of the call.
+ * \return The derived type of the call.
*/
-TVM_DLL StructInfo DeriveCallRetStructInfo(const FuncStructInfo& finfo, const Call& call,
- const BlockBuilder& ctx, const arith::Analyzer& ana);
+TVM_DLL Type DeriveCallRetType(const FuncType& finfo, const Call& call, const BlockBuilder& ctx,
+ const arith::Analyzer& ana);
/*!
* \brief Erase the info to a corresponding more coarse grained
- * struct info that is still well-defined(with all the vars in scope).
+ * type that is still well-defined(with all the vars in scope).
*
- * When we are returning a StructInfo to another scope,
- * it is important to remember that StructInfo may carry
+ * When we are returning a Type to another scope,
+ * it is important to remember that Type may carry
* dependencies on var that is not defined the other scope.
*
* In such cases, it is important to call EraseToWellDefined to get
- * another StructInfo that **only** contains the vars that are defined
+ * another Type that **only** contains the vars that are defined
* in the target scope.
*
* For example, consider the following function
@@ -150,10 +149,10 @@ TVM_DLL StructInfo DeriveCallRetStructInfo(const FuncStructInfo& finfo, const Ca
* will give us R.Tensor[(3, m)], where n get replaced by 2.
*
* Use this function in the following scenarios:
- * - Decide the struct_info of expr with sub-scopes, such as If, SeqExpr
- * - Decide the deduced return struct_info of a function that can be fully decided by params.
+ * - Decide the ty of expr with sub-scopes, such as If, SeqExpr
+ * - Decide the deduced return ty of a function that can be fully decided by params.
*
- * \param info The struct info.
+ * \param info The type.
* \param f_shape_var_map callback function to specify
* whether a symbolic shape var is defined and the value it maps to,
* return nullopt if var is undefined.
@@ -161,15 +160,15 @@ TVM_DLL StructInfo DeriveCallRetStructInfo(const FuncStructInfo& finfo, const Ca
* whether a var is defined in the target scope and the value it maps to,
* return nullopt if var is undefined.
*
- * \return the corresponding erased struct info.
+ * \return the corresponding erased type.
*/
-TVM_DLL StructInfo EraseToWellDefined(
- const StructInfo& info,
+TVM_DLL Type EraseToWellDefined(
+ const Type& info,
std::function<ffi::Optional<PrimExpr>(const tirx::Var& var)> f_shape_var_map = nullptr,
std::function<ffi::Optional<Expr>(const Var& var)> f_var_map = nullptr);
/*!
* \brief EraseToWellDefined overload using a caller-provided analyzer.
- * \param info The struct info.
+ * \param info The type.
* \param f_shape_var_map callback function to specify
* whether a symbolic shape var is defined and the value it maps to,
* return nullopt if var is undefined.
@@ -177,16 +176,15 @@ TVM_DLL StructInfo EraseToWellDefined(
* whether a var is defined in the target scope and the value it maps to,
* return nullopt if var is undefined.
* \param ana Context analyzer to prove symbolic expression equality.
- * \return the corresponding erased struct info.
+ * \return the corresponding erased type.
*/
-TVM_DLL StructInfo EraseToWellDefined(
- const StructInfo& info,
- std::function<ffi::Optional<PrimExpr>(const tirx::Var& var)> f_shape_var_map,
+TVM_DLL Type EraseToWellDefined(
+ const Type& info, std::function<ffi::Optional<PrimExpr>(const tirx::Var& var)> f_shape_var_map,
std::function<ffi::Optional<Expr>(const Var& var)> f_var_map, const arith::Analyzer& ana);
/*!
* \brief EraseToWellDefined variant with map.
- * \param info The struct info.
+ * \param info The type.
* \param shape_var_map map to specify
* whether a symbolic shape var is defined and the value it maps to,
* return nullopt if var is undefined.
@@ -194,14 +192,13 @@ TVM_DLL StructInfo EraseToWellDefined(
* whether a var is defined in the target scope and the value it maps to,
* return nullopt if var is undefined.
*
- * \return the corresponding erased struct info.
+ * \return the corresponding erased type.
*/
-TVM_DLL StructInfo EraseToWellDefined(const StructInfo& info,
- ffi::Map<tirx::Var, PrimExpr> shape_var_map,
- ffi::Map<Var, Expr> var_map);
+TVM_DLL Type EraseToWellDefined(const Type& info, ffi::Map<tirx::Var, PrimExpr> shape_var_map,
+ ffi::Map<Var, Expr> var_map);
/*!
* \brief EraseToWellDefined map overload using a caller-provided analyzer.
- * \param info The struct info.
+ * \param info The type.
* \param shape_var_map map to specify
* whether a symbolic shape var is defined and the value it maps to,
* return nullopt if var is undefined.
@@ -209,11 +206,10 @@ TVM_DLL StructInfo EraseToWellDefined(const StructInfo& info,
* whether a var is defined in the target scope and the value it maps to,
* return nullopt if var is undefined.
* \param ana Context analyzer to prove symbolic expression equality.
- * \return the corresponding erased struct info.
+ * \return the corresponding erased type.
*/
-TVM_DLL StructInfo EraseToWellDefined(const StructInfo& info,
- ffi::Map<tirx::Var, PrimExpr> shape_var_map,
- ffi::Map<Var, Expr> var_map, const arith::Analyzer& ana);
+TVM_DLL Type EraseToWellDefined(const Type& info, ffi::Map<tirx::Var, PrimExpr> shape_var_map,
+ ffi::Map<Var, Expr> var_map, const arith::Analyzer& ana);
/*!
* \brief Fine grained result of base check.
@@ -221,10 +217,10 @@ TVM_DLL StructInfo EraseToWellDefined(const StructInfo& info,
* This analysis comes with different levels of checking failures
* that can help to customize the compilation decisions.
*
- * For a given pair of lhs_struct_info, rhs_struct_info. We adopt
+ * For a given pair of lhs_ty, rhs_ty. We adopt
* the following terminology:
- * - LSet = {value | value matches lhs_struct_info}
- * - RSet = {value | value matches rhs_struct_info}
+ * - LSet = {value | value matches lhs_ty}
+ * - RSet = {value | value matches rhs_ty}
*
* See the definition of each level below.
*/
@@ -242,13 +238,13 @@ enum class BaseCheckResult {
/*!
* \brief WLSet is not superset of RSet because of mismatch in value information.
*
- * L1-level mismatches in params of FuncStructInfo is categorized as
- * If lhs is FuncStructInfo, then L1-level mismatch in its params
+ * L1-level mismatches in params of FuncType is categorized as
+ * If lhs is FuncType, then L1-level mismatch in its params
* is categorized as L2-level mismatch for lhs.
*
* Design considerations for functions:
* - (a) We want to be able to erase type/value in function signature
- * when we unify function struct info and preserve simpler representations.
+ * when we unify function type and preserve simpler representations.
* - (b) We automatically insert match_cast at function boundary, so
* we can erase (int)->int argument as (object)->int.
* The input shape/type mismatch will be detected by runtime checks at function boundary.
@@ -267,47 +263,46 @@ enum class BaseCheckResult {
*
* This function returns fine-grained base-check result on reasons of failure.
*
- * \param base The base struct info.
- * \param derived The derived struct info.
+ * \param base The base type.
+ * \param derived The derived type.
* \return Whether the relation holds.
*
* \sa BaseCheckResult
*/
-TVM_DLL BaseCheckResult StructInfoBaseCheck(const StructInfo& base, const StructInfo& derived);
+TVM_DLL BaseCheckResult TypeBaseCheck(const Type& base, const Type& derived);
/*!
* \brief Run a base check using a caller-provided analyzer.
- * \param base The base struct info.
- * \param derived The derived struct info.
+ * \param base The base type.
+ * \param derived The derived type.
* \param ana Context analyzer to prove symbolic expression equality.
* \return Whether the relation holds.
*
* \sa BaseCheckResult
*/
-TVM_DLL BaseCheckResult StructInfoBaseCheck(const StructInfo& base, const StructInfo& derived,
- const arith::Analyzer& ana);
+TVM_DLL BaseCheckResult TypeBaseCheck(const Type& base, const Type& derived,
+ const arith::Analyzer& ana);
/*!
- * \brief Check the relation of two struct info to see if one subsumes another one.
+ * \brief Check the relation of two type to see if one subsumes another one.
*
- * \param base The base struct info.
- * \param derived The derived struct info.
+ * \param base The base type.
+ * \param derived The derived type.
* \return Whether the relation holds.
*/
-TVM_DLL bool IsBaseOf(const StructInfo& base, const StructInfo& derived);
+TVM_DLL bool IsBaseOf(const Type& base, const Type& derived);
/*!
- * \brief Check whether one struct info subsumes another using a caller-provided analyzer.
- * \param base The base struct info.
- * \param derived The derived struct info.
+ * \brief Check whether one type subsumes another using a caller-provided analyzer.
+ * \param base The base type.
+ * \param derived The derived type.
* \param ana Context analyzer to prove symbolic expression equality.
* \return Whether the relation holds.
*/
-TVM_DLL bool IsBaseOf(const StructInfo& base, const StructInfo& derived,
- const arith::Analyzer& ana);
+TVM_DLL bool IsBaseOf(const Type& base, const Type& derived, const arith::Analyzer& ana);
/*!
* \brief Return the condition for which base is a superset of derived
*
- * This function returns finer-grained conditions for kFailL2 than StructInfoBaseCheck
+ * This function returns finer-grained conditions for kFailL2 than TypeBaseCheck
*
* If the returned expression is true, or simplifies to true, then
* base is a superset of derived. If the returned expression is
@@ -318,53 +313,52 @@ TVM_DLL bool IsBaseOf(const StructInfo& base, const StructInfo& derived,
* expression in terms of the symbolic variables available in `base`
* and `derived`.
*
- * \param base The base struct info.
- * \param derived The derived struct info.
+ * \param base The base type.
+ * \param derived The derived type.
* \return Whether base is a base of derived.
*
* \sa BaseCheckResult
*/
-TVM_DLL PrimExpr StructInfoBaseCheckPrecondition(const StructInfo& base, const StructInfo& derived);
+TVM_DLL PrimExpr TypeBaseCheckPrecondition(const Type& base, const Type& derived);
/*!
- * \brief Unify the two struct info to their least common ancestor.
+ * \brief Unify the two type to their least common ancestor.
*
* \param lhs The left operand.
* \param rhs The right operand.
* \return The unified information.
*/
-TVM_DLL StructInfo StructInfoLCA(const StructInfo& lhs, const StructInfo& rhs);
+TVM_DLL Type TypeLCA(const Type& lhs, const Type& rhs);
/*!
- * \brief Unify two struct infos using a caller-provided analyzer.
+ * \brief Unify two types using a caller-provided analyzer.
* \param lhs The left operand.
* \param rhs The right operand.
* \param ana Context analyzer to prove symbolic expression equality.
* \return The unified information.
*/
-TVM_DLL StructInfo StructInfoLCA(const StructInfo& lhs, const StructInfo& rhs,
- const arith::Analyzer& ana);
+TVM_DLL Type TypeLCA(const Type& lhs, const Type& rhs, const arith::Analyzer& ana);
/*!
- * \brief Get the TIR variables that appear in the input struct info.
+ * \brief Get the TIR variables that appear in the input type.
* The returned list is deduplicated - each TIR variable will appear at most once.
- * \param sinfo The struct info object to be analyzed.
- * \return The list of TIR variables that appear in the input struct info.
+ * \param ty The type object to be analyzed.
+ * \return The list of TIR variables that appear in the input type.
*/
-TVM_DLL ffi::Array<tirx::Var> TIRVarsInStructInfo(const StructInfo& sinfo);
+TVM_DLL ffi::Array<tirx::Var> TIRVarsInType(const Type& ty);
/*!
- * \brief Get the TIR variables that appear in the input struct info.
+ * \brief Get the TIR variables that appear in the input type.
*
* Returns all symbolic variables that are definable based on, and
- * used within, the StructInfo.
+ * used within, the Type.
*
- * \param sinfo The struct info object to be analyzed.
+ * \param ty The type object to be analyzed.
*
* \return A tuple of (definable,used) TIR variables. Both lists are
* deduplicated, each TIR variable will appear at most once, and in
* order of occurrence.
*/
-TVM_DLL ffi::Array<tirx::Var> DefinableTIRVarsInStructInfo(const StructInfo& sinfo);
+TVM_DLL ffi::Array<tirx::Var> DefinableTIRVarsInType(const Type& ty);
/*! \brief Collect expressions whose usage requires them to be non-negative
*
@@ -373,11 +367,11 @@ TVM_DLL ffi::Array<tirx::Var> DefinableTIRVarsInStructInfo(const StructInfo& sin
* to generate assertions prior to calling a kernel, or to provide
* assumptions within a kernel that may be useful for simplification.
*
- * \param sinfo The struct info to be analyzed
+ * \param ty The type to be analyzed
*
* \return A list of non-negative expressions.
*/
-TVM_DLL ffi::Array<PrimExpr> CollectNonNegativeExpressions(const StructInfo& sinfo);
+TVM_DLL ffi::Array<PrimExpr> CollectNonNegativeExpressions(const Type& ty);
/*!
* \brief Get the TIR variables that defined in the input function.
@@ -609,7 +603,7 @@ TVM_DLL bool HasReshapePattern(const tirx::PrimFunc& func);
* can be ignored in the check (must be a Var or GlobalVar).
* \return The impure expression, if one exists within the given
* expression. Otherwise, std::nullopt.
- * \note Relies on StructInfo annotations, so ensure that the module has been normalized first.
+ * \note Relies on Type annotations, so ensure that the module has been normalized first.
* Also, an impure call in a *nested* function does *not* mean that the outer expression contains
* an impure call--it only does if the nested function is *later called*.
*/
@@ -623,7 +617,7 @@ TVM_DLL ffi::Optional<Expr> FindImpureCall(
* the caller can pass the function's name so recursive calls
* can be ignored in the check (must be a Var or GlobalVar).
* \return A boolean indicating if the expression contains any impure calls.
- * \note Relies on StructInfo annotations, so ensure that the module has been normalized first.
+ * \note Relies on Type annotations, so ensure that the module has been normalized first.
* Also, an impure call in a *nested* function does *not* mean that the outer expression contains
* an impure call--it only does if the nested function is *later called*.
*/
@@ -638,12 +632,12 @@ TVM_DLL bool ContainsImpureCall(
* path. Use \ref CheckWellFormed for a boolean answer.
*
* \param obj The IRModule or relax::Function to check.
- * \param check_struct_info If true, verify that every Expr has struct_info populated.
- * \note By default the structure info is always checked. It is only in test cases
- * where `check_struct_info` might be false, so that other well-formed requirements
- * will be well tested and will not be blocked by not having structure info.
+ * \param check_ty If true, verify that every Expr has ty populated.
+ * \note By default the type information is always checked. It is only in test cases
+ * where `check_ty` might be false, so that other well-formed requirements
+ * will be well tested and will not be blocked by not having type information.
*/
-TVM_DLL void WellFormed(ffi::Variant<IRModule, Function> obj, bool check_struct_info = true);
+TVM_DLL void WellFormed(ffi::Variant<IRModule, Function> obj, bool check_ty = true);
/*!
* \brief Return whether an IRModule or Function is well-formed.
@@ -652,10 +646,10 @@ TVM_DLL void WellFormed(ffi::Variant<IRModule, Function> obj, bool check_struct_
* violation.
*
* \param obj The IRModule or relax::Function to check.
- * \param check_struct_info If true, verify that every Expr has struct_info populated.
+ * \param check_ty If true, verify that every Expr has ty populated.
* \return true if the object is well-formed, false otherwise.
*/
-TVM_DLL bool CheckWellFormed(ffi::Variant<IRModule, Function> obj, bool check_struct_info = true);
+TVM_DLL bool CheckWellFormed(ffi::Variant<IRModule, Function> obj, bool check_ty = true);
/*!
* \brief Using the layout transforms on the outputs, suggest layout transformation on the blocks
diff --git a/include/tvm/relax/attrs/distributed.h b/include/tvm/relax/attrs/distributed.h
index 23b698eb36..78e63b3306 100644
--- a/include/tvm/relax/attrs/distributed.h
+++ b/include/tvm/relax/attrs/distributed.h
@@ -25,7 +25,7 @@
#define TVM_RELAX_ATTRS_DISTRIBUTED_H_
#include <tvm/relax/distributed/global_info.h>
-#include <tvm/relax/distributed/struct_info.h>
+#include <tvm/relax/distributed/type.h>
#include <tvm/relax/expr.h>
namespace tvm {
diff --git a/include/tvm/relax/binding_rewrite.h b/include/tvm/relax/binding_rewrite.h
index 740e8ed01f..6adefa9a43 100644
--- a/include/tvm/relax/binding_rewrite.h
+++ b/include/tvm/relax/binding_rewrite.h
@@ -47,8 +47,8 @@ class DataflowBlockRewriteNode : public ffi::Object {
void Add(Binding binding);
/*! \brief Insert an expression as VarBinding with variable name. */
void Add(ffi::String var_name, Expr expr, bool is_dfvar = false) {
- auto var = is_dfvar ? DataflowVar(var_name, GetStructInfo(expr)) //
- : Var(var_name, GetStructInfo(expr));
+ auto var = is_dfvar ? DataflowVar(var_name, GetType(expr)) //
+ : Var(var_name, GetType(expr));
Add(VarBinding(std::move(var), std::move(expr)));
}
/*! \brief Insert an expression as VarBinding with automatic variable name. */
diff --git a/include/tvm/relax/block_builder.h b/include/tvm/relax/block_builder.h
index 8413686dc9..ac20e2ad8f 100644
--- a/include/tvm/relax/block_builder.h
+++ b/include/tvm/relax/block_builder.h
@@ -202,11 +202,11 @@ class BlockBuilderNode : public ffi::Object {
/*!
* \brief Emit a MatchCast.
* \param value The input value.
- * \param struct_info The struct info to be matched.
+ * \param ty The type to be matched.
* \param name_hint Name hint for the bound variable.
* \return The variable bound to the MatchCast.
*/
- virtual Var EmitMatchCast(Expr value, StructInfo struct_info, ffi::String name_hint = "") = 0;
+ virtual Var EmitMatchCast(Expr value, Type ty, ffi::String name_hint = "") = 0;
/*!
* \brief Generate an output for the current dataflow block.
@@ -230,7 +230,7 @@ class BlockBuilderNode : public ffi::Object {
* \param expr The input expression.
* \return The normalized expression.
*
- * \note Invariant: If any of the sub expr have struct_info field.
+ * \note Invariant: If any of the sub expr have ty field.
* they must have already been normalized.
*/
virtual Expr Normalize(const Expr& expr) = 0;
diff --git a/include/tvm/relax/dataflow_pattern.h b/include/tvm/relax/dataflow_pattern.h
index 58d46f0438..27894da3ad 100644
--- a/include/tvm/relax/dataflow_pattern.h
+++ b/include/tvm/relax/dataflow_pattern.h
@@ -56,7 +56,7 @@ class OrPattern;
class AndPattern;
class NotPattern;
class ShapePattern;
-class StructInfoPattern;
+class TypePattern;
class DataTypePattern;
class AttrPattern;
class SameShapeConstraint;
@@ -114,8 +114,8 @@ class DFPattern : public ffi::ObjectRef {
TVM_DLL NotPattern operator~() const;
/*! \brief Syntatic Sugar for creating an AttrPattern */
TVM_DLL AttrPattern HasAttr(const ffi::Map<ffi::String, Any>& attrs) const;
- /*! \brief Syntatic Sugar for creating a StructInfoPattern */
- TVM_DLL StructInfoPattern HasStructInfo(const StructInfo& struct_info) const;
+ /*! \brief Syntatic Sugar for creating a TypePattern */
+ TVM_DLL TypePattern HasType(const Type& ty) const;
/*! \brief Syntatic Sugar for creating a DataTypePattern with a DataType */
TVM_DLL DataTypePattern HasDtype(const DataType& dtype) const;
/*! \brief Syntatic Sugar for creating a DataTypePattern with a data type's name */
@@ -484,7 +484,7 @@ class CallPatternNode : public DFPatternNode {
*/
bool varg_default_wildcard; /*!< N(args) can be < N(real args) by the padding of Wildcard */
- // Todo(relax-team): Dataflow pattern for StructInfo, and match sinfo_args
+ // Todo(relax-team): Dataflow pattern for Type, and match ty_args
static void RegisterReflection() {
namespace refl = tvm::ffi::reflection;
@@ -770,28 +770,27 @@ class WildcardPattern : public DFPattern {
};
/*!
- * \brief Pattern for matching a certain struct info.
- * \sa StructInfoPattern
+ * \brief Pattern for matching a certain type.
+ * \sa TypePattern
*/
-class StructInfoPatternNode : public DFPatternNode {
+class TypePatternNode : public DFPatternNode {
public:
- DFPattern pattern; /*!< The pattern to match */
- StructInfo struct_info; /*!< The type to match */
+ DFPattern pattern; /*!< The pattern to match */
+ Type ty; /*!< The type to match */
static void RegisterReflection() {
namespace refl = tvm::ffi::reflection;
- refl::ObjectDef<StructInfoPatternNode>()
- .def_ro("pattern", &StructInfoPatternNode::pattern)
- .def_ro("struct_info", &StructInfoPatternNode::struct_info);
+ refl::ObjectDef<TypePatternNode>()
+ .def_ro("pattern", &TypePatternNode::pattern)
+ .def_ro("ty", &TypePatternNode::ty);
}
- TVM_FFI_DECLARE_OBJECT_INFO_FINAL("relax.dpl.StructInfoPattern", StructInfoPatternNode,
- DFPatternNode);
+ TVM_FFI_DECLARE_OBJECT_INFO_FINAL("relax.dpl.TypePattern", TypePatternNode, DFPatternNode);
};
-class StructInfoPattern : public DFPattern {
+class TypePattern : public DFPattern {
public:
- TVM_DLL StructInfoPattern(DFPattern pattern, StructInfo struct_info);
- TVM_FFI_DEFINE_OBJECT_REF_METHODS_NULLABLE(StructInfoPattern, DFPattern, StructInfoPatternNode);
+ TVM_DLL TypePattern(DFPattern pattern, Type ty);
+ TVM_FFI_DEFINE_OBJECT_REF_METHODS_NULLABLE(TypePattern, DFPattern, TypePatternNode);
};
/*!
@@ -953,7 +952,7 @@ ExprPattern IsExpr(const Expr& expr);
/*! \brief Syntatic Sugar for creating a ExprPattern base on an Op */
ExprPattern IsOp(const ffi::String& op_name);
/*! \brief Syntatic Sugar for call_tir (return a tensor) */
-// Todo(relax-team): Dataflow pattern for StructInfo, and match out_sinfo
+// Todo(relax-team): Dataflow pattern for Type, and match out_ty
CallPattern IsCallTIR(const ffi::String& name, ffi::Optional<TuplePattern> args = std::nullopt);
/*! \brief Syntatic Sugar for call_tir (return a tuple of tensor) */
CallPattern IsCallTIR(const ffi::String& name, TuplePattern var_args);
diff --git a/include/tvm/relax/dataflow_pattern_functor.h b/include/tvm/relax/dataflow_pattern_functor.h
index 6816dbb92c..dfe130a09c 100644
--- a/include/tvm/relax/dataflow_pattern_functor.h
+++ b/include/tvm/relax/dataflow_pattern_functor.h
@@ -96,8 +96,7 @@ class DFPatternFunctor<R(const DFPattern& n, Args...)> {
virtual R VisitDFPattern_(const TupleGetItemPatternNode* op,
Args... args) DFPATTERN_FUNCTOR_DEFAULT;
virtual R VisitDFPattern_(const TuplePatternNode* op, Args... args) DFPATTERN_FUNCTOR_DEFAULT;
- virtual R VisitDFPattern_(const StructInfoPatternNode* op,
- Args... args) DFPATTERN_FUNCTOR_DEFAULT;
+ virtual R VisitDFPattern_(const TypePatternNode* op, Args... args) DFPATTERN_FUNCTOR_DEFAULT;
virtual R VisitDFPattern_(const WildcardPatternNode* op, Args... args) DFPATTERN_FUNCTOR_DEFAULT;
virtual R VisitDFPattern_(const VarPatternNode* op, Args... args) DFPATTERN_FUNCTOR_DEFAULT;
@@ -132,7 +131,7 @@ class DFPatternFunctor<R(const DFPattern& n, Args...)> {
RELAX_DFPATTERN_FUNCTOR_DISPATCH(ShapePatternNode);
RELAX_DFPATTERN_FUNCTOR_DISPATCH(TupleGetItemPatternNode);
RELAX_DFPATTERN_FUNCTOR_DISPATCH(TuplePatternNode);
- RELAX_DFPATTERN_FUNCTOR_DISPATCH(StructInfoPatternNode);
+ RELAX_DFPATTERN_FUNCTOR_DISPATCH(TypePatternNode);
RELAX_DFPATTERN_FUNCTOR_DISPATCH(WildcardPatternNode);
RELAX_DFPATTERN_FUNCTOR_DISPATCH(VarPatternNode);
RELAX_DFPATTERN_FUNCTOR_DISPATCH(DataflowVarPatternNode);
@@ -166,7 +165,7 @@ class DFPatternVisitor : public DFPatternFunctor<void(const DFPattern&)> {
void VisitDFPattern_(const ShapePatternNode* op) override;
void VisitDFPattern_(const TupleGetItemPatternNode* op) override;
void VisitDFPattern_(const TuplePatternNode* op) override;
- void VisitDFPattern_(const StructInfoPatternNode* op) override;
+ void VisitDFPattern_(const TypePatternNode* op) override;
void VisitDFPattern_(const WildcardPatternNode* op) override;
void VisitDFPattern_(const VarPatternNode* op) override;
diff --git a/include/tvm/relax/distributed/axis_group_graph.h b/include/tvm/relax/distributed/axis_group_graph.h
index 86b34b7135..0e15b20164 100644
--- a/include/tvm/relax/distributed/axis_group_graph.h
+++ b/include/tvm/relax/distributed/axis_group_graph.h
@@ -21,7 +21,7 @@
#define TVM_RELAX_DISTRIBUTED_AXIS_GROUP_GRAPH_H_
#include <tvm/arith/iter_affine_map.h>
-#include <tvm/relax/distributed/struct_info.h>
+#include <tvm/relax/distributed/type.h>
#include <tvm/relax/expr.h>
#include <tvm/tirx/function.h>
#include <tvm/tirx/stmt_functor.h>
diff --git a/include/tvm/relax/distributed/struct_info.h b/include/tvm/relax/distributed/type.h
similarity index 74%
rename from include/tvm/relax/distributed/struct_info.h
rename to include/tvm/relax/distributed/type.h
index 81fdf0fb3f..5ffbff890e 100644
--- a/include/tvm/relax/distributed/struct_info.h
+++ b/include/tvm/relax/distributed/type.h
@@ -18,15 +18,18 @@
*/
/*!
- * \file tvm/relax/distributed/struct_info.h
- * \brief Struct info for DTensor (Distributed Tensor)
+ * \file tvm/relax/distributed/type.h
+ * \brief Type definitions for DTensor (Distributed Tensor)
*/
-#ifndef TVM_RELAX_DISTRIBUTED_STRUCT_INFO_H_
-#define TVM_RELAX_DISTRIBUTED_STRUCT_INFO_H_
+#ifndef TVM_RELAX_DISTRIBUTED_TYPE_H_
+#define TVM_RELAX_DISTRIBUTED_TYPE_H_
#include <tvm/relax/distributed/global_info.h>
-#include <tvm/relax/struct_info.h>
+#include <tvm/relax/type.h>
+
+#include <utility>
+
namespace tvm {
namespace relax {
namespace distributed {
@@ -111,14 +114,22 @@ class Placement : public ffi::ObjectRef {
};
/*!
- * \brief StructInfo of DTensor (Distributed Tensor).
+ * \brief Type of DTensor (Distributed Tensor).
*/
-class DTensorStructInfoNode : public StructInfoNode {
+class DTensorTypeNode : public TypeNode {
public:
+ explicit DTensorTypeNode(ffi::UnsafeInit)
+ : tensor_ty(ffi::UnsafeInit{}), device_mesh(), placement() {}
+
+ DTensorTypeNode(TensorType tensor_ty, DeviceMesh device_mesh, Placement placement)
+ : tensor_ty(std::move(tensor_ty)),
+ device_mesh(std::move(device_mesh)),
+ placement(std::move(placement)) {}
+
/*!
- * \brief The struct info inherited from TensorStructInfo
+ * \brief The tensor type carried by the DTensor type.
*/
- TensorStructInfo tensor_sinfo;
+ TensorType tensor_ty;
/*!
* \brief The device mesh of the tensor.
*/
@@ -130,36 +141,35 @@ class DTensorStructInfoNode : public StructInfoNode {
static void RegisterReflection() {
namespace refl = tvm::ffi::reflection;
- refl::ObjectDef<DTensorStructInfoNode>()
- .def_ro("device_mesh", &DTensorStructInfoNode::device_mesh)
- .def_ro("placement", &DTensorStructInfoNode::placement)
- .def_ro("tensor_sinfo", &DTensorStructInfoNode::tensor_sinfo);
+ refl::ObjectDef<DTensorTypeNode>()
+ .def_ro("device_mesh", &DTensorTypeNode::device_mesh)
+ .def_ro("placement", &DTensorTypeNode::placement)
+ .def_ro("tensor_ty", &DTensorTypeNode::tensor_ty);
}
- TVM_FFI_DECLARE_OBJECT_INFO_FINAL("relax.DTensorStructInfo", DTensorStructInfoNode,
- StructInfoNode);
+ TVM_FFI_DECLARE_OBJECT_INFO_FINAL("relax.DTensorType", DTensorTypeNode, TypeNode);
};
/*!
- * \brief Managed reference to DTensorStructInfoNode.
- * \sa DTensorStructInfoNode
+ * \brief Managed reference to DTensorTypeNode.
+ * \sa DTensorTypeNode
*/
-class DTensorStructInfo : public StructInfo {
+class DTensorType : public Type {
public:
/*!
* \brief Construction with device mesh and placement.
- * \param tensor_sinfo The struct info inherited from TensorStructInfo
+ * \param tensor_ty The tensor type carried by the DTensor type.
* \param device_mesh The device mesh of the tensor.
* \param placement The placement of the tensor among the device mesh.
* \param span The span of the AST.
*/
- TVM_DLL DTensorStructInfo(TensorStructInfo tensor_sinfo, DeviceMesh device_mesh,
- Placement placement, Span span = Span());
+ TVM_DLL DTensorType(TensorType tensor_ty, DeviceMesh device_mesh, Placement placement,
+ Span span = Span());
- TVM_FFI_DEFINE_OBJECT_REF_METHODS_NULLABLE(DTensorStructInfo, StructInfo, DTensorStructInfoNode);
+ TVM_FFI_DEFINE_OBJECT_REF_METHODS_NOTNULLABLE(DTensorType, Type, DTensorTypeNode);
};
} // namespace distributed
} // namespace relax
} // namespace tvm
-#endif // TVM_RELAX_DISTRIBUTED_STRUCT_INFO_H_
+#endif // TVM_RELAX_DISTRIBUTED_TYPE_H_
diff --git a/include/tvm/relax/expr.h b/include/tvm/relax/expr.h
index 6da9cb1692..937091255b 100644
--- a/include/tvm/relax/expr.h
+++ b/include/tvm/relax/expr.h
@@ -36,8 +36,6 @@
namespace tvm {
namespace relax {
-using Expr = RelaxExpr;
-using ExprNode = RelaxExprNode;
/*!
* \brief The unique identifier of variables.
*
@@ -76,64 +74,6 @@ class Id : public ffi::ObjectRef {
TVM_FFI_DEFINE_OBJECT_REF_METHODS_NULLABLE(Id, ffi::ObjectRef, IdNode);
};
-/*!
- * \brief Base type of all structure information.
- *
- * StructInfo stores possible structure information
- * deduced during compile-time. It encapsulates
- * both static type and runtime information such
- * as shape.
- *
- * StructInfo of each non-primitive Expr can be
- * deduced during compilation in a "best-effort" manner.
- *
- * When struct_info appears in function parameter and return
- * signatures. They will imply a runtime check that matches
- * the structure information with the value.
- *
- * When it appears in Expr, they follow "assume-semantics",
- * which means the compiler will take the deduced information as it is
- * and only do best effort prove and checks.
- *
- * Each struct info can be uniquely erased to a static-type.
- * The compiler will still compile the code(with less information)
- * when we erase to the static type.
- *
- * If an StructInfo contains an Expr field, then that field
- * must be normalized already through NormalizeArg.
- * This invariant will be checked in constructors
- * and help us to simplify our assumption
- * during struct info deduction.
- */
-class StructInfoNode : public ffi::Object {
- public:
- /*!
- * \brief Span that points to the original source code.
- * Reserved debug information.
- */
- mutable Span span;
-
- static void RegisterReflection() {
- namespace refl = tvm::ffi::reflection;
- refl::ObjectDef<StructInfoNode>().def_ro("span", &StructInfoNode::span,
- refl::AttachFieldFlag::SEqHashIgnore());
- }
-
- static constexpr TVMFFISEqHashKind _type_s_eq_hash_kind = kTVMFFISEqHashKindTreeNode;
-
- static constexpr const uint32_t _type_child_slots = 7;
- TVM_FFI_DECLARE_OBJECT_INFO("ir.StructInfo", StructInfoNode, ffi::Object);
-};
-
-/*!
- * \brief Managed reference to StructInfoNode.
- * \sa StructInfoNode
- */
-class StructInfo : public ffi::ObjectRef {
- public:
- TVM_FFI_DEFINE_OBJECT_REF_METHODS_NULLABLE(StructInfo, ffi::ObjectRef, StructInfoNode);
-};
-
/*!
* \brief Call corresponds to callable invocation.
* Corresponds to operation in computational graph terminology.
@@ -155,16 +95,16 @@ class CallNode : public ExprNode {
Attrs attrs;
/*!
- * \brief The structure info arguments of a CallNode.
- * sinfo_args is by default designed to be non-empty only for intrinsic op (e.g.,
+ * \brief The type information arguments of a CallNode.
+ * ty_args is by default designed to be non-empty only for intrinsic op (e.g.,
* call_tir, call_builtin_with_ctx, etc.) and calls to ExternFuncs, with the main
- * usage of structure info inference.
+ * usage of type information inference.
*
- * Regular ops also at times may have sinfo_args defined to specialize partial
- * or complete structure info. Like VDevice customization with mixed input memory_scopes.
+ * Regular ops also at times may have ty_args defined to specialize partial
+ * or complete type information. Like VDevice customization with mixed input memory_scopes.
* The customized pass can set this info and operator specific inference will respect it.
*/
- ffi::Array<StructInfo> sinfo_args;
+ ffi::Array<Type> ty_args;
static void RegisterReflection() {
namespace refl = tvm::ffi::reflection;
@@ -172,7 +112,7 @@ class CallNode : public ExprNode {
.def_ro("op", &CallNode::op)
.def_ro("args", &CallNode::args)
.def_ro("attrs", &CallNode::attrs)
- .def_ro("sinfo_args", &CallNode::sinfo_args);
+ .def_ro("ty_args", &CallNode::ty_args);
}
TVM_FFI_DECLARE_OBJECT_INFO_FINAL("relax.expr.Call", CallNode, ExprNode);
};
@@ -184,11 +124,11 @@ class Call : public Expr {
* \param op The operator to be invoked.
* \param args The arguments of the call.
* \param attrs The attributes of the call node.
- * \param sinfo_args The structure info arguments passed to a function.
+ * \param ty_args The type information arguments passed to a function.
* \param span The source span of the expression.
*/
TVM_DLL Call(Expr op, ffi::Array<Expr> args, Attrs attrs = Attrs(),
- ffi::Array<StructInfo> sinfo_args = ffi::Array<StructInfo>(), Span span = Span());
+ ffi::Array<Type> ty_args = ffi::Array<Type>(), Span span = Span());
TVM_FFI_DEFINE_OBJECT_REF_METHODS_NULLABLE(Call, Expr, CallNode);
TVM_DEFINE_OBJECT_REF_COW_METHOD(CallNode);
@@ -199,12 +139,11 @@ class Call : public Expr {
* Returns \p call if all properties are unchanged. Otherwise, returns a copy with the new
* fields.
*/
-Call WithFields(
- Call call, ffi::Optional<Expr> opt_op = ffi::Optional<Expr>(),
- ffi::Optional<ffi::Array<Expr>> opt_args = ffi::Optional<ffi::Array<Expr>>(),
- ffi::Optional<Attrs> opt_attrs = ffi::Optional<Attrs>(),
- ffi::Optional<ffi::Array<StructInfo>> opt_sinfo_args = ffi::Optional<ffi::Array<StructInfo>>(),
- ffi::Optional<Span> opt_span = ffi::Optional<Span>());
+Call WithFields(Call call, ffi::Optional<Expr> opt_op = ffi::Optional<Expr>(),
+ ffi::Optional<ffi::Array<Expr>> opt_args = ffi::Optional<ffi::Array<Expr>>(),
+ ffi::Optional<Attrs> opt_attrs = ffi::Optional<Attrs>(),
+ ffi::Optional<ffi::Array<Type>> opt_ty_args = ffi::Optional<ffi::Array<Type>>(),
+ ffi::Optional<Span> opt_span = ffi::Optional<Span>());
/*! \brief Tuple container */
class TupleNode : public ExprNode {
@@ -353,7 +292,7 @@ class VarNode : public LeafExprNode {
static void RegisterReflection() {
namespace refl = tvm::ffi::reflection;
refl::ObjectDef<VarNode>().def_ro("vid", &VarNode::vid);
- // customize structural equal and hash to include struct_info_
+ // customize structural equal and hash to include ty
refl::TypeAttrDef<VarNode>()
.def("__s_equal__", &VarNode::SEqual)
.def("__s_hash__", &VarNode::SHash);
@@ -361,14 +300,13 @@ class VarNode : public LeafExprNode {
bool SEqual(const VarNode* other,
ffi::TypedFunction<bool(AnyView, AnyView, bool, AnyView)> equal) const {
- return equal(vid, other->vid, false, "vid") &&
- equal(struct_info_, other->struct_info_, false, "struct_info_");
+ return equal(vid, other->vid, false, "vid") && equal(ty, other->ty, false, "ty");
}
int64_t SHash(int64_t init_hash, ffi::TypedFunction<int64_t(AnyView, int64_t, bool)> hash) const {
int64_t hash_value = init_hash;
hash_value = hash(vid, hash_value, false);
- hash_value = hash(struct_info_, hash_value, false);
+ hash_value = hash(ty, hash_value, false);
return hash_value;
}
@@ -379,12 +317,10 @@ class VarNode : public LeafExprNode {
class Var : public LeafExpr {
public:
- TVM_DLL explicit Var(ffi::String name_hint, ffi::Optional<StructInfo> struct_info_annotation,
- Span span = Span())
- : Var(Id(name_hint), struct_info_annotation, span) {}
+ TVM_DLL explicit Var(ffi::String name_hint, ffi::Optional<Type> ty_annotation, Span span = Span())
+ : Var(Id(name_hint), ty_annotation, span) {}
- TVM_DLL explicit Var(Id vid, ffi::Optional<StructInfo> struct_info_annotation,
- Span span = Span());
+ TVM_DLL explicit Var(Id vid, ffi::Optional<Type> ty_annotation, Span span = Span());
TVM_FFI_DEFINE_OBJECT_REF_METHODS_NULLABLE(Var, LeafExpr, VarNode);
VarNode* CopyOnWrite();
@@ -406,12 +342,11 @@ class DataflowVarNode : public VarNode {
class DataflowVar : public Var {
public:
- TVM_DLL explicit DataflowVar(ffi::String name_hint,
- ffi::Optional<StructInfo> struct_info_annotation, Span span = Span())
- : DataflowVar(Id(name_hint), struct_info_annotation, span) {}
+ TVM_DLL explicit DataflowVar(ffi::String name_hint, ffi::Optional<Type> ty_annotation,
+ Span span = Span())
+ : DataflowVar(Id(name_hint), ty_annotation, span) {}
- TVM_DLL explicit DataflowVar(Id vid, ffi::Optional<StructInfo> struct_info_annotation,
- Span span = Span());
+ TVM_DLL explicit DataflowVar(Id vid, ffi::Optional<Type> ty_annotation, Span span = Span());
TVM_FFI_DEFINE_OBJECT_REF_METHODS_NULLABLE(DataflowVar, Var, DataflowVarNode);
TVM_DEFINE_OBJECT_REF_COW_METHOD(DataflowVarNode);
@@ -445,12 +380,11 @@ class Constant : public LeafExpr {
/*!
* \brief The constructor
* \param data The data of the constant tensor.
- * \param struct_info_annotation The struct info of the constant tensor.
+ * \param ty_annotation The type of the constant tensor.
* If not specified, infer it from data.
* \param span The source span of the expression.
*/
- TVM_DLL explicit Constant(runtime::Tensor data,
- ffi::Optional<StructInfo> struct_info_annotation = std::nullopt,
+ TVM_DLL explicit Constant(runtime::Tensor data, ffi::Optional<Type> ty_annotation = std::nullopt,
Span span = Span());
TVM_FFI_DEFINE_OBJECT_REF_METHODS_NULLABLE(Constant, LeafExpr, ConstantNode);
@@ -600,26 +534,25 @@ class Binding : public ffi::ObjectRef {
};
/*!
- * \brief Runtime-match the value to the struct info.
+ * \brief Runtime-match the value to the type.
*
* This operation does runtime check, populates the un-defined symbolic shape vars
- * and vars in struct_info in first occurance, and insert equality assertions in
+ * and vars in ty in first occurance, and insert equality assertions in
* other cases.
*/
class MatchCastNode : public BindingNode {
public:
/*! \brief The input value to match cast. */
Expr value;
- /*! \brief The struct info pattern to match to. */
- StructInfo struct_info;
+ /*! \brief The type pattern to match to. */
+ Type ty;
static void RegisterReflection() {
namespace refl = tvm::ffi::reflection;
refl::ObjectDef<MatchCastNode>()
.def_ro("value", &MatchCastNode::value)
// TODO(tqchen): use SEqHashDefNonRecursive after the next pypi tvm-ffi release
- .def_ro("struct_info", &MatchCastNode::struct_info,
- refl::AttachFieldFlag::SEqHashDefRecursive());
+ .def_ro("ty", &MatchCastNode::ty, refl::AttachFieldFlag::SEqHashDefRecursive());
}
TVM_FFI_DECLARE_OBJECT_INFO_FINAL("relax.expr.MatchCast", MatchCastNode, BindingNode);
};
@@ -630,7 +563,7 @@ class MatchCastNode : public BindingNode {
*/
class MatchCast : public Binding {
public:
- TVM_DLL explicit MatchCast(Var var, Expr value, StructInfo struct_info, Span span = Span());
+ TVM_DLL explicit MatchCast(Var var, Expr value, Type ty, Span span = Span());
TVM_FFI_DEFINE_OBJECT_REF_METHODS_NULLABLE(MatchCast, Binding, MatchCastNode);
TVM_DEFINE_OBJECT_REF_COW_METHOD(MatchCastNode);
@@ -818,7 +751,7 @@ class FunctionNode : public BaseFuncNode {
/*! \brief The body of the function. */
SeqExpr body;
/*! \brief The return type of the function. */
- StructInfo ret_struct_info;
+ Type ret_ty;
/*! \brief Whether the function is annotated as pure or not. */
bool is_pure;
@@ -827,7 +760,7 @@ class FunctionNode : public BaseFuncNode {
refl::ObjectDef<FunctionNode>()
.def_ro("params", &FunctionNode::params, refl::AttachFieldFlag::SEqHashDefRecursive())
.def_ro("body", &FunctionNode::body)
- .def_ro("ret_struct_info", &FunctionNode::ret_struct_info)
+ .def_ro("ret_ty", &FunctionNode::ret_ty)
.def_ro("is_pure", &FunctionNode::is_pure);
}
@@ -847,8 +780,8 @@ class Function : public BaseFunc {
* Relax IR requirement that all scopes be contained in a
* SeqExpr.
*
- * \param ret_struct_info The StructInfo returned by the function.
- * If std::nullopt, will be inferred from the StructInfo of the
+ * \param ret_ty The Type returned by the function.
+ * If std::nullopt, will be inferred from the Type of the
* function's body.
*
* \param is_pure The purity of the function.
@@ -858,17 +791,15 @@ class Function : public BaseFunc {
*
* \param span The source span of the expression.
*/
- TVM_DLL explicit Function(ffi::Array<Var> params, Expr body,
- ffi::Optional<StructInfo> ret_struct_info, bool is_pure = true,
- DictAttrs attrs = DictAttrs(), Span span = Span());
+ TVM_DLL explicit Function(ffi::Array<Var> params, Expr body, ffi::Optional<Type> ret_ty,
+ bool is_pure = true, DictAttrs attrs = DictAttrs(), Span span = Span());
/*!
* \brief Mimics the constructor but without body Expr.
- * \note ret_struct_info is required, since it can not deduced by the body.
+ * \note ret_ty is required, since it can not deduced by the body.
*/
- TVM_DLL static Function CreateEmpty(ffi::Array<Var> params, StructInfo ret_struct_info,
- bool is_pure = true, DictAttrs attrs = DictAttrs(),
- Span span = Span());
+ TVM_DLL static Function CreateEmpty(ffi::Array<Var> params, Type ret_ty, bool is_pure = true,
+ DictAttrs attrs = DictAttrs(), Span span = Span());
TVM_FFI_DEFINE_OBJECT_REF_METHODS_NULLABLE(Function, BaseFunc, FunctionNode);
TVM_DEFINE_OBJECT_REF_COW_METHOD(FunctionNode);
@@ -921,7 +852,7 @@ class ExternFuncNode : public BaseFuncNode {
class ExternFunc : public BaseFunc {
public:
TVM_DLL ExternFunc(ffi::String global_symbol, Span span = Span());
- TVM_DLL ExternFunc(ffi::String global_symbol, StructInfo struct_info, Span span = Span());
+ TVM_DLL ExternFunc(ffi::String global_symbol, Type ty, Span span = Span());
TVM_FFI_DEFINE_OBJECT_REF_METHODS_NULLABLE(ExternFunc, BaseFunc, ExternFuncNode);
TVM_DEFINE_OBJECT_REF_COW_METHOD(ExternFuncNode);
@@ -933,7 +864,7 @@ class ExternFunc : public BaseFunc {
* \return The corresonding shape.
*
* \note This function requires expr to be normalized.
- * The function will report an error if expr's StructInfo is not TensorStructInfo.
+ * The function will report an error if expr's Type is not TensorType.
* It will try to return symbolic function when possible. If the tensor do not
* have a compile-time symbolic shape, the function will then choose to return
* Call(relax.op.shape_of, [expr]).
diff --git a/include/tvm/relax/expr_functor.h b/include/tvm/relax/expr_functor.h
index 9189d5d1fe..eb00f3bbdd 100644
--- a/include/tvm/relax/expr_functor.h
+++ b/include/tvm/relax/expr_functor.h
@@ -28,8 +28,8 @@
#include <tvm/ir/node_functor.h>
#include <tvm/relax/block_builder.h>
#include <tvm/relax/expr.h>
-#include <tvm/relax/struct_info.h>
-#include <tvm/relax/struct_info_functor.h>
+#include <tvm/relax/type.h>
+#include <tvm/relax/type_functor.h>
#include <tvm/tirx/function.h>
#include <unordered_map>
@@ -256,20 +256,20 @@ class ExprVisitor : public ExprFunctor<void(const Expr&)> {
virtual void VisitVarDef(const Var& var);
/*!
- * \brief Visit struct_info may recursively contain Expr/PrimExpr.
+ * \brief Visit ty may recursively contain Expr/PrimExpr.
*
- * By default, this function recurse into struct info such as
- * TensorStructInfo and ShapeStructInfo and call VisitExpr/VisitPrimExpr
- * accordingly. It does not recurse into FunctionStructInfo as it does
+ * By default, this function recurse into type such as
+ * TensorType and ShapeType and call VisitExpr/VisitPrimExpr
+ * accordingly. It does not recurse into FunctionType as it does
* not contain Expr defined in the current scope.
*
* Pass writers can overload this function to change to other behaviors.
- * For example, if we are not interested in Expr in StructInfo, we can
+ * For example, if we are not interested in Expr in Type, we can
* override this function by a no-op.
*
- * \param struct_info Input struct info field.
+ * \param ty Input type field.
*/
- virtual void VisitExprDepStructInfoField(const StructInfo& struct_info);
+ virtual void VisitExprDepTypeField(const Type& ty);
// specific leaf level visitor functions
virtual void VisitVarDef_(const VarNode* var);
@@ -285,29 +285,29 @@ class ExprVisitor : public ExprFunctor<void(const Expr&)> {
// initialize the vtable.
static VisitBindingVTable InitVisitBindingVTable();
/*!
- * \brief Private internal struct info field visitor.
+ * \brief Private internal type field visitor.
*
- * Support default visiting of struct info field and recursive into
+ * Support default visiting of type field and recursive into
* their Expr fields.
*
* We use component instead of sub-classing so there can be other
- * joint inheritance between ExprVisitor and StructInfoVisitor.
+ * joint inheritance between ExprVisitor and TypeVisitor.
*/
- class DefaultStructInfoFieldVisitor : public StructInfoVisitor {
+ class DefaultTypeFieldVisitor : public TypeVisitor {
public:
- explicit DefaultStructInfoFieldVisitor(ExprVisitor* parent);
+ explicit DefaultTypeFieldVisitor(ExprVisitor* parent);
- // Override defaults in struct info visitor.
- void VisitStructInfoExprField(const Expr& expr) final;
- void VisitStructInfoExprField(const PrimExpr& expr) final;
- void VisitStructInfo_(const FuncStructInfoNode* op) final;
+ // Override defaults in type visitor.
+ void VisitTypeExprField(const Expr& expr) final;
+ void VisitTypeExprField(const PrimExpr& expr) final;
+ void VisitType_(const FuncTypeNode* op) final;
private:
ExprVisitor* parent_;
};
// This visitor is not visible to child classes and only
// used to supported default visiting behavior.
- DefaultStructInfoFieldVisitor default_struct_info_field_visitor_{this};
+ DefaultTypeFieldVisitor default_tyfield_visitor_{this};
};
void PostOrderVisit(const Expr& node, std::function<void(const Expr&)> fvisit);
@@ -315,7 +315,7 @@ void PostOrderVisit(const Expr& node, std::function<void(const Expr&)> fvisit);
/*!
* \brief A mutator works in unnormalized form.
*
- * ExprMutatorBase expects input AST to be in the unnormalized form, i.e., struct_info_
+ * ExprMutatorBase expects input AST to be in the unnormalized form, i.e., ty
* of expressions can be nullptr, and the expressions may nest(and as a result the AST is not in
* ANF).
*/
@@ -355,34 +355,34 @@ class ExprMutatorBase : public ExprFunctor<Expr(const Expr&)> {
virtual PrimExpr VisitPrimExpr(const PrimExpr& expr);
/*!
- * \brief Visit struct_info that may recursively contain Expr/PrimExpr.
+ * \brief Visit ty that may recursively contain Expr/PrimExpr.
*
- * By default, this function recurse into struct info such as
- * TensorStructInfo and ShapeStructInfo and call VisitExpr/VisitPrimExpr
- * accordingly. It does not recurse into FunctionStructInfo as it does
+ * By default, this function recurse into type such as
+ * TensorType and ShapeType and call VisitExpr/VisitPrimExpr
+ * accordingly. It does not recurse into FunctionType as it does
* not contain Expr defined in the current scope.
*
* Pass writers can overload this function to change to other behaviors.
- * For example, if in Expr in StructInfo won't change, we can
+ * For example, if in Expr in Type won't change, we can
* override this function by an identity function.
*
- * \param struct_info Input struct info field.
- * \return The updated struct info.
+ * \param ty Input type field.
+ * \return The updated type.
*/
- virtual StructInfo VisitExprDepStructInfoField(const StructInfo& struct_info);
+ virtual Type VisitExprDepTypeField(const Type& ty);
protected:
/*!
- * \brief Check whether VisitExprDepStructInfoField change struct_info.
- * \return Whether struct info changed.
+ * \brief Check whether VisitExprDepTypeField change ty.
+ * \return Whether type changed.
* \note This function is used by mutator implementations to check if
- * previous Expr update will trigger a change in struct_info.
+ * previous Expr update will trigger a change in ty.
* If change is detected, the implementation can generate a fresh
- * node without struct_info, and trigger normalizer to re-derive.
+ * node without ty, and trigger normalizer to re-derive.
*/
- bool VisitAndCheckStructInfoFieldUnchanged(const ffi::ObjectRef& struct_info) {
- if (const StructInfoNode* sinfo = struct_info.as<StructInfoNode>()) {
- return this->VisitExprDepStructInfoField(ffi::GetRef<StructInfo>(sinfo)).same_as(struct_info);
+ bool VisitAndCheckTypeFieldUnchanged(const ffi::ObjectRef& ty) {
+ if (const TypeNode* ty_node = ty.as<TypeNode>()) {
+ return this->VisitExprDepTypeField(ffi::GetRef<Type>(ty_node)).same_as(ty);
} else {
return true;
}
@@ -390,34 +390,34 @@ class ExprMutatorBase : public ExprFunctor<Expr(const Expr&)> {
private:
/*!
- * \brief Private internal struct info field visitor to support
- * Default visiting of struct info field and recursive into their Expr fields.
+ * \brief Private internal type field visitor to support
+ * Default visiting of type field and recursive into their Expr fields.
*
* We use component instead of sub-classing so there can be other
- * joint inheritance between ExprMutator and StructInfoMutator.
+ * joint inheritance between ExprMutator and TypeMutator.
*/
- class DefaultStructInfoFieldMutator : public StructInfoMutator {
+ class DefaultTypeFieldMutator : public TypeMutator {
public:
- explicit DefaultStructInfoFieldMutator(ExprMutatorBase* parent);
+ explicit DefaultTypeFieldMutator(ExprMutatorBase* parent);
- // Override defaults in struct info visitor.
- Expr VisitStructInfoExprField(const Expr& expr) final;
- PrimExpr VisitStructInfoExprField(const PrimExpr& expr) final;
- StructInfo VisitStructInfo_(const FuncStructInfoNode* op) final;
+ // Override defaults in type visitor.
+ Expr VisitTypeExprField(const Expr& expr) final;
+ PrimExpr VisitTypeExprField(const PrimExpr& expr) final;
+ Type VisitType_(const FuncTypeNode* op) final;
private:
ExprMutatorBase* parent_;
};
// This visitor is not visible to child classes and only
// used to supported default visiting behavior.
- DefaultStructInfoFieldMutator default_struct_info_field_mutator_{this};
+ DefaultTypeFieldMutator default_tyfield_mutator_{this};
};
/*!
* \brief A mutator works in normal form.
*
* ExprMutator expects input AST to be in the normal form, i.e., the expressions are normalized(no
- * nesting and hence the AST is in ANF), and all struct_info_ of expressions are
+ * nesting and hence the AST is in ANF), and all ty of expressions are
* available.
*/
class ExprMutator : public ExprMutatorBase {
@@ -544,13 +544,13 @@ class ExprMutator : public ExprMutatorBase {
}
/*!
- * \brief Create a new var with specified struct_info if the original var's shape or type does
- * not match with the specified ones.
+ * \brief Create a new var with specified type if the original var's shape or type does not
+ * match with the specified ones.
* \param var The var to be updated.
- * \param struct_info The struct info to be updated.
- * \return The var filled with struct_info
+ * \param ty The type to be updated.
+ * \return The var filled with type information.
*/
- Var WithStructInfo(Var var, StructInfo struct_info);
+ Var WithType(Var var, Type ty);
/*! \brief Internal block builder to emit bindings during rewriting. */
BlockBuilder builder_;
diff --git a/include/tvm/relax/nested_msg.h b/include/tvm/relax/nested_msg.h
index 4b11e9d2b0..400f73d1a2 100644
--- a/include/tvm/relax/nested_msg.h
+++ b/include/tvm/relax/nested_msg.h
@@ -31,7 +31,7 @@
#include <tvm/ffi/container/array.h>
#include <tvm/ffi/optional.h>
#include <tvm/relax/expr.h>
-#include <tvm/relax/struct_info.h>
+#include <tvm/relax/type.h>
#include <string>
#include <utility>
@@ -277,50 +277,50 @@ NestedMsg<T> MapToNestedMsg(Expr expr, FType fmapleaf) {
}
/*!
- * \brief Map structinfo with possible nested-sinfo to nested message.
+ * \brief Map structinfo with possible nested-ty to nested message.
*
- * This function will unpack recursive sinfo and run fmapleaf for each leaf,
+ * This function will unpack recursive ty and run fmapleaf for each leaf,
* then recursively combines the results together into a NestedMsg.
*
* The nesting structure will corresponds to the tuple structure.
*
- * \param sinfo The input struct info.
- * \param fmapleaf The mapping function for each leaf with signature `NestedMsg<T> fmap(StructInfo)`
+ * \param ty The input type.
+ * \param fmapleaf The mapping function for each leaf with signature `NestedMsg<T> fmap(Type)`
* \tparam T the content type of nested msg
* \tparam FType The mapping function type
*/
template <typename T, typename FType>
-NestedMsg<T> MapToNestedMsg(StructInfo sinfo, FType fmapleaf) {
- if (auto* tuple = sinfo.as<TupleStructInfoNode>()) {
+NestedMsg<T> MapToNestedMsg(Type ty, FType fmapleaf) {
+ if (auto* tuple = ty.as<TupleTypeNode>()) {
ffi::Array<NestedMsg<T>> res;
res.reserve(tuple->fields.size());
- for (StructInfo x : tuple->fields) {
+ for (Type x : tuple->fields) {
res.push_back(MapToNestedMsg<T, FType>(x, fmapleaf));
}
return res;
} else {
- return fmapleaf(sinfo);
+ return fmapleaf(ty);
}
}
/*!
* \brief Map expr with possible nested-tuple to nested message.
*
- * This function will unpack recursive expr by its struct info and
+ * This function will unpack recursive expr by its type and
* run fmapleaf for each leaf, then recursively combines the results
* together into a NestedMsg.
*
- * The nesting structure will corresponds to the struct info of expr.
+ * The nesting structure will corresponds to the type of expr.
*
- * \param expr The input expression which should have struct info.
+ * \param expr The input expression which should have type.
* \param fmapleaf The mapping function for each leaf with signature `NestedMsg<T> fmapleaf(Expr)`
* \tparam T the content type of nested msg
* \tparam FType The mapping function type
*/
template <typename T, typename FType>
-NestedMsg<T> MapToNestedMsgBySInfo(Expr expr, FType fmapleaf) {
- auto sinfo = GetStructInfo(expr);
- if (auto* tuple = sinfo.as<TupleStructInfoNode>()) {
+NestedMsg<T> MapToNestedMsgByType(Expr expr, FType fmapleaf) {
+ auto ty = GetType(expr);
+ if (auto* tuple = ty.as<TupleTypeNode>()) {
ffi::Array<NestedMsg<T>> res;
res.reserve(tuple->fields.size());
for (size_t i = 0; i < tuple->fields.size(); ++i) {
@@ -330,7 +330,7 @@ NestedMsg<T> MapToNestedMsgBySInfo(Expr expr, FType fmapleaf) {
} else {
field = TupleGetItem(expr, i);
}
- res.push_back(MapToNestedMsgBySInfo<T, FType>(field, fmapleaf));
+ res.push_back(MapToNestedMsgByType<T, FType>(field, fmapleaf));
}
return res;
} else {
@@ -520,8 +520,8 @@ void DecomposeNestedMsg(Expr expr, NestedMsg<T> msg, FType fvisitleaf) {
*/
template <typename T, std::size_t N, typename FType>
Expr TransformTupleLeaf(Expr expr, std::array<NestedMsg<T>, N> msgs, FType ftransleaf) {
- StructInfo sinfo = GetStructInfo(expr);
- if (const auto* tuple = sinfo.as<TupleStructInfoNode>()) {
+ Type ty = GetType(expr);
+ if (const auto* tuple = ty.as<TupleTypeNode>()) {
std::array<ffi::Array<NestedMsg<T>>, N> msg_arrays;
for (size_t i = 0; i < N; ++i) {
TVM_FFI_ICHECK(msgs[i].IsNested()) << "Expected nested to match tuple";
@@ -554,33 +554,32 @@ Expr TransformTupleLeaf(Expr expr, std::array<NestedMsg<T>, N> msgs, FType ftran
}
/*!
- * \brief Recursively transform the tuple structure in sinfo and msgs along with it.
+ * \brief Recursively transform the tuple structure in ty and msgs along with it.
*
- * This function will call ftransleaf for each leaf sinfo in sinfo.
+ * This function will call ftransleaf for each leaf ty in ty.
* This function will throw an error if the nesting structure in msg does not
- * match the tuple nesting structure in sinfo.
+ * match the tuple nesting structure in ty.
*
- * \param sinfo The input sinfo to be transform.
+ * \param ty The input ty to be transform.
* \param msgs The input messages to guide the transformation.
- * \param ftransleaf with signature ftransleaf(StructInfo, ffi::Array<NestedMsg<T>>)->StructInfo
+ * \param ftransleaf with signature ftransleaf(Type, ffi::Array<NestedMsg<T>>)->Type
* \tparam T the content type of nested msg
* \tparam N the number of messages
* \tparam FType The visit function type.
*/
template <typename T, std::size_t N, typename FType>
-StructInfo TransformTupleLeaf(StructInfo sinfo, std::array<NestedMsg<T>, N> msgs,
- FType ftransleaf) {
- if (const auto* tuple = sinfo.as<TupleStructInfoNode>()) {
+Type TransformTupleLeaf(Type ty, std::array<NestedMsg<T>, N> msgs, FType ftransleaf) {
+ if (const auto* tuple = ty.as<TupleTypeNode>()) {
std::array<ffi::Array<NestedMsg<T>>, N> msg_arrays;
for (size_t i = 0; i < N; ++i) {
TVM_FFI_ICHECK(msgs[i].IsNested()) << "Expected nested to match tuple";
msg_arrays[i] = msgs[i].NestedArray();
}
bool same = true;
- ffi::Array<StructInfo> fields;
+ ffi::Array<Type> fields;
fields.reserve(tuple->fields.size());
for (size_t i = 0; i < tuple->fields.size(); ++i) {
- StructInfo field = tuple->fields[i];
+ Type field = tuple->fields[i];
std::array<NestedMsg<T>, N> sub_msgs;
for (size_t j = 0; j < N; ++j) {
sub_msgs[j] = msg_arrays[j][i];
@@ -588,12 +587,12 @@ StructInfo TransformTupleLeaf(StructInfo sinfo, std::array<NestedMsg<T>, N> msgs
fields.push_back(TransformTupleLeaf(field, std::move(sub_msgs), ftransleaf));
same &= (fields.back().same_as(field));
}
- return same ? sinfo : TupleStructInfo(fields);
+ return same ? ty : TupleType(fields);
} else {
for (const auto& msg : msgs) {
TVM_FFI_ICHECK(msg.IsLeaf()) << "Expected leaf to match non-tuple";
}
- return ftransleaf(sinfo, msgs);
+ return ftransleaf(ty, msgs);
}
}
diff --git a/include/tvm/relax/op_attr_types.h b/include/tvm/relax/op_attr_types.h
index 1fd9b45c32..034b0dd2fe 100644
--- a/include/tvm/relax/op_attr_types.h
+++ b/include/tvm/relax/op_attr_types.h
@@ -24,8 +24,9 @@
#ifndef TVM_RELAX_OP_ATTR_TYPES_H_
#define TVM_RELAX_OP_ATTR_TYPES_H_
+#include <tvm/relax/block_builder.h>
#include <tvm/relax/expr.h>
-#include <tvm/relax/struct_info.h>
+#include <tvm/relax/type.h>
#include <tvm/te/tensor.h>
namespace tvm {
@@ -60,12 +61,12 @@ enum OpPatternKind {
using FCallPacked = ffi::String;
/*!
- * \brief Infer output struct info given the call
+ * \brief Infer output type given the call
*
* \param call The call expression to be derived.
* \param ctx The builder context.
*/
-using FInferStructInfo = ffi::TypedFunction<StructInfo(const Call& call, const BlockBuilder& ctx)>;
+using FInferType = ffi::TypedFunction<Type(const Call& call, const BlockBuilder& ctx)>;
/*!
* \brief The function type of a normalization function.
diff --git a/include/tvm/relax/script/builder/frame.h b/include/tvm/relax/script/builder/frame.h
index ab87aaf778..7a20ceed1a 100644
--- a/include/tvm/relax/script/builder/frame.h
+++ b/include/tvm/relax/script/builder/frame.h
@@ -100,15 +100,15 @@ class FunctionFrameNode : public SeqExprFrameNode {
/*! \brief The function params. */
ffi::Array<tvm::relax::Var> params;
/*!
- * \brief The function return struct info.
+ * \brief The function return type.
* \note Usually the function return type can be deduced by the function body.
* But we can use this field to specify a more "accurate" return type.
- * i.e. If the `ret_struct_info` is None, try to use the deduced type from body
- * If the `ret_struct_info` is not None, we can still take body.struct_info
- * if we ret_struct_info is base of body.struct_info. If not, we will
- * take the specified `ret_struct_info`.
+ * i.e. If the `ret_ty` is None, try to use the deduced type from body
+ * If the `ret_ty` is not None, we can still take body.ty
+ * if we ret_ty is base of body.ty. If not, we will
+ * take the specified `ret_ty`.
*/
- ffi::Optional<tvm::relax::StructInfo> ret_struct_info;
+ ffi::Optional<tvm::Type> ret_ty;
/*! \brief Whether the function is annotated as pure */
ffi::Optional<bool> is_pure;
/*! \brief Whether the function is annotated as private */
@@ -123,7 +123,7 @@ class FunctionFrameNode : public SeqExprFrameNode {
refl::ObjectDef<FunctionFrameNode>()
.def_ro("name", &FunctionFrameNode::name)
.def_ro("params", &FunctionFrameNode::params)
- .def_ro("ret_struct_info", &FunctionFrameNode::ret_struct_info)
+ .def_ro("ret_ty", &FunctionFrameNode::ret_ty)
.def_ro("is_pure", &FunctionFrameNode::is_pure)
.def_ro("attrs", &FunctionFrameNode::attrs);
// `binding_blocks` and `output` are inherited from SeqExprFrameNode.
diff --git a/include/tvm/relax/script/builder/ir.h b/include/tvm/relax/script/builder/ir.h
index 48318c8918..6a5df22245 100644
--- a/include/tvm/relax/script/builder/ir.h
+++ b/include/tvm/relax/script/builder/ir.h
@@ -21,7 +21,7 @@
#include <tvm/relax/expr.h>
#include <tvm/relax/script/builder/frame.h>
-#include <tvm/relax/struct_info.h>
+#include <tvm/relax/type.h>
#include <tvm/script/ir_builder/base.h>
namespace tvm {
@@ -42,10 +42,10 @@ TVM_DLL FunctionFrame Function(bool is_pure, bool is_private);
/*!
* \brief Add a parameter to the last function frame.
* \param name The name of the parameter.
- * \param struct_info The struct_info of the parameter.
+ * \param ty The ty of the parameter.
* \return The created function parameter var.
*/
-TVM_DLL tvm::relax::Var Arg(const ffi::String& name, const tvm::relax::StructInfo& struct_info);
+TVM_DLL tvm::relax::Var Arg(const ffi::String& name, const tvm::Type& ty);
/*!
* \brief Specify the name of the last function frame.
@@ -60,10 +60,10 @@ TVM_DLL void FuncName(const ffi::String& name);
TVM_DLL void FuncAttrs(ffi::Map<ffi::String, Any> attrs);
/*!
- * \brief Specify the return struct info of the last function frame.
- * \param ret_sinfo The return struct info.
+ * \brief Specify the return type of the last function frame.
+ * \param ret_ty The return type.
*/
-TVM_DLL void FuncRetStructInfo(const tvm::relax::StructInfo& ret_sinfo);
+TVM_DLL void FuncRetType(const tvm::Type& ret_ty);
/*!
* \brief Specify the return value of the last function frame.
@@ -96,21 +96,19 @@ TVM_DLL void DataflowBlockOutput(const ffi::Array<tvm::relax::Var>& vars);
/*!
* \brief Emit a binding to the last binding block frame.
* \param value The right side value of the bindings to be emitted.
- * \param annotate_struct_info The optional struct info annotation for the emitted value.
+ * \param annotate_ty The optional type annotation for the emitted value.
* \return The left side var of the emitted binding.
*/
-TVM_DLL tvm::relax::Var Emit(
- const tvm::relax::Expr& value,
- const ffi::Optional<tvm::relax::StructInfo>& annotate_struct_info = std::nullopt);
+TVM_DLL tvm::relax::Var Emit(const tvm::relax::Expr& value,
+ const ffi::Optional<tvm::Type>& annotate_ty = std::nullopt);
/*!
* \brief Emit a match_cast binding to the last binding block frame.
* \param value The value of the MatchCast to be emitted.
- * \param struct_info The struct info of the MatchCast to be emitted.
+ * \param ty The type of the MatchCast to be emitted.
* \return The left side var of the emitted binding.
*/
-TVM_DLL tvm::relax::Var EmitMatchCast(const tvm::relax::Expr& value,
- const tvm::relax::StructInfo& struct_info);
+TVM_DLL tvm::relax::Var EmitMatchCast(const tvm::relax::Expr& value, const tvm::Type& ty);
/*!
* \brief Emit a binding to the last binding block frame.
diff --git a/include/tvm/relax/struct_info.h b/include/tvm/relax/struct_info.h
deleted file mode 100644
index 049469027b..0000000000
--- a/include/tvm/relax/struct_info.h
+++ /dev/null
@@ -1,424 +0,0 @@
-/*
- * Licensed to the Apache Software Foundation (ASF) under one
- * or more contributor license agreements. See the NOTICE file
- * distributed with this work for additional information
- * regarding copyright ownership. The ASF licenses this file
- * to you under the Apache License, Version 2.0 (the
- * "License"); you may not use this file except in compliance
- * with the License. You may obtain a copy of the License at
- *
- * http://www.apache.org/licenses/LICENSE-2.0
- *
- * Unless required by applicable law or agreed to in writing,
- * software distributed under the License is distributed on an
- * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
- * KIND, either express or implied. See the License for the
- * specific language governing permissions and limitations
- * under the License.
- */
-#ifndef TVM_RELAX_STRUCT_INFO_H_
-#define TVM_RELAX_STRUCT_INFO_H_
-
-#include <tvm/ffi/reflection/registry.h>
-#include <tvm/ir/cast.h>
-#include <tvm/ir/env_func.h>
-#include <tvm/ir/source_map.h>
-#include <tvm/relax/block_builder.h>
-#include <tvm/relax/expr.h>
-#include <tvm/relax/type.h>
-#include <tvm/runtime/base.h>
-
-#include <utility>
-
-namespace tvm {
-namespace relax {
-
-/*!
- * \brief Opaque object.
- */
-class ObjectStructInfoNode : public StructInfoNode {
- public:
- static void RegisterReflection() {
- namespace refl = tvm::ffi::reflection;
- refl::ObjectDef<ObjectStructInfoNode>();
- }
- TVM_FFI_DECLARE_OBJECT_INFO_FINAL("relax.ObjectStructInfo", ObjectStructInfoNode, StructInfoNode);
-};
-
-/*!
- * \brief Managed reference to ObjectStructInfoNode.
- * \sa ObjectStructInfoNode
- */
-class ObjectStructInfo : public StructInfo {
- public:
- TVM_DLL ObjectStructInfo(Span span = Span());
-
- TVM_FFI_DEFINE_OBJECT_REF_METHODS_NOTNULLABLE(ObjectStructInfo, StructInfo, ObjectStructInfoNode);
-};
-
-/*!
- * \brief Primitive value.
- */
-class PrimStructInfoNode : public StructInfoNode {
- public:
- /*! \brief Underlying primitive value, if known */
- ffi::Optional<PrimExpr> value;
-
- /*! \brief Underlying data type of the primitive value */
- DataType dtype;
-
- static void RegisterReflection() {
- namespace refl = tvm::ffi::reflection;
- refl::ObjectDef<PrimStructInfoNode>()
- .def_ro("value", &PrimStructInfoNode::value)
- .def_ro("dtype", &PrimStructInfoNode::dtype);
- }
- TVM_FFI_DECLARE_OBJECT_INFO_FINAL("relax.PrimStructInfo", PrimStructInfoNode, StructInfoNode);
-};
-
-/*!
- * \brief Managed reference to PrimStructInfoNode.
- * \sa PrimStructInfoNode
- */
-class PrimStructInfo : public StructInfo {
- public:
- /* Construct a PrimStructInfo with a known dtype, but unknown value */
- TVM_DLL PrimStructInfo(DataType dtype, Span span = Span());
-
- /* Construct a PrimStructInfo with a known value */
- TVM_DLL PrimStructInfo(PrimExpr value, Span span = Span());
-
- TVM_FFI_DEFINE_OBJECT_REF_METHODS_NOTNULLABLE(PrimStructInfo, StructInfo, PrimStructInfoNode);
-};
-
-/*!
- * \brief StructInfo of shape value.
- */
-class ShapeStructInfoNode : public StructInfoNode {
- public:
- /*! \brief optionally stores the symbolic value patterns of the shape */
- ffi::Optional<ffi::Array<PrimExpr>> values;
- /*!
- * \brief The number of dimension of the shape, can be unknown.
- * \sa kUnknownNDim
- */
- int ndim;
-
- /*! \return Whether the struct info contains unknown ndim. */
- bool IsUnknownNdim() const { return ndim == kUnknownNDim; }
-
- static void RegisterReflection() {
- namespace refl = tvm::ffi::reflection;
- refl::ObjectDef<ShapeStructInfoNode>()
- .def_ro("values", &ShapeStructInfoNode::values)
- .def_ro("ndim", &ShapeStructInfoNode::ndim);
- }
- TVM_FFI_DECLARE_OBJECT_INFO_FINAL("relax.ShapeStructInfo", ShapeStructInfoNode, StructInfoNode);
-};
-
-/*!
- * \brief Managed reference to ShapeStructInfoNode.
- * \sa ShapeStructInfoNode
- */
-class ShapeStructInfo : public StructInfo {
- public:
- /*!
- * \brief Construction with known symbolic shape patterns
- * \param values The symbolic shape values
- * \param span The span of the AST.
- */
- TVM_DLL ShapeStructInfo(ffi::Array<PrimExpr> values, Span span = Span());
- /*!
- * \brief Construction with known unknown symbolic shape patterns.
- * \param ndim Number of dimensions -- can be kUnknownNDim
- * \param span The span of the AST.
- */
- TVM_DLL ShapeStructInfo(int ndim, Span span = Span());
-
- TVM_FFI_DEFINE_OBJECT_REF_METHODS_NOTNULLABLE(ShapeStructInfo, StructInfo, ShapeStructInfoNode);
-};
-
-/*!
- * \brief StructInfo of Tensor.
- */
-class TensorStructInfoNode : public StructInfoNode {
- public:
- /*!
- * \brief optionally store the shape expression of the tensor.
- * \note shape must be normalized: it can only be std::nullopt or ShapeExpr or Var.
- */
- ffi::Optional<Expr> shape;
- /*! \brief The virtual device, indicates where the tensor
- * is expected to be executed.
- */
- ffi::Optional<VDevice> vdevice;
- /*! \brief The content data type, use void to denote the dtype is unknown. */
- DataType dtype;
- /*!
- * \brief The number of dimension of the tensor, can be unknown.
- * \sa kUnknownNDim
- */
- int ndim;
-
- /*! \return Whether the struct info contains unknown ndim. */
- bool IsUnknownNdim() const { return ndim == kUnknownNDim; }
-
- /*! \return Whether the struct info contains unknown dtype. */
- bool IsUnknownDtype() const { return dtype.is_void(); }
-
- /*! \return Shape if it is known. */
- ffi::Optional<ffi::Array<PrimExpr>> GetShape() const {
- if (!shape.defined()) return {};
- ShapeStructInfo shape_sinfo = Downcast<ShapeStructInfo>(this->shape.value()->struct_info_);
- return shape_sinfo->values;
- }
-
- static void RegisterReflection() {
- namespace refl = tvm::ffi::reflection;
- refl::ObjectDef<TensorStructInfoNode>()
- .def_ro("shape", &TensorStructInfoNode::shape)
- .def_ro("dtype", &TensorStructInfoNode::dtype)
- .def_ro("vdevice", &TensorStructInfoNode::vdevice)
- .def_ro("ndim", &TensorStructInfoNode::ndim);
- }
- TVM_FFI_DECLARE_OBJECT_INFO_FINAL("relax.TensorStructInfo", TensorStructInfoNode, StructInfoNode);
-};
-
-/*!
- * \brief Managed reference to TensorStructInfoNode.
- * \sa TensorStructInfoNode
- */
-class TensorStructInfo : public StructInfo {
- public:
- /*!
- * \brief Construction with a known shape expression.
- * \param shape The shape of the tensor.
- * \param dtype The data type of tensor's elements.
- * \param vdevice The virtual device.
- * \param span The span of the AST.
- *
- * \note shape must already be normalized.
- */
- TVM_DLL TensorStructInfo(Expr shape, DataType dtype,
- ffi::Optional<VDevice> vdevice = std::nullopt, Span span = Span());
-
- /*!
- * \brief Construction with an unknown shape expression.
- * \param dtype The data type of tensor's elements.
- * \param ndim The number of dimensions
- * \param vdevice The virtual device.
- * \param span The span of the AST.
- */
- TVM_DLL TensorStructInfo(DataType dtype, int ndim, ffi::Optional<VDevice> vdevice = std::nullopt,
- Span span = Span());
-
- TVM_FFI_DEFINE_OBJECT_REF_METHODS_NULLABLE(TensorStructInfo, StructInfo, TensorStructInfoNode);
-};
-
-/*!
- * \brief StructInfo of Tuple.
- */
-class TupleStructInfoNode : public StructInfoNode {
- public:
- /*! \brief The struct info of tuple fields. */
- ffi::Array<StructInfo> fields;
-
- static void RegisterReflection() {
- namespace refl = tvm::ffi::reflection;
- refl::ObjectDef<TupleStructInfoNode>().def_ro("fields", &TupleStructInfoNode::fields);
- }
- TVM_FFI_DECLARE_OBJECT_INFO_FINAL("relax.TupleStructInfo", TupleStructInfoNode, StructInfoNode);
-};
-
-/*!
- * \brief Managed reference to TupleStructInfoNode.
- * \sa TupleStructInfoNode
- */
-class TupleStructInfo : public StructInfo {
- public:
- /*!
- * \brief Constructor
- * \param fields Struct info of tuple fields.
- * \param span The span of the AST.
- */
- TVM_DLL TupleStructInfo(ffi::Array<StructInfo> fields, Span span = Span());
-
- TVM_FFI_DEFINE_OBJECT_REF_METHODS_NOTNULLABLE(TupleStructInfo, StructInfo, TupleStructInfoNode);
-};
-
-/*!
- * \brief custom-defined StructInfo derivation function.
- * \param call The call expression to be derived.
- * \param ctx The builder context.
- * \return The derived struct info of the call.
- */
-using StructInfoDeriveFunc = TypedEnvFunc<StructInfo(const Call& call, const BlockBuilder& ctx)>;
-
-/*!
- * \brief Structure information about function.
- *
- * This data structure contains enough information for us to
- * do best-effort structure information deduction.
- */
-class FuncStructInfoNode : public StructInfoNode {
- public:
- /*!
- * \brief The parameter struct info of the function.
- * \note When params is std::nullopt means the function can take arbitrary number of arguments.
- * We define such functions as Opaque function.
- */
- ffi::Optional<ffi::Array<StructInfo>> params;
- /*!
- * \brief The struct info of the function's return value.
- */
- StructInfo ret;
- /*!
- * \brief Derivation function of opaque functions that may take any number of parameters.
- * \note When derive_func is not empty, then params should be std::nullopt,
- * ret should be ObjectStructInfo()
- */
- ffi::Optional<StructInfoDeriveFunc> derive_func;
- /*!
- * \brief Whether the function is pure.
- * \note This parameter should be set to true only if the function is pure on all inputs.
- * If the function _may_ have visible side effects, set it to false.
- */
- bool purity;
-
- /*!
- * \return Whether the func struct info is opaque.
- * \note We define a function as opaque we have no constraints on params.
- */
- bool IsOpaque() const { return !params.defined(); }
-
- static void RegisterReflection() {
- namespace refl = tvm::ffi::reflection;
- refl::ObjectDef<FuncStructInfoNode>()
- .def_ro("params", &FuncStructInfoNode::params, refl::AttachFieldFlag::SEqHashDefRecursive())
- .def_ro("ret", &FuncStructInfoNode::ret)
- .def_ro("derive_func", &FuncStructInfoNode::derive_func)
- .def_ro("purity", &FuncStructInfoNode::purity);
- }
- TVM_FFI_DECLARE_OBJECT_INFO_FINAL("relax.FuncStructInfo", FuncStructInfoNode, StructInfoNode);
-};
-
-/*!
- * \brief Managed reference to FuncStructInfoNode.
- * \sa FuncStructInfoNode
- */
-class FuncStructInfo : public StructInfo {
- public:
- explicit FuncStructInfo(ffi::ObjectPtr<FuncStructInfoNode> data) : StructInfo(ffi::UnsafeInit{}) {
- TVM_FFI_ICHECK(data != nullptr);
- data_ = std::move(data);
- }
- /*!
- * \brief Constructor from parameter struct info and return value struct info.
- * \param params The struct info of function parameters.
- * \param ret The return value struct info.
- * \param purity The purity of the function (true by default).
- * \param span The span of the AST.
- *
- * \note If the ret contains variables(tirx::Var and relax::Var), they must be deducible from
- * params. If you are unsure, you can always erase ret to static.
- */
- TVM_DLL FuncStructInfo(ffi::Array<StructInfo> params, StructInfo ret, bool purity = true,
- Span span = Span());
-
- /*!
- * \brief Constructing an opaque function struct info using derive_func.
- *
- * \param derive_func Derivation function.
- * \param purity The purity of the function
- * (false by default: most external functions are not pure).
- * \param span The span of the AST.
- *
- * \return The FuncStructInfo for opaque packedfunc.
- * \note Defaults to an derive func that always return ObjectStructInfo if not specified.
- */
- TVM_DLL static FuncStructInfo OpaqueFunc(StructInfoDeriveFunc derive_func, bool purity = false,
- Span span = Span());
-
- /*!
- * \brief Construct an opaque function using from return struct info.
- *
- * \param ret The struct info of the return value.
- * \param purity The purity of the function
- * (false by default: most external functions are not pure).
- * \param span The span of the AST.
- *
- * \return The FuncStructInfo for opaque packedfunc.
- * \note Defaults to an derive func that always return ObjectStructInfo if not specified.
- */
- TVM_DLL static FuncStructInfo OpaqueFunc(StructInfo ret = ObjectStructInfo(), bool purity = false,
- Span span = Span());
-
- TVM_FFI_DEFINE_OBJECT_REF_METHODS_NOTNULLABLE(FuncStructInfo, StructInfo, FuncStructInfoNode);
-};
-
-/*!
- * \brief Match and check if expr have StructInfo T and return it.
- *
- * \param expr The input expression.
- * \return The result of match.
- * \tparam T the underlying structure info type
- */
-template <typename T>
-inline ffi::Optional<T> MatchStructInfo(const Expr& expr) {
- using TNode = typename T::ContainerType;
- if (const TNode* ptr = expr->struct_info_.as<TNode>()) {
- return ffi::GetRef<T>(ptr);
- } else {
- return std::nullopt;
- }
-}
-
-/*!
- * \brief Get the structure info of a given expr and try to cast it as const T*.
- *
- * \param expr The input expression.
- * \return The pointer. Returns nullptr if the type does not match
- * \tparam T the underlying structure info type
- */
-template <typename T>
-inline const T* GetStructInfoAs(const Expr& expr) {
- TVM_FFI_ICHECK(expr->struct_info_.defined())
- << "The struct_info is not populated, check if you have normalized the expr";
- return expr->struct_info_.as<T>();
-}
-
-/*!
- * \brief Get the underlying structure info of expr.
- *
- * \param expr The input expression.
- * \return underlying struct info.
- */
-inline StructInfo GetStructInfo(const Expr& expr) {
- auto* ptr = expr->struct_info_.as<StructInfoNode>();
- TVM_FFI_ICHECK(ptr) << "The struct_info is not populated, check if you have normalized the expr";
- return ffi::GetRef<StructInfo>(ptr);
-}
-
-/*!
- * \brief Whether the expr has void struct info.
- *
- * \param expr The input expression.
- * \return Whether the expr has void struct info.
- */
-inline bool HasVoidStructInfo(const Expr& expr) {
- auto* ptr = expr->struct_info_.as<TupleStructInfoNode>();
- return ptr != nullptr && ptr->fields.size() == 0;
-}
-
-/*!
- * \brief Update the struct info of an Expr.
- * \param expr The Expr whose struct info to be updated.
- * \param struct_info The struct_info assigned.
- * \note We ensure idempotence, that is we can only update the struct_info of an Expr only
- * if the original one is nullptr.
- */
-TVM_DLL void UpdateStructInfo(Expr expr, StructInfo struct_info);
-
-} // namespace relax
-} // namespace tvm
-#endif // TVM_RELAX_STRUCT_INFO_H_
diff --git a/include/tvm/relax/struct_info_functor.h b/include/tvm/relax/struct_info_functor.h
deleted file mode 100644
index c4b3ea31a1..0000000000
--- a/include/tvm/relax/struct_info_functor.h
+++ /dev/null
@@ -1,160 +0,0 @@
-/*
- * Licensed to the Apache Software Foundation (ASF) under one
- * or more contributor license agreements. See the NOTICE file
- * distributed with this work for additional information
- * regarding copyright ownership. The ASF licenses this file
- * to you under the Apache License, Version 2.0 (the
- * "License"); you may not use this file except in compliance
- * with the License. You may obtain a copy of the License at
- *
- * http://www.apache.org/licenses/LICENSE-2.0
- *
- * Unless required by applicable law or agreed to in writing,
- * software distributed under the License is distributed on an
- * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
- * KIND, either express or implied. See the License for the
- * specific language governing permissions and limitations
- * under the License.
- */
-
-/*!
- * \file tvm/relax/struct_info_functor.h
- * \brief Functors and visitors for struct info.
- */
-#ifndef TVM_RELAX_STRUCT_INFO_FUNCTOR_H_
-#define TVM_RELAX_STRUCT_INFO_FUNCTOR_H_
-
-#include <tvm/ir/node_functor.h>
-#include <tvm/relax/distributed/struct_info.h>
-#include <tvm/relax/struct_info.h>
-
-#include <utility>
-
-namespace tvm {
-namespace relax {
-
-template <typename FStructInfo>
-class StructInfoFunctor;
-
-// functions to be overriden.
-#define STRUCT_INFO_FUNCTOR_DEFAULT \
- { \
- return VisitStructInfoDefault_(op, std::forward<Args>(args)...); \
- }
-
-#define TVM_STRUCT_INFO_FUNCTOR_DISPATCH(OP) \
- vtable.template set_dispatch<OP>([](const ffi::ObjectRef& n, TSelf* self, Args... args) { \
- return self->VisitStructInfo_(static_cast<const OP*>(n.get()), std::forward<Args>(args)...); \
- });
-
-template <typename R, typename... Args>
-class StructInfoFunctor<R(const StructInfo& n, Args...)> {
- private:
- using TSelf = StructInfoFunctor<R(const StructInfo& n, Args...)>;
- using FStructInfo = tvm::NodeFunctor<R(const ffi::ObjectRef& n, TSelf* self, Args...)>;
-
- public:
- /*! \brief the result type of this functor */
- using result_type = R;
- /*! \brief virtual destructor */
- virtual ~StructInfoFunctor() {}
- /*!
- * \brief Same as call.
- * \param n The expression node.
- * \param args Additional arguments.
- * \return The result of the call
- */
- R operator()(const StructInfo& n, Args... args) {
- return VisitStructInfo(n, std::forward<Args>(args)...);
- }
- /*!
- * \brief The functor call.
- * \param n The expression node.
- * \param args Additional arguments.
- * \return The result of the call
- */
- virtual R VisitStructInfo(const StructInfo& n, Args... args) {
- TVM_FFI_ICHECK(n.defined());
- static FStructInfo vtable = InitVTable();
- return vtable(n, this, std::forward<Args>(args)...);
- }
- // Functions that can be overriden by subclass
- virtual R VisitStructInfo_(const ObjectStructInfoNode* op,
- Args... args) STRUCT_INFO_FUNCTOR_DEFAULT;
- virtual R VisitStructInfo_(const PrimStructInfoNode* op,
- Args... args) STRUCT_INFO_FUNCTOR_DEFAULT;
- virtual R VisitStructInfo_(const ShapeStructInfoNode* op,
- Args... args) STRUCT_INFO_FUNCTOR_DEFAULT;
- virtual R VisitStructInfo_(const TensorStructInfoNode* op,
- Args... args) STRUCT_INFO_FUNCTOR_DEFAULT;
- virtual R VisitStructInfo_(const distributed::DTensorStructInfoNode* op,
- Args... args) STRUCT_INFO_FUNCTOR_DEFAULT;
- virtual R VisitStructInfo_(const TupleStructInfoNode* op,
- Args... args) STRUCT_INFO_FUNCTOR_DEFAULT;
- virtual R VisitStructInfo_(const FuncStructInfoNode* op,
- Args... args) STRUCT_INFO_FUNCTOR_DEFAULT;
- virtual R VisitStructInfoDefault_(const ffi::Object* op, Args...) {
- TVM_FFI_THROW(InternalError) << "Do not have a default for " << op->GetTypeKey();
- throw; // unreachable, written to stop compiler warning
- }
-
- private:
- // initialize the vtable.
- static FStructInfo InitVTable() {
- FStructInfo vtable;
- // Set dispatch
- TVM_STRUCT_INFO_FUNCTOR_DISPATCH(ObjectStructInfoNode);
- TVM_STRUCT_INFO_FUNCTOR_DISPATCH(PrimStructInfoNode);
- TVM_STRUCT_INFO_FUNCTOR_DISPATCH(ShapeStructInfoNode);
- TVM_STRUCT_INFO_FUNCTOR_DISPATCH(TensorStructInfoNode);
- TVM_STRUCT_INFO_FUNCTOR_DISPATCH(distributed::DTensorStructInfoNode);
- TVM_STRUCT_INFO_FUNCTOR_DISPATCH(TupleStructInfoNode);
- TVM_STRUCT_INFO_FUNCTOR_DISPATCH(FuncStructInfoNode);
- vtable.Finalize();
- return vtable;
- }
-};
-
-#undef TVM_STRUCT_INFO_FUNCTOR_DISPATCH
-
-/*!
- * \brief A struct info visitor.
- */
-class TVM_DLL StructInfoVisitor : public StructInfoFunctor<void(const StructInfo& n)> {
- public:
- void VisitStructInfo_(const ObjectStructInfoNode* op) override;
- void VisitStructInfo_(const PrimStructInfoNode* op) override;
- void VisitStructInfo_(const ShapeStructInfoNode* op) override;
- void VisitStructInfo_(const TensorStructInfoNode* op) override;
- void VisitStructInfo_(const distributed::DTensorStructInfoNode* op) override;
- void VisitStructInfo_(const TupleStructInfoNode* op) override;
- void VisitStructInfo_(const FuncStructInfoNode* op) override;
-
- protected:
- // two functions to override when visit expr fields in struct info.
- virtual void VisitStructInfoExprField(const Expr& expr) {}
- virtual void VisitStructInfoExprField(const PrimExpr& expr) {}
-};
-
-/*!
- * \brief StructInfoMutator that mutates struct info.
- */
-class TVM_DLL StructInfoMutator : public StructInfoFunctor<StructInfo(const StructInfo& n)> {
- public:
- StructInfo VisitStructInfo_(const ObjectStructInfoNode* op) override;
- StructInfo VisitStructInfo_(const PrimStructInfoNode* op) override;
- StructInfo VisitStructInfo_(const ShapeStructInfoNode* op) override;
- StructInfo VisitStructInfo_(const TensorStructInfoNode* op) override;
- StructInfo VisitStructInfo_(const distributed::DTensorStructInfoNode* op) override;
- StructInfo VisitStructInfo_(const TupleStructInfoNode* op) override;
- StructInfo VisitStructInfo_(const FuncStructInfoNode* op) override;
-
- protected:
- // two functions to override when visit expr fields in struct info.
- virtual Expr VisitStructInfoExprField(const Expr& expr) { return expr; }
- virtual PrimExpr VisitStructInfoExprField(const PrimExpr& expr) { return expr; }
-};
-
-} // namespace relax
-} // namespace tvm
-#endif // TVM_RELAX_STRUCT_INFO_FUNCTOR_H_
diff --git a/include/tvm/relax/transform.h b/include/tvm/relax/transform.h
index 493e51ef50..d0d0d1bb54 100644
--- a/include/tvm/relax/transform.h
+++ b/include/tvm/relax/transform.h
@@ -153,7 +153,7 @@ TVM_DLL Pass AttachGlobalSymbol();
/*!
* \brief Transform Relax IR to normal form: transform AST to A-normal form, and fill the
- * struct_info_ of expressions.
+ * ty of expressions.
*
* \return The Pass.
*/
diff --git a/include/tvm/relax/type.h b/include/tvm/relax/type.h
index b70a2756b7..fcaddc3ab5 100644
--- a/include/tvm/relax/type.h
+++ b/include/tvm/relax/type.h
@@ -19,7 +19,7 @@
/*!
* \file tvm/relax/type.h
- * \brief Relax Types.
+ * \brief Relax types, including the richer dependent Relax type nodes.
*/
#ifndef TVM_RELAX_TYPE_H_
#define TVM_RELAX_TYPE_H_
@@ -27,122 +27,437 @@
#include <tvm/ffi/function.h>
#include <tvm/ffi/reflection/registry.h>
#include <tvm/ir/attrs.h>
+#include <tvm/ir/cast.h>
#include <tvm/ir/env_func.h>
+#include <tvm/ir/global_info.h>
#include <tvm/ir/type.h>
#include <tvm/tirx/expr.h>
#include <string>
+#include <utility>
namespace tvm {
namespace relax {
+using Expr = RelaxExpr;
+using ExprNode = RelaxExprNode;
+
+class BlockBuilder;
+class Call;
+
/*! \brief Indicates the number of dimensions of a tensor is unknown at compile time. */
static constexpr int kUnknownNDim = -1;
+using tvm::TupleType;
+using tvm::TupleTypeNode;
+
+class PackedFuncTypeNode : public TypeNode {
+ public:
+ static void RegisterReflection() {
+ namespace refl = tvm::ffi::reflection;
+ refl::ObjectDef<PackedFuncTypeNode>();
+ }
+ TVM_FFI_DECLARE_OBJECT_INFO_FINAL("relax.PackedFuncType", PackedFuncTypeNode, TypeNode);
+};
+
+class PackedFuncType : public Type {
+ public:
+ TVM_DLL PackedFuncType(Span span = Span());
+
+ TVM_FFI_DEFINE_OBJECT_REF_METHODS_NOTNULLABLE(PackedFuncType, Type, PackedFuncTypeNode);
+};
+
+/*!
+ * \brief Base type of all Relax type information.
+ *
+ * Type stores possible type information deduced during compile-time.
+ * It encapsulates both static type and runtime information such as shape.
+ *
+ * Type of each non-primitive Expr can be deduced during compilation in a
+ * "best-effort" manner.
+ *
+ * When ty appears in function parameter and return signatures, it
+ * implies a runtime check that matches the type information with the value.
+ *
+ * When it appears in Expr, it follows "assume-semantics", which means the
+ * compiler will take the deduced information as it is and only do best effort
+ * proofs and checks.
+ *
+ * Each type can be uniquely erased to a static-type. The compiler will
+ * still compile the code, with less information, when we erase to the static
+ * type.
+ *
+ * If a Type contains an Expr field, then that field must already be
+ * normalized through NormalizeArg. This invariant is checked in constructors
+ * and simplifies assumptions during type deduction.
+ */
+/*!
+ * \brief Opaque object.
+ */
+class ObjectTypeNode : public TypeNode {
+ public:
+ static void RegisterReflection() {
+ namespace refl = tvm::ffi::reflection;
+ refl::ObjectDef<ObjectTypeNode>();
+ }
+ TVM_FFI_DECLARE_OBJECT_INFO_FINAL("relax.ObjectType", ObjectTypeNode, TypeNode);
+};
+
+/*!
+ * \brief Managed reference to ObjectTypeNode.
+ * \sa ObjectTypeNode
+ */
+class ObjectType : public Type {
+ public:
+ TVM_DLL ObjectType(Span span = Span());
+
+ TVM_FFI_DEFINE_OBJECT_REF_METHODS_NOTNULLABLE(ObjectType, Type, ObjectTypeNode);
+};
+
+/*!
+ * \brief Primitive value.
+ */
+class PrimTypeNode : public TypeNode {
+ public:
+ /*! \brief Underlying primitive value, if known */
+ ffi::Optional<PrimExpr> value;
+
+ /*! \brief Underlying data type of the primitive value */
+ DataType dtype;
+
+ static void RegisterReflection() {
+ namespace refl = tvm::ffi::reflection;
+ refl::ObjectDef<PrimTypeNode>()
+ .def_ro("value", &PrimTypeNode::value)
+ .def_ro("dtype", &PrimTypeNode::dtype);
+ }
+ TVM_FFI_DECLARE_OBJECT_INFO_FINAL("relax.PrimType", PrimTypeNode, TypeNode);
+};
+
+/*!
+ * \brief Managed reference to PrimTypeNode.
+ * \sa PrimTypeNode
+ */
+class PrimType : public Type {
+ public:
+ /* Construct a PrimType with a known dtype, but unknown value */
+ TVM_DLL PrimType(DataType dtype, Span span = Span());
+
+ /* Construct a PrimType with a known value */
+ TVM_DLL PrimType(PrimExpr value, Span span = Span());
+
+ TVM_FFI_DEFINE_OBJECT_REF_METHODS_NOTNULLABLE(PrimType, Type, PrimTypeNode);
+};
+
+/*!
+ * \brief Type of shape value.
+ */
class ShapeTypeNode : public TypeNode {
public:
- /*! \brief size of the shape. */
+ /*! \brief optionally stores the symbolic value patterns of the shape */
+ ffi::Optional<ffi::Array<PrimExpr>> values;
+ /*!
+ * \brief The number of dimension of the shape, can be unknown.
+ * \sa kUnknownNDim
+ */
int ndim;
+ /*! \return Whether the type contains unknown ndim. */
+ bool IsUnknownNdim() const { return ndim == kUnknownNDim; }
+
static void RegisterReflection() {
namespace refl = tvm::ffi::reflection;
- refl::ObjectDef<ShapeTypeNode>().def_ro("ndim", &ShapeTypeNode::ndim);
+ refl::ObjectDef<ShapeTypeNode>()
+ .def_ro("values", &ShapeTypeNode::values)
+ .def_ro("ndim", &ShapeTypeNode::ndim);
}
TVM_FFI_DECLARE_OBJECT_INFO_FINAL("relax.ShapeType", ShapeTypeNode, TypeNode);
};
+/*!
+ * \brief Managed reference to ShapeTypeNode.
+ * \sa ShapeTypeNode
+ */
class ShapeType : public Type {
public:
+ /*!
+ * \brief Construction with known symbolic shape patterns
+ * \param values The symbolic shape values
+ * \param span The span of the AST.
+ */
+ TVM_DLL ShapeType(ffi::Array<PrimExpr> values, Span span = Span());
+ /*!
+ * \brief Construction with known unknown symbolic shape patterns.
+ * \param ndim Number of dimensions -- can be kUnknownNDim
+ * \param span The span of the AST.
+ */
TVM_DLL ShapeType(int ndim, Span span = Span());
TVM_FFI_DEFINE_OBJECT_REF_METHODS_NOTNULLABLE(ShapeType, Type, ShapeTypeNode);
};
/*!
- * \brief Dynamic version of TensorType
- *
- * Use relax::TensorStructInfo for more detailed (possibly dynamic) shape constrains
+ * \brief Type of Tensor.
*/
class TensorTypeNode : public TypeNode {
public:
/*!
- * \brief The number of dimensions of the tensor, use -1 to denote tensor with unknown number of
- * dimensions.
+ * \brief optionally store the shape expression of the tensor.
+ * \note shape must be normalized: it can only be std::nullopt or ShapeExpr or Var.
*/
- int ndim;
+ ffi::Optional<Expr> shape;
+ /*! \brief The virtual device, indicates where the tensor
+ * is expected to be executed.
+ */
+ ffi::Optional<VDevice> vdevice;
/*! \brief The content data type, use void to denote the dtype is unknown. */
DataType dtype;
+ /*!
+ * \brief The number of dimension of the tensor, can be unknown.
+ * \sa kUnknownNDim
+ */
+ int ndim;
+
+ /*! \return Whether the type contains unknown ndim. */
+ bool IsUnknownNdim() const { return ndim == kUnknownNDim; }
+
+ /*! \return Whether the type contains unknown dtype. */
+ bool IsUnknownDtype() const { return dtype.is_void(); }
+
+ /*! \return Shape if it is known. */
+ ffi::Optional<ffi::Array<PrimExpr>> GetShape() const {
+ if (!shape.defined()) return {};
+ const Expr& shape_expr = this->shape.value();
+ if (!shape_expr->ty.defined()) return {};
+ if (const auto* shape_ty = shape_expr->ty.as<ShapeTypeNode>()) {
+ return shape_ty->values;
+ }
+ return {};
+ }
static void RegisterReflection() {
namespace refl = tvm::ffi::reflection;
refl::ObjectDef<TensorTypeNode>()
- .def_ro("ndim", &TensorTypeNode::ndim)
- .def_ro("dtype", &TensorTypeNode::dtype);
+ .def_ro("shape", &TensorTypeNode::shape)
+ .def_ro("dtype", &TensorTypeNode::dtype)
+ .def_ro("vdevice", &TensorTypeNode::vdevice)
+ .def_ro("ndim", &TensorTypeNode::ndim);
}
-
- inline bool IsUnknownNdim() const { return ndim == kUnknownNDim; }
-
- inline bool IsUnknownDtype() const { return dtype.is_void(); }
- TVM_FFI_DECLARE_OBJECT_INFO_FINAL("relax.DynTensorType", TensorTypeNode, TypeNode);
+ TVM_FFI_DECLARE_OBJECT_INFO_FINAL("relax.TensorType", TensorTypeNode, TypeNode);
};
/*!
* \brief Managed reference to TensorTypeNode.
- * \sa TensorTypeNode.
+ * \sa TensorTypeNode
*/
class TensorType : public Type {
public:
+ explicit TensorType(ffi::ObjectPtr<TensorTypeNode> data) : Type(ffi::UnsafeInit{}) {
+ TVM_FFI_ICHECK(data != nullptr);
+ data_ = std::move(data);
+ }
+
/*!
- * \brief Constructor.
- * \param ndim The number of dimensions of the tensor.
- * \param dtype The runtime dtype of the tensor's elements.
- * \param span The span.
+ * \brief Construction with a known shape expression.
+ * \param shape The shape of the tensor.
+ * \param dtype The data type of tensor's elements.
+ * \param vdevice The virtual device.
+ * \param span The span of the AST.
+ *
+ * \note shape must already be normalized.
*/
- TVM_DLL TensorType(int ndim, DataType dtype, Span span = Span());
+ TVM_DLL TensorType(Expr shape, DataType dtype, ffi::Optional<VDevice> vdevice = std::nullopt,
+ Span span = Span());
/*!
- * \brief Create a TensorType with unknown ndim.
+ * \brief Construction with an unknown shape expression.
+ * \param dtype The data type of tensor's elements.
+ * \param ndim The number of dimensions
+ * \param vdevice The virtual device.
+ * \param span The span of the AST.
*/
- TVM_DLL static TensorType CreateUnknownNDim(DataType dtype, Span span = Span());
+ TVM_DLL TensorType(DataType dtype, int ndim, ffi::Optional<VDevice> vdevice = std::nullopt,
+ Span span = Span());
- TVM_FFI_DEFINE_OBJECT_REF_METHODS_NULLABLE(TensorType, Type, TensorTypeNode);
+ TVM_FFI_DEFINE_OBJECT_REF_METHODS_NOTNULLABLE(TensorType, Type, TensorTypeNode);
};
-using TensorTypeNode = TensorTypeNode;
-using TensorType = TensorType;
+/*!
+ * \brief custom-defined Type derivation function.
+ * \param call The call expression to be derived.
+ * \param ctx The builder context.
+ * \return The derived type of the call.
+ */
+using TypeDeriveFunc = TypedEnvFunc<Type(const Call& call, const BlockBuilder& ctx)>;
-class ObjectTypeNode : public TypeNode {
+/*!
+ * \brief Function type information.
+ *
+ * This data structure contains enough information for us to do best-effort
+ * type deduction.
+ */
+class FuncTypeNode : public TypeNode {
public:
+ /*!
+ * \brief The parameter type of the function.
+ * \note When params is std::nullopt means the function can take arbitrary number of arguments.
+ * We define such functions as Opaque function.
+ */
+ ffi::Optional<ffi::Array<Type>> params;
+ /*!
+ * \brief The type of the function's return value.
+ */
+ Type ret;
+ /*!
+ * \brief Derivation function of opaque functions that may take any number of parameters.
+ * \note When derive_func is not empty, then params should be std::nullopt,
+ * ret should be ObjectType()
+ */
+ ffi::Optional<TypeDeriveFunc> derive_func;
+ /*!
+ * \brief Whether the function is pure.
+ * \note This parameter should be set to true only if the function is pure on all inputs.
+ * If the function _may_ have visible side effects, set it to false.
+ */
+ bool purity;
+
+ /*!
+ * \return Whether the func type is opaque.
+ * \note We define a function as opaque we have no constraints on params.
+ */
+ bool IsOpaque() const { return !params.defined(); }
+
static void RegisterReflection() {
namespace refl = tvm::ffi::reflection;
- refl::ObjectDef<ObjectTypeNode>();
+ refl::ObjectDef<FuncTypeNode>()
+ .def_ro("params", &FuncTypeNode::params, refl::AttachFieldFlag::SEqHashDefRecursive())
+ .def_ro("ret", &FuncTypeNode::ret)
+ .def_ro("derive_func", &FuncTypeNode::derive_func)
+ .def_ro("purity", &FuncTypeNode::purity);
}
- TVM_FFI_DECLARE_OBJECT_INFO_FINAL("relax.ObjectType", ObjectTypeNode, TypeNode);
+ TVM_FFI_DECLARE_OBJECT_INFO_FINAL("relax.FuncType", FuncTypeNode, TypeNode);
};
-class ObjectType : public Type {
+/*!
+ * \brief Managed reference to FuncTypeNode.
+ * \sa FuncTypeNode
+ */
+class FuncType : public Type {
public:
- TVM_DLL ObjectType(Span span = Span());
+ explicit FuncType(ffi::ObjectPtr<FuncTypeNode> data) : Type(ffi::UnsafeInit{}) {
+ TVM_FFI_ICHECK(data != nullptr);
+ data_ = std::move(data);
+ }
+ /*!
+ * \brief Constructor from parameter type and return value type.
+ * \param params The type of function parameters.
+ * \param ret The return value type.
+ * \param purity The purity of the function (true by default).
+ * \param span The span of the AST.
+ *
+ * \note If the ret contains variables(tirx::Var and relax::Var), they must be deducible from
+ * params. If you are unsure, you can always erase ret to static.
+ */
+ TVM_DLL FuncType(ffi::Array<Type> params, Type ret, bool purity = true, Span span = Span());
- TVM_FFI_DEFINE_OBJECT_REF_METHODS_NOTNULLABLE(ObjectType, Type, ObjectTypeNode);
+ /*!
+ * \brief Constructing an opaque function type using derive_func.
+ *
+ * \param derive_func Derivation function.
+ * \param purity The purity of the function
+ * (false by default: most external functions are not pure).
+ * \param span The span of the AST.
+ *
+ * \return The FuncType for opaque packedfunc.
+ * \note Defaults to an derive func that always return ObjectType if not specified.
+ */
+ TVM_DLL static FuncType OpaqueFunc(TypeDeriveFunc derive_func, bool purity = false,
+ Span span = Span());
+
+ /*!
+ * \brief Construct an opaque function using from return type.
+ *
+ * \param ret The type of the return value.
+ * \param purity The purity of the function
+ * (false by default: most external functions are not pure).
+ * \param span The span of the AST.
+ *
+ * \return The FuncType for opaque packedfunc.
+ * \note Defaults to an derive func that always return ObjectType if not specified.
+ */
+ TVM_DLL static FuncType OpaqueFunc(Type ret = ObjectType(), bool purity = false,
+ Span span = Span());
+
+ TVM_FFI_DEFINE_OBJECT_REF_METHODS_NOTNULLABLE(FuncType, Type, FuncTypeNode);
};
-class PackedFuncTypeNode : public TypeNode {
- public:
- static void RegisterReflection() {
- namespace refl = tvm::ffi::reflection;
- refl::ObjectDef<PackedFuncTypeNode>();
+/*!
+ * \brief Match and check if expr has Relax type T and return it.
+ *
+ * \param expr The input expression.
+ * \return The result of match.
+ * \tparam T the underlying Relax type
+ */
+template <typename T>
+inline ffi::Optional<T> MatchType(const Expr& expr) {
+ if (!expr.defined()) {
+ return std::nullopt;
}
- TVM_FFI_DECLARE_OBJECT_INFO_FINAL("relax.PackedFuncType", PackedFuncTypeNode, TypeNode);
-};
+ using TNode = typename T::ContainerType;
+ if (const TNode* ptr = expr->ty.as<TNode>()) {
+ return ffi::GetRef<T>(ptr);
+ } else {
+ return std::nullopt;
+ }
+}
-class PackedFuncType : public Type {
- public:
- TVM_DLL PackedFuncType(Span span = Span());
+/*!
+ * \brief Get the type of a given expr and try to cast it as const T*.
+ *
+ * \param expr The input expression.
+ * \return The pointer. Returns nullptr if the type does not match.
+ * \tparam T the underlying Relax type node
+ */
+template <typename T>
+inline const T* GetTypeAs(const Expr& expr) {
+ TVM_FFI_ICHECK(expr->ty.defined())
+ << "The type is not populated, check if you have normalized the expr";
+ return expr->ty.as<T>();
+}
- TVM_FFI_DEFINE_OBJECT_REF_METHODS_NOTNULLABLE(PackedFuncType, Type, PackedFuncTypeNode);
-};
+/*!
+ * \brief Get the underlying Relax type of expr.
+ *
+ * \param expr The input expression.
+ * \return underlying Relax type.
+ */
+inline Type GetType(const Expr& expr) {
+ TVM_FFI_ICHECK(expr->ty.defined())
+ << "The type is not populated, check if you have normalized the expr";
+ return expr->ty;
+}
+
+/*!
+ * \brief Whether the expr has void type.
+ *
+ * \param expr The input expression.
+ * \return Whether the expr has void type.
+ */
+inline bool HasVoidType(const Expr& expr) {
+ auto* ptr = expr->ty.as<TupleTypeNode>();
+ return ptr != nullptr && ptr->fields.size() == 0;
+}
+
+/*!
+ * \brief Update the type of an Expr.
+ * \param expr The Expr whose type to be updated.
+ * \param ty The type assigned.
+ * \note We ensure idempotence, that is we can only update the type of an Expr only
+ * if the original one is nullptr.
+ */
+TVM_DLL void UpdateType(Expr expr, Type ty);
} // namespace relax
} // namespace tvm
+
#endif // TVM_RELAX_TYPE_H_
diff --git a/include/tvm/ir/type_functor.h b/include/tvm/relax/type_functor.h
similarity index 55%
rename from include/tvm/ir/type_functor.h
rename to include/tvm/relax/type_functor.h
index a56690e817..4751316144 100644
--- a/include/tvm/ir/type_functor.h
+++ b/include/tvm/relax/type_functor.h
@@ -18,31 +18,31 @@
*/
/*!
- * \file tvm/ir/type_functor.h
- * \brief A way to defined arbitrary function signature with dispatch on types.
+ * \file tvm/relax/type_functor.h
+ * \brief Functors and visitors for Relax type nodes.
*/
-#ifndef TVM_IR_TYPE_FUNCTOR_H_
-#define TVM_IR_TYPE_FUNCTOR_H_
+#ifndef TVM_RELAX_TYPE_FUNCTOR_H_
+#define TVM_RELAX_TYPE_FUNCTOR_H_
#include <tvm/ir/node_functor.h>
-#include <tvm/ir/type.h>
+#include <tvm/relax/distributed/type.h>
+#include <tvm/relax/type.h>
-#include <string>
#include <utility>
-#include <vector>
namespace tvm {
+namespace relax {
template <typename FType>
class TypeFunctor;
// functions to be overriden.
-#define TYPE_FUNCTOR_DEFAULT \
+#define RELAX_TYPE_FUNCTOR_DEFAULT \
{ \
return VisitTypeDefault_(op, std::forward<Args>(args)...); \
}
-#define TVM_TYPE_FUNCTOR_DISPATCH(OP) \
+#define TVM_RELAX_TYPE_FUNCTOR_DISPATCH(OP) \
vtable.template set_dispatch<OP>([](const ffi::ObjectRef& n, TSelf* self, Args... args) { \
return self->VisitType_(static_cast<const OP*>(n.get()), std::forward<Args>(args)...); \
});
@@ -60,14 +60,14 @@ class TypeFunctor<R(const Type& n, Args...)> {
virtual ~TypeFunctor() {}
/*!
* \brief Same as call.
- * \param n The expression node.
+ * \param n The type node.
* \param args Additional arguments.
* \return The result of the call
*/
R operator()(const Type& n, Args... args) { return VisitType(n, std::forward<Args>(args)...); }
/*!
* \brief The functor call.
- * \param n The expression node.
+ * \param n The type node.
* \param args Additional arguments.
* \return The result of the call
*/
@@ -77,10 +77,14 @@ class TypeFunctor<R(const Type& n, Args...)> {
return vtable(n, this, std::forward<Args>(args)...);
}
// Functions that can be overriden by subclass
- virtual R VisitType_(const FuncTypeNode* op, Args... args) TYPE_FUNCTOR_DEFAULT;
- virtual R VisitType_(const TupleTypeNode* op, Args... args) TYPE_FUNCTOR_DEFAULT;
- virtual R VisitType_(const PrimTypeNode* op, Args... args) TYPE_FUNCTOR_DEFAULT;
- virtual R VisitType_(const PointerTypeNode* op, Args... args) TYPE_FUNCTOR_DEFAULT;
+ virtual R VisitType_(const ObjectTypeNode* op, Args... args) RELAX_TYPE_FUNCTOR_DEFAULT;
+ virtual R VisitType_(const PrimTypeNode* op, Args... args) RELAX_TYPE_FUNCTOR_DEFAULT;
+ virtual R VisitType_(const ShapeTypeNode* op, Args... args) RELAX_TYPE_FUNCTOR_DEFAULT;
+ virtual R VisitType_(const TensorTypeNode* op, Args... args) RELAX_TYPE_FUNCTOR_DEFAULT;
+ virtual R VisitType_(const distributed::DTensorTypeNode* op,
+ Args... args) RELAX_TYPE_FUNCTOR_DEFAULT;
+ virtual R VisitType_(const TupleTypeNode* op, Args... args) RELAX_TYPE_FUNCTOR_DEFAULT;
+ virtual R VisitType_(const FuncTypeNode* op, Args... args) RELAX_TYPE_FUNCTOR_DEFAULT;
virtual R VisitTypeDefault_(const ffi::Object* op, Args...) {
TVM_FFI_THROW(InternalError) << "Do not have a default for " << op->GetTypeKey();
throw; // unreachable, written to stop compiler warning
@@ -91,42 +95,58 @@ class TypeFunctor<R(const Type& n, Args...)> {
static FType InitVTable() {
FType vtable;
// Set dispatch
- TVM_TYPE_FUNCTOR_DISPATCH(FuncTypeNode);
- TVM_TYPE_FUNCTOR_DISPATCH(TupleTypeNode);
- TVM_TYPE_FUNCTOR_DISPATCH(PrimTypeNode);
- TVM_TYPE_FUNCTOR_DISPATCH(PointerTypeNode);
+ TVM_RELAX_TYPE_FUNCTOR_DISPATCH(ObjectTypeNode);
+ TVM_RELAX_TYPE_FUNCTOR_DISPATCH(PrimTypeNode);
+ TVM_RELAX_TYPE_FUNCTOR_DISPATCH(ShapeTypeNode);
+ TVM_RELAX_TYPE_FUNCTOR_DISPATCH(TensorTypeNode);
+ TVM_RELAX_TYPE_FUNCTOR_DISPATCH(distributed::DTensorTypeNode);
+ TVM_RELAX_TYPE_FUNCTOR_DISPATCH(TupleTypeNode);
+ TVM_RELAX_TYPE_FUNCTOR_DISPATCH(FuncTypeNode);
vtable.Finalize();
return vtable;
}
};
-#undef TVM_TYPE_FUNCTOR_DISPATCH
+#undef TVM_RELAX_TYPE_FUNCTOR_DISPATCH
/*!
- * \brief A type visitor that recursively visit types.
+ * \brief A type visitor.
*/
class TVM_DLL TypeVisitor : public TypeFunctor<void(const Type& n)> {
public:
- void VisitType_(const FuncTypeNode* op) override;
- void VisitType_(const TupleTypeNode* op) override;
+ void VisitType_(const ObjectTypeNode* op) override;
void VisitType_(const PrimTypeNode* op) override;
- void VisitType_(const PointerTypeNode* op) override;
+ void VisitType_(const ShapeTypeNode* op) override;
+ void VisitType_(const TensorTypeNode* op) override;
+ void VisitType_(const distributed::DTensorTypeNode* op) override;
+ void VisitType_(const TupleTypeNode* op) override;
+ void VisitType_(const FuncTypeNode* op) override;
+
+ protected:
+ // two functions to override when visit expr fields in type nodes.
+ virtual void VisitTypeExprField(const Expr& expr) {}
+ virtual void VisitTypeExprField(const PrimExpr& expr) {}
};
/*!
- * \brief TypeMutator that mutates expressions.
+ * \brief TypeMutator that mutates Relax type nodes.
*/
class TVM_DLL TypeMutator : public TypeFunctor<Type(const Type& n)> {
public:
- Type VisitType(const Type& t) override;
- Type VisitType_(const FuncTypeNode* op) override;
- Type VisitType_(const TupleTypeNode* op) override;
+ Type VisitType_(const ObjectTypeNode* op) override;
Type VisitType_(const PrimTypeNode* op) override;
- Type VisitType_(const PointerTypeNode* op) override;
+ Type VisitType_(const ShapeTypeNode* op) override;
+ Type VisitType_(const TensorTypeNode* op) override;
+ Type VisitType_(const distributed::DTensorTypeNode* op) override;
+ Type VisitType_(const TupleTypeNode* op) override;
+ Type VisitType_(const FuncTypeNode* op) override;
- private:
- ffi::Array<Type> MutateArray(ffi::Array<Type> arr);
+ protected:
+ // two functions to override when visit expr fields in type nodes.
+ virtual Expr VisitTypeExprField(const Expr& expr) { return expr; }
+ virtual PrimExpr VisitTypeExprField(const PrimExpr& expr) { return expr; }
};
+} // namespace relax
} // namespace tvm
-#endif // TVM_IR_TYPE_FUNCTOR_H_
+#endif // TVM_RELAX_TYPE_FUNCTOR_H_
diff --git a/include/tvm/relax/utils.h b/include/tvm/relax/utils.h
index 77f8bab555..f36da70f9c 100644
--- a/include/tvm/relax/utils.h
+++ b/include/tvm/relax/utils.h
@@ -51,11 +51,10 @@ TVM_DLL Expr Bind(const Expr& expr, const tvm::ffi::Map<Var, Expr>& binds,
const tvm::ffi::Map<tirx::Var, PrimExpr>& symbolic_var_map = {});
/*!
- * \brief Bind the symbolic variables to a StructInfo. This is a helper function usually called by
+ * \brief Bind the symbolic variables to a Type. This is a helper function usually called by
* other pass functions to help optimizations.
*/
-TVM_DLL StructInfo Bind(const StructInfo& sinfo,
- const tvm::ffi::Map<tirx::Var, PrimExpr>& symbolic_var_map);
+TVM_DLL Type Bind(const Type& ty, const tvm::ffi::Map<tirx::Var, PrimExpr>& symbolic_var_map);
/*!
* \brief Infer a binding map for symbolic variables
@@ -78,10 +77,10 @@ TVM_DLL tvm::ffi::Map<tirx::Var, PrimExpr> InferSymbolicVarMap(
const tvm::ffi::Map<relax::Var, relax::Expr>& binds, const arith::Analyzer& analyzer);
/*!
- * \brief Check if the given StructInfo is for a boolean scalar (tensor of rank 0 with a boolean
+ * \brief Check if the given Type is for a boolean scalar (tensor of rank 0 with a boolean
* dtype).
*
- * \param sinfo The input StructInfo.
+ * \param ty The input Type.
* \param permit_unknown_rank If true, it will permit the input type to have unknown rank
* (ndim of -1), which will require a dynamic check.
* \param permit_unknown_dtype If true, it will permit the input type to have an unknown dtype
@@ -90,8 +89,8 @@ TVM_DLL tvm::ffi::Map<tirx::Var, PrimExpr> InferSymbolicVarMap(
* \return True iff the input type is a boolean scalar type (or, depending on options, has unknown
* rank or dtype)
*/
-TVM_DLL bool IsBoolStructInfo(const StructInfo& sinfo, bool permit_unknown_rank = true,
- bool permit_unknown_dtype = true);
+TVM_DLL bool IsBoolType(const Type& ty, bool permit_unknown_rank = true,
+ bool permit_unknown_dtype = true);
/*!
* \brief Check if the given expression is a "leaf" node or tuple node for normalization purposes.
@@ -112,7 +111,7 @@ TVM_DLL bool IsLeafOrTuple(const Expr& expr);
/*!
* \brief Check if the given Call node is an impure operation. If the callee is a general
- * expression, this simply requires checking the purity field of the FuncStructInfo. If it is an Op,
+ * expression, this simply requires checking the purity field of the FuncType. If it is an Op,
* then this checks the `fPurity` field.
*
* \param call The input call
diff --git a/include/tvm/script/printer/config.h b/include/tvm/script/printer/config.h
index 541d66f635..5bba247b44 100644
--- a/include/tvm/script/printer/config.h
+++ b/include/tvm/script/printer/config.h
@@ -91,7 +91,7 @@ class PrinterConfigNode : public ffi::Object {
* Keys are conventionally namespaced as "<dialect>.<knob>", e.g.:
* "tirx.prefix" — the TIR prefix (default "T")
* "relax.prefix" — the Relax prefix (default "R")
- * "relax.show_all_struct_info" — whether to show all struct info (default true)
+ * "relax.show_all_ty" — whether to show all Relax type annotations (default true)
*
* Use GetExtraConfig<T>(key, fallback) to read values with a typed fallback.
*/
@@ -101,7 +101,7 @@ class PrinterConfigNode : public ffi::Object {
* \brief Look up a value in extra_config with type cast and fallback.
*
* Keys are conventionally namespaced as "<dialect>.<knob>"
- * (e.g. "tirx.prefix", "relax.show_all_struct_info").
+ * (e.g. "tirx.prefix", "relax.show_all_ty").
*/
template <typename T>
T GetExtraConfig(const ffi::String& key, T fallback) const {
diff --git a/python/tvm/contrib/cutlass/build.py b/python/tvm/contrib/cutlass/build.py
index 4ff3f0812a..aa07b7d43a 100644
--- a/python/tvm/contrib/cutlass/build.py
+++ b/python/tvm/contrib/cutlass/build.py
@@ -379,21 +379,21 @@ def _extract_relax_function_signature(f):
signature = {}
for i, arg in enumerate(f.params):
- sinfo = arg.struct_info
- if isinstance(sinfo, relax.TensorStructInfo):
- signature[f"arg{i}_shape"] = get_const_tuple(sinfo.shape)
- signature[f"arg{i}_dtype"] = sinfo.dtype
- elif isinstance(sinfo, relax.ShapeStructInfo):
- signature[f"arg{i}_shape"] = get_const_tuple(sinfo.values)
+ ty = arg.ty
+ if isinstance(ty, relax.TensorType):
+ signature[f"arg{i}_shape"] = get_const_tuple(ty.shape)
+ signature[f"arg{i}_dtype"] = ty.dtype
+ elif isinstance(ty, relax.ShapeType):
+ signature[f"arg{i}_shape"] = get_const_tuple(ty.values)
else:
raise NotImplementedError()
- ret_sinfo = f.ret_struct_info
- if ret_sinfo.shape is not None:
- signature["ret_shape"] = get_const_tuple(ret_sinfo.shape)
+ ret_ty = f.ret_ty
+ if ret_ty.shape is not None:
+ signature["ret_shape"] = get_const_tuple(ret_ty.shape)
else:
signature["ret_shape"] = None
- signature["ret_dtype"] = ret_sinfo.dtype
+ signature["ret_dtype"] = ret_ty.dtype
return signature
@@ -714,12 +714,12 @@ class CutlassRelaxFunctionAnnotator(relax.PyExprMutator):
if "stacked_attention" in op_type:
arg["arg0_dtype"] = signature["arg0_dtype"]
- q_shape = get_const_tuple(attention_node.args[0].struct_info.shape)
- k_shape = get_const_tuple(attention_node.args[1].struct_info.shape)
- v_shape = get_const_tuple(attention_node.args[2].struct_info.shape)
+ q_shape = get_const_tuple(attention_node.args[0].ty.shape)
+ k_shape = get_const_tuple(attention_node.args[1].ty.shape)
+ v_shape = get_const_tuple(attention_node.args[2].ty.shape)
if len(attention_node.args) == 4:
- arg["bias_shape"] = get_const_tuple(attention_node.args[3].struct_info.shape)
- arg["bias_dtype"] = attention_node.args[3].struct_info.dtype
+ arg["bias_shape"] = get_const_tuple(attention_node.args[3].ty.shape)
+ arg["bias_dtype"] = attention_node.args[3].ty.dtype
qkv_layout = "qkv_stacked"
else:
@@ -803,7 +803,7 @@ class CutlassRelaxFunctionAnnotator(relax.PyExprMutator):
def visit_function_(self, f):
if "Composite" not in f.attrs:
body = super().visit_expr(f.body)
- return relax.Function(f.params, body, f.ret_struct_info, f.is_pure, f.attrs, f.span)
+ return relax.Function(f.params, body, f.ret_ty, f.is_pure, f.attrs, f.span)
op_type = f.attrs["Composite"]
diff --git a/python/tvm/contrib/hexagon/generate_take_op.py b/python/tvm/contrib/hexagon/generate_take_op.py
index 080a7d6a19..372d719989 100644
--- a/python/tvm/contrib/hexagon/generate_take_op.py
+++ b/python/tvm/contrib/hexagon/generate_take_op.py
@@ -70,7 +70,7 @@ class Tanh2TakeReplace(tvm.relax.PyExprMutator):
var = call_node.args[0]
func = self.mod_[var]
- if call_node.args[1][0].struct_info.dtype == "uint8":
+ if call_node.args[1][0].ty.dtype == "uint8":
if op_replace(call_node, func):
inp, inp_scale, inp_zp, out_scale, out_zp = [x for x in call_node.args[1]]
# LUT node creation
@@ -78,7 +78,7 @@ class Tanh2TakeReplace(tvm.relax.PyExprMutator):
inp_scale, inp_zp, out_scale, out_zp, call_node.args[0].name_hint
)
# Take operation node creation
- take_func = hexagon_unary_ops.generate_take_primfunc(inp, call_node.struct_info)
+ take_func = hexagon_unary_ops.generate_take_primfunc(inp, call_node.ty)
take_func = take_func.without_attr("global_symbol")
take_func_gv = self.builder_.add_func(take_func, "take")
take_node = relax.call_tir(
@@ -86,7 +86,7 @@ class Tanh2TakeReplace(tvm.relax.PyExprMutator):
relax.expr.Tuple(
[call_node.args[1][0], relax.expr.Constant(tvm.runtime.tensor(LUT))]
),
- call_node.struct_info,
+ call_node.ty,
)
return take_node
return call_node
diff --git a/python/tvm/contrib/hexagon/hexagon_unary_ops.py b/python/tvm/contrib/hexagon/hexagon_unary_ops.py
index 92fdeed353..2a4aecaa02 100644
--- a/python/tvm/contrib/hexagon/hexagon_unary_ops.py
+++ b/python/tvm/contrib/hexagon/hexagon_unary_ops.py
@@ -83,16 +83,16 @@ def LUT_generation(inp_scale, inp_zp, out_scale, out_zp, op_name) -> None:
return LUT
-def generate_take_primfunc(inp, struct_info):
+def generate_take_primfunc(inp, ty):
# Generating the take op
- N, H, W, C = inp.struct_info.shape
- data = te.placeholder((N, H, W, C), dtype=struct_info.dtype, name="data")
+ N, H, W, C = inp.ty.shape
+ data = te.placeholder((N, H, W, C), dtype=ty.dtype, name="data")
LUT_func = te.placeholder((256,), dtype="uint8", name="LUT")
take = te.compute(
- struct_info.shape,
- lambda *indices: saturate(
- (LUT_func[data[indices].astype("uint8")]), struct_info.dtype
- ).astype(struct_info.dtype),
+ ty.shape,
+ lambda *indices: saturate((LUT_func[data[indices].astype("uint8")]), ty.dtype).astype(
+ ty.dtype
+ ),
name="take_op",
)
mod = te.create_prim_func([data, LUT_func, take])
diff --git a/python/tvm/ir/expr.py b/python/tvm/ir/expr.py
index c2107c9a8f..4fbebeddd0 100644
--- a/python/tvm/ir/expr.py
+++ b/python/tvm/ir/expr.py
@@ -17,7 +17,6 @@
"""Common expressions data structures in the IR."""
from numbers import Number
-from typing import Optional
import tvm_ffi
@@ -33,6 +32,7 @@ class BaseExpr(Node):
"""Base class of all the expressions."""
span: Span | None
+ ty: "tvm.ir.Type | None"
@tvm_ffi.register_object("ir.PrimExpr")
@@ -50,17 +50,6 @@ class PrimExpr(BaseExpr):
class RelaxExpr(BaseExpr):
"""Base class of all non-primitive expressions."""
- @property
- def struct_info(self) -> Optional["tvm.relax.StructInfo"]:
- """Get the struct info field
-
- Returns
- -------
- struct_info : tvm.relax.StructInfo
- The struct info if available.
- """
- return _ffi_api.ExprStructInfo(self)
-
@tvm_ffi.register_object("ir.GlobalVar")
class GlobalVar(RelaxExpr):
diff --git a/python/tvm/ir/global_info.py b/python/tvm/ir/global_info.py
index c6754e747a..a9198afac7 100644
--- a/python/tvm/ir/global_info.py
+++ b/python/tvm/ir/global_info.py
@@ -29,7 +29,7 @@ class GlobalInfo(Object):
"""Base node for all global info that can appear in the IR"""
def __eq__(self, other):
- """Compare two struct info for structural equivalence."""
+ """Compare two global info objects for structural equivalence."""
return tvm_ffi.structural_equal(self, other)
def __ne__(self, other):
diff --git a/python/tvm/ir/type.py b/python/tvm/ir/type.py
index 3ade4b80fc..567ebafa2d 100644
--- a/python/tvm/ir/type.py
+++ b/python/tvm/ir/type.py
@@ -81,8 +81,8 @@ class TupleType(Type):
The fields in the tuple
"""
- def __init__(self, fields):
- self.__init_handle_by_constructor__(_ffi_api.TupleType, fields)
+ def __init__(self, fields, span=None):
+ self.__init_handle_by_constructor__(_ffi_api.TupleType, fields, span)
@tvm_ffi.register_object("ir.FuncType")
diff --git a/python/tvm/relax/__init__.py b/python/tvm/relax/__init__.py
index 3784c6a63d..3eea8b0b02 100644
--- a/python/tvm/relax/__init__.py
+++ b/python/tvm/relax/__init__.py
@@ -51,9 +51,10 @@ from .expr import (
from .expr import const, extern, get_shape_of
# Type
-from .ty import (
+from .type import (
Type,
ObjectType,
+ PrimType,
ShapeType,
TensorType,
TupleType,
@@ -79,17 +80,6 @@ from .block_builder import BlockBuilder
# ExprFunctor
from .expr_functor import ExprFunctor, PyExprVisitor, PyExprMutator
-# StructInfo
-from .struct_info import (
- StructInfo,
- ObjectStructInfo,
- PrimStructInfo,
- ShapeStructInfo,
- TensorStructInfo,
- TupleStructInfo,
- FuncStructInfo,
-)
-
# pipeline
from .pipeline import get_default_pipeline
from .pipeline import get_pipeline
@@ -105,11 +95,11 @@ from .base_py_module import BasePyModule
from . import exec_builder
from . import expr
from . import ty
+from . import type
from . import analysis
from . import transform
from . import block_builder
from . import op
-from . import struct_info
from . import backend
from . import training
from . import distributed
diff --git a/python/tvm/relax/analysis/__init__.py b/python/tvm/relax/analysis/__init__.py
index fa466d93d1..d9f9fa72e6 100644
--- a/python/tvm/relax/analysis/__init__.py
+++ b/python/tvm/relax/analysis/__init__.py
@@ -26,9 +26,9 @@ from .analysis import (
collect_non_negative_expressions,
computable_at_compile_time,
contains_impure_call,
- definable_tir_vars_in_struct_info,
+ definable_tir_vars_in_type,
defined_symbolic_vars,
- derive_call_ret_struct_info,
+ derive_call_ret_type,
detect_recursion,
erase_to_well_defined,
free_symbolic_vars,
@@ -40,10 +40,10 @@ from .analysis import (
name_to_binding,
post_order_visit,
remove_all_unused,
- struct_info_base_check,
- struct_info_lca,
+ type_base_check,
+ type_lca,
suggest_layout_transforms,
- tir_vars_in_struct_info,
+ tir_vars_in_type,
udchain,
well_formed,
)
diff --git a/python/tvm/relax/analysis/analysis.py b/python/tvm/relax/analysis/analysis.py
index 5ffb2f5ec8..3abd47d521 100644
--- a/python/tvm/relax/analysis/analysis.py
+++ b/python/tvm/relax/analysis/analysis.py
@@ -26,44 +26,44 @@ from enum import IntEnum
import tvm
from tvm import IRModule, tirx
+from tvm.ir import Type
from tvm.relax.expr import Binding, Call, DataflowBlock, Expr, Function, GlobalVar, Var
-from tvm.relax.struct_info import FuncStructInfo, StructInfo
-from tvm.relax.ty import Type
+from tvm.relax.type import FuncType
from tvm.tirx import Buffer, IndexMap, PrimFunc, SBlock
from . import _ffi_api
-def get_static_type(sinfo: StructInfo) -> Type:
- """Get the corresponding static type from a StructInfo.
+def get_static_type(ty: Type) -> Type:
+ """Get the corresponding static type from a Type.
Parameters
----------
- sinfo : StructInfo
- The input struct info.
+ ty : Type
+ The input type.
Returns
-------
ret : Type
The corresponding static type.
"""
- return _ffi_api.GetStaticType(sinfo) # type: ignore
+ return _ffi_api.GetStaticType(ty) # type: ignore
def erase_to_well_defined(
- sinfo: StructInfo,
+ ty: Type,
shape_var_map: dict[tirx.Var, tirx.PrimExpr] | None = None,
var_map: dict[Var, Expr] | None = None,
-) -> StructInfo:
- """Erase sinfo into a well defined form.
+) -> Type:
+ """Erase ty into a well defined form.
- This function removes the StructInfo's dependencies on shape and vars that
+ This function removes the Type's dependencies on shape and vars that
are not defined in given maps.
Parameters
----------
- sinfo : StructInfo
- The input struct info.
+ ty : Type
+ The input type.
shape_var_map : Dict[tirx.Var, tirx.PrimExpr]
Specifies the defined shape vars and the values they should map to.
@@ -73,13 +73,13 @@ def erase_to_well_defined(
Returns
-------
- ret : StructInfo
- The corresponding erased struct info.
+ ret : Type
+ The corresponding erased type.
"""
shape_var_map = {} if shape_var_map is None else shape_var_map
var_map = {} if var_map is None else var_map
- return _ffi_api.EraseToWellDefined(sinfo, shape_var_map, var_map) # type: ignore
+ return _ffi_api.EraseToWellDefined(ty, shape_var_map, var_map) # type: ignore
class BaseCheckResult(IntEnum):
@@ -100,33 +100,31 @@ class BaseCheckResult(IntEnum):
PASS = 3
-def struct_info_base_check(base: StructInfo, derived: StructInfo) -> BaseCheckResult:
+def type_base_check(base: Type, derived: Type) -> BaseCheckResult:
"""Run a base check to see if base subsumes derived.
Parameters
----------
- base: StructInfo
- The base struct info.
+ base: Type
+ The base type.
- derived: StructInfo
- The derived struct info.
+ derived: Type
+ The derived type.
Returns
-------
- ret : StructInfo
- The derived return value struct info.
+ ret : Type
+ The derived return value type.
"""
- return _ffi_api.StructInfoBaseCheck(base, derived) # type: ignore
+ return _ffi_api.TypeBaseCheck(base, derived) # type: ignore
-def derive_call_ret_struct_info(
- func_sinfo: FuncStructInfo, call: Call, ctx: "tvm.relax.BlockBuilder"
-) -> StructInfo:
- """Derive the call's ret value struct info from inputs.
+def derive_call_ret_type(func_ty: FuncType, call: Call, ctx: "tvm.relax.BlockBuilder") -> Type:
+ """Derive the call's ret value type from inputs.
Parameters
----------
- func_sinfo: FuncStructInfo
+ func_ty: FuncType
The call's function signature.
call: Call
@@ -137,96 +135,96 @@ def derive_call_ret_struct_info(
Returns
-------
- ret : StructInfo
- The derived return value struct info.
+ ret : Type
+ The derived return value type.
Note
----
This is an internal derivation function, call.op field is
- ignored in this case and the derivation only depends on func_sinfo.
+ ignored in this case and the derivation only depends on func_ty.
"""
- return _ffi_api.DeriveCallRetStructInfo(func_sinfo, call, ctx) # type: ignore
+ return _ffi_api.DeriveCallRetType(func_ty, call, ctx) # type: ignore
-def struct_info_lca(lhs: StructInfo, rhs: StructInfo) -> StructInfo:
- """Unify the two struct info to their least common ancestor.
+def type_lca(lhs: Type, rhs: Type) -> Type:
+ """Unify the two type to their least common ancestor.
Parameters
----------
- lhs: StructInfo
+ lhs: Type
The left operand.
- rhs: StructInfo
+ rhs: Type
The right operand.
Returns
-------
- ret : StructInfo
+ ret : Type
The corresponding lca result.
"""
- return _ffi_api.StructInfoLCA(lhs, rhs) # type: ignore
+ return _ffi_api.TypeLCA(lhs, rhs) # type: ignore
-def tir_vars_in_struct_info(sinfo: StructInfo) -> list[tirx.Var]:
- """Get the TIR variables that appear in the input struct info.
+def tir_vars_in_type(ty: Type) -> list[tirx.Var]:
+ """Get the TIR variables that appear in the input type.
The returned list is deduplicated - each TIR variable will appear at most once.
Parameters
----------
- sinfo : StructInfo
- The struct info object to be analyzed.
+ ty : Type
+ The type object to be analyzed.
Returns
-------
ret : List[tirx.Var]
- The list of TIR variables that appear in the input struct info.
+ The list of TIR variables that appear in the input type.
"""
- return _ffi_api.TIRVarsInStructInfo(sinfo) # type: ignore
+ return _ffi_api.TIRVarsInType(ty) # type: ignore
-def definable_tir_vars_in_struct_info(sinfo: StructInfo) -> list[tirx.Var]:
- """Get the TIR variables that may be defined from input struct info.
+def definable_tir_vars_in_type(ty: Type) -> list[tirx.Var]:
+ """Get the TIR variables that may be defined from input type.
The returned list is deduplicated - each TIR variable will appear at most once.
Parameters
----------
- sinfo : StructInfo
- The struct info object to be analyzed.
+ ty : Type
+ The type object to be analyzed.
Returns
-------
ret : List[tirx.Var]
- The list of TIR variables that can be defined from the StructInfo
+ The list of TIR variables that can be defined from the Type
"""
- return _ffi_api.DefinableTIRVarsInStructInfo(sinfo) # type: ignore
+ return _ffi_api.DefinableTIRVarsInType(ty) # type: ignore
-def collect_non_negative_expressions(sinfo: StructInfo) -> list[tirx.PrimExpr]:
+def collect_non_negative_expressions(ty: Type) -> list[tirx.PrimExpr]:
"""Collect TIR expressions used in non-negative contexts
Get TIR variables that are non-negative within the context where
- the struct info is used. For example, any expression used as a
+ the type is used. For example, any expression used as a
tensor shape.
The returned list is deduplicated - each TIR expression will
appear at most once. The order of the list is in the order of
- occurrence within the struct info.
+ occurrence within the type.
Parameters
----------
- sinfo : StructInfo
- The struct info object to be analyzed.
+ ty : Type
+ The type object to be analyzed.
Returns
-------
ret : List[tirx.Var]
- The list of TIR variables that can be defined from the StructInfo
+ The list of TIR variables that can be defined from the Type
"""
- return _ffi_api.CollectNonNegativeExpressions(sinfo) # type: ignore
+ return _ffi_api.CollectNonNegativeExpressions(ty) # type: ignore
def defined_symbolic_vars(func: Function) -> list[Var]:
@@ -413,7 +411,7 @@ def contains_impure_call(expr: Expr, own_name: Var | GlobalVar | None = None) ->
Notes
-----
- Relies on StructInfo annotations, so ensure that the module has been normalized first.
+ Relies on Type annotations, so ensure that the module has been normalized first.
Also, an impure call in a *nested* function does *not* mean that the outer expression contains
an impure call--it only does if the nested function is *later called*.
"""
@@ -481,7 +479,7 @@ def remove_all_unused(func: Function) -> Function:
return _ffi_api.remove_all_unused(func) # type: ignore
-def well_formed(obj: IRModule | Function, check_struct_info: bool = True) -> None:
+def well_formed(obj: IRModule | Function, check_ty: bool = True) -> None:
"""Check if the IRModule is well formed, raising on the first violation.
Raises an error (seeded with the offending node so a pass runner can report a
@@ -493,20 +491,20 @@ def well_formed(obj: IRModule | Function, check_struct_info: bool = True) -> Non
obj : Union[tvm.IRModule, Function]
The input IRModule or relax.Function.
- check_struct_info : bool
+ check_ty : bool
A boolean flag indicating if the property "every Expr must
- have defined structure info" will be checked.
+ have defined type information" will be checked.
Note
----
- By default the structure info is always checked. It is only in test cases
- where `check_struct_info` might be false, so that other well-formed requirements
- will be well tested and will not be blocked by not having structure info.
+ By default the type information is always checked. It is only in test cases
+ where `check_ty` might be false, so that other well-formed requirements
+ will be well tested and will not be blocked by not having type information.
"""
- _ffi_api.well_formed(obj, check_struct_info) # type: ignore
+ _ffi_api.well_formed(obj, check_ty) # type: ignore
-def check_well_formed(obj: IRModule | Function, check_struct_info: bool = True) -> bool:
+def check_well_formed(obj: IRModule | Function, check_ty: bool = True) -> bool:
"""Return whether the IRModule or Function is well formed.
Wraps :func:`well_formed`, returning False instead of raising on the first violation.
@@ -516,16 +514,16 @@ def check_well_formed(obj: IRModule | Function, check_struct_info: bool = True)
obj : Union[tvm.IRModule, Function]
The input IRModule or relax.Function.
- check_struct_info : bool
+ check_ty : bool
A boolean flag indicating if the property "every Expr must
- have defined structure info" will be checked.
+ have defined type information" will be checked.
Returns
-------
ret: bool
True if the IRModule is well formed, False if not.
"""
- return _ffi_api.check_well_formed(obj, check_struct_info) # type: ignore
+ return _ffi_api.check_well_formed(obj, check_ty) # type: ignore
def _get_prim_func_default_dtype(func: PrimFunc):
diff --git a/python/tvm/relax/backend/adreno/clml.py b/python/tvm/relax/backend/adreno/clml.py
index 2b3c78d6a6..8797fe2fa9 100644
--- a/python/tvm/relax/backend/adreno/clml.py
+++ b/python/tvm/relax/backend/adreno/clml.py
@@ -64,7 +64,7 @@ class AppendReshapeToBNRewriter(PyExprMutator):
bn_call = self.bn_vars[tuple_value]
if op.index == 0:
bn_out = relax.TupleGetItem(bn_call, 0)
- input_shape = bn_call.args[0].struct_info.shape
+ input_shape = bn_call.args[0].ty.shape
return relax.Call(reshape_op, [bn_out, input_shape])
return super().visit_tuple_getitem_(op)
@@ -137,13 +137,13 @@ def clml_pattern_table():
if "data" in context.annotated_expr:
input_expr = context.annotated_expr["data"]
- input_dtype = input_expr.struct_info.dtype
+ input_dtype = input_expr.ty.dtype
if input_dtype not in ["float32", "float16"]:
return False
if "weight" in context.annotated_expr:
weight_expr = context.annotated_expr["weight"]
- weight_dtype = weight_expr.struct_info.dtype
+ weight_dtype = weight_expr.ty.dtype
if weight_dtype not in ["float32", "float16"]:
return False
@@ -244,7 +244,7 @@ def clml_pattern_table():
return False
data = context.annotated_expr["data"]
- input_shape = data.struct_info.shape
+ input_shape = data.ty.shape
if len(input_shape) != 4:
return False
@@ -311,7 +311,7 @@ def clml_pattern_table():
return False
data = context.annotated_expr["data"]
- input_shape = data.struct_info.shape
+ input_shape = data.ty.shape
if len(input_shape) != 4:
return False
@@ -371,7 +371,7 @@ def clml_pattern_table():
return False
data = context.annotated_expr["data"]
- input_shape = data.struct_info.shape
+ input_shape = data.ty.shape
if len(input_shape) != 4:
return False
@@ -451,8 +451,8 @@ def clml_pattern_table():
base_shape = None
for param in params.values():
- shape = param.struct_info.shape
- dtype = param.struct_info.dtype
+ shape = param.ty.shape
+ dtype = param.ty.dtype
if dtype not in {"float32"}:
return False
@@ -496,8 +496,8 @@ def clml_pattern_table():
def _check_binary_op(context: PatternCheckContext) -> bool:
def _check_arg(input_expr):
- input_dtype = input_expr.struct_info.dtype
- input_shape = input_expr.struct_info.shape
+ input_dtype = input_expr.ty.dtype
+ input_shape = input_expr.ty.shape
if len(input_shape) == 0:
return False
@@ -523,13 +523,13 @@ def clml_pattern_table():
rhs_shape = None
if "lhs" in context.annotated_expr:
lhs = context.annotated_expr["lhs"]
- lhs_shape = lhs.struct_info.shape
+ lhs_shape = lhs.ty.shape
if not _check_arg(lhs):
return False
if "rhs" in context.annotated_expr:
rhs = context.annotated_expr["rhs"]
- rhs_shape = rhs.struct_info.shape
+ rhs_shape = rhs.ty.shape
if not _check_arg(rhs):
return False
@@ -630,7 +630,7 @@ def _check_dequantize_matmul(ctx: relax.transform.PatternCheckContext) -> bool:
wdq = ctx.annotated_expr["w_decoded"]
w_pack = ctx.annotated_expr["w_encoded"]
- if ctx.annotated_expr["lhs"].struct_info.dtype != "float16":
+ if ctx.annotated_expr["lhs"].ty.dtype != "float16":
return False
if not isinstance(wdq, relax.Call):
return False
@@ -639,17 +639,17 @@ def _check_dequantize_matmul(ctx: relax.transform.PatternCheckContext) -> bool:
return False
if not (
- (len(root.struct_info.shape) == 3)
- and isinstance(root.struct_info.shape[0], tirx.IntImm)
- and (root.struct_info.dtype == "float16")
- and (root.struct_info.shape[0] == 1)
+ (len(root.ty.shape) == 3)
+ and isinstance(root.ty.shape[0], tirx.IntImm)
+ and (root.ty.dtype == "float16")
+ and (root.ty.shape[0] == 1)
):
return False
if not (
- (len(wdq.struct_info.shape) == 2)
- and (w_pack.struct_info.shape[-1] == root.struct_info.shape[-1])
- and (wdq.struct_info.shape[-2] == _input.struct_info.shape[-1])
+ (len(wdq.ty.shape) == 2)
+ and (w_pack.ty.shape[-1] == root.ty.shape[-1])
+ and (wdq.ty.shape[-2] == _input.ty.shape[-1])
):
return False
diff --git a/python/tvm/relax/backend/contrib/example_npu/patterns.py b/python/tvm/relax/backend/contrib/example_npu/patterns.py
index 17d6656ef8..41ae185357 100644
--- a/python/tvm/relax/backend/contrib/example_npu/patterns.py
+++ b/python/tvm/relax/backend/contrib/example_npu/patterns.py
@@ -63,7 +63,7 @@ def _check_npu_memory_constraints(
Placeholder for NPU memory hierarchy constraint checking.
A real implementation would inspect the annotated expression's
- TensorStructInfo to verify the tensor fits within the NPU's
+ TensorType to verify the tensor fits within the NPU's
on-chip SRAM (L1) or compute memory (L2/CMX). Tensors that
exceed on-chip capacity require tiling before offload.
"""
diff --git a/python/tvm/relax/backend/cuda/cublas.py b/python/tvm/relax/backend/cuda/cublas.py
index a0907a7678..eef6c70ad9 100644
--- a/python/tvm/relax/backend/cuda/cublas.py
+++ b/python/tvm/relax/backend/cuda/cublas.py
@@ -59,20 +59,20 @@ def _check_matmul(context: PatternCheckContext) -> bool:
scale = context.annotated_expr["scale"]
zero_point = context.annotated_expr["zp"]
# Only scalar values for scale and zero_point are supported.
- if scale.struct_info.ndim != 0 or zero_point.struct_info.ndim != 0:
+ if scale.ty.ndim != 0 or zero_point.ty.ndim != 0:
return False
# Only zero_point == 0.0 is supported.
if zero_point.data.numpy()[()].item() != 0.0:
return False
- lhs_dtype = lhs.struct_info.dtype
- rhs_dtype = rhs.struct_info.dtype
- out_dtype = matmul_call.struct_info.dtype
+ lhs_dtype = lhs.ty.dtype
+ rhs_dtype = rhs.ty.dtype
+ out_dtype = matmul_call.ty.dtype
if not _is_supported_dtype(lhs_dtype, rhs_dtype, out_dtype):
return False
- lhs_shape = lhs.struct_info.shape.values
- rhs_shape = rhs.struct_info.shape.values
+ lhs_shape = lhs.ty.shape.values
+ rhs_shape = rhs.ty.shape.values
if not isinstance(lhs_shape[-1], tvm.tirx.expr.IntImm | int):
# Reduction axis must be constant
@@ -120,7 +120,7 @@ def _check_matmul(context: PatternCheckContext) -> bool:
# Non-default epilogue not supported for IGEMM
return False
bias = context.annotated_expr["bias"]
- bias_shape = bias.struct_info.shape.values
+ bias_shape = bias.ty.shape.values
bias_batches = reduce(operator.mul, bias_shape[:-1], 1)
if not isinstance(bias_batches, tvm.tirx.expr.IntImm | int) or int(bias_batches) > 1:
# cuBLAS only supports bias vector
diff --git a/python/tvm/relax/backend/cuda/cudnn.py b/python/tvm/relax/backend/cuda/cudnn.py
index 2be11fd04d..2df31ca62f 100644
--- a/python/tvm/relax/backend/cuda/cudnn.py
+++ b/python/tvm/relax/backend/cuda/cudnn.py
@@ -53,8 +53,8 @@ def _check_conv2d(context: PatternCheckContext) -> bool:
weight_expr = context.annotated_expr["weight"]
# Check if the data types of input and weights are supported by cuDNN BYOC
- input_dtype = input_expr.struct_info.dtype
- weight_dtype = weight_expr.struct_info.dtype
+ input_dtype = input_expr.ty.dtype
+ weight_dtype = weight_expr.ty.dtype
if not _is_supported_dtype(input_dtype, weight_dtype):
return False
@@ -71,14 +71,14 @@ def _check_stacked_attention(context: PatternCheckContext, layout: str) -> bool:
if has_leaking_intermediate_variables(context):
return False
if layout == "BS3NH":
- if not context.annotated_expr["stacked_qkv"].struct_info.ndim == 3:
+ if not context.annotated_expr["stacked_qkv"].ty.ndim == 3:
return False
if "split" in context.annotated_expr:
split_op = context.annotated_expr["split"]
if not split_op.attrs.axis == 2:
return False
elif layout == "SBN3H":
- if not context.annotated_expr["stacked_qkv"].struct_info.ndim == 4:
+ if not context.annotated_expr["stacked_qkv"].ty.ndim == 4:
return False
if "split" in context.annotated_expr:
split_op = context.annotated_expr["split"]
@@ -167,7 +167,7 @@ class WorkspaceAnnotator(PyExprMutator):
def visit_function_(self, f):
if "Composite" not in f.attrs:
body = super().visit_expr(f.body)
- new_f = relax.Function(f.params, body, f.ret_struct_info, f.is_pure, f.attrs, f.span)
+ new_f = relax.Function(f.params, body, f.ret_ty, f.is_pure, f.attrs, f.span)
if "global_symbol" in f.attrs and "cudnn" in f.attrs["global_symbol"]:
composite_func = body.blocks[0].bindings[0].value
@@ -178,8 +178,8 @@ class WorkspaceAnnotator(PyExprMutator):
if "attention" in f.attrs["Composite"] and "cudnn" in f.attrs["Composite"]:
# Workspace is needed only for larger head sizes, but for simplicity we always allocate.
- out_dtype = f.ret_struct_info.dtype
- out_size_1d = _shape_1d(f.ret_struct_info.shape)
+ out_dtype = f.ret_ty.dtype
+ out_size_1d = _shape_1d(f.ret_ty.shape)
# This needs to be in sync with the actual value that the kernel expects.
workspace_size_bytes = out_size_1d * {"float16": 2, "float32": 4}[out_dtype]
if not isinstance(workspace_size_bytes, int | tvm.tirx.expr.IntImm):
diff --git a/python/tvm/relax/backend/cuda/cutlass.py b/python/tvm/relax/backend/cuda/cutlass.py
index b15b3b6985..3e6e936d8c 100644
--- a/python/tvm/relax/backend/cuda/cutlass.py
+++ b/python/tvm/relax/backend/cuda/cutlass.py
@@ -105,8 +105,8 @@ def _check_residual(root_call: Call, context: PatternCheckContext) -> bool:
# If residual depends on the result of the root call, this cannot be handled by cutlass.
return False
- shape1 = root_var.struct_info.shape
- shape2 = residual.struct_info.shape
+ shape1 = root_var.ty.shape
+ shape2 = residual.ty.shape
out_channel = shape1[-1]
if not _is_same_shape(shape1, shape2) and not _is_bias_like(shape2, out_channel):
@@ -127,7 +127,7 @@ def _check_conv2d(context: PatternCheckContext) -> bool:
if (
data_layout != "NHWC"
or kernel_layout != "OHWI"
- or not _is_supported_dtype(data.struct_info.dtype, weight.struct_info.dtype)
+ or not _is_supported_dtype(data.ty.dtype, weight.ty.dtype)
):
return False
@@ -135,13 +135,13 @@ def _check_conv2d(context: PatternCheckContext) -> bool:
return False
# Check if any dimensions are symbolic.
- for dim in data.struct_info.shape.values:
+ for dim in data.ty.shape.values:
if isinstance(dim, tvm.tirx.Var):
return False
# pylint: disable=invalid-name
- IC = data.struct_info.shape.values[3]
- OC = weight.struct_info.shape.values[0]
+ IC = data.ty.shape.values[3]
+ OC = weight.ty.shape.values[0]
# not depthwise conv2d
return not IC == OC == conv2d_call.attrs.groups
@@ -154,16 +154,16 @@ def _check_matmul(context: PatternCheckContext) -> bool:
lhs = context.annotated_expr["lhs"]
rhs = context.annotated_expr["rhs"]
- lhs_dtype = lhs.struct_info.dtype
- rhs_dtype = rhs.struct_info.dtype
+ lhs_dtype = lhs.ty.dtype
+ rhs_dtype = rhs.ty.dtype
if not _is_supported_dtype(lhs_dtype, rhs_dtype):
return False
if not _check_residual(context.annotated_expr["root"], context):
return False
- lhs_shape = lhs.struct_info.shape.values
- rhs_shape = rhs.struct_info.shape.values
+ lhs_shape = lhs.ty.shape.values
+ rhs_shape = rhs.ty.shape.values
return is_shape_valid_for_cutlass_matmul(lhs_shape, rhs_shape)
@@ -223,22 +223,22 @@ def _check_decode_matmul(ctx):
return False
# out_dtype = "float32" not supported unless matmul is followed by cast to fp16.
- if root.struct_info.dtype == "float32":
+ if root.ty.dtype == "float32":
return False
call_tir_decode = ctx.annotated_expr["w_decoded"]
if "decode" not in call_tir_decode.args[0].name_hint:
return False
- N = root.struct_info.shape[-1]
+ N = root.ty.shape[-1]
- if ctx.annotated_expr["lhs"].struct_info.dtype != "float16":
+ if ctx.annotated_expr["lhs"].ty.dtype != "float16":
return False
# weight needs to be packed to int8.
packed_weight = ctx.annotated_expr["w_encoded"]
- if packed_weight.struct_info.dtype != "int8":
+ if packed_weight.ty.dtype != "int8":
return False
# The kernel expects the weight to be preprocessed by this packed function.
@@ -251,16 +251,16 @@ def _check_decode_matmul(ctx):
scales = ctx.annotated_expr["scales"]
- if scales.struct_info.dtype != "float16":
+ if scales.ty.dtype != "float16":
return False
# scale shape needs to be (N,) or (1, N) or (K // group_size, N)
- if len(scales.struct_info.shape) > 2 or scales.struct_info.shape[-1] != N:
+ if len(scales.ty.shape) > 2 or scales.ty.shape[-1] != N:
return False
if "bias" in ctx.annotated_expr:
- out_shape = root.struct_info.shape
- bias_shape = ctx.annotated_expr["bias"].struct_info.shape
+ out_shape = root.ty.shape
+ bias_shape = ctx.annotated_expr["bias"].ty.shape
# bias shape needs to be (N,), possibly with additional axes on the front.
# It can also have the same shape as the output.
@@ -378,7 +378,7 @@ def _check_stacked_attention(context: PatternCheckContext) -> bool:
"""Check if the given stacked attention workload can be offloaded to CUTLASS."""
if has_leaking_intermediate_variables(context):
return False
- if not context.annotated_expr["stacked_qkv"].struct_info.ndim == 3:
+ if not context.annotated_expr["stacked_qkv"].ty.ndim == 3:
return False
if "split" in context.annotated_expr:
split_op = context.annotated_expr["split"]
@@ -458,7 +458,7 @@ def _check_layer_norm(context: PatternCheckContext) -> bool:
return False
axis = int(attrs.axes[0])
- rank = len(context.matched_expr.struct_info.shape)
+ rank = len(context.matched_expr.ty.shape)
if axis < 0:
axis += rank
@@ -536,7 +536,7 @@ class WorkspaceAnnotator(PyExprMutator):
def visit_function_(self, f):
if "Composite" not in f.attrs:
body = super().visit_expr(f.body)
- new_f = Function(f.params, body, f.ret_struct_info, f.is_pure, f.attrs, f.span)
+ new_f = Function(f.params, body, f.ret_ty, f.is_pure, f.attrs, f.span)
if "global_symbol" in f.attrs and "cutlass" in f.attrs["global_symbol"]:
composite_func = body.blocks[0].bindings[0].value
@@ -547,8 +547,8 @@ class WorkspaceAnnotator(PyExprMutator):
if "attention" in f.attrs["Composite"] and "cutlass" in f.attrs["Composite"]:
# Workspace is needed only for larger head sizes, but for simplicity we always allocate.
- out_dtype = f.ret_struct_info.dtype
- out_size_1d = _shape_1d(f.ret_struct_info.shape)
+ out_dtype = f.ret_ty.dtype
+ out_size_1d = _shape_1d(f.ret_ty.shape)
# This needs to be in sync with the actual value that the kernel expects.
workspace_size_bytes = out_size_1d * {"float16": 2, "float32": 4}[out_dtype]
if not isinstance(workspace_size_bytes, int | tvm.tirx.expr.IntImm):
diff --git a/python/tvm/relax/backend/dispatch_sampling.py b/python/tvm/relax/backend/dispatch_sampling.py
index 780fcd0cc0..a82b8fd837 100644
--- a/python/tvm/relax/backend/dispatch_sampling.py
+++ b/python/tvm/relax/backend/dispatch_sampling.py
@@ -41,7 +41,7 @@ class SamplingDispatcher(BackendDispatcher):
)
prob, uniform_sample, sample_indices = call.args
- tgt = self._get_target(call.struct_info)
+ tgt = self._get_target(call.ty)
dtype = call.attrs.dtype
_, prob_dtype = self.get_shape_dtype(prob)
sample_shape, sample_dtype = self.get_shape_dtype(uniform_sample)
@@ -63,7 +63,7 @@ class SamplingDispatcher(BackendDispatcher):
return relax.call_tir(
gv,
[prob, uniform_sample, sample_indices],
- out_sinfo=call.struct_info,
+ out_ty=call.ty,
)
else:
cumsum_prob = relax.op.cumsum(prob, axis=1, dtype=prob_dtype, exclusive=False)
@@ -74,7 +74,7 @@ class SamplingDispatcher(BackendDispatcher):
return relax.call_tir(
gv,
[cumsum_prob, uniform_sample, sample_indices],
- out_sinfo=call.struct_info,
+ out_ty=call.ty,
)
return super().visit_call_(call)
diff --git a/python/tvm/relax/backend/dispatch_sort_scan.py b/python/tvm/relax/backend/dispatch_sort_scan.py
index c901aceed3..0656f65a43 100644
--- a/python/tvm/relax/backend/dispatch_sort_scan.py
+++ b/python/tvm/relax/backend/dispatch_sort_scan.py
@@ -80,16 +80,16 @@ class SortScanDispatcher(BackendDispatcher):
input_tensor = call.args[0]
boundaries = call.args[1]
right = call.attrs.right
- tgt = self._get_target(call.struct_info)
+ tgt = self._get_target(call.ty)
te_func = topi.searchsorted
with tgt:
if self.is_gpu_target(tgt):
te_func = topi.gpu.searchsorted
return self.builder_.call_te(
- te_func, boundaries, input_tensor, right, input_tensor.struct_info.dtype
+ te_func, boundaries, input_tensor, right, input_tensor.ty.dtype
)
if call.op.name == "relax.sort":
- tgt = self._get_target(call.struct_info)
+ tgt = self._get_target(call.ty)
te_func = topi.sort
kwargs = {}
with tgt:
@@ -102,7 +102,7 @@ class SortScanDispatcher(BackendDispatcher):
te_func, call.args[0], call.attrs.axis, not call.attrs.descending, **kwargs
)
if call.op.name == "relax.argsort":
- tgt = self._get_target(call.struct_info)
+ tgt = self._get_target(call.ty)
te_func = topi.argsort
kwargs = {}
with tgt:
@@ -120,7 +120,7 @@ class SortScanDispatcher(BackendDispatcher):
**kwargs,
)
if call.op.name == "relax.topk":
- tgt = self._get_target(call.struct_info)
+ tgt = self._get_target(call.ty)
te_func = topi.topk
kwargs = {}
if can_use_thrust(tgt, "tvm.contrib.thrust.sort"):
@@ -141,9 +141,9 @@ class SortScanDispatcher(BackendDispatcher):
self._append_calls_to_update(tir_call, tgt)
return tir_call
if call.op.name in ("relax.cumprod", "relax.cumsum"):
- tgt = self._get_target(call.struct_info)
+ tgt = self._get_target(call.ty)
axis = int(call.attrs.axis) if call.attrs.axis is not None else call.attrs.axis
- shape = call.struct_info.shape
+ shape = call.ty.shape
# TODO(tvm-team): Support fully dynamic case with `shape=None`
if shape is None:
raise ValueError("non-symbolic shape is not supported for now")
@@ -163,7 +163,7 @@ class SortScanDispatcher(BackendDispatcher):
dim = 1
for i in range(len(shape) - 1):
dim *= shape[i]
- in_dtype = call.args[0].struct_info.dtype
+ in_dtype = call.args[0].ty.dtype
out_dtype = call.attrs.dtype
out_dtype = out_dtype or in_dtype
cumsum_2d_shape = relax.ShapeExpr([dim, shape[-1]])
@@ -171,7 +171,7 @@ class SortScanDispatcher(BackendDispatcher):
"vm.builtin.reshape",
call.args[0],
cumsum_2d_shape,
- sinfo_args=relax.TensorStructInfo(cumsum_2d_shape, out_dtype),
+ ty_args=relax.TensorType(cumsum_2d_shape, out_dtype),
)
gv = self.builder_.add_func(
gpu_2d_continuous_cumsum(in_dtype=in_dtype, out_dtype=out_dtype),
@@ -180,13 +180,13 @@ class SortScanDispatcher(BackendDispatcher):
cumsum = relax.call_tir(
gv,
reshape,
- out_sinfo=relax.TensorStructInfo(cumsum_2d_shape, out_dtype),
+ out_ty=relax.TensorType(cumsum_2d_shape, out_dtype),
)
return relax.call_pure_packed(
"vm.builtin.reshape",
cumsum,
shape,
- sinfo_args=call.struct_info,
+ ty_args=call.ty,
)
with tgt:
@@ -214,8 +214,8 @@ class SortScanDispatcher(BackendDispatcher):
"""
Estimate the workspace size for thrust sort/argsort/topk/cumsum
"""
- input_shape = call.args[0].struct_info.shape
- input_byte_per_elem = DataType(call.args[0].struct_info.dtype).bits // 8
+ input_shape = call.args[0].ty.shape
+ input_byte_per_elem = DataType(call.args[0].ty.dtype).bits // 8
int64_byte_per_elem = DataType("int64").bits // 8
int32_byte_per_elem = DataType("int32").bits // 8
num_elem = reduce(mul, input_shape, 1)
diff --git a/python/tvm/relax/backend/metal/coreml.py b/python/tvm/relax/backend/metal/coreml.py
index dab375478a..d0b7ea3fc8 100644
--- a/python/tvm/relax/backend/metal/coreml.py
+++ b/python/tvm/relax/backend/metal/coreml.py
@@ -36,8 +36,8 @@ from tvm.relax.expr import (
Var,
VarBinding,
)
-from tvm.relax.struct_info import PrimStructInfo, TensorStructInfo
from tvm.relax.transform import PatternCheckContext
+from tvm.relax.type import PrimType, TensorType
from tvm.support.xcode import compile_coreml
from ...expr_functor import PyExprVisitor, visitor
@@ -355,14 +355,14 @@ class CodegenCoreML(PyExprVisitor):
def visit_function_(self, op) -> None:
for var in op.params:
name = var.name_hint
- sinfo = var.struct_info
- if isinstance(sinfo, TensorStructInfo):
- shape = [int(v) for v in list(sinfo.shape)]
- elif isinstance(sinfo, PrimStructInfo):
+ ty = var.ty
+ if isinstance(ty, TensorType):
+ shape = [int(v) for v in list(ty.shape)]
+ elif isinstance(ty, PrimType):
shape = []
else:
- raise Exception("Currently not supported: ", type(sinfo))
- dtype = sinfo.dtype
+ raise Exception("Currently not supported: ", type(ty))
+ dtype = ty.dtype
self.model_inputs_.append((name, shape, dtype))
self.visit_expr(op.body)
@@ -456,12 +456,12 @@ class CodegenCoreML(PyExprVisitor):
input_desc = self.builder.spec.description.input
input_desc[i].type.multiArrayType.dataType = FEATURE_TYPE_MAP[dtype]
- output_dim = [int(n) for n in self.function.struct_info.ret.shape]
+ output_dim = [int(n) for n in self.function.ty.ret.shape]
last_binding_var = self.function.body.blocks[0].bindings[-1].var
self.builder.set_output(self.out_map[last_binding_var], [output_dim])
- for i, dtype in enumerate([self.function.struct_info.ret.dtype]):
+ for i, dtype in enumerate([self.function.ty.ret.dtype]):
assert dtype in FEATURE_TYPE_MAP
output_desc = self.builder.spec.description.output
output_desc[i].type.multiArrayType.dataType = FEATURE_TYPE_MAP[dtype]
diff --git a/python/tvm/relax/backend/patterns.py b/python/tvm/relax/backend/patterns.py
index a328684b7c..06011d6f6e 100644
--- a/python/tvm/relax/backend/patterns.py
+++ b/python/tvm/relax/backend/patterns.py
@@ -489,12 +489,12 @@ def make_attention_rewrite_pattern(
transposed = is_op("relax.permute_dims")(reshaped)
def rewriter(matchings, x):
- if matchings[tensor].struct_info.ndim != 4:
+ if matchings[tensor].ty.ndim != 4:
return None
if list(matchings[permuted].attrs.axes) != [0, 2, 1, 3]:
return None
- before_reshape = matchings[permuted].struct_info.shape.values
- after_reshape = matchings[shape].struct_info.values
+ before_reshape = matchings[permuted].ty.shape.values
+ after_reshape = matchings[shape].ty.values
if not (
len(before_reshape) == 4
and len(after_reshape) == 3
@@ -503,7 +503,7 @@ def make_attention_rewrite_pattern(
return None
if transpose and list(matchings[transposed].attrs.axes) != [0, 2, 1]:
return None
- return x, x.struct_info.shape
+ return x, x.ty.shape
if transpose:
return transposed, rewriter
@@ -514,11 +514,11 @@ def make_attention_rewrite_pattern(
transposed = is_op("relax.permute_dims")(tensor)
def rewriter(matchings, x):
- if matchings[tensor].struct_info.ndim != 3:
+ if matchings[tensor].ty.ndim != 3:
return None
if transpose and list(matchings[transposed].attrs.axes) != [0, 2, 1]:
return None
- before_reshape = x.struct_info.shape.values
+ before_reshape = x.ty.shape.values
after_reshape = [before_reshape[0], before_reshape[1], 1, before_reshape[2]]
return R.reshape(x, after_reshape), after_reshape
@@ -536,10 +536,10 @@ def make_attention_rewrite_pattern(
permuted = is_op("relax.permute_dims")(reshaped)
def rewriter(matchings, x):
- if matchings[tensor].struct_info.ndim != 3:
+ if matchings[tensor].ty.ndim != 3:
return None
- before_reshape = matchings[tensor].struct_info.shape.values
- after_reshape = matchings[shape].struct_info.values
+ before_reshape = matchings[tensor].ty.shape.values
+ after_reshape = matchings[shape].ty.values
if not (
len(before_reshape) == 3
and len(after_reshape) == 4
@@ -554,9 +554,9 @@ def make_attention_rewrite_pattern(
elif layout == "BSH":
def rewriter(matchings, x):
- if matchings[tensor].struct_info.ndim != 3:
+ if matchings[tensor].ty.ndim != 3:
return None
- return R.reshape(x, matchings[tensor].struct_info.shape.values)
+ return R.reshape(x, matchings[tensor].ty.shape.values)
return tensor, rewriter
else:
@@ -602,7 +602,7 @@ def make_attention_rewrite_pattern(
if query is None or key is None or value is None:
return original
softmax_axis = matchings[softmax].attrs.axis
- softmax_input_rank = len(matchings[softmax].struct_info.shape)
+ softmax_input_rank = len(matchings[softmax].ty.shape)
if softmax_axis == -1:
softmax_axis += softmax_input_rank
if softmax_axis != softmax_input_rank - 1:
@@ -611,7 +611,7 @@ def make_attention_rewrite_pattern(
_, s_kv, _, _ = key_shape
if with_bias:
bias = matchings[bias_raw]
- bias_shape = list(bias.struct_info.shape)
+ bias_shape = list(bias.ty.shape)
if bias_shape == [b * n, s, s_kv]:
bias = R.reshape(bias, [b, n, s, s_kv])
elif bias_shape == [b * n, 1, s_kv]:
diff --git a/python/tvm/relax/backend/rocm/hipblas.py b/python/tvm/relax/backend/rocm/hipblas.py
index 15eaf80ab0..1a52102b48 100644
--- a/python/tvm/relax/backend/rocm/hipblas.py
+++ b/python/tvm/relax/backend/rocm/hipblas.py
@@ -47,14 +47,14 @@ def _check_matmul(context: PatternCheckContext) -> bool:
rhs = context.annotated_expr["rhs"]
matmul_call = context.annotated_expr["root"]
- lhs_dtype = lhs.struct_info.dtype
- rhs_dtype = rhs.struct_info.dtype
- out_dtype = matmul_call.struct_info.dtype
+ lhs_dtype = lhs.ty.dtype
+ rhs_dtype = rhs.ty.dtype
+ out_dtype = matmul_call.ty.dtype
if not _is_supported_dtype(lhs_dtype, rhs_dtype, out_dtype):
return False
- lhs_shape = lhs.struct_info.shape.values
- rhs_shape = rhs.struct_info.shape.values
+ lhs_shape = lhs.ty.shape.values
+ rhs_shape = rhs.ty.shape.values
if not isinstance(lhs_shape[-1], tvm.tirx.expr.IntImm | int):
# Reduction axis must be constant
@@ -73,7 +73,7 @@ def _check_matmul(context: PatternCheckContext) -> bool:
# Non-default epilogue not supported for IGEMM
return False
bias = context.annotated_expr["bias"]
- bias_shape = bias.struct_info.shape.values
+ bias_shape = bias.ty.shape.values
bias_batches = reduce(operator.mul, bias_shape[:-1], 1)
if not isinstance(bias_batches, tvm.tirx.expr.IntImm | int) or int(bias_batches) > 1:
# hipblas only supports bias vector
diff --git a/python/tvm/relax/backend/utils.py b/python/tvm/relax/backend/utils.py
index 9d4598e669..47c01ac431 100644
--- a/python/tvm/relax/backend/utils.py
+++ b/python/tvm/relax/backend/utils.py
@@ -38,13 +38,11 @@ class BackendDispatcher(PyExprMutator):
def get_shape_dtype(expr: relax.Expr) -> tuple[relax.ShapeExpr, str]:
"""Get shape and dtype from an expression.
If the shape and dtype is unknown, raise an error."""
- sinfo = expr.struct_info
- if not isinstance(expr.struct_info, relax.TensorStructInfo):
- raise ValueError(
- f"Expecting a expr with TensorStructInfo, but got {expr} with {expr.struct_info}"
- )
+ ty = expr.ty
+ if not isinstance(expr.ty, relax.TensorType):
+ raise ValueError(f"Expecting a expr with TensorType, but got {expr} with {expr.ty}")
- shape, dtype = sinfo.shape, sinfo.dtype
+ shape, dtype = ty.shape, ty.dtype
if shape is None:
raise ValueError(
f"Expecting a expr with known shape, but got {expr} with unknown shape"
@@ -52,14 +50,14 @@ class BackendDispatcher(PyExprMutator):
return shape, dtype
- def _get_target(self, sinfo: relax.StructInfo) -> Target:
- # Get target information from TensorStructInfo
- if isinstance(sinfo, relax.TensorStructInfo):
- vdevice = sinfo.vdevice
+ def _get_target(self, ty: relax.Type) -> Target:
+ # Get target information from TensorType
+ if isinstance(ty, relax.TensorType):
+ vdevice = ty.vdevice
if vdevice is not None:
return vdevice.target
- elif isinstance(sinfo, relax.TupleStructInfo):
- for f in sinfo.fields:
+ elif isinstance(ty, relax.TupleType):
+ for f in ty.fields:
tgt = self._get_target(f)
if tgt != Target.current():
return tgt
diff --git a/python/tvm/relax/base_py_module.py b/python/tvm/relax/base_py_module.py
index 5dd8a107ae..825b532364 100644
--- a/python/tvm/relax/base_py_module.py
+++ b/python/tvm/relax/base_py_module.py
@@ -218,7 +218,7 @@ class BasePyModule:
wrapped_func = create_py_func_wrapper(func_name, py_func)
register_py_func(func_name, wrapped_func)
- def call_tir(self, tir_func, args, out_sinfo):
+ def call_tir(self, tir_func, args, out_ty):
"""Call a TIR function with PyTorch tensors."""
# Try to get function name from different sources
if isinstance(tir_func, str):
@@ -244,7 +244,7 @@ class BasePyModule:
)
func = self.compiled_tir_funcs[func_name]
- out = self._create_output_tensors(out_sinfo, args)
+ out = self._create_output_tensors(out_ty, args)
tvm_args = self._convert_pytorch_to_tvm(args)
tvm_out = self._convert_pytorch_to_tvm(out)
@@ -253,7 +253,7 @@ class BasePyModule:
result = self._convert_tvm_to_pytorch(tvm_out)
return result[0] if len(result) == 1 else result
- def call_dps_packed(self, func_name: str, args, out_sinfo):
+ def call_dps_packed(self, func_name: str, args, out_ty):
"""Call a packed function with PyTorch tensors, converting TVM Tensors via DLPack."""
if hasattr(self, func_name) and callable(getattr(self, func_name)):
return getattr(self, func_name)(*args)
@@ -268,7 +268,7 @@ class BasePyModule:
) from error
func = self.extern_funcs[func_name]
- out = self._create_output_tensors(out_sinfo, args)
+ out = self._create_output_tensors(out_ty, args)
tvm_args = self._convert_pytorch_to_tvm(args)
tvm_out = self._convert_pytorch_to_tvm(out)
func(*tvm_args, *tvm_out)
@@ -281,22 +281,20 @@ class BasePyModule:
py_func = self.pyfuncs[func_name]
return py_func(self, *args)
- def _create_output_tensors(self, out_sinfo, in_args=None):
+ def _create_output_tensors(self, out_ty, in_args=None):
# pylint: disable=import-outside-toplevel
import torch
- sinfo_list = out_sinfo if isinstance(out_sinfo, list) else [out_sinfo]
+ ty_list = out_ty if isinstance(out_ty, list) else [out_ty]
out_tensors = []
- for sinfo in sinfo_list:
- if isinstance(sinfo, tuple | list) and all(
- isinstance(x, int | np.integer) for x in sinfo
- ):
- out_tensors.append(torch.zeros(list(map(int, sinfo)), dtype=torch.float32))
+ for ty in ty_list:
+ if isinstance(ty, tuple | list) and all(isinstance(x, int | np.integer) for x in ty):
+ out_tensors.append(torch.zeros(list(map(int, ty)), dtype=torch.float32))
continue
- if hasattr(sinfo, "shape") and hasattr(sinfo, "dtype"):
- concrete_shape = self._infer_concrete_shape_from_args(sinfo.shape, in_args)
- torch_dtype = self._convert_tvm_dtype_to_torch(sinfo.dtype)
+ if hasattr(ty, "shape") and hasattr(ty, "dtype"):
+ concrete_shape = self._infer_concrete_shape_from_args(ty.shape, in_args)
+ torch_dtype = self._convert_tvm_dtype_to_torch(ty.dtype)
out_tensors.append(torch.zeros(concrete_shape, dtype=torch_dtype))
continue
@@ -341,7 +339,7 @@ class BasePyModule:
raise ValueError(
"Cannot infer concrete output shape from symbolic shape and inputs. "
- "Please provide a concrete `out_sinfo` (e.g., a tuple/list of ints) "
+ "Please provide a concrete `out_ty` (e.g., a tuple/list of ints) "
"or ensure input tensors carry shapes that determine output extents."
)
@@ -510,7 +508,7 @@ class BasePyModule:
num_context_lines: int = -1,
syntax_sugar: bool = True,
show_object_address: bool = False,
- show_all_struct_info: bool = True,
+ show_all_ty: bool = True,
extra_config: dict | None = None,
) -> str:
"""Print TVM IR into TVMScript text format with Python function support.
@@ -532,7 +530,7 @@ class BasePyModule:
num_context_lines=num_context_lines,
syntax_sugar=syntax_sugar,
show_object_address=show_object_address,
- show_all_struct_info=show_all_struct_info,
+ show_all_ty=show_all_ty,
extra_config=extra_config,
)
diff --git a/python/tvm/relax/block_builder.py b/python/tvm/relax/block_builder.py
index f347f05f15..ef7f6d58f8 100644
--- a/python/tvm/relax/block_builder.py
+++ b/python/tvm/relax/block_builder.py
@@ -32,7 +32,7 @@ from tvm.runtime import Object
from . import _ffi_api
from .expr import BaseFunc, Binding, BindingBlock, Expr, GlobalVar, Tuple, Var
from .op.base import call_tir, call_tir_with_grad
-from .struct_info import StructInfo
+from .type import Type
from .utils import gen_call_tir_inputs
@@ -92,7 +92,7 @@ class TestingScope:
else:
raise ValueError("def_vars only can take tirx.Var")
# setup a dummy var so shape is in scope.
- sparam = rx.Var("sparam", rx.ShapeStructInfo(shape_vars))
+ sparam = rx.Var("sparam", rx.ShapeType(shape_vars))
self._scope_params = [sparam]
def __enter__(self):
@@ -114,8 +114,8 @@ class BlockBuilder(Object):
m = tirx.Var("m", "int32")
n = tirx.Var("n", "int32")
- x = rx.Var("x", rx.TensorStructInfo([m, n], "float16"))
- y = rx.Var("y", rx.TensorStructInfo([n], "float16")
+ x = rx.Var("x", rx.TensorType([m, n], "float16"))
+ y = rx.Var("y", rx.TensorType([n], "float16"))
bb = rx.BlockBuilder()
with bb.function([x, y], "func"):
with bb.dataflow() as df:
@@ -361,13 +361,13 @@ class BlockBuilder(Object):
"""
primfunc_name = kwargs.pop("primfunc_name_hint", None)
- tir_func, call_args, output_sinfo, tir_vars = gen_call_tir_inputs(func, *args, **kwargs)
+ tir_func, call_args, output_ty, tir_vars = gen_call_tir_inputs(func, *args, **kwargs)
if not primfunc_name:
primfunc_name = func.__name__
gvar = self.add_func(tir_func, primfunc_name)
- return call_tir(gvar, call_args, output_sinfo, tir_vars)
+ return call_tir(gvar, call_args, output_ty, tir_vars)
def call_te_with_grad(
self,
@@ -413,7 +413,7 @@ class BlockBuilder(Object):
"""
primfunc_name = kwargs.pop("primfunc_name_hint", None)
- tir_func, call_args, output_sinfo, tir_vars = gen_call_tir_inputs(func, *args, **kwargs)
+ tir_func, call_args, output_ty, tir_vars = gen_call_tir_inputs(func, *args, **kwargs)
if te_grad_kwargs is None:
te_grad_kwargs = {}
@@ -423,7 +423,7 @@ class BlockBuilder(Object):
gvar = self.add_func(tir_func, primfunc_name)
return call_tir_with_grad(
- gvar, call_args, output_sinfo, te_grad_name, te_grad_kwargs, tir_vars
+ gvar, call_args, output_ty, te_grad_name, te_grad_kwargs, tir_vars
)
def emit_te(self, func: Callable, *args: Any, **kwargs: Any) -> Var:
@@ -456,8 +456,8 @@ class BlockBuilder(Object):
bb = rx.BlockBuilder()
n, m = tirx.Var("n", "int64"), tirx.Var("m", "int64")
- x = rx.Var("x", rx.TensorStructInfo([n, m], "float32"))
- y = rx.Var("y", rx.TensorStructInfo([n, m], "float32"))
+ x = rx.Var("x", rx.TensorType([n, m], "float32"))
+ y = rx.Var("y", rx.TensorType([n, m], "float32"))
def te_func(args, args_dict, msg):
A = args[0]
@@ -506,8 +506,8 @@ class BlockBuilder(Object):
bb = relax.BlockBuilder()
n = tirx.Var("n", "int64")
- x = relax.Var("x", relax.TensorStructInfo([n], "float32"))
- y = relax.Var("y", relax.TensorStructInfo([n + 1], "float32"))
+ x = relax.Var("x", relax.TensorType([n], "float32"))
+ y = relax.Var("y", relax.TensorType([n + 1], "float32"))
def te_func(A):
C = te.compute((n + 1), lambda i: A[i])
@@ -547,7 +547,7 @@ class BlockBuilder(Object):
name_hint = kwargs.pop("name_hint", "")
return self.emit(self.call_te(func, *args, **kwargs), name_hint=name_hint)
- def match_cast(self, value: Expr, struct_info: StructInfo, name_hint: str = "") -> Var:
+ def match_cast(self, value: Expr, ty: Type, name_hint: str = "") -> Var:
"""Emit a MatchCast.
Parameters
@@ -555,8 +555,8 @@ class BlockBuilder(Object):
value : tvm.relax.Expr
The value of the MatchCast to be emitted.
- struct_info : StructInfo
- The struct info to be matched.
+ ty : Type
+ The type to be matched.
name_hint : str
The name of the match cast
@@ -569,7 +569,7 @@ class BlockBuilder(Object):
return _ffi_api.BlockBuilderEmitMatchCast(
self,
value,
- struct_info,
+ ty,
name_hint,
) # type: ignore
@@ -651,8 +651,8 @@ class BlockBuilder(Object):
finally:
self.end_scope()
- # do not specify ret_struct_info and let constructor deduce
- # from seqe.struct_info
+ # do not specify ret_ty and let constructor deduce
+ # from seqe.ty
func = rx.Function(self._func._params, seqe, is_pure=self._func._is_pure)
for key, value in self._func._attrs.items():
func = func.with_attr(key, value)
diff --git a/python/tvm/relax/distributed/__init__.py b/python/tvm/relax/distributed/__init__.py
index 7996b2a25b..aa7876fd35 100644
--- a/python/tvm/relax/distributed/__init__.py
+++ b/python/tvm/relax/distributed/__init__.py
@@ -19,6 +19,6 @@
"""The infrastructure for distributed inference on Relax."""
from .global_info import DeviceMesh, device_mesh
-from .struct_info import Placement, DTensorStructInfo, PlacementSpec
+from .type import Placement, DTensorType, PlacementSpec
from . import transform
diff --git a/python/tvm/relax/distributed/struct_info.py b/python/tvm/relax/distributed/type.py
similarity index 88%
rename from python/tvm/relax/distributed/struct_info.py
rename to python/tvm/relax/distributed/type.py
index 0c94d3dda9..01fb8385a9 100644
--- a/python/tvm/relax/distributed/struct_info.py
+++ b/python/tvm/relax/distributed/type.py
@@ -15,14 +15,14 @@
# specific language governing permissions and limitations
# under the License.
# pylint: disable=redefined-builtin, invalid-name
-"""Struct Info for distributed tensor."""
+"""Types for distributed tensor."""
import enum
import tvm_ffi
from tvm.ir import Span
-from tvm.relax.struct_info import StructInfo, TensorStructInfo
+from tvm.relax.type import TensorType, Type
from tvm.runtime import Object
from . import _ffi_api
@@ -111,14 +111,14 @@ class Placement(Object):
return _ffi_api.PlacementFromText(text)
-@tvm_ffi.register_object("relax.DTensorStructInfo")
-class DTensorStructInfo(StructInfo):
- """StructInfo of a Distributed Tensor value.
+@tvm_ffi.register_object("relax.DTensorType")
+class DTensorType(Type):
+ """Type of a Distributed Tensor value.
Parameters
----------
- tensor_sinfo: TensorStructInfo
- The struct info inherited from TensorStructInfo
+ tensor_ty: TensorType
+ The tensor type carried by the distributed tensor.
device_mesh: DeviceMesh
The device mesh of the tensor.
placement: Placement
@@ -126,20 +126,20 @@ class DTensorStructInfo(StructInfo):
"""
- tensor_sinfo: TensorStructInfo
+ tensor_ty: TensorType
device_mesh: DeviceMesh
placement: Placement
def __init__(
self,
- tensor_sinfo: TensorStructInfo,
+ tensor_ty: TensorType,
device_mesh: DeviceMesh,
placement: Placement,
span: Span = None,
) -> None:
self.__init_handle_by_constructor__(
- _ffi_api.DTensorStructInfo,
- tensor_sinfo,
+ _ffi_api.DTensorType,
+ tensor_ty,
device_mesh,
placement,
span, # type: ignore
diff --git a/python/tvm/relax/dpl/pattern.py b/python/tvm/relax/dpl/pattern.py
index 89feac1bce..47b718981e 100644
--- a/python/tvm/relax/dpl/pattern.py
+++ b/python/tvm/relax/dpl/pattern.py
@@ -123,8 +123,8 @@ class DFPattern(Node):
attrs = make_node("ir.DictAttrs", **attrs)
return AttrPattern(self, attrs)
- def has_struct_info(self, struct_info: "StructInfo") -> "StructInfoPattern":
- return StructInfoPattern(self, struct_info)
+ def has_ty(self, ty: "Type") -> "TypePattern":
+ return TypePattern(self, ty)
def has_dtype(self, dtype: str) -> "DataTypePattern":
"""
@@ -568,23 +568,23 @@ class WildcardPattern(DFPattern):
@register_df_node
-class StructInfoPattern(DFPattern):
- """A pattern that matches another pattern with a certain StructInfo
+class TypePattern(DFPattern):
+ """A pattern that matches another pattern with a certain Type
Parameters
----------
pattern: tvm.relax.dpl.DFPattern
The input pattern that needs type annotation.
- struct_info: tvm.relax.StructInfo
- The struct info to match against
+ ty: tvm.relax.Type
+ The type to match against
"""
- def __init__(self, pattern: "DFPattern", struct_info: "StructInfo"):
+ def __init__(self, pattern: "DFPattern", ty: "Type"):
self.__init_handle_by_constructor__(
- ffi.StructInfoPattern,
+ ffi.TypePattern,
pattern,
- struct_info,
+ ty,
) # type: ignore
@@ -861,7 +861,7 @@ def is_shape(shape: list[tvm.ir.PrimExpr]) -> "PrimArrPattern":
return PrimArrPattern(shape)
-# Todo(relax-team): Dataflow pattern for StructInfo, and match out_sinfo
+# Todo(relax-team): Dataflow pattern for Type, and match out_ty
def _is_call_tir(
func_pattern: DFPattern,
args: list | tuple | TuplePattern = None,
@@ -877,7 +877,7 @@ def _is_call_tir(
return is_op("relax.call_tir")(func_pattern, args, tir_vars, add_constraint=False)
-# Todo(relax-team): Dataflow pattern for StructInfo, and match out_sinfo
+# Todo(relax-team): Dataflow pattern for Type, and match out_ty
def is_call_tir(
func_name: str,
args: list | tuple | TuplePattern = None,
diff --git a/python/tvm/relax/dpl/rewrite.py b/python/tvm/relax/dpl/rewrite.py
index 5cc928ff25..a0ac3665a3 100644
--- a/python/tvm/relax/dpl/rewrite.py
+++ b/python/tvm/relax/dpl/rewrite.py
@@ -274,7 +274,7 @@ def rewrite_bindings(
Q_weight = matchings[Q_weight_pat]
K_weight = matchings[K_weight_pat]
V_weight = matchings[V_weight_pat]
- width = Q_weight.struct_info.shape[1]
+ width = Q_weight.ty.shape[1]
concat = R.concat([Q_weight, K_weight, V_weight], axis=1)
matmul = R.matmul(inp, concat)
diff --git a/python/tvm/relax/expr.py b/python/tvm/relax/expr.py
index 734d7bde67..e02bbb51ca 100644
--- a/python/tvm/relax/expr.py
+++ b/python/tvm/relax/expr.py
@@ -57,41 +57,13 @@ class Id(Object):
raise RuntimeError("Cannot directly construct Id")
-# NOTE: place base struct info in expr to avoid cyclic dep
-# from expr to struct info.
-@tvm_ffi.register_object("ir.StructInfo")
-class StructInfo(Node, Scriptable):
- """The base class of all StructInfo.
-
- StructInfo contains both the static type
- and runtime structural information.
- """
-
- def __eq__(self, other):
- """Compare two struct info for structural equivalence."""
- return tvm_ffi.structural_equal(self, other)
+def _relax_type_is_base_of(self: Type, derived: Type) -> bool:
+ """Check if this Relax type is a base of another Relax type."""
- def __ne__(self, other):
- return not self.__eq__(other)
+ return _ffi_api.TypeIsBaseOf(self, derived) # type: ignore
- def same_as(self, other):
- """Overload with structural equality."""
- return super().__eq__(other)
- def is_base_of(self, derived: "StructInfo") -> bool:
- """Check if self is base of another derived struct info.
-
- Parameters
- ----------
- derived : StructInfo
- The derived struct info to be checked.
-
- Returns
- -------
- result : bool
- The check result.
- """
- return _ffi_api.StructInfoIsBaseOf(self, derived) # type: ignore
+Type.is_base_of = _relax_type_is_base_of # type: ignore[attr-defined]
# will be registered afterwards in python/tvm/relax/op/init.py
@@ -155,7 +127,7 @@ class ExprWithOp(Expr, Scriptable):
# NOTE: Cannot override __eq__ and __ne__, which will influence object equal
def __add__(self, other: Expr) -> "ExprWithOp":
- if isinstance(self.struct_info_, tvm.relax.TupleStructInfo) and isinstance(other, tuple):
+ if isinstance(self.ty, tvm.relax.TupleType) and isinstance(other, tuple):
return tuple([*self, *other])
return _binary_op_helper(self, other, _op_ffi_api.add) # type: ignore
@@ -246,7 +218,7 @@ class ExprWithOp(Expr, Scriptable):
raise IndexError from err
raise
- def _check_for_tensor_struct_info(self):
+ def _check_for_tensor_ty(self):
"""Raise an error if this is something other than a Tensor
Used for early checks in `expr.dtype` and `expr.shape`
@@ -254,10 +226,10 @@ class ExprWithOp(Expr, Scriptable):
raised during shape inference, an earlier check makes it
easier to find the invalid usage.
"""
- if self.struct_info_ is None:
+ if self.ty is None:
return
- if not isinstance(self.struct_info_, tvm.relax.TensorStructInfo):
+ if not isinstance(self.ty, tvm.relax.TensorType):
raise TypeError(
f"Runtime unpacking of DLDataType is only implemented for tensors, "
f"but was applied to object {self} of type {type(self)}."
@@ -266,32 +238,32 @@ class ExprWithOp(Expr, Scriptable):
@property
def dtype(self) -> "_DLTensorDTypeProxy":
"""Returns a proxy object for accessing DLTensor::dtype"""
- self._check_for_tensor_struct_info()
+ self._check_for_tensor_ty()
return _DLTensorDTypeProxy(self)
@property
def ndim(self) -> "Expr":
"""Returns the runtime value of DLTensor::ndim"""
- self._check_for_tensor_struct_info()
+ self._check_for_tensor_ty()
op = tvm.ir.Op.get("relax.inspect.tensor_ndim")
return tvm.relax.Call(op, [self])
@property
def shape(self) -> "_DLTensorShapeProxy":
"""Returns a proxy object for accessing DLTensor::shape"""
- self._check_for_tensor_struct_info()
+ self._check_for_tensor_ty()
return _DLTensorShapeProxy(self)
@property
def strides(self) -> "_DLTensorStrideProxy":
"""Returns a proxy object for accessing DLTensor::strides"""
- self._check_for_tensor_struct_info()
+ self._check_for_tensor_ty()
return _DLTensorStrideProxy(self)
@property
def byte_offset(self) -> "Expr":
"""Returns a proxy object for accessing DLTensor::byte_offset"""
- self._check_for_tensor_struct_info()
+ self._check_for_tensor_ty()
op = tvm.ir.Op.get("relax.inspect.tensor_byte_offset")
return tvm.relax.Call(op, [self])
@@ -305,7 +277,7 @@ class ExprWithOp(Expr, Scriptable):
`tirx::BufferNode::elem_offset` field when interacting with TIR
buffers.
"""
- self._check_for_tensor_struct_info()
+ self._check_for_tensor_ty()
op = tvm.ir.Op.get("relax.inspect.tensor_elem_offset")
return tvm.relax.Call(op, [self])
@@ -447,13 +419,11 @@ class _DLTensorShapeProxy(tvm.runtime.ObjectConvertible):
if not isinstance(axis, tvm.relax.Expr):
axis = tvm.relax.PrimValue(axis)
- if axis.struct_info_ is not None and not isinstance(
- axis.struct_info_, tvm.relax.PrimStructInfo
- ):
+ if axis.ty is not None and not isinstance(axis.ty, tvm.relax.PrimType):
raise TypeError(
f"The index used to access {self.tensor}.shape "
- f'must have struct info R.Prim("int64"), '
- f"but index {axis} had struct info {axis.struct_info_}."
+ f'must have type R.Prim("int64"), '
+ f"but index {axis} had type {axis.ty}."
)
op = tvm.ir.Op.get("relax.inspect.tensor_shape_i")
@@ -517,13 +487,11 @@ class _DLTensorStrideProxy(tvm.runtime.ObjectConvertible):
if not isinstance(axis, tvm.relax.Expr):
axis = tvm.relax.PrimValue(axis)
- if axis.struct_info_ is not None and not isinstance(
- axis.struct_info_, tvm.relax.PrimStructInfo
- ):
+ if axis.ty is not None and not isinstance(axis.ty, tvm.relax.PrimType):
raise TypeError(
f"The index used to access {self.tensor}.strides "
- f'must have struct info R.Prim("int64"), '
- f"but index {axis} had struct info {axis.struct_info_}."
+ f'must have type R.Prim("int64"), '
+ f"but index {axis} had type {axis.ty}."
)
op = tvm.ir.Op.get("relax.inspect.tensor_stride_i")
@@ -548,11 +516,11 @@ class Call(ExprWithOp):
attrs: Optional[tvm.ir.Attrs]
Attributes to the call, can be None
- sinfo_args: Optional[Union[List[StructInfo], typing.Tuple[StructInfo, ...]]]
- The structure info arguments of a CallNode.
- sinfo_args is designed to be non-empty only for intrinsic op (e.g.,
+ ty_args: Optional[Union[List[Type], typing.Tuple[Type, ...]]]
+ The type information arguments of a CallNode.
+ ty_args is designed to be non-empty only for intrinsic op (e.g.,
call_tir, call_builtin_with_ctx, etc.) and calls to ExternFuncs, with the main
- usage of structure info inference.
+ usage of type information inference.
span: Optional[Span]
Span that points to original source code
@@ -561,7 +529,7 @@ class Call(ExprWithOp):
op: Expr
args: list[Expr]
attrs: tvm.ir.Attrs
- sinfo_args: list[StructInfo]
+ ty_args: list[Type]
span: Span | None
def __init__(
@@ -569,17 +537,17 @@ class Call(ExprWithOp):
op: Expr | tvm.ir.Op,
args: list[Expr] | tuple[Expr, ...],
attrs: tvm.ir.Attrs | None = None,
- sinfo_args: list[StructInfo] | tuple[StructInfo, ...] | None = None,
+ ty_args: list[Type] | tuple[Type, ...] | None = None,
span: Span | None = None,
):
- if not sinfo_args:
- sinfo_args = []
+ if not ty_args:
+ ty_args = []
self.__init_handle_by_constructor__(
_ffi_api.Call,
op,
args,
attrs,
- sinfo_args,
+ ty_args,
span, # type: ignore
)
@@ -637,7 +605,7 @@ class Tuple(ExprWithOp):
def __init__(self, fields: list[Expr] | tuple[Expr, ...], span: Span | None = None):
if isinstance(fields, tvm.relax.Tuple):
fields = fields.fields
- elif isinstance(getattr(fields, "struct_info_", None), tvm.relax.TupleStructInfo):
+ elif isinstance(getattr(fields, "ty", None), tvm.relax.TupleType):
fields = [*fields]
self.__init_handle_by_constructor__(_ffi_api.Tuple, fields, span) # type: ignore
@@ -730,8 +698,8 @@ class Constant(ExprWithOp):
data: tvm.runtime.Tensor
The data of the constant tensor.
- struct_info: Optional[StructInfo]
- The struct info of the constant tensor. If not specified, infer it from data.
+ ty: Optional[Type]
+ The type of the constant tensor. If not specified, infer it from data.
span: Optional[Span]
Span that points to original source code
@@ -747,13 +715,13 @@ class Constant(ExprWithOp):
def __init__(
self,
data: tvm.runtime.Tensor,
- struct_info: StructInfo | None = None,
+ ty: Type | None = None,
span: Span | None = None,
) -> None:
self.__init_handle_by_constructor__(
_ffi_api.Constant,
data,
- struct_info,
+ ty,
span, # type: ignore
)
@@ -767,8 +735,8 @@ class Var(ExprWithOp):
name_hint: str | Id
The name hint of the variable.
- struct_info: Optional[StructInfo]
- The struct info annotation of the variable.
+ ty: Optional[Type]
+ The type annotation of the variable.
span: Optional[Span]
Span that points to original source code
@@ -780,21 +748,21 @@ class Var(ExprWithOp):
def __init__(
self,
name_hint: str | Id,
- struct_info: StructInfo | None = None,
+ ty: Type | None = None,
span: Span | None = None,
) -> None:
- if struct_info is not None:
- struct_info = tvm.runtime.convert(struct_info)
- if not isinstance(struct_info, StructInfo):
+ if ty is not None:
+ ty = tvm.runtime.convert(ty)
+ if not isinstance(ty, Type):
raise TypeError(
- "struct_info needs to be an instance of StructInfo. "
+ "ty needs to be an instance of Type. "
"If you attempt to pass in shape, "
- "use relax.TensorStructInfo(shape, dtype)."
+ "use relax.TensorType(shape, dtype)."
)
self.__init_handle_by_constructor__(
_ffi_api.Var if isinstance(name_hint, str) else _ffi_api.VarFromId, # type: ignore
name_hint,
- struct_info,
+ ty,
span,
)
@@ -816,8 +784,8 @@ class DataflowVar(Var):
name_hint: str | Id
The name hint of the variable.
- struct_info: Optional[StructInfo]
- The struct info annotation of the variable.
+ ty: Optional[Type]
+ The type annotation of the variable.
span: Optional[Span]
Span that points to original source code
@@ -829,17 +797,17 @@ class DataflowVar(Var):
def __init__(
self,
name_hint: str | Id,
- struct_info: StructInfo | None = None,
+ ty: Type | None = None,
span: Span | None = None,
) -> None:
# pylint: disable=super-init-not-called
- if struct_info is not None:
- struct_info = tvm.runtime.convert(struct_info)
- if not isinstance(struct_info, StructInfo):
+ if ty is not None:
+ ty = tvm.runtime.convert(ty)
+ if not isinstance(ty, Type):
raise TypeError(
- "struct_info needs to be an instance of StructInfo. "
+ "ty needs to be an instance of Type. "
"If you attempt to pass in shape, "
- "use relax.TensorStructInfo(shape, dtype)."
+ "use relax.TensorType(shape, dtype)."
)
self.__init_handle_by_constructor__(
@@ -849,7 +817,7 @@ class DataflowVar(Var):
else _ffi_api.DataflowVarFromId
), # type: ignore
name_hint,
- struct_info,
+ ty,
span,
)
@@ -898,10 +866,10 @@ class Binding(Node, Scriptable):
@tvm_ffi.register_object("relax.expr.MatchCast")
class MatchCast(Binding):
- """Runtime-match the value to the struct info.
+ """Runtime-match the value to the type.
This operation does runtime check, populates the un-defined symbolic shape vars
- and vars in struct_info in the first occurrence, and insert equality assertions in
+ and vars in ty in the first occurrence, and insert equality assertions in
other cases.
Parameters
@@ -912,22 +880,20 @@ class MatchCast(Binding):
value: Expr
The input value expression.
- struct_info: tvm.relax.StructInfo
- The struct info to match cast to.
+ ty: tvm.relax.Type
+ The type to match cast to.
"""
- struct_info: StructInfo
+ ty: Type
value: Expr
span: Span | None
- def __init__(
- self, var: Var, value: Expr, struct_info: StructInfo, span: Span | None = None
- ) -> None:
+ def __init__(self, var: Var, value: Expr, ty: Type, span: Span | None = None) -> None:
self.__init_handle_by_constructor__(
_ffi_api.MatchCast,
var,
value,
- struct_info,
+ ty,
span, # type: ignore
)
@@ -996,7 +962,7 @@ class Function(BaseFunc, Scriptable):
params: list[Var]
body: Expr
- ret_struct_info: StructInfo
+ ret_ty: Type
is_pure: bool
attrs: tvm.ir.DictAttrs
span: Span | None
@@ -1005,7 +971,7 @@ class Function(BaseFunc, Scriptable):
self,
params: list[Var],
body: Expr,
- ret_struct_info: StructInfo | None = None,
+ ret_ty: Type | None = None,
is_pure: bool | None = True,
attrs: tvm.ir.DictAttrs | None = None,
span: Span | None = None,
@@ -1016,7 +982,7 @@ class Function(BaseFunc, Scriptable):
_ffi_api.Function,
params,
body,
- ret_struct_info,
+ ret_ty,
is_pure,
attrs,
span,
@@ -1025,7 +991,7 @@ class Function(BaseFunc, Scriptable):
@staticmethod
def create_empty(
params: list[Var],
- ret_struct_info: StructInfo,
+ ret_ty: Type,
is_pure: bool | None = True,
attrs: tvm.ir.DictAttrs | None = None,
span: Span | None = None,
@@ -1033,7 +999,7 @@ class Function(BaseFunc, Scriptable):
"""Construct a relax.Function but without body"""
if attrs is None:
attrs = tvm.ir.DictAttrs({})
- return _ffi_api.FunctionCreateEmpty(params, ret_struct_info, is_pure, attrs, span) # type: ignore
+ return _ffi_api.FunctionCreateEmpty(params, ret_ty, is_pure, attrs, span) # type: ignore
def __call__(self, *args):
"""Invoke the global function.
@@ -1139,20 +1105,20 @@ class ExternFunc(BaseFunc, ExprWithOp):
def __init__(
self,
global_symbol: String,
- struct_info: StructInfo | None = None,
+ ty: Type | None = None,
span: Span | None = None,
) -> None:
self.__init_handle_by_constructor__(
_ffi_api.ExternFunc,
global_symbol,
- struct_info,
+ ty,
span, # type: ignore
)
-def extern(name: str, struct_info: StructInfo | None = None, span: Span | None = None):
+def extern(name: str, ty: Type | None = None, span: Span | None = None):
"""Create extern function."""
- return ExternFunc(name, struct_info, span)
+ return ExternFunc(name, ty, span)
def const(
@@ -1218,7 +1184,7 @@ def te_tensor(
Parameters
----------
value : Expr
- The relax expression, which is required to have TensorStructInfo.
+ The relax expression, which is required to have TensorType.
tir_var_map : Dict[tvm.tirx.Var, tvm.tirx.PrimExpr]
The mapping to substitute the TIR variables appeared in the
@@ -1246,7 +1212,7 @@ def get_shape_of(expr: Expr) -> Expr:
Note
----
This function requires expr to be normalized.
- The function will report an error if expr's StructInfo is not TensorStructInfo.
+ The function will report an error if expr's Type is not TensorType.
It will try to return symbolic function when possible. If the tensor do not
have a compile-time symbolic shape, the function will then choose to return
`Call(relax.op.shape_of, [expr])`.
@@ -1254,5 +1220,5 @@ def get_shape_of(expr: Expr) -> Expr:
return _ffi_api.GetShapeOf(expr) # type: ignore
-def _update_struct_info(expr: Expr, struct_info: StructInfo | None) -> None:
- _ffi_api.UpdateStructInfo(expr, struct_info) # type: ignore
+def _update_type(expr: Expr, ty: Type | None) -> None:
+ _ffi_api.UpdateType(expr, ty) # type: ignore
diff --git a/python/tvm/relax/expr_functor.py b/python/tvm/relax/expr_functor.py
index c9ea88d111..e84e37b02a 100644
--- a/python/tvm/relax/expr_functor.py
+++ b/python/tvm/relax/expr_functor.py
@@ -54,7 +54,7 @@ from .expr import (
Var,
VarBinding,
)
-from .struct_info import StructInfo
+from .type import Type
visitor = derived_object
"""
@@ -1521,7 +1521,7 @@ class PyExprMutator:
# Using self._outer() to ref _PyExprMutator
return _ffi_api.PyExprMutatorLookupBinding(self._outer(), var) # type: ignore
- def with_struct_info(self, var: Var, struct_info: StructInfo) -> Var:
+ def with_type(self, var: Var, ty: Type) -> Var:
"""Create a new var with specified shape and type if the original var's shape or type does
not match with the specified ones.
@@ -1529,8 +1529,8 @@ class PyExprMutator:
----------
var : Var
The var to be updated.
- struct_info : StructInfo
- The struct info.
+ ty : Type
+ The type.
Returns
-------
@@ -1538,4 +1538,4 @@ class PyExprMutator:
The var filled with shape and type.
"""
# Using self._outer() to ref _PyExprMutator
- return _ffi_api.PyExprMutatorWithStructInfo(self._outer(), var, struct_info) # type: ignore
+ return _ffi_api.PyExprMutatorWithType(self._outer(), var, ty) # type: ignore
diff --git a/python/tvm/relax/frontend/common.py b/python/tvm/relax/frontend/common.py
index 63b58cd8f1..b4aad8f778 100644
--- a/python/tvm/relax/frontend/common.py
+++ b/python/tvm/relax/frontend/common.py
@@ -75,8 +75,8 @@ def autopad(
[(kernel - 1) * dilation + 1 for kernel, dilation in zip(kernel_shape, dilations)]
)
# get input shape
- ndim = data.struct_info.ndim
- data_shape = list(data.struct_info.shape)
+ ndim = data.ty.ndim
+ data_shape = list(data.ty.shape)
shape = data_shape[2:ndim]
# set up integer constants
diff --git a/python/tvm/relax/frontend/nn/core.py b/python/tvm/relax/frontend/nn/core.py
index 1301fa471f..c3ef1fb268 100644
--- a/python/tvm/relax/frontend/nn/core.py
+++ b/python/tvm/relax/frontend/nn/core.py
@@ -15,7 +15,7 @@
# specific language governing permissions and limitations
# under the License.
"""The core infra for nn.Module, which includes the following pieces:
-- Tensor, a wrapper on top of relax.Expr whose struct_info is a TensorStructInfo,
+- Tensor, a wrapper on top of relax.Expr whose ty is a TensorType,
providing more convenient access shape and dtype information.
Tensor is always symbolic and not bound to any concrete values.
- Parameter, a special tensor which could be bound or not bound to concrete values.
@@ -45,11 +45,11 @@ from tvm.target import Target
from .... import relax as rx
from ...block_builder import BlockBuilder
-from ...struct_info import (
- ObjectStructInfo,
- ShapeStructInfo,
- TensorStructInfo,
- TupleStructInfo,
+from ...type import (
+ ObjectType,
+ ShapeType,
+ TensorType,
+ TupleType,
)
from ._tensor_op import _TensorOp
from .subroutine import SubroutineMixin
@@ -88,7 +88,7 @@ def set_default_dtype(dtype: str) -> None:
class Tensor(_TensorOp):
- """A wrapper on top of relax.Expr whose struct_info is a TensorStructInfo, providing more
+ """A wrapper on top of relax.Expr whose ty is a TensorType, providing more
convenient access shape and dtype information. Tensor is always symbolc and not bound to any
concrete values. Shape and dtype inference is done eagerly upon tensor creation, i.e. when
operators are applied on tensors, the shape and dtype information is already available.
@@ -100,13 +100,13 @@ class Tensor(_TensorOp):
"""Private constructor. Tensor is never supposed to be constructed directly by users."""
def _check_tensor(expr: rx.Expr) -> None:
- assert expr.struct_info_ is not None
- assert isinstance(expr.struct_info, TensorStructInfo)
- assert expr.struct_info.ndim != -1
- assert expr.struct_info.shape is not None
- assert expr.struct_info.shape.struct_info_ is not None
- assert isinstance(expr.struct_info.shape.struct_info, ShapeStructInfo)
- assert expr.struct_info.shape.struct_info.values is not None
+ assert expr.ty is not None
+ assert isinstance(expr.ty, TensorType)
+ assert expr.ty.ndim != -1
+ assert expr.ty.shape is not None
+ assert expr.ty.shape.ty is not None
+ assert isinstance(expr.ty.shape.ty, ShapeType)
+ assert expr.ty.shape.ty.values is not None
_check_tensor(_expr)
self._expr = _expr
@@ -122,17 +122,17 @@ class Tensor(_TensorOp):
return Tensor(_expr=rx.const(data, dtype=dtype))
@staticmethod
- def from_struct_info(struct_info: rx.TensorStructInfo, name: str = "tensor") -> "Tensor":
- """Construct a nn.Tensor from a Relax TensorStructInfo.
+ def from_ty(ty: rx.TensorType, name: str = "tensor") -> "Tensor":
+ """Construct a nn.Tensor from a Relax TensorType.
- TensorStructInfo is the Relax type-level description of a tensor, carrying its shape
+ TensorType is the Relax type-level description of a tensor, carrying its shape
and dtype without holding actual data. This factory creates an unbound placeholder
``nn.Tensor`` that can be used as a symbolic input when tracing an ``nn.Module``.
Parameters
----------
- struct_info : rx.TensorStructInfo
- The struct info describing the tensor's shape and dtype.
+ ty : rx.TensorType
+ The type describing the tensor's shape and dtype.
name : str
Name hint for the underlying Relax variable.
@@ -140,12 +140,12 @@ class Tensor(_TensorOp):
Returns
-------
tensor : Tensor
- A symbolic ``nn.Tensor`` backed by a ``relax.Var`` with the given struct info.
+ A symbolic ``nn.Tensor`` backed by a ``relax.Var`` with the given type.
"""
return Tensor(
_expr=rx.Var(
name_hint=name,
- struct_info=struct_info,
+ ty=ty,
)
)
@@ -179,7 +179,7 @@ class Tensor(_TensorOp):
return Tensor(
_expr=rx.Var(
name_hint=name,
- struct_info=TensorStructInfo(
+ ty=TensorType(
shape=new_shape, # type: ignore[arg-type]
dtype=dtype,
),
@@ -203,8 +203,8 @@ class Tensor(_TensorOp):
def _simplify(expr: tirx.PrimExpr):
return expr.value if isinstance(expr, tirx.IntImm) else expr
- shape_sinfo: ShapeStructInfo = self._expr.struct_info.shape.struct_info
- return [_simplify(x) for x in shape_sinfo.values]
+ shape_ty: ShapeType = self._expr.ty.shape.ty
+ return [_simplify(x) for x in shape_ty.values]
@property
def ndim(self) -> int:
@@ -215,7 +215,7 @@ class Tensor(_TensorOp):
ndim : int
The number of dimensions of the tensor
"""
- return self._expr.struct_info.ndim
+ return self._expr.ty.ndim
@property
def dtype(self) -> str:
@@ -226,7 +226,7 @@ class Tensor(_TensorOp):
dtype : str
The data type of the tensor
"""
- return self._expr.struct_info.dtype
+ return self._expr.ty.dtype
def __repr__(self) -> str:
return f'Tensor({self.shape}, "{self.dtype}")'
@@ -310,8 +310,8 @@ class Parameter(Tensor):
class Object:
- """A wrapper on top of relax.Expr whose struct_info is the base
- ObjectStructInfo (rather than any its subclass). Object effectively
+ """A wrapper on top of relax.Expr whose ty is the base
+ ObjectType (rather than any its subclass). Object effectively
represents non-tensor frontend components such as KV caches.
"""
@@ -322,7 +322,7 @@ class Object:
if not isinstance(_expr, rx.Var):
_expr = BlockBuilder.current().emit(_expr, _name)
self._expr = _expr
- assert isinstance(self._expr.struct_info, ObjectStructInfo)
+ assert isinstance(self._expr.ty, ObjectType)
class Effect:
@@ -778,17 +778,17 @@ def wrap_nested(expr: rx.Expr, name: str) -> Tensor | Sequence[Tensor]:
"""
if not isinstance(expr, rx.DataflowVar):
expr = BlockBuilder.current().emit(expr, name)
- if isinstance(expr.struct_info_, TensorStructInfo):
+ if isinstance(expr.ty, TensorType):
return Tensor(_expr=expr)
- if isinstance(expr.struct_info_, TupleStructInfo):
+ if isinstance(expr.ty, TupleType):
return tuple(
wrap_nested( # type: ignore
rx.TupleGetItem(expr, i),
name=f"{name}.{i}",
)
- for i in range(len(expr.struct_info_.fields))
+ for i in range(len(expr.ty.fields))
)
- raise TypeError(f"Unsupported return type: {expr.struct_info_}")
+ raise TypeError(f"Unsupported return type: {expr.ty}")
def _attribute_finder(root: Module, prefix: str, condition_yield: Callable[[Any], bool]):
diff --git a/python/tvm/relax/frontend/nn/exporter.py b/python/tvm/relax/frontend/nn/exporter.py
index 8b68bb6e70..627bbe1506 100644
--- a/python/tvm/relax/frontend/nn/exporter.py
+++ b/python/tvm/relax/frontend/nn/exporter.py
@@ -26,7 +26,7 @@ from tvm.ir import IRModule
from .... import relax as rx
from ...block_builder import BlockBuilder
-from ...struct_info import ObjectStructInfo, ShapeStructInfo, TupleStructInfo
+from ...type import ObjectType, ShapeType, TupleType
from . import core, extern
from . import spec as _spec
from .modules import IOEffect
@@ -178,15 +178,13 @@ def _emit_method( # pylint: disable=too-many-locals,too-many-branches,too-many-
def _convert_input(arg):
if isinstance(arg, tirx.Var):
- return rx.Var(arg.name, struct_info=ShapeStructInfo(values=[arg]))
+ return rx.Var(arg.name, ty=ShapeType(values=[arg]))
if isinstance(arg, core.Tensor | core.Object):
return arg._expr # pylint: disable=protected-access
if isinstance(arg, _spec.Tuple):
return rx.Var(
arg.name,
- struct_info=TupleStructInfo(
- [_convert_input(arg_i).struct_info for arg_i in arg.elements]
- ),
+ ty=TupleType([_convert_input(arg_i).ty for arg_i in arg.elements]),
)
raise TypeError(f"Unsupported input type: {type(arg)}")
@@ -215,7 +213,7 @@ def _emit_method( # pylint: disable=too-many-locals,too-many-branches,too-many-
if mode == "packed":
input_var = rx.Var(
"packed_params",
- TupleStructInfo(fields=[x.struct_info for x in inputs]),
+ TupleType(fields=[x.ty for x in inputs]),
)
for i, (name, param) in enumerate(params):
param._expr = builder.emit(rx.TupleGetItem(input_var, i), name_hint=name)
@@ -236,7 +234,7 @@ def _emit_method( # pylint: disable=too-many-locals,too-many-branches,too-many-
if mode == "packed":
input_var = rx.Var(
"packed_effects",
- TupleStructInfo(fields=[x.struct_info for x in inputs]),
+ TupleType(fields=[x.ty for x in inputs]),
)
i = 0
for effect_input, (_, effect) in zip(unflat_inputs, effects):
@@ -313,7 +311,7 @@ def _method_spec_to_inputs(
name=arg_name,
)
elif isinstance(arg_spec, _spec.Object):
- arg = arg_spec.object_type(_expr=rx.Var(arg_name, ObjectStructInfo()), _name=arg_name)
+ arg = arg_spec.object_type(_expr=rx.Var(arg_name, ObjectType()), _name=arg_name)
elif isinstance(arg_spec, _spec.Tuple):
elements = type(arg_spec.elements)(
[
diff --git a/python/tvm/relax/frontend/nn/extern.py b/python/tvm/relax/frontend/nn/extern.py
index 866da6c4ab..9c8efce690 100644
--- a/python/tvm/relax/frontend/nn/extern.py
+++ b/python/tvm/relax/frontend/nn/extern.py
@@ -67,8 +67,8 @@ class ExternModule:
raise TypeError(f"Unsupported input type: {type(arg)}")
rx_inputs = _convert(input_args, "input")
- rx_outputs_sinfo = _convert(_inference_function(*input_args), "dummy").struct_info
- return wrap_nested(call_dps_packed(func_name, rx_inputs, rx_outputs_sinfo), func_name)
+ rx_outputs_ty = _convert(_inference_function(*input_args), "dummy").ty
+ return wrap_nested(call_dps_packed(func_name, rx_inputs, rx_outputs_ty), func_name)
return _call
diff --git a/python/tvm/relax/frontend/nn/llm/kv_cache.py b/python/tvm/relax/frontend/nn/llm/kv_cache.py
index 0e0828c75d..1da233a78a 100644
--- a/python/tvm/relax/frontend/nn/llm/kv_cache.py
+++ b/python/tvm/relax/frontend/nn/llm/kv_cache.py
@@ -141,7 +141,7 @@ class PagedKVCache(Object): # pylint: disable=too-few-public-methods
- The input qkv and output tensor have `head_dim` at the last dim.
"""
# pylint: disable=protected-access
- b, s, _, d = qkv._expr.struct_info.shape
+ b, s, _, d = qkv._expr.ty.shape
qkv = qkv.reshape(b * s, qkv.shape[2], d)
return Tensor(
_expr=rx.BlockBuilder.current().emit(
@@ -153,7 +153,7 @@ class PagedKVCache(Object): # pylint: disable=too-few-public-methods
rx.PrimValue(sm_scale),
qkv._expr,
],
- out_sinfo=rx.TensorStructInfo((b * s, num_qo_heads, d), qkv.dtype),
+ out_ty=rx.TensorType((b * s, num_qo_heads, d), qkv.dtype),
)
)
).reshape(b, s, num_qo_heads, d)
@@ -168,8 +168,8 @@ class PagedKVCache(Object): # pylint: disable=too-few-public-methods
) -> tuple[Tensor, Tensor]:
"""Fine-grained API that computes ragged self attention with Q/K/V data."""
# pylint: disable=protected-access
- b, s, h_qo, d_qk = q._expr.struct_info.shape
- _, _, h_kv, d_v = v._expr.struct_info.shape
+ b, s, h_qo, d_qk = q._expr.ty.shape
+ _, _, h_kv, d_v = v._expr.ty.shape
q = q.reshape(b * s, h_qo, d_qk)
k = k.reshape(b * s, h_kv, d_qk)
v = v.reshape(b * s, h_kv, d_v)
@@ -185,14 +185,14 @@ class PagedKVCache(Object): # pylint: disable=too-few-public-methods
k._expr,
v._expr,
],
- out_sinfo=[
- rx.TensorStructInfo((b * s, h_qo, d_v), q.dtype),
- rx.TensorStructInfo((b * s, h_qo), "float32"),
+ out_ty=[
+ rx.TensorType((b * s, h_qo, d_v), q.dtype),
+ rx.TensorType((b * s, h_qo), "float32"),
],
)
)
- assert isinstance(attn_results.struct_info, rx.TupleStructInfo)
- assert len(attn_results.struct_info.fields) == 2
+ assert isinstance(attn_results.ty, rx.TupleType)
+ assert len(attn_results.ty.fields) == 2
o = Tensor(_expr=bb.emit(rx.TupleGetItem(attn_results, 0))).reshape(b, s, h_qo, d_v)
lse = Tensor(_expr=bb.emit(rx.TupleGetItem(attn_results, 1))).reshape(b, s, h_qo)
return o, lse
@@ -206,7 +206,7 @@ class PagedKVCache(Object): # pylint: disable=too-few-public-methods
) -> tuple[Tensor, Tensor]:
"""Fine-grained API that computes paged cross attention with Q and in-cache KV data."""
# pylint: disable=protected-access
- b, s, h_qo, d_qk = q._expr.struct_info.shape
+ b, s, h_qo, d_qk = q._expr.ty.shape
q = q.reshape(b * s, h_qo, d_qk)
bb = rx.BlockBuilder.current()
attn_results = bb.emit(
@@ -218,14 +218,14 @@ class PagedKVCache(Object): # pylint: disable=too-few-public-methods
rx.PrimValue(sm_scale),
q._expr,
],
- out_sinfo=[
- rx.TensorStructInfo((b * s, h_qo, v_head_dim), q.dtype),
- rx.TensorStructInfo((b * s, h_qo), "float32"),
+ out_ty=[
+ rx.TensorType((b * s, h_qo, v_head_dim), q.dtype),
+ rx.TensorType((b * s, h_qo), "float32"),
],
)
)
- assert isinstance(attn_results.struct_info, rx.TupleStructInfo)
- assert len(attn_results.struct_info.fields) == 2
+ assert isinstance(attn_results.ty, rx.TupleType)
+ assert len(attn_results.ty.fields) == 2
o = Tensor(_expr=bb.emit(rx.TupleGetItem(attn_results, 0))).reshape(b, s, h_qo, v_head_dim)
lse = Tensor(_expr=bb.emit(rx.TupleGetItem(attn_results, 1))).reshape(b, s, h_qo)
return o, lse
@@ -233,7 +233,7 @@ class PagedKVCache(Object): # pylint: disable=too-few-public-methods
def append_mla_kv(self, layer_id: int, kv: Tensor) -> "PagedKVCache":
"""Fine-grained API that appends the MLA K/V data to KV cache."""
# pylint: disable=protected-access
- b, s, _, d_qk = kv._expr.struct_info.shape
+ b, s, _, d_qk = kv._expr.ty.shape
kv = kv.reshape(b * s, d_qk)
return PagedKVCache(
_expr=rx.call_pure_packed(
@@ -241,7 +241,7 @@ class PagedKVCache(Object): # pylint: disable=too-few-public-methods
self._expr,
rx.PrimValue(layer_id), # type: ignore[arg-type]
kv._expr,
- sinfo_args=rx.ObjectStructInfo(),
+ ty_args=rx.ObjectType(),
),
_name="paged_kv_cache",
)
@@ -257,7 +257,7 @@ class PagedKVCache(Object): # pylint: disable=too-few-public-methods
The first two tensors will be inplace updated.
"""
# pylint: disable=protected-access
- b, s, h_qo, d_v = o_self_attn._expr.struct_info.shape
+ b, s, h_qo, d_v = o_self_attn._expr.ty.shape
o_self_attn = o_self_attn.reshape(b * s, h_qo, d_v)
lse_self_attn = lse_self_attn.reshape(b * s, h_qo)
o_cross_attn = o_cross_attn.reshape(b * s, h_qo, d_v)
@@ -271,13 +271,13 @@ class PagedKVCache(Object): # pylint: disable=too-few-public-methods
lse_self_attn._expr,
o_cross_attn._expr,
lse_cross_attn._expr,
- sinfo_args=rx.TupleStructInfo(
- [o_self_attn._expr.struct_info, lse_self_attn._expr.struct_info]
+ ty_args=rx.TupleType(
+ [o_self_attn._expr.ty, lse_self_attn._expr.ty]
),
)
)
- assert isinstance(merge_results.struct_info, rx.TupleStructInfo)
- assert len(merge_results.struct_info.fields) == 2
+ assert isinstance(merge_results.ty, rx.TupleType)
+ assert len(merge_results.ty.fields) == 2
o_self_attn = Tensor(_expr=bb.emit(rx.TupleGetItem(merge_results, 0))).reshape(
b, s, h_qo, d_v
)
@@ -304,7 +304,7 @@ class PagedKVCache(Object): # pylint: disable=too-few-public-methods
rx.call_pure_packed(
"vm.builtin.attention_kv_cache_get_query_positions",
self._expr,
- sinfo_args=rx.TensorStructInfo((total_length,), "int32"),
+ ty_args=rx.TensorType((total_length,), "int32"),
)
)
)
@@ -505,7 +505,7 @@ class FlashInferPagedKVCache(PagedKVCache): # pylint: disable=too-few-public-me
_expr=rx.call_pure_packed(
"vm.builtin.paged_attention_kv_cache_create",
*args,
- sinfo_args=rx.ObjectStructInfo(),
+ ty_args=rx.ObjectType(),
),
_name=name,
)
@@ -680,7 +680,7 @@ class TIRPagedKVCache(PagedKVCache): # pylint: disable=too-few-public-methods
_expr=rx.call_pure_packed(
"vm.builtin.paged_attention_kv_cache_create",
*args,
- sinfo_args=rx.ObjectStructInfo(),
+ ty_args=rx.ObjectType(),
),
_name=name,
)
diff --git a/python/tvm/relax/frontend/nn/modules.py b/python/tvm/relax/frontend/nn/modules.py
index 8753d37bac..6bd90c9c4a 100644
--- a/python/tvm/relax/frontend/nn/modules.py
+++ b/python/tvm/relax/frontend/nn/modules.py
@@ -43,7 +43,7 @@ class IOEffect(Effect):
def create(self, name_hint: str) -> list[rx.Var]:
assert self.effect is None
- effect = rx.Var(f"{name_hint}.io", struct_info=rx.ObjectStructInfo())
+ effect = rx.Var(f"{name_hint}.io", ty=rx.ObjectType())
return [effect]
def set_state(self, state_vars: list[rx.Var]) -> None:
@@ -812,7 +812,7 @@ class KVCache(Effect):
rx.op.zeros(init_shape, self.dtype),
init_shape,
rx.PrimValue(0),
- sinfo_args=rx.ObjectStructInfo(),
+ ty_args=rx.ObjectType(),
),
name_hint=name_hint,
)
@@ -832,7 +832,7 @@ class KVCache(Effect):
ret : List[relax.Var]
The relax.Var for KVCache.
"""
- cache = rx.Var(name_hint, struct_info=rx.ObjectStructInfo())
+ cache = rx.Var(name_hint, ty=rx.ObjectType())
return [cache]
def set_state(self, state_vars: list[rx.Var]) -> None:
@@ -884,7 +884,7 @@ class KVCache(Effect):
"vm.builtin.attention_kv_cache_view",
self.cache,
shape,
- sinfo_args=rx.TensorStructInfo(shape, self.dtype),
+ ty_args=rx.TensorType(shape, self.dtype),
)
)
)
@@ -908,7 +908,7 @@ class KVCache(Effect):
self.cache,
new_element._expr,
inplace_indices=[0],
- sinfo_args=rx.ObjectStructInfo(),
+ ty_args=rx.ObjectType(),
)
)
diff --git a/python/tvm/relax/frontend/nn/op.py b/python/tvm/relax/frontend/nn/op.py
index 80108e317e..f31c276ca0 100644
--- a/python/tvm/relax/frontend/nn/op.py
+++ b/python/tvm/relax/frontend/nn/op.py
@@ -1362,7 +1362,7 @@ def layer_norm(
normalized_shape = [normalized_shape]
dim_num = len(normalized_shape)
axes = list(range(-dim_num, 0))
- dtype = x._expr.struct_info.dtype
+ dtype = x._expr.ty.dtype
if weight is not None:
weight = weight._expr
@@ -1480,7 +1480,7 @@ def group_norm(
weight = weight._expr
if bias is not None:
bias = bias._expr
- dim = len(x._expr.struct_info.shape)
+ dim = len(x._expr.ty.shape)
if axes is None:
axes = list(range(2, dim))
return wrap_nested(
@@ -2084,9 +2084,9 @@ def tensor_ir_op(
)
if isinstance(out, Tensor):
- out_sinfo = [out._expr.struct_info]
+ out_ty = [out._expr.ty]
else:
- out_sinfo = [x._expr.struct_info for x in out]
+ out_ty = [x._expr.ty for x in out]
bb = BlockBuilder.current()
global_var = bb.add_func(func, name_hint)
@@ -2095,7 +2095,7 @@ def tensor_ir_op(
tir_vars = None
return wrap_nested(
- bb.emit(rx.call_tir(global_var, call_tir_args, out_sinfo, tir_vars=tir_vars)),
+ bb.emit(rx.call_tir(global_var, call_tir_args, out_ty, tir_vars=tir_vars)),
name=name_hint,
)
@@ -2154,17 +2154,15 @@ def tensor_ir_inplace_op(
)
if isinstance(out, Tensor):
- out_sinfo = [out._expr.struct_info]
+ out_ty = [out._expr.ty]
else:
- out_sinfo = [x._expr.struct_info for x in out]
+ out_ty = [x._expr.ty for x in out]
bb = BlockBuilder.current()
global_var = bb.add_func(func, name_hint)
return wrap_nested(
- bb.emit(
- rx.call_tir_inplace(global_var, call_tir_args, inplace_indices, out_sinfo, tir_vars)
- ),
+ bb.emit(rx.call_tir_inplace(global_var, call_tir_args, inplace_indices, out_ty, tir_vars)),
name=name_hint,
)
@@ -2211,12 +2209,12 @@ def extern(
raise TypeError(f"Unsupported input type: {type(arg)}")
rx_inputs = _convert(args, "input")
- rx_outputs_sinfo = _convert(out, "dummy").struct_info
+ rx_outputs_ty = _convert(out, "dummy").ty
return wrap_nested(
_op.call_dps_packed(
name,
args=rx_inputs,
- out_sinfo=rx_outputs_sinfo,
+ out_ty=rx_outputs_ty,
),
name,
) # type: ignore
@@ -2282,7 +2280,7 @@ def debug_func(
rx.StringImm(name),
rx.StringImm(_line_info),
*converted_args,
- sinfo_args=[rx.ObjectStructInfo()],
+ ty_args=[rx.ObjectType()],
),
name_hint=io.effect.name_hint,
)
diff --git a/python/tvm/relax/frontend/nn/subroutine.py b/python/tvm/relax/frontend/nn/subroutine.py
index e821756e8d..d197355998 100644
--- a/python/tvm/relax/frontend/nn/subroutine.py
+++ b/python/tvm/relax/frontend/nn/subroutine.py
@@ -41,11 +41,11 @@ def _camel_to_snake(name):
def _normalize_expr(block_builder, arg, as_relax_expr=False):
- """Ensure that an argument is a relax.Expr with struct info"""
+ """Ensure that an argument is a relax.Expr with type"""
if isinstance(arg, tuple):
arg = relax.Tuple([_normalize_expr(block_builder, element) for element in arg])
- if isinstance(arg, relax.Expr) and getattr(arg, "struct_info_", None) is None:
+ if isinstance(arg, relax.Expr) and getattr(arg, "ty", None) is None:
arg = block_builder.emit(arg)
if isinstance(arg, nn.Tensor) and as_relax_expr:
@@ -54,15 +54,15 @@ def _normalize_expr(block_builder, arg, as_relax_expr=False):
return arg
-def _get_struct_info(arg):
+def _get_ty(arg):
if isinstance(arg, relax.Expr):
- return arg.struct_info_
+ return arg.ty
elif isinstance(arg, nn.Tensor):
- return arg._expr.struct_info_
+ return arg._expr.ty
elif isinstance(arg, tuple | list | tvm_ffi.Array):
- return relax.TupleStructInfo([_get_struct_info(field) for field in arg])
+ return relax.TupleType([_get_ty(field) for field in arg])
else:
- raise TypeError(f"Cannot find struct info for {arg} of type {type(arg)}")
+ raise TypeError(f"Cannot find type for {arg} of type {type(arg)}")
class SubroutineMixin:
@@ -108,7 +108,7 @@ class SubroutineMixin:
out = subroutine(*subroutine_args)
if is_nn_tensor_output:
- if out.struct_info_ is None:
+ if out.ty is None:
out = block_builder.emit(out, name_hint=f"{subroutine.name_hint}_output")
out = nn.Tensor(_expr=out)
return out
@@ -141,20 +141,20 @@ class SubroutineMixin:
param._expr if isinstance(param, nn.Tensor) else param for param in self.parameters()
]
- arg_sinfo = _get_struct_info([*func_args.values(), *model_params])
+ arg_ty = _get_ty([*func_args.values(), *model_params])
is_dataflow = block_builder.current_block_is_dataflow()
lookup_key = (
old_forward,
- tvm_ffi.structural_hash(arg_sinfo, map_free_vars=True),
+ tvm_ffi.structural_hash(arg_ty, map_free_vars=True),
is_dataflow,
)
- for cached_sinfo, cached_result in cls._gvar.get(lookup_key, []):
- if tvm_ffi.structural_equal(cached_sinfo, arg_sinfo, map_free_vars=True):
+ for cached_ty, cached_result in cls._gvar.get(lookup_key, []):
+ if tvm_ffi.structural_equal(cached_ty, arg_ty, map_free_vars=True):
return cached_result
func_name = _camel_to_snake(cls.__name__)
- func_params = [relax.Var(name, sinfo) for name, sinfo in zip(func_args, arg_sinfo.fields)]
+ func_params = [relax.Var(name, ty) for name, ty in zip(func_args, arg_ty.fields)]
old_forward_args = [
nn.Tensor(_expr=param) if isinstance(old_arg, nn.Tensor) else param
for param, old_arg in zip(func_params, func_args.values())
@@ -175,7 +175,7 @@ class SubroutineMixin:
gvar = block_builder.emit_func_output(out)
# The relax.Var instances in model_params, along with any
- # tirx.Var instances in the struct info, appear in both the
+ # tirx.Var instances in the type, appear in both the
# calling scope and as parameters for the subroutine. To
# maintain SSA, replace all relax and TIR variables in the
# subroutine.
@@ -184,5 +184,5 @@ class SubroutineMixin:
result = (gvar, is_nn_tensor_output)
bucket = cls._gvar.setdefault(lookup_key, [])
- bucket.append((arg_sinfo, result))
+ bucket.append((arg_ty, result))
return result
diff --git a/python/tvm/relax/frontend/onnx/onnx_frontend.py b/python/tvm/relax/frontend/onnx/onnx_frontend.py
index 430e7cef7a..d550d3bc00 100644
--- a/python/tvm/relax/frontend/onnx/onnx_frontend.py
+++ b/python/tvm/relax/frontend/onnx/onnx_frontend.py
@@ -340,9 +340,9 @@ class QuantizeLinear(OnnxOpConverter):
x, scale = inputs[0], inputs[1]
zp = inputs[2] if len(inputs) > 2 and inputs[2] is not None else None
axis = attr.get("axis", 1)
- if hasattr(x.struct_info, "ndim") and x.struct_info.ndim <= 1 and axis == 1:
+ if hasattr(x.ty, "ndim") and x.ty.ndim <= 1 and axis == 1:
axis = 0
- out_dtype = "uint8" if zp is None else zp.struct_info.dtype
+ out_dtype = "uint8" if zp is None else zp.ty.dtype
if zp is None:
zp = relax.const(0, out_dtype)
return relax.op.quantize(x, scale, zp, axis=axis, out_dtype=out_dtype)
@@ -352,9 +352,9 @@ class QuantizeLinear(OnnxOpConverter):
x, scale = inputs[0], inputs[1]
zp = inputs[2] if len(inputs) > 2 and inputs[2] is not None else None
axis = attr.get("axis", 1)
- if hasattr(x.struct_info, "ndim") and x.struct_info.ndim <= 1 and axis == 1:
+ if hasattr(x.ty, "ndim") and x.ty.ndim <= 1 and axis == 1:
axis = 0
- out_dtype = "uint8" if zp is None else zp.struct_info.dtype
+ out_dtype = "uint8" if zp is None else zp.ty.dtype
if zp is None:
zp = relax.const(0, out_dtype)
return relax.op.quantize(x, scale, zp, axis=axis, out_dtype=out_dtype)
@@ -366,10 +366,10 @@ class DequantizeLinear(OnnxOpConverter):
x, scale = inputs[0], inputs[1]
zp = inputs[2] if len(inputs) > 2 and inputs[2] is not None else None
axis = attr.get("axis", 1)
- if hasattr(x.struct_info, "ndim") and x.struct_info.ndim <= 1 and axis == 1:
+ if hasattr(x.ty, "ndim") and x.ty.ndim <= 1 and axis == 1:
axis = 0
if zp is None:
- zp = relax.const(0, x.struct_info.dtype)
+ zp = relax.const(0, x.ty.dtype)
return relax.op.dequantize(x, scale, zp, axis=axis, out_dtype="float32")
@classmethod
@@ -377,10 +377,10 @@ class DequantizeLinear(OnnxOpConverter):
x, scale = inputs[0], inputs[1]
zp = inputs[2] if len(inputs) > 2 and inputs[2] is not None else None
axis = attr.get("axis", 1)
- if hasattr(x.struct_info, "ndim") and x.struct_info.ndim <= 1 and axis == 1:
+ if hasattr(x.ty, "ndim") and x.ty.ndim <= 1 and axis == 1:
axis = 0
if zp is None:
- zp = relax.const(0, x.struct_info.dtype)
+ zp = relax.const(0, x.ty.dtype)
return relax.op.dequantize(x, scale, zp, axis=axis, out_dtype="float32")
@@ -388,7 +388,7 @@ class DynamicQuantizeLinear(OnnxOpConverter):
@classmethod
def _impl_v11(cls, bb, inputs, attr, params):
x = inputs[0]
- x_dtype = x.struct_info.dtype
+ x_dtype = x.ty.dtype
qmin = relax.const(0, x_dtype)
qmax = relax.const(255, x_dtype)
@@ -420,22 +420,18 @@ class MatMulInteger16(OnnxOpConverter):
raise ValueError(f"MatMulInteger16 expects two inputs, but got {len(inputs)}")
a, b = inputs
valid_types = ["int16", "uint16"]
- if a.struct_info.dtype not in valid_types:
+ if a.ty.dtype not in valid_types:
raise ValueError(
"MatMulInteger16 expects input A to have int16 or uint16 dtype, "
- f"but got {a.struct_info.dtype}"
+ f"but got {a.ty.dtype}"
)
- if b.struct_info.dtype not in valid_types:
+ if b.ty.dtype not in valid_types:
raise ValueError(
"MatMulInteger16 expects input B to have int16 or uint16 dtype, "
- f"but got {b.struct_info.dtype}"
+ f"but got {b.ty.dtype}"
)
- out_dtype = (
- "uint32"
- if a.struct_info.dtype == "uint16" and b.struct_info.dtype == "uint16"
- else "int32"
- )
+ out_dtype = "uint32" if a.ty.dtype == "uint16" and b.ty.dtype == "uint16" else "int32"
return relax.op.matmul(
relax.op.astype(a, out_dtype),
relax.op.astype(b, out_dtype),
@@ -535,8 +531,8 @@ class Div(BinaryBase):
@classmethod
def _impl_v7(cls, bb, inputs, attr, params):
try:
- lhs_code = DataType(inputs[0].struct_info.dtype).type_code
- rhs_code = DataType(inputs[1].struct_info.dtype).type_code
+ lhs_code = DataType(inputs[0].ty.dtype).type_code
+ rhs_code = DataType(inputs[1].ty.dtype).type_code
except (AttributeError, ValueError, TypeError, RuntimeError):
return cls.base_impl(bb, inputs, attr, params)
@@ -682,10 +678,10 @@ class BitwiseBase(BinaryBase):
"""Base implementation for bitwise operations."""
valid_types = ["int8", "int16", "int32", "int64", "uint8", "uint16", "uint32", "uint64"]
for num, inp in enumerate(inputs):
- if inp.struct_info.dtype not in valid_types:
+ if inp.ty.dtype not in valid_types:
raise ValueError(
f"Bitwise operations expect all inputs to have integer types, "
- f"got {inp.struct_info.dtype} for input {num}"
+ f"got {inp.ty.dtype} for input {num}"
)
return super().base_impl(bb, inputs, attr, params)
@@ -729,7 +725,7 @@ class BitwiseNot(OnnxOpConverter):
@classmethod
def _impl_v18(cls, bb, inputs, attr, params):
if isinstance(inputs[0], relax.Constant):
- return relax.const(_np.bitwise_not(inputs[0].data.numpy()), inputs[0].struct_info.dtype)
+ return relax.const(_np.bitwise_not(inputs[0].data.numpy()), inputs[0].ty.dtype)
return relax.op.bitwise_not(inputs[0])
@@ -798,13 +794,13 @@ def _legacy_softmax_prepare(
return None
axis = _normalize_legacy_softmax_axis(axis, rank, op_name)
- struct_info = data.struct_info
- if not isinstance(struct_info, relax.TensorStructInfo):
+ ty = data.ty
+ if not isinstance(ty, relax.TensorType):
return None
- if not isinstance(struct_info.shape, relax.ShapeExpr):
+ if not isinstance(ty.shape, relax.ShapeExpr):
return None
- original_shape = list(struct_info.shape.values)
+ original_shape = list(ty.shape.values)
if len(original_shape) != rank:
return None
@@ -822,11 +818,9 @@ def _get_axis_extent(data: relax.Expr, axis: int, op_name: str) -> tuple[int, in
raise ValueError(f"{op_name} requires a statically known input rank.")
normalized_axis = _normalize_constant_axes([axis], rank, op_name)[0]
- struct_info = data.struct_info
- if isinstance(struct_info, relax.TensorStructInfo) and isinstance(
- struct_info.shape, relax.ShapeExpr
- ):
- axis_extent = struct_info.shape.values[normalized_axis]
+ ty = data.ty
+ if isinstance(ty, relax.TensorType) and isinstance(ty.shape, relax.ShapeExpr):
+ axis_extent = ty.shape.values[normalized_axis]
if isinstance(axis_extent, tirx.IntImm):
axis_extent = int(axis_extent.value)
return normalized_axis, axis_extent
@@ -908,7 +902,7 @@ class Hardmax(OnnxOpConverter):
if bb is not None:
data = bb.normalize(data)
normalized_axis, axis_extent = _get_axis_extent(data, axis, "Hardmax")
- dtype = data.struct_info.dtype
+ dtype = data.ty.dtype
argmax = relax.op.argmax(data, axis=normalized_axis)
on_value = relax.PrimValue(tvm.tirx.const(1.0, dtype))
off_value = relax.PrimValue(tvm.tirx.const(0.0, dtype))
@@ -949,10 +943,10 @@ class Transpose(OnnxOpConverter):
data = inputs[0]
axes = attr.get("perm", None)
- if hasattr(data.struct_info, "ndim"):
- input_ndim = data.struct_info.ndim
- elif hasattr(data.struct_info, "shape") and data.struct_info.shape:
- input_ndim = len(data.struct_info.shape)
+ if hasattr(data.ty, "ndim"):
+ input_ndim = data.ty.ndim
+ elif hasattr(data.ty, "shape") and data.ty.shape:
+ input_ndim = len(data.ty.shape)
else:
if isinstance(data, relax.Constant):
input_ndim = data.data.numpy().ndim
@@ -1015,7 +1009,7 @@ class Unsqueeze(OnnxOpConverter):
else:
new_shape.append(next(input_dims_iter))
expanded = expanded.reshape(new_shape)
- return relax.const(expanded, data.struct_info.dtype)
+ return relax.const(expanded, data.ty.dtype)
if isinstance(axes, relax.Constant):
if data_ndim is None:
@@ -1056,7 +1050,7 @@ class Concat(OnnxOpConverter):
if isinstance(x, relax.ShapeExpr):
return True
elif isinstance(x, relax.Constant):
- return x.struct_info.ndim == 1 and x.struct_info.dtype == "int64"
+ return x.ty.ndim == 1 and x.ty.dtype == "int64"
else:
return False
@@ -1090,7 +1084,7 @@ class Concat(OnnxOpConverter):
for inp in inputs:
const_inputs.append(inp.data.numpy())
out = _np.concatenate(const_inputs, axis=axis)
- dtype = inputs[0].struct_info.dtype
+ dtype = inputs[0].ty.dtype
return relax.const(out, dtype)
return relax.op.concat(inputs, axis=axis)
@@ -1120,7 +1114,7 @@ class Cast(OnnxOpConverter):
if np_dst.kind in ("i", "u"):
src = inputs[0]
- src_dtype = getattr(getattr(src, "struct_info", None), "dtype", None) or getattr(
+ src_dtype = getattr(getattr(src, "ty", None), "dtype", None) or getattr(
src, "dtype", None
)
if src_dtype is not None and _relax_dtype_is_floating_point(src_dtype):
@@ -1199,7 +1193,7 @@ class Gather(OnnxOpConverter):
shape_val = data[np_index]
return relax.PrimValue(shape_val)
- indices_dtype = indices.struct_info.dtype
+ indices_dtype = indices.ty.dtype
if not indices_dtype.startswith("uint"):
data_shape = bb.normalize(relax.op.shape_of(data))
data_shape_tensor = bb.normalize(relax.op.shape_to_tensor(data_shape))
@@ -1320,13 +1314,13 @@ class Compress(OnnxOpConverter):
axis = attr.get("axis", None)
# Change one hot tensor to indices e.g. [0, 1, 1, 0, 1] -> [1, 2, 4]
- if condition.struct_info.dtype != "bool":
+ if condition.ty.dtype != "bool":
raise ValueError("Condition tensor is expected to be a boolean tensor")
- if condition.struct_info.ndim != 1:
+ if condition.ty.ndim != 1:
raise ValueError("Condition tensor is expected to be a 1D boolean tensor")
indices = relax.op.nonzero(condition)
num_nonzero = tirx.Var("num_nonzero", "int64")
- indices = bb.match_cast(indices, relax.TensorStructInfo([1, num_nonzero], "int64"))
+ indices = bb.match_cast(indices, relax.TensorType([1, num_nonzero], "int64"))
indices = relax.op.reshape(indices, [-1])
if axis is not None:
@@ -1351,7 +1345,7 @@ class EyeLike(OnnxOpConverter):
@classmethod
def _impl_v9(cls, bb, inputs, attr, params):
k = attr.get("k", 0)
- input_dtype = inputs[0].struct_info.dtype
+ input_dtype = inputs[0].ty.dtype
if "dtype" in attr and get_type(attr["dtype"]) != input_dtype:
raise ValueError(
f"dtype mismatch between input ({input_dtype}) and attribute ({attr['dtype']})"
@@ -1371,7 +1365,7 @@ class Gemm(OnnxOpConverter):
A = inputs[0]
B = inputs[1]
C = inputs[2]
- dtype = A.struct_info.dtype
+ dtype = A.ty.dtype
# Compute Y = alpha * A X B + beta * C
@@ -1444,7 +1438,7 @@ class Clip(OnnxOpConverter):
@staticmethod
def _sanitize_nan_clip_bound(bb, bound: relax.Expr, *, for_min: bool) -> relax.Expr:
"""ONNX/ORT treat NaN clip bounds as unbounded; plain max/min with NaN poisons output."""
- dtype = bound.struct_info.dtype
+ dtype = bound.ty.dtype
if not _relax_dtype_is_floating_point(dtype):
return bound
repl = -_np.inf if for_min else _np.inf
@@ -1482,14 +1476,14 @@ class Shape(OnnxOpConverter):
@classmethod
def _impl_v13(cls, bb, inputs, attr, params):
- data_info = inputs[0].struct_info
+ data_info = inputs[0].ty
- if isinstance(data_info, relax.ShapeStructInfo):
+ if isinstance(data_info, relax.ShapeType):
if data_info.ndim == -1:
raise ValueError("The ndim of ShapeExpr is expected to a real number, but got -1.")
return relax.ShapeExpr([data_info.ndim])
- # If no shape is defined in the struct info, it must be computed at runtime.
+ # If no shape is defined in the type, it must be computed at runtime.
if not data_info.shape:
data_shape = bb.normalize(relax.op.shape_of(inputs[0]))
return data_shape
@@ -1563,7 +1557,7 @@ class Mish(OnnxOpConverter):
@classmethod
def _impl_v18(cls, bb, inputs, attr, params):
- dtype = inputs[0].struct_info.dtype
+ dtype = inputs[0].ty.dtype
return inputs[0] * relax.op.tanh(
relax.op.log(relax.const(1.0, dtype) + relax.op.exp(inputs[0]))
)
@@ -1580,8 +1574,8 @@ class PRelu(OnnxOpConverter):
x = inputs[0]
slope = inputs[1]
- x_shape = x.struct_info.shape
- slope_shape = slope.struct_info.shape
+ x_shape = x.ty.shape
+ slope_shape = slope.ty.shape
ndim = len(x_shape)
s_ndim = len(slope_shape)
@@ -1673,12 +1667,12 @@ class FastGelu(OnnxOpConverter):
x = inputs[0]
if len(inputs) > 1 and inputs[1] is not None:
bias = inputs[1]
- bias_shape = bias.struct_info.shape
+ bias_shape = bias.ty.shape
assert len(bias_shape) == 1, "bias term must be a 1D tensor"
x = bb.emit(relax.op.add(x, bias))
# Declare consts
- const_dtype = x.struct_info.dtype
+ const_dtype = x.ty.dtype
half = relax.const(0.5, dtype=const_dtype)
one = relax.const(1.0, dtype=const_dtype)
const1 = relax.const(math.sqrt(2 / math.pi), dtype=const_dtype)
@@ -1715,7 +1709,7 @@ class Shrink(OnnxOpConverter):
@classmethod
def _impl_v9(cls, bb, inputs, attr, params):
x = inputs[0]
- dtype = x.struct_info.dtype
+ dtype = x.ty.dtype
lambd = relax.const(attr.get("lambd", 0.5), dtype)
bias = relax.const(attr.get("bias", 0.0), dtype)
zeros = relax.op.zeros_like(x)
@@ -1730,13 +1724,13 @@ class Conv(OnnxOpConverter):
@classmethod
def _impl_v11(cls, bb, inputs, attr, params):
data = inputs[0]
- if hasattr(inputs[0].struct_info, "ndim"):
- ndim = inputs[0].struct_info.ndim
+ if hasattr(inputs[0].ty, "ndim"):
+ ndim = inputs[0].ty.ndim
else:
- ndim = len(inputs[0].struct_info.shape)
+ ndim = len(inputs[0].ty.shape)
if "kernel_shape" not in attr:
- attr["kernel_shape"] = inputs[1].struct_info.shape.values[2:]
+ attr["kernel_shape"] = inputs[1].ty.shape.values[2:]
if ndim == 3:
op = relax.op.nn.conv1d
@@ -1800,10 +1794,10 @@ class ConvTranspose(OnnxOpConverter):
@classmethod
def _impl_v1(cls, bb, inputs, attr, params):
- if hasattr(inputs[0].struct_info, "ndim"):
- ndim = inputs[0].struct_info.ndim
+ if hasattr(inputs[0].ty, "ndim"):
+ ndim = inputs[0].ty.ndim
else:
- ndim = len(inputs[0].struct_info.shape)
+ ndim = len(inputs[0].ty.shape)
if ndim == 3:
op = relax.op.nn.conv1d_transpose
@@ -1827,7 +1821,7 @@ class ConvTranspose(OnnxOpConverter):
if "kernel_shape" in attr:
kernel_shape = list(attr["kernel_shape"])
else:
- kernel_shape = [int(s) for s in inputs[1].struct_info.shape.values[2:]]
+ kernel_shape = [int(s) for s in inputs[1].ty.shape.values[2:]]
# Resolve `auto_pad` per ONNX ConvTranspose spec. Unlike Conv, the spec
# derives `pads` from `output_shape`/`strides` when auto_pad is SAME_*,
@@ -1917,9 +1911,7 @@ class CumSum(OnnxOpConverter):
f"got shape {axis_data.shape}"
)
elif isinstance(axis_input, relax.Var):
- axis_shape = (
- axis_input.struct_info.shape if hasattr(axis_input.struct_info, "shape") else None
- )
+ axis_shape = axis_input.ty.shape if hasattr(axis_input.ty, "shape") else None
raise ValueError(
"CumSum with non-constant axis input is not supported yet. "
"ONNX permits runtime axis tensors, but Relax/TE currently requires a compile-time "
@@ -1957,7 +1949,7 @@ class Squeeze(OnnxOpConverter):
else:
raise NotImplementedError("Squeeze with symbolic axes not supported")
- return relax.const(out_data, data.struct_info.dtype)
+ return relax.const(out_data, data.ty.dtype)
if isinstance(data, relax.ShapeExpr):
shape_tensor_ndim = 1
@@ -2146,7 +2138,7 @@ class Neg(OnnxOpConverter):
def _impl_v13(cls, bb, inputs, attr, params):
if isinstance(inputs[0], relax.Constant):
data_np = inputs[0].data.numpy()
- return relax.const(_np.negative(data_np), inputs[0].struct_info.dtype)
+ return relax.const(_np.negative(data_np), inputs[0].ty.dtype)
if isinstance(inputs[0], relax.PrimValue):
return relax.PrimValue(-inputs[0].value)
return relax.op.negative(inputs[0])
@@ -2168,7 +2160,7 @@ class Reciprocal(OnnxOpConverter):
@classmethod
def _impl_v13(cls, bb, inputs, attr, params):
- input_dtype = inputs[0].struct_info.dtype
+ input_dtype = inputs[0].ty.dtype
return relax.op.divide(relax.const(1, dtype=input_dtype), inputs[0])
@@ -2251,7 +2243,7 @@ class MultiInputBase(OnnxOpConverter):
output = cls.numpy_op(*np_inputs) # pylint: disable=not-callable
return relax.const(output, output.dtype)
- input_shapes = [inp.struct_info.shape for inp in inputs]
+ input_shapes = [inp.ty.shape for inp in inputs]
target_shape = functools.reduce(compute_broadcast_shape, input_shapes)
# broadcast_to, stack them, then perform minimum over the new axis.
@@ -2294,7 +2286,7 @@ class Log(OnnxOpConverter):
@classmethod
def _impl_v13(cls, bb, inputs, attr, params):
if isinstance(inputs[0], relax.Constant):
- return relax.const(_np.log(inputs[0].data.numpy()), inputs[0].struct_info.dtype)
+ return relax.const(_np.log(inputs[0].data.numpy()), inputs[0].ty.dtype)
return relax.op.log(inputs[0])
@@ -2309,7 +2301,7 @@ class Exp(OnnxOpConverter):
def _impl_v1(cls, bb, inputs, attr, params):
data = inputs[0]
valid_types = ["float", "float32", "double", "float64", "float16"]
- cls._check_type(data.struct_info.dtype, valid_types)
+ cls._check_type(data.ty.dtype, valid_types)
return relax.op.exp(data)
@@ -2317,7 +2309,7 @@ class Exp(OnnxOpConverter):
def _impl_v13(cls, bb, inputs, attr, params):
data = inputs[0]
valid_types = ["float", "float32", "double", "float64", "float16", "bfloat16"]
- cls._check_type(data.struct_info.dtype, valid_types)
+ cls._check_type(data.ty.dtype, valid_types)
return relax.op.exp(data)
@@ -2327,7 +2319,7 @@ class Softplus(OnnxOpConverter):
@classmethod
def _impl_v1(cls, bb, inputs, attr, params):
- dtype = inputs[0].struct_info.dtype
+ dtype = inputs[0].ty.dtype
threshold = 10.0 if dtype == "float16" else 20.0
return relax.op.nn.softplus(inputs[0], threshold=threshold)
@@ -2337,7 +2329,7 @@ class Softsign(OnnxOpConverter):
@classmethod
def _impl_v1(cls, bb, inputs, attr, params):
- dtype = inputs[0].struct_info.dtype
+ dtype = inputs[0].ty.dtype
return inputs[0] / (relax.op.abs(inputs[0]) + relax.const(1, dtype=dtype))
@@ -2363,7 +2355,7 @@ class Split(OnnxOpConverter):
splits = inputs[1]
splits_rank = None
if splits is not None:
- splits_rank = splits.struct_info.ndim
+ splits_rank = splits.ty.ndim
if splits is not None and splits_rank > 0:
if isinstance(splits, relax.Constant):
splits = splits.data.numpy()
@@ -2399,9 +2391,9 @@ def _get_known_tensor_rank(expr: relax.Expr) -> int | None:
return 1
if isinstance(expr, relax.PrimValue):
return 0
- struct_info = expr.struct_info
- if isinstance(struct_info, relax.TensorStructInfo):
- return None if struct_info.ndim == -1 else struct_info.ndim
+ ty = expr.ty
+ if isinstance(ty, relax.TensorType):
+ return None if ty.ndim == -1 else ty.ndim
return None
@@ -2419,15 +2411,15 @@ def _get_known_tensor_length(expr: relax.Expr | None) -> int | None:
return len(expr.values)
if isinstance(expr, relax.PrimValue):
return 1
- struct_info = expr.struct_info
- if not isinstance(struct_info, relax.TensorStructInfo):
+ ty = expr.ty
+ if not isinstance(ty, relax.TensorType):
return None
- if struct_info.ndim == -1:
+ if ty.ndim == -1:
return None
- if struct_info.ndim != 1:
- raise ValueError(f"Expected a 1-D tensor, but got ndim={struct_info.ndim}.")
- if isinstance(struct_info.shape, relax.ShapeExpr):
- dim = struct_info.shape.values[0]
+ if ty.ndim != 1:
+ raise ValueError(f"Expected a 1-D tensor, but got ndim={ty.ndim}.")
+ if isinstance(ty.shape, relax.ShapeExpr):
+ dim = ty.shape.values[0]
if isinstance(dim, tirx.IntImm):
return int(dim.value)
if isinstance(dim, int):
@@ -2459,10 +2451,10 @@ def _as_int64_tensor(bb: relax.BlockBuilder, expr: relax.Expr) -> relax.Expr:
if isinstance(expr, relax.PrimValue):
return bb.normalize(relax.op.full((1,), expr, dtype="int64"))
if isinstance(expr, relax.Constant):
- if expr.struct_info.dtype == "int64":
+ if expr.ty.dtype == "int64":
return expr
return bb.normalize(relax.op.astype(expr, "int64"))
- if isinstance(expr.struct_info, relax.TensorStructInfo) and expr.struct_info.dtype != "int64":
+ if isinstance(expr.ty, relax.TensorType) and expr.ty.dtype != "int64":
return bb.normalize(relax.op.astype(expr, "int64"))
return expr
@@ -2472,10 +2464,10 @@ def _tensor_to_shape_expr(
) -> relax.ShapeExpr:
"""Convert a statically sized int64 tensor into a ShapeExpr."""
- shape_tensor = bb.match_cast(shape_tensor, relax.TensorStructInfo([shape_ndim], "int64"))
+ shape_tensor = bb.match_cast(shape_tensor, relax.TensorType([shape_ndim], "int64"))
shape_dataflow_var = bb.emit(relax.op.tensor_to_shape(shape_tensor))
shape_vars = [tirx.Var(f"{prefix}_{i}", "int64") for i in range(shape_ndim)]
- bb.match_cast(shape_dataflow_var, relax.ShapeStructInfo(shape_vars))
+ bb.match_cast(shape_dataflow_var, relax.ShapeType(shape_vars))
return relax.ShapeExpr(shape_vars)
@@ -2546,7 +2538,7 @@ def _build_squeezed_shape_tensor(
keep_mask = bb.normalize(relax.op.equal(remove_mask, relax.const(0, "int64")))
keep_indices = bb.normalize(relax.op.nonzero(keep_mask))
num_keep_dims = tirx.Var("squeeze_num_keep_dims", "int64")
- keep_indices = bb.match_cast(keep_indices, relax.TensorStructInfo([1, num_keep_dims], "int64"))
+ keep_indices = bb.match_cast(keep_indices, relax.TensorType([1, num_keep_dims], "int64"))
keep_indices = bb.normalize(relax.op.reshape(keep_indices, [-1]))
return bb.normalize(relax.op.take(data_shape_tensor, keep_indices, axis=0))
@@ -2701,7 +2693,7 @@ class Pad(OnnxOpConverter):
@classmethod
def _impl_v2(cls, bb, inputs, attr, params):
pads = attr.get("pads")
- pads = relax.const(_np.array(pads), inputs[0].struct_info.shape[0].dtype)
+ pads = relax.const(_np.array(pads), inputs[0].ty.shape[0].dtype)
constant_value = attr.get("value")
if constant_value is None:
constant_value = 0.0
@@ -2763,7 +2755,7 @@ class Tile(OnnxOpConverter):
@staticmethod
def _tensor_length(expr):
- shape = expr.struct_info.shape
+ shape = expr.ty.shape
if not isinstance(shape, relax.ShapeExpr):
return None
@@ -2780,12 +2772,12 @@ class Tile(OnnxOpConverter):
return bb.emit_te(topi.tile, inputs[0], reps)
data = inputs[0]
- data_ndim = data.struct_info.ndim
+ data_ndim = data.ty.ndim
reps_len = cls._tensor_length(reps)
if data_ndim == -1 or reps_len is None:
raise ValueError("Dynamic Tile requires known input rank and repeats length.")
- if reps.struct_info.dtype != "int64":
+ if reps.ty.dtype != "int64":
reps = bb.normalize(relax.op.astype(reps, "int64"))
data_shape = bb.normalize(relax.op.shape_of(data))
@@ -2810,7 +2802,7 @@ class Tile(OnnxOpConverter):
output_shape_vars = [
tirx.Var(f"tile_dim_{i}", "int64") for i in range(max(data_ndim, reps_len))
]
- bb.match_cast(output_shape, relax.ShapeStructInfo(output_shape_vars))
+ bb.match_cast(output_shape, relax.ShapeType(output_shape_vars))
return bb.emit_te(topi.dyn_tile, data, output_shape_vars, reps_len)
@@ -2822,7 +2814,7 @@ class Expand(OnnxOpConverter):
data = inputs[0]
shape = inputs[1]
if isinstance(shape, relax.ShapeExpr):
- data_shape = list(data.struct_info.shape)
+ data_shape = list(data.ty.shape)
target_shape = list(shape.values)
original_data_shape = [
dim.value if hasattr(dim, "value") else str(dim) for dim in data_shape
@@ -2875,7 +2867,7 @@ class Expand(OnnxOpConverter):
new_shape = shape.data.numpy().tolist()
# ONNX Expand operator requires preserving target rank and broadcasting
# according to standard rules. Dimensions are right-aligned.
- data_shape = [dim.value for dim in data.struct_info.shape]
+ data_shape = [dim.value for dim in data.ty.shape]
original_data_shape = data_shape.copy()
original_new_shape = new_shape.copy()
@@ -2916,22 +2908,22 @@ class Expand(OnnxOpConverter):
return relax.op.broadcast_to(data, relax.ShapeExpr(new_shape))
# Otherwise handle dynamic shapes.
- shape_ndim = next(dim.value for dim in shape.struct_info.shape.values)
+ shape_ndim = next(dim.value for dim in shape.ty.shape.values)
shape_dataflow_var = bb.emit(
relax.Call(
relax.ExternFunc("vm.builtin.tensor_to_shape"),
[shape],
- sinfo_args=[relax.ShapeStructInfo(ndim=shape_ndim)],
+ ty_args=[relax.ShapeType(ndim=shape_ndim)],
)
)
shape_vars = []
for i in range(shape_ndim):
shape_vars.append(tvm.tirx.Var(f"x_{i}", "int64"))
- bb.match_cast(shape_dataflow_var, relax.ShapeStructInfo(shape_vars))
+ bb.match_cast(shape_dataflow_var, relax.ShapeType(shape_vars))
# Applying broadcasting rules for dynamic shapes
- data_shape = list(data.struct_info.shape)
+ data_shape = list(data.ty.shape)
data_ndim = len(data_shape)
target_ndim = shape_ndim
padded_data = data
@@ -2993,10 +2985,8 @@ class Attention(OnnxOpConverter):
assert inputs[6] is None, "past_sequence_length is not currently supported"
- (batch_size, seq_len, input_hidden_size) = [
- val.value for val in input_emb.struct_info.shape.values
- ]
- weight_shape = [val.value for val in weight.struct_info.shape.values]
+ (batch_size, seq_len, input_hidden_size) = [val.value for val in input_emb.ty.shape.values]
+ weight_shape = [val.value for val in weight.ty.shape.values]
assert weight_shape[0] == input_hidden_size, (
"input and weight should share the same input hiden size"
@@ -3017,7 +3007,7 @@ class Attention(OnnxOpConverter):
head_size_v = hidden_size_v // num_heads
if mask_index is not None:
- mask_index_shape = [val.value for val in mask_index.struct_info.shape.values]
+ mask_index_shape = [val.value for val in mask_index.ty.shape.values]
assert mask_index_shape in (
[batch_size, seq_len],
[
@@ -3027,14 +3017,12 @@ class Attention(OnnxOpConverter):
],
), """mask index should be in shape of (batch_size, seq_len),
or (batch_size, seq_len, seq_len)"""
- mask_bias = relax.op.subtract(
- relax.const(1, dtype=mask_index.struct_info.dtype), mask_index
- )
- mask_bias = relax.op.astype(mask_bias, dtype=input_emb.struct_info.dtype)
+ mask_bias = relax.op.subtract(relax.const(1, dtype=mask_index.ty.dtype), mask_index)
+ mask_bias = relax.op.astype(mask_bias, dtype=input_emb.ty.dtype)
mask_bias = bb.normalize(
relax.op.multiply(
mask_bias,
- relax.const(mask_filter_value, dtype=input_emb.struct_info.dtype),
+ relax.const(mask_filter_value, dtype=input_emb.ty.dtype),
)
)
if qk_bias is None:
@@ -3053,7 +3041,7 @@ class Attention(OnnxOpConverter):
QKV = relax.op.matmul(input_emb, weight)
if bias:
- bias_shape = [val.value for val in bias.struct_info.shape.values]
+ bias_shape = [val.value for val in bias.ty.shape.values]
assert bias_shape[0] == weight_shape[1], (
"bias and weight should share the same hidden size sum"
)
@@ -3206,7 +3194,7 @@ class Resize(OnnxOpConverter):
roi = get_constant(inputs[1], params) if len(inputs) > 1 and inputs[1] is not None else None
scales = get_constant(inputs[2], params) if len(inputs) > 2 else None
sizes = get_constant(inputs[3], params) if len(inputs) > 3 else None
- ndims = len(x.struct_info.shape)
+ ndims = len(x.ty.shape)
assert ndims in (3, 4, 5), "Only resize1d/resize2d/resize3d are supported."
assert scales is None or sizes is None, (
@@ -3245,7 +3233,7 @@ class Resize(OnnxOpConverter):
raise ValueError(f"Type {type(scales)} for scale is currently unsupported.")
sizes = []
- for i, dim in enumerate(x.struct_info.shape):
+ for i, dim in enumerate(x.ty.shape):
sizes.append(cast(scales[i] * dim, "int64"))
sizes = sizes[2:]
else:
@@ -3365,7 +3353,7 @@ class RoiAlign(OnnxOpConverter):
data = inputs[0]
rois = inputs[1]
batch_indices = inputs[2]
- rois_dtype = rois.struct_info.dtype
+ rois_dtype = rois.ty.dtype
mode = attr.get("mode", b"avg")
if isinstance(mode, bytes):
@@ -3449,7 +3437,7 @@ class Range(OnnxOpConverter):
start = get_constant(inputs[0], params)
limit = get_constant(inputs[1], params)
delta = get_constant(inputs[2], params)
- out_dtype = start.struct_info.dtype
+ out_dtype = start.ty.dtype
if isinstance(start, relax.Constant):
start = start.data.numpy().tolist()
@@ -3478,9 +3466,9 @@ class InstanceNormalization(OnnxOpConverter):
scale = inputs[1]
B = inputs[2]
epsilon = attr.get("epsilon", 1e-05)
- epsilon = relax.const(epsilon, dtype=data.struct_info.dtype)
+ epsilon = relax.const(epsilon, dtype=data.ty.dtype)
- ndim = len(data.struct_info.shape)
+ ndim = len(data.ty.shape)
redux_axes = list(range(2, ndim))
mean = relax.op.mean(data, axis=redux_axes, keepdims=True)
@@ -3551,10 +3539,10 @@ class LocalResponseNormalization(OnnxOpConverter):
beta = attr.get("beta", 0.75)
bias = attr.get("bias", 1.0)
- if hasattr(data.struct_info, "ndim"):
- ndim = data.struct_info.ndim
+ if hasattr(data.ty, "ndim"):
+ ndim = data.ty.ndim
else:
- ndim = len(data.struct_info.shape)
+ ndim = len(data.ty.shape)
if ndim not in [3, 4]:
raise ValueError(f"LRN only supports 3D or 4D input, got {ndim}D.")
@@ -3619,7 +3607,7 @@ class Pool(OnnxOpConverter):
def _impl_v1(cls, bb, inputs, attr, params):
# Unpack inputs and attributes.
data = inputs[0]
- input_shape = data.struct_info.shape
+ input_shape = data.ty.shape
ndim = len(input_shape)
auto_pad = attr.get("auto_pad", b"NOTSET").decode("utf-8")
@@ -3678,7 +3666,7 @@ class Pool(OnnxOpConverter):
@classmethod
def _get_input_spatial_shape(cls, tensor):
# shape is (N x C x D1 x D2 ... Dn)
- return _np.array([int(d) for d in tensor.struct_info.shape], dtype="int64")[2:]
+ return _np.array([int(d) for d in tensor.ty.shape], dtype="int64")[2:]
class MaxPool(Pool):
@@ -3698,10 +3686,10 @@ class LpPool(OnnxOpConverter):
@classmethod
def _impl_v1(cls, bb, inputs, attr, params):
- dtype = inputs[0].struct_info.dtype
+ dtype = inputs[0].ty.dtype
p = attr.get("p", 2.0)
reci_p = relax.const(1.0 / p, dtype=dtype)
- # emit for get struct_info
+ # emit for get ty
data = bb.emit(relax.op.power(inputs[0], relax.const(p, dtype=dtype)))
attr.update({"count_include_pad": True})
avg_pool = AveragePool._impl_v1(bb, [data], attr, params)
@@ -3715,7 +3703,7 @@ class GlobalAveragePool(OnnxOpConverter):
@classmethod
def _impl_v1(cls, bb, inputs, attr, params):
- rank = len(inputs[0].struct_info.shape)
+ rank = len(inputs[0].ty.shape)
axes = list(range(2, rank))
return relax.op.mean(inputs[0], axis=axes, keepdims=True)
@@ -3725,7 +3713,7 @@ class GlobalMaxPool(OnnxOpConverter):
@classmethod
def _impl_v1(cls, bb, inputs, attr, params):
- rank = len(inputs[0].struct_info.shape)
+ rank = len(inputs[0].ty.shape)
axes = list(range(2, rank))
return relax.op.max(inputs[0], axis=axes, keepdims=True)
@@ -3736,8 +3724,8 @@ class GlobalLpPool(OnnxOpConverter):
@classmethod
def _impl_v2(cls, bb, inputs, attr, params):
p = attr.get("p", 2.0)
- dtype = inputs[0].struct_info.dtype
- rank = len(inputs[0].struct_info.shape)
+ dtype = inputs[0].ty.dtype
+ rank = len(inputs[0].ty.shape)
axes = list(range(2, rank))
x_abs = relax.op.abs(inputs[0])
x_p = relax.op.power(x_abs, relax.const(p, dtype=dtype))
@@ -3758,7 +3746,7 @@ class MaxUnpool(OnnxOpConverter):
strides = attr.get("strides", [1] * len(kernel_shape))
multiplier = _np.concatenate([[1, 1], list(strides)])
- shape = [v.value for v in data.struct_info.shape]
+ shape = [v.value for v in data.ty.shape]
total_output_shape = multiplier * shape
# Add extra dimensions from kernel size and stride mismatch
total_output_shape += _np.concatenate([[0, 0], list(kernel_shape)], axis=0)
@@ -3778,7 +3766,7 @@ class MaxUnpool(OnnxOpConverter):
# Create a tensor of zeros then scatter our data through it.
relax_shape = relax.ShapeExpr(total_output_shape.tolist())
- zeros_tensor = bb.emit(relax.op.zeros(relax_shape, data.struct_info.dtype))
+ zeros_tensor = bb.emit(relax.op.zeros(relax_shape, data.ty.dtype))
# We need to flatten all our tensors before scattering.
flat_tensor = relax.op.scatter_elements(
relax.op.reshape(zeros_tensor, [-1]),
@@ -3797,7 +3785,7 @@ class Flatten(OnnxOpConverter):
@classmethod
def _impl_v13(cls, bb, inputs, attr, params):
axis = attr.get("axis", 1)
- data_shape = list(inputs[0].struct_info.shape)
+ data_shape = list(inputs[0].ty.shape)
if axis == 0:
new_shape = (1, -1)
@@ -3829,12 +3817,12 @@ class LayerNormalization(OnnxOpConverter):
axis = attr.get("axis", -1)
epsilon = attr.get("epsilon", 1e-05)
- gamma_shape = get_const_tuple(scale.struct_info.shape)
+ gamma_shape = get_const_tuple(scale.ty.shape)
if bias is None:
- bias = relax.const(_np.zeros(gamma_shape), dtype=scale.struct_info.dtype)
+ bias = relax.const(_np.zeros(gamma_shape), dtype=scale.ty.dtype)
else:
- beta_shape = get_const_tuple(bias.struct_info.shape)
+ beta_shape = get_const_tuple(bias.ty.shape)
if gamma_shape != beta_shape:
raise ValueError("gamma and beta shapes do not match")
@@ -3868,7 +3856,7 @@ class RMSNormalization(OnnxOpConverter):
axes = list(range(axis, ndim))
# If stash_type requires float32 computation and input is not float32, cast
- input_dtype = data.struct_info.dtype
+ input_dtype = data.ty.dtype
if stash_type == 1 and input_dtype != "float32":
data_compute = relax.op.astype(data, "float32")
scale_compute = relax.op.astype(scale, "float32")
@@ -4249,7 +4237,7 @@ def _argreduce_select_last_index(bb, data, axis, keepdims, op):
"""
data_flipped = relax.op.flip(data, axis=axis)
flipped_idx = bb.normalize(op(data_flipped, axis, keepdims))
- axis_size = data.struct_info.shape[axis]
+ axis_size = data.ty.shape[axis]
if isinstance(axis_size, tirx.IntImm):
offset = relax.const(int(axis_size) - 1, "int64")
else:
@@ -4269,7 +4257,7 @@ class ArgMax(OnnxOpConverter):
@classmethod
def _check_attrs(cls, data, attr, shift_axis=True):
- dims_num = len(data.struct_info.shape)
+ dims_num = len(data.ty.shape)
axis = attr.get("axis", 0)
if shift_axis and axis < 0:
axis += dims_num
@@ -4304,7 +4292,7 @@ class ArgMin(OnnxOpConverter):
@classmethod
def _check_attrs(cls, data, attr, shift_axis=True):
- dims_num = len(data.struct_info.shape)
+ dims_num = len(data.ty.shape)
axis = attr.get("axis", 0)
if shift_axis and axis < 0:
axis += dims_num
@@ -4405,7 +4393,7 @@ class EmbedLayerNormalization(OnnxOpConverter):
epsilon = attr.get("epsilon", 1e-12)
- (batch_size, seq_len) = [dim.value for dim in input_ids.struct_info.shape]
+ (batch_size, seq_len) = [dim.value for dim in input_ids.ty.shape]
if segment_ids:
assert segment_emb
@@ -4473,8 +4461,8 @@ class Unique(OnnxOpConverter):
)
unique_numbers = tirx.Var("unique_numbers", "int64")
- input_shape = data.struct_info.shape
- dtype = data.struct_info.dtype
+ input_shape = data.ty.shape
+ dtype = data.ty.dtype
if axis is None:
output_shape = (unique_numbers,)
@@ -4487,15 +4475,15 @@ class Unique(OnnxOpConverter):
]
if num_outputs == 1:
- return bb.match_cast(unique, relax.TensorStructInfo(output_shape, dtype))
+ return bb.match_cast(unique, relax.TensorType(output_shape, dtype))
- outputs = [bb.match_cast(unique[0], relax.TensorStructInfo(output_shape, dtype))]
+ outputs = [bb.match_cast(unique[0], relax.TensorType(output_shape, dtype))]
tuple_idx = 1 # Track which index in the tuple we're at
if return_index:
index_shape = (unique_numbers,)
- index_sinfo = relax.TensorStructInfo(index_shape, "int64")
- outputs.append(bb.match_cast(unique[tuple_idx], index_sinfo))
+ index_ty = relax.TensorType(index_shape, "int64")
+ outputs.append(bb.match_cast(unique[tuple_idx], index_ty))
tuple_idx += 1
if return_inverse:
@@ -4503,14 +4491,14 @@ class Unique(OnnxOpConverter):
# When axis is None: shape is [X.size]
# When axis is specified: shape is [X.shape[axis]]
inverse_shape = (tirx.Var("inverse_numbers", "int64"),)
- inverse_sinfo = relax.TensorStructInfo(inverse_shape, "int64")
- outputs.append(bb.match_cast(unique[tuple_idx], inverse_sinfo))
+ inverse_ty = relax.TensorType(inverse_shape, "int64")
+ outputs.append(bb.match_cast(unique[tuple_idx], inverse_ty))
tuple_idx += 1
if return_counts:
count_shape = (unique_numbers,)
- count_sinfo = relax.TensorStructInfo(count_shape, "int64")
- outputs.append(bb.match_cast(unique[tuple_idx], count_sinfo))
+ count_ty = relax.TensorType(count_shape, "int64")
+ outputs.append(bb.match_cast(unique[tuple_idx], count_ty))
return relax.Tuple(outputs)
@@ -4520,11 +4508,11 @@ class NonZero(OnnxOpConverter):
@classmethod
def _impl_v9(cls, bb, inputs, attr, params):
- ndim = inputs[0].struct_info.ndim
+ ndim = inputs[0].ty.ndim
ndim = 1 if ndim == 0 else ndim
nonzero_numbers = tirx.Var("nonzero_numbers", "int64")
return bb.match_cast(
- relax.op.nonzero(inputs[0]), relax.TensorStructInfo((ndim, nonzero_numbers), "int64")
+ relax.op.nonzero(inputs[0]), relax.TensorType((ndim, nonzero_numbers), "int64")
)
@@ -4537,7 +4525,7 @@ class Upsample(OnnxOpConverter):
assert len(scales) == 4
assert scales[0] == scales[1] == 1
- inp_shape = [int(x) for x in inputs[0].struct_info.shape]
+ inp_shape = [int(x) for x in inputs[0].ty.shape]
assert len(inp_shape) == 4
out_shape2d = [int(dim * scale) for dim, scale in zip(inp_shape[2:], scales[2:])]
@@ -4563,7 +4551,7 @@ class HardSigmoid(OnnxOpConverter):
@classmethod
def _impl_v1(cls, bb, inputs, attr, params):
x = inputs[0]
- dtype = x.struct_info.dtype
+ dtype = x.ty.dtype
alpha = float(attr.get("alpha", 0.2))
alpha = relax.const(alpha, dtype=dtype)
beta = float(attr.get("beta", 0.5))
@@ -4577,7 +4565,7 @@ class HardSwish(OnnxOpConverter):
@classmethod
def _impl_v14(cls, bb, inputs, attr, params):
x = inputs[0]
- dtype = x.struct_info.dtype
+ dtype = x.ty.dtype
return relax.op.multiply(
x,
relax.op.divide(
@@ -4610,7 +4598,7 @@ class DepthToSpace(OnnxOpConverter):
def _impl_v11(cls, bb, inputs, attr, params):
block_size = int(attr["blocksize"])
mode = attr.get("mode", b"DCR").decode("utf-8")
- b, c, h, w = inputs[0].struct_info.shape
+ b, c, h, w = inputs[0].ty.shape
if mode == "DCR":
x = relax.op.reshape(inputs[0], (b, block_size, block_size, c // (block_size**2), h, w))
x = relax.op.permute_dims(x, [0, 3, 4, 1, 5, 2])
@@ -4629,7 +4617,7 @@ class SpaceToDepth(OnnxOpConverter):
@classmethod
def _impl_v1(cls, bb, inputs, attr, params):
block_size = int(attr["blocksize"])
- b, c, h, w = inputs[0].struct_info.shape
+ b, c, h, w = inputs[0].ty.shape
x = relax.op.reshape(
inputs[0], (b, c, h // block_size, block_size, w // block_size, block_size)
)
@@ -4800,7 +4788,7 @@ class SplitToSequence(OnnxOpConverter):
keepdims = attr.get("keepdims", 1)
input_tensor = inputs[0]
- input_shape = input_tensor.struct_info.shape
+ input_shape = input_tensor.ty.shape
if len(inputs) == 1:
split = _np.array(1)
@@ -4831,7 +4819,7 @@ class SplitToSequence(OnnxOpConverter):
# Per ONNX spec: "If input 'split' is specified, this attribute is ignored."
if not keepdims and len(inputs) == 1:
output = bb.emit(output)
- n = len(output.struct_info.fields)
+ n = len(output.ty.fields)
squeezed = [
relax.op.squeeze(bb.emit(relax.TupleGetItem(output, i)), axis=[axis])
for i in range(n)
@@ -4927,8 +4915,8 @@ class NonMaxSuppression(OnnxOpConverter):
yc = split_result[1]
w = split_result[2]
h = split_result[3]
- half_w = w / relax.const(2.0, boxes.struct_info.dtype)
- half_h = h / relax.const(2.0, boxes.struct_info.dtype)
+ half_w = w / relax.const(2.0, boxes.ty.dtype)
+ half_h = h / relax.const(2.0, boxes.ty.dtype)
x1 = xc - half_w
x2 = xc + half_w
y1 = yc - half_h
@@ -5023,8 +5011,8 @@ class AllClassNMS(OnnxOpConverter):
yc = split_result[1]
w = split_result[2]
h = split_result[3]
- half_w = w / relax.const(2.0, boxes.struct_info.dtype)
- half_h = h / relax.const(2.0, boxes.struct_info.dtype)
+ half_w = w / relax.const(2.0, boxes.ty.dtype)
+ half_h = h / relax.const(2.0, boxes.ty.dtype)
x1 = xc - half_w
x2 = xc + half_w
y1 = yc - half_h
@@ -5069,10 +5057,10 @@ class GridSample(OnnxOpConverter):
align_corners = bool(attr.get("align_corners", 0))
- if hasattr(data.struct_info, "ndim"):
- ndim = data.struct_info.ndim
+ if hasattr(data.ty, "ndim"):
+ ndim = data.ty.ndim
else:
- ndim = len(data.struct_info.shape)
+ ndim = len(data.ty.shape)
if ndim == 5 and method == "bicubic":
raise NotImplementedError(
@@ -5132,8 +5120,8 @@ class MatMulInteger(OnnxOpConverter):
a_zp = relax.op.astype(
a_zero_point, "int32"
) # Ensure zero point is int32 for subtraction
- a_zp = bb.normalize(a_zp) # Normalize the expr so struct_info gets populated
- a_zp_ndim = len(a_zp.struct_info.shape)
+ a_zp = bb.normalize(a_zp) # Normalize the expr so ty gets populated
+ a_zp_ndim = len(a_zp.ty.shape)
# Per-row case: [M] -> [M, 1] so it broadcasts over [M, K] row-wise
# N-D case: spec says shape is [D1, D2, M, 1], which already broadcasts correctly (no need to reshape)
@@ -5145,7 +5133,7 @@ class MatMulInteger(OnnxOpConverter):
if b_zero_point is not None:
b_zp = relax.op.astype(b_zero_point, "int32")
b_zp = bb.normalize(b_zp)
- b_zp_ndim = len(b_zp.struct_info.shape)
+ b_zp_ndim = len(b_zp.ty.shape)
# Per-col case: [N] -> [1, N] so it broadcasts over [K, N] column-wise
# N-D case: [D1, D2, 1, N] already broadcasts correctly
@@ -5481,9 +5469,7 @@ class ONNXGraphImporter:
def _new_var(self, var_name: str, shape: list, dtype: str = "float32"):
"""Creates a new Relax variable."""
- return relax.Var(
- name_hint=var_name, struct_info=relax.TensorStructInfo(shape=shape, dtype=dtype)
- )
+ return relax.Var(name_hint=var_name, ty=relax.TensorType(shape=shape, dtype=dtype))
def _parse_graph_input(self, graph: onnx.onnx_ml_pb2.GraphProto):
"""Parse model inputs to Relax parameters."""
@@ -5593,13 +5579,13 @@ class ONNXGraphImporter:
if (
inp is not None
and isinstance(inp, relax.Expr)
- and isinstance(inp.struct_info, relax.ShapeStructInfo)
+ and isinstance(inp.ty, relax.ShapeType)
and op_name not in shape_compatible_ops
):
raise ValueError(f"Node {node.name} cannot handle ShapeExpr inputs.")
try:
op = self._convert_operator(op_name, inputs, attr, self.opset)
- # Create struct information for the new operator.
+ # Create type information for the new operator.
if isinstance(op, relax.Expr):
op = self.bb.normalize(op)
except Exception as err: # pylint: disable=broad-exception-caught
@@ -5609,11 +5595,11 @@ class ONNXGraphImporter:
if op_name in return_tuple_ops:
outputs_num = 1
elif not isinstance(op, relax.Tuple):
- if isinstance(op.struct_info, relax.TupleStructInfo):
+ if isinstance(op.ty, relax.TupleType):
# This is a var bound to a tuple. We need to unpack it and create
# a new tuple.
tuple_items = []
- for i in range(len(op.struct_info.fields)):
+ for i in range(len(op.ty.fields)):
tuple_items.append(self.bb.emit(relax.TupleGetItem(op, i)))
op = relax.Tuple(tuple_items)
outputs_num = len(tuple_items)
@@ -5749,19 +5735,17 @@ class ONNXGraphImporter:
op = self._convert_operator(op_name, inputs, attr, self.opset)
try:
- _ = op.struct_info
- has_struct_info = True
+ _ = op.ty
+ has_ty = True
except tvm.error.InternalError:
- has_struct_info = False
+ has_ty = False
- if not has_struct_info:
+ if not has_ty:
op = bb.normalize(op)
if not isinstance(op, relax.Tuple):
- if isinstance(op.struct_info, relax.TupleStructInfo):
- tuple_items = [
- relax.TupleGetItem(op, i) for i in range(len(op.struct_info.fields))
- ]
+ if isinstance(op.ty, relax.TupleType):
+ tuple_items = [relax.TupleGetItem(op, i) for i in range(len(op.ty.fields))]
op = relax.Tuple(tuple_items)
outputs = node.output
diff --git a/python/tvm/relax/frontend/stablehlo/stablehlo_translator.py b/python/tvm/relax/frontend/stablehlo/stablehlo_translator.py
index a5196095ee..51994122d7 100644
--- a/python/tvm/relax/frontend/stablehlo/stablehlo_translator.py
+++ b/python/tvm/relax/frontend/stablehlo/stablehlo_translator.py
@@ -145,10 +145,10 @@ class StableHLOImporter:
if isinstance(lhs, relax.Expr) and isinstance(rhs, relax.Expr):
return lhs, rhs
if isinstance(lhs, relax.Expr):
- assert isinstance(lhs.struct_info, relax.TensorStructInfo)
- return lhs, relax.const(rhs, lhs.struct_info.dtype)
- assert isinstance(rhs.struct_info, relax.TensorStructInfo)
- return relax.const(lhs, rhs.struct_info.dtype), rhs
+ assert isinstance(lhs.ty, relax.TensorType)
+ return lhs, relax.const(rhs, lhs.ty.dtype)
+ assert isinstance(rhs.ty, relax.TensorType)
+ return relax.const(lhs, rhs.ty.dtype), rhs
def _call_binary_op(self, op, lhs, rhs):
lhs, rhs = StableHLOImporter._promote_binary_op_args(lhs, rhs)
@@ -381,7 +381,7 @@ class StableHLOImporter:
ipt_shape = self.get_shape(arg_shape)
ipt_dtype = self._convert_data_type(arg_shape.element_type)
ipt_name = "arg" + str(idx)
- ipt_var = relax.Var(f"arg{idx}", relax.TensorStructInfo(ipt_shape, ipt_dtype))
+ ipt_var = relax.Var(f"arg{idx}", relax.TensorType(ipt_shape, ipt_dtype))
self._nodes[ipt_name] = ipt_var
inputs.append(ipt_var)
diff --git a/python/tvm/relax/frontend/tflite/tflite_frontend.py b/python/tvm/relax/frontend/tflite/tflite_frontend.py
index 35e75e89bc..a087ce5d36 100644
--- a/python/tvm/relax/frontend/tflite/tflite_frontend.py
+++ b/python/tvm/relax/frontend/tflite/tflite_frontend.py
@@ -528,7 +528,7 @@ class OperatorConverter:
return relax.Tuple(relax.Tuple(ret), selections)
def _infer_shape(self, arg):
- return self.bb.normalize(arg).struct_info.shape
+ return self.bb.normalize(arg).ty.shape
def convert_op_to_relax(self):
"""Convert TFLite ops to relax ops"""
@@ -1120,13 +1120,11 @@ class OperatorConverter:
dims_expr = self.get_expr(shape_tensor.tensor_idx)
dims_ndim = int(self.get_tensor_shape(shape_tensor)[0])
dims_dtype = self.get_tensor_type_str(shape_tensor.tensor.Type())
- dims_expr = self.bb.match_cast(
- dims_expr, relax.TensorStructInfo([dims_ndim], dims_dtype)
- )
+ dims_expr = self.bb.match_cast(dims_expr, relax.TensorType([dims_ndim], dims_dtype))
dims_expr = self.bb.normalize(relax.op.astype(dims_expr, "int64"))
shape_dataflow_var = self.bb.emit(relax.op.tensor_to_shape(dims_expr))
shape_vars = [tirx.Var(f"{prefix}_{i}", "int64") for i in range(dims_ndim)]
- self.bb.match_cast(shape_dataflow_var, relax.ShapeStructInfo(shape_vars))
+ self.bb.match_cast(shape_dataflow_var, relax.ShapeType(shape_vars))
return relax.ShapeExpr(shape_vars), shape_vars
dims = to_int_list(self.get_tensor_value(shape_tensor))
@@ -1134,7 +1132,7 @@ class OperatorConverter:
def flatten_to_nd(self, x, nd=3):
"""Flatten input tensor to nd rank"""
- shape = x.struct_info.shape
+ shape = x.ty.shape
ndims = len(shape)
if ndims == nd:
return x
@@ -1899,7 +1897,7 @@ class OperatorConverter:
lhs = self.get_tensor_expr(input_tensors[0])
rhs = self.get_tensor_expr(input_tensors[1])
- dtype = lhs.struct_info.dtype
+ dtype = lhs.ty.dtype
if dtype == "bool":
op_fn = _op.logical_and
elif dtype.startswith(("int", "uint")):
@@ -1917,7 +1915,7 @@ class OperatorConverter:
lhs = self.get_tensor_expr(input_tensors[0])
rhs = self.get_tensor_expr(input_tensors[1])
- dtype = lhs.struct_info.dtype
+ dtype = lhs.ty.dtype
if dtype == "bool":
op_fn = _op.logical_or
elif dtype.startswith(("int", "uint")):
@@ -1987,7 +1985,7 @@ class OperatorConverter:
def _ensure_stablehlo_float_dtype(self, expr, op_name):
"""Return expr dtype if the StableHLO subset supports it."""
- dtype = expr.struct_info.dtype
+ dtype = expr.ty.dtype
if not dtype.startswith("float"):
raise tvm.error.OpNotImplemented(f"{op_name} with dtype {dtype} is not supported")
return dtype
@@ -2473,8 +2471,8 @@ class OperatorConverter:
gv,
[state_expr],
[
- relax.TensorStructInfo(tuple(state_shape), "uint64"),
- relax.TensorStructInfo(out_shape, out_dtype),
+ relax.TensorType(tuple(state_shape), "uint64"),
+ relax.TensorType(out_shape, out_dtype),
],
)
return self.bb.normalize(call)
@@ -2617,7 +2615,7 @@ class OperatorConverter:
input_name = get_tensor_name(subgraph, int(input_index))
shape = self._get_relax_tensor_shape(tensor)
dtype = self._get_relax_tensor_dtype(tensor)
- param = relax.Var(input_name, relax.TensorStructInfo(shape=shape, dtype=dtype))
+ param = relax.Var(input_name, relax.TensorType(shape=shape, dtype=dtype))
exp_tab.set_expr(input_name, param)
params.append(param)
return params, exp_tab
@@ -2627,7 +2625,7 @@ class OperatorConverter:
name = get_tensor_name(self.subgraph, tensor_wrapper.tensor_idx)
shape = self._get_relax_tensor_shape(tensor_wrapper)
dtype = self._get_relax_tensor_dtype(tensor_wrapper)
- return relax.Var(name, relax.TensorStructInfo(shape=shape, dtype=dtype))
+ return relax.Var(name, relax.TensorType(shape=shape, dtype=dtype))
def _lower_subgraph_to_function(self, subgraph_index, function_name_hint, op_name="CALL"):
"""Lower a TFLite subgraph into a private Relax function."""
@@ -3852,7 +3850,7 @@ class OperatorConverter:
stride = [int(i) for i in stride]
axes = list(range(len(begin)))
out = relax.op.strided_slice(data_expr, axes=axes, begin=begin, end=end, strides=stride)
- out_shape = self.bb.normalize(out).struct_info.shape
+ out_shape = self.bb.normalize(out).ty.shape
if not fshape_indices:
fshape_indices = range(len(out_shape))
@@ -3941,7 +3939,7 @@ class OperatorConverter:
return relax.op.call_dps_packed(
"tvm.contrib.random.uniform",
(seed, seed2, 0.0, 1.0),
- out_sinfo=relax.TensorStructInfo(out_shape, output_dtype),
+ out_ty=relax.TensorType(out_shape, output_dtype),
)
def convert_random_standard_normal(self, op):
@@ -3962,7 +3960,7 @@ class OperatorConverter:
return relax.op.call_dps_packed(
"tvm.contrib.random.normal",
(seed, seed2, 0.0, 1.0),
- out_sinfo=relax.TensorStructInfo(out_shape, output_dtype),
+ out_ty=relax.TensorType(out_shape, output_dtype),
)
def convert_multinomial(self, op):
@@ -3976,12 +3974,12 @@ class OperatorConverter:
if self.has_expr(num_samples_tensor.tensor_idx):
scalar_expr = self.get_expr(num_samples_tensor.tensor_idx)
scalar_dtype = self.get_tensor_type_str(num_samples_tensor.tensor.Type())
- scalar_expr = self.bb.match_cast(scalar_expr, relax.TensorStructInfo([], scalar_dtype))
+ scalar_expr = self.bb.match_cast(scalar_expr, relax.TensorType([], scalar_dtype))
scalar_expr = self.bb.normalize(relax.op.astype(scalar_expr, "int64"))
scalar_expr = self.bb.normalize(relax.op.reshape(scalar_expr, [1]))
shape_dataflow_var = self.bb.emit(relax.op.tensor_to_shape(scalar_expr))
num_samples = tirx.Var("multinomial_num_samples", "int64")
- self.bb.match_cast(shape_dataflow_var, relax.ShapeStructInfo([num_samples]))
+ self.bb.match_cast(shape_dataflow_var, relax.ShapeType([num_samples]))
else:
value = self.get_tensor_value(num_samples_tensor)
assert value.size == 1, (
@@ -4001,7 +3999,7 @@ class OperatorConverter:
uniform_sample = relax.op.call_dps_packed(
"tvm.contrib.random.uniform",
(seed, seed2, 0.0, 1.0),
- out_sinfo=relax.TensorStructInfo([output_batch, 1], "float32"),
+ out_ty=relax.TensorType([output_batch, 1], "float32"),
)
sample_indices = relax.op.reshape(
relax.op.broadcast_to(
@@ -4879,7 +4877,7 @@ class OperatorConverter:
# TFLite fixes the tuple arity in the graph, even when the split
# sizes themselves are supplied at runtime.
num_splits = len(output_tensors)
- rank = len(in_expr.struct_info.shape)
+ rank = len(in_expr.ty.shape)
# end_base is the full input shape; only split_axis changes per slice.
end_base = relax.op.shape_to_tensor(relax.op.shape_of(in_expr))
@@ -5256,7 +5254,7 @@ class OperatorConverter:
call = relax.call_tir(
gv,
[data_expr],
- relax.TensorStructInfo(relax_output_shape, "float32"),
+ relax.TensorType(relax_output_shape, "float32"),
)
return self.bb.normalize(call)
@@ -6543,7 +6541,7 @@ class OperatorConverter:
relax.ShapeExpr(crop_begin),
relax.ShapeExpr(crop_end),
),
- out_sinfo=relax.TensorStructInfo(output_shape, output_dtype),
+ out_ty=relax.TensorType(output_shape, output_dtype),
)
return out
@@ -6711,8 +6709,8 @@ class OperatorConverter:
input_a = self.get_expr(input_tensors[0].tensor_idx)
input_b = self.get_expr(input_tensors[1].tensor_idx)
- shape_a = list(input_a.struct_info.shape)
- shape_b = list(input_b.struct_info.shape)
+ shape_a = list(input_a.ty.shape)
+ shape_b = list(input_b.ty.shape)
rank_a = len(shape_a)
rank_b = len(shape_b)
@@ -6787,7 +6785,7 @@ class OperatorConverter:
relax.ShapeExpr(pad_after),
0.0,
),
- out_sinfo=relax.TensorStructInfo(output_shape, output_dtype),
+ out_ty=relax.TensorType(output_shape, output_dtype),
)
return out
@@ -6882,7 +6880,7 @@ class OperatorConverter:
out = relax.op.call_dps_packed(
"topi.sparse_to_dense",
(indices_expr, output_shape_expr, values_expr, default_value_expr),
- out_sinfo=relax.TensorStructInfo(output_shape_val, output_dtype),
+ out_ty=relax.TensorType(output_shape_val, output_dtype),
)
return out
@@ -7134,13 +7132,11 @@ class OperatorConverter:
# per-axis math.
if self.has_expr(dilations_tensor.tensor_idx):
dilations_expr = self.get_expr(dilations_tensor.tensor_idx)
- dilations_expr = self.bb.match_cast(
- dilations_expr, relax.TensorStructInfo([n_dims], "int32")
- )
+ dilations_expr = self.bb.match_cast(dilations_expr, relax.TensorType([n_dims], "int32"))
dilations_int64 = self.bb.normalize(relax.op.astype(dilations_expr, "int64"))
shape_var = self.bb.emit(relax.op.tensor_to_shape(dilations_int64))
stride_vars = [tirx.Var(f"dilate_stride_{i}", "int64") for i in range(n_dims)]
- self.bb.match_cast(shape_var, relax.ShapeStructInfo(stride_vars))
+ self.bb.match_cast(shape_var, relax.ShapeType(stride_vars))
strides = stride_vars
else:
strides = to_int_list(self.get_tensor_value(dilations_tensor))
@@ -7328,7 +7324,7 @@ class OperatorConverter:
num_detections = self.bb.emit(relax.TupleGetItem(nms_out, 2))
class_id_from_score = relax.op.squeeze(class_id_from_score, axis=[1])
- selected_score_slots = selected_scores.struct_info.shape.values[1]
+ selected_score_slots = selected_scores.ty.shape.values[1]
selected_detection_positions = relax.op.expand_dims(
relax.op.arange(selected_score_slots, dtype="int64"), axis=0
)
@@ -7693,7 +7689,7 @@ class OperatorConverter:
relax.const(False),
relax.const(False),
),
- out_sinfo=relax.TensorStructInfo(output_shape, output_dtype),
+ out_ty=relax.TensorType(output_shape, output_dtype),
)
return out
@@ -7734,7 +7730,7 @@ class OperatorConverter:
relax.const(False),
relax.const(False),
),
- out_sinfo=relax.TensorStructInfo(output_shape, output_dtype),
+ out_ty=relax.TensorType(output_shape, output_dtype),
)
return out
@@ -8776,7 +8772,7 @@ def from_tflite(
shape = tuple(shape) + (2,)
input_var = relax.Var(
name_hint=model_input_name,
- struct_info=relax.TensorStructInfo(shape=shape, dtype=dtype),
+ ty=relax.TensorType(shape=shape, dtype=dtype),
)
exp_tab.set_expr(model_input_name, input_var)
input_list.append(input_var)
diff --git a/python/tvm/relax/frontend/torch/base_fx_graph_translator.py b/python/tvm/relax/frontend/torch/base_fx_graph_translator.py
index 4c3cdd464f..f987f48d42 100644
--- a/python/tvm/relax/frontend/torch/base_fx_graph_translator.py
+++ b/python/tvm/relax/frontend/torch/base_fx_graph_translator.py
@@ -115,9 +115,9 @@ class BaseFXGraphImporter(metaclass=abc.ABCMeta):
import torch # type: ignore
if isinstance(tensor, relax.Expr):
- if not isinstance(tensor.struct_info, relax.TensorStructInfo):
+ if not isinstance(tensor.ty, relax.TensorType):
raise TypeError("The input Expr of shape_of should be a Tensor")
- return tensor.struct_info.shape
+ return tensor.ty.shape
elif isinstance(tensor, torch.Tensor):
return tensor.shape
raise ValueError(f"Unsupported type: {type(tensor)}")
@@ -210,7 +210,7 @@ class BaseFXGraphImporter(metaclass=abc.ABCMeta):
def _celu(self, node: fx.Node) -> relax.Var:
x = self.env[node.args[0]]
alpha = node.args[1] if len(node.args) > 1 else node.kwargs.get("alpha", 1.0)
- dtype = x.struct_info.dtype
+ dtype = x.ty.dtype
if isinstance(alpha, int | float):
alpha = relax.const(alpha, dtype)
@@ -329,7 +329,7 @@ class BaseFXGraphImporter(metaclass=abc.ABCMeta):
def _elu(self, node: fx.Node) -> relax.Var:
x = self.env[node.args[0]]
alpha = node.args[1] if len(node.args) > 1 else node.kwargs.get("alpha", 1.0)
- dtype = x.struct_info.dtype
+ dtype = x.ty.dtype
if isinstance(alpha, int | float):
alpha = relax.const(-alpha, dtype)
@@ -360,7 +360,7 @@ class BaseFXGraphImporter(metaclass=abc.ABCMeta):
def _hardsigmoid(self, node: fx.Node) -> relax.Var:
args = self.retrieve_args(node)
x = args[0]
- dtype = x.struct_info.dtype
+ dtype = x.ty.dtype
x0 = relax.op.add(x, relax.const(3, dtype))
x1 = relax.op.clip(x0, 0, 6)
return self.block_builder.emit(relax.op.divide(x1, relax.const(6, dtype)))
@@ -368,7 +368,7 @@ class BaseFXGraphImporter(metaclass=abc.ABCMeta):
def _hardswish(self, node: fx.Node) -> relax.Var:
args = self.retrieve_args(node)
x = args[0]
- dtype = x.struct_info.dtype
+ dtype = x.ty.dtype
x0 = relax.op.add(x, relax.const(3, dtype))
x1 = relax.op.clip(x0, 0, 6)
x2 = relax.op.divide(x1, relax.const(6, dtype))
@@ -396,9 +396,9 @@ class BaseFXGraphImporter(metaclass=abc.ABCMeta):
rhs = self.env[node.args[1]]
# torch.logical_and accepts any dtype (treating nonzero as True) and returns bool, but
# relax.op.logical_and requires boolean inputs, so cast non-bool inputs to bool first.
- if lhs.struct_info.dtype != "bool":
+ if lhs.ty.dtype != "bool":
lhs = self.block_builder.emit(relax.op.astype(lhs, "bool"))
- if rhs.struct_info.dtype != "bool":
+ if rhs.ty.dtype != "bool":
rhs = self.block_builder.emit(relax.op.astype(rhs, "bool"))
return self.block_builder.emit(relax.op.logical_and(lhs, rhs))
@@ -406,7 +406,7 @@ class BaseFXGraphImporter(metaclass=abc.ABCMeta):
x = self.env[node.args[0]]
# torch.logical_not accepts any dtype (treating nonzero as True) and returns bool, but
# relax.op.logical_not requires a boolean input, so cast non-bool inputs to bool first.
- if x.struct_info.dtype != "bool":
+ if x.ty.dtype != "bool":
x = self.block_builder.emit(relax.op.astype(x, "bool"))
return self.block_builder.emit(relax.op.logical_not(x))
@@ -415,9 +415,9 @@ class BaseFXGraphImporter(metaclass=abc.ABCMeta):
rhs = self.env[node.args[1]]
# torch.logical_or accepts any dtype (treating nonzero as True) and returns bool, but
# relax.op.logical_or requires boolean inputs, so cast non-bool inputs to bool first.
- if lhs.struct_info.dtype != "bool":
+ if lhs.ty.dtype != "bool":
lhs = self.block_builder.emit(relax.op.astype(lhs, "bool"))
- if rhs.struct_info.dtype != "bool":
+ if rhs.ty.dtype != "bool":
rhs = self.block_builder.emit(relax.op.astype(rhs, "bool"))
return self.block_builder.emit(relax.op.logical_or(lhs, rhs))
@@ -426,16 +426,16 @@ class BaseFXGraphImporter(metaclass=abc.ABCMeta):
rhs = self.env[node.args[1]]
# torch.logical_xor accepts any dtype (treating nonzero as True) and returns bool, but
# relax.op.logical_xor requires boolean inputs, so cast non-bool inputs to bool first.
- if lhs.struct_info.dtype != "bool":
+ if lhs.ty.dtype != "bool":
lhs = self.block_builder.emit(relax.op.astype(lhs, "bool"))
- if rhs.struct_info.dtype != "bool":
+ if rhs.ty.dtype != "bool":
rhs = self.block_builder.emit(relax.op.astype(rhs, "bool"))
return self.block_builder.emit(relax.op.logical_xor(lhs, rhs))
def _prelu(self, node: fx.Node) -> relax.Var:
x = self.env[node.args[0]]
alpha = self.env[node.args[1]]
- axis = 0 if len(x.struct_info.shape) == 1 else 1
+ axis = 0 if len(x.ty.shape) == 1 else 1
return self.block_builder.emit(relax.op.nn.prelu(x, alpha, axis))
def _round(self, node: fx.Node) -> relax.Expr:
@@ -446,7 +446,7 @@ class BaseFXGraphImporter(metaclass=abc.ABCMeta):
return self.block_builder.emit(relax.op.round(arg))
# For decimals != 0, use: round(x * 10^decimals) / 10^decimals
- dtype = arg.struct_info.dtype
+ dtype = arg.ty.dtype
scale = relax.const(10**decimals, dtype)
scaled = relax.op.multiply(arg, scale)
rounded = relax.op.round(scaled)
@@ -487,17 +487,17 @@ class BaseFXGraphImporter(metaclass=abc.ABCMeta):
"""
args = self.retrieve_args(node)
x = args[0]
- lambd = relax.const(args[1] if len(args) > 1 else 0.5, x.struct_info.dtype)
+ lambd = relax.const(args[1] if len(args) > 1 else 0.5, x.ty.dtype)
# Apply Softshrink transformation with masking
shrink_pos = relax.op.multiply(
relax.op.subtract(x, lambd),
- relax.op.astype(relax.op.greater(x, lambd), x.struct_info.dtype),
+ relax.op.astype(relax.op.greater(x, lambd), x.ty.dtype),
)
shrink_neg = relax.op.multiply(
relax.op.add(x, lambd),
- relax.op.astype(relax.op.less(x, relax.op.negative(lambd)), x.struct_info.dtype),
+ relax.op.astype(relax.op.less(x, relax.op.negative(lambd)), x.ty.dtype),
)
# Combine the positive and negative shrink results
@@ -522,10 +522,10 @@ class BaseFXGraphImporter(metaclass=abc.ABCMeta):
def convert(node: fx.Node) -> relax.Var:
def promote_binary_op_args(lhs, rhs):
if isinstance(lhs, relax.Expr) and isinstance(rhs, relax.Expr):
- lhs_si = getattr(lhs, "struct_info", None)
- rhs_si = getattr(rhs, "struct_info", None)
- if isinstance(lhs_si, relax.TensorStructInfo) and isinstance(
- rhs_si, relax.TensorStructInfo
+ lhs_si = getattr(lhs, "ty", None)
+ rhs_si = getattr(rhs, "ty", None)
+ if isinstance(lhs_si, relax.TensorType) and isinstance(
+ rhs_si, relax.TensorType
):
target_dtype = self._promote_common_dtype(lhs_si.dtype, rhs_si.dtype)
if target_dtype is not None:
@@ -535,11 +535,11 @@ class BaseFXGraphImporter(metaclass=abc.ABCMeta):
rhs = self.block_builder.emit(relax.op.astype(rhs, target_dtype))
return lhs, rhs
elif isinstance(lhs, relax.Expr):
- assert isinstance(lhs.struct_info, relax.TensorStructInfo)
- return lhs, relax.const(rhs, lhs.struct_info.dtype)
+ assert isinstance(lhs.ty, relax.TensorType)
+ return lhs, relax.const(rhs, lhs.ty.dtype)
elif isinstance(rhs, relax.Expr):
- assert isinstance(rhs.struct_info, relax.TensorStructInfo)
- return relax.const(lhs, rhs.struct_info.dtype), rhs
+ assert isinstance(rhs.ty, relax.TensorType)
+ return relax.const(lhs, rhs.ty.dtype), rhs
else:
assert False
@@ -551,9 +551,9 @@ class BaseFXGraphImporter(metaclass=abc.ABCMeta):
if isinstance(lhs, relax.Var) or isinstance(rhs, relax.Var):
return call_binary_op(relax_op, lhs, rhs)
elif isinstance(lhs, relax.expr.Constant) and not isinstance(rhs, relax.expr.Constant):
- return call_binary_op(relax_op, lhs, relax.const(rhs, dtype=lhs.struct_info.dtype))
+ return call_binary_op(relax_op, lhs, relax.const(rhs, dtype=lhs.ty.dtype))
elif isinstance(rhs, relax.expr.Constant) and not isinstance(lhs, relax.expr.Constant):
- return call_binary_op(relax_op, relax.const(lhs, dtype=rhs.struct_info.dtype), rhs)
+ return call_binary_op(relax_op, relax.const(lhs, dtype=rhs.ty.dtype), rhs)
return intrinsic_op(lhs, rhs)
return convert
@@ -565,8 +565,8 @@ class BaseFXGraphImporter(metaclass=abc.ABCMeta):
# a constant non-negative integer exponent into repeated multiplication instead.
if (
isinstance(lhs, relax.Expr)
- and isinstance(lhs.struct_info, relax.TensorStructInfo)
- and "int" in lhs.struct_info.dtype
+ and isinstance(lhs.ty, relax.TensorType)
+ and "int" in lhs.ty.dtype
and isinstance(rhs, int)
and not isinstance(rhs, bool)
and rhs >= 0
@@ -612,9 +612,9 @@ class BaseFXGraphImporter(metaclass=abc.ABCMeta):
if isinstance(lhs, relax.Expr) and isinstance(rhs, relax.Expr):
return self.block_builder.emit(relax.op.mod(lhs, rhs))
elif isinstance(lhs, relax.Expr):
- rhs = relax.const(rhs, lhs.struct_info.dtype)
+ rhs = relax.const(rhs, lhs.ty.dtype)
elif isinstance(rhs, relax.Expr):
- lhs = relax.const(lhs, rhs.struct_info.dtype)
+ lhs = relax.const(lhs, rhs.ty.dtype)
else:
assert False
return self.block_builder.emit(relax.op.mod(lhs, rhs))
@@ -639,7 +639,7 @@ class BaseFXGraphImporter(metaclass=abc.ABCMeta):
comparison = relax.op.equal(expanded_elements, flattened_test_elements)
summed = relax.op.sum(comparison, axis=-1)
- result = relax.op.greater(summed, relax.const(0, dtype=elements.struct_info.dtype))
+ result = relax.op.greater(summed, relax.const(0, dtype=elements.ty.dtype))
return self.block_builder.emit(result)
@@ -656,7 +656,7 @@ class BaseFXGraphImporter(metaclass=abc.ABCMeta):
# If ord_val is a Python float/int, wrap it in a Relax const
# so that it matches data's dtype.
- dtype = data.struct_info.dtype
+ dtype = data.ty.dtype
ord_expr = (
ord_val if isinstance(ord_val, relax.Expr) else relax.const(float(ord_val), dtype)
)
@@ -684,7 +684,7 @@ class BaseFXGraphImporter(metaclass=abc.ABCMeta):
x = self.env[node.args[0]]
output_size = node.args[1] if len(node.args) > 1 else node.kwargs["output_size"]
# Expand to 3D by adding batch dim if input is 2D
- x_ndim = x.struct_info.ndim
+ x_ndim = x.ty.ndim
if x_ndim == 2:
x = relax.op.expand_dims(x, axis=0)
@@ -700,7 +700,7 @@ class BaseFXGraphImporter(metaclass=abc.ABCMeta):
x = self.env[node.args[0]]
output_size = node.args[1]
# Expand to 4D by adding batch dim if input is 3D
- x_ndim = x.struct_info.ndim
+ x_ndim = x.ty.ndim
if x_ndim == 3:
x = relax.op.expand_dims(x, axis=0)
@@ -716,7 +716,7 @@ class BaseFXGraphImporter(metaclass=abc.ABCMeta):
x = self.env[node.args[0]]
output_size = node.args[1]
# Expand to 5D by adding batch dim if input is 4D
- x_ndim = x.struct_info.ndim
+ x_ndim = x.ty.ndim
if x_ndim == 4:
x = relax.op.expand_dims(x, axis=0)
@@ -739,10 +739,10 @@ class BaseFXGraphImporter(metaclass=abc.ABCMeta):
if alpha != 0:
res = self.block_builder.emit(relax.op.linear_algebra.matmul(y, z, out_dtype="float32"))
if alpha != 1:
- dtype = res.struct_info.dtype
+ dtype = res.ty.dtype
res = self.block_builder.emit(relax.op.multiply(res, relax.const(alpha, dtype)))
if beta != 0:
- dtype = x.struct_info.dtype
+ dtype = x.ty.dtype
if beta != 1:
bias = self.block_builder.emit(relax.op.multiply(x, relax.const(beta, dtype)))
else:
@@ -760,7 +760,7 @@ class BaseFXGraphImporter(metaclass=abc.ABCMeta):
count_include_pad: bool | None = True,
) -> relax.Var:
# Expand to 3D by adding batch dim if input is 2D
- x_ndim = x.struct_info.ndim
+ x_ndim = x.ty.ndim
if x_ndim == 2:
x = relax.op.expand_dims(x, axis=0)
stride = kernel_size if stride is None or stride == [] else stride
@@ -802,7 +802,7 @@ class BaseFXGraphImporter(metaclass=abc.ABCMeta):
count_include_pad: bool | None = True,
) -> relax.Var:
# Expand to 4D by adding batch dim if input is 3D
- x_ndim = x.struct_info.ndim
+ x_ndim = x.ty.ndim
if x_ndim == 3:
x = relax.op.expand_dims(x, axis=0)
stride = kernel_size if stride is None or stride == [] else stride
@@ -843,7 +843,7 @@ class BaseFXGraphImporter(metaclass=abc.ABCMeta):
count_include_pad: bool | None = True,
) -> relax.Var:
# Expand to 5D by adding batch dim if input is 4D
- x_ndim = x.struct_info.ndim
+ x_ndim = x.ty.ndim
if x_ndim == 4:
x = relax.op.expand_dims(x, axis=0)
stride = kernel_size if stride is None or stride == [] else stride
@@ -886,10 +886,10 @@ class BaseFXGraphImporter(metaclass=abc.ABCMeta):
if alpha != 0:
res = self.block_builder.emit(relax.op.matmul(batch1, batch2))
if alpha != 1:
- dtype = res.struct_info.dtype
+ dtype = res.ty.dtype
res = self.block_builder.emit(relax.op.multiply(res, relax.const(alpha, dtype)))
if beta != 0:
- dtype = x.struct_info.dtype
+ dtype = x.ty.dtype
if beta != 1:
bias = self.block_builder.emit(relax.op.multiply(x, relax.const(beta, dtype)))
else:
@@ -1259,12 +1259,12 @@ class BaseFXGraphImporter(metaclass=abc.ABCMeta):
) -> relax.Var:
x = self.block_builder.emit(relax.op.astype(x, "int32"))
- ndim = x.struct_info.ndim
+ ndim = x.ty.ndim
if ndim == 1:
return self.block_builder.emit(relax.op.take(weight, x, axis=0))
else:
- x_shape = x.struct_info.shape.values
- emb_size = weight.struct_info.shape.values[-1]
+ x_shape = x.ty.shape.values
+ emb_size = weight.ty.shape.values[-1]
x = self.block_builder.emit(relax.op.reshape(x, shape=[-1]))
embedding = self.block_builder.emit(relax.op.take(weight, x, axis=0))
return self.block_builder.emit(relax.op.reshape(embedding, [*x_shape, emb_size]))
@@ -1286,10 +1286,10 @@ class BaseFXGraphImporter(metaclass=abc.ABCMeta):
if gamma is None:
shape_tuple = [int(s) for s in normalized_shape]
- gamma = relax.const(np.ones(shape_tuple), x.struct_info.dtype)
+ gamma = relax.const(np.ones(shape_tuple), x.ty.dtype)
if beta is None:
shape_tuple = [int(s) for s in normalized_shape]
- beta = relax.const(np.zeros(shape_tuple), x.struct_info.dtype)
+ beta = relax.const(np.zeros(shape_tuple), x.ty.dtype)
return self.block_builder.emit(
relax.op.nn.layer_norm(
@@ -1319,8 +1319,8 @@ class BaseFXGraphImporter(metaclass=abc.ABCMeta):
gamma = self.params[module.weight]
beta = self.params[module.bias]
else:
- gamma = relax.const(torch.ones_like(module.normalized_shape), x.struct_info.dtype)
- beta = relax.const(torch.zeros_like(module.normalized_shape), x.struct_info.dtype)
+ gamma = relax.const(torch.ones_like(module.normalized_shape), x.ty.dtype)
+ beta = relax.const(torch.zeros_like(module.normalized_shape), x.ty.dtype)
eps = module.eps
return self._layer_norm_impl(x, gamma, beta, eps, normalized_shape)
@@ -1341,7 +1341,7 @@ class BaseFXGraphImporter(metaclass=abc.ABCMeta):
ceil_mode: bool | None = False,
) -> relax.Var:
# Expand to 3D by adding batch dim if input is 2D
- x_ndim = x.struct_info.ndim
+ x_ndim = x.ty.ndim
if x_ndim == 2:
x = relax.op.expand_dims(x, axis=0)
@@ -1385,7 +1385,7 @@ class BaseFXGraphImporter(metaclass=abc.ABCMeta):
ceil_mode: bool | None = False,
) -> relax.Var:
# Expand to 4D by adding batch dim if input is 3D
- x_ndim = x.struct_info.ndim
+ x_ndim = x.ty.ndim
if x_ndim == 3:
x = relax.op.expand_dims(x, axis=0)
@@ -1429,7 +1429,7 @@ class BaseFXGraphImporter(metaclass=abc.ABCMeta):
ceil_mode: bool | None = False,
) -> relax.Var:
# Expand to 5D by adding batch dim if input is 4D
- x_ndim = x.struct_info.ndim
+ x_ndim = x.ty.ndim
if x_ndim == 4:
x = relax.op.expand_dims(x, axis=0)
@@ -1519,7 +1519,7 @@ class BaseFXGraphImporter(metaclass=abc.ABCMeta):
# Calculate symmetric padding width for each dimension
# and applying them in reverse order to match the input dimensions.
- input_ndim = x.struct_info.ndim
+ input_ndim = x.ty.ndim
pad_width = [0] * (input_ndim * 2)
pad_pairs = [pad[i : i + 2] for i in range(0, len(pad), 2)]
reversed_pairs = list(reversed(pad_pairs))
@@ -1536,7 +1536,7 @@ class BaseFXGraphImporter(metaclass=abc.ABCMeta):
# Calculate symmetric padding width for each dimension
# and applying them in reverse order to match the input dimensions.
- input_ndim = x.struct_info.ndim
+ input_ndim = x.ty.ndim
pad_width = [0] * (input_ndim * 2)
pad_pairs = [pad[i : i + 2] for i in range(0, len(pad), 2)]
reversed_pairs = list(reversed(pad_pairs))
@@ -1560,7 +1560,7 @@ class BaseFXGraphImporter(metaclass=abc.ABCMeta):
value_tensor = self.env[node.args[2]]
# Check the dimensionality of the input tensors
- query_ndim = len(query_tensor.struct_info.shape)
+ query_ndim = len(query_tensor.ty.shape)
# TVM's nn.attention requires 4D inputs in format (batch, num_heads, seq_len, head_dim)
# For 2D inputs (seq_len, head_dim), we need to reshape to 4D first
@@ -1607,7 +1607,7 @@ class BaseFXGraphImporter(metaclass=abc.ABCMeta):
if attn_mask is not None:
attn_mask = self.env[attn_mask]
msg = "Only a float mask is supported for the attn_mask input."
- assert "float" in attn_mask.struct_info.dtype, msg
+ assert "float" in attn_mask.ty.dtype, msg
attention_output = self.block_builder.emit(
relax.op.nn.attention(query, key, value, bias=attn_mask, causal_mask=causal_mask)
@@ -1647,7 +1647,7 @@ class BaseFXGraphImporter(metaclass=abc.ABCMeta):
def _norm(self, node: fx.Node) -> relax.Var:
data = self.env[node.args[0]]
- dtype = data.struct_info.dtype
+ dtype = data.ty.dtype
order = node.args[1] if len(node.args) > 1 else node.kwargs.get("p", 2)
axis = node.args[2] if len(node.args) > 2 else None
keepdims = node.args[3] if len(node.args) > 3 else False
@@ -1714,7 +1714,7 @@ class BaseFXGraphImporter(metaclass=abc.ABCMeta):
else:
# Match PyTorch type promotion: summing bool or integer tensors
# accumulates in int64 unless an explicit dtype is given.
- input_dtype = x.struct_info.dtype
+ input_dtype = x.ty.dtype
if input_dtype == "bool" or (
(input_dtype.startswith("int") or input_dtype.startswith("uint"))
and input_dtype != "int64"
@@ -1761,14 +1761,12 @@ class BaseFXGraphImporter(metaclass=abc.ABCMeta):
scale = float("nan")
else:
scale = float(n) / float(n - correction)
- return self.block_builder.emit(
- relax.op.multiply(var, relax.const(scale, x.struct_info.dtype))
- )
+ return self.block_builder.emit(relax.op.multiply(var, relax.const(scale, x.ty.dtype)))
@staticmethod
def _reduction_size(x: relax.Expr, dim) -> int | None:
"""Static product of reduced-axis sizes; None if any axis is dynamic."""
- shape = x.struct_info.shape
+ shape = x.ty.shape
if shape is None:
return None
rank = len(shape)
@@ -1796,7 +1794,7 @@ class BaseFXGraphImporter(metaclass=abc.ABCMeta):
keepdim = args[2] if len(node.args) > 2 else node.kwargs.get("keepdim", False)
# max doesn't support boolean tensors directly, so we compute it in int8 and cast back
- if x.struct_info.dtype == "bool":
+ if x.ty.dtype == "bool":
x = relax.op.astype(x, "int8")
ret = relax.op.max(x, dim, keepdims=keepdim)
return self.block_builder.emit(relax.op.astype(ret, "bool"))
@@ -1977,7 +1975,7 @@ class BaseFXGraphImporter(metaclass=abc.ABCMeta):
data_shape = self.shape_of(tensor)
processed_indices = []
- max_ndim = max((idx.struct_info.ndim for _, idx in non_none_indices), default=1)
+ max_ndim = max((idx.ty.ndim for _, idx in non_none_indices), default=1)
for i, idx in enumerate(indices):
if idx is None:
@@ -2045,10 +2043,10 @@ class BaseFXGraphImporter(metaclass=abc.ABCMeta):
# Check if all indices can be squeezed to 1D for sequential take
def is_squeezable(idx):
- if idx.struct_info.ndim == 1:
+ if idx.ty.ndim == 1:
return True
- if idx.struct_info.ndim == 2:
- shape = idx.struct_info.shape
+ if idx.ty.ndim == 2:
+ shape = idx.ty.shape
for d in shape:
if isinstance(d, int) and d == 1:
return True
@@ -2061,13 +2059,13 @@ class BaseFXGraphImporter(metaclass=abc.ABCMeta):
if all_squeezable:
result = data
for axis, idx in reversed(non_none_indices):
- if idx.struct_info.ndim > 1:
+ if idx.ty.ndim > 1:
idx = self.block_builder.emit(relax.op.squeeze(idx))
result = self.block_builder.emit(relax.op.take(result, idx, axis=axis))
return result
# General case: replace None with arange, reshaped for broadcasting
- max_ndim = max((idx.struct_info.ndim for _, idx in non_none_indices), default=1)
+ max_ndim = max((idx.ty.ndim for _, idx in non_none_indices), default=1)
processed_indices = []
data_shape = self.shape_of(data)
@@ -2096,9 +2094,9 @@ class BaseFXGraphImporter(metaclass=abc.ABCMeta):
return input_list
new_inputs = []
for i, item in enumerate(input_list):
- if item.struct_info.ndim == 1:
+ if item.ty.ndim == 1:
new_inputs.append(item)
- elif item.struct_info.ndim == 0: # Change scalar value into 1D
+ elif item.ty.ndim == 0: # Change scalar value into 1D
const_tensor = relax.op.reshape(item, (1,))
new_inputs.append(const_tensor)
else:
@@ -2385,8 +2383,8 @@ class BaseFXGraphImporter(metaclass=abc.ABCMeta):
src = self.env[node.args[1]]
# Match PyTorch semantics: cast to destination dtype and broadcast to destination shape.
- if src.struct_info.dtype != dest.struct_info.dtype:
- src = self.block_builder.emit(relax.op.astype(src, dest.struct_info.dtype))
+ if src.ty.dtype != dest.ty.dtype:
+ src = self.block_builder.emit(relax.op.astype(src, dest.ty.dtype))
dest_shape = self.shape_of(dest)
src_shape = self.shape_of(src)
@@ -2476,16 +2474,16 @@ class BaseFXGraphImporter(metaclass=abc.ABCMeta):
def _fill(self, node: fx.Node) -> relax.Var:
args = self.retrieve_args(node)
x = args[0]
- dtype = x.struct_info.dtype
+ dtype = x.ty.dtype
value = args[1] if isinstance(args[1], relax.Expr) else relax.const(args[1], dtype)
- return self.block_builder.emit(relax.op.full(x.struct_info.shape, value, dtype))
+ return self.block_builder.emit(relax.op.full(x.ty.shape, value, dtype))
def _inplace_fill(self, node: fx.Node) -> relax.Var:
args = self.retrieve_args(node)
x = args[0]
- dtype = x.struct_info.dtype
+ dtype = x.ty.dtype
value = args[1] if isinstance(args[1], relax.Expr) else relax.const(args[1], dtype)
- filled = self.block_builder.emit(relax.op.full(x.struct_info.shape, value, dtype))
+ filled = self.block_builder.emit(relax.op.full(x.ty.shape, value, dtype))
self.env[node.args[0]] = filled
return filled
@@ -2511,7 +2509,7 @@ class BaseFXGraphImporter(metaclass=abc.ABCMeta):
value = node.args[1]
fill_value = relax.const(value)
- x_dtype = x.struct_info.dtype
+ x_dtype = x.ty.dtype
fill_dtype = None
if isinstance(value, int | float) and (math.isinf(value) or math.isnan(value)):
if not ("float" in x_dtype or "bfloat16" in x_dtype):
@@ -2531,7 +2529,7 @@ class BaseFXGraphImporter(metaclass=abc.ABCMeta):
value = node.args[2]
rx_value = relax.const(value)
- x_dtype = x.struct_info.dtype
+ x_dtype = x.ty.dtype
fill_dtype = None
if isinstance(value, int | float) and (math.isinf(value) or math.isnan(value)):
if not ("float" in x_dtype or "bfloat16" in x_dtype):
@@ -2576,7 +2574,7 @@ class BaseFXGraphImporter(metaclass=abc.ABCMeta):
value = node.args[2]
rx_value = relax.const(value)
- x_dtype = x.struct_info.dtype
+ x_dtype = x.ty.dtype
fill_dtype = None
if isinstance(value, int | float) and (math.isinf(value) or math.isnan(value)):
if not ("float" in x_dtype or "bfloat16" in x_dtype):
@@ -2620,8 +2618,8 @@ class BaseFXGraphImporter(metaclass=abc.ABCMeta):
return self.block_builder.emit(
relax.op.full(
size,
- relax.const(1, self_var.struct_info.dtype),
- self_var.struct_info.dtype,
+ relax.const(1, self_var.ty.dtype),
+ self_var.ty.dtype,
)
)
@@ -2639,8 +2637,8 @@ class BaseFXGraphImporter(metaclass=abc.ABCMeta):
return self.block_builder.emit(
relax.op.full(
size,
- relax.const(0, input_tensor.struct_info.dtype),
- input_tensor.struct_info.dtype,
+ relax.const(0, input_tensor.ty.dtype),
+ input_tensor.ty.dtype,
)
)
@@ -2678,7 +2676,7 @@ class BaseFXGraphImporter(metaclass=abc.ABCMeta):
def _type_as(self, node: fx.Node) -> relax.Var:
x = self.env[node.args[0]]
other = self.env[node.args[1]]
- dtype = other.struct_info.dtype
+ dtype = other.ty.dtype
return self.block_builder.emit(relax.op.astype(x, dtype))
########## Others ##########
@@ -2690,10 +2688,10 @@ class BaseFXGraphImporter(metaclass=abc.ABCMeta):
if isinstance(x, list | tuple | relax.ShapeExpr | relax.Tuple):
return x[node.args[1]]
elif isinstance(x, relax.Var):
- if isinstance(x.struct_info, relax.TupleStructInfo):
+ if isinstance(x.ty, relax.TupleType):
return self.block_builder.emit(relax.TupleGetItem(x, node.args[1]))
- assert isinstance(x.struct_info, relax.TensorStructInfo)
+ assert isinstance(x.ty, relax.TensorType)
if isinstance(node.args[1], int):
return x
if not isinstance(node.args[1], list | tuple):
@@ -2765,7 +2763,7 @@ class BaseFXGraphImporter(metaclass=abc.ABCMeta):
sliced_shape.insert(i, 1)
return self.block_builder.emit(relax.op.reshape(sliced, sliced_shape))
elif isinstance(x, relax.Constant):
- dtype = x.struct_info.dtype
+ dtype = x.ty.dtype
return relax.const(x.data.numpy()[node.args[1]], dtype)
else:
assert False
diff --git a/python/tvm/relax/frontend/torch/exported_program_translator.py b/python/tvm/relax/frontend/torch/exported_program_translator.py
index b96316adee..ac07ad2ef2 100644
--- a/python/tvm/relax/frontend/torch/exported_program_translator.py
+++ b/python/tvm/relax/frontend/torch/exported_program_translator.py
@@ -84,27 +84,27 @@ class ExportedProgramImporter(BaseFXGraphImporter):
def _log2(self, node: fx.Node) -> relax.Var:
x = self.env[node.args[0]]
return self.block_builder.emit(
- relax.op.divide(relax.op.log(x), relax.const(0.6931471805599453, x.struct_info.dtype))
+ relax.op.divide(relax.op.log(x), relax.const(0.6931471805599453, x.ty.dtype))
)
def _log10(self, node: fx.Node) -> relax.Var:
x = self.env[node.args[0]]
return self.block_builder.emit(
- relax.op.divide(relax.op.log(x), relax.const(2.302585092994046, x.struct_info.dtype))
+ relax.op.divide(relax.op.log(x), relax.const(2.302585092994046, x.ty.dtype))
)
def _log1p(self, node: fx.Node) -> relax.Var:
x = self.env[node.args[0]]
- one = relax.const(1, x.struct_info.dtype)
+ one = relax.const(1, x.ty.dtype)
return self.block_builder.emit(relax.op.log(relax.op.add(x, one)))
def _reciprocal(self, node: fx.Node) -> relax.Var:
x = self.env[node.args[0]]
- return self.block_builder.emit(relax.op.divide(relax.const(1.0, x.struct_info.dtype), x))
+ return self.block_builder.emit(relax.op.divide(relax.const(1.0, x.ty.dtype), x))
def _sqrt(self, node: fx.Node) -> relax.Var:
x = self.env[node.args[0]]
- dtype = x.struct_info.dtype
+ dtype = x.ty.dtype
# Check if input is integer type and convert to float32 if needed
if dtype in ("int8", "int16", "int32", "int64", "uint8", "uint16", "uint32", "uint64"):
@@ -114,7 +114,7 @@ class ExportedProgramImporter(BaseFXGraphImporter):
def _rsqrt(self, node: fx.Node) -> relax.Var:
x = self.env[node.args[0]]
- dtype = x.struct_info.dtype
+ dtype = x.ty.dtype
# Check if input is integer type and convert to float32 if needed
if dtype in ("int8", "int16", "int32", "int64", "uint8", "uint16", "uint32", "uint64"):
@@ -134,7 +134,7 @@ class ExportedProgramImporter(BaseFXGraphImporter):
x = self.env[node.args[0]]
channel = int(self.shape_of(x)[1])
- dtype = x.struct_info.dtype
+ dtype = x.ty.dtype
scale = node.args[1] is not None
center = node.args[2] is not None
weight = self.env.get(node.args[1], relax.const(np.ones(channel), dtype=dtype))
@@ -192,7 +192,7 @@ class ExportedProgramImporter(BaseFXGraphImporter):
x = self.env[node.args[0]]
channel = int(self.shape_of(x)[1])
- dtype = x.struct_info.dtype
+ dtype = x.ty.dtype
output = self.block_builder.emit(bn_tuple[0])
new_running_mean = self.block_builder.emit(bn_tuple[1])
@@ -210,7 +210,7 @@ class ExportedProgramImporter(BaseFXGraphImporter):
x = self.env[node.args[0]]
channel = int(self.shape_of(x)[1])
- dtype = x.struct_info.dtype
+ dtype = x.ty.dtype
weight = self.env.get(node.args[1], relax.const(np.ones(channel), dtype=dtype))
bias = self.env.get(node.args[2], relax.const(np.zeros(channel), dtype=dtype))
eps = node.args[5] if len(node.args) > 5 else node.kwargs.get("eps", 1e-05)
@@ -508,7 +508,7 @@ class ExportedProgramImporter(BaseFXGraphImporter):
# o_t = sigmoid(W_io * x_t + b_io + W_ho * h_{t-1} + b_ho)
# c_t = f_t * c_{t-1} + i_t * g_t
# h_t = o_t * tanh(c_t)
- dtype = input_tensor.struct_info.dtype
+ dtype = input_tensor.ty.dtype
params_per_direction = 4 if has_biases else 2
# Extract or create forward direction weights
@@ -807,7 +807,7 @@ class ExportedProgramImporter(BaseFXGraphImporter):
# Fallback to a default hidden size
hidden_size = 16
- dtype = input_tensor.struct_info.dtype
+ dtype = input_tensor.ty.dtype
# Extract forward direction weights
if params and len(params) >= params_per_direction:
@@ -1045,9 +1045,9 @@ class ExportedProgramImporter(BaseFXGraphImporter):
import numpy as np
x = self.env[node.args[0]]
- x_sinfo = x.struct_info
- shape = [int(s) for s in x_sinfo.shape]
- dtype = self._convert_data_type(node.kwargs.get("dtype", None) or x_sinfo.dtype, self.env)
+ x_ty = x.ty
+ shape = [int(s) for s in x_ty.shape]
+ dtype = self._convert_data_type(node.kwargs.get("dtype", None) or x_ty.dtype, self.env)
data = np.random.randn(*shape).astype(dtype)
return self.block_builder.emit(relax.const(data, dtype))
@@ -1088,13 +1088,13 @@ class ExportedProgramImporter(BaseFXGraphImporter):
)
if alpha != 1.0:
- alpha_const = relax.const(alpha, matmul_result.struct_info.dtype)
+ alpha_const = relax.const(alpha, matmul_result.ty.dtype)
matmul_result = self.block_builder.emit(relax.op.multiply(matmul_result, alpha_const))
# Compute beta * input + alpha * matmul_result
if beta != 0.0:
if beta != 1.0:
- beta_const = relax.const(beta, input_tensor.struct_info.dtype)
+ beta_const = relax.const(beta, input_tensor.ty.dtype)
input_scaled = self.block_builder.emit(relax.op.multiply(input_tensor, beta_const))
else:
input_scaled = input_tensor
@@ -1168,7 +1168,7 @@ class ExportedProgramImporter(BaseFXGraphImporter):
relax.op.strided_slice(rois, axes=[1], begin=[1], end=[5])
)
boxes = self.block_builder.emit(
- relax.op.subtract(boxes, relax.const(0.5, rois.struct_info.dtype))
+ relax.op.subtract(boxes, relax.const(0.5, rois.ty.dtype))
)
rois = self.block_builder.emit(relax.op.concat([batch_indices, boxes], axis=1))
@@ -1198,7 +1198,7 @@ class ExportedProgramImporter(BaseFXGraphImporter):
x = self.env[node.args[0]]
channel = int(self.shape_of(x)[1])
- dtype = x.struct_info.dtype
+ dtype = x.ty.dtype
gamma = self.env.get(node.args[1], relax.const(np.ones(channel), dtype=dtype))
beta = self.env.get(node.args[2], relax.const(np.zeros(channel), dtype=dtype))
eps = node.args[4] if node.args[4] else 1e-05
@@ -1247,7 +1247,7 @@ class ExportedProgramImporter(BaseFXGraphImporter):
index = self.env[node.args[2]]
value = node.args[3]
- value_const = relax.const(value, x.struct_info.dtype)
+ value_const = relax.const(value, x.ty.dtype)
src = self.block_builder.emit(relax.op.broadcast_to(value_const, self.shape_of(index)))
return self.block_builder.emit(relax.op.scatter_elements(x, index, src, axis=dim))
@@ -1386,7 +1386,7 @@ class ExportedProgramImporter(BaseFXGraphImporter):
graph_module : torch.fx.GraphModule
The branch subgraph (e.g. true_graph_0 / false_graph_0).
operands : list[relax.Expr]
- The operands passed to the cond; used to derive parameter struct_info.
+ The operands passed to the cond; used to derive parameter ty.
name_hint : str
A hint for the function name (e.g. "cond_true_branch_0").
@@ -1417,10 +1417,8 @@ class ExportedProgramImporter(BaseFXGraphImporter):
placeholders = [n for n in nodes if n.op == "placeholder"]
params = []
for ph, operand in zip(placeholders, operands):
- if hasattr(operand, "struct_info") and isinstance(
- operand.struct_info, relax.TensorStructInfo
- ):
- orig_si = operand.struct_info
+ if hasattr(operand, "ty") and isinstance(operand.ty, relax.TensorType):
+ orig_si = operand.ty
# Create fresh SizeVars to avoid sharing with the caller function.
if orig_si.shape is not None:
new_shape = [
@@ -1429,13 +1427,13 @@ class ExportedProgramImporter(BaseFXGraphImporter):
else s
for s in orig_si.shape
]
- si = relax.TensorStructInfo(new_shape, orig_si.dtype)
+ si = relax.TensorType(new_shape, orig_si.dtype)
else:
si = orig_si
- elif hasattr(operand, "struct_info"):
- si = operand.struct_info
+ elif hasattr(operand, "ty"):
+ si = operand.ty
else:
- si = relax.ObjectStructInfo()
+ si = relax.ObjectType()
param = relax.Var(ph.name, si)
params.append(param)
self.env[ph] = param
@@ -1536,7 +1534,7 @@ class ExportedProgramImporter(BaseFXGraphImporter):
"expm1.default": lambda node: self.block_builder.emit(
relax.op.subtract(
relax.op.exp(self.env[node.args[0]]),
- relax.const(1.0, self.env[node.args[0]].struct_info.dtype),
+ relax.const(1.0, self.env[node.args[0]].ty.dtype),
)
),
"floor.default": self._unary_op(relax.op.floor),
@@ -1960,7 +1958,7 @@ class ExportedProgramImporter(BaseFXGraphImporter):
relax_shape.append(s)
dtype = self._convert_data_type(torch_dtype)
- relax_var = relax.Var(name_hint, relax.TensorStructInfo(relax_shape, dtype))
+ relax_var = relax.Var(name_hint, relax.TensorType(relax_shape, dtype))
if spec.kind is torch.export.graph_signature.InputKind.USER_INPUT:
user_inputs[name_hint] = relax_var
else:
diff --git a/python/tvm/relax/frontend/torch/fx_translator.py b/python/tvm/relax/frontend/torch/fx_translator.py
index 4932871bad..40215f07ac 100644
--- a/python/tvm/relax/frontend/torch/fx_translator.py
+++ b/python/tvm/relax/frontend/torch/fx_translator.py
@@ -66,7 +66,7 @@ class TorchFXImporter(BaseFXGraphImporter):
def _reciprocal(self, node: fx.Node) -> relax.Var:
x = self.env[node.args[0]]
- return self.block_builder.emit(relax.op.divide(relax.const(1.0, x.struct_info.dtype), x))
+ return self.block_builder.emit(relax.op.divide(relax.const(1.0, x.ty.dtype), x))
def _leakyrelu_module(self, node: fx.Node) -> relax.Var:
x = self.env[node.args[0]]
@@ -84,23 +84,23 @@ class TorchFXImporter(BaseFXGraphImporter):
def _log2(self, node: fx.Node) -> relax.Var:
x = self.env[node.args[0]]
return self.block_builder.emit(
- relax.op.divide(relax.op.log(x), relax.const(0.6931471805599453, x.struct_info.dtype))
+ relax.op.divide(relax.op.log(x), relax.const(0.6931471805599453, x.ty.dtype))
)
def _log10(self, node: fx.Node) -> relax.Var:
x = self.env[node.args[0]]
return self.block_builder.emit(
- relax.op.divide(relax.op.log(x), relax.const(2.302585092994046, x.struct_info.dtype))
+ relax.op.divide(relax.op.log(x), relax.const(2.302585092994046, x.ty.dtype))
)
def _log1p(self, node: fx.Node) -> relax.Var:
x = self.env[node.args[0]]
- one = relax.const(1, x.struct_info.dtype)
+ one = relax.const(1, x.ty.dtype)
return self.block_builder.emit(relax.op.log(relax.op.add(x, one)))
def _sqrt(self, node: fx.Node) -> relax.Var:
x = self.env[node.args[0]]
- dtype = x.struct_info.dtype
+ dtype = x.ty.dtype
# Check if input is integer type and convert to float32 if needed
if dtype in ["int8", "int16", "int32", "int64", "uint8", "uint16", "uint32", "uint64"]:
@@ -110,7 +110,7 @@ class TorchFXImporter(BaseFXGraphImporter):
def _rsqrt(self, node: fx.Node) -> relax.Var:
x = self.env[node.args[0]]
- dtype = x.struct_info.dtype
+ dtype = x.ty.dtype
# Check if input is integer type and convert to float32 if needed
if dtype in ["int8", "int16", "int32", "int64", "uint8", "uint16", "uint32", "uint64"]:
@@ -130,7 +130,7 @@ class TorchFXImporter(BaseFXGraphImporter):
module = self.named_modules[node.target]
alpha_tensor = module.weight.numpy()
alpha = relax.const(alpha_tensor, dtype="float32")
- axis = 0 if len(x.struct_info.shape) == 1 else 1 # Extract Channel size
+ axis = 0 if len(x.ty.shape) == 1 else 1 # Extract Channel size
return self.block_builder.emit(relax.op.nn.prelu(x, alpha, axis))
def _softmax_module(self, node: fx.Node) -> relax.Var:
@@ -164,11 +164,11 @@ class TorchFXImporter(BaseFXGraphImporter):
if isinstance(lhs, relax.Expr) and isinstance(rhs, relax.Expr):
return lhs, rhs
elif isinstance(lhs, relax.Expr):
- assert isinstance(lhs.struct_info, relax.TensorStructInfo)
- return lhs, relax.const(rhs, lhs.struct_info.dtype)
+ assert isinstance(lhs.ty, relax.TensorType)
+ return lhs, relax.const(rhs, lhs.ty.dtype)
elif isinstance(rhs, relax.Expr):
- assert isinstance(rhs.struct_info, relax.TensorStructInfo)
- return relax.const(lhs, rhs.struct_info.dtype), rhs
+ assert isinstance(rhs.ty, relax.TensorType)
+ return relax.const(lhs, rhs.ty.dtype), rhs
else:
assert False
@@ -183,16 +183,12 @@ class TorchFXImporter(BaseFXGraphImporter):
return output
elif isinstance(lhs, relax.expr.Constant):
- output = call_binary_op(
- relax_op, lhs, relax.const(rhs, dtype=lhs.struct_info.dtype)
- )
+ output = call_binary_op(relax_op, lhs, relax.const(rhs, dtype=lhs.ty.dtype))
self.env[node.args[0]] = output
return output
elif isinstance(rhs, relax.expr.Constant):
- output = call_binary_op(
- relax_op, relax.const(lhs, dtype=rhs.struct_info.dtype), rhs
- )
+ output = call_binary_op(relax_op, relax.const(lhs, dtype=rhs.ty.dtype), rhs)
self.env[node.args[0]] = output
return output
@@ -209,7 +205,7 @@ class TorchFXImporter(BaseFXGraphImporter):
x = self.env[node.args[0]]
output_size = module.output_size
# Expand to 3D by adding batch dim if input is 2D
- x_ndim = x.struct_info.ndim
+ x_ndim = x.ty.ndim
if x_ndim == 2:
x = relax.op.expand_dims(x, axis=0)
result = self.block_builder.emit(
@@ -225,7 +221,7 @@ class TorchFXImporter(BaseFXGraphImporter):
x = self.env[node.args[0]]
output_size = module.output_size
# Expand to 4D by adding batch dim if input is 3D
- x_ndim = x.struct_info.ndim
+ x_ndim = x.ty.ndim
if x_ndim == 3:
x = relax.op.expand_dims(x, axis=0)
result = self.block_builder.emit(
@@ -241,7 +237,7 @@ class TorchFXImporter(BaseFXGraphImporter):
x = self.env[node.args[0]]
output_size = module.output_size
# Expand to 5D by adding batch dim if input is 4D
- x_ndim = x.struct_info.ndim
+ x_ndim = x.ty.ndim
if x_ndim == 4:
x = relax.op.expand_dims(x, axis=0)
result = self.block_builder.emit(
@@ -312,7 +308,7 @@ class TorchFXImporter(BaseFXGraphImporter):
else:
import numpy as np
- dtype = x.struct_info.dtype
+ dtype = x.ty.dtype
channel = int(self.shape_of(x)[1])
weight = relax.const(np.ones(channel), dtype=dtype)
bias = relax.const(np.zeros(channel), dtype=dtype)
@@ -432,7 +428,7 @@ class TorchFXImporter(BaseFXGraphImporter):
if weights in self.params:
weights = self.params[weights]
else:
- weights = relax.const(weights.numpy(), preds.struct_info.dtype)
+ weights = relax.const(weights.numpy(), preds.ty.dtype)
reduction = module.reduction
ignore_index = module.ignore_index
@@ -461,8 +457,8 @@ class TorchFXImporter(BaseFXGraphImporter):
gamma = self.params[module.weight]
beta = self.params[module.bias]
else:
- gamma = relax.const(torch.ones_like(module.num_channels), x.struct_info.dtype)
- beta = relax.const(torch.zeros_like(module.num_channels), x.struct_info.dtype)
+ gamma = relax.const(torch.ones_like(module.num_channels), x.ty.dtype)
+ beta = relax.const(torch.zeros_like(module.num_channels), x.ty.dtype)
eps = module.eps
dim = len(self.shape_of(x))
@@ -559,7 +555,7 @@ class TorchFXImporter(BaseFXGraphImporter):
else:
coord_trans = "half_pixel"
- if data.struct_info.ndim == 5:
+ if data.ty.ndim == 5:
if self.default_image_layout == "NDHWC":
layout_3d = "NDHWC"
else:
@@ -691,8 +687,8 @@ class TorchFXImporter(BaseFXGraphImporter):
dest = self.env[node.args[0]]
src = self.env[node.args[1]]
- if src.struct_info.dtype != dest.struct_info.dtype:
- src = self.block_builder.emit(relax.op.astype(src, dest.struct_info.dtype))
+ if src.ty.dtype != dest.ty.dtype:
+ src = self.block_builder.emit(relax.op.astype(src, dest.ty.dtype))
dest_shape = self.shape_of(dest)
src_shape = self.shape_of(src)
@@ -706,7 +702,7 @@ class TorchFXImporter(BaseFXGraphImporter):
x = self.env[node.args[0]]
mask = self.env[node.args[1]]
source = self.env[node.args[2]]
- ndim = len(mask.struct_info.shape)
+ ndim = len(mask.ty.shape)
if ndim == 1:
index = self.block_builder.emit(relax.op.cumsum(mask, 0, dtype="int32"))
index = self.block_builder.emit(relax.op.subtract(index, relax.const(1, "int32")))
@@ -715,16 +711,12 @@ class TorchFXImporter(BaseFXGraphImporter):
f_mask = self.block_builder.emit(relax.op.reshape(mask, [-1]))
index = self.block_builder.emit(relax.op.cumsum(f_mask, 0, dtype="int32"))
index = self.block_builder.emit(relax.op.subtract(index, relax.const(1, "int32")))
- source_shape = [-1] + [
- s for idx, s in enumerate(source.struct_info.shape) if idx >= ndim
- ]
+ source_shape = [-1] + [s for idx, s in enumerate(source.ty.shape) if idx >= ndim]
f_source = self.block_builder.emit(relax.op.reshape(source, source_shape))
gathered_source = self.block_builder.emit(relax.op.take(f_source, index, axis=0))
- gathered_source = self.block_builder.emit(
- relax.op.reshape(gathered_source, x.struct_info.shape)
- )
- if ndim != len(x.struct_info.shape):
- mask = self.block_builder.emit(relax.op.broadcast_to(mask, x.struct_info.shape))
+ gathered_source = self.block_builder.emit(relax.op.reshape(gathered_source, x.ty.shape))
+ if ndim != len(x.ty.shape):
+ mask = self.block_builder.emit(relax.op.broadcast_to(mask, x.ty.shape))
return self.block_builder.emit(relax.op.where(mask, gathered_source, x))
def _one_hot(self, node: fx.Node) -> relax.Var:
@@ -757,9 +749,7 @@ class TorchFXImporter(BaseFXGraphImporter):
def _is_floating_point(self, node: fx.Node) -> relax.Var:
x = self.env[node.args[0]]
- return relax.const(
- x.struct_info.dtype in ["float16", "float32", "float64", "bfloat16"], "bool"
- )
+ return relax.const(x.ty.dtype in ["float16", "float32", "float64", "bfloat16"], "bool")
def _type(self, node: fx.Node) -> relax.Var:
x = self.env[node.args[0]]
@@ -771,7 +761,7 @@ class TorchFXImporter(BaseFXGraphImporter):
def _getattr(self, node: fx.Node) -> relax.Var:
if isinstance(self.env[node.args[0]], relax.Expr):
if node.args[1] == "dtype":
- return self.env[node.args[0]].struct_info.dtype
+ return self.env[node.args[0]].ty.dtype
elif node.args[1] == "shape":
return self.shape_of(self.env[node.args[0]])
return getattr(self.env[node.args[0]], node.args[1])
@@ -780,9 +770,7 @@ class TorchFXImporter(BaseFXGraphImporter):
inputs = list()
for idx, (shape, dtype) in enumerate(input_info):
inputs.append(
- relax.Var(
- f"inp_{idx}", relax.TensorStructInfo(shape, self._convert_data_type(dtype))
- )
+ relax.Var(f"inp_{idx}", relax.TensorType(shape, self._convert_data_type(dtype)))
)
return inputs
@@ -1094,7 +1082,7 @@ class TorchFXImporter(BaseFXGraphImporter):
for name, param in sorted(model.named_parameters(), key=lambda x: x[0]):
shape = param.data.shape
dtype = self._convert_data_type(str(param.data.dtype))
- inputs.append(relax.Var(name, relax.TensorStructInfo(shape, dtype)))
+ inputs.append(relax.Var(name, relax.TensorType(shape, dtype)))
self.params[param] = inputs[-1]
params.append(tvm.runtime.tensor(param.data.cpu().numpy()))
else:
diff --git a/python/tvm/relax/ir/instrument.py b/python/tvm/relax/ir/instrument.py
index 6254a75a8e..711ec10e6a 100644
--- a/python/tvm/relax/ir/instrument.py
+++ b/python/tvm/relax/ir/instrument.py
@@ -27,9 +27,9 @@ class WellFormedInstrument:
Parameters
----------
- check_struct_info: bool
+ check_ty: bool
- If True, validate the struct info in the module. If False,
+ If True, validate the type in the module. If False,
skip these checks.
validate_before_transform: bool
@@ -39,9 +39,9 @@ class WellFormedInstrument:
after running a transform.
"""
- def __init__(self, check_struct_info: bool = True, validate_before_transform: bool = True):
+ def __init__(self, check_ty: bool = True, validate_before_transform: bool = True):
self.skip_pass_name = ["Normalize", "NormalizeGlobalVar", "ResolveGlobals"]
- self.check_struct_info = check_struct_info
+ self.check_ty = check_ty
self.validate_before_transform = validate_before_transform
def run_before_pass(self, mod, pass_info):
@@ -53,7 +53,7 @@ class WellFormedInstrument:
def _check(self, mod, pass_name, name_prefix):
if pass_name not in self.skip_pass_name:
- is_well_formed = relax.analysis.check_well_formed(mod, self.check_struct_info)
+ is_well_formed = relax.analysis.check_well_formed(mod, self.check_ty)
if not is_well_formed:
mod.show(name=f"{name_prefix}{pass_name}")
assert is_well_formed
diff --git a/python/tvm/relax/op/_op_gradient.py b/python/tvm/relax/op/_op_gradient.py
index 2d949ee276..5aae26b75f 100644
--- a/python/tvm/relax/op/_op_gradient.py
+++ b/python/tvm/relax/op/_op_gradient.py
@@ -22,7 +22,7 @@ import operator
from tvm import relax
from tvm.arith import Analyzer
-from tvm.relax.struct_info import ShapeStructInfo
+from tvm.relax.type import ShapeType
from ...tirx import PrimExpr
from ..block_builder import BlockBuilder
@@ -65,7 +65,7 @@ from .unary import cos, exp, log, sigmoid, sin
def _get_shape(expr: Expr) -> ShapeExpr:
"""Get the shape from a Tensor expr."""
try:
- shape = expr.struct_info.shape
+ shape = expr.ty.shape
except Exception as error:
raise RuntimeError(
f"Get the shape of {expr} failed. Please normalize it first and ensure it is a Tensor."
@@ -76,7 +76,7 @@ def _get_shape(expr: Expr) -> ShapeExpr:
def _get_dtype(expr: Expr) -> str:
"""Get the dtype from a Tensor expr."""
try:
- dtype = expr.struct_info.dtype
+ dtype = expr.ty.dtype
except Exception as error:
raise RuntimeError(
f"Get the dtype of {expr} failed. Please normalize it first and ensure it is a Tensor."
@@ -86,17 +86,17 @@ def _get_dtype(expr: Expr) -> str:
def _fit_shape(bb: BlockBuilder, input_grad: Expr, input: Expr) -> Expr:
"""When expr and target has the same shape, return expr;
- otherwise return `collapse_sum_to(expr, target.struct_info.shape)`.
+ otherwise return `collapse_sum_to(expr, target.ty.shape)`.
Will use BlockBuilder to normalize expr first.
"""
target_shape = _get_shape(input)
- expr_sinfo = _get_shape(bb.normalize(input_grad)).struct_info
- target_sinfo = target_shape.struct_info
- assert isinstance(expr_sinfo, ShapeStructInfo)
- assert isinstance(target_sinfo, ShapeStructInfo)
+ expr_ty = _get_shape(bb.normalize(input_grad)).ty
+ target_ty = target_shape.ty
+ assert isinstance(expr_ty, ShapeType)
+ assert isinstance(target_ty, ShapeType)
- def _check_shape_equal(lhs: ShapeStructInfo, rhs: ShapeStructInfo):
+ def _check_shape_equal(lhs: ShapeType, rhs: ShapeType):
if len(lhs.values) != len(rhs.values):
return False
analyzer = Analyzer()
@@ -107,7 +107,7 @@ def _fit_shape(bb: BlockBuilder, input_grad: Expr, input: Expr) -> Expr:
return (
input_grad
- if _check_shape_equal(expr_sinfo, target_sinfo)
+ if _check_shape_equal(expr_ty, target_ty)
else collapse_sum_to(input_grad, target_shape)
)
@@ -736,13 +736,13 @@ def concat_grad(
assert axis is not None
axis = int(axis)
split_indices: list[PrimExpr] = []
- sinfo = orig_call.args[0].struct_info
- assert isinstance(sinfo, relax.TupleStructInfo)
- for i in range(len(sinfo.fields) - 1):
- tensor_sinfo = sinfo.fields[i]
- assert isinstance(tensor_sinfo, relax.TensorStructInfo)
- assert tensor_sinfo.shape is not None
- index = tensor_sinfo.shape[axis]
+ ty = orig_call.args[0].ty
+ assert isinstance(ty, relax.TupleType)
+ for i in range(len(ty.fields) - 1):
+ tensor_ty = ty.fields[i]
+ assert isinstance(tensor_ty, relax.TensorType)
+ assert tensor_ty.shape is not None
+ index = tensor_ty.shape[axis]
if i > 0:
index += split_indices[i - 1]
split_indices.append(index)
@@ -1108,7 +1108,7 @@ def cross_entropy_with_logits_grad(
"""
x, y = orig_call.args
- if x.struct_info.ndim > 1:
+ if x.ty.ndim > 1:
batch_size = int(_get_shape(x)[0])
output_grad = output_grad / relax.const(batch_size, _get_dtype(output_grad))
diff --git a/python/tvm/relax/op/base.py b/python/tvm/relax/op/base.py
index 04f12d087f..4534c3efc4 100644
--- a/python/tvm/relax/op/base.py
+++ b/python/tvm/relax/op/base.py
@@ -27,7 +27,7 @@ from tvm.runtime import Object, ObjectConvertible
from ...ir import PrimExpr
from ..expr import Call, Expr, ExternFunc, GlobalVar, ShapeExpr, StringImm, Var
-from ..struct_info import StructInfo, TensorStructInfo
+from ..type import TensorType, Type
from ..utils import convert_to_expr
from . import _ffi_api
@@ -82,10 +82,7 @@ def _wrap_inline_arg_tuple(args) -> Expr:
elif (
isinstance(args, Expr)
and not isinstance(args, tvm.relax.Tuple)
- and (
- args.struct_info_ is None
- or not isinstance(args.struct_info_, tvm.relax.TupleStructInfo)
- )
+ and (args.ty is None or not isinstance(args.ty, tvm.relax.TupleType))
):
return tvm.relax.Tuple([args])
else:
@@ -95,7 +92,7 @@ def _wrap_inline_arg_tuple(args) -> Expr:
def call_tir(
gvar: GlobalVar,
args: Expr,
- out_sinfo: TensorStructInfo | list[TensorStructInfo],
+ out_ty: TensorType | list[TensorType],
tir_vars: ShapeExpr | tuple[PrimExpr] | list[PrimExpr] | None = None,
) -> Call:
"""
@@ -109,10 +106,10 @@ def call_tir(
args : Expr
The input arguments.
- out_sinfo : Union[TensorStructInfo, List[TensorStructInfo]]
- The structure info of the call_tir output.
- It should be a single or a list of TensorStructInfo. Each one denotes the
- structure info of a returned tensor.
+ out_ty : Union[TensorType, List[TensorType]]
+ The type information of the call_tir output.
+ It should be a single or a list of TensorType. Each one denotes the
+ type information of a returned tensor.
tir_vars : Optional[Union[ShapeExpr, Tuple[PrimExpr], List[PrimExpr]]]
ShapeExpr representing a tuple of integers to unpack when calling func. Is null if not used
@@ -124,19 +121,19 @@ def call_tir(
"""
args = _wrap_inline_arg_tuple(args)
- if not isinstance(out_sinfo, list):
- out_sinfo = [out_sinfo]
+ if not isinstance(out_ty, list):
+ out_ty = [out_ty]
if isinstance(tir_vars, list | tuple):
tir_vars = ShapeExpr(tir_vars)
- return _ffi_api.call_tir(gvar, args, out_sinfo, tir_vars) # type: ignore
+ return _ffi_api.call_tir(gvar, args, out_ty, tir_vars) # type: ignore
def call_tir_with_grad(
gvar: GlobalVar,
args: Expr,
- out_sinfo: TensorStructInfo | list[TensorStructInfo],
+ out_ty: TensorType | list[TensorType],
te_grad_name: str,
te_grad_kwargs: dict[str, Object] | None = None,
tir_vars: ShapeExpr | tuple[PrimExpr] | list[PrimExpr] | None = None,
@@ -154,10 +151,10 @@ def call_tir_with_grad(
args : Expr
The input arguments.
- out_sinfo : Union[TensorStructInfo, List[TensorStructInfo]]
- The structure info of the call_tir_with_grad output.
- It should be a single or a list of TensorStructInfo. Each one denotes the
- structure info of a returned tensor.
+ out_ty : Union[TensorType, List[TensorType]]
+ The type information of the call_tir_with_grad output.
+ It should be a single or a list of TensorType. Each one denotes the
+ type information of a returned tensor.
te_grad_name : str
The registered name of the te gradient function associated with the call_tir_with_grad
@@ -177,8 +174,8 @@ def call_tir_with_grad(
"""
args = _wrap_inline_arg_tuple(args)
- if not isinstance(out_sinfo, list):
- out_sinfo = [out_sinfo]
+ if not isinstance(out_ty, list):
+ out_ty = [out_ty]
if isinstance(tir_vars, list | tuple):
tir_vars = ShapeExpr(tir_vars)
@@ -187,7 +184,7 @@ def call_tir_with_grad(
te_grad_kwargs = {}
return _ffi_api.call_tir_with_grad( # type: ignore
- gvar, args, out_sinfo, te_grad_name, te_grad_kwargs, tir_vars
+ gvar, args, out_ty, te_grad_name, te_grad_kwargs, tir_vars
)
@@ -195,7 +192,7 @@ def call_tir_inplace(
gvar: GlobalVar,
args: Expr,
inplace_indices: int | list[int],
- out_sinfo: TensorStructInfo | list[TensorStructInfo],
+ out_ty: TensorType | list[TensorType],
tir_vars: ShapeExpr | tuple[PrimExpr] | list[PrimExpr] | None = None,
) -> Call:
"""
@@ -227,11 +224,11 @@ def call_tir_inplace(
If `inplace_indices[i] = -1`, then the `i`th output will be a freshly allocated tensor.
At least one member of `inplace_indices` must not be -1.
- out_sinfo : Union[TensorStructInfo, List[TensorStructInfo]]
- The structure info of the call_tir_inplace output.
- It should be a single `TensorStructInfo` or a list of `TensorStructInfo`.
- Each one denotes the structure info of a returned tensor.
- If a list of `TensorStructInfo` is given, the result will be a tuple of `TensorStructInfo`.
+ out_ty : Union[TensorType, List[TensorType]]
+ The type information of the call_tir_inplace output.
+ It should be a single `TensorType` or a list of `TensorType`.
+ Each one denotes the type information of a returned tensor.
+ If a list of `TensorType` is given, the result will be a tuple of `TensorType`.
tir_vars : Optional[Union[ShapeExpr, Tuple[PrimExpr], List[PrimExpr]]]
ShapeExpr representing a tuple of integers to unpack when calling func. Is null if not used
@@ -246,8 +243,8 @@ def call_tir_inplace(
if not isinstance(inplace_indices, list):
inplace_indices = [inplace_indices]
- if not isinstance(out_sinfo, list):
- out_sinfo = [out_sinfo]
+ if not isinstance(out_ty, list):
+ out_ty = [out_ty]
if isinstance(tir_vars, list | tuple):
tir_vars = ShapeExpr(tir_vars)
@@ -256,7 +253,7 @@ def call_tir_inplace(
gvar,
args,
inplace_indices,
- out_sinfo,
+ out_ty,
tir_vars,
)
@@ -264,7 +261,7 @@ def call_tir_inplace(
def call_dps_packed(
func: str | Expr,
args: Expr,
- out_sinfo: TensorStructInfo | list[TensorStructInfo],
+ out_ty: TensorType | list[TensorType],
) -> Call:
"""
Call a destination-passing-style packed function and return the output.
@@ -281,10 +278,10 @@ def call_dps_packed(
args : Expr
The input arguments.
- out_sinfo : Union[TensorStructInfo, List[TensorStructInfo]]
- The structure info of the call_dps_packed output.
- It should be a single or a list of TensorStructInfo. Each one denotes the
- structure info of a returned tensor.
+ out_ty : Union[TensorType, List[TensorType]]
+ The type information of the call_dps_packed output.
+ It should be a single or a list of TensorType. Each one denotes the
+ type information of a returned tensor.
Returns
-------
@@ -296,16 +293,16 @@ def call_dps_packed(
args = _wrap_inline_arg_tuple(args)
- if not isinstance(out_sinfo, list):
- out_sinfo = [out_sinfo]
+ if not isinstance(out_ty, list):
+ out_ty = [out_ty]
- return _ffi_api.call_dps_packed(func, args, out_sinfo) # type: ignore
+ return _ffi_api.call_dps_packed(func, args, out_ty) # type: ignore
def call_py_func(
func_name: str,
args: Expr,
- out_sinfo: TensorStructInfo | list[TensorStructInfo],
+ out_ty: TensorType | list[TensorType],
) -> Call:
"""
Call a Python function and return the output.
@@ -319,10 +316,10 @@ def call_py_func(
args : Expr
The input arguments.
- out_sinfo : Union[TensorStructInfo, List[TensorStructInfo]]
- The structure info of the call_py_func output.
- It should be a single or a list of TensorStructInfo. Each one denotes the
- structure info of a returned tensor.
+ out_ty : Union[TensorType, List[TensorType]]
+ The type information of the call_py_func output.
+ It should be a single or a list of TensorType. Each one denotes the
+ type information of a returned tensor.
Returns
-------
@@ -331,17 +328,17 @@ def call_py_func(
"""
args = _wrap_inline_arg_tuple(args)
- if not isinstance(out_sinfo, list):
- out_sinfo = [out_sinfo]
+ if not isinstance(out_ty, list):
+ out_ty = [out_ty]
- return _ffi_api.call_py_func(func_name, args, out_sinfo) # type: ignore
+ return _ffi_api.call_py_func(func_name, args, out_ty) # type: ignore
def call_builtin_with_ctx(
func: str | Expr,
args: Expr,
*,
- sinfo_args: StructInfo | list[StructInfo] | None = None,
+ ty_args: Type | list[Type] | None = None,
) -> Call:
"""Call a builtin function func.
@@ -353,8 +350,8 @@ def call_builtin_with_ctx(
args : Expr
The input arguments.
- sinfo_args: Optional[Union[StructInfo, List[StructInfo]]]
- The struct info arguments to the call node.
+ ty_args: Optional[Union[Type, List[Type]]]
+ The type arguments to the call node.
Returns
-------
@@ -366,13 +363,13 @@ def call_builtin_with_ctx(
args = _wrap_inline_arg_tuple(args)
- if sinfo_args is not None and not isinstance(sinfo_args, list | tuple):
- sinfo_args = [sinfo_args]
+ 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,
- sinfo_args, # type: ignore
+ ty_args, # type: ignore
)
@@ -406,7 +403,7 @@ def make_closure(
def invoke_closure(
closure: Expr,
args: Expr,
- sinfo_args: list[StructInfo] | StructInfo,
+ ty_args: list[Type] | Type,
) -> Call:
"""
Invoke a closure.
@@ -419,8 +416,8 @@ def invoke_closure(
args : Expr
The input arguments.
- type_args: Union[List[StructInfo], StructInfo]
- The structure info arguments of the CallNode
+ type_args: Union[List[Type], Type]
+ The type information arguments of the CallNode
Returns
-------
@@ -429,10 +426,10 @@ def invoke_closure(
"""
args = _wrap_inline_arg_tuple(args)
- if not isinstance(sinfo_args, list | tuple):
- sinfo_args = [sinfo_args]
+ if not isinstance(ty_args, list | tuple):
+ ty_args = [ty_args]
- return _ffi_api.invoke_closure(closure, args, sinfo_args) # type: ignore
+ return _ffi_api.invoke_closure(closure, args, ty_args) # type: ignore
def render_object(val: tvm.Object) -> str:
@@ -681,7 +678,7 @@ def call_inplace_packed(
func: str | ExternFunc | GlobalVar,
*args: Expr,
inplace_indices: int | list[int],
- sinfo_args: StructInfo | list[StructInfo],
+ ty_args: Type | list[Type],
) -> Expr:
"""
Construct a call to a packed function that consumes some of its arguments "in-place"
@@ -717,36 +714,36 @@ def call_inplace_packed(
If `inplace_indices[i] = -1`, then the `i`th output will be a freshly allocated tensor.
At least one member of `inplace_indices` must not be -1.
- sinfo_args: Union[StructInfo, List[StructInfo]]
- The list of structure info arguments (giving the structural info for the returned value).
+ ty_args: Union[Type, List[Type]]
+ The list of type information arguments (giving the type information for the returned value).
Returns
-------
result : Expr
A Relax call, corresponding to
- `call_pure_packed(ExternFunc(func), args, DictAttrs(kwargs), sinfo_args)`
+ `call_pure_packed(ExternFunc(func), args, DictAttrs(kwargs), ty_args)`
"""
if isinstance(func, ExternFunc):
func = func.global_symbol
op = ExternFunc(func)
args = tuple(convert_to_expr(a) for a in args)
- if sinfo_args is None:
+ if ty_args is None:
raise ValueError("R.call_pure_packed is required to have type_args")
- if isinstance(sinfo_args, tuple): # type: ignore
- sinfo_args = list(sinfo_args)
- elif not isinstance(sinfo_args, list):
- sinfo_args = [sinfo_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):
inplace_indices = [inplace_indices]
- return _ffi_api.call_inplace_packed(op, args, inplace_indices, sinfo_args) # type: ignore # pylint: disable=no-member
+ return _ffi_api.call_inplace_packed(op, args, inplace_indices, ty_args) # type: ignore # pylint: disable=no-member
def call_pure_packed(
func: str | ExternFunc | GlobalVar,
*args: Expr,
- sinfo_args: StructInfo | list[StructInfo],
+ ty_args: Type | list[Type],
) -> Expr:
"""
Construct a call to a packed function that should be treated as pure,
@@ -768,14 +765,14 @@ def call_pure_packed(
args: Expr
The arguments for the PackedFunc.
- sinfo_args: Union[StructInfo, List[StructInfo]]
- The list of structure info arguments (giving the structural info for the returned value).
+ ty_args: Union[Type, List[Type]]
+ The list of type information arguments (giving the type information for the returned value).
Returns
-------
result : Expr
A Relax call, corresponding to
- `call_pure_packed(ExternFunc(func), args, DictAttrs(kwargs), sinfo_args)`
+ `call_pure_packed(ExternFunc(func), args, DictAttrs(kwargs), ty_args)`
"""
if isinstance(func, ExternFunc):
func = func.global_symbol
@@ -783,34 +780,28 @@ def call_pure_packed(
op = ExternFunc(func)
args = tuple(convert_to_expr(a) for a in args)
- if sinfo_args is None:
+ if ty_args is None:
raise ValueError("R.call_pure_packed is required to have type_args")
- if isinstance(sinfo_args, tuple): # type: ignore
- sinfo_args = list(sinfo_args)
- elif not isinstance(sinfo_args, list):
- sinfo_args = [sinfo_args]
-
- sinfo_args = [
- (
- sinfo()
- if callable(sinfo)
- else sinfo.asobject()
- if isinstance(sinfo, ObjectConvertible)
- else sinfo
- )
- for sinfo in sinfo_args
+ if isinstance(ty_args, tuple): # type: ignore
+ ty_args = list(ty_args)
+ elif not isinstance(ty_args, list):
+ ty_args = [ty_args]
+
+ ty_args = [
+ (ty() if callable(ty) else ty.asobject() if isinstance(ty, ObjectConvertible) else ty)
+ for ty in ty_args
]
# note: if we need attributes, we can also take them here
- return _ffi_api.call_pure_packed(op, args, None, sinfo_args) # type: ignore # pylint: disable=no-member
+ return _ffi_api.call_pure_packed(op, args, None, ty_args) # type: ignore # pylint: disable=no-member
def invoke_pure_closure(
closure: Expr,
args: Expr,
- sinfo_args: list[StructInfo] | StructInfo,
+ ty_args: list[Type] | Type,
) -> Call:
"""
Invoke a closure and indicate to the compiler that it is pure.
@@ -829,8 +820,8 @@ def invoke_pure_closure(
args : Expr
The input arguments.
- type_args: Union[List[StructInfo], StructInfo]
- The structure info arguments of the CallNode
+ type_args: Union[List[Type], Type]
+ The type information arguments of the CallNode
Returns
-------
@@ -839,10 +830,10 @@ def invoke_pure_closure(
"""
args = _wrap_inline_arg_tuple(args)
- if not isinstance(sinfo_args, list | tuple):
- sinfo_args = [sinfo_args]
+ if not isinstance(ty_args, list | tuple):
+ ty_args = [ty_args]
- return _ffi_api.invoke_pure_closure(closure, args, sinfo_args) # type: ignore
+ return _ffi_api.invoke_pure_closure(closure, args, ty_args) # type: ignore
def to_vdevice(data, dst_vdevice) -> Expr:
diff --git a/python/tvm/relax/op/binary.py b/python/tvm/relax/op/binary.py
index 9480612e6f..4b26226276 100644
--- a/python/tvm/relax/op/binary.py
+++ b/python/tvm/relax/op/binary.py
@@ -43,9 +43,9 @@ def add(x1: Expr, x2: Expr) -> Expr:
.. code:: python
bb = relax.BlockBuilder()
- a = relax.Var("a", relax.TensorStructInfo(shape=(2, 3), dtype="float32"))
- b = relax.Var("b", relax.TensorStructInfo(shape=(2, 1), dtype="float32"))
- c = bb.normalize(relax.op.add(a, b)) # c has TensorStructInfo(shape=(2, 3), dtype="float32")
+ a = relax.Var("a", relax.TensorType(shape=(2, 3), dtype="float32"))
+ 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
diff --git a/python/tvm/relax/op/distributed/distributed.py b/python/tvm/relax/op/distributed/distributed.py
index 07ff674dd0..aa35125257 100644
--- a/python/tvm/relax/op/distributed/distributed.py
+++ b/python/tvm/relax/op/distributed/distributed.py
@@ -18,8 +18,7 @@
"""Operators for distributed Relax."""
from tvm.ir import PrimExpr
-from tvm.relax.distributed import DTensorStructInfo
-from tvm.relax.distributed.struct_info import DeviceMesh, Placement
+from tvm.relax.distributed import DeviceMesh, DTensorType, Placement
from ...expr import Call, Expr, GlobalVar, ShapeExpr
from ...expr import Tuple as RxTuple
@@ -69,7 +68,7 @@ def redistribute(input: Expr, device_mesh: DeviceMesh, placement: Placement) ->
def call_tir_local_view(
gvar: GlobalVar,
args: Expr,
- out_sinfo: DTensorStructInfo | list[DTensorStructInfo],
+ out_ty: DTensorType | list[DTensorType],
tir_vars: ShapeExpr | tuple[PrimExpr] | list[PrimExpr] | None = None,
) -> Call:
"""
@@ -85,10 +84,10 @@ def call_tir_local_view(
args : Expr
The input arguments.
- out_sinfo : Union[DTensorStructInfo, List[DTensorStructInfo]]
- The structure info of the call_tir output.
- It should be a single or a list of DTensorStructInfo. Each one denotes the
- structure info of a returned tensor.
+ out_ty : Union[DTensorType, List[DTensorType]]
+ The type information of the call_tir output.
+ It should be a single or a list of DTensorType. Each one denotes the
+ type information of a returned tensor.
tir_vars : Optional[Union[ShapeExpr, Tuple[PrimExpr], List[PrimExpr]]]
ShapeExpr representing a tuple of integers to unpack when calling func. Is null if not used
@@ -103,19 +102,19 @@ def call_tir_local_view(
elif isinstance(args, Expr) and not isinstance(args, RxTuple): # type: ignore
args = RxTuple((args,))
- if not isinstance(out_sinfo, list):
- out_sinfo = [out_sinfo]
+ if not isinstance(out_ty, list):
+ out_ty = [out_ty]
if isinstance(tir_vars, list | tuple):
tir_vars = ShapeExpr(tir_vars)
- return _ffi_api.call_tir_local_view(gvar, args, out_sinfo, tir_vars) # type: ignore
+ return _ffi_api.call_tir_local_view(gvar, args, out_ty, tir_vars) # type: ignore
def redistribute_replica_to_shard(input: Expr, num_workers: int, axis: int) -> Expr:
"""Slice tensor into several parts along one axis,
and each worker takes one part.
- input.struct_info.shape[axis] % num_workers == 0 is required.
+ input.ty.shape[axis] % num_workers == 0 is required.
Each worker must have an identical copy of the input.
This is a specialized version of redistribute op.
diff --git a/python/tvm/relax/op/grad/grad.py b/python/tvm/relax/op/grad/grad.py
index 45ceda6fdb..86c0792b54 100644
--- a/python/tvm/relax/op/grad/grad.py
+++ b/python/tvm/relax/op/grad/grad.py
@@ -18,8 +18,8 @@
"""Operators to implement operaor gradients. Used in `_op_gradient.py`.
We are trying to keep grad operators as simple as possible, and hope they are only used for finding
-gradients for forward operators. The struct_info inference for grad operators just returns the
-struct_info of the input.
+gradients for forward operators. The ty inference for grad operators just returns the
+ty of the input.
"""
from ...expr import Expr
@@ -52,8 +52,8 @@ def start_checkpoint(input: Expr) -> Expr:
For instance,
```
- a = relax.Var("a", relax.TensorStructInfo((2, 2), "float32"))
- b = relax.Var("b", relax.TensorStructInfo((2, 2), "float32"))
+ a = relax.Var("a", relax.TensorType((2, 2), "float32"))
+ b = relax.Var("b", relax.TensorType((2, 2), "float32"))
c = a * 2
d = b * 2
c_cp = start_checkpoint(c)
diff --git a/python/tvm/relax/op/manipulate.py b/python/tvm/relax/op/manipulate.py
index 21fd7b565c..e4814bc62a 100644
--- a/python/tvm/relax/op/manipulate.py
+++ b/python/tvm/relax/op/manipulate.py
@@ -144,7 +144,7 @@ def layout_transform(
if callable(index_map):
index_map = IndexMap.from_func(index_map, index_dtype=default_index_dtype)
- x_dtype = x.struct_info.dtype
+ x_dtype = x.ty.dtype
# Explicitly convert python int/float pad_value to the x's type. If the default behavior
# is applied, it would be converted to int32/float32, which may not match the x's type.
diff --git a/python/tvm/relax/op/mask.py b/python/tvm/relax/op/mask.py
index 3c013768a2..a4f8218849 100644
--- a/python/tvm/relax/op/mask.py
+++ b/python/tvm/relax/op/mask.py
@@ -35,5 +35,5 @@ def masked_fill(x: Expr, mask: Expr, value: Expr):
result : relax.Expr
The filled tensor.
"""
- values = _ffi_api.full_like(x, value, value.struct_info.dtype) # type: ignore
+ values = _ffi_api.full_like(x, value, value.ty.dtype) # type: ignore
return _ffi_api.where(mask, values, x) # 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 6830b1dc63..e4e41a9873 100644
--- a/python/tvm/relax/op/vision/multibox_transform_loc.py
+++ b/python/tvm/relax/op/vision/multibox_transform_loc.py
@@ -60,7 +60,7 @@ def multibox_transform_loc(
Notes
-----
- **Shape/dtype (checked in ``FInferStructInfo`` when static):**
+ **Shape/dtype (checked in ``FInferType`` when static):**
- ``cls_pred``: 3-D; ``loc_pred``: 2-D; ``anchor``: 3-D.
- ``cls_pred``, ``loc_pred``, ``anchor`` dtypes must match.
@@ -69,7 +69,7 @@ def multibox_transform_loc(
- ``cls_pred.shape[0]`` must equal ``loc_pred.shape[0]`` (batch).
If ``cls_pred`` has **unknown** shape, inference only returns generic rank-3 tensor
- struct info for the two outputs; it does **not** verify ``4*N`` vs ``loc_pred`` or
+ type for the two outputs; it does **not** verify ``4*N`` vs ``loc_pred`` or
``anchor.shape[1]`` vs ``N``, because ``N`` is not available statically. Other checks
(ranks, dtypes, ``loc_pred.shape[1] % 4 == 0`` when known, batch match when both batch
axes are known, etc.) still run where applicable.
diff --git a/python/tvm/relax/relax_to_pyfunc_converter.py b/python/tvm/relax/relax_to_pyfunc_converter.py
index bd15a85610..970e51360e 100644
--- a/python/tvm/relax/relax_to_pyfunc_converter.py
+++ b/python/tvm/relax/relax_to_pyfunc_converter.py
@@ -426,8 +426,8 @@ class RelaxExpressionConverter:
return self.variable_map[var_name]
# Try to infer shape from var's type annotation
- if hasattr(var, "struct_info") and hasattr(var.struct_info, "shape"):
- shape = var.struct_info.shape
+ if hasattr(var, "ty") and hasattr(var.ty, "shape"):
+ shape = var.ty.shape
if shape and len(shape) > 0:
# Convert symbolic shapes to concrete values
concrete_shape = []
@@ -599,7 +599,7 @@ class RelaxExpressionConverter:
# Extract TIR function name and arguments
tir_func = call.args[0]
tir_args = call.args[1] if len(call.args) > 1 else []
- out_sinfo = call.attrs.get("out_sinfo") if call.attrs else None
+ out_ty = call.attrs.get("out_ty") if call.attrs else None
# Get function name
if isinstance(tir_func, relax.GlobalVar):
@@ -660,8 +660,8 @@ class RelaxExpressionConverter:
# For call_tir, we need to allocate output tensor
output_shape = None
- if out_sinfo and hasattr(out_sinfo, "shape"):
- output_shape = out_sinfo.shape
+ if out_ty and hasattr(out_ty, "shape"):
+ output_shape = out_ty.shape
elif converted_args:
# Use the shape of the first input tensor
first_arg = converted_args[0]
@@ -713,7 +713,7 @@ class RelaxExpressionConverter:
# Extract packed function name and arguments
packed_func = call.args[0]
packed_args = call.args[1] if len(call.args) > 1 else []
- _out_sinfo = call.attrs.get("out_sinfo") if call.attrs else None
+ _out_ty = call.attrs.get("out_ty") if call.attrs else None
# Get function name
if isinstance(packed_func, relax.GlobalVar):
diff --git a/python/tvm/relax/script/builder/distributed/ir.py b/python/tvm/relax/script/builder/distributed/ir.py
index a74727c995..798fb99c67 100644
--- a/python/tvm/relax/script/builder/distributed/ir.py
+++ b/python/tvm/relax/script/builder/distributed/ir.py
@@ -27,7 +27,7 @@ import numpy as _np # type: ignore
import tvm
from tvm import base as _base
from tvm.ir import PrimExpr
-from tvm.relax.distributed import DeviceMesh, DTensorStructInfo, Placement
+from tvm.relax.distributed import DeviceMesh, DTensorType, Placement
from tvm.relax.expr import Call, Constant, Expr, ExternFunc, ShapeExpr
from tvm.relax.expr import Tuple as RxTuple
from tvm.relax.op.distributed import (
@@ -52,7 +52,7 @@ from . import _ffi_api
def call_tir(
func: str | Expr,
args: Expr,
- out_sinfo: DTensorStructInfo | list[DTensorStructInfo],
+ out_ty: DTensorType | list[DTensorType],
... 58244 lines suppressed ...