| # mypy: allow-untyped-decorators |
| # mypy: allow-untyped-defs |
| import math |
| import os |
| import weakref |
| from functools import lru_cache |
| from typing import Optional, Tuple |
| |
| import torch |
| from torch._dynamo.utils import warn_once |
| from torch.utils._triton import has_triton |
| |
| from ._triton_ops_meta import get_meta |
| |
| |
| 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, |
| _version=0, |
| **extra, |
| ): |
| # Specifying _version is useful for situations when one wants to |
| # discard existing triton kernel tuning results, say, in testing |
| # bsr_dense_addmm_meta functionality. |
| 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=(_version, dtype, sparsity) |
| ) |
| if meta is None and sparsity != 0.5: |
| meta = get_meta( |
| "bsr_dense_addmm", key, device_name, version=(_version, dtype, 0.5) |
| ) |
| if meta is None: |
| # find approximate meta such that N % SPLIT_N == 0. |
| matching_meta = get_meta( |
| "bsr_dense_addmm", |
| (*key[:2], "*", *key[3:]), |
| device_name, |
| version=(_version, dtype, 0.5), |
| ) |
| for mkey in sorted(matching_meta or {}): |
| meta_ = matching_meta[mkey] |
| n = mkey[2] |
| split_n = meta_["SPLIT_N"] |
| c = n // split_n |
| if N % c == 0 and n <= N: |
| meta = dict(meta_) |
| meta["SPLIT_N"] = N // c |
| if meta is not None: |
| meta.update(**extra) |
| return meta |
| else: |
| # see [Computing optimal kernel parameters] in |
| # _triton_ops_meta.py for ways to avoid this warning |
| # message |
| warn_once( |
| f"bsr_dense_addmm uses non-optimal triton kernel parameters for {M=} {K=} {N=} {Ms=}, {Ks=} {beta=} {alpha=}" |
| ) |
| |
| 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 _int_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, |
| ): |
| if out is None and dense.dtype is torch.int8: |
| f_name = "_int_bsr_dense_addmm" |
| crow_indices = bsr.crow_indices() |
| batch_ndim = crow_indices.dim() - 1 |
| M = bsr.shape[batch_ndim] |
| N = dense.shape[-1] |
| original_batch_dims_broadcasted = broadcast_batch_dims(f_name, bsr, dense) |
| out = torch.empty( |
| original_batch_dims_broadcasted + (M, N), |
| dtype=torch.int32, |
| device=dense.device, |
| ) |
| return bsr_dense_addmm( |
| input, |
| bsr, |
| dense, |
| beta=beta, |
| alpha=alpha, |
| out=out, |
| skip_checks=skip_checks, |
| max_grid=max_grid, |
| meta=meta, |
| ) |
| |
| |
| 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, |
| torch.int8: tl.int32, |
| torch.int32: tl.int32, |
| }[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, (torch.int8,)) |
| 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] |
