You are viewing a plain text version of this content. The canonical link for it is here.
Posted to commits@tvm.apache.org by tq...@apache.org on 2026/09/16 23:05:19 UTC
(tvm) branch main updated: [REFACTOR][IR] Share primitive helpers and Python expressions (#20364)
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 4d478c41ff [REFACTOR][IR] Share primitive helpers and Python expressions (#20364)
4d478c41ff is described below
commit 4d478c41ffac2ba62e0cd4fe2d58ac5ba2b30efb
Author: Tianqi Chen <tq...@users.noreply.github.com>
AuthorDate: Wed Sep 16 19:05:13 2026 -0400
[REFACTOR][IR] Share primitive helpers and Python expressions (#20364)
Shared primitive expressions need construction and constant helpers that
arithmetic can use without depending on TIRX. Add a small shared
operation header and implementation for construction, constant folding,
integer helpers, and paired value limits; retain basic C++ expression
operations and overloaded operators in the root `tvm` namespace.
Use canonical `prim.*` registrations for shared nodes and selected
intrinsics, including `ceil` and `log2`, while keeping richer intrinsic
identities in TIRX. Preserve existing operand conversion, unsigned
constants, vector behavior, and lowering semantics.
Move canonical Python primitive classes, operators, and construction
helpers into `tvm.ir.prim`. Route arithmetic and generic conversion
through those definitions and retain TIRX compatibility exports without
duplicate registrations or shared-module back-imports. JSON compact
migration is unchanged.
---
include/tvm/ir/prim/builtin.h | 6 +
include/tvm/ir/prim/expr.h | 44 +-
include/tvm/{tirx => ir/prim}/op.h | 543 ++--------
include/tvm/ir/prim/vector_expr.h | 6 +-
include/tvm/s_tir/analysis.h | 3 -
include/tvm/s_tir/backend/adreno/transform.h | 1 -
include/tvm/s_tir/meta_schedule/arg_info.h | 1 -
include/tvm/s_tir/meta_schedule/builder.h | 1 -
include/tvm/s_tir/meta_schedule/cost_model.h | 1 -
include/tvm/s_tir/meta_schedule/database.h | 1 -
include/tvm/s_tir/meta_schedule/extracted_task.h | 1 -
.../tvm/s_tir/meta_schedule/feature_extractor.h | 1 -
include/tvm/s_tir/meta_schedule/measure_callback.h | 1 -
.../tvm/s_tir/meta_schedule/measure_candidate.h | 1 -
include/tvm/s_tir/meta_schedule/mutator.h | 1 -
include/tvm/s_tir/meta_schedule/postproc.h | 1 -
include/tvm/s_tir/meta_schedule/profiler.h | 1 -
include/tvm/s_tir/meta_schedule/runner.h | 1 -
.../meta_schedule/schedule/cuda/thread_bind.h | 1 -
.../meta_schedule/schedule/generic/winograd.h | 1 -
include/tvm/s_tir/meta_schedule/schedule_rule.h | 1 -
include/tvm/s_tir/meta_schedule/search_strategy.h | 1 -
include/tvm/s_tir/meta_schedule/space_generator.h | 1 -
include/tvm/s_tir/meta_schedule/task_scheduler.h | 1 -
include/tvm/s_tir/meta_schedule/tune_context.h | 1 -
include/tvm/s_tir/random_engine.h | 2 -
include/tvm/s_tir/schedule/instruction.h | 1 -
include/tvm/s_tir/schedule/schedule.h | 1 -
include/tvm/s_tir/schedule/state.h | 1 -
include/tvm/s_tir/schedule/trace.h | 1 -
include/tvm/s_tir/stmt.h | 1 -
include/tvm/s_tir/transform.h | 2 -
include/tvm/te/tensor.h | 1 -
include/tvm/tirx/op.h | 764 +-------------
include/tvm/tirx/script/builder/ir.h | 2 +-
include/tvm/topi/broadcast.h | 15 +-
include/tvm/topi/detail/broadcast.h | 2 +-
include/tvm/topi/detail/extern.h | 1 +
include/tvm/topi/detail/strided_slice.h | 3 +-
include/tvm/topi/detail/tensor_utils.h | 8 +-
include/tvm/topi/elemwise.h | 97 +-
include/tvm/topi/nn.h | 52 +-
include/tvm/topi/nn/bnn.h | 7 +-
include/tvm/topi/nn/dense.h | 5 +-
include/tvm/topi/nn/dilate.h | 1 +
include/tvm/topi/nn/group_norm.h | 3 +-
include/tvm/topi/nn/instance_norm.h | 4 +-
include/tvm/topi/nn/layer_norm.h | 7 +-
include/tvm/topi/nn/local_response_norm.h | 12 +-
include/tvm/topi/nn/pooling.h | 20 +-
include/tvm/topi/nn/rms_norm.h | 5 +-
include/tvm/topi/nn/softmax.h | 12 +-
include/tvm/topi/reduction.h | 29 +-
include/tvm/topi/transform.h | 53 +-
python/tvm/arith/analyzer.py | 42 +-
python/tvm/arith/bound.py | 6 +-
python/tvm/arith/int_set.py | 10 +-
python/tvm/arith/iter_affine_map.py | 12 +-
python/tvm/arith/pattern.py | 4 +-
python/tvm/ir/__init__.py | 1 +
python/tvm/ir/_overload_prim_expr.py | 416 ++++++--
python/tvm/ir/expr.py | 11 +-
python/tvm/ir/prim/__init__.py | 83 +-
python/tvm/ir/prim/_ffi_api.py | 2 +-
python/tvm/ir/prim/expr.py | 699 +++++++++++++
python/tvm/{arith/pattern.py => ir/prim/op.py} | 49 +-
python/tvm/tirx/expr.py | 1079 +-------------------
python/tvm/tirx/op.py | 90 +-
python/tvm/tirx/script/builder/ir.py | 6 +-
python/tvm/topi/_te_tensor_overload.py | 5 +-
python/tvm/topi/math.py | 2 +-
src/arith/analyzer.cc | 14 +-
src/arith/canonical_simplify.cc | 20 +-
src/arith/conjunctive_normal_form.cc | 2 +-
src/arith/const_fold.h | 490 +--------
src/arith/const_int_bound.cc | 4 +-
src/arith/detect_linear_equation.cc | 7 +-
src/arith/int_operator.h | 48 +-
src/arith/int_set.cc | 27 +-
src/arith/interval_set.h | 2 +-
src/arith/iter_affine_map.cc | 7 +-
src/arith/modular_set.cc | 1 +
src/arith/pattern_match.h | 3 +-
src/arith/presburger_set.cc | 1 +
src/arith/presburger_set.h | 2 +-
src/arith/product_normal_form.h | 2 +-
src/arith/rewrite_simplify.cc | 44 +-
src/arith/rewrite_simplify.h | 2 +-
src/arith/simplify_base.cc | 10 +-
src/arith/transitive_comparison_analyzer.cc | 3 +
src/arith/z3_prover.cc | 7 +-
src/backend/cuda/codegen/codegen_cuda.cc | 9 +-
src/backend/cuda/codegen/intrin_rule_cuda.cc | 4 +-
src/backend/cuda/codegen/llvm/intrin_rule_nvptx.cc | 7 +-
src/backend/cuda/transforms/lower_iket.cc | 11 +-
.../hexagon/codegen/llvm/intrin_rule_hexagon.cc | 21 +-
src/backend/metal/codegen/codegen_metal.cc | 9 +-
src/backend/metal/codegen/intrin_rule_metal.cc | 4 +-
src/backend/opencl/codegen/codegen_opencl.cc | 23 +-
src/backend/opencl/codegen/intrin_rule_opencl.cc | 4 +-
src/backend/rocm/codegen/llvm/intrin_rule_rocm.cc | 22 +-
src/backend/trn/codegen/codegen_trn.cc | 4 +-
src/backend/trn/transform/lower_trainium_layout.cc | 4 +-
src/backend/vulkan/codegen/codegen_spirv.cc | 29 +-
src/backend/vulkan/codegen/intrin_rule_spirv.cc | 4 +-
src/backend/webgpu/codegen/codegen_webgpu.cc | 9 +-
src/backend/webgpu/codegen/intrin_rule_webgpu.cc | 4 +-
src/ir/expr.cc | 4 +-
src/ir/prim/builtin.cc | 42 +-
src/{arith => ir/prim}/const_fold.h | 88 +-
src/ir/prim/deep_equal.cc | 2 +-
src/ir/prim/expr.cc | 51 +-
src/ir/prim/int_operator.h | 77 ++
src/ir/prim/op.cc | 773 ++++++++++++++
src/ir/prim/op_utils.h | 99 ++
.../winograd.h => src/ir/prim/symbolic_limits.cc | 28 +-
src/ir/prim/vector_expr.cc | 13 +-
src/relax/analysis/tir_op_pattern_kind.cc | 1 +
src/relax/analysis/type_analysis.cc | 3 +-
src/relax/analysis/well_formed.cc | 1 +
.../backend/adreno/annotate_custom_storage.cc | 2 +
.../backend/contrib/codegen_json/codegen_json.h | 2 +
src/relax/backend/contrib/tensorrt/codegen.cc | 2 +
src/relax/backend/contrib/utils.h | 3 +-
src/relax/backend/vm/lower_runtime_builtin.cc | 1 +
src/relax/backend/vm/vm_shape_lower.cc | 1 +
src/relax/distributed/axis_group_graph.cc | 2 +
src/relax/distributed/transform/lower_distir.cc | 2 +
.../transform/lower_global_view_to_local_view.cc | 2 +
src/relax/ir/block_builder.cc | 1 +
src/relax/ir/dataflow_block_rewriter.cc | 2 +-
src/relax/ir/dataflow_expr_rewriter.cc | 1 +
src/relax/ir/dataflow_matcher.cc | 3 +-
src/relax/ir/dependent_type.cc | 3 +-
src/relax/ir/expr.cc | 2 +-
src/relax/ir/expr_functor.cc | 1 +
src/relax/op/ccl/ccl.cc | 1 +
src/relax/op/distributed/distributed.cc | 1 +
src/relax/op/memory/view.cc | 1 +
src/relax/op/nn/attention.cc | 2 +-
src/relax/op/nn/convolution.cc | 1 +
src/relax/op/nn/pooling.cc | 1 +
src/relax/op/op.cc | 7 +-
src/relax/op/op_common.cc | 3 +-
src/relax/op/op_common.h | 1 +
src/relax/op/tensor/binary.cc | 1 +
src/relax/op/tensor/create.cc | 5 +-
src/relax/op/tensor/index.cc | 1 +
src/relax/op/tensor/inspect.cc | 4 +-
src/relax/op/tensor/manipulate.cc | 2 +-
src/relax/op/tensor/sampling.cc | 4 +-
src/relax/script/builder/ir.cc | 2 +
src/relax/transform/adjust_matmul_order.cc | 1 +
src/relax/transform/alter_op_impl.cc | 2 +-
src/relax/transform/canonicalize_bindings.cc | 1 +
src/relax/transform/combine_parallel_matmul.cc | 1 +
src/relax/transform/dataflow_inplace.cc | 1 +
src/relax/transform/decompose_ops.cc | 1 +
src/relax/transform/fold_constant.cc | 1 +
src/relax/transform/fuse_ops.cc | 1 +
src/relax/transform/fuse_tir.cc | 1 +
src/relax/transform/lazy_transform_params.cc | 1 +
src/relax/transform/lift_transform_params.cc | 1 +
src/relax/transform/lower_alloc_tensor.cc | 1 +
src/relax/transform/merge_composite_functions.cc | 1 +
src/relax/transform/remove_unused_outputs.cc | 1 +
src/relax/transform/remove_unused_parameters.cc | 1 +
.../specialize_primfunc_based_on_callsite.cc | 1 +
src/relax/transform/split_call_tir_by_pattern.cc | 1 +
.../transform/split_layout_rewrite_preproc.cc | 2 +
src/relax/transform/static_plan_block_memory.cc | 7 +-
src/relax/transform/to_non_dataflow.cc | 1 +
src/relax/utils.cc | 1 +
src/s_tir/analysis/calculate_allocated_memory.cc | 1 -
src/s_tir/analysis/conditional_bounds.cc | 13 +-
src/s_tir/analysis/estimate_flops.cc | 43 +-
src/s_tir/analysis/identify_memcpy.cc | 3 +-
src/s_tir/analysis/is_pure_function.cc | 1 -
src/s_tir/analysis/oob_checker.cc | 1 -
.../analysis/sblock_access_region_detector.cc | 2 +-
src/s_tir/backend/adreno/inject_texture_alloc.cc | 3 +-
src/s_tir/backend/adreno/texture_flatten.cc | 6 +-
src/s_tir/data_layout.cc | 3 +-
src/s_tir/meta_schedule/arg_info.cc | 1 -
src/s_tir/meta_schedule/builder/builder.cc | 1 -
src/s_tir/meta_schedule/cost_model/cost_model.cc | 1 -
src/s_tir/meta_schedule/database/database.cc | 1 -
src/s_tir/meta_schedule/database/database_utils.cc | 1 -
src/s_tir/meta_schedule/database/json_database.cc | 1 -
.../meta_schedule/database/memory_database.cc | 1 -
.../database/ordered_union_database.cc | 1 -
.../meta_schedule/database/schedule_fn_database.cc | 1 -
src/s_tir/meta_schedule/database/union_database.cc | 1 -
src/s_tir/meta_schedule/extracted_task.cc | 1 -
.../feature_extractor/feature_extractor.cc | 1 -
.../feature_extractor/per_store_feature.cc | 3 +-
.../measure_callback/add_to_database.cc | 1 -
.../measure_callback/measure_callback.cc | 1 -
.../measure_callback/remove_build_artifact.cc | 1 -
.../measure_callback/update_cost_model.cc | 1 -
src/s_tir/meta_schedule/module_equality.cc | 1 -
src/s_tir/meta_schedule/module_equality.h | 1 -
.../mutator/mutate_compute_location.cc | 1 -
src/s_tir/meta_schedule/mutator/mutate_parallel.cc | 2 -
.../meta_schedule/mutator/mutate_thread_binding.cc | 1 -
.../meta_schedule/mutator/mutate_tile_size.cc | 1 -
src/s_tir/meta_schedule/mutator/mutate_unroll.cc | 2 -
src/s_tir/meta_schedule/mutator/mutator.cc | 1 -
.../postproc/disallow_async_strided_mem_copy.cc | 2 -
.../postproc/disallow_dynamic_loop.cc | 2 -
src/s_tir/meta_schedule/postproc/postproc.cc | 1 -
.../postproc/rewrite_cooperative_fetch.cc | 2 -
src/s_tir/meta_schedule/postproc/rewrite_layout.cc | 2 -
.../postproc/rewrite_parallel_vectorize_unroll.cc | 2 -
.../postproc/rewrite_reduction_block.cc | 2 -
.../meta_schedule/postproc/rewrite_tensorize.cc | 1 -
.../postproc/rewrite_unbound_block.cc | 2 -
.../meta_schedule/postproc/verify_gpu_code.cc | 2 -
.../meta_schedule/postproc/verify_vtcm_limit.cc | 1 -
src/s_tir/meta_schedule/profiler.cc | 1 -
src/s_tir/meta_schedule/runner/runner.cc | 1 -
src/s_tir/meta_schedule/schedule/cpu/winograd.cc | 1 -
src/s_tir/meta_schedule/schedule/cuda/winograd.cc | 1 -
.../meta_schedule/schedule/generic/winograd.cc | 1 -
.../meta_schedule/schedule_rule/add_rfactor.cc | 1 -
.../schedule_rule/apply_custom_rule.cc | 1 -
src/s_tir/meta_schedule/schedule_rule/auto_bind.cc | 1 -
.../meta_schedule/schedule_rule/auto_inline.cc | 1 -
.../schedule_rule/cross_thread_reduction.cc | 1 -
.../schedule_rule/multi_level_tiling.cc | 2 -
.../schedule_rule/multi_level_tiling.h | 2 -
.../multi_level_tiling_tensor_core.cc | 3 +-
.../multi_level_tiling_wide_vector.cc | 1 -
.../multi_level_tiling_with_intrin.cc | 1 -
.../schedule_rule/parallel_vectorize_unroll.cc | 2 -
.../schedule_rule/random_compute_location.cc | 1 -
.../meta_schedule/schedule_rule/schedule_rule.cc | 1 -
.../search_strategy/evolutionary_search.cc | 1 -
.../meta_schedule/search_strategy/replay_func.cc | 1 -
.../meta_schedule/search_strategy/replay_trace.cc | 1 -
.../search_strategy/search_strategy.cc | 1 -
.../space_generator/post_order_apply.cc | 1 -
.../meta_schedule/space_generator/schedule_fn.cc | 1 -
.../space_generator/space_generator.cc | 1 -
.../space_generator/space_generator_union.cc | 1 -
.../meta_schedule/task_scheduler/gradient_based.cc | 1 -
.../meta_schedule/task_scheduler/round_robin.cc | 1 -
.../meta_schedule/task_scheduler/task_scheduler.cc | 1 -
src/s_tir/meta_schedule/trace_apply.cc | 1 -
src/s_tir/meta_schedule/trace_apply.h | 1 -
src/s_tir/meta_schedule/tune_context.cc | 1 -
src/s_tir/meta_schedule/utils.h | 1 -
src/s_tir/schedule/analysis.h | 1 -
src/s_tir/schedule/analysis/analysis.cc | 2 +-
src/s_tir/schedule/analysis/reducer.cc | 1 -
src/s_tir/schedule/analysis/verify.cc | 1 -
src/s_tir/schedule/concrete_schedule.h | 1 -
src/s_tir/schedule/error.cc | 1 -
src/s_tir/schedule/error.h | 1 -
src/s_tir/schedule/instruction.cc | 1 -
src/s_tir/schedule/instruction_traits.h | 1 -
src/s_tir/schedule/ir_comparator.h | 1 -
src/s_tir/schedule/primitive.h | 1 -
src/s_tir/schedule/primitive/annotate.cc | 1 -
.../schedule/primitive/annotate_buffer_access.cc | 1 -
src/s_tir/schedule/primitive/block_annotate.cc | 1 -
src/s_tir/schedule/primitive/cache_index.cc | 1 -
src/s_tir/schedule/primitive/compute_inline.cc | 2 +-
src/s_tir/schedule/primitive/for_kind.cc | 1 -
src/s_tir/schedule/primitive/get_block_loop.cc | 1 -
src/s_tir/schedule/primitive/hide_buffer_access.cc | 1 -
.../schedule/primitive/layout_transformation.cc | 8 +-
.../schedule/primitive/loop_transformation.cc | 7 +-
src/s_tir/schedule/primitive/pad_einsum.cc | 7 +-
src/s_tir/schedule/primitive/read_write_at.cc | 1 -
src/s_tir/schedule/primitive/reduction.cc | 22 +-
.../schedule/primitive/reorder_block_iter_var.cc | 1 -
src/s_tir/schedule/primitive/sampling.cc | 1 -
src/s_tir/schedule/schedule.cc | 1 -
src/s_tir/schedule/state.cc | 1 -
src/s_tir/schedule/trace.cc | 1 -
src/s_tir/schedule/traced_schedule.cc | 1 -
src/s_tir/schedule/traced_schedule.h | 1 -
src/s_tir/schedule/transform.h | 1 -
src/s_tir/transform/annotate_irregular_loop.cc | 1 -
src/s_tir/transform/compact_buffer_region.cc | 1 -
src/s_tir/transform/convert_blocks_to_opaque.cc | 1 -
src/s_tir/transform/decorate_device_scope.cc | 1 -
src/s_tir/transform/default_gpu_schedule.cc | 1 -
src/s_tir/transform/hoist_expression.cc | 9 +-
src/s_tir/transform/inject_permuted_layout.cc | 1 -
src/s_tir/transform/inject_ptx_async_copy.cc | 41 +-
src/s_tir/transform/inject_ptx_ldg32.cc | 5 +-
src/s_tir/transform/inject_virtual_thread.cc | 4 +-
src/s_tir/transform/lift_thread_binding.cc | 1 -
src/s_tir/transform/lower_async_dma.cc | 1 -
src/s_tir/transform/lower_init_block.cc | 1 -
src/s_tir/transform/lower_match_buffer.cc | 2 +-
src/s_tir/transform/lower_opaque_block.cc | 2 +-
src/s_tir/transform/lower_thread_allreduce.cc | 3 +-
src/s_tir/transform/lower_vtcm_alloc.cc | 3 +-
.../manifest_shared_memory_local_stage.cc | 1 -
.../transform/memhammer_intermediate_stage.cc | 1 -
src/s_tir/transform/memhammer_lower_auto_copy.cc | 1 -
.../transform/memhammer_tensorcore_rewrite.cc | 1 -
.../transform/merge_shared_memory_allocations.cc | 1 -
.../plan_update_buffer_allocation_location.cc | 1 -
src/s_tir/transform/profile_instrumentation.cc | 1 -
src/s_tir/transform/remove_store_undef.cc | 1 -
src/s_tir/transform/renew_defs.cc | 1 -
src/s_tir/transform/renormalize_split_pattern.cc | 31 +-
src/s_tir/transform/rewrite_unsafe_select.cc | 59 +-
src/s_tir/transform/storage_access.cc | 3 +-
src/s_tir/transform/storage_access.h | 1 -
src/s_tir/transform/tensorcore_infer_fragment.cc | 8 +-
src/s_tir/transform/thread_storage_sync.cc | 5 +-
src/s_tir/transform/transform_mma_buffer_layout.cc | 1 -
.../transform/using_assume_to_reduce_branches.cc | 1 -
src/target/intrin_rule.cc | 10 +-
src/target/intrin_rule.h | 3 +-
src/target/llvm/codegen_arm.cc | 1 +
src/target/llvm/codegen_cpu.cc | 57 +-
src/target/llvm/codegen_llvm.cc | 47 +-
src/target/llvm/intrin_rule_llvm.cc | 10 +-
src/target/source/codegen_c.cc | 79 +-
src/te/operation/create_primfunc.cc | 1 +
src/te/operation/scan_op.cc | 2 +
src/tirx/ir/buffer.cc | 7 +-
src/tirx/ir/data_type_rewriter.cc | 49 +-
src/tirx/ir/exec_scope.cc | 1 +
src/tirx/ir/function.cc | 1 +
src/tirx/ir/index_map.cc | 5 +-
src/tirx/ir/layout/compose_layout.cc | 1 +
src/tirx/ir/layout/tile_canonicalize.cc | 1 +
src/tirx/ir/layout/tile_core.cc | 1 +
src/tirx/ir/layout/tile_direct_sum_ops.cc | 1 +
src/tirx/ir/layout/tile_tile_ops.cc | 1 +
src/tirx/ir/stmt.cc | 1 +
src/tirx/ir_mutator_with_analyzer.cc | 5 +-
src/tirx/ir_visitor_with_analyzer.cc | 2 +-
src/tirx/op/builtin.cc | 22 +
src/tirx/op/op.cc | 869 +---------------
src/tirx/script/builder/frame.cc | 3 +-
src/tirx/script/builder/ir.cc | 4 +-
src/tirx/script/printer/block.cc | 5 +-
src/tirx/script/printer/buffer.cc | 7 +-
src/tirx/script/printer/for_loop.cc | 5 +-
src/tirx/script/printer/function.cc | 3 +-
src/tirx/script/printer/utils.h | 3 +-
src/tirx/transform/dtype_conversion.cc | 1 +
src/tirx/transform/flatten_buffer.cc | 1 +
src/tirx/transform/force_narrow_index_to_i32.cc | 1 +
src/tirx/transform/ir_utils.cc | 1 +
src/tirx/transform/lower_intrin.cc | 9 +-
src/tirx/transform/lower_tirx_dedup_tensormap.cc | 1 +
src/tirx/transform/lower_tirx_opaque.cc | 3 +-
src/tirx/transform/lower_tvm_builtin.cc | 17 +-
src/tirx/transform/lower_warp_memory.cc | 1 +
src/tirx/transform/make_packed_api.cc | 2 +-
src/tirx/transform/narrow_datatype.cc | 1 +
src/tirx/transform/remove_no_op.cc | 1 +
src/tirx/transform/stmt_simplify.cc | 1 +
src/tirx/transform/storage_rewrite.cc | 1 +
src/tirx/transform/tvm_ffi_binder.cc | 1 +
src/tirx/transform/unsupported_dtype_legalize.cc | 5 +-
src/tirx/transform/vectorize_loop.cc | 1 +
src/topi/einsum.cc | 1 +
src/topi/elemwise.cc | 1 +
tests/cpp/arith_simplify_test.cc | 18 +-
tests/cpp/ir_functor_test.cc | 5 +-
tests/cpp/pattern_match_test.cc | 8 +-
tests/cpp/te_compute_test.cc | 1 +
tests/cpp/tir_analysis_side_effect.cc | 1 +
tests/python/arith/test_arith_rewrite_simplify.py | 2 +-
tests/python/ir/test_node_reflection.py | 2 +-
.../test_s_tir_transform_rewrite_unsafe_select.py | 2 +-
tests/python/tirx-base/test_tir_constructor.py | 4 +-
tests/python/tirx-base/test_tir_op_types.py | 14 +-
.../tirx-transform/test_tir_transform_vectorize.py | 2 +-
.../operator/tile_primitive/cuda/copy/test_reg.py | 8 +-
tests/python/tirx/test_layout.py | 8 +-
.../tvmscript/test_tvmscript_printer_metadata.py | 8 +-
382 files changed, 3320 insertions(+), 4845 deletions(-)
diff --git a/include/tvm/ir/prim/builtin.h b/include/tvm/ir/prim/builtin.h
index dd9d2f7e50..3c475cda71 100644
--- a/include/tvm/ir/prim/builtin.h
+++ b/include/tvm/ir/prim/builtin.h
@@ -31,6 +31,12 @@ namespace builtin {
*/
TVM_DLL const Op& vscale();
+/*! \brief Round up to the nearest integral value. */
+TVM_DLL const Op& ceil();
+
+/*! \brief Base-two logarithm. */
+TVM_DLL const Op& log2();
+
/*! \brief Left shift. */
TVM_DLL const Op& shift_left();
diff --git a/include/tvm/ir/prim/expr.h b/include/tvm/ir/prim/expr.h
index 99740c8288..6920c5f626 100644
--- a/include/tvm/ir/prim/expr.h
+++ b/include/tvm/ir/prim/expr.h
@@ -58,7 +58,7 @@ class StringImmNode : public ExprNode {
namespace refl = tvm::ffi::reflection;
refl::ObjectDef<StringImmNode>().def_ro("value", &StringImmNode::value);
}
- TVM_FFI_DECLARE_OBJECT_INFO_FINAL("ir.prim.StringImm", StringImmNode, ExprNode);
+ TVM_FFI_DECLARE_OBJECT_INFO_FINAL("prim.StringImm", StringImmNode, ExprNode);
};
/*!
@@ -85,7 +85,7 @@ class CastNode : public ExprNode {
namespace refl = tvm::ffi::reflection;
refl::ObjectDef<CastNode>().def_ro("value", &CastNode::value);
}
- TVM_FFI_DECLARE_OBJECT_INFO_FINAL("ir.prim.Cast", CastNode, ExprNode);
+ TVM_FFI_DECLARE_OBJECT_INFO_FINAL("prim.Cast", CastNode, ExprNode);
};
/*!
@@ -123,7 +123,7 @@ class BinaryOpNode : public ExprNode {
/*! \brief a + b */
class AddNode : public BinaryOpNode<AddNode> {
public:
- static constexpr const char* _type_key = "ir.prim.Add";
+ static constexpr const char* _type_key = "prim.Add";
};
/*!
@@ -141,7 +141,7 @@ class Add : public PrimExpr {
/*! \brief a - b */
class SubNode : public BinaryOpNode<SubNode> {
public:
- static constexpr const char* _type_key = "ir.prim.Sub";
+ static constexpr const char* _type_key = "prim.Sub";
};
/*!
@@ -160,7 +160,7 @@ class Sub : public PrimExpr {
/*! \brief a * b */
class MulNode : public BinaryOpNode<MulNode> {
public:
- static constexpr const char* _type_key = "ir.prim.Mul";
+ static constexpr const char* _type_key = "prim.Mul";
};
/*!
@@ -181,7 +181,7 @@ class Mul : public PrimExpr {
*/
class DivNode : public BinaryOpNode<DivNode> {
public:
- static constexpr const char* _type_key = "ir.prim.Div";
+ static constexpr const char* _type_key = "prim.Div";
};
/*!
@@ -202,7 +202,7 @@ class Div : public PrimExpr {
*/
class ModNode : public BinaryOpNode<ModNode> {
public:
- static constexpr const char* _type_key = "ir.prim.Mod";
+ static constexpr const char* _type_key = "prim.Mod";
};
/*!
@@ -220,7 +220,7 @@ class Mod : public PrimExpr {
/*! \brief Floor division, floor(a/b) */
class FloorDivNode : public BinaryOpNode<FloorDivNode> {
public:
- static constexpr const char* _type_key = "ir.prim.FloorDiv";
+ static constexpr const char* _type_key = "prim.FloorDiv";
};
/*!
@@ -238,7 +238,7 @@ class FloorDiv : public PrimExpr {
/*! \brief The remainder of the floordiv */
class FloorModNode : public BinaryOpNode<FloorModNode> {
public:
- static constexpr const char* _type_key = "ir.prim.FloorMod";
+ static constexpr const char* _type_key = "prim.FloorMod";
};
/*!
@@ -256,7 +256,7 @@ class FloorMod : public PrimExpr {
/*! \brief min(a, b) */
class MinNode : public BinaryOpNode<MinNode> {
public:
- static constexpr const char* _type_key = "ir.prim.Min";
+ static constexpr const char* _type_key = "prim.Min";
};
/*!
@@ -274,7 +274,7 @@ class Min : public PrimExpr {
/*! \brief max(a, b) */
class MaxNode : public BinaryOpNode<MaxNode> {
public:
- static constexpr const char* _type_key = "ir.prim.Max";
+ static constexpr const char* _type_key = "prim.Max";
};
/*!
@@ -312,7 +312,7 @@ class CmpOpNode : public ExprNode {
/*! \brief a == b */
class EQNode : public CmpOpNode<EQNode> {
public:
- static constexpr const char* _type_key = "ir.prim.EQ";
+ static constexpr const char* _type_key = "prim.EQ";
};
/*!
@@ -330,7 +330,7 @@ class EQ : public PrimExpr {
/*! \brief a != b */
class NENode : public CmpOpNode<NENode> {
public:
- static constexpr const char* _type_key = "ir.prim.NE";
+ static constexpr const char* _type_key = "prim.NE";
};
/*!
@@ -348,7 +348,7 @@ class NE : public PrimExpr {
/*! \brief a < b */
class LTNode : public CmpOpNode<LTNode> {
public:
- static constexpr const char* _type_key = "ir.prim.LT";
+ static constexpr const char* _type_key = "prim.LT";
};
/*!
@@ -366,7 +366,7 @@ class LT : public PrimExpr {
/*! \brief a <= b */
struct LENode : public CmpOpNode<LENode> {
public:
- static constexpr const char* _type_key = "ir.prim.LE";
+ static constexpr const char* _type_key = "prim.LE";
};
/*!
@@ -384,7 +384,7 @@ class LE : public PrimExpr {
/*! \brief a > b */
class GTNode : public CmpOpNode<GTNode> {
public:
- static constexpr const char* _type_key = "ir.prim.GT";
+ static constexpr const char* _type_key = "prim.GT";
};
/*!
@@ -402,7 +402,7 @@ class GT : public PrimExpr {
/*! \brief a >= b */
class GENode : public CmpOpNode<GENode> {
public:
- static constexpr const char* _type_key = "ir.prim.GE";
+ static constexpr const char* _type_key = "prim.GE";
};
/*!
@@ -428,7 +428,7 @@ class AndNode : public ExprNode {
namespace refl = tvm::ffi::reflection;
refl::ObjectDef<AndNode>().def_ro("a", &AndNode::a).def_ro("b", &AndNode::b);
}
- TVM_FFI_DECLARE_OBJECT_INFO_FINAL("ir.prim.And", AndNode, ExprNode);
+ TVM_FFI_DECLARE_OBJECT_INFO_FINAL("prim.And", AndNode, ExprNode);
};
/*!
@@ -454,7 +454,7 @@ class OrNode : public ExprNode {
namespace refl = tvm::ffi::reflection;
refl::ObjectDef<OrNode>().def_ro("a", &OrNode::a).def_ro("b", &OrNode::b);
}
- TVM_FFI_DECLARE_OBJECT_INFO_FINAL("ir.prim.Or", OrNode, ExprNode);
+ TVM_FFI_DECLARE_OBJECT_INFO_FINAL("prim.Or", OrNode, ExprNode);
};
/*!
@@ -478,7 +478,7 @@ class NotNode : public ExprNode {
namespace refl = tvm::ffi::reflection;
refl::ObjectDef<NotNode>().def_ro("a", &NotNode::a);
}
- TVM_FFI_DECLARE_OBJECT_INFO_FINAL("ir.prim.Not", NotNode, ExprNode);
+ TVM_FFI_DECLARE_OBJECT_INFO_FINAL("prim.Not", NotNode, ExprNode);
};
/*!
@@ -515,7 +515,7 @@ class SelectNode : public ExprNode {
.def_ro("true_value", &SelectNode::true_value)
.def_ro("false_value", &SelectNode::false_value);
}
- TVM_FFI_DECLARE_OBJECT_INFO_FINAL("ir.prim.Select", SelectNode, ExprNode);
+ TVM_FFI_DECLARE_OBJECT_INFO_FINAL("prim.Select", SelectNode, ExprNode);
};
/*!
@@ -549,7 +549,7 @@ class LetNode : public ExprNode {
.def_ro("value", &LetNode::value)
.def_ro("body", &LetNode::body);
}
- TVM_FFI_DECLARE_OBJECT_INFO_FINAL("ir.prim.Let", LetNode, ExprNode);
+ TVM_FFI_DECLARE_OBJECT_INFO_FINAL("prim.Let", LetNode, ExprNode);
};
/*!
diff --git a/include/tvm/tirx/op.h b/include/tvm/ir/prim/op.h
similarity index 62%
copy from include/tvm/tirx/op.h
copy to include/tvm/ir/prim/op.h
index 3e1aa55f0b..7d12045496 100644
--- a/include/tvm/tirx/op.h
+++ b/include/tvm/ir/prim/op.h
@@ -17,136 +17,19 @@
* under the License.
*/
-/*!
- * \file tvm/tirx/op.h
- * \brief Common operators defined for Expr.
- *
- * \note Most of the operator defined here perform simple constant folding
- * when the type is int32 or int64 for simplifying the index expressions.
+/*! \file tvm/ir/prim/op.h
+ * \brief Shared primitive construction and constant helpers.
*/
-// Acknowledgement: Most operator APIs originate from Halide.
-#ifndef TVM_TIRX_OP_H_
-#define TVM_TIRX_OP_H_
-
-#include <tvm/ir/expr.h>
-#include <tvm/ir/op.h>
+#ifndef TVM_IR_PRIM_OP_H_
+#define TVM_IR_PRIM_OP_H_
#include <tvm/ir/prim/builtin.h>
#include <tvm/ir/prim/expr.h>
-#include <tvm/ir/type.h>
-#include <tvm/tirx/builtin.h>
-#include <tvm/tirx/op_attr_types.h>
-#include <tvm/tirx/stmt.h>
#include <algorithm>
#include <limits>
#include <type_traits>
#include <utility>
-
namespace tvm {
-
-#define TVM_TIR_REGISTER_OP(OpName) \
- TVM_REGISTER_OP("tirx." OpName) \
- .set_attr<TScriptPrinterName>("TScriptPrinterName", OpName) \
- .set_attr<TIRxOpCategory>("TIRxOpCategory", ffi::String("builtin"), /*plevel=*/1)
-
-#define TVM_TIRX_REGISTER_OP(OpName) TVM_TIR_REGISTER_OP(OpName)
-
-// Most common operators can be overloaded by argument type(PrimExpr).
-// So we put them under the root namespace.
-//
-// We put more developer oriented APIs -- MakeConst and is_const under tirx
-// as they are more specific to the tirx namespace.
-
-/*!
- * \brief Get the type of the expression under the unified type system.
- *
- * This function could return a more refined type than the runtime dtype
- * implied by PrimExpr::ty().
- *
- * \param expr The input parameter.
- * \return The result type.
- *
- * \sa tvm/ir/type.h for discussion about the relation between Type and DLPack dtype.
- */
-TVM_DLL Type GetType(const PrimExpr& expr);
-
-/*!
- * \brief Get the type corresponding to a runtime DLPack dtype.
- * \param dtype The runtime dtype.
- * \return The result type
- *
- * \sa tvm/ir/type.h for discussion about the relation between Type and DLPack dtype.
- */
-TVM_DLL Type GetTypeFromRuntimeDataType(DLDataType dtype);
-
-/*!
- * \brief Return from a thread.
- *
- * \param span The location of this operation in the source.
- * \return The return expression.
- */
-TVM_DLL PrimExpr thread_return(Span span = Span());
-
-/*!
- * \brief Continue current loop.
- * \param span The location of this operation in the source.
- * \return The continue loop expression.
- */
-TVM_DLL PrimExpr continue_loop(Span span = Span());
-
-/*!
- * \brief Break current loop.
- * \param span The location of this operation in the source.
- * \return The break loop expression.
- */
-TVM_DLL PrimExpr break_loop(Span span = Span());
-
-/*!
- * Query the maximum possible value of dtype.
- * \param dtype The primitive type.
- * \param span The location of this operation in the source.
- * \return the maximum possible value in this format.
- */
-TVM_DLL PrimExpr max_value(PrimType dtype, Span span = Span());
-
-/*!
- * Query the minimum possible value of dtype.
- * \param dtype The primitive type.
- * \param span The location of this operation in the source.
- * \return the minimum possible value in this format.
- */
-TVM_DLL PrimExpr min_value(PrimType dtype, Span span = Span());
-
-/*!
- * Get the value of infinity.
- * \param dtype The primitive type.
- * \param span The location of this operation in the source.
- * \return the infinity value in this format.
- */
-TVM_DLL PrimExpr infinity(PrimType dtype, Span span = Span());
-
-/*!
- * \brief cast value to type.
- *
- * \param t the target type.
- * \param value The value
- * \param span The location of this operation in the source.
- * \return The result expression.
- * \note This function may return value if the type is the same.
- */
-TVM_DLL PrimExpr cast(PrimType t, PrimExpr value, Span span = Span());
-/*!
- * \brief perform reinterpret cast value to type.
- *
- * \param t the target type.
- * \param value The value
- * \param span The location of this operation in the source.
- * \return The result expression.
- * \note This function may return value if the type is the same.
- */
-TVM_DLL PrimExpr reinterpret(PrimType t, PrimExpr value, Span span = Span());
-/*! \brief Perform a reinterpret cast involving an exact primitive or pointer type. */
-TVM_DLL Expr reinterpret(Type target_ty, Expr value, Span span = Span());
/*!
* \brief add operator
*
@@ -158,6 +41,7 @@ TVM_DLL Expr reinterpret(Type target_ty, Expr value, Span span = Span());
* index types(int32, int64) when possible.
*/
TVM_DLL PrimExpr add(PrimExpr a, PrimExpr b, Span span = Span());
+
/*!
* \brief subtraction operator
*
@@ -169,6 +53,7 @@ TVM_DLL PrimExpr add(PrimExpr a, PrimExpr b, Span span = Span());
* index types(int32, int64) when possible.
*/
TVM_DLL PrimExpr sub(PrimExpr a, PrimExpr b, Span span = Span());
+
/*!
* \brief negation.
*
@@ -179,6 +64,7 @@ TVM_DLL PrimExpr sub(PrimExpr a, PrimExpr b, Span span = Span());
* index types(int32, int64) when possible.
*/
TVM_DLL PrimExpr neg(PrimExpr a, Span span = Span());
+
/*!
* \brief multiplication operator
*
@@ -190,6 +76,7 @@ TVM_DLL PrimExpr neg(PrimExpr a, Span span = Span());
* index types(int32, int64) when possible.
*/
TVM_DLL PrimExpr mul(PrimExpr a, PrimExpr b, Span span = Span());
+
/*!
* \brief left shift operator
*
@@ -201,6 +88,7 @@ TVM_DLL PrimExpr mul(PrimExpr a, PrimExpr b, Span span = Span());
* index types(int32, int64) when possible.
*/
TVM_DLL PrimExpr left_shift(PrimExpr a, PrimExpr b, Span span = Span());
+
/*!
* \brief right shift operator
*
@@ -212,6 +100,7 @@ TVM_DLL PrimExpr left_shift(PrimExpr a, PrimExpr b, Span span = Span());
* index types(int32, int64) when possible.
*/
TVM_DLL PrimExpr right_shift(PrimExpr a, PrimExpr b, Span span = Span());
+
/*!
* \brief greater
*
@@ -223,6 +112,7 @@ TVM_DLL PrimExpr right_shift(PrimExpr a, PrimExpr b, Span span = Span());
* index types(int32, int64) when possible.
*/
TVM_DLL PrimExpr greater(PrimExpr a, PrimExpr b, Span span = Span());
+
/*!
* \brief greater_equal
*
@@ -234,6 +124,7 @@ TVM_DLL PrimExpr greater(PrimExpr a, PrimExpr b, Span span = Span());
* index types(int32, int64) when possible.
*/
TVM_DLL PrimExpr greater_equal(PrimExpr a, PrimExpr b, Span span = Span());
+
/*!
* \brief less
*
@@ -245,6 +136,7 @@ TVM_DLL PrimExpr greater_equal(PrimExpr a, PrimExpr b, Span span = Span());
* index types(int32, int64) when possible.
*/
TVM_DLL PrimExpr less(PrimExpr a, PrimExpr b, Span span = Span());
+
/*!
* \brief less_equal
*
@@ -256,6 +148,7 @@ TVM_DLL PrimExpr less(PrimExpr a, PrimExpr b, Span span = Span());
* index types(int32, int64) when possible.
*/
TVM_DLL PrimExpr less_equal(PrimExpr a, PrimExpr b, Span span = Span());
+
/*!
* \brief equal
*
@@ -267,6 +160,7 @@ TVM_DLL PrimExpr less_equal(PrimExpr a, PrimExpr b, Span span = Span());
* index types(int32, int64) when possible.
*/
TVM_DLL PrimExpr equal(PrimExpr a, PrimExpr b, Span span = Span());
+
/*!
* \brief not_equal
*
@@ -278,6 +172,7 @@ TVM_DLL PrimExpr equal(PrimExpr a, PrimExpr b, Span span = Span());
* index types(int32, int64) when possible.
*/
TVM_DLL PrimExpr not_equal(PrimExpr a, PrimExpr b, Span span = Span());
+
/*!
* \brief and
*
@@ -288,6 +183,7 @@ TVM_DLL PrimExpr not_equal(PrimExpr a, PrimExpr b, Span span = Span());
* \note This operator does eager constant folding.
*/
TVM_DLL PrimExpr logical_and(PrimExpr a, PrimExpr b, Span span = Span());
+
/*!
* \brief or
*
@@ -298,6 +194,7 @@ TVM_DLL PrimExpr logical_and(PrimExpr a, PrimExpr b, Span span = Span());
* \note This operator does eager constant folding.
*/
TVM_DLL PrimExpr logical_or(PrimExpr a, PrimExpr b, Span span = Span());
+
/*!
* \brief not
*
@@ -307,6 +204,7 @@ TVM_DLL PrimExpr logical_or(PrimExpr a, PrimExpr b, Span span = Span());
* \note This operator does eager constant folding.
*/
TVM_DLL PrimExpr logical_not(PrimExpr a, Span span = Span());
+
/*!
* \brief compute division in C semantics.
*
@@ -322,6 +220,7 @@ TVM_DLL PrimExpr logical_not(PrimExpr a, Span span = Span());
* index types(int32, int64) when possible.
*/
TVM_DLL PrimExpr div(PrimExpr a, PrimExpr b, Span span = Span());
+
/*!
* \brief compute trunc(a / b)
*
@@ -335,6 +234,7 @@ TVM_DLL PrimExpr div(PrimExpr a, PrimExpr b, Span span = Span());
* index types(int32, int64) when possible.
*/
TVM_DLL PrimExpr truncdiv(PrimExpr a, PrimExpr b, Span span = Span());
+
/*!
* \brief compute the remainder of truncdiv
*
@@ -348,6 +248,7 @@ TVM_DLL PrimExpr truncdiv(PrimExpr a, PrimExpr b, Span span = Span());
* index types(int32, int64) when possible.
*/
TVM_DLL PrimExpr truncmod(PrimExpr a, PrimExpr b, Span span = Span());
+
/*!
* \brief compute floor(a / b) where a and b are non-negative.
*
@@ -364,6 +265,7 @@ TVM_DLL PrimExpr truncmod(PrimExpr a, PrimExpr b, Span span = Span());
* index types(int32, int64) when possible.
*/
TVM_DLL PrimExpr indexdiv(PrimExpr a, PrimExpr b, Span span = Span());
+
/*!
* \brief compute ceil(a / b) where a and b are non-negative.
*
@@ -380,6 +282,7 @@ TVM_DLL PrimExpr indexdiv(PrimExpr a, PrimExpr b, Span span = Span());
* shape types(int32, int64) when possible.
*/
TVM_DLL PrimExpr shapediv(PrimExpr a, PrimExpr b, Span span = Span());
+
/*!
* \brief compute the remainder floor(a / b) where a and b are non-negative.
*
@@ -395,6 +298,7 @@ TVM_DLL PrimExpr shapediv(PrimExpr a, PrimExpr b, Span span = Span());
* index types(int32, int64) when possible.
*/
TVM_DLL PrimExpr indexmod(PrimExpr a, PrimExpr b, Span span = Span());
+
/*!
* \brief compute floor(a / b)
*
@@ -406,15 +310,7 @@ TVM_DLL PrimExpr indexmod(PrimExpr a, PrimExpr b, Span span = Span());
* index types(int32, int64) when possible.
*/
TVM_DLL PrimExpr floordiv(PrimExpr a, PrimExpr b, Span span = Span());
-/*!
- * \brief Compute log(exp(a) + exp(b)).
- *
- * \param a Left operand.
- * \param b Right operand.
- * \param span The location of this operation in the source.
- * \return The result expression.
- */
-TVM_DLL PrimExpr logaddexp(PrimExpr a, PrimExpr b, Span span = Span());
+
/*!
* \brief compute ceil(a / b)
*
@@ -427,6 +323,7 @@ TVM_DLL PrimExpr logaddexp(PrimExpr a, PrimExpr b, Span span = Span());
*/
TVM_DLL PrimExpr ceildiv(PrimExpr a, PrimExpr b, Span span = Span());
+
/*!
* \brief compute the remainder of floordiv
*
@@ -438,6 +335,7 @@ TVM_DLL PrimExpr ceildiv(PrimExpr a, PrimExpr b, Span span = Span());
* index types(int32, int64) when possible.
*/
TVM_DLL PrimExpr floormod(PrimExpr a, PrimExpr b, Span span = Span());
+
/*!
* \brief take maximum of two values
*
@@ -449,6 +347,7 @@ TVM_DLL PrimExpr floormod(PrimExpr a, PrimExpr b, Span span = Span());
* index types(int32, int64) when possible.
*/
TVM_DLL PrimExpr max(PrimExpr a, PrimExpr b, Span span = Span());
+
/*!
* \brief take minimum of two values
*
@@ -460,6 +359,7 @@ TVM_DLL PrimExpr max(PrimExpr a, PrimExpr b, Span span = Span());
* index types(int32, int64) when possible.
*/
TVM_DLL PrimExpr min(PrimExpr a, PrimExpr b, Span span = Span());
+
/*!
* \brief take bitwise and of two values
*
@@ -471,6 +371,7 @@ TVM_DLL PrimExpr min(PrimExpr a, PrimExpr b, Span span = Span());
* index types(int32, int64) when possible.
*/
TVM_DLL PrimExpr bitwise_and(PrimExpr a, PrimExpr b, Span span = Span());
+
/*!
* \brief take bitwise or of two values
*
@@ -482,6 +383,7 @@ TVM_DLL PrimExpr bitwise_and(PrimExpr a, PrimExpr b, Span span = Span());
* index types(int32, int64) when possible.
*/
TVM_DLL PrimExpr bitwise_or(PrimExpr a, PrimExpr b, Span span = Span());
+
/*!
* \brief take bitwise xor of two values
*
@@ -493,6 +395,7 @@ TVM_DLL PrimExpr bitwise_or(PrimExpr a, PrimExpr b, Span span = Span());
* index types(int32, int64) when possible.
*/
TVM_DLL PrimExpr bitwise_xor(PrimExpr a, PrimExpr b, Span span = Span());
+
/*!
* \brief take bitwise negation of two values
*
@@ -503,6 +406,7 @@ TVM_DLL PrimExpr bitwise_xor(PrimExpr a, PrimExpr b, Span span = Span());
* index types(int32, int64) when possible.
*/
TVM_DLL PrimExpr bitwise_neg(PrimExpr a, Span span = Span());
+
/*!
* \brief Conditional expression.
*
@@ -516,6 +420,7 @@ TVM_DLL PrimExpr bitwise_neg(PrimExpr a, Span span = Span());
*/
TVM_DLL PrimExpr if_then_else(PrimExpr cond, PrimExpr true_value, PrimExpr false_value,
Span span = Span());
+
/*!
* \brief Mark condition as likely.
* \param cond The condition
@@ -523,288 +428,41 @@ TVM_DLL PrimExpr if_then_else(PrimExpr cond, PrimExpr true_value, PrimExpr false
* \return The marked expression.
*/
TVM_DLL PrimExpr likely(PrimExpr cond, Span span = Span());
-/*!
- * \brief Calculate power(x, y)
- * \param x The left operand.
- * \param y The right operand.
- * \param span The location of this operation in the source.
- */
-TVM_DLL PrimExpr pow(PrimExpr x, PrimExpr y, Span span = Span());
-/*!
- * \brief Calculate absolute value of x.
- * \param x The input data
- * \param span The location of this operation in the source.
- *
- * \return The absolute value of input data x
- */
-TVM_DLL PrimExpr abs(PrimExpr x, Span span = Span());
-/*!
- * \brief Check if x is NaN.
- * \param x The input data
- * \param span The location of this operation in the source.
- * \return The result expression.
- */
-TVM_DLL PrimExpr isnan(PrimExpr x, Span span = Span());
-
-/*!
- * \brief Check if x is finite.
- * \param x The input data
- * \param span The location of this operation in the source.
- * \return The result expression.
- */
-TVM_DLL PrimExpr isfinite(PrimExpr x, Span span = Span());
-
-/*!
- * \brief Check if x is infinite.
- * \param x The input data
- * \param span The location of this operation in the source.
- * \return The result expression.
- */
-TVM_DLL PrimExpr isinf(PrimExpr x, Span span = Span());
-
-/*!
- * \brief sum of source expression over axis
- * \param source The source expression.
- * \param axis List of iteration variables that will be used for reduction.
- * \param init The value with which to initialize the output.
- * \param span The location of this operation in the source.
- * \return The result.
- */
-TVM_DLL PrimExpr sum(PrimExpr source, ffi::Array<tirx::IterVar> axis,
- ffi::Array<PrimExpr> init = {}, Span span = Span());
-
-/*!
- * \brief logical And of source expression over axis
- * \param source The source expression.
- * \param axis List of iteration variables that will be used for reduction.
- * \param init The value with which to initialize the output.
- * \param span The location of this operation in the source.
- */
-TVM_DLL PrimExpr all(PrimExpr source, ffi::Array<tirx::IterVar> axis,
- ffi::Array<PrimExpr> init = {}, Span span = Span());
-
-/*!
- * \brief logical Or of source expression over axis
- * \param source The source expression.
- * \param axis List of iteration variables that will be used for reduction.
- * \param init The value with which to initialize the output.
- * \param span The location of this operation in the source.
- * \return The result.
- */
-TVM_DLL PrimExpr any(PrimExpr source, ffi::Array<tirx::IterVar> axis,
- ffi::Array<PrimExpr> init = {}, Span span = Span());
-
-/*!
- * \brief max of source expression over axis
- * \param source The source expression.
- * \param axis List of iteration variables that will be used for reduction.
- * \param init The value with which to initialize the output.
- * \param span The location of this operation in the source.
- * \return The result.
- */
-TVM_DLL PrimExpr max(PrimExpr source, ffi::Array<tirx::IterVar> axis,
- ffi::Array<PrimExpr> init = {}, Span span = Span());
-
-/*!
- * \brief max of source expression over axis
- * \param source The source expression.
- * \param axis List of iteration variables that will be used for reduction.
- * \param init The value with which to initialize the output.
- * \param span The location of this operation in the source.
- * \return The result.
- */
-TVM_DLL PrimExpr min(PrimExpr source, ffi::Array<tirx::IterVar> axis,
- ffi::Array<PrimExpr> init = {}, Span span = Span());
-
-/*!
- * \brief product of source expression over axis
- * \param source The source expression.
- * \param axis List of iteration variables that will be used for reduction.
- * \param init The value with which to initialize the output.
- * \param span The location of this operation in the source.
- * \return The result.
- */
-TVM_DLL PrimExpr prod(PrimExpr source, ffi::Array<tirx::IterVar> axis,
- ffi::Array<PrimExpr> init = {}, Span span = Span());
-
-/*!
- * \brief Calculate floor(x)
- * \param x The input expression.
- * \param span The location of this operation in the source.
- * \return The result expression.
- */
-TVM_DLL PrimExpr floor(PrimExpr x, Span span = Span());
-
-/*!
- * \brief Calculate ceil(x)
- * \param x The input expression.
- * \param span The location of this operation in the source.
- * \return The result expression.
- */
+/*! \brief Round up to the nearest integral value, preserving the input type. */
TVM_DLL PrimExpr ceil(PrimExpr x, Span span = Span());
-
-/*!
- * \brief Round x to the nearest integer, ties to even.
- *
- * Uses IEEE 754 default rounding mode (ties-to-even / banker's rounding).
- * Constant-folding and all backends consistently use std::nearbyint semantics.
- *
- * \param x The input expression.
- * \param span The location of this operation in the source.
- * \return The result expression.
- */
-TVM_DLL PrimExpr round(PrimExpr x, Span span = Span());
-
+/*! \brief Construct a base-two logarithm with the input's primitive type. */
+TVM_DLL PrimExpr log2(PrimExpr x, Span span = Span());
+namespace prim {
/*!
- * \brief Round x to the nearest integer, ties to even.
- *
- * Equivalent to round(). Both use IEEE 754 default rounding mode (ties-to-even).
- *
- * \param x The input expression.
- * \param span The location of this operation in the source.
- * \return The result expression.
- */
-TVM_DLL PrimExpr nearbyint(PrimExpr x, Span span = Span());
-
-/*!
- * \brief Calculate trunc(x)
- * \param x The input expression.
+ * Query the minimum possible value of dtype.
+ * \param dtype The primitive type.
* \param span The location of this operation in the source.
- * \return The result expression.
+ * \return the minimum possible value in this format.
*/
-TVM_DLL PrimExpr trunc(PrimExpr x, Span span = Span());
+TVM_DLL PrimExpr min_value(PrimType dtype, Span span = Span());
/*!
- * \brief Construct a large uint constant by its low 32 bits and high 32bits.
- * \param value_ty The final primitive type.
- * \param low The lower 32 bits.
- * \param high The higher 32 bits.
+ * Query the maximum possible value of dtype.
+ * \param dtype The primitive type.
* \param span The location of this operation in the source.
- * \return The constructed expression.
+ * \return the maximum possible value in this format.
*/
-TVM_DLL PrimExpr LargeUIntImm(PrimType value_ty, int64_t low, int64_t high, Span span = Span());
+TVM_DLL PrimExpr max_value(PrimType dtype, Span span = Span());
/*!
- * \brief Execute a multiplication between two Q-numbers x and y
- * followed by a right shift s. The mathematical expression is:
- *
- * out = round(x*y*2^-s)
- *
- * Please note that the two Q-numbers x and y are supposed to have
- * the same number of fractional bits q.
- *
- * More about Q-numbers here: https://en.wikipedia.org/wiki/Q_(number_format)
+ * \brief cast value to type.
*
- * The rounding rule is to the nearest value, rounding half up
- * (i.e., round(x.1) = x and round (x.5) = x+1)
- * \param x first Q-number
- * \param y second Q-number
- * \param q number of fractional bits in x and y. Needs to be > 0
- * \param s integer right shift
+ * \param t the target type.
+ * \param value The value
* \param span The location of this operation in the source.
- * \return The constructed expression.
- */
-TVM_DLL PrimExpr q_multiply_shift(PrimExpr x, PrimExpr y, PrimExpr q, PrimExpr s,
- Span span = Span());
-
-/*!
- * \brief Fast_erf_float expression from Eigen
- *
- * \param arg The input expression.
- * \param bits The number of bits in the type.
- * \return The constructed expression.
- */
-TVM_DLL PrimExpr fast_erf_float_expr(PrimExpr arg, int bits);
-
-inline void CheckMathUnaryOpInputDType(const char* op_name, const PrimType& dtype) {
- TVM_FFI_CHECK(dtype.code() == DLDataTypeCode::kDLFloat ||
- dtype.MatchesElementType(DLDataTypeCode::kDLBfloat, 16),
- TypeError)
- << "tirx." << op_name << " only supports floating-point inputs, but got " << dtype;
-}
-
-// Intrinsic operators
-#define TVM_DECLARE_INTRIN_UNARY_WITH_CHECK(OpName, CheckInputDType) \
- inline PrimExpr OpName(PrimExpr x, Span span = Span()) { \
- static const Op op = Op::Get("tirx." #OpName); \
- PrimType x_ty = x.ty(); \
- CheckInputDType(#OpName, x_ty); \
- if (x_ty.MatchesElementType(DLDataTypeCode::kDLBfloat, 16)) { \
- PrimType bf16_ty = x_ty; \
- PrimType f32_ty = \
- x_ty.IsScalableVector() \
- ? PrimType::ScalableVector(DLDataTypeCode::kDLFloat, 32, x_ty.VScaleFactor()) \
- : PrimType::Float(32, x_ty.lanes()); \
- PrimExpr x_fp32 = prim::Cast(f32_ty, x, span); \
- PrimExpr result_fp32 = Call(f32_ty, op, {x_fp32}, {}, {}, span).as_or_throw<PrimExpr>(); \
- return prim::Cast(bf16_ty, result_fp32, span); \
- } else { \
- return Call(x_ty, op, {x}, {}, {}, span).as_or_throw<PrimExpr>(); \
- } \
- }
-
-#define TVM_DECLARE_INTRIN_UNARY(OpName) \
- TVM_DECLARE_INTRIN_UNARY_WITH_CHECK(OpName, [](const char*, const PrimType&) {})
-
-#define TVM_DECLARE_FLOAT_INTRIN_UNARY(OpName) \
- TVM_DECLARE_INTRIN_UNARY_WITH_CHECK(OpName, CheckMathUnaryOpInputDType)
-
-TVM_DECLARE_INTRIN_UNARY(exp);
-TVM_DECLARE_INTRIN_UNARY(exp2);
-TVM_DECLARE_INTRIN_UNARY(exp10);
-TVM_DECLARE_INTRIN_UNARY(erf);
-TVM_DECLARE_FLOAT_INTRIN_UNARY(tanh);
-TVM_DECLARE_INTRIN_UNARY(sigmoid);
-TVM_DECLARE_INTRIN_UNARY(sqrt);
-TVM_DECLARE_INTRIN_UNARY(rsqrt);
-TVM_DECLARE_INTRIN_UNARY(log);
-TVM_DECLARE_INTRIN_UNARY(log2);
-TVM_DECLARE_INTRIN_UNARY(log10);
-TVM_DECLARE_INTRIN_UNARY(log1p);
-TVM_DECLARE_INTRIN_UNARY(popcount);
-TVM_DECLARE_FLOAT_INTRIN_UNARY(tan);
-TVM_DECLARE_FLOAT_INTRIN_UNARY(cos);
-TVM_DECLARE_FLOAT_INTRIN_UNARY(cosh);
-TVM_DECLARE_FLOAT_INTRIN_UNARY(sin);
-TVM_DECLARE_FLOAT_INTRIN_UNARY(sinh);
-TVM_DECLARE_FLOAT_INTRIN_UNARY(asin);
-TVM_DECLARE_FLOAT_INTRIN_UNARY(acos);
-TVM_DECLARE_FLOAT_INTRIN_UNARY(atan);
-TVM_DECLARE_FLOAT_INTRIN_UNARY(acosh);
-TVM_DECLARE_FLOAT_INTRIN_UNARY(asinh);
-TVM_DECLARE_FLOAT_INTRIN_UNARY(atanh);
-TVM_DECLARE_INTRIN_UNARY(clz);
-
-#define TVM_DECLARE_INTRIN_BINARY(OpName) \
- inline PrimExpr OpName(PrimExpr x, PrimExpr y, Span span = Span()) { \
- static const Op op = Op::Get("tirx." #OpName); \
- return Call(x.ty(), op, {x, y}, {}, {}, span).as_or_throw<PrimExpr>(); \
- }
-
-TVM_DECLARE_INTRIN_BINARY(atan2);
-TVM_DECLARE_INTRIN_BINARY(nextafter);
-TVM_DECLARE_INTRIN_BINARY(copysign);
-TVM_DECLARE_INTRIN_BINARY(hypot);
-TVM_DECLARE_INTRIN_BINARY(ldexp);
-
-namespace tirx {
-
-/*!
- * \brief Check if type is a pointer to a runtime element type.
- * \param type The type to be checked.
- * \param element_type The corresponding element type.
- * \return The check results
+ * \return The result expression.
+ * \note This function may return value if the type is the same. Constant folding
+ * uses MakeConst for scalar and vector constants.
*/
-inline bool IsPointerType(const Type& type, DLDataType element_type) {
- if (type.IsMissing()) return false;
- if (const auto* ptr_type = type.as<PointerTypeNode>()) {
- if (const auto* prim_type = ptr_type->element_type.as<PrimTypeNode>()) {
- return prim_type->dtype == element_type;
- }
- }
- return false;
-}
-
+TVM_DLL PrimExpr cast(PrimType t, PrimExpr value, Span span = Span());
+TVM_DLL PrimExpr cast(DLDataType t, PrimExpr value, Span span = Span());
+/*! \brief Construct integer absolute value; floating absolute value belongs to TIRX. */
+TVM_DLL PrimExpr IntegerAbs(PrimExpr x, Span span = Span());
/*!
* \brief Make a const value with certain data type.
*
@@ -824,13 +482,7 @@ template <typename ValueType,
typename = typename std::enable_if<std::is_standard_layout<ValueType>::value &&
std::is_trivial<ValueType>::value>::type>
inline PrimExpr MakeConst(PrimType dtype, ValueType value, Span span = Span());
-/*!
- * \brief Make a constant opaque-pointer value.
- * \param value The integer payload to reinterpret as a handle.
- * \param span The location of this operation in the source.
- * \return The result expression.
- */
-inline Expr ConstHandle(int64_t value, Span span = Span());
+
/*!
* \brief Get x as constant int expression.
* \param x The expression
@@ -854,13 +506,6 @@ inline const int64_t* as_const_int(const PrimExpr& x) {
*/
inline bool is_const_int(const PrimExpr& x, int64_t value);
-/*!
- * \brief Check whether stmt is nop.
- * \param stmt The input statement
- * \return whether stmt is nop
- */
-inline bool is_no_op(const tirx::Stmt& stmt);
-
/*!
* \brief Check whether x is a constant integer 1
* \param x The input argument.
@@ -891,24 +536,6 @@ inline bool is_const_int(const PrimExpr& x);
*/
inline bool is_const_number(const PrimExpr& x);
-/*!
- * \brief Left fold.
- * \param freduce The reduction function.
- * \param init_value The initial value.
- * \param values The values to be folded.
- * \param span The location of the fold in the source.
- * \return The result.
- * \tparam FReduce The type of the reduction.
- */
-template <typename FReduce>
-inline PrimExpr foldl(FReduce freduce, PrimExpr init_value, const ffi::Array<PrimExpr>& values,
- Span span = Span()) {
- for (PrimExpr val : values) {
- init_value = freduce(init_value, val, span);
- }
- return init_value;
-}
-
/*!
* \brief Check whether x is a constant power of two
* If x is power of two, write the power to the shift.
@@ -948,17 +575,15 @@ inline bool is_const_int(const PrimExpr& x, int64_t value) {
return as_int && (*as_int == value);
}
-inline bool is_no_op(const tirx::Stmt& stmt) {
- if (!stmt.defined()) return true;
- if (const auto* op = stmt.as<tirx::EvaluateNode>()) {
- auto value = op->value.as<PrimExpr>();
- return value && is_const_int(value.value());
- }
- if (const auto* op = stmt.as<tirx::SeqStmtNode>()) {
- return op->seq.size() == 0;
- }
- return false;
-}
+/*!
+ * \brief Construct a large uint constant by its low 32 bits and high 32bits.
+ * \param value_ty The final primitive type.
+ * \param low The lower 32 bits.
+ * \param high The higher 32 bits.
+ * \param span The location of this operation in the source.
+ * \return The constructed expression.
+ */
+TVM_DLL PrimExpr LargeUIntImm(PrimType value_ty, int64_t low, int64_t high, Span span = Span());
template <typename ValueType>
inline PrimExpr MakeConstScalar(PrimType dtype, ValueType value, Span span = Span()) {
@@ -974,10 +599,8 @@ inline PrimExpr MakeConstScalar(PrimType dtype, ValueType value, Span span = Spa
} else if (uval <= static_cast<uint64_t>(std::numeric_limits<int64_t>::max())) {
return IntImm(dtype, static_cast<int64_t>(value), span);
} else {
- uint64_t mask = (static_cast<uint64_t>(1) << 32U) - 1U;
- uint64_t low = uval & mask;
- uint64_t high = uval >> 32U;
- return LargeUIntImm(dtype, static_cast<int64_t>(low), static_cast<int64_t>(high), span);
+ return LargeUIntImm(dtype, static_cast<int64_t>(uval & 0xffffffffULL),
+ static_cast<int64_t>(uval >> 32U), span);
}
}
if (dtype.MatchesCode(DLDataTypeCode::kDLFloat, DLDataTypeCode::kDLFloat8_e3m4,
@@ -1013,11 +636,7 @@ inline PrimExpr MakeConst(PrimType dtype, ValueType value, Span span) {
return prim::Broadcast(MakeConstScalar(elem_ty, value, span), lanes, span);
}
-inline Expr ConstHandle(int64_t value, Span span) {
- return reinterpret(PointerType::VoidPointerTy(), IntImm(PrimType::UInt(64), value, span), span);
-}
-
-} // namespace tirx
+} // namespace prim
// additional const expression overloading
#define TVM_DEFINE_ASSIGN_OP_OVERLOAD(Name, OpFunc) \
@@ -1029,8 +648,8 @@ inline Expr ConstHandle(int64_t value, Span span) {
#define TVM_DEFINE_BINOP_CONST_VAL_OVERLOAD(Name) \
inline PrimExpr Name(const PrimExpr& a, float b) { return Name(a, PrimExpr(b)); } \
inline PrimExpr Name(float a, const PrimExpr& b) { return Name(PrimExpr(a), b); } \
- inline PrimExpr Name(int a, const PrimExpr& b) { return Name(tirx::MakeConst(b.ty(), a), b); } \
- inline PrimExpr Name(const PrimExpr& a, int b) { return Name(a, tirx::MakeConst(a.ty(), b)); } \
+ inline PrimExpr Name(int a, const PrimExpr& b) { return Name(prim::MakeConst(b.ty(), a), b); } \
+ inline PrimExpr Name(const PrimExpr& a, int b) { return Name(a, prim::MakeConst(a.ty(), b)); } \
inline PrimExpr Name(const PrimExpr& a, double b) { \
return Name(a, FloatImm(PrimType::Float(64), b)); \
}
@@ -1043,10 +662,10 @@ inline Expr ConstHandle(int64_t value, Span span) {
return Name(PrimExpr(a), b, span); \
} \
inline PrimExpr Name(int a, const PrimExpr& b, Span span = Span()) { \
- return Name(tirx::MakeConst(b.ty(), a), b, span); \
+ return Name(prim::MakeConst(b.ty(), a), b, span); \
} \
inline PrimExpr Name(const PrimExpr& a, int b, Span span = Span()) { \
- return Name(a, tirx::MakeConst(a.ty(), b), span); \
+ return Name(a, prim::MakeConst(a.ty(), b), span); \
} \
inline PrimExpr Name(const PrimExpr& a, double b, Span span = Span()) { \
return Name(a, FloatImm(PrimType::Float(64), b), span); \
@@ -1065,15 +684,15 @@ inline Expr ConstHandle(int64_t value, Span span) {
}
#define TVM_DEFINE_INT_OP_CONST_VAL_OVERLOAD(Name) \
- inline PrimExpr Name(const PrimExpr& a, int b) { return Name(a, tirx::MakeConst(a.ty(), b)); } \
- inline PrimExpr Name(int a, const PrimExpr& b) { return Name(tirx::MakeConst(b.ty(), a), b); }
+ inline PrimExpr Name(const PrimExpr& a, int b) { return Name(a, prim::MakeConst(a.ty(), b)); } \
+ inline PrimExpr Name(int a, const PrimExpr& b) { return Name(prim::MakeConst(b.ty(), a), b); }
#define TVM_DEFINE_INT_OP_CONST_VAL_OVERLOAD_SPANNED(Name) \
inline PrimExpr Name(const PrimExpr& a, int b, Span span = Span()) { \
- return Name(a, tirx::MakeConst(a.ty(), b), span); \
+ return Name(a, prim::MakeConst(a.ty(), b), span); \
} \
inline PrimExpr Name(int a, const PrimExpr& b, Span span = Span()) { \
- return Name(tirx::MakeConst(b.ty(), a), b, span); \
+ return Name(prim::MakeConst(b.ty(), a), b, span); \
}
TVM_DEFINE_ASSIGN_OP_OVERLOAD(operator+=, operator+);
@@ -1102,7 +721,6 @@ TVM_DEFINE_INT_OP_CONST_VAL_OVERLOAD_SPANNED(indexmod);
TVM_DEFINE_INT_OP_CONST_VAL_OVERLOAD_SPANNED(truncdiv);
TVM_DEFINE_INT_OP_CONST_VAL_OVERLOAD_SPANNED(truncmod);
TVM_DEFINE_INT_OP_CONST_VAL_OVERLOAD_SPANNED(floordiv);
-TVM_DEFINE_INT_OP_CONST_VAL_OVERLOAD_SPANNED(logaddexp);
TVM_DEFINE_INT_OP_CONST_VAL_OVERLOAD_SPANNED(floormod);
TVM_DEFINE_INT_OP_CONST_VAL_OVERLOAD_SPANNED(right_shift); // NOLINT(*)
TVM_DEFINE_INT_OP_CONST_VAL_OVERLOAD_SPANNED(left_shift); // NOLINT(*)
@@ -1133,7 +751,7 @@ inline void DivAmbiguityError(const TA& a) {
"please call div, indexdiv/indexmod, "
"floordiv/floormod or truncdiv/truncmod directly "
"to avoid ambiguity in the code. "
- "Checkout these functions in tirx/op.h.");
+ "Checkout these functions in ir/prim/op.h.");
}
// The following code are not intended to be used in the codebase.
@@ -1158,5 +776,6 @@ inline PrimExpr operator%(const PrimExpr& a, const TB& b) {
DivAmbiguityError(a);
return a;
}
+
} // namespace tvm
-#endif // TVM_TIR_OP_H_
+#endif // TVM_IR_PRIM_OP_H_
diff --git a/include/tvm/ir/prim/vector_expr.h b/include/tvm/ir/prim/vector_expr.h
index 28177467e9..0ccb51aa1d 100644
--- a/include/tvm/ir/prim/vector_expr.h
+++ b/include/tvm/ir/prim/vector_expr.h
@@ -54,7 +54,7 @@ class RampNode : public ExprNode {
.def_ro("stride", &RampNode::stride)
.def_ro("lanes", &RampNode::lanes);
}
- TVM_FFI_DECLARE_OBJECT_INFO_FINAL("ir.prim.Ramp", RampNode, ExprNode);
+ TVM_FFI_DECLARE_OBJECT_INFO_FINAL("prim.Ramp", RampNode, ExprNode);
};
/*!
@@ -82,7 +82,7 @@ class BroadcastNode : public ExprNode {
.def_ro("value", &BroadcastNode::value)
.def_ro("lanes", &BroadcastNode::lanes);
}
- TVM_FFI_DECLARE_OBJECT_INFO_FINAL("ir.prim.Broadcast", BroadcastNode, ExprNode);
+ TVM_FFI_DECLARE_OBJECT_INFO_FINAL("prim.Broadcast", BroadcastNode, ExprNode);
};
/*!
@@ -114,7 +114,7 @@ class ShuffleNode : public ExprNode {
.def_ro("vectors", &ShuffleNode::vectors)
.def_ro("indices", &ShuffleNode::indices);
}
- TVM_FFI_DECLARE_OBJECT_INFO_FINAL("ir.prim.Shuffle", ShuffleNode, ExprNode);
+ TVM_FFI_DECLARE_OBJECT_INFO_FINAL("prim.Shuffle", ShuffleNode, ExprNode);
};
/*!
diff --git a/include/tvm/s_tir/analysis.h b/include/tvm/s_tir/analysis.h
index a44c9fed82..9c5640900f 100644
--- a/include/tvm/s_tir/analysis.h
+++ b/include/tvm/s_tir/analysis.h
@@ -35,7 +35,6 @@
namespace tvm {
namespace tirx {
-
/*!
* \brief Auto detect the block access region according to its body stmt
* It will detect the access region as an array in order of appearance in AST
@@ -96,8 +95,6 @@ class Analyzer;
} // namespace arith
namespace s_tir {
-using namespace tvm::prim;
-
using namespace tvm::tirx;
/*!
diff --git a/include/tvm/s_tir/backend/adreno/transform.h b/include/tvm/s_tir/backend/adreno/transform.h
index 4fe4b6c641..a4265dd1a2 100644
--- a/include/tvm/s_tir/backend/adreno/transform.h
+++ b/include/tvm/s_tir/backend/adreno/transform.h
@@ -35,7 +35,6 @@
namespace tvm {
namespace s_tir {
-using namespace tvm::prim;
namespace backend {
namespace adreno {
namespace transform {
diff --git a/include/tvm/s_tir/meta_schedule/arg_info.h b/include/tvm/s_tir/meta_schedule/arg_info.h
index a01e089386..586f5009a8 100644
--- a/include/tvm/s_tir/meta_schedule/arg_info.h
+++ b/include/tvm/s_tir/meta_schedule/arg_info.h
@@ -28,7 +28,6 @@
namespace tvm {
namespace s_tir {
-using namespace tvm::prim;
namespace meta_schedule {
/*! \brief The argument information. */
diff --git a/include/tvm/s_tir/meta_schedule/builder.h b/include/tvm/s_tir/meta_schedule/builder.h
index d65daaa218..866bbc08df 100644
--- a/include/tvm/s_tir/meta_schedule/builder.h
+++ b/include/tvm/s_tir/meta_schedule/builder.h
@@ -32,7 +32,6 @@
namespace tvm {
namespace s_tir {
-using namespace tvm::prim;
namespace meta_schedule {
/*! \brief The builder's input, containing an IRModule and the target. */
diff --git a/include/tvm/s_tir/meta_schedule/cost_model.h b/include/tvm/s_tir/meta_schedule/cost_model.h
index 22f6ab7ebf..9997179ddc 100644
--- a/include/tvm/s_tir/meta_schedule/cost_model.h
+++ b/include/tvm/s_tir/meta_schedule/cost_model.h
@@ -35,7 +35,6 @@
namespace tvm {
namespace s_tir {
-using namespace tvm::prim;
namespace meta_schedule {
class TuneContext;
diff --git a/include/tvm/s_tir/meta_schedule/database.h b/include/tvm/s_tir/meta_schedule/database.h
index 9975516da5..c7d184c7c2 100644
--- a/include/tvm/s_tir/meta_schedule/database.h
+++ b/include/tvm/s_tir/meta_schedule/database.h
@@ -36,7 +36,6 @@
namespace tvm {
namespace s_tir {
-using namespace tvm::prim;
namespace meta_schedule {
class ModuleEquality;
diff --git a/include/tvm/s_tir/meta_schedule/extracted_task.h b/include/tvm/s_tir/meta_schedule/extracted_task.h
index b66237f564..82fd7a58cd 100644
--- a/include/tvm/s_tir/meta_schedule/extracted_task.h
+++ b/include/tvm/s_tir/meta_schedule/extracted_task.h
@@ -37,7 +37,6 @@ class Tensor;
namespace tvm {
namespace s_tir {
-using namespace tvm::prim;
namespace meta_schedule {
/*! \brief A tuning task extracted from the high-level IR */
diff --git a/include/tvm/s_tir/meta_schedule/feature_extractor.h b/include/tvm/s_tir/meta_schedule/feature_extractor.h
index c288a59a86..70b9977136 100644
--- a/include/tvm/s_tir/meta_schedule/feature_extractor.h
+++ b/include/tvm/s_tir/meta_schedule/feature_extractor.h
@@ -30,7 +30,6 @@
namespace tvm {
namespace s_tir {
-using namespace tvm::prim;
namespace meta_schedule {
class TuneContext;
diff --git a/include/tvm/s_tir/meta_schedule/measure_callback.h b/include/tvm/s_tir/meta_schedule/measure_callback.h
index 9698ad7d19..c896856a68 100644
--- a/include/tvm/s_tir/meta_schedule/measure_callback.h
+++ b/include/tvm/s_tir/meta_schedule/measure_callback.h
@@ -34,7 +34,6 @@
namespace tvm {
namespace s_tir {
-using namespace tvm::prim;
namespace meta_schedule {
class TaskScheduler;
diff --git a/include/tvm/s_tir/meta_schedule/measure_candidate.h b/include/tvm/s_tir/meta_schedule/measure_candidate.h
index eda1831df5..4e33c8c252 100644
--- a/include/tvm/s_tir/meta_schedule/measure_candidate.h
+++ b/include/tvm/s_tir/meta_schedule/measure_candidate.h
@@ -29,7 +29,6 @@
namespace tvm {
namespace s_tir {
-using namespace tvm::prim;
namespace meta_schedule {
/*! \brief The schedule (with input shapes) to be measured. */
diff --git a/include/tvm/s_tir/meta_schedule/mutator.h b/include/tvm/s_tir/meta_schedule/mutator.h
index b2cdc1824d..4c64f3d0e5 100644
--- a/include/tvm/s_tir/meta_schedule/mutator.h
+++ b/include/tvm/s_tir/meta_schedule/mutator.h
@@ -31,7 +31,6 @@
namespace tvm {
namespace s_tir {
-using namespace tvm::prim;
namespace meta_schedule {
class TuneContext;
diff --git a/include/tvm/s_tir/meta_schedule/postproc.h b/include/tvm/s_tir/meta_schedule/postproc.h
index f66f8f533f..0e5b6f1593 100644
--- a/include/tvm/s_tir/meta_schedule/postproc.h
+++ b/include/tvm/s_tir/meta_schedule/postproc.h
@@ -28,7 +28,6 @@
namespace tvm {
namespace s_tir {
-using namespace tvm::prim;
namespace meta_schedule {
class TuneContext;
diff --git a/include/tvm/s_tir/meta_schedule/profiler.h b/include/tvm/s_tir/meta_schedule/profiler.h
index ee3cfdbc29..060c0d4766 100644
--- a/include/tvm/s_tir/meta_schedule/profiler.h
+++ b/include/tvm/s_tir/meta_schedule/profiler.h
@@ -35,7 +35,6 @@
namespace tvm {
namespace s_tir {
-using namespace tvm::prim;
namespace meta_schedule {
class ScopedTimer {
diff --git a/include/tvm/s_tir/meta_schedule/runner.h b/include/tvm/s_tir/meta_schedule/runner.h
index 085ac493e6..847fd5a24d 100644
--- a/include/tvm/s_tir/meta_schedule/runner.h
+++ b/include/tvm/s_tir/meta_schedule/runner.h
@@ -30,7 +30,6 @@
namespace tvm {
namespace s_tir {
-using namespace tvm::prim;
namespace meta_schedule {
/*! \brief Runner's input containing path of artifact, type of device and argument info. */
diff --git a/include/tvm/s_tir/meta_schedule/schedule/cuda/thread_bind.h b/include/tvm/s_tir/meta_schedule/schedule/cuda/thread_bind.h
index 71cf2b280f..d4e9c0942a 100644
--- a/include/tvm/s_tir/meta_schedule/schedule/cuda/thread_bind.h
+++ b/include/tvm/s_tir/meta_schedule/schedule/cuda/thread_bind.h
@@ -28,7 +28,6 @@
namespace tvm {
namespace s_tir {
-using namespace tvm::prim;
namespace meta_schedule {
/*!
diff --git a/include/tvm/s_tir/meta_schedule/schedule/generic/winograd.h b/include/tvm/s_tir/meta_schedule/schedule/generic/winograd.h
index fe6d0bf5da..5ddffc7207 100644
--- a/include/tvm/s_tir/meta_schedule/schedule/generic/winograd.h
+++ b/include/tvm/s_tir/meta_schedule/schedule/generic/winograd.h
@@ -24,7 +24,6 @@
namespace tvm {
namespace s_tir {
-using namespace tvm::prim;
namespace meta_schedule {
/*!
diff --git a/include/tvm/s_tir/meta_schedule/schedule_rule.h b/include/tvm/s_tir/meta_schedule/schedule_rule.h
index b6b164df7b..5d30c44ed7 100644
--- a/include/tvm/s_tir/meta_schedule/schedule_rule.h
+++ b/include/tvm/s_tir/meta_schedule/schedule_rule.h
@@ -32,7 +32,6 @@
namespace tvm {
namespace s_tir {
-using namespace tvm::prim;
namespace meta_schedule {
class TuneContext;
diff --git a/include/tvm/s_tir/meta_schedule/search_strategy.h b/include/tvm/s_tir/meta_schedule/search_strategy.h
index 77be5e8cb9..6eb07cd739 100644
--- a/include/tvm/s_tir/meta_schedule/search_strategy.h
+++ b/include/tvm/s_tir/meta_schedule/search_strategy.h
@@ -34,7 +34,6 @@
namespace tvm {
namespace s_tir {
-using namespace tvm::prim;
namespace meta_schedule {
// Forward declaration
diff --git a/include/tvm/s_tir/meta_schedule/space_generator.h b/include/tvm/s_tir/meta_schedule/space_generator.h
index d3e68ef48b..607d9e71a6 100644
--- a/include/tvm/s_tir/meta_schedule/space_generator.h
+++ b/include/tvm/s_tir/meta_schedule/space_generator.h
@@ -32,7 +32,6 @@
namespace tvm {
namespace s_tir {
-using namespace tvm::prim;
namespace meta_schedule {
// Forward declaration
diff --git a/include/tvm/s_tir/meta_schedule/task_scheduler.h b/include/tvm/s_tir/meta_schedule/task_scheduler.h
index 8ee6727e9e..e7fdb09bc1 100644
--- a/include/tvm/s_tir/meta_schedule/task_scheduler.h
+++ b/include/tvm/s_tir/meta_schedule/task_scheduler.h
@@ -37,7 +37,6 @@
namespace tvm {
namespace s_tir {
-using namespace tvm::prim;
namespace meta_schedule {
class TaskRecordNode : public ffi::Object {
diff --git a/include/tvm/s_tir/meta_schedule/tune_context.h b/include/tvm/s_tir/meta_schedule/tune_context.h
index 85a7481241..1e14fc4fca 100644
--- a/include/tvm/s_tir/meta_schedule/tune_context.h
+++ b/include/tvm/s_tir/meta_schedule/tune_context.h
@@ -37,7 +37,6 @@
namespace tvm {
namespace s_tir {
-using namespace tvm::prim;
namespace meta_schedule {
class TaskSchedulerNode;
diff --git a/include/tvm/s_tir/random_engine.h b/include/tvm/s_tir/random_engine.h
index cc12b5bcc7..e1c7b6fae9 100644
--- a/include/tvm/s_tir/random_engine.h
+++ b/include/tvm/s_tir/random_engine.h
@@ -32,8 +32,6 @@
namespace tvm {
namespace s_tir {
-using namespace tvm::prim;
-
/*!
* \brief This linear congruential engine is a drop-in replacement for std::minstd_rand. It strictly
* corresponds to std::minstd_rand and is designed to be platform-independent.
diff --git a/include/tvm/s_tir/schedule/instruction.h b/include/tvm/s_tir/schedule/instruction.h
index 0c17231a41..56fdb0ef0c 100644
--- a/include/tvm/s_tir/schedule/instruction.h
+++ b/include/tvm/s_tir/schedule/instruction.h
@@ -31,7 +31,6 @@ template <typename, typename>
class AttrRegistry;
namespace s_tir {
-using namespace tvm::prim;
using namespace tvm::tirx;
// Forward declaration
diff --git a/include/tvm/s_tir/schedule/schedule.h b/include/tvm/s_tir/schedule/schedule.h
index 7865b8cfb9..b149eb3b2c 100644
--- a/include/tvm/s_tir/schedule/schedule.h
+++ b/include/tvm/s_tir/schedule/schedule.h
@@ -27,7 +27,6 @@
namespace tvm {
namespace s_tir {
-using namespace tvm::prim;
using namespace tvm::tirx;
/*! \brief The level of detailed error message rendering */
diff --git a/include/tvm/s_tir/schedule/state.h b/include/tvm/s_tir/schedule/state.h
index 001a9e51d6..bbaeb30396 100644
--- a/include/tvm/s_tir/schedule/state.h
+++ b/include/tvm/s_tir/schedule/state.h
@@ -34,7 +34,6 @@
namespace tvm {
namespace s_tir {
-using namespace tvm::prim;
using namespace tvm::tirx;
/*!
diff --git a/include/tvm/s_tir/schedule/trace.h b/include/tvm/s_tir/schedule/trace.h
index 33202c9db5..dba4f6c80c 100644
--- a/include/tvm/s_tir/schedule/trace.h
+++ b/include/tvm/s_tir/schedule/trace.h
@@ -24,7 +24,6 @@
namespace tvm {
namespace s_tir {
-using namespace tvm::prim;
using namespace tvm::tirx;
// Forward declaration
diff --git a/include/tvm/s_tir/stmt.h b/include/tvm/s_tir/stmt.h
index 18ee044789..57dcd70aab 100644
--- a/include/tvm/s_tir/stmt.h
+++ b/include/tvm/s_tir/stmt.h
@@ -31,7 +31,6 @@
namespace tvm {
namespace s_tir {
-using namespace tvm::prim;
namespace attr {
/*!
diff --git a/include/tvm/s_tir/transform.h b/include/tvm/s_tir/transform.h
index 9752cebfac..6bbc91edfb 100644
--- a/include/tvm/s_tir/transform.h
+++ b/include/tvm/s_tir/transform.h
@@ -34,8 +34,6 @@
namespace tvm {
namespace s_tir {
-using namespace tvm::prim;
-
/*!
* \brief Renew the definition nodes for a TIR, including Var, Buffer and IterVar.
* This pass works as a simple DeepCopy to duplicate a function with different Vars and
diff --git a/include/tvm/te/tensor.h b/include/tvm/te/tensor.h
index cbec00fd00..da66e034a6 100644
--- a/include/tvm/te/tensor.h
+++ b/include/tvm/te/tensor.h
@@ -36,7 +36,6 @@
namespace tvm {
namespace te {
-
using arith::IntSet;
using namespace tvm::tirx;
diff --git a/include/tvm/tirx/op.h b/include/tvm/tirx/op.h
index 3e1aa55f0b..874719a4e3 100644
--- a/include/tvm/tirx/op.h
+++ b/include/tvm/tirx/op.h
@@ -32,6 +32,7 @@
#include <tvm/ir/op.h>
#include <tvm/ir/prim/builtin.h>
#include <tvm/ir/prim/expr.h>
+#include <tvm/ir/prim/op.h>
#include <tvm/ir/type.h>
#include <tvm/tirx/builtin.h>
#include <tvm/tirx/op_attr_types.h>
@@ -42,7 +43,7 @@
#include <type_traits>
#include <utility>
-namespace tvm {
+namespace tvm::prim {
#define TVM_TIR_REGISTER_OP(OpName) \
TVM_REGISTER_OP("tirx." OpName) \
@@ -51,11 +52,7 @@ namespace tvm {
#define TVM_TIRX_REGISTER_OP(OpName) TVM_TIR_REGISTER_OP(OpName)
-// Most common operators can be overloaded by argument type(PrimExpr).
-// So we put them under the root namespace.
-//
-// We put more developer oriented APIs -- MakeConst and is_const under tirx
-// as they are more specific to the tirx namespace.
+// Shared primitive construction and constants are declared in ir/prim/op.h.
/*!
* \brief Get the type of the expression under the unified type system.
@@ -101,22 +98,6 @@ TVM_DLL PrimExpr continue_loop(Span span = Span());
*/
TVM_DLL PrimExpr break_loop(Span span = Span());
-/*!
- * Query the maximum possible value of dtype.
- * \param dtype The primitive type.
- * \param span The location of this operation in the source.
- * \return the maximum possible value in this format.
- */
-TVM_DLL PrimExpr max_value(PrimType dtype, Span span = Span());
-
-/*!
- * Query the minimum possible value of dtype.
- * \param dtype The primitive type.
- * \param span The location of this operation in the source.
- * \return the minimum possible value in this format.
- */
-TVM_DLL PrimExpr min_value(PrimType dtype, Span span = Span());
-
/*!
* Get the value of infinity.
* \param dtype The primitive type.
@@ -125,16 +106,6 @@ TVM_DLL PrimExpr min_value(PrimType dtype, Span span = Span());
*/
TVM_DLL PrimExpr infinity(PrimType dtype, Span span = Span());
-/*!
- * \brief cast value to type.
- *
- * \param t the target type.
- * \param value The value
- * \param span The location of this operation in the source.
- * \return The result expression.
- * \note This function may return value if the type is the same.
- */
-TVM_DLL PrimExpr cast(PrimType t, PrimExpr value, Span span = Span());
/*!
* \brief perform reinterpret cast value to type.
*
@@ -145,267 +116,10 @@ TVM_DLL PrimExpr cast(PrimType t, PrimExpr value, Span span = Span());
* \note This function may return value if the type is the same.
*/
TVM_DLL PrimExpr reinterpret(PrimType t, PrimExpr value, Span span = Span());
+TVM_DLL PrimExpr reinterpret(DLDataType t, PrimExpr value, Span span = Span());
/*! \brief Perform a reinterpret cast involving an exact primitive or pointer type. */
TVM_DLL Expr reinterpret(Type target_ty, Expr value, Span span = Span());
-/*!
- * \brief add operator
- *
- * \param a left operand
- * \param b right operand
- * \param span The location of this operation in the source.
- * \return The result expression.
- * \note this function does eager constant folding for
- * index types(int32, int64) when possible.
- */
-TVM_DLL PrimExpr add(PrimExpr a, PrimExpr b, Span span = Span());
-/*!
- * \brief subtraction operator
- *
- * \param a left operand
- * \param b right operand
- * \param span The location of this operation in the source.
- * \return The result expression.
- * \note this function does eager constant folding for
- * index types(int32, int64) when possible.
- */
-TVM_DLL PrimExpr sub(PrimExpr a, PrimExpr b, Span span = Span());
-/*!
- * \brief negation.
- *
- * \param a input.
- * \param span The location of this operation in the source.
- * \return The result expression.
- * \note this function does eager constant folding for
- * index types(int32, int64) when possible.
- */
-TVM_DLL PrimExpr neg(PrimExpr a, Span span = Span());
-/*!
- * \brief multiplication operator
- *
- * \param a left operand
- * \param b right operand
- * \param span The location of this operation in the source.
- * \return The result expression.
- * \note this function does eager constant folding for
- * index types(int32, int64) when possible.
- */
-TVM_DLL PrimExpr mul(PrimExpr a, PrimExpr b, Span span = Span());
-/*!
- * \brief left shift operator
- *
- * \param a left operand
- * \param b right operand
- * \param span The location of this operation in the source.
- * \return The result expression.
- * \note this function does eager constant folding for
- * index types(int32, int64) when possible.
- */
-TVM_DLL PrimExpr left_shift(PrimExpr a, PrimExpr b, Span span = Span());
-/*!
- * \brief right shift operator
- *
- * \param a left operand
- * \param b right operand
- * \param span The location of this operation in the source.
- * \return The result expression.
- * \note this function does eager constant folding for
- * index types(int32, int64) when possible.
- */
-TVM_DLL PrimExpr right_shift(PrimExpr a, PrimExpr b, Span span = Span());
-/*!
- * \brief greater
- *
- * \param a left operand
- * \param b right operand
- * \param span The location of this operation in the source.
- * \return The result expression.
- * \note this function does eager constant folding for
- * index types(int32, int64) when possible.
- */
-TVM_DLL PrimExpr greater(PrimExpr a, PrimExpr b, Span span = Span());
-/*!
- * \brief greater_equal
- *
- * \param a left operand
- * \param b right operand
- * \param span The location of this operation in the source.
- * \return The result expression.
- * \note this function does eager constant folding for
- * index types(int32, int64) when possible.
- */
-TVM_DLL PrimExpr greater_equal(PrimExpr a, PrimExpr b, Span span = Span());
-/*!
- * \brief less
- *
- * \param a left operand
- * \param b right operand
- * \param span The location of this operation in the source.
- * \return The result expression.
- * \note this function does eager constant folding for
- * index types(int32, int64) when possible.
- */
-TVM_DLL PrimExpr less(PrimExpr a, PrimExpr b, Span span = Span());
-/*!
- * \brief less_equal
- *
- * \param a left operand
- * \param b right operand
- * \param span The location of this operation in the source.
- * \return The result expression.
- * \note this function does eager constant folding for
- * index types(int32, int64) when possible.
- */
-TVM_DLL PrimExpr less_equal(PrimExpr a, PrimExpr b, Span span = Span());
-/*!
- * \brief equal
- *
- * \param a left operand
- * \param b right operand
- * \param span The location of this operation in the source.
- * \return The result expression.
- * \note this function does eager constant folding for
- * index types(int32, int64) when possible.
- */
-TVM_DLL PrimExpr equal(PrimExpr a, PrimExpr b, Span span = Span());
-/*!
- * \brief not_equal
- *
- * \param a left operand
- * \param b right operand
- * \param span The location of this operation in the source.
- * \return The result expression.
- * \note this function does eager constant folding for
- * index types(int32, int64) when possible.
- */
-TVM_DLL PrimExpr not_equal(PrimExpr a, PrimExpr b, Span span = Span());
-/*!
- * \brief and
- *
- * \param a left operand
- * \param b right operand
- * \param span The location of this operation in the source.
- * \return The result expression.
- * \note This operator does eager constant folding.
- */
-TVM_DLL PrimExpr logical_and(PrimExpr a, PrimExpr b, Span span = Span());
-/*!
- * \brief or
- *
- * \param a left operand
- * \param b right operand
- * \param span The location of this operation in the source.
- * \return The result expression.
- * \note This operator does eager constant folding.
- */
-TVM_DLL PrimExpr logical_or(PrimExpr a, PrimExpr b, Span span = Span());
-/*!
- * \brief not
- *
- * \param a left operand
- * \param span The location of this operation in the source.
- * \return The result expression.
- * \note This operator does eager constant folding.
- */
-TVM_DLL PrimExpr logical_not(PrimExpr a, Span span = Span());
-/*!
- * \brief compute division in C semantics.
- *
- * a / b as in C/C++.
- *
- * When operands are integers, it directly corresponds to truncdiv.
- *
- * \param a left operand
- * \param b right operand
- * \param span The location of this operation in the source.
- * \return The result expression.
- * \note this function does eager constant folding for
- * index types(int32, int64) when possible.
- */
-TVM_DLL PrimExpr div(PrimExpr a, PrimExpr b, Span span = Span());
-/*!
- * \brief compute trunc(a / b)
- *
- * This is the default integer division behavior in C.
- *
- * \param a left operand
- * \param b right operand
- * \param span The location of this operation in the source.
- * \return The result expression.
- * \note this function does eager constant folding for
- * index types(int32, int64) when possible.
- */
-TVM_DLL PrimExpr truncdiv(PrimExpr a, PrimExpr b, Span span = Span());
-/*!
- * \brief compute the remainder of truncdiv
- *
- * This is the default integer division behavior in C.
- *
- * \param a left operand
- * \param b right operand
- * \param span The location of this operation in the source.
- * \return The result expression.
- * \note this function does eager constant folding for
- * index types(int32, int64) when possible.
- */
-TVM_DLL PrimExpr truncmod(PrimExpr a, PrimExpr b, Span span = Span());
-/*!
- * \brief compute floor(a / b) where a and b are non-negative.
- *
- * Use this function for index split calculation.
- *
- * This function might take advantage of the fact
- * that a and b are non-negative.
- *
- * \param a left operand
- * \param b right operand
- * \param span The location of this operation in the source.
- * \return The result expression.
- * \note this function does eager constant folding for
- * index types(int32, int64) when possible.
- */
-TVM_DLL PrimExpr indexdiv(PrimExpr a, PrimExpr b, Span span = Span());
-/*!
- * \brief compute ceil(a / b) where a and b are non-negative.
- *
- * Use this function for shape split calculation.
- *
- * This function might take advantage of the fact
- * that a and b are non-negative.
- *
- * \param a left operand
- * \param b right operand
- * \param span The location of this operation in the source.
- * \return The result expression.
- * \note this function does eager constant folding for
- * shape types(int32, int64) when possible.
- */
-TVM_DLL PrimExpr shapediv(PrimExpr a, PrimExpr b, Span span = Span());
-/*!
- * \brief compute the remainder floor(a / b) where a and b are non-negative.
- *
- * Use this function for index split calculation.
- * This function might take advantage of the fact
- * that a and b are non-negative.
- *
- * \param a left operand
- * \param b right operand
- * \param span The location of this operation in the source.
- * \return The result expression.
- * \note this function does eager constant folding for
- * index types(int32, int64) when possible.
- */
-TVM_DLL PrimExpr indexmod(PrimExpr a, PrimExpr b, Span span = Span());
-/*!
- * \brief compute floor(a / b)
- *
- * \param a left operand
- * \param b right operand
- * \param span The location of this operation in the source.
- * \return The result expression.
- * \note this function does eager constant folding for
- * index types(int32, int64) when possible.
- */
-TVM_DLL PrimExpr floordiv(PrimExpr a, PrimExpr b, Span span = Span());
+
/*!
* \brief Compute log(exp(a) + exp(b)).
*
@@ -415,114 +129,7 @@ TVM_DLL PrimExpr floordiv(PrimExpr a, PrimExpr b, Span span = Span());
* \return The result expression.
*/
TVM_DLL PrimExpr logaddexp(PrimExpr a, PrimExpr b, Span span = Span());
-/*!
- * \brief compute ceil(a / b)
- *
- * \param a left operand
- * \param b right operand
- * \param span The location of this operation in the source.
- * \return The result expression.
- * \note this function does eager constant folding for
- * index types(int32, int64) when possible.
- */
-TVM_DLL PrimExpr ceildiv(PrimExpr a, PrimExpr b, Span span = Span());
-/*!
- * \brief compute the remainder of floordiv
- *
- * \param a left operand
- * \param b right operand
- * \param span The location of this operation in the source.
- * \return The result expression.
- * \note this function does eager constant folding for
- * index types(int32, int64) when possible.
- */
-TVM_DLL PrimExpr floormod(PrimExpr a, PrimExpr b, Span span = Span());
-/*!
- * \brief take maximum of two values
- *
- * \param a left operand
- * \param b right operand
- * \param span The location of this operation in the source.
- * \return The result expression.
- * \note this function does eager constant folding for
- * index types(int32, int64) when possible.
- */
-TVM_DLL PrimExpr max(PrimExpr a, PrimExpr b, Span span = Span());
-/*!
- * \brief take minimum of two values
- *
- * \param a left operand
- * \param b right operand
- * \param span The location of this operation in the source.
- * \return The result expression.
- * \note this function does eager constant folding for
- * index types(int32, int64) when possible.
- */
-TVM_DLL PrimExpr min(PrimExpr a, PrimExpr b, Span span = Span());
-/*!
- * \brief take bitwise and of two values
- *
- * \param a left operand
- * \param b right operand
- * \param span The location of this operation in the source.
- * \return The result expression.
- * \note this function does eager constant folding for
- * index types(int32, int64) when possible.
- */
-TVM_DLL PrimExpr bitwise_and(PrimExpr a, PrimExpr b, Span span = Span());
-/*!
- * \brief take bitwise or of two values
- *
- * \param a left operand
- * \param b right operand
- * \param span The location of this operation in the source.
- * \return The result expression.
- * \note this function does eager constant folding for
- * index types(int32, int64) when possible.
- */
-TVM_DLL PrimExpr bitwise_or(PrimExpr a, PrimExpr b, Span span = Span());
-/*!
- * \brief take bitwise xor of two values
- *
- * \param a left operand
- * \param b right operand
- * \param span The location of this operation in the source.
- * \return The result expression.
- * \note this function does eager constant folding for
- * index types(int32, int64) when possible.
- */
-TVM_DLL PrimExpr bitwise_xor(PrimExpr a, PrimExpr b, Span span = Span());
-/*!
- * \brief take bitwise negation of two values
- *
- * \param a the input expression.
- * \param span The location of this operation in the source.
- * \return The result expression.
- * \note this function does eager constant folding for
- * index types(int32, int64) when possible.
- */
-TVM_DLL PrimExpr bitwise_neg(PrimExpr a, Span span = Span());
-/*!
- * \brief Conditional expression.
- *
- * \param cond The condition
- * \param true_value The value when results are true.
- * \param false_value The value when results are false.
- * \param span The location of this operation in the source.
- * \return The result expression.
- * \note this function does eager constant folding for
- * index types(int32, int64) when possible.
- */
-TVM_DLL PrimExpr if_then_else(PrimExpr cond, PrimExpr true_value, PrimExpr false_value,
- Span span = Span());
-/*!
- * \brief Mark condition as likely.
- * \param cond The condition
- * \param span The location of this operation in the source.
- * \return The marked expression.
- */
-TVM_DLL PrimExpr likely(PrimExpr cond, Span span = Span());
/*!
* \brief Calculate power(x, y)
* \param x The left operand.
@@ -594,28 +201,6 @@ TVM_DLL PrimExpr all(PrimExpr source, ffi::Array<tirx::IterVar> axis,
TVM_DLL PrimExpr any(PrimExpr source, ffi::Array<tirx::IterVar> axis,
ffi::Array<PrimExpr> init = {}, Span span = Span());
-/*!
- * \brief max of source expression over axis
- * \param source The source expression.
- * \param axis List of iteration variables that will be used for reduction.
- * \param init The value with which to initialize the output.
- * \param span The location of this operation in the source.
- * \return The result.
- */
-TVM_DLL PrimExpr max(PrimExpr source, ffi::Array<tirx::IterVar> axis,
- ffi::Array<PrimExpr> init = {}, Span span = Span());
-
-/*!
- * \brief max of source expression over axis
- * \param source The source expression.
- * \param axis List of iteration variables that will be used for reduction.
- * \param init The value with which to initialize the output.
- * \param span The location of this operation in the source.
- * \return The result.
- */
-TVM_DLL PrimExpr min(PrimExpr source, ffi::Array<tirx::IterVar> axis,
- ffi::Array<PrimExpr> init = {}, Span span = Span());
-
/*!
* \brief product of source expression over axis
* \param source The source expression.
@@ -635,14 +220,6 @@ TVM_DLL PrimExpr prod(PrimExpr source, ffi::Array<tirx::IterVar> axis,
*/
TVM_DLL PrimExpr floor(PrimExpr x, Span span = Span());
-/*!
- * \brief Calculate ceil(x)
- * \param x The input expression.
- * \param span The location of this operation in the source.
- * \return The result expression.
- */
-TVM_DLL PrimExpr ceil(PrimExpr x, Span span = Span());
-
/*!
* \brief Round x to the nearest integer, ties to even.
*
@@ -674,16 +251,6 @@ TVM_DLL PrimExpr nearbyint(PrimExpr x, Span span = Span());
*/
TVM_DLL PrimExpr trunc(PrimExpr x, Span span = Span());
-/*!
- * \brief Construct a large uint constant by its low 32 bits and high 32bits.
- * \param value_ty The final primitive type.
- * \param low The lower 32 bits.
- * \param high The higher 32 bits.
- * \param span The location of this operation in the source.
- * \return The constructed expression.
- */
-TVM_DLL PrimExpr LargeUIntImm(PrimType value_ty, int64_t low, int64_t high, Span span = Span());
-
/*!
* \brief Execute a multiplication between two Q-numbers x and y
* followed by a right shift s. The mathematical expression is:
@@ -758,7 +325,6 @@ TVM_DECLARE_INTRIN_UNARY(sigmoid);
TVM_DECLARE_INTRIN_UNARY(sqrt);
TVM_DECLARE_INTRIN_UNARY(rsqrt);
TVM_DECLARE_INTRIN_UNARY(log);
-TVM_DECLARE_INTRIN_UNARY(log2);
TVM_DECLARE_INTRIN_UNARY(log10);
TVM_DECLARE_INTRIN_UNARY(log1p);
TVM_DECLARE_INTRIN_UNARY(popcount);
@@ -787,8 +353,6 @@ TVM_DECLARE_INTRIN_BINARY(copysign);
TVM_DECLARE_INTRIN_BINARY(hypot);
TVM_DECLARE_INTRIN_BINARY(ldexp);
-namespace tirx {
-
/*!
* \brief Check if type is a pointer to a runtime element type.
* \param type The type to be checked.
@@ -805,25 +369,6 @@ inline bool IsPointerType(const Type& type, DLDataType element_type) {
return false;
}
-/*!
- * \brief Make a const value with certain data type.
- *
- * Prefer direct IntImm or FloatImm construction when dtype is known to be
- * scalar integer or floating point. This makes the compiled code more compact
- * and efficient. Keep MakeConst for generic overload cases where dtype can be
- * integer, floating point, or vector-valued and the caller needs its
- * scalar/vector dispatch.
- *
- * \param dtype The target type.
- * \param value The input value
- * \return the result expression.
- * \tparam ValueType The constant value type
- * \param span The location of this operation in the source.
- */
-template <typename ValueType,
- typename = typename std::enable_if<std::is_standard_layout<ValueType>::value &&
- std::is_trivial<ValueType>::value>::type>
-inline PrimExpr MakeConst(PrimType dtype, ValueType value, Span span = Span());
/*!
* \brief Make a constant opaque-pointer value.
* \param value The integer payload to reinterpret as a handle.
@@ -831,28 +376,6 @@ inline PrimExpr MakeConst(PrimType dtype, ValueType value, Span span = Span());
* \return The result expression.
*/
inline Expr ConstHandle(int64_t value, Span span = Span());
-/*!
- * \brief Get x as constant int expression.
- * \param x The expression
- * \return the address to the int expression,
- * return nullptr, if x is not IntImm.
- */
-inline const int64_t* as_const_int(const PrimExpr& x) {
- if (!x.defined()) return nullptr;
- if (const IntImmNode* op = x.as<IntImmNode>()) {
- return &(op->value);
- }
-
- return nullptr;
-}
-
-/*!
- * \brief Check whether x is a constant integer expression.
- * \param x The input argument
- * \param value the value to be compared against.
- * \return whether x is constant expression.
- */
-inline bool is_const_int(const PrimExpr& x, int64_t value);
/*!
* \brief Check whether stmt is nop.
@@ -861,36 +384,6 @@ inline bool is_const_int(const PrimExpr& x, int64_t value);
*/
inline bool is_no_op(const tirx::Stmt& stmt);
-/*!
- * \brief Check whether x is a constant integer 1
- * \param x The input argument.
- * \note This only return true for integer types.
- * \return whether x is constant 1
- */
-inline bool is_one(const PrimExpr& x) { return is_const_int(x, 1); }
-
-/*!
- * \brief Check whether x is a constant integer 0
- * \param x The input argument
- * \return whether x is constant 0
- * \note This only return true for integer types.
- */
-inline bool is_zero(const PrimExpr& x) { return is_const_int(x, 0); }
-
-/*!
- * \brief Check whether x is an integer constant.
- * \note This only return true for integer types.
- * \return whether x is constant
- */
-inline bool is_const_int(const PrimExpr& x);
-
-/*!
- * \brief Check whether x is an integer/float constant.
- * \note This only return true for integer types.
- * \return whether x is constant
- */
-inline bool is_const_number(const PrimExpr& x);
-
/*!
* \brief Left fold.
* \param freduce The reduction function.
@@ -909,45 +402,6 @@ inline PrimExpr foldl(FReduce freduce, PrimExpr init_value, const ffi::Array<Pri
return init_value;
}
-/*!
- * \brief Check whether x is a constant power of two
- * If x is power of two, write the power to the shift.
- *
- * \param x The input expression.
- * \param shift The output shift if x is power of two.
- * \return whether x is constant power of two
- */
-TVM_DLL bool is_const_power_of_two_integer(const PrimExpr& x, int* shift);
-
-// Implementation details after this
-inline bool is_const_int(const PrimExpr& x) { return as_const_int(x); }
-
-inline bool is_const_number(const PrimExpr& x) {
- if (x.as<IntImmNode>()) {
- return true;
- } else if (x.as<FloatImmNode>()) {
- return true;
- } else if (const auto* op = x.as<prim::BroadcastNode>()) {
- return (op->value->IsInstance<IntImmNode>() || op->value->IsInstance<FloatImmNode>());
- }
- return false;
-}
-
-inline bool is_positive_const(const PrimExpr& a) {
- const int64_t* as_int = as_const_int(a);
- return as_int && (*as_int > 0);
-}
-
-inline bool is_negative_const(const PrimExpr& a) {
- const int64_t* as_int = as_const_int(a);
- return as_int && (*as_int < 0);
-}
-
-inline bool is_const_int(const PrimExpr& x, int64_t value) {
- const int64_t* as_int = as_const_int(x);
- return as_int && (*as_int == value);
-}
-
inline bool is_no_op(const tirx::Stmt& stmt) {
if (!stmt.defined()) return true;
if (const auto* op = stmt.as<tirx::EvaluateNode>()) {
@@ -960,203 +414,35 @@ inline bool is_no_op(const tirx::Stmt& stmt) {
return false;
}
-template <typename ValueType>
-inline PrimExpr MakeConstScalar(PrimType dtype, ValueType value, Span span = Span()) {
- DLDataTypeCode code = dtype.code();
- if (code == DLDataTypeCode::kDLInt || code == DLDataTypeCode::kDLBool) {
- return IntImm(dtype, static_cast<int64_t>(value), span);
- }
- if (code == DLDataTypeCode::kDLUInt) {
- // Use IntImm if it is a small integer
- uint64_t uval = static_cast<uint64_t>(value);
- if (value < static_cast<ValueType>(0)) {
- TVM_FFI_THROW(InternalError) << "cannot make uint from negative value " << value;
- } else if (uval <= static_cast<uint64_t>(std::numeric_limits<int64_t>::max())) {
- return IntImm(dtype, static_cast<int64_t>(value), span);
- } else {
- uint64_t mask = (static_cast<uint64_t>(1) << 32U) - 1U;
- uint64_t low = uval & mask;
- uint64_t high = uval >> 32U;
- return LargeUIntImm(dtype, static_cast<int64_t>(low), static_cast<int64_t>(high), span);
- }
- }
- if (dtype.MatchesCode(DLDataTypeCode::kDLFloat, DLDataTypeCode::kDLFloat8_e3m4,
- DLDataTypeCode::kDLFloat8_e4m3, DLDataTypeCode::kDLFloat8_e4m3b11fnuz,
- DLDataTypeCode::kDLFloat8_e4m3fn, DLDataTypeCode::kDLFloat8_e4m3fnuz,
- DLDataTypeCode::kDLFloat8_e5m2, DLDataTypeCode::kDLFloat8_e5m2fnuz,
- DLDataTypeCode::kDLFloat8_e8m0fnu, DLDataTypeCode::kDLFloat6_e2m3fn,
- DLDataTypeCode::kDLFloat6_e3m2fn, DLDataTypeCode::kDLFloat4_e2m1fn) ||
- dtype.MatchesElementType(DLDataTypeCode::kDLBfloat, 16)) {
- return FloatImm(dtype, static_cast<double>(value), span);
- }
- TVM_FFI_THROW(InternalError) << "cannot make const for type " << dtype;
- throw;
-}
-
-template <>
-inline PrimExpr MakeConstScalar(PrimType dtype, bool value, Span span) {
- return MakeConstScalar(dtype, static_cast<int>(value), span);
-}
-
-template <typename ValueType, typename>
-inline PrimExpr MakeConst(PrimType dtype, ValueType value, Span span) {
- if (!dtype.IsScalableVector() && !dtype.IsFixedLengthVector()) {
- return MakeConstScalar(dtype, value, span);
- }
- PrimType elem_ty = dtype.WithLanes(1);
- if (dtype.IsFixedLengthVector()) {
- return prim::Broadcast(MakeConstScalar(elem_ty, value, span), dtype.lanes(), span);
- }
- PrimExpr lanes =
- prim::Mul(Call(PrimType::Int(32), prim::builtin::vscale(), {}).as_or_throw<PrimExpr>(),
- dtype.VScaleFactor());
- return prim::Broadcast(MakeConstScalar(elem_ty, value, span), lanes, span);
-}
-
inline Expr ConstHandle(int64_t value, Span span) {
return reinterpret(PointerType::VoidPointerTy(), IntImm(PrimType::UInt(64), value, span), span);
}
-} // namespace tirx
-
-// additional const expression overloading
-#define TVM_DEFINE_ASSIGN_OP_OVERLOAD(Name, OpFunc) \
- inline PrimExpr Name(PrimExpr& a, PrimExpr b) { \
- a = OpFunc(a, b); \
- return a; \
- }
-
-#define TVM_DEFINE_BINOP_CONST_VAL_OVERLOAD(Name) \
- inline PrimExpr Name(const PrimExpr& a, float b) { return Name(a, PrimExpr(b)); } \
- inline PrimExpr Name(float a, const PrimExpr& b) { return Name(PrimExpr(a), b); } \
- inline PrimExpr Name(int a, const PrimExpr& b) { return Name(tirx::MakeConst(b.ty(), a), b); } \
- inline PrimExpr Name(const PrimExpr& a, int b) { return Name(a, tirx::MakeConst(a.ty(), b)); } \
- inline PrimExpr Name(const PrimExpr& a, double b) { \
- return Name(a, FloatImm(PrimType::Float(64), b)); \
- }
-
-#define TVM_DEFINE_BINOP_CONST_VAL_OVERLOAD_SPANNED(Name) \
- inline PrimExpr Name(const PrimExpr& a, float b, Span span = Span()) { \
- return Name(a, PrimExpr(b), span); \
- } \
- inline PrimExpr Name(float a, const PrimExpr& b, Span span = Span()) { \
- return Name(PrimExpr(a), b, span); \
- } \
- inline PrimExpr Name(int a, const PrimExpr& b, Span span = Span()) { \
- return Name(tirx::MakeConst(b.ty(), a), b, span); \
- } \
- inline PrimExpr Name(const PrimExpr& a, int b, Span span = Span()) { \
- return Name(a, tirx::MakeConst(a.ty(), b), span); \
- } \
- inline PrimExpr Name(const PrimExpr& a, double b, Span span = Span()) { \
- return Name(a, FloatImm(PrimType::Float(64), b), span); \
- }
-
-#define TVM_DEFINE_LOGICAL_OP_CONST_VAL_OVERLOAD(Name) \
- inline PrimExpr Name(const PrimExpr& a, bool b) { return Name(a, PrimExpr(b)); } \
- inline PrimExpr Name(bool a, const PrimExpr& b) { return Name(PrimExpr(a), b); }
-
-#define TVM_DEFINE_LOGICAL_OP_CONST_VAL_OVERLOAD_SPANNED(Name) \
- inline PrimExpr Name(const PrimExpr& a, bool b, Span span = Span()) { \
- return Name(a, PrimExpr(b), span); \
- } \
- inline PrimExpr Name(bool a, const PrimExpr& b, Span span = Span()) { \
- return Name(PrimExpr(a), b, span); \
- }
-
-#define TVM_DEFINE_INT_OP_CONST_VAL_OVERLOAD(Name) \
- inline PrimExpr Name(const PrimExpr& a, int b) { return Name(a, tirx::MakeConst(a.ty(), b)); } \
- inline PrimExpr Name(int a, const PrimExpr& b) { return Name(tirx::MakeConst(b.ty(), a), b); }
-
-#define TVM_DEFINE_INT_OP_CONST_VAL_OVERLOAD_SPANNED(Name) \
- inline PrimExpr Name(const PrimExpr& a, int b, Span span = Span()) { \
- return Name(a, tirx::MakeConst(a.ty(), b), span); \
- } \
- inline PrimExpr Name(int a, const PrimExpr& b, Span span = Span()) { \
- return Name(tirx::MakeConst(b.ty(), a), b, span); \
- }
-
-TVM_DEFINE_ASSIGN_OP_OVERLOAD(operator+=, operator+);
-TVM_DEFINE_ASSIGN_OP_OVERLOAD(operator-=, operator-);
-TVM_DEFINE_ASSIGN_OP_OVERLOAD(operator*=, operator*);
-TVM_DEFINE_BINOP_CONST_VAL_OVERLOAD(operator+);
-TVM_DEFINE_BINOP_CONST_VAL_OVERLOAD(operator-);
-TVM_DEFINE_BINOP_CONST_VAL_OVERLOAD(operator*);
-TVM_DEFINE_BINOP_CONST_VAL_OVERLOAD(operator>); // NOLINT(*)
-TVM_DEFINE_BINOP_CONST_VAL_OVERLOAD(operator>=);
-TVM_DEFINE_BINOP_CONST_VAL_OVERLOAD(operator<); // NOLINT(*)
-TVM_DEFINE_BINOP_CONST_VAL_OVERLOAD(operator<=);
-TVM_DEFINE_BINOP_CONST_VAL_OVERLOAD_SPANNED(max);
-TVM_DEFINE_BINOP_CONST_VAL_OVERLOAD_SPANNED(min);
-TVM_DEFINE_BINOP_CONST_VAL_OVERLOAD_SPANNED(div);
-TVM_DEFINE_BINOP_CONST_VAL_OVERLOAD_SPANNED(add);
-TVM_DEFINE_BINOP_CONST_VAL_OVERLOAD_SPANNED(sub);
-TVM_DEFINE_BINOP_CONST_VAL_OVERLOAD_SPANNED(mul);
-TVM_DEFINE_BINOP_CONST_VAL_OVERLOAD_SPANNED(greater);
-TVM_DEFINE_BINOP_CONST_VAL_OVERLOAD_SPANNED(greater_equal);
-TVM_DEFINE_BINOP_CONST_VAL_OVERLOAD_SPANNED(less);
-TVM_DEFINE_BINOP_CONST_VAL_OVERLOAD_SPANNED(less_equal);
-// integer related ops
-TVM_DEFINE_INT_OP_CONST_VAL_OVERLOAD_SPANNED(indexdiv);
-TVM_DEFINE_INT_OP_CONST_VAL_OVERLOAD_SPANNED(indexmod);
-TVM_DEFINE_INT_OP_CONST_VAL_OVERLOAD_SPANNED(truncdiv);
-TVM_DEFINE_INT_OP_CONST_VAL_OVERLOAD_SPANNED(truncmod);
-TVM_DEFINE_INT_OP_CONST_VAL_OVERLOAD_SPANNED(floordiv);
TVM_DEFINE_INT_OP_CONST_VAL_OVERLOAD_SPANNED(logaddexp);
-TVM_DEFINE_INT_OP_CONST_VAL_OVERLOAD_SPANNED(floormod);
-TVM_DEFINE_INT_OP_CONST_VAL_OVERLOAD_SPANNED(right_shift); // NOLINT(*)
-TVM_DEFINE_INT_OP_CONST_VAL_OVERLOAD_SPANNED(left_shift); // NOLINT(*)
-TVM_DEFINE_INT_OP_CONST_VAL_OVERLOAD_SPANNED(bitwise_and);
-TVM_DEFINE_INT_OP_CONST_VAL_OVERLOAD_SPANNED(bitwise_or);
-TVM_DEFINE_INT_OP_CONST_VAL_OVERLOAD_SPANNED(bitwise_xor);
-TVM_DEFINE_INT_OP_CONST_VAL_OVERLOAD(operator>>); // NOLINT(*)
-TVM_DEFINE_INT_OP_CONST_VAL_OVERLOAD(operator<<); // NOLINT(*)
-TVM_DEFINE_INT_OP_CONST_VAL_OVERLOAD(operator&);
-TVM_DEFINE_INT_OP_CONST_VAL_OVERLOAD(operator|);
-TVM_DEFINE_INT_OP_CONST_VAL_OVERLOAD(operator^);
-// logical ops
-TVM_DEFINE_LOGICAL_OP_CONST_VAL_OVERLOAD(operator&&);
-TVM_DEFINE_LOGICAL_OP_CONST_VAL_OVERLOAD(operator||);
-TVM_DEFINE_LOGICAL_OP_CONST_VAL_OVERLOAD_SPANNED(logical_and);
-TVM_DEFINE_LOGICAL_OP_CONST_VAL_OVERLOAD_SPANNED(logical_or);
+} // namespace tvm::prim
+namespace tvm {
/*!
- * \brief Helper function to raise a compiler error about division ambiguity.
- * \note The call to this function will always results in a compiler error.
- * \tparam TA Any class type.
+ * \brief max of source expression over axis
+ * \param source The source expression.
+ * \param axis List of iteration variables that will be used for reduction.
+ * \param init The value with which to initialize the output.
+ * \param span The location of this operation in the source.
+ * \return The result.
*/
-template <typename TA>
-inline void DivAmbiguityError(const TA& a) {
- constexpr bool div_ambiguity = !std::is_class<TA>::value;
- static_assert(div_ambiguity,
- "TVM supports multiple types of integer divisions, "
- "please call div, indexdiv/indexmod, "
- "floordiv/floormod or truncdiv/truncmod directly "
- "to avoid ambiguity in the code. "
- "Checkout these functions in tirx/op.h.");
-}
-
-// The following code are not intended to be used in the codebase.
-// Instead, they generate clear compiler errors that ask developers
-// to use the specific division function.
-// The second template argument is necessary to make sure the
-// code compiles lazily by the compiler during invocation.
-template <typename TB>
-inline PrimExpr operator/(const PrimExpr& a, const TB& b) {
- DivAmbiguityError(a);
- return a;
-}
+TVM_DLL PrimExpr max(PrimExpr source, ffi::Array<tirx::IterVar> axis,
+ ffi::Array<PrimExpr> init = {}, Span span = Span());
-template <typename TB>
-inline PrimExpr operator/=(const PrimExpr& a, const TB& b) {
- DivAmbiguityError(a);
- return a;
-}
+/*!
+ * \brief max of source expression over axis
+ * \param source The source expression.
+ * \param axis List of iteration variables that will be used for reduction.
+ * \param init The value with which to initialize the output.
+ * \param span The location of this operation in the source.
+ * \return The result.
+ */
+TVM_DLL PrimExpr min(PrimExpr source, ffi::Array<tirx::IterVar> axis,
+ ffi::Array<PrimExpr> init = {}, Span span = Span());
-template <typename TB>
-inline PrimExpr operator%(const PrimExpr& a, const TB& b) {
- DivAmbiguityError(a);
- return a;
-}
} // namespace tvm
#endif // TVM_TIR_OP_H_
diff --git a/include/tvm/tirx/script/builder/ir.h b/include/tvm/tirx/script/builder/ir.h
index affdbfa690..57a8d1098f 100644
--- a/include/tvm/tirx/script/builder/ir.h
+++ b/include/tvm/tirx/script/builder/ir.h
@@ -529,7 +529,7 @@ inline Var TensorMap() { return tvm::tirx::Var("", PointerType(TensorMapType()))
#define TVM_TIRX_IR_BUILDER_DEF_DTYPE_CAST(FuncName, DType) \
inline PrimExpr FuncName(ffi::Optional<PrimExpr> expr = std::nullopt) { \
PrimType dtype = DType; \
- return expr.has_value() ? tvm::cast(dtype, expr.value()) \
+ return expr.has_value() ? tvm::prim::cast(dtype, expr.value()) \
: tvm::tirx::Var("", dtype).as_or_throw<PrimExpr>(); \
}
diff --git a/include/tvm/topi/broadcast.h b/include/tvm/topi/broadcast.h
index a114bf3287..6120046dac 100644
--- a/include/tvm/topi/broadcast.h
+++ b/include/tvm/topi/broadcast.h
@@ -33,7 +33,6 @@
namespace tvm {
namespace topi {
-
/*!
* \brief Creates an operation that broadcasts a tensor into a compatible
* shape according to numpy's rules
@@ -242,7 +241,7 @@ TOPI_DEFINE_BCAST_OP(divide, { return div(a, b); });
/*!
* \fn floor divide
- * \brief Compute floor(A / B) with auto-broadcasting.
+ * \brief Compute prim::floor(A / B) with auto-broadcasting.
*
* \param A The first tensor, or Expr
* \param B The second tensor, or Expr
@@ -256,7 +255,7 @@ TOPI_DEFINE_BCAST_OP(floor_divide, {
if (a_ty.MatchesCode(DLDataTypeCode::kDLInt, DLDataTypeCode::kDLUInt)) {
return floordiv(a, b);
} else {
- return floor(div(a, b));
+ return prim::floor(div(a, b));
}
});
@@ -274,11 +273,11 @@ TOPI_DEFINE_BCAST_OP(floor_divide, {
*
* \return The computed log-sum-exp result.
*/
-TOPI_DEFINE_BCAST_OP(log_add_exp, { return logaddexp(a, b); });
+TOPI_DEFINE_BCAST_OP(log_add_exp, { return prim::logaddexp(a, b); });
/*!
* \fn trunc divide
- * \brief Compute trunc(A / B) with auto-broadcasting.
+ * \brief Compute prim::trunc(A / B) with auto-broadcasting.
*
* \param A The first tensor, or Expr
* \param B The second tensor, or Expr
@@ -292,7 +291,7 @@ TOPI_DEFINE_BCAST_OP(trunc_divide, {
if (a_ty.MatchesCode(DLDataTypeCode::kDLInt, DLDataTypeCode::kDLUInt)) {
return truncdiv(a, b);
} else {
- return trunc(div(a, b));
+ return prim::trunc(div(a, b));
}
});
@@ -386,7 +385,7 @@ TOPI_DEFINE_BCAST_OP(minimum, { return tvm::min(a, b); });
*
* \return The result.
*/
-TOPI_DEFINE_BCAST_OP(power, { return tvm::pow(a, b); });
+TOPI_DEFINE_BCAST_OP(power, { return tvm::prim::pow(a, b); });
/*!
* \fn atan2
@@ -399,7 +398,7 @@ TOPI_DEFINE_BCAST_OP(power, { return tvm::pow(a, b); });
*
* \return The result.
*/
-TOPI_DEFINE_BCAST_OP(atan2, { return tvm::atan2(a, b); });
+TOPI_DEFINE_BCAST_OP(atan2, { return tvm::prim::atan2(a, b); });
/*!
* \fn left_shift
diff --git a/include/tvm/topi/detail/broadcast.h b/include/tvm/topi/detail/broadcast.h
index 2d6b873bcd..a3bf9fa7d2 100644
--- a/include/tvm/topi/detail/broadcast.h
+++ b/include/tvm/topi/detail/broadcast.h
@@ -57,7 +57,7 @@ inline BroadcastHelper BroadcastShape(const tvm::ffi::Array<tvm::PrimExpr>& shap
int i;
auto cast_if_needed = [](PrimType to_type, PrimExpr expr) {
- return to_type == expr.ty() ? expr : cast(to_type, expr);
+ return to_type == expr.ty() ? expr : prim::cast(to_type, expr);
};
for (i = 1; i <= std::min(s1_size, s2_size); ++i) {
diff --git a/include/tvm/topi/detail/extern.h b/include/tvm/topi/detail/extern.h
index 578de8e4c6..14b16ace93 100644
--- a/include/tvm/topi/detail/extern.h
+++ b/include/tvm/topi/detail/extern.h
@@ -100,6 +100,7 @@ inline ffi::Array<Tensor> make_extern(const ffi::Array<ffi::Array<PrimExpr>>& ou
* \return An expression representing the pack operation
*/
inline Expr pack_buffer(BufferVar buf) {
+ using namespace tvm::prim;
TVM_FFI_ICHECK_GT(buf->shape.size(), 0) << "buf shape must have at least one element";
Expr shape =
Call(PointerType(PrimType::Int(64)), tvm::tirx::builtin::tvm_stack_make_shape(), buf->shape);
diff --git a/include/tvm/topi/detail/strided_slice.h b/include/tvm/topi/detail/strided_slice.h
index d5a9b2f854..44df556be5 100644
--- a/include/tvm/topi/detail/strided_slice.h
+++ b/include/tvm/topi/detail/strided_slice.h
@@ -93,6 +93,7 @@ inline ffi::Array<PrimExpr> StridedSliceCanonicalizeBegin(const ffi::Array<PrimE
const ffi::Array<int64_t>& axes,
PrimType dtype,
std::string slice_mode = "end") {
+ using namespace tvm::prim;
ffi::Array<PrimExpr> begin_expr;
for (size_t i = 0; i < axes.size(); ++i) {
int64_t ax = axes[i];
@@ -140,7 +141,7 @@ inline ffi::Array<PrimExpr> StridedSliceOutputShape(
static_cast<int>((interval + std::abs(strides[i]) - 1) / std::abs(strides[i]));
TVM_FFI_ICHECK(strides[i] < 0 ? (end_i <= begin_i) : (begin_i <= end_i))
<< ": Input [Begin=" << begin[i] << ", End=" << end[i] << "] is invalid for axis=" << i;
- out_shape.Set(ax, cast(out_shape[i].ty(), PrimExpr(slice_size)));
+ out_shape.Set(ax, prim::cast(out_shape[i].ty(), PrimExpr(slice_size)));
} else {
out_shape.Set(ax, tvm::PrimVar("dim", out_shape[i].ty()));
}
diff --git a/include/tvm/topi/detail/tensor_utils.h b/include/tvm/topi/detail/tensor_utils.h
index 82649cd0b3..1512464ed8 100644
--- a/include/tvm/topi/detail/tensor_utils.h
+++ b/include/tvm/topi/detail/tensor_utils.h
@@ -70,10 +70,10 @@ inline PrimExpr bilinear_sample_nchw(const Tensor& input, const ffi::Array<PrimE
auto in_y = indices[2];
auto in_x = indices[3];
- auto y_low = tvm::cast(PrimType::Int(32), tvm::floor(in_y));
+ auto y_low = tvm::prim::cast(PrimType::Int(32), tvm::prim::floor(in_y));
auto y_high = y_low + 1;
- auto x_low = tvm::cast(PrimType::Int(32), tvm::floor(in_x));
+ auto x_low = tvm::prim::cast(PrimType::Int(32), tvm::prim::floor(in_x));
auto x_high = x_low + 1;
auto wy_h = in_y - y_low;
@@ -114,10 +114,10 @@ inline PrimExpr bilinear_sample_nhwc(const Tensor& input, const ffi::Array<PrimE
auto in_y = indices[1];
auto in_x = indices[2];
- auto y_low = tvm::cast(PrimType::Int(32), tvm::floor(in_y));
+ auto y_low = tvm::prim::cast(PrimType::Int(32), tvm::prim::floor(in_y));
auto y_high = y_low + 1;
- auto x_low = tvm::cast(PrimType::Int(32), tvm::floor(in_x));
+ auto x_low = tvm::prim::cast(PrimType::Int(32), tvm::prim::floor(in_x));
auto x_high = x_low + 1;
auto wy_h = in_y - y_low;
diff --git a/include/tvm/topi/elemwise.h b/include/tvm/topi/elemwise.h
index 002b4f1e9e..457fc4e3da 100644
--- a/include/tvm/topi/elemwise.h
+++ b/include/tvm/topi/elemwise.h
@@ -37,50 +37,50 @@
namespace tvm {
namespace topi {
-
using namespace tvm::te;
// Unary intrinsic operators
-#define TOPI_DECLARE_UNARY_OP(OpName) \
- inline Tensor OpName(const Tensor& x, std::string name = "T_" #OpName, \
- std::string tag = kElementWise) { \
- return compute( \
- x->shape, [&](const ffi::Array<PrimVar>& i) { return ::tvm::OpName(x(i)); }, name, tag); \
+#define TOPI_DECLARE_UNARY_OP(OpName, ExprFunc) \
+ inline Tensor OpName(const Tensor& x, std::string name = "T_" #OpName, \
+ std::string tag = kElementWise) { \
+ return compute( \
+ x->shape, [&](const ffi::Array<PrimVar>& i) { return ExprFunc(x(i)); }, name, tag); \
}
-TOPI_DECLARE_UNARY_OP(exp);
-TOPI_DECLARE_UNARY_OP(erf);
-TOPI_DECLARE_UNARY_OP(sigmoid);
-TOPI_DECLARE_UNARY_OP(sqrt);
-TOPI_DECLARE_UNARY_OP(log);
-TOPI_DECLARE_UNARY_OP(log2);
-TOPI_DECLARE_UNARY_OP(log10);
-TOPI_DECLARE_UNARY_OP(floor);
-TOPI_DECLARE_UNARY_OP(ceil);
-TOPI_DECLARE_UNARY_OP(round);
-TOPI_DECLARE_UNARY_OP(trunc);
-TOPI_DECLARE_UNARY_OP(abs);
-TOPI_DECLARE_UNARY_OP(cos);
-TOPI_DECLARE_UNARY_OP(cosh);
-TOPI_DECLARE_UNARY_OP(tan);
-TOPI_DECLARE_UNARY_OP(sin);
-TOPI_DECLARE_UNARY_OP(sinh);
-TOPI_DECLARE_UNARY_OP(acos);
-TOPI_DECLARE_UNARY_OP(acosh);
-TOPI_DECLARE_UNARY_OP(asin);
-TOPI_DECLARE_UNARY_OP(asinh);
-TOPI_DECLARE_UNARY_OP(atan);
-TOPI_DECLARE_UNARY_OP(atanh);
-TOPI_DECLARE_UNARY_OP(isnan);
-TOPI_DECLARE_UNARY_OP(tanh);
-TOPI_DECLARE_UNARY_OP(isfinite);
-TOPI_DECLARE_UNARY_OP(isinf);
+TOPI_DECLARE_UNARY_OP(exp, ::tvm::prim::exp);
+TOPI_DECLARE_UNARY_OP(erf, ::tvm::prim::erf);
+TOPI_DECLARE_UNARY_OP(sigmoid, ::tvm::prim::sigmoid);
+TOPI_DECLARE_UNARY_OP(sqrt, ::tvm::prim::sqrt);
+TOPI_DECLARE_UNARY_OP(log, ::tvm::prim::log);
+TOPI_DECLARE_UNARY_OP(log2, ::tvm::log2);
+TOPI_DECLARE_UNARY_OP(log10, ::tvm::prim::log10);
+TOPI_DECLARE_UNARY_OP(floor, ::tvm::prim::floor);
+TOPI_DECLARE_UNARY_OP(ceil, ::tvm::ceil);
+TOPI_DECLARE_UNARY_OP(round, ::tvm::prim::round);
+TOPI_DECLARE_UNARY_OP(trunc, ::tvm::prim::trunc);
+TOPI_DECLARE_UNARY_OP(abs, ::tvm::prim::abs);
+TOPI_DECLARE_UNARY_OP(cos, ::tvm::prim::cos);
+TOPI_DECLARE_UNARY_OP(cosh, ::tvm::prim::cosh);
+TOPI_DECLARE_UNARY_OP(tan, ::tvm::prim::tan);
+TOPI_DECLARE_UNARY_OP(sin, ::tvm::prim::sin);
+TOPI_DECLARE_UNARY_OP(sinh, ::tvm::prim::sinh);
+TOPI_DECLARE_UNARY_OP(acos, ::tvm::prim::acos);
+TOPI_DECLARE_UNARY_OP(acosh, ::tvm::prim::acosh);
+TOPI_DECLARE_UNARY_OP(asin, ::tvm::prim::asin);
+TOPI_DECLARE_UNARY_OP(asinh, ::tvm::prim::asinh);
+TOPI_DECLARE_UNARY_OP(atan, ::tvm::prim::atan);
+TOPI_DECLARE_UNARY_OP(atanh, ::tvm::prim::atanh);
+TOPI_DECLARE_UNARY_OP(isnan, ::tvm::prim::isnan);
+TOPI_DECLARE_UNARY_OP(tanh, ::tvm::prim::tanh);
+TOPI_DECLARE_UNARY_OP(isfinite, ::tvm::prim::isfinite);
+TOPI_DECLARE_UNARY_OP(isinf, ::tvm::prim::isinf);
/*!
* \brief Fast_tanh_float implementation from Eigen
* https://github.com/eigenteam/eigen-git-mirror/blob/master/Eigen/src/Core/MathFunctionsImpl.h#L26
*/
inline Tensor fast_tanh_float(const Tensor& in, std::string name, std::string tag) {
+ using namespace tvm::prim;
// Clamp the inputs to the range [-9, 9] since anything outside
// this range is +/-1.0f in single-precision.
PrimType input_type = in->GetDataType();
@@ -138,7 +138,7 @@ inline Tensor fast_tanh(const Tensor& x, std::string name = "T_fast_tanh",
} else {
// fallback to default implementation
return compute(
- x->shape, [&](const ffi::Array<PrimVar>& i) { return ::tvm::tanh(x(i)); }, name, tag);
+ x->shape, [&](const ffi::Array<PrimVar>& i) { return ::tvm::prim::tanh(x(i)); }, name, tag);
}
}
@@ -208,6 +208,7 @@ inline Tensor bitwise_not(const Tensor& x, std::string name = "T_bitwise_not",
* \return A Tensor whose op member is the sign
*/
inline Tensor sign(const Tensor& x, std::string name = "T_sign", std::string tag = kElementWise) {
+ using namespace tvm::prim;
return compute(
x->shape,
[&](const ffi::Array<PrimVar>& i) {
@@ -232,11 +233,12 @@ inline Tensor sign(const Tensor& x, std::string name = "T_sign", std::string tag
* \return A Tensor whose op member is the rsqrt operation
*/
inline Tensor rsqrt(const Tensor& x, std::string name = "tensor", std::string tag = kElementWise) {
+ using namespace tvm::prim;
return compute(
x->shape,
[&](const ffi::Array<PrimVar>& i) {
PrimExpr one = MakeConst(x->GetDataType(), 1);
- return one / tvm::sqrt(x(i));
+ return one / tvm::prim::sqrt(x(i));
},
name, tag);
}
@@ -259,8 +261,8 @@ inline Tensor clip(const Tensor& x, const PrimExpr& a_min, const PrimExpr& a_max
x->shape,
[&](const ffi::Array<PrimVar>& i) {
PrimType x_type(x->GetDataType());
- auto min_val = tvm::cast(x_type, a_min);
- auto max_val = tvm::cast(x_type, a_max);
+ auto min_val = tvm::prim::cast(x_type, a_min);
+ auto max_val = tvm::prim::cast(x_type, a_max);
return tvm::max(tvm::min(x(i), max_val), min_val); // NOLINT(*)
},
name, tag);
@@ -282,6 +284,7 @@ inline Tensor cast(const Tensor& x, PrimType type, std::string name, std::string
inline Tensor cast(const Tensor& x, DLDataType type, std::string name = "T_cast",
std::string tag = kElementWise) {
+ using namespace tvm::prim;
return cast(x, PrimType(type), std::move(name), std::move(tag));
}
@@ -300,7 +303,7 @@ inline Tensor cast(const Tensor& x, PrimType type, std::string name = "T_cast",
}
}
- return tvm::cast(type, x(i));
+ return tvm::prim::cast(type, x(i));
},
name, tag);
}
@@ -319,13 +322,15 @@ inline Tensor reinterpret(const Tensor& x, PrimType type, std::string name, std:
inline Tensor reinterpret(const Tensor& x, DLDataType type, std::string name = "tensor",
std::string tag = kElementWise) {
+ using namespace tvm::prim;
return reinterpret(x, PrimType(type), std::move(name), std::move(tag));
}
inline Tensor reinterpret(const Tensor& x, PrimType type, std::string name = "tensor",
std::string tag = kElementWise) {
return compute(
- x->shape, [&](const ffi::Array<PrimVar>& i) { return reinterpret(type, x(i)); }, name, tag);
+ x->shape, [&](const ffi::Array<PrimVar>& i) { return prim::reinterpret(type, x(i)); }, name,
+ tag);
}
/*!
@@ -373,7 +378,7 @@ inline Tensor full(const ffi::Array<PrimExpr>& shape, DLDataType dtype, const Pr
inline Tensor full(const ffi::Array<PrimExpr>& shape, PrimType dtype, const PrimExpr fill_value,
std::string name = "T_full", std::string tag = kElementWise) {
- PrimExpr ev = cast(dtype, fill_value);
+ PrimExpr ev = prim::cast(dtype, fill_value);
if (!ev.defined()) {
LOG(ERROR) << "Can't cast fill_value to " << dtype;
}
@@ -393,7 +398,7 @@ inline Tensor full(const ffi::Array<PrimExpr>& shape, PrimType dtype, const Prim
*/
inline Tensor full_like(const Tensor& x, const PrimExpr fill_value,
std::string name = "T_full_like", std::string tag = kElementWise) {
- PrimExpr ev = cast(x->GetDataType(), fill_value);
+ PrimExpr ev = prim::cast(x->GetDataType(), fill_value);
return compute(x->shape, [&](const ffi::Array<PrimVar>& i) { return ev; }, name, tag);
}
@@ -437,14 +442,14 @@ inline Tensor fast_exp_float32(const Tensor& _x, std::string name, std::string t
// clamp x
auto x = ::tvm::max(::tvm::min(_x(i), x_hi), x_lo);
// integer part
- auto n = ::tvm::floor(x * log2e + one_half);
+ auto n = ::tvm::prim::floor(x * log2e + one_half);
// fractional part
auto f = x - n * ln2;
auto y =
(((((p[0] * f + p[1]) * f + p[2]) * f + p[3]) * f + p[4]) * f + p[5]) * f * f + f + one;
// Return 2^m * exp(r).
- auto ef =
- tvm::reinterpret(PrimType::Float(32), ::tvm::cast(PrimType::Int(32), n + b) << 23);
+ auto ef = tvm::prim::reinterpret(PrimType::Float(32),
+ ::tvm::prim::cast(PrimType::Int(32), n + b) << 23);
return ::tvm::max(ef * y, _x(i)); // NOLINT(*)
},
name, tag);
@@ -467,7 +472,7 @@ inline Tensor fast_exp(const Tensor& x, std::string name = "T_fast_exp",
return ret;
} else {
return compute(
- x->shape, [&](const ffi::Array<PrimVar>& i) { return ::tvm::exp(x(i)); }, name, tag);
+ x->shape, [&](const ffi::Array<PrimVar>& i) { return ::tvm::prim::exp(x(i)); }, name, tag);
}
}
@@ -475,6 +480,7 @@ inline Tensor fast_exp(const Tensor& x, std::string name = "T_fast_exp",
* \brief Fast_erf_float expression from Eigen
*/
inline Tensor fast_erf_float32(const Tensor& data, std::string name, std::string tag) {
+ using namespace tvm::prim;
return compute(
data->shape, [&](const ffi::Array<PrimVar>& i) { return fast_erf_float_expr(data(i), 32); },
name, tag);
@@ -484,6 +490,7 @@ inline Tensor fast_erf_float32(const Tensor& data, std::string name, std::string
* \brief Fast_erf_float expression from Eigen for float16.
*/
inline Tensor fast_erf_float16(const Tensor& data, std::string name, std::string tag) {
+ using namespace tvm::prim;
return compute(
data->shape, [&](const ffi::Array<PrimVar>& i) { return fast_erf_float_expr(data(i), 16); },
name, tag);
diff --git a/include/tvm/topi/nn.h b/include/tvm/topi/nn.h
index c7a778cc75..8878959acf 100644
--- a/include/tvm/topi/nn.h
+++ b/include/tvm/topi/nn.h
@@ -38,7 +38,6 @@
namespace tvm {
namespace topi {
-
using namespace tvm::te;
/*!
@@ -57,7 +56,7 @@ inline tvm::te::Tensor relu(const tvm::te::Tensor& t, T threshold = static_cast<
return tvm::te::compute(
t->shape,
[&](const tvm::ffi::Array<tvm::PrimVar>& i) {
- auto threshold_const = tvm::tirx::MakeConst(tvm::PrimType(t->dtype), threshold);
+ auto threshold_const = tvm::prim::MakeConst(tvm::PrimType(t->dtype), threshold);
return tvm::max(t(i), threshold_const);
},
name, tag);
@@ -80,7 +79,7 @@ inline tvm::te::Tensor leaky_relu(const tvm::te::Tensor& t, double alpha = 0.1,
t->shape,
[&](const tvm::ffi::Array<tvm::PrimVar>& i) {
auto value = t(i);
- auto calpha = tvm::tirx::MakeConst(value.ty(), alpha);
+ auto calpha = tvm::prim::MakeConst(value.ty(), alpha);
return tvm::prim::Select(value > 0, value, value * calpha);
},
name, tag);
@@ -158,6 +157,7 @@ inline tvm::te::Tensor pad(
tvm::ffi::Array<tvm::PrimExpr> pad_after = tvm::ffi::Array<tvm::PrimExpr>(),
PrimExpr pad_value = PrimExpr(), std::string name = "T_pad", std::string tag = kElementWise,
std::string pad_mode = "constant", const ffi::Array<PrimExpr>* dyn_output_shape = nullptr) {
+ using namespace tvm::prim;
if (pad_after.size() < pad_before.size()) {
for (size_t i = pad_after.size(); i < pad_before.size(); ++i) {
pad_after.push_back(pad_before[i]);
@@ -171,10 +171,10 @@ inline tvm::te::Tensor pad(
tvm::ffi::Array<tvm::PrimExpr> pad_after_int32;
for (const auto& ele : pad_before) {
- pad_before_int32.push_back(tvm::cast(tvm::PrimType::Int(32), ele));
+ pad_before_int32.push_back(tvm::prim::cast(tvm::PrimType::Int(32), ele));
}
for (const auto& ele : pad_after) {
- pad_after_int32.push_back(tvm::cast(tvm::PrimType::Int(32), ele));
+ pad_after_int32.push_back(tvm::prim::cast(tvm::PrimType::Int(32), ele));
}
tvm::ffi::Array<tvm::PrimExpr> output_shape;
@@ -194,7 +194,7 @@ inline tvm::te::Tensor pad(
}
if (!pad_value.defined()) {
- pad_value = tvm::tirx::MakeConst(tvm::PrimType(t->dtype), 0);
+ pad_value = tvm::prim::MakeConst(tvm::PrimType(t->dtype), 0);
}
auto l = [&](tvm::ffi::Array<tvm::PrimVar> ovars) {
@@ -287,7 +287,8 @@ inline tvm::te::Tensor conv2d_nchw(const tvm::te::Tensor& I, const tvm::te::Tens
auto T =
(pad_h == 0 && pad_w == 0) ? I : pad(I, {tvm::PrimExpr(0), tvm::PrimExpr(0), pad_h, pad_w});
auto l = [&](tvm::PrimVar b, tvm::PrimVar o, tvm::PrimVar h, tvm::PrimVar w) {
- return tvm::sum(T(b, i, stride_h * h + kh, stride_w * w + kw) * W(o, i, kh, kw), {i, kh, kw});
+ return tvm::prim::sum(T(b, i, stride_h * h + kh, stride_w * w + kw) * W(o, i, kh, kw),
+ {i, kh, kw});
};
return tvm::te::compute(output_shape, l, name, tag);
}
@@ -330,7 +331,8 @@ inline tvm::te::Tensor conv2d_hwcn(const tvm::te::Tensor& I, const tvm::te::Tens
auto kw = tvm::te::reduce_axis(tvm::Range{0, W->shape[1]}, "kw");
auto T = (pad_h == 0 && pad_w == 0) ? I : pad(I, {pad_h, pad_w});
auto l = [&](tvm::PrimVar b, tvm::PrimVar o, tvm::PrimVar h, tvm::PrimVar w) {
- return tvm::sum(T(stride_h * h + kh, stride_w * w + kw, i, b) * W(kh, kw, i, o), {i, kh, kw});
+ return tvm::prim::sum(T(stride_h * h + kh, stride_w * w + kw, i, b) * W(kh, kw, i, o),
+ {i, kh, kw});
};
return tvm::te::compute(output_shape, l, name, tag);
}
@@ -377,9 +379,9 @@ inline tvm::te::Tensor depthwise_conv2d_nchw(const tvm::te::Tensor& I, const tvm
auto T =
(pad_h == 0 && pad_w == 0) ? I : pad(I, {tvm::PrimExpr(0), tvm::PrimExpr(0), pad_h, pad_w});
auto l = [&](tvm::PrimVar b, tvm::PrimVar o, tvm::PrimVar h, tvm::PrimVar w) {
- return tvm::sum(T(b, indexdiv(i, pCM), stride_h * h + kh, stride_w * w + kw) *
- W(indexdiv(i, pCM), indexmod(o, pCM), kh, kw),
- {i, kh, kw});
+ return tvm::prim::sum(T(b, indexdiv(i, pCM), stride_h * h + kh, stride_w * w + kw) *
+ W(indexdiv(i, pCM), indexmod(o, pCM), kh, kw),
+ {i, kh, kw});
};
return tvm::te::compute(output_shape, l, name, tag);
}
@@ -406,9 +408,9 @@ inline tvm::te::Tensor depthwise_conv2d_nhwc(const tvm::te::Tensor& I, const tvm
auto T =
(pad_h == 0 && pad_w == 0) ? I : pad(I, {tvm::PrimExpr(0), pad_h, pad_w, tvm::PrimExpr(0)});
auto l = [&](tvm::PrimVar b, tvm::PrimVar h, tvm::PrimVar w, tvm::PrimVar o) {
- return tvm::sum(T(b, stride_h * h + kh, stride_w * w + kw, indexdiv(i, pCM)) *
- W(kh, kw, indexdiv(i, pCM), indexmod(o, pCM)),
- {kh, kw, i});
+ return tvm::prim::sum(T(b, stride_h * h + kh, stride_w * w + kw, indexdiv(i, pCM)) *
+ W(kh, kw, indexdiv(i, pCM), indexmod(o, pCM)),
+ {kh, kw, i});
};
return tvm::te::compute(output_shape, l, name, tag);
}
@@ -462,8 +464,8 @@ inline tvm::te::Tensor group_conv2d_ngchw(const tvm::te::Tensor& I, const tvm::t
tvm::PrimVar o = args[2];
tvm::PrimVar h = args[3];
tvm::PrimVar w = args[4];
- return tvm::sum(I(b, g, i, stride_h * h + kh, stride_w * w + kw) * W(g, i, o, kh, kw),
- {i, kh, kw});
+ return tvm::prim::sum(I(b, g, i, stride_h * h + kh, stride_w * w + kw) * W(g, i, o, kh, kw),
+ {i, kh, kw});
};
return tvm::te::compute(output_shape, l, name, tag);
}
@@ -496,19 +498,19 @@ inline tvm::te::Tensor space_to_batch_nd(const tvm::te::Tensor& data,
tvm::ffi::Array<tvm::PrimExpr> pad_after_int32;
// pad size for batch dimension is 0
- pad_before_int32.push_back(tvm::cast(tvm::PrimType::Int(32), 0));
- pad_after_int32.push_back(tvm::cast(tvm::PrimType::Int(32), 0));
+ pad_before_int32.push_back(tvm::prim::cast(tvm::PrimType::Int(32), 0));
+ pad_after_int32.push_back(tvm::prim::cast(tvm::PrimType::Int(32), 0));
// insert pad sizes given for spatial dimensions
for (const auto& ele : pad_before) {
- pad_before_int32.push_back(tvm::cast(tvm::PrimType::Int(32), ele));
+ pad_before_int32.push_back(tvm::prim::cast(tvm::PrimType::Int(32), ele));
}
for (const auto& ele : pad_after) {
- pad_after_int32.push_back(tvm::cast(tvm::PrimType::Int(32), ele));
+ pad_after_int32.push_back(tvm::prim::cast(tvm::PrimType::Int(32), ele));
}
// pad the input with paddings provided
if (!pad_value.defined()) {
- pad_value = tvm::tirx::MakeConst(tvm::PrimType(data->dtype), 0);
+ pad_value = tvm::prim::MakeConst(tvm::PrimType(data->dtype), 0);
}
padded_t = pad(data, pad_before_int32, pad_after_int32, pad_value);
@@ -678,7 +680,7 @@ inline Tensor nll_loss(const Tensor& predictions, const Tensor& targets, const T
[&](const tvm::ffi::Array<tvm::PrimVar>& target_indices) {
auto c = targets();
return tvm::prim::Select(c != ignore_index, -predictions(c) * weights(c),
- tvm::tirx::MakeConst(tvm::PrimType(predictions->dtype), 0));
+ tvm::prim::MakeConst(tvm::PrimType(predictions->dtype), 0));
},
name, tag);
if (reduction == "mean") {
@@ -687,7 +689,7 @@ inline Tensor nll_loss(const Tensor& predictions, const Tensor& targets, const T
[&](const tvm::ffi::Array<tvm::PrimVar>& target_indices) {
auto c = targets();
return tvm::prim::Select(c != ignore_index, weights(c),
- tvm::tirx::MakeConst(tvm::PrimType(predictions->dtype), 0));
+ tvm::prim::MakeConst(tvm::PrimType(predictions->dtype), 0));
},
name, tag);
return topi::divide(T, W);
@@ -706,7 +708,7 @@ inline Tensor nll_loss(const Tensor& predictions, const Tensor& targets, const T
pred_indices.push_back(target_indices[i]); // indices for multidimensional loss
}
return tvm::prim::Select(c != ignore_index, -predictions(pred_indices) * weights(c),
- tvm::tirx::MakeConst(tvm::PrimType(predictions->dtype), 0));
+ tvm::prim::MakeConst(tvm::PrimType(predictions->dtype), 0));
},
name, tag);
TVM_FFI_ICHECK(T->shape.size() != 0);
@@ -716,7 +718,7 @@ inline Tensor nll_loss(const Tensor& predictions, const Tensor& targets, const T
[&](const tvm::ffi::Array<tvm::PrimVar>& target_indices) {
auto c = targets(target_indices);
return tvm::prim::Select(c != ignore_index, weights(c),
- tvm::tirx::MakeConst(tvm::PrimType(predictions->dtype), 0));
+ tvm::prim::MakeConst(tvm::PrimType(predictions->dtype), 0));
},
name, tag);
return topi::divide(topi::sum(T, tvm::ffi::Array<int64_t>(nullptr)),
diff --git a/include/tvm/topi/nn/bnn.h b/include/tvm/topi/nn/bnn.h
index 4072055191..f114c4dc9f 100644
--- a/include/tvm/topi/nn/bnn.h
+++ b/include/tvm/topi/nn/bnn.h
@@ -78,7 +78,7 @@ inline tvm::te::Tensor binarize_pack(const tvm::te::Tensor& data, int axis,
idx.push_back(i == static_cast<size_t>(axis) ? start_idx[i] + static_cast<int>(j)
: start_idx[i]);
}
- auto sign = tvm::cast(PrimType::UInt(32), data(idx) >= 0);
+ auto sign = tvm::prim::cast(PrimType::UInt(32), data(idx) >= 0);
packed = (packed | sign);
if (j == 31) {
return packed;
@@ -99,6 +99,7 @@ inline tvm::te::Tensor binarize_pack(const tvm::te::Tensor& data, int axis,
* \return Tensor with shape [batch, out_dim], dtype is float32
*/
inline tvm::te::Tensor binary_dense(const tvm::te::Tensor& data, const tvm::te::Tensor& weight) {
+ using namespace tvm::prim;
TVM_FFI_ICHECK_EQ(data->shape.size(), 2) << "binary_dense requires 2-D data";
TVM_FFI_ICHECK_EQ(weight->shape.size(), 2) << "binary_dense requires 2-D weight";
TVM_FFI_ICHECK_EQ(data->dtype, PrimType::UInt(32)) << "binary_dense requires uint32 data";
@@ -111,7 +112,9 @@ inline tvm::te::Tensor binary_dense(const tvm::te::Tensor& data, const tvm::te::
auto k = tvm::te::reduce_axis(Range(0, in_dim), "k");
auto matmul = tvm::te::compute(
{batch, out_dim},
- [&](PrimVar i, PrimVar j) { return tvm::sum(popcount(data(i, k) ^ weight(j, k)), {k}); },
+ [&](PrimVar i, PrimVar j) {
+ return tvm::prim::sum(popcount(data(i, k) ^ weight(j, k)), {k});
+ },
"tensor", "binary_dense");
return tvm::te::compute(
diff --git a/include/tvm/topi/nn/dense.h b/include/tvm/topi/nn/dense.h
index d6e86fbac3..55e4e43b72 100644
--- a/include/tvm/topi/nn/dense.h
+++ b/include/tvm/topi/nn/dense.h
@@ -61,14 +61,15 @@ inline tvm::te::Tensor dense(const tvm::te::Tensor& data, const tvm::te::Tensor&
auto matmul = tvm::te::compute(
{batch, out_dim},
[&](PrimVar i, PrimVar j) {
- return tvm::sum(tvm::cast(out_dtype, data(i, k)) * tvm::cast(out_dtype, weight(j, k)), {k});
+ return tvm::prim::sum(
+ tvm::prim::cast(out_dtype, data(i, k)) * tvm::prim::cast(out_dtype, weight(j, k)), {k});
},
"tensor", "dense");
if (bias.defined()) {
matmul = tvm::te::compute(
{batch, out_dim},
- [&](PrimVar i, PrimVar j) { return matmul(i, j) + tvm::cast(out_dtype, bias(j)); },
+ [&](PrimVar i, PrimVar j) { return matmul(i, j) + tvm::prim::cast(out_dtype, bias(j)); },
"tensor", kBroadcast);
}
diff --git a/include/tvm/topi/nn/dilate.h b/include/tvm/topi/nn/dilate.h
index 18d5ab7b7b..ce59700a48 100644
--- a/include/tvm/topi/nn/dilate.h
+++ b/include/tvm/topi/nn/dilate.h
@@ -69,6 +69,7 @@ PrimExpr all(ffi::Array<PrimExpr> args) {
*/
inline Tensor dilate(const Tensor& x, ffi::Array<PrimExpr> strides, double dilation_value,
std::string name = "tensor", std::string tag = kInjective) {
+ using namespace tvm::prim;
auto n = x->shape.size();
TVM_FFI_ICHECK_EQ(n, strides.size())
<< "strides size (" << strides.size() << ") must match dimension of x (" << n << ")";
diff --git a/include/tvm/topi/nn/group_norm.h b/include/tvm/topi/nn/group_norm.h
index b54904ec65..222cc4131b 100644
--- a/include/tvm/topi/nn/group_norm.h
+++ b/include/tvm/topi/nn/group_norm.h
@@ -40,6 +40,7 @@ inline Tensor group_norm(const Tensor& data, const Tensor& gamma, const Tensor&
int num_groups, int channel_axis, const ffi::Array<int64_t>& axes,
double epsilon, std::string name = "T_group_norm",
std::string tag = kInjective) {
+ using namespace tvm::prim;
const auto& data_type = data->dtype;
const auto& gamma_type = gamma.defined() ? gamma->dtype : data_type;
const auto& beta_type = beta.defined() ? beta->dtype : data_type;
@@ -143,7 +144,7 @@ inline Tensor group_norm(const Tensor& data, const Tensor& gamma, const Tensor&
auto mean = temp_x(non_reduce_indices) / reduce_extent;
auto var = temp_x2(non_reduce_indices) / reduce_extent - mean * mean;
PrimExpr group_norm = (data_reshaped(indices) - mean) *
- tvm::rsqrt(var + MakeConst(PrimType(data->dtype), epsilon));
+ tvm::prim::rsqrt(var + MakeConst(PrimType(data->dtype), epsilon));
if (is_float16) {
group_norm = prim::Cast(PrimType::Float(16), group_norm);
}
diff --git a/include/tvm/topi/nn/instance_norm.h b/include/tvm/topi/nn/instance_norm.h
index 60c738fb69..519bfba345 100644
--- a/include/tvm/topi/nn/instance_norm.h
+++ b/include/tvm/topi/nn/instance_norm.h
@@ -53,6 +53,7 @@ using namespace tvm::te;
inline Tensor instance_norm(const Tensor& data, const Tensor& gamma, const Tensor& beta,
int channel_axis, const ffi::Array<int64_t>& axis, double epsilon,
std::string name = "T_instance_norm", std::string tag = kInjective) {
+ using namespace tvm::prim;
const auto& data_type = data->dtype;
const auto& gamma_type = gamma.defined() ? gamma->dtype : data_type;
const auto& beta_type = beta.defined() ? beta->dtype : data_type;
@@ -125,7 +126,8 @@ inline Tensor instance_norm(const Tensor& data, const Tensor& gamma, const Tenso
channel = indices[channel_axis];
auto mean = temp_x(non_reduce_indices) / reduce_extent;
auto var = temp_x2(non_reduce_indices) / reduce_extent - mean * mean;
- auto instance_norm = (data(indices) - mean) * tvm::rsqrt(var + MakeConst(var.ty(), epsilon));
+ auto instance_norm =
+ (data(indices) - mean) * tvm::prim::rsqrt(var + MakeConst(var.ty(), epsilon));
if (is_float16) {
instance_norm = prim::Cast(PrimType::Float(16), instance_norm);
}
diff --git a/include/tvm/topi/nn/layer_norm.h b/include/tvm/topi/nn/layer_norm.h
index 27d7f51887..acc800219e 100644
--- a/include/tvm/topi/nn/layer_norm.h
+++ b/include/tvm/topi/nn/layer_norm.h
@@ -52,6 +52,7 @@ using namespace tvm::te;
inline Tensor layer_norm(const Tensor& data, const Tensor& gamma, const Tensor& beta,
const ffi::Array<int64_t>& axis, double epsilon,
std::string name = "T_layer_norm", std::string tag = kInjective) {
+ using namespace tvm::prim;
const auto& data_type = data->dtype;
const auto& gamma_type = gamma.defined() ? gamma->dtype : data_type;
const auto& beta_type = beta.defined() ? beta->dtype : data_type;
@@ -99,7 +100,7 @@ inline Tensor layer_norm(const Tensor& data, const Tensor& gamma, const Tensor&
if (is_float16) {
x = prim::Cast(f32_ty, x);
}
- return sum(x, reduce_axes);
+ return prim::sum(x, reduce_axes);
},
data->op->name + "_sum", kCommReduce);
@@ -125,7 +126,7 @@ inline Tensor layer_norm(const Tensor& data, const Tensor& gamma, const Tensor&
x = prim::Cast(f32_ty, x);
}
PrimExpr diff = x - temp_mean(indices);
- return sum(diff * diff, reduce_axes);
+ return prim::sum(diff * diff, reduce_axes);
},
data->op->name + "_var_sum", kCommReduce);
@@ -140,7 +141,7 @@ inline Tensor layer_norm(const Tensor& data, const Tensor& gamma, const Tensor&
}
auto mean = temp_mean(non_reduce_indices);
auto var = temp_var_sum(non_reduce_indices) / reduce_extent;
- auto layer_norm = (data(indices) - mean) * rsqrt(var + MakeConst(var.ty(), epsilon));
+ auto layer_norm = (data(indices) - mean) * prim::rsqrt(var + MakeConst(var.ty(), epsilon));
if (is_float16) {
layer_norm = prim::Cast(PrimType::Float(16), layer_norm);
}
diff --git a/include/tvm/topi/nn/local_response_norm.h b/include/tvm/topi/nn/local_response_norm.h
index 35e32d4178..07990d8205 100644
--- a/include/tvm/topi/nn/local_response_norm.h
+++ b/include/tvm/topi/nn/local_response_norm.h
@@ -69,24 +69,24 @@ inline Tensor lrn(const Tensor& data, int size, int axis = 1, float alpha = 0.00
sqr_sum = tvm::te::compute(
input_shape,
[&](PrimVar i, PrimVar l, PrimVar j, PrimVar k) {
- return tvm::sum(pad_data(i, l + rxs, j, k) * pad_data(i, l + rxs, j, k), {rxs});
+ return tvm::prim::sum(pad_data(i, l + rxs, j, k) * pad_data(i, l + rxs, j, k), {rxs});
},
"tensor", "sqr_sum");
} else if (axis == 3) {
sqr_sum = tvm::te::compute(
input_shape,
[&](PrimVar i, PrimVar l, PrimVar j, PrimVar k) {
- return tvm::sum(pad_data(i, l, j, k + rxs) * pad_data(i, l, j, k + rxs), {rxs});
+ return tvm::prim::sum(pad_data(i, l, j, k + rxs) * pad_data(i, l, j, k + rxs), {rxs});
},
"tensor", "sqr_sum");
}
- PrimExpr alpha_imm = tvm::te::MakeConst(PrimType(data->dtype), alpha);
- PrimExpr beta_imm = tvm::te::MakeConst(PrimType(data->dtype), beta);
- PrimExpr bias_imm = tvm::te::MakeConst(PrimType(data->dtype), bias);
+ PrimExpr alpha_imm = tvm::prim::MakeConst(PrimType(data->dtype), alpha);
+ PrimExpr beta_imm = tvm::prim::MakeConst(PrimType(data->dtype), beta);
+ PrimExpr bias_imm = tvm::prim::MakeConst(PrimType(data->dtype), bias);
auto sqrt_sum_up = tvm::te::compute(
input_shape,
[&](PrimVar i, PrimVar j, PrimVar k, PrimVar l) {
- return tvm::pow(bias_imm + (div(alpha_imm * sqr_sum(i, j, k, l), size)), beta_imm);
+ return tvm::prim::pow(bias_imm + (div(alpha_imm * sqr_sum(i, j, k, l), size)), beta_imm);
},
"tensor", kElementWise);
return topi::divide(data, sqrt_sum_up);
diff --git a/include/tvm/topi/nn/pooling.h b/include/tvm/topi/nn/pooling.h
index 2471bb4680..d0c4ad2840 100644
--- a/include/tvm/topi/nn/pooling.h
+++ b/include/tvm/topi/nn/pooling.h
@@ -52,6 +52,7 @@ inline Tensor pool_grad_impl(const Tensor& out_grad, const Tensor& x,
const ffi::Array<PrimExpr>& padding_size, PoolType pool_type,
bool ceil_mode, const size_t height_axis, const size_t width_axis,
bool count_include_pad) {
+ using namespace tvm::prim;
TVM_FFI_ICHECK(out_grad->shape.size() >= 2) << "Pooling grad output must >= 2-D (H, W)";
TVM_FFI_ICHECK(x->shape.size() >= 2) << "Pooling input must >= 2-D (H, W)";
TVM_FFI_ICHECK_EQ(kernel_size.size(), 2) << "Pooling kernel_size must have 2 elements";
@@ -118,7 +119,8 @@ inline Tensor pool_grad_impl(const Tensor& out_grad, const Tensor& x,
auto argmax = MakeArgmaxReducer();
auto pad_x =
- do_pad ? pad(x, pad_before, pad_after, tvm::min_value(PrimType(x->dtype)), "pad_temp") : x;
+ do_pad ? pad(x, pad_before, pad_after, tvm::prim::min_value(PrimType(x->dtype)), "pad_temp")
+ : x;
auto mp_argmax = tvm::te::compute(
out_shape,
@@ -155,7 +157,7 @@ inline Tensor pool_grad_impl(const Tensor& out_grad, const Tensor& x,
pad_inds[width_axis] < kernel_width, IntImm(pad_inds[width_axis].ty(), 0),
(pad_inds[width_axis] - kernel_width) / stride_width + 1);
- return tvm::sum(
+ return tvm::prim::sum(
tvm::if_then_else(prim::And(prim::And(out_idx[height_axis] >= out_idx_lower_h,
out_idx[width_axis] >= out_idx_lower_w),
mp_inds(out_idx) == idx),
@@ -200,7 +202,7 @@ inline Tensor pool_grad_impl(const Tensor& out_grad, const Tensor& x,
w_start = max(w_start, IntImm(w_start.ty(), 0));
divide_factor = max((h_end - h_start) * (w_end - w_start), IntImm(h_end.ty(), 1));
}
- return tvm::sum(
+ return tvm::prim::sum(
tvm::if_then_else(prim::And(prim::And(out_idx[height_axis] >= out_idx_lower_h,
out_idx[height_axis] < out_height),
prim::And(out_idx[width_axis] >= out_idx_lower_w,
@@ -379,7 +381,7 @@ inline Tensor adaptive_pool_impl(const Tensor& x, const ffi::Array<PrimExpr>& ou
ffi::Array<PrimExpr> indices;
ffi::Array<tirx::IterVar> reduce_axes;
std::tie(indices, reduce_axes) = get_iter_vars(output, true);
- return tvm::sum(x(indices), reduce_axes);
+ return tvm::prim::sum(x(indices), reduce_axes);
},
"adaptive_pool_sum", "adaptive_pool_sum");
@@ -390,9 +392,9 @@ inline Tensor adaptive_pool_impl(const Tensor& x, const ffi::Array<PrimExpr>& ou
ffi::Array<tirx::IterVar> reduce_axes;
std::tie(indices, reduce_axes) = get_iter_vars(output, false);
- PrimExpr divide_factor = tvm::cast(PrimType(x->dtype), 1);
+ PrimExpr divide_factor = tvm::prim::cast(PrimType(x->dtype), 1);
for (size_t i = 0; i < n_dim; ++i) {
- divide_factor *= tvm::cast(PrimType::Int(32), reduce_axes[i]->dom->extent);
+ divide_factor *= tvm::prim::cast(PrimType::Int(32), reduce_axes[i]->dom->extent);
}
return div(pool_sum(indices), divide_factor);
@@ -519,6 +521,7 @@ inline Tensor pool_impl_nd(const Tensor& x, const ffi::Array<PrimExpr>& kernel_s
const ffi::Array<PrimExpr>& dilation_size,
const ffi::Array<PrimExpr>& padding_size, PoolType pool_type,
bool ceil_mode, const std::vector<int>& axis, bool count_include_pad) {
+ using namespace tvm::prim;
int k_size = kernel_size.size();
int x_size = x->shape.size();
TVM_FFI_ICHECK_EQ(stride_size.size(), k_size)
@@ -589,7 +592,8 @@ inline Tensor pool_impl_nd(const Tensor& x, const ffi::Array<PrimExpr>& kernel_s
ffi::Map<ffi::String, ffi::Any> attrs;
if (pool_type == kMaxPool) {
auto temp =
- do_pad ? pad(x, pad_before, pad_after, tvm::min_value(PrimType(x->dtype)), "pad_temp") : x;
+ do_pad ? pad(x, pad_before, pad_after, tvm::prim::min_value(PrimType(x->dtype)), "pad_temp")
+ : x;
attrs.Set("schedule_rule", tvm::ffi::String("meta_schedule.pool_max"));
return tvm::te::compute(
out_shape,
@@ -620,7 +624,7 @@ inline Tensor pool_impl_nd(const Tensor& x, const ffi::Array<PrimExpr>& kernel_s
int ii = axis[i];
indices.Set(ii, output[ii] * stride[i] + daxis[i] * dilation[i]);
}
- return tvm::sum(temp(indices), daxis);
+ return tvm::prim::sum(temp(indices), daxis);
},
"pool_sum", "pool_sum");
diff --git a/include/tvm/topi/nn/rms_norm.h b/include/tvm/topi/nn/rms_norm.h
index 8036d5373d..a4c6152a7a 100644
--- a/include/tvm/topi/nn/rms_norm.h
+++ b/include/tvm/topi/nn/rms_norm.h
@@ -50,6 +50,7 @@ using namespace tvm::te;
inline Tensor rms_norm(const Tensor& data, const Tensor& weight, const ffi::Array<int64_t>& axis,
double epsilon, std::string name = "T_rms_norm",
std::string tag = kInjective) {
+ using namespace tvm::prim;
const auto& data_type = data->dtype;
const auto& weight_type = weight.defined() ? weight->dtype : data_type;
TVM_FFI_ICHECK(data_type == weight_type) << "rms_norm: data and weight must have the same type";
@@ -74,8 +75,8 @@ inline Tensor rms_norm(const Tensor& data, const Tensor& weight, const ffi::Arra
non_reduce_indices.push_back(indices[i]);
}
}
- auto output = tvm::rsqrt(square_sum(non_reduce_indices) / reduce_extent +
- MakeConst(PrimType(data_type), epsilon));
+ auto output = tvm::prim::rsqrt(square_sum(non_reduce_indices) / reduce_extent +
+ MakeConst(PrimType(data_type), epsilon));
return output;
};
auto rsqrt_shape = ffi::Array<PrimExpr>();
diff --git a/include/tvm/topi/nn/softmax.h b/include/tvm/topi/nn/softmax.h
index dcf7e35e1c..518a113d97 100644
--- a/include/tvm/topi/nn/softmax.h
+++ b/include/tvm/topi/nn/softmax.h
@@ -49,6 +49,7 @@ using namespace tvm::te;
*/
inline Tensor softmax(const Tensor& x, int axis = -1, std::string name = "tensor",
std::string tag = "softmax_output") {
+ using namespace tvm::prim;
auto input_shape = x->shape;
auto ndim = input_shape.size();
if (axis < 0) {
@@ -92,12 +93,12 @@ inline Tensor softmax(const Tensor& x, int axis = -1, std::string name = "tensor
auto _compute_exp = [&](const Tensor& max_elem, const ffi::Array<PrimVar>& indices) {
auto non_reduce_indices = get_non_reduce_indices(indices);
- return tvm::exp(x(indices) - max_elem(non_reduce_indices));
+ return tvm::prim::exp(x(indices) - max_elem(non_reduce_indices));
};
auto _compute_expsum = [&](const Tensor& exp, const ffi::Array<PrimVar>& indices) {
auto eval_range = insert_reduce_index(indices, k2);
- return tvm::sum(exp(eval_range), {k2});
+ return tvm::prim::sum(exp(eval_range), {k2});
};
auto _normalize = [&](const Tensor& exp, const Tensor& expsum,
@@ -141,11 +142,12 @@ inline Tensor log_softmax(const Tensor& x, std::string name = "tensor",
k = tvm::te::reduce_axis(Range(0, n), "k");
auto expsum = tvm::te::compute(
- {m}, [&](PrimVar i) { return tvm::sum(tvm::exp(x(i, k) - max_elem(i)), {k}); });
+ {m}, [&](PrimVar i) { return tvm::prim::sum(tvm::prim::exp(x(i, k) - max_elem(i)), {k}); });
return tvm::te::compute(
- x->shape, [&](PrimVar i, PrimVar j) { return x(i, j) - max_elem(i) - tvm::log(expsum(i)); },
- name, tag);
+ x->shape,
+ [&](PrimVar i, PrimVar j) { return x(i, j) - max_elem(i) - tvm::prim::log(expsum(i)); }, name,
+ tag);
}
} // namespace nn
diff --git a/include/tvm/topi/reduction.h b/include/tvm/topi/reduction.h
index 9a34a09098..ef8d53c9e2 100644
--- a/include/tvm/topi/reduction.h
+++ b/include/tvm/topi/reduction.h
@@ -39,7 +39,6 @@
namespace tvm {
namespace topi {
-
using namespace tvm::te;
/*! \brief The operation to use for CommReduce */
@@ -128,7 +127,7 @@ inline ffi::Array<PrimExpr> MakeReduceTargetShape(const std::vector<int>& real_a
* \brief Create a reduction operation.
*
* \param data The input tensor.
- * \param func The reduction function eg. tvm::sum
+ * \param func The reduction function eg. tvm::prim::sum
* \param target_shape The output Tensor shape.
* \param reduce_axes The real axes along which the reduction is performed.
* \param squeeze_axes The real axes to squeeze. Unsqueezed, reduced axes will
@@ -173,7 +172,7 @@ inline Tensor DoCommReduce(const Tensor& data, FReduce func,
*
* \param data The input tensor.
* \param axis The axes along which the reduction is performed.
- * \param func The reduction function eg. tvm::sum
+ * \param func The reduction function eg. tvm::prim::sum
* \param keepdims If this is set to true, the axes which are reduced are
* left in the result as dimensions with size one. This enables the result
* to broadcast correctly against the input array.
@@ -315,10 +314,10 @@ inline PrimExpr MaxOp(PrimExpr source, ffi::Array<IterVar> axis, ffi::Array<Prim
return tvm::max(source, axis, init, span); // NOLINT(*)
}
-/*! \brief Wrap tvm::prod to ensure we get the correct overload */
+/*! \brief Wrap tvm::prim::prod to ensure we get the correct overload */
inline PrimExpr ProdOp(PrimExpr source, ffi::Array<IterVar> axis, ffi::Array<PrimExpr> init = {},
Span span = Span()) {
- return tvm::prod(source, axis, init, span); // NOLINT(*)
+ return tvm::prim::prod(source, axis, init, span); // NOLINT(*)
}
/*!
@@ -338,9 +337,9 @@ inline Tensor sum(const Tensor& data, const ffi::Optional<ffi::Array<int64_t>>&
bool keepdims = false, bool atleast1d = false) {
// Reduction dispatch only depends on boolean element kind; lane encoding is irrelevant here.
if (data->dtype.code() == DLDataTypeCode::kDLBool) {
- return CommReduce(data, axis, tvm::any, keepdims, atleast1d);
+ return CommReduce(data, axis, tvm::prim::any, keepdims, atleast1d);
} else {
- return CommReduce(data, axis, tvm::sum, keepdims, atleast1d);
+ return CommReduce(data, axis, tvm::prim::sum, keepdims, atleast1d);
}
}
@@ -375,7 +374,7 @@ inline Tensor collapse_sum(const Tensor& data, ffi::Array<PrimExpr> target_shape
std::reverse(reduce_axes.begin(), reduce_axes.end());
std::reverse(squeeze_axes.begin(), squeeze_axes.end());
- return DoCommReduce(data, tvm::sum, target_shape, reduce_axes, squeeze_axes);
+ return DoCommReduce(data, tvm::prim::sum, target_shape, reduce_axes, squeeze_axes);
}
/*!
@@ -394,7 +393,7 @@ inline Tensor collapse_sum(const Tensor& data, ffi::Array<PrimExpr> target_shape
*/
inline Tensor all(const Tensor& data, const ffi::Optional<ffi::Array<int64_t>>& axis,
bool keepdims = false, bool atleast1d = false) {
- return CommReduce(data, axis, tvm::all, keepdims, atleast1d);
+ return CommReduce(data, axis, tvm::prim::all, keepdims, atleast1d);
}
/*!
@@ -413,7 +412,7 @@ inline Tensor all(const Tensor& data, const ffi::Optional<ffi::Array<int64_t>>&
*/
inline Tensor any(const Tensor& data, const ffi::Optional<ffi::Array<int64_t>>& axis,
bool keepdims = false, bool atleast1d = false) {
- return CommReduce(data, axis, tvm::any, keepdims, atleast1d);
+ return CommReduce(data, axis, tvm::prim::any, keepdims, atleast1d);
}
/*!
@@ -486,8 +485,8 @@ inline FCommReduce MakeArgminReducer(bool select_last_index = false) {
};
auto fidentity = [&](std::vector<PrimType> types) {
ffi::Array<PrimExpr> result;
- result.push_back(tvm::tirx::MakeConst(types[0], -1)); // idx
- result.push_back(tvm::max_value(types[1])); // val
+ result.push_back(tvm::prim::MakeConst(types[0], -1)); // idx
+ result.push_back(tvm::prim::max_value(types[1])); // val
return result;
};
return MakeCommReducer(fcombine, fidentity, "argmin");
@@ -548,8 +547,8 @@ inline FCommReduce MakeArgmaxReducer(bool select_last_index = false) {
};
auto fidentity = [&](std::vector<PrimType> types) {
ffi::Array<PrimExpr> result;
- result.push_back(tvm::tirx::MakeConst(types[0], -1)); // idx
- result.push_back(tvm::min_value(types[1])); // val
+ result.push_back(tvm::prim::MakeConst(types[0], -1)); // idx
+ result.push_back(tvm::prim::min_value(types[1])); // val
return result;
};
return MakeCommReducer(fcombine, fidentity, "argmax");
@@ -611,7 +610,7 @@ inline FCommReduce MakeTupleSumReducer() {
auto fidentity = [](std::vector<PrimType> types) {
ffi::Array<PrimExpr> result;
for (size_t i = 0; i < types.size(); ++i) {
- result.push_back(tvm::tirx::MakeConst(types[i], 0));
+ result.push_back(tvm::prim::MakeConst(types[i], 0));
}
return result;
};
diff --git a/include/tvm/topi/transform.h b/include/tvm/topi/transform.h
index 6ac02f7844..c340c3b150 100644
--- a/include/tvm/topi/transform.h
+++ b/include/tvm/topi/transform.h
@@ -52,7 +52,6 @@
namespace tvm {
namespace topi {
-
using namespace tvm::te;
using namespace topi::detail;
@@ -339,8 +338,8 @@ inline Tensor reshape(const Tensor& x, ffi::Array<PrimExpr> newshape,
if (is_empty_shape(target_shape) || is_empty_shape(x->shape)) {
return compute(
target_shape,
- [&](const ffi::Array<PrimVar>& indices) { return tvm::cast(PrimType(x->dtype), 0); }, name,
- tag);
+ [&](const ffi::Array<PrimVar>& indices) { return tvm::prim::cast(PrimType(x->dtype), 0); },
+ name, tag);
} else {
return compute(
target_shape,
@@ -970,6 +969,7 @@ inline Tensor strided_slice(const Tensor& x, const ffi::Array<ffi::Optional<IntI
const ffi::Array<ffi::Optional<IntImm>>& end,
const ffi::Array<IntImm>& strides, std::string slice_mode = "end",
std::string name = "T_strided_slice", std::string tag = kInjective) {
+ using namespace tvm::prim;
size_t src_tensor_dim = static_cast<size_t>(x->shape.size());
ffi::Array<int64_t> axes;
for (size_t i = 0; i < src_tensor_dim; ++i) axes.push_back(i);
@@ -1123,8 +1123,8 @@ inline Tensor sequence_mask(const Tensor& data, const Tensor& valid_length, doub
auto bid = out_index[1 - axis];
len_index.push_back(bid);
PrimExpr ret = tvm::if_then_else(
- tvm::cast(PrimType(valid_length->dtype), tid) >= valid_length(len_index),
- tvm::tirx::MakeConst(PrimType(data->dtype), mask_value), data(out_index));
+ tvm::prim::cast(PrimType(valid_length->dtype), tid) >= valid_length(len_index),
+ tvm::prim::MakeConst(PrimType(data->dtype), mask_value), data(out_index));
return ret;
},
name, tag);
@@ -1205,12 +1205,11 @@ inline Tensor take(const Tensor& a, ffi::Variant<Tensor, PrimExpr> indices, int
auto get_index = [&](const ffi::Array<PrimExpr>& indices_position) -> PrimExpr {
if (auto tensor = indices.as<Tensor>()) {
return tensor.value()(indices_position);
- } else if (auto prim = indices.as<PrimExpr>()) {
- TVM_FFI_ICHECK_EQ(indices_position.size(), 0);
- return prim.value();
- } else {
- TVM_FFI_THROW(InternalError) << "Variant did not contain either allowed type";
}
+ auto prim = indices.as<PrimExpr>();
+ TVM_FFI_ICHECK(prim.has_value()) << "Variant did not contain either allowed type";
+ TVM_FFI_ICHECK_EQ(indices_position.size(), 0);
+ return prim.value();
};
if (mode == "clip") {
@@ -1299,7 +1298,7 @@ inline Tensor take(const Tensor& a, ffi::Variant<Tensor, PrimExpr> indices, int
PrimExpr in_bounds = idx >= 0 && idx < axis_dim;
return tvm::if_then_else(
in_bounds, a(real_indices),
- tvm::tirx::MakeConst(PrimType(a->dtype), std::numeric_limits<float>::quiet_NaN()));
+ tvm::prim::MakeConst(PrimType(a->dtype), std::numeric_limits<float>::quiet_NaN()));
},
name, tag);
} else { // mode == "wrap"
@@ -1450,8 +1449,8 @@ inline Tensor tile(const Tensor& x, ffi::Array<int64_t> reps, std::string name =
if (is_empty_shape(new_shape)) {
return compute(
new_shape,
- [&](const ffi::Array<PrimVar>& indices) { return tvm::cast(PrimType(x->dtype), 0); }, name,
- tag);
+ [&](const ffi::Array<PrimVar>& indices) { return tvm::prim::cast(PrimType(x->dtype), 0); },
+ name, tag);
} else {
return compute(
new_shape,
@@ -1486,8 +1485,8 @@ inline Tensor dyn_tile(const Tensor& x, ffi::Array<PrimExpr> new_shape, size_t r
if (is_empty_shape(new_shape)) {
return compute(
new_shape,
- [&](const ffi::Array<PrimVar>& indices) { return tvm::cast(PrimType(x->dtype), 0); }, name,
- tag);
+ [&](const ffi::Array<PrimVar>& indices) { return tvm::prim::cast(PrimType(x->dtype), 0); },
+ name, tag);
} else {
return compute(
new_shape,
@@ -1609,7 +1608,8 @@ inline Tensor gather_nd(const Tensor& data, const Tensor& indices, int batch_dim
if (indices_ty.MatchesCode(DLDataTypeCode::kDLInt, DLDataTypeCode::kDLUInt)) {
real_indices.push_back(indices(indices_position));
} else {
- real_indices.push_back(tvm::cast(tvm::PrimType::Int(32), indices(indices_position)));
+ real_indices.push_back(
+ tvm::prim::cast(tvm::PrimType::Int(32), indices(indices_position)));
}
}
if (real_indices.size() == ndim_d) {
@@ -1644,7 +1644,7 @@ inline tvm::te::Tensor matmul(const tvm::te::Tensor& A, const tvm::te::Tensor& B
tvm::ffi::Array<tvm::PrimExpr> output_shape{A->shape[trans_a ? 1 : 0], B->shape[trans_b ? 0 : 1]};
auto k = tvm::te::reduce_axis(tvm::Range{0, A->shape[trans_a ? 0 : 1]}, "k");
auto l = [&](tvm::PrimVar i, tvm::PrimVar j) {
- return tvm::sum((trans_a ? A[k][i] : A[i][k]) * (trans_b ? B[j][k] : B[k][j]), {k});
+ return tvm::prim::sum((trans_a ? A[k][i] : A[i][k]) * (trans_b ? B[j][k] : B[k][j]), {k});
};
return tvm::te::compute(output_shape, l, name, tag);
}
@@ -1662,6 +1662,7 @@ inline tvm::te::Tensor matmul(const tvm::te::Tensor& A, const tvm::te::Tensor& B
*/
inline Tensor tensordot(const Tensor& A, const tvm::te::Tensor& B, int axes = 2,
std::string name = "T_tensordot", std::string tag = kMatMul) {
+ using namespace tvm::prim;
TVM_FFI_ICHECK_GE(A->shape.size(), axes);
TVM_FFI_ICHECK_GE(B->shape.size(), axes);
@@ -1714,6 +1715,7 @@ inline Tensor tensordot(const Tensor& A, const tvm::te::Tensor& B, int axes = 2,
inline Tensor tensordot(const Tensor& A, const tvm::te::Tensor& B, ffi::Array<PrimExpr> A_axes,
ffi::Array<PrimExpr> B_axes, std::string name = "T_tensordot",
std::string tag = kMatMul) {
+ using namespace tvm::prim;
TVM_FFI_ICHECK_EQ(A_axes.size(), B_axes.size());
auto A_axes_val = GetConstIntValues(A_axes, "A_axes");
@@ -1776,15 +1778,16 @@ inline Tensor arange(const PrimExpr& start, const PrimExpr& stop, const PrimExpr
num_elem = tvm::floordiv((start - stop - step - 1), -step);
} else {
// fallback path for non-integer or step of unknown sign
- num_elem = tvm::cast(PrimType(DefaultIndexType()),
- tvm::ceil(tvm::cast(tvm::PrimType::Float(32), stop - start) / step));
+ num_elem =
+ tvm::prim::cast(PrimType(DefaultIndexType()),
+ tvm::ceil(tvm::prim::cast(tvm::PrimType::Float(32), stop - start) / step));
}
num_elem = analyzer->Simplify(num_elem);
return compute(
{num_elem},
[&](const ffi::Array<PrimVar>& indices) {
- return tvm::cast(dtype, start + step * indices[0]);
+ return tvm::prim::cast(dtype, start + step * indices[0]);
},
name, tag);
}
@@ -1873,7 +1876,7 @@ inline Tensor layout_transform(const Tensor& src, const std::string& src_layout,
in_range = in_range && (src_indices[i] < src->shape[i]);
}
return if_then_else(in_range, src(src_indices),
- tvm::cast(PrimType(src->dtype), PrimExpr(0)));
+ tvm::prim::cast(PrimType(src->dtype), PrimExpr(0)));
},
name, tag, attrs);
}
@@ -2023,7 +2026,7 @@ inline Tensor shape(const Tensor& src, PrimType dtype, const std::string name =
for (int i = 0; i < ndim; ++i) {
ret = tvm::if_then_else(idx == i, src->shape[i], ret);
}
- return tvm::cast(dtype, ret);
+ return tvm::prim::cast(dtype, ret);
},
name, tag);
}
@@ -2053,7 +2056,7 @@ inline te::Tensor tensor_size(const te::Tensor& src, PrimType dtype,
for (int i = 0; i < ndim; ++i) {
ret *= src->shape[i];
}
- return tvm::cast(dtype, ret);
+ return tvm::prim::cast(dtype, ret);
},
name, tag);
}
@@ -2095,8 +2098,8 @@ inline Tensor one_hot(const Tensor& indices, const PrimExpr on_value, const Prim
}
}
- PrimExpr on_value_cast = cast(dtype, on_value);
- PrimExpr off_value_cast = cast(dtype, off_value);
+ PrimExpr on_value_cast = prim::cast(dtype, on_value);
+ PrimExpr off_value_cast = prim::cast(dtype, off_value);
return compute(
oshape,
[&](const ffi::Array<PrimVar>& iter_vars) {
diff --git a/python/tvm/arith/analyzer.py b/python/tvm/arith/analyzer.py
index ad086518fb..59f8ba4e5d 100644
--- a/python/tvm/arith/analyzer.py
+++ b/python/tvm/arith/analyzer.py
@@ -22,7 +22,7 @@ from contextlib import contextmanager
import tvm_ffi
-from tvm import ir, tirx
+from tvm import ir
from tvm.arith import IntSet
from tvm.runtime import Object
@@ -161,7 +161,7 @@ class Analyzer(Object):
"Rebuild TVM with USE_Z3=ON to use Z3-specific Analyzer APIs."
)
- def get_smtlib2(self, expr: tirx.Expr | None = None) -> str:
+ def get_smtlib2(self, expr: ir.Expr | None = None) -> str:
"""Get the current Z3 problem in SMT-LIB2 format.
Raises
@@ -229,7 +229,7 @@ class Analyzer(Object):
self._check_z3_enabled()
return _ffi_api.AnalyzerGetZ3Stats(self)
- def const_int_bound(self, expr: tirx.Expr) -> ConstIntBound:
+ def const_int_bound(self, expr: ir.Expr) -> ConstIntBound:
"""Find constant integer bound for expr.
Parameters
@@ -244,12 +244,12 @@ class Analyzer(Object):
"""
return _ffi_api.AnalyzerConstIntBound(self, expr)
- def const_int_bound_is_bound(self, var: tirx.Var) -> bool:
+ def const_int_bound_is_bound(self, var: ir.Var) -> bool:
"""Check if a variable is bound to a range.
Parameters
----------
- var : tvm.tirx.Var
+ var : tvm.ir.Var
The variable.
Returns
@@ -259,7 +259,7 @@ class Analyzer(Object):
"""
return _ffi_api.AnalyzerConstIntBoundIsBound(self, var)
- def modular_set(self, expr: tirx.Expr) -> ModularSet:
+ def modular_set(self, expr: ir.Expr) -> ModularSet:
"""Find a modular set that expr belongs to.
Parameters
@@ -274,7 +274,7 @@ class Analyzer(Object):
"""
return _ffi_api.AnalyzerModularSet(self, expr)
- def simplify(self, expr: tirx.Expr, steps: int = 2) -> tirx.Expr:
+ def simplify(self, expr: ir.Expr, steps: int = 2) -> ir.Expr:
"""Simplify expression via both rewrite and canonicalization.
Parameters
@@ -312,7 +312,7 @@ class Analyzer(Object):
"""
return _ffi_api.AnalyzerClone(self)
- def rewrite_simplify(self, expr: tirx.Expr) -> tirx.Expr:
+ def rewrite_simplify(self, expr: ir.Expr) -> ir.Expr:
"""Simplify expression via rewriting rules.
Parameters
@@ -334,7 +334,7 @@ class Analyzer(Object):
def reset_rewrite_simplify_stats(self):
_ffi_api.AnalyzerResetRewriteSimplifyStats(self)
- def canonical_simplify(self, expr: tirx.Expr) -> tirx.Expr:
+ def canonical_simplify(self, expr: ir.Expr) -> ir.Expr:
"""Simplify expression via canonicalization.
Parameters
@@ -349,7 +349,7 @@ class Analyzer(Object):
"""
return _ffi_api.AnalyzerCanonicalSimplify(self, expr)
- def int_set(self, expr: tirx.Expr, dom_map: dict[tirx.Var, IntSet] | None = None) -> IntSet:
+ def int_set(self, expr: ir.Expr, dom_map: dict[ir.Var, IntSet] | None = None) -> IntSet:
"""Compute a symbolic IntSet that covers expr for all values in dom_map.
Parameters
@@ -357,7 +357,7 @@ class Analyzer(Object):
expr : Expr
The expression.
- dom_map : Optional[Dict[tvm.tirx.Var, tvm.arith.IntSet]]
+ dom_map : Optional[Dict[tvm.ir.Var, tvm.arith.IntSet]]
The domain for variables to be relaxed. When omitted, the analyzer
uses the domains of the variables already bound to it.
@@ -368,7 +368,7 @@ class Analyzer(Object):
"""
return _ffi_api.AnalyzerIntSet(self, expr, dom_map)
- def can_prove(self, expr: tirx.Expr, strength: ProofStrength = ProofStrength.DEFAULT) -> bool:
+ def can_prove(self, expr: ir.Expr, strength: ProofStrength = ProofStrength.DEFAULT) -> bool:
"""Check whether we can prove expr to be true.
Parameters
@@ -405,18 +405,18 @@ class Analyzer(Object):
def bind(
self,
- var: tirx.Var,
- expr: tirx.Expr | ir.Range,
+ var: ir.Var,
+ expr: ir.Expr | ir.Range,
allow_override: bool = False,
) -> None:
"""Bind a variable to the expression.
Parameters
----------
- var : tvm.tirx.Var
+ var : tvm.ir.Var
The variable.
- expr : Union[tirx.Expr, ir.Range]
+ expr : Union[ir.Expr, ir.Range]
The expression or the range to bind to.
allow_override : bool
@@ -424,7 +424,7 @@ class Analyzer(Object):
"""
return _ffi_api.AnalyzerBind(self, var, expr, allow_override)
- def constraint_scope(self, constraint: tirx.Expr) -> ConstraintScope:
+ def constraint_scope(self, constraint: ir.Expr) -> ConstraintScope:
"""Create a constraint scope.
Parameters
@@ -456,13 +456,13 @@ class Analyzer(Object):
return ConstraintScope(_fenter)
def update(
- self, var: tirx.Var, info: ConstIntBound | ModularSet | IntSet, override: bool = False
+ self, var: ir.Var, info: ConstIntBound | ModularSet | IntSet, override: bool = False
) -> None:
"""Update information about var.
Parameters
----------
- var : tvm.tirx.Var
+ var : tvm.ir.Var
The variable.
info : Union[ConstIntBound, ModularSet, IntSet]
@@ -482,7 +482,7 @@ class Analyzer(Object):
else:
raise TypeError(f"Do not know how to handle type {type(info)}")
- def can_prove_equal(self, lhs: tirx.Expr, rhs: tirx.Expr) -> bool:
+ def can_prove_equal(self, lhs: ir.Expr, rhs: ir.Expr) -> bool:
"""Whether we can prove that lhs == rhs
Parameters
@@ -501,7 +501,7 @@ class Analyzer(Object):
return _ffi_api.AnalyzerCanProveEqual(self, lhs, rhs)
def try_compare(
- self, lhs: tirx.Expr, rhs: tirx.Expr, propagate_inequalities: bool = True
+ self, lhs: ir.Expr, rhs: ir.Expr, propagate_inequalities: bool = True
) -> CompareResult:
"""Compare lhs and rhs using previously provided known comparisons.
diff --git a/python/tvm/arith/bound.py b/python/tvm/arith/bound.py
index 44ab752514..3f59afc7bb 100644
--- a/python/tvm/arith/bound.py
+++ b/python/tvm/arith/bound.py
@@ -24,16 +24,16 @@ def deduce_bound(var, cond, hint_map, relax_map):
Parameters
----------
- var : tvm.tirx.Var
+ var : tvm.ir.Var
The target variable to be deduced.
cond : Expr
The condition
- hint_map : Map[tvm.tirx.Var, IntSet]
+ hint_map : Map[tvm.ir.Var, IntSet]
Domain of variables used to help deduction.
- relax_map : Map[tvm.tirx.Var, IntSet]
+ relax_map : Map[tvm.ir.Var, IntSet]
The fomain of the variables to be relaxed
using the provided domain.
"""
diff --git a/python/tvm/arith/int_set.py b/python/tvm/arith/int_set.py
index db6dde2ff9..c7e79f6795 100644
--- a/python/tvm/arith/int_set.py
+++ b/python/tvm/arith/int_set.py
@@ -102,7 +102,7 @@ def estimate_region_lower_bound(region, var_dom, predicate, analyzer=None):
region : List[Range]
The region to be analyzed.
- var_dom : Dict[tvm.tirx.Var, Range]
+ var_dom : Dict[tvm.ir.Var, Range]
The ranges of the variables
predicate : Expr
@@ -129,7 +129,7 @@ def estimate_region_strict_bound(region, var_dom, predicate, analyzer=None):
region : List[Range]
The region to be analyzed.
- var_dom : Dict[tvm.tirx.Var, Range]
+ var_dom : Dict[tvm.ir.Var, Range]
The ranges of the variables
predicate : Expr
@@ -157,7 +157,7 @@ def estimate_region_upper_bound(region, var_dom, predicate, analyzer=None):
region : List[Range]
The region to be analyzed.
- var_dom : Dict[tvm.tirx.Var, Range]
+ var_dom : Dict[tvm.ir.Var, Range]
The ranges of the variables
predicate : Expr
@@ -180,7 +180,7 @@ def pos_inf():
Returns
----------
- pos_inf : tvm.tirx.Var
+ pos_inf : tvm.ir.Var
A symbolic var that indicates positive infinity
"""
return _ffi_api.PosInf()
@@ -191,7 +191,7 @@ def neg_inf():
Returns
----------
- neg_inf : tvm.tirx.Var
+ neg_inf : tvm.ir.Var
A symbolic var that indicates positive infinity
"""
return _ffi_api.NegInf()
diff --git a/python/tvm/arith/iter_affine_map.py b/python/tvm/arith/iter_affine_map.py
index e65d1d01fc..bac1b030b1 100644
--- a/python/tvm/arith/iter_affine_map.py
+++ b/python/tvm/arith/iter_affine_map.py
@@ -138,7 +138,7 @@ def detect_iter_map(
indices : List[Expr]
The input indices
- input_iters : Map[tvm.tirx.Var, Range]
+ input_iters : Map[tvm.ir.Var, Range]
The domain of each input iterators.
predicate : Expr
@@ -183,7 +183,7 @@ def normalize_to_iter_sum(index, input_iters, analyzer=None):
index : Expr
The input index
- input_iters : Map[tvm.tirx.Var, Range]
+ input_iters : Map[tvm.ir.Var, Range]
The domain of each input iterators.
analyzer : Optional[tvm.arith.Analyzer]
@@ -221,7 +221,7 @@ def iter_map_simplify(
indices : List[Expr]
The input indices
- input_iters : Map[tvm.tirx.Var, Range]
+ input_iters : Map[tvm.ir.Var, Range]
The domain of each input iterators.
predicate : Expr
@@ -304,10 +304,10 @@ def subspace_divide(
bindings : List[Expr]
The input bindings
- input_iters : Map[tvm.tirx.Var, Range]
+ input_iters : Map[tvm.ir.Var, Range]
The domain of input iterator, which is the basis of the whole space
- sub_iters : Array[tvm.tirx.Var]
+ sub_iters : Array[tvm.ir.Var]
The subset of input_iters, which is the basis of the subspace
predicate : Expr
@@ -369,7 +369,7 @@ def inverse_affine_iter_map(iter_map, outputs):
Returns
-------
- results : Map[tvm.tirx.Var, Expr]
+ results : Map[tvm.ir.Var, Expr]
The map from the input to the transformed result.
"""
return _ffi_api.InverseAffineIterMap(iter_map, outputs)
diff --git a/python/tvm/arith/pattern.py b/python/tvm/arith/pattern.py
index 8582b2f576..1840634993 100644
--- a/python/tvm/arith/pattern.py
+++ b/python/tvm/arith/pattern.py
@@ -29,7 +29,7 @@ def detect_linear_equation(expr, var_list):
expr : Expr
The expression to be matched.
- var_list : List[tvm.tirx.Var]
+ var_list : List[tvm.ir.Var]
A list of variables.
Returns
@@ -49,7 +49,7 @@ def detect_clip_bound(expr, var_list):
expr : Expr
The expression to be matched.
- var_list : List[tvm.tirx.Var]
+ var_list : List[tvm.ir.Var]
A list of variables.
Returns
diff --git a/python/tvm/ir/__init__.py b/python/tvm/ir/__init__.py
index bfb378d3bd..a02e77cdd6 100644
--- a/python/tvm/ir/__init__.py
+++ b/python/tvm/ir/__init__.py
@@ -49,6 +49,7 @@ from .expr import (
is_prim_expr,
is_prim_var,
)
+from . import prim
from .function import BaseFunc, CallingConv
from .global_info import GlobalInfo
from .module import IRModule
diff --git a/python/tvm/ir/_overload_prim_expr.py b/python/tvm/ir/_overload_prim_expr.py
index e6566ee786..8ccd3c46ea 100644
--- a/python/tvm/ir/_overload_prim_expr.py
+++ b/python/tvm/ir/_overload_prim_expr.py
@@ -16,138 +16,354 @@
# under the License.
"""Primitive-expression overloads for shared IR expressions."""
+from ..runtime import DataTypeCode, ObjectConvertible, const
+from .base import Span
+from .expr import Expr, is_prim_expr
+from .prim import _ffi_api as _prim_ffi_api
+from .type import PrimType
+
+
+def div_ambiguity_error() -> RuntimeError:
+ return RuntimeError(
+ "TVM supports multiple types of integer divisions, "
+ + "please call div, indexdiv/indexmod, floordiv/floormod "
+ + " or truncdiv/truncmod directly to avoid ambiguity in the code."
+ )
+
+
+def _dtype_is_int(value):
+ if isinstance(value, int):
+ return True
+ if isinstance(value, ExprOp):
+ return value.expr_ty().matches_code(DataTypeCode.INT)
+ if is_prim_expr(value):
+ return value.ty.matches_code(DataTypeCode.INT)
+ return False
+
+
+def _dtype_is_float(value):
+ if isinstance(value, float):
+ return True
+ if isinstance(value, ExprOp):
+ return value.expr_ty().matches_code(DataTypeCode.FLOAT)
+ if is_prim_expr(value):
+ return value.ty.matches_code(DataTypeCode.FLOAT)
+ return False
+
+
+def _is_scalar_operand(value):
+ type_info = getattr(type(value), "__tvm_ffi_type_info__", None)
+ if type_info is not None and type_info.type_key == "tirx.BufferRegion":
+ raise TypeError(
+ "BufferRegion is not a primitive operand; construct a BufferLoad explicitly"
+ )
+ return isinstance(value, ExprOp | int | float) or is_prim_expr(value)
+
+
+class ExprOp:
+ """Operator overloading for Expr like expressions."""
+
+ # TODO(tkonolige): use inspect to add source information to these objects
+
+ def expr_ty(self) -> PrimType:
+ """Return the compile-time primitive type for expression operators."""
+ ty = getattr(self, "ty", None)
+ if isinstance(ty, PrimType):
+ return ty
+ raise TypeError(f"Cannot determine PrimType for {type(self).__name__}")
+
+ def __add__(self, other: Expr) -> Expr:
+ if not _is_scalar_operand(other):
+ return NotImplemented
+ return _prim_ffi_api._OpAdd(self, other, None) # type: ignore
+
+ def __radd__(self, other: Expr) -> Expr:
+ if not _is_scalar_operand(other):
+ return NotImplemented
+ return _prim_ffi_api._OpAdd(other, self, None) # type: ignore
+
+ def __sub__(self, other: Expr) -> Expr:
+ if not _is_scalar_operand(other):
+ return NotImplemented
+ return _prim_ffi_api._OpSub(self, other, None) # type: ignore
+
+ def __rsub__(self, other: Expr) -> Expr:
+ if not _is_scalar_operand(other):
+ return NotImplemented
+ return _prim_ffi_api._OpSub(other, self, None) # type: ignore
+
+ def __mul__(self, other: Expr) -> Expr:
+ if not _is_scalar_operand(other):
+ return NotImplemented
+ return _prim_ffi_api._OpMul(self, other, None) # type: ignore
+
+ def __rmul__(self, other: Expr) -> Expr:
+ if not _is_scalar_operand(other):
+ return NotImplemented
+ return _prim_ffi_api._OpMul(other, self, None) # type: ignore
+
+ def __div__(self, other: Expr) -> Expr:
+ if not _is_scalar_operand(other):
+ return NotImplemented
+ if _dtype_is_int(self) and _dtype_is_int(other):
+ raise div_ambiguity_error()
+ return _prim_ffi_api._OpDiv(self, other, None) # type: ignore
+
+ def __rdiv__(self, other: Expr) -> Expr:
+ if not _is_scalar_operand(other):
+ return NotImplemented
+ if _dtype_is_int(self) and _dtype_is_int(other):
+ raise div_ambiguity_error()
+ return _prim_ffi_api._OpDiv(other, self, None) # type: ignore
+
+ def __truediv__(self, other: Expr) -> Expr:
+ if not _is_scalar_operand(other):
+ return NotImplemented
+ if _dtype_is_int(self) and _dtype_is_int(other):
+ raise div_ambiguity_error()
+ return _prim_ffi_api._OpDiv(self, other, None) # type: ignore
+
+ def __rtruediv__(self, other: Expr) -> Expr:
+ if not _is_scalar_operand(other):
+ return NotImplemented
+ if _dtype_is_int(self) and _dtype_is_int(other):
+ raise div_ambiguity_error()
+ return _prim_ffi_api._OpDiv(other, self, None) # type: ignore
-def __add__(_lhs, _rhs):
- return NotImplemented
+ def __floordiv__(self, other: Expr) -> Expr:
+ return _prim_ffi_api._OpFloorDiv(self, other, None) # type: ignore
+
+ def __rfloordiv__(self, other: Expr) -> Expr:
+ return _prim_ffi_api._OpFloorDiv(other, self, None) # type: ignore
+
+ def __mod__(self, other: Expr) -> Expr:
+ return _prim_ffi_api._OpFloorMod(self, other, None) # type: ignore
+ def __rmod__(self, other: Expr) -> Expr:
+ return _prim_ffi_api._OpFloorMod(other, self, None) # type: ignore
+
+ def __neg__(self) -> Expr:
+ neg_one = const(-1, self.expr_ty().dtype)
+ return self.__mul__(neg_one)
-def __radd__(_lhs, _rhs):
- return NotImplemented
+ def __lshift__(self, other: Expr) -> Expr:
+ return _prim_ffi_api.left_shift(self, other, None) # type: ignore
+ def __rlshift__(self, other: Expr) -> Expr:
+ return _prim_ffi_api.left_shift(other, self, None) # type: ignore
-def __sub__(_lhs, _rhs):
- return NotImplemented
+ def __rshift__(self, other: Expr) -> Expr:
+ return _prim_ffi_api.right_shift(self, other, None) # type: ignore
+ def __rrshift__(self, other: Expr) -> Expr:
+ return _prim_ffi_api.right_shift(other, self, None) # type: ignore
-def __rsub__(_lhs, _rhs):
- return NotImplemented
+ def __and__(self, other: Expr) -> Expr:
+ return _prim_ffi_api.bitwise_and(self, other, None) # type: ignore
+ def __rand__(self, other: Expr) -> Expr:
+ return _prim_ffi_api.bitwise_and(other, self, None) # type: ignore
-def __mul__(_lhs, _rhs):
- return NotImplemented
+ def __or__(self, other: Expr) -> Expr:
+ return _prim_ffi_api.bitwise_or(self, other, None) # type: ignore
+ def __ror__(self, other: Expr) -> Expr:
+ return _prim_ffi_api.bitwise_or(other, self, None) # type: ignore
-def __rmul__(_lhs, _rhs):
- return NotImplemented
+ def __xor__(self, other: Expr) -> Expr:
+ return _prim_ffi_api.bitwise_xor(self, other, None) # type: ignore
+ def __rxor__(self, other: Expr) -> Expr:
+ return _prim_ffi_api.bitwise_xor(other, self, None) # type: ignore
-def __div__(_lhs, _rhs):
- return NotImplemented
+ def __invert__(self) -> Expr:
+ if _dtype_is_float(self):
+ raise RuntimeError("Cannot use ~ operator on float type Expr.")
+ return _prim_ffi_api.bitwise_not(self, None) # type: ignore
+ def __lt__(self, other: Expr) -> Expr:
+ return _prim_ffi_api._OpLT(self, other, None) # type: ignore
-def __rdiv__(_lhs, _rhs):
- return NotImplemented
+ def __le__(self, other: Expr) -> Expr:
+ return _prim_ffi_api._OpLE(self, other, None) # type: ignore
+ def __eq__(self, other: Expr) -> Expr:
+ return EqualOp(self, other)
-def __truediv__(_lhs, _rhs):
- return NotImplemented
+ def __ne__(self, other: Expr) -> Expr:
+ return NotEqualOp(self, other)
+ def __gt__(self, other: Expr) -> Expr:
+ return _prim_ffi_api._OpGT(self, other, None) # type: ignore
-def __rtruediv__(_lhs, _rhs):
- return NotImplemented
+ def __ge__(self, other: Expr) -> Expr:
+ return _prim_ffi_api._OpGE(self, other, None) # type: ignore
+ def __nonzero__(self):
+ raise ValueError(
+ "Cannot use and / or / not operator to Expr, hint: use tvm.tirx.all / "
+ "tvm.tirx.any, if it is None checking, use node is not None"
+ )
-def __floordiv__(_lhs, _rhs):
- return NotImplemented
+ def __bool__(self) -> bool:
+ return self.__nonzero__()
+ def equal(self, other: Expr, span: Span | None = None) -> bool:
+ """Build an equal check expression with other expr.
-def __rfloordiv__(_lhs, _rhs):
- return NotImplemented
+ Parameters
+ ----------
+ other : Expr
+ The other expression
+ span : Optional[Span]
+ The location of the cast in the source.
-def __mod__(_lhs, _rhs):
- return NotImplemented
+ Returns
+ -------
+ ret : Expr
+ The equality expression.
+ """
+ return _prim_ffi_api._OpEQ(self, other, span) # type: ignore
+ def astype(self, dtype: str | PrimType, span: Span | None = None) -> Expr:
+ """Cast the expression to other type.
-def __rmod__(_lhs, _rhs):
- return NotImplemented
+ Parameters
+ ----------
+ dtype : str
+ The type of new expression
+ span : Optional[Span]
+ The location of the cast in the source.
-def __neg__(_value):
- return NotImplemented
+ Returns
+ -------
+ expr : Expr
+ Expression with new type
+ """
+ return _prim_ffi_api._cast(dtype, self, span) # type: ignore
-def __lshift__(_lhs, _rhs):
- return NotImplemented
+class EqualOp(ObjectConvertible, ExprOp):
+ """Deferred equal operator.
+ This is used to support sugar that a == b can either
+ mean Object.same_as or Object.equal.
-def __rlshift__(_lhs, _rhs):
- return NotImplemented
+ Parameters
+ ----------
+ a : Expr
+ Left operand.
+ b : Expr
+ Right operand.
-def __rshift__(_lhs, _rhs):
- return NotImplemented
-
-
-def __rrshift__(_lhs, _rhs):
- return NotImplemented
-
-
-def __and__(_lhs, _rhs):
- return NotImplemented
-
-
-def __rand__(_lhs, _rhs):
- return NotImplemented
-
-
-def __or__(_lhs, _rhs):
- return NotImplemented
-
-
-def __ror__(_lhs, _rhs):
- return NotImplemented
-
-
-def __xor__(_lhs, _rhs):
- return NotImplemented
-
-
-def __rxor__(_lhs, _rhs):
- return NotImplemented
-
-
-def __invert__(_value):
- return NotImplemented
-
-
-def __lt__(_lhs, _rhs):
- return NotImplemented
-
-
-def __le__(_lhs, _rhs):
- return NotImplemented
-
-
-def __eq__(_lhs, _rhs):
- return NotImplemented
-
-
-def __ne__(_lhs, _rhs):
- return NotImplemented
-
-
-def __gt__(_lhs, _rhs):
- return NotImplemented
-
-
-def __ge__(_lhs, _rhs):
- return NotImplemented
-
-
-def equal(_lhs, _rhs, _span=None):
- return NotImplemented
-
-
-def astype(_value, _dtype, _span=None):
- return NotImplemented
+ span : Optional[Span]
+ The location of the cast in the source.
+ """
+
+ # This class is not manipulated by C++. So use python's identity check function is sufficient
+ same_as = object.__eq__
+
+ def __init__(self, a: Expr, b: Expr, span: Span | None = None):
+ self.a = a
+ self.b = b
+ self.span = span
+
+ def __nonzero__(self) -> bool:
+ return self.a.same_as(self.b)
+
+ def __bool__(self) -> bool:
+ return self.__nonzero__()
+
+ def asobject(self) -> Expr:
+ """Convert object."""
+ return _prim_ffi_api._OpEQ(self.a, self.b, self.span) # type: ignore
+
+ def expr_ty(self) -> PrimType:
+ """Compile-time type of the equality result."""
+ return PrimType("bool")
+
+ def __repr__(self) -> str:
+ return f"EqualOp({self.a!r}, {self.b!r})"
+
+
+class NotEqualOp(ObjectConvertible, ExprOp):
+ """Deferred NE operator.
+
+ This is used to support sugar that a != b can either
+ mean not Object.same_as or make.NE.
+
+ Parameters
+ ----------
+ a : Expr
+ Left operand.
+
+ b : Expr
+ Right operand.
+
+ span : Optional[Span]
+ The location of the cast in the source.
+ """
+
+ # This class is not manipulated by C++. So use python's identity check function is sufficient
+ same_as = object.__eq__
+
+ def __init__(self, a: Expr, b: Expr, span: Span | None = None) -> None:
+ self.a = a
+ self.b = b
+ self.span = span
+
+ def __nonzero__(self) -> bool:
+ return not self.a.same_as(self.b)
+
+ def __bool__(self) -> bool:
+ return self.__nonzero__()
+
+ def asobject(self) -> Expr:
+ """Convert object."""
+ return _prim_ffi_api._OpNE(self.a, self.b, self.span) # type: ignore
+
+ def expr_ty(self) -> PrimType:
+ """Compile-time type of the inequality result."""
+ return PrimType("bool")
+
+ def __repr__(self) -> str:
+ return f"NotEqualOp({self.a!r}, {self.b!r})"
+
+
+__add__ = ExprOp.__add__
+__radd__ = ExprOp.__radd__
+__sub__ = ExprOp.__sub__
+__rsub__ = ExprOp.__rsub__
+__mul__ = ExprOp.__mul__
+__rmul__ = ExprOp.__rmul__
+__div__ = ExprOp.__div__
+__rdiv__ = ExprOp.__rdiv__
+__truediv__ = ExprOp.__truediv__
+__rtruediv__ = ExprOp.__rtruediv__
+__floordiv__ = ExprOp.__floordiv__
+__rfloordiv__ = ExprOp.__rfloordiv__
+__mod__ = ExprOp.__mod__
+__rmod__ = ExprOp.__rmod__
+__neg__ = ExprOp.__neg__
+__lshift__ = ExprOp.__lshift__
+__rlshift__ = ExprOp.__rlshift__
+__rshift__ = ExprOp.__rshift__
+__rrshift__ = ExprOp.__rrshift__
+__and__ = ExprOp.__and__
+__rand__ = ExprOp.__rand__
+__or__ = ExprOp.__or__
+__ror__ = ExprOp.__ror__
+__xor__ = ExprOp.__xor__
+__rxor__ = ExprOp.__rxor__
+__invert__ = ExprOp.__invert__
+__lt__ = ExprOp.__lt__
+__le__ = ExprOp.__le__
+__eq__ = ExprOp.__eq__
+__ne__ = ExprOp.__ne__
+__gt__ = ExprOp.__gt__
+__ge__ = ExprOp.__ge__
+equal = ExprOp.equal
+astype = ExprOp.astype
diff --git a/python/tvm/ir/expr.py b/python/tvm/ir/expr.py
index 7bcddca35e..2ae414f758 100644
--- a/python/tvm/ir/expr.py
+++ b/python/tvm/ir/expr.py
@@ -22,8 +22,8 @@ import tvm_ffi
import tvm
-from ..runtime import Object, Scriptable
-from . import _ffi_api, _overload_prim_expr, _tensor_expr_overload
+from ..runtime import Object, Scriptable, const
+from . import _ffi_api, _tensor_expr_overload
from .base import Node, Span
@@ -33,7 +33,7 @@ def _convert_subscript_index(index):
def convert(value):
if value is None or is_prim_expr(value):
return value
- return tvm.tirx.const(value)
+ return const(value)
if isinstance(index, slice):
return (convert(index.start), convert(index.stop), convert(index.step))
@@ -612,3 +612,8 @@ class Range(Node, Scriptable):
def __ne__(self, other: Object) -> bool:
return not self.__eq__(other)
+
+
+# Primitive overloads also initialize the concrete primitive nodes, whose bases
+# must be defined before importing them.
+from . import _overload_prim_expr # noqa: E402 # pylint: disable=wrong-import-position
diff --git a/python/tvm/ir/prim/__init__.py b/python/tvm/ir/prim/__init__.py
index f554ec655d..9a8f388efd 100644
--- a/python/tvm/ir/prim/__init__.py
+++ b/python/tvm/ir/prim/__init__.py
@@ -14,10 +14,46 @@
# KIND, either express or implied. See the License for the
# specific language governing permissions and limitations
# under the License.
-"""Primitive expression nodes shared by TVM IR dialects."""
+# pylint: disable=unused-import
+"""Primitive expressions and construction helpers shared by TVM IR dialects."""
+from ...runtime import const
from ..expr import Expr
from . import _ffi_api
+from .expr import (
+ EQ,
+ GE,
+ GT,
+ LE,
+ LT,
+ NE,
+ Add,
+ And,
+ BinaryOpExpr,
+ Broadcast,
+ Cast,
+ CmpExpr,
+ ConstExpr,
+ Div,
+ FloatImm,
+ FloorDiv,
+ FloorMod,
+ IntImm,
+ Let,
+ LogicalExpr,
+ Max,
+ Min,
+ Mod,
+ Mul,
+ Not,
+ Or,
+ Ramp,
+ Select,
+ Shuffle,
+ StringImm,
+ Sub,
+)
+from .op import convert, max_value, min_value
def expr_deep_equal(lhs: Expr, rhs: Expr) -> bool:
@@ -55,48 +91,3 @@ def expr_deep_equal(lhs: Expr, rhs: Expr) -> bool:
tvm_ffi.structural_equal
"""
return _ffi_api.expr_deep_equal(lhs, rhs) # type: ignore
-
-
-_EXPR_NAMES = {
- "StringImm",
- "Cast",
- "Add",
- "Sub",
- "Mul",
- "Div",
- "Mod",
- "FloorDiv",
- "FloorMod",
- "Min",
- "Max",
- "EQ",
- "NE",
- "LT",
- "LE",
- "GT",
- "GE",
- "And",
- "Or",
- "Not",
- "Select",
- "Let",
- "Ramp",
- "Broadcast",
- "Shuffle",
-}
-
-
-def __getattr__(name):
- # Keep the historical tvm.tirx classes as the single Python definitions
- # during this mechanical C++ ownership move. Importing lazily avoids the
- # tvm.ir <-> tvm.tirx initialization cycle while exposing the new public
- # tvm.ir.prim spelling.
- if name in _EXPR_NAMES:
- from tvm.tirx import expr # pylint: disable=import-outside-toplevel
-
- return getattr(expr, name)
- raise AttributeError(name)
-
-
-def __dir__():
- return sorted(set(globals()) | _EXPR_NAMES)
diff --git a/python/tvm/ir/prim/_ffi_api.py b/python/tvm/ir/prim/_ffi_api.py
index 6bbd068200..8b2fb8894c 100644
--- a/python/tvm/ir/prim/_ffi_api.py
+++ b/python/tvm/ir/prim/_ffi_api.py
@@ -18,4 +18,4 @@
import tvm_ffi
-tvm_ffi.init_ffi_api("ir.prim", __name__)
+tvm_ffi.init_ffi_api("prim", __name__)
diff --git a/python/tvm/ir/prim/expr.py b/python/tvm/ir/prim/expr.py
new file mode 100644
index 0000000000..ba07c5f289
--- /dev/null
+++ b/python/tvm/ir/prim/expr.py
@@ -0,0 +1,699 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements. See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership. The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied. See the License for the
+# specific language governing permissions and limitations
+# under the License.
+"""Primitive expression nodes shared by TVM IR dialects."""
+
+import tvm_ffi
+
+from .. import _ffi_api as _ir_ffi_api
+from ..base import Span
+from ..expr import Expr, ExprWithOp, Var
+from ..type import PrimType
+from . import _ffi_api as _prim_ffi_api
+
+
+class ConstExpr(ExprWithOp):
+ pass
+
+
+class BinaryOpExpr(ExprWithOp):
+ a: Expr
+ b: Expr
+
+
+class CmpExpr(ExprWithOp):
+ a: Expr
+ b: Expr
+
+
+class LogicalExpr(ExprWithOp):
+ pass
+
+
+@tvm_ffi.register_object("ir.FloatImm")
+class FloatImm(ConstExpr):
+ """Float constant.
+
+ Parameters
+ ----------
+ dtype : str
+ The data type
+
+ value : float
+ The constant value.
+
+ span : Optional[Span]
+ The location of this expression in the source code.
+ """
+
+ value: float
+
+ def __init__(self, dtype: str | PrimType, value: float, span: Span | None = None) -> None:
+ if isinstance(dtype, PrimType):
+ dtype = dtype.dtype
+ self.__init_handle_by_constructor__(
+ _ir_ffi_api.FloatImm,
+ dtype,
+ value,
+ span, # type: ignore
+ )
+
+ def __float__(self) -> float:
+ return self.value
+
+
+@tvm_ffi.register_object("ir.IntImm")
+class IntImm(ConstExpr):
+ """Int constant.
+
+ Parameters
+ ----------
+ dtype : str
+ The data type
+
+ value : int
+ The constant value.
+
+ span : Optional[Span]
+ The location of this expression in the source code.
+ """
+
+ value: int
+
+ def __init__(self, dtype: str | PrimType, value: int, span: Span | None = None) -> None:
+ if isinstance(dtype, PrimType):
+ dtype = dtype.dtype
+ self.__init_handle_by_constructor__(
+ _ir_ffi_api.IntImm,
+ dtype,
+ value,
+ span, # type: ignore
+ )
+
+ def __hash__(self) -> int:
+ return self.value
+
+ def __int__(self) -> int:
+ return self.value
+
+ def __nonzero__(self) -> bool:
+ return self.value != 0
+
+ def __eq__(self, other: Expr) -> Expr:
+ return _prim_ffi_api._OpEQ(self, other, None) # type: ignore
+
+ def __ne__(self, other: Expr) -> Expr:
+ return _prim_ffi_api._OpNE(self, other, None) # type: ignore
+
+ def __bool__(self) -> bool:
+ return self.__nonzero__()
+
+
+@tvm_ffi.register_object("prim.StringImm") # type: ignore
+class StringImm(ConstExpr):
+ """String constant.
+
+ Parameters
+ ----------
+ value : str
+ The value of the function.
+
+ span : Optional[Span]
+ The location of this expression in the source code.
+ """
+
+ value: str
+
+ def __init__(self, value: str, span: Span | None = None) -> None:
+ self.__init_handle_by_constructor__(_prim_ffi_api.StringImm, value, span) # type: ignore
+
+ def __eq__(self, other: Expr) -> bool:
+ if isinstance(other, ConstExpr):
+ return self.value == other.value
+ return self.value == other
+
+ def __ne__(self, other: Expr) -> bool:
+ if isinstance(other, ConstExpr):
+ return self.value != other.value
+ return self.value != other
+
+ def __hash__(self) -> int:
+ return Expr.__hash__(self)
+
+
+@tvm_ffi.register_object("prim.Cast")
+class Cast(ExprWithOp):
+ """Cast expression.
+
+ Parameters
+ ----------
+ dtype : str
+ The data type
+
+ value : Expr
+ The value of the function.
+
+ span : Optional[Span]
+ The location of this expression in the source code.
+ """
+
+ value: Expr
+
+ def __init__(self, dtype: str | PrimType, value, span: Span | None = None) -> None:
+ if isinstance(dtype, PrimType):
+ dtype = dtype.dtype
+ self.__init_handle_by_constructor__(_prim_ffi_api.Cast, dtype, value, span) # type: ignore
+
+
+@tvm_ffi.register_object("prim.Add")
+class Add(BinaryOpExpr):
+ """Add node.
+
+ Parameters
+ ----------
+ a : Expr
+ The left hand operand.
+
+ b : Expr
+ The right hand operand.
+
+ span : Optional[Span]
+ The location of this expression in the source code.
+ """
+
+ def __init__(self, a: Expr, b: Expr, span: Span | None = None) -> None:
+ self.__init_handle_by_constructor__(_prim_ffi_api.Add, a, b, span) # type: ignore
+
+
+@tvm_ffi.register_object("prim.Sub")
+class Sub(BinaryOpExpr):
+ """Sub node.
+
+ Parameters
+ ----------
+ a : Expr
+ The left hand operand.
+
+ b : Expr
+ The right hand operand.
+
+ span : Optional[Span]
+ The location of this expression in the source code.
+ """
+
+ def __init__(self, a: Expr, b: Expr, span: Span | None = None) -> None:
+ self.__init_handle_by_constructor__(_prim_ffi_api.Sub, a, b, span) # type: ignore
+
+
+@tvm_ffi.register_object("prim.Mul")
+class Mul(BinaryOpExpr):
+ """Mul node.
+
+ Parameters
+ ----------
+ a : Expr
+ The left hand operand.
+
+ b : Expr
+ The right hand operand.
+
+ span : Optional[Span]
+ The location of this expression in the source code.
+ """
+
+ def __init__(self, a: Expr, b: Expr, span: Span | None = None) -> None:
+ self.__init_handle_by_constructor__(_prim_ffi_api.Mul, a, b, span) # type: ignore
+
+
+@tvm_ffi.register_object("prim.Div")
+class Div(BinaryOpExpr):
+ """Div node.
+
+ Parameters
+ ----------
+ a : Expr
+ The left hand operand.
+
+ b : Expr
+ The right hand operand.
+
+ span : Optional[Span]
+ The location of this expression in the source code.
+ """
+
+ def __init__(self, a: Expr, b: Expr, span: Span | None = None) -> None:
+ self.__init_handle_by_constructor__(_prim_ffi_api.Div, a, b, span) # type: ignore
+
+
+@tvm_ffi.register_object("prim.Mod")
+class Mod(BinaryOpExpr):
+ """Mod node.
+
+ Parameters
+ ----------
+ a : Expr
+ The left hand operand.
+
+ b : Expr
+ The right hand operand.
+
+ span : Optional[Span]
+ The location of this expression in the source code.
+ """
+
+ def __init__(self, a: Expr, b: Expr, span: Span | None = None) -> None:
+ self.__init_handle_by_constructor__(_prim_ffi_api.Mod, a, b, span) # type: ignore
+
+
+@tvm_ffi.register_object("prim.FloorDiv")
+class FloorDiv(BinaryOpExpr):
+ """FloorDiv node.
+
+ Parameters
+ ----------
+ a : Expr
+ The left hand operand.
+
+ b : Expr
+ The right hand operand.
+
+ span : Optional[Span]
+ The location of this expression in the source code.
+ """
+
+ def __init__(self, a: Expr, b: Expr, span: Span | None = None) -> None:
+ self.__init_handle_by_constructor__(_prim_ffi_api.FloorDiv, a, b, span) # type: ignore
+
+
+@tvm_ffi.register_object("prim.FloorMod")
+class FloorMod(BinaryOpExpr):
+ """FloorMod node.
+
+ Parameters
+ ----------
+ a : Expr
+ The left hand operand.
+
+ b : Expr
+ The right hand operand.
+
+ span : Optional[Span]
+ The location of this expression in the source code.
+ """
+
+ def __init__(self, a: Expr, b: Expr, span: Span | None = None) -> None:
+ self.__init_handle_by_constructor__(_prim_ffi_api.FloorMod, a, b, span) # type: ignore
+
+
+@tvm_ffi.register_object("prim.Min")
+class Min(BinaryOpExpr):
+ """Min node.
+
+ Parameters
+ ----------
+ a : Expr
+ The left hand operand.
+
+ b : Expr
+ The right hand operand.
+
+ span : Optional[Span]
+ The location of this expression in the source code.
+ """
+
+ def __init__(self, a: Expr, b: Expr, span: Span | None = None) -> None:
+ self.__init_handle_by_constructor__(_prim_ffi_api.Min, a, b, span) # type: ignore
+
+
+@tvm_ffi.register_object("prim.Max")
+class Max(BinaryOpExpr):
+ """Max node.
+
+ Parameters
+ ----------
+ a : Expr
+ The left hand operand.
+
+ b : Expr
+ The right hand operand.
+
+ span : Optional[Span]
+ The location of this expression in the source code.
+ """
+
+ def __init__(self, a: Expr, b: Expr, span: Span | None = None) -> None:
+ self.__init_handle_by_constructor__(_prim_ffi_api.Max, a, b, span) # type: ignore
+
+
+@tvm_ffi.register_object("prim.EQ")
+class EQ(CmpExpr):
+ """EQ node.
+
+ Parameters
+ ----------
+ a : Expr
+ The left hand operand.
+
+ b : Expr
+ The right hand operand.
+
+ span : Optional[Span]
+ The location of this expression in the source code.
+ """
+
+ def __init__(self, a: Expr, b: Expr, span: Span | None = None) -> None:
+ self.__init_handle_by_constructor__(_prim_ffi_api.EQ, a, b, span) # type: ignore
+
+
+@tvm_ffi.register_object("prim.NE")
+class NE(CmpExpr):
+ """NE node.
+
+ Parameters
+ ----------
+ a : Expr
+ The left hand operand.
+
+ b : Expr
+ The right hand operand.
+
+ span : Optional[Span]
+ The location of this expression in the source code.
+ """
+
+ def __init__(self, a: Expr, b: Expr, span: Span | None = None) -> None:
+ self.__init_handle_by_constructor__(_prim_ffi_api.NE, a, b, span) # type: ignore
+
+
+@tvm_ffi.register_object("prim.LT")
+class LT(CmpExpr):
+ """LT node.
+
+ Parameters
+ ----------
+ a : Expr
+ The left hand operand.
+
+ b : Expr
+ The right hand operand.
+
+ span : Optional[Span]
+ The location of this expression in the source code.
+ """
+
+ def __init__(self, a: Expr, b: Expr, span: Span | None = None) -> None:
+ self.__init_handle_by_constructor__(_prim_ffi_api.LT, a, b, span) # type: ignore
+
+
+@tvm_ffi.register_object("prim.LE")
+class LE(CmpExpr):
+ """LE node.
+
+ Parameters
+ ----------
+ a : Expr
+ The left hand operand.
+
+ b : Expr
+ The right hand operand.
+
+ span : Optional[Span]
+ The location of this expression in the source code.
+ """
+
+ def __init__(self, a: Expr, b: Expr, span: Span | None = None) -> None:
+ self.__init_handle_by_constructor__(_prim_ffi_api.LE, a, b, span) # type: ignore
+
+
+@tvm_ffi.register_object("prim.GT")
+class GT(CmpExpr):
+ """GT node.
+
+ Parameters
+ ----------
+ a : Expr
+ The left hand operand.
+
+ b : Expr
+ The right hand operand.
+
+ span : Optional[Span]
+ The location of this expression in the source code.
+ """
+
+ def __init__(self, a: Expr, b: Expr, span: Span | None = None) -> None:
+ self.__init_handle_by_constructor__(_prim_ffi_api.GT, a, b, span) # type: ignore
+
+
+@tvm_ffi.register_object("prim.GE")
+class GE(CmpExpr):
+ """GE node.
+
+ Parameters
+ ----------
+ a : Expr
+ The left hand operand.
+
+ b : Expr
+ The right hand operand.
+
+ span : Optional[Span]
+ The location of this expression in the source code.
+ """
+
+ def __init__(self, a: Expr, b: Expr, span: Span | None = None) -> None:
+ self.__init_handle_by_constructor__(_prim_ffi_api.GE, a, b, span) # type: ignore
+
+
+@tvm_ffi.register_object("prim.And")
+class And(LogicalExpr):
+ """And node.
+
+ Parameters
+ ----------
+ a : Expr
+ The left hand operand.
+
+ b : Expr
+ The right hand operand.
+
+ span : Optional[Span]
+ The location of this expression in the source code.
+ """
+
+ def __init__(self, a: Expr, b: Expr, span: Span | None = None) -> None:
+ self.__init_handle_by_constructor__(_prim_ffi_api.And, a, b, span) # type: ignore
+
+
+@tvm_ffi.register_object("prim.Or")
+class Or(LogicalExpr):
+ """Or node.
+
+ Parameters
+ ----------
+ a : Expr
+ The left hand operand.
+
+ b : Expr
+ The right hand operand.
+
+ span : Optional[Span]
+ The location of this expression in the source code.
+ """
+
+ a: Expr
+ b: Expr
+
+ def __init__(self, a: Expr, b: Expr, span: Span | None = None) -> None:
+ self.__init_handle_by_constructor__(_prim_ffi_api.Or, a, b, span) # type: ignore
+
+
+@tvm_ffi.register_object("prim.Not")
+class Not(LogicalExpr):
+ """Not node.
+
+ Parameters
+ ----------
+ a : Expr
+ The input value
+
+ span : Optional[Span]
+ The location of this expression in the source code.
+ """
+
+ a: Expr
+
+ def __init__(self, a: Expr, span: Span | None = None) -> None:
+ self.__init_handle_by_constructor__(_prim_ffi_api.Not, a, span) # type: ignore
+
+
+@tvm_ffi.register_object("prim.Select")
+class Select(ExprWithOp):
+ """Select node.
+
+ Note
+ ----
+ Select may compute both true_value and false_value.
+ Use :py:class:`tvm.tirx.if_then_else` instead if you want to
+ get a conditional expression that only evaluates
+ the correct branch.
+
+ Parameters
+ ----------
+ condition : Expr
+ The condition expression.
+
+ true_value : Expr
+ The value to take when condition is true.
+
+ false_value : Expr
+ The value to take when condition is false.
+
+ span : Optional[Span]
+ The location of this expression in the source code.
+ """
+
+ condition: Expr
+ true_value: Expr
+ false_value: Expr
+
+ def __init__(
+ self,
+ condition: Expr,
+ true_value: Expr,
+ false_value: Expr,
+ span: Span | None = None,
+ ) -> None:
+ if isinstance(condition, bool):
+ condition = IntImm("bool", condition)
+ self.__init_handle_by_constructor__(
+ _prim_ffi_api.Select,
+ condition,
+ true_value,
+ false_value,
+ span, # type: ignore
+ )
+
+
+@tvm_ffi.register_object("prim.Ramp")
+class Ramp(ExprWithOp):
+ """Ramp node.
+
+ Parameters
+ ----------
+ base : Expr
+ The base expression.
+
+ stride : Expr
+ The stride of the ramp.
+
+ lanes : Expr
+ The lanes of the expression.
+
+ span : Optional[Span]
+ The location of this expression in the source code.
+ """
+
+ base: Expr
+ stride: Expr
+ lanes: Expr
+
+ def __init__(self, base: Expr, stride: Expr, lanes: Expr, span: Span | None = None) -> None:
+ self.__init_handle_by_constructor__(
+ _prim_ffi_api.Ramp,
+ base,
+ stride,
+ lanes,
+ span, # type: ignore
+ )
+
+
+@tvm_ffi.register_object("prim.Broadcast")
+class Broadcast(ExprWithOp):
+ """Broadcast node.
+
+ Parameters
+ ----------
+ value : Expr
+ The value of the expression.
+
+ lanes : Expr
+ The lanes of the expression.
+
+ span : Optional[Span]
+ The location of this expression in the source code.
+ """
+
+ value: Expr
+ lanes: Expr
+
+ def __init__(self, value: Expr, lanes: Expr, span: Span | None = None) -> None:
+ self.__init_handle_by_constructor__(_prim_ffi_api.Broadcast, value, lanes, span) # type: ignore
+
+
+@tvm_ffi.register_object("prim.Shuffle")
+class Shuffle(ExprWithOp):
+ """Shuffle node.
+
+ Parameters
+ ----------
+ vectors : List[Expr]
+ The vectors
+
+ indices : List[Expr]
+ The indices
+
+ span : Optional[Span]
+ The location of this expression in the source code.
+ """
+
+ vectors: list[Expr]
+ indices: list[Expr]
+
+ def __init__(self, vectors: list[Expr], indices: list[Expr], span: Span | None = None) -> None:
+ self.__init_handle_by_constructor__(
+ _prim_ffi_api.Shuffle,
+ vectors,
+ indices,
+ span, # type: ignore
+ )
+
+
+@tvm_ffi.register_object("prim.Let")
+class Let(ExprWithOp):
+ """Let node.
+
+ Parameters
+ ----------
+ var : Var
+ The variable in the binding.
+
+ value : Expr
+ The value in to be bound.
+
+ body : Expr
+ The body expression.
+
+ span : Optional[Span]
+ The location of this expression in the source code.
+ """
+
+ var: Var
+ value: Expr
+ body: Expr
+
+ def __init__(self, var: Var, value: Expr, body: Expr, span: Span | None = None) -> None:
+ self.__init_handle_by_constructor__(_prim_ffi_api.Let, var, value, body, span) # type: ignore
diff --git a/python/tvm/arith/pattern.py b/python/tvm/ir/prim/op.py
similarity index 50%
copy from python/tvm/arith/pattern.py
copy to python/tvm/ir/prim/op.py
index 8582b2f576..c56ca21ace 100644
--- a/python/tvm/arith/pattern.py
+++ b/python/tvm/ir/prim/op.py
@@ -14,48 +14,53 @@
# KIND, either express or implied. See the License for the
# specific language governing permissions and limitations
# under the License.
-"""Detect common patterns."""
+"""Construction helpers for shared primitive expressions."""
+from typing import Any
+
+from ..base import Span
+from ..expr import Expr
from . import _ffi_api
-def detect_linear_equation(expr, var_list):
- """Match `expr = sum_{i=0}^{n-1} var[i] * coeff[i] + coeff[n]`
+def convert(expr) -> Expr:
+ """Convert a scalar or sequence to primitive expressions."""
+ return _ffi_api.convert(expr)
+
- Where coeff[i] and base are invariant of var[j] for all i and j.
+def min_value(dtype, span=None):
+ """minimum value of dtype
Parameters
----------
- expr : Expr
- The expression to be matched.
+ dtype : str
+ The data type.
- var_list : List[tvm.tirx.Var]
- A list of variables.
+ span : Optional[Span]
+ The location of this operator in the source code.
Returns
-------
- coeff : List[Expr]
- A list of co-efficients if the match is successful.
- An empty list if the match failed.
+ value : tvm.ir.Expr
+ The minimum value of dtype.
"""
- return _ffi_api.DetectLinearEquation(expr, var_list)
+ return _ffi_api.min_value(dtype, span) # type: ignore
-def detect_clip_bound(expr, var_list):
- """Detect if expression corresponds to clip bound of the vars
+def max_value(dtype: str, span: Span | None = None) -> Any:
+ """maximum value of dtype
Parameters
----------
- expr : Expr
- The expression to be matched.
+ dtype : str
+ The data type.
- var_list : List[tvm.tirx.Var]
- A list of variables.
+ span : Optional[Span]
+ The location of this operator in the source code.
Returns
-------
- coeff : List[Expr]
- `concat([min_value[i], max_value[i]] for i, v in enumerate(var_list))`
- An empty list if the match failed.
+ value : tvm.ir.Expr
+ The maximum value of dtype.
"""
- return _ffi_api.DetectClipBound(expr, var_list)
+ return _ffi_api.max_value(dtype, span) # type: ignore
diff --git a/python/tvm/tirx/expr.py b/python/tvm/tirx/expr.py
index 6c1b6cde94..ba119388ec 100644
--- a/python/tvm/tirx/expr.py
+++ b/python/tvm/tirx/expr.py
@@ -32,13 +32,53 @@ import functools
import tvm_ffi
-import tvm.ir._ffi_api
-import tvm.ir._overload_prim_expr as _overload_prim_expr
-import tvm.ir.prim._ffi_api as _prim_ffi_api
from tvm import ir
from tvm.ir import Expr
+from tvm.ir._overload_prim_expr import ( # noqa: F401
+ EqualOp,
+ ExprOp,
+ NotEqualOp,
+ div_ambiguity_error,
+)
from tvm.ir.base import Span
-from tvm.runtime import DataTypeCode, Object, ObjectConvertible, Scriptable, const
+from tvm.ir.prim import convert as convert
+
+# Retain historical imports as aliases of the canonical shared definitions.
+from tvm.ir.prim.expr import ( # noqa: F401
+ EQ,
+ GE,
+ GT,
+ LE,
+ LT,
+ NE,
+ Add,
+ And,
+ BinaryOpExpr,
+ Broadcast,
+ Cast,
+ CmpExpr,
+ ConstExpr,
+ Div,
+ FloatImm,
+ FloorDiv,
+ FloorMod,
+ IntImm,
+ Let,
+ LogicalExpr,
+ Max,
+ Min,
+ Mod,
+ Mul,
+ Not,
+ Or,
+ Ramp,
+ Select,
+ Shuffle,
+ StringImm,
+ Sub,
+)
+from tvm.runtime import Object, ObjectConvertible, Scriptable
+from tvm.runtime import const as const
from . import _ffi_api
from .buffer import Buffer
@@ -50,357 +90,6 @@ def _get_te_constructor(name: str):
return tvm_ffi.get_global_func(f"te.{name}")
-def convert(expr) -> Expr:
- return _ffi_api.convert(expr)
-
-
-def div_ambiguity_error() -> RuntimeError:
- return RuntimeError(
- "TVM supports multiple types of integer divisions, "
- + "please call div, indexdiv/indexmod, floordiv/floormod "
- + " or truncdiv/truncmod directly to avoid ambiguity in the code."
- )
-
-
-def _dtype_is_int(value):
- if isinstance(value, int):
- return True
- if isinstance(value, ExprOp):
- return value.expr_ty().matches_code(DataTypeCode.INT)
- if ir.is_prim_expr(value):
- return value.ty.matches_code(DataTypeCode.INT)
- return False
-
-
-def _dtype_is_float(value):
- if isinstance(value, float):
- return True
- if isinstance(value, ExprOp):
- return value.expr_ty().matches_code(DataTypeCode.FLOAT)
- if ir.is_prim_expr(value):
- return value.ty.matches_code(DataTypeCode.FLOAT)
- return False
-
-
-def _is_scalar_operand(value):
- from .stmt import BufferRegion # pylint: disable=import-outside-toplevel
-
- if isinstance(value, BufferRegion):
- raise TypeError(
- "BufferRegion is not a primitive operand; construct a BufferLoad explicitly"
- )
- return isinstance(value, ExprOp | int | float) or ir.is_prim_expr(value)
-
-
-class ExprOp:
- """Operator overloading for Expr like expressions."""
-
- # TODO(tkonolige): use inspect to add source information to these objects
-
- def expr_ty(self) -> ir.PrimType:
- """Return the compile-time primitive type for expression operators."""
- ty = getattr(self, "ty", None)
- if isinstance(ty, ir.PrimType):
- return ty
- raise TypeError(f"Cannot determine PrimType for {type(self).__name__}")
-
- def __add__(self, other: Expr) -> Expr:
- if not _is_scalar_operand(other):
- return NotImplemented
- return _ffi_api._OpAdd(self, other, None) # type: ignore
-
- def __radd__(self, other: Expr) -> Expr:
- if not _is_scalar_operand(other):
- return NotImplemented
- return _ffi_api._OpAdd(other, self, None) # type: ignore
-
- def __sub__(self, other: Expr) -> Expr:
- if not _is_scalar_operand(other):
- return NotImplemented
- return _ffi_api._OpSub(self, other, None) # type: ignore
-
- def __rsub__(self, other: Expr) -> Expr:
- if not _is_scalar_operand(other):
- return NotImplemented
- return _ffi_api._OpSub(other, self, None) # type: ignore
-
- def __mul__(self, other: Expr) -> Expr:
- if not _is_scalar_operand(other):
- return NotImplemented
- return _ffi_api._OpMul(self, other, None) # type: ignore
-
- def __rmul__(self, other: Expr) -> Expr:
- if not _is_scalar_operand(other):
- return NotImplemented
- return _ffi_api._OpMul(other, self, None) # type: ignore
-
- def __div__(self, other: Expr) -> Expr:
- if not _is_scalar_operand(other):
- return NotImplemented
- if _dtype_is_int(self) and _dtype_is_int(other):
- raise div_ambiguity_error()
- return _ffi_api._OpDiv(self, other, None) # type: ignore
-
- def __rdiv__(self, other: Expr) -> Expr:
- if not _is_scalar_operand(other):
- return NotImplemented
- if _dtype_is_int(self) and _dtype_is_int(other):
- raise div_ambiguity_error()
- return _ffi_api._OpDiv(other, self, None) # type: ignore
-
- def __truediv__(self, other: Expr) -> Expr:
- if not _is_scalar_operand(other):
- return NotImplemented
- if _dtype_is_int(self) and _dtype_is_int(other):
- raise div_ambiguity_error()
- return _ffi_api._OpDiv(self, other, None) # type: ignore
-
- def __rtruediv__(self, other: Expr) -> Expr:
- if not _is_scalar_operand(other):
- return NotImplemented
- if _dtype_is_int(self) and _dtype_is_int(other):
- raise div_ambiguity_error()
- return _ffi_api._OpDiv(other, self, None) # type: ignore
-
- def __floordiv__(self, other: Expr) -> Expr:
- return _ffi_api._OpFloorDiv(self, other, None) # type: ignore
-
- def __rfloordiv__(self, other: Expr) -> Expr:
- return _ffi_api._OpFloorDiv(other, self, None) # type: ignore
-
- def __mod__(self, other: Expr) -> Expr:
- return _ffi_api._OpFloorMod(self, other, None) # type: ignore
-
- def __rmod__(self, other: Expr) -> Expr:
- return _ffi_api._OpFloorMod(other, self, None) # type: ignore
-
- def __neg__(self) -> Expr:
- neg_one = const(-1, self.expr_ty().dtype)
- return self.__mul__(neg_one)
-
- def __lshift__(self, other: Expr) -> Expr:
- return _ffi_api.left_shift(self, other, None) # type: ignore
-
- def __rlshift__(self, other: Expr) -> Expr:
- return _ffi_api.left_shift(other, self, None) # type: ignore
-
- def __rshift__(self, other: Expr) -> Expr:
- return _ffi_api.right_shift(self, other, None) # type: ignore
-
- def __rrshift__(self, other: Expr) -> Expr:
- return _ffi_api.right_shift(other, self, None) # type: ignore
-
- def __and__(self, other: Expr) -> Expr:
- return _ffi_api.bitwise_and(self, other, None) # type: ignore
-
- def __rand__(self, other: Expr) -> Expr:
- return _ffi_api.bitwise_and(other, self, None) # type: ignore
-
- def __or__(self, other: Expr) -> Expr:
- return _ffi_api.bitwise_or(self, other, None) # type: ignore
-
- def __ror__(self, other: Expr) -> Expr:
- return _ffi_api.bitwise_or(other, self, None) # type: ignore
-
- def __xor__(self, other: Expr) -> Expr:
- return _ffi_api.bitwise_xor(self, other, None) # type: ignore
-
- def __rxor__(self, other: Expr) -> Expr:
- return _ffi_api.bitwise_xor(other, self, None) # type: ignore
-
- def __invert__(self) -> Expr:
- if _dtype_is_float(self):
- raise RuntimeError("Cannot use ~ operator on float type Expr.")
- return _ffi_api.bitwise_not(self, None) # type: ignore
-
- def __lt__(self, other: Expr) -> Expr:
- return _ffi_api._OpLT(self, other, None) # type: ignore
-
- def __le__(self, other: Expr) -> Expr:
- return _ffi_api._OpLE(self, other, None) # type: ignore
-
- def __eq__(self, other: Expr) -> Expr:
- return EqualOp(self, other)
-
- def __ne__(self, other: Expr) -> Expr:
- return NotEqualOp(self, other)
-
- def __gt__(self, other: Expr) -> Expr:
- return _ffi_api._OpGT(self, other, None) # type: ignore
-
- def __ge__(self, other: Expr) -> Expr:
- return _ffi_api._OpGE(self, other, None) # type: ignore
-
- def __nonzero__(self):
- raise ValueError(
- "Cannot use and / or / not operator to Expr, hint: use tvm.tirx.all / "
- "tvm.tirx.any, if it is None checking, use node is not None"
- )
-
- def __bool__(self) -> bool:
- return self.__nonzero__()
-
- def equal(self, other: Expr, span: Span | None = None) -> bool:
- """Build an equal check expression with other expr.
-
- Parameters
- ----------
- other : Expr
- The other expression
-
- span : Optional[Span]
- The location of the cast in the source.
-
- Returns
- -------
- ret : Expr
- The equality expression.
- """
- return _ffi_api._OpEQ(self, other, span) # type: ignore
-
- def astype(self, dtype: str | ir.PrimType, span: Span | None = None) -> Expr:
- """Cast the expression to other type.
-
- Parameters
- ----------
- dtype : str
- The type of new expression
-
- span : Optional[Span]
- The location of the cast in the source.
-
- Returns
- -------
- expr : Expr
- Expression with new type
- """
- return _ffi_api._cast(dtype, self, span) # type: ignore
-
-
-_overload_prim_expr.__add__ = ExprOp.__add__
-_overload_prim_expr.__radd__ = ExprOp.__radd__
-_overload_prim_expr.__sub__ = ExprOp.__sub__
-_overload_prim_expr.__rsub__ = ExprOp.__rsub__
-_overload_prim_expr.__mul__ = ExprOp.__mul__
-_overload_prim_expr.__rmul__ = ExprOp.__rmul__
-_overload_prim_expr.__div__ = ExprOp.__div__
-_overload_prim_expr.__rdiv__ = ExprOp.__rdiv__
-_overload_prim_expr.__truediv__ = ExprOp.__truediv__
-_overload_prim_expr.__rtruediv__ = ExprOp.__rtruediv__
-_overload_prim_expr.__floordiv__ = ExprOp.__floordiv__
-_overload_prim_expr.__rfloordiv__ = ExprOp.__rfloordiv__
-_overload_prim_expr.__mod__ = ExprOp.__mod__
-_overload_prim_expr.__rmod__ = ExprOp.__rmod__
-_overload_prim_expr.__neg__ = ExprOp.__neg__
-_overload_prim_expr.__lshift__ = ExprOp.__lshift__
-_overload_prim_expr.__rlshift__ = ExprOp.__rlshift__
-_overload_prim_expr.__rshift__ = ExprOp.__rshift__
-_overload_prim_expr.__rrshift__ = ExprOp.__rrshift__
-_overload_prim_expr.__and__ = ExprOp.__and__
-_overload_prim_expr.__rand__ = ExprOp.__rand__
-_overload_prim_expr.__or__ = ExprOp.__or__
-_overload_prim_expr.__ror__ = ExprOp.__ror__
-_overload_prim_expr.__xor__ = ExprOp.__xor__
-_overload_prim_expr.__rxor__ = ExprOp.__rxor__
-_overload_prim_expr.__invert__ = ExprOp.__invert__
-_overload_prim_expr.__lt__ = ExprOp.__lt__
-_overload_prim_expr.__le__ = ExprOp.__le__
-_overload_prim_expr.__eq__ = ExprOp.__eq__
-_overload_prim_expr.__ne__ = ExprOp.__ne__
-_overload_prim_expr.__gt__ = ExprOp.__gt__
-_overload_prim_expr.__ge__ = ExprOp.__ge__
-_overload_prim_expr.equal = ExprOp.equal
-_overload_prim_expr.astype = ExprOp.astype
-
-
-class EqualOp(ObjectConvertible, ExprOp):
- """Deferred equal operator.
-
- This is used to support sugar that a == b can either
- mean Object.same_as or Object.equal.
-
- Parameters
- ----------
- a : Expr
- Left operand.
-
- b : Expr
- Right operand.
-
- span : Optional[Span]
- The location of the cast in the source.
- """
-
- # This class is not manipulated by C++. So use python's identity check function is sufficient
- same_as = object.__eq__
-
- def __init__(self, a: Expr, b: Expr, span: Span | None = None):
- self.a = a
- self.b = b
- self.span = span
-
- def __nonzero__(self) -> bool:
- return self.a.same_as(self.b)
-
- def __bool__(self) -> bool:
- return self.__nonzero__()
-
- def asobject(self) -> Expr:
- """Convert object."""
- return _ffi_api._OpEQ(self.a, self.b, self.span) # type: ignore
-
- def expr_ty(self) -> ir.PrimType:
- """Compile-time type of the equality result."""
- return ir.PrimType("bool")
-
- def __repr__(self) -> str:
- return f"EqualOp({self.a!r}, {self.b!r})"
-
-
-class NotEqualOp(ObjectConvertible, ExprOp):
- """Deferred NE operator.
-
- This is used to support sugar that a != b can either
- mean not Object.same_as or make.NE.
-
- Parameters
- ----------
- a : Expr
- Left operand.
-
- b : Expr
- Right operand.
-
- span : Optional[Span]
- The location of the cast in the source.
- """
-
- # This class is not manipulated by C++. So use python's identity check function is sufficient
- same_as = object.__eq__
-
- def __init__(self, a: Expr, b: Expr, span: Span | None = None) -> None:
- self.a = a
- self.b = b
- self.span = span
-
- def __nonzero__(self) -> bool:
- return not self.a.same_as(self.b)
-
- def __bool__(self) -> bool:
- return self.__nonzero__()
-
- def asobject(self) -> Expr:
- """Convert object."""
- return _ffi_api._OpNE(self.a, self.b, self.span) # type: ignore
-
- def expr_ty(self) -> ir.PrimType:
- """Compile-time type of the inequality result."""
- return ir.PrimType("bool")
-
- def __repr__(self) -> str:
- return f"NotEqualOp({self.a!r}, {self.b!r})"
-
-
class IntImmEnum(ObjectConvertible):
"""Lazily evaluate an IntImm in case
the constructor is not available in runtime.
@@ -423,24 +112,6 @@ class IntImmEnum(ObjectConvertible):
return IntImm("int32", self.value, self.span) # type: ignore
-class ConstExpr(ir.ExprWithOp):
- pass
-
-
-class BinaryOpExpr(ir.ExprWithOp):
- a: Expr
- b: Expr
-
-
-class CmpExpr(ir.ExprWithOp):
- a: Expr
- b: Expr
-
-
-class LogicalExpr(ir.ExprWithOp):
- pass
-
-
Var = ir.Var
@@ -500,9 +171,9 @@ class IterVar(ExprOp, Object, Scriptable):
if isinstance(dom, list | tuple):
if len(dom) != 2:
raise TypeError("need to be list of ranges")
- dom = tvm.ir.Range(dom[0], dom[1])
+ dom = ir.Range(dom[0], dom[1])
- if not isinstance(dom, tvm.ir.Range):
+ if not isinstance(dom, ir.Range):
raise TypeError("dom need to be Range")
name = var if var is not None else "iter"
@@ -627,552 +298,7 @@ class Reduce(ir.ExprWithOp):
)
-@tvm_ffi.register_object("ir.FloatImm")
-class FloatImm(ConstExpr):
- """Float constant.
-
- Parameters
- ----------
- dtype : str
- The data type
-
- value : float
- The constant value.
-
- span : Optional[Span]
- The location of this expression in the source code.
- """
-
- value: float
-
- def __init__(self, dtype: str | ir.PrimType, value: float, span: Span | None = None) -> None:
- if isinstance(dtype, ir.PrimType):
- dtype = dtype.dtype
- self.__init_handle_by_constructor__(
- tvm.ir._ffi_api.FloatImm,
- dtype,
- value,
- span, # type: ignore
- )
-
- def __float__(self) -> float:
- return self.value
-
-
-@tvm_ffi.register_object("ir.IntImm")
-class IntImm(ConstExpr):
- """Int constant.
-
- Parameters
- ----------
- dtype : str
- The data type
-
- value : int
- The constant value.
-
- span : Optional[Span]
- The location of this expression in the source code.
- """
-
- value: int
-
- def __init__(self, dtype: str | ir.PrimType, value: int, span: Span | None = None) -> None:
- if isinstance(dtype, ir.PrimType):
- dtype = dtype.dtype
- self.__init_handle_by_constructor__(
- tvm.ir._ffi_api.IntImm,
- dtype,
- value,
- span, # type: ignore
- )
-
- def __hash__(self) -> int:
- return self.value
-
- def __int__(self) -> int:
- return self.value
-
- def __nonzero__(self) -> bool:
- return self.value != 0
-
- def __eq__(self, other: Expr) -> Expr:
- return _ffi_api._OpEQ(self, other, None) # type: ignore
-
- def __ne__(self, other: Expr) -> Expr:
- return _ffi_api._OpNE(self, other, None) # type: ignore
-
- def __bool__(self) -> bool:
- return self.__nonzero__()
-
-
-@tvm_ffi.register_object("ir.prim.StringImm") # type: ignore
-class StringImm(ConstExpr):
- """String constant.
-
- Parameters
- ----------
- value : str
- The value of the function.
-
- span : Optional[Span]
- The location of this expression in the source code.
- """
-
- value: str
-
- def __init__(self, value: str, span: Span | None = None) -> None:
- self.__init_handle_by_constructor__(_prim_ffi_api.StringImm, value, span) # type: ignore
-
- def __eq__(self, other: Expr) -> bool:
- if isinstance(other, ConstExpr):
- return self.value == other.value
- return self.value == other
-
- def __ne__(self, other: Expr) -> bool:
- if isinstance(other, ConstExpr):
- return self.value != other.value
- return self.value != other
-
- def __hash__(self) -> int:
- return Expr.__hash__(self)
-
-
-@tvm_ffi.register_object("ir.prim.Cast")
-class Cast(ir.ExprWithOp):
- """Cast expression.
-
- Parameters
- ----------
- dtype : str
- The data type
-
- value : Expr
- The value of the function.
-
- span : Optional[Span]
- The location of this expression in the source code.
- """
-
- value: Expr
-
- def __init__(self, dtype: str | ir.PrimType, value, span: Span | None = None) -> None:
- if isinstance(dtype, ir.PrimType):
- dtype = dtype.dtype
- self.__init_handle_by_constructor__(_prim_ffi_api.Cast, dtype, value, span) # type: ignore
-
-
-@tvm_ffi.register_object("ir.prim.Add")
-class Add(BinaryOpExpr):
- """Add node.
-
- Parameters
- ----------
- a : Expr
- The left hand operand.
-
- b : Expr
- The right hand operand.
-
- span : Optional[Span]
- The location of this expression in the source code.
- """
-
- def __init__(self, a: Expr, b: Expr, span: Span | None = None) -> None:
- self.__init_handle_by_constructor__(_prim_ffi_api.Add, a, b, span) # type: ignore
-
-
-@tvm_ffi.register_object("ir.prim.Sub")
-class Sub(BinaryOpExpr):
- """Sub node.
-
- Parameters
- ----------
- a : Expr
- The left hand operand.
-
- b : Expr
- The right hand operand.
-
- span : Optional[Span]
- The location of this expression in the source code.
- """
-
- def __init__(self, a: Expr, b: Expr, span: Span | None = None) -> None:
- self.__init_handle_by_constructor__(_prim_ffi_api.Sub, a, b, span) # type: ignore
-
-
-@tvm_ffi.register_object("ir.prim.Mul")
-class Mul(BinaryOpExpr):
- """Mul node.
-
- Parameters
- ----------
- a : Expr
- The left hand operand.
-
- b : Expr
- The right hand operand.
-
- span : Optional[Span]
- The location of this expression in the source code.
- """
-
- def __init__(self, a: Expr, b: Expr, span: Span | None = None) -> None:
- self.__init_handle_by_constructor__(_prim_ffi_api.Mul, a, b, span) # type: ignore
-
-
-@tvm_ffi.register_object("ir.prim.Div")
-class Div(BinaryOpExpr):
- """Div node.
-
- Parameters
- ----------
- a : Expr
- The left hand operand.
-
- b : Expr
- The right hand operand.
-
- span : Optional[Span]
- The location of this expression in the source code.
- """
-
- def __init__(self, a: Expr, b: Expr, span: Span | None = None) -> None:
- self.__init_handle_by_constructor__(_prim_ffi_api.Div, a, b, span) # type: ignore
-
-
-@tvm_ffi.register_object("ir.prim.Mod")
-class Mod(BinaryOpExpr):
- """Mod node.
-
- Parameters
- ----------
- a : Expr
- The left hand operand.
-
- b : Expr
- The right hand operand.
-
- span : Optional[Span]
- The location of this expression in the source code.
- """
-
- def __init__(self, a: Expr, b: Expr, span: Span | None = None) -> None:
- self.__init_handle_by_constructor__(_prim_ffi_api.Mod, a, b, span) # type: ignore
-
-
-@tvm_ffi.register_object("ir.prim.FloorDiv")
-class FloorDiv(BinaryOpExpr):
- """FloorDiv node.
-
- Parameters
- ----------
- a : Expr
- The left hand operand.
-
- b : Expr
- The right hand operand.
-
- span : Optional[Span]
- The location of this expression in the source code.
- """
-
- def __init__(self, a: Expr, b: Expr, span: Span | None = None) -> None:
- self.__init_handle_by_constructor__(_prim_ffi_api.FloorDiv, a, b, span) # type: ignore
-
-
-@tvm_ffi.register_object("ir.prim.FloorMod")
-class FloorMod(BinaryOpExpr):
- """FloorMod node.
-
- Parameters
- ----------
- a : Expr
- The left hand operand.
-
- b : Expr
- The right hand operand.
-
- span : Optional[Span]
- The location of this expression in the source code.
- """
-
- def __init__(self, a: Expr, b: Expr, span: Span | None = None) -> None:
- self.__init_handle_by_constructor__(_prim_ffi_api.FloorMod, a, b, span) # type: ignore
-
-
-@tvm_ffi.register_object("ir.prim.Min")
-class Min(BinaryOpExpr):
- """Min node.
-
- Parameters
- ----------
- a : Expr
- The left hand operand.
-
- b : Expr
- The right hand operand.
-
- span : Optional[Span]
- The location of this expression in the source code.
- """
-
- def __init__(self, a: Expr, b: Expr, span: Span | None = None) -> None:
- self.__init_handle_by_constructor__(_prim_ffi_api.Min, a, b, span) # type: ignore
-
-
-@tvm_ffi.register_object("ir.prim.Max")
-class Max(BinaryOpExpr):
- """Max node.
-
- Parameters
- ----------
- a : Expr
- The left hand operand.
-
- b : Expr
- The right hand operand.
-
- span : Optional[Span]
- The location of this expression in the source code.
- """
-
- def __init__(self, a: Expr, b: Expr, span: Span | None = None) -> None:
- self.__init_handle_by_constructor__(_prim_ffi_api.Max, a, b, span) # type: ignore
-
-
-@tvm_ffi.register_object("ir.prim.EQ")
-class EQ(CmpExpr):
- """EQ node.
-
- Parameters
- ----------
- a : Expr
- The left hand operand.
-
- b : Expr
- The right hand operand.
-
- span : Optional[Span]
- The location of this expression in the source code.
- """
-
- def __init__(self, a: Expr, b: Expr, span: Span | None = None) -> None:
- self.__init_handle_by_constructor__(_prim_ffi_api.EQ, a, b, span) # type: ignore
-
-
-@tvm_ffi.register_object("ir.prim.NE")
-class NE(CmpExpr):
- """NE node.
-
- Parameters
- ----------
- a : Expr
- The left hand operand.
-
- b : Expr
- The right hand operand.
-
- span : Optional[Span]
- The location of this expression in the source code.
- """
-
- def __init__(self, a: Expr, b: Expr, span: Span | None = None) -> None:
- self.__init_handle_by_constructor__(_prim_ffi_api.NE, a, b, span) # type: ignore
-
-
-@tvm_ffi.register_object("ir.prim.LT")
-class LT(CmpExpr):
- """LT node.
-
- Parameters
- ----------
- a : Expr
- The left hand operand.
-
- b : Expr
- The right hand operand.
-
- span : Optional[Span]
- The location of this expression in the source code.
- """
-
- def __init__(self, a: Expr, b: Expr, span: Span | None = None) -> None:
- self.__init_handle_by_constructor__(_prim_ffi_api.LT, a, b, span) # type: ignore
-
-
-@tvm_ffi.register_object("ir.prim.LE")
-class LE(CmpExpr):
- """LE node.
-
- Parameters
- ----------
- a : Expr
- The left hand operand.
-
- b : Expr
- The right hand operand.
-
- span : Optional[Span]
- The location of this expression in the source code.
- """
-
- def __init__(self, a: Expr, b: Expr, span: Span | None = None) -> None:
- self.__init_handle_by_constructor__(_prim_ffi_api.LE, a, b, span) # type: ignore
-
-
-@tvm_ffi.register_object("ir.prim.GT")
-class GT(CmpExpr):
- """GT node.
-
- Parameters
- ----------
- a : Expr
- The left hand operand.
-
- b : Expr
- The right hand operand.
-
- span : Optional[Span]
- The location of this expression in the source code.
- """
-
- def __init__(self, a: Expr, b: Expr, span: Span | None = None) -> None:
- self.__init_handle_by_constructor__(_prim_ffi_api.GT, a, b, span) # type: ignore
-
-
-@tvm_ffi.register_object("ir.prim.GE")
-class GE(CmpExpr):
- """GE node.
-
- Parameters
- ----------
- a : Expr
- The left hand operand.
-
- b : Expr
- The right hand operand.
-
- span : Optional[Span]
- The location of this expression in the source code.
- """
-
- def __init__(self, a: Expr, b: Expr, span: Span | None = None) -> None:
- self.__init_handle_by_constructor__(_prim_ffi_api.GE, a, b, span) # type: ignore
-
-
-@tvm_ffi.register_object("ir.prim.And")
-class And(LogicalExpr):
- """And node.
-
- Parameters
- ----------
- a : Expr
- The left hand operand.
-
- b : Expr
- The right hand operand.
-
- span : Optional[Span]
- The location of this expression in the source code.
- """
-
- def __init__(self, a: Expr, b: Expr, span: Span | None = None) -> None:
- self.__init_handle_by_constructor__(_prim_ffi_api.And, a, b, span) # type: ignore
-
-
-@tvm_ffi.register_object("ir.prim.Or")
-class Or(LogicalExpr):
- """Or node.
-
- Parameters
- ----------
- a : Expr
- The left hand operand.
-
- b : Expr
- The right hand operand.
-
- span : Optional[Span]
- The location of this expression in the source code.
- """
-
- a: Expr
- b: Expr
-
- def __init__(self, a: Expr, b: Expr, span: Span | None = None) -> None:
- self.__init_handle_by_constructor__(_prim_ffi_api.Or, a, b, span) # type: ignore
-
-
-@tvm_ffi.register_object("ir.prim.Not")
-class Not(LogicalExpr):
- """Not node.
-
- Parameters
- ----------
- a : Expr
- The input value
-
- span : Optional[Span]
- The location of this expression in the source code.
- """
-
- a: Expr
-
- def __init__(self, a: Expr, span: Span | None = None) -> None:
- self.__init_handle_by_constructor__(_prim_ffi_api.Not, a, span) # type: ignore
-
-
-@tvm_ffi.register_object("ir.prim.Select")
-class Select(ir.ExprWithOp):
- """Select node.
-
- Note
- ----
- Select may compute both true_value and false_value.
- Use :py:class:`tvm.tirx.if_then_else` instead if you want to
- get a conditional expression that only evaluates
- the correct branch.
-
- Parameters
- ----------
- condition : Expr
- The condition expression.
-
- true_value : Expr
- The value to take when condition is true.
-
- false_value : Expr
- The value to take when condition is false.
-
- span : Optional[Span]
- The location of this expression in the source code.
- """
-
- condition: Expr
- true_value: Expr
- false_value: Expr
-
- def __init__(
- self,
- condition: Expr,
- true_value: Expr,
- false_value: Expr,
- span: Span | None = None,
- ) -> None:
- if isinstance(condition, bool):
- condition = IntImm("bool", condition)
- self.__init_handle_by_constructor__(
- _prim_ffi_api.Select,
- condition,
- true_value,
- false_value,
- span, # type: ignore
- )
-
-
-def BufferLoad(buffer: Buffer, indices: list[Expr], span: Span | None = None) -> tvm.ir.TensorLoad:
+def BufferLoad(buffer: Buffer, indices: list[Expr], span: Span | None = None) -> ir.TensorLoad:
"""Construct a validated buffer load.
Parameters
@@ -1191,90 +317,6 @@ def BufferLoad(buffer: Buffer, indices: list[Expr], span: Span | None = None) ->
return _ffi_api.BufferLoad(buffer, indices, span)
-@tvm_ffi.register_object("ir.prim.Ramp")
-class Ramp(ir.ExprWithOp):
- """Ramp node.
-
- Parameters
- ----------
- base : Expr
- The base expression.
-
- stride : Expr
- The stride of the ramp.
-
- lanes : Expr
- The lanes of the expression.
-
- span : Optional[Span]
- The location of this expression in the source code.
- """
-
- base: Expr
- stride: Expr
- lanes: Expr
-
- def __init__(self, base: Expr, stride: Expr, lanes: Expr, span: Span | None = None) -> None:
- self.__init_handle_by_constructor__(
- _prim_ffi_api.Ramp,
- base,
- stride,
- lanes,
- span, # type: ignore
- )
-
-
-@tvm_ffi.register_object("ir.prim.Broadcast")
-class Broadcast(ir.ExprWithOp):
- """Broadcast node.
-
- Parameters
- ----------
- value : Expr
- The value of the expression.
-
- lanes : Expr
- The lanes of the expression.
-
- span : Optional[Span]
- The location of this expression in the source code.
- """
-
- value: Expr
- lanes: Expr
-
- def __init__(self, value: Expr, lanes: Expr, span: Span | None = None) -> None:
- self.__init_handle_by_constructor__(_prim_ffi_api.Broadcast, value, lanes, span) # type: ignore
-
-
-@tvm_ffi.register_object("ir.prim.Shuffle")
-class Shuffle(ir.ExprWithOp):
- """Shuffle node.
-
- Parameters
- ----------
- vectors : List[Expr]
- The vectors
-
- indices : List[Expr]
- The indices
-
- span : Optional[Span]
- The location of this expression in the source code.
- """
-
- vectors: list[Expr]
- indices: list[Expr]
-
- def __init__(self, vectors: list[Expr], indices: list[Expr], span: Span | None = None) -> None:
- self.__init_handle_by_constructor__(
- _prim_ffi_api.Shuffle,
- vectors,
- indices,
- span, # type: ignore
- )
-
-
class CallEffectKind:
"""Possible kinds of Call effects."""
@@ -1284,30 +326,3 @@ class CallEffectKind:
ReadState = IntImmEnum(2)
UpdateState = IntImmEnum(3)
Opaque = UpdateState
-
-
-@tvm_ffi.register_object("ir.prim.Let")
-class Let(ir.ExprWithOp):
- """Let node.
-
- Parameters
- ----------
- var : Var
- The variable in the binding.
-
- value : Expr
- The value in to be bound.
-
- body : Expr
- The body expression.
-
- span : Optional[Span]
- The location of this expression in the source code.
- """
-
- var: Var
- value: Expr
- body: Expr
-
- def __init__(self, var: Var, value: Expr, body: Expr, span: Span | None = None) -> None:
- self.__init_handle_by_constructor__(_prim_ffi_api.Let, var, value, body, span) # type: ignore
diff --git a/python/tvm/tirx/op.py b/python/tvm/tirx/op.py
index 73422f82db..81a97ee891 100644
--- a/python/tvm/tirx/op.py
+++ b/python/tvm/tirx/op.py
@@ -23,9 +23,11 @@ import tvm_ffi
from tvm_ffi import Array
import tvm
+import tvm.ir.prim._ffi_api as _prim_ffi_api
from tvm import tirx
from tvm.ir import Call, Expr, ExprWithOp, Op, PointerType, PrimType, TensorLoad
from tvm.ir.base import Span
+from tvm.ir.prim import max_value, min_value
from tvm.ir.type import TensorMapType
from tvm.runtime import const
@@ -1247,9 +1249,9 @@ def any(*args, span=None):
raise ValueError("Any must take at least 1 argument")
if len(args) == 1:
return args[0]
- val = _ffi_api._OpOr(args[0], args[1], span) # type: ignore
+ val = _prim_ffi_api._OpOr(args[0], args[1], span) # type: ignore
for i in range(2, len(args)):
- val = _ffi_api._OpOr(val, args[i], span) # type: ignore
+ val = _prim_ffi_api._OpOr(val, args[i], span) # type: ignore
return val
@@ -1274,9 +1276,9 @@ def all(*args, span=None):
raise ValueError("Any must take at least 1 argument")
if len(args) == 1:
return args[0]
- val = _ffi_api._OpAnd(args[0], args[1], span) # type: ignore
+ val = _prim_ffi_api._OpAnd(args[0], args[1], span) # type: ignore
for i in range(2, len(args)):
- val = _ffi_api._OpAnd(val, args[i], span) # type: ignore
+ val = _prim_ffi_api._OpAnd(val, args[i], span) # type: ignore
return val
@@ -1321,44 +1323,6 @@ def trace(args, trace_action="tvm.default_trace_action"):
return tvm.ir.Call(Op.get("tirx.tvm_call_trace_packed"), call_args, ret_ty=ret_ty)
-def min_value(dtype, span=None):
- """minimum value of dtype
-
- Parameters
- ----------
- dtype : str
- The data type.
-
- span : Optional[Span]
- The location of this operator in the source code.
-
- Returns
- -------
- value : tvm.Expr
- The minimum value of dtype.
- """
- return _ffi_api.min_value(dtype, span) # type: ignore
-
-
-def max_value(dtype: str, span: Span | None = None) -> Any:
- """maximum value of dtype
-
- Parameters
- ----------
- dtype : str
- The data type.
-
- span : Optional[Span]
- The location of this operator in the source code.
-
- Returns
- -------
- value : tvm.Expr
- The maximum value of dtype.
- """
- return _ffi_api.max_value(dtype, span) # type: ignore
-
-
def infinity(dtype: str, span: Span | None = None) -> Any:
"""infinity value of dtype
@@ -1562,7 +1526,7 @@ def log2(x):
The result.
"""
x = tir.convert(x)
- return call_intrin(_primexpr_ty(x), "tirx.log2", x)
+ return call_intrin(_primexpr_ty(x), "prim.log2", x)
def log10(x):
@@ -1893,7 +1857,7 @@ def ceil(x, span=None):
y : Expr
The result.
"""
- return _ffi_api.ceil(x, span) # type: ignore
+ return _prim_ffi_api.ceil(x, span) # type: ignore
def trunc(x, span=None):
@@ -1956,7 +1920,7 @@ def bitwise_and(x, y, span=None):
res : Expr
The result.
"""
- return _ffi_api.bitwise_and(x, y, span)
+ return _prim_ffi_api.bitwise_and(x, y, span)
def bitwise_not(x, span=None):
@@ -1975,7 +1939,7 @@ def bitwise_not(x, span=None):
res : Expr
The result.
"""
- return _ffi_api.bitwise_not(x, span)
+ return _prim_ffi_api.bitwise_not(x, span)
def bitwise_or(x, y, span=None):
@@ -1997,7 +1961,7 @@ def bitwise_or(x, y, span=None):
res : Expr
The result.
"""
- return _ffi_api.bitwise_or(x, y, span)
+ return _prim_ffi_api.bitwise_or(x, y, span)
def bitwise_xor(x, y, span=None):
@@ -2019,7 +1983,7 @@ def bitwise_xor(x, y, span=None):
res : Expr
The result.
"""
- return _ffi_api.bitwise_xor(x, y, span)
+ return _prim_ffi_api.bitwise_xor(x, y, span)
def round(x, span=None):
@@ -2168,7 +2132,7 @@ def likely(cond, span=None):
y : Expr
The marked expression.
"""
- return _ffi_api.likely(cond, span) # type: ignore
+ return _prim_ffi_api.likely(cond, span) # type: ignore
def filter(var, pred, *, span=None): # pylint: disable=redefined-builtin
@@ -2431,7 +2395,7 @@ def shift_left(x, y, span=None):
z : Expr
The result.
"""
- return _ffi_api.left_shift(x, y, span)
+ return _prim_ffi_api.left_shift(x, y, span)
def shift_right(x, y, span=None):
@@ -2450,7 +2414,7 @@ def shift_right(x, y, span=None):
z : Expr
The result.
"""
- return _ffi_api.right_shift(x, y, span)
+ return _prim_ffi_api.right_shift(x, y, span)
def fmod(x, y):
@@ -2503,7 +2467,7 @@ def if_then_else(cond, t, f, span=None):
Unlike Select, if_then_else cannot be vectorized
if some lanes in the vector have different conditions.
"""
- return _ffi_api._OpIfThenElse(cond, t, f, span) # type: ignore
+ return _prim_ffi_api._OpIfThenElse(cond, t, f, span) # type: ignore
def div(a, b, span=None):
@@ -2528,7 +2492,7 @@ def div(a, b, span=None):
----
When operands are integers, returns truncdiv(a, b, span).
"""
- return _ffi_api._OpDiv(a, b, span) # type: ignore
+ return _prim_ffi_api._OpDiv(a, b, span) # type: ignore
def indexdiv(a, b, span=None):
@@ -2556,7 +2520,7 @@ def indexdiv(a, b, span=None):
This function may take advantage of operands'
non-negativeness.
"""
- return _ffi_api._OpIndexDiv(a, b, span) # type: ignore
+ return _prim_ffi_api._OpIndexDiv(a, b, span) # type: ignore
def indexmod(a, b, span=None):
@@ -2584,7 +2548,7 @@ def indexmod(a, b, span=None):
This function may take advantage of operands'
non-negativeness.
"""
- return _ffi_api._OpIndexMod(a, b, span) # type: ignore
+ return _prim_ffi_api._OpIndexMod(a, b, span) # type: ignore
def truncdiv(a, b, span=None):
@@ -2610,7 +2574,7 @@ def truncdiv(a, b, span=None):
----
This is the default integer division behavior in C.
"""
- return _ffi_api._OpTruncDiv(a, b, span) # type: ignore
+ return _prim_ffi_api._OpTruncDiv(a, b, span) # type: ignore
def truncmod(a, b, span=None):
@@ -2636,7 +2600,7 @@ def truncmod(a, b, span=None):
----
This is the default integer division behavior in C.
"""
- return _ffi_api._OpTruncMod(a, b, span) # type: ignore
+ return _prim_ffi_api._OpTruncMod(a, b, span) # type: ignore
def floordiv(a, b, span=None):
@@ -2658,7 +2622,7 @@ def floordiv(a, b, span=None):
res : Expr
The result expression.
"""
- return _ffi_api._OpFloorDiv(a, b, span) # type: ignore
+ return _prim_ffi_api._OpFloorDiv(a, b, span) # type: ignore
def logaddexp(a, b, span=None):
@@ -2702,7 +2666,7 @@ def floormod(a, b, span=None):
res : Expr
The result expression.
"""
- return _ffi_api._OpFloorMod(a, b, span) # type: ignore
+ return _prim_ffi_api._OpFloorMod(a, b, span) # type: ignore
def ceildiv(lhs, rhs, span=None):
@@ -2722,7 +2686,7 @@ def ceildiv(lhs, rhs, span=None):
op : tvm.Expr
The result Expr of ceildiv operaton.
"""
- return _ffi_api._OpCeilDiv(lhs, rhs, span) # type: ignore
+ return _prim_ffi_api._OpCeilDiv(lhs, rhs, span) # type: ignore
def comm_reducer(fcombine, fidentity, name="reduce"):
@@ -3005,7 +2969,7 @@ def vscale():
call : Expr
Call to the vscale intrinsic
"""
- return call_intrin("int32", "ir.prim.vscale")
+ return call_intrin("int32", "prim.vscale")
def get_active_lane_mask(dtype, base, limit):
@@ -3105,8 +3069,8 @@ def ignore_loop_partition(predicate) -> Expr:
# pylint: disable=unnecessary-lambda
sum = comm_reducer(lambda x, y: x + y, lambda t: const(0, dtype=t), name="sum")
-min = comm_reducer(lambda x, y: _ffi_api._OpMin(x, y, None), max_value, name="min") # type: ignore
-max = comm_reducer(lambda x, y: _ffi_api._OpMax(x, y, None), min_value, name="max") # type: ignore
+min = comm_reducer(lambda x, y: _prim_ffi_api._OpMin(x, y, None), max_value, name="min") # type: ignore
+max = comm_reducer(lambda x, y: _prim_ffi_api._OpMax(x, y, None), min_value, name="max") # type: ignore
def tvm_load_matrix_sync(fragment, m, n, k, index, buffer_ptr, stride, layout):
diff --git a/python/tvm/tirx/script/builder/ir.py b/python/tvm/tirx/script/builder/ir.py
index 1782c5ff6a..3cd5d45e88 100644
--- a/python/tvm/tirx/script/builder/ir.py
+++ b/python/tvm/tirx/script/builder/ir.py
@@ -37,6 +37,7 @@ from tvm import tirx as tir
from tvm.ir import Call, TensorLoad, Type, is_prim_expr
from tvm.ir import register_op_attr as _register_op_attr
from tvm.ir.base import deprecated
+from tvm.ir.prim import _ffi_api as _prim_ffi_api
from tvm.runtime import convert
from tvm.script.ir_builder.base import IRBuilder
from tvm.script.ir_builder.ir import meta_var
@@ -45,7 +46,6 @@ from tvm.target import Target
# pylint: disable=unused-import
from tvm.target.codegen import llvm_lookup_intrinsic_id
from tvm.tirx import Buffer, BufferRegion, Expr, IndexMap, is_buffer_var, type_annotation
-from tvm.tirx import _ffi_api as _tirx_ffi_api
from tvm.tirx import op as _tir_op
from tvm.tirx.exec_scope import ExecScope, ScopeIdDef, Var
@@ -101,7 +101,7 @@ from .external_kernel import call_kernel
def cast(value, dtype, span=None):
"""Cast an expression to the requested data type."""
- return _tirx_ffi_api._cast(dtype, value, span) # type: ignore[attr-defined]
+ return _prim_ffi_api._cast(dtype, value, span) # type: ignore[attr-defined]
def _current_s_tir() -> bool:
@@ -1601,7 +1601,7 @@ class DtypeConstructor:
def __call__(
self,
- expr: "None | Expr | Literal['inf', '-inf', 'nan'] | int | float" = None,
+ expr: "Expr | Literal['inf', '-inf', 'nan'] | int | float | None" = None,
) -> "Expr":
if isinstance(expr, str):
expr = float(expr)
diff --git a/python/tvm/topi/_te_tensor_overload.py b/python/tvm/topi/_te_tensor_overload.py
index 67a5347438..c4e973c75e 100644
--- a/python/tvm/topi/_te_tensor_overload.py
+++ b/python/tvm/topi/_te_tensor_overload.py
@@ -17,8 +17,9 @@
"""Register TOPI implementations for TE tensor overload hooks."""
from tvm import te
+from tvm.ir import _overload_prim_expr as _expr
+from tvm.runtime import DataTypeCode
from tvm.te import _te_tensor_overload as _overload
-from tvm.tirx import expr as _expr
from . import broadcast as _broadcast
from . import math as _math
@@ -26,7 +27,7 @@ from . import math as _math
def _is_integer(value):
if isinstance(value, te.Tensor | te.TensorSlice):
- return value.dtype.matches_code(_expr.DataTypeCode.INT)
+ return value.dtype.matches_code(DataTypeCode.INT)
return _expr._dtype_is_int(value)
diff --git a/python/tvm/topi/math.py b/python/tvm/topi/math.py
index e24e9bc8f8..6fd03db7b2 100644
--- a/python/tvm/topi/math.py
+++ b/python/tvm/topi/math.py
@@ -764,7 +764,7 @@ def cast(x, dtype, span=None):
if isinstance(x, te.tensor.Tensor):
return te.compute(x.shape, lambda *i: x(*i).astype(dtype), tag=tag.ELEMWISE)
# pylint: disable=import-outside-toplevel
- from tvm.tirx import _ffi_api
+ from tvm.ir.prim import _ffi_api
return _ffi_api._cast(dtype, x, span)
diff --git a/src/arith/analyzer.cc b/src/arith/analyzer.cc
index 3819017b46..755fbfe4f0 100644
--- a/src/arith/analyzer.cc
+++ b/src/arith/analyzer.cc
@@ -25,8 +25,8 @@
#include <tvm/ffi/function.h>
#include <tvm/ffi/reflection/registry.h>
#include <tvm/ir/prim/expr.h>
+#include <tvm/ir/prim/op.h>
#include <tvm/runtime/logging.h>
-#include <tvm/tirx/op.h>
#include "const_fold.h"
#include "product_normal_form.h"
@@ -58,7 +58,7 @@ void AnalyzerObj::Bind(const Var& var, const PrimExpr& expr, bool allow_override
void AnalyzerObj::Bind(const Var& var, const Range& range, bool allow_override) {
TVM_FFI_ICHECK(range.defined());
- if (tirx::is_one(range->extent)) {
+ if (prim::is_one(range->extent)) {
this->Bind(var, range->min, allow_override);
} else {
this->const_int_bound.Bind(var, range, allow_override);
@@ -74,7 +74,7 @@ void AnalyzerObj::MarkGlobalNonNegValue(const PrimExpr& value) {
// decompose value as symbol * scale + offset
int64_t offset = 0;
PrimType value_ty = value.ty();
- PrimExpr symbol_scale = tirx::MakeConst(value_ty, 0);
+ PrimExpr symbol_scale = prim::MakeConst(value_ty, 0);
auto fcollect_sum = [&](PrimExpr val, int sign) {
if (const auto* intimm = val.as<IntImmNode>()) {
@@ -91,7 +91,7 @@ void AnalyzerObj::MarkGlobalNonNegValue(const PrimExpr& value) {
// split out the symbol and non-symbolic part
int64_t cscale = 1;
- PrimExpr symbol = tirx::MakeConst(value_ty, 1);
+ PrimExpr symbol = prim::MakeConst(value_ty, 1);
auto fcollect_prod = [&](PrimExpr val) {
if (const auto* intimm = val.as<IntImmNode>()) {
cscale *= intimm->value;
@@ -181,7 +181,7 @@ bool AnalyzerObj::CanProveLessEqualThanSymbolicShapeValue(const PrimExpr& lhs,
const PrimExpr& shape) {
if (this->CanProve(lhs <= shape, ProofStrength::kSymbolicBound)) return true;
// no need to do further attempt if shape is already a constant.
- if (tirx::is_const_int(shape)) return false;
+ if (prim::is_const_int(shape)) return false;
// collect constant scale and ignore symbolic part
// so 32 * n => cscale = 32
int64_t cscale = 1;
@@ -202,7 +202,7 @@ bool AnalyzerObj::CanProve(const PrimExpr& expr, ProofStrength strength) {
return ptr->value != 0;
}
PrimExpr simplified = Simplify(expr);
- const int64_t* as_int = tirx::as_const_int(simplified);
+ const int64_t* as_int = prim::as_const_int(simplified);
if (as_int && *as_int) return true;
if (strength >= ProofStrength::kSymbolicBound) {
// NOTE: we intentionally only pattern match common bound predicate i < bound
@@ -254,7 +254,7 @@ PrimExpr AnalyzerObj::Simplify(const PrimExpr& expr, int steps) {
res = this->canonical_simplify(res);
for (int i = 0; i < steps; ++i) {
- if (tirx::is_const_int(res)) {
+ if (prim::is_const_int(res)) {
return res;
}
if (i % 2 == 0) {
diff --git a/src/arith/canonical_simplify.cc b/src/arith/canonical_simplify.cc
index f0bea1a5a9..93fd0b2c7d 100644
--- a/src/arith/canonical_simplify.cc
+++ b/src/arith/canonical_simplify.cc
@@ -26,7 +26,7 @@
#include <tvm/ffi/expected.h>
#include <tvm/ir/cow.h>
#include <tvm/ir/prim/expr.h>
-#include <tvm/tirx/op.h>
+#include <tvm/ir/prim/op.h>
#include "const_fold.h"
#include "pattern_match.h"
@@ -107,8 +107,8 @@ bool CastIsSafe(PrimType dtype, PrimExpr value, AnalyzerObj* analyzer) {
return false;
}
ConstIntBound bound = analyzer->const_int_bound(value);
- int64_t ubound = max_value(dtype).as_or_throw<IntImm>()->value;
- int64_t lbound = min_value(dtype).as_or_throw<IntImm>()->value;
+ int64_t ubound = prim::max_value(dtype).as_or_throw<IntImm>()->value;
+ int64_t lbound = prim::min_value(dtype).as_or_throw<IntImm>()->value;
if (value.ty().bits() <= dtype.bits() || // upcast is safe
(bound->max_value <= ubound && bound->min_value >= lbound)) {
return true;
@@ -176,8 +176,8 @@ class SplitExprNode : public CanonicalExprNode {
* \return whether the cast can be safely pushed to children
*/
bool CanPushCastToChildren(PrimType dtype, AnalyzerObj* analyzer) const {
- // cast(dtype, index % upper_factor / lower_factor * scale) ==
- // cast(dtype, index) % upper_factor / lower_factor * scale
+ // prim::cast(dtype, index % upper_factor / lower_factor * scale) ==
+ // prim::cast(dtype, index) % upper_factor / lower_factor * scale
// iff it is an upcast (dtype.bits >= self.dtype.bits) or all of
// its intermediate results fit in the range of dtype
PrimType self_dtype = this->ExprNode::ty.as_or_throw<PrimType>();
@@ -214,11 +214,11 @@ class SplitExprNode : public CanonicalExprNode {
}
/*!
- * \brief self = cast(dtype, self)
+ * \brief self = prim::cast(dtype, self)
* \param dtype The target datatype
*/
void PushCastToChildren(PrimType dtype) {
- this->index = cast(dtype, this->index);
+ this->index = prim::cast(dtype, this->index);
this->ExprNode::ty = dtype;
}
@@ -359,8 +359,8 @@ class SumExprNode : public CanonicalExprNode {
bool CanPushCastToChildren(PrimType dtype, AnalyzerObj* analyzer) const {
bool is_min_value = dtype.bits() == 64 ? base == std::numeric_limits<int64_t>::lowest()
: base == -(1LL << (dtype.bits() - 1));
- // cast(dtype, arg_1 + arg_2 + ... arg_n) ==
- // cast(dtype, arg_1) + ... + cast(dtype, arg_n)
+ // prim::cast(dtype, arg_1 + arg_2 + ... arg_n) ==
+ // prim::cast(dtype, arg_1) + ... + prim::cast(dtype, arg_n)
// iff it is an upcast (dtype.bits >= self.dtype.bits) or all of
// its intermediate results fit in the range of dtype
PrimType self_dtype = this->ExprNode::ty.as_or_throw<PrimType>();
@@ -406,7 +406,7 @@ class SumExprNode : public CanonicalExprNode {
}
/*!
- * \brief self = cast(dtype, self)
+ * \brief self = prim::cast(dtype, self)
* \param dtype The target datatype
*/
void PushCastToChildren(PrimType dtype) {
diff --git a/src/arith/conjunctive_normal_form.cc b/src/arith/conjunctive_normal_form.cc
index 37c9a4aa7a..3194c0b33e 100644
--- a/src/arith/conjunctive_normal_form.cc
+++ b/src/arith/conjunctive_normal_form.cc
@@ -25,7 +25,7 @@
#include <tvm/arith/analyzer.h>
#include <tvm/ir/prim/expr.h>
-#include <tvm/tirx/op.h>
+#include <tvm/ir/prim/op.h>
#include <optional>
#include <unordered_map>
diff --git a/src/arith/const_fold.h b/src/arith/const_fold.h
index 0c2b72a45c..454fe45a5a 100644
--- a/src/arith/const_fold.h
+++ b/src/arith/const_fold.h
@@ -19,25 +19,31 @@
/*!
* \file const_fold.h
- * \brief Centralized location for constant folding.
+ * \brief Arithmetic aliases for shared constant folding and scoped index analysis.
*/
#ifndef TVM_ARITH_CONST_FOLD_H_
#define TVM_ARITH_CONST_FOLD_H_
-#include <tvm/ffi/optional.h>
-#include <tvm/ir/prim/expr.h>
-#include <tvm/runtime/logging.h>
-#include <tvm/tirx/op.h>
-
-#include <algorithm>
-#include <cmath>
-#include <limits>
-
+#include "../ir/prim/const_fold.h"
#include "int_operator.h"
+#define TVM_ARITH_CONST_PROPAGATION(BODY) TVM_PRIM_CONST_PROPAGATION(BODY)
+#define TVM_INDEX_CONST_PROPAGATION(BODY) TVM_PRIM_INDEX_CONST_PROPAGATION(BODY)
+
namespace tvm {
namespace arith {
+using prim::detail::GetFoldResultDoubleRepr;
+using prim::detail::GetFoldResultInt64Repr;
+using prim::detail::is_neg_inf;
+using prim::detail::is_pos_inf;
+using prim::detail::IsIndexType;
+using prim::detail::IsIndexTypedExpr;
+using prim::detail::neg_inf;
+using prim::detail::pos_inf;
+using prim::detail::SymbolicLimits;
+using prim::detail::TryConstFold;
+
/*!
* \brief Scoped opt-in flag for unsigned index analysis (allow_u32).
*
@@ -65,470 +71,6 @@ class AllowUintAsIndexGuard {
AllowUintAsIndexGuard& operator=(const AllowUintAsIndexGuard&) = delete;
};
-/*!
- * \brief Try to run binary compute with constant folding.
- *
- * \param a The left operand.
- * \param b The right operand.
- * \tparam Op The operator type.
- *
- * \note a and b Must already matched data types with each other.
- * \return std::nullopt if constant fold fails, otherwise return folded result.
- */
-template <typename Op>
-inline ffi::Optional<PrimExpr> TryConstFold(PrimExpr a, PrimExpr b);
-
-/*!
- * \brief Try to run unary compute with constant folding.
- *
- * \param a The left operand.
- * \tparam Op The operator type.
- *
- * \note a and b Must already matched data types with each other.
- * \return std::nullopt if constant fold fails, otherwise return folded result.
- */
-template <typename Op>
-inline ffi::Optional<PrimExpr> TryConstFold(PrimExpr a);
-
-/*!
- * \brief Check whether type is used to represent index.
- *
- * Index types are frequently used in shape computation
- * and need to be aggressively constant-folded.
- *
- * \param type The type to represent index.
- * \return the checked result.
- */
-inline bool IsIndexType(DLDataType type) {
- return type.code == static_cast<uint8_t>(DLDataTypeCode::kDLInt) &&
- (type.bits == 32 || type.bits == 64) && type.lanes == 1;
-}
-
-inline bool IsIndexTypedExpr(const ExprNode* expr) {
- TVM_FFI_DCHECK(expr != nullptr);
- TVM_FFI_DCHECK(!expr->ExprNode::ty.IsMissing());
- const auto* prim_ty = expr->ExprNode::ty.as<PrimTypeNode>();
- TVM_FFI_DCHECK(prim_ty != nullptr);
- return IsIndexType(prim_ty->dtype);
-}
-
-inline bool IsIndexTypedExpr(const PrimExpr& expr) { return IsIndexTypedExpr(expr.get()); }
-
-/*! \brief Helper to get const folding result repr in int64. */
-inline int64_t GetFoldResultInt64Repr(int64_t x, const PrimType& dtype) {
- if (dtype.bits() < 64) {
- x &= (1LL << dtype.bits()) - 1;
- }
- if (dtype.MatchesCode(DLDataTypeCode::kDLInt)) {
- int64_t m = 1LL << (dtype.bits() - 1);
- x = (x ^ m) - m;
- }
- return x;
-}
-
-/*! \brief Helper to get fp32 const folding result repr in double. */
-inline double GetFoldResultDoubleRepr(float x) {
- double res = static_cast<double>(x);
- if (std::isinf(res) || std::isnan(res)) {
- return res;
- }
- // certain platform (eg, on gcc7-i386) do the folding arithmetic
- // on float and write back to double is optimized to double
- // precision arithmetic, this is legal and we check the output
- // range thus to ensure consistency when the float result is inf.
- if (res < std::numeric_limits<float>::lowest()) {
- LOG(WARNING) << "underlying float value overflow";
- return -std::numeric_limits<double>::infinity();
- } else if (res > std::numeric_limits<float>::max()) {
- LOG(WARNING) << "underlying float value overflow";
- return std::numeric_limits<double>::infinity();
- }
- return res;
-}
-
-#define TVM_ARITH_CONST_PROPAGATION(BODY) \
- using tvm::FloatImmNode; \
- const IntImmNode* pa = a.as<IntImmNode>(); \
- const IntImmNode* pb = b.as<IntImmNode>(); \
- const FloatImmNode* fa = a.as<FloatImmNode>(); \
- const FloatImmNode* fb = b.as<FloatImmNode>(); \
- BODY;
-
-#define TVM_INDEX_CONST_PROPAGATION(BODY) \
- const IntImmNode* pa = a.as<IntImmNode>(); \
- const IntImmNode* pb = b.as<IntImmNode>(); \
- if (arith::IsIndexTypedExpr(a) && arith::IsIndexTypedExpr(b)) { \
- BODY; \
- }
-
-// specialization of constant folders.
-template <>
-inline ffi::Optional<PrimExpr> TryConstFold<prim::Add>(PrimExpr a, PrimExpr b) {
- TVM_ARITH_CONST_PROPAGATION({
- PrimType result_ty = a.ty();
- if (pa && pb) {
- int64_t res = pa->value + pb->value;
- return IntImm(result_ty, GetFoldResultInt64Repr(res, result_ty));
- }
- if (pa && pa->value == 0) return b;
- if (pb && pb->value == 0) return a;
- if (fa && fb) {
- if (result_ty.bits() == 32) {
- return FloatImm(result_ty, GetFoldResultDoubleRepr(static_cast<float>(fa->value) +
- static_cast<float>(fb->value)));
- } else if (result_ty.bits() == 64) {
- return FloatImm(result_ty, fa->value + fb->value);
- }
- }
- if (fa && fa->value == 0) return b;
- if (fb && fb->value == 0) return a;
- });
- return std::nullopt;
-}
-
-template <>
-inline ffi::Optional<PrimExpr> TryConstFold<prim::Sub>(PrimExpr a, PrimExpr b) {
- TVM_ARITH_CONST_PROPAGATION({
- TVM_FFI_ICHECK(!((pa && pa->ty.as_or_throw<PrimType>().MatchesCode(DLDataTypeCode::kDLUInt) &&
- pa->value == 0U) &&
- (pb && pb->ty.as_or_throw<PrimType>().MatchesCode(DLDataTypeCode::kDLUInt) &&
- pb->value > 0U)))
- << "Checked failed. Minuend 's value is 0U and it's dtype is uint "
- << "while Subtrahend's dtype is uint; which will cause a negative uint";
- PrimType result_ty = a.ty();
- if (pa && pb) {
- int64_t res = pa->value - pb->value;
- return IntImm(result_ty, GetFoldResultInt64Repr(res, result_ty));
- }
- if (pb && pb->value == 0) return a;
- if (fa && fb) {
- if (result_ty.bits() == 32) {
- return FloatImm(result_ty, GetFoldResultDoubleRepr(static_cast<float>(fa->value) -
- static_cast<float>(fb->value)));
- } else if (result_ty.bits() == 64) {
- return FloatImm(result_ty, fa->value - fb->value);
- }
- }
- if (fb && fb->value == 0) return a;
- });
- return std::nullopt;
-}
-
-template <>
-inline ffi::Optional<PrimExpr> TryConstFold<prim::Mul>(PrimExpr a, PrimExpr b) {
- TVM_ARITH_CONST_PROPAGATION({
- PrimType result_ty = a.ty();
- if (pa && pb) {
- int64_t res = pa->value * pb->value;
- return IntImm(result_ty, GetFoldResultInt64Repr(res, result_ty));
- }
- if (pa) {
- if (pa->value == 1) return b;
- if (pa->value == 0) return a;
- }
- if (pb) {
- if (pb->value == 1) return a;
- if (pb->value == 0) return b;
- }
- if (fa && fb) {
- if (result_ty.bits() == 32) {
- return FloatImm(result_ty, GetFoldResultDoubleRepr(static_cast<float>(fa->value) *
- static_cast<float>(fb->value)));
- } else if (result_ty.bits() == 64) {
- return FloatImm(result_ty, fa->value * fb->value);
- }
- }
- if (fa) {
- if (fa->value == 1) return b;
- if (fa->value == 0) return a;
- }
- if (fb) {
- if (fb->value == 1) return a;
- if (fb->value == 0) return b;
- }
- });
- return std::nullopt;
-}
-
-template <>
-inline ffi::Optional<PrimExpr> TryConstFold<prim::Div>(PrimExpr a, PrimExpr b) {
- TVM_ARITH_CONST_PROPAGATION({
- PrimType result_ty = a.ty();
- if (pa && pb) {
- // due to division and mod can have different modes
- // NOTE: this will assumes truc div.
- TVM_FFI_ICHECK_NE(pb->value, 0) << "Divide by zero";
- int64_t res = pa->value / pb->value;
- return IntImm(result_ty, GetFoldResultInt64Repr(res, result_ty));
- }
- if (pa) {
- if (pa->value == 0) return a;
- }
- if (pb) {
- if (pb->value == 1) return a;
- TVM_FFI_ICHECK_NE(pb->value, 0) << "Divide by zero";
- }
- if (fa && fb) {
- TVM_FFI_ICHECK_NE(fb->value, 0) << "Divide by zero";
- if (result_ty.bits() == 32) {
- return FloatImm(result_ty, GetFoldResultDoubleRepr(static_cast<float>(fa->value) /
- static_cast<float>(fb->value)));
- } else if (result_ty.bits() == 64) {
- return FloatImm(result_ty, fa->value / fb->value);
- }
- }
- if (fa && fa->value == 0) return a;
- if (fb) {
- if (fb->value == 1) return a;
- TVM_FFI_ICHECK_NE(fb->value, 0) << "Divide by zero";
- }
- });
- return std::nullopt;
-}
-
-template <>
-inline ffi::Optional<PrimExpr> TryConstFold<prim::Mod>(PrimExpr a, PrimExpr b) {
- TVM_INDEX_CONST_PROPAGATION({
- PrimType result_ty = a.ty();
- if (pa && pb) {
- TVM_FFI_ICHECK_NE(pb->value, 0) << "Divide by zero";
- int64_t res = pa->value % pb->value;
- return IntImm(result_ty, GetFoldResultInt64Repr(res, result_ty));
- }
- if (pa) {
- if (pa->value == 0) return a;
- }
- if (pb) {
- if (pb->value == 1) return IntImm(result_ty, 0);
- TVM_FFI_ICHECK_NE(pb->value, 0) << "Divide by zero";
- }
- });
- return std::nullopt;
-}
-
-template <>
-inline ffi::Optional<PrimExpr> TryConstFold<prim::FloorDiv>(PrimExpr a, PrimExpr b) {
- TVM_ARITH_CONST_PROPAGATION({
- PrimType result_ty = a.ty();
- if (pa && pb) {
- TVM_FFI_ICHECK_NE(pb->value, 0) << "Divide by zero";
- int64_t res = arith::floordiv(pa->value, pb->value);
- return IntImm(result_ty, GetFoldResultInt64Repr(res, result_ty));
- }
- if (pa) {
- if (pa->value == 0) return a;
- }
- if (pb) {
- if (pb->value == 1) return a;
- TVM_FFI_ICHECK_NE(pb->value, 0) << "Divide by zero";
- }
- if (fa && fb && fb->value != 0) {
- if (result_ty.bits() == 32) {
- return FloatImm(result_ty,
- GetFoldResultDoubleRepr(std::floor(static_cast<float>(fa->value) /
- static_cast<float>(fb->value))));
- } else if (result_ty.bits() == 64) {
- return FloatImm(result_ty, std::floor(fa->value / fb->value));
- } else {
- return std::nullopt;
- }
- }
- if (fa && fa->value == 0) return a;
- if (fb) {
- if (fb->value == 1) return a;
- TVM_FFI_ICHECK_NE(fb->value, 0) << "Divide by zero";
- }
- });
- return std::nullopt;
-}
-
-template <>
-inline ffi::Optional<PrimExpr> TryConstFold<prim::FloorMod>(PrimExpr a, PrimExpr b) {
- const IntImmNode* ua = a.as<IntImmNode>();
- const IntImmNode* ub = b.as<IntImmNode>();
- PrimType utype = a.ty();
- if (utype.MatchesCode(DLDataTypeCode::kDLUInt) && utype == b.ty() && ua && ub) {
- auto as_uint = [](int64_t value, const PrimType& dtype) -> uint64_t {
- uint64_t result = static_cast<uint64_t>(value);
- if (dtype.bits() < 64) {
- result &= (uint64_t{1} << dtype.bits()) - 1;
- }
- return result;
- };
- uint64_t lhs = as_uint(ua->value, utype);
- uint64_t rhs = as_uint(ub->value, b.ty());
- TVM_FFI_ICHECK_NE(rhs, 0U) << "Divide by zero";
- return IntImm(utype, static_cast<int64_t>(lhs % rhs));
- }
- TVM_INDEX_CONST_PROPAGATION({
- PrimType result_ty = a.ty();
- if (pa && pb) {
- TVM_FFI_ICHECK_NE(pb->value, 0) << "Divide by zero";
- int64_t res = arith::floormod(pa->value, pb->value);
- return IntImm(result_ty, GetFoldResultInt64Repr(res, result_ty));
- }
- if (pa) {
- if (pa->value == 0) return a;
- }
- if (pb) {
- if (pb->value == 1) return IntImm(result_ty, 0);
- TVM_FFI_ICHECK_NE(pb->value, 0) << "Divide by zero";
- }
- });
- return std::nullopt;
-}
-
-template <>
-inline ffi::Optional<PrimExpr> TryConstFold<prim::Min>(PrimExpr a, PrimExpr b) {
- TVM_ARITH_CONST_PROPAGATION({
- PrimType result_ty = a.ty();
- if (pa && pb) return IntImm(result_ty, std::min(pa->value, pb->value));
- if (fa && fb) return FloatImm(result_ty, std::min(fa->value, fb->value));
- });
- if (a.same_as(b)) return a;
- return std::nullopt;
-}
-
-template <>
-inline ffi::Optional<PrimExpr> TryConstFold<prim::Max>(PrimExpr a, PrimExpr b) {
- TVM_ARITH_CONST_PROPAGATION({
- PrimType result_ty = a.ty();
- if (pa && pb) return IntImm(result_ty, std::max(pa->value, pb->value));
- if (fa && fb) return FloatImm(result_ty, std::max(fa->value, fb->value));
- });
- if (a.same_as(b)) return a;
- return std::nullopt;
-}
-
-template <>
-inline ffi::Optional<PrimExpr> TryConstFold<prim::GT>(PrimExpr a, PrimExpr b) {
- TVM_ARITH_CONST_PROPAGATION({
- if (pa && pb) return IntImm::Bool(pa->value > pb->value);
- if (fa && fb) return IntImm::Bool(fa->value > fb->value);
- });
- return std::nullopt;
-}
-
-template <>
-inline ffi::Optional<PrimExpr> TryConstFold<prim::GE>(PrimExpr a, PrimExpr b) {
- TVM_ARITH_CONST_PROPAGATION({
- if (pa && pb) return IntImm::Bool(pa->value >= pb->value);
- if (fa && fb) return IntImm::Bool(fa->value >= fb->value);
- });
- return std::nullopt;
-}
-
-template <>
-inline ffi::Optional<PrimExpr> TryConstFold<prim::LT>(PrimExpr a, PrimExpr b) {
- TVM_ARITH_CONST_PROPAGATION({
- if (pa && pb) return IntImm::Bool(pa->value < pb->value);
- if (fa && fb) return IntImm::Bool(fa->value < fb->value);
- });
- return std::nullopt;
-}
-
-template <>
-inline ffi::Optional<PrimExpr> TryConstFold<prim::LE>(PrimExpr a, PrimExpr b) {
- TVM_ARITH_CONST_PROPAGATION({
- if (pa && pb) return IntImm::Bool(pa->value <= pb->value);
- if (fa && fb) return IntImm::Bool(fa->value <= fb->value);
- });
- return std::nullopt;
-}
-
-template <>
-inline ffi::Optional<PrimExpr> TryConstFold<prim::EQ>(PrimExpr a, PrimExpr b) {
- TVM_ARITH_CONST_PROPAGATION({
- if (pa && pb) return IntImm::Bool(pa->value == pb->value);
- if (fa && fb) return IntImm::Bool(fa->value == fb->value);
- });
- return std::nullopt;
-}
-
-template <>
-inline ffi::Optional<PrimExpr> TryConstFold<prim::NE>(PrimExpr a, PrimExpr b) {
- TVM_ARITH_CONST_PROPAGATION({
- if (pa && pb) return IntImm::Bool(pa->value != pb->value);
- if (fa && fb) return IntImm::Bool(fa->value != fb->value);
- });
- return std::nullopt;
-}
-
-template <>
-inline ffi::Optional<PrimExpr> TryConstFold<prim::And>(PrimExpr a, PrimExpr b) {
- const IntImmNode* pa = a.as<IntImmNode>();
- const IntImmNode* pb = b.as<IntImmNode>();
- if (pa && pa->value) return b;
- if (pa && !pa->value) return a;
- if (pb && pb->value) return a;
- if (pb && !pb->value) return b;
- return std::nullopt;
-}
-
-template <>
-inline ffi::Optional<PrimExpr> TryConstFold<prim::Or>(PrimExpr a, PrimExpr b) {
- const IntImmNode* pa = a.as<IntImmNode>();
- const IntImmNode* pb = b.as<IntImmNode>();
- if (pa && pa->value) return a;
- if (pa && !pa->value) return b;
- if (pb && pb->value) return b;
- if (pb && !pb->value) return a;
- return std::nullopt;
-}
-
-template <>
-inline ffi::Optional<PrimExpr> TryConstFold<prim::Not>(PrimExpr a) {
- const IntImmNode* pa = a.as<IntImmNode>();
- if (pa) {
- return IntImm::Bool(!(pa->value));
- }
- return std::nullopt;
-}
-
-/*! \brief Helper namespace for symbolic value limits */
-struct SymbolicLimits {
- /*! \brief positive infinity */
- static PrimExpr pos_inf_;
- /*! \brief negative infinity */
- static PrimExpr neg_inf_;
-};
-
-/*!
- * \brief Opaque expression representing positive infinity.
- *
- * It can only be used as parameter of by min/max
- * for integer analysis and cannot be used in normal expressions.
- *
- * \return positive infinity.
- */
-inline PrimExpr pos_inf() { return SymbolicLimits::pos_inf_; }
-
-/*!
- * \brief Check if value is positive infinity.
- * \param value The value to be checked.
- *
- * \return The check result.
- */
-inline bool is_pos_inf(const PrimExpr& value) { return value.same_as(SymbolicLimits::pos_inf_); }
-
-/*!
- * \brief Opaque expression representing negative infinity.
- *
- * It can only be used as parameter of by min/max
- * for integer analysis and cannot be used in normal expressions.
- *
- * \return negative infinity.
- */
-inline PrimExpr neg_inf() { return SymbolicLimits::neg_inf_; }
-
-/*!
- * \brief Check if value is negative infinity.
- * \param value The value to be checked.
- *
- * \return The check result.
- */
-inline bool is_neg_inf(const PrimExpr& value) { return value.same_as(SymbolicLimits::neg_inf_); }
-
} // namespace arith
} // namespace tvm
#endif // TVM_ARITH_CONST_FOLD_H_
diff --git a/src/arith/const_int_bound.cc b/src/arith/const_int_bound.cc
index 678603189d..5e1cc86663 100644
--- a/src/arith/const_int_bound.cc
+++ b/src/arith/const_int_bound.cc
@@ -847,8 +847,8 @@ class ConstIntBoundAnalyzer::Impl
* topi.math.ceil_log2, and can appear in iteration bounds.
*/
static ffi::Optional<PrimExpr> FindCeilLog2Arg(const prim::CastNode* op) {
- static const Op& ceil_op = Op::Get("tirx.ceil");
- static const Op& log2_op = Op::Get("tirx.log2");
+ static const Op& ceil_op = prim::builtin::ceil();
+ static const Op& log2_op = prim::builtin::log2();
if (op->ty.as_or_throw<PrimType>().code() == DLDataTypeCode::kDLInt) {
if (auto as_call = op->value.as<CallNode>()) {
if (as_call->op.same_as(ceil_op)) {
diff --git a/src/arith/detect_linear_equation.cc b/src/arith/detect_linear_equation.cc
index fb12f54674..daf58918ce 100644
--- a/src/arith/detect_linear_equation.cc
+++ b/src/arith/detect_linear_equation.cc
@@ -32,6 +32,7 @@
namespace tvm {
namespace arith {
+using namespace tvm::prim;
using namespace tirx;
@@ -102,7 +103,7 @@ class LinearEqDetector : public tvm::ExprFunctor<LinearEqEntry(const Expr&, cons
LinearEqEntry ret;
if (op == var_.get()) {
PrimType dtype = op->ty.as_or_throw<PrimType>();
- ret.coeff = MakeConst(PrimType::Int(dtype.bits(), dtype.lanes()), 1);
+ ret.coeff = prim::MakeConst(PrimType::Int(dtype.bits(), dtype.lanes()), 1);
} else {
ret.base = e;
}
@@ -206,14 +207,14 @@ bool DetectClipBound(const PrimExpr& cond,
if (const prim::LTNode* op = cond.as<prim::LTNode>()) {
PrimType a_ty = op->a.ty();
if (!a_ty.MatchesCode(DLDataTypeCode::kDLInt)) return false;
- canonical = op->b - op->a - MakeConst(a_ty, 1);
+ canonical = op->b - op->a - prim::MakeConst(a_ty, 1);
} else if (const prim::LENode* op = cond.as<prim::LENode>()) {
if (!op->a.ty().MatchesCode(DLDataTypeCode::kDLInt)) return false;
canonical = op->b - op->a;
} else if (const prim::GTNode* op = cond.as<prim::GTNode>()) {
PrimType a_ty = op->a.ty();
if (!a_ty.MatchesCode(DLDataTypeCode::kDLInt)) return false;
- canonical = op->a - op->b - MakeConst(a_ty, 1);
+ canonical = op->a - op->b - prim::MakeConst(a_ty, 1);
} else if (const prim::GENode* op = cond.as<prim::GENode>()) {
if (!op->a.ty().MatchesCode(DLDataTypeCode::kDLInt)) return false;
canonical = op->a - op->b;
diff --git a/src/arith/int_operator.h b/src/arith/int_operator.h
index 1bb09bb394..bea3ea92af 100644
--- a/src/arith/int_operator.h
+++ b/src/arith/int_operator.h
@@ -24,9 +24,13 @@
#ifndef TVM_ARITH_INT_OPERATOR_H_
#define TVM_ARITH_INT_OPERATOR_H_
+#include <tvm/ir/prim/expr.h>
+
#include <limits>
#include <utility>
+#include "../ir/prim/int_operator.h"
+
namespace tvm {
namespace arith {
@@ -81,46 +85,10 @@ inline bool WillOverflow<prim::ModNode>(int64_t x, int64_t y, int64_t min_value,
return y == 0;
}
-/*!
- * \brief Perform trunc division of two integers.
- * \param x The left operand.
- * \param y The right operand.
- * \return the result.
- */
-inline int64_t truncdiv(int64_t x, int64_t y) { return x / y; }
-
-/*!
- * \brief Compute the truncdiv remainder of two integers.
- * \param x The left operand.
- * \param y The right operand.
- * \return the result.
- */
-inline int64_t truncmod(int64_t x, int64_t y) { return x % y; }
-
-/*!
- * \brief Perform floor division of two integers.
- * \param x The left operand.
- * \param y The right operand.
- * \return the result.
- */
-inline int64_t floordiv(int64_t x, int64_t y) {
- int64_t rdiv = x / y;
- int64_t rmod = x % y;
- bool is_floor_div = (y >= 0 && rmod >= 0) || (y < 0 && rmod <= 0);
- return is_floor_div ? rdiv : (rdiv - 1);
-}
-
-/*!
- * \brief Compute the floordiv remainder of two integers.
- * \param x The left operand.
- * \param y The right operand.
- * \return the result.
- */
-inline int64_t floormod(int64_t x, int64_t y) {
- int64_t rmod = x % y;
- bool is_floor_div = (y >= 0 && rmod >= 0) || (y < 0 && rmod <= 0);
- return is_floor_div ? rmod : rmod + y;
-}
+using prim::detail::floordiv;
+using prim::detail::floormod;
+using prim::detail::truncdiv;
+using prim::detail::truncmod;
/*!
* \brief Use Extended Euclidean algorithm to solve ax + by = gcd(a, b)
diff --git a/src/arith/int_set.cc b/src/arith/int_set.cc
index 3fe11eb61d..2e7655fc05 100644
--- a/src/arith/int_set.cc
+++ b/src/arith/int_set.cc
@@ -29,8 +29,8 @@
#include <tvm/ffi/reflection/registry.h>
#include <tvm/ir/expr_functor.h>
#include <tvm/ir/prim/expr.h>
+#include <tvm/ir/prim/op.h>
#include <tvm/runtime/logging.h>
-#include <tvm/tirx/op.h>
#include <algorithm>
#include <unordered_map>
@@ -44,16 +44,14 @@
namespace tvm {
namespace arith {
+using namespace tvm::prim;
-using tirx::is_one;
-using tirx::is_zero;
-using tirx::MakeConst;
+using prim::is_one;
+using prim::is_zero;
+using prim::MakeConst;
TVM_FFI_STATIC_INIT_BLOCK() { IntervalSetNode::RegisterReflection(); }
-PrimExpr SymbolicLimits::pos_inf_ = PrimVar("pos_inf", PrimType::Int(64));
-PrimExpr SymbolicLimits::neg_inf_ = PrimVar("neg_inf", PrimType::Int(64));
-
IntervalSet::IntervalSet(PrimExpr min_value, PrimExpr max_value) {
auto node = ffi::make_object<IntervalSetNode>();
node->min_value = std::move(min_value);
@@ -261,7 +259,7 @@ inline IntervalSet Combine<prim::Mod>(AnalyzerObj* analyzer, IntervalSet a, Inte
if (analyzer->CanProveGreaterEqual(divisor, 0)) {
return IntervalSet(IntImm(divisor.ty(), 0), divisor - 1);
} else {
- PrimExpr bound = abs(divisor) - 1;
+ PrimExpr bound = prim::IntegerAbs(divisor) - 1;
return IntervalSet(-bound, bound);
}
}
@@ -358,7 +356,7 @@ inline IntervalSet Combine<prim::FloorMod>(AnalyzerObj* analyzer, IntervalSet a,
}
return IntervalSet(IntImm(divisor.ty(), 0), divisor - 1);
} else {
- PrimExpr bound = abs(divisor) - 1;
+ PrimExpr bound = prim::IntegerAbs(divisor) - 1;
return IntervalSet(-bound, bound);
}
}
@@ -534,13 +532,13 @@ class IntervalSetEvaluator : public tvm::ExprFunctor<IntervalSet(const Expr&)> {
if (op->lanes->IsInstance<IntImmNode>()) {
int lanes = static_cast<int>(op->lanes.as_or_throw<IntImm>()->value);
if (vstride > 0) {
- PrimExpr stride_expr = MakeConst(t, vstride * (lanes - 1));
+ PrimExpr stride_expr = prim::MakeConst(t, vstride * (lanes - 1));
auto add_op = prim::Add(op->base, stride_expr);
auto add_node = add_op.as<prim::AddNode>();
return Combine<prim::Add>(analyzer_, base, IntervalSet(IntImm(t, 0), stride_expr),
add_node);
} else {
- PrimExpr stride_expr = MakeConst(t, vstride * (lanes - 1));
+ PrimExpr stride_expr = prim::MakeConst(t, vstride * (lanes - 1));
auto add_op = prim::Add(op->base, stride_expr);
auto add_node = add_op.as<prim::AddNode>();
return Combine<prim::Add>(analyzer_, base, IntervalSet(stride_expr, IntImm(t, 0)),
@@ -568,13 +566,14 @@ class IntervalSetEvaluator : public tvm::ExprFunctor<IntervalSet(const Expr&)> {
// short cut for the int set.
if (value_set->min_value.same_as(value_set->max_value)) {
if (value_set->IsEmpty()) return value_set;
- return IntervalSet::SinglePoint(cast(op->ty.as_or_throw<PrimType>(), value_set->min_value));
+ return IntervalSet::SinglePoint(
+ prim::cast(op->ty.as_or_throw<PrimType>(), value_set->min_value));
}
PrimExpr min_value = value_set->HasLowerBound()
- ? cast(op->ty.as_or_throw<PrimType>(), value_set->min_value)
+ ? prim::cast(op->ty.as_or_throw<PrimType>(), value_set->min_value)
: neg_inf();
PrimExpr max_value = value_set->HasUpperBound()
- ? cast(op->ty.as_or_throw<PrimType>(), value_set->max_value)
+ ? prim::cast(op->ty.as_or_throw<PrimType>(), value_set->max_value)
: pos_inf();
return IntervalSet(min_value, max_value);
}
diff --git a/src/arith/interval_set.h b/src/arith/interval_set.h
index 72820c178e..b0fa9edf02 100644
--- a/src/arith/interval_set.h
+++ b/src/arith/interval_set.h
@@ -27,7 +27,7 @@
#include <tvm/arith/analyzer.h>
#include <tvm/ffi/reflection/registry.h>
#include <tvm/ir/cow.h>
-#include <tvm/tirx/op.h>
+#include <tvm/ir/prim/op.h>
#include <limits>
diff --git a/src/arith/iter_affine_map.cc b/src/arith/iter_affine_map.cc
index 16a14412df..b4f25a2b13 100644
--- a/src/arith/iter_affine_map.cc
+++ b/src/arith/iter_affine_map.cc
@@ -40,6 +40,7 @@
namespace tvm {
namespace arith {
+using namespace tvm::prim;
using namespace tirx;
@@ -67,7 +68,7 @@ TVM_FFI_STATIC_INIT_BLOCK() {
IterSplitExpr::IterSplitExpr(IterMark source) {
auto n = ffi::make_object<IterSplitExprNode>();
- auto one = MakeConst(source->source.ty(), 1);
+ auto one = prim::MakeConst(source->source.ty(), 1);
n->ExprNode::ty = source->source.ty();
n->source = std::move(source);
n->extent = n->source->extent;
@@ -78,7 +79,7 @@ IterSplitExpr::IterSplitExpr(IterMark source) {
IterSplitExpr::IterSplitExpr(IterMark source, PrimExpr scale) {
auto n = ffi::make_object<IterSplitExprNode>();
- auto one = MakeConst(source->source.ty(), 1);
+ auto one = prim::MakeConst(source->source.ty(), 1);
n->ExprNode::ty = source->source.ty();
n->source = std::move(source);
n->extent = n->source->extent;
@@ -2678,7 +2679,7 @@ class InverseAffineIterMapTransformer {
// Case 1: Propagate to the input node directly when the sum expression has only one components
if (iter_map_expr->args.size() == 1) {
const auto& source = iter_map_expr->args[0];
- TVM_FFI_ICHECK(analyzer_->CanProveEqual(abs(source->scale), 1));
+ TVM_FFI_ICHECK(analyzer_->CanProveEqual(prim::IntegerAbs(source->scale), 1));
backprop_.Set(source, (backprop_.at(source) + input) * source->scale);
return;
}
diff --git a/src/arith/modular_set.cc b/src/arith/modular_set.cc
index 6b8de6f36f..159bd7b4e6 100644
--- a/src/arith/modular_set.cc
+++ b/src/arith/modular_set.cc
@@ -38,6 +38,7 @@
namespace tvm {
namespace arith {
+using namespace tvm::prim;
using namespace tirx;
diff --git a/src/arith/pattern_match.h b/src/arith/pattern_match.h
index fa105b158e..0b71ff0f6d 100644
--- a/src/arith/pattern_match.h
+++ b/src/arith/pattern_match.h
@@ -68,6 +68,7 @@
#include <tvm/ffi/cast.h>
#include <tvm/ir/prim/builtin.h>
#include <tvm/ir/prim/expr.h>
+#include <tvm/ir/prim/op.h>
#include <tvm/tirx/builtin.h>
#include <cmath>
@@ -379,7 +380,7 @@ class PConstWithTypeLike : public Pattern<PConstWithTypeLike<TA>> {
}
}
- PrimExpr Eval() const { return tirx::MakeConst(ref_.Eval().ty(), value_); }
+ PrimExpr Eval() const { return prim::MakeConst(ref_.Eval().ty(), value_); }
private:
typename TA::Nested ref_;
diff --git a/src/arith/presburger_set.cc b/src/arith/presburger_set.cc
index be79bd3d10..35064eac36 100644
--- a/src/arith/presburger_set.cc
+++ b/src/arith/presburger_set.cc
@@ -41,6 +41,7 @@
namespace tvm {
namespace arith {
+using namespace tvm::prim;
#if defined(TVM_MLIR_VERSION) && TVM_MLIR_VERSION >= 150
diff --git a/src/arith/presburger_set.h b/src/arith/presburger_set.h
index 61a39e710b..a3fc969290 100644
--- a/src/arith/presburger_set.h
+++ b/src/arith/presburger_set.h
@@ -35,7 +35,7 @@
#include <tvm/arith/analyzer.h>
#include <tvm/ffi/reflection/registry.h>
#include <tvm/ir/cow.h>
-#include <tvm/tirx/op.h>
+#include <tvm/ir/prim/op.h>
#include <limits>
#include <vector>
diff --git a/src/arith/product_normal_form.h b/src/arith/product_normal_form.h
index a7da34403f..03c9dd98b3 100644
--- a/src/arith/product_normal_form.h
+++ b/src/arith/product_normal_form.h
@@ -25,7 +25,7 @@
#define TVM_ARITH_PRODUCT_NORMAL_FORM_H_
#include <tvm/ir/prim/expr.h>
-#include <tvm/tirx/op.h>
+#include <tvm/ir/prim/op.h>
namespace tvm {
namespace arith {
diff --git a/src/arith/rewrite_simplify.cc b/src/arith/rewrite_simplify.cc
index c6054a81ad..b3be3128a5 100644
--- a/src/arith/rewrite_simplify.cc
+++ b/src/arith/rewrite_simplify.cc
@@ -1756,10 +1756,10 @@ ffi::Optional<PrimExpr> RewriteSimplifier::Impl::TryMatchLiteralConstraint(
prim::ExprDeepEqual expr_equal;
for (const auto& constraint : literal_constraints_) {
if (expr_equal(constraint, expr)) {
- return MakeConst(expr.ty(), true);
+ return prim::MakeConst(expr.ty(), true);
}
if (expr_equal(constraint, negation)) {
- return MakeConst(expr.ty(), false);
+ return prim::MakeConst(expr.ty(), false);
}
}
return std::nullopt;
@@ -1790,7 +1790,7 @@ PrimExpr RewriteSimplifier::Impl::ApplyRewriteRules(prim::EQ ret, InplaceMode in
// Pattern var match IntImm
PVar<IntImm> c1, c2;
PVar<PrimExpr> lanes;
- PConst<PrimExpr> ctrue(MakeConst(ret->ty.as_or_throw<PrimType>(), true));
+ PConst<PrimExpr> ctrue(prim::MakeConst(ret->ty.as_or_throw<PrimType>(), true));
// vector rule
if (IsVectorExpr(ret.get())) {
@@ -1800,10 +1800,10 @@ PrimExpr RewriteSimplifier::Impl::ApplyRewriteRules(prim::EQ ret, InplaceMode in
if (IsIndexTypedExpr(ret->a)) {
CompareResult result = TryCompare(ret->a, ret->b);
if (result == CompareResult::kEQ) {
- return MakeConst(ret->ty.as_or_throw<PrimType>(), true);
+ return prim::MakeConst(ret->ty.as_or_throw<PrimType>(), true);
} else if (result == CompareResult::kNE || result == CompareResult::kGT ||
result == CompareResult::kLT) {
- return MakeConst(ret->ty.as_or_throw<PrimType>(), false);
+ return prim::MakeConst(ret->ty.as_or_throw<PrimType>(), false);
}
TVM_TRY_REWRITE(c1 == x, x == c1);
@@ -1839,9 +1839,9 @@ UnchangedOr<PrimExpr> RewriteSimplifier::Impl::Mutate_(const prim::NENode* op,
CompareResult result = TryCompare(op->a, op->b);
if (result == CompareResult::kNE || result == CompareResult::kGT ||
result == CompareResult::kLT) {
- return MakeConst(op->ty.as_or_throw<PrimType>(), true);
+ return prim::MakeConst(op->ty.as_or_throw<PrimType>(), true);
} else if (result == CompareResult::kEQ) {
- return MakeConst(op->ty.as_or_throw<PrimType>(), false);
+ return prim::MakeConst(op->ty.as_or_throw<PrimType>(), false);
} else if (result == CompareResult::kGE) {
// Known: a >= b
//
@@ -1886,9 +1886,9 @@ UnchangedOr<PrimExpr> RewriteSimplifier::Impl::Mutate_(const prim::LENode* op,
CompareResult result = TryCompare(op->a, op->b);
if (result == CompareResult::kLE || result == CompareResult::kLT ||
result == CompareResult::kEQ) {
- return MakeConst(op->ty.as_or_throw<PrimType>(), true);
+ return prim::MakeConst(op->ty.as_or_throw<PrimType>(), true);
} else if (result == CompareResult::kGT) {
- return MakeConst(op->ty.as_or_throw<PrimType>(), false);
+ return prim::MakeConst(op->ty.as_or_throw<PrimType>(), false);
} else if (result == CompareResult::kNE) {
// Known: a != b
//
@@ -1954,11 +1954,11 @@ PrimExpr RewriteSimplifier::Impl::ApplyRewriteRules(prim::LT ret, InplaceMode in
if (IsIndexTypedExpr(ret->a)) {
CompareResult result = TryCompare(ret->a, ret->b);
if (result == CompareResult::kLT) {
- return MakeConst(ret->ty.as_or_throw<PrimType>(), true);
+ return prim::MakeConst(ret->ty.as_or_throw<PrimType>(), true);
}
if (result == CompareResult::kEQ || result == CompareResult::kGT ||
result == CompareResult::kGE) {
- return MakeConst(ret->ty.as_or_throw<PrimType>(), false);
+ return prim::MakeConst(ret->ty.as_or_throw<PrimType>(), false);
}
// clang-format off
@@ -2201,7 +2201,7 @@ UnchangedOr<PrimExpr> RewriteSimplifier::Impl::Mutate_(const prim::AndNode* op,
TVM_TRY_REWRITE(broadcast(x, lanes) && broadcast(y, lanes), broadcast(x && y, lanes));
}
- auto cfalse = PConst<PrimExpr>(MakeConst(op->ty.as_or_throw<PrimType>(), false));
+ auto cfalse = PConst<PrimExpr>(prim::MakeConst(op->ty.as_or_throw<PrimType>(), false));
TVM_TRY_REWRITE(x == y && x != y, cfalse);
TVM_TRY_REWRITE(x != y && x == y, cfalse);
TVM_TRY_REWRITE(x && !x, cfalse);
@@ -2351,7 +2351,7 @@ UnchangedOr<PrimExpr> RewriteSimplifier::Impl::Mutate_(const prim::OrNode* op,
TVM_TRY_REWRITE(broadcast(x, lanes) || broadcast(y, lanes), broadcast(x || y, lanes));
}
- auto ctrue = PConst<PrimExpr>(MakeConst(op->ty.as_or_throw<PrimType>(), true));
+ auto ctrue = PConst<PrimExpr>(prim::MakeConst(op->ty.as_or_throw<PrimType>(), true));
TVM_TRY_REWRITE(x == y || x != y, ctrue);
TVM_TRY_REWRITE(x != y || x == y, ctrue);
@@ -2410,7 +2410,7 @@ UnchangedOr<Expr> RewriteSimplifier::Impl::Mutate_(const CallNode* op, InplaceMo
if (op == nullptr) return ret;
if (op->op.same_as(prim::builtin::likely()) &&
- is_const_int(op->args[0].as_or_throw<PrimExpr>())) {
+ prim::is_const_int(op->args[0].as_or_throw<PrimExpr>())) {
return op->args[0].as_or_throw<PrimExpr>();
} else if (op->op.same_as(prim::builtin::shift_right())) {
if (op->args[0].as<IntImmNode>() && op->args[1].as<IntImmNode>()) {
@@ -2423,17 +2423,17 @@ UnchangedOr<Expr> RewriteSimplifier::Impl::Mutate_(const CallNode* op, InplaceMo
return op->args[0].as_or_throw<PrimExpr>() << op->args[1].as_or_throw<PrimExpr>();
}
}
- static const Op& ceil_op = Op::Get("tirx.ceil");
- static const Op& log2_op = Op::Get("tirx.log2");
+ static const Op& ceil_op = prim::builtin::ceil();
+ static const Op& log2_op = prim::builtin::log2();
static const Op& clz_op = Op::Get("tirx.clz");
PrimType ret_ty = op->ty.as_or_throw<PrimType>();
if (op->op.same_as(ceil_op)) {
PrimExpr ceil_arg = op->args[0].as_or_throw<PrimExpr>();
if (auto arg_int = op->args[0].as<IntImmNode>()) {
- return cast(ret_ty, IntImm(arg_int->ty.as_or_throw<PrimType>(), arg_int->value));
+ return prim::cast(ret_ty, IntImm(arg_int->ty.as_or_throw<PrimType>(), arg_int->value));
} else if (auto arg_float = ceil_arg.as<FloatImmNode>()) {
- return cast(ret_ty,
- FloatImm(arg_float->ty.as_or_throw<PrimType>(), std::ceil(arg_float->value)));
+ return prim::cast(
+ ret_ty, FloatImm(arg_float->ty.as_or_throw<PrimType>(), std::ceil(arg_float->value)));
} else if (auto arg_call = ceil_arg.as<CallNode>()) {
// ceil(log2(cast(n,"float64"))) is used as the implementation of
// topi.math.ceil_log2, and appears in iteration bounds.
@@ -2480,7 +2480,7 @@ UnchangedOr<Expr> RewriteSimplifier::Impl::Mutate_(const CallNode* op, InplaceMo
PrimExpr inner_then_expr = inner_call->args[1].as_or_throw<PrimExpr>();
PrimExpr inner_else_expr = inner_call->args[2].as_or_throw<PrimExpr>();
// Only check constant cases to avoid recursion
- if (is_const_number(inner_else_expr) && is_const_number(else_expr) &&
+ if (prim::is_const_number(inner_else_expr) && prim::is_const_number(else_expr) &&
analyzer_->CanProve(inner_else_expr == else_expr)) {
return Call(ret_ty, op->op, {cond && inner_cond, inner_then_expr, else_expr}, op->attrs,
op->ty_args, op->span)
@@ -2520,13 +2520,13 @@ UnchangedOr<PrimExpr> RewriteSimplifier::Impl::Mutate_(const prim::CastNode* op,
PrimExpr ret =
SimplifierBase::Mutate_(op, inplace_mode).ValueOrUnchanged(ffi::GetRef<PrimExpr>(op));
op = ret.as<prim::CastNode>();
- return cast(ret.ty(), op->value);
+ return prim::cast(ret.ty(), op->value);
}
bool RewriteSimplifier::Impl::CanInlineLet(const prim::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 (prim::is_const_number(op->value)) return true;
if (op->value.as<PrimVar>()) return true;
return false;
}
diff --git a/src/arith/rewrite_simplify.h b/src/arith/rewrite_simplify.h
index 78cf06e403..23b9edde1d 100644
--- a/src/arith/rewrite_simplify.h
+++ b/src/arith/rewrite_simplify.h
@@ -27,7 +27,7 @@
#include <tvm/arith/analyzer.h>
#include <tvm/ffi/reflection/registry.h>
#include <tvm/ir/cow.h>
-#include <tvm/tirx/op.h>
+#include <tvm/ir/prim/op.h>
#include <algorithm>
#include <unordered_map>
diff --git a/src/arith/simplify_base.cc b/src/arith/simplify_base.cc
index 29d8aaf24b..3427c7f6b0 100644
--- a/src/arith/simplify_base.cc
+++ b/src/arith/simplify_base.cc
@@ -21,7 +21,7 @@
#include <tvm/ir/expr.h>
#include <tvm/ir/prim/builtin.h>
-#include <tvm/tirx/op.h>
+#include <tvm/ir/prim/op.h>
#include "constraint_helpers.h"
@@ -101,8 +101,8 @@ UnchangedOr<Expr> SimplifierBase::Mutate_(const CallNode* op, InplaceMode inplac
.ValueOrUnchanged((*args)[2])
.as_or_throw<Expr>();
}
- if (tirx::is_zero(cond)) return false_value;
- if (tirx::is_one(cond)) return true_value;
+ if (prim::is_zero(cond)) return false_value;
+ if (prim::is_one(cond)) return true_value;
if (cond.same_as(op->args[0]) && true_value.same_as(op->args[1]) &&
false_value.same_as(op->args[2])) {
// Reuse the original node identity; there is no replacement to process.
@@ -153,8 +153,8 @@ UnchangedOr<PrimExpr> SimplifierBase::Mutate_(const prim::SelectNode* op,
})
.ValueOrUnchanged(op->false_value);
}
- if (tirx::is_zero(cond)) return false_value;
- if (tirx::is_one(cond)) return true_value;
+ if (prim::is_zero(cond)) return false_value;
+ if (prim::is_one(cond)) return true_value;
if (cond.same_as(op->condition) && true_value.same_as(op->true_value) &&
false_value.same_as(op->false_value)) {
// Reuse the original node identity; there is no replacement to process.
diff --git a/src/arith/transitive_comparison_analyzer.cc b/src/arith/transitive_comparison_analyzer.cc
index ee96ca427d..6b524efcaf 100644
--- a/src/arith/transitive_comparison_analyzer.cc
+++ b/src/arith/transitive_comparison_analyzer.cc
@@ -32,6 +32,9 @@
namespace tvm {
namespace arith {
+using namespace tvm::prim;
+
+using prim::is_const_int;
using namespace tirx;
diff --git a/src/arith/z3_prover.cc b/src/arith/z3_prover.cc
index ba14e09727..4e9556a51a 100644
--- a/src/arith/z3_prover.cc
+++ b/src/arith/z3_prover.cc
@@ -58,6 +58,7 @@
#include "z3++.h"
namespace tvm::arith {
+using namespace tvm::prim;
using namespace tirx;
using namespace ffi;
@@ -434,9 +435,9 @@ class Z3Prover::Impl : tvm::ExprFunctor<z3::expr(const Expr&)> {
// test is_const_int on min and extent individually and add the two constants
// in C++. Otherwise this fast path is never taken and we always emit the more expensive
// symbolic constraint below.
- if (tirx::is_const_int(min) && tirx::is_const_int(extent)) {
- int64_t min_value = *tirx::as_const_int(min);
- int64_t extent_value = *tirx::as_const_int(extent);
+ if (prim::is_const_int(min) && prim::is_const_int(extent)) {
+ int64_t min_value = *prim::as_const_int(min);
+ int64_t extent_value = *prim::as_const_int(extent);
int64_t max_value = min_value + extent_value;
if (min_value < max_value) {
solver->add(ctx->int_val(min_value) <= var_expr);
diff --git a/src/backend/cuda/codegen/codegen_cuda.cc b/src/backend/cuda/codegen/codegen_cuda.cc
index eb0e573094..d60eea8d29 100644
--- a/src/backend/cuda/codegen/codegen_cuda.cc
+++ b/src/backend/cuda/codegen/codegen_cuda.cc
@@ -46,6 +46,7 @@
namespace tvm {
namespace codegen {
+using namespace tvm::prim;
namespace {
@@ -1341,7 +1342,7 @@ void CodeGenCUDA::Dispatch_(const CallNode* op, std::ostream& os) {
stream << ": \"l\"((void*)(" << global_buffer << "+" << global_addr << ")), \"r\"((int)"
<< guard << ")\n";
stream << ");\n";
- } else if (op->op.same_as(builtin::reinterpret())) {
+ } else if (op->op.same_as(tirx::builtin::reinterpret())) {
// Compile-time pointer reinterpret of a literal (e.g. a tcgen05 descriptor
// template encoded at address 0): emit C++-style reinterpret_cast<T*>(...)
// to match the encoded template. Runtime pointer reinterprets fall through
@@ -1454,7 +1455,7 @@ void CodeGenCUDA::Dispatch_(const CallNode* op, std::ostream& os) {
<< "Invalid number of lanes for float4_e2m1fn reinterpret: " << lanes;
}
EndScope(ssa_scope);
- } else if (op->op.same_as(builtin::print_buffer())) {
+ } else if (op->op.same_as(tirx::builtin::print_buffer())) {
TVM_FFI_ICHECK_GE(op->args.size(), 5U) << "Print operation expects at least 5 arguments";
Expr arg = op->args[0];
@@ -1604,7 +1605,7 @@ void CodeGenCUDA::Dispatch_(const CallNode* op, std::ostream& os) {
} else if (op->op.same_as(cuda_func_call_op) ||
(op->op.as<Op>() && op->op.as<Op>().value()->name == "tirx.cuda.func_call")) {
print_cuda_func_call(op, os);
- } else if (op->op.same_as(builtin::thread_return())) {
+ } else if (op->op.same_as(tirx::builtin::thread_return())) {
os << "return";
} else {
CodeGenC::Dispatch_(op, os);
@@ -1750,7 +1751,7 @@ void CodeGenCUDA::VisitStmt_(const AllocBufferNode* op) {
void CodeGenCUDA::VisitStmt_(const EvaluateNode* op) {
if (auto value = op->value.as<PrimExpr>(); value && is_const_int(value.value())) return;
const CallNode* call = op->value.as<CallNode>();
- if (call && call->op.same_as(builtin::tvm_global_barrier_kinit())) {
+ if (call && call->op.same_as(tirx::builtin::tvm_global_barrier_kinit())) {
PrintIndent();
stream << "__shared__ unsigned " << vid_global_barrier_expect_ << ";\n";
PrintIndent();
diff --git a/src/backend/cuda/codegen/intrin_rule_cuda.cc b/src/backend/cuda/codegen/intrin_rule_cuda.cc
index bd1bd8308c..4a57f2f469 100644
--- a/src/backend/cuda/codegen/intrin_rule_cuda.cc
+++ b/src/backend/cuda/codegen/intrin_rule_cuda.cc
@@ -172,7 +172,7 @@ TVM_REGISTER_OP("tirx.clz")
TVM_REGISTER_OP("tirx.floor")
.set_attr<FLowerIntrinsic>("cuda.FLowerIntrinsic", DispatchPureExtern<CUDAMath>);
-TVM_REGISTER_OP("tirx.ceil")
+TVM_REGISTER_OP("prim.ceil")
.set_attr<FLowerIntrinsic>("cuda.FLowerIntrinsic", DispatchPureExtern<CUDAMath>);
TVM_REGISTER_OP("tirx.trunc")
@@ -205,7 +205,7 @@ TVM_REGISTER_OP("tirx.log")
.set_attr<FLowerIntrinsic>("cuda.fastmath.FLowerIntrinsic", DispatchPureExtern<CUDAFastMath>)
.set_attr<FLowerIntrinsic>("cuda.FLowerIntrinsic", DispatchPureExtern<CUDAMath>);
-TVM_REGISTER_OP("tirx.log2")
+TVM_REGISTER_OP("prim.log2")
.set_attr<FLowerIntrinsic>("cuda.fastmath.FLowerIntrinsic", DispatchPureExtern<CUDAFastMath>)
.set_attr<FLowerIntrinsic>("cuda.FLowerIntrinsic", DispatchPureExtern<CUDAMath>);
diff --git a/src/backend/cuda/codegen/llvm/intrin_rule_nvptx.cc b/src/backend/cuda/codegen/llvm/intrin_rule_nvptx.cc
index 5f76ee5a35..1758cdd874 100644
--- a/src/backend/cuda/codegen/llvm/intrin_rule_nvptx.cc
+++ b/src/backend/cuda/codegen/llvm/intrin_rule_nvptx.cc
@@ -46,7 +46,8 @@ inline PrimExpr DispatchPureExternLibDevice(const PrimExpr& e) {
const OpNode* op = call->op.as<OpNode>();
TVM_FFI_ICHECK(op != nullptr);
std::string name = op->name;
- TVM_FFI_ICHECK_EQ(name.substr(0, 5), "tirx.");
+ TVM_FFI_ICHECK(name.substr(0, 5) == "tirx." || name == "prim.ceil" || name == "prim.log2")
+ << "Unexpected intrinsic name: " << name;
std::ostringstream intrinsic_name;
intrinsic_name << "__nv_" << name.substr(5);
@@ -63,7 +64,7 @@ using tirx::FLowerIntrinsic;
TVM_REGISTER_OP("tirx.floor")
.set_attr<FLowerIntrinsic>("nvptx.FLowerIntrinsic", DispatchPureExternLibDevice);
-TVM_REGISTER_OP("tirx.ceil")
+TVM_REGISTER_OP("prim.ceil")
.set_attr<FLowerIntrinsic>("nvptx.FLowerIntrinsic", DispatchPureExternLibDevice);
TVM_REGISTER_OP("tirx.round")
@@ -105,7 +106,7 @@ TVM_REGISTER_OP("tirx.fma")
TVM_REGISTER_OP("tirx.log")
.set_attr<FLowerIntrinsic>("nvptx.FLowerIntrinsic", DispatchPureExternLibDevice);
-TVM_REGISTER_OP("tirx.log2")
+TVM_REGISTER_OP("prim.log2")
.set_attr<FLowerIntrinsic>("nvptx.FLowerIntrinsic", DispatchPureExternLibDevice);
TVM_REGISTER_OP("tirx.log10")
diff --git a/src/backend/cuda/transforms/lower_iket.cc b/src/backend/cuda/transforms/lower_iket.cc
index a4324f85c8..b24c3899b6 100644
--- a/src/backend/cuda/transforms/lower_iket.cc
+++ b/src/backend/cuda/transforms/lower_iket.cc
@@ -45,6 +45,7 @@
namespace tvm {
namespace tirx {
+
namespace transform {
namespace {
@@ -1118,14 +1119,14 @@ class InstrumentOfficialKernel : public StmtExprMutator {
PrimExpr Event(PrimExpr event_id) const {
static const Op& event_op = Op::Get("tirx.cuda.iket_official_event");
return Call(PrimType::UInt(32), event_op,
- {cast(PrimType::UInt(32), event_id), prim::StringImm(device_source_)});
+ {prim::cast(PrimType::UInt(32), event_id), prim::StringImm(device_source_)});
}
PrimExpr Event(PrimExpr event_id, PrimExpr payload) const {
static const Op& event_op = Op::Get("tirx.cuda.iket_official_event");
- return Call(
- PrimType::UInt(32), event_op,
- {cast(PrimType::UInt(32), event_id), prim::StringImm(device_source_), std::move(payload)});
+ return Call(PrimType::UInt(32), event_op,
+ {prim::cast(PrimType::UInt(32), event_id), prim::StringImm(device_source_),
+ std::move(payload)});
}
PrimExpr NormalizePayload(PrimExpr payload, PayloadType type) const {
@@ -1143,7 +1144,7 @@ class InstrumentOfficialKernel : public StmtExprMutator {
PrimExpr payload = NormalizePayload(
Mutate(call->args[1]).ValueOrUnchanged(call->args[1]).as_or_throw<PrimExpr>(),
payload_type);
- return IfThenElse(token != 0, Evaluate(Event(token, std::move(payload))));
+ return IfThenElse(not_equal(token, 0), Evaluate(Event(token, std::move(payload))));
}
return Evaluate(Event(token));
}
diff --git a/src/backend/hexagon/codegen/llvm/intrin_rule_hexagon.cc b/src/backend/hexagon/codegen/llvm/intrin_rule_hexagon.cc
index 6782be8ba9..2bb41c40fb 100644
--- a/src/backend/hexagon/codegen/llvm/intrin_rule_hexagon.cc
+++ b/src/backend/hexagon/codegen/llvm/intrin_rule_hexagon.cc
@@ -34,6 +34,8 @@
namespace tvm {
namespace codegen {
+using namespace tvm::prim;
+
namespace llvm {
using tirx::FLowerIntrinsic;
@@ -142,9 +144,9 @@ TVM_REGISTER_OP("tirx.tanh")
return TVMExternCall(call, tvm_wrapper);
}
#endif
- PrimExpr one = tirx::MakeConst(x_ty, 1);
- PrimExpr two = tirx::MakeConst(x_ty, 2);
- PrimExpr neg_two = tirx::MakeConst(x_ty, -2);
+ PrimExpr one = tvm::prim::MakeConst(x_ty, 1);
+ PrimExpr two = tvm::prim::MakeConst(x_ty, 2);
+ PrimExpr neg_two = tvm::prim::MakeConst(x_ty, -2);
PrimExpr exp_neg2x = exp(neg_two * x);
PrimExpr exp_pos2x = exp(two * x);
@@ -152,7 +154,7 @@ TVM_REGISTER_OP("tirx.tanh")
PrimExpr tanh_pos = (one - exp_neg2x) / (one + exp_neg2x);
PrimExpr tanh_neg = (exp_pos2x - one) / (exp_pos2x + one);
// MakeConst can handle both vector and scalar types.
- PrimExpr tanh_x = prim::Select(x >= tirx::MakeConst(x_ty, 0), tanh_pos, tanh_neg);
+ PrimExpr tanh_x = prim::Select(x >= tvm::prim::MakeConst(x_ty, 0), tanh_pos, tanh_neg);
return tanh_x;
});
@@ -205,8 +207,8 @@ TVM_REGISTER_OP("tirx.sigmoid")
useqhl = tstring.find("+hvx-qfloat") != std::string::npos;
}
- PrimExpr MinBound = tirx::MakeConst(x_ty, -8);
- PrimExpr MaxBound = tirx::MakeConst(x_ty, 8);
+ PrimExpr MinBound = tvm::prim::MakeConst(x_ty, -8);
+ PrimExpr MaxBound = tvm::prim::MakeConst(x_ty, 8);
const PrimExpr v1 = prim::Max(x, MinBound);
const PrimExpr v2 = prim::Min(v1, MaxBound);
@@ -221,11 +223,14 @@ TVM_REGISTER_OP("tirx.sigmoid")
return TVMExternCall(new_call.get(), tvm_wrapper);
}
#endif
- PrimExpr one = tirx::MakeConst(x_ty, 1);
+ PrimExpr one = tvm::prim::MakeConst(x_ty, 1);
return one / (one + exp(-x));
});
-TVM_REGISTER_QHL_OP_FP16(ceil, 1)
+TVM_REGISTER_OP("prim.ceil")
+ .set_attr<FLowerIntrinsic>(
+ "hexagon.FLowerIntrinsic",
+ DispatchTVMQHLWrapperFp16<tvm_qhl_ahf_ceil, ::llvm::Intrinsic::ceil, 1>);
TVM_REGISTER_QHL_OP_FP16(cos, 1)
diff --git a/src/backend/metal/codegen/codegen_metal.cc b/src/backend/metal/codegen/codegen_metal.cc
index f6755b7a85..e25d33db26 100644
--- a/src/backend/metal/codegen/codegen_metal.cc
+++ b/src/backend/metal/codegen/codegen_metal.cc
@@ -44,6 +44,7 @@
namespace tvm {
namespace codegen {
+using namespace tvm::prim;
namespace {
@@ -487,9 +488,9 @@ void CodeGenMetal::Dispatch_(const CallNode* op, std::ostream& os) { // NOLINT(
<< PrintExpr(a) << "[" << PrintExpr(op->args[3]) << "], " //
<< PrintExpr(b) << "[" << PrintExpr(op->args[5]) << "], " //
<< PrintExpr(c) << "[" << PrintExpr(op->args[7]) << "])";
- } else if (op->op.same_as(builtin::ptr_byte_offset()) ||
- op->op.same_as(builtin::handle_add_byte_offset())) {
- bool is_typed_offset = op->op.same_as(builtin::ptr_byte_offset());
+ } else if (op->op.same_as(tirx::builtin::ptr_byte_offset()) ||
+ op->op.same_as(tirx::builtin::handle_add_byte_offset())) {
+ bool is_typed_offset = op->op.same_as(tirx::builtin::ptr_byte_offset());
TVM_FFI_ICHECK_EQ(op->args.size(), is_typed_offset ? 3U : 2U);
const auto* pointer_type = op->ty.as<PointerTypeNode>();
TVM_FFI_ICHECK(pointer_type)
@@ -508,7 +509,7 @@ void CodeGenMetal::Dispatch_(const CallNode* op, std::ostream& os) { // NOLINT(
os << ") + ";
PrintExpr(op->args[1], os);
os << "))";
- } else if (op->op.same_as(builtin::reinterpret())) {
+ } else if (op->op.same_as(tirx::builtin::reinterpret())) {
if (!op->ty.as<PrimTypeNode>() || !op->args[0]->ty.as<PrimTypeNode>()) {
return CodeGenC::Dispatch_(op, os);
}
diff --git a/src/backend/metal/codegen/intrin_rule_metal.cc b/src/backend/metal/codegen/intrin_rule_metal.cc
index 21ec9ffadc..947ab191b6 100644
--- a/src/backend/metal/codegen/intrin_rule_metal.cc
+++ b/src/backend/metal/codegen/intrin_rule_metal.cc
@@ -64,7 +64,7 @@ TVM_REGISTER_OP("tirx.clz")
TVM_REGISTER_OP("tirx.floor")
.set_attr<FLowerIntrinsic>("metal.FLowerIntrinsic", DispatchPureExtern<Direct>);
-TVM_REGISTER_OP("tirx.ceil")
+TVM_REGISTER_OP("prim.ceil")
.set_attr<FLowerIntrinsic>("metal.FLowerIntrinsic", DispatchPureExtern<Direct>);
TVM_REGISTER_OP("tirx.trunc")
@@ -100,7 +100,7 @@ TVM_REGISTER_OP("tirx.exp10")
TVM_REGISTER_OP("tirx.log")
.set_attr<FLowerIntrinsic>("metal.FLowerIntrinsic", DispatchPureExtern<Direct>);
-TVM_REGISTER_OP("tirx.log2")
+TVM_REGISTER_OP("prim.log2")
.set_attr<FLowerIntrinsic>("metal.FLowerIntrinsic", DispatchPureExtern<Direct>);
TVM_REGISTER_OP("tirx.log10")
diff --git a/src/backend/opencl/codegen/codegen_opencl.cc b/src/backend/opencl/codegen/codegen_opencl.cc
index 8c28b3192b..692b7cd0a0 100644
--- a/src/backend/opencl/codegen/codegen_opencl.cc
+++ b/src/backend/opencl/codegen/codegen_opencl.cc
@@ -43,7 +43,8 @@ const VarNode* TryUnwrapTextureVar(const Expr& texture) {
if (const auto* var = texture.as<VarNode>()) {
return var;
}
- if (const auto* call = texture.as<CallNode>(); call && call->op.same_as(builtin::buffer_data())) {
+ if (const auto* call = texture.as<CallNode>();
+ call && call->op.same_as(tirx::builtin::buffer_data())) {
TVM_FFI_ICHECK_EQ(call->args.size(), 1U);
const auto* buffer = call->args[0].as<VarNode>();
TVM_FFI_ICHECK(buffer && buffer->ty.as<BufferTypeNode>())
@@ -99,11 +100,11 @@ class InferTextureAccess : public StmtExprVisitor {
return StmtExprVisitor::Visit_(op);
}
ffi::Optional<VisitInterrupt> Visit_(const CallNode* op) final {
- if (op->op.same_as(builtin::texture2d_load())) {
+ if (op->op.same_as(tirx::builtin::texture2d_load())) {
const VarNode* texture = UnwrapTextureArgument(op->args[0]).var;
auto it = buffer_data_map_.find(texture);
var_access_map_[it == buffer_data_map_.end() ? texture : it->second] |= kReadAccess;
- } else if (op->op.same_as(builtin::texture2d_store())) {
+ } else if (op->op.same_as(tirx::builtin::texture2d_store())) {
const VarNode* texture = UnwrapTextureArgument(op->args[0]).var;
auto it = buffer_data_map_.find(texture);
var_access_map_[it == buffer_data_map_.end() ? texture : it->second] |= kWriteAccess;
@@ -457,7 +458,7 @@ void CodeGenOpenCL::VisitStmt_(const AllocBufferNode* op) {
}
void CodeGenOpenCL::Dispatch_(const CallNode* op, std::ostream& os) {
- if (op->op.same_as(builtin::address_of())) {
+ if (op->op.same_as(tirx::builtin::address_of())) {
// Overload tvm_address_of to add storage scope (e.g. __global).
const TensorLoadNode* load = op->args[0].as<TensorLoadNode>();
TVM_FFI_ICHECK(op->args.size() == 1 && load);
@@ -472,7 +473,7 @@ void CodeGenOpenCL::Dispatch_(const CallNode* op, std::ostream& os) {
os << " *)" << this->GetVarID(load->source.as_or_throw<tvm::tirx::BufferVar>().get()) << " + ";
this->PrintExpr(load->indices[0], os);
os << ')';
- } else if (op->op.same_as(builtin::texture2d_store())) {
+ } else if (op->op.same_as(tirx::builtin::texture2d_store())) {
TextureArgument texture = UnwrapTextureArgument(op->args[0]);
const int channel_size = op->args[4].as_or_throw<IntImm>()->value;
TVM_FFI_ICHECK(channel_size == 64 || channel_size == 128)
@@ -506,7 +507,7 @@ void CodeGenOpenCL::Dispatch_(const CallNode* op, std::ostream& os) {
this->PrintType(channel_type, os);
os << "(" << value << ")";
os << ")";
- } else if (op->op.same_as(builtin::texture2d_load())) {
+ } else if (op->op.same_as(tirx::builtin::texture2d_load())) {
TextureArgument texture = UnwrapTextureArgument(op->args[0]);
enable_compliant_texture_reads_ = true;
std::stringstream ss;
@@ -540,13 +541,13 @@ void CodeGenOpenCL::Dispatch_(const CallNode* op, std::ostream& os) {
std::string rhs = SSAGetID(ss.str(), op_ty.WithLanes(data_lanes));
if (auto ramp = op->args.back().as<prim::RampNode>()) {
- if (ramp->base.as<IntImmNode>() && *tirx::as_const_int(ramp->base) == 0 &&
- *tirx::as_const_int(ramp->lanes) == data_lanes &&
- *tirx::as_const_int(ramp->stride) == 1) {
+ if (ramp->base.as<IntImmNode>() && *tvm::prim::as_const_int(ramp->base) == 0 &&
+ *tvm::prim::as_const_int(ramp->lanes) == data_lanes &&
+ *tvm::prim::as_const_int(ramp->stride) == 1) {
os << rhs;
- } else if (*tirx::as_const_int(ramp->stride) == 1) {
+ } else if (*tvm::prim::as_const_int(ramp->stride) == 1) {
os << "(*(";
- this->PrintType(op_ty.WithLanes(*tirx::as_const_int(ramp->lanes)), os);
+ this->PrintType(op_ty.WithLanes(*tvm::prim::as_const_int(ramp->lanes)), os);
os << "*)";
os << "((";
this->PrintType(op_ty.WithLanes(1), os);
diff --git a/src/backend/opencl/codegen/intrin_rule_opencl.cc b/src/backend/opencl/codegen/intrin_rule_opencl.cc
index 4c88dcec6f..9b3ee23f2e 100644
--- a/src/backend/opencl/codegen/intrin_rule_opencl.cc
+++ b/src/backend/opencl/codegen/intrin_rule_opencl.cc
@@ -59,7 +59,7 @@ TVM_REGISTER_OP("tirx.clz")
TVM_REGISTER_OP("tirx.floor")
.set_attr<FLowerIntrinsic>("opencl.FLowerIntrinsic", DispatchPureExtern<Direct>);
-TVM_REGISTER_OP("tirx.ceil")
+TVM_REGISTER_OP("prim.ceil")
.set_attr<FLowerIntrinsic>("opencl.FLowerIntrinsic", DispatchPureExtern<Direct>);
TVM_REGISTER_OP("tirx.trunc")
@@ -98,7 +98,7 @@ TVM_REGISTER_OP("tirx.exp10")
TVM_REGISTER_OP("tirx.log")
.set_attr<FLowerIntrinsic>("opencl.FLowerIntrinsic", DispatchPureExtern<Direct>);
-TVM_REGISTER_OP("tirx.log2")
+TVM_REGISTER_OP("prim.log2")
.set_attr<FLowerIntrinsic>("opencl.FLowerIntrinsic", DispatchPureExtern<Direct>);
TVM_REGISTER_OP("tirx.log10")
diff --git a/src/backend/rocm/codegen/llvm/intrin_rule_rocm.cc b/src/backend/rocm/codegen/llvm/intrin_rule_rocm.cc
index 1f13d4e4b0..6698cd9240 100644
--- a/src/backend/rocm/codegen/llvm/intrin_rule_rocm.cc
+++ b/src/backend/rocm/codegen/llvm/intrin_rule_rocm.cc
@@ -37,6 +37,7 @@
namespace tvm {
namespace codegen {
+using namespace tvm::prim;
inline PrimExpr DispatchPureExternOCML(const PrimExpr& e) {
// NOTE: OCML dispatch fails to work properly with vectorization, and thus should be used with
@@ -48,7 +49,8 @@ inline PrimExpr DispatchPureExternOCML(const PrimExpr& e) {
const OpNode* op = call->op.as<OpNode>();
TVM_FFI_ICHECK(op != nullptr);
std::string name = op->name;
- TVM_FFI_ICHECK_EQ(name.substr(0, 5), "tirx.");
+ TVM_FFI_ICHECK(name.substr(0, 5) == "tirx." || name == "prim.ceil" || name == "prim.log2")
+ << "Unexpected intrinsic name: " << name;
std::ostringstream intrinsic_name;
PrimType call_ty = call->ty.as_or_throw<PrimType>();
@@ -59,7 +61,7 @@ inline PrimExpr DispatchPureExternOCML(const PrimExpr& e) {
new_args.push_back(arg);
}
- return Call(call_ty, builtin::call_pure_extern(), new_args).as_or_throw<PrimExpr>();
+ return Call(call_ty, tirx::builtin::call_pure_extern(), new_args).as_or_throw<PrimExpr>();
}
inline PrimExpr DispatchShuffle(const PrimExpr& e) {
@@ -76,25 +78,25 @@ inline PrimExpr DispatchShuffle(const PrimExpr& e) {
PrimExpr minus_one = IntImm::Int32(-1);
PrimExpr zero = IntImm::Int32(0);
PrimType i32_ty = PrimType::Int(32);
- PrimExpr lo = Call(i32_ty, builtin::call_pure_extern(),
+ PrimExpr lo = Call(i32_ty, tirx::builtin::call_pure_extern(),
ffi::Array<PrimExpr>{prim::StringImm("llvm.amdgcn.mbcnt.lo"), minus_one, zero})
.as_or_throw<PrimExpr>();
- PrimExpr self = Call(i32_ty, builtin::call_pure_extern(),
+ PrimExpr self = Call(i32_ty, tirx::builtin::call_pure_extern(),
ffi::Array<PrimExpr>{prim::StringImm("llvm.amdgcn.mbcnt.hi"), minus_one, lo})
.as_or_throw<PrimExpr>();
// compute lane to get from
PrimExpr width = args[3];
PrimExpr index;
- if (call->op.same_as(builtin::tvm_warp_shuffle())) {
+ if (call->op.same_as(tirx::builtin::tvm_warp_shuffle())) {
PrimExpr src_lane = args[2];
index = src_lane + (self & ~(width - 1));
- } else if (call->op.same_as(builtin::tvm_warp_shuffle_up())) {
+ } else if (call->op.same_as(tirx::builtin::tvm_warp_shuffle_up())) {
PrimExpr delta = args[2];
index = self - delta;
index = prim::Select(index < (self & ~(width - 1)), self, index);
} else {
- TVM_FFI_ICHECK(call->op.same_as(builtin::tvm_warp_shuffle_down()));
+ TVM_FFI_ICHECK(call->op.same_as(tirx::builtin::tvm_warp_shuffle_down()));
PrimExpr delta = args[2];
index = self + delta;
index = prim::Select((self & (width - 1)) + delta >= width, self, index);
@@ -103,7 +105,7 @@ inline PrimExpr DispatchShuffle(const PrimExpr& e) {
bool is_int32 = var_ty.MatchesElementType(DLDataTypeCode::kDLInt, 32);
PrimExpr source = is_int32 ? var : reinterpret(PrimType::Int(32), var);
PrimExpr res =
- Call(i32_ty, builtin::call_pure_extern(),
+ Call(i32_ty, tirx::builtin::call_pure_extern(),
ffi::Array<PrimExpr>{prim::StringImm("llvm.amdgcn.ds.bpermute"), index << 2, source})
.as_or_throw<PrimExpr>();
if (!is_int32) {
@@ -137,7 +139,7 @@ TVM_REGISTER_OP("tirx.floor")
.set_attr<FLowerIntrinsic>("rocm.FLowerIntrinsic",
DispatchLLVMPureIntrin<::llvm::Intrinsic::floor, 1>);
-TVM_REGISTER_OP("tirx.ceil")
+TVM_REGISTER_OP("prim.ceil")
.set_attr<FLowerIntrinsic>("rocm.FLowerIntrinsic",
DispatchLLVMPureIntrin<::llvm::Intrinsic::ceil, 1>);
@@ -173,7 +175,7 @@ TVM_REGISTER_OP("tirx.log")
.set_attr<FLowerIntrinsic>("rocm.FLowerIntrinsic",
DispatchLLVMPureIntrin<::llvm::Intrinsic::log, 1>);
-TVM_REGISTER_OP("tirx.log2")
+TVM_REGISTER_OP("prim.log2")
.set_attr<FLowerIntrinsic>("rocm.FLowerIntrinsic",
DispatchLLVMPureIntrin<::llvm::Intrinsic::log2, 1>);
diff --git a/src/backend/trn/codegen/codegen_trn.cc b/src/backend/trn/codegen/codegen_trn.cc
index e198c29bef..bf36859bf6 100644
--- a/src/backend/trn/codegen/codegen_trn.cc
+++ b/src/backend/trn/codegen/codegen_trn.cc
@@ -39,6 +39,8 @@
namespace tvm {
namespace codegen {
+using namespace tvm::prim;
+
namespace {
std::string PrintShapeAsList(const ffi::Array<PrimExpr>& shape) {
std::ostringstream os;
@@ -618,7 +620,7 @@ void CodeGenTrainium::VisitStmt_(const DeclBufferNode* op) {
}
const VarNode* data = op->data.as<VarNode>();
if (const auto* call = op->data.as<CallNode>();
- call && call->op.same_as(builtin::buffer_data()) && call->args.size() == 1) {
+ call && call->op.same_as(tirx::builtin::buffer_data()) && call->args.size() == 1) {
data = call->args[0].as<VarNode>();
}
TVM_FFI_ICHECK(data) << "Trainium codegen expects DeclBuffer data to be a buffer variable";
diff --git a/src/backend/trn/transform/lower_trainium_layout.cc b/src/backend/trn/transform/lower_trainium_layout.cc
index 252f3bcc97..30ed8a5bf4 100644
--- a/src/backend/trn/transform/lower_trainium_layout.cc
+++ b/src/backend/trn/transform/lower_trainium_layout.cc
@@ -211,7 +211,7 @@ class TrainiumLayoutApplier : public tirx::IRMutatorWithAnalyzer {
TVM_FFI_ICHECK_EQ(store->buffer->dtype->dtype, (DLDataType{kDLInt, 8, 1}))
<< "Expected int8 backing array for boolean tensor";
auto writer = store.CopyOnWrite();
- writer->value = tvm::cast(PrimType::Int(8), store->value);
+ writer->value = tvm::prim::cast(PrimType::Int(8), store->value);
return std::move(store);
}
return std::move(store);
@@ -230,7 +230,7 @@ class TrainiumLayoutApplier : public tirx::IRMutatorWithAnalyzer {
(DLDataType{kDLInt, 8, 1}))
<< "Expected int8 backing array for boolean tensor";
load.CopyOnWrite()->ExprNode::ty = PrimType::Int(8);
- return tvm::cast(PrimType::Bool(), load);
+ return tvm::prim::cast(PrimType::Bool(), load);
} else {
return std::move(load);
}
diff --git a/src/backend/vulkan/codegen/codegen_spirv.cc b/src/backend/vulkan/codegen/codegen_spirv.cc
index 5c222ad117..45bfe016b2 100644
--- a/src/backend/vulkan/codegen/codegen_spirv.cc
+++ b/src/backend/vulkan/codegen/codegen_spirv.cc
@@ -37,6 +37,7 @@
namespace tvm {
namespace codegen {
+using namespace tvm::prim;
namespace {
@@ -63,7 +64,7 @@ const VarNode* AsBufferVarNode(const Expr& expr) {
return var;
}
if (const auto* call = expr.as<CallNode>();
- call && call->op.same_as(builtin::buffer_data()) && call->args.size() == 1) {
+ call && call->op.same_as(tirx::builtin::buffer_data()) && call->args.size() == 1) {
return call->args[0].as<VarNode>();
}
return nullptr;
@@ -332,14 +333,14 @@ spirv::Value CodeGenSPIRV::Dispatch_(const prim::LetNode* op) {
}
spirv::Value CodeGenSPIRV::Dispatch_(const CallNode* op) {
- TVM_FFI_ICHECK(!op->op.same_as(builtin::masked_load()))
+ TVM_FFI_ICHECK(!op->op.same_as(tirx::builtin::masked_load()))
<< "Predicated buffer load is not supported.";
- TVM_FFI_ICHECK(!op->op.same_as(builtin::masked_store()))
+ TVM_FFI_ICHECK(!op->op.same_as(tirx::builtin::masked_store()))
<< "Predicated buffer store is not supported.";
- if (op->op.same_as(builtin::buffer_data())) {
+ if (op->op.same_as(tirx::builtin::buffer_data())) {
TVM_FFI_ICHECK_EQ(op->args.size(), 1U);
return MakeValue(op->args[0]);
- } else if (op->op.same_as(builtin::call_spirv_pure_glsl450())) {
+ } else if (op->op.same_as(tirx::builtin::call_spirv_pure_glsl450())) {
TVM_FFI_ICHECK_GE(op->args.size(), 2U);
uint32_t inst_id = static_cast<uint32_t>(op->args[0].as<IntImmNode>()->value);
std::vector<spirv::Value> values;
@@ -381,16 +382,16 @@ spirv::Value CodeGenSPIRV::Dispatch_(const CallNode* op) {
} else {
return builder_->MakeValue(spv::OpShiftRightLogical, a.stype, a, b);
}
- } else if (op->op.same_as(builtin::reinterpret())) {
+ } else if (op->op.same_as(tirx::builtin::reinterpret())) {
return builder_->MakeValue(spv::OpBitcast, builder_->GetSType(op->ty.as_or_throw<PrimType>()),
MakeValue(op->args[0]));
- } else if (op->op.same_as(builtin::large_uint_imm())) {
+ } else if (op->op.same_as(tirx::builtin::large_uint_imm())) {
TVM_FFI_ICHECK_EQ(op->args.size(), 2U);
uint64_t low = static_cast<uint64_t>(AsIntImmNode(op->args[0])->value);
uint64_t high = static_cast<uint64_t>(AsIntImmNode(op->args[1])->value);
uint64_t val = (high << 32U) | low;
return builder_->UIntImm(builder_->GetSType(op->ty.as_or_throw<PrimType>()), val);
- } else if (op->op.same_as(builtin::tvm_storage_sync())) {
+ } else if (op->op.same_as(tirx::builtin::tvm_storage_sync())) {
return this->CreateStorageSync(op);
} else if (op->op.same_as(prim::builtin::if_then_else())) {
TVM_FFI_ICHECK_EQ(op->args.size(), 3U);
@@ -416,10 +417,10 @@ spirv::Value CodeGenSPIRV::Dispatch_(const CallNode* op) {
phi.SetIncoming(0, then_value, then_value_label);
phi.SetIncoming(1, else_value, else_value_label);
return phi;
- } else if (op->op.same_as(builtin::popcount())) {
+ } else if (op->op.same_as(tirx::builtin::popcount())) {
return builder_->MakeValue(spv::OpBitCount, builder_->GetSType(op->ty.as_or_throw<PrimType>()),
MakeValue(op->args[0]));
- } else if (op->op.same_as(builtin::call_pure_extern())) {
+ } else if (op->op.same_as(tirx::builtin::call_pure_extern())) {
TVM_FFI_ICHECK_GE(op->args.size(), 1U);
const std::string& func_name = op->args[0].as<prim::StringImmNode>()->value;
if (func_name == "__dp4a") {
@@ -435,7 +436,7 @@ spirv::Value CodeGenSPIRV::Dispatch_(const CallNode* op) {
<< AsStringImmNode(op->args[0])->value << "\"";
return spirv::Value();
}
- } else if (op->op.same_as(builtin::call_extern())) {
+ } else if (op->op.same_as(tirx::builtin::call_extern())) {
TVM_FFI_ICHECK_GE(op->args.size(), 1U);
TVM_FFI_THROW(InternalError)
<< "SPIR-V shader cannot make extern calls. Graph contains extern \""
@@ -553,7 +554,7 @@ spirv::Value CodeGenSPIRV::Dispatch_(const CallNode* op) {
builder_->MakeInst(spv::OpCooperativeMatrixStoreNV, dst_ptr, loaded, stride_val,
(layout != "row_major") ? t_val : f_val);
return spirv::Value();
- } else if (op->op.same_as(builtin::address_of())) {
+ } else if (op->op.same_as(tirx::builtin::address_of())) {
const TensorLoadNode* load = op->args[0].as<TensorLoadNode>();
Var buffer_var = load->source.as_or_throw<tvm::tirx::BufferVar>().var();
const VarNode* buffer_node = buffer_var.get();
@@ -564,7 +565,7 @@ spirv::Value CodeGenSPIRV::Dispatch_(const CallNode* op) {
spirv::SType ptr_type = builder_->GetPointerType(ele_stype, buffer_val.stype.storage_class);
TVM_FFI_ICHECK(var_map_.count(buffer_node));
return builder_->StructArrayAccess(ptr_type, var_map_[buffer_node], MakeValue(index));
- } else if (op->op.same_as(builtin::tvm_thread_invariant())) {
+ } else if (op->op.same_as(tirx::builtin::tvm_thread_invariant())) {
return MakeValue(op->args[0]);
} else {
TVM_FFI_THROW(InternalError) << "Unresolved call " << op->op;
@@ -745,7 +746,7 @@ void CodeGenSPIRV::VisitStmt_(const ForNode* op) {
? builder_->IntImm(init_value.stype, 1)
: builder_->UIntImm(init_value.stype, 1);
} else {
- step = MakeValue(tvm::cast(end.ty(), *op->step));
+ step = MakeValue(tvm::prim::cast(end.ty(), *op->step));
}
// Must get init label after making value(to make sure they are correct)
diff --git a/src/backend/vulkan/codegen/intrin_rule_spirv.cc b/src/backend/vulkan/codegen/intrin_rule_spirv.cc
index 6d473fa8c6..13bf738626 100644
--- a/src/backend/vulkan/codegen/intrin_rule_spirv.cc
+++ b/src/backend/vulkan/codegen/intrin_rule_spirv.cc
@@ -74,7 +74,7 @@ void RegisterVulkanLowerIntrinRules() {
TVM_REGISTER_OP("tirx.floor")
.set_attr<FLowerIntrinsic>("vulkan.FLowerIntrinsic", DispatchGLSLPureIntrin<GLSLstd450Floor>);
-TVM_REGISTER_OP("tirx.ceil")
+TVM_REGISTER_OP("prim.ceil")
.set_attr<FLowerIntrinsic>("vulkan.FLowerIntrinsic", DispatchGLSLPureIntrin<GLSLstd450Ceil>);
TVM_REGISTER_OP("tirx.round")
@@ -139,7 +139,7 @@ TVM_REGISTER_OP("tirx.atan2")
TVM_REGISTER_OP("tirx.log")
.set_attr<FLowerIntrinsic>("vulkan.FLowerIntrinsic", DispatchGLSLPureIntrin<GLSLstd450Log>);
-TVM_REGISTER_OP("tirx.log2")
+TVM_REGISTER_OP("prim.log2")
.set_attr<FLowerIntrinsic>("vulkan.FLowerIntrinsic", DispatchGLSLPureIntrin<GLSLstd450Log2>);
TVM_REGISTER_OP("tirx.sqrt")
diff --git a/src/backend/webgpu/codegen/codegen_webgpu.cc b/src/backend/webgpu/codegen/codegen_webgpu.cc
index 5bb9a7f743..9d345c27fd 100644
--- a/src/backend/webgpu/codegen/codegen_webgpu.cc
+++ b/src/backend/webgpu/codegen/codegen_webgpu.cc
@@ -50,6 +50,7 @@
namespace tvm {
namespace codegen {
+using namespace tvm::prim;
namespace {
@@ -456,11 +457,11 @@ PrimExpr CodeGenWebGPU::EnforceU32(PrimExpr value) {
}
void CodeGenWebGPU::Dispatch_(const CallNode* op, std::ostream& os) { // NOLINT(*)
- TVM_FFI_ICHECK(!op->op.same_as(builtin::masked_load()))
+ TVM_FFI_ICHECK(!op->op.same_as(tirx::builtin::masked_load()))
<< "Predicated buffer load is not supported.";
- TVM_FFI_ICHECK(!op->op.same_as(builtin::masked_store()))
+ TVM_FFI_ICHECK(!op->op.same_as(tirx::builtin::masked_store()))
<< "Predicated buffer store is not supported.";
- if (op->op.same_as(builtin::reinterpret())) {
+ if (op->op.same_as(tirx::builtin::reinterpret())) {
// generate bitcast<TYPE>(ARG)
os << "bitcast<";
this->PrintType(op->ty.as_or_throw<PrimType>(), os);
@@ -506,7 +507,7 @@ void CodeGenWebGPU::Dispatch_(const CallNode* op, std::ostream& os) { // NOLINT
this->EndScope(else_scope);
}
os << result;
- } else if (op->op.same_as(builtin::dp4a())) {
+ } else if (op->op.same_as(tirx::builtin::dp4a())) {
// generate `dot4I8Packed(vec1, vec2) + acc` for the builtin `dp4a`
os << "dot4I8Packed(";
this->PrintExpr(op->args[0], os);
diff --git a/src/backend/webgpu/codegen/intrin_rule_webgpu.cc b/src/backend/webgpu/codegen/intrin_rule_webgpu.cc
index 82883fee17..a275c4a984 100644
--- a/src/backend/webgpu/codegen/intrin_rule_webgpu.cc
+++ b/src/backend/webgpu/codegen/intrin_rule_webgpu.cc
@@ -94,7 +94,7 @@ TVM_REGISTER_OP("tirx.atan")
TVM_REGISTER_OP("tirx.atan2")
.set_attr<FLowerIntrinsic>("webgpu.FLowerIntrinsic", DispatchPureExtern<Direct>);
-TVM_REGISTER_OP("tirx.ceil")
+TVM_REGISTER_OP("prim.ceil")
.set_attr<FLowerIntrinsic>("webgpu.FLowerIntrinsic", DispatchPureExtern<Direct>);
TVM_REGISTER_OP("tirx.cos")
@@ -118,7 +118,7 @@ TVM_REGISTER_OP("tirx.fma")
TVM_REGISTER_OP("tirx.log")
.set_attr<FLowerIntrinsic>("webgpu.FLowerIntrinsic", DispatchPureExtern<Direct>);
-TVM_REGISTER_OP("tirx.log2")
+TVM_REGISTER_OP("prim.log2")
.set_attr<FLowerIntrinsic>("webgpu.FLowerIntrinsic", DispatchPureExtern<Direct>);
TVM_REGISTER_OP("tirx.pow")
diff --git a/src/ir/expr.cc b/src/ir/expr.cc
index 37bd552959..0f4502271c 100644
--- a/src/ir/expr.cc
+++ b/src/ir/expr.cc
@@ -30,6 +30,7 @@
#include <tvm/ir/function.h>
#include <tvm/ir/op.h>
#include <tvm/ir/prim/expr.h>
+#include <tvm/ir/prim/op.h>
#include <tvm/ir/type.h>
#include <tvm/te/tensor.h>
@@ -814,7 +815,8 @@ TVM_FFI_STATIC_INIT_BLOCK() {
// Range
Range::Range(PrimExpr begin, PrimExpr end, Span span)
- : Range(ffi::make_object<RangeNode>(begin, tirx::is_zero(begin) ? end : (end - begin), span)) {}
+ : Range(ffi::make_object<RangeNode>(begin, tvm::prim::is_zero(begin) ? end : (end - begin),
+ span)) {}
TVM_FFI_STATIC_INIT_BLOCK() {
namespace refl = tvm::ffi::reflection;
diff --git a/src/ir/prim/builtin.cc b/src/ir/prim/builtin.cc
index 805a064b4b..d8b15a69ad 100644
--- a/src/ir/prim/builtin.cc
+++ b/src/ir/prim/builtin.cc
@@ -16,59 +16,63 @@
* specific language governing permissions and limitations
* under the License.
*/
+#include <tvm/ir/op_attr_types.h>
#include <tvm/ir/prim/builtin.h>
-#include <tvm/tirx/op.h>
-#include <tvm/tirx/op_attr_types.h>
namespace tvm {
namespace prim {
namespace builtin {
-using namespace tvm::tirx;
-
-#define PRIM_DEFINE_BUILTIN_FUNC(OpName) \
- const Op& OpName() { \
- static const Op& op = Op::Get("ir.prim." #OpName); \
- return op; \
- } \
- TVM_REGISTER_OP("ir.prim." #OpName) \
- .set_attr<TScriptPrinterName>("TScriptPrinterName", #OpName) \
- .set_attr<TIRxOpCategory>("TIRxOpCategory", ffi::String("builtin"), 1)
+#define PRIM_DEFINE_BUILTIN_FUNC(OpName) \
+ const Op& OpName() { \
+ static const Op& op = Op::Get("prim." #OpName); \
+ return op; \
+ } \
+ TVM_REGISTER_OP("prim." #OpName)
PRIM_DEFINE_BUILTIN_FUNC(likely)
.set_num_inputs(1)
.set_attr<TCallEffectKind>("TCallEffectKind",
static_cast<int64_t>(CallEffectKind::kExprAnnotation))
- .set_attr<TVectorizable>("TVectorizable", true);
+ .set_attr<bool>("TVectorizable", true);
PRIM_DEFINE_BUILTIN_FUNC(bitwise_and)
.set_num_inputs(2)
.set_attr<TCallEffectKind>("TCallEffectKind", static_cast<int64_t>(CallEffectKind::kPure))
- .set_attr<TVectorizable>("TVectorizable", true);
+ .set_attr<bool>("TVectorizable", true);
PRIM_DEFINE_BUILTIN_FUNC(bitwise_or)
.set_num_inputs(2)
.set_attr<TCallEffectKind>("TCallEffectKind", static_cast<int64_t>(CallEffectKind::kPure))
- .set_attr<TVectorizable>("TVectorizable", true);
+ .set_attr<bool>("TVectorizable", true);
PRIM_DEFINE_BUILTIN_FUNC(bitwise_xor)
.set_num_inputs(2)
.set_attr<TCallEffectKind>("TCallEffectKind", static_cast<int64_t>(CallEffectKind::kPure))
- .set_attr<TVectorizable>("TVectorizable", true);
+ .set_attr<bool>("TVectorizable", true);
PRIM_DEFINE_BUILTIN_FUNC(bitwise_not)
.set_num_inputs(1)
.set_attr<TCallEffectKind>("TCallEffectKind", static_cast<int64_t>(CallEffectKind::kPure))
- .set_attr<TVectorizable>("TVectorizable", true);
+ .set_attr<bool>("TVectorizable", true);
PRIM_DEFINE_BUILTIN_FUNC(shift_left)
.set_num_inputs(2)
.set_attr<TCallEffectKind>("TCallEffectKind", static_cast<int64_t>(CallEffectKind::kPure))
- .set_attr<TVectorizable>("TVectorizable", true);
+ .set_attr<bool>("TVectorizable", true);
PRIM_DEFINE_BUILTIN_FUNC(shift_right)
.set_num_inputs(2)
.set_attr<TCallEffectKind>("TCallEffectKind", static_cast<int64_t>(CallEffectKind::kPure))
- .set_attr<TVectorizable>("TVectorizable", true);
+ .set_attr<bool>("TVectorizable", true);
PRIM_DEFINE_BUILTIN_FUNC(if_then_else)
.set_num_inputs(3)
.set_attr<TCallEffectKind>("TCallEffectKind", static_cast<int64_t>(CallEffectKind::kPure));
PRIM_DEFINE_BUILTIN_FUNC(vscale).set_attr<TCallEffectKind>(
"TCallEffectKind", static_cast<int64_t>(CallEffectKind::kPure));
+PRIM_DEFINE_BUILTIN_FUNC(ceil)
+ .set_num_inputs(1)
+ .set_attr<TCallEffectKind>("TCallEffectKind", static_cast<int64_t>(CallEffectKind::kPure))
+ .set_attr<bool>("TVectorizable", true);
+PRIM_DEFINE_BUILTIN_FUNC(log2)
+ .set_num_inputs(1)
+ .set_attr<TCallEffectKind>("TCallEffectKind", static_cast<int64_t>(CallEffectKind::kPure))
+ .set_attr<bool>("TVectorizable", true);
+
#undef PRIM_DEFINE_BUILTIN_FUNC
} // namespace builtin
} // namespace prim
diff --git a/src/arith/const_fold.h b/src/ir/prim/const_fold.h
similarity index 88%
copy from src/arith/const_fold.h
copy to src/ir/prim/const_fold.h
index 0c2b72a45c..158470a081 100644
--- a/src/arith/const_fold.h
+++ b/src/ir/prim/const_fold.h
@@ -21,13 +21,12 @@
* \file const_fold.h
* \brief Centralized location for constant folding.
*/
-#ifndef TVM_ARITH_CONST_FOLD_H_
-#define TVM_ARITH_CONST_FOLD_H_
+#ifndef TVM_IR_PRIM_CONST_FOLD_H_
+#define TVM_IR_PRIM_CONST_FOLD_H_
#include <tvm/ffi/optional.h>
#include <tvm/ir/prim/expr.h>
#include <tvm/runtime/logging.h>
-#include <tvm/tirx/op.h>
#include <algorithm>
#include <cmath>
@@ -36,34 +35,8 @@
#include "int_operator.h"
namespace tvm {
-namespace arith {
-
-/*!
- * \brief Scoped opt-in flag for unsigned index analysis (allow_u32).
- *
- * Unsigned arithmetic wraps modulo 2^bits, so most index rewrite rules are
- * unsound for it in general. When the caller can guarantee values never
- * approach the type limit (e.g. compile-time layout proofs over SMEM
- * offsets), enabling this flag activates a small set of otherwise-unsafe
- * rewrite rules for unsigned operands (see the unsigned blocks in
- * rewrite_simplify.cc). It is OFF by default; thread-local and nested-safe.
- */
-namespace uint_as_index {
-
-inline thread_local int g_depth = 0;
-
-inline bool Enabled() { return g_depth > 0; }
-
-} // namespace uint_as_index
-
-/*! \brief RAII guard enabling uint_as_index mode in the current scope. */
-class AllowUintAsIndexGuard {
- public:
- AllowUintAsIndexGuard() { ++uint_as_index::g_depth; }
- ~AllowUintAsIndexGuard() { --uint_as_index::g_depth; }
- AllowUintAsIndexGuard(const AllowUintAsIndexGuard&) = delete;
- AllowUintAsIndexGuard& operator=(const AllowUintAsIndexGuard&) = delete;
-};
+namespace prim {
+namespace detail {
/*!
* \brief Try to run binary compute with constant folding.
@@ -146,7 +119,7 @@ inline double GetFoldResultDoubleRepr(float x) {
return res;
}
-#define TVM_ARITH_CONST_PROPAGATION(BODY) \
+#define TVM_PRIM_CONST_PROPAGATION(BODY) \
using tvm::FloatImmNode; \
const IntImmNode* pa = a.as<IntImmNode>(); \
const IntImmNode* pb = b.as<IntImmNode>(); \
@@ -154,17 +127,17 @@ inline double GetFoldResultDoubleRepr(float x) {
const FloatImmNode* fb = b.as<FloatImmNode>(); \
BODY;
-#define TVM_INDEX_CONST_PROPAGATION(BODY) \
- const IntImmNode* pa = a.as<IntImmNode>(); \
- const IntImmNode* pb = b.as<IntImmNode>(); \
- if (arith::IsIndexTypedExpr(a) && arith::IsIndexTypedExpr(b)) { \
- BODY; \
+#define TVM_PRIM_INDEX_CONST_PROPAGATION(BODY) \
+ const IntImmNode* pa = a.as<IntImmNode>(); \
+ const IntImmNode* pb = b.as<IntImmNode>(); \
+ if (::tvm::prim::detail::IsIndexTypedExpr(a) && ::tvm::prim::detail::IsIndexTypedExpr(b)) { \
+ BODY; \
}
// specialization of constant folders.
template <>
inline ffi::Optional<PrimExpr> TryConstFold<prim::Add>(PrimExpr a, PrimExpr b) {
- TVM_ARITH_CONST_PROPAGATION({
+ TVM_PRIM_CONST_PROPAGATION({
PrimType result_ty = a.ty();
if (pa && pb) {
int64_t res = pa->value + pb->value;
@@ -188,7 +161,7 @@ inline ffi::Optional<PrimExpr> TryConstFold<prim::Add>(PrimExpr a, PrimExpr b) {
template <>
inline ffi::Optional<PrimExpr> TryConstFold<prim::Sub>(PrimExpr a, PrimExpr b) {
- TVM_ARITH_CONST_PROPAGATION({
+ TVM_PRIM_CONST_PROPAGATION({
TVM_FFI_ICHECK(!((pa && pa->ty.as_or_throw<PrimType>().MatchesCode(DLDataTypeCode::kDLUInt) &&
pa->value == 0U) &&
(pb && pb->ty.as_or_throw<PrimType>().MatchesCode(DLDataTypeCode::kDLUInt) &&
@@ -216,7 +189,7 @@ inline ffi::Optional<PrimExpr> TryConstFold<prim::Sub>(PrimExpr a, PrimExpr b) {
template <>
inline ffi::Optional<PrimExpr> TryConstFold<prim::Mul>(PrimExpr a, PrimExpr b) {
- TVM_ARITH_CONST_PROPAGATION({
+ TVM_PRIM_CONST_PROPAGATION({
PrimType result_ty = a.ty();
if (pa && pb) {
int64_t res = pa->value * pb->value;
@@ -252,7 +225,7 @@ inline ffi::Optional<PrimExpr> TryConstFold<prim::Mul>(PrimExpr a, PrimExpr b) {
template <>
inline ffi::Optional<PrimExpr> TryConstFold<prim::Div>(PrimExpr a, PrimExpr b) {
- TVM_ARITH_CONST_PROPAGATION({
+ TVM_PRIM_CONST_PROPAGATION({
PrimType result_ty = a.ty();
if (pa && pb) {
// due to division and mod can have different modes
@@ -288,7 +261,7 @@ inline ffi::Optional<PrimExpr> TryConstFold<prim::Div>(PrimExpr a, PrimExpr b) {
template <>
inline ffi::Optional<PrimExpr> TryConstFold<prim::Mod>(PrimExpr a, PrimExpr b) {
- TVM_INDEX_CONST_PROPAGATION({
+ TVM_PRIM_INDEX_CONST_PROPAGATION({
PrimType result_ty = a.ty();
if (pa && pb) {
TVM_FFI_ICHECK_NE(pb->value, 0) << "Divide by zero";
@@ -308,11 +281,11 @@ inline ffi::Optional<PrimExpr> TryConstFold<prim::Mod>(PrimExpr a, PrimExpr b) {
template <>
inline ffi::Optional<PrimExpr> TryConstFold<prim::FloorDiv>(PrimExpr a, PrimExpr b) {
- TVM_ARITH_CONST_PROPAGATION({
+ TVM_PRIM_CONST_PROPAGATION({
PrimType result_ty = a.ty();
if (pa && pb) {
TVM_FFI_ICHECK_NE(pb->value, 0) << "Divide by zero";
- int64_t res = arith::floordiv(pa->value, pb->value);
+ int64_t res = detail::floordiv(pa->value, pb->value);
return IntImm(result_ty, GetFoldResultInt64Repr(res, result_ty));
}
if (pa) {
@@ -360,11 +333,11 @@ inline ffi::Optional<PrimExpr> TryConstFold<prim::FloorMod>(PrimExpr a, PrimExpr
TVM_FFI_ICHECK_NE(rhs, 0U) << "Divide by zero";
return IntImm(utype, static_cast<int64_t>(lhs % rhs));
}
- TVM_INDEX_CONST_PROPAGATION({
+ TVM_PRIM_INDEX_CONST_PROPAGATION({
PrimType result_ty = a.ty();
if (pa && pb) {
TVM_FFI_ICHECK_NE(pb->value, 0) << "Divide by zero";
- int64_t res = arith::floormod(pa->value, pb->value);
+ int64_t res = detail::floormod(pa->value, pb->value);
return IntImm(result_ty, GetFoldResultInt64Repr(res, result_ty));
}
if (pa) {
@@ -380,7 +353,7 @@ inline ffi::Optional<PrimExpr> TryConstFold<prim::FloorMod>(PrimExpr a, PrimExpr
template <>
inline ffi::Optional<PrimExpr> TryConstFold<prim::Min>(PrimExpr a, PrimExpr b) {
- TVM_ARITH_CONST_PROPAGATION({
+ TVM_PRIM_CONST_PROPAGATION({
PrimType result_ty = a.ty();
if (pa && pb) return IntImm(result_ty, std::min(pa->value, pb->value));
if (fa && fb) return FloatImm(result_ty, std::min(fa->value, fb->value));
@@ -391,7 +364,7 @@ inline ffi::Optional<PrimExpr> TryConstFold<prim::Min>(PrimExpr a, PrimExpr b) {
template <>
inline ffi::Optional<PrimExpr> TryConstFold<prim::Max>(PrimExpr a, PrimExpr b) {
- TVM_ARITH_CONST_PROPAGATION({
+ TVM_PRIM_CONST_PROPAGATION({
PrimType result_ty = a.ty();
if (pa && pb) return IntImm(result_ty, std::max(pa->value, pb->value));
if (fa && fb) return FloatImm(result_ty, std::max(fa->value, fb->value));
@@ -402,7 +375,7 @@ inline ffi::Optional<PrimExpr> TryConstFold<prim::Max>(PrimExpr a, PrimExpr b) {
template <>
inline ffi::Optional<PrimExpr> TryConstFold<prim::GT>(PrimExpr a, PrimExpr b) {
- TVM_ARITH_CONST_PROPAGATION({
+ TVM_PRIM_CONST_PROPAGATION({
if (pa && pb) return IntImm::Bool(pa->value > pb->value);
if (fa && fb) return IntImm::Bool(fa->value > fb->value);
});
@@ -411,7 +384,7 @@ inline ffi::Optional<PrimExpr> TryConstFold<prim::GT>(PrimExpr a, PrimExpr b) {
template <>
inline ffi::Optional<PrimExpr> TryConstFold<prim::GE>(PrimExpr a, PrimExpr b) {
- TVM_ARITH_CONST_PROPAGATION({
+ TVM_PRIM_CONST_PROPAGATION({
if (pa && pb) return IntImm::Bool(pa->value >= pb->value);
if (fa && fb) return IntImm::Bool(fa->value >= fb->value);
});
@@ -420,7 +393,7 @@ inline ffi::Optional<PrimExpr> TryConstFold<prim::GE>(PrimExpr a, PrimExpr b) {
template <>
inline ffi::Optional<PrimExpr> TryConstFold<prim::LT>(PrimExpr a, PrimExpr b) {
- TVM_ARITH_CONST_PROPAGATION({
+ TVM_PRIM_CONST_PROPAGATION({
if (pa && pb) return IntImm::Bool(pa->value < pb->value);
if (fa && fb) return IntImm::Bool(fa->value < fb->value);
});
@@ -429,7 +402,7 @@ inline ffi::Optional<PrimExpr> TryConstFold<prim::LT>(PrimExpr a, PrimExpr b) {
template <>
inline ffi::Optional<PrimExpr> TryConstFold<prim::LE>(PrimExpr a, PrimExpr b) {
- TVM_ARITH_CONST_PROPAGATION({
+ TVM_PRIM_CONST_PROPAGATION({
if (pa && pb) return IntImm::Bool(pa->value <= pb->value);
if (fa && fb) return IntImm::Bool(fa->value <= fb->value);
});
@@ -438,7 +411,7 @@ inline ffi::Optional<PrimExpr> TryConstFold<prim::LE>(PrimExpr a, PrimExpr b) {
template <>
inline ffi::Optional<PrimExpr> TryConstFold<prim::EQ>(PrimExpr a, PrimExpr b) {
- TVM_ARITH_CONST_PROPAGATION({
+ TVM_PRIM_CONST_PROPAGATION({
if (pa && pb) return IntImm::Bool(pa->value == pb->value);
if (fa && fb) return IntImm::Bool(fa->value == fb->value);
});
@@ -447,7 +420,7 @@ inline ffi::Optional<PrimExpr> TryConstFold<prim::EQ>(PrimExpr a, PrimExpr b) {
template <>
inline ffi::Optional<PrimExpr> TryConstFold<prim::NE>(PrimExpr a, PrimExpr b) {
- TVM_ARITH_CONST_PROPAGATION({
+ TVM_PRIM_CONST_PROPAGATION({
if (pa && pb) return IntImm::Bool(pa->value != pb->value);
if (fa && fb) return IntImm::Bool(fa->value != fb->value);
});
@@ -485,7 +458,7 @@ inline ffi::Optional<PrimExpr> TryConstFold<prim::Not>(PrimExpr a) {
return std::nullopt;
}
-/*! \brief Helper namespace for symbolic value limits */
+/*! \brief Shared symbolic value limits used by primitive folding and integer analysis. */
struct SymbolicLimits {
/*! \brief positive infinity */
static PrimExpr pos_inf_;
@@ -529,6 +502,7 @@ inline PrimExpr neg_inf() { return SymbolicLimits::neg_inf_; }
*/
inline bool is_neg_inf(const PrimExpr& value) { return value.same_as(SymbolicLimits::neg_inf_); }
-} // namespace arith
+} // namespace detail
+} // namespace prim
} // namespace tvm
-#endif // TVM_ARITH_CONST_FOLD_H_
+#endif // TVM_IR_PRIM_CONST_FOLD_H_
diff --git a/src/ir/prim/deep_equal.cc b/src/ir/prim/deep_equal.cc
index cff330402d..0fe7de4bf7 100644
--- a/src/ir/prim/deep_equal.cc
+++ b/src/ir/prim/deep_equal.cc
@@ -218,7 +218,7 @@ bool ExprDeepEqual::operator()(const PrimExpr& lhs, const PrimExpr& rhs) const {
TVM_FFI_STATIC_INIT_BLOCK() {
namespace refl = tvm::ffi::reflection;
- refl::GlobalDef().def("ir.prim.expr_deep_equal", [](const PrimExpr& lhs, const PrimExpr& rhs) {
+ refl::GlobalDef().def("prim.expr_deep_equal", [](const PrimExpr& lhs, const PrimExpr& rhs) {
return ExprDeepEqual()(lhs, rhs);
});
}
diff --git a/src/ir/prim/expr.cc b/src/ir/prim/expr.cc
index 48de984613..17ea8b735f 100644
--- a/src/ir/prim/expr.cc
+++ b/src/ir/prim/expr.cc
@@ -354,7 +354,7 @@ TVMFFIAny LetMaybeInplaceMutate(ffi::StructuralMutatorObj* mutator, ffi::AnyView
TVM_FFI_STATIC_INIT_BLOCK() {
namespace refl = tvm::ffi::reflection;
- refl::GlobalDef().def("tirx.convert",
+ refl::GlobalDef().def("prim.convert",
[](ffi::Variant<PrimExpr, ffi::Array<PrimExpr>> expr) { return expr; });
// Note: kRepr for VarNode is registered via TVM_REGISTER_SCRIPT_AS_REPR in
// src/script/printer/tirx/expr.cc (-> ReprPrintTIR which delegates to TVMScriptPrinter).
@@ -378,7 +378,7 @@ TVM_FFI_STATIC_INIT_BLOCK() {
.attr(refl::type_attr::kStructuralMaybeInplaceMutate,
reinterpret_cast<void*>(&StringImmMaybeInplaceMutate));
- refl::GlobalDef().def("ir.prim.StringImm",
+ refl::GlobalDef().def("prim.StringImm",
[](ffi::String value, Span span) { return StringImm(value, span); });
}
@@ -403,7 +403,7 @@ TVM_FFI_STATIC_INIT_BLOCK() {
.attr(refl::type_attr::kStructuralMaybeInplaceMutate,
reinterpret_cast<void*>(&CastMaybeInplaceMutate));
- refl::GlobalDef().def("ir.prim.Cast", [](PrimType dtype, PrimExpr value, Span span) {
+ refl::GlobalDef().def("prim.Cast", [](PrimType dtype, PrimExpr value, Span span) {
return Cast(dtype, value, span);
});
}
@@ -414,7 +414,7 @@ TVM_DEFINE_BINOP_CONSTRUCTOR(Add);
TVM_FFI_STATIC_INIT_BLOCK() {
namespace refl = tvm::ffi::reflection;
AddNode::RegisterReflection();
- refl::GlobalDef().def("ir.prim.Add",
+ refl::GlobalDef().def("prim.Add",
[](PrimExpr a, PrimExpr b, Span span) { return Add(a, b, span); });
refl::TypeAttrDef<AddNode>()
.attr(refl::type_attr::kStructuralVisit, reinterpret_cast<void*>(&BinaryVisit<AddNode>))
@@ -429,7 +429,7 @@ TVM_DEFINE_BINOP_CONSTRUCTOR(Sub);
TVM_FFI_STATIC_INIT_BLOCK() {
namespace refl = tvm::ffi::reflection;
SubNode::RegisterReflection();
- refl::GlobalDef().def("ir.prim.Sub",
+ refl::GlobalDef().def("prim.Sub",
[](PrimExpr a, PrimExpr b, Span span) { return Sub(a, b, span); });
refl::TypeAttrDef<SubNode>()
.attr(refl::type_attr::kStructuralVisit, reinterpret_cast<void*>(&BinaryVisit<SubNode>))
@@ -444,7 +444,7 @@ TVM_DEFINE_BINOP_CONSTRUCTOR(Mul);
TVM_FFI_STATIC_INIT_BLOCK() {
namespace refl = tvm::ffi::reflection;
MulNode::RegisterReflection();
- refl::GlobalDef().def("ir.prim.Mul",
+ refl::GlobalDef().def("prim.Mul",
[](PrimExpr a, PrimExpr b, Span span) { return Mul(a, b, span); });
refl::TypeAttrDef<MulNode>()
.attr(refl::type_attr::kStructuralVisit, reinterpret_cast<void*>(&BinaryVisit<MulNode>))
@@ -459,7 +459,7 @@ TVM_DEFINE_BINOP_CONSTRUCTOR(Div);
TVM_FFI_STATIC_INIT_BLOCK() {
namespace refl = tvm::ffi::reflection;
DivNode::RegisterReflection();
- refl::GlobalDef().def("ir.prim.Div",
+ refl::GlobalDef().def("prim.Div",
[](PrimExpr a, PrimExpr b, Span span) { return Div(a, b, span); });
refl::TypeAttrDef<DivNode>()
.attr(refl::type_attr::kStructuralVisit, reinterpret_cast<void*>(&BinaryVisit<DivNode>))
@@ -474,7 +474,7 @@ TVM_DEFINE_BINOP_CONSTRUCTOR(Mod);
TVM_FFI_STATIC_INIT_BLOCK() {
namespace refl = tvm::ffi::reflection;
ModNode::RegisterReflection();
- refl::GlobalDef().def("ir.prim.Mod",
+ refl::GlobalDef().def("prim.Mod",
[](PrimExpr a, PrimExpr b, Span span) { return Mod(a, b, span); });
refl::TypeAttrDef<ModNode>()
.attr(refl::type_attr::kStructuralVisit, reinterpret_cast<void*>(&BinaryVisit<ModNode>))
@@ -489,7 +489,7 @@ TVM_DEFINE_BINOP_CONSTRUCTOR(FloorDiv);
TVM_FFI_STATIC_INIT_BLOCK() {
namespace refl = tvm::ffi::reflection;
FloorDivNode::RegisterReflection();
- refl::GlobalDef().def("ir.prim.FloorDiv",
+ refl::GlobalDef().def("prim.FloorDiv",
[](PrimExpr a, PrimExpr b, Span span) { return FloorDiv(a, b, span); });
refl::TypeAttrDef<FloorDivNode>()
.attr(refl::type_attr::kStructuralVisit, reinterpret_cast<void*>(&BinaryVisit<FloorDivNode>))
@@ -505,7 +505,7 @@ TVM_DEFINE_BINOP_CONSTRUCTOR(FloorMod);
TVM_FFI_STATIC_INIT_BLOCK() {
namespace refl = tvm::ffi::reflection;
FloorModNode::RegisterReflection();
- refl::GlobalDef().def("ir.prim.FloorMod",
+ refl::GlobalDef().def("prim.FloorMod",
[](PrimExpr a, PrimExpr b, Span span) { return FloorMod(a, b, span); });
refl::TypeAttrDef<FloorModNode>()
.attr(refl::type_attr::kStructuralVisit, reinterpret_cast<void*>(&BinaryVisit<FloorModNode>))
@@ -521,7 +521,7 @@ TVM_DEFINE_BINOP_CONSTRUCTOR(Min);
TVM_FFI_STATIC_INIT_BLOCK() {
namespace refl = tvm::ffi::reflection;
MinNode::RegisterReflection();
- refl::GlobalDef().def("ir.prim.Min",
+ refl::GlobalDef().def("prim.Min",
[](PrimExpr a, PrimExpr b, Span span) { return Min(a, b, span); });
refl::TypeAttrDef<MinNode>()
.attr(refl::type_attr::kStructuralVisit, reinterpret_cast<void*>(&BinaryVisit<MinNode>))
@@ -536,7 +536,7 @@ TVM_DEFINE_BINOP_CONSTRUCTOR(Max);
TVM_FFI_STATIC_INIT_BLOCK() {
namespace refl = tvm::ffi::reflection;
MaxNode::RegisterReflection();
- refl::GlobalDef().def("ir.prim.Max",
+ refl::GlobalDef().def("prim.Max",
[](PrimExpr a, PrimExpr b, Span span) { return Max(a, b, span); });
refl::TypeAttrDef<MaxNode>()
.attr(refl::type_attr::kStructuralVisit, reinterpret_cast<void*>(&BinaryVisit<MaxNode>))
@@ -551,7 +551,7 @@ TVM_DEFINE_CMPOP_CONSTRUCTOR(EQ);
TVM_FFI_STATIC_INIT_BLOCK() {
namespace refl = tvm::ffi::reflection;
EQNode::RegisterReflection();
- refl::GlobalDef().def("ir.prim.EQ",
+ refl::GlobalDef().def("prim.EQ",
[](PrimExpr a, PrimExpr b, Span span) { return EQ(a, b, span); });
refl::TypeAttrDef<EQNode>()
.attr(refl::type_attr::kStructuralVisit, reinterpret_cast<void*>(&BinaryVisit<EQNode>))
@@ -566,7 +566,7 @@ TVM_DEFINE_CMPOP_CONSTRUCTOR(NE);
TVM_FFI_STATIC_INIT_BLOCK() {
namespace refl = tvm::ffi::reflection;
NENode::RegisterReflection();
- refl::GlobalDef().def("ir.prim.NE",
+ refl::GlobalDef().def("prim.NE",
[](PrimExpr a, PrimExpr b, Span span) { return NE(a, b, span); });
refl::TypeAttrDef<NENode>()
.attr(refl::type_attr::kStructuralVisit, reinterpret_cast<void*>(&BinaryVisit<NENode>))
@@ -581,7 +581,7 @@ TVM_DEFINE_CMPOP_CONSTRUCTOR(LT);
TVM_FFI_STATIC_INIT_BLOCK() {
namespace refl = tvm::ffi::reflection;
LTNode::RegisterReflection();
- refl::GlobalDef().def("ir.prim.LT",
+ refl::GlobalDef().def("prim.LT",
[](PrimExpr a, PrimExpr b, Span span) { return LT(a, b, span); });
refl::TypeAttrDef<LTNode>()
.attr(refl::type_attr::kStructuralVisit, reinterpret_cast<void*>(&BinaryVisit<LTNode>))
@@ -596,7 +596,7 @@ TVM_DEFINE_CMPOP_CONSTRUCTOR(LE);
TVM_FFI_STATIC_INIT_BLOCK() {
namespace refl = tvm::ffi::reflection;
LENode::RegisterReflection();
- refl::GlobalDef().def("ir.prim.LE",
+ refl::GlobalDef().def("prim.LE",
[](PrimExpr a, PrimExpr b, Span span) { return LE(a, b, span); });
refl::TypeAttrDef<LENode>()
.attr(refl::type_attr::kStructuralVisit, reinterpret_cast<void*>(&BinaryVisit<LENode>))
@@ -611,7 +611,7 @@ TVM_DEFINE_CMPOP_CONSTRUCTOR(GT);
TVM_FFI_STATIC_INIT_BLOCK() {
namespace refl = tvm::ffi::reflection;
GTNode::RegisterReflection();
- refl::GlobalDef().def("ir.prim.GT",
+ refl::GlobalDef().def("prim.GT",
[](PrimExpr a, PrimExpr b, Span span) { return GT(a, b, span); });
refl::TypeAttrDef<GTNode>()
.attr(refl::type_attr::kStructuralVisit, reinterpret_cast<void*>(&BinaryVisit<GTNode>))
@@ -626,7 +626,7 @@ TVM_DEFINE_CMPOP_CONSTRUCTOR(GE);
TVM_FFI_STATIC_INIT_BLOCK() {
namespace refl = tvm::ffi::reflection;
GENode::RegisterReflection();
- refl::GlobalDef().def("ir.prim.GE",
+ refl::GlobalDef().def("prim.GE",
[](PrimExpr a, PrimExpr b, Span span) { return GE(a, b, span); });
refl::TypeAttrDef<GENode>()
.attr(refl::type_attr::kStructuralVisit, reinterpret_cast<void*>(&BinaryVisit<GENode>))
@@ -656,7 +656,7 @@ And::And(PrimExpr a, PrimExpr b, Span span) {
TVM_FFI_STATIC_INIT_BLOCK() {
namespace refl = tvm::ffi::reflection;
AndNode::RegisterReflection();
- refl::GlobalDef().def("ir.prim.And",
+ refl::GlobalDef().def("prim.And",
[](PrimExpr a, PrimExpr b, Span span) { return And(a, b, span); });
refl::TypeAttrDef<AndNode>()
.attr(refl::type_attr::kStructuralVisit, reinterpret_cast<void*>(&BinaryVisit<AndNode>))
@@ -686,7 +686,7 @@ Or::Or(PrimExpr a, PrimExpr b, Span span) {
TVM_FFI_STATIC_INIT_BLOCK() {
namespace refl = tvm::ffi::reflection;
OrNode::RegisterReflection();
- refl::GlobalDef().def("ir.prim.Or",
+ refl::GlobalDef().def("prim.Or",
[](PrimExpr a, PrimExpr b, Span span) { return Or(a, b, span); });
refl::TypeAttrDef<OrNode>()
.attr(refl::type_attr::kStructuralVisit, reinterpret_cast<void*>(&BinaryVisit<OrNode>))
@@ -717,7 +717,7 @@ TVM_FFI_STATIC_INIT_BLOCK() {
.attr(refl::type_attr::kStructuralMaybeInplaceMutate,
reinterpret_cast<void*>(&NotMaybeInplaceMutate));
- refl::GlobalDef().def("ir.prim.Not", [](PrimExpr a, Span span) { return Not(a, span); });
+ refl::GlobalDef().def("prim.Not", [](PrimExpr a, Span span) { return Not(a, span); });
}
// Select
@@ -753,9 +753,10 @@ TVM_FFI_STATIC_INIT_BLOCK() {
.attr(refl::type_attr::kStructuralMaybeInplaceMutate,
reinterpret_cast<void*>(&SelectMaybeInplaceMutate));
- refl::GlobalDef().def("ir.prim.Select",
- [](PrimExpr condition, PrimExpr true_value, PrimExpr false_value,
- Span span) { return Select(condition, true_value, false_value, span); });
+ refl::GlobalDef().def(
+ "prim.Select", [](PrimExpr condition, PrimExpr true_value, PrimExpr false_value, Span span) {
+ return Select(condition, true_value, false_value, span);
+ });
}
// Let
@@ -782,7 +783,7 @@ TVM_FFI_STATIC_INIT_BLOCK() {
.attr(refl::type_attr::kStructuralMaybeInplaceMutate,
reinterpret_cast<void*>(&LetMaybeInplaceMutate));
- refl::GlobalDef().def("ir.prim.Let", [](Var var, PrimExpr value, PrimExpr body, Span span) {
+ refl::GlobalDef().def("prim.Let", [](Var var, PrimExpr value, PrimExpr body, Span span) {
return Let(var, value, body, span);
});
}
diff --git a/src/ir/prim/int_operator.h b/src/ir/prim/int_operator.h
new file mode 100644
index 0000000000..ece6ac2e50
--- /dev/null
+++ b/src/ir/prim/int_operator.h
@@ -0,0 +1,77 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied. See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+
+/*!
+ * \file int_operator.h
+ * \brief Integer division and remainder helpers for primitive operations.
+ */
+#ifndef TVM_IR_PRIM_INT_OPERATOR_H_
+#define TVM_IR_PRIM_INT_OPERATOR_H_
+
+#include <cstdint>
+
+namespace tvm {
+namespace prim {
+namespace detail {
+
+/*!
+ * \brief Perform trunc division of two integers.
+ * \param x The left operand.
+ * \param y The right operand.
+ * \return the result.
+ */
+inline int64_t truncdiv(int64_t x, int64_t y) { return x / y; }
+
+/*!
+ * \brief Compute the truncdiv remainder of two integers.
+ * \param x The left operand.
+ * \param y The right operand.
+ * \return the result.
+ */
+inline int64_t truncmod(int64_t x, int64_t y) { return x % y; }
+
+/*!
+ * \brief Perform floor division of two integers.
+ * \param x The left operand.
+ * \param y The right operand.
+ * \return the result.
+ */
+inline int64_t floordiv(int64_t x, int64_t y) {
+ int64_t rdiv = x / y;
+ int64_t rmod = x % y;
+ bool is_floor_div = (y >= 0 && rmod >= 0) || (y < 0 && rmod <= 0);
+ return is_floor_div ? rdiv : (rdiv - 1);
+}
+
+/*!
+ * \brief Compute the floordiv remainder of two integers.
+ * \param x The left operand.
+ * \param y The right operand.
+ * \return the result.
+ */
+inline int64_t floormod(int64_t x, int64_t y) {
+ int64_t rmod = x % y;
+ bool is_floor_div = (y >= 0 && rmod >= 0) || (y < 0 && rmod <= 0);
+ return is_floor_div ? rmod : rmod + y;
+}
+
+} // namespace detail
+} // namespace prim
+} // namespace tvm
+#endif // TVM_IR_PRIM_INT_OPERATOR_H_
diff --git a/src/ir/prim/op.cc b/src/ir/prim/op.cc
new file mode 100644
index 0000000000..2b86131b34
--- /dev/null
+++ b/src/ir/prim/op.cc
@@ -0,0 +1,773 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied. See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+#include <tvm/ffi/reflection/registry.h>
+#include <tvm/ir/prim/builtin.h>
+#include <tvm/ir/prim/op.h>
+
+#include <cmath>
+
+#include "const_fold.h"
+#include "op_utils.h"
+
+namespace tvm {
+using prim::is_const_int;
+using prim::MakeConst;
+using namespace prim::detail;
+namespace {
+// File-local helper: true if `expr` is a call to prim::builtin::vscale().
+bool IsVScaleCall(const PrimExpr& expr) {
+ if (const auto* call = expr.as<CallNode>()) {
+ return call->op.same_as(prim::builtin::vscale());
+ }
+ return false;
+}
+
+} // namespace
+namespace prim::detail {
+void BroadcastToMatchLanes(PrimExpr& op_a, PrimExpr& op_b) { // NOLINT(*)
+ PrimType ty_a = op_a.ty();
+ PrimType ty_b = op_b.ty();
+
+ if (!ty_a.IsScalableVector() && !ty_a.IsFixedLengthVector() &&
+ (ty_b.IsScalableVector() || ty_b.IsFixedLengthVector())) {
+ if (ty_b.IsScalableVector()) {
+ PrimType i32_ty = PrimType::Int(32);
+ op_a = prim::Broadcast(
+ op_a, prim::Mul(ty_b.VScaleFactor(),
+ Call(i32_ty, prim::builtin::vscale(), {}).as_or_throw<PrimExpr>()));
+ } else {
+ op_a = prim::Broadcast(op_a, ty_b.lanes());
+ }
+ }
+}
+
+PrimType PromoteBinaryOpType(PrimType lhs_ty, PrimType rhs_ty) {
+ if (lhs_ty->dtype == rhs_ty->dtype) {
+ return lhs_ty;
+ }
+
+ // Keep conversion behavior consistent with the previous DataType-based path.
+ if (IsFloatType(lhs_ty) && IsFloatType(rhs_ty)) {
+ return lhs_ty.bits() < rhs_ty.bits() ? rhs_ty : lhs_ty;
+ } else if (!IsFloatType(lhs_ty) && IsFloatType(rhs_ty)) {
+ return rhs_ty;
+ } else if (IsFloatType(lhs_ty) && !IsFloatType(rhs_ty)) {
+ return lhs_ty;
+ } else if (!IsBFloat16Type(lhs_ty) && IsBFloat16Type(rhs_ty)) {
+ return rhs_ty;
+ } else if (IsBFloat16Type(lhs_ty) && !IsBFloat16Type(rhs_ty)) {
+ return lhs_ty;
+ } else if (!IsFloat8Type(lhs_ty) && IsFloat8Type(rhs_ty)) {
+ return rhs_ty;
+ } else if (IsFloat8Type(lhs_ty) && !IsFloat8Type(rhs_ty)) {
+ return lhs_ty;
+ } else if (!IsFloat6Type(lhs_ty) && IsFloat6Type(rhs_ty)) {
+ return rhs_ty;
+ } else if (IsFloat6Type(lhs_ty) && !IsFloat6Type(rhs_ty)) {
+ return lhs_ty;
+ } else if (!IsFloat4Type(lhs_ty) && IsFloat4Type(rhs_ty)) {
+ return rhs_ty;
+ } else if (IsFloat4Type(lhs_ty) && !IsFloat4Type(rhs_ty)) {
+ return lhs_ty;
+ } else if (lhs_ty.MatchesCode(DLDataTypeCode::kDLBool) &&
+ rhs_ty.MatchesCode(DLDataTypeCode::kDLInt, DLDataTypeCode::kDLUInt)) {
+ return rhs_ty;
+ } else if (lhs_ty.MatchesCode(DLDataTypeCode::kDLInt, DLDataTypeCode::kDLUInt) &&
+ rhs_ty.MatchesCode(DLDataTypeCode::kDLBool)) {
+ return lhs_ty;
+ } else if ((lhs_ty.MatchesCode(DLDataTypeCode::kDLInt) &&
+ rhs_ty.MatchesCode(DLDataTypeCode::kDLInt)) ||
+ (lhs_ty.MatchesCode(DLDataTypeCode::kDLUInt) &&
+ rhs_ty.MatchesCode(DLDataTypeCode::kDLUInt))) {
+ return lhs_ty.bits() < rhs_ty.bits() ? rhs_ty : lhs_ty;
+ } else if ((lhs_ty.MatchesCode(DLDataTypeCode::kDLInt) &&
+ rhs_ty.MatchesCode(DLDataTypeCode::kDLUInt)) ||
+ (lhs_ty.MatchesCode(DLDataTypeCode::kDLUInt) &&
+ rhs_ty.MatchesCode(DLDataTypeCode::kDLInt))) {
+ if (lhs_ty.bits() < rhs_ty.bits()) {
+ return rhs_ty;
+ } else if (lhs_ty.bits() > rhs_ty.bits()) {
+ return lhs_ty;
+ } else {
+ return lhs_ty.MatchesCode(DLDataTypeCode::kDLUInt) ? lhs_ty
+ : lhs_ty.WithCode(DLDataTypeCode::kDLUInt);
+ }
+ } else {
+ TVM_FFI_THROW(InternalError) << "Cannot match type " << lhs_ty->dtype << " vs "
+ << rhs_ty->dtype;
+ }
+ return lhs_ty;
+}
+
+// The public function with a quick checking path.
+void BinaryOpMatchTypes(PrimExpr& lhs, PrimExpr& rhs, Span span) { // NOLINT(*)
+ TVM_FFI_CHECK(lhs.defined(), ValueError) << "`lhs` is null in the binary operator";
+ TVM_FFI_CHECK(rhs.defined(), ValueError) << "`rhs` is null in the binary operator";
+ const PrimTypeNode* lhs_ty_node = GetPrimTypeNode(lhs);
+ const PrimTypeNode* rhs_ty_node = GetPrimTypeNode(rhs);
+ if (lhs_ty_node == rhs_ty_node || lhs_ty_node->dtype == rhs_ty_node->dtype) return;
+
+ BroadcastToMatchLanes(lhs, rhs);
+ BroadcastToMatchLanes(rhs, lhs);
+
+ PrimType lhs_ty = lhs.ty();
+ PrimType rhs_ty = rhs.ty();
+
+ TVM_FFI_ICHECK(lhs_ty.IsScalableVector() == rhs_ty.IsScalableVector())
+ << "Can't match scalable and fixed length vectors";
+
+ bool lanes_match = false;
+
+ if (lhs_ty.IsScalableVector()) {
+ lanes_match = lhs_ty.VScaleFactor() == rhs_ty.VScaleFactor();
+ } else {
+ lanes_match = lhs_ty.lanes() == rhs_ty.lanes();
+ }
+
+ TVM_FFI_ICHECK(lanes_match) << "Cannot match type " << lhs_ty->dtype << " vs " << rhs_ty->dtype;
+
+ PrimType promoted_ty = PromoteBinaryOpType(lhs_ty, rhs_ty);
+ if (lhs_ty->dtype != promoted_ty->dtype) {
+ lhs = prim::cast(promoted_ty, lhs, span);
+ }
+ if (rhs_ty->dtype != promoted_ty->dtype) {
+ rhs = prim::cast(promoted_ty, rhs, span);
+ }
+}
+
+} // namespace prim::detail
+namespace prim {
+// LargeUIntImm
+PrimExpr LargeUIntImm(PrimType value_ty, int64_t low, int64_t high, Span span) {
+ return Call(value_ty, Op::Get("tirx.large_uint_imm"),
+ {IntImm(PrimType::UInt(32), low, span), IntImm(PrimType::UInt(32), high, span)}, {},
+ {}, span)
+ .as_or_throw<PrimExpr>();
+}
+
+// maximum and min limits
+PrimExpr max_value(PrimType value_ty, Span span) {
+ PrimType dtype = value_ty;
+ TVM_FFI_ICHECK_EQ(dtype.lanes(), 1);
+ if (dtype.MatchesCode(DLDataTypeCode::kDLInt)) {
+ if (dtype.bits() == 64) {
+ return IntImm(value_ty, std::numeric_limits<int64_t>::max(), span);
+ } else if (dtype.bits() < 64) {
+ int64_t val = 1;
+ val = (val << (dtype.bits() - 1)) - 1;
+ return IntImm(value_ty, val, span);
+ }
+ } else if (dtype.MatchesCode(DLDataTypeCode::kDLUInt)) {
+ if (dtype.bits() == 64) {
+ return MakeConst(dtype, std::numeric_limits<uint64_t>::max(), span);
+ } else if (dtype.bits() < 64) {
+ uint64_t val = 1;
+ val = (val << static_cast<uint64_t>(dtype.bits())) - 1;
+ return IntImm(value_ty, static_cast<int64_t>(val), span);
+ }
+ } else if (IsFloatType(dtype)) {
+ if (dtype.bits() == 64) {
+ return FloatImm(value_ty, std::numeric_limits<double>::max(), span);
+ } else if (dtype.bits() == 32) {
+ return FloatImm(value_ty, std::numeric_limits<float>::max(), span);
+ } else if (dtype.bits() == 16) {
+ return FloatImm(value_ty, 65504.0, span);
+ }
+ } else if (IsBFloat16Type(dtype)) {
+ return FloatImm(value_ty, std::numeric_limits<float>::max(), span);
+ } else if (IsFloat8Type(dtype)) {
+ // according to https://arxiv.org/pdf/2209.05433.pdf
+ if (dtype.code() == DLDataTypeCode::kDLFloat8_e5m2) {
+ return FloatImm(value_ty, 57344.0, span);
+ } else if (dtype.code() == DLDataTypeCode::kDLFloat8_e5m2fnuz) {
+ return FloatImm(value_ty, 57344.0, span);
+ } else if (dtype.code() == DLDataTypeCode::kDLFloat8_e4m3fn) {
... 7158 lines suppressed ...