How to reuse buffer space for `pto.AllocTileOp` with varying valid_row/valid_col ?

[learning-chip]

In huawei-csl/pto-dsl#12 I made an arbitrary-dynamic-shape kernel example by masking-out out-of-bound tiles and elements.

However, I had to adjust valid_col inside the for-loop:

for i in scf.for_(c0, tiles_to_process, c1):
    ...
    # change`c_tile_actual` according to remaining elements
    tb0 = pto.AllocTileOp(
        tile_buf_dynamic, valid_row=c1, valid_col=c_tile_actual
    ).result

Because AllocTileOp is lowered to new Tile allocation + TASSIGN, VEC UB space is wasted, not reused across iterations.

Note that buffer-reuse is doable if I don't need to adjust valid_shape, ref this less-flexible version of dynamic shape kernel: huawei-csl/pto-dsl#11. But supporting arbitrary shape is very important for practical use cases, such as for dynamic batch dim and sequence dim.

Proposed solution

Provide a set_valid_shape API to modify the valid_shape attribute of an already-allocated tile?

[HecreReed]

Fixed and verified in PR #110: tile lowering now preserves runtime valid_row/valid_col so dynamic valid shapes are not dropped (auto-closes on merge).

[learning-chip]

tile lowering now preserves runtime valid_row/valid_col so dynamic valid shapes are not dropped (auto-closes on merge).

Thanks, I think the issue is more than just preserving the valid shape attribute, but about in-place modifying the valid shape of an already-allocated Tile buffer.

For full example, see those lines of add_builder.py from huawei-csl/pto-dsl#12. Can be reproduced by this README.md.

In short, the current code pattern is:

for i in scf.for_(c0, tiles_to_process, c1):
    ...
    # change`c_tile_actual` according to remaining elements
    tb0 = pto.AllocTileOp(
        tile_buf_dynamic, valid_row=c1, valid_col=c_tile_actual
    ).result  # allocates new UB space everytime!
    pto.TLoadOp(None, sv0, tb0)
    pto.TAddOp(tb0, tb1, tb2)  # only apply to valid subset

To reuse the allocated tb0 space across iters, and still mask-out tail-elements in the last iteration, the code should become:

tb0 = pto.AllocTileOp(
    tile_buf_dynamic, valid_row=-1, valid_col=-1
).result   # only allocating UB once here!
for i in scf.for_(c0, tiles_to_process, c1):
    ...
    # change`c_tile_actual` according to remaining elements
   tb0.set_valid_shape(valid_row=c1, valid_col=c_tile_actual)  # non-existing API
   pto.TLoadOp(None, sv0, tb0)
   pto.TAddOp(tb0, tb1, tb2)  # only apply to valid subset

Tile buffer reuse is possible because valid shape basically means calling set_vector_mask before calling vector instructions. It doesn't need a new buffer allocation...

In comparison, PartitionViewOp has the more desired behavior, as it creates a non-copy view:

for i in scf.for_(c0, tiles_to_process, c1):
    ...
    # change`c_tile_actual` according to remaining elements
    sv0 = pto.PartitionViewOp(
        tile_view,
        tv0,
        offsets=[offset_global],
        sizes=[c_tile_actual],
    ).result

[learning-chip]

well, this use case (dynamic any-shape kernel) actually needs an PartitionViewOp-equivalent that works on Tile type rather GlobalTensor. Maybe same as #107 ?

Code pattern with a mock pto.TileViewOp API:

tb0 = pto.AllocTileOp(
    tile_buf_dynamic, valid_row=-1, valid_col=-1
).result   # only allocating UB once here!
for i in scf.for_(c0, tiles_to_process, c1):
    ...
    # change`c_tile_actual` according to remaining elements
    sv0 = pto.PartitionViewOp(
        tile_view,
        tv0,
        offsets=[offset_global],
        sizes=[c_tile_actual],
    ).result  # zero-copy view on GM tensor
   tb0_view = pto.TileViewOp(tb0, offsets=[c0, c0], sizes=[c1, c_tile_actual]).result  # zero-copy view on UB tile
   pto.TLoadOp(None, sv0, tb0_view)  # works on subset of both GM and UB data
   pto.TAddOp(tb0_view, tb1_view, tb2_view)  # works on subset of both GM and UB data

[learning-chip]

Maybe same as #107 ?

Ok, the pto.SubsetOp used in test/samples/Subset seems exactly what I need.

No, the sizes parameter to SubsetOp must be static, but c_tile_actual is dynamically computed at run-time.

PTOAS/test/samples/Subset/vadd_pto_pingpong.py

Lines 58 to 60 in a301aa4

	# Subset sizes must be static (I64ArrayAttr); offsets must be SSA.
	ping = pto.SubsetOp(workspace, [c0, c0], sizes=[32, 32]).result
	pong = pto.SubsetOp(workspace, [c0, c32], sizes=[32, 32]).result

@HecreReed any suggestions?

[HecreReed]

You are right: pto.subset only supports static sizes today, so dynamic c_tile_actual is not possible yet; I will keep this issue open and follow up with a UB Tile view / in-place valid-shape update proposal.

[learning-chip]

You are right: pto.subset only supports static sizes today, so dynamic c_tile_actual is not possible yet; I will keep this issue open and follow up with a UB Tile view / in-place valid-shape update proposal.

ok, for now I'll limit the input shapes to be [batch, hidden_dim] where hidden_dim needs to be a nice multiple of tile_size, while batch can be any integer, to temporarily avoid the tail element issue.

[zhangstevenunity]

#236, this pr add pto.set_validshape to support this issue requirement