diff --git a/src/ascend/op/copy.cc b/src/ascend/op/copy.cc index 302d832b99..d3689ad7c6 100644 --- a/src/ascend/op/copy.cc +++ b/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 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(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 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 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 diff --git a/src/ascend/transform/normalize_fractal_storage.cc b/src/ascend/transform/normalize_fractal_storage.cc index 7fb7228b82..1a19a98a01 100644 --- a/src/ascend/transform/normalize_fractal_storage.cc +++ b/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 diff --git a/testing/ascend/layout/test_ascend_fractal_storage_integration.py b/testing/ascend/layout/test_ascend_fractal_storage_integration.py index f8ccf45c7c..c7c04820ba 100644 --- a/testing/ascend/layout/test_ascend_fractal_storage_integration.py +++ b/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")): diff --git a/testing/ascend/transform/test_ascend_lower_copy.py b/testing/ascend/transform/test_ascend_lower_copy.py index 525f4848df..206a488a24 100644 --- a/testing/ascend/transform/test_ascend_lower_copy.py +++ b/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)) diff --git a/tilelang/layout/ascend.py b/tilelang/layout/ascend.py index cfcaf8560f..6409f341b3 100644 --- a/tilelang/layout/ascend.py +++ b/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)