feat(ptodsl): add multi-tile allocation helpers

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 |

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",

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,

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(