[Ascend] Preserve physical views and pointer units during layout lowering

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<Map<Buffer, Layout>>()
                             .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);
@@ -585,12 +594,33 @@ class LowerTileOpPass : arith::IRMutatorWithAnalyzer {
       // Get original and new buffer shapes.
       Array<PrimExpr> old_shape = original_buffer->shape;
       Array<PrimExpr> 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<PrimExpr> zero_indices;
       for (size_t i = 0; i < old_shape.size(); ++i) {
         zero_indices.push_back(make_zero(elem_offset.dtype()));
       }
       Array<PrimExpr> 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<int>(new_shape.size()) - 1; i >= 0; --i) {
@@ -1081,6 +1111,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<BufferLoadNode>()) {
       bind_var_to_expr_[op->var] = value;
@@ -1089,10 +1120,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<Stmt>(op);
     } else {
       auto n = this->CopyOnWrite(op);
+      n->var = var;
       n->value = value;
       return Stmt(n);
     }

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)

testing/ascend/transform/test_ascend_lower_copy.py

@@ -566,6 +566,86 @@ 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(
+    "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")