[Ascend] Support native NZ output for L0C-to-UB copies

src/ascend/op/copy.cc

@@ -663,8 +663,15 @@ Stmt LowerDMACopy(const AscendCopyNode &op, const LowerArgs &T,
   } else if (dma_path == DMAPath::kL0CToUB) {
     StridedLayout src_layout = NormalizeTrailingMTE2DLayout(
         op.src, src_range, analyzer, "Ascend L0C->UB source");
-    StridedLayout dst_layout = NormalizeMTE2DLayout(
-        op.dst, dst_range, analyzer, "Ascend L0C->UB destination");
+    Optional<Layout> ub_layout = FindLayoutForBuffer(T.layout_map, op.dst);
+    bool nz_output =
+        ub_layout.defined() && IsAscendFractalLayout(ub_layout.value());
+    StridedLayout dst_layout =
+        nz_output
+            ? NormalizeTrailingMTE2DLayout(op.dst, dst_range, analyzer,
+                                           "Ascend NZ L0C->UB destination")
+            : NormalizeMTE2DLayout(op.dst, dst_range, analyzer,
+                                   "Ascend L0C->UB destination");
     const Mode &src_inner = src_layout.modes[0];
     const Mode &src_row = src_layout.modes[1];
     const Mode &dst_inner = dst_layout.modes[0];
@@ -721,6 +728,43 @@ Stmt LowerDMACopy(const AscendCopyNode &op, const LowerArgs &T,
       copy_inner = dst_inner.size;
       copy_rows = dst_row.size;
     }
+    if (nz_output) {
+      AscendFractalLayoutInfo dst_info;
+      ICHECK(
+          TryExtractAscendFractalLayout(ub_layout.value(), op.dst, &dst_info));
+      TVM_FFI_CHECK(dst_info.c0_axis == LogicalAxis::kCol &&
+                        dst_info.row16_axis == LogicalAxis::kRow &&
+                        analyzer->CanProveEqual(dst_info.c0, 16),
+                    ValueError)
+          << "Ascend NZ L0C->UB requires 16-column NZ panels on "
+          << op.dst->name << "; use make_ascend_l0c_layout.";
+      if (!op.dst->strides.empty()) {
+        PrimExpr stride = I(1);
+        for (int i = static_cast<int>(op.dst->shape.size()) - 1; i >= 0; --i) {
+          TVM_FFI_CHECK(analyzer->CanProveEqual(op.dst->shape[i], 1) ||
+                            analyzer->CanProveEqual(op.dst->strides[i], stride),
+                        ValueError)
+              << "Ascend NZ L0C->UB requires dense logical strides; express "
+                 "padding in the UB buffer shape instead of explicit strides.";
+          stride = stride * op.dst->shape[i];
+        }
+      }
+      With<arith::ConstraintContext> nonempty_copy(analyzer, has_data);
+      TVM_FFI_CHECK(
+          analyzer->CanProveEqual(FloorMod(dst_range.back()->min, I(16)), 0) &&
+              analyzer->CanProveEqual(FloorMod(dst_inner.size, I(16)), 0) &&
+              analyzer->CanProveEqual(FloorMod(copy_inner, I(16)), 0),
+          ValueError)
+          << "Ascend NZ L0C->UB requires a destination column origin and "
+             "copy widths aligned to 16 elements; got "
+          << dst_range.back() << " and source width " << copy_inner << ".";
+      // NZ pitch advances between 16-column panels, in destination elements.
+      // Use the allocation's physical shape, not this copy's row extent.
+      Array<PrimExpr> physical_shape = ub_layout.value()->OutputShape();
+      dst_row_stride = I(1);
+      for (size_t i = physical_shape.size() - 3; i < physical_shape.size(); ++i)
+        dst_row_stride = dst_row_stride * physical_shape[i];
+    }
     PrimExpr dual_dst_ctl_val = I(dual_dst_ctl);
     PrimExpr unit_flag_ctl_val = cast(DataType::Int(32), op.unit_flag_ctl);
     PrimExpr sub_blockid_val = cast(DataType::Int(32), op.sub_blockid);
@@ -739,7 +783,7 @@ Stmt LowerDMACopy(const AscendCopyNode &op, const LowerArgs &T,
                  I(GetCCQuantPre(op.src->dtype, op.dst->dtype)), // quant_pre
                  zero,                                           // relu_pre
                  zero,                                           // split_en
-                 I(1),                                           // NZ2ND_en
+                 I(nz_output ? 0 : 1),                           // NZ2ND_en
                  zero,                                           // quant_post
                  zero,                                           // relu_post
                  zero,   // clip_relu_post

src/ascend/transform/normalize_fractal_storage.cc

@@ -276,6 +276,11 @@ class FractalStorageNormalizer : public arith::IRMutatorWithAnalyzer {
   Call NormalizeFixpipeCopy_(Call call) {
     constexpr int kFixpipeAlignmentBytes = 32;
     AscendCopy rewritten(call->args, call->annotations);
+    // NZ output has a panel pitch and whole-column-panel requirements, not
+    // ND's row pitch and 32-byte row padding. Copy lowering validates it.
+    if (auto layout = canonical_layouts_.Get(rewritten->dst);
+        layout && IsAscendFractalLayout(layout.value()))
+      return call;
     ICHECK_GE(rewritten->src_range.size(), 2U);
     ICHECK_EQ(rewritten->src_range.size(), rewritten->src->shape.size());
     // Lowering skips the DMA if any source or destination extent is

testing/ascend/layout/test_ascend_fractal_storage_integration.py

@@ -5,10 +5,85 @@
 import tilelang.ascend.language as T
 import tilelang.testing
 from tilelang import tvm
+from tilelang.layout import make_ascend_l0c_layout
 from tvm import tirx
 from testing.ascend._ir import allocated_buffer, calls, nodes, statements
 
 
+def _nz_fixpipe_kernel(split, dtype="float32", n=64, alias=False):
+    rows = 16 if split == "m" else 32
+    cols = n // 2 if split == "n" else n
+
+    @T.prim_func
+    def main(A: T.Tensor((32, 16), "bfloat16"), B: T.Tensor((n, 16), "bfloat16"), C: T.Tensor((32, n), dtype)):
+        with T.MixedKernel(1, sids=1 if split == "single" else 2) as (_, sid):
+            a = T.alloc_l1((32, 16), "bfloat16")
+            b = T.alloc_l1((n, 16), "bfloat16")
+            acc = T.alloc_l0c((32, n), "float32")
+            if alias:
+                storage = T.alloc_shared((2 * 48 * 80,), dtype)
+                ub = T.view(storage, (2, 48, 80))
+            else:
+                ub = T.alloc_shared((2, 48, 80), dtype)
+            T.annotate_layout({ub: make_ascend_l0c_layout(ub)})
+            T.copy(A, a)
+            T.copy(B, b)
+            T.gemm(a, b, acc, transpose_B=True, clear_accum=True)
+            if split == "single":
+                T.copy(acc, ub[1, 16 : 16 + rows, 16 : 16 + cols])
+            else:
+                T.dual_copy(acc, ub[1, 16 : 16 + rows, 16 : 16 + cols])
+            with T.SimtVF(threads=128):
+                for i, j in T.Parallel(rows, cols):
+                    C[i + (sid * rows if split == "m" else 0), j + (sid * cols if split == "n" else 0)] = ub[1, 16 + i, 16 + j]
+
+    return main
+
+
+@pytest.mark.parametrize("alias", [False, True])
+def test_nz_fixpipe_layout_survives_full_lowering(alias):
+    snapshots = {}
+
+    @tvm.ir.instrument.pass_instrument
+    class Capture:
+        def run_after_pass(self, mod, info):
+            if info.name == "tl.AscendLowerTileOp":
+                snapshots[info.name] = mod
+
+    with tvm.target.Target("ascend"), tvm.transform.PassContext(instruments=[Capture()]):
+        tilelang.lower(_nz_fixpipe_kernel("m", alias=alias), target="ascend")
+    (copy,) = calls(snapshots["tl.AscendLowerTileOp"], "tl.ascend_copy_matrix_cc_to_ub")
+    assert int(copy.args[5]) == 48 * 16
+    assert int(copy.args[7]) == 1
+    assert int(copy.args[14]) == 0
+
+
+def test_nz_fixpipe_does_not_apply_nd_tail_padding():
+    with tvm.target.Target("ascend"), pytest.raises(ValueError, match="aligned to 16 elements"):
+        tilelang.lower(_nz_fixpipe_kernel("single", n=9), target="ascend")
+
+
+@tilelang.testing.requires_ascend
+@pytest.mark.parametrize(
+    "split,dtype,alias",
+    [("single", "float32", False), ("m", "float32", False), ("n", "float32", False), ("single", "bfloat16", False), ("m", "float32", True)],
+)
+def test_nz_fixpipe_logical_reads_npu(split, dtype, alias):
+    import torch
+
+    # Exact small inputs make addressing errors distinguishable from rounding.
+    generator = torch.Generator().manual_seed(42)
+    a = torch.randint(-4, 5, (32, 16), generator=generator).to(torch.bfloat16)
+    b = torch.randint(-4, 5, (64, 16), generator=generator).to(torch.bfloat16)
+    with tvm.target.Target("ascend"):
+        kernel = tilelang.compile(_nz_fixpipe_kernel(split, dtype, alias=alias), out_idx=-1, target="ascend")
+    expected = (a.float() @ b.float().T).to(getattr(torch, dtype))
+    for _ in range(2):
+        actual = kernel(a.npu(), b.npu())
+        torch.npu.synchronize()
+        torch.testing.assert_close(actual.cpu(), expected, atol=0, rtol=0)
+
+
 def test_padded_alias_storage_sets_the_physical_version_pitch():
     @T.prim_func
     def main(A: T.Tensor((60, 15), "bfloat16")):

testing/ascend/transform/test_ascend_lower_copy.py

@@ -159,6 +159,50 @@ def test_fix_uses_destination_extent_and_preserves_source_pitch(case):
     _equal(copy.args[3:7], [12, 8, 40, 16] if case == "grouped" else [16, 72, 16, 128])
 
 
+@pytest.mark.parametrize("dual,src_shape", [(0, (32, 64)), (1, (64, 64)), (2, (32, 128))], ids=["single", "split-m", "split-n"])
+def test_fix_nz_uses_physical_panel_pitch_and_slice_origin(dual, src_shape):
+    src = tirx.decl_buffer(src_shape, "float32", name="src", scope="shared.l0c")
+    dst = tirx.decl_buffer((2, 64, 96), "float32", name="dst", scope="shared")
+    after = _lower(
+        src,
+        dst,
+        dst_region=_region(dst, [1, 32, 64], [1, 16, 16]),
+        layouts={src: make_ascend_l0c_layout(src), dst: make_ascend_l0c_layout(dst)},
+        annotations={"dual_dst_ctl": dual},
+        unit_flag_ctrl=3,
+    )
+    (copy,) = calls(after, "tl.ascend_copy_matrix_cc_to_ub")
+    _equal(copy.args[3:8], [src_shape[1], src_shape[0], 64 * 16, src_shape[0], dual])
+    _equal([copy.args[0].args[2], copy.args[10], copy.args[14]], [64 * 96 + 64 * 16 + 16 * 16, 3, 0])
+
+
+@pytest.mark.parametrize("column,width", [(8, 32), (0, 24)])
+def test_fix_nz_rejects_partial_column_panels(column, width):
+    src = tirx.decl_buffer((32, width), "float32", name="src", scope="shared.l0c")
+    dst = tirx.decl_buffer((32, 64), "float32", name="dst", scope="shared")
+    with pytest.raises(ValueError, match="aligned to 16 elements"):
+        _lower(
+            src,
+            dst,
+            dst_region=_region(dst, [32, width], [0, column]),
+            layouts={src: make_ascend_l0c_layout(src), dst: make_ascend_l0c_layout(dst)},
+        )
+
+
+def test_fix_nz_rejects_l1_float32_panel_width():
+    src = tirx.decl_buffer((32, 64), "float32", name="src", scope="shared.l0c")
+    dst = tirx.decl_buffer((32, 64), "float32", name="dst", scope="shared")
+    with pytest.raises(ValueError, match="use make_ascend_l0c_layout"):
+        _lower(src, dst, layouts={src: make_ascend_l0c_layout(src), dst: make_ascend_major_k_layout(dst)})
+
+
+def test_fix_nz_rejects_noncompact_logical_strides():
+    src = tirx.decl_buffer((32, 64), "float32", name="src", scope="shared.l0c")
+    dst = tirx.decl_buffer((32, 64), "float32", name="dst", scope="shared", strides=(80, 1))
+    with pytest.raises(ValueError, match="requires dense logical strides"):
+        _lower(src, dst, layouts={src: make_ascend_l0c_layout(src), dst: make_ascend_l0c_layout(dst)})
+
+
 @pytest.mark.parametrize("scope", ["shared.l1", "shared.dyn"])
 def test_nonrepresentable_dma_geometry_is_rejected(scope):
     src = tirx.decl_buffer((2, 3, 4), "bfloat16", name="src", strides=(20, 5, 1) if scope == "shared.l1" else (14, 4, 1))

tilelang/layout/ascend.py

@@ -87,6 +87,11 @@ def make_ascend_l0c_layout(buffer: BufferLikeType):
 
     The physical shape is ``[.., N / 16, M / 16, 16, 16]`` and the forward map
     is ``(.., m, n) -> [.., n / 16, m / 16, m % 16, n % 16]``.
+
+    This layout can also annotate a UB buffer to preserve NZ in L0C-to-UB
+    ``T.copy`` and ``T.dual_copy``. Declare padding in the buffer shape;
+    copy regions retain the physical panel pitch of that allocation and
+    must cover whole 16-column panels. Buffer indices remain logical.
     """
     buf, _, _ = _get_buffer_info(buffer)
     return _ffi_api.make_ascend_l0c_layout(buf)