blksprs 1.3__tar.gz → 1.4.1__tar.gz

{blksprs-1.3 → blksprs-1.4.1}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: blksprs
-Version: 1.3
+Version: 1.4.1
 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
@@ -14,10 +14,8 @@ Requires-Dist: pytest-xdist; extra == "test"
 Requires-Dist: pytest-cov; extra == "test"
 Requires-Dist: coverage; extra == "test"
 Requires-Dist: matplotlib; extra == "test"
-Provides-Extra: deploy
-Requires-Dist: build; extra == "deploy"
-Requires-Dist: twine; extra == "deploy"
-Requires-Dist: pdoc3; extra == "deploy"
+Provides-Extra: build
+Requires-Dist: build; extra == "build"
 
 # blksprs
 
@@ -83,14 +81,7 @@ the [test cases](https://github.com/FelixSchoen/blksprs/blob/main/test/cases/tes
 
 ```python
 import torch
-
-from blksprs.layouting.sparsity_layout import build_sparsity_layout
-from blksprs.ops.conversion import to_sparse, to_dense
-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
-from blksprs.utils.tools import do_shape_blocksparse, undo_shape_blocksparse
+import blksprs as bs
 
 
 def test_readme():
@@ -112,47 +103,57 @@ def test_readme():
     y = torch.randn(size=(b, h, n, k), device="cuda").transpose(-1, -2).contiguous()
 
     # Convert tensors to three-dimensional (dense) tensors since Triton can only handle tensors of exactly three dimensions
-    x_dense, x_shape_original = do_shape_blocksparse(x)
-    y_dense, y_shape_original = do_shape_blocksparse(y)
+    x_dense, x_shape_original = bs.util.do_shape_blocksparse(x)
+    y_dense, y_shape_original = bs.util.do_shape_blocksparse(y)
 
     # Create sparsity layouts from existing tensors
-    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)
+    sparsity_layout_x = bs.layout.build_sparsity_layout(x_dense, sparsity_block_size,
+                                                        triton_block_size=triton_block_size)
+    sparsity_layout_y = bs.layout.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)
 
     # Convert tensors to sparse tensors for matrix multiplication
-    x_sparse = to_sparse(x_dense, sparsity_layout_x, sparsity_block_size, triton_block_size=triton_block_size)
-    y_sparse = to_sparse(y_dense, sparsity_layout_y, sparsity_block_size, triton_block_size=triton_block_size)
+    x_sparse = bs.to_sparse(x_dense, sparsity_layout_x, sparsity_block_size, triton_block_size=triton_block_size)
+    y_sparse = bs.to_sparse(y_dense, sparsity_layout_y, sparsity_block_size, triton_block_size=triton_block_size)
 
     # Perform matrix multiplication
-    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)
+    o_sparse = bs.matmul(x_sparse, sparsity_layout_x, y_sparse, sparsity_layout_y, sparsity_layout_o,
+                         sparsity_block_size,
+                         triton_block_size=triton_block_size)
+
+    # Apply element-wise operation
+    o_sparse = torch.add(o_sparse, 1)
+
+    o_dense = bs.to_dense(o_sparse, sparsity_layout_o, sparsity_block_size, triton_block_size=triton_block_size)
 
     # Sanity check
     o_torch = torch.matmul(x_dense, y_dense)
+    o_torch = torch.add(o_torch, 1)
 
     # Perform round trip to set sparse blocks to 0
-    o_torch_round_trip = to_dense(
-        to_sparse(o_torch, sparsity_layout_o, sparsity_block_size, triton_block_size=triton_block_size),
+    o_torch_round_trip = bs.to_dense(
+        bs.to_sparse(o_torch, sparsity_layout_o, sparsity_block_size, triton_block_size=triton_block_size),
         sparsity_layout_o, sparsity_block_size, fill_value=0, triton_block_size=triton_block_size)
 
     # Assert that the output is correct
     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 = build_sparsity_layout(o_dense, sparsity_block_size, triton_block_size=triton_block_size)
-    assert torch.allclose(actual_sparsity_layout_o, sparsity_layout_o)
+    actual_sparsity_layout_o = bs.layout.build_sparsity_layout(o_dense, sparsity_block_size,
+                                                               triton_block_size=triton_block_size)
+    assert torch.allclose(actual_sparsity_layout_o.to(torch.int), sparsity_layout_o)
 
     # Convert output tensor back to original shape
-    o = undo_shape_blocksparse(o_dense, x_shape_original)
+    o = bs.util.undo_shape_blocksparse(o_dense, x_shape_original)
 
     # Other available functions
-    transpose(o_sparse, sparsity_layout_o, sparsity_block_size, triton_block_size=triton_block_size)
-    softmax(o_sparse, sparsity_layout_o, sparsity_block_size, triton_block_size=triton_block_size)
-    row_wise_sum(o_sparse, sparsity_layout_o, sparsity_block_size, triton_block_size=triton_block_size)
+    bs.transpose(o_sparse, sparsity_layout_o, sparsity_block_size, triton_block_size=triton_block_size)
+    bs.softmax(o_sparse, sparsity_layout_o, sparsity_block_size, triton_block_size=triton_block_size)
+    bs.misc.row_wise_sum(o_sparse, sparsity_layout_o, sparsity_block_size, triton_block_size=triton_block_size)
+    bs.misc.row_wise_max(o_sparse, sparsity_layout_o, sparsity_block_size, triton_block_size=triton_block_size)
 
 
 def _get_random_sparsity_layout(b, m, n, sparsity_block_size, sparsity_percentage):

{blksprs-1.3 → blksprs-1.4.1}/README.md

@@ -62,14 +62,7 @@ the [test cases](https://github.com/FelixSchoen/blksprs/blob/main/test/cases/tes
 
 ```python
 import torch
-
-from blksprs.layouting.sparsity_layout import build_sparsity_layout
-from blksprs.ops.conversion import to_sparse, to_dense
-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
-from blksprs.utils.tools import do_shape_blocksparse, undo_shape_blocksparse
+import blksprs as bs
 
 
 def test_readme():
@@ -91,47 +84,57 @@ def test_readme():
     y = torch.randn(size=(b, h, n, k), device="cuda").transpose(-1, -2).contiguous()
 
     # Convert tensors to three-dimensional (dense) tensors since Triton can only handle tensors of exactly three dimensions
-    x_dense, x_shape_original = do_shape_blocksparse(x)
-    y_dense, y_shape_original = do_shape_blocksparse(y)
+    x_dense, x_shape_original = bs.util.do_shape_blocksparse(x)
+    y_dense, y_shape_original = bs.util.do_shape_blocksparse(y)
 
     # Create sparsity layouts from existing tensors
-    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)
+    sparsity_layout_x = bs.layout.build_sparsity_layout(x_dense, sparsity_block_size,
+                                                        triton_block_size=triton_block_size)
+    sparsity_layout_y = bs.layout.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)
 
     # Convert tensors to sparse tensors for matrix multiplication
-    x_sparse = to_sparse(x_dense, sparsity_layout_x, sparsity_block_size, triton_block_size=triton_block_size)
-    y_sparse = to_sparse(y_dense, sparsity_layout_y, sparsity_block_size, triton_block_size=triton_block_size)
+    x_sparse = bs.to_sparse(x_dense, sparsity_layout_x, sparsity_block_size, triton_block_size=triton_block_size)
+    y_sparse = bs.to_sparse(y_dense, sparsity_layout_y, sparsity_block_size, triton_block_size=triton_block_size)
 
     # Perform matrix multiplication
-    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)
+    o_sparse = bs.matmul(x_sparse, sparsity_layout_x, y_sparse, sparsity_layout_y, sparsity_layout_o,
+                         sparsity_block_size,
+                         triton_block_size=triton_block_size)
+
+    # Apply element-wise operation
+    o_sparse = torch.add(o_sparse, 1)
+
+    o_dense = bs.to_dense(o_sparse, sparsity_layout_o, sparsity_block_size, triton_block_size=triton_block_size)
 
     # Sanity check
     o_torch = torch.matmul(x_dense, y_dense)
+    o_torch = torch.add(o_torch, 1)
 
     # Perform round trip to set sparse blocks to 0
-    o_torch_round_trip = to_dense(
-        to_sparse(o_torch, sparsity_layout_o, sparsity_block_size, triton_block_size=triton_block_size),
+    o_torch_round_trip = bs.to_dense(
+        bs.to_sparse(o_torch, sparsity_layout_o, sparsity_block_size, triton_block_size=triton_block_size),
         sparsity_layout_o, sparsity_block_size, fill_value=0, triton_block_size=triton_block_size)
 
     # Assert that the output is correct
     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 = build_sparsity_layout(o_dense, sparsity_block_size, triton_block_size=triton_block_size)
-    assert torch.allclose(actual_sparsity_layout_o, sparsity_layout_o)
+    actual_sparsity_layout_o = bs.layout.build_sparsity_layout(o_dense, sparsity_block_size,
+                                                               triton_block_size=triton_block_size)
+    assert torch.allclose(actual_sparsity_layout_o.to(torch.int), sparsity_layout_o)
 
     # Convert output tensor back to original shape
-    o = undo_shape_blocksparse(o_dense, x_shape_original)
+    o = bs.util.undo_shape_blocksparse(o_dense, x_shape_original)
 
     # Other available functions
-    transpose(o_sparse, sparsity_layout_o, sparsity_block_size, triton_block_size=triton_block_size)
-    softmax(o_sparse, sparsity_layout_o, sparsity_block_size, triton_block_size=triton_block_size)
-    row_wise_sum(o_sparse, sparsity_layout_o, sparsity_block_size, triton_block_size=triton_block_size)
+    bs.transpose(o_sparse, sparsity_layout_o, sparsity_block_size, triton_block_size=triton_block_size)
+    bs.softmax(o_sparse, sparsity_layout_o, sparsity_block_size, triton_block_size=triton_block_size)
+    bs.misc.row_wise_sum(o_sparse, sparsity_layout_o, sparsity_block_size, triton_block_size=triton_block_size)
+    bs.misc.row_wise_max(o_sparse, sparsity_layout_o, sparsity_block_size, triton_block_size=triton_block_size)
 
 
 def _get_random_sparsity_layout(b, m, n, sparsity_block_size, sparsity_percentage):

blksprs-1.4.1/blksprs/__init__.py

@@ -0,0 +1,18 @@
+from blksprs.ops.conversion import to_dense, to_sparse
+from blksprs.ops.distribution import gather, scatter, scatter_reduce
+from blksprs.ops.exp import exp
+from blksprs.ops.matmul import matmul
+from blksprs.ops.softmax import softmax
+from blksprs.ops.transpose import transpose
+
+class layout:
+    from blksprs.layouting.distribution_layout import build_distribution_layout
+    from blksprs.layouting.sparsity_layout import build_sparsity_layout, build_sparsity_layout_adaption
+
+class misc:
+    from blksprs.misc.broadcast_ops import broadcast_add, broadcast_sub
+    from blksprs.misc.repeat_interleave import repeat_interleave
+    from blksprs.misc.row_wise import row_wise_sum, row_wise_max, row_wise_add, row_wise_sub
+
+class util:
+    from blksprs.utils.tools import do_shape_blocksparse, undo_shape_blocksparse, disable_validation

{blksprs-1.3 → blksprs-1.4.1}/blksprs/layouting/distribution_layout.py

@@ -31,7 +31,7 @@ def build_distribution_layout(indices: Tensor, sparsity_layout_indices: Tensor,
     sparsity_lut_i = torch.nonzero(sparsity_layout_indices).contiguous()
 
     output = torch.zeros(size_target[0], size_target[1] // sparsity_block_size, size_target[2] // sparsity_block_size,
-                         device=indices.device, dtype=torch.int32)
+                         dtype=torch.bool, device=indices.device)
 
     i_b, i_r, i_c = indices.size()
     i_b_s, i_r_s, i_c_s = indices.stride()

{blksprs-1.3 → blksprs-1.4.1}/blksprs/layouting/sparsity_layout.py

@@ -27,7 +27,7 @@ def build_sparsity_layout(x: Tensor, sparsity_block_size: int, triton_block_size
     validate_device(x)
 
     output = torch.zeros(x.size(0), x.size(1) // sparsity_block_size, x.size(2) // sparsity_block_size,
-                         device=x.device, dtype=torch.int32)
+                         dtype=torch.bool, device=x.device)
 
     x_b, x_r, x_c = x.size()
     x_b_s, x_r_s, x_c_s = x.stride()
@@ -117,7 +117,7 @@ def build_sparsity_layout_adaption(x: Tensor, sparsity_layout_from: Tensor,
     o_r = math.ceil(sparsity_layout_from.size(1) * sparsity_block_size_from // sparsity_block_size_to)
     o_c = math.ceil(sparsity_layout_from.size(2) * sparsity_block_size_from // sparsity_block_size_to)
 
-    output = torch.zeros(o_b, o_r, o_c, device=x.device, dtype=torch.int32)
+    output = torch.zeros(o_b, o_r, o_c, dtype=torch.bool, device=x.device)
 
     x_b, x_r, x_c = x.size()
     x_b_s, x_r_s, x_c_s = x.stride()

blksprs-1.3/blksprs/misc/broadcast_addition.py → blksprs-1.4.1/blksprs/misc/broadcast_ops.py

@@ -8,8 +8,8 @@ from blksprs.utils.validation import validate_contiguous, validate_device, \
     validate_sparsity_block_size, validate_triton_block_size
 
 
-def broadcast_addition(x: Tensor, y: Tensor, sparsity_layout_output: Tensor,
-                       sparsity_block_size: int, triton_block_size: int = None) -> Tensor:
+def broadcast_add(x: Tensor, y: Tensor, sparsity_layout_output: Tensor,
+                  sparsity_block_size: int, triton_block_size: int = None) -> Tensor:
     """Performs a broadcast and subsequent addition of two dense tensors x and y. Returns a block-sparse tensor in
         compressed form.
 
@@ -25,6 +25,9 @@ def broadcast_addition(x: Tensor, y: Tensor, sparsity_layout_output: Tensor,
             output tensor corresponds to x(i) + y(j).
 
     """
+    x = x.contiguous()
+    y = y.contiguous()
+
     validate_device(x, y)
     validate_contiguous(x, y)
     if x.size(-1) != y.size(-1):
@@ -70,12 +73,12 @@ def broadcast_addition(x: Tensor, y: Tensor, sparsity_layout_output: Tensor,
     return output
 
 
-def broadcast_subtraction(x: Tensor, y: Tensor, sparsity_layout_output: Tensor,
-                          sparsity_block_size: int, triton_block_size: int = None) -> Tensor:
-    """Wrapper for ``broadcast_addition`` with negated y.
+def broadcast_sub(x: Tensor, y: Tensor, sparsity_layout_output: Tensor,
+                  sparsity_block_size: int, triton_block_size: int = None) -> Tensor:
+    """Wrapper for ``broadcast_add`` with negated y.
 
     """
-    return broadcast_addition(x, torch.neg(y), sparsity_layout_output, sparsity_block_size, triton_block_size)
+    return broadcast_add(x, torch.neg(y), sparsity_layout_output, sparsity_block_size, triton_block_size)
 
 
 @triton.jit

{blksprs-1.3 → blksprs-1.4.1}/blksprs/misc/repeat_interleave.py

@@ -27,6 +27,8 @@ def repeat_interleave(x: Tensor, sparsity_layout: Tensor, repeats: int,
         Tensor: The sparsity layout of the resulting output tensor.
 
     """
+    x = x.contiguous()
+
     validate_dimensions(x)
     validate_contiguous(x)
     validate_device(x)