This is an automated email from the ASF dual-hosted git repository.
spectrometerHBH pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/tvm.git
The following commit(s) were added to refs/heads/main by this push:
new e6d5252f25 [REFACTOR][TIR] Phase out redundant TIRx attr names (#20017)
e6d5252f25 is described below
commit e6d5252f25ba0a7283a12c46c142796cba6eeb68
Author: Tianqi Chen <[email protected]>
AuthorDate: Fri Jul 17 05:27:07 2026 +0800
[REFACTOR][TIR] Phase out redundant TIRx attr names (#20017)
TIRx currently re-declares many attributes that are owned by S-TIR,
along with several legacy attributes that no longer have consumers.
Keeping duplicate names makes ownership ambiguous and lets consumers
accidentally depend on transitive declarations.
This change leaves TIRx with only the attributes needed by low-level
lowering and code generation, removes unused legacy declarations, moves
`buffer_bound` to S-TIR, and updates the remaining cross-layer consumers
to include and use the S-TIR declarations explicitly. Attribute string
values remain unchanged.
---
include/tvm/s_tir/stmt.h | 3 +
include/tvm/tirx/stmt.h | 233 -------------------------------
src/backend/cuda/codegen/codegen_cuda.cc | 9 +-
src/s_tir/transform/bound_checker.cc | 3 +-
src/te/operation/create_primfunc.cc | 1 +
src/tirx/transform/lower_tirx_opaque.cc | 2 +-
6 files changed, 12 insertions(+), 239 deletions(-)
diff --git a/include/tvm/s_tir/stmt.h b/include/tvm/s_tir/stmt.h
index c0dcf942cf..72aae3642d 100644
--- a/include/tvm/s_tir/stmt.h
+++ b/include/tvm/s_tir/stmt.h
@@ -43,6 +43,9 @@ constexpr const char* async_wait_inflight_count =
"async_wait_inflight_count";
*/
constexpr const char* async_scope = "async_scope";
+/*! \brief Mark stores/loads with their bounds. */
+constexpr const char* buffer_bound = "buffer_bound";
+
/*!
* \brief Marks production of double buffer data
*/
diff --git a/include/tvm/tirx/stmt.h b/include/tvm/tirx/stmt.h
index 5c6779b257..1d06f1568c 100644
--- a/include/tvm/tirx/stmt.h
+++ b/include/tvm/tirx/stmt.h
@@ -983,8 +983,6 @@ class ScopeIdDefStmt : public Stmt {
/*! \brief namespace of possible attributes in AttrStmt.attr_key */
namespace attr {
-/*! \brief Mark stores/loads with their bounds. */
-constexpr const char* buffer_bound = "buffer_bound";
/*!
* \brief Mark the scope as when computation start to happen.
* This can hint some code generator to create a new function for compute.
@@ -1008,10 +1006,6 @@ constexpr const char* pragma_auto_unroll_max_step =
"pragma_auto_unroll_max_step
constexpr const char* pragma_import_c = "pragma_import_c";
/*! \brief Import llvm source or file into the final code gen module */
constexpr const char* pragma_import_llvm = "pragma_import_llvm";
-/*! \brief Mark region is guarded by the pragma extension */
-constexpr const char* pragma_scope_prefix = "pragma_";
-/*! \brief Try to modify the AST to support Tensor Core */
-constexpr const char* pragma_tensor_core = "pragma_tensor_core";
/*! \brief Pragma: unroll explicit */
constexpr const char* pragma_unroll_explicit = "pragma_unroll_explicit";
/*! \brief Mark storage alignment requirement of buffers */
@@ -1020,239 +1014,12 @@ constexpr const char* storage_alignment =
"storage_alignment";
constexpr const char* thread_extent = "thread_extent";
/*! \brief Annotation key on AllocBuffer marking the allocation as volatile. */
constexpr const char* kVolatile = "tirx.volatile";
-/*!
- * \brief Marks the layout transforms to be used for a tensor.
- *
- * Only applies to a DataProducer, as it should be made part of the
- * PrimFunc attributes for TIR.
- */
-constexpr const char* layout_transforms = "layout_transforms";
-/*!
- * \brief Marks the physical axis separators
- *
- * Only applies to a DataProducer, as it should be made part of the
- * Buffer definition in a PrimFunc. See `BufferNode::axis_separators`
- * for more details.
- */
-constexpr const char* axis_separators = "axis_separators";
-/*!
- * \brief Marks production of double buffer data
- */
-constexpr const char* double_buffer_scope = "double_buffer_scope";
-/*!
- * \brief Marks region used by double buffer write
- */
-constexpr const char* double_buffer_write = "double_buffer_write";
-/*! \brief Mark of scan update scope */
-constexpr const char* scan_update_scope = "scan_update_scope";
-/*! \brief Mark of scan init scope */
-constexpr const char* scan_init_scope = "scan_init_scope";
-/*!
- * \brief Mark alignment of buffer dimension
- * stmt.node is Tensor
- * stmt.value is tvm_tuple(dim, align, offset)
- * This gives hint to require stride of dim to be k * align + offset.
- */
-constexpr const char* buffer_dim_align = "buffer_dim_align";
/*! \brief Mark buffer initial addr alignment in bytes */
constexpr const char* buffer_data_alignment = "buffer_data_alignment";
/*! \brief Mark buffer allocated addr in bytes */
constexpr const char* buffer_allocated_addr = "buffer_allocated_addr";
-/*!
- * \brief Bind the buffer specification to the region of the op
- * When this scope occurs, the stmt.node is a ffi::Array<NodeRef> = [buffer,
tensor]
- * stmt.value is a tvm_tuple(min0, extent0, min1, extent1, ...).
- * The scope represents that we need to bind the storage region of tensor to
buffer.
- * This will affect replacement of some variables inside the scope that
- * corresponds to field of buffer to be the actual expressions of tensor
during
- * storage flattening phase.
- */
-constexpr const char* buffer_bind_scope = "buffer_bind_scope";
-// Pipeline related attributes
-/*! \brief channel read scope */
-constexpr const char* channel_read_scope = "channel_read_scope";
-/*! \brief Advance step of channel after end of scope */
-constexpr const char* channel_read_advance = "channel_read_advance";
-/*! \brief channel write scope */
-constexpr const char* channel_write_scope = "channel_write_scope";
-/*! \brief Advance step of channel after end of scope */
-constexpr const char* channel_write_advance = "channel_write_advance";
-/*! \brief pipeline stage scope, implies always execution */
-constexpr const char* pipeline_stage_scope = "pipeline_stage_scope";
-/*! \brief pipeline execution scope, implies the scope can be pipelined. */
-constexpr const char* pipeline_exec_scope = "pipeline_exec_scope";
-
-/*!
- * \brief Mark that the attached statement runs asynchronously.
- */
-constexpr const char* async_scope = "async_scope";
-
-/*!
- * \brief Annotations for invoking and synchronizing asynchronous operations.
-
- * Synchronization is done in terms of "queue": It is an abstract entity
associated
- * with each asynchronous unit, and it tracks invocations and completions of
asynchronous
- * operations in the FIFO order.
- *
- * Similarly to PTX instructions commit_group and wait_group, these
annotations express
- * synchronization by "counting":
- *
- * async_commit_queue(i): Group one or more invocations of async operations in
the given scope,
- * and "commit" (or push) them to the queue i. A group of operations committed
together is
- * awaited as one chunk. Groups committed to the same queue complete in the
FIFO order.
- *
- * async_wait_queue(i, N): Block until only N most recent committed groups are
still in-flight at
- * the queue i. N does not have to be a constant, but some backends may
require a constant count.
-*/
-constexpr const char* async_commit_queue_scope = "async_commit_queue_scope";
-constexpr const char* async_wait_queue_scope = "async_wait_queue_scope";
-constexpr const char* async_wait_inflight_count = "async_wait_inflight_count";
-
-/*!
- * \brief Mark that the shape of TensorCore fragment
- */
-constexpr const char* fragment_shape = "fragment_shape";
-
-/*!
- * \brief Mark that the layout of TensorCore fragment
- */
-constexpr const char* fragment_layout = "fragment_layout";
-
-/*!
- * \brief Mark that the kernel is hand threaded and doesn't need syncs inserted
- */
-constexpr const char* hand_threaded = "hand_threaded";
-
-/*!
- * \brief Mark whether the script-completer need to fill in missing access
region
- * during script parsing.
- * \note The result should be a integer mask with range [0, 4).
- * if (mask & 1) the read region should be detected,
- * if (mask & 2) the write region should be detected.
- */
-constexpr const char* script_parsing_detect_access =
"tirx.script_parsing_detect_access";
-
-/*!
- * \brief Mark that the loop should be partitioned.
- */
-constexpr const char* pragma_loop_partition_hint =
"pragma_loop_partition_hint";
-
-/*! \brief Mark the stage of a statement in the software pipeline */
-constexpr const char* software_pipeline_stage = "software_pipeline_stage";
-
-/*! \brief Mark the order of a statement in the software pipeline */
-constexpr const char* software_pipeline_order = "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* software_pipeline_async_stages =
"software_pipeline_async_stages";
-
-/*! \brief Mark the buffers which is const access and can be transformed
layout. */
-constexpr const char* layout_free_buffers = "layout_free_buffers";
-
-/*! \brief Mark the local stage for the shared memory access should be added.
*/
-constexpr const char* manifest_shared_memory_local_stage =
- "tirx.manifest_shared_memory_local_stage";
-
-/*! \brief Mark the tiling structure of blocks that are applied by rule
Multi-Level-Tiling */
-constexpr const char* meta_schedule_tiling_structure =
"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* meta_schedule_cooperative_fetch =
"meta_schedule.cooperative_fetch";
-
-/*! \brief The allowed range of thread extent in thread bindings */
-constexpr const char* meta_schedule_thread_extent_low_inclusive =
- "meta_schedule.thread_extent_low_inclusive";
-
-/*! \brief The allowed range of thread extent in thread bindings */
-constexpr const char* meta_schedule_thread_extent_high_inclusive =
- "meta_schedule.thread_extent_high_inclusive";
-
-/*! \brief Mark the block whose producer needs to be applied by rule
Random-Compute-Location */
-constexpr const char* meta_schedule_random_compute_producer =
- "meta_schedule.random_compute_producer";
-
-/*! \brief Mark auto-parallel setting on the block. */
-constexpr const char* meta_schedule_parallel = "meta_schedule.parallel";
-
-/*! \brief Mark auto-vectorize setting on the block. */
-constexpr const char* meta_schedule_vectorize = "meta_schedule.vectorize";
-
-/*! \brief Mark auto-unroll setting on the block. */
-constexpr const char* meta_schedule_unroll_explicit =
"meta_schedule.unroll_explicit";
-
-/*! \brief Mark auto-unroll setting on the block. */
-constexpr const char* meta_schedule_unroll_implicit =
"meta_schedule.unroll_implicit";
-
-/*! \brief Mark that a block should be further rewritten using tensorization.
*/
-constexpr const char* meta_schedule_auto_tensorize =
"meta_schedule.auto_tensorize";
-
-/*! \brief Mark that a block is a preprocessor block for layout rewrite. */
-constexpr const char* meta_schedule_layout_rewrite_preproc =
"meta_schedule.layout_rewrite_preproc";
-/*!
- * \brief Mark that the init statement of a block should be further rewritten
using tensorization.
- */
-constexpr const char* meta_schedule_auto_tensorize_init =
"meta_schedule.auto_tensorize_init";
-
-/*!
- * \brief Mark that the block need to add predicate for block var bounds
during lowering
- */
-constexpr const char* require_block_var_bound_predicate =
"require_bound_predicate";
-
-/*! \brief Mark that tensor core is enabled in the PrimExpr */
-constexpr const char* meta_schedule_tensor_core_enabled =
"meta_schedule.tensor_core_enabled";
-
-/*!
- * \brief Mark a block as generated by cache_read or cache_write block.
- * 0 means cache_read; 1 means cache_write.
- * \sa meta_schedule_cache_type_read
- * \sa meta_schedule_cache_type_write
- */
-constexpr const char* meta_schedule_cache_type = "meta_schedule.cache_type";
-
-/*! \sa meta_schedule_cache_type */
-constexpr const int meta_schedule_cache_type_read = 0;
-
-/*! \sa meta_schedule_cache_type */
-constexpr const int meta_schedule_cache_type_write = 1;
-
-/*! \brief Mark auto copy for memhammer */
-constexpr const char* auto_copy = "auto_copy";
-
-/*! \brief Mark local stage constraint on data copy */
-constexpr const char* local_stage = "local_stage";
-
-/*! \brief Mark vectorization length constraint on block */
-constexpr const char* vector_bytes = "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* warp_execution = "warp_execution";
-
-/*! \brief Mark that a block is disallowed in auto inline. */
-constexpr const char* meta_schedule_inline_rule = "meta_schedule.inline_rule";
-
-/*! \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* explicit_read_region = "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* explicit_write_region = "explicit_write_region";
constexpr const char* tensorized_nki_instruction =
"tensorized_nki_instruction";
-/*! \brief ,ark a ForNode represent an irregular loop of non-structural
control flow edges. */
-constexpr const char* irregular_loop_mark = "irregular_loop_mark";
-
/*!
* \brief Mark the kernel as persistent.
*/
diff --git a/src/backend/cuda/codegen/codegen_cuda.cc
b/src/backend/cuda/codegen/codegen_cuda.cc
index f2d6b410af..5576d8a578 100644
--- a/src/backend/cuda/codegen/codegen_cuda.cc
+++ b/src/backend/cuda/codegen/codegen_cuda.cc
@@ -26,6 +26,7 @@
#include <tvm/arith/analyzer.h>
#include <tvm/ffi/function.h>
#include <tvm/ffi/reflection/registry.h>
+#include <tvm/s_tir/stmt.h>
#include <tvm/tirx/index_map.h>
#include <tvm/tirx/stmt_functor.h>
@@ -1578,15 +1579,15 @@ void CodeGenCUDA::VisitExpr_(const CallNode* op,
std::ostream& os) {
}
void CodeGenCUDA::VisitStmt_(const AttrStmtNode* op) {
- if (op->attr_key == tirx::attr::fragment_shape) {
+ if (op->attr_key == s_tir::attr::fragment_shape) {
const VarNode* buffer = op->node.as<VarNode>();
const StringImmNode* shape_str = op->value.as<StringImmNode>();
fragment_shapes[buffer] = shape_str->value;
- } else if (op->attr_key == tirx::attr::fragment_layout) {
+ } else if (op->attr_key == s_tir::attr::fragment_layout) {
const VarNode* buffer = op->node.as<VarNode>();
const StringImmNode* layout_str = op->value.as<StringImmNode>();
fragment_layouts[buffer] = layout_str->value;
- } else if (op->attr_key == tirx::attr::async_commit_queue_scope) {
+ } else if (op->attr_key == s_tir::attr::async_commit_queue_scope) {
const IntImmNode* queue_id = op->value.as<IntImmNode>();
TVM_FFI_ICHECK(queue_id && queue_id->value == 0)
<< "For CUDA, the index of an async queue must be 0.";
@@ -1598,7 +1599,7 @@ void CodeGenCUDA::VisitStmt_(const AttrStmtNode* op) {
this->VisitExpr(commit_group, this->stream);
this->stream << ";\n";
return;
- } else if (op->attr_key == tirx::attr::async_wait_queue_scope) {
+ } else if (op->attr_key == s_tir::attr::async_wait_queue_scope) {
auto wait_attrs = GetAsyncWaitAttributes(op);
auto queue_id = wait_attrs.first.as<IntImmNode>();
TVM_FFI_ICHECK(queue_id && queue_id->value == 0)
diff --git a/src/s_tir/transform/bound_checker.cc
b/src/s_tir/transform/bound_checker.cc
index 51a4300ca4..73436123d0 100644
--- a/src/s_tir/transform/bound_checker.cc
+++ b/src/s_tir/transform/bound_checker.cc
@@ -26,6 +26,7 @@
#include <tvm/ffi/cast.h>
#include <tvm/ffi/function.h>
#include <tvm/ffi/reflection/registry.h>
+#include <tvm/s_tir/stmt.h>
#include <tvm/s_tir/transform.h>
#include <tvm/tirx/builtin.h>
#include <tvm/tirx/expr.h>
@@ -50,7 +51,7 @@ class BoundCollector : public StmtVisitor {
BoundCollector() {}
void VisitStmt_(const AttrStmtNode* op) final {
- if (op->attr_key == tirx::attr::buffer_bound) {
+ if (op->attr_key == s_tir::attr::buffer_bound) {
const VarNode* key = op->node.as<VarNode>();
const CallNode* container = op->value.as<CallNode>();
if (key && container) {
diff --git a/src/te/operation/create_primfunc.cc
b/src/te/operation/create_primfunc.cc
index d4c1de9ec0..926676fbe5 100644
--- a/src/te/operation/create_primfunc.cc
+++ b/src/te/operation/create_primfunc.cc
@@ -24,6 +24,7 @@
#include <tvm/ffi/function.h>
#include <tvm/ffi/reflection/registry.h>
#include <tvm/ir/unique_name_supply.h>
+#include <tvm/s_tir/stmt.h>
#include <tvm/te/operation.h>
#include <tvm/tirx/analysis.h>
#include <tvm/tirx/function.h>
diff --git a/src/tirx/transform/lower_tirx_opaque.cc
b/src/tirx/transform/lower_tirx_opaque.cc
index 685b9968b5..d2bfec19d8 100644
--- a/src/tirx/transform/lower_tirx_opaque.cc
+++ b/src/tirx/transform/lower_tirx_opaque.cc
@@ -124,7 +124,7 @@ class TIRxOpaqueLower : public StmtExprMutator {
ffi::String thread_tag = op->thread_binding.value()->thread_tag;
body = MakeLaunchThread(min, extent, op->loop_var, thread_tag, body);
} else if (is_one(extent) && op->annotations.empty() &&
- !op->annotations.count(tirx::attr::irregular_loop_mark)) {
+ !op->annotations.count(s_tir::attr::irregular_loop_mark)) {
// Case 2. Unit loop elimination
return body;
} else {
