From fb879ca7ccd206183fb467e02f2011a33cd4952e Mon Sep 17 00:00:00 2001 From: Denver Jin Date: Tue, 13 Oct 2026 16:52:11 +0800 Subject: [PATCH 1/6] [Ascend] Add repeat and stride access region overlap framework --- .../auto_schedule/dependency_analysis.cc | 68 +++++++++++----- .../auto_schedule/dependency_analysis.h | 8 +- .../transform/auto_schedule/ir_structure.cc | 4 +- .../transform/auto_schedule/ir_structure.h | 77 +++++++++++-------- .../transform/auto_schedule/memory_detector.h | 44 +++++++++++ .../transform/auto_schedule/task_analysis.h | 25 +++--- src/ascend/transform/resolve_core.cc | 3 +- 7 files changed, 161 insertions(+), 68 deletions(-) diff --git a/src/ascend/transform/auto_schedule/dependency_analysis.cc b/src/ascend/transform/auto_schedule/dependency_analysis.cc index 1e2e3bcf4a..fbeaae38b6 100644 --- a/src/ascend/transform/auto_schedule/dependency_analysis.cc +++ b/src/ascend/transform/auto_schedule/dependency_analysis.cc @@ -112,11 +112,13 @@ std::vector StorageIterationNest(ControlNode *loop, return {controls.begin() + first, controls.end()}; } -bool RegionsMayConflict(const ConstrSet &a_ctx, const BufferRegion &a_region, - const ConstrSet &b_ctx, const BufferRegion &b_region, +bool RegionsMayConflict(const ConstrSet &a_ctx, const AccessRegion &a_access, + const ConstrSet &b_ctx, const AccessRegion &b_access, ControlNode *loop, int offset, size_t num_storage_owners, const ConflictHintList &root_conflicts) { + const BufferRegion &a_region = a_access.logical; + const BufferRegion &b_region = b_access.logical; const bool cross = offset != 0; if (cross && loop == nullptr) return false; @@ -150,8 +152,8 @@ bool RegionsMayConflict(const ConstrSet &a_ctx, const BufferRegion &a_region, if (!a_region->buffer.same_as(b_region->buffer)) return true; - const Array &a_rng = a_region->region; - const Array &b_rng = b_region->region; + const Array &a_rng = a_access.bounds; + const Array &b_rng = b_access.bounds; if (a_rng.size() != b_rng.size()) return true; // shape mismatch — be safe const size_t ndim = a_rng.size(); @@ -208,21 +210,49 @@ bool RegionsMayConflict(const ConstrSet &a_ctx, const BufferRegion &a_region, } loop_ctx.Merge(a_local_ctx).Merge(b_local_ctx).Populate(ana); - // Freshen mutable reads (BufferLoads etc.) inline in the region bounds, per - // side: a read written verbatim in both bounds is not provably equal across - // the two access points (a store may sit between them), so it must become an - // independent unknown on each side. + // Snapshot bounds, repeat, and stride together on each side. Identical + // mutable reads at two access points need not have the same value. FreshenMutableReads freshen_a(FreshenMutableReads::Mode::kSnapshot); FreshenMutableReads freshen_b(FreshenMutableReads::Mode::kSnapshot); - PrimExpr any_dim_disjoint = Bool(false); + PrimExpr a_repeat = freshen_a(Substitute(a_access.repeat, a_sub)); + PrimExpr b_repeat = freshen_b(Substitute(b_access.repeat, b_sub)); + + // Each side chooses an independent interval from its repeated last axis. + // Proving disjointness for every permitted pair is equivalent to proving + // that no pair overlaps; no enumeration or explicit quantifier is needed. + auto repeat_min = [&](PrimExpr min, const PrimExpr &repeat, + const PrimExpr &stride, const char *name) -> PrimExpr { + if (is_one(ana.Simplify(repeat))) + return min; + // Widen before multiplying so int32 counts/pitches do not truncate the + // address calculation. Zero and negative strides need no special case. + Var index(name, DataType::Int(64)); + ana.EnterConstraint(index >= 0 && index < cast(index.dtype(), repeat)); + return cast(index.dtype(), min) + index * cast(index.dtype(), stride); + }; + + PrimExpr any_dim_disjoint = (a_repeat <= 0) || (b_repeat <= 0); for (size_t i = 0; i < ndim; ++i) { PrimExpr a_min = freshen_a(Substitute(a_rng[i]->min, a_sub)); - PrimExpr a_max = - freshen_a(Substitute(a_rng[i]->min + a_rng[i]->extent - 1, a_sub)); + PrimExpr a_extent = freshen_a(Substitute(a_rng[i]->extent, a_sub)); PrimExpr b_min = freshen_b(Substitute(b_rng[i]->min, b_sub)); - PrimExpr b_max = - freshen_b(Substitute(b_rng[i]->min + b_rng[i]->extent - 1, b_sub)); - any_dim_disjoint = any_dim_disjoint || (a_max < b_min) || (b_max < a_min); + PrimExpr b_extent = freshen_b(Substitute(b_rng[i]->extent, b_sub)); + if (i + 1 == ndim) { + a_min = repeat_min(a_min, a_repeat, + freshen_a(Substitute(a_access.stride, a_sub)), + "region_repeat_p"); + b_min = repeat_min(b_min, b_repeat, + freshen_b(Substitute(b_access.stride, b_sub)), + "region_repeat_c"); + } + // A half-open endpoint can exceed the int32 index range by one even + // for an ordinary single-element access. + PrimExpr a_end = + cast(DataType::Int(64), a_min) + cast(DataType::Int(64), a_extent); + PrimExpr b_end = + cast(DataType::Int(64), b_min) + cast(DataType::Int(64), b_extent); + any_dim_disjoint = any_dim_disjoint || (a_extent <= 0) || (b_extent <= 0) || + (a_end <= b_min) || (b_end <= a_min); } any_dim_disjoint = ana.Simplify(any_dim_disjoint); // Raw Z3 as a fallback catches the infeasible-context case that Analyzer's @@ -254,7 +284,7 @@ AnalyzeDependencies(std::vector nodes, ControlNode *loop, struct RegionAccess { TaskNode *task; - BufferRegion region; + AccessRegion region; }; struct NodeAccesses { std::vector reads; @@ -265,9 +295,9 @@ AnalyzeDependencies(std::vector nodes, ControlNode *loop, std::vector tasks; CollectAllTaskNodes(nodes[i], tasks); for (TaskNode *task : tasks) { - for (const BufferRegion ®ion : task->GetReadRegions()) + for (const AccessRegion ®ion : task->GetReadAccesses()) accesses[i].reads.push_back({task, region}); - for (const BufferRegion ®ion : task->GetWriteRegions()) + for (const AccessRegion ®ion : task->GetWriteAccesses()) accesses[i].writes.push_back({task, region}); } } @@ -372,8 +402,8 @@ AnalyzeDependencies(std::vector nodes, ControlNode *loop, const std::vector &rhs_accesses) { for (const RegionAccess &lhs : lhs_accesses) { for (const RegionAccess &rhs : rhs_accesses) { - const Var &lhs_storage = lhs.region->buffer->data; - const Var &rhs_storage = rhs.region->buffer->data; + const Var &lhs_storage = lhs.region.logical->buffer->data; + const Var &rhs_storage = rhs.region.logical->buffer->data; ffi::Optional storage; if (lhs_storage.same_as(rhs_storage)) storage = lhs_storage; diff --git a/src/ascend/transform/auto_schedule/dependency_analysis.h b/src/ascend/transform/auto_schedule/dependency_analysis.h index e92ecbb5d1..441429d707 100644 --- a/src/ascend/transform/auto_schedule/dependency_analysis.h +++ b/src/ascend/transform/auto_schedule/dependency_analysis.h @@ -109,7 +109,9 @@ AnalyzeDependencies(std::vector nodes, const ConflictHintList &root_conflicts = {}); // Check whether two task regions may overlap under their respective symbolic -// contexts. `offset == 0` compares accesses in the same iteration. A positive +// contexts, including independent repeats of their last conflict axis. Only a +// proof of disjointness removes a conflict; solver unknowns remain conflicts. +// `offset == 0` compares accesses in the same iteration. A positive // offset constrains the current consumer loop coordinate to equal the producer // coordinate plus that offset. A negative offset denotes an automatic // cross-iteration check. Normally it renames the loop suffix below the nearest @@ -119,8 +121,8 @@ AnalyzeDependencies(std::vector nodes, // an identically-shaped iteration in another owner can still be a distinct // epoch. Shared by scheduling-time IRStructure analysis and the scheduled-TIR // synchronization pass. -bool RegionsMayConflict(const ConstrSet &a_ctx, const BufferRegion &a_region, - const ConstrSet &b_ctx, const BufferRegion &b_region, +bool RegionsMayConflict(const ConstrSet &a_ctx, const AccessRegion &a_access, + const ConstrSet &b_ctx, const AccessRegion &b_access, ControlNode *loop, int offset, size_t num_storage_owners = 0, const ConflictHintList &root_conflicts = {}); diff --git a/src/ascend/transform/auto_schedule/ir_structure.cc b/src/ascend/transform/auto_schedule/ir_structure.cc index 774510534e..92fbb10c08 100644 --- a/src/ascend/transform/auto_schedule/ir_structure.cc +++ b/src/ascend/transform/auto_schedule/ir_structure.cc @@ -297,8 +297,8 @@ std::shared_ptr TaskNode::Clone() const { if (IsPerCoreTask()) new_task->MarkPerCoreTask(); // Copy memory access regions - new_task->SetReadRegions(GetReadRegions()); - new_task->SetWriteRegions(GetWriteRegions()); + new_task->SetReadAccesses(GetReadAccesses()); + new_task->SetWriteAccesses(GetWriteAccesses()); new_task->SetReadVars(GetReadVars()); new_task->SetWriteVars(GetWriteVars()); new_task->outer_ctx = outer_ctx; diff --git a/src/ascend/transform/auto_schedule/ir_structure.h b/src/ascend/transform/auto_schedule/ir_structure.h index c749e4258b..9c7e516c08 100644 --- a/src/ascend/transform/auto_schedule/ir_structure.h +++ b/src/ascend/transform/auto_schedule/ir_structure.h @@ -25,6 +25,7 @@ #include "../ascend_pipe.h" #include "../buffer_version.h" #include "../core_mask.h" +#include "./memory_detector.h" #include "op/utils.h" namespace tvm { @@ -110,6 +111,15 @@ inline bool RegionsEqual(const Region &a, const Region &b) { return true; } +inline bool AccessRegionsEqual(const AccessRegion &a, const AccessRegion &b) { + arith::Analyzer analyzer; + return a.logical->buffer.same_as(b.logical->buffer) && + RegionsEqual(a.logical->region, b.logical->region) && + RegionsEqual(a.bounds, b.bounds) && + analyzer.CanProveEqual(a.repeat, b.repeat) && + analyzer.CanProveEqual(a.stride, b.stride); +} + // Return true only when every point in `inner` is provably contained in // `outer`. This is stricter than overlap and is suitable for checking whether // one task's declared write footprint fully defines another task's read. @@ -375,7 +385,8 @@ class IRStructure { // only TaskNode stores a settable mask. virtual uint16_t GetHbmMask() const = 0; - // Logical access regions. Physical footprints are derived from current IR. + // Logical projections for storage identity and core analysis. Dependency + // analysis consumes the leaf TaskNode's complete access records. virtual std::vector GetReadRegions() const = 0; virtual std::vector GetWriteRegions() const = 0; @@ -529,13 +540,20 @@ class TaskNode : public IRStructure { // A copy that stores to GM in this leaf reads the atomic-mode register. bool ReadsAtomicRegister() const override { return reads_atomic_; } - // Logical access regions. Physical footprints are derived from current IR. - std::vector GetReadRegions() const override { + const std::vector &GetReadAccesses() const { return read_regions_; } - std::vector GetWriteRegions() const override { + const std::vector &GetWriteAccesses() const { return write_regions_; } + + // Logical projections retain the union's individual records. + std::vector GetReadRegions() const override { + return LogicalRegions(read_regions_); + } + std::vector GetWriteRegions() const override { + return LogicalRegions(write_regions_); + } std::vector GetReadVars() const override { return read_vars_; } std::vector GetWriteVars() const override { return write_vars_; } @@ -550,10 +568,10 @@ class TaskNode : public IRStructure { void AddSpecialWrite(SpecialRegister reg) { special_write_mask_ |= static_cast(reg); } - void SetReadRegions(const std::vector ®ions) { + void SetReadAccesses(const std::vector ®ions) { read_regions_ = regions; } - void SetWriteRegions(const std::vector ®ions) { + void SetWriteAccesses(const std::vector ®ions) { write_regions_ = regions; } void SetReadVars(const std::vector &vars) { read_vars_ = vars; } @@ -568,9 +586,9 @@ class TaskNode : public IRStructure { }; std::unordered_map buffer_remap; - auto substitute_regions = [&](std::vector ®ions) { - for (auto ®ion : regions) { - Buffer buffer = region->buffer; + auto substitute_regions = [&](std::vector ®ions) { + for (AccessRegion &access : regions) { + Buffer buffer = access.logical->buffer; auto buffer_it = buffer_remap.find(buffer); if (buffer_it != buffer_remap.end()) { buffer = buffer_it->second; @@ -580,8 +598,12 @@ class TaskNode : public IRStructure { buffer_remap.emplace(buffer, remapped); buffer = std::move(remapped); } - region = BufferRegion(std::move(buffer), - Substitute(region->region, substitution)); + access.logical = + BufferRegion(std::move(buffer), + Substitute(access.logical->region, substitution)); + access.bounds = Substitute(access.bounds, substitution); + access.repeat = Substitute(access.repeat, substitution); + access.stride = Substitute(access.stride, substitution); } }; substitute_vars(read_vars_); @@ -602,27 +624,21 @@ class TaskNode : public IRStructure { // Clone method std::shared_ptr Clone() const override; - // Helper methods to add regions (for incremental analysis) - void AddReadRegion(const BufferRegion ®ion) { - // Check for duplicate regions - for (const auto &existing : read_regions_) { - if (existing->buffer.same_as(region->buffer) && - RegionsEqual(existing->region, region->region)) { - return; // Region already exists - } + // Keep the existing deduplication, comparing the complete access record. + void AddReadAccess(const AccessRegion &access) { + for (const AccessRegion &existing : read_regions_) { + if (AccessRegionsEqual(existing, access)) + return; } - read_regions_.push_back(region); + read_regions_.push_back(access); } - void AddWriteRegion(const BufferRegion ®ion) { - // Check for duplicate regions - for (const auto &existing : write_regions_) { - if (existing->buffer.same_as(region->buffer) && - RegionsEqual(existing->region, region->region)) { - return; // Region already exists - } + void AddWriteAccess(const AccessRegion &access) { + for (const AccessRegion &existing : write_regions_) { + if (AccessRegionsEqual(existing, access)) + return; } - write_regions_.push_back(region); + write_regions_.push_back(access); } void AddReadVar(const Var &var) { @@ -654,9 +670,8 @@ class TaskNode : public IRStructure { // inferred or explicitly marked candidate. bool is_per_core_task_{false}; - // Logical access regions. Physical footprints are derived from current IR. - std::vector read_regions_; - std::vector write_regions_; + std::vector read_regions_; + std::vector write_regions_; std::vector read_vars_; std::vector write_vars_; diff --git a/src/ascend/transform/auto_schedule/memory_detector.h b/src/ascend/transform/auto_schedule/memory_detector.h index 681de2e96a..dd599bacf3 100644 --- a/src/ascend/transform/auto_schedule/memory_detector.h +++ b/src/ascend/transform/auto_schedule/memory_detector.h @@ -35,6 +35,35 @@ namespace ascend { using namespace tirx; +// One detected access. A list of these records denotes a union, including +// when several records refer to the same logical buffer. +struct AccessRegion { + explicit AccessRegion(BufferRegion logical) + : logical(std::move(logical)), bounds(this->logical->region), + repeat(Integer(1)), stride(Integer(0)) {} + + // Buffer identity, conflict hints, and logical cost accounting. + BufferRegion logical; + // Conflict coordinates, independent of the logical buffer's rank. All + // accesses to one buffer must use the same coordinate convention. + Region bounds; + // Repeat the last axis at [min + i * stride, min + i * stride + extent), + // 0 <= i < repeat. Ordinary rectangular accesses use repeat=1, stride=0. + // A non-positive repeat or axis extent denotes an empty access. Stride may + // be zero or negative. The detector currently emits ordinary accesses. + PrimExpr repeat; + PrimExpr stride; +}; + +inline std::vector +LogicalRegions(const std::vector &accesses) { + std::vector regions; + regions.reserve(accesses.size()); + for (const AccessRegion &access : accesses) + regions.push_back(access.logical); + return regions; +} + // MemoryAccessDetector: detect read/write regions in statements // Adapted from BlockReadWriteDetector in TVM class MemoryAccessDetector : public StmtExprVisitor { @@ -88,6 +117,21 @@ class MemoryAccessDetector : public StmtExprVisitor { operator()(expr); } + // Wrap the existing per-buffer rectangular summaries. Producing separate + // records for individual accesses is a later precision improvement. + std::vector GetReadAccesses() const { + std::vector accesses; + for (const BufferRegion ®ion : GetReadRegions()) + accesses.emplace_back(region); + return accesses; + } + std::vector GetWriteAccesses() const { + std::vector accesses; + for (const BufferRegion ®ion : GetWriteRegions()) + accesses.emplace_back(region); + return accesses; + } + // Return collected read regions std::vector GetReadRegions() const { return CollectRegions(read_buffers_, read_regions_); diff --git a/src/ascend/transform/auto_schedule/task_analysis.h b/src/ascend/transform/auto_schedule/task_analysis.h index 0169bc99d1..248700b297 100644 --- a/src/ascend/transform/auto_schedule/task_analysis.h +++ b/src/ascend/transform/auto_schedule/task_analysis.h @@ -174,8 +174,8 @@ struct TaskResourceUsage { // remain separate from hardware resource classification because AutoSchedule // also adds accesses from surrounding guards and loop expressions to a task. struct TaskAccessInfo { - std::vector read_regions; - std::vector write_regions; + std::vector read_regions; + std::vector write_regions; std::vector read_vars; std::vector write_vars; }; @@ -401,8 +401,8 @@ inline TaskAccessInfo AnalyzeTaskAccesses(const Stmt &stmt) { MemoryAccessDetector memory_detector; memory_detector.Analyze(stmt); TaskAccessInfo result; - result.read_regions = memory_detector.GetReadRegions(); - result.write_regions = memory_detector.GetWriteRegions(); + result.read_regions = memory_detector.GetReadAccesses(); + result.write_regions = memory_detector.GetWriteAccesses(); result.read_vars = memory_detector.GetReadVars(); result.write_vars = memory_detector.GetWriteVars(); return result; @@ -410,9 +410,9 @@ inline TaskAccessInfo AnalyzeTaskAccesses(const Stmt &stmt) { inline bool TaskAccessesStorage(const TaskAccessInfo &accesses, const Var &storage) { - auto regions_access = [&](const std::vector ®ions) { + auto regions_access = [&](const std::vector ®ions) { return std::any_of(regions.begin(), regions.end(), [&](const auto ®ion) { - return region->buffer->data.same_as(storage); + return region.logical->buffer->data.same_as(storage); }); }; if (regions_access(accesses.read_regions) || @@ -505,7 +505,8 @@ class CoreMaskAvailability { CoreMask result = kCoreBroadcast; for (const Var &var : accesses.read_vars) result &= GetValueCoreMask(var, reader); - for (const BufferRegion ®ion : accesses.read_regions) { + for (const AccessRegion &access : accesses.read_regions) { + const BufferRegion ®ion = access.logical; if (IsSharedBuffer(region->buffer)) { result &= kCoreVector; } else if (IsL1Buffer(region->buffer) || IsL0ABuffer(region->buffer) || @@ -521,7 +522,7 @@ class CoreMaskAvailability { CoreMask GetTaskReadCoreMask(const TaskNode *task) const { ICHECK(task != nullptr); TaskAccessInfo accesses; - accesses.read_regions = task->GetReadRegions(); + accesses.read_regions = task->GetReadAccesses(); accesses.read_vars = task->GetReadVars(); return GetAccessCoreMask(accesses, task); } @@ -625,10 +626,10 @@ inline void ApplyTaskResourceUsage(const TaskResourceUsage &usage, inline void ApplyTaskAccesses(const TaskAccessInfo &accesses, TaskNode *task) { ICHECK(task != nullptr); - for (const BufferRegion ®ion : accesses.read_regions) - task->AddReadRegion(region); - for (const BufferRegion ®ion : accesses.write_regions) - task->AddWriteRegion(region); + for (const AccessRegion &access : accesses.read_regions) + task->AddReadAccess(access); + for (const AccessRegion &access : accesses.write_regions) + task->AddWriteAccess(access); for (const Var &var : accesses.read_vars) task->AddReadVar(var); for (const Var &var : accesses.write_vars) diff --git a/src/ascend/transform/resolve_core.cc b/src/ascend/transform/resolve_core.cc index 5e40ecbeaf..3617b9175f 100644 --- a/src/ascend/transform/resolve_core.cc +++ b/src/ascend/transform/resolve_core.cc @@ -221,7 +221,8 @@ class CoreMaskResolver { bool changed = false; for (const Var &var : accesses.read_vars) changed |= AddDemand(&var_demand_, var, mask); - for (const BufferRegion ®ion : accesses.read_regions) { + for (const AccessRegion &access : accesses.read_regions) { + const BufferRegion ®ion = access.logical; if (IsRegisterRegion(region)) changed |= AddDemand(®ister_demand_, region->buffer->data, mask); } From 85931e2bb40664f9bc388f29246f220acbd391c3 Mon Sep 17 00:00:00 2001 From: Denver Jin Date: Tue, 13 Oct 2026 17:04:44 +0800 Subject: [PATCH 2/6] [Ascend] Qualify access region type in core resolution --- src/ascend/transform/resolve_core.cc | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/ascend/transform/resolve_core.cc b/src/ascend/transform/resolve_core.cc index 3617b9175f..f8cd38620e 100644 --- a/src/ascend/transform/resolve_core.cc +++ b/src/ascend/transform/resolve_core.cc @@ -221,7 +221,7 @@ class CoreMaskResolver { bool changed = false; for (const Var &var : accesses.read_vars) changed |= AddDemand(&var_demand_, var, mask); - for (const AccessRegion &access : accesses.read_regions) { + for (const ascend::AccessRegion &access : accesses.read_regions) { const BufferRegion ®ion = access.logical; if (IsRegisterRegion(region)) changed |= AddDemand(®ister_demand_, region->buffer->data, mask); From bb6b7fd42dd840b41ff0ace4c8b3e7259e8d1614 Mon Sep 17 00:00:00 2001 From: Denver Jin Date: Tue, 13 Oct 2026 17:08:48 +0800 Subject: [PATCH 3/6] [Ascend] Rename detector access records to MemoryAccess --- .../auto_schedule/dependency_analysis.cc | 10 +++---- .../auto_schedule/dependency_analysis.h | 4 +-- .../transform/auto_schedule/ir_structure.h | 30 +++++++++---------- .../transform/auto_schedule/memory_detector.h | 16 +++++----- .../transform/auto_schedule/task_analysis.h | 12 ++++---- src/ascend/transform/resolve_core.cc | 2 +- 6 files changed, 37 insertions(+), 37 deletions(-) diff --git a/src/ascend/transform/auto_schedule/dependency_analysis.cc b/src/ascend/transform/auto_schedule/dependency_analysis.cc index fbeaae38b6..055aada7c7 100644 --- a/src/ascend/transform/auto_schedule/dependency_analysis.cc +++ b/src/ascend/transform/auto_schedule/dependency_analysis.cc @@ -112,8 +112,8 @@ std::vector StorageIterationNest(ControlNode *loop, return {controls.begin() + first, controls.end()}; } -bool RegionsMayConflict(const ConstrSet &a_ctx, const AccessRegion &a_access, - const ConstrSet &b_ctx, const AccessRegion &b_access, +bool RegionsMayConflict(const ConstrSet &a_ctx, const MemoryAccess &a_access, + const ConstrSet &b_ctx, const MemoryAccess &b_access, ControlNode *loop, int offset, size_t num_storage_owners, const ConflictHintList &root_conflicts) { @@ -284,7 +284,7 @@ AnalyzeDependencies(std::vector nodes, ControlNode *loop, struct RegionAccess { TaskNode *task; - AccessRegion region; + MemoryAccess region; }; struct NodeAccesses { std::vector reads; @@ -295,9 +295,9 @@ AnalyzeDependencies(std::vector nodes, ControlNode *loop, std::vector tasks; CollectAllTaskNodes(nodes[i], tasks); for (TaskNode *task : tasks) { - for (const AccessRegion ®ion : task->GetReadAccesses()) + for (const MemoryAccess ®ion : task->GetReadAccesses()) accesses[i].reads.push_back({task, region}); - for (const AccessRegion ®ion : task->GetWriteAccesses()) + for (const MemoryAccess ®ion : task->GetWriteAccesses()) accesses[i].writes.push_back({task, region}); } } diff --git a/src/ascend/transform/auto_schedule/dependency_analysis.h b/src/ascend/transform/auto_schedule/dependency_analysis.h index 441429d707..3579a84662 100644 --- a/src/ascend/transform/auto_schedule/dependency_analysis.h +++ b/src/ascend/transform/auto_schedule/dependency_analysis.h @@ -121,8 +121,8 @@ AnalyzeDependencies(std::vector nodes, // an identically-shaped iteration in another owner can still be a distinct // epoch. Shared by scheduling-time IRStructure analysis and the scheduled-TIR // synchronization pass. -bool RegionsMayConflict(const ConstrSet &a_ctx, const AccessRegion &a_access, - const ConstrSet &b_ctx, const AccessRegion &b_access, +bool RegionsMayConflict(const ConstrSet &a_ctx, const MemoryAccess &a_access, + const ConstrSet &b_ctx, const MemoryAccess &b_access, ControlNode *loop, int offset, size_t num_storage_owners = 0, const ConflictHintList &root_conflicts = {}); diff --git a/src/ascend/transform/auto_schedule/ir_structure.h b/src/ascend/transform/auto_schedule/ir_structure.h index 9c7e516c08..3f3bf91d21 100644 --- a/src/ascend/transform/auto_schedule/ir_structure.h +++ b/src/ascend/transform/auto_schedule/ir_structure.h @@ -111,7 +111,7 @@ inline bool RegionsEqual(const Region &a, const Region &b) { return true; } -inline bool AccessRegionsEqual(const AccessRegion &a, const AccessRegion &b) { +inline bool MemoryAccessesEqual(const MemoryAccess &a, const MemoryAccess &b) { arith::Analyzer analyzer; return a.logical->buffer.same_as(b.logical->buffer) && RegionsEqual(a.logical->region, b.logical->region) && @@ -540,10 +540,10 @@ class TaskNode : public IRStructure { // A copy that stores to GM in this leaf reads the atomic-mode register. bool ReadsAtomicRegister() const override { return reads_atomic_; } - const std::vector &GetReadAccesses() const { + const std::vector &GetReadAccesses() const { return read_regions_; } - const std::vector &GetWriteAccesses() const { + const std::vector &GetWriteAccesses() const { return write_regions_; } @@ -568,10 +568,10 @@ class TaskNode : public IRStructure { void AddSpecialWrite(SpecialRegister reg) { special_write_mask_ |= static_cast(reg); } - void SetReadAccesses(const std::vector ®ions) { + void SetReadAccesses(const std::vector ®ions) { read_regions_ = regions; } - void SetWriteAccesses(const std::vector ®ions) { + void SetWriteAccesses(const std::vector ®ions) { write_regions_ = regions; } void SetReadVars(const std::vector &vars) { read_vars_ = vars; } @@ -586,8 +586,8 @@ class TaskNode : public IRStructure { }; std::unordered_map buffer_remap; - auto substitute_regions = [&](std::vector ®ions) { - for (AccessRegion &access : regions) { + auto substitute_regions = [&](std::vector ®ions) { + for (MemoryAccess &access : regions) { Buffer buffer = access.logical->buffer; auto buffer_it = buffer_remap.find(buffer); if (buffer_it != buffer_remap.end()) { @@ -625,17 +625,17 @@ class TaskNode : public IRStructure { std::shared_ptr Clone() const override; // Keep the existing deduplication, comparing the complete access record. - void AddReadAccess(const AccessRegion &access) { - for (const AccessRegion &existing : read_regions_) { - if (AccessRegionsEqual(existing, access)) + void AddReadAccess(const MemoryAccess &access) { + for (const MemoryAccess &existing : read_regions_) { + if (MemoryAccessesEqual(existing, access)) return; } read_regions_.push_back(access); } - void AddWriteAccess(const AccessRegion &access) { - for (const AccessRegion &existing : write_regions_) { - if (AccessRegionsEqual(existing, access)) + void AddWriteAccess(const MemoryAccess &access) { + for (const MemoryAccess &existing : write_regions_) { + if (MemoryAccessesEqual(existing, access)) return; } write_regions_.push_back(access); @@ -670,8 +670,8 @@ class TaskNode : public IRStructure { // inferred or explicitly marked candidate. bool is_per_core_task_{false}; - std::vector read_regions_; - std::vector write_regions_; + std::vector read_regions_; + std::vector write_regions_; std::vector read_vars_; std::vector write_vars_; diff --git a/src/ascend/transform/auto_schedule/memory_detector.h b/src/ascend/transform/auto_schedule/memory_detector.h index dd599bacf3..8c94deaf39 100644 --- a/src/ascend/transform/auto_schedule/memory_detector.h +++ b/src/ascend/transform/auto_schedule/memory_detector.h @@ -37,8 +37,8 @@ using namespace tirx; // One detected access. A list of these records denotes a union, including // when several records refer to the same logical buffer. -struct AccessRegion { - explicit AccessRegion(BufferRegion logical) +struct MemoryAccess { + explicit MemoryAccess(BufferRegion logical) : logical(std::move(logical)), bounds(this->logical->region), repeat(Integer(1)), stride(Integer(0)) {} @@ -56,10 +56,10 @@ struct AccessRegion { }; inline std::vector -LogicalRegions(const std::vector &accesses) { +LogicalRegions(const std::vector &accesses) { std::vector regions; regions.reserve(accesses.size()); - for (const AccessRegion &access : accesses) + for (const MemoryAccess &access : accesses) regions.push_back(access.logical); return regions; } @@ -119,14 +119,14 @@ class MemoryAccessDetector : public StmtExprVisitor { // Wrap the existing per-buffer rectangular summaries. Producing separate // records for individual accesses is a later precision improvement. - std::vector GetReadAccesses() const { - std::vector accesses; + std::vector GetReadAccesses() const { + std::vector accesses; for (const BufferRegion ®ion : GetReadRegions()) accesses.emplace_back(region); return accesses; } - std::vector GetWriteAccesses() const { - std::vector accesses; + std::vector GetWriteAccesses() const { + std::vector accesses; for (const BufferRegion ®ion : GetWriteRegions()) accesses.emplace_back(region); return accesses; diff --git a/src/ascend/transform/auto_schedule/task_analysis.h b/src/ascend/transform/auto_schedule/task_analysis.h index 248700b297..e807fabb82 100644 --- a/src/ascend/transform/auto_schedule/task_analysis.h +++ b/src/ascend/transform/auto_schedule/task_analysis.h @@ -174,8 +174,8 @@ struct TaskResourceUsage { // remain separate from hardware resource classification because AutoSchedule // also adds accesses from surrounding guards and loop expressions to a task. struct TaskAccessInfo { - std::vector read_regions; - std::vector write_regions; + std::vector read_regions; + std::vector write_regions; std::vector read_vars; std::vector write_vars; }; @@ -410,7 +410,7 @@ inline TaskAccessInfo AnalyzeTaskAccesses(const Stmt &stmt) { inline bool TaskAccessesStorage(const TaskAccessInfo &accesses, const Var &storage) { - auto regions_access = [&](const std::vector ®ions) { + auto regions_access = [&](const std::vector ®ions) { return std::any_of(regions.begin(), regions.end(), [&](const auto ®ion) { return region.logical->buffer->data.same_as(storage); }); @@ -505,7 +505,7 @@ class CoreMaskAvailability { CoreMask result = kCoreBroadcast; for (const Var &var : accesses.read_vars) result &= GetValueCoreMask(var, reader); - for (const AccessRegion &access : accesses.read_regions) { + for (const MemoryAccess &access : accesses.read_regions) { const BufferRegion ®ion = access.logical; if (IsSharedBuffer(region->buffer)) { result &= kCoreVector; @@ -626,9 +626,9 @@ inline void ApplyTaskResourceUsage(const TaskResourceUsage &usage, inline void ApplyTaskAccesses(const TaskAccessInfo &accesses, TaskNode *task) { ICHECK(task != nullptr); - for (const AccessRegion &access : accesses.read_regions) + for (const MemoryAccess &access : accesses.read_regions) task->AddReadAccess(access); - for (const AccessRegion &access : accesses.write_regions) + for (const MemoryAccess &access : accesses.write_regions) task->AddWriteAccess(access); for (const Var &var : accesses.read_vars) task->AddReadVar(var); diff --git a/src/ascend/transform/resolve_core.cc b/src/ascend/transform/resolve_core.cc index f8cd38620e..a44d96c3d0 100644 --- a/src/ascend/transform/resolve_core.cc +++ b/src/ascend/transform/resolve_core.cc @@ -221,7 +221,7 @@ class CoreMaskResolver { bool changed = false; for (const Var &var : accesses.read_vars) changed |= AddDemand(&var_demand_, var, mask); - for (const ascend::AccessRegion &access : accesses.read_regions) { + for (const ascend::MemoryAccess &access : accesses.read_regions) { const BufferRegion ®ion = access.logical; if (IsRegisterRegion(region)) changed |= AddDemand(®ister_demand_, region->buffer->data, mask); From 95883ef62583f7d4a6014af96794e1773b5246f3 Mon Sep 17 00:00:00 2001 From: Denver Jin Date: Tue, 13 Oct 2026 17:23:38 +0800 Subject: [PATCH 4/6] [Ascend] Support partial L0 copies with explicit stride --- .../developer_guide/ascend_partial_l0_copy.md | 113 +++++++ src/ascend/layout/ascend_layouts.cc | 24 ++ src/ascend/layout/ascend_layouts.h | 18 ++ src/ascend/op/copy.cc | 263 +++++++++------- src/ascend/op/copy.h | 3 +- .../auto_schedule/dependency_analysis.cc | 18 ++ .../transform/auto_schedule/ir_structure.h | 13 + .../transform/auto_schedule/memory_detector.h | 284 ++++++++++++++++-- .../transform/auto_schedule/scheduled_tir.cc | 14 +- .../transform/auto_schedule/task_analysis.h | 6 +- src/ascend/transform/lower_tile_op.cc | 11 + testing/ascend/_ir.py | 13 +- .../auto_schedule/test_access_footprints.py | 46 ++- .../auto_schedule/test_cross_core_flags.py | 7 +- .../auto_schedule/test_flag_allocation.py | 10 +- .../test_sync_access_footprints.py | 43 ++- .../auto_schedule/test_sync_dependencies.py | 152 +++++++++- .../layout/test_ascend_partial_l0_copy.py | 185 ++++++++++++ .../transform/test_ascend_lower_copy.py | 155 +++++++++- tilelang/ascend/language/copy_op.py | 35 +++ 20 files changed, 1256 insertions(+), 157 deletions(-) create mode 100644 docs/developer_guide/ascend_partial_l0_copy.md create mode 100644 testing/ascend/layout/test_ascend_partial_l0_copy.py diff --git a/docs/developer_guide/ascend_partial_l0_copy.md b/docs/developer_guide/ascend_partial_l0_copy.md new file mode 100644 index 0000000000..938a5f2b97 --- /dev/null +++ b/docs/developer_guide/ascend_partial_l0_copy.md @@ -0,0 +1,113 @@ +# Ascend：显式 copy stride 与部分 L0 装载 + +## 前端合约 + +`tilelang.ascend.language.copy` 增加可选的 `stride` 参数，表示 L0 矩阵的 +pitch，单位为 M/N 方向的 **16 元素分形数**，不是字节或分形间隙。 + +| 路径 | `stride` 控制 | 与普通 GEMM 配合时的值 | +|---|---|---| +| L1→L0A | 目标 pitch | `ceildiv(M, 16)` | +| L1→L0B | 目标 pitch，包括 MajorMN 转置装载 | `ceildiv(N, 16)` | +| L0C→UB/GM | 源 pitch | `ceildiv(M, 16)` | + +不传 `stride` 保留既有 compact region 行为：一次 copy 搬运一个完整 GEMM +操作数时，从这次 region 推导 pitch。多次 copy 拼接一个操作数时，每次显式 +传入完整 tile 的 pitch。非零矩阵目标起点要求显式 `stride`。 +其他搬运路径不接受此参数。 + +编译器不查找 copy 对应哪个 GEMM，也不要求同一 buffer 的所有 GEMM 使用相同 +pitch。调用者保证 copy 的 stride 与本次 GEMM 的有效 M/N 一致。 +`stride` 可以是整数表达式，随循环或分支改变；在 copy 执行位置求值。 +改变 pitch 不会重排旧数据，调用者必须重新准备相应输入和累加器。 + +MajorK/MajorMN 仍由既有 layout inference 决定，分配容量及 leading plane / +multibuffer slot 间距保持 allocation layout。显式 stride 仅改变 plane 内的寻址。 +MX scale-factor handle 保持专用搬运合约，不在此参数的支持路径中。 + +## 普通装载与尾块 + +```python +# a 的 allocation 为 [64,64]，本次 GEMM 的有效 M=24。 +T.copy(a_l1[:16, :], a[:16, :], stride=2) +T.copy(a_l1[16:32, :], a[16:32, :], stride=2) +T.gemm(a[:24, :], b, acc[:24, :], transpose_B=True, clear_accum=True) +# 完整输出 region 已经提供本次 compact pitch，无需标注。 +T.copy(acc[:24, :], out[:24, :]) +# 若只输出前 16 行，仍需按完整结果的 pitch 读取。 +T.copy(acc[:16, :], first_rows, stride=2) +``` + +M=64 时，两次各搬 32 行的 copy 都传 `stride=4`。动态尾块可以传 +`stride=T.ceildiv(effective_m, 16)`，region 选择包含 padding 的完整分形。 + +对 BF16 MajorK，逻辑坐标映射为 `[k//16, m//16, m%16, k%16]`。 +分形线性位置为 `outer0 * stride + outer1`，每个分形占 512 B。 +例如 allocation 为 [2,64,64]，选择 plane=1、stride=2、目标起点 [0,32]： + +```text +plane base = 64 * 64 个元素 +plane 内偏移 = (2 * 2 + 0) * 256 个元素 +``` + +## V 的转置装载 + +```python +# v_l1 / v 的逻辑形状是 [K,N]，此处 K=64。 +T.copy(v_l1[:32, :N_pad], v[:32, :N_pad], stride=T.ceildiv(N, 16)) +T.copy(v_l1[32:64, :N_pad], v[32:64, :N_pad], stride=T.ceildiv(N, 16)) +T.gemm(a, v[:, :N], acc[:, :N], clear_accum=True) +``` + +GEMM 为 V 推导 MajorMN，目标偏移按 MajorMN 和显式 stride 计算。 +硬件转置位仍是 `user_transpose XOR major_mismatch`。 +令 C0=256/dtype.bits，源 region 的分形步数为 (k_step,m_step)： + +```text +普通实际写入 = (k_step, m_step) +转置实际写入 = (m_step*16/C0, k_step*C0/16) +``` + +copy 检查完整转置组、实际写入范围、目标起点和容量。 +非零目标起点要求源和目标的起点、extent 均分形对齐；零起点保留既有 padded load。 +实际写入必须位于声明的目标 region 内，且不能越过 stride 或 allocation。 +空 copy guard 保持不变。 + +## Lowering 与同步 + +前端将 `stride` 存入 Ascend copy annotation；原有 IR 变量替换同时处理该表达式， +包括流水化迭代的变量替换。L1→L0 lowering 将它用于指令的 dst_stride， +并通过内部 access-pointer annotation 将它传给地址 remap。 +LowerTileOp 保留完整 Major 映射和 allocation 间距，只修正 matrix outer0 的 stride。 +L0C→UB/GM 将分形数乘以 16，得到 Fixpipe 的源行距。 +GEMM lowering、布局推导不新增 pitch 传播或匹配过程。 + +AutoSchedule 开关不影响 stride 语义。开启时仍按 RAW/WAR/WAW 关系加同步： +每条指令的 L0 访问先表示为 plane 内的字节 Range，附带 `(repeat, stride)`， +分别表示重复次数和重复段起点间的字节距离。不同 pitch 也直接比较实际地址范围； +仅在无法确定寻址方式时退化为整个 plane。保留 leading plane / slot 分离与条件、 +循环上下文。`MemoryAccessDetector` 在收集 Task 时逐次生成 `MemoryAccess`， +经 `TaskAccessInfo` 存入 `TaskNode`，依赖分析直接消费这些记录。 +同一 buffer 的多个 copy 分别保留各自的 pitch 和访问范围，不先合并为一个矩形。 +Task 内循环的变量范围随访问记录保存，并在比较两侧访问时分别实例化，保留 +地址与动态 pitch 的关联。Task 克隆和变量替换同时处理这些记录；无需额外扫描 +Task 或维护一份访问缓存。 +显式 conflict hints 仍匹配原始逻辑 region，不由编译器自动生成。 +这只是访问冲突判断，不给 copy 选择 GEMM 或推导 stride。 +关闭 AutoSchedule 时，由调用者设置各流水线之间的同步。 + +多个 writer 允许顺序覆盖，不新增 writer 并集覆盖证明。 +物理 lifetime 分析将 compact L0 storage 标记为 opaque，跳过跨 allocation 的重叠 +复用证明，并保留所选 plane 的 may-access；其他普通 operand 保留 must-write 信息。 + +## 支持边界与回归 + +沿用 GEMM 和 L0C drain 的矩阵零起点限制。支持 L1→L0A/L0B 的分形对齐非零 +目标起点、普通及显式/隐式转置、动态 M/N 尾块、K padding 和有效 K。 +不插入 repack，不自动补齐非对齐 offset copy，不校验 copy 与 GEMM 的布局一致性。 +FP4 packed-address 不在首批设备验证范围内。 + +主机回归覆盖默认 compact pitch、显式静态/动态 stride、普通/V 目标地址、 +leading plane、部分 L0C drain、实际边界、不同 pitch 的同步及 allocation 复用。 +NPU 回归覆盖两/三版本 ring 的部分装载、动态尾块、顺序覆盖，并覆盖关闭 +AutoSchedule、手动同步、分支内不同 GEMM pitch 共用同一 buffer 的情形。 diff --git a/src/ascend/layout/ascend_layouts.cc b/src/ascend/layout/ascend_layouts.cc index 2b52acdd07..110d657849 100644 --- a/src/ascend/layout/ascend_layouts.cc +++ b/src/ascend/layout/ascend_layouts.cc @@ -5,12 +5,14 @@ */ #include "support/check.h" +#include #include #include #include #include #include "ascend/layout/ascend_layouts.h" +#include "ascend/op/utils.h" #include "layout/layout.h" namespace tvm { @@ -380,6 +382,28 @@ bool TryExtractAscendFractalRegion(const Layout &layout, return true; } +BufferRegion ExpandAscendFractalMayAccess(const BufferRegion ®ion) { + const Buffer &buffer = region->buffer; + int bits = buffer->dtype.bits(); + if (!(IsL1Buffer(buffer) || IsL0ABuffer(buffer) || IsL0BBuffer(buffer) || + IsL0CBuffer(buffer)) || + region->region.size() < 2 || bits <= 0 || 256 % bits != 0) + return region; + int alignment = IsL0CBuffer(buffer) ? 16 : std::max(16, 256 / bits); + ffi::Array ranges = region->region; + for (const Range &range : ranges) + if (is_zero(range->extent)) + return region; + for (size_t i = ranges.size() - 2; i < ranges.size(); ++i) { + PrimExpr begin = floordiv(ranges[i]->min, alignment) * alignment; + PrimExpr end = floordiv(ranges[i]->min + ranges[i]->extent + alignment - 1, + alignment) * + alignment; + ranges.Set(i, Range::FromMinExtent(begin, end - begin)); + } + return BufferRegion(buffer, ranges); +} + // Ascend ND (row-major) identity layout for GM/UB operands. Forward is the // identity [i0, .., i_{n-1}] -> [i0, .., i_{n-1}], so OutputShape == // InputShape. Used only for copy-parameter derivation in lowering; diff --git a/src/ascend/layout/ascend_layouts.h b/src/ascend/layout/ascend_layouts.h index 43b06af9c3..420624678b 100644 --- a/src/ascend/layout/ascend_layouts.h +++ b/src/ascend/layout/ascend_layouts.h @@ -127,6 +127,24 @@ bool TryExtractAscendFractalRegion(const Layout &layout, const ffi::Array &logical_region, AscendFractalRegionInfo *out); +/// Conservatively expand an L1/L0 data region in logical coordinates to cover +/// whole fractals for either Cube major. Both trailing axes are rounded to +/// max(16, C0), or 16 for L0C; leading axes and statically empty regions remain +/// unchanged. Unsupported scopes, ranks, or element widths pass through. +/// This is a may-access bound only: it must not establish must-write coverage +/// or determine the actual instruction transfer geometry. +tirx::BufferRegion +ExpandAscendFractalMayAccess(const tirx::BufferRegion ®ion); + +namespace attr { + +// Copy lowering transfers the destination pitch to its tvm_access_ptr. +// LowerTileOp uses it to remap matrix offsets while retaining the allocation's +// leading-plane / buffer-version stride. +constexpr const char *kL0AccessPitch = "tl.ascend_l0_access_pitch"; + +} // namespace attr + } // namespace tl } // namespace tvm diff --git a/src/ascend/op/copy.cc b/src/ascend/op/copy.cc index fc87f4339c..c3b04487b4 100644 --- a/src/ascend/op/copy.cc +++ b/src/ascend/op/copy.cc @@ -418,125 +418,141 @@ Stmt LowerDMACopy(const AscendCopyNode &op, const LowerArgs &T, "Ascend L1->L0 source"); NormalizeTrailingMTE2DLayout(op.dst, dst_range, analyzer, "Ascend L1->L0 destination"); - CheckCompactL0Region(op.dst, dst_range, analyzer, - "Ascend L1->L0 compact destination"); - // M/K positions go through intrinsic args; leading dims (e.g. version) - // through the pointer offset. + // Source matrix coordinates are intrinsic arguments; only its leading + // plane goes into the pointer. The destination pointer keeps the logical + // slice origin. LowerTileOp maps it with the copy pitch while preserving + // the allocation's leading-plane (including buffer-version) stride. PrimExpr src_base_ptr = MakeAscendLeadingDimAccessPtr(BufferRegion(op.src, src_range), 1); - - AscendFractalLayoutInfo src_info; - AscendFractalLayoutInfo dst_info; + AscendFractalLayoutInfo src_info, dst_info; Optional src_layout = FindLayoutForBuffer(T.layout_map, op.src); Optional dst_layout = FindLayoutForBuffer(T.layout_map, op.dst); - bool have_src_layout = - src_layout.defined() && - TryExtractAscendFractalLayout(src_layout.value(), op.src, &src_info); - bool have_dst_layout = - dst_layout.defined() && - TryExtractAscendFractalLayout(dst_layout.value(), op.dst, &dst_info); - - PrimExpr m_start, k_start, m_step, k_step, src_stride, dst_stride; - - ICHECK(have_src_layout && have_dst_layout) - << "Ascend L1->L0 compact-region lowering requires canonical fractal " - "layouts on both buffers, got source " - << op.src->name << " layout=" << src_layout.defined() - << " and destination " << op.dst->name - << " layout=" << dst_layout.defined() << "."; - AscendFractalRegionInfo src_region; - AscendFractalRegionInfo dst_region; + ICHECK( + src_layout.defined() && dst_layout.defined() && + TryExtractAscendFractalLayout(src_layout.value(), op.src, &src_info) && + TryExtractAscendFractalLayout(dst_layout.value(), op.dst, &dst_info) && + IsAscendMajorKind(src_info.kind) && IsAscendMajorKind(dst_info.kind)) + << "Ascend L1->L0 lowering requires canonical MajorK/MajorMN layouts"; + ICHECK(op.src->dtype == op.dst->dtype) + << "Ascend L1->L0 copy requires matching source/destination dtypes"; + + AscendFractalRegionInfo src_region, dst_region; ICHECK(TryExtractAscendFractalRegion(src_layout.value(), src_range, - &src_region)) - << "Failed to map L1->L0 source region through layout for buffer " - << op.src->name; + &src_region)); ICHECK(TryExtractAscendFractalRegion(dst_layout.value(), dst_range, - &dst_region)) - << "Failed to map L1->L0 destination region through layout for buffer " - << op.dst->name; - m_start = src_region.outer1->min; - k_start = src_region.outer0->min; - m_step = src_region.outer1->extent; - k_step = src_region.outer0->extent; - src_stride = src_info.outer1; - // L0 regions are compact MAD tiles: the destination pitch is the mapped - // region pitch, while the L1 source keeps its allocation-wide NZ pitch. - dst_stride = dst_region.outer1->extent; - - // The user may request an explicit transpose via T.copy(.., - // transpose=True). Additionally, when the L1 source and the L0 - // destination carry different fractal majors (e.g. a K-major L1 tile - // feeding an MN-major L0 fractal, or vice versa), an implicit transpose is - // required to reconcile the two layouts. The effective transpose is the - // XOR of the user request and the major mismatch: two transposes cancel. - // - // L1 tiles are canonically K-major (reduce-K on the C0 axis); when the L1 - // source carries no explicit layout we assume K-major so the mismatch is - // driven purely by the inferred L0 major. - bool user_transpose = op.transpose != 0; - LogicalAxis src_c0_axis = - have_src_layout ? src_info.c0_axis : LogicalAxis::kCol; - LogicalAxis dst_c0_axis = - have_dst_layout ? dst_info.c0_axis : LogicalAxis::kCol; - bool major_mismatch = src_c0_axis != dst_c0_axis; - bool needs_transpose = user_transpose ^ major_mismatch; + &dst_region)); + PrimExpr m_start = src_region.outer1->min; + PrimExpr k_start = src_region.outer0->min; + PrimExpr m_step = src_region.outer1->extent; + PrimExpr k_step = src_region.outer0->extent; + PrimExpr src_stride = src_info.outer1; + PrimExpr dst_stride = op.stride.value_or(dst_region.outer1->extent); + Call ptr = Downcast(dst_ptr); + ptr.CopyOnWrite()->annotations.Set(attr::kL0AccessPitch, dst_stride); + dst_ptr = ptr; + With nonempty_copy(analyzer, has_data); + ICHECK(analyzer->CanProve(dst_stride > 0) && + analyzer->CanProve(dst_stride <= dst_info.outer1, + arith::ProofStrength::kSymbolicBound)) + << "Ascend L1->L0 stride must be positive and fit destination capacity"; + auto require_multiple = [&](const PrimExpr &value, int multiple, + const char *parameter) { + ICHECK(analyzer->CanProveEqual( + FloorMod(value, make_const(value.dtype(), multiple)), 0)) + << "Ascend L1->L0 copy of " << op.src->name << " (" << op.src->dtype + << ") requires " << parameter << " to be divisible by " << multiple + << ", got " << value + << ". Copy complete fractals/transpose groups; restrict GEMM to " + "the effective M/N/K."; + }; + bool nonzero_destination = false; + for (size_t i = dst_range.size() - 2; i < dst_range.size(); ++i) + nonzero_destination |= !analyzer->CanProveEqual(dst_range[i]->min, 0); + ICHECK(!nonzero_destination || op.stride.defined()) + << "Ascend partial L1->L0 copy with a nonzero destination origin " + "requires an explicit stride"; + auto check_region = [&](const Buffer &buffer, const Array &ranges, + const AscendFractalLayoutInfo &info, + const AscendFractalRegionInfo ®ion, + bool source) { + size_t row16_axis = + ranges.size() - 2 + (info.row16_axis == LogicalAxis::kCol); + size_t c0_axis = ranges.size() - 2 + (info.c0_axis == LogicalAxis::kCol); + require_multiple(ranges[row16_axis]->min, 16, + source ? "source row16 origin" + : "destination row16 origin"); + require_multiple(ranges[c0_axis]->min, AscendC0(buffer->dtype.bits()), + source ? "source C0 origin" : "destination C0 origin"); + // Keep existing zero-origin padded loads. Newly supported offset writes + // explicitly select complete fractals, including their padding lanes. + if (nonzero_destination) { + require_multiple(ranges[row16_axis]->extent, 16, + source ? "source row16 extent" + : "destination row16 extent"); + require_multiple(ranges[c0_axis]->extent, + AscendC0(buffer->dtype.bits()), + source ? "source C0 extent" : "destination C0 extent"); + } + for (size_t i = 0; i + 2 < ranges.size(); ++i) { + ICHECK(analyzer->CanProve(ranges[i]->min >= 0) && + analyzer->CanProve(ranges[i]->min + ranges[i]->extent <= + buffer->shape[i])) + << "Ascend L1->L0 leading plane is outside buffer " << buffer->name; + } + ICHECK(analyzer->CanProve(region.outer0->min >= 0) && + analyzer->CanProve(region.outer1->min >= 0) && + analyzer->CanProve(region.outer0->min + region.outer0->extent <= + info.outer0, + arith::ProofStrength::kSymbolicBound) && + analyzer->CanProve(region.outer1->min + region.outer1->extent <= + info.outer1, + arith::ProofStrength::kSymbolicBound)) + << "Ascend L1->L0 " << (source ? "source" : "destination") + << " region exceeds or cannot be proved within the physical " + "allocation of " + << buffer->name; + }; + check_region(op.src, src_range, src_info, src_region, true); + check_region(op.dst, dst_range, dst_info, dst_region, false); + + // Explicit logical transpose and the layout conversion cancel each other. + bool needs_transpose = + (op.transpose != 0) ^ (src_info.c0_axis != dst_info.c0_axis); + PrimExpr written_outer0 = k_step; + PrimExpr written_outer1 = m_step; if (needs_transpose) { - // Validate the emitted load after storage normalization and OOB - // clamping. Logical extents need not cover whole fractals (e.g. bf16 - // MN=24), but the hardware's transpose groups must be complete. int c0 = AscendC0(op.src->dtype.bits()); - auto require_multiple = [&](const PrimExpr &value, int multiple, - const char *parameter) { - if (multiple == 1) - return; - // Empty copies never issue the instruction. Scope this assumption - // to the proof so the runtime has_data guard is preserved below. - With nonempty_copy(analyzer, has_data); - ICHECK(analyzer->CanProveEqual( - FloorMod(value, make_const(value.dtype(), multiple)), 0)) - << "Ascend transposed L1->L0 copy of " << op.src->name << " (" - << op.src->dtype << ") requires " << parameter - << " to be divisible by " << multiple << ", got " << value - << ". Copy a padded source region covering complete transpose " - "groups; restrict the GEMM region to the effective M/N/K."; - }; require_multiple(m_step, c0 > 16 ? c0 / 16 : 1, "m_step"); require_multiple(k_step, c0 < 16 ? 16 / c0 : 1, "k_step"); - size_t ndim = src_range.size(); - require_multiple( - src_range[ndim - 2 + (src_info.row16_axis == LogicalAxis::kCol)]->min, - 16, "source row16 origin"); - require_multiple( - src_range[ndim - 2 + (src_info.c0_axis == LogicalAxis::kCol)]->min, - c0, "source C0 origin"); - - // Transposition exchanges the source's row16 and C0 dimensions. - // Compare physical regions, not logical extents or allocation bytes: - // the declared destination also describes the copy's write effects. - PrimExpr written_outer0 = FloorDiv(m_step * 16, c0); - PrimExpr written_outer1 = FloorDiv(k_step * c0, 16); - ICHECK(!analyzer->CanProve(written_outer0 > dst_region.outer0->extent) && - !analyzer->CanProve(written_outer1 > dst_region.outer1->extent)) - << "Ascend transposed L1->L0 copy writes " << written_outer0 << "x" - << written_outer1 << " fractals, but destination region " - << op.dst->name << " only covers " << dst_region.outer0->extent << "x" - << dst_region.outer1->extent - << ". Pad the L0 allocation and copy destination region to cover " - "the physical write; use the effective K only in T.gemm."; - ICHECK(!analyzer->CanProve(written_outer0 > dst_info.outer0) && - !analyzer->CanProve(written_outer1 > dst_info.outer1)) - << "Ascend transposed L1->L0 copy exceeds the physical L0 allocation " - << op.dst->name; + written_outer0 = FloorDiv(m_step * 16, c0); + written_outer1 = FloorDiv(k_step * c0, 16); + } + ICHECK(analyzer->CanProve(written_outer0 <= dst_region.outer0->extent) && + analyzer->CanProve(written_outer1 <= dst_region.outer1->extent)) + << "Ascend L1->L0 copy writes " << written_outer0 << "x" + << written_outer1 << " fractals, but destination region " + << op.dst->name << " only covers " << dst_region.outer0->extent << "x" + << dst_region.outer1->extent + << ". Pad the allocation and copy region; use effective K in T.gemm."; + ICHECK(analyzer->CanProve(dst_region.outer0->min + written_outer0 <= + dst_info.outer0) && + analyzer->CanProve(dst_region.outer1->min + written_outer1 <= + dst_stride)) + << "Ascend L1->L0 copy exceeds the physical L0 allocation " + << "or stride of " << op.dst->name; + if (nonzero_destination) { + ICHECK( + analyzer->CanProveEqual(written_outer0, dst_region.outer0->extent) && + analyzer->CanProveEqual(written_outer1, dst_region.outer1->extent)) + << "Ascend partial L1->L0 copy requires matching physical source and " + "destination extents"; } - auto intrinsic = dma_path == DMAPath::kL1ToL0A ? ascend_load_cbuf_to_ca() : ascend_load_cbuf_to_cb(); - Array ld_args = { - dst_ptr, src_base_ptr, m_start, - k_start, m_step, k_step, - src_stride, dst_stride, I(needs_transpose ? 1 : 0)}; - call = Call(DataType::Void(), intrinsic, ld_args); + call = Call(DataType::Void(), intrinsic, + {dst_ptr, src_base_ptr, m_start, k_start, m_step, k_step, + src_stride, dst_stride, I(needs_transpose ? 1 : 0)}); } else if (dma_path == DMAPath::kL0CToUB) { StridedLayout src_layout = NormalizeTrailingMTE2DLayout( op.src, src_range, analyzer, "Ascend L0C->UB source"); @@ -569,6 +585,17 @@ Stmt LowerDMACopy(const AscendCopyNode &op, const LowerArgs &T, << "Failed to map L0C->UB source region through layout for buffer " << op.src->name; loop_src_stride = compact_src_region.outer1->extent * info.row_frac; + if (op.stride.defined()) { + PrimExpr compact_pitch = op.stride.value(); + With nonempty_copy(analyzer, has_data); + ICHECK(analyzer->CanProve(compact_src_region.outer1->extent <= + compact_pitch) && + analyzer->CanProve(compact_pitch <= info.outer1, + arith::ProofStrength::kSymbolicBound)) + << "Ascend L0C copy stride must cover the source region and fit " + "capacity"; + loop_src_stride = compact_pitch * info.row_frac; + } int dual_dst_ctl = op.dual_dst_ctl; PrimExpr copy_inner = src_inner.size; @@ -687,6 +714,17 @@ Stmt LowerDMACopy(const AscendCopyNode &op, const LowerArgs &T, << "Failed to map L0C->GM source region through layout for buffer " << op.src->name; loop_src_stride = compact_src_region.outer1->extent * info.row_frac; + if (op.stride.defined()) { + PrimExpr compact_pitch = op.stride.value(); + With nonempty_copy(analyzer, has_data); + ICHECK(analyzer->CanProve(compact_src_region.outer1->extent <= + compact_pitch) && + analyzer->CanProve(compact_pitch <= info.outer1, + arith::ProofStrength::kSymbolicBound)) + << "Ascend L0C copy stride must cover the source region and fit " + "capacity"; + loop_src_stride = compact_pitch * info.row_frac; + } // L0C->GM needs the physical distance between adjacent destination rows, // not the logical last dimension. They are equal for compact tensors but // differ for padded/runtime StridedTensor outputs. @@ -1007,6 +1045,21 @@ void DecodeAscendCopyAnnotations(AscendCopyNode *node) { if (auto val = node->annotations.Get("sub_blockid")) { node->sub_blockid = Downcast(val.value()); } + node->stride = std::nullopt; + if (auto val = node->annotations.Get("stride")) { + node->stride = Downcast(val.value()); + ICHECK(node->stride.value().dtype().is_scalar() && + (node->stride.value().dtype().is_int() || + node->stride.value().dtype().is_uint())) + << "Ascend copy stride must be a scalar integer"; + auto path = ascend::GetDMAPath(node->src, node->dst); + ICHECK(path == ascend::DMAPath::kL1ToL0A || + path == ascend::DMAPath::kL1ToL0B || + path == ascend::DMAPath::kL0CToUB || + path == ascend::DMAPath::kL0CToGM) + << "Ascend copy stride is supported only for L1->L0A/L0B and " + "L0C->UB/GM"; + } node->pad_value = std::nullopt; if (auto val = node->annotations.Get("pad_value")) { node->pad_value = Downcast(val.value()); diff --git a/src/ascend/op/copy.h b/src/ascend/op/copy.h index 2a3cb69f37..dec0f01258 100644 --- a/src/ascend/op/copy.h +++ b/src/ascend/op/copy.h @@ -39,7 +39,8 @@ class AscendCopyNode : public CopyNode { PrimExpr unit_flag_ctl; // "unit_flag_ctrl", defaults to 0 PrimExpr sub_blockid; // "sub_blockid", defaults to 0 Optional l2_cache_ctrl; // "l2_cache_ctrl"; defaults differ per path - Optional pad_value; // "pad_value": GM->UB pad fill value + Optional stride; // "stride": L0 matrix pitch in fractal blocks + Optional pad_value; // "pad_value": GM->UB pad fill value /*! * \brief l2_cache_ctrl with a path-specific default (GM loads default to 0, diff --git a/src/ascend/transform/auto_schedule/dependency_analysis.cc b/src/ascend/transform/auto_schedule/dependency_analysis.cc index 055aada7c7..1ec3833ffa 100644 --- a/src/ascend/transform/auto_schedule/dependency_analysis.cc +++ b/src/ascend/transform/auto_schedule/dependency_analysis.cc @@ -214,6 +214,24 @@ bool RegionsMayConflict(const ConstrSet &a_ctx, const MemoryAccess &a_access, // mutable reads at two access points need not have the same value. FreshenMutableReads freshen_a(FreshenMutableReads::Mode::kSnapshot); FreshenMutableReads freshen_b(FreshenMutableReads::Mode::kSnapshot); + // Internal task loops are independent access choices on the two sides. + // Rename them together with their bounds and pitch before adding domains. + auto bind_domains = [&](const MemoryAccess &access, Map *sub, + FreshenMutableReads &freshen, const char *suffix) { + for (const auto &[var, range] : access.domains) + sub->Set(var, var.copy_with_suffix(suffix)); + for (const auto &[var, range] : access.domains) { + PrimExpr index = freshen(Substitute(var, *sub)); + PrimExpr min = freshen(Substitute(range->min, *sub)); + PrimExpr extent = freshen(Substitute(range->extent, *sub)); + ana.EnterConstraint(index >= min && + cast(DataType::Int(64), index) < + cast(DataType::Int(64), min) + + cast(DataType::Int(64), extent)); + } + }; + bind_domains(a_access, &a_sub, freshen_a, "_access_p"); + bind_domains(b_access, &b_sub, freshen_b, "_access_c"); PrimExpr a_repeat = freshen_a(Substitute(a_access.repeat, a_sub)); PrimExpr b_repeat = freshen_b(Substitute(b_access.repeat, b_sub)); diff --git a/src/ascend/transform/auto_schedule/ir_structure.h b/src/ascend/transform/auto_schedule/ir_structure.h index 3f3bf91d21..6c59d571a5 100644 --- a/src/ascend/transform/auto_schedule/ir_structure.h +++ b/src/ascend/transform/auto_schedule/ir_structure.h @@ -112,6 +112,13 @@ inline bool RegionsEqual(const Region &a, const Region &b) { } inline bool MemoryAccessesEqual(const MemoryAccess &a, const MemoryAccess &b) { + if (a.domains.size() != b.domains.size()) + return false; + for (size_t i = 0; i < a.domains.size(); ++i) { + if (!a.domains[i].first.same_as(b.domains[i].first) || + !RegionsEqual({a.domains[i].second}, {b.domains[i].second})) + return false; + } arith::Analyzer analyzer; return a.logical->buffer.same_as(b.logical->buffer) && RegionsEqual(a.logical->region, b.logical->region) && @@ -604,6 +611,12 @@ class TaskNode : public IRStructure { access.bounds = Substitute(access.bounds, substitution); access.repeat = Substitute(access.repeat, substitution); access.stride = Substitute(access.stride, substitution); + for (auto &[var, range] : access.domains) { + if (var.same_as(old_var)) + var = new_var; + range = Range::FromMinExtent(Substitute(range->min, substitution), + Substitute(range->extent, substitution)); + } } }; substitute_vars(read_vars_); diff --git a/src/ascend/transform/auto_schedule/memory_detector.h b/src/ascend/transform/auto_schedule/memory_detector.h index 8c94deaf39..67efd27c2f 100644 --- a/src/ascend/transform/auto_schedule/memory_detector.h +++ b/src/ascend/transform/auto_schedule/memory_detector.h @@ -10,6 +10,7 @@ #include #include +#include #include #include #include @@ -20,11 +21,15 @@ #include #include +#include "ascend/layout/ascend_layouts.h" #include "ascend/op/ascend_mte_plan.h" #include "ascend/op/builtin.h" #include "ascend/op/copy.h" +#include "ascend/op/oob_padding.h" #include "ascend/op/utils.h" #include "op/builtin.h" +#include "op/gemm.h" +#include "op/operator.h" #include "op/utils.h" #include "support/check.h" #include "transform/common/attr.h" @@ -50,9 +55,12 @@ struct MemoryAccess { // Repeat the last axis at [min + i * stride, min + i * stride + extent), // 0 <= i < repeat. Ordinary rectangular accesses use repeat=1, stride=0. // A non-positive repeat or axis extent denotes an empty access. Stride may - // be zero or negative. The detector currently emits ordinary accesses. + // be zero or negative. PrimExpr repeat; PrimExpr stride; + // Loops inside an atomic task. Keep address/pitch correlations until each + // side of an overlap query chooses its own iteration. + std::vector> domains; }; inline std::vector @@ -64,18 +72,140 @@ LogicalRegions(const std::vector &accesses) { return regions; } +using BufferMemoryAccesses = + std::unordered_map; + +// Unknown L0 addressing covers a whole selected plane. Other storage keeps +// its normal coordinate system, with instruction padding for L1 fractals. +inline MemoryAccess MakeMemoryAccess(const BufferRegion ®ion) { + MemoryAccess access(region); + access.bounds.clear(); + const Buffer &buffer = region->buffer; + if (IsL0ABuffer(buffer) || IsL0BBuffer(buffer) || IsL0CBuffer(buffer)) { + for (size_t i = 0; i + 2 < region->region.size(); ++i) + access.bounds.push_back(region->region[i]); + int alignment = + IsL0CBuffer(buffer) ? 16 : std::max(16, AscendC0(buffer->dtype.bits())); + PrimExpr bytes = make_const(DataType::Int(64), buffer->dtype.bits()); + for (size_t i = buffer->shape.size() - 2; i < buffer->shape.size(); ++i) + bytes *= cast(DataType::Int(64), + floordiv(buffer->shape[i] + alignment - 1, alignment) * + alignment); + access.bounds.push_back( + Range::FromMinExtent(make_zero(DataType::Int(64)), floordiv(bytes, 8))); + } else { + access.bounds = ExpandAscendFractalMayAccess(region)->region; + } + return access; +} + +inline MemoryAccess FractalAccess(const BufferRegion ®ion, + const AscendFractalLayoutInfo &info, + const AscendFractalRegionInfo &fractal, + PrimExpr pitch, PrimExpr count, + PrimExpr width) { + MemoryAccess access = MakeMemoryAccess(region); + // Address arithmetic must not overflow the logical index's int32 type. + auto wide = [](PrimExpr value) { return cast(DataType::Int(64), value); }; + PrimExpr block_bytes = + floordiv(wide(info.row_frac * info.c0) * region->buffer->dtype.bits(), 8); + PrimExpr base = + (wide(fractal.outer0->min) * wide(pitch) + wide(fractal.outer1->min)) * + block_bytes; + access.bounds.Set(access.bounds.size() - 1, + Range::FromMinExtent(base, wide(width) * block_bytes)); + access.repeat = wide(count); + access.stride = wide(pitch) * block_bytes; + return access; +} + +// Describe each instruction using its own fractal geometry and pitch. +inline BufferMemoryAccesses GetCallMemoryAccesses(const Call &call, + const LayoutMap &layouts) { + BufferMemoryAccesses result; + if (IsAscendCopyCall(call.get())) { + AscendCopy copy(call->args, call->annotations); + if (IsL1Buffer(copy->src) && + (IsL0ABuffer(copy->dst) || IsL0BBuffer(copy->dst))) { + auto src_layout = FindLayoutForBuffer(layouts, copy->src); + auto dst_layout = FindLayoutForBuffer(layouts, copy->dst); + AscendFractalLayoutInfo src_info, dst_info; + AscendFractalRegionInfo src, dst; + if (!src_layout.defined() || !dst_layout.defined() || + !TryExtractAscendFractalLayout(src_layout.value(), copy->src, + &src_info) || + !TryExtractAscendFractalLayout(dst_layout.value(), copy->dst, + &dst_info) || + !IsAscendMajorKind(src_info.kind) || + !IsAscendMajorKind(dst_info.kind) || + !TryExtractAscendFractalRegion(src_layout.value(), copy->src_range, + &src) || + !TryExtractAscendFractalRegion(dst_layout.value(), copy->dst_range, + &dst)) + return result; + PrimExpr count = src.outer0->extent; + PrimExpr width = src.outer1->extent; + if ((copy->transpose != 0) ^ (src_info.c0_axis != dst_info.c0_axis)) { + count = floordiv(src.outer1->extent * 16, src_info.c0); + width = floordiv(src.outer0->extent * src_info.c0, 16); + } + result.emplace(copy->dst, + FractalAccess(BufferRegion(copy->dst, copy->dst_range), + dst_info, dst, + copy->stride.value_or(dst.outer1->extent), + count, width)); + } else if (IsL0CBuffer(copy->src)) { + Layout layout = MakeAscendL0CLayout(copy->src); + AscendFractalLayoutInfo info; + AscendFractalRegionInfo fractal; + TryExtractAscendFractalLayout(layout, copy->src, &info); + TryExtractAscendFractalRegion(layout, copy->src_range, &fractal); + result.emplace( + copy->src, + FractalAccess(BufferRegion(copy->src, copy->src_range), info, fractal, + copy->stride.value_or(fractal.outer1->extent), + fractal.outer0->extent, fractal.outer1->extent)); + } + } else if (call->op.same_as(Op::Get("tl.tileop.gemm"))) { + Gemm gemm(call->args, call->annotations); + if (!IsL0ABuffer(gemm->a_) || !IsL0BBuffer(gemm->b_)) + return result; + auto record = [&](const BufferRegion ®ion, const Layout &layout) { + AscendFractalLayoutInfo info; + AscendFractalRegionInfo fractal; + TryExtractAscendFractalLayout(layout, region->buffer, &info); + TryExtractAscendFractalRegion(layout, region->region, &fractal); + result.emplace(region->buffer, FractalAccess(region, info, fractal, + fractal.outer1->extent, + fractal.outer0->extent, + fractal.outer1->extent)); + }; + record(gemm->aRegion_, gemm->transA_ ? MakeAscendMajorMNLayout(gemm->a_) + : MakeAscendMajorKLayout(gemm->a_)); + record(gemm->bRegion_, gemm->transB_ ? MakeAscendMajorKLayout(gemm->b_) + : MakeAscendMajorMNLayout(gemm->b_)); + record(gemm->cRegion_, MakeAscendL0CLayout(gemm->c_)); + } + return result; +} + // MemoryAccessDetector: detect read/write regions in statements // Adapted from BlockReadWriteDetector in TVM class MemoryAccessDetector : public StmtExprVisitor { public: - // Task/synchronization access semantics stay logical. Physical footprints - // are requested only by buffer-lifetime analysis and never stored in Task. + // Task accesses carry logical regions and instruction conflict coordinates. + // Lifetime analysis separately requests may/must-write coverage. enum class AccessMode { kLogical, kPhysical }; - explicit MemoryAccessDetector(AccessMode mode = AccessMode::kLogical) - : mode_(mode) {} + explicit MemoryAccessDetector(AccessMode mode = AccessMode::kLogical, + LayoutMap layouts = {}) + : layouts_(std::move(layouts)), mode_(mode) {} // Analyze a statement and collect read/write regions void Analyze(const Stmt &stmt) { + read_accesses_.clear(); + write_accesses_.clear(); + operation_accesses_.clear(); read_buffers_.clear(); write_buffers_.clear(); read_regions_.clear(); @@ -97,6 +227,9 @@ class MemoryAccessDetector : public StmtExprVisitor { } void Analyze(const PrimExpr &expr) { + read_accesses_.clear(); + write_accesses_.clear(); + operation_accesses_.clear(); read_buffers_.clear(); write_buffers_.clear(); read_regions_.clear(); @@ -117,21 +250,6 @@ class MemoryAccessDetector : public StmtExprVisitor { operator()(expr); } - // Wrap the existing per-buffer rectangular summaries. Producing separate - // records for individual accesses is a later precision improvement. - std::vector GetReadAccesses() const { - std::vector accesses; - for (const BufferRegion ®ion : GetReadRegions()) - accesses.emplace_back(region); - return accesses; - } - std::vector GetWriteAccesses() const { - std::vector accesses; - for (const BufferRegion ®ion : GetWriteRegions()) - accesses.emplace_back(region); - return accesses; - } - // Return collected read regions std::vector GetReadRegions() const { return CollectRegions(read_buffers_, read_regions_); @@ -142,6 +260,13 @@ class MemoryAccessDetector : public StmtExprVisitor { return CollectRegions(write_buffers_, write_regions_); } + const std::vector &GetReadAccesses() const { + return read_accesses_; + } + const std::vector &GetWriteAccesses() const { + return write_accesses_; + } + // Return all variables that are read from std::vector GetReadVars() const { return read_vars_; } // Return all variables that are written to @@ -161,6 +286,34 @@ class MemoryAccessDetector : public StmtExprVisitor { } private: + LayoutMap layouts_; + BufferMemoryAccesses operation_accesses_; + std::vector read_accesses_; + std::vector write_accesses_; + + MemoryAccess DescribeAccess(const Buffer &buffer, + const std::vector ®ion) { + MemoryAccess access = MakeMemoryAccess(MakeBufferRegion(buffer, region)); + auto it = operation_accesses_.find(buffer); + if (it == operation_accesses_.end()) + return access; + BufferRegion logical = access.logical; + access = it->second; + access.logical = std::move(logical); + Array bounds; + for (const Range &range : access.bounds) + bounds.push_back(Range::FromMinExtent(ResolveLets(range->min), + ResolveLets(range->extent))); + access.bounds = std::move(bounds); + access.repeat = ResolveLets(access.repeat); + access.stride = ResolveLets(access.stride); + for (const auto &[var, range] : loop_domains_) + access.domains.emplace_back( + var, Range::FromMinExtent(ResolveLets(range->min), + ResolveLets(range->extent))); + return access; + } + /*! \brief Iteration range for loop_vars */ std::unordered_map dom_map_; /*! \brief Extra iteration range hint for free vars */ @@ -224,6 +377,11 @@ class MemoryAccessDetector : public StmtExprVisitor { for (const PrimExpr &stride : buffer->strides) visit_mutable_layout(stride); visit_mutable_layout(buffer->elem_offset); + auto &accesses = + buffers == &read_buffers_ ? read_accesses_ : write_accesses_; + accesses.push_back(mode_ == AccessMode::kLogical + ? DescribeAccess(buffer, region) + : MemoryAccess(MakeBufferRegion(buffer, region))); MergeRegion(buffers, regions, buffer, region); } @@ -390,7 +548,7 @@ class MemoryAccessDetector : public StmtExprVisitor { } void UpdateOpaqueAccessVar(const Var &var) { - if (mode_ != AccessMode::kPhysical) + if (mode_ == AccessMode::kLogical) return; for (const Var &existing : opaque_access_vars_) { if (existing.same_as(var)) @@ -426,6 +584,36 @@ class MemoryAccessDetector : public StmtExprVisitor { } } + // Lifetime analysis includes instruction padding, without promoting the + // conservative may-write footprint into an initialization fact. + void RecordFractalMayAccess(const BufferRegion ®ion, bool is_read) { + BufferRegion expanded = ExpandAscendFractalMayAccess(region); + std::vector ranges; + for (const Range &range : expanded->region) + ranges.push_back(RelaxAccessIndex(range->min, range->extent)); + if (is_read) + Update(&read_buffers_, &read_regions_, expanded->buffer, + std::move(ranges)); + else + Update(&write_buffers_, &write_regions_, expanded->buffer, + std::move(ranges)); + } + + // Compact storage is opaque to allocation reuse; ordinary operands keep + // their must-write information. Keep the full selected plane as may-access. + void RecordL0MayAccess(const BufferRegion ®ion, bool is_read, + bool compact) { + if (!compact) { + RecordFractalMayAccess(region, is_read); + return; + } + UpdateOpaqueAccessVar(region->buffer->data); + Array ranges = region->region; + for (size_t i = ranges.size() - 2; i < ranges.size(); ++i) + ranges.Set(i, Range::FromMinExtent(0, region->buffer->shape[i])); + RecordFractalMayAccess(BufferRegion(region->buffer, ranges), is_read); + } + BufferRegion ExpandPaddedCopyWriteRegion(const AscendCopy ©, BufferRegion write) { if (copy->data_select == 0 || !IsGlobalBuffer(copy->src) || @@ -962,6 +1150,60 @@ class MemoryAccessDetector : public StmtExprVisitor { if (op->op.same_as(tl::loop_break())) has_loop_break_ = true; + if (mode_ == AccessMode::kLogical && + (IsAscendCopyCall(op) || op->op.same_as(Op::Get("tl.tileop.gemm")))) { + auto saved = std::move(operation_accesses_); + operation_accesses_ = + GetCallMemoryAccesses(ffi::GetRef(op), layouts_); + StmtExprVisitor::VisitExpr_(op); + operation_accesses_ = std::move(saved); + return; + } + + if (mode_ == AccessMode::kPhysical && IsAscendCopyCall(op)) { + AscendCopy copy(op->args, op->annotations); + if (copy->stride.defined() || + (IsL1Buffer(copy->src) && + (IsL0ABuffer(copy->dst) || IsL0BBuffer(copy->dst)))) { + StmtExprVisitor::VisitExpr_(op); + auto record = [&](const BufferRegion ®ion, bool is_read) { + bool l0 = IsL0ABuffer(region->buffer) || + IsL0BBuffer(region->buffer) || IsL0CBuffer(region->buffer); + bool compact = copy->stride.defined(); + for (size_t i = region->region.size() - 2; i < region->region.size(); + ++i) + compact |= !ana_.CanProveEqual(region->region[i]->extent, + region->buffer->shape[i]); + RecordL0MayAccess(region, is_read, l0 && compact); + }; + record(BufferRegion(copy->src, copy->src_range), true); + record(BufferRegion(copy->dst, copy->dst_range), false); + return; + } + } + if (mode_ == AccessMode::kPhysical && + op->op.same_as(Op::Get("tl.tileop.gemm"))) { + Gemm gemm(op->args, op->annotations); + if (IsL0ABuffer(gemm->a_) && IsL0BBuffer(gemm->b_)) { + StmtExprVisitor::VisitExpr_(op); + auto record = [&](const BufferRegion ®ion, bool is_read, + size_t mn_axis) { + bool compact = !ana_.CanProveEqual( + floordiv(region->region[mn_axis]->extent + 15, 16), + floordiv(region->buffer->shape[mn_axis] + 15, 16)); + RecordL0MayAccess(region, is_read, compact); + }; + record(gemm->aRegion_, true, + gemm->a_->shape.size() - (gemm->transA_ ? 1 : 2)); + record(gemm->bRegion_, true, + gemm->b_->shape.size() - (gemm->transB_ ? 2 : 1)); + record(gemm->cRegion_, false, gemm->c_->shape.size() - 2); + if (!is_one(gemm->clearAccum_)) + record(gemm->cRegion_, true, gemm->c_->shape.size() - 2); + return; + } + } + if (mode_ == AccessMode::kPhysical && IsAscendCopyCall(op)) { // Visit the original operands first so symbolic indices and the logical // source/destination regions retain their normal access accounting. diff --git a/src/ascend/transform/auto_schedule/scheduled_tir.cc b/src/ascend/transform/auto_schedule/scheduled_tir.cc index da24e3dc31..f0bf57a822 100644 --- a/src/ascend/transform/auto_schedule/scheduled_tir.cc +++ b/src/ascend/transform/auto_schedule/scheduled_tir.cc @@ -146,8 +146,9 @@ std::shared_ptr MakePlainControlNode() { class ScheduledTIRDecoder : public TaskAwareConstrVisitor { public: ScheduledTIRDecoder(const ConstrSet &outer_ctx, - ScheduledExtraInfoFactory make_extra_info) - : make_extra_info_(make_extra_info) { + ScheduledExtraInfoFactory make_extra_info, + LayoutMap layouts) + : make_extra_info_(make_extra_info), layouts_(std::move(layouts)) { constr_stack_ = outer_ctx.constrs_; } @@ -292,7 +293,7 @@ class ScheduledTIRDecoder : public TaskAwareConstrVisitor { task->SetGuards(std::move(guards)); task->outer_ctx = GetConstrSet(); ApplyTaskResourceUsage(AnalyzeTaskResourceUsage(body), task.get()); - ApplyTaskAccesses(AnalyzeTaskAccesses(body), task.get()); + ApplyTaskAccesses(AnalyzeTaskAccesses(body, layouts_), task.get()); uint16_t pipe_mask = task->GetPipeMask(); if (semantic_task->attr_key == attr::kAscendPerCoreTask) { ICHECK_NE(pipe_mask, 0) @@ -344,6 +345,7 @@ class ScheduledTIRDecoder : public TaskAwareConstrVisitor { std::vector> result_; ScheduledExtraInfoFactory make_extra_info_; + LayoutMap layouts_; }; void ValidateBufferVersions(const BufferVersionMap &buffer_versions, @@ -445,8 +447,12 @@ SBlock EncodeScheduledTIR(ScheduledTIR scheduled_tir) { ScheduledTIR DecodeScheduledTIR(const SBlock &root, const ConstrSet &outer_ctx, ScheduledExtraInfoFactory make_extra_info) { ScheduledTIR scheduled_tir; + LayoutMap layouts; + if (auto value = root->annotations.Get("layout_map")) + layouts = value.value().cast(); scheduled_tir.tree = - ScheduledTIRDecoder(outer_ctx, make_extra_info).Decode(root->body); + ScheduledTIRDecoder(outer_ctx, make_extra_info, std::move(layouts)) + .Decode(root->body); scheduled_tir.metadata.kernel_root = root; if (auto value = root->annotations.Get(kBufferVersionsMap)) { scheduled_tir.metadata.buffer_versions = diff --git a/src/ascend/transform/auto_schedule/task_analysis.h b/src/ascend/transform/auto_schedule/task_analysis.h index e807fabb82..fc5b3e1061 100644 --- a/src/ascend/transform/auto_schedule/task_analysis.h +++ b/src/ascend/transform/auto_schedule/task_analysis.h @@ -397,8 +397,10 @@ inline uint16_t AnalyzeSpecialRegisterWrites(const Stmt &stmt) { return task_analysis_detail::SpecialRegisterWriteAnalyzer::Analyze(stmt); } -inline TaskAccessInfo AnalyzeTaskAccesses(const Stmt &stmt) { - MemoryAccessDetector memory_detector; +inline TaskAccessInfo AnalyzeTaskAccesses(const Stmt &stmt, + const LayoutMap &layouts = {}) { + MemoryAccessDetector memory_detector( + MemoryAccessDetector::AccessMode::kLogical, layouts); memory_detector.Analyze(stmt); TaskAccessInfo result; result.read_regions = memory_detector.GetReadAccesses(); diff --git a/src/ascend/transform/lower_tile_op.cc b/src/ascend/transform/lower_tile_op.cc index 4555732a27..4b72e55e7b 100644 --- a/src/ascend/transform/lower_tile_op.cc +++ b/src/ascend/transform/lower_tile_op.cc @@ -17,6 +17,7 @@ #include #include +#include "ascend/layout/ascend_layouts.h" #include "ascend/op/copy.h" #include "ascend/op/utils.h" #include "cuda/op/builtin.h" @@ -596,6 +597,16 @@ class LowerTileOpPass : arith::IRMutatorWithAnalyzer { new_offset += forward_indices[i] * stride_offset; stride_offset *= new_shape[i]; } + if (auto pitch = access_ptr_call->annotations.Get(attr::kL0AccessPitch)) { + ICHECK(IsL0ABuffer(original_buffer) || IsL0BBuffer(original_buffer)); + ICHECK_GE(new_shape.size(), 4); + size_t n = new_shape.size(); + // Only the matrix's outer0 stride changes. Leading planes keep the + // full allocation size, so compact tails cannot alias adjacent slots. + new_offset += forward_indices[n - 4] * + (Downcast(pitch.value()) - new_shape[n - 3]) * + new_shape[n - 2] * new_shape[n - 1]; + } new_offset = analyzer_->Simplify(new_offset); Array new_indices; layout_remap_.Set(new_buffer, layout); diff --git a/testing/ascend/_ir.py b/testing/ascend/_ir.py index 88e0d44d1b..cc28286c53 100644 --- a/testing/ascend/_ir.py +++ b/testing/ascend/_ir.py @@ -58,12 +58,12 @@ def calls(ir, name): return [node for node in nodes(ir, tirx.Call) if isinstance(node.op, tvm.ir.Op) and node.op.name == name] -def schedule_snapshots(program, pass_names, *, auto_schedule=True): +def schedule_snapshots(program, pass_names, *, auto_schedule=True, infer_layouts=False): """Run the scheduling prerequisites and stop at the last requested boundary. This helper accepts frontend programs, including explicit unrolls and dual - copies. It does not run layout inference: tests needing physical layouts - must prepare those themselves, or belong in the lowering integration suite. + copies. Tests needing physical layouts can opt into layout inference or + supply an input whose layouts are already canonical. """ pipeline = ( ascend.NormalizeControlFlowForSchedule, @@ -116,6 +116,9 @@ def assign_stage(node): mod = ascend.RewriteDualCopy()(mod) mod = ascend.UnrollLoopSkipVF()(mod) mod = tilelang.transform.Simplify()(mod) + if infer_layouts: + mod = ascend.AscendLayoutInference()(mod) + mod = ascend.NormalizeAscendFractalStorage()(mod) for factory in pipeline: if not auto_schedule and factory == ascend.AutoSchedule: mod = use_source_order(mod) @@ -128,9 +131,9 @@ def assign_stage(node): return snapshots -def schedule(program, through="InsertSync", *, auto_schedule=True): +def schedule(program, through="InsertSync", *, auto_schedule=True, infer_layouts=False): name = "tl." + through - return schedule_snapshots(program, {name}, auto_schedule=auto_schedule)[name] + return schedule_snapshots(program, {name}, auto_schedule=auto_schedule, infer_layouts=infer_layouts)[name] def statements(ir, ancestors=()): diff --git a/testing/ascend/auto_schedule/test_access_footprints.py b/testing/ascend/auto_schedule/test_access_footprints.py index d5c4e99d84..da0c13d084 100644 --- a/testing/ascend/auto_schedule/test_access_footprints.py +++ b/testing/ascend/auto_schedule/test_access_footprints.py @@ -9,6 +9,16 @@ from testing.ascend.auto_schedule._reuse_utils import _insert_sync_alias_pairs, _make_cyclic_scalar_program +def _infer_fractal_layouts(program): + """Prepare the physical layouts expected at the scheduling boundary.""" + mod = tvm.IRModule({"main": program.with_attr("target", tvm.target.Target("ascend"))}) + mod = tilelang.transform.MaterializeKernelLaunch()(mod) + with tvm.target.Target("ascend"): + mod = tilelang.ascend.transform.AscendLayoutInference()(mod) + mod = tilelang.ascend.transform.NormalizeAscendFractalStorage()(mod) + return mod["main"] + + def _make_internal_conditional_write_program(): @T.prim_func def main( @@ -598,11 +608,8 @@ def main(x: T.Tensor((16, 64), "float32"), out: T.Tensor((1,), "float32")): T.simd.vsts(other[0, 0], T.simd.vdup(T.float32(200), "float32")) out[0] = other[0, 0] - mod = tvm.IRModule({"main": main.with_attr("target", tvm.target.Target("ascend"))}) - mod = tilelang.transform.MaterializeKernelLaunch()(mod) + mod = tvm.IRModule({"main": _infer_fractal_layouts(main)}) with tvm.target.Target("ascend"): - mod = tilelang.ascend.transform.AscendLayoutInference()(mod) - mod = tilelang.ascend.transform.NormalizeAscendFractalStorage()(mod) mod = tilelang.ascend.transform.InsertNd2Nz()(mod) assert ("dst", "src") in _insert_sync_alias_pairs(mod["main"]) @@ -624,7 +631,7 @@ def main(x: T.Tensor((16, 64), "float32"), out: T.Tensor((1,), "bfloat16")): T.copy(nz, other) out[0] = other[0, 0] - assert ("other", "src") in _insert_sync_alias_pairs(main) + assert ("other", "src") in _insert_sync_alias_pairs(_infer_fractal_layouts(main)) @pytest.mark.parametrize("dtype", ["float32", "float16"]) @@ -808,6 +815,35 @@ def main(A: T.Tensor((2, 4), "float32"), out: T.Tensor((1,), "float32")): assert ("a", "b") in _insert_sync_alias_pairs(main) +@pytest.mark.parametrize("compact", [False, True]) +def test_compact_l0_drain_preserves_ub_write_coverage(compact): + stride = 1 if compact else None + + @T.prim_func + def main(scratch: T.Tensor((16, 16), "float32"), out: T.Tensor((16, 16), "float32")): + with T.Kernel(1): + a = T.alloc_shared((16, 16), "float32") + b = T.alloc_shared((16, 16), "float32") + l1 = T.alloc_l1((16, 16), "float32") + lhs = T.alloc_l0a((16, 16), "float32") + rhs = T.alloc_l0b((16, 16), "float32") + acc = T.alloc_l0c((32 if compact else 16, 16), "float32") + with T.Stage(0): + T.fill(a, 100.0) + T.fill(b, 200.0) + # The round trip orders b's last read before a's new write. + T.copy(b, scratch) + T.copy(scratch, l1) + T.copy(l1, lhs) + T.copy(l1, rhs) + T.gemm(lhs, rhs, acc[:16, :], transpose_B=True, clear_accum=True) + # A compact source pitch must not hide the full UB write. + T.copy(acc[:16, :], a, stride=stride) + T.copy(a, out) + + assert ("a", "b") in _insert_sync_alias_pairs(_infer_fractal_layouts(main)) + + def test_padded_copy_overflow_excludes_storage_from_reuse(): @T.prim_func def main(A: T.Tensor((4, 30), "float32"), C: T.Tensor((4, 30), "float32")): diff --git a/testing/ascend/auto_schedule/test_cross_core_flags.py b/testing/ascend/auto_schedule/test_cross_core_flags.py index b6590b956e..ec9b74c30d 100644 --- a/testing/ascend/auto_schedule/test_cross_core_flags.py +++ b/testing/ascend/auto_schedule/test_cross_core_flags.py @@ -16,7 +16,8 @@ def main(C: T.Tensor((16, 16), "float32"), flag: T.int32): accum = T.alloc_l0c((16, 16), "float32") temp = T.alloc_shared((8, 16), "float32") # This is an InsertSync input: Fixpipe produces UB and MTE3 reads - # it. GEMM, layout inference and device execution are irrelevant. + # it. Layout inference prepares the scheduling input; no GEMM or + # device execution is needed for this flag-allocation contract. T.dual_copy(accum, temp) if reservation == "sparse": T.ascend_sync_inter_arrive("PIPE_FIX", 15) @@ -66,7 +67,7 @@ def _generated_slots(mod, operation): ) def test_generated_flags_do_not_collide_with_explicit_flags(reservation, reserved): with tvm.transform.PassContext(config={"tl.Simplify": {"enable_simplify_let_inline": False}}): - before = schedule(_program(reservation), through="ResolveCore", auto_schedule=False) + before = schedule(_program(reservation), through="ResolveCore", auto_schedule=False, infer_layouts=True) mod = transform.InsertSync()(before) sets = _generated_slots(mod, "set") waits = _generated_slots(mod, "wait") @@ -77,7 +78,7 @@ def test_generated_flags_do_not_collide_with_explicit_flags(reservation, reserve def test_unbounded_explicit_flag_is_rejected(): with pytest.raises(tvm.error.InternalError, match="integer range cannot be bounded"): - before = schedule(_program("unbounded"), through="ResolveCore", auto_schedule=False) + before = schedule(_program("unbounded"), through="ResolveCore", auto_schedule=False, infer_layouts=True) transform.InsertSync()(before) diff --git a/testing/ascend/auto_schedule/test_flag_allocation.py b/testing/ascend/auto_schedule/test_flag_allocation.py index 90e2e7cb0d..d459453ed5 100644 --- a/testing/ascend/auto_schedule/test_flag_allocation.py +++ b/testing/ascend/auto_schedule/test_flag_allocation.py @@ -247,8 +247,8 @@ def main(C: T.Tensor((8, 32, 16), "float32")): output = T.alloc_shared((16, 16), "float32") T.annotate_buffer_versions({output: (versions, "counter") if counter else versions}) for tile in T.serial(8): - # Pass input: Fixpipe produces UB. GEMM and layout inference - # are not prerequisites for this synchronization contract. + # Fixpipe produces UB. The fixture infers canonical layouts + # before scheduling; no GEMM or device execution is needed. T.dual_copy(accum, output, unit_flag_ctrl=3) if reserve_last_slot: T.ascend_sync_inter_arrive("PIPE_FIX", 15) @@ -281,8 +281,8 @@ def main(A: T.Tensor((4 * count * 64,), "float32"), C: T.Tensor((4 * count * 64, return main -def _insert_sync(program): - before = schedule(program, through="ResolveCore", auto_schedule=False) +def _insert_sync(program, *, infer_layouts=False): + before = schedule(program, through="ResolveCore", auto_schedule=False, infer_layouts=infer_layouts) return ascend_transform.InsertSync()(before) @@ -385,7 +385,7 @@ def test_legal_sparse_flags_remain_unchanged(): ) def test_cross_core_rings_fit_available_slots(versions, counter, vector_stage, reserve_slot): program = _make_cross_core_reuse_program(versions, vector_stage, reserve_slot, counter) - mod = _insert_sync(program) + mod = _insert_sync(program, infer_layouts=True) for operation in ("set", "wait"): generated = [call.args[2] for call in calls(mod, "tl.ascend_cross_core_" + operation + "_flag") if int(call.args[0]) == 4] assert generated diff --git a/testing/ascend/auto_schedule/test_sync_access_footprints.py b/testing/ascend/auto_schedule/test_sync_access_footprints.py index 750be74589..331003ceb5 100644 --- a/testing/ascend/auto_schedule/test_sync_access_footprints.py +++ b/testing/ascend/auto_schedule/test_sync_access_footprints.py @@ -5,7 +5,9 @@ import tilelang.ascend.language as T from tilelang import tvm from tilelang.ascend import transform -from testing.ascend._ir import calls, schedule +from testing.ascend._ir import calls, copy, kernel, schedule, seq +from testing.ascend.auto_schedule._scheduled_ir import unit +from tvm import tirx def _events(program, disable_reuse=False, auto_schedule=False): @@ -17,6 +19,45 @@ def _events(program, disable_reuse=False, auto_schedule=False): } +@pytest.mark.parametrize("fill_first", [False, True], ids=["protect-source-reuse", "wait-for-source-padding"]) +@pytest.mark.parametrize("row", [24, 32], ids=["same-fractal", "next-fractal"]) +def test_l0_tail_load_accounts_for_source_fractal_padding(fill_first, row): + src = tirx.decl_buffer((48, 16), "bfloat16", name="src", scope="shared.l1") + dst = tirx.decl_buffer((32, 16), "bfloat16", name="dst", scope="shared.l0a") + load = unit(copy(src, dst, src_shape=[24, 16], dst_shape=[24, 16]), core=2) + padding = tirx.BufferRegion(src, [tvm.ir.Range.from_min_extent(row, 8), tvm.ir.Range(0, 16)]) + fill = unit(tirx.Evaluate(T.fill(padding, 0)), core=2) + # L1 fill can address individual rows. The load reads two whole row16 + # fractals even though its effective region ends at row 24. + before = kernel(seq(fill, load) if fill_first else seq(load, fill), buffers=[src, dst]) + after = transform.InsertSync()(before) + pipe = "MTE2_MTE1" if fill_first else "MTE1_MTE2" + events = [call for op in ("set", "wait") for call in calls(after, "tl.ascend_" + op + "_flag") if call.args[0].value == pipe] + if row == 24: + assert len(events) == 2 + tvm.ir.assert_structural_equal(events[0].args[1], events[1].args[1]) + else: + assert not events + + +def test_l0_tail_load_preserves_explicit_logical_region_conflicts(): + @T.prim_func + def main(): + with T.Kernel(1): + left = T.alloc_l1((32, 16), "bfloat16") + right = T.alloc_l1((32, 16), "bfloat16") + dst = T.alloc_l0a((32, 16), "bfloat16") + T.assume_conflict(left[:24, :], right[:24, :], cross=False) + with T.Stage(0): + T.fill(left[:24, :], 0) + # The source is already prepared at this sync-only boundary. + T.copy(right[:24, :], dst[:24, :]) + + events = _events(main) + assert ("set", "MTE2_MTE1") in events + assert ("wait", "MTE2_MTE1") in events + + @pytest.mark.parametrize("disable_reuse", [False, True]) @pytest.mark.parametrize("strided", [False, True]) def test_pointer_store_past_view_keeps_scalar_read_dependency(disable_reuse, strided): diff --git a/testing/ascend/auto_schedule/test_sync_dependencies.py b/testing/ascend/auto_schedule/test_sync_dependencies.py index 48c40100ec..fdd9678819 100644 --- a/testing/ascend/auto_schedule/test_sync_dependencies.py +++ b/testing/ascend/auto_schedule/test_sync_dependencies.py @@ -2,7 +2,8 @@ import pytest from tilelang import tvm -from tilelang.ascend import transform +from tilelang.ascend import language as T, transform +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, calls, copy, kernel, nodes, seq, statements from testing.ascend.auto_schedule._scheduled_ir import copy_ring, unit @@ -135,6 +136,155 @@ def test_dma_row_writes_only_serialize_overlapping_regions(overlap): tvm.ir.assert_structural_equal(releases[0].args[1], waits[0].args[1]) +@pytest.mark.parametrize("overlap", [False, True], ids=["disjoint-halves", "overwrite-middle"]) +def test_partial_l0_loads_complete_before_gemm_and_order_overlapping_writes(overlap): + src = tirx.decl_buffer((64, 64), "bfloat16", name="src", scope="shared.l1") + a = tirx.decl_buffer((64, 64), "bfloat16", name="a", scope="shared.l0a") + b = tirx.decl_buffer((64, 64), "bfloat16", name="b", scope="shared.l0b") + c = tirx.decl_buffer((64, 64), "float32", name="c", scope="shared.l0c") + # B and the L1 source are already prepared at this pass boundary. + # The first destination can hold 48 rows; its actual write covers only + # the source's 32 or 48 rows. Only the latter overlaps the next copy. + body = seq( + unit( + copy( + src, + a, + src_shape=[48 if overlap else 32, 64], + dst_shape=[48, 64], + annotations={"stride": tirx.IntImm("int32", 4)}, + ), + core=2, + ), + unit( + copy( + src, + a, + src_indices=[32, 0], + dst_indices=[32, 0], + src_shape=[32, 64], + dst_shape=[32, 64], + annotations={"stride": tirx.IntImm("int32", 4)}, + ), + core=2, + ), + unit(tirx.Evaluate(T.gemm(a, b, c, transpose_B=True, clear_accum=True)), core=2), + ) + layouts = {src: make_ascend_major_k_layout(src), a: make_ascend_major_k_layout(a)} + after = transform.InsertSync()(kernel(body, buffers=[src, a, b, c], annotations={"layout_map": layouts})) + barriers = [call for call in calls(after, "tl.ascend_pipe_barrier") if call.args[0].value == "PIPE_MTE1"] + assert bool(barriers) == overlap + releases, waits = _events(after, "set", "MTE1_M"), _events(after, "wait", "MTE1_M") + assert releases and waits + for release in releases: + assert any(tvm.ir.structural_equal(release.args[1], wait.args[1]) for wait in waits) + operations = [stmt.value for stmt in nodes(after, tirx.Evaluate) if isinstance(stmt.value, tirx.Call)] + loads = [i for i, call in enumerate(operations) if call.op.name == "tl.tileop.ascend_copy"] + release_sites = [i for i, call in enumerate(operations) if any(call.same_as(release) for release in releases)] + wait_sites = [i for i, call in enumerate(operations) if any(call.same_as(wait) for wait in waits)] + gemm_site = next(i for i, call in enumerate(operations) if call.op.name == "tl.tileop.gemm") + assert len(loads) == 2 and max(loads) < max(release_sites) <= max(wait_sites) < gemm_site + + +@pytest.mark.parametrize("transposed", [False, True], ids=["major-k", "major-mn"]) +@pytest.mark.parametrize( + "plane, mn_start, k_start, mn_size, k_size, overlap", + [(1, 32, 0, 32, 64, False), (0, 32, 0, 32, 64, True), (0, 48, 48, 16, 16, False)], + ids=["separate-planes", "physical-overlap", "physical-disjoint"], +) +def test_different_l0_pitches_order_physical_overlap(transposed, plane, mn_start, k_start, mn_size, k_size, overlap): + src = tirx.decl_buffer((64, 64), "bfloat16", name="src", scope="shared.l1") + a = tirx.decl_buffer((2, 64, 64), "bfloat16", name="a", scope="shared.l0b" if transposed else "shared.l0a") + b = tirx.decl_buffer((64, 64), "bfloat16", name="b", scope="shared.l0a" if transposed else "shared.l0b") + c = tirx.decl_buffer((32, 64), "float32", name="c", scope="shared.l0c") + # GEMM reads rows [0,32) at pitch=2. A subsequent load of [32,64) + # at pitch=4 overlaps that compact storage, despite logical separation. + operand_shape = (1, 64, 32) if transposed else (1, 32, 64) + lhs = tirx.BufferRegion(a, [tvm.ir.Range.from_min_extent(0, e) for e in operand_shape]) + if transposed: + c = tirx.decl_buffer((64, 32), "float32", name="c", scope="shared.l0c") + gemm = T.gemm(b, lhs, c, clear_accum=True) + start, shape = [k_start, mn_start], [k_size, mn_size] + else: + gemm = T.gemm(lhs, b, c, transpose_B=True, clear_accum=True) + start, shape = [mn_start, k_start], [mn_size, k_size] + body = seq( + unit(tirx.Evaluate(gemm), core=2), + unit( + copy( + src, + a, + src_shape=shape, + dst_shape=[1, *shape], + dst_indices=[plane, *start], + annotations={"stride": tirx.IntImm("int32", 4)}, + ), + core=2, + ), + ) + layouts = { + src: make_ascend_major_k_layout(src), + a: (make_ascend_major_mn_layout if transposed else make_ascend_major_k_layout)(a), + } + after = transform.InsertSync()(kernel(body, buffers=[src, a, b, c], annotations={"layout_map": layouts})) + assert bool(_events(after, "set", "M_MTE1")) == overlap + assert bool(_events(after, "wait", "M_MTE1")) == overlap + + +@pytest.mark.parametrize("varying_pitch", [False, True], ids=["disjoint", "overlap-after-pitch-change"]) +def test_l0_copy_pitch_is_instantiated_per_iteration(varying_pitch): + src = tirx.decl_buffer((16, 16), "bfloat16", name="src", scope="shared.l1") + a = tirx.decl_buffer((32, 32), "bfloat16", name="a", scope="shared.l0a") + i = tirx.Var("i", "int32") + # Disjoint logical blocks (m,k)=(0,1),(1,0). With pitches 1,2 they + # both address physical block 1; at a fixed pitch 2 they address 2,1. + load = copy( + src, + a, + dst_indices=[16 * i, 16 * (1 - i)], + dst_shape=[16, 16], + annotations={"stride": 1 + i if varying_pitch else tirx.IntImm("int32", 2)}, + ) + body = unit(tirx.For(i, 0, 2, tirx.ForKind.SERIAL, unit(load, core=2)), core=None) + layouts = {src: make_ascend_major_k_layout(src), a: make_ascend_major_k_layout(a)} + after = transform.InsertSync()(kernel(body, buffers=[src, a], annotations={"layout_map": layouts})) + barriers = [call for call in calls(after, "tl.ascend_pipe_barrier") if call.args[0].value == "PIPE_MTE1"] + assert bool(barriers) == varying_pitch + + +@pytest.mark.parametrize("overlap", [False, True], ids=["disjoint", "overlap"]) +@pytest.mark.parametrize("unroll", [False, True], ids=["internal-loop", "multiple-calls"]) +def test_atomic_l0_copies_keep_address_and_pitch_correlated(overlap, unroll): + src = tirx.decl_buffer((16, 16), "bfloat16", name="src", scope="shared.l1") + a = tirx.decl_buffer((32, 32), "bfloat16", name="a", scope="shared.l0a") + b = tirx.decl_buffer((16, 32 if overlap else 16), "bfloat16", name="b", scope="shared.l0b") + c = tirx.decl_buffer((16, 16), "float32", name="c", scope="shared.l0c") + i = tirx.Var("i", "int32") + load = copy( + src, + a, + dst_indices=[16 * i, 16 * (1 - i)], + dst_shape=[16, 16], + annotations={"stride": 1 + i}, + ) + lhs = tirx.BufferRegion(a, [tvm.ir.Range.from_min_extent(0, e) for e in (16, 32 if overlap else 16)]) + # Both loop iterations write block 1. The consumer reads block 0 or + # blocks [0,2); taking the logical loop hull loses this distinction. + loads = ( + seq(*(tirx.stmt_functor.substitute(load, {i: value}) for value in (0, 1))) + if unroll + else tirx.For(i, 0, 2, tirx.ForKind.SERIAL, load) + ) + body = seq( + unit(loads, core=2), + unit(tirx.Evaluate(T.gemm(lhs, b, c, transpose_B=True, clear_accum=True)), core=2), + ) + layouts = {src: make_ascend_major_k_layout(src), a: make_ascend_major_k_layout(a)} + after = transform.InsertSync()(kernel(body, buffers=[src, a, b, c], annotations={"layout_map": layouts})) + assert bool(_events(after, "set", "MTE1_M")) == overlap + assert bool(_events(after, "wait", "MTE1_M")) == overlap + + def test_equivalent_lexical_epoch_serializes_counter_ring(): a = tirx.decl_buffer((64,), "float32", name="A") ub = tirx.decl_buffer((64,), "float32", name="ub", scope="shared.dyn") diff --git a/testing/ascend/layout/test_ascend_partial_l0_copy.py b/testing/ascend/layout/test_ascend_partial_l0_copy.py new file mode 100644 index 0000000000..a977e5df8d --- /dev/null +++ b/testing/ascend/layout/test_ascend_partial_l0_copy.py @@ -0,0 +1,185 @@ +"""Explicit copy stride assembles L0 tiles across compact tails and ring reuse.""" + +import pytest +import torch + +import tilelang +from tilelang.ascend import language as T + + +def _partial_copy_kernel(mode, dtype, versions): + steps, n, k_alloc = 7, 64, 64 + m, m_alloc = (40, 64) if mode == "m-tail" else (64, 64) + k = 56 if mode == "k-tail" else k_alloc + v_layout = mode in ("v", "k-tail", "v-tail") + + @T.prim_func + def main( + A: T.Tensor((steps, m_alloc, k_alloc), dtype), + B: T.Tensor((steps, n, k_alloc), dtype), + C: T.Tensor((steps, m, n), "float32"), + ): + with T.Kernel(1): + a_l1 = T.alloc_l1((m_alloc, k_alloc), dtype) + b_l1 = T.alloc_l1((n, k_alloc), dtype) + a = T.alloc_l0a((m_alloc, k_alloc), dtype) + b = T.alloc_l0b((n, k_alloc), dtype) + acc = T.alloc_l0c((m_alloc, n), "float32") + T.annotate_buffer_versions({a: versions, b: versions}) + for e in T.Pipelined(steps, num_stages=versions): + effective_m = 24 + 16 * (e % 3) if mode == "dynamic-tail" else m + effective_n = 24 + 16 * (e % 3) if mode == "v-tail" else n + padded_m = T.ceildiv(effective_m, 16) * 16 + padded_n = T.ceildiv(effective_n, 16) * 16 + T.copy(A[e, :, :], a_l1) + T.copy(B[e, :, :], b_l1) + if mode == "overwrite": + T.copy(a_l1, a) + T.copy(a_l1[:32, :], a[32:64, :], stride=T.ceildiv(effective_m, 16)) + elif mode in ("m-tail", "dynamic-tail"): + T.copy(a_l1[:16, :], a[:16, :], stride=T.ceildiv(effective_m, 16)) + T.copy(a_l1[16:padded_m, :], a[16:padded_m, :], stride=T.ceildiv(effective_m, 16)) + else: + for part in T.serial(2): + p = 1 - part if mode == "reverse" else part + if mode == "k-tail": + T.copy(a_l1[:, p * 32 : (p + 1) * 32], a[:, p * 32 : (p + 1) * 32], stride=T.ceildiv(effective_m, 16)) + else: + T.copy(a_l1[p * 32 : (p + 1) * 32, :], a[p * 32 : (p + 1) * 32, :], stride=T.ceildiv(effective_m, 16)) + for part in T.serial(2): + # B is [K,N] for V, [N,K] otherwise. GEMM anchors the + # MajorMN/MajorK layout; the copy has no explicit transpose. + if mode == "v-tail": + T.copy( + b_l1[part * 32 : (part + 1) * 32, :padded_n], + b[part * 32 : (part + 1) * 32, :padded_n], + stride=T.ceildiv(effective_n, 16), + ) + else: + T.copy(b_l1[part * 32 : (part + 1) * 32, :], b[part * 32 : (part + 1) * 32, :], stride=T.ceildiv(effective_n, 16)) + if v_layout: + T.gemm(a[:effective_m, :k], b[:k, :effective_n], acc[:effective_m, :effective_n], clear_accum=True) + else: + T.gemm(a[:effective_m, :k], b[:, :k], acc[:effective_m, :], transpose_B=True, clear_accum=True) + drain_m = 16 if mode == "partial-drain" else effective_m + T.copy( + acc[:drain_m, :effective_n], + C[e, :drain_m, :effective_n], + stride=T.ceildiv(effective_m, 16) if mode == "partial-drain" else None, + ) + + return main + + +@pytest.mark.parametrize( + "mode,dtype,versions", + [ + ("ordinary", "bfloat16", 2), + ("reverse", "bfloat16", 3), + ("v", "bfloat16", 2), + ("k-tail", "bfloat16", 3), + ("m-tail", "bfloat16", 2), + ("v-tail", "bfloat16", 3), + ("dynamic-tail", "bfloat16", 3), + ("overwrite", "float16", 2), + ("partial-drain", "float16", 2), + ], +) +def test_partial_copies_feed_full_gemm_across_versions(mode, dtype, versions): + if not torch.npu.is_available(): + pytest.skip("Ascend NPU unavailable") + m, m_alloc = (40, 64) if mode == "m-tail" else (64, 64) + k = 56 if mode == "k-tail" else 64 + generator = torch.Generator().manual_seed(23) + a = torch.randint(-4, 5, (7, m_alloc, 64), generator=generator).to(getattr(torch, dtype)) / 8 + b = torch.randint(-4, 5, (7, 64, 64), generator=generator).to(getattr(torch, dtype)) / 8 + # Different block/epoch values expose wrong pitch, missing loads and old slots. + a[:, 32:, :] += 1 + a += torch.arange(7, dtype=getattr(torch, dtype))[:, None, None] / 8 + if mode == "k-tail": + a[:, :, k:] = 2 + b[:, k:, :] = 3 + effective_a = torch.cat((a[:, :32], a[:, :32]), dim=1) if mode == "overwrite" else a + effective_b = b[:, :k, :] if mode in ("v", "k-tail", "v-tail") else b[:, :, :k].transpose(-1, -2) + expected = effective_a[:, :m, :k].float() @ effective_b.float() + compiled = tilelang.compile(_partial_copy_kernel(mode, dtype, versions), out_idx=-1, target="ascend") + actual = compiled(a.npu(), b.npu()) + torch.npu.synchronize() + actual = actual.cpu() + for e in range(7): + valid_m = 24 + 16 * (e % 3) if mode == "dynamic-tail" else m + if mode == "partial-drain": + valid_m = 16 + valid_n = 24 + 16 * (e % 3) if mode == "v-tail" else 64 + torch.testing.assert_close(actual[e, :valid_m, :valid_n], expected[e, :valid_m, :valid_n], atol=1e-5, rtol=1e-5) + + +@pytest.mark.parametrize("v_layout", [False, True], ids=["ordinary", "v-transpose"]) +def test_manual_partial_copies_with_branch_specific_gemm_pitch(v_layout): + if not torch.npu.is_available(): + pytest.skip("Ascend NPU unavailable") + + @T.prim_func + def main(A: T.Tensor((3, 64, 64), "bfloat16"), B: T.Tensor((3, 64, 64), "bfloat16"), C: T.Tensor((3, 64, 64), "float32")): + with T.Kernel(1): + a_l1 = T.alloc_l1((64, 64), "bfloat16") + b_l1 = T.alloc_l1((64, 64), "bfloat16") + a = T.alloc_l0a((64, 64), "bfloat16") + b = T.alloc_l0b((64, 64), "bfloat16") + acc = T.alloc_l0c((64, 64), "float32") + with T.Cube(): + for e in T.serial(3): + m = 24 + 32 * (e % 2) + n = 56 - 32 * (e % 2) + padded_m = T.ceildiv(m, 16) * 16 + padded_n = T.ceildiv(n, 16) * 16 + T.copy(A[e, :, :], a_l1) + T.copy(B[e, :, :], b_l1) + T.ascend_set_flag("MTE2_MTE1", 0) + T.ascend_wait_flag("MTE2_MTE1", 0) + for part in T.serial(2): + T.copy( + a_l1[:padded_m, part * 32 : (part + 1) * 32], a[:padded_m, part * 32 : (part + 1) * 32], stride=T.ceildiv(m, 16) + ) + if v_layout: + T.copy( + b_l1[part * 32 : (part + 1) * 32, :padded_n], + b[part * 32 : (part + 1) * 32, :padded_n], + stride=T.ceildiv(n, 16), + ) + else: + T.copy( + b_l1[:padded_n, part * 32 : (part + 1) * 32], + b[:padded_n, part * 32 : (part + 1) * 32], + stride=T.ceildiv(n, 16), + ) + T.ascend_set_flag("MTE1_M", 0) + T.ascend_wait_flag("MTE1_M", 0) + if e % 2 == 0: + if v_layout: + T.gemm(a[:24, :], b[:, :56], acc[:24, :56], clear_accum=True) + else: + T.gemm(a[:24, :], b[:56, :], acc[:24, :56], transpose_B=True, clear_accum=True) + else: + if v_layout: + T.gemm(a[:56, :], b[:, :24], acc[:56, :24], clear_accum=True) + else: + T.gemm(a[:56, :], b[:24, :], acc[:56, :24], transpose_B=True, clear_accum=True) + T.ascend_set_flag("M_FIX", 0) + T.ascend_wait_flag("M_FIX", 0) + # The complete result uses the original compact-copy default. + T.copy(acc[:m, :n], C[e, :m, :n]) + T.ascend_pipe_barrier("PIPE_ALL") + + generator = torch.Generator().manual_seed(31) + a = torch.randint(-4, 5, (3, 64, 64), generator=generator).to(torch.bfloat16) / 8 + b = torch.randint(-4, 5, (3, 64, 64), generator=generator).to(torch.bfloat16) / 8 + a[:, :, 32:] += 1 + expected = a.float() @ (b.float() if v_layout else b.float().transpose(-1, -2)) + compiled = tilelang.compile(main, out_idx=-1, target="ascend", pass_configs={tilelang.PassConfigKey.TL_ENABLE_AUTO_SCHEDULE: False}) + actual = compiled(a.npu(), b.npu()) + torch.npu.synchronize() + actual = actual.cpu() + for e in range(3): + m, n = 24 + 32 * (e % 2), 56 - 32 * (e % 2) + torch.testing.assert_close(actual[e, :m, :n], expected[e, :m, :n], atol=1e-5, rtol=1e-5) diff --git a/testing/ascend/transform/test_ascend_lower_copy.py b/testing/ascend/transform/test_ascend_lower_copy.py index 61054fd266..7136283726 100644 --- a/testing/ascend/transform/test_ascend_lower_copy.py +++ b/testing/ascend/transform/test_ascend_lower_copy.py @@ -1,5 +1,7 @@ """AscendLowerTileOp maps logical copy regions to DMA instruction geometry.""" +import math + import pytest from tilelang import tvm from tilelang.ascend import language as T, transform @@ -11,7 +13,7 @@ make_ascend_sf_layout, ) from tvm import tirx -from testing.ascend._ir import allocated_buffer, calls, nodes +from testing.ascend._ir import allocated_buffer, calls, nodes, kernel def _region(buffer, extents=None, mins=None): @@ -163,10 +165,11 @@ def test_nonrepresentable_dma_geometry_is_rejected(scope): _lower(src, dst, dst_region=_region(dst, [6, 4] if scope == "shared.l1" else [2, 3, 4]), layouts=layout) +@pytest.mark.parametrize("explicit", [False, True], ids=["compact-default", "explicit-stride"]) @pytest.mark.parametrize("symbolic", [False, True], ids=["static-region", "dynamic-region"]) @pytest.mark.parametrize("scope", ["shared.l0a", "shared.l0b"]) -def test_l1_to_l0_pitch_comes_from_the_compact_destination_region(symbolic, scope): - m, k = (tirx.Var("m", "int32"), tirx.Var("k", "int32")) if symbolic else (17, 18) +def test_l1_to_l0_compact_pitch_can_be_overridden(explicit, symbolic, scope): + m, k = (T.min(T.max(tirx.Var(name, "int32"), 1), 64) for name in ("m", "k")) if symbolic else (17, 18) src = tirx.decl_buffer((64, 64), "bfloat16", name="src", scope="shared.l1") dst = tirx.decl_buffer((64, 64), "bfloat16", name="dst", scope=scope) after = _lower( @@ -175,9 +178,153 @@ def test_l1_to_l0_pitch_comes_from_the_compact_destination_region(symbolic, scop _region(src, [m, k]), _region(dst, [m, k]), layouts={src: make_ascend_nz_layout(src), dst: make_ascend_major_k_layout(dst)}, + stride=4 if explicit else None, ) (copy,) = calls(after, "tl.ascend_load_cbuf_to_ca" if scope.endswith("a") else "tl.ascend_load_cbuf_to_cb") - _equal(copy.args[4:8], [T.ceildiv(m, 16), T.ceildiv(k, 16), 4, T.ceildiv(m, 16)]) + _equal(copy.args[4:8], [T.ceildiv(m, 16), T.ceildiv(k, 16), 4, 4 if explicit else T.ceildiv(m, 16)]) + + +@pytest.mark.parametrize( + "dtype,major_mn,transpose,leading,offset,geometry", + [ + ("bfloat16", False, False, False, 4608, [2, 4, 2, 4, 4, 4, 0]), + ("bfloat16", True, False, True, 5120, [2, 4, 2, 4, 4, 8, 1]), + ("bfloat16", False, True, False, 5120, [2, 4, 2, 4, 4, 8, 1]), + ("bfloat16", True, True, False, 4608, [2, 4, 2, 4, 4, 4, 0]), + ("float32", True, False, False, 4608, [2, 8, 2, 8, 4, 8, 1]), + ("float8_e4m3fn", True, False, False, 10240, [4, 2, 4, 2, 8, 8, 1]), + ], + ids=["ordinary", "v-leading-plane", "explicit-transpose", "transpose-cancels", "fp32", "fp8"], +) +def test_partial_l0_copy_uses_full_pitch_and_mapped_destination_origin(dtype, major_mn, transpose, leading, offset, geometry): + rows = 128 if dtype == "float8_e4m3fn" else 64 + plane = [2, 3] if leading else [] + index = [1, 2] if leading else [] + src = tirx.decl_buffer((*plane, rows, 128), dtype, name="src", scope="shared.l1") + dst = tirx.decl_buffer((*plane, *((128, rows) if transpose else (rows, 128))), dtype, name="dst", scope="shared.l0b") + shape, origin = [rows // 2, 64], [rows // 2, 64] + after = _lower( + src, + dst, + _region(src, [1] * len(plane) + shape, index + origin), + _region(dst, [1] * len(plane) + (shape[::-1] if transpose else shape), index + (origin[::-1] if transpose else origin)), + layouts={src: make_ascend_major_k_layout(src), dst: (make_ascend_major_mn_layout if major_mn else make_ascend_major_k_layout)(dst)}, + transpose=transpose, + stride=geometry[5], + ) + (load,) = calls(after, "tl.ascend_load_cbuf_to_cb") + base = 5 * rows * 128 if leading else 0 + _equal([load.args[0].args[2], load.args[1].args[2]], [base + offset, base]) + _equal(load.args[2:], geometry) + + +@pytest.mark.parametrize( + "src_min,dst_min,extent,stride,diagnostic", + [ + ([0, 0], [32, 0], [32, 32], None, "requires an explicit stride"), + ([0, 0], [8, 0], [16, 32], 4, "destination row16 origin"), + ([0, 0], [16, 8], [16, 32], 4, "destination C0 origin"), + ([8, 0], [16, 0], [16, 32], 4, "source row16 origin"), + ([0, 0], [48, 0], [32, 32], 4, "destination region exceeds"), + ([48, 0], [16, 0], [32, 32], 4, "source region exceeds"), + ([0, 0], [16, 0], [24, 32], 4, "source row16 extent"), + ], +) +def test_partial_l0_copy_rejects_misalignment_and_offset_overflow(src_min, dst_min, extent, stride, diagnostic): + src = tirx.decl_buffer((64, 64), "bfloat16", name="src", scope="shared.l1") + dst = tirx.decl_buffer((64, 64), "bfloat16", name="dst", scope="shared.l0a") + with pytest.raises(tvm.error.InternalError, match=diagnostic): + _lower( + src, + dst, + _region(src, extent, src_min), + _region(dst, extent, dst_min), + layouts={src: make_ascend_major_k_layout(src), dst: make_ascend_major_k_layout(dst)}, + stride=stride, + ) + + +@pytest.mark.parametrize( + "v_layout,oversized,full_m", + [(False, False, False), (True, False, False), (False, True, False), (False, False, True)], + ids=["q-compact-m", "v-compact-n", "oversized", "full-m-partial-drain"], +) +def test_explicit_stride_controls_offsets_and_keeps_plane_capacity(v_layout, oversized, full_m): + src = tirx.decl_buffer((2, 64, 64), "bfloat16", name="src", scope="shared.l1") + a = tirx.decl_buffer((2, 64, 64), "bfloat16", name="a", scope="shared.l0a") + b = tirx.decl_buffer((2, 64, 64), "bfloat16", name="b", scope="shared.l0b") + c = tirx.decl_buffer((2, 64, 64), "float32", name="c", scope="shared.l0c") + ub = tirx.decl_buffer((16, 64), "float32", name="ub", scope="shared") + dst = b if v_layout else a + m, n = (64, 24) if v_layout else (24, 64) + if full_m: + m = 64 + pitch = ((n if v_layout else m) + 15) // 16 + statements = [] + for part in range(2): + origin = [1, part * 32, 0] if v_layout else [1, 0, part * 32] + region = [1, 64 if oversized or full_m else 32, 32] + statements.append(tirx.Evaluate(T.copy(_region(src, region, origin), _region(dst, region, origin), stride=pitch))) + if not v_layout: + statements.append(tirx.Evaluate(T.copy(_region(c, [1, 16, 64], [1, 0, 0]), ub, stride=pitch))) + mod = kernel( + tirx.SeqStmt(statements), + buffers=[src, a, b, c, ub], + annotations={ + "layout_map": { + src: make_ascend_major_k_layout(src), + a: make_ascend_major_k_layout(a), + b: (make_ascend_major_mn_layout if v_layout else make_ascend_major_k_layout)(b), + c: make_ascend_l0c_layout(c), + }, + }, + ) + with tvm.target.Target("ascend"): + if oversized: + with pytest.raises(tvm.error.InternalError, match="stride"): + transform.AscendLowerTileOp()(mod) + return + after = 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]) + _equal([load.args[1].args[2] for load in loads], [4096, 4096]) + assert math.prod(int(x) for x in allocated_buffer(after, dst.name).shape) == 2 * 64 * 64 + if not v_layout: + (drain,) = calls(after, "tl.ascend_copy_matrix_cc_to_ub") + _equal([drain.args[6]], [pitch * 16]) # Only 16 rows are drained; stride describes the larger tile. + + +def test_explicit_dynamic_stride_maps_each_iteration_independently(): + src = tirx.decl_buffer((64, 64), "bfloat16", name="src", scope="shared.l1") + dst = tirx.decl_buffer((64, 64), "bfloat16", name="dst", scope="shared.l0a") + i = tirx.Var("i", "int32") + after = _lower( + src, + dst, + _region(src, [32 + 16 * i, 32], [0, 32]), + _region(dst, [32 + 16 * i, 32], [0, 32]), + layouts={src: make_ascend_major_k_layout(src), dst: make_ascend_major_k_layout(dst)}, + stride=2 + i, + wrap=lambda body: tirx.For(i, 0, 3, tirx.ForKind.SERIAL, body), + ) + (load,) = calls(after, "tl.ascend_load_cbuf_to_ca") + _equal([load.args[7], load.args[0].args[2]], [2 + i, 2 * (2 + i) * 256]) + + +@pytest.mark.parametrize("stride,error", [(0, ValueError), (-1, ValueError), (True, TypeError), (1.5, TypeError)]) +def test_copy_stride_rejects_invalid_values(stride, error): + src = tirx.decl_buffer((64, 64), "bfloat16", scope="shared.l1") + dst = tirx.decl_buffer((64, 64), "bfloat16", scope="shared.l0a") + with pytest.raises(error, match="stride"): + T.copy(src, dst, stride=stride) + + +def test_copy_stride_is_only_supported_on_l0_paths(): + src = tirx.decl_buffer((64, 64), "bfloat16") + dst = tirx.decl_buffer((64, 64), "bfloat16", scope="shared.l1") + with pytest.raises(ValueError, match="stride is supported only"): + T.copy(src, dst, stride=4) def test_pad_value_is_set_before_the_padded_dma(): diff --git a/tilelang/ascend/language/copy_op.py b/tilelang/ascend/language/copy_op.py index 25e6e42ad1..7129e931fb 100644 --- a/tilelang/ascend/language/copy_op.py +++ b/tilelang/ascend/language/copy_op.py @@ -233,6 +233,7 @@ def copy( # noqa: A001 *, coalesced_width: int | None = None, transpose: bool = False, + stride: int | tirx.PrimExpr | None = None, l2_cache_ctrl: int | str | None = None, unit_flag_ctrl: int | tirx.PrimExpr | None = None, sub_blockid: int | tirx.PrimExpr | None = None, @@ -244,6 +245,14 @@ def copy( # noqa: A001 """Copy data between memory regions, with Ascend DMA lowering hints. Uses the common region handling with Ascend-specific whole-buffer checks. + + L1-to-L0A/L0B copies default to a compact tile whose pitch comes from + the destination region. When assembling a tile with multiple copies, + pass ``stride`` on each copy to specify the full tile's fractal pitch. + Nonzero matrix origins require ``stride`` and complete, aligned fractals + (including any transpose groups). The allocation and version spacing + remain unchanged. Stride is explicit and independent of auto-scheduling; + the caller must match the consuming GEMM's effective M/N. Whole buffers must have equal element counts, but their shapes may differ for Ascend format conversion and transpose paths (for example, copying ``[K, M]`` into ``[M, K]``). The extra keywords steer how the Ascend backend @@ -267,6 +276,12 @@ def copy( # noqa: A001 memory access. Defaults to None. transpose (bool, keyword-only): Ascend GM→L1 only. Emit dn2nz (which transposes the N/D mapping) instead of nd2nz. Defaults to False. + stride (int | PrimExpr | None, keyword-only): L0 matrix pitch in + 16-element M/N fractal blocks, not bytes or an inter-block gap. + For L1→L0A/L0B, overrides the destination pitch; for L0C→UB/GM, + overrides the source pitch. Use ``ceildiv(M, 16)`` for A/C and + ``ceildiv(N, 16)`` for B, including MajorMN transposed loads. + May vary at runtime. ``None`` keeps the compact region pitch. l2_cache_ctrl (Optional[Union[int, str]], keyword-only): L2 cache control for the GM↔L1 and UB↔GM DMA paths. Accepts the raw integer or a cache-policy name (case-insensitive, suffix-matched), e.g. @@ -309,6 +324,26 @@ def copy( # noqa: A001 ann = _normalize_annotations(annotations) + if stride is not None and "stride" not in ann: + ann["stride"] = stride + if "stride" in ann: + value = ann["stride"] + if isinstance(value, bool) or not isinstance(value, (int, tirx.PrimExpr)): + raise TypeError("Ascend T.copy stride must be an integer or integer PrimExpr") + if isinstance(value, tirx.PrimExpr) and not str(value.dtype).startswith(("int", "uint")): + raise TypeError("Ascend T.copy stride must be an integer or integer PrimExpr") + if isinstance(value, (int, tirx.IntImm)) and int(value) <= 0: + raise ValueError("Ascend T.copy stride must be positive") + src_buf = src if isinstance(src, tirx.Buffer) else src.buffer + dst_buf = dst if isinstance(dst, tirx.Buffer) else dst.buffer + if not ( + src_buf.scope() == "shared.l1" + and dst_buf.scope() in ("shared.l0a", "shared.l0b") + or src_buf.scope() == "shared.l0c" + and dst_buf.scope() in ("shared", "shared.dyn", "global") + ): + raise ValueError("Ascend T.copy stride is supported only for L1→L0A/L0B and L0C→UB/GM") + # Ascend GM→L1: use dn2nz (transpose N/D mapping) instead of nd2nz if transpose and "transpose" not in ann: ann["transpose"] = tirx.IntImm("int32", 1) From f37d0a4aecdca5e2539de57961d643427a2036a9 Mon Sep 17 00:00:00 2001 From: Denver Jin Date: Tue, 13 Oct 2026 17:29:25 +0800 Subject: [PATCH 5/6] [Ascend] Remove standalone partial-copy guide --- .../developer_guide/ascend_partial_l0_copy.md | 113 ------------------ 1 file changed, 113 deletions(-) delete mode 100644 docs/developer_guide/ascend_partial_l0_copy.md diff --git a/docs/developer_guide/ascend_partial_l0_copy.md b/docs/developer_guide/ascend_partial_l0_copy.md deleted file mode 100644 index 938a5f2b97..0000000000 --- a/docs/developer_guide/ascend_partial_l0_copy.md +++ /dev/null @@ -1,113 +0,0 @@ -# Ascend：显式 copy stride 与部分 L0 装载 - -## 前端合约 - -`tilelang.ascend.language.copy` 增加可选的 `stride` 参数，表示 L0 矩阵的 -pitch，单位为 M/N 方向的 **16 元素分形数**，不是字节或分形间隙。 - -| 路径 | `stride` 控制 | 与普通 GEMM 配合时的值 | -|---|---|---| -| L1→L0A | 目标 pitch | `ceildiv(M, 16)` | -| L1→L0B | 目标 pitch，包括 MajorMN 转置装载 | `ceildiv(N, 16)` | -| L0C→UB/GM | 源 pitch | `ceildiv(M, 16)` | - -不传 `stride` 保留既有 compact region 行为：一次 copy 搬运一个完整 GEMM -操作数时，从这次 region 推导 pitch。多次 copy 拼接一个操作数时，每次显式 -传入完整 tile 的 pitch。非零矩阵目标起点要求显式 `stride`。 -其他搬运路径不接受此参数。 - -编译器不查找 copy 对应哪个 GEMM，也不要求同一 buffer 的所有 GEMM 使用相同 -pitch。调用者保证 copy 的 stride 与本次 GEMM 的有效 M/N 一致。 -`stride` 可以是整数表达式，随循环或分支改变；在 copy 执行位置求值。 -改变 pitch 不会重排旧数据，调用者必须重新准备相应输入和累加器。 - -MajorK/MajorMN 仍由既有 layout inference 决定，分配容量及 leading plane / -multibuffer slot 间距保持 allocation layout。显式 stride 仅改变 plane 内的寻址。 -MX scale-factor handle 保持专用搬运合约，不在此参数的支持路径中。 - -## 普通装载与尾块 - -```python -# a 的 allocation 为 [64,64]，本次 GEMM 的有效 M=24。 -T.copy(a_l1[:16, :], a[:16, :], stride=2) -T.copy(a_l1[16:32, :], a[16:32, :], stride=2) -T.gemm(a[:24, :], b, acc[:24, :], transpose_B=True, clear_accum=True) -# 完整输出 region 已经提供本次 compact pitch，无需标注。 -T.copy(acc[:24, :], out[:24, :]) -# 若只输出前 16 行，仍需按完整结果的 pitch 读取。 -T.copy(acc[:16, :], first_rows, stride=2) -``` - -M=64 时，两次各搬 32 行的 copy 都传 `stride=4`。动态尾块可以传 -`stride=T.ceildiv(effective_m, 16)`，region 选择包含 padding 的完整分形。 - -对 BF16 MajorK，逻辑坐标映射为 `[k//16, m//16, m%16, k%16]`。 -分形线性位置为 `outer0 * stride + outer1`，每个分形占 512 B。 -例如 allocation 为 [2,64,64]，选择 plane=1、stride=2、目标起点 [0,32]： - -```text -plane base = 64 * 64 个元素 -plane 内偏移 = (2 * 2 + 0) * 256 个元素 -``` - -## V 的转置装载 - -```python -# v_l1 / v 的逻辑形状是 [K,N]，此处 K=64。 -T.copy(v_l1[:32, :N_pad], v[:32, :N_pad], stride=T.ceildiv(N, 16)) -T.copy(v_l1[32:64, :N_pad], v[32:64, :N_pad], stride=T.ceildiv(N, 16)) -T.gemm(a, v[:, :N], acc[:, :N], clear_accum=True) -``` - -GEMM 为 V 推导 MajorMN，目标偏移按 MajorMN 和显式 stride 计算。 -硬件转置位仍是 `user_transpose XOR major_mismatch`。 -令 C0=256/dtype.bits，源 region 的分形步数为 (k_step,m_step)： - -```text -普通实际写入 = (k_step, m_step) -转置实际写入 = (m_step*16/C0, k_step*C0/16) -``` - -copy 检查完整转置组、实际写入范围、目标起点和容量。 -非零目标起点要求源和目标的起点、extent 均分形对齐；零起点保留既有 padded load。 -实际写入必须位于声明的目标 region 内，且不能越过 stride 或 allocation。 -空 copy guard 保持不变。 - -## Lowering 与同步 - -前端将 `stride` 存入 Ascend copy annotation；原有 IR 变量替换同时处理该表达式， -包括流水化迭代的变量替换。L1→L0 lowering 将它用于指令的 dst_stride， -并通过内部 access-pointer annotation 将它传给地址 remap。 -LowerTileOp 保留完整 Major 映射和 allocation 间距，只修正 matrix outer0 的 stride。 -L0C→UB/GM 将分形数乘以 16，得到 Fixpipe 的源行距。 -GEMM lowering、布局推导不新增 pitch 传播或匹配过程。 - -AutoSchedule 开关不影响 stride 语义。开启时仍按 RAW/WAR/WAW 关系加同步： -每条指令的 L0 访问先表示为 plane 内的字节 Range，附带 `(repeat, stride)`， -分别表示重复次数和重复段起点间的字节距离。不同 pitch 也直接比较实际地址范围； -仅在无法确定寻址方式时退化为整个 plane。保留 leading plane / slot 分离与条件、 -循环上下文。`MemoryAccessDetector` 在收集 Task 时逐次生成 `MemoryAccess`， -经 `TaskAccessInfo` 存入 `TaskNode`，依赖分析直接消费这些记录。 -同一 buffer 的多个 copy 分别保留各自的 pitch 和访问范围，不先合并为一个矩形。 -Task 内循环的变量范围随访问记录保存，并在比较两侧访问时分别实例化，保留 -地址与动态 pitch 的关联。Task 克隆和变量替换同时处理这些记录；无需额外扫描 -Task 或维护一份访问缓存。 -显式 conflict hints 仍匹配原始逻辑 region，不由编译器自动生成。 -这只是访问冲突判断，不给 copy 选择 GEMM 或推导 stride。 -关闭 AutoSchedule 时，由调用者设置各流水线之间的同步。 - -多个 writer 允许顺序覆盖，不新增 writer 并集覆盖证明。 -物理 lifetime 分析将 compact L0 storage 标记为 opaque，跳过跨 allocation 的重叠 -复用证明，并保留所选 plane 的 may-access；其他普通 operand 保留 must-write 信息。 - -## 支持边界与回归 - -沿用 GEMM 和 L0C drain 的矩阵零起点限制。支持 L1→L0A/L0B 的分形对齐非零 -目标起点、普通及显式/隐式转置、动态 M/N 尾块、K padding 和有效 K。 -不插入 repack，不自动补齐非对齐 offset copy，不校验 copy 与 GEMM 的布局一致性。 -FP4 packed-address 不在首批设备验证范围内。 - -主机回归覆盖默认 compact pitch、显式静态/动态 stride、普通/V 目标地址、 -leading plane、部分 L0C drain、实际边界、不同 pitch 的同步及 allocation 复用。 -NPU 回归覆盖两/三版本 ring 的部分装载、动态尾块、顺序覆盖，并覆盖关闭 -AutoSchedule、手动同步、分支内不同 GEMM pitch 共用同一 buffer 的情形。 From 9ec2c59080ac421cb142397b256df4865e702e62 Mon Sep 17 00:00:00 2001 From: Denver Jin Date: Wed, 14 Oct 2026 12:04:19 +0800 Subject: [PATCH 6/6] [Ascend] Simplify L0 copy dependency analysis --- .../auto_schedule/dependency_analysis.cc | 94 +++---- .../auto_schedule/dependency_analysis.h | 11 +- .../transform/auto_schedule/ir_structure.cc | 5 +- .../transform/auto_schedule/ir_structure.h | 109 ++++---- .../transform/auto_schedule/memory_detector.h | 242 +++++------------- .../transform/auto_schedule/task_analysis.h | 31 ++- src/ascend/transform/resolve_core.cc | 3 +- .../auto_schedule/test_sync_dependencies.py | 67 +++-- 8 files changed, 221 insertions(+), 341 deletions(-) diff --git a/src/ascend/transform/auto_schedule/dependency_analysis.cc b/src/ascend/transform/auto_schedule/dependency_analysis.cc index 1ec3833ffa..ec16c4300b 100644 --- a/src/ascend/transform/auto_schedule/dependency_analysis.cc +++ b/src/ascend/transform/auto_schedule/dependency_analysis.cc @@ -38,6 +38,7 @@ #include #include +#include "ascend/layout/ascend_layouts.h" #include "support/check.h" #include "transform/common/attr.h" #include "transform/common/constr_visitor.h" @@ -112,13 +113,13 @@ std::vector StorageIterationNest(ControlNode *loop, return {controls.begin() + first, controls.end()}; } -bool RegionsMayConflict(const ConstrSet &a_ctx, const MemoryAccess &a_access, - const ConstrSet &b_ctx, const MemoryAccess &b_access, +bool RegionsMayConflict(const ConstrSet &a_ctx, const BufferRegion &a_region, + const ConstrSet &b_ctx, const BufferRegion &b_region, ControlNode *loop, int offset, size_t num_storage_owners, - const ConflictHintList &root_conflicts) { - const BufferRegion &a_region = a_access.logical; - const BufferRegion &b_region = b_access.logical; + const ConflictHintList &root_conflicts, + Optional a_pitch, + Optional b_pitch) { const bool cross = offset != 0; if (cross && loop == nullptr) return false; @@ -152,8 +153,8 @@ bool RegionsMayConflict(const ConstrSet &a_ctx, const MemoryAccess &a_access, if (!a_region->buffer.same_as(b_region->buffer)) return true; - const Array &a_rng = a_access.bounds; - const Array &b_rng = b_access.bounds; + Array a_rng = ExpandAscendFractalMayAccess(a_region)->region; + Array b_rng = ExpandAscendFractalMayAccess(b_region)->region; if (a_rng.size() != b_rng.size()) return true; // shape mismatch — be safe const size_t ndim = a_rng.size(); @@ -210,59 +211,27 @@ bool RegionsMayConflict(const ConstrSet &a_ctx, const MemoryAccess &a_access, } loop_ctx.Merge(a_local_ctx).Merge(b_local_ctx).Populate(ana); - // Snapshot bounds, repeat, and stride together on each side. Identical - // mutable reads at two access points need not have the same value. FreshenMutableReads freshen_a(FreshenMutableReads::Mode::kSnapshot); FreshenMutableReads freshen_b(FreshenMutableReads::Mode::kSnapshot); - // Internal task loops are independent access choices on the two sides. - // Rename them together with their bounds and pitch before adding domains. - auto bind_domains = [&](const MemoryAccess &access, Map *sub, - FreshenMutableReads &freshen, const char *suffix) { - for (const auto &[var, range] : access.domains) - sub->Set(var, var.copy_with_suffix(suffix)); - for (const auto &[var, range] : access.domains) { - PrimExpr index = freshen(Substitute(var, *sub)); - PrimExpr min = freshen(Substitute(range->min, *sub)); - PrimExpr extent = freshen(Substitute(range->extent, *sub)); - ana.EnterConstraint(index >= min && - cast(DataType::Int(64), index) < - cast(DataType::Int(64), min) + - cast(DataType::Int(64), extent)); + size_t compared_dims = ndim; + const Buffer &buffer = a_region->buffer; + if (IsL0ABuffer(buffer) || IsL0BBuffer(buffer) || IsL0CBuffer(buffer)) { + bool same_pitch = false; + if (a_pitch.defined() && b_pitch.defined()) { + PrimExpr equal = freshen_a(Substitute(a_pitch.value(), a_sub)) == + freshen_b(Substitute(b_pitch.value(), b_sub)); + same_pitch = ana.CanProve(equal) || ana.z3_prover.CanProve(equal); } - }; - bind_domains(a_access, &a_sub, freshen_a, "_access_p"); - bind_domains(b_access, &b_sub, freshen_b, "_access_c"); - PrimExpr a_repeat = freshen_a(Substitute(a_access.repeat, a_sub)); - PrimExpr b_repeat = freshen_b(Substitute(b_access.repeat, b_sub)); - - // Each side chooses an independent interval from its repeated last axis. - // Proving disjointness for every permitted pair is equivalent to proving - // that no pair overlaps; no enumeration or explicit quantifier is needed. - auto repeat_min = [&](PrimExpr min, const PrimExpr &repeat, - const PrimExpr &stride, const char *name) -> PrimExpr { - if (is_one(ana.Simplify(repeat))) - return min; - // Widen before multiplying so int32 counts/pitches do not truncate the - // address calculation. Zero and negative strides need no special case. - Var index(name, DataType::Int(64)); - ana.EnterConstraint(index >= 0 && index < cast(index.dtype(), repeat)); - return cast(index.dtype(), min) + index * cast(index.dtype(), stride); - }; - - PrimExpr any_dim_disjoint = (a_repeat <= 0) || (b_repeat <= 0); - for (size_t i = 0; i < ndim; ++i) { + // L0 matrix regions may overlap unless their pitches are provably equal. + if (!same_pitch) + compared_dims = ndim > 2 ? ndim - 2 : 0; + } + PrimExpr any_dim_disjoint = Bool(false); + for (size_t i = 0; i < compared_dims; ++i) { PrimExpr a_min = freshen_a(Substitute(a_rng[i]->min, a_sub)); PrimExpr a_extent = freshen_a(Substitute(a_rng[i]->extent, a_sub)); PrimExpr b_min = freshen_b(Substitute(b_rng[i]->min, b_sub)); PrimExpr b_extent = freshen_b(Substitute(b_rng[i]->extent, b_sub)); - if (i + 1 == ndim) { - a_min = repeat_min(a_min, a_repeat, - freshen_a(Substitute(a_access.stride, a_sub)), - "region_repeat_p"); - b_min = repeat_min(b_min, b_repeat, - freshen_b(Substitute(b_access.stride, b_sub)), - "region_repeat_c"); - } // A half-open endpoint can exceed the int32 index range by one even // for an ordinary single-element access. PrimExpr a_end = @@ -302,7 +271,7 @@ AnalyzeDependencies(std::vector nodes, ControlNode *loop, struct RegionAccess { TaskNode *task; - MemoryAccess region; + BufferRegion region; }; struct NodeAccesses { std::vector reads; @@ -313,9 +282,9 @@ AnalyzeDependencies(std::vector nodes, ControlNode *loop, std::vector tasks; CollectAllTaskNodes(nodes[i], tasks); for (TaskNode *task : tasks) { - for (const MemoryAccess ®ion : task->GetReadAccesses()) + for (const BufferRegion ®ion : task->GetReadRegions()) accesses[i].reads.push_back({task, region}); - for (const MemoryAccess ®ion : task->GetWriteAccesses()) + for (const BufferRegion ®ion : task->GetWriteRegions()) accesses[i].writes.push_back({task, region}); } } @@ -420,8 +389,8 @@ AnalyzeDependencies(std::vector nodes, ControlNode *loop, const std::vector &rhs_accesses) { for (const RegionAccess &lhs : lhs_accesses) { for (const RegionAccess &rhs : rhs_accesses) { - const Var &lhs_storage = lhs.region.logical->buffer->data; - const Var &rhs_storage = rhs.region.logical->buffer->data; + const Var &lhs_storage = lhs.region->buffer->data; + const Var &rhs_storage = rhs.region->buffer->data; ffi::Optional storage; if (lhs_storage.same_as(rhs_storage)) storage = lhs_storage; @@ -437,8 +406,9 @@ AnalyzeDependencies(std::vector nodes, ControlNode *loop, for (int d = (i < j ? 0 : 1); d <= manual_versions; ++d) { if (RegionsMayConflict(lhs.task->outer_ctx, lhs.region, rhs.task->outer_ctx, rhs.region, loop, d, - /*num_storage_owners=*/0, - root_conflicts)) { + /*num_storage_owners=*/0, root_conflicts, + lhs.task->GetL0Pitch(lhs_storage), + rhs.task->GetL0Pitch(rhs_storage))) { add_pair(storage, d, lhs.task, rhs.task); break; } @@ -453,7 +423,9 @@ AnalyzeDependencies(std::vector nodes, ControlNode *loop, if (RegionsMayConflict(lhs.task->outer_ctx, lhs.region, rhs.task->outer_ctx, rhs.region, loop, distance, get_num_storage_owners(storage), - root_conflicts)) { + 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 3579a84662..9c08634113 100644 --- a/src/ascend/transform/auto_schedule/dependency_analysis.h +++ b/src/ascend/transform/auto_schedule/dependency_analysis.h @@ -109,8 +109,7 @@ AnalyzeDependencies(std::vector nodes, const ConflictHintList &root_conflicts = {}); // Check whether two task regions may overlap under their respective symbolic -// contexts, including independent repeats of their last conflict axis. Only a -// proof of disjointness removes a conflict; solver unknowns remain conflicts. +// contexts. Only a proof of disjointness removes a conflict. // `offset == 0` compares accesses in the same iteration. A positive // offset constrains the current consumer loop coordinate to equal the producer // coordinate plus that offset. A negative offset denotes an automatic @@ -121,11 +120,13 @@ AnalyzeDependencies(std::vector nodes, // an identically-shaped iteration in another owner can still be a distinct // epoch. Shared by scheduling-time IRStructure analysis and the scheduled-TIR // synchronization pass. -bool RegionsMayConflict(const ConstrSet &a_ctx, const MemoryAccess &a_access, - const ConstrSet &b_ctx, const MemoryAccess &b_access, +bool RegionsMayConflict(const ConstrSet &a_ctx, const BufferRegion &a_region, + const ConstrSet &b_ctx, const BufferRegion &b_region, ControlNode *loop, int offset, size_t num_storage_owners = 0, - const ConflictHintList &root_conflicts = {}); + const ConflictHintList &root_conflicts = {}, + Optional a_pitch = std::nullopt, + Optional b_pitch = std::nullopt); } // namespace ascend } // namespace tl diff --git a/src/ascend/transform/auto_schedule/ir_structure.cc b/src/ascend/transform/auto_schedule/ir_structure.cc index 92fbb10c08..e91280062a 100644 --- a/src/ascend/transform/auto_schedule/ir_structure.cc +++ b/src/ascend/transform/auto_schedule/ir_structure.cc @@ -297,8 +297,9 @@ std::shared_ptr TaskNode::Clone() const { if (IsPerCoreTask()) new_task->MarkPerCoreTask(); // Copy memory access regions - new_task->SetReadAccesses(GetReadAccesses()); - new_task->SetWriteAccesses(GetWriteAccesses()); + new_task->SetReadRegions(GetReadRegions()); + new_task->SetWriteRegions(GetWriteRegions()); + new_task->SetL0Pitches(GetL0Pitches()); new_task->SetReadVars(GetReadVars()); new_task->SetWriteVars(GetWriteVars()); new_task->outer_ctx = outer_ctx; diff --git a/src/ascend/transform/auto_schedule/ir_structure.h b/src/ascend/transform/auto_schedule/ir_structure.h index 6c59d571a5..36ac7ebfb1 100644 --- a/src/ascend/transform/auto_schedule/ir_structure.h +++ b/src/ascend/transform/auto_schedule/ir_structure.h @@ -25,7 +25,6 @@ #include "../ascend_pipe.h" #include "../buffer_version.h" #include "../core_mask.h" -#include "./memory_detector.h" #include "op/utils.h" namespace tvm { @@ -111,22 +110,6 @@ inline bool RegionsEqual(const Region &a, const Region &b) { return true; } -inline bool MemoryAccessesEqual(const MemoryAccess &a, const MemoryAccess &b) { - if (a.domains.size() != b.domains.size()) - return false; - for (size_t i = 0; i < a.domains.size(); ++i) { - if (!a.domains[i].first.same_as(b.domains[i].first) || - !RegionsEqual({a.domains[i].second}, {b.domains[i].second})) - return false; - } - arith::Analyzer analyzer; - return a.logical->buffer.same_as(b.logical->buffer) && - RegionsEqual(a.logical->region, b.logical->region) && - RegionsEqual(a.bounds, b.bounds) && - analyzer.CanProveEqual(a.repeat, b.repeat) && - analyzer.CanProveEqual(a.stride, b.stride); -} - // Return true only when every point in `inner` is provably contained in // `outer`. This is stricter than overlap and is suitable for checking whether // one task's declared write footprint fully defines another task's read. @@ -392,8 +375,7 @@ class IRStructure { // only TaskNode stores a settable mask. virtual uint16_t GetHbmMask() const = 0; - // Logical projections for storage identity and core analysis. Dependency - // analysis consumes the leaf TaskNode's complete access records. + // Logical access regions. Physical footprints are derived from current IR. virtual std::vector GetReadRegions() const = 0; virtual std::vector GetWriteRegions() const = 0; @@ -547,19 +529,12 @@ class TaskNode : public IRStructure { // A copy that stores to GM in this leaf reads the atomic-mode register. bool ReadsAtomicRegister() const override { return reads_atomic_; } - const std::vector &GetReadAccesses() const { - return read_regions_; - } - const std::vector &GetWriteAccesses() const { - return write_regions_; - } - - // Logical projections retain the union's individual records. + // Logical access regions. Physical footprints are derived from current IR. std::vector GetReadRegions() const override { - return LogicalRegions(read_regions_); + return read_regions_; } std::vector GetWriteRegions() const override { - return LogicalRegions(write_regions_); + return write_regions_; } std::vector GetReadVars() const override { return read_vars_; } std::vector GetWriteVars() const override { return write_vars_; } @@ -575,12 +550,22 @@ class TaskNode : public IRStructure { void AddSpecialWrite(SpecialRegister reg) { special_write_mask_ |= static_cast(reg); } - void SetReadAccesses(const std::vector ®ions) { + void SetReadRegions(const std::vector ®ions) { read_regions_ = regions; } - void SetWriteAccesses(const std::vector ®ions) { + void SetWriteRegions(const std::vector ®ions) { write_regions_ = regions; } + const Map> &GetL0Pitches() const { + return l0_pitches_; + } + Optional GetL0Pitch(const Var &storage) const { + auto it = l0_pitches_.find(storage); + return it == l0_pitches_.end() ? Optional() : (*it).second; + } + void SetL0Pitches(const Map> &pitches) { + l0_pitches_ = pitches; + } void SetReadVars(const std::vector &vars) { read_vars_ = vars; } void SetWriteVars(const std::vector &vars) { write_vars_ = vars; } void SubstituteVar(const Var &old_var, const Var &new_var) override { @@ -593,9 +578,9 @@ class TaskNode : public IRStructure { }; std::unordered_map buffer_remap; - auto substitute_regions = [&](std::vector ®ions) { - for (MemoryAccess &access : regions) { - Buffer buffer = access.logical->buffer; + auto substitute_regions = [&](std::vector ®ions) { + for (auto ®ion : regions) { + Buffer buffer = region->buffer; auto buffer_it = buffer_remap.find(buffer); if (buffer_it != buffer_remap.end()) { buffer = buffer_it->second; @@ -605,24 +590,22 @@ class TaskNode : public IRStructure { buffer_remap.emplace(buffer, remapped); buffer = std::move(remapped); } - access.logical = - BufferRegion(std::move(buffer), - Substitute(access.logical->region, substitution)); - access.bounds = Substitute(access.bounds, substitution); - access.repeat = Substitute(access.repeat, substitution); - access.stride = Substitute(access.stride, substitution); - for (auto &[var, range] : access.domains) { - if (var.same_as(old_var)) - var = new_var; - range = Range::FromMinExtent(Substitute(range->min, substitution), - Substitute(range->extent, substitution)); - } + region = BufferRegion(std::move(buffer), + Substitute(region->region, substitution)); } }; substitute_vars(read_vars_); substitute_vars(write_vars_); substitute_regions(read_regions_); substitute_regions(write_regions_); + Map> pitches; + for (const auto &[storage, pitch] : l0_pitches_) { + pitches.Set(storage.same_as(old_var) ? new_var : storage, + pitch.defined() ? Optional( + Substitute(pitch.value(), substitution)) + : std::nullopt); + } + l0_pitches_ = std::move(pitches); outer_ctx = outer_ctx.Substitute(substitution); stmt = Substitute(stmt, substitution); SubstituteVarInGuards(old_var, new_var); @@ -637,21 +620,27 @@ class TaskNode : public IRStructure { // Clone method std::shared_ptr Clone() const override; - // Keep the existing deduplication, comparing the complete access record. - void AddReadAccess(const MemoryAccess &access) { - for (const MemoryAccess &existing : read_regions_) { - if (MemoryAccessesEqual(existing, access)) - return; + // Helper methods to add regions (for incremental analysis) + void AddReadRegion(const BufferRegion ®ion) { + // Check for duplicate regions + for (const auto &existing : read_regions_) { + if (existing->buffer.same_as(region->buffer) && + RegionsEqual(existing->region, region->region)) { + return; // Region already exists + } } - read_regions_.push_back(access); + read_regions_.push_back(region); } - void AddWriteAccess(const MemoryAccess &access) { - for (const MemoryAccess &existing : write_regions_) { - if (MemoryAccessesEqual(existing, access)) - return; + void AddWriteRegion(const BufferRegion ®ion) { + // Check for duplicate regions + for (const auto &existing : write_regions_) { + if (existing->buffer.same_as(region->buffer) && + RegionsEqual(existing->region, region->region)) { + return; // Region already exists + } } - write_regions_.push_back(access); + write_regions_.push_back(region); } void AddReadVar(const Var &var) { @@ -683,8 +672,10 @@ class TaskNode : public IRStructure { // inferred or explicitly marked candidate. bool is_per_core_task_{false}; - std::vector read_regions_; - std::vector write_regions_; + // Logical access regions. Physical footprints are derived from current IR. + std::vector read_regions_; + std::vector write_regions_; + Map> l0_pitches_; std::vector read_vars_; std::vector write_vars_; diff --git a/src/ascend/transform/auto_schedule/memory_detector.h b/src/ascend/transform/auto_schedule/memory_detector.h index 67efd27c2f..85fa4ad945 100644 --- a/src/ascend/transform/auto_schedule/memory_detector.h +++ b/src/ascend/transform/auto_schedule/memory_detector.h @@ -40,152 +40,51 @@ namespace ascend { using namespace tirx; -// One detected access. A list of these records denotes a union, including -// when several records refer to the same logical buffer. -struct MemoryAccess { - explicit MemoryAccess(BufferRegion logical) - : logical(std::move(logical)), bounds(this->logical->region), - repeat(Integer(1)), stride(Integer(0)) {} - - // Buffer identity, conflict hints, and logical cost accounting. - BufferRegion logical; - // Conflict coordinates, independent of the logical buffer's rank. All - // accesses to one buffer must use the same coordinate convention. - Region bounds; - // Repeat the last axis at [min + i * stride, min + i * stride + extent), - // 0 <= i < repeat. Ordinary rectangular accesses use repeat=1, stride=0. - // A non-positive repeat or axis extent denotes an empty access. Stride may - // be zero or negative. - PrimExpr repeat; - PrimExpr stride; - // Loops inside an atomic task. Keep address/pitch correlations until each - // side of an overlap query chooses its own iteration. - std::vector> domains; -}; - -inline std::vector -LogicalRegions(const std::vector &accesses) { - std::vector regions; - regions.reserve(accesses.size()); - for (const MemoryAccess &access : accesses) - regions.push_back(access.logical); - return regions; -} - -using BufferMemoryAccesses = - std::unordered_map; - -// Unknown L0 addressing covers a whole selected plane. Other storage keeps -// its normal coordinate system, with instruction padding for L1 fractals. -inline MemoryAccess MakeMemoryAccess(const BufferRegion ®ion) { - MemoryAccess access(region); - access.bounds.clear(); - const Buffer &buffer = region->buffer; - if (IsL0ABuffer(buffer) || IsL0BBuffer(buffer) || IsL0CBuffer(buffer)) { - for (size_t i = 0; i + 2 < region->region.size(); ++i) - access.bounds.push_back(region->region[i]); - int alignment = - IsL0CBuffer(buffer) ? 16 : std::max(16, AscendC0(buffer->dtype.bits())); - PrimExpr bytes = make_const(DataType::Int(64), buffer->dtype.bits()); - for (size_t i = buffer->shape.size() - 2; i < buffer->shape.size(); ++i) - bytes *= cast(DataType::Int(64), - floordiv(buffer->shape[i] + alignment - 1, alignment) * - alignment); - access.bounds.push_back( - Range::FromMinExtent(make_zero(DataType::Int(64)), floordiv(bytes, 8))); - } else { - access.bounds = ExpandAscendFractalMayAccess(region)->region; +inline void MergeL0Pitch(Map> *pitches, + const Var &storage, Optional pitch) { + auto it = pitches->find(storage); + if (it != pitches->end()) { + arith::Analyzer analyzer; + if (!(*it).second.defined() || !pitch.defined() || + !analyzer.CanProveEqual((*it).second.value(), pitch.value())) + pitch = std::nullopt; } - return access; + pitches->Set(storage, pitch); } -inline MemoryAccess FractalAccess(const BufferRegion ®ion, - const AscendFractalLayoutInfo &info, - const AscendFractalRegionInfo &fractal, - PrimExpr pitch, PrimExpr count, - PrimExpr width) { - MemoryAccess access = MakeMemoryAccess(region); - // Address arithmetic must not overflow the logical index's int32 type. - auto wide = [](PrimExpr value) { return cast(DataType::Int(64), value); }; - PrimExpr block_bytes = - floordiv(wide(info.row_frac * info.c0) * region->buffer->dtype.bits(), 8); - PrimExpr base = - (wide(fractal.outer0->min) * wide(pitch) + wide(fractal.outer1->min)) * - block_bytes; - access.bounds.Set(access.bounds.size() - 1, - Range::FromMinExtent(base, wide(width) * block_bytes)); - access.repeat = wide(count); - access.stride = wide(pitch) * block_bytes; - return access; -} - -// Describe each instruction using its own fractal geometry and pitch. -inline BufferMemoryAccesses GetCallMemoryAccesses(const Call &call, - const LayoutMap &layouts) { - BufferMemoryAccesses result; +inline Map GetCallL0Pitches(const Call &call, + const LayoutMap &layouts) { + Map result; + auto record = [&](const BufferRegion ®ion, size_t mn_axis, + Optional stride = std::nullopt) { + PrimExpr pitch = + stride.value_or(floordiv(region->region[mn_axis]->extent + 15, 16)); + result.Set(region->buffer->data, pitch); + }; if (IsAscendCopyCall(call.get())) { AscendCopy copy(call->args, call->annotations); if (IsL1Buffer(copy->src) && (IsL0ABuffer(copy->dst) || IsL0BBuffer(copy->dst))) { - auto src_layout = FindLayoutForBuffer(layouts, copy->src); - auto dst_layout = FindLayoutForBuffer(layouts, copy->dst); - AscendFractalLayoutInfo src_info, dst_info; - AscendFractalRegionInfo src, dst; - if (!src_layout.defined() || !dst_layout.defined() || - !TryExtractAscendFractalLayout(src_layout.value(), copy->src, - &src_info) || - !TryExtractAscendFractalLayout(dst_layout.value(), copy->dst, - &dst_info) || - !IsAscendMajorKind(src_info.kind) || - !IsAscendMajorKind(dst_info.kind) || - !TryExtractAscendFractalRegion(src_layout.value(), copy->src_range, - &src) || - !TryExtractAscendFractalRegion(dst_layout.value(), copy->dst_range, - &dst)) - return result; - PrimExpr count = src.outer0->extent; - PrimExpr width = src.outer1->extent; - if ((copy->transpose != 0) ^ (src_info.c0_axis != dst_info.c0_axis)) { - count = floordiv(src.outer1->extent * 16, src_info.c0); - width = floordiv(src.outer0->extent * src_info.c0, 16); - } - result.emplace(copy->dst, - FractalAccess(BufferRegion(copy->dst, copy->dst_range), - dst_info, dst, - copy->stride.value_or(dst.outer1->extent), - count, width)); - } else if (IsL0CBuffer(copy->src)) { - Layout layout = MakeAscendL0CLayout(copy->src); + auto layout = FindLayoutForBuffer(layouts, copy->dst); AscendFractalLayoutInfo info; - AscendFractalRegionInfo fractal; - TryExtractAscendFractalLayout(layout, copy->src, &info); - TryExtractAscendFractalRegion(layout, copy->src_range, &fractal); - result.emplace( - copy->src, - FractalAccess(BufferRegion(copy->src, copy->src_range), info, fractal, - copy->stride.value_or(fractal.outer1->extent), - fractal.outer0->extent, fractal.outer1->extent)); + if (layout.defined() && + TryExtractAscendFractalLayout(layout.value(), copy->dst, &info) && + IsAscendMajorKind(info.kind)) + record(BufferRegion(copy->dst, copy->dst_range), + copy->dst_range.size() - 2 + + (info.row16_axis == LogicalAxis::kCol), + copy->stride); + } else if (IsL0CBuffer(copy->src)) { + record(BufferRegion(copy->src, copy->src_range), + copy->src_range.size() - 2, copy->stride); } } else if (call->op.same_as(Op::Get("tl.tileop.gemm"))) { Gemm gemm(call->args, call->annotations); - if (!IsL0ABuffer(gemm->a_) || !IsL0BBuffer(gemm->b_)) - return result; - auto record = [&](const BufferRegion ®ion, const Layout &layout) { - AscendFractalLayoutInfo info; - AscendFractalRegionInfo fractal; - TryExtractAscendFractalLayout(layout, region->buffer, &info); - TryExtractAscendFractalRegion(layout, region->region, &fractal); - result.emplace(region->buffer, FractalAccess(region, info, fractal, - fractal.outer1->extent, - fractal.outer0->extent, - fractal.outer1->extent)); - }; - record(gemm->aRegion_, gemm->transA_ ? MakeAscendMajorMNLayout(gemm->a_) - : MakeAscendMajorKLayout(gemm->a_)); - record(gemm->bRegion_, gemm->transB_ ? MakeAscendMajorKLayout(gemm->b_) - : MakeAscendMajorMNLayout(gemm->b_)); - record(gemm->cRegion_, MakeAscendL0CLayout(gemm->c_)); + if (IsL0ABuffer(gemm->a_) && IsL0BBuffer(gemm->b_)) { + record(gemm->aRegion_, gemm->a_->shape.size() - (gemm->transA_ ? 1 : 2)); + record(gemm->bRegion_, gemm->b_->shape.size() - (gemm->transB_ ? 2 : 1)); + record(gemm->cRegion_, gemm->c_->shape.size() - 2); + } } return result; } @@ -194,7 +93,6 @@ inline BufferMemoryAccesses GetCallMemoryAccesses(const Call &call, // Adapted from BlockReadWriteDetector in TVM class MemoryAccessDetector : public StmtExprVisitor { public: - // Task accesses carry logical regions and instruction conflict coordinates. // Lifetime analysis separately requests may/must-write coverage. enum class AccessMode { kLogical, kPhysical }; explicit MemoryAccessDetector(AccessMode mode = AccessMode::kLogical, @@ -203,9 +101,8 @@ class MemoryAccessDetector : public StmtExprVisitor { // Analyze a statement and collect read/write regions void Analyze(const Stmt &stmt) { - read_accesses_.clear(); - write_accesses_.clear(); - operation_accesses_.clear(); + l0_pitches_.clear(); + operation_pitches_.clear(); read_buffers_.clear(); write_buffers_.clear(); read_regions_.clear(); @@ -227,9 +124,8 @@ class MemoryAccessDetector : public StmtExprVisitor { } void Analyze(const PrimExpr &expr) { - read_accesses_.clear(); - write_accesses_.clear(); - operation_accesses_.clear(); + l0_pitches_.clear(); + operation_pitches_.clear(); read_buffers_.clear(); write_buffers_.clear(); read_regions_.clear(); @@ -260,11 +156,8 @@ class MemoryAccessDetector : public StmtExprVisitor { return CollectRegions(write_buffers_, write_regions_); } - const std::vector &GetReadAccesses() const { - return read_accesses_; - } - const std::vector &GetWriteAccesses() const { - return write_accesses_; + const Map> &GetL0Pitches() const { + return l0_pitches_; } // Return all variables that are read from @@ -287,32 +180,8 @@ class MemoryAccessDetector : public StmtExprVisitor { private: LayoutMap layouts_; - BufferMemoryAccesses operation_accesses_; - std::vector read_accesses_; - std::vector write_accesses_; - - MemoryAccess DescribeAccess(const Buffer &buffer, - const std::vector ®ion) { - MemoryAccess access = MakeMemoryAccess(MakeBufferRegion(buffer, region)); - auto it = operation_accesses_.find(buffer); - if (it == operation_accesses_.end()) - return access; - BufferRegion logical = access.logical; - access = it->second; - access.logical = std::move(logical); - Array bounds; - for (const Range &range : access.bounds) - bounds.push_back(Range::FromMinExtent(ResolveLets(range->min), - ResolveLets(range->extent))); - access.bounds = std::move(bounds); - access.repeat = ResolveLets(access.repeat); - access.stride = ResolveLets(access.stride); - for (const auto &[var, range] : loop_domains_) - access.domains.emplace_back( - var, Range::FromMinExtent(ResolveLets(range->min), - ResolveLets(range->extent))); - return access; - } + Map operation_pitches_; + Map> l0_pitches_; /*! \brief Iteration range for loop_vars */ std::unordered_map dom_map_; @@ -377,11 +246,21 @@ class MemoryAccessDetector : public StmtExprVisitor { for (const PrimExpr &stride : buffer->strides) visit_mutable_layout(stride); visit_mutable_layout(buffer->elem_offset); - auto &accesses = - buffers == &read_buffers_ ? read_accesses_ : write_accesses_; - accesses.push_back(mode_ == AccessMode::kLogical - ? DescribeAccess(buffer, region) - : MemoryAccess(MakeBufferRegion(buffer, region))); + if (mode_ == AccessMode::kLogical && + (IsL0ABuffer(buffer) || IsL0BBuffer(buffer) || IsL0CBuffer(buffer))) { + Optional pitch = operation_pitches_.Get(buffer->data); + if (pitch.defined()) { + PrimExpr value = ana_.Simplify(ResolveLets(pitch.value())); + if (SideEffect(value) > CallEffectKind::kPure || + UsesVar(value, [&](const VarNode *var) { + return dom_map_.count(var) != 0; + })) + pitch = std::nullopt; + else + pitch = value; + } + MergeL0Pitch(&l0_pitches_, buffer->data, pitch); + } MergeRegion(buffers, regions, buffer, region); } @@ -1152,11 +1031,10 @@ class MemoryAccessDetector : public StmtExprVisitor { if (mode_ == AccessMode::kLogical && (IsAscendCopyCall(op) || op->op.same_as(Op::Get("tl.tileop.gemm")))) { - auto saved = std::move(operation_accesses_); - operation_accesses_ = - GetCallMemoryAccesses(ffi::GetRef(op), layouts_); + auto saved_pitches = std::move(operation_pitches_); + operation_pitches_ = GetCallL0Pitches(ffi::GetRef(op), layouts_); StmtExprVisitor::VisitExpr_(op); - operation_accesses_ = std::move(saved); + operation_pitches_ = std::move(saved_pitches); return; } diff --git a/src/ascend/transform/auto_schedule/task_analysis.h b/src/ascend/transform/auto_schedule/task_analysis.h index fc5b3e1061..294930162e 100644 --- a/src/ascend/transform/auto_schedule/task_analysis.h +++ b/src/ascend/transform/auto_schedule/task_analysis.h @@ -174,8 +174,9 @@ struct TaskResourceUsage { // remain separate from hardware resource classification because AutoSchedule // also adds accesses from surrounding guards and loop expressions to a task. struct TaskAccessInfo { - std::vector read_regions; - std::vector write_regions; + std::vector read_regions; + std::vector write_regions; + Map> l0_pitches; std::vector read_vars; std::vector write_vars; }; @@ -403,8 +404,9 @@ inline TaskAccessInfo AnalyzeTaskAccesses(const Stmt &stmt, MemoryAccessDetector::AccessMode::kLogical, layouts); memory_detector.Analyze(stmt); TaskAccessInfo result; - result.read_regions = memory_detector.GetReadAccesses(); - result.write_regions = memory_detector.GetWriteAccesses(); + result.read_regions = memory_detector.GetReadRegions(); + result.write_regions = memory_detector.GetWriteRegions(); + result.l0_pitches = memory_detector.GetL0Pitches(); result.read_vars = memory_detector.GetReadVars(); result.write_vars = memory_detector.GetWriteVars(); return result; @@ -412,9 +414,9 @@ inline TaskAccessInfo AnalyzeTaskAccesses(const Stmt &stmt, inline bool TaskAccessesStorage(const TaskAccessInfo &accesses, const Var &storage) { - auto regions_access = [&](const std::vector ®ions) { + auto regions_access = [&](const std::vector ®ions) { return std::any_of(regions.begin(), regions.end(), [&](const auto ®ion) { - return region.logical->buffer->data.same_as(storage); + return region->buffer->data.same_as(storage); }); }; if (regions_access(accesses.read_regions) || @@ -507,8 +509,7 @@ class CoreMaskAvailability { CoreMask result = kCoreBroadcast; for (const Var &var : accesses.read_vars) result &= GetValueCoreMask(var, reader); - for (const MemoryAccess &access : accesses.read_regions) { - const BufferRegion ®ion = access.logical; + for (const BufferRegion ®ion : accesses.read_regions) { if (IsSharedBuffer(region->buffer)) { result &= kCoreVector; } else if (IsL1Buffer(region->buffer) || IsL0ABuffer(region->buffer) || @@ -524,7 +525,7 @@ class CoreMaskAvailability { CoreMask GetTaskReadCoreMask(const TaskNode *task) const { ICHECK(task != nullptr); TaskAccessInfo accesses; - accesses.read_regions = task->GetReadAccesses(); + accesses.read_regions = task->GetReadRegions(); accesses.read_vars = task->GetReadVars(); return GetAccessCoreMask(accesses, task); } @@ -628,10 +629,14 @@ inline void ApplyTaskResourceUsage(const TaskResourceUsage &usage, inline void ApplyTaskAccesses(const TaskAccessInfo &accesses, TaskNode *task) { ICHECK(task != nullptr); - for (const MemoryAccess &access : accesses.read_regions) - task->AddReadAccess(access); - for (const MemoryAccess &access : accesses.write_regions) - task->AddWriteAccess(access); + for (const BufferRegion ®ion : accesses.read_regions) + task->AddReadRegion(region); + for (const BufferRegion ®ion : accesses.write_regions) + task->AddWriteRegion(region); + auto pitches = task->GetL0Pitches(); + for (const auto &[storage, pitch] : accesses.l0_pitches) + MergeL0Pitch(&pitches, storage, pitch); + task->SetL0Pitches(pitches); for (const Var &var : accesses.read_vars) task->AddReadVar(var); for (const Var &var : accesses.write_vars) diff --git a/src/ascend/transform/resolve_core.cc b/src/ascend/transform/resolve_core.cc index a44d96c3d0..5e40ecbeaf 100644 --- a/src/ascend/transform/resolve_core.cc +++ b/src/ascend/transform/resolve_core.cc @@ -221,8 +221,7 @@ class CoreMaskResolver { bool changed = false; for (const Var &var : accesses.read_vars) changed |= AddDemand(&var_demand_, var, mask); - for (const ascend::MemoryAccess &access : accesses.read_regions) { - const BufferRegion ®ion = access.logical; + for (const BufferRegion ®ion : accesses.read_regions) { if (IsRegisterRegion(region)) changed |= AddDemand(®ister_demand_, region->buffer->data, mask); } diff --git a/testing/ascend/auto_schedule/test_sync_dependencies.py b/testing/ascend/auto_schedule/test_sync_dependencies.py index fdd9678819..ef5850b763 100644 --- a/testing/ascend/auto_schedule/test_sync_dependencies.py +++ b/testing/ascend/auto_schedule/test_sync_dependencies.py @@ -136,22 +136,25 @@ def test_dma_row_writes_only_serialize_overlapping_regions(overlap): tvm.ir.assert_structural_equal(releases[0].args[1], waits[0].args[1]) -@pytest.mark.parametrize("overlap", [False, True], ids=["disjoint-halves", "overwrite-middle"]) -def test_partial_l0_loads_complete_before_gemm_and_order_overlapping_writes(overlap): +@pytest.mark.parametrize( + "overlap, declared_rows", + [(False, 32), (False, 48), (True, 48)], + ids=["disjoint-halves", "declared-overlap", "overwrite-middle"], +) +def test_partial_l0_loads_complete_before_gemm_and_order_overlapping_writes(overlap, declared_rows): src = tirx.decl_buffer((64, 64), "bfloat16", name="src", scope="shared.l1") a = tirx.decl_buffer((64, 64), "bfloat16", name="a", scope="shared.l0a") b = tirx.decl_buffer((64, 64), "bfloat16", name="b", scope="shared.l0b") c = tirx.decl_buffer((64, 64), "float32", name="c", scope="shared.l0c") # B and the L1 source are already prepared at this pass boundary. - # The first destination can hold 48 rows; its actual write covers only - # the source's 32 or 48 rows. Only the latter overlaps the next copy. + # Dependencies conservatively cover the declared destination region. body = seq( unit( copy( src, a, src_shape=[48 if overlap else 32, 64], - dst_shape=[48, 64], + dst_shape=[declared_rows, 64], annotations={"stride": tirx.IntImm("int32", 4)}, ), core=2, @@ -173,7 +176,7 @@ def test_partial_l0_loads_complete_before_gemm_and_order_overlapping_writes(over layouts = {src: make_ascend_major_k_layout(src), a: make_ascend_major_k_layout(a)} after = transform.InsertSync()(kernel(body, buffers=[src, a, b, c], annotations={"layout_map": layouts})) barriers = [call for call in calls(after, "tl.ascend_pipe_barrier") if call.args[0].value == "PIPE_MTE1"] - assert bool(barriers) == overlap + assert bool(barriers) == (declared_rows > 32) releases, waits = _events(after, "set", "MTE1_M"), _events(after, "wait", "MTE1_M") assert releases and waits for release in releases: @@ -189,10 +192,10 @@ def test_partial_l0_loads_complete_before_gemm_and_order_overlapping_writes(over @pytest.mark.parametrize("transposed", [False, True], ids=["major-k", "major-mn"]) @pytest.mark.parametrize( "plane, mn_start, k_start, mn_size, k_size, overlap", - [(1, 32, 0, 32, 64, False), (0, 32, 0, 32, 64, True), (0, 48, 48, 16, 16, False)], - ids=["separate-planes", "physical-overlap", "physical-disjoint"], + [(1, 32, 0, 32, 64, False), (0, 32, 0, 32, 64, True), (0, 48, 48, 16, 16, True)], + ids=["separate-planes", "physical-overlap", "different-pitch-disjoint-regions"], ) -def test_different_l0_pitches_order_physical_overlap(transposed, plane, mn_start, k_start, mn_size, k_size, overlap): +def test_different_l0_pitches_conflict_within_same_plane(transposed, plane, mn_start, k_start, mn_size, k_size, overlap): src = tirx.decl_buffer((64, 64), "bfloat16", name="src", scope="shared.l1") a = tirx.decl_buffer((2, 64, 64), "bfloat16", name="a", scope="shared.l0b" if transposed else "shared.l0a") b = tirx.decl_buffer((64, 64), "bfloat16", name="b", scope="shared.l0a" if transposed else "shared.l0b") @@ -231,6 +234,36 @@ def test_different_l0_pitches_order_physical_overlap(transposed, plane, mn_start assert bool(_events(after, "wait", "M_MTE1")) == overlap +@pytest.mark.parametrize("pitch_relation", ["same", "equivalent", "unknown"]) +def test_partial_l0_copies_require_provably_equal_pitches(pitch_relation): + src = tirx.decl_buffer((64, 64), "bfloat16", name="src", scope="shared.l1") + a = tirx.decl_buffer((80, 64), "bfloat16", name="a", scope="shared.l0a") + p, q = tirx.Var("p", "int32"), tirx.Var("q", "int32") + first_pitch = 4 + p % 2 + second_pitch = {"same": first_pitch, "equivalent": 5 - (1 - p % 2), "unknown": 4 + q % 2}[pitch_relation] + body = seq( + *( + unit( + copy( + src, + a, + src_indices=[start, 0], + dst_indices=[start, 0], + src_shape=[32, 64], + dst_shape=[32, 64], + annotations={"stride": pitch}, + ), + core=2, + ) + for start, pitch in [(0, first_pitch), (32, second_pitch)] + ) + ) + layouts = {src: make_ascend_major_k_layout(src), a: make_ascend_major_k_layout(a)} + after = transform.InsertSync()(kernel(body, buffers=[src, a], params=[p, q], annotations={"layout_map": layouts})) + barriers = [call for call in calls(after, "tl.ascend_pipe_barrier") if call.args[0].value == "PIPE_MTE1"] + assert bool(barriers) == (pitch_relation == "unknown") + + @pytest.mark.parametrize("varying_pitch", [False, True], ids=["disjoint", "overlap-after-pitch-change"]) def test_l0_copy_pitch_is_instantiated_per_iteration(varying_pitch): src = tirx.decl_buffer((16, 16), "bfloat16", name="src", scope="shared.l1") @@ -254,22 +287,22 @@ def test_l0_copy_pitch_is_instantiated_per_iteration(varying_pitch): @pytest.mark.parametrize("overlap", [False, True], ids=["disjoint", "overlap"]) @pytest.mark.parametrize("unroll", [False, True], ids=["internal-loop", "multiple-calls"]) -def test_atomic_l0_copies_keep_address_and_pitch_correlated(overlap, unroll): +@pytest.mark.parametrize("varying_pitch", [False, True], ids=["same-pitch", "mixed-pitches"]) +def test_atomic_l0_copies_require_a_uniform_pitch(overlap, unroll, varying_pitch): src = tirx.decl_buffer((16, 16), "bfloat16", name="src", scope="shared.l1") - a = tirx.decl_buffer((32, 32), "bfloat16", name="a", scope="shared.l0a") + a = tirx.decl_buffer((32, 48), "bfloat16", name="a", scope="shared.l0a") b = tirx.decl_buffer((16, 32 if overlap else 16), "bfloat16", name="b", scope="shared.l0b") c = tirx.decl_buffer((16, 16), "float32", name="c", scope="shared.l0c") i = tirx.Var("i", "int32") load = copy( src, a, - dst_indices=[16 * i, 16 * (1 - i)], + dst_indices=[0, 16 * (i + 1)], dst_shape=[16, 16], - annotations={"stride": 1 + i}, + annotations={"stride": 1 + i if varying_pitch else tirx.IntImm("int32", 1)}, ) lhs = tirx.BufferRegion(a, [tvm.ir.Range.from_min_extent(0, e) for e in (16, 32 if overlap else 16)]) - # Both loop iterations write block 1. The consumer reads block 0 or - # blocks [0,2); taking the logical loop hull loses this distinction. + # Mixed pitches conservatively conflict even when the K regions are disjoint. loads = ( seq(*(tirx.stmt_functor.substitute(load, {i: value}) for value in (0, 1))) if unroll @@ -281,8 +314,8 @@ def test_atomic_l0_copies_keep_address_and_pitch_correlated(overlap, unroll): ) layouts = {src: make_ascend_major_k_layout(src), a: make_ascend_major_k_layout(a)} after = transform.InsertSync()(kernel(body, buffers=[src, a, b, c], annotations={"layout_map": layouts})) - assert bool(_events(after, "set", "MTE1_M")) == overlap - assert bool(_events(after, "wait", "MTE1_M")) == overlap + assert bool(_events(after, "set", "MTE1_M")) == (overlap or varying_pitch) + assert bool(_events(after, "wait", "MTE1_M")) == (overlap or varying_pitch) def test_equivalent_lexical_epoch_serializes_counter_ring():