From 9385dc137c4e0122b6fd3cdce5f5158fa358c9a9 Mon Sep 17 00:00:00 2001 From: Tianqi Chen Date: Sat, 17 Oct 2026 02:58:44 +0000 Subject: [PATCH 1/4] [REFACTOR] Keep attribute contracts with their owning interfaces Name operator attribute headers for their key declarations and type colocated keys using their existing value contracts. Preserve erased-map and callback compatibility boundaries, and expose optional typed Target backend configuration without changing serialized runtime configuration. Use tensor type alignment directly when lowering allocations and emitting CUDA storage declarations. Keep the Hexagon device constructor in the runtime namespace. --- docs/arch/fusion.rst | 2 +- include/tvm/relax/analysis.h | 2 +- include/tvm/relax/expr.h | 41 ++++---- .../tvm/relax/{op_attr_types.h => op_attr.h} | 10 +- include/tvm/s_tir/function.h | 12 ++- include/tvm/s_tir/meta_schedule/database.h | 2 +- include/tvm/s_tir/stmt.h | 95 +++++++++++-------- include/tvm/topi/nn/pooling.h | 10 +- include/tvm/topi/transform.h | 2 +- .../analysis/computable_at_compile_time.cc | 2 +- src/relax/analysis/graph_partitioner.h | 2 +- src/relax/analysis/tir_op_pattern_kind.cc | 2 +- src/relax/analysis/well_formed.cc | 15 ++- .../backend/adreno/annotate_custom_storage.cc | 9 +- .../adreno/fold_vdevice_scope_change.cc | 2 +- src/relax/backend/contrib/clml/codegen.cc | 6 +- .../contrib/codegen_json/codegen_json.h | 8 +- src/relax/backend/contrib/cublas/codegen.cc | 2 +- src/relax/backend/contrib/cudnn/codegen.cc | 2 +- src/relax/backend/contrib/cutlass/codegen.cc | 4 +- src/relax/backend/contrib/dnnl/codegen.cc | 2 +- .../backend/contrib/example_npu/codegen.cc | 2 +- src/relax/backend/contrib/hipblas/codegen.cc | 2 +- src/relax/backend/contrib/nnapi/codegen.cc | 2 +- src/relax/backend/contrib/tensorrt/codegen.cc | 2 +- src/relax/backend/vm/codegen_vm.cc | 9 +- src/relax/backend/vm/lower_runtime_builtin.cc | 2 +- src/relax/backend/vm/vm_shape_lower.cc | 13 +-- src/relax/ir/block_builder.cc | 2 +- src/relax/ir/dataflow_expr_rewriter.cc | 2 +- src/relax/ir/dependent_type.cc | 2 +- src/relax/op/ccl/ccl.cc | 2 +- src/relax/op/distributed/binary.cc | 2 +- src/relax/op/distributed/ccl.cc | 2 +- src/relax/op/distributed/distributed.cc | 2 +- src/relax/op/distributed/linear_algebra.cc | 2 +- src/relax/op/distributed/manipulate.cc | 2 +- src/relax/op/distributed/nn.cc | 2 +- src/relax/op/distributed/op.cc | 2 +- src/relax/op/distributed/statistical.h | 2 +- src/relax/op/distributed/unary.h | 2 +- src/relax/op/distributed/utils.h | 2 +- src/relax/op/image/resize.cc | 2 +- src/relax/op/memory/view.cc | 2 +- src/relax/op/nn/attention.cc | 2 +- src/relax/op/nn/convolution.cc | 2 +- src/relax/op/nn/nn.cc | 2 +- src/relax/op/nn/pooling.cc | 2 +- src/relax/op/op.cc | 2 +- src/relax/op/op_common.h | 2 +- src/relax/op/tensor/binary.h | 2 +- src/relax/op/tensor/create.cc | 2 +- src/relax/op/tensor/datatype.cc | 2 +- src/relax/op/tensor/grad.cc | 2 +- src/relax/op/tensor/index.cc | 2 +- src/relax/op/tensor/inspect.cc | 8 +- src/relax/op/tensor/linear_algebra.cc | 2 +- src/relax/op/tensor/manipulate.cc | 2 +- src/relax/op/tensor/qdq.cc | 2 +- src/relax/op/tensor/sampling.cc | 2 +- src/relax/op/tensor/search.cc | 2 +- src/relax/op/tensor/set.cc | 2 +- src/relax/op/tensor/sorting.cc | 2 +- src/relax/op/tensor/statistical.cc | 2 +- src/relax/op/tensor/ternary.cc | 2 +- src/relax/op/tensor/unary.cc | 2 +- src/relax/op/vision/multibox_transform_loc.cc | 2 +- src/relax/op/vision/nms.cc | 2 +- src/relax/op/vision/roi_align.cc | 2 +- src/relax/op/vision/roi_pool.cc | 2 +- src/relax/script/ir_builder/frame.cc | 5 +- src/relax/script/ir_builder/ir.cc | 4 +- src/relax/script/printer/binding.cc | 9 +- src/relax/script/printer/dependent_type.cc | 6 +- src/relax/script/printer/distributed.cc | 6 +- src/relax/script/printer/expr.cc | 6 +- src/relax/script/printer/function.cc | 23 ++--- src/relax/script/printer/global_info.cc | 4 +- src/relax/script/printer/region.cc | 6 +- src/relax/training/utils.cc | 2 +- src/relax/transform/allocate_workspace.cc | 23 +++-- src/relax/transform/alter_op_impl.cc | 10 +- .../transform/annotate_tir_op_pattern.cc | 4 +- .../attach_attr_layout_free_tensors.cc | 4 +- src/relax/transform/attach_global_symbol.cc | 20 ++-- src/relax/transform/bind_params.cc | 3 +- src/relax/transform/bundle_model_params.cc | 2 +- src/relax/transform/compute_prim_value.cc | 8 +- src/relax/transform/convert_layout.cc | 2 +- src/relax/transform/dataflow_inplace.cc | 2 +- src/relax/transform/decompose_ops.cc | 6 +- .../transform/eliminate_common_subexpr.cc | 2 +- src/relax/transform/expand_tuple_arguments.cc | 2 +- src/relax/transform/fold_constant.cc | 4 +- src/relax/transform/fuse_ops.cc | 50 +++++----- src/relax/transform/fuse_tir.cc | 12 +-- src/relax/transform/gradient.cc | 4 +- src/relax/transform/infer_amp_utils.h | 2 +- src/relax/transform/infer_layout_utils.h | 2 +- src/relax/transform/inline_functions.cc | 2 +- src/relax/transform/lambda_lift.cc | 8 +- src/relax/transform/lazy_transform_params.cc | 14 +-- src/relax/transform/legalize_ops.cc | 6 +- src/relax/transform/lift_transform_params.cc | 30 +++--- .../transform/merge_composite_functions.cc | 18 ++-- src/relax/transform/meta_schedule.cc | 3 +- src/relax/transform/normalize.cc | 6 +- src/relax/transform/remove_purity_checking.cc | 2 +- src/relax/transform/remove_unused_outputs.cc | 2 +- .../transform/remove_unused_parameters.cc | 2 +- src/relax/transform/rewrite_cuda_graph.cc | 9 +- src/relax/transform/run_codegen.cc | 19 ++-- .../specialize_function_based_on_callsite.cc | 7 +- .../transform/split_call_tir_by_pattern.cc | 20 ++-- .../transform/split_layout_rewrite_preproc.cc | 12 +-- .../transform/static_plan_block_memory.cc | 9 +- src/relax/transform/to_mixed_precision.cc | 2 +- src/relax/transform/utils.h | 2 +- src/relax/utils.cc | 2 +- .../analysis/calculate_allocated_memory.cc | 2 +- src/s_tir/analysis/check_well_formed.cc | 6 +- src/s_tir/analysis/estimate_flops.cc | 2 +- src/s_tir/meta_schedule/arg_info.cc | 6 +- .../database/schedule_fn_database.cc | 2 +- .../feature_extractor/per_store_feature.cc | 4 +- .../meta_schedule/mutator/mutate_parallel.cc | 2 +- .../meta_schedule/mutator/mutate_tile_size.cc | 2 +- .../meta_schedule/mutator/mutate_unroll.cc | 4 +- .../disallow_async_strided_mem_copy.cc | 2 +- .../postproc/rewrite_cooperative_fetch.cc | 6 +- .../meta_schedule/postproc/rewrite_layout.cc | 4 +- .../rewrite_parallel_vectorize_unroll.cc | 20 ++-- .../postproc/rewrite_reduction_block.cc | 12 +-- .../postproc/rewrite_tensorize.cc | 6 +- .../meta_schedule/postproc/verify_gpu_code.cc | 17 ++-- .../schedule_rule/add_rfactor.cc | 2 +- .../schedule_rule/apply_custom_rule.cc | 6 +- .../schedule_rule/auto_inline.cc | 4 +- .../schedule_rule/multi_level_tiling.cc | 20 ++-- .../multi_level_tiling_with_intrin.cc | 2 +- .../parallel_vectorize_unroll.cc | 9 +- .../schedule_rule/random_compute_location.cc | 6 +- .../space_generator/post_order_apply.cc | 3 +- src/s_tir/schedule/analysis/analysis.cc | 2 +- src/s_tir/schedule/primitive/annotate.cc | 4 +- .../primitive/annotate_tensor_access.cc | 4 +- .../schedule/primitive/block_annotate.cc | 8 +- .../schedule/primitive/compute_inline.cc | 4 +- src/s_tir/schedule/primitive/for_kind.cc | 6 +- src/s_tir/schedule/primitive/read_write_at.cc | 4 +- src/s_tir/script/ir_builder/frame.cc | 12 +-- .../script/ir_builder/script_complete.cc | 8 +- src/s_tir/script/printer/expr.cc | 2 +- src/s_tir/script/printer/stmt.cc | 10 +- src/s_tir/stmt.cc | 2 +- src/s_tir/stmt_functor.cc | 6 +- .../transform/annotate_irregular_loop.cc | 7 +- src/s_tir/transform/compact_tensor_region.cc | 4 +- src/s_tir/transform/default_gpu_schedule.cc | 12 ++- .../transform/inject_software_pipeline.cc | 20 ++-- src/s_tir/transform/ir_utils.cc | 4 +- src/s_tir/transform/lift_thread_binding.cc | 2 +- src/s_tir/transform/loop_partition.cc | 14 +-- .../transform/lower_cross_thread_reduction.cc | 2 +- src/s_tir/transform/lower_match_tensor.cc | 6 +- src/s_tir/transform/lower_opaque_block.cc | 17 ++-- src/s_tir/transform/lower_thread_binding.cc | 6 +- .../manifest_shared_memory_local_stage.cc | 4 +- src/s_tir/transform/memhammer_coalesce.cc | 6 +- .../transform/memhammer_lower_auto_copy.cc | 4 +- src/s_tir/transform/memhammer_rewrite_rule.h | 4 +- .../merge_shared_memory_allocations.cc | 4 +- .../remove_weight_layout_rewrite_block.cc | 2 +- src/s_tir/transform/rewrite_unsafe_select.cc | 2 +- .../transform/tensorcore_infer_fragment.cc | 4 +- src/s_tir/transform/thread_storage_sync.cc | 2 +- src/te/operation/create_function.cc | 16 ++-- src/te/script/printer/expr.cc | 4 +- ...test_s_tir_transform_lower_opaque_block.py | 50 +++++----- .../cuda/{op_attr_types.h => op_attr.h} | 8 +- tirx/include/tvm/codegen/intrin_rule.h | 2 +- tirx/include/tvm/ir/expr.h | 6 +- tirx/include/tvm/ir/function.h | 16 ++-- tirx/include/tvm/ir/module.h | 12 ++- tirx/include/tvm/ir/stmt.h | 6 +- .../tvm/script/printer/doc_translator.h | 5 +- tirx/include/tvm/target/target.h | 21 ++++ tirx/include/tvm/tirx/analysis.h | 4 +- tirx/include/tvm/tirx/function.h | 21 ++-- .../tvm/tirx/{op_attr_types.h => op/attr.h} | 10 +- tirx/include/tvm/tirx/stmt.h | 10 +- tirx/include/tvm/tirx/tile/context.h | 9 +- tirx/include/tvm/tirx/transform.h | 2 +- tirx/python/tvm_tirx/__init__.py | 2 +- tirx/python/tvm_tirx/runtime/__init__.py | 2 +- tirx/src/backend/cuda/codegen/codegen_cuda.cc | 20 ++-- tirx/src/backend/cuda/codegen/codegen_cuda.h | 2 +- .../backend/cuda/codegen/codegen_cuda_host.cc | 4 +- .../backend/cuda/codegen/intrin_rule_cuda.cc | 4 +- .../cuda/codegen/llvm/codegen_nvptx.cc | 2 +- .../cuda/codegen/llvm/intrin_rule_nvptx.cc | 2 +- tirx/src/backend/cuda/codegen/target_kind.cc | 2 +- tirx/src/backend/cuda/op/iket.cc | 2 +- tirx/src/backend/cuda/op/target_builtin.cc | 2 +- tirx/src/backend/cuda/op/tcgen05.cc | 2 +- tirx/src/backend/cuda/op/tensormap.cc | 2 +- tirx/src/backend/cuda/tile/exec_analysis.cc | 2 +- .../cuda/transforms/bind_backend_config.cc | 25 +++-- .../src/backend/cuda/transforms/lower_iket.cc | 8 +- .../hexagon/codegen/llvm/codegen_hexagon.cc | 4 +- .../codegen/llvm/intrin_rule_hexagon.cc | 2 +- .../backend/metal/codegen/codegen_metal.cc | 11 +-- .../metal/codegen/intrin_rule_metal.cc | 2 +- .../backend/metal/op/cooperative_tensor.cc | 2 +- tirx/src/backend/metal/op/target_builtin.cc | 2 +- .../backend/opencl/codegen/codegen_opencl.cc | 2 +- .../opencl/codegen/intrin_rule_opencl.cc | 2 +- tirx/src/backend/opencl/op/memory.cc | 2 +- .../rocm/codegen/llvm/codegen_amdgpu.cc | 2 +- .../rocm/codegen/llvm/intrin_rule_rocm.cc | 2 +- tirx/src/backend/trn/codegen/codegen_trn.cc | 10 +- tirx/src/backend/trn/codegen/codegen_trn.h | 2 +- tirx/src/backend/trn/op/target_builtin.cc | 2 +- .../backend/vulkan/codegen/codegen_spirv.cc | 4 +- .../vulkan/codegen/intrin_rule_spirv.cc | 2 +- .../src/backend/vulkan/codegen/spirv_utils.cc | 4 +- .../backend/webgpu/codegen/codegen_webgpu.cc | 8 +- .../webgpu/codegen/intrin_rule_webgpu.cc | 2 +- tirx/src/codegen/build_common.h | 4 +- tirx/src/codegen/c/codegen_c.cc | 6 +- tirx/src/codegen/c/codegen_c.h | 4 +- tirx/src/codegen/c/codegen_c_host.cc | 8 +- tirx/src/codegen/intrin_rule.cc | 2 +- tirx/src/codegen/llvm/codegen_aarch64.cc | 7 +- tirx/src/codegen/llvm/codegen_cpu.cc | 4 +- tirx/src/codegen/llvm/codegen_llvm.cc | 6 +- tirx/src/codegen/llvm/codegen_llvm.h | 2 +- tirx/src/codegen/llvm/intrin_rule_llvm.cc | 2 +- tirx/src/codegen/llvm/llvm_module.cc | 7 +- tirx/src/ir/expr.cc | 4 +- tirx/src/ir/module.cc | 2 +- tirx/src/ir/prim/builtin.cc | 2 +- tirx/src/ir/prim/math.cc | 2 +- tirx/src/ir/stmt.cc | 9 +- tirx/src/ir/subscript_expr.cc | 6 +- tirx/src/script/printer/doc_translator.cc | 2 +- tirx/src/script/printer/ir/call.cc | 2 +- tirx/src/script/printer/ir/ir.cc | 13 +-- tirx/src/script/printer/ir/misc.cc | 15 +-- tirx/src/script/printer/ir/prim_expr.cc | 95 ++++++++++--------- tirx/src/script/printer/ir/prim_type.cc | 17 ++-- tirx/src/script/printer/ir/stmt.cc | 23 +++-- tirx/src/script/printer/printer.cc | 5 +- tirx/src/tirx/analysis/check_well_formed.cc | 4 +- tirx/src/tirx/analysis/stmt_finding.cc | 4 +- tirx/src/tirx/analysis/verify_memory.cc | 4 +- tirx/src/tirx/ir/stmt.cc | 6 +- tirx/src/tirx/ir/tensor.cc | 4 +- tirx/src/tirx/ir/tensor_load.cc | 3 +- tirx/src/tirx/ir/transform.cc | 4 +- tirx/src/tirx/op/abi.cc | 2 +- tirx/src/tirx/op/annotation.cc | 2 +- tirx/src/tirx/op/cpu.cc | 2 +- tirx/src/tirx/op/gpu.cc | 2 +- tirx/src/tirx/op/memory.cc | 2 +- tirx/src/tirx/op/region.cc | 2 +- tirx/src/tirx/op/vector.cc | 2 +- tirx/src/tirx/script/ir_builder/frame.cc | 6 +- tirx/src/tirx/script/ir_builder/ir.cc | 6 +- tirx/src/tirx/script/printer/block.cc | 2 +- tirx/src/tirx/script/printer/expr.cc | 4 +- tirx/src/tirx/script/printer/for_loop.cc | 4 +- tirx/src/tirx/script/printer/function.cc | 24 ++--- tirx/src/tirx/script/printer/ir.cc | 5 +- tirx/src/tirx/script/printer/stmt.cc | 8 +- tirx/src/tirx/script/printer/tensor.cc | 10 +- tirx/src/tirx/tile/context.cc | 10 +- tirx/src/tirx/tile/lower.cc | 12 +-- tirx/src/tirx/transform/bind_target.cc | 36 ++++--- tirx/src/tirx/transform/function_utils.cc | 10 +- tirx/src/tirx/transform/lower_intrin.cc | 2 +- .../tirx/transform/lower_thread_allreduce.cc | 6 +- tirx/src/tirx/transform/lower_tirx_cleanup.cc | 2 +- tirx/src/tirx/transform/lower_tvm_builtin.cc | 6 +- tirx/src/tirx/transform/make_packed_api.cc | 16 ++-- tirx/src/tirx/transform/split_host_device.cc | 80 ++++++++-------- tirx/src/tirx/transform/storage_rewrite.cc | 16 ++-- tirx/src/tirx/transform/unroll_loop.cc | 12 +-- .../transform/unsupported_dtype_legalize.cc | 8 +- tirx/src/tirx/transform/vectorize_loop.cc | 6 +- tirx/tests/cpp/ir/tensor_load_test.cc | 4 +- tirx/tests/cpp/target/target_test.cc | 29 +++++- tirx/tests/python/tirx/codegen/test_cuda.py | 27 ++++-- 293 files changed, 1111 insertions(+), 993 deletions(-) rename include/tvm/relax/{op_attr_types.h => op_attr.h} (97%) rename tirx/include/tvm/backend/cuda/{op_attr_types.h => op_attr.h} (87%) rename tirx/include/tvm/tirx/{op_attr_types.h => op/attr.h} (89%) diff --git a/docs/arch/fusion.rst b/docs/arch/fusion.rst index 0ba228f43dda..4fed1521daa1 100644 --- a/docs/arch/fusion.rst +++ b/docs/arch/fusion.rst @@ -376,7 +376,7 @@ Source Code Map - IndexedForwardGraph, DominatorTree, GraphPartitioner (Union-Find) * - ``src/relax/transform/fuse_tir.cc`` - FuseTIR implementation, SymbolicMatcher - * - ``include/tvm/relax/op_attr_types.h`` + * - ``include/tvm/relax/op_attr.h`` - ``OpPatternKind`` enum definition * - ``python/tvm/relax/transform/transform.py`` - Python API: FuseOps, FuseTIR, FuseOpsByPattern, FusionPattern diff --git a/include/tvm/relax/analysis.h b/include/tvm/relax/analysis.h index 5daf98b081be..24360dfe245f 100644 --- a/include/tvm/relax/analysis.h +++ b/include/tvm/relax/analysis.h @@ -27,7 +27,7 @@ #include #include #include -#include +#include #include #include #include diff --git a/include/tvm/relax/expr.h b/include/tvm/relax/expr.h index 60de2c70d0ef..ec27709cb37e 100644 --- a/include/tvm/relax/expr.h +++ b/include/tvm/relax/expr.h @@ -496,54 +496,59 @@ TVM_DLL Expr GetShapeOf(const Expr& expr); // kPrimitive. namespace attr { /*! \brief Mark the function as a primitive function. */ -constexpr const char* kPrimitive = "Primitive"; +inline constexpr ffi::reflection::AttrKey kPrimitive{"Primitive"}; /*! * \brief Indicate the codegen that should be used for building this function. * When this is unset or set to "default", the default compilation pipeline will be used. */ -constexpr const char* kCodegen = "Codegen"; +inline constexpr ffi::reflection::AttrKey kCodegen{"Codegen"}; /*! \brief Treat the function as a composite operator. */ -constexpr const char* kComposite = "Composite"; +inline constexpr ffi::reflection::AttrKey kComposite{"Composite"}; /*! \brief Indicate the function was created by the Pattern Partitioning Pass. */ -constexpr const char* kPartitionedFromPattern = "PartitionedFromPattern"; +inline constexpr ffi::reflection::AttrKey kPartitionedFromPattern{ + "PartitionedFromPattern"}; /*! \brief The required workspace for an external function. */ -constexpr const char* kWorkspaceSize = "WorkspaceSize"; +inline constexpr ffi::reflection::AttrKey kWorkspaceSize{"WorkspaceSize"}; // Note: in the future, we prefer snake_case instead of CamelCase for attributes. // Past ones will be kept for backwards compatibility. /*! \brief Override checking purity for this function and treat as pure * (is_pure must be set to true) */ -constexpr const char* kForcePure = "relax.force_pure"; +inline constexpr ffi::reflection::AttrKey kForcePure{"relax.force_pure"}; /*! * \brief The number of inputs of a function. * If a function has the num_input attribute, the last func->params.size() - num_inputs * arguments are assumed to be weights that are fixed across invocations. */ -constexpr const char* kNumInput = "num_input"; -constexpr const char* kOpPattern = "op_pattern"; +inline constexpr ffi::reflection::AttrKey kNumInput{"num_input"}; +inline constexpr ffi::reflection::AttrKey kOpPattern{"op_pattern"}; -constexpr const char* kOpAttrs = "op_attrs"; +inline constexpr ffi::reflection::AttrKey kOpAttrs{"op_attrs"}; -constexpr const char* kOperatorName = "operator_name"; +inline constexpr ffi::reflection::AttrKey kOperatorName{"operator_name"}; -constexpr const char* kLibraryKernel = "library_kernel"; +inline constexpr ffi::reflection::AttrKey kLibraryKernel{"library_kernel"}; -constexpr const char* kCSource = "c_source"; +inline constexpr ffi::reflection::AttrKey kCSource{"c_source"}; -constexpr const char* kCSourceFmt = "c_source_fmt"; +inline constexpr ffi::reflection::AttrKey kCSourceFmt{"c_source_fmt"}; -constexpr const char* kTirVarUpperBound = "tir_var_upper_bound"; +inline constexpr ffi::reflection::AttrKey> kTirVarUpperBound{ + "tir_var_upper_bound"}; -constexpr const char* kTirVarLowerBound = "tir_var_lower_bound"; +inline constexpr ffi::reflection::AttrKey> kTirVarLowerBound{ + "tir_var_lower_bound"}; -constexpr const char* kTirNonNegativeVar = "tir_non_negative_var"; +inline constexpr ffi::reflection::AttrKey> kTirNonNegativeVar{ + "tir_non_negative_var"}; constexpr const char* kMemoryPlanDynamicFuncOutput = "relax.memory_plan_dynamic_func_output"; -constexpr const char* kCaptureSymbolicVars = "relax.rewrite_cuda_graph.capture_symbolic_vars"; +inline constexpr ffi::reflection::AttrKey> kCaptureSymbolicVars{ + "relax.rewrite_cuda_graph.capture_symbolic_vars"}; -constexpr const char* kScoped = "scoped"; +inline constexpr ffi::reflection::AttrKey kScoped{"scoped"}; } // namespace attr diff --git a/include/tvm/relax/op_attr_types.h b/include/tvm/relax/op_attr.h similarity index 97% rename from include/tvm/relax/op_attr_types.h rename to include/tvm/relax/op_attr.h index 4b222012b9f4..d543f1915473 100644 --- a/include/tvm/relax/op_attr_types.h +++ b/include/tvm/relax/op_attr.h @@ -18,11 +18,11 @@ */ /*! - * \file tvm/relax/op_attr_types.h - * \brief Data structures that can appear in operator attributes. + * \file tvm/relax/op_attr.h + * \brief Attribute keys and callback contracts for Relax operators. */ -#ifndef TVM_RELAX_OP_ATTR_TYPES_H_ -#define TVM_RELAX_OP_ATTR_TYPES_H_ +#ifndef TVM_RELAX_OP_ATTR_H_ +#define TVM_RELAX_OP_ATTR_H_ #include #include @@ -170,4 +170,4 @@ inline constexpr ffi::reflection::AttrKey kInferTypeWithB } // namespace relax } // namespace tvm -#endif // TVM_RELAX_OP_ATTR_TYPES_H_ +#endif // TVM_RELAX_OP_ATTR_H_ diff --git a/include/tvm/s_tir/function.h b/include/tvm/s_tir/function.h index 4be636364735..8c6b9674a956 100644 --- a/include/tvm/s_tir/function.h +++ b/include/tvm/s_tir/function.h @@ -19,18 +19,22 @@ #ifndef TVM_S_TIR_FUNCTION_H_ #define TVM_S_TIR_FUNCTION_H_ +#include +#include + namespace tvm { namespace s_tir { namespace attr { /*! * \brief Mark the function as scheduled, so the default schedule will pass will skip it. * - * Type: IntImm + * Type: bool */ -constexpr const char* kIsScheduled = "tirx.is_scheduled"; +inline constexpr ffi::reflection::AttrKey kIsScheduled{"tirx.is_scheduled"}; /*! \brief Mark the tensors which is const access and can be transformed layout. */ -constexpr const char* kLayoutFreeTensors = "layout_free_tensors"; +inline constexpr ffi::reflection::AttrKey> kLayoutFreeTensors{ + "layout_free_tensors"}; /*! * \brief Marks the layout transforms to be used for a tensor. * @@ -38,7 +42,7 @@ constexpr const char* kLayoutFreeTensors = "layout_free_tensors"; * Function attributes for TIR. */ constexpr const char* kLayoutTransforms = "layout_transforms"; -constexpr const char* kEstimatedFlops = "estimated_flops"; +inline constexpr ffi::reflection::AttrKey kEstimatedFlops{"estimated_flops"}; } // namespace attr } // namespace s_tir diff --git a/include/tvm/s_tir/meta_schedule/database.h b/include/tvm/s_tir/meta_schedule/database.h index 2748989e4177..564e41d385df 100644 --- a/include/tvm/s_tir/meta_schedule/database.h +++ b/include/tvm/s_tir/meta_schedule/database.h @@ -550,7 +550,7 @@ class Database : public ffi::ObjectRef { }; namespace attr { -constexpr const char* kTaskName = "task_name"; +inline constexpr ffi::reflection::AttrKey kTaskName{"task_name"}; } // namespace attr } // namespace meta_schedule diff --git a/include/tvm/s_tir/stmt.h b/include/tvm/s_tir/stmt.h index f534ffc6a242..2315dfe6cd11 100644 --- a/include/tvm/s_tir/stmt.h +++ b/include/tvm/s_tir/stmt.h @@ -239,85 +239,96 @@ TVM_DLL const Op& manual_sync(); namespace attr { /*! \brief Hardware thread axis bound to a parallel loop during S-TIR scheduling. */ -constexpr const char* kThreadBinding = "thread_binding"; +inline constexpr ffi::reflection::AttrKey kThreadBinding{"thread_binding"}; /*! * \brief SBlock annotation selecting a write-tensor index for double buffering in * InjectSoftwarePipeline. */ -constexpr const char* kDoubleBufferScope = "double_buffer_scope"; +inline constexpr ffi::reflection::AttrKey kDoubleBufferScope{"double_buffer_scope"}; /*! * \brief String-valued allocation Call attribute containing the TensorCore fragment shape */ -constexpr const char* kFragmentShape = "fragment_shape"; +inline constexpr ffi::reflection::AttrKey kFragmentShape{"fragment_shape"}; /*! * \brief String-valued allocation Call attribute containing the TensorCore fragment layout */ -constexpr const char* kFragmentLayout = "fragment_layout"; +inline constexpr ffi::reflection::AttrKey kFragmentLayout{"fragment_layout"}; /*! * \brief For annotation requesting partitioning when its PrimExpr value is provably true. */ -constexpr const char* kLoopPartitionHint = "loop_partition_hint"; +inline constexpr ffi::reflection::AttrKey kLoopPartitionHint{"loop_partition_hint"}; // ----------------------------------------------------------------------- // meta_schedule annotations // ----------------------------------------------------------------------- /*! \brief Mark the tiling structure of blocks that are applied by rule Multi-Level-Tiling */ -constexpr const char* kMetaScheduleTilingStructure = "meta_schedule.tiling_structure"; +inline constexpr ffi::reflection::AttrKey kMetaScheduleTilingStructure{ + "meta_schedule.tiling_structure"}; /*! * \brief Mark that the loop should be further skip and bound to environment threads to enable * cooperative fetching. */ -constexpr const char* kMetaScheduleCooperativeFetch = "meta_schedule.cooperative_fetch"; +inline constexpr ffi::reflection::AttrKey kMetaScheduleCooperativeFetch{ + "meta_schedule.cooperative_fetch"}; /*! \brief The allowed range of thread extent in thread bindings */ -constexpr const char* kMetaScheduleThreadExtentLowInclusive = - "meta_schedule.thread_extent_low_inclusive"; +inline constexpr ffi::reflection::AttrKey kMetaScheduleThreadExtentLowInclusive{ + "meta_schedule.thread_extent_low_inclusive"}; /*! \brief The allowed range of thread extent in thread bindings */ -constexpr const char* kMetaScheduleThreadExtentHighInclusive = - "meta_schedule.thread_extent_high_inclusive"; +inline constexpr ffi::reflection::AttrKey kMetaScheduleThreadExtentHighInclusive{ + "meta_schedule.thread_extent_high_inclusive"}; /*! \brief Mark the block whose producer needs to be applied by rule Random-Compute-Location */ -constexpr const char* kMetaScheduleRandomComputeProducer = "meta_schedule.random_compute_producer"; +inline constexpr ffi::reflection::AttrKey kMetaScheduleRandomComputeProducer{ + "meta_schedule.random_compute_producer"}; /*! \brief Mark auto-parallel setting on the block. */ -constexpr const char* kMetaScheduleParallel = "meta_schedule.parallel"; +inline constexpr ffi::reflection::AttrKey kMetaScheduleParallel{"meta_schedule.parallel"}; /*! \brief Mark auto-vectorize setting on the block. */ -constexpr const char* kMetaScheduleVectorize = "meta_schedule.vectorize"; +inline constexpr ffi::reflection::AttrKey kMetaScheduleVectorize{ + "meta_schedule.vectorize"}; /*! \brief Mark auto-unroll setting on the block. */ -constexpr const char* kMetaScheduleUnrollExplicit = "meta_schedule.unroll_explicit"; +inline constexpr ffi::reflection::AttrKey kMetaScheduleUnrollExplicit{ + "meta_schedule.unroll_explicit"}; /*! \brief Mark auto-unroll setting on the block. */ -constexpr const char* kMetaScheduleUnrollImplicit = "meta_schedule.unroll_implicit"; +inline constexpr ffi::reflection::AttrKey kMetaScheduleUnrollImplicit{ + "meta_schedule.unroll_implicit"}; /*! \brief Mark that a block should be further rewritten using tensorization. */ -constexpr const char* kMetaScheduleAutoTensorize = "meta_schedule.auto_tensorize"; +inline constexpr ffi::reflection::AttrKey kMetaScheduleAutoTensorize{ + "meta_schedule.auto_tensorize"}; /*! \brief Mark that a block is a preprocessor block for layout rewrite. */ -constexpr const char* kMetaScheduleLayoutRewritePreproc = "meta_schedule.layout_rewrite_preproc"; +inline constexpr ffi::reflection::AttrKey kMetaScheduleLayoutRewritePreproc{ + "meta_schedule.layout_rewrite_preproc"}; /*! * \brief Mark that the init statement of a block should be further rewritten using tensorization. */ -constexpr const char* kMetaScheduleAutoTensorizeInit = "meta_schedule.auto_tensorize_init"; +inline constexpr ffi::reflection::AttrKey kMetaScheduleAutoTensorizeInit{ + "meta_schedule.auto_tensorize_init"}; /*! \brief Mark that a block is disallowed in auto inline. */ -constexpr const char* kMetaScheduleInlineRule = "meta_schedule.inline_rule"; +inline constexpr ffi::reflection::AttrKey kMetaScheduleInlineRule{ + "meta_schedule.inline_rule"}; // ----------------------------------------------------------------------- // Schedule primitive / SBlock annotations // ----------------------------------------------------------------------- /*! \brief TensorAllocatedAddresses for block-owned allocations and match tensors. */ -constexpr const char* kTensorAllocatedAddr = "s_tir.tensor_allocated_addr"; +inline constexpr ffi::reflection::AttrKey kTensorAllocatedAddr{ + "s_tir.tensor_allocated_addr"}; /*! * \brief Mark whether the script-completer need to fill in missing access region @@ -326,7 +337,8 @@ constexpr const char* kTensorAllocatedAddr = "s_tir.tensor_allocated_addr"; * if (mask & 1) the read region should be detected, * if (mask & 2) the write region should be detected. */ -constexpr const char* kScriptParsingDetectAccess = "tirx.script_parsing_detect_access"; +inline constexpr ffi::reflection::AttrKey kScriptParsingDetectAccess{ + "tirx.script_parsing_detect_access"}; /*! * \brief Mark that the block need to add predicate for block var bounds during lowering @@ -334,19 +346,23 @@ constexpr const char* kScriptParsingDetectAccess = "tirx.script_parsing_detect_a constexpr const char* kRequireBlockVarBoundPredicate = "require_bound_predicate"; /*! \brief Mark the stage of a statement in the software pipeline */ -constexpr const char* kSoftwarePipelineStage = "software_pipeline_stage"; +inline constexpr ffi::reflection::AttrKey> kSoftwarePipelineStage{ + "software_pipeline_stage"}; /*! \brief Mark the order of a statement in the software pipeline */ -constexpr const char* kSoftwarePipelineOrder = "software_pipeline_order"; +inline constexpr ffi::reflection::AttrKey> kSoftwarePipelineOrder{ + "software_pipeline_order"}; /*! \brief List stages in the software pipeline that should run asynchronously * \note All statements in the provided stages are assumed to have asynchronous * semantics (e.g. CUDA async global to shared memory copy). */ -constexpr const char* kSoftwarePipelineAsyncStages = "software_pipeline_async_stages"; +inline constexpr ffi::reflection::AttrKey> kSoftwarePipelineAsyncStages{ + "software_pipeline_async_stages"}; /*! \brief Mark the local stage for the shared memory access should be added. */ -constexpr const char* kManifestSharedMemoryLocalStage = "tirx.manifest_shared_memory_local_stage"; +inline constexpr ffi::reflection::AttrKey kManifestSharedMemoryLocalStage{ + "tirx.manifest_shared_memory_local_stage"}; /*! * \brief Mark alignment of tensor dimension @@ -354,47 +370,52 @@ constexpr const char* kManifestSharedMemoryLocalStage = "tirx.manifest_shared_me * (tensor_index, axis, factor, offset). * This requires the stride of an axis to be k * factor + offset. */ -constexpr const char* kTensorDimAlign = "tensor_dim_align"; +inline constexpr ffi::reflection::AttrKey< + ffi::Array>> + kTensorDimAlign{"tensor_dim_align"}; /*! \brief Mark that a block has an explicitly specified read region. * This is used to override the default read region inference in TIR. */ -constexpr const char* kExplicitReadRegion = "explicit_read_region"; +inline constexpr ffi::reflection::AttrKey> kExplicitReadRegion{ + "explicit_read_region"}; /*! \brief Mark that a block has an explicitly specified write region. * This is used to override the default write region inference in TIR. */ -constexpr const char* kExplicitWriteRegion = "explicit_write_region"; +inline constexpr ffi::reflection::AttrKey> kExplicitWriteRegion{ + "explicit_write_region"}; /*! \brief ,ark a ForNode represent an irregular loop of non-structural control flow edges. */ -constexpr const char* kIrregularLoopMark = "irregular_loop_mark"; +inline constexpr ffi::reflection::AttrKey kIrregularLoopMark{"irregular_loop_mark"}; /*! \brief Mark auto copy for memhammer */ -constexpr const char* kAutoCopy = "auto_copy"; +inline constexpr ffi::reflection::AttrKey kAutoCopy{"auto_copy"}; /*! \brief Mark local stage constraint on data copy */ -constexpr const char* kLocalStage = "local_stage"; +inline constexpr ffi::reflection::AttrKey kLocalStage{"local_stage"}; /*! \brief Mark vectorization length constraint on block */ -constexpr const char* kVectorBytes = "vector_bytes"; +inline constexpr ffi::reflection::AttrKey kVectorBytes{"vector_bytes"}; /*! * \brief Mark that a block is executed by a warp. This implies the extend of threadIdx.x is * warp size. */ -constexpr const char* kWarpExecution = "warp_execution"; +inline constexpr ffi::reflection::AttrKey kWarpExecution{"warp_execution"}; constexpr const char* kPermutedLayout = "permuted_layout"; -constexpr const char* kScheduleRule = "schedule_rule"; +inline constexpr ffi::reflection::AttrKey kScheduleRule{"schedule_rule"}; -constexpr const char* kMetaScheduleWriteCacheLevel = "s_tir.meta_schedule.write_cache_level"; +inline constexpr ffi::reflection::AttrKey> kMetaScheduleWriteCacheLevel{ + "s_tir.meta_schedule.write_cache_level"}; } // namespace attr /*! \brief Interpret and validate S-TIR thread placement on a shared loop. */ inline ffi::Optional GetThreadBinding(const ForNode* loop) { - if (auto tag = loop->annotations.Get(tvm::s_tir::attr::kThreadBinding)) { + if (auto tag = loop->annotations.Get(tvm::s_tir::attr::kThreadBinding.name)) { TVM_FFI_CHECK(loop->kind == ForKind::kParallel, ValueError) << "thread_binding requires a parallel loop"; TVM_FFI_CHECK(tag->as().has_value(), TypeError) diff --git a/include/tvm/topi/nn/pooling.h b/include/tvm/topi/nn/pooling.h index be2b9d0e112d..3e50ffd0d82c 100644 --- a/include/tvm/topi/nn/pooling.h +++ b/include/tvm/topi/nn/pooling.h @@ -362,7 +362,8 @@ inline Tensor adaptive_pool_impl(const Tensor& x, const ffi::Array& ou ffi::Map attrs; if (pool_type == kMaxPool) { - attrs.Set(tvm::s_tir::attr::kScheduleRule, tvm::ffi::String("meta_schedule.adaptive_pool_max")); + attrs.Set(tvm::s_tir::attr::kScheduleRule.name, + tvm::ffi::String("meta_schedule.adaptive_pool_max")); return tvm::te::compute( out_shape, [&](const ffi::Array& output) { @@ -373,7 +374,8 @@ inline Tensor adaptive_pool_impl(const Tensor& x, const ffi::Array& ou }, "adaptive_pool_max", "adaptive_pool_max", attrs); } else if (pool_type == kAvgPool) { - attrs.Set(tvm::s_tir::attr::kScheduleRule, tvm::ffi::String("meta_schedule.adaptive_pool_avg")); + attrs.Set(tvm::s_tir::attr::kScheduleRule.name, + tvm::ffi::String("meta_schedule.adaptive_pool_avg")); auto pool_sum = tvm::te::compute( out_shape, [&](const ffi::Array& output) { @@ -589,7 +591,7 @@ inline Tensor pool_impl_nd(const Tensor& x, const ffi::Array& kernel_s auto temp = do_pad ? pad(x, pad_before, pad_after, tvm::prim::min_value(PrimType(x->dtype)), "pad_temp") : x; - attrs.Set(tvm::s_tir::attr::kScheduleRule, tvm::ffi::String("meta_schedule.pool_max")); + attrs.Set(tvm::s_tir::attr::kScheduleRule.name, tvm::ffi::String("meta_schedule.pool_max")); return tvm::te::compute( out_shape, [&](const ffi::Array& output) { @@ -604,7 +606,7 @@ inline Tensor pool_impl_nd(const Tensor& x, const ffi::Array& kernel_s }, "pool_max", "pool_max", attrs); } else if (pool_type == kAvgPool) { - attrs.Set(tvm::s_tir::attr::kScheduleRule, tvm::ffi::String("meta_schedule.pool_avg")); + attrs.Set(tvm::s_tir::attr::kScheduleRule.name, tvm::ffi::String("meta_schedule.pool_avg")); // Pad the inputs auto temp = do_pad ? pad(x, pad_before, pad_after, 0, "pad_temp") : x; diff --git a/include/tvm/topi/transform.h b/include/tvm/topi/transform.h index c668a484c34c..d1ae30b2cd47 100644 --- a/include/tvm/topi/transform.h +++ b/include/tvm/topi/transform.h @@ -1858,7 +1858,7 @@ inline Tensor layout_transform(const Tensor& src, const std::string& src_layout, ffi::Array dst_shape = layout_converter.value().ForwardShape(src->shape); ffi::Map attrs = { - {tvm::s_tir::attr::kScheduleRule, ffi::String(schedule_rule)}, + {tvm::s_tir::attr::kScheduleRule.name, ffi::String(schedule_rule)}, // Information about layouts needed for the schedule rule {tvm::topi::attr::kSrcLayout, ffi::String(src_layout)}, {tvm::topi::attr::kDstLayout, ffi::String(dst_layout)}, diff --git a/src/relax/analysis/computable_at_compile_time.cc b/src/relax/analysis/computable_at_compile_time.cc index 2873cb04650e..e23b78c9ffae 100644 --- a/src/relax/analysis/computable_at_compile_time.cc +++ b/src/relax/analysis/computable_at_compile_time.cc @@ -44,7 +44,7 @@ class CompileTimeCollector : ExprVisitor { private: void VisitExpr_(const FunctionNode* func) override { - if (auto opt_num_input = func->attrs.GetAttr(tvm::relax::attr::kNumInput)) { + if (auto opt_num_input = func->attrs.GetAttr(tvm::relax::attr::kNumInput)) { size_t num_input = opt_num_input.value(); for (size_t i = num_input; i < func->params.size(); i++) { MarkAsKnown(func->params[i]); diff --git a/src/relax/analysis/graph_partitioner.h b/src/relax/analysis/graph_partitioner.h index 1ae994842b1f..1dbd9207c206 100644 --- a/src/relax/analysis/graph_partitioner.h +++ b/src/relax/analysis/graph_partitioner.h @@ -26,7 +26,7 @@ #define TVM_RELAX_ANALYSIS_GRAPH_PARTITIONER_H_ #include -#include +#include #include #include diff --git a/src/relax/analysis/tir_op_pattern_kind.cc b/src/relax/analysis/tir_op_pattern_kind.cc index 503f07dea09d..3af3a8b9083d 100644 --- a/src/relax/analysis/tir_op_pattern_kind.cc +++ b/src/relax/analysis/tir_op_pattern_kind.cc @@ -24,7 +24,7 @@ #include #include #include -#include +#include #include #include #include diff --git a/src/relax/analysis/well_formed.cc b/src/relax/analysis/well_formed.cc index 61a5e086b076..1cbccc40dc23 100644 --- a/src/relax/analysis/well_formed.cc +++ b/src/relax/analysis/well_formed.cc @@ -73,7 +73,7 @@ #include #include #include -#include +#include #include #include #include @@ -179,7 +179,7 @@ class WellFormedChecker : public relax::ExprVisitor, public relax::TypeVisitor { // to check again // check name in global var and gsymbol - ffi::Optional gsymbol = func->GetAttr(tvm::attr::kGlobalSymbol); + ffi::Optional gsymbol = func->GetAttr(tvm::attr::kGlobalSymbol); if (gsymbol.has_value() && gsymbol != var->name_hint) { TVM_FFI_VISIT_THROW(ValueError, func->loc) << "Name in GlobalVar is not equal to name in gsymbol: " << var @@ -281,11 +281,11 @@ class WellFormedChecker : public relax::ExprVisitor, public relax::TypeVisitor { is_dataflow_ = false; // ensure the purity attributes are valid - if (op->GetAttr(tvm::relax::attr::kForcePure).value_or(false) && !op->is_pure) { + if (op->GetAttr(tvm::relax::attr::kForcePure).value_or(false) && !op->is_pure) { TVM_FFI_VISIT_THROW(ValueError, op->loc) << "Function " << ffi::GetRef(op) << " has true for " - << tvm::relax::attr::kForcePure << " but false for is_pure; " - << tvm::relax::attr::kForcePure << " should be true only if is_pure is also true."; + << tvm::relax::attr::kForcePure.name << " but false for is_pure; " + << tvm::relax::attr::kForcePure.name << " should be true only if is_pure is also true."; } // Bare symbolic occurrences in parameter types are signature definitions, @@ -338,12 +338,11 @@ class WellFormedChecker : public relax::ExprVisitor, public relax::TypeVisitor { // if we are not forcing purity and the function is annotated as pure, it must not contain an // impure call - if (check_ty && !op->GetAttr(tvm::relax::attr::kForcePure).value_or(false) && - op->is_pure) { + if (check_ty && !op->GetAttr(tvm::relax::attr::kForcePure).value_or(false) && op->is_pure) { if (auto impure = FindImpureCall(op->body)) { TVM_FFI_VISIT_THROW(ValueError, ffi::GetRef(op)) << "Function " << op << " is annotated as pure but contains an impure call: " << impure - << ". Please set " << tvm::relax::attr::kForcePure << " to true " + << ". Please set " << tvm::relax::attr::kForcePure.name << " to true " << "or use a pure operator variant (e.g., call_pure_packed) " << "if it is necessary to override this judgment."; } diff --git a/src/relax/backend/adreno/annotate_custom_storage.cc b/src/relax/backend/adreno/annotate_custom_storage.cc index 7914cce09a05..8b429aa30b7c 100644 --- a/src/relax/backend/adreno/annotate_custom_storage.cc +++ b/src/relax/backend/adreno/annotate_custom_storage.cc @@ -244,7 +244,7 @@ #include #include #include -#include +#include #include #include @@ -383,8 +383,7 @@ class CollectConsumerScopeInfo : public ExprVisitor { template ffi::Array ExtractAttrs(const T& func) { ffi::Array op_attrs; - ffi::Optional attrs = - func->template GetAttr(tvm::relax::attr::kOpAttrs); + ffi::Optional attrs = func->template GetAttr(tvm::relax::attr::kOpAttrs); if (attrs) { if (auto val = attrs.value().as()) { op_attrs.push_back(val.value()); @@ -397,7 +396,7 @@ class CollectConsumerScopeInfo : public ExprVisitor { template ffi::Optional ExtractPattern(const T& func) { - ffi::Optional op_pat = func->template GetAttr(tvm::relax::attr::kOpPattern); + ffi::Optional op_pat = func->template GetAttr(tvm::relax::attr::kOpPattern); return op_pat; } @@ -600,7 +599,7 @@ class DefineVDevice : ExprMutator { if (func->IsInstance()) { const auto& base_func = mod_->Lookup(gv); // Only non primitive relax functions - if (base_func->HasNonzeroAttr(tvm::relax::attr::kPrimitive)) { + if (base_func->HasNonzeroAttr(tvm::relax::attr::kPrimitive.name)) { continue; } auto info = diff --git a/src/relax/backend/adreno/fold_vdevice_scope_change.cc b/src/relax/backend/adreno/fold_vdevice_scope_change.cc index db147afa9a32..1312ea608ccf 100644 --- a/src/relax/backend/adreno/fold_vdevice_scope_change.cc +++ b/src/relax/backend/adreno/fold_vdevice_scope_change.cc @@ -29,7 +29,7 @@ #include #include #include -#include +#include #include #include diff --git a/src/relax/backend/contrib/clml/codegen.cc b/src/relax/backend/contrib/clml/codegen.cc index 3eec2cecb016..cb8a678366e5 100644 --- a/src/relax/backend/contrib/clml/codegen.cc +++ b/src/relax/backend/contrib/clml/codegen.cc @@ -140,7 +140,7 @@ class OpenCLMLJSONSerializer : public JSONSerializer { TVM_FFI_ICHECK(fn_var); const auto fn = bindings_[ffi::GetRef(fn_var)].as_or_throw(); - auto opt_composite = fn->GetAttr(tvm::relax::attr::kComposite); + auto opt_composite = fn->GetAttr(tvm::relax::attr::kComposite); TVM_FFI_ICHECK(opt_composite.has_value()); std::string name = opt_composite.value(); @@ -192,7 +192,7 @@ class OpenCLMLJSONSerializer : public JSONSerializer { const auto* fn_var = cn->op.as(); TVM_FFI_ICHECK(fn_var); const auto fn = bindings_[ffi::GetRef(fn_var)].as_or_throw(); - auto opt_composite = fn->GetAttr(tvm::relax::attr::kComposite); + auto opt_composite = fn->GetAttr(tvm::relax::attr::kComposite); TVM_FFI_ICHECK(opt_composite.has_value()); nodes.pad = backend::TryGetOpInFunction(fn, "relax.nn.pad"); @@ -221,7 +221,7 @@ class OpenCLMLJSONSerializer : public JSONSerializer { const auto* fn_var = cn->op.as(); TVM_FFI_ICHECK(fn_var); const auto fn = bindings_[ffi::GetRef(fn_var)].as_or_throw(); - auto opt_composite = fn->GetAttr(tvm::relax::attr::kComposite); + auto opt_composite = fn->GetAttr(tvm::relax::attr::kComposite); TVM_FFI_ICHECK(opt_composite.has_value()); std::string name = opt_composite.value(); diff --git a/src/relax/backend/contrib/codegen_json/codegen_json.h b/src/relax/backend/contrib/codegen_json/codegen_json.h index d035a2315c50..e6fe066d9a9a 100644 --- a/src/relax/backend/contrib/codegen_json/codegen_json.h +++ b/src/relax/backend/contrib/codegen_json/codegen_json.h @@ -316,9 +316,9 @@ class JSONSerializer : public relax::MemoizedExprTranslator { extractor.Extract(const_cast(call_attr)); } else if (const auto* fn = cn->op.as()) { TVM_FFI_ICHECK(false); - auto pattern = fn->GetAttr(tvm::relax::attr::kPartitionedFromPattern); + auto pattern = fn->GetAttr(tvm::relax::attr::kPartitionedFromPattern); TVM_FFI_ICHECK(pattern.has_value()); - node->SetAttr(tvm::relax::attr::kPartitionedFromPattern, pattern.value()); + node->SetAttr(tvm::relax::attr::kPartitionedFromPattern.name, pattern.value()); } } @@ -412,7 +412,7 @@ class JSONSerializer : public relax::MemoizedExprTranslator { if (const auto* op_node = cn->op.as()) { name = op_node->name; } else if (const auto* fn = cn->op.as()) { - auto comp = fn->GetAttr(tvm::relax::attr::kComposite); + auto comp = fn->GetAttr(tvm::relax::attr::kComposite); TVM_FFI_ICHECK(comp.has_value()) << "JSON runtime only supports composite functions."; name = comp.value(); } else { @@ -442,7 +442,7 @@ class JSONSerializer : public relax::MemoizedExprTranslator { } NodeEntries VisitExpr_(const FunctionNode* fn) { - TVM_FFI_ICHECK(fn->GetAttr(tvm::relax::attr::kComposite).has_value()) + TVM_FFI_ICHECK(fn->GetAttr(tvm::relax::attr::kComposite).has_value()) << "JSON runtime only supports composite functions"; // FunctionNode should be handled by the caller. diff --git a/src/relax/backend/contrib/cublas/codegen.cc b/src/relax/backend/contrib/cublas/codegen.cc index 72a110f94165..969ec1dc0ba4 100644 --- a/src/relax/backend/contrib/cublas/codegen.cc +++ b/src/relax/backend/contrib/cublas/codegen.cc @@ -56,7 +56,7 @@ class CublasJSONSerializer : public JSONSerializer { const auto fn = bindings_[ffi::GetRef(fn_var)].as_or_throw(); TVM_FFI_ICHECK(fn.defined()) << "Expects the callee to be a function."; - auto composite_opt = fn->GetAttr(tvm::relax::attr::kComposite); + auto composite_opt = fn->GetAttr(tvm::relax::attr::kComposite); TVM_FFI_ICHECK(composite_opt.has_value()) << "Only composite functions are supported."; std::string composite_name = composite_opt.value(); diff --git a/src/relax/backend/contrib/cudnn/codegen.cc b/src/relax/backend/contrib/cudnn/codegen.cc index b2259f42de6d..939d5e029ee1 100644 --- a/src/relax/backend/contrib/cudnn/codegen.cc +++ b/src/relax/backend/contrib/cudnn/codegen.cc @@ -54,7 +54,7 @@ class cuDNNJSONSerializer : public JSONSerializer { const auto fn = bindings_[ffi::GetRef(fn_var)].as_or_throw(); TVM_FFI_ICHECK(fn.defined()) << "Expects the callee to be a function."; - auto composite_opt = fn->GetAttr(tvm::relax::attr::kComposite); + auto composite_opt = fn->GetAttr(tvm::relax::attr::kComposite); TVM_FFI_ICHECK(composite_opt.has_value()) << "Only composite functions are supported."; std::string composite_name = composite_opt.value(); diff --git a/src/relax/backend/contrib/cutlass/codegen.cc b/src/relax/backend/contrib/cutlass/codegen.cc index 05a6ec1c5d3f..5539eaacdfe2 100644 --- a/src/relax/backend/contrib/cutlass/codegen.cc +++ b/src/relax/backend/contrib/cutlass/codegen.cc @@ -214,7 +214,7 @@ class CodegenCutlass : public relax::MemoizedExprTranslator, const auto* fn_var = call->op.as(); TVM_FFI_ICHECK(fn_var); const auto func = bindings_[ffi::GetRef(fn_var)].as_or_throw(); - const auto pattern_name_opt = func->GetAttr(tvm::relax::attr::kComposite); + const auto pattern_name_opt = func->GetAttr(tvm::relax::attr::kComposite); TVM_FFI_ICHECK(pattern_name_opt) << "Only composite function is supported for CUTLASS."; auto ret = GenerateBody(call, pattern_name_opt.value(), func->attrs->dict); ext_func_body_.push_back(ret.decl); @@ -223,7 +223,7 @@ class CodegenCutlass : public relax::MemoizedExprTranslator, } OutputType VisitExpr_(const FunctionNode* fn) final { - TVM_FFI_ICHECK(fn->GetAttr(tvm::relax::attr::kComposite).has_value()) + TVM_FFI_ICHECK(fn->GetAttr(tvm::relax::attr::kComposite).has_value()) << "JSON runtime only supports composite functions"; // FunctionNode should be handled by the caller. return {}; diff --git a/src/relax/backend/contrib/dnnl/codegen.cc b/src/relax/backend/contrib/dnnl/codegen.cc index ee40262a78a9..a3fb320227ff 100644 --- a/src/relax/backend/contrib/dnnl/codegen.cc +++ b/src/relax/backend/contrib/dnnl/codegen.cc @@ -54,7 +54,7 @@ class DNNLJSONSerializer : public JSONSerializer { const auto fn = bindings_[ffi::GetRef(fn_var)].as_or_throw(); TVM_FFI_ICHECK(fn.defined()) << "Expects the callee to be a function."; - auto composite_opt = fn->GetAttr(tvm::relax::attr::kComposite); + auto composite_opt = fn->GetAttr(tvm::relax::attr::kComposite); TVM_FFI_ICHECK(composite_opt.has_value()) << "Only composite functions are supported."; std::string composite_name = composite_opt.value(); diff --git a/src/relax/backend/contrib/example_npu/codegen.cc b/src/relax/backend/contrib/example_npu/codegen.cc index dcfc2b7a67f2..e731e3640ac4 100644 --- a/src/relax/backend/contrib/example_npu/codegen.cc +++ b/src/relax/backend/contrib/example_npu/codegen.cc @@ -54,7 +54,7 @@ class ExampleNPUJSONSerializer : public JSONSerializer { const auto fn = bindings_[ffi::GetRef(fn_var)].as_or_throw(); TVM_FFI_ICHECK(fn.defined()) << "Expects the callee to be a function."; - auto composite_opt = fn->GetAttr(tvm::relax::attr::kComposite); + auto composite_opt = fn->GetAttr(tvm::relax::attr::kComposite); TVM_FFI_ICHECK(composite_opt.has_value()) << "Only composite functions are supported."; std::string composite_name = composite_opt.value(); diff --git a/src/relax/backend/contrib/hipblas/codegen.cc b/src/relax/backend/contrib/hipblas/codegen.cc index 5f7addaf2187..905a9fb37f05 100644 --- a/src/relax/backend/contrib/hipblas/codegen.cc +++ b/src/relax/backend/contrib/hipblas/codegen.cc @@ -54,7 +54,7 @@ class HipblasJSONSerializer : public JSONSerializer { const auto fn = bindings_[ffi::GetRef(fn_var)].as_or_throw(); TVM_FFI_ICHECK(fn.defined()) << "Expects the callee to be a function."; - auto composite_opt = fn->GetAttr(tvm::relax::attr::kComposite); + auto composite_opt = fn->GetAttr(tvm::relax::attr::kComposite); TVM_FFI_ICHECK(composite_opt.has_value()) << "Only composite functions are supported."; std::string composite_name = composite_opt.value(); diff --git a/src/relax/backend/contrib/nnapi/codegen.cc b/src/relax/backend/contrib/nnapi/codegen.cc index 473ea1a5aab9..7cc485926c10 100644 --- a/src/relax/backend/contrib/nnapi/codegen.cc +++ b/src/relax/backend/contrib/nnapi/codegen.cc @@ -156,7 +156,7 @@ class NNAPIJSONSerializer : public JSONSerializer { const auto fn = bindings_[ffi::GetRef(fn_var)].as_or_throw(); TVM_FFI_ICHECK(fn.defined()) << "Expects the callee to be a function."; - auto composite_opt = fn->GetAttr(tvm::relax::attr::kComposite); + auto composite_opt = fn->GetAttr(tvm::relax::attr::kComposite); TVM_FFI_ICHECK(composite_opt.has_value()) << "Only composite functions are supported."; std::string composite_name = composite_opt.value(); diff --git a/src/relax/backend/contrib/tensorrt/codegen.cc b/src/relax/backend/contrib/tensorrt/codegen.cc index 5c28437e5dd5..2221678b6429 100644 --- a/src/relax/backend/contrib/tensorrt/codegen.cc +++ b/src/relax/backend/contrib/tensorrt/codegen.cc @@ -246,7 +246,7 @@ class TensorRTJSONSerializer : public JSONSerializer { TVM_FFI_ICHECK(fn_var); const auto fn = bindings_[ffi::GetRef(fn_var)].as_or_throw(); - auto opt_composite = fn->GetAttr(tvm::relax::attr::kComposite); + auto opt_composite = fn->GetAttr(tvm::relax::attr::kComposite); TVM_FFI_ICHECK(opt_composite.has_value()); std::string name = opt_composite.value(); diff --git a/src/relax/backend/vm/codegen_vm.cc b/src/relax/backend/vm/codegen_vm.cc index 85689602407d..1367b6aed37e 100644 --- a/src/relax/backend/vm/codegen_vm.cc +++ b/src/relax/backend/vm/codegen_vm.cc @@ -27,7 +27,7 @@ #include #include #include -#include +#include #include #include #include @@ -85,7 +85,7 @@ class CodeGenVM : public ExprFunctor { } void Codegen(const Function& func) { - ffi::Optional gsymbol = func->GetAttr(tvm::attr::kGlobalSymbol); + ffi::Optional gsymbol = func->GetAttr(tvm::attr::kGlobalSymbol); TVM_FFI_ICHECK(gsymbol.has_value()) << "there should be no local functions in Relax VM codegen phase. " "Did you forget to apply LambdaLift or AttachGlobalSymbol Pass?"; @@ -334,10 +334,9 @@ class CodeGenVM : public ExprFunctor { } Instruction::Arg VisitExpr_(const ExternFuncNode* op) final { - if (ffi::Optional opt_code = - op->attrs.GetAttr(tvm::relax::attr::kCSource)) { + if (ffi::Optional opt_code = op->attrs.GetAttr(tvm::relax::attr::kCSource)) { ffi::String sym = op->global_symbol; - ffi::String fmt = op->attrs.GetAttr(tvm::relax::attr::kCSourceFmt).value_or("c"); + ffi::String fmt = op->attrs.GetAttr(tvm::relax::attr::kCSourceFmt).value_or("c"); ffi::String code = opt_code.value(); // Grep hint: grep -rn 'runtime.CSourceModuleCreate' src/ static const auto create_source = diff --git a/src/relax/backend/vm/lower_runtime_builtin.cc b/src/relax/backend/vm/lower_runtime_builtin.cc index 8ad321fa9860..bd40b086ed19 100644 --- a/src/relax/backend/vm/lower_runtime_builtin.cc +++ b/src/relax/backend/vm/lower_runtime_builtin.cc @@ -29,7 +29,7 @@ #include #include #include -#include +#include #include namespace tvm { diff --git a/src/relax/backend/vm/vm_shape_lower.cc b/src/relax/backend/vm/vm_shape_lower.cc index 48368a9d5bee..0be9aa8f723e 100644 --- a/src/relax/backend/vm/vm_shape_lower.cc +++ b/src/relax/backend/vm/vm_shape_lower.cc @@ -32,7 +32,7 @@ #include #include #include -#include +#include #include #include #include @@ -368,7 +368,7 @@ class VMShapeLowerMutator this->builder_->EmitNormalized(shape_heap_binding); std::vector match_todos; size_t num_input = func->params.size(); - if (auto opt_num_input = func->attrs.GetAttr(tvm::relax::attr::kNumInput)) { + if (auto opt_num_input = func->attrs.GetAttr(tvm::relax::attr::kNumInput)) { // If the function has the attribute 'num_input', do shape checking on for the real inputs // and skip weights. num_input = static_cast(opt_num_input.value()); @@ -750,12 +750,13 @@ class VMShapeLowerMutator // the shape_func to indicate that this is a host function // This could require us to attach target to the relax function here. tirx::Function shape_func(params, body, ret_type); - shape_func = WithAttr(std::move(shape_func), tvm::attr::kSTir, true); - shape_func = WithAttr(std::move(shape_func), tvm::attr::kScriptNamespace, ffi::String("s_tir")); - if (!shape_func->attrs.GetAttr(tvm::attr::kTarget).has_value()) { + shape_func = WithAttr(std::move(shape_func), tvm::attr::kSTir.name, true); + shape_func = + WithAttr(std::move(shape_func), tvm::attr::kScriptNamespace.name, ffi::String("s_tir")); + if (!shape_func->attrs.GetAttr(tvm::attr::kTarget.name).has_value()) { // kTarget and kIsHostFunc are mutually exclusive shape_func = - WithAttr(std::move(shape_func), tvm::tirx::attr::kIsHostFunc, true); + WithAttr(std::move(shape_func), tvm::tirx::attr::kIsHostFunc.name, true); } GlobalVar shape_func_var = builder_->AddFunction(shape_func, "shape_func"); builder_->Emit(Call(Type::Missing(), Op::Get("relax.call_tir_packed"), diff --git a/src/relax/ir/block_builder.cc b/src/relax/ir/block_builder.cc index ef92f4cfaf38..13405a90be5e 100644 --- a/src/relax/ir/block_builder.cc +++ b/src/relax/ir/block_builder.cc @@ -28,7 +28,7 @@ #include #include #include -#include +#include #include #include #include diff --git a/src/relax/ir/dataflow_expr_rewriter.cc b/src/relax/ir/dataflow_expr_rewriter.cc index d8c105990a4f..6446ccff265d 100644 --- a/src/relax/ir/dataflow_expr_rewriter.cc +++ b/src/relax/ir/dataflow_expr_rewriter.cc @@ -689,7 +689,7 @@ PatternMatchingRewriter PatternMatchingRewriter::FromModule(IRModule mod) { ffi::Map new_subroutines; for (const auto& [gvar, func] : mod->functions) { if (gvar->name_hint != "pattern" && gvar->name_hint != "replacement") { - bool is_public = func->GetAttr(tvm::attr::kGlobalSymbol).has_value(); + bool is_public = func->GetAttr(tvm::attr::kGlobalSymbol).has_value(); TVM_FFI_CHECK(!is_public, ValueError) << "Expected module to have no publicly-exposed functions " << "other than 'pattern' and 'replacement'. " diff --git a/src/relax/ir/dependent_type.cc b/src/relax/ir/dependent_type.cc index 12f896f91e2c..4430f7e527bf 100644 --- a/src/relax/ir/dependent_type.cc +++ b/src/relax/ir/dependent_type.cc @@ -272,7 +272,7 @@ TVM_FFI_STATIC_INIT_BLOCK() { ffi::FStructuralMutate::FromNative<&TensorTypeMutate>()) .attr(ffi::type_attr::kStructuralMaybeInplaceMutate, ffi::FStructuralMutate::FromNative<&TensorTypeMaybeInplaceMutate>()) - .def(tvm::type_attr::kSubscriptExprRealize, + .def(tvm::type_attr::kSubscriptExprRealize.name, [](Expr value, ffi::Array, ffi::Optional, ffi::Optional>, diff --git a/src/relax/op/ccl/ccl.cc b/src/relax/op/ccl/ccl.cc index 3fccbf665c7d..5efe5c0d4d06 100644 --- a/src/relax/op/ccl/ccl.cc +++ b/src/relax/op/ccl/ccl.cc @@ -23,7 +23,7 @@ #include #include #include -#include +#include #include diff --git a/src/relax/op/distributed/binary.cc b/src/relax/op/distributed/binary.cc index 65d561584c68..3e3452736c45 100644 --- a/src/relax/op/distributed/binary.cc +++ b/src/relax/op/distributed/binary.cc @@ -19,7 +19,7 @@ #include "binary.h" -#include +#include namespace tvm { namespace relax { diff --git a/src/relax/op/distributed/ccl.cc b/src/relax/op/distributed/ccl.cc index d7e76b095eba..0806ac77ba66 100644 --- a/src/relax/op/distributed/ccl.cc +++ b/src/relax/op/distributed/ccl.cc @@ -19,7 +19,7 @@ #include "tvm/relax/op/ccl.h" #include -#include +#include #include "utils.h" diff --git a/src/relax/op/distributed/distributed.cc b/src/relax/op/distributed/distributed.cc index 9e87d74934b2..5791451d2083 100644 --- a/src/relax/op/distributed/distributed.cc +++ b/src/relax/op/distributed/distributed.cc @@ -27,7 +27,7 @@ #include #include #include -#include +#include #include #include diff --git a/src/relax/op/distributed/linear_algebra.cc b/src/relax/op/distributed/linear_algebra.cc index e4ae8fbb0cab..6c81a855a8cb 100644 --- a/src/relax/op/distributed/linear_algebra.cc +++ b/src/relax/op/distributed/linear_algebra.cc @@ -20,7 +20,7 @@ #include "linear_algebra.h" #include -#include +#include #include namespace tvm { diff --git a/src/relax/op/distributed/manipulate.cc b/src/relax/op/distributed/manipulate.cc index ca6e123901e3..58d954313974 100644 --- a/src/relax/op/distributed/manipulate.cc +++ b/src/relax/op/distributed/manipulate.cc @@ -20,7 +20,7 @@ #include "manipulate.h" #include -#include +#include #include #include diff --git a/src/relax/op/distributed/nn.cc b/src/relax/op/distributed/nn.cc index 4664ddbf8e55..9ff821b8e2af 100644 --- a/src/relax/op/distributed/nn.cc +++ b/src/relax/op/distributed/nn.cc @@ -20,7 +20,7 @@ #include "nn.h" #include -#include +#include namespace tvm { namespace relax { diff --git a/src/relax/op/distributed/op.cc b/src/relax/op/distributed/op.cc index 027dc8ed759e..2b13f20a7b2b 100644 --- a/src/relax/op/distributed/op.cc +++ b/src/relax/op/distributed/op.cc @@ -18,7 +18,7 @@ */ #include -#include +#include #include "utils.h" diff --git a/src/relax/op/distributed/statistical.h b/src/relax/op/distributed/statistical.h index 1b7f7fd3d3a4..0e3bbda78633 100644 --- a/src/relax/op/distributed/statistical.h +++ b/src/relax/op/distributed/statistical.h @@ -25,7 +25,7 @@ #ifndef TVM_RELAX_OP_DISTRIBUTED_STATISTICAL_H_ #define TVM_RELAX_OP_DISTRIBUTED_STATISTICAL_H_ -#include +#include #include "utils.h" diff --git a/src/relax/op/distributed/unary.h b/src/relax/op/distributed/unary.h index 3d6904ecd9fa..a5dc0d6747d6 100644 --- a/src/relax/op/distributed/unary.h +++ b/src/relax/op/distributed/unary.h @@ -26,7 +26,7 @@ #define TVM_RELAX_OP_DISTRIBUTED_UNARY_H_ #include -#include +#include #include "utils.h" diff --git a/src/relax/op/distributed/utils.h b/src/relax/op/distributed/utils.h index 36ac64880b4c..ada2d967eef2 100644 --- a/src/relax/op/distributed/utils.h +++ b/src/relax/op/distributed/utils.h @@ -26,7 +26,7 @@ #define TVM_RELAX_OP_DISTRIBUTED_UTILS_H_ #include -#include +#include #include #include "../op_common.h" diff --git a/src/relax/op/image/resize.cc b/src/relax/op/image/resize.cc index 135c555970e4..ae683aee09a5 100644 --- a/src/relax/op/image/resize.cc +++ b/src/relax/op/image/resize.cc @@ -28,7 +28,7 @@ #include #include #include -#include +#include #include diff --git a/src/relax/op/memory/view.cc b/src/relax/op/memory/view.cc index 779a8459a1f1..4da2756c2dd8 100644 --- a/src/relax/op/memory/view.cc +++ b/src/relax/op/memory/view.cc @@ -27,7 +27,7 @@ #include #include #include -#include +#include namespace tvm { namespace relax { diff --git a/src/relax/op/nn/attention.cc b/src/relax/op/nn/attention.cc index e2cfbccb2228..91145d6f4ad7 100644 --- a/src/relax/op/nn/attention.cc +++ b/src/relax/op/nn/attention.cc @@ -21,7 +21,7 @@ #include #include -#include +#include #include diff --git a/src/relax/op/nn/convolution.cc b/src/relax/op/nn/convolution.cc index dac76e9c89ff..497f3924008d 100644 --- a/src/relax/op/nn/convolution.cc +++ b/src/relax/op/nn/convolution.cc @@ -27,7 +27,7 @@ #include #include #include -#include +#include #include diff --git a/src/relax/op/nn/nn.cc b/src/relax/op/nn/nn.cc index be016e74fb07..b566cba643e8 100644 --- a/src/relax/op/nn/nn.cc +++ b/src/relax/op/nn/nn.cc @@ -23,7 +23,7 @@ #include #include #include -#include +#include #include #include diff --git a/src/relax/op/nn/pooling.cc b/src/relax/op/nn/pooling.cc index e25239673a03..2d9c8fc70363 100644 --- a/src/relax/op/nn/pooling.cc +++ b/src/relax/op/nn/pooling.cc @@ -22,7 +22,7 @@ #include #include #include -#include +#include #include #include diff --git a/src/relax/op/op.cc b/src/relax/op/op.cc index 505eacedc513..ff54bbf33887 100644 --- a/src/relax/op/op.cc +++ b/src/relax/op/op.cc @@ -25,7 +25,7 @@ #include #include #include -#include +#include #include #include #include diff --git a/src/relax/op/op_common.h b/src/relax/op/op_common.h index 5eeef011e1f9..ad3e68091d47 100644 --- a/src/relax/op/op_common.h +++ b/src/relax/op/op_common.h @@ -28,7 +28,7 @@ #include #include #include -#include +#include #include #include diff --git a/src/relax/op/tensor/binary.h b/src/relax/op/tensor/binary.h index 2eef273069c9..4887bf9fa1f7 100644 --- a/src/relax/op/tensor/binary.h +++ b/src/relax/op/tensor/binary.h @@ -24,7 +24,7 @@ #ifndef TVM_RELAX_OP_TENSOR_BINARY_H_ #define TVM_RELAX_OP_TENSOR_BINARY_H_ -#include +#include #include "../op_common.h" diff --git a/src/relax/op/tensor/create.cc b/src/relax/op/tensor/create.cc index faa43bfca8f8..a278839d113c 100644 --- a/src/relax/op/tensor/create.cc +++ b/src/relax/op/tensor/create.cc @@ -28,7 +28,7 @@ #include #include #include -#include +#include #include #include diff --git a/src/relax/op/tensor/datatype.cc b/src/relax/op/tensor/datatype.cc index 8ae8547040ff..c33efce461b6 100644 --- a/src/relax/op/tensor/datatype.cc +++ b/src/relax/op/tensor/datatype.cc @@ -26,7 +26,7 @@ #include #include -#include +#include #include diff --git a/src/relax/op/tensor/grad.cc b/src/relax/op/tensor/grad.cc index 23cc690330cf..1c7c4c5de4cf 100644 --- a/src/relax/op/tensor/grad.cc +++ b/src/relax/op/tensor/grad.cc @@ -27,7 +27,7 @@ #include #include #include -#include +#include #include diff --git a/src/relax/op/tensor/index.cc b/src/relax/op/tensor/index.cc index 3feb6fbb322e..8a58194d3f46 100644 --- a/src/relax/op/tensor/index.cc +++ b/src/relax/op/tensor/index.cc @@ -28,7 +28,7 @@ #include #include #include -#include +#include #include #include diff --git a/src/relax/op/tensor/inspect.cc b/src/relax/op/tensor/inspect.cc index 1505da6c767d..5ee3330f7772 100644 --- a/src/relax/op/tensor/inspect.cc +++ b/src/relax/op/tensor/inspect.cc @@ -26,7 +26,7 @@ #include #include -#include +#include #include #include #include @@ -103,7 +103,8 @@ tirx::Function GetDLTensorField(tirx::TVMStructFieldKind field, PrimType field_t .as_or_throw()), tvm::Return(value)}); - DictAttrs attrs({{tvm::s_tir::attr::kIsScheduled, true}, {tvm::tirx::attr::kIsHostFunc, true}}); + DictAttrs attrs( + {{tvm::s_tir::attr::kIsScheduled.name, true}, {tvm::tirx::attr::kIsHostFunc.name, true}}); tirx::Function func(ffi::Array{dlpack_handle}, body, field_ty, attrs); @@ -284,7 +285,8 @@ Expr LegalizeTensorShape(const BlockBuilder& bb, const Call& call) { tvm::Bind(extent, TensorLoad(shape_buffer, {axis.as_or_throw()})), tvm::Return(extent)}); - DictAttrs attrs({{tvm::s_tir::attr::kIsScheduled, true}, {tvm::tirx::attr::kIsHostFunc, true}}); + DictAttrs attrs( + {{tvm::s_tir::attr::kIsScheduled.name, true}, {tvm::tirx::attr::kIsHostFunc.name, true}}); tirx::Function func({dlpack_handle, axis}, body, field_ty, attrs); diff --git a/src/relax/op/tensor/linear_algebra.cc b/src/relax/op/tensor/linear_algebra.cc index e39e667b73cd..671e870334b0 100644 --- a/src/relax/op/tensor/linear_algebra.cc +++ b/src/relax/op/tensor/linear_algebra.cc @@ -27,7 +27,7 @@ #include #include #include -#include +#include #include #include diff --git a/src/relax/op/tensor/manipulate.cc b/src/relax/op/tensor/manipulate.cc index eb22596e2d6d..92b2950632d1 100644 --- a/src/relax/op/tensor/manipulate.cc +++ b/src/relax/op/tensor/manipulate.cc @@ -28,7 +28,7 @@ #include #include #include -#include +#include #include #include diff --git a/src/relax/op/tensor/qdq.cc b/src/relax/op/tensor/qdq.cc index 0d0205e87dbb..d188d5536fc2 100644 --- a/src/relax/op/tensor/qdq.cc +++ b/src/relax/op/tensor/qdq.cc @@ -26,7 +26,7 @@ #include #include -#include +#include #include diff --git a/src/relax/op/tensor/sampling.cc b/src/relax/op/tensor/sampling.cc index c233b2ef1c30..d552f89895c8 100644 --- a/src/relax/op/tensor/sampling.cc +++ b/src/relax/op/tensor/sampling.cc @@ -28,7 +28,7 @@ #include #include #include -#include +#include #include diff --git a/src/relax/op/tensor/search.cc b/src/relax/op/tensor/search.cc index cffb2ab0eff3..ba3a982f9b6b 100644 --- a/src/relax/op/tensor/search.cc +++ b/src/relax/op/tensor/search.cc @@ -27,7 +27,7 @@ #include #include #include -#include +#include #include #include diff --git a/src/relax/op/tensor/set.cc b/src/relax/op/tensor/set.cc index 9680b0864dd2..36f1bf1617c8 100644 --- a/src/relax/op/tensor/set.cc +++ b/src/relax/op/tensor/set.cc @@ -27,7 +27,7 @@ #include #include #include -#include +#include #include #include diff --git a/src/relax/op/tensor/sorting.cc b/src/relax/op/tensor/sorting.cc index 0a86c9eba8e2..57921126486c 100644 --- a/src/relax/op/tensor/sorting.cc +++ b/src/relax/op/tensor/sorting.cc @@ -26,7 +26,7 @@ #include #include -#include +#include #include diff --git a/src/relax/op/tensor/statistical.cc b/src/relax/op/tensor/statistical.cc index c8e003ae9a58..dc4ce4b6ea62 100644 --- a/src/relax/op/tensor/statistical.cc +++ b/src/relax/op/tensor/statistical.cc @@ -26,7 +26,7 @@ #include #include -#include +#include #include #include diff --git a/src/relax/op/tensor/ternary.cc b/src/relax/op/tensor/ternary.cc index c80f7c33fe35..94a3d508a74a 100644 --- a/src/relax/op/tensor/ternary.cc +++ b/src/relax/op/tensor/ternary.cc @@ -26,7 +26,7 @@ #include #include -#include +#include namespace tvm { namespace relax { diff --git a/src/relax/op/tensor/unary.cc b/src/relax/op/tensor/unary.cc index cc00ca72295d..ab82f0b7d1ea 100644 --- a/src/relax/op/tensor/unary.cc +++ b/src/relax/op/tensor/unary.cc @@ -26,7 +26,7 @@ #include #include -#include +#include #include diff --git a/src/relax/op/vision/multibox_transform_loc.cc b/src/relax/op/vision/multibox_transform_loc.cc index 3a80758ac33f..06ee3c8fcca0 100644 --- a/src/relax/op/vision/multibox_transform_loc.cc +++ b/src/relax/op/vision/multibox_transform_loc.cc @@ -27,7 +27,7 @@ #include #include #include -#include +#include #include #include diff --git a/src/relax/op/vision/nms.cc b/src/relax/op/vision/nms.cc index b66becf82210..290061a066c1 100644 --- a/src/relax/op/vision/nms.cc +++ b/src/relax/op/vision/nms.cc @@ -26,7 +26,7 @@ #include #include #include -#include +#include #include #include diff --git a/src/relax/op/vision/roi_align.cc b/src/relax/op/vision/roi_align.cc index 1594591f7e75..a2e3205b2d15 100644 --- a/src/relax/op/vision/roi_align.cc +++ b/src/relax/op/vision/roi_align.cc @@ -27,7 +27,7 @@ #include #include #include -#include +#include #include diff --git a/src/relax/op/vision/roi_pool.cc b/src/relax/op/vision/roi_pool.cc index 089336fe0d21..6c322b4b0eb9 100644 --- a/src/relax/op/vision/roi_pool.cc +++ b/src/relax/op/vision/roi_pool.cc @@ -27,7 +27,7 @@ #include #include #include -#include +#include #include diff --git a/src/relax/script/ir_builder/frame.cc b/src/relax/script/ir_builder/frame.cc index aaea861c1f29..d43c8e6a5466 100644 --- a/src/relax/script/ir_builder/frame.cc +++ b/src/relax/script/ir_builder/frame.cc @@ -108,8 +108,9 @@ void FunctionFrameObj::ExitWithScope() { Expr body = this->block_builder->Normalize(tvm::relax::SeqExpr(binding_blocks, output.value(), loc)); // if the function is not private, add a global symbol to its attributes - if (!is_private.value_or(false) && name.has_value() && !attrs.count(tvm::attr::kGlobalSymbol)) { - attrs.Set(tvm::attr::kGlobalSymbol, name.value()); + if (!is_private.value_or(false) && name.has_value() && + !attrs.count(tvm::attr::kGlobalSymbol.name)) { + attrs.Set(tvm::attr::kGlobalSymbol.name, name.value()); } this->block_builder->EndScope(); tvm::relax::Function func(/*params=*/params, diff --git a/src/relax/script/ir_builder/ir.cc b/src/relax/script/ir_builder/ir.cc index 106bf93cf6ca..b9847813ab02 100644 --- a/src/relax/script/ir_builder/ir.cc +++ b/src/relax/script/ir_builder/ir.cc @@ -139,9 +139,9 @@ void FuncName(const ffi::String& name) { void FuncAttrs(ffi::Map attrs) { FunctionFrame frame = FindFunctionFrame("R.func_attr"); for (const auto& [key, value] : attrs) { - if (key == tvm::attr::kGlobalSymbol && frame->is_private.value_or(false)) { + if (key == tvm::attr::kGlobalSymbol.name && frame->is_private.value_or(false)) { TVM_FFI_THROW(ValueError) << "A private function may not have the kGlobalSymbol (\"" - << tvm::attr::kGlobalSymbol << "\") attribute. " + << tvm::attr::kGlobalSymbol.name << "\") attribute. " << "However, a private function specified the global symbol as " << value; } diff --git a/src/relax/script/printer/binding.cc b/src/relax/script/printer/binding.cc index 19bb1716f8a6..cb6113d16dd0 100644 --- a/src/relax/script/printer/binding.cc +++ b/src/relax/script/printer/binding.cc @@ -18,7 +18,7 @@ */ #include #include -#include +#include #include #include @@ -54,7 +54,7 @@ ffi::Optional MatchCastDocTranslate(DocTranslatorObj* d, ffi::AnyView i TVM_FFI_STATIC_INIT_BLOCK() { ffi::reflection::TypeAttrDef().attr( - tvm::script::printer::type_attr::kDocTranslate, + tvm::script::printer::type_attr::kDocTranslate.name, FDocTranslate::FromNative<&MatchCastDocTranslate>()); } @@ -126,7 +126,7 @@ ffi::Optional VarBindingDocTranslate(DocTranslatorObj* d, ffi::AnyView TVM_FFI_STATIC_INIT_BLOCK() { ffi::reflection::TypeAttrDef().attr( - tvm::script::printer::type_attr::kDocTranslate, + tvm::script::printer::type_attr::kDocTranslate.name, FDocTranslate::FromNative<&VarBindingDocTranslate>()); } @@ -157,7 +157,8 @@ ffi::Optional IfDocTranslate(DocTranslatorObj* d, ffi::AnyView input, TVM_FFI_STATIC_INIT_BLOCK() { ffi::reflection::TypeAttrDef().attr( - tvm::script::printer::type_attr::kDocTranslate, FDocTranslate::FromNative<&IfDocTranslate>()); + tvm::script::printer::type_attr::kDocTranslate.name, + FDocTranslate::FromNative<&IfDocTranslate>()); } } // namespace diff --git a/src/relax/script/printer/dependent_type.cc b/src/relax/script/printer/dependent_type.cc index 398a26da6157..b12683dbd815 100644 --- a/src/relax/script/printer/dependent_type.cc +++ b/src/relax/script/printer/dependent_type.cc @@ -47,7 +47,7 @@ ffi::Optional ShapeTypeDocTranslate(DocTranslatorObj* d, ffi::AnyView i TVM_FFI_STATIC_INIT_BLOCK() { ffi::reflection::TypeAttrDef().attr( - tvm::script::printer::type_attr::kDocTranslate, + tvm::script::printer::type_attr::kDocTranslate.name, FDocTranslate::FromNative<&ShapeTypeDocTranslate>()); } @@ -89,7 +89,7 @@ ffi::Optional TensorTypeDocTranslate(DocTranslatorObj* d, ffi::AnyView TVM_FFI_STATIC_INIT_BLOCK() { ffi::reflection::TypeAttrDef().attr( - tvm::script::printer::type_attr::kDocTranslate, + tvm::script::printer::type_attr::kDocTranslate.name, FDocTranslate::FromNative<&TensorTypeDocTranslate>()); } @@ -125,7 +125,7 @@ ffi::Optional RelaxFuncTypeDocTranslate(DocTranslatorObj* d, ffi::AnyVi TVM_FFI_STATIC_INIT_BLOCK() { ffi::reflection::TypeAttrDef().attr( - tvm::script::printer::type_attr::kDocTranslate, + tvm::script::printer::type_attr::kDocTranslate.name, FDocTranslate::FromNative<&RelaxFuncTypeDocTranslate>()); } diff --git a/src/relax/script/printer/distributed.cc b/src/relax/script/printer/distributed.cc index bc624b4770d3..d3a3b58246eb 100644 --- a/src/relax/script/printer/distributed.cc +++ b/src/relax/script/printer/distributed.cc @@ -39,7 +39,7 @@ ffi::Optional PlacementDocTranslate(DocTranslatorObj*, ffi::AnyView inp TVM_FFI_STATIC_INIT_BLOCK() { ffi::reflection::TypeAttrDef().attr( - tvm::script::printer::type_attr::kDocTranslate, + tvm::script::printer::type_attr::kDocTranslate.name, FDocTranslate::FromNative<&PlacementDocTranslate>()); } @@ -95,7 +95,7 @@ ffi::Optional DTensorTypeDocTranslate(DocTranslatorObj* d, ffi::AnyView TVM_FFI_STATIC_INIT_BLOCK() { ffi::reflection::TypeAttrDef().attr( - tvm::script::printer::type_attr::kDocTranslate, + tvm::script::printer::type_attr::kDocTranslate.name, FDocTranslate::FromNative<&DTensorTypeDocTranslate>()); } @@ -120,7 +120,7 @@ ffi::Optional DeviceMeshDocTranslate(DocTranslatorObj* d, ffi::AnyView TVM_FFI_STATIC_INIT_BLOCK() { ffi::reflection::TypeAttrDef().attr( - tvm::script::printer::type_attr::kDocTranslate, + tvm::script::printer::type_attr::kDocTranslate.name, FDocTranslate::FromNative<&DeviceMeshDocTranslate>()); } diff --git a/src/relax/script/printer/expr.cc b/src/relax/script/printer/expr.cc index c8f3edb19fd4..dc1c5178d718 100644 --- a/src/relax/script/printer/expr.cc +++ b/src/relax/script/printer/expr.cc @@ -105,7 +105,7 @@ ffi::Optional GenericConstDocTranslate(DocTranslatorObj* d, ffi::AnyVie TVM_FFI_STATIC_INIT_BLOCK() { ffi::reflection::TypeAttrDef().attr( - tvm::script::printer::type_attr::kDocTranslate, + tvm::script::printer::type_attr::kDocTranslate.name, FDocTranslate::FromNative<&GenericConstDocTranslate>()); } @@ -120,13 +120,13 @@ ffi::Optional ShapeExprDocTranslate(DocTranslatorObj* d, ffi::AnyView i TVM_FFI_STATIC_INIT_BLOCK() { ffi::reflection::TypeAttrDef().attr( - tvm::script::printer::type_attr::kDocTranslate, + tvm::script::printer::type_attr::kDocTranslate.name, FDocTranslate::FromNative<&ShapeExprDocTranslate>()); } TVM_FFI_STATIC_INIT_BLOCK() { ffi::reflection::TypeAttrDef().attr( - tvm::script::printer::type_attr::kDocTranslate, + tvm::script::printer::type_attr::kDocTranslate.name, FDocTranslate::FromNative<&VarDocTranslate>()); } diff --git a/src/relax/script/printer/function.cc b/src/relax/script/printer/function.cc index 8cfa9055cbac..78014339e799 100644 --- a/src/relax/script/printer/function.cc +++ b/src/relax/script/printer/function.cc @@ -47,7 +47,7 @@ ffi::Optional FunctionDocTranslate(DocTranslatorObj* d, ffi::AnyView in auto bound_name = d->GetOrCreateExtraState>("ir.function_name"); d->WithExtraStateScope>("ir.function_name", std::nullopt, [&]() { auto completed = d->WithVarScope([&]() { - auto global_symbol = func->GetAttr(tvm::attr::kGlobalSymbol); + auto global_symbol = func->GetAttr(tvm::attr::kGlobalSymbol); ffi::String name = bound_name.value_or(global_symbol.value_or("main")); ffi::Array args; ffi::Array param_ids; @@ -77,7 +77,7 @@ ffi::Optional FunctionDocTranslate(DocTranslatorObj* d, ffi::AnyView in decorator_keys.push_back("pure"); decorator_values.push_back(LiteralDoc::Boolean(false, std::nullopt)); } - if (!func->GetAttr(tvm::attr::kGlobalSymbol)) { + if (!func->GetAttr(tvm::attr::kGlobalSymbol)) { decorator_keys.push_back("private"); decorator_values.push_back(LiteralDoc::Boolean(true, std::nullopt)); } @@ -87,7 +87,7 @@ ffi::Optional FunctionDocTranslate(DocTranslatorObj* d, ffi::AnyView in auto docs = d->WithDocScope([&]() { std::vector> attrs; for (const auto& [key, value] : func->attrs->dict) { - if (key != tvm::attr::kGlobalSymbol) attrs.emplace_back(key, value); + if (key != tvm::attr::kGlobalSymbol.name) attrs.emplace_back(key, value); } if (!attrs.empty()) { std::sort(attrs.begin(), attrs.end(), @@ -114,12 +114,13 @@ ffi::Optional FunctionDocTranslate(DocTranslatorObj* d, ffi::AnyView in }); auto body = ToStmtDocArray(docs); if (bound_name && global_symbol && global_symbol.value() != bound_name.value()) { - body.insert(body.begin(), - ExprStmtDoc(NamespaceDoc("relax") - ->Attr("func_attr") - ->Call({DictDoc( - {LiteralDoc::Str(tvm::attr::kGlobalSymbol, std::nullopt)}, - {LiteralDoc::Str(global_symbol.value(), std::nullopt)})}))); + body.insert( + body.begin(), + ExprStmtDoc( + NamespaceDoc("relax") + ->Attr("func_attr") + ->Call({DictDoc({LiteralDoc::Str(tvm::attr::kGlobalSymbol.name, std::nullopt)}, + {LiteralDoc::Str(global_symbol.value(), std::nullopt)})}))); } FunctionDoc function(IdDoc(name), args, {decorator}, ret_type, body); auto module = d->GetOrCreateExtraState>("ir.module"); @@ -146,7 +147,7 @@ ffi::Optional FunctionDocTranslate(DocTranslatorObj* d, ffi::AnyView in TVM_FFI_STATIC_INIT_BLOCK() { RegisterNamespaceAlias("relax.prefix", "R"); ffi::reflection::TypeAttrDef().attr( - tvm::script::printer::type_attr::kDocTranslate, + tvm::script::printer::type_attr::kDocTranslate.name, FDocTranslate::FromNative<&FunctionDocTranslate>()); } @@ -174,7 +175,7 @@ ffi::Optional ExternFuncDocTranslate(DocTranslatorObj* d, ffi::AnyView TVM_FFI_STATIC_INIT_BLOCK() { ffi::reflection::TypeAttrDef().attr( - tvm::script::printer::type_attr::kDocTranslate, + tvm::script::printer::type_attr::kDocTranslate.name, FDocTranslate::FromNative<&ExternFuncDocTranslate>()); } diff --git a/src/relax/script/printer/global_info.cc b/src/relax/script/printer/global_info.cc index 7acc8405207a..a60e729ea2e5 100644 --- a/src/relax/script/printer/global_info.cc +++ b/src/relax/script/printer/global_info.cc @@ -77,7 +77,7 @@ ffi::Optional VDeviceDocTranslate(DocTranslatorObj* d, ffi::AnyView inp TVM_FFI_STATIC_INIT_BLOCK() { ffi::reflection::TypeAttrDef().attr( - tvm::script::printer::type_attr::kDocTranslate, + tvm::script::printer::type_attr::kDocTranslate.name, FDocTranslate::FromNative<&VDeviceDocTranslate>()); } @@ -88,7 +88,7 @@ ffi::Optional DummyGlobalInfoDocTranslate(DocTranslatorObj*, ffi::AnyVi TVM_FFI_STATIC_INIT_BLOCK() { ffi::reflection::TypeAttrDef().attr( - tvm::script::printer::type_attr::kDocTranslate, + tvm::script::printer::type_attr::kDocTranslate.name, FDocTranslate::FromNative<&DummyGlobalInfoDocTranslate>()); } diff --git a/src/relax/script/printer/region.cc b/src/relax/script/printer/region.cc index 249a49208730..57e22c41016b 100644 --- a/src/relax/script/printer/region.cc +++ b/src/relax/script/printer/region.cc @@ -86,7 +86,7 @@ ffi::Optional SeqExprDocTranslate(DocTranslatorObj* d, ffi::AnyView inp TVM_FFI_STATIC_INIT_BLOCK() { ffi::reflection::TypeAttrDef().attr( - tvm::script::printer::type_attr::kDocTranslate, + tvm::script::printer::type_attr::kDocTranslate.name, FDocTranslate::FromNative<&SeqExprDocTranslate>()); } @@ -103,7 +103,7 @@ ffi::Optional BindingBlockDocTranslate(DocTranslatorObj* d, ffi::AnyVie TVM_FFI_STATIC_INIT_BLOCK() { ffi::reflection::TypeAttrDef().attr( - tvm::script::printer::type_attr::kDocTranslate, + tvm::script::printer::type_attr::kDocTranslate.name, FDocTranslate::FromNative<&BindingBlockDocTranslate>()); } @@ -133,7 +133,7 @@ ffi::Optional DataflowBlockDocTranslate(DocTranslatorObj* d, ffi::AnyVi TVM_FFI_STATIC_INIT_BLOCK() { ffi::reflection::TypeAttrDef().attr( - tvm::script::printer::type_attr::kDocTranslate, + tvm::script::printer::type_attr::kDocTranslate.name, FDocTranslate::FromNative<&DataflowBlockDocTranslate>()); } diff --git a/src/relax/training/utils.cc b/src/relax/training/utils.cc index ca083717e994..0fa7b29d4fd0 100644 --- a/src/relax/training/utils.cc +++ b/src/relax/training/utils.cc @@ -52,7 +52,7 @@ class AppendLossMutator : private ExprMutator { AppendLossMutator mutator(mod, new_loss_func, num_backbone_outputs); auto new_func_transformed = - WithAttr(mutator.VisitExpr(new_func).as_or_throw(), tvm::attr::kGlobalSymbol, + WithAttr(mutator.VisitExpr(new_func).as_or_throw(), tvm::attr::kGlobalSymbol.name, new_func_name.value_or(func_name + "_loss")); auto new_module = ffi::GetRef(mod.CopyOnWrite()); diff --git a/src/relax/transform/allocate_workspace.cc b/src/relax/transform/allocate_workspace.cc index 6c3b18629825..af3a34d50cf4 100644 --- a/src/relax/transform/allocate_workspace.cc +++ b/src/relax/transform/allocate_workspace.cc @@ -46,7 +46,7 @@ class ExternFunctionRewriter : ExprMutator { std::unordered_map Run() { std::unordered_map ret; for (const auto& [gvar, f] : builder_->GetContextIRModule()->functions) { - if (f->GetAttr(tvm::relax::attr::kWorkspaceSize)) { + if (f->GetAttr(tvm::relax::attr::kWorkspaceSize)) { ret.insert_or_assign(gvar.get(), VisitExpr(f).as_or_throw()); } } @@ -54,24 +54,24 @@ class ExternFunctionRewriter : ExprMutator { } Expr VisitExpr_(const FunctionNode* func_node) override { - if (!func_node->GetAttr(tvm::relax::attr::kCodegen) && - !func_node->GetAttr(tvm::relax::attr::kComposite)) { + if (!func_node->GetAttr(tvm::relax::attr::kCodegen) && + !func_node->GetAttr(tvm::relax::attr::kComposite)) { return ExprMutator::VisitExpr_(func_node); } - if (auto workspace = func_node->GetAttr(tvm::relax::attr::kWorkspaceSize)) { + if (auto workspace = func_node->GetAttr(tvm::relax::attr::kWorkspaceSize)) { // Append the workspace parameter to this function. ffi::Array new_params = func_node->params; auto ty = TensorType(ShapeExpr({IntImm::Int32(max_workspace_size_)}), PrimType::UInt(8)); Var workspace_param(name_sup_->FreshName("workspace"), ty); - if (func_node->GetAttr(tvm::relax::attr::kCodegen)) { + if (func_node->GetAttr(tvm::relax::attr::kCodegen)) { workspace_var_param_ = workspace_param; } new_params.push_back(workspace_param); auto new_attrs = func_node->attrs; - new_attrs.CopyOnWrite()->dict.erase(tvm::relax::attr::kWorkspaceSize); + new_attrs.CopyOnWrite()->dict.erase(tvm::relax::attr::kWorkspaceSize.name); return Function(new_params, VisitExpr(func_node->body), func_node->ret_ty, func_node->is_pure, new_attrs); } @@ -84,7 +84,7 @@ class ExternFunctionRewriter : ExprMutator { if (auto callee = builder_->LookupBinding(var.value()); callee && callee.value()->IsInstance() && callee.value().as_or_throw()->GetAttr( - tvm::relax::attr::kComposite)) { + tvm::relax::attr::kComposite.name)) { // Append the workspace argument to this call. The callee should have been updated to accept // a workspace as the last parameter. auto new_args = call_node->args; @@ -112,7 +112,7 @@ class WorkspaceProvider : ExprMutator { IRModule Run() { for (const auto& [gvar, f] : mod_->functions) { - if (auto workspace = f->GetAttr(tvm::relax::attr::kWorkspaceSize)) { + if (auto workspace = f->GetAttr(tvm::relax::attr::kWorkspaceSize)) { max_workspace_size_ = std::max(max_workspace_size_, workspace.value()); } } @@ -128,7 +128,7 @@ class WorkspaceProvider : ExprMutator { // This is only required since the well-formed check requires kGlobalSymbol to be the same // as the actual name of the global variable. builder_->UpdateFunction(new_gvar, - WithAttr(f, tvm::attr::kGlobalSymbol, new_gvar->name_hint)); + WithAttr(f, tvm::attr::kGlobalSymbol.name, new_gvar->name_hint)); gvar_map_.insert_or_assign(gvar, new_gvar); new_gvars_.insert(new_gvar); builder_->GetContextIRModule()->Remove(ffi::GetRef(gvar)); @@ -136,9 +136,8 @@ class WorkspaceProvider : ExprMutator { for (const auto& [gvar, f] : mod_->functions) { workspace_var_main_ = std::nullopt; - if (!f->IsInstance() || - f->GetAttr(tvm::relax::attr::kCodegen) || - f->GetAttr(tvm::relax::attr::kComposite)) { + if (!f->IsInstance() || f->GetAttr(tvm::relax::attr::kCodegen) || + f->GetAttr(tvm::relax::attr::kComposite)) { continue; } auto func = mod_->Lookup(gvar).as_or_throw(); diff --git a/src/relax/transform/alter_op_impl.cc b/src/relax/transform/alter_op_impl.cc index 23d92ebf2264..39a8ebf4c19a 100644 --- a/src/relax/transform/alter_op_impl.cc +++ b/src/relax/transform/alter_op_impl.cc @@ -119,7 +119,7 @@ class AlterOpImplMutator : public ExprMutator { const tirx::Function& old_func = mod_->Lookup(call->args[0].as_or_throw()).as_or_throw(); ffi::Optional maybe_op_kind = - old_func->attrs.GetAttr(tvm::relax::attr::kOperatorName); + old_func->attrs.GetAttr(tvm::relax::attr::kOperatorName); // If the callee does not have kOperatorName attribute or no replacement is requested for // it, nothing to do here. @@ -221,12 +221,12 @@ class AlterOpImplMutator : public ExprMutator { // Create tirx::Function and add op_name to func.attrs tirx::Function remove_pad_with_frozen_layout = WithAttr(CreateFunction({placeholder_tensor, output_tensor}), - tvm::relax::attr::kOperatorName, op_name); + tvm::relax::attr::kOperatorName.name, op_name); // Add tirx::Function to module GlobalVar gv_remove_pad = builder_->AddFunction(remove_pad_with_frozen_layout, op_name); // Mark the remove_pad tirx::Function as private by removing it from global scope - builder_->UpdateFunction(gv_remove_pad, - WithoutAttr(remove_pad_with_frozen_layout, tvm::attr::kGlobalSymbol)); + builder_->UpdateFunction( + gv_remove_pad, WithoutAttr(remove_pad_with_frozen_layout, tvm::attr::kGlobalSymbol.name)); remove_pad_map_.insert_or_assign(t_shape, gv_remove_pad); return gv_remove_pad; @@ -268,7 +268,7 @@ class AlterOpImplMutator : public ExprMutator { // future passes that use kOperatorName attribute to identify operator represented by a // tirx::Function. tirx::Function replacement_func_with_frozen_layout = - WithAttr(replacement_func, tvm::relax::attr::kOperatorName, op_kind); + WithAttr(replacement_func, tvm::relax::attr::kOperatorName.name, op_kind); GlobalVar gv_replacement = builder_->AddFunction(replacement_func_with_frozen_layout, op_kind + "_replacement"); diff --git a/src/relax/transform/annotate_tir_op_pattern.cc b/src/relax/transform/annotate_tir_op_pattern.cc index 442536d667a2..56d6deaf7180 100644 --- a/src/relax/transform/annotate_tir_op_pattern.cc +++ b/src/relax/transform/annotate_tir_op_pattern.cc @@ -32,11 +32,11 @@ namespace tvm { namespace relax { tirx::Function AnnotateOpPattern(tirx::Function f) { - if (f->HasNonzeroAttr(tvm::relax::attr::kOpPattern)) { + if (f->HasNonzeroAttr(tvm::relax::attr::kOpPattern.name)) { return f; } else { OpPatternKind kind = AnalyzeOpPatternKind(f); - return WithAttr(std::move(f), tvm::relax::attr::kOpPattern, static_cast(kind)); + return WithAttr(std::move(f), tvm::relax::attr::kOpPattern.name, static_cast(kind)); } } diff --git a/src/relax/transform/attach_attr_layout_free_tensors.cc b/src/relax/transform/attach_attr_layout_free_tensors.cc index d7e51a566d54..38750cc01a51 100644 --- a/src/relax/transform/attach_attr_layout_free_tensors.cc +++ b/src/relax/transform/attach_attr_layout_free_tensors.cc @@ -51,7 +51,7 @@ class AttrAttacher : public ExprMutator { using ExprMutator::VisitExpr_; Expr VisitExpr_(const FunctionNode* op) final { - if (auto opt_num_input = op->attrs.GetAttr(tvm::relax::attr::kNumInput)) { + if (auto opt_num_input = op->attrs.GetAttr(tvm::relax::attr::kNumInput)) { TVM_FFI_ICHECK(layout_free_exprs_.empty()) << "meet a non-global function with num_input attr"; size_t num_input = opt_num_input.value(); @@ -84,7 +84,7 @@ class AttrAttacher : public ExprMutator { } // Attach the layout free tensors to the tirx::Function tirx::Function func = WithAttr(mod_->Lookup(gv).as_or_throw(), - tvm::s_tir::attr::kLayoutFreeTensors, layout_free_tensors); + tvm::s_tir::attr::kLayoutFreeTensors.name, layout_free_tensors); // Renew defs func = tirx::RenewDef(func); // Add the updated tirx::Function in the IRModule diff --git a/src/relax/transform/attach_global_symbol.cc b/src/relax/transform/attach_global_symbol.cc index b789b8b2ce3b..09d0bfdc4e74 100644 --- a/src/relax/transform/attach_global_symbol.cc +++ b/src/relax/transform/attach_global_symbol.cc @@ -96,9 +96,9 @@ IRModule ReplaceGlobalVarsInModule(IRModule mod, ffi::Map func.CopyOnWrite()->body = new_body; } // Update kGlobalSymbol if the function is externally exposed and being renamed. - if (func->GetAttr(tvm::attr::kGlobalSymbol)) { + if (func->GetAttr(tvm::attr::kGlobalSymbol)) { if (new_gvar->name_hint != old_gvar->name_hint) { - func = WithAttr(func, tvm::attr::kGlobalSymbol, new_gvar->name_hint); + func = WithAttr(func, tvm::attr::kGlobalSymbol.name, new_gvar->name_hint); } } new_func = func; @@ -106,9 +106,10 @@ IRModule ReplaceGlobalVarsInModule(IRModule mod, ffi::Map RelaxGvarMutator mutator(replacements); auto new_relax_func = mutator(ffi::GetRef(relax_func_node)).as_or_throw(); // Update kGlobalSymbol if the function is externally exposed and being renamed. - if (new_relax_func->GetAttr(tvm::attr::kGlobalSymbol)) { + if (new_relax_func->GetAttr(tvm::attr::kGlobalSymbol)) { if (new_gvar->name_hint != old_gvar->name_hint) { - new_relax_func = WithAttr(new_relax_func, tvm::attr::kGlobalSymbol, new_gvar->name_hint); + new_relax_func = + WithAttr(new_relax_func, tvm::attr::kGlobalSymbol.name, new_gvar->name_hint); } } new_func = new_relax_func; @@ -141,12 +142,12 @@ IRModule ReplaceGlobalVarsInModule(IRModule mod, ffi::Map Pass AttachGlobalSymbol() { auto pass_func = [=](IRModule mod, PassContext pc) { - ffi::String c_prefix = mod->GetAttr(tvm::attr::kSystemLibPrefix).value_or(""); + ffi::String c_prefix = mod->GetAttr(tvm::attr::kSystemLibPrefix).value_or(""); IRModule updates; ffi::Map gvar_updates; for (const auto& [gvar, func] : mod->functions) { - ffi::Optional old_name = func->GetAttr(tvm::attr::kGlobalSymbol); + ffi::Optional old_name = func->GetAttr(tvm::attr::kGlobalSymbol); // TODO(tvm-team): re-enable once fix relax integration part // if (old_name) continue; @@ -156,11 +157,12 @@ Pass AttachGlobalSymbol() { if (auto* function = func.as()) { new_name = c_prefix + gvar->name_hint; - new_func = - WithAttr(ffi::GetRef(function), tvm::attr::kGlobalSymbol, new_name); + new_func = WithAttr(ffi::GetRef(function), tvm::attr::kGlobalSymbol.name, + new_name); } else if (auto* relax_func = func.as()) { new_name = gvar->name_hint; - new_func = WithAttr(ffi::GetRef(relax_func), tvm::attr::kGlobalSymbol, new_name); + new_func = + WithAttr(ffi::GetRef(relax_func), tvm::attr::kGlobalSymbol.name, new_name); } if (new_name.has_value() && (!old_name.has_value() || old_name.value() != new_name.value())) { diff --git a/src/relax/transform/bind_params.cc b/src/relax/transform/bind_params.cc index 63294a1bf292..e9fbab70bd14 100644 --- a/src/relax/transform/bind_params.cc +++ b/src/relax/transform/bind_params.cc @@ -121,8 +121,7 @@ IRModule BindParam(IRModule m, ffi::String func_name, ffi::Map()) { if (relax_f->GetLinkageType() == LinkageType::kExternal) { // Use global_symbol if it's external linkage - ffi::Optional gsymbol = - relax_f->GetAttr(tvm::attr::kGlobalSymbol); + ffi::Optional gsymbol = relax_f->GetAttr(tvm::attr::kGlobalSymbol); if (gsymbol.has_value() && gsymbol.value() == func_name) { Function f_after_bind = FunctionBindParams(ffi::GetRef(relax_f), bind_params); new_module->Update(func_pr.first, f_after_bind); diff --git a/src/relax/transform/bundle_model_params.cc b/src/relax/transform/bundle_model_params.cc index 6e4d5754aa53..5222ffd8486d 100644 --- a/src/relax/transform/bundle_model_params.cc +++ b/src/relax/transform/bundle_model_params.cc @@ -44,7 +44,7 @@ class ModelParamBundler : public ExprMutator { Expr VisitExpr_(const FunctionNode* op) override { Function func = ffi::GetRef(op); - auto opt_num_input = func->attrs.GetAttr(tvm::relax::attr::kNumInput); + auto opt_num_input = func->attrs.GetAttr(tvm::relax::attr::kNumInput); if (!opt_num_input) return func; auto signed_num_input = opt_num_input.value(); diff --git a/src/relax/transform/compute_prim_value.cc b/src/relax/transform/compute_prim_value.cc index 87146a1529e4..18b4d5733c42 100644 --- a/src/relax/transform/compute_prim_value.cc +++ b/src/relax/transform/compute_prim_value.cc @@ -23,7 +23,7 @@ #include #include #include -#include +#include #include #include #include @@ -117,9 +117,9 @@ class PrimExprComputeInjector : public ExprMutator { tvm::Stmt body = tvm::Return(node); tirx::Function func(param_vars, tvm::SeqStmt(body), ret_ty, - DictAttrs({{tvm::tirx::attr::kIsHostFunc, true}, - {tvm::attr::kSTir, true}, - {tvm::attr::kScriptNamespace, ffi::String("s_tir")}})); + DictAttrs({{tvm::tirx::attr::kIsHostFunc.name, true}, + {tvm::attr::kSTir.name, true}, + {tvm::attr::kScriptNamespace.name, ffi::String("s_tir")}})); func = tirx::RenewDef(func); auto callee = builder_->AddFunction(func, "compute_symbolic_expr"); diff --git a/src/relax/transform/convert_layout.cc b/src/relax/transform/convert_layout.cc index 697bb00dce51..be920b838a12 100644 --- a/src/relax/transform/convert_layout.cc +++ b/src/relax/transform/convert_layout.cc @@ -26,7 +26,7 @@ #include #include #include -#include +#include #include #include diff --git a/src/relax/transform/dataflow_inplace.cc b/src/relax/transform/dataflow_inplace.cc index f735a0a49ca5..ce3e7743c55e 100644 --- a/src/relax/transform/dataflow_inplace.cc +++ b/src/relax/transform/dataflow_inplace.cc @@ -32,7 +32,7 @@ #include #include #include -#include +#include #include #include #include diff --git a/src/relax/transform/decompose_ops.cc b/src/relax/transform/decompose_ops.cc index 054b12595f9c..ef720fdc7f2d 100644 --- a/src/relax/transform/decompose_ops.cc +++ b/src/relax/transform/decompose_ops.cc @@ -235,11 +235,11 @@ Pass ApplyDecomposeToFunction(Pass pass, ffi::String func_name) { for (auto [gvar, func] : mod->functions) { if (gvar->name_hint == func_name) { - if (!func->GetAttr(tvm::attr::kGlobalSymbol).has_value()) { + if (!func->GetAttr(tvm::attr::kGlobalSymbol).has_value()) { // Mark internal functions as externally-exposed so that // call-tracing transforms inside the pass do not remove them. internal_functions.insert(gvar->name_hint); - func = WithAttr(std::move(func), tvm::attr::kGlobalSymbol, gvar->name_hint); + func = WithAttr(std::move(func), tvm::attr::kGlobalSymbol.name, gvar->name_hint); } } else { // Replace non-target functions with stubs to keep references intact. @@ -263,7 +263,7 @@ Pass ApplyDecomposeToFunction(Pass pass, ffi::String func_name) { write_ptr->Remove((*it).second); } if (internal_functions.count(gvar->name_hint)) { - func = WithoutAttr(std::move(func), tvm::attr::kGlobalSymbol); + func = WithoutAttr(std::move(func), tvm::attr::kGlobalSymbol.name); } write_ptr->Add(gvar, func); } diff --git a/src/relax/transform/eliminate_common_subexpr.cc b/src/relax/transform/eliminate_common_subexpr.cc index ba89bfed3c52..1a0e603df6ce 100644 --- a/src/relax/transform/eliminate_common_subexpr.cc +++ b/src/relax/transform/eliminate_common_subexpr.cc @@ -31,7 +31,7 @@ #include #include #include -#include +#include #include #include #include diff --git a/src/relax/transform/expand_tuple_arguments.cc b/src/relax/transform/expand_tuple_arguments.cc index 772dbcb8b5c2..d25fa4bb88df 100644 --- a/src/relax/transform/expand_tuple_arguments.cc +++ b/src/relax/transform/expand_tuple_arguments.cc @@ -34,7 +34,7 @@ template using PMap = std::unordered_map; ffi::Optional ExpandParams(Function func) { - bool is_exposed = func->attrs.GetAttr(tvm::attr::kGlobalSymbol).has_value(); + bool is_exposed = func->attrs.GetAttr(tvm::attr::kGlobalSymbol).has_value(); if (is_exposed) return std::nullopt; bool has_tuple_param = diff --git a/src/relax/transform/fold_constant.cc b/src/relax/transform/fold_constant.cc index 4d36d0a1f4d8..3fb85faa7c9d 100644 --- a/src/relax/transform/fold_constant.cc +++ b/src/relax/transform/fold_constant.cc @@ -25,7 +25,7 @@ #include #include #include -#include +#include #include #include #include @@ -123,7 +123,7 @@ class ConstantFolder : public ExprMutator { // now // TODO(Hongyi): further check and narrow the scope of foldable function const auto pf = tvm::ffi::Function::GetGlobalRequired("tirx.compile"); - func = WithAttr(func, tvm::attr::kGlobalSymbol, ffi::String("tir_function")); + func = WithAttr(func, tvm::attr::kGlobalSymbol.name, ffi::String("tir_function")); ffi::Module rt_module = pf(func, eval_cpu_target).cast(); build_func = rt_module->GetFunction("tir_function"); } catch (const tvm::ffi::Error& err) { diff --git a/src/relax/transform/fuse_ops.cc b/src/relax/transform/fuse_ops.cc index 9195145beaf6..a681756227e5 100644 --- a/src/relax/transform/fuse_ops.cc +++ b/src/relax/transform/fuse_ops.cc @@ -124,16 +124,16 @@ class GraphCreator : public ExprVisitor { static IndexedForwardGraph Create(IRModule mod, support::Arena* arena) { GraphCreator creator(mod, arena); for (const auto& it : mod->functions) { - // Only visit Relax functions with neither attr::kPrimitive nor - // attr::kCodegen. Relax functions with `attr::kPrimitive` are + // Only visit Relax functions with neither attr::kPrimitive.name nor + // attr::kCodegen.name. Relax functions with `attr::kPrimitive.name` are // previously fused functions, potentially from a previous use // of `FuseOps` or `FuseOpsByPattern`. Relax functions with - // `attr::kCodegen` are previously fused functions from + // `attr::kCodegen.name` are previously fused functions from // `FuseOpsByPattern`, when the `annotate_codegen` option is // true. const auto* func = it.second.as(); - if (func == nullptr || func->HasNonzeroAttr(tvm::relax::attr::kPrimitive) || - func->GetAttr(tvm::relax::attr::kCodegen).has_value()) { + if (func == nullptr || func->HasNonzeroAttr(tvm::relax::attr::kPrimitive.name) || + func->GetAttr(tvm::relax::attr::kCodegen).has_value()) { continue; } creator(ffi::GetRef(func)); @@ -162,7 +162,7 @@ class GraphCreator : public ExprVisitor { SetNodePattern(param_node, OpPatternKind::kOpaque); AddToPostDFSOrder(param_node, param.get()); } - if (auto opt_num_input = func->GetAttr(tvm::relax::attr::kNumInput)) { + if (auto opt_num_input = func->GetAttr(tvm::relax::attr::kNumInput)) { for (int i = static_cast(opt_num_input.value()); i < static_cast(func->params.size()); ++i) { input_params_.insert(func->params[i].get()); @@ -222,7 +222,7 @@ class GraphCreator : public ExprVisitor { // Override args for call_tir args = call->args[1].as_or_throw()->fields; - ffi::Optional opt_pattern = func->GetAttr(tvm::relax::attr::kOpPattern); + ffi::Optional opt_pattern = func->GetAttr(tvm::relax::attr::kOpPattern); if (opt_pattern.has_value()) { pattern = static_cast(opt_pattern.value()); } else { @@ -582,7 +582,7 @@ class FunctionCreator : public ExprMutator { Expr body = outputs.size() == 1 ? outputs[0] : Tuple(outputs); body = builder_->Normalize(body); body = builder_->Normalize(SeqExpr({new_block}, body)); - group_attrs.Set(tvm::relax::attr::kPrimitive, true); + group_attrs.Set(tvm::relax::attr::kPrimitive.name, true); Function function = Function(/*params=*/params_, // /*body=*/body, // /*ret_ty=*/std::nullopt, // @@ -847,11 +847,11 @@ class OperatorFusor : public ExprMutator { } for (const auto& gv : entry_functions) { const auto& func = mod_->Lookup(gv); - // Only visit Relax functions with neither attr::kPrimitive nor - // attr::kCodegen. + // Only visit Relax functions with neither attr::kPrimitive.name nor + // attr::kCodegen.name. if (func->IsInstance() && - !func->HasNonzeroAttr(tvm::relax::attr::kPrimitive) && - !func->GetAttr(tvm::relax::attr::kCodegen).has_value()) { + !func->HasNonzeroAttr(tvm::relax::attr::kPrimitive.name) && + !func->GetAttr(tvm::relax::attr::kCodegen).has_value()) { outer_bindings_ = AnalyzeVar2Value(func); auto updated_func = VisitExpr(func).as_or_throw(); builder_->UpdateFunction(gv, updated_func); @@ -1266,7 +1266,7 @@ class PatternBasedPartitioner : ExprVisitor { // parent_group corresponds to the group of "conv1" above. auto parent_group = GetGroupForBoundVar(binding->var); TVM_FFI_ICHECK(parent_group); - parent_group->attrs.Set(tvm::relax::attr::kComposite, pat_name_); + parent_group->attrs.Set(tvm::relax::attr::kComposite.name, pat_name_); if (attrs_getter_ != nullptr) { const auto& custom_attrs = attrs_getter_(context->annotated_expr); for (const auto& pair : custom_attrs) { @@ -1378,8 +1378,8 @@ class CompositeFunctionAnnotator : public ExprMutator { } const auto& base_func = (*it).second; if (const auto* func = base_func.as()) { - if (func->GetAttr(tvm::relax::attr::kComposite).has_value() || - func->GetAttr(tvm::relax::attr::kCodegen).has_value()) { + if (func->GetAttr(tvm::relax::attr::kComposite).has_value() || + func->GetAttr(tvm::relax::attr::kCodegen).has_value()) { continue; } @@ -1400,13 +1400,13 @@ class CompositeFunctionAnnotator : public ExprMutator { return Call(Type::Missing(), it->second, call_node->args); } auto func = builder_->GetContextIRModule()->Lookup(ffi::GetRef(gvar)); - if (auto composite_name = func->GetAttr(tvm::relax::attr::kComposite)) { + if (auto composite_name = func->GetAttr(tvm::relax::attr::kComposite)) { auto new_func = VisitExpr(func).as_or_throw(); auto codegen_name = GetCodegenName(composite_name.value()); auto gsymbol = gvar->name_hint + "_" + codegen_name; - new_func = WithAttrs(new_func, {{tvm::relax::attr::kCodegen, codegen_name}, - {tvm::attr::kGlobalSymbol, gsymbol}}); - new_func = WithoutAttr(std::move(new_func), tvm::relax::attr::kPrimitive); + new_func = WithAttrs(new_func, {{tvm::relax::attr::kCodegen.name, codegen_name}, + {tvm::attr::kGlobalSymbol.name, gsymbol}}); + new_func = WithoutAttr(std::move(new_func), tvm::relax::attr::kPrimitive.name); builder_->GetContextIRModule()->Remove(ffi::GetRef(gvar)); auto new_gvar = builder_->AddFunction(new_func, gsymbol); gvar_map_.insert_or_assign(gvar, new_gvar); @@ -1419,13 +1419,13 @@ class CompositeFunctionAnnotator : public ExprMutator { Expr VisitExpr_(const FunctionNode* func_node) final { Function f_inner = ExprMutator::VisitExpr_(func_node).as_or_throw(); - if (!func_node->GetAttr(tvm::relax::attr::kComposite)) { - // This lambda function doesn't have `attr::kComposite`, so it + if (!func_node->GetAttr(tvm::relax::attr::kComposite)) { + // This lambda function doesn't have `attr::kComposite.name`, so it // was not produced by FuseOps. return f_inner; } - f_inner = WithoutAttr(std::move(f_inner), tvm::relax::attr::kPrimitive); + f_inner = WithoutAttr(std::move(f_inner), tvm::relax::attr::kPrimitive.name); ffi::Array param_vars; ffi::Array params; @@ -1478,9 +1478,9 @@ IRModule FuseOpsByPattern(const tvm::ffi::Array& patte continue; } const FunctionNode* function = base_func.as(); - if (function->GetAttr(tvm::relax::attr::kPrimitive).value_or(false) || - function->GetAttr(tvm::relax::attr::kComposite).has_value() || - function->GetAttr(tvm::relax::attr::kCodegen).has_value()) { + if (function->GetAttr(tvm::relax::attr::kPrimitive).value_or(false) || + function->GetAttr(tvm::relax::attr::kComposite).has_value() || + function->GetAttr(tvm::relax::attr::kCodegen).has_value()) { continue; } entry_functions.push_back(base_func.as_or_throw()); diff --git a/src/relax/transform/fuse_tir.cc b/src/relax/transform/fuse_tir.cc index 775b2cc59bf4..06d349590f07 100644 --- a/src/relax/transform/fuse_tir.cc +++ b/src/relax/transform/fuse_tir.cc @@ -436,7 +436,7 @@ class FusedTIRConstructor : public ExprVisitor { BaseFunc f = mod->Lookup(gv); TVM_FFI_ICHECK(f->IsInstance()) << "Expected relax functions, but got: " << f->GetTypeKey(); - TVM_FFI_ICHECK(f->HasNonzeroAttr(tvm::relax::attr::kPrimitive)) + TVM_FFI_ICHECK(f->HasNonzeroAttr(tvm::relax::attr::kPrimitive.name)) << "Expected a function with attr `kPrimitive`"; visitor(f.as_or_throw()); ffi::Array inplace_indices; @@ -865,9 +865,9 @@ class FusedTIRConstructor : public ExprVisitor { */ tirx::Function ConstructFunc() { ffi::Map attr_map; - attr_map.Set(tvm::tirx::attr::kNoAlias, true); - attr_map.Set(tvm::attr::kSTir, true); - attr_map.Set(tvm::attr::kScriptNamespace, ffi::String("s_tir")); + attr_map.Set(tvm::tirx::attr::kNoAlias.name, true); + attr_map.Set(tvm::attr::kSTir.name, true); + attr_map.Set(tvm::attr::kScriptNamespace.name, ffi::String("s_tir")); auto subst = ffi::make_object(func_info_.tensor_subst_map, func_info_.symbolic_var_remap); TVM_FFI_ICHECK(func_info_.global_name != "fused"); @@ -997,7 +997,7 @@ class TIRFuseMutator : public ExprMutator { for (const auto& gvar : mod->GetGlobalVars()) { const auto& base_func = mod->Lookup(gvar); // Only fuse primitive relax functions - if (base_func->HasNonzeroAttr(tvm::relax::attr::kPrimitive)) { + if (base_func->HasNonzeroAttr(tvm::relax::attr::kPrimitive.name)) { if (auto func = base_func.as()) { primitive_relax.Set(gvar, func.value()); } @@ -1034,7 +1034,7 @@ class TIRFuseMutator : public ExprMutator { for (const auto& [gv, func] : mod->functions) { if (func->IsInstance()) { - TVM_FFI_ICHECK(!func->HasNonzeroAttr(tvm::relax::attr::kPrimitive)) + TVM_FFI_ICHECK(!func->HasNonzeroAttr(tvm::relax::attr::kPrimitive.name)) << "Module should not contain any primitive relax functions at this point"; relax::Function update_func = mutator.VisitExpr(func).as_or_throw(); if (!update_func.same_as(func)) { diff --git a/src/relax/transform/gradient.cc b/src/relax/transform/gradient.cc index 3021dc9ea962..23f310d61ece 100644 --- a/src/relax/transform/gradient.cc +++ b/src/relax/transform/gradient.cc @@ -33,7 +33,7 @@ #include #include #include -#include +#include #include #include @@ -674,7 +674,7 @@ class GradientMutator : private ExprMutator { // Step 4.4 mark the transformed function as public // because the original function may be public, and have gsymbol attribute as func_name auto new_func_name = func_name + "_adjoint"; - auto new_func_with_gsymbol = WithAttr(new_func, tvm::attr::kGlobalSymbol, new_func_name); + auto new_func_with_gsymbol = WithAttr(new_func, tvm::attr::kGlobalSymbol.name, new_func_name); // Step 4.5 Add the transformed function to IRModule builder_->AddFunction(new_func_with_gsymbol, new_func_name); diff --git a/src/relax/transform/infer_amp_utils.h b/src/relax/transform/infer_amp_utils.h index 38604786b464..69ca165333f8 100644 --- a/src/relax/transform/infer_amp_utils.h +++ b/src/relax/transform/infer_amp_utils.h @@ -28,7 +28,7 @@ #include #include #include -#include +#include #include #include diff --git a/src/relax/transform/infer_layout_utils.h b/src/relax/transform/infer_layout_utils.h index 2a84271f04bb..176ee2e5a8a7 100644 --- a/src/relax/transform/infer_layout_utils.h +++ b/src/relax/transform/infer_layout_utils.h @@ -37,7 +37,7 @@ #include #include #include -#include +#include #include #include diff --git a/src/relax/transform/inline_functions.cc b/src/relax/transform/inline_functions.cc index f2cbb10f87cf..f5203c658fb2 100644 --- a/src/relax/transform/inline_functions.cc +++ b/src/relax/transform/inline_functions.cc @@ -180,7 +180,7 @@ Pass InlinePrivateFunctions() { for (const auto& [gvar, base_func] : mod->functions) { if (auto opt = base_func.as()) { auto func = opt.value(); - bool is_private = !func->GetAttr(tvm::attr::kGlobalSymbol).has_value(); + bool is_private = !func->GetAttr(tvm::attr::kGlobalSymbol).has_value(); if (is_private) { replacements.Set(gvar, func); } diff --git a/src/relax/transform/lambda_lift.cc b/src/relax/transform/lambda_lift.cc index d432589f34d3..75debf420147 100644 --- a/src/relax/transform/lambda_lift.cc +++ b/src/relax/transform/lambda_lift.cc @@ -29,7 +29,7 @@ #include #include #include -#include +#include #include #include @@ -63,7 +63,7 @@ class LambdaNameCollector : ExprVisitor { private: void VisitBinding_(const VarBindingNode* binding, const FunctionNode* func) override { - if (auto opt = func->GetAttr(tvm::attr::kGlobalSymbol)) { + if (auto opt = func->GetAttr(tvm::attr::kGlobalSymbol)) { ffi::String public_name = opt.value(); // If a kGlobalSymbol exists, we must use the name exactly as it @@ -75,7 +75,7 @@ class LambdaNameCollector : ExprVisitor { // the error. TVM_FFI_ICHECK(!previous_global_vars_.count(public_name)) << "Function " << name_stack_.front() << " contains a lambda with kGlobalSymbol (\"" - << tvm::attr::kGlobalSymbol << "\" attribute of \"" << public_name << "\". " + << tvm::attr::kGlobalSymbol.name << "\" attribute of \"" << public_name << "\". " << "However, the module already contains a GlobalVar with this name. " << "If present, the kGlobalSymbol attribute must match the name of the GlobalVar, " << "and GlobalVar names must be unique across an IRModule. " @@ -85,7 +85,7 @@ class LambdaNameCollector : ExprVisitor { auto it = new_public_names_.find(public_name); TVM_FFI_ICHECK(it == new_public_names_.end()) << "Function " << name_stack_.front() << " contains a lambda with kGlobalSymbol (\"" - << tvm::attr::kGlobalSymbol << "\" attribute of \"" << public_name << "\". " + << tvm::attr::kGlobalSymbol.name << "\" attribute of \"" << public_name << "\". " << "However, the function " << it->second.front() << " also contains a lambda with the same value for kGlobalSymbol. " << "If present, the kGlobalSymbol attribute must match the name of the GlobalVar, " diff --git a/src/relax/transform/lazy_transform_params.cc b/src/relax/transform/lazy_transform_params.cc index 38713a39ad69..bb51c7438dc8 100644 --- a/src/relax/transform/lazy_transform_params.cc +++ b/src/relax/transform/lazy_transform_params.cc @@ -38,13 +38,13 @@ using namespace tvm::prim; namespace { std::optional GetNumInputParams(const FunctionNode* func) { - if (auto opt_int_imm = func->GetAttr(tvm::relax::attr::kNumInput)) { + if (auto opt_int_imm = func->GetAttr(tvm::relax::attr::kNumInput.name)) { int64_t num_input_params = static_cast(opt_int_imm.value()->value); TVM_FFI_CHECK_GE(num_input_params, 0, ValueError) - << "Annotation for tvm::relax::attr::kNumInput (\"" << tvm::relax::attr::kNumInput + << "Annotation for tvm::relax::attr::kNumInput.name (\"" << tvm::relax::attr::kNumInput.name << "\") must be non-negative, but was " << num_input_params; TVM_FFI_CHECK_LE(static_cast(num_input_params), func->params.size(), ValueError) - << "Annotation for tvm::relax::attr::kNumInput (\"" << tvm::relax::attr::kNumInput + << "Annotation for tvm::relax::attr::kNumInput.name (\"" << tvm::relax::attr::kNumInput.name << "\") specifies " << num_input_params << " parameters to be provided at runtime, " << "but the function only accepts " << func->params.size() << " parameters in total"; return num_input_params; @@ -86,7 +86,7 @@ class LazyInputMutator : public ExprMutator { auto node = ffi::GetRef(func); node.CopyOnWrite()->params = new_params; node.CopyOnWrite()->ret_ty = new_ret_ty; - node = WithAttr(node, tvm::relax::attr::kNumInput, num_input_params + 1); + node = WithAttr(node, tvm::relax::attr::kNumInput.name, num_input_params + 1); plan_ = FunctionPlan{std::move(param_lookup), fget_param}; auto output = ExprMutator::VisitExpr_(node.get()).as_or_throw(); @@ -189,7 +189,7 @@ class LazyOutputMutator : public ExprMutator { write_ptr->is_pure = false; } if (num_input_params.has_value()) { - node = WithAttr(node, tvm::relax::attr::kNumInput, num_input_params.value() + 1); + node = WithAttr(node, tvm::relax::attr::kNumInput.name, num_input_params.value() + 1); } auto output = ExprMutator::VisitExpr_(node.get()).as_or_throw(); @@ -245,7 +245,7 @@ namespace transform { Pass LazyGetInput() { auto pass_func = [](Function func, IRModule, PassContext) -> Function { - if (!func->GetAttr(tvm::attr::kGlobalSymbol).has_value()) { + if (!func->GetAttr(tvm::attr::kGlobalSymbol).has_value()) { return func; } return WithLazyInputs(func); @@ -262,7 +262,7 @@ TVM_FFI_STATIC_INIT_BLOCK() { Pass LazySetOutput() { auto pass_func = [](Function func, IRModule, PassContext) -> Function { - if (!func->GetAttr(tvm::attr::kGlobalSymbol).has_value()) { + if (!func->GetAttr(tvm::attr::kGlobalSymbol).has_value()) { return func; } return WithLazyOutputs(func); diff --git a/src/relax/transform/legalize_ops.cc b/src/relax/transform/legalize_ops.cc index 5a9349ccc2f1..014f77e6c6fe 100644 --- a/src/relax/transform/legalize_ops.cc +++ b/src/relax/transform/legalize_ops.cc @@ -28,7 +28,7 @@ #include #include #include -#include +#include #include #include #include @@ -212,7 +212,7 @@ class LegalizeMutator : public ExprMutator { } auto function = opt_function.value(); - auto func_target = function->GetAttr(tvm::attr::kTarget); + auto func_target = function->GetAttr(tvm::attr::kTarget); if (func_target && func_target.value()->kind == vdevice_target.value()->kind) { // The function already has compatible annotations for the // target, so no modifications are required. @@ -225,7 +225,7 @@ class LegalizeMutator : public ExprMutator { // `tirx::Function` with the appropriate target annotation. In the // future, this may be treated as a bug in the FLegalize // implementation, rather than expected output from it. - auto new_function = WithAttr(function, tvm::attr::kTarget, vdevice_target.value()); + auto new_function = WithAttr(function, tvm::attr::kTarget.name, vdevice_target.value()); auto new_gvar_name = [&]() -> std::string { std::stringstream ss; ss << gvar.value()->name_hint; diff --git a/src/relax/transform/lift_transform_params.cc b/src/relax/transform/lift_transform_params.cc index dc1f855bd0e3..cab95099c072 100644 --- a/src/relax/transform/lift_transform_params.cc +++ b/src/relax/transform/lift_transform_params.cc @@ -57,7 +57,7 @@ struct BaseCollectInfo { * - All bindings in `computable_at_compile_time` are suitable for * use in a DataflowBlock. * - * - Do not depend on any parameter prior to attr::kNumInput. + * - Do not depend on any parameter prior to attr::kNumInput.name. * * - Does not include "relax.builtin.stop_lift_params" */ @@ -123,7 +123,7 @@ struct BaseCollectInfo { }, tuple_var); Function func(params, body, GetType(tuple_var)); - func = WithAttr(func, tvm::relax::attr::kNumInput, 0); + func = WithAttr(func, tvm::relax::attr::kNumInput.name, 0); func = CopyWithNewVars(func); func = BundleModelParams(func); func = CanonicalizeBindings(func).as_or_throw(); @@ -414,7 +414,7 @@ class LocalLiftableBindingCollector : public BaseLiftableBindingCollector { } void VisitExpr_(const FunctionNode* func) override { size_t num_runtime_params = func->params.size(); - if (auto opt = func->attrs.GetAttr(tvm::relax::attr::kNumInput)) { + if (auto opt = func->attrs.GetAttr(tvm::relax::attr::kNumInput)) { num_runtime_params = opt.value(); } @@ -494,10 +494,10 @@ class ParamRemapper : private ExprFunctor { static ffi::Map GetParamMapping(const ffi::Array& functions) { ParamRemapper mapper; if (functions.size()) { - auto num_inputs_0 = functions[0]->GetAttr(tvm::relax::attr::kNumInput).value(); + auto num_inputs_0 = functions[0]->GetAttr(tvm::relax::attr::kNumInput).value(); int num_params = static_cast(functions[0]->params.size()) - num_inputs_0; for (int i = 0; i < static_cast(functions.size()); i++) { - auto num_inputs_i = functions[i]->GetAttr(tvm::relax::attr::kNumInput).value(); + auto num_inputs_i = functions[i]->GetAttr(tvm::relax::attr::kNumInput).value(); TVM_FFI_ICHECK_EQ(num_params, static_cast(functions[i]->params.size()) - num_inputs_i) << "The number of parameters should be the same for all target functions"; @@ -547,15 +547,15 @@ class GlobalLiftableBindingCollector : public BaseLiftableBindingCollector { GlobalLiftableBindingCollector collector(var_remap); TVM_FFI_ICHECK(functions.size()); for (const auto& func : functions) { - int num_inputs = func->GetAttr(tvm::relax::attr::kNumInput).value(); + int num_inputs = func->GetAttr(tvm::relax::attr::kNumInput).value(); for (int i = num_inputs; i < static_cast(func->params.size()); i++) { collector.liftable_vars_.insert(func->params[i]); } collector(func); } - ffi::Array params(functions[0]->params.begin() + - functions[0]->GetAttr(tvm::relax::attr::kNumInput).value(), - functions[0]->params.end()); + ffi::Array params( + functions[0]->params.begin() + functions[0]->GetAttr(tvm::relax::attr::kNumInput).value(), + functions[0]->params.end()); // todo(@tvm-team): use c++20 designated initializers when windows CI supports it GlobalCollectInfo info = GlobalCollectInfo(); info.orig_functions = functions; @@ -660,7 +660,7 @@ class ConsumeBundledParams : public ExprMutator { } Expr VisitExpr_(const FunctionNode* func) final { - auto opt_num_input = func->GetAttr(tvm::relax::attr::kNumInput); + auto opt_num_input = func->GetAttr(tvm::relax::attr::kNumInput); TVM_FFI_ICHECK(opt_num_input.has_value()); auto num_input = opt_num_input.value(); TVM_FFI_ICHECK_EQ(func->params.size(), num_input + 1); @@ -697,10 +697,10 @@ std::vector> GetTargetFunctions( << "only functions in the list must be relax functions. " << "However, the function " << name << " is of type " << base_func.value()->GetTypeKey(); - TVM_FFI_ICHECK(func.value()->GetAttr(tvm::relax::attr::kNumInput)) + TVM_FFI_ICHECK(func.value()->GetAttr(tvm::relax::attr::kNumInput)) << "When LiftTransformParams is called with a list of function names, " - << "all functions in the list must have the kNumInput ('" << tvm::relax::attr::kNumInput - << "') attribute. " + << "all functions in the list must have the kNumInput ('" + << tvm::relax::attr::kNumInput.name << "') attribute. " << "However, the function " << name << " does not have the kNumInput attribute"; target_functions.push_back({gvar.value(), func.value()}); @@ -710,7 +710,7 @@ std::vector> GetTargetFunctions( // are not already the result of `LiftTransformParams`. for (const auto& [gvar, func] : mod->functions) { if (func->IsInstance()) { - auto opt_num_input = func->GetAttr(tvm::relax::attr::kNumInput); + auto opt_num_input = func->GetAttr(tvm::relax::attr::kNumInput); if (opt_num_input && !ends_with(gvar->name_hint, "transform_params")) { target_functions.emplace_back(gvar, func.as_or_throw()); } @@ -813,7 +813,7 @@ Pass LiftTransformParams(ffi::Variant> shared_tran std::string func_name = gvar->name_hint; if (ends_with(func_name, "transform_params")) { - func = WithAttr(func, tvm::attr::kGlobalSymbol, gvar->name_hint); + func = WithAttr(func, tvm::attr::kGlobalSymbol.name, gvar->name_hint); if (pc->GetConfig(kLiftTransformConsumeParams).value_or(false)) { func = ConsumeBundledParams()(func).as_or_throw(); } diff --git a/src/relax/transform/merge_composite_functions.cc b/src/relax/transform/merge_composite_functions.cc index 44e5220c6384..8b8b80fead5b 100644 --- a/src/relax/transform/merge_composite_functions.cc +++ b/src/relax/transform/merge_composite_functions.cc @@ -240,8 +240,8 @@ class CompositeGroupsBuilder : public MemoizedExprTranslator { return std::nullopt; } - auto composite_name_opt = mod_->Lookup(ffi::GetRef(gvar)) - ->GetAttr(tvm::relax::attr::kComposite); + auto composite_name_opt = + mod_->Lookup(ffi::GetRef(gvar))->GetAttr(tvm::relax::attr::kComposite); if (!composite_name_opt) { return std::nullopt; } @@ -250,7 +250,7 @@ class CompositeGroupsBuilder : public MemoizedExprTranslator { } ffi::Optional GetCodegenName(Group* group) { - if (auto opt_str = group->attrs.Get(tvm::relax::attr::kCodegen)) { + if (auto opt_str = group->attrs.Get(tvm::relax::attr::kCodegen.name)) { return opt_str.value().as_or_throw(); } return std::nullopt; @@ -332,7 +332,7 @@ class CompositeGroupsBuilder : public MemoizedExprTranslator { Group* CreateNewGroup(const CallNode* call) { Group* group = arena_->make(); if (ffi::Optional codegen_name = GetCodegenName(call->op)) { - group->attrs.Set(tvm::relax::attr::kCodegen, codegen_name.value()); + group->attrs.Set(tvm::relax::attr::kCodegen.name, codegen_name.value()); } return group; } @@ -465,10 +465,10 @@ class CompositeInliner : public ExprMutator { if (call->op->IsInstance()) { auto gvar = call->op.as_or_throw(); auto func = mod_->Lookup(gvar).as_or_throw(); - if (func->GetAttr(tvm::relax::attr::kComposite)) { + if (func->GetAttr(tvm::relax::attr::kComposite)) { if (!inlined_functions_.count(func)) { auto new_func = CopyWithNewVars(func); - new_func = WithoutAttr(new_func, tvm::relax::attr::kPrimitive); + new_func = WithoutAttr(new_func, tvm::relax::attr::kPrimitive.name); inlined_functions_.Set(func, new_func); } return Call(Type::Missing(), inlined_functions_[func], call->args); @@ -507,7 +507,7 @@ class CompositeFunctionAnnotator : public ExprMutator { if (call->op->IsInstance()) { GlobalVar cur_var = call->op.as_or_throw(); auto func = mod_->Lookup(cur_var).as_or_throw(); - if (auto codegen_name = func->GetAttr(tvm::relax::attr::kCodegen)) { + if (auto codegen_name = func->GetAttr(tvm::relax::attr::kCodegen)) { GlobalVar new_var = [&]() -> GlobalVar { if (var_map_.count(cur_var) > 0) { // if we visited before, we don't need to create the new function, @@ -522,8 +522,8 @@ class CompositeFunctionAnnotator : public ExprMutator { // rename the function. ffi::String new_func_name = cur_var->name_hint + "_" + codegen_name.value(); Function new_func = inliner.Run(func.as_or_throw()); - new_func = WithAttr(new_func, tvm::attr::kGlobalSymbol, new_func_name); - new_func = WithoutAttr(std::move(new_func), tvm::relax::attr::kPrimitive); + new_func = WithAttr(new_func, tvm::attr::kGlobalSymbol.name, new_func_name); + new_func = WithoutAttr(std::move(new_func), tvm::relax::attr::kPrimitive.name); // add a function with a new name. auto new_var = builder_->AddFunction(new_func, new_func_name); var_map_.insert_or_assign(cur_var, new_var); diff --git a/src/relax/transform/meta_schedule.cc b/src/relax/transform/meta_schedule.cc index f13fc2b2e0bd..13001c922814 100644 --- a/src/relax/transform/meta_schedule.cc +++ b/src/relax/transform/meta_schedule.cc @@ -133,7 +133,8 @@ Pass MetaScheduleApplyDatabase(ffi::Optional work_dir, bool enable_ /*body=*/tuned_function->body, /*ret_type=*/tuned_function->ret_type, /*attrs=*/function->attrs); - new_function = WithAttr(std::move(new_function), tvm::s_tir::attr::kIsScheduled, true); + new_function = + WithAttr(std::move(new_function), tvm::s_tir::attr::kIsScheduled.name, true); result.Set(gv, new_function); continue; } else if (enable_warning) { diff --git a/src/relax/transform/normalize.cc b/src/relax/transform/normalize.cc index 9ad107639f0b..21e74d5d3481 100644 --- a/src/relax/transform/normalize.cc +++ b/src/relax/transform/normalize.cc @@ -221,7 +221,7 @@ class GlobalVarNormalizer : private ExprMutator { /*! \brief Check if any function needs to be renamed. */ bool NeedRename() { for (const auto& [gvar, func] : module_->functions) { - auto global_symbol = func->GetAttr(tvm::attr::kGlobalSymbol); + auto global_symbol = func->GetAttr(tvm::attr::kGlobalSymbol); if (global_symbol && global_symbol.value() != gvar->name_hint) { return true; } @@ -232,7 +232,7 @@ class GlobalVarNormalizer : private ExprMutator { /*! \brief Add public functions to the builder, and update the name supplier. */ void AddPublicFunctions() { for (const auto& [gvar, func] : module_->functions) { - auto global_symbol = func->GetAttr(tvm::attr::kGlobalSymbol); + auto global_symbol = func->GetAttr(tvm::attr::kGlobalSymbol); if (!global_symbol) { continue; } @@ -252,7 +252,7 @@ class GlobalVarNormalizer : private ExprMutator { */ void AddPrivateFunctions() { for (auto [gvar, func] : module_->functions) { - auto global_symbol = func->GetAttr(tvm::attr::kGlobalSymbol); + auto global_symbol = func->GetAttr(tvm::attr::kGlobalSymbol); if (global_symbol) { continue; } diff --git a/src/relax/transform/remove_purity_checking.cc b/src/relax/transform/remove_purity_checking.cc index 5c31c11ae6bf..f86a368c33d4 100644 --- a/src/relax/transform/remove_purity_checking.cc +++ b/src/relax/transform/remove_purity_checking.cc @@ -39,7 +39,7 @@ class PurityRemover : public ExprMutator { bool purity = func->is_pure; auto ret = func; if (purity) { - ret = WithAttr(func, tvm::relax::attr::kForcePure, true); + ret = WithAttr(func, tvm::relax::attr::kForcePure.name, true); } auto new_body = VisitExpr(ret->body); if (!new_body.same_as(ret->body)) { diff --git a/src/relax/transform/remove_unused_outputs.cc b/src/relax/transform/remove_unused_outputs.cc index 0d209362f3d9..660e87cbd1d4 100644 --- a/src/relax/transform/remove_unused_outputs.cc +++ b/src/relax/transform/remove_unused_outputs.cc @@ -48,7 +48,7 @@ class PartialTupleUsageCollector : ExprVisitor { PMap num_outputs; for (const auto& [gvar, base_func] : mod->functions) { - bool is_exposed = base_func->attrs.GetAttr(tvm::attr::kGlobalSymbol).has_value(); + bool is_exposed = base_func->attrs.GetAttr(tvm::attr::kGlobalSymbol).has_value(); if (!is_exposed) { if (auto relax_func = base_func.as()) { diff --git a/src/relax/transform/remove_unused_parameters.cc b/src/relax/transform/remove_unused_parameters.cc index 9e318919d85f..b105e417800e 100644 --- a/src/relax/transform/remove_unused_parameters.cc +++ b/src/relax/transform/remove_unused_parameters.cc @@ -57,7 +57,7 @@ struct CalleeAnalysis { }; std::optional AnalyzeCallee(Function func) { - bool is_exposed = func->attrs.GetAttr(tvm::attr::kGlobalSymbol).has_value(); + bool is_exposed = func->attrs.GetAttr(tvm::attr::kGlobalSymbol).has_value(); if (is_exposed) return std::nullopt; auto free_relax_vars = [&]() -> PSet { diff --git a/src/relax/transform/rewrite_cuda_graph.cc b/src/relax/transform/rewrite_cuda_graph.cc index 3ad746c44a52..a32ba207fa59 100644 --- a/src/relax/transform/rewrite_cuda_graph.cc +++ b/src/relax/transform/rewrite_cuda_graph.cc @@ -162,7 +162,7 @@ class FuncBuilder : public ExprMutator { auto block = builder_->EndBlock(); auto body = builder_->Normalize(SeqExpr({block}, output)); ffi::Map attrs; - attrs.Set(tvm::relax::attr::kForcePure, true); + attrs.Set(tvm::relax::attr::kForcePure.name, true); auto func = Function(params, body, output->ty.as_or_throw(), /*is_pure=*/true, /*attrs=*/DictAttrs(attrs)); return func; @@ -247,7 +247,7 @@ class CUDAGraphRewritePlanner : public ExprVisitor { // these names are extracted from the type for the capturing. const auto& func = pair.second.as_or_throw(); int64_t num_inputs = - func->attrs.GetAttr(tvm::relax::attr::kNumInput).value_or(func->params.size()); + func->attrs.GetAttr(tvm::relax::attr::kNumInput).value_or(func->params.size()); auto capture_symbolic_var_name_hints = ExtractSymbolicVarHints(func); for (int i = 0; i < static_cast(func->params.size()); ++i) { ffi::Array symbolic_vars = @@ -308,9 +308,8 @@ class CUDAGraphRewritePlanner : public ExprVisitor { * from the function attributes. */ std::unordered_set ExtractSymbolicVarHints(const Function& func) { - auto symbolic_var_names = - func->attrs.GetAttr>(tvm::relax::attr::kCaptureSymbolicVars) - .value_or(ffi::Array()); + auto symbolic_var_names = func->attrs.GetAttr(tvm::relax::attr::kCaptureSymbolicVars) + .value_or(ffi::Array()); return {symbolic_var_names.begin(), symbolic_var_names.end()}; } diff --git a/src/relax/transform/run_codegen.cc b/src/relax/transform/run_codegen.cc index 226b161ef4be..b6b5b7b9e639 100644 --- a/src/relax/transform/run_codegen.cc +++ b/src/relax/transform/run_codegen.cc @@ -32,7 +32,7 @@ #include #include #include -#include +#include #include #include "../../support/ordered_set.h" @@ -65,8 +65,7 @@ class CodeGenRunner : ExprMutator { std::vector attr_entry_functions; for (const auto& [gv, func] : mod->functions) { if (func->GetLinkageType() == LinkageType::kExternal && - !func->GetAttr(tvm::relax::attr::kCodegen) && - func->IsInstance()) { + !func->GetAttr(tvm::relax::attr::kCodegen) && func->IsInstance()) { attr_entry_functions.push_back(gv); } } @@ -87,11 +86,11 @@ class CodeGenRunner : ExprMutator { auto out_mod = builder_->GetContextIRModule(); if (ext_mods.size()) { - if (auto opt_old_ext_mods = mod->GetAttr>(tvm::attr::kExternalMods)) { + if (auto opt_old_ext_mods = mod->GetAttr(tvm::attr::kExternalMods)) { auto old_ext_mods = opt_old_ext_mods.value(); ext_mods.insert(ext_mods.begin(), old_ext_mods.begin(), old_ext_mods.end()); } - out_mod = WithAttr(out_mod, tvm::attr::kExternalMods, std::move(ext_mods)); + out_mod = WithAttr(out_mod, tvm::attr::kExternalMods.name, std::move(ext_mods)); } if (constant_names.size()) { @@ -101,7 +100,7 @@ class CodeGenRunner : ExprMutator { TVM_FFI_ICHECK(!constants.count(name)) << "More than one constant with the name " << name; constants.Set(name, constant->value.cast()); } - out_mod = WithAttr(out_mod, tvm::attr::kConstNameToConstant, std::move(constants)); + out_mod = WithAttr(out_mod, tvm::attr::kConstNameToConstant.name, std::move(constants)); } // TODO(@tvm-team): Implicit pass dependency. Revisit when we have a better way to handle this. @@ -136,8 +135,8 @@ class CodeGenRunner : ExprMutator { extern_funcs_.insert_or_assign(gvar_node, new_func); // Remove the global symbol and codegen attributes from the function so that it can be // removed the module. - func = WithoutAttr(std::move(func), tvm::attr::kGlobalSymbol); - func = WithoutAttr(std::move(func), tvm::relax::attr::kCodegen); + func = WithoutAttr(std::move(func), tvm::attr::kGlobalSymbol.name); + func = WithoutAttr(std::move(func), tvm::relax::attr::kCodegen.name); builder_->UpdateFunction(gvar, func); return create_call_dps_packed(new_func, ret_ty); } @@ -165,7 +164,7 @@ class CodeGenRunner : ExprMutator { Expr VisitExpr_(const FunctionNode* func_node) override { Function func = ffi::GetRef(func_node); - auto opt_codegen = func->GetAttr(tvm::relax::attr::kCodegen); + auto opt_codegen = func->GetAttr(tvm::relax::attr::kCodegen); if (opt_codegen) { auto ext_symbol = GetExtSymbol(func); size_t count = 0; @@ -195,7 +194,7 @@ class CodeGenRunner : ExprMutator { PostOrderVisit(entry.second, [&target_functions](Expr e) { if (e->IsInstance()) { auto f = e.as_or_throw(); - if (auto target_opt = f->GetAttr(tvm::relax::attr::kCodegen)) { + if (auto target_opt = f->GetAttr(tvm::relax::attr::kCodegen)) { ffi::String target = target_opt.value(); target_functions[target].push_back(f); } diff --git a/src/relax/transform/specialize_function_based_on_callsite.cc b/src/relax/transform/specialize_function_based_on_callsite.cc index 3694581c618e..457a3d72ce18 100644 --- a/src/relax/transform/specialize_function_based_on_callsite.cc +++ b/src/relax/transform/specialize_function_based_on_callsite.cc @@ -25,7 +25,7 @@ #include #include #include -#include +#include #include #include @@ -55,7 +55,7 @@ class SpecializeTIRCallArgs : ExprMutator { if (func->IsInstance()) { const auto& base_func = mod->Lookup(gv); // Only non primitive relax functions - if (base_func->HasNonzeroAttr(tvm::relax::attr::kPrimitive)) { + if (base_func->HasNonzeroAttr(tvm::relax::attr::kPrimitive.name)) { continue; } relax::Function update_func = VisitExpr(func).as_or_throw(); @@ -144,7 +144,8 @@ class SpecializeTIRCallArgs : ExprMutator { } auto new_pfunc = Specialize(pfunc, param_map); - auto new_function = WithAttr(new_pfunc, tvm::relax::attr::kScoped, static_cast(1)); + auto new_function = + WithAttr(new_pfunc, tvm::relax::attr::kScoped.name, static_cast(1)); updates_->Add(gv, new_function); return call; } diff --git a/src/relax/transform/split_call_tir_by_pattern.cc b/src/relax/transform/split_call_tir_by_pattern.cc index 2efecb2182a0..b5836704f6d4 100644 --- a/src/relax/transform/split_call_tir_by_pattern.cc +++ b/src/relax/transform/split_call_tir_by_pattern.cc @@ -631,7 +631,7 @@ std::pair> SplitFunctions( } if (!has_second_func) { // No need to split the function. - return {WithAttr(func, tvm::relax::attr::kLibraryKernel, library_code), std::nullopt}; + return {WithAttr(func, tvm::relax::attr::kLibraryKernel.name, library_code), std::nullopt}; } // Step 2. Split the function into two functions. Stmt body1 = BlockRemover::RemoveBlockByPartition( @@ -656,7 +656,7 @@ std::pair> SplitFunctions( arg_partition->push_back(arg_partition1); new_params1.push_back(partitioner->intermediate_tensor.value()); tirx::Function func1 = tirx::Function(new_params1, SeqStmt(body1), func->ret_type, func->attrs); - func1 = WithAttr(func1, tvm::relax::attr::kLibraryKernel, library_code); + func1 = WithAttr(func1, tvm::relax::attr::kLibraryKernel.name, library_code); // Step 4. Craft the second function. ffi::Array new_params2; std::vector arg_partition2; @@ -687,19 +687,19 @@ void StringReplace(std::string* subject, const std::string& search, const std::s tvm::BaseFunc CodegenWithLibrary(const tirx::FunctionNode* pf, ffi::String global_symbol) { using namespace tvm::tirx; - ffi::Optional library_code = - pf->attrs.GetAttr(tvm::relax::attr::kLibraryKernel); + ffi::Optional library_code = pf->attrs.GetAttr(tvm::relax::attr::kLibraryKernel); if (!library_code.has_value()) { return ffi::GetRef(pf); } std::string source = library_code.value(); StringReplace(&source, "{global_symbol}", global_symbol); ExternFunc ret(global_symbol); - ret = WithAttrs(std::move(ret), - ffi::Map{ - {ffi::String(tvm::relax::attr::kCSource), ffi::String(source)}, - {ffi::String(tvm::relax::attr::kCSourceFmt), ffi::String(kCSourceFmtCuda)}, - }); + ret = + WithAttrs(std::move(ret), + ffi::Map{ + {ffi::String(tvm::relax::attr::kCSource.name), ffi::String(source)}, + {ffi::String(tvm::relax::attr::kCSourceFmt.name), ffi::String(kCSourceFmtCuda)}, + }); return ret; } @@ -785,7 +785,7 @@ class SplitMutator : public ExprMutator { } GlobalVar gv2 = builder_->AddFunction(func2, "unfused_epilogue"); Call call2(Type::Missing(), call_tir_op_, {gv2, Tuple(args2)}, call->attrs, call->ty_args); - builder_->UpdateFunction(gv, WithoutAttr(func, tvm::attr::kGlobalSymbol)); + builder_->UpdateFunction(gv, WithoutAttr(func, tvm::attr::kGlobalSymbol.name)); return call2; } diff --git a/src/relax/transform/split_layout_rewrite_preproc.cc b/src/relax/transform/split_layout_rewrite_preproc.cc index f2dbad3bcdb4..b8acefc2512f 100644 --- a/src/relax/transform/split_layout_rewrite_preproc.cc +++ b/src/relax/transform/split_layout_rewrite_preproc.cc @@ -94,9 +94,9 @@ class SplitFunctionLayoutRewrite : public s_tir::StmtExprMutator { ffi::Map dict; for (const auto& [key, original_value] : original_func_->attrs->dict) { - if (key == tvm::attr::kGlobalSymbol) { + if (key == tvm::attr::kGlobalSymbol.name) { dict.Set(key, original_value.as_or_throw() + "_weight_prepack"); - } else if (key != tvm::s_tir::attr::kLayoutFreeTensors) { + } else if (key != tvm::s_tir::attr::kLayoutFreeTensors.name) { dict.Set(key, original_value); } } @@ -140,9 +140,9 @@ class SplitFunctionLayoutRewrite : public s_tir::StmtExprMutator { ffi::Map dict; for (const auto& [key, original_value] : original_func_->attrs->dict) { - if (key == tvm::attr::kGlobalSymbol) { + if (key == tvm::attr::kGlobalSymbol.name) { dict.Set(key, original_value.as_or_throw() + "_prepacked"); - } else if (key != tvm::s_tir::attr::kLayoutFreeTensors) { + } else if (key != tvm::s_tir::attr::kLayoutFreeTensors.name) { dict.Set(key, original_value); } } @@ -170,7 +170,7 @@ class SplitFunctionLayoutRewrite : public s_tir::StmtExprMutator { s_tir::SBlock block = s_tir::StmtExprMutator::Mutate_(op, inplace_mode) .ValueOrUnchanged(ffi::GetRef(op)) .as_or_throw(); - auto it = op->annotations.find(tvm::s_tir::attr::kMetaScheduleLayoutRewritePreproc); + auto it = op->annotations.find(tvm::s_tir::attr::kMetaScheduleLayoutRewritePreproc.name); bool is_layout_rewrite_preproc = it != op->annotations.end() && IsOne((*it).second.cast()); @@ -209,7 +209,7 @@ class SplitFunctionLayoutRewrite : public s_tir::StmtExprMutator { op->writes[0]->source.as_or_throw()}); auto new_annotations = op->annotations; - new_annotations.erase(tvm::s_tir::attr::kMetaScheduleLayoutRewritePreproc); + new_annotations.erase(tvm::s_tir::attr::kMetaScheduleLayoutRewritePreproc.name); auto n = ffi::make_object(*block.get()); n->annotations = new_annotations; return s_tir::SBlock(n); diff --git a/src/relax/transform/static_plan_block_memory.cc b/src/relax/transform/static_plan_block_memory.cc index 3705d9bf31e7..cb54fc5cbf65 100644 --- a/src/relax/transform/static_plan_block_memory.cc +++ b/src/relax/transform/static_plan_block_memory.cc @@ -436,14 +436,11 @@ void SetTIRVarRangeConstraints(Function func, sym::AnalyzerObj* ana, // NOTE: we only apply the annotated bounds to the TIR variables that // appear in the **function signature**. ffi::Map var_upper_bound_attr_raw = - func->GetAttr>(tvm::relax::attr::kTirVarUpperBound) - .value_or(ffi::Map()); + func->GetAttr(tvm::relax::attr::kTirVarUpperBound).value_or(ffi::Map()); ffi::Map var_lower_bound_attr_raw = - func->GetAttr>(tvm::relax::attr::kTirVarLowerBound) - .value_or(ffi::Map()); + func->GetAttr(tvm::relax::attr::kTirVarLowerBound).value_or(ffi::Map()); ffi::Array non_negative_var_attr_raw = - func->GetAttr>(tvm::relax::attr::kTirNonNegativeVar) - .value_or(ffi::Array()); + func->GetAttr(tvm::relax::attr::kTirNonNegativeVar).value_or(ffi::Array()); std::unordered_map var_upper_bound_attr; std::unordered_map var_lower_bound_attr; std::unordered_set non_negative_var_attr; diff --git a/src/relax/transform/to_mixed_precision.cc b/src/relax/transform/to_mixed_precision.cc index b9f73f846f62..690da0fd4398 100644 --- a/src/relax/transform/to_mixed_precision.cc +++ b/src/relax/transform/to_mixed_precision.cc @@ -24,7 +24,7 @@ #include #include #include -#include +#include #include #include diff --git a/src/relax/transform/utils.h b/src/relax/transform/utils.h index 67d8a18f2c16..4a0246144bc8 100644 --- a/src/relax/transform/utils.h +++ b/src/relax/transform/utils.h @@ -131,7 +131,7 @@ TVM_DLL IRModule DeadCodeElimination(const IRModule& mod, ffi::ArrayGetAttr(tvm::attr::kGlobalSymbol); + const auto name_node = func->GetAttr(tvm::attr::kGlobalSymbol); TVM_FFI_ICHECK(name_node.has_value()) << "Fail to retrieve external symbol."; return std::string(name_node.value()); } diff --git a/src/relax/utils.cc b/src/relax/utils.cc index faa690dcdb55..4e1c6bbe6da9 100644 --- a/src/relax/utils.cc +++ b/src/relax/utils.cc @@ -26,7 +26,7 @@ #include #include #include -#include +#include #include #include diff --git a/src/s_tir/analysis/calculate_allocated_memory.cc b/src/s_tir/analysis/calculate_allocated_memory.cc index 1757f8eb40a1..81f9bf716a37 100644 --- a/src/s_tir/analysis/calculate_allocated_memory.cc +++ b/src/s_tir/analysis/calculate_allocated_memory.cc @@ -224,7 +224,7 @@ Pass VerifyVTCMLimit(ffi::Optional default_target) { auto func = opt.value(); std::optional limit = std::nullopt; - if (auto func_target = func->GetAttr(tvm::attr::kTarget)) { + if (auto func_target = func->GetAttr(tvm::attr::kTarget)) { limit = GetVTCMCapacity(func_target.value(), ctx); } else if (default_target) { limit = GetVTCMCapacity(default_target.value(), ctx); diff --git a/src/s_tir/analysis/check_well_formed.cc b/src/s_tir/analysis/check_well_formed.cc index 57c0be28c190..97beb877e349 100644 --- a/src/s_tir/analysis/check_well_formed.cc +++ b/src/s_tir/analysis/check_well_formed.cc @@ -376,8 +376,8 @@ class WellFormednessVerifier : public StmtExprVisitor { std::vector globals; for (const auto& [global, function] : op->functions) globals.push_back(global); std::sort(globals.begin(), globals.end(), [&](const GlobalVar& a, const GlobalVar& b) { - bool a_public = op->functions[a]->GetAttr(tvm::attr::kGlobalSymbol).has_value(); - bool b_public = op->functions[b]->GetAttr(tvm::attr::kGlobalSymbol).has_value(); + bool a_public = op->functions[a]->GetAttr(tvm::attr::kGlobalSymbol).has_value(); + bool b_public = op->functions[b]->GetAttr(tvm::attr::kGlobalSymbol).has_value(); return a_public != b_public ? a_public > b_public : a->name_hint < b->name_hint; }); for (const auto& global : globals) { @@ -464,7 +464,7 @@ class WellFormednessVerifier : public StmtExprVisitor { Result Visit_(const ForNode* op) override { if (check_ == Check::kThreadBinding) { - if (auto tag = op->annotations.Get(attr::kThreadBinding)) { + if (auto tag = op->annotations.Get(attr::kThreadBinding.name)) { Verify(op->kind == ForKind::kParallel) << "ValueError: thread_binding requires a parallel loop at "; Verify(tag->as().has_value()) diff --git a/src/s_tir/analysis/estimate_flops.cc b/src/s_tir/analysis/estimate_flops.cc index be6c5e289056..5a730bb739aa 100644 --- a/src/s_tir/analysis/estimate_flops.cc +++ b/src/s_tir/analysis/estimate_flops.cc @@ -268,7 +268,7 @@ double EstimateTIRFlops(const IRModule& mod) { double cached_result = 0; for (const auto& [_, base_func] : mod->functions) { if (const auto* func = base_func.as()) { - if (auto cached = func->attrs.GetAttr(tvm::s_tir::attr::kEstimatedFlops)) { + if (auto cached = func->attrs.GetAttr(tvm::s_tir::attr::kEstimatedFlops)) { cached_result += cached.value(); } else if (func->body.has_value()) { result += counter.Dispatch(func->body.value()); diff --git a/src/s_tir/meta_schedule/arg_info.cc b/src/s_tir/meta_schedule/arg_info.cc index 33212109379d..6c2c99ed322b 100644 --- a/src/s_tir/meta_schedule/arg_info.cc +++ b/src/s_tir/meta_schedule/arg_info.cc @@ -32,12 +32,12 @@ namespace meta_schedule { /*! * \brief Find the entry function of the given IRModule, i.e, functions marked by - * `tirx::attr::kIsEntryFunc`, whose name is `main` or being the only Function. + * `tirx::attr::kIsEntryFunc.name`, whose name is `main` or being the only Function. * \param mod The IRModule to find the entry function. * \return The entry function. */ inline tirx::Function FindEntryFunc(const IRModule& mod) { - // Priority 1: Function marked as `tirx::attr::kIsEntryFunc` + // Priority 1: Function marked as `tirx::attr::kIsEntryFunc.name` int num_function = 0; const tirx::FunctionNode* main_func = nullptr; const tirx::FunctionNode* last_func = nullptr; @@ -46,7 +46,7 @@ inline tirx::Function FindEntryFunc(const IRModule& mod) { BaseFunc base_func = kv.second; if (const auto* func = base_func.as()) { last_func = func; - if (func->HasNonzeroAttr(tvm::tirx::attr::kIsEntryFunc)) { + if (func->HasNonzeroAttr(tvm::tirx::attr::kIsEntryFunc.name)) { return ffi::GetRef(func); } if (gv->name_hint == "main") { diff --git a/src/s_tir/meta_schedule/database/schedule_fn_database.cc b/src/s_tir/meta_schedule/database/schedule_fn_database.cc index 008da83a8bda..d7a480a335cc 100644 --- a/src/s_tir/meta_schedule/database/schedule_fn_database.cc +++ b/src/s_tir/meta_schedule/database/schedule_fn_database.cc @@ -55,7 +55,7 @@ class ScheduleFnDatabaseObj : public DatabaseObj { ffi::Optional QuerySchedule(const IRModule& mod, const Target& target, const ffi::String& workload_name) final { s_tir::Schedule sch = s_tir::Schedule::Traced( - WithAttr(mod, tvm::s_tir::meta_schedule::attr::kTaskName, workload_name), + WithAttr(mod, tvm::s_tir::meta_schedule::attr::kTaskName.name, workload_name), /*rand_state=*/-1, /*debug_mode=*/0, /*error_render_level=*/s_tir::ScheduleErrorRenderLevel::kDetail); diff --git a/src/s_tir/meta_schedule/feature_extractor/per_store_feature.cc b/src/s_tir/meta_schedule/feature_extractor/per_store_feature.cc index f28645cedb4d..cb29db7129fe 100644 --- a/src/s_tir/meta_schedule/feature_extractor/per_store_feature.cc +++ b/src/s_tir/meta_schedule/feature_extractor/per_store_feature.cc @@ -95,7 +95,7 @@ std::vector GetTensorShape(const TensorVar& tensor, sym::AnalyzerObj* a */ int64_t GetAutoUnrollMaxStep(const ForNode* loop) { if (ffi::Optional auto_unroll = - GetAnn(loop, tvm::tirx::attr::kAutoUnrollMaxStep)) { + GetAnn(loop, tvm::tirx::attr::kAutoUnrollMaxStep.name)) { return static_cast(auto_unroll.value()->value); } return -1; @@ -340,7 +340,7 @@ Pass SimplifyForFeatureExtraction(bool normalize_thread_bindings = false) { .ValueOrUnchanged(loop->annotations); Stmt body = Mutate(loop->body, inplace_mode).ValueOrUnchanged(loop->body); VarRemapSet(loop->loop_var, previous_remap); - annotations.erase(tvm::s_tir::attr::kThreadBinding); + annotations.erase(tvm::s_tir::attr::kThreadBinding.name); if (!annotations.empty()) { PrimType ty = loop->loop_var.ty(); body = For(PrimVar("annotation", ty), IntImm(ty, 0), IntImm(ty, 1), ForKind::kDefault, diff --git a/src/s_tir/meta_schedule/mutator/mutate_parallel.cc b/src/s_tir/meta_schedule/mutator/mutate_parallel.cc index bb9b8145ef62..687b5a5ef012 100644 --- a/src/s_tir/meta_schedule/mutator/mutate_parallel.cc +++ b/src/s_tir/meta_schedule/mutator/mutate_parallel.cc @@ -42,7 +42,7 @@ bool IsAnnotateWithParallel(const Instruction& inst) { } TVM_FFI_ICHECK_EQ(inst->attrs.size(), 1); ffi::String ann_key = inst->attrs[0].as_or_throw(); - return ann_key == tvm::s_tir::attr::kMetaScheduleParallel; + return ann_key == tvm::s_tir::attr::kMetaScheduleParallel.name; } /*! diff --git a/src/s_tir/meta_schedule/mutator/mutate_tile_size.cc b/src/s_tir/meta_schedule/mutator/mutate_tile_size.cc index e615c0f2d1c2..9982f083ca05 100644 --- a/src/s_tir/meta_schedule/mutator/mutate_tile_size.cc +++ b/src/s_tir/meta_schedule/mutator/mutate_tile_size.cc @@ -127,7 +127,7 @@ void FindSampleVectorize(const Trace& trace, std::vector* inst, TVM_FFI_ICHECK_EQ(inst->attrs.size(), 1); TVM_FFI_ICHECK_EQ(inst->inputs.size(), 2); if (inst->attrs[0].as_or_throw() == - tvm::s_tir::attr::kMetaScheduleCooperativeFetch) { + tvm::s_tir::attr::kMetaScheduleCooperativeFetch.name) { const auto* ann_val = inst->inputs[1].as(); TVM_FFI_ICHECK(ann_val); annotated.insert(ann_val); diff --git a/src/s_tir/meta_schedule/mutator/mutate_unroll.cc b/src/s_tir/meta_schedule/mutator/mutate_unroll.cc index 6a9717f9fa0e..fce78165f2ab 100644 --- a/src/s_tir/meta_schedule/mutator/mutate_unroll.cc +++ b/src/s_tir/meta_schedule/mutator/mutate_unroll.cc @@ -39,8 +39,8 @@ bool IsAnnotateWithUnroll(const Instruction& inst) { } TVM_FFI_ICHECK_EQ(inst->attrs.size(), 1); ffi::String ann_key = inst->attrs[0].as_or_throw(); - return ann_key == tvm::s_tir::attr::kMetaScheduleUnrollExplicit || - ann_key == tvm::s_tir::attr::kMetaScheduleUnrollImplicit; + return ann_key == tvm::s_tir::attr::kMetaScheduleUnrollExplicit.name || + ann_key == tvm::s_tir::attr::kMetaScheduleUnrollImplicit.name; } } // namespace s_tir diff --git a/src/s_tir/meta_schedule/postproc/disallow_async_strided_mem_copy.cc b/src/s_tir/meta_schedule/postproc/disallow_async_strided_mem_copy.cc index 9fd4faca7e6d..b6519c90b1b4 100644 --- a/src/s_tir/meta_schedule/postproc/disallow_async_strided_mem_copy.cc +++ b/src/s_tir/meta_schedule/postproc/disallow_async_strided_mem_copy.cc @@ -123,7 +123,7 @@ class DisallowAsyncStridedMemCopyObj : public PostprocObj { pass_list.push_back(tirx::transform::VectorizeLoop(true)); pass_list.push_back(tirx::transform::StorageRewrite()); tirx::Function f = WithAttr(ffi::GetRef(function), - tvm::attr::kGlobalSymbol, ffi::String(g_var->name_hint)); + tvm::attr::kGlobalSymbol.name, ffi::String(g_var->name_hint)); IRModule mod = IRModule(ffi::Map({{GlobalVar(g_var->name_hint), f}})); lowered = tvm::transform::Sequential(pass_list)(std::move(mod)); diff --git a/src/s_tir/meta_schedule/postproc/rewrite_cooperative_fetch.cc b/src/s_tir/meta_schedule/postproc/rewrite_cooperative_fetch.cc index c9e8446d3fc8..554206aa5f06 100644 --- a/src/s_tir/meta_schedule/postproc/rewrite_cooperative_fetch.cc +++ b/src/s_tir/meta_schedule/postproc/rewrite_cooperative_fetch.cc @@ -66,7 +66,7 @@ ffi::Optional ParseAnnotate(const Schedule& sch, const Instruction& in TVM_FFI_ICHECK_EQ(inst->inputs.size(), 2); TVM_FFI_ICHECK_EQ(inst->attrs.size(), 1); ffi::String ann_key = inst->attrs[0].as_or_throw(); - if (ann_key != tvm::s_tir::attr::kMetaScheduleCooperativeFetch) { + if (ann_key != tvm::s_tir::attr::kMetaScheduleCooperativeFetch.name) { return std::nullopt; } *vector_lane = static_cast( @@ -88,7 +88,7 @@ bool ParseWarpExecutionAnn(const Schedule& sch, const Instruction& inst) { TVM_FFI_ICHECK_EQ(inst->inputs.size(), 2); TVM_FFI_ICHECK_EQ(inst->attrs.size(), 1); ffi::String ann_key = inst->attrs[0].as_or_throw(); - return ann_key == tvm::s_tir::attr::kWarpExecution; + return ann_key == tvm::s_tir::attr::kWarpExecution.name; } size_t GetMaxUsedDtypeBytes(SBlock block) { @@ -178,7 +178,7 @@ bool RewriteCooperativeFetchObj::Apply(const s_tir::Schedule& sch) { } auto task = [thread_extent_x, thread_extent_y, vector_lane, sch, block = opt_block_rv.value()]() mutable -> void { - sch->Unannotate(block, tvm::s_tir::attr::kMetaScheduleCooperativeFetch); + sch->Unannotate(block, tvm::s_tir::attr::kMetaScheduleCooperativeFetch.name); s_tir::LoopRV fused = sch->GetLoops(block).back(); int64_t fused_extent = -1; const auto* extent_imm = sch->Get(fused)->extent.as(); diff --git a/src/s_tir/meta_schedule/postproc/rewrite_layout.cc b/src/s_tir/meta_schedule/postproc/rewrite_layout.cc index 8291df1dea57..14a5747ca417 100644 --- a/src/s_tir/meta_schedule/postproc/rewrite_layout.cc +++ b/src/s_tir/meta_schedule/postproc/rewrite_layout.cc @@ -144,7 +144,7 @@ class LayoutFreeTensorCollector : public StmtExprVisitor { ffi::Array CollectLayoutFreeTensors(const FunctionNode* func) { // Only rewrite Functions with attr "layout_free_tensors" ffi::Array layout_free_tensor_index = - func->GetAttr(s_tir::attr::kLayoutFreeTensors, ffi::Array()).value(); + func->GetAttr(s_tir::attr::kLayoutFreeTensors.name, ffi::Array()).value(); ffi::Array layout_free_tensors; for (int64_t index : layout_free_tensor_index) { @@ -216,7 +216,7 @@ bool RewriteLayout(const Schedule& sch) { std::vector> results; auto add_layout_rewrite_block = [&sch](SBlockRV consumer_block_rv, int tensor_index) { SBlockRV rewrite_block_rv = sch->CacheRead(consumer_block_rv, tensor_index, "global"); - sch->Annotate(rewrite_block_rv, s_tir::attr::kMetaScheduleLayoutRewritePreproc, true); + sch->Annotate(rewrite_block_rv, s_tir::attr::kMetaScheduleLayoutRewritePreproc.name, true); }; for (const auto& [g_var, base_func] : sch->mod()->functions) { diff --git a/src/s_tir/meta_schedule/postproc/rewrite_parallel_vectorize_unroll.cc b/src/s_tir/meta_schedule/postproc/rewrite_parallel_vectorize_unroll.cc index 0941aa6084ca..2846334260e4 100644 --- a/src/s_tir/meta_schedule/postproc/rewrite_parallel_vectorize_unroll.cc +++ b/src/s_tir/meta_schedule/postproc/rewrite_parallel_vectorize_unroll.cc @@ -116,22 +116,22 @@ bool ParseAnnotation(const SBlock& block, ParsedAnnotation* parsed) { bool found = false; *parsed = ParsedAnnotation{-1, -1, -1, -1, -1, -1}; for (const auto& ann : block->annotations) { - if (ann.first == tvm::s_tir::attr::kMetaScheduleParallel) { + if (ann.first == tvm::s_tir::attr::kMetaScheduleParallel.name) { found = true; if (auto opt_int_imm = ann.second.try_cast()) { parsed->max_parallel_extent = (*opt_int_imm)->value.as().value(); } - } else if (ann.first == tvm::s_tir::attr::kMetaScheduleVectorize) { + } else if (ann.first == tvm::s_tir::attr::kMetaScheduleVectorize.name) { found = true; if (auto opt_int_imm = ann.second.try_cast()) { parsed->max_vectorize_extent = (*opt_int_imm)->value.as().value(); } - } else if (ann.first == tvm::s_tir::attr::kMetaScheduleUnrollExplicit) { + } else if (ann.first == tvm::s_tir::attr::kMetaScheduleUnrollExplicit.name) { found = true; if (auto opt_int_imm = ann.second.try_cast()) { parsed->unroll_explicit = (*opt_int_imm)->value.as().value(); } - } else if (ann.first == tvm::s_tir::attr::kMetaScheduleUnrollImplicit) { + } else if (ann.first == tvm::s_tir::attr::kMetaScheduleUnrollImplicit.name) { found = true; if (auto opt_int_imm = ann.second.try_cast()) { parsed->unroll_implicit = (*opt_int_imm)->value.as().value(); @@ -144,16 +144,16 @@ bool ParseAnnotation(const SBlock& block, ParsedAnnotation* parsed) { void RemoveParsedAnn(const Schedule& sch, const SBlockRV& block_rv, const ParsedAnnotation& parsed) { if (parsed.max_parallel_extent != -1) { - sch->Unannotate(block_rv, tvm::s_tir::attr::kMetaScheduleParallel); + sch->Unannotate(block_rv, tvm::s_tir::attr::kMetaScheduleParallel.name); } if (parsed.max_vectorize_extent != -1) { - sch->Unannotate(block_rv, tvm::s_tir::attr::kMetaScheduleVectorize); + sch->Unannotate(block_rv, tvm::s_tir::attr::kMetaScheduleVectorize.name); } if (parsed.unroll_explicit != -1) { - sch->Unannotate(block_rv, tvm::s_tir::attr::kMetaScheduleUnrollExplicit); + sch->Unannotate(block_rv, tvm::s_tir::attr::kMetaScheduleUnrollExplicit.name); } if (parsed.unroll_implicit != -1) { - sch->Unannotate(block_rv, tvm::s_tir::attr::kMetaScheduleUnrollImplicit); + sch->Unannotate(block_rv, tvm::s_tir::attr::kMetaScheduleUnrollImplicit.name); } } @@ -444,8 +444,8 @@ void RewriteUnroll(const Schedule& sch, int unroll_explicit, int max_step, const return; } - sch->Annotate(loop, tvm::tirx::attr::kAutoUnrollMaxStep, max_step); - sch->Annotate(loop, tvm::tirx::attr::kUnrollExplicit, unroll_explicit); + sch->Annotate(loop, tvm::tirx::attr::kAutoUnrollMaxStep.name, max_step); + sch->Annotate(loop, tvm::tirx::attr::kUnrollExplicit.name, unroll_explicit); } } // namespace s_tir diff --git a/src/s_tir/meta_schedule/postproc/rewrite_reduction_block.cc b/src/s_tir/meta_schedule/postproc/rewrite_reduction_block.cc index 7d7762cd5ffa..87befb35b76e 100644 --- a/src/s_tir/meta_schedule/postproc/rewrite_reduction_block.cc +++ b/src/s_tir/meta_schedule/postproc/rewrite_reduction_block.cc @@ -162,20 +162,20 @@ bool RewriteReductionBlockObj::Apply(const s_tir::Schedule& sch) { s_tir::SBlockRV init_block_rv = sch->DecomposeReduction(block_rv, loop_rvs[decompose_point]); // Rewrite auto tensorization related annotations - if (s_tir::GetAnn(block_sref, tvm::s_tir::attr::kMetaScheduleAutoTensorize) + if (s_tir::GetAnn(block_sref, tvm::s_tir::attr::kMetaScheduleAutoTensorize.name) .has_value()) { // Remove tensorization annotation as it shouldn't be propagated to the init block. - sch->Unannotate(init_block_rv, tvm::s_tir::attr::kMetaScheduleAutoTensorize); + sch->Unannotate(init_block_rv, tvm::s_tir::attr::kMetaScheduleAutoTensorize.name); ffi::Optional tensorize_init = s_tir::GetAnn( - block_sref, tvm::s_tir::attr::kMetaScheduleAutoTensorizeInit); + block_sref, tvm::s_tir::attr::kMetaScheduleAutoTensorizeInit.name); // The annotation of tensorization of the init statement should be moved to the init block // after 'DecomposeReduction'. // Annotate to hint `RewriteTensorize` postprocessor even if tensorize_init is std::nullopt. - sch->Annotate(init_block_rv, tvm::s_tir::attr::kMetaScheduleAutoTensorize, + sch->Annotate(init_block_rv, tvm::s_tir::attr::kMetaScheduleAutoTensorize.name, tensorize_init.value_or("")); if (tensorize_init.has_value()) { - sch->Unannotate(block_rv, tvm::s_tir::attr::kMetaScheduleAutoTensorizeInit); - sch->Unannotate(init_block_rv, tvm::s_tir::attr::kMetaScheduleAutoTensorizeInit); + sch->Unannotate(block_rv, tvm::s_tir::attr::kMetaScheduleAutoTensorizeInit.name); + sch->Unannotate(init_block_rv, tvm::s_tir::attr::kMetaScheduleAutoTensorizeInit.name); } } ++rewritten; diff --git a/src/s_tir/meta_schedule/postproc/rewrite_tensorize.cc b/src/s_tir/meta_schedule/postproc/rewrite_tensorize.cc index cbac9b847b25..71208c56a5c1 100644 --- a/src/s_tir/meta_schedule/postproc/rewrite_tensorize.cc +++ b/src/s_tir/meta_schedule/postproc/rewrite_tensorize.cc @@ -40,8 +40,8 @@ void CollectTensorizationJobs( auto walk_fn = [=, &jobs](const s_tir::SBlock& block) -> ffi::Expected { tirx::StmtSRef block_sref = sch->GetSRef(block.get()); std::string block_name = block_sref->StmtAs()->name_hint; - if (ffi::Optional intrin_name = - s_tir::GetAnn(block_sref, tvm::s_tir::attr::kMetaScheduleAutoTensorize)) { + if (ffi::Optional intrin_name = s_tir::GetAnn( + block_sref, tvm::s_tir::attr::kMetaScheduleAutoTensorize.name)) { if (intrin_name.value() != "") { jobs->emplace_back(block_name, func_name, [sch, intrin_name](s_tir::SBlockRV block) { try { @@ -102,7 +102,7 @@ bool RewriteTensorizeObj::Apply(const s_tir::Schedule& sch) { const ffi::String& func_name = std::get<1>(job); const auto& job_func = std::get<2>(job); SBlockRV block = sch->GetSBlock(block_name, func_name); - sch->Unannotate(block, tvm::s_tir::attr::kMetaScheduleAutoTensorize); + sch->Unannotate(block, tvm::s_tir::attr::kMetaScheduleAutoTensorize.name); job_func(block); } return true; diff --git a/src/s_tir/meta_schedule/postproc/verify_gpu_code.cc b/src/s_tir/meta_schedule/postproc/verify_gpu_code.cc index 6413e8e2582c..79f57b5881f2 100644 --- a/src/s_tir/meta_schedule/postproc/verify_gpu_code.cc +++ b/src/s_tir/meta_schedule/postproc/verify_gpu_code.cc @@ -86,13 +86,13 @@ class ThreadExtentChecker : public StmtExprVisitor { ffi::Optional Visit_(const SBlockNode* block) { int old_thread_idx_x = thread_idx_x; - if (block->annotations.count(tvm::s_tir::attr::kWarpExecution)) { + if (block->annotations.count(tvm::s_tir::attr::kWarpExecution.name)) { thread_idx_x = thread_warp_size_; } if (ffi::Optional low_inclusive = - GetAnn(block, tvm::s_tir::attr::kMetaScheduleThreadExtentLowInclusive)) { - if (ffi::Optional high_inclusive = - GetAnn(block, tvm::s_tir::attr::kMetaScheduleThreadExtentHighInclusive)) { + GetAnn(block, tvm::s_tir::attr::kMetaScheduleThreadExtentLowInclusive.name)) { + if (ffi::Optional high_inclusive = GetAnn( + block, tvm::s_tir::attr::kMetaScheduleThreadExtentHighInclusive.name)) { int64_t low = low_inclusive.value(); int64_t high = high_inclusive.value(); int64_t thread_extent_product = thread_idx_x * thread_idx_y * thread_idx_z; @@ -200,11 +200,12 @@ class VerifyGPUCodeObj : public PostprocObj { // Convert Function to IRModule tvm::transform::PassContext pass_ctx = tvm::transform::PassContext::Current(); tirx::Function f = WithAttr(ffi::GetRef(function), - tvm::attr::kGlobalSymbol, ffi::String(g_var->name_hint)); - f = WithAttr(f, tvm::attr::kTarget, this->target_.value()); // Required for LowerIntrin - bool noalias = pass_ctx->GetConfig(tvm::tirx::attr::kNoAlias).value_or(true); + tvm::attr::kGlobalSymbol.name, ffi::String(g_var->name_hint)); + f = WithAttr(f, tvm::attr::kTarget.name, + this->target_.value()); // Required for LowerIntrin + bool noalias = pass_ctx->GetConfig(tvm::tirx::attr::kNoAlias.name).value_or(true); if (noalias) { - f = WithAttr(std::move(f), tvm::tirx::attr::kNoAlias, true); + f = WithAttr(std::move(f), tvm::tirx::attr::kNoAlias.name, true); } IRModule mod = IRModule(ffi::Map({{GlobalVar(g_var->name_hint), f}})); diff --git a/src/s_tir/meta_schedule/schedule_rule/add_rfactor.cc b/src/s_tir/meta_schedule/schedule_rule/add_rfactor.cc index 6afd7c3d02db..48db6ccd0033 100644 --- a/src/s_tir/meta_schedule/schedule_rule/add_rfactor.cc +++ b/src/s_tir/meta_schedule/schedule_rule/add_rfactor.cc @@ -114,7 +114,7 @@ ffi::Array AddRFactorObj::Apply(const s_tir::Schedule& sch, // Annotate that the rfactor block, which is now the producer of the original block, needs to // be considered by the rule Random-Compute-Location. - sch_tmp->Annotate(block_rv, tvm::s_tir::attr::kMetaScheduleRandomComputeProducer, 1); + sch_tmp->Annotate(block_rv, tvm::s_tir::attr::kMetaScheduleRandomComputeProducer.name, 1); res.push_back(sch_tmp); } catch (const tvm::ffi::Error& e) { } diff --git a/src/s_tir/meta_schedule/schedule_rule/apply_custom_rule.cc b/src/s_tir/meta_schedule/schedule_rule/apply_custom_rule.cc index 58700635fa6d..2e4ff04b41b1 100644 --- a/src/s_tir/meta_schedule/schedule_rule/apply_custom_rule.cc +++ b/src/s_tir/meta_schedule/schedule_rule/apply_custom_rule.cc @@ -45,8 +45,8 @@ class ApplyCustomRuleObj : public ScheduleRuleObj { TVM_FFI_CHECK(this->target_.has_value(), ValueError) << "ApplyCustomRule is not initialized with TuneContext that has a Target."; ffi::Array keys = this->target_.value()->keys; - if (ffi::Optional ann = - s_tir::GetAnn(sch->GetSRef(block_rv), tvm::s_tir::attr::kScheduleRule)) { + if (ffi::Optional ann = s_tir::GetAnn( + sch->GetSRef(block_rv), tvm::s_tir::attr::kScheduleRule.name)) { if (ann.value() != "None") { for (const ffi::String& key : keys) { if (const auto custom_schedule_fn = @@ -63,7 +63,7 @@ class ApplyCustomRuleObj : public ScheduleRuleObj { os << "\n " << GetCustomRuleName(ann.value(), key); } LOG(WARNING) << os.str(); - sch->Unannotate(block_rv, tvm::s_tir::attr::kScheduleRule); + sch->Unannotate(block_rv, tvm::s_tir::attr::kScheduleRule.name); } } return {sch}; diff --git a/src/s_tir/meta_schedule/schedule_rule/auto_inline.cc b/src/s_tir/meta_schedule/schedule_rule/auto_inline.cc index 1ce5578d6fa2..05b054c75b02 100644 --- a/src/s_tir/meta_schedule/schedule_rule/auto_inline.cc +++ b/src/s_tir/meta_schedule/schedule_rule/auto_inline.cc @@ -161,7 +161,7 @@ inline InlineType AutoInlineObj::CheckInline(const s_tir::Schedule& sch, } // Cond 6. The block is disallowed for auto inline if (ffi::Optional ann = - s_tir::GetAnn(block_sref, tvm::s_tir::attr::kMetaScheduleInlineRule)) { + s_tir::GetAnn(block_sref, tvm::s_tir::attr::kMetaScheduleInlineRule.name)) { if (ann.value() == "disable") return InlineType::kNoInline; } // Last cond: Check inline into the consumers or the spatial producer @@ -189,7 +189,7 @@ inline InlineType AutoInlineObj::CheckInline(const s_tir::Schedule& sch, if (producer_srefs.size() == 1 && s_tir::IsCompleteBlock(sch->state(), producer_srefs[0], scope_block) && CanReverseComputeInline(state, block_sref) && - !GetAnn(producer_srefs[0], tvm::s_tir::attr::kMetaScheduleAutoTensorize) + !GetAnn(producer_srefs[0], tvm::s_tir::attr::kMetaScheduleAutoTensorize.name) .has_value()) { return InlineType::kInlineIntoProducer; } diff --git a/src/s_tir/meta_schedule/schedule_rule/multi_level_tiling.cc b/src/s_tir/meta_schedule/schedule_rule/multi_level_tiling.cc index 8e8f41ade948..cf496b21cbec 100644 --- a/src/s_tir/meta_schedule/schedule_rule/multi_level_tiling.cc +++ b/src/s_tir/meta_schedule/schedule_rule/multi_level_tiling.cc @@ -112,7 +112,7 @@ void MultiLevelTilingObj::InitializeWithTuneContext(const TuneContext& context) ffi::Array MultiLevelTilingObj::Apply(const Schedule& sch, const SBlockRV& block_rv) { if ((filter_fn_ && filter_fn_.value()(sch, sch->GetSRef(block_rv)).cast()) || NeedsMultiLevelTiling(sch->state(), sch->GetSRef(block_rv))) { - sch->Annotate(block_rv, tvm::s_tir::attr::kMetaScheduleTilingStructure, structure); + sch->Annotate(block_rv, tvm::s_tir::attr::kMetaScheduleTilingStructure.name, structure); ffi::Array results; for (auto&& state : ApplySubRules({State(sch, block_rv)})) { @@ -145,8 +145,9 @@ std::vector MultiLevelTilingObj::AddWriteReuse(State state) const { } std::vector levels = config.levels; ReuseType req = config.req; - if (ffi::Optional> ann = s_tir::GetAnn>( - state->sch->GetSRef(state->block_rv), tvm::s_tir::attr::kMetaScheduleWriteCacheLevel)) { + if (ffi::Optional> ann = + s_tir::GetAnn>(state->sch->GetSRef(state->block_rv), + tvm::s_tir::attr::kMetaScheduleWriteCacheLevel.name)) { req = ReuseType::kMustReuse; levels.clear(); std::transform(ann.value().begin(), ann.value().end(), std::back_inserter(levels), @@ -285,8 +286,9 @@ std::vector MultiLevelTilingObj::TileLoopNest(State state, if (spatial_loop_product > 2 * this->thread_warp_size_) { low_inclusive = this->thread_warp_size_; } - sch->Annotate(block_rv, tvm::s_tir::attr::kMetaScheduleThreadExtentLowInclusive, low_inclusive); - sch->Annotate(block_rv, tvm::s_tir::attr::kMetaScheduleThreadExtentHighInclusive, + sch->Annotate(block_rv, tvm::s_tir::attr::kMetaScheduleThreadExtentLowInclusive.name, + low_inclusive); + sch->Annotate(block_rv, tvm::s_tir::attr::kMetaScheduleThreadExtentHighInclusive.name, high_inclusive); } return {state}; @@ -354,11 +356,11 @@ std::vector MultiLevelTilingObj::AddAsyncPipeline(State state) const { for (int stage : this->stages) { State new_state = state->Copy(); LoopRV r_loop_fused = new_state->sch->Fuse(new_state->tiles[r_indices_[0]]); - new_state->sch->Annotate(r_loop_fused, tvm::s_tir::attr::kSoftwarePipelineStage, + new_state->sch->Annotate(r_loop_fused, tvm::s_tir::attr::kSoftwarePipelineStage.name, ffi::Array{0, 0, stage - 2}); - new_state->sch->Annotate(r_loop_fused, tvm::s_tir::attr::kSoftwarePipelineOrder, + new_state->sch->Annotate(r_loop_fused, tvm::s_tir::attr::kSoftwarePipelineOrder.name, ffi::Array{0, 1, 2}); - new_state->sch->Annotate(r_loop_fused, tvm::s_tir::attr::kSoftwarePipelineAsyncStages, + new_state->sch->Annotate(r_loop_fused, tvm::s_tir::attr::kSoftwarePipelineAsyncStages.name, ffi::Array{0}); ret.push_back(std::move(new_state)); } @@ -399,7 +401,7 @@ void MultiLevelTilingObj::AnnotateCooperativeFetching(Schedule* sch, for (int v : valid_vector_lens) valid_vector_lens_arr.push_back(static_cast(v)); s_tir::ExprRV vector_load_len = (*sch)->SampleCategorical( valid_vector_lens_arr, ffi::Array(n, FloatImm(PrimType::Float(32), prob))); - (*sch)->Annotate(block, tvm::s_tir::attr::kMetaScheduleCooperativeFetch, vector_load_len); + (*sch)->Annotate(block, tvm::s_tir::attr::kMetaScheduleCooperativeFetch.name, vector_load_len); } } diff --git a/src/s_tir/meta_schedule/schedule_rule/multi_level_tiling_with_intrin.cc b/src/s_tir/meta_schedule/schedule_rule/multi_level_tiling_with_intrin.cc index f022c53f4dcb..62aa15aa76f8 100644 --- a/src/s_tir/meta_schedule/schedule_rule/multi_level_tiling_with_intrin.cc +++ b/src/s_tir/meta_schedule/schedule_rule/multi_level_tiling_with_intrin.cc @@ -42,7 +42,7 @@ ffi::Optional TileForIntrin(s_tir::Schedule sch, s_tir::SBlockR } TVM_FFI_ICHECK(tiled_loop_rv.has_value()); s_tir::SBlockRV outer_block = sch->Blockize(tiled_loop_rv.value()); - sch->Annotate(outer_block, tvm::s_tir::attr::kMetaScheduleAutoTensorize, + sch->Annotate(outer_block, tvm::s_tir::attr::kMetaScheduleAutoTensorize.name, ffi::String(intrin_name)); return outer_block; } diff --git a/src/s_tir/meta_schedule/schedule_rule/parallel_vectorize_unroll.cc b/src/s_tir/meta_schedule/schedule_rule/parallel_vectorize_unroll.cc index 17cfa3c7f9b9..830d87da7433 100644 --- a/src/s_tir/meta_schedule/schedule_rule/parallel_vectorize_unroll.cc +++ b/src/s_tir/meta_schedule/schedule_rule/parallel_vectorize_unroll.cc @@ -63,11 +63,12 @@ class ParallelizeVectorizeUnrollObj : public ScheduleRuleObj { // Parallelization if (max_jobs_per_core != -1) { - sch->Annotate(root_rv, tvm::s_tir::attr::kMetaScheduleParallel, this->max_parallel_extent_); + sch->Annotate(root_rv, tvm::s_tir::attr::kMetaScheduleParallel.name, + this->max_parallel_extent_); } // Vectorization if (max_vectorize_extent != -1) { - sch->Annotate(root_rv, tvm::s_tir::attr::kMetaScheduleVectorize, max_vectorize_extent); + sch->Annotate(root_rv, tvm::s_tir::attr::kMetaScheduleVectorize.name, max_vectorize_extent); } // Unroll if (!unroll_max_steps.empty() && !s_tir::CheckSpatialFunction(sch, root_rv)) { @@ -76,9 +77,9 @@ class ParallelizeVectorizeUnrollObj : public ScheduleRuleObj { ffi::Array probs(n, FloatImm(PrimType::Float(32), prob)); PrimExpr max_step = sch->SampleCategorical(unroll_max_steps, probs); if (unroll_explicit) { - sch->Annotate(root_rv, tvm::s_tir::attr::kMetaScheduleUnrollExplicit, max_step); + sch->Annotate(root_rv, tvm::s_tir::attr::kMetaScheduleUnrollExplicit.name, max_step); } else { - sch->Annotate(root_rv, tvm::s_tir::attr::kMetaScheduleUnrollImplicit, max_step); + sch->Annotate(root_rv, tvm::s_tir::attr::kMetaScheduleUnrollImplicit.name, max_step); } } return {sch}; diff --git a/src/s_tir/meta_schedule/schedule_rule/random_compute_location.cc b/src/s_tir/meta_schedule/schedule_rule/random_compute_location.cc index 52c9541c1401..0ffc90528d31 100644 --- a/src/s_tir/meta_schedule/schedule_rule/random_compute_location.cc +++ b/src/s_tir/meta_schedule/schedule_rule/random_compute_location.cc @@ -44,10 +44,10 @@ class RandomComputeLocationObj : public ScheduleRuleObj { // access the input block. Hence we collect its producer ahead of time. // - Note that only single producer is allowed in this case. ffi::Array producers; - if (s_tir::HasAnn(sch->GetSRef(block_rv), tvm::s_tir::attr::kMetaScheduleRandomComputeProducer, - true)) { + if (s_tir::HasAnn(sch->GetSRef(block_rv), + tvm::s_tir::attr::kMetaScheduleRandomComputeProducer.name, true)) { producers = sch->GetProducers(block_rv); - sch->Unannotate(block_rv, tvm::s_tir::attr::kMetaScheduleRandomComputeProducer); + sch->Unannotate(block_rv, tvm::s_tir::attr::kMetaScheduleRandomComputeProducer.name); TVM_FFI_ICHECK_EQ(producers.size(), 1); } diff --git a/src/s_tir/meta_schedule/space_generator/post_order_apply.cc b/src/s_tir/meta_schedule/space_generator/post_order_apply.cc index f10f256ff402..700312e744f4 100644 --- a/src/s_tir/meta_schedule/space_generator/post_order_apply.cc +++ b/src/s_tir/meta_schedule/space_generator/post_order_apply.cc @@ -83,7 +83,8 @@ class PostOrderApplyObj : public SpaceGeneratorObj { continue; } if (!ScheduleRule::IsApplyCustomRule(sch_rule)) { - if (s_tir::GetAnn(sch->GetSRef(block_rv), tvm::s_tir::attr::kScheduleRule) + if (s_tir::GetAnn(sch->GetSRef(block_rv), + tvm::s_tir::attr::kScheduleRule.name) .has_value()) { stack.emplace_back(sch, blocks); continue; diff --git a/src/s_tir/schedule/analysis/analysis.cc b/src/s_tir/schedule/analysis/analysis.cc index 8a8a6ed70634..38b0c077f44a 100644 --- a/src/s_tir/schedule/analysis/analysis.cc +++ b/src/s_tir/schedule/analysis/analysis.cc @@ -997,7 +997,7 @@ StmtSRef GetSRefLowestCommonAncestor(const ffi::Array& srefs) { } bool HasBeenMultiLevelTiled(const StmtSRef& block_sref) { - return GetAnn(block_sref, tvm::s_tir::attr::kMetaScheduleTilingStructure) + return GetAnn(block_sref, tvm::s_tir::attr::kMetaScheduleTilingStructure.name) .has_value(); } diff --git a/src/s_tir/schedule/primitive/annotate.cc b/src/s_tir/schedule/primitive/annotate.cc index 3aa8dc1679d3..221b4463b1a8 100644 --- a/src/s_tir/schedule/primitive/annotate.cc +++ b/src/s_tir/schedule/primitive/annotate.cc @@ -31,7 +31,7 @@ void Annotate(ScheduleState self, const StmtSRef& sref, const ffi::String& ann_k // Extract annotation const ffi::Map* annotations = nullptr; if (const auto* loop = sref->StmtAs()) { - TVM_FFI_CHECK(ann_key != tvm::s_tir::attr::kThreadBinding, ValueError) + TVM_FFI_CHECK(ann_key != tvm::s_tir::attr::kThreadBinding.name, ValueError) << "Use loop scheduling primitives to change thread binding"; annotations = &loop->annotations; } else if (const auto* block = sref->StmtAs()) { @@ -66,7 +66,7 @@ void Unannotate(ScheduleState self, const StmtSRef& sref, const ffi::String& ann // Extract annotation const ffi::Map* annotations = nullptr; if (const auto* loop = sref->StmtAs()) { - TVM_FFI_CHECK(ann_key != tvm::s_tir::attr::kThreadBinding, ValueError) + TVM_FFI_CHECK(ann_key != tvm::s_tir::attr::kThreadBinding.name, ValueError) << "Use loop scheduling primitives to change thread binding"; annotations = &loop->annotations; } else if (const auto* block = sref->StmtAs()) { diff --git a/src/s_tir/schedule/primitive/annotate_tensor_access.cc b/src/s_tir/schedule/primitive/annotate_tensor_access.cc index e39528097441..663456379303 100644 --- a/src/s_tir/schedule/primitive/annotate_tensor_access.cc +++ b/src/s_tir/schedule/primitive/annotate_tensor_access.cc @@ -59,8 +59,8 @@ class AnnotateRegionRewriter : public StmtExprMutator { // Annotate the block with explicit_read_region or explicit_write_region ffi::Map new_annotations = n->annotations; ffi::String annotation_key = tensor_index_type_ == TensorIndexType::kWrite - ? s_tir::attr::kExplicitWriteRegion - : s_tir::attr::kExplicitReadRegion; + ? s_tir::attr::kExplicitWriteRegion.name + : s_tir::attr::kExplicitReadRegion.name; if (new_annotations.count(annotation_key)) { ffi::Array tensor_indices = new_annotations[annotation_key].as_or_throw>(); diff --git a/src/s_tir/schedule/primitive/block_annotate.cc b/src/s_tir/schedule/primitive/block_annotate.cc index 1058caa3d90f..0caaea5065eb 100644 --- a/src/s_tir/schedule/primitive/block_annotate.cc +++ b/src/s_tir/schedule/primitive/block_annotate.cc @@ -152,14 +152,14 @@ class StorageAlignInvalidAnnotationError : public ScheduleErrorContextObj { std::ostringstream os; os << "The block annotation for storage align is expected to be an array of 4-integer-tuples " "(tensor_index, axis, factor, offset). However, the block annotation with key " - << s_tir::attr::kTensorDimAlign << " of the block {0} is " - << block_->annotations.at(s_tir::attr::kTensorDimAlign) << ", which is unexpected."; + << s_tir::attr::kTensorDimAlign.name << " of the block {0} is " + << block_->annotations.at(s_tir::attr::kTensorDimAlign.name) << ", which is unexpected."; return os.str(); } static StorageAlignAnnotation CheckAndGetAnnotation(const IRModule& mod, const SBlock& block) { // Get existing annotation value. - auto it = block->annotations.find(s_tir::attr::kTensorDimAlign); + auto it = block->annotations.find(s_tir::attr::kTensorDimAlign.name); if (it != block->annotations.end()) { if (!IsValidAnnotation(block, (*it).second)) { throw MakeScheduleError(mod, block); @@ -261,7 +261,7 @@ void StorageAlign(ScheduleState self, const StmtSRef& block_sref, int tensor_ind // Step 3: Replace the block with the new annotation SBlock new_block = - WithAnnotation(block_ptr, s_tir::attr::kTensorDimAlign, storage_align_annotation); + WithAnnotation(block_ptr, s_tir::attr::kTensorDimAlign.name, storage_align_annotation); self->Replace(block_sref, new_block, {{ffi::GetRef(block_ptr), new_block}}); } diff --git a/src/s_tir/schedule/primitive/compute_inline.cc b/src/s_tir/schedule/primitive/compute_inline.cc index 5babfda443da..bb1aec47cafc 100644 --- a/src/s_tir/schedule/primitive/compute_inline.cc +++ b/src/s_tir/schedule/primitive/compute_inline.cc @@ -782,7 +782,7 @@ class ReverseComputeInliner : public BaseInliner { : nullptr; } else { producer_store = producer_body.as(); - if (producer_block_->annotations.count(tvm::s_tir::attr::kAutoCopy) != 0) { + if (producer_block_->annotations.count(tvm::s_tir::attr::kAutoCopy.name) != 0) { const ForNode* producer_inner_loop = producer_body.as(); if (producer_inner_loop == nullptr) return false; for (;;) { @@ -823,7 +823,7 @@ class ReverseComputeInliner : public BaseInliner { analyzer_->Bind(iter->var, Range::FromMinExtent(iter->dom.value()->min, iter->dom.value()->extent)); } - if (producer_block->annotations.count(tvm::s_tir::attr::kAutoCopy) != 0) { + if (producer_block->annotations.count(tvm::s_tir::attr::kAutoCopy.name) != 0) { auto bind = [&](const ForNode* loop) { analyzer_->Bind(loop->loop_var, Range::FromMinExtent(IntImm(loop->extent.ty(), 0), loop->extent)); diff --git a/src/s_tir/schedule/primitive/for_kind.cc b/src/s_tir/schedule/primitive/for_kind.cc index 3b1c9b73953f..0eb9fbf0821f 100644 --- a/src/s_tir/schedule/primitive/for_kind.cc +++ b/src/s_tir/schedule/primitive/for_kind.cc @@ -183,9 +183,9 @@ void ParallelizeComputation(const ScheduleState& self, const StmtSRef& loop_sref ffi::ObjectPtr new_loop = ffi::make_object(*loop); new_loop->kind = for_kind; if (thread_axis.has_value()) { - new_loop->annotations.Set(tvm::s_tir::attr::kThreadBinding, thread_axis.value()); + new_loop->annotations.Set(tvm::s_tir::attr::kThreadBinding.name, thread_axis.value()); } else { - new_loop->annotations.erase(tvm::s_tir::attr::kThreadBinding); + new_loop->annotations.erase(tvm::s_tir::attr::kThreadBinding.name); } self->Replace(loop_sref, For(new_loop), {}); } @@ -206,7 +206,7 @@ void Unroll(ScheduleState self, const StmtSRef& loop_sref) { const ForNode* loop = TVM_SREF_TO_FOR(loop_sref); ffi::ObjectPtr new_loop = ffi::make_object(*loop); new_loop->kind = ForKind::kUnrolled; - new_loop->annotations.erase(tvm::s_tir::attr::kThreadBinding); + new_loop->annotations.erase(tvm::s_tir::attr::kThreadBinding.name); self->Replace(loop_sref, For(new_loop), {}); } diff --git a/src/s_tir/schedule/primitive/read_write_at.cc b/src/s_tir/schedule/primitive/read_write_at.cc index e1a85b8313ed..b7290285b3fa 100644 --- a/src/s_tir/schedule/primitive/read_write_at.cc +++ b/src/s_tir/schedule/primitive/read_write_at.cc @@ -379,13 +379,13 @@ struct ReadWriteAtImpl { StmtSRef ReadAt(ScheduleState self, const StmtSRef& loop_sref, const StmtSRef& block_sref, int read_tensor_index, const ffi::String& storage_scope) { return ReadWriteAtImpl::Main(self, loop_sref, block_sref, read_tensor_index, storage_scope, - {{s_tir::attr::kAutoCopy, true}}); + {{s_tir::attr::kAutoCopy.name, true}}); } StmtSRef WriteAt(ScheduleState self, const StmtSRef& loop_sref, const StmtSRef& block_sref, int write_tensor_index, const ffi::String& storage_scope) { return ReadWriteAtImpl::Main(self, loop_sref, block_sref, write_tensor_index, - storage_scope, {{s_tir::attr::kAutoCopy, true}}); + storage_scope, {{s_tir::attr::kAutoCopy.name, true}}); } /******** Instruction Registration ********/ diff --git a/src/s_tir/script/ir_builder/frame.cc b/src/s_tir/script/ir_builder/frame.cc index e1f6ae5af824..c0da42aaf5b5 100644 --- a/src/s_tir/script/ir_builder/frame.cc +++ b/src/s_tir/script/ir_builder/frame.cc @@ -34,7 +34,7 @@ TVM_FFI_STATIC_INIT_BLOCK() { SBlockFrameObj::RegisterReflection(); BlockInitFrameObj::RegisterReflection(); tirx::FunctionFrameObj::RegisterAttrValidator( - tvm::attr::kSTir, [](const tirx::FunctionFrameObj* frame, const ffi::Any& value) { + tvm::attr::kSTir.name, [](const tirx::FunctionFrameObj* frame, const ffi::Any& value) { TVM_FFI_CHECK(frame->IsInstance(), ValueError) << "The s_tir attribute requires Ts.function"; bool enabled = false; @@ -50,11 +50,11 @@ TVM_FFI_STATIC_INIT_BLOCK() { } tvm::tirx::Function FunctionFrameObj::FinalizeFunction(tvm::tirx::Function func) { - func = WithAttr(std::move(func), tvm::attr::kScriptNamespace, ffi::String("s_tir")); + func = WithAttr(std::move(func), tvm::attr::kScriptNamespace.name, ffi::String("s_tir")); TVM_FFI_CHECK(!is_declaration || root_alloc_tensors.empty(), ValueError) << "A function declaration cannot allocate tensors"; if (!is_declaration) { - func = WithAttr(std::move(func), tvm::attr::kSTir, true); + func = WithAttr(std::move(func), tvm::attr::kSTir.name, true); func = tvm::s_tir::ScriptComplete(std::move(func), root_alloc_tensors, root_allocated_addresses); } @@ -91,12 +91,12 @@ void SBlockFrameObj::ExitWithScope() { } ffi::Map attrs = annotations.value_or({}); if (!allocated_addresses.empty()) { - TVM_FFI_CHECK(!attrs.count(tvm::s_tir::attr::kTensorAllocatedAddr), ValueError) + TVM_FFI_CHECK(!attrs.count(tvm::s_tir::attr::kTensorAllocatedAddr.name), ValueError) << "Tensor placement must be specified on its allocation or match_tensor"; - attrs.Set(tvm::s_tir::attr::kTensorAllocatedAddr, allocated_addresses); + attrs.Set(tvm::s_tir::attr::kTensorAllocatedAddr.name, allocated_addresses); } if (int detect_access = (!reads.has_value()) | (!writes.has_value() << 1)) { - attrs.Set(tvm::s_tir::attr::kScriptParsingDetectAccess, detect_access); + attrs.Set(tvm::s_tir::attr::kScriptParsingDetectAccess.name, detect_access); } tvm::s_tir::SBlock block(iter_vars, reads.value_or(ffi::Array()), writes.value_or(ffi::Array()), name, diff --git a/src/s_tir/script/ir_builder/script_complete.cc b/src/s_tir/script/ir_builder/script_complete.cc index 4f4975173f69..b5ac212cefd2 100644 --- a/src/s_tir/script/ir_builder/script_complete.cc +++ b/src/s_tir/script/ir_builder/script_complete.cc @@ -96,7 +96,7 @@ class ScriptCompleter : public s_tir::StmtExprMutator { // Get access detection mask // 0 for provided region, 1 and 3 for need detect read, 2 and 3 for need detect write int mask = 0; - auto it = op->annotations.find(tvm::s_tir::attr::kScriptParsingDetectAccess); + auto it = op->annotations.find(tvm::s_tir::attr::kScriptParsingDetectAccess.name); if (it != op->annotations.end()) { mask = (*it).second.cast(); } @@ -115,7 +115,7 @@ class ScriptCompleter : public s_tir::StmtExprMutator { if (mask & 2) n->writes = writes; } n->annotations = op->annotations; - n->annotations.erase(tvm::s_tir::attr::kScriptParsingDetectAccess); + n->annotations.erase(tvm::s_tir::attr::kScriptParsingDetectAccess.name); return block; } else { return block; @@ -138,7 +138,7 @@ class ScriptCompleter : public s_tir::StmtExprMutator { Function ScriptComplete(Function func, const ffi::Array& root_allocates, const TensorAllocatedAddresses& root_addresses) { - func = WithAttr(std::move(func), tvm::attr::kScriptNamespace, ffi::String("s_tir")); + func = WithAttr(std::move(func), tvm::attr::kScriptNamespace.name, ffi::String("s_tir")); if (!func->body.has_value()) return func; ffi::Map tensor_var_map; for (const Var& param : func->params) { @@ -174,7 +174,7 @@ Function ScriptComplete(Function func, const ffi::Array& root_allocat if (should_insert_root) { ffi::Map annotations; if (!root_addresses.empty()) { - annotations.Set(s_tir::attr::kTensorAllocatedAddr, root_addresses); + annotations.Set(s_tir::attr::kTensorAllocatedAddr.name, root_addresses); } s_tir::SBlock root_block({}, {}, {}, "root", std::move(res), std::nullopt, root_allocates, {}, annotations); diff --git a/src/s_tir/script/printer/expr.cc b/src/s_tir/script/printer/expr.cc index 8fa1d82f0b97..81ea8f3727ae 100644 --- a/src/s_tir/script/printer/expr.cc +++ b/src/s_tir/script/printer/expr.cc @@ -47,7 +47,7 @@ ffi::Optional IterVarDocTranslate(DocTranslatorObj* d, ffi::AnyView inp TVM_FFI_STATIC_INIT_BLOCK() { ffi::reflection::TypeAttrDef().attr( - tvm::script::printer::type_attr::kDocTranslate, + tvm::script::printer::type_attr::kDocTranslate.name, FDocTranslate::FromNative<&IterVarDocTranslate>()); } diff --git a/src/s_tir/script/printer/stmt.cc b/src/s_tir/script/printer/stmt.cc index 8c4aa89a39be..1846097f17c1 100644 --- a/src/s_tir/script/printer/stmt.cc +++ b/src/s_tir/script/printer/stmt.cc @@ -131,11 +131,11 @@ ffi::Array SBlockBody(DocTranslatorObj* d, const s_tir::SBlockNode* blo d->Emit(ExprStmtDoc(NamespaceDoc("s_tir")->Attr("writes")->Call(writes)), block->writes); auto annotations = block->annotations; ffi::Map> addresses; - if (auto value = annotations.Get(s_tir::attr::kTensorAllocatedAddr)) { + if (auto value = annotations.Get(s_tir::attr::kTensorAllocatedAddr.name)) { for (const auto& entry : value.value().cast()) { addresses.Set(entry.get<0>(), entry.get<1>()); } - annotations.erase(s_tir::attr::kTensorAllocatedAddr); + annotations.erase(s_tir::attr::kTensorAllocatedAddr.name); } if (!annotations.empty()) { d->Emit(ExprStmtDoc(NamespaceDoc("s_tir") @@ -194,7 +194,7 @@ ffi::Optional SBlockRealizeDocTranslate(DocTranslatorObj* d, ffi::AnyVi TVM_FFI_STATIC_INIT_BLOCK() { ffi::reflection::TypeAttrDef().attr( - tvm::script::printer::type_attr::kDocTranslate, + tvm::script::printer::type_attr::kDocTranslate.name, FDocTranslate::FromNative<&SBlockRealizeDocTranslate>()); } @@ -215,7 +215,7 @@ ffi::Optional SBlockDocTranslate(DocTranslatorObj* d, ffi::AnyView inpu TVM_FFI_STATIC_INIT_BLOCK() { ffi::reflection::TypeAttrDef().attr( - tvm::script::printer::type_attr::kDocTranslate, + tvm::script::printer::type_attr::kDocTranslate.name, FDocTranslate::FromNative<&SBlockDocTranslate>()); } @@ -231,7 +231,7 @@ ffi::Optional MatchTensorRegionDocTranslate(DocTranslatorObj* d, ffi::A TVM_FFI_STATIC_INIT_BLOCK() { ffi::reflection::TypeAttrDef().attr( - tvm::script::printer::type_attr::kDocTranslate, + tvm::script::printer::type_attr::kDocTranslate.name, FDocTranslate::FromNative<&MatchTensorRegionDocTranslate>()); } diff --git a/src/s_tir/stmt.cc b/src/s_tir/stmt.cc index f8d1d50becc4..0b196a9c0f47 100644 --- a/src/s_tir/stmt.cc +++ b/src/s_tir/stmt.cc @@ -30,7 +30,7 @@ #include #include #include -#include +#include namespace tvm { namespace s_tir { diff --git a/src/s_tir/stmt_functor.cc b/src/s_tir/stmt_functor.cc index 46149d3c5048..9589ad31fd44 100644 --- a/src/s_tir/stmt_functor.cc +++ b/src/s_tir/stmt_functor.cc @@ -64,7 +64,7 @@ ffi::Optional StmtExprVisitor::VisitBlock(tirx::StmtExprVisitor* kTVMFFIDefRegionKindSimple, [&]() { return visitor->Visit(match_tensor_region->tensor); })); TVM_FFI_S_VISIT_MAYBE_EARLY_RETURN(visitor->Visit(match_tensor_region->source)); } - if (auto addresses = op->annotations.Get(attr::kTensorAllocatedAddr)) { + if (auto addresses = op->annotations.Get(attr::kTensorAllocatedAddr.name)) { // The owner reference identifies the allocation; it is not a data access. for (const auto& entry : addresses.value().cast()) { TVM_FFI_S_VISIT_MAYBE_EARLY_RETURN(visitor->Visit(entry.get<1>())); @@ -139,10 +139,10 @@ UnchangedOr StmtExprMutator::MutateBlock(tirx::StmtExprMutator* mutator, c auto match_tensors = mutator->Mutate(op->match_tensors, inplace_mode) .as_or_throw>>(); auto annotations = op->annotations; - if (auto addresses = annotations.Get(attr::kTensorAllocatedAddr)) { + if (auto addresses = annotations.Get(attr::kTensorAllocatedAddr.name)) { auto updated = mutator->Mutate(addresses.value(), InplaceMode::kDisallow); if (!updated.IsUnchanged()) { - annotations.Set(attr::kTensorAllocatedAddr, std::move(updated).ValueUnchecked()); + annotations.Set(attr::kTensorAllocatedAddr.name, std::move(updated).ValueUnchecked()); } } auto reads = diff --git a/src/s_tir/transform/annotate_irregular_loop.cc b/src/s_tir/transform/annotate_irregular_loop.cc index 6693e227da12..a8c315b4d549 100644 --- a/src/s_tir/transform/annotate_irregular_loop.cc +++ b/src/s_tir/transform/annotate_irregular_loop.cc @@ -58,12 +58,13 @@ class IrregularLoopAnnotator : public StmtExprMutator { TVM_FFI_ICHECK(op->kind == ForKind::kDefault) << "Loop kind " << op->kind << " is invalid for irregular loop " << op->loop_var; for (const char* key : - {tvm::tirx::attr::kAutoUnrollMaxStep, tvm::tirx::attr::kUnrollExplicit, - tvm::s_tir::attr::kLoopPartitionHint, tvm::s_tir::attr::kSoftwarePipelineStage}) { + {tvm::tirx::attr::kAutoUnrollMaxStep.name, tvm::tirx::attr::kUnrollExplicit.name, + tvm::s_tir::attr::kLoopPartitionHint.name, + tvm::s_tir::attr::kSoftwarePipelineStage.name}) { TVM_FFI_ICHECK(!res->annotations.count(key)) << "Annotation `" << key << "` is invalid for irregular loop " << op->loop_var; } - res.CopyOnWrite()->annotations.Set(tvm::s_tir::attr::kIrregularLoopMark, 1); + res.CopyOnWrite()->annotations.Set(tvm::s_tir::attr::kIrregularLoopMark.name, 1); } std::swap(cur_has_jump, has_jump_); return res; diff --git a/src/s_tir/transform/compact_tensor_region.cc b/src/s_tir/transform/compact_tensor_region.cc index 48a56fa7364a..6ef4db27bbf7 100644 --- a/src/s_tir/transform/compact_tensor_region.cc +++ b/src/s_tir/transform/compact_tensor_region.cc @@ -315,8 +315,8 @@ class TensorAccessRegionCollector : public StmtExprVisitor { } }; - record_explicit_region(s_tir::attr::kExplicitReadRegion, TensorIndexType::kRead); - record_explicit_region(s_tir::attr::kExplicitWriteRegion, TensorIndexType::kWrite); + record_explicit_region(s_tir::attr::kExplicitReadRegion.name, TensorIndexType::kRead); + record_explicit_region(s_tir::attr::kExplicitWriteRegion.name, TensorIndexType::kWrite); // Step 3. Record relax position of ancestor_loops_ for (const TensorVar& tensor : op->alloc_tensors) { diff --git a/src/s_tir/transform/default_gpu_schedule.cc b/src/s_tir/transform/default_gpu_schedule.cc index 4b4efdce162e..bcb0fbf69496 100644 --- a/src/s_tir/transform/default_gpu_schedule.cc +++ b/src/s_tir/transform/default_gpu_schedule.cc @@ -95,7 +95,7 @@ IRModule MarkScheduled(const IRModule& mod) { if (const auto* function_node = base_func.as()) { tirx::Function function = ffi::GetRef(function_node); tirx::Function new_function = - WithAttr(std::move(function), tvm::s_tir::attr::kIsScheduled, true); + WithAttr(std::move(function), tvm::s_tir::attr::kIsScheduled.name, true); result.Set(gv, new_function); } else { result.Set(gv, base_func); @@ -164,7 +164,8 @@ bool IsScheduledOnGPU(const BaseFunc& func) { // the target from context. ffi::Optional target = tvm::Target::Current(); // the Target in kTarget attribute of Function - ffi::Optional func_target = func->attrs.GetAttr(tvm::attr::kTarget); + ffi::Optional func_target = + func->attrs.GetAttr(tvm::attr::kTarget.name); if (func_target.has_value()) { target = func_target.value(); } @@ -189,7 +190,7 @@ Pass DefaultGPUSchedule() { for (const auto& [gv, base_func] : m->functions) { if (const auto* function_node = base_func.as(); function_node != nullptr && IsScheduledOnGPU(base_func) && - !base_func->HasNonzeroAttr(tvm::s_tir::attr::kIsScheduled)) { + !base_func->HasNonzeroAttr(tvm::s_tir::attr::kIsScheduled.name)) { tirx::Function func = ffi::GetRef(function_node); tirx::Function new_func = WrapBareSBlockBody(func); if (!new_func.same_as(func)) { @@ -207,12 +208,13 @@ Pass DefaultGPUSchedule() { s_tir::ScheduleErrorRenderLevel::kDetail); for (const auto& [gv, func] : m->functions) { if (func->IsInstance() && - !func->HasNonzeroAttr(tvm::s_tir::attr::kIsScheduled) && IsScheduledOnGPU(func)) { + !func->HasNonzeroAttr(tvm::s_tir::attr::kIsScheduled.name) && + IsScheduledOnGPU(func)) { // get the target from context. ffi::Optional target = tvm::Target::Current(); // get the target from kTarget attribute ffi::Optional func_target = - func->attrs.GetAttr(tvm::attr::kTarget); + func->attrs.GetAttr(tvm::attr::kTarget.name); if (func_target.has_value()) { target = func_target.value(); } diff --git a/src/s_tir/transform/inject_software_pipeline.cc b/src/s_tir/transform/inject_software_pipeline.cc index 090da95f51dd..3f4cd6e7c338 100644 --- a/src/s_tir/transform/inject_software_pipeline.cc +++ b/src/s_tir/transform/inject_software_pipeline.cc @@ -1078,7 +1078,7 @@ class PipelineInjector : public StmtExprMutator { using StmtExprMutator::Mutate_; static Stmt Inject(const Function& func) { - auto global_symbol = func->GetAttr(tvm::attr::kGlobalSymbol); + auto global_symbol = func->GetAttr(tvm::attr::kGlobalSymbol); auto injector = ffi::make_object(global_symbol); for (const Var& param : func->params) { if (auto tensor = param.as()) { @@ -1194,9 +1194,9 @@ class PipelineInjector : public StmtExprMutator { } } - auto pipeline_stages = op->annotations.at(tvm::s_tir::attr::kSoftwarePipelineStage) + auto pipeline_stages = op->annotations.at(tvm::s_tir::attr::kSoftwarePipelineStage.name) .as_or_throw>(); - auto pipeline_orders = op->annotations.at(tvm::s_tir::attr::kSoftwarePipelineOrder) + auto pipeline_orders = op->annotations.at(tvm::s_tir::attr::kSoftwarePipelineOrder.name) .as_or_throw>(); TVM_FFI_ICHECK_EQ(pipeline_stages.size(), original_order.size()) << "Function " << global_symbol_ << " has original order " @@ -1208,7 +1208,7 @@ class PipelineInjector : public StmtExprMutator { << ", but pipeline annotation is " << pipeline_orders << " with different size"; std::unordered_set pipeline_async_stages; - if (auto annot = op->annotations.Get(tvm::s_tir::attr::kSoftwarePipelineAsyncStages)) { + if (auto annot = op->annotations.Get(tvm::s_tir::attr::kSoftwarePipelineAsyncStages.name)) { for (int64_t s : annot.value().as_or_throw>()) { pipeline_async_stages.insert(static_cast(s)); } @@ -1217,9 +1217,9 @@ class PipelineInjector : public StmtExprMutator { ffi::Map preserved_annotations; for (const auto& kv : op->annotations) { const ffi::String& key = kv.first; - if (kv.first != tvm::s_tir::attr::kSoftwarePipelineStage && - kv.first != tvm::s_tir::attr::kSoftwarePipelineOrder && - kv.first != tvm::s_tir::attr::kSoftwarePipelineAsyncStages) { + if (kv.first != tvm::s_tir::attr::kSoftwarePipelineStage.name && + kv.first != tvm::s_tir::attr::kSoftwarePipelineOrder.name && + kv.first != tvm::s_tir::attr::kSoftwarePipelineAsyncStages.name) { preserved_annotations.Set(key, kv.second); } } @@ -1271,7 +1271,7 @@ class PipelineInjector : public StmtExprMutator { tensor_data_to_tensor_.Set(tensor, tensor); } - auto it = op->annotations.find(tvm::s_tir::attr::kDoubleBufferScope); + auto it = op->annotations.find(tvm::s_tir::attr::kDoubleBufferScope.name); if (it != op->annotations.end()) { int tensor_index = (*it).second.cast()->value.as().value(); TVM_FFI_CHECK(tensor_index >= 0 && static_cast(tensor_index) < op->writes.size(), @@ -1291,8 +1291,8 @@ class PipelineInjector : public StmtExprMutator { } bool HasPipelineAnnotation(const ForNode* op) const { - auto it1 = op->annotations.find(tvm::s_tir::attr::kSoftwarePipelineStage); - auto it2 = op->annotations.find(tvm::s_tir::attr::kSoftwarePipelineOrder); + auto it1 = op->annotations.find(tvm::s_tir::attr::kSoftwarePipelineStage.name); + auto it2 = op->annotations.find(tvm::s_tir::attr::kSoftwarePipelineOrder.name); bool has_stage = it1 != op->annotations.end(); bool has_order = it2 != op->annotations.end(); if (has_stage && has_order) { diff --git a/src/s_tir/transform/ir_utils.cc b/src/s_tir/transform/ir_utils.cc index e7899c21d8a1..b7086544c379 100644 --- a/src/s_tir/transform/ir_utils.cc +++ b/src/s_tir/transform/ir_utils.cc @@ -359,7 +359,7 @@ class StorageAlignCollector : public StmtExprVisitor { /*! \brief SBlock: resolve each annotation's tensor index through the write regions. */ ffi::Optional Visit_(const SBlockNode* op) final { - auto it = op->annotations.find(attr::kTensorDimAlign); + auto it = op->annotations.find(attr::kTensorDimAlign.name); if (it != op->annotations.end()) { auto annotation = (*it).second.as_or_throw(); for (const auto& item : annotation) { @@ -380,7 +380,7 @@ class StorageAlignCollector : public StmtExprVisitor { ffi::Optional DispatchAllocTensor(const BindNode* op, const CallNode* call) { DictAttrs annotations = call->attrs.as_or_throw(); - auto it = annotations->dict.find(attr::kTensorDimAlign); + auto it = annotations->dict.find(attr::kTensorDimAlign.name); if (it != annotations->dict.end()) { auto storage_align_annotation = (*it).second.as_or_throw(); for (const auto& storage_align_tuple : storage_align_annotation) { diff --git a/src/s_tir/transform/lift_thread_binding.cc b/src/s_tir/transform/lift_thread_binding.cc index b1ecdc66b54c..4d9106e41f44 100644 --- a/src/s_tir/transform/lift_thread_binding.cc +++ b/src/s_tir/transform/lift_thread_binding.cc @@ -137,7 +137,7 @@ class ThreadBindingLifter : public StmtExprMutator { if (auto it = iter_lca.find(op); it != iter_lca.end()) { for (const auto& [iter_var, annotation] : it->second) { auto annotations = annotation; - annotations.Set(tvm::s_tir::attr::kThreadBinding, iter_var->thread_tag); + annotations.Set(tvm::s_tir::attr::kThreadBinding.name, iter_var->thread_tag); body = For(iter_var->var, iter_var->dom.value()->min, iter_var->dom.value()->extent, ForKind::kParallel, std::move(body), std::move(annotations), std::nullopt); } diff --git a/src/s_tir/transform/loop_partition.cc b/src/s_tir/transform/loop_partition.cc index 9921b5406062..4cc6ec81dd25 100644 --- a/src/s_tir/transform/loop_partition.cc +++ b/src/s_tir/transform/loop_partition.cc @@ -118,7 +118,7 @@ static bool IsVirtualThread(const ForNode* op) { } static bool HasPartitionHint(const ForNode* op) { - auto it = op->annotations.find(tvm::s_tir::attr::kLoopPartitionHint); + auto it = op->annotations.find(tvm::s_tir::attr::kLoopPartitionHint.name); return it != op->annotations.end() && IsOne((*it).second.as_or_throw()); } @@ -133,14 +133,14 @@ class NormalizePartitionHints : public StmtExprMutator { For loop = StmtExprMutator::Mutate_(op, inplace_mode) .ValueOrUnchanged(ffi::GetRef(op)) .as_or_throw(); - auto it = loop->annotations.find(tvm::s_tir::attr::kLoopPartitionHint); + auto it = loop->annotations.find(tvm::s_tir::attr::kLoopPartitionHint.name); if (it != loop->annotations.end()) { auto value = (*it).second; if (value == nullptr) { - loop.CopyOnWrite()->annotations.erase(tvm::s_tir::attr::kLoopPartitionHint); + loop.CopyOnWrite()->annotations.erase(tvm::s_tir::attr::kLoopPartitionHint.name); } else { bool enabled = analyzer_->CanProve(value.cast()); - loop.CopyOnWrite()->annotations.Set(tvm::s_tir::attr::kLoopPartitionHint, + loop.CopyOnWrite()->annotations.Set(tvm::s_tir::attr::kLoopPartitionHint.name, IntImm::Bool(enabled)); } } @@ -917,7 +917,7 @@ inline Stmt LoopPartitioner::MakeFor(const ffi::Object* node, PrimExpr extent, S if (analyzer_->CanProve(extent == IntImm::Int32(1)) && !no_unroll_loop_with_extent_one_ && for_node->annotations.size() == - for_node->annotations.count(tvm::s_tir::attr::kLoopPartitionHint)) { + for_node->annotations.count(tvm::s_tir::attr::kLoopPartitionHint.name)) { // If the loop extent is 1, do not create the loop anymore auto f_substitute = [loop_var = for_node->loop_var]( const Var& var) -> ffi::Expected> { @@ -956,8 +956,8 @@ class RemoveLikelyTagsAndHints : public StmtExprMutator { For loop = StmtExprMutator::Mutate_(op, inplace_mode) .ValueOrUnchanged(ffi::GetRef(op)) .as_or_throw(); - if (loop->annotations.count(tvm::s_tir::attr::kLoopPartitionHint)) { - loop.CopyOnWrite()->annotations.erase(tvm::s_tir::attr::kLoopPartitionHint); + if (loop->annotations.count(tvm::s_tir::attr::kLoopPartitionHint.name)) { + loop.CopyOnWrite()->annotations.erase(tvm::s_tir::attr::kLoopPartitionHint.name); } return loop; } diff --git a/src/s_tir/transform/lower_cross_thread_reduction.cc b/src/s_tir/transform/lower_cross_thread_reduction.cc index 48cf504e3804..82f847ff8d99 100644 --- a/src/s_tir/transform/lower_cross_thread_reduction.cc +++ b/src/s_tir/transform/lower_cross_thread_reduction.cc @@ -952,7 +952,7 @@ class CrossThreadReductionTransformer : public StmtExprMutator { /*kind=*/ForKind::kParallel, // /*body=*/body, // /*annotations=*/ - {{tvm::s_tir::attr::kThreadBinding, ffi::String("threadIdx." + dim_index)}}, + {{tvm::s_tir::attr::kThreadBinding.name, ffi::String("threadIdx." + dim_index)}}, /*step=*/std::nullopt); } return body; diff --git a/src/s_tir/transform/lower_match_tensor.cc b/src/s_tir/transform/lower_match_tensor.cc index cbaccf90fc38..f064ab9c259c 100644 --- a/src/s_tir/transform/lower_match_tensor.cc +++ b/src/s_tir/transform/lower_match_tensor.cc @@ -90,7 +90,7 @@ class MatchTensorLower : public StmtExprMutator { auto* n = stmt.CopyOnWrite(); // Match tensors are aliases of their source region. Their placement belongs // to that alias definition, which disappears along with the match tensor. - if (auto value = n->annotations.Get(s_tir::attr::kTensorAllocatedAddr)) { + if (auto value = n->annotations.Get(s_tir::attr::kTensorAllocatedAddr.name)) { TensorAllocatedAddresses addresses; for (const auto& entry : value.value().cast()) { bool is_alias = false; @@ -100,9 +100,9 @@ class MatchTensorLower : public StmtExprMutator { if (!is_alias) addresses.push_back(entry); } if (addresses.empty()) { - n->annotations.erase(s_tir::attr::kTensorAllocatedAddr); + n->annotations.erase(s_tir::attr::kTensorAllocatedAddr.name); } else { - n->annotations.Set(s_tir::attr::kTensorAllocatedAddr, addresses); + n->annotations.Set(s_tir::attr::kTensorAllocatedAddr.name, addresses); } } n->match_tensors = {}; diff --git a/src/s_tir/transform/lower_opaque_block.cc b/src/s_tir/transform/lower_opaque_block.cc index 35ef5052fb56..ad8c3d02952e 100644 --- a/src/s_tir/transform/lower_opaque_block.cc +++ b/src/s_tir/transform/lower_opaque_block.cc @@ -82,7 +82,7 @@ class OpaqueBlockLower : public StmtExprMutator { } // Step 3. Handle allocations in reverse order ffi::Map> addresses; - if (auto value = new_block->annotations.Get(s_tir::attr::kTensorAllocatedAddr)) { + if (auto value = new_block->annotations.Get(s_tir::attr::kTensorAllocatedAddr.name)) { for (const auto& entry : value.value().cast()) { addresses.Set(entry.get<0>(), entry.get<1>()); } @@ -97,10 +97,8 @@ class OpaqueBlockLower : public StmtExprMutator { tuple.Set<0>(-1); allocate_aligns.push_back(tuple); } - allocate_annotations.Set(s_tir::attr::kTensorDimAlign, allocate_aligns); + allocate_annotations.Set(s_tir::attr::kTensorDimAlign.name, allocate_aligns); } - allocate_annotations.Set(tvm::tirx::attr::kTensorDataAlignment, - IntImm::Int32(tensor.ty()->data_alignment)); ffi::Array args; if (auto address = addresses.Get(tensor)) args.push_back(tvm::Tuple(address.value())); body = SeqStmt({Bind(tensor, Call(tensor.ty(), tirx::alloc_tensor_op(), args, @@ -132,10 +130,11 @@ class OpaqueBlockLower : public StmtExprMutator { UpdateUnrollPolicy(op->annotations); const auto& policy = unroll_policy_.Current(); if (policy.auto_unroll_max_step.has_value()) { - annotations.Set(tvm::tirx::attr::kAutoUnrollMaxStep, policy.auto_unroll_max_step.value()); + annotations.Set(tvm::tirx::attr::kAutoUnrollMaxStep.name, + policy.auto_unroll_max_step.value()); } if (policy.unroll_explicit.has_value()) { - annotations.Set(tvm::tirx::attr::kUnrollExplicit, policy.unroll_explicit.value()); + annotations.Set(tvm::tirx::attr::kUnrollExplicit.name, policy.unroll_explicit.value()); } // Rewrite annotations before visiting body-local definitions. new_annotations = HandleAnnotations(annotations); @@ -145,7 +144,7 @@ class OpaqueBlockLower : public StmtExprMutator { // Step 2. Keep thread-binding loops until LowerThreadBinding. if (!tvm::s_tir::GetThreadBinding(op).has_value() && IsOne(extent) && op->annotations.empty() && - !op->annotations.count(tvm::s_tir::attr::kIrregularLoopMark)) { + !op->annotations.count(tvm::s_tir::attr::kIrregularLoopMark.name)) { return body; } return For(op->loop_var, std::move(min), std::move(extent), op->kind, std::move(body), @@ -154,11 +153,11 @@ class OpaqueBlockLower : public StmtExprMutator { void UpdateUnrollPolicy(const ffi::Map& annotations) { auto& policy = unroll_policy_.Current(); - if (auto value = annotations.Get(tvm::tirx::attr::kAutoUnrollMaxStep); + if (auto value = annotations.Get(tvm::tirx::attr::kAutoUnrollMaxStep.name); value.has_value() && value.value() != nullptr) { policy.auto_unroll_max_step = value.value(); } - if (auto value = annotations.Get(tvm::tirx::attr::kUnrollExplicit); + if (auto value = annotations.Get(tvm::tirx::attr::kUnrollExplicit.name); value.has_value() && value.value() != nullptr) { policy.unroll_explicit = value.value(); } diff --git a/src/s_tir/transform/lower_thread_binding.cc b/src/s_tir/transform/lower_thread_binding.cc index 07a3f2fbc8cc..89a4459701eb 100644 --- a/src/s_tir/transform/lower_thread_binding.cc +++ b/src/s_tir/transform/lower_thread_binding.cc @@ -60,8 +60,8 @@ class ThreadBindingLowerer : public StmtExprMutator { PrimExpr extent = Mutate(op->extent, inplace_mode).ValueOrUnchanged(op->extent); TVM_FFI_ICHECK(IsZero(min)) << "Thread binding loops must start at zero"; TVM_FFI_ICHECK(tvm::s_tir::GetThreadBinding(op).has_value()); - TVM_FFI_ICHECK(!op->annotations.count(tvm::s_tir::attr::kLoopPartitionHint) || - op->annotations.at(tvm::s_tir::attr::kLoopPartitionHint) == nullptr) + TVM_FFI_ICHECK(!op->annotations.count(tvm::s_tir::attr::kLoopPartitionHint.name) || + op->annotations.at(tvm::s_tir::attr::kLoopPartitionHint.name) == nullptr) << "Run LoopPartition before LowerThreadBinding"; PrimVar launch_var(op->loop_var->name, extent.ty()); auto previous_remap = VarRemapGet(op->loop_var); @@ -71,7 +71,7 @@ class ThreadBindingLowerer : public StmtExprMutator { .ValueOrUnchanged(op->annotations); Stmt body = Mutate(op->body, inplace_mode).ValueOrUnchanged(op->body); VarRemapSet(op->loop_var, previous_remap); - annotations.erase(tvm::s_tir::attr::kThreadBinding); + annotations.erase(tvm::s_tir::attr::kThreadBinding.name); if (!annotations.empty()) { PrimType ty = op->loop_var.ty(); body = For(PrimVar("annotation", ty), IntImm(ty, 0), IntImm(ty, 1), ForKind::kDefault, diff --git a/src/s_tir/transform/manifest_shared_memory_local_stage.cc b/src/s_tir/transform/manifest_shared_memory_local_stage.cc index a20c047a27fd..fab275217196 100644 --- a/src/s_tir/transform/manifest_shared_memory_local_stage.cc +++ b/src/s_tir/transform/manifest_shared_memory_local_stage.cc @@ -217,7 +217,7 @@ class SharedMemoryLocalStageInserter : public StmtExprMutator { } UnchangedOr Mutate_(const SBlockNode* op, InplaceMode inplace_mode) final { - if (op->annotations.count(s_tir::attr::kManifestSharedMemoryLocalStage)) { + if (op->annotations.count(s_tir::attr::kManifestSharedMemoryLocalStage.name)) { // Rewrite the shared memory access to load from the intermediate tensor. // The annotated block must be a leaf block (will be checked during rewriting). No need to // visit its body recursively. @@ -226,7 +226,7 @@ class SharedMemoryLocalStageInserter : public StmtExprMutator { auto [target_tensor, new_tensor, new_block, local_stage] = rewriter.Rewrite(op); tensor_remap_.Set(target_tensor, new_tensor); - new_block.CopyOnWrite()->annotations.erase(s_tir::attr::kManifestSharedMemoryLocalStage); + new_block.CopyOnWrite()->annotations.erase(s_tir::attr::kManifestSharedMemoryLocalStage.name); tensor_local_stage_.Set(target_tensor, local_stage); target_tensors_.push_back(target_tensor); diff --git a/src/s_tir/transform/memhammer_coalesce.cc b/src/s_tir/transform/memhammer_coalesce.cc index 2e0b2ad1ca6c..cba52981a045 100644 --- a/src/s_tir/transform/memhammer_coalesce.cc +++ b/src/s_tir/transform/memhammer_coalesce.cc @@ -135,9 +135,9 @@ Stmt SplitBindVectorize(const Stmt& stmt, const ConstraintSet& constraints) { body = For(new_loop_vars.back().as_or_throw(), 0, vector_len, ForKind::kVectorized, std::move(body)); for (int i = n - 2; i >= 1; i--) { - body = For(new_loop_vars[i].as_or_throw(), 0, factors[i], ForKind::kParallel, - std::move(body), - {{tvm::s_tir::attr::kThreadBinding, ffi::String(thread_axis[i - 1])}}, std::nullopt); + body = For( + new_loop_vars[i].as_or_throw(), 0, factors[i], ForKind::kParallel, std::move(body), + {{tvm::s_tir::attr::kThreadBinding.name, ffi::String(thread_axis[i - 1])}}, std::nullopt); } return For(new_loop_vars[0].as_or_throw(), 0, factors[0], ForKind::kDefault, std::move(body)); diff --git a/src/s_tir/transform/memhammer_lower_auto_copy.cc b/src/s_tir/transform/memhammer_lower_auto_copy.cc index c436ce1301e2..ca59b327b8a9 100644 --- a/src/s_tir/transform/memhammer_lower_auto_copy.cc +++ b/src/s_tir/transform/memhammer_lower_auto_copy.cc @@ -748,7 +748,7 @@ class AutoCopyMutator : public StmtExprMutator { .ValueOrUnchanged(ffi::GetRef(op)) .as_or_throw(); // only rewrite the block annotated with "auto_copy" - if (!GetAnn(op, tvm::s_tir::attr::kAutoCopy).value_or(false)) { + if (!GetAnn(op, tvm::s_tir::attr::kAutoCopy.name).value_or(false)) { SBlockNode* n = block.CopyOnWrite(); n->alloc_tensors = padder.PadSharedMemory(std::move(n->alloc_tensors)); return block; @@ -840,7 +840,7 @@ class ThreadExtentCollector : public StmtExprVisitor { private: ffi::Optional Visit_(const SBlockNode* op) final { if (ffi::Optional warp_execution = - GetAnn(op, tvm::s_tir::attr::kWarpExecution)) { + GetAnn(op, tvm::s_tir::attr::kWarpExecution.name)) { if (warp_execution.value() != 0) { thread_extent_.Set("threadIdx.x", 32); } diff --git a/src/s_tir/transform/memhammer_rewrite_rule.h b/src/s_tir/transform/memhammer_rewrite_rule.h index d9cf40c6c365..fea6db24fca7 100644 --- a/src/s_tir/transform/memhammer_rewrite_rule.h +++ b/src/s_tir/transform/memhammer_rewrite_rule.h @@ -66,10 +66,10 @@ struct ConstraintSet { read_region(read_region), write_region(write_region), data_bits(data_bits) { - if (auto add_local_stage = ann.Get(tvm::s_tir::attr::kLocalStage)) { + if (auto add_local_stage = ann.Get(tvm::s_tir::attr::kLocalStage.name)) { this->add_local_stage = static_cast(add_local_stage.value().cast()->value); } - if (auto vector_bytes = ann.Get(tvm::s_tir::attr::kVectorBytes)) { + if (auto vector_bytes = ann.Get(tvm::s_tir::attr::kVectorBytes.name)) { this->vector_bytes = vector_bytes.value().cast()->value.as().value(); } } diff --git a/src/s_tir/transform/merge_shared_memory_allocations.cc b/src/s_tir/transform/merge_shared_memory_allocations.cc index cbd11c616e1c..e1f018702240 100644 --- a/src/s_tir/transform/merge_shared_memory_allocations.cc +++ b/src/s_tir/transform/merge_shared_memory_allocations.cc @@ -539,7 +539,7 @@ class SharedMemoryRewriter : public StmtExprMutator { // 7. Wrap with the merged-buffer AllocTensor. ffi::Map annotations; if (scope.has_volatile_alloc) { - annotations.Set(tvm::tirx::attr::kVolatile, true); + annotations.Set(tvm::tirx::attr::kVolatile.name, true); } Stmt alloc_stmt = Bind(scope.merged_tensor.value(), Call(scope.merged_tensor.value().ty(), tirx::alloc_tensor_op(), {}, @@ -566,7 +566,7 @@ class SharedMemoryRewriter : public StmtExprMutator { if (!scope_stack_.empty()) { KernelScope& scope = scope_stack_.back(); if (scope.shmem_allocs.count(op->var.as_or_throw().get())) { - if (call->attrs.as()->dict.count(tvm::tirx::attr::kVolatile)) { + if (call->attrs.as()->dict.count(tvm::tirx::attr::kVolatile.name)) { scope.has_volatile_alloc = true; } return Evaluate(0); diff --git a/src/s_tir/transform/remove_weight_layout_rewrite_block.cc b/src/s_tir/transform/remove_weight_layout_rewrite_block.cc index 489a79e25a63..aa3f703b6783 100644 --- a/src/s_tir/transform/remove_weight_layout_rewrite_block.cc +++ b/src/s_tir/transform/remove_weight_layout_rewrite_block.cc @@ -63,7 +63,7 @@ class RemoveLayoutRewriteBlock : public StmtExprMutator { .ValueOrUnchanged(ffi::GetRef(op)) .as_or_throw(); - auto it = block->annotations.find(tvm::s_tir::attr::kMetaScheduleLayoutRewritePreproc); + auto it = block->annotations.find(tvm::s_tir::attr::kMetaScheduleLayoutRewritePreproc.name); if (it == block->annotations.end() || !IsOne((*it).second.cast())) { // The block is not a weight layout block // Remove allocates if needed diff --git a/src/s_tir/transform/rewrite_unsafe_select.cc b/src/s_tir/transform/rewrite_unsafe_select.cc index f42b4ab94594..ff2e90f4f556 100644 --- a/src/s_tir/transform/rewrite_unsafe_select.cc +++ b/src/s_tir/transform/rewrite_unsafe_select.cc @@ -28,8 +28,8 @@ #include #include #include +#include #include -#include namespace tvm { namespace s_tir { diff --git a/src/s_tir/transform/tensorcore_infer_fragment.cc b/src/s_tir/transform/tensorcore_infer_fragment.cc index 343e4cdd6e50..f41f1046a769 100644 --- a/src/s_tir/transform/tensorcore_infer_fragment.cc +++ b/src/s_tir/transform/tensorcore_infer_fragment.cc @@ -227,9 +227,9 @@ class InferFragmenter : public s_tir::StmtExprMutator { } }; set_annotation( - tvm::s_tir::attr::kFragmentShape, + tvm::s_tir::attr::kFragmentShape.name, std::to_string(info.m) + ", " + std::to_string(info.n) + ", " + std::to_string(info.k)); - if (!info.layout.empty()) set_annotation(tvm::s_tir::attr::kFragmentLayout, info.layout); + if (!info.layout.empty()) set_annotation(tvm::s_tir::attr::kFragmentLayout.name, info.layout); if (annotations.same_as(attrs->dict)) return ffi::Unchanged(); auto value = ffi::GetRef(call); value.CopyOnWrite()->attrs = DictAttrs(std::move(annotations)); diff --git a/src/s_tir/transform/thread_storage_sync.cc b/src/s_tir/transform/thread_storage_sync.cc index 667b56fc5827..87314c8e5f56 100644 --- a/src/s_tir/transform/thread_storage_sync.cc +++ b/src/s_tir/transform/thread_storage_sync.cc @@ -411,7 +411,7 @@ Pass ThreadSync(runtime::MemoryScope storage_scope) { Pass LowerSynchronization() { auto pass_func = [](Function f, IRModule, PassContext) { if (!f->body.has_value()) return f; - auto target = f->GetAttr(tvm::attr::kTarget); + auto target = f->GetAttr(tvm::attr::kTarget); bool is_cuda = target && target.value()->kind->name == "cuda"; auto* n = f.CopyOnWrite(); n->body = ffi::make_object(is_cuda) diff --git a/src/te/operation/create_function.cc b/src/te/operation/create_function.cc index 8ac3f162d8f0..4be1405c2ed0 100644 --- a/src/te/operation/create_function.cc +++ b/src/te/operation/create_function.cc @@ -246,7 +246,7 @@ class LayoutFreePlaceholdersNormalizer : public s_tir::StmtExprMutator { for (int i : this->layout_free_tensor_indices_) { indices.push_back(i); } - return WithAttr(std::move(func), tvm::s_tir::attr::kLayoutFreeTensors, indices); + return WithAttr(std::move(func), tvm::s_tir::attr::kLayoutFreeTensors.name, indices); } UnchangedOr Mutate_(const s_tir::SBlockNode* _block, InplaceMode inplace_mode) final { @@ -426,7 +426,7 @@ ffi::Map GenerateBlockAnnotations(const te::ComputeOp& co } } // Set script_parsing_detect_access - annotations.Set(tvm::s_tir::attr::kScriptParsingDetectAccess, 3); + annotations.Set(tvm::s_tir::attr::kScriptParsingDetectAccess.name, 3); return annotations; } @@ -906,9 +906,9 @@ Function GenerateAndCompleteFunction(const ffi::Array& arg_list, Function func = WithAttrs(Function(/*params=*/std::move(parameters), /*body=*/std::move(body), /*ret_type=*/VoidType()), - {{tvm::attr::kGlobalSymbol, ffi::String("main")}, - {tvm::tirx::attr::kNoAlias, true}, - {tvm::attr::kSTir, true}}); + {{tvm::attr::kGlobalSymbol.name, ffi::String("main")}, + {tvm::tirx::attr::kNoAlias.name, true}, + {tvm::attr::kSTir.name, true}}); const auto fcomplete = tvm::ffi::Function::GetGlobal("s_tir.script.Complete"); TVM_FFI_ICHECK(fcomplete.has_value()); func = (*fcomplete)(std::move(func), info->root_alloc).cast(); @@ -978,9 +978,9 @@ Function GenerateAndCompleteFunction(const ffi::Array& arg_tir_v Function func = WithAttrs(Function(/*params=*/std::move(parameters), /*body=*/std::move(body), /*ret_type=*/VoidType()), - {{tvm::attr::kGlobalSymbol, ffi::String("main")}, - {tvm::tirx::attr::kNoAlias, true}, - {tvm::attr::kSTir, true}}); + {{tvm::attr::kGlobalSymbol.name, ffi::String("main")}, + {tvm::tirx::attr::kNoAlias.name, true}, + {tvm::attr::kSTir.name, true}}); const auto fcomplete = tvm::ffi::Function::GetGlobal("s_tir.script.Complete"); TVM_FFI_ICHECK(fcomplete.has_value()); func = (*fcomplete)(std::move(func), info->root_alloc).cast(); diff --git a/src/te/script/printer/expr.cc b/src/te/script/printer/expr.cc index 37c0f0b31889..dbc0e5030257 100644 --- a/src/te/script/printer/expr.cc +++ b/src/te/script/printer/expr.cc @@ -51,7 +51,7 @@ ffi::Optional CommReducerDocTranslate(DocTranslatorObj* d, ffi::AnyView TVM_FFI_STATIC_INIT_BLOCK() { ffi::reflection::TypeAttrDef().attr( - tvm::script::printer::type_attr::kDocTranslate, + tvm::script::printer::type_attr::kDocTranslate.name, FDocTranslate::FromNative<&CommReducerDocTranslate>()); } @@ -86,7 +86,7 @@ ffi::Optional ReduceDocTranslate(DocTranslatorObj* d, ffi::AnyView inpu TVM_FFI_STATIC_INIT_BLOCK() { ffi::reflection::TypeAttrDef().attr( - tvm::script::printer::type_attr::kDocTranslate, + tvm::script::printer::type_attr::kDocTranslate.name, FDocTranslate::FromNative<&ReduceDocTranslate>()); } diff --git a/tests/python/s_tir/transform/test_s_tir_transform_lower_opaque_block.py b/tests/python/s_tir/transform/test_s_tir_transform_lower_opaque_block.py index 21616e37103b..7b8de2334e19 100644 --- a/tests/python/s_tir/transform/test_s_tir_transform_lower_opaque_block.py +++ b/tests/python/s_tir/transform/test_s_tir_transform_lower_opaque_block.py @@ -62,9 +62,7 @@ def transformed_elementwise_func( A: T.Tensor((16, 16), "float32"), C: T.Tensor((16, 16), "float32") ) -> None: for i in T.serial(0, 16): - B_new = T.alloc_tensor( - ty_args=[T.Tensor([1, 16], "float32")], annotations={"tirx.tensor_data_alignment": 64} - ) + B_new = T.alloc_tensor(ty_args=[T.Tensor([1, 16], "float32")]) for j in T.serial(0, 16): B_new[0, j] = A[i, j] + 1.0 for j in T.serial(0, 16): @@ -99,10 +97,7 @@ def transformed_gpu_func( i0 = T.launch_thread("blockIdx.x", 4) i1 = T.launch_thread("threadIdx.x", 2) i2 = T.launch_thread("vthread", 2) - B = T.alloc_tensor( - ty_args=[T.Tensor([1, 16], "float32", scope="local")], - annotations={"tirx.tensor_data_alignment": 64}, - ) + B = T.alloc_tensor(ty_args=[T.Tensor([1, 16], "float32", scope="local")]) for j in range(0, 16): B[0, j] = A[i0 * 4 + i1 * 2 + i2, j] + 1.0 for j in range(0, 16): @@ -141,9 +136,7 @@ def transformed_symbolic_func( m: T.int32, ) -> None: for i in range(0, n): - B = T.alloc_tensor( - ty_args=[T.Tensor([m], "float32")], annotations={"tirx.tensor_data_alignment": 64} - ) + B = T.alloc_tensor(ty_args=[T.Tensor([m], "float32")]) for j in range(0, m): B[j] = A[i, j] + 1.0 for j in range(0, m): @@ -198,12 +191,8 @@ def compacted_multi_alloc_func(A: T.Tensor(32, "float32"), D: T.Tensor(32, "floa @Ts.function def transformed_multi_alloc_func(A: T.Tensor(32, "float32"), D: T.Tensor(32, "float32")) -> None: for i in range(0, 32): - B = T.alloc_tensor( - ty_args=[T.Tensor((32,), "float32")], annotations={"tirx.tensor_data_alignment": 64} - ) - C = T.alloc_tensor( - ty_args=[T.Tensor((32,), "float32")], annotations={"tirx.tensor_data_alignment": 64} - ) + B = T.alloc_tensor(ty_args=[T.Tensor((32,), "float32")]) + C = T.alloc_tensor(ty_args=[T.Tensor((32,), "float32")]) B[i] = A[i] + 1.0 C[i] = A[i] + B[i] D[i] = C[i] * 2.0 @@ -238,10 +227,7 @@ def transformed_strided_tensor_func( ) -> None: # body for i0 in T.serial(4): - B = T.alloc_tensor( - ty_args=[T.Tensor([4, 16], "float32", strides=[17, 1])], - annotations={"tirx.tensor_data_alignment": 64}, - ) + B = T.alloc_tensor(ty_args=[T.Tensor([4, 16], "float32", strides=[17, 1])]) for i1, j in T.grid(4, 16): B[i1, j] = A[i0 * 4 + i1, j] + T.float32(1) for i1, j in T.grid(4, 16): @@ -285,8 +271,7 @@ def transformed_symbolic_strided_tensor_func(A: T.Tensor((1, n, 10240), "float32 strides=(72 * padded_size, 72, 1), scope="shared.dyn", ) - ], - annotations={"tirx.tensor_data_alignment": 64}, + ] ) for ax0, ax1 in T.grid(96, 64): if i * 128 + j * 32 + ax0 < (n + 63) // 64 * 64: @@ -383,6 +368,27 @@ def after(A: T.Tensor(8, "float32"), B: T.Tensor(8, "float32")): ) +def test_allocation_type_alignment(): + @Ts.function + def before(A: T.Tensor(4, "float32")): + with Ts.sblock("root"): + B = Ts.sblock_alloc_tensor((4,), "float32", scope="shared", align=128) + for i in range(4): + B[i] = A[i] + for i in range(4): + A[i] = B[i] + + @Ts.function + def after(A: T.Tensor(4, "float32")): + B = T.alloc_tensor(ty_args=[T.Tensor((4,), "float32", scope="shared", align=128)]) + for i in range(4): + B[i] = A[i] + for i in range(4): + A[i] = B[i] + + _check(before, after) + + def test_boolean_handling(): _check(boolean_handling_before, boolean_handling_after) diff --git a/tirx/include/tvm/backend/cuda/op_attr_types.h b/tirx/include/tvm/backend/cuda/op_attr.h similarity index 87% rename from tirx/include/tvm/backend/cuda/op_attr_types.h rename to tirx/include/tvm/backend/cuda/op_attr.h index e5f38bad0975..7788fff025eb 100644 --- a/tirx/include/tvm/backend/cuda/op_attr_types.h +++ b/tirx/include/tvm/backend/cuda/op_attr.h @@ -17,11 +17,11 @@ * under the License. */ /*! - * \file tvm/backend/cuda/op_attr_types.h + * \file tvm/backend/cuda/op_attr.h * \brief CUDA backend operator attribute keys. */ -#ifndef TVM_BACKEND_CUDA_OP_ATTR_TYPES_H_ -#define TVM_BACKEND_CUDA_OP_ATTR_TYPES_H_ +#ifndef TVM_BACKEND_CUDA_OP_ATTR_H_ +#define TVM_BACKEND_CUDA_OP_ATTR_H_ #include @@ -35,4 +35,4 @@ inline constexpr ffi::reflection::AttrKey kNeedWarpShuffle{"cuda.need_warp } // namespace backend } // namespace tvm -#endif // TVM_BACKEND_CUDA_OP_ATTR_TYPES_H_ +#endif // TVM_BACKEND_CUDA_OP_ATTR_H_ diff --git a/tirx/include/tvm/codegen/intrin_rule.h b/tirx/include/tvm/codegen/intrin_rule.h index 45d534ccb44b..247393e3fd5a 100644 --- a/tirx/include/tvm/codegen/intrin_rule.h +++ b/tirx/include/tvm/codegen/intrin_rule.h @@ -65,7 +65,7 @@ struct Direct { * For example, an off-tree backend can register a C math mapping with * `OpDef("prim.exp").set_attr( * "my_backend.FLowerIntrinsic", DispatchPureExtern)` after - * including tvm/tirx/op_attr_types.h. + * including tvm/tirx/op/attr.h. * \tparam dtype_from_arg Use the argument type of a unary call instead of its result type. */ template diff --git a/tirx/include/tvm/ir/expr.h b/tirx/include/tvm/ir/expr.h index 7a30a8735998..313c1409d768 100644 --- a/tirx/include/tvm/ir/expr.h +++ b/tirx/include/tvm/ir/expr.h @@ -1190,8 +1190,10 @@ inline constexpr ffi::reflection::AttrKey kCallEffectKind{"TCallEffectK using TCallEffectKind = decltype(op_attr::kCallEffectKind)::ValueType; namespace type_attr { -inline constexpr const char* kTensorLoadInferType = "__tvm_tensor_load_infer_type__"; -inline constexpr const char* kSubscriptExprRealize = "__tvm_subscript_expr_realize__"; +inline constexpr ffi::reflection::AttrKey kTensorLoadInferType{ + "__tvm_tensor_load_infer_type__"}; +inline constexpr ffi::reflection::AttrKey kSubscriptExprRealize{ + "__tvm_subscript_expr_realize__"}; } // namespace type_attr } // namespace tvm diff --git a/tirx/include/tvm/ir/function.h b/tirx/include/tvm/ir/function.h index 0c1f989c772c..051132452646 100644 --- a/tirx/include/tvm/ir/function.h +++ b/tirx/include/tvm/ir/function.h @@ -35,6 +35,8 @@ namespace tvm { +class Target; + /*! * \brief Possible Calling conventions. * @@ -217,7 +219,7 @@ namespace attr { * * \sa tvm::CallingConv */ -constexpr const char* kCallingConv = "calling_conv"; +inline constexpr ffi::reflection::AttrKey kCallingConv{"calling_conv"}; /*! * \brief Compilation target of the function. @@ -226,7 +228,7 @@ constexpr const char* kCallingConv = "calling_conv"; * * \sa tvm::Target */ -constexpr const char* kTarget = "target"; +inline constexpr ffi::reflection::AttrKey kTarget{"target"}; /*! * \brief Global linker symbol of the function in generated code. @@ -247,7 +249,7 @@ constexpr const char* kTarget = "target"; * * Type: String */ -constexpr const char* kGlobalSymbol = "global_symbol"; +inline constexpr ffi::reflection::AttrKey kGlobalSymbol{"global_symbol"}; /*! * \brief The function uses s_tir (apache-derived TIR) semantics: @@ -257,7 +259,7 @@ constexpr const char* kGlobalSymbol = "global_symbol"; * * Type: IntImm (bool dtype) */ -constexpr const char* kSTir = "s_tir"; +inline constexpr ffi::reflection::AttrKey kSTir{"s_tir"}; /*! * \brief Script construction namespace used to select the function decorator. @@ -266,19 +268,19 @@ constexpr const char* kSTir = "s_tir"; * * Type: String */ -constexpr const char* kScriptNamespace = "script.namespace"; +inline constexpr ffi::reflection::AttrKey kScriptNamespace{"script.namespace"}; /*! * \brief Number of inputs of the Function * * Type: Int */ -constexpr const char* kNumInputs = "num_inputs"; +inline constexpr ffi::reflection::AttrKey kNumInputs{"num_inputs"}; } // namespace attr inline LinkageType BaseFuncNode::GetLinkageType() const { - if (GetAttr(tvm::attr::kGlobalSymbol)) + if (GetAttr(tvm::attr::kGlobalSymbol)) return LinkageType::kExternal; else return LinkageType::kInternal; diff --git a/tirx/include/tvm/ir/module.h b/tirx/include/tvm/ir/module.h index 25dbf9457577..c0a157fcd7dd 100644 --- a/tirx/include/tvm/ir/module.h +++ b/tirx/include/tvm/ir/module.h @@ -26,6 +26,8 @@ #include #include +#include +#include #include #include #include @@ -320,7 +322,7 @@ namespace attr { * * Type: ffi::Array */ -constexpr const char* kExternalMods = "external_mods"; +inline constexpr ffi::reflection::AttrKey> kExternalMods{"external_mods"}; /*! * \brief A prefix for generating C symbols system lib creation. @@ -349,7 +351,7 @@ constexpr const char* kExternalMods = "external_mods"; * * Having this attribute implies system lib generation linkage. */ -constexpr const char* kSystemLibPrefix = "system_lib_prefix"; +inline constexpr ffi::reflection::AttrKey kSystemLibPrefix{"system_lib_prefix"}; /*! * \brief All the named Tensors accumulated during compilation by external codegen. @@ -358,9 +360,11 @@ constexpr const char* kSystemLibPrefix = "system_lib_prefix"; * * Type: ffi::Map */ -constexpr const char* kConstNameToConstant = "const_name_to_constant"; +inline constexpr ffi::reflection::AttrKey> kConstNameToConstant{ + "const_name_to_constant"}; -constexpr const char* kDeviceContexts = "device_contexts"; +inline constexpr ffi::reflection::AttrKey> kDeviceContexts{ + "device_contexts"}; } // namespace attr } // namespace tvm diff --git a/tirx/include/tvm/ir/stmt.h b/tirx/include/tvm/ir/stmt.h index d35cc2026e41..6813bd44d842 100644 --- a/tirx/include/tvm/ir/stmt.h +++ b/tirx/include/tvm/ir/stmt.h @@ -681,8 +681,10 @@ using FRegionGetBodyParams = decltype(op_attr::kRegionGetBodyParams)::ValueType; using FRegionValidate = decltype(op_attr::kRegionValidate)::ValueType; namespace type_attr { -inline constexpr const char* kEvaluateValidate = "__tvm_evaluate_validate__"; -inline constexpr const char* kTensorStoreValidate = "__tvm_tensor_store_validate__"; +inline constexpr ffi::reflection::AttrKey kEvaluateValidate{ + "__tvm_evaluate_validate__"}; +inline constexpr ffi::reflection::AttrKey kTensorStoreValidate{ + "__tvm_tensor_store_validate__"}; } // namespace type_attr } // namespace tvm diff --git a/tirx/include/tvm/script/printer/doc_translator.h b/tirx/include/tvm/script/printer/doc_translator.h index 202c85938b13..37e992614683 100644 --- a/tirx/include/tvm/script/printer/doc_translator.h +++ b/tirx/include/tvm/script/printer/doc_translator.h @@ -434,7 +434,7 @@ namespace type_attr { /*! * \brief Type attribute containing a native or packed translation hook. */ -inline constexpr const char* kDocTranslate = "__tvm_doc_translate__"; +inline constexpr ffi::reflection::AttrKey kDocTranslate{"__tvm_doc_translate__"}; } // namespace type_attr namespace op_attr { @@ -447,7 +447,8 @@ inline constexpr ffi::reflection::AttrKey kDocTranslateOpCall{ inline ffi::Optional DocTranslatorObj::DefaultTranslate(ffi::AnyView value, const ffi::Object* destination) { - static ffi::reflection::TypeAttrColumn column(tvm::script::printer::type_attr::kDocTranslate); + static ffi::reflection::TypeAttrColumn column( + tvm::script::printer::type_attr::kDocTranslate.name); ffi::AnyView attr = column[value.type_index()]; if (attr.type_index() == ffi::TypeIndex::kTVMFFIOpaquePtr) { return ffi::details::ExpectedUnsafe::MoveFromTVMFFIAny>( diff --git a/tirx/include/tvm/target/target.h b/tirx/include/tvm/target/target.h index 52ed87c73154..e6a9e2105284 100644 --- a/tirx/include/tvm/target/target.h +++ b/tirx/include/tvm/target/target.h @@ -35,6 +35,12 @@ namespace tvm { +namespace backend::attr { +/*! \brief Optional backend configuration on a Target, distinct from serialized IR/runtime JSON. */ +inline constexpr ffi::reflection::AttrKey> kBackendConfig{ + "backend_config"}; +} // namespace backend::attr + class TargetInternal; class Target; @@ -119,6 +125,21 @@ class TargetObj : public ffi::Object { return GetAttr(attr_key, ffi::Optional(default_value)); } + /*! \brief Get an optional Target attribute using the value type declared by its key. */ + template + ffi::Optional GetAttr(ffi::reflection::AttrKey key, + ffi::Optional::ValueType> + default_value = std::nullopt) const { + return GetAttr(key.name, std::move(default_value)); + } + + /*! \brief Use a default value when the typed Target attribute is absent. */ + template + ffi::Optional GetAttr(ffi::reflection::AttrKey key, + typename ffi::reflection::AttrKey::ValueType default_value) const { + return GetAttr(key, ffi::Optional(std::move(default_value))); + } + static constexpr TVMFFISEqHashKind _type_s_eq_hash_kind = kTVMFFISEqHashKindTreeNode; TVM_FFI_DECLARE_OBJECT_INFO_FINAL("target.Target", TargetObj, ffi::Object); diff --git a/tirx/include/tvm/tirx/analysis.h b/tirx/include/tvm/tirx/analysis.h index 424766849af1..d574282a5d18 100644 --- a/tirx/include/tvm/tirx/analysis.h +++ b/tirx/include/tvm/tirx/analysis.h @@ -29,7 +29,7 @@ #include #include #include -#include +#include #include #include @@ -93,7 +93,7 @@ TVM_TIRX_DLL bool VerifyWellFormed(const IRModule& mod, bool assert_mode = true) /*! * \brief Find the entry function of the given IRModule, i.e, functions marked by - * `tirx::attr::kIsEntryFunc`, whose name is `main` or being the only Function. + * `tirx::attr::kIsEntryFunc.name`, whose name is `main` or being the only Function. * \param mod The IRModule to find the entry function. * \param result_g_var The result GlobalVar of the entry function. * \return The entry function. diff --git a/tirx/include/tvm/tirx/function.h b/tirx/include/tvm/tirx/function.h index 615e5cec27a8..411412c23f6e 100644 --- a/tirx/include/tvm/tirx/function.h +++ b/tirx/include/tvm/tirx/function.h @@ -234,43 +234,44 @@ namespace attr { * * \sa tvm::CallingConv::kDeviceKernelLaunch */ -constexpr const char* kKernelLaunchParams = "tirx.kernel_launch_params"; +inline constexpr ffi::reflection::AttrKey> kKernelLaunchParams{ + "tirx.kernel_launch_params"}; /*! * \brief Whether to set noalias rule on the function arguments. * - * Type: IntImm + * Type: bool */ -constexpr const char* kNoAlias = "tirx.noalias"; +inline constexpr ffi::reflection::AttrKey kNoAlias{"tirx.noalias"}; /*! * \brief Mark the function as the entry function of * the final generated runtime module. * - * Type: IntImm + * Type: bool * * \note There can only be one entry function per module. */ -constexpr const char* kIsEntryFunc = "tirx.is_entry_func"; +inline constexpr ffi::reflection::AttrKey kIsEntryFunc{"tirx.is_entry_func"}; /*! * \brief Mark the function as the global function called from the host. * - * Type: IntImm + * Type: bool */ -constexpr const char* kIsGlobalFunc = "tirx.is_global_func"; +inline constexpr ffi::reflection::AttrKey kIsGlobalFunc{"tirx.is_global_func"}; /*! * \brief Mark the function as run on the host, mutually exclusive with kTarget. * - * Type: IntImm + * Type: bool */ -constexpr const char* kIsHostFunc = "tirx.is_host_func"; +inline constexpr ffi::reflection::AttrKey kIsHostFunc{"tirx.is_host_func"}; /*! * \brief Mark the kernel as persistent. */ -constexpr const char* kPersistentKernel = "tirx.persistent_kernel"; +inline constexpr ffi::reflection::AttrKey kPersistentKernel{"tirx.persistent_kernel"}; } // namespace attr } // namespace tirx diff --git a/tirx/include/tvm/tirx/op_attr_types.h b/tirx/include/tvm/tirx/op/attr.h similarity index 89% rename from tirx/include/tvm/tirx/op_attr_types.h rename to tirx/include/tvm/tirx/op/attr.h index 27a82e2618f2..ef352a9036d2 100644 --- a/tirx/include/tvm/tirx/op_attr_types.h +++ b/tirx/include/tvm/tirx/op/attr.h @@ -18,15 +18,15 @@ */ /*! - * \file tvm/tirx/op_attr_types.h - * \brief Attribute types in the Op registry for TIR ops. + * \file tvm/tirx/op/attr.h + * \brief Attribute keys and callback contracts for TIRx operators. * * These attributes can be set via OpDef::set_attr * * \sa tvm/ir/op.h */ -#ifndef TVM_TIRX_OP_ATTR_TYPES_H_ -#define TVM_TIRX_OP_ATTR_TYPES_H_ +#ifndef TVM_TIRX_OP_ATTR_H_ +#define TVM_TIRX_OP_ATTR_H_ #include #include @@ -60,4 +60,4 @@ inline constexpr ffi::reflection::AttrKey kVectorizable{"TVectorizable"}; } // namespace tirx } // namespace tvm -#endif // TVM_TIRX_OP_ATTR_TYPES_H_ +#endif // TVM_TIRX_OP_ATTR_H_ diff --git a/tirx/include/tvm/tirx/stmt.h b/tirx/include/tvm/tirx/stmt.h index c0a0fe729282..9142e725a0f3 100644 --- a/tirx/include/tvm/tirx/stmt.h +++ b/tirx/include/tvm/tirx/stmt.h @@ -41,24 +41,22 @@ inline bool IsNoOp(const Stmt& stmt) { namespace attr { constexpr const char* kPragmaUnroll = "tirx.pragma_unroll"; -constexpr const char* kDisableUnroll = "tirx.disable_unroll"; +inline constexpr ffi::reflection::AttrKey kDisableUnroll{"tirx.disable_unroll"}; /*! * \brief For annotation: maximum work for automatic unrolling. * * Integer policy inherited by nested loops unless they override it. Consumed by UnrollLoop. */ -constexpr const char* kAutoUnrollMaxStep = "tirx.auto_unroll_max_step"; +inline constexpr ffi::reflection::AttrKey kAutoUnrollMaxStep{"tirx.auto_unroll_max_step"}; /*! * \brief For annotation: expand unrolled bodies instead of preserving unrolled loops. * * Integer policy inherited by nested loops unless they override it. Consumed by UnrollLoop. */ -constexpr const char* kUnrollExplicit = "tirx.unroll_explicit"; +inline constexpr ffi::reflection::AttrKey kUnrollExplicit{"tirx.unroll_explicit"}; /*! \brief Annotation key on AllocTensor marking the allocation as volatile. */ -constexpr const char* kVolatile = "tirx.volatile"; -/*! \brief Mark tensor initial addr alignment in bytes */ -constexpr const char* kTensorDataAlignment = "tirx.tensor_data_alignment"; +inline constexpr ffi::reflection::AttrKey kVolatile{"tirx.volatile"}; } // namespace attr diff --git a/tirx/include/tvm/tirx/tile/context.h b/tirx/include/tvm/tirx/tile/context.h index c514548ee15c..dc44797f5fb6 100644 --- a/tirx/include/tvm/tirx/tile/context.h +++ b/tirx/include/tvm/tirx/tile/context.h @@ -33,19 +33,20 @@ namespace tile { namespace callback { /*! \brief The allocation Bind statements requested by the operator. */ -constexpr const char* kPrivateAlloc = "private_alloc"; +inline constexpr ffi::reflection::AttrKey> kPrivateAlloc{"private_alloc"}; /*! \brief The initialization statement of the operator. * which will be inserted at the beginning of the kernel */ -constexpr const char* kDeviceInitStmt = "device_init_stmt"; +inline constexpr ffi::reflection::AttrKey> kDeviceInitStmt{"device_init_stmt"}; /*! \brief The initialization statement of the operator. * which will be inserted at the beginning of the kernel */ -constexpr const char* kHostInitStmt = "host_init_stmt"; +inline constexpr ffi::reflection::AttrKey> kHostInitStmt{"host_init_stmt"}; /*! \brief Statements to be inserted after a specific tensor's definition (DeclTensor/AllocTensor). * Stored as Map>. */ -constexpr const char* kPostTensorDefStmt = "post_tensor_def_stmt"; +inline constexpr ffi::reflection::AttrKey>> kPostTensorDefStmt{ + "post_tensor_def_stmt"}; } // namespace callback /*! diff --git a/tirx/include/tvm/tirx/transform.h b/tirx/include/tvm/tirx/transform.h index 2956e0ebe560..9d13a742c03c 100644 --- a/tirx/include/tvm/tirx/transform.h +++ b/tirx/include/tvm/tirx/transform.h @@ -172,7 +172,7 @@ TVM_TIRX_DLL Pass MakePackedAPI(); * lowers host-to-device calls into the device kernel launch ABI. * * The resulting host-side function will keep the same - * `tvm::attr::kTarget` attribute (e.g. `T.target("cuda", + * `tvm::attr::kTarget.name` attribute (e.g. `T.target("cuda", * host=T.target("llvm"))`). This ensures that `MakePackedAPI` knows * which device type should be used for the input tensors. * diff --git a/tirx/python/tvm_tirx/__init__.py b/tirx/python/tvm_tirx/__init__.py index a6e35fd6c689..f42166e8a447 100644 --- a/tirx/python/tvm_tirx/__init__.py +++ b/tirx/python/tvm_tirx/__init__.py @@ -45,7 +45,7 @@ def _is_config_mode(): # tvm.runtime from .runtime._tensor import cpu, cuda, opencl, vulkan, metal - from .runtime._tensor import vpi, rocm, ext_dev, hexagon + from .runtime._tensor import vpi, rocm, ext_dev from .runtime import DataTypeCode # tvm.ir diff --git a/tirx/python/tvm_tirx/runtime/__init__.py b/tirx/python/tvm_tirx/runtime/__init__.py index 2722d8f6c141..b917f11a1323 100644 --- a/tirx/python/tvm_tirx/runtime/__init__.py +++ b/tirx/python/tvm_tirx/runtime/__init__.py @@ -27,7 +27,7 @@ # function exposures from ._tensor import cpu, cuda, opencl, vulkan, metal -from ._tensor import vpi, rocm, ext_dev +from ._tensor import vpi, rocm, ext_dev, hexagon from . import _ffi_api diff --git a/tirx/src/backend/cuda/codegen/codegen_cuda.cc b/tirx/src/backend/cuda/codegen/codegen_cuda.cc index 1c8b718f48fd..7ada9676e111 100644 --- a/tirx/src/backend/cuda/codegen/codegen_cuda.cc +++ b/tirx/src/backend/cuda/codegen/codegen_cuda.cc @@ -175,8 +175,7 @@ void CodeGenCUDA::PrintFunctionSignature(const ffi::String& function_name, const } void CodeGenCUDA::PrintFunctionPrefix(const Function& func, std::ostream& os) { - CallingConv calling_conv = - func->GetAttr(tvm::attr::kCallingConv, CallingConv::kDefault).value(); + CallingConv calling_conv = func->GetAttr(tvm::attr::kCallingConv, CallingConv::kDefault).value(); in_kernel_launch_ = (calling_conv == CallingConv::kDeviceKernelLaunch); if (calling_conv == CallingConv::kDeviceKernelLaunch) { os << "extern \"C\" __global__ "; @@ -242,9 +241,8 @@ void CodeGenCUDA::InitFuncState(const Function& func) { max_blocks_per_cluster_.reset(); max_registers_per_thread_.reset(); required_block_size_.reset(); - in_kernel_launch_ = - func->GetAttr(tvm::attr::kCallingConv, CallingConv::kDefault).value() == - CallingConv::kDeviceKernelLaunch; + in_kernel_launch_ = func->GetAttr(tvm::attr::kCallingConv, CallingConv::kDefault).value() == + CallingConv::kDeviceKernelLaunch; // Thread metadata is also needed while emitting cluster index accesses. auto extractor = ffi::make_object(); extractor->Visit(func->body); @@ -398,7 +396,7 @@ void CodeGenCUDA::Dispatch_(const ForNode* op) { PrimExpr end = IsZero(op->min) ? op->extent : sym::Analyzer()->Simplify(op->min + op->extent); std::string end_str = PrintExpr(end); std::string step_str = op->step.has_value() ? PrintExpr(*op->step) : ""; - if (op->annotations.count(tvm::tirx::attr::kDisableUnroll)) { + if (op->annotations.count(tvm::tirx::attr::kDisableUnroll.name)) { PrintIndent(); stream << "#pragma unroll 1\n"; } else if (op->kind == ForKind::kUnrolled) { @@ -1263,12 +1261,6 @@ void CodeGenCUDA::DispatchAllocTensor(const BindNode* op, const CallNode* buffer this->PrintIndent(); PrintStorageScope(scope, stream); int align = buffer.ty()->data_alignment; - auto it = annotations->dict.find(tvm::tirx::attr::kTensorDataAlignment); - if (it != annotations->dict.end()) { - if (const auto* n = (*it).second.as()) { - align = n->value.as().value(); - } - } if (align > 0 && scope == scope_shared_dyn) { stream << "__align__(" << align << ") "; } else if (align > 0) { @@ -1298,7 +1290,7 @@ void CodeGenCUDA::DispatchAllocTensor(const BindNode* op, const CallNode* buffer } RegisterHandleType(buffer.get(), PrimType(dtype)); - if (annotations->dict.count(tvm::tirx::attr::kVolatile)) { + if (annotations->dict.count(tvm::tirx::attr::kVolatile.name)) { MarkVolatile(buffer.get()); } } @@ -1668,7 +1660,7 @@ ffi::Module BuildCUDA(IRModule mod, Target target) { TVM_FFI_ICHECK(base_func->IsInstance()) << "CodeGenCUDA: Can only take Function"; auto function = base_func.as_or_throw(); CallingConv calling_conv = - function->GetAttr(tvm::attr::kCallingConv, CallingConv::kDefault).value(); + function->GetAttr(tvm::attr::kCallingConv, CallingConv::kDefault).value(); TVM_FFI_ICHECK(calling_conv == CallingConv::kDeviceKernelLaunch || calling_conv == CallingConv::kDefault) << "CodeGenCUDA: expect calling_conv equals CallingConv::kDeviceKernelLaunch or " diff --git a/tirx/src/backend/cuda/codegen/codegen_cuda.h b/tirx/src/backend/cuda/codegen/codegen_cuda.h index d133e450e0a0..eabecc1b5e39 100644 --- a/tirx/src/backend/cuda/codegen/codegen_cuda.h +++ b/tirx/src/backend/cuda/codegen/codegen_cuda.h @@ -24,7 +24,7 @@ #ifndef TVM_TARGET_SOURCE_CODEGEN_CUDA_H_ #define TVM_TARGET_SOURCE_CODEGEN_CUDA_H_ -#include +#include #include #include #include diff --git a/tirx/src/backend/cuda/codegen/codegen_cuda_host.cc b/tirx/src/backend/cuda/codegen/codegen_cuda_host.cc index 968061a9fccd..867ec3712390 100644 --- a/tirx/src/backend/cuda/codegen/codegen_cuda_host.cc +++ b/tirx/src/backend/cuda/codegen/codegen_cuda_host.cc @@ -154,9 +154,9 @@ class CodeGenCUDAHost : public CodeGenCHost { // last-error state so a subsequent successful launch is not blamed for it. stream << " } catch (...) {\n cudaGetLastError();\n throw;\n }\n" << " TVM_FFI_SAFE_CALL_END();\n}\n\n"; - if (auto symbol = func->GetAttr(tvm::attr::kGlobalSymbol)) { + if (auto symbol = func->GetAttr(tvm::attr::kGlobalSymbol)) { function_names_.push_back(symbol.value()); - if (func->HasNonzeroAttr(tvm::tirx::attr::kIsEntryFunc) && !has_tvm_ffi_main_func_) { + if (func->HasNonzeroAttr(tvm::tirx::attr::kIsEntryFunc.name) && !has_tvm_ffi_main_func_) { function_names_.push_back(ffi::symbol::tvm_ffi_main); PrintFuncPrefix(stream); PrintType(func->ret_type, stream); diff --git a/tirx/src/backend/cuda/codegen/intrin_rule_cuda.cc b/tirx/src/backend/cuda/codegen/intrin_rule_cuda.cc index 4d52596a468f..c34bff568ef5 100644 --- a/tirx/src/backend/cuda/codegen/intrin_rule_cuda.cc +++ b/tirx/src/backend/cuda/codegen/intrin_rule_cuda.cc @@ -21,14 +21,14 @@ * \file intrin_rule_cuda.cc * \brief CUDA intrinsic rules. */ -#include +#include #include #include #include #include #include +#include #include -#include namespace tvm { namespace codegen { diff --git a/tirx/src/backend/cuda/codegen/llvm/codegen_nvptx.cc b/tirx/src/backend/cuda/codegen/llvm/codegen_nvptx.cc index 692ae3b44f18..cfea6c3699f5 100644 --- a/tirx/src/backend/cuda/codegen/llvm/codegen_nvptx.cc +++ b/tirx/src/backend/cuda/codegen/llvm/codegen_nvptx.cc @@ -145,7 +145,7 @@ class CodeGenNVPTX : public CodeGenLLVM { buf->getType()->getPointerAddressSpace())); TVM_FFI_ICHECK(!var_map_.count(buffer.get())); var_map_[buffer.get()] = buf; - if (annotations->dict.count(tvm::tirx::attr::kVolatile)) { + if (annotations->dict.count(tvm::tirx::attr::kVolatile.name)) { volatile_buf_.insert(buffer.get()); } } diff --git a/tirx/src/backend/cuda/codegen/llvm/intrin_rule_nvptx.cc b/tirx/src/backend/cuda/codegen/llvm/intrin_rule_nvptx.cc index f40d3dab184b..e75eef20787e 100644 --- a/tirx/src/backend/cuda/codegen/llvm/intrin_rule_nvptx.cc +++ b/tirx/src/backend/cuda/codegen/llvm/intrin_rule_nvptx.cc @@ -27,7 +27,7 @@ #include #include #include -#include +#include #include diff --git a/tirx/src/backend/cuda/codegen/target_kind.cc b/tirx/src/backend/cuda/codegen/target_kind.cc index 5add58a024c1..9bd9e4c4f4df 100644 --- a/tirx/src/backend/cuda/codegen/target_kind.cc +++ b/tirx/src/backend/cuda/codegen/target_kind.cc @@ -145,7 +145,7 @@ TVM_FFI_STATIC_INIT_BLOCK() { TargetKindDef("cuda") .set_default_device_type(kDLCUDA) - .def_option>("backend_config") + .def_option>(backend::attr::kBackendConfig.name) .def_option("mcpu") .def_option("arch") .def_option("max_shared_memory_per_block") diff --git a/tirx/src/backend/cuda/op/iket.cc b/tirx/src/backend/cuda/op/iket.cc index 525ded779c2a..f6a3c877636f 100644 --- a/tirx/src/backend/cuda/op/iket.cc +++ b/tirx/src/backend/cuda/op/iket.cc @@ -24,7 +24,7 @@ #include #include #include -#include +#include namespace tvm { namespace backend { diff --git a/tirx/src/backend/cuda/op/target_builtin.cc b/tirx/src/backend/cuda/op/target_builtin.cc index d69f2131a146..0dbaf36fca3d 100644 --- a/tirx/src/backend/cuda/op/target_builtin.cc +++ b/tirx/src/backend/cuda/op/target_builtin.cc @@ -29,7 +29,7 @@ #include #include #include -#include +#include #include #include diff --git a/tirx/src/backend/cuda/op/tcgen05.cc b/tirx/src/backend/cuda/op/tcgen05.cc index a3b4f1f94b92..8e3def9559f7 100644 --- a/tirx/src/backend/cuda/op/tcgen05.cc +++ b/tirx/src/backend/cuda/op/tcgen05.cc @@ -20,7 +20,7 @@ #include #include #include -#include +#include namespace tvm { namespace backend { diff --git a/tirx/src/backend/cuda/op/tensormap.cc b/tirx/src/backend/cuda/op/tensormap.cc index 4b05aa70ef01..3f138b27f894 100644 --- a/tirx/src/backend/cuda/op/tensormap.cc +++ b/tirx/src/backend/cuda/op/tensormap.cc @@ -20,7 +20,7 @@ #include #include #include -#include +#include namespace tvm { namespace backend { diff --git a/tirx/src/backend/cuda/tile/exec_analysis.cc b/tirx/src/backend/cuda/tile/exec_analysis.cc index 2368586a4032..66a3b1000233 100644 --- a/tirx/src/backend/cuda/tile/exec_analysis.cc +++ b/tirx/src/backend/cuda/tile/exec_analysis.cc @@ -29,10 +29,10 @@ #include #include #include +#include #include #include #include -#include #include #include #include diff --git a/tirx/src/backend/cuda/transforms/bind_backend_config.cc b/tirx/src/backend/cuda/transforms/bind_backend_config.cc index 4acc81980e6c..94fa1b94b4e8 100644 --- a/tirx/src/backend/cuda/transforms/bind_backend_config.cc +++ b/tirx/src/backend/cuda/transforms/bind_backend_config.cc @@ -54,7 +54,7 @@ class BackendConfigBinder : public StmtExprMutator { auto body = Mutate(func->body.value(), InplaceMode::kDisallow).ValueOrUnchanged(func->body.value()); func.CopyOnWrite()->body = body; - return WithAttrs(func, {{tvm::attr::kTarget, target_}, {"backend_config", defaults_}}); + return WithAttrs(func, {{tvm::attr::kTarget.name, target_}, {"backend_config", defaults_}}); } private: @@ -67,7 +67,7 @@ class BackendConfigBinder : public StmtExprMutator { defaults_, target_) .cast>(); auto attrs = op->attrs->dict; - attrs.Set(tvm::attr::kTarget, resolved[0]); + attrs.Set(tvm::attr::kTarget.name, resolved[0]); attrs.Set("backend_config", resolved[1]); auto body = Mutate(op->body, InplaceMode::kDisallow).ValueOrUnchanged(op->body); return RegionStmt(op->op, op->args, op->body_params, DictAttrs(attrs), body, op->result_vars, @@ -80,7 +80,7 @@ class BackendConfigBinder : public StmtExprMutator { Pass BindBackendConfig(ffi::String defaults) { auto transform = [defaults](Function func, IRModule mod, PassContext ctx) -> Function { - auto target = func->GetAttr(tvm::attr::kTarget); + auto target = func->GetAttr(tvm::attr::kTarget); if (!target.has_value() || target.value()->kind->name != "cuda") return func; return BackendConfigBinder(target.value(), defaults).Apply(func); }; @@ -118,7 +118,7 @@ class DeviceHelperSpecializer : public StmtExprMutator { if (op->op->name != "tirx.device_entry") { return StmtExprMutator::Mutate_(op, inplace_mode); } - auto entry_target = op->attrs->dict.Get(tvm::attr::kTarget); + auto entry_target = op->attrs->dict.Get(tvm::attr::kTarget.name); if (!entry_target || entry_target.value().cast()->kind->name != "cuda") { return StmtExprMutator::Mutate_(op, inplace_mode); } @@ -146,14 +146,13 @@ class DeviceHelperSpecializer : public StmtExprMutator { if (!active_ || !callee.has_value() || !input_->functions.count(callee.value())) return result; originals.insert(callee.value().get()); auto original = input_->Lookup(callee.value()).as_or_throw(); - auto original_target = original->GetAttr(tvm::attr::kTarget); + auto original_target = original->GetAttr(tvm::attr::kTarget); TVM_FFI_CHECK(original_target.has_value() && original_target.value()->kind->name == "cuda", ValueError) << "CUDA device function cannot call a non-CUDA helper"; - TVM_FFI_CHECK( - original->GetAttr(tvm::attr::kCallingConv, CallingConv::kDefault).value() == - CallingConv::kDefault, - ValueError) + TVM_FFI_CHECK(original->GetAttr(tvm::attr::kCallingConv, CallingConv::kDefault).value() == + CallingConv::kDefault, + ValueError) << "CUDA device calls must refer to private device helpers"; GlobalVar replacement(ffi::UnsafeInit{}); if (auto existing = replacements_.Get(callee.value())) { @@ -166,7 +165,7 @@ class DeviceHelperSpecializer : public StmtExprMutator { : GlobalVar(names_->FreshName(callee.value()->name_hint + "_cuda", false)); replacements_.Set(callee.value(), replacement); auto helper = - WithAttrs(original, {{tvm::attr::kTarget, target_}, {"backend_config", config_}}); + WithAttrs(original, {{tvm::attr::kTarget.name, target_}, {"backend_config", config_}}); helper = Rewrite(helper); output_->Add(replacement, helper); } @@ -200,7 +199,7 @@ Pass SpecializeEntryHelpers() { std::unordered_set originals; for (auto [gv, base] : mod->functions) { auto func = base.as(); - auto target = func ? func.value()->GetAttr(tvm::attr::kTarget) : std::nullopt; + auto target = func ? func.value()->GetAttr(tvm::attr::kTarget) : std::nullopt; if (target && target.value()->kind->name == "cuda") { auto rewriter = ffi::make_object( mod, result, names, target.value(), ffi::String(), false); @@ -249,12 +248,12 @@ Pass SpecializeDeviceHelpers() { result->Add(gv, base); continue; } - auto target = func.value()->GetAttr(tvm::attr::kTarget); + auto target = func.value()->GetAttr(tvm::attr::kTarget); if (!target.has_value() || target.value()->kind->name != "cuda") { result->Add(gv, base); continue; } - auto cc = func.value()->GetAttr(tvm::attr::kCallingConv, CallingConv::kDefault); + auto cc = func.value()->GetAttr(tvm::attr::kCallingConv, CallingConv::kDefault); if (cc.value() != CallingConv::kDeviceKernelLaunch) { if (target.value()->host.has_value()) result->Add(gv, base); continue; diff --git a/tirx/src/backend/cuda/transforms/lower_iket.cc b/tirx/src/backend/cuda/transforms/lower_iket.cc index b51dabc7fe75..fab8df2ba000 100644 --- a/tirx/src/backend/cuda/transforms/lower_iket.cc +++ b/tirx/src/backend/cuda/transforms/lower_iket.cc @@ -652,15 +652,15 @@ class RemoveStrippedIketNoOps : public StmtExprMutator { }; bool IsCudaDeviceFunction(const Function& function) { - auto target = function->GetAttr(tvm::attr::kTarget); + auto target = function->GetAttr(tvm::attr::kTarget); CallingConv calling_conv = - function->GetAttr(tvm::attr::kCallingConv, CallingConv::kDefault).value(); + function->GetAttr(tvm::attr::kCallingConv, CallingConv::kDefault).value(); return target.has_value() && target.value()->kind->name == "cuda" && calling_conv == CallingConv::kDeviceKernelLaunch; } bool IsSm90OrNewer(const Function& function) { - auto target = function->GetAttr(tvm::attr::kTarget); + auto target = function->GetAttr(tvm::attr::kTarget); if (!target.has_value()) return false; auto arch = target.value()->GetAttr("arch"); if (!arch.has_value()) return false; @@ -673,7 +673,7 @@ bool IsSm90OrNewer(const Function& function) { } std::string FunctionName(const GlobalVar& global_var, const Function& function) { - if (auto symbol = function->GetAttr(tvm::attr::kGlobalSymbol)) { + if (auto symbol = function->GetAttr(tvm::attr::kGlobalSymbol)) { return symbol.value(); } return global_var->name_hint; diff --git a/tirx/src/backend/hexagon/codegen/llvm/codegen_hexagon.cc b/tirx/src/backend/hexagon/codegen/llvm/codegen_hexagon.cc index 3a3ae3308f1c..d20683e5d94f 100644 --- a/tirx/src/backend/hexagon/codegen/llvm/codegen_hexagon.cc +++ b/tirx/src/backend/hexagon/codegen/llvm/codegen_hexagon.cc @@ -479,8 +479,8 @@ ffi::Module BuildHexagon(IRModule mod, Target target) { continue; } auto f = kv.second.as_or_throw(); - if (f->HasNonzeroAttr(tvm::tirx::attr::kIsEntryFunc)) { - auto global_symbol = f->GetAttr(tvm::attr::kGlobalSymbol); + if (f->HasNonzeroAttr(tvm::tirx::attr::kIsEntryFunc.name)) { + auto global_symbol = f->GetAttr(tvm::attr::kGlobalSymbol); TVM_FFI_ICHECK(global_symbol.has_value()); entry_func = global_symbol.value(); } diff --git a/tirx/src/backend/hexagon/codegen/llvm/intrin_rule_hexagon.cc b/tirx/src/backend/hexagon/codegen/llvm/intrin_rule_hexagon.cc index 24a43e8b65a1..184e029b5959 100644 --- a/tirx/src/backend/hexagon/codegen/llvm/intrin_rule_hexagon.cc +++ b/tirx/src/backend/hexagon/codegen/llvm/intrin_rule_hexagon.cc @@ -22,7 +22,7 @@ #include #include #include -#include +#include #include "../../../../codegen/llvm/intrin_rule_llvm.h" diff --git a/tirx/src/backend/metal/codegen/codegen_metal.cc b/tirx/src/backend/metal/codegen/codegen_metal.cc index eb15444eb778..231c8239a8c1 100644 --- a/tirx/src/backend/metal/codegen/codegen_metal.cc +++ b/tirx/src/backend/metal/codegen/codegen_metal.cc @@ -113,7 +113,7 @@ void CodeGenMetal::AddFunction(const GlobalVar& gvar, const Function& func) { name_supply_->FreshName("v_"); // add to alloc buffer type. - auto global_symbol = func->GetAttr(tvm::attr::kGlobalSymbol); + auto global_symbol = func->GetAttr(tvm::attr::kGlobalSymbol); TVM_FFI_ICHECK(global_symbol.has_value()) << "CodeGenC: Expect Function to have the global_symbol attribute"; @@ -186,8 +186,7 @@ void CodeGenMetal::AddFunction(const GlobalVar& gvar, const Function& func) { TVM_FFI_ICHECK_EQ(name_supply_->FreshName("threadIdx"), "threadIdx"); TVM_FFI_ICHECK_EQ(name_supply_->FreshName("blockIdx"), "blockIdx"); int work_dim = 0; - auto launch_params = - func->GetAttr>(tvm::tirx::attr::kKernelLaunchParams).value(); + auto launch_params = func->GetAttr(tvm::tirx::attr::kKernelLaunchParams).value(); for (const auto& tag : launch_params) { if (tag != runtime::launch_param::kUseDynamicSharedMemoryTag) { support::GPUThreadTag scope = support::GPUThreadTag::Parse(tag); @@ -439,7 +438,7 @@ void CodeGenMetal::DispatchAllocTensor(const BindNode* op, const CallNode* buffe } RegisterHandleType(buffer.get(), PrimType(dtype)); - if (annotations->dict.count(tvm::tirx::attr::kVolatile)) { + if (annotations->dict.count(tvm::tirx::attr::kVolatile.name)) { MarkVolatile(buffer.get()); } } @@ -590,7 +589,7 @@ ffi::Module BuildMetal(IRModule mod, Target target) { for (auto kv : mod->functions) { TVM_FFI_ICHECK(kv.second->IsInstance()) << "CodeGenMetal: Can only take Function"; - auto global_symbol = kv.second->GetAttr(tvm::attr::kGlobalSymbol); + auto global_symbol = kv.second->GetAttr(tvm::attr::kGlobalSymbol); TVM_FFI_ICHECK(global_symbol.has_value()); std::string func_name = global_symbol.value(); @@ -598,7 +597,7 @@ ffi::Module BuildMetal(IRModule mod, Target target) { CodeGenMetal cg(target); cg.Init(output_ssa); auto f = kv.second.as_or_throw(); - auto calling_conv = f->GetAttr(tvm::attr::kCallingConv); + auto calling_conv = f->GetAttr(tvm::attr::kCallingConv); TVM_FFI_ICHECK(calling_conv.has_value()) << "CodeGenMetal: expected kCallingConv attribute to be set."; TVM_FFI_ICHECK(calling_conv.value() == CallingConv::kDeviceKernelLaunch) diff --git a/tirx/src/backend/metal/codegen/intrin_rule_metal.cc b/tirx/src/backend/metal/codegen/intrin_rule_metal.cc index fc0b7e10ca76..766743b297e6 100644 --- a/tirx/src/backend/metal/codegen/intrin_rule_metal.cc +++ b/tirx/src/backend/metal/codegen/intrin_rule_metal.cc @@ -25,8 +25,8 @@ #include #include #include +#include #include -#include namespace tvm { namespace codegen { diff --git a/tirx/src/backend/metal/op/cooperative_tensor.cc b/tirx/src/backend/metal/op/cooperative_tensor.cc index 60ae754986e4..a3a202ef181c 100644 --- a/tirx/src/backend/metal/op/cooperative_tensor.cc +++ b/tirx/src/backend/metal/op/cooperative_tensor.cc @@ -20,7 +20,7 @@ #include #include #include -#include +#include namespace tvm { namespace backend { diff --git a/tirx/src/backend/metal/op/target_builtin.cc b/tirx/src/backend/metal/op/target_builtin.cc index 527ac50e57f9..2c9bb888b955 100644 --- a/tirx/src/backend/metal/op/target_builtin.cc +++ b/tirx/src/backend/metal/op/target_builtin.cc @@ -27,7 +27,7 @@ #include #include #include -#include +#include namespace tvm { namespace tirx { diff --git a/tirx/src/backend/opencl/codegen/codegen_opencl.cc b/tirx/src/backend/opencl/codegen/codegen_opencl.cc index 365aaa19217c..397860b87192 100644 --- a/tirx/src/backend/opencl/codegen/codegen_opencl.cc +++ b/tirx/src/backend/opencl/codegen/codegen_opencl.cc @@ -775,7 +775,7 @@ ffi::Module BuildOpenCL(IRModule mod, Target target) { TVM_FFI_ICHECK(base_func->IsInstance()) << "CodeGenOpenCL: Can only take Function"; auto function = base_func.as_or_throw(); - auto calling_conv = function->GetAttr(tvm::attr::kCallingConv); + auto calling_conv = function->GetAttr(tvm::attr::kCallingConv); TVM_FFI_ICHECK(calling_conv.has_value()) << "CodeGenOpenCL: expected kCallingConv attribute to be set."; TVM_FFI_ICHECK(calling_conv.value() == CallingConv::kDeviceKernelLaunch) diff --git a/tirx/src/backend/opencl/codegen/intrin_rule_opencl.cc b/tirx/src/backend/opencl/codegen/intrin_rule_opencl.cc index 66eb4c0e8f8a..99b749c1ccd9 100644 --- a/tirx/src/backend/opencl/codegen/intrin_rule_opencl.cc +++ b/tirx/src/backend/opencl/codegen/intrin_rule_opencl.cc @@ -24,7 +24,7 @@ #include #include #include -#include +#include namespace tvm { namespace codegen { diff --git a/tirx/src/backend/opencl/op/memory.cc b/tirx/src/backend/opencl/op/memory.cc index da952705bd64..e37d618e5ec9 100644 --- a/tirx/src/backend/opencl/op/memory.cc +++ b/tirx/src/backend/opencl/op/memory.cc @@ -19,7 +19,7 @@ #include #include -#include +#include namespace tvm { namespace backend { diff --git a/tirx/src/backend/rocm/codegen/llvm/codegen_amdgpu.cc b/tirx/src/backend/rocm/codegen/llvm/codegen_amdgpu.cc index fcf71bfcec17..62141ad59d65 100644 --- a/tirx/src/backend/rocm/codegen/llvm/codegen_amdgpu.cc +++ b/tirx/src/backend/rocm/codegen/llvm/codegen_amdgpu.cc @@ -159,7 +159,7 @@ class CodeGenAMDGPU : public CodeGenLLVM { buf->getType()->getPointerAddressSpace())); TVM_FFI_ICHECK(!var_map_.count(buffer.get())); var_map_[buffer.get()] = buf; - if (annotations->dict.count(tvm::tirx::attr::kVolatile)) { + if (annotations->dict.count(tvm::tirx::attr::kVolatile.name)) { volatile_buf_.insert(buffer.get()); } } diff --git a/tirx/src/backend/rocm/codegen/llvm/intrin_rule_rocm.cc b/tirx/src/backend/rocm/codegen/llvm/intrin_rule_rocm.cc index bb57508b96b4..3322fc40355e 100644 --- a/tirx/src/backend/rocm/codegen/llvm/intrin_rule_rocm.cc +++ b/tirx/src/backend/rocm/codegen/llvm/intrin_rule_rocm.cc @@ -28,9 +28,9 @@ #include #include #include +#include #include #include -#include #include diff --git a/tirx/src/backend/trn/codegen/codegen_trn.cc b/tirx/src/backend/trn/codegen/codegen_trn.cc index 54d45071b5e1..ad64919928ac 100644 --- a/tirx/src/backend/trn/codegen/codegen_trn.cc +++ b/tirx/src/backend/trn/codegen/codegen_trn.cc @@ -106,7 +106,7 @@ void CodeGenTrainium::AddFunction(const GlobalVar& gvar, const Function& func) { name_supply_->FreshName("v_"); // add to alloc buffer type. - auto global_symbol = func->GetAttr(tvm::attr::kGlobalSymbol); + auto global_symbol = func->GetAttr(tvm::attr::kGlobalSymbol); TVM_FFI_ICHECK(global_symbol.has_value()) << "CodeGenC: Expect Function to have the global_symbol attribute"; @@ -114,7 +114,7 @@ void CodeGenTrainium::AddFunction(const GlobalVar& gvar, const Function& func) { this->stream << "def " << static_cast(global_symbol.value()) << "("; // Buffer arguments - auto num_inputs = func->GetAttr(tvm::attr::kNumInputs); + auto num_inputs = func->GetAttr(tvm::attr::kNumInputs); TVM_FFI_ICHECK(num_inputs.has_value()); std::vector output_vids; size_t num_buffer = 0; @@ -303,9 +303,9 @@ void CodeGenTrainium::Dispatch_(const ForNode* op) { TVM_FFI_ICHECK(IsZero(op->min)); if (ctx_.tensorizing) { stream << vid << " = nl.arange(" << extent << ")\n"; - if (op->annotations.count(tvm::backend::trn::attr::kNkiDim)) { + if (op->annotations.count(tvm::backend::trn::attr::kNkiDim.name)) { ctx_.loopvar2dim[op->loop_var.get()] = - op->annotations[tvm::backend::trn::attr::kNkiDim].as_or_throw(); + op->annotations[tvm::backend::trn::attr::kNkiDim.name].as_or_throw(); } ctx_.tensorized_loop_vars.insert(op->loop_var.get()); TVM_FFI_ICHECK(ctx_.loopvar2dim.empty() || @@ -694,7 +694,7 @@ ffi::Module BuildTrainium(IRModule mod, Target target) { for (auto kv : mod->functions) { TVM_FFI_ICHECK(kv.second->IsInstance()) << "CodeGenTrainium: Can only take Function"; - auto global_symbol = kv.second->GetAttr(tvm::attr::kGlobalSymbol); + auto global_symbol = kv.second->GetAttr(tvm::attr::kGlobalSymbol); TVM_FFI_ICHECK(global_symbol.has_value()); std::string func_name = global_symbol.value(); source_maker << "# Function: " << func_name << "\n"; diff --git a/tirx/src/backend/trn/codegen/codegen_trn.h b/tirx/src/backend/trn/codegen/codegen_trn.h index eadd20fc1769..636ea8b6816b 100644 --- a/tirx/src/backend/trn/codegen/codegen_trn.h +++ b/tirx/src/backend/trn/codegen/codegen_trn.h @@ -93,7 +93,7 @@ namespace backend { namespace trn { namespace attr { -constexpr const char* kNkiDim = "nki_dim"; +inline constexpr ffi::reflection::AttrKey kNkiDim{"nki_dim"}; } // namespace attr } // namespace trn } // namespace backend diff --git a/tirx/src/backend/trn/op/target_builtin.cc b/tirx/src/backend/trn/op/target_builtin.cc index c14d86cafa76..3b1674583ceb 100644 --- a/tirx/src/backend/trn/op/target_builtin.cc +++ b/tirx/src/backend/trn/op/target_builtin.cc @@ -29,7 +29,7 @@ #include #include #include -#include +#include #include #include diff --git a/tirx/src/backend/vulkan/codegen/codegen_spirv.cc b/tirx/src/backend/vulkan/codegen/codegen_spirv.cc index 912e0aa1f34f..2b2bfa8fa9e7 100644 --- a/tirx/src/backend/vulkan/codegen/codegen_spirv.cc +++ b/tirx/src/backend/vulkan/codegen/codegen_spirv.cc @@ -117,7 +117,7 @@ runtime::SPIRVShader CodeGenSPIRV::BuildFunction(const Function& f, const std::s TVM_FFI_CHECK(f->body.has_value(), ValueError) << "Kernel code generation requires a function body"; this->InitFuncState(); - TVM_FFI_ICHECK(f->HasNonzeroAttr(tvm::tirx::attr::kNoAlias)) + TVM_FFI_ICHECK(f->HasNonzeroAttr(tvm::tirx::attr::kNoAlias.name)) << "SPIRV only takes restricted memory model"; std::vector pod_args; uint32_t i_buffer = 0; @@ -982,7 +982,7 @@ void CodeGenSPIRV::DispatchAllocTensor(const BindNode* op, const CallNode* buffe TVM_FFI_ICHECK(!var_map_.count(var_node)); var_map_[var_node] = buf; - if (annotations->dict.count(tvm::tirx::attr::kVolatile)) { + if (annotations->dict.count(tvm::tirx::attr::kVolatile.name)) { storage_info_[var_node].is_volatile = true; } } diff --git a/tirx/src/backend/vulkan/codegen/intrin_rule_spirv.cc b/tirx/src/backend/vulkan/codegen/intrin_rule_spirv.cc index 984eb439f87c..578bd90334e6 100644 --- a/tirx/src/backend/vulkan/codegen/intrin_rule_spirv.cc +++ b/tirx/src/backend/vulkan/codegen/intrin_rule_spirv.cc @@ -26,7 +26,7 @@ #include #include #include -#include +#include namespace tvm { namespace codegen { diff --git a/tirx/src/backend/vulkan/codegen/spirv_utils.cc b/tirx/src/backend/vulkan/codegen/spirv_utils.cc index 56b7dab88b97..d8239a4770f5 100644 --- a/tirx/src/backend/vulkan/codegen/spirv_utils.cc +++ b/tirx/src/backend/vulkan/codegen/spirv_utils.cc @@ -128,13 +128,13 @@ std::pair, std::string> Lo auto func = kv.second.as(); TVM_FFI_ICHECK(func); Function f = func.value(); - auto calling_conv = f->GetAttr(tvm::attr::kCallingConv); + auto calling_conv = f->GetAttr(tvm::attr::kCallingConv); TVM_FFI_ICHECK(calling_conv.has_value()) << "CodeGenSPIRV: expected kCallingConv attribute to be set."; TVM_FFI_ICHECK(calling_conv.value() == CallingConv::kDeviceKernelLaunch) << "CodeGenSPIRV: expect calling_conv equals CallingConv::kDeviceKernelLaunch, but got " << static_cast(calling_conv.value()); - auto global_symbol = f->GetAttr(tvm::attr::kGlobalSymbol); + auto global_symbol = f->GetAttr(tvm::attr::kGlobalSymbol); TVM_FFI_ICHECK(global_symbol.has_value()) << "CodeGenSPIRV: Expect Function to have the global_symbol attribute"; diff --git a/tirx/src/backend/webgpu/codegen/codegen_webgpu.cc b/tirx/src/backend/webgpu/codegen/codegen_webgpu.cc index d541a7827c1a..87043b378eeb 100644 --- a/tirx/src/backend/webgpu/codegen/codegen_webgpu.cc +++ b/tirx/src/backend/webgpu/codegen/codegen_webgpu.cc @@ -226,7 +226,7 @@ runtime::FunctionInfo CodeGenWebGPU::AddFunction(const Function& f, bool skip_re TVM_FFI_ICHECK_EQ(name_supply_->FreshName("gridDim"), "gridDim"); // add to alloc buffer type. - auto global_symbol = f->GetAttr(tvm::attr::kGlobalSymbol); + auto global_symbol = f->GetAttr(tvm::attr::kGlobalSymbol); TVM_FFI_ICHECK(global_symbol.has_value()) << "CodeGenWebGPU: Expect Function to have the global_symbol attribute"; @@ -328,7 +328,7 @@ runtime::FunctionInfo CodeGenWebGPU::AddFunction(const Function& f, bool skip_re << "var " << val_pod_args << " : " << type_pod_args << ";\n\n"; // setup thread tags and param access in launch param tags; - if (auto opt = f->GetAttr>(tvm::tirx::attr::kKernelLaunchParams)) { + if (auto opt = f->GetAttr(tvm::tirx::attr::kKernelLaunchParams)) { for (const auto& thread_tag : opt.value()) { func_launch_param_tags.push_back(thread_tag); } @@ -923,13 +923,13 @@ ffi::Module BuildWebGPU(IRModule mod, Target target) { TVM_FFI_ICHECK(kv.second->IsInstance()) << "CodeGenWebGPU: Can only take Function"; auto f = kv.second.as_or_throw(); - auto calling_conv = f->GetAttr(tvm::attr::kCallingConv); + auto calling_conv = f->GetAttr(tvm::attr::kCallingConv); TVM_FFI_ICHECK(calling_conv.has_value()) << "CodeGenWebGPU: expected kCallingConv attribute to be set."; TVM_FFI_ICHECK(calling_conv.value() == CallingConv::kDeviceKernelLaunch) << "CodeGenWebGPU: expect calling_conv equals CallingConv::kDeviceKernelLaunch, but got " << static_cast(calling_conv.value()); - auto global_symbol = f->GetAttr(tvm::attr::kGlobalSymbol); + auto global_symbol = f->GetAttr(tvm::attr::kGlobalSymbol); TVM_FFI_ICHECK(global_symbol.has_value()) << "CodeGenWebGPU: Expect Function to have the global_symbol attribute"; std::string f_name = global_symbol.value(); diff --git a/tirx/src/backend/webgpu/codegen/intrin_rule_webgpu.cc b/tirx/src/backend/webgpu/codegen/intrin_rule_webgpu.cc index f5cee178b919..012a02f3b0bd 100644 --- a/tirx/src/backend/webgpu/codegen/intrin_rule_webgpu.cc +++ b/tirx/src/backend/webgpu/codegen/intrin_rule_webgpu.cc @@ -25,8 +25,8 @@ #include #include #include +#include #include -#include namespace tvm { namespace codegen { diff --git a/tirx/src/codegen/build_common.h b/tirx/src/codegen/build_common.h index 0fb32f3c1909..60a148b52c30 100644 --- a/tirx/src/codegen/build_common.h +++ b/tirx/src/codegen/build_common.h @@ -73,12 +73,12 @@ inline ffi::Map ExtractFuncInfo(const IRModu : runtime::ArgExtraTags::kNone); } ffi::Array launch_param_tags; - if (auto opt = f->GetAttr>(tvm::tirx::attr::kKernelLaunchParams)) { + if (auto opt = f->GetAttr(tvm::tirx::attr::kKernelLaunchParams)) { for (const auto& tag : opt.value()) { launch_param_tags.push_back(tag); } } - auto global_symbol = f->GetAttr(tvm::attr::kGlobalSymbol); + auto global_symbol = f->GetAttr(tvm::attr::kGlobalSymbol); if (global_symbol) { fmap.Set( global_symbol.value(), diff --git a/tirx/src/codegen/c/codegen_c.cc b/tirx/src/codegen/c/codegen_c.cc index 71c25ee35e6f..23a74b55686b 100644 --- a/tirx/src/codegen/c/codegen_c.cc +++ b/tirx/src/codegen/c/codegen_c.cc @@ -124,7 +124,7 @@ void CodeGenC::PrintFunctionParameters(const Function& func, std::ostream& os) { PrintType(v->ty, os); } - bool no_alias = func->HasNonzeroAttr(tvm::tirx::attr::kNoAlias); + bool no_alias = func->HasNonzeroAttr(tvm::tirx::attr::kNoAlias.name); bool is_handle = v->ty.as(); auto* ptr = v->ty.as(); if (ptr && ptr->element_type.as()) { @@ -154,7 +154,7 @@ bool CodeGenC::RegisterFunctionName(const GlobalVar& gvar, const Function& func) if (internal_functions_.count(gvar)) return false; auto function_name = [&]() -> ffi::String { - if (auto global_symbol = func->GetAttr(tvm::attr::kGlobalSymbol)) { + if (auto global_symbol = func->GetAttr(tvm::attr::kGlobalSymbol)) { auto name = global_symbol.value(); TVM_FFI_ICHECK(!func_name_supply_->ContainsName(name)) << "Function " << gvar << " must use global symbol " << name @@ -1307,7 +1307,7 @@ void CodeGenC::DispatchAllocTensor(const BindNode* op, const CallNode* buffer_ca stream << ' ' << vid << '[' << constant_size << "];\n"; RegisterHandleType(buffer.get(), PrimType(dtype)); - if (annotations->dict.count(tvm::tirx::attr::kVolatile)) { + if (annotations->dict.count(tvm::tirx::attr::kVolatile.name)) { MarkVolatile(buffer.get()); } } diff --git a/tirx/src/codegen/c/codegen_c.h b/tirx/src/codegen/c/codegen_c.h index 68a97e54a56c..4eca8d29353b 100644 --- a/tirx/src/codegen/c/codegen_c.h +++ b/tirx/src/codegen/c/codegen_c.h @@ -33,7 +33,7 @@ #include #include #include -#include +#include #include #include @@ -384,7 +384,7 @@ class CodeGenC : public tvm::ExprFunctor, /* \brief Map of GlobalVar to their symbol. * * For externally-exposed functions, this is given by the - * tvm::attr::kTarget attribute of the Function. For internal + * tvm::attr::kTarget.name attribute of the Function. For internal * functions, this is the name of the function's GlobalVar, possibly * altered to prevent duplicate names. */ diff --git a/tirx/src/codegen/c/codegen_c_host.cc b/tirx/src/codegen/c/codegen_c_host.cc index 0d0097bcdaa3..cc37cb66c8eb 100644 --- a/tirx/src/codegen/c/codegen_c_host.cc +++ b/tirx/src/codegen/c/codegen_c_host.cc @@ -72,14 +72,14 @@ void CodeGenCHost::AddFunction(const GlobalVar& gvar, const Function& func) { void CodeGenCHost::AddFunction(const GlobalVar& gvar, const Function& func, bool emit_fwd_func_decl) { - auto global_symbol = func->GetAttr(tvm::attr::kGlobalSymbol); + auto global_symbol = func->GetAttr(tvm::attr::kGlobalSymbol); if (global_symbol) { function_names_.push_back(global_symbol.value()); } emit_fwd_func_decl_ = emit_fwd_func_decl; CodeGenC::AddFunction(gvar, func); - if (func->HasNonzeroAttr(tvm::tirx::attr::kIsEntryFunc) && !has_tvm_ffi_main_func_) { + if (func->HasNonzeroAttr(tvm::tirx::attr::kIsEntryFunc.name) && !has_tvm_ffi_main_func_) { TVM_FFI_ICHECK(global_symbol.has_value()) << "CodeGenCHost: The entry func must have the global_symbol attribute, " << "but function " << gvar << " only has attributes " << func->attrs; @@ -379,9 +379,9 @@ ffi::Module BuildCHost(IRModule mod, Target target) { bool emit_fwd_func_decl = true; std::unordered_set devices; - if (mod->GetAttr>(tvm::attr::kDeviceContexts) != nullptr) { + if (mod->GetAttr(tvm::attr::kDeviceContexts) != nullptr) { ffi::Map device_contexts = - mod->GetAttr>(tvm::attr::kDeviceContexts).value(); + mod->GetAttr(tvm::attr::kDeviceContexts).value(); for (auto const& context : device_contexts) { devices.insert(context.second.data()); } diff --git a/tirx/src/codegen/intrin_rule.cc b/tirx/src/codegen/intrin_rule.cc index 7bad8e63284f..0e94a1e8daf7 100644 --- a/tirx/src/codegen/intrin_rule.cc +++ b/tirx/src/codegen/intrin_rule.cc @@ -25,8 +25,8 @@ #include #include #include +#include #include -#include namespace tvm { namespace codegen { diff --git a/tirx/src/codegen/llvm/codegen_aarch64.cc b/tirx/src/codegen/llvm/codegen_aarch64.cc index 407605e30a60..c3fb8037d13b 100644 --- a/tirx/src/codegen/llvm/codegen_aarch64.cc +++ b/tirx/src/codegen/llvm/codegen_aarch64.cc @@ -56,7 +56,7 @@ class CodeGenAArch64 final : public CodeGenCPU { llvm::Function* CodeGenAArch64::DeclareFunction(const GlobalVar& gvar, const Function& f) { llvm::Function* func = CodeGenCPU::DeclareFunction(gvar, f); - if (f->GetAttr(tvm::attr::kCallingConv) != CallingConv::kCPackedFunc) { + if (f->GetAttr(tvm::attr::kCallingConv) != CallingConv::kCPackedFunc) { SetPStateAttributes(func, f); } return func; @@ -65,9 +65,8 @@ llvm::Function* CodeGenAArch64::DeclareFunction(const GlobalVar& gvar, const Fun void CodeGenAArch64::AddFunction(const GlobalVar& gvar, const Function& f) { llvm::Function* func = DeclareFunction(gvar, f); current_function_ = f; - packed_function_ = f->GetAttr(tvm::attr::kCallingConv) == CallingConv::kCPackedFunc - ? func - : nullptr; + packed_function_ = + f->GetAttr(tvm::attr::kCallingConv) == CallingConv::kCPackedFunc ? func : nullptr; CodeGenCPU::AddFunction(gvar, f); current_function_ = std::nullopt; packed_function_ = nullptr; diff --git a/tirx/src/codegen/llvm/codegen_cpu.cc b/tirx/src/codegen/llvm/codegen_cpu.cc index 60688548cdb9..1a43d90c97ce 100644 --- a/tirx/src/codegen/llvm/codegen_cpu.cc +++ b/tirx/src/codegen/llvm/codegen_cpu.cc @@ -226,7 +226,7 @@ llvm::DISubprogram* CodeGenCPU::CreateDebugFunction(llvm::StringRef name, } llvm::DISubprogram* CodeGenCPU::CreateDebugFunction(const GlobalVar& gvar, const Function& func) { - std::string name = func->GetAttr(tvm::attr::kGlobalSymbol).value_or(gvar->name_hint); + std::string name = func->GetAttr(tvm::attr::kGlobalSymbol).value_or(gvar->name_hint); return CreateDebugFunction(name, func->params.Map([](const Var& var) { return var->ty; }), func->ret_type); } @@ -236,7 +236,7 @@ void CodeGenCPU::AddFunction(const GlobalVar& gvar, const Function& func) { EmitDebugLocation(func->loc); CodeGenLLVM::AddFunction(gvar, func); if (f_tvm_register_system_symbol_ != nullptr) { - if (auto global_symbol = func->GetAttr(tvm::attr::kGlobalSymbol)) { + if (auto global_symbol = func->GetAttr(tvm::attr::kGlobalSymbol)) { export_system_symbols_.emplace_back( std::make_pair(global_symbol.value().operator std::string(), function_)); } diff --git a/tirx/src/codegen/llvm/codegen_llvm.cc b/tirx/src/codegen/llvm/codegen_llvm.cc index 3634ff7dd118..a602811097ba 100644 --- a/tirx/src/codegen/llvm/codegen_llvm.cc +++ b/tirx/src/codegen/llvm/codegen_llvm.cc @@ -277,7 +277,7 @@ void CodeGenLLVM::InitFuncState() { std::tuple CodeGenLLVM::GetLinkage( const GlobalVar& gvar, const Function& func) { - if (auto global_symbol = func->GetAttr(tvm::attr::kGlobalSymbol)) { + if (auto global_symbol = func->GetAttr(tvm::attr::kGlobalSymbol)) { return {global_symbol.value(), llvm::Function::ExternalLinkage}; } @@ -302,7 +302,7 @@ llvm::Function* CodeGenLLVM::DeclareFunctionInternal(const GlobalVar& gvar, cons } std::vector param_types; - is_restricted_ = func->HasNonzeroAttr(tvm::tirx::attr::kNoAlias); + is_restricted_ = func->HasNonzeroAttr(tvm::tirx::attr::kNoAlias.name); for (Var param : func->params) { param_types.push_back(GetLLVMType(param->ty)); if (!is_restricted_ && param->ty.as()) { @@ -2261,7 +2261,7 @@ void CodeGenLLVM::DispatchAllocTensor(const BindNode* op, const CallNode* buffer TVM_FFI_ICHECK(!var_map_.count(buffer.get())); var_map_[buffer.get()] = buf; - if (annotations->dict.count(tvm::tirx::attr::kVolatile)) { + if (annotations->dict.count(tvm::tirx::attr::kVolatile.name)) { volatile_buf_.insert(buffer.get()); } } diff --git a/tirx/src/codegen/llvm/codegen_llvm.h b/tirx/src/codegen/llvm/codegen_llvm.h index f9bdbf8e001b..9d5a05d4e719 100644 --- a/tirx/src/codegen/llvm/codegen_llvm.h +++ b/tirx/src/codegen/llvm/codegen_llvm.h @@ -49,7 +49,7 @@ #include #include #include -#include +#include #include #include diff --git a/tirx/src/codegen/llvm/intrin_rule_llvm.cc b/tirx/src/codegen/llvm/intrin_rule_llvm.cc index 1a840ef635a6..6fd4c3e19f52 100644 --- a/tirx/src/codegen/llvm/intrin_rule_llvm.cc +++ b/tirx/src/codegen/llvm/intrin_rule_llvm.cc @@ -30,7 +30,7 @@ #include #include #include -#include +#include #include diff --git a/tirx/src/codegen/llvm/llvm_module.cc b/tirx/src/codegen/llvm/llvm_module.cc index e9cbacbeebb4..02cb1da865da 100644 --- a/tirx/src/codegen/llvm/llvm_module.cc +++ b/tirx/src/codegen/llvm/llvm_module.cc @@ -255,8 +255,7 @@ void LLVMModuleObj::Init(const IRModule& mod, const Target& target) { std::string entry_func; - ffi::Optional system_lib_prefix = - mod->GetAttr(tvm::attr::kSystemLibPrefix); + ffi::Optional system_lib_prefix = mod->GetAttr(tvm::attr::kSystemLibPrefix); for (auto kv : mod->functions) { if (!kv.second->IsInstance()) { @@ -264,8 +263,8 @@ void LLVMModuleObj::Init(const IRModule& mod, const Target& target) { continue; } auto f = kv.second.as_or_throw(); - auto global_symbol = f->GetAttr(tvm::attr::kGlobalSymbol); - bool is_entry_func = f->HasNonzeroAttr(tvm::tirx::attr::kIsEntryFunc); + auto global_symbol = f->GetAttr(tvm::attr::kGlobalSymbol); + bool is_entry_func = f->HasNonzeroAttr(tvm::tirx::attr::kIsEntryFunc.name); TVM_FFI_ICHECK(global_symbol || !is_entry_func) << "The entry func must be exposed externally."; diff --git a/tirx/src/ir/expr.cc b/tirx/src/ir/expr.cc index 7c43a89c3da2..4cc801ae5213 100644 --- a/tirx/src/ir/expr.cc +++ b/tirx/src/ir/expr.cc @@ -701,7 +701,7 @@ TensorLoad::TensorLoad(Expr source, ffi::Array indices, Location loc) TVM_FFI_CHECK(source.defined() && source->ty.defined(), TypeError) << "TensorLoad requires a source with a defined type"; namespace refl = ffi::reflection; - static const refl::TypeAttrColumn infer_column(type_attr::kTensorLoadInferType); + static const refl::TypeAttrColumn infer_column(type_attr::kTensorLoadInferType.name); auto infer = infer_column[source->ty->type_index()]; TVM_FFI_CHECK(infer != nullptr, TypeError) << "Type " << source->ty->GetTypeKey() << " does not support TensorLoad"; @@ -729,7 +729,7 @@ TensorLoad TensorLoad::Unchecked(PrimType ty, Expr source, ffi::Array TVM_FFI_STATIC_INIT_BLOCK() { namespace refl = tvm::ffi::reflection; - refl::EnsureTypeAttrColumn(type_attr::kTensorLoadInferType); + refl::EnsureTypeAttrColumn(type_attr::kTensorLoadInferType.name); refl::GlobalDef().def("ir.TensorLoad", [](Expr source, ffi::Array indices, Location loc) { return TensorLoad(source, indices, loc); diff --git a/tirx/src/ir/module.cc b/tirx/src/ir/module.cc index b140a46589a3..a1f3d2f24bc5 100644 --- a/tirx/src/ir/module.cc +++ b/tirx/src/ir/module.cc @@ -230,7 +230,7 @@ IRModule IRModule::FromExpr(const Expr& expr, // All global definitions must be functions. BaseFunc func = expr.as_or_throw(); { - if (auto opt = func->GetAttr(tvm::attr::kGlobalSymbol)) { + if (auto opt = func->GetAttr(tvm::attr::kGlobalSymbol)) { // Function literal has been annotated with it's required global symbol. gv_name = opt.value(); } diff --git a/tirx/src/ir/prim/builtin.cc b/tirx/src/ir/prim/builtin.cc index 692d3232b231..5ef75f691b9d 100644 --- a/tirx/src/ir/prim/builtin.cc +++ b/tirx/src/ir/prim/builtin.cc @@ -19,7 +19,7 @@ #include #include #include -#include +#include namespace tvm { namespace prim { diff --git a/tirx/src/ir/prim/math.cc b/tirx/src/ir/prim/math.cc index ffd57eb52b67..96a98d1fecbf 100644 --- a/tirx/src/ir/prim/math.cc +++ b/tirx/src/ir/prim/math.cc @@ -22,7 +22,7 @@ #include #include #include -#include +#include #include diff --git a/tirx/src/ir/stmt.cc b/tirx/src/ir/stmt.cc index 94b8b7673e42..1396b7eb69a3 100644 --- a/tirx/src/ir/stmt.cc +++ b/tirx/src/ir/stmt.cc @@ -650,8 +650,8 @@ TVM_FFI_INLINE ffi::Expected> TensorStoreMaybeInplace TVM_FFI_STATIC_INIT_BLOCK() { StmtNode::RegisterReflection(); - ffi::reflection::EnsureTypeAttrColumn(tvm::type_attr::kTensorStoreValidate); - ffi::reflection::EnsureTypeAttrColumn(tvm::type_attr::kEvaluateValidate); + ffi::reflection::EnsureTypeAttrColumn(tvm::type_attr::kTensorStoreValidate.name); + ffi::reflection::EnsureTypeAttrColumn(tvm::type_attr::kEvaluateValidate.name); } // Bind @@ -1065,7 +1065,8 @@ TVM_FFI_STATIC_INIT_BLOCK() { // Evaluate Evaluate::Evaluate(Expr value, Location loc) : Stmt(ffi::UnsafeInit{}) { TVM_FFI_ICHECK(value.defined()); - static const ffi::reflection::TypeAttrColumn validate_column(tvm::type_attr::kEvaluateValidate); + static const ffi::reflection::TypeAttrColumn validate_column( + tvm::type_attr::kEvaluateValidate.name); if (auto validate = validate_column[value->ty->type_index()]; validate != nullptr) { validate.cast()(value); } @@ -1093,7 +1094,7 @@ TVM_FFI_STATIC_INIT_BLOCK() { TensorStore::TensorStore(Expr dest, ffi::Array indices, PrimExpr value, Location loc) : Stmt(ffi::UnsafeInit{}) { namespace refl = ffi::reflection; - static const refl::TypeAttrColumn validate_column(tvm::type_attr::kTensorStoreValidate); + static const refl::TypeAttrColumn validate_column(tvm::type_attr::kTensorStoreValidate.name); auto validate = validate_column[dest->ty->type_index()]; TVM_FFI_CHECK(validate != nullptr, TypeError) << "Type " << dest->ty->GetTypeKey() << " does not support TensorStore"; diff --git a/tirx/src/ir/subscript_expr.cc b/tirx/src/ir/subscript_expr.cc index 24ef8ae529a0..7409d1a3056b 100644 --- a/tirx/src/ir/subscript_expr.cc +++ b/tirx/src/ir/subscript_expr.cc @@ -34,9 +34,9 @@ using SubscriptSlice = ffi::Array().def( - tvm::type_attr::kSubscriptExprRealize, + tvm::type_attr::kSubscriptExprRealize.name, [](Expr value, SubscriptSlice slice, Location loc) -> ffi::ObjectRef { TVM_FFI_CHECK_EQ(slice.size(), 1, IndexError) << "A tuple expression requires exactly one index"; @@ -51,7 +51,7 @@ TVM_FFI_STATIC_INIT_BLOCK() { "ir.SubscriptExprRealize", [](Expr value, SubscriptSlice slice, Location loc) -> ffi::ObjectRef { TVM_FFI_CHECK(value.defined(), TypeError) << "Cannot subscript an undefined expression"; - static refl::TypeAttrColumn realize_column(tvm::type_attr::kSubscriptExprRealize); + static refl::TypeAttrColumn realize_column(tvm::type_attr::kSubscriptExprRealize.name); ffi::AnyView packed_realize = realize_column[value->ty->type_index()]; TVM_FFI_CHECK(packed_realize != nullptr, TypeError) << "Type " << value->ty->GetTypeKey() << " does not support subscript"; diff --git a/tirx/src/script/printer/doc_translator.cc b/tirx/src/script/printer/doc_translator.cc index cddaf50fe7ea..c6fd4511243f 100644 --- a/tirx/src/script/printer/doc_translator.cc +++ b/tirx/src/script/printer/doc_translator.cc @@ -239,7 +239,7 @@ Doc DocTranslate(ffi::AnyView ir, ffi::Dict* doc_origins, TVM_FFI_STATIC_INIT_BLOCK() { namespace refl = ffi::reflection; refl::ObjectDef(); - refl::EnsureTypeAttrColumn(tvm::script::printer::type_attr::kDocTranslate); + refl::EnsureTypeAttrColumn(tvm::script::printer::type_attr::kDocTranslate.name); } } // namespace printer diff --git a/tirx/src/script/printer/ir/call.cc b/tirx/src/script/printer/ir/call.cc index 7d6b557c0e16..a66749b24773 100644 --- a/tirx/src/script/printer/ir/call.cc +++ b/tirx/src/script/printer/ir/call.cc @@ -62,7 +62,7 @@ ffi::Optional CallDocTranslate(DocTranslatorObj* d, ffi::AnyView input, } TVM_FFI_STATIC_INIT_BLOCK() { - ffi::reflection::TypeAttrDef().attr(tvm::script::printer::type_attr::kDocTranslate, + ffi::reflection::TypeAttrDef().attr(tvm::script::printer::type_attr::kDocTranslate.name, FDocTranslate::FromNative<&CallDocTranslate>()); } diff --git a/tirx/src/script/printer/ir/ir.cc b/tirx/src/script/printer/ir/ir.cc index 06ddd275b85c..da5ab2551b54 100644 --- a/tirx/src/script/printer/ir/ir.cc +++ b/tirx/src/script/printer/ir/ir.cc @@ -110,7 +110,7 @@ ffi::Optional IRModuleDocTranslate(DocTranslatorObj* d, ffi::AnyView in TVM_FFI_STATIC_INIT_BLOCK() { RegisterNamespaceAlias("ir.prefix", "I"); ffi::reflection::TypeAttrDef().attr( - tvm::script::printer::type_attr::kDocTranslate, + tvm::script::printer::type_attr::kDocTranslate.name, FDocTranslate::FromNative<&IRModuleDocTranslate>()); } @@ -123,7 +123,7 @@ ffi::Optional DictAttrsDocTranslate(DocTranslatorObj* d, ffi::AnyView i TVM_FFI_STATIC_INIT_BLOCK() { ffi::reflection::TypeAttrDef().attr( - tvm::script::printer::type_attr::kDocTranslate, + tvm::script::printer::type_attr::kDocTranslate.name, FDocTranslate::FromNative<&DictAttrsDocTranslate>()); } @@ -136,7 +136,7 @@ ffi::Optional GlobalVarDocTranslate(DocTranslatorObj* d, ffi::AnyView i TVM_FFI_STATIC_INIT_BLOCK() { ffi::reflection::TypeAttrDef().attr( - tvm::script::printer::type_attr::kDocTranslate, + tvm::script::printer::type_attr::kDocTranslate.name, FDocTranslate::FromNative<&GlobalVarDocTranslate>()); } @@ -147,7 +147,7 @@ ffi::Optional OpDocTranslate(DocTranslatorObj*, ffi::AnyView input, con } TVM_FFI_STATIC_INIT_BLOCK() { - ffi::reflection::TypeAttrDef().attr(tvm::script::printer::type_attr::kDocTranslate, + ffi::reflection::TypeAttrDef().attr(tvm::script::printer::type_attr::kDocTranslate.name, FDocTranslate::FromNative<&OpDocTranslate>()); } @@ -161,8 +161,9 @@ ffi::Optional RangeDocTranslate(DocTranslatorObj* d, ffi::AnyView input } TVM_FFI_STATIC_INIT_BLOCK() { - ffi::reflection::TypeAttrDef().attr(tvm::script::printer::type_attr::kDocTranslate, - FDocTranslate::FromNative<&RangeDocTranslate>()); + ffi::reflection::TypeAttrDef().attr( + tvm::script::printer::type_attr::kDocTranslate.name, + FDocTranslate::FromNative<&RangeDocTranslate>()); } } // namespace diff --git a/tirx/src/script/printer/ir/misc.cc b/tirx/src/script/printer/ir/misc.cc index 1d5d20752cf0..554c83901234 100644 --- a/tirx/src/script/printer/ir/misc.cc +++ b/tirx/src/script/printer/ir/misc.cc @@ -47,7 +47,7 @@ ffi::Optional DataTypeImmDocTranslate(DocTranslatorObj*, ffi::AnyView i TVM_FFI_STATIC_INIT_BLOCK() { ffi::reflection::TypeAttrDef().attr( - tvm::script::printer::type_attr::kDocTranslate, + tvm::script::printer::type_attr::kDocTranslate.name, FDocTranslate::FromNative<&DataTypeImmDocTranslate>()); } @@ -98,15 +98,16 @@ ffi::Optional MetadataDocTranslate(DocTranslatorObj* d, ffi::AnyView in TVM_FFI_STATIC_INIT_BLOCK() { ffi::reflection::TypeAttrDef().attr( - tvm::script::printer::type_attr::kDocTranslate, + tvm::script::printer::type_attr::kDocTranslate.name, FDocTranslate::FromNative<&ArrayDocTranslate>()); - ffi::reflection::TypeAttrDef().attr(tvm::script::printer::type_attr::kDocTranslate, - FDocTranslate::FromNative<&MapDocTranslate>()); + ffi::reflection::TypeAttrDef().attr( + tvm::script::printer::type_attr::kDocTranslate.name, + FDocTranslate::FromNative<&MapDocTranslate>()); ffi::reflection::TypeAttrDef().attr( - tvm::script::printer::type_attr::kDocTranslate, + tvm::script::printer::type_attr::kDocTranslate.name, FDocTranslate::FromNative<&MetadataDocTranslate>()); ffi::reflection::TypeAttrDef().attr( - tvm::script::printer::type_attr::kDocTranslate, + tvm::script::printer::type_attr::kDocTranslate.name, FDocTranslate::FromNative<&MetadataDocTranslate>()); } @@ -127,7 +128,7 @@ ffi::Optional ShapeDocTranslate(DocTranslatorObj*, ffi::AnyView input, TVM_FFI_STATIC_INIT_BLOCK() { ffi::reflection::TypeAttrDef().attr( - tvm::script::printer::type_attr::kDocTranslate, + tvm::script::printer::type_attr::kDocTranslate.name, FDocTranslate::FromNative<&ShapeDocTranslate>()); } diff --git a/tirx/src/script/printer/ir/prim_expr.cc b/tirx/src/script/printer/ir/prim_expr.cc index b6e83f854db9..4536993a20e8 100644 --- a/tirx/src/script/printer/ir/prim_expr.cc +++ b/tirx/src/script/printer/ir/prim_expr.cc @@ -48,7 +48,7 @@ void EmitVarDefinition(DocTranslatorObj* d, const Var& var, ExprDoc rhs) { } TVM_FFI_STATIC_INIT_BLOCK() { - ffi::reflection::TypeAttrDef().attr(tvm::script::printer::type_attr::kDocTranslate, + ffi::reflection::TypeAttrDef().attr(tvm::script::printer::type_attr::kDocTranslate.name, FDocTranslate::FromNative<&VarDocTranslate>()); } @@ -70,7 +70,7 @@ ffi::Optional LambdaExprDocTranslate(DocTranslatorObj* d, ffi::AnyView TVM_FFI_STATIC_INIT_BLOCK() { ffi::reflection::TypeAttrDef().attr( - tvm::script::printer::type_attr::kDocTranslate, + tvm::script::printer::type_attr::kDocTranslate.name, FDocTranslate::FromNative<&LambdaExprDocTranslate>()); } @@ -171,7 +171,7 @@ ffi::Optional TensorRegionDocTranslate(DocTranslatorObj* d, ffi::AnyVie TVM_FFI_STATIC_INIT_BLOCK() { ffi::reflection::TypeAttrDef().attr( - tvm::script::printer::type_attr::kDocTranslate, + tvm::script::printer::type_attr::kDocTranslate.name, FDocTranslate::FromNative<&TensorRegionDocTranslate>()); } @@ -184,7 +184,7 @@ ffi::Optional TensorLoadDocTranslate(DocTranslatorObj* d, ffi::AnyView TVM_FFI_STATIC_INIT_BLOCK() { ffi::reflection::TypeAttrDef().attr( - tvm::script::printer::type_attr::kDocTranslate, + tvm::script::printer::type_attr::kDocTranslate.name, FDocTranslate::FromNative<&TensorLoadDocTranslate>()); } @@ -198,8 +198,9 @@ ffi::Optional TupleDocTranslate(DocTranslatorObj* d, ffi::AnyView input } TVM_FFI_STATIC_INIT_BLOCK() { - ffi::reflection::TypeAttrDef().attr(tvm::script::printer::type_attr::kDocTranslate, - FDocTranslate::FromNative<&TupleDocTranslate>()); + ffi::reflection::TypeAttrDef().attr( + tvm::script::printer::type_attr::kDocTranslate.name, + FDocTranslate::FromNative<&TupleDocTranslate>()); } ffi::Optional TupleGetItemDocTranslate(DocTranslatorObj* d, ffi::AnyView input, @@ -211,7 +212,7 @@ ffi::Optional TupleGetItemDocTranslate(DocTranslatorObj* d, ffi::AnyVie TVM_FFI_STATIC_INIT_BLOCK() { ffi::reflection::TypeAttrDef().attr( - tvm::script::printer::type_attr::kDocTranslate, + tvm::script::printer::type_attr::kDocTranslate.name, FDocTranslate::FromNative<&TupleGetItemDocTranslate>()); } @@ -280,7 +281,7 @@ ffi::Optional BitwiseNotDocTranslate(DocTranslatorObj* d, ffi::AnyView TVM_FFI_STATIC_INIT_BLOCK() { ffi::reflection::TypeAttrDef().attr( - tvm::script::printer::type_attr::kDocTranslate, + tvm::script::printer::type_attr::kDocTranslate.name, FDocTranslate::FromNative<&BitwiseNotDocTranslate>()); } @@ -293,8 +294,9 @@ ffi::Optional NotDocTranslate(DocTranslatorObj* d, ffi::AnyView input, } TVM_FFI_STATIC_INIT_BLOCK() { - ffi::reflection::TypeAttrDef().attr(tvm::script::printer::type_attr::kDocTranslate, - FDocTranslate::FromNative<&NotDocTranslate>()); + ffi::reflection::TypeAttrDef().attr( + tvm::script::printer::type_attr::kDocTranslate.name, + FDocTranslate::FromNative<&NotDocTranslate>()); } ffi::Optional StringImmDocTranslate(DocTranslatorObj* d, ffi::AnyView input, @@ -310,7 +312,7 @@ ffi::Optional StringImmDocTranslate(DocTranslatorObj* d, ffi::AnyView i TVM_FFI_STATIC_INIT_BLOCK() { ffi::reflection::TypeAttrDef().attr( - tvm::script::printer::type_attr::kDocTranslate, + tvm::script::printer::type_attr::kDocTranslate.name, FDocTranslate::FromNative<&StringImmDocTranslate>()); } @@ -325,7 +327,7 @@ ffi::Optional CastDocTranslate(DocTranslatorObj* d, ffi::AnyView input, TVM_FFI_STATIC_INIT_BLOCK() { ffi::reflection::TypeAttrDef().attr( - tvm::script::printer::type_attr::kDocTranslate, + tvm::script::printer::type_attr::kDocTranslate.name, FDocTranslate::FromNative<&CastDocTranslate>()); } @@ -340,7 +342,7 @@ ffi::Optional SelectDocTranslate(DocTranslatorObj* d, ffi::AnyView inpu TVM_FFI_STATIC_INIT_BLOCK() { ffi::reflection::TypeAttrDef().attr( - tvm::script::printer::type_attr::kDocTranslate, + tvm::script::printer::type_attr::kDocTranslate.name, FDocTranslate::FromNative<&SelectDocTranslate>()); } @@ -355,7 +357,7 @@ ffi::Optional RampDocTranslate(DocTranslatorObj* d, ffi::AnyView input, TVM_FFI_STATIC_INIT_BLOCK() { ffi::reflection::TypeAttrDef().attr( - tvm::script::printer::type_attr::kDocTranslate, + tvm::script::printer::type_attr::kDocTranslate.name, FDocTranslate::FromNative<&RampDocTranslate>()); } @@ -370,7 +372,7 @@ ffi::Optional BroadcastDocTranslate(DocTranslatorObj* d, ffi::AnyView i TVM_FFI_STATIC_INIT_BLOCK() { ffi::reflection::TypeAttrDef().attr( - tvm::script::printer::type_attr::kDocTranslate, + tvm::script::printer::type_attr::kDocTranslate.name, FDocTranslate::FromNative<&BroadcastDocTranslate>()); } @@ -387,7 +389,7 @@ ffi::Optional ShuffleDocTranslate(DocTranslatorObj* d, ffi::AnyView inp TVM_FFI_STATIC_INIT_BLOCK() { ffi::reflection::TypeAttrDef().attr( - tvm::script::printer::type_attr::kDocTranslate, + tvm::script::printer::type_attr::kDocTranslate.name, FDocTranslate::FromNative<&ShuffleDocTranslate>()); } @@ -408,8 +410,9 @@ ffi::Optional LetDocTranslate(DocTranslatorObj* d, ffi::AnyView input, } TVM_FFI_STATIC_INIT_BLOCK() { - ffi::reflection::TypeAttrDef().attr(tvm::script::printer::type_attr::kDocTranslate, - FDocTranslate::FromNative<&LetDocTranslate>()); + ffi::reflection::TypeAttrDef().attr( + tvm::script::printer::type_attr::kDocTranslate.name, + FDocTranslate::FromNative<&LetDocTranslate>()); } ffi::Optional DivDocTranslate(DocTranslatorObj* d, ffi::AnyView input, @@ -428,78 +431,79 @@ ffi::Optional DivDocTranslate(DocTranslatorObj* d, ffi::AnyView input, } TVM_FFI_STATIC_INIT_BLOCK() { - ffi::reflection::TypeAttrDef().attr(tvm::script::printer::type_attr::kDocTranslate, - FDocTranslate::FromNative<&DivDocTranslate>()); + ffi::reflection::TypeAttrDef().attr( + tvm::script::printer::type_attr::kDocTranslate.name, + FDocTranslate::FromNative<&DivDocTranslate>()); ffi::reflection::TypeAttrDef().attr( - tvm::script::printer::type_attr::kDocTranslate, + tvm::script::printer::type_attr::kDocTranslate.name, FDocTranslate::FromNative< &BinaryOpDocTranslate>()); ffi::reflection::TypeAttrDef().attr( - tvm::script::printer::type_attr::kDocTranslate, + tvm::script::printer::type_attr::kDocTranslate.name, FDocTranslate::FromNative< &BinaryOpDocTranslate>()); ffi::reflection::TypeAttrDef().attr( - tvm::script::printer::type_attr::kDocTranslate, + tvm::script::printer::type_attr::kDocTranslate.name, FDocTranslate::FromNative< &BinaryOpDocTranslate>()); ffi::reflection::TypeAttrDef().attr( - tvm::script::printer::type_attr::kDocTranslate, + tvm::script::printer::type_attr::kDocTranslate.name, FDocTranslate::FromNative<&BinaryOpDocTranslate< prim::FloorDivNode, OperationDocNode::Kind::kFloorDiv, tvm::floordiv>>()); ffi::reflection::TypeAttrDef().attr( - tvm::script::printer::type_attr::kDocTranslate, + tvm::script::printer::type_attr::kDocTranslate.name, FDocTranslate::FromNative<&BinaryOpDocTranslate< prim::FloorModNode, OperationDocNode::Kind::kMod, tvm::floormod>>()); ffi::reflection::TypeAttrDef().attr( - tvm::script::printer::type_attr::kDocTranslate, + tvm::script::printer::type_attr::kDocTranslate.name, FDocTranslate::FromNative<&BinaryOpDocTranslate< prim::LShiftNode, OperationDocNode::Kind::kLShift, tvm::left_shift>>()); ffi::reflection::TypeAttrDef().attr( - tvm::script::printer::type_attr::kDocTranslate, + tvm::script::printer::type_attr::kDocTranslate.name, FDocTranslate::FromNative<&BinaryOpDocTranslate< prim::RShiftNode, OperationDocNode::Kind::kRShift, tvm::right_shift>>()); ffi::reflection::TypeAttrDef().attr( - tvm::script::printer::type_attr::kDocTranslate, + tvm::script::printer::type_attr::kDocTranslate.name, FDocTranslate::FromNative<&BinaryOpDocTranslate< prim::BitwiseAndNode, OperationDocNode::Kind::kBitAnd, tvm::bitwise_and>>()); ffi::reflection::TypeAttrDef().attr( - tvm::script::printer::type_attr::kDocTranslate, + tvm::script::printer::type_attr::kDocTranslate.name, FDocTranslate::FromNative<&BinaryOpDocTranslate< prim::BitwiseOrNode, OperationDocNode::Kind::kBitOr, tvm::bitwise_or>>()); ffi::reflection::TypeAttrDef().attr( - tvm::script::printer::type_attr::kDocTranslate, + tvm::script::printer::type_attr::kDocTranslate.name, FDocTranslate::FromNative<&BinaryOpDocTranslate< prim::BitwiseXorNode, OperationDocNode::Kind::kBitXor, tvm::bitwise_xor>>()); ffi::reflection::TypeAttrDef().attr( - tvm::script::printer::type_attr::kDocTranslate, + tvm::script::printer::type_attr::kDocTranslate.name, FDocTranslate::FromNative< &BinaryOpDocTranslate>()); ffi::reflection::TypeAttrDef().attr( - tvm::script::printer::type_attr::kDocTranslate, + tvm::script::printer::type_attr::kDocTranslate.name, FDocTranslate::FromNative< &BinaryOpDocTranslate>()); ffi::reflection::TypeAttrDef().attr( - tvm::script::printer::type_attr::kDocTranslate, + tvm::script::printer::type_attr::kDocTranslate.name, FDocTranslate::FromNative< &BinaryOpDocTranslate>()); ffi::reflection::TypeAttrDef().attr( - tvm::script::printer::type_attr::kDocTranslate, + tvm::script::printer::type_attr::kDocTranslate.name, FDocTranslate::FromNative< &BinaryOpDocTranslate>()); ffi::reflection::TypeAttrDef().attr( - tvm::script::printer::type_attr::kDocTranslate, + tvm::script::printer::type_attr::kDocTranslate.name, FDocTranslate::FromNative< &BinaryOpDocTranslate>()); ffi::reflection::TypeAttrDef().attr( - tvm::script::printer::type_attr::kDocTranslate, + tvm::script::printer::type_attr::kDocTranslate.name, FDocTranslate::FromNative< &BinaryOpDocTranslate>()); ffi::reflection::TypeAttrDef().attr( - tvm::script::printer::type_attr::kDocTranslate, + tvm::script::printer::type_attr::kDocTranslate.name, FDocTranslate::FromNative< &BinaryOpDocTranslate>()); ffi::reflection::TypeAttrDef().attr( - tvm::script::printer::type_attr::kDocTranslate, + tvm::script::printer::type_attr::kDocTranslate.name, FDocTranslate::FromNative< &BinaryOpDocTranslate>()); } @@ -512,8 +516,9 @@ ffi::Optional ModDocTranslate(DocTranslatorObj* d, ffi::AnyView input, } TVM_FFI_STATIC_INIT_BLOCK() { - ffi::reflection::TypeAttrDef().attr(tvm::script::printer::type_attr::kDocTranslate, - FDocTranslate::FromNative<&ModDocTranslate>()); + ffi::reflection::TypeAttrDef().attr( + tvm::script::printer::type_attr::kDocTranslate.name, + FDocTranslate::FromNative<&ModDocTranslate>()); } ffi::Optional MinDocTranslate(DocTranslatorObj* d, ffi::AnyView input, @@ -524,8 +529,9 @@ ffi::Optional MinDocTranslate(DocTranslatorObj* d, ffi::AnyView input, } TVM_FFI_STATIC_INIT_BLOCK() { - ffi::reflection::TypeAttrDef().attr(tvm::script::printer::type_attr::kDocTranslate, - FDocTranslate::FromNative<&MinDocTranslate>()); + ffi::reflection::TypeAttrDef().attr( + tvm::script::printer::type_attr::kDocTranslate.name, + FDocTranslate::FromNative<&MinDocTranslate>()); } ffi::Optional MaxDocTranslate(DocTranslatorObj* d, ffi::AnyView input, @@ -536,8 +542,9 @@ ffi::Optional MaxDocTranslate(DocTranslatorObj* d, ffi::AnyView input, } TVM_FFI_STATIC_INIT_BLOCK() { - ffi::reflection::TypeAttrDef().attr(tvm::script::printer::type_attr::kDocTranslate, - FDocTranslate::FromNative<&MaxDocTranslate>()); + ffi::reflection::TypeAttrDef().attr( + tvm::script::printer::type_attr::kDocTranslate.name, + FDocTranslate::FromNative<&MaxDocTranslate>()); } } // namespace diff --git a/tirx/src/script/printer/ir/prim_type.cc b/tirx/src/script/printer/ir/prim_type.cc index ad0cdcc069a9..d1f8f06bcf1a 100644 --- a/tirx/src/script/printer/ir/prim_type.cc +++ b/tirx/src/script/printer/ir/prim_type.cc @@ -44,7 +44,7 @@ ffi::Optional MemoryScopeDocTranslate(DocTranslatorObj*, ffi::AnyView i TVM_FFI_STATIC_INIT_BLOCK() { ffi::reflection::TypeAttrDef().attr( - tvm::script::printer::type_attr::kDocTranslate, + tvm::script::printer::type_attr::kDocTranslate.name, FDocTranslate::FromNative<&MemoryScopeDocTranslate>()); } @@ -87,8 +87,9 @@ ffi::Optional IntImmDocTranslate(DocTranslatorObj* d, ffi::AnyView inpu } TVM_FFI_STATIC_INIT_BLOCK() { - ffi::reflection::TypeAttrDef().attr(tvm::script::printer::type_attr::kDocTranslate, - FDocTranslate::FromNative<&IntImmDocTranslate>()); + ffi::reflection::TypeAttrDef().attr( + tvm::script::printer::type_attr::kDocTranslate.name, + FDocTranslate::FromNative<&IntImmDocTranslate>()); } ffi::Optional FloatImmDocTranslate(DocTranslatorObj* d, ffi::AnyView input, @@ -115,7 +116,7 @@ ffi::Optional FloatImmDocTranslate(DocTranslatorObj* d, ffi::AnyView in TVM_FFI_STATIC_INIT_BLOCK() { ffi::reflection::TypeAttrDef().attr( - tvm::script::printer::type_attr::kDocTranslate, + tvm::script::printer::type_attr::kDocTranslate.name, FDocTranslate::FromNative<&FloatImmDocTranslate>()); } @@ -135,7 +136,7 @@ ffi::Optional PrimTypeDocTranslate(DocTranslatorObj*, ffi::AnyView inpu TVM_FFI_STATIC_INIT_BLOCK() { ffi::reflection::TypeAttrDef().attr( - tvm::script::printer::type_attr::kDocTranslate, + tvm::script::printer::type_attr::kDocTranslate.name, FDocTranslate::FromNative<&PrimTypeDocTranslate>()); } @@ -154,7 +155,7 @@ ffi::Optional PtrTypeDocTranslate(DocTranslatorObj* d, ffi::AnyView inp TVM_FFI_STATIC_INIT_BLOCK() { ffi::reflection::TypeAttrDef().attr( - tvm::script::printer::type_attr::kDocTranslate, + tvm::script::printer::type_attr::kDocTranslate.name, FDocTranslate::FromNative<&PtrTypeDocTranslate>()); } @@ -167,7 +168,7 @@ ffi::Optional TensorRegionTypeDocTranslate(DocTranslatorObj*, ffi::AnyV TVM_FFI_STATIC_INIT_BLOCK() { ffi::reflection::TypeAttrDef().attr( - tvm::script::printer::type_attr::kDocTranslate, + tvm::script::printer::type_attr::kDocTranslate.name, FDocTranslate::FromNative<&TensorRegionTypeDocTranslate>()); } @@ -183,7 +184,7 @@ ffi::Optional TupleTypeDocTranslate(DocTranslatorObj* d, ffi::AnyView i TVM_FFI_STATIC_INIT_BLOCK() { ffi::reflection::TypeAttrDef().attr( - tvm::script::printer::type_attr::kDocTranslate, + tvm::script::printer::type_attr::kDocTranslate.name, FDocTranslate::FromNative<&TupleTypeDocTranslate>()); } diff --git a/tirx/src/script/printer/ir/stmt.cc b/tirx/src/script/printer/ir/stmt.cc index 76802059219b..572a194105e2 100644 --- a/tirx/src/script/printer/ir/stmt.cc +++ b/tirx/src/script/printer/ir/stmt.cc @@ -60,7 +60,7 @@ ffi::Optional BindDocTranslate(DocTranslatorObj* d, ffi::AnyView input, } TVM_FFI_STATIC_INIT_BLOCK() { - ffi::reflection::TypeAttrDef().attr(tvm::script::printer::type_attr::kDocTranslate, + ffi::reflection::TypeAttrDef().attr(tvm::script::printer::type_attr::kDocTranslate.name, FDocTranslate::FromNative<&BindDocTranslate>()); } @@ -75,8 +75,9 @@ ffi::Optional ReturnDocTranslate(DocTranslatorObj* d, ffi::AnyView inpu } TVM_FFI_STATIC_INIT_BLOCK() { - ffi::reflection::TypeAttrDef().attr(tvm::script::printer::type_attr::kDocTranslate, - FDocTranslate::FromNative<&ReturnDocTranslate>()); + ffi::reflection::TypeAttrDef().attr( + tvm::script::printer::type_attr::kDocTranslate.name, + FDocTranslate::FromNative<&ReturnDocTranslate>()); } ffi::Optional AssertStmtDocTranslate(DocTranslatorObj* d, ffi::AnyView input, @@ -98,7 +99,7 @@ ffi::Optional AssertStmtDocTranslate(DocTranslatorObj* d, ffi::AnyView TVM_FFI_STATIC_INIT_BLOCK() { ffi::reflection::TypeAttrDef().attr( - tvm::script::printer::type_attr::kDocTranslate, + tvm::script::printer::type_attr::kDocTranslate.name, FDocTranslate::FromNative<&AssertStmtDocTranslate>()); } @@ -114,8 +115,9 @@ ffi::Optional WhileDocTranslate(DocTranslatorObj* d, ffi::AnyView input } TVM_FFI_STATIC_INIT_BLOCK() { - ffi::reflection::TypeAttrDef().attr(tvm::script::printer::type_attr::kDocTranslate, - FDocTranslate::FromNative<&WhileDocTranslate>()); + ffi::reflection::TypeAttrDef().attr( + tvm::script::printer::type_attr::kDocTranslate.name, + FDocTranslate::FromNative<&WhileDocTranslate>()); } ffi::Optional BreakDocTranslate(DocTranslatorObj* d, ffi::AnyView input, @@ -129,8 +131,9 @@ ffi::Optional BreakDocTranslate(DocTranslatorObj* d, ffi::AnyView input } TVM_FFI_STATIC_INIT_BLOCK() { - ffi::reflection::TypeAttrDef().attr(tvm::script::printer::type_attr::kDocTranslate, - FDocTranslate::FromNative<&BreakDocTranslate>()); + ffi::reflection::TypeAttrDef().attr( + tvm::script::printer::type_attr::kDocTranslate.name, + FDocTranslate::FromNative<&BreakDocTranslate>()); } ffi::Optional ContinueDocTranslate(DocTranslatorObj* d, ffi::AnyView input, @@ -145,7 +148,7 @@ ffi::Optional ContinueDocTranslate(DocTranslatorObj* d, ffi::AnyView in TVM_FFI_STATIC_INIT_BLOCK() { ffi::reflection::TypeAttrDef().attr( - tvm::script::printer::type_attr::kDocTranslate, + tvm::script::printer::type_attr::kDocTranslate.name, FDocTranslate::FromNative<&ContinueDocTranslate>()); } @@ -164,7 +167,7 @@ ffi::Optional IfDocTranslate(DocTranslatorObj* d, ffi::AnyView input, } TVM_FFI_STATIC_INIT_BLOCK() { - ffi::reflection::TypeAttrDef().attr(tvm::script::printer::type_attr::kDocTranslate, + ffi::reflection::TypeAttrDef().attr(tvm::script::printer::type_attr::kDocTranslate.name, FDocTranslate::FromNative<&IfDocTranslate>()); } diff --git a/tirx/src/script/printer/printer.cc b/tirx/src/script/printer/printer.cc index 2d4f2d3eba74..58d8f52ad0df 100644 --- a/tirx/src/script/printer/printer.cc +++ b/tirx/src/script/printer/printer.cc @@ -133,7 +133,7 @@ void ApplyDisplayNames(const Doc& root, const PrinterConfig& config, if (auto original = origins.Get(definition)) { auto attrs = ReflectedField(ReflectedField(original.value(), "attrs"), "dict"); if (auto dict = attrs.as>()) { - if (auto symbol = dict->Get(tvm::attr::kGlobalSymbol)) { + if (auto symbol = dict->Get(tvm::attr::kGlobalSymbol.name)) { TVM_FFI_CHECK(symbol.value().cast() == name.value(), TypeError) << "printer function global_symbol must match its definition name"; } @@ -459,7 +459,8 @@ std::string RenderFallbackWithInvisiblePathInfo(const ffi::String& script, std::string Script(const ffi::ObjectRef& node, const ffi::Optional& cfg) { PrinterConfig config = cfg.value_or(PrinterConfig()); - static ffi::reflection::TypeAttrColumn translate(tvm::script::printer::type_attr::kDocTranslate); + static ffi::reflection::TypeAttrColumn translate( + tvm::script::printer::type_attr::kDocTranslate.name); // Builtin runtime roots keep their native repr; hooks still translate them within IR. if (!node.defined() || node->type_index() < ffi::TypeIndex::kTVMFFIDynObjectBegin || (translate[node->type_index()].type_index() == ffi::TypeIndex::kTVMFFINone && diff --git a/tirx/src/tirx/analysis/check_well_formed.cc b/tirx/src/tirx/analysis/check_well_formed.cc index 0efa2fad1d5a..2dce2f9fecdb 100644 --- a/tirx/src/tirx/analysis/check_well_formed.cc +++ b/tirx/src/tirx/analysis/check_well_formed.cc @@ -269,8 +269,8 @@ class WellFormednessVerifier : public StmtExprVisitor { std::vector globals; for (const auto& [global, function] : op->functions) globals.push_back(global); std::sort(globals.begin(), globals.end(), [&](const GlobalVar& a, const GlobalVar& b) { - bool a_public = op->functions[a]->GetAttr(tvm::attr::kGlobalSymbol).has_value(); - bool b_public = op->functions[b]->GetAttr(tvm::attr::kGlobalSymbol).has_value(); + bool a_public = op->functions[a]->GetAttr(tvm::attr::kGlobalSymbol).has_value(); + bool b_public = op->functions[b]->GetAttr(tvm::attr::kGlobalSymbol).has_value(); return a_public != b_public ? a_public > b_public : a->name_hint < b->name_hint; }); for (const auto& global : globals) { diff --git a/tirx/src/tirx/analysis/stmt_finding.cc b/tirx/src/tirx/analysis/stmt_finding.cc index 8e2461e1b034..44370456d4a8 100644 --- a/tirx/src/tirx/analysis/stmt_finding.cc +++ b/tirx/src/tirx/analysis/stmt_finding.cc @@ -26,7 +26,7 @@ namespace tirx { const FunctionNode* FindEntryFunc(const IRModule& mod, GlobalVar* result_g_var) { ffi::Optional result, last_gvar; - // Priority 1: Function marked as `tirx::attr::kIsEntryFunc` + // Priority 1: Function marked as `tirx::attr::kIsEntryFunc.name` int num_function = 0; const tirx::FunctionNode* main_func = nullptr; const tirx::FunctionNode* last_func = nullptr; @@ -36,7 +36,7 @@ const FunctionNode* FindEntryFunc(const IRModule& mod, GlobalVar* result_g_var) if (const auto* func = base_func.as()) { last_func = func; last_gvar = gv; - if (func->HasNonzeroAttr(tvm::tirx::attr::kIsEntryFunc)) { + if (func->HasNonzeroAttr(tvm::tirx::attr::kIsEntryFunc.name)) { if (result_g_var != nullptr) { *result_g_var = gv; } diff --git a/tirx/src/tirx/analysis/verify_memory.cc b/tirx/src/tirx/analysis/verify_memory.cc index 6ccf246df4dc..4544e19482aa 100644 --- a/tirx/src/tirx/analysis/verify_memory.cc +++ b/tirx/src/tirx/analysis/verify_memory.cc @@ -186,7 +186,7 @@ class MemoryAccessVerifier final : public StmtExprVisitor { /// Interface of VerifyMemory pass std::vector VerifyMemory_(const Function& func) { - auto target = func->GetAttr(tvm::attr::kTarget); + auto target = func->GetAttr(tvm::attr::kTarget); // Skip verification for functions without a target attribute, as they are // typically host-only helper functions that do not have device-memory constraints. if (!target.has_value()) return {}; @@ -195,7 +195,7 @@ std::vector VerifyMemory_(const Function& func) { << "' for primitive:" << std::endl << func; - if (func->GetAttr(tvm::attr::kCallingConv, CallingConv::kDefault).value() == + if (func->GetAttr(tvm::attr::kCallingConv, CallingConv::kDefault).value() == CallingConv::kDefault) { auto v = ffi::make_object(func, target.value()->GetTargetDeviceType()); v->Run(); diff --git a/tirx/src/tirx/ir/stmt.cc b/tirx/src/tirx/ir/stmt.cc index eac0cd59e455..d8f7e412eb31 100644 --- a/tirx/src/tirx/ir/stmt.cc +++ b/tirx/src/tirx/ir/stmt.cc @@ -28,8 +28,8 @@ #include #include #include +#include #include -#include #include #include @@ -105,8 +105,8 @@ void ValidateTensorEvaluate(Expr value) { TVM_FFI_STATIC_INIT_BLOCK() { namespace refl = ffi::reflection; refl::TypeAttrDef() - .def(tvm::type_attr::kTensorStoreValidate, ValidateTensorStore) - .def(tvm::type_attr::kEvaluateValidate, ValidateTensorEvaluate); + .def(tvm::type_attr::kTensorStoreValidate.name, ValidateTensorStore) + .def(tvm::type_attr::kEvaluateValidate.name, ValidateTensorEvaluate); } } // namespace tirx diff --git a/tirx/src/tirx/ir/tensor.cc b/tirx/src/tirx/ir/tensor.cc index 8864a345bb10..c1a8fe13f470 100644 --- a/tirx/src/tirx/ir/tensor.cc +++ b/tirx/src/tirx/ir/tensor.cc @@ -192,9 +192,9 @@ TVM_FFI_STATIC_INIT_BLOCK() { TVM_FFI_STATIC_INIT_BLOCK() { namespace refl = tvm::ffi::reflection; - refl::TypeAttrDef().def(tvm::type_attr::kSubscriptExprRealize, + refl::TypeAttrDef().def(tvm::type_attr::kSubscriptExprRealize.name, RealizeTensorSubscript); - refl::TypeAttrDef().def(tvm::type_attr::kSubscriptExprRealize, + refl::TypeAttrDef().def(tvm::type_attr::kSubscriptExprRealize.name, RealizeTensorRegionSubscript); } diff --git a/tirx/src/tirx/ir/tensor_load.cc b/tirx/src/tirx/ir/tensor_load.cc index 02ccb216292d..3129f59d1609 100644 --- a/tirx/src/tirx/ir/tensor_load.cc +++ b/tirx/src/tirx/ir/tensor_load.cc @@ -71,7 +71,8 @@ PrimType InferTensorLoadType(const TensorLoadNode* load) { TVM_FFI_STATIC_INIT_BLOCK() { namespace refl = tvm::ffi::reflection; refl::TypeAttrDef().attr( - type_attr::kTensorLoadInferType, FTensorLoadInferType::FromNative<&InferTensorLoadType>()); + type_attr::kTensorLoadInferType.name, + FTensorLoadInferType::FromNative<&InferTensorLoadType>()); } } // namespace tirx diff --git a/tirx/src/tirx/ir/transform.cc b/tirx/src/tirx/ir/transform.cc index bc613a2cf678..051d1595ee32 100644 --- a/tirx/src/tirx/ir/transform.cc +++ b/tirx/src/tirx/ir/transform.cc @@ -34,14 +34,14 @@ namespace transform { // Register build pipeline related options TVM_FFI_STATIC_INIT_BLOCK() { - ::tvm::transform::PassContext::RegisterConfigOption(tvm::tirx::attr::kNoAlias); + ::tvm::transform::PassContext::RegisterConfigOption(tvm::tirx::attr::kNoAlias.name); ::tvm::transform::PassContext::RegisterConfigOption("tirx.disable_assert"); ::tvm::transform::PassContext::RegisterConfigOption("tirx.disable_vectorize"); ::tvm::transform::PassContext::RegisterConfigOption("tirx.enable_tensor_level_predication"); ::tvm::transform::PassContext::RegisterConfigOption("tirx.disable_cse_tir"); ::tvm::transform::PassContext::RegisterConfigOption("tirx.enable_debug"); ::tvm::transform::PassContext::RegisterConfigOption("tirx.disable_storage_rewrite"); - ::tvm::transform::PassContext::RegisterConfigOption(tvm::tirx::attr::kIsEntryFunc); + ::tvm::transform::PassContext::RegisterConfigOption(tvm::tirx::attr::kIsEntryFunc.name); ::tvm::transform::PassContext::RegisterConfigOption>>( "tirx.add_lower_pass"); ::tvm::transform::PassContext::RegisterConfigOption("tirx.debug_keep_trivial_loop"); diff --git a/tirx/src/tirx/op/abi.cc b/tirx/src/tirx/op/abi.cc index 49ee57ced9c1..9f25bd8ca34d 100644 --- a/tirx/src/tirx/op/abi.cc +++ b/tirx/src/tirx/op/abi.cc @@ -26,7 +26,7 @@ #include #include #include -#include +#include #include namespace tvm { diff --git a/tirx/src/tirx/op/annotation.cc b/tirx/src/tirx/op/annotation.cc index 2069b1ffaae6..63d36d635080 100644 --- a/tirx/src/tirx/op/annotation.cc +++ b/tirx/src/tirx/op/annotation.cc @@ -26,7 +26,7 @@ #include #include #include -#include +#include namespace tvm { namespace tirx { diff --git a/tirx/src/tirx/op/cpu.cc b/tirx/src/tirx/op/cpu.cc index e01fba581f9c..2d1347601bfc 100644 --- a/tirx/src/tirx/op/cpu.cc +++ b/tirx/src/tirx/op/cpu.cc @@ -25,8 +25,8 @@ #include #include #include +#include #include -#include namespace tvm { namespace tirx { diff --git a/tirx/src/tirx/op/gpu.cc b/tirx/src/tirx/op/gpu.cc index d2ac0bb14f40..d943f39a867d 100644 --- a/tirx/src/tirx/op/gpu.cc +++ b/tirx/src/tirx/op/gpu.cc @@ -25,8 +25,8 @@ #include #include #include +#include #include -#include namespace tvm { namespace tirx { diff --git a/tirx/src/tirx/op/memory.cc b/tirx/src/tirx/op/memory.cc index ac1be17b8d48..a0373ff38fbb 100644 --- a/tirx/src/tirx/op/memory.cc +++ b/tirx/src/tirx/op/memory.cc @@ -24,8 +24,8 @@ #include #include #include +#include #include -#include #include namespace tvm { diff --git a/tirx/src/tirx/op/region.cc b/tirx/src/tirx/op/region.cc index 76e92f69de30..a642b79c3a35 100644 --- a/tirx/src/tirx/op/region.cc +++ b/tirx/src/tirx/op/region.cc @@ -25,8 +25,8 @@ #include #include #include +#include #include -#include namespace tvm { namespace tirx { diff --git a/tirx/src/tirx/op/vector.cc b/tirx/src/tirx/op/vector.cc index b2ea27e9596b..46c5f72fb559 100644 --- a/tirx/src/tirx/op/vector.cc +++ b/tirx/src/tirx/op/vector.cc @@ -25,8 +25,8 @@ #include #include #include +#include #include -#include namespace tvm { namespace tirx { diff --git a/tirx/src/tirx/script/ir_builder/frame.cc b/tirx/src/tirx/script/ir_builder/frame.cc index d1559cc26b56..46362e910d03 100644 --- a/tirx/src/tirx/script/ir_builder/frame.cc +++ b/tirx/src/tirx/script/ir_builder/frame.cc @@ -80,11 +80,11 @@ void FunctionFrameObj::ExitWithScope() { // Default attributes belong to the completed function. A declaration must // leave body-level func_attr free to supply those values on the same frame. if (!is_declaration && !is_private && name.has_value() && - !attrs.count(tvm::attr::kGlobalSymbol)) { - insert_attr(tvm::attr::kGlobalSymbol, name.value()); + !attrs.count(tvm::attr::kGlobalSymbol.name)) { + insert_attr(tvm::attr::kGlobalSymbol.name, name.value()); } if (!is_declaration && persistent) { - insert_attr(tvm::tirx::attr::kPersistentKernel, true); + insert_attr(tvm::tirx::attr::kPersistentKernel.name, true); } TVM_FFI_CHECK(!is_declaration || stmts.empty(), ValueError) << "A function declaration cannot contain body statements"; diff --git a/tirx/src/tirx/script/ir_builder/ir.cc b/tirx/src/tirx/script/ir_builder/ir.cc index 98dd69a7593a..6c6f6908690d 100644 --- a/tirx/src/tirx/script/ir_builder/ir.cc +++ b/tirx/src/tirx/script/ir_builder/ir.cc @@ -97,11 +97,11 @@ void FuncAttrs(ffi::Map new_attrs) { using namespace tvm::tirx; FunctionFrame frame = FindFunctionFrame("T.func_attr"); for (const auto& [key, value] : new_attrs) { - if (key == tvm::attr::kGlobalSymbol && frame->is_private) { + if (key == tvm::attr::kGlobalSymbol.name && frame->is_private) { TVM_FFI_THROW(InternalError) << "ValueError: " - << "A private function may not have the kGlobalSymbol (\"" << tvm::attr::kGlobalSymbol - << "\") attribute. " + << "A private function may not have the kGlobalSymbol (\"" + << tvm::attr::kGlobalSymbol.name << "\") attribute. " << "However, a private function specified the global symbol as " << value; } diff --git a/tirx/src/tirx/script/printer/block.cc b/tirx/src/tirx/script/printer/block.cc index 41ab52804a8d..e899d47350ee 100644 --- a/tirx/src/tirx/script/printer/block.cc +++ b/tirx/src/tirx/script/printer/block.cc @@ -44,7 +44,7 @@ ffi::Optional ThreadScopeDocTranslate(DocTranslatorObj*, ffi::AnyView i TVM_FFI_STATIC_INIT_BLOCK() { ffi::reflection::TypeAttrDef().attr( - tvm::script::printer::type_attr::kDocTranslate, + tvm::script::printer::type_attr::kDocTranslate.name, FDocTranslate::FromNative<&ThreadScopeDocTranslate>()); } diff --git a/tirx/src/tirx/script/printer/expr.cc b/tirx/src/tirx/script/printer/expr.cc index 51cc83f69c0b..8cfa996845b9 100644 --- a/tirx/src/tirx/script/printer/expr.cc +++ b/tirx/src/tirx/script/printer/expr.cc @@ -26,7 +26,7 @@ #include #include #include -#include +#include #include #include @@ -70,7 +70,7 @@ ffi::Optional IndexMapDocTranslate(DocTranslatorObj* d, ffi::AnyView in TVM_FFI_STATIC_INIT_BLOCK() { ffi::reflection::TypeAttrDef().attr( - tvm::script::printer::type_attr::kDocTranslate, + tvm::script::printer::type_attr::kDocTranslate.name, FDocTranslate::FromNative<&IndexMapDocTranslate>()); } diff --git a/tirx/src/tirx/script/printer/for_loop.cc b/tirx/src/tirx/script/printer/for_loop.cc index bfaf22161a60..609ff36b6e15 100644 --- a/tirx/src/tirx/script/printer/for_loop.cc +++ b/tirx/src/tirx/script/printer/for_loop.cc @@ -76,7 +76,7 @@ ffi::Optional ForDocTranslate(DocTranslatorObj* d, ffi::AnyView input, const auto& [key, value] = *loop->annotations.begin(); auto boolean = value.as(); auto integer = value.as(); - if (key == tvm::tirx::attr::kDisableUnroll && boolean.value_or(false)) { + if (key == tvm::tirx::attr::kDisableUnroll.name && boolean.value_or(false)) { keys.push_back("unroll"); values.push_back(LiteralDoc::Boolean(false, std::nullopt)); unroll_option = true; @@ -134,7 +134,7 @@ ffi::Optional ForDocTranslate(DocTranslatorObj* d, ffi::AnyView input, } TVM_FFI_STATIC_INIT_BLOCK() { - ffi::reflection::TypeAttrDef().attr(tvm::script::printer::type_attr::kDocTranslate, + ffi::reflection::TypeAttrDef().attr(tvm::script::printer::type_attr::kDocTranslate.name, FDocTranslate::FromNative<&ForDocTranslate>()); } diff --git a/tirx/src/tirx/script/printer/function.cc b/tirx/src/tirx/script/printer/function.cc index 8033519a1f55..faa8ac9b62a8 100644 --- a/tirx/src/tirx/script/printer/function.cc +++ b/tirx/src/tirx/script/printer/function.cc @@ -46,7 +46,7 @@ void PrintFunction(DocTranslatorObj* d, const tirx::FunctionNode* func, ExprDoc // A nested function must not inherit the enclosing module member's binding. d->WithExtraStateScope>("ir.function_name", std::nullopt, [&]() { auto completed = d->WithVarScope([&]() { - auto global_symbol = func->GetAttr(tvm::attr::kGlobalSymbol); + auto global_symbol = func->GetAttr(tvm::attr::kGlobalSymbol); ffi::String name = bound_name.value_or(global_symbol.value_or("main")); FunctionDoc doc(ffi::UnsafeInit{}); { @@ -84,11 +84,11 @@ void PrintFunction(DocTranslatorObj* d, const tirx::FunctionNode* func, ExprDoc for (const Var& var : func->params) signature_candidates.erase(var); ffi::Array decorator_keys; ffi::Array decorator_values; - if (!func->GetAttr(tvm::attr::kGlobalSymbol)) { + if (!func->GetAttr(tvm::attr::kGlobalSymbol)) { decorator_keys.push_back("private"); decorator_values.push_back(LiteralDoc::Boolean(true, std::nullopt)); } - if (func->attrs->dict.count(tvm::tirx::attr::kPersistentKernel)) { + if (func->attrs->dict.count(tvm::tirx::attr::kPersistentKernel.name)) { decorator_keys.push_back("persistent"); decorator_values.push_back(LiteralDoc::Boolean(true, std::nullopt)); } @@ -99,7 +99,8 @@ void PrintFunction(DocTranslatorObj* d, const tirx::FunctionNode* func, ExprDoc body = ToStmtDocArray(d->WithDocScope([&]() { d->Translate(func->body.value()); })); std::vector> attrs; for (const auto& [key, value] : func->attrs->dict) { - if (key != tvm::attr::kGlobalSymbol && key != tvm::tirx::attr::kPersistentKernel) + if (key != tvm::attr::kGlobalSymbol.name && + key != tvm::tirx::attr::kPersistentKernel.name) attrs.emplace_back(key, value); } if (!attrs.empty()) { @@ -119,11 +120,11 @@ void PrintFunction(DocTranslatorObj* d, const tirx::FunctionNode* func, ExprDoc if (bound_name && global_symbol && global_symbol.value() != bound_name.value()) { body.insert( body.begin(), - ExprStmtDoc( - NamespaceDoc("tirx") - ->Attr("func_attr") - ->Call({DictDoc({LiteralDoc::Str(tvm::attr::kGlobalSymbol, std::nullopt)}, - {LiteralDoc::Str(global_symbol.value(), std::nullopt)})}))); + ExprStmtDoc(NamespaceDoc("tirx") + ->Attr("func_attr") + ->Call({DictDoc( + {LiteralDoc::Str(tvm::attr::kGlobalSymbol.name, std::nullopt)}, + {LiteralDoc::Str(global_symbol.value(), std::nullopt)})}))); } ffi::Optional ret_type = std::nullopt; if (!func->ret_type.as().has_value() && !IsVoidType(func->ret_type)) { @@ -157,8 +158,7 @@ ffi::Optional TirxFunctionDocTranslate(DocTranslatorObj* d, ffi::AnyVie ffi::details::AnyUnsafe::RawObjectPtrFromAnyViewAfterCheck(input); TVM_FFI_CHECK(destination == nullptr, TypeError) << "printer statement-only node cannot fulfill a destination"; - auto construction_namespace = - func->GetAttr(tvm::attr::kScriptNamespace).value_or("tirx"); + auto construction_namespace = func->GetAttr(tvm::attr::kScriptNamespace).value_or("tirx"); PrintFunction(d, func, NamespaceDoc(construction_namespace)->Attr("function")); return std::nullopt; } @@ -166,7 +166,7 @@ ffi::Optional TirxFunctionDocTranslate(DocTranslatorObj* d, ffi::AnyVie TVM_FFI_STATIC_INIT_BLOCK() { RegisterNamespaceAlias("tirx.prefix", "T"); ffi::reflection::TypeAttrDef().attr( - tvm::script::printer::type_attr::kDocTranslate, + tvm::script::printer::type_attr::kDocTranslate.name, FDocTranslate::FromNative<&TirxFunctionDocTranslate>()); } diff --git a/tirx/src/tirx/script/printer/ir.cc b/tirx/src/tirx/script/printer/ir.cc index 95cd71e508a2..1c04aaa2ef77 100644 --- a/tirx/src/tirx/script/printer/ir.cc +++ b/tirx/src/tirx/script/printer/ir.cc @@ -41,8 +41,9 @@ ffi::Optional TargetDocTranslate(DocTranslatorObj* d, ffi::AnyView inpu } TVM_FFI_STATIC_INIT_BLOCK() { - ffi::reflection::TypeAttrDef().attr(tvm::script::printer::type_attr::kDocTranslate, - FDocTranslate::FromNative<&TargetDocTranslate>()); + ffi::reflection::TypeAttrDef().attr( + tvm::script::printer::type_attr::kDocTranslate.name, + FDocTranslate::FromNative<&TargetDocTranslate>()); } } // namespace diff --git a/tirx/src/tirx/script/printer/stmt.cc b/tirx/src/tirx/script/printer/stmt.cc index 6114009ebb5b..a64421d289ac 100644 --- a/tirx/src/tirx/script/printer/stmt.cc +++ b/tirx/src/tirx/script/printer/stmt.cc @@ -20,9 +20,9 @@ #include #include #include +#include #include #include -#include #include #include @@ -60,7 +60,7 @@ ffi::Optional EvaluateDocTranslate(DocTranslatorObj* d, ffi::AnyView in TVM_FFI_STATIC_INIT_BLOCK() { ffi::reflection::TypeAttrDef().attr( - tvm::script::printer::type_attr::kDocTranslate, + tvm::script::printer::type_attr::kDocTranslate.name, FDocTranslate::FromNative<&EvaluateDocTranslate>()); } @@ -78,7 +78,7 @@ ffi::Optional SeqStmtDocTranslate(DocTranslatorObj* d, ffi::AnyView inp TVM_FFI_STATIC_INIT_BLOCK() { ffi::reflection::TypeAttrDef().attr( - tvm::script::printer::type_attr::kDocTranslate, + tvm::script::printer::type_attr::kDocTranslate.name, FDocTranslate::FromNative<&SeqStmtDocTranslate>()); } @@ -121,7 +121,7 @@ ffi::Optional RegionStmtDocTranslate(DocTranslatorObj* d, ffi::AnyView TVM_FFI_STATIC_INIT_BLOCK() { ffi::reflection::TypeAttrDef().attr( - tvm::script::printer::type_attr::kDocTranslate, + tvm::script::printer::type_attr::kDocTranslate.name, FDocTranslate::FromNative<&RegionStmtDocTranslate>()); } diff --git a/tirx/src/tirx/script/printer/tensor.cc b/tirx/src/tirx/script/printer/tensor.cc index c550e7766fc3..2a7abcb91a23 100644 --- a/tirx/src/tirx/script/printer/tensor.cc +++ b/tirx/src/tirx/script/printer/tensor.cc @@ -104,7 +104,7 @@ ffi::Optional TensorTypeDocTranslate(DocTranslatorObj* d, ffi::AnyView TVM_FFI_STATIC_INIT_BLOCK() { ffi::reflection::TypeAttrDef().attr( - tvm::script::printer::type_attr::kDocTranslate, + tvm::script::printer::type_attr::kDocTranslate.name, FDocTranslate::FromNative<&TensorTypeDocTranslate>()); } @@ -127,7 +127,7 @@ ffi::Optional TensorStoreDocTranslate(DocTranslatorObj* d, ffi::AnyView TVM_FFI_STATIC_INIT_BLOCK() { ffi::reflection::TypeAttrDef().attr( - tvm::script::printer::type_attr::kDocTranslate, + tvm::script::printer::type_attr::kDocTranslate.name, FDocTranslate::FromNative<&TensorStoreDocTranslate>()); } @@ -143,7 +143,7 @@ ffi::Optional IterDocTranslate(DocTranslatorObj* d, ffi::AnyView input, TVM_FFI_STATIC_INIT_BLOCK() { ffi::reflection::TypeAttrDef().attr( - tvm::script::printer::type_attr::kDocTranslate, + tvm::script::printer::type_attr::kDocTranslate.name, FDocTranslate::FromNative<&IterDocTranslate>()); } @@ -179,7 +179,7 @@ ffi::Optional TileLayoutDocTranslate(DocTranslatorObj* d, ffi::AnyView TVM_FFI_STATIC_INIT_BLOCK() { ffi::reflection::TypeAttrDef().attr( - tvm::script::printer::type_attr::kDocTranslate, + tvm::script::printer::type_attr::kDocTranslate.name, FDocTranslate::FromNative<&TileLayoutDocTranslate>()); } @@ -205,7 +205,7 @@ ffi::Optional ComposeLayoutDocTranslate(DocTranslatorObj* d, ffi::AnyVi TVM_FFI_STATIC_INIT_BLOCK() { ffi::reflection::TypeAttrDef().attr( - tvm::script::printer::type_attr::kDocTranslate, + tvm::script::printer::type_attr::kDocTranslate.name, FDocTranslate::FromNative<&ComposeLayoutDocTranslate>()); } diff --git a/tirx/src/tirx/tile/context.cc b/tirx/src/tirx/tile/context.cc index c6531a438517..65b8d0149eaf 100644 --- a/tirx/src/tirx/tile/context.cc +++ b/tirx/src/tirx/tile/context.cc @@ -45,23 +45,23 @@ Value getOrSetDefault(ffi::Map& m, const Key& key, void DispatchContextObj::AddAllocTensor(TensorVar tensor, ffi::Array allocated_addr, ffi::Map annotations) { - auto tensors = getOrSetDefault(callbacks, callback::kPrivateAlloc, ffi::Array()); + auto tensors = getOrSetDefault(callbacks, callback::kPrivateAlloc.name, ffi::Array()); ffi::Array args; if (!allocated_addr.empty()) args.push_back(tvm::Tuple(allocated_addr)); tensors.push_back(Bind( tensor, Call(tensor.ty(), alloc_tensor_op(), args, DictAttrs(annotations), {tensor.ty()}))); - callbacks.Set(callback::kPrivateAlloc, tensors); + callbacks.Set(callback::kPrivateAlloc.name, tensors); } void DispatchContextObj::AddInitStmt(Stmt stmt, bool host) { - auto tag = host ? callback::kHostInitStmt : callback::kDeviceInitStmt; + auto tag = host ? callback::kHostInitStmt.name : callback::kDeviceInitStmt.name; auto stmts = getOrSetDefault(callbacks, tag, ffi::Array()); stmts.push_back(stmt); callbacks.Set(tag, stmts); } void DispatchContextObj::AddPostTensorDefStmt(TensorVar tensor, Stmt stmt) { - auto mapping = getOrSetDefault(callbacks, callback::kPostTensorDefStmt, + auto mapping = getOrSetDefault(callbacks, callback::kPostTensorDefStmt.name, ffi::Map>()); auto it = mapping.find(tensor); ffi::Array stmts; @@ -70,7 +70,7 @@ void DispatchContextObj::AddPostTensorDefStmt(TensorVar tensor, Stmt stmt) { } stmts.push_back(stmt); mapping.Set(tensor, stmts); - callbacks.Set(callback::kPostTensorDefStmt, mapping); + callbacks.Set(callback::kPostTensorDefStmt.name, mapping); } void DispatchContextObj::SharedStateSet(ffi::String key, ffi::ObjectRef value) { diff --git a/tirx/src/tirx/tile/lower.cc b/tirx/src/tirx/tile/lower.cc index bf4488980349..e82fe8164dc4 100644 --- a/tirx/src/tirx/tile/lower.cc +++ b/tirx/src/tirx/tile/lower.cc @@ -231,7 +231,7 @@ class TileDispatcher : public StmtExprMutator { Stmt ProcessDeviceEntry(const RegionStmtNode* entry_node) { Stmt body_to_visit = entry_node->body; Target saved_target = target_; - if (auto target = entry_node->attrs->dict.Get(tvm::attr::kTarget)) { + if (auto target = entry_node->attrs->dict.Get(tvm::attr::kTarget.name)) { target_ = target.value().cast(); } With target_scope(target_); @@ -503,19 +503,19 @@ class TileDispatcher : public StmtExprMutator { intra_map, scope_kind); Stmt res = tile::Dispatch(ffi::GetRef(op), sctx); // Implementation found, handle callbacks - if (auto bufs = sctx->callbacks.Get(tirx::tile::callback::kPrivateAlloc)) { + if (auto bufs = sctx->callbacks.Get(tirx::tile::callback::kPrivateAlloc.name)) { auto buf_list = bufs.value().as>().value(); alloc_tensors_.insert(alloc_tensors_.end(), buf_list.begin(), buf_list.end()); } - if (auto stmts = sctx->callbacks.Get(tirx::tile::callback::kDeviceInitStmt)) { + if (auto stmts = sctx->callbacks.Get(tirx::tile::callback::kDeviceInitStmt.name)) { auto stmt_list = stmts.value().as>().value(); device_init_stmts_.insert(device_init_stmts_.end(), stmt_list.begin(), stmt_list.end()); } - if (auto stmts = sctx->callbacks.Get(tirx::tile::callback::kHostInitStmt)) { + if (auto stmts = sctx->callbacks.Get(tirx::tile::callback::kHostInitStmt.name)) { auto stmt_list = stmts.value().as>().value(); host_init_stmts_.insert(host_init_stmts_.end(), stmt_list.begin(), stmt_list.end()); } - if (auto mapping = sctx->callbacks.Get(tirx::tile::callback::kPostTensorDefStmt)) { + if (auto mapping = sctx->callbacks.Get(tirx::tile::callback::kPostTensorDefStmt.name)) { auto map = mapping.value().as_or_throw>>(); for (const auto& [tensor, stmts] : map) { auto& vec = post_tensor_def_stmts_[tensor]; @@ -619,7 +619,7 @@ class TileDispatcher : public StmtExprMutator { namespace { Target ResolveTarget(const Function& f) { - auto target = f->GetAttr(tvm::attr::kTarget); + auto target = f->GetAttr(tvm::attr::kTarget); if (!target.has_value()) { target = Target::Current(false).value(); } diff --git a/tirx/src/tirx/transform/bind_target.cc b/tirx/src/tirx/transform/bind_target.cc index 89fffe63d65a..f810714ccbb7 100644 --- a/tirx/src/tirx/transform/bind_target.cc +++ b/tirx/src/tirx/transform/bind_target.cc @@ -75,7 +75,7 @@ class FunctionClassifierVisitor : public StmtExprVisitor { // Only analyze externally exposed functions as potential callers // since they represent the entry points where host/device calls originate for (const auto& [gvar, func] : mod->functions) { - bool is_externally_exposed = func->GetAttr(tvm::attr::kGlobalSymbol).has_value(); + bool is_externally_exposed = func->GetAttr(tvm::attr::kGlobalSymbol).has_value(); const auto* function = func.as(); if (is_externally_exposed && function != nullptr) { @@ -247,10 +247,9 @@ IRModule BindTarget(IRModule mod, const Target& target) { } auto function = ffi::GetRef(function_node); - bool is_externally_exposed = - function->GetAttr(tvm::attr::kGlobalSymbol).has_value(); + bool is_externally_exposed = function->GetAttr(tvm::attr::kGlobalSymbol).has_value(); - if (auto func_target = func->GetAttr(tvm::attr::kTarget)) { + if (auto func_target = func->GetAttr(tvm::attr::kTarget)) { // Rule 1: If the function has a target, and the target has a host, and the function does not // have a host, then add the host to the function target auto func_target_host = func_target.value()->GetHost(); @@ -258,22 +257,22 @@ IRModule BindTarget(IRModule mod, const Target& target) { if (target_host && !func_target_host && is_externally_exposed) { auto new_target = Target::WithHost(func_target.value(), target_host.value()); - new_mod->Update(gvar, WithAttr(std::move(function), tvm::attr::kTarget, new_target)); + new_mod->Update(gvar, WithAttr(std::move(function), tvm::attr::kTarget.name, new_target)); } continue; } - if (function->HasNonzeroAttr(tvm::tirx::attr::kIsHostFunc)) { + if (function->HasNonzeroAttr(tvm::tirx::attr::kIsHostFunc.name)) { // Rule 2: If the function is marked as host function, bind the host target to the function - function = WithAttr(std::move(function), tvm::attr::kTarget, + function = WithAttr(std::move(function), tvm::attr::kTarget.name, Target::WithHost(target_host, target_host)); - new_mod->Update(gvar, WithoutAttr(std::move(function), tvm::tirx::attr::kIsHostFunc)); + new_mod->Update(gvar, WithoutAttr(std::move(function), tvm::tirx::attr::kIsHostFunc.name)); continue; } if (is_externally_exposed) { // Rule 3: Externally exposed functions get the full target - new_mod->Update(gvar, WithAttr(std::move(function), tvm::attr::kTarget, target)); + new_mod->Update(gvar, WithAttr(std::move(function), tvm::attr::kTarget.name, target)); } else { const auto* gvar_node = gvar.get(); bool called_by_host = host_called_global_vars.count(gvar_node); @@ -283,11 +282,11 @@ IRModule BindTarget(IRModule mod, const Target& target) { // Rule 4.1: Called by both host and device // Bind device target to current function Function host_func = tirx::RenewDef(function); - new_mod->Update(gvar, - WithAttr(std::move(function), tvm::attr::kTarget, target_without_host)); + new_mod->Update( + gvar, WithAttr(std::move(function), tvm::attr::kTarget.name, target_without_host)); // Create duplicate with host target for host callers - host_func = WithAttr(std::move(host_func), tvm::attr::kTarget, target_host); + host_func = WithAttr(std::move(host_func), tvm::attr::kTarget.name, target_host); ffi::String host_func_name = gvar->name_hint + "_host"; GlobalVar host_gvar = GlobalVar(global_names->FreshName(host_func_name, false)); @@ -296,17 +295,17 @@ IRModule BindTarget(IRModule mod, const Target& target) { } else if (called_by_host) { // Rule 4.2: Called by host only - new_mod->Update(gvar, WithAttr(std::move(function), tvm::attr::kTarget, target_host)); + new_mod->Update(gvar, WithAttr(std::move(function), tvm::attr::kTarget.name, target_host)); } else if (called_by_device) { // Rule 4.3: Called by device only - new_mod->Update(gvar, - WithAttr(std::move(function), tvm::attr::kTarget, target_without_host)); + new_mod->Update( + gvar, WithAttr(std::move(function), tvm::attr::kTarget.name, target_without_host)); } else { // Rule 4.4: Not called by any context // NOTE: To keep the current behavior, we bind the target to the full target, but it needs // further check - new_mod->Update(gvar, - WithAttr(std::move(function), tvm::attr::kTarget, target_without_host)); + new_mod->Update( + gvar, WithAttr(std::move(function), tvm::attr::kTarget.name, target_without_host)); } } } @@ -321,8 +320,7 @@ IRModule BindTarget(IRModule mod, const Target& target) { continue; } - bool is_externally_exposed = - function->GetAttr(tvm::attr::kGlobalSymbol).has_value(); + bool is_externally_exposed = function->GetAttr(tvm::attr::kGlobalSymbol).has_value(); if (is_externally_exposed) { // Update calls in externally exposed functions to use host duplicates Function new_func = substitutor->Substitute(func.as_or_throw()); diff --git a/tirx/src/tirx/transform/function_utils.cc b/tirx/src/tirx/transform/function_utils.cc index 4471029f21c6..b46a74cb8437 100644 --- a/tirx/src/tirx/transform/function_utils.cc +++ b/tirx/src/tirx/transform/function_utils.cc @@ -37,10 +37,10 @@ transform::Pass AnnotateEntryFunc() { // If only a single function exists, that function must be the entry if (mod->functions.size() == 1) { auto [gvar, base_func] = *mod->functions.begin(); - if (!base_func->HasNonzeroAttr(tvm::tirx::attr::kIsEntryFunc)) { + if (!base_func->HasNonzeroAttr(tvm::tirx::attr::kIsEntryFunc.name)) { if (auto ptr = base_func.as()) { - mod->Update(gvar, - WithAttr(ffi::GetRef(ptr), tvm::tirx::attr::kIsEntryFunc, true)); + mod->Update( + gvar, WithAttr(ffi::GetRef(ptr), tvm::tirx::attr::kIsEntryFunc.name, true)); } } return mod; @@ -51,11 +51,11 @@ transform::Pass AnnotateEntryFunc() { bool has_external_non_functions = false; IRModule with_annotations; for (const auto& [gvar, base_func] : mod->functions) { - bool is_external = base_func->GetAttr(tvm::attr::kGlobalSymbol).has_value(); + bool is_external = base_func->GetAttr(tvm::attr::kGlobalSymbol).has_value(); if (is_external) { if (auto ptr = base_func.as()) { with_annotations->Add( - gvar, WithAttr(ffi::GetRef(ptr), tvm::tirx::attr::kIsEntryFunc, true)); + gvar, WithAttr(ffi::GetRef(ptr), tvm::tirx::attr::kIsEntryFunc.name, true)); } else { has_external_non_functions = true; } diff --git a/tirx/src/tirx/transform/lower_intrin.cc b/tirx/src/tirx/transform/lower_intrin.cc index c1f3a5f65896..73f4dc240654 100644 --- a/tirx/src/tirx/transform/lower_intrin.cc +++ b/tirx/src/tirx/transform/lower_intrin.cc @@ -408,7 +408,7 @@ namespace transform { Pass LowerIntrin() { auto pass_func = [](Function f, IRModule m, PassContext ctx) { auto* n = f.CopyOnWrite(); - auto target = f->GetAttr(tvm::attr::kTarget); + auto target = f->GetAttr(tvm::attr::kTarget); TVM_FFI_ICHECK(target.has_value()) << "LowerIntrin: Require the target attribute"; sym::Analyzer analyzer; bool enable_fast_math = ctx->GetConfig("tirx.enable_fast_math").value_or(false); diff --git a/tirx/src/tirx/transform/lower_thread_allreduce.cc b/tirx/src/tirx/transform/lower_thread_allreduce.cc index 825fb5c81623..2f0e8126dca9 100644 --- a/tirx/src/tirx/transform/lower_thread_allreduce.cc +++ b/tirx/src/tirx/transform/lower_thread_allreduce.cc @@ -128,7 +128,7 @@ class ThreadAllreduceBuilder final : public StmtExprMutator { DictAttrs annotations = call->attrs.value().as_or_throw(); static const auto shared = runtime::MemoryScope::Get("shared"); if (replacement.scope() == shared) { - annotations.CopyOnWrite()->dict.Set(tvm::tirx::attr::kVolatile, true); + annotations.CopyOnWrite()->dict.Set(tvm::tirx::attr::kVolatile.name, true); } ffi::Array args = call->args; @@ -988,7 +988,7 @@ class DeferredRemapper : public StmtExprMutator { DictAttrs annotations = call->attrs.value().as_or_throw(); static const auto shared = runtime::MemoryScope::Get("shared"); if (replacement.scope() == shared) { - annotations.CopyOnWrite()->dict.Set(tvm::tirx::attr::kVolatile, true); + annotations.CopyOnWrite()->dict.Set(tvm::tirx::attr::kVolatile.name, true); } ffi::Array args = call->args; @@ -1036,7 +1036,7 @@ class DeferredRemapper : public StmtExprMutator { Function LowerThreadAllreduce(Function f) { auto* n = f.CopyOnWrite(); - auto target = f->GetAttr(tvm::attr::kTarget); + auto target = f->GetAttr(tvm::attr::kTarget); TVM_FFI_ICHECK(target.has_value()) << "LowerThreadAllreduce: Require the target attribute"; const TargetObj* target_node = target.as(); auto thread_all_reduce = ffi::make_object(target_node, f->params); diff --git a/tirx/src/tirx/transform/lower_tirx_cleanup.cc b/tirx/src/tirx/transform/lower_tirx_cleanup.cc index 2fd560423db8..fa14e940b173 100644 --- a/tirx/src/tirx/transform/lower_tirx_cleanup.cc +++ b/tirx/src/tirx/transform/lower_tirx_cleanup.cc @@ -355,7 +355,7 @@ class LayoutApplier : public StmtExprMutatorWithAnalyzer { namespace { Target ResolveTarget(const Function& f) { - auto target = f->GetAttr(tvm::attr::kTarget); + auto target = f->GetAttr(tvm::attr::kTarget); if (!target.has_value()) { target = Target::Current(false).value(); } diff --git a/tirx/src/tirx/transform/lower_tvm_builtin.cc b/tirx/src/tirx/transform/lower_tvm_builtin.cc index a2e80dab0edb..0993422bbb63 100644 --- a/tirx/src/tirx/transform/lower_tvm_builtin.cc +++ b/tirx/src/tirx/transform/lower_tvm_builtin.cc @@ -64,7 +64,7 @@ class BuiltinLower : public StmtExprMutator { if (!func->body.has_value()) return func; ffi::Optional device_type = std::nullopt; bool preserve_ffi_kernel = false; - if (auto target = func->GetAttr(tvm::attr::kTarget)) { + if (auto target = func->GetAttr(tvm::attr::kTarget)) { device_type = IntImm::Int32(target.value()->kind->default_device_type); auto host = target.value()->GetHost().value_or(target.value()); preserve_ffi_kernel = host->kind->name == "cuda_host"; @@ -805,9 +805,9 @@ namespace transform { Pass LowerTVMBuiltin() { auto pass_func = [](Function func, IRModule m, PassContext ctx) { - bool is_host = func->HasNonzeroAttr(tvm::tirx::attr::kIsHostFunc); + bool is_host = func->HasNonzeroAttr(tvm::tirx::attr::kIsHostFunc.name); if (!is_host) { - if (auto target = func->GetAttr(tvm::attr::kTarget)) { + if (auto target = func->GetAttr(tvm::attr::kTarget)) { is_host = target.value()->HasKey("cpu"); } } diff --git a/tirx/src/tirx/transform/make_packed_api.cc b/tirx/src/tirx/transform/make_packed_api.cc index 3fbfd66e586f..038a4d8ec406 100644 --- a/tirx/src/tirx/transform/make_packed_api.cc +++ b/tirx/src/tirx/transform/make_packed_api.cc @@ -205,14 +205,14 @@ class SubroutineCallRewriter : public StmtExprMutator { ffi::Optional RequiresPackedAPI(const Function& func) { // A function with an explicit calling convention has already been // lowered, and should not be modified. - if (auto opt = func->GetAttr(tvm::attr::kCallingConv)) { + if (auto opt = func->GetAttr(tvm::attr::kCallingConv)) { if (opt.value() != CallingConv::kDefault) { return std::nullopt; } } // Internal function calls do not need the ffi::Function API - auto global_symbol = func->GetAttr(tvm::attr::kGlobalSymbol); + auto global_symbol = func->GetAttr(tvm::attr::kGlobalSymbol); if (!global_symbol.has_value()) { return std::nullopt; } @@ -229,10 +229,10 @@ Function MakePackedAPI(Function func) { std::string name_hint = global_symbol.value(); Target target = [&]() { - auto opt = func->GetAttr(tvm::attr::kTarget); + auto opt = func->GetAttr(tvm::attr::kTarget); TVM_FFI_ICHECK(opt) - << "MakePackedAPI required the function to be annotated with tvm::attr::kTarget (" - << tvm::attr::kTarget << "), but the function only has attributes " << func->attrs; + << "MakePackedAPI required the function to be annotated with tvm::attr::kTarget.name (" + << tvm::attr::kTarget.name << "), but the function only has attributes " << func->attrs; return opt.value(); }(); int target_device_type = target->GetTargetDeviceType(); @@ -266,9 +266,9 @@ Function MakePackedAPI(Function func) { ffi::Array args{v_self_handle, v_packed_args, v_num_packed_args, v_result}; // reset global symbol to attach prefix - func = WithAttrs(std::move(func), {{tvm::attr::kCallingConv, CallingConv::kCPackedFunc}, - {tvm::attr::kTarget, target_host}, - {tvm::attr::kGlobalSymbol, + func = WithAttrs(std::move(func), {{tvm::attr::kCallingConv.name, CallingConv::kCPackedFunc}, + {tvm::attr::kTarget.name, target_host}, + {tvm::attr::kGlobalSymbol.name, ffi::symbol::tvm_ffi_symbol_prefix + global_symbol.value()}}); Stmt body = ffi::make_object(v_result) diff --git a/tirx/src/tirx/transform/split_host_device.cc b/tirx/src/tirx/transform/split_host_device.cc index 2842f695e4aa..bb6875a12dde 100644 --- a/tirx/src/tirx/transform/split_host_device.cc +++ b/tirx/src/tirx/transform/split_host_device.cc @@ -39,9 +39,9 @@ #include #include #include +#include #include #include -#include #include #include @@ -68,14 +68,15 @@ class DeviceRegionAnnotater : public StmtExprMutator { UnchangedOr Mutate_(const RegionStmtNode* op, InplaceMode inplace_mode) final { static const Op device_scope = Op::Get("tirx.device_scope"); if (op->op.same_as(device_scope)) { - if (op->attrs->dict.count(tvm::attr::kTarget)) return ffi::Unchanged(); + if (op->attrs->dict.count(tvm::attr::kTarget.name)) return ffi::Unchanged(); auto attrs = op->attrs->dict; - attrs.Set(tvm::attr::kTarget, device_target_); + attrs.Set(tvm::attr::kTarget.name, device_target_); return RegionStmt(op->op, op->args, op->body_params, DictAttrs(attrs), op->body, op->result_vars, op->loc); } if (op->op.same_as(tirx::launch_thread_op())) { - return RegionStmt(device_scope, {}, {}, DictAttrs({{tvm::attr::kTarget, device_target_}}), + return RegionStmt(device_scope, {}, {}, + DictAttrs({{tvm::attr::kTarget.name, device_target_}}), ffi::GetRef(op)); } return StmtExprMutator::Mutate_(op, inplace_mode); @@ -86,7 +87,7 @@ class DeviceRegionAnnotater : public StmtExprMutator { }; Function AnnotateDeviceRegionsForSplit(Function func) { - auto opt_target = func->GetAttr(tvm::attr::kTarget); + auto opt_target = func->GetAttr(tvm::attr::kTarget); TVM_FFI_ICHECK(opt_target) << "SplitHostDevice: Require the target attribute"; Target target = opt_target.value(); @@ -124,7 +125,7 @@ class HostDeviceSplitter : public StmtExprMutator { UnchangedOr Mutate_(const RegionStmtNode* op, InplaceMode inplace_mode) final { static const Op device_scope = Op::Get("tirx.device_scope"); if (op->op.same_as(device_scope)) { - auto target = op->attrs->dict.Get(tvm::attr::kTarget); + auto target = op->attrs->dict.Get(tvm::attr::kTarget.name); if (!target) return Mutate(op->body, inplace_mode).ValueOrUnchanged(op->body); Target device_target = target.value().as_or_throw(); return SplitDeviceFunc(op->body, device_target.WithoutHost(), op); @@ -233,30 +234,30 @@ class HostDeviceSplitter : public StmtExprMutator { std::move(body)}); } Function device_func(kernel_params, SeqStmt(body), kernel_ret_type); - device_func = WithAttrs(std::move(device_func), {{tvm::attr::kTarget, device_target}, - {tvm::tirx::attr::kNoAlias, true}, - {tvm::tirx::attr::kIsGlobalFunc, true}}); + device_func = WithAttrs(std::move(device_func), {{tvm::attr::kTarget.name, device_target}, + {tvm::tirx::attr::kNoAlias.name, true}, + {tvm::tirx::attr::kIsGlobalFunc.name, true}}); if (auto config = region->attrs->dict.Get("backend_config")) { device_func = WithAttr(std::move(device_func), "backend_config", config.value()); } else if (auto config = cur_func_->GetAttr("backend_config")) { device_func = WithAttr(std::move(device_func), "backend_config", config.value()); } - if (auto script_namespace = cur_func_->GetAttr(tvm::attr::kScriptNamespace)) { - device_func = - WithAttr(std::move(device_func), tvm::attr::kScriptNamespace, script_namespace.value()); + if (auto script_namespace = cur_func_->GetAttr(tvm::attr::kScriptNamespace)) { + device_func = WithAttr(std::move(device_func), tvm::attr::kScriptNamespace.name, + script_namespace.value()); } - bool is_stir = cur_func_->attrs->dict.count(tvm::attr::kSTir); + bool is_stir = cur_func_->attrs->dict.count(tvm::attr::kSTir.name); if (is_stir) { - device_func = WithAttr(std::move(device_func), tvm::attr::kSTir, true); + device_func = WithAttr(std::move(device_func), tvm::attr::kSTir.name, true); } - if (auto launch_params = - cur_func_->GetAttr>(tvm::tirx::attr::kKernelLaunchParams)) { - device_func = WithAttr(std::move(device_func), tvm::tirx::attr::kKernelLaunchParams, + if (auto launch_params = cur_func_->GetAttr>( + tvm::tirx::attr::kKernelLaunchParams.name)) { + device_func = WithAttr(std::move(device_func), tvm::tirx::attr::kKernelLaunchParams.name, launch_params.value()); } - auto num_inputs = cur_func_->GetAttr(tvm::attr::kNumInputs); + auto num_inputs = cur_func_->GetAttr(tvm::attr::kNumInputs); if (num_inputs.has_value()) { - device_func = WithAttr(std::move(device_func), tvm::attr::kNumInputs, num_inputs); + device_func = WithAttr(std::move(device_func), tvm::attr::kNumInputs.name, num_inputs); } GlobalVar kernel_symbol_global = var_supply_(); if (region->attrs->dict.count(tvm::backend::cuda::attr::kLaunchFields.name)) { @@ -343,16 +344,15 @@ class DeviceInfoCollector : public StmtExprVisitor { static KernelInfo Collect(const GlobalVar& gvar, const Function& func, bool allow_placeholder = false) { auto collector = ffi::make_object( - func->GetAttr(tvm::attr::kTarget).value().WithoutHost()); + func->GetAttr(tvm::attr::kTarget).value().WithoutHost()); collector->info_.params = func->params; if (func->GetAttr>(tvm::backend::cuda::attr::kLaunchFields)) { collector->info_.global_symbol = - func->GetAttr(tvm::attr::kGlobalSymbol).value_or(gvar->name_hint); + func->GetAttr(tvm::attr::kGlobalSymbol).value_or(gvar->name_hint); return collector->info_; } - if (auto requested = - func->GetAttr>(tvm::tirx::attr::kKernelLaunchParams)) { + if (auto requested = func->GetAttr(tvm::tirx::attr::kKernelLaunchParams)) { for (const ffi::String& tag : requested.value()) { if (tag == tvm::runtime::launch_param::kUseProgramaticDependentLaunch) { collector->use_programmatic_dependent_launch_ = true; @@ -386,7 +386,7 @@ class DeviceInfoCollector : public StmtExprVisitor { } collector->info_.global_symbol = - func->GetAttr(tvm::attr::kGlobalSymbol).value_or(gvar->name_hint); + func->GetAttr(tvm::attr::kGlobalSymbol).value_or(gvar->name_hint); for (const ffi::String& param : collector->info_.launch_params) { if (param == tvm::runtime::launch_param::kUseProgramaticDependentLaunch || @@ -628,12 +628,12 @@ Stmt MakeCudaKernelLaunch(const GlobalVar& symbol, Function* func, ffi::Arrayname_hint}, - {tvm::backend::cuda::attr::kLaunchFields.name, fields}, - {tvm::backend::cuda::attr::kKernelAttrs.name, kernel_attrs}, - {tvm::backend::cuda::attr::kLaunchDimensions.name, dimensions}}); + *func = WithAttrs(std::move(*func), + {{tvm::attr::kCallingConv.name, tvm::CallingConv::kDeviceKernelLaunch}, + {tvm::attr::kGlobalSymbol.name, symbol->name_hint}, + {tvm::backend::cuda::attr::kLaunchFields.name, fields}, + {tvm::backend::cuda::attr::kKernelAttrs.name, kernel_attrs}, + {tvm::backend::cuda::attr::kLaunchDimensions.name, dimensions}}); auto attrs = ffi::make_object(); attrs->launch_fields = fields; attrs->kernel_attrs = kernel_attrs; @@ -666,7 +666,7 @@ class DeviceKernelMutator : public StmtExprMutator { // helper (e.g. a same-target subroutine that SplitHostDevice // emitted on the host side) should compare against the host // target, not the device target stripped by WithoutHost(). - auto full_target = func->GetAttr(tvm::attr::kTarget).value(); + auto full_target = func->GetAttr(tvm::attr::kTarget).value(); current_target_ = full_target.WithoutHost(); if (full_target->GetHost().has_value()) { current_caller_host_target_ = full_target->GetHost().value(); @@ -696,7 +696,7 @@ class DeviceKernelMutator : public StmtExprMutator { << "This case is not yet supported."; if (is_kernel_launch || is_call_extern) { - func = WithAttr(std::move(func), tvm::tirx::attr::kIsGlobalFunc, true); + func = WithAttr(std::move(func), tvm::tirx::attr::kIsGlobalFunc.name, true); } if (is_kernel_launch) { @@ -707,18 +707,18 @@ class DeviceKernelMutator : public StmtExprMutator { { auto write_ptr = func.CopyOnWrite(); write_ptr->ret_type = VoidType(); - Target target = func->GetAttr(tvm::attr::kTarget).value(); + Target target = func->GetAttr(tvm::attr::kTarget).value(); bool preserve_early_returns = target->kind->name == "cuda"; write_ptr->body = ReturnRemover::Apply(write_ptr->body.value(), !preserve_early_returns); } func = WithAttrs(std::move(func), - {{tvm::attr::kCallingConv, tvm::CallingConv::kDeviceKernelLaunch}, - {tvm::tirx::attr::kKernelLaunchParams, info.launch_params}, - {tvm::attr::kGlobalSymbol, info.global_symbol}}); + {{tvm::attr::kCallingConv.name, tvm::CallingConv::kDeviceKernelLaunch}, + {tvm::tirx::attr::kKernelLaunchParams.name, info.launch_params}, + {tvm::attr::kGlobalSymbol.name, info.global_symbol}}); - } else if (is_call_extern && !func->GetAttr(tvm::attr::kGlobalSymbol)) { - func = WithAttr(func, tvm::attr::kGlobalSymbol, gvar->name_hint); + } else if (is_call_extern && !func->GetAttr(tvm::attr::kGlobalSymbol)) { + func = WithAttr(func, tvm::attr::kGlobalSymbol.name, gvar->name_hint); } return func; @@ -806,7 +806,7 @@ class DeviceKernelMutator : public StmtExprMutator { TVM_FFI_ICHECK(dev_info.launch_params.defined()) << "CallNode attempted kernel launch to " << gvar->name_hint << " on target " << dev_info.target << ", but subroutine " << gvar->name_hint - << " did not have the tvm::tirx::attr::kKernelLaunchParams attribute " + << " did not have the tvm::tirx::attr::kKernelLaunchParams.name attribute " << "required for cross-target kernel launch"; // Collected kernel information may be in terms of the callee's @@ -929,7 +929,7 @@ Pass SplitHostDevice() { Function func = opt.value(); func = AnnotateDeviceRegionsForSplit(std::move(func)); - auto global_symbol = func->GetAttr(tvm::attr::kGlobalSymbol); + auto global_symbol = func->GetAttr(tvm::attr::kGlobalSymbol); auto name_prefix = global_symbol.value_or(gvar->name_hint); auto kernel_name = name_prefix + "_kernel"; auto var_supply = [&global_names, &kernel_name]() -> GlobalVar { diff --git a/tirx/src/tirx/transform/storage_rewrite.cc b/tirx/src/tirx/transform/storage_rewrite.cc index 36c996143e17..6e0a2dc17c47 100644 --- a/tirx/src/tirx/transform/storage_rewrite.cc +++ b/tirx/src/tirx/transform/storage_rewrite.cc @@ -432,7 +432,7 @@ class InplaceOpVerifier : public StmtExprVisitor { const CallNode* tensor_call) { // reject inplace for volatile buffers if ((tensor_call->attrs.has_value() ? tensor_call->attrs.value().as() : nullptr) - ->dict.count(tvm::tirx::attr::kVolatile)) { + ->dict.count(tvm::tirx::attr::kVolatile.name)) { result_ = false; return std::nullopt; } @@ -832,7 +832,7 @@ class StoragePlanRewriter : public StmtExprMutator { // Rebuilt allocations only preserve volatile metadata. Other owners may attach // constraints that cannot be transferred to merged or resized storage. for (const auto& [key, value] : annotations) { - if (key != tvm::tirx::attr::kVolatile) return true; + if (key != tvm::tirx::attr::kVolatile.name) return true; } return false; } @@ -976,7 +976,7 @@ class StoragePlanRewriter : public StmtExprMutator { TensorVar buf = RemapTensor(e->allocs[0]->var.as_or_throw(), e->alloc_var); ffi::Map annotations; if (e->is_volatile) { - annotations.Set(tvm::tirx::attr::kVolatile, true); + annotations.Set(tvm::tirx::attr::kVolatile.name, true); } e->alloc_nest.push_back(Bind(buf, Call(buf.ty(), tirx::alloc_tensor_op(), {}, DictAttrs(annotations), {buf.ty()}))); @@ -1022,7 +1022,7 @@ class StoragePlanRewriter : public StmtExprMutator { TensorVar buf = RemapTensor(e->allocs[0]->var.as_or_throw(), e->alloc_var); ffi::Map annotations; if (e->is_volatile) { - annotations.Set(tvm::tirx::attr::kVolatile, true); + annotations.Set(tvm::tirx::attr::kVolatile.name, true); } e->alloc_nest.push_back(Bind(buf, Call(buf.ty(), tirx::alloc_tensor_op(), {}, DictAttrs(annotations), {buf.ty()}))); @@ -1077,7 +1077,7 @@ class StoragePlanRewriter : public StmtExprMutator { e->alloc_var = buf; ffi::Map annotations; if (e->is_volatile) { - annotations.Set(tvm::tirx::attr::kVolatile, true); + annotations.Set(tvm::tirx::attr::kVolatile.name, true); } e->alloc_nest.push_back(Bind(buf, Call(buf.ty(), tirx::alloc_tensor_op(), {}, DictAttrs(annotations), {buf.ty()}))); @@ -1128,7 +1128,7 @@ class StoragePlanRewriter : public StmtExprMutator { } ffi::Map annotations; if (any_volatile) { - annotations.Set(tvm::tirx::attr::kVolatile, true); + annotations.Set(tvm::tirx::attr::kVolatile.name, true); } e->alloc_nest.push_back( Bind(buf, Call(buf.ty(), tirx::alloc_tensor_op(), {}, DictAttrs(annotations), {buf.ty()}))); @@ -1253,7 +1253,7 @@ class StoragePlanRewriter : public StmtExprMutator { } dst_entry->allocs.emplace_back(alloc); DictAttrs annotations = call->attrs.value().as_or_throw(); - if (annotations->dict.count(tvm::tirx::attr::kVolatile)) { + if (annotations->dict.count(tvm::tirx::attr::kVolatile.name)) { dst_entry->is_volatile = true; } alloc_map_[var] = dst_entry; @@ -2291,7 +2291,7 @@ Pass StorageRewrite() { enable_reuse = false; } - ffi::Optional target = f->GetAttr(tvm::attr::kTarget); + ffi::Optional target = f->GetAttr(tvm::attr::kTarget); if (target.has_value() && (target.value()->kind->name == "vulkan" || target.value()->kind->name == "webgpu")) { // Require exactly same-dtype matching in smem reuse for Vulkan and WebGPU diff --git a/tirx/src/tirx/transform/unroll_loop.cc b/tirx/src/tirx/transform/unroll_loop.cc index f6bcb0248781..f1d7d89456e2 100644 --- a/tirx/src/tirx/transform/unroll_loop.cc +++ b/tirx/src/tirx/transform/unroll_loop.cc @@ -116,11 +116,11 @@ class LoopUnroller : public StmtExprMutator { return unroll_policy_.WithNewScope([&]() { auto& policy = unroll_policy_.Current(); policy = parent_policy; - if (auto value = op->annotations.Get(tvm::tirx::attr::kAutoUnrollMaxStep); + if (auto value = op->annotations.Get(tvm::tirx::attr::kAutoUnrollMaxStep.name); value.has_value() && value.value() != nullptr) { policy.auto_unroll_max_step = value.value().cast()->value.as().value(); } - if (auto value = op->annotations.Get(tvm::tirx::attr::kUnrollExplicit); + if (auto value = op->annotations.Get(tvm::tirx::attr::kUnrollExplicit.name); value.has_value() && value.value() != nullptr) { policy.unroll_explicit = static_cast(value.value().cast()->value); } @@ -137,8 +137,8 @@ class LoopUnroller : public StmtExprMutator { inplace_mode = InplaceMode::kDisallow; } } - if (op->annotations.count(tvm::tirx::attr::kAutoUnrollMaxStep) || - op->annotations.count(tvm::tirx::attr::kUnrollExplicit)) { + if (op->annotations.count(tvm::tirx::attr::kAutoUnrollMaxStep.name) || + op->annotations.count(tvm::tirx::attr::kUnrollExplicit.name)) { ForNode* node; if (inplace_mode == InplaceMode::kAllow) { node = const_cast(op); @@ -147,8 +147,8 @@ class LoopUnroller : public StmtExprMutator { node = copy.get(); result = For(std::move(copy)); } - node->annotations.erase(tvm::tirx::attr::kAutoUnrollMaxStep); - node->annotations.erase(tvm::tirx::attr::kUnrollExplicit); + node->annotations.erase(tvm::tirx::attr::kAutoUnrollMaxStep.name); + node->annotations.erase(tvm::tirx::attr::kUnrollExplicit.name); op = node; } int value = GetExtent(op); diff --git a/tirx/src/tirx/transform/unsupported_dtype_legalize.cc b/tirx/src/tirx/transform/unsupported_dtype_legalize.cc index 26b3c8987bb6..17f54bc27f18 100644 --- a/tirx/src/tirx/transform/unsupported_dtype_legalize.cc +++ b/tirx/src/tirx/transform/unsupported_dtype_legalize.cc @@ -755,7 +755,7 @@ bool CheckDataTypeSupport(const Target& target, const std::string& support_func_ Pass BF16ComputeLegalize() { auto pass_func = [](Function f, IRModule m, PassContext ctx) { - auto opt_target = f->GetAttr(tvm::attr::kTarget); + auto opt_target = f->GetAttr(tvm::attr::kTarget); if (opt_target.has_value() && CheckDataTypeSupport(opt_target.value(), "tvm.support.nvcc.supports_bf16")) { return f; @@ -772,7 +772,7 @@ TVM_FFI_STATIC_INIT_BLOCK() { Pass BF16StorageLegalize() { auto pass_func = [](Function f, IRModule m, PassContext ctx) { - auto opt_target = f->GetAttr(tvm::attr::kTarget); + auto opt_target = f->GetAttr(tvm::attr::kTarget); if (opt_target.has_value() && CheckDataTypeSupport(opt_target.value(), "tvm.support.nvcc.supports_bf16")) { return f; @@ -789,7 +789,7 @@ TVM_FFI_STATIC_INIT_BLOCK() { Pass FP8ComputeLegalize(ffi::String promote_dtype) { auto pass_func = [=](Function f, IRModule m, PassContext ctx) { - auto opt_target = f->GetAttr(tvm::attr::kTarget); + auto opt_target = f->GetAttr(tvm::attr::kTarget); if (opt_target.has_value() && CheckDataTypeSupport(opt_target.value(), "tvm.support.nvcc.supports_fp8")) { return f; @@ -807,7 +807,7 @@ TVM_FFI_STATIC_INIT_BLOCK() { Pass FP8StorageLegalize() { auto pass_func = [=](Function f, IRModule m, PassContext ctx) { - auto opt_target = f->GetAttr(tvm::attr::kTarget); + auto opt_target = f->GetAttr(tvm::attr::kTarget); if (opt_target.has_value() && CheckDataTypeSupport(opt_target.value(), "tvm.support.nvcc.supports_fp8")) { return f; diff --git a/tirx/src/tirx/transform/vectorize_loop.cc b/tirx/src/tirx/transform/vectorize_loop.cc index ff092c953049..4fb101215592 100644 --- a/tirx/src/tirx/transform/vectorize_loop.cc +++ b/tirx/src/tirx/transform/vectorize_loop.cc @@ -34,9 +34,9 @@ #include #include #include +#include #include #include -#include #include #include @@ -1340,7 +1340,7 @@ Stmt PreserveUnrollPolicy(Stmt body, const ffi::Map& anno } }; auto transfer = ffi::make_object(); - for (const char* key : {attr::kAutoUnrollMaxStep, attr::kUnrollExplicit}) { + for (const char* key : {attr::kAutoUnrollMaxStep.name, attr::kUnrollExplicit.name}) { if (auto value = annotations.Get(key); value.has_value() && value.value() != nullptr) { transfer->policy.Set(key, value.value()); } @@ -1357,7 +1357,7 @@ class LoopVectorizer : public StmtExprMutator { return StmtExprMutator::Mutate(value, inplace_mode); } - explicit LoopVectorizer(DictAttrs attrs) : target_(attrs.GetAttr(tvm::attr::kTarget)) {} + explicit LoopVectorizer(DictAttrs attrs) : target_(attrs.GetAttr(tvm::attr::kTarget)) {} UnchangedOr Mutate_(const ForNode* op, InplaceMode inplace_mode) final { if (op->kind == ForKind::kVectorized) { diff --git a/tirx/tests/cpp/ir/tensor_load_test.cc b/tirx/tests/cpp/ir/tensor_load_test.cc index 577e127dcbe8..a74ad753d38d 100644 --- a/tirx/tests/cpp/ir/tensor_load_test.cc +++ b/tirx/tests/cpp/ir/tensor_load_test.cc @@ -51,12 +51,12 @@ TVM_FFI_STATIC_INIT_BLOCK() { // A packed hook can return an arbitrary value; native hooks must obey their // non-nullable PrimType return contract before entering the native ABI. refl::TypeAttrDef().attr( - type_attr::kTensorLoadInferType, + type_attr::kTensorLoadInferType.name, ffi::Function::FromTyped([](const TensorLoadNode* load) -> ffi::Any { if (load->indices.empty()) return Type::Missing(); return nullptr; })); - refl::TypeAttrDef().attr(type_attr::kTensorLoadInferType, + refl::TypeAttrDef().attr(type_attr::kTensorLoadInferType.name, FTensorLoadInferType::FromNative<&InferTestLoad>()); } diff --git a/tirx/tests/cpp/target/target_test.cc b/tirx/tests/cpp/target/target_test.cc index 773ddc311e70..d8c3eb974e2d 100644 --- a/tirx/tests/cpp/target/target_test.cc +++ b/tirx/tests/cpp/target/target_test.cc @@ -24,7 +24,34 @@ using namespace tvm; TVM_FFI_STATIC_INIT_BLOCK() { - TargetKindDef("TestTargetKind").set_default_device_type(kDLCPU).def_option("my_bool"); + TargetKindDef("TestTargetKind") + .set_default_device_type(kDLCPU) + .def_option("my_bool") + .def_option>(backend::attr::kBackendConfig.name); +} + +TEST(TargetCreation, TypedBackendConfig) { + constexpr auto key = backend::attr::kBackendConfig; + using Config = decltype(key)::ValueType; + Config fallback{{"cuda", Config{{"arch", ffi::String("sm_80")}}}}; + Target absent("TestTargetKind"); + EXPECT_FALSE(absent->GetAttr(key).has_value()); + EXPECT_TRUE(absent->GetAttr(key, fallback).value().same_as(fallback)); + EXPECT_TRUE(absent->GetAttr(key, ffi::Optional(fallback)).value().same_as(fallback)); + EXPECT_FALSE(absent->GetAttr(key, std::nullopt).has_value()); + + Config config{{"cuda", Config{{"arch", ffi::String("sm_100a")}}}}; + Target present( + ffi::Map{{"kind", ffi::String("TestTargetKind")}, {key.name, config}}); + auto result = present->GetAttr(key, fallback).value(); + EXPECT_EQ(result.at("cuda").cast().at("arch").cast(), "sm_100a"); + Target round_trip(present->ToConfig()); + EXPECT_EQ( + round_trip->GetAttr(key).value().at("cuda").cast().at("arch").cast(), + "sm_100a"); + constexpr ffi::reflection::AttrKey wrong_type{key.name}; + EXPECT_THROW(present->GetAttr(wrong_type), ffi::Error); + EXPECT_THROW(present->GetAttr(wrong_type, ffi::String("fallback")), ffi::Error); } ffi::Map TestTargetParser(ffi::Map target) { diff --git a/tirx/tests/python/tirx/codegen/test_cuda.py b/tirx/tests/python/tirx/codegen/test_cuda.py index 629880d3d708..d34b68520c97 100644 --- a/tirx/tests/python/tirx/codegen/test_cuda.py +++ b/tirx/tests/python/tirx/codegen/test_cuda.py @@ -209,17 +209,28 @@ def main(data: T.handle("float32")): ) -def test_typed_buffer_bindings(monkeypatch): +@pytest.mark.parametrize("alignment", [16, 64, 128]) +def test_typed_buffer_bindings(monkeypatch, alignment): @T.function def main(data: T.handle("float32"), output: T.handle): T.func_attr({"global_symbol": "main", "calling_conv": 2, "tirx.kernel_launch_params": []}) A = T.decl_tensor(data, ty_args=[T.Tensor((1,), "float32")]) B = T.decl_tensor(output, ty_args=[T.Tensor((1,), "float32")]) - storage = T.alloc_tensor(ty_args=[T.Tensor((1,), "uint32", scope="shared")]) - S = T.decl_tensor(storage.data, ty_args=[T.Tensor((1,), "float32", scope="shared")]) - pool = T.alloc_tensor(ty_args=[T.Tensor((4,), "uint8", scope="shared.dyn")]) - D = T.decl_tensor(pool.data, ty_args=[T.Tensor((1,), "float32", scope="shared.dyn")]) - alias = T.decl_tensor(D.data, ty_args=[T.Tensor((1,), "float32", scope="shared.dyn")]) + storage = T.alloc_tensor( + ty_args=[T.Tensor((1,), "uint32", scope="shared", align=alignment)] + ) + S = T.decl_tensor( + storage.data, ty_args=[T.Tensor((1,), "float32", scope="shared", align=alignment)] + ) + pool = T.alloc_tensor( + ty_args=[T.Tensor((4,), "uint8", scope="shared.dyn", align=alignment)] + ) + D = T.decl_tensor( + pool.data, ty_args=[T.Tensor((1,), "float32", scope="shared.dyn", align=alignment)] + ) + alias = T.decl_tensor( + D.data, ty_args=[T.Tensor((1,), "float32", scope="shared.dyn", align=alignment)] + ) S[0] = A[0] D[0] = S[0] B[0] = alias[0] @@ -232,9 +243,9 @@ def main(data: T.handle("float32"), output: T.handle): ).inspect_source() assert "float* data, void* output" in source assert "float* B = ((float*)output);" in source - assert "__shared__ alignas(64) uint storage_ptr[1];" in source + assert f"__shared__ alignas({alignment}) uint storage_ptr[1];" in source assert "float* S = ((float*)storage_ptr);" in source - assert "extern __shared__ __align__(64) uchar pool_ptr[];" in source + assert f"extern __shared__ __align__({alignment}) uchar pool_ptr[];" in source assert "float* D = ((float*)pool_ptr);" in source assert "S[0] = data[0];" in source assert "D[0] = S[0];" in source From d30464470afa840ca8c4d33a567bcec8c7b16694 Mon Sep 17 00:00:00 2001 From: Tianqi Chen Date: Sat, 17 Oct 2026 03:19:45 +0000 Subject: [PATCH 2/4] [FIX] Deduce typed attribute calls in Adreno helpers Remove template disambiguators from dependent GetAttr calls whose value types are inferred from their attribute keys. --- src/relax/backend/adreno/annotate_custom_storage.cc | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/src/relax/backend/adreno/annotate_custom_storage.cc b/src/relax/backend/adreno/annotate_custom_storage.cc index 8b429aa30b7c..a07e403bcbd6 100644 --- a/src/relax/backend/adreno/annotate_custom_storage.cc +++ b/src/relax/backend/adreno/annotate_custom_storage.cc @@ -383,7 +383,7 @@ class CollectConsumerScopeInfo : public ExprVisitor { template ffi::Array ExtractAttrs(const T& func) { ffi::Array op_attrs; - ffi::Optional attrs = func->template GetAttr(tvm::relax::attr::kOpAttrs); + ffi::Optional attrs = func->GetAttr(tvm::relax::attr::kOpAttrs); if (attrs) { if (auto val = attrs.value().as()) { op_attrs.push_back(val.value()); @@ -396,7 +396,7 @@ class CollectConsumerScopeInfo : public ExprVisitor { template ffi::Optional ExtractPattern(const T& func) { - ffi::Optional op_pat = func->template GetAttr(tvm::relax::attr::kOpPattern); + ffi::Optional op_pat = func->GetAttr(tvm::relax::attr::kOpPattern); return op_pat; } From a25c43700c5c89f979d5cd499c068612b3496621 Mon Sep 17 00:00:00 2001 From: Tianqi Chen Date: Sat, 17 Oct 2026 11:21:26 +0000 Subject: [PATCH 3/4] [FIX][Web] Use FFI visibility for WebAssembly runtime exports --- web/emcc/tvmjs_support.cc | 8 ++++---- 1 file changed, 4 insertions(+), 4 deletions(-) diff --git a/web/emcc/tvmjs_support.cc b/web/emcc/tvmjs_support.cc index 30411263f3cc..8183e3200feb 100644 --- a/web/emcc/tvmjs_support.cc +++ b/web/emcc/tvmjs_support.cc @@ -45,13 +45,13 @@ extern "C" { * \param size The size of the space. * \return The allocated space. */ -TVM_DLL void* TVMWasmAllocSpace(int size); +TVM_FFI_DLL void* TVMWasmAllocSpace(int size); /*! * \brief Free the space allocated by TVMWasmAllocSpace. * \param data The data pointer. */ -TVM_DLL void TVMWasmFreeSpace(void* data); +TVM_FFI_DLL void TVMWasmFreeSpace(void* data); /*! * \brief Create ffi::Function from a resource handle. @@ -60,13 +60,13 @@ TVM_DLL void TVMWasmFreeSpace(void* data); * \sa TVMWasmPackedCFunc, TVMWasmPackedCFuncFinalizer 3A * \return 0 if success. */ -TVM_DLL int TVMFFIWasmFunctionCreate(void* resource_handle, TVMFFIObjectHandle* out); +TVM_FFI_DLL int TVMFFIWasmFunctionCreate(void* resource_handle, TVMFFIObjectHandle* out); /*! * \brief Get the last error message. * \return The last error message. */ -TVM_DLL const char* TVMFFIWasmGetLastError(); +TVM_FFI_DLL const char* TVMFFIWasmGetLastError(); // --- APIs to be implemented by the frontend. --- From cd13a16d1c259755b38f66865264c35bceadd9e6 Mon Sep 17 00:00:00 2001 From: Tianqi Chen Date: Sat, 17 Oct 2026 11:21:26 +0000 Subject: [PATCH 4/4] [DOCS] Keep consecutive BYOC prose in one gallery text block --- docs/how_to/tutorials/bring_your_own_codegen.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/docs/how_to/tutorials/bring_your_own_codegen.py b/docs/how_to/tutorials/bring_your_own_codegen.py index e5978aaaa51f..7b3bc9bfffaf 100644 --- a/docs/how_to/tutorials/bring_your_own_codegen.py +++ b/docs/how_to/tutorials/bring_your_own_codegen.py @@ -65,7 +65,7 @@ # Steps 1 and 2 are pure Python and run anywhere; steps 3 and 4 need the # backend's codegen and runtime compiled into TVM, which is why the # build-and-run cells below are guarded. -###################################################################### +# # Step 1: Import the backends to register their patterns # ------------------------------------------------------ #