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/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 # ------------------------------------------------------ # 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..a07e403bcbd6 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->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->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 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. ---