PTODSL New Kernel Surface Design

[Zhendong404]

Problem

The current kernel story exposes too many layers:

@pto.jit
  -> Tile Ops
  -> @pto.ukernel          # extra orchestration layer
       -> MTE + sync
       -> @pto.simd / @pto.simt / @pto.cube

1. Mental model too deep. A user writing one custom vector step must learn
   Tile Ops, the hardware unit, plus an orchestration layer.
2. Hierarchy doesn't match workflow. Most users want tile orchestration
   with occasional custom compute, not low-level DMA and sync.
3. Mixes compute authorship with scheduling ownership. A user may want
   custom SIMD code while still expecting the compiler to manage staging and
   sync.
4. Expert path leaks into beginner path. The public @pto.ukernel
   suggests explicit MTE/sync is the default authoring model.

Design Summary

One entry point (@pto.jit), two modes. The distinction is what the compiler
sees:

basic (tile-atomic, compiler optimizes):
  @pto.jit -> Tile Ops -> @pto.simd / @pto.simt / @pto.cube

advance (flexible, no compiler optimization):
  @pto.jit -> Tile Ops + MTE Ops + sync ops -> @pto.simd / @pto.simt / @pto.cube

• basic: the compiler sees tile-level atomic semantics — tile.load,
  tile.store, and compute-unit calls as indivisible units. This enables
  scheduling, fusion, and memory planning optimizations.
• advance: the user writes explicit DMA and sync. The compiler treats these
  as opaque user-ordered operations and does not optimize across them.

No separate public @pto.ukernel.

Basic Mode

Default mode. The compiler sees tile-level atomic semantics: tile.load,
tile.store, and compute-unit calls are indivisible units. This contract
enables the compiler to perform scheduling, fusion, and memory planning
optimizations.

Allowed: Tile Ops, named @pto.simd/@pto.simt/@pto.cube sub-kernels,
inline with pto.simd(): scopes, scalar ops.

Not allowed: MTE Ops, explicit DMA scheduling.

Parameter contract: compute-unit sub-kernel boundaries prefer Tile + scalar.

Example

@pto.simd
def add_rows(a_tile: pto.Tile, b_tile: pto.Tile, o_tile: pto.Tile,
             rows: pto.index, cols: pto.index):
    VEC = pto.elements_per_vreg(pto.f32)
    remained0 = scalar.index_cast(pto.i32, cols)
    with pto.for_(0, rows, step=1) as r:
        col_loop = pto.for_(0, cols, step=VEC).carry(remained=remained0)
        with col_loop:
            mask, remained = pto.make_mask(pto.f32, col_loop.remained)
            a_vec = pto.vlds(a_tile[r, col_loop.iv:])
            b_vec = pto.vlds(b_tile[r, col_loop.iv:])
            o_vec = pto.vadd(a_vec, b_vec, mask)
            pto.vsts(o_vec, o_tile[r, col_loop.iv:], mask)
            col_loop.update(remained=remained)

@pto.jit(target="a5")
def kernel(A: pto.tensor_spec(rank=2, dtype=pto.f32),
           B: pto.tensor_spec(rank=2, dtype=pto.f32),
           O: pto.tensor_spec(rank=2, dtype=pto.f32)):
    a_view = pto.make_tensor_view(A, shape=A.shape, strides=A.strides)
    b_view = pto.make_tensor_view(B, shape=B.shape, strides=B.strides)
    o_view = pto.make_tensor_view(O, shape=O.shape, strides=O.strides)

    a_tile = pto.alloc_tile(shape=[1, 128], dtype=pto.f32)
    b_tile = pto.alloc_tile(shape=[1, 128], dtype=pto.f32)
    o_tile = pto.alloc_tile(shape=[1, 128], dtype=pto.f32)

    pto.tile.load(pto.partition_view(a_view, offsets=[0,0], sizes=[1,128]), a_tile)
    pto.tile.load(pto.partition_view(b_view, offsets=[0,0], sizes=[1,128]), b_tile)
    add_rows(a_tile, b_tile, o_tile, 1, 128)
    pto.tile.store(o_tile, pto.partition_view(o_view, offsets=[0,0], sizes=[1,128]))

Advance Mode

Expert mode. User writes explicit DMA, barriers, and pipe ordering for maximum
flexibility. The compiler treats these as opaque user-ordered operations and
does not optimize across them.

Allowed: everything in basic plus MTE Ops, explicit sync, mixed
Tile+MTE scheduling.

Synchronization contract: user owns correctness in regions with explicit
MTE/sync. Rule: presence of user-authored MTE or sync = user-managed
orchestration. The compiler doesn't silently "fix" expert scheduling.

Parameter contract: broader PTODSL-traceable values allowed beyond
Tile + scalar, but still constrained to values the tracer can reason about.

Example

@pto.jit(target="a5")
def kernel_expert(part, tile, out_part, rows: pto.i32, cols: pto.i32):
    row_bytes = cols * pto.bytewidth(pto.f32)

    pto.mte_load(part.as_ptr(), tile.as_ptr(), 0, row_bytes,
                 nburst=(rows, part.strides[0] * pto.bytewidth(pto.f32), row_bytes))
    pto.pipe_barrier(pto.Pipe.ALL)

    with pto.simd():
        VEC = pto.elements_per_vreg(pto.f32)
        with pto.for_(0, rows, step=1) as r:
            col_loop = pto.for_(0, cols, step=VEC).carry(remained=cols)
            with col_loop:
                mask, remained = pto.make_mask(pto.f32, col_loop.remained)
                x = pto.vlds(tile[r, col_loop.iv:])
                pto.vsts(pto.vmul(x, x, mask), tile[r, col_loop.iv:], mask)
                col_loop.update(remained=remained)

    pto.pipe_barrier(pto.Pipe.ALL)
    pto.mte_store(tile.as_ptr(), out_part.as_ptr(), row_bytes,
                  nburst=(rows, row_bytes, out_part.strides[0] * pto.bytewidth(pto.f32)))

Boundary Summary

Topic	basic	advance
Entry point	@pto.jit	@pto.jit
Style	Tile-first	Orchestration-first
MTE Ops	Not allowed	Allowed
Compiler optimization	Yes (tile-atomic semantics)	No (opaque user ops)
Compute-unit calls / inline	Allowed	Allowed
Compute-unit boundary	Tile + scalar	Broader PTODSL-traceable values

Both modes support named @pto.simd/@pto.simt/@pto.cube sub-kernels and
inline with pto.simd(): scopes equally. The distinction is tile-atomic
semantics vs. raw flexibility, not reuse vs. inline.

[learning-chip]

def kernel(A: pto.tensor_spec(rank=2, dtype=pto.f32),
           B: pto.tensor_spec(rank=2, dtype=pto.f32),
           O: pto.tensor_spec(rank=2, dtype=pto.f32)):
    a_view = pto.make_tensor_view(A, shape=A.shape, strides=A.strides)
    b_view = pto.make_tensor_view(B, shape=B.shape, strides=B.strides)
    o_view = pto.make_tensor_view(O, shape=O.shape, strides=O.strides)

Why need to specify rank=2? The rank can be inferred from the shape later.

• basic: the compiler sees tile-level atomic semantics — tile.load,
  tile.store, and compute-unit calls as indivisible units. This enables
  scheduling, fusion, and memory planning optimizations.
• advance: the user writes explicit DMA and sync. The compiler treats these
  as opaque user-ordered operations and does not optimize across them.

Only two modes basic advance can be too restricted for long-term evolution. How about put each compile passes as @jit options:

# inserts set/wait barrier, no need for user to write
@jit(compile_pass={"auto_sync": True})
def my_kernel():

# if only few options, then flatten the args
@jit(auto_sync=True)
def my_kernel():

If some APIs are not allowed in auto-sync mode (for example is pto.mte_load is too low-level for InsertSync pass??), then just raise error auto_sync mode does not allow low-level "mte_load", use "tile.load" instead when trying to use @jit(auto_sync=True)

No separate public @pto.ukernel.

This is good, we should reduce redundant decorators, unless they are absolutely necessary.