blksprs 2.1.4__py3-none-any.whl → 2.1.6__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.6"
 
 
 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

@@ -9,7 +9,7 @@ from blksprs.utils.blksprs_tensor import BlksprsTensor
 from blksprs.utils.tools import stride
 from blksprs.utils.autotuning import get_autotune_configs, prune_autotune_configs, prune_autotune_configs_conversion
 from blksprs.utils.validation import validate_contiguous, validate_dimensions, validate_device, \
-    validate_sparsity, validate_sparsity_block_size, validate_sparsity_dense
+    validate_sparsity, validate_sparsity_block_size, validate_sparsity_dense, ensure_contiguous
 
 
 def to_blksprs(x: Tensor, sparsity_layout: Tensor, sparsity_block_size: int) -> BlksprsTensor:
@@ -35,7 +35,7 @@ def to_sparse(x: Tensor, sparsity_layout: Tensor,
         BlksprsTensor: The block-sparse tensor converted to compressed form.
 
     """
-    x = x.contiguous()
+    x = ensure_contiguous(x)
 
     validate_dimensions(x)
     validate_contiguous(x)
@@ -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 +
@@ -191,7 +187,7 @@ def to_dense(x: BlksprsTensor, sparsity_layout: Tensor,
         Tensor: The block-sparse tensor converted to regular form.
 
     """
-    x = x.contiguous()
+    x = ensure_contiguous(x)
 
     validate_dimensions(x)
     validate_contiguous(x, sparsity_layout)
@@ -339,7 +335,7 @@ def adapt_layout(x: BlksprsTensor, sparsity_layout_from: Tensor, sparsity_block_
         Tensor: The sparsity layout of the resulting output tensor.
 
     """
-    x = x.contiguous()
+    x = ensure_contiguous(x)
 
     validate_dimensions(x)
     validate_contiguous(x, sparsity_layout_from)
@@ -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

@@ -9,7 +9,7 @@ from blksprs.utils.blksprs_tensor import BlksprsTensor
 from blksprs.utils.tools import stride
 from blksprs.utils.autotuning import get_autotune_configs, prune_autotune_configs
 from blksprs.utils.validation import validate_contiguous, validate_dimensions, validate_device, \
-    validate_sparsity, validate_dtype_int, validate_sparsity_block_size
+    validate_sparsity, validate_dtype_int, validate_sparsity_block_size, ensure_contiguous
 
 
 @torch.amp.custom_fwd(device_type="cuda", cast_inputs=torch.float16)
@@ -32,8 +32,7 @@ def gather(src: BlksprsTensor, sparsity_layout_src: Tensor,
         BlksprsTensor: The result of the gather operation as a block-sparse tensor in compressed form.
 
     """
-    src = src.contiguous()
-    idx = idx.contiguous()
+    src, idx = ensure_contiguous(src, idx)
 
     validate_dimensions(src, idx)
     validate_contiguous(src, idx)
@@ -125,17 +124,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) +
@@ -265,8 +260,7 @@ def scatter_reduce(src: BlksprsTensor, sparsity_layout_src: Tensor,
         BlksprsTensor: The result of the scatter operation as a block-sparse tensor in compressed form.
 
     """
-    src = src.contiguous()
-    idx = idx.contiguous()
+    src, idx = ensure_contiguous(src, idx)
 
     validate_dimensions(src, idx)
     validate_contiguous(src, idx)
@@ -374,17 +368,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

@@ -9,7 +9,7 @@ from blksprs.utils.blksprs_tensor import BlksprsTensor
 from blksprs.utils.tools import stride
 from blksprs.utils.autotuning import get_autotune_configs, prune_autotune_configs
 from blksprs.utils.validation import validate_contiguous, validate_dimensions, validate_device, \
-    validate_sparsity, validate_sparsity_block_size, validate_dtype_float
+    validate_sparsity, validate_sparsity_block_size, validate_dtype_float, ensure_contiguous
 
 
 @torch.amp.custom_fwd(device_type="cuda", cast_inputs=torch.float16)
@@ -34,8 +34,7 @@ def matmul(x: BlksprsTensor, sparsity_layout_x: Tensor,
         BlksprsTensor: The result of the matrix multiplication as a block-sparse tensor in compressed form.
 
     """
-    x = x.contiguous()
-    y = y.contiguous()
+    x, y = ensure_contiguous(x, y)
 
     validate_dimensions(x, y)
     validate_contiguous(x, y)
@@ -145,17 +144,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

@@ -9,7 +9,7 @@ from blksprs.utils.blksprs_tensor import BlksprsTensor
 from blksprs.utils.tools import stride
 from blksprs.utils.autotuning import get_autotune_configs, prune_autotune_configs
 from blksprs.utils.validation import validate_contiguous, validate_device, \
-    validate_sparsity_block_size
+    validate_sparsity_block_size, ensure_contiguous
 
 
 @torch.amp.custom_fwd(device_type="cuda", cast_inputs=torch.float16)
@@ -29,8 +29,7 @@ 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()
+    x, y = ensure_contiguous(x, y)
 
     validate_device(x, y)
     validate_contiguous(x, y)
@@ -110,17 +109,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

@@ -8,7 +8,7 @@ from blksprs.utils.autotuning import get_autotune_configs, prune_autotune_config
 from blksprs.utils.blksprs_tensor import BlksprsTensor
 from blksprs.utils.tools import stride
 from blksprs.utils.validation import validate_dimensions, validate_contiguous, validate_device, validate_sparsity, \
-    validate_sparsity_block_size
+    validate_sparsity_block_size, ensure_contiguous
 
 
 @torch.amp.custom_fwd(device_type="cuda", cast_inputs=torch.float32)
@@ -34,7 +34,7 @@ def row_wise_sum(x: BlksprsTensor, sparsity_layout: Tensor, sparsity_block_size:
             of the input and the sparsity layout of the output tensor.
 
     """
-    x = x.contiguous()
+    x = ensure_contiguous(x)
 
     validate_dimensions(x)
     validate_contiguous(x)
@@ -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)
 
@@ -176,7 +176,7 @@ def row_wise_max(x: BlksprsTensor, sparsity_layout: Tensor, sparsity_block_size:
     """
     # TODO Fix for triton bug, see https://github.com/triton-lang/triton/issues/6376, should be fixed with the upcoming 3.4.0 release
     x = torch.where(x == -0.0, torch.tensor(0.0), x)
-    x = x.contiguous()
+    x = ensure_contiguous(x)
 
     validate_dimensions(x)
     validate_contiguous(x)
@@ -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)
 
@@ -311,6 +311,8 @@ def row_wise_add(x: BlksprsTensor, sparsity_layout_x: Tensor, y: Tensor,
             compressed form.
 
     """
+    x = ensure_contiguous(x)
+
     validate_dimensions(x)
     validate_contiguous(x)
     validate_device(x)
@@ -361,7 +363,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 +385,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/partitioning.py

@@ -5,7 +5,7 @@ from torch._library import triton_op
 from blksprs.ops.flow import flow_pull_forward
 from blksprs.utils.blksprs_tensor import BlksprsTensor
 from blksprs.utils.validation import validate_dimensions, validate_contiguous, validate_device, \
-    validate_sparsity, validate_sparsity_block_size
+    validate_sparsity, validate_sparsity_block_size, ensure_contiguous
 
 
 @torch.amp.custom_fwd(device_type="cuda", cast_inputs=torch.float16)
@@ -27,7 +27,7 @@ def split(x: BlksprsTensor, sparsity_layout: Tensor, partitions: int,
         Tensor: The sparsity layout of the output tensor.
 
     """
-    x = x.contiguous()
+    x = ensure_contiguous(x)
 
     validate_dimensions(x)
     validate_contiguous(x)
@@ -132,7 +132,7 @@ def merge(x: BlksprsTensor, sparsity_layout: Tensor, partitions: int,
         Tensor: The sparsity layout of the output tensor.
 
     """
-    x = x.contiguous()
+    x = ensure_contiguous(x)
 
     validate_dimensions(x)
     validate_contiguous(x)

blksprs/ops/repeat.py

@@ -5,7 +5,7 @@ from torch._library import triton_op
 from blksprs.ops.flow import flow_pull_forward, flow_push_forward
 from blksprs.utils.blksprs_tensor import BlksprsTensor
 from blksprs.utils.validation import validate_dimensions, validate_contiguous, validate_device, \
-    validate_sparsity, validate_sparsity_block_size
+    validate_sparsity, validate_sparsity_block_size, ensure_contiguous
 
 
 @torch.amp.custom_fwd(device_type="cuda", cast_inputs=torch.float16)
@@ -36,7 +36,7 @@ def repeat(x: BlksprsTensor, sparsity_layout_x: Tensor, repeats: tuple[int, int,
         Tensor: The sparsity layout of the resulting output tensor.
 
     """
-    x = x.contiguous()
+    x = ensure_contiguous(x)
 
     validate_dimensions(x)
     validate_contiguous(x)
@@ -77,7 +77,7 @@ def repeat_interleave(x: BlksprsTensor, sparsity_layout_x: Tensor, repeats: int,
         Tensor: The sparsity layout of the resulting output tensor.
 
     """
-    x = x.contiguous()
+    x = ensure_contiguous(x)
 
     validate_dimensions(x)
     validate_contiguous(x)
@@ -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

@@ -12,7 +12,7 @@ from blksprs.utils.blksprs_tensor import BlksprsTensor
 from blksprs.utils.tools import stride, ceil_pow2
 from blksprs.utils.autotuning import get_autotune_configs, prune_autotune_configs
 from blksprs.utils.validation import validate_contiguous, validate_dimensions, validate_device, \
-    validate_sparsity, validate_sparsity_block_size, validate_dtype_float_32
+    validate_sparsity, validate_sparsity_block_size, validate_dtype_float_32, ensure_contiguous
 
 
 def softmax(x: BlksprsTensor, sparsity_layout: Tensor, sparsity_block_size: int, flag_fused: bool = True,
@@ -44,7 +44,7 @@ def softmax_regular(x: BlksprsTensor, sparsity_layout: Tensor, sparsity_block_si
         BlksprsTensor: The result of the softmax operation as a block-sparse tensor in compressed form.
 
     """
-    x = x.contiguous()
+    x = ensure_contiguous(x)
 
     validate_dimensions(x)
     validate_contiguous(x)
@@ -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)
@@ -337,7 +335,7 @@ def softmax_fused(x: BlksprsTensor, sparsity_layout: Tensor, sparsity_block_size
         BlksprsTensor: The result of the softmax operation as a block-sparse tensor in compressed form.
 
     """
-    x = x.contiguous()
+    x = ensure_contiguous(x)
 
     validate_dimensions(x)
     validate_contiguous(x)

blksprs/ops/transpose.py

@@ -5,7 +5,7 @@ from torch._library import triton_op
 from blksprs.ops.flow import flow_pull_forward
 from blksprs.utils.blksprs_tensor import BlksprsTensor
 from blksprs.utils.validation import validate_dimensions, validate_contiguous, validate_device, \
-    validate_sparsity, validate_sparsity_block_size
+    validate_sparsity, validate_sparsity_block_size, ensure_contiguous
 
 
 @torch.amp.custom_fwd(device_type="cuda", cast_inputs=torch.float16)
@@ -27,7 +27,7 @@ def transpose(x: BlksprsTensor, sparsity_layout: Tensor,
         Tensor: The sparsity layout of the transposed tensor.
 
     """
-    x = x.contiguous()
+    x = ensure_contiguous(x)
 
     validate_dimensions(x)
     validate_contiguous(x)

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,18 @@
 import torch
 from torch import Tensor
 
+CONTIGUOUS = True
 VALIDATION = True
 
 
+def ensure_contiguous(*tensors: Tensor) -> tuple[Tensor, ...]:
+    if _check_skip_contiguous():
+        return tensors
+
+    transformed = tuple(tensor.contiguous() for tensor in tensors)
+    return transformed[0] if len(transformed) == 1 else transformed
+
+
 def validate_dimensions(*tensors: Tensor, dims=3) -> None:
     if _check_skip_validation():
         return
@@ -124,6 +133,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.6.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: blksprs
-Version: 2.1.4
+Version: 2.1.6
 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.6.dist-info/RECORD

@@ -0,0 +1,23 @@
+blksprs/__init__.py,sha256=F3JKF_GqJrXYtM31r9eT8kSamz8I8bav-6lbQUiLHfA,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=NcBxWWWzMkjQx_fEfh14RWt688X6J82FzDqByAd3Pj4,21405
+blksprs/ops/distribution.py,sha256=pabgyw0m3A4A0osfnOoKffk-b2BKXCn-lC6BU26ocKY,20180
+blksprs/ops/flow.py,sha256=JEGES5ZbMqxR02rwi2Ym4j3VDxkcRxhFO1f-5nNUlM8,7760
+blksprs/ops/matmul.py,sha256=9XPsKbYBw0cdmZY6i4T3Phbx00LXIuA6KI0EIcyGo9U,11584
+blksprs/ops/partitioning.py,sha256=67_a9a5ZpsRmB4BVTOks0stFWp34cb0nk28zQFkXEZc,9985
+blksprs/ops/repeat.py,sha256=Eo7L-TcrrXb_I6xKXLVklp1EuCuA0sfhPaOzw_8y1eU,9080
+blksprs/ops/softmax.py,sha256=YcoZpdC1BdL4zKRQOSjIRtfGgDoQvUZabgNmjbeY8-4,23470
+blksprs/ops/transpose.py,sha256=AyIPuiMAtUAPJPs9eK-Apz6vjZdmnJO9RF6_yH6u6Fk,4097
+blksprs/ops/misc/broadcast_ops.py,sha256=ro7K2ZMOsscxNEp2HY_6efqJ4Wrf-QCFL4NLeDqvah8,5692
+blksprs/ops/misc/row_wise.py,sha256=dfhuXexyFBaNvfZjOt9w3s29ih19JhWIy04_FhUnHgk,19420
+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=XmDMAVSg7SHd7KZswFGU_2kshSTWe0dI6yB5iSGj6cQ,4850
+blksprs-2.1.6.dist-info/METADATA,sha256=jALZxAvt1JAvlQc219KI5mRqHsCq624d0P8LzyLxe9Q,9590
+blksprs-2.1.6.dist-info/WHEEL,sha256=_zCd3N1l69ArxyTb8rzEoP9TpbYXkqRFSNOD5OuxnTs,91
+blksprs-2.1.6.dist-info/top_level.txt,sha256=qyp0IHeY3H2GQA97i4hk_To5rRBS2YcE1HRPSLy04fk,8
+blksprs-2.1.6.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,,