diff --git a/src/ascend/op/copy.cc b/src/ascend/op/copy.cc index fc87f4339c..59ec404011 100644 --- a/src/ascend/op/copy.cc +++ b/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")) { diff --git a/src/ascend/op/copy.h b/src/ascend/op/copy.h index 2a3cb69f37..9ca8b46baa 100644 --- a/src/ascend/op/copy.h +++ b/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 mx_sf_data; diff --git a/testing/ascend/language/test_tilelang_ascend_dma_copy.py b/testing/ascend/language/test_tilelang_ascend_dma_copy.py index 4881abef77..4830ea730a 100644 --- a/testing/ascend/language/test_tilelang_ascend_dma_copy.py +++ b/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 diff --git a/testing/ascend/transform/test_ascend_lower_copy.py b/testing/ascend/transform/test_ascend_lower_copy.py index 61054fd266..53c86204f3 100644 --- a/testing/ascend/transform/test_ascend_lower_copy.py +++ b/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)) diff --git a/tilelang/ascend/language/copy_op.py b/tilelang/ascend/language/copy_op.py index 25e6e42ad1..c1a89bad86 100644 --- a/tilelang/ascend/language/copy_op.py +++ b/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.