From 8e528e57370d6b2c972b7341f0eb56dea5eb8949 Mon Sep 17 00:00:00 2001 From: Denver Jin Date: Wed, 21 Oct 2026 11:51:18 +0800 Subject: [PATCH 1/2] [Ascend] Preserve physical allocation shapes and address-backed views --- src/ascend/transform/lower_tile_op.cc | 17 ++++++++-- .../test_schedule_integration.py | 10 +++++- .../transform/test_ascend_lower_copy.py | 32 +++++++++++++++++++ 3 files changed, 55 insertions(+), 4 deletions(-) diff --git a/src/ascend/transform/lower_tile_op.cc b/src/ascend/transform/lower_tile_op.cc index 82a52a4389..ef5b3735e0 100644 --- a/src/ascend/transform/lower_tile_op.cc +++ b/src/ascend/transform/lower_tile_op.cc @@ -347,10 +347,14 @@ class LowerTileOpPass : arith::IRMutatorWithAnalyzer { // Record the mapping from buffer data var to buffer for later lookup for (auto buffer : op->alloc_buffers) { - buffer_map_.insert({buffer->data, buffer}); + // An enclosing layout annotation can still name the logical buffer + // before multi-buffer materialization. The allocation owns the physical + // shape, including version axes, used to decode flat pointer offsets. + buffer_map_.insert_or_assign(buffer->data, buffer); } for (auto match_buffer : op->match_buffers) { - buffer_map_.insert({match_buffer->buffer->data, match_buffer->buffer}); + buffer_map_.insert_or_assign(match_buffer->buffer->data, + match_buffer->buffer); buffer_data_to_buffer_.Set(match_buffer->buffer->data, match_buffer->buffer); } @@ -369,6 +373,11 @@ class LowerTileOpPass : arith::IRMutatorWithAnalyzer { .as>() .value(); for (auto [buffer, layout] : layout_map) { + // Address-backed views have layouts but no allocation or match_buffer. + // Register them so DMA pointers receive the same layout/pitch remap. + buffer_map_.insert({buffer->data, buffer}); + if (!buffer_data_to_buffer_.count(buffer->data)) + buffer_data_to_buffer_.Set(buffer->data, buffer); buffer_remap_.Set(buffer, makeBufferWithLayout(buffer, layout, var_remap_)); layout_map_.Set(buffer, layout); @@ -1081,6 +1090,7 @@ class LowerTileOpPass : arith::IRMutatorWithAnalyzer { Stmt VisitStmt_(const BindNode *op) final { PrimExpr value = this->VisitExpr(op->value); + Var var = var_remap_.Get(op->var).value_or(op->var); bool recorded = false; if (value->IsInstance()) { bind_var_to_expr_[op->var] = value; @@ -1089,10 +1099,11 @@ class LowerTileOpPass : arith::IRMutatorWithAnalyzer { if (SideEffect(value) <= CallEffectKind::kPure) { analyzer_->Bind(op->var, value); } - if (value.same_as(op->value)) { + if (value.same_as(op->value) && var.same_as(op->var)) { return GetRef(op); } else { auto n = this->CopyOnWrite(op); + n->var = var; n->value = value; return Stmt(n); } diff --git a/testing/ascend/auto_schedule/test_schedule_integration.py b/testing/ascend/auto_schedule/test_schedule_integration.py index 141a51872b..3b6b210632 100644 --- a/testing/ascend/auto_schedule/test_schedule_integration.py +++ b/testing/ascend/auto_schedule/test_schedule_integration.py @@ -131,7 +131,7 @@ def test_sibling_l0_rings_survive_layout_lowering(k_tiles): @tvm.ir.instrument.pass_instrument class Capture: def run_after_pass(self, mod, info): - if info.name in ("tl.PrepareMultiBuffer", "tl.MaterializeMultiBuffer"): + if info.name in ("tl.PrepareMultiBuffer", "tl.MaterializeMultiBuffer", "tl.AscendLowerTileOp"): snapshots[info.name] = mod with tvm.transform.PassContext(instruments=[Capture()]): @@ -145,6 +145,12 @@ def run_after_pass(self, mod, info): assert all(loop.annotations["tl.multi_buffer_counter_map"][b.data].same_as(counters[0]) for loop in owners) materialized = snapshots["tl.MaterializeMultiBuffer"] assert int(allocated_buffer(materialized, "l0a").shape[0]) == 2 + lowered = snapshots["tl.AscendLowerTileOp"] + physical_a = allocated_buffer(lowered, "l0a") + offsets = [call.args[2] for call in calls(lowered, "tirx.tvm_access_ptr") if call.args[1].same_as(physical_a.data)] + # Stale logical layout entries must not erase the automatic version axis + # while decoding flat DMA/MMAD pointers into physical fractal coordinates. + assert offsets and any(not tvm.arith.Analyzer().can_prove_equal(offset, 0) for offset in offsets) assert artifact.kernel_source and artifact.device_mod @@ -548,6 +554,8 @@ def main(C: T.Tensor((5, 16, 16), "float32")): T.copy(ub, C[i, :, :]) progress += 1 + # L0C copies attach frontend layout hints; scheduling consumes the + # normalized Buffer-keyed map produced by layout inference. resolved = schedule(main, through="ResolveCore", auto_schedule=False, infer_layouts=True) progress = [buffer for block in nodes(resolved, tirx.SBlock) for buffer in block.alloc_buffers if buffer.name.startswith("progress")] assert len(progress) == 2 and not progress[0].data.same_as(progress[1].data) diff --git a/testing/ascend/transform/test_ascend_lower_copy.py b/testing/ascend/transform/test_ascend_lower_copy.py index 525f4848df..6acf0ddceb 100644 --- a/testing/ascend/transform/test_ascend_lower_copy.py +++ b/testing/ascend/transform/test_ascend_lower_copy.py @@ -566,6 +566,38 @@ def lower(): _equal(load.args[4:6], [tirx.floordiv(n + 15, 16), 1] if axis == "mn" else [1, n]) +@pytest.mark.parametrize("major_mn,pitch,offset", [(False, 4, 8704), (True, 8, 12288)], ids=["qk", "pv"]) +def test_address_backed_l0_view_remaps_partial_copy_origin(major_mn, pitch, offset): + """A view has no alloc_buffer entry, but its DMA origin still uses NZ.""" + import re + import tilelang + + @T.prim_func + def main(): + with T.Kernel(1): + src = T.alloc_l1((32, 128), "bfloat16") + storage = T.alloc_l0b((16384,), "bfloat16") + dst = T.make_tensor_from_addr( + T.access_ptr(storage[0], "rw", extent=16384), (2, 64, 128), "bfloat16", storage_scope="shared.l0b" + ) + T.annotate_layout( + { + src: make_ascend_major_k_layout(src), + dst: make_ascend_major_mn_layout(dst) if major_mn else make_ascend_major_k_layout(dst), + } + ) + with T.Cube(): + T.copy(src, dst[1, 32:, :], l0_pitch=pitch) + + with tvm.transform.PassContext(config={tilelang.PassConfigKey.TL_ENABLE_AUTO_SCHEDULE: False}): + lowered = tilelang.lower(main, target="ascend") + source = lowered.kernel_source + instruction = "asc_copy_l12l0b_transpose" if major_mn else "asc_copy_l12l0b" + assert re.search(rf"{instruction}\(.*dst \+ {offset}\)", source), source + # The remapped data variable must still be defined by the pointer binding. + assert re.search(r"__cb__ bfloat16_t \*dst =", source), source + + @pytest.mark.parametrize("normalize", [False, True]) def test_sf_tail_source_fits_its_fractal_allocation(normalize): src = tirx.decl_buffer((24, 4), "uint8", name="src", scope="shared.l1") From f3fad906edd3298221a533418b7eb07d8c6b476b Mon Sep 17 00:00:00 2001 From: Denver Jin Date: Sun, 18 Oct 2026 11:15:51 +0800 Subject: [PATCH 2/2] [Ascend] Preserve pointer element units when remapping layouts --- src/ascend/transform/lower_tile_op.cc | 23 ++++++++- .../transform/test_ascend_lower_copy.py | 48 +++++++++++++++++++ 2 files changed, 70 insertions(+), 1 deletion(-) diff --git a/src/ascend/transform/lower_tile_op.cc b/src/ascend/transform/lower_tile_op.cc index ef5b3735e0..d9ea5dedd7 100644 --- a/src/ascend/transform/lower_tile_op.cc +++ b/src/ascend/transform/lower_tile_op.cc @@ -594,12 +594,33 @@ class LowerTileOpPass : arith::IRMutatorWithAnalyzer { // Get original and new buffer shapes. Array old_shape = original_buffer->shape; Array new_shape = new_buffer->shape; + Layout access_layout = layout; + DataType access_dtype = access_ptr_call->args[0].dtype(); + int buffer_bits = + original_buffer->dtype.bits() * original_buffer->dtype.lanes(); + int access_bits = access_dtype.bits() * access_dtype.lanes(); + if (access_bits != buffer_bits) { + // A raw access pointer retains its element type but not its Buffer + // identity. Its offset is in that type's elements, even when the data + // Var resolves to a differently typed allocation. Reinterpret the + // logical layout before mapping the offset, preserving sub-element + // lanes and expressing the physical result in the pointer's units. + PrimExpr elements = Integer(1); + for (const PrimExpr &extent : old_shape) { + elements *= extent; + } + old_shape = {analyzer_->Simplify( + floordiv(elements * Integer(buffer_bits), Integer(access_bits)))}; + access_layout = layout->Reshape( + old_shape, analyzer_, Integer(buffer_bits), Integer(access_bits)); + new_shape = access_layout->OutputShape(); + } Array zero_indices; for (size_t i = 0; i < old_shape.size(); ++i) { zero_indices.push_back(make_zero(elem_offset.dtype())); } Array forward_indices = RemapAccessIndices( - zero_indices, old_shape, new_shape, layout, elem_offset); + zero_indices, old_shape, new_shape, access_layout, elem_offset); PrimExpr new_offset = 0; PrimExpr stride_offset = 1; for (int i = static_cast(new_shape.size()) - 1; i >= 0; --i) { diff --git a/testing/ascend/transform/test_ascend_lower_copy.py b/testing/ascend/transform/test_ascend_lower_copy.py index 6acf0ddceb..fbdd8d1482 100644 --- a/testing/ascend/transform/test_ascend_lower_copy.py +++ b/testing/ascend/transform/test_ascend_lower_copy.py @@ -598,6 +598,54 @@ def main(): assert re.search(r"__cb__ bfloat16_t \*dst =", source), source +@pytest.mark.parametrize( + "storage_dtype,view_dtype,view_cols", + [("float32", "bfloat16", 256), ("bfloat16", "float32", 64), ("float32", "float32", 64)], + ids=["narrow", "widen", "same-width-reshape"], +) +@pytest.mark.parametrize("prefix", [0, 64]) +@pytest.mark.parametrize("row,col", [(0, 0), (1, 16)]) +def test_dma_reinterpreted_view_prefix_uses_pointer_element_units(storage_dtype, view_dtype, view_cols, prefix, row, col): + """Both DMA directions must agree with the view's byte address.""" + import tilelang + + storage_bytes = tvm.DataType(storage_dtype).bits // 8 + view_bytes = tvm.DataType(view_dtype).bits // 8 + view_rows = 512 * storage_bytes // (view_cols * view_bytes) + + @T.prim_func + def main(A: T.Tensor((2, 32), view_dtype), O: T.Tensor((2, 32), view_dtype)): + with T.Kernel(1): + storage = T.alloc_shared((4, 128), storage_dtype) + view = T.view(storage, shape=(view_rows, view_cols), dtype=view_dtype) + T.annotate_layout({storage: tilelang.layout.Layout((4, 128), lambda r, c: r * 128 + c + prefix)}) + with T.Vector(): + T.copy(A, view[row : row + 2, col : col + 32]) + T.copy(view[row : row + 2, col : col + 32], O) + + captured = [] + + @tvm.instrument.pass_instrument + class Capture: + def run_after_pass(self, mod, info): + if info.name == "tl.AscendLowerTileOp": + captured.append(mod) + + with tvm.transform.PassContext(config={tilelang.PassConfigKey.TL_ENABLE_AUTO_SCHEDULE: False}, instruments=[Capture()]): + tilelang.lower(main, target="ascend") + + assert len(captured) == 1 + (load,) = calls(captured[0], "tl.ascend_copy_gm_to_ubuf") + (store,) = calls(captured[0], "tl.ascend_copy_ubuf_to_gm") + # This byte address comes directly from the storage prefix and the view's + # row/column coordinates, independently of the layout reshape machinery. + expected_bytes = prefix * storage_bytes + (row * view_cols + col) * view_bytes + for ptr in (load.args[0], store.args[1]): + assert ptr.op.name == "tirx.tvm_access_ptr" + assert ptr.args[0].dtype == view_dtype + _equal([ptr.args[2] * view_bytes], [expected_bytes]) + + @pytest.mark.parametrize("normalize", [False, True]) def test_sf_tail_source_fits_its_fractal_allocation(normalize): src = tirx.decl_buffer((24, 4), "uint8", name="src", scope="shared.l1")