Reduction 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(...).

Sums, means and extrema

tileops.reduction.SumFwdOp

Sum over dim, following torch.sum.

__init__

__init__(
    dim=None,
    keepdim=False,
    *,
    dtype=None,
    target=None
)

Construct the op.

Parameters:

• dim (Dim, default: None ) –

  Axes to reduce: an int, a sequence of them, or None for all.

• keepdim (bool, default: False ) –

  Whether a reduced axis stays as a length-1 axis.

• dtype (Optional[dtype], default: None ) –

  The dtype the input is cast to before the reduction, and the output's; None keeps the input's.

• target (Target, default: None ) –

  Which set of kernels serves this op — a target name, BUILTIN for the in-tree kernels, or None to decide from the input device.

forward

forward(
    x,
)

Reduce x over the configured axes.

Parameters:

• x (Tensor) –

  Input tensor of any rank.

Returns:

• Tensor –

  The reduction, shaped by dim and keepdim.

tileops.reduction.MeanFwdOp

Mean over dim, following torch.mean; an empty reduction is NaN.

__init__

__init__(
    dim=None,
    keepdim=False,
    *,
    dtype=None,
    target=None
)

Construct the op.

Parameters:

• dim (Dim, default: None ) –

  Axes to reduce: an int, a sequence of them, or None for all.

• keepdim (bool, default: False ) –

  Whether a reduced axis stays as a length-1 axis.

• dtype (Optional[dtype], default: None ) –

  The dtype the input is cast to before the reduction, and the output's; None keeps the input's.

• target (Target, default: None ) –

  Which set of kernels serves this op — a target name, BUILTIN for the in-tree kernels, or None to decide from the input device.

forward

forward(
    x,
)

Reduce x over the configured axes.

Parameters:

• x (Tensor) –

  Input tensor of any rank.

Returns:

• Tensor –

  The reduction, shaped by dim and keepdim.

tileops.reduction.ProdFwdOp

Product over one axis, following torch.prod(input, dim, keepdim, *, dtype).

__init__

__init__(
    dim,
    keepdim=False,
    *,
    dtype=None,
    target=None
)

Construct ProdFwdOp.

Parameters:

• dim (int) –

  The axis to reduce.

• keepdim (bool, default: False ) –

  Whether the reduced axis stays as a length-1 axis.

• dtype (Optional[dtype], default: None ) –

  The dtype the input is cast to before the reduction, and the output's; None keeps the input's.

• target (Target, default: None ) –

  Which set of kernels serves this op — a target name, BUILTIN for the in-tree kernels, or None to decide from the input device.

forward

forward(
    x,
)

Reduce x over the configured axes.

Parameters:

• x (Tensor) –

  Input tensor of any rank.

Returns:

• Tensor –

  The reduction, shaped by dim and keepdim.

tileops.reduction.AmaxFwdOp

Maximum over dim, following torch.amax.

__init__

__init__(
    dim=None,
    keepdim=False,
    *,
    target=None
)

Construct a reduce op.

Parameters:

• dim (Dim, default: None ) –

  Axes to reduce: an int, a sequence of them, or None for all.

• keepdim (bool, default: False ) –

  Whether a reduced axis stays as a length-1 axis.

• target (Target, default: None ) –

  Which set of kernels serves this op — a target name, BUILTIN for the in-tree kernels, or None to decide from the input device.

forward

forward(
    x,
)

Reduce x over the configured axes.

Parameters:

• x (Tensor) –

  Input tensor of any rank.

Returns:

• Tensor –

  The reduction, shaped by dim and keepdim.

tileops.reduction.AminFwdOp

Minimum over dim, following torch.amin.

__init__

__init__(
    dim=None,
    keepdim=False,
    *,
    target=None
)

Construct a reduce op.

Parameters:

• dim (Dim, default: None ) –

  Axes to reduce: an int, a sequence of them, or None for all.

• keepdim (bool, default: False ) –

  Whether a reduced axis stays as a length-1 axis.

• target (Target, default: None ) –

  Which set of kernels serves this op — a target name, BUILTIN for the in-tree kernels, or None to decide from the input device.

forward

forward(
    x,
)

Reduce x over the configured axes.

Parameters:

• x (Tensor) –

  Input tensor of any rank.

Returns:

• Tensor –

  The reduction, shaped by dim and keepdim.

Variance and deviation

tileops.reduction.VarFwdOp

Variance over dim, following torch.var.

__init__

__init__(
    dim=None,
    *,
    correction=1,
    keepdim=False,
    target=None
)

Construct a variance-family op.

Parameters:

• dim (Dim, default: None ) –

  Axes to reduce: an int, a sequence of them, or None for all.

• correction (float | None, default: 1 ) –

  Difference between the sample size and the degrees of freedom; None means 1, as in torch.

• keepdim (bool, default: False ) –

  Whether a reduced axis stays as a length-1 axis.

• target (Target, default: None ) –

  Which set of kernels serves this op — a target name, BUILTIN for the in-tree kernels, or None to decide from the input device.

forward

forward(
    x,
)

Reduce x over the configured axes.

Parameters:

• x (Tensor) –

  Input tensor of any rank.

Returns:

• Tensor –

  The reduction, shaped by dim and keepdim.

tileops.reduction.VarMeanFwdOp

Variance and mean over dim, following torch.var_mean.

__init__

__init__(
    dim=None,
    *,
    correction=1,
    keepdim=False,
    target=None
)

Construct a variance-family op.

Parameters:

• dim (Dim, default: None ) –

  Axes to reduce: an int, a sequence of them, or None for all.

• correction (float | None, default: 1 ) –

  Difference between the sample size and the degrees of freedom; None means 1, as in torch.

• keepdim (bool, default: False ) –

  Whether a reduced axis stays as a length-1 axis.

• target (Target, default: None ) –

  Which set of kernels serves this op — a target name, BUILTIN for the in-tree kernels, or None to decide from the input device.

forward

forward(
    x,
)

Reduce x over the configured axes.

Parameters:

• x (Tensor) –

  Input tensor of any rank.

Returns:

• Tensor –

  The reduction, shaped by dim and keepdim.

tileops.reduction.StdFwdOp

Standard deviation over dim, following torch.std.

__init__

__init__(
    dim=None,
    *,
    correction=1,
    keepdim=False,
    target=None
)

Construct a variance-family op.

Parameters:

• dim (Dim, default: None ) –

  Axes to reduce: an int, a sequence of them, or None for all.

• correction (float | None, default: 1 ) –

  Difference between the sample size and the degrees of freedom; None means 1, as in torch.

• keepdim (bool, default: False ) –

  Whether a reduced axis stays as a length-1 axis.

• target (Target, default: None ) –

  Which set of kernels serves this op — a target name, BUILTIN for the in-tree kernels, or None to decide from the input device.

forward

forward(
    x,
)

Reduce x over the configured axes.

Parameters:

• x (Tensor) –

  Input tensor of any rank.

Returns:

• Tensor –

  The reduction, shaped by dim and keepdim.

Arg reductions

tileops.reduction.ArgmaxFwdOp

Index of the maximum along dim, following torch.argmax; returns int64.

__init__

__init__(
    dim=None,
    keepdim=False,
    *,
    target=None
)

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

Parameters:

• dim (Optional[int], default: None ) –

  Reduction axis. None (the default) returns the index into the flattened input, as torch.argmax(x) does.

• keepdim (bool, default: False ) –

  Whether to retain the reduced dimension as size 1.

• target (Target, default: None ) –

  Which set of kernels serves this op — a target name, BUILTIN for the in-tree kernels, or None to decide from the input device.

forward

forward(
    x,
)

Reduce x over the configured axes.

Parameters:

• x (Tensor) –

  Input tensor of any rank.

Returns:

• Tensor –

  The reduction, shaped by dim and keepdim.

tileops.reduction.ArgminFwdOp

Index of the minimum along dim, following torch.argmin; returns int64.

__init__

__init__(
    dim=None,
    keepdim=False,
    *,
    target=None
)

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

Parameters:

• dim (Optional[int], default: None ) –

  Reduction axis. None (the default) returns the index into the flattened input, as torch.argmax(x) does.

• keepdim (bool, default: False ) –

  Whether to retain the reduced dimension as size 1.

• target (Target, default: None ) –

  Which set of kernels serves this op — a target name, BUILTIN for the in-tree kernels, or None to decide from the input device.

forward

forward(
    x,
)

Reduce x over the configured axes.

Parameters:

• x (Tensor) –

  Input tensor of any rank.

Returns:

• Tensor –

  The reduction, shaped by dim and keepdim.

Softmax

tileops.reduction.SoftmaxFwdOp

Softmax along dim, following torch.nn.functional.softmax.

__init__

__init__(
    dim=None,
    *,
    dtype=None,
    target=None
)

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

Parameters:

• dim (Optional[int], default: None ) –

  The axis to normalize along. None takes torch's implicit axis (0 at rank 0, 1 or 3, else 1) and warns, as torch does.

• dtype (Optional[dtype], default: None ) –

  The dtype the input is cast to first, and the output's; None keeps the input's.

• target (Target, default: None ) –

  Which set of kernels serves this op — a target name, BUILTIN for the in-tree kernels, or None to decide from the input device.

forward

forward(
    x,
)

Normalize x along the configured axis.

Parameters:

• x (Tensor) –

  Input tensor of any rank.

Returns:

• Tensor –

  A tensor of x's shape, in dtype when one was passed.

tileops.reduction.LogSoftmaxFwdOp

Log-softmax along dim, following torch.nn.functional.log_softmax.

__init__

__init__(
    dim=None,
    *,
    dtype=None,
    target=None
)

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

Parameters:

• dim (Optional[int], default: None ) –

  The axis to normalize along. None takes torch's implicit axis (0 at rank 0, 1 or 3, else 1) and warns, as torch does.

• dtype (Optional[dtype], default: None ) –

  The dtype the input is cast to first, and the output's; None keeps the input's.

• target (Target, default: None ) –

  Which set of kernels serves this op — a target name, BUILTIN for the in-tree kernels, or None to decide from the input device.

forward

forward(
    x,
)

Normalize x along the configured axis.

Parameters:

• x (Tensor) –

  Input tensor of any rank.

Returns:

• Tensor –

  A tensor of x's shape, in dtype when one was passed.

tileops.reduction.LogSumExpFwdOp

LogSumExp over dim, following torch.logsumexp; an empty reduction is -inf.

__init__

__init__(
    dim,
    keepdim=False,
    *,
    target=None
)

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

Parameters:

• dim (Union[int, List[int], Tuple[int, ...]]) –

  Axes to reduce: an int or a non-empty sequence of them.

• keepdim (bool, default: False ) –

  Whether a reduced axis stays as a length-1 axis.

• target (Target, default: None ) –

  Which set of kernels serves this op — a target name, BUILTIN for the in-tree kernels, or None to decide from the input device.

forward

forward(
    x,
)

Reduce x over the configured axes.

Parameters:

• x (Tensor) –

  Input tensor of any rank.

Returns:

• Tensor –

  The reduction, shaped by dim and keepdim.

Vector norms

One op serves every order torch.linalg.vector_norm is called at here: pass ord as 1, 2 or inf.

tileops.reduction.VectorNormFwdOp

torch.linalg.vector_norm(x, ord, dim, keepdim, dtype=dtype).

ord selects the fold: 1 sums magnitudes, 2 takes the root of the sum of squares, and inf takes the largest magnitude. A row holding a NaN yields NaN at ord=inf, as in torch.

__init__

__init__(
    ord=2,
    dim=None,
    keepdim=False,
    *,
    dtype=None,
    target=None
)

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

Parameters:

• ord (Union[int, float], default: 2 ) –

  Norm order; 1, 2 or inf.

• dim (Union[int, List[int], Tuple[int, ...], None], default: None ) –

  Axes to reduce: an int, a sequence of them, or None for all.

• keepdim (bool, default: False ) –

  Whether a reduced axis stays as a length-1 axis.

• dtype (Optional[dtype], default: None ) –

  The dtype the input is cast to before the reduction, and the output's; it may not narrow the input's. None keeps the input's.

• target (Target, default: None ) –

  Which set of kernels serves this op — a target name, BUILTIN for the in-tree kernels, or None to decide from the input device.

forward

forward(
    x,
)

Reduce x over the configured axes.

Parameters:

• x (Tensor) –

  Input tensor of any rank.

Returns:

• Tensor –

  The reduction, shaped by dim and keepdim.

Cumulative

tileops.reduction.CumsumFwdOp

Cumulative sum operator: y = cumsum(x, dim).

Output has the same shape and dtype as x. Alignment padding is handled inside the kernel via masked loads.

A row one thread block can stage in shared memory takes the whole-row scan; other rows with M < 128 and N > 8192 take a three-pass parallel scan for SM utilization; every other shape takes the tiled scan.

Parameters:

• dim (int, default: -1 ) –

  Reduction axis (default -1). Negative values are normalized at forward time.

• target (Target, default: None ) –

  Which set of kernels serves this op — a target name, BUILTIN for the in-tree kernels, or None to decide from the input device.

Example

op = CumsumFwdOp()
x = torch.randn(1024, 4096, dtype=torch.float16, device="cuda")
y = op(x)  # shape: (1024, 4096)

__init__

__init__(
    dim=-1,
    *,
    target=None
)

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

Parameters:

• dim (int, default: -1 ) –

  Reduction axis (default -1). Negative values are normalized at forward time (dim % x.ndim).

• target (Target, default: None ) –

  Which set of kernels serves this op — a target name, BUILTIN for the in-tree kernels, or None to decide from the input device.

forward

forward(
    x,
)

Run the scan.

One call to the operator this op registers: this is as far as dynamo traces.

tileops.reduction.CumprodFwdOp

Cumulative product operator: y = cumprod(x, dim).

Output has the same shape and dtype as x. Alignment padding is handled inside the kernel via masked loads.

Parameters:

• dim (int, default: -1 ) –

  Reduction axis (default -1). Negative values are normalized at forward time.

• target (Target, default: None ) –

  Which set of kernels serves this op — a target name, BUILTIN for the in-tree kernels, or None to decide from the input device.

Example

op = CumprodFwdOp()
x = torch.randn(1024, 4096, dtype=torch.float16, device="cuda") * 0.01 + 0.99
y = op(x)  # shape: (1024, 4096)

__init__

__init__(
    dim=-1,
    *,
    target=None
)

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

Parameters:

• dim (int, default: -1 ) –

  Reduction axis (default -1). Negative values are normalized at forward time (dim % x.ndim).

• target (Target, default: None ) –

  Which set of kernels serves this op — a target name, BUILTIN for the in-tree kernels, or None to decide from the input device.

forward

forward(
    x,
)

Run the scan.

One call to the operator this op registers: this is as far as dynamo traces.

Logical reductions

tileops.reduction.AllFwdOp

Whether every element along dim is nonzero, following torch.all; returns bool.

An empty dim reduces nothing: the output is x != 0 with the input's shape.

__init__

__init__(
    dim=None,
    keepdim=False,
    *,
    target=None
)

Construct a reduce op.

Parameters:

• dim (Dim, default: None ) –

  Axes to reduce: an int, a sequence of them, or None for all.

• keepdim (bool, default: False ) –

  Whether a reduced axis stays as a length-1 axis.

• target (Target, default: None ) –

  Which set of kernels serves this op — a target name, BUILTIN for the in-tree kernels, or None to decide from the input device.

forward

forward(
    x,
)

Reduce x over the configured axes.

Parameters:

• x (Tensor) –

  Input tensor of any rank.

Returns:

• Tensor –

  The reduction, shaped by dim and keepdim.

tileops.reduction.AnyFwdOp

Whether any element along dim is nonzero, following torch.any; returns bool.

An empty dim reduces nothing: the output is x != 0 with the input's shape.

__init__

__init__(
    dim=None,
    keepdim=False,
    *,
    target=None
)

Construct a reduce op.

Parameters:

• dim (Dim, default: None ) –

  Axes to reduce: an int, a sequence of them, or None for all.

• keepdim (bool, default: False ) –

  Whether a reduced axis stays as a length-1 axis.

• target (Target, default: None ) –

  Which set of kernels serves this op — a target name, BUILTIN for the in-tree kernels, or None to decide from the input device.

forward

forward(
    x,
)

Reduce x over the configured axes.

Parameters:

• x (Tensor) –

  Input tensor of any rank.

Returns:

• Tensor –

  The reduction, shaped by dim and keepdim.

tileops.reduction.CountNonzeroFwdOp

Count of nonzero elements along dim, following torch.count_nonzero; returns int64.

There is no keepdim: a reduced axis always goes, as in torch.

__init__

__init__(
    dim=None,
    *,
    target=None
)

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

Parameters:

• dim (Union[int, List[int], Tuple[int, ...], None], default: None ) –

  Axes to reduce: an int, a sequence of them, or None for all.

• target (Target, default: None ) –

  Which set of kernels serves this op — a target name, BUILTIN for the in-tree kernels, or None to decide from the input device.

forward

forward(
    x,
)

Reduce x over the configured axes.

Parameters:

• x (Tensor) –

  Input tensor of any rank.

Returns:

• Tensor –

  The reduction, shaped by dim and keepdim.