blksprs 2.1.4__py3-none-any.whl → 2.1.5__py3-none-any.whl

blksprs/__init__.py

@@ -1,6 +1,6 @@
 from blksprs.utils.blksprs_tensor import BlksprsTensor
 
-__version__ = "2.1.4"
+__version__ = "2.1.5"
 
 
 class ops:
@@ -27,9 +27,9 @@ class utils:
     from blksprs.utils.processing import apply_torch_linear, apply_torch_normalisation, apply_torch_dropout, \
         apply_function_applicable_row_wise
     from blksprs.utils.tools import do_shape_blocksparse, undo_shape_blocksparse
+    from blksprs.utils.validation import disable_contiguous, disable_validation
 
     class validation:
-        from blksprs.utils.validation import disable_validation
         from blksprs.utils.validation import validate_dimensions, validate_contiguous, validate_dtype_float, \
             validate_dtype_int, validate_device, validate_sparsity, validate_sparsity_dense, \
             validate_sparsity_block_size

blksprs/ops/conversion.py

@@ -106,17 +106,13 @@ def to_sparse_kernel(x,
     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_val_idx = pid_blk * s_lut_r_s + tl.arange(0, 4) * s_lut_c_s
+    spa_val_msk = (tl.arange(0, 4) < 3)
+    spa_val = tl.load(s_lut + spa_val_idx, mask=spa_val_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)
+    spa_bat = tl.sum(spa_val * (tl.arange(0, 4) == 0))
+    spa_row = tl.sum(spa_val * (tl.arange(0, 4) == 1))
+    spa_col = tl.sum(spa_val * (tl.arange(0, 4) == 2))
 
     # Load block from dense tensor
     blk_d_idx = (spa_bat * x_b_s +
@@ -445,17 +441,13 @@ def adapt_layout_kernel(x,
     pid_col = tl.program_id(axis=2)
 
     # Get position of current sparsity block consisting of its batch, row, and column index
-    spa_bat_o_idx = (pid_blk * s_lut_o_r_s + 0 * s_lut_o_c_s)
-    spa_bat_o_msk = (spa_bat_o_idx >= 0 and spa_bat_o_idx < s_lut_o_r * s_lut_o_r_s)
-    spa_bat_o = tl.load(s_lut_o + spa_bat_o_idx, mask=spa_bat_o_msk)
-
-    spa_row_o_idx = (pid_blk * s_lut_o_r_s + 1 * s_lut_o_c_s)
-    spa_row_o_msk = (spa_row_o_idx >= 0 and spa_row_o_idx < s_lut_o_r * s_lut_o_r_s)
-    spa_row_o = tl.load(s_lut_o + spa_row_o_idx, mask=spa_row_o_msk)
+    spa_val_idx = pid_blk * s_lut_o_r_s + tl.arange(0, 4) * s_lut_o_c_s
+    spa_val_msk = (tl.arange(0, 4) < 3)
+    spa_val = tl.load(s_lut_o + spa_val_idx, mask=spa_val_msk)
 
-    spa_col_o_idx = (pid_blk * s_lut_o_r_s + 2 * s_lut_o_c_s)
-    spa_col_o_msk = (spa_col_o_idx >= 0 and spa_col_o_idx < s_lut_o_r * s_lut_o_r_s)
-    spa_col_o = tl.load(s_lut_o + spa_col_o_idx, mask=spa_col_o_msk)
+    spa_bat_o = tl.sum(spa_val * (tl.arange(0, 4) == 0))
+    spa_row_o = tl.sum(spa_val * (tl.arange(0, 4) == 1))
+    spa_col_o = tl.sum(spa_val * (tl.arange(0, 4) == 2))
 
     # Get equivalent sparsity block in from layout
     spa_bat_x = spa_bat_o

blksprs/ops/distribution.py

@@ -125,17 +125,13 @@ def gather_kernel(x,
     pid_col = tl.program_id(axis=2)
 
     # Get position of current sparsity block consisting of its batch, row, and column index
-    spa_bat_o_idx = (pid_blk * s_lut_o_r_s + 0 * s_lut_o_c_s)
-    spa_bat_o_msk = (spa_bat_o_idx >= 0 and spa_bat_o_idx < s_lut_o_r * s_lut_o_r_s)
-    spa_bat_o = tl.load(s_lut_o + spa_bat_o_idx, mask=spa_bat_o_msk)
+    spa_val_idx = pid_blk * s_lut_o_r_s + tl.arange(0, 4) * s_lut_o_c_s
+    spa_val_msk = (tl.arange(0, 4) < 3)
+    spa_val = tl.load(s_lut_o + spa_val_idx, mask=spa_val_msk)
 
-    spa_row_o_idx = (pid_blk * s_lut_o_r_s + 1 * s_lut_o_c_s)
-    spa_row_o_msk = (spa_row_o_idx >= 0 and spa_row_o_idx < s_lut_o_r * s_lut_o_r_s)
-    spa_row_o = tl.load(s_lut_o + spa_row_o_idx, mask=spa_row_o_msk)
-
-    spa_col_o_idx = (pid_blk * s_lut_o_r_s + 2 * s_lut_o_c_s)
-    spa_col_o_msk = (spa_col_o_idx >= 0 and spa_col_o_idx < s_lut_o_r * s_lut_o_r_s)
-    spa_col_o = tl.load(s_lut_o + spa_col_o_idx, mask=spa_col_o_msk)
+    spa_bat_o = tl.sum(spa_val * (tl.arange(0, 4) == 0))
+    spa_row_o = tl.sum(spa_val * (tl.arange(0, 4) == 1))
+    spa_col_o = tl.sum(spa_val * (tl.arange(0, 4) == 2))
 
     # Load index values
     blk_i_idx = ((pid_blk * i_b_s) +
@@ -374,17 +370,13 @@ def scatter_reduce_kernel(x,
     pid_col = tl.program_id(axis=2)
 
     # Get position of current sparsity block consisting of its batch, row, and column index
-    spa_bat_x_idx = (pid_blk * s_lut_x_r_s + 0 * s_lut_x_c_s)
-    spa_bat_x_msk = (spa_bat_x_idx >= 0 and spa_bat_x_idx < s_lut_x_r * s_lut_x_r_s)
-    spa_bat_x = tl.load(s_lut_x + spa_bat_x_idx, mask=spa_bat_x_msk)
-
-    spa_row_x_idx = (pid_blk * s_lut_x_r_s + 1 * s_lut_x_c_s)
-    spa_row_x_msk = (spa_row_x_idx >= 0 and spa_row_x_idx < s_lut_x_r * s_lut_x_r_s)
-    spa_row_x = tl.load(s_lut_x + spa_row_x_idx, mask=spa_row_x_msk)
+    spa_val_idx = pid_blk * s_lut_x_r_s + tl.arange(0, 4) * s_lut_x_c_s
+    spa_val_msk = (tl.arange(0, 4) < 3)
+    spa_val = tl.load(s_lut_x + spa_val_idx, mask=spa_val_msk)
 
-    spa_col_x_idx = (pid_blk * s_lut_x_r_s + 2 * s_lut_x_c_s)
-    spa_col_x_msk = (spa_col_x_idx >= 0 and spa_col_x_idx < s_lut_x_r * s_lut_x_r_s)
-    spa_col_x = tl.load(s_lut_x + spa_col_x_idx, mask=spa_col_x_msk)
+    spa_bat_x = tl.sum(spa_val * (tl.arange(0, 4) == 0))
+    spa_row_x = tl.sum(spa_val * (tl.arange(0, 4) == 1))
+    spa_col_x = tl.sum(spa_val * (tl.arange(0, 4) == 2))
 
     # Load x values
     blk_x_idx = ((pid_blk * x_b_s) +

blksprs/ops/flow.py

@@ -66,17 +66,13 @@ def flow_pull_kernel(x,
     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_val_idx = pid_blk * s_lut_r_s + tl.arange(0, 4) * s_lut_c_s
+    spa_val_msk = (tl.arange(0, 4) < 3)
+    spa_val = tl.load(s_lut + spa_val_idx, mask=spa_val_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)
+    spa_bat = tl.sum(spa_val * (tl.arange(0, 4) == 0))
+    spa_row = tl.sum(spa_val * (tl.arange(0, 4) == 1))
+    spa_col = tl.sum(spa_val * (tl.arange(0, 4) == 2))
 
     # Load reverse sparsity index
     rev_idx_spa_idx = (spa_bat * s_l_o_b_s +
@@ -157,17 +153,13 @@ def flow_push_kernel(x,
     pid_col = tl.program_id(axis=2)
 
     # Get sparsity index of current input 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_val_idx = pid_blk * s_lut_r_s + tl.arange(0, 4) * s_lut_c_s
+    spa_val_msk = (tl.arange(0, 4) < 3)
+    spa_val = tl.load(s_lut + spa_val_idx, mask=spa_val_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)
+    spa_bat = tl.sum(spa_val * (tl.arange(0, 4) == 0))
+    spa_row = tl.sum(spa_val * (tl.arange(0, 4) == 1))
+    spa_col = tl.sum(spa_val * (tl.arange(0, 4) == 2))
 
     # Get reverse sparsity index
     rev_idx_spa_idx = (spa_bat * s_l_x_b_s +

blksprs/ops/matmul.py

@@ -145,17 +145,13 @@ def matmul_kernel(x,
     pid_col = tl.program_id(axis=2)
 
     # Get position of current sparsity block consisting of its batch, row, and column index
-    spa_bat_o_idx = (pid_blk * s_lut_o_r_s + 0 * s_lut_o_c_s)
-    spa_bat_o_msk = (spa_bat_o_idx >= 0 and spa_bat_o_idx < s_lut_o_r * s_lut_o_r_s)
-    spa_bat_o = tl.load(s_lut_o + spa_bat_o_idx, mask=spa_bat_o_msk)
+    spa_val_idx = pid_blk * s_lut_o_r_s + tl.arange(0, 4) * s_lut_o_c_s
+    spa_val_msk = (tl.arange(0, 4) < 3)
+    spa_val = tl.load(s_lut_o + spa_val_idx, mask=spa_val_msk)
 
-    spa_row_o_idx = (pid_blk * s_lut_o_r_s + 1 * s_lut_o_c_s)
-    spa_row_o_msk = (spa_row_o_idx >= 0 and spa_row_o_idx < s_lut_o_r * s_lut_o_r_s)
-    spa_row_o = tl.load(s_lut_o + spa_row_o_idx, mask=spa_row_o_msk)
-
-    spa_col_o_idx = (pid_blk * s_lut_o_r_s + 2 * s_lut_o_c_s)
-    spa_col_o_msk = (spa_col_o_idx >= 0 and spa_col_o_idx < s_lut_o_r * s_lut_o_r_s)
-    spa_col_o = tl.load(s_lut_o + spa_col_o_idx, mask=spa_col_o_msk)
+    spa_bat_o = tl.sum(spa_val * (tl.arange(0, 4) == 0))
+    spa_row_o = tl.sum(spa_val * (tl.arange(0, 4) == 1))
+    spa_col_o = tl.sum(spa_val * (tl.arange(0, 4) == 2))
 
     # Setup buffer
     buf = tl.zeros(shape=(TRITON_BLOCK_SIZE, TRITON_BLOCK_SIZE), dtype=tl.float32)

blksprs/ops/misc/broadcast_ops.py

@@ -110,17 +110,13 @@ def broadcast_add_kernel(x,
     pid_col = tl.program_id(axis=2)
 
     # Get position of current sparsity block consisting of its batch, row, and column index
-    spa_bat_o_idx = (pid_blk * s_lut_o_r_s + 0 * s_lut_o_c_s)
-    spa_bat_o_msk = (spa_bat_o_idx >= 0 and spa_bat_o_idx < s_lut_o_r * s_lut_o_r_s)
-    spa_bat_o = tl.load(s_lut_o + spa_bat_o_idx, mask=spa_bat_o_msk)
+    spa_val_idx = pid_blk * s_lut_o_r_s + tl.arange(0, 4) * s_lut_o_c_s
+    spa_val_msk = (tl.arange(0, 4) < 3)
+    spa_val = tl.load(s_lut_o + spa_val_idx, mask=spa_val_msk)
 
-    spa_row_o_idx = (pid_blk * s_lut_o_r_s + 1 * s_lut_o_c_s)
-    spa_row_o_msk = (spa_row_o_idx >= 0 and spa_row_o_idx < s_lut_o_r * s_lut_o_r_s)
-    spa_row_o = tl.load(s_lut_o + spa_row_o_idx, mask=spa_row_o_msk)
-
-    spa_col_o_idx = (pid_blk * s_lut_o_r_s + 2 * s_lut_o_c_s)
-    spa_col_o_msk = (spa_col_o_idx >= 0 and spa_col_o_idx < s_lut_o_r * s_lut_o_r_s)
-    spa_col_o = tl.load(s_lut_o + spa_col_o_idx, mask=spa_col_o_msk)
+    spa_bat_o = tl.sum(spa_val * (tl.arange(0, 4) == 0))
+    spa_row_o = tl.sum(spa_val * (tl.arange(0, 4) == 1))
+    spa_col_o = tl.sum(spa_val * (tl.arange(0, 4) == 2))
 
     # Load x block
     blk_x_idx = (spa_bat_o * x_b_s +

blksprs/ops/misc/row_wise.py

@@ -119,17 +119,17 @@ def row_wise_sum_kernel(x,
     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 >= 0 and 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_val_idx = pid_blk * s_lut_x_r_s + tl.arange(0, 4) * s_lut_x_c_s
+    spa_val_msk = (tl.arange(0, 4) < 3)
+    spa_val = tl.load(s_lut_x + spa_val_idx, mask=spa_val_msk)
 
-    spa_row_idx = (pid_blk * s_lut_x_r_s + 1 * s_lut_x_c_s)
-    spa_row_msk = (spa_row_idx >= 0 and 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)
+    spa_bat_x = tl.sum(spa_val * (tl.arange(0, 4) == 0))
+    spa_row_x = tl.sum(spa_val * (tl.arange(0, 4) == 1))
+    spa_col_x = tl.sum(spa_val * (tl.arange(0, 4) == 2))
 
     # 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_idx = (spa_bat_x * s_l_o_b_s +
+                       spa_row_x * s_l_o_r_s)
     rev_idx_spa_msk = (rev_idx_spa_idx >= 0 and 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)
 
@@ -263,17 +263,17 @@ def row_wise_max_kernel(x,
     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 >= 0 and 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_val_idx = pid_blk * s_lut_x_r_s + tl.arange(0, 4) * s_lut_x_c_s
+    spa_val_msk = (tl.arange(0, 4) < 3)
+    spa_val = tl.load(s_lut_x + spa_val_idx, mask=spa_val_msk)
 
-    spa_row_idx = (pid_blk * s_lut_x_r_s + 1 * s_lut_x_c_s)
-    spa_row_msk = (spa_row_idx >= 0 and 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)
+    spa_bat_x = tl.sum(spa_val * (tl.arange(0, 4) == 0))
+    spa_row_x = tl.sum(spa_val * (tl.arange(0, 4) == 1))
+    spa_col_x = tl.sum(spa_val * (tl.arange(0, 4) == 2))
 
     # 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_idx = (spa_bat_x * s_l_o_b_s +
+                       spa_row_x * s_l_o_r_s)
     rev_idx_spa_msk = (rev_idx_spa_idx >= 0 and 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)
 
@@ -361,7 +361,7 @@ def row_wise_add_forward(x: Tensor, sparsity_lut_x: Tensor,
                                     triton.cdiv(o_r, meta["TRITON_BLOCK_SIZE"]),
                                     triton.cdiv(o_c, meta["TRITON_BLOCK_SIZE"])]
 
-        (wrap_triton(kernel_blocksparse_row_wise_add)[triton_grid]
+        (wrap_triton(row_wise_add_kernel)[triton_grid]
          (x,
           x_b, x_b_s, x_r_s, x_c_s,
           sparsity_lut_x, s_lut_r, s_lut_r_s, s_lut_c_s,
@@ -383,33 +383,33 @@ def row_wise_add_forward(x: Tensor, sparsity_lut_x: Tensor,
     reset_to_zero=["o"]
 )
 @triton.jit
-def kernel_blocksparse_row_wise_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,
-                                    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,
-                                    sparsity_block_size,
-                                    TRITON_BLOCK_SIZE: tl.constexpr) -> None:
+def row_wise_add_kernel(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,
+                        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,
+                        sparsity_block_size,
+                        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_x_r_s + 0 * s_lut_x_c_s)
-    spa_bat_msk = (spa_bat_idx >= 0 and 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_val_idx = pid_blk * s_lut_x_r_s + tl.arange(0, 4) * s_lut_x_c_s
+    spa_val_msk = (tl.arange(0, 4) < 3)
+    spa_val = tl.load(s_lut_x + spa_val_idx, mask=spa_val_msk)
 
-    spa_row_idx = (pid_blk * s_lut_x_r_s + 1 * s_lut_x_c_s)
-    spa_row_msk = (spa_row_idx >= 0 and 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)
+    spa_bat_x = tl.sum(spa_val * (tl.arange(0, 4) == 0))
+    spa_row_x = tl.sum(spa_val * (tl.arange(0, 4) == 1))
+    spa_col_x = tl.sum(spa_val * (tl.arange(0, 4) == 2))
 
     # 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_idx = (spa_bat_x * s_l_y_b_s +
+                         spa_row_x * s_l_y_r_s)
     rev_idx_spa_s_msk = (rev_idx_spa_s_idx >= 0 and 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)
 

blksprs/ops/repeat.py

@@ -142,7 +142,7 @@ def repeat_build_lut(lut: dict, sparsity_layout_x: Tensor, repeats: tuple[int, i
         n_sparse_blocks = torch.sum(lut["sparsity_layout_o"].to(torch.int)).item()
         lut["n_sparse_blocks"] = n_sparse_blocks
 
-    validate_contiguous(sparsity_layout_o, lut["sparsity_lut"], lut["sparsity_reverse_lut"])
+    validate_contiguous(lut["sparsity_layout_o"], lut["sparsity_lut"], lut["sparsity_reverse_lut"])
 
     return lut
 
@@ -178,7 +178,7 @@ def repeat_interleave_build_lut(lut: dict, sparsity_layout_x: Tensor, repeats: i
         n_sparse_blocks = torch.sum(lut["sparsity_layout_o"].to(torch.int)).item()
         lut["n_sparse_blocks"] = n_sparse_blocks
 
-    validate_contiguous(sparsity_layout_o, lut["sparsity_lut"], lut["sparsity_reverse_lut"])
+    validate_contiguous(lut["sparsity_layout_o"], lut["sparsity_lut"], lut["sparsity_reverse_lut"])
 
     return lut
 

blksprs/ops/softmax.py

@@ -176,13 +176,12 @@ def softmax_kernel(x,
     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 >= 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_val_idx = pid_blk * s_lut_r_s + tl.arange(0, 4) * s_lut_c_s
+    spa_val_msk = (tl.arange(0, 4) < 3)
+    spa_val = tl.load(s_lut + spa_val_idx, mask=spa_val_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_bat = tl.sum(spa_val * (tl.arange(0, 4) == 0))
+    spa_row = tl.sum(spa_val * (tl.arange(0, 4) == 1))
 
     # Get reverse sparsity indices for s
     rev_idx_spa_s_idx = (spa_bat * s_l_s_b_s +
@@ -241,13 +240,12 @@ def softmax_kernel_grad(g,
     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 >= 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_val_idx = pid_blk * s_lut_r_s + tl.arange(0, 4) * s_lut_c_s
+    spa_val_msk = (tl.arange(0, 4) < 3)
+    spa_val = tl.load(s_lut + spa_val_idx, mask=spa_val_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_bat = tl.sum(spa_val * (tl.arange(0, 4) == 0))
+    spa_row = tl.sum(spa_val * (tl.arange(0, 4) == 1))
 
     rev_idx_spa_s_idx = (spa_bat * s_l_s_b_s +
                          spa_row * s_l_s_r_s)

blksprs/utils/autotuning.py

@@ -14,11 +14,11 @@ if blksprs_autotune_mode == "DEFAULT":
 
         (64, 3, 8),
         (64, 4, 4),
-        (64, 5, 2),
+        (64, 4, 8),
 
         (128, 3, 8),
         (128, 4, 4),
-        (128, 5, 2),
+        (128, 4, 8),
     ]
 elif blksprs_autotune_mode == "TEST":
     autotune_parameters = [

blksprs/utils/validation.py

@@ -1,9 +1,17 @@
 import torch
 from torch import Tensor
 
+CONTIGUOUS = True
 VALIDATION = True
 
 
+def ensure_contiguous(*tensors: Tensor) -> tuple[Tensor, ...]:
+    if _check_skip_contiguous():
+        return tensors
+
+    return tuple(tensor.contiguous() for tensor in tensors)
+
+
 def validate_dimensions(*tensors: Tensor, dims=3) -> None:
     if _check_skip_validation():
         return
@@ -124,6 +132,19 @@ def validate_sparsity_block_size(sparsity_block_size: int, *tensors):
             raise ValueError("Tensor sizes must be divisible by sparsity block size")
 
 
+def _check_skip_contiguous():
+    return not CONTIGUOUS
+
+
+def _set_skip_contiguous(skip_contiguous: bool):
+    global CONTIGUOUS
+    CONTIGUOUS = not skip_contiguous
+
+
+def disable_contiguous():
+    _set_skip_contiguous(True)
+
+
 def _check_skip_validation():
     return not VALIDATION
 

{blksprs-2.1.4.dist-info → blksprs-2.1.5.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: blksprs
-Version: 2.1.4
+Version: 2.1.5
 Summary: A lightweight library for operations on block-sparse matrices in PyTorch.
 Author-email: Felix Schön <schoen@kr.tuwien.ac.at>
 Project-URL: Homepage, https://github.com/FelixSchoen/blksprs

blksprs-2.1.5.dist-info/RECORD

@@ -0,0 +1,23 @@
+blksprs/__init__.py,sha256=xlrLL9EgaiEnGQsdzFScy4SVZMN9g_5nvX-LkWxVKCw,1631
+blksprs/layouting/distribution_layout.py,sha256=ur1ty_2U-Hfj78hMWsLZvu7ZuGhzW3qGLKMc72DfTZM,5861
+blksprs/layouting/sparsity_layout.py,sha256=eXHmu2h7K5Q-YUpfOxocJoeP_5ZoQFZf_eHLxRZQbYU,11207
+blksprs/ops/conversion.py,sha256=_LgkT-6aSLPO2FXeMA2lE26g9qAEzxhMWcagenMedFU,21368
+blksprs/ops/distribution.py,sha256=HcFKcB1x59cP8Im_LuxKeJXTknZNM2Kx8hz3nu1GpvE,20183
+blksprs/ops/flow.py,sha256=JEGES5ZbMqxR02rwi2Ym4j3VDxkcRxhFO1f-5nNUlM8,7760
+blksprs/ops/matmul.py,sha256=ZYOv8Qeb7pBpbMsMnndk7IR2WO8rEXfL_KtYhbVeFdw,11576
+blksprs/ops/partitioning.py,sha256=cfQmY9BZqGTvvJorIhtb-EyuGRJGPraWR-wTKdb47aI,9954
+blksprs/ops/repeat.py,sha256=-rFC-u2eytmFxKi7vZTXpvyxReHOPZeRz4SvuO07NxE,9049
+blksprs/ops/softmax.py,sha256=iJ8GniyM83iKM3J9BXTpLdqqEVeRjxeU2rAKP553VPM,23439
+blksprs/ops/transpose.py,sha256=U-VAyLRT6_NDv9qYSFzBqfVlDeIpTqAMEXkqto0VF6w,4072
+blksprs/ops/misc/broadcast_ops.py,sha256=TD7wjBJIMn-4SUdYy7e_5bpf0UQ4Sga4QEipQFaaVPM,5684
+blksprs/ops/misc/row_wise.py,sha256=kKKpDfpq92UU5P7HVuK9gh2MNPvHOB2KQ6ijKE1RmHM,19359
+blksprs/utils/autotuning.py,sha256=xalNP3sWdRn8XiVG4jE1-_iy2QhUmIJvTGM83YwgKA0,2052
+blksprs/utils/benchmarking.py,sha256=dLabDscTFn5NkmOI1g7DnKeTneUYW3RIVv9MDF-8BKc,1271
+blksprs/utils/blksprs_tensor.py,sha256=pfoz59aJixj_fIoFx76ySiygwRQUemmgjMKepZ2c4j0,244
+blksprs/utils/processing.py,sha256=RNkEDc0g-sNHRuMPkRzNWU13d3_lIkXMJdoqES4yQTM,3738
+blksprs/utils/tools.py,sha256=TKygEKge4wJtJnXXDg8BTL8vzBpqIJsQ_A3_5FmLpcE,859
+blksprs/utils/validation.py,sha256=hME6hf5t7-IxM1rHypqlzk7IE1kYEQACqCZ9KEtW6N0,4775
+blksprs-2.1.5.dist-info/METADATA,sha256=MnA7fThFWn_mrMk0BFEkBm29rtWtPY-npQecpGF1P7c,9590
+blksprs-2.1.5.dist-info/WHEEL,sha256=_zCd3N1l69ArxyTb8rzEoP9TpbYXkqRFSNOD5OuxnTs,91
+blksprs-2.1.5.dist-info/top_level.txt,sha256=qyp0IHeY3H2GQA97i4hk_To5rRBS2YcE1HRPSLy04fk,8
+blksprs-2.1.5.dist-info/RECORD,,

blksprs-2.1.4.dist-info/RECORD

@@ -1,23 +0,0 @@
-blksprs/__init__.py,sha256=XERzTtkiElDeBppOO8rNrF6OktUQf_yozDiA4DUXqTY,1615
-blksprs/layouting/distribution_layout.py,sha256=ur1ty_2U-Hfj78hMWsLZvu7ZuGhzW3qGLKMc72DfTZM,5861
-blksprs/layouting/sparsity_layout.py,sha256=eXHmu2h7K5Q-YUpfOxocJoeP_5ZoQFZf_eHLxRZQbYU,11207
-blksprs/ops/conversion.py,sha256=nv5gXiyZkUtk1kCIlPr0Vpaj4G8G6dJdW7StlbV3nDw,21914
-blksprs/ops/distribution.py,sha256=0tPldv0ARzmCV1CU2jvfqpHBgOuHPrDFiCtqsLs7CZc,20789
-blksprs/ops/flow.py,sha256=oUn_xDT74220-EmnBnB8bRNtbS1mjbxWpm76PFsK22o,8246
-blksprs/ops/matmul.py,sha256=ES9bpiCIRBxaynNIL5ftDP0c9LSArbj8YJqkPEzBaIU,11879
-blksprs/ops/partitioning.py,sha256=cfQmY9BZqGTvvJorIhtb-EyuGRJGPraWR-wTKdb47aI,9954
-blksprs/ops/repeat.py,sha256=TLYNxwPuT9y5K9xyM41WK5gnggAJF3lI61Q2K7zWjns,9035
-blksprs/ops/softmax.py,sha256=tfC_jaAKrA956rxGeb57klMuYRKTiyMCd5Zg5DIH3fc,23649
-blksprs/ops/transpose.py,sha256=U-VAyLRT6_NDv9qYSFzBqfVlDeIpTqAMEXkqto0VF6w,4072
-blksprs/ops/misc/broadcast_ops.py,sha256=-PrHiSJikZh8nXUmXxSCtFEP27TTxFr4wcrNxBjnimk,5987
-blksprs/ops/misc/row_wise.py,sha256=n5FJjAuOd8BHBJQx4bsQwr-HmXkR9PYVAqfk77wjOFU,19653
-blksprs/utils/autotuning.py,sha256=a-kmWRjJ3eED2XbjkQeOJSyW8bdIs27HgKMPvAKqWeU,2052
-blksprs/utils/benchmarking.py,sha256=dLabDscTFn5NkmOI1g7DnKeTneUYW3RIVv9MDF-8BKc,1271
-blksprs/utils/blksprs_tensor.py,sha256=pfoz59aJixj_fIoFx76ySiygwRQUemmgjMKepZ2c4j0,244
-blksprs/utils/processing.py,sha256=RNkEDc0g-sNHRuMPkRzNWU13d3_lIkXMJdoqES4yQTM,3738
-blksprs/utils/tools.py,sha256=TKygEKge4wJtJnXXDg8BTL8vzBpqIJsQ_A3_5FmLpcE,859
-blksprs/utils/validation.py,sha256=G8eQlvJVMKfEX3k2AwBD0A6Ck-gFoRLpLNY6HXsB3fA,4348
-blksprs-2.1.4.dist-info/METADATA,sha256=qGLQunHEIoHlmRvFnM0TVDjOSApwGzBglpZezmfhHLU,9590
-blksprs-2.1.4.dist-info/WHEEL,sha256=_zCd3N1l69ArxyTb8rzEoP9TpbYXkqRFSNOD5OuxnTs,91
-blksprs-2.1.4.dist-info/top_level.txt,sha256=qyp0IHeY3H2GQA97i4hk_To5rRBS2YcE1HRPSLy04fk,8
-blksprs-2.1.4.dist-info/RECORD,,