From 0f79092d8da3f5d237f06cdefce816d7fa8a3aaa Mon Sep 17 00:00:00 2001 From: Denver Jin Date: Tue, 20 Oct 2026 11:25:59 +0800 Subject: [PATCH 1/7] [Ascend] Allow storage reuse for explicit conflict hints --- .../references/synchronization.md | 8 ++ .../auto_schedule/buffer_alias_analysis.h | 34 +++++++- src/ascend/transform/buffer_alias.h | 3 +- src/ascend/transform/insert_sync.cc | 2 +- .../auto_schedule/test_conflict_hints.py | 43 +++++++++- .../auto_schedule/test_reuse_integration.py | 78 ++++++++++++++++++- tilelang/ascend/language/schedule_hint.py | 11 ++- 7 files changed, 171 insertions(+), 8 deletions(-) diff --git a/.agents/skills/tilelang-ascend/references/synchronization.md b/.agents/skills/tilelang-ascend/references/synchronization.md index b8d1cc51e9..ced59faa27 100644 --- a/.agents/skills/tilelang-ascend/references/synchronization.md +++ b/.agents/skills/tilelang-ascend/references/synchronization.md @@ -326,6 +326,14 @@ to be treated as conflicting while retaining the usual RAW/WAR/WAW rules; read-only pairs still have no dependency. If contradictory hints match, conflict wins. See `tilelang/ascend/language/schedule_hint.py` for the API. +A positive conflict hint also explicitly permits the two backing storages to +reuse memory, even when lifetime inference cannot prove that reuse. This +permission is pairwise, with no transitive closure. The region, `level`, and +`cross` qualifiers continue to govern dependency matching; reuse permission is +for the storage pair. The caller must cover all accesses affected by reuse. +Allocation scope, alignment, and the shared-memory-reuse disable option still +apply. Buffer identities, layouts, and version counters remain independent. + ## Trace the complete storage protocol For a suspicious access or wait, build a table with: diff --git a/src/ascend/transform/auto_schedule/buffer_alias_analysis.h b/src/ascend/transform/auto_schedule/buffer_alias_analysis.h index bb7e8c987d..e5b38fa877 100644 --- a/src/ascend/transform/auto_schedule/buffer_alias_analysis.h +++ b/src/ascend/transform/auto_schedule/buffer_alias_analysis.h @@ -41,6 +41,7 @@ #include "../buffer_alias.h" #include "./buffer_lifetime_analysis.h" +#include "./dependency_analysis.h" #include "./ir_structure.h" #include "./memory_detector.h" #include "ascend/op/utils.h" @@ -406,15 +407,37 @@ class BufferAliasAnalyzerImpl { } } - void Analyze(const std::vector> &root) { + void Analyze(const std::vector> &root, + const ConflictHintList &root_conflicts) { lifetime_summaries_.clear(); AnalyzeScope(root, /*loop=*/nullptr); + AddConflictAliases_(root_conflicts); } const BufferAliasMap &GetAliases() const { return buffer_aliases_; } void Validate() const { ValidateBufferAliasMap(buffer_aliases_); } private: + void AddConflictAliases_(const ConflictHintList &hints) { + auto storage = [](const ffi::Any &operand) -> Var { + if (auto var = operand.as()) + return var.value(); + return operand.cast()->buffer->data; + }; + for (const ffi::Any &entry : hints) { + Array hint = entry.cast>(); + ICHECK_EQ(hint.size(), 4U); + if (!hint[3].cast()->value) + continue; + Var lhs = groups_.Representative(storage(hint[0])); + Var rhs = groups_.Representative(storage(hint[1])); + // A positive hint explicitly permits storage reuse. Its region and + // iteration qualifiers still control dependency analysis independently. + if (!lhs.same_as(rhs)) + AddBufferAlias(&buffer_aliases_, lhs, rhs); + } + } + HappensBeforeQuery MakeHappensBeforeQuery(ControlNode *scope) const { return [this, scope](size_t src, size_t dst, int64_t distance) { return context_->HappensBefore(scope, src, dst, distance); @@ -429,6 +452,10 @@ class BufferAliasAnalyzerImpl { void AnalyzeScope(const std::vector> &nodes, ControlNode *loop) { + if (loop != nullptr) { + if (auto hints = loop->control->annotations.Get("conflict_hint")) + AddConflictAliases_(hints.value().cast()); + } // Cache every child scope before summarizing this scope. A parent can then // promote the child's already-built lifetime instead of recursively // rediscovering the same nested control chain. @@ -1500,8 +1527,9 @@ class BufferAliasAnalyzer { void Collect(const std::vector> &root) { impl_.Collect(root); } - void Analyze(const std::vector> &root) { - impl_.Analyze(root); + void Analyze(const std::vector> &root, + const ConflictHintList &root_conflicts = {}) { + impl_.Analyze(root, root_conflicts); } const BufferAliasMap &GetAliases() const { return impl_.GetAliases(); } diff --git a/src/ascend/transform/buffer_alias.h b/src/ascend/transform/buffer_alias.h index cbeb34b731..768db66adf 100644 --- a/src/ascend/transform/buffer_alias.h +++ b/src/ascend/transform/buffer_alias.h @@ -47,7 +47,8 @@ constexpr const char *kAivStorageMap = "tl.aiv_storage_map"; constexpr const char *kSplitSubcores = "tl.split_subcores"; using AivStorageMap = ffi::Map; -/*! \brief Storage data Var -> storages proven safe to alias with it. */ +/*! \brief Storage data Var -> inferred or explicitly permitted alias partners. + */ using BufferAliasMap = ffi::Map>; inline void AddBufferAliasStorage(BufferAliasMap *aliases, diff --git a/src/ascend/transform/insert_sync.cc b/src/ascend/transform/insert_sync.cc index 1ded68e606..9dd4f78903 100644 --- a/src/ascend/transform/insert_sync.cc +++ b/src/ascend/transform/insert_sync.cc @@ -831,7 +831,7 @@ class SyncAnalyzer : public BufferAliasAnalysisContext { BufferAliasAnalyzer buffer_alias_analyzer(*this, storage_groups_); buffer_alias_analyzer.Collect(root); - buffer_alias_analyzer.Analyze(root); + buffer_alias_analyzer.Analyze(root, root_conflicts_); buffer_alias_analyzer.Validate(); buffer_aliases_ = buffer_alias_analyzer.GetAliases(); } diff --git a/testing/ascend/auto_schedule/test_conflict_hints.py b/testing/ascend/auto_schedule/test_conflict_hints.py index 5a12b8d7d8..4ec1032f82 100644 --- a/testing/ascend/auto_schedule/test_conflict_hints.py +++ b/testing/ascend/auto_schedule/test_conflict_hints.py @@ -6,7 +6,8 @@ from tilelang import tvm from tilelang.ascend import transform from tvm import tirx -from testing.ascend._ir import calls, schedule, nodes +from testing.ascend._ir import calls, kernel, schedule, nodes, seq +from testing.ascend.auto_schedule._scheduled_ir import unit def _insert_sync(program): @@ -18,6 +19,46 @@ def _event_count(mod, operation, pipe): return sum(call.args[0].value == pipe for call in calls(mod, "tl.ascend_" + operation + "_flag")) +@pytest.mark.parametrize("scope,cross", [("root", False), ("loop", False), ("loop", True), ("nested", None)]) +@pytest.mark.parametrize("regions", [False, True], ids=["storage", "region"]) +@pytest.mark.parametrize("disable_reuse", [False, True]) +def test_conflict_hints_permit_pairwise_storage_reuse(scope, cross, regions, disable_reuse): + a, b, c = [tirx.decl_buffer((8,), "int32", name=name, scope="shared.dyn") for name in ("a", "b", "c")] + out = tirx.decl_buffer((1,), "int32", name="out") + + def operand(buffer): + return tirx.BufferRegion(buffer, [tvm.ir.Range(0, 4)]) if regions else buffer.data + + def hint(left, right, positive=True): + return [operand(left), operand(right), tirx.IntImm("int32", -1 if cross is None else int(cross)), tirx.IntImm("bool", positive)] + + hints = [hint(a, b), hint(b, c), hint(a, b), hint(b, b), hint(a, c, False)] + # All three values stay live until the final read, so ordinary lifetime + # inference cannot authorize reuse. Equal values make the declared reuse + # valid without changing the result. + body = seq( + *(unit(tirx.BufferStore(buffer, 1, [0]), core=1) for buffer in (a, b, c)), + unit(tirx.BufferStore(out, a[0] + b[0] + c[0], [0]), core=1), + ) + annotations = {"tl.root_conflict_hints": hints} if scope == "root" else {} + if scope != "root": + loop = tirx.For(tirx.Var("i", "int32"), 0, 2, tirx.ForKind.SERIAL, body, annotations={"conflict_hint": hints}) + body = unit(loop, core=None) + if scope == "nested": + body = unit(tirx.For(tirx.Var("j", "int32"), 0, 2, tirx.ForKind.SERIAL, body), core=None) + before = kernel(body, buffers=[a, b, c], params=[out], annotations=annotations) + with tvm.transform.PassContext(config={"tl.disable_shared_memory_reuse": disable_reuse}): + after = transform.InsertSync()(before) + (aliases,) = [ + block.annotations["tl.buffer_alias_map"] for block in nodes(after, tirx.SBlock) if "tl.buffer_alias_map" in block.annotations + ] + expected = [] if disable_reuse else [(a.data, b.data), (b.data, a.data), (b.data, c.data), (c.data, b.data)] + actual = [(storage, other) for storage, partners in aliases.items() for other in partners] + assert len(actual) == len(expected) + for lhs, rhs in expected: + assert any(storage.same_as(lhs) and other.same_as(rhs) for storage, other in actual) + + def _hint_program(scope, exact, cross, explicit_root): @T.prim_func def main(A: T.Tensor((128,), "float32"), C: T.Tensor((128,), "float32")): diff --git a/testing/ascend/auto_schedule/test_reuse_integration.py b/testing/ascend/auto_schedule/test_reuse_integration.py index 1252b1ca49..dbfc79a965 100644 --- a/testing/ascend/auto_schedule/test_reuse_integration.py +++ b/testing/ascend/auto_schedule/test_reuse_integration.py @@ -1,13 +1,89 @@ """The reuse contract survives scheduled lowering and is consumed by packing.""" +import pytest import tilelang import tilelang.ascend.language as T from tilelang import tvm from tvm import tirx -from testing.ascend._ir import nodes +from testing.ascend._ir import calls, nodes from testing.ascend.auto_schedule._reuse_utils import _make_cyclic_scalar_program +@pytest.mark.parametrize("disable_reuse", [False, True]) +def test_conflict_hint_reuses_independent_l0_layouts(disable_reuse): + @T.prim_func + def main( + Q: T.Tensor((64, 128), "bfloat16"), + K: T.Tensor((64, 128), "bfloat16"), + CQ: T.Tensor((64, 64), "float32"), + CV: T.Tensor((64, 128), "float32"), + ): + with T.Kernel(1): + q = T.alloc_l1((64, 128), "bfloat16") + k0 = T.alloc_l1((32, 128), "bfloat16") + k1 = T.alloc_l1((32, 128), "bfloat16") + aq = T.alloc_l0a((64, 128), "bfloat16") + av = T.alloc_l0a((64, 64), "bfloat16") + bq = T.alloc_l0b((64, 128), "bfloat16") + bv = T.alloc_l0b((64, 128), "bfloat16") + cq = T.alloc_l0c((64, 64), "float32") + cv = T.alloc_l0c((64, 128), "float32") + T.annotate_buffer_versions({q: 1, k0: 1, k1: 1, aq: 2, av: 1, bq: 4, bv: 4, cq: 1, cv: 1}) + T.assume_conflict(bq, bv, level=-1, cross=False) + T.copy(Q, q) + T.copy(K[:32, :], k0) + T.copy(K[32:, :], k1) + T.set_mmad_direction("m") + for _d in T.Pipelined(4, num_stages=4): + with T.Stage(0): + T.copy(q, aq) + with T.Task(): + T.copy(k0, bq[:32, :], l0_pitch=4) + T.copy(k1, bq[32:, :], l0_pitch=4) + T.gemm(aq, bq, cq, transpose_B=True, clear_accum=True) + T.copy(cq, CQ) + T.copy(q[:, :64], av) + T.set_mmad_direction("n") + for _d in T.Pipelined(4, num_stages=4): + with T.Stage(0): + with T.Task(): + T.copy(k0, bv[:32, :], l0_pitch=8) + T.copy(k1, bv[32:, :], l0_pitch=8) + T.gemm(av, bv, cv, clear_accum=True) + T.copy(cv, CV) + + snapshots = {} + + @tvm.instrument.pass_instrument + class Capture: + def run_after_pass(self, mod, info): + if info.name in ("tl.AscendLowerTileOp", "tl.MergeUBAllocations"): + snapshots[info.name] = mod + + with ( + tvm.target.Target("ascend"), + tvm.transform.PassContext(config={"tl.disable_shared_memory_reuse": disable_reuse}, instruments=[Capture()]), + ): + tilelang.lower(main, target="ascend", enable_device_compile=False) + (allocation,) = [ + node.buffer + for func in snapshots["tl.MergeUBAllocations"].functions.values() + for node in nodes(func, tirx.AllocBuffer) + if node.buffer.scope() in ("shared.l0b", "shared.l0b.dyn") + ] + allocated_bytes = int(allocation.shape[0]) * tvm.DataType(allocation.dtype).bits // 8 + assert allocated_bytes == (131072 if disable_reuse else 65536) + + offsets = {0: [], 1: []} + analyzer = tvm.arith.Analyzer() + for call in calls(snapshots["tl.AscendLowerTileOp"], "tl.ascend_load_cbuf_to_cb"): + pointer = call.args[0] + assert pointer.op.name == "tirx.tvm_access_ptr" + offsets[int(call.args[-1])].append(int(analyzer.simplify(pointer.args[2] % 8192))) + assert sorted(offsets[0]) == [0, 512] + assert sorted(offsets[1]) == [0, 4096] + + def test_alias_contract_lifetime(): snapshots = {} wanted = { diff --git a/tilelang/ascend/language/schedule_hint.py b/tilelang/ascend/language/schedule_hint.py index 51e0fdb371..2c56067937 100644 --- a/tilelang/ascend/language/schedule_hint.py +++ b/tilelang/ascend/language/schedule_hint.py @@ -244,7 +244,7 @@ def assume_no_conflict(a, b=None, *, level=None, cross=None, group=None): def assume_conflict(a, b=None, *, level=None, cross=None, group=None): - """Assert that two buffer regions conflict for dependency analysis. + """Assert a dependency conflict and permit the two storages to be reused. This is the conservative counterpart of :func:`assume_no_conflict`. It makes the auto-scheduler treat matching accesses as if they used the same @@ -253,6 +253,15 @@ def assume_conflict(a, b=None, *, level=None, cross=None, group=None): accesses do not form a dependency. The marker is consumed before auto-scheduling and is a runtime no-op. + A positive hint also explicitly permits physical reuse of the two backing + storages, including when lifetime inference cannot establish that reuse. + The permission is pairwise and does not imply transitive compatibility. + Region, ``level``, and ``cross`` restrict the dependency matching; the reuse + permission applies to the whole storage pair. The caller must ensure that + the declared dependencies cover every access affected by that reuse. + Existing allocation scope, alignment, and reuse-disable rules still apply. + This does not merge buffer identities, layouts, or version counters. + Args: a: first region -- ``buffer[indices...]`` or a bare ``buffer`` (the whole buffer). A concrete region must exactly match the footprint From 9a99639f1805def9ebe3fb2bafbae4cc6ab49b81 Mon Sep 17 00:00:00 2001 From: Denver Jin Date: Tue, 20 Oct 2026 14:34:11 +0800 Subject: [PATCH 2/7] [Ascend] Share L0B ring ownership for explicit conflict hints Keep independent buffer layouts while an isolated compatible L0B pair uses one version clock and physical owner dependencies. Preserve the ordinary path for incompatible or nested owners. Validate retimed and unequal owner lengths, allocation reuse controls, and repeated device launches. --- .../references/synchronization.md | 6 +- .../annotate_multi_buffer_eligible.cc | 2 +- src/ascend/transform/auto_schedule.cc | 174 +++++++++++++++++- .../auto_schedule/buffer_alias_analysis.h | 8 +- .../auto_schedule/dependency_analysis.cc | 35 ++-- .../auto_schedule/dependency_analysis.h | 6 +- src/ascend/transform/buffer_version.h | 45 ++++- src/ascend/transform/insert_sync.cc | 10 +- .../transform/materialize_multi_buffer.cc | 4 +- src/ascend/transform/prepare_multi_buffer.cc | 2 +- src/ascend/transform/resolve_core.cc | 9 +- .../test_conflict_ring_owners.py | 169 +++++++++++++++++ tilelang/ascend/language/schedule_hint.py | 6 +- 13 files changed, 431 insertions(+), 45 deletions(-) create mode 100644 testing/ascend/auto_schedule/test_conflict_ring_owners.py diff --git a/.agents/skills/tilelang-ascend/references/synchronization.md b/.agents/skills/tilelang-ascend/references/synchronization.md index ced59faa27..6f25576827 100644 --- a/.agents/skills/tilelang-ascend/references/synchronization.md +++ b/.agents/skills/tilelang-ascend/references/synchronization.md @@ -332,7 +332,11 @@ permission is pairwise, with no transitive closure. The region, `level`, and `cross` qualifiers continue to govern dependency matching; reuse permission is for the storage pair. The caller must cover all accesses affected by reuse. Allocation scope, alignment, and the shared-memory-reuse disable option still -apply. Buffer identities, layouts, and version counters remain independent. +apply. Buffer identities and layouts remain independent. An isolated, +equal-sized L0B pair with matching fixed automatic version counts can share a +ring clock when a whole-buffer hint uses the default iteration matching and all +accesses belong to disjoint automatic owners. Explicit iteration indexing, manual +versions, strides, and offsets keep the ordinary dependency path. ## Trace the complete storage protocol diff --git a/src/ascend/transform/annotate_multi_buffer_eligible.cc b/src/ascend/transform/annotate_multi_buffer_eligible.cc index 4ce5a0f4b4..5cd44cef63 100644 --- a/src/ascend/transform/annotate_multi_buffer_eligible.cc +++ b/src/ascend/transform/annotate_multi_buffer_eligible.cc @@ -808,7 +808,7 @@ tvm::transform::Pass AnnotateMultiBufferEligible() { [](const TilelangKernelContext &context) { ScheduledTIR scheduled_tir = DecodeScheduledTIR(context.root, context.outer_ctx); - L0StorageGroups groups(CollectL0SFBindings(context.root)); + L0StorageGroups groups = CollectL0StorageGroups(context.root); auto &metadata = scheduled_tir.metadata; metadata.buffer_versions = ExpandL0StorageGroupValues( metadata.buffer_versions, groups, "version counts"); diff --git a/src/ascend/transform/auto_schedule.cc b/src/ascend/transform/auto_schedule.cc index 1687d6b936..2be86bd77a 100644 --- a/src/ascend/transform/auto_schedule.cc +++ b/src/ascend/transform/auto_schedule.cc @@ -40,6 +40,7 @@ #include #include +#include #include #include #include @@ -65,6 +66,7 @@ #include "ascend/op/utils.h" #include "ascend/transform/buffer_version.h" #include "backend/common/target_utils.h" +#include "layout/layout.h" #include "op/utils.h" #include "runtime/thread_storage_scope.h" #include "support/check.h" @@ -583,7 +585,7 @@ std::vector> ScheduleBuilder::Z3SchedulePython( // Buffer-region data dependencies auto deps = AnalyzeDependencies( nodes, /*loop=*/nullptr, manual_buffer_versions_, multi_buffer_owners_, - &dependency_cache_, root_conflict_hints_); + &dependency_cache_, root_conflict_hints_, storage_groups_); std::unordered_map node_idx; for (size_t i = 0; i < n; ++i) { node_idx[nodes[i]] = i; @@ -987,7 +989,7 @@ void ScheduleBuilder::Z3SchedulePythonLoop(ControlNode *ctrl, // Buffer-region data dependencies auto deps = AnalyzeDependencies(nodes, ctrl, manual_buffer_versions_, multi_buffer_owners_, &dependency_cache_, - root_conflict_hints_); + root_conflict_hints_, storage_groups_); std::unordered_map node_idx; for (size_t i = 0; i < n; ++i) { node_idx[nodes[i]] = i; @@ -1351,6 +1353,171 @@ void ScheduleBuilder::Z3SchedulePythonLoop(ControlNode *ctrl, // Pass orchestration. +// Select an isolated, equal-sized L0B pair whose complete automatic epochs +// can use one ring clock. Other conflict hints retain their dependency-only +// interpretation during scheduling. In particular, do not take a transitive +// closure of pairwise reuse permissions. +static void PlanExplicitL0BReuseRing(ScheduledTIR *scheduled_tir) { + auto &metadata = scheduled_tir->metadata; + if (metadata.root_conflict_hints.empty() || + metadata.kernel_root->annotations.count(kReuseRingBindings)) + return; + + class L0BCollector : public StmtExprVisitor { + public: + std::vector buffers; + + private: + std::unordered_set seen; + void Record(const Buffer &buffer) { + if (IsL0BBuffer(buffer) && seen.insert(buffer->data).second) + buffers.push_back(buffer); + } + void VisitStmt_(const SBlockNode *op) final { + for (const Buffer &buffer : op->alloc_buffers) + Record(buffer); + StmtExprVisitor::VisitStmt_(op); + } + void VisitStmt_(const AllocBufferNode *op) final { Record(op->buffer); } + } collector; + collector(metadata.kernel_root); + if (collector.buffers.size() != 2) + return; + const Buffer &a = collector.buffers[0]; + const Buffer &b = collector.buffers[1]; + if (a->dtype != b->dtype || !a->strides.empty() || !b->strides.empty() || + !is_zero(a->elem_offset) || !is_zero(b->elem_offset)) + return; + auto a_versions = metadata.buffer_versions.Get(a->data); + auto b_versions = metadata.buffer_versions.Get(b->data); + if (!a_versions.has_value() || !b_versions.has_value() || + a_versions.value() <= 1 || a_versions.value() != b_versions.value() || + metadata.manual_buffer_versions.count(a->data) || + metadata.manual_buffer_versions.count(b->data)) + return; + for (const Var &storage : {a->data, b->data}) { + if (auto mode = metadata.buffer_version_modes.Get(storage)) { + if (mode.value() == "iteration") + return; + } + } + + bool declared = false; + for (const ffi::Any &entry : metadata.root_conflict_hints) { + auto hint = entry.cast>(); + if (!hint[3].cast()->value || hint[2].cast()->value != -1) + continue; + auto lhs = hint[0].as(); + auto rhs = hint[1].as(); + if (lhs.has_value() && rhs.has_value() && + ((lhs.value().same_as(a->data) && rhs.value().same_as(b->data)) || + (lhs.value().same_as(b->data) && rhs.value().same_as(a->data)))) + declared = true; + } + if (!declared) + return; + + Map layouts; + if (auto value = + metadata.kernel_root->annotations.Get(tvm::tl::attr::kLayoutMap)) + layouts = value.value().cast>(); + auto storage_bytes = [&](const Buffer &buffer) -> int64_t { + Array shape = buffer->shape; + bool has_layout = false; + bool found_layout = false; + for (const auto &[candidate, layout] : layouts) { + if (candidate->data.same_as(buffer->data)) { + has_layout = true; + if (candidate->dtype == buffer->dtype && candidate->strides.empty() && + is_zero(candidate->elem_offset) && + ffi::StructuralEqual()(candidate->shape, buffer->shape)) { + shape = layout->OutputShape(); + found_layout = true; + break; + } + } + } + if (has_layout && !found_layout) + return 0; + PrimExpr bits = + IntImm(DataType::Int(64), buffer->dtype.bits() * buffer->dtype.lanes()); + for (const PrimExpr &extent : shape) + bits *= extent; + arith::Analyzer analyzer; + PrimExpr bytes = analyzer.Simplify(floordiv(bits + 7, 8)); + const auto *value = bytes.as(); + return value != nullptr ? value->value : 0; + }; + int64_t bytes = storage_bytes(a); + if (bytes <= 0 || bytes != storage_bytes(b)) + return; + + L0StorageGroups existing = CollectL0StorageGroups(metadata.kernel_root); + if (existing.Members(a->data).size() != 1 || + existing.Members(b->data).size() != 1) + return; + MultiBufferOwnerMap owners = + CollectMultiBufferOwners(scheduled_tir->tree, existing); + auto a_owners = owners.find(a->data); + auto b_owners = owners.find(b->data); + if (a_owners == owners.end() || b_owners == owners.end()) + return; + for (ControlNode *a_owner : a_owners->second) { + for (ControlNode *b_owner : b_owners->second) { + if (a_owner == b_owner || a_owner->IsWithin(b_owner) || + b_owner->IsWithin(a_owner)) + return; + } + } + std::vector tasks; + CollectAllTaskNodes(scheduled_tir->tree, tasks); + for (const Var &storage : {a->data, b->data}) { + auto found = owners.find(storage); + if (found == owners.end() || found->second.empty()) + return; + for (TaskNode *task : tasks) { + auto touches = [&](const auto ®ions) { + return std::any_of(regions.begin(), regions.end(), + [&](const BufferRegion ®ion) { + return region->buffer->data.same_as(storage); + }); + }; + if (!touches(task->GetReadRegions()) && !touches(task->GetWriteRegions())) + continue; + if (std::count_if( + found->second.begin(), found->second.end(), + [&](ControlNode *owner) { return task->IsWithin(owner); }) != 1) + return; + } + } + + ReuseRingBindings bindings{{b->data, a->data}}; + metadata.kernel_root.CopyOnWrite()->annotations.Set(kReuseRingBindings, + bindings); + L0StorageGroups groups = CollectL0StorageGroups(metadata.kernel_root); + std::function expand = [&](IRStructure *node) { + if (!node->IsControl()) + return; + auto *control = static_cast(node); + if (auto value = control->control->annotations.Get(kMultiBufferEligible)) { + Array eligible; + std::unordered_set seen; + for (const Var &storage : value.value().cast>()) { + for (const Var &member : groups.Members(storage)) { + if (seen.insert(member).second) + eligible.push_back(member); + } + } + control->control.CopyOnWrite()->annotations.Set(kMultiBufferEligible, + eligible); + } + for (const auto &child : control->children) + expand(child.get()); + }; + for (const auto &node : scheduled_tir->tree) + expand(node.get()); +} + static void WarnOnUnappliedBufferVersionOverrides( const BufferVersionMap &buffer_version_overrides, const BufferVersionMap &selected_buffer_versions) { @@ -1386,7 +1553,8 @@ static void ScheduleSingleKernel(const Stmt &kernel_body, Target target, // Schedule IRStructure in place with the Z3 scheduler. ScheduleBuilder unit_builder; - L0StorageGroups groups(CollectL0SFBindings(metadata.kernel_root)); + PlanExplicitL0BReuseRing(scheduled_tir); + L0StorageGroups groups = CollectL0StorageGroups(metadata.kernel_root); metadata.buffer_versions = ExpandL0StorageGroupValues( metadata.buffer_versions, groups, "version counts"); metadata.manual_buffer_versions = ExpandL0StorageGroupValues( diff --git a/src/ascend/transform/auto_schedule/buffer_alias_analysis.h b/src/ascend/transform/auto_schedule/buffer_alias_analysis.h index e5b38fa877..85468cc271 100644 --- a/src/ascend/transform/auto_schedule/buffer_alias_analysis.h +++ b/src/ascend/transform/auto_schedule/buffer_alias_analysis.h @@ -429,8 +429,12 @@ class BufferAliasAnalyzerImpl { ICHECK_EQ(hint.size(), 4U); if (!hint[3].cast()->value) continue; - Var lhs = groups_.Representative(storage(hint[0])); - Var rhs = groups_.Representative(storage(hint[1])); + Var lhs = storage(hint[0]); + Var rhs = storage(hint[1]); + if (!groups_.SharesReuseRing(lhs, rhs)) { + lhs = groups_.Representative(lhs); + rhs = groups_.Representative(rhs); + } // A positive hint explicitly permits storage reuse. Its region and // iteration qualifiers still control dependency analysis independently. if (!lhs.same_as(rhs)) diff --git a/src/ascend/transform/auto_schedule/dependency_analysis.cc b/src/ascend/transform/auto_schedule/dependency_analysis.cc index 3bfa1efb9d..5f18f6751f 100644 --- a/src/ascend/transform/auto_schedule/dependency_analysis.cc +++ b/src/ascend/transform/auto_schedule/dependency_analysis.cc @@ -120,8 +120,8 @@ bool RegionsMayConflict(const ConstrSet &a_ctx, const BufferRegion &a_region, ControlNode *loop, int offset, size_t num_storage_owners, const ConflictHintList &root_conflicts, - Optional a_pitch, - Optional b_pitch) { + Optional a_pitch, Optional b_pitch, + bool shared_reuse_ring) { const bool cross = offset != 0; if (cross && loop == nullptr) return false; @@ -147,7 +147,8 @@ bool RegionsMayConflict(const ConstrSet &a_ctx, const BufferRegion &a_region, if (declared_no_conflict) return false; - if (!a_region->buffer->data.same_as(b_region->buffer->data)) + if (!shared_reuse_ring && + !a_region->buffer->data.same_as(b_region->buffer->data)) return false; // Same storage but different logical buffer (T.view/T.reshape): regions may @@ -255,7 +256,8 @@ AnalyzeDependencies(std::vector nodes, ControlNode *loop, const BufferVersionMap &manual_buffer_versions, const MultiBufferOwnerMap &multi_buffer_owners, DependencyCache *dependency_cache, - const ConflictHintList &root_conflicts) { + const ConflictHintList &root_conflicts, + const L0StorageGroups &storage_groups) { DependencyCache local_cache; DependencyCache &cache = dependency_cache ? *dependency_cache : local_cache; if (loop) { @@ -299,9 +301,9 @@ AnalyzeDependencies(std::vector nodes, ControlNode *loop, std::vector children; for (const auto &child : ctrl->children) children.push_back(child.get()); - auto child_deps = - AnalyzeDependencies(children, ctrl, manual_buffer_versions, - multi_buffer_owners, &cache, root_conflicts); + auto child_deps = AnalyzeDependencies( + children, ctrl, manual_buffer_versions, multi_buffer_owners, &cache, + root_conflicts, storage_groups); for (const auto &cdep : child_deps) { auto &covered_pairs = out[cdep.storage]; covered_pairs.insert(cdep.task_pairs.begin(), cdep.task_pairs.end()); @@ -407,8 +409,12 @@ AnalyzeDependencies(std::vector nodes, ControlNode *loop, const Var &lhs_storage = lhs.region->buffer->data; const Var &rhs_storage = rhs.region->buffer->data; ffi::Optional storage; + bool shared_reuse_ring = + storage_groups.SharesReuseRing(lhs_storage, rhs_storage); if (lhs_storage.same_as(rhs_storage)) storage = lhs_storage; + else if (shared_reuse_ring) + storage = storage_groups.Representative(lhs_storage); bool owner_exclusion = is_owner_exclusion(storage, lhs.task, rhs.task); // The i nodes, ControlNode *loop, rhs.task->outer_ctx, rhs.region, loop, d, /*num_storage_owners=*/0, root_conflicts, lhs.task->GetL0Pitch(lhs_storage), - rhs.task->GetL0Pitch(rhs_storage))) { + rhs.task->GetL0Pitch(rhs_storage), + shared_reuse_ring)) { add_pair(storage, d, lhs.task, rhs.task); break; } @@ -435,12 +442,12 @@ AnalyzeDependencies(std::vector nodes, ControlNode *loop, continue; if (distance < 0 && loop == nullptr) continue; - if (RegionsMayConflict(lhs.task->outer_ctx, lhs.region, - rhs.task->outer_ctx, rhs.region, loop, - distance, get_num_storage_owners(storage), - root_conflicts, - lhs.task->GetL0Pitch(lhs_storage), - rhs.task->GetL0Pitch(rhs_storage))) { + if (RegionsMayConflict( + lhs.task->outer_ctx, lhs.region, rhs.task->outer_ctx, + rhs.region, loop, distance, + get_num_storage_owners(storage), root_conflicts, + lhs.task->GetL0Pitch(lhs_storage), + rhs.task->GetL0Pitch(rhs_storage), shared_reuse_ring)) { add_pair(storage, owner_exclusion ? 0 : distance, lhs.task, rhs.task, owner_exclusion ? DependencyKind::kOwnerExclusion diff --git a/src/ascend/transform/auto_schedule/dependency_analysis.h b/src/ascend/transform/auto_schedule/dependency_analysis.h index 6eb96f486d..d3054ab71e 100644 --- a/src/ascend/transform/auto_schedule/dependency_analysis.h +++ b/src/ascend/transform/auto_schedule/dependency_analysis.h @@ -106,7 +106,8 @@ AnalyzeDependencies(std::vector nodes, const BufferVersionMap &manual_buffer_versions = {}, const MultiBufferOwnerMap &multi_buffer_owners = {}, DependencyCache *dependency_cache = nullptr, - const ConflictHintList &root_conflicts = {}); + const ConflictHintList &root_conflicts = {}, + const L0StorageGroups &storage_groups = L0StorageGroups()); // Check whether two task regions may overlap under their respective symbolic // contexts. Only a proof of disjointness removes a conflict. @@ -128,7 +129,8 @@ bool RegionsMayConflict(const ConstrSet &a_ctx, const BufferRegion &a_region, size_t num_storage_owners = 0, const ConflictHintList &root_conflicts = {}, Optional a_pitch = std::nullopt, - Optional b_pitch = std::nullopt); + Optional b_pitch = std::nullopt, + bool shared_reuse_ring = false); } // namespace ascend } // namespace tl diff --git a/src/ascend/transform/buffer_version.h b/src/ascend/transform/buffer_version.h index 641bc0fcf7..c070f9e0dc 100644 --- a/src/ascend/transform/buffer_version.h +++ b/src/ascend/transform/buffer_version.h @@ -24,8 +24,10 @@ #pragma once +#include #include #include +#include #include "ascend/op/utils.h" #include "transform/common/attr.h" @@ -46,31 +48,56 @@ constexpr const char *kMultiBufferEligible = "multi_buffer_eligible"; /*! \brief Storage data Var -> positive compile-time version count. */ using BufferVersionMap = ffi::Map; +// Selected physical L0B reuse pairs. Layouts and logical buffers remain +// distinct, but the selected pair advances one storage-ring protocol. +constexpr const char *kReuseRingBindings = "tl.reuse_ring_bindings"; +using ReuseRingBindings = ffi::Map; + // Logical reads/writes remain separate. Ownership, physical versions and // address reuse use the same allocation group, represented by the data Var. class L0StorageGroups { public: - explicit L0StorageGroups(L0SFBindings bindings = {}) - : bindings_(std::move(bindings)) { - for (const auto &[sf, data] : bindings_) - scales_.Set(data, sf); + explicit L0StorageGroups(L0SFBindings bindings = {}, + ReuseRingBindings reuse_bindings = {}) + : sf_bindings_(std::move(bindings)), + reuse_bindings_(std::move(reuse_bindings)), bindings_(sf_bindings_) { + for (const auto &[member, data] : reuse_bindings_) + bindings_.Set(member, data); + for (const auto &[member, data] : bindings_) { + auto members = members_.Get(data).value_or(ffi::Array{data}); + members.push_back(member); + members_.Set(data, members); + } } tirx::Var Representative(const tirx::Var &storage) const { return bindings_.Get(storage).value_or(storage); } ffi::Array Members(const tirx::Var &storage) const { tirx::Var data = Representative(storage); - if (auto sf = scales_.Get(data)) - return {data, sf.value()}; - return {data}; + return members_.Get(data).value_or(ffi::Array{data}); + } + const L0SFBindings &Bindings() const { return sf_bindings_; } + bool SharesReuseRing(const tirx::Var &lhs, const tirx::Var &rhs) const { + auto lhs_data = reuse_bindings_.Get(lhs); + auto rhs_data = reuse_bindings_.Get(rhs); + return (lhs_data.has_value() && lhs_data.value().same_as(rhs)) || + (rhs_data.has_value() && rhs_data.value().same_as(lhs)); } - const L0SFBindings &Bindings() const { return bindings_; } private: + L0SFBindings sf_bindings_; + ReuseRingBindings reuse_bindings_; L0SFBindings bindings_; - L0SFBindings scales_; + ffi::Map> members_; }; +inline L0StorageGroups CollectL0StorageGroups(const tirx::SBlock &root) { + ReuseRingBindings reuse; + if (auto value = root->annotations.Get(kReuseRingBindings)) + reuse = value.value().cast(); + return L0StorageGroups(CollectL0SFBindings(root), std::move(reuse)); +} + // User overrides and prepared metadata must describe the entire group. // Contradictory member requests are errors, rather than last-writer wins. template diff --git a/src/ascend/transform/insert_sync.cc b/src/ascend/transform/insert_sync.cc index 9dd4f78903..98317c3a06 100644 --- a/src/ascend/transform/insert_sync.cc +++ b/src/ascend/transform/insert_sync.cc @@ -1250,9 +1250,9 @@ class SyncAnalyzer : public BufferAliasAnalysisContext { std::vector CollectSyncPoints(const std::vector &nodes, ControlNode *loop) { - std::vector deps = - AnalyzeDependencies(nodes, loop, manual_, multi_buffer_owners_, - &dependency_cache_, root_conflicts_); + std::vector deps = AnalyzeDependencies( + nodes, loop, manual_, multi_buffer_owners_, &dependency_cache_, + root_conflicts_, storage_groups_); if (loop != nullptr) { std::vector directed_deps; directed_deps.reserve(deps.size() * 2); @@ -2532,8 +2532,8 @@ SBlock InsertKernelSync(const Stmt &kernel_body, ScheduledTIR scheduled_tir) { std::vector> &root = scheduled_tir.tree; int num_aiv_subcores = scheduled_tir.metadata.num_aiv_subcores.value_or(2); CoreMask vector_mask = VectorCoreMask(num_aiv_subcores); - L0StorageGroups groups( - CollectL0SFBindings(scheduled_tir.metadata.kernel_root)); + L0StorageGroups groups = + CollectL0StorageGroups(scheduled_tir.metadata.kernel_root); MultiBufferPlan multi_buffer_plan = ReadMultiBufferPlan(root, scheduled_tir.metadata.buffer_versions, groups); std::vector tasks; diff --git a/src/ascend/transform/materialize_multi_buffer.cc b/src/ascend/transform/materialize_multi_buffer.cc index fb3d83c885..15ea1a3a5a 100644 --- a/src/ascend/transform/materialize_multi_buffer.cc +++ b/src/ascend/transform/materialize_multi_buffer.cc @@ -517,7 +517,7 @@ void ElideUnitVersionCounterGroups(ScheduledTIR *scheduled_tir, SBlock MaterializeKernel(ScheduledTIR scheduled_tir) { MultiBufferPlan plan = ReadMultiBufferPlan( scheduled_tir.tree, scheduled_tir.metadata.buffer_versions, - L0StorageGroups(CollectL0SFBindings(scheduled_tir.metadata.kernel_root))); + CollectL0StorageGroups(scheduled_tir.metadata.kernel_root)); ElideUnitVersionCounterGroups(&scheduled_tir, plan); VersionedBufferRegistry buffers(plan, scheduled_tir.metadata.kernel_root); RewriteScheduledTree(&scheduled_tir.tree, plan, &buffers); @@ -531,6 +531,8 @@ SBlock MaterializeKernel(ScheduledTIR scheduled_tir) { AddBufferVersion(&physical_versions, info.storage, info.num_versions); scheduled_tir.metadata.buffer_versions = {}; scheduled_tir.metadata.manual_buffer_versions = std::move(physical_versions); + scheduled_tir.metadata.kernel_root.CopyOnWrite()->annotations.erase( + kReuseRingBindings); Stmt serialized = EncodeScheduledTIR(std::move(scheduled_tir)); serialized = AllocationRewriter(&buffers)(std::move(serialized)); diff --git a/src/ascend/transform/prepare_multi_buffer.cc b/src/ascend/transform/prepare_multi_buffer.cc index 434a8a08e6..ae19c60c6f 100644 --- a/src/ascend/transform/prepare_multi_buffer.cc +++ b/src/ascend/transform/prepare_multi_buffer.cc @@ -1222,7 +1222,7 @@ tvm::transform::Pass PrepareMultiBuffer() { ICHECK_NE(task->GetCoreMask(), kCoreUnassigned) << "PrepareMultiBuffer requires AssignCore to run first"; } - L0StorageGroups groups(CollectL0SFBindings(context.root)); + L0StorageGroups groups = CollectL0StorageGroups(context.root); auto &metadata = scheduled_tir.metadata; metadata.buffer_versions = ExpandL0StorageGroupValues( metadata.buffer_versions, groups, "version counts"); diff --git a/src/ascend/transform/resolve_core.cc b/src/ascend/transform/resolve_core.cc index 9b87518bcb..d1d36b0a70 100644 --- a/src/ascend/transform/resolve_core.cc +++ b/src/ascend/transform/resolve_core.cc @@ -639,11 +639,10 @@ tvm::transform::Pass ResolveCore() { [](const TilelangKernelContext &context) { ScheduledTIR scheduled_tir = DecodeScheduledTIR(context.root, context.outer_ctx); - ResolveCoreMasks(scheduled_tir.tree, - scheduled_tir.metadata.buffer_versions, - context.outer_sid, - L0StorageGroups(CollectL0SFBindings(context.root)), - scheduled_tir.metadata.num_aiv_subcores.value_or(2)); + ResolveCoreMasks( + scheduled_tir.tree, scheduled_tir.metadata.buffer_versions, + context.outer_sid, CollectL0StorageGroups(context.root), + scheduled_tir.metadata.num_aiv_subcores.value_or(2)); return EncodeScheduledTIR(std::move(scheduled_tir)); }); }; diff --git a/testing/ascend/auto_schedule/test_conflict_ring_owners.py b/testing/ascend/auto_schedule/test_conflict_ring_owners.py new file mode 100644 index 0000000000..ce7fdf8d9c --- /dev/null +++ b/testing/ascend/auto_schedule/test_conflict_ring_owners.py @@ -0,0 +1,169 @@ +"""Explicit layout aliases share the physical clock of retimed L0B owners.""" + +import pytest +import tilelang +import tilelang.ascend.language as T +import tilelang.testing +from tilelang import tvm +from tilelang.layout import make_ascend_major_k_layout, make_ascend_major_mn_layout +from tvm import tirx +from testing.ascend._ir import allocated_buffer, nodes + + +def _retimed_layout_pair(qk_steps=4, pv_steps=4): + @T.prim_func + def main( + Q: T.Tensor((64, 128), "bfloat16"), + K: T.Tensor((64, 128), "bfloat16"), + CQ: T.Tensor((3, 64, 64), "float32"), + CV: T.Tensor((3, pv_steps, 64, 128), "float32"), + ): + with T.Kernel(1): + q = T.alloc_l1((64, 128), "bfloat16") + k0 = T.alloc_l1((32, 128), "bfloat16") + k1 = T.alloc_l1((32, 128), "bfloat16") + aq = T.alloc_l0a((64, 128), "bfloat16") + bq = T.alloc_l0b((64, 128), "bfloat16") + bv = T.alloc_l0b((64, 128), "bfloat16") + cq = T.alloc_l0c((64, 64), "float32") + cv = T.alloc_l0c((pv_steps, 64, 128), "float32") + T.annotate_buffer_versions({q: 1, k0: 1, k1: 1, aq: 4, bq: 4, bv: 4, cq: 1, cv: 1}) + T.annotate_layout({bq: make_ascend_major_k_layout(bq), bv: make_ascend_major_mn_layout(bv)}) + T.copy(Q, q) + T.copy(K[:32, :], k0) + T.copy(K[32:, :], k1) + for step in T.Pipelined(3, num_stages=2, annotations={"enable_offset": True}): + T.assume_conflict(bq, bv) + with T.Stage(0): + T.set_mmad_direction("m") + for d in T.Pipelined(qk_steps, num_stages=4): + with T.Stage(0): + T.copy(q, aq) + with T.Task(): + T.copy(k0, bq[:32, :], l0_pitch=4) + T.copy(k1, bq[32:, :], l0_pitch=4) + T.gemm(aq, bq, cq, transpose_B=True, clear_accum=d == 0) + T.copy(cq, CQ[step, :, :]) + with T.Stage(1), T.IssuePhase({"MTE1": 1, "M": 1}): + T.set_mmad_direction("n") + for _region in T.Pipelined(1, num_stages=4): + with T.Stage(0): + T.copy(q[:, :64], aq[:, :64]) + for d in T.Pipelined(pv_steps, num_stages=2): + with T.Stage(0): + with T.Task(): + T.copy(k0, bv[:32, :], l0_pitch=8) + T.copy(k1, bv[32:, :], l0_pitch=8) + T.gemm(aq[:, :64], bv, cv[d, :, :], clear_accum=True) + for d in T.Pipelined(pv_steps, num_stages=2): + with T.Stage(0): + T.copy(cv[d, :, :], CV[step, d, :, :]) + + return main + + +@pytest.mark.parametrize("qk_steps,pv_steps", [(4, 4), (3, 2)]) +@pytest.mark.parametrize("disable_reuse", [False, True]) +def test_default_conflict_hint_shares_retimed_owner_clock(qk_steps, pv_steps, disable_reuse): + snapshots = {} + + @tvm.instrument.pass_instrument + class Capture: + def run_after_pass(self, mod, info): + if info.name in ("tl.PrepareMultiBuffer", "tl.MergeUBAllocations"): + snapshots[info.name] = mod + + with ( + tvm.target.Target("ascend"), + tvm.transform.PassContext(config={"tl.disable_shared_memory_reuse": disable_reuse}, instruments=[Capture()]), + ): + tilelang.lower(_retimed_layout_pair(qk_steps, pv_steps), target="ascend", enable_device_compile=False) + + prepared = snapshots["tl.PrepareMultiBuffer"] + bq, bv = [allocated_buffer(prepared, name) for name in ("bq", "bv")] + assert not bq.data.same_as(bv.data) + clocks = [] + for loop in nodes(prepared, tirx.For): + if "tl.multi_buffer_counter_map" not in loop.annotations: + continue + counters = [ + counter + for storage, counter in loop.annotations["tl.multi_buffer_counter_map"].items() + if storage.same_as(bq.data) or storage.same_as(bv.data) + ] + if not counters: + continue + assert len(counters) == 2 + assert counters[0].same_as(counters[1]) + clocks.append(counters[0]) + assert len(clocks) == 2 + assert all(clock.same_as(clocks[0]) for clock in clocks) + advances = [ + store + for store in nodes(prepared, tirx.BufferStore) + if store.buffer.same_as(clocks[0]) and tvm.arith.Analyzer().can_prove_equal(store.value, clocks[0][0] + 1) + ] + assert len(advances) == 2 + + allocations = [ + node.buffer + for func in snapshots["tl.MergeUBAllocations"].functions.values() + for node in nodes(func, tirx.AllocBuffer) + if node.buffer.scope() in ("shared.l0b", "shared.l0b.dyn") + ] + assert len(allocations) == 1 + allocated_bytes = int(allocations[0].shape[0]) * tvm.DataType(allocations[0].dtype).bits // 8 + assert allocated_bytes == (131072 if disable_reuse else 65536) + + +def test_conflict_hint_does_not_join_nested_storage_owners(): + @T.prim_func + def main(Q: T.Tensor((64, 128), "bfloat16"), K: T.Tensor((64, 128), "bfloat16"), C: T.Tensor((64, 128), "float32")): + with T.Kernel(1): + q = T.alloc_l1((64, 128), "bfloat16") + k = T.alloc_l1((64, 128), "bfloat16") + aq = T.alloc_l0a((64, 128), "bfloat16") + av = T.alloc_l0a((64, 64), "bfloat16") + bq = T.alloc_l0b((64, 128), "bfloat16") + bv = T.alloc_l0b((64, 128), "bfloat16") + cq = T.alloc_l0c((64, 64), "float32") + cv = T.alloc_l0c((64, 128), "float32") + T.annotate_buffer_versions({q: 1, k: 1, aq: 1, av: 1, bq: 4, bv: 4, cq: 1, cv: 1}) + T.copy(Q, q) + T.copy(K, k) + for _i in T.Pipelined(2, num_stages=2): + T.assume_conflict(bq, bv) + with T.Stage(0): + T.copy(q, aq) + T.copy(k, bq) + T.gemm(aq, bq, cq, transpose_B=True, clear_accum=True) + T.copy(cq, C[:, :64]) + for _j in T.Pipelined(2, num_stages=2): + with T.Stage(0): + T.copy(q[:, :64], av) + T.copy(k, bv) + T.gemm(av, bv, cv, clear_accum=True) + T.copy(cv, C) + + with tvm.target.Target("ascend"): + tilelang.lower(main, target="ascend", enable_device_compile=False) + + +@tilelang.testing.requires_ascend +def test_retimed_layout_pair_values_across_owner_lengths_and_launches(): + import torch + + torch.npu.set_device(0) + torch.npu.set_op_timeout_ms(30000) + with tvm.target.Target("ascend"): + kernel = tilelang.compile(_retimed_layout_pair(3, 2), target="ascend", out_idx=[2, 3]) + for seed in (7, 11, 19): + torch.manual_seed(seed) + q = torch.randint(-2, 3, (64, 128), dtype=torch.int32).to(dtype=torch.bfloat16, device="npu") + k = torch.randint(-2, 3, (64, 128), dtype=torch.int32).to(dtype=torch.bfloat16, device="npu") + cq, cv = kernel(q, k) + torch.npu.synchronize() + expected_qk = 3 * (q.cpu().float() @ k.cpu().float().T) + expected_pv = q.cpu().float()[:, :64] @ k.cpu().float() + torch.testing.assert_close(cq.cpu(), expected_qk.expand(3, 64, 64), rtol=0, atol=0) + torch.testing.assert_close(cv.cpu(), expected_pv.expand(3, 2, 64, 128), rtol=0, atol=0) diff --git a/tilelang/ascend/language/schedule_hint.py b/tilelang/ascend/language/schedule_hint.py index 2c56067937..44497a699a 100644 --- a/tilelang/ascend/language/schedule_hint.py +++ b/tilelang/ascend/language/schedule_hint.py @@ -260,7 +260,11 @@ def assume_conflict(a, b=None, *, level=None, cross=None, group=None): permission applies to the whole storage pair. The caller must ensure that the declared dependencies cover every access affected by that reuse. Existing allocation scope, alignment, and reuse-disable rules still apply. - This does not merge buffer identities, layouts, or version counters. + Buffer identities and layouts remain separate. For an isolated, equal-sized + L0B pair with the same fixed automatic version count, a whole-buffer hint + with the default iteration matching can also select one shared ring clock. + All accesses must belong to disjoint automatic owners; explicit iteration indexing, + manual versions, strides, or offsets retain the ordinary dependency path. Args: a: first region -- ``buffer[indices...]`` or a bare ``buffer`` (the From cd1f1ffdcc4fecefe02b8a956e2f502cf0c3671c Mon Sep 17 00:00:00 2001 From: Denver Jin Date: Wed, 21 Oct 2026 12:02:53 +0800 Subject: [PATCH 3/7] [Ascend] Isolate shared-ring regressions from issue-phase APIs --- testing/ascend/auto_schedule/test_conflict_ring_owners.py | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/testing/ascend/auto_schedule/test_conflict_ring_owners.py b/testing/ascend/auto_schedule/test_conflict_ring_owners.py index ce7fdf8d9c..155f96cdb6 100644 --- a/testing/ascend/auto_schedule/test_conflict_ring_owners.py +++ b/testing/ascend/auto_schedule/test_conflict_ring_owners.py @@ -32,10 +32,11 @@ def main( T.copy(Q, q) T.copy(K[:32, :], k0) T.copy(K[32:, :], k1) + # Keep the Cube mode fixed so this regression isolates storage + # ownership from separately scheduled configuration changes. for step in T.Pipelined(3, num_stages=2, annotations={"enable_offset": True}): T.assume_conflict(bq, bv) with T.Stage(0): - T.set_mmad_direction("m") for d in T.Pipelined(qk_steps, num_stages=4): with T.Stage(0): T.copy(q, aq) @@ -44,8 +45,7 @@ def main( T.copy(k1, bq[32:, :], l0_pitch=4) T.gemm(aq, bq, cq, transpose_B=True, clear_accum=d == 0) T.copy(cq, CQ[step, :, :]) - with T.Stage(1), T.IssuePhase({"MTE1": 1, "M": 1}): - T.set_mmad_direction("n") + with T.Stage(1): for _region in T.Pipelined(1, num_stages=4): with T.Stage(0): T.copy(q[:, :64], aq[:, :64]) From 4c85ff7c026fb49d43ec50518790453204d62644 Mon Sep 17 00:00:00 2001 From: Denver Jin Date: Wed, 21 Oct 2026 17:47:52 +0800 Subject: [PATCH 4/7] [Ascend] Keep incompatible L0B owners on independent rings --- .../references/synchronization.md | 3 + src/ascend/transform/auto_schedule.cc | 23 +++++++ .../test_conflict_ring_owners.py | 65 +++++++++++++++++++ 3 files changed, 91 insertions(+) diff --git a/.agents/skills/tilelang-ascend/references/synchronization.md b/.agents/skills/tilelang-ascend/references/synchronization.md index 6f25576827..2b7d9a326e 100644 --- a/.agents/skills/tilelang-ascend/references/synchronization.md +++ b/.agents/skills/tilelang-ascend/references/synchronization.md @@ -337,6 +337,9 @@ equal-sized L0B pair with matching fixed automatic version counts can share a ring clock when a whole-buffer hint uses the default iteration matching and all accesses belong to disjoint automatic owners. Explicit iteration indexing, manual versions, strides, and offsets keep the ordinary dependency path. +Each owner's accesses must also be guaranteed to use a single schedule stage. +Cross-stage manual accesses and automatic scheduling with `enable_offset=True` +keep independent rings because the shared counter cannot span those stages. ## Trace the complete storage protocol diff --git a/src/ascend/transform/auto_schedule.cc b/src/ascend/transform/auto_schedule.cc index 2be86bd77a..cff87f17f5 100644 --- a/src/ascend/transform/auto_schedule.cc +++ b/src/ascend/transform/auto_schedule.cc @@ -1475,6 +1475,29 @@ static void PlanExplicitL0BReuseRing(ScheduledTIR *scheduled_tir) { auto found = owners.find(storage); if (found == owners.end() || found->second.empty()) return; + // Joining disjoint owners requires a counter. PrepareMultiBuffer can + // advance that counter only when each owner's storage accesses share a + // stage, so keep independent rings unless scheduling preserves this. + for (ControlNode *owner : found->second) { + std::optional storage_stage; + for (const auto &child : owner->children) { + if (!child->TouchesStorage(storage)) + continue; + int stage = child->GetStage(); + if (storage_stage.has_value() && storage_stage.value() != stage) + return; + storage_stage = stage; + } + if (!storage_stage.has_value()) + return; + if (storage_stage.value() == kUnscheduledStage) { + // Automatic offsets are not known yet. Without offsets, scheduling + // collapses all automatic tasks to stage zero. + auto enable_offset = owner->control->annotations.Get("enable_offset"); + if (enable_offset.has_value() && enable_offset.value().cast()) + return; + } + } for (TaskNode *task : tasks) { auto touches = [&](const auto ®ions) { return std::any_of(regions.begin(), regions.end(), diff --git a/testing/ascend/auto_schedule/test_conflict_ring_owners.py b/testing/ascend/auto_schedule/test_conflict_ring_owners.py index 155f96cdb6..93e6e419a2 100644 --- a/testing/ascend/auto_schedule/test_conflict_ring_owners.py +++ b/testing/ascend/auto_schedule/test_conflict_ring_owners.py @@ -149,6 +149,71 @@ def main(Q: T.Tensor((64, 128), "bfloat16"), K: T.Tensor((64, 128), "bfloat16"), tilelang.lower(main, target="ascend", enable_device_compile=False) +@pytest.mark.parametrize( + "automatic,staged_second", + [(False, False), (False, True), (True, False)], + ids=["manual-first", "manual-second", "automatic-offset"], +) +@pytest.mark.parametrize("disable_reuse", [False, True]) +def test_conflict_hint_keeps_independent_owners_when_stages_can_differ(automatic, staged_second, disable_reuse): + @T.prim_func + def main(Q: T.Tensor((16, 16), "bfloat16"), K: T.Tensor((16, 16), "bfloat16"), C: T.Tensor((2, 16, 16), "float32")): + with T.Kernel(1): + q = T.alloc_l1((16, 16), "bfloat16") + k = T.alloc_l1((16, 16), "bfloat16") + aq = T.alloc_l0a((16, 16), "bfloat16") + av = T.alloc_l0a((16, 16), "bfloat16") + bq = T.alloc_l0b((16, 16), "bfloat16") + bv = T.alloc_l0b((16, 16), "bfloat16") + cq = T.alloc_l0c((16, 16), "float32") + cv = T.alloc_l0c((16, 16), "float32") + T.annotate_buffer_versions({q: 1, k: 1, aq: 4, av: 1, bq: 4, bv: 4, cq: 1, cv: 1}) + T.assume_conflict(bq, bv) + T.copy(Q, q) + T.copy(K, k) + T.copy(q, av) + staged = bv if staged_second else bq + other = bq if staged_second else bv + for _i in T.Pipelined(4, num_stages=4, annotations={"enable_offset": True}): + if automatic: + T.copy(q, aq) + T.copy(k, staged) + T.gemm(aq, staged, cq, transpose_B=True, clear_accum=True) + else: + with T.Stage(0): + T.copy(q, aq) + T.copy(k, staged) + with T.Stage(1): + T.gemm(aq, staged, cq, transpose_B=True, clear_accum=True) + T.copy(cq, C[0, :, :]) + for _j in T.Pipelined(4, num_stages=4): + with T.Stage(0): + T.copy(k, other) + T.gemm(av, other, cv, transpose_B=True, clear_accum=True) + T.copy(cv, C[1, :, :]) + + snapshots = {} + + @tvm.instrument.pass_instrument + class Capture: + def run_after_pass(self, mod, info): + if info.name in ("tl.PrepareMultiBuffer", "tl.MaterializeMultiBuffer"): + snapshots[info.name] = mod + + with tvm.transform.PassContext(config={"tl.disable_shared_memory_reuse": disable_reuse}, instruments=[Capture()]): + artifact = tilelang.lower(main, target="ascend", enable_device_compile=False) + + prepared = snapshots["tl.PrepareMultiBuffer"] + owners = [] + for name in ("staged", "other"): + buffer = allocated_buffer(prepared, name) + (owner,) = [loop for loop in nodes(prepared, tirx.For) if buffer.data in loop.annotations.get("multi_buffer_eligible", [])] + owners.append(owner) + assert int(allocated_buffer(snapshots["tl.MaterializeMultiBuffer"], name).shape[0]) == 4 + assert not owners[0].same_as(owners[1]) + assert artifact.kernel_source + + @tilelang.testing.requires_ascend def test_retimed_layout_pair_values_across_owner_lengths_and_launches(): import torch From 393a0bde46d60a3b7cfaabeeb65ebc67a7a14ec7 Mon Sep 17 00:00:00 2001 From: Denver Jin Date: Thu, 22 Oct 2026 15:04:42 +0800 Subject: [PATCH 5/7] [Ascend] Preserve independent Cube layouts on shared buffer views --- .agents/skills/tilelang-ascend/SKILL.md | 15 +++ src/ascend/op/copy.cc | 10 ++ src/ascend/op/oob_padding.cc | 4 + src/ascend/op/oob_padding.h | 5 +- src/ascend/op/utils.h | 28 ++++-- src/ascend/transform/layout_inference.cc | 24 +++-- src/ascend/transform/lower_tile_op.cc | 38 +++++++- .../transform/materialize_multi_buffer.cc | 36 ++----- .../transform/normalize_fractal_storage.cc | 70 ++++++++++---- .../test_conflict_ring_owners.py | 95 +++++++++++++++++-- .../test_ascend_normalize_fractal_storage.py | 68 +++++++++++++ tilelang/ascend/language/__init__.py | 1 + tilelang/ascend/language/annotations.py | 14 +++ tilelang/ascend/op/gemm/gemm_mad.py | 56 +++-------- .../ascend/op/gemm/gemm_mad_blockscaled.py | 36 +------ tilelang/language/annotations.py | 14 ++- 16 files changed, 358 insertions(+), 156 deletions(-) diff --git a/.agents/skills/tilelang-ascend/SKILL.md b/.agents/skills/tilelang-ascend/SKILL.md index 705102a3ff..9cb632d6fe 100644 --- a/.agents/skills/tilelang-ascend/SKILL.md +++ b/.agents/skills/tilelang-ascend/SKILL.md @@ -347,6 +347,21 @@ when changing manual kernels. For capacity or launch failures, record the generated allocation sizes, buffer versions, VF mode, and compiler options. Check SDK reservations as well as user buffers before changing tile sizes. +For Cube storage (L1, L0A, L0B, L0C), `T.view` and `T.reshape` aliases can have +different canonical layouts inferred from their operations or supplied with +`T.annotate_layout`. For example, a `(64, 128)` L0B allocation can have a +`(128, 64)` view filled with `T.copy(..., transpose=True)` and consumed with +`T.gemm(..., transpose_B=True)`. A same-shape view can instead use a different +`MajorK`/`MajorMN` layout. Each view keeps its own address mapping while the +shared `.data` keeps one allocation and version clock; `T.assume_conflict` is +unnecessary for these aliases. + +These views require dense logical strides, zero base offset, the same accessed +dtype, and equal padded physical sizes. They reinterpret storage without moving +existing values: the producer must write the representation its consumer uses. +Use `T.access_ptr(view[...], ...)` when passing a view to an intrinsic, so layout +lowering can select its address mapping. + `T.annotate_unlimit_memory("shared", "shared.l1")` inside a kernel removes the scheduler's capacity constraint for those pools. Other supported scopes are `shared.l0a`, `shared.l0b`, and `shared.l0c`. This is useful for investigating a diff --git a/src/ascend/op/copy.cc b/src/ascend/op/copy.cc index 302d832b99..e45d9c8848 100644 --- a/src/ascend/op/copy.cc +++ b/src/ascend/op/copy.cc @@ -425,6 +425,16 @@ Stmt LowerDMACopy(const AscendCopyNode &op, const LowerArgs &T, auto make_access_ptr = [&](const Buffer &buf, PrimExpr offset, PrimExpr extent, int rw_mask) -> PrimExpr { + if (IsAscendCubeBuffer(buf)) { + Array indices(buf->shape.size(), I(0)); + PrimExpr remaining = offset; + for (int i = static_cast(buf->shape.size()) - 1; i >= 0; --i) { + indices.Set(i, i == 0 ? remaining : floormod(remaining, buf->shape[i])); + remaining = floordiv(remaining, buf->shape[i]); + } + return Call(DataType::Handle(), tl::access_ptr(), + {BufferLoad(buf, indices), extent, I(rw_mask)}); + } return Call( DataType::Handle(), builtin::tvm_access_ptr(), {TypeAnnotation(buf->dtype), buf->data, offset, extent, I(rw_mask)}); diff --git a/src/ascend/op/oob_padding.cc b/src/ascend/op/oob_padding.cc index 03759cb935..cc595482be 100644 --- a/src/ascend/op/oob_padding.cc +++ b/src/ascend/op/oob_padding.cc @@ -72,6 +72,10 @@ Optional FindLayoutForBuffer(const LayoutMap &layout_map, if (layout_map.count(buffer)) { return layout_map[buffer]; } + // A Cube alias may have a different canonical matrix interpretation. + // Borrowing a sibling's layout here suppresses its own inference constraints. + if (IsAscendCubeBuffer(buffer)) + return std::nullopt; for (const auto &kv : layout_map) { if (kv.first->data.same_as(buffer->data)) { return kv.second; diff --git a/src/ascend/op/oob_padding.h b/src/ascend/op/oob_padding.h index eaf0e2b6e4..c6eaeef050 100644 --- a/src/ascend/op/oob_padding.h +++ b/src/ascend/op/oob_padding.h @@ -53,8 +53,9 @@ enum class DMAPath { DMAPath GetDMAPath(const Buffer &src, const Buffer &dst); /*! - * \brief Look up a buffer's inferred Layout in a layout map, matching either - * the buffer object or its underlying data Var. Returns nullopt if absent. + * \brief Look up a buffer's inferred Layout by Buffer identity. Non-Cube + * buffers may fall back to an alias with the same data Var; Cube views can + * have distinct layouts and require an exact match. Returns nullopt if absent. */ Optional FindLayoutForBuffer(const LayoutMap &layout_map, const Buffer &buffer); diff --git a/src/ascend/op/utils.h b/src/ascend/op/utils.h index e56a22f330..dba7d262bb 100644 --- a/src/ascend/op/utils.h +++ b/src/ascend/op/utils.h @@ -13,6 +13,7 @@ #include +#include "op/builtin.h" #include "op/utils.h" namespace tvm { @@ -74,8 +75,13 @@ inline bool IsAscendOnChipBuffer(const Buffer &buffer) { IsL0BBuffer(buffer) || IsL0CBuffer(buffer) || IsSSBuffer(buffer); } +inline bool IsAscendCubeBuffer(const Buffer &buffer) { + return IsL1Buffer(buffer) || IsL0ABuffer(buffer) || IsL0BBuffer(buffer) || + IsL0CBuffer(buffer); +} + /*! - * \brief tvm_access_ptr for a region whose innermost 2-D position travels as + * \brief A view-preserving pointer whose innermost 2-D position travels as * separate intrinsic arguments. * * The Ascend L1->L0 loads pass the M/K start positions as their own arguments @@ -85,13 +91,8 @@ inline bool IsAscendOnChipBuffer(const Buffer &buffer) { * example -- and counting the last two dims into the offset as well would place * that position twice. * - * This builds exactly that pointer, on top of the neutral - * MakeAccessPtrFromRegion: the innermost two dims are re-based to zero before - * delegating, which leaves their extents (and hence the reported extent) alone - * and drops them from the offset. - * - * The buffer keeps ``require_2d == false`` semantics, so 1-D regions still take - * the neutral function's dedicated path and are unaffected. + * Keep the Buffer identity until layout lowering: a data Var alone cannot + * identify which canonical matrix view supplies the physical addresses. */ inline PrimExpr MakeAscendLeadingDimAccessPtr(const BufferRegion ®ion, int rw_mask) { @@ -104,8 +105,15 @@ inline PrimExpr MakeAscendLeadingDimAccessPtr(const BufferRegion ®ion, const auto &r = leading[i]; leading.Set(i, Range::FromMinExtent(make_zero(r->min.dtype()), r->extent)); } - return MakeAccessPtrFromRegion(BufferRegion(region->buffer, leading), - rw_mask); + ffi::Array indices; + for (const Range &range : leading) + indices.push_back(range->min); + PrimExpr extent = leading.back()->extent; + if (n >= 2) + extent = extent * leading[n - 2]->extent; + return Call(DataType::Handle(), tl::access_ptr(), + {BufferLoad(region->buffer, indices), extent, + IntImm(DataType::Int(32), rw_mask)}); } } // namespace tl diff --git a/src/ascend/transform/layout_inference.cc b/src/ascend/transform/layout_inference.cc index 17425a378a..46dbe20c1b 100644 --- a/src/ascend/transform/layout_inference.cc +++ b/src/ascend/transform/layout_inference.cc @@ -316,6 +316,10 @@ class BufferUseDefCollector : public IRVisitorWithAnalyzer { // Helper: propagate inferred layout to alias buffers (same data Var) auto propagate_alias = [&](const Buffer &src_buffer, const Layout &src_layout) { + // Cube operands constrain their own views. Only aliases without a + // constraint inherit a reshaped layout after inference completes. + if (IsAscendCubeBuffer(src_buffer)) + return; if (!buffer_data_to_buffers_.count(src_buffer->data)) return; const auto &siblings = buffer_data_to_buffers_[src_buffer->data]; @@ -1003,9 +1007,10 @@ class BufferUseDefCollector : public IRVisitorWithAnalyzer { // populated. if (pushed_region_scope && op->annotations.count(attr::kLayoutMap)) { auto map = op->annotations.Get(attr::kLayoutMap) - ->as>() - .value(); - for (const auto &[var, layout] : map) { + ->cast>(); + for (const auto &[key, layout] : map) { + Var var = key.as() ? Downcast(key)->data + : Downcast(key); if (buffer_data_to_buffers_.count(var)) { for (const auto &buffer : buffer_data_to_buffers_[var]) { annotated_layout_map_.Set(buffer, layout); @@ -1038,13 +1043,18 @@ class BufferUseDefCollector : public IRVisitorWithAnalyzer { // After visiting, apply layouts to all collected buffers if (op->annotations.count(attr::kLayoutMap)) { - // Check if the layout map is Map auto map = - op->annotations.Get(attr::kLayoutMap)->as>().value(); - for (const auto &[var, layout] : map) { + op->annotations.Get(attr::kLayoutMap)->cast>(); + for (const auto &[key, layout] : map) { + Optional annotated = key.as(); + Var var = + annotated.defined() ? annotated.value()->data : Downcast(key); ICHECK(buffer_data_to_buffers_.count(var)) << "buffer " << var << " is not found in the block"; - const auto &buffers = buffer_data_to_buffers_[var]; + Array buffers = + annotated.defined() && IsAscendCubeBuffer(annotated.value()) + ? Array{annotated.value()} + : buffer_data_to_buffers_[var]; ICHECK(!buffers.empty()) << "buffer list for " << var << " is empty"; for (const auto &buffer : buffers) { if (!IsReducerV2Buffer(buffer)) { diff --git a/src/ascend/transform/lower_tile_op.cc b/src/ascend/transform/lower_tile_op.cc index 82a52a4389..3bca5bbe07 100644 --- a/src/ascend/transform/lower_tile_op.cc +++ b/src/ascend/transform/lower_tile_op.cc @@ -695,10 +695,11 @@ class LowerTileOpPass : arith::IRMutatorWithAnalyzer { } BufferLoad load = Downcast(access_ptr_call->args[0]); - PrimExpr extent = access_ptr_call->args[1]; - PrimExpr rw_mask = access_ptr_call->args[2]; + PrimExpr extent = VisitExpr(access_ptr_call->args[1]); + PrimExpr rw_mask = VisitExpr(access_ptr_call->args[2]); - Array indices = load->indices; + Array indices = load->indices.Map( + [&](const PrimExpr &index) { return VisitExpr(index); }); Array old_shape = load->buffer->shape; ICHECK_EQ(indices.size(), old_shape.size()) @@ -738,6 +739,28 @@ class LowerTileOpPass : arith::IRMutatorWithAnalyzer { Array new_args = {BufferLoad(new_buffer, new_indices), extent, rw_mask}; + if (IsAscendCubeBuffer(remap_key)) { + size_t n = new_shape.size(); + PrimExpr physical_offset = 0; + for (size_t i = 0; i < n; ++i) + physical_offset = physical_offset * new_shape[i] + new_indices[i]; + if (auto pitch = + access_ptr_call->annotations.Get(attr::kL0AccessPitch)) { + ICHECK(IsL0ABuffer(remap_key) || IsL0BBuffer(remap_key)); + ICHECK_GE(n, 4); + physical_offset += + new_indices[n - 4] * + (Downcast(pitch.value()) - new_shape[n - 3]) * + new_shape[n - 2] * new_shape[n - 1]; + } + layout_remap_.Set(new_buffer, layout.value()); + result.rewritten = true; + result.expr = + Call(access_ptr_call->dtype, builtin::tvm_access_ptr(), + {TypeAnnotation(new_buffer->dtype), new_buffer->data, + analyzer_->Simplify(physical_offset), extent, rw_mask}); + return result; + } if (buffer_remap_.count(remap_key)) { layout_remap_.Set(new_buffer, layout.value()); } @@ -1008,6 +1031,15 @@ class LowerTileOpPass : arith::IRMutatorWithAnalyzer { return call; } + // Preserve the typed base until its view-specific layout and compact L0 + // pitch have both been applied. Visiting BufferLoad first loses that view. + if (op->op.same_as(tl::access_ptr())) { + auto rewritten = HandleAccessPtrAndOffset(GetRef(op)); + if (rewritten.rewritten) + return rewritten.expr; + return IRMutatorWithAnalyzer::VisitExpr_(op); + } + // Handle standalone tvm_access_ptr calls with layout transformation if (op->op.same_as(builtin::tvm_access_ptr())) { auto call = Downcast(IRMutatorWithAnalyzer::VisitExpr_(op)); diff --git a/src/ascend/transform/materialize_multi_buffer.cc b/src/ascend/transform/materialize_multi_buffer.cc index 15ea1a3a5a..baa423bd70 100644 --- a/src/ascend/transform/materialize_multi_buffer.cc +++ b/src/ascend/transform/materialize_multi_buffer.cc @@ -354,30 +354,6 @@ Layout ExpandLayoutToBuffer(const Buffer &layout_buffer, return layout->Expand(leading_shape); } -Layout FindLayoutForBuffer(const Map &layout_map, - const Buffer &buffer) { - if (layout_map.count(buffer)) - return layout_map[buffer]; - for (const auto &[layout_buffer, layout] : layout_map) { - if (layout_buffer->data.same_as(buffer->data)) { - Layout expanded = ExpandLayoutToBuffer(layout_buffer, buffer, layout); - if (expanded.defined()) - return expanded; - } - } - return Layout(); -} - -bool AddLayoutForBuffer(Map *layout_map, const Buffer &buffer) { - if (layout_map->count(buffer)) - return false; - Layout layout = FindLayoutForBuffer(*layout_map, buffer); - if (!layout.defined()) - return false; - layout_map->Set(buffer, layout); - return true; -} - class AllocationRewriter : public StmtMutator { public: explicit AllocationRewriter(VersionedBufferRegistry *buffers) @@ -406,8 +382,16 @@ class AllocationRewriter : public StmtMutator { if (auto value = op->annotations.Get(attr::kLayoutMap)) { if (auto layout_map = value.value().as>()) { Map rewritten_layouts = layout_map.value(); - for (const Buffer &buffer : alloc_buffers) - layout_changed |= AddLayoutForBuffer(&rewritten_layouts, buffer); + // Expand each view from its own layout. Data-Var lookup is ambiguous + // when several canonical matrix views share a versioned allocation. + for (const auto &[buffer, layout] : layout_map.value()) { + Buffer expanded = buffers_->Resolve(buffer); + if (!expanded.same_as(buffer)) { + rewritten_layouts.Set( + expanded, ExpandLayoutToBuffer(buffer, expanded, layout)); + layout_changed = true; + } + } if (layout_changed) annotations.Set(attr::kLayoutMap, rewritten_layouts); } diff --git a/src/ascend/transform/normalize_fractal_storage.cc b/src/ascend/transform/normalize_fractal_storage.cc index 7fb7228b82..79bd979008 100644 --- a/src/ascend/transform/normalize_fractal_storage.cc +++ b/src/ascend/transform/normalize_fractal_storage.cc @@ -61,10 +61,10 @@ class BufferCollector : public StmtExprVisitor { } }; -// Every alias has the same padded matrix shape and canonical layout. Its old -// reshape is absorbed into indices/regions here, not a special layout or hidden -// pitch. Ordinary capacity accounting and version expansion therefore see the -// complete storage through every alias. GM/UB and compact NZ are untouched. +// Ordinary reshapes inherit a canonical matrix interpretation and have their +// indices normalized to it. Independently constrained Cube views retain their +// own interpretation, but share one allocation and an equal physical pitch. +// GM/UB and compact NZ are untouched. class FractalStorageNormalizer : public arith::IRMutatorWithAnalyzer { public: static PrimFunc Rewrite(PrimFunc func) { @@ -86,12 +86,27 @@ class FractalStorageNormalizer : public arith::IRMutatorWithAnalyzer { ICHECK(TryExtractAscendFractalLayout(collector.layouts[anchor.value()], anchor.value(), &lhs)); ICHECK(TryExtractAscendFractalLayout(layout, buffer, &rhs)); - ICHECK(lhs.kind == rhs.kind && - collector.layouts[anchor.value()]->IsEqual(layout.get()) && - anchor.value()->dtype.bits() == buffer->dtype.bits()) - << "Conflicting canonical Cube views of " << buffer->name; - if (collector.used_buffers.count(buffer) && - !collector.used_buffers.count(anchor.value())) + if (!collector.layouts[anchor.value()]->IsEqual(layout.get())) { + PrimExpr lhs_bits = Integer(anchor.value()->dtype.bits()); + PrimExpr rhs_bits = Integer(buffer->dtype.bits()); + for (const PrimExpr &extent : + collector.layouts[anchor.value()]->OutputShape()) + lhs_bits = lhs_bits * extent; + for (const PrimExpr &extent : layout->OutputShape()) + rhs_bits = rhs_bits * extent; + ICHECK(analyzer.CanProveEqual(lhs_bits, rhs_bits)) + << "Cube views sharing storage require equal padded physical " + "sizes: " + << anchor.value()->name << " and " << buffer->name; + rewriter.independent_views_.insert(buffer->data); + } + ICHECK_EQ(anchor.value()->dtype.bits(), buffer->dtype.bits()) + << "Cube aliases must preserve element width: " << buffer->name; + if ((collector.used_buffers.count(buffer) && + !collector.used_buffers.count(anchor.value())) || + (collector.used_buffers.count(buffer) == + collector.used_buffers.count(anchor.value()) && + buffer.same_as(collector.owners[buffer->data]))) rewriter.anchors_.Set(buffer->data, buffer); } else { rewriter.anchors_.Set(buffer->data, buffer); @@ -104,6 +119,17 @@ class FractalStorageNormalizer : public arith::IRMutatorWithAnalyzer { << "Cube alias accesses must use the canonical operand dtype: " << buffer->name << " has " << buffer->dtype << ", expected " << anchor.value()->dtype; + if (rewriter.independent_views_.count(buffer->data)) { + auto layout = collector.layouts.Get(buffer); + if (layout.has_value() && IsAscendFractalLayout(layout.value())) { + anchor = buffer; + } else { + ICHECK(!collector.used_buffers.count(buffer)) + << "Each accessed Cube view sharing storage with distinct " + "layouts requires its own canonical layout: " + << buffer->name; + } + } rewriter.RegisterBuffer_(buffer, anchor.value(), collector.owners[buffer->data], collector.layouts[anchor.value()]); @@ -156,10 +182,13 @@ class FractalStorageNormalizer : public arith::IRMutatorWithAnalyzer { shape.Set(axis, analyzer_->Simplify( floordiv(shape[axis] + align - 1, align) * align)); } - // Coalesce aliases to one Buffer identity and geometry so versioning - // expands the allocation and every use together. An unused equal-width - // backing allocation may adopt the canonical operand dtype. - Buffer padded = canonical_buffers_.Get(buffer->data).value_or(owner); + // Ordinary aliases coalesce to one Buffer identity and geometry. Distinct + // canonical views retain their own geometry on the same data Var. An unused + // equal-width backing allocation may adopt the canonical operand dtype. + bool independent = independent_views_.count(buffer->data); + Buffer padded = independent + ? buffer + : canonical_buffers_.Get(buffer->data).value_or(owner); if (!StructuralEqual()(padded->shape, shape) || !padded->strides.empty() || padded->dtype != anchor->dtype) { auto *writer = padded.CopyOnWrite(); @@ -183,7 +212,9 @@ class FractalStorageNormalizer : public arith::IRMutatorWithAnalyzer { break; } buffers_.Set(buffer, padded); - canonical_buffers_.Set(buffer->data, padded); + buffer_anchors_.Set(buffer, anchor); + if (!independent || buffer.same_as(owner)) + canonical_buffers_.Set(buffer->data, padded); canonical_layouts_.Set(padded, padded_layout); } @@ -204,7 +235,7 @@ class FractalStorageNormalizer : public arith::IRMutatorWithAnalyzer { }); if (!buffers_.count(buffer)) return BufferRegion(buffer, ranges); - if (!StructuralEqual()(buffer->shape, anchors_[buffer->data]->shape)) { + if (!StructuralEqual()(buffer->shape, buffer_anchors_[buffer]->shape)) { // A dense reshape is injective only while trailing region extents do // not wrap across the original view's strides. Otherwise equal bounding // box volume alone can hide duplicated logical elements. @@ -456,7 +487,7 @@ class FractalStorageNormalizer : public arith::IRMutatorWithAnalyzer { // GM region unchanged: extra storage does not authorize an extra read. Copy rewritten(call->args, call->annotations); Array ranges = rewritten->dst_range; - const auto &logical_shape = anchors_[copy->dst->data]->shape; + const auto &logical_shape = buffer_anchors_[copy->dst]->shape; for (size_t i = 0; i < ranges.size(); ++i) { if (analyzer_->CanProveEqual(ranges[i]->min, 0) && analyzer_->CanProveEqual(ranges[i]->extent, logical_shape[i])) @@ -472,6 +503,9 @@ class FractalStorageNormalizer : public arith::IRMutatorWithAnalyzer { if (op->op.same_as(builtin::tvm_access_ptr())) { auto data = op->args[1].as(); if (data && anchors_.count(data.value())) { + ICHECK(!independent_views_.count(data.value())) + << "Raw Cube pointers cannot select among distinct view layouts; " + "use T.access_ptr(view[...], ...)"; Buffer anchor = anchors_[data.value()]; ICHECK(op->args[0].dtype() == anchor->dtype) << "Cube raw pointer must use the canonical operand dtype " @@ -520,6 +554,8 @@ class FractalStorageNormalizer : public arith::IRMutatorWithAnalyzer { return block; } Map anchors_; + Map buffer_anchors_; + std::unordered_set independent_views_; Map canonical_buffers_; Map buffers_; LayoutMap canonical_layouts_; diff --git a/testing/ascend/auto_schedule/test_conflict_ring_owners.py b/testing/ascend/auto_schedule/test_conflict_ring_owners.py index 93e6e419a2..306a10e07d 100644 --- a/testing/ascend/auto_schedule/test_conflict_ring_owners.py +++ b/testing/ascend/auto_schedule/test_conflict_ring_owners.py @@ -10,7 +10,7 @@ from testing.ascend._ir import allocated_buffer, nodes -def _retimed_layout_pair(qk_steps=4, pv_steps=4): +def _retimed_layout_pair(qk_steps=4, pv_steps=4, alias_kind=None, transpose_pv=False, annotated=True): @T.prim_func def main( Q: T.Tensor((64, 128), "bfloat16"), @@ -24,18 +24,30 @@ def main( k1 = T.alloc_l1((32, 128), "bfloat16") aq = T.alloc_l0a((64, 128), "bfloat16") bq = T.alloc_l0b((64, 128), "bfloat16") - bv = T.alloc_l0b((64, 128), "bfloat16") + if alias_kind == "view": + bv = T.view(bq, (128, 64) if transpose_pv else (64, 128)) + elif alias_kind == "reshape": + bv = T.reshape(bq, (128, 64) if transpose_pv else (64, 128)) + else: + bv = T.alloc_l0b((64, 128), "bfloat16") cq = T.alloc_l0c((64, 64), "float32") cv = T.alloc_l0c((pv_steps, 64, 128), "float32") T.annotate_buffer_versions({q: 1, k0: 1, k1: 1, aq: 4, bq: 4, bv: 4, cq: 1, cv: 1}) - T.annotate_layout({bq: make_ascend_major_k_layout(bq), bv: make_ascend_major_mn_layout(bv)}) + if annotated: + T.annotate_layout( + { + bq: make_ascend_major_k_layout(bq), + bv: make_ascend_major_k_layout(bv) if transpose_pv else make_ascend_major_mn_layout(bv), + } + ) T.copy(Q, q) T.copy(K[:32, :], k0) T.copy(K[32:, :], k1) # Keep the Cube mode fixed so this regression isolates storage # ownership from separately scheduled configuration changes. for step in T.Pipelined(3, num_stages=2, annotations={"enable_offset": True}): - T.assume_conflict(bq, bv) + if alias_kind is None: + T.assume_conflict(bq, bv) with T.Stage(0): for d in T.Pipelined(qk_steps, num_stages=4): with T.Stage(0): @@ -52,9 +64,13 @@ def main( for d in T.Pipelined(pv_steps, num_stages=2): with T.Stage(0): with T.Task(): - T.copy(k0, bv[:32, :], l0_pitch=8) - T.copy(k1, bv[32:, :], l0_pitch=8) - T.gemm(aq[:, :64], bv, cv[d, :, :], clear_accum=True) + if transpose_pv: + T.copy(k0, bv[:, :32], transpose=True, l0_pitch=8) + T.copy(k1, bv[:, 32:], transpose=True, l0_pitch=8) + else: + T.copy(k0, bv[:32, :], l0_pitch=8) + T.copy(k1, bv[32:, :], l0_pitch=8) + T.gemm(aq[:, :64], bv, cv[d, :, :], transpose_B=transpose_pv, clear_accum=True) for d in T.Pipelined(pv_steps, num_stages=2): with T.Stage(0): T.copy(cv[d, :, :], CV[step, d, :, :]) @@ -116,6 +132,64 @@ def run_after_pass(self, mod, info): assert allocated_bytes == (131072 if disable_reuse else 65536) +@pytest.mark.parametrize("alias_kind", ["view", "reshape"]) +@pytest.mark.parametrize("transpose_pv,annotated", [(False, True), (True, True), (False, False), (True, False)]) +def test_storage_views_keep_layouts_and_share_one_version_clock(alias_kind, transpose_pv, annotated): + snapshots = {} + + @tvm.instrument.pass_instrument + class Capture: + def run_after_pass(self, mod, info): + if info.name in ( + "tl.NormalizeAscendFractalStorage", + "tl.PrepareMultiBuffer", + "tl.MaterializeMultiBuffer", + "tl.MergeUBAllocations", + ): + snapshots[info.name] = mod + + with tvm.target.Target("ascend"), tvm.transform.PassContext(config={"tl.disable_shared_memory_reuse": True}, instruments=[Capture()]): + tilelang.lower(_retimed_layout_pair(3, 2, alias_kind, transpose_pv, annotated), target="ascend", enable_device_compile=False) + + def l0b_layouts(mod): + return { + buffer: layout + for block in nodes(mod, tirx.SBlock) + for buffer, layout in block.annotations.get("layout_map", {}).items() + if buffer.scope() == "shared.l0b" + } + + layouts = l0b_layouts(snapshots["tl.NormalizeAscendFractalStorage"]) + assert len(layouts) == 2 + first, second = layouts + assert first.data.same_as(second.data) + assert not layouts[first].is_equal(layouts[second]) + expected_shapes = {(64, 128), (128, 64)} if transpose_pv else {(64, 128)} + assert {tuple(int(x) for x in buffer.shape) for buffer in layouts} == expected_shapes + + clocks = [ + counter + for loop in nodes(snapshots["tl.PrepareMultiBuffer"], tirx.For) + for storage, counter in loop.annotations.get("tl.multi_buffer_counter_map", {}).items() + if storage.same_as(first.data) + ] + assert len(clocks) == 2 and clocks[0].same_as(clocks[1]) + expanded = {buffer: layout for buffer, layout in l0b_layouts(snapshots["tl.MaterializeMultiBuffer"]).items() if len(buffer.shape) == 3} + assert len(expanded) == 2 + for buffer, layout in expanded.items(): + assert int(buffer.shape[0]) == 4 and int(buffer.strides[0]) == 8192 + expected = make_ascend_major_k_layout(buffer) if transpose_pv or buffer.name == "bq" else make_ascend_major_mn_layout(buffer) + assert layout.is_equal(expected) + allocations = [ + node.buffer + for func in snapshots["tl.MergeUBAllocations"].functions.values() + for node in nodes(func, tirx.AllocBuffer) + if node.buffer.scope().startswith("shared.l0b") + ] + assert len(allocations) == 1 + assert int(allocations[0].shape[0]) * tvm.DataType(allocations[0].dtype).bits // 8 == 65536 + + def test_conflict_hint_does_not_join_nested_storage_owners(): @T.prim_func def main(Q: T.Tensor((64, 128), "bfloat16"), K: T.Tensor((64, 128), "bfloat16"), C: T.Tensor((64, 128), "float32")): @@ -214,14 +288,17 @@ def run_after_pass(self, mod, info): assert artifact.kernel_source +@pytest.mark.parametrize( + "alias_kind,transpose_pv,annotated", [(None, False, True), ("view", False, True), ("reshape", True, True), ("view", True, False)] +) @tilelang.testing.requires_ascend -def test_retimed_layout_pair_values_across_owner_lengths_and_launches(): +def test_retimed_layout_pair_values_across_owner_lengths_and_launches(alias_kind, transpose_pv, annotated): import torch torch.npu.set_device(0) torch.npu.set_op_timeout_ms(30000) with tvm.target.Target("ascend"): - kernel = tilelang.compile(_retimed_layout_pair(3, 2), target="ascend", out_idx=[2, 3]) + kernel = tilelang.compile(_retimed_layout_pair(3, 2, alias_kind, transpose_pv, annotated), target="ascend", out_idx=[2, 3]) for seed in (7, 11, 19): torch.manual_seed(seed) q = torch.randint(-2, 3, (64, 128), dtype=torch.int32).to(dtype=torch.bfloat16, device="npu") diff --git a/testing/ascend/transform/test_ascend_normalize_fractal_storage.py b/testing/ascend/transform/test_ascend_normalize_fractal_storage.py index 869669c45d..27d1093e52 100644 --- a/testing/ascend/transform/test_ascend_normalize_fractal_storage.py +++ b/testing/ascend/transform/test_ascend_normalize_fractal_storage.py @@ -161,6 +161,74 @@ def test_alias_normalization_preserves_allocation_properties(): assert allocation.same_as(normalized["main"].body.block.body.value.buffer) +@pytest.mark.parametrize( + "shape,view_shape,transpose", [((64, 128), (64, 128), False), ((64, 128), (128, 64), True), ((17, 32), (32, 17), False)] +) +def test_independent_cube_views_preserve_their_matrix_geometry(shape, view_shape, transpose): + owner = tirx.decl_buffer(shape, "bfloat16", name="owner", scope="shared.l0b") + view = tirx.decl_buffer(view_shape, "bfloat16", name="view", data=owner.data) + layouts = { + owner: tilelang.layout.make_ascend_major_k_layout(owner), + view: (tilelang.layout.make_ascend_major_k_layout if transpose else tilelang.layout.make_ascend_major_mn_layout)(view), + } + body = tirx.SeqStmt([tirx.Evaluate(owner[1, 0]), tirx.BufferStore(view, tirx.const(7, "bfloat16"), [0, 1])]) + block = tirx.SBlock([], [], [], "root", body, alloc_buffers=[owner], annotations={"layout_map": layouts}) + before = tvm.IRModule({"main": tirx.PrimFunc([], tirx.SBlockRealize([], True, block))}) + after = transform.NormalizeAscendFractalStorage()(before) + root = after["main"].body.block + read, write = root.body.seq[0].value, root.body.seq[1] + assert not read.buffer.same_as(write.buffer) + assert read.buffer.data.same_as(write.buffer.data) + assert tuple(int(x) for x in read.indices) == (1, 0) + assert tuple(int(x) for x in write.indices) == (0, 1) + assert tuple(int(x) for x in read.buffer.shape) == tuple((x + 15) // 16 * 16 for x in shape) + assert tuple(int(x) for x in write.buffer.shape) == tuple((x + 15) // 16 * 16 for x in view_shape) + assert len(root.alloc_buffers) == 1 + assert not root.annotations["layout_map"][read.buffer].is_equal(root.annotations["layout_map"][write.buffer]) + tvm.ir.assert_structural_equal(after, transform.NormalizeAscendFractalStorage()(after)) + + +def test_cube_views_with_different_padded_sizes_are_rejected(): + owner = tirx.decl_buffer((15, 16), "bfloat16", scope="shared.l0b") + view = tirx.decl_buffer((8, 30), "bfloat16", data=owner.data) + block = tirx.SBlock( + [], + [], + [], + "root", + tirx.Evaluate(view[0, 0]), + alloc_buffers=[owner], + annotations={ + "layout_map": {owner: tilelang.layout.make_ascend_major_k_layout(owner), view: tilelang.layout.make_ascend_major_k_layout(view)} + }, + ) + before = tvm.IRModule({"main": tirx.PrimFunc([], tirx.SBlockRealize([], True, block))}) + with pytest.raises(tvm.error.InternalError, match="equal padded physical sizes"): + transform.NormalizeAscendFractalStorage()(before) + + +@pytest.mark.parametrize("raw", [False, True], ids=["typed-view", "ambiguous-data-pointer"]) +def test_distinct_cube_layouts_require_typed_pointers(raw): + owner = tirx.decl_buffer((64, 128), "bfloat16", scope="shared.l0b") + view = tirx.decl_buffer((128, 64), "bfloat16", data=owner.data) + pointer = owner.access_ptr("r", offset=64, extent=1) if raw else T.access_ptr(view[1, 0], "r", extent=1) + before = kernel( + tirx.Evaluate(tirx.call_extern("int32", "consume", pointer)), + buffers=[owner], + annotations={ + "layout_map": {owner: tilelang.layout.make_ascend_major_k_layout(owner), view: tilelang.layout.make_ascend_major_k_layout(view)} + }, + ) + if raw: + with pytest.raises(tvm.error.InternalError, match="Raw Cube pointers cannot select among distinct view layouts"): + transform.NormalizeAscendFractalStorage()(before) + else: + after = transform.NormalizeAscendFractalStorage()(before) + (pointer,) = calls(after, "tl.access_ptr") + assert tuple(int(x) for x in pointer.args[0].buffer.shape) == (128, 64) + assert tuple(int(x) for x in pointer.args[0].indices) == (1, 0) + + def test_non_rectangular_alias_region_is_rejected(): # Three logical elements in different rows cannot become an 11-row DMA. with pytest.raises(tvm.error.InternalError, match="not a contiguous rectangular matrix"): diff --git a/tilelang/ascend/language/__init__.py b/tilelang/ascend/language/__init__.py index 423d829d51..3bd88062d3 100644 --- a/tilelang/ascend/language/__init__.py +++ b/tilelang/ascend/language/__init__.py @@ -18,6 +18,7 @@ alloc_l0c, ) from .annotations import ( # noqa: F401 + annotate_layout, annotate_buffer_versions, annotate_manual_multi_buffer, annotate_unlimit_memory, diff --git a/tilelang/ascend/language/annotations.py b/tilelang/ascend/language/annotations.py index d53216bb86..5bde41391d 100644 --- a/tilelang/ascend/language/annotations.py +++ b/tilelang/ascend/language/annotations.py @@ -4,9 +4,12 @@ from tvm.tirx import IntImm from tvm.tirx.script.builder.ir import sblock_attr +from tilelang.language.annotations import _normalize_layout_map + _VALID_MEMORY_SCOPES = ("shared", "shared.l1", "shared.l0c", "shared.l0a", "shared.l0b") __all__ = [ + "annotate_layout", "annotate_buffer_versions", "annotate_manual_multi_buffer", "annotate_unlimit_memory", @@ -16,6 +19,17 @@ _VALID_BUFFER_VERSION_MODES = ("auto", "iteration", "counter") +def annotate_layout(layout_map: dict): + """Annotate each Buffer view, preserving aliases with distinct Cube layouts. + + Views share storage, not a layout conversion. A producer must write the + representation consumed through each view; annotating a view never moves + or transposes existing values. Canonical Cube views sharing storage must + have the same dtype and padded physical size. + """ + return sblock_attr({"layout_map": _normalize_layout_map(layout_map)}) + + def annotate_buffer_versions(buffer_versions_map: dict): """Control automatic multi-buffer version counts and indexing modes. diff --git a/tilelang/ascend/op/gemm/gemm_mad.py b/tilelang/ascend/op/gemm/gemm_mad.py index 781b2731aa..aeda0808ed 100644 --- a/tilelang/ascend/op/gemm/gemm_mad.py +++ b/tilelang/ascend/op/gemm/gemm_mad.py @@ -20,29 +20,6 @@ def _to_int32(value): return tirx.IntImm("int32", value) if isinstance(value, int) else value -def _compute_flat_offset(buf, region): - offset = tirx.const(0, "int32") - ndim = len(region.region) - for i, rng in enumerate(region.region): - stride = tirx.const(1, "int32") - for j in range(i + 1, ndim): - stride = stride * buf.shape[j] - offset = offset + tirx.Cast("int32", rng.min) * stride - return offset - - -def _compute_flat_offset_excluding_last(buf, region): - offset = tirx.const(0, "int32") - ndim = len(region.region) - for i in range(ndim - 1): - rng = region.region[i] - stride = tirx.const(1, "int32") - for j in range(i + 1, len(buf.shape)): - stride = stride * buf.shape[j] - offset = offset + tirx.Cast("int32", rng.min) * stride - return offset - - def _compute_extent(region): extent = tirx.const(1, "int32") for rng in region.region: @@ -50,11 +27,14 @@ def _compute_extent(region): return extent -def _make_access_ptr(buf, region, rw_mask): - return tirx.op.tvm_access_ptr( - tirx.op.type_annotation(buf.dtype), - buf.data, - _compute_flat_offset(buf, region), +def _make_access_ptr(buf, region, rw_mask, exclude_last=False): + indices = [tirx.Cast("int32", rng.min) for rng in region.region] + if exclude_last: + indices[-1] = tirx.const(0, "int32") + return tirx.call_intrin( + "handle", + tirx.op.Op.get("tl.access_ptr"), + tirx.BufferLoad(buf, indices), _compute_extent(region), tirx.IntImm("int32", rw_mask), ) @@ -63,7 +43,7 @@ def _make_access_ptr(buf, region, rw_mask): def _find_layout(layout_map, buf): if layout_map is None: return None - for key in (buf.data if hasattr(buf, "data") else None, buf): + for key in (buf, buf.data if hasattr(buf, "data") else None): if key is not None: layout = layout_map.get(key) if layout is not None: @@ -251,22 +231,8 @@ def _lower_l1(self): tile_k_sub = self._compute_tile_k_sub() c_ptr = _make_access_ptr(self.CRegion.buffer, self.CRegion, 2) - a_buf = self.ARegion.buffer - b_buf = self.BRegion.buffer - a_ptr = tirx.op.tvm_access_ptr( - tirx.op.type_annotation(a_buf.dtype), - a_buf.data, - _compute_flat_offset_excluding_last(a_buf, self.ARegion), - _compute_extent(self.ARegion), - tirx.IntImm("int32", 1), - ) - b_ptr = tirx.op.tvm_access_ptr( - tirx.op.type_annotation(b_buf.dtype), - b_buf.data, - _compute_flat_offset_excluding_last(b_buf, self.BRegion), - _compute_extent(self.BRegion), - tirx.IntImm("int32", 1), - ) + a_ptr = _make_access_ptr(self.ARegion.buffer, self.ARegion, 1, exclude_last=True) + b_ptr = _make_access_ptr(self.BRegion.buffer, self.BRegion, 1, exclude_last=True) input_dtype = self._input_dtype() dtype_str = _ASCEND_DTYPE_MAP.get(input_dtype) diff --git a/tilelang/ascend/op/gemm/gemm_mad_blockscaled.py b/tilelang/ascend/op/gemm/gemm_mad_blockscaled.py index 2661a0ac79..5e0fd8a91c 100644 --- a/tilelang/ascend/op/gemm/gemm_mad_blockscaled.py +++ b/tilelang/ascend/op/gemm/gemm_mad_blockscaled.py @@ -13,8 +13,6 @@ from .gemm_mad import ( _ASCEND_DTYPE_MAP, - _compute_extent, - _compute_flat_offset_excluding_last, _find_layout, _make_access_ptr, _to_int32, @@ -102,41 +100,15 @@ def _lower_l1_blockscaled(self): tile_k_sub = self._compute_tile_k_sub() c_ptr = _make_access_ptr(self.CRegion.buffer, self.CRegion, 2) - a_buf = self.ARegion.buffer - b_buf = self.BRegion.buffer - a_ptr = tirx.op.tvm_access_ptr( - tirx.op.type_annotation(a_buf.dtype), - a_buf.data, - _compute_flat_offset_excluding_last(a_buf, self.ARegion), - _compute_extent(self.ARegion), - tirx.IntImm("int32", 1), - ) - b_ptr = tirx.op.tvm_access_ptr( - tirx.op.type_annotation(b_buf.dtype), - b_buf.data, - _compute_flat_offset_excluding_last(b_buf, self.BRegion), - _compute_extent(self.BRegion), - tirx.IntImm("int32", 1), - ) + a_ptr = _make_access_ptr(self.ARegion.buffer, self.ARegion, 1, exclude_last=True) + b_ptr = _make_access_ptr(self.BRegion.buffer, self.BRegion, 1, exclude_last=True) sfa_region = self.SFARegion sfb_region = self.SFBRegion sfa_buf = sfa_region.buffer sfb_buf = sfb_region.buffer - sfa_ptr = tirx.op.tvm_access_ptr( - tirx.op.type_annotation(sfa_buf.dtype), - sfa_buf.data, - _compute_flat_offset_excluding_last(sfa_buf, sfa_region), - _compute_extent(sfa_region), - tirx.IntImm("int32", 1), - ) - sfb_ptr = tirx.op.tvm_access_ptr( - tirx.op.type_annotation(sfb_buf.dtype), - sfb_buf.data, - _compute_flat_offset_excluding_last(sfb_buf, sfb_region), - _compute_extent(sfb_region), - tirx.IntImm("int32", 1), - ) + sfa_ptr = _make_access_ptr(sfa_buf, sfa_region, 1, exclude_last=True) + sfb_ptr = _make_access_ptr(sfb_buf, sfb_region, 1, exclude_last=True) input_dtype = self._input_dtype() in_dtype_str = _ASCEND_DTYPE_MAP.get(input_dtype) diff --git a/tilelang/language/annotations.py b/tilelang/language/annotations.py index 65965637ec..9d72b9b6bc 100644 --- a/tilelang/language/annotations.py +++ b/tilelang/language/annotations.py @@ -42,8 +42,7 @@ def use_swizzle(panel_size: int, order: str = "row", enable: bool = True): return attr(None, "threadblock_swizzle_pattern", tvm_tuple(device_func, panel_size)) -def annotate_layout(layout_map: dict): - """Annotate the layout of the buffer.""" +def _normalize_layout_map(layout_map: dict): _layout_map = {} for buffer, layout in layout_map.items(): if is_reducer(buffer): @@ -57,13 +56,18 @@ def annotate_layout(layout_map: dict): elif is_fragment(buffer): assert isinstance(layout, Fragment), f"for Fragment {buffer}, layout must be a Fragment, but got {type(layout)}" if isinstance(layout, Layout): - _layout_map[buffer.data] = layout + _layout_map[buffer] = layout elif isinstance(layout, Callable): - _layout_map[buffer.data] = Layout(buffer.shape, layout) + _layout_map[buffer] = Layout(buffer.shape, layout) else: raise ValueError(f"Invalid layout: {layout}") - return sblock_attr({"layout_map": _layout_map}) + return _layout_map + + +def annotate_layout(layout_map: dict): + """Annotate the layout of the buffer.""" + return sblock_attr({"layout_map": {buffer.data: layout for buffer, layout in _normalize_layout_map(layout_map).items()}}) def annotate_safe_value(safe_value_map: dict): From d477048d34c1f55d875a8726a007413f9b4003cf Mon Sep 17 00:00:00 2001 From: Denver Jin Date: Thu, 22 Oct 2026 15:51:24 +0800 Subject: [PATCH 6/7] [Ascend] Limit shared-layout support to buffer views --- .../references/synchronization.md | 15 - .../annotate_multi_buffer_eligible.cc | 2 +- src/ascend/transform/auto_schedule.cc | 197 +---------- .../auto_schedule/buffer_alias_analysis.h | 38 +-- .../auto_schedule/dependency_analysis.cc | 35 +- .../auto_schedule/dependency_analysis.h | 6 +- src/ascend/transform/buffer_alias.h | 3 +- src/ascend/transform/buffer_version.h | 45 +-- src/ascend/transform/insert_sync.cc | 12 +- .../transform/materialize_multi_buffer.cc | 4 +- src/ascend/transform/prepare_multi_buffer.cc | 2 +- src/ascend/transform/resolve_core.cc | 9 +- .../auto_schedule/test_conflict_hints.py | 43 +-- .../test_conflict_ring_owners.py | 311 ------------------ .../auto_schedule/test_cube_storage_views.py | 153 +++++++++ .../auto_schedule/test_reuse_integration.py | 78 +---- tilelang/ascend/language/schedule_hint.py | 15 +- 17 files changed, 202 insertions(+), 766 deletions(-) delete mode 100644 testing/ascend/auto_schedule/test_conflict_ring_owners.py create mode 100644 testing/ascend/auto_schedule/test_cube_storage_views.py diff --git a/.agents/skills/tilelang-ascend/references/synchronization.md b/.agents/skills/tilelang-ascend/references/synchronization.md index 2b7d9a326e..b8d1cc51e9 100644 --- a/.agents/skills/tilelang-ascend/references/synchronization.md +++ b/.agents/skills/tilelang-ascend/references/synchronization.md @@ -326,21 +326,6 @@ to be treated as conflicting while retaining the usual RAW/WAR/WAW rules; read-only pairs still have no dependency. If contradictory hints match, conflict wins. See `tilelang/ascend/language/schedule_hint.py` for the API. -A positive conflict hint also explicitly permits the two backing storages to -reuse memory, even when lifetime inference cannot prove that reuse. This -permission is pairwise, with no transitive closure. The region, `level`, and -`cross` qualifiers continue to govern dependency matching; reuse permission is -for the storage pair. The caller must cover all accesses affected by reuse. -Allocation scope, alignment, and the shared-memory-reuse disable option still -apply. Buffer identities and layouts remain independent. An isolated, -equal-sized L0B pair with matching fixed automatic version counts can share a -ring clock when a whole-buffer hint uses the default iteration matching and all -accesses belong to disjoint automatic owners. Explicit iteration indexing, manual -versions, strides, and offsets keep the ordinary dependency path. -Each owner's accesses must also be guaranteed to use a single schedule stage. -Cross-stage manual accesses and automatic scheduling with `enable_offset=True` -keep independent rings because the shared counter cannot span those stages. - ## Trace the complete storage protocol For a suspicious access or wait, build a table with: diff --git a/src/ascend/transform/annotate_multi_buffer_eligible.cc b/src/ascend/transform/annotate_multi_buffer_eligible.cc index 5cd44cef63..4ce5a0f4b4 100644 --- a/src/ascend/transform/annotate_multi_buffer_eligible.cc +++ b/src/ascend/transform/annotate_multi_buffer_eligible.cc @@ -808,7 +808,7 @@ tvm::transform::Pass AnnotateMultiBufferEligible() { [](const TilelangKernelContext &context) { ScheduledTIR scheduled_tir = DecodeScheduledTIR(context.root, context.outer_ctx); - L0StorageGroups groups = CollectL0StorageGroups(context.root); + L0StorageGroups groups(CollectL0SFBindings(context.root)); auto &metadata = scheduled_tir.metadata; metadata.buffer_versions = ExpandL0StorageGroupValues( metadata.buffer_versions, groups, "version counts"); diff --git a/src/ascend/transform/auto_schedule.cc b/src/ascend/transform/auto_schedule.cc index cff87f17f5..1687d6b936 100644 --- a/src/ascend/transform/auto_schedule.cc +++ b/src/ascend/transform/auto_schedule.cc @@ -40,7 +40,6 @@ #include #include -#include #include #include #include @@ -66,7 +65,6 @@ #include "ascend/op/utils.h" #include "ascend/transform/buffer_version.h" #include "backend/common/target_utils.h" -#include "layout/layout.h" #include "op/utils.h" #include "runtime/thread_storage_scope.h" #include "support/check.h" @@ -585,7 +583,7 @@ std::vector> ScheduleBuilder::Z3SchedulePython( // Buffer-region data dependencies auto deps = AnalyzeDependencies( nodes, /*loop=*/nullptr, manual_buffer_versions_, multi_buffer_owners_, - &dependency_cache_, root_conflict_hints_, storage_groups_); + &dependency_cache_, root_conflict_hints_); std::unordered_map node_idx; for (size_t i = 0; i < n; ++i) { node_idx[nodes[i]] = i; @@ -989,7 +987,7 @@ void ScheduleBuilder::Z3SchedulePythonLoop(ControlNode *ctrl, // Buffer-region data dependencies auto deps = AnalyzeDependencies(nodes, ctrl, manual_buffer_versions_, multi_buffer_owners_, &dependency_cache_, - root_conflict_hints_, storage_groups_); + root_conflict_hints_); std::unordered_map node_idx; for (size_t i = 0; i < n; ++i) { node_idx[nodes[i]] = i; @@ -1353,194 +1351,6 @@ void ScheduleBuilder::Z3SchedulePythonLoop(ControlNode *ctrl, // Pass orchestration. -// Select an isolated, equal-sized L0B pair whose complete automatic epochs -// can use one ring clock. Other conflict hints retain their dependency-only -// interpretation during scheduling. In particular, do not take a transitive -// closure of pairwise reuse permissions. -static void PlanExplicitL0BReuseRing(ScheduledTIR *scheduled_tir) { - auto &metadata = scheduled_tir->metadata; - if (metadata.root_conflict_hints.empty() || - metadata.kernel_root->annotations.count(kReuseRingBindings)) - return; - - class L0BCollector : public StmtExprVisitor { - public: - std::vector buffers; - - private: - std::unordered_set seen; - void Record(const Buffer &buffer) { - if (IsL0BBuffer(buffer) && seen.insert(buffer->data).second) - buffers.push_back(buffer); - } - void VisitStmt_(const SBlockNode *op) final { - for (const Buffer &buffer : op->alloc_buffers) - Record(buffer); - StmtExprVisitor::VisitStmt_(op); - } - void VisitStmt_(const AllocBufferNode *op) final { Record(op->buffer); } - } collector; - collector(metadata.kernel_root); - if (collector.buffers.size() != 2) - return; - const Buffer &a = collector.buffers[0]; - const Buffer &b = collector.buffers[1]; - if (a->dtype != b->dtype || !a->strides.empty() || !b->strides.empty() || - !is_zero(a->elem_offset) || !is_zero(b->elem_offset)) - return; - auto a_versions = metadata.buffer_versions.Get(a->data); - auto b_versions = metadata.buffer_versions.Get(b->data); - if (!a_versions.has_value() || !b_versions.has_value() || - a_versions.value() <= 1 || a_versions.value() != b_versions.value() || - metadata.manual_buffer_versions.count(a->data) || - metadata.manual_buffer_versions.count(b->data)) - return; - for (const Var &storage : {a->data, b->data}) { - if (auto mode = metadata.buffer_version_modes.Get(storage)) { - if (mode.value() == "iteration") - return; - } - } - - bool declared = false; - for (const ffi::Any &entry : metadata.root_conflict_hints) { - auto hint = entry.cast>(); - if (!hint[3].cast()->value || hint[2].cast()->value != -1) - continue; - auto lhs = hint[0].as(); - auto rhs = hint[1].as(); - if (lhs.has_value() && rhs.has_value() && - ((lhs.value().same_as(a->data) && rhs.value().same_as(b->data)) || - (lhs.value().same_as(b->data) && rhs.value().same_as(a->data)))) - declared = true; - } - if (!declared) - return; - - Map layouts; - if (auto value = - metadata.kernel_root->annotations.Get(tvm::tl::attr::kLayoutMap)) - layouts = value.value().cast>(); - auto storage_bytes = [&](const Buffer &buffer) -> int64_t { - Array shape = buffer->shape; - bool has_layout = false; - bool found_layout = false; - for (const auto &[candidate, layout] : layouts) { - if (candidate->data.same_as(buffer->data)) { - has_layout = true; - if (candidate->dtype == buffer->dtype && candidate->strides.empty() && - is_zero(candidate->elem_offset) && - ffi::StructuralEqual()(candidate->shape, buffer->shape)) { - shape = layout->OutputShape(); - found_layout = true; - break; - } - } - } - if (has_layout && !found_layout) - return 0; - PrimExpr bits = - IntImm(DataType::Int(64), buffer->dtype.bits() * buffer->dtype.lanes()); - for (const PrimExpr &extent : shape) - bits *= extent; - arith::Analyzer analyzer; - PrimExpr bytes = analyzer.Simplify(floordiv(bits + 7, 8)); - const auto *value = bytes.as(); - return value != nullptr ? value->value : 0; - }; - int64_t bytes = storage_bytes(a); - if (bytes <= 0 || bytes != storage_bytes(b)) - return; - - L0StorageGroups existing = CollectL0StorageGroups(metadata.kernel_root); - if (existing.Members(a->data).size() != 1 || - existing.Members(b->data).size() != 1) - return; - MultiBufferOwnerMap owners = - CollectMultiBufferOwners(scheduled_tir->tree, existing); - auto a_owners = owners.find(a->data); - auto b_owners = owners.find(b->data); - if (a_owners == owners.end() || b_owners == owners.end()) - return; - for (ControlNode *a_owner : a_owners->second) { - for (ControlNode *b_owner : b_owners->second) { - if (a_owner == b_owner || a_owner->IsWithin(b_owner) || - b_owner->IsWithin(a_owner)) - return; - } - } - std::vector tasks; - CollectAllTaskNodes(scheduled_tir->tree, tasks); - for (const Var &storage : {a->data, b->data}) { - auto found = owners.find(storage); - if (found == owners.end() || found->second.empty()) - return; - // Joining disjoint owners requires a counter. PrepareMultiBuffer can - // advance that counter only when each owner's storage accesses share a - // stage, so keep independent rings unless scheduling preserves this. - for (ControlNode *owner : found->second) { - std::optional storage_stage; - for (const auto &child : owner->children) { - if (!child->TouchesStorage(storage)) - continue; - int stage = child->GetStage(); - if (storage_stage.has_value() && storage_stage.value() != stage) - return; - storage_stage = stage; - } - if (!storage_stage.has_value()) - return; - if (storage_stage.value() == kUnscheduledStage) { - // Automatic offsets are not known yet. Without offsets, scheduling - // collapses all automatic tasks to stage zero. - auto enable_offset = owner->control->annotations.Get("enable_offset"); - if (enable_offset.has_value() && enable_offset.value().cast()) - return; - } - } - for (TaskNode *task : tasks) { - auto touches = [&](const auto ®ions) { - return std::any_of(regions.begin(), regions.end(), - [&](const BufferRegion ®ion) { - return region->buffer->data.same_as(storage); - }); - }; - if (!touches(task->GetReadRegions()) && !touches(task->GetWriteRegions())) - continue; - if (std::count_if( - found->second.begin(), found->second.end(), - [&](ControlNode *owner) { return task->IsWithin(owner); }) != 1) - return; - } - } - - ReuseRingBindings bindings{{b->data, a->data}}; - metadata.kernel_root.CopyOnWrite()->annotations.Set(kReuseRingBindings, - bindings); - L0StorageGroups groups = CollectL0StorageGroups(metadata.kernel_root); - std::function expand = [&](IRStructure *node) { - if (!node->IsControl()) - return; - auto *control = static_cast(node); - if (auto value = control->control->annotations.Get(kMultiBufferEligible)) { - Array eligible; - std::unordered_set seen; - for (const Var &storage : value.value().cast>()) { - for (const Var &member : groups.Members(storage)) { - if (seen.insert(member).second) - eligible.push_back(member); - } - } - control->control.CopyOnWrite()->annotations.Set(kMultiBufferEligible, - eligible); - } - for (const auto &child : control->children) - expand(child.get()); - }; - for (const auto &node : scheduled_tir->tree) - expand(node.get()); -} - static void WarnOnUnappliedBufferVersionOverrides( const BufferVersionMap &buffer_version_overrides, const BufferVersionMap &selected_buffer_versions) { @@ -1576,8 +1386,7 @@ static void ScheduleSingleKernel(const Stmt &kernel_body, Target target, // Schedule IRStructure in place with the Z3 scheduler. ScheduleBuilder unit_builder; - PlanExplicitL0BReuseRing(scheduled_tir); - L0StorageGroups groups = CollectL0StorageGroups(metadata.kernel_root); + L0StorageGroups groups(CollectL0SFBindings(metadata.kernel_root)); metadata.buffer_versions = ExpandL0StorageGroupValues( metadata.buffer_versions, groups, "version counts"); metadata.manual_buffer_versions = ExpandL0StorageGroupValues( diff --git a/src/ascend/transform/auto_schedule/buffer_alias_analysis.h b/src/ascend/transform/auto_schedule/buffer_alias_analysis.h index 85468cc271..bb7e8c987d 100644 --- a/src/ascend/transform/auto_schedule/buffer_alias_analysis.h +++ b/src/ascend/transform/auto_schedule/buffer_alias_analysis.h @@ -41,7 +41,6 @@ #include "../buffer_alias.h" #include "./buffer_lifetime_analysis.h" -#include "./dependency_analysis.h" #include "./ir_structure.h" #include "./memory_detector.h" #include "ascend/op/utils.h" @@ -407,41 +406,15 @@ class BufferAliasAnalyzerImpl { } } - void Analyze(const std::vector> &root, - const ConflictHintList &root_conflicts) { + void Analyze(const std::vector> &root) { lifetime_summaries_.clear(); AnalyzeScope(root, /*loop=*/nullptr); - AddConflictAliases_(root_conflicts); } const BufferAliasMap &GetAliases() const { return buffer_aliases_; } void Validate() const { ValidateBufferAliasMap(buffer_aliases_); } private: - void AddConflictAliases_(const ConflictHintList &hints) { - auto storage = [](const ffi::Any &operand) -> Var { - if (auto var = operand.as()) - return var.value(); - return operand.cast()->buffer->data; - }; - for (const ffi::Any &entry : hints) { - Array hint = entry.cast>(); - ICHECK_EQ(hint.size(), 4U); - if (!hint[3].cast()->value) - continue; - Var lhs = storage(hint[0]); - Var rhs = storage(hint[1]); - if (!groups_.SharesReuseRing(lhs, rhs)) { - lhs = groups_.Representative(lhs); - rhs = groups_.Representative(rhs); - } - // A positive hint explicitly permits storage reuse. Its region and - // iteration qualifiers still control dependency analysis independently. - if (!lhs.same_as(rhs)) - AddBufferAlias(&buffer_aliases_, lhs, rhs); - } - } - HappensBeforeQuery MakeHappensBeforeQuery(ControlNode *scope) const { return [this, scope](size_t src, size_t dst, int64_t distance) { return context_->HappensBefore(scope, src, dst, distance); @@ -456,10 +429,6 @@ class BufferAliasAnalyzerImpl { void AnalyzeScope(const std::vector> &nodes, ControlNode *loop) { - if (loop != nullptr) { - if (auto hints = loop->control->annotations.Get("conflict_hint")) - AddConflictAliases_(hints.value().cast()); - } // Cache every child scope before summarizing this scope. A parent can then // promote the child's already-built lifetime instead of recursively // rediscovering the same nested control chain. @@ -1531,9 +1500,8 @@ class BufferAliasAnalyzer { void Collect(const std::vector> &root) { impl_.Collect(root); } - void Analyze(const std::vector> &root, - const ConflictHintList &root_conflicts = {}) { - impl_.Analyze(root, root_conflicts); + void Analyze(const std::vector> &root) { + impl_.Analyze(root); } const BufferAliasMap &GetAliases() const { return impl_.GetAliases(); } diff --git a/src/ascend/transform/auto_schedule/dependency_analysis.cc b/src/ascend/transform/auto_schedule/dependency_analysis.cc index 5f18f6751f..3bfa1efb9d 100644 --- a/src/ascend/transform/auto_schedule/dependency_analysis.cc +++ b/src/ascend/transform/auto_schedule/dependency_analysis.cc @@ -120,8 +120,8 @@ bool RegionsMayConflict(const ConstrSet &a_ctx, const BufferRegion &a_region, ControlNode *loop, int offset, size_t num_storage_owners, const ConflictHintList &root_conflicts, - Optional a_pitch, Optional b_pitch, - bool shared_reuse_ring) { + Optional a_pitch, + Optional b_pitch) { const bool cross = offset != 0; if (cross && loop == nullptr) return false; @@ -147,8 +147,7 @@ bool RegionsMayConflict(const ConstrSet &a_ctx, const BufferRegion &a_region, if (declared_no_conflict) return false; - if (!shared_reuse_ring && - !a_region->buffer->data.same_as(b_region->buffer->data)) + if (!a_region->buffer->data.same_as(b_region->buffer->data)) return false; // Same storage but different logical buffer (T.view/T.reshape): regions may @@ -256,8 +255,7 @@ AnalyzeDependencies(std::vector nodes, ControlNode *loop, const BufferVersionMap &manual_buffer_versions, const MultiBufferOwnerMap &multi_buffer_owners, DependencyCache *dependency_cache, - const ConflictHintList &root_conflicts, - const L0StorageGroups &storage_groups) { + const ConflictHintList &root_conflicts) { DependencyCache local_cache; DependencyCache &cache = dependency_cache ? *dependency_cache : local_cache; if (loop) { @@ -301,9 +299,9 @@ AnalyzeDependencies(std::vector nodes, ControlNode *loop, std::vector children; for (const auto &child : ctrl->children) children.push_back(child.get()); - auto child_deps = AnalyzeDependencies( - children, ctrl, manual_buffer_versions, multi_buffer_owners, &cache, - root_conflicts, storage_groups); + auto child_deps = + AnalyzeDependencies(children, ctrl, manual_buffer_versions, + multi_buffer_owners, &cache, root_conflicts); for (const auto &cdep : child_deps) { auto &covered_pairs = out[cdep.storage]; covered_pairs.insert(cdep.task_pairs.begin(), cdep.task_pairs.end()); @@ -409,12 +407,8 @@ AnalyzeDependencies(std::vector nodes, ControlNode *loop, const Var &lhs_storage = lhs.region->buffer->data; const Var &rhs_storage = rhs.region->buffer->data; ffi::Optional storage; - bool shared_reuse_ring = - storage_groups.SharesReuseRing(lhs_storage, rhs_storage); if (lhs_storage.same_as(rhs_storage)) storage = lhs_storage; - else if (shared_reuse_ring) - storage = storage_groups.Representative(lhs_storage); bool owner_exclusion = is_owner_exclusion(storage, lhs.task, rhs.task); // The i nodes, ControlNode *loop, rhs.task->outer_ctx, rhs.region, loop, d, /*num_storage_owners=*/0, root_conflicts, lhs.task->GetL0Pitch(lhs_storage), - rhs.task->GetL0Pitch(rhs_storage), - shared_reuse_ring)) { + rhs.task->GetL0Pitch(rhs_storage))) { add_pair(storage, d, lhs.task, rhs.task); break; } @@ -442,12 +435,12 @@ AnalyzeDependencies(std::vector nodes, ControlNode *loop, continue; if (distance < 0 && loop == nullptr) continue; - if (RegionsMayConflict( - lhs.task->outer_ctx, lhs.region, rhs.task->outer_ctx, - rhs.region, loop, distance, - get_num_storage_owners(storage), root_conflicts, - lhs.task->GetL0Pitch(lhs_storage), - rhs.task->GetL0Pitch(rhs_storage), shared_reuse_ring)) { + if (RegionsMayConflict(lhs.task->outer_ctx, lhs.region, + rhs.task->outer_ctx, rhs.region, loop, + distance, get_num_storage_owners(storage), + root_conflicts, + lhs.task->GetL0Pitch(lhs_storage), + rhs.task->GetL0Pitch(rhs_storage))) { add_pair(storage, owner_exclusion ? 0 : distance, lhs.task, rhs.task, owner_exclusion ? DependencyKind::kOwnerExclusion diff --git a/src/ascend/transform/auto_schedule/dependency_analysis.h b/src/ascend/transform/auto_schedule/dependency_analysis.h index d3054ab71e..6eb96f486d 100644 --- a/src/ascend/transform/auto_schedule/dependency_analysis.h +++ b/src/ascend/transform/auto_schedule/dependency_analysis.h @@ -106,8 +106,7 @@ AnalyzeDependencies(std::vector nodes, const BufferVersionMap &manual_buffer_versions = {}, const MultiBufferOwnerMap &multi_buffer_owners = {}, DependencyCache *dependency_cache = nullptr, - const ConflictHintList &root_conflicts = {}, - const L0StorageGroups &storage_groups = L0StorageGroups()); + const ConflictHintList &root_conflicts = {}); // Check whether two task regions may overlap under their respective symbolic // contexts. Only a proof of disjointness removes a conflict. @@ -129,8 +128,7 @@ bool RegionsMayConflict(const ConstrSet &a_ctx, const BufferRegion &a_region, size_t num_storage_owners = 0, const ConflictHintList &root_conflicts = {}, Optional a_pitch = std::nullopt, - Optional b_pitch = std::nullopt, - bool shared_reuse_ring = false); + Optional b_pitch = std::nullopt); } // namespace ascend } // namespace tl diff --git a/src/ascend/transform/buffer_alias.h b/src/ascend/transform/buffer_alias.h index 768db66adf..cbeb34b731 100644 --- a/src/ascend/transform/buffer_alias.h +++ b/src/ascend/transform/buffer_alias.h @@ -47,8 +47,7 @@ constexpr const char *kAivStorageMap = "tl.aiv_storage_map"; constexpr const char *kSplitSubcores = "tl.split_subcores"; using AivStorageMap = ffi::Map; -/*! \brief Storage data Var -> inferred or explicitly permitted alias partners. - */ +/*! \brief Storage data Var -> storages proven safe to alias with it. */ using BufferAliasMap = ffi::Map>; inline void AddBufferAliasStorage(BufferAliasMap *aliases, diff --git a/src/ascend/transform/buffer_version.h b/src/ascend/transform/buffer_version.h index c070f9e0dc..641bc0fcf7 100644 --- a/src/ascend/transform/buffer_version.h +++ b/src/ascend/transform/buffer_version.h @@ -24,10 +24,8 @@ #pragma once -#include #include #include -#include #include "ascend/op/utils.h" #include "transform/common/attr.h" @@ -48,56 +46,31 @@ constexpr const char *kMultiBufferEligible = "multi_buffer_eligible"; /*! \brief Storage data Var -> positive compile-time version count. */ using BufferVersionMap = ffi::Map; -// Selected physical L0B reuse pairs. Layouts and logical buffers remain -// distinct, but the selected pair advances one storage-ring protocol. -constexpr const char *kReuseRingBindings = "tl.reuse_ring_bindings"; -using ReuseRingBindings = ffi::Map; - // Logical reads/writes remain separate. Ownership, physical versions and // address reuse use the same allocation group, represented by the data Var. class L0StorageGroups { public: - explicit L0StorageGroups(L0SFBindings bindings = {}, - ReuseRingBindings reuse_bindings = {}) - : sf_bindings_(std::move(bindings)), - reuse_bindings_(std::move(reuse_bindings)), bindings_(sf_bindings_) { - for (const auto &[member, data] : reuse_bindings_) - bindings_.Set(member, data); - for (const auto &[member, data] : bindings_) { - auto members = members_.Get(data).value_or(ffi::Array{data}); - members.push_back(member); - members_.Set(data, members); - } + explicit L0StorageGroups(L0SFBindings bindings = {}) + : bindings_(std::move(bindings)) { + for (const auto &[sf, data] : bindings_) + scales_.Set(data, sf); } tirx::Var Representative(const tirx::Var &storage) const { return bindings_.Get(storage).value_or(storage); } ffi::Array Members(const tirx::Var &storage) const { tirx::Var data = Representative(storage); - return members_.Get(data).value_or(ffi::Array{data}); - } - const L0SFBindings &Bindings() const { return sf_bindings_; } - bool SharesReuseRing(const tirx::Var &lhs, const tirx::Var &rhs) const { - auto lhs_data = reuse_bindings_.Get(lhs); - auto rhs_data = reuse_bindings_.Get(rhs); - return (lhs_data.has_value() && lhs_data.value().same_as(rhs)) || - (rhs_data.has_value() && rhs_data.value().same_as(lhs)); + if (auto sf = scales_.Get(data)) + return {data, sf.value()}; + return {data}; } + const L0SFBindings &Bindings() const { return bindings_; } private: - L0SFBindings sf_bindings_; - ReuseRingBindings reuse_bindings_; L0SFBindings bindings_; - ffi::Map> members_; + L0SFBindings scales_; }; -inline L0StorageGroups CollectL0StorageGroups(const tirx::SBlock &root) { - ReuseRingBindings reuse; - if (auto value = root->annotations.Get(kReuseRingBindings)) - reuse = value.value().cast(); - return L0StorageGroups(CollectL0SFBindings(root), std::move(reuse)); -} - // User overrides and prepared metadata must describe the entire group. // Contradictory member requests are errors, rather than last-writer wins. template diff --git a/src/ascend/transform/insert_sync.cc b/src/ascend/transform/insert_sync.cc index 98317c3a06..1ded68e606 100644 --- a/src/ascend/transform/insert_sync.cc +++ b/src/ascend/transform/insert_sync.cc @@ -831,7 +831,7 @@ class SyncAnalyzer : public BufferAliasAnalysisContext { BufferAliasAnalyzer buffer_alias_analyzer(*this, storage_groups_); buffer_alias_analyzer.Collect(root); - buffer_alias_analyzer.Analyze(root, root_conflicts_); + buffer_alias_analyzer.Analyze(root); buffer_alias_analyzer.Validate(); buffer_aliases_ = buffer_alias_analyzer.GetAliases(); } @@ -1250,9 +1250,9 @@ class SyncAnalyzer : public BufferAliasAnalysisContext { std::vector CollectSyncPoints(const std::vector &nodes, ControlNode *loop) { - std::vector deps = AnalyzeDependencies( - nodes, loop, manual_, multi_buffer_owners_, &dependency_cache_, - root_conflicts_, storage_groups_); + std::vector deps = + AnalyzeDependencies(nodes, loop, manual_, multi_buffer_owners_, + &dependency_cache_, root_conflicts_); if (loop != nullptr) { std::vector directed_deps; directed_deps.reserve(deps.size() * 2); @@ -2532,8 +2532,8 @@ SBlock InsertKernelSync(const Stmt &kernel_body, ScheduledTIR scheduled_tir) { std::vector> &root = scheduled_tir.tree; int num_aiv_subcores = scheduled_tir.metadata.num_aiv_subcores.value_or(2); CoreMask vector_mask = VectorCoreMask(num_aiv_subcores); - L0StorageGroups groups = - CollectL0StorageGroups(scheduled_tir.metadata.kernel_root); + L0StorageGroups groups( + CollectL0SFBindings(scheduled_tir.metadata.kernel_root)); MultiBufferPlan multi_buffer_plan = ReadMultiBufferPlan(root, scheduled_tir.metadata.buffer_versions, groups); std::vector tasks; diff --git a/src/ascend/transform/materialize_multi_buffer.cc b/src/ascend/transform/materialize_multi_buffer.cc index baa423bd70..3ddca42422 100644 --- a/src/ascend/transform/materialize_multi_buffer.cc +++ b/src/ascend/transform/materialize_multi_buffer.cc @@ -501,7 +501,7 @@ void ElideUnitVersionCounterGroups(ScheduledTIR *scheduled_tir, SBlock MaterializeKernel(ScheduledTIR scheduled_tir) { MultiBufferPlan plan = ReadMultiBufferPlan( scheduled_tir.tree, scheduled_tir.metadata.buffer_versions, - CollectL0StorageGroups(scheduled_tir.metadata.kernel_root)); + L0StorageGroups(CollectL0SFBindings(scheduled_tir.metadata.kernel_root))); ElideUnitVersionCounterGroups(&scheduled_tir, plan); VersionedBufferRegistry buffers(plan, scheduled_tir.metadata.kernel_root); RewriteScheduledTree(&scheduled_tir.tree, plan, &buffers); @@ -515,8 +515,6 @@ SBlock MaterializeKernel(ScheduledTIR scheduled_tir) { AddBufferVersion(&physical_versions, info.storage, info.num_versions); scheduled_tir.metadata.buffer_versions = {}; scheduled_tir.metadata.manual_buffer_versions = std::move(physical_versions); - scheduled_tir.metadata.kernel_root.CopyOnWrite()->annotations.erase( - kReuseRingBindings); Stmt serialized = EncodeScheduledTIR(std::move(scheduled_tir)); serialized = AllocationRewriter(&buffers)(std::move(serialized)); diff --git a/src/ascend/transform/prepare_multi_buffer.cc b/src/ascend/transform/prepare_multi_buffer.cc index ae19c60c6f..434a8a08e6 100644 --- a/src/ascend/transform/prepare_multi_buffer.cc +++ b/src/ascend/transform/prepare_multi_buffer.cc @@ -1222,7 +1222,7 @@ tvm::transform::Pass PrepareMultiBuffer() { ICHECK_NE(task->GetCoreMask(), kCoreUnassigned) << "PrepareMultiBuffer requires AssignCore to run first"; } - L0StorageGroups groups = CollectL0StorageGroups(context.root); + L0StorageGroups groups(CollectL0SFBindings(context.root)); auto &metadata = scheduled_tir.metadata; metadata.buffer_versions = ExpandL0StorageGroupValues( metadata.buffer_versions, groups, "version counts"); diff --git a/src/ascend/transform/resolve_core.cc b/src/ascend/transform/resolve_core.cc index d1d36b0a70..9b87518bcb 100644 --- a/src/ascend/transform/resolve_core.cc +++ b/src/ascend/transform/resolve_core.cc @@ -639,10 +639,11 @@ tvm::transform::Pass ResolveCore() { [](const TilelangKernelContext &context) { ScheduledTIR scheduled_tir = DecodeScheduledTIR(context.root, context.outer_ctx); - ResolveCoreMasks( - scheduled_tir.tree, scheduled_tir.metadata.buffer_versions, - context.outer_sid, CollectL0StorageGroups(context.root), - scheduled_tir.metadata.num_aiv_subcores.value_or(2)); + ResolveCoreMasks(scheduled_tir.tree, + scheduled_tir.metadata.buffer_versions, + context.outer_sid, + L0StorageGroups(CollectL0SFBindings(context.root)), + scheduled_tir.metadata.num_aiv_subcores.value_or(2)); return EncodeScheduledTIR(std::move(scheduled_tir)); }); }; diff --git a/testing/ascend/auto_schedule/test_conflict_hints.py b/testing/ascend/auto_schedule/test_conflict_hints.py index 4ec1032f82..5a12b8d7d8 100644 --- a/testing/ascend/auto_schedule/test_conflict_hints.py +++ b/testing/ascend/auto_schedule/test_conflict_hints.py @@ -6,8 +6,7 @@ from tilelang import tvm from tilelang.ascend import transform from tvm import tirx -from testing.ascend._ir import calls, kernel, schedule, nodes, seq -from testing.ascend.auto_schedule._scheduled_ir import unit +from testing.ascend._ir import calls, schedule, nodes def _insert_sync(program): @@ -19,46 +18,6 @@ def _event_count(mod, operation, pipe): return sum(call.args[0].value == pipe for call in calls(mod, "tl.ascend_" + operation + "_flag")) -@pytest.mark.parametrize("scope,cross", [("root", False), ("loop", False), ("loop", True), ("nested", None)]) -@pytest.mark.parametrize("regions", [False, True], ids=["storage", "region"]) -@pytest.mark.parametrize("disable_reuse", [False, True]) -def test_conflict_hints_permit_pairwise_storage_reuse(scope, cross, regions, disable_reuse): - a, b, c = [tirx.decl_buffer((8,), "int32", name=name, scope="shared.dyn") for name in ("a", "b", "c")] - out = tirx.decl_buffer((1,), "int32", name="out") - - def operand(buffer): - return tirx.BufferRegion(buffer, [tvm.ir.Range(0, 4)]) if regions else buffer.data - - def hint(left, right, positive=True): - return [operand(left), operand(right), tirx.IntImm("int32", -1 if cross is None else int(cross)), tirx.IntImm("bool", positive)] - - hints = [hint(a, b), hint(b, c), hint(a, b), hint(b, b), hint(a, c, False)] - # All three values stay live until the final read, so ordinary lifetime - # inference cannot authorize reuse. Equal values make the declared reuse - # valid without changing the result. - body = seq( - *(unit(tirx.BufferStore(buffer, 1, [0]), core=1) for buffer in (a, b, c)), - unit(tirx.BufferStore(out, a[0] + b[0] + c[0], [0]), core=1), - ) - annotations = {"tl.root_conflict_hints": hints} if scope == "root" else {} - if scope != "root": - loop = tirx.For(tirx.Var("i", "int32"), 0, 2, tirx.ForKind.SERIAL, body, annotations={"conflict_hint": hints}) - body = unit(loop, core=None) - if scope == "nested": - body = unit(tirx.For(tirx.Var("j", "int32"), 0, 2, tirx.ForKind.SERIAL, body), core=None) - before = kernel(body, buffers=[a, b, c], params=[out], annotations=annotations) - with tvm.transform.PassContext(config={"tl.disable_shared_memory_reuse": disable_reuse}): - after = transform.InsertSync()(before) - (aliases,) = [ - block.annotations["tl.buffer_alias_map"] for block in nodes(after, tirx.SBlock) if "tl.buffer_alias_map" in block.annotations - ] - expected = [] if disable_reuse else [(a.data, b.data), (b.data, a.data), (b.data, c.data), (c.data, b.data)] - actual = [(storage, other) for storage, partners in aliases.items() for other in partners] - assert len(actual) == len(expected) - for lhs, rhs in expected: - assert any(storage.same_as(lhs) and other.same_as(rhs) for storage, other in actual) - - def _hint_program(scope, exact, cross, explicit_root): @T.prim_func def main(A: T.Tensor((128,), "float32"), C: T.Tensor((128,), "float32")): diff --git a/testing/ascend/auto_schedule/test_conflict_ring_owners.py b/testing/ascend/auto_schedule/test_conflict_ring_owners.py deleted file mode 100644 index 306a10e07d..0000000000 --- a/testing/ascend/auto_schedule/test_conflict_ring_owners.py +++ /dev/null @@ -1,311 +0,0 @@ -"""Explicit layout aliases share the physical clock of retimed L0B owners.""" - -import pytest -import tilelang -import tilelang.ascend.language as T -import tilelang.testing -from tilelang import tvm -from tilelang.layout import make_ascend_major_k_layout, make_ascend_major_mn_layout -from tvm import tirx -from testing.ascend._ir import allocated_buffer, nodes - - -def _retimed_layout_pair(qk_steps=4, pv_steps=4, alias_kind=None, transpose_pv=False, annotated=True): - @T.prim_func - def main( - Q: T.Tensor((64, 128), "bfloat16"), - K: T.Tensor((64, 128), "bfloat16"), - CQ: T.Tensor((3, 64, 64), "float32"), - CV: T.Tensor((3, pv_steps, 64, 128), "float32"), - ): - with T.Kernel(1): - q = T.alloc_l1((64, 128), "bfloat16") - k0 = T.alloc_l1((32, 128), "bfloat16") - k1 = T.alloc_l1((32, 128), "bfloat16") - aq = T.alloc_l0a((64, 128), "bfloat16") - bq = T.alloc_l0b((64, 128), "bfloat16") - if alias_kind == "view": - bv = T.view(bq, (128, 64) if transpose_pv else (64, 128)) - elif alias_kind == "reshape": - bv = T.reshape(bq, (128, 64) if transpose_pv else (64, 128)) - else: - bv = T.alloc_l0b((64, 128), "bfloat16") - cq = T.alloc_l0c((64, 64), "float32") - cv = T.alloc_l0c((pv_steps, 64, 128), "float32") - T.annotate_buffer_versions({q: 1, k0: 1, k1: 1, aq: 4, bq: 4, bv: 4, cq: 1, cv: 1}) - if annotated: - T.annotate_layout( - { - bq: make_ascend_major_k_layout(bq), - bv: make_ascend_major_k_layout(bv) if transpose_pv else make_ascend_major_mn_layout(bv), - } - ) - T.copy(Q, q) - T.copy(K[:32, :], k0) - T.copy(K[32:, :], k1) - # Keep the Cube mode fixed so this regression isolates storage - # ownership from separately scheduled configuration changes. - for step in T.Pipelined(3, num_stages=2, annotations={"enable_offset": True}): - if alias_kind is None: - T.assume_conflict(bq, bv) - with T.Stage(0): - for d in T.Pipelined(qk_steps, num_stages=4): - with T.Stage(0): - T.copy(q, aq) - with T.Task(): - T.copy(k0, bq[:32, :], l0_pitch=4) - T.copy(k1, bq[32:, :], l0_pitch=4) - T.gemm(aq, bq, cq, transpose_B=True, clear_accum=d == 0) - T.copy(cq, CQ[step, :, :]) - with T.Stage(1): - for _region in T.Pipelined(1, num_stages=4): - with T.Stage(0): - T.copy(q[:, :64], aq[:, :64]) - for d in T.Pipelined(pv_steps, num_stages=2): - with T.Stage(0): - with T.Task(): - if transpose_pv: - T.copy(k0, bv[:, :32], transpose=True, l0_pitch=8) - T.copy(k1, bv[:, 32:], transpose=True, l0_pitch=8) - else: - T.copy(k0, bv[:32, :], l0_pitch=8) - T.copy(k1, bv[32:, :], l0_pitch=8) - T.gemm(aq[:, :64], bv, cv[d, :, :], transpose_B=transpose_pv, clear_accum=True) - for d in T.Pipelined(pv_steps, num_stages=2): - with T.Stage(0): - T.copy(cv[d, :, :], CV[step, d, :, :]) - - return main - - -@pytest.mark.parametrize("qk_steps,pv_steps", [(4, 4), (3, 2)]) -@pytest.mark.parametrize("disable_reuse", [False, True]) -def test_default_conflict_hint_shares_retimed_owner_clock(qk_steps, pv_steps, disable_reuse): - snapshots = {} - - @tvm.instrument.pass_instrument - class Capture: - def run_after_pass(self, mod, info): - if info.name in ("tl.PrepareMultiBuffer", "tl.MergeUBAllocations"): - snapshots[info.name] = mod - - with ( - tvm.target.Target("ascend"), - tvm.transform.PassContext(config={"tl.disable_shared_memory_reuse": disable_reuse}, instruments=[Capture()]), - ): - tilelang.lower(_retimed_layout_pair(qk_steps, pv_steps), target="ascend", enable_device_compile=False) - - prepared = snapshots["tl.PrepareMultiBuffer"] - bq, bv = [allocated_buffer(prepared, name) for name in ("bq", "bv")] - assert not bq.data.same_as(bv.data) - clocks = [] - for loop in nodes(prepared, tirx.For): - if "tl.multi_buffer_counter_map" not in loop.annotations: - continue - counters = [ - counter - for storage, counter in loop.annotations["tl.multi_buffer_counter_map"].items() - if storage.same_as(bq.data) or storage.same_as(bv.data) - ] - if not counters: - continue - assert len(counters) == 2 - assert counters[0].same_as(counters[1]) - clocks.append(counters[0]) - assert len(clocks) == 2 - assert all(clock.same_as(clocks[0]) for clock in clocks) - advances = [ - store - for store in nodes(prepared, tirx.BufferStore) - if store.buffer.same_as(clocks[0]) and tvm.arith.Analyzer().can_prove_equal(store.value, clocks[0][0] + 1) - ] - assert len(advances) == 2 - - allocations = [ - node.buffer - for func in snapshots["tl.MergeUBAllocations"].functions.values() - for node in nodes(func, tirx.AllocBuffer) - if node.buffer.scope() in ("shared.l0b", "shared.l0b.dyn") - ] - assert len(allocations) == 1 - allocated_bytes = int(allocations[0].shape[0]) * tvm.DataType(allocations[0].dtype).bits // 8 - assert allocated_bytes == (131072 if disable_reuse else 65536) - - -@pytest.mark.parametrize("alias_kind", ["view", "reshape"]) -@pytest.mark.parametrize("transpose_pv,annotated", [(False, True), (True, True), (False, False), (True, False)]) -def test_storage_views_keep_layouts_and_share_one_version_clock(alias_kind, transpose_pv, annotated): - snapshots = {} - - @tvm.instrument.pass_instrument - class Capture: - def run_after_pass(self, mod, info): - if info.name in ( - "tl.NormalizeAscendFractalStorage", - "tl.PrepareMultiBuffer", - "tl.MaterializeMultiBuffer", - "tl.MergeUBAllocations", - ): - snapshots[info.name] = mod - - with tvm.target.Target("ascend"), tvm.transform.PassContext(config={"tl.disable_shared_memory_reuse": True}, instruments=[Capture()]): - tilelang.lower(_retimed_layout_pair(3, 2, alias_kind, transpose_pv, annotated), target="ascend", enable_device_compile=False) - - def l0b_layouts(mod): - return { - buffer: layout - for block in nodes(mod, tirx.SBlock) - for buffer, layout in block.annotations.get("layout_map", {}).items() - if buffer.scope() == "shared.l0b" - } - - layouts = l0b_layouts(snapshots["tl.NormalizeAscendFractalStorage"]) - assert len(layouts) == 2 - first, second = layouts - assert first.data.same_as(second.data) - assert not layouts[first].is_equal(layouts[second]) - expected_shapes = {(64, 128), (128, 64)} if transpose_pv else {(64, 128)} - assert {tuple(int(x) for x in buffer.shape) for buffer in layouts} == expected_shapes - - clocks = [ - counter - for loop in nodes(snapshots["tl.PrepareMultiBuffer"], tirx.For) - for storage, counter in loop.annotations.get("tl.multi_buffer_counter_map", {}).items() - if storage.same_as(first.data) - ] - assert len(clocks) == 2 and clocks[0].same_as(clocks[1]) - expanded = {buffer: layout for buffer, layout in l0b_layouts(snapshots["tl.MaterializeMultiBuffer"]).items() if len(buffer.shape) == 3} - assert len(expanded) == 2 - for buffer, layout in expanded.items(): - assert int(buffer.shape[0]) == 4 and int(buffer.strides[0]) == 8192 - expected = make_ascend_major_k_layout(buffer) if transpose_pv or buffer.name == "bq" else make_ascend_major_mn_layout(buffer) - assert layout.is_equal(expected) - allocations = [ - node.buffer - for func in snapshots["tl.MergeUBAllocations"].functions.values() - for node in nodes(func, tirx.AllocBuffer) - if node.buffer.scope().startswith("shared.l0b") - ] - assert len(allocations) == 1 - assert int(allocations[0].shape[0]) * tvm.DataType(allocations[0].dtype).bits // 8 == 65536 - - -def test_conflict_hint_does_not_join_nested_storage_owners(): - @T.prim_func - def main(Q: T.Tensor((64, 128), "bfloat16"), K: T.Tensor((64, 128), "bfloat16"), C: T.Tensor((64, 128), "float32")): - with T.Kernel(1): - q = T.alloc_l1((64, 128), "bfloat16") - k = T.alloc_l1((64, 128), "bfloat16") - aq = T.alloc_l0a((64, 128), "bfloat16") - av = T.alloc_l0a((64, 64), "bfloat16") - bq = T.alloc_l0b((64, 128), "bfloat16") - bv = T.alloc_l0b((64, 128), "bfloat16") - cq = T.alloc_l0c((64, 64), "float32") - cv = T.alloc_l0c((64, 128), "float32") - T.annotate_buffer_versions({q: 1, k: 1, aq: 1, av: 1, bq: 4, bv: 4, cq: 1, cv: 1}) - T.copy(Q, q) - T.copy(K, k) - for _i in T.Pipelined(2, num_stages=2): - T.assume_conflict(bq, bv) - with T.Stage(0): - T.copy(q, aq) - T.copy(k, bq) - T.gemm(aq, bq, cq, transpose_B=True, clear_accum=True) - T.copy(cq, C[:, :64]) - for _j in T.Pipelined(2, num_stages=2): - with T.Stage(0): - T.copy(q[:, :64], av) - T.copy(k, bv) - T.gemm(av, bv, cv, clear_accum=True) - T.copy(cv, C) - - with tvm.target.Target("ascend"): - tilelang.lower(main, target="ascend", enable_device_compile=False) - - -@pytest.mark.parametrize( - "automatic,staged_second", - [(False, False), (False, True), (True, False)], - ids=["manual-first", "manual-second", "automatic-offset"], -) -@pytest.mark.parametrize("disable_reuse", [False, True]) -def test_conflict_hint_keeps_independent_owners_when_stages_can_differ(automatic, staged_second, disable_reuse): - @T.prim_func - def main(Q: T.Tensor((16, 16), "bfloat16"), K: T.Tensor((16, 16), "bfloat16"), C: T.Tensor((2, 16, 16), "float32")): - with T.Kernel(1): - q = T.alloc_l1((16, 16), "bfloat16") - k = T.alloc_l1((16, 16), "bfloat16") - aq = T.alloc_l0a((16, 16), "bfloat16") - av = T.alloc_l0a((16, 16), "bfloat16") - bq = T.alloc_l0b((16, 16), "bfloat16") - bv = T.alloc_l0b((16, 16), "bfloat16") - cq = T.alloc_l0c((16, 16), "float32") - cv = T.alloc_l0c((16, 16), "float32") - T.annotate_buffer_versions({q: 1, k: 1, aq: 4, av: 1, bq: 4, bv: 4, cq: 1, cv: 1}) - T.assume_conflict(bq, bv) - T.copy(Q, q) - T.copy(K, k) - T.copy(q, av) - staged = bv if staged_second else bq - other = bq if staged_second else bv - for _i in T.Pipelined(4, num_stages=4, annotations={"enable_offset": True}): - if automatic: - T.copy(q, aq) - T.copy(k, staged) - T.gemm(aq, staged, cq, transpose_B=True, clear_accum=True) - else: - with T.Stage(0): - T.copy(q, aq) - T.copy(k, staged) - with T.Stage(1): - T.gemm(aq, staged, cq, transpose_B=True, clear_accum=True) - T.copy(cq, C[0, :, :]) - for _j in T.Pipelined(4, num_stages=4): - with T.Stage(0): - T.copy(k, other) - T.gemm(av, other, cv, transpose_B=True, clear_accum=True) - T.copy(cv, C[1, :, :]) - - snapshots = {} - - @tvm.instrument.pass_instrument - class Capture: - def run_after_pass(self, mod, info): - if info.name in ("tl.PrepareMultiBuffer", "tl.MaterializeMultiBuffer"): - snapshots[info.name] = mod - - with tvm.transform.PassContext(config={"tl.disable_shared_memory_reuse": disable_reuse}, instruments=[Capture()]): - artifact = tilelang.lower(main, target="ascend", enable_device_compile=False) - - prepared = snapshots["tl.PrepareMultiBuffer"] - owners = [] - for name in ("staged", "other"): - buffer = allocated_buffer(prepared, name) - (owner,) = [loop for loop in nodes(prepared, tirx.For) if buffer.data in loop.annotations.get("multi_buffer_eligible", [])] - owners.append(owner) - assert int(allocated_buffer(snapshots["tl.MaterializeMultiBuffer"], name).shape[0]) == 4 - assert not owners[0].same_as(owners[1]) - assert artifact.kernel_source - - -@pytest.mark.parametrize( - "alias_kind,transpose_pv,annotated", [(None, False, True), ("view", False, True), ("reshape", True, True), ("view", True, False)] -) -@tilelang.testing.requires_ascend -def test_retimed_layout_pair_values_across_owner_lengths_and_launches(alias_kind, transpose_pv, annotated): - import torch - - torch.npu.set_device(0) - torch.npu.set_op_timeout_ms(30000) - with tvm.target.Target("ascend"): - kernel = tilelang.compile(_retimed_layout_pair(3, 2, alias_kind, transpose_pv, annotated), target="ascend", out_idx=[2, 3]) - for seed in (7, 11, 19): - torch.manual_seed(seed) - q = torch.randint(-2, 3, (64, 128), dtype=torch.int32).to(dtype=torch.bfloat16, device="npu") - k = torch.randint(-2, 3, (64, 128), dtype=torch.int32).to(dtype=torch.bfloat16, device="npu") - cq, cv = kernel(q, k) - torch.npu.synchronize() - expected_qk = 3 * (q.cpu().float() @ k.cpu().float().T) - expected_pv = q.cpu().float()[:, :64] @ k.cpu().float() - torch.testing.assert_close(cq.cpu(), expected_qk.expand(3, 64, 64), rtol=0, atol=0) - torch.testing.assert_close(cv.cpu(), expected_pv.expand(3, 2, 64, 128), rtol=0, atol=0) diff --git a/testing/ascend/auto_schedule/test_cube_storage_views.py b/testing/ascend/auto_schedule/test_cube_storage_views.py new file mode 100644 index 0000000000..b93370591d --- /dev/null +++ b/testing/ascend/auto_schedule/test_cube_storage_views.py @@ -0,0 +1,153 @@ +"""Cube views preserve distinct layouts on one versioned storage allocation.""" + +import pytest +import tilelang +import tilelang.ascend.language as T +import tilelang.testing +from tilelang import tvm +from tilelang.layout import make_ascend_major_k_layout, make_ascend_major_mn_layout +from tvm import tirx +from testing.ascend._ir import nodes + + +def _retimed_layout_views(qk_steps=4, pv_steps=4, alias_kind="view", transpose_pv=False, annotated=True): + @T.prim_func + def main( + Q: T.Tensor((64, 128), "bfloat16"), + K: T.Tensor((64, 128), "bfloat16"), + CQ: T.Tensor((3, 64, 64), "float32"), + CV: T.Tensor((3, pv_steps, 64, 128), "float32"), + ): + with T.Kernel(1): + q = T.alloc_l1((64, 128), "bfloat16") + k0 = T.alloc_l1((32, 128), "bfloat16") + k1 = T.alloc_l1((32, 128), "bfloat16") + aq = T.alloc_l0a((64, 128), "bfloat16") + bq = T.alloc_l0b((64, 128), "bfloat16") + if alias_kind == "view": + bv = T.view(bq, (128, 64) if transpose_pv else (64, 128)) + else: + bv = T.reshape(bq, (128, 64) if transpose_pv else (64, 128)) + cq = T.alloc_l0c((64, 64), "float32") + cv = T.alloc_l0c((pv_steps, 64, 128), "float32") + T.annotate_buffer_versions({q: 1, k0: 1, k1: 1, aq: 4, bq: 4, cq: 1, cv: 1}) + if annotated: + T.annotate_layout( + { + bq: make_ascend_major_k_layout(bq), + bv: make_ascend_major_k_layout(bv) if transpose_pv else make_ascend_major_mn_layout(bv), + } + ) + T.copy(Q, q) + T.copy(K[:32, :], k0) + T.copy(K[32:, :], k1) + # Keep the Cube mode fixed so this regression isolates storage + # ownership from separately scheduled configuration changes. + for step in T.Pipelined(3, num_stages=2, annotations={"enable_offset": True}): + with T.Stage(0): + for d in T.Pipelined(qk_steps, num_stages=4): + with T.Stage(0): + T.copy(q, aq) + with T.Task(): + T.copy(k0, bq[:32, :], l0_pitch=4) + T.copy(k1, bq[32:, :], l0_pitch=4) + T.gemm(aq, bq, cq, transpose_B=True, clear_accum=d == 0) + T.copy(cq, CQ[step, :, :]) + with T.Stage(1): + for _region in T.Pipelined(1, num_stages=4): + with T.Stage(0): + T.copy(q[:, :64], aq[:, :64]) + for d in T.Pipelined(pv_steps, num_stages=2): + with T.Stage(0): + with T.Task(): + if transpose_pv: + T.copy(k0, bv[:, :32], transpose=True, l0_pitch=8) + T.copy(k1, bv[:, 32:], transpose=True, l0_pitch=8) + else: + T.copy(k0, bv[:32, :], l0_pitch=8) + T.copy(k1, bv[32:, :], l0_pitch=8) + T.gemm(aq[:, :64], bv, cv[d, :, :], transpose_B=transpose_pv, clear_accum=True) + for d in T.Pipelined(pv_steps, num_stages=2): + with T.Stage(0): + T.copy(cv[d, :, :], CV[step, d, :, :]) + + return main + + +@pytest.mark.parametrize("alias_kind", ["view", "reshape"]) +@pytest.mark.parametrize("transpose_pv,annotated", [(False, True), (True, True), (False, False), (True, False)]) +def test_storage_views_keep_layouts_and_share_one_version_clock(alias_kind, transpose_pv, annotated): + snapshots = {} + + @tvm.instrument.pass_instrument + class Capture: + def run_after_pass(self, mod, info): + if info.name in ( + "tl.NormalizeAscendFractalStorage", + "tl.PrepareMultiBuffer", + "tl.MaterializeMultiBuffer", + "tl.MergeUBAllocations", + ): + snapshots[info.name] = mod + + with tvm.target.Target("ascend"), tvm.transform.PassContext(config={"tl.disable_shared_memory_reuse": True}, instruments=[Capture()]): + tilelang.lower(_retimed_layout_views(3, 2, alias_kind, transpose_pv, annotated), target="ascend", enable_device_compile=False) + + def l0b_layouts(mod): + return { + buffer: layout + for block in nodes(mod, tirx.SBlock) + for buffer, layout in block.annotations.get("layout_map", {}).items() + if buffer.scope() == "shared.l0b" + } + + layouts = l0b_layouts(snapshots["tl.NormalizeAscendFractalStorage"]) + assert len(layouts) == 2 + first, second = layouts + assert first.data.same_as(second.data) + assert not layouts[first].is_equal(layouts[second]) + expected_shapes = {(64, 128), (128, 64)} if transpose_pv else {(64, 128)} + assert {tuple(int(x) for x in buffer.shape) for buffer in layouts} == expected_shapes + + clocks = [ + counter + for loop in nodes(snapshots["tl.PrepareMultiBuffer"], tirx.For) + for storage, counter in loop.annotations.get("tl.multi_buffer_counter_map", {}).items() + if storage.same_as(first.data) + ] + assert len(clocks) == 2 and clocks[0].same_as(clocks[1]) + expanded = {buffer: layout for buffer, layout in l0b_layouts(snapshots["tl.MaterializeMultiBuffer"]).items() if len(buffer.shape) == 3} + assert len(expanded) == 2 + for buffer, layout in expanded.items(): + assert int(buffer.shape[0]) == 4 and int(buffer.strides[0]) == 8192 + expected = make_ascend_major_k_layout(buffer) if transpose_pv or buffer.name == "bq" else make_ascend_major_mn_layout(buffer) + assert layout.is_equal(expected) + allocations = [ + node.buffer + for func in snapshots["tl.MergeUBAllocations"].functions.values() + for node in nodes(func, tirx.AllocBuffer) + if node.buffer.scope().startswith("shared.l0b") + ] + assert len(allocations) == 1 + assert int(allocations[0].shape[0]) * tvm.DataType(allocations[0].dtype).bits // 8 == 65536 + + +@pytest.mark.parametrize("alias_kind,transpose_pv,annotated", [("view", False, True), ("reshape", True, True), ("view", True, False)]) +@tilelang.testing.requires_ascend +def test_view_values_across_owner_lengths_and_launches(alias_kind, transpose_pv, annotated): + import torch + + torch.npu.set_device(0) + torch.npu.set_op_timeout_ms(30000) + with tvm.target.Target("ascend"): + kernel = tilelang.compile(_retimed_layout_views(3, 2, alias_kind, transpose_pv, annotated), target="ascend", out_idx=[2, 3]) + for seed in (7, 11, 19): + torch.manual_seed(seed) + q = torch.randint(-2, 3, (64, 128), dtype=torch.int32).to(dtype=torch.bfloat16, device="npu") + k = torch.randint(-2, 3, (64, 128), dtype=torch.int32).to(dtype=torch.bfloat16, device="npu") + cq, cv = kernel(q, k) + torch.npu.synchronize() + expected_qk = 3 * (q.cpu().float() @ k.cpu().float().T) + expected_pv = q.cpu().float()[:, :64] @ k.cpu().float() + torch.testing.assert_close(cq.cpu(), expected_qk.expand(3, 64, 64), rtol=0, atol=0) + torch.testing.assert_close(cv.cpu(), expected_pv.expand(3, 2, 64, 128), rtol=0, atol=0) diff --git a/testing/ascend/auto_schedule/test_reuse_integration.py b/testing/ascend/auto_schedule/test_reuse_integration.py index dbfc79a965..1252b1ca49 100644 --- a/testing/ascend/auto_schedule/test_reuse_integration.py +++ b/testing/ascend/auto_schedule/test_reuse_integration.py @@ -1,89 +1,13 @@ """The reuse contract survives scheduled lowering and is consumed by packing.""" -import pytest import tilelang import tilelang.ascend.language as T from tilelang import tvm from tvm import tirx -from testing.ascend._ir import calls, nodes +from testing.ascend._ir import nodes from testing.ascend.auto_schedule._reuse_utils import _make_cyclic_scalar_program -@pytest.mark.parametrize("disable_reuse", [False, True]) -def test_conflict_hint_reuses_independent_l0_layouts(disable_reuse): - @T.prim_func - def main( - Q: T.Tensor((64, 128), "bfloat16"), - K: T.Tensor((64, 128), "bfloat16"), - CQ: T.Tensor((64, 64), "float32"), - CV: T.Tensor((64, 128), "float32"), - ): - with T.Kernel(1): - q = T.alloc_l1((64, 128), "bfloat16") - k0 = T.alloc_l1((32, 128), "bfloat16") - k1 = T.alloc_l1((32, 128), "bfloat16") - aq = T.alloc_l0a((64, 128), "bfloat16") - av = T.alloc_l0a((64, 64), "bfloat16") - bq = T.alloc_l0b((64, 128), "bfloat16") - bv = T.alloc_l0b((64, 128), "bfloat16") - cq = T.alloc_l0c((64, 64), "float32") - cv = T.alloc_l0c((64, 128), "float32") - T.annotate_buffer_versions({q: 1, k0: 1, k1: 1, aq: 2, av: 1, bq: 4, bv: 4, cq: 1, cv: 1}) - T.assume_conflict(bq, bv, level=-1, cross=False) - T.copy(Q, q) - T.copy(K[:32, :], k0) - T.copy(K[32:, :], k1) - T.set_mmad_direction("m") - for _d in T.Pipelined(4, num_stages=4): - with T.Stage(0): - T.copy(q, aq) - with T.Task(): - T.copy(k0, bq[:32, :], l0_pitch=4) - T.copy(k1, bq[32:, :], l0_pitch=4) - T.gemm(aq, bq, cq, transpose_B=True, clear_accum=True) - T.copy(cq, CQ) - T.copy(q[:, :64], av) - T.set_mmad_direction("n") - for _d in T.Pipelined(4, num_stages=4): - with T.Stage(0): - with T.Task(): - T.copy(k0, bv[:32, :], l0_pitch=8) - T.copy(k1, bv[32:, :], l0_pitch=8) - T.gemm(av, bv, cv, clear_accum=True) - T.copy(cv, CV) - - snapshots = {} - - @tvm.instrument.pass_instrument - class Capture: - def run_after_pass(self, mod, info): - if info.name in ("tl.AscendLowerTileOp", "tl.MergeUBAllocations"): - snapshots[info.name] = mod - - with ( - tvm.target.Target("ascend"), - tvm.transform.PassContext(config={"tl.disable_shared_memory_reuse": disable_reuse}, instruments=[Capture()]), - ): - tilelang.lower(main, target="ascend", enable_device_compile=False) - (allocation,) = [ - node.buffer - for func in snapshots["tl.MergeUBAllocations"].functions.values() - for node in nodes(func, tirx.AllocBuffer) - if node.buffer.scope() in ("shared.l0b", "shared.l0b.dyn") - ] - allocated_bytes = int(allocation.shape[0]) * tvm.DataType(allocation.dtype).bits // 8 - assert allocated_bytes == (131072 if disable_reuse else 65536) - - offsets = {0: [], 1: []} - analyzer = tvm.arith.Analyzer() - for call in calls(snapshots["tl.AscendLowerTileOp"], "tl.ascend_load_cbuf_to_cb"): - pointer = call.args[0] - assert pointer.op.name == "tirx.tvm_access_ptr" - offsets[int(call.args[-1])].append(int(analyzer.simplify(pointer.args[2] % 8192))) - assert sorted(offsets[0]) == [0, 512] - assert sorted(offsets[1]) == [0, 4096] - - def test_alias_contract_lifetime(): snapshots = {} wanted = { diff --git a/tilelang/ascend/language/schedule_hint.py b/tilelang/ascend/language/schedule_hint.py index 44497a699a..51e0fdb371 100644 --- a/tilelang/ascend/language/schedule_hint.py +++ b/tilelang/ascend/language/schedule_hint.py @@ -244,7 +244,7 @@ def assume_no_conflict(a, b=None, *, level=None, cross=None, group=None): def assume_conflict(a, b=None, *, level=None, cross=None, group=None): - """Assert a dependency conflict and permit the two storages to be reused. + """Assert that two buffer regions conflict for dependency analysis. This is the conservative counterpart of :func:`assume_no_conflict`. It makes the auto-scheduler treat matching accesses as if they used the same @@ -253,19 +253,6 @@ def assume_conflict(a, b=None, *, level=None, cross=None, group=None): accesses do not form a dependency. The marker is consumed before auto-scheduling and is a runtime no-op. - A positive hint also explicitly permits physical reuse of the two backing - storages, including when lifetime inference cannot establish that reuse. - The permission is pairwise and does not imply transitive compatibility. - Region, ``level``, and ``cross`` restrict the dependency matching; the reuse - permission applies to the whole storage pair. The caller must ensure that - the declared dependencies cover every access affected by that reuse. - Existing allocation scope, alignment, and reuse-disable rules still apply. - Buffer identities and layouts remain separate. For an isolated, equal-sized - L0B pair with the same fixed automatic version count, a whole-buffer hint - with the default iteration matching can also select one shared ring clock. - All accesses must belong to disjoint automatic owners; explicit iteration indexing, - manual versions, strides, or offsets retain the ordinary dependency path. - Args: a: first region -- ``buffer[indices...]`` or a bare ``buffer`` (the whole buffer). A concrete region must exactly match the footprint From 17bbbece35820d58b220fe82cdea4ca2edbf1b44 Mon Sep 17 00:00:00 2001 From: Denver Jin Date: Thu, 22 Oct 2026 17:20:49 +0800 Subject: [PATCH 7/7] [Ascend] Normalize each buffer view independently --- src/ascend/op/copy.cc | 30 +-- src/ascend/op/oob_padding.cc | 14 +- src/ascend/op/oob_padding.h | 5 +- src/ascend/op/utils.h | 5 - src/ascend/transform/layout_inference.cc | 105 +++++---- src/ascend/transform/lower_tile_op.cc | 82 ++----- .../transform/normalize_fractal_storage.cc | 200 +++++++++--------- .../test_ascend_buffer_view_layouts.py | 114 ++++++++++ .../transform/test_ascend_lower_copy.py | 6 +- .../test_ascend_normalize_fractal_storage.py | 60 ++++-- tilelang/ascend/op/gemm/gemm_mad.py | 14 +- 11 files changed, 343 insertions(+), 292 deletions(-) create mode 100644 testing/ascend/transform/test_ascend_buffer_view_layouts.py diff --git a/src/ascend/op/copy.cc b/src/ascend/op/copy.cc index e45d9c8848..bfc0259ae9 100644 --- a/src/ascend/op/copy.cc +++ b/src/ascend/op/copy.cc @@ -423,21 +423,13 @@ Stmt LowerDMACopy(const AscendCopyNode &op, const LowerArgs &T, Array src_range = NormalizeEmptyUnitAxesForMTE(op.src_range, analyzer); Array dst_range = NormalizeEmptyUnitAxesForMTE(op.dst_range, analyzer); - auto make_access_ptr = [&](const Buffer &buf, PrimExpr offset, + auto make_access_ptr = [&](const Buffer &buf, const Array &ranges, PrimExpr extent, int rw_mask) -> PrimExpr { - if (IsAscendCubeBuffer(buf)) { - Array indices(buf->shape.size(), I(0)); - PrimExpr remaining = offset; - for (int i = static_cast(buf->shape.size()) - 1; i >= 0; --i) { - indices.Set(i, i == 0 ? remaining : floormod(remaining, buf->shape[i])); - remaining = floordiv(remaining, buf->shape[i]); - } - return Call(DataType::Handle(), tl::access_ptr(), - {BufferLoad(buf, indices), extent, I(rw_mask)}); - } - return Call( - DataType::Handle(), builtin::tvm_access_ptr(), - {TypeAnnotation(buf->dtype), buf->data, offset, extent, I(rw_mask)}); + Array indices; + for (const Range &range : ranges) + indices.push_back(range->min); + return Call(DataType::Handle(), tl::access_ptr(), + {BufferLoad(buf, indices), extent, I(rw_mask)}); }; auto region_elements = [&](const Array &ranges) { PrimExpr result = I(1); @@ -452,10 +444,8 @@ Stmt LowerDMACopy(const AscendCopyNode &op, const LowerArgs &T, op.dst, dst_range, op.dst->dtype.bits() * op.dst->dtype.lanes()); PrimExpr src_region_elements = region_elements(src_range); PrimExpr dst_region_elements = region_elements(dst_range); - PrimExpr src_ptr = - make_access_ptr(op.src, src_region.offset, src_region_elements, 1); - PrimExpr dst_ptr = - make_access_ptr(op.dst, dst_region.offset, dst_region_elements, 2); + PrimExpr src_ptr = make_access_ptr(op.src, src_range, src_region_elements, 1); + PrimExpr dst_ptr = make_access_ptr(op.dst, dst_range, dst_region_elements, 2); PrimExpr sid = I(0); PrimExpr zero = I(0); @@ -474,9 +464,9 @@ Stmt LowerDMACopy(const AscendCopyNode &op, const LowerArgs &T, PrimExpr plan_total_elems = plan.row_elems * plan.n_rows; PrimExpr plan_src_ptr = - make_access_ptr(op.src, plan.src_offset, plan_total_elems, 1); + make_access_ptr(op.src, src_range, plan_total_elems, 1); PrimExpr plan_dst_ptr = - make_access_ptr(op.dst, plan.dst_offset, plan_total_elems, 2); + make_access_ptr(op.dst, dst_range, plan_total_elems, 2); PrimExpr plan_total_bytes = plan_row_bytes * plan.n_rows; // Ascend GM->UB padding (data_select): an unaligned row is right-padded up diff --git a/src/ascend/op/oob_padding.cc b/src/ascend/op/oob_padding.cc index cc595482be..f2a78519d4 100644 --- a/src/ascend/op/oob_padding.cc +++ b/src/ascend/op/oob_padding.cc @@ -69,19 +69,7 @@ DMAPath GetDMAPath(const Buffer &src, const Buffer &dst) { Optional FindLayoutForBuffer(const LayoutMap &layout_map, const Buffer &buffer) { - if (layout_map.count(buffer)) { - return layout_map[buffer]; - } - // A Cube alias may have a different canonical matrix interpretation. - // Borrowing a sibling's layout here suppresses its own inference constraints. - if (IsAscendCubeBuffer(buffer)) - return std::nullopt; - for (const auto &kv : layout_map) { - if (kv.first->data.same_as(buffer->data)) { - return kv.second; - } - } - return std::nullopt; + return layout_map.Get(buffer); } BoundedDMACopyRanges ClampDMACopyTail(const AscendCopyNode &op, diff --git a/src/ascend/op/oob_padding.h b/src/ascend/op/oob_padding.h index c6eaeef050..de79ebb5cc 100644 --- a/src/ascend/op/oob_padding.h +++ b/src/ascend/op/oob_padding.h @@ -53,9 +53,8 @@ enum class DMAPath { DMAPath GetDMAPath(const Buffer &src, const Buffer &dst); /*! - * \brief Look up a buffer's inferred Layout by Buffer identity. Non-Cube - * buffers may fall back to an alias with the same data Var; Cube views can - * have distinct layouts and require an exact match. Returns nullopt if absent. + * \brief Look up a view's inferred Layout by Buffer identity. Storage aliases + * may have distinct layouts; inference resolves any layout inheritance. */ Optional FindLayoutForBuffer(const LayoutMap &layout_map, const Buffer &buffer); diff --git a/src/ascend/op/utils.h b/src/ascend/op/utils.h index dba7d262bb..ca4d9a9870 100644 --- a/src/ascend/op/utils.h +++ b/src/ascend/op/utils.h @@ -75,11 +75,6 @@ inline bool IsAscendOnChipBuffer(const Buffer &buffer) { IsL0BBuffer(buffer) || IsL0CBuffer(buffer) || IsSSBuffer(buffer); } -inline bool IsAscendCubeBuffer(const Buffer &buffer) { - return IsL1Buffer(buffer) || IsL0ABuffer(buffer) || IsL0BBuffer(buffer) || - IsL0CBuffer(buffer); -} - /*! * \brief A view-preserving pointer whose innermost 2-D position travels as * separate intrinsic arguments. diff --git a/src/ascend/transform/layout_inference.cc b/src/ascend/transform/layout_inference.cc index 46dbe20c1b..d3b1c716fc 100644 --- a/src/ascend/transform/layout_inference.cc +++ b/src/ascend/transform/layout_inference.cc @@ -316,9 +316,10 @@ class BufferUseDefCollector : public IRVisitorWithAnalyzer { // Helper: propagate inferred layout to alias buffers (same data Var) auto propagate_alias = [&](const Buffer &src_buffer, const Layout &src_layout) { - // Cube operands constrain their own views. Only aliases without a - // constraint inherit a reshaped layout after inference completes. - if (IsAscendCubeBuffer(src_buffer)) + // Fragment aliases participate in the thread-ownership solver. + // Memory layouts belong to individual views; fill missing layouts + // only after each view's own constraints have been inferred. + if (!src_layout.as()) return; if (!buffer_data_to_buffers_.count(src_buffer->data)) return; @@ -627,45 +628,41 @@ class BufferUseDefCollector : public IRVisitorWithAnalyzer { in_queue); // step 3: relax constraints to free and re-run InferInFreeMode(layout_map, strict_layout_map); - // step 4: finalize alias layouts by Var - // For each storage var, if any buffer in the group has a layout, - // propagate (reshape if needed) to the rest to ensure completeness. + // step 4: fill unconstrained views only when all known interpretations of + // their storage agree. A shared data Var does not impose layout equality. + const LayoutMap inferred_layouts = layout_map; for (const auto &[var, buffers] : buffer_data_to_buffers_) { - // Find a representative with existing layout - Optional rep; - Optional rep_layout; for (const auto &buf : buffers) { - if (layout_map.count(buf)) { - rep = buf; - rep_layout = layout_map[buf]; - break; - } - } - if (!rep_layout.defined()) - continue; - for (const auto &buf : buffers) { - if (!layout_map.count(buf)) { - bool shapes_equal = - rep_layout.value()->InputShape().size() == buf->shape.size(); - if (shapes_equal) { - for (size_t i = 0; i < rep_layout.value()->InputShape().size(); - ++i) { - if (!analyzer_.CanProveEqual(rep_layout.value()->InputShape()[i], - buf->shape[i])) { - shapes_equal = false; - break; - } - } - } - Layout reshaped = - shapes_equal - ? rep_layout.value() - : rep_layout.value()->Reshape( + if (layout_map.count(buf)) + continue; + Optional inherited; + bool ambiguous = false; + for (const auto &source : buffers) { + Optional known = inferred_layouts.Get(source); + if (!known.defined()) + continue; + Layout candidate = + ShapesEqual(known.value()->InputShape(), buf->shape, &analyzer_) + ? known.value() + : known.value()->Reshape( buf->shape, &analyzer_, - Integer(GetElementStorageBits(rep.value()->dtype)), + Integer(GetElementStorageBits(source->dtype)), Integer(GetElementStorageBits(buf->dtype))); - layout_map.Set(buf, reshaped); + if (inherited.defined() && + !inherited.value()->IsEqual(candidate.get())) { + ambiguous = true; + break; + } + inherited = candidate; } + ICHECK(!ambiguous || !accessed_buffers_.count(buf)) + << "Ambiguous layout for buffer view " << buf->name << ": storage " + << var->name_hint + << " has distinct layouts. Annotate this view explicitly."; + // An unused backing Buffer only supplies allocation geometry. Any + // candidate suffices; storage normalization checks physical footprints. + if (inherited.defined()) + layout_map.Set(buf, inherited.value()); } } @@ -873,6 +870,7 @@ class BufferUseDefCollector : public IRVisitorWithAnalyzer { } void addToUseList(const Buffer &buffer) { + accessed_buffers_.insert(buffer); // Fragment buffers and reducers (whose PartialFragment travels the same // use_list_ edges: commits notify users, shared buffers union free-mode // components) both register; anything else has no layout to solve. @@ -1005,19 +1003,7 @@ class BufferUseDefCollector : public IRVisitorWithAnalyzer { // annotated layouts are available during InferLayout. // Must happen after visiting body so buffer_data_to_buffers_ is // populated. - if (pushed_region_scope && op->annotations.count(attr::kLayoutMap)) { - auto map = op->annotations.Get(attr::kLayoutMap) - ->cast>(); - for (const auto &[key, layout] : map) { - Var var = key.as() ? Downcast(key)->data - : Downcast(key); - if (buffer_data_to_buffers_.count(var)) { - for (const auto &buffer : buffer_data_to_buffers_[var]) { - annotated_layout_map_.Set(buffer, layout); - } - } - } - } + RecordLayoutAnnotations(GetRef(op)); // Regions contribute nothing to the inference worklist themselves: // nested tile operators were already collected while visiting the // body, and the region block is consumed by codegen as-is. @@ -1040,21 +1026,23 @@ class BufferUseDefCollector : public IRVisitorWithAnalyzer { // First, visit the block body to collect all buffers from // BufferLoad/BufferStore IRVisitorWithAnalyzer::VisitStmt_(op); + RecordLayoutAnnotations(GetRef(op)); + } + void RecordLayoutAnnotations(const SBlock &block) { // After visiting, apply layouts to all collected buffers - if (op->annotations.count(attr::kLayoutMap)) { - auto map = - op->annotations.Get(attr::kLayoutMap)->cast>(); + if (block->annotations.count(attr::kLayoutMap)) { + auto map = block->annotations.Get(attr::kLayoutMap) + ->cast>(); for (const auto &[key, layout] : map) { Optional annotated = key.as(); Var var = annotated.defined() ? annotated.value()->data : Downcast(key); ICHECK(buffer_data_to_buffers_.count(var)) << "buffer " << var << " is not found in the block"; - Array buffers = - annotated.defined() && IsAscendCubeBuffer(annotated.value()) - ? Array{annotated.value()} - : buffer_data_to_buffers_[var]; + Array buffers = annotated.defined() + ? Array{annotated.value()} + : buffer_data_to_buffers_[var]; ICHECK(!buffers.empty()) << "buffer list for " << var << " is empty"; for (const auto &buffer : buffers) { if (!IsReducerV2Buffer(buffer)) { @@ -1188,6 +1176,7 @@ class BufferUseDefCollector : public IRVisitorWithAnalyzer { } void VisitExpr_(const BufferLoadNode *op) final { + accessed_buffers_.insert(op->buffer); // Collect buffer from BufferLoad if (op->buffer.defined() && op->buffer->data.defined()) { if (buffer_data_to_buffers_.count(op->buffer->data)) { @@ -1218,6 +1207,7 @@ class BufferUseDefCollector : public IRVisitorWithAnalyzer { } void VisitStmt_(const BufferStoreNode *op) final { + accessed_buffers_.insert(op->buffer); // Collect buffer from BufferStore if (op->buffer.defined() && op->buffer->data.defined()) { if (buffer_data_to_buffers_.count(op->buffer->data)) { @@ -1346,6 +1336,7 @@ class BufferUseDefCollector : public IRVisitorWithAnalyzer { } Map> buffer_data_to_buffers_; + std::unordered_set accessed_buffers_; // Map from Bind variable to its bound expression Map bind_var_to_expr_; std::vector infer_list_stmt_; diff --git a/src/ascend/transform/lower_tile_op.cc b/src/ascend/transform/lower_tile_op.cc index 3bca5bbe07..63c1ca14d9 100644 --- a/src/ascend/transform/lower_tile_op.cc +++ b/src/ascend/transform/lower_tile_op.cc @@ -641,8 +641,8 @@ class LowerTileOpPass : arith::IRMutatorWithAnalyzer { << "but got indices size: " << indices.size() << " and shape size: " << old_shape.size(); - Buffer remap_key = FindRemapBuffer(load->buffer).value_or(load->buffer); - Optional layout = FindLayout(remap_key); + Buffer remap_key = load->buffer; + Optional layout = layout_map_.Get(remap_key); if (!layout.defined() || !buffer_map_.count(remap_key->data)) { return result; } @@ -708,8 +708,8 @@ class LowerTileOpPass : arith::IRMutatorWithAnalyzer { << "but got indices size: " << indices.size() << " and shape size: " << old_shape.size(); - Buffer remap_key = FindRemapBuffer(load->buffer).value_or(load->buffer); - Optional layout = FindLayout(remap_key); + Buffer remap_key = load->buffer; + Optional layout = layout_map_.Get(remap_key); if (!layout.defined() || !buffer_map_.count(remap_key->data)) { return result; } @@ -737,30 +737,24 @@ class LowerTileOpPass : arith::IRMutatorWithAnalyzer { Array new_indices = RemapAccessIndices( indices, old_shape, new_shape, layout.value(), offset); - Array new_args = {BufferLoad(new_buffer, new_indices), extent, - rw_mask}; - if (IsAscendCubeBuffer(remap_key)) { + if (auto pitch = access_ptr_call->annotations.Get(attr::kL0AccessPitch)) { + ICHECK(IsL0ABuffer(remap_key) || IsL0BBuffer(remap_key)); size_t n = new_shape.size(); - PrimExpr physical_offset = 0; - for (size_t i = 0; i < n; ++i) - physical_offset = physical_offset * new_shape[i] + new_indices[i]; - if (auto pitch = - access_ptr_call->annotations.Get(attr::kL0AccessPitch)) { - ICHECK(IsL0ABuffer(remap_key) || IsL0BBuffer(remap_key)); - ICHECK_GE(n, 4); - physical_offset += - new_indices[n - 4] * - (Downcast(pitch.value()) - new_shape[n - 3]) * - new_shape[n - 2] * new_shape[n - 1]; + ICHECK_GE(n, 4); + // A compact L0 copy changes this access's matrix row stride, not the + // allocation or the distance between storage versions. + Array strides(n, Integer(1)); + PrimExpr stride = Integer(1); + for (size_t i = n; i > 0; --i) { + strides.Set(i - 1, stride); + stride = stride * new_shape[i - 1]; } - layout_remap_.Set(new_buffer, layout.value()); - result.rewritten = true; - result.expr = - Call(access_ptr_call->dtype, builtin::tvm_access_ptr(), - {TypeAnnotation(new_buffer->dtype), new_buffer->data, - analyzer_->Simplify(physical_offset), extent, rw_mask}); - return result; + strides.Set(n - 4, Downcast(pitch.value()) * + new_shape[n - 2] * new_shape[n - 1]); + new_buffer.CopyOnWrite()->strides = std::move(strides); } + Array new_args = {BufferLoad(new_buffer, new_indices), extent, + rw_mask}; if (buffer_remap_.count(remap_key)) { layout_remap_.Set(new_buffer, layout.value()); } @@ -789,44 +783,6 @@ class LowerTileOpPass : arith::IRMutatorWithAnalyzer { return Optional(); } - Optional FindRemapBuffer(const Buffer &buffer) const { - if (buffer_remap_.count(buffer)) { - return buffer; - } - auto it = buffer_map_.find(buffer->data); - if (it != buffer_map_.end() && buffer_remap_.count(it->second)) { - return it->second; - } - for (const auto &kv : buffer_remap_) { - if (kv.first->data.same_as(buffer->data)) { - return kv.first; - } - if (kv.first->name == buffer->name) { - return kv.first; - } - } - return Optional(); - } - - Optional FindLayout(const Buffer &buffer) const { - if (layout_map_.count(buffer)) { - return layout_map_[buffer]; - } - auto it = buffer_map_.find(buffer->data); - if (it != buffer_map_.end() && layout_map_.count(it->second)) { - return layout_map_[it->second]; - } - for (const auto &kv : layout_map_) { - if (kv.first->data.same_as(buffer->data)) { - return kv.second; - } - if (kv.first->name == buffer->name) { - return kv.second; - } - } - return Optional(); - } - PrimExpr VisitExpr_(const tirx::CallNode *op) final { if (op->op.same_as(tl::tma_load()) || op->op.same_as(tl::tma_load_im2col()) || diff --git a/src/ascend/transform/normalize_fractal_storage.cc b/src/ascend/transform/normalize_fractal_storage.cc index 79bd979008..bdb554be6b 100644 --- a/src/ascend/transform/normalize_fractal_storage.cc +++ b/src/ascend/transform/normalize_fractal_storage.cc @@ -61,10 +61,9 @@ class BufferCollector : public StmtExprVisitor { } }; -// Ordinary reshapes inherit a canonical matrix interpretation and have their -// indices normalized to it. Independently constrained Cube views retain their -// own interpretation, but share one allocation and an equal physical pitch. -// GM/UB and compact NZ are untouched. +// Layouts and padded matrix coordinates belong to each Buffer. The data Var +// identifies only the allocation, whose physical size must agree across views. +// GM/UB and compact NZ geometry are untouched. class FractalStorageNormalizer : public arith::IRMutatorWithAnalyzer { public: static PrimFunc Rewrite(PrimFunc func) { @@ -73,7 +72,7 @@ class FractalStorageNormalizer : public arith::IRMutatorWithAnalyzer { arith::Analyzer analyzer; analyzer.z3_prover.SetRLimit(50000); FractalStorageNormalizer rewriter(&analyzer); - rewriter.canonical_layouts_ = collector.layouts; + rewriter.layouts_ = collector.layouts; rewriter.MarkBufferMapShapes(func); for (const auto &[buffer, layout] : collector.layouts) { if (!IsAscendFractalLayout(layout) || @@ -81,59 +80,37 @@ class FractalStorageNormalizer : public arith::IRMutatorWithAnalyzer { IsL0CBuffer(buffer)) || !collector.owners.count(buffer->data)) continue; - if (auto anchor = rewriter.anchors_.Get(buffer->data)) { - AscendFractalLayoutInfo lhs, rhs; - ICHECK(TryExtractAscendFractalLayout(collector.layouts[anchor.value()], - anchor.value(), &lhs)); - ICHECK(TryExtractAscendFractalLayout(layout, buffer, &rhs)); - if (!collector.layouts[anchor.value()]->IsEqual(layout.get())) { - PrimExpr lhs_bits = Integer(anchor.value()->dtype.bits()); - PrimExpr rhs_bits = Integer(buffer->dtype.bits()); - for (const PrimExpr &extent : - collector.layouts[anchor.value()]->OutputShape()) - lhs_bits = lhs_bits * extent; - for (const PrimExpr &extent : layout->OutputShape()) - rhs_bits = rhs_bits * extent; - ICHECK(analyzer.CanProveEqual(lhs_bits, rhs_bits)) - << "Cube views sharing storage require equal padded physical " - "sizes: " - << anchor.value()->name << " and " << buffer->name; - rewriter.independent_views_.insert(buffer->data); - } - ICHECK_EQ(anchor.value()->dtype.bits(), buffer->dtype.bits()) - << "Cube aliases must preserve element width: " << buffer->name; - if ((collector.used_buffers.count(buffer) && - !collector.used_buffers.count(anchor.value())) || - (collector.used_buffers.count(buffer) == - collector.used_buffers.count(anchor.value()) && - buffer.same_as(collector.owners[buffer->data]))) - rewriter.anchors_.Set(buffer->data, buffer); - } else { - rewriter.anchors_.Set(buffer->data, buffer); - } + auto views = + rewriter.matrix_views_.Get(buffer->data).value_or(Array{}); + views.push_back(buffer); + rewriter.matrix_views_.Set(buffer->data, views); } + Map storage_bits; for (const Buffer &buffer : collector.buffers) { - if (auto anchor = rewriter.anchors_.Get(buffer->data)) { + if (auto views = rewriter.matrix_views_.Get(buffer->data)) { + Buffer matrix = rewriter.FindMatrixView_( + buffer, views.value(), collector.used_buffers.count(buffer)); ICHECK(!collector.used_buffers.count(buffer) || - buffer->dtype == anchor.value()->dtype) + buffer->dtype == matrix->dtype) << "Cube alias accesses must use the canonical operand dtype: " << buffer->name << " has " << buffer->dtype << ", expected " - << anchor.value()->dtype; - if (rewriter.independent_views_.count(buffer->data)) { - auto layout = collector.layouts.Get(buffer); - if (layout.has_value() && IsAscendFractalLayout(layout.value())) { - anchor = buffer; - } else { - ICHECK(!collector.used_buffers.count(buffer)) - << "Each accessed Cube view sharing storage with distinct " - "layouts requires its own canonical layout: " - << buffer->name; - } - } - rewriter.RegisterBuffer_(buffer, anchor.value(), - collector.owners[buffer->data], - collector.layouts[anchor.value()]); + << matrix->dtype; + rewriter.RegisterBuffer_(buffer, matrix, collector.layouts[matrix]); } + Buffer padded = rewriter.buffers_.Get(buffer).value_or(buffer); + Optional layout = rewriter.layouts_.Get(padded); + if (!collector.owners.count(buffer->data) || !layout.defined() || + layout.value().as()) + continue; + PrimExpr bits = Integer(padded->dtype.bits() * padded->dtype.lanes()); + for (const PrimExpr &extent : layout.value()->OutputShape()) + bits = bits * extent; + if (auto previous = storage_bits.Get(buffer->data)) + ICHECK(analyzer.CanProveEqual(previous.value(), bits)) + << "Buffer views sharing storage require equal physical sizes " + "including padding: " + << buffer->data->name_hint; + storage_bits.Set(buffer->data, bits); } if (rewriter.buffers_.empty()) return func; @@ -146,13 +123,38 @@ class FractalStorageNormalizer : public arith::IRMutatorWithAnalyzer { explicit FractalStorageNormalizer(arith::Analyzer *analyzer) : Parent(analyzer) {} - void RegisterBuffer_(const Buffer &buffer, const Buffer &anchor, - const Buffer &owner, const Layout &layout) { - ICHECK_EQ(buffer->dtype.lanes(), 1); - ICHECK_EQ(buffer->dtype.bits(), anchor->dtype.bits()) + Buffer FindMatrixView_(const Buffer &buffer, const Array &views, + bool accessed) { + ICHECK_EQ(buffer->dtype.bits(), views[0]->dtype.bits()) << "Cube aliases must preserve element width: " << buffer->name - << " has dtype " << buffer->dtype << ", canonical view " << anchor->name - << " has dtype " << anchor->dtype; + << " has dtype " << buffer->dtype; + Optional layout = layouts_.Get(buffer); + if (layout.defined() && IsAscendFractalLayout(layout.value())) + return buffer; + Optional matrix; + for (const Buffer &view : views) { + Layout candidate = layouts_[view]->Reshape(buffer->shape, analyzer_); + if (layout.defined() && !layout.value()->IsEqual(candidate.get())) + continue; + if (!matrix.defined()) + matrix = view; + if (layout.defined() || !accessed) + break; + ICHECK(layouts_[matrix.value()] + ->Reshape(buffer->shape, analyzer_) + ->IsEqual(candidate.get())) + << "Ambiguous layout for Cube view " << buffer->name + << "; annotate this view explicitly"; + } + ICHECK(matrix.defined()) + << "Cube view layout must be a reshape of a canonical matrix layout: " + << buffer->name; + return matrix.value(); + } + + void RegisterBuffer_(const Buffer &buffer, const Buffer &matrix, + const Layout &layout) { + ICHECK_EQ(buffer->dtype.lanes(), 1); ICHECK(buffer->axis_separators.empty() && is_zero(buffer->elem_offset)) << "Ascend fractal buffers require a zero base and no axis separators"; if (!buffer->strides.empty()) { @@ -166,14 +168,14 @@ class FractalStorageNormalizer : public arith::IRMutatorWithAnalyzer { } } Array identity; - for (size_t i = 0; i < anchor->shape.size(); ++i) + for (size_t i = 0; i < matrix->shape.size(); ++i) identity.push_back(InputPlaceholder(i)); logical_maps_.Set( buffer, - Layout(anchor->shape, identity)->Reshape(buffer->shape, analyzer_)); + Layout(matrix->shape, identity)->Reshape(buffer->shape, analyzer_)); AscendFractalLayoutInfo info; - ICHECK(TryExtractAscendFractalLayout(layout, anchor, &info)); - Array shape = anchor->shape; + ICHECK(TryExtractAscendFractalLayout(layout, matrix, &info)); + Array shape = matrix->shape; for (size_t axis = shape.size() - 2; axis < shape.size(); ++axis) { bool is_row = axis == shape.size() - 2; PrimExpr align = (is_row == (info.c0_axis == LogicalAxis::kRow)) @@ -182,19 +184,15 @@ class FractalStorageNormalizer : public arith::IRMutatorWithAnalyzer { shape.Set(axis, analyzer_->Simplify( floordiv(shape[axis] + align - 1, align) * align)); } - // Ordinary aliases coalesce to one Buffer identity and geometry. Distinct - // canonical views retain their own geometry on the same data Var. An unused - // equal-width backing allocation may adopt the canonical operand dtype. - bool independent = independent_views_.count(buffer->data); - Buffer padded = independent - ? buffer - : canonical_buffers_.Get(buffer->data).value_or(owner); + // Preserve each view and the allocation owner's properties. An unused + // equal-width backing allocation may adopt the operand dtype. + Buffer padded = buffer; if (!StructuralEqual()(padded->shape, shape) || !padded->strides.empty() || - padded->dtype != anchor->dtype) { + padded->dtype != matrix->dtype) { auto *writer = padded.CopyOnWrite(); writer->shape = shape; writer->strides = {}; - writer->dtype = anchor->dtype; + writer->dtype = matrix->dtype; } Layout padded_layout; switch (info.kind) { @@ -212,10 +210,8 @@ class FractalStorageNormalizer : public arith::IRMutatorWithAnalyzer { break; } buffers_.Set(buffer, padded); - buffer_anchors_.Set(buffer, anchor); - if (!independent || buffer.same_as(owner)) - canonical_buffers_.Set(buffer->data, padded); - canonical_layouts_.Set(padded, padded_layout); + logical_shapes_.Set(buffer, matrix->shape); + layouts_.Set(padded, padded_layout); } Array RewriteIndices_(const Buffer &buffer, @@ -235,7 +231,7 @@ class FractalStorageNormalizer : public arith::IRMutatorWithAnalyzer { }); if (!buffers_.count(buffer)) return BufferRegion(buffer, ranges); - if (!StructuralEqual()(buffer->shape, buffer_anchors_[buffer]->shape)) { + if (!StructuralEqual()(buffer->shape, logical_shapes_[buffer])) { // A dense reshape is injective only while trailing region extents do // not wrap across the original view's strides. Otherwise equal bounding // box volume alone can hide duplicated logical elements. @@ -278,11 +274,11 @@ class FractalStorageNormalizer : public arith::IRMutatorWithAnalyzer { if (IsL0SFBuffer(copy->dst)) { // SF handles carry the effective shape of their bound data tile and // have no independent storage. Only the L1 source is a transfer region. - source = ascend::GetL1ToL0SFSourceRegion(*copy.get(), canonical_layouts_, - analyzer_); + source = + ascend::GetL1ToL0SFSourceRegion(*copy.get(), layouts_, analyzer_); } else { - auto regions = ascend::GetL1ToL0CopyRegions( - *copy.get(), canonical_layouts_, analyzer_); + auto regions = + ascend::GetL1ToL0CopyRegions(*copy.get(), layouts_, analyzer_); source = regions.source; destination = regions.destination; pitch = regions.pitch; @@ -296,9 +292,8 @@ class FractalStorageNormalizer : public arith::IRMutatorWithAnalyzer { // than the actual transfer. This also makes repeated normalization // stable. AscendFractalRegionInfo normalized; - ICHECK(TryExtractAscendFractalRegion(canonical_layouts_[copy->dst], - destination.value()->region, - &normalized)); + ICHECK(TryExtractAscendFractalRegion( + layouts_[copy->dst], destination.value()->region, &normalized)); if (!analyzer_->CanProveEqual(normalized.outer1->extent, pitch.value())) writer->annotations.Set("l0_pitch", pitch.value()); } @@ -487,7 +482,7 @@ class FractalStorageNormalizer : public arith::IRMutatorWithAnalyzer { // GM region unchanged: extra storage does not authorize an extra read. Copy rewritten(call->args, call->annotations); Array ranges = rewritten->dst_range; - const auto &logical_shape = buffer_anchors_[copy->dst]->shape; + const auto &logical_shape = logical_shapes_[copy->dst]; for (size_t i = 0; i < ranges.size(); ++i) { if (analyzer_->CanProveEqual(ranges[i]->min, 0) && analyzer_->CanProveEqual(ranges[i]->extent, logical_shape[i])) @@ -502,21 +497,23 @@ class FractalStorageNormalizer : public arith::IRMutatorWithAnalyzer { } if (op->op.same_as(builtin::tvm_access_ptr())) { auto data = op->args[1].as(); - if (data && anchors_.count(data.value())) { - ICHECK(!independent_views_.count(data.value())) - << "Raw Cube pointers cannot select among distinct view layouts; " - "use T.access_ptr(view[...], ...)"; - Buffer anchor = anchors_[data.value()]; - ICHECK(op->args[0].dtype() == anchor->dtype) + if (data && matrix_views_.count(data.value())) { + const auto &views = matrix_views_[data.value()]; + Buffer matrix = views[0]; + for (const Buffer &view : views) + ICHECK(layouts_[matrix]->IsEqual(layouts_[view].get())) + << "Raw Cube pointers cannot select among distinct view layouts; " + "use T.access_ptr(view[...], ...)"; + ICHECK(op->args[0].dtype() == matrix->dtype) << "Cube raw pointer must use the canonical operand dtype " - << anchor->dtype << "; take the pointer through that view"; - Buffer target = canonical_buffers_[data.value()]; + << matrix->dtype << "; take the pointer through that view"; + Buffer target = buffers_[matrix]; PrimExpr remaining = VisitExpr(op->args[2]); - Array indices(anchor->shape.size(), Integer(0)); - for (int i = static_cast(anchor->shape.size()) - 1; i >= 0; --i) { + Array indices(matrix->shape.size(), Integer(0)); + for (int i = static_cast(matrix->shape.size()) - 1; i >= 0; --i) { indices.Set(i, i == 0 ? remaining - : floormod(remaining, anchor->shape[i])); - remaining = floordiv(remaining, anchor->shape[i]); + : floormod(remaining, matrix->shape[i])); + remaining = floordiv(remaining, matrix->shape[i]); } // Preserve the caller's extent; only the logical base address changes. return Call(op->dtype, tl::access_ptr(), @@ -546,19 +543,16 @@ class FractalStorageNormalizer : public arith::IRMutatorWithAnalyzer { for (const auto &[buffer, layout] : value.value().cast>()) { Buffer rewritten = buffers_.Get(buffer).value_or(buffer); - layouts.Set(rewritten, - canonical_layouts_.Get(rewritten).value_or(layout)); + layouts.Set(rewritten, layouts_.Get(rewritten).value_or(layout)); } block.CopyOnWrite()->annotations.Set(attr::kLayoutMap, layouts); } return block; } - Map anchors_; - Map buffer_anchors_; - std::unordered_set independent_views_; - Map canonical_buffers_; + Map> matrix_views_; + Map> logical_shapes_; Map buffers_; - LayoutMap canonical_layouts_; + LayoutMap layouts_; Map logical_maps_; }; } // namespace diff --git a/testing/ascend/transform/test_ascend_buffer_view_layouts.py b/testing/ascend/transform/test_ascend_buffer_view_layouts.py new file mode 100644 index 0000000000..8ff7ae1ad1 --- /dev/null +++ b/testing/ascend/transform/test_ascend_buffer_view_layouts.py @@ -0,0 +1,114 @@ +"""Layouts belong to buffer views; pointers and allocations retain that split.""" + +import pytest +import tilelang +import tilelang.testing +from tilelang import tvm +from tilelang.ascend import language as T, transform +from tilelang.layout import Layout +from tvm import tirx + +from testing.ascend._ir import calls, kernel + + +def _views(): + owner = tirx.decl_buffer((8, 16), "float32", name="owner", scope="shared") + view = tirx.decl_buffer(owner.shape, owner.dtype, name="view", data=owner.data) + return owner, view, Layout(owner.shape, lambda i, j: (i, j)), Layout(view.shape, lambda i, j: (j, i)) + + +@pytest.mark.parametrize("region", [None, "SIMT_VF", "SIMD_VF"]) +def test_explicit_layouts_remain_attached_to_each_view(region): + owner, view, first, second = _views() + body = tirx.SeqStmt([tirx.Evaluate(owner[3, 5]), tirx.Evaluate(view[3, 5])]) + annotations = {"layout_map": {owner: first, view: second}} + if region: + if region == "SIMT_VF": + tx = tirx.IterVar(tvm.ir.Range(0, 32), tirx.Var("tx", "int32"), tirx.IterVar.ThreadIndex, "threadIdx.x") + body = tirx.AttrStmt(tx, "thread_extent", 32, body) + body = tirx.SBlockRealize([], True, tirx.SBlock([], [], [], region, body, annotations=annotations)) + annotations = {} + mod = kernel(body, buffers=[owner], annotations=annotations) + with tvm.target.Target("ascend"): + inferred = transform.AscendLayoutInference()(mod) + layouts = inferred["main"].body.block.annotations["layout_map"] + assert layouts[owner].is_equal(first) + assert layouts[view].is_equal(second) + + +def test_unconstrained_reshape_inherits_the_unique_layout(): + owner, _, layout, _ = _views() + view = tirx.decl_buffer((4, 32), owner.dtype, name="reshaped", data=owner.data) + mod = kernel(tirx.Evaluate(view[1, 7]), buffers=[owner], annotations={"layout_map": {owner: layout}}) + with tvm.target.Target("ascend"): + inferred = transform.AscendLayoutInference()(mod) + layouts = inferred["main"].body.block.annotations["layout_map"] + assert layouts[view].is_equal(layout.reshape(view.shape)) + + +def test_unconstrained_access_cannot_choose_between_distinct_layouts(): + owner, view, first, second = _views() + unknown = tirx.decl_buffer(owner.shape, owner.dtype, name="unknown", data=owner.data) + body = tirx.SeqStmt([tirx.Evaluate(buf[0, 0]) for buf in (owner, view, unknown)]) + mod = kernel(body, buffers=[owner], annotations={"layout_map": {owner: first, view: second}}) + with tvm.target.Target("ascend"), pytest.raises(tvm.error.InternalError, match="Ambiguous layout for buffer view unknown"): + transform.AscendLayoutInference()(mod) + + +def test_unused_storage_owner_does_not_need_a_view_layout(): + owner, first_view, _, _ = _views() + second_view = tirx.decl_buffer(owner.shape, owner.dtype, name="second", data=owner.data) + first = Layout(first_view.shape, lambda i, j: (i, j)) + second = Layout(second_view.shape, lambda i, j: (j, i)) + body = tirx.SeqStmt([tirx.Evaluate(buf[0, 0]) for buf in (first_view, second_view)]) + mod = kernel(body, buffers=[owner], annotations={"layout_map": {first_view: first, second_view: second}}) + with tvm.target.Target("ascend"): + inferred = transform.AscendLayoutInference()(mod) + layouts = inferred["main"].body.block.annotations["layout_map"] + assert layouts[first_view].is_equal(first) + assert layouts[second_view].is_equal(second) + + +def test_view_layouts_must_preserve_the_storage_footprint(): + owner, view, first, _ = _views() + second = Layout(view.shape, lambda i, j: (i, j * 2)) + body = tirx.SeqStmt([tirx.Evaluate(buf[0, 0]) for buf in (owner, view)]) + mod = kernel(body, buffers=[owner], annotations={"layout_map": {owner: first, view: second}}) + with tvm.target.Target("ascend"), pytest.raises(tvm.error.InternalError, match="equal physical sizes"): + transform.NormalizeAscendFractalStorage()(transform.AscendLayoutInference()(mod)) + + +def test_typed_pointers_use_each_view_layout_and_one_storage(): + owner, view, first, second = _views() + body = tirx.SeqStmt([tirx.Evaluate(T.access_ptr(buf[3, 5], "r", extent=1)) for buf in (owner, view)]) + mod = kernel(body, buffers=[owner], annotations={"layout_map": {owner: first, view: second}}) + with tvm.target.Target("ascend"): + mod = transform.AscendLowerTileOp()(mod) + mod = tilelang.transform.LowerAccessPtr()(mod) + pointers = calls(mod, "tirx.tvm_access_ptr") + assert len(pointers) == 2 + assert pointers[0].args[1].same_as(pointers[1].args[1]) + analyzer = tvm.arith.Analyzer() + assert analyzer.can_prove_equal(pointers[0].args[2], 3 * 16 + 5) + assert analyzer.can_prove_equal(pointers[1].args[2], 5 * 8 + 3) + + +@tilelang.testing.requires_ascend +def test_ub_view_reads_the_same_bytes_through_its_own_layout(): + import torch + + @T.prim_func + def transpose(A: T.Tensor((8, 16), "float32"), B: T.Tensor((16, 8), "float32")): + with T.Kernel(1): + storage = T.alloc_shared((8, 16), "float32") + view = T.view(storage, (16, 8)) + T.annotate_layout({storage: Layout((8, 16), lambda i, j: (i, j)), view: Layout((16, 8), lambda i, j: (j, i))}) + with T.SimtVF(threads=128): + for i, j in T.Parallel(8, 16): + storage[i, j] = A[i, j] + for i, j in T.Parallel(16, 8): + B[i, j] = view[i, j] + + compiled = tilelang.compile(transpose, target="ascend", out_idx=-1) + source = torch.arange(128, dtype=torch.float32, device="npu").reshape(8, 16) + torch.testing.assert_close(compiled(source), source.T, rtol=0, atol=0) diff --git a/testing/ascend/transform/test_ascend_lower_copy.py b/testing/ascend/transform/test_ascend_lower_copy.py index 525f4848df..29b0170db4 100644 --- a/testing/ascend/transform/test_ascend_lower_copy.py +++ b/testing/ascend/transform/test_ascend_lower_copy.py @@ -3,7 +3,7 @@ import math import pytest -from tilelang import tvm +from tilelang import tvm, transform as tl_transform from tilelang.ascend import language as T, transform from tilelang.layout import ( make_ascend_nz_layout, @@ -52,7 +52,7 @@ def _lower(src, dst, src_region=None, dst_region=None, *, layouts=None, bindings mod = tvm.IRModule({"main": func}) if normalize: mod = transform.NormalizeAscendFractalStorage()(mod) - return transform.AscendLowerTileOp()(mod) + return tl_transform.LowerAccessPtr()(transform.AscendLowerTileOp()(mod)) def _equal(actual, expected): @@ -297,7 +297,7 @@ def test_explicit_l0_pitch_controls_offsets_and_keeps_plane_capacity(v_layout, o with pytest.raises(ValueError, match="l0_pitch"): transform.AscendLowerTileOp()(mod) return - after = transform.AscendLowerTileOp()(mod) + after = tl_transform.LowerAccessPtr()(transform.AscendLowerTileOp()(mod)) loads = calls(after, "tl.ascend_load_cbuf_to_cb" if v_layout else "tl.ascend_load_cbuf_to_ca") _equal([load.args[7] for load in loads], [pitch, pitch]) _equal([load.args[0].args[2] for load in loads], [4096, 4096 + 2 * pitch * 256]) diff --git a/testing/ascend/transform/test_ascend_normalize_fractal_storage.py b/testing/ascend/transform/test_ascend_normalize_fractal_storage.py index 27d1093e52..3c7c9612fb 100644 --- a/testing/ascend/transform/test_ascend_normalize_fractal_storage.py +++ b/testing/ascend/transform/test_ascend_normalize_fractal_storage.py @@ -35,7 +35,7 @@ def test_canonical_allocation_rounds_only_fractal_axes(shape, dtype, padded): @pytest.mark.parametrize("kind", ["view", "reshape", "reinterpret"]) -def test_alias_accesses_use_one_canonical_allocation(kind): +def test_alias_accesses_preserve_views_on_one_allocation(kind): owner_shape = (15, 15) if kind == "view" else (225,) owner = tirx.decl_buffer(owner_shape, "uint16" if kind == "reinterpret" else "bfloat16", name="owner", scope="shared.l1") view = tirx.decl_buffer((15, 15), "bfloat16", data=owner.data) @@ -43,7 +43,8 @@ def test_alias_accesses_use_one_canonical_allocation(kind): after = transform.NormalizeAscendFractalStorage()(before) allocation = after["main"].body.block.alloc_buffers[0] (store,) = nodes(after, tirx.BufferStore) - assert store.buffer.same_as(allocation) + assert not store.buffer.same_as(allocation) + assert store.buffer.data.same_as(allocation.data) assert allocation.data.same_as(owner.data) assert allocation.dtype == "bfloat16" assert tuple(int(x) for x in allocation.shape) == (16, 16) @@ -61,7 +62,7 @@ def test_pointer_base_is_reexpressed_in_canonical_indices(raw): after = transform.NormalizeAscendFractalStorage()(before) (pointer,) = calls(after, "tl.access_ptr") load, extent, _ = pointer.args - assert load.buffer.same_as(after["main"].body.block.alloc_buffers[0]) + assert load.buffer.data.same_as(after["main"].body.block.alloc_buffers[0].data) analyzer = tvm.arith.Analyzer() assert analyzer.can_prove_equal(load.indices[0], 1 if raw else row * 5) assert int(load.indices[1]) == 0 and int(extent) == 1 @@ -126,7 +127,7 @@ def test_alias_indices_and_access_regions_use_the_same_padded_storage(): root = normalized["main"].body.block owner = root.alloc_buffers[0] alias = root.reads[0].buffer - assert owner.same_as(alias) + assert not owner.same_as(alias) assert owner.data.same_as(alias.data) assert tuple(int(x) for x in owner.shape) == tuple(int(x) for x in alias.shape) == (16, 16) assert [(int(r.min), int(r.extent)) for r in root.reads[0].region] == [(1, 2), (0, 15)] @@ -158,7 +159,10 @@ def test_alias_normalization_preserves_allocation_properties(): assert allocation.offset_factor == 8 assert allocation.dtype == "bfloat16" assert tuple(int(x) for x in allocation.shape) == (16, 16) - assert allocation.same_as(normalized["main"].body.block.body.value.buffer) + view_buffer = normalized["main"].body.block.body.value.buffer + assert not allocation.same_as(view_buffer) + assert allocation.data.same_as(view_buffer.data) + assert allocation.name == "owner" and view_buffer.name == "view" @pytest.mark.parametrize( @@ -188,9 +192,17 @@ def test_independent_cube_views_preserve_their_matrix_geometry(shape, view_shape tvm.ir.assert_structural_equal(after, transform.NormalizeAscendFractalStorage()(after)) -def test_cube_views_with_different_padded_sizes_are_rejected(): - owner = tirx.decl_buffer((15, 16), "bfloat16", scope="shared.l0b") - view = tirx.decl_buffer((8, 30), "bfloat16", data=owner.data) +@pytest.mark.parametrize( + "dtype,shape,view_shape,view_layout", + [ + ("bfloat16", (15, 16), (8, 30), make_ascend_nz_layout), + # Both unpadded layouts cover 128 elements, but their Cube padding differs. + ("float32", (8, 16), (8, 16), make_ascend_major_mn_layout), + ], +) +def test_cube_views_with_different_padded_sizes_are_rejected(dtype, shape, view_shape, view_layout): + owner = tirx.decl_buffer(shape, dtype, scope="shared.l0b") + view = tirx.decl_buffer(view_shape, dtype, data=owner.data) block = tirx.SBlock( [], [], @@ -198,15 +210,39 @@ def test_cube_views_with_different_padded_sizes_are_rejected(): "root", tirx.Evaluate(view[0, 0]), alloc_buffers=[owner], - annotations={ - "layout_map": {owner: tilelang.layout.make_ascend_major_k_layout(owner), view: tilelang.layout.make_ascend_major_k_layout(view)} - }, + annotations={"layout_map": {owner: make_ascend_nz_layout(owner), view: view_layout(view)}}, ) before = tvm.IRModule({"main": tirx.PrimFunc([], tirx.SBlockRealize([], True, block))}) - with pytest.raises(tvm.error.InternalError, match="equal padded physical sizes"): + with pytest.raises(tvm.error.InternalError, match="equal physical sizes including padding"): transform.NormalizeAscendFractalStorage()(before) +@pytest.mark.parametrize("source", ["owner", "view"]) +@pytest.mark.parametrize("reverse_annotations", [False, True]) +def test_reshape_uses_its_layout_when_storage_has_distinct_matrix_views(source, reverse_annotations): + owner = tirx.decl_buffer((17, 32), "bfloat16", name="owner", scope="shared.l0b") + view = tirx.decl_buffer((32, 17), "bfloat16", name="view", data=owner.data) + reshaped = tirx.decl_buffer((544,), "bfloat16", name="reshaped", data=owner.data) + layouts = {owner: make_ascend_nz_layout(owner), view: make_ascend_major_mn_layout(view)} + matrix = owner if source == "owner" else view + layouts[reshaped] = layouts[matrix].reshape(reshaped.shape) + if reverse_annotations: + layouts = dict(reversed(list(layouts.items()))) + body = tirx.SeqStmt([tirx.Evaluate(buf[1, 0]) for buf in (owner, view)] + [tirx.Evaluate(reshaped[35])]) + before = kernel(body, buffers=[owner], annotations={"layout_map": layouts}) + after = transform.NormalizeAscendFractalStorage()(before) + root = after["main"].body.block + read = root.body.seq[2].value + assert tuple(int(x) for x in read.indices) == ((1, 3) if source == "owner" else (2, 1)) + assert tuple(int(x) for x in read.buffer.shape) == (32, 32) + assert read.buffer.name == "reshaped" + assert not read.buffer.same_as(root.alloc_buffers[0]) + assert all(stmt.value.buffer.data.same_as(root.alloc_buffers[0].data) for stmt in root.body.seq) + expected = (make_ascend_nz_layout if source == "owner" else make_ascend_major_mn_layout)(read.buffer) + assert root.annotations["layout_map"][read.buffer].is_equal(expected) + tvm.ir.assert_structural_equal(after, transform.NormalizeAscendFractalStorage()(after)) + + @pytest.mark.parametrize("raw", [False, True], ids=["typed-view", "ambiguous-data-pointer"]) def test_distinct_cube_layouts_require_typed_pointers(raw): owner = tirx.decl_buffer((64, 128), "bfloat16", scope="shared.l0b") diff --git a/tilelang/ascend/op/gemm/gemm_mad.py b/tilelang/ascend/op/gemm/gemm_mad.py index aeda0808ed..004250390b 100644 --- a/tilelang/ascend/op/gemm/gemm_mad.py +++ b/tilelang/ascend/op/gemm/gemm_mad.py @@ -41,19 +41,7 @@ def _make_access_ptr(buf, region, rw_mask, exclude_last=False): def _find_layout(layout_map, buf): - if layout_map is None: - return None - for key in (buf, buf.data if hasattr(buf, "data") else None): - if key is not None: - layout = layout_map.get(key) - if layout is not None: - return layout - for key, layout in layout_map.items(): - if hasattr(key, "data") and hasattr(buf, "data") and key.data.same_as(buf.data): - return layout - if hasattr(key, "name") and getattr(key, "name", None) == getattr(buf, "name", None): - return layout - return None + return None if layout_map is None else layout_map.get(buf) _ASCEND_DTYPE_MAP = {