[Ascend] Add opt-in zero-burst DMA emission for T.copy

src/ascend/op/copy.cc

@@ -262,6 +262,17 @@ Stmt LowerDMACopy(const AscendCopyNode &op, const LowerArgs &T,
       return -1;
     };
     int64_t rb_plan = as_int_plan(plan_row_bytes);
+    if (op.emit_zero_burst) {
+      ICHECK_GT(rb_plan, 0)
+          << "Ascend T.copy emit_zero_burst requires a statically-known "
+             "positive DMA row byte length; use a strided 2D copy region.";
+      // Removing this guard must not discard another axis's OOB condition.
+      ICHECK(analyzer->CanProveEqual(
+          has_data, plan.n_rows > make_zero(plan.n_rows.dtype())))
+          << "Ascend T.copy emit_zero_burst requires the DMA row count to "
+             "represent every possibly-empty axis; another axis still "
+             "requires an empty-region guard.";
+    }
     bool data_select = dma_path == DMAPath::kGMToUB && op.data_select != 0;
     // Padding right-pads each row up to the next 32B boundary, so it needs a
     // statically-known row byte length to compute the pad count. A symbolic
@@ -732,7 +743,8 @@ Stmt LowerDMACopy(const AscendCopyNode &op, const LowerArgs &T,
   }
 
   Stmt copy_stmt = Evaluate(call);
-  if (!is_one(has_data) && !analyzer->CanProve(has_data)) {
+  if (!op.emit_zero_burst && !is_one(has_data) &&
+      !analyzer->CanProve(has_data)) {
     copy_stmt = IfThenElse(has_data, copy_stmt);
   }
   return copy_stmt;
@@ -949,6 +961,11 @@ struct AscendCopyImpl {
 
   static Stmt Lower(const AscendCopyNode &op, const LowerArgs &T,
                     arith::Analyzer *analyzer) {
+    ICHECK(!op.emit_zero_burst ||
+           (GetDMAPath(op.src, op.dst) == DMAPath::kGMToUB &&
+            !IsInsideSimtVF(T) && !op.data_select && !op.pad_value.defined()))
+        << "Ascend T.copy emit_zero_burst requires unpadded GM->UB DMA "
+           "outside SimtVF.";
     ICHECK(!IsL0SFBuffer(op.src))
         << "L0 SF handles can only be read by gemm_blockscaled";
     if (IsL0SFBuffer(op.dst)) {
@@ -992,6 +1009,7 @@ void DecodeAscendCopyAnnotations(AscendCopyNode *node) {
   node->dual_dst_ctl = int_or("dual_dst_ctl", 0);
   node->transpose = int_or("transpose", 0);
   node->data_select = int_or("data_select", 0);
+  node->emit_zero_burst = int_or("emit_zero_burst", 0) != 0;
   node->mx_sf_data = std::nullopt;
   node->l2_cache_ctrl = std::nullopt;
   if (auto val = node->annotations.Get("l2_cache_ctrl")) {

src/ascend/op/copy.h

@@ -33,6 +33,8 @@ class AscendCopyNode : public CopyNode {
   int dual_dst_ctl{0}; // "dual_dst_ctl": L0C->UB dual-destination mode
   int transpose{0};    // "transpose": GM->L1 dn2nz / L1->L0 transpose request
   int data_select{0};  // "data_select": GM->UB hardware right-padding
+  // "emit_zero_burst": unguarded GM->UB row count.
+  bool emit_zero_burst{false};
   // Canonical allocation buffer selected from alloc_l0*_sf's IR binding.
   // Leading indices already include any materialized physical versions.
   Optional<Buffer> mx_sf_data;

testing/ascend/language/test_tilelang_ascend_dma_copy.py

@@ -40,6 +40,36 @@ def copy(src: T.Tensor((4, 2, 32), "float32"), dst: T.Tensor((4, 32), "float32")
     torch.testing.assert_close(dst, expected)
 
 
+@tilelang.testing.requires_ascend
+def test_gm_to_ub_emit_zero_burst_preserves_destination():
+    src_pitch = T.dynamic("src_pitch", dtype="int64")
+
+    @T.prim_func
+    def copy(
+        src: T.StridedTensor((2, 16), (src_pitch, 1), "int32"),
+        seed: T.Tensor((2, 24), "int32"),
+        dst: T.Tensor((2, 24), "int32"),
+        valid: T.int32,
+    ):
+        with T.Kernel(1):
+            ub = T.alloc_shared((2, 24), "int32")
+            T.copy(seed, ub)
+            rows = T.if_then_else(valid != 0, 2, 0)
+            T.copy(src[:rows, :], ub[:rows, :16], annotations={"emit_zero_burst": True})
+            T.copy(ub, dst)
+
+    kernel = tilelang.compile(copy, out_idx=2, target="ascend")
+    seed = torch.arange(48, dtype=torch.int32, device="npu").reshape(2, 24) + 1000
+    for width in (32, 48):
+        storage = torch.arange(2 * width, dtype=torch.int32, device="npu").reshape(2, width)
+        src = storage[:, :16]
+        for valid in (1, 0, 1):
+            expected = seed.clone()
+            if valid:
+                expected[:, :16] = src
+            torch.testing.assert_close(kernel(src, seed, valid), expected)
+
+
 def l0c_to_strided_gm_copy():
     ldd = T.dynamic("ldd", dtype="int64")
     m, n, k = 16, 16, 64

testing/ascend/transform/test_ascend_lower_copy.py

@@ -131,6 +131,66 @@ def test_dynamic_rows_keep_the_stride_across_a_fixed_middle_axis():
         assert bool(tvm.arith.Analyzer().simplify(condition)) == (count > 0)
 
 
+@pytest.mark.parametrize("enabled", [False, True])
+@pytest.mark.parametrize("count_expr", ["select", "scalar"])
+def test_zero_burst_annotation_controls_only_the_dma_guard(enabled, count_expr):
+    valid = tirx.Var("valid", "bool")
+    rows = tirx.if_then_else(valid, 2, 0) if count_expr == "select" else tirx.Var("rows", "int32")
+    src_pitch = tirx.Var("src_pitch", "int64")
+    src = tirx.decl_buffer((4, 32), "float32", name="src", strides=(src_pitch, 1))
+    dst = tirx.decl_buffer((4, 64), "float32", name="dst", scope="shared.dyn")
+    after = _lower(
+        src,
+        dst,
+        _region(src, [rows, 32], [1, 0]),
+        _region(dst, [rows, 32], [1, 0]),
+        annotations={"emit_zero_burst": enabled, "l2_cache_ctrl": 5},
+    )
+    (copy,) = calls(after, "tl.ascend_copy_gm_to_ubuf")
+    _equal([copy.args[i] for i in (3, 4, 8, 9, 10)], [rows, 128, 5, src_pitch * 4, 256])
+    _equal([copy.args[0].args[2], copy.args[1].args[2]], [64, src_pitch])
+    guards = nodes(after, tirx.IfThenElse)
+    assert len(guards) == (0 if enabled else 1)
+    if guards:
+        _equal([guards[0].condition], [rows > 0])
+
+
+@pytest.mark.parametrize(
+    "case,diagnostic",
+    [
+        ("store", "requires unpadded GM->UB DMA"),
+        ("padding", "requires unpadded GM->UB DMA"),
+        ("pad-value", "requires unpadded GM->UB DMA"),
+        ("simtvf", "outside SimtVF"),
+        ("symbolic-width", "statically-known positive DMA row byte length"),
+        ("extra-empty-axis", "every possibly-empty axis"),
+        ("normalized-zero-or-one", "every possibly-empty axis"),
+    ],
+)
+def test_zero_burst_rejects_unsupported_requests(case, diagnostic):
+    rows = tirx.Select(tirx.Var("valid", "bool"), 1 if case == "normalized-zero-or-one" else 2, 0)
+    cols = tirx.Var("cols", "int32") if case == "symbolic-width" else 32
+    src = tirx.decl_buffer((4, 64), "float32", name="src", scope="shared.dyn" if case == "store" else "global")
+    dst_scope = "global" if case == "store" else "shared.dyn"
+    dst = tirx.decl_buffer((4, 64), "float32", name="dst", scope=dst_scope)
+    src_region = _region(src, [rows, cols])
+    if case == "extra-empty-axis":
+        src = tirx.decl_buffer((4, 2, 64), "float32", name="src")
+        selector_extent = tirx.Select(tirx.Var("in_bounds", "bool"), 1, 0)
+        src_region = _region(src, [rows, selector_extent, cols])
+    annotations = {"emit_zero_burst": True}
+    if case == "padding":
+        annotations["data_select"] = 1
+    if case == "pad-value":
+        annotations["pad_value"] = 0
+    wrap = None
+    if case == "simtvf":
+        tx = tirx.IterVar(tvm.ir.Range(0, 32), tirx.Var("simtvf_thread", "int32"), tirx.IterVar.ThreadIndex, "threadIdx.x")
+        wrap = lambda body: tirx.AttrStmt(tx, "thread_extent", 32, body)  # noqa: E731
+    with pytest.raises(tvm.error.InternalError, match=diagnostic):
+        _lower(src, dst, src_region, _region(dst, [rows, cols]), annotations=annotations, wrap=wrap)
+
+
 def test_fixed_outer_axis_affects_the_pointer_but_not_the_l1_row_pitch():
     group = tirx.Var("group", "int32")
     src = tirx.decl_buffer((4, 16, 64), "bfloat16", name="src", strides=(2000, 80, 1))

tilelang/ascend/language/copy_op.py

@@ -287,6 +287,14 @@ def copy(  # noqa: A001
             ``T.ascend_set_copy_pad_value(...)`` beforehand.
         annotations (Optional[dict], keyword-only): Additional annotations
             dict; values in it take precedence over the individual keywords.
+            ``{"emit_zero_burst": True}`` skips the empty-region guard for
+            an unpadded GM→UB DMA outside SimtVF. Requires fixed positive row
+            bytes and a DMA row count that captures all empty-region conditions;
+            unsupported geometry raises a lowering error. The caller must
+            supply nonnegative, hardware-valid counts and address/stride
+            expressions safe to evaluate even for zero rows. Zero bursts
+            preserve UB contents. Compile-time-empty copies may be eliminated.
+            Defaults to False.
         loop_layout (Optional[Fragment], keyword-only): Parallel loop layout
             hint for the SIMT copy path.