blksprs 1.3__py3-none-any.whl → 1.4.1__py3-none-any.whl

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/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/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/misc/{broadcast_addition.py → 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/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)

blksprs/misc/row_wise.py

@@ -0,0 +1,390 @@
+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.
+
+    """
+    x = x.contiguous()
+
+    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.
+
+    """
+    x = x.contiguous()
+
+    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/ops/conversion.py

@@ -28,6 +28,8 @@ def to_dense(x: Tensor, sparsity_layout: Tensor, sparsity_block_size: int, fill_
         Tensor: The block-sparse tensor converted to regular form.
 
     """
+    x = x.contiguous()
+
     validate_dimensions(x)
     validate_contiguous(x, sparsity_layout)
     validate_device(x)
@@ -156,6 +158,8 @@ def to_sparse(x: Tensor, sparsity_layout: Tensor, sparsity_block_size: int, trit
         Tensor: The block-sparse tensor converted to compressed form.
 
     """
+    x = x.contiguous()
+
     validate_dimensions(x)
     validate_contiguous(x)
     validate_device(x)
@@ -282,6 +286,8 @@ def adapt_layout(x: Tensor, sparsity_layout_from: Tensor, sparsity_block_size_fr
         Tensor: The block-sparse tensor in compressed form with the adapted sparsity layout and sparsity block size.
 
     """
+    x = x.contiguous()
+
     validate_dimensions(x)
     validate_contiguous(x, sparsity_layout_from)
     validate_device(x)

blksprs/ops/distribution.py

@@ -24,6 +24,9 @@ def gather(src: Tensor, sparsity_layout_src: Tensor, idx: Tensor, sparsity_layou
         Tensor: The result of the gather operation as a block-sparse tensor in compressed form.
 
     """
+    src = src.contiguous()
+    idx = idx.contiguous()
+
     validate_dimensions(src, idx)
     validate_contiguous(src, idx)
     validate_dtype_int(idx)
@@ -200,6 +203,9 @@ def scatter_reduce(src: Tensor, sparsity_layout_src: Tensor,
         Tensor: The result of the scatter operation as a block-sparse tensor in compressed form.
 
     """
+    src = src.contiguous()
+    idx = idx.contiguous()
+
     validate_dimensions(src, idx)
     validate_contiguous(src, idx)
     validate_dtype_int(idx)

blksprs/ops/exp.py

@@ -25,6 +25,8 @@ def exp(x: Tensor, sparsity_block_size: int, triton_block_size: int = None) -> T
             compressed form.
 
     """
+    x = x.contiguous()
+
     validate_dimensions(x)
     validate_contiguous(x)
     validate_device(x)

blksprs/ops/matmul.py

@@ -6,7 +6,7 @@ from triton import language as tl
 from blksprs.ops.transpose import transpose
 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
+    validate_sparsity, validate_sparsity_block_size, validate_triton_block_size, validate_dtype_float
 
 
 def matmul(x: Tensor, sparsity_layout_x: Tensor,
@@ -30,8 +30,12 @@ def matmul(x: Tensor, sparsity_layout_x: Tensor,
         Tensor: The result of the matrix multiplication as a block-sparse tensor in compressed form.
 
     """
+    x = x.contiguous()
+    y = y.contiguous()
+
     validate_dimensions(x, y)
     validate_contiguous(x, y)
+    validate_dtype_float(x, y)
     validate_device(x, y)
     validate_sparsity(sparsity_block_size, (x, sparsity_layout_x), (y, sparsity_layout_y))
     if sparsity_layout_x.size(-1) != sparsity_layout_y.size(-2):
@@ -211,7 +215,7 @@ class _BlocksparseMatmulSSS(torch.autograd.Function):
                 blk_y = tl.load(y + blk_y_idx, mask=blk_y_msk)
 
                 # Perform matrix multiplication
-                buf += tl.dot(blk_x, blk_y)
+                buf += tl.dot(blk_x, blk_y, input_precision="tf32")
 
         # Store output
         blk_o_idx = ((pid_blk * o_b_s) +

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
@@ -26,6 +26,8 @@ def softmax(x: Tensor, sparsity_layout: Tensor, sparsity_block_size: int, triton
         Tensor: The result of the softmax operation as a block-sparse tensor in compressed form.
 
     """
+    x = x.contiguous()
+
     validate_dimensions(x)
     validate_contiguous(x)
     validate_device(x)
@@ -33,12 +35,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 +45,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 +60,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 +174,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/ops/transpose.py

@@ -26,6 +26,8 @@ def transpose(x: Tensor, sparsity_layout: Tensor, sparsity_block_size: int, trit
         Tensor: The sparsity layout of the transposed tensor.
 
     """
+    x = x.contiguous()
+
     validate_dimensions(x)
     validate_contiguous(x)
     validate_device(x)

blksprs/utils/tools.py

@@ -1,10 +1,12 @@
 import torch
 from torch import Tensor, Size
 
+from blksprs.utils.validation import _set_skip_validation
+
 
 def do_shape_blocksparse(x: Tensor):
     if x.dim() == 3:
-        return x, x.size()
+        return x.contiguous(), x.size()
 
     return x.reshape(-1, x.size(-2), x.size(-1)), x.size()
 
@@ -18,3 +20,7 @@ def undo_shape_blocksparse(x: Tensor, shape: Size):
 
 def get_triton_block_size(sparsity_block_size: int, limit: int = 128):
     return min(sparsity_block_size, limit)
+
+
+def disable_validation():
+    _set_skip_validation(True)

blksprs/utils/validation.py

@@ -1,9 +1,10 @@
 import torch
 from torch import Tensor
 
+VALIDATION = True
 
 def validate_dimensions(*tensors: Tensor) -> None:
-    if _skip_validation():
+    if _check_skip_validation():
         return
 
     for tensor in tensors:
@@ -12,7 +13,7 @@ def validate_dimensions(*tensors: Tensor) -> None:
 
 
 def validate_contiguous(*tensors: Tensor) -> None:
-    if _skip_validation():
+    if _check_skip_validation():
         return
 
     for tensor in tensors:
@@ -21,7 +22,7 @@ def validate_contiguous(*tensors: Tensor) -> None:
 
 
 def validate_dtype_float(*tensors: Tensor) -> None:
-    if _skip_validation():
+    if _check_skip_validation():
         return
 
     for tensor in tensors:
@@ -30,7 +31,7 @@ def validate_dtype_float(*tensors: Tensor) -> None:
 
 
 def validate_dtype_int(*tensors: Tensor) -> None:
-    if _skip_validation():
+    if _check_skip_validation():
         return
 
     for tensor in tensors:
@@ -39,7 +40,7 @@ def validate_dtype_int(*tensors: Tensor) -> None:
 
 
 def validate_device(*tensors: Tensor) -> None:
-    if _skip_validation():
+    if _check_skip_validation():
         return
 
     device = None
@@ -56,7 +57,7 @@ def validate_device(*tensors: Tensor) -> None:
 
 
 def validate_sparsity(sparsity_block_size: int, *tensor_sparsity_layout_tuples: tuple[Tensor, Tensor]) -> None:
-    if _skip_validation():
+    if _check_skip_validation():
         return
 
     for (tensor, sparsity_layout) in tensor_sparsity_layout_tuples:
@@ -73,7 +74,7 @@ def _validate_sparsity_layout_values(sparsity_layout: Tensor):
         raise ValueError("Sparsity layout values must be either 0 or 1")
 
 def validate_sparsity_block_size(sparsity_block_size: int, *tensors):
-    if _skip_validation():
+    if _check_skip_validation():
         return
 
     if not (sparsity_block_size & (sparsity_block_size - 1)) == 0:
@@ -84,7 +85,7 @@ def validate_sparsity_block_size(sparsity_block_size: int, *tensors):
             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():
+    if _check_skip_validation():
         return
 
     if triton_block_size is None:
@@ -93,5 +94,9 @@ def validate_triton_block_size(triton_block_size: int, sparsity_block_size: int)
     if triton_block_size > sparsity_block_size:
         raise ValueError("Triton block size cannot be larger than sparsity block size")
 
-def _skip_validation():
-    return False
+def _check_skip_validation():
+    return not VALIDATION
+
+def _set_skip_validation(skip_validation: bool):
+    global VALIDATION
+    VALIDATION = not skip_validation

{blksprs-1.3.dist-info → blksprs-1.4.1.dist-info}/METADATA

@@ -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
@@ -8,10 +8,8 @@ Project-URL: Bugtracker, https://github.com/FelixSchoen/blksprs/issues
 Requires-Python: >=3.11
 Description-Content-Type: text/markdown
 Requires-Dist: torch
-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"
 Provides-Extra: test
 Requires-Dist: pytest; extra == "test"
 Requires-Dist: pytest-xdist; extra == "test"
@@ -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.4.1.dist-info/RECORD

@@ -0,0 +1,19 @@
+blksprs/__init__.py,sha256=ORAVhGR91G1wyIOs9Wecv-xfmjju3bJ4Jynq_SGOVY4,833
+blksprs/layouting/distribution_layout.py,sha256=Xd8KjZwI87L9EL1Bw5SGUW9YztFD5q0Ygr99sffvdak,4939
+blksprs/layouting/sparsity_layout.py,sha256=vZL8r5LkMwILYYqTYPZcN_NYFJuVFIB6mmBkdtRyXmI,7893
+blksprs/misc/broadcast_ops.py,sha256=RTcqvx6X_THRBb55jipeEe63YSLIAh27jdpuze0aSek,5308
+blksprs/misc/repeat_interleave.py,sha256=KJeapmxbpA7zGFfa5hUhCGrk4aFmhOhlMw-hbTh9PLI,5668
+blksprs/misc/row_wise.py,sha256=KCDO5ry5TkjI88LLD_QINZwBkzfmjoQpOOvYLfpUn5I,16853
+blksprs/ops/conversion.py,sha256=h1c5T74rQjqYgY9dwWXfPTXRpgzy0dtAhCmtUp8-6uo,21332
+blksprs/ops/distribution.py,sha256=KhtHRVcv4_woyNlldAjIWF-7021-KX-xyIcN6rE-UgE,16879
+blksprs/ops/exp.py,sha256=CVWVq_emO2CnS_xk6Unx67P7EI7IL26dwtsmBJZOLzQ,3698
+blksprs/ops/matmul.py,sha256=6DaYxecJgwiW8L-UISkgyNyzQ31AAkmDL-Oq1EjHt98,11210
+blksprs/ops/softmax.py,sha256=cSTxDnNmMRlJGOlCSpdg1U5KUIFpVtHulz8fteJFeh0,11972
+blksprs/ops/transpose.py,sha256=et8R124L29TUqihci18ms_hBoYXTtPu5LXgEA8sxk_w,6744
+blksprs/utils/benchmarking.py,sha256=4pLVlnPW_2EM-NT3n4SClaRznVYEljztLbJcccz8kZE,1360
+blksprs/utils/tools.py,sha256=RKGWCGd5h1qFOIoShsdJObx4-QsS0RxCyzFie0geNxo,596
+blksprs/utils/validation.py,sha256=Gsx3aah6355bWXRPpbFuZ1p0fOrYduIqaM3ON9d5NiI,3197
+blksprs-1.4.1.dist-info/METADATA,sha256=3xRmBFHv2U2KnrW3_QX3003SHLkQ1JCaSqh4AUBsJD4,7609
+blksprs-1.4.1.dist-info/WHEEL,sha256=GV9aMThwP_4oNCtvEC2ec3qUYutgWeAzklro_0m4WJQ,91
+blksprs-1.4.1.dist-info/top_level.txt,sha256=qyp0IHeY3H2GQA97i4hk_To5rRBS2YcE1HRPSLy04fk,8
+blksprs-1.4.1.dist-info/RECORD,,

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)

blksprs-1.3.dist-info/RECORD

@@ -1,18 +0,0 @@
-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/misc/repeat_interleave.py,sha256=WrIp7uJsnvjIhFeLYPfkL2j5vXyKmDQGrJ69b3Y0lQ8,5644
-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=cX_E3b-QMhsUDNn9D8HVkYesc2JBc-EcVBUZfCWExM8,6720
-blksprs/utils/benchmarking.py,sha256=4pLVlnPW_2EM-NT3n4SClaRznVYEljztLbJcccz8kZE,1360
-blksprs/utils/tools.py,sha256=DwophH01AeNTZAo0B1uWbKFSGBQjI5z0WmFnYKh-BBk,465
-blksprs/utils/validation.py,sha256=gJYZO5C48YUrXV3Fy_Z_lCaOpiFj951FT-Od7sKfprg,3007
-blksprs-1.3.dist-info/METADATA,sha256=bs4_e4DjSYyAQ354tLVNIKcGLkww_-C2AfHnJIMdjA8,7515
-blksprs-1.3.dist-info/WHEEL,sha256=GV9aMThwP_4oNCtvEC2ec3qUYutgWeAzklro_0m4WJQ,91
-blksprs-1.3.dist-info/top_level.txt,sha256=qyp0IHeY3H2GQA97i4hk_To5rRBS2YcE1HRPSLy04fk,8
-blksprs-1.3.dist-info/RECORD,,