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

blksprs/__init__.py

@@ -15,4 +15,4 @@ class misc:
     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
+    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_ops.py

@@ -25,6 +25,9 @@ def broadcast_add(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):

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

@@ -31,6 +31,8 @@ def row_wise_sum(x: Tensor, sparsity_layout: Tensor, sparsity_block_size: int,
             of the input and the sparsity layout of the output tensor.
 
     """
+    x = x.contiguous()
+
     validate_dimensions(x)
     validate_contiguous(x)
     validate_device(x)
@@ -151,6 +153,8 @@ def row_wise_max(x: Tensor, sparsity_layout: Tensor, sparsity_block_size: int,
             of the input and the sparsity layout of the output tensor.
 
     """
+    x = x.contiguous()
+
     validate_dimensions(x)
     validate_contiguous(x)
     validate_device(x)

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

@@ -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)

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.4.dist-info → blksprs-1.4.1.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: blksprs
-Version: 1.4
+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"
@@ -146,7 +144,7 @@ def test_readme():
     # Assert that the output has the correct sparsity layout
     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)
+    assert torch.allclose(actual_sparsity_layout_o.to(torch.int), sparsity_layout_o)
 
     # Convert output tensor back to original shape
     o = bs.util.undo_shape_blocksparse(o_dense, x_shape_original)

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-1.4.dist-info/RECORD

@@ -1,19 +0,0 @@
-blksprs/__init__.py,sha256=vUthykoYgmHqo2rNYgfrKTNMq7IDalRpCa1nVdFEOqA,813
-blksprs/layouting/distribution_layout.py,sha256=GQ-ZRXbeImiLcbaqnL2FuUZ6DoFwmB0naT_YrOpD84Q,4940
-blksprs/layouting/sparsity_layout.py,sha256=TtADT_WWcZpW3zyGy6KAgkAo44gDryXZqdJLZGEX2V8,7895
-blksprs/misc/broadcast_ops.py,sha256=xLj7CH5yEBihI5gT8SRFqQta1DXvl3iSskhbHsOX_EM,5261
-blksprs/misc/repeat_interleave.py,sha256=WrIp7uJsnvjIhFeLYPfkL2j5vXyKmDQGrJ69b3Y0lQ8,5644
-blksprs/misc/row_wise.py,sha256=Fa57BVfmneXT_8Ms-Vao8H8fh89sT3Z0b_gtN-7gano,16805
-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/softmax.py,sha256=1lxgS12oJ5UcRkDxq13OOjp9AHwhgzSfBosEO1GzKvs,11948
-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.4.dist-info/METADATA,sha256=mC9Vql8wtF_gLYwnGXx8_p9aKL7PnxrQSoZNkQegxic,7675
-blksprs-1.4.dist-info/WHEEL,sha256=GV9aMThwP_4oNCtvEC2ec3qUYutgWeAzklro_0m4WJQ,91
-blksprs-1.4.dist-info/top_level.txt,sha256=qyp0IHeY3H2GQA97i4hk_To5rRBS2YcE1HRPSLy04fk,8
-blksprs-1.4.dist-info/RECORD,,