Pooling Operators¶

Every op on this page is used the same way: construct it once, then call it. The constructor takes what the kernel is compiled with; the call takes the tensors. Both are documented under each op — __init__ and forward, where forward is what runs when you call op(...).

Average pooling¶

tileops.pool.AvgPool1dFwdOp ¶

Average pooling over PyTorch-compatible NCL inputs.

init ¶

__init__(
    kernel_size,
    stride=None,
    padding=0,
    ceil_mode=False,
    count_include_pad=True,
    *,
    target=None
)

Build the op. Shapes and dtype are taken from the first call.

Parameters:

kernel_size (int | Tuple[int]) –

Manifest params.kernel_size, int | tuple[int].
stride (Optional[int | Tuple[int]], default: None ) –

Manifest params.stride, int | tuple[int] | None, default None.
padding (int | Tuple[int], default: 0 ) –

Manifest params.padding, int | tuple[int], default 0.
ceil_mode (bool, default: False ) –

Manifest params.ceil_mode, bool, default False.
count_include_pad (bool, default: True ) –

Manifest params.count_include_pad, bool, default True.
target (Target, default: None ) –

Backend target to serve this op, or None to decide from the input device.

forward ¶

forward(
    input,
)

Run the op on input.

tileops.pool.AvgPool2dFwdOp ¶

Average pooling over PyTorch-compatible NCHW inputs.

init ¶

__init__(
    kernel_size,
    stride=None,
    padding=0,
    ceil_mode=False,
    count_include_pad=True,
    divisor_override=None,
    *,
    target=None
)

Build the op. Shapes and dtype are taken from the first call.

Parameters:

kernel_size (int | Tuple[int, int]) –

Manifest params.kernel_size, int | tuple[int, int].
stride (Optional[int | Tuple[int, int]], default: None ) –

Manifest params.stride, int | tuple[int, int] | None, default None.
padding (int | Tuple[int, int], default: 0 ) –

Manifest params.padding, int | tuple[int, int], default 0.
ceil_mode (bool, default: False ) –

Manifest params.ceil_mode, bool, default False.
count_include_pad (bool, default: True ) –

Manifest params.count_include_pad, bool, default True.
divisor_override (Optional[int], default: None ) –

Manifest params.divisor_override, int | None, default None.
target (Target, default: None ) –

Backend target to serve this op, or None to decide from the input device.

forward ¶

forward(
    input,
)

Run the op on input.

tileops.pool.AvgPool3dFwdOp ¶

Average pooling over PyTorch-compatible NCDHW inputs.

init ¶

__init__(
    kernel_size,
    stride=None,
    padding=0,
    ceil_mode=False,
    count_include_pad=True,
    divisor_override=None,
    *,
    target=None
)

Build the op. Shapes and dtype are taken from the first call.

Parameters:

kernel_size (int | Tuple[int, int, int]) –

Manifest params.kernel_size, int | tuple[int, int, int].
stride (Optional[int | Tuple[int, int, int]], default: None ) –

Manifest params.stride, int | tuple[int, int, int] | None, default None.
padding (int | Tuple[int, int, int], default: 0 ) –

Manifest params.padding, int | tuple[int, int, int], default 0.
ceil_mode (bool, default: False ) –

Manifest params.ceil_mode, bool, default False.
count_include_pad (bool, default: True ) –

Manifest params.count_include_pad, bool, default True.
divisor_override (Optional[int], default: None ) –

Manifest params.divisor_override, int | None, default None.
target (Target, default: None ) –

Backend target to serve this op, or None to decide from the input device.

forward ¶

forward(
    input,
)

Run the op on input.

Max pooling¶

tileops.pool.MaxPool1dFwdOp ¶

Max pooling over PyTorch-compatible NCL inputs (return_indices=False).

init ¶

__init__(
    kernel_size,
    stride=None,
    padding=0,
    dilation=1,
    ceil_mode=False,
    *,
    target=None
)

Build the op. Shapes and dtype are taken from the first call.

Parameters:

kernel_size (int | Tuple[int]) –

Manifest params.kernel_size, int | tuple[int].
stride (Optional[int | Tuple[int]], default: None ) –

Manifest params.stride, int | tuple[int] | None, default None.
padding (int | Tuple[int], default: 0 ) –

Manifest params.padding, int | tuple[int], default 0.
dilation (int | Tuple[int], default: 1 ) –

Manifest params.dilation, int | tuple[int], default 1.
ceil_mode (bool, default: False ) –

Manifest params.ceil_mode, bool, default False.
target (Target, default: None ) –

Backend target to serve this op, or None to decide from the input device.

forward ¶

forward(
    input,
)

Run the op on input.

tileops.pool.MaxPool1dIndicesFwdOp ¶

Max pooling over PyTorch-compatible NCL inputs (return_indices=True).

init ¶

__init__(
    kernel_size,
    stride=None,
    padding=0,
    dilation=1,
    ceil_mode=False,
    *,
    target=None
)

Build the op. Shapes and dtype are taken from the first call.

Parameters:

kernel_size (int | Tuple[int]) –

Manifest params.kernel_size, int | tuple[int].
stride (Optional[int | Tuple[int]], default: None ) –

Manifest params.stride, int | tuple[int] | None, default None.
padding (int | Tuple[int], default: 0 ) –

Manifest params.padding, int | tuple[int], default 0.
dilation (int | Tuple[int], default: 1 ) –

Manifest params.dilation, int | tuple[int], default 1.
ceil_mode (bool, default: False ) –

Manifest params.ceil_mode, bool, default False.
target (Target, default: None ) –

Backend target to serve this op, or None to decide from the input device.

forward ¶

forward(
    input,
)

Run the op on the inputs the manifest declares.

Parameters:

input (Tensor) –

Input tensor, dtype float16 | bfloat16 | float32.

Returns:

Tuple[Tensor, Tensor] –

output, indices, as the manifest declares.

tileops.pool.MaxPool2dFwdOp ¶

Max pooling over PyTorch-compatible NCHW inputs (return_indices=False).

init ¶

__init__(
    kernel_size,
    stride=None,
    padding=0,
    dilation=1,
    ceil_mode=False,
    *,
    target=None
)

Build the op. Shapes and dtype are taken from the first call.

Parameters:

kernel_size (int | Tuple[int, int]) –

Manifest params.kernel_size, int | tuple[int, int].
stride (Optional[int | Tuple[int, int]], default: None ) –

Manifest params.stride, int | tuple[int, int] | None, default None.
padding (int | Tuple[int, int], default: 0 ) –

Manifest params.padding, int | tuple[int, int], default 0.
dilation (int | Tuple[int, int], default: 1 ) –

Manifest params.dilation, int | tuple[int, int], default 1.
ceil_mode (bool, default: False ) –

Manifest params.ceil_mode, bool, default False.
target (Target, default: None ) –

Backend target to serve this op, or None to decide from the input device.

forward ¶

forward(
    input,
)

Run the op on input.

tileops.pool.MaxPool2dIndicesFwdOp ¶

Max pooling over PyTorch-compatible NCHW inputs (return_indices=True).

init ¶

__init__(
    kernel_size,
    stride=None,
    padding=0,
    dilation=1,
    ceil_mode=False,
    *,
    target=None
)

Build the op. Shapes and dtype are taken from the first call.

Parameters:

kernel_size (int | Tuple[int, int]) –

Manifest params.kernel_size, int | tuple[int, int].
stride (Optional[int | Tuple[int, int]], default: None ) –

Manifest params.stride, int | tuple[int, int] | None, default None.
padding (int | Tuple[int, int], default: 0 ) –

Manifest params.padding, int | tuple[int, int], default 0.
dilation (int | Tuple[int, int], default: 1 ) –

Manifest params.dilation, int | tuple[int, int], default 1.
ceil_mode (bool, default: False ) –

Manifest params.ceil_mode, bool, default False.
target (Target, default: None ) –

Backend target to serve this op, or None to decide from the input device.

forward ¶

forward(
    input,
)

Run the op on the inputs the manifest declares.

Parameters:

input (Tensor) –

Input tensor, dtype float16 | bfloat16 | float32.

Returns:

Tuple[Tensor, Tensor] –

output, indices, as the manifest declares.

tileops.pool.MaxPool3dFwdOp ¶

Max pooling over PyTorch-compatible NCDHW inputs (return_indices=False).

init ¶

__init__(
    kernel_size,
    stride=None,
    padding=0,
    dilation=1,
    ceil_mode=False,
    *,
    target=None
)

Build the op. Shapes and dtype are taken from the first call.

Parameters:

kernel_size (int | Tuple[int, int, int]) –

Manifest params.kernel_size, int | tuple[int, int, int].
stride (Optional[int | Tuple[int, int, int]], default: None ) –

Manifest params.stride, int | tuple[int, int, int] | None, default None.
padding (int | Tuple[int, int, int], default: 0 ) –

Manifest params.padding, int | tuple[int, int, int], default 0.
dilation (int | Tuple[int, int, int], default: 1 ) –

Manifest params.dilation, int | tuple[int, int, int], default 1.
ceil_mode (bool, default: False ) –

Manifest params.ceil_mode, bool, default False.
target (Target, default: None ) –

Backend target to serve this op, or None to decide from the input device.

forward ¶

forward(
    input,
)

Run the op on input.

tileops.pool.MaxPool3dIndicesFwdOp ¶

Max pooling over PyTorch-compatible NCDHW inputs (return_indices=True).

init ¶

__init__(
    kernel_size,
    stride=None,
    padding=0,
    dilation=1,
    ceil_mode=False,
    *,
    target=None
)

Build the op. Shapes and dtype are taken from the first call.

Parameters:

kernel_size (int | Tuple[int, int, int]) –

Manifest params.kernel_size, int | tuple[int, int, int].
stride (Optional[int | Tuple[int, int, int]], default: None ) –

Manifest params.stride, int | tuple[int, int, int] | None, default None.
padding (int | Tuple[int, int, int], default: 0 ) –

Manifest params.padding, int | tuple[int, int, int], default 0.
dilation (int | Tuple[int, int, int], default: 1 ) –

Manifest params.dilation, int | tuple[int, int, int], default 1.
ceil_mode (bool, default: False ) –

Manifest params.ceil_mode, bool, default False.
target (Target, default: None ) –

Backend target to serve this op, or None to decide from the input device.

forward ¶

forward(
    input,
)

Run the op on the inputs the manifest declares.

Parameters:

input (Tensor) –

Input tensor, dtype float16 | bfloat16 | float32.

Returns:

Tuple[Tensor, Tensor] –

output, indices, as the manifest declares.

Adaptive pooling¶

tileops.pool.AdaptiveAvgPool2dFwdOp ¶

Adaptive average pooling over PyTorch-compatible CHW/NCHW inputs.

init ¶

__init__(
    output_size,
    *,
    target=None
)

Build the op. Shapes and dtype are taken from the first call.

Parameters:

output_size (int | None | Tuple[Optional[int], Optional[int]] | list[Optional[int]]) –

Manifest params.output_size, int | None | tuple[int | None, int | None] | list[int | None].
target (Target, default: None ) –

Backend target to serve this op, or None to decide from the input device.

forward ¶

forward(
    input,
)

Run the op on input.

tileops.pool.AdaptiveMaxPool2dFwdOp ¶

Adaptive max pooling over CHW/NCHW inputs (return_indices=False).

init ¶

__init__(
    output_size,
    *,
    target=None
)

Build the op. Shapes and dtype are taken from the first call.

Parameters:

output_size (int | None | Tuple[Optional[int], Optional[int]] | list[Optional[int]]) –

Manifest params.output_size, int | None | tuple[int | None, int | None] | list[int | None].
target (Target, default: None ) –

Backend target to serve this op, or None to decide from the input device.

forward ¶

forward(
    input,
)

Run the op on input.

tileops.pool.AdaptiveMaxPool2dIndicesFwdOp ¶

Adaptive max pooling over CHW/NCHW inputs (return_indices=True).

init ¶

__init__(
    output_size,
    *,
    target=None
)

Build the op. Shapes and dtype are taken from the first call.

Parameters:

output_size (int | None | Tuple[Optional[int], Optional[int]] | list[Optional[int]]) –

Manifest params.output_size, int | None | tuple[int | None, int | None] | list[int | None].
target (Target, default: None ) –

Backend target to serve this op, or None to decide from the input device.

forward ¶

forward(
    input,
)

Run the op on the inputs the manifest declares.

Parameters:

input (Tensor) –

Input tensor, dtype float16 | bfloat16.

Returns:

Tuple[Tensor, Tensor] –

output, indices, as the manifest declares.

Chunked sequence mean¶

The one op on this page with no PyTorch counterpart. It averages the sequence axis of a [batch, seq, heads, dim] tensor in fixed-size chunks, and can follow ragged sequence boundaries rather than a uniform split.

tileops.pool.MeanPoolingFwdOp ¶

Chunked mean over the sequence axis of a [batch, seq, heads, dim] tensor.

Not a PyTorch pooling op, and torch has no counterpart. The sequence axis is cut into chunks of chunk_size and each chunk is averaged, giving one output row per chunk. Pass offsets and indices and the chunks follow the ragged sequence boundaries offsets describes instead of a uniform split; indices then names the (sequence, chunk-within-sequence) pair each output row belongs to, one row per chunk.

A sequence's last chunk may be shorter than chunk_size. It is divided by the count it actually holds, so no padding is averaged in. On a uniform split that makes the op equal to torch.nn.functional.avg_pool1d(kernel_size=chunk_size, stride=chunk_size, ceil_mode=True) over a view with the sequence axis last. Chunk sums accumulate in accum_dtype and are cast back to the input dtype at the boundary, so a float16 input with accum_dtype=torch.float32 does not lose the sum to rounding.

By default the op does not check the contents of offsets and indices. The kernel indexes offsets by indices and x by offsets without bounds checks, so an inconsistent map can cause out-of-bounds accesses or incorrect results. The caller guarantees a consistent map, or passes validate_inputs=True, which checks it on every call at the cost of device synchronizations and cannot run inside CUDA Graph capture.

Example

op = MeanPoolingFwdOp(chunk_size=64, accum_dtype=torch.float32)
chunk_means = op(x)                        # uniform split
chunk_means = op(x, offsets, indices)      # ragged, one row of indices per chunk

init ¶

__init__(
    chunk_size,
    accum_dtype,
    *,
    validate_inputs=False,
    target=None
)

Build the op. Shapes and dtype are taken from the first call.

Parameters:

chunk_size (int) –

Manifest params.chunk_size, int, a positive multiple of 32.
accum_dtype (dtype) –

Manifest params.accum_dtype, torch.dtype — what a chunk sum accumulates in.
validate_inputs (bool, default: False ) –

Check ragged metadata contents synchronously on each call. Disable during CUDA Graph capture; shapes and dtypes are always checked.
target (Target, default: None ) –

Backend target to serve this op, or None to decide from the input device.

forward ¶

forward(
    x,
    offsets=None,
    indices=None,
)

Average each chunk of x's sequence axis.

Parameters:

x (Tensor) –

Input tensor, [batch, seq, heads, dim], dtype float16, bfloat16 or float32. dim is at most 128 or a multiple of it, which is what the manifest declares rather than what the kernel needs.
offsets (Optional[Tensor], default: None ) –

Sequence boundaries, [seq_num + 1], dtype int32. Passing it selects the ragged split, and it comes with indices.
indices (Optional[Tensor], default: None ) –

One (sequence, chunk-within-sequence) pair per output chunk, [chunks, 2], dtype int32.

Returns:

Tensor –

output, [batch, chunks, heads, dim], dtype as x. chunks is
Tensor –

ceil(seq / chunk_size) for a uniform split and indices.shape[0] for a
Tensor –

ragged one.

Raises:

ValueError –

With validate_inputs=True, indices disagrees with offsets or offsets does not partition the sequence axis.

Pooling Operators¶

Average pooling¶

tileops.pool.AvgPool1dFwdOp ¶

__init__ ¶

forward ¶

tileops.pool.AvgPool2dFwdOp ¶

__init__ ¶

forward ¶

tileops.pool.AvgPool3dFwdOp ¶

__init__ ¶

forward ¶

Max pooling¶

tileops.pool.MaxPool1dFwdOp ¶

__init__ ¶

forward ¶

tileops.pool.MaxPool1dIndicesFwdOp ¶

__init__ ¶

forward ¶

tileops.pool.MaxPool2dFwdOp ¶

__init__ ¶

forward ¶

tileops.pool.MaxPool2dIndicesFwdOp ¶

__init__ ¶

forward ¶

tileops.pool.MaxPool3dFwdOp ¶

__init__ ¶

forward ¶

tileops.pool.MaxPool3dIndicesFwdOp ¶

__init__ ¶

forward ¶

Adaptive pooling¶

tileops.pool.AdaptiveAvgPool2dFwdOp ¶

__init__ ¶

forward ¶

tileops.pool.AdaptiveMaxPool2dFwdOp ¶

__init__ ¶

forward ¶

tileops.pool.AdaptiveMaxPool2dIndicesFwdOp ¶

__init__ ¶

forward ¶

Chunked sequence mean¶

tileops.pool.MeanPoolingFwdOp ¶

__init__ ¶

forward ¶

init ¶

init ¶

init ¶

init ¶

init ¶

init ¶

init ¶

init ¶

init ¶

init ¶

init ¶

init ¶

init ¶