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 1e2e3bcf4a..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" @@ -116,7 +117,9 @@ 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 ConflictHintList &root_conflicts, + Optional a_pitch, + Optional b_pitch) { const bool cross = offset != 0; if (cross && loop == nullptr) return false; @@ -150,8 +153,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; + 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(); @@ -208,21 +211,35 @@ 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. FreshenMutableReads freshen_a(FreshenMutableReads::Mode::kSnapshot); FreshenMutableReads freshen_b(FreshenMutableReads::Mode::kSnapshot); + 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); + } + // 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 < ndim; ++i) { + for (size_t i = 0; i < compared_dims; ++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)); + // 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 @@ -389,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; } @@ -405,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 e92ecbb5d1..9c08634113 100644 --- a/src/ascend/transform/auto_schedule/dependency_analysis.h +++ b/src/ascend/transform/auto_schedule/dependency_analysis.h @@ -109,7 +109,8 @@ 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. 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 // cross-iteration check. Normally it renames the loop suffix below the nearest @@ -123,7 +124,9 @@ 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 774510534e..e91280062a 100644 --- a/src/ascend/transform/auto_schedule/ir_structure.cc +++ b/src/ascend/transform/auto_schedule/ir_structure.cc @@ -299,6 +299,7 @@ std::shared_ptr TaskNode::Clone() const { // Copy memory access regions 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 c749e4258b..36ac7ebfb1 100644 --- a/src/ascend/transform/auto_schedule/ir_structure.h +++ b/src/ascend/transform/auto_schedule/ir_structure.h @@ -556,6 +556,16 @@ class TaskNode : public IRStructure { 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 { @@ -588,6 +598,14 @@ class TaskNode : public IRStructure { 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); @@ -657,6 +675,7 @@ class TaskNode : public IRStructure { // 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 681de2e96a..85fa4ad945 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" @@ -35,18 +40,69 @@ namespace ascend { using namespace tirx; +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; + } + pitches->Set(storage, pitch); +} + +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 layout = FindLayoutForBuffer(layouts, copy->dst); + AscendFractalLayoutInfo info; + 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_)) { + 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; +} + // 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. + // 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) { + l0_pitches_.clear(); + operation_pitches_.clear(); read_buffers_.clear(); write_buffers_.clear(); read_regions_.clear(); @@ -68,6 +124,8 @@ class MemoryAccessDetector : public StmtExprVisitor { } void Analyze(const PrimExpr &expr) { + l0_pitches_.clear(); + operation_pitches_.clear(); read_buffers_.clear(); write_buffers_.clear(); read_regions_.clear(); @@ -98,6 +156,10 @@ class MemoryAccessDetector : public StmtExprVisitor { return CollectRegions(write_buffers_, write_regions_); } + const Map> &GetL0Pitches() const { + return l0_pitches_; + } + // Return all variables that are read from std::vector GetReadVars() const { return read_vars_; } // Return all variables that are written to @@ -117,6 +179,10 @@ class MemoryAccessDetector : public StmtExprVisitor { } private: + LayoutMap layouts_; + Map operation_pitches_; + Map> l0_pitches_; + /*! \brief Iteration range for loop_vars */ std::unordered_map dom_map_; /*! \brief Extra iteration range hint for free vars */ @@ -180,6 +246,21 @@ class MemoryAccessDetector : public StmtExprVisitor { for (const PrimExpr &stride : buffer->strides) visit_mutable_layout(stride); visit_mutable_layout(buffer->elem_offset); + 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); } @@ -346,7 +427,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)) @@ -382,6 +463,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) || @@ -918,6 +1029,59 @@ 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_pitches = std::move(operation_pitches_); + operation_pitches_ = GetCallL0Pitches(ffi::GetRef(op), layouts_); + StmtExprVisitor::VisitExpr_(op); + operation_pitches_ = std::move(saved_pitches); + 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 0169bc99d1..294930162e 100644 --- a/src/ascend/transform/auto_schedule/task_analysis.h +++ b/src/ascend/transform/auto_schedule/task_analysis.h @@ -176,6 +176,7 @@ struct TaskResourceUsage { struct TaskAccessInfo { std::vector read_regions; std::vector write_regions; + Map> l0_pitches; std::vector read_vars; std::vector write_vars; }; @@ -397,12 +398,15 @@ 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.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; @@ -629,6 +633,10 @@ inline void ApplyTaskAccesses(const TaskAccessInfo &accesses, TaskNode *task) { 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/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..ef5850b763 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,188 @@ 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, 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. + # Dependencies conservatively cover the declared destination region. + body = seq( + unit( + copy( + src, + a, + src_shape=[48 if overlap else 32, 64], + dst_shape=[declared_rows, 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) == (declared_rows > 32) + 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, True)], + ids=["separate-planes", "physical-overlap", "different-pitch-disjoint-regions"], +) +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") + 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("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") + 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"]) +@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, 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=[0, 16 * (i + 1)], + dst_shape=[16, 16], + 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)]) + # 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 + 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 or varying_pitch) + assert bool(_events(after, "wait", "MTE1_M")) == (overlap or varying_pitch) + + 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)