saad
/
ai-content-maker
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity
ai-content-maker/.venv/Lib/site-packages/torch/sparse/_triton_ops.py

1929 lines
73 KiB
Python
Raw Blame History

import math
import os
import torch
import weakref
from functools import lru_cache
from torch.utils._triton import has_triton
from ._triton_ops_meta import get_meta
from typing import Optional, Tuple
TORCH_SPARSE_BSR_SCATTER_MM_LRU_CACHE_SIZE = int(os.getenv('TORCH_SPARSE_BSR_SCATTER_MM_LRU_CACHE_SIZE', 2))
def check(cond, msg):
if not cond:
raise ValueError(msg)
def check_bsr_layout(f_name, t):
check(
t.layout == torch.sparse_bsr,
f"{f_name}(): only BSR sparse format is supported for the sparse argument.",
)
def check_device(f_name, t, device):
check(
t.device == device and t.device.type == "cuda",
f"{f_name}(): all inputs are expected to be on the same GPU device.",
)
def check_mm_compatible_shapes(f_name, lhs, rhs):
check(
lhs.dim() >= 2 and rhs.dim() >= 2,
f"{f_name}(): all inputs involved in the matrix product are expected to be at least 2D, "
f"but got lhs.dim() == {lhs.dim()} and rhs.dim() == {rhs.dim()}."
)
m, kl = lhs.shape[-2:]
kr, n = rhs.shape[-2:]
check(
kl == kr,
f"{f_name}(): arguments' sizes involved in the matrix product are not compatible for matrix multiplication, "
f"got lhs.shape[-1] == {kl} which is not equal to rhs.shape[-2] == {kr}.",
)
def check_dtype(f_name, t, dtype, *additional_dtypes):
check(
t.dtype == dtype
and t.dtype in ((torch.half, torch.bfloat16, torch.float) + tuple(*additional_dtypes)),
f"{f_name}(): all inputs are expected to be of the same dtype "
f"and one of (half, bfloat16, float32) or {additional_dtypes}, "
f"but got dtype == {t.dtype}.",
)
def check_blocksize(f_name, blocksize):
assert len(blocksize) == 2
def is_power_of_two(v):
return not (v & (v - 1))
def is_compatible_blocksize(b):
res = True
for blocksize in b:
# Triton loads only blocks which are at least 16 and powers of 2.
res = (blocksize >= 16 and is_power_of_two(blocksize)) and res
return res
check(
is_compatible_blocksize(blocksize),
f"{f_name}(): sparse inputs' blocksize ({blocksize[0]}, {blocksize[1]}) "
"should be at least 16 and a power of 2 in each dimension.",
)
def make_triton_contiguous(t):
"""Return input as a triton-contiguous tensor.
A triton-contiguous tensor is defined as a tensor that has strides
with minimal value equal to 1.
While triton kernels support triton-non-contiguous tensors (all
strides being greater than 1 or having 0 strides) arguments, a
considerable slow-down occurs because tensor data is copied
element-wise rather than chunk-wise.
"""
if min(t.stride()) != 1:
# TODO: investigate if contiguity along other axes than the
# last one can be beneficial for performance
return t.contiguous()
else:
return t
def broadcast_batch_dims(f_name, *tensors):
try:
return torch.broadcast_shapes(*(t.shape[:-2] for t in tensors))
except Exception:
check(False, f"{f_name}(): inputs' batch dimensions are not broadcastable!")
def slicer(dim, slice_range, *tensors):
for t in tensors:
slices = [slice(None)] * t.dim()
slices[dim] = slice_range
yield t[slices]
def multidim_slicer(dims, slices, *tensors):
for t in tensors:
s = [slice(None)] * t.dim()
for d, d_slice in zip(dims, slices):
if d is not None:
s[d] = d_slice
yield t[s]
def ptr_stride_extractor(*tensors):
for t in tensors:
yield t
yield from t.stride()
def grid_partitioner(full_grid, grid_blocks, tensor_dims_map):
assert 0 <= len(full_grid) <= 3
assert 0 <= len(grid_blocks) <= 3
import itertools
def generate_grid_points():
for fg, mg in zip(full_grid, grid_blocks):
yield range(0, fg, mg)
def generate_sliced_tensors(slices):
for t, t_dims in tensor_dims_map.items():
yield next(multidim_slicer(t_dims, slices, t))
for grid_point in itertools.product(*generate_grid_points()):
grid = [min(fg - gp, mg) for fg, gp, mg in zip(full_grid, grid_point, grid_blocks)]
slices = [slice(gp, gp + g) for gp, g in zip(grid_point, grid)]
# grid_points are iterated in a "contiguous" order, i.e.
# left dimensions traversed slower than right dimensions.
# This order is reversed for CUDA grids.
yield grid[::-1], *generate_sliced_tensors(slices)
def launch_kernel(kernel, tensor_dims_map, full_grid, grid_blocks=None):
# cuda_max_grid = (2 ** 31 - 1, 2 ** 16 - 1, 2 ** 16 - 1)
cuda_max_grid = (2147483647, 65535, 65535)[::-1]
if grid_blocks is None:
grid_blocks = cuda_max_grid
else:
def valid_grid_dim(g, mg):
if g is None:
return mg
else:
# grid must be at least 1 and no greater than mg
return max(1, min(g, mg))
grid_blocks = tuple(
valid_grid_dim(g, mg) for g, mg in zip(grid_blocks, cuda_max_grid)
) # type: ignore[assignment]
for grid, *sliced_tensors in grid_partitioner(full_grid, grid_blocks, tensor_dims_map):
kernel(grid, *sliced_tensors)
def prepare_inputs(bsr, *dense_tensors):
# Introduce fake batch dimension if not present for convenience.
crow_indices = bsr.crow_indices().unsqueeze(0)
col_indices = bsr.col_indices().unsqueeze(0)
values = make_triton_contiguous(bsr.values().unsqueeze(0))
tensors = [make_triton_contiguous(t.unsqueeze(0)) for t in dense_tensors]
# Compute broadcasted batch dimension
batch_dims_broadcasted = torch.broadcast_shapes(values.shape[:-3], *(t.shape[:-2] for t in tensors))
# Broadcast batch dimensions and squash.
# The result can be either a view or a copy.
def batch_broadcast_and_squash(t, batch_dims, invariant_dims):
return t.broadcast_to(batch_dims + invariant_dims).flatten(
0, len(batch_dims) - 1
)
crow_indices = batch_broadcast_and_squash(
crow_indices, batch_dims_broadcasted, (-1,)
)
col_indices = batch_broadcast_and_squash(
col_indices, batch_dims_broadcasted, (-1,)
)
values = batch_broadcast_and_squash(
values, batch_dims_broadcasted, values.shape[-3:]
)
tensors = [
batch_broadcast_and_squash(t, batch_dims_broadcasted, t.shape[-2:]) for t in tensors
]
return crow_indices, col_indices, values, *tensors
def broadcast_batch_dims_bsr(f_name, bsr, *tensors):
batch_shape = broadcast_batch_dims(f_name, bsr, *tensors)
crow_indices = bsr.crow_indices().broadcast_to(batch_shape + (-1,))
col_indices = bsr.col_indices().broadcast_to(batch_shape + (-1,))
values = bsr.values().broadcast_to(batch_shape + bsr.values().shape[-3:])
size = batch_shape + bsr.shape[-2:]
return torch.sparse_compressed_tensor(crow_indices, col_indices, values, size=size, layout=bsr.layout)
# NOTE: this function will ALWAYS create a view
def tile_to_blocksize(t, blocksize):
*rest, m, n = t.shape
new_shape = rest + [
m // blocksize[0],
blocksize[0],
n // blocksize[1],
blocksize[1],
]
# using .view instead of .reshape to ensure that the result is
# indeed a view:
return t.view(new_shape).transpose(-3, -2)
def as1Dbatch(tensor):
"""Return tensor as 3D tensor by either prepending new dimensions to
the tensor shape (when ``tensor.ndim < 3``), or by collapsing
starting dimensions into the first dimension (when ``tensor.ndim >
3``).
"""
while tensor.ndim < 3:
tensor = tensor.unsqueeze(0)
if tensor.ndim > 3:
tensor = tensor.flatten(0, tensor.ndim - 3)
assert tensor.ndim == 3, tensor.shape
return tensor
def scatter_mm(blocks, others, indices_data, *, accumulators=None):
"""Scattered matrix multiplication of tensors.
A scattered matrix multiplication is defined as a series of matrix
multiplications applied to input tensors according to the input
and output mappings specified by indices data.
The following indices data formats are supported for defining a
scattered matrix multiplication operation (:attr:`indices_data[0]`
holds the name of the indices data format as specified below):
- ``"scatter_mm"`` - matrix multiplications scattered in batches
of tensors.
If :attr:`blocks` is a :math:`(* \times M \times K) tensor,
:attr:`others` is a :math:`(* \times K \times N)` tensor,
:attr:`accumulators` is a :math:`(* \times M \times N)` tensor,
and :attr:`indices = indices_data['indices']` is a :math:`(*
\times 3)` tensor, then the operation is equivalent to the
following code::
c_offsets, pq = indices_data[1:]
for r in range(len(c_offsets) - 1):
for g in range(c_offsets[r], c_offsets[r + 1]):
p, q = pq[g]
accumulators[r] += blocks[p] @ others[q]
- ``"bsr_strided_mm"`` - matrix multiplications scattered in
batches of tensors and a tensor.
If :attr:`blocks` is a :math:`(Ms \times Ks) tensor,
:attr:`others` is a :math:`(* \times K \times N)` tensor,
:attr:`accumulators` is a :math:`(* \times M \times N)` tensor, then
the operation is equivalent to the following code::
c_indices, r_offsets, p_offsets, q_offsets, meta = indices_data[1:]
for b in range(nbatches):
for i, r in enumerate(r_offsets):
r0, r1 = divmod(r, N)
acc = accumulators[b, r0:r0 + Ms, r1:r1 + Ns]
for g in range(c_indices[i], c_indices[i+1]):
p = p_offsets[g]
q0, q1 = divmod(q_offsets[g], N)
acc += blocks[p] @ others[b, q0:q0 + Ks, q1:q1 + Ns]
where ``Ns = N // meta['SPLIT_N']``, and ``M`` and ``K`` are
integer multiples of ``Ms`` and ``Ks``, respectively.
- ``"bsr_strided_mm_compressed"`` - matrix multiplications
scattered in batches of tensors and a tensor. A memory and
processor efficient version of ``"bsr_strided_mm"`` format. If
:attr:`blocks` is a :math:`(Ms \times Ks) tensor, :attr:`others`
is a :math:`(* \times K \times N)` tensor, :attr:`accumulators`
is a :math:`(* \times M \times N)` tensor, then the operation is
equivalent to the following code::
c_indices, r_offsets, q_offsets, meta = indices_data[1:]
for b in range(nbatches):
for r in r_offsets:
m = (r // N) // Ms
n = (r % N) // Ns
r0, r1 = divmod(r, N)
c0, c1 = c_indices[m], c_indices[m + 1]
acc = accumulators[b, r0:r0 + Ms, r1:r1 + Ns]
for i, p in enumerate(range(c0, c1)):
q = q_offsets[n * c1 + (SPLIT_N - n) * c0 + i]
q0, q1 = divmod(q, N)
acc += blocks[p] @ others[b, q0:q0 + Ks, q1:q1 + Ns]
where ``Ns = N // meta['SPLIT_N']``, and ``M`` and ``K`` are
integer multiples of ``Ms`` and ``Ks``, respectively.
Notice that the order of ``r_offsets`` items can be arbitrary;
this property enables defining swizzle operators via
rearrangements of ``r_offsets`` items..
Auxilary functions are provided for pre-computing
:attr:`indices_data`. For example,
:func:`bsr_scatter_mm_indices_data` is used to define indices data
for matrix multiplication of BSR and strided tensors.
Parameters
----------
blocks (Tensor): a 3-D tensor of first matrices to be multiplied
others (Tensor): a tensor of second matrices to be multiplied. If
``indices_data[0]=="scatter_mm"``, the tensor is a 1-D batch
tensor of second input matrices to be multiplied. Otherwise, the
second input matrices are slices of the :attr:`others` tensor.
indices_data (tuple): a format data that defines the inputs and
outputs of scattered matrix multiplications.
Keyword arguments
-----------------
accumulators (Tensor, optional): a tensor of matrix product
accumulators. If ``indices_data[0]=="scatter_mm"``, the tensor
is a 1-D batch tensor of output matrices. Otherwise, output
matrices are slices of the :attr:`accumulators` tensor.
"""
indices_format = indices_data[0]
assert blocks.ndim == 3
P, Ms, Ks = blocks.shape
if indices_format == 'scatter_mm':
c_offsets, pq = indices_data[1:]
assert others.ndim == 3
Q, Ks_, Ns = others.shape
assert Ks == Ks_
if accumulators is None:
R = c_offsets.shape[0] - 1
accumulators = torch.zeros((R, Ms, Ns), dtype=blocks.dtype, device=blocks.device)
else:
R, Ms_, Ns_ = accumulators.shape
assert Ms_ == Ms
assert Ns_ == Ns
if Ms % 16 or Ks % 16 or Ns % 16 or _scatter_mm2 is None:
for r in range(c_offsets.shape[0] - 1):
g0 = c_offsets[r]
g1 = c_offsets[r + 1]
for g in range(g0, g1):
p, q = pq[g]
accumulators[r] += blocks[p] @ others[q]
else:
_scatter_mm2(blocks, others, c_offsets, pq, accumulators)
return accumulators
elif indices_format == 'bsr_strided_mm':
others_shape = others.shape
others = as1Dbatch(others)
B, K, N = others.shape
assert K % Ks == 0
c_indices, r_offsets, p_offsets, q_offsets, meta = indices_data[1:]
SPLIT_N = meta['SPLIT_N']
if accumulators is None:
M = Ms + (r_offsets.max().item() + 1) // N
accumulators = torch.zeros((*others_shape[:-2], M, N), dtype=blocks.dtype, device=blocks.device)
else:
M, N_ = accumulators.shape[-2:]
assert N_ == N
accumulators_shape = accumulators.shape
accumulators = as1Dbatch(accumulators)
Ns = N // SPLIT_N
if Ms % 16 or Ks % 16 or Ns % 16 or _scatter_mm6 is None:
accumulators.zero_()
for b in range(B):
for r in range(r_offsets.shape[0]):
r_ = r_offsets[r].item()
g0 = c_indices[r].item()
g1 = c_indices[r + 1].item()
r0, r1 = divmod(r_, N)
acc = accumulators[b, r0:r0 + Ms, r1:r1 + Ns]
for g in range(g0, g1):
p, q = p_offsets[g], q_offsets[g]
q0, q1 = divmod(q.item(), N)
acc += blocks[p] @ others[b, q0:q0 + Ks, q1:q1 + Ns]
else:
_scatter_mm6(blocks, others, c_indices, r_offsets, p_offsets, q_offsets, meta, accumulators)
return accumulators.view(accumulators_shape)
elif indices_format == 'bsr_strided_mm_compressed':
others_shape = others.shape
others = as1Dbatch(others)
B, K, N = others.shape
assert K % Ks == 0
c_indices, r_offsets, q_offsets, meta = indices_data[1:]
SPLIT_N = meta['SPLIT_N']
if accumulators is None:
M = Ms + (r_offsets.max().item() + 1) // N
accumulators = torch.zeros((*others_shape[:-2], M, N), dtype=blocks.dtype, device=blocks.device)
else:
M, N_ = accumulators.shape[-2:]
assert N_ == N
accumulators_shape = accumulators.shape
accumulators = as1Dbatch(accumulators)
Ns = N // SPLIT_N
if Ms % 16 or Ks % 16 or Ns % 16 or _scatter_mm6 is None:
for b in range(B):
for j in range(len(r_offsets)):
r0, r1 = divmod(r_offsets[j].item(), N)
m = r0 // Ms
n = r1 // Ns
c0 = c_indices[m].item()
c1 = c_indices[m + 1].item()
acc = accumulators[b, r0:r0 + Ms, r1:r1 + Ns]
for i, p in enumerate(range(c0, c1)):
q = q_offsets[n * c1 + (SPLIT_N - n) * c0 + i].item()
q0, q1 = divmod(q, N)
acc += blocks[p] @ others[b, q0:q0 + Ks, q1:q1 + Ns]
else:
p_offsets = torch.empty((0, ), dtype=q_offsets.dtype, device=q_offsets.device)
_scatter_mm6(blocks, others, c_indices, r_offsets, p_offsets, q_offsets, meta, accumulators)
return accumulators.view(accumulators_shape)
else:
raise NotImplementedError(indices_format)
def scatter_mm_meta(M, K, N, Ms, Ks,
GROUP_SIZE=None, TILE_M=None, TILE_N=None, SPLIT_N=None, num_warps=None, num_stages=None, **extra):
if {TILE_M, TILE_N, SPLIT_N, num_warps, num_stages, GROUP_SIZE} == {None}:
device_name = torch.cuda.get_device_name()
meta = get_meta('scatter_mm', (M, K, N, Ms, Ks), device_name,
version=(0, torch.float16, 0.5))
if meta is not None:
meta.update(**extra)
return meta
# The following parameters are optimized for the performance
# equilibrium points of bsr-dense and dense-dense matrix
# multiplications when using GPU card NVIDIA GeForce RTX 2060
# SUPER. For points far from the performance equilibrium
# points as well as for other GPU cards, the optimal
# parameters are likely different from what specified below.
if (M, K, N) == (256,) * 3:
if (Ms, Ks) == (16, 16):
SPLIT_N=1;TILE_M=16;TILE_N=16;GROUP_SIZE=4;num_stages=1;num_warps=4 # noqa: E225,E231,E702
elif (Ms, Ks) == (32, 32):
SPLIT_N=2;TILE_M=32;TILE_N=16;GROUP_SIZE=4;num_stages=1;num_warps=4 # noqa: E225,E231,E702
elif (Ms, Ks) == (64, 64):
SPLIT_N=1;TILE_M=32;TILE_N=32;GROUP_SIZE=4;num_stages=1;num_warps=4 # noqa: E225,E231,E702
elif (Ms, Ks) == (128, 128):
SPLIT_N=1;TILE_M=32;TILE_N=32;GROUP_SIZE=2;num_stages=1;num_warps=4 # noqa: E225,E231,E702
elif (M, K, N) == (512,) * 3:
if (Ms, Ks) == (16, 16):
SPLIT_N=8;TILE_M=16;TILE_N=64;GROUP_SIZE=2;num_stages=1;num_warps=2 # noqa: E225,E231,E702
elif (Ms, Ks) == (32, 32):
SPLIT_N=8;TILE_M=32;TILE_N=64;GROUP_SIZE=4;num_stages=1;num_warps=2 # noqa: E225,E231,E702
elif (Ms, Ks) == (64, 64):
SPLIT_N=4;TILE_M=32;TILE_N=128;GROUP_SIZE=4;num_stages=1;num_warps=4 # noqa: E225,E231,E702
elif (Ms, Ks) == (128, 128):
SPLIT_N=8;TILE_M=64;TILE_N=64;GROUP_SIZE=4;num_stages=1;num_warps=4 # noqa: E225,E231,E702
elif (M, K, N) == (1024,) * 3:
if (Ms, Ks) == (16, 16):
SPLIT_N=4;TILE_M=16;TILE_N=128;GROUP_SIZE=2;num_stages=1;num_warps=1 # noqa: E225,E231,E702
elif (Ms, Ks) == (32, 32):
SPLIT_N=8;TILE_M=32;TILE_N=64;GROUP_SIZE=2;num_stages=1;num_warps=1 # noqa: E225,E231,E702
elif (Ms, Ks) == (64, 64):
SPLIT_N=16;TILE_M=64;TILE_N=64;GROUP_SIZE=4;num_stages=1;num_warps=2 # noqa: E225,E231,E702
elif (Ms, Ks) == (128, 128):
SPLIT_N=16;TILE_M=64;TILE_N=64;GROUP_SIZE=4;num_stages=1;num_warps=4 # noqa: E225,E231,E702
elif (Ms, Ks) == (256, 256):
SPLIT_N=16;TILE_M=64;TILE_N=64;GROUP_SIZE=2;num_stages=1;num_warps=4 # noqa: E225,E231,E702
elif (M, K, N) == (2048,) * 3:
if (Ms, Ks) == (16, 16):
SPLIT_N=4;TILE_M=16;TILE_N=128;GROUP_SIZE=8;num_stages=1;num_warps=1 # noqa: E225,E231,E702
elif (Ms, Ks) == (32, 32):
SPLIT_N=4;TILE_M=32;TILE_N=64;GROUP_SIZE=4;num_stages=1;num_warps=1 # noqa: E225,E231,E702
elif (Ms, Ks) == (64, 64):
SPLIT_N=4;TILE_M=64;TILE_N=128;GROUP_SIZE=4;num_stages=1;num_warps=4 # noqa: E225,E231,E702
elif (Ms, Ks) == (128, 128):
SPLIT_N=8;TILE_M=64;TILE_N=64;GROUP_SIZE=4;num_stages=1;num_warps=4 # noqa: E225,E231,E702
elif (Ms, Ks) == (256, 256):
SPLIT_N=4;TILE_M=64;TILE_N=64;GROUP_SIZE=2;num_stages=1;num_warps=4 # noqa: E225,E231,E702
elif (M, K, N) == (4096,) * 3:
if (Ms, Ks) == (16, 16):
SPLIT_N=2;TILE_M=16;TILE_N=256;GROUP_SIZE=2;num_stages=1;num_warps=2 # noqa: E225,E231,E702
elif (Ms, Ks) == (32, 32):
SPLIT_N=2;TILE_M=32;TILE_N=64;GROUP_SIZE=2;num_stages=1;num_warps=1 # noqa: E225,E231,E702
elif (Ms, Ks) == (64, 64):
SPLIT_N=2;TILE_M=64;TILE_N=128;GROUP_SIZE=2;num_stages=1;num_warps=4 # noqa: E225,E231,E702
if SPLIT_N is None:
# Assume NVIDIA GeForce RTX 2060 SUPER:
# With the probality of 92% (99.9% when N > 512), the
# performance will not be worse more than 2% from the
# performance when using an optimal value. Otherwise, when N
# <= 512, using the following heuristics may give upto 15%
# lower performance.
SPLIT_N = {16: 1, 32: 2, 64: 4, 128: 8, 256: 16, 512: 8, 1024: 16, 4096: 32, 8192: 64}.get(N, 16)
if Ms >= 512 and N >= 2048:
SPLIT_N = 1
Ns = N // SPLIT_N
if TILE_M is None:
TILE_M = min(64 if Ns < 512 else 32, Ms)
if TILE_N is None:
TILE_N = min(64 if Ns < 512 else 32, Ns)
num_stages = num_stages or 1
if num_warps is None:
if min(M, N) > 1024:
num_warps = {16: 1, 32: 1, 64: 2}.get(Ms, 4)
elif min(M, N) == 1024:
num_warps = {16: 1, 32: 1, 64: 2}.get(Ms, 4)
elif min(M, N) == 256:
num_warps = {16: 1, 32: 4}.get(Ms, 4)
else:
num_warps = {16: 1, 32: 2}.get(Ms, 4)
GROUP_SIZE = GROUP_SIZE or 4
assert TILE_M <= Ms, dict(TILE_M=TILE_M, Ms=Ms)
assert TILE_N <= Ns, dict(TILE_N=TILE_N, Ns=Ns)
assert Ms <= M, dict(M=M, Ms=Ms)
assert Ns <= N, dict(N=N, Ns=Ns)
assert Ks <= K, dict(K=K, Ks=Ks)
return dict(TILE_M=TILE_M, TILE_N=TILE_N, GROUP_SIZE=GROUP_SIZE,
num_stages=num_stages, num_warps=num_warps, SPLIT_N=SPLIT_N, **extra)
def bsr_dense_addmm_meta(M, K, N, Ms, Ks, beta, alpha,
SPLIT_N=None, GROUP_SIZE_ROW=None, num_warps=None, num_stages=None, sparsity=None, dtype=None, **extra):
if dtype is None:
dtype = torch.float16
if sparsity is None:
sparsity = 0.5
if {SPLIT_N, num_warps, num_stages, GROUP_SIZE_ROW} == {None}:
device_name = torch.cuda.get_device_name()
key = (M, K, N, Ms, Ks, beta == 0, beta == 1, alpha == 1)
meta = get_meta('bsr_dense_addmm', key,
device_name, version=(0, dtype, sparsity))
if meta is None and sparsity != 0.5:
meta = get_meta('bsr_dense_addmm', key,
device_name, version=(0, dtype, 0.5))
if meta is not None:
meta.update(**extra)
return meta
SPLIT_N = SPLIT_N or max(N // Ms, 1)
GROUP_SIZE_ROW = GROUP_SIZE_ROW or 4
num_stages = num_stages or 1
num_warps = num_warps or 4
return dict(SPLIT_N=SPLIT_N, GROUP_SIZE_ROW=GROUP_SIZE_ROW, num_stages=num_stages, num_warps=num_warps, **extra)
class TensorAsKey:
"""A light-weight wrapper of a tensor that enables storing tensors as
keys with efficient memory reference based comparision as an
approximation to data equality based keys.
Motivation: the hash value of a torch tensor is tensor instance
based that does not use data equality and makes the usage of
tensors as keys less useful. For instance, the result of
``len({a.crow_indices(), a.crow_indices()})`` is `2`, although,
the tensor results from `crow_indices` method call are equal, in
fact, these share the same data storage.
On the other hand, for efficient caching of tensors we want to
avoid calling torch.equal that compares tensors item-wise.
TensorAsKey offers a compromise in that it guarantees key equality
of tensors that references data in the same storage in the same
manner and without accessing underlying data. However, this
approach does not always guarantee correctness. For instance, for
a complex tensor ``x``, we have ``TensorAsKey(x) ==
TensorAsKey(x.conj())`` while ``torch.equal(x, x.conj())`` would
return False.
"""
def __init__(self, obj):
def get_tensor_key(obj):
# Warning: TensorAsKey does not track negative nor
# conjugate bits of its input object because in the use
# case of wrapping compressed/plain indices of compressed
# sparse tensors (that are always integer tensors with
# non-negative items) these bits are never set. However,
# when extending the use of TensorAsKey to float or
# complex tensors, the values of these bits (see is_neg
# and is_conj methods) must be included in the key as
# well.
assert not (obj.dtype.is_floating_point or obj.dtype.is_complex), obj.dtype
return (obj.data_ptr(), obj.storage_offset(), obj.shape, obj.stride(), obj.dtype)
self._obj_ref = weakref.ref(obj)
if obj.layout is torch.strided:
self.key = get_tensor_key(obj)
elif obj.layout in {torch.sparse_csr, torch.sparse_bsr}:
self.key = (get_tensor_key(obj.crow_indices()), get_tensor_key(obj.col_indices()))
elif obj.layout in {torch.sparse_csc, torch.sparse_bsc}:
self.key = (get_tensor_key(obj.ccol_indices()), get_tensor_key(obj.row_indices()))
else:
raise NotImplementedError(obj.layout)
self._hash = hash(self.key)
def __hash__(self):
return self._hash
def __eq__(self, other):
if not isinstance(other, TensorAsKey):
return False
if self.obj is None or other.obj is None:
# dead objects always compare unequal unless these are
# same objects
return self is other
return self.key == other.key
@property
def obj(self):
"""Return object if alive, otherwise None."""
return self._obj_ref()
@lru_cache(maxsize=TORCH_SPARSE_BSR_SCATTER_MM_LRU_CACHE_SIZE)
def _bsr_scatter_mm_indices_data(indices_format, M, K, N, Ms, Ks, nbatches, SPLIT_N, compressed_sparse_tensor_as_key):
bsr = compressed_sparse_tensor_as_key.obj
assert bsr is not None
crow_indices, col_indices = bsr.crow_indices(), bsr.col_indices()
device = crow_indices.device
indices_dtype = torch.int32
if indices_format == 'bsr_strided_mm_compressed':
Ns = N // SPLIT_N
q_offsets_lst = []
b = torch.arange(SPLIT_N, dtype=indices_dtype, device=device) * Ns
for m in range(M // Ms):
r0 = crow_indices[m].item()
r1 = crow_indices[m + 1].item()
if r1 == r0:
continue
q_offsets_lst.append((col_indices[r0:r1] * (Ks * N)).repeat(SPLIT_N) + b.repeat_interleave(r1 - r0))
q_offsets = torch.cat(q_offsets_lst)
crow_indices_diff = crow_indices.diff()
non_zero_row_indices = crow_indices_diff.nonzero()
a = non_zero_row_indices * (Ms * N)
r_offsets = (a + b).view(-1)
c_indices = crow_indices
# swizzle operation: mm elements with longer sums are computed first:
nnz_per_row = crow_indices_diff[non_zero_row_indices].repeat_interleave(SPLIT_N)
nnz_per_row, indices = nnz_per_row.sort(descending=True, stable=True)
r_offsets = r_offsets[indices]
return (indices_format, c_indices, r_offsets, q_offsets)
elif indices_format == 'bsr_strided_mm':
Ns = N // SPLIT_N
p_offsets_lst = []
q_offsets_lst = []
b = torch.arange(SPLIT_N, dtype=indices_dtype, device=device) * Ns
for m in range(M // Ms):
r0 = crow_indices[m].item()
r1 = crow_indices[m + 1].item()
if r1 == r0:
continue
p_offsets_lst.append(torch.arange(r0, r1, dtype=indices_dtype, device=device).repeat(SPLIT_N))
q_offsets_lst.append((col_indices[r0:r1] * (Ks * N)).repeat(SPLIT_N) + b.repeat_interleave(r1 - r0))
q_offsets = torch.cat(q_offsets_lst)
crow_indices_diff = crow_indices.diff()
non_zero_row_indices = crow_indices_diff.nonzero()
a = non_zero_row_indices * (Ms * N)
r_offsets = (a + b).view(-1)
c_indices = torch.cat((crow_indices[:1],
torch.cumsum(crow_indices_diff[non_zero_row_indices].repeat_interleave(SPLIT_N), 0)))
p_offsets = torch.cat(p_offsets_lst)
return (indices_format, c_indices, r_offsets, p_offsets, q_offsets)
elif indices_format == 'scatter_mm':
Ns = Ms
c_indices = [0]
pq_offsets = []
# todo: eliminate inner for-loops for efficiency
for b in range(nbatches):
for m in range(M // Ms):
r0 = crow_indices[m].item()
r1 = crow_indices[m + 1].item()
for n in range(N // Ns):
c_indices.append(c_indices[-1] + r1 - r0)
for t in range(r1 - r0):
p = r0 + t
q = (col_indices[p].item() + b * (K // Ks)) * (N // Ns) + n
pq_offsets.append([p, q])
return (indices_format,
torch.tensor(c_indices, dtype=indices_dtype, device=device),
torch.tensor(pq_offsets, dtype=indices_dtype, device=device))
else:
raise ValueError(f'Invalid {indices_format=}. Expected bsr_strided_mm_compressed|bsr_strided_mm|scatter_mm')
def bsr_scatter_mm_indices_data(bsr, other, indices_format='bsr_strided_mm_compressed', **meta_input):
"""Computes indices data for :func:`scatter_mm` used in BSR and
strided tensor matrix multiplication.
"""
assert bsr.dense_dim() == 0
assert bsr.ndim == 2 # no batch dims
crow_indices = bsr.crow_indices()
col_indices = bsr.col_indices()
blocksize = bsr.values().shape[-2:]
M, K = bsr.shape
Ms, Ks = blocksize
K_, N = other.shape[-2:]
assert K_ == K
nbatches = other.shape[:-2].numel()
meta = scatter_mm_meta(M, K, N, Ms, Ks, **meta_input)
if 'allow_tf32' not in meta_input:
meta.update(allow_tf32=bsr.dtype in {torch.float16, torch.bfloat16})
SPLIT_N = meta['SPLIT_N']
indices_data = _bsr_scatter_mm_indices_data(
indices_format, M, K, N, Ms, Ks, nbatches, SPLIT_N, TensorAsKey(bsr))
if indices_format == 'bsr_strided_mm_compressed':
meta.update(is_compressed=True)
return indices_data + (meta,)
elif indices_format == 'bsr_strided_mm':
meta.update(is_compressed=False)
return indices_data + (meta,)
else:
return indices_data
def bsr_scatter_mm(bsr, other, indices_data=None, out=None):
"""BSR @ strided -> strided
"""
assert bsr.ndim == 2
assert other.ndim >= 2
Ms, Ks, Ns = bsr.shape[-2], bsr.shape[-1], other.shape[-1]
blocksize = bsr.values().shape[-2:]
if indices_data is None:
indices_data = bsr_scatter_mm_indices_data(bsr, other, indices_format='bsr_strided_mm_compressed')
indices_format = indices_data[0]
if out is None:
out = torch.empty((*other.shape[:-2], Ms, Ns), dtype=bsr.dtype, device=bsr.device)
out_shape = out.shape
out = as1Dbatch(out)
if bsr._nnz() == 0:
out.zero_()
elif indices_format in {'bsr_strided_mm_compressed', 'bsr_strided_mm'}:
out.zero_()
scatter_mm(bsr.values(), other, indices_data, accumulators=out)
elif indices_format == 'scatter_mm':
nbatches = other.shape[:-2].numel()
accumulators = torch.zeros((nbatches * Ms // blocksize[0] * Ns // blocksize[0], blocksize[0], blocksize[0]),
dtype=bsr.dtype, device=bsr.device)
others = (as1Dbatch(other)
.transpose(-2, -1)
.view(nbatches, Ns // blocksize[0], blocksize[0], Ks // blocksize[1], blocksize[1])
.movedim((3, 1, 4, 2), (1, 2, 3, 4)) # equivalent to .transpose(-3, -2).transpose(-2, -1).transpose(-4, -3)
.flatten(0, 2)
)
scatter_mm(bsr.values(), others, indices_data, accumulators=accumulators)
out.copy_(accumulators
.unflatten(0, (nbatches, Ms // blocksize[0], Ns // blocksize[0]))
.movedim((1, 2, 3, 4), (3, 1, 4, 2)) # equivalent to .transpose(-4, -3).transpose(-2, -1).transpose(-3, -2)
.reshape(nbatches, Ns, Ms)
.transpose(-2, -1))
else:
raise NotImplementedError(indices_format)
return out.view(out_shape)
def bsr_dense_addmm(
input: torch.Tensor,
bsr: torch.Tensor,
dense: torch.Tensor,
*,
beta=1,
alpha=1,
out: Optional[torch.Tensor] = None,
skip_checks: bool = False,
max_grid: Optional[Tuple[Optional[int], Optional[int], Optional[int]]] = None,
meta: Optional[dict] = None):
f_name = 'bsr_dense_addmm'
values = bsr.values()
crow_indices = bsr.crow_indices()
col_indices = bsr.col_indices()
batch_ndim = crow_indices.dim() - 1
M, K = bsr.shape[batch_ndim:batch_ndim + 2]
blocksize = values.shape[batch_ndim + 1:batch_ndim + 3]
N = dense.shape[-1]
# todo: implement checks
if out is None:
original_batch_dims_broadcasted = broadcast_batch_dims(f_name, bsr, dense)
out = dense.new_empty(original_batch_dims_broadcasted + (M, N))
if bsr._nnz() == 0 or alpha == 0 or N == 0 or M == 0 or K == 0:
if beta == 0:
out.zero_()
else:
out.copy_(input)
if beta != 1:
out.mul_(beta)
return out
if meta is None:
sparsity = round(1 - bsr._nnz() * blocksize[0] * blocksize[1] / (M * K), 2)
meta = bsr_dense_addmm_meta(M, K, N, blocksize[0], blocksize[1], beta, alpha, sparsity=sparsity, dtype=out.dtype)
out_backup = out
crow_indices, col_indices, values, input, dense, out = prepare_inputs(bsr, input, dense, out)
BM, BK = blocksize
SPLIT_N = meta.get('SPLIT_N', N // BM)
BN = N // SPLIT_N
out_untiled = out
out = tile_to_blocksize(out, (BM, BN))
dense = tile_to_blocksize(dense, (BK, BN))
input = tile_to_blocksize(input, (BM, BN))
dot_out_dtype = {torch.float16: tl.float32,
torch.bfloat16: tl.float32,
torch.float32: tl.float64,
torch.float64: tl.float64}[out.dtype]
n_batches = dense.size(0)
n_block_rows = crow_indices.size(-1) - 1
n_block_cols = dense.size(-3)
full_grid = (n_batches, n_block_cols, n_block_rows)
if max_grid is not None:
grid_blocks = tuple(max_grid[:3][::-1]) + (None,) * (3 - len(max_grid[:3]))
else:
grid_blocks = None
tensor_dims_map = {
values: (0, None, None),
crow_indices: (0, None, -1),
col_indices: (0, None, None),
input: (0, -3, -4),
dense: (0, -3, None),
out: (0, -3, -4),
}
assert alpha != 0
def kernel(grid, *sliced_tensors):
_bsr_strided_addmm_kernel[grid](
*ptr_stride_extractor(*sliced_tensors),
beta, alpha,
beta_is_one=beta == 1,
beta_is_nonzero=beta != 0,
alpha_is_one=alpha == 1,
BLOCKSIZE_ROW=BM,
BLOCKSIZE_INNER=BK,
BLOCKSIZE_COL=BN,
allow_tf32=dot_out_dtype == tl.float32,
acc_dtype=dot_out_dtype,
**meta)
launch_kernel(kernel, tensor_dims_map, full_grid, grid_blocks)
if out.data_ptr() != out_backup.data_ptr():
# prepare_inputs has made a copy of out, copy its content back
# to out_backup:
out_backup.copy_(out_untiled.view(out_backup.shape))
return out_backup
if has_triton():
import triton
import triton.language as tl
@triton.jit
def _sampled_addmm_kernel(
alpha,
beta,
IS_BETA_ZERO: tl.constexpr,
BLOCKSIZE_ROW: tl.constexpr,
BLOCKSIZE_COL: tl.constexpr,
k,
TILE_K: tl.constexpr,
values_ptr,
values_batch_stride,
values_nnz_stride,
values_row_block_stride,
values_col_block_stride,
crow_indices_ptr,
crow_indices_batch_stride,
crow_indices_stride,
col_indices_ptr,
col_indices_batch_stride,
col_indices_stride,
mat1_ptr,
mat1_batch_stride,
mat1_tiled_row_stride,
mat1_tiled_col_stride,
mat1_row_block_stride,
mat1_col_block_stride,
mat2_ptr,
mat2_batch_stride,
mat2_tiled_row_stride,
mat2_tiled_col_stride,
mat2_row_block_stride,
mat2_col_block_stride,
acc_dtype: tl.constexpr,
allow_tf32: tl.constexpr,
):
batch_pid = tl.program_id(axis=1)
row_block_pid = tl.program_id(axis=0)
crow_indices_offset_ptr = (
crow_indices_ptr
+ crow_indices_batch_stride * batch_pid
+ crow_indices_stride * row_block_pid
)
nnz_offset = tl.load(crow_indices_offset_ptr)
nnz_offset_next = tl.load(crow_indices_offset_ptr + crow_indices_stride)
# Compute nnz for the row with number row_block_pid.
# If it is zero, skip the row.
row_nnz = nnz_offset_next - nnz_offset
if row_nnz == 0:
return
row_block_arange = tl.arange(0, BLOCKSIZE_ROW)
col_block_arange = tl.arange(0, BLOCKSIZE_COL)
# Pointers are set to the first block of the current row.
values_block_ptrs = (
values_ptr
+ values_batch_stride * batch_pid
+ values_nnz_stride * nnz_offset
+ values_row_block_stride * row_block_arange[:, None]
+ values_col_block_stride * col_block_arange[None, :]
)
col_index_nnz_ptr = (
col_indices_ptr
+ col_indices_batch_stride * batch_pid
+ col_indices_stride * nnz_offset
)
# Advance mat1 to the current tiled row, ignore columns.
mat1_block_ptrs = (
mat1_ptr
+ mat1_batch_stride * batch_pid
+ mat1_tiled_row_stride * row_block_pid
+ mat1_row_block_stride * row_block_arange[:, None]
)
# Advance mat2 in batch and block col dimension.
mat2_block_ptrs = (
mat2_ptr
+ mat2_batch_stride * batch_pid
+ mat2_col_block_stride * col_block_arange[None, :]
)
k_tile_arange = tl.arange(0, TILE_K)
for _ in range(row_nnz):
acc_block = tl.zeros((BLOCKSIZE_ROW, BLOCKSIZE_COL), dtype=acc_dtype)
# find column block index
col_block = tl.load(col_index_nnz_ptr)
for k_tile in range(0, k, TILE_K):
k_offsets = k_tile + k_tile_arange
mask_k = k_offsets < k
mat1_block = tl.load(
mat1_block_ptrs
+ mat1_col_block_stride * k_offsets[None, :],
mask=mask_k[None, :], other=0.0
)
mat2_block = tl.load(
mat2_block_ptrs
+ mat2_tiled_col_stride * col_block
+ mat2_row_block_stride * k_offsets[:, None],
mask=mask_k[:, None], other=0.0
)
acc_block += tl.dot(mat1_block, mat2_block, allow_tf32=allow_tf32, out_dtype=acc_dtype)
if IS_BETA_ZERO:
acc_block *= alpha
else:
acc_block = alpha * acc_block + beta * tl.load(values_block_ptrs)
# write result
tl.store(values_block_ptrs, acc_block.to(values_ptr.dtype.element_ty))
# advance val/col_index ptrs to the next block in the row.
values_block_ptrs += values_nnz_stride
col_index_nnz_ptr += col_indices_stride
@triton.jit
def _bsr_strided_dense_rowspace_kernel(
# values prologue
values_ptr,
values_batch_stride,
values_nnz_stride,
values_row_block_stride,
values_col_block_stride,
# values epilogue
# crow_indices prologue
crow_indices_ptr,
crow_indices_batch_stride,
crow_indices_stride,
# crow_indices epilogue
# col_indices prologue
col_indices_ptr,
col_indices_batch_stride,
col_indices_stride,
# col_indices epilogue
# dense prologue
dense_ptr,
dense_batch_stride,
dense_tiled_row_stride,
dense_tiled_col_stride,
dense_row_block_stride,
dense_col_block_stride,
# dense epilogue
# output prologue
output_ptr,
output_batch_stride,
output_tiled_row_stride,
output_tiled_col_stride,
output_row_block_stride,
output_col_block_stride,
# output epilogue
#
# gh-113754: Always keep all constexpr arguments at the end of
# triton kernel arguments list because with triton 2.1 or
# earlier non-contiguous outputs will corrupt CUDA state due
# to a triton bug (fixed in openai/triton#2262).
BLOCKSIZE_ROW: tl.constexpr,
BLOCKSIZE_COL: tl.constexpr,
acc_dtype: tl.constexpr,
allow_tf32: tl.constexpr,
GROUP_SIZE_ROW: tl.constexpr,
):
batch_pid = tl.program_id(axis=2)
row_block_pid = tl.program_id(axis=0)
col_block_pid = tl.program_id(axis=1)
n_block_rows = tl.num_programs(axis=0)
n_block_cols = tl.num_programs(axis=1)
row_block_pid, col_block_pid = tl.swizzle2d(
row_block_pid, col_block_pid, n_block_rows, n_block_cols, GROUP_SIZE_ROW
)
crow_indices_offset_ptr = (
crow_indices_ptr
+ crow_indices_batch_stride * batch_pid
+ crow_indices_stride * row_block_pid
)
nnz_offset = tl.load(crow_indices_offset_ptr)
nnz_offset_next = tl.load(crow_indices_offset_ptr + crow_indices_stride)
# Compute nnz for the row with number row_block_pid.
# If it is zero, skip the row.
row_nnz = nnz_offset_next - nnz_offset
if row_nnz == 0:
return
row_block_arange = tl.arange(0, BLOCKSIZE_ROW)
col_block_arange = tl.arange(0, BLOCKSIZE_COL)
# Pointers are set to the first block of the current row.
values_block_ptrs = (
values_ptr
+ values_batch_stride * batch_pid
+ values_nnz_stride * nnz_offset
+ values_row_block_stride * row_block_arange[:, None]
+ values_col_block_stride * col_block_arange[None, :]
)
# NOTE: dense is advanced into all dimensions but the tiled row one.
# That will be advanced in the loop according to values in col_indices.
dense_block_ptrs = (
dense_ptr
+ dense_batch_stride * batch_pid
+ dense_tiled_col_stride * col_block_pid
+ dense_row_block_stride * col_block_arange[:, None]
+ dense_col_block_stride * row_block_arange[None, :]
)
# Pointers are set to exact write-to locations
output_ptrs = (
output_ptr
+ output_batch_stride * batch_pid
+ output_tiled_row_stride * row_block_pid
+ output_tiled_col_stride * col_block_pid
+ output_row_block_stride * row_block_arange[:, None]
+ output_col_block_stride * row_block_arange[None, :]
)
# Set pointer to the first nonzero element in the current row
col_index_nnz_ptr = (
col_indices_ptr
+ col_indices_batch_stride * batch_pid
+ col_indices_stride * nnz_offset
)
output_acc_block = tl.zeros((BLOCKSIZE_ROW, BLOCKSIZE_COL), dtype=acc_dtype)
for _ in range(row_nnz):
values_block = tl.load(values_block_ptrs)
# find which row of dense needs to get loaded
# for multiplication with values_block.
dense_row_idx = tl.load(col_index_nnz_ptr)
dense_block = tl.load(dense_block_ptrs + dense_tiled_row_stride * dense_row_idx)
# do block mm
output_acc_block += tl.dot(values_block, dense_block, allow_tf32=allow_tf32, out_dtype=acc_dtype)
# move val/col_index ptrs to the next block in the row
values_block_ptrs += values_nnz_stride
col_index_nnz_ptr += col_indices_stride
# write back the result
tl.store(output_ptrs, output_acc_block.to(output_ptr.dtype.element_ty))
def _run_sampled_addmm_kernel(
alpha, beta, is_beta_zero,
blocksize, k, tile_k,
values, crow_indices, col_indices,
mat1, mat2,
max_grid
):
n_batches = values.size(0)
n_block_rows = crow_indices.size(-1) - 1
full_grid = (n_batches, n_block_rows)
if max_grid is not None:
grid_blocks = tuple(max_grid[:2][::-1]) + (None,) * (2 - len(max_grid[:2]))
else:
grid_blocks = None
tensor_dims_map = {
values: (0, None),
crow_indices: (0, -1),
col_indices: (0, None),
mat1: (0, -4),
mat2: (0, None),
}
if values.dtype in (torch.half, torch.bfloat16):
acc_dtype = tl.float32
allow_tf32 = True
else:
acc_dtype = tl.float64
allow_tf32 = False
def kernel(grid, *sliced_tensors):
_sampled_addmm_kernel[grid](
alpha, beta, is_beta_zero,
*blocksize, k, tile_k,
*ptr_stride_extractor(*sliced_tensors),
acc_dtype=acc_dtype,
allow_tf32=allow_tf32,
num_stages=1,
num_warps=4
)
launch_kernel(kernel, tensor_dims_map, full_grid, grid_blocks)
def sampled_addmm(
input: torch.Tensor,
mat1: torch.Tensor,
mat2: torch.Tensor,
*,
beta=1.0,
alpha=1.0,
out: Optional[torch.Tensor] = None,
skip_checks: bool = False,
max_grid: Optional[Tuple[Optional[int], Optional[int], Optional[int]]] = None,
):
f_name = "sampled_addmm"
check_bsr_layout(f_name, input)
input_broadcasted = broadcast_batch_dims_bsr(f_name, input, mat1, mat2)
if not skip_checks:
check_device(f_name, mat1, input.device)
check_device(f_name, mat2, input.device)
if beta != 0.0 and input.dtype is torch.bool:
check(
False,
f"{f_name}(): having beta == {beta} not equal to 0.0 with boolean mask is not allowed."
)
if input.dtype is not torch.bool:
check_dtype(f_name, mat1, input.dtype)
check_dtype(f_name, mat2, input.dtype)
else:
check_dtype(f_name, mat1, mat2.dtype)
check_mm_compatible_shapes(f_name, mat1, mat2)
if out is not None:
check_bsr_layout(f_name, out)
check_device(f_name, out, mat1.device)
check_dtype(f_name, out, input.dtype)
check(
out.shape == input_broadcasted.shape
and out._nnz() == input._nnz(),
f"{f_name}(): Expects `out` to be of shape {input_broadcasted.shape} "
f"and with nnz equal to {input_broadcasted._nnz()} "
f"but got out.shape = {out.shape} and out.nnz = {out._nnz()}"
)
if out is None:
out = input_broadcasted.to(mat1.dtype, copy=True)
else:
out.copy_(input_broadcasted)
if out.numel() == 0 or out._nnz() == 0:
return out
blocksize = out.values().shape[-2:]
m = mat1.size(-2)
n = mat2.size(-1)
k = mat1.size(-1)
# NOTE: (m, 0) @ (0, n) == zeros(m, n)
if alpha == 0.0 or k == 0:
out.values().mul_(beta)
return out
# prepare inputs by reshaping them to be kernel-compatible
out_backup = out
crow_indices, col_indices, values, mat1, mat2 = prepare_inputs(out, mat1, mat2)
mat1 = tile_to_blocksize(mat1, (blocksize[0], k))
mat2 = tile_to_blocksize(mat2, (k, blocksize[1]))
tile_k = max(*blocksize)
_run_sampled_addmm_kernel(
alpha, beta, beta == 0.0,
blocksize, k, tile_k,
values, crow_indices, col_indices,
mat1, mat2,
max_grid
)
# If nnz x block strides are not the same in out_backup.values and values,
# it means that out_backup.values and values are not the views of each other,
# so we have to copy.
if out_backup.values().stride()[-3:] != values.stride()[-3:]:
out_backup.values().copy_(values.reshape(out_backup.values().shape))
return out_backup
def bsr_dense_mm(
bsr: torch.Tensor,
dense: torch.Tensor,
*,
out: Optional[torch.Tensor] = None,
skip_checks: bool = False,
max_grid: Optional[Tuple[Optional[int], Optional[int], Optional[int]]] = None,
meta: Optional[dict] = None
):
f_name = "bsr_dense_mm"
m, kl = bsr.shape[-2:]
if not skip_checks:
check_bsr_layout(f_name, bsr)
check_device(f_name, bsr, dense.device)
check_dtype(f_name, bsr, dense.dtype)
check_mm_compatible_shapes(f_name, bsr, dense)
n = dense.size(-1)
row_block, col_block = bsr.values().shape[-2:]
check_blocksize(f_name, (row_block, col_block))
check(
not n % 16,
f"{f_name}(): dense.size(-1) == {n} should be divisible by 16"
)
else:
kr, n = dense.shape[-2:]
original_batch_dims_broadcasted = broadcast_batch_dims(f_name, bsr, dense)
if out is not None and not skip_checks:
expected_out_shape = original_batch_dims_broadcasted + (m, n)
check(
out.shape == expected_out_shape,
"bsr_dense_mm(): `out` argument has wrong shape, "
f"expected {expected_out_shape}, but got {out.shape}.",
)
check(
out.is_contiguous() or out.transpose(-2, -1).is_contiguous(),
"bsr_dense_mm(): only row-major/col-major `out` arguments are supported, "
"i.e. (out.is_contiguous() or out.transpose(-2, -1).is_contiguous()) "
"should be True.",
)
# Allocate out
if out is None:
out = dense.new_empty(original_batch_dims_broadcasted + (m, n))
# Short circuit if lhs is zero
if bsr._nnz() == 0:
return out.zero_()
# with beta==0, addmm ignores input content, so we can use out
# as a placeholder for input because their shapes match:
return bsr_dense_addmm(out, bsr, dense, alpha=1, beta=0, out=out)
@triton.jit
def _bsr_softmax_kernel(
crow_indices_ptr,
crow_indices_batch_stride,
crow_indices_stride,
values_ptr,
values_batch_stride,
values_row_block_stride,
values_nnz_col_block_stride,
row_block, col_block,
MAX_ROW_NNZ: tl.constexpr,
TILE: tl.constexpr
):
batch_pid = tl.program_id(axis=2)
row_block_offset_pid = tl.program_id(axis=1)
row_block_pid = tl.program_id(axis=0)
crow_indices_offset_ptr = (
crow_indices_ptr
+ crow_indices_batch_stride * batch_pid
+ crow_indices_stride * row_block_pid
)
nnz_offset = tl.load(crow_indices_offset_ptr)
nnz_offset_next = tl.load(crow_indices_offset_ptr + crow_indices_stride)
# Compute nnz for the row with number row_block_pid.
# If it is zero, skip the row.
row_nnz = nnz_offset_next - nnz_offset
if row_nnz == 0:
return
row_arange = tl.arange(0, TILE)
mask = row_arange < row_nnz * col_block
curr_row_values_ptrs = (
values_ptr
+ values_batch_stride * batch_pid
+ values_row_block_stride * row_block_offset_pid
+ nnz_offset * col_block
)
# find max in the row
row_tile = tl.load(curr_row_values_ptrs + row_arange, mask=mask, other=-float('inf')).to(tl.float32)
max_row_value = tl.max(row_tile, axis=0)
for _ in range(TILE, MAX_ROW_NNZ, TILE):
row_arange += TILE
mask = row_arange < row_nnz * col_block
row_tile = tl.load(curr_row_values_ptrs + row_arange, mask=mask, other=-float('inf')).to(tl.float32)
curr_max_row_value = tl.max(row_tile, axis=0)
max_row_value = tl.where(max_row_value > curr_max_row_value, max_row_value, curr_max_row_value)
# find denominator for stable softmax
num = tl.exp(row_tile - max_row_value)
denom = tl.sum(num, axis=0)
for _ in range(TILE, MAX_ROW_NNZ, TILE):
row_arange -= TILE
mask = row_arange < row_nnz * col_block
row_tile = tl.load(curr_row_values_ptrs + row_arange, mask=mask, other=-float('inf')).to(tl.float32)
num = tl.exp(row_tile - max_row_value)
denom += tl.sum(num, axis=0)
# populate output
tl.store(curr_row_values_ptrs + row_arange, (num / denom).to(values_ptr.dtype.element_ty), mask=mask)
for _ in range(TILE, MAX_ROW_NNZ, TILE):
row_arange += TILE
mask = row_arange < row_nnz * col_block
row_tile = tl.load(curr_row_values_ptrs + row_arange, mask=mask, other=-float('inf')).to(tl.float32)
num = tl.exp(row_tile - max_row_value)
tl.store(curr_row_values_ptrs + row_arange, (num / denom).to(values_ptr.dtype.element_ty), mask=mask)
def bsr_softmax(input, max_row_nnz=None):
f_name = "bsr_softmax"
check_bsr_layout(f_name, input)
check_dtype(f_name, input, input.dtype)
if input._nnz() == 0 or input.numel() == 0:
return input.clone()
m, n = input.shape[-2:]
nnz = input._nnz()
row_block, col_block = input.values().shape[-2:]
if max_row_nnz is None:
max_row_nnz = triton.next_power_of_2(n)
else:
max_row_nnz = triton.next_power_of_2(max_row_nnz)
crow_indices = input.crow_indices().unsqueeze(0).flatten(0, -2)
# reshape values from
# (b1, ..., bn, nnz, row_block, col_block) to
# (b1 * ... * bn, row_block, nnz * col_block).
# This simplifies batch dim manipulation and unlocks
# the possibility to access all nnzs in any given row.
if input.values().transpose(-3, -2).is_contiguous():
# Need to clone to avoid `contiguous` returning a view.
values = input.values().clone()
else:
values = input.values()
values = values.transpose(-3, -2).contiguous().unsqueeze(0).flatten(0, -4).reshape(-1, row_block, nnz * col_block)
full_grid = (values.shape[0], row_block, m // row_block)
grid_blocks = None
tensor_dims_map = {
# We span nnz number of blocks, not nnz + 1,
# hence crow_indices[..., :-1]
crow_indices[..., :-1]: (0, None, -1),
values: (0, None, None),
}
def kernel(grid, *sliced_tensors):
_bsr_softmax_kernel[grid](
*ptr_stride_extractor(*sliced_tensors),
row_block, col_block,
max_row_nnz,
# Triton's max numel is bounded by 2 ** 17.
min(2 ** 17, max_row_nnz)
)
launch_kernel(kernel, tensor_dims_map, full_grid, grid_blocks)
values = values.reshape(-1, row_block, nnz, col_block).transpose(-3, -2).reshape(*input.values().shape)
return torch.sparse_compressed_tensor(
input.crow_indices().clone(),
input.col_indices().clone(),
values,
size=input.shape,
layout=input.layout
)
def _scaled_dot_product_attention(
query: torch.Tensor,
key: torch.Tensor,
value: torch.Tensor,
attn_mask: Optional[torch.Tensor],
dropout_p: float = 0.0,
is_causal: bool = False,
scale: Optional[float] = None
):
f_name = "_scaled_dot_product_attention"
check(
not is_causal,
f"{f_name}(): is_causal == True is not supported."
)
check(
attn_mask is not None,
f"{f_name}(): attn_mask == None is not supported."
)
assert attn_mask is not None
check(
attn_mask.layout == torch.sparse_bsr,
f"{f_name}(): "
f"attn_mask.layout must be {torch.sparse_bsr}, but got "
f"attn_mask.layout == {attn_mask.layout}."
)
check_device(f_name, key, query.device)
check_device(f_name, value, query.device)
check_device(f_name, attn_mask, query.device)
check_dtype(f_name, key, query.dtype)
check_dtype(f_name, value, query.dtype)
if attn_mask.dtype is not torch.bool:
check_dtype(f_name, attn_mask, query.dtype)
sdpa = sampled_addmm(attn_mask, query, key.transpose(-2, -1), beta=0.0, skip_checks=False)
if scale is None and query.size(-1) == 0 or scale == 0.0:
check(
False,
f"{f_name}(): current value of scale == {scale} "
"results in division by zero."
)
scale_factor = 1 / math.sqrt(query.size(-1)) if scale is None else scale
sdpa.values().mul_(scale_factor)
sdpa = bsr_softmax(sdpa)
torch.nn.functional.dropout(sdpa.values(), p=dropout_p, inplace=True)
sdpa = bsr_dense_mm(sdpa, value)
return sdpa
@triton.jit
def _scatter_mm2_kernel(
M: tl.constexpr, K: tl.constexpr, N: tl.constexpr,
blocks_ptr, blocks_stride_P, blocks_stride_M, blocks_stride_K,
others_ptr, others_stride_Q, others_stride_K, others_stride_N,
accumulators_ptr, accumulators_stride_R, accumulators_stride_M, accumulators_stride_N,
pq_offsets_ptr, pq_offsets_stride,
pq_ptr, pq_stride_T, pq_stride_1,
dot_out_dtype: tl.constexpr,
TILE_M: tl.constexpr,
TILE_N: tl.constexpr,
allow_tf32: tl.constexpr):
Ms = M // TILE_M
Ns = N // TILE_N
pid_t = tl.program_id(axis=0)
pid = tl.program_id(axis=1)
pid_m = pid // Ms
pid_n = pid % Ms
rm = (pid_m * TILE_M + tl.arange(0, TILE_M))
rn = (pid_n * TILE_N + tl.arange(0, TILE_N))
rk = tl.arange(0, K)
A_ptr = blocks_ptr + (rm[:, None] * blocks_stride_M + rk[None, :] * blocks_stride_K)
B_ptr = others_ptr + (rk[:, None] * others_stride_K + rn[None, :] * others_stride_N)
g0 = tl.load(pq_offsets_ptr + pid_t * pq_offsets_stride)
g1 = tl.load(pq_offsets_ptr + (pid_t + 1) * pq_offsets_stride)
if g0 == g1:
return
acc_block = tl.zeros((TILE_M, TILE_N), dtype=dot_out_dtype)
for i in range(g0, g1):
p = tl.load(pq_ptr + i * pq_stride_T)
q = tl.load(pq_ptr + i * pq_stride_T + pq_stride_1)
A = tl.load(A_ptr + p * blocks_stride_P)
B = tl.load(B_ptr + q * others_stride_Q)
acc_block += tl.dot(A, B, out_dtype=dot_out_dtype, allow_tf32=allow_tf32)
C_ptr = accumulators_ptr + pid_t * accumulators_stride_R + (
rm[:, None] * accumulators_stride_M + rn[None, :] * accumulators_stride_N)
tl.store(C_ptr, acc_block.to(accumulators_ptr.dtype.element_ty))
def _scatter_mm2(
blocks: torch.Tensor,
others: torch.Tensor,
pq_offsets: torch.Tensor,
pq_indices: torch.Tensor,
accumulators: torch.Tensor
):
P, M, K = blocks.shape
Q, _, N = others.shape
R, _, _ = accumulators.shape
meta = dict(TILE_M=max(16, M // 4), TILE_N=max(16, N // 4), num_stages=1, num_warps=2)
def grid(META):
return (pq_offsets.shape[0] - 1, triton.cdiv(M, META['TILE_M']) * triton.cdiv(N, META['TILE_N']), 1)
dot_out_dtype = {torch.float16: tl.float32,
torch.bfloat16: tl.float32,
torch.float32: tl.float64,
torch.float64: tl.float64}[accumulators.dtype]
if 'allow_tf32' not in meta:
meta.update(allow_tf32=dot_out_dtype == tl.float32)
_scatter_mm2_kernel[grid](
M, K, N,
blocks, blocks.stride(0), blocks.stride(1), blocks.stride(2),
others, others.stride(0), others.stride(1), others.stride(2),
accumulators, accumulators.stride(0), accumulators.stride(1), accumulators.stride(2),
pq_offsets, pq_offsets.stride(0),
pq_indices, pq_indices.stride(0), pq_indices.stride(1),
dot_out_dtype=dot_out_dtype,
**meta
)
@triton.jit
def _scatter_mm6_kernel(
nbatches, Ms, Ks: tl.constexpr, N,
blocks_ptr, blocks_stride_P, blocks_stride_M, blocks_stride_K,
others_ptr, others_stride_B, others_stride_K, others_stride_N,
accumulators_ptr, accumulators_stride_B, accumulators_stride_M, accumulators_stride_N,
c_indices_ptr, r_offsets_ptr,
p_offsets_ptr, q_offsets_ptr,
is_compressed: tl.constexpr,
dot_out_dtype: tl.constexpr,
SPLIT_N: tl.constexpr,
TILE_M: tl.constexpr,
TILE_N: tl.constexpr,
GROUP_SIZE: tl.constexpr,
allow_tf32: tl.constexpr):
Ns = N // SPLIT_N
BLOCKS_M = Ms // TILE_M
BLOCKS_N = Ns // TILE_N
pid_t_ = tl.program_id(axis=0)
pid = tl.program_id(axis=1)
pid_b = pid_t_ % nbatches
pid_t = pid_t_ // nbatches
num_pid_in_group = GROUP_SIZE * BLOCKS_N
group_id = pid // num_pid_in_group
first_pid_m = group_id * GROUP_SIZE
group_size_m = min(BLOCKS_M - first_pid_m, GROUP_SIZE)
pid_m = first_pid_m + (pid % group_size_m)
pid_n = (pid % num_pid_in_group) // group_size_m
rm = (pid_m * TILE_M + tl.arange(0, TILE_M))
rn = (pid_n * TILE_N + tl.arange(0, TILE_N))
rk = tl.arange(0, Ks)
A_ptr = blocks_ptr + (rm[:, None] * blocks_stride_M + rk[None, :] * blocks_stride_K)
B_ptr = others_ptr + pid_b * others_stride_B + (rk[:, None] * others_stride_K + rn[None, :] * others_stride_N)
# When is_compressed is True, r is the only variable that
# depends on pid_t. This property allows sorting r values
# before calling the kernel. The sorting of r is equivalent to
# defining swizzle operator outside of the kernel.
r = tl.load(r_offsets_ptr + pid_t)
if is_compressed:
m = (r // N) // Ms
n = (r % N) // Ns
r0 = tl.load(c_indices_ptr + m)
r1 = tl.load(c_indices_ptr + m + 1)
g0 = n * r1 + (SPLIT_N - n) * r0
nnz = r1 - r0
else:
g0 = tl.load(c_indices_ptr + pid_t)
g1 = tl.load(c_indices_ptr + pid_t + 1)
nnz = g1 - g0
q_ptr = q_offsets_ptr + g0
acc_block = tl.zeros((TILE_M, TILE_N), dtype=dot_out_dtype)
if is_compressed:
A_ptr += r0 * blocks_stride_P # type: ignore[possibly-undefined]
for _ in range(nnz):
q = tl.load(q_ptr)
B = tl.load(B_ptr + q)
A = tl.load(A_ptr)
acc_block += tl.dot(A, B, out_dtype=dot_out_dtype, allow_tf32=allow_tf32)
A_ptr += blocks_stride_P
q_ptr += 1
else:
p_ptr = p_offsets_ptr + g0
for _ in range(nnz):
q = tl.load(q_ptr)
B = tl.load(B_ptr + q)
p = tl.load(p_ptr)
A = tl.load(A_ptr + p * blocks_stride_P)
p_ptr += 1
q_ptr += 1
acc_block += tl.dot(A, B, out_dtype=dot_out_dtype, allow_tf32=allow_tf32)
C_ptr = accumulators_ptr + r + pid_b * accumulators_stride_B + (
rm[:, None] * accumulators_stride_M + rn[None, :] * accumulators_stride_N)
tl.store(C_ptr, acc_block.to(accumulators_ptr.dtype.element_ty))
def _scatter_mm6(
blocks: torch.Tensor,
others: torch.Tensor,
c_indices: torch.Tensor,
r_offsets: torch.Tensor,
p_offsets: torch.Tensor,
q_offsets: torch.Tensor,
meta: dict,
accumulators: torch.Tensor,
force_contiguous: bool = True,
):
SPLIT_N = meta['SPLIT_N']
P, Ms, Ks = blocks.shape
B, K_, N = others.shape
B_, M, N_ = accumulators.shape
assert N_ == N
Ns = N // SPLIT_N
assert B_ == B
def grid(META):
return (r_offsets.shape[0] * B, triton.cdiv(Ms, META['TILE_M']) * triton.cdiv(Ns, META['TILE_N']))
dot_out_dtype = {torch.float16: tl.float32,
torch.bfloat16: tl.float32,
torch.float32: tl.float64,
torch.float64: tl.float64}[accumulators.dtype]
if 'allow_tf32' not in meta:
meta.update(allow_tf32=dot_out_dtype == tl.float32)
assert c_indices.stride(0) == 1
assert r_offsets.stride(0) == 1
assert p_offsets.stride(0) == 1
assert q_offsets.stride(0) == 1
# Re non-contiguous tensor arguments. Sometimes triton kernel
# launches may fail with
#
# RuntimeError: Triton Error [CUDA]: an illegal memory access was encountered
#
# that appears to be case when the size of a non-contiguous
# tensor argument is larger than a certain threshold. Could
# this be related to shared memory or L1 cache size of a GPU
# card? In anycase, ensuring that tensor arguments are
# contiguous seems to avoid the above exception. So, in the
# following we'll always convert tensor arguments to
# C-contiguous tensors.
if force_contiguous:
blocks = blocks.contiguous()
others = others.contiguous()
if not accumulators.is_contiguous():
accumulators_ = accumulators.contiguous()
else:
accumulators_ = accumulators
else:
accumulators_ = accumulators
_scatter_mm6_kernel[grid](
B, Ms, Ks, N,
blocks, blocks.stride(0), blocks.stride(1), blocks.stride(2),
others, others.stride(0), others.stride(1), others.stride(2),
accumulators_, accumulators_.stride(0), accumulators_.stride(1), accumulators_.stride(2),
c_indices,
r_offsets,
p_offsets,
q_offsets,
dot_out_dtype=dot_out_dtype,
**meta
)
if force_contiguous and not accumulators.is_contiguous():
accumulators.copy_(accumulators_)
@triton.jit
def _bsr_strided_addmm_kernel(
# values prologue
values_ptr,
values_batch_stride,
values_nnz_stride,
values_row_block_stride,
values_col_block_stride,
# values epilogue
# crow_indices prologue
crow_indices_ptr,
crow_indices_batch_stride,
crow_indices_stride,
# crow_indices epilogue
# col_indices prologue
col_indices_ptr,
col_indices_batch_stride,
col_indices_stride,
# col_indices epilogue
# input prologue
input_ptr,
input_batch_stride,
input_tiled_row_stride,
input_tiled_col_stride,
input_row_block_stride,
input_col_block_stride,
# input epilogue
# dense prologue
dense_ptr,
dense_batch_stride,
dense_tiled_row_stride,
dense_tiled_col_stride,
dense_row_block_stride,
dense_col_block_stride,
# dense epilogue
# output prologue
output_ptr,
output_batch_stride,
output_tiled_row_stride,
output_tiled_col_stride,
output_row_block_stride,
output_col_block_stride,
# output epilogue
beta,
alpha,
beta_is_one: tl.constexpr,
beta_is_nonzero: tl.constexpr,
alpha_is_one: tl.constexpr,
BLOCKSIZE_ROW: tl.constexpr,
BLOCKSIZE_COL: tl.constexpr,
BLOCKSIZE_INNER: tl.constexpr,
acc_dtype: tl.constexpr,
allow_tf32: tl.constexpr,
GROUP_SIZE_ROW: tl.constexpr,
SPLIT_N: tl.constexpr
):
batch_pid = tl.program_id(axis=2)
row_block_pid = tl.program_id(axis=0)
col_block_pid = tl.program_id(axis=1)
n_block_rows = tl.num_programs(axis=0)
n_block_cols = tl.num_programs(axis=1)
row_block_pid, col_block_pid = tl.swizzle2d(
row_block_pid, col_block_pid, n_block_rows, n_block_cols, GROUP_SIZE_ROW
)
crow_indices_offset_ptr = (
crow_indices_ptr
+ crow_indices_batch_stride * batch_pid
+ crow_indices_stride * row_block_pid
)
nnz_offset = tl.load(crow_indices_offset_ptr)
nnz_offset_next = tl.load(crow_indices_offset_ptr + crow_indices_stride)
# Compute nnz for the row with number row_block_pid.
row_nnz = nnz_offset_next - nnz_offset
row_block_arange = tl.arange(0, BLOCKSIZE_ROW)
inner_block_arange = tl.arange(0, BLOCKSIZE_INNER)
col_block_arange = tl.arange(0, BLOCKSIZE_COL)
if beta_is_nonzero:
# Pointers are set to exact write-to locations
input_ptrs = (
input_ptr
+ input_batch_stride * batch_pid
+ input_tiled_row_stride * row_block_pid
+ input_tiled_col_stride * col_block_pid
+ input_row_block_stride * row_block_arange[:, None]
+ input_col_block_stride * col_block_arange[None, :]
)
# Pointers are set to the first block of the current row.
values_block_ptrs = (
values_ptr
+ values_batch_stride * batch_pid
+ values_nnz_stride * nnz_offset
+ values_row_block_stride * row_block_arange[:, None]
+ values_col_block_stride * inner_block_arange[None, :]
)
# NOTE: dense is advanced into all dimensions but the tiled row one.
# That will be advanced in the loop according to values in col_indices.
dense_block_ptrs = (
dense_ptr
+ dense_batch_stride * batch_pid
+ dense_tiled_col_stride * col_block_pid
+ dense_row_block_stride * inner_block_arange[:, None]
+ dense_col_block_stride * col_block_arange[None, :]
)
# Pointers are set to exact write-to locations
output_ptrs = (
output_ptr
+ output_batch_stride * batch_pid
+ output_tiled_row_stride * row_block_pid
+ output_tiled_col_stride * col_block_pid
+ output_row_block_stride * row_block_arange[:, None]
+ output_col_block_stride * col_block_arange[None, :]
)
# Set pointer to the first nonzero element in the current row
col_index_nnz_ptr = (
col_indices_ptr
+ col_indices_batch_stride * batch_pid
+ col_indices_stride * nnz_offset
)
# alpha is never 0
if beta_is_nonzero:
output_acc_block = tl.load(input_ptrs).to(acc_dtype) # type: ignore[possibly-undefined]
if not (beta_is_one and alpha_is_one):
beta_alpha = beta / alpha
output_acc_block *= beta_alpha
else:
output_acc_block = tl.zeros((BLOCKSIZE_ROW, BLOCKSIZE_COL), dtype=acc_dtype)
for _ in range(row_nnz):
values_block = tl.load(values_block_ptrs)
# find which row of dense needs to get loaded
# for multiplication with values_block.
dense_row_idx = tl.load(col_index_nnz_ptr)
dense_block = tl.load(dense_block_ptrs + dense_tiled_row_stride * dense_row_idx)
# do block mm
output_acc_block += tl.dot(values_block, dense_block, allow_tf32=allow_tf32, out_dtype=acc_dtype)
# move val/col_index ptrs to the next block in the row
values_block_ptrs += values_nnz_stride
col_index_nnz_ptr += col_indices_stride
if not alpha_is_one:
output_acc_block *= alpha
# write back the result
tl.store(output_ptrs, output_acc_block.to(output_ptr.dtype.element_ty))
else:
bsr_softmax = None # type: ignore[assignment]
bsr_dense_mm = None # type: ignore[assignment]
sampled_addmm = None # type: ignore[assignment]
_scaled_dot_product_attention = None # type: ignore[assignment]
_scatter_mm2 = None # type: ignore[assignment]
_scatter_mm6 = None # type: ignore[assignment]
_bsr_strided_addmm_kernel = None # type: ignore[assignment]
Reference in New Issue View Git Blame Copy Permalink