You are viewing a plain text version of this content. The canonical link for it is here.
Posted to commits@tvm.apache.org by tq...@apache.org on 2026/07/15 09:12:48 UTC
(tvm) branch main updated: [IR][Relax][TIRx] Unify Var identity (#20004)
This is an automated email from the ASF dual-hosted git repository.
tqchen pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/tvm.git
The following commit(s) were added to refs/heads/main by this push:
new c717c5b217 [IR][Relax][TIRx] Unify Var identity (#20004)
c717c5b217 is described below
commit c717c5b217f5360c52f55f1ff0eb5be4c209b2b3
Author: Tianqi Chen <tq...@users.noreply.github.com>
AuthorDate: Wed Jul 15 17:12:40 2026 +0800
[IR][Relax][TIRx] Unify Var identity (#20004)
---
docs/reference/api/python/relax/relax.rst | 2 +-
docs/tirx/api/tirx.rst | 2 +-
docs/tirx/native_basics/cuda/data_types.rst | 2 +-
include/tvm/arith/int_solver.h | 7 +-
include/tvm/arith/iter_affine_map.h | 13 +-
include/tvm/arith/pattern.h | 4 +-
include/tvm/ir/expr.h | 57 ++++-
include/tvm/ir/node_functor.h | 24 +-
include/tvm/relax/analysis.h | 43 ++--
include/tvm/relax/distributed/axis_group_graph.h | 7 +-
include/tvm/relax/expr.h | 57 +----
include/tvm/relax/expr_functor.h | 131 ++++++++++-
include/tvm/relax/script/builder/frame.h | 8 +-
include/tvm/relax/script/builder/ir.h | 12 +-
include/tvm/relax/transform.h | 5 +-
include/tvm/relax/type.h | 2 +-
include/tvm/relax/utils.h | 16 +-
include/tvm/tirx/stmt_functor.h | 22 +-
include/tvm/tirx/var.h | 118 +---------
include/tvm/topi/transform.h | 6 +-
.../cuda/operator/tile_primitive/copy/fallback.py | 4 +-
.../cuda/operator/tile_primitive/copy_async/tma.py | 4 +-
python/tvm/ir/__init__.py | 2 +-
python/tvm/ir/expr.py | 141 ++++++++----
python/tvm/ir/json_compact.py | 12 +
python/tvm/relax/backend/cuda/cublas.py | 4 +-
python/tvm/relax/backend/cuda/cutlass.py | 2 +-
python/tvm/relax/backend/metal/coreml.py | 4 +-
python/tvm/relax/backend/rocm/hipblas.py | 4 +-
python/tvm/relax/block_builder.py | 4 +-
python/tvm/relax/expr.py | 74 ++-----
python/tvm/relax/expr_functor.py | 142 +++++++++++-
python/tvm/relax/frontend/nn/core.py | 8 +-
python/tvm/relax/frontend/nn/exporter.py | 9 +-
python/tvm/relax/frontend/nn/op.py | 4 +-
python/tvm/relax/frontend/onnx/onnx_frontend.py | 44 ++--
.../frontend/stablehlo/stablehlo_translator.py | 13 +-
.../tvm/relax/frontend/tflite/tflite_frontend.py | 6 +-
.../frontend/torch/base_fx_graph_translator.py | 11 +-
.../frontend/torch/exported_program_translator.py | 6 +-
python/tvm/relax/frontend/torch/fx_translator.py | 11 +-
python/tvm/relax/relax_to_pyfunc_converter.py | 58 +++--
python/tvm/relax/script/parser/parser.py | 202 ++++++++++-------
python/tvm/relax/testing/ast_printer.py | 9 +-
python/tvm/relax/testing/nn.py | 10 +-
python/tvm/relax/training/loss.py | 8 +-
python/tvm/relax/training/optimizer.py | 14 +-
python/tvm/relax/training/trainer.py | 4 +-
python/tvm/relax/utils.py | 25 ++-
.../tvm/s_tir/dlight/analysis/common_analysis.py | 12 +-
python/tvm/s_tir/dlight/benchmark/extract.py | 12 +-
python/tvm/s_tir/dlight/benchmark/utils.py | 7 +-
python/tvm/s_tir/dlight/gpu/general_reduction.py | 4 +-
python/tvm/script/ir_builder/ir/__init__.py | 1 +
python/tvm/script/ir_builder/ir/ir.py | 31 +++
python/tvm/script/parser/core/parser.py | 9 +-
python/tvm/script/parser/ir/__init__.py | 1 +
python/tvm/tirx/expr.py | 25 +--
python/tvm/tirx/function.py | 13 +-
python/tvm/tirx/script/builder/ir.py | 36 +--
python/tvm/tirx/script/builder/triton.py | 4 +-
python/tvm/tirx/script/parser/parser.py | 40 ++--
python/tvm/tirx/transform/transform.py | 6 +-
python/tvm/topi/nn/batch_matmul.py | 4 +-
python/tvm/topi/nn/dense.py | 15 +-
python/tvm/topi/scatter.py | 7 +-
python/tvm/topi/signal.py | 4 +-
python/tvm/topi/utils.py | 5 +-
src/arith/analyzer.cc | 4 +-
src/arith/detect_linear_equation.cc | 26 ++-
src/arith/int_constraints.cc | 21 +-
src/arith/int_set.cc | 12 +-
src/arith/ir_mutator_with_analyzer.cc | 2 +-
src/arith/ir_mutator_with_analyzer.h | 11 +-
src/arith/iter_affine_map.cc | 42 ++--
src/arith/presburger_set.cc | 14 +-
src/arith/presburger_set.h | 15 +-
src/arith/rewrite_simplify.cc | 2 +-
src/arith/solve_linear_equation.cc | 18 +-
src/arith/solve_linear_inequality.cc | 22 +-
src/backend/cuda/codegen/codegen_cuda.cc | 2 +-
src/ir/expr.cc | 33 +++
src/relax/analysis/analysis.cc | 16 ++
src/relax/analysis/computable_at_compile_time.cc | 20 +-
src/relax/analysis/layout_transformation.cc | 12 +-
src/relax/analysis/tir_op_pattern_kind.cc | 20 +-
src/relax/analysis/type_analysis.cc | 179 +++++++--------
src/relax/analysis/udchain.cc | 1 +
src/relax/analysis/well_formed.cc | 168 +++++++-------
src/relax/backend/contrib/cutlass/codegen.cc | 2 +
src/relax/backend/contrib/tensorrt/codegen.cc | 6 +-
src/relax/backend/vm/codegen_vm.cc | 22 +-
src/relax/backend/vm/codegen_vm_tir.cc | 39 +++-
src/relax/backend/vm/vm_shape_lower.cc | 128 +++++++++--
src/relax/distributed/axis_group_graph.cc | 18 +-
.../transform/lower_global_view_to_local_view.cc | 5 +-
src/relax/ir/block_builder.cc | 35 +--
src/relax/ir/dataflow_expr_rewriter.cc | 4 +-
src/relax/ir/dataflow_matcher.cc | 4 +-
src/relax/ir/expr.cc | 30 +--
src/relax/ir/expr_functor.cc | 243 ++++++++++++++++++---
src/relax/ir/py_expr_functor.cc | 4 +-
src/relax/ir/transform.cc | 86 +++-----
src/relax/op/memory/view.cc | 6 +-
src/relax/op/tensor/manipulate.cc | 4 +-
src/relax/script/builder/frame.cc | 12 +-
src/relax/script/builder/ir.cc | 32 +--
src/relax/script/builder/utils.h | 4 +-
src/relax/script/printer/binding.cc | 14 +-
src/relax/script/printer/call.cc | 5 +-
src/relax/script/printer/dependent_type.cc | 4 +-
src/relax/script/printer/expr.cc | 18 +-
src/relax/script/printer/function.cc | 14 +-
src/relax/script/printer/tir.cc | 28 +--
src/relax/script/printer/utils.h | 10 +-
src/relax/training/utils.cc | 6 +-
src/relax/transform/bind_params.cc | 84 +------
src/relax/transform/bind_symbolic_vars.cc | 21 +-
src/relax/transform/bundle_model_params.cc | 89 +++++++-
src/relax/transform/canonicalize_bindings.cc | 114 +++++++---
src/relax/transform/compute_prim_value.cc | 47 +++-
src/relax/transform/decompose_ops.cc | 2 +-
src/relax/transform/fuse_ops.cc | 3 +-
src/relax/transform/fuse_tir.cc | 32 ++-
src/relax/transform/lazy_transform_params.cc | 15 +-
src/relax/transform/lift_transform_params.cc | 115 ++++------
src/relax/transform/realize_vdevice.cc | 7 +-
src/relax/transform/remove_unused_parameters.cc | 2 +-
src/relax/transform/rewrite_cuda_graph.cc | 86 ++++----
.../specialize_primfunc_based_on_callsite.cc | 2 +-
src/relax/transform/split_call_tir_by_pattern.cc | 44 ++--
src/relax/transform/topological_sort.cc | 23 +-
src/relax/transform/utils.h | 126 ++++++-----
src/relax/utils.cc | 104 +++++----
src/s_tir/analysis/identify_memcpy.cc | 2 +-
.../analysis/sblock_buffer_access_lca_detector.cc | 6 +-
.../feature_extractor/per_store_feature.cc | 4 +-
.../postproc/disallow_async_strided_mem_copy.cc | 2 +-
.../multi_level_tiling_tensor_core.cc | 38 ++--
.../multi_level_tiling_wide_vector.cc | 2 +-
src/s_tir/schedule/analysis/analysis.cc | 22 +-
src/s_tir/schedule/analysis/layout.cc | 14 +-
src/s_tir/schedule/analysis/reducer.cc | 5 +-
src/s_tir/schedule/instruction.cc | 4 +-
src/s_tir/schedule/ir_comparator.cc | 6 +-
.../schedule/primitive/annotate_buffer_access.cc | 17 +-
src/s_tir/schedule/primitive/blockize_tensorize.cc | 20 +-
src/s_tir/schedule/primitive/cache_index.cc | 8 +-
.../schedule/primitive/cache_index_helpers.cc | 4 +-
src/s_tir/schedule/primitive/compute_at.cc | 6 +-
src/s_tir/schedule/primitive/compute_inline.cc | 10 +-
src/s_tir/schedule/primitive/decompose_padding.cc | 8 +-
.../schedule/primitive/layout_transformation.cc | 4 +-
.../schedule/primitive/loop_transformation.cc | 18 +-
src/s_tir/schedule/primitive/pad_einsum.cc | 22 +-
src/s_tir/schedule/primitive/read_write_at.cc | 11 +-
src/s_tir/schedule/transform.cc | 18 +-
src/s_tir/schedule/utils.h | 4 +-
src/s_tir/transform/inject_double_buffer.cc | 2 +-
.../transform/lower_cross_thread_reduction.cc | 7 +-
src/s_tir/transform/lower_thread_allreduce.cc | 3 +-
src/s_tir/transform/memhammer_coalesce.cc | 2 +-
src/s_tir/transform/memhammer_lower_auto_copy.cc | 4 +-
.../remove_weight_layout_rewrite_block.cc | 2 +-
src/s_tir/transform/renew_defs.cc | 8 +-
src/s_tir/transform/renormalize_split_pattern.cc | 3 +-
src/s_tir/transform/storage_access.cc | 4 +-
src/s_tir/transform/storage_access.h | 2 +-
.../transform/using_assume_to_reduce_branches.cc | 3 +-
src/script/ir_builder/ir/ir.cc | 8 +
src/te/operation/create_primfunc.cc | 6 +-
src/te/tensor.cc | 6 +-
src/tirx/analysis/verify_well_formed.cc | 12 +-
src/tirx/ir/buffer.cc | 2 +-
src/tirx/ir/expr.cc | 47 ----
src/tirx/ir/index_map.cc | 29 ++-
src/tirx/ir/specialize.cc | 10 +-
src/tirx/ir/stmt.cc | 2 +-
src/tirx/ir/stmt_functor.cc | 10 +-
src/tirx/op/op.cc | 7 +-
src/tirx/op/tirx.cc | 1 +
src/tirx/script/builder/ir.cc | 28 +--
src/tirx/script/printer/block.cc | 6 +-
src/tirx/script/printer/expr.cc | 12 +-
src/tirx/transform/ir_utils.cc | 10 +-
src/tirx/transform/lower_warp_memory.cc | 3 +-
src/tirx/transform/make_packed_api.cc | 2 +-
src/tirx/transform/stmt_simplify.cc | 3 +-
src/tirx/transform/tile_primitive_dispatch.cc | 21 +-
src/tirx/transform/tvm_ffi_binder.cc | 6 +-
tests/cpp/ir_functor_test.cc | 14 +-
tests/cpp/nested_msg_test.cc | 22 +-
tests/python/arith/test_arith_deduce_bound.py | 20 +-
tests/python/arith/test_arith_simplify.py | 5 +-
tests/python/ir/test_node_reflection.py | 107 ++++++++-
.../adreno/test_transform_annotate_custom_scope.py | 9 +-
tests/python/relax/test_analysis.py | 26 ++-
.../test_analysis_computable_at_compile_time.py | 2 +-
tests/python/relax/test_analysis_well_formed.py | 26 +++
tests/python/relax/test_ast_printer.py | 46 ++--
tests/python/relax/test_bind_symbolic_vars.py | 4 +-
tests/python/relax/test_blockbuilder_core.py | 30 +--
tests/python/relax/test_blockbuilder_emit_te.py | 7 +-
tests/python/relax/test_dataflow_inplace.py | 14 +-
tests/python/relax/test_dataflow_pattern.py | 10 +-
tests/python/relax/test_dataflow_rewriter.py | 19 +-
tests/python/relax/test_expr.py | 44 ++--
tests/python/relax/test_expr_functor.py | 22 +-
tests/python/relax/test_frontend_nn_packing.py | 2 +-
tests/python/relax/test_frontend_onnx.py | 12 +-
tests/python/relax/test_frontend_onnx_backend.py | 2 +-
tests/python/relax/test_frontend_tflite.py | 33 ++-
tests/python/relax/test_op_manipulate.py | 2 +-
tests/python/relax/test_transform.py | 52 ++++-
tests/python/relax/test_transform_bind_params.py | 68 +++++-
.../relax/test_transform_bind_symbolic_vars.py | 20 +-
.../relax/test_transform_bundle_model_params.py | 143 +++++++++++-
.../relax/test_transform_canonicalize_bindings.py | 27 ++-
.../test_transform_combine_parallel_matmul.py | 2 +-
.../python/relax/test_transform_convert_layout.py | 54 ++++-
.../relax/test_transform_error_enrichment.py | 4 +-
.../relax/test_transform_fuse_ops_by_pattern.py | 4 +-
tests/python/relax/test_transform_fuse_tir.py | 33 +++
tests/python/relax/test_transform_lambda_lift.py | 28 +++
..._transform_legalize_ops_index_linear_algebra.py | 4 +-
.../test_transform_legalize_ops_manipulate.py | 12 +-
.../python/relax/test_transform_legalize_ops_nn.py | 6 +-
.../relax/test_transform_lift_transform_params.py | 51 ++++-
.../relax/test_transform_update_param_type.py | 2 +-
tests/python/relax/test_tvmscript_ir_builder.py | 12 +-
tests/python/relax/test_tvmscript_parser.py | 229 +++++++++++++++++--
tests/python/relax/test_tvmscript_printer_relax.py | 17 +-
.../s_tir/schedule/test_tir_schedule_analysis.py | 2 +-
tests/python/tirx-base/test_tir_base.py | 6 +-
tests/python/tirx-base/test_tir_constructor.py | 2 +-
tests/python/tirx-base/test_tir_nodes.py | 10 +-
tests/python/tirx-base/test_tir_op_types.py | 42 ++--
tests/python/tirx-transform/test_tir_functor.py | 42 ++--
.../tirx-transform/test_tir_transform_vectorize.py | 4 +-
.../cuda/gemm_async/test_gemm_async.py | 8 +-
tests/python/tirx/test_layout.py | 2 +-
tests/python/tirx/test_parser_printer.py | 24 +-
242 files changed, 3743 insertions(+), 2123 deletions(-)
diff --git a/docs/reference/api/python/relax/relax.rst b/docs/reference/api/python/relax/relax.rst
index 7b060932aa..fefd074f00 100644
--- a/docs/reference/api/python/relax/relax.rst
+++ b/docs/reference/api/python/relax/relax.rst
@@ -20,4 +20,4 @@ tvm.relax
.. automodule:: tvm.relax
:members:
:imported-members:
- :exclude-members: BlockBuilder, Call, Span, GlobalVar, SourceName, TupleType, Type, FuncType
+ :exclude-members: BlockBuilder, Call, Var, Span, GlobalVar, SourceName, TupleType, Type, FuncType
diff --git a/docs/tirx/api/tirx.rst b/docs/tirx/api/tirx.rst
index df5716f911..af54c83fb1 100644
--- a/docs/tirx/api/tirx.rst
+++ b/docs/tirx/api/tirx.rst
@@ -20,4 +20,4 @@ tvm.tirx
.. automodule:: tvm.tirx
:members:
:imported-members:
- :exclude-members: Expr, PrimExpr, Op, Call, const
+ :exclude-members: Expr, PrimExpr, Op, Call, Var, const
diff --git a/docs/tirx/native_basics/cuda/data_types.rst b/docs/tirx/native_basics/cuda/data_types.rst
index 472ac8f0ef..4b17543566 100644
--- a/docs/tirx/native_basics/cuda/data_types.rst
+++ b/docs/tirx/native_basics/cuda/data_types.rst
@@ -111,6 +111,6 @@ A buffer's ``data`` — its pointer — is a ``Var`` of pointer type, and it is
from tvm.ir.type import PointerType, PrimType
- ptr: T.let[T.Var(name="ptr", dtype=PointerType(PrimType("uint64")))] = \
+ ptr: T.let[T.Var(name="ptr", ty=PointerType(PrimType("uint64")))] = \
T.reinterpret("handle", T.ptx.map_shared_rank(mbar.ptr_to([0]), 0))
remote_mbar = T.decl_buffer([1], "uint64", data=ptr, scope="shared")
diff --git a/include/tvm/arith/int_solver.h b/include/tvm/arith/int_solver.h
index fa4847b003..a01fe589ce 100644
--- a/include/tvm/arith/int_solver.h
+++ b/include/tvm/arith/int_solver.h
@@ -38,6 +38,7 @@ namespace tvm {
namespace arith {
using tirx::IterVar;
+using tirx::PrimVar;
using tirx::Var;
using tirx::VarNode;
@@ -132,7 +133,7 @@ class IntGroupBounds : public ffi::ObjectRef {
class IntConstraintsNode : public ffi::Object {
public:
// e.g., \alpha, \beta, must be integers
- ffi::Array<Var> variables;
+ ffi::Array<PrimVar> variables;
// e.g., 1 <= \alpha <= N, etc.
// it is absolutely ok to include ranges for parameters
// (variables that are not in this->variables) in this map
@@ -166,7 +167,7 @@ class IntConstraints : public ffi::ObjectRef {
* \param relations The linear relations between the variables
* (either equations or inequalities)
*/
- TVM_DLL IntConstraints(ffi::Array<Var> variables, ffi::Map<Var, Range> ranges,
+ TVM_DLL IntConstraints(ffi::Array<PrimVar> variables, ffi::Map<Var, Range> ranges,
ffi::Array<PrimExpr> relations);
TVM_FFI_DEFINE_OBJECT_REF_METHODS_NULLABLE(IntConstraints, ffi::ObjectRef, IntConstraintsNode);
@@ -299,7 +300,7 @@ PartialSolvedInequalities SolveLinearInequalities(const IntConstraints& system_t
* \param bounds grouped boundary of the variables.
* \param relations other relations.
*/
-ffi::Array<PrimExpr> AsConditions(const ffi::Array<Var>& variables,
+ffi::Array<PrimExpr> AsConditions(const ffi::Array<PrimVar>& variables,
const ffi::Map<Var, IntGroupBounds>& bounds,
const ffi::Array<PrimExpr>& relations);
diff --git a/include/tvm/arith/iter_affine_map.h b/include/tvm/arith/iter_affine_map.h
index adae57b761..6e5e14b236 100644
--- a/include/tvm/arith/iter_affine_map.h
+++ b/include/tvm/arith/iter_affine_map.h
@@ -316,8 +316,9 @@ class IterMapResult : public ffi::ObjectRef {
* The return object's .indices is empty on failure.
*/
IterMapResult DetectIterMap(const ffi::Array<PrimExpr>& indices,
- const ffi::Map<Var, Range>& input_iters, const PrimExpr& predicate,
- IterMapLevel check_level, const arith::Analyzer& analyzer,
+ const ffi::Map<tirx::PrimVar, Range>& input_iters,
+ const PrimExpr& predicate, IterMapLevel check_level,
+ const arith::Analyzer& analyzer,
bool simplify_trivial_iterators = true);
/*!
@@ -332,7 +333,7 @@ IterMapResult DetectIterMap(const ffi::Array<PrimExpr>& indices,
* \return The indices after rewrite
*/
ffi::Array<PrimExpr> IterMapSimplify(const ffi::Array<PrimExpr>& indices,
- const ffi::Map<Var, Range>& input_iters,
+ const ffi::Map<tirx::PrimVar, Range>& input_iters,
const PrimExpr& input_pred, IterMapLevel check_level,
const arith::Analyzer& analyzer,
bool simplify_trivial_iterators = true);
@@ -388,8 +389,8 @@ ffi::Map<Var, PrimExpr> InverseAffineIterMap(const ffi::Array<IterSumExpr>& iter
Empty array if no match can be found.
*/
ffi::Array<ffi::Array<IterMark>> SubspaceDivide(const ffi::Array<PrimExpr>& bindings,
- const ffi::Map<Var, Range>& input_iters,
- const ffi::Array<Var>& sub_iters,
+ const ffi::Map<tirx::PrimVar, Range>& input_iters,
+ const ffi::Array<tirx::PrimVar>& sub_iters,
const PrimExpr& predicate, IterMapLevel check_level,
const arith::Analyzer& analyzer,
bool simplify_trivial_iterators = true);
@@ -417,7 +418,7 @@ PrimExpr NormalizeIterMapToExpr(const PrimExpr& expr);
* \param analyzer The input analyzer.
* \note This function is useful to detect iterator stride patterns.
*/
-IterSumExpr NormalizeToIterSum(PrimExpr index, const ffi::Map<Var, Range>& input_iters,
+IterSumExpr NormalizeToIterSum(PrimExpr index, const ffi::Map<tirx::PrimVar, Range>& input_iters,
const arith::Analyzer& analyzer);
} // namespace arith
diff --git a/include/tvm/arith/pattern.h b/include/tvm/arith/pattern.h
index c476e2f01a..ed2e139ab7 100644
--- a/include/tvm/arith/pattern.h
+++ b/include/tvm/arith/pattern.h
@@ -37,7 +37,7 @@ namespace arith {
* \param vars List of variables to be used in detection.
* \return [coeff[i]] if it is possible, empty array if it is not.
*/
-ffi::Array<PrimExpr> DetectLinearEquation(const PrimExpr& e, const ffi::Array<tirx::Var>& vars);
+ffi::Array<PrimExpr> DetectLinearEquation(const PrimExpr& e, const ffi::Array<tirx::PrimVar>& vars);
/*!
* \brief Detect if expression corresponds to clip bound of the vars
@@ -47,7 +47,7 @@ ffi::Array<PrimExpr> DetectLinearEquation(const PrimExpr& e, const ffi::Array<ti
* \return concat([min_value[i], max_value[i]]), None is returned if there is no min or max value
* return empty if the e does not match the pattern.
*/
-ffi::Array<PrimExpr> DetectClipBound(const PrimExpr& e, const ffi::Array<tirx::Var>& vars);
+ffi::Array<PrimExpr> DetectClipBound(const PrimExpr& e, const ffi::Array<tirx::PrimVar>& vars);
} // namespace arith
} // namespace tvm
diff --git a/include/tvm/ir/expr.h b/include/tvm/ir/expr.h
index dd168a2462..30fe958b95 100644
--- a/include/tvm/ir/expr.h
+++ b/include/tvm/ir/expr.h
@@ -259,6 +259,47 @@ TVM_DLL PrimExpr operator^(PrimExpr a, PrimExpr b);
*/
TVM_DLL PrimExpr operator~(PrimExpr a);
+/*!
+ * \brief A local variable in the IR.
+ *
+ * Variables are uniquely identified by their object identity. The name is a
+ * hint for printing and does not participate in structural equality or
+ * hashing.
+ */
+class VarNode : public ExprNode {
+ public:
+ /*! \brief The hint to the variable name. */
+ ffi::String name_hint;
+
+ static void RegisterReflection() {
+ namespace refl = tvm::ffi::reflection;
+ refl::ObjectDef<VarNode>().def_ro("name_hint", &VarNode::name_hint,
+ refl::AttachFieldFlag::SEqHashIgnore());
+ }
+
+ static constexpr TVMFFISEqHashKind _type_s_eq_hash_kind = kTVMFFISEqHashKindFreeVar;
+ static constexpr const uint32_t _type_child_slots = 1;
+ TVM_FFI_DECLARE_OBJECT_INFO("ir.Var", VarNode, ExprNode);
+};
+
+/*! \brief Managed reference to VarNode. */
+class Var : public Expr {
+ public:
+ TVM_DLL explicit Var(ffi::String name_hint, ffi::Optional<Type> ty_annotation,
+ Span span = Span());
+
+ /*! \brief Return a fresh ordinary Var with the same type and a new name. */
+ TVM_DLL Var copy_with_name(const ffi::String& name) const;
+
+ /*! \brief Return a fresh ordinary Var with a suffix appended to its name. */
+ TVM_DLL Var CopyWithSuffix(const ffi::String& suffix) const;
+
+ /*! \brief Return a fresh ordinary Var with a new primitive type. */
+ TVM_DLL Var copy_with_dtype(PrimType dtype) const;
+
+ TVM_FFI_DEFINE_OBJECT_REF_METHODS_NULLABLE(Var, Expr, VarNode);
+};
+
class GlobalVar;
/*!
* \brief Global variable that lives in the top-level module.
@@ -540,7 +581,7 @@ struct TypeTraits<FloatImm> : public ObjectRefWithFallbackTraitsBase<FloatImm, d
} // namespace ffi
} // namespace tvm
-/* \brief Allow tvm.GLobalVar as key in STL tables
+/* \brief Allow tvm.Var and tvm.GlobalVar as keys in STL tables
*
* For most IR expressions, it would be ambiguous whether the
* expression should follow reference equality or structural equality.
@@ -549,11 +590,23 @@ struct TypeTraits<FloatImm> : public ObjectRefWithFallbackTraitsBase<FloatImm, d
* tables.
*
* Providing `std::hash` and `std::equal_to` specializations for
- * `tvm::GlobalVar` allows it to be used as a key in STL tables. For
+ * `tvm::Var` and `tvm::GlobalVar` allows them to be used as keys in STL tables. For
* other IR expressions, the user must specify the type of equality
* used (e.g. `std::unordered_set<T, StructuralHash, StructuralEqual>`
* or `std::unordered_set<T, ffi::ObjectPtrHash, ffi::ObjectPtrEqual>`).
*/
+template <>
+struct std::hash<tvm::Var> {
+ std::size_t operator()(const tvm::Var& var) const { return tvm::ffi::ObjectPtrHash()(var); }
+};
+
+template <>
+struct std::equal_to<tvm::Var> {
+ bool operator()(const tvm::Var& var_a, const tvm::Var& var_b) const {
+ return tvm::ffi::ObjectPtrEqual()(var_a, var_b);
+ }
+};
+
template <>
struct std::hash<tvm::GlobalVar> {
std::size_t operator()(const tvm::GlobalVar& var) const { return tvm::ffi::ObjectPtrHash()(var); }
diff --git a/include/tvm/ir/node_functor.h b/include/tvm/ir/node_functor.h
index fc6c925da3..5eea5ea0e7 100644
--- a/include/tvm/ir/node_functor.h
+++ b/include/tvm/ir/node_functor.h
@@ -94,9 +94,27 @@ class NodeFunctor<R(const ffi::ObjectRef& n, Args...)> {
* \return The result.
*/
R operator()(const ffi::ObjectRef& n, Args... args) const {
- TVM_FFI_ICHECK(can_dispatch(n))
- << "NodeFunctor calls un-registered function on type " << n->GetTypeKey();
- return (*func_[n->type_index() - begin_type_index_])(n, std::forward<Args>(args)...);
+ uint32_t type_index = n->type_index();
+ if (type_index >= begin_type_index_) {
+ uint32_t index = type_index - begin_type_index_;
+ if (index < func_.size() && func_[index] != nullptr) {
+ return (*func_[index])(n, std::forward<Args>(args)...);
+ }
+ }
+
+ const TVMFFITypeInfo* type_info = TVMFFIGetTypeInfo(type_index);
+ for (int32_t i = type_info->type_depth - 1; i >= 0; --i) {
+ type_index = type_info->type_ancestors[i]->type_index;
+ if (type_index >= begin_type_index_) {
+ uint32_t index = type_index - begin_type_index_;
+ if (index < func_.size() && func_[index] != nullptr) {
+ return (*func_[index])(n, std::forward<Args>(args)...);
+ }
+ }
+ }
+ TVM_FFI_THROW(InternalError) << "NodeFunctor calls un-registered function on type "
+ << n->GetTypeKey();
+ throw;
}
/*!
* \brief set the dispatcher for type TNode
diff --git a/include/tvm/relax/analysis.h b/include/tvm/relax/analysis.h
index 248efd66a4..454e218b94 100644
--- a/include/tvm/relax/analysis.h
+++ b/include/tvm/relax/analysis.h
@@ -137,15 +137,15 @@ TVM_DLL Type DeriveCallRetType(const FuncType& finfo, const Call& call, const Bl
* However, at the level of function signature, only n, m are defined,
* k is undefined here.
*
- * When we call EraseToWellDefined(R.Tensor[(n + 1, k + 2)], fshape_var_map={n: n, m: m}),
+ * When we call EraseToWellDefined(R.Tensor[(n + 1, k + 2)], var_map={n: n, m: m}),
* we will obtain R.Tensor(ndim=2), which is an erased info that does not depend
* on k(which is undefined from parameter signature).
*
- * However, if we call EraseToWellDefined(R.Tensor[(n + 1, m)], fshape_var_map={n: n, m: m}),
+ * However, if we call EraseToWellDefined(R.Tensor[(n + 1, m)], var_map={n: n, m: m}),
* Then the return value will be R.Tensor[(n + 1, m)], because both n and m are defined.
*
* We can also make these var map to return a different expression.
- * For example, EraseToWellDefined(R.Tensor[(n + 1, m)], fshape_var_map={n: 2, m: m})
+ * For example, EraseToWellDefined(R.Tensor[(n + 1, m)], var_map={n: 2, m: m})
* will give us R.Tensor[(3, m)], where n get replaced by 2.
*
* Use this function in the following scenarios:
@@ -153,63 +153,48 @@ TVM_DLL Type DeriveCallRetType(const FuncType& finfo, const Call& call, const Bl
* - Decide the deduced return ty of a function that can be fully decided by params.
*
* \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.
* \param f_var_map callback function to specify
- * whether a var is defined in the target scope and the value it maps to,
+ * 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 type.
*/
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);
+ const Type& info, std::function<ffi::Optional<Expr>(const Var& var)> f_var_map = nullptr);
/*!
* \brief EraseToWellDefined overload using a caller-provided analyzer.
* \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.
* \param f_var_map callback function to specify
- * whether a var is defined in the target scope and the value it maps to,
+ * 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 type.
*/
-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);
+TVM_DLL Type EraseToWellDefined(const Type& info,
+ std::function<ffi::Optional<Expr>(const Var& var)> f_var_map,
+ const arith::Analyzer& ana);
/*!
* \brief EraseToWellDefined variant with map.
* \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.
* \param var_map map to specify
- * whether a var is defined in the target scope and the value it maps to,
+ * 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 type.
*/
-TVM_DLL Type EraseToWellDefined(const Type& info, ffi::Map<tirx::Var, PrimExpr> shape_var_map,
- ffi::Map<Var, Expr> var_map);
+TVM_DLL Type EraseToWellDefined(const Type& info, ffi::Map<Var, Expr> var_map);
/*!
* \brief EraseToWellDefined map overload using a caller-provided analyzer.
* \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.
* \param var_map map to specify
- * whether a var is defined in the target scope and the value it maps to,
+ * 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 type.
*/
-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);
+TVM_DLL Type EraseToWellDefined(const Type& info, ffi::Map<Var, Expr> var_map,
+ const arith::Analyzer& ana);
/*!
* \brief Fine grained result of base check.
diff --git a/include/tvm/relax/distributed/axis_group_graph.h b/include/tvm/relax/distributed/axis_group_graph.h
index 8417087190..e7945a1b65 100644
--- a/include/tvm/relax/distributed/axis_group_graph.h
+++ b/include/tvm/relax/distributed/axis_group_graph.h
@@ -126,14 +126,15 @@ class BufferAxisGraphExtractor : public StmtExprVisitor {
bool Match(PrimExpr a, PrimExpr buffer_shape_a, PrimExpr b, PrimExpr buffer_shape_b,
const arith::Analyzer& analyzer) {
- if (b.as<VarNode>()) {
+ if (b.as<PrimVar>()) {
std::swap(a, b);
std::swap(buffer_shape_a, buffer_shape_b);
}
- if (!a.as<VarNode>()) {
+ auto prim_var = a.as<PrimVar>();
+ if (!prim_var) {
return false;
}
- Var var = a.as_or_throw<Var>();
+ Var var = *prim_var;
analyzer->Bind(iter_var_range_);
b = analyzer->Simplify(b);
// index var `a` must access whole range of a specific buffer dimension
diff --git a/include/tvm/relax/expr.h b/include/tvm/relax/expr.h
index b8ed8cfc23..bb1248016a 100644
--- a/include/tvm/relax/expr.h
+++ b/include/tvm/relax/expr.h
@@ -62,7 +62,7 @@ class Tuple : public Expr {
* \brief Utility constructor to handle conversion to relax::Expr
*
* If the calling scope already has an array of a specific type of
- * relax expression (e.g. `ffi::Array<relax::Var>`), it must be converted
+ * relax expression (e.g. `ffi::Array<Var>`), it must be converted
* into an array of base type. This constructor handles the
* conversion to the base `ffi::Array<relax::Expr>`.
*
@@ -132,33 +132,6 @@ class ShapeExpr : public Expr {
TVM_DEFINE_OBJECT_REF_COW_METHOD(ShapeExprNode);
};
-/*! \brief The variable class for all Relax bindings. */
-class VarNode : public ExprNode {
- public:
- /*!
- * \brief The hint to the variable name.
- * \note Each variable is uniquely identified by its address.
- */
- ffi::String name_hint;
-
- static void RegisterReflection() {
- namespace refl = tvm::ffi::reflection;
- refl::ObjectDef<VarNode>().def_ro("name_hint", &VarNode::name_hint,
- refl::AttachFieldFlag::SEqHashIgnore());
- }
-
- static constexpr TVMFFISEqHashKind _type_s_eq_hash_kind = kTVMFFISEqHashKindFreeVar;
- static constexpr const uint32_t _type_child_slots = 1;
- TVM_FFI_DECLARE_OBJECT_INFO("relax.expr.Var", VarNode, ExprNode);
-};
-
-class Var : public Expr {
- public:
- TVM_DLL explicit Var(ffi::String name_hint, ffi::Optional<Type> ty_annotation,
- Span span = Span());
- TVM_FFI_DEFINE_OBJECT_REF_METHODS_NULLABLE(Var, Expr, VarNode);
-};
-
/*! \brief A sub-type of the variable node used to mark dataflow variables from
* normal visible "function local" bindings.
*/
@@ -670,32 +643,4 @@ TVM_DLL Expr GetShapeOf(const Expr& expr);
} // namespace relax
} // namespace tvm
-/* \brief Allow relax.Var as key in STL tables
- *
- * For most Relax expressions, it would be ambiguous whether the
- * expression should follow reference equality or structural equality.
- * This is not the case for variables, which do not contain nested
- * internal structure, and are frequently used as keys in lookup
- * tables.
- *
- * Providing `std::hash` and `std::equal_to` specializations for
- * `relax::Var` allows it to be used as a key in STL tables. For
- * `relax::Expr`, the user must specify the type of equality used
- * (e.g. `std::unordered_set<T, StructuralHash, StructuralEqual>` or
- * `std::unordered_set<T, ffi::ObjectPtrHash, ffi::ObjectPtrEqual>`).
- */
-template <>
-struct std::hash<tvm::relax::Var> {
- std::size_t operator()(const tvm::relax::Var& var) const {
- return tvm::ffi::ObjectPtrHash()(var);
- }
-};
-
-template <>
-struct std::equal_to<tvm::relax::Var> {
- bool operator()(const tvm::relax::Var& var_a, const tvm::relax::Var& var_b) const {
- return tvm::ffi::ObjectPtrEqual()(var_a, var_b);
- }
-};
-
#endif // TVM_RELAX_EXPR_H_
diff --git a/include/tvm/relax/expr_functor.h b/include/tvm/relax/expr_functor.h
index 7f1388803d..98fd2fa149 100644
--- a/include/tvm/relax/expr_functor.h
+++ b/include/tvm/relax/expr_functor.h
@@ -57,6 +57,12 @@ class ExprFunctor;
return VisitExprDefault_(op, std::forward<Args>(args)...); \
}
+#define EXPR_FUNCTOR_DISABLED \
+ final { \
+ TVM_FFI_THROW(TypeError) << "Relax does not support " << op->GetTypeKey() << " expressions"; \
+ throw; \
+ }
+
#define RELAX_EXPR_FUNCTOR_DISPATCH(OP) \
vtable.template set_dispatch<OP>([](const ffi::ObjectRef& n, TSelf* self, Args... args) { \
return self->VisitExpr_(static_cast<const OP*>(n.get()), std::forward<Args>(args)...); \
@@ -149,6 +155,36 @@ class ExprFunctor<R(const Expr& n, Args...)> {
virtual R VisitExpr_(const GlobalVarNode* op, Args... args) EXPR_FUNCTOR_DEFAULT;
virtual R VisitExpr_(const FunctionNode* op, Args... args) EXPR_FUNCTOR_DEFAULT;
virtual R VisitExpr_(const CallNode* op, Args... args) EXPR_FUNCTOR_DEFAULT;
+ virtual R VisitExpr_(const tirx::BufferLoadNode* op, Args... args) EXPR_FUNCTOR_DEFAULT;
+ virtual R VisitExpr_(const tirx::ProducerLoadNode* op, Args...) EXPR_FUNCTOR_DISABLED;
+ virtual R VisitExpr_(const tirx::LetNode* op, Args...) EXPR_FUNCTOR_DISABLED;
+ virtual R VisitExpr_(const tirx::ReduceNode* op, Args...) EXPR_FUNCTOR_DISABLED;
+ virtual R VisitExpr_(const tirx::AddNode* op, Args... args) EXPR_FUNCTOR_DEFAULT;
+ virtual R VisitExpr_(const tirx::SubNode* op, Args... args) EXPR_FUNCTOR_DEFAULT;
+ virtual R VisitExpr_(const tirx::MulNode* op, Args... args) EXPR_FUNCTOR_DEFAULT;
+ virtual R VisitExpr_(const tirx::DivNode* op, Args... args) EXPR_FUNCTOR_DEFAULT;
+ virtual R VisitExpr_(const tirx::ModNode* op, Args... args) EXPR_FUNCTOR_DEFAULT;
+ virtual R VisitExpr_(const tirx::FloorDivNode* op, Args... args) EXPR_FUNCTOR_DEFAULT;
+ virtual R VisitExpr_(const tirx::FloorModNode* op, Args... args) EXPR_FUNCTOR_DEFAULT;
+ virtual R VisitExpr_(const tirx::MinNode* op, Args... args) EXPR_FUNCTOR_DEFAULT;
+ virtual R VisitExpr_(const tirx::MaxNode* op, Args... args) EXPR_FUNCTOR_DEFAULT;
+ virtual R VisitExpr_(const tirx::EQNode* op, Args... args) EXPR_FUNCTOR_DEFAULT;
+ virtual R VisitExpr_(const tirx::NENode* op, Args... args) EXPR_FUNCTOR_DEFAULT;
+ virtual R VisitExpr_(const tirx::LTNode* op, Args... args) EXPR_FUNCTOR_DEFAULT;
+ virtual R VisitExpr_(const tirx::LENode* op, Args... args) EXPR_FUNCTOR_DEFAULT;
+ virtual R VisitExpr_(const tirx::GTNode* op, Args... args) EXPR_FUNCTOR_DEFAULT;
+ virtual R VisitExpr_(const tirx::GENode* op, Args... args) EXPR_FUNCTOR_DEFAULT;
+ virtual R VisitExpr_(const tirx::AndNode* op, Args... args) EXPR_FUNCTOR_DEFAULT;
+ virtual R VisitExpr_(const tirx::OrNode* op, Args... args) EXPR_FUNCTOR_DEFAULT;
+ virtual R VisitExpr_(const tirx::CastNode* op, Args... args) EXPR_FUNCTOR_DEFAULT;
+ virtual R VisitExpr_(const tirx::NotNode* op, Args... args) EXPR_FUNCTOR_DEFAULT;
+ virtual R VisitExpr_(const tirx::SelectNode* op, Args... args) EXPR_FUNCTOR_DEFAULT;
+ virtual R VisitExpr_(const tirx::RampNode* op, Args... args) EXPR_FUNCTOR_DEFAULT;
+ virtual R VisitExpr_(const tirx::BroadcastNode* op, Args... args) EXPR_FUNCTOR_DEFAULT;
+ virtual R VisitExpr_(const tirx::ShuffleNode* op, Args... args) EXPR_FUNCTOR_DEFAULT;
+ virtual R VisitExpr_(const tvm::IntImmNode* op, Args... args) EXPR_FUNCTOR_DEFAULT;
+ virtual R VisitExpr_(const tvm::FloatImmNode* op, Args... args) EXPR_FUNCTOR_DEFAULT;
+ virtual R VisitExpr_(const tirx::StringImmNode* op, Args... args) EXPR_FUNCTOR_DEFAULT;
virtual R VisitExpr_(const SeqExprNode* op, Args... args) EXPR_FUNCTOR_DEFAULT;
virtual R VisitExpr_(const IfNode* op, Args... args) EXPR_FUNCTOR_DEFAULT;
virtual R VisitExpr_(const OpNode* op, Args... args) EXPR_FUNCTOR_DEFAULT;
@@ -175,6 +211,36 @@ class ExprFunctor<R(const Expr& n, Args...)> {
RELAX_EXPR_FUNCTOR_DISPATCH(GlobalVarNode);
RELAX_EXPR_FUNCTOR_DISPATCH(FunctionNode);
RELAX_EXPR_FUNCTOR_DISPATCH(CallNode);
+ RELAX_EXPR_FUNCTOR_DISPATCH(tirx::ProducerLoadNode);
+ RELAX_EXPR_FUNCTOR_DISPATCH(tirx::LetNode);
+ RELAX_EXPR_FUNCTOR_DISPATCH(tirx::ReduceNode);
+ RELAX_EXPR_FUNCTOR_DISPATCH(tirx::BufferLoadNode);
+ RELAX_EXPR_FUNCTOR_DISPATCH(tirx::AddNode);
+ RELAX_EXPR_FUNCTOR_DISPATCH(tirx::SubNode);
+ RELAX_EXPR_FUNCTOR_DISPATCH(tirx::MulNode);
+ RELAX_EXPR_FUNCTOR_DISPATCH(tirx::DivNode);
+ RELAX_EXPR_FUNCTOR_DISPATCH(tirx::ModNode);
+ RELAX_EXPR_FUNCTOR_DISPATCH(tirx::FloorDivNode);
+ RELAX_EXPR_FUNCTOR_DISPATCH(tirx::FloorModNode);
+ RELAX_EXPR_FUNCTOR_DISPATCH(tirx::MinNode);
+ RELAX_EXPR_FUNCTOR_DISPATCH(tirx::MaxNode);
+ RELAX_EXPR_FUNCTOR_DISPATCH(tirx::EQNode);
+ RELAX_EXPR_FUNCTOR_DISPATCH(tirx::NENode);
+ RELAX_EXPR_FUNCTOR_DISPATCH(tirx::LTNode);
+ RELAX_EXPR_FUNCTOR_DISPATCH(tirx::LENode);
+ RELAX_EXPR_FUNCTOR_DISPATCH(tirx::GTNode);
+ RELAX_EXPR_FUNCTOR_DISPATCH(tirx::GENode);
+ RELAX_EXPR_FUNCTOR_DISPATCH(tirx::AndNode);
+ RELAX_EXPR_FUNCTOR_DISPATCH(tirx::OrNode);
+ RELAX_EXPR_FUNCTOR_DISPATCH(tirx::CastNode);
+ RELAX_EXPR_FUNCTOR_DISPATCH(tirx::NotNode);
+ RELAX_EXPR_FUNCTOR_DISPATCH(tirx::SelectNode);
+ RELAX_EXPR_FUNCTOR_DISPATCH(tirx::RampNode);
+ RELAX_EXPR_FUNCTOR_DISPATCH(tirx::BroadcastNode);
+ RELAX_EXPR_FUNCTOR_DISPATCH(tirx::ShuffleNode);
+ RELAX_EXPR_FUNCTOR_DISPATCH(tvm::IntImmNode);
+ RELAX_EXPR_FUNCTOR_DISPATCH(tvm::FloatImmNode);
+ RELAX_EXPR_FUNCTOR_DISPATCH(tirx::StringImmNode);
RELAX_EXPR_FUNCTOR_DISPATCH(SeqExprNode);
RELAX_EXPR_FUNCTOR_DISPATCH(IfNode);
RELAX_EXPR_FUNCTOR_DISPATCH(OpNode);
@@ -207,6 +273,33 @@ class ExprVisitor : public ExprFunctor<void(const Expr&)> {
void VisitExpr_(const GlobalVarNode* op) override;
void VisitExpr_(const FunctionNode* op) override;
void VisitExpr_(const CallNode* op) override;
+ void VisitExpr_(const tirx::BufferLoadNode* op) override;
+ void VisitExpr_(const tirx::AddNode* op) override;
+ void VisitExpr_(const tirx::SubNode* op) override;
+ void VisitExpr_(const tirx::MulNode* op) override;
+ void VisitExpr_(const tirx::DivNode* op) override;
+ void VisitExpr_(const tirx::ModNode* op) override;
+ void VisitExpr_(const tirx::FloorDivNode* op) override;
+ void VisitExpr_(const tirx::FloorModNode* op) override;
+ void VisitExpr_(const tirx::MinNode* op) override;
+ void VisitExpr_(const tirx::MaxNode* op) override;
+ void VisitExpr_(const tirx::EQNode* op) override;
+ void VisitExpr_(const tirx::NENode* op) override;
+ void VisitExpr_(const tirx::LTNode* op) override;
+ void VisitExpr_(const tirx::LENode* op) override;
+ void VisitExpr_(const tirx::GTNode* op) override;
+ void VisitExpr_(const tirx::GENode* op) override;
+ void VisitExpr_(const tirx::AndNode* op) override;
+ void VisitExpr_(const tirx::OrNode* op) override;
+ void VisitExpr_(const tirx::CastNode* op) override;
+ void VisitExpr_(const tirx::NotNode* op) override;
+ void VisitExpr_(const tirx::SelectNode* op) override;
+ void VisitExpr_(const tirx::RampNode* op) override;
+ void VisitExpr_(const tirx::BroadcastNode* op) override;
+ void VisitExpr_(const tirx::ShuffleNode* op) override;
+ void VisitExpr_(const tvm::IntImmNode* op) override;
+ void VisitExpr_(const tvm::FloatImmNode* op) override;
+ void VisitExpr_(const tirx::StringImmNode* op) override;
void VisitExpr_(const SeqExprNode* op) override;
void VisitExpr_(const IfNode* op) override;
void VisitExpr_(const OpNode* op) override;
@@ -262,7 +355,8 @@ class ExprVisitor : public ExprFunctor<void(const Expr&)> {
*
* By default, this function recurse into type such as
* TensorType and ShapeType and call VisitExpr/VisitTypePrimExprField
- * accordingly. It does not recurse into FunctionType as it does
+ * 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.
@@ -334,6 +428,33 @@ class ExprMutatorBase : public ExprFunctor<Expr(const Expr&)> {
Expr VisitExpr_(const GlobalVarNode* op) override;
Expr VisitExpr_(const FunctionNode* op) override;
Expr VisitExpr_(const CallNode* op) override;
+ Expr VisitExpr_(const tirx::BufferLoadNode* op) override;
+ Expr VisitExpr_(const tirx::AddNode* op) override;
+ Expr VisitExpr_(const tirx::SubNode* op) override;
+ Expr VisitExpr_(const tirx::MulNode* op) override;
+ Expr VisitExpr_(const tirx::DivNode* op) override;
+ Expr VisitExpr_(const tirx::ModNode* op) override;
+ Expr VisitExpr_(const tirx::FloorDivNode* op) override;
+ Expr VisitExpr_(const tirx::FloorModNode* op) override;
+ Expr VisitExpr_(const tirx::MinNode* op) override;
+ Expr VisitExpr_(const tirx::MaxNode* op) override;
+ Expr VisitExpr_(const tirx::EQNode* op) override;
+ Expr VisitExpr_(const tirx::NENode* op) override;
+ Expr VisitExpr_(const tirx::LTNode* op) override;
+ Expr VisitExpr_(const tirx::LENode* op) override;
+ Expr VisitExpr_(const tirx::GTNode* op) override;
+ Expr VisitExpr_(const tirx::GENode* op) override;
+ Expr VisitExpr_(const tirx::AndNode* op) override;
+ Expr VisitExpr_(const tirx::OrNode* op) override;
+ Expr VisitExpr_(const tirx::CastNode* op) override;
+ Expr VisitExpr_(const tirx::NotNode* op) override;
+ Expr VisitExpr_(const tirx::SelectNode* op) override;
+ Expr VisitExpr_(const tirx::RampNode* op) override;
+ Expr VisitExpr_(const tirx::BroadcastNode* op) override;
+ Expr VisitExpr_(const tirx::ShuffleNode* op) override;
+ Expr VisitExpr_(const tvm::IntImmNode* op) override;
+ Expr VisitExpr_(const tvm::FloatImmNode* op) override;
+ Expr VisitExpr_(const tirx::StringImmNode* op) override;
Expr VisitExpr_(const SeqExprNode* op) override;
Expr VisitExpr_(const IfNode* op) override;
Expr VisitExpr_(const OpNode* op) override;
@@ -350,9 +471,10 @@ class ExprMutatorBase : public ExprFunctor<Expr(const Expr&)> {
virtual BindingBlock VisitBindingBlock(const BindingBlock& block);
/*!
- * \brief Used to visit the PrimExpr inside of expressions.
+ * \brief Used to visit the PrimExpr inside of dependent type fields.
*
- * Can be overloaded to transform the shape expressions.
+ * Can be overloaded to transform shape expressions without changing
+ * runtime PrimExpr traversal.
*/
virtual PrimExpr VisitTypePrimExprField(const PrimExpr& expr);
@@ -361,7 +483,8 @@ class ExprMutatorBase : public ExprFunctor<Expr(const Expr&)> {
*
* By default, this function recurse into type such as
* TensorType and ShapeType and call VisitExpr/VisitTypePrimExprField
- * accordingly. It does not recurse into FunctionType as it does
+ * 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.
diff --git a/include/tvm/relax/script/builder/frame.h b/include/tvm/relax/script/builder/frame.h
index 7a20ceed1a..a2f663fb50 100644
--- a/include/tvm/relax/script/builder/frame.h
+++ b/include/tvm/relax/script/builder/frame.h
@@ -98,7 +98,7 @@ class FunctionFrameNode : public SeqExprFrameNode {
*/
ffi::Optional<ffi::String> name;
/*! \brief The function params. */
- ffi::Array<tvm::relax::Var> params;
+ ffi::Array<tvm::Var> params;
/*!
* \brief The function return type.
* \note Usually the function return type can be deduced by the function body.
@@ -151,7 +151,7 @@ class BindingBlockFrameNode : public RelaxFrameNode {
/*! \brief The flag that indicates whether the block is a dataflow block. */
bool is_dataflow;
/*! \brief The variables emitted in this block. */
- ffi::Array<tvm::relax::Var> emitted_vars;
+ ffi::Array<tvm::Var> emitted_vars;
/*!
* \brief A boolean indicating if the dataflow block is ended of construction.
* If it is true, any new binding trying to be emitted into this block will cause an error.
@@ -162,7 +162,7 @@ class BindingBlockFrameNode : public RelaxFrameNode {
* \brief The output vars of the dataflow block.
* \note Only used for a dataflow block.
*/
- ffi::Array<tvm::relax::Var> output_vars;
+ ffi::Array<tvm::Var> output_vars;
static void RegisterReflection() {
namespace refl = tvm::ffi::reflection;
@@ -203,7 +203,7 @@ class IfFrameNode : public RelaxFrameNode {
/*! \brief The Bindings in the false branch. */
ffi::Optional<tvm::relax::Expr> else_expr;
/*! \brief The Binding var. */
- tvm::relax::Var var;
+ tvm::Var var;
/*! \brief The binding var name. */
ffi::String var_name;
diff --git a/include/tvm/relax/script/builder/ir.h b/include/tvm/relax/script/builder/ir.h
index 6a5df22245..2516dc0134 100644
--- a/include/tvm/relax/script/builder/ir.h
+++ b/include/tvm/relax/script/builder/ir.h
@@ -45,7 +45,7 @@ TVM_DLL FunctionFrame Function(bool is_pure, bool is_private);
* \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::Type& ty);
+TVM_DLL tvm::Var Arg(const ffi::String& name, const tvm::Type& ty);
/*!
* \brief Specify the name of the last function frame.
@@ -89,7 +89,7 @@ TVM_DLL BindingBlockFrame Dataflow();
* \brief Expose the dataflow block output variables as global ones
* \param vars The output variables of a dataflow block
*/
-TVM_DLL void DataflowBlockOutput(const ffi::Array<tvm::relax::Var>& vars);
+TVM_DLL void DataflowBlockOutput(const ffi::Array<tvm::Var>& vars);
////////////////////////////// Bindings ////////////////////////////////
@@ -99,8 +99,8 @@ TVM_DLL void DataflowBlockOutput(const ffi::Array<tvm::relax::Var>& vars);
* \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::Type>& annotate_ty = std::nullopt);
+TVM_DLL tvm::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.
@@ -108,14 +108,14 @@ TVM_DLL tvm::relax::Var Emit(const tvm::relax::Expr& value,
* \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::Type& ty);
+TVM_DLL tvm::Var EmitMatchCast(const tvm::relax::Expr& value, const tvm::Type& ty);
/*!
* \brief Emit a binding to the last binding block frame.
* \param binding The binding to be emitted.
* \return The left side var of the emitted binding.
*/
-TVM_DLL tvm::relax::Var EmitVarBinding(const tvm::relax::VarBinding& binding);
+TVM_DLL tvm::Var EmitVarBinding(const tvm::relax::VarBinding& binding);
///////////////////////////// If Then Else /////////////////////////////
diff --git a/include/tvm/relax/transform.h b/include/tvm/relax/transform.h
index 5c757ba151..94382cf602 100644
--- a/include/tvm/relax/transform.h
+++ b/include/tvm/relax/transform.h
@@ -215,8 +215,9 @@ TVM_DLL Pass BindParams(ffi::String func_name, ffi::Map<Any, ffi::ObjectRef> par
*
* \return The Pass.
*/
-TVM_DLL Pass BindSymbolicVars(ffi::Map<ffi::Variant<tirx::Var, ffi::String>, PrimExpr> binding_map,
- ffi::Optional<ffi::String> func_name = std::nullopt);
+TVM_DLL Pass
+BindSymbolicVars(ffi::Map<ffi::Variant<tirx::PrimVar, ffi::String>, PrimExpr> binding_map,
+ ffi::Optional<ffi::String> func_name = std::nullopt);
/*!
* \brief Fold constant expressions within dataflow blocks.
diff --git a/include/tvm/relax/type.h b/include/tvm/relax/type.h
index 9858cfea31..246487f9d4 100644
--- a/include/tvm/relax/type.h
+++ b/include/tvm/relax/type.h
@@ -322,7 +322,7 @@ class FuncType : public 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
+ * \note If the ret contains canonical tvm::Var values, 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());
diff --git a/include/tvm/relax/utils.h b/include/tvm/relax/utils.h
index f36da70f9c..d081d99166 100644
--- a/include/tvm/relax/utils.h
+++ b/include/tvm/relax/utils.h
@@ -43,18 +43,15 @@ namespace relax {
* \param expr The input expression.
* \param binds The variable to expression map that will be used to help the
* binding.
- * \param symbolic_var_map The map from symbolic var to the expr it binds to.
- *
* \return The updated expression.
*/
-TVM_DLL Expr Bind(const Expr& expr, const tvm::ffi::Map<Var, Expr>& binds,
- const tvm::ffi::Map<tirx::Var, PrimExpr>& symbolic_var_map = {});
+TVM_DLL Expr Bind(const Expr& expr, const tvm::ffi::Map<Var, Expr>& binds);
/*!
* \brief Bind the symbolic variables to a Type. This is a helper function usually called by
* other pass functions to help optimizations.
*/
-TVM_DLL Type Bind(const Type& ty, const tvm::ffi::Map<tirx::Var, PrimExpr>& symbolic_var_map);
+TVM_DLL Type Bind(const Type& ty, const tvm::ffi::Map<Var, Expr>& binds);
/*!
* \brief Infer a binding map for symbolic variables
@@ -71,10 +68,10 @@ TVM_DLL Type Bind(const Type& ty, const tvm::ffi::Map<tirx::Var, PrimExpr>& symb
*
* \param analyzer The analyzer to use for simplifications
*
- * \return A map of TIR variables to TIR expressions
+ * \return The input binding map augmented with inferred symbolic bindings.
*/
-TVM_DLL tvm::ffi::Map<tirx::Var, PrimExpr> InferSymbolicVarMap(
- const tvm::ffi::Map<relax::Var, relax::Expr>& binds, const arith::Analyzer& analyzer);
+TVM_DLL tvm::ffi::Map<Var, Expr> InferSymbolicVarMap(
+ const tvm::ffi::Map<tvm::Var, relax::Expr>& binds, const arith::Analyzer& analyzer);
/*!
* \brief Check if the given Type is for a boolean scalar (tensor of rank 0 with a boolean
@@ -112,7 +109,8 @@ 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 FuncType. If it is an Op,
- * then this checks the `fPurity` field.
+ * then this checks the Relax `FPurity` attribute, falling back to the TIR
+ * `TCallEffectKind` attribute for primitive operators.
*
* \param call The input call
*
diff --git a/include/tvm/tirx/stmt_functor.h b/include/tvm/tirx/stmt_functor.h
index 01563a9c15..d4cf5a7f9c 100644
--- a/include/tvm/tirx/stmt_functor.h
+++ b/include/tvm/tirx/stmt_functor.h
@@ -406,6 +406,17 @@ inline PrimExpr Substitute(PrimExpr expr, std::function<ffi::Optional<Expr>(cons
return Substitute(Expr(expr), std::move(vmap)).as_or_throw<PrimExpr>();
}
+/*!
+ * \brief Substitute the vars specified by vmap.
+ * \param range The array of Stmt/PrimExpr to be substituted
+ * \param vmap returns a new value if re-mapping is needed, otherwise returns nullptr.
+ * \return The modified Range.
+ */
+inline Range Substitute(const Range& range,
+ std::function<ffi::Optional<Expr>(const Var& var)> vmap) {
+ return Range::FromMinExtent(Substitute(range->min, vmap), Substitute(range->extent, vmap));
+}
+
/*!
* \brief Substitute the var specified by vmap.
* \param arr The array of Stmt/PrimExpr to be substituted
@@ -418,17 +429,6 @@ ffi::Array<T> Substitute(const ffi::Array<T>& arr,
return arr.Map([&vmap](const auto& elem) { return Substitute(elem, vmap); });
}
-/*!
- * \brief Substitute the vars specified by vmap.
- * \param range The array of Stmt/PrimExpr to be substituted
- * \param vmap returns a new value if re-mapping is needed, otherwise returns nullptr.
- * \return The modified Range.
- */
-inline Range Substitute(const Range& range,
- std::function<ffi::Optional<Expr>(const Var& var)> vmap) {
- return Range::FromMinExtent(Substitute(range->min, vmap), Substitute(range->extent, vmap));
-}
-
/*!
* \brief Substitute the vars specified by vmap.
*
diff --git a/include/tvm/tirx/var.h b/include/tvm/tirx/var.h
index 5f957b5d53..e295358b83 100644
--- a/include/tvm/tirx/var.h
+++ b/include/tvm/tirx/var.h
@@ -29,94 +29,13 @@
#include <tvm/ir/expr.h>
#include <tvm/ir/type.h>
-#include <functional>
-#include <string>
+#include <utility>
namespace tvm {
namespace tirx {
-/*!
- * \brief A variable node in the IR.
- *
- * A variable is uniquely identified by its address.
- *
- * Each variable is only bound once in the following nodes:
- * - Allocate
- * - For
- * - Let
- * - Bind
- */
-class VarNode : public ExprNode {
- public:
- /*!
- * \brief The hint to the variable name.
- * \note Each variable is uniquely identified by its address.
- */
- ffi::String name_hint;
- static void RegisterReflection() {
- namespace refl = tvm::ffi::reflection;
- refl::ObjectDef<VarNode>().def_ro("name", &VarNode::name_hint,
- refl::AttachFieldFlag::SEqHashIgnore());
- }
-
- static constexpr TVMFFISEqHashKind _type_s_eq_hash_kind = kTVMFFISEqHashKindFreeVar;
- static constexpr const uint32_t _type_child_slots = 1;
- TVM_FFI_DECLARE_OBJECT_INFO("tirx.Var", VarNode, ExprNode);
-};
-
-/*! \brief a named variable in TIR */
-class Var : public Expr {
- public:
- explicit Var(ffi::UnsafeInit tag) : Expr(tag) {}
- explicit Var(ffi::ObjectPtr<VarNode> n) : Expr(n) {}
- /*!
- * \brief Constructor
- * \param name_hint variable name
- * \param dtype data type
- * \param span The location of this object in the source code.
- */
- TVM_DLL explicit Var(ffi::String name_hint = "v", PrimType dtype = PrimType::Int(32),
- Span span = Span());
- /*!
- * \brief Constructor which provides a more detailed type annotation.
- * \param name_hint variable name.
- * \param type_annotation The type annotation.
- * \param span The location of this object in the source code.
- */
- TVM_DLL explicit Var(ffi::String name_hint, Type type_annotation, Span span = Span());
- /*!
- * \brief Make a new copy of var with same type, but a different nam
- * \param name The new name to be used.
- * \return the new Var copy
- */
- TVM_DLL Var copy_with_name(const ffi::String& name) const;
- /*!
- * \brief Make a new copy of var with same type, append suffix
- * \param suffix The suffix to be appended.
- * \return the new Var copy
- */
- TVM_DLL Var CopyWithSuffix(const ffi::String& suffix) const;
- /*!
- * \brief Make a new copy of the variable with specified dtype
- * \param dtype The specified dtype
- * \return The new variable
- */
- TVM_DLL Var copy_with_dtype(PrimType dtype) const;
-
- /*!
- * \brief Get pointer to the internal value.
- * \return the corresponding Variable.
- */
- const VarNode* operator->() const { return get(); }
- /*!
- * \brief Get pointer to the internal value.
- * \return the corresponding Variable.
- */
- const VarNode* get() const { return static_cast<const VarNode*>(data_.get()); }
- /*! \brief type indicate the container type */
- using ContainerType = VarNode;
- static constexpr bool _type_container_is_exact = true;
-};
+using VarNode = tvm::VarNode;
+using Var = tvm::Var;
/*!
* \brief Checked scalar view over a VarNode.
@@ -147,7 +66,7 @@ class PrimVar : public PrimExpr {
}
TVM_FFI_DEFINE_OBJECT_REF_METHODS_NULLABLE(PrimVar, PrimExpr, VarNode);
- static constexpr bool _type_container_is_exact = true;
+ static constexpr bool _type_container_is_exact = false;
};
using Region = ffi::Array<Range>;
@@ -334,8 +253,7 @@ struct TypeTraits<tirx::PrimVar> : public ObjectRefTypeTraitsBase<tirx::PrimVar>
if (src->type_index == TypeIndex::kTVMFFINone) {
return tirx::PrimVar::_type_is_nullable;
}
- if (src->type_index < TypeIndex::kTVMFFIStaticObjectBegin ||
- !details::IsObjectInstance<tirx::VarNode>(src->type_index)) {
+ if (src->type_index != tirx::VarNode::RuntimeTypeIndex()) {
return false;
}
const auto* var = static_cast<const tirx::VarNode*>(
@@ -357,30 +275,4 @@ struct TypeTraits<tirx::PrimVar> : public ObjectRefTypeTraitsBase<tirx::PrimVar>
} // namespace tvm::ffi
-/* \brief Allow tirx.Var as key in STL tables
- *
- * For most TIR expressions, it would be ambiguous whether the
- * expression should follow reference equality or structural equality.
- * This is not the case for variables, which do not contain nested
- * internal structure, and are frequently used as keys in lookup
- * tables.
- *
- * Providing `std::hash` and `std::equal_to` specializations for
- * `tirx::Var` allows it to be used as a key in STL tables. For
- * `PrimExpr`, the user must specify the type of equality used
- * (e.g. `std::unordered_set<T, StructuralHash, StructuralEqual>` or
- * `std::unordered_set<T, ffi::ObjectPtrHash, ffi::ObjectPtrEqual>`).
- */
-template <>
-struct std::hash<tvm::tirx::Var> {
- std::size_t operator()(const tvm::tirx::Var& var) const { return tvm::ffi::ObjectPtrHash()(var); }
-};
-
-template <>
-struct std::equal_to<tvm::tirx::Var> {
- bool operator()(const tvm::tirx::Var& var_a, const tvm::tirx::Var& var_b) const {
- return tvm::ffi::ObjectPtrEqual()(var_a, var_b);
- }
-};
-
#endif // TVM_TIR_VAR_H_
diff --git a/include/tvm/topi/transform.h b/include/tvm/topi/transform.h
index e9f9249425..5d56a0ca64 100644
--- a/include/tvm/topi/transform.h
+++ b/include/tvm/topi/transform.h
@@ -650,10 +650,10 @@ inline ffi::Array<Tensor> split_indices_array(const Tensor& x, ffi::Array<PrimEx
}
inline PrimExpr DynamicCanonicalizeIndex(PrimExpr index, PrimExpr extent, PrimExpr stride) {
- auto idx_var = index.as<tvm::tirx::VarNode>();
- auto extent_var = extent.as<tvm::tirx::VarNode>();
+ auto idx_var = index.as<tvm::tirx::PrimVar>();
+ auto extent_var = extent.as<tvm::tirx::PrimVar>();
- if (idx_var && extent_var && idx_var->name_hint == extent_var->name_hint) {
+ if (idx_var && extent_var && (*idx_var)->name_hint == (*extent_var)->name_hint) {
return index;
}
diff --git a/python/tvm/backend/cuda/operator/tile_primitive/copy/fallback.py b/python/tvm/backend/cuda/operator/tile_primitive/copy/fallback.py
index ab69d39244..622a8e7a15 100644
--- a/python/tvm/backend/cuda/operator/tile_primitive/copy/fallback.py
+++ b/python/tvm/backend/cuda/operator/tile_primitive/copy/fallback.py
@@ -60,7 +60,7 @@ def _emit_fallback(op_call: TilePrimitiveCall, sctx: DispatchContext) -> PrimFun
copy_extents = [dst_extent[i] for i in dst_indices]
def _dst_coord(lvs):
- if isinstance(lvs, tvm.tirx.Var):
+ if isinstance(lvs, tvm.ir.Var):
lvs = [lvs]
coord = list(dst_st)
for k, lv in enumerate(lvs):
@@ -68,7 +68,7 @@ def _emit_fallback(op_call: TilePrimitiveCall, sctx: DispatchContext) -> PrimFun
return coord
def _src_coord(lvs):
- if isinstance(lvs, tvm.tirx.Var):
+ if isinstance(lvs, tvm.ir.Var):
lvs = [lvs]
coord = list(src_st)
for k, lv in enumerate(lvs):
diff --git a/python/tvm/backend/cuda/operator/tile_primitive/copy_async/tma.py b/python/tvm/backend/cuda/operator/tile_primitive/copy_async/tma.py
index 9d79eb2c49..8f78679de2 100644
--- a/python/tvm/backend/cuda/operator/tile_primitive/copy_async/tma.py
+++ b/python/tvm/backend/cuda/operator/tile_primitive/copy_async/tma.py
@@ -284,7 +284,7 @@ def _simplify_with_var_ranges(exprs, var_ranges, sctx: DispatchContext):
for var, value_range in sctx.var_range_map.items():
local_analyzer.bind(var, value_range)
for var, extent in var_ranges:
- if isinstance(var, tvm.tirx.Var):
+ if tvm.ir.is_prim_var(var):
local_analyzer.bind(var, tvm.ir.Range.from_min_extent(0, extent))
return [local_analyzer.simplify(expr) for expr in exprs]
@@ -1179,7 +1179,7 @@ def copy_tma_impl(op_call: TilePrimitiveCall, sctx: DispatchContext) -> PrimFunc
tensor_map = cached_tensormap
tensormap_is_cached = True
else:
- tensor_map = T.Var(g_buf.data.name + "_tensormap", dtype=T.handle("tensormap").ty)
+ tensor_map = T.Var(g_buf.data.name + "_tensormap", ty=T.handle("tensormap").ty)
tensormap_is_cached = False
# fmt: off
diff --git a/python/tvm/ir/__init__.py b/python/tvm/ir/__init__.py
index 010e946a27..ff05d229f5 100644
--- a/python/tvm/ir/__init__.py
+++ b/python/tvm/ir/__init__.py
@@ -30,7 +30,7 @@ from .base import (
load_json,
save_json,
)
-from .expr import Call, Expr, GlobalVar, Range, is_prim_expr
+from .expr import Call, Expr, GlobalVar, Range, Var, is_prim_expr, is_prim_var
from .function import BaseFunc, CallingConv
from .global_info import GlobalInfo, DummyGlobalInfo, VDevice
from .module import IRModule
diff --git a/python/tvm/ir/expr.py b/python/tvm/ir/expr.py
index 05071696ad..c1f146da87 100644
--- a/python/tvm/ir/expr.py
+++ b/python/tvm/ir/expr.py
@@ -40,6 +40,11 @@ def is_prim_expr(value: object) -> bool:
return isinstance(value, Expr) and isinstance(value.ty, tvm.ir.PrimType)
+def is_prim_var(value: object) -> bool:
+ """Return whether a value is an ordinary variable with a primitive type."""
+ return isinstance(value, Var) and type(value) is Var and is_prim_expr(value)
+
+
@tvm_ffi.register_object("ir.GlobalVar")
class GlobalVar(Expr):
"""A global variable in the IR.
@@ -90,51 +95,16 @@ class GlobalVar(Expr):
raise RuntimeError(f"Do not know how to handle GlobalVar.__call__ for types {arg_types}")
-@tvm_ffi.register_object("ir.Call")
-class Call(Expr, Scriptable):
- """Core function call node."""
+class _ExprWithOp(Expr, Scriptable):
+ """Common type-directed operator behavior for core expressions."""
__hash__ = Expr.__hash__
- op: Expr
- args: list[Expr]
- attrs: "tvm.ir.Attrs | None"
- ty_args: list["tvm.ir.Type"]
- span: Span | None
-
- def __init__(
- self,
- op: Expr | str,
- args: list[Expr] | tuple[Expr, ...],
- attrs: "tvm.ir.Attrs | dict | None" = None,
- ty_args: list["tvm.ir.Type"] | tuple["tvm.ir.Type", ...] | None = None,
- span: Span | None = None,
- ret_ty: "tvm.ir.Type | str | None" = None,
- ) -> None:
- # pylint: disable=import-outside-toplevel
- from .attrs import DictAttrs
- from .op import Op
- from .type import PointerType, PrimType, Type
-
- if isinstance(op, str):
- op = Op.get(op)
- if attrs is not None and isinstance(attrs, dict):
- attrs = DictAttrs(attrs)
- if ret_ty is None:
- ret_ty = Type.missing()
- if isinstance(ret_ty, str) and ret_ty == "handle":
- ret_ty = PointerType(PrimType("void"))
- elif ret_ty is not None and not isinstance(ret_ty, Type):
- ret_ty = PrimType(ret_ty)
- if ty_args is None:
- ty_args = []
- self.__init_handle_by_constructor__(_ffi_api.Call, ret_ty, op, args, attrs, ty_args, span)
-
def expr_ty(self):
"""Return this expression's primitive result type."""
if is_prim_expr(self):
return self.ty
- raise TypeError(f"Expected primitive-valued Call, but result type is {self.ty}")
+ raise TypeError(f"Expected a primitive-valued expression, but result type is {self.ty}")
def __add__(self, other):
if is_prim_expr(self):
@@ -343,7 +313,7 @@ class Call(Expr, Scriptable):
def __call__(self, *args, attrs=None):
if is_prim_expr(self):
- raise TypeError("A primitive-valued Call cannot be called")
+ raise TypeError("A primitive-valued expression cannot be called")
result = _tensor_expr_overload.__call__(self, *args, attrs=attrs)
if result is NotImplemented:
raise TypeError("Tensor expression overload __call__ is not registered")
@@ -351,13 +321,104 @@ class Call(Expr, Scriptable):
def __getitem__(self, index):
if is_prim_expr(self):
- raise TypeError("A primitive-valued Call cannot be indexed")
+ raise TypeError("A primitive-valued expression cannot be indexed")
result = _tensor_expr_overload.__getitem__(self, index)
if result is NotImplemented:
raise TypeError("Tensor expression overload __getitem__ is not registered")
return result
+@tvm_ffi.register_object("ir.Call")
+class Call(_ExprWithOp):
+ """Core function call node."""
+
+ op: Expr
+ args: list[Expr]
+ attrs: "tvm.ir.Attrs | None"
+ ty_args: list["tvm.ir.Type"]
+ span: Span | None
+
+ def __init__(
+ self,
+ op: Expr | str,
+ args: list[Expr] | tuple[Expr, ...],
+ attrs: "tvm.ir.Attrs | dict | None" = None,
+ ty_args: list["tvm.ir.Type"] | tuple["tvm.ir.Type", ...] | None = None,
+ span: Span | None = None,
+ ret_ty: "tvm.ir.Type | str | None" = None,
+ ) -> None:
+ # pylint: disable=import-outside-toplevel
+ from .attrs import DictAttrs
+ from .op import Op
+ from .type import PointerType, PrimType, Type
+
+ if isinstance(op, str):
+ op = Op.get(op)
+ if attrs is not None and isinstance(attrs, dict):
+ attrs = DictAttrs(attrs)
+ if ret_ty is None:
+ ret_ty = Type.missing()
+ if isinstance(ret_ty, str) and ret_ty == "handle":
+ ret_ty = PointerType(PrimType("void"))
+ elif ret_ty is not None and not isinstance(ret_ty, Type):
+ ret_ty = PrimType(ret_ty)
+ if ty_args is None:
+ ty_args = []
+ self.__init_handle_by_constructor__(_ffi_api.Call, ret_ty, op, args, attrs, ty_args, span)
+
+
+@tvm_ffi.register_object("ir.Var")
+class Var(_ExprWithOp):
+ """A canonical local variable in the IR.
+
+ Parameters
+ ----------
+ name : str
+ The API name of the variable. It is stored internally as a name hint.
+
+ ty : Optional[Type or str]
+ The exact type of the variable. A string denotes a primitive dtype.
+
+ span : Optional[Span]
+ Span that points to the original source code.
+
+ """
+
+ name_hint: str
+ span: Span | None
+
+ def __init__(
+ self,
+ name: str | None = None,
+ ty: "tvm.ir.Type | str | None" = None,
+ span: Span | None = None,
+ *,
+ name_hint: str | None = None,
+ ) -> None:
+ if name is None:
+ name = name_hint
+ elif name_hint is not None:
+ raise TypeError("Specify either name or name_hint, not both")
+ if not isinstance(name, str):
+ raise TypeError("name must be a str")
+
+ # pylint: disable=import-outside-toplevel
+ from .type import PointerType, PrimType, Type
+
+ if isinstance(ty, str):
+ ty = PointerType(PrimType("void")) if ty == "handle" else PrimType(ty)
+ elif ty is not None:
+ ty = tvm.runtime.convert(ty)
+ if not isinstance(ty, Type):
+ raise TypeError("ty must be a Type or primitive dtype string")
+ self.__init_handle_by_constructor__(_ffi_api.Var, name, ty, span)
+
+ @property
+ def name(self) -> str:
+ """The user-facing variable name."""
+ return self.name_hint
+
+
@tvm_ffi.register_object("ir.Range")
class Range(Node, Scriptable):
"""Represent a range in TVM.
diff --git a/python/tvm/ir/json_compact.py b/python/tvm/ir/json_compact.py
index 848de6b54d..948c064e64 100644
--- a/python/tvm/ir/json_compact.py
+++ b/python/tvm/ir/json_compact.py
@@ -84,4 +84,16 @@ def upgrade_json(json_str):
data = json.loads(json_str)
if "metadata" not in data and "attrs" in data:
raise ValueError("Legacy json graph format detected, we don't support it anymore.")
+
+ # `ir.Var` is the sole runtime variable node. Keep `tvm.ir.load_json`
+ # compatible with the exact pre-unification Relax and TIRx schemas.
+ # Rewriting nodes in place preserves node indices and shared references.
+ for node in data.get("nodes", []):
+ if node.get("type") == "relax.expr.Var":
+ node["type"] = "ir.Var"
+ elif node.get("type") == "tirx.Var":
+ node["type"] = "ir.Var"
+ fields = node.get("data", {})
+ if "name" in fields and "name_hint" not in fields:
+ fields["name_hint"] = fields.pop("name")
return json.dumps(data, indent=2)
diff --git a/python/tvm/relax/backend/cuda/cublas.py b/python/tvm/relax/backend/cuda/cublas.py
index eef6c70ad9..54bc7679c6 100644
--- a/python/tvm/relax/backend/cuda/cublas.py
+++ b/python/tvm/relax/backend/cuda/cublas.py
@@ -133,8 +133,8 @@ def _check_matmul(context: PatternCheckContext) -> bool:
# must be equal. If lhs is batched but rhs is not, we can use the regular GEMM by
# flattening all batch axes into the M axis.
return (
- isinstance(lhs_batches, tvm.tirx.Var)
- or isinstance(rhs_batches, tvm.tirx.Var)
+ isinstance(lhs_batches, tvm.ir.Var)
+ or isinstance(rhs_batches, tvm.ir.Var)
or (analyzer.can_prove_equal(lhs_batches, rhs_batches))
or (analyzer.can_prove(lhs_batches >= 1) and analyzer.can_prove(rhs_batches == 1))
)
diff --git a/python/tvm/relax/backend/cuda/cutlass.py b/python/tvm/relax/backend/cuda/cutlass.py
index 3e6e936d8c..2ce5a3de60 100644
--- a/python/tvm/relax/backend/cuda/cutlass.py
+++ b/python/tvm/relax/backend/cuda/cutlass.py
@@ -136,7 +136,7 @@ def _check_conv2d(context: PatternCheckContext) -> bool:
# Check if any dimensions are symbolic.
for dim in data.ty.shape.values:
- if isinstance(dim, tvm.tirx.Var):
+ if isinstance(dim, tvm.ir.Var):
return False
# pylint: disable=invalid-name
diff --git a/python/tvm/relax/backend/metal/coreml.py b/python/tvm/relax/backend/metal/coreml.py
index c9ee23a004..68266574a2 100644
--- a/python/tvm/relax/backend/metal/coreml.py
+++ b/python/tvm/relax/backend/metal/coreml.py
@@ -353,7 +353,7 @@ class CodegenCoreML(PyExprVisitor):
def visit_function_(self, op) -> None:
for var in op.params:
- name = var.name_hint
+ name = var.name
ty = var.ty
if isinstance(ty, TensorType):
shape = [int(v) for v in list(ty.shape)]
@@ -367,7 +367,7 @@ class CodegenCoreML(PyExprVisitor):
self.visit_expr(op.body)
def visit_var_(self, var):
- self.out_map[var] = [var.name_hint]
+ self.out_map[var] = [var.name]
prev_binding_var = self.cur_binding_var
self.cur_binding_var = var
if var in self.var2val:
diff --git a/python/tvm/relax/backend/rocm/hipblas.py b/python/tvm/relax/backend/rocm/hipblas.py
index 1a52102b48..75c47321ee 100644
--- a/python/tvm/relax/backend/rocm/hipblas.py
+++ b/python/tvm/relax/backend/rocm/hipblas.py
@@ -84,8 +84,8 @@ def _check_matmul(context: PatternCheckContext) -> bool:
# must be equal. If lhs is batched but rhs is not, we can use the regular GEMM by
# flattening all batch axes into the M axis.
return (
- isinstance(lhs_batches, tvm.tirx.Var)
- or isinstance(rhs_batches, tvm.tirx.Var)
+ isinstance(lhs_batches, tvm.ir.Var)
+ or isinstance(rhs_batches, tvm.ir.Var)
or (int(lhs_batches) == int(rhs_batches))
or (lhs_batches >= 1 and rhs_batches == 1)
)
diff --git a/python/tvm/relax/block_builder.py b/python/tvm/relax/block_builder.py
index ef7f6d58f8..3a2cd4178d 100644
--- a/python/tvm/relax/block_builder.py
+++ b/python/tvm/relax/block_builder.py
@@ -87,10 +87,10 @@ class TestingScope:
self._bb = block_builder
shape_vars = []
for var in def_vars:
- if isinstance(var, tvm.tirx.Var):
+ if isinstance(var, tvm.ir.Var):
shape_vars.append(var)
else:
- raise ValueError("def_vars only can take tirx.Var")
+ raise ValueError("def_vars can only contain Vars")
# setup a dummy var so shape is in scope.
sparam = rx.Var("sparam", rx.ShapeType(shape_vars))
self._scope_params = [sparam]
diff --git a/python/tvm/relax/expr.py b/python/tvm/relax/expr.py
index c92a08e868..41b49dd2c5 100644
--- a/python/tvm/relax/expr.py
+++ b/python/tvm/relax/expr.py
@@ -410,45 +410,8 @@ class Constant(ExprWithOp):
)
-@tvm_ffi.register_object("relax.expr.Var")
-class Var(ExprWithOp):
- """The variable class for all Relax bindings.
-
- Parameters
- ----------
- name_hint: str
- The name hint of the variable.
-
- ty: Optional[Type]
- The type annotation of the variable.
-
- span: Optional[Span]
- Span that points to original source code
- """
-
- name_hint: str
- span: Span | None
-
- def __init__(
- self,
- name_hint: str,
- ty: Type | None = None,
- span: Span | None = None,
- ) -> None:
- if ty is not None:
- ty = tvm.runtime.convert(ty)
- if not isinstance(ty, Type):
- raise TypeError(
- "ty needs to be an instance of Type. "
- "If you attempt to pass in shape, "
- "use relax.TensorType(shape, dtype)."
- )
- self.__init_handle_by_constructor__(
- _ffi_api.Var, # type: ignore
- name_hint,
- ty,
- span,
- )
+# Ordinary Relax bindings use the canonical IR Var directly.
+Var = tvm.ir.Var
@tvm_ffi.register_object("relax.expr.DataflowVar")
@@ -459,8 +422,8 @@ class DataflowVar(Var):
Parameters
----------
- name_hint: str
- The name hint of the variable.
+ name: str
+ The API name of the variable. It is stored internally as a name hint.
ty: Optional[Type]
The type annotation of the variable.
@@ -469,16 +432,25 @@ class DataflowVar(Var):
Span that points to original source code
"""
- name_hint: str
+ name: str
span: Span | None
def __init__(
self,
- name_hint: str,
+ name: str | None = None,
ty: Type | None = None,
span: Span | None = None,
+ *,
+ name_hint: str | None = None,
) -> None:
# pylint: disable=super-init-not-called
+ if name is None:
+ name = name_hint
+ elif name_hint is not None:
+ raise TypeError("Specify either name or name_hint, not both")
+ if not isinstance(name, str):
+ raise TypeError("name must be a str")
+
if ty is not None:
ty = tvm.runtime.convert(ty)
if not isinstance(ty, Type):
@@ -488,7 +460,7 @@ class DataflowVar(Var):
"use relax.TensorType(shape, dtype)."
)
- self.__init_handle_by_constructor__(_ffi_api.DataflowVar, name_hint, ty, span) # type: ignore
+ self.__init_handle_by_constructor__(_ffi_api.DataflowVar, name, ty, span) # type: ignore
@tvm_ffi.register_object("relax.expr.StringImm")
@@ -836,19 +808,21 @@ class TEPlaceholderOp(tvm.te.tensor.Operation):
def te_tensor(
- value: Expr, tir_var_map: dict[tvm.tirx.Var, tvm.tirx.Expr], name: str = "rxplaceholder"
+ value: Expr,
+ tir_var_map: dict[tvm.ir.Var, tvm.tirx.Expr],
+ name: str = "rxplaceholder",
):
- """Create a TE tensor from relax expression, with TIR variables in the
- tensor shape substituted by the given mapping
+ """Create a TE tensor from a Relax expression, with primitive variables
+ in the tensor shape substituted by the given mapping.
Parameters
----------
value : Expr
The relax expression, which is required to have TensorType.
- tir_var_map : Dict[tvm.tirx.Var, tvm.tirx.Expr]
- The mapping to substitute the TIR variables appeared in the
- shape of the input Expr.
+ tir_var_map : Dict[tvm.ir.Var, tvm.tirx.Expr]
+ The exact-identity mapping used to refresh canonical primitive Vars
+ that appear in the input shape.
name : str
The name of the created tensor.
diff --git a/python/tvm/relax/expr_functor.py b/python/tvm/relax/expr_functor.py
index 39da280eee..688fb6fd6e 100644
--- a/python/tvm/relax/expr_functor.py
+++ b/python/tvm/relax/expr_functor.py
@@ -22,7 +22,8 @@ from collections.abc import Callable
import tvm_ffi
-from tvm.ir import Call, Op
+from tvm import tirx as _tirx
+from tvm.ir import Call, Op, is_prim_expr
from tvm.ir.utils import derived_object
from tvm.runtime import Object
@@ -151,6 +152,60 @@ class ExprFunctor:
ret = self.visit_function_(expr)
elif isinstance(expr, Call): # type: ignore
ret = self.visit_call_(expr)
+ elif isinstance(expr, _tirx.BufferLoad):
+ ret = self.visit_buffer_load_(expr)
+ elif isinstance(expr, _tirx.Add):
+ ret = self.visit_add_(expr)
+ elif isinstance(expr, _tirx.Sub):
+ ret = self.visit_sub_(expr)
+ elif isinstance(expr, _tirx.Mul):
+ ret = self.visit_mul_(expr)
+ elif isinstance(expr, _tirx.Div):
+ ret = self.visit_div_(expr)
+ elif isinstance(expr, _tirx.Mod):
+ ret = self.visit_mod_(expr)
+ elif isinstance(expr, _tirx.FloorDiv):
+ ret = self.visit_floor_div_(expr)
+ elif isinstance(expr, _tirx.FloorMod):
+ ret = self.visit_floor_mod_(expr)
+ elif isinstance(expr, _tirx.Min):
+ ret = self.visit_min_(expr)
+ elif isinstance(expr, _tirx.Max):
+ ret = self.visit_max_(expr)
+ elif isinstance(expr, _tirx.EQ):
+ ret = self.visit_eq_(expr)
+ elif isinstance(expr, _tirx.NE):
+ ret = self.visit_ne_(expr)
+ elif isinstance(expr, _tirx.LT):
+ ret = self.visit_lt_(expr)
+ elif isinstance(expr, _tirx.LE):
+ ret = self.visit_le_(expr)
+ elif isinstance(expr, _tirx.GT):
+ ret = self.visit_gt_(expr)
+ elif isinstance(expr, _tirx.GE):
+ ret = self.visit_ge_(expr)
+ elif isinstance(expr, _tirx.And):
+ ret = self.visit_and_(expr)
+ elif isinstance(expr, _tirx.Or):
+ ret = self.visit_or_(expr)
+ elif isinstance(expr, _tirx.Cast):
+ ret = self.visit_cast_(expr)
+ elif isinstance(expr, _tirx.Not):
+ ret = self.visit_not_(expr)
+ elif isinstance(expr, _tirx.Select):
+ ret = self.visit_select_(expr)
+ elif isinstance(expr, _tirx.Ramp):
+ ret = self.visit_ramp_(expr)
+ elif isinstance(expr, _tirx.Broadcast):
+ ret = self.visit_broadcast_(expr)
+ elif isinstance(expr, _tirx.Shuffle):
+ ret = self.visit_shuffle_(expr)
+ elif isinstance(expr, _tirx.IntImm):
+ ret = self.visit_int_imm_(expr)
+ elif isinstance(expr, _tirx.FloatImm):
+ ret = self.visit_float_imm_(expr)
+ elif isinstance(expr, _tirx.StringImm):
+ ret = self.visit_tirx_string_imm_(expr)
elif isinstance(expr, SeqExpr):
ret = self.visit_seq_expr_(expr)
elif isinstance(expr, If): # type: ignore
@@ -163,7 +218,11 @@ class ExprFunctor:
ret = self.visit_string_imm_(expr)
elif isinstance(expr, DataTypeImm):
ret = self.visit_data_type_imm_(expr)
+ elif isinstance(expr, _tirx.Let | _tirx.Reduce | _tirx.ProducerLoad):
+ raise TypeError(f"Relax does not support {type(expr).__name__} expressions")
elif isinstance(expr, Expr):
+ if is_prim_expr(expr):
+ raise TypeError(f"Invalid primitive expression type: {type(expr)}")
ret = self.visit_expr_fallback_(expr)
else:
raise TypeError(f"Invalid type: {type(expr)}")
@@ -197,6 +256,87 @@ class ExprFunctor:
def visit_call_(self, op: Call):
raise NotImplementedError()
+ def visit_buffer_load_(self, op: _tirx.BufferLoad):
+ return self.visit_expr_fallback_(op)
+
+ def visit_add_(self, op: _tirx.Add):
+ return self.visit_expr_fallback_(op)
+
+ def visit_sub_(self, op: _tirx.Sub):
+ return self.visit_expr_fallback_(op)
+
+ def visit_mul_(self, op: _tirx.Mul):
+ return self.visit_expr_fallback_(op)
+
+ def visit_div_(self, op: _tirx.Div):
+ return self.visit_expr_fallback_(op)
+
+ def visit_mod_(self, op: _tirx.Mod):
+ return self.visit_expr_fallback_(op)
+
+ def visit_floor_div_(self, op: _tirx.FloorDiv):
+ return self.visit_expr_fallback_(op)
+
+ def visit_floor_mod_(self, op: _tirx.FloorMod):
+ return self.visit_expr_fallback_(op)
+
+ def visit_min_(self, op: _tirx.Min):
+ return self.visit_expr_fallback_(op)
+
+ def visit_max_(self, op: _tirx.Max):
+ return self.visit_expr_fallback_(op)
+
+ def visit_eq_(self, op: _tirx.EQ):
+ return self.visit_expr_fallback_(op)
+
+ def visit_ne_(self, op: _tirx.NE):
+ return self.visit_expr_fallback_(op)
+
+ def visit_lt_(self, op: _tirx.LT):
+ return self.visit_expr_fallback_(op)
+
+ def visit_le_(self, op: _tirx.LE):
+ return self.visit_expr_fallback_(op)
+
+ def visit_gt_(self, op: _tirx.GT):
+ return self.visit_expr_fallback_(op)
+
+ def visit_ge_(self, op: _tirx.GE):
+ return self.visit_expr_fallback_(op)
+
+ def visit_and_(self, op: _tirx.And):
+ return self.visit_expr_fallback_(op)
+
+ def visit_or_(self, op: _tirx.Or):
+ return self.visit_expr_fallback_(op)
+
+ def visit_cast_(self, op: _tirx.Cast):
+ return self.visit_expr_fallback_(op)
+
+ def visit_not_(self, op: _tirx.Not):
+ return self.visit_expr_fallback_(op)
+
+ def visit_select_(self, op: _tirx.Select):
+ return self.visit_expr_fallback_(op)
+
+ def visit_ramp_(self, op: _tirx.Ramp):
+ return self.visit_expr_fallback_(op)
+
+ def visit_broadcast_(self, op: _tirx.Broadcast):
+ return self.visit_expr_fallback_(op)
+
+ def visit_shuffle_(self, op: _tirx.Shuffle):
+ return self.visit_expr_fallback_(op)
+
+ def visit_int_imm_(self, op: _tirx.IntImm):
+ return self.visit_expr_fallback_(op)
+
+ def visit_float_imm_(self, op: _tirx.FloatImm):
+ return self.visit_expr_fallback_(op)
+
+ def visit_tirx_string_imm_(self, op: _tirx.StringImm):
+ return self.visit_expr_fallback_(op)
+
def visit_seq_expr_(self, op: SeqExpr):
raise NotImplementedError()
diff --git a/python/tvm/relax/frontend/nn/core.py b/python/tvm/relax/frontend/nn/core.py
index 73ae8e9c95..3e4a315aee 100644
--- a/python/tvm/relax/frontend/nn/core.py
+++ b/python/tvm/relax/frontend/nn/core.py
@@ -144,7 +144,7 @@ class Tensor(_TensorOp):
"""
return Tensor(
_expr=rx.Var(
- name_hint=name,
+ name=name,
ty=ty,
)
)
@@ -178,7 +178,7 @@ class Tensor(_TensorOp):
new_shape.append(expr)
return Tensor(
_expr=rx.Var(
- name_hint=name,
+ name=name,
ty=TensorType(
shape=new_shape, # type: ignore[arg-type]
dtype=dtype,
@@ -191,8 +191,8 @@ class Tensor(_TensorOp):
"""Returns the shape of the tensor as a list of integers.
An integer can be a python int or tvm.tirx.Expr, depending on whether the shape is
- fully static, for example, [1, 2, tvm.tirx.Var("n")] is a valid shape where the last
- dimension is dynamic while the first two dimensions are always static constants.
+ fully static, for example, [1, 2, tvm.tirx.Var("n", "int64")] is a valid shape where
+ the last dimension is dynamic while the first two dimensions are always static constants.
Returns
-------
diff --git a/python/tvm/relax/frontend/nn/exporter.py b/python/tvm/relax/frontend/nn/exporter.py
index 23406c5b17..a271ad9fea 100644
--- a/python/tvm/relax/frontend/nn/exporter.py
+++ b/python/tvm/relax/frontend/nn/exporter.py
@@ -21,6 +21,7 @@ import operator
import threading
import typing
+import tvm
from tvm import tirx
from tvm.ir import IRModule
@@ -177,7 +178,7 @@ def _emit_method( # pylint: disable=too-many-locals,too-many-branches,too-many-
raise TypeError(f"Unsupported return type: {type(expr)}")
def _convert_input(arg):
- if isinstance(arg, tirx.Var):
+ if isinstance(arg, tvm.ir.Var):
return rx.Var(arg.name, ty=ShapeType(values=[arg]))
if isinstance(arg, core.Tensor | core.Object):
return arg._expr # pylint: disable=protected-access
@@ -202,7 +203,7 @@ def _emit_method( # pylint: disable=too-many-locals,too-many-branches,too-many-
for name, param in params:
# Make sure the a symbolic shape is not re-registered (same as _method_spec_to_inputs)
# e.g. we do not see `vocab_size` for `lm_head` and `vocab_size_1` for `embed_tokens`
- new_shape = [_get_var(x) if isinstance(x, tirx.Var) else x for x in param.shape]
+ new_shape = [_get_var(x) if isinstance(x, tvm.ir.Var) else x for x in param.shape]
var = core.Tensor.placeholder(new_shape, param.dtype, name)._expr
inputs.append(var)
param._expr = var
@@ -243,7 +244,7 @@ def _emit_method( # pylint: disable=too-many-locals,too-many-branches,too-many-
updated_effect_input.append(
builder.emit(
rx.TupleGetItem(input_var, i),
- name_hint=effect_input_i.name_hint,
+ name_hint=effect_input_i.name,
)
)
i += 1
@@ -263,7 +264,7 @@ def _emit_method( # pylint: disable=too-many-locals,too-many-branches,too-many-
return type(arg.elements)(ret)
if isinstance(arg, core.Tensor):
return core.Tensor(_expr=var)
- if isinstance(arg, tirx.Var):
+ if isinstance(arg, tvm.ir.Var):
return arg
raise TypeError(f"Unsupported input type: {type(arg)}")
diff --git a/python/tvm/relax/frontend/nn/op.py b/python/tvm/relax/frontend/nn/op.py
index 18842c2f94..a12c96cdbe 100644
--- a/python/tvm/relax/frontend/nn/op.py
+++ b/python/tvm/relax/frontend/nn/op.py
@@ -748,7 +748,7 @@ def permute_dims(x: Tensor, axes: list[int] | None = None, name: str | None = No
The transposed result.
"""
if name is None:
- x_name = getattr(getattr(x, "_expr", None), "name_hint", None)
+ x_name = getattr(getattr(x, "_expr", None), "name", None)
if x_name is not None and "linear" in x_name:
name = x_name.replace("linear", "matmul")
else:
@@ -2283,7 +2283,7 @@ def debug_func(
*converted_args,
ty_args=[rx.AnyType()],
),
- name_hint=io.effect.name_hint,
+ name_hint=io.effect.name,
)
diff --git a/python/tvm/relax/frontend/onnx/onnx_frontend.py b/python/tvm/relax/frontend/onnx/onnx_frontend.py
index f195f6276a..45cd747e98 100644
--- a/python/tvm/relax/frontend/onnx/onnx_frontend.py
+++ b/python/tvm/relax/frontend/onnx/onnx_frontend.py
@@ -121,11 +121,11 @@ def get_constant(
# Params is actually both the graph nodes and param dictionary, unpack them.
graph_nodes, params = params
# Convert if possible
- if isinstance(var, relax.Var) and var.name_hint in params:
+ if isinstance(var, relax.Var) and var.name in params:
# When converting a parameter to a constant, update references to it as well.
- _, value = params[var.name_hint]
+ _, value = params[var.name]
const_value = relax.const(value)
- graph_nodes[var.name_hint] = const_value
+ graph_nodes[var.name] = const_value
return const_value
# Otherwise return variable.
else:
@@ -583,8 +583,12 @@ class BinaryBase(OnnxOpConverter):
"""Base implementation for binary operations."""
if cls.numpy_op is None or cls.relax_op is None:
raise ValueError("Numpy and Relax operators must be defined for BinaryBase.")
- if all([not isinstance(inp, tvm.ir.Call | relax.Var) for inp in inputs]):
- has_prim_expr = any([tvm.ir.is_prim_expr(inp) for inp in inputs])
+ is_prim_var = [tvm.ir.is_prim_var(inp) for inp in inputs]
+ if all(
+ is_var or not isinstance(inp, tvm.ir.Call | tvm.ir.Var)
+ for inp, is_var in zip(inputs, is_prim_var)
+ ):
+ has_prim_expr = any(tvm.ir.is_prim_expr(inp) for inp in inputs)
x = _to_numpy(inputs[0])
y = _to_numpy(inputs[1])
output = cls.numpy_op(x, y) # pylint: disable=not-callable
@@ -1204,8 +1208,8 @@ class Concat(OnnxOpConverter):
# fast path; tensor fallback keeps the original Vars so runtime
# weights aren't folded under keep_params_in_input=True.
def resolve(x):
- if isinstance(x, relax.Var) and x.name_hint in param_dict:
- arr = param_dict[x.name_hint][1].numpy()
+ if isinstance(x, relax.Var) and x.name in param_dict:
+ arr = param_dict[x.name][1].numpy()
if arr.ndim == 1 and arr.dtype == _np.int64:
return relax.const(arr, "int64")
return x
@@ -3252,7 +3256,7 @@ class Attention(OnnxOpConverter):
QKV = relax.op.matmul(input_emb, weight)
- if bias:
+ if bias is not None:
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"
@@ -4812,24 +4816,24 @@ class EmbedLayerNormalization(OnnxOpConverter):
(batch_size, seq_len) = [dim.value for dim in input_ids.ty.shape]
- if segment_ids:
- assert segment_emb
+ if segment_ids is not None:
+ assert segment_emb is not None
if pos_ids is None:
pos_ids = relax.const([list(range(seq_len))] * batch_size, dtype="int64")
word_vec = relax.op.take(word_emb, input_ids, axis=0)
- if segment_ids:
+ if segment_ids is not None:
segment_vec = relax.op.take(segment_emb, segment_ids, axis=0)
pos_vec = relax.op.take(pos_emb, pos_ids, axis=0)
vec_sum = relax.op.add(word_vec, pos_vec)
- if segment_ids:
+ if segment_ids is not None:
vec_sum = relax.op.add(vec_sum, segment_vec)
ln = relax.op.nn.layer_norm(vec_sum, gamma, beta, axes=-1, epsilon=epsilon)
mask_index = relax.const(_np.zeros((batch_size,), dtype="int64"))
- if mask:
+ if mask is not None:
# Caculate number of words per sentence.
mask_index = relax.op.sum(mask, axis=1)
@@ -5296,7 +5300,7 @@ class NonMaxSuppression(OnnxOpConverter):
elif max_output_boxes_per_class is not None and isinstance(
max_output_boxes_per_class, relax.Var
):
- var_name = max_output_boxes_per_class.name_hint
+ var_name = max_output_boxes_per_class.name
if var_name in params[1]:
_, param_value = params[1][var_name]
max_output_boxes_per_class = int(param_value.numpy().item())
@@ -5308,7 +5312,7 @@ class NonMaxSuppression(OnnxOpConverter):
if iou_threshold is not None and isinstance(iou_threshold, relax.Constant):
iou_threshold = float(iou_threshold.data.numpy().item())
elif iou_threshold is not None and isinstance(iou_threshold, relax.Var):
- var_name = iou_threshold.name_hint
+ var_name = iou_threshold.name
if var_name in params[1]:
_, param_value = params[1][var_name]
iou_threshold = float(param_value.numpy().item())
@@ -5320,7 +5324,7 @@ class NonMaxSuppression(OnnxOpConverter):
if score_threshold is not None and isinstance(score_threshold, relax.Constant):
score_threshold = float(score_threshold.data.numpy().item())
elif score_threshold is not None and isinstance(score_threshold, relax.Var):
- var_name = score_threshold.name_hint
+ var_name = score_threshold.name
if var_name in params[1]:
_, param_value = params[1][var_name]
score_threshold = float(param_value.numpy().item())
@@ -5392,7 +5396,7 @@ class AllClassNMS(OnnxOpConverter):
elif max_output_boxes_per_class is not None and isinstance(
max_output_boxes_per_class, relax.Var
):
- var_name = max_output_boxes_per_class.name_hint
+ var_name = max_output_boxes_per_class.name
if var_name in params[1]:
_, param_value = params[1][var_name]
max_output_boxes_per_class = int(param_value.numpy().item())
@@ -5404,7 +5408,7 @@ class AllClassNMS(OnnxOpConverter):
if iou_threshold is not None and isinstance(iou_threshold, relax.Constant):
iou_threshold = float(iou_threshold.data.numpy().item())
elif iou_threshold is not None and isinstance(iou_threshold, relax.Var):
- var_name = iou_threshold.name_hint
+ var_name = iou_threshold.name
if var_name in params[1]:
_, param_value = params[1][var_name]
iou_threshold = float(param_value.numpy().item())
@@ -5416,7 +5420,7 @@ class AllClassNMS(OnnxOpConverter):
if score_threshold is not None and isinstance(score_threshold, relax.Constant):
score_threshold = float(score_threshold.data.numpy().item())
elif score_threshold is not None and isinstance(score_threshold, relax.Var):
- var_name = score_threshold.name_hint
+ var_name = score_threshold.name
if var_name in params[1]:
_, param_value = params[1][var_name]
score_threshold = float(param_value.numpy().item())
@@ -5895,7 +5899,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, ty=relax.TensorType(shape=shape, dtype=dtype))
+ return relax.Var(name=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."""
diff --git a/python/tvm/relax/frontend/stablehlo/stablehlo_translator.py b/python/tvm/relax/frontend/stablehlo/stablehlo_translator.py
index 51994122d7..d6330b4058 100644
--- a/python/tvm/relax/frontend/stablehlo/stablehlo_translator.py
+++ b/python/tvm/relax/frontend/stablehlo/stablehlo_translator.py
@@ -25,6 +25,11 @@ import tvm
from tvm import relax, tirx
+def _is_tensor_var(value) -> bool:
+ """Whether ``value`` belongs to Relax's tensor domain."""
+ return isinstance(value, tvm.ir.Var) and isinstance(value.ty, relax.TensorType)
+
+
class StableHLOImporter:
"""An importer from StableHLO to Relax."""
@@ -156,7 +161,7 @@ class StableHLOImporter:
def _add(self, node: mlir.ir.Operation) -> relax.Expr:
lhs, rhs = self.retrieve_operands(node)
- if isinstance(lhs, relax.Var) or isinstance(rhs, relax.Var):
+ if _is_tensor_var(lhs) or _is_tensor_var(rhs):
return self._call_binary_op(relax.op.add, lhs, rhs)
return lhs + rhs
@@ -170,19 +175,19 @@ class StableHLOImporter:
def _divide(self, node: mlir.ir.Operation) -> relax.Expr:
lhs, rhs = self.retrieve_operands(node)
- if isinstance(lhs, relax.Var) or isinstance(rhs, relax.Var):
+ if _is_tensor_var(lhs) or _is_tensor_var(rhs):
return self._call_binary_op(relax.op.divide, lhs, rhs)
return lhs / rhs
def _multiply(self, node: mlir.ir.Operation) -> relax.Expr:
lhs, rhs = self.retrieve_operands(node)
- if isinstance(lhs, relax.Var) or isinstance(rhs, relax.Var):
+ if _is_tensor_var(lhs) or _is_tensor_var(rhs):
return self._call_binary_op(relax.op.multiply, lhs, rhs)
return lhs * rhs
def _subtract(self, node: mlir.ir.Operation) -> relax.Expr:
lhs, rhs = self.retrieve_operands(node)
- if isinstance(lhs, relax.Var) or isinstance(rhs, relax.Var):
+ if _is_tensor_var(lhs) or _is_tensor_var(rhs):
return self._call_binary_op(relax.op.subtract, lhs, rhs)
return lhs - rhs
diff --git a/python/tvm/relax/frontend/tflite/tflite_frontend.py b/python/tvm/relax/frontend/tflite/tflite_frontend.py
index a2db135466..d271221fc7 100644
--- a/python/tvm/relax/frontend/tflite/tflite_frontend.py
+++ b/python/tvm/relax/frontend/tflite/tflite_frontend.py
@@ -5180,7 +5180,7 @@ class OperatorConverter:
if self.has_expr(shape_tensor.tensor_idx):
shape_expr = self.get_expr(shape_tensor.tensor_idx)
shape_expr = self.bb.normalize(relax.op.astype(shape_expr, "int64"))
- shape = self.bb.emit(relax.op.tensor_to_shape(shape_expr))
+ shape = self.bb.emit_output(relax.op.tensor_to_shape(shape_expr))
else:
shape = to_int_list(self.get_tensor_value(shape_tensor))
@@ -6780,7 +6780,7 @@ class OperatorConverter:
shape_tensor = input_tensors[1]
if self.has_expr(shape_tensor.tensor_idx):
shape_expr = self.get_expr(shape_tensor.tensor_idx)
- shape = self.bb.emit(relax.op.tensor_to_shape(shape_expr))
+ shape = self.bb.emit_output(relax.op.tensor_to_shape(shape_expr))
else:
shape = to_int_list(self.get_tensor_value(shape_tensor))
return relax.op.broadcast_to(data, shape)
@@ -9006,7 +9006,7 @@ def from_tflite(
"(e.g. runtime string queries)"
)
input_var = relax.Var(
- name_hint=model_input_name,
+ name=model_input_name,
ty=relax.TensorType(shape=shape, dtype=dtype),
)
exp_tab.set_expr(model_input_name, 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 c62fbf2ace..85a0d956ca 100644
--- a/python/tvm/relax/frontend/torch/base_fx_graph_translator.py
+++ b/python/tvm/relax/frontend/torch/base_fx_graph_translator.py
@@ -557,7 +557,10 @@ class BaseFXGraphImporter(metaclass=abc.ABCMeta):
return self.block_builder.emit(op(lhs, rhs))
lhs, rhs = self.retrieve_args(node)
- if isinstance(lhs, relax.Var) or isinstance(rhs, relax.Var):
+
+ if (isinstance(lhs, tvm.ir.Var) and isinstance(lhs.ty, relax.TensorType)) or (
+ isinstance(rhs, tvm.ir.Var) and isinstance(rhs.ty, relax.TensorType)
+ ):
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.ty.dtype))
@@ -2623,9 +2626,9 @@ class BaseFXGraphImporter(metaclass=abc.ABCMeta):
if tensor_meta is not None and len(tensor_meta.shape) == 1:
num_selected = tensor_meta.shape[0]
if not isinstance(num_selected, int):
- num_selected = tirx.Var(str(num_selected), "int64")
+ num_selected = tvm.ir.Var(str(num_selected), "int64")
else:
- num_selected = tirx.Var(f"{node.name}_num_selected", "int64")
+ num_selected = tvm.ir.Var(f"{node.name}_num_selected", "int64")
indices = self.block_builder.match_cast(
indices, relax.TensorType([1, num_selected], "int64")
)
@@ -2714,7 +2717,7 @@ class BaseFXGraphImporter(metaclass=abc.ABCMeta):
x = self.env[node.args[0]]
if isinstance(x, list | tuple | relax.ShapeExpr | relax.Tuple):
return x[node.args[1]]
- elif isinstance(x, relax.Var):
+ elif isinstance(x, tvm.ir.Var) and isinstance(x.ty, relax.TupleType | relax.TensorType):
if isinstance(x.ty, relax.TupleType):
return self.block_builder.emit(relax.TupleGetItem(x, node.args[1]))
diff --git a/python/tvm/relax/frontend/torch/exported_program_translator.py b/python/tvm/relax/frontend/torch/exported_program_translator.py
index b11356ec77..ea604c21cc 100644
--- a/python/tvm/relax/frontend/torch/exported_program_translator.py
+++ b/python/tvm/relax/frontend/torch/exported_program_translator.py
@@ -1618,7 +1618,7 @@ class ExportedProgramImporter(BaseFXGraphImporter):
# Create fresh symbolic variables to avoid sharing with the caller function.
if orig_si.shape is not None:
new_shape = [
- tvm.tirx.Var(s.name, s.ty) if isinstance(s, tvm.tirx.Var) else s
+ tvm.ir.Var(s.name, s.ty) if tvm.ir.is_prim_var(s) else s
for s in orig_si.shape
]
si = relax.TensorType(new_shape, orig_si.dtype)
@@ -2077,12 +2077,12 @@ class ExportedProgramImporter(BaseFXGraphImporter):
if isinstance(tir_expr, tvm.tirx.Add):
for const, var in [(tir_expr.a, tir_expr.b), (tir_expr.b, tir_expr.a)]:
- if isinstance(const, tvm.tirx.IntImm) and isinstance(var, tvm.tirx.Var):
+ if isinstance(const, tvm.tirx.IntImm) and tvm.ir.is_prim_var(var):
return f"{var.name}___{const.value}", tir_expr
if isinstance(tir_expr, tvm.tirx.Mul):
for const, var in [(tir_expr.a, tir_expr.b), (tir_expr.b, tir_expr.a)]:
- if isinstance(const, tvm.tirx.IntImm) and isinstance(var, tvm.tirx.Var):
+ if isinstance(const, tvm.tirx.IntImm) and tvm.ir.is_prim_var(var):
return f"{var.name}_{const.value}", tir_expr
return str(symbol), tir_expr
diff --git a/python/tvm/relax/frontend/torch/fx_translator.py b/python/tvm/relax/frontend/torch/fx_translator.py
index 75098ee9af..bef0b58f08 100644
--- a/python/tvm/relax/frontend/torch/fx_translator.py
+++ b/python/tvm/relax/frontend/torch/fx_translator.py
@@ -177,7 +177,10 @@ class TorchFXImporter(BaseFXGraphImporter):
return self.block_builder.emit(op(lhs, rhs))
lhs, rhs = self.retrieve_args(node)
- if isinstance(lhs, relax.Var) or isinstance(rhs, relax.Var):
+
+ if (isinstance(lhs, tvm.ir.Var) and isinstance(lhs.ty, relax.TensorType)) or (
+ isinstance(rhs, tvm.ir.Var) and isinstance(rhs.ty, relax.TensorType)
+ ):
output = call_binary_op(relax_op, lhs, rhs)
self.env[node.args[0]] = output
return output
@@ -1093,9 +1096,9 @@ class TorchFXImporter(BaseFXGraphImporter):
# Find all the missing function types
self._check_unsupported_func_type(graph.nodes)
- from tvm import tirx
-
- sym_vars = {v.name: v for shape, _ in input_info for v in shape if isinstance(v, tirx.Var)}
+ sym_vars = {
+ v.name: v for shape, _ in input_info for v in shape if isinstance(v, tvm.ir.Var)
+ }
with self.block_builder.function(name=func_name, params=inputs.copy(), attrs=func_attrs):
output = None
diff --git a/python/tvm/relax/relax_to_pyfunc_converter.py b/python/tvm/relax/relax_to_pyfunc_converter.py
index 0099b774a8..a90628bc34 100644
--- a/python/tvm/relax/relax_to_pyfunc_converter.py
+++ b/python/tvm/relax/relax_to_pyfunc_converter.py
@@ -413,36 +413,34 @@ class RelaxExpressionConverter:
def _convert_var(self, var: relax.Var, args: list[Any]) -> Any:
"""Convert a Relax variable to Python equivalent."""
- if hasattr(var, "name_hint"):
- var_name = var.name_hint
-
- # Check if it's a function parameter
- for i, param in enumerate(self.current_params):
- if hasattr(param, "name_hint") and param.name_hint == var_name:
- return args[i]
-
- # Check if it's a bound variable
- if var_name in self.variable_map:
- return self.variable_map[var_name]
-
- # Try to infer shape from var's type annotation
- 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 = []
- for dim in shape:
- if isinstance(dim, int):
- concrete_shape.append(dim)
- else:
- # For symbolic dimensions, use a reasonable default
- concrete_shape.append(1)
- return torch.zeros(concrete_shape, dtype=torch.float32)
+ var_name = var.name
+
+ # Check if it's a function parameter
+ for i, param in enumerate(self.current_params):
+ if param.name == var_name:
+ return args[i]
+
+ # Check if it's a bound variable
+ if var_name in self.variable_map:
+ return self.variable_map[var_name]
+
+ # Try to infer shape from var's type annotation
+ 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 = []
+ for dim in shape:
+ if isinstance(dim, int):
+ concrete_shape.append(dim)
+ else:
+ # For symbolic dimensions, use a reasonable default
+ concrete_shape.append(1)
+ return torch.zeros(concrete_shape, dtype=torch.float32)
- if args and isinstance(args[0], torch.Tensor):
- return torch.zeros_like(args[0])
- # Use fallback tensor with shape inference
- return self._create_fallback_tensor()
+ if args and isinstance(args[0], torch.Tensor):
+ return torch.zeros_like(args[0])
+ # Use fallback tensor with shape inference
return self._create_fallback_tensor()
def _convert_call(self, call: relax.Call, args: list[Any]) -> Any:
@@ -806,7 +804,7 @@ class RelaxExpressionConverter:
if hasattr(block, "bindings"):
for binding in block.bindings:
if isinstance(binding, relax.VarBinding):
- var_name = binding.var.name_hint
+ var_name = binding.var.name
value = self.convert_expr(binding.value, args)
self.variable_map[var_name] = value
diff --git a/python/tvm/relax/script/parser/parser.py b/python/tvm/relax/script/parser/parser.py
index ced4054e7f..bef126a884 100644
--- a/python/tvm/relax/script/parser/parser.py
+++ b/python/tvm/relax/script/parser/parser.py
@@ -20,10 +20,11 @@ import functools
import numbers
from typing import Any
+import numpy as np
import tvm_ffi
import tvm
-from tvm import relax, tirx
+from tvm import relax
from tvm.ir import GlobalVar
from tvm.relax import Expr, Type
from tvm.relax.script import builder as R
@@ -32,9 +33,11 @@ from tvm.relax.utils import convert_to_expr
from tvm.script.ir_builder import ir as I
from tvm.script.ir_builder.base import IRBuilder
from tvm.script.parser._core import Parser, dispatch, doc
+from tvm.tirx.script import builder as T
from .entry import (
MatchCastPair,
+ PrimProxy,
TypeProxy,
_normalize_ty,
_normalize_ty_proxy,
@@ -53,49 +56,54 @@ def bind_assign_value(
var_name: str,
value: Any,
anno_ty: Type | None = None,
- emit_prim_expr: bool = False,
+ symbolic_declarations: set[str] | None = None,
) -> Any:
var_table = self.var_table.get()
- if isinstance(value, tirx.Var):
+ # Primitive assignments emit Relax bindings by default, reversing the old default.
+ # Explicit I.meta_var values and literal T.dtype() declarations are the two
+ # parser-time exceptions.
+ if isinstance(value, I.meta_var):
+ return value.value
+
+ if tvm.ir.is_prim_var(value) and var_name in (symbolic_declarations or set()):
+ if anno_ty is not None and not tvm_ffi.structural_equal(anno_ty, value.ty):
+ self.report_error(
+ node,
+ f"Expected the declared symbolic variable to have type {anno_ty}, "
+ f"but got {value.ty}",
+ )
if value.name and var_name != value.name:
self.report_error(
node,
- "Cannot define TIR variables with different names. The LHS of binding should "
- "has the same name provided in RHS.",
+ "Cannot define PrimType values with different names. The LHS of binding "
+ "must have the same name provided in the RHS.",
)
if var_name in var_table:
prev_value = var_table[var_name]
- if not isinstance(prev_value, tirx.Var):
+ if not tvm.ir.is_prim_var(prev_value):
self.report_error(
node,
- "Cannot redefine a non-TIR-variable object to a TIR variable. Please "
- "define the TIR variable with another name.",
+ "Cannot redefine a non-PrimType object to a PrimType value. Please "
+ "define the PrimType value with another name.",
)
if prev_value.ty != value.ty:
self.report_error(
node,
- f"Expected the same dtype for TIR vars but got {value.ty} vs {prev_value.ty}",
- )
- if not isinstance(value, type(prev_value)):
- self.report_error(
- node,
- f"Expected the same IR type for TIR vars "
- f"but existing value {type(value)} is mismatched "
- f"to previous {type(prev_value)}",
+ f"Expected the same dtype for PrimType values but got "
+ f"{value.ty} vs {prev_value.ty}",
)
value = prev_value
IRBuilder.name(var_name, value)
return value
- if tvm.ir.is_prim_expr(value):
- if not emit_prim_expr:
- return value
-
if isinstance(value, tuple):
value = convert_to_expr(value)
- if isinstance(value, numbers.Number):
- value = R.const(value)
+ if isinstance(value, numbers.Number | np.bool_):
+ if isinstance(anno_ty, tvm.ir.PrimType):
+ value = relax.prim_value(value, dtype=anno_ty.dtype)
+ else:
+ value = R.const(value)
if isinstance(value, relax.Expr):
var = R.emit(value, anno_ty)
@@ -112,8 +120,46 @@ def bind_assign_value(
return var
-def is_prim_value_call(node: doc.expr) -> bool:
- return isinstance(node, doc.Call) and getattr(node.func, "attr", None) == "prim_value"
+def is_symbolic_var_declaration(node: doc.expr) -> bool:
+ """Return whether an expression is literal ``T.dtype()`` declaration syntax.
+
+ Detection is syntactic so declaration semantics do not depend on evaluating its RHS.
+ """
+ if not (
+ isinstance(node, doc.Call)
+ and not node.args
+ and not node.keywords
+ and isinstance(node.func, doc.Attribute)
+ and isinstance(node.func.value, doc.Name)
+ and node.func.value.id == "T"
+ ):
+ return False
+ constructor = getattr(T, node.func.attr, None)
+ return isinstance(constructor, T.DtypeConstructor) or constructor is T.bool
+
+
+def collect_symbolic_var_declaration_nodes(
+ target: doc.expr, value: doc.expr
+) -> dict[str, doc.expr]:
+ """Pair targets with literal ``T.dtype()`` declarations.
+
+ Declaration-vs-binding is per target, allowing tuple unpacking to mix both.
+ """
+ if isinstance(target, doc.Name):
+ return {target.id: value} if is_symbolic_var_declaration(value) else {}
+ if isinstance(target, doc.Tuple | doc.List) and isinstance(value, doc.Tuple | doc.List):
+ if len(target.elts) != len(value.elts):
+ return {}
+ declarations = {}
+ for lhs, rhs in zip(target.elts, value.elts):
+ declarations.update(collect_symbolic_var_declaration_nodes(lhs, rhs))
+ return declarations
+ return {}
+
+
+def collect_symbolic_var_declarations(target: doc.expr, value: doc.expr) -> set[str]:
+ """Collect assignment targets whose matching RHS is ``T.dtype()``."""
+ return set(collect_symbolic_var_declaration_nodes(target, value))
def eval_ty_proxy(self: Parser, node: doc.expr) -> TypeProxy:
@@ -168,41 +214,42 @@ def is_recursive(node: doc.FunctionDef) -> bool:
def collect_symbolic_var_from_prelude(
- self: Parser, node: doc.FunctionDef, symbolic_vars: dict[str, tirx.Var]
-) -> dict[str, tirx.Var]:
+ self: Parser, node: doc.FunctionDef, symbolic_vars: dict[str, tvm.ir.Var]
+) -> dict[str, tvm.ir.Var]:
prelude_vars = {}
for stmt in node.body:
- if isinstance(stmt, doc.Assign) and all(
- isinstance(target, doc.Name) and target.id in symbolic_vars for target in stmt.targets
- ):
- values = self.eval_expr(stmt.value)
-
- try:
- iter(values)
- except TypeError:
- values = [values]
-
- assert len(stmt.targets) == len(values)
- for target, value in zip(stmt.targets, values):
- name = target.id
- prelude_vars[name] = value
+ if isinstance(stmt, doc.Assign) and len(stmt.targets) == 1:
+ declarations = collect_symbolic_var_declaration_nodes(stmt.targets[0], stmt.value)
+ for name, value_node in declarations.items():
+ if name not in symbolic_vars:
+ continue
+ declared_var = self.eval_expr(value_node)
+ if tvm.ir.is_prim_var(declared_var) and declared_var.ty == symbolic_vars[name].ty:
+ prelude_vars[name] = declared_var
return {**symbolic_vars, **prelude_vars}
def collect_symbolic_var_from_params(self: Parser, node: doc.FunctionDef) -> None:
- # Collect symbolic vars from parameters
symbolic_vars = {}
- for arg in node.args.args:
- if arg.annotation is None:
- self.report_error(arg, "Type annotation is required for function parameters.")
- param_ty_proxy = eval_ty_proxy(self, arg.annotation)
+ prim_params = set()
+ with self.var_table.with_frame():
+ for arg in node.args.args:
+ if arg.annotation is None:
+ self.report_error(arg, "Type annotation is required for function parameters.")
+ param_ty_proxy = eval_ty_proxy(self, arg.annotation)
+
+ for var_name in param_ty_proxy.get_symbolic_vars():
+ if var_name not in prim_params and var_name not in symbolic_vars:
+ symbolic_vars[var_name] = tvm.ir.Var(var_name, "int64")
- for var_name in param_ty_proxy.get_symbolic_vars():
- if var_name not in symbolic_vars:
- symbolic_vars[var_name] = tirx.Var(var_name, "int64")
+ if isinstance(param_ty_proxy, PrimProxy):
+ temp_param = tvm.ir.Var(arg.arg, param_ty_proxy.as_ty())
+ self.var_table.add(arg.arg, temp_param)
+ prim_params.add(arg.arg)
- # Update symbolic vars based on
+ # Prelude declarations select each scope symbol's canonical Var object.
+ # Per-assignment declaration classification separately controls emission.
symbolic_vars = collect_symbolic_var_from_prelude(self, node, symbolic_vars)
# Define symbolic vars to the current var_table frame
@@ -210,6 +257,20 @@ def collect_symbolic_var_from_params(self: Parser, node: doc.FunctionDef) -> Non
self.var_table.add(var_name, var, allow_shadowing=False)
+def parse_function_params(self: Parser, node: doc.FunctionDef, param_factory) -> list[tvm.ir.Var]:
+ collect_symbolic_var_from_params(self, node)
+ params = []
+ for arg in node.args.args:
+ if arg.annotation is None:
+ self.report_error(arg, "Type annotation is required for function parameters.")
+ param_ty = eval_ty(self, arg.annotation, eval_str=True)
+ param = param_factory(arg.arg, param_ty)
+ self.var_table.add(arg.arg, param)
+ params.append(param)
+
+ return params
+
+
@dispatch.register(token="relax", type_name="FunctionDef")
def visit_function_def(self: Parser, node: doc.FunctionDef) -> None:
is_inner_function = self.inside_function
@@ -217,20 +278,16 @@ def visit_function_def(self: Parser, node: doc.FunctionDef) -> None:
# reserve a var for local function
func_val = self.var_table.get().get(node.name)
- if not func_val and is_recursive(node):
- collect_symbolic_var_from_params(self, node)
- if node.returns is None:
- ret_ty = relax.TupleType([])
- else:
- ret_ty = eval_ty(self, node.returns, eval_str=True)
- params_ty = []
- for arg in node.args.args:
- if arg.annotation is None:
- self.report_error(arg, "Type annotation is required for function parameters.")
- param_ty = eval_ty(self, arg.annotation, eval_str=True)
- params_ty.append(param_ty)
+ if func_val is None and is_recursive(node):
+ with self.var_table.with_frame():
+ provisional_params = parse_function_params(self, node, tvm.ir.Var)
+ if node.returns is None:
+ ret_ty = relax.TupleType([])
+ else:
+ ret_ty = eval_ty(self, node.returns, eval_str=True)
+ params_ty = [param.ty for param in provisional_params]
# created a var for the local function, the same var could be used for recursive call
- local_func_var = relax.Var(node.name, relax.FuncType(params_ty, ret_ty))
+ local_func_var = tvm.ir.Var(node.name, relax.FuncType(params_ty, ret_ty))
self.var_table.add(node.name, local_func_var)
purity = find_decorator_annotation(node, "pure")
@@ -242,14 +299,12 @@ def visit_function_def(self: Parser, node: doc.FunctionDef) -> None:
with self.with_dispatch_token("relax"):
with R.function(is_pure=purity, is_private=privacy):
R.func_name(node.name)
- collect_symbolic_var_from_params(self, node)
+ parse_function_params(self, node, R.arg)
if node.returns is not None:
ann_ty = eval_ty(self, node.returns, eval_str=True)
R.func_ret_ty(ann_ty)
- self.visit(node.args)
-
for stmt in node.body:
if isinstance(stmt, doc.FunctionDef):
if not stmt.decorator_list:
@@ -281,7 +336,7 @@ def find_decorator_annotation(node: doc.FunctionDef, annotation: str, default: b
@dispatch.register(token="relax", type_name="tvm_declare_function")
def visit_tvm_declare_function(self: Parser, node: doc.FunctionDef) -> GlobalVar:
with self.var_table.with_frame():
- collect_symbolic_var_from_params(self, node)
+ params = parse_function_params(self, node, tvm.ir.Var)
if node.returns is None:
# Use AnyType as unknown return type
@@ -289,13 +344,6 @@ def visit_tvm_declare_function(self: Parser, node: doc.FunctionDef) -> GlobalVar
ret_ty = relax.AnyType()
else:
ret_ty = eval_ty(self, node.returns, eval_str=True)
- params = []
- for arg in node.args.args:
- if arg.annotation is None:
- self.report_error(arg, "Type annotation is required for function parameters.")
- param_ty = eval_ty(self, arg.annotation, eval_str=True)
- params.append(relax.Var(arg.arg, param_ty))
-
is_pure = find_decorator_annotation(node, "pure")
func_signature = relax.Function.create_empty(params, ret_ty, is_pure=is_pure)
@@ -315,7 +363,7 @@ def post_visit_local_function(self: Parser, node: doc.Expr) -> None:
ir_builder.__exit__(None, None, None)
# reuse var if it is reserved
reserved_var = self.var_table.get().get(node.name)
- if reserved_var:
+ if reserved_var is not None:
var = R.emit_var_binding(relax.VarBinding(reserved_var, result))
else:
var = R.emit(result)
@@ -377,7 +425,7 @@ def visit_with(self: Parser, node: doc.With) -> None:
if isinstance(frame, BindingBlockFrame) and frame.is_dataflow:
output_vars = frame.output_vars
for var in output_vars:
- self.var_table.add(var.name_hint, var, allow_shadowing=True)
+ self.var_table.add(var.name, var, allow_shadowing=True)
@dispatch.register(token="relax", type_name="Assign")
@@ -391,7 +439,7 @@ def visit_assign(self: Parser, node: doc.Assign) -> None:
source=rhs,
bind_value=functools.partial(
bind_assign_value,
- emit_prim_expr=is_prim_value_call(node.value),
+ symbolic_declarations=collect_symbolic_var_declarations(lhs, node.value),
),
allow_shadowing=True,
)
@@ -408,7 +456,7 @@ def visit_ann_assign(self: Parser, node: doc.AnnAssign) -> None:
bind_value=functools.partial(
bind_assign_value,
anno_ty=anno_ty,
- emit_prim_expr=is_prim_value_call(node.value),
+ symbolic_declarations=collect_symbolic_var_declarations(lhs, node.value),
),
allow_shadowing=True,
)
diff --git a/python/tvm/relax/testing/ast_printer.py b/python/tvm/relax/testing/ast_printer.py
index 970bc6e8cf..55a0048ea7 100644
--- a/python/tvm/relax/testing/ast_printer.py
+++ b/python/tvm/relax/testing/ast_printer.py
@@ -122,10 +122,10 @@ class ASTPrinter(ExprFunctor):
return self.build_expr(op, "Tuple", fields=self.build_list(map(self.visit_expr, op.fields)))
def visit_dataflow_var_(self, op: relax.DataflowVar) -> str:
- return self.build_expr(op, "DataflowVar", name_hint=wrap_quotes(op.name_hint))
+ return self.build_expr(op, "DataflowVar", name=wrap_quotes(op.name))
def visit_var_(self, op: relax.Var) -> str:
- return self.build_expr(op, "Var", name_hint=wrap_quotes(op.name_hint))
+ return self.build_expr(op, "Var", name=wrap_quotes(op.name))
def visit_shape_expr_(self, op: relax.ShapeExpr) -> str:
return self.build_expr(
@@ -221,7 +221,8 @@ class ASTPrinter(ExprFunctor):
def visit_prim_expr_field_(self, prim_expr: Expr) -> str:
# TODO: We may want to print Expr ASTs, but this is a simplification for now
- return self.build_ast_node("Expr", value=f"`{prim_expr!s}`")
+ value = prim_expr.name if isinstance(prim_expr, tvm.ir.Var) else str(prim_expr)
+ return self.build_ast_node("Expr", value=f"`{value}`")
def visit_expr_fallback_(self, op: Expr) -> str:
if not tvm.ir.is_prim_expr(op):
@@ -288,7 +289,7 @@ class ASTPrinter(ExprFunctor):
elif isinstance(ty_node, relax.TensorType):
fields = {}
fields["dtype"] = ty_node.dtype
- if ty_node.shape:
+ if ty_node.shape is not None:
fields["shape"] = self.visit_expr(ty_node.shape)
else:
fields["ndim"] = str(ty_node.ndim)
diff --git a/python/tvm/relax/testing/nn.py b/python/tvm/relax/testing/nn.py
index d2297f3766..5b0d98d326 100644
--- a/python/tvm/relax/testing/nn.py
+++ b/python/tvm/relax/testing/nn.py
@@ -69,11 +69,11 @@ def emit_checkpoint(
"""
bb = relax.BlockBuilder.current()
args = [
- bb.emit(start_checkpoint(v), v.name_hint + "_scp") if isinstance(v, relax.Var) else v
+ bb.emit(start_checkpoint(v), v.name + "_scp") if isinstance(v, relax.Var) else v
for v in args
]
kwargs = {
- k: bb.emit(start_checkpoint(v), v.name_hint + "_scp") if isinstance(v, relax.Var) else v
+ k: bb.emit(start_checkpoint(v), v.name + "_scp") if isinstance(v, relax.Var) else v
for k, v in kwargs.items()
}
result = func(*args, **kwargs)
@@ -83,11 +83,11 @@ def emit_checkpoint(
if isinstance(v, relax.Expr):
if not isinstance(v, relax.Var):
v = bb.emit(v)
- result[i] = bb.emit(end_checkpoint(v), v.name_hint + "_ecp")
+ result[i] = bb.emit(end_checkpoint(v), v.name + "_ecp")
else:
assert isinstance(result, relax.Expr)
result_emit = bb.emit(result)
- result = bb.emit(end_checkpoint(result_emit), result_emit.name_hint + "_ecp")
+ result = bb.emit(end_checkpoint(result_emit), result_emit.name + "_ecp")
return result
@@ -279,7 +279,7 @@ def _unpack_params(value: object) -> list[relax.Var]:
def init_params(mod: tvm.IRModule) -> list[tvm.runtime.Tensor]:
"""Utility function to initialize model's parameters."""
- shape_dict = {v.name_hint: v.ty.shape for v in mod["main"].params}
+ shape_dict = {v.name: v.ty.shape for v in mod["main"].params}
params = []
for k, v in shape_dict.items():
if k.startswith("data"):
diff --git a/python/tvm/relax/training/loss.py b/python/tvm/relax/training/loss.py
index 79d07b1bc2..469b7dba89 100644
--- a/python/tvm/relax/training/loss.py
+++ b/python/tvm/relax/training/loss.py
@@ -37,7 +37,7 @@ def _create_param_var(param: Var | Type, param_name: str) -> Var:
param = Var(param_name, param)
if not isinstance(param, Var):
raise TypeError("The type of param should be Var or Type, but got " + type(param))
- return Var(param.name_hint, param.ty)
+ return Var(param.name, param.ty)
class Loss:
@@ -275,7 +275,7 @@ class CrossEntropyLoss(Loss):
targets = _create_param_var(targets, "targets")
arg_list = [predictions, targets]
- if weights:
+ if weights is not None:
weights = _create_param_var(weights, "weights")
arg_list.append(weights)
@@ -355,7 +355,7 @@ class CategoricalCrossEntropyLoss(Loss):
targets = _create_param_var(targets, "targets")
arg_list = [predictions, targets]
- if weights:
+ if weights is not None:
weights = _create_param_var(weights, "weights")
arg_list.append(weights)
@@ -374,7 +374,7 @@ class CategoricalCrossEntropyLoss(Loss):
)
else:
lv = bb.emit(-logits * targets.astype("float32"))
- if weights:
+ if weights is not None:
lv = bb.emit(lv * weights)
loss = bb.emit_output(self._with_reduction(lv))
bb.emit_func_output(loss)
diff --git a/python/tvm/relax/training/optimizer.py b/python/tvm/relax/training/optimizer.py
index f35767e67f..25758bc795 100644
--- a/python/tvm/relax/training/optimizer.py
+++ b/python/tvm/relax/training/optimizer.py
@@ -146,26 +146,26 @@ class Optimizer:
raise ValueError(f"Parameter {x} is not a Var")
if not isinstance(x.ty, TensorType):
raise ValueError(
- f"Optimizers only support Tensor parameters, but parameter {x.name_hint} has "
+ f"Optimizers only support Tensor parameters, but parameter {x.name} has "
f"type {x.ty}"
)
data_type = tvm.DataType(x.ty.dtype.dtype)
if data_type.type_code not in (tvm.DataTypeCode.BFLOAT, tvm.DataTypeCode.FLOAT):
raise ValueError(
f"Optimizers only support Tensor parameters of floating point dtype, but dtype "
- f"of {x.name_hint} is {x.ty.dtype}"
+ f"of {x.name} is {x.ty.dtype}"
)
if dtype is None:
dtype = x.ty.dtype
else:
if dtype != x.ty.dtype:
raise ValueError(
- f"All parameters should have the same dtype, but parameter {x.name_hint} "
+ f"All parameters should have the same dtype, but parameter {x.name} "
f"has dtype {x.ty.dtype}, which differs from the previous dtype "
f"{dtype}"
)
if x in params_set:
- raise ValueError(f"Parameter {x.name_hint} appears more than once")
+ raise ValueError(f"Parameter {x.name} appears more than once")
params_set.add(x)
self.param_list = params
self.dtype = dtype
@@ -338,7 +338,7 @@ class SGD(Optimizer):
# computation logics
for i in range(len_param):
- name = self.param_list[i].name_hint
+ name = self.param_list[i].name
p = builder.emit(TupleGetItem(param_var, i), name)
g = builder.emit(TupleGetItem(grad_var, i), name + "_grad")
if self.weight_decay:
@@ -490,7 +490,7 @@ class MomentumSGD(Optimizer):
# computation logics
for i in range(len_param):
- name = self.param_list[i].name_hint
+ name = self.param_list[i].name
p = builder.emit(TupleGetItem(param_var, i), name)
g = builder.emit(TupleGetItem(grad_var, i), name + "_grad")
v = builder.emit(TupleGetItem(state_var, i + 1), name + "_v")
@@ -687,7 +687,7 @@ class Adam(Optimizer):
# computation logics
for i in range(len_param):
- name = self.param_list[i].name_hint
+ name = self.param_list[i].name
p = builder.emit(TupleGetItem(param_var, i), name)
g = builder.emit(TupleGetItem(grad_var, i), name + "_grad")
m = builder.emit(TupleGetItem(state_var, i + 3), name + "_m")
diff --git a/python/tvm/relax/training/trainer.py b/python/tvm/relax/training/trainer.py
index 4dfbafb809..997b684405 100644
--- a/python/tvm/relax/training/trainer.py
+++ b/python/tvm/relax/training/trainer.py
@@ -102,12 +102,12 @@ class Trainer:
self._params: list[Tensor | None] = [None] * self._param_num
self._param_name_to_pos: dict[str, int] = {
- p.name_hint: i for i, p in enumerate(self._param_vars)
+ p.name: i for i, p in enumerate(self._param_vars)
}
self._states: list[Tensor | None] = [None] * self._state_num
self._state_name_to_pos: dict[str, int] = {
- s.name_hint: i for i, s in enumerate(self._state_vars)
+ s.name: i for i, s in enumerate(self._state_vars)
}
if zero_init_param_state:
diff --git a/python/tvm/relax/utils.py b/python/tvm/relax/utils.py
index 12532c1245..a7cc1403a8 100644
--- a/python/tvm/relax/utils.py
+++ b/python/tvm/relax/utils.py
@@ -170,7 +170,7 @@ def gen_call_tir_inputs(
out_ty, and tir_vars.
"""
- tir_var_map: dict[tirx.Var, tirx.Expr] = {}
+ tir_var_map: dict[tvm.ir.Var, tirx.Var] = {}
call_tir_args = []
create_primfunc_args = []
@@ -180,9 +180,8 @@ def gen_call_tir_inputs(
def _copy_undefined_var(expr: tirx.Expr):
def _visit_expr(e: tirx.Expr):
- if isinstance(e, tirx.Var) and e not in tir_var_map:
- new_var = tirx.Var(e.name, e.ty)
- tir_var_map[e] = new_var
+ if isinstance(e, tvm.ir.Var) and e not in tir_var_map:
+ tir_var_map[e] = tvm.ir.Var(e.name, e.ty)
tirx.stmt_functor.post_order_visit(expr, _visit_expr)
@@ -207,9 +206,9 @@ def gen_call_tir_inputs(
te_args : Any
Argument to convert to TE
- tir_var_map : Dict[tirx.Var, tirx.Expr]
- The TIR variable mapping, which maps TIR variables on the Relax function
- side to the new set of variables used on the PrimFunc side.
+ tir_var_map : Dict[tvm.ir.Var, tirx.Var]
+ The variable mapping from caller-side canonical Vars to exact ordinary
+ Vars used on the PrimFunc side.
Returns
-------
@@ -219,8 +218,12 @@ def gen_call_tir_inputs(
"""
def _convert_te_arg_helper(arg):
- if isinstance(arg, tvm.relax.Var) and tvm.ir.is_prim_expr(arg):
- name = arg.name_hint or f"scalar_input_{len(create_primfunc_args)}"
+ if (
+ isinstance(arg, tvm.ir.Var)
+ and tvm.ir.is_prim_expr(arg)
+ and not (tvm.ir.is_prim_var(arg) and arg.ty.dtype == "int64")
+ ):
+ name = arg.name or f"scalar_input_{len(create_primfunc_args)}"
tir_param = tirx.Var(name, arg.ty.dtype)
call_tir_args.append(arg)
create_primfunc_args.append(tir_param)
@@ -242,7 +245,7 @@ def gen_call_tir_inputs(
n_args = len(create_primfunc_args)
if isinstance(arg, tvm.relax.Var):
- name = arg.name_hint
+ name = arg.name
elif n_args < len(string.ascii_uppercase):
name = string.ascii_uppercase[n_args]
else:
@@ -289,7 +292,7 @@ def gen_call_tir_inputs(
if isinstance(expr, te_Tensor):
for dim in expr.shape:
_populate_bound_vars(dim)
- elif isinstance(expr, tirx.Var):
+ elif isinstance(expr, tvm.ir.Var):
bound_vars.add(expr)
def _populate_used_vars(expr):
diff --git a/python/tvm/s_tir/dlight/analysis/common_analysis.py b/python/tvm/s_tir/dlight/analysis/common_analysis.py
index adb541dc48..17061895e3 100644
--- a/python/tvm/s_tir/dlight/analysis/common_analysis.py
+++ b/python/tvm/s_tir/dlight/analysis/common_analysis.py
@@ -104,11 +104,11 @@ class BufferInfo:
for expr in buf.region:
expr = expr.min
dim = None
- if isinstance(expr, tirx.expr.Add) and isinstance(expr.b, tirx.expr.Var):
+ if isinstance(expr, tirx.expr.Add) and ir.is_prim_var(expr.b):
var_add = expr.b
if (
isinstance(expr, tirx.expr.Mul)
- and isinstance(expr.a, tirx.expr.Var)
+ and ir.is_prim_var(expr.a)
and isinstance(expr.b, tirx.expr.IntImm)
):
mul = expr.b
@@ -116,17 +116,17 @@ class BufferInfo:
dim = MergeIndex(var_mul, mul, var_add)
elif (
isinstance(expr, tirx.expr.FloorMod)
- and isinstance(expr.a, tirx.expr.Var)
+ and ir.is_prim_var(expr.a)
and isinstance(expr.b, tirx.expr.IntImm)
):
dim = RemIndex(expr.a, expr.b)
elif (
isinstance(expr, tirx.expr.FloorDiv)
- and isinstance(expr.a, tirx.expr.Var)
+ and ir.is_prim_var(expr.a)
and isinstance(expr.b, tirx.expr.IntImm)
):
dim = DivIndex(expr.a, expr.b)
- elif isinstance(expr, tirx.expr.Var):
+ elif ir.is_prim_var(expr):
dim = Index(expr)
buf_index.append(dim)
return buf_index
@@ -420,7 +420,7 @@ def collect_vars_used_in_prim_expr(expr: tirx.Expr) -> set[tirx.Var]:
tir_vars = set()
def _collect_tir_var(expr):
- if isinstance(expr, tirx.Var):
+ if ir.is_prim_var(expr):
tir_vars.add(expr)
tirx.stmt_functor.post_order_visit(expr, _collect_tir_var)
diff --git a/python/tvm/s_tir/dlight/benchmark/extract.py b/python/tvm/s_tir/dlight/benchmark/extract.py
index bc09e62b4b..a1eb56361e 100644
--- a/python/tvm/s_tir/dlight/benchmark/extract.py
+++ b/python/tvm/s_tir/dlight/benchmark/extract.py
@@ -129,12 +129,12 @@ def extract_dynamic_var(
for arg in flattened_arg_list:
if isinstance(arg, relax.TensorType):
for val in arg.shape.values:
- if isinstance(val, tvm.tirx.Var):
- dym_var_dict[gv][str(val)] = str(val.ty)
+ if tvm.ir.is_prim_var(val):
+ dym_var_dict[gv][val.name] = str(val.ty)
elif isinstance(arg, relax.ShapeType):
for val in arg.values:
- if isinstance(val, tvm.tirx.Var):
- dym_var_dict[gv][str(val)] = str(val.ty)
+ if tvm.ir.is_prim_var(val):
+ dym_var_dict[gv][val.name] = str(val.ty)
else:
raise NotImplementedError
return dym_var_dict
@@ -185,8 +185,8 @@ def extract_func_info_from_prim_func(
for dim in buffer.shape:
if isinstance(dim, tvm.tirx.IntImm):
shape.append(dim.value)
- elif isinstance(dim, tvm.tirx.Var):
- dym_var[str(dim)] = str(dim.ty)
+ elif tvm.ir.is_prim_var(dim):
+ dym_var[dim.name] = str(dim.ty)
shape.append(dim)
else:
raise ValueError(f"Unknown shape: {buffer.shape}")
diff --git a/python/tvm/s_tir/dlight/benchmark/utils.py b/python/tvm/s_tir/dlight/benchmark/utils.py
index 66e50f2934..64449daa76 100644
--- a/python/tvm/s_tir/dlight/benchmark/utils.py
+++ b/python/tvm/s_tir/dlight/benchmark/utils.py
@@ -89,7 +89,9 @@ def populuate_input_shape(
shape = []
if isinstance(input_info, relax.ShapeType):
# scalar input
- results.append(((dym_var_sample[str(input_info.values[0])],), "scalar"))
+ value = input_info.values[0]
+ key = value.name if isinstance(value, tvm.ir.Var) else str(value)
+ results.append(((dym_var_sample[key],), "scalar"))
else:
if isinstance(input_info, relax.TensorType):
tensor_shape = input_info.shape
@@ -102,7 +104,8 @@ def populuate_input_shape(
elif isinstance(dim, tvm.tirx.IntImm):
shape.append(dim.value)
else:
- shape.append(dym_var_sample[str(dim)])
+ key = dim.name if isinstance(dim, tvm.ir.Var) else str(dim)
+ shape.append(dym_var_sample[key])
results.append(((*shape,), _dtype_str(tensor_dtype)))
return results
diff --git a/python/tvm/s_tir/dlight/gpu/general_reduction.py b/python/tvm/s_tir/dlight/gpu/general_reduction.py
index d240573a53..d83d8fe98b 100644
--- a/python/tvm/s_tir/dlight/gpu/general_reduction.py
+++ b/python/tvm/s_tir/dlight/gpu/general_reduction.py
@@ -17,7 +17,7 @@
# pylint: disable=invalid-name
"""Reduction rule for operators including softmax, layer norm, RMS norm, etc"""
-from tvm import arith, s_tir, tirx
+from tvm import arith, ir, s_tir, tirx
from tvm.target import Target
from ..analysis import normalize_prim_func
@@ -135,7 +135,7 @@ class GeneralReduction(GPUScheduleRule):
block_var_to_loop_var[block_iter.var] = sch.get(loop_rv).loop_var
def _visit_expr(e: tirx.Expr):
- if isinstance(e, tirx.Var) and e in block_var_to_loop_var:
+ if ir.is_prim_var(e) and e in block_var_to_loop_var:
spatial_loops.add(block_var_to_loop_var[e])
for buffer_read in spatial_block.reads:
diff --git a/python/tvm/script/ir_builder/ir/__init__.py b/python/tvm/script/ir_builder/ir/__init__.py
index d157aae556..86c6f1bf27 100644
--- a/python/tvm/script/ir_builder/ir/__init__.py
+++ b/python/tvm/script/ir_builder/ir/__init__.py
@@ -28,6 +28,7 @@ from .ir import (
module_global_infos,
lookup_vdevice,
lookup_name,
+ meta_var,
vdevice,
dummy_global_info,
)
diff --git a/python/tvm/script/ir_builder/ir/ir.py b/python/tvm/script/ir_builder/ir/ir.py
index a9987b2f79..a61ad3c344 100644
--- a/python/tvm/script/ir_builder/ir/ir.py
+++ b/python/tvm/script/ir_builder/ir/ir.py
@@ -16,12 +16,43 @@
# under the License.
"""Package tvm.script.ir_builder.ir.ir"""
+from typing import TYPE_CHECKING, Any, TypeVar
+
from tvm.ir import BaseFunc, DummyGlobalInfo, GlobalInfo, GlobalVar, VDevice
from tvm.runtime import Object as tvm_Object
from . import _ffi_api
from .frame import IRModuleFrame
+if TYPE_CHECKING:
+ T = TypeVar("T")
+
+ def meta_var(value: T) -> T:
+ """Mark a value for parser-time metaprogramming."""
+ return value
+
+else:
+
+ class meta_var: # pylint: disable=invalid-name
+ """A value used only for TVMScript parser-time metaprogramming.
+
+ Assignments unwrap this object without emitting an IR binding. The
+ shared wrapper is exposed as ``I.meta_var``; dialect namespaces may
+ provide compatibility aliases to the same implementation.
+ For Relax, this is the explicit opt-out from default primitive binding emission.
+
+ Parameters
+ ----------
+ value : Any
+ The parser-time value.
+ """
+
+ def __init__(self, value: Any) -> None:
+ self.value = value
+
+ def __iter__(self):
+ return (meta_var(item) for item in self.value)
+
def ir_module() -> IRModuleFrame:
"""Start a ir_module frame.
diff --git a/python/tvm/script/parser/core/parser.py b/python/tvm/script/parser/core/parser.py
index f776e66cce..89ccdcddaf 100644
--- a/python/tvm/script/parser/core/parser.py
+++ b/python/tvm/script/parser/core/parser.py
@@ -27,6 +27,7 @@ import numpy as np
from tvm.error import DiagnosticError
from tvm.ir import GlobalVar
+from tvm.runtime import Object
from . import dispatch, doc
from .diagnostics import Diagnostics, Source
@@ -262,9 +263,13 @@ class VarTable:
"""
# Skip if the key and value are equal to those in the var_table
if self.name2value[var] and isinstance(self.name2value[var][-1], type(value)):
- if isinstance(value, np.ndarray) and (self.name2value[var][-1] == value).all():
+ old_value = self.name2value[var][-1]
+ if isinstance(value, np.ndarray) and (old_value == value).all():
return
- elif self.name2value[var][-1] == value:
+ if isinstance(old_value, Object):
+ if old_value.same_as(value):
+ return
+ elif old_value == value:
return
if allow_shadowing and var in self.frames[-1].vars:
# Shadowing
diff --git a/python/tvm/script/parser/ir/__init__.py b/python/tvm/script/parser/ir/__init__.py
index 0d81bd75e2..7553df5b71 100644
--- a/python/tvm/script/parser/ir/__init__.py
+++ b/python/tvm/script/parser/ir/__init__.py
@@ -28,6 +28,7 @@ __all__ = [
"dummy_global_info",
"ir_module",
"lookup_vdevice",
+ "meta_var",
"module_attrs",
"module_global_infos",
"pyfunc",
diff --git a/python/tvm/tirx/expr.py b/python/tvm/tirx/expr.py
index 9b3189f1ca..3e4902c376 100644
--- a/python/tvm/tirx/expr.py
+++ b/python/tvm/tirx/expr.py
@@ -440,28 +440,7 @@ class LogicalExpr(ExprWithOp):
pass
-@tvm_ffi.register_object("tirx.Var")
-class Var(ExprWithOp):
- """Symbolic variable.
-
- Parameters
- ----------
- name : str
- The name
-
- dtype : Union[str, ir.Type]
- The data type
-
- span : Optional[Span]
- The location of this expression in the source code.
- """
-
- name_hint: str
-
- def __init__(self, name: str, dtype: str | ir.Type, span: Span | None = None) -> None:
- if isinstance(dtype, str) and dtype == "handle":
- dtype = ir.PointerType(ir.PrimType("void"))
- self.__init_handle_by_constructor__(_ffi_api.Var, name, dtype, span) # type: ignore
+Var = ir.Var
@tvm_ffi.register_object("tirx.IterVar")
@@ -527,7 +506,7 @@ class IterVar(ExprOp, Object, Scriptable):
name = var if var is not None else "iter"
dtype = "int32" if dom is None else dom.extent.ty
- var = Var(name, dtype=dtype, span=span) if not isinstance(var, Var) else var
+ var = Var(name, ty=dtype, span=span) if not isinstance(var, Var) else var
if dom is not None:
assert var.ty == dom.extent.ty, "IterVar's Var type must match its domain's extent type"
self.__init_handle_by_constructor__(
diff --git a/python/tvm/tirx/function.py b/python/tvm/tirx/function.py
index f7ba3fcd2c..ab7b6435a2 100644
--- a/python/tvm/tirx/function.py
+++ b/python/tvm/tirx/function.py
@@ -73,7 +73,7 @@ class PrimFunc(BaseFunc, Scriptable):
for x in params:
x = tvm.runtime.convert(x) if not isinstance(x, Object) else x
if isinstance(x, Buffer):
- var = Var(x.name, dtype="handle")
+ var = Var(x.name, ty="handle")
param_list.append(var)
buffer_map[var] = x
elif isinstance(x, Var):
@@ -405,11 +405,14 @@ class IndexMap(Object):
final_indices = []
axis_separators = []
- try:
- iter(mapping)
- is_iterable = True
- except TypeError:
+ if tvm.ir.is_prim_expr(mapping):
is_iterable = False
+ else:
+ try:
+ iter(mapping)
+ is_iterable = True
+ except TypeError:
+ is_iterable = False
if is_iterable:
for val in mapping:
diff --git a/python/tvm/tirx/script/builder/ir.py b/python/tvm/tirx/script/builder/ir.py
index 6ce2934d06..4e839565a8 100644
--- a/python/tvm/tirx/script/builder/ir.py
+++ b/python/tvm/tirx/script/builder/ir.py
@@ -39,6 +39,7 @@ from tvm.ir import register_op_attr as _register_op_attr
from tvm.ir.base import deprecated
from tvm.runtime import convert
from tvm.script.ir_builder.base import IRBuilder
+from tvm.script.ir_builder.ir import meta_var
from tvm.target import Target
# pylint: disable=unused-import
@@ -1471,9 +1472,11 @@ def Bind( # pylint: disable=invalid-name
The bound variable.
"""
if type_annotation is not None:
- if callable(type_annotation):
+ # Canonical Vars are callable when they denote functions. Here a Var is
+ # already a resolved type annotation, rather than a deferred annotation factory.
+ if callable(type_annotation) and not isinstance(type_annotation, Expr):
type_annotation = type_annotation()
- if isinstance(type_annotation, Var):
+ if isinstance(type_annotation, ir.Var):
type_annotation = type_annotation.ty
return _ffi_api.Bind(value, type_annotation, var) # type: ignore[attr-defined] # pylint: disable=no-member
@@ -1510,7 +1513,7 @@ class LetAnnotation:
def as_var(self, rhs_dtype=None):
"""Resolve to a tir.Var."""
if self.type_spec is not None:
- if isinstance(self.type_spec, Var):
+ if isinstance(self.type_spec, ir.Var):
return self.type_spec # Already a Var (e.g. T.handle(...))
elif callable(self.type_spec):
return self.type_spec() # e.g. T.int32() -> Var
@@ -2123,7 +2126,7 @@ def _name_meta_value(
IRBuilder.name(prefix, resource)
named_resources.append(resource)
return
- if isinstance(value, Var | IterVar):
+ if isinstance(value, ir.Var | IterVar):
if owned_resources is not None:
return
IRBuilder.name(prefix, value)
@@ -2872,14 +2875,8 @@ def Range(begin: Expr, end: Expr) -> ir.Range: # pylint: disable=invalid-name
if TYPE_CHECKING:
- T = TypeVar("T")
C = TypeVar("C")
- # When type checking (and by extension, for linters like Pylint), treat
- # meta_var as an identity function.
- def meta_var(x: T) -> T:
- return x
-
def meta_class(cls: C) -> C:
return cls
@@ -2924,25 +2921,6 @@ else:
"""
return _install_meta_class(cls)
- class meta_var:
- """A meta variable used in TVMScript metaprogramming.
-
- The value does not appear in the final TIR and only exists in the parser.
-
- Parameters
- ----------
- value: Any
- The meta variable.
- """
-
- def __init__(self, value: Any) -> None:
- self.value = value
-
- def __iter__(self):
- # Return a generator that yields wrapped items.
- return (meta_var(i) for i in self.value)
-
-
# pylint: disable=invalid-name
diff --git a/python/tvm/tirx/script/builder/triton.py b/python/tvm/tirx/script/builder/triton.py
index d0991e5427..95690b7b1c 100644
--- a/python/tvm/tirx/script/builder/triton.py
+++ b/python/tvm/tirx/script/builder/triton.py
@@ -24,7 +24,7 @@ from packaging import version
from triton.runtime.jit import type_canonicalisation_dict
from tvm import tirx
-from tvm.ir import PointerType, PrimType, is_prim_expr
+from tvm.ir import PointerType, PrimType, Var, is_prim_expr
from tvm.runtime import Module
from tvm.topi.utils import get_const_int
@@ -120,7 +120,7 @@ class TritonKernel(BaseKernel):
kernel_args.append(arg)
continue
if isinstance(arg.ty, PointerType):
- assert isinstance(arg, tirx.Var)
+ assert isinstance(arg, Var)
assert isinstance(arg.ty.element_type, PrimType)
elem_type = arg.ty.element_type.dtype
pointer_type = "*" + type_canonicalisation_dict[elem_type]
diff --git a/python/tvm/tirx/script/parser/parser.py b/python/tvm/tirx/script/parser/parser.py
index 766373309f..ed262825e2 100644
--- a/python/tvm/tirx/script/parser/parser.py
+++ b/python/tvm/tirx/script/parser/parser.py
@@ -29,7 +29,7 @@ from tvm.script.ir_builder.base import IRBuilder
from tvm.script.ir_builder.base import IRBuilderFrame as Frame
from tvm.script.parser._core import Parser, dispatch, doc
from tvm.script.parser.core.doc import from_doc
-from tvm.tirx import Buffer, IterVar, Layout, Var
+from tvm.tirx import Buffer, IterVar, Layout
from tvm.tirx.script import builder as T
from tvm.tirx.script.builder.ir import name_meta_class_value
from tvm.tirx.stmt import BufferRegion
@@ -123,7 +123,7 @@ def bind_with_value(self: Parser, node: doc.expr, var_name: str, value: Any) ->
for i, v in enumerate(value):
bind_with_value(self, node, f"{var_name}_{i}", v)
return value
- elif isinstance(value, Buffer | Var):
+ elif isinstance(value, Buffer | tvm.ir.Var):
IRBuilder.name(var_name, value)
return value
else:
@@ -159,7 +159,7 @@ def bind_for_value(self: Parser, node: doc.expr, var_name: str, value: Any) -> A
for i, v in enumerate(value):
bind_for_value(self, node, f"{var_name}_{i}", v)
return value
- elif isinstance(value, Var):
+ elif isinstance(value, tvm.ir.Var):
IRBuilder.name(var_name, value)
return value
else:
@@ -196,7 +196,7 @@ def bind_assign_value(self: Parser, node: doc.expr, var_name: str, value: Any) -
assert isinstance(value.scalar, T.BufferLoad)
IRBuilder.name(var_name, value.scalar.buffer)
return value.scalar
- if isinstance(value, T.meta_var):
+ if isinstance(value, I.meta_var):
return value.value
elif getattr(type(value), "_is_meta_class", False):
name_meta_class_value(var_name, value)
@@ -216,7 +216,7 @@ def bind_assign_value(self: Parser, node: doc.expr, var_name: str, value: Any) -
IRBuilder.name(var_name, res)
return res
elif isinstance(value, Buffer | IterVar | Layout) or (
- isinstance(value, Var) and not self.var_table.exist(value)
+ isinstance(value, tvm.ir.Var) and not self.var_table.exist(value)
):
IRBuilder.name(var_name, value)
return value
@@ -540,13 +540,13 @@ def visit_ann_assign(self: Parser, node: doc.AnnAssign) -> None:
ann_var = raw_ann.as_var()
else:
ann_var = raw_ann.as_var(rhs_dtype=rhs.ty)
- if not isinstance(ann_var, Var):
+ if not isinstance(ann_var, tvm.ir.Var):
self.report_error(node.annotation, "Annotation should resolve to Var")
self.eval_assign(target=lhs, source=ann_var, bind_value=bind_assign_value)
T.Bind(rhs, var=ann_var)
else:
- ann_var = raw_ann() if callable(raw_ann) else raw_ann
- if not isinstance(ann_var, Var):
+ ann_var = raw_ann() if callable(raw_ann) and not isinstance(raw_ann, Expr) else raw_ann
+ if not isinstance(ann_var, tvm.ir.Var):
self.report_error(node.annotation, "Annotation should resolve to Var")
if not isinstance(ann_var.ty, PrimType):
self.report_error(
@@ -618,8 +618,10 @@ def visit_function_def(self: Parser, node: doc.FunctionDef) -> None:
T.func_name(node.name)
if node.returns is not None:
ret_type = self.eval_expr(node.returns)
- if callable(ret_type):
- ret_type = ret_type().ty
+ if callable(ret_type) and not isinstance(ret_type, Expr):
+ ret_type = ret_type()
+ if isinstance(ret_type, Expr):
+ ret_type = ret_type.ty
T.func_ret(ret_type)
with self.with_dispatch_token("tirx"):
# TODO: handle different types of arguments:
@@ -634,7 +636,11 @@ def visit_function_def(self: Parser, node: doc.FunctionDef) -> None:
self.report_error(arg, "Type annotation required for function parameters.")
try:
ann = self.eval_expr(arg.annotation)
- if callable(ann) and ann is not _constexpr_sentinel:
+ if (
+ callable(ann)
+ and not isinstance(ann, Expr)
+ and ann is not _constexpr_sentinel
+ ):
ann = ann()
except Exception: # pylint: disable=broad-except
ann = func_annotation.get(arg.arg, None)
@@ -706,7 +712,7 @@ def visit_tvm_annotation(self: Parser, node: doc.expr):
The doc AST expr node.
"""
annotation = self.eval_expr(node)
- if callable(annotation):
+ if callable(annotation) and not isinstance(annotation, Expr):
annotation = annotation()
return annotation
@@ -735,7 +741,7 @@ def visit_expr_stmt(self: Parser, node: doc.Expr) -> None:
for f in res.frames:
f.add_callback(partial(f.__exit__, None, None, None))
f.__enter__()
- elif isinstance(res, Var):
+ elif isinstance(res, tvm.ir.Var):
# Standalone Var expression (e.g. from T.bind(value, var=v)) --
# the Bind statement was already emitted to the parent frame by the FFI call,
# so just discard the returned Var.
@@ -891,8 +897,10 @@ def visit_tvm_declare_function(self: Parser, node: doc.FunctionDef) -> GlobalVar
with self.var_table.with_frame():
if node.returns is not None:
ret_type = self.eval_expr(node.returns)
- if callable(ret_type):
- ret_type = ret_type().ty
+ if callable(ret_type) and not isinstance(ret_type, Expr):
+ ret_type = ret_type()
+ if isinstance(ret_type, Expr):
+ ret_type = ret_type.ty
arg_annotations = []
for arg in node.args.args:
@@ -900,7 +908,7 @@ def visit_tvm_declare_function(self: Parser, node: doc.FunctionDef) -> GlobalVar
self.report_error(arg, "Type annotation required for function parameters.")
try:
ann = self.eval_expr(arg.annotation)
- if callable(ann):
+ if callable(ann) and not isinstance(ann, Expr):
ann = ann()
except Exception: # pylint: disable=broad-except
ann = func_annotation.get(arg.arg, None)
diff --git a/python/tvm/tirx/transform/transform.py b/python/tvm/tirx/transform/transform.py
index 98a16fd79d..7eb26b7225 100644
--- a/python/tvm/tirx/transform/transform.py
+++ b/python/tvm/tirx/transform/transform.py
@@ -264,9 +264,9 @@ def MakePackedAPI():
variables use `tirx.Var` that are not defined by other PrimFunc
parameters, these are instead used to define the variables based on
the corresponding `DLTensor` members. (e.g. A PrimFunc that accepts a
- buffer of shape `[tirx.Var("n"), tirx.Var("m")]`, when passed a
- `DLTensor` of shape `[16,32]`, will define `n = 16` and `n=32`, based
- on the argument's shape.
+ buffer of shape `[tirx.Var("n", "int64"), tirx.Var("m", "int64")]`,
+ when passed a `DLTensor` of shape `[16, 32]`, will define `n = 16` and
+ `m = 32`, based on the argument's shape.
Returns
-------
diff --git a/python/tvm/topi/nn/batch_matmul.py b/python/tvm/topi/nn/batch_matmul.py
index 713942affb..d677a5d03b 100644
--- a/python/tvm/topi/nn/batch_matmul.py
+++ b/python/tvm/topi/nn/batch_matmul.py
@@ -91,13 +91,13 @@ def batch_matmul(
else:
YB, YK, YJ = get_const_tuple(tensor_b.shape)
- assert XK == YK or isinstance(YK, tvm.tirx.expr.Var), "shapes of x and y are inconsistent"
+ assert XK == YK or tvm.ir.is_prim_var(YK), "shapes of x and y are inconsistent"
k = te.reduce_axis((0, XK), name="k")
if oshape is None:
assert XB == YB or XB == 1 or YB == 1, "batch dimension doesn't match"
batch = (
tvm.tirx.expr.Var("batch", "int32")
- if isinstance(XB, tvm.tirx.expr.Var) or isinstance(YB, tvm.tirx.expr.Var)
+ if tvm.ir.is_prim_var(XB) or tvm.ir.is_prim_var(YB)
else te.max(XB, YB)
)
oshape = (batch, XI, YJ)
diff --git a/python/tvm/topi/nn/dense.py b/python/tvm/topi/nn/dense.py
index ceac8d29e1..c0b4350285 100644
--- a/python/tvm/topi/nn/dense.py
+++ b/python/tvm/topi/nn/dense.py
@@ -95,7 +95,7 @@ def matmul(
reduce_dim_b, out_dim = tensor_b.shape[-2:]
batch_dims_b = tensor_b.shape[:-2]
- if not isinstance(reduce_dim_a, tvm.tirx.Var) and not isinstance(reduce_dim_b, tvm.tirx.Var):
+ if not tvm.ir.is_prim_var(reduce_dim_a) and not tvm.ir.is_prim_var(reduce_dim_b):
assert int(reduce_dim_a) == int(reduce_dim_b), (
f"Reduction dimensions of dense do not match. {reduce_dim_a} vs {reduce_dim_b}."
)
@@ -105,17 +105,12 @@ def matmul(
batch_dims_b = [1] * (result_ndim - len(batch_dims_b)) + batch_dims_b
for idx, (l, r) in enumerate(zip(batch_dims_a, batch_dims_b)):
- if (
- not isinstance(l, tvm.tirx.Var)
- and not isinstance(r, tvm.tirx.Var)
- and int(l) != 1
- and int(r) != 1
- ):
+ if not tvm.ir.is_prim_var(l) and not tvm.ir.is_prim_var(r) and int(l) != 1 and int(r) != 1:
assert int(l) == int(r), (
"Batch dimensions of dense do not match: "
f"{tensor_a.shape[:-2]} vs {tensor_b.shape[:-2]}."
)
- if not isinstance(l, tvm.tirx.Var) and int(l) == 1:
+ if not tvm.ir.is_prim_var(l) and int(l) == 1:
batch_dims_a[idx] = batch_dims_b[idx]
k = te.reduce_axis((0, reduce_dim_a), name="k")
@@ -123,12 +118,12 @@ def matmul(
def compute(*indices):
batch_indices_a = indices[-len(tensor_a.shape) : -2]
batch_indices_a = [
- i if isinstance(dim, tvm.tirx.Var) or int(dim) != 1 else 0
+ i if tvm.ir.is_prim_var(dim) or int(dim) != 1 else 0
for i, dim in zip(batch_indices_a, tensor_a.shape[:-2])
]
batch_indices_b = indices[-len(tensor_b.shape) : -2]
batch_indices_b = [
- i if isinstance(dim, tvm.tirx.Var) or int(dim) != 1 else 0
+ i if tvm.ir.is_prim_var(dim) or int(dim) != 1 else 0
for i, dim in zip(batch_indices_b, tensor_b.shape[:-2])
]
i, j = indices[-2:]
diff --git a/python/tvm/topi/scatter.py b/python/tvm/topi/scatter.py
index de35577c4d..453cfbedf4 100644
--- a/python/tvm/topi/scatter.py
+++ b/python/tvm/topi/scatter.py
@@ -18,11 +18,10 @@
# ruff: noqa: E741
"""ScatterND operator"""
-from tvm import DataTypeCode, te, tirx # hide redefinition of min and max
+from tvm import DataTypeCode, ir, te, tirx # hide redefinition of min and max
from tvm.arith.analyzer import Analyzer
from tvm.script.ir_builder import IRBuilder
from tvm.script.ir_builder import tirx as T
-from tvm.tirx import expr
def _verify_scatter_nd_inputs(data, indices, updates):
@@ -33,7 +32,7 @@ def _verify_scatter_nd_inputs(data, indices, updates):
f"the length of the shape of the output ({len(data.shape)})."
)
for i in range(len(indices.shape) - 1):
- if isinstance(indices.shape[i + 1], expr.Var) or isinstance(updates.shape[i], expr.Var):
+ if ir.is_prim_var(indices.shape[i + 1]) or ir.is_prim_var(updates.shape[i]):
continue
assert analyzer.can_prove_equal(indices.shape[i + 1], updates.shape[i]), (
@@ -42,7 +41,7 @@ def _verify_scatter_nd_inputs(data, indices, updates):
)
for i in range(mdim, len(data.shape)):
data_ind = i - mdim + len(indices.shape) - 1
- if isinstance(updates.shape[data_ind], expr.Var) or isinstance(data.shape[i], expr.Var):
+ if ir.is_prim_var(updates.shape[data_ind]) or ir.is_prim_var(data.shape[i]):
continue
assert updates.shape[data_ind] == data.shape[i], (
f"Dimension of updates[{data_ind}] ({updates.shape[data_ind]}) must equal dimension "
diff --git a/python/tvm/topi/signal.py b/python/tvm/topi/signal.py
index a3d1f63409..c3c38854b7 100644
--- a/python/tvm/topi/signal.py
+++ b/python/tvm/topi/signal.py
@@ -19,7 +19,7 @@
from math import pi
-from tvm import te, tirx
+from tvm import ir, te, tirx
from tvm.script.ir_builder import IRBuilder
from tvm.script.ir_builder import tirx as T
@@ -112,7 +112,7 @@ def stft(
output_buf = tirx.decl_buffer(output_shape, data.dtype, "output_buf", layout=None)
loop_kind = "vectorize"
- if isinstance(output_shape[2], tirx.expr.Var): # any_dim
+ if ir.is_prim_var(output_shape[2]): # any_dim
loop_kind = "serial"
return te.extern(
diff --git a/python/tvm/topi/utils.py b/python/tvm/topi/utils.py
index fb02d40d61..b05bb8c2bb 100644
--- a/python/tvm/topi/utils.py
+++ b/python/tvm/topi/utils.py
@@ -25,7 +25,6 @@ import numpy as np
import tvm
from tvm import te
from tvm.s_tir import sbijective_layout, slayout
-from tvm.tirx import Var
from . import cpp, tag
@@ -194,7 +193,7 @@ def get_const_tuple(in_tuple):
ret = []
ana = None
for elem in in_tuple:
- if isinstance(elem, tvm.tirx.Var):
+ if tvm.ir.is_prim_var(elem):
ret.append(elem)
elif not isinstance(elem, tvm.tirx.IntImm | int):
ana = tvm.arith.Analyzer() if ana is None else ana
@@ -542,4 +541,4 @@ def is_target(names):
def is_dynamic_shape(shape):
"""Checks if any part of a shape is dynamic"""
- return any(isinstance(x, Var) for x in shape)
+ return any(tvm.ir.is_prim_var(x) for x in shape)
diff --git a/src/arith/analyzer.cc b/src/arith/analyzer.cc
index 92ce8952aa..0332d23bdd 100644
--- a/src/arith/analyzer.cc
+++ b/src/arith/analyzer.cc
@@ -107,8 +107,8 @@ void AnalyzerObj::MarkGlobalNonNegValue(const PrimExpr& value) {
//
// We may consider enhance the sub analyzer to directly take
// MarkPositiveVar so their bounds do not overlap
- if (const auto* var_ptr = symbol.as<VarNode>()) {
- Var var = ffi::GetRef<Var>(var_ptr);
+ if (auto prim_var = symbol.as<tirx::PrimVar>()) {
+ Var var = *prim_var;
// skip non-index type, keep it to be compatible
// with any_dim that do not represent any value
if (!IsIndexTypedExpr(var.as_or_throw<PrimExpr>())) return;
diff --git a/src/arith/detect_linear_equation.cc b/src/arith/detect_linear_equation.cc
index 79afbdf5fe..312cc98c23 100644
--- a/src/arith/detect_linear_equation.cc
+++ b/src/arith/detect_linear_equation.cc
@@ -48,7 +48,7 @@ struct IntervalEntry {
class LinearEqDetector : public ExprFunctor<LinearEqEntry(const Expr&, const PrimExpr&)> {
public:
- explicit LinearEqDetector(Var var) : var_(var) {}
+ explicit LinearEqDetector(PrimVar var) : var_(var) {}
bool Detect(const PrimExpr& e, LinearEqEntry* ret) {
*ret = VisitExpr(e, e);
@@ -121,7 +121,7 @@ class LinearEqDetector : public ExprFunctor<LinearEqEntry(const Expr&, const Pri
}
private:
- Var var_;
+ PrimVar var_;
bool fail_{false};
// Combine by add
PrimExpr AddCombine(PrimExpr a, PrimExpr b) {
@@ -142,11 +142,11 @@ class LinearEqDetector : public ExprFunctor<LinearEqEntry(const Expr&, const Pri
}
};
-ffi::Array<PrimExpr> DetectLinearEquation(const PrimExpr& e, const ffi::Array<Var>& vars) {
+ffi::Array<PrimExpr> DetectLinearEquation(const PrimExpr& e, const ffi::Array<PrimVar>& vars) {
PrimExpr base = e;
ffi::Array<PrimExpr> coeff;
- for (Var v : vars) {
+ for (PrimVar v : vars) {
LinearEqEntry ret;
if (!LinearEqDetector(v).Detect(base, &ret)) {
return ffi::Array<PrimExpr>();
@@ -173,12 +173,13 @@ ffi::Array<PrimExpr> DetectLinearEquation(const PrimExpr& e, const ffi::Array<Va
bool DetectClipBound(const PrimExpr& cond,
std::unordered_map<const VarNode*, IntervalEntry>* bmap) {
int flag = 0;
- Var var;
+ PrimVar var;
auto fvisit = [&bmap, &flag, &var](const ffi::ObjectRef& n) {
- if (const VarNode* v = n.as<VarNode>()) {
+ if (auto prim_var = n.as<PrimVar>()) {
+ const VarNode* v = prim_var->get();
if (bmap->count(v)) {
if (flag == 0) {
- var = n.as_or_throw<Var>();
+ var = *prim_var;
flag = 1;
} else if (flag == 1) {
if (!var.same_as(n)) {
@@ -267,19 +268,19 @@ void SplitCommExpr(const PrimExpr& e, std::vector<PrimExpr>* ret) {
// Detect the lower and upper bound from the expression.
// e must be connected by and.
-ffi::Array<PrimExpr> DetectClipBound(const PrimExpr& e, const ffi::Array<Var>& vars) {
+ffi::Array<PrimExpr> DetectClipBound(const PrimExpr& e, const ffi::Array<PrimVar>& vars) {
std::vector<PrimExpr> splits;
Analyzer analyzer;
SplitCommExpr<tirx::AndNode>(analyzer->Simplify(e), &splits);
std::unordered_map<const VarNode*, IntervalEntry> rmap;
- for (Var v : vars) {
+ for (PrimVar v : vars) {
rmap[v.get()] = IntervalEntry();
}
for (PrimExpr cond : splits) {
if (!DetectClipBound(cond, &rmap)) return ffi::Array<PrimExpr>();
}
ffi::Array<PrimExpr> ret;
- for (Var v : vars) {
+ for (PrimVar v : vars) {
IntervalEntry e = rmap[v.get()];
if (e.min_value.defined()) {
e.min_value = analyzer->Simplify(e.min_value);
@@ -297,8 +298,9 @@ TVM_FFI_STATIC_INIT_BLOCK() {
namespace refl = tvm::ffi::reflection;
refl::GlobalDef()
.def("arith.DetectLinearEquation", DetectLinearEquation)
- .def("arith.DetectClipBound",
- [](const PrimExpr& e, const ffi::Array<Var>& vars) { return DetectClipBound(e, vars); });
+ .def("arith.DetectClipBound", [](const PrimExpr& e, const ffi::Array<PrimVar>& vars) {
+ return DetectClipBound(e, vars);
+ });
}
} // namespace arith
} // namespace tvm
diff --git a/src/arith/int_constraints.cc b/src/arith/int_constraints.cc
index 442299704d..2d8490bc5e 100644
--- a/src/arith/int_constraints.cc
+++ b/src/arith/int_constraints.cc
@@ -45,7 +45,7 @@ TVM_FFI_STATIC_INIT_BLOCK() {
IntConstraintsTransformNode::RegisterReflection();
}
-ffi::Array<PrimExpr> AsConditions(const ffi::Array<Var>& variables,
+ffi::Array<PrimExpr> AsConditions(const ffi::Array<PrimVar>& variables,
const ffi::Map<Var, IntGroupBounds>& bounds,
const ffi::Array<PrimExpr>& relations) {
ffi::Array<PrimExpr> res;
@@ -222,19 +222,18 @@ TVM_FFI_STATIC_INIT_BLOCK() {
// Pattern A (RM): auto-default repr from reflection.
-IntConstraints::IntConstraints(ffi::Array<Var> variables, ffi::Map<Var, Range> ranges,
+IntConstraints::IntConstraints(ffi::Array<PrimVar> variables, ffi::Map<Var, Range> ranges,
ffi::Array<PrimExpr> relations) {
ffi::ObjectPtr<IntConstraintsNode> node = ffi::make_object<IntConstraintsNode>();
if (!variables.defined()) {
- variables = ffi::Array<Var>();
+ variables = ffi::Array<PrimVar>();
}
if (!ranges.defined()) {
ranges = ffi::Map<Var, Range>();
}
TVM_FFI_ICHECK(relations.defined());
- for (const auto& var : variables) {
- PrimType var_ty = var->ty.as_or_throw<PrimType>();
- TVM_FFI_ICHECK(var_ty.MatchesCode(DLDataTypeCode::kDLInt, DLDataTypeCode::kDLUInt))
+ for (const PrimVar& var : variables) {
+ TVM_FFI_CHECK(var.ty().MatchesCode(DLDataTypeCode::kDLInt, DLDataTypeCode::kDLUInt), TypeError)
<< "Variables in IntConstraints must be integers";
}
node->variables = std::move(variables);
@@ -247,9 +246,8 @@ TVM_FFI_STATIC_INIT_BLOCK() {
namespace refl = tvm::ffi::reflection;
refl::GlobalDef().def(
"arith.IntConstraints",
- [](ffi::Array<Var> variables, ffi::Map<Var, Range> ranges, ffi::Array<PrimExpr> relations) {
- return IntConstraints(variables, ranges, relations);
- });
+ [](ffi::Array<PrimVar> variables, ffi::Map<Var, Range> ranges,
+ ffi::Array<PrimExpr> relations) { return IntConstraints(variables, ranges, relations); });
}
// Pattern A (RM): auto-default repr from reflection.
@@ -275,13 +273,14 @@ IntConstraintsTransform IntConstraintsTransform::operator+(
Analyzer ana_first;
ana_first->Bind(operator->()->src->ranges);
for (auto p : other->dst_to_src) {
- dst_to_src.Set(p.first, ana_first->Simplify(Substitute(p.second, operator->()->dst_to_src)));
+ dst_to_src.Set(p.first,
+ ana_first->Simplify(tirx::Substitute(p.second, operator->()->dst_to_src)));
}
Analyzer ana_second;
ana_second->Bind(other->dst->ranges);
for (auto p : operator->()->src_to_dst) {
- src_to_dst.Set(p.first, ana_second->Simplify(Substitute(p.second, other->src_to_dst)));
+ src_to_dst.Set(p.first, ana_second->Simplify(tirx::Substitute(p.second, other->src_to_dst)));
}
return IntConstraintsTransform(operator->()->src, other->dst, src_to_dst, dst_to_src);
}
diff --git a/src/arith/int_set.cc b/src/arith/int_set.cc
index 1e6f7497ea..45a561610e 100644
--- a/src/arith/int_set.cc
+++ b/src/arith/int_set.cc
@@ -1191,6 +1191,10 @@ ffi::Optional<ffi::Array<IntSet>> EstimateRegionStrictBound(const ffi::Array<Ran
const ffi::Map<Var, Range>& var_dom,
const PrimExpr& predicate,
const Analyzer& analyzer) {
+ ffi::Map<PrimVar, Range> input_iters;
+ for (const auto& [var, range] : var_dom) {
+ input_iters.Set(var.as_or_throw<PrimVar>(), range);
+ }
AnalyzerObj* analyzer_ptr = analyzer.get();
int ndim = region.size();
ffi::Array<IterSumExpr> iter_sum_exprs{nullptr};
@@ -1205,7 +1209,7 @@ ffi::Optional<ffi::Array<IntSet>> EstimateRegionStrictBound(const ffi::Array<Ran
affine_indices.push_back(range->min);
}
auto res = DetectIterMap(
- /*indices=*/affine_indices, /*input_iters=*/var_dom,
+ /*indices=*/affine_indices, /*input_iters=*/input_iters,
/*predicate=*/predicate, /*check_level=*/IterMapLevel::Surjective, analyzer);
iter_sum_exprs = res->indices;
}
@@ -1247,10 +1251,14 @@ ffi::Array<IntSet> EstimateRegionUpperBound(const ffi::Array<Range>& region,
}
ffi::Array<IntSet> result;
result.reserve(region.size());
+ ffi::Map<PrimVar, Range> input_iters;
+ for (const auto& [var, range] : var_dom) {
+ input_iters.Set(var.as_or_throw<PrimVar>(), range);
+ }
// try estimate each dimension independently
for (const Range& range : region) {
auto res = DetectIterMap(
- /*indices=*/{range->min}, /*input_iters=*/var_dom,
+ /*indices=*/{range->min}, /*input_iters=*/input_iters,
/*predicate=*/predicate, /*check_level=*/IterMapLevel::Surjective, analyzer);
if (!res->indices.empty()) {
TVM_FFI_ICHECK_EQ(res->indices.size(), 1U);
diff --git a/src/arith/ir_mutator_with_analyzer.cc b/src/arith/ir_mutator_with_analyzer.cc
index 98576efa2c..65f20e0429 100644
--- a/src/arith/ir_mutator_with_analyzer.cc
+++ b/src/arith/ir_mutator_with_analyzer.cc
@@ -174,7 +174,7 @@ ffi::Array<PrimExpr> IRMutatorWithAnalyzer::IterMapSimplifyWithContext(
indices, this->iter_vars_, pred, arith::IterMapLevel::Surjective, analyzer_ref);
if (non_trivial_only) {
for (int i = 0; i < n; ++i) {
- if (simplified[i]->IsInstance<IntImmNode>() && indices[i]->IsInstance<VarNode>()) {
+ if (simplified[i]->IsInstance<IntImmNode>() && indices[i].as<PrimVar>()) {
simplified.Set(i, indices[i]);
}
}
diff --git a/src/arith/ir_mutator_with_analyzer.h b/src/arith/ir_mutator_with_analyzer.h
index 8998fafd0f..47b52da78c 100644
--- a/src/arith/ir_mutator_with_analyzer.h
+++ b/src/arith/ir_mutator_with_analyzer.h
@@ -31,6 +31,7 @@
#include <tvm/tirx/analysis.h>
#include <tvm/tirx/stmt_functor.h>
+#include <unordered_set>
#include <utility>
namespace tvm {
@@ -91,7 +92,7 @@ class IRMutatorWithAnalyzer : public tirx::StmtExprMutator {
// expensive and we only encourage doing them during
// necessary cases like layout remapping
/*! \brief Recorded loop iterators */
- ffi::Map<Var, Range> iter_vars_;
+ ffi::Map<tirx::PrimVar, Range> iter_vars_;
/*! \brief iterator predicates */
ffi::Array<PrimExpr> iter_predicates_;
/*!
@@ -101,8 +102,12 @@ class IRMutatorWithAnalyzer : public tirx::StmtExprMutator {
*/
template <typename FLambda>
void WithRecordIterPredicate(PrimExpr condition, FLambda callback) {
- auto f_use_itervar = [this](const tirx::VarNode* v) {
- return iter_vars_.count(ffi::GetRef<tirx::Var>(v));
+ std::unordered_set<const tirx::VarNode*> iter_var_nodes;
+ for (const auto& [var, _] : iter_vars_) {
+ iter_var_nodes.insert(var.get());
+ }
+ auto f_use_itervar = [&iter_var_nodes](const tirx::VarNode* v) {
+ return iter_var_nodes.count(v);
};
// simple heuristics for detecting predicate
if (tirx::UsesVar(condition, f_use_itervar)) {
diff --git a/src/arith/iter_affine_map.cc b/src/arith/iter_affine_map.cc
index b1718f5988..c76729c4ef 100644
--- a/src/arith/iter_affine_map.cc
+++ b/src/arith/iter_affine_map.cc
@@ -174,7 +174,7 @@ class IterMapRewriter : public ExprMutator {
public:
using Parent = ExprMutator;
- explicit IterMapRewriter(AnalyzerObj* analyzer, const ffi::Map<Var, Range>& input_iters,
+ explicit IterMapRewriter(AnalyzerObj* analyzer, const ffi::Map<PrimVar, Range>& input_iters,
IterMapLevel check_level, bool simplify_trivial_iterators,
ffi::Array<ffi::String>* errors)
: analyzer_(analyzer),
@@ -182,7 +182,7 @@ class IterMapRewriter : public ExprMutator {
errors_(*errors),
padding_predicate_(IntImm::Bool(false)) {
for (auto kv : input_iters) {
- const Var& var = kv.first;
+ const PrimVar& var = kv.first;
const Range& vrng = kv.second;
if (simplify_trivial_iterators && is_one(vrng->extent)) {
var_map_[var] = IterSumExpr({}, vrng->min);
@@ -1299,9 +1299,13 @@ struct IterConstraint {
* \param result The result of predicate split.
* \return A list of IterConstraint, empty if the split failed.
*/
-bool MatchBoundConstraints(PrimExpr pred, ffi::Map<Var, Range>* input_iters,
+bool MatchBoundConstraints(PrimExpr pred, ffi::Map<PrimVar, Range>* input_iters,
std::vector<IterConstraint>* result) {
arith::PVar<PrimExpr> lhs, rhs, rest;
+ std::unordered_set<const VarNode*> input_iter_nodes;
+ for (const auto& [var, _] : *input_iters) {
+ input_iter_nodes.insert(var.get());
+ }
for (;;) {
// try extract comparisions
bool is_finish = false;
@@ -1342,8 +1346,8 @@ bool MatchBoundConstraints(PrimExpr pred, ffi::Map<Var, Range>* input_iters,
}
// determine iter and bound, if we can not distinguish them simply,
// try divide (lhs - rhs) into itervar aware and itervar free parts
- auto f_use_itervar = [&input_iters](const VarNode* v) {
- return input_iters->count(ffi::GetRef<Var>(v));
+ auto f_use_itervar = [&input_iter_nodes](const VarNode* v) {
+ return input_iter_nodes.count(v);
};
bool bound_at_left;
if (UsesVar(lhs_expr, f_use_itervar) || UsesVar(rhs_expr, f_use_itervar)) {
@@ -1400,7 +1404,7 @@ bool MatchBoundConstraints(PrimExpr pred, ffi::Map<Var, Range>* input_iters,
}
}
// If it is a predicate for a single input iter
- if (auto opt = iter.as<Var>()) {
+ if (auto opt = iter.as<PrimVar>()) {
auto var = opt.value();
auto it = input_iters->find(var);
if (it != input_iters->end()) {
@@ -1422,7 +1426,7 @@ bool MatchBoundConstraints(PrimExpr pred, ffi::Map<Var, Range>* input_iters,
return true;
}
-bool IterRangeSanityCheck(const ffi::Map<Var, Range>& iter_ranges) {
+bool IterRangeSanityCheck(const ffi::Map<PrimVar, Range>& iter_ranges) {
std::unordered_set<Var> iters;
for (const auto& it : iter_ranges) iters.insert(it.first);
auto f = [&](const VarNode* var) { return iters.count(ffi::GetRef<Var>(var)); };
@@ -1433,7 +1437,7 @@ bool IterRangeSanityCheck(const ffi::Map<Var, Range>& iter_ranges) {
}
IterMapResult DetectIterMap(const ffi::Array<PrimExpr>& indices,
- const ffi::Map<Var, Range>& input_iters, const PrimExpr& predicate,
+ const ffi::Map<PrimVar, Range>& input_iters, const PrimExpr& predicate,
IterMapLevel check_level, const arith::Analyzer& analyzer,
bool simplify_trivial_iterators) {
arith::AnalyzerObj* analyzer_ptr = analyzer.get();
@@ -1446,7 +1450,7 @@ IterMapResult DetectIterMap(const ffi::Array<PrimExpr>& indices,
result->errors.push_back("Invalid iterators. Iterators may not be expressions of each other.");
return result;
}
- ffi::Map<Var, Range> constrained_input_iters = input_iters;
+ ffi::Map<PrimVar, Range> constrained_input_iters = input_iters;
std::vector<IterConstraint> constraints;
if (!is_one(predicate) &&
!MatchBoundConstraints(predicate, &constrained_input_iters, &constraints)) {
@@ -1523,7 +1527,7 @@ TVM_FFI_STATIC_INIT_BLOCK() {
namespace refl = tvm::ffi::reflection;
refl::GlobalDef().def(
"arith.DetectIterMap",
- [](const ffi::Array<PrimExpr>& indices, const ffi::Map<Var, Range>& input_iters,
+ [](const ffi::Array<PrimExpr>& indices, const ffi::Map<PrimVar, Range>& input_iters,
const PrimExpr& input_pred, int check_level, bool simplify_trivial_iterators,
ffi::Optional<Analyzer> opt_analyzer) {
Analyzer ana = opt_analyzer.has_value() ? opt_analyzer.value() : Analyzer();
@@ -1532,7 +1536,7 @@ TVM_FFI_STATIC_INIT_BLOCK() {
});
}
-IterSumExpr NormalizeToIterSum(PrimExpr index, const ffi::Map<Var, Range>& input_iters,
+IterSumExpr NormalizeToIterSum(PrimExpr index, const ffi::Map<PrimVar, Range>& input_iters,
const arith::Analyzer& analyzer) {
arith::AnalyzerObj* analyzer_ptr = analyzer.get();
IterMapResult result;
@@ -1552,7 +1556,7 @@ IterSumExpr NormalizeToIterSum(PrimExpr index, const ffi::Map<Var, Range>& input
TVM_FFI_STATIC_INIT_BLOCK() {
namespace refl = tvm::ffi::reflection;
refl::GlobalDef().def(
- "arith.NormalizeToIterSum", [](PrimExpr index, const ffi::Map<Var, Range>& input_iters,
+ "arith.NormalizeToIterSum", [](PrimExpr index, const ffi::Map<PrimVar, Range>& input_iters,
ffi::Optional<Analyzer> opt_analyzer) {
Analyzer ana = opt_analyzer.has_value() ? opt_analyzer.value() : Analyzer();
return NormalizeToIterSum(index, input_iters, ana);
@@ -2159,7 +2163,7 @@ TVM_FFI_STATIC_INIT_BLOCK() {
}
ffi::Array<PrimExpr> IterMapSimplify(const ffi::Array<PrimExpr>& indices,
- const ffi::Map<Var, Range>& input_iters,
+ const ffi::Map<PrimVar, Range>& input_iters,
const PrimExpr& input_pred, IterMapLevel check_level,
const arith::Analyzer& ana, bool simplify_trivial_iterators) {
arith::AnalyzerObj* ana_ptr = ana.get();
@@ -2191,7 +2195,7 @@ TVM_FFI_STATIC_INIT_BLOCK() {
namespace refl = tvm::ffi::reflection;
refl::GlobalDef().def(
"arith.IterMapSimplify",
- [](const ffi::Array<PrimExpr>& indices, const ffi::Map<Var, Range>& input_iters,
+ [](const ffi::Array<PrimExpr>& indices, const ffi::Map<PrimVar, Range>& input_iters,
const PrimExpr& input_pred, int check_level, bool simplify_trivial_iterators,
ffi::Optional<Analyzer> opt_analyzer) {
Analyzer ana = opt_analyzer.has_value() ? opt_analyzer.value() : Analyzer();
@@ -2495,8 +2499,8 @@ class SubspaceDivider {
};
ffi::Array<ffi::Array<IterMark>> SubspaceDivide(const ffi::Array<PrimExpr>& bindings,
- const ffi::Map<Var, Range>& input_iters,
- const ffi::Array<Var>& sub_iters,
+ const ffi::Map<PrimVar, Range>& input_iters,
+ const ffi::Array<PrimVar>& sub_iters,
const PrimExpr& predicate, IterMapLevel check_level,
const arith::Analyzer& analyzer,
bool simplify_trivial_iterators) {
@@ -2508,7 +2512,7 @@ ffi::Array<ffi::Array<IterMark>> SubspaceDivide(const ffi::Array<PrimExpr>& bind
if (maps.empty()) return {};
std::unordered_set<Var> inner_iter_set;
- for (const Var& inner_iter : sub_iters) {
+ for (const PrimVar& inner_iter : sub_iters) {
inner_iter_set.insert(inner_iter);
}
@@ -2533,8 +2537,8 @@ TVM_FFI_STATIC_INIT_BLOCK() {
namespace refl = tvm::ffi::reflection;
refl::GlobalDef().def(
"arith.SubspaceDivide",
- [](const ffi::Array<PrimExpr>& bindings, const ffi::Map<Var, Range>& root_iters,
- const ffi::Array<Var>& sub_iters, const PrimExpr& predicate, int check_level,
+ [](const ffi::Array<PrimExpr>& bindings, const ffi::Map<PrimVar, Range>& root_iters,
+ const ffi::Array<PrimVar>& sub_iters, const PrimExpr& predicate, int check_level,
bool simplify_trivial_iterators, ffi::Optional<Analyzer> opt_analyzer) {
Analyzer ana = opt_analyzer.has_value() ? opt_analyzer.value() : Analyzer();
return SubspaceDivide(bindings, root_iters, sub_iters, predicate, IterMapLevel(check_level),
diff --git a/src/arith/presburger_set.cc b/src/arith/presburger_set.cc
index ba6b5564d9..5c065e9430 100644
--- a/src/arith/presburger_set.cc
+++ b/src/arith/presburger_set.cc
@@ -93,11 +93,12 @@ static void Update(const PrimExpr& constraint, PresburgerSetNode* intset) {
}
PresburgerSet::PresburgerSet(const PrimExpr& constraint) {
- ffi::Array<Var> vars;
+ ffi::Array<PrimVar> vars;
PostOrderVisit(constraint, [&vars](const ffi::ObjectRef& obj) {
- if (const VarNode* new_var = obj.as<VarNode>()) {
- auto var = ffi::GetRef<Var>(new_var);
- if (!std::any_of(vars.begin(), vars.end(), [&var](const Var& v) { return v.same_as(var); })) {
+ if (auto prim_var = obj.as<PrimVar>()) {
+ PrimVar var = *prim_var;
+ if (!std::any_of(vars.begin(), vars.end(),
+ [&var](const PrimVar& v) { return v.same_as(var); })) {
vars.push_back(var);
}
}
@@ -113,12 +114,13 @@ PresburgerSet::PresburgerSet(const PrimExpr& constraint) {
}
PresburgerSet::PresburgerSet(const std::vector<IntegerRelation>& disjuncts,
- const ffi::Array<Var>& vars) {
+ const ffi::Array<PrimVar>& vars) {
auto node = ffi::make_object<PresburgerSetNode>(disjuncts, disjuncts[0].getSpace(), vars);
data_ = std::move(node);
}
-void PresburgerSetNode::UpdateConstraint(const PrimExpr& constraint, const ffi::Array<Var>& vars) {
+void PresburgerSetNode::UpdateConstraint(const PrimExpr& constraint,
+ const ffi::Array<PrimVar>& vars) {
Analyzer analyzer;
PrimExpr simplified_constraint = analyzer->Simplify(constraint, kSimplifyRewriteCanonicalRewrite);
Update(simplified_constraint, this);
diff --git a/src/arith/presburger_set.h b/src/arith/presburger_set.h
index 08205d6f16..b56d9c09f2 100644
--- a/src/arith/presburger_set.h
+++ b/src/arith/presburger_set.h
@@ -61,10 +61,10 @@ using namespace presburger;
class PresburgerSetNode : public IntSetNode {
public:
PresburgerSetNode() : space(PresburgerSpace::getRelationSpace()) {}
- explicit PresburgerSetNode(const PresburgerSpace& space, const ffi::Array<Var>& vars)
+ explicit PresburgerSetNode(const PresburgerSpace& space, const ffi::Array<tirx::PrimVar>& vars)
: disjuncts({}), space(space), vars(vars) {}
explicit PresburgerSetNode(const std::vector<IntegerRelation>& disjuncts,
- const PresburgerSpace& space, const ffi::Array<Var>& vars)
+ const PresburgerSpace& space, const ffi::Array<tirx::PrimVar>& vars)
: disjuncts(disjuncts), space(space), vars(vars) {}
/*! \brief Represent the union of multiple IntegerRelation */
@@ -92,7 +92,7 @@ class PresburgerSetNode : public IntSetNode {
* \param constraint The added constraint to the PresburgerSet.
* \param vars The specified domain vars in constraint expression.
*/
- void UpdateConstraint(const PrimExpr& constraint, const ffi::Array<Var>& vars);
+ void UpdateConstraint(const PrimExpr& constraint, const ffi::Array<tirx::PrimVar>& vars);
/*!
* \brief Generate expression that represents the constraint
@@ -104,13 +104,13 @@ class PresburgerSetNode : public IntSetNode {
* \brief Set domain vars
* \param new_vars Vars that will be taken as the domain vars
*/
- void SetVars(const ffi::Array<Var>& new_vars) { vars = new_vars; }
+ void SetVars(const ffi::Array<tirx::PrimVar>& new_vars) { vars = new_vars; }
/*!
* \brief Get the current domain vars
* \return The current doamin vars
*/
- ffi::Array<Var> GetVars() const { return vars; }
+ ffi::Array<tirx::PrimVar> GetVars() const { return vars; }
/*! \return whether integer set is empty */
bool IsEmpty() const {
@@ -120,7 +120,7 @@ class PresburgerSetNode : public IntSetNode {
TVM_FFI_DECLARE_OBJECT_INFO_FINAL("arith.PresburgerSet", PresburgerSetNode, IntSetNode);
private:
- ffi::Array<Var> vars;
+ ffi::Array<tirx::PrimVar> vars;
};
/*!
@@ -135,7 +135,8 @@ class PresburgerSet : public IntSet {
* \param vars The variables that the constraint describes about.
* \return The created PresburgerSet.
*/
- TVM_DLL PresburgerSet(const std::vector<IntegerRelation>& disjuncts, const ffi::Array<Var>& vars);
+ TVM_DLL PresburgerSet(const std::vector<IntegerRelation>& disjuncts,
+ const ffi::Array<tirx::PrimVar>& vars);
/*!
* \brief Make a new instance of PresburgerSet, collect all vars as space vars.
diff --git a/src/arith/rewrite_simplify.cc b/src/arith/rewrite_simplify.cc
index c3ac370ffc..461bb31543 100644
--- a/src/arith/rewrite_simplify.cc
+++ b/src/arith/rewrite_simplify.cc
@@ -2450,7 +2450,7 @@ bool RewriteSimplifier::Impl::CanInlineLet(const LetNode* op) {
// Only inline trivial bindings to avoid deep expression explosion
// when we need let to construct complicated expressions.
if (is_const_number(op->value)) return true;
- if (op->value.as<VarNode>()) return true;
+ if (op->value.as<PrimVar>()) return true;
return false;
}
diff --git a/src/arith/solve_linear_equation.cc b/src/arith/solve_linear_equation.cc
index 7e48d56376..382c166305 100644
--- a/src/arith/solve_linear_equation.cc
+++ b/src/arith/solve_linear_equation.cc
@@ -213,13 +213,13 @@ void SmithNormalFormDiag(std::vector<std::vector<int64_t>>* S, std::vector<std::
}
ffi::Map<Var, Range> InferRange(const ffi::Map<Var, PrimExpr>& vars_to_infer,
- const ffi::Array<Var>& ori_vars,
+ const ffi::Array<PrimVar>& ori_vars,
const ffi::Map<Var, Range>& ori_ranges) {
// The resulting ranges
ffi::Map<Var, Range> new_ranges;
std::unordered_set<const VarNode*> ori_vset;
- for (const Var& v : ori_vars) {
+ for (const PrimVar& v : ori_vars) {
ori_vset.insert(v.get());
}
@@ -341,7 +341,7 @@ IntConstraintsTransform SolveLinearEquations(const IntConstraints& system_to_sol
// Uy is U \times y
SmithNormalFormDiag(&S, &V, &V_inv_x, &Uy);
- ffi::Array<Var> new_vars;
+ ffi::Array<PrimVar> new_vars;
ffi::Array<PrimExpr> new_relations;
ffi::Map<Var, PrimExpr> new_to_old_map;
ffi::Map<Var, PrimExpr> old_to_new_map;
@@ -392,11 +392,11 @@ IntConstraintsTransform SolveLinearEquations(const IntConstraints& system_to_sol
// The j-th variable can take any integer value, create a tvm variable for it
PrimExpr to_old = analyzer_problem->Simplify(V_inv_x[j]);
std::string name_hint = "n" + std::to_string(new_vars.size());
- if (const VarNode* v_old = to_old.as<VarNode>()) {
- name_hint += "_" + v_old->name_hint;
+ if (auto old_var = to_old.as<tirx::PrimVar>()) {
+ name_hint += "_" + (*old_var)->name_hint;
}
- Var v = Var(name_hint, V_inv_x[j].ty());
- solution_for_V_inv_x.push_back(v.as_or_throw<PrimExpr>());
+ PrimVar v(name_hint, V_inv_x[j].ty());
+ solution_for_V_inv_x.push_back(v);
new_vars.push_back(v);
new_to_old_map.Set(v, to_old);
} else {
@@ -450,7 +450,7 @@ IntConstraintsTransform SolveLinearEquations(const IntConstraints& system_to_sol
// Add the rest conditions
for (const PrimExpr& cond : rest) {
- new_relations.push_back(Substitute(cond, old_to_new_map));
+ new_relations.push_back(tirx::Substitute(cond, old_to_new_map));
}
IntConstraints solution(new_vars, new_ranges, new_relations);
@@ -466,7 +466,7 @@ TVM_FFI_STATIC_INIT_BLOCK() {
if (args.size() == 1) {
*ret = SolveLinearEquations(args[0].cast<IntConstraints>());
} else if (args.size() == 3) {
- auto opt_vars = args[0].cast<ffi::Optional<ffi::Array<Var>>>();
+ auto opt_vars = args[0].cast<ffi::Optional<ffi::Array<PrimVar>>>();
auto opt_map = args[1].cast<ffi::Optional<ffi::Map<Var, Range>>>();
auto opt_relations = args[2].cast<ffi::Optional<ffi::Array<PrimExpr>>>();
IntConstraints problem(opt_vars.value_or({}), opt_map.value_or({}),
diff --git a/src/arith/solve_linear_inequality.cc b/src/arith/solve_linear_inequality.cc
index 014ff267e6..69ec115379 100644
--- a/src/arith/solve_linear_inequality.cc
+++ b/src/arith/solve_linear_inequality.cc
@@ -127,7 +127,7 @@ void AddInequality(std::vector<PrimExpr>* inequality_set, const PrimExpr& new_in
inequality_set->push_back(new_ineq);
}
-void ClassifyByPolarity(const Var& var, const std::vector<PrimExpr>& current_ineq_set,
+void ClassifyByPolarity(const PrimVar& var, const std::vector<PrimExpr>& current_ineq_set,
std::vector<PrimExpr>* next_ineq_set, std::vector<PrimExpr>* rest,
std::vector<std::pair<int64_t, PrimExpr>>* coef_pos,
std::vector<std::pair<int64_t, PrimExpr>>* coef_neg,
@@ -223,8 +223,7 @@ PartialSolvedInequalities SolveLinearInequalities(const IntConstraints& system_t
}
ffi::Map<Var, IntGroupBounds> res_bounds;
- for (const Var& var : system_to_solve->variables) {
- PrimVar v = var.as_or_throw<PrimVar>();
+ for (const PrimVar& v : system_to_solve->variables) {
TVM_FFI_ICHECK(!res_bounds.count(v))
<< "Variable " << v
<< " appears more than one time in the `variables` which might be a bug";
@@ -394,7 +393,7 @@ IntConstraints SolveInequalitiesToRange(const IntConstraints& inequalities) {
arith::Analyzer analyzer;
analyzer->Bind(vranges);
- const Var& var = *it;
+ const PrimVar& var = *it;
TVM_FFI_ICHECK(solved_bounds.count(var));
auto bnd = solved_bounds[var];
if (is_one(bnd->coef) && !bnd->equal.empty()) {
@@ -458,7 +457,7 @@ IntConstraintsTransform SolveInequalitiesDeskewRange(const IntConstraints& inequ
ffi::Map<Var, PrimExpr> res_src_to_dst;
ffi::Map<Var, PrimExpr> res_dst_to_src;
- ffi::Array<Var> res_variables;
+ ffi::Array<PrimVar> res_variables;
ffi::Array<PrimExpr> res_relations;
// this keeps being updated during determining the range of each variable.
@@ -471,7 +470,7 @@ IntConstraintsTransform SolveInequalitiesDeskewRange(const IntConstraints& inequ
// We process variables in the reverse direction to start with the most independent one.
// This order is needed to compute new ranges.
for (auto it = inequalities->variables.rbegin(); it != inequalities->variables.rend(); ++it) {
- const Var& var = *it;
+ const PrimVar& var = *it;
auto bnd = solved_bounds[var];
// Note that we replace old vars with new ones
bnd = bnd.Substitute(res_src_to_dst);
@@ -489,7 +488,7 @@ IntConstraintsTransform SolveInequalitiesDeskewRange(const IntConstraints& inequ
auto best_range = bnd.FindBestRange(vranges);
- Var new_var = var.CopyWithSuffix(".shifted");
+ PrimVar new_var = var.CopyWithSuffix(".shifted");
if (!best_range.defined()) {
res_src_to_dst.Set(var, var.as_or_throw<PrimExpr>());
res_dst_to_src.Set(var, var.as_or_throw<PrimExpr>());
@@ -536,7 +535,7 @@ IntConstraintsTransform SolveInequalitiesDeskewRange(const IntConstraints& inequ
}
// Reverse the axis so that it matches the order of the original variables
- res_variables = ffi::Array<Var>(res_variables.rbegin(), res_variables.rend());
+ res_variables = ffi::Array<PrimVar>(res_variables.rbegin(), res_variables.rend());
IntConstraints new_inequalities(res_variables, res_ranges, res_relations);
IntConstraintsTransform transform(inequalities, new_inequalities, res_src_to_dst, res_dst_to_src);
@@ -555,7 +554,7 @@ TVM_FFI_STATIC_INIT_BLOCK() {
problem = args[0].cast<IntConstraints>();
ret_ineq = SolveLinearInequalities(problem);
} else if (args.size() == 3) {
- problem = IntConstraints(args[0].cast<ffi::Array<Var>>(),
+ problem = IntConstraints(args[0].cast<ffi::Array<PrimVar>>(),
args[1].cast<ffi::Map<Var, Range>>(),
args[2].cast<ffi::Array<PrimExpr>>());
ret_ineq = SolveLinearInequalities(problem);
@@ -572,7 +571,8 @@ TVM_FFI_STATIC_INIT_BLOCK() {
*ret = SolveInequalitiesToRange(args[0].cast<IntConstraints>());
} else if (args.size() == 3) {
auto opt_map = args[1].cast<ffi::Optional<ffi::Map<Var, Range>>>();
- IntConstraints problem(args[0].cast<ffi::Array<Var>>(), opt_map.value_or({}),
+ IntConstraints problem(args[0].cast<ffi::Array<PrimVar>>(),
+ opt_map.value_or({}),
args[2].cast<ffi::Array<PrimExpr>>());
*ret = SolveInequalitiesToRange(problem);
} else {
@@ -586,7 +586,7 @@ TVM_FFI_STATIC_INIT_BLOCK() {
*ret = SolveInequalitiesDeskewRange(args[0].cast<IntConstraints>());
} else if (args.size() == 3) {
auto opt_map = args[1].cast<ffi::Optional<ffi::Map<Var, Range>>>();
- IntConstraints problem(args[0].cast<ffi::Array<Var>>(), opt_map.value_or({}),
+ IntConstraints problem(args[0].cast<ffi::Array<PrimVar>>(), opt_map.value_or({}),
args[2].cast<ffi::Array<PrimExpr>>());
*ret = SolveInequalitiesDeskewRange(problem);
} else {
diff --git a/src/backend/cuda/codegen/codegen_cuda.cc b/src/backend/cuda/codegen/codegen_cuda.cc
index 5e65eaa034..f2d6b410af 100644
--- a/src/backend/cuda/codegen/codegen_cuda.cc
+++ b/src/backend/cuda/codegen/codegen_cuda.cc
@@ -1635,7 +1635,7 @@ void CodeGenCUDA::VisitStmt_(const AllocBufferNode* op) {
this->PrintIndent();
std::string scope = GetPtrStorageScope(op->buffer->data);
- const VarNode* buffer = op->buffer->data.as<VarNode>();
+ const VarNode* buffer = op->buffer->data.get();
PrimType dtype = op->buffer->dtype;
if (scope.find("wmma.") == 0) {
diff --git a/src/ir/expr.cc b/src/ir/expr.cc
index c2b8619832..21b345587c 100644
--- a/src/ir/expr.cc
+++ b/src/ir/expr.cc
@@ -39,6 +39,7 @@ namespace tvm {
TVM_FFI_STATIC_INIT_BLOCK() {
ExprNode::RegisterReflection();
BaseFuncNode::RegisterReflection();
+ VarNode::RegisterReflection();
GlobalVarNode::RegisterReflection();
CallNode::RegisterReflection();
IntImmNode::RegisterReflection();
@@ -250,6 +251,36 @@ TVM_FFI_STATIC_INIT_BLOCK() {
});
}
+Var::Var(ffi::String name_hint, ffi::Optional<Type> ty_annotation, Span span) {
+ ffi::ObjectPtr<VarNode> n = ffi::make_object<VarNode>();
+ n->name_hint = std::move(name_hint);
+ if (ty_annotation.has_value()) {
+ n->ty = ty_annotation.value();
+ }
+ n->span = std::move(span);
+ data_ = std::move(n);
+}
+
+Var Var::copy_with_name(const ffi::String& name) const {
+ TVM_FFI_CHECK_EQ(type_index(), VarNode::RuntimeTypeIndex(), TypeError)
+ << "Cannot copy a Var runtime subtype as an ordinary Var";
+ ffi::ObjectPtr<VarNode> copy = ffi::make_object<VarNode>(*get());
+ copy->name_hint = name;
+ return Var(std::move(copy));
+}
+
+Var Var::CopyWithSuffix(const ffi::String& suffix) const {
+ return copy_with_name(get()->name_hint + suffix);
+}
+
+Var Var::copy_with_dtype(PrimType dtype) const {
+ TVM_FFI_CHECK_EQ(type_index(), VarNode::RuntimeTypeIndex(), TypeError)
+ << "Cannot copy a Var runtime subtype as an ordinary Var";
+ ffi::ObjectPtr<VarNode> copy = ffi::make_object<VarNode>(*get());
+ copy->ExprNode::ty = std::move(dtype);
+ return Var(std::move(copy));
+}
+
GlobalVar::GlobalVar(ffi::String name_hint, Span span) {
ffi::ObjectPtr<GlobalVarNode> n = ffi::make_object<GlobalVarNode>();
n->name_hint = std::move(name_hint);
@@ -274,6 +305,8 @@ Call::Call(Type ret_ty, Expr op, ffi::Array<Expr> args, Attrs attrs, ffi::Array<
TVM_FFI_STATIC_INIT_BLOCK() {
namespace refl = tvm::ffi::reflection;
refl::GlobalDef()
+ .def("ir.Var", [](ffi::String name_hint, ffi::Optional<Type> ty_annotation,
+ Span span) { return Var(name_hint, ty_annotation, span); })
.def("ir.GlobalVar", [](ffi::String name) { return GlobalVar(name); })
.def("ir.Call",
[](Type ret_ty, Expr op, ffi::Array<Expr> args, Attrs attrs, ffi::Array<Type> ty_args,
diff --git a/src/relax/analysis/analysis.cc b/src/relax/analysis/analysis.cc
index c53c2c4cd4..3a39d0d004 100644
--- a/src/relax/analysis/analysis.cc
+++ b/src/relax/analysis/analysis.cc
@@ -94,11 +94,20 @@ class VarVisitor : protected ExprVisitor {
vars_.Insert(v);
}
+ void MarkTypeDefinitionsBound(const Type& type) {
+ for (const Var& var : DefinableTIRVarsInType(type)) {
+ MarkBounded(var);
+ }
+ }
+
+ void VisitExprDepTypeField(const Type&) final {}
+
void VisitExpr_(const VarNode* var) final { vars_.Insert(ffi::GetRef<Var>(var)); }
void VisitExpr_(const FunctionNode* op) final {
for (const auto& param : op->params) {
MarkBounded(param);
+ MarkTypeDefinitionsBound(GetType(param));
}
VisitExpr(op->body);
}
@@ -128,6 +137,7 @@ class VarVisitor : protected ExprVisitor {
void VisitBinding_(const MatchCastNode* binding) final {
MarkBounded(binding->var);
+ MarkTypeDefinitionsBound(binding->ty);
ExprVisitor::VisitBinding_(binding);
}
@@ -171,6 +181,12 @@ ffi::Optional<Expr> FindImpureCall(const Expr& expr, const ffi::Optional<Expr>&
// does *not* mean the outer function contains an impure call
}
+ void VisitExprDepTypeField(const Type&) final {
+ // Purity is a property of runtime value computation. Dependent type
+ // expressions may contain TIRX calls, but those calls are not executed
+ // as part of the Relax expression being checked.
+ }
+
void VisitExpr_(const CallNode* call) override {
// ignore recursive calls if we find one
bool is_recursive = (own_name_ && own_name_.value().same_as(call->op));
diff --git a/src/relax/analysis/computable_at_compile_time.cc b/src/relax/analysis/computable_at_compile_time.cc
index 3610fb49bb..2faf44652d 100644
--- a/src/relax/analysis/computable_at_compile_time.cc
+++ b/src/relax/analysis/computable_at_compile_time.cc
@@ -38,7 +38,7 @@ class CompileTimeCollector : ExprVisitor {
static ffi::Array<Var> Collect(const Function& func) {
CompileTimeCollector visitor;
visitor(func);
- return ffi::Array<Var>(visitor.known_relax_vars_.begin(), visitor.known_relax_vars_.end());
+ return ffi::Array<Var>(visitor.computable_vars_.begin(), visitor.computable_vars_.end());
}
private:
@@ -56,13 +56,8 @@ class CompileTimeCollector : ExprVisitor {
void VisitBinding(const Binding& binding) override {
Expr value = GetBoundValue(binding);
bool can_compute_at_compile_time = [&]() {
- for (const auto& relax_var : FreeVars(value)) {
- if (!known_relax_vars_.count(relax_var)) {
- return false;
- }
- }
- for (const auto& tir_var : FreeSymbolicVars(value)) {
- if (!known_tir_vars_.count(tir_var)) {
+ for (const Var& var : FreeVars(value)) {
+ if (!known_vars_.count(var)) {
return false;
}
}
@@ -78,14 +73,15 @@ class CompileTimeCollector : ExprVisitor {
}
void MarkAsKnown(const Var& var) {
- known_relax_vars_.insert(var);
+ known_vars_.insert(var);
+ computable_vars_.insert(var);
for (const auto& tir_var : DefinableTIRVarsInType(GetType(var))) {
- known_tir_vars_.insert(tir_var);
+ known_vars_.insert(tir_var);
}
}
- support::OrderedSet<Var, ffi::ObjectPtrHash, ffi::ObjectPtrEqual> known_relax_vars_;
- std::unordered_set<tirx::Var> known_tir_vars_;
+ support::OrderedSet<Var, ffi::ObjectPtrHash, ffi::ObjectPtrEqual> known_vars_;
+ support::OrderedSet<Var, ffi::ObjectPtrHash, ffi::ObjectPtrEqual> computable_vars_;
};
} // namespace
diff --git a/src/relax/analysis/layout_transformation.cc b/src/relax/analysis/layout_transformation.cc
index f9a9cc341c..b0140a643c 100644
--- a/src/relax/analysis/layout_transformation.cc
+++ b/src/relax/analysis/layout_transformation.cc
@@ -41,7 +41,7 @@ using namespace tirx;
/*! \brief Checks if a transformation is bijective affine over the given ranges */
static bool IsBijectiveAffine(const IndexMap& m, const ffi::Array<Range>& ranges) {
- ffi::Map<tirx::Var, Range> input_iters;
+ ffi::Map<tirx::PrimVar, Range> input_iters;
TVM_FFI_ICHECK_EQ(m->initial_indices.size(), ranges.size());
for (size_t i = 0; i < ranges.size(); i++) {
input_iters.Set(m->initial_indices[i], ranges[i]);
@@ -85,8 +85,8 @@ class IndexAnalyzer : public ExprVisitor {
}
void VisitIterMark(const arith::IterMark& op) {
- if (const auto* var = op->source.as<tirx::VarNode>())
- iterators_.push_back(ffi::GetRef<tirx::Var>(var));
+ if (auto var = op->source.as<tirx::PrimVar>())
+ iterators_.push_back(var.value());
else
VisitExpr(op->source);
VisitExpr(op->extent);
@@ -298,8 +298,8 @@ static ffi::Optional<IndexMap> InferLayoutTransformation(const SpatialLayout& sr
continue;
}
- auto new_dim = tirx::Var("d");
- PrimExpr new_dim_expr = new_dim.as_or_throw<PrimExpr>();
+ tirx::PrimVar new_dim("d");
+ PrimExpr new_dim_expr = new_dim;
initial_indices_it = initial_indices.insert(initial_indices_it, new_dim);
final_indices_it = final_indices.insert(final_indices_it, new_dim_expr);
// Advance past the newly inserted element so subsequent iterations start correctly.
@@ -531,7 +531,7 @@ class BlockAnalyzer : public StmtExprVisitor {
private:
bool can_transform_block_;
IndexMap write_transformation_;
- ffi::Map<tirx::Var, Range> spatial_dom_;
+ ffi::Map<tirx::PrimVar, Range> spatial_dom_;
arith::Analyzer arith_analyzer_;
SBlock block_;
diff --git a/src/relax/analysis/tir_op_pattern_kind.cc b/src/relax/analysis/tir_op_pattern_kind.cc
index 89c40c1e4e..912fb58305 100644
--- a/src/relax/analysis/tir_op_pattern_kind.cc
+++ b/src/relax/analysis/tir_op_pattern_kind.cc
@@ -228,8 +228,8 @@ class PatternKindAnalyzer : public StmtExprVisitor {
static bool IsInjectivePattern(const BufferStore& store, const BufferLoad& load) {
std::unordered_set<const tirx::VarNode*> vars;
for (const PrimExpr& store_index : store->indices) {
- if (const auto* v = store_index.as<tirx::VarNode>()) {
- vars.insert(v);
+ if (auto var = store_index.as<tirx::PrimVar>()) {
+ vars.insert(var.value().get());
} else {
return false;
}
@@ -253,17 +253,17 @@ class PatternKindAnalyzer : public StmtExprVisitor {
static bool IsAllowReusePattern(const BufferStore& store, const BufferLoad& load) {
std::unordered_set<const tirx::VarNode*> vars;
for (const PrimExpr& index : store->indices) {
- if (const auto* v = index.as<tirx::VarNode>()) {
- vars.insert(v);
+ if (auto var = index.as<tirx::PrimVar>()) {
+ vars.insert(var.value().get());
} else {
return false;
}
}
for (const PrimExpr& index : load->indices) {
PreOrderVisit(index, [&](const ffi::ObjectRef& node) {
- if (const auto* v = node.as<tirx::VarNode>()) {
- if (vars.count(v)) {
- vars.erase(v);
+ if (auto var = node.as<tirx::PrimVar>()) {
+ if (vars.count(var.value().get())) {
+ vars.erase(var.value().get());
}
}
return true;
@@ -408,7 +408,7 @@ bool HasReshapePattern(const PrimFunc& func) {
return;
}
- ffi::Map<tirx::Var, Range> var_range;
+ ffi::Map<tirx::PrimVar, Range> var_range;
for (const IterVar& v : block->iter_vars) {
ana_->Bind(v->var, Range::FromMinExtent(v->dom->min, v->dom->extent));
var_range.Set(v->var, Range::FromMinExtent(v->dom->min, v->dom->extent));
@@ -487,7 +487,7 @@ bool HasReshapePattern(const PrimFunc& func) {
if (nontrivial_indices.defined() && !has_zero_extent) {
PrimType dtype =
!block->iter_vars.empty() ? block->iter_vars[0]->var.ty() : PrimType::Int(64);
- tirx::Var fused_var("fused", dtype);
+ tirx::PrimVar fused_var("fused", dtype);
ffi::Map<tirx::Var, PrimExpr> inverse_indices_map;
PrimExpr stride = IntImm(dtype, /*value=*/1);
for (int i = static_cast<int>(block->iter_vars.size()) - 1; i >= 0; --i) {
@@ -502,7 +502,7 @@ bool HasReshapePattern(const PrimFunc& func) {
ffi::Array<PrimExpr> simplify_res = arith::IterMapSimplify(
/*indices=*/{flattened_idx},
/*input_iters=*/
- ffi::Map<tirx::Var, Range>{{fused_var, Range(IntImm(dtype, /*value=*/0), stride)}},
+ ffi::Map<tirx::PrimVar, Range>{{fused_var, Range(IntImm(dtype, /*value=*/0), stride)}},
/*input_pred=*/IntImm::Bool(true),
/*check_level=*/arith::IterMapLevel::Surjective,
/*analyzer=*/this->ana_,
diff --git a/src/relax/analysis/type_analysis.cc b/src/relax/analysis/type_analysis.cc
index 6a596e383e..82c37fb574 100644
--- a/src/relax/analysis/type_analysis.cc
+++ b/src/relax/analysis/type_analysis.cc
@@ -30,8 +30,8 @@
#include <tvm/relax/expr_functor.h>
#include <tvm/relax/type_functor.h>
#include <tvm/tirx/analysis.h>
-#include <tvm/tirx/expr_functor.h>
#include <tvm/tirx/op.h>
+#include <tvm/tirx/stmt_functor.h>
namespace tvm {
namespace relax {
@@ -120,12 +120,11 @@ Type TypeFromStaticType(const Type& type) {
//--------------------------
// EraseToWellDefined
//--------------------------
-class WellDefinedEraser : public TypeMutator, public ExprMutatorBase, public tirx::ExprMutator {
+class WellDefinedEraser : public TypeMutator, public ExprMutatorBase {
public:
- WellDefinedEraser(std::function<ffi::Optional<PrimExpr>(const tirx::Var& var)> f_shape_var_map,
- std::function<ffi::Optional<Expr>(const Var& var)> f_var_map,
+ WellDefinedEraser(std::function<ffi::Optional<Expr>(const Var& var)> f_var_map,
arith::AnalyzerObj* ana)
- : f_shape_var_map_(f_shape_var_map), f_var_map_(f_var_map), ana_(ana) {}
+ : f_var_map_(f_var_map), ana_(ana) {}
Type VisitType_(const PrimTypeNode* op) final { return ffi::GetRef<Type>(op); }
@@ -135,8 +134,7 @@ class WellDefinedEraser : public TypeMutator, public ExprMutatorBase, public tir
if (op->values.has_value()) {
std::swap(has_undefined_, has_undefined);
- values =
- op->values.value().Map([&](PrimExpr val) { return this->VisitTypePrimExprField(val); });
+ values = op->values.value().Map([&](PrimExpr val) { return VisitPrimitiveExpr(val); });
std::swap(has_undefined_, has_undefined);
}
// erase symbolic shape if we have undefined.
@@ -188,11 +186,25 @@ class WellDefinedEraser : public TypeMutator, public ExprMutatorBase, public tir
}
using relax::ExprMutatorBase::VisitExpr_;
- using tirx::ExprMutator::VisitExpr_;
PrimExpr VisitPrimitiveExpr(const PrimExpr& expr) {
- // apply eager simplification
- PrimExpr val = tirx::ExprMutator::VisitExpr(expr).as_or_throw<PrimExpr>();
+ PrimExpr val = tirx::Substitute(expr, [this](const Var& var) -> ffi::Optional<Expr> {
+ if (var.as<DataflowVarNode>()) {
+ has_undefined_ = true;
+ return std::nullopt;
+ }
+ ffi::Optional<Expr> ret = f_var_map_ == nullptr ? std::nullopt : f_var_map_(var);
+ has_undefined_ = has_undefined_ || !ret.has_value();
+ if (!ret.has_value()) return std::nullopt;
+
+ PrimExpr value = ret.value().as_or_throw<PrimExpr>();
+ if (value->IsInstance<IntImmNode>()) {
+ return tvm::cast(PrimType::Int(64), value);
+ }
+ TVM_FFI_ICHECK(value.ty().MatchesElementType(DLDataTypeCode::kDLInt, 64))
+ << "Can only provide i64 expressions in shape";
+ return value;
+ });
if (!val.same_as(expr)) {
return ana_->Simplify(val);
} else {
@@ -200,87 +212,58 @@ class WellDefinedEraser : public TypeMutator, public ExprMutatorBase, public tir
}
}
- Expr VisitExprFallback_(const ExprNode* op) final {
- if (op->ty.as<PrimTypeNode>()) {
- return VisitPrimitiveExpr(ffi::GetRef<Expr>(op).as_or_throw<PrimExpr>());
- }
- return ExprMutatorBase::VisitExprFallback_(op);
+ Expr VisitExpr_(const ShapeExprNode* op) final {
+ ffi::Array<PrimExpr> values =
+ op->values.Map([this](const PrimExpr& expr) { return VisitPrimitiveExpr(expr); });
+ return values.same_as(op->values) ? ffi::GetRef<Expr>(op) : ShapeExpr(values, op->span);
}
- PrimExpr VisitTypePrimExprField(const PrimExpr& expr) final { return VisitPrimitiveExpr(expr); }
-
Expr VisitExpr_(const VarNode* var) final {
+ Var id = ffi::GetRef<Var>(var);
ffi::Optional<Expr> ret;
if (f_var_map_ != nullptr) {
- ret = f_var_map_(ffi::GetRef<Var>(var));
+ ret = f_var_map_(id);
}
+
has_undefined_ = has_undefined_ || !ret.has_value();
if (ret.has_value()) {
- TVM_FFI_ICHECK(ret.as<VarNode>() || ret.as<ShapeExprNode>())
+ TVM_FFI_ICHECK((ret.as<VarNode>() && !ret.as<DataflowVarNode>()) || ret.as<ShapeExprNode>())
<< "Only allow Expr in Type to be ShapeExpr or Var";
}
return ret.value_or(ffi::GetRef<Expr>(var));
}
- Expr VisitExpr_(const tirx::VarNode* var) final {
- ffi::Optional<PrimExpr> ret;
- if (f_shape_var_map_ != nullptr) {
- ret = f_shape_var_map_(ffi::GetRef<tirx::Var>(var));
- }
- has_undefined_ = has_undefined_ || !ret.has_value();
-
- if (ret.has_value()) {
- PrimExpr value = ret.value();
- if (value->IsInstance<IntImmNode>()) {
- return tvm::cast(PrimType::Int(64), value);
- }
- TVM_FFI_ICHECK(value.ty().MatchesElementType(DLDataTypeCode::kDLInt, 64))
- << "Can only provide i64 expressions in shape";
- return value;
- } else {
- return ffi::GetRef<tirx::Var>(var);
- }
+ Expr VisitExpr_(const DataflowVarNode* var) final {
+ has_undefined_ = true;
+ return ffi::GetRef<Expr>(var);
}
private:
bool has_undefined_ = false;
- std::function<ffi::Optional<PrimExpr>(const tirx::Var& var)> f_shape_var_map_;
std::function<ffi::Optional<Expr>(const Var& var)> f_var_map_;
arith::AnalyzerObj* ana_;
};
-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) {
+Type EraseToWellDefined(const Type& info,
+ std::function<ffi::Optional<Expr>(const Var& var)> f_var_map) {
arith::Analyzer analyzer;
- return EraseToWellDefined(info, f_shape_var_map, f_var_map, analyzer);
+ return EraseToWellDefined(info, f_var_map, analyzer);
}
-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) {
- return WellDefinedEraser(f_shape_var_map, f_var_map, ana.get()).VisitType(info);
+Type EraseToWellDefined(const Type& info,
+ std::function<ffi::Optional<Expr>(const Var& var)> f_var_map,
+ const arith::Analyzer& ana) {
+ return WellDefinedEraser(f_var_map, ana.get()).VisitType(info);
}
-Type EraseToWellDefined(const Type& info, ffi::Map<tirx::Var, PrimExpr> shape_var_map,
- ffi::Map<Var, Expr> var_map) {
+Type EraseToWellDefined(const Type& info, ffi::Map<Var, Expr> var_map) {
arith::Analyzer analyzer;
- return EraseToWellDefined(info, shape_var_map, var_map, analyzer);
+ return EraseToWellDefined(info, var_map, analyzer);
}
-Type EraseToWellDefined(const Type& info, ffi::Map<tirx::Var, PrimExpr> shape_var_map,
- ffi::Map<Var, Expr> var_map, const arith::Analyzer& ana) {
- std::function<ffi::Optional<PrimExpr>(const tirx::Var& var)> f_shape_var_map = nullptr;
+Type EraseToWellDefined(const Type& info, ffi::Map<Var, Expr> var_map, const arith::Analyzer& ana) {
std::function<ffi::Optional<Expr>(const Var& var)> f_var_map = nullptr;
- if (!shape_var_map.empty()) {
- f_shape_var_map = [&](const tirx::Var& var) -> ffi::Optional<PrimExpr> {
- auto it = shape_var_map.find(var);
- if (it != shape_var_map.end()) return (*it).second;
- return std::nullopt;
- };
- }
-
if (!var_map.empty()) {
f_var_map = [&](const Var& var) -> ffi::Optional<Expr> {
auto it = var_map.find(var);
@@ -289,15 +272,21 @@ Type EraseToWellDefined(const Type& info, ffi::Map<tirx::Var, PrimExpr> shape_va
};
}
- return EraseToWellDefined(info, f_shape_var_map, f_var_map, ana);
+ return EraseToWellDefined(info, f_var_map, ana);
}
TVM_FFI_STATIC_INIT_BLOCK() {
namespace refl = tvm::ffi::reflection;
- refl::GlobalDef().def(
- "relax.analysis.EraseToWellDefined",
- [](const Type& info, ffi::Map<tirx::Var, PrimExpr> shape_var_map,
- ffi::Map<Var, Expr> var_map) { return EraseToWellDefined(info, shape_var_map, var_map); });
+ refl::GlobalDef().def("relax.analysis.EraseToWellDefined",
+ [](const Type& info, ffi::Map<tirx::Var, PrimExpr> shape_var_map,
+ ffi::Map<Var, Expr> var_map) {
+ for (const auto& [var, value] : shape_var_map) {
+ TVM_FFI_CHECK(var.as<tirx::PrimVar>(), TypeError)
+ << "Expected an exact primitive Var, but received " << var;
+ var_map.Set(var, value);
+ }
+ return EraseToWellDefined(info, var_map);
+ });
}
//--------------------------
@@ -871,14 +860,13 @@ class CallRetTypeDeriver : public TypeBaseChecker {
TVM_FFI_VISIT_END(call->args[i]);
}
// map the ret using the populated var map.
- return EraseToWellDefined(finfo->ret, shape_var_map_, var_map_);
+ return EraseToWellDefined(finfo->ret, var_map_);
}
protected:
// Whether to populate map in params.
bool populate_mapping_{true};
// for simplicity, we make these fields public so the user can access them.
- ffi::Map<tirx::Var, PrimExpr> shape_var_map_;
ffi::Map<Var, Expr> var_map_;
using TypeBaseChecker::ShapeMatchCheck;
@@ -889,16 +877,15 @@ class CallRetTypeDeriver : public TypeBaseChecker {
return TypeBaseChecker::PrimExprMatchCheck(param, arg);
}
- if (auto* ptr = param.as<tirx::VarNode>()) {
- auto var = ffi::GetRef<tirx::Var>(ptr);
- auto it = shape_var_map_.find(var);
+ if (auto var = param.as<tirx::PrimVar>()) {
+ auto it = var_map_.find(var.value());
// not populated
- if (it == shape_var_map_.end()) {
- shape_var_map_.Set(var, arg);
+ if (it == var_map_.end()) {
+ var_map_.Set(var.value(), arg);
return BaseCheckResult::kPass;
} else {
// Best effort prove.
- PrimExpr mapped_value = (*it).second;
+ PrimExpr mapped_value = (*it).second.as_or_throw<PrimExpr>();
if (analyzer_->CanProveEqual(mapped_value, arg)) return BaseCheckResult::kPass;
return BaseCheckResult::kFailL2;
}
@@ -916,7 +903,8 @@ class CallRetTypeDeriver : public TypeBaseChecker {
return TypeBaseChecker::ShapeMatchCheck(lhs, rhs);
}
- if (auto* ptr = lhs.as<VarNode>()) {
+ if (auto* ptr = lhs.as<VarNode>();
+ ptr && !lhs.as<DataflowVarNode>() && !lhs.as<tirx::PrimVar>()) {
auto var = ffi::GetRef<Var>(ptr);
auto it = var_map_.find(var);
// not populated
@@ -1191,12 +1179,14 @@ class TIRVarsDetector : public TypeVisitor {
private:
void VisitTypePrimExprField(PrimExpr expr) {
if (collection_type == VarType::Definition) {
- if (auto opt = expr.as<tirx::Var>()) {
+ if (auto opt = expr.as<tirx::PrimVar>()) {
RecordTIRVar(opt.value());
}
} else if (collection_type == VarType::Usage) {
for (const tirx::Var& tir_var : tirx::UndefinedVars(expr)) {
- RecordTIRVar(tir_var);
+ if (auto prim_var = tir_var.as<tirx::PrimVar>()) {
+ RecordTIRVar(prim_var.value());
+ }
}
} else {
TVM_FFI_THROW(InternalError)
@@ -1304,9 +1294,7 @@ TVM_FFI_STATIC_INIT_BLOCK() {
CollectNonNegativeExpressions);
}
-class SymbolicVarCollector : public relax::ExprVisitor,
- public relax::TypeVisitor,
- public tirx::ExprVisitor {
+class SymbolicVarCollector : public relax::ExprVisitor, public relax::TypeVisitor {
public:
static ffi::Array<tirx::Var> Free(const Expr& expr) {
SymbolicVarCollector collector;
@@ -1327,8 +1315,6 @@ class SymbolicVarCollector : public relax::ExprVisitor,
private:
using relax::ExprVisitor::VisitExpr;
using relax::ExprVisitor::VisitExpr_;
- using tirx::ExprVisitor::VisitExpr;
- using tirx::ExprVisitor::VisitExpr_;
// Possible mode of visitor, used as bit-flags
enum VisitMode {
@@ -1397,36 +1383,53 @@ class SymbolicVarCollector : public relax::ExprVisitor,
}
void VisitTypeExprField(const Expr& expr) final {
- relax::ExprVisitor::VisitExpr(expr);
if (auto* shape = expr.as<relax::ShapeExprNode>()) {
for (const auto& val : shape->values) {
this->VisitTypeExprField(val);
}
- }
- if (auto prim_value = expr.as<PrimExpr>()) {
+ return;
+ } else if (auto prim_value = expr.as<PrimExpr>()) {
this->VisitTypeExprField(prim_value.value());
+ return;
}
+ relax::ExprVisitor::VisitExpr(expr);
}
void VisitTypeExprField(const PrimExpr& expr) final {
if (mode_ & VisitMode::kProvideDefinition) {
- if (auto var = expr.as<tirx::Var>()) {
+ if (auto var = expr.as<tirx::PrimVar>()) {
defined_symbolic_var_.insert(var.value());
}
}
if (mode_ & VisitMode::kRequireDefinition) {
- tirx::ExprVisitor::VisitExpr(expr);
+ relax::ExprVisitor::VisitExpr(expr);
}
}
- void VisitExpr_(const tirx::VarNode* op) final {
- tirx::Var var = ffi::GetRef<tirx::Var>(op);
+ void VisitExpr_(const VarNode* op) final {
+ if (!op->ty.as<PrimTypeNode>()) {
+ return;
+ }
+ tirx::PrimVar var = ffi::GetRef<Var>(op).as_or_throw<tirx::PrimVar>();
// default mode, check defined.
if (defined_symbolic_var_.count(var) == 0) {
free_symbolic_var_.insert(var);
}
}
+ void VisitExpr_(const DataflowVarNode*) final {}
+
+ void VisitVarDef_(const VarNode* op) final {
+ if (op->ty.as<PrimTypeNode>()) {
+ defined_symbolic_var_.insert(ffi::GetRef<Var>(op).as_or_throw<tirx::PrimVar>());
+ }
+ relax::ExprVisitor::VisitVarDef_(op);
+ }
+
+ void VisitVarDef_(const DataflowVarNode* op) final {
+ relax::ExprVisitor::VisitVarDef_(static_cast<const VarNode*>(op));
+ }
+
// Run callback with mode.
template <typename FType>
void WithMode(VisitMode mode, FType callback) {
diff --git a/src/relax/analysis/udchain.cc b/src/relax/analysis/udchain.cc
index 3efd39cb7f..f993a27c89 100644
--- a/src/relax/analysis/udchain.cc
+++ b/src/relax/analysis/udchain.cc
@@ -142,6 +142,7 @@ std::set<const VarNode*> GetUsedVars(const Expr& expr) {
class UsedVars : public ExprVisitor {
public:
std::set<const VarNode*> used_vars;
+ void VisitExprDepTypeField(const Type&) final {}
void VisitExpr_(const VarNode* op) override { used_vars.insert(op); }
} visitor;
visitor.VisitExpr(expr);
diff --git a/src/relax/analysis/well_formed.cc b/src/relax/analysis/well_formed.cc
index d28d5167a4..78ff8e187c 100644
--- a/src/relax/analysis/well_formed.cc
+++ b/src/relax/analysis/well_formed.cc
@@ -33,23 +33,21 @@
* attribute, the name of the GlobalVar must equal the value of the
* `global_symbol` attribute value.
* 4. Any variable cannot used as different function parameters in the same IRModule.
- * 5. Any symbolic var cannot present across different functions in the same IRModule.
- * 6. Vars are defined before use.
- * 7. Vars are defined exactly once.
- * 8. Symbolic Vars are defined before use.
- * 9. DataflowVars cannot be defined inside BindingBlock.
- * 10. Vars defined in IfNode, except the return Var, are invisible
+ * 5. Vars, including those introduced in definitional type positions, are defined before use.
+ * 6. Vars are defined exactly once.
+ * 7. DataflowVars cannot be defined inside BindingBlock.
+ * 8. Vars defined in IfNode, except the return Var, are invisible
* out of the If body.(May change for new AST designs)
- * 11. SeqExpr only serves as function body, or in the true and
+ * 9. SeqExpr only serves as function body, or in the true and
* false branches in IfNode.
- * 12. The IR is in ANF:
+ * 10. The IR is in ANF:
* (a) Expressions cannot contain nested complex expressions.
* Here are the expressions that may be nested inside other expressions:
* Var, DataflowVar, GlobalVar, Constant, ShapeExpr,
* Op, Tuple (we call these "leaf" expressions).
* (b) The right-hand side of a binding may contain a non-leaf expression
* (where all expressions nested in it are leaf expressions),
- * other than SeqExprs (see rule 6)
+ * other than SeqExprs (see rule 9)
* (c) Exceptions: The body of a Function node and the true branch
* and false branch of If nodes *must* be SeqExprs.
* (d) Places where non-leaf expressions cannot appear:
@@ -57,13 +55,13 @@
* * The cond field of If nodes
* * The op or args fields of Call nodes
* * Inside the fields of Tuple nodes
- * 13. Expr always has ty (with the exception of Op).
- * 14. DataflowBlocks may not contain If nodes.
- * 15. DataflowBlocks may not contain calls to impure functions or operators
+ * 11. Expr always has ty (with the exception of Op).
+ * 12. DataflowBlocks may not contain If nodes.
+ * 13. DataflowBlocks may not contain calls to impure functions or operators
* (only checked if check_ty is true).
- * 16. If a function has is_pure set to true and the kForcePure attribute is not set,
+ * 14. If a function has is_pure set to true and the kForcePure attribute is not set,
* the body may not contain any impure call (only checked if check_ty is true).
- * 17. If the kForcePure attribute is set for a function,
+ * 15. If the kForcePure attribute is set for a function,
* that function's is_pure field must be true.
*/
#include <tvm/ffi/cast.h>
@@ -90,9 +88,7 @@ namespace relax {
// Scope Frame to store manage the var context.
//
/*! \brief Helper to implement well formed check.*/
-class WellFormedChecker : public relax::ExprVisitor,
- public relax::TypeVisitor,
- public tirx::ExprVisitor {
+class WellFormedChecker : public relax::ExprVisitor, public relax::TypeVisitor {
public:
// Throws ffi::Error on the first well-formedness violation, seeded with the
// offending node so the caller can resolve an access path. Returns normally
@@ -123,8 +119,6 @@ class WellFormedChecker : public relax::ExprVisitor,
: mod_(std::move(mod)), check_ty(check_ty), cur_visited_func_(nullptr) {}
using relax::ExprVisitor::VisitExpr_;
- using tirx::ExprVisitor::VisitExpr;
- using tirx::ExprVisitor::VisitExpr_;
// Possible mode of visitor
enum class VisitMode {
@@ -139,6 +133,16 @@ class WellFormedChecker : public relax::ExprVisitor,
kMatchVarDef
};
+ class PrimitiveExprChecker : public tirx::ExprVisitor {
+ public:
+ explicit PrimitiveExprChecker(WellFormedChecker* parent) : parent_(parent) {}
+
+ private:
+ void VisitExpr_(const tvm::VarNode* op) final { parent_->VisitExpr(ffi::GetRef<Expr>(op)); }
+
+ WellFormedChecker* parent_;
+ };
+
/*! \brief Get the name of a function for use in error messages. */
std::string FuncName(const FunctionNode* func) const {
auto it = func_name_map_.find(func);
@@ -246,20 +250,9 @@ class WellFormedChecker : public relax::ExprVisitor,
// save the var_set_ for local function
auto prev_var_set = var_set_;
auto prev_dataflow_var_set = dataflow_var_set_;
- auto prev_symbolic_var_set = symbolic_var_set_;
bool old_dataflow_state = is_dataflow_;
- // symbolic var is not captured across function boundaries
- symbolic_var_set_.clear();
is_dataflow_ = false;
- // first populate defs in params
- WithMode(VisitMode::kMatchVarDef, [&]() {
- TVM_FFI_ICHECK(mode_ == VisitMode::kMatchVarDef);
- for (Var param : op->params) {
- relax::TypeVisitor::VisitType(GetType(param));
- }
- });
-
// ensure the purity attributes are valid
if (op->GetAttr<bool>(relax::attr::kForcePure).value_or(false) && !op->is_pure) {
TVM_FFI_VISIT_THROW(ValueError, op->span)
@@ -268,17 +261,36 @@ class WellFormedChecker : public relax::ExprVisitor,
<< " should be true only if is_pure is also true.";
}
- // check all expr are well defined.
+ // Bare symbolic occurrences in parameter types are signature definitions,
+ // including when a later parameter type defines a symbol used by an
+ // earlier compound expression.
+ WithMode(VisitMode::kMatchVarDef, [&]() {
+ for (Var param : op->params) {
+ relax::TypeVisitor::VisitType(GetType(param));
+ }
+ });
+ for (Var param : op->params) {
+ if (!prev_var_set.count(param)) {
+ var_set_.erase(param);
+ }
+ }
+
+ // Runtime parameters themselves are definitions only in source order.
for (Var param : op->params) {
- this->VisitVarDef(param);
+ if (auto* dataflow_var = param.as<DataflowVarNode>()) {
+ TVM_FFI_VISIT_THROW(ValueError, ffi::GetRef<DataflowVar>(dataflow_var))
+ << "DataflowVar " << param << " is defined outside DataflowBlock.";
+ }
+ RegisterVarDefinition(param);
auto it = param_var_func_map_.find(param);
- if (it != param_var_func_map_.end()) {
+ if (it != param_var_func_map_.end() && it->second != cur_visited_func_) {
TVM_FFI_VISIT_THROW(ValueError, param->span)
<< "Relax variable " << param << " is used as a parameter in both function "
<< FuncName(it->second) << " and function " << FuncName(cur_visited_func_) << ".";
}
param_var_func_map_.insert({param, cur_visited_func_});
+ CheckType(param.get());
}
// check function ret_ty
if (!op->ret_ty.IsMissing()) {
@@ -304,7 +316,6 @@ class WellFormedChecker : public relax::ExprVisitor,
is_dataflow_ = old_dataflow_state;
dataflow_var_set_ = prev_dataflow_var_set;
var_set_ = prev_var_set;
- symbolic_var_set_ = prev_symbolic_var_set;
if (cur_visited_func_ == op) {
cur_visited_func_ = nullptr;
@@ -316,7 +327,7 @@ class WellFormedChecker : public relax::ExprVisitor,
TVM_FFI_VISIT_BEGIN();
if (IsLeafOrTuple(call->op)) {
const FunctionNode* prev_visited_func = cur_visited_func_;
- cur_visited_func_ = nullptr; // close the symbolic var dup check
+ cur_visited_func_ = nullptr; // do not attribute the callee to its caller
this->VisitExpr(call->op);
cur_visited_func_ = prev_visited_func;
} else {
@@ -445,13 +456,10 @@ class WellFormedChecker : public relax::ExprVisitor,
}
std::unordered_set<Var> previous_var_set = var_set_;
- std::unordered_set<tirx::Var> previous_symbolic_var_set = symbolic_var_set_;
this->VisitSeqExpr(op->true_branch.get());
var_set_ = previous_var_set;
- symbolic_var_set_ = previous_symbolic_var_set;
this->VisitSeqExpr(op->false_branch.get());
var_set_ = previous_var_set;
- symbolic_var_set_ = previous_symbolic_var_set;
CheckType(op);
TVM_FFI_VISIT_END(ffi::GetRef<Expr>(op));
@@ -460,7 +468,7 @@ class WellFormedChecker : public relax::ExprVisitor,
void VisitExpr_(const ShapeExprNode* op) final {
for (PrimExpr expr : op->values) {
// check if the symbolic vars in the expr are defined, e.g, 2 * m
- tirx::ExprVisitor::VisitExpr(expr);
+ VisitPrimitiveExpr(expr);
if (expr.ty().code() != DLDataTypeCode::kDLInt) {
TVM_FFI_VISIT_THROW(TypeError, expr)
<< "Shape expressions must be of integer type, but got " << expr.ty()->dtype;
@@ -556,37 +564,21 @@ class WellFormedChecker : public relax::ExprVisitor,
void VisitVarDef_(const VarNode* var) final {
Var gv = ffi::GetRef<Var>(var);
- if (var_set_.count(gv) == 1) {
- TVM_FFI_VISIT_THROW(ValueError, gv) << "Var " << gv << " is defined more than once.";
- }
- // register Var
- var_set_.insert(gv);
+ RegisterVarDefinition(gv);
CheckType(var);
}
- void VisitExpr_(const tirx::VarNode* op) final {
- tirx::Var var = ffi::GetRef<tirx::Var>(op);
- // default mode, check defined.
- if (symbolic_var_set_.count(var) == 0) {
- TVM_FFI_VISIT_THROW(ValueError, var) << "Symbolic Var " << var << " is not defined.";
+ void RegisterVarDefinition(const Var& var) {
+ if (var_set_.count(var) == 1) {
+ TVM_FFI_VISIT_THROW(ValueError, var) << "Var " << var << " is defined more than once.";
}
-
- // don't perform the check
- if (cur_visited_func_ == nullptr) {
- return;
- }
-
- // check across functions presence
- auto it = symbolic_var_func_map_.find(var);
- if (it != symbolic_var_func_map_.end() && it->second != cur_visited_func_) {
- TVM_FFI_VISIT_THROW(ValueError, var->span)
- << "Symbolic Var " << var << " is present in both function " << FuncName(it->second)
- << " and function " << FuncName(cur_visited_func_) << " in the same Module.";
- }
- symbolic_var_func_map_.insert({var, cur_visited_func_});
+ var_set_.insert(var);
}
void VisitType_(const FuncTypeNode* op) final {
+ auto prev_var_set = var_set_;
+ const FunctionNode* prev_visited_func = cur_visited_func_;
+ cur_visited_func_ = nullptr;
if (op->params.has_value()) {
WithMode(VisitMode::kMatchVarDef, [&]() {
TVM_FFI_ICHECK(mode_ == VisitMode::kMatchVarDef);
@@ -596,16 +588,19 @@ class WellFormedChecker : public relax::ExprVisitor,
});
}
this->VisitType(op->ret);
+ cur_visited_func_ = prev_visited_func;
+ var_set_ = prev_var_set;
}
void VisitTypeExprField(const Expr& expr) final {
+ if (auto* op = expr.as<DataflowVarNode>()) {
+ TVM_FFI_VISIT_THROW(ValueError, ffi::GetRef<DataflowVar>(op))
+ << "DataflowVar cannot be used as a dependent type expression.";
+ }
if (mode_ == VisitMode::kMatchVarDef) {
- // populate symbolic var in first occurrence
- if (auto* op = expr.as<relax::VarNode>()) {
- auto var = ffi::GetRef<relax::Var>(op);
- if (var_set_.count(var) == 0) {
- var_set_.insert(var);
- }
+ // Populate a canonical Var definition on first occurrence.
+ if (auto* op = expr.as<VarNode>()) {
+ MarkTypeVarDefinition(ffi::GetRef<Var>(op));
}
if (auto* shape = expr.as<relax::ShapeExprNode>()) {
for (auto val : shape->values) {
@@ -618,17 +613,36 @@ class WellFormedChecker : public relax::ExprVisitor,
}
void VisitTypeExprField(const PrimExpr& expr) final {
+ if (auto* op = expr.as<DataflowVarNode>()) {
+ TVM_FFI_VISIT_THROW(ValueError, ffi::GetRef<DataflowVar>(op))
+ << "DataflowVar cannot be used as a dependent type expression.";
+ }
if (mode_ == VisitMode::kMatchVarDef) {
- // populate symbolic var in first occurrence
- if (auto* op = expr.as<tirx::VarNode>()) {
- auto var = ffi::GetRef<tirx::Var>(op);
- if (symbolic_var_set_.count(var) == 0) {
- symbolic_var_set_.insert(var);
- }
+ // A primitive Var is the same canonical Var as a Relax binding. A
+ // first occurrence in a definitional type position enters the one
+ // lexical definition set used by all Var occurrences.
+ if (auto* op = expr.as<VarNode>()) {
+ MarkTypeVarDefinition(ffi::GetRef<Var>(op));
}
} else {
- tirx::ExprVisitor::VisitExpr(expr);
+ VisitPrimitiveExpr(expr);
+ }
+ }
+
+ void VisitPrimitiveExpr(const PrimExpr& expr) { PrimitiveExprChecker(this)(expr); }
+
+ void MarkTypeVarDefinition(const Var& var) {
+ var_set_.insert(var);
+ if (cur_visited_func_ == nullptr) {
+ return;
+ }
+ auto it = param_var_func_map_.find(var);
+ if (it != param_var_func_map_.end() && it->second != cur_visited_func_) {
+ TVM_FFI_VISIT_THROW(ValueError, var->span)
+ << "Var " << var << " is defined in both function " << FuncName(it->second)
+ << " and function " << FuncName(cur_visited_func_) << ".";
}
+ param_var_func_map_.insert({var, cur_visited_func_});
}
void CheckType(const ExprNode* op) {
@@ -666,9 +680,7 @@ class WellFormedChecker : public relax::ExprVisitor,
std::unordered_set<Var> var_set_;
std::unordered_set<Var> recur_vars_;
std::unordered_set<DataflowVar, ffi::ObjectPtrHash, ffi::ObjectPtrEqual> dataflow_var_set_;
- std::unordered_set<tirx::Var> symbolic_var_set_;
std::unordered_map<Var, const FunctionNode*> param_var_func_map_;
- std::unordered_map<tirx::Var, const FunctionNode*> symbolic_var_func_map_;
tvm::OpAttrMap<FNormalize> op_map_normalize_ = Op::GetAttrMap<FNormalize>("FNormalize");
tvm::OpAttrMap<FValidate> op_map_validate_ = Op::GetAttrMap<FValidate>("FValidate");
diff --git a/src/relax/backend/contrib/cutlass/codegen.cc b/src/relax/backend/contrib/cutlass/codegen.cc
index a299289266..4dacc019d8 100644
--- a/src/relax/backend/contrib/cutlass/codegen.cc
+++ b/src/relax/backend/contrib/cutlass/codegen.cc
@@ -152,6 +152,8 @@ using OutputType = std::vector<Output>;
class CodegenCutlass : public relax::MemoizedExprTranslator<OutputType>,
public relax::contrib::CodegenCBase {
public:
+ using relax::MemoizedExprTranslator<OutputType>::VisitExpr_;
+
CodegenCutlass(const std::string& id, const ffi::Map<Var, Expr>& bindings)
: ext_func_id_(id), bindings_(bindings) {}
diff --git a/src/relax/backend/contrib/tensorrt/codegen.cc b/src/relax/backend/contrib/tensorrt/codegen.cc
index c5338e2d27..b09c988931 100644
--- a/src/relax/backend/contrib/tensorrt/codegen.cc
+++ b/src/relax/backend/contrib/tensorrt/codegen.cc
@@ -196,11 +196,11 @@ class CollectFromCompositeFunctionBody : public ExprVisitor {
if (initial.size() != final_indices.size()) return true;
ffi::Array<int64_t> permutation;
for (const PrimExpr& expr : final_indices) {
- const auto* var = expr.as<tirx::VarNode>();
- if (var == nullptr) return true;
+ auto var = expr.as<tirx::PrimVar>();
+ if (!var.has_value()) return true;
int64_t pos = -1;
for (size_t j = 0; j < initial.size(); ++j) {
- if (initial[j].get() == var) {
+ if (initial[j].get() == var.value().get()) {
pos = static_cast<int64_t>(j);
break;
}
diff --git a/src/relax/backend/vm/codegen_vm.cc b/src/relax/backend/vm/codegen_vm.cc
index c27a39279f..176d28eabe 100644
--- a/src/relax/backend/vm/codegen_vm.cc
+++ b/src/relax/backend/vm/codegen_vm.cc
@@ -217,6 +217,10 @@ class CodeGenVM : public ExprFunctor<Instruction::Arg(const Expr&)> {
return it->second;
}
+ Instruction::Arg VisitExpr_(const DataflowVarNode* op) final {
+ return VisitExpr_(static_cast<const VarNode*>(op));
+ }
+
Instruction::Arg VisitExpr_(const ConstantNode* op) final {
auto arg = builder_->ConvertConstant(op->data);
@@ -242,18 +246,12 @@ class CodeGenVM : public ExprFunctor<Instruction::Arg(const Expr&)> {
return builder_->ConvertConstant(ffi::Shape(shape));
}
- Instruction::Arg VisitExprFallback_(const ExprNode* op) final {
- PrimExpr value = ffi::GetRef<Expr>(op).as_or_throw<PrimExpr>();
- if (auto* int_imm = value.as<IntImmNode>()) {
- return builder_->ConvertConstant(int_imm->value);
- } else if (auto* float_imm = value.as<FloatImmNode>()) {
- return builder_->ConvertConstant(float_imm->value);
- } else {
- TVM_FFI_THROW(InternalError)
- << "PrimExpr should only contain constant after VMShapeLower, "
- << "but received " << value << " with type " << value->GetTypeKey();
- TVM_FFI_UNREACHABLE();
- }
+ Instruction::Arg VisitExpr_(const IntImmNode* op) final {
+ return builder_->ConvertConstant(op->value);
+ }
+
+ Instruction::Arg VisitExpr_(const FloatImmNode* op) final {
+ return builder_->ConvertConstant(op->value);
}
Instruction::Arg VisitExpr_(const StringImmNode* op) final {
diff --git a/src/relax/backend/vm/codegen_vm_tir.cc b/src/relax/backend/vm/codegen_vm_tir.cc
index ee515489da..81482b0293 100644
--- a/src/relax/backend/vm/codegen_vm_tir.cc
+++ b/src/relax/backend/vm/codegen_vm_tir.cc
@@ -292,6 +292,43 @@ class CodeGenVMTIR : public ExprFunctor<ffi::Optional<Expr>(const Expr&)> {
return it->second;
}
+ ffi::Optional<Expr> VisitExpr_(const DataflowVarNode* op) final {
+ return VisitExpr_(static_cast<const VarNode*>(op));
+ }
+
+#define VM_TIR_PRIM_EXPR(OP) \
+ ffi::Optional<Expr> VisitExpr_(const OP* op) final { return ffi::GetRef<Expr>(op); }
+
+ VM_TIR_PRIM_EXPR(tirx::BufferLoadNode);
+ VM_TIR_PRIM_EXPR(tirx::AddNode);
+ VM_TIR_PRIM_EXPR(tirx::SubNode);
+ VM_TIR_PRIM_EXPR(tirx::MulNode);
+ VM_TIR_PRIM_EXPR(tirx::DivNode);
+ VM_TIR_PRIM_EXPR(tirx::ModNode);
+ VM_TIR_PRIM_EXPR(tirx::FloorDivNode);
+ VM_TIR_PRIM_EXPR(tirx::FloorModNode);
+ VM_TIR_PRIM_EXPR(tirx::MinNode);
+ VM_TIR_PRIM_EXPR(tirx::MaxNode);
+ VM_TIR_PRIM_EXPR(tirx::EQNode);
+ VM_TIR_PRIM_EXPR(tirx::NENode);
+ VM_TIR_PRIM_EXPR(tirx::LTNode);
+ VM_TIR_PRIM_EXPR(tirx::LENode);
+ VM_TIR_PRIM_EXPR(tirx::GTNode);
+ VM_TIR_PRIM_EXPR(tirx::GENode);
+ VM_TIR_PRIM_EXPR(tirx::AndNode);
+ VM_TIR_PRIM_EXPR(tirx::OrNode);
+ VM_TIR_PRIM_EXPR(tirx::CastNode);
+ VM_TIR_PRIM_EXPR(tirx::NotNode);
+ VM_TIR_PRIM_EXPR(tirx::SelectNode);
+ VM_TIR_PRIM_EXPR(tirx::RampNode);
+ VM_TIR_PRIM_EXPR(tirx::BroadcastNode);
+ VM_TIR_PRIM_EXPR(tirx::ShuffleNode);
+ VM_TIR_PRIM_EXPR(tvm::IntImmNode);
+ VM_TIR_PRIM_EXPR(tvm::FloatImmNode);
+ VM_TIR_PRIM_EXPR(tirx::StringImmNode);
+
+#undef VM_TIR_PRIM_EXPR
+
ffi::Optional<Expr> VisitExpr_(const ConstantNode* op) final {
return ConstListGet(builder_->ConvertConstant(op->data).value());
}
@@ -309,8 +346,6 @@ class CodeGenVMTIR : public ExprFunctor<ffi::Optional<Expr>(const Expr&)> {
return ConstListGet(builder_->ConvertConstant(ffi::Shape(shape)).value());
}
- ffi::Optional<Expr> VisitExprFallback_(const ExprNode* op) final { return ffi::GetRef<Expr>(op); }
-
ffi::Optional<Expr> VisitExpr_(const StringImmNode* op) final {
return ConstListGet(builder_->ConvertConstant(op->value).value());
}
diff --git a/src/relax/backend/vm/vm_shape_lower.cc b/src/relax/backend/vm/vm_shape_lower.cc
index 3fb8ea0434..12ff275784 100644
--- a/src/relax/backend/vm/vm_shape_lower.cc
+++ b/src/relax/backend/vm/vm_shape_lower.cc
@@ -34,6 +34,8 @@
#include <tvm/tirx/op.h>
#include <tvm/tirx/stmt_functor.h>
+#include <unordered_set>
+
namespace tvm {
namespace relax {
@@ -73,6 +75,8 @@ struct MatchShapeTodoItem {
using PrimExprSlotMap =
std::unordered_map<PrimExpr, PrimExprSlot*, ffi::StructuralHash, tirx::ExprDeepEqual>;
+using LiveVarSet = std::unordered_set<Var, ffi::ObjectPtrHash, ffi::ObjectPtrEqual>;
+
static bool IsRelaxOwnedCall(const CallNode* call) {
auto op = call->op.as<Op>();
if (!op) return true;
@@ -90,32 +94,63 @@ class PrimExprSlotCollector : public ExprVisitor, public TypeVisitor {
PrimExprSlotCollector collector;
collector.slot_vec_ = slot_vec;
collector.slot_map_ = slot_map;
- // collect shape declaration in func params
- for (auto param : func->params) {
+ VarUsageInfo usage = CollectVarUsage(func);
+ collector.bound_values_ = &usage.bound_values;
+ for (const Var& output : usage.outputs) collector.MarkLive(output);
+ for (const Var& param : func->params) collector.VisitType(GetType(param));
+ collector.VisitType(func->ret_ty);
+ collector.VisitExpr(func->body);
+ collector.collect_slots_ = true;
+ for (const Var& param : func->params) {
collector.VisitType(GetType(param));
- collector.VisitExpr(param);
}
collector.VisitExpr(func->body);
collector.VisitType(func->ret_ty);
}
private:
- void VisitExpr_(const CallNode* op) final {
- if (auto prim_expr = ffi::GetRef<Call>(op).as<PrimExpr>(); prim_expr && !IsRelaxOwnedCall(op)) {
- CollectPrimExprSlot(prim_expr.value());
+ void VisitBinding_(const VarBindingNode* op) final {
+ if (!collect_slots_) {
+ VisitVarDef(op->var);
return;
}
- ExprVisitor::VisitExpr_(op);
+ bool collect_scalar = collect_scalar_;
+ collect_scalar_ = live_vars_.count(op->var);
+ ExprVisitor::VisitExpr(op->value);
+ collect_scalar_ = collect_scalar;
+ VisitVarDef(op->var);
}
- void VisitExprFallback_(const ExprNode* op) final {
- if (auto prim_expr = ffi::GetRef<Expr>(op).as<PrimExpr>()) {
- CollectPrimExprSlot(prim_expr.value());
- } else {
- ExprVisitor::VisitExprFallback_(op);
+ void MarkLive(const PrimExpr& expr) {
+ for (const Var& var : tirx::UndefinedVars(expr)) MarkLive(var);
+ }
+
+ void MarkLive(const Var& var) {
+ if (!live_vars_.insert(var).second) return;
+ if (auto it = bound_values_->find(var); it != bound_values_->end()) {
+ for (const VarNode* used : GetUsedVars((*it).second)) MarkLive(ffi::GetRef<Var>(used));
}
}
+ void VisitExpr_(const VarNode* op) final {
+ Var var = ffi::GetRef<Var>(op);
+ if (collect_scalar_ && !var.as<DataflowVarNode>()) {
+ if (auto prim_var = var.as<tirx::PrimVar>();
+ prim_var && prim_var.value().ty()->dtype == DLDataType{kDLInt, 64, 1}) {
+ HandlePrimExpr(prim_var.value());
+ }
+ }
+ ExprVisitor::VisitExpr_(op);
+ }
+
+ void VisitExpr_(const ShapeExprNode* op) final {
+ for (const PrimExpr& value : op->values) HandlePrimExpr(value);
+ }
+
+ void HandlePrimExpr(const PrimExpr& expr) {
+ collect_slots_ ? CollectPrimExprSlot(expr) : MarkLive(expr);
+ }
+
void CollectPrimExprSlot(const PrimExpr& expr) {
if (expr->IsInstance<IntImmNode>()) return;
if (const auto* call = expr.as<CallNode>()) {
@@ -150,14 +185,18 @@ class PrimExprSlotCollector : public ExprVisitor, public TypeVisitor {
// Do not recurse into function type as it is self-contained
}
- void VisitTypePrimExprField(const PrimExpr& expr) final { CollectPrimExprSlot(expr); }
+ void VisitTypePrimExprField(const PrimExpr& expr) final { HandlePrimExpr(expr); }
- void VisitTypeExprField(const PrimExpr& expr) final { CollectPrimExprSlot(expr); }
+ void VisitTypeExprField(const PrimExpr& expr) final { HandlePrimExpr(expr); }
void VisitTypeExprField(const Expr& expr) final { ExprVisitor::VisitExpr(expr); }
std::vector<std::unique_ptr<PrimExprSlot>>* slot_vec_;
PrimExprSlotMap* slot_map_;
+ const ffi::Map<Var, Expr>* bound_values_;
+ LiveVarSet live_vars_;
+ bool collect_slots_{false};
+ bool collect_scalar_{true};
};
/*!
@@ -258,20 +297,44 @@ class VMShapeLowerMutator
using ExprMutator::VisitExpr_;
+ Expr VisitExpr_(const VarNode* op) final {
+ Var var = ffi::GetRef<Var>(op);
+ if (!var.as<DataflowVarNode>()) {
+ if (auto prim_var = var.as<tirx::PrimVar>(); prim_var && slot_map_.count(*prim_var)) {
+ return RewritePrimValue(*prim_var);
+ }
+ }
+ return ExprMutator::VisitExpr_(op);
+ }
+
Expr VisitExpr_(const CallNode* op) final {
- if (auto prim_expr = ffi::GetRef<Call>(op).as<PrimExpr>(); prim_expr && !IsRelaxOwnedCall(op)) {
+ if (auto prim_expr = ffi::GetRef<Call>(op).as<PrimExpr>();
+ prim_expr && slot_map_.count(prim_expr.value()) && !IsRelaxOwnedCall(op)) {
return RewritePrimValue(prim_expr.value());
}
return ExprMutator::VisitExpr_(op);
}
Expr VisitExprFallback_(const ExprNode* op) final {
- if (auto prim_expr = ffi::GetRef<Expr>(op).as<PrimExpr>()) {
+ if (auto prim_expr = ffi::GetRef<Expr>(op).as<PrimExpr>();
+ prim_expr && slot_map_.count(prim_expr.value())) {
return RewritePrimValue(prim_expr.value());
}
return ExprMutator::VisitExprFallback_(op);
}
+ void VisitBinding_(const VarBindingNode* binding) final {
+ if (binding->var.as<DataflowVarNode>()) {
+ ExprMutator::VisitBinding_(binding);
+ return;
+ }
+ ExprMutator::VisitBinding_(binding);
+ PrimExprSlot* result_slot = GetPrimValueSlot(binding->var);
+ if (result_slot != nullptr) {
+ EmitRuntimePrimValueMatch(binding->var, result_slot, "");
+ }
+ }
+
// Unit rewrite function per function.
Function Rewrite(GlobalVar gvar, Function func) {
// prepare mapping and heap var
@@ -308,6 +371,9 @@ class VMShapeLowerMutator
<< "], param=" << func->params[i]->name_hint << ", annotation=" << ty << ") ";
this->CheckMatchCast(ty, func->params[i], true, i >= num_input, err_ctx.str(),
&match_todos);
+ if (PrimExprSlot* slot = GetPrimValueSlot(func->params[i])) {
+ match_todos.push_back(MatchShapeTodoItem{func->params[i], {slot->expr}, err_ctx.str()});
+ }
}
// insert heap generation logic.
match_todos = this->RunMatch(match_todos, false);
@@ -350,6 +416,29 @@ class VMShapeLowerMutator
//-------------------------------------------------------
static DLDataType ShapeDType() { return DLDataType{kDLInt, 64, 1}; }
+ PrimExprSlot* GetPrimValueSlot(const Var& var) const {
+ if (var.as<DataflowVarNode>()) return nullptr;
+ auto prim_var = var.as<tirx::PrimVar>();
+ if (!prim_var) return nullptr;
+ auto it = slot_map_.find(PrimExpr(*prim_var));
+ return it == slot_map_.end() ? nullptr : it->second;
+ }
+
+ /*! \brief Match one runtime register against its demanded leaf slot. */
+ void EmitRuntimePrimValueMatch(const Var& var, PrimExprSlot* slot, const ffi::String& err_ctx) {
+ Var runtime_var = var;
+ if (auto it = var_remap_.find(var); it != var_remap_.end()) {
+ runtime_var = it->second;
+ }
+
+ auto [code, rvalue] = MakeMatchArgs(slot->expr, false);
+ ffi::Array<Expr> args = {runtime_var, shape_heap_, IntImm::Int64(static_cast<int>(code)),
+ rvalue, GetErrContext(err_ctx)};
+ builder_->Emit(Call(Type::Missing(), builtin_match_prim_value_, args, Attrs(), {void_ty_}),
+ "_");
+ this->EmitOutstandingPrimExprCompute();
+ }
+
/*! \brief populate additional information in the slot. */
void PopulateSlotInfo() {
for (auto& kv : slot_map_) {
@@ -462,10 +551,10 @@ class VMShapeLowerMutator
}
void VisitBinding_(const MatchCastNode* binding) final {
- Expr value = this->VisitExpr(binding->value);
- std::vector<MatchShapeTodoItem> match_todos;
std::ostringstream err_ctx;
err_ctx << "ErrorContext(match_cast, ty=" << binding->ty << ") ";
+ Expr value = this->VisitExpr(binding->value);
+ std::vector<MatchShapeTodoItem> match_todos;
// always_check=false
this->CheckMatchCast(binding->ty, value, false, false, err_ctx.str(), &match_todos);
@@ -476,6 +565,9 @@ class VMShapeLowerMutator
// These checks are emitted as extra, in codegen
// match-cast is simply ignored and treated as a normal binding.
ExprMutator::VisitBinding_(binding);
+ if (PrimExprSlot* slot = GetPrimValueSlot(binding->var)) {
+ EmitRuntimePrimValueMatch(binding->var, slot, err_ctx.str());
+ }
}
// Do not override shape in type fields
diff --git a/src/relax/distributed/axis_group_graph.cc b/src/relax/distributed/axis_group_graph.cc
index c1359d8f22..042d459386 100644
--- a/src/relax/distributed/axis_group_graph.cc
+++ b/src/relax/distributed/axis_group_graph.cc
@@ -32,10 +32,14 @@ namespace tvm {
namespace tirx {
Var GetShardingVarFromIndex(PrimExpr index, ffi::Map<Var, Range> var_range,
const arith::Analyzer& analyzer) {
- if (index.as<VarNode>()) {
- return index.as_or_throw<Var>();
+ if (auto prim_var = index.as<PrimVar>()) {
+ return prim_var.value();
}
- arith::IterSumExpr iter_sum = arith::NormalizeToIterSum(index, var_range, analyzer);
+ ffi::Map<PrimVar, Range> primitive_var_range;
+ for (const auto& [var, range] : var_range) {
+ primitive_var_range.Set(var.as_or_throw<PrimVar>(), range);
+ }
+ arith::IterSumExpr iter_sum = arith::NormalizeToIterSum(index, primitive_var_range, analyzer);
if (!is_zero(iter_sum->base)) {
return Var();
}
@@ -44,17 +48,17 @@ Var GetShardingVarFromIndex(PrimExpr index, ffi::Map<Var, Range> var_range,
}
// floormod(floordiv(source, lower_factor), extent) * scale
arith::IterSplitExpr highest_iter_split = iter_sum->args[0];
- const auto* source_var = highest_iter_split->source->source.as<VarNode>();
+ auto source_var = highest_iter_split->source->source.as<PrimVar>();
if (!source_var) {
return Var();
}
+ Var var = source_var.value();
// the floormod must take no effect
- if (!analyzer->CanProve(floordiv(var_range[ffi::GetRef<Var>(source_var)]->extent,
- highest_iter_split->lower_factor) <=
+ if (!analyzer->CanProve(floordiv(var_range[var]->extent, highest_iter_split->lower_factor) <=
highest_iter_split->extent)) {
return Var();
}
- return ffi::GetRef<Var>(source_var);
+ return var;
}
} // namespace tirx
} // namespace tvm
diff --git a/src/relax/distributed/transform/lower_global_view_to_local_view.cc b/src/relax/distributed/transform/lower_global_view_to_local_view.cc
index c04970f631..8c9ca99fb6 100644
--- a/src/relax/distributed/transform/lower_global_view_to_local_view.cc
+++ b/src/relax/distributed/transform/lower_global_view_to_local_view.cc
@@ -321,8 +321,9 @@ class DistributedBufferCompactor : StmtExprMutator {
if (!iter_var_shards_.count(iter_var->var)) {
continue;
}
- TVM_FFI_ICHECK(iter_value.as<VarNode>());
- loop_var_shards_[iter_value.as_or_throw<Var>()] = iter_var_shards_[iter_var->var];
+ auto loop_var = iter_value.as<PrimVar>();
+ TVM_FFI_ICHECK(loop_var);
+ loop_var_shards_[loop_var.value()] = iter_var_shards_[iter_var->var];
}
return realize;
}
diff --git a/src/relax/ir/block_builder.cc b/src/relax/ir/block_builder.cc
index 60cefd081d..f62851e445 100644
--- a/src/relax/ir/block_builder.cc
+++ b/src/relax/ir/block_builder.cc
@@ -32,6 +32,7 @@
#include <tvm/relax/transform.h>
#include <tvm/relax/type.h>
#include <tvm/relax/type_functor.h>
+#include <tvm/relax/utils.h>
#include <tvm/runtime/logging.h>
#include <tvm/tirx/function.h>
@@ -196,9 +197,9 @@ class BlockBuilderImpl : public BlockBuilderNode {
// defined in parameter type annotations. The implementation
// is correct (since we will simply erase all relax Vars in EraseToWellDefined),
// but can be further improved.
- ffi::Map<tirx::Var, PrimExpr> var_map = TypeVarCollector::Collect(GetType(var));
+ ffi::Map<tirx::PrimVar, PrimExpr> var_map = TypeVarCollector::Collect(GetType(var));
for (const auto& kv : var_map) {
- const tirx::Var& shape_var = kv.first;
+ const tirx::PrimVar& shape_var = kv.first;
const PrimExpr& shape_expr = kv.second;
auto it = shape_var_map.find(shape_var);
if (it == shape_var_map.end()) {
@@ -207,7 +208,7 @@ class BlockBuilderImpl : public BlockBuilderNode {
// of shape inference. In many cases, knowning that the
// shape variable is non-negative allows for simpler
// expressions for dynamic shapes.
- analyzer_->MarkGlobalNonNegValue(shape_var.as_or_throw<PrimExpr>());
+ analyzer_->MarkGlobalNonNegValue(shape_var);
} else {
const PrimExpr& old_shape_expr = (*it).second;
TVM_FFI_ICHECK(old_shape_expr.same_as(shape_expr) ||
@@ -332,7 +333,7 @@ class BlockBuilderImpl : public BlockBuilderNode {
//
// TODO(relax-team) tracks the var defined also through match-cast.
/*! \brief set of defined symbolic vars, value as themself. */
- ffi::Map<tirx::Var, PrimExpr> shape_var_map;
+ ffi::Map<tirx::PrimVar, PrimExpr> shape_var_map;
};
/*! \brief A stack to store block frames. */
@@ -471,7 +472,7 @@ class BlockBuilderImpl : public BlockBuilderNode {
// shape vars as defined when calling BeginScope(params)
class TypeVarCollector : public TypeVisitor {
public:
- static ffi::Map<tirx::Var, PrimExpr> Collect(const Type& ty) {
+ static ffi::Map<tirx::PrimVar, PrimExpr> Collect(const Type& ty) {
TypeVarCollector collector;
collector(ty);
return collector.shape_var_map_;
@@ -482,8 +483,8 @@ class BlockBuilderImpl : public BlockBuilderNode {
if (const auto* shape_expr = op->shape.as<ShapeExprNode>()) {
for (const PrimExpr& s : shape_expr->values) {
// Only collect single var defined shape. Ignore something like `R.Tensor((m + 1, n + 1))
- if (const auto* var = s.as<tirx::VarNode>()) {
- shape_var_map_.Set(ffi::GetRef<tirx::Var>(var), s);
+ if (auto var = s.as<tirx::PrimVar>()) {
+ shape_var_map_.Set(var.value(), s);
}
}
}
@@ -492,14 +493,14 @@ class BlockBuilderImpl : public BlockBuilderNode {
void VisitType_(const ShapeTypeNode* op) final {
for (const PrimExpr& s : op->values.value_or(ffi::Array<PrimExpr>())) {
// Only collect single var defined shape. Ignore something like `R.Shape((m + 1, n + 1))
- if (const auto* var = s.as<tirx::VarNode>()) {
- shape_var_map_.Set(ffi::GetRef<tirx::Var>(var), s);
+ if (auto var = s.as<tirx::PrimVar>()) {
+ shape_var_map_.Set(var.value(), s);
}
}
}
private:
- ffi::Map<tirx::Var, PrimExpr> shape_var_map_;
+ ffi::Map<tirx::PrimVar, PrimExpr> shape_var_map_;
};
};
@@ -578,7 +579,11 @@ class Normalizer : public BlockBuilderImpl, private ExprFunctor<Expr(const Expr&
RELAX_EXPR_NORMALIZER_LEAF(StringImmNode);
RELAX_EXPR_NORMALIZER_LEAF(DataTypeImmNode);
- Expr VisitExprFallback_(const ExprNode* op) final { return ffi::GetRef<Expr>(op); }
+ Expr VisitExprDefault_(const ffi::Object* op) final {
+ Expr expr = ffi::GetRef<Expr>(static_cast<const ExprNode*>(op));
+ if (expr.as<PrimExpr>()) return expr;
+ return ExprFunctor::VisitExprDefault_(op);
+ }
template <typename T>
Expr VisitVar_(const typename T::ContainerType* var) {
@@ -858,12 +863,14 @@ class Normalizer : public BlockBuilderImpl, private ExprFunctor<Expr(const Expr&
return info;
}
auto* curr_scope = CurrentScopeFrame();
- auto f_shape_var_map = [curr_scope](tirx::Var var) -> ffi::Optional<PrimExpr> {
- auto it = curr_scope->shape_var_map.find(var);
+ auto f_var_map = [curr_scope](const Var& var) -> ffi::Optional<Expr> {
+ auto prim_var = var.as<tirx::PrimVar>();
+ if (!prim_var) return std::nullopt;
+ auto it = curr_scope->shape_var_map.find(prim_var.value());
if (it != curr_scope->shape_var_map.end()) return (*it).second;
return std::nullopt;
};
- return EraseToWellDefined(info, f_shape_var_map);
+ return EraseToWellDefined(info, f_var_map);
}
Expr VisitWithNewScope(const Expr& expr, ffi::Optional<ffi::Array<Var>> params = std::nullopt) {
diff --git a/src/relax/ir/dataflow_expr_rewriter.cc b/src/relax/ir/dataflow_expr_rewriter.cc
index c5e3a77a12..ca8567bf80 100644
--- a/src/relax/ir/dataflow_expr_rewriter.cc
+++ b/src/relax/ir/dataflow_expr_rewriter.cc
@@ -774,8 +774,8 @@ PatternMatchingRewriter PatternMatchingRewriter::FromModule(IRModule mod) {
// wouldn't normally be normalized into a variable.
Var intermediate_var("intermediate_var", GetType(matched_expr));
wildcard_bindings.push_back(VarBinding(intermediate_var, matched_expr));
- wildcard_bindings.push_back(
- MatchCast(func_replacement->params[i], intermediate_var, GetType(matched_expr)));
+ wildcard_bindings.push_back(MatchCast(func_replacement->params[i], intermediate_var,
+ GetType(func_replacement->params[i])));
}
new_blocks.push_back(DataflowBlock(wildcard_bindings));
diff --git a/src/relax/ir/dataflow_matcher.cc b/src/relax/ir/dataflow_matcher.cc
index ef881e317b..21717bdb8c 100644
--- a/src/relax/ir/dataflow_matcher.cc
+++ b/src/relax/ir/dataflow_matcher.cc
@@ -462,8 +462,8 @@ PrimExpr DFPatternMatcher::SimplifyCondition(PrimExpr condition) {
auto sort_key = [](PrimExpr expr) -> ffi::String {
if (const auto* equal = expr.as<tirx::EQNode>()) {
- if (const auto* var = equal->a.as<tirx::VarNode>()) {
- return var->name_hint;
+ if (auto var = equal->a.as<tirx::PrimVar>()) {
+ return var.value()->name_hint;
}
}
return "";
diff --git a/src/relax/ir/expr.cc b/src/relax/ir/expr.cc
index 3da61a48fb..050b2dc8b0 100644
--- a/src/relax/ir/expr.cc
+++ b/src/relax/ir/expr.cc
@@ -31,7 +31,6 @@ TVM_FFI_STATIC_INIT_BLOCK() {
TupleNode::RegisterReflection();
TupleGetItemNode::RegisterReflection();
ShapeExprNode::RegisterReflection();
- VarNode::RegisterReflection();
BindingNode::RegisterReflection();
DataflowVarNode::RegisterReflection();
ConstantNode::RegisterReflection();
@@ -139,22 +138,6 @@ TVM_FFI_STATIC_INIT_BLOCK() {
});
}
-Var::Var(ffi::String name_hint, ffi::Optional<Type> ty_annotation, Span span) {
- ffi::ObjectPtr<VarNode> n = ffi::make_object<VarNode>();
- n->name_hint = std::move(name_hint);
- if (ty_annotation.has_value()) {
- n->ty = ty_annotation.value();
- }
- n->span = std::move(span);
- data_ = std::move(n);
-}
-
-TVM_FFI_STATIC_INIT_BLOCK() {
- namespace refl = tvm::ffi::reflection;
- refl::GlobalDef().def("relax.Var", [](ffi::String name_hint, ffi::Optional<Type> ty_annotation,
- Span span) { return Var(name_hint, ty_annotation, span); });
-}
-
DataflowVar::DataflowVar(ffi::String name_hint, ffi::Optional<Type> ty_annotation, Span span) {
ffi::ObjectPtr<DataflowVarNode> n = ffi::make_object<DataflowVarNode>();
n->name_hint = std::move(name_hint);
@@ -393,18 +376,17 @@ Function::Function(ffi::Array<Var> params, Expr body, ffi::Optional<Type> ret_ty
// for the function's return type. When hoisting the body's type
// to the function return type, symbolic variables may only be
// used if they were defined by the function's parameters.
- auto f_shape_var_map = [&] {
+ auto f_var_map = [&] {
auto tir_vars = DefinableTIRVarsInType(TupleType(params.Map(GetType)));
std::unordered_set<tirx::Var> lookup(tir_vars.begin(), tir_vars.end());
- return [lookup = std::move(lookup)](const tirx::Var& var) -> ffi::Optional<PrimExpr> {
- if (lookup.count(var)) {
- return var.as_or_throw<PrimExpr>();
- } else {
- return std::nullopt;
+ return [lookup = std::move(lookup)](const Var& var) -> ffi::Optional<Expr> {
+ if (auto prim_var = var.as<tirx::PrimVar>(); prim_var && lookup.count(prim_var.value())) {
+ return prim_var.value().as_or_throw<PrimExpr>();
}
+ return std::nullopt;
};
}();
- ret_ty = EraseToWellDefined(body_ty.value(), f_shape_var_map);
+ ret_ty = EraseToWellDefined(body_ty.value(), f_var_map);
}
FuncType func_ty(param_ty, ret_ty.value(), is_pure);
diff --git a/src/relax/ir/expr_functor.cc b/src/relax/ir/expr_functor.cc
index 963d302936..4647d51070 100644
--- a/src/relax/ir/expr_functor.cc
+++ b/src/relax/ir/expr_functor.cc
@@ -175,6 +175,97 @@ void ExprVisitor::VisitExpr_(const CallNode* op) {
VisitExprDepTypeFieldIfNeeded(this, op->ty);
}
+void ExprVisitor::VisitExpr_(const tirx::BufferLoadNode* op) {
+ this->VisitSpan(op->span);
+ for (const PrimExpr& index : op->indices) {
+ this->VisitExpr(index);
+ }
+ if (op->predicate.has_value()) {
+ this->VisitExpr(op->predicate.value());
+ }
+ VisitExprDepTypeFieldIfNeeded(this, op->ty);
+}
+
+#define RELAX_VISIT_TIRX_BINOP(OP) \
+ void ExprVisitor::VisitExpr_(const tirx::OP* op) { \
+ this->VisitSpan(op->span); \
+ this->VisitExpr(op->a); \
+ this->VisitExpr(op->b); \
+ VisitExprDepTypeFieldIfNeeded(this, op->ty); \
+ }
+
+RELAX_VISIT_TIRX_BINOP(AddNode);
+RELAX_VISIT_TIRX_BINOP(SubNode);
+RELAX_VISIT_TIRX_BINOP(MulNode);
+RELAX_VISIT_TIRX_BINOP(DivNode);
+RELAX_VISIT_TIRX_BINOP(ModNode);
+RELAX_VISIT_TIRX_BINOP(FloorDivNode);
+RELAX_VISIT_TIRX_BINOP(FloorModNode);
+RELAX_VISIT_TIRX_BINOP(MinNode);
+RELAX_VISIT_TIRX_BINOP(MaxNode);
+RELAX_VISIT_TIRX_BINOP(EQNode);
+RELAX_VISIT_TIRX_BINOP(NENode);
+RELAX_VISIT_TIRX_BINOP(LTNode);
+RELAX_VISIT_TIRX_BINOP(LENode);
+RELAX_VISIT_TIRX_BINOP(GTNode);
+RELAX_VISIT_TIRX_BINOP(GENode);
+RELAX_VISIT_TIRX_BINOP(AndNode);
+RELAX_VISIT_TIRX_BINOP(OrNode);
+
+#undef RELAX_VISIT_TIRX_BINOP
+
+void ExprVisitor::VisitExpr_(const tirx::CastNode* op) {
+ this->VisitSpan(op->span);
+ this->VisitExpr(op->value);
+ VisitExprDepTypeFieldIfNeeded(this, op->ty);
+}
+
+void ExprVisitor::VisitExpr_(const tirx::NotNode* op) {
+ this->VisitSpan(op->span);
+ this->VisitExpr(op->a);
+ VisitExprDepTypeFieldIfNeeded(this, op->ty);
+}
+
+void ExprVisitor::VisitExpr_(const tirx::SelectNode* op) {
+ this->VisitSpan(op->span);
+ this->VisitExpr(op->condition);
+ this->VisitExpr(op->true_value);
+ this->VisitExpr(op->false_value);
+ VisitExprDepTypeFieldIfNeeded(this, op->ty);
+}
+
+void ExprVisitor::VisitExpr_(const tirx::RampNode* op) {
+ this->VisitSpan(op->span);
+ this->VisitExpr(op->base);
+ this->VisitExpr(op->stride);
+ this->VisitExpr(op->lanes);
+ VisitExprDepTypeFieldIfNeeded(this, op->ty);
+}
+
+void ExprVisitor::VisitExpr_(const tirx::BroadcastNode* op) {
+ this->VisitSpan(op->span);
+ this->VisitExpr(op->value);
+ this->VisitExpr(op->lanes);
+ VisitExprDepTypeFieldIfNeeded(this, op->ty);
+}
+
+void ExprVisitor::VisitExpr_(const tirx::ShuffleNode* op) {
+ this->VisitSpan(op->span);
+ for (const PrimExpr& vector : op->vectors) {
+ this->VisitExpr(vector);
+ }
+ for (const PrimExpr& index : op->indices) {
+ this->VisitExpr(index);
+ }
+ VisitExprDepTypeFieldIfNeeded(this, op->ty);
+}
+
+void ExprVisitor::VisitExpr_(const tvm::IntImmNode* op) { this->VisitSpan(op->span); }
+
+void ExprVisitor::VisitExpr_(const tvm::FloatImmNode* op) { this->VisitSpan(op->span); }
+
+void ExprVisitor::VisitExpr_(const tirx::StringImmNode* op) { this->VisitSpan(op->span); }
+
void ExprVisitor::VisitExpr_(const IfNode* op) {
this->VisitSpan(op->span);
this->VisitExpr(op->cond);
@@ -198,8 +289,6 @@ void ExprVisitor::VisitExpr_(const ShapeExprNode* op) {
this->VisitExpr(val);
}
this->VisitSpan(op->span);
-
- VisitExprDepTypeFieldIfNeeded(this, op->ty);
}
void ExprVisitor::VisitExpr_(const ExternFuncNode* op) {
@@ -217,15 +306,7 @@ void ExprVisitor::VisitExpr_(const SeqExprNode* op) {
VisitExprDepTypeFieldIfNeeded(this, op->ty);
}
-void ExprVisitor::VisitExprFallback_(const ExprNode* op) {
- if (op->ty.IsMissing() || !op->ty.as<PrimTypeNode>()) {
- this->VisitExprDefault_(op);
- return;
- }
-
- VisitExprDepTypeFieldIfNeeded(this, op->ty);
- this->VisitSpan(op->span);
-}
+void ExprVisitor::VisitExprFallback_(const ExprNode* op) { this->VisitExprDefault_(op); }
void ExprVisitor::VisitExpr_(const StringImmNode* op) { this->VisitSpan(op->span); }
@@ -391,11 +472,7 @@ Expr ExprMutatorBase::VisitExpr_(const TupleNode* op) {
}
// Visit the use-site of a defined Var
-Expr ExprMutatorBase::VisitExpr_(const VarNode* op) {
- // type of var-use should remain stable
- // or the var itself will get replaced
- return ffi::GetRef<Expr>(op);
-}
+Expr ExprMutatorBase::VisitExpr_(const VarNode* op) { return ffi::GetRef<Expr>(op); }
// Visit the use-site of a defined DataflowVar
Expr ExprMutatorBase::VisitExpr_(const DataflowVarNode* op) {
@@ -432,13 +509,129 @@ Expr ExprMutatorBase::VisitExpr_(const CallNode* call_node) {
unchanged &= new_arg.same_as(arg);
}
- if (unchanged && VisitAndCheckTypeFieldUnchanged(call_node->ty)) {
+ Type ret_ty = call_node->ty;
+ if (!ret_ty.IsMissing()) {
+ ret_ty = this->VisitExprDepTypeField(ret_ty);
+ }
+ bool ret_ty_unchanged = ret_ty.same_as(call_node->ty);
+
+ if (unchanged && ret_ty_unchanged) {
return ffi::GetRef<Expr>(call_node);
} else {
- return Call(Type::Missing(), new_op, call_args, call_node->attrs, ty_args, call_node->span);
+ // A result-type rewrite is authoritative only while the semantic inputs
+ // are unchanged. Once the callee, arguments, or type arguments change,
+ // the old result may contain both rewritten dependent dimensions and
+ // stale derived metadata (for example an unknown dtype). Invalidate the
+ // complete result so BlockBuilder can infer it from the rebuilt call.
+ Type rebuilt_ret_ty = unchanged ? ret_ty : Type::Missing();
+ return Call(rebuilt_ret_ty, new_op, call_args, call_node->attrs, ty_args, call_node->span);
}
}
+Expr ExprMutatorBase::VisitExpr_(const tirx::BufferLoadNode* op) {
+ ffi::Array<PrimExpr> indices = op->indices.Map(
+ [this](const PrimExpr& e) { return this->VisitExpr(e).as_or_throw<PrimExpr>(); });
+ ffi::Optional<PrimExpr> predicate = op->predicate;
+ if (predicate.has_value()) {
+ predicate = this->VisitExpr(predicate.value()).as_or_throw<PrimExpr>();
+ }
+ bool predicate_unchanged =
+ !op->predicate.has_value() || predicate.value().same_as(op->predicate.value());
+ if (indices.same_as(op->indices) && predicate_unchanged) {
+ return ffi::GetRef<Expr>(op);
+ }
+ return tirx::BufferLoad(op->buffer, indices, predicate, op->span);
+}
+
+#define RELAX_MUTATE_TIRX_BINOP(OP) \
+ Expr ExprMutatorBase::VisitExpr_(const tirx::OP##Node* op) { \
+ PrimExpr a = this->VisitExpr(op->a).as_or_throw<PrimExpr>(); \
+ PrimExpr b = this->VisitExpr(op->b).as_or_throw<PrimExpr>(); \
+ if (a.same_as(op->a) && b.same_as(op->b)) { \
+ return ffi::GetRef<Expr>(op); \
+ } \
+ return tirx::OP(a, b, op->span); \
+ }
+
+RELAX_MUTATE_TIRX_BINOP(Add);
+RELAX_MUTATE_TIRX_BINOP(Sub);
+RELAX_MUTATE_TIRX_BINOP(Mul);
+RELAX_MUTATE_TIRX_BINOP(Div);
+RELAX_MUTATE_TIRX_BINOP(Mod);
+RELAX_MUTATE_TIRX_BINOP(FloorDiv);
+RELAX_MUTATE_TIRX_BINOP(FloorMod);
+RELAX_MUTATE_TIRX_BINOP(Min);
+RELAX_MUTATE_TIRX_BINOP(Max);
+RELAX_MUTATE_TIRX_BINOP(EQ);
+RELAX_MUTATE_TIRX_BINOP(NE);
+RELAX_MUTATE_TIRX_BINOP(LT);
+RELAX_MUTATE_TIRX_BINOP(LE);
+RELAX_MUTATE_TIRX_BINOP(GT);
+RELAX_MUTATE_TIRX_BINOP(GE);
+RELAX_MUTATE_TIRX_BINOP(And);
+RELAX_MUTATE_TIRX_BINOP(Or);
+
+#undef RELAX_MUTATE_TIRX_BINOP
+
+Expr ExprMutatorBase::VisitExpr_(const tirx::CastNode* op) {
+ PrimExpr value = this->VisitExpr(op->value).as_or_throw<PrimExpr>();
+ return value.same_as(op->value)
+ ? ffi::GetRef<Expr>(op)
+ : Expr(tirx::Cast(op->ty.as_or_throw<PrimType>(), value, op->span));
+}
+
+Expr ExprMutatorBase::VisitExpr_(const tirx::NotNode* op) {
+ PrimExpr a = this->VisitExpr(op->a).as_or_throw<PrimExpr>();
+ return a.same_as(op->a) ? ffi::GetRef<Expr>(op) : Expr(tirx::Not(a, op->span));
+}
+
+Expr ExprMutatorBase::VisitExpr_(const tirx::SelectNode* op) {
+ PrimExpr condition = this->VisitExpr(op->condition).as_or_throw<PrimExpr>();
+ PrimExpr true_value = this->VisitExpr(op->true_value).as_or_throw<PrimExpr>();
+ PrimExpr false_value = this->VisitExpr(op->false_value).as_or_throw<PrimExpr>();
+ if (condition.same_as(op->condition) && true_value.same_as(op->true_value) &&
+ false_value.same_as(op->false_value)) {
+ return ffi::GetRef<Expr>(op);
+ }
+ return tirx::Select(condition, true_value, false_value, op->span);
+}
+
+Expr ExprMutatorBase::VisitExpr_(const tirx::RampNode* op) {
+ PrimExpr base = this->VisitExpr(op->base).as_or_throw<PrimExpr>();
+ PrimExpr stride = this->VisitExpr(op->stride).as_or_throw<PrimExpr>();
+ PrimExpr lanes = this->VisitExpr(op->lanes).as_or_throw<PrimExpr>();
+ if (base.same_as(op->base) && stride.same_as(op->stride) && lanes.same_as(op->lanes)) {
+ return ffi::GetRef<Expr>(op);
+ }
+ return tirx::Ramp(base, stride, lanes, op->span);
+}
+
+Expr ExprMutatorBase::VisitExpr_(const tirx::BroadcastNode* op) {
+ PrimExpr value = this->VisitExpr(op->value).as_or_throw<PrimExpr>();
+ PrimExpr lanes = this->VisitExpr(op->lanes).as_or_throw<PrimExpr>();
+ if (value.same_as(op->value) && lanes.same_as(op->lanes)) {
+ return ffi::GetRef<Expr>(op);
+ }
+ return tirx::Broadcast(value, lanes, op->span);
+}
+
+Expr ExprMutatorBase::VisitExpr_(const tirx::ShuffleNode* op) {
+ ffi::Array<PrimExpr> vectors = op->vectors.Map(
+ [this](const PrimExpr& e) { return this->VisitExpr(e).as_or_throw<PrimExpr>(); });
+ ffi::Array<PrimExpr> indices = op->indices.Map(
+ [this](const PrimExpr& e) { return this->VisitExpr(e).as_or_throw<PrimExpr>(); });
+ if (vectors.same_as(op->vectors) && indices.same_as(op->indices)) {
+ return ffi::GetRef<Expr>(op);
+ }
+ return tirx::Shuffle(vectors, indices, op->span);
+}
+
+Expr ExprMutatorBase::VisitExpr_(const tvm::IntImmNode* op) { return ffi::GetRef<Expr>(op); }
+
+Expr ExprMutatorBase::VisitExpr_(const tvm::FloatImmNode* op) { return ffi::GetRef<Expr>(op); }
+
+Expr ExprMutatorBase::VisitExpr_(const tirx::StringImmNode* op) { return ffi::GetRef<Expr>(op); }
+
Expr ExprMutatorBase::VisitExpr_(const IfNode* op) {
Expr guard = this->VisitExpr(op->cond);
Expr true_b = this->VisitExpr(op->true_branch);
@@ -464,13 +657,7 @@ Expr ExprMutatorBase::VisitExpr_(const TupleGetItemNode* op) {
}
}
-Expr ExprMutatorBase::VisitExprFallback_(const ExprNode* op) {
- if (op->ty.IsMissing() || !op->ty.as<PrimTypeNode>()) {
- return this->VisitExprDefault_(op);
- }
-
- return ffi::GetRef<Expr>(op);
-}
+Expr ExprMutatorBase::VisitExprFallback_(const ExprNode* op) { return this->VisitExprDefault_(op); }
Expr ExprMutatorBase::VisitExpr_(const StringImmNode* op) { return ffi::GetRef<Expr>(op); }
@@ -577,12 +764,14 @@ Expr ExprMutator::VisitExpr_(const FunctionNode* op) {
}
}
+ Type ret_ty = all_params_unchanged ? op->ret_ty : this->VisitExprDepTypeField(op->ret_ty);
+
Expr body = this->VisitWithNewScope(op->body, params);
if (all_params_unchanged && body.same_as(op->body)) {
// No changes to the function, return the original object
return ffi::GetRef<Expr>(op);
- } else if (IsBaseOf(GetType(body), op->ret_ty)) {
+ } else if (IsBaseOf(GetType(body), ret_ty)) {
// If the function was mutated into a form that can no longer
// propagate shape information all the way to the return value, we
// may keep the return type. This is only allowed when the
@@ -591,7 +780,7 @@ Expr ExprMutator::VisitExpr_(const FunctionNode* op) {
// the previous return value was `TensorType(shape=[16,16])`
// but the body only produced `TensorType(ndim=2)`, we can
// keep the more specific information.
- return Function(params, body, op->ret_ty, op->is_pure, op->attrs);
+ return Function(params, body, ret_ty, op->is_pure, op->attrs);
} else {
// If the function was mutated such that the body produces an
// output that is incompatible with the original return struct
diff --git a/src/relax/ir/py_expr_functor.cc b/src/relax/ir/py_expr_functor.cc
index be64ae6459..3bbacd8dc7 100644
--- a/src/relax/ir/py_expr_functor.cc
+++ b/src/relax/ir/py_expr_functor.cc
@@ -106,7 +106,7 @@ class PyExprVisitorNode : public ffi::Object, public ExprVisitor {
if (vtable.can_dispatch(expr)) {
vtable(expr, this);
} else {
- VisitExprFallback_(expr.get());
+ ExprVisitor::VisitExpr(expr);
}
}
}
@@ -350,7 +350,7 @@ class PyExprMutatorNode : public ffi::Object, public ExprMutator {
if (vtable.can_dispatch(expr)) {
return builder_->Normalize(vtable(expr, this));
}
- return builder_->Normalize(VisitExprFallback_(expr.get()));
+ return ExprMutator::VisitExpr(expr);
}
}
diff --git a/src/relax/ir/transform.cc b/src/relax/ir/transform.cc
index 75f3873ce7..b77a2be9e8 100644
--- a/src/relax/ir/transform.cc
+++ b/src/relax/ir/transform.cc
@@ -29,9 +29,10 @@
#include <tvm/relax/analysis.h>
#include <tvm/relax/expr_functor.h>
#include <tvm/relax/transform.h>
-#include <tvm/relax/type_functor.h>
#include <tvm/runtime/logging.h>
+#include <unordered_set>
+
namespace tvm {
namespace relax {
namespace transform {
@@ -214,75 +215,52 @@ class DataflowBlockMutator : public ExprMutator {
/*!
* \brief Rewrite the DataflowBlockNode with pass_func_
*
- * This function will check that there are no rewrites of the global scope Vars
- * and symbolic shape Vars defined inside the dataflow block.
+ * This function will check that there are no rewrites of the global-scope
+ * Vars or symbolic Vars defined inside the dataflow block.
*/
BindingBlock VisitBindingBlock_(const DataflowBlockNode* n) final {
- // collect Global Scope Vars and Symbolic Vars inside the DataflowBlock
- ffi::Map<ffi::String, Var> global_scope_vars;
- ffi::Map<ffi::String, tirx::Var> symbolic_vars;
- for (const Binding& binding : n->bindings) {
- Var var = binding->var;
- if (const auto* match_cast = binding.as<MatchCastNode>()) {
- auto collected_vars = SymbolicVarCollector::Collect(match_cast->ty);
- for (const tirx::VarNode* var : collected_vars) {
- symbolic_vars.Set(var->name_hint, ffi::GetRef<tirx::Var>(var));
- }
- }
- if (!var.as<DataflowVarNode>()) {
- global_scope_vars.Set(var->name_hint, var);
- }
- }
+ PreservedVars preserved_vars = CollectPreservedVars(n->bindings);
// apply pass_func_ to the DataflowBlock
DataflowBlock block = ffi::GetRef<DataflowBlock>(n);
DataflowBlock updated_block = pass_func_(block, mod_, pass_ctx_);
- // raise error if there are updates of recorded Global Scope Vars and Symbolic Vars
- for (const Binding& binding : updated_block->bindings) {
- Var var = binding->var;
- if (const auto* match_cast = binding.as<MatchCastNode>()) {
- auto collected_vars = SymbolicVarCollector::Collect(match_cast->ty);
- for (const tirx::VarNode* var : collected_vars) {
- if (symbolic_vars.count(var->name_hint) > 0) {
- tirx::Var old_var = symbolic_vars[var->name_hint];
- TVM_FFI_ICHECK(var == old_var.get())
- << "Error: DataflowBlock Pass should not rewrite any Symbolic Var.";
- symbolic_vars.erase(var->name_hint);
- }
- }
- }
- if (!var.as<DataflowVarNode>() && global_scope_vars.count(var->name_hint) > 0) {
- TVM_FFI_ICHECK(var.same_as(global_scope_vars[var->name_hint]))
- << "Error: DataflowBlock Pass should not rewrite any GlobalScope Var.";
- global_scope_vars.erase(var->name_hint);
- }
+ PreservedVars updated_vars = CollectPreservedVars(updated_block->bindings);
+ for (const VarNode* var : preserved_vars.binding_vars) {
+ TVM_FFI_ICHECK(updated_vars.binding_vars.count(var))
+ << "Error: DataflowBlock Pass should not rewrite or delete any global-scope Var.";
+ }
+ for (const VarNode* var : preserved_vars.match_cast_symbolic_vars) {
+ TVM_FFI_ICHECK(updated_vars.match_cast_symbolic_vars.count(var))
+ << "Error: DataflowBlock Pass should not rewrite or delete any symbolic Var declared "
+ "by a MatchCast.";
}
- TVM_FFI_ICHECK(global_scope_vars.empty() && symbolic_vars.empty())
- << "Error: DataflowBlock Pass should not delete any GlobalScope/Symbolic Var.";
return updated_block;
}
private:
- class SymbolicVarCollector : public TypeVisitor {
- public:
- static std::unordered_set<const tirx::VarNode*> Collect(const Type& info) {
- SymbolicVarCollector collector;
- collector.VisitType(info);
- return std::move(collector.symbolic_vars_);
- }
+ using VarSet = std::unordered_set<const VarNode*>;
- private:
- void VisitTypeExprField(const PrimExpr& expr) final {
- if (const tirx::VarNode* sym_var = expr.as<tirx::VarNode>()) {
- symbolic_vars_.insert(sym_var);
+ struct PreservedVars {
+ VarSet binding_vars;
+ VarSet match_cast_symbolic_vars;
+ };
+
+ static PreservedVars CollectPreservedVars(const ffi::Array<Binding>& bindings) {
+ PreservedVars vars;
+ for (const Binding& binding : bindings) {
+ if (const auto* match_cast = binding.as<MatchCastNode>()) {
+ for (const tirx::Var& var : DefinableTIRVarsInType(match_cast->ty)) {
+ vars.match_cast_symbolic_vars.insert(var.get());
+ }
+ }
+ if (!binding->var.as<DataflowVarNode>()) {
+ vars.binding_vars.insert(binding->var.get());
}
}
-
- private:
- std::unordered_set<const tirx::VarNode*> symbolic_vars_;
- };
+ return vars;
+ }
std::function<DataflowBlock(DataflowBlock, IRModule, PassContext)> pass_func_;
IRModule mod_;
diff --git a/src/relax/op/memory/view.cc b/src/relax/op/memory/view.cc
index ea11fedafb..5a81a90f8d 100644
--- a/src/relax/op/memory/view.cc
+++ b/src/relax/op/memory/view.cc
@@ -152,9 +152,9 @@ Type InferTypeView(const Call& call, const BlockBuilder& ctx) {
<< "Operator " << call->op
<< " expects the relative_byte_offset to be a 64-bit integer, but received "
<< arg_relative_byte_offset << ", which has type " << ty;
- if (arg_relative_byte_offset.as<VarNode>()) {
- // A scalar Relax variable has an unknown value. Although it has a
- // PrimType, it is not a TIRX expression that can be analyzed.
+ if (arg_relative_byte_offset.as<DataflowVarNode>()) {
+ // A primitive DataflowVar is a runtime value, not a symbolic value
+ // that can be analyzed.
return std::nullopt;
} else if (auto prim_value = arg_relative_byte_offset.as<PrimExpr>()) {
// An offset of known value is applied. The known value may
diff --git a/src/relax/op/tensor/manipulate.cc b/src/relax/op/tensor/manipulate.cc
index c9e9322809..6421769905 100644
--- a/src/relax/op/tensor/manipulate.cc
+++ b/src/relax/op/tensor/manipulate.cc
@@ -918,8 +918,8 @@ Expr ConvertNewShapeToExpr(const Expr& data,
"Array of PrimExprs. However, the given new shape is "
<< shape;
PrimExpr len = prim_len.value();
- TVM_FFI_CHECK(!len.as<VarNode>(), TypeError)
- << "Reshape shape dimensions must be TIRX expressions, but received " << len;
+ TVM_FFI_CHECK(!len.as<DataflowVarNode>(), TypeError)
+ << "Reshape shape dimensions must be symbolic primitive expressions, but received " << len;
TVM_FFI_ICHECK(len.ty().code() == DLDataTypeCode::kDLInt)
<< "Reshape requires the new shape values to be all "
"integers. However, the give new shape is "
diff --git a/src/relax/script/builder/frame.cc b/src/relax/script/builder/frame.cc
index 80fec50fce..edee981cc6 100644
--- a/src/relax/script/builder/frame.cc
+++ b/src/relax/script/builder/frame.cc
@@ -136,12 +136,11 @@ void BindingBlockFrameNode::EnterWithScope() {
class VarReplacer : public tvm::relax::ExprMutator {
public:
explicit VarReplacer(
- std::unordered_map<tvm::relax::Var, tvm::relax::Var, ffi::ObjectPtrHash, ffi::ObjectPtrEqual>
- var_remap) {
+ std::unordered_map<tvm::Var, tvm::Var, ffi::ObjectPtrHash, ffi::ObjectPtrEqual> var_remap) {
var_remap_ = std::move(var_remap);
}
- tvm::relax::Var VisitVarDef(const tvm::relax::Var& var) override {
+ tvm::Var VisitVarDef(const tvm::Var& var) override {
// ExprMutator only applies var_remap_ at usage sites. This
// applies var_remap_ at each definition site as well.
if (auto it = var_remap_.find(var); it != var_remap_.end()) {
@@ -167,11 +166,10 @@ void BindingBlockFrameNode::ExitWithScope() {
// Step 3. Rewrite the dataflow block.
if (is_dataflow) {
// Step 3.0. Define a map to replace variables
- ffi::Array<tvm::relax::Var> new_output_vars;
- std::unordered_map<tvm::relax::Var, tvm::relax::Var, ffi::ObjectPtrHash, ffi::ObjectPtrEqual>
- var_remap;
+ ffi::Array<tvm::Var> new_output_vars;
+ std::unordered_map<tvm::Var, tvm::Var, ffi::ObjectPtrHash, ffi::ObjectPtrEqual> var_remap;
for (const auto& output_var : output_vars) {
- tvm::relax::Var new_output_var(output_var->name_hint, GetType(output_var));
+ tvm::Var new_output_var(output_var->name_hint, tvm::relax::GetType(output_var));
new_output_vars.push_back(new_output_var);
var_remap[output_var] = new_output_var;
}
diff --git a/src/relax/script/builder/ir.cc b/src/relax/script/builder/ir.cc
index 32f9f66d59..87115eca5f 100644
--- a/src/relax/script/builder/ir.cc
+++ b/src/relax/script/builder/ir.cc
@@ -33,13 +33,6 @@ namespace relax {
using tvm::script::ir_builder::details::Namer;
-TVM_STATIC_IR_FUNCTOR(Namer, vtable)
- .set_dispatch<tvm::relax::VarNode>([](const ffi::ObjectRef& node, ffi::String name) -> void {
- using tvm::relax::VarNode;
- VarNode* var = const_cast<VarNode*>(node.as<VarNode>());
- var->name_hint = name;
- });
-
TVM_STATIC_IR_FUNCTOR(Namer, vtable)
.set_dispatch<tvm::relax::DataflowVarNode>([](const ffi::ObjectRef& node,
ffi::String name) -> void {
@@ -65,9 +58,9 @@ FunctionFrame Function(bool is_pure, bool is_private) {
return FunctionFrame(n);
}
-tvm::relax::Var Arg(const ffi::String& name, const tvm::Type& ty) {
+tvm::Var Arg(const ffi::String& name, const tvm::Type& ty) {
FunctionFrame frame = FindFunctionFrame("R.Arg");
- tvm::relax::Var var(name, ty);
+ tvm::Var var(name, ty);
frame->params.push_back(var);
frame->block_builder->AddDefinitionToScope(var);
@@ -169,7 +162,7 @@ BindingBlockFrame BindingBlock() {
return BindingBlockFrame(n);
}
-void DataflowBlockOutput(const ffi::Array<tvm::relax::Var>& vars) {
+void DataflowBlockOutput(const ffi::Array<tvm::Var>& vars) {
// Step 1. Check that we're in a Dataflow block that is not ended.
ffi::Optional<BindingBlockFrame> block_frame =
IRBuilder::Current()->GetLastFrame<BindingBlockFrame>();
@@ -185,12 +178,11 @@ void DataflowBlockOutput(const ffi::Array<tvm::relax::Var>& vars) {
// Step 3. All the output variables must be global variables and must be emitted by this dataflow
// block.
- const ffi::Array<tvm::relax::Var>& emitted_vars = block_frame.value()->emitted_vars;
- for (const tvm::relax::Var& var : vars) {
+ const ffi::Array<tvm::Var>& emitted_vars = block_frame.value()->emitted_vars;
+ for (const tvm::Var& var : vars) {
TVM_FFI_CHECK(std::find_if(emitted_vars.begin(), emitted_vars.end(),
- [&](const tvm::relax::Var& emitted) {
- return emitted.same_as(var);
- }) != emitted_vars.end(),
+ [&](const tvm::Var& emitted) { return emitted.same_as(var); }) !=
+ emitted_vars.end(),
ValueError)
<< "An output variable is not emitted by this dataflow block. Please make sure "
"all dataflow block output variables are emitted exactly by this block.";
@@ -208,7 +200,7 @@ TVM_FFI_STATIC_INIT_BLOCK() {
/////////////////////////////// Bindings ///////////////////////////////
-tvm::relax::Var Emit(const tvm::relax::Expr& expr, const ffi::Optional<tvm::Type>& annotate_ty) {
+tvm::Var Emit(const tvm::relax::Expr& expr, const ffi::Optional<tvm::Type>& annotate_ty) {
using tvm::relax::GetType;
BindingBlockFrame block_frame = CheckBindingBlockFrameExistAndUnended();
const tvm::relax::BlockBuilder& block_builder = GetBlockBuilder();
@@ -222,21 +214,21 @@ tvm::relax::Var Emit(const tvm::relax::Expr& expr, const ffi::Optional<tvm::Type
<< "Invalid annotation. Got rhs value type: " << GetType(expr) << ", given type: " << ty;
}
}
- tvm::relax::Var var = block_builder->Emit(expr);
+ tvm::Var var = block_builder->Emit(expr);
block_frame->emitted_vars.push_back(var);
return var;
}
-tvm::relax::Var EmitMatchCast(const tvm::relax::Expr& value, const tvm::Type& ty) {
+tvm::Var EmitMatchCast(const tvm::relax::Expr& value, const tvm::Type& ty) {
BindingBlockFrame block_frame = CheckBindingBlockFrameExistAndUnended();
const tvm::relax::BlockBuilder& block_builder = GetBlockBuilder();
- tvm::relax::Var var = block_builder->EmitMatchCast(value, ty);
+ tvm::Var var = block_builder->EmitMatchCast(value, ty);
block_frame->emitted_vars.push_back(var);
return var;
}
-tvm::relax::Var EmitVarBinding(const tvm::relax::VarBinding& binding) {
+tvm::Var EmitVarBinding(const tvm::relax::VarBinding& binding) {
BindingBlockFrame block_frame = CheckBindingBlockFrameExistAndUnended();
const tvm::relax::BlockBuilder& block_builder = GetBlockBuilder();
block_builder->EmitNormalized(binding);
diff --git a/src/relax/script/builder/utils.h b/src/relax/script/builder/utils.h
index f353465ee3..bc8f087a0d 100644
--- a/src/relax/script/builder/utils.h
+++ b/src/relax/script/builder/utils.h
@@ -107,8 +107,8 @@ inline tvm::relax::SeqExpr GetSeqExprForBranch(const SeqExprFrame& frame, ffi::S
ffi::Array<tvm::relax::Binding> last_block_bindings(last_block->bindings.begin(),
last_block->bindings.end() - 1);
- tvm::relax::Var new_var =
- tvm::relax::Var(last_binding->var->name_hint + output_var_suffix, GetType(last_binding->var));
+ tvm::Var new_var(last_binding->var->name_hint + output_var_suffix,
+ tvm::relax::GetType(last_binding->var));
tvm::relax::Expr body;
const auto* var_binding = last_binding.as<tvm::relax::VarBindingNode>();
diff --git a/src/relax/script/printer/binding.cc b/src/relax/script/printer/binding.cc
index b0984e9576..c4a26c0204 100644
--- a/src/relax/script/printer/binding.cc
+++ b/src/relax/script/printer/binding.cc
@@ -54,7 +54,7 @@ TVM_STATIC_IR_FUNCTOR(IRDocsifier, vtable)
ExprDoc rhs = Relax(d, "match_cast")
->Call({d->AsDoc<ExprDoc>(n->value, n_p->Attr("value")),
d->AsDoc<ExprDoc>(n->ty, n_p->Attr("ty"))});
- ExprDoc lhs = DefineVar(n->var, d->frames.back(), d);
+ ExprDoc lhs = DefineRelaxVar(n->var, d->frames.back(), d);
return AssignDoc(lhs, rhs, ann);
});
@@ -63,11 +63,14 @@ TVM_STATIC_IR_FUNCTOR(IRDocsifier, vtable)
"", [](relax::VarBinding n, AccessPath n_p, IRDocsifier d) -> Doc {
if (const auto if_ = n->value.as<relax::IfNode>()) {
ffi::Optional<ExprDoc> ann = TypeAsAnn(n->var, n_p->Attr("var"), d, n->value);
- ExprDoc lhs = DefineVar(n->var, d->frames.back(), d);
+ if (!ann.has_value() && n->var->ty.as<PrimTypeNode>()) {
+ ann = d->AsDoc<ExprDoc>(n->var->ty, n_p->Attr("var")->Attr("ty"));
+ }
+ ExprDoc lhs = DefineRelaxVar(n->var, d->frames.back(), d);
return PrintIfExpr(ffi::GetRef<relax::If>(if_), n_p->Attr("value"), d, lhs, ann);
} else if (n->value->IsInstance<tvm::BaseFuncNode>() &&
!n->value->IsInstance<relax::ExternFuncNode>()) {
- IdDoc lhs = DefineVar(n->var, d->frames.back(), d);
+ IdDoc lhs = DefineRelaxVar(n->var, d->frames.back(), d);
d->cfg->binding_names.push_back(lhs->name);
Doc ret = d->AsDoc(n->value, n_p->Attr("value"));
d->cfg->binding_names.pop_back();
@@ -79,7 +82,10 @@ TVM_STATIC_IR_FUNCTOR(IRDocsifier, vtable)
} else {
ExprDoc rhs = d->AsDoc<ExprDoc>(n->value, n_p->Attr("value"));
ffi::Optional<ExprDoc> ann = TypeAsAnn(n->var, n_p->Attr("var"), d, n->value);
- ExprDoc lhs = DefineVar(n->var, d->frames.back(), d);
+ if (!ann.has_value() && n->var->ty.as<PrimTypeNode>()) {
+ ann = d->AsDoc<ExprDoc>(n->var->ty, n_p->Attr("var")->Attr("ty"));
+ }
+ ExprDoc lhs = DefineRelaxVar(n->var, d->frames.back(), d);
return AssignDoc(lhs, rhs, ann);
}
});
diff --git a/src/relax/script/printer/call.cc b/src/relax/script/printer/call.cc
index a44c92ff65..9b3b41e565 100644
--- a/src/relax/script/printer/call.cc
+++ b/src/relax/script/printer/call.cc
@@ -248,7 +248,7 @@ bool ShouldPrintAsTIR(const Call& call) {
if (!call->ty.as<PrimTypeNode>()) {
return false;
}
- if (call->op->ty.as<relax::FuncTypeNode>() || call->op.as<relax::VarNode>() ||
+ if (call->op->ty.as<relax::FuncTypeNode>() || call->op.as<tvm::VarNode>() ||
call->op.as<relax::FunctionNode>() || call->op.as<relax::ExternFuncNode>()) {
return false;
}
@@ -301,8 +301,7 @@ TVM_STATIC_IR_FUNCTOR(IRDocsifier, vtable)
prefix = IdDoc(name);
}
prefix->source_paths.push_back(n_p->Attr("op"));
- } else if (n->op->IsInstance<relax::VarNode>() ||
- n->op->IsInstance<tvm::GlobalVarNode>()) {
+ } else if (n->op->IsInstance<tvm::VarNode>() || n->op->IsInstance<tvm::GlobalVarNode>()) {
prefix = d->AsDoc<ExprDoc>(n->op, n_p->Attr("op"));
} else {
TVM_FFI_THROW(TypeError) << "Unsupported op: " << n->op->GetTypeKey();
diff --git a/src/relax/script/printer/dependent_type.cc b/src/relax/script/printer/dependent_type.cc
index 938a5a6d2c..87f158d94b 100644
--- a/src/relax/script/printer/dependent_type.cc
+++ b/src/relax/script/printer/dependent_type.cc
@@ -45,8 +45,8 @@ ExprDoc PrintShapeVar(const PrimExpr& e, const AccessPath& e_p, const IRDocsifie
bool func_var_mode = false;
if (f != nullptr) {
tirx::PostOrderVisit(e, [f, &func_var_mode](const ffi::ObjectRef& obj) -> void {
- if (const auto* var = obj.as<tirx::VarNode>()) {
- if (f->func_vars->count(var)) {
+ if (auto var = obj.as<tirx::PrimVar>()) {
+ if (f->func_vars->count(var.value().get())) {
func_var_mode = true;
}
}
diff --git a/src/relax/script/printer/expr.cc b/src/relax/script/printer/expr.cc
index 229260cdf7..3a633259e6 100644
--- a/src/relax/script/printer/expr.cc
+++ b/src/relax/script/printer/expr.cc
@@ -22,6 +22,7 @@
#include <cmath>
#include <limits>
+#include "../../../tirx/script/printer/utils.h"
#include "./utils.h"
namespace tvm {
@@ -143,25 +144,32 @@ TVM_STATIC_IR_FUNCTOR(IRDocsifier, vtable)
return d->AddMetadata(n);
});
-Doc PrintRelaxVar(relax::Var n, AccessPath p, IRDocsifier d) {
+Doc PrintRelaxVar(tvm::Var n, AccessPath p, IRDocsifier d) {
if (!d->IsVarDefined(n)) {
ExprDoc ann = d->AsDoc<ExprDoc>(n->ty, p->Attr("ty"));
Frame f = d->frames.back();
- ExprDoc var = DefineVar(n, f, d);
+ ExprDoc var = DefineRelaxVar(n, f, d);
f->stmts.push_back(AssignDoc(var, std::nullopt, ann));
}
return d->GetVarDoc(n).value();
}
-TVM_STATIC_IR_FUNCTOR(IRDocsifier, vtable).set_dispatch<relax::Var>("", PrintRelaxVar);
-TVM_STATIC_IR_FUNCTOR(IRDocsifier, vtable).set_dispatch<relax::DataflowVar>("", PrintRelaxVar);
+TVM_STATIC_IR_FUNCTOR(IRDocsifier, vtable).set_dispatch<relax::DataflowVar>("relax", PrintRelaxVar);
+
+std::string ReprPrintVar(const ffi::ObjectRef& obj, const PrinterConfig& cfg) {
+ Var var = obj.as_or_throw<Var>();
+ if (var->ty.as<PrimTypeNode>() || var->ty.as<PointerTypeNode>()) {
+ return ReprPrintTIR(obj, cfg);
+ }
+ return ReprPrintRelax(obj, cfg);
+}
TVM_REGISTER_SCRIPT_AS_REPR(relax::StringImmNode, ReprPrintRelax);
TVM_REGISTER_SCRIPT_AS_REPR(relax::DataTypeImmNode, ReprPrintRelax);
TVM_REGISTER_SCRIPT_AS_REPR(relax::TupleNode, ReprPrintRelax);
TVM_REGISTER_SCRIPT_AS_REPR(relax::TupleGetItemNode, ReprPrintRelax);
TVM_REGISTER_SCRIPT_AS_REPR(relax::ShapeExprNode, ReprPrintRelax);
-TVM_REGISTER_SCRIPT_AS_REPR(relax::VarNode, ReprPrintRelax);
+TVM_REGISTER_SCRIPT_AS_REPR(VarNode, ReprPrintVar);
TVM_REGISTER_SCRIPT_AS_REPR(relax::DataflowVarNode, ReprPrintRelax);
TVM_REGISTER_SCRIPT_AS_REPR(relax::ConstantNode, ReprPrintRelax);
diff --git a/src/relax/script/printer/function.cc b/src/relax/script/printer/function.cc
index d395131710..fafdad135b 100644
--- a/src/relax/script/printer/function.cc
+++ b/src/relax/script/printer/function.cc
@@ -50,7 +50,7 @@ TVM_FFI_STATIC_INIT_BLOCK() { RelaxFrameNode::RegisterReflection(); }
TVM_STATIC_IR_FUNCTOR(IRDocsifier, vtable)
.set_dispatch<relax::Function>("", [](relax::Function n, AccessPath n_p, IRDocsifier d) -> Doc {
- std::unordered_set<const tirx::VarNode*> func_vars;
+ std::unordered_set<const VarNode*> func_vars;
With<RelaxFrame> f(d);
IdDoc func_name("");
@@ -64,23 +64,19 @@ TVM_STATIC_IR_FUNCTOR(IRDocsifier, vtable)
(*f)->AddDispatchToken(d, "relax");
(*f)->is_func = true;
(*f)->func_vars = &func_vars;
- // Step 1. Print the return type
- ffi::Optional<ExprDoc> ret_type = std::nullopt;
- if (const auto& func_ty = relax::MatchType<relax::FuncType>(n)) {
- ret_type = d->AsDoc<ExprDoc>(func_ty.value()->ret, //
- n_p->Attr("ty")->Attr("ret"));
- }
- // Step 2. Print params
+ // Step 1. Print params
ffi::Array<AssignDoc> params;
{
AccessPath params_p = n_p->Attr("params");
for (int i = 0, l = n->params.size(); i < l; ++i) {
params.push_back(AssignDoc(
- /*lhs=*/DefineVar(n->params[i], *f, d),
+ /*lhs=*/DefineRelaxVar(n->params[i], *f, d),
/*rhs=*/std::nullopt,
TypeAsAnn(n->params[i], params_p->ArrayItem(i), d, std::nullopt)));
}
}
+ // Step 2. Print the return type
+ ffi::Optional<ExprDoc> ret_type = d->AsDoc<ExprDoc>(n->ret_ty, n_p->Attr("ret_ty"));
// Step 3. Clean up func variables
(*f)->func_vars = nullptr;
// Step 4. Print attributes
diff --git a/src/relax/script/printer/tir.cc b/src/relax/script/printer/tir.cc
index cea616e21d..4b989958a5 100644
--- a/src/relax/script/printer/tir.cc
+++ b/src/relax/script/printer/tir.cc
@@ -26,14 +26,13 @@ namespace tvm {
namespace script {
namespace printer {
-/*! \brief Find the outmost Relax function frame. If not exist, the outmost Relax frame. */
+/*! \brief Find the innermost Relax function frame, or the innermost Relax frame. */
RelaxFrameNode* GetRelaxFrame(IRDocsifier d) {
RelaxFrameNode* f = nullptr;
- for (const Frame& frame : d->frames) {
- if (const auto* relax_frame = frame.as<RelaxFrameNode>()) {
+ for (auto it = d->frames.rbegin(); it != d->frames.rend(); ++it) {
+ if (const auto* relax_frame = (*it).as<RelaxFrameNode>()) {
if (relax_frame->is_func) {
- f = const_cast<RelaxFrameNode*>(relax_frame);
- break;
+ return const_cast<RelaxFrameNode*>(relax_frame);
} else if (f == nullptr) {
f = const_cast<RelaxFrameNode*>(relax_frame);
}
@@ -42,13 +41,16 @@ RelaxFrameNode* GetRelaxFrame(IRDocsifier d) {
return f;
}
-Doc PrintTIRVar(tirx::Var n, AccessPath n_p, IRDocsifier d) {
- PrimType n_ty = n->ty.as_or_throw<PrimType>();
- TVM_FFI_CHECK(!n_ty.IsScalableVector() && !n_ty.IsFixedLengthVector(), TypeError)
- << "Relax only uses scalar TIR variables,"
- << "but received TIR variable " << n << " with dtype " << n_ty->dtype;
-
+Doc PrintCanonicalVar(Var n, AccessPath n_p, IRDocsifier d) {
+ if (!n->ty.as<PrimTypeNode>()) {
+ return PrintRelaxVar(n, n_p, d);
+ }
+ tirx::PrimVar prim_var = n.as_or_throw<tirx::PrimVar>();
if (!d->IsVarDefined(n)) {
+ PrimType n_ty = n->ty.as_or_throw<PrimType>();
+ TVM_FFI_CHECK(!n_ty.IsScalableVector() && !n_ty.IsFixedLengthVector(), TypeError)
+ << "Relax only uses scalar TIR variables,"
+ << "but received TIR variable " << n << " with dtype " << n_ty->dtype;
RelaxFrameNode* f = GetRelaxFrame(d);
// There should be at least one Relax frame
if (f == nullptr) {
@@ -62,7 +64,7 @@ Doc PrintTIRVar(tirx::Var n, AccessPath n_p, IRDocsifier d) {
}
IdDoc var = d->Define(n, ffi::GetRef<Frame>(f), n->name_hint.empty() ? "v" : n->name_hint);
var->source_paths.push_back(n_p);
- f->stmts.push_back(AssignDoc(var, PrintVarCreation(n, n_p, d), std::nullopt));
+ f->stmts.push_back(AssignDoc(var, PrintVarCreation(prim_var, n_p, d), std::nullopt));
}
if (ffi::Optional<ExprDoc> doc = d->GetVarDoc(n)) {
return doc.value();
@@ -71,7 +73,7 @@ Doc PrintTIRVar(tirx::Var n, AccessPath n_p, IRDocsifier d) {
TVM_FFI_UNREACHABLE();
}
-TVM_STATIC_IR_FUNCTOR(IRDocsifier, vtable).set_dispatch<tirx::Var>("relax", PrintTIRVar);
+TVM_STATIC_IR_FUNCTOR(IRDocsifier, vtable).set_dispatch<Var>("relax", PrintCanonicalVar);
TVM_STATIC_IR_FUNCTOR(IRDocsifier, vtable)
.set_dispatch<tvm::IntImm>( //
diff --git a/src/relax/script/printer/utils.h b/src/relax/script/printer/utils.h
index b5d575666a..646f8157f3 100644
--- a/src/relax/script/printer/utils.h
+++ b/src/relax/script/printer/utils.h
@@ -41,7 +41,7 @@ class RelaxFrameNode : public FrameNode {
public:
bool is_func = false;
bool module_alias_printed = false;
- std::unordered_set<const tirx::VarNode*>* func_vars = nullptr;
+ std::unordered_set<const VarNode*>* func_vars = nullptr;
static void RegisterReflection() {
namespace refl = tvm::ffi::reflection;
@@ -74,11 +74,11 @@ inline std::string ReprPrintRelax(const ffi::ObjectRef& obj, const PrinterConfig
return Docsify(obj, d, *f, cfg);
}
-inline IdDoc DefineVar(const relax::Var& var, const Frame& frame, const IRDocsifier& d) {
+inline IdDoc DefineRelaxVar(const tvm::Var& var, const Frame& frame, const IRDocsifier& d) {
return d->Define(var, frame, var->name_hint.empty() ? "v" : var->name_hint);
}
-inline ffi::Optional<ExprDoc> TypeAsAnn(const relax::Var& v, const AccessPath& v_p,
+inline ffi::Optional<ExprDoc> TypeAsAnn(const tvm::Var& v, const AccessPath& v_p,
const IRDocsifier& d,
const ffi::Optional<relax::Expr>& rhs) {
if (v->ty.IsMissing()) {
@@ -121,7 +121,7 @@ inline ffi::Optional<ExprDoc> TypeAsAnn(const relax::Var& v, const AccessPath& v
inferred_ty = ptr->fields[get_item->index];
}
- } else if (const auto* trivial_binding = rhs.as<relax::VarNode>()) {
+ } else if (const auto* trivial_binding = rhs.as<tvm::VarNode>()) {
inferred_ty = trivial_binding->ty.as<tvm::Type>();
}
@@ -135,6 +135,8 @@ inline ffi::Optional<ExprDoc> TypeAsAnn(const relax::Var& v, const AccessPath& v
ffi::Array<StmtDoc> PrintSeqExpr(const relax::SeqExpr& n, const AccessPath& n_p,
const IRDocsifier& d, bool use_ret);
+Doc PrintRelaxVar(tvm::Var n, AccessPath p, IRDocsifier d);
+
ExprDoc PrintShapeVar(const PrimExpr& e, const AccessPath& e_p, const IRDocsifier& d);
inline int FindVDeviceIndexByTargetKind(const VDevice& vdevice, const IRDocsifier& d) {
diff --git a/src/relax/training/utils.cc b/src/relax/training/utils.cc
index a279972e95..aca5ed16a9 100644
--- a/src/relax/training/utils.cc
+++ b/src/relax/training/utils.cc
@@ -178,8 +178,10 @@ class AppendLossMutator : private ExprMutator {
*
* 1. Is used in prediction_outputs of the backbone function,
* 2. Is not used in other_outputs of the backbone function.
+ * 3. Does not have PrimType.
*
- * Because such Vars are no longer the outputs of the new function.
+ * Because such Vars are no longer the outputs of the new function. Primitive Vars remain
+ * ordinary bindings because they may define dependent types in the appended loss function.
*/
void CheckAndRemapBackboneReturn() {
TVM_FFI_ICHECK(static_cast<int>(backbone_return_arr_.size()) >= num_backbone_outputs_)
@@ -189,7 +191,7 @@ class AppendLossMutator : private ExprMutator {
backbone_return_arr_.begin() + num_backbone_outputs_, backbone_return_arr_.end());
for (int i = 0; i < num_backbone_outputs_; ++i) {
auto var = backbone_return_arr_[i];
- if (other_outputs_var.count(var) == 0) {
+ if (other_outputs_var.count(var) == 0 && !var->ty.as<PrimTypeNode>()) {
auto new_var = DataflowVar(var->name_hint, GetType(var), var->span);
this->var_remap_[var] = new_var;
backbone_return_arr_.Set(i, new_var);
diff --git a/src/relax/transform/bind_params.cc b/src/relax/transform/bind_params.cc
index 04d5f33c66..1866584041 100644
--- a/src/relax/transform/bind_params.cc
+++ b/src/relax/transform/bind_params.cc
@@ -24,81 +24,28 @@
#include <tvm/relax/expr_functor.h>
#include <tvm/relax/transform.h>
#include <tvm/relax/type.h>
-#include <tvm/tirx/op.h>
-#include <tuple>
#include <utility>
namespace tvm {
namespace relax {
-void MatchSymbolicVar(const Expr& arg, const Expr& constant,
- ffi::Map<tirx::Var, PrimExpr>* symbolic_var_map,
- arith::AnalyzerObj* analyzer_) {
- auto opt_arg_ty = MatchType<TensorType>(arg);
- TVM_FFI_ICHECK(opt_arg_ty)
- << "The type of the bound parameter is expected to be TensorType, but got: " << GetType(arg);
- auto opt_const_ty = MatchType<TensorType>(constant);
- // As the constant is generated by internal codes, we use TVM_FFI_ICHECK here.
- TVM_FFI_ICHECK(opt_const_ty)
- << "The type of the bound weight is expected to be TensorType, but got: "
- << GetType(constant);
-
- TensorType arg_ty = opt_arg_ty.value();
- TensorType const_ty = opt_const_ty.value();
- TVM_FFI_ICHECK(!const_ty->IsUnknownDtype());
- TVM_FFI_ICHECK(!const_ty->IsUnknownNdim());
- TVM_FFI_ICHECK(const_ty->shape.has_value());
-
- // dtype mismatch
- if (!arg_ty->IsUnknownDtype() && arg_ty->dtype != const_ty->dtype) {
- TVM_FFI_THROW(InternalError) << "The dtype of the bound parameter is expected to be "
- << arg_ty->dtype << ", but got: " << const_ty->dtype;
- }
- // ndim mismatch
- if (!arg_ty->IsUnknownNdim() && arg_ty->ndim != const_ty->ndim) {
- TVM_FFI_THROW(InternalError) << "The ndim of the bound parameter is expected to be "
- << arg_ty->ndim << ", but got: " << const_ty->ndim;
- }
- if (!arg_ty->shape.has_value()) return;
- const auto* arg_shape = arg_ty->shape.value().as<ShapeExprNode>();
- const auto* const_shape = const_ty->shape.value().as<ShapeExprNode>();
-
- TVM_FFI_ICHECK(arg_shape && const_shape)
- << "The shape of the bound parameter and weight is expected to be ShapeExprNode for now";
-
- for (int i = 0; i < arg_ty->ndim; ++i) {
- const PrimExpr& const_dim = const_shape->values[i];
- TVM_FFI_ICHECK(tirx::is_const_int(const_dim));
- if (const auto* shape_var = arg_shape->values[i].as<tirx::VarNode>()) {
- auto it = symbolic_var_map->find(ffi::GetRef<tirx::Var>(shape_var));
- if (it == symbolic_var_map->end()) {
- symbolic_var_map->Set(ffi::GetRef<tirx::Var>(shape_var), const_dim);
- } else {
- TVM_FFI_ICHECK(analyzer_->CanProveEqual((*it).second, const_dim))
- << "The shape of the bound parameter is expected to be " << (*it).second
- << ", but got: " << const_dim;
- }
- }
- }
-}
-
-std::tuple<ffi::Map<Var, Expr>, ffi::Map<tirx::Var, PrimExpr>> NormalizeBindings(
- const Function& func, const ffi::Map<Any, ffi::ObjectRef>& untyped_params) {
+ffi::Map<Var, Expr> NormalizeBindings(const Function& func,
+ const ffi::Map<Any, ffi::ObjectRef>& untyped_params) {
TVM_FFI_ICHECK(func.defined());
TVM_FFI_ICHECK(untyped_params.defined());
// Map from string to the variable(s) with that name.
- std::unordered_map<std::string, ffi::Array<relax::Var>> string_lookup;
- std::unordered_set<const relax::VarNode*> var_set;
+ std::unordered_map<std::string, ffi::Array<tvm::Var>> string_lookup;
+ std::unordered_set<const tvm::VarNode*> var_set;
for (const auto& param : func->params) {
string_lookup[param->name_hint].push_back(param);
var_set.insert(param.get());
}
- ffi::Map<relax::Var, relax::Expr> relax_var_remap;
+ ffi::Map<tvm::Var, relax::Expr> relax_var_remap;
- auto normalize_key = [&](ffi::Any obj) -> relax::Var {
+ auto normalize_key = [&](ffi::Any obj) -> tvm::Var {
if (auto opt_str = obj.as<ffi::String>()) {
std::string str = opt_str.value();
auto it = string_lookup.find(str);
@@ -114,7 +61,7 @@ std::tuple<ffi::Map<Var, Expr>, ffi::Map<tirx::Var, PrimExpr>> NormalizeBindings
<< "Remap of variable " << var << " was defined multiple times";
return var;
- } else if (auto opt_var = obj.as<relax::Var>()) {
+ } else if (auto opt_var = obj.as<tvm::Var>()) {
auto var = opt_var.value();
TVM_FFI_ICHECK(!relax_var_remap.count(var))
<< "Remap of variable " << var << " was defined multiple times";
@@ -124,8 +71,8 @@ std::tuple<ffi::Map<Var, Expr>, ffi::Map<tirx::Var, PrimExpr>> NormalizeBindings
return var;
} else {
TVM_FFI_THROW(InternalError)
- << "Expected bound parameter to be a relax::Var, "
- << " or a string that uniquely identifies a relax::Var param within the function. "
+ << "Expected bound parameter to be a tvm::Var, "
+ << " or a string that uniquely identifies a tvm::Var param within the function. "
<< "However, received object " << obj << " of type " << obj.GetTypeKey();
}
};
@@ -145,13 +92,7 @@ std::tuple<ffi::Map<Var, Expr>, ffi::Map<tirx::Var, PrimExpr>> NormalizeBindings
}
arith::Analyzer analyzer;
- ffi::Map<tirx::Var, PrimExpr> symbolic_var_map = InferSymbolicVarMap(relax_var_remap, analyzer);
-
- // for (const auto& [bind_param, bind_expr] : relax_var_remap) {
- // MatchSymbolicVar(bind_param, bind_expr, &symbolic_var_map, &analyzer);
- // }
-
- return {relax_var_remap, symbolic_var_map};
+ return InferSymbolicVarMap(relax_var_remap, analyzer);
}
/*!
@@ -161,9 +102,8 @@ std::tuple<ffi::Map<Var, Expr>, ffi::Map<tirx::Var, PrimExpr>> NormalizeBindings
* \return Function
*/
Function FunctionBindParams(Function func, const ffi::Map<Any, ffi::ObjectRef>& untyped_params) {
- auto [bind_dict, symbolic_var_map] = NormalizeBindings(func, untyped_params);
-
- Expr bound_expr = Bind(func, bind_dict, symbolic_var_map);
+ ffi::Map<Var, Expr> bind_dict = NormalizeBindings(func, untyped_params);
+ Expr bound_expr = Bind(func, bind_dict);
return bound_expr.as_or_throw<Function>();
}
diff --git a/src/relax/transform/bind_symbolic_vars.cc b/src/relax/transform/bind_symbolic_vars.cc
index c6d5a61722..b0ed4a3674 100644
--- a/src/relax/transform/bind_symbolic_vars.cc
+++ b/src/relax/transform/bind_symbolic_vars.cc
@@ -31,7 +31,7 @@ namespace tvm {
namespace relax {
Function FunctionBindSymbolicVars(
- Function func, ffi::Map<ffi::Variant<tirx::Var, ffi::String>, PrimExpr> obj_remap) {
+ Function func, ffi::Map<ffi::Variant<tirx::PrimVar, ffi::String>, PrimExpr> obj_remap) {
// Early bail-out if no updates need to be made.
if (obj_remap.empty()) {
return func;
@@ -48,7 +48,7 @@ Function FunctionBindSymbolicVars(
}
// Replacement map to be used when rewriting the function.
- ffi::Map<tirx::Var, PrimExpr> var_remap;
+ ffi::Map<Var, Expr> var_remap;
for (const auto& [key, replacement] : obj_remap) {
if (auto opt = key.as<ffi::String>()) {
ffi::String string_key = opt.value();
@@ -65,7 +65,7 @@ Function FunctionBindSymbolicVars(
TVM_FFI_ICHECK(!var_remap.count(var))
<< "Remap of variable " << var << " was defined multiple times";
var_remap.Set(var, replacement);
- } else if (auto opt = key.as<tirx::Var>()) {
+ } else if (auto opt = key.as<tirx::PrimVar>()) {
auto var = opt.value();
TVM_FFI_ICHECK(!var_remap.count(var))
@@ -81,7 +81,7 @@ Function FunctionBindSymbolicVars(
}
}
- auto new_func = Bind(func, {}, var_remap).as_or_throw<Function>();
+ auto new_func = Bind(func, var_remap).as_or_throw<Function>();
auto free_symbolic_vars = FreeSymbolicVars(new_func);
@@ -94,7 +94,7 @@ Function FunctionBindSymbolicVars(
namespace {
IRModule ModuleBindSymbolicVars(
- IRModule mod, ffi::Map<ffi::Variant<tirx::Var, ffi::String>, PrimExpr> binding_map) {
+ IRModule mod, ffi::Map<ffi::Variant<tirx::PrimVar, ffi::String>, PrimExpr> binding_map) {
std::unordered_set<ffi::Any, ffi::AnyHash, ffi::AnyEqual> used;
IRModule updates;
for (const auto& [gvar, base_func] : mod->functions) {
@@ -102,7 +102,8 @@ IRModule ModuleBindSymbolicVars(
auto func = opt.value();
// Collect bindings that are used by this function.
- auto func_binding_map = [&]() -> ffi::Map<ffi::Variant<tirx::Var, ffi::String>, PrimExpr> {
+ auto func_binding_map =
+ [&]() -> ffi::Map<ffi::Variant<tirx::PrimVar, ffi::String>, PrimExpr> {
std::unordered_set<std::string> var_names;
std::unordered_set<const tirx::VarNode*> vars;
for (const auto& var : DefinedSymbolicVars(func)) {
@@ -110,13 +111,13 @@ IRModule ModuleBindSymbolicVars(
vars.insert(var.get());
}
- ffi::Map<ffi::Variant<tirx::Var, ffi::String>, PrimExpr> out;
+ ffi::Map<ffi::Variant<tirx::PrimVar, ffi::String>, PrimExpr> out;
for (const auto& [key, replacement] : binding_map) {
bool used_by_function = false;
if (auto opt = key.as<ffi::String>()) {
used_by_function = var_names.count(opt.value());
- } else if (auto ptr = key.as<tirx::VarNode>()) {
- used_by_function = vars.count(ptr);
+ } else if (auto var = key.as<tirx::PrimVar>()) {
+ used_by_function = vars.count(var.value().get());
} else {
TVM_FFI_THROW(InternalError)
<< "Expected symbolic variable to be a tirx::Var "
@@ -161,7 +162,7 @@ TVM_FFI_STATIC_INIT_BLOCK() {
namespace transform {
-Pass BindSymbolicVars(ffi::Map<ffi::Variant<tirx::Var, ffi::String>, PrimExpr> binding_map,
+Pass BindSymbolicVars(ffi::Map<ffi::Variant<tirx::PrimVar, ffi::String>, PrimExpr> binding_map,
ffi::Optional<ffi::String> func_name) {
auto pass_func = [=](IRModule mod, PassContext context) -> IRModule {
if (func_name) {
diff --git a/src/relax/transform/bundle_model_params.cc b/src/relax/transform/bundle_model_params.cc
index 7aa4c8acbe..c2de36d9d3 100644
--- a/src/relax/transform/bundle_model_params.cc
+++ b/src/relax/transform/bundle_model_params.cc
@@ -30,6 +30,8 @@
#include <tvm/relax/expr_functor.h>
#include <tvm/relax/transform.h>
+#include <unordered_set>
+
#include "utils.h"
namespace tvm {
@@ -57,9 +59,34 @@ class ModelParamBundler : public ExprMutator {
params.push_back(func->params[i]);
}
+ std::unordered_set<const VarNode*> signature_vars;
+ for (const tirx::Var& var : DefinableTIRVarsInType(TupleType(params.Map(GetType)))) {
+ signature_vars.insert(var.get());
+ }
+
+ std::unordered_set<const VarNode*> bundled_prim_params;
+ ffi::Array<Var> bundled_prim_params_in_order;
+ for (size_t i = num_input; i < func->params.size(); i++) {
+ if (func->params[i].as<tirx::PrimVar>()) {
+ bundled_prim_params.insert(func->params[i].get());
+ bundled_prim_params_in_order.push_back(func->params[i]);
+ }
+ }
+
+ auto erase_removed_prim_params = [&](const Type& type) {
+ return EraseToWellDefined(type, [&](const Var& var) -> ffi::Optional<Expr> {
+ if (bundled_prim_params.count(var.get()) && !signature_vars.count(var.get())) {
+ return std::nullopt;
+ }
+ return var;
+ });
+ };
+
ffi::Array<Type> param_tuple;
for (size_t i = num_input; i < func->params.size(); i++) {
- param_tuple.push_back(GetType(func->params[i]));
+ Type field_type = erase_removed_prim_params(GetType(func->params[i]));
+ param_tuple.push_back(field_type);
+ var_to_field_type_.Set(func->params[i], field_type);
}
Var var_param_tuple(param_tuple_name_.value_or("model_params"), TupleType(param_tuple));
@@ -69,23 +96,73 @@ class ModelParamBundler : public ExprMutator {
var_to_expr_.Set(func->params[i], TupleGetItem(var_param_tuple, i - num_input));
}
- func.CopyOnWrite()->params = params;
+ Type ret_ty = erase_removed_prim_params(func->ret_ty);
+ bool previous_rewrite_model_params = rewrite_model_params_;
+ ffi::Array<Var> previous_pending_prim_params = std::move(pending_prim_params_);
+ rewrite_model_params_ = true;
+ pending_prim_params_ = std::move(bundled_prim_params_in_order);
+ Expr body = VisitWithNewScope(func->body, params);
+ rewrite_model_params_ = previous_rewrite_model_params;
+ pending_prim_params_ = std::move(previous_pending_prim_params);
+ return Function(params, body, ret_ty, func->is_pure, func->attrs);
+ }
+
+ Expr VisitExpr_(const SeqExprNode* op) override {
+ if (pending_prim_params_.empty()) {
+ return ExprMutator::VisitExpr_(op);
+ }
+
+ ffi::Array<Var> prim_params = std::move(pending_prim_params_);
+ pending_prim_params_.clear();
- return ExprMutator::VisitExpr_(func.get());
+ builder_->BeginBindingBlock();
+ for (const Var& var : prim_params) {
+ auto it = var_to_expr_.find(var);
+ TVM_FFI_ICHECK(it != var_to_expr_.end());
+ var_remap_[var] = builder_->Emit((*it).second, var->name_hint);
+ }
+ BindingBlock prologue = builder_->EndBlock();
+
+ SeqExpr body = ExprMutator::VisitExpr_(op).as_or_throw<SeqExpr>();
+ ffi::Array<BindingBlock> blocks;
+ blocks.push_back(prologue);
+ for (const BindingBlock& block : body->blocks) {
+ blocks.push_back(block);
+ }
+ return SeqExpr(blocks, body->body);
}
Expr VisitExpr_(const VarNode* op) override {
auto var = ffi::GetRef<Var>(op);
- if (auto it = var_to_expr_.find(var); it != var_to_expr_.end()) {
- return builder_->Emit((*it).second, op->name_hint);
- } else {
+ if (!rewrite_model_params_) {
return ExprMutator::VisitExpr_(op);
}
+ if (auto it = var_to_expr_.find(var); it != var_to_expr_.end()) {
+ bool is_prim_param = var.as<tirx::PrimVar>().has_value();
+ if (is_prim_param) {
+ if (auto cached = var_remap_.find(var); cached != var_remap_.end()) {
+ return cached->second;
+ }
+ TVM_FFI_THROW(InternalError)
+ << "Bundled primitive parameters must be materialized in the function prologue";
+ }
+ auto field_type = var_to_field_type_.find(var);
+ TVM_FFI_ICHECK(field_type != var_to_field_type_.end());
+ Type rebound_type = VisitExprDepTypeField(GetType(var));
+ Var replacement = (*field_type).second.same_as(rebound_type)
+ ? builder_->Emit((*it).second, op->name_hint)
+ : builder_->EmitMatchCast((*it).second, rebound_type, op->name_hint);
+ return replacement;
+ }
+ return ExprMutator::VisitExpr_(op);
}
private:
ffi::Optional<ffi::String> param_tuple_name_;
ffi::Map<Var, Expr> var_to_expr_;
+ ffi::Map<Var, Type> var_to_field_type_;
+ ffi::Array<Var> pending_prim_params_;
+ bool rewrite_model_params_{false};
};
Function BundleModelParams(const Function& func, ffi::Optional<ffi::String> param_tuple_name) {
diff --git a/src/relax/transform/canonicalize_bindings.cc b/src/relax/transform/canonicalize_bindings.cc
index cc12031d5f..65851b0cb3 100644
--- a/src/relax/transform/canonicalize_bindings.cc
+++ b/src/relax/transform/canonicalize_bindings.cc
@@ -31,6 +31,7 @@
#include <tvm/relax/expr_functor.h>
#include <tvm/relax/transform.h>
#include <tvm/relax/type.h>
+#include <tvm/relax/utils.h>
#include <tvm/tirx/stmt_functor.h>
namespace tvm {
@@ -39,30 +40,63 @@ namespace relax {
namespace {
class SymbolicVarCanonicalizer : public ExprMutator {
+ class RuntimePrimVarCollector;
+
public:
+ static Expr Apply(Expr expr);
+
+ PrimExpr VisitTypePrimExprField(const PrimExpr& expr) final {
+ if (!canonicalize_shape_values_) return expr;
+ return CanonicalizeShapeValue(expr);
+ }
+
+ Expr VisitExpr_(const ShapeExprNode* op) final {
+ if (!canonicalize_shape_values_) return ffi::GetRef<Expr>(op);
+ ffi::Array<PrimExpr> values =
+ op->values.Map([this](const PrimExpr& value) { return CanonicalizeShapeValue(value); });
+ if (values.same_as(op->values)) {
+ return ffi::GetRef<Expr>(op);
+ }
+ return ShapeExpr(values, op->span);
+ }
+
Expr VisitExpr_(const FunctionNode* func) override {
auto cached = known_values_;
+ auto cached_runtime_uses = runtime_prim_var_uses_;
+ runtime_prim_var_uses_ = RuntimePrimVarCollector::Collect(func->body);
auto output = ExprMutator::VisitExpr_(func);
known_values_ = cached;
+ runtime_prim_var_uses_ = cached_runtime_uses;
return output;
}
void VisitBinding_(const MatchCastNode* binding) override {
auto tir_var_map =
InferSymbolicVarMap({{binding->var, binding->value}}, builder_->GetAnalyzer());
- for (const auto& [tir_var, prim_expr] : tir_var_map) {
- if (auto it = known_values_.find(tir_var); it != known_values_.end()) {
+ bool has_runtime_use = false;
+ for (const auto& [var, value] : tir_var_map) {
+ if (var.same_as(binding->var)) continue;
+ auto tir_var = var.as<tirx::PrimVar>();
+ if (!tir_var) continue;
+ has_runtime_use = has_runtime_use || runtime_prim_var_uses_.count(*tir_var);
+ PrimExpr prim_expr = value.as_or_throw<PrimExpr>();
+ if (auto it = known_values_.find(tir_var.value()); it != known_values_.end()) {
TVM_FFI_CHECK(!builder_->GetAnalyzer()->CanProve(it->second.expr != prim_expr), ValueError)
<< "MatchCast statements must be consistent. "
<< "However, the definition of Relax variable " << it->second.source->var
- << " implies that TIR variable " << tir_var << " is " << it->second.expr
+ << " implies that TIR variable " << tir_var.value() << " is " << it->second.expr
<< ", while the later definition of Relax variable " << binding->var
- << " instead implies that TIR variable " << tir_var << " is " << prim_expr;
+ << " instead implies that TIR variable " << tir_var.value() << " is " << prim_expr;
} else {
- known_values_[tir_var] = KnownValue{prim_expr, ffi::GetRef<MatchCast>(binding)};
+ known_values_[tir_var.value()] = KnownValue{prim_expr, ffi::GetRef<MatchCast>(binding)};
}
}
+ // A MatchCast that defines a runtime-used symbolic variable cannot be folded away.
+ // Preserve its symbolic pattern while still canonicalizing downstream shape annotations.
+ bool cached = canonicalize_shape_values_;
+ canonicalize_shape_values_ = !has_runtime_use;
ExprMutator::VisitBinding_(binding);
+ canonicalize_shape_values_ = cached;
}
Expr VisitExpr_(const IfNode* op) override {
@@ -115,42 +149,62 @@ class SymbolicVarCanonicalizer : public ExprMutator {
return If(guard, true_b, false_b, op->span);
}
- PrimExpr VisitTypePrimExprField(const PrimExpr& expr) override {
- if (known_values_.empty()) {
- return expr;
- }
- PrimExpr output =
- tirx::Substitute(expr, [this](const tirx::Var& var) -> ffi::Optional<tvm::Expr> {
- if (auto it = known_values_.find(var); it != known_values_.end()) {
- return tvm::Expr(it->second.expr);
- } else {
- return std::nullopt;
- }
- });
- if (output.same_as(expr)) {
- return expr;
+ private:
+ class RuntimePrimVarCollector : public ExprVisitor {
+ public:
+ static std::unordered_set<tirx::Var, ffi::ObjectPtrHash, ffi::ObjectPtrEqual> Collect(
+ const Expr& expr) {
+ RuntimePrimVarCollector collector;
+ collector.VisitExpr(expr);
+ return collector.uses_;
}
- output = builder_->GetAnalyzer()->Simplify(output);
- return output;
- }
+ void VisitExprDepTypeField(const Type&) final {}
+
+ void VisitExpr_(const ShapeExprNode*) final {}
- Expr VisitExprFallback_(const ExprNode* op) final {
- if (op->ty.as<PrimTypeNode>()) {
- return VisitTypePrimExprField(ffi::GetRef<Expr>(op).as_or_throw<PrimExpr>());
+ void VisitExpr_(const VarNode* op) final {
+ Var var = ffi::GetRef<Var>(op);
+ if (auto prim_var = var.as<tirx::PrimVar>()) {
+ uses_.insert(*prim_var);
+ }
}
- return ExprMutator::VisitExprFallback_(op);
+
+ private:
+ std::unordered_set<tirx::Var, ffi::ObjectPtrHash, ffi::ObjectPtrEqual> uses_;
+ };
+
+ PrimExpr CanonicalizeShapeValue(const PrimExpr& expr) {
+ PrimExpr output = tirx::Substitute(expr, [this](const Var& var) -> ffi::Optional<Expr> {
+ auto prim_var = var.as<tirx::PrimVar>();
+ if (!prim_var) return std::nullopt;
+ auto it = known_values_.find(*prim_var);
+ if (it == known_values_.end()) return std::nullopt;
+ return it->second.expr;
+ });
+ return output.same_as(expr) ? expr : builder_->GetAnalyzer()->Simplify(output);
}
- private:
struct KnownValue {
PrimExpr expr;
MatchCast source;
};
- std::unordered_map<tirx::Var, KnownValue> known_values_;
+ std::unordered_map<tirx::Var, KnownValue, ffi::ObjectPtrHash, ffi::ObjectPtrEqual> known_values_;
+ std::unordered_set<tirx::Var, ffi::ObjectPtrHash, ffi::ObjectPtrEqual> runtime_prim_var_uses_;
+ bool canonicalize_shape_values_ = true;
};
+Expr SymbolicVarCanonicalizer::Apply(Expr expr) {
+ SymbolicVarCanonicalizer mutator;
+ if (!expr->IsInstance<FunctionNode>()) {
+ // Dataflow rewriting canonicalizes a SeqExpr directly, without entering
+ // VisitExpr_(FunctionNode), so initialize the same runtime-use context.
+ mutator.runtime_prim_var_uses_ = RuntimePrimVarCollector::Collect(expr);
+ }
+ return mutator(std::move(expr));
+}
+
struct CanonicalizationPlan {
ffi::Map<Var, Var> replace_usage;
ffi::Map<Var, Var> replace_binding;
@@ -565,7 +619,9 @@ class BindingCanonicalizer : public ExprMutator {
};
} // namespace
-Expr CanonicalizeTIRVariables(Expr expr) { return SymbolicVarCanonicalizer()(std::move(expr)); }
+Expr CanonicalizeTIRVariables(Expr expr) {
+ return SymbolicVarCanonicalizer::Apply(std::move(expr));
+}
Expr CanonicalizeRelaxBindings(Expr expr) { return BindingCanonicalizer::Apply(std::move(expr)); }
diff --git a/src/relax/transform/compute_prim_value.cc b/src/relax/transform/compute_prim_value.cc
index 49beccc88f..0d73745124 100644
--- a/src/relax/transform/compute_prim_value.cc
+++ b/src/relax/transform/compute_prim_value.cc
@@ -49,29 +49,54 @@ class PrimExprComputeInjector : public ExprMutator {
Expr VisitExpr_(const CallNode* op) final {
Call call = ffi::GetRef<Call>(op);
if (auto prim_expr = call.as<PrimExpr>()) {
- if (call->op.as<Op>()) {
- if (!HasRelaxCallCapabilities(op)) {
- return LiftPrimValue(prim_expr.value());
- }
+ if (call->op.as<Op>() && !HasRelaxCallCapabilities(op)) {
+ return LiftPrimValue(prim_expr.value());
}
}
return ExprMutator::VisitExpr_(op);
}
- Expr VisitExprFallback_(const ExprNode* op) final {
- Expr expr = ffi::GetRef<Expr>(op);
- if (auto prim_expr = expr.as<PrimExpr>()) {
- return LiftPrimValue(prim_expr.value());
- }
- return ExprMutator::VisitExprFallback_(op);
+#define RELAX_LIFT_PRIM_EXPR(OP) \
+ Expr VisitExpr_(const OP* op) final { \
+ return LiftPrimValue(ffi::GetRef<Expr>(op).as_or_throw<PrimExpr>()); \
}
+ RELAX_LIFT_PRIM_EXPR(tirx::BufferLoadNode);
+ RELAX_LIFT_PRIM_EXPR(tirx::AddNode);
+ RELAX_LIFT_PRIM_EXPR(tirx::SubNode);
+ RELAX_LIFT_PRIM_EXPR(tirx::MulNode);
+ RELAX_LIFT_PRIM_EXPR(tirx::DivNode);
+ RELAX_LIFT_PRIM_EXPR(tirx::ModNode);
+ RELAX_LIFT_PRIM_EXPR(tirx::FloorDivNode);
+ RELAX_LIFT_PRIM_EXPR(tirx::FloorModNode);
+ RELAX_LIFT_PRIM_EXPR(tirx::MinNode);
+ RELAX_LIFT_PRIM_EXPR(tirx::MaxNode);
+ RELAX_LIFT_PRIM_EXPR(tirx::EQNode);
+ RELAX_LIFT_PRIM_EXPR(tirx::NENode);
+ RELAX_LIFT_PRIM_EXPR(tirx::LTNode);
+ RELAX_LIFT_PRIM_EXPR(tirx::LENode);
+ RELAX_LIFT_PRIM_EXPR(tirx::GTNode);
+ RELAX_LIFT_PRIM_EXPR(tirx::GENode);
+ RELAX_LIFT_PRIM_EXPR(tirx::AndNode);
+ RELAX_LIFT_PRIM_EXPR(tirx::OrNode);
+ RELAX_LIFT_PRIM_EXPR(tirx::CastNode);
+ RELAX_LIFT_PRIM_EXPR(tirx::NotNode);
+ RELAX_LIFT_PRIM_EXPR(tirx::SelectNode);
+ RELAX_LIFT_PRIM_EXPR(tirx::RampNode);
+ RELAX_LIFT_PRIM_EXPR(tirx::BroadcastNode);
+ RELAX_LIFT_PRIM_EXPR(tirx::ShuffleNode);
+ RELAX_LIFT_PRIM_EXPR(tvm::IntImmNode);
+ RELAX_LIFT_PRIM_EXPR(tvm::FloatImmNode);
+ RELAX_LIFT_PRIM_EXPR(tirx::StringImmNode);
+
+#undef RELAX_LIFT_PRIM_EXPR
+
Expr VisitExpr_(const ShapeExprNode* op) final { return ffi::GetRef<Expr>(op); }
PrimExpr VisitTypePrimExprField(const PrimExpr& expr) final { return expr; }
Expr LiftPrimValue(const PrimExpr& node) {
- if (node->IsInstance<tirx::IntImmNode>() || node->IsInstance<tirx::VarNode>()) {
+ if (node->IsInstance<tirx::IntImmNode>() || node->IsInstance<VarNode>()) {
return node;
}
diff --git a/src/relax/transform/decompose_ops.cc b/src/relax/transform/decompose_ops.cc
index ed270d6e66..8d6da8c046 100644
--- a/src/relax/transform/decompose_ops.cc
+++ b/src/relax/transform/decompose_ops.cc
@@ -162,7 +162,7 @@ Expr TensorToShape(const Call& call_node, const BlockBuilder& builder) {
shape_var.push_back(tirx::Var("x", PrimType::Int(64)).as_or_throw<PrimExpr>());
}
// bind symbolic variables to the shape tuple
- relax::Var var("y", ShapeType(shape_var));
+ tvm::Var var("y", ShapeType(shape_var));
builder->EmitNormalized(MatchCast(var, call, ShapeType(shape_var)));
return ShapeExpr(shape_var);
}
diff --git a/src/relax/transform/fuse_ops.cc b/src/relax/transform/fuse_ops.cc
index 15d8493092..b76e129043 100644
--- a/src/relax/transform/fuse_ops.cc
+++ b/src/relax/transform/fuse_ops.cc
@@ -35,6 +35,7 @@
#include <tvm/relax/expr_functor.h>
#include <tvm/relax/transform.h>
#include <tvm/relax/type.h>
+#include <tvm/relax/utils.h>
#include <tvm/runtime/logging.h>
#include <tvm/tirx/analysis.h>
#include <tvm/tirx/expr_functor.h>
@@ -923,7 +924,7 @@ class OperatorFusor : public ExprMutator {
// - If the var's group is different with the binding's, the var must be the output from
// another group. Mark it to be the group output.
auto update_boundary = [this, binding, &cur_group](const Expr& e) {
- if (e->IsInstance<VarNode>()) {
+ if (e->IsInstance<VarNode>() && obj2group_.count(e.get())) {
const Var& used_var = e.as_or_throw<Var>();
Group* producer_group = GetGroupFromVar(used_var);
// Only check those group defined before.
diff --git a/src/relax/transform/fuse_tir.cc b/src/relax/transform/fuse_tir.cc
index 6aa6da3c1a..54f7e5af35 100644
--- a/src/relax/transform/fuse_tir.cc
+++ b/src/relax/transform/fuse_tir.cc
@@ -560,7 +560,7 @@ class FusedTIRConstructor : public ExprVisitor {
void VisitExpr_(const FunctionNode* func) final {
auto relax_to_tir_var_map =
RelaxToTIRVarMapCollector::Collect(mod_, ffi::GetRef<Function>(func));
- std::vector<ffi::Variant<tirx::Var, tirx::Buffer>> prim_func_params;
+ std::vector<ffi::Variant<tirx::PrimVar, tirx::Buffer>> prim_func_params;
for (const Var& relax_param : func->params) {
size_t size_before = prim_func_params.size();
CollectPrimFuncParams(relax_param, &prim_func_params, relax_to_tir_var_map.Get(relax_param));
@@ -583,8 +583,8 @@ class FusedTIRConstructor : public ExprVisitor {
// std::stable_sort is used instead of std::sort.
std::stable_sort(prim_func_params.begin(), prim_func_params.end(),
[](const auto& a, const auto& b) {
- bool a_is_var = a.template as<tirx::VarNode>();
- bool b_is_var = b.template as<tirx::VarNode>();
+ bool a_is_var = a.template as<tirx::PrimVar>().has_value();
+ bool b_is_var = b.template as<tirx::PrimVar>().has_value();
return a_is_var < b_is_var;
});
@@ -651,7 +651,7 @@ class FusedTIRConstructor : public ExprVisitor {
// Step 4. Append symbolic vars
for (const auto& param : prim_func_params) {
- if (auto var = param.as<tirx::Var>()) {
+ if (auto var = param.as<tirx::PrimVar>()) {
func_info_.params.push_back(var.value());
}
}
@@ -960,7 +960,7 @@ class FusedTIRConstructor : public ExprVisitor {
* \param out The vector into which to collect the params/buffers
*/
static void CollectPrimFuncParams(const Var& relax_param,
- std::vector<ffi::Variant<tirx::Var, tirx::Buffer>>* out,
+ std::vector<ffi::Variant<tirx::PrimVar, tirx::Buffer>>* out,
const ffi::Optional<tirx::Buffer>& tir_buffer_param) {
auto ty = GetType(relax_param);
@@ -985,15 +985,16 @@ class FusedTIRConstructor : public ExprVisitor {
}
out->push_back(std::move(buffer));
- } else if (const auto* prim_value = ty.as<PrimTypeNode>()) {
- // Case 2. The relax param is a scalar, we directly create a tirx var
- out->push_back(tirx::Var(name_hint, tvm::PrimType(prim_value->dtype)));
+ } else if (ty.as<PrimTypeNode>()) {
+ // Case 2. The relax param is a scalar, so its canonical Var is a TIR parameter.
+ out->push_back(relax_param.as_or_throw<tirx::PrimVar>());
} else if (const auto* shape_expr = ty.as<ShapeTypeNode>()) {
// Case 3. The relax param is a tuple of scalars, each represented as a tirx var
for (const auto& var : shape_expr->values.value()) {
- TVM_FFI_ICHECK(var->IsInstance<tirx::VarNode>());
- out->push_back(var.as_or_throw<tirx::Var>());
+ auto prim_var = var.as<tirx::PrimVar>();
+ TVM_FFI_ICHECK(prim_var.has_value());
+ out->push_back(prim_var.value());
}
} else {
TVM_FFI_THROW(TypeError) << "The param type of PrimFunc is expected to be "
@@ -1257,18 +1258,15 @@ class TIRFuseMutator : public ExprMutator {
TVM_FFI_ICHECK(shape->values.has_value())
<< "FuseTIR requires all shape input has ty value.";
for (const PrimExpr& prim_value : shape->values.value()) {
- TVM_FFI_ICHECK(prim_value->IsInstance<tirx::VarNode>())
+ TVM_FFI_ICHECK(prim_value.as<tirx::PrimVar>())
<< "All shape inputs are expected to be single tirx var.";
tir_vars.push_back(prim_value);
}
- } else if (const auto* prim_value = ty.as<PrimTypeNode>()) {
- if (const auto* var = arg.as<VarNode>()) {
- tir_vars.push_back(
- tirx::Var(var->name_hint, tvm::PrimType(prim_value->dtype)).as_or_throw<PrimExpr>());
- } else if (auto literal = arg.as<PrimExpr>()) {
+ } else if (ty.as<PrimTypeNode>()) {
+ if (auto literal = arg.as<PrimExpr>()) {
tir_vars.push_back(literal.value());
} else {
- TVM_FFI_THROW(TypeError) << "FuseTIR expects scalar arguments to be PrimExpr or Var, "
+ TVM_FFI_THROW(TypeError) << "FuseTIR expects scalar arguments to be PrimExpr, "
<< "but received " << arg;
}
diff --git a/src/relax/transform/lazy_transform_params.cc b/src/relax/transform/lazy_transform_params.cc
index 02fd26c05f..027d796245 100644
--- a/src/relax/transform/lazy_transform_params.cc
+++ b/src/relax/transform/lazy_transform_params.cc
@@ -73,14 +73,13 @@ class LazyInputMutator : public ExprMutator {
auto array_externally_visible_vars = DefinableTIRVarsInType(TupleType(new_params.Map(GetType)));
std::unordered_set<tirx::Var> externally_visible_vars(array_externally_visible_vars.begin(),
array_externally_visible_vars.end());
- Type new_ret_ty =
- EraseToWellDefined(func->ret_ty, [&](const tirx::Var& var) -> ffi::Optional<PrimExpr> {
- if (externally_visible_vars.count(var)) {
- return var.as_or_throw<PrimExpr>();
- } else {
- return std::nullopt;
- }
- });
+ Type new_ret_ty = EraseToWellDefined(func->ret_ty, [&](const Var& var) -> ffi::Optional<Expr> {
+ if (auto prim_var = var.as<tirx::PrimVar>();
+ prim_var && externally_visible_vars.count(prim_var.value())) {
+ return prim_var.value().as_or_throw<PrimExpr>();
+ }
+ return std::nullopt;
+ });
auto node = ffi::GetRef<Function>(func);
node.CopyOnWrite()->params = new_params;
diff --git a/src/relax/transform/lift_transform_params.cc b/src/relax/transform/lift_transform_params.cc
index 3907219968..0a4214b878 100644
--- a/src/relax/transform/lift_transform_params.cc
+++ b/src/relax/transform/lift_transform_params.cc
@@ -65,12 +65,10 @@ struct BaseCollectInfo {
* model weights, and computed tensors that require neither model
* weights nor runtime arguments (e.g. `R.zeros([16], "float16")`).
*/
- std::unordered_set<ffi::Variant<relax::Var, tirx::Var>, ffi::ObjectPtrHash, ffi::ObjectPtrEqual>
- requires_compile_time_param;
+ std::unordered_set<Var, ffi::ObjectPtrHash, ffi::ObjectPtrEqual> requires_compile_time_param;
/*! \brief Variables that are required at runtime */
- std::unordered_set<ffi::Variant<relax::Var, tirx::Var>, ffi::ObjectPtrHash, ffi::ObjectPtrEqual>
- required_at_runtime;
+ std::unordered_set<Var, ffi::ObjectPtrHash, ffi::ObjectPtrEqual> required_at_runtime;
protected:
ffi::Array<Var> GetCompileTimeOutputsHelper(const ffi::Array<Var>& params) const {
@@ -137,9 +135,7 @@ struct GlobalCollectInfo : public BaseCollectInfo {
// The parameters of the compile-time function.
ffi::Array<Var> params;
// The cross-function mapping between variables.
- ffi::Map<relax::Var, Expr> var_remap;
- // The cross-function between between TIR variables.
- ffi::Map<tirx::Var, PrimExpr> tir_var_remap;
+ ffi::Map<Var, Expr> var_remap;
ffi::Array<tirx::Var> GetPropagatedSymbolicVariables() const {
auto vars_from_original_params = DefinableTIRVarsInType(TupleType(params.Map(GetType)));
auto vars_from_transformed_params = [&]() -> std::unordered_set<tirx::Var> {
@@ -222,24 +218,29 @@ struct LocalCollectInfo : public BaseCollectInfo {
// serve as the parameter. This trivial binding will later be
// removed with CanonicalizeBindings.
ffi::Array<Var> params = GetRuntimeInputs();
+ ffi::Array<tirx::Var> local_tir_vars = GetPropagatedSymbolicVariables();
+ ffi::Array<Var> local_outputs = GetCompileTimeOutputs();
+ ffi::Map<Var, Var> global_to_local;
+ if (global_info) {
+ auto record_reverse = [&](const Var& local_var) {
+ if (auto it = global_info->var_remap.find(local_var); it != global_info->var_remap.end()) {
+ global_to_local.Set((*it).second.as_or_throw<Var>(), local_var);
+ }
+ };
+ for (const auto& var : local_tir_vars) record_reverse(var);
+ for (const auto& var : local_outputs) record_reverse(var);
+ }
+
auto propagated_tir_vars = [&]() {
- ffi::Array<tirx::Var> local_tir_vars = GetPropagatedSymbolicVariables();
if (!global_info) {
return local_tir_vars;
}
// When global lifting is enabled, the compile-time outputs are the global outputs, but the
// variables in the global outputs to the local variables.
- ffi::Map<tirx::Var, tirx::Var> reverse_map;
- for (const auto& var : local_tir_vars) {
- if (auto it = global_info->tir_var_remap.find(var);
- it != global_info->tir_var_remap.end()) {
- reverse_map.Set((*it).second.as_or_throw<tirx::Var>(), var);
- }
- }
ffi::Array<tirx::Var> global_tir_vars = global_info->GetPropagatedSymbolicVariables();
- global_tir_vars = global_tir_vars.Map([&](const tirx::Var& var) {
- if (auto it = reverse_map.find(var); it != reverse_map.end()) {
- return (*it).second.as_or_throw<tirx::Var>();
+ global_tir_vars = global_tir_vars.Map([&](const tirx::Var& var) -> tirx::Var {
+ if (auto it = global_to_local.find(var); it != global_to_local.end()) {
+ return (*it).second.as_or_throw<tirx::PrimVar>();
} else {
// This is the case when the some of the outputs of the shared transform is not used in
// this function.
@@ -249,27 +250,20 @@ struct LocalCollectInfo : public BaseCollectInfo {
return global_tir_vars;
}();
if (propagated_tir_vars.size()) {
- ShapeType shape_ty(
- propagated_tir_vars.Map([](tirx::Var var) { return var.as_or_throw<PrimExpr>(); }));
+ ShapeType shape_ty(propagated_tir_vars.Map(
+ [](tirx::Var var) { return var.as_or_throw<tirx::PrimVar>().as_or_throw<PrimExpr>(); }));
Var shape_expr("vars_from_compile_time_params", shape_ty);
params.push_back(shape_expr);
}
ffi::Array<Var> compile_time_outputs = [&]() {
- ffi::Array<Var> local_outputs = GetCompileTimeOutputs();
if (!global_info) {
return local_outputs;
}
// When global lifting is enabled, the compile-time outputs are the global outputs, but the
// variables in the global outputs to the local variables.
- ffi::Map<Var, Var> reverse_map;
- for (const auto& var : local_outputs) {
- if (auto it = global_info->var_remap.find(var); it != global_info->var_remap.end()) {
- reverse_map.Set((*it).second.as_or_throw<Var>(), var);
- }
- }
ffi::Array<Var> global_outputs = global_info->GetCompileTimeOutputs();
global_outputs = global_outputs.Map([&](const Var& var) {
- if (auto it = reverse_map.find(var); it != reverse_map.end()) {
+ if (auto it = global_to_local.find(var); it != global_to_local.end()) {
return (*it).second.as_or_throw<Var>();
} else {
// This is the case when the some of the outputs of the shared transform is not used in
@@ -376,8 +370,7 @@ class BaseLiftableBindingCollector : public ExprVisitor {
return true;
}
- std::unordered_set<ffi::Variant<Var, tirx::Var>, ffi::ObjectPtrHash, ffi::ObjectPtrEqual>
- liftable_vars_;
+ std::unordered_set<Var, ffi::ObjectPtrHash, ffi::ObjectPtrEqual> liftable_vars_;
bool is_in_dataflow_block_{false};
};
@@ -389,37 +382,23 @@ class LocalLiftableBindingCollector : public BaseLiftableBindingCollector {
visitor.info_.orig_func = func;
auto set_union =
- [&](std::unordered_set<ffi::Variant<relax::Var, tirx::Var>, ffi::ObjectPtrHash,
- ffi::ObjectPtrEqual>& target_set,
- const std::unordered_set<ffi::Variant<relax::Var, tirx::Var>, ffi::ObjectPtrHash,
- ffi::ObjectPtrEqual>& source_set,
- const ffi::Map<relax::Var, Expr>& var_remap,
- const ffi::Map<tirx::Var, PrimExpr>& tir_var_remap) {
- // In-place update the set in global info by unioning with the local set, variable
- // mappings are applied.
- for (const auto& relax_or_tir_var : source_set) {
- if (relax_or_tir_var.as<relax::VarNode>()) {
- if (auto it = var_remap.find(relax_or_tir_var.get<Var>()); it != var_remap.end()) {
- target_set.insert((*it).second.as_or_throw<relax::Var>());
- } else {
- target_set.insert(relax_or_tir_var.get<relax::Var>());
- }
+ [&](std::unordered_set<Var, ffi::ObjectPtrHash, ffi::ObjectPtrEqual>& target_set,
+ const std::unordered_set<Var, ffi::ObjectPtrHash, ffi::ObjectPtrEqual>& source_set,
+ const ffi::Map<Var, Expr>& var_remap) {
+ for (const Var& var : source_set) {
+ if (auto it = var_remap.find(var); it != var_remap.end()) {
+ target_set.insert((*it).second.as_or_throw<Var>());
} else {
- if (auto it = tir_var_remap.find(relax_or_tir_var.get<tirx::Var>());
- it != tir_var_remap.end()) {
- target_set.insert((*it).second.as_or_throw<tirx::Var>());
- } else {
- target_set.insert(relax_or_tir_var.get<tirx::Var>());
- }
+ target_set.insert(var);
}
}
};
if (global_info) {
set_union(global_info->requires_compile_time_param, visitor.info_.requires_compile_time_param,
- global_info->var_remap, global_info->tir_var_remap);
+ global_info->var_remap);
set_union(global_info->required_at_runtime, visitor.info_.required_at_runtime,
- global_info->var_remap, global_info->tir_var_remap);
+ global_info->var_remap);
}
return visitor.info_;
}
@@ -507,8 +486,7 @@ class LocalLiftableBindingCollector : public BaseLiftableBindingCollector {
/*! \brief Visitor to find the correspondence between parameters in multiple functions. */
class ParamRemapper : private ExprFunctor<void(const Expr&, const Expr&)> {
public:
- static std::pair<ffi::Map<Var, Expr>, ffi::Map<tirx::Var, PrimExpr>> GetParamMapping(
- const ffi::Array<Function>& functions) {
+ static ffi::Map<Var, Expr> GetParamMapping(const ffi::Array<Function>& functions) {
ParamRemapper mapper;
if (functions.size()) {
auto num_inputs_0 = functions[0]->GetAttr<int64_t>(attr::kNumInput).value();
@@ -528,7 +506,7 @@ class ParamRemapper : private ExprFunctor<void(const Expr&, const Expr&)> {
}
}
}
- return {mapper.var_remap_, mapper.tir_var_remap_};
+ return mapper.var_remap_;
}
private:
@@ -546,24 +524,22 @@ class ParamRemapper : private ExprFunctor<void(const Expr&, const Expr&)> {
auto rhs_tir_vars = DefinableTIRVarsInType(GetType(rhs_expr));
TVM_FFI_ICHECK_EQ(lhs_tir_vars.size(), rhs_tir_vars.size());
for (size_t i = 0; i < lhs_tir_vars.size(); i++) {
- if (auto it = tir_var_remap_.find(lhs_tir_vars[i]); it != tir_var_remap_.end()) {
+ if (auto it = var_remap_.find(lhs_tir_vars[i]); it != var_remap_.end()) {
TVM_FFI_ICHECK((*it).second.same_as(rhs_tir_vars[i]));
} else {
- tir_var_remap_.Set(lhs_tir_vars[i], rhs_tir_vars[i].as_or_throw<PrimExpr>());
+ var_remap_.Set(lhs_tir_vars[i], rhs_tir_vars[i]);
}
}
}
ffi::Map<Var, Expr> var_remap_;
- ffi::Map<tirx::Var, PrimExpr> tir_var_remap_;
};
class GlobalLiftableBindingCollector : public BaseLiftableBindingCollector {
public:
static GlobalCollectInfo Collect(const ffi::Array<Function>& functions,
- const ffi::Map<Var, Expr>& var_remap,
- const ffi::Map<tirx::Var, PrimExpr>& tir_var_remap) {
- GlobalLiftableBindingCollector collector(var_remap, tir_var_remap);
+ const ffi::Map<Var, Expr>& var_remap) {
+ GlobalLiftableBindingCollector collector(var_remap);
TVM_FFI_ICHECK(functions.size());
for (const auto& func : functions) {
int num_inputs = func->GetAttr<int64_t>(attr::kNumInput).value();
@@ -580,7 +556,6 @@ class GlobalLiftableBindingCollector : public BaseLiftableBindingCollector {
info.orig_functions = functions;
info.params = std::move(params);
info.var_remap = var_remap;
- info.tir_var_remap = tir_var_remap;
// Find shared bindings among transform_params. Re-compute var_remap based on the shared
// bindings as collector.var_remap_ may contain invalid mappings.
for (const auto& unified_binding : collector.unified_bindings_) {
@@ -609,16 +584,15 @@ class GlobalLiftableBindingCollector : public BaseLiftableBindingCollector {
}
private:
- GlobalLiftableBindingCollector(const ffi::Map<Var, Expr>& var_remap,
- const ffi::Map<tirx::Var, PrimExpr> tir_var_remap)
- : var_remap_(var_remap), tir_var_remap_(tir_var_remap) {}
+ explicit GlobalLiftableBindingCollector(const ffi::Map<Var, Expr>& var_remap)
+ : var_remap_(var_remap) {}
void VisitBinding(const Binding& binding) override {
TVM_FFI_ICHECK(!binding->IsInstance<MatchCastNode>())
<< "MatchCast is not supported in global lifting";
if (CanLiftBinding(binding)) {
liftable_vars_.insert(binding->var);
auto bound_value = GetBoundValue(binding);
- auto new_value = Bind(bound_value, var_remap_, tir_var_remap_);
+ auto new_value = Bind(bound_value, var_remap_);
if (auto it = original_bindings_.find(new_value); it != original_bindings_.end()) {
it->second.push_back(binding);
} else {
@@ -633,8 +607,6 @@ class GlobalLiftableBindingCollector : public BaseLiftableBindingCollector {
// function parameters, and is updated with the mapping between binding variables asthe collector
// visits the bindings.
ffi::Map<Var, Expr> var_remap_;
- // The cross-function between between TIR variables.
- ffi::Map<tirx::Var, PrimExpr> tir_var_remap_;
std::vector<Binding> unified_bindings_;
// The mapping between the unified bindings and the original bindings in different functions.
// The unified binding is the binding with all variables replaced by the unified variables as
@@ -645,9 +617,8 @@ class GlobalLiftableBindingCollector : public BaseLiftableBindingCollector {
GlobalCollectInfo MakeGlobalLiftPlan(const IRModule& mod,
const std::vector<Function>& target_functions) {
- ParamRemapper remapper;
- auto [var_remap, tir_var_remap] = ParamRemapper::GetParamMapping(target_functions);
- return GlobalLiftableBindingCollector::Collect(target_functions, var_remap, tir_var_remap);
+ auto var_remap = ParamRemapper::GetParamMapping(target_functions);
+ return GlobalLiftableBindingCollector::Collect(target_functions, var_remap);
}
// Adapted from https://stackoverflow.com/a/2072890
diff --git a/src/relax/transform/realize_vdevice.cc b/src/relax/transform/realize_vdevice.cc
index 5e0a990b7d..d2ff9aa5d1 100644
--- a/src/relax/transform/realize_vdevice.cc
+++ b/src/relax/transform/realize_vdevice.cc
@@ -247,9 +247,14 @@ class VDeviceSetCollector : ExprVisitor {
}
}
+ void VisitExprDepTypeField(const Type&) final {}
+
void VisitExpr_(const VarNode* op) override {
- if (current_binding_) {
+ if (current_binding_ && !GetType(current_binding_.value()).as<PrimTypeNode>()) {
auto var = ffi::GetRef<Var>(op);
+ if (GetType(var).as<PrimTypeNode>()) {
+ return;
+ }
var_to_co_located_vars_[current_binding_.value()].push_back(var);
var_to_co_located_vars_[var].push_back(current_binding_.value());
}
diff --git a/src/relax/transform/remove_unused_parameters.cc b/src/relax/transform/remove_unused_parameters.cc
index afce3df33f..61ead90cf7 100644
--- a/src/relax/transform/remove_unused_parameters.cc
+++ b/src/relax/transform/remove_unused_parameters.cc
@@ -134,7 +134,7 @@ std::optional<CalleeAnalysis> AnalyzeCallee(Function func) {
for (const auto& tir_var : free_tir_vars) {
// Pass the symbolic var value as a 1-D shape, matching the ShapeType
// param that now defines the var in the callee.
- new_args.push_back(ShapeExpr({tir_binding.at(tir_var)}));
+ new_args.push_back(ShapeExpr({tir_binding.at(tir_var).as_or_throw<PrimExpr>()}));
}
}
diff --git a/src/relax/transform/rewrite_cuda_graph.cc b/src/relax/transform/rewrite_cuda_graph.cc
index 585ff6bdeb..c993a83748 100644
--- a/src/relax/transform/rewrite_cuda_graph.cc
+++ b/src/relax/transform/rewrite_cuda_graph.cc
@@ -54,6 +54,7 @@
#include <tvm/relax/analysis.h>
#include <tvm/relax/backend.h>
#include <tvm/relax/expr_functor.h>
+#include <tvm/relax/utils.h>
#include <tvm/tirx/expr.h>
#include <tvm/tirx/expr_functor.h>
#include <tvm/tirx/stmt_functor.h>
@@ -111,7 +112,7 @@ class FuncBuilder : public ExprMutator {
* \brief Mark a TIR variable as the ShapeExpr input of the new function.
* \param var The variable to mark as input
*/
- void MarkShapeExprInput(const tirx::VarNode* var) { shape_expr_inputs_.push_back(var); }
+ void MarkShapeExprInput(const tirx::PrimVar& var) { shape_expr_inputs_.push_back(var); }
/*!
* \brief Mark a variable as the output of the new function. The variable must be the LHS of an
* existing binding in the new function.
@@ -128,10 +129,10 @@ class FuncBuilder : public ExprMutator {
ffi::Optional<Var> shape_expr = std::nullopt;
if (shape_expr_inputs_.size()) {
ffi::Array<PrimExpr> tir_vars;
- for (const auto* var : shape_expr_inputs_) {
- auto new_var = ffi::GetRef<tirx::Var>(var).CopyWithSuffix("");
- tir_var_remap_.Set(ffi::GetRef<tirx::Var>(var), new_var.as_or_throw<PrimExpr>());
- tir_vars.push_back(new_var.as_or_throw<PrimExpr>());
+ for (const tirx::PrimVar& var : shape_expr_inputs_) {
+ tirx::PrimVar new_var = var.CopyWithSuffix("");
+ var_remap_[var] = new_var;
+ tir_vars.push_back(new_var);
}
shape_expr = Var("shape_expr", ShapeType(tir_vars));
}
@@ -164,22 +165,10 @@ class FuncBuilder : public ExprMutator {
return func;
}
- PrimExpr VisitTypePrimExprField(const PrimExpr& expr) {
- return tirx::Substitute(expr, tir_var_remap_);
- }
-
- Expr VisitExprFallback_(const ExprNode* op) final {
- if (op->ty.as<PrimTypeNode>()) {
- return VisitTypePrimExprField(ffi::GetRef<Expr>(op).as_or_throw<PrimExpr>());
- }
- return ExprMutator::VisitExprFallback_(op);
- }
-
support::OrderedSet<const VarNode*> inputs_;
support::OrderedSet<const VarNode*> outputs_;
- support::OrderedSet<const tirx::VarNode*> shape_expr_inputs_;
+ support::OrderedSet<tirx::PrimVar, ffi::ObjectPtrHash, ffi::ObjectPtrEqual> shape_expr_inputs_;
std::vector<const VarBindingNode*> bindings_;
- ffi::Map<tirx::Var, PrimExpr> tir_var_remap_;
};
// Collect the storage objects that are used as the function output
@@ -262,14 +251,15 @@ class CUDAGraphRewritePlanner : public ExprVisitor {
DefinableTIRVarsInType(func->params[i]->ty.as_or_throw<Type>());
if (i < num_inputs) {
for (const auto& symbolic_var : symbolic_vars) {
+ auto prim_var = symbolic_var.as_or_throw<tirx::PrimVar>();
if (capture_symbolic_var_name_hints.count(symbolic_var->name_hint)) {
- capture_symbolic_vars_.insert(symbolic_var.get());
+ capture_symbolic_vars_.insert(prim_var);
}
}
} else {
static_vars_.insert(func->params[i].get());
for (const auto& symbolic_var : symbolic_vars) {
- capture_symbolic_vars_.insert(symbolic_var.get());
+ capture_symbolic_vars_.insert(symbolic_var.as_or_throw<tirx::PrimVar>());
}
}
}
@@ -287,8 +277,8 @@ class CUDAGraphRewritePlanner : public ExprVisitor {
plan->lifted_bindings = std::move(region->bindings_);
if (region->shape_expr_inputs_.size()) {
ffi::Array<PrimExpr> tir_vars;
- for (const auto* var : region->shape_expr_inputs_) {
- tir_vars.push_back(ffi::GetRef<tirx::Var>(var).as_or_throw<PrimExpr>());
+ for (const tirx::PrimVar& var : region->shape_expr_inputs_) {
+ tir_vars.push_back(var);
}
plan->propogated_tir_vars = ShapeExpr(tir_vars);
}
@@ -380,7 +370,7 @@ class CUDAGraphRewritePlanner : public ExprVisitor {
// Check whether the call can be lifted to the capture function. It requires all the arguments
// to be static and the call to be a kernel launch or a pure operation (e.g. memory view).
std::vector<const VarNode*> args;
- std::vector<const tirx::VarNode*> tir_vars;
+ std::vector<tirx::PrimVar> tir_vars;
bool is_all_static = [&]() {
if (!IsStatic(call->args, &args, &tir_vars)) {
return false;
@@ -427,7 +417,7 @@ class CUDAGraphRewritePlanner : public ExprVisitor {
}
void MarkAsFuncInput(const std::vector<const VarNode*>& vars,
- const std::vector<const tirx::VarNode*>& tir_vars = {}) {
+ const std::vector<tirx::PrimVar>& tir_vars = {}) {
if (current_block_scope_.capture_builder == nullptr) {
return;
}
@@ -437,7 +427,7 @@ class CUDAGraphRewritePlanner : public ExprVisitor {
current_block_scope_.capture_builder->MarkInput(var);
}
}
- for (const tirx::VarNode* tir_var : tir_vars) {
+ for (const tirx::PrimVar& tir_var : tir_vars) {
current_block_scope_.capture_builder->MarkShapeExprInput(tir_var);
}
}
@@ -467,7 +457,7 @@ class CUDAGraphRewritePlanner : public ExprVisitor {
void VisitBinding_(const VarBindingNode* binding, const TupleNode* tuple) final {
std::vector<const VarNode*> args;
- std::vector<const tirx::VarNode*> tir_vars;
+ std::vector<tirx::PrimVar> tir_vars;
if (IsStatic(tuple->fields, &args, &tir_vars)) {
AddStaticBinding(binding, false);
MarkAsFuncInput(args, tir_vars);
@@ -491,16 +481,16 @@ class CUDAGraphRewritePlanner : public ExprVisitor {
bool IsStatic(const PrimExpr& expr,
[[maybe_unused]] std::vector<const VarNode*>* vars_collector = nullptr,
- std::vector<const tirx::VarNode*>* tir_vars_collector = nullptr) {
+ std::vector<tirx::PrimVar>* tir_vars_collector = nullptr) {
bool is_static = true;
tirx::PostOrderVisit(expr, [&](const ffi::ObjectRef& e) {
- if (auto var = e.as<tirx::VarNode>()) {
- if (!capture_symbolic_vars_.count(var)) {
+ if (auto var = e.as<tirx::PrimVar>()) {
+ if (!capture_symbolic_vars_.count(var.value())) {
is_static = false;
return;
}
if (tir_vars_collector != nullptr) {
- tir_vars_collector->push_back(var);
+ tir_vars_collector->push_back(var.value());
}
}
});
@@ -508,21 +498,20 @@ class CUDAGraphRewritePlanner : public ExprVisitor {
}
bool IsStatic(const Expr& expr, std::vector<const VarNode*>* vars_collector = nullptr,
- std::vector<const tirx::VarNode*>* tir_vars_collector = nullptr) {
+ std::vector<tirx::PrimVar>* tir_vars_collector = nullptr) {
if (expr->IsInstance<ConstantNode>() || expr->IsInstance<DataTypeImmNode>() ||
expr->IsInstance<StringImmNode>() || expr->IsInstance<GlobalVarNode>()) {
return true;
}
+ if (const auto* var = expr.as<VarNode>();
+ var && (expr.as<DataflowVarNode>() || static_vars_.count(var))) {
+ return IsStaticRuntimeVar(var, vars_collector, tir_vars_collector);
+ }
if (auto prim_value = expr.as<PrimExpr>()) {
return IsStatic(prim_value.value(), vars_collector, tir_vars_collector);
}
if (const auto* var = expr.as<VarNode>()) {
- if (vars_collector != nullptr) {
- vars_collector->push_back(var);
- }
- // recursively check the type to collect the symbolic TIR vars
- return static_vars_.count(var) &&
- IsStatic(var->ty.as_or_throw<Type>(), vars_collector, tir_vars_collector);
+ return IsStaticRuntimeVar(var, vars_collector, tir_vars_collector);
}
if (const auto* shape = expr.as<ShapeExprNode>()) {
@@ -534,9 +523,18 @@ class CUDAGraphRewritePlanner : public ExprVisitor {
return false;
}
+ bool IsStaticRuntimeVar(const VarNode* var, std::vector<const VarNode*>* vars_collector,
+ std::vector<tirx::PrimVar>* tir_vars_collector) {
+ if (vars_collector != nullptr) {
+ vars_collector->push_back(var);
+ }
+ return static_vars_.count(var) &&
+ IsStatic(var->ty.as_or_throw<Type>(), vars_collector, tir_vars_collector);
+ }
+
template <typename T>
bool IsStatic(const ffi::Array<T>& exprs, std::vector<const VarNode*>* vars_collector = nullptr,
- std::vector<const tirx::VarNode*>* tir_vars_collector = nullptr) {
+ std::vector<tirx::PrimVar>* tir_vars_collector = nullptr) {
bool result = true;
for (const auto& expr : exprs) {
// If vars_collector is provided, we will collect all the vars in the exprs and we should
@@ -550,7 +548,7 @@ class CUDAGraphRewritePlanner : public ExprVisitor {
}
bool IsStatic(const Type& ty, std::vector<const VarNode*>* vars_collector = nullptr,
- std::vector<const tirx::VarNode*>* tir_vars_collector = nullptr) {
+ std::vector<tirx::PrimVar>* tir_vars_collector = nullptr) {
if (const auto* tensor_ty = ty.as<TensorTypeNode>()) {
if (auto shape = tensor_ty->GetShape()) {
return IsStatic(shape.value(), vars_collector, tir_vars_collector);
@@ -626,7 +624,7 @@ class CUDAGraphRewritePlanner : public ExprVisitor {
std::unordered_set<const VarNode*> static_vars_;
// Symbolic variables that are allowed to be captured. This can come from symbolic shapes of
// weights or hints in the function annotations.
- std::unordered_set<const tirx::VarNode*> capture_symbolic_vars_;
+ std::unordered_set<tirx::PrimVar, ffi::ObjectPtrHash, ffi::ObjectPtrEqual> capture_symbolic_vars_;
// Binding to the FuncBuilder if the binding is lifted. This is used to update the inputs/outputs
// of the lifted function when its binding is used outside.
std::unordered_map<const VarNode*, FuncBuilder*> binding_to_region_;
@@ -812,13 +810,13 @@ class CUDAGraphRewriter : public ExprMutator {
// bind the symbolic parameters to the actual values.
const auto& shape_expr = plan->func->params.back();
auto symbolic_params = shape_expr->ty.as_or_throw<ShapeType>()->values.value();
- ffi::Map<tirx::Var, PrimExpr> tir_var_remap;
+ ffi::Map<Var, Expr> var_remap;
TVM_FFI_ICHECK_EQ(symbolic_params.size(), propogated_tir_vars->values.size());
for (int i = 0; i < static_cast<int>(symbolic_params.size()); ++i) {
- tir_var_remap.Set(symbolic_params[i].as_or_throw<tirx::Var>(),
- propogated_tir_vars->values[i]);
+ var_remap.Set(symbolic_params[i].as_or_throw<tirx::PrimVar>(),
+ propogated_tir_vars->values[i]);
}
- call_ty = Bind(call_ty, tir_var_remap);
+ call_ty = Bind(call_ty, var_remap);
}
// Arguments of builtin_run_or_capture
ffi::Array<Expr> tuple_arg_fields{gv_func, Tuple(args),
diff --git a/src/relax/transform/specialize_primfunc_based_on_callsite.cc b/src/relax/transform/specialize_primfunc_based_on_callsite.cc
index 8445446e68..e9563d43cd 100644
--- a/src/relax/transform/specialize_primfunc_based_on_callsite.cc
+++ b/src/relax/transform/specialize_primfunc_based_on_callsite.cc
@@ -93,7 +93,7 @@ class SpecializeTIRCallArgs : ExprMutator {
scope = tensor_ty->vdevice.value()->memory_scope;
}
ffi::String name;
- if (args[i]->IsInstance<relax::VarNode>()) {
+ if (args[i]->IsInstance<tvm::VarNode>()) {
name = args[i].as_or_throw<Var>()->name_hint;
} else {
name = std::string({static_cast<char>('A' + i)});
diff --git a/src/relax/transform/split_call_tir_by_pattern.cc b/src/relax/transform/split_call_tir_by_pattern.cc
index 10686d262f..14373e0433 100644
--- a/src/relax/transform/split_call_tir_by_pattern.cc
+++ b/src/relax/transform/split_call_tir_by_pattern.cc
@@ -98,17 +98,17 @@ class ForMatcher : public TensorizeComparator {
bool VisitExpr(const Expr& expr, const PrimExpr& rhs) final {
PrimExpr lhs = expr.as_or_throw<PrimExpr>();
- if (const auto* op = rhs.as<VarNode>()) {
- if (pattern_vars_.count(ffi::GetRef<Var>(op))) {
+ if (auto rhs_prim_var = rhs.as<PrimVar>()) {
+ Var rhs_var = rhs_prim_var.value();
+ if (pattern_vars_.count(rhs_var)) {
// special case for pattern vars
- const auto* lhs_ptr = lhs.as<VarNode>();
- if (lhs_ptr == nullptr) {
+ if (!lhs.as<PrimVar>()) {
if (lhs->IsInstance<tirx::IntImmNode>() || lhs->IsInstance<tirx::FloatImmNode>()) {
- ffi::Optional<PrimExpr> value = QueryEvaluatedSymbols(ffi::GetRef<Var>(op));
+ ffi::Optional<PrimExpr> value = QueryEvaluatedSymbols(rhs_var);
if (value.has_value()) {
if (!analyzer_->CanProveEqual(lhs, value.value())) return false;
} else {
- evaluated_symbols.back()[ffi::GetRef<Var>(op)] = lhs;
+ evaluated_symbols.back()[rhs_var] = lhs;
}
return true;
} else {
@@ -119,9 +119,10 @@ class ForMatcher : public TensorizeComparator {
}
// pattern_var * expr
if (const auto* rhs_ptr = rhs.as<MulNode>()) {
- const auto* operand_a = rhs_ptr->a.as<VarNode>();
- const auto* operand_b = rhs_ptr->b.as<VarNode>();
- if (operand_a != nullptr && pattern_vars_.count(ffi::GetRef<Var>(operand_a))) {
+ auto operand_a = rhs_ptr->a.as<PrimVar>();
+ auto operand_b = rhs_ptr->b.as<PrimVar>();
+ if (operand_a && pattern_vars_.count(operand_a.value())) {
+ Var pattern_var = operand_a.value();
// pattern var is on the left
evaluated_symbols.push_back(SymbolMap());
bool match = VisitExpr(lhs, rhs_ptr->b);
@@ -129,12 +130,12 @@ class ForMatcher : public TensorizeComparator {
evaluated_symbols.pop_back();
if (match) {
evaluated_symbols.back().insert(symbol_map.begin(), symbol_map.end());
- evaluated_symbols.back()[ffi::GetRef<Var>(operand_a)] =
- MakeConstScalar(rhs_ptr->b.ty(), 1);
+ evaluated_symbols.back()[pattern_var] = MakeConstScalar(rhs_ptr->b.ty(), 1);
return true;
}
}
- if (operand_b != nullptr && pattern_vars_.count(ffi::GetRef<Var>(operand_b))) {
+ if (operand_b && pattern_vars_.count(operand_b.value())) {
+ Var pattern_var = operand_b.value();
// pattern var is on the right
evaluated_symbols.push_back(SymbolMap());
bool match = VisitExpr(lhs, rhs_ptr->a);
@@ -142,17 +143,17 @@ class ForMatcher : public TensorizeComparator {
evaluated_symbols.pop_back();
if (match) {
evaluated_symbols.back().insert(symbol_map.begin(), symbol_map.end());
- evaluated_symbols.back()[ffi::GetRef<Var>(operand_b)] =
- MakeConstScalar(rhs_ptr->a.ty(), 1);
+ evaluated_symbols.back()[pattern_var] = MakeConstScalar(rhs_ptr->a.ty(), 1);
return true;
}
}
}
// pattern_Var + expr
if (const auto* rhs_ptr = rhs.as<AddNode>()) {
- const auto* operand_a = rhs_ptr->a.as<VarNode>();
- const auto* operand_b = rhs_ptr->b.as<VarNode>();
- if (operand_a != nullptr && pattern_vars_.count(ffi::GetRef<Var>(operand_a))) {
+ auto operand_a = rhs_ptr->a.as<PrimVar>();
+ auto operand_b = rhs_ptr->b.as<PrimVar>();
+ if (operand_a && pattern_vars_.count(operand_a.value())) {
+ Var pattern_var = operand_a.value();
// pattern var is on the left
evaluated_symbols.push_back(SymbolMap());
bool match = VisitExpr(lhs, rhs_ptr->b);
@@ -160,12 +161,12 @@ class ForMatcher : public TensorizeComparator {
evaluated_symbols.pop_back();
if (match) {
evaluated_symbols.back().insert(symbol_map.begin(), symbol_map.end());
- evaluated_symbols.back()[ffi::GetRef<Var>(operand_a)] =
- MakeConstScalar(rhs_ptr->b.ty(), 0);
+ evaluated_symbols.back()[pattern_var] = MakeConstScalar(rhs_ptr->b.ty(), 0);
return true;
}
}
- if (operand_b != nullptr && pattern_vars_.count(ffi::GetRef<Var>(operand_b))) {
+ if (operand_b && pattern_vars_.count(operand_b.value())) {
+ Var pattern_var = operand_b.value();
// pattern var is on the right
evaluated_symbols.push_back(SymbolMap());
bool match = VisitExpr(lhs, rhs_ptr->a);
@@ -173,8 +174,7 @@ class ForMatcher : public TensorizeComparator {
evaluated_symbols.pop_back();
if (match) {
evaluated_symbols.back().insert(symbol_map.begin(), symbol_map.end());
- evaluated_symbols.back()[ffi::GetRef<Var>(operand_b)] =
- MakeConstScalar(rhs_ptr->a.ty(), 0);
+ evaluated_symbols.back()[pattern_var] = MakeConstScalar(rhs_ptr->a.ty(), 0);
return true;
}
}
diff --git a/src/relax/transform/topological_sort.cc b/src/relax/transform/topological_sort.cc
index 3f24715382..f50dceeb92 100644
--- a/src/relax/transform/topological_sort.cc
+++ b/src/relax/transform/topological_sort.cc
@@ -38,13 +38,13 @@ namespace {
struct InputNode {};
struct OutputNode {};
-using DataflowNode = std::variant<InputNode, OutputNode, tvm::relax::Var>;
+using DataflowNode = std::variant<InputNode, OutputNode, tvm::Var>;
bool operator==(const DataflowNode& a, const DataflowNode& b) {
- if (const tvm::relax::Var* var_a = std::get_if<tvm::relax::Var>(&a)) {
- if (const tvm::relax::Var* var_b = std::get_if<tvm::relax::Var>(&b)) {
- const tvm::relax::VarNode* ptr_a = var_a->get();
- const tvm::relax::VarNode* ptr_b = var_b->get();
+ if (const tvm::Var* var_a = std::get_if<tvm::Var>(&a)) {
+ if (const tvm::Var* var_b = std::get_if<tvm::Var>(&b)) {
+ const tvm::VarNode* ptr_a = var_a->get();
+ const tvm::VarNode* ptr_b = var_b->get();
return ptr_a == ptr_b;
}
}
@@ -57,8 +57,8 @@ bool operator==(const DataflowNode& a, const DataflowNode& b) {
template <>
struct std::hash<DataflowNode> {
std::size_t operator()(const DataflowNode& node) const noexcept {
- if (const tvm::relax::Var* var = std::get_if<tvm::relax::Var>(&node)) {
- const tvm::relax::VarNode* ptr = var->get();
+ if (const tvm::Var* var = std::get_if<tvm::Var>(&node)) {
+ const tvm::VarNode* ptr = var->get();
std::hash<decltype(ptr)> hasher;
return hasher(ptr);
} else {
@@ -132,10 +132,11 @@ class BindingOrderCollector : ExprVisitor {
}
private:
- void VisitVarDef(const Var& var) override { dependencies_.binding_order.push_back(var); }
+ void VisitVarDef(const Var&) override {}
void VisitExpr_(const FunctionNode* op) override {
for (const auto& var : op->params) {
+ graph_vars_.insert(var);
dependencies_.downstream_users[InputNode()].push_back(var);
dependencies_.upstream_requirements[var].push_back(InputNode());
}
@@ -145,12 +146,17 @@ class BindingOrderCollector : ExprVisitor {
void VisitBinding(const Binding& binding) override {
auto cache = current_binding_;
current_binding_ = binding->var;
+ graph_vars_.insert(binding->var);
+ dependencies_.binding_order.push_back(binding->var);
ExprVisitor::VisitBinding(binding);
current_binding_ = cache;
}
void VisitExpr_(const VarNode* op) override {
Var upstream_requirement = ffi::GetRef<Var>(op);
+ if (!graph_vars_.count(upstream_requirement)) {
+ return;
+ }
auto downstream_user = current_binding_;
dependencies_.downstream_users[upstream_requirement].push_back(downstream_user);
@@ -158,6 +164,7 @@ class BindingOrderCollector : ExprVisitor {
}
DataflowNode current_binding_ = OutputNode();
+ std::unordered_set<Var, ffi::ObjectPtrHash, ffi::ObjectPtrEqual> graph_vars_;
Dependencies dependencies_;
};
diff --git a/src/relax/transform/utils.h b/src/relax/transform/utils.h
index 105dbff5e9..54b6d9f5a0 100644
--- a/src/relax/transform/utils.h
+++ b/src/relax/transform/utils.h
@@ -89,6 +89,10 @@ class MemoizedExprTranslator : public ExprFunctor<OutputType(const Expr&)> {
return memo_[ffi::GetRef<Expr>(vn)];
}
+ virtual OutputType VisitExpr_(const DataflowVarNode* vn) {
+ return VisitExpr_(static_cast<const VarNode*>(vn));
+ }
+
virtual OutputType VisitBinding_(const VarBindingNode* binding) {
TVM_FFI_ICHECK_EQ(memo_.count(binding->var), 0);
auto v = VisitExpr(binding->value);
@@ -219,50 +223,89 @@ class VarReplacer : public ExprMutator {
};
/*!
- * \brief Renew the definition of symbolic vars in Relax.
- * \details This mutator is used to prevent the same symbolic var from being used in different
- * functions, which is malformed.
+ * \brief Renew the definition of dependent primitive vars in Relax.
+ * \details This mutator is used to prevent the same primitive Var object from being defined in
+ * different functions, which is malformed.
*/
-class SymbolicVarRenewMutator : public ExprMutator, tirx::ExprMutator {
+class SymbolicVarRenewMutator : public ExprMutator {
public:
static Function Renew(const Function& function) {
SymbolicVarRenewMutator mutator;
return mutator.VisitExpr(function).as_or_throw<Function>();
}
- SymbolicVarRenewMutator() = default;
+ SymbolicVarRenewMutator() : SymbolicVarRenewMutator(false) {}
protected:
+ explicit SymbolicVarRenewMutator(bool renew_all_var_definitions)
+ : renew_all_var_definitions_(renew_all_var_definitions) {}
+
using relax::ExprMutator::VisitExpr;
using relax::ExprMutator::VisitExpr_;
- using tirx::ExprMutator::VisitExpr_;
- PrimExpr VisitTypePrimExprField(const PrimExpr& expr) final {
- return tirx::ExprMutator::VisitExpr(expr).as_or_throw<PrimExpr>();
+ static Var CopyVar(const VarNode* op, Type ty) {
+ ffi::Optional<Type> ty_annotation = ty.IsMissing() ? std::nullopt : ffi::Optional<Type>(ty);
+ if (op->IsInstance<DataflowVarNode>()) {
+ return DataflowVar(op->name_hint, std::move(ty_annotation), op->span);
+ }
+ return Var(op->name_hint, std::move(ty_annotation), op->span);
+ }
+
+ Type RenewType(const VarNode* op) {
+ return op->ty.IsMissing() ? op->ty : this->VisitExprDepTypeField(op->ty);
}
- Expr VisitExprFallback_(const ExprNode* op) final {
- if (op->ty.as<PrimTypeNode>()) {
- return VisitTypePrimExprField(ffi::GetRef<Expr>(op).as_or_throw<PrimExpr>());
+ Var RenewVarDefinition(const VarNode* op) {
+ Var old_var = ffi::GetRef<Var>(op);
+ if (auto it = var_remap_.find(old_var); it != var_remap_.end()) {
+ return it->second;
+ }
+
+ Type new_ty = RenewType(op);
+ bool is_dataflow = op->IsInstance<DataflowVarNode>();
+ bool renew = renew_all_var_definitions_ || (!is_dataflow && op->ty.as<PrimTypeNode>()) ||
+ !new_ty.same_as(op->ty);
+ if (!renew) {
+ return old_var;
}
- return relax::ExprMutator::VisitExprFallback_(op);
+
+ Var renewed = CopyVar(op, std::move(new_ty));
+ var_remap_[old_var] = renewed;
+ return renewed;
}
- // TODO(Siyuan): enhance the method to the following steps:
- // 1. Visit and replace all tirx::Vars at the definition point
- // 2. Revisit the function again and update the use side.
- Expr VisitExpr_(const tirx::VarNode* op) final {
- auto it = var_map_.find(ffi::GetRef<tirx::Var>(op));
- if (it != var_map_.end()) {
- return (*it).second;
- } else {
- auto n = ffi::make_object<tirx::VarNode>(*op);
- tirx::Var v(n);
- var_map_.Set(ffi::GetRef<tirx::Var>(op), v);
- return v;
+ Expr VisitExpr_(const VarNode* op) final {
+ Var old_var = ffi::GetRef<Var>(op);
+ if (auto it = var_remap_.find(old_var); it != var_remap_.end()) {
+ return it->second;
+ }
+ if (!op->ty.as<PrimTypeNode>()) {
+ return ExprMutator::VisitExpr_(op);
}
+
+ Var renewed = CopyVar(op, op->ty);
+ var_remap_[old_var] = renewed;
+ return renewed;
}
- Expr VisitExpr_(const FunctionNode* op) {
+ Expr VisitExpr_(const DataflowVarNode* op) final {
+ Var old_var = ffi::GetRef<Var>(op);
+ if (auto it = var_remap_.find(old_var); it != var_remap_.end()) {
+ return it->second;
+ }
+ if (!renew_all_var_definitions_) {
+ return old_var;
+ }
+
+ Var renewed = CopyVar(op, op->ty);
+ var_remap_[old_var] = renewed;
+ return renewed;
+ }
+
+ Var VisitVarDef_(const VarNode* op) override { return RenewVarDefinition(op); }
+
+ Var VisitVarDef_(const DataflowVarNode* op) override { return RenewVarDefinition(op); }
+
+ Expr VisitExpr_(const FunctionNode* op) final {
tvm::ffi::Array<Var> params;
bool all_params_unchanged = true;
for (Var param : op->params) {
@@ -284,40 +327,25 @@ class SymbolicVarRenewMutator : public ExprMutator, tirx::ExprMutator {
}
}
- ffi::Map<tirx::Var, tirx::Var> var_map_;
+ bool renew_all_var_definitions_{false};
};
/*!
- * \brief Copy a function while renewing the relax Vars and the tirx Vars.
+ * \brief Copy a function while renewing every Var object, including dependent primitive uses.
* \details All variables that are bound inside the original function would be copied to satisfy
* the restriction in the well-formed check: Variables in Relax must be bound exactly once.
*/
class FunctionCopier : public SymbolicVarRenewMutator {
public:
- FunctionCopier() = default;
+ FunctionCopier() : SymbolicVarRenewMutator(true) {}
Function Copy(Function func) { return VisitExpr(func).as_or_throw<Function>(); }
- ffi::Map<Var, Var> GetVarMap() { return relax_var_map_; }
-
- private:
- using relax::ExprMutator::VisitExpr;
-
- Var VisitVarDef_(const DataflowVarNode* var) override {
- Var new_var = SymbolicVarRenewMutator::VisitVarDef_(var);
- Var copied_var = DataflowVar(new_var->name_hint, GetType(new_var), new_var->span);
- var_remap_[ffi::GetRef<Var>(var)] = copied_var;
- relax_var_map_.Set(ffi::GetRef<Var>(var), copied_var);
- return copied_var;
- }
-
- Var VisitVarDef_(const VarNode* var) override {
- Var new_var = SymbolicVarRenewMutator::VisitVarDef_(var);
- Var copied_var = Var(new_var->name_hint, GetType(new_var), new_var->span);
- var_remap_[ffi::GetRef<Var>(var)] = copied_var;
- relax_var_map_.Set(ffi::GetRef<Var>(var), copied_var);
- return copied_var;
+ ffi::Map<Var, Var> GetVarMap() {
+ ffi::Map<Var, Var> result;
+ for (const auto& [old_var, new_var] : var_remap_) {
+ result.Set(old_var, new_var);
+ }
+ return result;
}
-
- ffi::Map<Var, Var> relax_var_map_;
};
/*!
diff --git a/src/relax/utils.cc b/src/relax/utils.cc
index 63f0d08636..2e004a59e5 100644
--- a/src/relax/utils.cc
+++ b/src/relax/utils.cc
@@ -24,6 +24,8 @@
#include <tvm/relax/analysis.h>
#include <tvm/relax/attrs/index.h>
#include <tvm/relax/expr_functor.h>
+#include <tvm/relax/utils.h>
+#include <tvm/tirx/op_attr_types.h>
#include <tvm/tirx/stmt_functor.h>
namespace tvm {
@@ -32,18 +34,22 @@ namespace relax {
/*! \brief Helper to implement bind params.*/
class ExprBinder : public ExprMutator {
public:
- explicit ExprBinder(const tvm::ffi::Map<Var, Expr>& args_map,
- const tvm::ffi::Map<tirx::Var, PrimExpr>& symbolic_var_map)
- : args_map_(args_map), symbolic_var_map_(symbolic_var_map) {}
+ explicit ExprBinder(const tvm::ffi::Map<Var, Expr>& bindings) : bindings_(bindings) {}
private:
using ExprMutator::VisitExpr_;
+ Expr VisitExpr_(const ShapeExprNode* op) final {
+ ffi::Array<PrimExpr> values =
+ op->values.Map([this](const PrimExpr& value) { return BindShapeValue(value); });
+ return values.same_as(op->values) ? ffi::GetRef<Expr>(op) : ShapeExpr(values, op->span);
+ }
+
Expr VisitExpr_(const FunctionNode* op) final {
tvm::ffi::Array<Var> params;
bool all_params_unchanged = true;
for (const Var& param : op->params) {
- if (args_map_.count(param)) {
+ if (bindings_.count(param)) {
all_params_unchanged = false;
} else {
Var new_param = this->VisitVarDef(param);
@@ -68,61 +74,72 @@ class ExprBinder : public ExprMutator {
Expr VisitExpr_(const VarNode* op) final {
auto id = ffi::GetRef<Var>(op);
- auto it = args_map_.find(id);
- if (it != args_map_.end()) {
+ auto it = bindings_.find(id);
+ if (it != bindings_.end()) {
return (*it).second;
} else {
return ExprMutator::VisitExpr_(op);
}
}
- PrimExpr VisitTypePrimExprField(const PrimExpr& expr) final {
- auto new_expr = tirx::Substitute(expr, symbolic_var_map_);
- if (!expr.same_as(new_expr)) {
- arith::Analyzer analyzer;
- new_expr = analyzer->Simplify(new_expr);
- }
- return new_expr;
- }
+ PrimExpr VisitTypePrimExprField(const PrimExpr& expr) final { return BindShapeValue(expr); }
- Expr VisitExprFallback_(const ExprNode* op) final {
- if (op->ty.as<PrimTypeNode>()) {
- return VisitTypePrimExprField(ffi::GetRef<Expr>(op).as_or_throw<PrimExpr>());
- }
- return ExprMutator::VisitExprFallback_(op);
+ PrimExpr BindShapeValue(const PrimExpr& expr) {
+ PrimExpr output = tirx::Substitute(expr, [this](const Var& var) -> ffi::Optional<Expr> {
+ auto it = bindings_.find(var);
+ if (it == bindings_.end()) return std::nullopt;
+ if (auto value = (*it).second.as<PrimExpr>()) {
+ return ffi::Optional<Expr>(*value);
+ }
+ return std::nullopt;
+ });
+ return output.same_as(expr) ? expr : analyzer_->Simplify(output);
}
- private:
- const tvm::ffi::Map<Var, Expr>& args_map_;
- const tvm::ffi::Map<tirx::Var, PrimExpr>& symbolic_var_map_;
+ const tvm::ffi::Map<Var, Expr>& bindings_;
+ arith::Analyzer analyzer_;
};
/*!
* \brief Bind params on expr
* \param expr The expr where to bind params
* \param binds The map from param var to the expr it binds to
- * \param symbolic_var_map The map from symbolic var to the expr it binds to
* \return The result expr after bind params
*/
-Expr Bind(const Expr& expr, const tvm::ffi::Map<Var, Expr>& binds,
- const tvm::ffi::Map<tirx::Var, PrimExpr>& symbolic_var_map) {
- return ExprBinder(binds, symbolic_var_map).VisitExpr(expr);
+Expr Bind(const Expr& expr, const tvm::ffi::Map<Var, Expr>& binds) {
+ return ExprBinder(binds).VisitExpr(expr);
}
-Type Bind(const Type& ty, const tvm::ffi::Map<tirx::Var, PrimExpr>& symbolic_var_map) {
- return ExprBinder({}, symbolic_var_map).VisitExprDepTypeField(ty);
+Type Bind(const Type& ty, const tvm::ffi::Map<Var, Expr>& binds) {
+ return ExprBinder(binds).VisitExprDepTypeField(ty);
}
-tvm::ffi::Map<tirx::Var, PrimExpr> InferSymbolicVarMap(
- const tvm::ffi::Map<relax::Var, relax::Expr>& relax_var_remap,
- const arith::Analyzer& analyzer) {
- (void)analyzer;
- tvm::ffi::Map<tirx::Var, PrimExpr> tir_var_remap;
+tvm::ffi::Map<Var, Expr> InferSymbolicVarMap(
+ const tvm::ffi::Map<tvm::Var, relax::Expr>& relax_var_remap, const arith::Analyzer& analyzer) {
+ tvm::ffi::Map<Var, Expr> var_remap = relax_var_remap;
+
+ for (const auto& [var, value] : relax_var_remap) {
+ if (!var.as<tirx::PrimVar>()) continue;
+ TVM_FFI_CHECK(value.as<PrimExpr>().has_value(), ValueError)
+ << "Explicit binding for symbolic variable " << var
+ << " must be a primitive expression, but received " << value;
+ }
- auto bind_from_prim_expr = [&tir_var_remap](const PrimExpr& var_shape,
- const PrimExpr& expr_shape) {
- if (auto var = var_shape.as<tirx::Var>()) {
- tir_var_remap.Set(var.value(), expr_shape);
+ auto bind_from_prim_expr = [&relax_var_remap, &var_remap, &analyzer](const PrimExpr& var_shape,
+ const PrimExpr& expr_shape) {
+ if (auto var = var_shape.as<tirx::PrimVar>()) {
+ if (auto it = relax_var_remap.find(var.value()); it != relax_var_remap.end()) {
+ auto explicit_value = (*it).second.as<PrimExpr>();
+ TVM_FFI_CHECK(explicit_value.has_value(), ValueError)
+ << "Explicit binding for symbolic variable " << var.value()
+ << " must be a primitive expression, but received " << (*it).second;
+ TVM_FFI_CHECK(analyzer->CanProveEqual(explicit_value.value(), expr_shape), ValueError)
+ << "Explicit binding for symbolic variable " << var.value()
+ << " conflicts with the value inferred from a parameter shape: explicit value "
+ << explicit_value.value() << ", inferred value " << expr_shape;
+ return;
+ }
+ var_remap.Set(var.value(), expr_shape);
}
};
@@ -182,7 +199,7 @@ tvm::ffi::Map<tirx::Var, PrimExpr> InferSymbolicVarMap(
bind_from_ty(var_ty, expr_ty);
}
- return tir_var_remap;
+ return var_remap;
}
bool IsBoolType(const Type& ty, bool permit_unknown_rank, bool permit_unknown_dtype) {
@@ -219,9 +236,14 @@ bool IsImpureCall(const Call& call) {
if (auto op_ptr = call->op.as<OpNode>()) {
auto op = ffi::GetRef<Op>(op_ptr);
static auto purity_map = Op::GetAttrMap<bool>("FPurity");
- TVM_FFI_ICHECK(purity_map.count(op))
- << "Cannot find the registered purity of this op: " << op->name;
- return !(purity_map[op]);
+ if (purity_map.count(op)) {
+ return !(purity_map[op]);
+ }
+ static auto effect_map = Op::GetAttrMap<tirx::TCallEffectKind>("TCallEffectKind");
+ TVM_FFI_ICHECK(effect_map.count(op))
+ << "Cannot find the registered purity or call effect of this op: " << op->name;
+ auto effect = static_cast<tirx::CallEffectKind>(effect_map[op]);
+ return effect > tirx::CallEffectKind::kPure;
}
// the Type must be FuncType
auto func_ty = GetTypeAs<FuncTypeNode>(call->op);
diff --git a/src/s_tir/analysis/identify_memcpy.cc b/src/s_tir/analysis/identify_memcpy.cc
index 8a77c5b3eb..8d2a393589 100644
--- a/src/s_tir/analysis/identify_memcpy.cc
+++ b/src/s_tir/analysis/identify_memcpy.cc
@@ -46,7 +46,7 @@ using namespace tvm::tirx;
std::variant<MemCpyDetails, std::string> IdentifyMemCpyImpl(const For& loop,
arith::AnalyzerObj* analyzer) {
ffi::Map<Var, arith::IntSet> loop_intervals;
- ffi::Map<Var, Range> loop_ranges;
+ ffi::Map<PrimVar, Range> loop_ranges;
PrimExpr total_loop_iterations = 1;
// Walk through the loop nest, stopping at the first loop whose body
diff --git a/src/s_tir/analysis/sblock_buffer_access_lca_detector.cc b/src/s_tir/analysis/sblock_buffer_access_lca_detector.cc
index 72e44a79a4..6b07bc2469 100644
--- a/src/s_tir/analysis/sblock_buffer_access_lca_detector.cc
+++ b/src/s_tir/analysis/sblock_buffer_access_lca_detector.cc
@@ -150,7 +150,8 @@ class LCADetector : public StmtExprVisitor {
const PrimExpr& binding) -> const ScopeInfo* {
const ScopeInfo* highest_scope = nullptr;
PostOrderVisit(binding, [this, &highest_scope](const ffi::ObjectRef& obj) {
- if (const VarNode* loop_var = obj.as<VarNode>()) {
+ if (auto var = obj.as<PrimVar>()) {
+ const VarNode* loop_var = var.value().get();
auto it = loop_scope_map_.find(loop_var);
if (it == loop_scope_map_.end()) {
return;
@@ -199,7 +200,8 @@ class LCADetector : public StmtExprVisitor {
const ScopeInfo* scope = ancestor_scopes_.back();
auto handle_itervar = [&opaque_var_scope, &scope](const ffi::ObjectRef& obj) {
- if (const VarNode* iter_var = obj.as<VarNode>()) {
+ if (auto var = obj.as<PrimVar>()) {
+ const VarNode* iter_var = var.value().get();
auto dom_scope_it = opaque_var_scope.find(iter_var);
if (dom_scope_it == opaque_var_scope.end()) {
return;
diff --git a/src/s_tir/meta_schedule/feature_extractor/per_store_feature.cc b/src/s_tir/meta_schedule/feature_extractor/per_store_feature.cc
index 1e45f836de..530da5c15f 100644
--- a/src/s_tir/meta_schedule/feature_extractor/per_store_feature.cc
+++ b/src/s_tir/meta_schedule/feature_extractor/per_store_feature.cc
@@ -919,8 +919,8 @@ void Feature::SubFeature::SetReuse(const LoopNest& loop_nest, int64_t top_loop_t
for (const MultiIndex& multi_index : this->multi_indices) {
for (const PrimExpr& index : multi_index) {
PostOrderVisit(index, [®ion_vars](const ffi::ObjectRef& obj) -> void {
- if (const auto* var = obj.as<VarNode>()) {
- region_vars.insert(var);
+ if (auto var = obj.as<PrimVar>()) {
+ region_vars.insert(var.value().get());
}
});
}
diff --git a/src/s_tir/meta_schedule/postproc/disallow_async_strided_mem_copy.cc b/src/s_tir/meta_schedule/postproc/disallow_async_strided_mem_copy.cc
index cfa7393203..be93894375 100644
--- a/src/s_tir/meta_schedule/postproc/disallow_async_strided_mem_copy.cc
+++ b/src/s_tir/meta_schedule/postproc/disallow_async_strided_mem_copy.cc
@@ -114,7 +114,7 @@ struct AsyncStridedMemCopyFinder : private StmtExprVisitor {
}
bool found_ = false;
- ffi::Map<Var, Range> input_iters = ffi::Map<Var, Range>();
+ ffi::Map<PrimVar, Range> input_iters = ffi::Map<PrimVar, Range>();
};
} // namespace s_tir
diff --git a/src/s_tir/meta_schedule/schedule_rule/multi_level_tiling_tensor_core.cc b/src/s_tir/meta_schedule/schedule_rule/multi_level_tiling_tensor_core.cc
index 89d0ffe1a9..c966358d5c 100644
--- a/src/s_tir/meta_schedule/schedule_rule/multi_level_tiling_tensor_core.cc
+++ b/src/s_tir/meta_schedule/schedule_rule/multi_level_tiling_tensor_core.cc
@@ -798,9 +798,11 @@ ffi::Optional<LoopRV> MultiLevelTilingTensorCoreNode::TransformWithTensorIntrin(
const tirx::IndexMap& index_map = mapping_info->mappings[0];
// Find the correspondence between block iters and the iters in the index map.
- std::unordered_map<tirx::Var, tirx::Var> lhs_to_index_map_src;
- std::unordered_map<tirx::Var, PrimExpr> rhs_to_index_map_tgt;
- std::unordered_set<tirx::Var> unmapped_index_map_src;
+ std::unordered_map<tirx::PrimVar, tirx::PrimVar, ffi::ObjectPtrHash, ffi::ObjectPtrEqual>
+ lhs_to_index_map_src;
+ std::unordered_map<tirx::PrimVar, PrimExpr, ffi::ObjectPtrHash, ffi::ObjectPtrEqual>
+ rhs_to_index_map_tgt;
+ std::unordered_set<tirx::PrimVar, ffi::ObjectPtrHash, ffi::ObjectPtrEqual> unmapped_index_map_src;
TVM_FFI_ICHECK_EQ(mapping_info->lhs_iters.size(), index_map->initial_indices.size());
for (int i = 0; i < static_cast<int>(mapping_info->lhs_iters.size()); ++i) {
lhs_to_index_map_src[mapping_info->lhs_iters[i]->var] = index_map->initial_indices[i];
@@ -814,9 +816,9 @@ ffi::Optional<LoopRV> MultiLevelTilingTensorCoreNode::TransformWithTensorIntrin(
static_cast<int>(mapping_info->rhs_iters.size());
TVM_FFI_ICHECK_GE(offset, 0);
for (int i = 0; i < offset; ++i) {
- const tirx::VarNode* var_ptr = index_map->final_indices[i].as<tirx::VarNode>();
- TVM_FFI_ICHECK(var_ptr != nullptr);
- unmapped_index_map_src.insert(ffi::GetRef<tirx::Var>(var_ptr));
+ auto var = index_map->final_indices[i].as<tirx::PrimVar>();
+ TVM_FFI_ICHECK(var.has_value());
+ unmapped_index_map_src.insert(var.value());
}
for (int i = offset; i < static_cast<int>(index_map->final_indices.size()); ++i) {
rhs_to_index_map_tgt[mapping_info->rhs_iters[i - offset]->var] = index_map->final_indices[i];
@@ -824,31 +826,25 @@ ffi::Optional<LoopRV> MultiLevelTilingTensorCoreNode::TransformWithTensorIntrin(
auto f_get_sub_index_map = [&](const tirx::Buffer& lhs_buffer,
const ffi::Array<Range>& lhs_region) {
- std::vector<tirx::Var> sub_index_map_src;
+ std::vector<tirx::PrimVar> sub_index_map_src;
std::vector<PrimExpr> sub_index_map_tgt;
const tirx::Buffer& rhs_buffer = mapping_info->lhs_buffer_map[lhs_buffer];
for (const Range& range : lhs_region) {
TVM_FFI_ICHECK(tirx::is_one(range->extent));
- const tirx::VarNode* var_ptr = range->min.as<tirx::VarNode>();
- TVM_FFI_ICHECK(var_ptr != nullptr);
- const tirx::Var& lhs_representer = lhs_to_index_map_src[ffi::GetRef<tirx::Var>(var_ptr)];
+ auto var = range->min.as<tirx::PrimVar>();
+ TVM_FFI_ICHECK(var.has_value());
+ const tirx::PrimVar& lhs_representer = lhs_to_index_map_src[var.value()];
sub_index_map_src.push_back(lhs_representer);
if (unmapped_index_map_src.count(lhs_representer)) {
- sub_index_map_tgt.push_back(lhs_representer.as_or_throw<PrimExpr>());
+ sub_index_map_tgt.push_back(lhs_representer);
}
}
for (size_t i = 0; i < mapping_info->rhs_buffer_indices[rhs_buffer].size(); ++i) {
- const tirx::VarNode* var =
- mapping_info->rhs_buffer_indices[rhs_buffer][i].as<tirx::VarNode>();
- TVM_FFI_ICHECK(var != nullptr);
- sub_index_map_tgt.push_back(rhs_to_index_map_tgt[ffi::GetRef<tirx::Var>(var)]);
+ auto var = mapping_info->rhs_buffer_indices[rhs_buffer][i].as<tirx::PrimVar>();
+ TVM_FFI_ICHECK(var.has_value());
+ sub_index_map_tgt.push_back(rhs_to_index_map_tgt[var.value()]);
}
- ffi::Array<tirx::PrimVar> prim_sub_index_map_src;
- prim_sub_index_map_src.reserve(sub_index_map_src.size());
- for (const tirx::Var& var : sub_index_map_src) {
- prim_sub_index_map_src.push_back(var.as_or_throw<tirx::PrimVar>());
- }
- return tirx::IndexMap(prim_sub_index_map_src, sub_index_map_tgt);
+ return tirx::IndexMap(sub_index_map_src, sub_index_map_tgt);
};
std::unordered_set<tirx::Buffer, ffi::ObjectPtrHash, ffi::ObjectPtrEqual> visited_buffers;
diff --git a/src/s_tir/meta_schedule/schedule_rule/multi_level_tiling_wide_vector.cc b/src/s_tir/meta_schedule/schedule_rule/multi_level_tiling_wide_vector.cc
index d9f49538b2..e11ad4ded1 100644
--- a/src/s_tir/meta_schedule/schedule_rule/multi_level_tiling_wide_vector.cc
+++ b/src/s_tir/meta_schedule/schedule_rule/multi_level_tiling_wide_vector.cc
@@ -95,7 +95,7 @@ MultiLevelTilingWideVectorNode::SplitLoop(const Schedule& sch, SBlockRV block_rv
const size_t innermost_axis = block_node->writes[0]->region.size() - 1;
const PrimExpr innermost_iter_value = block_realize->iter_values[innermost_axis];
- if (!arith::Analyzer()->CanProve(loop->loop_var == innermost_iter_value)) {
+ if (!arith::Analyzer()->CanProve(static_cast<PrimExpr>(loop->loop_var) == innermost_iter_value)) {
// If this is not the innermost spatial loop, split the loop in the normal way.
return MultiLevelTilingNode::SplitLoop(sch, block_rv, loop_rv, n_tiles);
} else {
diff --git a/src/s_tir/schedule/analysis/analysis.cc b/src/s_tir/schedule/analysis/analysis.cc
index 5eee45c482..d1cbb409a2 100644
--- a/src/s_tir/schedule/analysis/analysis.cc
+++ b/src/s_tir/schedule/analysis/analysis.cc
@@ -143,8 +143,8 @@ ScopeBlockLoopInfo GetScopeBlockLoopInfo(const SBlock& scope_block) {
vars = &result.non_spatial_vars;
}
PostOrderVisit(iter_value, [vars](const ffi::ObjectRef& obj) {
- if (const VarNode* var = obj.as<VarNode>()) {
- vars->insert(var);
+ if (auto var = obj.as<PrimVar>()) {
+ vars->insert(var.value().get());
}
});
}
@@ -560,9 +560,13 @@ bool IsAffineBinding(const SBlockRealize& realize, const ffi::Map<Var, Range>& l
if (loop_var_ranges.empty()) {
return true;
}
+ ffi::Map<PrimVar, Range> primitive_loop_var_ranges;
+ for (const auto& [var, range] : loop_var_ranges) {
+ primitive_loop_var_ranges.Set(var.as_or_throw<PrimVar>(), range);
+ }
auto res = arith::DetectIterMap(
/*indices=*/realize->iter_values,
- /*input_iters=*/loop_var_ranges,
+ /*input_iters=*/primitive_loop_var_ranges,
/*predicate=*/realize->predicate,
/*check_level=*/arith::IterMapLevel::Surjective,
/*analyzer=*/ffi::GetRef<arith::Analyzer>(analyzer),
@@ -1387,8 +1391,8 @@ AnalyzeReadWritePattern(const BufferRegion& read_region, const BufferRegion& wri
if (as_const_int(dom->extent) == nullptr) {
return kNotExist;
}
- if (const auto* v = dom->min.as<VarNode>()) {
- var2idx.emplace(v, i);
+ if (auto var = dom->min.as<PrimVar>()) {
+ var2idx.emplace(var.value().get(), i);
} else {
return kNotExist;
}
@@ -1998,8 +2002,8 @@ class AutoTensorizeMappingProposer {
for (const auto& it : extractor_->rhs_buffer_indices_map_) {
const Buffer& rhs_buffer = it.first;
for (const PrimExpr& rhs_index : it.second) {
- if (const VarNode* var_node = rhs_index.as<VarNode>()) {
- update_mask(var_node, &rhs_buffer_masks, rhs_buffer_index.at(rhs_buffer));
+ if (auto var = rhs_index.as<PrimVar>()) {
+ update_mask(var.value().get(), &rhs_buffer_masks, rhs_buffer_index.at(rhs_buffer));
} else {
TVM_FFI_THROW(ValueError)
<< "Buffer index " << rhs_index
@@ -2012,8 +2016,8 @@ class AutoTensorizeMappingProposer {
const Buffer& lhs_buffer = lhs_buffer_it->second;
for (const PrimExpr& index : extractor_->lhs_buffer_indices_map_.at(lhs_buffer)) {
PreOrderVisit(index, [&](const ffi::ObjectRef& obj) -> bool {
- if (const VarNode* var = obj.as<VarNode>()) {
- update_mask(var, &lhs_buffer_masks, lhs_buffer_index.at(lhs_buffer));
+ if (auto var = obj.as<PrimVar>()) {
+ update_mask(var.value().get(), &lhs_buffer_masks, lhs_buffer_index.at(lhs_buffer));
}
return true;
});
diff --git a/src/s_tir/schedule/analysis/layout.cc b/src/s_tir/schedule/analysis/layout.cc
index c46e02079a..5619422ff2 100644
--- a/src/s_tir/schedule/analysis/layout.cc
+++ b/src/s_tir/schedule/analysis/layout.cc
@@ -60,7 +60,7 @@ class SplitExprCollector {
*/
struct SplitExpr {
/*! \brief The source variable */
- Var source;
+ PrimVar source;
/*! \brief The lower factor of the split expression */
int64_t lower_factor;
/*! \brief The extent of the split expression */
@@ -77,9 +77,9 @@ class SplitExprCollector {
* \return The collected split expressions
*/
static std::vector<SplitExpr> Collect(const PrimExpr& index,
- const ffi::Map<Var, Range>& input_iters, //
- const PrimExpr& predicate, //
- arith::IterMapLevel check_level, //
+ const ffi::Map<PrimVar, Range>& input_iters, //
+ const PrimExpr& predicate, //
+ arith::IterMapLevel check_level, //
arith::AnalyzerObj* analyzer) {
arith::Analyzer analyzer_ref = ffi::GetRef<arith::Analyzer>(analyzer);
arith::IterMapResult res = arith::DetectIterMap({analyzer->Simplify(index)}, input_iters,
@@ -102,14 +102,14 @@ class SplitExprCollector {
private:
void Visit(const arith::IterSplitExpr& expr) {
- if (const auto* var = expr->source->source.as<tirx::VarNode>()) {
+ if (auto var = expr->source->source.as<PrimVar>()) {
const int64_t* lower_factor = as_const_int(expr->lower_factor);
const int64_t* extent = as_const_int(expr->extent);
if (lower_factor == nullptr || extent == nullptr) {
failed_ = true;
return;
}
- exprs_.push_back(SplitExpr{ffi::GetRef<Var>(var), *lower_factor, *extent});
+ exprs_.push_back(SplitExpr{var.value(), *lower_factor, *extent});
} else if (auto iter_sum_expr = expr->source->source.as<arith::IterSumExpr>()) {
Visit(iter_sum_expr.value());
} else {
@@ -135,7 +135,7 @@ ffi::Optional<IndexMap> SuggestIndexMap(const Buffer& buffer, const ffi::Array<P
int ndim = buffer->shape.size();
int n_loops = loops.size();
// Step 1. Collect the domains and indices of loop variables
- ffi::Map<Var, Range> input_iters;
+ ffi::Map<PrimVar, Range> input_iters;
std::unordered_map<const VarNode*, int> var2id;
var2id.reserve(n_loops);
for (int i = 0; i < n_loops; ++i) {
diff --git a/src/s_tir/schedule/analysis/reducer.cc b/src/s_tir/schedule/analysis/reducer.cc
index 7fb24c4cde..4eb5affe70 100644
--- a/src/s_tir/schedule/analysis/reducer.cc
+++ b/src/s_tir/schedule/analysis/reducer.cc
@@ -408,10 +408,11 @@ void ExtractReductionUpdates(const ffi::Optional<ScheduleState>& self, SBlock bl
if (buf_store == nullptr) {
ErrorRFactorCrossThreadReductionNotApplicable(self, std::move(block), /*violated_cond=*/5);
}
- const auto* var = buf_store->value.as<VarNode>();
- if (var == nullptr) {
+ auto var_ref = buf_store->value.as<PrimVar>();
+ if (!var_ref.has_value()) {
ErrorRFactorCrossThreadReductionNotApplicable(self, std::move(block), /*violated_cond=*/7);
}
+ const VarNode* var = var_ref.value().get();
auto it = var2index.find(var);
if (it == var2index.end()) {
ErrorRFactorCrossThreadReductionNotApplicable(self, std::move(block), /*violated_cond=*/7);
diff --git a/src/s_tir/schedule/instruction.cc b/src/s_tir/schedule/instruction.cc
index e1e546ed38..99b2c68a84 100644
--- a/src/s_tir/schedule/instruction.cc
+++ b/src/s_tir/schedule/instruction.cc
@@ -81,9 +81,7 @@ ffi::String InstructionAsPythonRepr(const InstructionNode* self) {
inputs.push_back(obj);
} else if (auto expr = obj.as<PrimExpr>()) {
PrimExpr new_expr = Substitute(expr.value(), [](const Var& var) -> ffi::Optional<Expr> {
- ffi::ObjectPtr<VarNode> new_var = ffi::make_object<VarNode>(*var.get());
- new_var->name_hint = "_";
- return Var(new_var).as_or_throw<PrimExpr>();
+ return Var("_", var->ty, var->span).as_or_throw<PrimExpr>();
});
std::ostringstream os;
os << new_expr;
diff --git a/src/s_tir/schedule/ir_comparator.cc b/src/s_tir/schedule/ir_comparator.cc
index 4b674c4b18..3917601f51 100644
--- a/src/s_tir/schedule/ir_comparator.cc
+++ b/src/s_tir/schedule/ir_comparator.cc
@@ -365,7 +365,9 @@ bool TensorizeComparator::VisitExpr_(const CastNode* op, const PrimExpr& other)
}
bool TensorizeComparator::VisitExpr_(const VarNode* op, const PrimExpr& other) {
- const auto* rhs = other.as<VarNode>();
+ auto rhs_ref = other.as<PrimVar>();
+ if (!rhs_ref.has_value()) return false;
+ const auto* rhs = rhs_ref.value().get();
auto lhs = ffi::GetRef<Var>(op);
if (lhs.same_as(other)) return true;
PrimType lhs_ty = op->ty.as_or_throw<PrimType>();
@@ -795,7 +797,7 @@ bool AutoTensorizeComparator::CompareBufferAccess(const T* lhs, const T* rhs) {
};
for (const auto& index : rhs->indices) {
- if (!index.template as<VarNode>() && !is_scalar_access(rhs->indices, index)) return false;
+ if (!index.template as<PrimVar>() && !is_scalar_access(rhs->indices, index)) return false;
}
lhs_buffer_indices_map_[lhs->buffer] = lhs_indices;
rhs_buffer_indices_map_[rhs->buffer] = rhs->indices;
diff --git a/src/s_tir/schedule/primitive/annotate_buffer_access.cc b/src/s_tir/schedule/primitive/annotate_buffer_access.cc
index 7041f1ce83..10887e2c4d 100644
--- a/src/s_tir/schedule/primitive/annotate_buffer_access.cc
+++ b/src/s_tir/schedule/primitive/annotate_buffer_access.cc
@@ -132,18 +132,29 @@ struct AnnotateBufferAccessTraits : public UnpackedInstTraits<AnnotateBufferAcce
static ffi::String IndexMap2GenNewRangesLambda(const IndexMap& index_map) {
std::ostringstream oss;
+ auto print_expr = [&oss](const PrimExpr& expr) {
+ if (auto var = expr.as<PrimVar>()) {
+ oss << var.value()->name_hint;
+ } else {
+ oss << expr;
+ }
+ };
oss << "lambda ";
for (size_t i = 0; i < index_map->initial_indices.size(); ++i) {
if (i != 0) oss << ", ";
- oss << index_map->initial_indices[i];
+ oss << index_map->initial_indices[i]->name_hint;
}
oss << ": [";
for (size_t i = 0; i < index_map->final_indices.size(); i += 2) {
if (i != 0) oss << ", ";
if (index_map->final_indices[i].same_as(index_map->final_indices[i + 1])) {
- oss << index_map->final_indices[i];
+ print_expr(index_map->final_indices[i]);
} else {
- oss << "(" << index_map->final_indices[i] << ", " << index_map->final_indices[i + 1] << ")";
+ oss << "(";
+ print_expr(index_map->final_indices[i]);
+ oss << ", ";
+ print_expr(index_map->final_indices[i + 1]);
+ oss << ")";
}
}
oss << "]";
diff --git a/src/s_tir/schedule/primitive/blockize_tensorize.cc b/src/s_tir/schedule/primitive/blockize_tensorize.cc
index 8206a99cfd..4054edb9e9 100644
--- a/src/s_tir/schedule/primitive/blockize_tensorize.cc
+++ b/src/s_tir/schedule/primitive/blockize_tensorize.cc
@@ -122,7 +122,7 @@ ffi::Array<ffi::Array<arith::IterMark>> TrivialSubspaceDivision(
bool outer = use_outer_loop_vars(bindings[i]);
bool inner = use_inner_loop_vars(bindings[i]);
arith::IterMark iter_mark;
- if (bindings[i]->IsInstance<VarNode>()) {
+ if (bindings[i].as<PrimVar>()) {
iter_mark = arith::IterMark(
arith::IterSplitExpr(arith::IterMark(bindings[i], iter_vars[i]->dom->extent)),
iter_vars[i]->dom->extent);
@@ -169,7 +169,8 @@ ffi::Array<ffi::Array<arith::IterMark>> SubspaceDivide(const SBlockRealize& real
bool loop_sref_as_outer = false) {
ffi::Array<Var> inner_vars;
ffi::Array<Var> outer_vars;
- ffi::Map<Var, Range> loop_var_domain;
+ ffi::Array<PrimVar> primitive_inner_vars;
+ ffi::Map<PrimVar, Range> primitive_loop_var_domain;
bool inner = true;
for (StmtSRefNode* sref = block_sref->parent; //
sref && sref->stmt->IsInstance<ForNode>(); //
@@ -178,18 +179,19 @@ ffi::Array<ffi::Array<arith::IterMark>> SubspaceDivide(const SBlockRealize& real
if (inner) {
loops->push_back(loop);
inner_vars.push_back(loop->loop_var);
+ primitive_inner_vars.push_back(loop->loop_var);
} else {
outer_vars.push_back(loop->loop_var);
}
- loop_var_domain.Set(loop->loop_var, Range::FromMinExtent(loop->min, loop->extent));
+ primitive_loop_var_domain.Set(loop->loop_var, Range::FromMinExtent(loop->min, loop->extent));
if ((loop_sref_as_outer && sref->parent == loop_sref.get()) || sref == loop_sref.get()) {
inner = false;
}
}
- ffi::Array<ffi::Array<arith::IterMark>> result =
- arith::SubspaceDivide(realize->iter_values, loop_var_domain, inner_vars, realize->predicate,
- arith::IterMapLevel::Surjective, ffi::GetRef<arith::Analyzer>(analyzer),
- /*simplify_trivial_iterators=*/!preserve_unit_iters);
+ ffi::Array<ffi::Array<arith::IterMark>> result = arith::SubspaceDivide(
+ realize->iter_values, primitive_loop_var_domain, primitive_inner_vars, realize->predicate,
+ arith::IterMapLevel::Surjective, ffi::GetRef<arith::Analyzer>(analyzer),
+ /*simplify_trivial_iterators=*/!preserve_unit_iters);
if (!result.empty()) {
return result;
}
@@ -615,7 +617,7 @@ SBlockRealize BlockizeBlocks(const ScheduleState& self, const ffi::Array<StmtSRe
}
ffi::Map<Var, arith::IntSet> inner_iter_dom;
for (const IterVar& iter : inner_iter_vars) {
- Range dom = Substitute(iter->dom, loop_var_subst);
+ Range dom = tirx::Substitute(iter->dom, loop_var_subst);
inner_iter_dom.Set(iter->var, arith::IntSet::FromRange(dom));
analyzer->Bind(iter->var, dom);
}
@@ -648,7 +650,7 @@ SBlockRealize BlockizeBlocks(const ScheduleState& self, const ffi::Array<StmtSRe
for (const ForNode* loop : loops) {
ffi::ObjectPtr<ForNode> new_loop = ffi::make_object<ForNode>(*loop);
new_loop->body = std::move(stmt);
- new_loop->extent = Substitute(new_loop->extent, loop_var_subst);
+ new_loop->extent = tirx::Substitute(new_loop->extent, loop_var_subst);
stmt = For(new_loop);
}
seq_body.push_back(stmt);
diff --git a/src/s_tir/schedule/primitive/cache_index.cc b/src/s_tir/schedule/primitive/cache_index.cc
index 4a273e5678..9871f18fd3 100644
--- a/src/s_tir/schedule/primitive/cache_index.cc
+++ b/src/s_tir/schedule/primitive/cache_index.cc
@@ -235,8 +235,8 @@ ffi::Array<SBlock> MakeIndexCacheStage(IndexInfo* info, const ffi::String& stora
// Collect the block vars in original index computation
info->origin_block_vars.push_back({});
PostOrderVisit(index_expr, [&info, &expr_index](const ffi::ObjectRef& node) {
- if (node->IsInstance<VarNode>()) {
- Var iter_var = node.as_or_throw<Var>();
+ if (auto var = node.as<PrimVar>()) {
+ Var iter_var = var.value();
const ffi::Array<Var>& origin_block_var = info->origin_block_vars[expr_index];
auto find_result = std::find_if(origin_block_var.begin(), origin_block_var.end(),
[&](Var it) { return it.get() == iter_var.get(); });
@@ -251,8 +251,8 @@ ffi::Array<SBlock> MakeIndexCacheStage(IndexInfo* info, const ffi::String& stora
std::vector<Var> iter_vars;
for (const Var& it : info->origin_block_vars[expr_index]) {
PostOrderVisit(info->var_binding.at(it), [/*&info,*/ &iter_vars](const ffi::ObjectRef& node) {
- if (node->IsInstance<VarNode>()) {
- Var iter_var = node.as_or_throw<Var>();
+ if (auto var = node.as<PrimVar>()) {
+ Var iter_var = var.value();
if (std::find_if(iter_vars.begin(), iter_vars.end(),
[&](Var it) { return it.get() == iter_var.get(); }) == iter_vars.end()) {
iter_vars.push_back(iter_var);
diff --git a/src/s_tir/schedule/primitive/cache_index_helpers.cc b/src/s_tir/schedule/primitive/cache_index_helpers.cc
index e71702b8f1..f865dd17c5 100644
--- a/src/s_tir/schedule/primitive/cache_index_helpers.cc
+++ b/src/s_tir/schedule/primitive/cache_index_helpers.cc
@@ -150,7 +150,7 @@ ComputationTable ComputationsDoneBy::GetComputationsDoneBy(
const PrimExpr& expr, std::function<bool(const PrimExpr&)> is_eligible_computation,
std::function<bool(const PrimExpr&)> can_contain_computations) {
if (expr.as<IntImmNode>() != nullptr || expr.as<FloatImmNode>() != nullptr ||
- expr.as<StringImmNode>() != nullptr || expr.as<VarNode>() != nullptr) {
+ expr.as<StringImmNode>() != nullptr || expr.as<PrimVar>()) {
return {};
}
@@ -204,7 +204,7 @@ void ComputationsDoneBy::VisitExpr(const Expr& expr_value) {
}
PrimExpr expr = opt_expr.value();
if (expr.as<IntImmNode>() != nullptr || expr.as<FloatImmNode>() != nullptr ||
- expr.as<StringImmNode>() != nullptr || expr.as<VarNode>() != nullptr) {
+ expr.as<StringImmNode>() != nullptr || expr.as<PrimVar>()) {
return;
}
diff --git a/src/s_tir/schedule/primitive/compute_at.cc b/src/s_tir/schedule/primitive/compute_at.cc
index df1f5ed480..e622602a5f 100644
--- a/src/s_tir/schedule/primitive/compute_at.cc
+++ b/src/s_tir/schedule/primitive/compute_at.cc
@@ -435,7 +435,7 @@ std::pair<Var, BlockVarDomainInfo> SolveBlockVarDomain(const arith::IntSet& prov
PrimExpr var_expr = p_f1.Eval();
PrimExpr fac = p_f2.Eval();
if (analyzer->CanProveGreaterEqual(fac, 1)) {
- if (var_expr->IsInstance<VarNode>()) {
+ if (var_expr.as<PrimVar>()) {
// a <= (x // factor) <= b, fac > 0 ==> (a * fac) <= x <= (b * fac + fac - 1)
var = var_expr.as_or_throw<Var>();
var_dom = arith::IntSet::Interval(required_min * fac,
@@ -450,7 +450,7 @@ std::pair<Var, BlockVarDomainInfo> SolveBlockVarDomain(const arith::IntSet& prov
}
} else if ((floormod(p_f1, p_f2).Match(provided_min))) {
PrimExpr var_expr = p_f1.Eval();
- if (var_expr->IsInstance<VarNode>()) {
+ if (var_expr.as<PrimVar>()) {
// generally domain of (x % fac) enforce no constraints to domain of x
Var var_mod = var_expr.as_or_throw<Var>();
return {var_mod, BlockVarDomainInfo()};
@@ -561,7 +561,7 @@ bool UpdateBlockVarDomainAffine(const BufferNode* buffer, const ffi::Array<IterV
if (!intset.CanProveSinglePoint(analyzer_ref)) return false;
}
// calculate forward mapping (block vars -> provided region point)
- ffi::Map<Var, Range> dom_map;
+ ffi::Map<PrimVar, Range> dom_map;
for (const IterVar& iter_var : iter_vars) {
dom_map.Set(iter_var->var, iter_var->dom);
}
diff --git a/src/s_tir/schedule/primitive/compute_inline.cc b/src/s_tir/schedule/primitive/compute_inline.cc
index 73ef0f0032..1d90e5e274 100644
--- a/src/s_tir/schedule/primitive/compute_inline.cc
+++ b/src/s_tir/schedule/primitive/compute_inline.cc
@@ -496,7 +496,7 @@ class ComputeInliner : public BaseInliner {
// If the mapping for store indices is non-trivial
// check bijective mapping from producer iter var to store indices
- ffi::Map<Var, Range> producer_iter_doms;
+ ffi::Map<PrimVar, Range> producer_iter_doms;
for (const auto& iter : producer_block->iter_vars) {
producer_iter_doms.Set(iter->var, iter->dom);
}
@@ -629,9 +629,9 @@ class ReverseComputeInliner : public BaseInliner {
// Initialize the predicates to ensure consumer block iters are in-bound
consumer_iter_in_bound_ = IntImm::Bool(true);
for (const IterVar& iter : consumer_block_realize->block->iter_vars) {
- consumer_iter_in_bound_ =
- consumer_iter_in_bound_ &&
- (iter->var >= iter->dom->min && iter->var < iter->dom->min + iter->dom->extent);
+ consumer_iter_in_bound_ = consumer_iter_in_bound_ && (iter->var >= iter->dom->min &&
+ static_cast<PrimExpr>(iter->var) <
+ iter->dom->min + iter->dom->extent);
}
}
@@ -653,7 +653,7 @@ class ReverseComputeInliner : public BaseInliner {
}
// Collect block iter domains and update the substition map
- ffi::Map<Var, Range> consumer_iter_doms;
+ ffi::Map<PrimVar, Range> consumer_iter_doms;
for (const auto& iter_var : consumer_block->iter_vars) {
consumer_iter_doms.Set(iter_var->var, iter_var->dom);
// Set default mapping for unit iters
diff --git a/src/s_tir/schedule/primitive/decompose_padding.cc b/src/s_tir/schedule/primitive/decompose_padding.cc
index b0e8ed0ddc..8973fbea40 100644
--- a/src/s_tir/schedule/primitive/decompose_padding.cc
+++ b/src/s_tir/schedule/primitive/decompose_padding.cc
@@ -70,7 +70,7 @@ class PaddingPatternMatchError : public ScheduleError {
class PaddingInfoAnalyzer {
public:
static PaddingSBlockInfo CheckAndGetPaddingInfo(IRModule mod, const SBlockRealizeNode* realize,
- const ffi::Map<Var, Range>& dom_map,
+ const ffi::Map<PrimVar, Range>& dom_map,
arith::AnalyzerObj* analyzer) {
PaddingInfoAnalyzer padding_analyzer(analyzer);
if (!padding_analyzer.MatchPadding(realize, dom_map)) {
@@ -83,7 +83,7 @@ class PaddingInfoAnalyzer {
explicit PaddingInfoAnalyzer(arith::AnalyzerObj* analyzer) : analyzer_(analyzer) {}
/*! \brief Detect padding pattern and update result. */
- bool MatchPadding(const SBlockRealizeNode* realize, const ffi::Map<Var, Range>& dom_map) {
+ bool MatchPadding(const SBlockRealizeNode* realize, const ffi::Map<PrimVar, Range>& dom_map) {
// Step 1. Check match padding computation pattern.
// A[...] = T.if_then_else(predicate, B[...], imm)
SBlock block = realize->block;
@@ -160,7 +160,7 @@ class PaddingInfoAnalyzer {
/*! \brief Return iteration region of block vars where the padding predicate evals to true. */
ffi::Array<Range> EstimateInBoundRegion(const ffi::Array<PrimExpr>& iter_values,
- const ffi::Map<Var, Range>& dom_map,
+ const ffi::Map<PrimVar, Range>& dom_map,
const PrimExpr& in_bound_predicate) {
ffi::Array<Range> region;
@@ -416,7 +416,7 @@ StmtSRef DecomposePaddingImpl(ScheduleState self, const StmtSRef& block_sref,
// Condition Checks and Information Collection
const SBlockNode* block = TVM_SREF_TO_SBLOCK(block_sref);
const SBlockRealizeNode* realize = GetSBlockRealize(self, block_sref).get();
- ffi::Map<Var, Range> dom_map;
+ ffi::Map<PrimVar, Range> dom_map;
arith::Analyzer analyzer;
// Check 1. check the block is complete.
diff --git a/src/s_tir/schedule/primitive/layout_transformation.cc b/src/s_tir/schedule/primitive/layout_transformation.cc
index 62f1699a25..de0ed83e28 100644
--- a/src/s_tir/schedule/primitive/layout_transformation.cc
+++ b/src/s_tir/schedule/primitive/layout_transformation.cc
@@ -1267,8 +1267,8 @@ IterVarType DetectNewBlockIterType(
IterVarType result{kOpaque};
bool found = false;
PostOrderVisit(expr, [&](const ffi::ObjectRef& obj) {
- if (const VarNode* var = obj.as<VarNode>()) {
- auto it = block_iter_type_map.find(var);
+ if (auto var = obj.as<PrimVar>()) {
+ auto it = block_iter_type_map.find(var.value().get());
if (it != block_iter_type_map.end()) {
if (!found) {
found = true;
diff --git a/src/s_tir/schedule/primitive/loop_transformation.cc b/src/s_tir/schedule/primitive/loop_transformation.cc
index 97bb5216e7..7a007058e9 100644
--- a/src/s_tir/schedule/primitive/loop_transformation.cc
+++ b/src/s_tir/schedule/primitive/loop_transformation.cc
@@ -81,7 +81,7 @@ class SubstituteVarAndCollectOpaqueBlock : public StmtExprMutator {
class IterMapSimplifyBlockBinding : public StmtExprMutator {
public:
explicit IterMapSimplifyBlockBinding(ffi::MapObj* opaque_blocks,
- ffi::Map<Var, Range> loop_var2extent,
+ ffi::Map<PrimVar, Range> loop_var2extent,
bool preserve_unit_iters)
: opaque_blocks_(opaque_blocks),
loop_var2extent_(loop_var2extent),
@@ -89,7 +89,7 @@ class IterMapSimplifyBlockBinding : public StmtExprMutator {
static For SimplifyBindings(Stmt stmt, const ffi::Array<StmtSRef>& loop_srefs,
ffi::MapObj* opaque_blocks, bool preserve_unit_iters) {
- ffi::Map<Var, Range> loop_var2extent;
+ ffi::Map<PrimVar, Range> loop_var2extent;
for (const StmtSRef& sref : loop_srefs) {
const ForNode* loop = TVM_SREF_TO_FOR(sref);
loop_var2extent.Set(loop->loop_var, Range::FromMinExtent(loop->min, loop->extent));
@@ -139,7 +139,7 @@ class IterMapSimplifyBlockBinding : public StmtExprMutator {
/*! \brief The reuse mapping */
ffi::MapObj* opaque_blocks_;
/*! \brief The range of loops */
- ffi::Map<Var, Range> loop_var2extent_;
+ ffi::Map<PrimVar, Range> loop_var2extent_;
/*! \brief Internal analyzer */
arith::Analyzer analzyer_;
/*! \brief Whether or not to simplify unit iterators */
@@ -480,12 +480,12 @@ class BufferIndicesMapExtractor : public StmtExprVisitor {
ffi::Array<ffi::String> indices;
bool check_ = false;
for (size_t i = 0; i < store->indices.size(); i++) {
- const VarNode* var_node = store->indices[i].as<VarNode>();
- if (var_node == nullptr) {
+ auto var = store->indices[i].as<PrimVar>();
+ if (!var.has_value()) {
check_ = true;
break;
}
- indices.push_back(var_node->name_hint);
+ indices.push_back(var.value()->name_hint);
}
if (buffer_indices_map.find(store->buffer->name) == buffer_indices_map.end() && !check_)
buffer_indices_map.Set(store->buffer->name, indices);
@@ -496,12 +496,12 @@ class BufferIndicesMapExtractor : public StmtExprVisitor {
ffi::Array<ffi::String> indices;
bool check_ = false;
for (size_t i = 0; i < load->indices.size(); i++) {
- const VarNode* var_node = load->indices[i].as<VarNode>();
- if (var_node == nullptr) {
+ auto var = load->indices[i].as<PrimVar>();
+ if (!var.has_value()) {
check_ = true;
break;
}
- indices.push_back(var_node->name_hint);
+ indices.push_back(var.value()->name_hint);
}
if (buffer_indices_map.find(load->buffer->name) == buffer_indices_map.end() && !check_)
buffer_indices_map.Set(load->buffer->name, indices);
diff --git a/src/s_tir/schedule/primitive/pad_einsum.cc b/src/s_tir/schedule/primitive/pad_einsum.cc
index d39cbfa01c..e10f8ff7d7 100644
--- a/src/s_tir/schedule/primitive/pad_einsum.cc
+++ b/src/s_tir/schedule/primitive/pad_einsum.cc
@@ -38,11 +38,11 @@ ffi::Optional<ffi::Array<Var>> CheckTrivialBufferIndices(
if (index->IsInstance<IntImmNode>()) {
continue;
}
- const VarNode* var = index.as<VarNode>();
- if (var == nullptr) {
+ auto var = index.as<PrimVar>();
+ if (!var.has_value()) {
return std::nullopt;
}
- indices.push_back(ffi::GetRef<Var>(var));
+ indices.push_back(var.value());
}
return indices;
}
@@ -57,8 +57,8 @@ ffi::Optional<ffi::Array<Var>> CheckTrivialBufferAccess(const BufferRegion& buff
if (range->min->IsInstance<IntImmNode>()) {
continue;
}
- if (const auto* var = range->min.as<VarNode>()) {
- indices.push_back(ffi::GetRef<Var>(var));
+ if (auto var = range->min.as<PrimVar>()) {
+ indices.push_back(var.value());
} else {
return std::nullopt;
}
@@ -145,7 +145,7 @@ struct BufferPadding {
int ndim = buffer_region->region.size();
for (int i = 0; i < ndim; ++i) {
PrimExpr pos = buffer_region->region[i]->min;
- TVM_FFI_ICHECK(pos->IsInstance<IntImmNode>() || pos->IsInstance<VarNode>());
+ TVM_FFI_ICHECK(pos->IsInstance<IntImmNode>() || pos.as<PrimVar>());
if (pos->IsInstance<IntImmNode>()) {
shape.push_back(IntImm(pos.ty(), 1));
} else if (ffi::Optional<PrimExpr> extent = iter_extents.Get(pos.as_or_throw<Var>())) {
@@ -395,16 +395,16 @@ void PadEinsum(ScheduleState self, const StmtSRef& block_sref, const ffi::Array<
PrimExpr new_dom = analyzer->Simplify(ceildiv(dom, pad_imm) * pad_imm);
if (!analyzer->CanProveEqual(new_dom, dom)) {
replacer.iter2padded_extents.Set(iter->var, new_dom);
- if (const auto* loop_var = realize->iter_values[i].as<VarNode>()) {
- replacer.iter2padded_extents.Set(ffi::GetRef<Var>(loop_var), new_dom);
- replacer.loop_var2padded_extent.Set(ffi::GetRef<Var>(loop_var), new_dom);
+ if (auto loop_var = realize->iter_values[i].as<PrimVar>()) {
+ replacer.iter2padded_extents.Set(loop_var.value(), new_dom);
+ replacer.loop_var2padded_extent.Set(loop_var.value(), new_dom);
}
}
}
auto f_needs_padding = [&replacer](const ffi::Array<Range>& region) {
for (const Range& range : region) {
- if (const auto* var = range->min.as<VarNode>()) {
- if (replacer.iter2padded_extents.count(ffi::GetRef<Var>(var))) {
+ if (auto var = range->min.as<PrimVar>()) {
+ if (replacer.iter2padded_extents.count(var.value())) {
return true;
}
}
diff --git a/src/s_tir/schedule/primitive/read_write_at.cc b/src/s_tir/schedule/primitive/read_write_at.cc
index 789ac8f9bf..8151894e39 100644
--- a/src/s_tir/schedule/primitive/read_write_at.cc
+++ b/src/s_tir/schedule/primitive/read_write_at.cc
@@ -272,7 +272,7 @@ struct ReadWriteAtImpl {
std::vector<Var> loop_vars;
loop_vars.reserve(n);
for (int i = 0; i < n; ++i) {
- loop_vars.push_back(Var("ax" + std::to_string(i)));
+ loop_vars.push_back(PrimVar("ax" + std::to_string(i)));
}
ffi::Map<Var, PrimExpr> bindings;
ffi::Array<IterVar> iter_vars;
@@ -289,12 +289,11 @@ struct ReadWriteAtImpl {
return (*it).second;
}
Range range = loop_domain.at(var);
- ffi::ObjectPtr<VarNode> v = ffi::make_object<VarNode>(*var.get());
- v->name_hint = "v" + std::to_string(iter_vars.size());
- bindings.Set(var, Var(v).as_or_throw<PrimExpr>());
+ Var v("v" + std::to_string(iter_vars.size()), var->ty, var->span);
+ bindings.Set(var, v.as_or_throw<PrimExpr>());
iter_values.push_back(var.as_or_throw<PrimExpr>());
- iter_vars.push_back(IterVar(range, Var(v).as_or_throw<PrimVar>(), IterVarType::kDataPar));
- return Var(v).as_or_throw<PrimExpr>();
+ iter_vars.push_back(IterVar(range, v.as_or_throw<PrimVar>(), IterVarType::kDataPar));
+ return v.as_or_throw<PrimExpr>();
};
ffi::ObjectPtr<RangeNode> dom = ffi::make_object<RangeNode>(*domain[i].get());
dom->min = Substitute(std::move(dom->min), f_substitute);
diff --git a/src/s_tir/schedule/transform.cc b/src/s_tir/schedule/transform.cc
index a0bf3f6ae7..ea67a44ce8 100644
--- a/src/s_tir/schedule/transform.cc
+++ b/src/s_tir/schedule/transform.cc
@@ -41,10 +41,9 @@ SBlock WithAnnotation(const SBlockNode* block, const ffi::String& attr_key,
/******** Buffer Related ********/
Buffer WithScope(const Buffer& buffer, const ffi::String& scope) {
ffi::ObjectPtr<BufferNode> new_buffer = ffi::make_object<BufferNode>(*buffer.get());
- ffi::ObjectPtr<VarNode> new_var = ffi::make_object<VarNode>(*buffer->data.get());
const auto* ptr_type = TVM_TYPE_AS(buffer->data->ty, PointerTypeNode);
- new_var->ty = PointerType(ptr_type->element_type, scope);
- new_buffer->data = Var(new_var->name_hint + "_" + scope, new_var->ty);
+ Type new_type = PointerType(ptr_type->element_type, scope);
+ new_buffer->data = tirx::Var(buffer->data->name_hint + "_" + scope, new_type);
new_buffer->name = buffer->name + "_" + scope;
return Buffer(new_buffer);
}
@@ -53,7 +52,8 @@ Buffer WithDType(const Buffer& buffer, PrimType dtype) {
ffi::ObjectPtr<BufferNode> new_buffer = ffi::make_object<BufferNode>(*buffer.get());
new_buffer->dtype = dtype;
const auto* ptr_type = TVM_TYPE_AS(buffer->data->ty, PointerTypeNode);
- new_buffer->data = Var(buffer->data->name_hint, PointerType(dtype, ptr_type->storage_scope));
+ new_buffer->data =
+ tirx::Var(buffer->data->name_hint, PointerType(dtype, ptr_type->storage_scope));
new_buffer->name = buffer->name;
return Buffer(new_buffer);
}
@@ -523,13 +523,13 @@ ffi::Optional<ffi::ObjectRef> NormalizePrimFunc(Schedule sch) {
for (const SBlockRV& block : leaf_blocks) {
ffi::Array<IterVar> iters = sch->Get(block)->iter_vars;
bool has_spatial_iter = false;
- ffi::Array<Var> index_map_inputs;
+ ffi::Array<PrimVar> index_map_inputs;
ffi::Array<PrimExpr> index_map_outputs;
for (const IterVar& iter : sch->Get(block)->iter_vars) {
- Var var = iter->var.CopyWithSuffix("");
+ PrimVar var = iter->var.CopyWithSuffix("");
index_map_inputs.push_back(var);
if (!is_one(iter->dom->extent)) {
- index_map_outputs.push_back(var.as_or_throw<PrimExpr>());
+ index_map_outputs.push_back(var);
if (iter->iter_type == IterVarType::kDataPar) {
has_spatial_iter = true;
}
@@ -539,9 +539,7 @@ ffi::Optional<ffi::ObjectRef> NormalizePrimFunc(Schedule sch) {
index_map_outputs.insert(index_map_outputs.begin(), IntImm::Int64(0));
}
try {
- sch->TransformBlockLayout(
- block, IndexMap(index_map_inputs.Map([](Var var) { return var.as_or_throw<PrimVar>(); }),
- index_map_outputs));
+ sch->TransformBlockLayout(block, IndexMap(index_map_inputs, index_map_outputs));
} catch (tvm::ffi::Error& e) {
// Skip layout transformation when not transformable.
}
diff --git a/src/s_tir/schedule/utils.h b/src/s_tir/schedule/utils.h
index b49c17f74a..e667f76e48 100644
--- a/src/s_tir/schedule/utils.h
+++ b/src/s_tir/schedule/utils.h
@@ -227,9 +227,9 @@ inline const int64_t* GetLoopIntExtent(const StmtSRef& loop_sref) {
*/
inline ffi::Optional<Var> AnalyzeVarWithShift(const PrimExpr& expr,
ffi::Optional<IntImm>* constant) {
- if (const auto* var = expr.as<VarNode>()) {
+ if (auto var = expr.as<PrimVar>()) {
*constant = std::nullopt;
- return ffi::GetRef<Var>(var);
+ return static_cast<Var>(var.value());
}
arith::PVar<Var> var;
arith::PVar<IntImm> shift;
diff --git a/src/s_tir/transform/inject_double_buffer.cc b/src/s_tir/transform/inject_double_buffer.cc
index 79d07737a4..151bc11f55 100644
--- a/src/s_tir/transform/inject_double_buffer.cc
+++ b/src/s_tir/transform/inject_double_buffer.cc
@@ -114,7 +114,7 @@ class DoubleBufferInjector : public StmtExprMutator {
}
Stmt VisitStmt_(const AllocBufferNode* op) final {
- const VarNode* buf = op->buffer->data.as<VarNode>();
+ const VarNode* buf = op->buffer->data.get();
auto it = dbuffer_info_.find(buf);
if (it != dbuffer_info_.end()) {
StorageEntry& entry = it->second;
diff --git a/src/s_tir/transform/lower_cross_thread_reduction.cc b/src/s_tir/transform/lower_cross_thread_reduction.cc
index c52d500547..c1fefb5c2e 100644
--- a/src/s_tir/transform/lower_cross_thread_reduction.cc
+++ b/src/s_tir/transform/lower_cross_thread_reduction.cc
@@ -444,8 +444,8 @@ Stmt TransformReductionBlock(const SBlockRealizeNode* realize,
IterVar new_iter_var{nullptr};
{
ffi::ObjectPtr<IterVarNode> n = ffi::make_object<IterVarNode>(*iter_var.get());
- ffi::ObjectPtr<VarNode> v = ffi::make_object<VarNode>(*iter_var->var.get());
- n->var = Var(v).as_or_throw<PrimVar>();
+ Var v(iter_var->var->name_hint, iter_var->var->ty, iter_var->var->span);
+ n->var = v.as_or_throw<PrimVar>();
new_iter_var = IterVar(n);
}
iter_vars.push_back(new_iter_var);
@@ -507,7 +507,8 @@ Stmt TransformReductionBlock(const SBlockRealizeNode* realize,
if (wb_buffers[0].scope() != "local") {
for (const ForNode* loop : reduction_loops) {
if (loop->thread_binding.has_value()) {
- wb_predicate = wb_predicate && (loop->loop_var == IntImm(loop->loop_var.ty(), 0));
+ wb_predicate = wb_predicate &&
+ (static_cast<PrimExpr>(loop->loop_var) == IntImm(loop->loop_var.ty(), 0));
}
}
}
diff --git a/src/s_tir/transform/lower_thread_allreduce.cc b/src/s_tir/transform/lower_thread_allreduce.cc
index 4fde26d529..898b9cf573 100644
--- a/src/s_tir/transform/lower_thread_allreduce.cc
+++ b/src/s_tir/transform/lower_thread_allreduce.cc
@@ -203,7 +203,8 @@ class ThreadAllreduceBuilder final : public StmtExprMutator {
std::unordered_set<const VarNode*> reduce_set;
for (size_t i = 2 + 2 * size; i < call->args.size(); ++i) {
- const VarNode* v = call->args[i].as<VarNode>();
+ auto var = call->args[i].as<PrimVar>();
+ const VarNode* v = var.has_value() ? var.value().get() : nullptr;
// The simply optimization replace a iteration variable with a constant
// when extent of the iteration is 1. As threaded IterVar always started from 0,
// we can just ignore this variable in this case.
diff --git a/src/s_tir/transform/memhammer_coalesce.cc b/src/s_tir/transform/memhammer_coalesce.cc
index 06fae25572..57388f4e81 100644
--- a/src/s_tir/transform/memhammer_coalesce.cc
+++ b/src/s_tir/transform/memhammer_coalesce.cc
@@ -185,7 +185,7 @@ ffi::Array<PrimExpr> GetMapping(const Stmt& stmt, const ConstraintSet& constrain
Stmt InverseMapping::Rewrite(const Stmt& stmt, const ConstraintSet& constraints,
OutputSet* output) const {
Stmt body = stmt;
- ffi::Map<Var, Range> var_range;
+ ffi::Map<PrimVar, Range> var_range;
ffi::Array<PrimExpr> loop_vars;
// Step 1. Get index mapping
ffi::Array<PrimExpr> mapping_pattern = GetMapping(stmt, constraints);
diff --git a/src/s_tir/transform/memhammer_lower_auto_copy.cc b/src/s_tir/transform/memhammer_lower_auto_copy.cc
index 5442172a39..95123836a0 100644
--- a/src/s_tir/transform/memhammer_lower_auto_copy.cc
+++ b/src/s_tir/transform/memhammer_lower_auto_copy.cc
@@ -579,8 +579,8 @@ class AutoPadder {
Region region = r->source->region;
ffi::Array<PrimExpr> indices;
for (int i = 0; i < static_cast<int>(region.size()); i++) {
- Var var("region" + std::to_string(i));
- indices.push_back(region[i]->min + var.as_or_throw<PrimExpr>());
+ PrimVar var("region" + std::to_string(i));
+ indices.push_back(region[i]->min + static_cast<PrimExpr>(var));
var_range_.Set(var, Range::FromMinExtent(0, region[i]->extent));
}
ffi::Array<PrimExpr> substitued_indices;
diff --git a/src/s_tir/transform/remove_weight_layout_rewrite_block.cc b/src/s_tir/transform/remove_weight_layout_rewrite_block.cc
index 4f56f93f5c..cfab35ad34 100644
--- a/src/s_tir/transform/remove_weight_layout_rewrite_block.cc
+++ b/src/s_tir/transform/remove_weight_layout_rewrite_block.cc
@@ -96,7 +96,7 @@ class RemoveLayoutRewriteBlock : public StmtMutator {
ffi::Array<PrimVar> load_indices;
for (auto ind : load->indices) {
- TVM_FFI_ICHECK(ind->IsInstance<VarNode>());
+ TVM_FFI_ICHECK(ind.as<PrimVar>());
load_indices.push_back(ind.as_or_throw<PrimVar>());
}
buffer_var_to_index_map_[load->buffer->data.get()] = IndexMap(load_indices, store->indices);
diff --git a/src/s_tir/transform/renew_defs.cc b/src/s_tir/transform/renew_defs.cc
index 91473fa232..392e89b86e 100644
--- a/src/s_tir/transform/renew_defs.cc
+++ b/src/s_tir/transform/renew_defs.cc
@@ -58,9 +58,9 @@ class RenewDefMutator : public StmtExprMutator {
if (it != func->buffer_map.end()) {
const Buffer& buffer = (*it).second;
for (const PrimExpr& e : buffer->shape) {
- if (const auto* v = e.as<VarNode>()) {
- if (generator.remap_.count(ffi::GetRef<Var>(v)) == 0) {
- generator.ReDefineVar(ffi::GetRef<Var>(v));
+ if (auto var = e.as<PrimVar>()) {
+ if (generator.remap_.count(var.value()) == 0) {
+ generator.ReDefineVar(var.value());
}
}
}
@@ -156,7 +156,7 @@ class RenewDefMutator : public StmtExprMutator {
private:
Var ReDefineVar(const Var& var) {
- Var new_var = Var(ffi::make_object<VarNode>(*var.get()));
+ Var new_var(var->name_hint, var->ty, var->span);
this->AddDefRemap(var, new_var);
return new_var;
}
diff --git a/src/s_tir/transform/renormalize_split_pattern.cc b/src/s_tir/transform/renormalize_split_pattern.cc
index 6866cd4a96..8789a58b76 100644
--- a/src/s_tir/transform/renormalize_split_pattern.cc
+++ b/src/s_tir/transform/renormalize_split_pattern.cc
@@ -173,7 +173,8 @@ class SplitPatternReNormalizer : public IRMutatorWithAnalyzer {
Stmt VisitStmt_(const ForNode* op) final {
analyzer_->Bind(op->loop_var, Range::FromMinExtent(op->min, op->extent));
With<ConstraintContext> ctx1(analyzer_, op->loop_var >= op->min);
- With<ConstraintContext> ctx2(analyzer_, op->loop_var < op->min + op->extent);
+ With<ConstraintContext> ctx2(analyzer_,
+ static_cast<PrimExpr>(op->loop_var) < op->min + op->extent);
return IRMutatorWithAnalyzer::VisitStmt_(op);
}
diff --git a/src/s_tir/transform/storage_access.cc b/src/s_tir/transform/storage_access.cc
index 6dffe5531c..ce6bfa62ef 100644
--- a/src/s_tir/transform/storage_access.cc
+++ b/src/s_tir/transform/storage_access.cc
@@ -253,7 +253,7 @@ void StorageAccessVisitor::VisitExpr_(const CallNode* op) {
// args[1] is not a raw Var — e.g. a nested tvm_access_ptr or some
// other PrimExpr. Recurse into sub-exprs so any inner buffer var
// refs still get visited, but don't try to record an access entry
- // here (GetScope(Var(nullptr)) would deref a null pointer).
+ // here (GetScope on a null Var would dereference a null pointer).
StmtExprVisitor::VisitExpr_(op);
return;
}
@@ -267,7 +267,7 @@ void StorageAccessVisitor::VisitExpr_(const CallNode* op) {
AccessEntry e;
e.threads = env_threads();
e.dtype = dtype;
- e.buffer = op->args[1].as_or_throw<Var>();
+ e.buffer = ffi::GetRef<Var>(buffer);
e.touched = {arith::IntSet::FromRange(Range::FromMinExtent(offset, extent))};
e.scope = scope;
if (flag->value & 1) {
diff --git a/src/s_tir/transform/storage_access.h b/src/s_tir/transform/storage_access.h
index 4b5d69bb35..8442947930 100644
--- a/src/s_tir/transform/storage_access.h
+++ b/src/s_tir/transform/storage_access.h
@@ -59,7 +59,7 @@ class StorageAccessVisitor : public StmtExprVisitor {
/*! \brief The thread index that access this entry */
ffi::Array<IterVar> threads;
/*! \brief The buffer variable, if any */
- Var buffer = Var(ffi::ObjectPtr<VarNode>(nullptr));
+ Var buffer{ffi::UnsafeInit{}};
/*! \brief The access data type */
PrimType dtype = PrimType::Void();
/*! \brief The touched access range
diff --git a/src/s_tir/transform/using_assume_to_reduce_branches.cc b/src/s_tir/transform/using_assume_to_reduce_branches.cc
index 4ffe2f460e..2a1e6b443c 100644
--- a/src/s_tir/transform/using_assume_to_reduce_branches.cc
+++ b/src/s_tir/transform/using_assume_to_reduce_branches.cc
@@ -190,7 +190,8 @@ class ParseAssumeAndOvercompute : public IRMutatorWithAnalyzer {
InternalConstraintContext */
analyzer_->Bind(op->loop_var, Range::FromMinExtent(op->min, op->extent));
InternalConstraintContext ctx1(this, op->loop_var >= op->min);
- InternalConstraintContext ctx2(this, op->loop_var < op->min + op->extent);
+ InternalConstraintContext ctx2(this,
+ static_cast<PrimExpr>(op->loop_var) < op->min + op->extent);
return Parent::VisitStmt_(op);
}
diff --git a/src/script/ir_builder/ir/ir.cc b/src/script/ir_builder/ir/ir.cc
index 32edb847ae..d1840e2fef 100644
--- a/src/script/ir_builder/ir/ir.cc
+++ b/src/script/ir_builder/ir/ir.cc
@@ -29,6 +29,14 @@ namespace script {
namespace ir_builder {
namespace ir {
+using tvm::script::ir_builder::details::Namer;
+
+TVM_STATIC_IR_FUNCTOR(Namer, vtable)
+ .set_dispatch<tvm::VarNode>([](const ffi::ObjectRef& node, ffi::String name) -> void {
+ VarNode* var = const_cast<VarNode*>(node.as<VarNode>());
+ var->name_hint = name;
+ });
+
IRModuleFrame IRModule() {
ffi::ObjectPtr<IRModuleFrameNode> n = ffi::make_object<IRModuleFrameNode>();
n->global_var_map.clear();
diff --git a/src/te/operation/create_primfunc.cc b/src/te/operation/create_primfunc.cc
index eb3eae86a7..d4c1de9ec0 100644
--- a/src/te/operation/create_primfunc.cc
+++ b/src/te/operation/create_primfunc.cc
@@ -644,8 +644,8 @@ Stmt GenerateStmtFromExternOp(const te::ExternOp& extern_op, CreateFuncInfo* inf
Buffer output_buffer = placeholder;
if (!info->IsArg(output_tensor)) {
PrimExpr zero_offset = IntImm(placeholder->elem_offset.ty(), 0);
- if (const auto* offset_var = placeholder->elem_offset.as<VarNode>()) {
- var_map[offset_var] = zero_offset;
+ if (auto offset_var = placeholder->elem_offset.as<PrimVar>()) {
+ var_map[offset_var.value().get()] = zero_offset;
}
output_buffer.CopyOnWrite()->elem_offset = zero_offset;
input_buffer_map[placeholder.get()] = output_buffer;
@@ -823,7 +823,7 @@ PrimFunc GenerateAndCompletePrimFunc(const ffi::Array<ffi::ObjectRef>& arg_tir_v
Var param("var_" + tensor->GetNameHint(), PointerType::VoidPointerTy());
parameters.push_back(param);
buffer_map.Set(param, it->second);
- } else if (auto var = arg.as<tirx::Var>()) {
+ } else if (auto var = arg.as<tirx::PrimVar>()) {
parameters.push_back(var.value());
}
}
diff --git a/src/te/tensor.cc b/src/te/tensor.cc
index 57e48e5e7a..8eeb08b89f 100644
--- a/src/te/tensor.cc
+++ b/src/te/tensor.cc
@@ -73,7 +73,8 @@ inline PrimExpr Tensor::IndexTensor(ffi::Array<PrimExpr> indices,
}
PrimExpr Tensor::operator()(ffi::Array<PrimVar> indices) const {
- ffi::Array<PrimExpr> arr(indices.begin(), indices.end());
+ ffi::Array<PrimExpr> arr =
+ indices.Map([](const PrimVar& var) { return static_cast<PrimExpr>(var); });
return operator()(arr);
}
@@ -82,7 +83,8 @@ PrimExpr Tensor::operator()(ffi::Array<PrimExpr> indices) const {
}
PrimExpr Tensor::IndexWithNegativeIndices(ffi::Array<PrimVar> indices) const {
- ffi::Array<PrimExpr> arr(indices.begin(), indices.end());
+ ffi::Array<PrimExpr> arr =
+ indices.Map([](const PrimVar& var) { return static_cast<PrimExpr>(var); });
return IndexWithNegativeIndices(arr);
}
diff --git a/src/tirx/analysis/verify_well_formed.cc b/src/tirx/analysis/verify_well_formed.cc
index e8eb5671cf..8603d01eb8 100644
--- a/src/tirx/analysis/verify_well_formed.cc
+++ b/src/tirx/analysis/verify_well_formed.cc
@@ -184,7 +184,7 @@ class UndefinedVarVerifier : public Verifier<UndefinedVarVerifier> {
auto verify = Verify(it == currently_defined_.end() || redefine_is_allowed);
verify << "ValueError: "
<< "TIR is ill-formed, "
- << "due to multiple nested definitions of variable " << var << ".";
+ << "due to multiple nested definitions of variable " << var->name_hint << ".";
if (it != currently_defined_.end()) {
verify << " It was first defined at " << it->second << ", and was re-defined at " << path;
}
@@ -195,7 +195,7 @@ class UndefinedVarVerifier : public Verifier<UndefinedVarVerifier> {
auto verify = Verify(it == previously_defined_.end() || redefine_is_allowed);
verify << "ValueError: "
<< "TIR is ill-formed, "
- << "due to multiple definitions of variable " << var << ".";
+ << "due to multiple definitions of variable " << var->name_hint << ".";
if (it != previously_defined_.end()) {
verify << " It was first defined at " << it->second << ", and was later re-defined at "
<< path;
@@ -218,7 +218,7 @@ class UndefinedVarVerifier : public Verifier<UndefinedVarVerifier> {
auto active_def = currently_defined_.find(var);
auto verify = Verify(active_def != currently_defined_.end());
verify << "ValueError: "
- << "Invalid use of undefined variable " << var << " at " << path << ".";
+ << "Invalid use of undefined variable " << var->name_hint << " at " << path << ".";
// Check if there was a previous definition, and append the
// location to the error message if there was. This is to aid in
@@ -339,9 +339,9 @@ class SingleEnvThreadVerifier : public Verifier<SingleEnvThreadVerifier> {
<< "While multiple tirx::AttrStmt may define the same environment thread, "
<< "all definitions within a single PrimFunc must share the same tirx::Var. "
<< "Binding of environment thread \"" << iter_var->thread_tag
- << "\" to the TIR variable " << iter_var->var << " at " << path
- << " conflicts with the previous binding to the TIR variable " << prev_var << " at "
- << path;
+ << "\" to the TIR variable " << iter_var->var->name_hint << " at " << path
+ << " conflicts with the previous binding to the TIR variable " << prev_var->name_hint
+ << " at " << path;
} else {
env_thread_vars_.insert({iter_var->thread_tag, {iter_var->var, path}});
}
diff --git a/src/tirx/ir/buffer.cc b/src/tirx/ir/buffer.cc
index 368d697474..6374c53aac 100644
--- a/src/tirx/ir/buffer.cc
+++ b/src/tirx/ir/buffer.cc
@@ -656,7 +656,7 @@ tirx::Buffer BufferWithOffsetAlignment(ffi::Array<PrimExpr> shape, PrimType dtyp
bool has_any = false;
if (!compact) {
for (const auto& it : shape) {
- if (it.as<tirx::VarNode>()) {
+ if (it.as<tirx::PrimVar>()) {
has_any = true;
break;
}
diff --git a/src/tirx/ir/expr.cc b/src/tirx/ir/expr.cc
index d3d77844a1..19516d5ff0 100644
--- a/src/tirx/ir/expr.cc
+++ b/src/tirx/ir/expr.cc
@@ -72,7 +72,6 @@ TVM_FFI_INLINE const PrimTypeNode* GetPrimTypeNode(const PrimExpr& expr) {
} // namespace
TVM_FFI_STATIC_INIT_BLOCK() {
- VarNode::RegisterReflection();
IterVarNode::RegisterReflection();
StringImmNode::RegisterReflection();
CastNode::RegisterReflection();
@@ -155,52 +154,6 @@ TVM_FFI_STATIC_INIT_BLOCK() {
data_ = std::move(node); \
}
-// Var
-Var::Var(ffi::String name_hint, PrimType dtype, Span span) {
- auto n = ffi::make_object<VarNode>();
- n->name_hint = std::move(name_hint);
- n->ExprNode::ty = dtype;
- n->span = std::move(span);
- data_ = std::move(n);
-}
-
-Var::Var(ffi::String name_hint, Type type_annotation, Span span) {
- auto n = ffi::make_object<VarNode>();
- n->name_hint = std::move(name_hint);
- n->ExprNode::ty = std::move(type_annotation);
- n->span = std::move(span);
- data_ = std::move(n);
-}
-
-Var Var::copy_with_name(const ffi::String& name) const {
- const VarNode* node = get();
- ffi::ObjectPtr<VarNode> new_ptr = ffi::make_object<VarNode>(*node);
- new_ptr->name_hint = name;
- return Var(new_ptr);
-}
-
-Var Var::CopyWithSuffix(const ffi::String& suffix) const {
- return this->copy_with_name(get()->name_hint + suffix);
-}
-
-Var Var::copy_with_dtype(PrimType dtype) const {
- const VarNode* node = get();
- ffi::ObjectPtr<VarNode> new_ptr = ffi::make_object<VarNode>(*node);
- new_ptr->ExprNode::ty = dtype;
- return Var(new_ptr);
-}
-
-TVM_FFI_STATIC_INIT_BLOCK() {
- namespace refl = tvm::ffi::reflection;
- refl::GlobalDef().def("tirx.Var", [](ffi::String name_hint, ffi::AnyView type, Span span) {
- if (type.as<Type>()) {
- return Var(name_hint, type.cast<Type>(), span);
- } else {
- return Var(name_hint, type.cast<PrimType>(), span);
- }
- });
-}
-
// IterVar
IterVar::IterVar(Range dom, PrimVar var, IterVarType t, ffi::String thread_tag, Span span) {
ffi::ObjectPtr<IterVarNode> n = ffi::make_object<IterVarNode>();
diff --git a/src/tirx/ir/index_map.cc b/src/tirx/ir/index_map.cc
index 02e156c50d..5ddd6defa9 100644
--- a/src/tirx/ir/index_map.cc
+++ b/src/tirx/ir/index_map.cc
@@ -86,8 +86,8 @@ std::pair<IndexMap, PrimExpr> IndexMapInverseImpl(const IndexMap& self,
// TODO(Lunderberg): Name a pair of output indices split from a single
// input index as (X.outer, X.inner).
std::string name;
- if (auto* var = index.as<VarNode>()) {
- name = var->name_hint;
+ if (auto var = index.as<PrimVar>()) {
+ name = var.value()->name_hint;
} else {
name = "axis" + std::to_string(i);
}
@@ -96,7 +96,7 @@ std::pair<IndexMap, PrimExpr> IndexMapInverseImpl(const IndexMap& self,
}
// Dummy ranges for the extent of each input.
- ffi::Map<Var, Range> input_iters;
+ ffi::Map<PrimVar, Range> input_iters;
TVM_FFI_ICHECK_EQ(self->initial_indices.size(), initial_ranges.size());
for (size_t i = 0; i < initial_ranges.size(); i++) {
input_iters.Set(self->initial_indices[i], initial_ranges[i]);
@@ -213,7 +213,7 @@ ffi::Array<Range> IndexMapNode::MapRanges(const ffi::Array<Range>& ranges,
arith::AnalyzerObj* analyzer_ptr = analyzer.get();
TVM_FFI_ICHECK_EQ(ranges.size(), initial_indices.size());
- ffi::Map<Var, Range> input_iters;
+ ffi::Map<PrimVar, Range> input_iters;
for (size_t i = 0; i < initial_indices.size(); i++) {
input_iters.Set(initial_indices[i], ranges[i]);
}
@@ -364,19 +364,19 @@ IndexMap IndexMap::RenameVariables(
std::unordered_set<const ffi::Object*> visited;
std::for_each(n->final_indices.begin(), n->final_indices.end(), [&](const PrimExpr& expr) {
PostOrderVisit(expr, [&](const ffi::ObjectRef& obj) {
- if (!obj->IsInstance<VarNode>()) {
+ auto var = obj.as<PrimVar>();
+ if (!var) {
... 3372 lines suppressed ...