diff --git a/.agents/skills/tilelang-ascend/SKILL.md b/.agents/skills/tilelang-ascend/SKILL.md index 705102a3ff..9cb632d6fe 100644 --- a/.agents/skills/tilelang-ascend/SKILL.md +++ b/.agents/skills/tilelang-ascend/SKILL.md @@ -347,6 +347,21 @@ when changing manual kernels. For capacity or launch failures, record the generated allocation sizes, buffer versions, VF mode, and compiler options. Check SDK reservations as well as user buffers before changing tile sizes. +For Cube storage (L1, L0A, L0B, L0C), `T.view` and `T.reshape` aliases can have +different canonical layouts inferred from their operations or supplied with +`T.annotate_layout`. For example, a `(64, 128)` L0B allocation can have a +`(128, 64)` view filled with `T.copy(..., transpose=True)` and consumed with +`T.gemm(..., transpose_B=True)`. A same-shape view can instead use a different +`MajorK`/`MajorMN` layout. Each view keeps its own address mapping while the +shared `.data` keeps one allocation and version clock; `T.assume_conflict` is +unnecessary for these aliases. + +These views require dense logical strides, zero base offset, the same accessed +dtype, and equal padded physical sizes. They reinterpret storage without moving +existing values: the producer must write the representation its consumer uses. +Use `T.access_ptr(view[...], ...)` when passing a view to an intrinsic, so layout +lowering can select its address mapping. + `T.annotate_unlimit_memory("shared", "shared.l1")` inside a kernel removes the scheduler's capacity constraint for those pools. Other supported scopes are `shared.l0a`, `shared.l0b`, and `shared.l0c`. This is useful for investigating a diff --git a/src/ascend/op/copy.cc b/src/ascend/op/copy.cc index 302d832b99..bfc0259ae9 100644 --- a/src/ascend/op/copy.cc +++ b/src/ascend/op/copy.cc @@ -423,11 +423,13 @@ Stmt LowerDMACopy(const AscendCopyNode &op, const LowerArgs &T, Array src_range = NormalizeEmptyUnitAxesForMTE(op.src_range, analyzer); Array dst_range = NormalizeEmptyUnitAxesForMTE(op.dst_range, analyzer); - auto make_access_ptr = [&](const Buffer &buf, PrimExpr offset, + auto make_access_ptr = [&](const Buffer &buf, const Array &ranges, PrimExpr extent, int rw_mask) -> PrimExpr { - return Call( - DataType::Handle(), builtin::tvm_access_ptr(), - {TypeAnnotation(buf->dtype), buf->data, offset, extent, I(rw_mask)}); + Array indices; + for (const Range &range : ranges) + indices.push_back(range->min); + return Call(DataType::Handle(), tl::access_ptr(), + {BufferLoad(buf, indices), extent, I(rw_mask)}); }; auto region_elements = [&](const Array &ranges) { PrimExpr result = I(1); @@ -442,10 +444,8 @@ Stmt LowerDMACopy(const AscendCopyNode &op, const LowerArgs &T, op.dst, dst_range, op.dst->dtype.bits() * op.dst->dtype.lanes()); PrimExpr src_region_elements = region_elements(src_range); PrimExpr dst_region_elements = region_elements(dst_range); - PrimExpr src_ptr = - make_access_ptr(op.src, src_region.offset, src_region_elements, 1); - PrimExpr dst_ptr = - make_access_ptr(op.dst, dst_region.offset, dst_region_elements, 2); + PrimExpr src_ptr = make_access_ptr(op.src, src_range, src_region_elements, 1); + PrimExpr dst_ptr = make_access_ptr(op.dst, dst_range, dst_region_elements, 2); PrimExpr sid = I(0); PrimExpr zero = I(0); @@ -464,9 +464,9 @@ Stmt LowerDMACopy(const AscendCopyNode &op, const LowerArgs &T, PrimExpr plan_total_elems = plan.row_elems * plan.n_rows; PrimExpr plan_src_ptr = - make_access_ptr(op.src, plan.src_offset, plan_total_elems, 1); + make_access_ptr(op.src, src_range, plan_total_elems, 1); PrimExpr plan_dst_ptr = - make_access_ptr(op.dst, plan.dst_offset, plan_total_elems, 2); + make_access_ptr(op.dst, dst_range, plan_total_elems, 2); PrimExpr plan_total_bytes = plan_row_bytes * plan.n_rows; // Ascend GM->UB padding (data_select): an unaligned row is right-padded up diff --git a/src/ascend/op/oob_padding.cc b/src/ascend/op/oob_padding.cc index 03759cb935..f2a78519d4 100644 --- a/src/ascend/op/oob_padding.cc +++ b/src/ascend/op/oob_padding.cc @@ -69,15 +69,7 @@ DMAPath GetDMAPath(const Buffer &src, const Buffer &dst) { Optional FindLayoutForBuffer(const LayoutMap &layout_map, const Buffer &buffer) { - if (layout_map.count(buffer)) { - return layout_map[buffer]; - } - for (const auto &kv : layout_map) { - if (kv.first->data.same_as(buffer->data)) { - return kv.second; - } - } - return std::nullopt; + return layout_map.Get(buffer); } BoundedDMACopyRanges ClampDMACopyTail(const AscendCopyNode &op, diff --git a/src/ascend/op/oob_padding.h b/src/ascend/op/oob_padding.h index eaf0e2b6e4..de79ebb5cc 100644 --- a/src/ascend/op/oob_padding.h +++ b/src/ascend/op/oob_padding.h @@ -53,8 +53,8 @@ enum class DMAPath { DMAPath GetDMAPath(const Buffer &src, const Buffer &dst); /*! - * \brief Look up a buffer's inferred Layout in a layout map, matching either - * the buffer object or its underlying data Var. Returns nullopt if absent. + * \brief Look up a view's inferred Layout by Buffer identity. Storage aliases + * may have distinct layouts; inference resolves any layout inheritance. */ Optional FindLayoutForBuffer(const LayoutMap &layout_map, const Buffer &buffer); diff --git a/src/ascend/op/utils.h b/src/ascend/op/utils.h index e56a22f330..ca4d9a9870 100644 --- a/src/ascend/op/utils.h +++ b/src/ascend/op/utils.h @@ -13,6 +13,7 @@ #include +#include "op/builtin.h" #include "op/utils.h" namespace tvm { @@ -75,7 +76,7 @@ inline bool IsAscendOnChipBuffer(const Buffer &buffer) { } /*! - * \brief tvm_access_ptr for a region whose innermost 2-D position travels as + * \brief A view-preserving pointer whose innermost 2-D position travels as * separate intrinsic arguments. * * The Ascend L1->L0 loads pass the M/K start positions as their own arguments @@ -85,13 +86,8 @@ inline bool IsAscendOnChipBuffer(const Buffer &buffer) { * example -- and counting the last two dims into the offset as well would place * that position twice. * - * This builds exactly that pointer, on top of the neutral - * MakeAccessPtrFromRegion: the innermost two dims are re-based to zero before - * delegating, which leaves their extents (and hence the reported extent) alone - * and drops them from the offset. - * - * The buffer keeps ``require_2d == false`` semantics, so 1-D regions still take - * the neutral function's dedicated path and are unaffected. + * Keep the Buffer identity until layout lowering: a data Var alone cannot + * identify which canonical matrix view supplies the physical addresses. */ inline PrimExpr MakeAscendLeadingDimAccessPtr(const BufferRegion ®ion, int rw_mask) { @@ -104,8 +100,15 @@ inline PrimExpr MakeAscendLeadingDimAccessPtr(const BufferRegion ®ion, const auto &r = leading[i]; leading.Set(i, Range::FromMinExtent(make_zero(r->min.dtype()), r->extent)); } - return MakeAccessPtrFromRegion(BufferRegion(region->buffer, leading), - rw_mask); + ffi::Array indices; + for (const Range &range : leading) + indices.push_back(range->min); + PrimExpr extent = leading.back()->extent; + if (n >= 2) + extent = extent * leading[n - 2]->extent; + return Call(DataType::Handle(), tl::access_ptr(), + {BufferLoad(region->buffer, indices), extent, + IntImm(DataType::Int(32), rw_mask)}); } } // namespace tl diff --git a/src/ascend/transform/layout_inference.cc b/src/ascend/transform/layout_inference.cc index 17425a378a..d3b1c716fc 100644 --- a/src/ascend/transform/layout_inference.cc +++ b/src/ascend/transform/layout_inference.cc @@ -316,6 +316,11 @@ class BufferUseDefCollector : public IRVisitorWithAnalyzer { // Helper: propagate inferred layout to alias buffers (same data Var) auto propagate_alias = [&](const Buffer &src_buffer, const Layout &src_layout) { + // Fragment aliases participate in the thread-ownership solver. + // Memory layouts belong to individual views; fill missing layouts + // only after each view's own constraints have been inferred. + if (!src_layout.as()) + return; if (!buffer_data_to_buffers_.count(src_buffer->data)) return; const auto &siblings = buffer_data_to_buffers_[src_buffer->data]; @@ -623,45 +628,41 @@ class BufferUseDefCollector : public IRVisitorWithAnalyzer { in_queue); // step 3: relax constraints to free and re-run InferInFreeMode(layout_map, strict_layout_map); - // step 4: finalize alias layouts by Var - // For each storage var, if any buffer in the group has a layout, - // propagate (reshape if needed) to the rest to ensure completeness. + // step 4: fill unconstrained views only when all known interpretations of + // their storage agree. A shared data Var does not impose layout equality. + const LayoutMap inferred_layouts = layout_map; for (const auto &[var, buffers] : buffer_data_to_buffers_) { - // Find a representative with existing layout - Optional rep; - Optional rep_layout; - for (const auto &buf : buffers) { - if (layout_map.count(buf)) { - rep = buf; - rep_layout = layout_map[buf]; - break; - } - } - if (!rep_layout.defined()) - continue; for (const auto &buf : buffers) { - if (!layout_map.count(buf)) { - bool shapes_equal = - rep_layout.value()->InputShape().size() == buf->shape.size(); - if (shapes_equal) { - for (size_t i = 0; i < rep_layout.value()->InputShape().size(); - ++i) { - if (!analyzer_.CanProveEqual(rep_layout.value()->InputShape()[i], - buf->shape[i])) { - shapes_equal = false; - break; - } - } - } - Layout reshaped = - shapes_equal - ? rep_layout.value() - : rep_layout.value()->Reshape( + if (layout_map.count(buf)) + continue; + Optional inherited; + bool ambiguous = false; + for (const auto &source : buffers) { + Optional known = inferred_layouts.Get(source); + if (!known.defined()) + continue; + Layout candidate = + ShapesEqual(known.value()->InputShape(), buf->shape, &analyzer_) + ? known.value() + : known.value()->Reshape( buf->shape, &analyzer_, - Integer(GetElementStorageBits(rep.value()->dtype)), + Integer(GetElementStorageBits(source->dtype)), Integer(GetElementStorageBits(buf->dtype))); - layout_map.Set(buf, reshaped); + if (inherited.defined() && + !inherited.value()->IsEqual(candidate.get())) { + ambiguous = true; + break; + } + inherited = candidate; } + ICHECK(!ambiguous || !accessed_buffers_.count(buf)) + << "Ambiguous layout for buffer view " << buf->name << ": storage " + << var->name_hint + << " has distinct layouts. Annotate this view explicitly."; + // An unused backing Buffer only supplies allocation geometry. Any + // candidate suffices; storage normalization checks physical footprints. + if (inherited.defined()) + layout_map.Set(buf, inherited.value()); } } @@ -869,6 +870,7 @@ class BufferUseDefCollector : public IRVisitorWithAnalyzer { } void addToUseList(const Buffer &buffer) { + accessed_buffers_.insert(buffer); // Fragment buffers and reducers (whose PartialFragment travels the same // use_list_ edges: commits notify users, shared buffers union free-mode // components) both register; anything else has no layout to solve. @@ -1001,18 +1003,7 @@ class BufferUseDefCollector : public IRVisitorWithAnalyzer { // annotated layouts are available during InferLayout. // Must happen after visiting body so buffer_data_to_buffers_ is // populated. - if (pushed_region_scope && op->annotations.count(attr::kLayoutMap)) { - auto map = op->annotations.Get(attr::kLayoutMap) - ->as>() - .value(); - for (const auto &[var, layout] : map) { - if (buffer_data_to_buffers_.count(var)) { - for (const auto &buffer : buffer_data_to_buffers_[var]) { - annotated_layout_map_.Set(buffer, layout); - } - } - } - } + RecordLayoutAnnotations(GetRef(op)); // Regions contribute nothing to the inference worklist themselves: // nested tile operators were already collected while visiting the // body, and the region block is consumed by codegen as-is. @@ -1035,16 +1026,23 @@ class BufferUseDefCollector : public IRVisitorWithAnalyzer { // First, visit the block body to collect all buffers from // BufferLoad/BufferStore IRVisitorWithAnalyzer::VisitStmt_(op); + RecordLayoutAnnotations(GetRef(op)); + } + void RecordLayoutAnnotations(const SBlock &block) { // After visiting, apply layouts to all collected buffers - if (op->annotations.count(attr::kLayoutMap)) { - // Check if the layout map is Map - auto map = - op->annotations.Get(attr::kLayoutMap)->as>().value(); - for (const auto &[var, layout] : map) { + if (block->annotations.count(attr::kLayoutMap)) { + auto map = block->annotations.Get(attr::kLayoutMap) + ->cast>(); + for (const auto &[key, layout] : map) { + Optional annotated = key.as(); + Var var = + annotated.defined() ? annotated.value()->data : Downcast(key); ICHECK(buffer_data_to_buffers_.count(var)) << "buffer " << var << " is not found in the block"; - const auto &buffers = buffer_data_to_buffers_[var]; + Array buffers = annotated.defined() + ? Array{annotated.value()} + : buffer_data_to_buffers_[var]; ICHECK(!buffers.empty()) << "buffer list for " << var << " is empty"; for (const auto &buffer : buffers) { if (!IsReducerV2Buffer(buffer)) { @@ -1178,6 +1176,7 @@ class BufferUseDefCollector : public IRVisitorWithAnalyzer { } void VisitExpr_(const BufferLoadNode *op) final { + accessed_buffers_.insert(op->buffer); // Collect buffer from BufferLoad if (op->buffer.defined() && op->buffer->data.defined()) { if (buffer_data_to_buffers_.count(op->buffer->data)) { @@ -1208,6 +1207,7 @@ class BufferUseDefCollector : public IRVisitorWithAnalyzer { } void VisitStmt_(const BufferStoreNode *op) final { + accessed_buffers_.insert(op->buffer); // Collect buffer from BufferStore if (op->buffer.defined() && op->buffer->data.defined()) { if (buffer_data_to_buffers_.count(op->buffer->data)) { @@ -1336,6 +1336,7 @@ class BufferUseDefCollector : public IRVisitorWithAnalyzer { } Map> buffer_data_to_buffers_; + std::unordered_set accessed_buffers_; // Map from Bind variable to its bound expression Map bind_var_to_expr_; std::vector infer_list_stmt_; diff --git a/src/ascend/transform/lower_tile_op.cc b/src/ascend/transform/lower_tile_op.cc index 82a52a4389..63c1ca14d9 100644 --- a/src/ascend/transform/lower_tile_op.cc +++ b/src/ascend/transform/lower_tile_op.cc @@ -641,8 +641,8 @@ class LowerTileOpPass : arith::IRMutatorWithAnalyzer { << "but got indices size: " << indices.size() << " and shape size: " << old_shape.size(); - Buffer remap_key = FindRemapBuffer(load->buffer).value_or(load->buffer); - Optional layout = FindLayout(remap_key); + Buffer remap_key = load->buffer; + Optional layout = layout_map_.Get(remap_key); if (!layout.defined() || !buffer_map_.count(remap_key->data)) { return result; } @@ -695,10 +695,11 @@ class LowerTileOpPass : arith::IRMutatorWithAnalyzer { } BufferLoad load = Downcast(access_ptr_call->args[0]); - PrimExpr extent = access_ptr_call->args[1]; - PrimExpr rw_mask = access_ptr_call->args[2]; + PrimExpr extent = VisitExpr(access_ptr_call->args[1]); + PrimExpr rw_mask = VisitExpr(access_ptr_call->args[2]); - Array indices = load->indices; + Array indices = load->indices.Map( + [&](const PrimExpr &index) { return VisitExpr(index); }); Array old_shape = load->buffer->shape; ICHECK_EQ(indices.size(), old_shape.size()) @@ -707,8 +708,8 @@ class LowerTileOpPass : arith::IRMutatorWithAnalyzer { << "but got indices size: " << indices.size() << " and shape size: " << old_shape.size(); - Buffer remap_key = FindRemapBuffer(load->buffer).value_or(load->buffer); - Optional layout = FindLayout(remap_key); + Buffer remap_key = load->buffer; + Optional layout = layout_map_.Get(remap_key); if (!layout.defined() || !buffer_map_.count(remap_key->data)) { return result; } @@ -736,6 +737,22 @@ class LowerTileOpPass : arith::IRMutatorWithAnalyzer { Array new_indices = RemapAccessIndices( indices, old_shape, new_shape, layout.value(), offset); + if (auto pitch = access_ptr_call->annotations.Get(attr::kL0AccessPitch)) { + ICHECK(IsL0ABuffer(remap_key) || IsL0BBuffer(remap_key)); + size_t n = new_shape.size(); + ICHECK_GE(n, 4); + // A compact L0 copy changes this access's matrix row stride, not the + // allocation or the distance between storage versions. + Array strides(n, Integer(1)); + PrimExpr stride = Integer(1); + for (size_t i = n; i > 0; --i) { + strides.Set(i - 1, stride); + stride = stride * new_shape[i - 1]; + } + strides.Set(n - 4, Downcast(pitch.value()) * + new_shape[n - 2] * new_shape[n - 1]); + new_buffer.CopyOnWrite()->strides = std::move(strides); + } Array new_args = {BufferLoad(new_buffer, new_indices), extent, rw_mask}; if (buffer_remap_.count(remap_key)) { @@ -766,44 +783,6 @@ class LowerTileOpPass : arith::IRMutatorWithAnalyzer { return Optional(); } - Optional FindRemapBuffer(const Buffer &buffer) const { - if (buffer_remap_.count(buffer)) { - return buffer; - } - auto it = buffer_map_.find(buffer->data); - if (it != buffer_map_.end() && buffer_remap_.count(it->second)) { - return it->second; - } - for (const auto &kv : buffer_remap_) { - if (kv.first->data.same_as(buffer->data)) { - return kv.first; - } - if (kv.first->name == buffer->name) { - return kv.first; - } - } - return Optional(); - } - - Optional FindLayout(const Buffer &buffer) const { - if (layout_map_.count(buffer)) { - return layout_map_[buffer]; - } - auto it = buffer_map_.find(buffer->data); - if (it != buffer_map_.end() && layout_map_.count(it->second)) { - return layout_map_[it->second]; - } - for (const auto &kv : layout_map_) { - if (kv.first->data.same_as(buffer->data)) { - return kv.second; - } - if (kv.first->name == buffer->name) { - return kv.second; - } - } - return Optional(); - } - PrimExpr VisitExpr_(const tirx::CallNode *op) final { if (op->op.same_as(tl::tma_load()) || op->op.same_as(tl::tma_load_im2col()) || @@ -1008,6 +987,15 @@ class LowerTileOpPass : arith::IRMutatorWithAnalyzer { return call; } + // Preserve the typed base until its view-specific layout and compact L0 + // pitch have both been applied. Visiting BufferLoad first loses that view. + if (op->op.same_as(tl::access_ptr())) { + auto rewritten = HandleAccessPtrAndOffset(GetRef(op)); + if (rewritten.rewritten) + return rewritten.expr; + return IRMutatorWithAnalyzer::VisitExpr_(op); + } + // Handle standalone tvm_access_ptr calls with layout transformation if (op->op.same_as(builtin::tvm_access_ptr())) { auto call = Downcast(IRMutatorWithAnalyzer::VisitExpr_(op)); diff --git a/src/ascend/transform/materialize_multi_buffer.cc b/src/ascend/transform/materialize_multi_buffer.cc index fb3d83c885..3ddca42422 100644 --- a/src/ascend/transform/materialize_multi_buffer.cc +++ b/src/ascend/transform/materialize_multi_buffer.cc @@ -354,30 +354,6 @@ Layout ExpandLayoutToBuffer(const Buffer &layout_buffer, return layout->Expand(leading_shape); } -Layout FindLayoutForBuffer(const Map &layout_map, - const Buffer &buffer) { - if (layout_map.count(buffer)) - return layout_map[buffer]; - for (const auto &[layout_buffer, layout] : layout_map) { - if (layout_buffer->data.same_as(buffer->data)) { - Layout expanded = ExpandLayoutToBuffer(layout_buffer, buffer, layout); - if (expanded.defined()) - return expanded; - } - } - return Layout(); -} - -bool AddLayoutForBuffer(Map *layout_map, const Buffer &buffer) { - if (layout_map->count(buffer)) - return false; - Layout layout = FindLayoutForBuffer(*layout_map, buffer); - if (!layout.defined()) - return false; - layout_map->Set(buffer, layout); - return true; -} - class AllocationRewriter : public StmtMutator { public: explicit AllocationRewriter(VersionedBufferRegistry *buffers) @@ -406,8 +382,16 @@ class AllocationRewriter : public StmtMutator { if (auto value = op->annotations.Get(attr::kLayoutMap)) { if (auto layout_map = value.value().as>()) { Map rewritten_layouts = layout_map.value(); - for (const Buffer &buffer : alloc_buffers) - layout_changed |= AddLayoutForBuffer(&rewritten_layouts, buffer); + // Expand each view from its own layout. Data-Var lookup is ambiguous + // when several canonical matrix views share a versioned allocation. + for (const auto &[buffer, layout] : layout_map.value()) { + Buffer expanded = buffers_->Resolve(buffer); + if (!expanded.same_as(buffer)) { + rewritten_layouts.Set( + expanded, ExpandLayoutToBuffer(buffer, expanded, layout)); + layout_changed = true; + } + } if (layout_changed) annotations.Set(attr::kLayoutMap, rewritten_layouts); } diff --git a/src/ascend/transform/normalize_fractal_storage.cc b/src/ascend/transform/normalize_fractal_storage.cc index 7fb7228b82..bdb554be6b 100644 --- a/src/ascend/transform/normalize_fractal_storage.cc +++ b/src/ascend/transform/normalize_fractal_storage.cc @@ -61,10 +61,9 @@ class BufferCollector : public StmtExprVisitor { } }; -// Every alias has the same padded matrix shape and canonical layout. Its old -// reshape is absorbed into indices/regions here, not a special layout or hidden -// pitch. Ordinary capacity accounting and version expansion therefore see the -// complete storage through every alias. GM/UB and compact NZ are untouched. +// Layouts and padded matrix coordinates belong to each Buffer. The data Var +// identifies only the allocation, whose physical size must agree across views. +// GM/UB and compact NZ geometry are untouched. class FractalStorageNormalizer : public arith::IRMutatorWithAnalyzer { public: static PrimFunc Rewrite(PrimFunc func) { @@ -73,7 +72,7 @@ class FractalStorageNormalizer : public arith::IRMutatorWithAnalyzer { arith::Analyzer analyzer; analyzer.z3_prover.SetRLimit(50000); FractalStorageNormalizer rewriter(&analyzer); - rewriter.canonical_layouts_ = collector.layouts; + rewriter.layouts_ = collector.layouts; rewriter.MarkBufferMapShapes(func); for (const auto &[buffer, layout] : collector.layouts) { if (!IsAscendFractalLayout(layout) || @@ -81,33 +80,37 @@ class FractalStorageNormalizer : public arith::IRMutatorWithAnalyzer { IsL0CBuffer(buffer)) || !collector.owners.count(buffer->data)) continue; - if (auto anchor = rewriter.anchors_.Get(buffer->data)) { - AscendFractalLayoutInfo lhs, rhs; - ICHECK(TryExtractAscendFractalLayout(collector.layouts[anchor.value()], - anchor.value(), &lhs)); - ICHECK(TryExtractAscendFractalLayout(layout, buffer, &rhs)); - ICHECK(lhs.kind == rhs.kind && - collector.layouts[anchor.value()]->IsEqual(layout.get()) && - anchor.value()->dtype.bits() == buffer->dtype.bits()) - << "Conflicting canonical Cube views of " << buffer->name; - if (collector.used_buffers.count(buffer) && - !collector.used_buffers.count(anchor.value())) - rewriter.anchors_.Set(buffer->data, buffer); - } else { - rewriter.anchors_.Set(buffer->data, buffer); - } + auto views = + rewriter.matrix_views_.Get(buffer->data).value_or(Array{}); + views.push_back(buffer); + rewriter.matrix_views_.Set(buffer->data, views); } + Map storage_bits; for (const Buffer &buffer : collector.buffers) { - if (auto anchor = rewriter.anchors_.Get(buffer->data)) { + if (auto views = rewriter.matrix_views_.Get(buffer->data)) { + Buffer matrix = rewriter.FindMatrixView_( + buffer, views.value(), collector.used_buffers.count(buffer)); ICHECK(!collector.used_buffers.count(buffer) || - buffer->dtype == anchor.value()->dtype) + buffer->dtype == matrix->dtype) << "Cube alias accesses must use the canonical operand dtype: " << buffer->name << " has " << buffer->dtype << ", expected " - << anchor.value()->dtype; - rewriter.RegisterBuffer_(buffer, anchor.value(), - collector.owners[buffer->data], - collector.layouts[anchor.value()]); + << matrix->dtype; + rewriter.RegisterBuffer_(buffer, matrix, collector.layouts[matrix]); } + Buffer padded = rewriter.buffers_.Get(buffer).value_or(buffer); + Optional layout = rewriter.layouts_.Get(padded); + if (!collector.owners.count(buffer->data) || !layout.defined() || + layout.value().as()) + continue; + PrimExpr bits = Integer(padded->dtype.bits() * padded->dtype.lanes()); + for (const PrimExpr &extent : layout.value()->OutputShape()) + bits = bits * extent; + if (auto previous = storage_bits.Get(buffer->data)) + ICHECK(analyzer.CanProveEqual(previous.value(), bits)) + << "Buffer views sharing storage require equal physical sizes " + "including padding: " + << buffer->data->name_hint; + storage_bits.Set(buffer->data, bits); } if (rewriter.buffers_.empty()) return func; @@ -120,13 +123,38 @@ class FractalStorageNormalizer : public arith::IRMutatorWithAnalyzer { explicit FractalStorageNormalizer(arith::Analyzer *analyzer) : Parent(analyzer) {} - void RegisterBuffer_(const Buffer &buffer, const Buffer &anchor, - const Buffer &owner, const Layout &layout) { - ICHECK_EQ(buffer->dtype.lanes(), 1); - ICHECK_EQ(buffer->dtype.bits(), anchor->dtype.bits()) + Buffer FindMatrixView_(const Buffer &buffer, const Array &views, + bool accessed) { + ICHECK_EQ(buffer->dtype.bits(), views[0]->dtype.bits()) << "Cube aliases must preserve element width: " << buffer->name - << " has dtype " << buffer->dtype << ", canonical view " << anchor->name - << " has dtype " << anchor->dtype; + << " has dtype " << buffer->dtype; + Optional layout = layouts_.Get(buffer); + if (layout.defined() && IsAscendFractalLayout(layout.value())) + return buffer; + Optional matrix; + for (const Buffer &view : views) { + Layout candidate = layouts_[view]->Reshape(buffer->shape, analyzer_); + if (layout.defined() && !layout.value()->IsEqual(candidate.get())) + continue; + if (!matrix.defined()) + matrix = view; + if (layout.defined() || !accessed) + break; + ICHECK(layouts_[matrix.value()] + ->Reshape(buffer->shape, analyzer_) + ->IsEqual(candidate.get())) + << "Ambiguous layout for Cube view " << buffer->name + << "; annotate this view explicitly"; + } + ICHECK(matrix.defined()) + << "Cube view layout must be a reshape of a canonical matrix layout: " + << buffer->name; + return matrix.value(); + } + + void RegisterBuffer_(const Buffer &buffer, const Buffer &matrix, + const Layout &layout) { + ICHECK_EQ(buffer->dtype.lanes(), 1); ICHECK(buffer->axis_separators.empty() && is_zero(buffer->elem_offset)) << "Ascend fractal buffers require a zero base and no axis separators"; if (!buffer->strides.empty()) { @@ -140,14 +168,14 @@ class FractalStorageNormalizer : public arith::IRMutatorWithAnalyzer { } } Array identity; - for (size_t i = 0; i < anchor->shape.size(); ++i) + for (size_t i = 0; i < matrix->shape.size(); ++i) identity.push_back(InputPlaceholder(i)); logical_maps_.Set( buffer, - Layout(anchor->shape, identity)->Reshape(buffer->shape, analyzer_)); + Layout(matrix->shape, identity)->Reshape(buffer->shape, analyzer_)); AscendFractalLayoutInfo info; - ICHECK(TryExtractAscendFractalLayout(layout, anchor, &info)); - Array shape = anchor->shape; + ICHECK(TryExtractAscendFractalLayout(layout, matrix, &info)); + Array shape = matrix->shape; for (size_t axis = shape.size() - 2; axis < shape.size(); ++axis) { bool is_row = axis == shape.size() - 2; PrimExpr align = (is_row == (info.c0_axis == LogicalAxis::kRow)) @@ -156,16 +184,15 @@ class FractalStorageNormalizer : public arith::IRMutatorWithAnalyzer { shape.Set(axis, analyzer_->Simplify( floordiv(shape[axis] + align - 1, align) * align)); } - // Coalesce aliases to one Buffer identity and geometry so versioning - // expands the allocation and every use together. An unused equal-width - // backing allocation may adopt the canonical operand dtype. - Buffer padded = canonical_buffers_.Get(buffer->data).value_or(owner); + // Preserve each view and the allocation owner's properties. An unused + // equal-width backing allocation may adopt the operand dtype. + Buffer padded = buffer; if (!StructuralEqual()(padded->shape, shape) || !padded->strides.empty() || - padded->dtype != anchor->dtype) { + padded->dtype != matrix->dtype) { auto *writer = padded.CopyOnWrite(); writer->shape = shape; writer->strides = {}; - writer->dtype = anchor->dtype; + writer->dtype = matrix->dtype; } Layout padded_layout; switch (info.kind) { @@ -183,8 +210,8 @@ class FractalStorageNormalizer : public arith::IRMutatorWithAnalyzer { break; } buffers_.Set(buffer, padded); - canonical_buffers_.Set(buffer->data, padded); - canonical_layouts_.Set(padded, padded_layout); + logical_shapes_.Set(buffer, matrix->shape); + layouts_.Set(padded, padded_layout); } Array RewriteIndices_(const Buffer &buffer, @@ -204,7 +231,7 @@ class FractalStorageNormalizer : public arith::IRMutatorWithAnalyzer { }); if (!buffers_.count(buffer)) return BufferRegion(buffer, ranges); - if (!StructuralEqual()(buffer->shape, anchors_[buffer->data]->shape)) { + if (!StructuralEqual()(buffer->shape, logical_shapes_[buffer])) { // A dense reshape is injective only while trailing region extents do // not wrap across the original view's strides. Otherwise equal bounding // box volume alone can hide duplicated logical elements. @@ -247,11 +274,11 @@ class FractalStorageNormalizer : public arith::IRMutatorWithAnalyzer { if (IsL0SFBuffer(copy->dst)) { // SF handles carry the effective shape of their bound data tile and // have no independent storage. Only the L1 source is a transfer region. - source = ascend::GetL1ToL0SFSourceRegion(*copy.get(), canonical_layouts_, - analyzer_); + source = + ascend::GetL1ToL0SFSourceRegion(*copy.get(), layouts_, analyzer_); } else { - auto regions = ascend::GetL1ToL0CopyRegions( - *copy.get(), canonical_layouts_, analyzer_); + auto regions = + ascend::GetL1ToL0CopyRegions(*copy.get(), layouts_, analyzer_); source = regions.source; destination = regions.destination; pitch = regions.pitch; @@ -265,9 +292,8 @@ class FractalStorageNormalizer : public arith::IRMutatorWithAnalyzer { // than the actual transfer. This also makes repeated normalization // stable. AscendFractalRegionInfo normalized; - ICHECK(TryExtractAscendFractalRegion(canonical_layouts_[copy->dst], - destination.value()->region, - &normalized)); + ICHECK(TryExtractAscendFractalRegion( + layouts_[copy->dst], destination.value()->region, &normalized)); if (!analyzer_->CanProveEqual(normalized.outer1->extent, pitch.value())) writer->annotations.Set("l0_pitch", pitch.value()); } @@ -456,7 +482,7 @@ class FractalStorageNormalizer : public arith::IRMutatorWithAnalyzer { // GM region unchanged: extra storage does not authorize an extra read. Copy rewritten(call->args, call->annotations); Array ranges = rewritten->dst_range; - const auto &logical_shape = anchors_[copy->dst->data]->shape; + const auto &logical_shape = logical_shapes_[copy->dst]; for (size_t i = 0; i < ranges.size(); ++i) { if (analyzer_->CanProveEqual(ranges[i]->min, 0) && analyzer_->CanProveEqual(ranges[i]->extent, logical_shape[i])) @@ -471,18 +497,23 @@ class FractalStorageNormalizer : public arith::IRMutatorWithAnalyzer { } if (op->op.same_as(builtin::tvm_access_ptr())) { auto data = op->args[1].as(); - if (data && anchors_.count(data.value())) { - Buffer anchor = anchors_[data.value()]; - ICHECK(op->args[0].dtype() == anchor->dtype) + if (data && matrix_views_.count(data.value())) { + const auto &views = matrix_views_[data.value()]; + Buffer matrix = views[0]; + for (const Buffer &view : views) + ICHECK(layouts_[matrix]->IsEqual(layouts_[view].get())) + << "Raw Cube pointers cannot select among distinct view layouts; " + "use T.access_ptr(view[...], ...)"; + ICHECK(op->args[0].dtype() == matrix->dtype) << "Cube raw pointer must use the canonical operand dtype " - << anchor->dtype << "; take the pointer through that view"; - Buffer target = canonical_buffers_[data.value()]; + << matrix->dtype << "; take the pointer through that view"; + Buffer target = buffers_[matrix]; PrimExpr remaining = VisitExpr(op->args[2]); - Array indices(anchor->shape.size(), Integer(0)); - for (int i = static_cast(anchor->shape.size()) - 1; i >= 0; --i) { + Array indices(matrix->shape.size(), Integer(0)); + for (int i = static_cast(matrix->shape.size()) - 1; i >= 0; --i) { indices.Set(i, i == 0 ? remaining - : floormod(remaining, anchor->shape[i])); - remaining = floordiv(remaining, anchor->shape[i]); + : floormod(remaining, matrix->shape[i])); + remaining = floordiv(remaining, matrix->shape[i]); } // Preserve the caller's extent; only the logical base address changes. return Call(op->dtype, tl::access_ptr(), @@ -512,17 +543,16 @@ class FractalStorageNormalizer : public arith::IRMutatorWithAnalyzer { for (const auto &[buffer, layout] : value.value().cast>()) { Buffer rewritten = buffers_.Get(buffer).value_or(buffer); - layouts.Set(rewritten, - canonical_layouts_.Get(rewritten).value_or(layout)); + layouts.Set(rewritten, layouts_.Get(rewritten).value_or(layout)); } block.CopyOnWrite()->annotations.Set(attr::kLayoutMap, layouts); } return block; } - Map anchors_; - Map canonical_buffers_; + Map> matrix_views_; + Map> logical_shapes_; Map buffers_; - LayoutMap canonical_layouts_; + LayoutMap layouts_; Map logical_maps_; }; } // namespace diff --git a/testing/ascend/auto_schedule/test_cube_storage_views.py b/testing/ascend/auto_schedule/test_cube_storage_views.py new file mode 100644 index 0000000000..b93370591d --- /dev/null +++ b/testing/ascend/auto_schedule/test_cube_storage_views.py @@ -0,0 +1,153 @@ +"""Cube views preserve distinct layouts on one versioned storage allocation.""" + +import pytest +import tilelang +import tilelang.ascend.language as T +import tilelang.testing +from tilelang import tvm +from tilelang.layout import make_ascend_major_k_layout, make_ascend_major_mn_layout +from tvm import tirx +from testing.ascend._ir import nodes + + +def _retimed_layout_views(qk_steps=4, pv_steps=4, alias_kind="view", transpose_pv=False, annotated=True): + @T.prim_func + def main( + Q: T.Tensor((64, 128), "bfloat16"), + K: T.Tensor((64, 128), "bfloat16"), + CQ: T.Tensor((3, 64, 64), "float32"), + CV: T.Tensor((3, pv_steps, 64, 128), "float32"), + ): + with T.Kernel(1): + q = T.alloc_l1((64, 128), "bfloat16") + k0 = T.alloc_l1((32, 128), "bfloat16") + k1 = T.alloc_l1((32, 128), "bfloat16") + aq = T.alloc_l0a((64, 128), "bfloat16") + bq = T.alloc_l0b((64, 128), "bfloat16") + if alias_kind == "view": + bv = T.view(bq, (128, 64) if transpose_pv else (64, 128)) + else: + bv = T.reshape(bq, (128, 64) if transpose_pv else (64, 128)) + cq = T.alloc_l0c((64, 64), "float32") + cv = T.alloc_l0c((pv_steps, 64, 128), "float32") + T.annotate_buffer_versions({q: 1, k0: 1, k1: 1, aq: 4, bq: 4, cq: 1, cv: 1}) + if annotated: + T.annotate_layout( + { + bq: make_ascend_major_k_layout(bq), + bv: make_ascend_major_k_layout(bv) if transpose_pv else make_ascend_major_mn_layout(bv), + } + ) + T.copy(Q, q) + T.copy(K[:32, :], k0) + T.copy(K[32:, :], k1) + # Keep the Cube mode fixed so this regression isolates storage + # ownership from separately scheduled configuration changes. + for step in T.Pipelined(3, num_stages=2, annotations={"enable_offset": True}): + with T.Stage(0): + for d in T.Pipelined(qk_steps, num_stages=4): + with T.Stage(0): + T.copy(q, aq) + with T.Task(): + T.copy(k0, bq[:32, :], l0_pitch=4) + T.copy(k1, bq[32:, :], l0_pitch=4) + T.gemm(aq, bq, cq, transpose_B=True, clear_accum=d == 0) + T.copy(cq, CQ[step, :, :]) + with T.Stage(1): + for _region in T.Pipelined(1, num_stages=4): + with T.Stage(0): + T.copy(q[:, :64], aq[:, :64]) + for d in T.Pipelined(pv_steps, num_stages=2): + with T.Stage(0): + with T.Task(): + if transpose_pv: + T.copy(k0, bv[:, :32], transpose=True, l0_pitch=8) + T.copy(k1, bv[:, 32:], transpose=True, l0_pitch=8) + else: + T.copy(k0, bv[:32, :], l0_pitch=8) + T.copy(k1, bv[32:, :], l0_pitch=8) + T.gemm(aq[:, :64], bv, cv[d, :, :], transpose_B=transpose_pv, clear_accum=True) + for d in T.Pipelined(pv_steps, num_stages=2): + with T.Stage(0): + T.copy(cv[d, :, :], CV[step, d, :, :]) + + return main + + +@pytest.mark.parametrize("alias_kind", ["view", "reshape"]) +@pytest.mark.parametrize("transpose_pv,annotated", [(False, True), (True, True), (False, False), (True, False)]) +def test_storage_views_keep_layouts_and_share_one_version_clock(alias_kind, transpose_pv, annotated): + snapshots = {} + + @tvm.instrument.pass_instrument + class Capture: + def run_after_pass(self, mod, info): + if info.name in ( + "tl.NormalizeAscendFractalStorage", + "tl.PrepareMultiBuffer", + "tl.MaterializeMultiBuffer", + "tl.MergeUBAllocations", + ): + snapshots[info.name] = mod + + with tvm.target.Target("ascend"), tvm.transform.PassContext(config={"tl.disable_shared_memory_reuse": True}, instruments=[Capture()]): + tilelang.lower(_retimed_layout_views(3, 2, alias_kind, transpose_pv, annotated), target="ascend", enable_device_compile=False) + + def l0b_layouts(mod): + return { + buffer: layout + for block in nodes(mod, tirx.SBlock) + for buffer, layout in block.annotations.get("layout_map", {}).items() + if buffer.scope() == "shared.l0b" + } + + layouts = l0b_layouts(snapshots["tl.NormalizeAscendFractalStorage"]) + assert len(layouts) == 2 + first, second = layouts + assert first.data.same_as(second.data) + assert not layouts[first].is_equal(layouts[second]) + expected_shapes = {(64, 128), (128, 64)} if transpose_pv else {(64, 128)} + assert {tuple(int(x) for x in buffer.shape) for buffer in layouts} == expected_shapes + + clocks = [ + counter + for loop in nodes(snapshots["tl.PrepareMultiBuffer"], tirx.For) + for storage, counter in loop.annotations.get("tl.multi_buffer_counter_map", {}).items() + if storage.same_as(first.data) + ] + assert len(clocks) == 2 and clocks[0].same_as(clocks[1]) + expanded = {buffer: layout for buffer, layout in l0b_layouts(snapshots["tl.MaterializeMultiBuffer"]).items() if len(buffer.shape) == 3} + assert len(expanded) == 2 + for buffer, layout in expanded.items(): + assert int(buffer.shape[0]) == 4 and int(buffer.strides[0]) == 8192 + expected = make_ascend_major_k_layout(buffer) if transpose_pv or buffer.name == "bq" else make_ascend_major_mn_layout(buffer) + assert layout.is_equal(expected) + allocations = [ + node.buffer + for func in snapshots["tl.MergeUBAllocations"].functions.values() + for node in nodes(func, tirx.AllocBuffer) + if node.buffer.scope().startswith("shared.l0b") + ] + assert len(allocations) == 1 + assert int(allocations[0].shape[0]) * tvm.DataType(allocations[0].dtype).bits // 8 == 65536 + + +@pytest.mark.parametrize("alias_kind,transpose_pv,annotated", [("view", False, True), ("reshape", True, True), ("view", True, False)]) +@tilelang.testing.requires_ascend +def test_view_values_across_owner_lengths_and_launches(alias_kind, transpose_pv, annotated): + import torch + + torch.npu.set_device(0) + torch.npu.set_op_timeout_ms(30000) + with tvm.target.Target("ascend"): + kernel = tilelang.compile(_retimed_layout_views(3, 2, alias_kind, transpose_pv, annotated), target="ascend", out_idx=[2, 3]) + for seed in (7, 11, 19): + torch.manual_seed(seed) + q = torch.randint(-2, 3, (64, 128), dtype=torch.int32).to(dtype=torch.bfloat16, device="npu") + k = torch.randint(-2, 3, (64, 128), dtype=torch.int32).to(dtype=torch.bfloat16, device="npu") + cq, cv = kernel(q, k) + torch.npu.synchronize() + expected_qk = 3 * (q.cpu().float() @ k.cpu().float().T) + expected_pv = q.cpu().float()[:, :64] @ k.cpu().float() + torch.testing.assert_close(cq.cpu(), expected_qk.expand(3, 64, 64), rtol=0, atol=0) + torch.testing.assert_close(cv.cpu(), expected_pv.expand(3, 2, 64, 128), rtol=0, atol=0) diff --git a/testing/ascend/transform/test_ascend_buffer_view_layouts.py b/testing/ascend/transform/test_ascend_buffer_view_layouts.py new file mode 100644 index 0000000000..8ff7ae1ad1 --- /dev/null +++ b/testing/ascend/transform/test_ascend_buffer_view_layouts.py @@ -0,0 +1,114 @@ +"""Layouts belong to buffer views; pointers and allocations retain that split.""" + +import pytest +import tilelang +import tilelang.testing +from tilelang import tvm +from tilelang.ascend import language as T, transform +from tilelang.layout import Layout +from tvm import tirx + +from testing.ascend._ir import calls, kernel + + +def _views(): + owner = tirx.decl_buffer((8, 16), "float32", name="owner", scope="shared") + view = tirx.decl_buffer(owner.shape, owner.dtype, name="view", data=owner.data) + return owner, view, Layout(owner.shape, lambda i, j: (i, j)), Layout(view.shape, lambda i, j: (j, i)) + + +@pytest.mark.parametrize("region", [None, "SIMT_VF", "SIMD_VF"]) +def test_explicit_layouts_remain_attached_to_each_view(region): + owner, view, first, second = _views() + body = tirx.SeqStmt([tirx.Evaluate(owner[3, 5]), tirx.Evaluate(view[3, 5])]) + annotations = {"layout_map": {owner: first, view: second}} + if region: + if region == "SIMT_VF": + tx = tirx.IterVar(tvm.ir.Range(0, 32), tirx.Var("tx", "int32"), tirx.IterVar.ThreadIndex, "threadIdx.x") + body = tirx.AttrStmt(tx, "thread_extent", 32, body) + body = tirx.SBlockRealize([], True, tirx.SBlock([], [], [], region, body, annotations=annotations)) + annotations = {} + mod = kernel(body, buffers=[owner], annotations=annotations) + with tvm.target.Target("ascend"): + inferred = transform.AscendLayoutInference()(mod) + layouts = inferred["main"].body.block.annotations["layout_map"] + assert layouts[owner].is_equal(first) + assert layouts[view].is_equal(second) + + +def test_unconstrained_reshape_inherits_the_unique_layout(): + owner, _, layout, _ = _views() + view = tirx.decl_buffer((4, 32), owner.dtype, name="reshaped", data=owner.data) + mod = kernel(tirx.Evaluate(view[1, 7]), buffers=[owner], annotations={"layout_map": {owner: layout}}) + with tvm.target.Target("ascend"): + inferred = transform.AscendLayoutInference()(mod) + layouts = inferred["main"].body.block.annotations["layout_map"] + assert layouts[view].is_equal(layout.reshape(view.shape)) + + +def test_unconstrained_access_cannot_choose_between_distinct_layouts(): + owner, view, first, second = _views() + unknown = tirx.decl_buffer(owner.shape, owner.dtype, name="unknown", data=owner.data) + body = tirx.SeqStmt([tirx.Evaluate(buf[0, 0]) for buf in (owner, view, unknown)]) + mod = kernel(body, buffers=[owner], annotations={"layout_map": {owner: first, view: second}}) + with tvm.target.Target("ascend"), pytest.raises(tvm.error.InternalError, match="Ambiguous layout for buffer view unknown"): + transform.AscendLayoutInference()(mod) + + +def test_unused_storage_owner_does_not_need_a_view_layout(): + owner, first_view, _, _ = _views() + second_view = tirx.decl_buffer(owner.shape, owner.dtype, name="second", data=owner.data) + first = Layout(first_view.shape, lambda i, j: (i, j)) + second = Layout(second_view.shape, lambda i, j: (j, i)) + body = tirx.SeqStmt([tirx.Evaluate(buf[0, 0]) for buf in (first_view, second_view)]) + mod = kernel(body, buffers=[owner], annotations={"layout_map": {first_view: first, second_view: second}}) + with tvm.target.Target("ascend"): + inferred = transform.AscendLayoutInference()(mod) + layouts = inferred["main"].body.block.annotations["layout_map"] + assert layouts[first_view].is_equal(first) + assert layouts[second_view].is_equal(second) + + +def test_view_layouts_must_preserve_the_storage_footprint(): + owner, view, first, _ = _views() + second = Layout(view.shape, lambda i, j: (i, j * 2)) + body = tirx.SeqStmt([tirx.Evaluate(buf[0, 0]) for buf in (owner, view)]) + mod = kernel(body, buffers=[owner], annotations={"layout_map": {owner: first, view: second}}) + with tvm.target.Target("ascend"), pytest.raises(tvm.error.InternalError, match="equal physical sizes"): + transform.NormalizeAscendFractalStorage()(transform.AscendLayoutInference()(mod)) + + +def test_typed_pointers_use_each_view_layout_and_one_storage(): + owner, view, first, second = _views() + body = tirx.SeqStmt([tirx.Evaluate(T.access_ptr(buf[3, 5], "r", extent=1)) for buf in (owner, view)]) + mod = kernel(body, buffers=[owner], annotations={"layout_map": {owner: first, view: second}}) + with tvm.target.Target("ascend"): + mod = transform.AscendLowerTileOp()(mod) + mod = tilelang.transform.LowerAccessPtr()(mod) + pointers = calls(mod, "tirx.tvm_access_ptr") + assert len(pointers) == 2 + assert pointers[0].args[1].same_as(pointers[1].args[1]) + analyzer = tvm.arith.Analyzer() + assert analyzer.can_prove_equal(pointers[0].args[2], 3 * 16 + 5) + assert analyzer.can_prove_equal(pointers[1].args[2], 5 * 8 + 3) + + +@tilelang.testing.requires_ascend +def test_ub_view_reads_the_same_bytes_through_its_own_layout(): + import torch + + @T.prim_func + def transpose(A: T.Tensor((8, 16), "float32"), B: T.Tensor((16, 8), "float32")): + with T.Kernel(1): + storage = T.alloc_shared((8, 16), "float32") + view = T.view(storage, (16, 8)) + T.annotate_layout({storage: Layout((8, 16), lambda i, j: (i, j)), view: Layout((16, 8), lambda i, j: (j, i))}) + with T.SimtVF(threads=128): + for i, j in T.Parallel(8, 16): + storage[i, j] = A[i, j] + for i, j in T.Parallel(16, 8): + B[i, j] = view[i, j] + + compiled = tilelang.compile(transpose, target="ascend", out_idx=-1) + source = torch.arange(128, dtype=torch.float32, device="npu").reshape(8, 16) + torch.testing.assert_close(compiled(source), source.T, rtol=0, atol=0) diff --git a/testing/ascend/transform/test_ascend_lower_copy.py b/testing/ascend/transform/test_ascend_lower_copy.py index 525f4848df..29b0170db4 100644 --- a/testing/ascend/transform/test_ascend_lower_copy.py +++ b/testing/ascend/transform/test_ascend_lower_copy.py @@ -3,7 +3,7 @@ import math import pytest -from tilelang import tvm +from tilelang import tvm, transform as tl_transform from tilelang.ascend import language as T, transform from tilelang.layout import ( make_ascend_nz_layout, @@ -52,7 +52,7 @@ def _lower(src, dst, src_region=None, dst_region=None, *, layouts=None, bindings mod = tvm.IRModule({"main": func}) if normalize: mod = transform.NormalizeAscendFractalStorage()(mod) - return transform.AscendLowerTileOp()(mod) + return tl_transform.LowerAccessPtr()(transform.AscendLowerTileOp()(mod)) def _equal(actual, expected): @@ -297,7 +297,7 @@ def test_explicit_l0_pitch_controls_offsets_and_keeps_plane_capacity(v_layout, o with pytest.raises(ValueError, match="l0_pitch"): transform.AscendLowerTileOp()(mod) return - after = transform.AscendLowerTileOp()(mod) + after = tl_transform.LowerAccessPtr()(transform.AscendLowerTileOp()(mod)) loads = calls(after, "tl.ascend_load_cbuf_to_cb" if v_layout else "tl.ascend_load_cbuf_to_ca") _equal([load.args[7] for load in loads], [pitch, pitch]) _equal([load.args[0].args[2] for load in loads], [4096, 4096 + 2 * pitch * 256]) diff --git a/testing/ascend/transform/test_ascend_normalize_fractal_storage.py b/testing/ascend/transform/test_ascend_normalize_fractal_storage.py index 869669c45d..3c7c9612fb 100644 --- a/testing/ascend/transform/test_ascend_normalize_fractal_storage.py +++ b/testing/ascend/transform/test_ascend_normalize_fractal_storage.py @@ -35,7 +35,7 @@ def test_canonical_allocation_rounds_only_fractal_axes(shape, dtype, padded): @pytest.mark.parametrize("kind", ["view", "reshape", "reinterpret"]) -def test_alias_accesses_use_one_canonical_allocation(kind): +def test_alias_accesses_preserve_views_on_one_allocation(kind): owner_shape = (15, 15) if kind == "view" else (225,) owner = tirx.decl_buffer(owner_shape, "uint16" if kind == "reinterpret" else "bfloat16", name="owner", scope="shared.l1") view = tirx.decl_buffer((15, 15), "bfloat16", data=owner.data) @@ -43,7 +43,8 @@ def test_alias_accesses_use_one_canonical_allocation(kind): after = transform.NormalizeAscendFractalStorage()(before) allocation = after["main"].body.block.alloc_buffers[0] (store,) = nodes(after, tirx.BufferStore) - assert store.buffer.same_as(allocation) + assert not store.buffer.same_as(allocation) + assert store.buffer.data.same_as(allocation.data) assert allocation.data.same_as(owner.data) assert allocation.dtype == "bfloat16" assert tuple(int(x) for x in allocation.shape) == (16, 16) @@ -61,7 +62,7 @@ def test_pointer_base_is_reexpressed_in_canonical_indices(raw): after = transform.NormalizeAscendFractalStorage()(before) (pointer,) = calls(after, "tl.access_ptr") load, extent, _ = pointer.args - assert load.buffer.same_as(after["main"].body.block.alloc_buffers[0]) + assert load.buffer.data.same_as(after["main"].body.block.alloc_buffers[0].data) analyzer = tvm.arith.Analyzer() assert analyzer.can_prove_equal(load.indices[0], 1 if raw else row * 5) assert int(load.indices[1]) == 0 and int(extent) == 1 @@ -126,7 +127,7 @@ def test_alias_indices_and_access_regions_use_the_same_padded_storage(): root = normalized["main"].body.block owner = root.alloc_buffers[0] alias = root.reads[0].buffer - assert owner.same_as(alias) + assert not owner.same_as(alias) assert owner.data.same_as(alias.data) assert tuple(int(x) for x in owner.shape) == tuple(int(x) for x in alias.shape) == (16, 16) assert [(int(r.min), int(r.extent)) for r in root.reads[0].region] == [(1, 2), (0, 15)] @@ -158,7 +159,110 @@ def test_alias_normalization_preserves_allocation_properties(): assert allocation.offset_factor == 8 assert allocation.dtype == "bfloat16" assert tuple(int(x) for x in allocation.shape) == (16, 16) - assert allocation.same_as(normalized["main"].body.block.body.value.buffer) + view_buffer = normalized["main"].body.block.body.value.buffer + assert not allocation.same_as(view_buffer) + assert allocation.data.same_as(view_buffer.data) + assert allocation.name == "owner" and view_buffer.name == "view" + + +@pytest.mark.parametrize( + "shape,view_shape,transpose", [((64, 128), (64, 128), False), ((64, 128), (128, 64), True), ((17, 32), (32, 17), False)] +) +def test_independent_cube_views_preserve_their_matrix_geometry(shape, view_shape, transpose): + owner = tirx.decl_buffer(shape, "bfloat16", name="owner", scope="shared.l0b") + view = tirx.decl_buffer(view_shape, "bfloat16", name="view", data=owner.data) + layouts = { + owner: tilelang.layout.make_ascend_major_k_layout(owner), + view: (tilelang.layout.make_ascend_major_k_layout if transpose else tilelang.layout.make_ascend_major_mn_layout)(view), + } + body = tirx.SeqStmt([tirx.Evaluate(owner[1, 0]), tirx.BufferStore(view, tirx.const(7, "bfloat16"), [0, 1])]) + block = tirx.SBlock([], [], [], "root", body, alloc_buffers=[owner], annotations={"layout_map": layouts}) + before = tvm.IRModule({"main": tirx.PrimFunc([], tirx.SBlockRealize([], True, block))}) + after = transform.NormalizeAscendFractalStorage()(before) + root = after["main"].body.block + read, write = root.body.seq[0].value, root.body.seq[1] + assert not read.buffer.same_as(write.buffer) + assert read.buffer.data.same_as(write.buffer.data) + assert tuple(int(x) for x in read.indices) == (1, 0) + assert tuple(int(x) for x in write.indices) == (0, 1) + assert tuple(int(x) for x in read.buffer.shape) == tuple((x + 15) // 16 * 16 for x in shape) + assert tuple(int(x) for x in write.buffer.shape) == tuple((x + 15) // 16 * 16 for x in view_shape) + assert len(root.alloc_buffers) == 1 + assert not root.annotations["layout_map"][read.buffer].is_equal(root.annotations["layout_map"][write.buffer]) + tvm.ir.assert_structural_equal(after, transform.NormalizeAscendFractalStorage()(after)) + + +@pytest.mark.parametrize( + "dtype,shape,view_shape,view_layout", + [ + ("bfloat16", (15, 16), (8, 30), make_ascend_nz_layout), + # Both unpadded layouts cover 128 elements, but their Cube padding differs. + ("float32", (8, 16), (8, 16), make_ascend_major_mn_layout), + ], +) +def test_cube_views_with_different_padded_sizes_are_rejected(dtype, shape, view_shape, view_layout): + owner = tirx.decl_buffer(shape, dtype, scope="shared.l0b") + view = tirx.decl_buffer(view_shape, dtype, data=owner.data) + block = tirx.SBlock( + [], + [], + [], + "root", + tirx.Evaluate(view[0, 0]), + alloc_buffers=[owner], + annotations={"layout_map": {owner: make_ascend_nz_layout(owner), view: view_layout(view)}}, + ) + before = tvm.IRModule({"main": tirx.PrimFunc([], tirx.SBlockRealize([], True, block))}) + with pytest.raises(tvm.error.InternalError, match="equal physical sizes including padding"): + transform.NormalizeAscendFractalStorage()(before) + + +@pytest.mark.parametrize("source", ["owner", "view"]) +@pytest.mark.parametrize("reverse_annotations", [False, True]) +def test_reshape_uses_its_layout_when_storage_has_distinct_matrix_views(source, reverse_annotations): + owner = tirx.decl_buffer((17, 32), "bfloat16", name="owner", scope="shared.l0b") + view = tirx.decl_buffer((32, 17), "bfloat16", name="view", data=owner.data) + reshaped = tirx.decl_buffer((544,), "bfloat16", name="reshaped", data=owner.data) + layouts = {owner: make_ascend_nz_layout(owner), view: make_ascend_major_mn_layout(view)} + matrix = owner if source == "owner" else view + layouts[reshaped] = layouts[matrix].reshape(reshaped.shape) + if reverse_annotations: + layouts = dict(reversed(list(layouts.items()))) + body = tirx.SeqStmt([tirx.Evaluate(buf[1, 0]) for buf in (owner, view)] + [tirx.Evaluate(reshaped[35])]) + before = kernel(body, buffers=[owner], annotations={"layout_map": layouts}) + after = transform.NormalizeAscendFractalStorage()(before) + root = after["main"].body.block + read = root.body.seq[2].value + assert tuple(int(x) for x in read.indices) == ((1, 3) if source == "owner" else (2, 1)) + assert tuple(int(x) for x in read.buffer.shape) == (32, 32) + assert read.buffer.name == "reshaped" + assert not read.buffer.same_as(root.alloc_buffers[0]) + assert all(stmt.value.buffer.data.same_as(root.alloc_buffers[0].data) for stmt in root.body.seq) + expected = (make_ascend_nz_layout if source == "owner" else make_ascend_major_mn_layout)(read.buffer) + assert root.annotations["layout_map"][read.buffer].is_equal(expected) + tvm.ir.assert_structural_equal(after, transform.NormalizeAscendFractalStorage()(after)) + + +@pytest.mark.parametrize("raw", [False, True], ids=["typed-view", "ambiguous-data-pointer"]) +def test_distinct_cube_layouts_require_typed_pointers(raw): + owner = tirx.decl_buffer((64, 128), "bfloat16", scope="shared.l0b") + view = tirx.decl_buffer((128, 64), "bfloat16", data=owner.data) + pointer = owner.access_ptr("r", offset=64, extent=1) if raw else T.access_ptr(view[1, 0], "r", extent=1) + before = kernel( + tirx.Evaluate(tirx.call_extern("int32", "consume", pointer)), + buffers=[owner], + annotations={ + "layout_map": {owner: tilelang.layout.make_ascend_major_k_layout(owner), view: tilelang.layout.make_ascend_major_k_layout(view)} + }, + ) + if raw: + with pytest.raises(tvm.error.InternalError, match="Raw Cube pointers cannot select among distinct view layouts"): + transform.NormalizeAscendFractalStorage()(before) + else: + after = transform.NormalizeAscendFractalStorage()(before) + (pointer,) = calls(after, "tl.access_ptr") + assert tuple(int(x) for x in pointer.args[0].buffer.shape) == (128, 64) + assert tuple(int(x) for x in pointer.args[0].indices) == (1, 0) def test_non_rectangular_alias_region_is_rejected(): diff --git a/tilelang/ascend/language/__init__.py b/tilelang/ascend/language/__init__.py index 423d829d51..3bd88062d3 100644 --- a/tilelang/ascend/language/__init__.py +++ b/tilelang/ascend/language/__init__.py @@ -18,6 +18,7 @@ alloc_l0c, ) from .annotations import ( # noqa: F401 + annotate_layout, annotate_buffer_versions, annotate_manual_multi_buffer, annotate_unlimit_memory, diff --git a/tilelang/ascend/language/annotations.py b/tilelang/ascend/language/annotations.py index d53216bb86..5bde41391d 100644 --- a/tilelang/ascend/language/annotations.py +++ b/tilelang/ascend/language/annotations.py @@ -4,9 +4,12 @@ from tvm.tirx import IntImm from tvm.tirx.script.builder.ir import sblock_attr +from tilelang.language.annotations import _normalize_layout_map + _VALID_MEMORY_SCOPES = ("shared", "shared.l1", "shared.l0c", "shared.l0a", "shared.l0b") __all__ = [ + "annotate_layout", "annotate_buffer_versions", "annotate_manual_multi_buffer", "annotate_unlimit_memory", @@ -16,6 +19,17 @@ _VALID_BUFFER_VERSION_MODES = ("auto", "iteration", "counter") +def annotate_layout(layout_map: dict): + """Annotate each Buffer view, preserving aliases with distinct Cube layouts. + + Views share storage, not a layout conversion. A producer must write the + representation consumed through each view; annotating a view never moves + or transposes existing values. Canonical Cube views sharing storage must + have the same dtype and padded physical size. + """ + return sblock_attr({"layout_map": _normalize_layout_map(layout_map)}) + + def annotate_buffer_versions(buffer_versions_map: dict): """Control automatic multi-buffer version counts and indexing modes. diff --git a/tilelang/ascend/op/gemm/gemm_mad.py b/tilelang/ascend/op/gemm/gemm_mad.py index 781b2731aa..004250390b 100644 --- a/tilelang/ascend/op/gemm/gemm_mad.py +++ b/tilelang/ascend/op/gemm/gemm_mad.py @@ -20,29 +20,6 @@ def _to_int32(value): return tirx.IntImm("int32", value) if isinstance(value, int) else value -def _compute_flat_offset(buf, region): - offset = tirx.const(0, "int32") - ndim = len(region.region) - for i, rng in enumerate(region.region): - stride = tirx.const(1, "int32") - for j in range(i + 1, ndim): - stride = stride * buf.shape[j] - offset = offset + tirx.Cast("int32", rng.min) * stride - return offset - - -def _compute_flat_offset_excluding_last(buf, region): - offset = tirx.const(0, "int32") - ndim = len(region.region) - for i in range(ndim - 1): - rng = region.region[i] - stride = tirx.const(1, "int32") - for j in range(i + 1, len(buf.shape)): - stride = stride * buf.shape[j] - offset = offset + tirx.Cast("int32", rng.min) * stride - return offset - - def _compute_extent(region): extent = tirx.const(1, "int32") for rng in region.region: @@ -50,30 +27,21 @@ def _compute_extent(region): return extent -def _make_access_ptr(buf, region, rw_mask): - return tirx.op.tvm_access_ptr( - tirx.op.type_annotation(buf.dtype), - buf.data, - _compute_flat_offset(buf, region), +def _make_access_ptr(buf, region, rw_mask, exclude_last=False): + indices = [tirx.Cast("int32", rng.min) for rng in region.region] + if exclude_last: + indices[-1] = tirx.const(0, "int32") + return tirx.call_intrin( + "handle", + tirx.op.Op.get("tl.access_ptr"), + tirx.BufferLoad(buf, indices), _compute_extent(region), tirx.IntImm("int32", rw_mask), ) def _find_layout(layout_map, buf): - if layout_map is None: - return None - for key in (buf.data if hasattr(buf, "data") else None, buf): - if key is not None: - layout = layout_map.get(key) - if layout is not None: - return layout - for key, layout in layout_map.items(): - if hasattr(key, "data") and hasattr(buf, "data") and key.data.same_as(buf.data): - return layout - if hasattr(key, "name") and getattr(key, "name", None) == getattr(buf, "name", None): - return layout - return None + return None if layout_map is None else layout_map.get(buf) _ASCEND_DTYPE_MAP = { @@ -251,22 +219,8 @@ def _lower_l1(self): tile_k_sub = self._compute_tile_k_sub() c_ptr = _make_access_ptr(self.CRegion.buffer, self.CRegion, 2) - a_buf = self.ARegion.buffer - b_buf = self.BRegion.buffer - a_ptr = tirx.op.tvm_access_ptr( - tirx.op.type_annotation(a_buf.dtype), - a_buf.data, - _compute_flat_offset_excluding_last(a_buf, self.ARegion), - _compute_extent(self.ARegion), - tirx.IntImm("int32", 1), - ) - b_ptr = tirx.op.tvm_access_ptr( - tirx.op.type_annotation(b_buf.dtype), - b_buf.data, - _compute_flat_offset_excluding_last(b_buf, self.BRegion), - _compute_extent(self.BRegion), - tirx.IntImm("int32", 1), - ) + a_ptr = _make_access_ptr(self.ARegion.buffer, self.ARegion, 1, exclude_last=True) + b_ptr = _make_access_ptr(self.BRegion.buffer, self.BRegion, 1, exclude_last=True) input_dtype = self._input_dtype() dtype_str = _ASCEND_DTYPE_MAP.get(input_dtype) diff --git a/tilelang/ascend/op/gemm/gemm_mad_blockscaled.py b/tilelang/ascend/op/gemm/gemm_mad_blockscaled.py index 2661a0ac79..5e0fd8a91c 100644 --- a/tilelang/ascend/op/gemm/gemm_mad_blockscaled.py +++ b/tilelang/ascend/op/gemm/gemm_mad_blockscaled.py @@ -13,8 +13,6 @@ from .gemm_mad import ( _ASCEND_DTYPE_MAP, - _compute_extent, - _compute_flat_offset_excluding_last, _find_layout, _make_access_ptr, _to_int32, @@ -102,41 +100,15 @@ def _lower_l1_blockscaled(self): tile_k_sub = self._compute_tile_k_sub() c_ptr = _make_access_ptr(self.CRegion.buffer, self.CRegion, 2) - a_buf = self.ARegion.buffer - b_buf = self.BRegion.buffer - a_ptr = tirx.op.tvm_access_ptr( - tirx.op.type_annotation(a_buf.dtype), - a_buf.data, - _compute_flat_offset_excluding_last(a_buf, self.ARegion), - _compute_extent(self.ARegion), - tirx.IntImm("int32", 1), - ) - b_ptr = tirx.op.tvm_access_ptr( - tirx.op.type_annotation(b_buf.dtype), - b_buf.data, - _compute_flat_offset_excluding_last(b_buf, self.BRegion), - _compute_extent(self.BRegion), - tirx.IntImm("int32", 1), - ) + a_ptr = _make_access_ptr(self.ARegion.buffer, self.ARegion, 1, exclude_last=True) + b_ptr = _make_access_ptr(self.BRegion.buffer, self.BRegion, 1, exclude_last=True) sfa_region = self.SFARegion sfb_region = self.SFBRegion sfa_buf = sfa_region.buffer sfb_buf = sfb_region.buffer - sfa_ptr = tirx.op.tvm_access_ptr( - tirx.op.type_annotation(sfa_buf.dtype), - sfa_buf.data, - _compute_flat_offset_excluding_last(sfa_buf, sfa_region), - _compute_extent(sfa_region), - tirx.IntImm("int32", 1), - ) - sfb_ptr = tirx.op.tvm_access_ptr( - tirx.op.type_annotation(sfb_buf.dtype), - sfb_buf.data, - _compute_flat_offset_excluding_last(sfb_buf, sfb_region), - _compute_extent(sfb_region), - tirx.IntImm("int32", 1), - ) + sfa_ptr = _make_access_ptr(sfa_buf, sfa_region, 1, exclude_last=True) + sfb_ptr = _make_access_ptr(sfb_buf, sfb_region, 1, exclude_last=True) input_dtype = self._input_dtype() in_dtype_str = _ASCEND_DTYPE_MAP.get(input_dtype) diff --git a/tilelang/language/annotations.py b/tilelang/language/annotations.py index 65965637ec..9d72b9b6bc 100644 --- a/tilelang/language/annotations.py +++ b/tilelang/language/annotations.py @@ -42,8 +42,7 @@ def use_swizzle(panel_size: int, order: str = "row", enable: bool = True): return attr(None, "threadblock_swizzle_pattern", tvm_tuple(device_func, panel_size)) -def annotate_layout(layout_map: dict): - """Annotate the layout of the buffer.""" +def _normalize_layout_map(layout_map: dict): _layout_map = {} for buffer, layout in layout_map.items(): if is_reducer(buffer): @@ -57,13 +56,18 @@ def annotate_layout(layout_map: dict): elif is_fragment(buffer): assert isinstance(layout, Fragment), f"for Fragment {buffer}, layout must be a Fragment, but got {type(layout)}" if isinstance(layout, Layout): - _layout_map[buffer.data] = layout + _layout_map[buffer] = layout elif isinstance(layout, Callable): - _layout_map[buffer.data] = Layout(buffer.shape, layout) + _layout_map[buffer] = Layout(buffer.shape, layout) else: raise ValueError(f"Invalid layout: {layout}") - return sblock_attr({"layout_map": _layout_map}) + return _layout_map + + +def annotate_layout(layout_map: dict): + """Annotate the layout of the buffer.""" + return sblock_attr({"layout_map": {buffer.data: layout for buffer, layout in _normalize_layout_map(layout_map).items()}}) def annotate_safe_value(safe_value_map: dict):