blksprs 1.2.1__tar.gz → 1.4__tar.gz

{blksprs-1.2.1 → blksprs-1.4}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: blksprs
-Version: 1.2.1
+Version: 1.4
 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
@@ -83,14 +83,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 +105,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)
+    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, 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.2.1 → blksprs-1.4}/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)
+    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, 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/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

blksprs-1.2.1/blksprs/misc/broadcast_addition.py → blksprs-1.4/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.
 
@@ -70,12 +70,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.4/blksprs/misc/repeat_interleave.py

@@ -0,0 +1,130 @@
+import torch
+import triton
+from torch import Tensor
+from triton import language as tl
+
+from blksprs.utils.tools import get_triton_block_size
+from blksprs.utils.validation import validate_contiguous, validate_device, \
+    validate_sparsity_block_size, validate_triton_block_size, validate_dimensions
+
+
+def repeat_interleave(x: Tensor, sparsity_layout: Tensor, repeats: int,
+                      sparsity_block_size: int, triton_block_size: int = None) -> tuple[Tensor, Tensor]:
+    """Repeats and interleaves the block-sparse tensor in compressed form.
+
+    Repeats each matrix contained in the tensors by ``repeats`` amount and places them consecutively in the output
+        tensor.
+
+    Args:
+        x (Tensor): A block-sparse tensor in compressed form.
+        sparsity_layout (Tensor): The sparsity layout of the block-sparse tensor.
+        repeats (int): The number of times to repeat the matrices.
+        sparsity_block_size (int): The size of the sparsity blocks.
+        triton_block_size (int): The block size to use for the triton kernel (default ``None``).
+
+    Returns:
+        Tensor: A block-sparse tensor in compressed form containing the repeated and interleaved matrices.
+        Tensor: The sparsity layout of the resulting output tensor.
+
+    """
+    validate_dimensions(x)
+    validate_contiguous(x)
+    validate_device(x)
+    validate_sparsity_block_size(sparsity_block_size, x)
+    validate_triton_block_size(triton_block_size, sparsity_block_size)
+
+    sparsity_layout_output = torch.repeat_interleave(sparsity_layout, 3, dim=0).contiguous()
+
+    sparsity_lut = torch.nonzero(sparsity_layout).contiguous()
+
+    sparsity_layout_output_flat = sparsity_layout_output.reshape(-1)
+    sparsity_output_reverse_lut = ((torch.cumsum(sparsity_layout_output_flat, dim=-1) - 1) *
+                                   (sparsity_layout_output_flat == 1) -
+                                   (1 * (sparsity_layout_output_flat == 0)))
+
+    n_sparse_blocks = torch.sum(sparsity_layout.to(torch.int)).item()
+
+    validate_contiguous(sparsity_layout, sparsity_lut, sparsity_layout_output, sparsity_output_reverse_lut)
+
+    output = torch.empty(n_sparse_blocks * repeats, sparsity_block_size, sparsity_block_size,
+                         dtype=x.dtype, device=x.device)
+
+    x_b, x_r, x_c = x.size()
+    x_b_s, x_r_s, x_c_s = x.stride()
+    s_lut_r, s_lut_c = sparsity_lut.size()
+    s_lut_r_s, s_lut_c_s = sparsity_lut.stride()
+    o_b, o_r, o_c = output.size()
+    o_b_s, o_r_s, o_c_s = output.stride()
+    s_l_o_b, s_l_o_r, s_l_o_c = sparsity_layout_output.size()
+    s_l_o_b_s, s_l_o_r_s, s_l_o_c_s = sparsity_layout_output.stride()
+
+    if triton_block_size is None:
+        triton_block_size = get_triton_block_size(sparsity_block_size)
+
+    triton_grid = lambda meta: [x_b,
+                                triton.cdiv(x_r, meta["TRITON_BLOCK_SIZE"]),
+                                triton.cdiv(x_c, meta["TRITON_BLOCK_SIZE"])]
+
+    (kernel_repeat_interleave[triton_grid]
+     (x,
+      x_b, x_b_s, x_r_s, x_c_s,
+      sparsity_lut, s_lut_r, s_lut_r_s, s_lut_c_s,
+      output,
+      o_b, o_b_s, o_r_s, o_c_s,
+      s_l_o_b, s_l_o_b_s, s_l_o_r_s, s_l_o_c_s,
+      sparsity_output_reverse_lut,
+      repeats,
+      triton_block_size))
+
+    return output, sparsity_layout_output
+
+
+@triton.jit
+def kernel_repeat_interleave(x,
+                             x_b, x_b_s, x_r_s, x_c_s,
+                             s_lut, s_lut_r, s_lut_r_s, s_lut_c_s,
+                             o,
+                             o_b, o_b_s, o_r_s, o_c_s,
+                             s_l_o_b, s_l_o_b_s, s_l_o_r_s, s_l_o_c_s,
+                             r_lut_o,
+                             repeats,
+                             TRITON_BLOCK_SIZE: tl.constexpr) -> None:
+    # Get triton block indices
+    pid_blk = tl.program_id(axis=0)
+    pid_row = tl.program_id(axis=1)
+    pid_col = tl.program_id(axis=2)
+
+    # Get sparsity index of current output block consisting of its batch, row, and column index
+    spa_bat_idx = (pid_blk * s_lut_r_s + 0 * s_lut_c_s)
+    spa_bat_msk = (spa_bat_idx < s_lut_r * s_lut_r_s)
+    spa_bat = tl.load(s_lut + spa_bat_idx, mask=spa_bat_msk)
+
+    spa_row_idx = (pid_blk * s_lut_r_s + 1 * s_lut_c_s)
+    spa_row_msk = (spa_row_idx < s_lut_r * s_lut_r_s)
+    spa_row = tl.load(s_lut + spa_row_idx, mask=spa_row_msk)
+
+    spa_col_idx = (pid_blk * s_lut_r_s + 2 * s_lut_c_s)
+    spa_col_msk = (spa_col_idx < s_lut_r * s_lut_r_s)
+    spa_col = tl.load(s_lut + spa_col_idx, mask=spa_col_msk)
+
+    # Load block
+    blk_x_idx = ((pid_blk * x_b_s) +
+                 ((pid_row * TRITON_BLOCK_SIZE + tl.arange(0, TRITON_BLOCK_SIZE)) * x_r_s)[:, None] +
+                 ((pid_col * TRITON_BLOCK_SIZE + tl.arange(0, TRITON_BLOCK_SIZE)) * x_c_s)[None, :])
+    blk_x_msk = (blk_x_idx < x_b * x_b_s)
+    blk_x = tl.load(x + blk_x_idx, mask=blk_x_msk)
+
+    for repeat in range(repeats):
+        # Get reverse sparsity index
+        rev_idx_spa_idx = ((spa_bat * repeats + repeat) * s_l_o_b_s +
+                           spa_row * s_l_o_r_s +
+                           spa_col * s_l_o_c_s)
+        rev_idx_spa_msk = (rev_idx_spa_idx < s_l_o_b * s_l_o_b_s)
+        rev_idx_spa = tl.load(r_lut_o + rev_idx_spa_idx, mask=rev_idx_spa_msk).to(tl.int32)
+
+        # Store block
+        blk_o_idx = ((rev_idx_spa * o_b_s) +
+                     ((pid_row * TRITON_BLOCK_SIZE + tl.arange(0, TRITON_BLOCK_SIZE)) * o_r_s)[:, None] +
+                     ((pid_col * TRITON_BLOCK_SIZE + tl.arange(0, TRITON_BLOCK_SIZE)) * o_c_s)[None, :])
+        blk_o_msk = (blk_o_idx < o_b * o_b_s)
+        tl.store(o + blk_o_idx, blk_x, mask=blk_o_msk)

blksprs-1.4/blksprs/misc/row_wise.py

@@ -0,0 +1,386 @@
+import torch
+import triton
+from torch import Tensor
+from triton import language as tl
+
+from blksprs.utils.tools import get_triton_block_size
+from blksprs.utils.validation import validate_dimensions, validate_contiguous, validate_device, validate_sparsity, \
+    validate_sparsity_block_size, validate_triton_block_size
+
+
+def row_wise_sum(x: Tensor, sparsity_layout: Tensor, sparsity_block_size: int,
+                 flag_slice_only: bool = False, triton_block_size: int = None) -> tuple[Tensor, Tensor]:
+    """Computes the row-wise sum of a block-sparse tensor.
+
+    Returns a block-sparse tensor in compressed form with only one block per row, where the first entry contains the sum
+        of the corresponding row.
+
+    Note:
+        If ``flag_slice_only`` is set the output will be of shape ``[x.size(0), x.size(1), 1]``.
+
+    Args:
+        x (Tensor): A block-sparse tensor in compressed form.
+        sparsity_layout (Tensor): The sparsity layout of the block-sparse tensor.
+        sparsity_block_size (int): The size of the sparsity blocks.
+        flag_slice_only (bool, optional): If set the output will be of shape ``[x.size(0), x.size(1), 1]``
+            (default ``False``).
+        triton_block_size (int): The block size to use for the triton kernel (default ``None``).
+
+    Returns:
+        tuple[Tensor, Tensor]: A tuple containing a block-sparse tensor in compressed form containing the row-wise sum
+            of the input and the sparsity layout of the output tensor.
+
+    """
+    validate_dimensions(x)
+    validate_contiguous(x)
+    validate_device(x)
+    validate_sparsity(sparsity_block_size, (x, sparsity_layout))
+    validate_sparsity_block_size(sparsity_block_size, x)
+    validate_triton_block_size(triton_block_size, sparsity_block_size)
+
+    sparsity_lut = torch.nonzero(sparsity_layout).contiguous()
+
+    sparsity_layout_output, _ = torch.max(sparsity_layout, dim=-1, keepdim=True)
+    sparsity_layout_output_flat = sparsity_layout_output.reshape(-1)
+    sparsity_reverse_lut_output = ((torch.cumsum(sparsity_layout_output_flat, dim=-1) - 1) *
+                                   (sparsity_layout_output_flat == 1) -
+                                   (1 * (sparsity_layout_output_flat == 0)))
+
+    n_sparse_blocks_output = torch.sum(sparsity_layout_output.to(torch.int)).item()
+
+    validate_contiguous(sparsity_layout, sparsity_lut,
+                        sparsity_layout_output, sparsity_reverse_lut_output)
+
+    output = torch.zeros(size=(n_sparse_blocks_output,
+                               sparsity_block_size,
+                               1 if flag_slice_only else sparsity_block_size),
+                         device=x.device)
+
+    x_b, x_r, x_c = x.size()
+    x_b_s, x_r_s, x_c_s = x.stride()
+    s_lut_x_r, s_lut_x_c = sparsity_lut.size()
+    s_lut_x_r_s, s_lut_x_c_s = sparsity_lut.stride()
+    o_b, o_r, o_c = output.size()
+    o_b_s, o_r_s, o_c_s = output.stride()
+    s_l_o_b, s_l_o_r, s_l_o_c = sparsity_layout_output.size()
+    s_l_o_b_s, s_l_o_r_s, s_l_o_c_s = sparsity_layout_output.stride()
+
+    if triton_block_size is None:
+        triton_block_size = get_triton_block_size(sparsity_block_size)
+
+    triton_grid = lambda meta: [x_b,
+                                triton.cdiv(x_r, meta["TRITON_BLOCK_SIZE"]),
+                                triton.cdiv(x_c, meta["TRITON_BLOCK_SIZE"])]
+
+    (kernel_blocksparse_row_wise_sum[triton_grid]
+     (x,
+      x_b, x_b_s, x_r_s, x_c_s,
+      sparsity_lut, s_lut_x_r, s_lut_x_r_s, s_lut_x_c_s,
+      output,
+      o_b, o_b_s, o_r_s,
+      s_l_o_b, s_l_o_b_s, s_l_o_r_s,
+      sparsity_reverse_lut_output,
+      triton_block_size))
+
+    return (output, sparsity_layout_output)
+
+
+@triton.jit
+def kernel_blocksparse_row_wise_sum(x,
+                                    x_b, x_b_s, x_r_s, x_c_s,
+                                    s_lut_x, s_lut_x_r, s_lut_x_r_s, s_lut_x_c_s,
+                                    o,
+                                    o_b, o_b_s, o_r_s,
+                                    s_l_o_b, s_l_o_b_s, s_l_o_r_s,
+                                    r_lut_o,
+                                    TRITON_BLOCK_SIZE: tl.constexpr) -> None:
+    pid_blk = tl.program_id(axis=0)
+    pid_row = tl.program_id(axis=1)
+    pid_col = tl.program_id(axis=2)
+
+    # Get position of current sparsity block consisting of its batch and row index
+    spa_bat_idx = (pid_blk * s_lut_x_r_s + 0 * s_lut_x_c_s)
+    spa_bat_msk = (spa_bat_idx < s_lut_x_r * s_lut_x_r_s)
+    spa_bat = tl.load(s_lut_x + spa_bat_idx, mask=spa_bat_msk)
+
+    spa_row_idx = (pid_blk * s_lut_x_r_s + 1 * s_lut_x_c_s)
+    spa_row_msk = (spa_row_idx < s_lut_x_r * s_lut_x_r_s)
+    spa_row = tl.load(s_lut_x + spa_row_idx, mask=spa_row_msk)
+
+    # Load reverse sparsity index for current block
+    rev_idx_spa_idx = (spa_bat * s_l_o_b_s +
+                       spa_row * s_l_o_r_s)
+    rev_idx_spa_msk = (rev_idx_spa_idx < s_l_o_b * s_l_o_b_s)
+    rev_idx_spa = tl.load(r_lut_o + rev_idx_spa_idx, mask=rev_idx_spa_msk).to(tl.int32)
+
+    blk_idx = ((pid_blk * x_b_s) +
+               ((pid_row * TRITON_BLOCK_SIZE + tl.arange(0, TRITON_BLOCK_SIZE)) * x_r_s)[:, None] +
+               ((pid_col * TRITON_BLOCK_SIZE + tl.arange(0, TRITON_BLOCK_SIZE)) * x_c_s)[None, :])
+    blk_msk = (blk_idx < x_b * x_b_s)
+    blk = tl.load(x + blk_idx, mask=blk_msk)
+
+    buf = tl.reshape(tl.sum(blk, axis=-1), (TRITON_BLOCK_SIZE, 1))
+
+    o_idx = (rev_idx_spa * o_b_s +
+             ((pid_row * TRITON_BLOCK_SIZE + tl.arange(0, TRITON_BLOCK_SIZE)) * o_r_s)[:, None] +
+             (tl.arange(0, 1))[None, :])
+    o_msk = (o_idx < o_b * o_b_s)
+    tl.atomic_add(o + o_idx, buf, o_msk)
+
+
+def row_wise_max(x: Tensor, sparsity_layout: Tensor, sparsity_block_size: int,
+                 flag_slice_only: bool = False, triton_block_size: int = None) -> tuple[Tensor, Tensor]:
+    """Computes the row-wise max of a block-sparse tensor.
+
+    Returns a block-sparse tensor in compressed form with only one block per row, where the first entry contains the
+        maximum of the corresponding row.
+
+    Note:
+        If ``flag_slice_only`` is set the output will be of shape ``[x.size(0), x.size(1), 1]``.
+
+    Args:
+        x (Tensor): A block-sparse tensor in compressed form.
+        sparsity_layout (Tensor): The sparsity layout of the block-sparse tensor.
+        sparsity_block_size (int): The size of the sparsity blocks.
+        flag_slice_only (bool, optional): If set the output will be of shape ``[x.size(0), x.size(1), 1]``
+            (default ``False``).
+        triton_block_size (int): The block size to use for the triton kernel (default ``None``).
+
+    Returns:
+        tuple[Tensor, Tensor]: A tuple containing a block-sparse tensor in compressed form containing the row-wise max
+            of the input and the sparsity layout of the output tensor.
+
+    """
+    validate_dimensions(x)
+    validate_contiguous(x)
+    validate_device(x)
+    validate_sparsity(sparsity_block_size, (x, sparsity_layout))
+    validate_sparsity_block_size(sparsity_block_size, x)
+    validate_triton_block_size(triton_block_size, sparsity_block_size)
+
+    sparsity_lut = torch.nonzero(sparsity_layout).contiguous()
+
+    sparsity_layout_output, _ = torch.max(sparsity_layout, dim=-1, keepdim=True)
+    sparsity_layout_output_flat = sparsity_layout_output.reshape(-1)
+    sparsity_reverse_lut_output = ((torch.cumsum(sparsity_layout_output_flat, dim=-1) - 1) *
+                                   (sparsity_layout_output_flat == 1) -
+                                   (1 * (sparsity_layout_output_flat == 0)))
+
+    n_sparse_blocks_output = torch.sum(sparsity_layout_output.to(torch.int)).item()
+
+    validate_contiguous(sparsity_layout, sparsity_lut,
+                        sparsity_layout_output, sparsity_reverse_lut_output)
+
+    output = torch.full(size=(n_sparse_blocks_output,
+                              sparsity_block_size,
+                              1 if flag_slice_only else sparsity_block_size),
+                        fill_value=float("-inf"),
+                        device=x.device)
+
+    x_b, x_r, x_c = x.size()
+    x_b_s, x_r_s, x_c_s = x.stride()
+    s_lut_x_r, s_lut_x_c = sparsity_lut.size()
+    s_lut_x_r_s, s_lut_x_c_s = sparsity_lut.stride()
+    o_b, o_r, o_c = output.size()
+    o_b_s, o_r_s, o_c_s = output.stride()
+    s_l_o_b, s_l_o_r, s_l_o_c = sparsity_layout_output.size()
+    s_l_o_b_s, s_l_o_r_s, s_l_o_c_s = sparsity_layout_output.stride()
+
+    if triton_block_size is None:
+        triton_block_size = get_triton_block_size(sparsity_block_size)
+
+    triton_grid = lambda meta: [x_b,
+                                triton.cdiv(x_r, meta["TRITON_BLOCK_SIZE"]),
+                                triton.cdiv(x_c, meta["TRITON_BLOCK_SIZE"])]
+
+    (kernel_blocksparse_row_wise_max[triton_grid]
+     (x,
+      x_b, x_b_s, x_r_s, x_c_s,
+      sparsity_lut, s_lut_x_r, s_lut_x_r_s, s_lut_x_c_s,
+      output,
+      o_b, o_b_s, o_r_s,
+      s_l_o_b, s_l_o_b_s, s_l_o_r_s,
+      sparsity_reverse_lut_output,
+      triton_block_size))
+
+    return output, sparsity_layout_output
+
+
+@triton.jit
+def kernel_blocksparse_row_wise_max(x,
+                                    x_b, x_b_s, x_r_s, x_c_s,
+                                    s_lut_x, s_lut_x_r, s_lut_x_r_s, s_lut_x_c_s,
+                                    o,
+                                    o_b, o_b_s, o_r_s,
+                                    s_l_o_b, s_l_o_b_s, s_l_o_r_s,
+                                    r_lut_o,
+                                    TRITON_BLOCK_SIZE: tl.constexpr) -> None:
+    pid_blk = tl.program_id(axis=0)
+    pid_row = tl.program_id(axis=1)
+    pid_col = tl.program_id(axis=2)
+
+    # Get position of current sparsity block consisting of its batch and row index
+    spa_bat_idx = (pid_blk * s_lut_x_r_s + 0 * s_lut_x_c_s)
+    spa_bat_msk = (spa_bat_idx < s_lut_x_r * s_lut_x_r_s)
+    spa_bat = tl.load(s_lut_x + spa_bat_idx, mask=spa_bat_msk)
+
+    spa_row_idx = (pid_blk * s_lut_x_r_s + 1 * s_lut_x_c_s)
+    spa_row_msk = (spa_row_idx < s_lut_x_r * s_lut_x_r_s)
+    spa_row = tl.load(s_lut_x + spa_row_idx, mask=spa_row_msk)
+
+    # Load reverse sparsity index for current block
+    rev_idx_spa_idx = (spa_bat * s_l_o_b_s +
+                       spa_row * s_l_o_r_s)
+    rev_idx_spa_msk = (rev_idx_spa_idx < s_l_o_b * s_l_o_b_s)
+    rev_idx_spa = tl.load(r_lut_o + rev_idx_spa_idx, mask=rev_idx_spa_msk).to(tl.int32)
+
+    blk_idx = ((pid_blk * x_b_s) +
+               ((pid_row * TRITON_BLOCK_SIZE + tl.arange(0, TRITON_BLOCK_SIZE)) * x_r_s)[:, None] +
+               ((pid_col * TRITON_BLOCK_SIZE + tl.arange(0, TRITON_BLOCK_SIZE)) * x_c_s)[None, :])
+    blk_msk = (blk_idx < x_b * x_b_s)
+    blk = tl.load(x + blk_idx, mask=blk_msk)
+
+    buf = tl.reshape(tl.max(blk, axis=-1), (TRITON_BLOCK_SIZE, 1))
+
+    o_idx = (rev_idx_spa * o_b_s +
+             ((pid_row * TRITON_BLOCK_SIZE + tl.arange(0, TRITON_BLOCK_SIZE)) * o_r_s)[:, None] +
+             (tl.arange(0, 1))[None, :])
+    o_msk = (o_idx < o_b * o_b_s)
+    tl.atomic_max(o + o_idx, buf, o_msk)
+
+
+def row_wise_add(x: Tensor, sparsity_layout_x: Tensor, y: Tensor,
+                 sparsity_block_size: int, triton_block_size: int = None) -> Tensor:
+    """For each row in ``y`` adds the value to each value in the corresponding row of the block-sparse tensor ``x``.
+
+    Args:
+        x (Tensor): A block-sparse tensor in compressed form.
+        sparsity_layout_x (Tensor): The sparsity layout of the block-sparse tensor.
+        y (Tensor): A block-sparse tensor in compressed form with only one value per row and a single column of sparse blocks.
+        sparsity_block_size (int): The size of the sparsity blocks.
+        triton_block_size (int): The block size to use for the triton kernel (default ``None``).
+
+    Returns:
+        Tensor: The values of ``x`` with the first value of ``y`` in each row added to them as a block-sparse tensor in
+            compressed form.
+
+    """
+    validate_dimensions(x)
+    validate_contiguous(x)
+    validate_device(x)
+    validate_sparsity(sparsity_block_size, (x, sparsity_layout_x))
+    validate_sparsity_block_size(sparsity_block_size, x)
+    validate_triton_block_size(triton_block_size, sparsity_block_size)
+
+    sparsity_lut = torch.nonzero(sparsity_layout_x).contiguous()
+
+    sparsity_layout_rwm, _ = torch.max(sparsity_layout_x, dim=-1, keepdim=True)
+    sparsity_layout_rwm_flat = sparsity_layout_rwm.reshape(-1)
+    sparsity_reverse_lut_rwm = ((torch.cumsum(sparsity_layout_rwm_flat, dim=-1) - 1) *
+                                (sparsity_layout_rwm_flat == 1) -
+                                (1 * (sparsity_layout_rwm_flat == 0)))
+
+    validate_contiguous(sparsity_layout_x, sparsity_lut, sparsity_reverse_lut_rwm)
+
+    output = torch.empty_like(x)
+
+    x_b, x_r, x_c = x.size()
+    x_b_s, x_r_s, x_c_s = x.stride()
+    s_lut_r, s_lut_c = sparsity_lut.size()
+    s_lut_r_s, s_lut_c_s = sparsity_lut.stride()
+    y_b, y_r, y_c = y.size()
+    y_b_s, y_r_s, y_c_s = y.stride()
+    s_l_y_b, s_l_y_r, s_l_y_c = sparsity_layout_rwm.size()
+    s_l_y_b_s, s_l_y_r_s, s_l_y_c_s = sparsity_layout_rwm.stride()
+    o_b, o_r, o_c = output.size()
+    o_b_s, o_r_s, o_c_s = output.stride()
+
+    if triton_block_size is None:
+        triton_block_size = get_triton_block_size(sparsity_block_size)
+
+    triton_grid = lambda meta: [o_b,
+                                triton.cdiv(o_r, meta["TRITON_BLOCK_SIZE"]),
+                                triton.cdiv(o_c, meta["TRITON_BLOCK_SIZE"])]
+
+    (kernel_blocksparse_row_wise_add[triton_grid]
+     (x,
+      x_b, x_b_s, x_r_s, x_c_s,
+      sparsity_lut, s_lut_r, s_lut_r_s, s_lut_c_s,
+      y, y_b, y_b_s, y_r_s, y_c_s,
+      s_l_y_b, s_l_y_b_s, s_l_y_r_s,
+      sparsity_reverse_lut_rwm,
+      output,
+      o_b, o_b_s, o_r_s, o_c_s,
+      triton_block_size
+      ))
+
+    return output
+
+
+def row_wise_sub(x: Tensor, sparsity_layout_x: Tensor, y: Tensor,
+                 sparsity_block_size: int, triton_block_size: int = None) -> Tensor:
+    """Wrapper for ``row_wise_add`` with negated y.
+
+    """
+    return row_wise_add(x, sparsity_layout_x, torch.neg(y), sparsity_block_size, triton_block_size)
+
+
+@triton.jit
+def kernel_blocksparse_row_wise_add(x,
+                                    x_b, x_b_s, x_r_s, x_c_s,
+                                    s_lut, s_lut_r, s_lut_r_s, s_lut_c_s,
+                                    y, y_b, y_b_s, y_r_s, y_c_s,
+                                    s_l_y_b, s_l_y_b_s, s_l_y_r_s,
+                                    r_lut_y,
+                                    o,
+                                    o_b, o_b_s, o_r_s, o_c_s,
+                                    TRITON_BLOCK_SIZE: tl.constexpr) -> None:
+    # Get triton block indices
+    pid_blk = tl.program_id(axis=0)
+    pid_row = tl.program_id(axis=1)
+    pid_col = tl.program_id(axis=2)
+
+    # Get position of current sparsity block consisting of its batch and row index
+    spa_bat_idx = (pid_blk * s_lut_r_s + 0 * s_lut_c_s)
+    spa_bat_msk = (spa_bat_idx < s_lut_r * s_lut_r_s)
+    spa_bat = tl.load(s_lut + spa_bat_idx, mask=spa_bat_msk)
+
+    spa_row_idx = (pid_blk * s_lut_r_s + 1 * s_lut_c_s)
+    spa_row_msk = (spa_row_idx < s_lut_r * s_lut_r_s)
+    spa_row = tl.load(s_lut + spa_row_idx, mask=spa_row_msk)
+
+    # Get reverse sparsity indices for s
+    rev_idx_spa_s_idx = (spa_bat * s_l_y_b_s +
+                         spa_row * s_l_y_r_s)
+    rev_idx_spa_s_msk = (rev_idx_spa_s_idx < s_l_y_b * s_l_y_b_s)
+    rev_idx_spa_s = tl.load(r_lut_y + rev_idx_spa_s_idx, mask=rev_idx_spa_s_msk).to(tl.int32)
+
+    if rev_idx_spa_s == -1:
+        assert False, "Invalid sparsity block"
+
+    # Load x block
+    blk_x_idx = ((pid_blk * x_b_s) +
+                 ((pid_row * TRITON_BLOCK_SIZE + tl.arange(0, TRITON_BLOCK_SIZE)) * x_r_s)[:, None] +
+                 ((pid_col * TRITON_BLOCK_SIZE + tl.arange(0, TRITON_BLOCK_SIZE)) * x_c_s)[None, :])
+    blk_x_msk = (blk_x_idx < x_b * x_b_s)
+    blk_x = tl.load(x + blk_x_idx, mask=blk_x_msk)
+
+    # Load sum block
+    blk_s_idx = (rev_idx_spa_s * y_b_s +
+                 ((pid_row * TRITON_BLOCK_SIZE + tl.arange(0, TRITON_BLOCK_SIZE)) * y_r_s)[:, None] +
+                 (tl.arange(0, 1) * y_c_s)[None, :])
+    blk_s_msk = (blk_s_idx < y_b * y_b_s)
+    blk_s = tl.load(y + blk_s_idx, mask=blk_s_msk)
+
+    # Compute exp
+    buf = blk_x + tl.broadcast_to(blk_s, (TRITON_BLOCK_SIZE, TRITON_BLOCK_SIZE))
+
+    # debug
+    asdf = tl.full((TRITON_BLOCK_SIZE, TRITON_BLOCK_SIZE), 1.0, dtype=tl.float32)
+
+    # Store block
+    blk_o_idx = ((pid_blk * o_b_s) +
+                 ((pid_row * TRITON_BLOCK_SIZE + tl.arange(0, TRITON_BLOCK_SIZE)) * o_r_s)[:, None] +
+                 ((pid_col * TRITON_BLOCK_SIZE + tl.arange(0, TRITON_BLOCK_SIZE)) * o_c_s)[None, :])
+    blk_o_msk = (blk_o_idx < o_b * o_b_s)
+    tl.store(o + blk_o_idx, buf, mask=blk_o_msk)

{blksprs-1.2.1 → blksprs-1.4}/blksprs/ops/softmax.py

@@ -4,7 +4,7 @@ from torch import Tensor
 from triton import language as tl
 
 from blksprs.ops.exp import exp
-from blksprs.ops.row_wise_sum import row_wise_sum
+from blksprs.misc.row_wise import row_wise_sum, row_wise_max, row_wise_sub
 from blksprs.utils.tools import get_triton_block_size
 from blksprs.utils.validation import validate_contiguous, validate_dimensions, validate_device, \
     validate_sparsity, validate_sparsity_block_size, validate_triton_block_size
@@ -33,12 +33,6 @@ def softmax(x: Tensor, sparsity_layout: Tensor, sparsity_block_size: int, triton
     validate_sparsity_block_size(sparsity_block_size, x)
     validate_triton_block_size(triton_block_size, sparsity_block_size)
 
-    if x.size(0) != 0:
-        max_val = torch.max(x).item()
-    else:
-        max_val = 0
-    x_scaled = x - max_val
-
     sparsity_lut = torch.nonzero(sparsity_layout).contiguous()
 
     sparsity_layout_rws, _ = torch.max(sparsity_layout, dim=-1, keepdim=True)
@@ -49,7 +43,7 @@ def softmax(x: Tensor, sparsity_layout: Tensor, sparsity_block_size: int, triton
 
     validate_contiguous(sparsity_layout, sparsity_lut, sparsity_reverse_lut_rws)
 
-    return _BlocksparseSoftmax.apply(x_scaled, sparsity_layout,
+    return _BlocksparseSoftmax.apply(x, sparsity_layout,
                                      sparsity_lut,
                                      sparsity_reverse_lut_rws,
                                      sparsity_block_size, triton_block_size)
@@ -64,13 +58,17 @@ class _BlocksparseSoftmax(torch.autograd.Function):
                 sparsity_block_size: int, triton_block_size: int) -> Tensor:
         output = torch.empty_like(x)
 
-        x_b, x_r, x_c = x.shape
+        x_b, x_r, x_c = x.size()
         x_b_s, x_r_s, x_c_s = x.stride()
-        s_lut_r, s_lut_c = sparsity_lut.shape
+        s_lut_r, s_lut_c = sparsity_lut.size()
         s_lut_r_s, s_lut_c_s = sparsity_lut.stride()
-        o_b, o_r, o_c = output.shape
+        o_b, o_r, o_c = output.size()
 
-        x_exp = exp(x, sparsity_block_size, triton_block_size=triton_block_size)
+        x_row_wise_max, sparsity_layout_rwm = row_wise_max(x, sparsity_layout, sparsity_block_size,
+                                                           flag_slice_only=True,
+                                                           triton_block_size=triton_block_size)
+        x_scaled = row_wise_sub(x, sparsity_layout, x_row_wise_max, sparsity_block_size, triton_block_size)
+        x_exp = exp(x_scaled, sparsity_block_size, triton_block_size=triton_block_size)
         x_exp_row_wise_sum, sparsity_layout_rws = row_wise_sum(x_exp, sparsity_layout, sparsity_block_size,
                                                                flag_slice_only=True,
                                                                triton_block_size=triton_block_size)
@@ -174,7 +172,7 @@ class _BlocksparseSoftmax(torch.autograd.Function):
         spa_row_msk = (spa_row_idx < s_lut_r * s_lut_r_s)
         spa_row = tl.load(s_lut + spa_row_idx, mask=spa_row_msk)
 
-        # Get reverse sparsity indices for x
+        # Get reverse sparsity indices for s
         rev_idx_spa_s_idx = (spa_bat * s_l_s_b_s +
                              spa_row * s_l_s_r_s)
         rev_idx_spa_s_msk = (rev_idx_spa_s_idx < s_l_s_b * s_l_s_b_s)

{blksprs-1.2.1 → blksprs-1.4}/blksprs/ops/transpose.py

@@ -129,7 +129,7 @@ class _BlocksparseTranspose(torch.autograd.Function):
         spa_col_msk = (spa_col_idx < s_lut_r * s_lut_r_s)
         spa_col = tl.load(s_lut + spa_col_idx, mask=spa_col_msk)
 
-        # Get reverse sparsity indices
+        # Get reverse sparsity index
         rev_idx_spa_idx = (spa_bat * s_l_b_s +
                            spa_row * s_l_r_s +
                            spa_col * s_l_c_s)

{blksprs-1.2.1 → blksprs-1.4}/blksprs/utils/tools.py

@@ -10,7 +10,7 @@ def do_shape_blocksparse(x: Tensor):
 
 
 def undo_shape_blocksparse(x: Tensor, shape: Size):
-    if x.shape[-2:] == shape[-2:]:
+    if x.shape[:-2] == shape[:-2]:
         return x
 
     return x.reshape((*shape[:-2], *x.shape[-2:]))

{blksprs-1.2.1 → blksprs-1.4}/blksprs.egg-info/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: blksprs
-Version: 1.2.1
+Version: 1.4
 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
@@ -83,14 +83,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 +105,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)
+    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, 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.2.1 → blksprs-1.4}/blksprs.egg-info/SOURCES.txt

@@ -1,5 +1,6 @@
 README.md
 pyproject.toml
+blksprs/__init__.py
 blksprs.egg-info/PKG-INFO
 blksprs.egg-info/SOURCES.txt
 blksprs.egg-info/dependency_links.txt
@@ -7,12 +8,13 @@ blksprs.egg-info/requires.txt
 blksprs.egg-info/top_level.txt
 blksprs/layouting/distribution_layout.py
 blksprs/layouting/sparsity_layout.py
-blksprs/misc/broadcast_addition.py
+blksprs/misc/broadcast_ops.py
+blksprs/misc/repeat_interleave.py
+blksprs/misc/row_wise.py
 blksprs/ops/conversion.py
 blksprs/ops/distribution.py
 blksprs/ops/exp.py
 blksprs/ops/matmul.py
-blksprs/ops/row_wise_sum.py
 blksprs/ops/softmax.py
 blksprs/ops/transpose.py
 blksprs/utils/benchmarking.py

{blksprs-1.2.1 → blksprs-1.4}/pyproject.toml

@@ -1,6 +1,6 @@
 [project]
 name = "blksprs"
-version = "1.2.1"
+version = "1.4"
 authors = [{ name = "Felix Schön", email = "schoen@kr.tuwien.ac.at" }]
 description = "A lightweight library for operations on blocksparse matrices in PyTorch."
 readme = "README.md"

blksprs-1.2.1/blksprs/ops/row_wise_sum.py

@@ -1,231 +0,0 @@
-import torch
-import triton
-from torch import Tensor
-from triton import language as tl
-
-from blksprs.utils.tools import get_triton_block_size
-from blksprs.utils.validation import validate_contiguous, validate_dimensions, validate_device, \
-    validate_sparsity, validate_sparsity_block_size, validate_triton_block_size
-
-
-def row_wise_sum(x: Tensor, sparsity_layout: Tensor, sparsity_block_size: int,
-                 flag_slice_only: bool = False, triton_block_size: int = None) -> tuple[Tensor, Tensor]:
-    """Computes the row-wise sum of a block-sparse tensor.
-
-    Returns a block-sparse tensor in compressed form with only one block per row, where the first entry contains the sum
-        of the corresponding row.
-
-    Note:
-        If ``flag_slice_only`` is set the output will be of shape ``[x.size(0), x.size(1), 1]``.
-
-    Args:
-        x (Tensor): A block-sparse tensor in compressed form.
-        sparsity_layout (Tensor): The sparsity layout of the block-sparse tensor.
-        sparsity_block_size (int): The size of the sparsity blocks.
-        flag_slice_only (bool, optional): If set the output will be of shape ``[x.size(0), x.size(1), 1]``
-            (default ``False``).
-        triton_block_size (int): The block size to use for the triton kernel (default ``None``).
-
-    Returns:
-        tuple[Tensor, Tensor]: A tuple containing a block-sparse tensor in compressed form containing the row-wise sum
-            of the input and the sparsity layout of the output tensor.
-
-    """
-    validate_dimensions(x)
-    validate_contiguous(x)
-    validate_device(x)
-    validate_sparsity(sparsity_block_size, (x, sparsity_layout))
-    validate_sparsity_block_size(sparsity_block_size, x)
-    validate_triton_block_size(triton_block_size, sparsity_block_size)
-
-    sparsity_lut = torch.nonzero(sparsity_layout).contiguous()
-    sparsity_layout_flat = sparsity_layout.reshape(-1)
-    sparsity_reverse_lut = ((torch.cumsum(sparsity_layout_flat, dim=-1) - 1) *
-                            (sparsity_layout_flat == 1) -
-                            (1 * (sparsity_layout_flat == 0)))
-
-    sparsity_layout_output, _ = torch.max(sparsity_layout, dim=-1, keepdim=True)
-    sparsity_lut_output = torch.nonzero(sparsity_layout_output).contiguous()
-    sparsity_layout_output_flat = sparsity_layout_output.reshape(-1)
-    sparsity_reverse_lut_output = ((torch.cumsum(sparsity_layout_output_flat, dim=-1) - 1) *
-                                   (sparsity_layout_output_flat == 1) -
-                                   (1 * (sparsity_layout_output_flat == 0)))
-
-    n_sparse_blocks_output = torch.sum(sparsity_layout_output.to(torch.int)).item()
-
-    validate_contiguous(sparsity_layout, sparsity_lut, sparsity_reverse_lut,
-                        sparsity_layout_output, sparsity_lut_output, sparsity_reverse_lut_output)
-
-    return (_BlocksparseRowWiseSum.apply(x,
-                                         sparsity_layout, sparsity_lut, sparsity_reverse_lut,
-                                         sparsity_layout_output, sparsity_lut_output, sparsity_reverse_lut_output,
-                                         n_sparse_blocks_output,
-                                         flag_slice_only,
-                                         sparsity_block_size, triton_block_size),
-            sparsity_layout_output)
-
-
-class _BlocksparseRowWiseSum(torch.autograd.Function):
-    IMPLEMENTATION = "atomic_add"
-
-    @staticmethod
-    def forward(ctx, x: Tensor,
-                sparsity_layout: Tensor, sparsity_lut: Tensor, sparsity_reverse_lut: Tensor,
-                sparsity_layout_output: Tensor, sparsity_lut_output: Tensor, sparsity_reverse_lut_output: Tensor,
-                n_sparse_blocks_output: int,
-                flag_slice_only: bool,
-                sparsity_block_size: int, triton_block_size: int) -> Tensor:
-        output = torch.zeros(size=(n_sparse_blocks_output,
-                                   sparsity_block_size,
-                                   1 if flag_slice_only else sparsity_block_size),
-                             device=x.device)
-
-        x_b, x_r, x_c = x.size()
-        x_b_s, x_r_s, x_c_s = x.stride()
-        s_l_x_b, s_l_x_r, s_l_x_c = sparsity_layout.size()
-        s_l_x_b_s, s_l_x_r_s, s_l_x_c_s = sparsity_layout.stride()
-        s_lut_x_r, s_lut_x_c = sparsity_lut.size()
-        s_lut_x_r_s, s_lut_x_c_s = sparsity_lut.stride()
-        o_b, o_r, o_c = output.size()
-        o_b_s, o_r_s, o_c_s = output.stride()
-        s_l_o_b, s_l_o_r, s_l_o_c = sparsity_layout_output.size()
-        s_l_o_b_s, s_l_o_r_s, s_l_o_c_s = sparsity_layout_output.stride()
-        s_lut_o_r, s_lut_o_c = sparsity_lut_output.size()
-        s_lut_o_r_s, s_lut_o_c_s = sparsity_lut_output.stride()
-
-        if triton_block_size is None:
-            triton_block_size = get_triton_block_size(sparsity_block_size)
-
-        if _BlocksparseRowWiseSum.IMPLEMENTATION == "basic":
-            triton_grid = lambda meta: [o_b,
-                                        triton.cdiv(o_r, meta["TRITON_BLOCK_SIZE"])]
-
-            (_BlocksparseRowWiseSum.kernel_blocksparse_row_wise_sum[triton_grid]
-             (x,
-              x_b, x_b_s, x_r_s, x_c_s,
-              s_l_x_b, s_l_x_b_s, s_l_x_r_s, s_l_x_c, s_l_x_c_s,
-              sparsity_reverse_lut,
-              output,
-              o_b, o_b_s, o_r_s,
-              sparsity_lut_output, s_lut_o_r, s_lut_o_r_s, s_lut_o_c_s,
-              sparsity_block_size,
-              triton_block_size))
-        elif _BlocksparseRowWiseSum.IMPLEMENTATION == "atomic_add":
-            triton_grid = lambda meta: [x_b,
-                                        triton.cdiv(x_r, meta["TRITON_BLOCK_SIZE"]),
-                                        triton.cdiv(x_c, meta["TRITON_BLOCK_SIZE"])]
-
-            (_BlocksparseRowWiseSum.kernel_blocksparse_row_wise_sum_atomic_add[triton_grid]
-             (x,
-              x_b, x_b_s, x_r_s, x_c_s,
-              sparsity_lut, s_lut_x_r, s_lut_x_r_s, s_lut_x_c_s,
-              output,
-              o_b, o_b_s, o_r_s,
-              s_l_o_b, s_l_o_b_s, s_l_o_r_s,
-              sparsity_reverse_lut_output,
-              triton_block_size))
-
-        return output
-
-    @staticmethod
-    def backward(ctx, grad_output):
-        raise NotImplementedError
-
-    @staticmethod
-    @triton.jit
-    def kernel_blocksparse_row_wise_sum(x,
-                                        x_b, x_b_s, x_r_s, x_c_s,
-                                        s_l_x_b, s_l_x_b_s, s_l_x_r_s, s_l_x_c, s_l_x_c_s,
-                                        r_lut_x,
-                                        o,
-                                        o_b, o_b_s, o_r_s,
-                                        s_lut_o, s_lut_o_r, s_lut_o_r_s, s_lut_o_c_s,
-                                        sparsity_block_size,
-                                        TRITON_BLOCK_SIZE: tl.constexpr) -> None:
-        pid_blk = tl.program_id(axis=0)
-        pid_row = tl.program_id(axis=1)
-
-        # Get position of current sparsity block consisting of its batch and row index
-        spa_bat_idx = (pid_blk * s_lut_o_r_s + 0 * s_lut_o_c_s)
-        spa_bat_msk = (spa_bat_idx < s_lut_o_r * s_lut_o_r_s)
-        spa_bat = tl.load(s_lut_o + spa_bat_idx, mask=spa_bat_msk)
-
-        spa_row_idx = (pid_blk * s_lut_o_r_s + 1 * s_lut_o_c_s)
-        spa_row_msk = (spa_row_idx < s_lut_o_r * s_lut_o_r_s)
-        spa_row = tl.load(s_lut_o + spa_row_idx, mask=spa_row_msk)
-
-        buf = tl.zeros(shape=(TRITON_BLOCK_SIZE, 1), dtype=tl.float32)
-
-        # Slide over triton block sized segments of input tensor
-        for i_seg_tri in range(0, tl.cdiv(s_l_x_c * sparsity_block_size, TRITON_BLOCK_SIZE)):
-            # Convert to segment index of sparsity layout
-            i_seg_spa = (i_seg_tri * TRITON_BLOCK_SIZE) // sparsity_block_size
-            # Calculate the triton segment index within a block
-            i_seg_tri_mod = i_seg_tri % (sparsity_block_size // TRITON_BLOCK_SIZE)
-
-            # Load reverse sparsity index for current block
-            rev_idx_spa_idx = (spa_bat * s_l_x_b_s +
-                               spa_row * s_l_x_r_s +
-                               i_seg_spa * s_l_x_c_s)
-            rev_idx_spa_msk = (rev_idx_spa_idx < s_l_x_b * s_l_x_b_s)
-            rev_idx_spa = tl.load(r_lut_x + rev_idx_spa_idx, mask=rev_idx_spa_msk).to(tl.int32)
-
-            # If block is present commence operations
-            if rev_idx_spa >= 0:
-                blk_idx = ((rev_idx_spa * x_b_s) +
-                           ((pid_row * TRITON_BLOCK_SIZE + tl.arange(0, TRITON_BLOCK_SIZE)) * x_r_s)[:, None] +
-                           ((i_seg_tri_mod * TRITON_BLOCK_SIZE +
-                             tl.arange(0, TRITON_BLOCK_SIZE)) * x_c_s)[None, :])
-                blk_msk = (blk_idx < x_b * x_b_s)
-                blk = tl.load(x + blk_idx, mask=blk_msk)
-
-                buf = buf + tl.reshape(tl.sum(blk, axis=-1), (TRITON_BLOCK_SIZE, 1))
-
-        o_idx = (pid_blk * o_b_s +
-                 ((pid_row * TRITON_BLOCK_SIZE + tl.arange(0, TRITON_BLOCK_SIZE)) * o_r_s)[:, None] +
-                 (tl.arange(0, 1))[None, :])
-        o_msk = (o_idx < o_b * o_b_s)
-        tl.store(o + o_idx, buf, o_msk)
-
-    @staticmethod
-    @triton.jit
-    def kernel_blocksparse_row_wise_sum_atomic_add(x,
-                                                   x_b, x_b_s, x_r_s, x_c_s,
-                                                   s_lut_x, s_lut_x_r, s_lut_x_r_s, s_lut_x_c_s,
-                                                   o,
-                                                   o_b, o_b_s, o_r_s,
-                                                   s_l_o_b, s_l_o_b_s, s_l_o_r_s,
-                                                   r_lut_o,
-                                                   TRITON_BLOCK_SIZE: tl.constexpr) -> None:
-        pid_blk = tl.program_id(axis=0)
-        pid_row = tl.program_id(axis=1)
-        pid_col = tl.program_id(axis=2)
-
-        # Get position of current sparsity block consisting of its batch and row index
-        spa_bat_idx = (pid_blk * s_lut_x_r_s + 0 * s_lut_x_c_s)
-        spa_bat_msk = (spa_bat_idx < s_lut_x_r * s_lut_x_r_s)
-        spa_bat = tl.load(s_lut_x + spa_bat_idx, mask=spa_bat_msk)
-
-        spa_row_idx = (pid_blk * s_lut_x_r_s + 1 * s_lut_x_c_s)
-        spa_row_msk = (spa_row_idx < s_lut_x_r * s_lut_x_r_s)
-        spa_row = tl.load(s_lut_x + spa_row_idx, mask=spa_row_msk)
-
-        # Load reverse sparsity index for current block
-        rev_idx_spa_idx = (spa_bat * s_l_o_b_s +
-                           spa_row * s_l_o_r_s)
-        rev_idx_spa_msk = (rev_idx_spa_idx < s_l_o_b * s_l_o_b_s)
-        rev_idx_spa = tl.load(r_lut_o + rev_idx_spa_idx, mask=rev_idx_spa_msk).to(tl.int32)
-
-        blk_idx = ((pid_blk * x_b_s) +
-                   ((pid_row * TRITON_BLOCK_SIZE + tl.arange(0, TRITON_BLOCK_SIZE)) * x_r_s)[:, None] +
-                   ((pid_col * TRITON_BLOCK_SIZE + tl.arange(0, TRITON_BLOCK_SIZE)) * x_c_s)[None, :])
-        blk_msk = (blk_idx < x_b * x_b_s)
-        blk = tl.load(x + blk_idx, mask=blk_msk)
-
-        buf = tl.reshape(tl.sum(blk, axis=-1), (TRITON_BLOCK_SIZE, 1))
-
-        o_idx = (rev_idx_spa * o_b_s +
-                 ((pid_row * TRITON_BLOCK_SIZE + tl.arange(0, TRITON_BLOCK_SIZE)) * o_r_s)[:, None] +
-                 (tl.arange(0, 1))[None, :])
-        o_msk = (o_idx < o_b * o_b_s)
-        tl.atomic_add(o + o_idx, buf, o_msk)