diff --git a/ptodsl/docs/user_guide/04-type-system-and-buffer.md b/ptodsl/docs/user_guide/04-type-system-and-buffer.md index cd89807e43..925367f035 100644 --- a/ptodsl/docs/user_guide/04-type-system-and-buffer.md +++ b/ptodsl/docs/user_guide/04-type-system-and-buffer.md @@ -267,6 +267,46 @@ physical row-alignment rule. For packed types (`pto.f4e1m2x2`, `pto.f4e2m1x2`), `shape` dimensions refer to the number of **packed** elements, each containing 2 f4 values. For example, `alloc_tile(shape=[128, 64], dtype=pto.f4e1m2x2)` allocates a 128×64 tile of packed elements, holding 128×64×2 individual 4-bit floats. The same applies to TensorView shapes when the tensor spec uses a packed dtype. +### Multi-buffer tiles + +Use `pto.alloc_multi_tile` to allocate several independent physical slots with +the same tile shape, element type, layout, and memory space. It is useful for +double buffering and deeper software pipelines: one slot can hold the tile +currently consumed by an operation while another slot is prepared for a later +iteration. + +#### `pto.alloc_multi_tile(*, shape, dtype, memory_space="ub", count, valid_shape=None, blayout="RowMajor", slayout="NoneBox", fractal_size=512, pad="Null", addr=None) -> MultiTile` + +`shape`, `dtype`, `memory_space`, `valid_shape`, and layout arguments have the +same meaning and constraints as for `pto.alloc_tile`. `count` is a compile-time +Python integer and must be at least `2`. The allocation owns `count` distinct +tile slots, so its total storage requirement is `count` times the footprint of +one slot. + +#### `pto.multi_tile_get(multi_tile, slot) -> Tile` + +Selects one slot and returns it as a normal `Tile`. The returned tile can be +passed directly to existing tile data-movement and compute APIs. Selection does +not copy, move, or transform data; it only chooses which physical slot later +operations access. + +`slot` may be a constant or a device-side `index` expression. It must satisfy +`0 <= slot < count`; PTODSL does not automatically apply modulo arithmetic. +For a double-buffered loop, write the intended slot expression explicitly: + +```python +tiles = pto.alloc_multi_tile(shape=[BLOCK, dim], dtype=pto.f32, count=2) + +with pto.for_(0, BLOCK, step=1) as iv: + slot = iv % 2 + current_tile = pto.multi_tile_get(tiles, slot) + current_tile.fill(0.0) +``` + +In this example, even iterations select slot `0` and odd iterations select slot +`1`. Coordinate data movement and computation with the synchronization needed +by the surrounding pipeline before reusing a slot. + ### Tile attributes | Attribute | Type | Description | diff --git a/ptodsl/ptodsl/_ops.py b/ptodsl/ptodsl/_ops.py index 365ec64b04..1ee0bab2e0 100644 --- a/ptodsl/ptodsl/_ops.py +++ b/ptodsl/ptodsl/_ops.py @@ -68,6 +68,7 @@ part_tensor_view_type, part_tensor_view_type_from_dims, ptr, + tile_buf_type, tensor_view_type, tensor_view_type_from_dims, vreg_type, @@ -3003,6 +3004,84 @@ def alloc_tile( ) +def alloc_multi_tile( + *, + shape, + dtype, + memory_space="ub", + count, + valid_shape=None, + blayout: str = "RowMajor", + slayout: str = "NoneBox", + fractal_size: int = 512, + pad: str = "Null", + addr=None, +): + """Allocate a logical tile with ``count`` independently planned slots. + + A slot is selected with :func:`multi_tile_get` before it is passed to a + DMA or compute operation. The slot selector remains visible in PTOIR so + memory planning and synchronization can distinguish pipeline versions. + """ + if isinstance(count, bool) or not isinstance(count, int): + raise TypeError("alloc_multi_tile(count=...) expects a positive Python int") + if count < 2: + raise ValueError("alloc_multi_tile(count=...) requires at least two slots") + + logical_shape = _normalize_static_tile_shape(shape) + physical_shape = _authored_tile_physical_shape(logical_shape) + _validate_authored_tile_row_alignment( + physical_shape, dtype, blayout=blayout, slayout=slayout + ) + type_valid_shape, valid_row, valid_col, surface_valid_shape = _split_valid_shape( + logical_shape, valid_shape + ) + slot_type = tile_buf_type( + physical_shape, + dtype, + type_valid_shape, + blayout=blayout, + address_space=memory_space, + slayout=slayout, + fractal_size=fractal_size, + pad=pad, + ) + multi_type = Type.parse(f"!pto.multi_tile_buf<{slot_type}, count={count}>") + value = _pto.AllocMultiTileOp( + multi_type, + addr=_coerce_i64(addr, context="alloc_multi_tile(addr)") if addr is not None else None, + valid_row=_coerce_index(valid_row, context="alloc_multi_tile(valid_row)") if valid_row is not None else None, + valid_col=_coerce_index(valid_col, context="alloc_multi_tile(valid_col)") if valid_col is not None else None, + ).result + result = wrap_surface_value(value) + result.tile_metadata = { + "shape": logical_shape, + "physical_shape": physical_shape, + "dtype": dtype, + "memory_space": memory_space, + "valid_shape": surface_valid_shape, + "count": count, + } + result._multi_tile_slot_type = slot_type + return result + + +def multi_tile_get(multi_tile, slot): + """Return the tile handle for one selected multi-buffer slot.""" + source = unwrap_surface_value(multi_tile) + slot_value = _coerce_index(slot, context="multi_tile_get(slot)") + result_type = getattr(multi_tile, "_multi_tile_slot_type", None) + if result_type is None: + raise TypeError( + "multi_tile_get(multi_tile, slot) expects a value returned by " + "alloc_multi_tile()" + ) + value = _pto.MultiTileGetOp(result_type, source, slot_value).result + metadata = dict(getattr(multi_tile, "tile_metadata", {})) + metadata.pop("count", None) + return wrap_surface_value(value, tile_metadata=metadata) + + def set_tile_valid_shape(tile, valid_shape): """Update the runtime valid-shape metadata of an authored dynamic tile.""" parsed_tile_type = parse_tile_type_metadata(unwrap_surface_value(tile).type) @@ -6443,7 +6522,7 @@ def import_reserved_buffer(name, *, peer_func): "vaxpy", "vmula", "vci", "vaddrelu", "vsubrelu", "vsel", "make_tensor_view", "partition_view", - "alloc_buffer", "alloc_tile", + "alloc_buffer", "alloc_tile", "alloc_multi_tile", "multi_tile_get", "tload", "tstore", "tmov", "tinsert", "tconcat", "tmatmul", "tmatmul_acc", "tmatmul_mx", "tmatmul_mx_acc", "tmatmul_mx_bias", "tgemv_mx", "tgemv_mx_acc", "tgemv_mx_bias", diff --git a/ptodsl/ptodsl/pto.py b/ptodsl/ptodsl/pto.py index c45d521700..db908d09e6 100644 --- a/ptodsl/ptodsl/pto.py +++ b/ptodsl/ptodsl/pto.py @@ -119,7 +119,7 @@ vmula, vmadd, vsel, make_tensor_view, partition_view, - alloc_buffer, alloc_tile, + alloc_buffer, alloc_tile, alloc_multi_tile, multi_tile_get, tsort32, tmrgsort, tgather, tscatter, mte_load, mte_store, mte_gm_ub, mte_ub_gm, mte_ub_ub, mte_ub_l1, mte_gm_l1, mte_l1_ub, mte_gm_l1_frac, mte_l1_bt, mte_l1_fb, mem_bar, diff --git a/ptodsl/tests/test_jit_compile.py b/ptodsl/tests/test_jit_compile.py index 8666a82a02..21f63a36c3 100644 --- a/ptodsl/tests/test_jit_compile.py +++ b/ptodsl/tests/test_jit_compile.py @@ -594,6 +594,14 @@ def dynamic_addr_tile_surface_probe( _ = tile +@pto.jit(target="a5") +def multi_tile_surface_probe(*, BLOCK: pto.const_expr = 128, dim: pto.const_expr = 16): + tiles = pto.alloc_multi_tile(shape=[BLOCK, dim], dtype=pto.f32, count=2) + with pto.for_(0, BLOCK, step=1) as iv: + tile = pto.multi_tile_get(tiles, iv % 2) + tile.fill(0.0) + + @pto.jit(target="a5") def tile_sort_gather_surface_probe(): src = pto.alloc_tile(shape=[1, 32], dtype=pto.f32) @@ -3695,6 +3703,8 @@ def main() -> None: "mad_mx_acc", "mad_mx_bias", "empty_like", + "alloc_multi_tile", + "multi_tile_get", ] for name in expected_public_exports: expect(hasattr(pto, name), f"pto.{name} should be exported from the public namespace") @@ -5019,6 +5029,21 @@ def fake_run_ptoas_cmd(cmd, *, cwd=None): "alloc_tile(shape=..., dtype=..., addr=runtime value, valid_shape=...) should accept dynamic i64-like operands", ) + multi_tile_text = multi_tile_surface_probe.compile().mlir_text() + expect_parse_roundtrip_and_verify(multi_tile_text, "multi-tile surface specialization") + expect( + "pto.alloc_multi_tile" in multi_tile_text, + "alloc_multi_tile(...) should lower to pto.alloc_multi_tile", + ) + expect( + "pto.multi_tile_get" in multi_tile_text, + "multi_tile_get(...) should lower to pto.multi_tile_get", + ) + expect( + "count=2" in multi_tile_text, + "alloc_multi_tile(count=2) should preserve the slot count in MLIR", + ) + tile_valid_shape_text = tile_valid_shape_update_probe.compile().mlir_text() expect_parse_roundtrip_and_verify(tile_valid_shape_text, "tile valid-shape update specialization") expect(