This is an automated email from the ASF dual-hosted git repository.
tqchen pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/tvm.git
The following commit(s) were added to refs/heads/main by this push:
new 1dc56a6bd1 [REFACTOR][TIR] Attach fragment metadata to allocation
calls (#20520)
1dc56a6bd1 is described below
commit 1dc56a6bd1c0c9a0c3f666ffefdc9238214a48f0
Author: Tianqi Chen <[email protected]>
AuthorDate: Fri Oct 2 14:30:45 2026 -0400
[REFACTOR][TIR] Attach fragment metadata to allocation calls (#20520)
Move TensorCore fragment shape and layout metadata from AttrStmt
wrappers onto their owning allocation calls. InferFragment writes direct
attributes, and CUDA and SPIR-V read them when emitting each allocation.
Preserve annotated allocation ownership through storage rewriting,
including custom pass ordering, while keeping unrelated allocation and
external-storage protocols unchanged.
---
include/tvm/s_tir/stmt.h | 4 +--
src/backend/cuda/codegen/codegen_cuda.cc | 37 +++++++++----------
src/backend/cuda/codegen/codegen_cuda.h | 8 ++---
src/backend/vulkan/codegen/codegen_spirv.cc | 9 +++--
src/s_tir/transform/tensorcore_infer_fragment.cc | 45 ++++++++++++------------
src/tirx/transform/storage_rewrite.cc | 28 ++++++++++++---
6 files changed, 70 insertions(+), 61 deletions(-)
diff --git a/include/tvm/s_tir/stmt.h b/include/tvm/s_tir/stmt.h
index 810435142d..d79674ff3e 100644
--- a/include/tvm/s_tir/stmt.h
+++ b/include/tvm/s_tir/stmt.h
@@ -226,12 +226,12 @@ constexpr const char* double_buffer_scope =
"double_buffer_scope";
constexpr const char* double_buffer_write = "double_buffer_write";
/*!
- * \brief Mark that the shape of TensorCore fragment
+ * \brief String-valued allocation Call attribute containing the TensorCore
fragment shape
*/
constexpr const char* fragment_shape = "fragment_shape";
/*!
- * \brief Mark that the layout of TensorCore fragment
+ * \brief String-valued allocation Call attribute containing the TensorCore
fragment layout
*/
constexpr const char* fragment_layout = "fragment_layout";
diff --git a/src/backend/cuda/codegen/codegen_cuda.cc
b/src/backend/cuda/codegen/codegen_cuda.cc
index 31ce6f3692..8945d5bce3 100644
--- a/src/backend/cuda/codegen/codegen_cuda.cc
+++ b/src/backend/cuda/codegen/codegen_cuda.cc
@@ -1544,15 +1544,7 @@ void CodeGenCUDA::Dispatch_(const CallNode* op,
std::ostream& os) {
}
void CodeGenCUDA::Dispatch_(const AttrStmtNode* op) {
- 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 == 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 == s_tir::attr::async_commit_queue_scope) {
+ 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.";
@@ -1607,6 +1599,16 @@ void CodeGenCUDA::DispatchAllocBuffer(const BindNode*
op, const CallNode* buffer
TVM_FFI_ICHECK(buffer.defined());
std::string vid = AllocVarID(buffer.get(), buffer.name() + "_ptr");
+ std::string fragment_shape;
+ std::string fragment_layout;
+ if (std::string(scope).find("wmma.") == 0) {
+ auto shape = annotations->dict.Get(s_tir::attr::fragment_shape);
+ TVM_FFI_ICHECK(shape.has_value()) << "Cannot find shape of the wmma
fragment " << buffer.name();
+ fragment_shape = shape.value().as_or_throw<ffi::String>();
+ if (auto layout = annotations->dict.Get(s_tir::attr::fragment_layout)) {
+ fragment_layout = layout.value().as_or_throw<ffi::String>();
+ }
+ }
this->PrintIndent();
if (std::string(scope).find("wmma.") == 0) {
if (scope == "wmma.matrix_a" || scope == "wmma.matrix_b") {
@@ -1625,7 +1627,7 @@ void CodeGenCUDA::DispatchAllocBuffer(const BindNode* op,
const CallNode* buffer
TVM_FFI_ICHECK(supported_wmma_accumulator_dtype)
<< "Accumulator only support half, float and int type for now";
}
- PrintWmmaScope(scope, PrimType(dtype), buffer.get(), stream);
+ PrintWmmaScope(scope, PrimType(dtype), fragment_shape, fragment_layout,
stream);
} else {
PrintStorageScope(scope, stream);
int align = buffer->data_alignment;
@@ -1656,7 +1658,7 @@ void CodeGenCUDA::DispatchAllocBuffer(const BindNode* op,
const CallNode* buffer
TVM_FFI_ICHECK_GT(constant_size, 0) << "Can only handle constant size
stack allocation for now";
if (std::string(scope).find("wmma.") == 0) {
- constant_size = GetWmmaFragmentSize(scope, buffer.get(), constant_size);
+ constant_size = GetWmmaFragmentSize(scope, fragment_shape,
constant_size);
}
bool is_packed_integer_dtype = PrimType(dtype) == PrimType::Int(4) ||
PrimType(dtype) == PrimType::UInt(4) ||
@@ -1952,12 +1954,10 @@ void CodeGenCUDA::Dispatch_(const FloatImmNode* op,
std::ostream& os) { // NOLI
}
void CodeGenCUDA::PrintWmmaScope(const std::string& scope, const PrimType& t,
- const VarNode* variable, std::ostream& os) {
+ const std::string& shape_str, const
std::string& layout_str,
+ std::ostream& os) {
std::stringstream type;
PrintType(t, type);
- TVM_FFI_ICHECK(fragment_shapes.count(variable))
- << "Cannot find shape of the wmma fragment " << variable->name;
- std::string shape_str = fragment_shapes.at(variable);
if ((t.MatchesCode(DLDataTypeCode::kDLInt, DLDataTypeCode::kDLUInt)) &&
t.bits() < 8 &&
t.lanes() == 1) {
type.str(std::string());
@@ -1979,13 +1979,11 @@ void CodeGenCUDA::PrintWmmaScope(const std::string&
scope, const PrimType& t,
}
if (scope == "wmma.matrix_a") {
codegen_tags_.insert("mma");
- std::string layout_str = fragment_layouts[variable];
TVM_FFI_ICHECK_NE(layout_str, "") << "Layout must be defined for matrix_a";
os << "nvcuda::wmma::fragment<nvcuda::wmma::matrix_a, " << shape_str << ",
" << type.str()
<< ", nvcuda::wmma::" << layout_str << ">";
} else if (scope == "wmma.matrix_b") {
codegen_tags_.insert("mma");
- std::string layout_str = fragment_layouts[variable];
TVM_FFI_ICHECK_NE(layout_str, "") << "Layout must be defined for matrix_b";
os << "nvcuda::wmma::fragment<nvcuda::wmma::matrix_b, " << shape_str << ",
" << type.str()
<< ", nvcuda::wmma::" << layout_str << ">";
@@ -2005,11 +2003,8 @@ int stoi(const std::string& str) {
}
}
-int32_t CodeGenCUDA::GetWmmaFragmentSize(const std::string& scope, const
VarNode* variable,
+int32_t CodeGenCUDA::GetWmmaFragmentSize(const std::string& scope, const
std::string& shape_str,
int32_t size) {
- TVM_FFI_ICHECK(fragment_shapes.count(variable))
- << "Cannot find shape of the wmma fragment " << variable->name;
- std::string shape_str = fragment_shapes.at(variable);
std::pair<int32_t, int32_t> dim = GetWmmaFragmentDimSize(shape_str, scope);
if (dim.first * dim.second != 0)
return size / dim.first / dim.second;
diff --git a/src/backend/cuda/codegen/codegen_cuda.h
b/src/backend/cuda/codegen/codegen_cuda.h
index 933bc2a6a7..d383a67426 100644
--- a/src/backend/cuda/codegen/codegen_cuda.h
+++ b/src/backend/cuda/codegen/codegen_cuda.h
@@ -124,12 +124,10 @@ class CodeGenCUDA final : public CodeGenC {
// The name prefix of the cuda::barrier::arrival_token array in registers
const std::string cuda_barrier_arrival_token_name_ = "cubar_tok";
- std::unordered_map<const VarNode*, std::string> fragment_shapes;
- std::unordered_map<const VarNode*, std::string> fragment_layouts;
friend void PrintConst(const FloatImmNode* op, std::ostream& os,
CodeGenCUDA* p);
- void PrintWmmaScope(const std::string& scope, const PrimType& t, const
VarNode* variable,
- std::ostream& os);
- int32_t GetWmmaFragmentSize(const std::string& scope, const VarNode*
variable, int32_t size);
+ void PrintWmmaScope(const std::string& scope, const PrimType& t, const
std::string& shape,
+ const std::string& layout, std::ostream& os);
+ int32_t GetWmmaFragmentSize(const std::string& scope, const std::string&
shape, int32_t size);
};
} // namespace codegen
diff --git a/src/backend/vulkan/codegen/codegen_spirv.cc
b/src/backend/vulkan/codegen/codegen_spirv.cc
index c9821b239f..0e48715677 100644
--- a/src/backend/vulkan/codegen/codegen_spirv.cc
+++ b/src/backend/vulkan/codegen/codegen_spirv.cc
@@ -898,7 +898,10 @@ void CodeGenSPIRV::DispatchAllocBuffer(const BindNode* op,
const CallNode* buffe
case runtime::StorageRank::kWMMAMatrixA:
case runtime::StorageRank::kWMMAMatrixB:
case runtime::StorageRank::kWMMAAccumulator: {
- TVM_FFI_ICHECK(fragment_info_.count(var_node));
+ auto shape = annotations->dict.Get(s_tir::attr::fragment_shape);
+ TVM_FFI_ICHECK(shape.has_value())
+ << "Cannot find shape of the wmma fragment " << buffer.name();
+ fragment_info_[var_node] = {shape.value().as_or_throw<ffi::String>()};
fragment_info_[var_node].scope = scope;
etype = GetFragmentSType(var_node, PrimType(dtype));
storage_class = spv::StorageClassFunction;
@@ -989,10 +992,6 @@ void CodeGenSPIRV::Dispatch_(const AttrStmtNode* op) {
var_map_[iv->var.get()] = GetThreadIndex(iv,
op->value.as_or_throw<PrimExpr>());
}
}
- } else 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_info_[buffer] = {shape_str->value};
}
this->Dispatch(op->body);
}
diff --git a/src/s_tir/transform/tensorcore_infer_fragment.cc
b/src/s_tir/transform/tensorcore_infer_fragment.cc
index 32f5bc3ab0..a2faffa6c8 100644
--- a/src/s_tir/transform/tensorcore_infer_fragment.cc
+++ b/src/s_tir/transform/tensorcore_infer_fragment.cc
@@ -209,33 +209,32 @@ class InferFragmenter : public s_tir::StmtExprMutator {
UnchangedOr<Stmt> Mutate_(const BindNode* op, InplaceMode inplace_mode)
final {
if (const auto* call = op->value.as<CallNode>();
call && call->op.same_as(tirx::builtin::alloc_buffer())) {
- return Mutate_AllocBuffer(op, inplace_mode);
+ auto it = fragment_getter.fragments.find(op->var.get());
+ if (it == fragment_getter.fragments.end()) return ffi::Unchanged();
+ const FragmentInfo& info = it->second;
+ DictAttrs attrs = call->attrs.as_or_throw<DictAttrs>();
+ auto annotations = attrs->dict;
+ auto set_annotation = [&](const char* key, const std::string& value) {
+ if (auto existing = annotations.Get(key)) {
+ auto str = existing.value().as<ffi::String>();
+ TVM_FFI_CHECK(str && str.value() == value, ValueError)
+ << "Conflicting " << key << " on fragment " << op->var->name;
+ } else {
+ annotations.Set(key, ffi::String(value));
+ }
+ };
+ set_annotation(
+ s_tir::attr::fragment_shape,
+ std::to_string(info.m) + ", " + std::to_string(info.n) + ", " +
std::to_string(info.k));
+ if (!info.layout.empty()) set_annotation(s_tir::attr::fragment_layout,
info.layout);
+ if (annotations.same_as(attrs->dict)) return ffi::Unchanged();
+ auto value = ffi::GetRef<Call>(call);
+ value.CopyOnWrite()->attrs = DictAttrs(std::move(annotations));
+ return Bind(op->var, std::move(value), op->span);
}
return s_tir::StmtExprMutator::Mutate_(op, inplace_mode);
}
- UnchangedOr<Stmt> Mutate_AllocBuffer(const BindNode* op, InplaceMode
inplace_mode) {
- Stmt stmt =
- s_tir::StmtExprMutator::Mutate_(op,
inplace_mode).ValueOrUnchanged(ffi::GetRef<Stmt>(op));
- const VarNode* buffer = op->var.get();
- if (fragment_getter.fragments.count(buffer)) {
- FragmentInfo info = fragment_getter.fragments.at(buffer);
-
- std::string shape =
- std::to_string(info.m) + ", " + std::to_string(info.n) + ", " +
std::to_string(info.k);
- Expr shape_expr = StringImm(shape);
- Stmt shape_attr = AttrStmt(op->var, s_tir::attr::fragment_shape,
shape_expr, stmt);
- if (info.layout != "") {
- Stmt layout_attr =
- AttrStmt(op->var, s_tir::attr::fragment_layout,
StringImm(info.layout), shape_attr);
- return layout_attr;
- } else {
- return shape_attr;
- }
- }
- return stmt;
- }
-
private:
// Fragment infomation
const FragmentGetter& fragment_getter;
diff --git a/src/tirx/transform/storage_rewrite.cc
b/src/tirx/transform/storage_rewrite.cc
index be31a612b1..3442a16a65 100644
--- a/src/tirx/transform/storage_rewrite.cc
+++ b/src/tirx/transform/storage_rewrite.cc
@@ -780,6 +780,8 @@ class StoragePlanRewriter : public StmtExprMutator {
PrimType elem_type = PrimType::Void();
// Whether any constituent allocation was marked volatile.
bool is_volatile{false};
+ // Fragment metadata belongs to this allocation, not interchangeable
storage.
+ bool has_fragment_metadata{false};
// This is non-zero if this alloc_buffer is folded into another one
// the address(in bits) becomes alloc_var + bits_offset;
// can be effectively converted to the element type.
@@ -794,6 +796,12 @@ class StoragePlanRewriter : public StmtExprMutator {
uint64_t bits_offset{0};
};
+ static bool HasFragmentMetadata(const CallNode* call) {
+ const auto& annotations = call->attrs.as_or_throw<DictAttrs>()->dict;
+ return annotations.count(s_tir::attr::fragment_shape) ||
+ annotations.count(s_tir::attr::fragment_layout);
+ }
+
// Checks whether the storage_scope is especially tagged for a specific
memory.
// Special memory is all combined into a single allocation.
bool IsSpecialTaggedMemory(const StorageScope& scope) {
@@ -836,10 +844,10 @@ class StoragePlanRewriter : public StmtExprMutator {
// try to find merge, for tagged memory
for (size_t i = 0; i < vec.size(); ++i) {
StorageEntry* e = vec[i];
- if (IsSpecialTaggedMemory(e->scope)) {
+ if (!e->has_fragment_metadata && IsSpecialTaggedMemory(e->scope)) {
TVM_FFI_ICHECK_NE(e->const_nbits, 0U) << "Special tagged memory must
be const size";
for (size_t j = 0; j < i; ++j) {
- if (e->scope == vec[j]->scope) {
+ if (!vec[j]->has_fragment_metadata && e->scope == vec[j]->scope) {
vec[j]->merged_children.push_back(e);
break;
}
@@ -849,6 +857,13 @@ class StoragePlanRewriter : public StmtExprMutator {
// Start allocation
for (size_t i = 0; i < vec.size(); ++i) {
StorageEntry* e = vec[i];
+ if (e->has_fragment_metadata) {
+ TVM_FFI_ICHECK_EQ(e->allocs.size(), 1U);
+ TVM_FFI_ICHECK(e->merged_children.empty());
+ e->alloc_var = e->allocs[0]->var;
+ e->alloc_nest.push_back(ffi::GetRef<Bind>(e->allocs[0]));
+ continue;
+ }
// already merged
if (e->bits_offset != 0) continue;
if (e->merged_children.size() != 0) {
@@ -1106,7 +1121,7 @@ class StoragePlanRewriter : public StmtExprMutator {
auto storage_scope = StorageScope::Create(scope);
StorageEntry* dst_entry = nullptr;
// inplace detection
- if (detect_inplace) {
+ if (detect_inplace && !HasFragmentMetadata(call)) {
DLDataType dtype = call->args[1].as_or_throw<DataTypeImm>()->value;
PrimType element_type(dtype);
// only one inplace var for s.stmt
@@ -1115,7 +1130,7 @@ class StoragePlanRewriter : public StmtExprMutator {
if (!inplace_flag.count(src) && alloc_map_.count(src)) {
auto visitor = ffi::make_object<InplaceOpVerifier>();
StorageEntry* src_entry = alloc_map_.at(src);
- if (src_entry->scope == storage_scope &&
+ if (!src_entry->has_fragment_metadata && src_entry->scope ==
storage_scope &&
src_entry->attach_scope_ == thread_scope_ &&
!element_type.IsScalableVector() &&
src_entry->elem_type == element_type.WithLanes(1) &&
visitor->Check(s.stmt, var, src)) {
@@ -1189,6 +1204,7 @@ class StoragePlanRewriter : public StmtExprMutator {
entry->scope = scope;
entry->elem_type = element_type.WithLanes(1);
entry->const_nbits = const_nbits;
+ entry->has_fragment_metadata = HasFragmentMetadata(call);
StorageEntry* e = entry.get();
alloc_vec_.emplace_back(std::move(entry));
return e;
@@ -1225,7 +1241,8 @@ class StoragePlanRewriter : public StmtExprMutator {
bool is_small_array = (scope.tag.length() == 0) && (scope.rank >=
StorageRank::kWarp ||
(is_known_size &&
const_nbits <= 32));
- if (is_scalable_vector || !enable_reuse || is_small_array ||
!is_flat_memory_space) {
+ if (HasFragmentMetadata(call) || is_scalable_vector || !enable_reuse ||
is_small_array ||
+ !is_flat_memory_space) {
return NewAlloc(op, attach_scope, scope, const_nbits);
}
@@ -1281,6 +1298,7 @@ class StoragePlanRewriter : public StmtExprMutator {
TVM_FFI_ICHECK(it != alloc_map_.end());
StorageEntry* e = it->second;
TVM_FFI_ICHECK_NE(e->allocs.size(), 0U);
+ if (e->has_fragment_metadata) return;
// disable reuse of small arrays, they will be lowered to registers in LLVM
// This rules only apply if we are using non special memory
