blksprs 1.0__py3-none-any.whl → 1.2__py3-none-any.whl

blksprs/utils/tools.py

@@ -15,33 +15,6 @@ def undo_shape_blocksparse(x: Tensor, shape: Size):
 
     return x.reshape((*shape[:-2], *x.shape[-2:]))
 
+
 def get_triton_block_size(sparsity_block_size: int, limit: int = 128):
     return min(sparsity_block_size, limit)
-
-#
-
-def slow_to_dense(x: Tensor, sparsity_layout: Tensor, sparsity_block_size: int):
-    output = torch.zeros(size=(sparsity_layout.size(0), sparsity_layout.size(1) * sparsity_block_size,
-                               sparsity_layout.size(2) * sparsity_block_size), device=x.device)
-    indices_sparse_blocks = sparsity_layout.nonzero(as_tuple=True)
-
-    for idx, (b, r, c) in enumerate(zip(*indices_sparse_blocks)):
-        t_r = r * sparsity_block_size
-        t_c = c * sparsity_block_size
-        to_insert = x[idx]
-        output[b, t_r:t_r + sparsity_block_size, t_c:t_c + sparsity_block_size] = to_insert
-
-    return output
-
-def slow_to_sparse(x: Tensor, sparsity_layout: Tensor, sparsity_block_size: int):
-    num_sparse_blocks = torch.sum(sparsity_layout.to(torch.int)).item()
-    output = torch.zeros(size=(num_sparse_blocks, sparsity_block_size, sparsity_block_size), device=x.device)
-    indices_sparse_blocks = sparsity_layout.nonzero(as_tuple=True)
-
-    for idx, (b, r, c) in enumerate(zip(*indices_sparse_blocks)):
-        t_r = r * sparsity_block_size
-        t_c = c * sparsity_block_size
-        to_insert = x[b, t_r:t_r + sparsity_block_size, t_c:t_c + sparsity_block_size]
-        output[idx] = to_insert
-
-    return output

blksprs/utils/validation.py

@@ -3,24 +3,45 @@ from torch import Tensor
 
 
 def validate_dimensions(*tensors: Tensor) -> None:
+    if _skip_validation():
+        return
+
     for tensor in tensors:
         if tensor.dim() != 3:
             raise ValueError("Tensor must have 3 dimensions")
 
 
 def validate_contiguous(*tensors: Tensor) -> None:
+    if _skip_validation():
+        return
+
     for tensor in tensors:
         if not tensor.is_contiguous():
             raise ValueError("Tensor must be contiguous")
 
 
 def validate_dtype_float(*tensors: Tensor) -> None:
+    if _skip_validation():
+        return
+
     for tensor in tensors:
         if tensor.dtype != torch.float32:
             raise ValueError("Tensor must have float32 dtype")
 
 
+def validate_dtype_int(*tensors: Tensor) -> None:
+    if _skip_validation():
+        return
+
+    for tensor in tensors:
+        if tensor.dtype != torch.int32 and tensor.dtype != torch.int64:
+            raise ValueError("Tensor must have int32 or int64 dtype")
+
+
 def validate_device(*tensors: Tensor) -> None:
+    if _skip_validation():
+        return
+
     device = None
 
     for i, tensor in enumerate(tensors):
@@ -33,13 +54,44 @@ def validate_device(*tensors: Tensor) -> None:
         if tensor.device != device:
             raise ValueError("Tensors must be on same device")
 
+
 def validate_sparsity(sparsity_block_size: int, *tensor_sparsity_layout_tuples: tuple[Tensor, Tensor]) -> None:
+    if _skip_validation():
+        return
+
     for (tensor, sparsity_layout) in tensor_sparsity_layout_tuples:
+        _validate_sparsity_layout_values(sparsity_layout)
+
         if not (tensor.size(-1) == tensor.size(-2) == sparsity_block_size):
             raise ValueError("Blocks not conforming to sparsity block size")
         if not tensor.size(0) == torch.sum(sparsity_layout.reshape(-1)):
             raise ValueError("Mismatch between sparsity layout and blocks")
 
+
+def _validate_sparsity_layout_values(sparsity_layout: Tensor):
+    if not torch.all(torch.logical_or(sparsity_layout == 0, sparsity_layout == 1)):
+        raise ValueError("Sparsity layout values must be either 0 or 1")
+
+def validate_sparsity_block_size(sparsity_block_size: int, *tensors):
+    if _skip_validation():
+        return
+
+    if not (sparsity_block_size & (sparsity_block_size - 1)) == 0:
+        raise ValueError("Sparsity block size must be a power of 2")
+
+    for tensor in tensors:
+        if not (tensor.size(-1) % sparsity_block_size == 0 and tensor.size(-2) % sparsity_block_size == 0):
+            raise ValueError("Tensor sizes must be divisible by sparsity block size")
+
 def validate_triton_block_size(triton_block_size: int, sparsity_block_size: int):
+    if _skip_validation():
+        return
+
+    if triton_block_size is None:
+        return
+
     if triton_block_size > sparsity_block_size:
-        raise ValueError("Triton block size cannot be larger than sparsity block size")
+        raise ValueError("Triton block size cannot be larger than sparsity block size")
+
+def _skip_validation():
+    return False

{blksprs-1.0.dist-info → blksprs-1.2.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: blksprs
-Version: 1.0
+Version: 1.2
 Summary: A lightweight library for operations on blocksparse matrices in PyTorch.
 Author-email: Felix Schön <schoen@kr.tuwien.ac.at>
 Project-URL: Homepage, https://github.com/FelixSchoen/blksprs
@@ -23,16 +23,21 @@ Requires-Dist: matplotlib; extra == "test"
 
 ## Overview
 
-A lightweight library for operations on blocksparse matrices in PyTorch.
+A lightweight and efficient library for operations on block-sparse matrices in PyTorch using Triton.
 
 Currently supported operations (includes gradient calculation):
 
-- Sparse matrix multiplication (_supports any combination of sparse and dense matrices due to support for `sparse = sparse @ sparse` matmul_)
+- Sparse matrix multiplication (_supports any combination of sparse and dense matrices due to support
+  for `sparse = sparse @ sparse` matmul_)
 - Softmax
 - Transposition
-- Conversion from and to sparse form
+- Gather
+- Scatter (_supports either no reduction or summation, gradients are only available for summation_)
+- Conversion to and from sparse form
+- Conversion to different sparsity layouts and different sparsity block sizes
 
-As with this library sparse matrices are represented using a tuple of `(matrix, sparsity_layout, sparsity_block_size)`, any element-wise operations can be applied in regular torch-like fashion.
+As with this library sparse matrices are represented using a tuple of `(matrix, sparsity_layout, sparsity_block_size)`,
+any element-wise operations can be applied in regular torch-like fashion.
 These include, e.g.,
 
 - Element-wise addition and subtraction
@@ -40,24 +45,45 @@ These include, e.g.,
 - Element-wise exponentiation
 - ...
 
+Note that in order to correctly apply element-wise operations between two sparse tensors their sparsity layouts have to
+match.
+
+Furthermore, the library provides a set of utility functions for the creation of sparsity layouts based on existing
+dense tensors.
+
 ## Installation
 
+Note that due to the dependency on [Triton](https://github.com/triton-lang/triton) this library is only compatible with
+the Linux platform.
+
 We recommend installing blksprs from [PyPI](https://pypi.org/project/blksprs/) using pip:
 
 ```pip install blksprs```
 
+### Dependencies
+
+- [PyTorch](https://pytorch.org/) (built with v2.4.0)
+- _[Triton](https://github.com/triton-lang/triton) (included with PyTorch)_
+
 ## Changelog
 
 See [`CHANGELOG.md`](https://github.com/FelixSchoen/blksprs/blob/main/CHANGELOG.md) for a detailed changelog.
 
 ## Usage
 
+We provide an example below to demonstrate the usage of the library.
+For more detailed examples, please refer to
+the [test cases](https://github.com/FelixSchoen/blksprs/blob/main/test/cases/test_blocksparse.py) which cover all
+implemented operations and functions.
+The example below can also be found in
+the [test cases](https://github.com/FelixSchoen/blksprs/blob/main/test/cases/test_readme.py).
+
 ```python
 import torch
 
-from blksprs.layouting.sparsity_layout import create_sparsity_layout
+from blksprs.layouting.sparsity_layout import build_sparsity_layout
 from blksprs.ops.conversion import to_sparse, to_dense
-from blksprs.ops.matmul_sss import matmul_sss
+from blksprs.ops.matmul import matmul
 from blksprs.ops.row_wise_sum import row_wise_sum
 from blksprs.ops.softmax import softmax
 from blksprs.ops.transpose import transpose
@@ -65,7 +91,7 @@ from blksprs.utils.tools import do_shape_blocksparse, undo_shape_blocksparse
 
 
 def test_readme():
-    # Set up parameters
+    # Set up parameters (batch size, number of heads, dimensions for matrices (m, k) and (n, k))
     b, h, m, n, k = 2, 4, 64, 64, 16
 
     # Percentage of blocks that will be sparse in the output for demonstration purposes
@@ -78,7 +104,6 @@ def test_readme():
     # If it is set to ``none`` a value will be chosen automatically
     triton_block_size = None
 
-
     # Initialise random (dense) tensors
     x = torch.randn(size=(b, h, m, k), device="cuda")
     y = torch.randn(size=(b, h, n, k), device="cuda").transpose(-1, -2).contiguous()
@@ -88,8 +113,8 @@ def test_readme():
     y_dense, y_shape_original = do_shape_blocksparse(y)
 
     # Create sparsity layouts from existing tensors
-    sparsity_layout_x = create_sparsity_layout(x_dense, sparsity_block_size, triton_block_size=triton_block_size)
-    sparsity_layout_y = create_sparsity_layout(y_dense, sparsity_block_size, triton_block_size=triton_block_size)
+    sparsity_layout_x = build_sparsity_layout(x_dense, sparsity_block_size, triton_block_size=triton_block_size)
+    sparsity_layout_y = build_sparsity_layout(y_dense, sparsity_block_size, triton_block_size=triton_block_size)
 
     # Create random sparsity layout for output tensor
     sparsity_layout_o = _get_random_sparsity_layout(b * h, m, n, sparsity_block_size, sparsity_percentage)
@@ -99,8 +124,8 @@ def test_readme():
     y_sparse = to_sparse(y_dense, sparsity_layout_y, sparsity_block_size, triton_block_size=triton_block_size)
 
     # Perform matrix multiplication
-    o_sparse = matmul_sss(x_sparse, y_sparse, sparsity_layout_x, sparsity_layout_y, sparsity_layout_o,
-                          sparsity_block_size, triton_block_size=triton_block_size)
+    o_sparse = matmul(x_sparse, sparsity_layout_x, y_sparse, sparsity_layout_y, sparsity_layout_o, sparsity_block_size,
+                      triton_block_size=triton_block_size)
     o_dense = to_dense(o_sparse, sparsity_layout_o, sparsity_block_size, triton_block_size=triton_block_size)
 
     # Sanity check
@@ -115,7 +140,7 @@ def test_readme():
     assert torch.allclose(o_dense, o_torch_round_trip, atol=2e-2)  # Note that small numerical differences are expected
 
     # Assert that the output has the correct sparsity layout
-    actual_sparsity_layout_o = create_sparsity_layout(o_dense, sparsity_block_size, triton_block_size=triton_block_size)
+    actual_sparsity_layout_o = build_sparsity_layout(o_dense, sparsity_block_size, triton_block_size=triton_block_size)
     assert torch.allclose(actual_sparsity_layout_o, sparsity_layout_o)
 
     # Convert output tensor back to original shape

blksprs-1.2.dist-info/RECORD

@@ -0,0 +1,17 @@
+blksprs/layouting/distribution_layout.py,sha256=GQ-ZRXbeImiLcbaqnL2FuUZ6DoFwmB0naT_YrOpD84Q,4940
+blksprs/layouting/sparsity_layout.py,sha256=TtADT_WWcZpW3zyGy6KAgkAo44gDryXZqdJLZGEX2V8,7895
+blksprs/misc/broadcast_addition.py,sha256=vf1Hdqz9Uyqykto3DCjmdyepMzpMXL238SpANQqRAwI,5297
+blksprs/ops/conversion.py,sha256=-AOzj_j3WrBLGIgd2oVPvYS8XKfzlvGtSIWzW_qP1lk,21260
+blksprs/ops/distribution.py,sha256=_fQb6fWpLxocAh86D74ATahChi0EK0eBb4eUOUEBVps,16769
+blksprs/ops/exp.py,sha256=qs8fVtCzxl4CKT4GepaqurjEL62jyi8VjMY12JFrFAU,3674
+blksprs/ops/matmul.py,sha256=x3lrYg4g8fIf5PeMtZY_SEpi11kP9RFcRoemCIxcSDE,11086
+blksprs/ops/row_wise_sum.py,sha256=ojuSejV37cLtRNS3lBfknA5KY3TEg8EHxOqVT6JZzoM,11387
+blksprs/ops/softmax.py,sha256=ZyeAVqmG_VzJ72FArGrpUSFfoSM4GPxyubrmNKERVIA,11654
+blksprs/ops/transpose.py,sha256=DVEXoxo2MoTNL3NZrjxsukMDrzk2vnEXL1uRnKFWkn0,6722
+blksprs/utils/benchmarking.py,sha256=4pLVlnPW_2EM-NT3n4SClaRznVYEljztLbJcccz8kZE,1360
+blksprs/utils/tools.py,sha256=P2UALvccRjJJ7w05YGuaxB3qmNObgct4idfM0jlE2wg,465
+blksprs/utils/validation.py,sha256=gJYZO5C48YUrXV3Fy_Z_lCaOpiFj951FT-Od7sKfprg,3007
+blksprs-1.2.dist-info/METADATA,sha256=4sbWg-lZK8DuRnkh3kh8toQRGMcBK9UlQtNLh4cU6mY,7209
+blksprs-1.2.dist-info/WHEEL,sha256=GV9aMThwP_4oNCtvEC2ec3qUYutgWeAzklro_0m4WJQ,91
+blksprs-1.2.dist-info/top_level.txt,sha256=qyp0IHeY3H2GQA97i4hk_To5rRBS2YcE1HRPSLy04fk,8
+blksprs-1.2.dist-info/RECORD,,

{blksprs-1.0.dist-info → blksprs-1.2.dist-info}/WHEEL

@@ -1,5 +1,5 @@
 Wheel-Version: 1.0
-Generator: setuptools (74.1.2)
+Generator: setuptools (75.1.0)
 Root-Is-Purelib: true
 Tag: py3-none-any
 

blksprs-1.0.dist-info/RECORD

@@ -1,14 +0,0 @@
-blksprs/layouting/sparsity_layout.py,sha256=fmzp0vasDDJuNwM-sPk3IBg1svCB2-ItELscppoSLPE,2553
-blksprs/ops/conversion.py,sha256=AYceIsv_g7xswoBqP3TGR2vjOjiLjTpwWG0vZ7XKCa8,10062
-blksprs/ops/exp.py,sha256=b0IuUVA_UoKNYDNT4Q3EFuXm7EEv_J2-DR7hfgCeT1Q,3222
-blksprs/ops/matmul_sss.py,sha256=34JSkO_9OOnQXB4KZHraElGDjbCx8p0dr9J5JebVdhY,10639
-blksprs/ops/row_wise_sum.py,sha256=ltoZpGVIApQBt_rbmknhsd-7MnibZRX6lkIzDnAC9k8,10462
-blksprs/ops/softmax.py,sha256=fJGKFshFMIVxaYJ_pgPD7EC0ooG-31XH9-qzRQ-xY5A,11018
-blksprs/ops/transpose.py,sha256=WEsXWRYDTWk2U36mt0aEQeReR3TG9TFrWAwVh3NNVYk,5985
-blksprs/utils/benchmarking.py,sha256=4pLVlnPW_2EM-NT3n4SClaRznVYEljztLbJcccz8kZE,1360
-blksprs/utils/tools.py,sha256=bqJtXXUKmcUxQsQ4ZkrpDZ5P7gnCburSW5hVd9U2M3E,1708
-blksprs/utils/validation.py,sha256=d4BFxzX-zVa5mUv_t3IW_bZbbP3vzSWan_KC1lyw7bs,1639
-blksprs-1.0.dist-info/METADATA,sha256=hrQQ8iK3-F2b38ogT8gp0O-sxm7UFn3i6QNv6vsS2so,5991
-blksprs-1.0.dist-info/WHEEL,sha256=cVxcB9AmuTcXqmwrtPhNK88dr7IR_b6qagTj0UvIEbY,91
-blksprs-1.0.dist-info/top_level.txt,sha256=qyp0IHeY3H2GQA97i4hk_To5rRBS2YcE1HRPSLy04fk,8
-blksprs-1.0.dist-info/RECORD,,