From 33f29b01b7473c0a3d34643148683a69e481074a Mon Sep 17 00:00:00 2001 From: Denver Jin Date: Mon, 12 Oct 2026 14:57:20 +0800 Subject: [PATCH 1/3] [Ascend] Support NZ layouts for L0C-to-UB copies --- src/ascend/op/copy.cc | 50 ++++++++++++- .../transform/normalize_fractal_storage.cc | 5 ++ ...test_ascend_fractal_storage_integration.py | 75 +++++++++++++++++++ ...est_tilelang_ascend_fractal_layout_tags.py | 5 ++ .../transform/test_ascend_lower_copy.py | 45 +++++++++++ tilelang/layout/__init__.py | 1 + tilelang/layout/ascend.py | 21 ++++++ 7 files changed, 199 insertions(+), 3 deletions(-) diff --git a/src/ascend/op/copy.cc b/src/ascend/op/copy.cc index 302d832b99..8025d66abc 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_ub_nz_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..7c93948c25 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_ub_nz_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_ub_nz_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/layout/test_tilelang_ascend_fractal_layout_tags.py b/testing/ascend/layout/test_tilelang_ascend_fractal_layout_tags.py index 2e5584e20a..d12e92f523 100644 --- a/testing/ascend/layout/test_tilelang_ascend_fractal_layout_tags.py +++ b/testing/ascend/layout/test_tilelang_ascend_fractal_layout_tags.py @@ -8,6 +8,7 @@ make_ascend_major_k_layout, make_ascend_major_mn_layout, make_ascend_sf_layout, + make_ascend_ub_nz_layout, try_extract_fractal_layout, ) @@ -18,6 +19,7 @@ (make_ascend_major_k_layout, (16, 32), "float8_e4m3fn", (0, 1, 0)), (make_ascend_major_mn_layout, (32, 16), "float8_e4m3fn", (3, 0, 1)), (make_ascend_l0c_layout, (16, 16), "float32", (1, 1, 0)), + (make_ascend_ub_nz_layout, (2, 32, 64), "float32", (1, 1, 0)), (make_ascend_sf_layout, (16, 1), "uint16", (2, 1, 0)), (make_ascend_sf_layout, (32, 1), "uint16", (2, 1, 0)), # The same small shape has different semantics despite identical indices. @@ -29,6 +31,9 @@ def test_canonical_layout_tags(factory, shape, dtype, expected): info = try_extract_fractal_layout(factory(buf), buf) assert (info.kind, info.c0_axis, info.row16_axis) == expected + if factory == make_ascend_ub_nz_layout: + assert int(info.c0) == 16 + assert list(factory(buf).get_output_shape()) == [2, 4, 2, 16, 16] def test_expand_preserves_sf_tag(): diff --git a/testing/ascend/transform/test_ascend_lower_copy.py b/testing/ascend/transform/test_ascend_lower_copy.py index 525f4848df..3aab93d5b9 100644 --- a/testing/ascend/transform/test_ascend_lower_copy.py +++ b/testing/ascend/transform/test_ascend_lower_copy.py @@ -11,6 +11,7 @@ make_ascend_major_k_layout, make_ascend_major_mn_layout, make_ascend_sf_layout, + make_ascend_ub_nz_layout, ) from tvm import tirx from testing.ascend._ir import allocated_buffer, calls, nodes, kernel @@ -159,6 +160,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_ub_nz_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_ub_nz_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_ub_nz_layout"): + _lower(src, dst, layouts={src: make_ascend_l0c_layout(src), dst: make_ascend_nz_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_ub_nz_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/__init__.py b/tilelang/layout/__init__.py index d13766127b..4389b72ba7 100644 --- a/tilelang/layout/__init__.py +++ b/tilelang/layout/__init__.py @@ -26,6 +26,7 @@ make_ascend_major_k_layout, # noqa: F401 make_ascend_major_mn_layout, # noqa: F401 make_ascend_l0c_layout, # noqa: F401 + make_ascend_ub_nz_layout, # noqa: F401 make_ascend_sf_layout, # noqa: F401 make_strided_slice, # noqa: F401 try_extract_fractal_layout, # noqa: F401 diff --git a/tilelang/layout/ascend.py b/tilelang/layout/ascend.py index cfcaf8560f..078d60afcc 100644 --- a/tilelang/layout/ascend.py +++ b/tilelang/layout/ascend.py @@ -92,6 +92,27 @@ def make_ascend_l0c_layout(buffer: BufferLikeType): return _ffi_api.make_ascend_l0c_layout(buf) +def make_ascend_ub_nz_layout(buffer: BufferLikeType): + """Declare the UB NZ format produced by an L0C-to-UB copy. + + Use with ``T.annotate_layout({ub: make_ascend_ub_nz_layout(ub)})``. + ``T.copy`` and ``T.dual_copy`` then preserve NZ instead of converting to + ND. Logical ``[..., m, n]`` maps to + ``[..., n // 16, m // 16, m % 16, n % 16]``. Unlike the L1 NZ layout, + the column block has 16 elements even for float32. + + Declare any padding in the buffer shape; copy regions retain the + allocation's physical panel pitch. NZ copies require whole 16-column + panels. Ordinary indexing uses logical coordinates; raw SIMD pointers + still access consecutive physical elements. + """ + buf, _, _ = _get_buffer_info(buffer) + if buf.scope() not in ("shared", "shared.dyn"): + raise ValueError("make_ascend_ub_nz_layout requires a UB buffer") + # Fixpipe preserves the accumulator's 16x16 ordering in UB. + return _ffi_api.make_ascend_l0c_layout(buf) + + def make_ascend_sf_layout(buffer: BufferLikeType): """ Create an Ascend SF layout for MX scale-factor L1 buffers. From e1a13a6f21303765954ad77acc8a05f06bc253b6 Mon Sep 17 00:00:00 2001 From: Denver Jin Date: Mon, 12 Oct 2026 15:11:58 +0800 Subject: [PATCH 2/3] [Ascend] Select NZ layouts from buffer scope --- src/ascend/layout/ascend_layouts.cc | 8 +++- src/ascend/layout/ascend_layouts.h | 5 ++- src/ascend/op/copy.cc | 2 +- src/ascend/transform/insert_nd2nz.cc | 5 +++ ...test_ascend_fractal_storage_integration.py | 4 +- ...est_tilelang_ascend_fractal_layout_tags.py | 17 +++++-- .../transform/test_ascend_insert_nd2nz.py | 25 ++++++++++- .../transform/test_ascend_lower_copy.py | 11 +++-- tilelang/layout/__init__.py | 1 - tilelang/layout/ascend.py | 44 +++++++------------ 10 files changed, 76 insertions(+), 46 deletions(-) diff --git a/src/ascend/layout/ascend_layouts.cc b/src/ascend/layout/ascend_layouts.cc index f141c20e70..489a2f9bbc 100644 --- a/src/ascend/layout/ascend_layouts.cc +++ b/src/ascend/layout/ascend_layouts.cc @@ -38,8 +38,8 @@ static Layout ExpandLayout2D(const Layout &base, const Buffer &buffer) { // ND2NZ convention where the contiguous axis (cols) is grouped by C0 and is the // outer physical major. Within a fractal the row lane remains stride C0 and the // column lane remains stride 1. MajorK means cols is the reduce-K axis; MajorMN -// means cols is the output M/N axis. The legacy NZ API remains an alias of -// MajorK while the frontend migrates to explicit major names. +// means cols is the output M/N axis. The NZ API selects this layout for Cube +// inputs and the fixed accumulator layout for L0C and UB Fixpipe output. static constexpr int kAscendRowFractal = 16; namespace { @@ -198,6 +198,10 @@ Layout MakeAscendMajorMNLayout(const Buffer &buffer, int k_align) { } Layout MakeAscendNZLayout(const Buffer &buffer) { + ICHECK(buffer.defined()) << "NZ layout expects a defined buffer"; + const String scope = buffer.scope(); + if (scope == "shared" || scope == "shared.dyn" || scope == "shared.l0c") + return MakeAscendL0CLayout(buffer); return MakeAscendMajorKLayout(buffer); } diff --git a/src/ascend/layout/ascend_layouts.h b/src/ascend/layout/ascend_layouts.h index 24d731191a..4f93c125ad 100644 --- a/src/ascend/layout/ascend_layouts.h +++ b/src/ascend/layout/ascend_layouts.h @@ -18,12 +18,13 @@ namespace tl { // tile maps to physical [.., cols / C0, rows / 16, 16, C0] (a grid of 16xC0 // fractals, row-major within a fractal) where C0 = 32 bytes / element_size // (16 for bf16/fp16, 32 for fp8, 8 for fp32). In MajorK, cols is the reduce-K -// axis; in MajorMN, cols is the output M/N axis. The legacy NZ name is kept as -// an alias of MajorK during the transition. +// axis; in MajorMN, cols is the output M/N axis. Layout MakeAscendMajorKLayout(const tirx::Buffer &buffer, int k_align = 1); Layout MakeAscendMajorMNLayout(const tirx::Buffer &buffer, int k_align = 1); Layout MakeAscendL0CLayout(const tirx::Buffer &buffer); Layout MakeAscendSFLayout(const tirx::Buffer &buffer); +// NZ selects the fixed 16x16 accumulator layout for L0C and UB Fixpipe +// output, and MajorK for other scopes. Compact UB packing remains explicit. Layout MakeAscendNZLayout(const tirx::Buffer &buffer); // Compact padded-NZ layout used by software UB ND->NZ scatter: diff --git a/src/ascend/op/copy.cc b/src/ascend/op/copy.cc index 8025d66abc..a8bd518002 100644 --- a/src/ascend/op/copy.cc +++ b/src/ascend/op/copy.cc @@ -737,7 +737,7 @@ Stmt LowerDMACopy(const AscendCopyNode &op, const LowerArgs &T, analyzer->CanProveEqual(dst_info.c0, 16), ValueError) << "Ascend NZ L0C->UB requires 16-column NZ panels on " - << op.dst->name << "; use make_ascend_ub_nz_layout."; + << op.dst->name << "; use make_ascend_nz_layout."; if (!op.dst->strides.empty()) { PrimExpr stride = I(1); for (int i = static_cast(op.dst->shape.size()) - 1; i >= 0; --i) { diff --git a/src/ascend/transform/insert_nd2nz.cc b/src/ascend/transform/insert_nd2nz.cc index a4e7ec9c5d..c38d83cddc 100644 --- a/src/ascend/transform/insert_nd2nz.cc +++ b/src/ascend/transform/insert_nd2nz.cc @@ -227,6 +227,11 @@ class Nd2NzInserter : public arith::IRMutatorWithAnalyzer { Array ValidatePrepackedNZSource(const CopyNode ©, PrimExpr rows, PrimExpr cols, PrimExpr full_rows) const { + TVM_FFI_CHECK( + IsAscendCompactNZLayout(LayoutFor(copy.src).value(), copy.src), + ValueError) + << "UB-to-L1 NZ post-copy requires compact NZ storage; use " + "make_ascend_compact_nz_layout for UB input tiles."; RequireStaticInt(rows, "NZ post-copy rows"); RequireStaticInt(cols, "NZ post-copy cols"); RequireStaticInt(full_rows, "NZ post-copy full_rows"); diff --git a/testing/ascend/layout/test_ascend_fractal_storage_integration.py b/testing/ascend/layout/test_ascend_fractal_storage_integration.py index 7c93948c25..06862d983a 100644 --- a/testing/ascend/layout/test_ascend_fractal_storage_integration.py +++ b/testing/ascend/layout/test_ascend_fractal_storage_integration.py @@ -5,7 +5,7 @@ import tilelang.ascend.language as T import tilelang.testing from tilelang import tvm -from tilelang.layout import make_ascend_ub_nz_layout +from tilelang.layout import make_ascend_nz_layout from tvm import tirx from testing.ascend._ir import allocated_buffer, calls, nodes, statements @@ -25,7 +25,7 @@ def main(A: T.Tensor((32, 16), "bfloat16"), B: T.Tensor((n, 16), "bfloat16"), C: ub = T.view(storage, (2, 48, 80)) else: ub = T.alloc_shared((2, 48, 80), dtype) - T.annotate_layout({ub: make_ascend_ub_nz_layout(ub)}) + T.annotate_layout({ub: make_ascend_nz_layout(ub)}) T.copy(A, a) T.copy(B, b) T.gemm(a, b, acc, transpose_B=True, clear_accum=True) diff --git a/testing/ascend/layout/test_tilelang_ascend_fractal_layout_tags.py b/testing/ascend/layout/test_tilelang_ascend_fractal_layout_tags.py index d12e92f523..8692dee50a 100644 --- a/testing/ascend/layout/test_tilelang_ascend_fractal_layout_tags.py +++ b/testing/ascend/layout/test_tilelang_ascend_fractal_layout_tags.py @@ -8,7 +8,7 @@ make_ascend_major_k_layout, make_ascend_major_mn_layout, make_ascend_sf_layout, - make_ascend_ub_nz_layout, + make_ascend_nz_layout, try_extract_fractal_layout, ) @@ -19,7 +19,7 @@ (make_ascend_major_k_layout, (16, 32), "float8_e4m3fn", (0, 1, 0)), (make_ascend_major_mn_layout, (32, 16), "float8_e4m3fn", (3, 0, 1)), (make_ascend_l0c_layout, (16, 16), "float32", (1, 1, 0)), - (make_ascend_ub_nz_layout, (2, 32, 64), "float32", (1, 1, 0)), + (make_ascend_nz_layout, (2, 32, 64), "float32", (1, 1, 0)), (make_ascend_sf_layout, (16, 1), "uint16", (2, 1, 0)), (make_ascend_sf_layout, (32, 1), "uint16", (2, 1, 0)), # The same small shape has different semantics despite identical indices. @@ -31,11 +31,22 @@ def test_canonical_layout_tags(factory, shape, dtype, expected): info = try_extract_fractal_layout(factory(buf), buf) assert (info.kind, info.c0_axis, info.row16_axis) == expected - if factory == make_ascend_ub_nz_layout: + if factory == make_ascend_nz_layout: assert int(info.c0) == 16 assert list(factory(buf).get_output_shape()) == [2, 4, 2, 16, 16] +@pytest.mark.parametrize("scope,kind", [("shared", 1), ("shared.dyn", 1), ("shared.l0c", 1), ("shared.l1", 0)]) +@pytest.mark.parametrize("dtype", ["float32", "bfloat16"]) +def test_nz_layout_uses_buffer_scope(scope, kind, dtype): + buf = tirx.decl_buffer((2, 32, 64), dtype, name="A", scope=scope) + layout = make_ascend_nz_layout(buf) + info = try_extract_fractal_layout(layout, buf) + c0 = 8 if scope == "shared.l1" and dtype == "float32" else 16 + assert (info.kind, int(info.c0), info.c0_axis, info.row16_axis) == (kind, c0, 1, 0) + assert list(layout.get_output_shape()) == [2, 64 // c0, 2, 16, c0] + + def test_expand_preserves_sf_tag(): base_buf = tirx.decl_buffer((16, 1), "uint16", name="SF2", scope="shared") expanded_buf = tirx.decl_buffer((2, 16, 1), "uint16", name="SF3", scope="shared") diff --git a/testing/ascend/transform/test_ascend_insert_nd2nz.py b/testing/ascend/transform/test_ascend_insert_nd2nz.py index 101173b966..3e28614bd6 100644 --- a/testing/ascend/transform/test_ascend_insert_nd2nz.py +++ b/testing/ascend/transform/test_ascend_insert_nd2nz.py @@ -3,7 +3,7 @@ import pytest from tilelang import tvm from tilelang.ascend import language as T, transform -from tilelang.layout import make_ascend_compact_nz_layout +from tilelang.layout import make_ascend_compact_nz_layout, make_ascend_nz_layout from tvm import tirx from testing.ascend._ir import calls @@ -66,3 +66,26 @@ def test_scatter_storage_requirements(kwargs, error): transform.InsertNd2Nz()(before) else: assert len(calls(transform.InsertNd2Nz()(before), "tl.ascend_nd2nz_scatter")) == 1 + + +@pytest.mark.parametrize("factory", [make_ascend_nz_layout, make_ascend_compact_nz_layout]) +@pytest.mark.parametrize("dtype", ["float32", "bfloat16"]) +def test_post_copy_requires_compact_nz_panel_pitch(factory, dtype): + src = tirx.decl_buffer((33, 128), dtype, name="src", scope="shared.dyn") + dst = tirx.decl_buffer((32, 128), dtype, name="dst", scope="shared.l1") + region = tirx.BufferRegion(src, [tvm.ir.Range(0, 32), tvm.ir.Range(0, 128)]) + root = tirx.SBlock( + [], + [], + [], + "root", + tirx.Evaluate(T.copy(region, dst)), + alloc_buffers=[src, dst], + annotations={"layout_map": {src: factory(src), dst: make_ascend_nz_layout(dst)}}, + ) + before = tvm.IRModule({"main": tirx.PrimFunc([], tirx.SBlockRealize([], True, root))}) + if factory == make_ascend_nz_layout: + with pytest.raises(ValueError, match="requires compact NZ storage"): + transform.InsertNd2Nz()(before) + else: + assert len(calls(transform.InsertNd2Nz()(before), "tl.ascend_nd2nz_post_copy")) == 1 diff --git a/testing/ascend/transform/test_ascend_lower_copy.py b/testing/ascend/transform/test_ascend_lower_copy.py index 3aab93d5b9..faa0d7522c 100644 --- a/testing/ascend/transform/test_ascend_lower_copy.py +++ b/testing/ascend/transform/test_ascend_lower_copy.py @@ -11,7 +11,6 @@ make_ascend_major_k_layout, make_ascend_major_mn_layout, make_ascend_sf_layout, - make_ascend_ub_nz_layout, ) from tvm import tirx from testing.ascend._ir import allocated_buffer, calls, nodes, kernel @@ -168,7 +167,7 @@ def test_fix_nz_uses_physical_panel_pitch_and_slice_origin(dual, src_shape): src, dst, dst_region=_region(dst, [1, 32, 64], [1, 16, 16]), - layouts={src: make_ascend_l0c_layout(src), dst: make_ascend_ub_nz_layout(dst)}, + layouts={src: make_ascend_l0c_layout(src), dst: make_ascend_nz_layout(dst)}, annotations={"dual_dst_ctl": dual}, unit_flag_ctrl=3, ) @@ -186,22 +185,22 @@ def test_fix_nz_rejects_partial_column_panels(column, width): src, dst, dst_region=_region(dst, [32, width], [0, column]), - layouts={src: make_ascend_l0c_layout(src), dst: make_ascend_ub_nz_layout(dst)}, + layouts={src: make_ascend_l0c_layout(src), dst: make_ascend_nz_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_ub_nz_layout"): - _lower(src, dst, layouts={src: make_ascend_l0c_layout(src), dst: make_ascend_nz_layout(dst)}) + with pytest.raises(ValueError, match="use make_ascend_nz_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_ub_nz_layout(dst)}) + _lower(src, dst, layouts={src: make_ascend_l0c_layout(src), dst: make_ascend_nz_layout(dst)}) @pytest.mark.parametrize("scope", ["shared.l1", "shared.dyn"]) diff --git a/tilelang/layout/__init__.py b/tilelang/layout/__init__.py index 4389b72ba7..d13766127b 100644 --- a/tilelang/layout/__init__.py +++ b/tilelang/layout/__init__.py @@ -26,7 +26,6 @@ make_ascend_major_k_layout, # noqa: F401 make_ascend_major_mn_layout, # noqa: F401 make_ascend_l0c_layout, # noqa: F401 - make_ascend_ub_nz_layout, # noqa: F401 make_ascend_sf_layout, # noqa: F401 make_strided_slice, # noqa: F401 try_extract_fractal_layout, # noqa: F401 diff --git a/tilelang/layout/ascend.py b/tilelang/layout/ascend.py index 078d60afcc..0316543e9e 100644 --- a/tilelang/layout/ascend.py +++ b/tilelang/layout/ascend.py @@ -18,16 +18,25 @@ class _AscendFractalLayout(Layout): def make_ascend_nz_layout(buffer: BufferLikeType): - """ - Create an Ascend NZ (zN fractal) layout for a Cube operand buffer. + """Create the Ascend NZ layout appropriate for the buffer's scope. + + Logical ``[..., m, n]`` maps to + ``[..., n // C0, m // 16, m % 16, n % C0]``. L0C and UB + (``shared``/``shared.dyn``) use the fixed 16-column accumulator format. + Other scopes use the Cube input format, with + ``C0 = 32 bytes / element_size``. + + With ``T.annotate_layout({ub: make_ascend_nz_layout(ub)})``, L0C-to-UB + ``T.copy`` and ``T.dual_copy`` preserve NZ. Declare padding in the UB + shape; copy regions retain the allocation's physical panel pitch and + must cover whole 16-column panels. Buffer indices remain logical; + SIMD pointer increments access consecutive physical elements. - A logical ``[.., rows, cols]`` tile maps to physical - ``[.., rows / 16, cols / C0, 16, C0]`` -- a grid of 16xC0 fractals -- where - ``C0 = 32 bytes / element_size`` is the Cube unit's fractal width and the - forward map is ``(.., r, c) -> [.., r / 16, c / C0, r % 16, c % C0]``. + For compact UB packing before UB-to-L1 copies, use + :func:`make_ascend_compact_nz_layout` explicitly. Args: - buffer: BufferLikeType (rank >= 2, constant last-2 dims) + buffer: BufferLikeType (rank >= 2). Returns: Layout: The NZ fractal layout for the buffer. @@ -92,27 +101,6 @@ def make_ascend_l0c_layout(buffer: BufferLikeType): return _ffi_api.make_ascend_l0c_layout(buf) -def make_ascend_ub_nz_layout(buffer: BufferLikeType): - """Declare the UB NZ format produced by an L0C-to-UB copy. - - Use with ``T.annotate_layout({ub: make_ascend_ub_nz_layout(ub)})``. - ``T.copy`` and ``T.dual_copy`` then preserve NZ instead of converting to - ND. Logical ``[..., m, n]`` maps to - ``[..., n // 16, m // 16, m % 16, n % 16]``. Unlike the L1 NZ layout, - the column block has 16 elements even for float32. - - Declare any padding in the buffer shape; copy regions retain the - allocation's physical panel pitch. NZ copies require whole 16-column - panels. Ordinary indexing uses logical coordinates; raw SIMD pointers - still access consecutive physical elements. - """ - buf, _, _ = _get_buffer_info(buffer) - if buf.scope() not in ("shared", "shared.dyn"): - raise ValueError("make_ascend_ub_nz_layout requires a UB buffer") - # Fixpipe preserves the accumulator's 16x16 ordering in UB. - return _ffi_api.make_ascend_l0c_layout(buf) - - def make_ascend_sf_layout(buffer: BufferLikeType): """ Create an Ascend SF layout for MX scale-factor L1 buffers. From e9caa1855f39d6907332fe58d2f4096e0b088337 Mon Sep 17 00:00:00 2001 From: Denver Jin Date: Mon, 12 Oct 2026 15:28:10 +0800 Subject: [PATCH 3/3] [Ascend] Reuse explicit L0C layouts for UB NZ output --- src/ascend/layout/ascend_layouts.cc | 8 ++---- src/ascend/layout/ascend_layouts.h | 5 ++-- src/ascend/op/copy.cc | 2 +- src/ascend/transform/insert_nd2nz.cc | 5 ---- ...test_ascend_fractal_storage_integration.py | 4 +-- ...est_tilelang_ascend_fractal_layout_tags.py | 16 ----------- .../transform/test_ascend_insert_nd2nz.py | 25 +---------------- .../transform/test_ascend_lower_copy.py | 8 +++--- tilelang/layout/ascend.py | 28 ++++++++----------- 9 files changed, 24 insertions(+), 77 deletions(-) diff --git a/src/ascend/layout/ascend_layouts.cc b/src/ascend/layout/ascend_layouts.cc index 489a2f9bbc..f141c20e70 100644 --- a/src/ascend/layout/ascend_layouts.cc +++ b/src/ascend/layout/ascend_layouts.cc @@ -38,8 +38,8 @@ static Layout ExpandLayout2D(const Layout &base, const Buffer &buffer) { // ND2NZ convention where the contiguous axis (cols) is grouped by C0 and is the // outer physical major. Within a fractal the row lane remains stride C0 and the // column lane remains stride 1. MajorK means cols is the reduce-K axis; MajorMN -// means cols is the output M/N axis. The NZ API selects this layout for Cube -// inputs and the fixed accumulator layout for L0C and UB Fixpipe output. +// means cols is the output M/N axis. The legacy NZ API remains an alias of +// MajorK while the frontend migrates to explicit major names. static constexpr int kAscendRowFractal = 16; namespace { @@ -198,10 +198,6 @@ Layout MakeAscendMajorMNLayout(const Buffer &buffer, int k_align) { } Layout MakeAscendNZLayout(const Buffer &buffer) { - ICHECK(buffer.defined()) << "NZ layout expects a defined buffer"; - const String scope = buffer.scope(); - if (scope == "shared" || scope == "shared.dyn" || scope == "shared.l0c") - return MakeAscendL0CLayout(buffer); return MakeAscendMajorKLayout(buffer); } diff --git a/src/ascend/layout/ascend_layouts.h b/src/ascend/layout/ascend_layouts.h index 4f93c125ad..24d731191a 100644 --- a/src/ascend/layout/ascend_layouts.h +++ b/src/ascend/layout/ascend_layouts.h @@ -18,13 +18,12 @@ namespace tl { // tile maps to physical [.., cols / C0, rows / 16, 16, C0] (a grid of 16xC0 // fractals, row-major within a fractal) where C0 = 32 bytes / element_size // (16 for bf16/fp16, 32 for fp8, 8 for fp32). In MajorK, cols is the reduce-K -// axis; in MajorMN, cols is the output M/N axis. +// axis; in MajorMN, cols is the output M/N axis. The legacy NZ name is kept as +// an alias of MajorK during the transition. Layout MakeAscendMajorKLayout(const tirx::Buffer &buffer, int k_align = 1); Layout MakeAscendMajorMNLayout(const tirx::Buffer &buffer, int k_align = 1); Layout MakeAscendL0CLayout(const tirx::Buffer &buffer); Layout MakeAscendSFLayout(const tirx::Buffer &buffer); -// NZ selects the fixed 16x16 accumulator layout for L0C and UB Fixpipe -// output, and MajorK for other scopes. Compact UB packing remains explicit. Layout MakeAscendNZLayout(const tirx::Buffer &buffer); // Compact padded-NZ layout used by software UB ND->NZ scatter: diff --git a/src/ascend/op/copy.cc b/src/ascend/op/copy.cc index a8bd518002..d3689ad7c6 100644 --- a/src/ascend/op/copy.cc +++ b/src/ascend/op/copy.cc @@ -737,7 +737,7 @@ Stmt LowerDMACopy(const AscendCopyNode &op, const LowerArgs &T, analyzer->CanProveEqual(dst_info.c0, 16), ValueError) << "Ascend NZ L0C->UB requires 16-column NZ panels on " - << op.dst->name << "; use make_ascend_nz_layout."; + << 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) { diff --git a/src/ascend/transform/insert_nd2nz.cc b/src/ascend/transform/insert_nd2nz.cc index c38d83cddc..a4e7ec9c5d 100644 --- a/src/ascend/transform/insert_nd2nz.cc +++ b/src/ascend/transform/insert_nd2nz.cc @@ -227,11 +227,6 @@ class Nd2NzInserter : public arith::IRMutatorWithAnalyzer { Array ValidatePrepackedNZSource(const CopyNode ©, PrimExpr rows, PrimExpr cols, PrimExpr full_rows) const { - TVM_FFI_CHECK( - IsAscendCompactNZLayout(LayoutFor(copy.src).value(), copy.src), - ValueError) - << "UB-to-L1 NZ post-copy requires compact NZ storage; use " - "make_ascend_compact_nz_layout for UB input tiles."; RequireStaticInt(rows, "NZ post-copy rows"); RequireStaticInt(cols, "NZ post-copy cols"); RequireStaticInt(full_rows, "NZ post-copy full_rows"); diff --git a/testing/ascend/layout/test_ascend_fractal_storage_integration.py b/testing/ascend/layout/test_ascend_fractal_storage_integration.py index 06862d983a..c7c04820ba 100644 --- a/testing/ascend/layout/test_ascend_fractal_storage_integration.py +++ b/testing/ascend/layout/test_ascend_fractal_storage_integration.py @@ -5,7 +5,7 @@ import tilelang.ascend.language as T import tilelang.testing from tilelang import tvm -from tilelang.layout import make_ascend_nz_layout +from tilelang.layout import make_ascend_l0c_layout from tvm import tirx from testing.ascend._ir import allocated_buffer, calls, nodes, statements @@ -25,7 +25,7 @@ def main(A: T.Tensor((32, 16), "bfloat16"), B: T.Tensor((n, 16), "bfloat16"), C: ub = T.view(storage, (2, 48, 80)) else: ub = T.alloc_shared((2, 48, 80), dtype) - T.annotate_layout({ub: make_ascend_nz_layout(ub)}) + 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) diff --git a/testing/ascend/layout/test_tilelang_ascend_fractal_layout_tags.py b/testing/ascend/layout/test_tilelang_ascend_fractal_layout_tags.py index 8692dee50a..2e5584e20a 100644 --- a/testing/ascend/layout/test_tilelang_ascend_fractal_layout_tags.py +++ b/testing/ascend/layout/test_tilelang_ascend_fractal_layout_tags.py @@ -8,7 +8,6 @@ make_ascend_major_k_layout, make_ascend_major_mn_layout, make_ascend_sf_layout, - make_ascend_nz_layout, try_extract_fractal_layout, ) @@ -19,7 +18,6 @@ (make_ascend_major_k_layout, (16, 32), "float8_e4m3fn", (0, 1, 0)), (make_ascend_major_mn_layout, (32, 16), "float8_e4m3fn", (3, 0, 1)), (make_ascend_l0c_layout, (16, 16), "float32", (1, 1, 0)), - (make_ascend_nz_layout, (2, 32, 64), "float32", (1, 1, 0)), (make_ascend_sf_layout, (16, 1), "uint16", (2, 1, 0)), (make_ascend_sf_layout, (32, 1), "uint16", (2, 1, 0)), # The same small shape has different semantics despite identical indices. @@ -31,20 +29,6 @@ def test_canonical_layout_tags(factory, shape, dtype, expected): info = try_extract_fractal_layout(factory(buf), buf) assert (info.kind, info.c0_axis, info.row16_axis) == expected - if factory == make_ascend_nz_layout: - assert int(info.c0) == 16 - assert list(factory(buf).get_output_shape()) == [2, 4, 2, 16, 16] - - -@pytest.mark.parametrize("scope,kind", [("shared", 1), ("shared.dyn", 1), ("shared.l0c", 1), ("shared.l1", 0)]) -@pytest.mark.parametrize("dtype", ["float32", "bfloat16"]) -def test_nz_layout_uses_buffer_scope(scope, kind, dtype): - buf = tirx.decl_buffer((2, 32, 64), dtype, name="A", scope=scope) - layout = make_ascend_nz_layout(buf) - info = try_extract_fractal_layout(layout, buf) - c0 = 8 if scope == "shared.l1" and dtype == "float32" else 16 - assert (info.kind, int(info.c0), info.c0_axis, info.row16_axis) == (kind, c0, 1, 0) - assert list(layout.get_output_shape()) == [2, 64 // c0, 2, 16, c0] def test_expand_preserves_sf_tag(): diff --git a/testing/ascend/transform/test_ascend_insert_nd2nz.py b/testing/ascend/transform/test_ascend_insert_nd2nz.py index 3e28614bd6..101173b966 100644 --- a/testing/ascend/transform/test_ascend_insert_nd2nz.py +++ b/testing/ascend/transform/test_ascend_insert_nd2nz.py @@ -3,7 +3,7 @@ import pytest from tilelang import tvm from tilelang.ascend import language as T, transform -from tilelang.layout import make_ascend_compact_nz_layout, make_ascend_nz_layout +from tilelang.layout import make_ascend_compact_nz_layout from tvm import tirx from testing.ascend._ir import calls @@ -66,26 +66,3 @@ def test_scatter_storage_requirements(kwargs, error): transform.InsertNd2Nz()(before) else: assert len(calls(transform.InsertNd2Nz()(before), "tl.ascend_nd2nz_scatter")) == 1 - - -@pytest.mark.parametrize("factory", [make_ascend_nz_layout, make_ascend_compact_nz_layout]) -@pytest.mark.parametrize("dtype", ["float32", "bfloat16"]) -def test_post_copy_requires_compact_nz_panel_pitch(factory, dtype): - src = tirx.decl_buffer((33, 128), dtype, name="src", scope="shared.dyn") - dst = tirx.decl_buffer((32, 128), dtype, name="dst", scope="shared.l1") - region = tirx.BufferRegion(src, [tvm.ir.Range(0, 32), tvm.ir.Range(0, 128)]) - root = tirx.SBlock( - [], - [], - [], - "root", - tirx.Evaluate(T.copy(region, dst)), - alloc_buffers=[src, dst], - annotations={"layout_map": {src: factory(src), dst: make_ascend_nz_layout(dst)}}, - ) - before = tvm.IRModule({"main": tirx.PrimFunc([], tirx.SBlockRealize([], True, root))}) - if factory == make_ascend_nz_layout: - with pytest.raises(ValueError, match="requires compact NZ storage"): - transform.InsertNd2Nz()(before) - else: - assert len(calls(transform.InsertNd2Nz()(before), "tl.ascend_nd2nz_post_copy")) == 1 diff --git a/testing/ascend/transform/test_ascend_lower_copy.py b/testing/ascend/transform/test_ascend_lower_copy.py index faa0d7522c..206a488a24 100644 --- a/testing/ascend/transform/test_ascend_lower_copy.py +++ b/testing/ascend/transform/test_ascend_lower_copy.py @@ -167,7 +167,7 @@ def test_fix_nz_uses_physical_panel_pitch_and_slice_origin(dual, src_shape): src, dst, dst_region=_region(dst, [1, 32, 64], [1, 16, 16]), - layouts={src: make_ascend_l0c_layout(src), dst: make_ascend_nz_layout(dst)}, + layouts={src: make_ascend_l0c_layout(src), dst: make_ascend_l0c_layout(dst)}, annotations={"dual_dst_ctl": dual}, unit_flag_ctrl=3, ) @@ -185,14 +185,14 @@ def test_fix_nz_rejects_partial_column_panels(column, width): src, dst, dst_region=_region(dst, [32, width], [0, column]), - layouts={src: make_ascend_l0c_layout(src), dst: make_ascend_nz_layout(dst)}, + 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_nz_layout"): + 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)}) @@ -200,7 +200,7 @@ 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_nz_layout(dst)}) + _lower(src, dst, layouts={src: make_ascend_l0c_layout(src), dst: make_ascend_l0c_layout(dst)}) @pytest.mark.parametrize("scope", ["shared.l1", "shared.dyn"]) diff --git a/tilelang/layout/ascend.py b/tilelang/layout/ascend.py index 0316543e9e..6409f341b3 100644 --- a/tilelang/layout/ascend.py +++ b/tilelang/layout/ascend.py @@ -18,25 +18,16 @@ class _AscendFractalLayout(Layout): def make_ascend_nz_layout(buffer: BufferLikeType): - """Create the Ascend NZ layout appropriate for the buffer's scope. - - Logical ``[..., m, n]`` maps to - ``[..., n // C0, m // 16, m % 16, n % C0]``. L0C and UB - (``shared``/``shared.dyn``) use the fixed 16-column accumulator format. - Other scopes use the Cube input format, with - ``C0 = 32 bytes / element_size``. - - With ``T.annotate_layout({ub: make_ascend_nz_layout(ub)})``, L0C-to-UB - ``T.copy`` and ``T.dual_copy`` preserve NZ. Declare padding in the UB - shape; copy regions retain the allocation's physical panel pitch and - must cover whole 16-column panels. Buffer indices remain logical; - SIMD pointer increments access consecutive physical elements. + """ + Create an Ascend NZ (zN fractal) layout for a Cube operand buffer. - For compact UB packing before UB-to-L1 copies, use - :func:`make_ascend_compact_nz_layout` explicitly. + A logical ``[.., rows, cols]`` tile maps to physical + ``[.., rows / 16, cols / C0, 16, C0]`` -- a grid of 16xC0 fractals -- where + ``C0 = 32 bytes / element_size`` is the Cube unit's fractal width and the + forward map is ``(.., r, c) -> [.., r / 16, c / C0, r % 16, c % C0]``. Args: - buffer: BufferLikeType (rank >= 2). + buffer: BufferLikeType (rank >= 2, constant last-2 dims) Returns: Layout: The NZ fractal layout for the buffer. @@ -96,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)