mirror of
https://github.com/pytorch/pytorch.git
synced 2025-10-21 13:44:15 +08:00
14b2273b0c
This PR adds in fast semi-structured sparsification kernels to PyTorch. These kernels allow for accelerated semi-structured sparsification kernels in PyTorch. The kernels have been added as aten native functions In particular, three new functions have been added: * `torch._sparse_semi_structured_tile` This function will return the packed representation and metadata for both X and X', as well as the thread masks. Note that this applies 2:4 sparsity in a 4x4 tile instead of a 1x4 strip as usual. * `torch._sparse_semi_structured_apply` This function takes in an input tensor and thread masks from the above function and returns a packed representation and metadata from applying thread masks to the input tensor. * `torch._sparse_semi_structured_apply_dense` This function does the same thing as above but instead of returning the tensor in the sparse representation it returns it in the dense representation The subclasses have also been updated to add a new `prune_dense_static_sort` classmethod to create sparse tensors with this format. I've added some additional documentatino on how to calculate the compressed tensors needed to create a SparseSemiStructuredTensor oneself. To this end, there are two new helper functions added: `sparse_semi_structured_tile` `compute_compressed_swizzled_bitmask` Differential Revision: [D56190801](https://our.internmc.facebook.com/intern/diff/D56190801) Pull Request resolved: https://github.com/pytorch/pytorch/pull/122350 Approved by: https://github.com/cpuhrsch
352 lines
14 KiB
Python
352 lines
14 KiB
Python
import torch
|
|
|
|
|
|
def _calculate_meta_reordering_scatter_offsets(m, meta_ncols, meta_dtype, device):
|
|
"""
|
|
This is PyTorch implementation of main part of reorder_meta()
|
|
function, from tools/util/include/cutlass/util/host_reorder.h file
|
|
of CUTLASS source tree. Furthermore, CUTLASS template for sparse
|
|
GEMM decides upon layout of this matrix, and at the moment for the
|
|
sparse GEMM executed on tensor cores, this is layout described by
|
|
ColumnMajorInterleaved<2> data structure, in
|
|
include/cutlass/layout/matrix.h of CUTLASS source tree. The
|
|
reordering of meta matrix into meta_reordered matrix calculated
|
|
according to these segments of CUTLASS code is re-implemented here.
|
|
Note that this calculation produces offsets for scattering metadata
|
|
matrix elements into reordered metadata matrix elements (or,
|
|
equivalently, for gathering reordered metadata matrix element back
|
|
into metadata matrix elements).
|
|
"""
|
|
dst_rows = torch.arange(0, m, device=device)[:, None].repeat(1, meta_ncols)
|
|
dst_cols = torch.arange(0, meta_ncols, device=device).repeat(m, 1)
|
|
|
|
# Reorder the rows, then swizzle the 2x2 blocks.
|
|
group = 32 if meta_dtype.itemsize == 2 else 16
|
|
interweave = 4 if meta_dtype.itemsize == 2 else 2
|
|
dst_rows = (
|
|
dst_rows // group * group
|
|
+ (dst_rows % 8) * interweave
|
|
+ (dst_rows % group) // 8
|
|
)
|
|
|
|
topright = ((dst_rows % 2 == 0) & (dst_cols % 2 == 1)).to(torch.int8)
|
|
bottomleft = ((dst_rows % 2 == 1) & (dst_cols % 2 == 0)).to(torch.int8)
|
|
dst_rows += topright - bottomleft
|
|
dst_cols -= topright - bottomleft
|
|
|
|
# Assumed that meta tensor is to be stored in CUTLASS
|
|
# InterleavedColumnMajor layout, and reverse engineered
|
|
# corresponding code to store values into this tensor.
|
|
interleave = 2
|
|
cols_maj = dst_cols // interleave
|
|
cols_min = dst_cols % interleave
|
|
return (cols_maj * m * interleave + dst_rows * interleave + cols_min).view(-1)
|
|
|
|
|
|
def sparse_semi_structured_from_dense_cutlass(dense):
|
|
"""
|
|
This function converts dense matrix into sparse semi-structured
|
|
representation, producing "compressed" matrix, in the layout used by
|
|
CUTLASS backend, and corresponding metadata matrix.
|
|
"""
|
|
if dense.dim() != 2:
|
|
raise RuntimeError(
|
|
f"Expected 2-dimensional dense tensor, got {dense.dim()}-dimensional tensor"
|
|
)
|
|
|
|
m, k = dense.shape
|
|
device = dense.device
|
|
|
|
meta_dtype = torch.int8
|
|
if dense.dtype == torch.int8:
|
|
meta_dtype = torch.int32
|
|
elif dense.dtype in [torch.half, torch.bfloat16, torch.float]:
|
|
meta_dtype = torch.int16
|
|
else:
|
|
raise RuntimeError(f"Invalid datatype {dense.dtype} of dense matrix")
|
|
quadbits_per_meta_elem = meta_dtype.itemsize * 8 // 4
|
|
if quadbits_per_meta_elem not in (4, 8):
|
|
raise RuntimeError("Invalid number of elements per meta element calculated")
|
|
|
|
if meta_dtype == torch.int32:
|
|
if m % 16 != 0:
|
|
raise RuntimeError(
|
|
f"Number of rows of dense matrix {m} must be divisible by 16"
|
|
)
|
|
else:
|
|
if m % 32 != 0:
|
|
raise RuntimeError(
|
|
f"Number of rows of dense matrix {m} must be divisible by 32"
|
|
)
|
|
if k % (4 * quadbits_per_meta_elem) != 0:
|
|
raise RuntimeError(
|
|
f"Number of columns of dense matrix {k} must be divisible by {4 * quadbits_per_meta_elem}"
|
|
)
|
|
|
|
if dense.dtype != torch.float:
|
|
ksparse = 4
|
|
dense_4 = dense.view(-1, k // ksparse, ksparse)
|
|
m0, m1, m2, m3 = (dense_4 != 0).unbind(-1)
|
|
else:
|
|
ksparse = 2
|
|
dense_2 = dense.view(-1, k // ksparse, ksparse)
|
|
m0, m2 = m1, m3 = (dense_2 != 0).unbind(-1)
|
|
meta_ncols = k // (ksparse * quadbits_per_meta_elem)
|
|
|
|
# Encoding quadruples of True/False values as follows:
|
|
# [True, True, False, False] -> 0b0100
|
|
# [True, False, True, False] -> 0b1000
|
|
# [False, True, True, False] -> 0b1001
|
|
# [True, False, False, True ] -> 0b1100
|
|
# [False, True, False, True ] -> 0b1101
|
|
# [False, False, True, True ] -> 0b1110
|
|
# Thus, lower two bits in the encoding are index of the True value
|
|
# at the lowest index in the quadruple, and the higher two bits in
|
|
# the encoding are index of the other True value in the quadruple.
|
|
# In case there are less than two True values, than False value or
|
|
# values at some index or indices are considered True for the
|
|
# encoding. In case there are more than two True values, then the
|
|
# excess True value(s) at some indices are considered False for
|
|
# the encoding. The exact encodings used for these cases are as
|
|
# follows:
|
|
# [False, False, False, False] -> 0b1110
|
|
# [False, False, False, True ] -> 0b1110
|
|
# [False, False, True, False] -> 0b1110
|
|
# [False, True, False, False] -> 0b1001
|
|
# [False, True, True, True ] -> 0b1101
|
|
# [True, False, False, False] -> 0b1000
|
|
# [True, False, True, True ] -> 0b1100
|
|
# [True, True, False, True ] -> 0b0100
|
|
# [True, True, True, False] -> 0b0100
|
|
# [True, True, True, True ] -> 0b0100
|
|
# These particular encodings are chosen, with the help of Espresso
|
|
# logic minimizer software, for the purpose of minimization of
|
|
# corresponding Boolean functions, that translate non-zero flags
|
|
# into encoding bits. Note also possible choices for the first
|
|
# and last of these encodings were limited only to (0b0100,
|
|
# 0b1110), in order to produce valid encodings for 1:2 sparsity
|
|
# case.
|
|
|
|
expr0 = m0 & m1
|
|
expr1 = ~m0 & m1
|
|
expr2 = ~m0 & ~m1
|
|
bit0 = expr1
|
|
bit1 = expr2
|
|
bit2 = expr0 | expr2 | m3
|
|
bit3 = expr1 | ~m1
|
|
idxs0 = bit0 | (bit1.to(torch.int64) << 1)
|
|
idxs1 = bit2 | (bit3.to(torch.int64) << 1)
|
|
|
|
if dense.dtype != torch.float:
|
|
sparse0 = dense_4.gather(-1, idxs0.unsqueeze(-1)) # type: ignore[possibly-undefined]
|
|
sparse1 = dense_4.gather(-1, idxs1.unsqueeze(-1))
|
|
sparse = torch.stack((sparse0, sparse1), dim=-1).view(m, k // 2)
|
|
else:
|
|
sparse = dense_2.gather(-1, idxs0.unsqueeze(-1) // 2).view(m, k // 2) # type: ignore[possibly-undefined]
|
|
|
|
meta_4 = idxs0 | (idxs1 << 2)
|
|
meta_n = meta_4.view((-1, meta_ncols, quadbits_per_meta_elem)).to(meta_dtype)
|
|
|
|
if quadbits_per_meta_elem == 4:
|
|
meta = (
|
|
meta_n[:, :, 0]
|
|
| (meta_n[:, :, 1] << 4)
|
|
| (meta_n[:, :, 2] << 8)
|
|
| (meta_n[:, :, 3] << 12)
|
|
)
|
|
elif quadbits_per_meta_elem == 8:
|
|
meta = (
|
|
meta_n[:, :, 0]
|
|
| (meta_n[:, :, 1] << 4)
|
|
| (meta_n[:, :, 2] << 8)
|
|
| (meta_n[:, :, 3] << 12)
|
|
| (meta_n[:, :, 4] << 16)
|
|
| (meta_n[:, :, 5] << 20)
|
|
| (meta_n[:, :, 6] << 24)
|
|
| (meta_n[:, :, 7] << 28)
|
|
)
|
|
|
|
# Reorder meta tensor elements.
|
|
meta_reordered = meta.new_empty((m * meta_ncols,)) # type: ignore[possibly-undefined]
|
|
meta_offsets = _calculate_meta_reordering_scatter_offsets(
|
|
m, meta_ncols, meta_dtype, device
|
|
)
|
|
meta_reordered.scatter_(0, meta_offsets, meta.view(-1))
|
|
|
|
return (sparse, meta_reordered.view(m, meta_ncols))
|
|
|
|
|
|
def sparse_semi_structured_to_dense_cutlass(sparse, meta_reordered):
|
|
"""
|
|
This function performs reverse of the function above - it
|
|
reconstructs dense matrix from a pair of "compressed" matrix, given
|
|
in the layout used by CUTLASS backend, and accompanying metadata
|
|
matrix.
|
|
"""
|
|
if sparse.dim() != 2:
|
|
raise RuntimeError(
|
|
f"Expected 2-dimensional sparse tensor, got {sparse.dim()}-dimensional tensor"
|
|
)
|
|
|
|
m, k = sparse.shape
|
|
device = sparse.device
|
|
|
|
if meta_reordered.dim() != 2:
|
|
raise RuntimeError(
|
|
f"Expected 2-dimensional meta tensor, got {meta_reordered.dim()}-dimensional tensor"
|
|
)
|
|
if meta_reordered.device != device:
|
|
raise RuntimeError(
|
|
f"Expected meta matrix to be on {device} device, got matrix on {meta_reordered.device} device"
|
|
)
|
|
|
|
meta_dtype = meta_reordered.dtype
|
|
if meta_dtype not in (torch.int16, torch.int32):
|
|
raise RuntimeError(f"Invalid datatype {meta_dtype} of meta matrix")
|
|
quadbits_per_meta_elem = meta_dtype.itemsize * 8 // 4
|
|
|
|
if sparse.dtype != torch.float:
|
|
ksparse = 4
|
|
else:
|
|
ksparse = 2
|
|
|
|
meta_nrows, meta_ncols = meta_reordered.shape
|
|
if meta_nrows != m:
|
|
raise RuntimeError(
|
|
f"Number of rows of meta matrix {meta_nrows} must be equal to number of columns of spase matrix {m}"
|
|
)
|
|
if meta_ncols * ksparse * quadbits_per_meta_elem != 2 * k:
|
|
raise RuntimeError(
|
|
f"Number of columns of sparse matrix {k} different from the {meta_ncols * ksparse * quadbits_per_meta_elem // 2}, "
|
|
"expected according to the number of columns of meta matrix"
|
|
)
|
|
|
|
# Undo meta tensor elements reordering.
|
|
meta_offsets = _calculate_meta_reordering_scatter_offsets(
|
|
m, meta_ncols, meta_dtype, device
|
|
)
|
|
meta = torch.gather(meta_reordered.view(-1), 0, meta_offsets).view(m, meta_ncols)
|
|
|
|
# Unpack sparse tensor back to original dense tensor, using
|
|
# information provided by meta tensor. Note that torch.float
|
|
# datatype is handled pretty much the same as
|
|
# torch.half/torch.bfloat16, as metadata for a pair of torch.float
|
|
# value is encoded as if underlying 8 bytes contain four
|
|
# torch.half/torch.bfloat16 values, where either first two or last
|
|
# two are zeros.
|
|
meta_2 = torch.empty(
|
|
(m, meta_ncols, 2 * quadbits_per_meta_elem),
|
|
dtype=meta_dtype,
|
|
device=device,
|
|
)
|
|
if quadbits_per_meta_elem == 4:
|
|
meta_2[:, :, 0] = meta & 0b11
|
|
meta_2[:, :, 1] = (meta >> 2) & 0b11
|
|
meta_2[:, :, 2] = (meta >> 4) & 0b11
|
|
meta_2[:, :, 3] = (meta >> 6) & 0b11
|
|
meta_2[:, :, 4] = (meta >> 8) & 0b11
|
|
meta_2[:, :, 5] = (meta >> 10) & 0b11
|
|
meta_2[:, :, 6] = (meta >> 12) & 0b11
|
|
meta_2[:, :, 7] = (meta >> 14) & 0b11
|
|
elif quadbits_per_meta_elem == 8:
|
|
meta_2[:, :, 0] = meta & 0b11
|
|
meta_2[:, :, 1] = (meta >> 2) & 0b11
|
|
meta_2[:, :, 2] = (meta >> 4) & 0b11
|
|
meta_2[:, :, 3] = (meta >> 6) & 0b11
|
|
meta_2[:, :, 4] = (meta >> 8) & 0b11
|
|
meta_2[:, :, 5] = (meta >> 10) & 0b11
|
|
meta_2[:, :, 6] = (meta >> 12) & 0b11
|
|
meta_2[:, :, 7] = (meta >> 14) & 0b11
|
|
meta_2[:, :, 8] = (meta >> 16) & 0b11
|
|
meta_2[:, :, 9] = (meta >> 18) & 0b11
|
|
meta_2[:, :, 10] = (meta >> 20) & 0b11
|
|
meta_2[:, :, 11] = (meta >> 22) & 0b11
|
|
meta_2[:, :, 12] = (meta >> 24) & 0b11
|
|
meta_2[:, :, 13] = (meta >> 26) & 0b11
|
|
meta_2[:, :, 14] = (meta >> 28) & 0b11
|
|
meta_2[:, :, 15] = (meta >> 30) & 0b11
|
|
|
|
dense_offsets = meta_2.view(-1) + (
|
|
torch.arange(0, 2 * m * k // ksparse, device=device) * 4
|
|
).view(-1, 1).repeat(1, 2).view(-1)
|
|
|
|
dense = torch.zeros((m * 2 * k,), dtype=sparse.dtype, device=device)
|
|
if sparse.dtype != torch.float:
|
|
dense.scatter_(0, dense_offsets, sparse.view(-1))
|
|
else:
|
|
dense.view(torch.half).scatter_(
|
|
0, dense_offsets, sparse.view(torch.half).view(-1)
|
|
)
|
|
|
|
return dense.view(m, 2 * k)
|
|
|
|
|
|
def _sparse_semi_structured_tile(dense):
|
|
"""
|
|
This function computes a 2:4 sparse tile by greedily taking the largest values.
|
|
|
|
Since we take the largest values greedily, how the sorting algorithm handles duplicates affects
|
|
the ultimate sparsity pattern.
|
|
|
|
Note that this function does not have the same sorting semantics as our CUDA backend,
|
|
which is exposed via `torch._sparse_semi_structured_tile` and thus returns a different pattern.
|
|
"""
|
|
|
|
def greedy_prune_tile(tile):
|
|
num_kept_row = [0, 0, 0, 0]
|
|
num_kept_col = [0, 0, 0, 0]
|
|
|
|
for x in tile.flatten().sort(descending=True, stable=True).indices:
|
|
r, c = x // 4, x % 4
|
|
if num_kept_row[r] < 2 and num_kept_col[c] < 2:
|
|
num_kept_row[r] += 1
|
|
num_kept_col[c] += 1
|
|
else:
|
|
tile[r, c] = 0
|
|
|
|
for batch in dense.unfold(0, 4, 4).unfold(1, 4, 4):
|
|
for tile in batch:
|
|
greedy_prune_tile(tile)
|
|
|
|
return dense
|
|
|
|
|
|
def _compute_compressed_swizzled_bitmask(dense):
|
|
"""
|
|
Calculates the compressed swizzled bitmask from a dense tensor
|
|
"""
|
|
|
|
# first we need to convert the dense tensor to a bitmask
|
|
int_bitmask = dense.bool().to(torch.uint8)
|
|
|
|
# Each thread is responsible for an 8x8 tile, which contains 4 4x4 tiles:
|
|
# A, B, C and D, as displayed in the following schema:
|
|
# +---+---+
|
|
# | A | B |
|
|
# +---+---+
|
|
# | C | D |
|
|
# +---+---+
|
|
|
|
# we first need to split into the 8x8 tiles
|
|
bitmask_8x8_chunks = int_bitmask.unfold(0, 8, 8).unfold(1, 8, 8)
|
|
|
|
# then we unfold again to get our indivdual 4x4 tiles
|
|
bitmask_4x4_chunks = bitmask_8x8_chunks.unfold(2, 4, 4).unfold(3, 4, 4)
|
|
|
|
# Each 4x4 bitmask defines two 8-bit integers, which encode the sparsity pattern
|
|
# of that tile. Note that the least siginificant bit is stored first.
|
|
# [1 1 0 0]
|
|
# [1 1 0 0] -> 0011 0011 -> 51
|
|
# [0 0 1 1] 1100 1100 204
|
|
# [0 0 1 1]
|
|
|
|
# reshape tensor to expand tiles into 8-bit vectors
|
|
bitmask_binary_representation = bitmask_4x4_chunks.reshape(*bitmask_4x4_chunks.shape[:2], 4, 2, 8)
|
|
|
|
# to convert from binary representaiton, we can do a matmul with powers of two
|
|
powers_of_two = 2**torch.arange(8, dtype=torch.float, device="cuda")
|
|
# To run on GPU: cast to float to do matmul and then cast back
|
|
compressed_swizzled_bitmask = (bitmask_binary_representation.to(torch.float) @ powers_of_two).to(torch.uint8)
|
|
|
|
return compressed_swizzled_bitmask
|