blksprs 1.10.2__py3-none-any.whl → 2.0rc1__py3-none-any.whl

blksprs/ops/transpose.py

@@ -1,16 +1,15 @@
 import torch
-import triton
 from torch import Tensor
-from triton import language as tl
+from torch._library import triton_op
 
+from blksprs.ops.flow import flow_pull_forward
 from blksprs.utils.blksprs_tensor import BlksprsTensor
-from blksprs.utils.tools import get_triton_block_size, stride
 from blksprs.utils.validation import validate_dimensions, validate_contiguous, validate_device, \
-    validate_sparsity, validate_sparsity_block_size, validate_triton_block_size
+    validate_sparsity, validate_sparsity_block_size
 
 
-def transpose(x: BlksprsTensor, sparsity_layout: Tensor, sparsity_block_size: int, triton_block_size: int = None) -> (
-        BlksprsTensor, Tensor):
+def transpose(x: BlksprsTensor, sparsity_layout: Tensor,
+              sparsity_block_size: int, lut: dict = None) -> (BlksprsTensor, Tensor):
     """Transposes a block-sparse tensor in compressed form.
 
     Note:
@@ -20,7 +19,7 @@ def transpose(x: BlksprsTensor, sparsity_layout: Tensor, sparsity_block_size: in
         x (BlksprsTensor): 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.
-        triton_block_size (int): The block size to use for the triton kernel (default ``None``).
+        lut (dict, optional): A dictionary containing the look-up tables for the operation (default ``None``).
 
     Returns:
         BlksprsTensor: The transposed block-sparse tensor in compressed form.
@@ -28,133 +27,72 @@ def transpose(x: BlksprsTensor, sparsity_layout: Tensor, sparsity_block_size: in
 
     """
     x = x.contiguous()
+    x_t = x.transpose(-1, -2).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_layout_t = sparsity_layout.transpose(-1, -2).contiguous()
-
-    sparsity_lut = torch.nonzero(sparsity_layout_t).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)))
-                            .reshape(sparsity_layout.size()).transpose(-1, -2).contiguous().reshape(-1))
-
-    n_sparse_blocks = torch.sum(sparsity_layout.to(torch.int)).item()
-
-    validate_contiguous(sparsity_layout_t, sparsity_lut, sparsity_reverse_lut)
-
-    return BlksprsTensor(
-        _BlocksparseTranspose.apply(x, sparsity_layout_t, sparsity_lut, sparsity_reverse_lut, sparsity_block_size,
-                                    n_sparse_blocks, triton_block_size)), sparsity_layout_t
-
-
-class _BlocksparseTranspose(torch.autograd.Function):
-
-    @staticmethod
-    def forward(ctx, x: Tensor, sparsity_layout_o: Tensor, sparsity_lut: Tensor, sparsity_reverse_lut: Tensor,
-                sparsity_block_size: int,
-                n_sparse_blocks: int, triton_block_size: int) -> Tensor:
-        output = torch.empty(size=(n_sparse_blocks, 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 = stride(x)
-        s_l_b, s_l_r, s_l_c = sparsity_layout_o.size()
-        s_l_b_s, s_l_r_s, s_l_c_s = stride(sparsity_layout_o)
-        s_lut_r, s_lut_c = sparsity_lut.shape
-        s_lut_r_s, s_lut_c_s = stride(sparsity_lut)
-        o_b, o_r, o_c = output.size()
-        o_b_s, o_r_s, o_c_s = stride(output)
-
-        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"])]
-
-        (_BlocksparseTranspose.kernel_blocksparse_transpose[triton_grid]
-         (x,
-          x_b, x_b_s, x_r_s, x_c_s,
-          s_l_b, s_l_b_s, s_l_r_s, s_l_c_s,
-          sparsity_lut, s_lut_r, s_lut_r_s, s_lut_c_s,
-          sparsity_reverse_lut,
-          output,
-          o_b, o_b_s,
-          triton_block_size))
-
-        # Save for backward pass
-        ctx.save_for_backward(sparsity_layout_o)
-        ctx.sparsity_block_size = sparsity_block_size
-        ctx.triton_block_size = triton_block_size
-
-        return output
-
-    @staticmethod
-    def backward(ctx, grad_output):
-        sparsity_layout = ctx.saved_tensors[0]
-        sparsity_block_size = ctx.sparsity_block_size
-        triton_block_size = ctx.triton_block_size
-
-        return transpose(grad_output, sparsity_layout, sparsity_block_size, triton_block_size)[
-            0], None, None, None, None, None, None
-
-    @staticmethod
-    @triton.jit
-    def kernel_blocksparse_transpose(x,
-                                     x_b, x_b_s, x_r_s, x_c_s,
-                                     s_l_b, s_l_b_s, s_l_r_s, s_l_c_s,
-                                     s_lut, s_lut_r, s_lut_r_s, s_lut_c_s,
-                                     r_lut,
-                                     o,
-                                     o_b, o_b_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 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 >= 0 and 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 >= 0 and 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 >= 0 and 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 index
-        rev_idx_spa_idx = (spa_bat * s_l_b_s +
-                           spa_row * s_l_r_s +
-                           spa_col * s_l_c_s)
-        rev_idx_spa_msk = (rev_idx_spa_idx >= 0 and rev_idx_spa_idx < s_l_b * s_l_b_s)
-        rev_idx_spa = tl.load(r_lut + rev_idx_spa_idx, mask=rev_idx_spa_msk).to(tl.int32)
-
-        if rev_idx_spa == -1:
-            tl.device_assert(False)
-            return
-
-        blk_x_idx = (rev_idx_spa * 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 >= 0 and blk_x_idx < x_b * x_b_s)
-        blk_x = tl.load(x + blk_x_idx, mask=blk_x_msk)
-
-        blk_x_t = tl.trans(blk_x)
-
-        blk_o_idx = (pid_blk * o_b_s +
-                     ((pid_col * TRITON_BLOCK_SIZE + tl.arange(0, TRITON_BLOCK_SIZE)) * x_r_s)[:, None] +
-                     ((pid_row * TRITON_BLOCK_SIZE + tl.arange(0, TRITON_BLOCK_SIZE)) * x_c_s)[None, :])
-        blk_o_msk = (blk_o_idx >= 0 and blk_o_idx < o_b * o_b_s)
-        tl.store(o + blk_o_idx, blk_x_t, mask=blk_o_msk)
+
+    lut = transpose_build_lut(lut, sparsity_layout)
+
+    return BlksprsTensor(transpose_forward(x_t, lut["sparsity_layout_t"],
+                                           lut["sparsity_lut"], lut["sparsity_reverse_lut"],
+                                           sparsity_block_size, lut["n_sparse_blocks"])), lut["sparsity_layout_t"]
+
+
+@triton_op("blksprs::transpose", mutates_args={})
+def transpose_forward(x: Tensor, sparsity_layout_o: Tensor,
+                      sparsity_lut: Tensor, sparsity_reverse_lut: Tensor,
+                      sparsity_block_size: int, n_sparse_blocks: int) -> Tensor:
+    return flow_pull_forward(x, sparsity_layout_o, sparsity_lut, sparsity_reverse_lut,
+                             sparsity_block_size, n_sparse_blocks)
+
+
+def transpose_backward(ctx, grad_output):
+    sparsity_layout = ctx.saved_tensors[0]
+    sparsity_block_size = ctx.sparsity_block_size
+
+    return transpose(grad_output, sparsity_layout, sparsity_block_size)[
+        0], None, None, None, None, None
+
+
+def transpose_build_lut(lut: dict, sparsity_layout: Tensor):
+    if lut is None:
+        lut = dict()
+
+    if "sparsity_layout_t" not in lut:
+        sparsity_layout_t = sparsity_layout.transpose(-1, -2).contiguous()
+        lut["sparsity_layout_t"] = sparsity_layout_t
+
+    if "sparsity_lut" not in lut:
+        sparsity_lut = torch.nonzero(lut["sparsity_layout_t"]).contiguous()
+        lut["sparsity_lut"] = sparsity_lut
+
+    if "sparsity_reverse_lut" not in lut:
+        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)))
+                                .reshape(sparsity_layout.size()).transpose(-1, -2).contiguous().reshape(-1))
+        lut["sparsity_reverse_lut"] = sparsity_reverse_lut
+
+    if "n_sparse_blocks" not in lut:
+        n_sparse_blocks = torch.sum(sparsity_layout.to(torch.int)).item()
+        lut["n_sparse_blocks"] = n_sparse_blocks
+
+    validate_contiguous(lut["sparsity_layout_t"], lut["sparsity_lut"], lut["sparsity_reverse_lut"])
+
+    return lut
+
+
+# noinspection PyUnusedLocal
+def transpose_setup_context(ctx, inputs, output):
+    (_, sparsity_layout_o, _, _, sparsity_block_size, _) = inputs
+
+    ctx.save_for_backward(sparsity_layout_o)
+    ctx.sparsity_block_size = sparsity_block_size
+
+
+transpose_forward.register_autograd(transpose_backward, setup_context=transpose_setup_context)

blksprs/utils/benchmarking.py

@@ -5,13 +5,13 @@ from matplotlib import pyplot as plt
 
 
 def benchmark(method_labels: list[str], func_input_generator: Callable,
-              matrix_sizes: list[int], sparsity_block_sizes: list[int], triton_block_sizes: list[int],
+              matrix_sizes: list[int], sparsity_block_sizes: list[int],
               *funcs_test_subject: Callable, y_lim_top: int = None):
     quantiles = [0.5, 0.2, 0.8]
     results = {}
 
-    for matrix_size, sparsity_block_size, triton_block_size in zip(matrix_sizes, sparsity_block_sizes, triton_block_sizes):
-        arguments = func_input_generator(matrix_size, sparsity_block_size, triton_block_size)
+    for matrix_size, sparsity_block_size in zip(matrix_sizes, sparsity_block_sizes):
+        arguments = func_input_generator(matrix_size, sparsity_block_size)
 
         for i, func_test_subject in enumerate(funcs_test_subject):
             func_ms_avg, func_ms_min, func_ms_max = triton.testing.do_bench(

blksprs/utils/tools.py

@@ -1,6 +1,10 @@
 import torch
+import triton
 from torch import Tensor, Size
 
+# Capture scalar outputs for JIT compilation
+torch._dynamo.config.capture_scalar_outputs = True
+
 
 def do_shape_blocksparse(x: Tensor):
     if x.dim() == 3:
@@ -16,10 +20,6 @@ def undo_shape_blocksparse(x: Tensor, shape: Size):
     return x.reshape((*shape[:-2], *x.shape[-2:]))
 
 
-def get_triton_block_size(sparsity_block_size: int, limit: int = 128):
-    return min(sparsity_block_size, limit)
-
-
 def stride(x: Tensor):
     if x.dim() == 2:
         return x.size(1), 1
@@ -27,3 +27,30 @@ def stride(x: Tensor):
         return x.size(1) * x.size(2), x.size(2), 1
     else:
         raise NotImplementedError
+
+
+@torch.compile
+def get_autotune_configs():
+    configs = []
+    config_parameters = [
+        (16, 3, 8),
+        (16, 4, 4),
+        (16, 5, 2),
+
+        (32, 3, 8),
+        (32, 4, 4),
+        (32, 5, 2),
+
+        (64, 3, 8),
+        (64, 4, 4),
+        (64, 5, 2),
+
+        (128, 3, 8),
+        (128, 4, 4),
+        (128, 5, 2),
+    ]
+
+    for block_size, num_stages, num_warps in config_parameters:
+        configs.append(triton.Config({"TRITON_BLOCK_SIZE": block_size}, num_stages=num_stages, num_warps=num_warps))
+
+    return configs

blksprs/utils/validation.py

@@ -104,20 +104,6 @@ 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 _check_skip_validation():
-        return
-
-    if triton_block_size is None:
-        return
-
-    if not (triton_block_size & (triton_block_size - 1)) == 0:
-        raise ValueError("Triton block size must be a power of 2")
-
-    if triton_block_size > sparsity_block_size:
-        raise ValueError("Triton block size cannot be larger than sparsity block size")
-
-
 def _check_skip_validation():
     return not VALIDATION
 

{blksprs-1.10.2.dist-info → blksprs-2.0rc1.dist-info}/METADATA

@@ -1,6 +1,6 @@
-Metadata-Version: 2.1
+Metadata-Version: 2.4
 Name: blksprs
-Version: 1.10.2
+Version: 2.0rc1
 Summary: A lightweight library for operations on blocksparse matrices in PyTorch.
 Author-email: Felix Schön <schoen@kr.tuwien.ac.at>
 Project-URL: Homepage, https://github.com/FelixSchoen/blksprs
@@ -14,9 +14,8 @@ Requires-Dist: pytest; extra == "test"
 Requires-Dist: pytest-xdist; extra == "test"
 Requires-Dist: pytest-cov; extra == "test"
 Requires-Dist: coverage; extra == "test"
+Requires-Dist: build; extra == "test"
 Requires-Dist: matplotlib; extra == "test"
-Provides-Extra: build
-Requires-Dist: build; extra == "build"
 
 # blksprs
 
@@ -25,6 +24,13 @@ Requires-Dist: build; extra == "build"
 
 ## Overview
 
+  ### News
+  
+  🎉 ***Version 2.0 released***. blksprs now supports kernel auto-tuning, JIT compilation, specification of pre-calculated
+  LUTs, and makes use of `torch.library.triton_op()`!
+
+---
+
 A lightweight and efficient library for operations on block-sparse matrices in PyTorch using Triton.
 
 Currently supported operations (includes gradient calculation):
@@ -52,23 +58,25 @@ These include, e.g.,
 Note that in order to correctly apply element-wise operations between two sparse tensors their sparsity layouts have to
 match.
 
-Further helpful operations (included in the ``bs.ops.misc`` module) that do **not** support gradient calculation include:
+Further helpful operations (included in the ``bs.ops.misc`` module) that do **not** support gradient calculation
+include:
 
 - Row-wise sum, max, addition, and subtraction
 - Broadcast addition and subtraction between slices
 
-Furthermore, the library provides a set of utility functions 
+Furthermore, the library provides a set of utility functions
 
 - for the creation of sparsity layouts based on existing
-dense tensors and for the scatter operation (module ``bs.layouting``),
+  dense tensors and for the scatter operation (module ``bs.layouting``),
 - for the application of ``nn.Linear``, ``nn.Dropout``, and ``nn.LayerNorm`` layers to block-sparse tensors,
 - as well as utility functions to ensure correct input dimensionality, and validate input (module ``bs.utils``).
 
-_* see the [Roadmap](#roadmap) section for more information_ 
+_* see the [Roadmap](#roadmap) section for more information_
 
 ## Installation
 
-Note that due to the dependency on [Triton](https://github.com/triton-lang/triton) this library is **only compatible with
+Note that due to the dependency on [Triton](https://github.com/triton-lang/triton) this library is **only compatible
+with
 the Linux platform**.
 Keep track of this [issue](https://github.com/triton-lang/triton/issues/1640) for updates.
 
@@ -78,8 +86,8 @@ We recommend installing blksprs from [PyPI](https://pypi.org/project/blksprs/) u
 
 ### Dependencies
 
-- [PyTorch](https://pytorch.org/) (built with v2.5.1)
-- _[NumPy](https://numpy.org/) (to get rid of warnings, built with v2.2.0)_
+- [PyTorch](https://pytorch.org/) (built with v2.6)
+- _[NumPy](https://numpy.org/) (to get rid of warnings, built with v2.2.4)_
 - _[Triton](https://github.com/triton-lang/triton) (included with PyTorch)_
 
 ## Changelog
@@ -89,12 +97,14 @@ See [`CHANGELOG.md`](https://github.com/FelixSchoen/blksprs/blob/main/CHANGELOG.
 ## Roadmap
 
 Note that since this library covers all our current needs it is in a **bugfix-only** state.
-This means that there are no plans to add new features, e.g., support for dimension specification of the ``split`` and ``merge`` operations.
+This means that there are no plans to add new features, e.g., support for dimension specification of the ``split`` and
+``merge`` operations.
 We will continue to maintain the library and fix any issues that arise.
 Should you find any bugs please open an [issue](https://github.com/FelixSchoen/blksprs/issues).
 We also encourage [pull requests](https://github.com/FelixSchoen/blksprs/pulls).
 
-It might be that this changes with future projects, but as of December 2024, we are content with the current state of the library.
+It might be that this changes with future projects, but as of March 2025, we are content with the current state of the
+library.
 
 ## Usage
 
@@ -120,10 +130,6 @@ def test_readme():
     # Must be a power of two, greater than or equal to 16 for matmul, and divide m, n, and k
     sparsity_block_size = 16
 
-    # Must be a power of two and smaller than or equal to sparsity_block_size
-    # If it is set to ``none`` a value will be chosen automatically
-    triton_block_size = None
-
     # Initialise random (dense) tensors
     x = torch.randn(size=(b, h, m, k), device="cuda")
     y = torch.randn(size=(b, h, n, k), device="cuda").transpose(-1, -2).contiguous()
@@ -133,53 +139,53 @@ def test_readme():
     y_dense, y_shape_original = bs.utils.do_shape_blocksparse(y)
 
     # Create sparsity layouts from existing tensors
-    sparsity_layout_x = bs.layouting.build_sparsity_layout(x_dense, sparsity_block_size,
-                                                           triton_block_size=triton_block_size)
-    sparsity_layout_y = bs.layouting.build_sparsity_layout(y_dense, sparsity_block_size,
-                                                           triton_block_size=triton_block_size)
+    sparsity_layout_x = bs.layouting.build_sparsity_layout(x_dense, sparsity_block_size)
+    sparsity_layout_y = bs.layouting.build_sparsity_layout(y_dense, sparsity_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 = 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)
+    x_sparse = bs.ops.to_sparse(x_dense, sparsity_layout_x, sparsity_block_size)
+    y_sparse = bs.ops.to_sparse(y_dense, sparsity_layout_y, sparsity_block_size)
+
+    # As of version 2.0, blksprs supports JIT compilation
+    matmul_compiled = torch.compile(bs.ops.matmul)
 
     # Perform matrix multiplication
-    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)
+    o_sparse = matmul_compiled(x_sparse, sparsity_layout_x,
+                               y_sparse, sparsity_layout_y,
+                               sparsity_layout_o, sparsity_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)
+    o_dense = bs.ops.to_dense(o_sparse, sparsity_layout_o, sparsity_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 = 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)
+    o_torch_round_trip = bs.ops.to_dense(
+        bs.ops.to_sparse(o_torch, sparsity_layout_o, sparsity_block_size),
+        sparsity_layout_o, sparsity_block_size, fill_value=0)
 
     # 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 = bs.layouting.build_sparsity_layout(o_dense, sparsity_block_size,
-                                                                  triton_block_size=triton_block_size)
+    actual_sparsity_layout_o = bs.layouting.build_sparsity_layout(o_dense, sparsity_block_size)
     assert torch.allclose(actual_sparsity_layout_o.to(torch.int), sparsity_layout_o)
 
     # Convert output tensor back to original shape
     o = bs.utils.undo_shape_blocksparse(o_dense, x_shape_original)
 
     # Other available functions
-    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)
+    bs.ops.transpose(o_sparse, sparsity_layout_o, sparsity_block_size)
+    bs.ops.softmax(o_sparse, sparsity_layout_o, sparsity_block_size)
+    bs.ops.misc.row_wise_sum(o_sparse, sparsity_layout_o, sparsity_block_size)
+    bs.ops.misc.row_wise_max(o_sparse, sparsity_layout_o, sparsity_block_size)
 
 
 def _get_random_sparsity_layout(b, m, n, sparsity_block_size, sparsity_percentage):

blksprs-2.0rc1.dist-info/RECORD

@@ -0,0 +1,22 @@
+blksprs/__init__.py,sha256=OHfpwJCZWGUfpT-DVfC1YSaeZl4aCMNt9CrzMPymywU,1577
+blksprs/layouting/distribution_layout.py,sha256=0glIteoY5oDkiEu5rjLIC-BB_oC4sa3rFWVkohsAG00,5329
+blksprs/layouting/sparsity_layout.py,sha256=UzMcdW7l4zoiLB_LMEbBR1JBdqVSgINDGYvoCYIOulk,10283
+blksprs/ops/conversion.py,sha256=_JKOovDZOmYJLcurJGhgNt5iQB9kOKp3fufFxD8QCZs,22204
+blksprs/ops/distribution.py,sha256=5gE19kPQGQljVbRpDZeqNaOe8ehRhxdQS7PiJp6mMug,21352
+blksprs/ops/flow.py,sha256=G8L_sMAWIM77gv-YLJtyutEzXqyaaofnSX2QKvmDr44,8409
+blksprs/ops/matmul.py,sha256=YAurJcXa_39gRdh2nWUOmbhm8h99arLoO-SN-l134II,11879
+blksprs/ops/partitioning.py,sha256=AooYZOw0oZgA9zXSu09O60hkJcnpWT1OTosr2T2wdQo,9700
+blksprs/ops/repeat.py,sha256=qty0qIFcfiWzROV2A2FB2KiPCC2Pe4q5TwJyGuDBAQE,8839
+blksprs/ops/softmax.py,sha256=eaZ8pfCpNZCX6Gk5Tk-lhNIrBQDhvfHqNNPltqxp91k,12793
+blksprs/ops/transpose.py,sha256=30pGCSjZs42Sg6TEXUdJNCDgmlN1n8aN88uNbV5wOtA,3941
+blksprs/ops/misc/broadcast_ops.py,sha256=lZ5bBIftUKffzeYz77SWB1xmtZTRGMvjF-tG9rqkOXA,6018
+blksprs/ops/misc/row_wise.py,sha256=iwOrHU8HiJGxq2hEmgJGZ60asRm72WLi10-PrpNrdeQ,19532
+blksprs/utils/benchmarking.py,sha256=dLabDscTFn5NkmOI1g7DnKeTneUYW3RIVv9MDF-8BKc,1271
+blksprs/utils/blksprs_tensor.py,sha256=pfoz59aJixj_fIoFx76ySiygwRQUemmgjMKepZ2c4j0,244
+blksprs/utils/processing.py,sha256=WLuMJQ8v-YovXwcDjhlDn3N31WMZXrtyeeyKSgq_zn4,3642
+blksprs/utils/tools.py,sha256=RL18P4NAj7d8gXTTKbMZt4SHCynsw1wPu9yvlrnBQlo,1220
+blksprs/utils/validation.py,sha256=_Ee6bqu7CxdYLFSy4WZOFoXJgd0p_RBMumCwGCk2_Hw,3763
+blksprs-2.0rc1.dist-info/METADATA,sha256=zXzVOvuwgYSyx-lCBycdFvRUmHUD_qYbK8sFkKWZnp8,8601
+blksprs-2.0rc1.dist-info/WHEEL,sha256=1tXe9gY0PYatrMPMDd6jXqjfpz_B-Wqm32CPfRC58XU,91
+blksprs-2.0rc1.dist-info/top_level.txt,sha256=qyp0IHeY3H2GQA97i4hk_To5rRBS2YcE1HRPSLy04fk,8
+blksprs-2.0rc1.dist-info/RECORD,,

{blksprs-1.10.2.dist-info → blksprs-2.0rc1.dist-info}/WHEEL

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

blksprs/ops/misc/exp.py

@@ -1,104 +0,0 @@
-import torch
-import triton
-from torch import Tensor
-from triton import language as tl
-
-from blksprs.utils.blksprs_tensor import BlksprsTensor
-from blksprs.utils.tools import get_triton_block_size, stride
-from blksprs.utils.validation import validate_contiguous, validate_dimensions, validate_device, \
-    validate_sparsity_block_size, validate_triton_block_size
-
-
-def exp(x: BlksprsTensor, sparsity_block_size: int, triton_block_size: int = None) -> BlksprsTensor:
-    """Applies the element-wise exponential function to a block-sparse tensor.
-
-    Note:
-        This operation does not consider sparse blocks, i.e., these will not be set to ``e^0``.
-        Consider this when converting back to tensors in regular form.
-
-    Args:
-        x (BlksprsTensor): A block-sparse tensor in compressed form.
-        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:
-        BlksprsTensor: The exponential function applied to all elements of the input tensor as a block-sparse tensor in
-            compressed form.
-
-    """
-    x = x.contiguous()
-
-    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)
-
-    return BlksprsTensor(_BlocksparseExp.apply(x, sparsity_block_size, triton_block_size))
-
-
-class _BlocksparseExp(torch.autograd.Function):
-
-    @staticmethod
-    def forward(ctx, x: Tensor, sparsity_block_size: int, triton_block_size: int) -> Tensor:
-        output = torch.empty_like(x)
-
-        x_b, x_r, x_c = x.shape
-        x_b_s, x_r_s, x_c_s = stride(x)
-        o_b, o_r, o_c = output.shape
-        o_b_s, o_r_s, o_c_s = stride(output)
-
-        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"])]
-
-        (_BlocksparseExp.kernel_blocksparse_exp[triton_grid]
-         (x,
-          x_b, x_b_s, x_r_s, x_c_s,
-          output,
-          o_b, o_b_s, o_r_s, o_c_s,
-          triton_block_size))
-
-        ctx.save_for_backward(output)
-
-        return output
-
-    @staticmethod
-    def backward(ctx, grad_output):
-        o = ctx.saved_tensors[0]
-
-        grad_x = torch.mul(grad_output, o)
-
-        return grad_x, None, None
-
-    @staticmethod
-    @triton.jit
-    def kernel_blocksparse_exp(x,
-                               x_b, x_b_s, x_r_s, x_c_s,
-                               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)
-
-        # 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 >= 0 and blk_x_idx < x_b * x_b_s)
-        blk_x = tl.load(x + blk_x_idx, mask=blk_x_msk)
-
-        # Compute exp
-        buf = tl.exp(blk_x)
-
-        # 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 >= 0 and blk_o_idx < o_b * o_b_s)
-        tl.store(o + blk_o_idx, buf, mask=blk_o_msk)

blksprs/utils/layout_utils.py

@@ -1,17 +0,0 @@
-import math
-
-import torch
-import triton
-from torch import Tensor
-from torch.xpu import device
-from triton import language as tl
-
-from blksprs.utils.blksprs_tensor import BlksprsTensor
-from blksprs.utils.tools import get_triton_block_size, stride
-from blksprs.utils.validation import validate_triton_block_size, validate_dimensions, validate_device, \
-    validate_contiguous, validate_sparsity, validate_sparsity_block_size
-
-
-def build_full_sparsity_layout(x: Tensor, sparsity_block_size: int) -> Tensor:
-    return torch.ones(size=(x.size(0), x.size(1) // sparsity_block_size, x.size(2) // sparsity_block_size),
-                      dtype=torch.bool, device=x.device)