diff --git a/src/ascend/target_utils.h b/src/ascend/target_utils.h index f124339a7e..fffb36a7dc 100644 --- a/src/ascend/target_utils.h +++ b/src/ascend/target_utils.h @@ -6,12 +6,21 @@ #ifndef TVM_TL_ASCEND_TARGET_UTILS_H_ #define TVM_TL_ASCEND_TARGET_UTILS_H_ +#include + #include #include namespace tvm { namespace tl { +constexpr int64_t kAscendUBMemoryBytes = 248 * 1024; +constexpr int64_t kAscendSimtContextBytes = 32 * 1024; +constexpr int64_t kAscendL1MemoryBytes = 512 * 1024; +constexpr int64_t kAscendL0ABMemoryBytes = 64 * 1024; +constexpr int64_t kAscendL0CMemoryBytes = 256 * 1024; +constexpr int64_t kAscendSSMemoryBytes = 3 * 1024; + bool TargetIsAscend(Target target); bool IsAscendVectorizableFP8(DataType dtype); diff --git a/src/ascend/transform/auto_schedule.cc b/src/ascend/transform/auto_schedule.cc index 1687d6b936..81d7566f9f 100644 --- a/src/ascend/transform/auto_schedule.cc +++ b/src/ascend/transform/auto_schedule.cc @@ -63,6 +63,7 @@ #include "./auto_schedule/task_annotations.h" #include "ascend/op/copy.h" #include "ascend/op/utils.h" +#include "ascend/transform/buffer_offsets.h" #include "ascend/transform/buffer_version.h" #include "backend/common/target_utils.h" #include "op/utils.h" @@ -104,37 +105,34 @@ inline bool HasSimtVF(const tirx::Stmt &stmt) { } inline int64_t GetSharedMemoryLimit(const tirx::Stmt &kernel_body) { - constexpr int64_t kPhysicalUnifiedBufferSize = 248 * 1024; - constexpr int64_t kSimtVFThreadContextSize = 32 * 1024; - return HasSimtVF(kernel_body) - ? kPhysicalUnifiedBufferSize - kSimtVFThreadContextSize - : kPhysicalUnifiedBufferSize; + return HasSimtVF(kernel_body) ? kAscendUBMemoryBytes - kAscendSimtContextBytes + : kAscendUBMemoryBytes; } inline int64_t GetL1MemoryLimit(Target target) { if (TargetIsAscend(target)) { - return 512 * 1024; // 512 KB L1 cache on Ascend + return kAscendL1MemoryBytes; } return 0; } inline int64_t GetL0CMemoryLimit(Target target) { if (TargetIsAscend(target)) { - return 256 * 1024; // 256 KB L0C buffer on Ascend + return kAscendL0CMemoryBytes; } return 0; } inline int64_t GetL0AMemoryLimit(Target target) { if (TargetIsAscend(target)) { - return 64 * 1024; // 64 KB L0A buffer on Ascend + return kAscendL0ABMemoryBytes; } return 0; } inline int64_t GetL0BMemoryLimit(Target target) { if (TargetIsAscend(target)) { - return 64 * 1024; // 64 KB L0B buffer on Ascend + return kAscendL0ABMemoryBytes; } return 0; } @@ -1404,7 +1402,7 @@ static void ScheduleSingleKernel(const Stmt &kernel_body, Target target, unit_builder.SetMemoryLimit("shared.l0a", GetL0AMemoryLimit(target)); unit_builder.SetMemoryLimit("shared.l0b", GetL0BMemoryLimit(target)); // One SS pool is shared by AIC and both AIVs, unlike per-AIV UB storage. - unit_builder.SetMemoryLimit("shared.ss", 3 * 1024); + unit_builder.SetMemoryLimit("shared.ss", kAscendSSMemoryBytes); constexpr int64_t kUnlimitedMemory = 1LL << 40; for (const String &scope : metadata.unlimit_memory_scopes) { unit_builder.SetMemoryLimit(scope, kUnlimitedMemory); @@ -1450,7 +1448,30 @@ tvm::transform::Pass AutoSchedule() { context.root, context.outer_ctx, MakeAutoScheduleExtraInfo); BufferVersionMap buffer_version_overrides = scheduled_tir.metadata.buffer_versions; + BufferOffsets offsets; + if (auto value = context.root->annotations.Get(kBufferOffsets)) { + offsets = ParseBufferOffsets(value.value()); + // Validate before AutoSchedule replaces frontend overrides with + // inferred counts. A mode-only annotation does not fix storage + // size. + for (const auto &[storage, offset] : offsets) { + auto versions = buffer_version_overrides.Get(storage); + ICHECK(versions.has_value() && versions.value() > 0) + << kBufferOffsets << " requires an explicit version count in " + << "T.annotate_buffer_versions for storage " << storage; + } + } ScheduleSingleKernel(context.root->body, target, &scheduled_tir); + for (const auto &[storage, offset] : offsets) { + int requested = buffer_version_overrides.at(storage); + int selected = + scheduled_tir.metadata.buffer_versions.Get(storage).value_or(1); + ICHECK_EQ(selected, requested) + << "Cannot place storage " << storage << ": requested " + << requested << " versions but AutoSchedule selected " + << selected + << "; the buffer must be eligible for the requested versions"; + } WarnOnUnappliedBufferVersionOverrides( buffer_version_overrides, scheduled_tir.metadata.buffer_versions); return EncodeScheduledTIR(std::move(scheduled_tir)); diff --git a/src/ascend/transform/buffer_offsets.h b/src/ascend/transform/buffer_offsets.h new file mode 100644 index 0000000000..f5de20fbd0 --- /dev/null +++ b/src/ascend/transform/buffer_offsets.h @@ -0,0 +1,85 @@ +/*! + * \file buffer_offsets.h + * \brief Exact byte placement of Ascend on-chip storages. + */ + +#ifndef TVM_TL_ASCEND_TRANSFORM_BUFFER_OFFSETS_H_ +#define TVM_TL_ASCEND_TRANSFORM_BUFFER_OFFSETS_H_ + +#include + +#include +#include + +#include "ascend/target_utils.h" +#include "support/check.h" +#include "tir/transforms/ir_utils.h" + +namespace tvm { +namespace tl { + +constexpr const char *kBufferOffsets = "tl.ascend_buffer_offsets"; +using BufferOffsets = ffi::Map; + +struct BufferPlacementInfo { + int alignment; + int64_t capacity; +}; + +inline BufferPlacementInfo GetBufferPlacementInfo(const std::string &scope) { + if (scope == "shared.dyn") + return {32, kAscendUBMemoryBytes}; + if (scope == "shared.l1") + return {512, kAscendL1MemoryBytes}; + if (scope == "shared.l0a" || scope == "shared.l0b") + return {512, kAscendL0ABMemoryBytes}; + if (scope == "shared.l0c") + return {1024, kAscendL0CMemoryBytes}; + if (scope == "shared.ss") + return {32, kAscendSSMemoryBytes}; + LOG(FATAL) + << kBufferOffsets << " does not support storage scope " << scope + << "; expected UB, L1, L0A, L0B, L0C or SS allocations. " + << "Place bound L0 scale-factor handles through their data buffer."; + return {}; +} + +inline BufferOffsets ParseBufferOffsets(const ffi::Any &value) { + auto offsets = value.try_cast(); + ICHECK(offsets.has_value()) + << kBufferOffsets << " must map storage Vars to constant byte offsets"; + ICHECK(!offsets.value().empty()) << kBufferOffsets << " must not be empty"; + for (const auto &[storage, offset] : offsets.value()) { + std::string scope = GetPtrStorageScope(storage); + auto info = GetBufferPlacementInfo(scope); + ICHECK(offset->dtype.is_int() && offset->value >= 0 && + offset->value % info.alignment == 0) + << kBufferOffsets << " requires non-negative, " << info.alignment + << "-byte aligned integer offsets for " << scope << "; storage " + << storage << " has " << offset; + ICHECK_LT(offset->value, info.capacity) + << "Buffer placement for " << storage << " exceeds " << scope + << " capacity of " << info.capacity << " bytes"; + } + return offsets.value(); +} + +inline BufferOffsets +RemapBufferOffsets(const BufferOffsets &offsets, + const ffi::Map &storage_remap) { + BufferOffsets result; + for (const auto &[storage, offset] : offsets) { + tirx::Var remapped = storage_remap.Get(storage).value_or(storage); + auto existing = result.Get(remapped); + ICHECK(!existing.has_value() || existing.value()->value == offset->value) + << "Conflicting " << kBufferOffsets << " after remapping storage " + << remapped; + result.Set(remapped, offset); + } + return result; +} + +} // namespace tl +} // namespace tvm + +#endif // TVM_TL_ASCEND_TRANSFORM_BUFFER_OFFSETS_H_ diff --git a/src/ascend/transform/lower_scheduled_tir.cc b/src/ascend/transform/lower_scheduled_tir.cc index a9a6aba8c7..9a6163d24c 100644 --- a/src/ascend/transform/lower_scheduled_tir.cc +++ b/src/ascend/transform/lower_scheduled_tir.cc @@ -25,6 +25,7 @@ #include "./auto_schedule/kernel_rewriter.h" #include "./auto_schedule/scheduled_tir.h" #include "buffer_alias.h" +#include "buffer_offsets.h" #include #include @@ -763,6 +764,11 @@ tvm::transform::Pass LowerScheduledTIR() { body); writer->annotations.erase(kAivStorageMap); } + if (auto offsets = root->annotations.Get(kBufferOffsets)) { + body = AttrStmt(ParseBufferOffsets(offsets.value()), kBufferOffsets, + Integer(1), std::move(body)); + writer->annotations.erase(kBufferOffsets); + } writer->annotations.erase(kSplitSubcores); writer->body = std::move(body); return root; diff --git a/src/ascend/transform/lower_tile_op.cc b/src/ascend/transform/lower_tile_op.cc index 82a52a4389..0c14f19f30 100644 --- a/src/ascend/transform/lower_tile_op.cc +++ b/src/ascend/transform/lower_tile_op.cc @@ -38,6 +38,7 @@ #include "transform/loop_partition.h" #include "ascend/transform/buffer_alias.h" +#include "ascend/transform/buffer_offsets.h" #include "ascend/transform/buffer_version.h" #include "ascend/transform/vf_regions.h" @@ -333,6 +334,9 @@ class LowerTileOpPass : arith::IRMutatorWithAnalyzer { } Stmt VisitStmt_(const SBlockNode *op) final { + ICHECK(!op->annotations.count(kBufferOffsets)) + << kBufferOffsets << " requires the Ascend scheduled pipeline; " + << "enable TL_ENABLE_AUTO_SCHEDULE"; Map previous_block_annotations = block_annotations_; Map previous_safe_value_map = safe_value_map_; block_annotations_ = op->annotations; @@ -1279,6 +1283,12 @@ class LowerTileOpPass : arith::IRMutatorWithAnalyzer { } } Stmt stmt = arith::IRMutatorWithAnalyzer::VisitStmt_(op); + if (op->attr_key == tl::kBufferOffsets) { + AttrStmt attribute = Downcast(std::move(stmt)); + attribute.CopyOnWrite()->node = + RemapBufferOffsets(ParseBufferOffsets(op->node), var_remap_); + return attribute; + } if (op->attr_key == tl::kAivStorageMap) { AivStorageMap remapped; for (const auto &[storage, sid] : op->node.cast()) { diff --git a/src/ascend/transform/merge_ub_allocations.cc b/src/ascend/transform/merge_ub_allocations.cc index ecc82a15fa..e758440dfe 100644 --- a/src/ascend/transform/merge_ub_allocations.cc +++ b/src/ascend/transform/merge_ub_allocations.cc @@ -22,7 +22,9 @@ * \brief Merge Ascend on-chip allocations from a buffer-alias contract. */ +#include "ascend/target_utils.h" #include "buffer_alias.h" +#include "buffer_offsets.h" #include "merge_ub_common.h" #include @@ -372,15 +374,165 @@ MemoryPlan BuildAliasMemoryPlan(const AllocationMap &allocations, return plan; } +class BufferPlacementValidator : public StmtVisitor { +public: + explicit BufferPlacementValidator(std::unordered_set scopes) + : scopes_(std::move(scopes)) {} + + bool HasSimtVF() const { return has_simt_vf_; } + + void VisitStmt_(const SBlockNode *op) final { + has_simt_vf_ |= op->name_hint == "SIMT_VF"; + StmtVisitor::VisitStmt_(op); + } + + void VisitStmt_(const AllocBufferNode *op) final { + std::string scope = GetPtrStorageScope(op->buffer->data); + if (scope.size() > 4 && scope.substr(scope.size() - 4) == ".dyn" && + scope != "shared.dyn") { + ICHECK(!scopes_.count(scope.substr(0, scope.size() - 4))) + << kBufferOffsets << " cannot coexist with an existing " << scope + << " arena: " << op->buffer; + } + if (scopes_.count(scope)) { + ICHECK(storages_.insert(op->buffer->data).second) + << kBufferOffsets + << " requires one allocation per storage: " << op->buffer; + } + StmtVisitor::VisitStmt_(op); + } + +private: + std::unordered_set scopes_; + std::unordered_set storages_; + bool has_simt_vf_{false}; +}; + +// Allocations are flattened and include every physical version and padding +// element. Reserve exact intervals first, then first-fit the remaining buffers. +// Placement never establishes permission to reuse storage, even when lifetimes +// are disjoint. Private V0/V1 UB intervals occupy different physical memories. +MemoryPlan BuildPlacedMemoryPlan(const AllocationMap &allocations, + const AllocationOrder &allocation_order, + const BufferOffsets &offsets, + const std::string &scope, int align_bytes, + bool has_simt_vf, + const AivStorageMap &aiv_storages) { + auto info = GetBufferPlacementInfo(scope); + int alignment = std::max(align_bytes, info.alignment); + int64_t capacity = info.capacity; + if (scope == "shared.dyn" && has_simt_vf) + capacity -= kAscendSimtContextBytes; + + struct Interval { + Var storage; + int64_t begin; + int64_t end; + int core_mask; + }; + std::vector intervals; + MemoryPlan plan; + int64_t peak = 0; + auto core_mask = [&](const Var &storage) { + if (scope == "shared.dyn") { + if (auto sid = aiv_storages.Get(storage)) + return 1 << sid.value()->value; + } + return 3; + }; + auto size_bytes = [&](const Var &storage) { + auto allocation = allocations.find(storage.get()); + ICHECK(allocation != allocations.end()) + << kBufferOffsets << " references unknown allocation " << storage; + const Buffer &buffer = allocation->second->buffer; + ICHECK(buffer->strides.empty() && is_zero(buffer->elem_offset)) + << kBufferOffsets + << " requires flattened dense allocations: " << buffer; + int64_t elem_bits = + static_cast(buffer->dtype.bits()) * buffer->dtype.lanes(); + ICHECK_GT(elem_bits, 0); + int64_t max_elements = capacity * 8 / elem_bits; + int64_t elements = 1; + for (const PrimExpr &extent : buffer->shape) { + const auto *constant = extent.as(); + ICHECK(constant && constant->value > 0) + << kBufferOffsets << " requires positive constant allocation " + << "sizes; storage " << storage << " has extent " << extent; + ICHECK_LE(constant->value, max_elements / elements) + << "Allocation " << storage << " exceeds " << scope << " capacity of " + << capacity << " bytes"; + elements *= constant->value; + } + return (elements * elem_bits + 7) / 8; + }; + auto reserve = [&](const Var &storage, int64_t begin, int64_t bytes) { + ICHECK_EQ(begin % alignment, 0) + << "Buffer placement for " << storage << " requires " << alignment + << "-byte alignment"; + ICHECK(begin >= 0 && begin <= capacity && bytes <= capacity - begin) + << "Buffer placement for " << storage << " at byte " << begin + << " with size " << bytes << " exceeds " << scope << " capacity of " + << capacity << " bytes"; + int64_t end = begin + bytes; + int mask = core_mask(storage); + for (const Interval &other : intervals) { + ICHECK(!(mask & other.core_mask) || end <= other.begin || + other.end <= begin) + << "Buffer placements overlap in " << scope << ": " << other.storage + << " [" << other.begin << ", " << other.end << ") and " << storage + << " [" << begin << ", " << end << ")"; + } + intervals.push_back({storage, begin, end, mask}); + plan.byte_offsets[storage.get()] = IntImm(DataType::Int(32), begin); + peak = std::max(peak, end); + }; + for (const auto &[storage, offset] : offsets) + reserve(storage, offset->value, size_bytes(storage)); + + for (const VarNode *node : allocation_order) { + Var storage = GetRef(node); + if (offsets.count(storage)) + continue; + int64_t bytes = size_bytes(storage); + int64_t begin = 0; + int mask = core_mask(storage); + std::sort( + intervals.begin(), intervals.end(), + [](const Interval &a, const Interval &b) { return a.begin < b.begin; }); + for (const Interval &other : intervals) { + if (!(mask & other.core_mask) || other.end <= begin) + continue; + if (bytes <= other.begin - begin) + break; + begin = AlignUp(other.end, alignment); + } + reserve(storage, begin, bytes); + } + ICHECK_EQ(plan.byte_offsets.size(), allocations.size()); + int64_t total = AlignUp(peak, alignment); + ICHECK_LE(total, capacity) + << "Buffer placements exceed " << scope << " capacity of " << capacity + << " bytes after alignment"; + plan.total_size = IntImm(DataType::Int(32), total); + return plan; +} + class BufferAliasContractExtractor : public StmtMutator { public: Stmt Extract(const Stmt &stmt) { return VisitStmt(stmt); } bool FoundContract() const { return found_contract_; } const BufferAliasMap &Aliases() const { return aliases_; } const AivStorageMap &AivStorages() const { return aiv_storages_; } + const Optional &Offsets() const { return offsets_; } private: Stmt VisitStmt_(const AttrStmtNode *op) final { + if (op->attr_key == kBufferOffsets) { + ICHECK(!offsets_.has_value()) + << "A PrimFunc must contain one " << kBufferOffsets << " contract"; + offsets_ = ParseBufferOffsets(op->node); + return VisitStmt(op->body); + } if (op->attr_key == kAivStorageMap) { for (const auto &[storage, sid] : op->node.cast()) { ICHECK(sid->value == 0 || sid->value == 1); @@ -411,6 +563,7 @@ class BufferAliasContractExtractor : public StmtMutator { BufferAliasMap aliases_; AivStorageMap aiv_storages_; + Optional offsets_; bool found_contract_{false}; }; @@ -418,16 +571,50 @@ Stmt MergeUBAllocations(Stmt stmt, int align_bytes, bool verbose, bool disable_reuse) { BufferAliasContractExtractor extractor; stmt = extractor.Extract(stmt); - if (!extractor.FoundContract()) + if (!extractor.FoundContract()) { + ICHECK(!extractor.Offsets().has_value()) + << kBufferOffsets << " requires the Ascend scheduled allocation " + << "pipeline and its buffer-alias contract"; return stmt; + } BufferAliasMap aliases = extractor.Aliases(); ValidateBufferAliasMap(aliases); AivStorageMap aiv_storages = extractor.AivStorages(); - auto planner = [aliases = std::move(aliases), aiv_storages, disable_reuse]( - const Stmt &, const AllocationMap &allocations, - const AllocationOrder &allocation_order, - const std::string &scope, int alignment, bool debug) { + std::unordered_map placed_plans; + std::unordered_set placed_scopes; + if (extractor.Offsets().has_value()) { + std::unordered_map scope_offsets; + for (const auto &[storage, offset] : extractor.Offsets().value()) { + std::string scope = GetPtrStorageScope(storage); + placed_scopes.insert(scope); + scope_offsets[scope].Set(storage, offset); + } + BufferPlacementValidator validator(placed_scopes); + validator(stmt); + AllocateCollector collector; + collector(stmt); + for (const auto &[scope, offsets] : scope_offsets) { + const AllocationMap &allocations = + scope == "shared.dyn" ? collector.dyn_shmem_allocs + : collector.ascend_scope_allocs[scope]; + const AllocationOrder &order = scope == "shared.dyn" + ? collector.dyn_shmem_order + : collector.ascend_scope_order[scope]; + placed_plans.emplace( + scope, + BuildPlacedMemoryPlan(allocations, order, offsets, scope, align_bytes, + validator.HasSimtVF(), aiv_storages)); + } + } + auto planner = [aliases = std::move(aliases), aiv_storages, disable_reuse, + placed_plans](const Stmt &, const AllocationMap &allocations, + const AllocationOrder &allocation_order, + const std::string &scope, int alignment, + bool debug) { + auto placed = placed_plans.find(scope.empty() ? "shared.dyn" : scope); + if (placed != placed_plans.end()) + return placed->second; auto plan = [&](const AllocationMap &group, const AllocationOrder &order) { return disable_reuse ? BuildSequentialPlan(group, alignment) : BuildAliasMemoryPlan(group, order, alignment, @@ -463,7 +650,7 @@ Stmt MergeUBAllocations(Stmt stmt, int align_bytes, bool verbose, return result; }; return merge_ub::MergeAllocations(std::move(stmt), align_bytes, verbose, - false, planner); + false, planner, placed_scopes); } } // namespace diff --git a/src/ascend/transform/merge_ub_common.h b/src/ascend/transform/merge_ub_common.h index 310220df6c..c24969c168 100644 --- a/src/ascend/transform/merge_ub_common.h +++ b/src/ascend/transform/merge_ub_common.h @@ -36,6 +36,7 @@ #include #include #include +#include #include #include @@ -292,7 +293,7 @@ class SharedMemoryRewriter : public tirx::StmtExprMutator { } if (sequence.empty()) return tirx::Evaluate(0); - return tirx::SeqStmt(std::move(sequence)); + return tirx::SeqStmt::Flatten(std::move(sequence)); } tirx::Stmt VisitStmt_(const tirx::AllocBufferNode *op) final { @@ -410,15 +411,18 @@ class SharedMemoryRewriter : public tirx::StmtExprMutator { }; template -tirx::Stmt MergeAllocations(tirx::Stmt stmt, int align_bytes, bool verbose, - bool disable_reuse, const Planner &planner) { +tirx::Stmt +MergeAllocations(tirx::Stmt stmt, int align_bytes, bool verbose, + bool disable_reuse, const Planner &planner, + const std::unordered_set &required_scopes = {}) { AllocateCollector collector; collector(stmt); auto run = [&](SharedMemoryRewriter &rewriter, const AllocationMap &allocations, const AllocationOrder &allocation_order, const std::string &target_scope) { - if (disable_reuse || target_scope == "shared.ss") { + if (disable_reuse || + (target_scope == "shared.ss" && !required_scopes.count(target_scope))) { rewriter.PlanNoReuse(); } else { rewriter.SetPlan(planner(stmt, allocations, allocation_order, @@ -427,7 +431,8 @@ tirx::Stmt MergeAllocations(tirx::Stmt stmt, int align_bytes, bool verbose, stmt = rewriter(std::move(stmt)); }; - if (collector.dyn_shmem_allocs.size() > 1) { + if (collector.dyn_shmem_allocs.size() > 1 || + required_scopes.count("shared.dyn")) { SharedMemoryRewriter rewriter(collector.dyn_shmem_allocs, verbose, align_bytes); run(rewriter, collector.dyn_shmem_allocs, collector.dyn_shmem_order, ""); @@ -440,7 +445,7 @@ tirx::Stmt MergeAllocations(tirx::Stmt stmt, int align_bytes, bool verbose, std::sort(scopes.begin(), scopes.end()); for (const std::string &scope : scopes) { const AllocationMap &allocations = collector.ascend_scope_allocs.at(scope); - if (allocations.size() <= 1) + if (allocations.size() <= 1 && !required_scopes.count(scope)) continue; auto order = collector.ascend_scope_order.find(scope); const AllocationOrder empty_order; diff --git a/testing/ascend/auto_schedule/test_buffer_version_annotations.py b/testing/ascend/auto_schedule/test_buffer_version_annotations.py index baf3357a8e..95b3393857 100644 --- a/testing/ascend/auto_schedule/test_buffer_version_annotations.py +++ b/testing/ascend/auto_schedule/test_buffer_version_annotations.py @@ -67,5 +67,44 @@ def main(): T.annotate_buffer_versions({ub: first, alias: second}) +@pytest.mark.parametrize("specification", [None, "auto", "iteration", "counter", "manual"]) +def test_placement_requires_an_explicit_version_count(specification): + from tilelang import tvm + from testing.ascend._ir import schedule_snapshots + + @T.prim_func + def main(A: T.Tensor((4, 64), "float32"), C: T.Tensor((4, 64), "float32")): + with T.Kernel(1): + ub = T.alloc_shared((64,), "float32") + T.annotate_buffer_offsets({ub: 32}) + if specification == "manual": + T.annotate_manual_multi_buffer({ub: 1}) + elif specification is not None: + T.annotate_buffer_versions({ub: specification}) + for i in T.Pipelined(4, num_stages=2): + T.copy(A[i, :], ub) + T.copy(ub, C[i, :]) + + with pytest.raises(tvm.error.InternalError, match="explicit version count.*annotate_buffer_versions"): + schedule_snapshots(main, ["tl.AutoSchedule"]) + + +def test_placement_rejects_unapplied_fixed_versions(): + from tilelang import tvm + from testing.ascend._ir import schedule_snapshots + + @T.prim_func + def main(A: T.Tensor((64,), "float32"), C: T.Tensor((64,), "float32")): + with T.Kernel(1): + ub = T.alloc_shared((64,), "float32") + T.annotate_buffer_versions({ub: 2}) + T.annotate_buffer_offsets({ub: 32}) + T.copy(A, ub) + T.copy(ub, C) + + with pytest.raises(tvm.error.InternalError, match="requested 2 versions.*selected 1"): + schedule_snapshots(main, ["tl.AutoSchedule"]) + + if __name__ == "__main__": tilelang.testing.main() diff --git a/testing/ascend/auto_schedule/test_schedule_integration.py b/testing/ascend/auto_schedule/test_schedule_integration.py index 141a51872b..5bbdf20632 100644 --- a/testing/ascend/auto_schedule/test_schedule_integration.py +++ b/testing/ascend/auto_schedule/test_schedule_integration.py @@ -834,3 +834,160 @@ def main(A: T.Tensor((16,), "float32"), O: T.Tensor((16,), "float32")): a = torch.arange(16, dtype=torch.float32) result = tilelang.compile(main, target="ascend", out_idx=-1)(a.npu()) torch.testing.assert_close(result.cpu(), 2 * (a.cumsum(0) - a), rtol=0, atol=0) + + +def _placed_versioned_kernel(b_offset, unlimit=False, rows=16): + @T.prim_func + def main(A: T.Tensor((4, rows, 16), "bfloat16"), B: T.Tensor((4, 16, 16), "bfloat16"), C: T.Tensor((4, rows, 16), "float32")): + with T.Kernel(1): + a = T.alloc_l1((rows, 16), "bfloat16") + b = T.alloc_l1((16, 16), "bfloat16") + a_view = T.view(a, shape=(rows, 16)) + aq = T.alloc_l0a((rows, 16), "bfloat16") + bk = T.alloc_l0b((16, 16), "bfloat16") + c = T.alloc_l0c((rows, 16), "float32") + T.annotate_buffer_versions({a: 2, b: 2, aq: 2, bk: 2, c: 2}) + T.annotate_buffer_offsets({a: 0, a_view: 0, b: b_offset, aq: 512, bk: 1024, c: 2048}) + if unlimit: + T.annotate_unlimit_memory("shared.l1") + for i in T.Pipelined(4, num_stages=2, annotations={"enable_offset": True}): + with T.Stage(0), T.Task(): + T.copy(A[i, :, :], a) + T.copy(B[i, :, :], b) + with T.Stage(1): + T.copy(a_view, aq) + T.copy(b, bk) + T.gemm(aq, bk, c, transpose_B=True, clear_accum=True) + T.copy(c, C[i, :, :]) + + return main + + +def test_pipeline_preserves_offsets_across_versions_layout_remap_and_flatten(): + snapshots = {} + + @tvm.ir.instrument.pass_instrument + class Capture: + def run_after_pass(self, mod, info): + if info.name in ("tl.MaterializeMultiBuffer", "tl.LowerScheduledTIR", "tl.AscendLowerTileOp", "tl.MergeUBAllocations"): + snapshots[info.name] = mod + + with tvm.transform.PassContext(instruments=[Capture()]): + artifact = tilelang.lower(_placed_versioned_kernel(262144), target="ascend") + physical = snapshots["tl.MaterializeMultiBuffer"] + buffers = [buffer for block in nodes(physical, tirx.SBlock) for buffer in block.alloc_buffers if buffer.scope() == "shared.l1"] + assert len(buffers) == 2 and all(int(buffer.shape[0]) == 2 for buffer in buffers) + for stage in ("tl.LowerScheduledTIR", "tl.AscendLowerTileOp"): + mod = snapshots[stage] + (placement,) = [node for node in nodes(mod, tirx.AttrStmt) if node.attr_key == "tl.ascend_buffer_offsets"] + storages = [ + buffer.data + for block in nodes(mod, tirx.SBlock) + for buffer in block.alloc_buffers + if buffer.scope() in ("shared.l1", "shared.l0a", "shared.l0b", "shared.l0c") + ] + assert len(placement.node) == 5 + assert all(any(key.same_as(storage) for storage in storages) for key in placement.node) + assert sorted(int(offset) for offset in placement.node.values()) == [0, 512, 1024, 2048, 262144] + merged = snapshots["tl.MergeUBAllocations"] + allocations = [node for func in merged.functions.values() for node in nodes(func, tirx.AllocBuffer)] + (allocation,) = [node for node in allocations if node.buffer.scope() == "shared.l1.dyn"] + assert int(allocation.buffer.shape[0]) == 262144 + 1024 + for scope, expected_bytes in [("shared.l0a.dyn", 1536), ("shared.l0b.dyn", 2048), ("shared.l0c.dyn", 4096)]: + (allocation,) = [node for node in allocations if node.buffer.scope() == scope] + assert int(allocation.buffer.shape[0]) == expected_bytes + assert not any(node.attr_key == "tl.ascend_buffer_offsets" for func in merged.functions.values() for node in nodes(func, tirx.AttrStmt)) + assert artifact.kernel_source and artifact.device_mod + + +def test_overlap_is_checked_after_automatic_versions_expand(): + # Each logical tile is only 512 bytes, but two physical versions occupy 1024. + with pytest.raises(tvm.error.InternalError, match="Buffer placements overlap"): + tilelang.lower(_placed_versioned_kernel(512), target="ascend") + + +def test_fractal_padding_is_included_in_the_physical_footprint(): + # Two 17x16 BF16 logical tiles take 1088 bytes, but padded 32x16 tiles + # take 2048. An offset of 1536 therefore overlaps only after padding. + with pytest.raises(tvm.error.InternalError, match="Buffer placements overlap"): + tilelang.lower(_placed_versioned_kernel(1536, rows=17), target="ascend") + + +@pytest.mark.parametrize("unlimit", [False, True]) +def test_holes_and_versions_cannot_exceed_physical_capacity(unlimit): + # The original tile fits; the final second version crosses the L1 boundary. + with pytest.raises(tvm.error.InternalError, match="capacity of 524288 bytes"): + tilelang.lower(_placed_versioned_kernel(524288 - 512, unlimit), target="ascend") + + +def test_unscheduled_pipeline_rejects_placement_instead_of_dropping_it(): + with ( + pytest.raises(tvm.error.InternalError, match="requires the Ascend scheduled pipeline"), + tvm.transform.PassContext(config={tilelang.PassConfigKey.TL_ENABLE_AUTO_SCHEDULE: False}), + ): + tilelang.lower(_placed_versioned_kernel(262144), target="ascend") + + +def _placed_ub_kernel(offset=64, versions=2, simt=False): + @T.prim_func + def main(A: T.Tensor((4, 64), "float32"), C: T.Tensor((4, 64), "float32")): + with T.Kernel(1): + ub = T.alloc_shared((64,), "float32") + alias = T.view(ub, shape=(64,)) + # Placement may precede the version declaration and use an alias. + T.annotate_buffer_offsets({alias: offset}) + if versions == 1: + T.annotate_buffer_versions({ub: 1}) + else: + T.annotate_buffer_versions({ub: (versions, "counter")}) + for i in T.Pipelined(4, num_stages=2): + T.copy(A[i, :], ub) + if simt: + with T.SimtVF(threads=32): + for j in T.Parallel(64): + ub[j] = ub[j] + 1.0 + T.copy(alias, C[i, :]) + + return main + + +@pytest.mark.parametrize("versions,simt", [(1, False), (2, False), (2, True)]) +def test_ub_placement_survives_version_and_vf_lowering(versions, simt): + with tvm.target.Target("ascend"): + artifact = tilelang.lower(_placed_ub_kernel(versions=versions, simt=simt), target="ascend") + allocations = [node for func in artifact.device_mod.functions.values() for node in nodes(func, tirx.AllocBuffer)] + (allocation,) = [node for node in allocations if node.buffer.scope() == "shared.dyn"] + assert int(allocation.buffer.shape[0]) == 64 + versions * 256 + assert "__ubuf__" in artifact.kernel_source + + +def test_simt_context_limits_placement_after_full_lowering(): + with pytest.raises(tvm.error.InternalError, match="capacity of 221184 bytes"), tvm.target.Target("ascend"): + tilelang.lower(_placed_ub_kernel(offset=216 * 1024 - 256, simt=True), target="ascend") + + +def _placed_ss_kernel(): + @T.prim_func + def main(seed: T.Tensor((16, 16), "float16"), out: T.Tensor((4, 2), "uint32")): + with T.MixedKernel(1) as (_, sid): + words = T.alloc_ss((2,), "uint32") + cube = T.alloc_l1((16, 16), "float16") + T.annotate_buffer_versions({words: 2}) + T.annotate_buffer_offsets({words: 96}) + T.copy(seed, cube) + for i in T.Pipelined(4, num_stages=2, annotations={"enable_offset": True}): + with T.Stage(0): + words[sid] = T.Cast("uint32", (i + 1) * (100 + sid)) + with T.Stage(1): + out[i, 0] = words[0] + out[i, 1] = words[1] + + return main + + +def test_ss_placement_survives_cross_core_version_lowering(): + artifact = tilelang.lower(_placed_ss_kernel(), target="ascend") + allocations = [node for func in artifact.device_mod.functions.values() for node in nodes(func, tirx.AllocBuffer)] + (allocation,) = [node for node in allocations if node.buffer.scope() == "shared.ss.dyn"] + assert int(allocation.buffer.shape[0]) == 128 # 96-byte offset + two 8-byte versions, rounded to 32. + assert "volatile __ssbuf__ uint8_t" in artifact.kernel_source diff --git a/testing/ascend/transform/test_ascend_buffer_offsets.py b/testing/ascend/transform/test_ascend_buffer_offsets.py new file mode 100644 index 0000000000..8185b85302 --- /dev/null +++ b/testing/ascend/transform/test_ascend_buffer_offsets.py @@ -0,0 +1,194 @@ +"""Exact placement of physical on-chip allocations after version/layout lowering.""" + +import pytest + +import tilelang.ascend.language as T +from tilelang import tvm +from tilelang.ascend import transform +from tvm import tirx +from testing.ascend._ir import nodes, seq + + +SCOPES = [ + pytest.param("shared.dyn", 32, 248 * 1024, id="ub"), + pytest.param("shared.l1", 512, 512 * 1024, id="l1"), + pytest.param("shared.l0a", 512, 64 * 1024, id="l0a"), + pytest.param("shared.l0b", 512, 64 * 1024, id="l0b"), + pytest.param("shared.l0c", 1024, 256 * 1024, id="l0c"), + pytest.param("shared.ss", 32, 3 * 1024, id="ss"), +] + + +def _allocation_program(buffers, offsets=None, *, aliases=True, aiv_storages=None, simt=False): + # Input to the allocator already has physical versions and flattened sizes. + # Frontend version-count validation belongs to AutoSchedule. + stores = seq(*[tirx.BufferStore(buffer, tirx.const(index, buffer.dtype), [0]) for index, buffer in enumerate(buffers)]) + if simt: + stores = tirx.SBlockRealize([], True, tirx.SBlock([], [], [], "SIMT_VF", stores)) + body = tirx.SeqStmt([*[tirx.AllocBuffer(buffer) for buffer in buffers], stores]) + if aliases: + # Even an explicit positive alias permission cannot authorize overlap. + body = tirx.AttrStmt( + {buffer.data: [other.data for other in buffers if other != buffer] for buffer in buffers}, "tl.buffer_alias_map", 1, body + ) + if offsets is not None: + body = tirx.AttrStmt(offsets, "tl.ascend_buffer_offsets", 1, body) + if aiv_storages is not None: + body = tirx.AttrStmt(aiv_storages, "tl.aiv_storage_map", 1, body) + thread = tirx.IterVar(None, tirx.Var("bx", "int32"), tirx.IterVar.ThreadIndex, "blockIdx.x") + return tvm.IRModule({"main": tirx.PrimFunc([], tirx.AttrStmt(thread, "thread_extent", 1, body))}) + + +def _buffer(name, size=512, scope="shared.l1"): + return tirx.decl_buffer((size // 4,), "uint32", name=name, scope=scope) + + +def _store_offsets(after): + return [int(tvm.arith.Analyzer().simplify(store.indices[0])) * 4 for store in nodes(after, tirx.BufferStore)] + + +@pytest.mark.parametrize("scope,alignment,capacity", SCOPES) +@pytest.mark.parametrize("disable_reuse", [False, True]) +def test_exact_plan_keeps_holes_and_physical_offsets(scope, alignment, capacity, disable_reuse): + a, b = _buffer("a", alignment, scope), _buffer("b", 2 * alignment, scope) + before = _allocation_program([a, b], {a.data: 4 * alignment, b.data: 0}) + after = transform.MergeUBAllocations(disable_reuse=disable_reuse)(before) + assert _store_offsets(after) == [4 * alignment, 0] + (allocation,) = nodes(after, tirx.AllocBuffer) + assert allocation.buffer.dtype == "uint8" + assert int(allocation.buffer.shape[0]) == 5 * alignment + assert not any(node.attr_key == "tl.ascend_buffer_offsets" for node in nodes(after, tirx.AttrStmt)) + + +@pytest.mark.parametrize("scope,alignment,capacity", SCOPES) +def test_single_allocation_can_end_at_capacity(scope, alignment, capacity): + a = _buffer("a", alignment, scope) + after = transform.MergeUBAllocations()(_allocation_program([a], {a.data: capacity - alignment})) + (allocation,) = nodes(after, tirx.AllocBuffer) + assert int(allocation.buffer.shape[0]) == capacity + assert _store_offsets(after) == [capacity - alignment] + + +@pytest.mark.parametrize("scope,alignment,capacity", SCOPES) +def test_partial_plan_first_fits_unplaced_buffers_without_reuse(scope, alignment, capacity): + large, fixed, small = [_buffer(name, size * alignment, scope) for name, size in [("large", 2), ("fixed", 1), ("small", 1)]] + after = transform.MergeUBAllocations()(_allocation_program([large, fixed, small], {fixed.data: alignment})) + assert _store_offsets(after) == [2 * alignment, alignment, 0] + (allocation,) = nodes(after, tirx.AllocBuffer) + assert int(allocation.buffer.shape[0]) == 4 * alignment + + +@pytest.mark.parametrize("scope,alignment,capacity", SCOPES) +@pytest.mark.parametrize("case", ["overlap", "capacity", "unaligned"]) +def test_scope_specific_physical_constraints(scope, alignment, capacity, case): + a, b = _buffer("a", 2 * alignment, scope), _buffer("b", alignment, scope) + offsets = {a.data: 0, b.data: alignment} + message = "overlap" + if case == "capacity": + offsets = {a.data: capacity - alignment} + message = "capacity" + elif case == "unaligned": + offsets = {a.data: alignment // 2} + message = "aligned integer offsets" + with pytest.raises(tvm.error.InternalError, match=message): + transform.MergeUBAllocations()(_allocation_program([a, b], offsets)) + + +@pytest.mark.parametrize( + "case,message", + [ + ("unknown", "unknown allocation"), + ("scope", "does not support storage scope"), + ("negative", "aligned integer offsets"), + ("boolean", "integer offsets"), + ("overflow", "capacity"), + ("dynamic", "positive constant"), + ("empty", "must not be empty"), + ("duplicate", "one allocation per storage"), + ("no-contract", "scheduled allocation"), + ], +) +def test_invalid_physical_plans_fail(case, message): + a, b = _buffer("a"), _buffer("b") + if case == "dynamic": + a = tirx.decl_buffer((tirx.Var("n", "int32"),), "uint32", name="a", scope="shared.l1") + elif case == "overflow": + a = _buffer("a", 2**62) + elif case == "scope": + a = _buffer("a", scope="local") + offsets = {a.data: 0, b.data: 1024} + if case == "unknown": + offsets[_buffer("unknown").data] = 2048 + elif case == "negative": + offsets[a.data] = -512 + elif case == "boolean": + offsets[a.data] = tirx.const(True, "bool") + elif case == "empty": + offsets = {} + buffers = [a, b, a] if case == "duplicate" else [a, b] + with pytest.raises(tvm.error.InternalError, match=message): + transform.MergeUBAllocations()(_allocation_program(buffers, offsets, aliases=case != "no-contract")) + + +def test_existing_dynamic_arena_is_rejected(): + a = _buffer("a") + arena = _buffer("arena", scope="shared.l1.dyn") + with pytest.raises(tvm.error.InternalError, match="existing shared.l1.dyn arena"): + transform.MergeUBAllocations()(_allocation_program([a, arena], {a.data: 0})) + + +@pytest.mark.parametrize("simt", [False, True]) +def test_ub_placement_accounts_for_simt_context(simt): + a = _buffer("a", 64, "shared.dyn") + before = _allocation_program([a], {a.data: 216 * 1024 - 32}, simt=simt) + if simt: + with pytest.raises(tvm.error.InternalError, match="capacity of 221184 bytes"): + transform.MergeUBAllocations()(before) + else: + assert _store_offsets(transform.MergeUBAllocations()(before)) == [216 * 1024 - 32] + + +@pytest.mark.parametrize("conflict", [None, "same-aiv", "common"]) +def test_ub_overlap_uses_physical_aiv_address_spaces(conflict): + common, v0, v1 = [_buffer(name, 32, "shared.dyn") for name in ("common", "v0", "v1")] + offsets = {v0.data: 0, v1.data: 0} + if conflict == "common": + offsets[common.data] = 0 + before = _allocation_program([common, v0, v1], offsets, aiv_storages={v0.data: 0, v1.data: 0 if conflict == "same-aiv" else 1}) + if conflict: + with pytest.raises(tvm.error.InternalError, match="overlap"): + transform.MergeUBAllocations()(before) + else: + assert _store_offsets(transform.MergeUBAllocations()(before)) == [32, 0, 0] + + +def test_different_pools_use_independent_offsets(): + buffers = [_buffer(str(i), alignment, scope) for i, (scope, alignment, _) in enumerate(item.values for item in SCOPES)] + before = _allocation_program(buffers, {buffer.data: 1024 for buffer in buffers}) + after = transform.MergeUBAllocations()(before) + assert _store_offsets(after) == [1024] * len(buffers) + assert len(nodes(after, tirx.AllocBuffer)) == len(buffers) + + +@pytest.mark.parametrize("offset", [True, -512, 1, 1.5, 524288, 2**64]) +def test_frontend_rejects_invalid_offsets(offset): + with pytest.raises(ValueError, match="byte offset"): + T.annotate_buffer_offsets({_buffer("a"): offset}) + + +@pytest.mark.parametrize("scope", ["shared", "global", "local.fragment", "shared.l0a.sf"]) +def test_frontend_rejects_scopes_without_independent_placement(scope): + with pytest.raises(ValueError, match="does not support storage scope"): + T.annotate_buffer_offsets({_buffer("a", scope=scope): 0}) + + +def test_frontend_aliases_share_one_storage_offset(): + a = _buffer("a") + view = tirx.decl_buffer((512,), "uint8", data=a.data, name="view", scope="shared.l1") + with pytest.raises(ValueError, match="same byte offset"): + T.annotate_buffer_offsets({a: 0, view: 512}) + + +def test_absent_placement_preserves_existing_no_contract_behavior(): + before = _allocation_program([_buffer("a")], aliases=False) + tvm.ir.assert_structural_equal(before, transform.MergeUBAllocations()(before)) diff --git a/tilelang/ascend/language/__init__.py b/tilelang/ascend/language/__init__.py index 423d829d51..2e8e7c5bc7 100644 --- a/tilelang/ascend/language/__init__.py +++ b/tilelang/ascend/language/__init__.py @@ -19,6 +19,7 @@ ) from .annotations import ( # noqa: F401 annotate_buffer_versions, + annotate_buffer_offsets, annotate_manual_multi_buffer, annotate_unlimit_memory, ) @@ -111,6 +112,7 @@ "alloc_l1", "alloc_ss", "annotate_buffer_versions", + "annotate_buffer_offsets", "annotate_manual_multi_buffer", "annotate_unlimit_memory", "assume_conflict", diff --git a/tilelang/ascend/language/annotations.py b/tilelang/ascend/language/annotations.py index d53216bb86..72946cb0ff 100644 --- a/tilelang/ascend/language/annotations.py +++ b/tilelang/ascend/language/annotations.py @@ -1,5 +1,4 @@ -"""Ascend auto-schedule annotations: multi-buffer version control and memory -capacity overrides consumed by the Z3 scheduler.""" +"""Ascend storage placement, multi-buffer versions and scheduling capacity annotations.""" from tvm.tirx import IntImm from tvm.tirx.script.builder.ir import sblock_attr @@ -8,6 +7,7 @@ __all__ = [ "annotate_buffer_versions", + "annotate_buffer_offsets", "annotate_manual_multi_buffer", "annotate_unlimit_memory", ] @@ -16,6 +16,68 @@ _VALID_BUFFER_VERSION_MODES = ("auto", "iteration", "counter") +# Physical pools on dav-3510: (base alignment, capacity), in bytes. +_BUFFER_PLACEMENT_SCOPES = { + "shared.dyn": (32, 248 * 1024), + "shared.l1": (512, 512 * 1024), + "shared.l0a": (512, 64 * 1024), + "shared.l0b": (512, 64 * 1024), + "shared.l0c": (1024, 256 * 1024), + "shared.ss": (32, 3 * 1024), +} + + +def annotate_buffer_offsets(buffer_offsets: dict): + """Place on-chip buffers at exact byte offsets within their memory pools. + + Supports UB, L1, L0A, L0B, L0C and SS in the Ascend scheduled pipeline. + Each placed storage must also have an explicit positive version count in + :func:`annotate_buffer_versions`, including ``1`` for unversioned storage. + A mode without a count, inferred versions, and manual multi-buffer metadata + alone do not satisfy this requirement. Annotation order does not matter. + + Offsets are non-negative Python integers, aligned to 32 bytes for UB/SS, + 512 for L1/L0A/L0B, and 1024 for L0C. A placement includes all physical + versions and layout padding; views of one storage share the same base. + Bound L0 scale-factor handles follow their data buffer's placement. + + Partial maps are allowed. Other allocations in a placed pool use the first + aligned free interval, without storage reuse. Overlap is rejected even for + disjoint lifetimes or positive conflict hints. Different pools and private + AIV UBs have independent addresses; common UB storage reserves both UBs. + Placement does not change synchronization or ownership. + + Every allocation in a placed pool must have a positive static physical size + and fit its hardware capacity, including the 32 KiB UB reservation for SIMT. + Pools without placements retain their ordinary allocation behavior. + + Example + ------- + >>> T.annotate_buffer_versions({left: 2, accum: 1}) + >>> T.annotate_buffer_offsets({left: 262144, accum: 1024}) + """ + if not buffer_offsets: + raise ValueError("annotate_buffer_offsets requires a non-empty placement map") + offsets = {} + for buffer, offset in buffer_offsets.items(): + scope = buffer.scope() + if scope not in _BUFFER_PLACEMENT_SCOPES: + raise ValueError( + f"annotate_buffer_offsets does not support storage scope {scope!r}; " + "expected UB, L1, L0A, L0B, L0C or SS. Place bound L0 scale-factor handles through their data buffer." + ) + alignment, capacity = _BUFFER_PLACEMENT_SCOPES[scope] + if isinstance(offset, bool) or not isinstance(offset, int) or offset < 0 or offset % alignment: + raise ValueError(f"{scope} byte offset must be a non-negative integer divisible by {alignment}, got {offset!r}") + if offset >= capacity: + raise ValueError(f"{scope} byte offset must be below the {capacity}-byte capacity, got {offset}") + data = buffer.data + if data in offsets and int(offsets[data]) != offset: + raise ValueError(f"Aliases of one storage require the same byte offset: {buffer}") + offsets[data] = IntImm("int64", offset) + return sblock_attr({"tl.ascend_buffer_offsets": offsets}) + + def annotate_buffer_versions(buffer_versions_map: dict): """Control automatic multi-buffer version counts and indexing modes.