Is this page helpful?

core.fusions.fused_row_copy#

Contiguous copy of a row-strided tensor: contiguous rows at a uniform stride, such as a slice of the last dimension.

Tensor.contiguous copies such a view with PyTorch’s generic strided copy kernel, which does not vectorize its accesses. This kernel copies 2D tiles instead. The values are copied unchanged and the gradient is passed through, as for Tensor.contiguous.

Module Contents#

Classes#

_RowContiguous

Functions#

_row_copy_kernel

_row_layout

Return (rows, columns, row stride) if x is contiguous rows at a uniform stride.

contiguous_rows

Return x.contiguous(), copying row-strided CUDA tensors with a 2D-tiled kernel.

API#

core.fusions.fused_row_copy._row_copy_kernel(
x_ptr,
y_ptr,
n_rows,
n_cols,
row_stride,
BLOCK_R: triton.language.constexpr,
BLOCK_C: triton.language.constexpr,
)#
core.fusions.fused_row_copy._row_layout(
x: torch.Tensor,
) → Optional[Tuple[int, int, int]]#

Return (rows, columns, row stride) if x is contiguous rows at a uniform stride.

class core.fusions.fused_row_copy._RowContiguous#

Bases: torch.autograd.Function

static forward(ctx, x, n_rows, n_cols, row_stride)#

Copy the rows of x into a contiguous tensor.

static backward(ctx, grad)#

Pass the gradient through, as the backward of Tensor.contiguous does.

core.fusions.fused_row_copy.contiguous_rows(x: torch.Tensor) → torch.Tensor#

Return x.contiguous(), copying row-strided CUDA tensors with a 2D-tiled kernel.

Parameters:

x (torch.Tensor) – Input tensor.

Returns:

x itself if it is contiguous, otherwise a contiguous copy.

Return type:

torch.Tensor