blksprs 1.11__py3-none-any.whl → 2.0__py3-none-any.whl

blksprs/__init__.py

@@ -1,5 +1,7 @@
 from blksprs.utils.blksprs_tensor import BlksprsTensor
 
+__version__ = "2.0"
+
 
 class ops:
     from blksprs.ops.conversion import to_dense, to_sparse, from_blksprs, to_blksprs, adapt_layout
@@ -18,19 +20,16 @@ class ops:
 class layouting:
     from blksprs.layouting.distribution_layout import build_distribution_layout
     from blksprs.layouting.sparsity_layout import build_sparsity_layout, build_sparsity_layout_adaption, \
-        build_sparsity_layout_matmul, build_sparsity_layout_matmul_fast
-    from blksprs.utils.layout_utils import build_full_sparsity_layout
+        build_sparsity_layout_matmul, build_sparsity_layout_matmul_fast, build_sparsity_layout_full
 
 
 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_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, \
-            validate_triton_block_size
+            validate_sparsity_block_size

blksprs/layouting/distribution_layout.py

@@ -1,17 +1,23 @@
+import typing
+
 import torch
 import triton
 from torch import Tensor
+from torch._library import triton_op
+from torch._library.triton import wrap_triton
 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, \
+from blksprs.utils.tools import stride
+from blksprs.utils.autotuning import get_autotune_configs, prune_autotune_configs
+from blksprs.utils.validation import validate_dimensions, validate_device, \
     validate_contiguous
 
 
+@torch.amp.custom_fwd(device_type="cuda", cast_inputs=torch.float16)
 def build_distribution_layout(indices: BlksprsTensor, sparsity_layout_indices: Tensor,
                               dim: int, size_target: torch.Size,
-                              sparsity_block_size: int, triton_block_size: int = None) -> Tensor:
+                              sparsity_block_size: int) -> Tensor:
     """Builds the sparsity layout of either the source of a gather or the target of a scatter operation.
 
     Args:
@@ -20,7 +26,6 @@ def build_distribution_layout(indices: BlksprsTensor, sparsity_layout_indices: T
         dim (int): The dimension along which the operation is conducted.
         size_target (torch.Size): The size of the block-sparse target tensor in regular form.
         sparsity_block_size (int): The size of the sparsity blocks.
-        triton_block_size (int, optional): The block size to use for the triton kernel (default ``None``).
 
     Returns:
         Tensor: The sparsity layout of the source or target tensor.
@@ -34,49 +39,58 @@ def build_distribution_layout(indices: BlksprsTensor, sparsity_layout_indices: T
 
     adjusted_dim = dim % 3
 
-    output = torch.zeros(size_target[0], size_target[1] // sparsity_block_size, size_target[2] // sparsity_block_size,
-                         dtype=torch.bool, device=indices.device)
-
-    i_b, i_r, i_c = indices.size()
-    i_b_s, i_r_s, i_c_s = stride(indices)
-    s_lut_i_r, s_lut_i_c = sparsity_lut_i.size()
-    s_lut_i_r_s, s_lut_i_c_s = stride(sparsity_lut_i)
-    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)
-
-    validate_triton_block_size(triton_block_size, sparsity_block_size)
-
-    triton_grid = lambda meta: [i_b,
-                                triton.cdiv(i_r, meta["TRITON_BLOCK_SIZE"]),
-                                triton.cdiv(i_c, meta["TRITON_BLOCK_SIZE"])]
-
-    (kernel_distribution_layout[triton_grid]
-     (indices,
-      i_b, i_b_s, i_r_s, i_c_s,
-      sparsity_lut_i,
-      s_lut_i_r, s_lut_i_r_s, s_lut_i_c_s,
-      adjusted_dim,
-      output,
-      o_b, o_b_s, o_r_s, o_c_s,
-      sparsity_block_size,
-      triton_block_size))
-
-    return output
-
-
+    return build_distribution_layout_operation(indices, sparsity_lut_i, adjusted_dim, size_target, sparsity_block_size)
+
+
+@triton_op("blksprs::build_distribution_layout", mutates_args={})
+def build_distribution_layout_operation(indices: Tensor, sparsity_lut_i: Tensor,
+                                        adjusted_dim: int, size_target: typing.List[int],
+                                        sparsity_block_size: int) -> Tensor:
+    with torch.no_grad():
+        output = torch.zeros(size_target[0], size_target[1] // sparsity_block_size,
+                             size_target[2] // sparsity_block_size,
+                             dtype=torch.bool, device=indices.device)
+
+        i_b, i_r, i_c = indices.size()
+        i_b_s, i_r_s, i_c_s = stride(indices)
+        s_lut_i_r, s_lut_i_c = sparsity_lut_i.size()
+        s_lut_i_r_s, s_lut_i_c_s = stride(sparsity_lut_i)
+        o_b, o_r, o_c = output.size()
+        o_b_s, o_r_s, o_c_s = stride(output)
+
+        triton_grid = lambda meta: [i_b,
+                                    triton.cdiv(i_r, meta["TRITON_BLOCK_SIZE"]),
+                                    triton.cdiv(i_c, meta["TRITON_BLOCK_SIZE"])]
+
+        (wrap_triton(build_distribution_layout_kernel)[triton_grid]
+         (indices,
+          i_b, i_b_s, i_r_s, i_c_s,
+          sparsity_lut_i,
+          s_lut_i_r, s_lut_i_r_s, s_lut_i_c_s,
+          adjusted_dim,
+          output,
+          o_b, o_b_s, o_r_s, o_c_s,
+          sparsity_block_size))
+
+        return output
+
+
+@triton.autotune(
+    configs=get_autotune_configs(),
+    key=["sparsity_block_size"],
+    prune_configs_by={"early_config_prune": prune_autotune_configs},
+    reset_to_zero=["o"]
+)
 @triton.jit
-def kernel_distribution_layout(i,
-                               i_b, i_b_s, i_r_s, i_c_s,
-                               s_lut_i,
-                               s_lut_i_r, s_lut_i_r_s, s_lut_i_c_s,
-                               dim,
-                               o,
-                               o_b, o_b_s, o_r_s, o_c_s,
-                               sparsity_block_size,
-                               TRITON_BLOCK_SIZE: tl.constexpr) -> None:
+def build_distribution_layout_kernel(i,
+                                     i_b, i_b_s, i_r_s, i_c_s,
+                                     s_lut_i,
+                                     s_lut_i_r, s_lut_i_r_s, s_lut_i_c_s,
+                                     dim,
+                                     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)
@@ -98,7 +112,8 @@ def kernel_distribution_layout(i,
     blk_i_idx = (pid_blk * i_b_s +
                  ((pid_row * TRITON_BLOCK_SIZE + tl.arange(0, TRITON_BLOCK_SIZE)) * i_r_s)[:, None] +
                  ((pid_col * TRITON_BLOCK_SIZE + tl.arange(0, TRITON_BLOCK_SIZE)) * i_c_s)[None, :])
-    blk_i_msk = (blk_i_idx >= 0 and blk_i_idx < i_b * i_b_s)
+    blk_i_msk = (blk_i_idx >= 0 and
+                 blk_i_idx < i_b * i_b_s)
     blk_i = tl.load(i + blk_i_idx, mask=blk_i_msk)
 
     dst_bat_idx = tl.full((TRITON_BLOCK_SIZE, TRITON_BLOCK_SIZE), spa_bat_i, dtype=tl.int32)
@@ -116,5 +131,6 @@ def kernel_distribution_layout(i,
     blk_o_idx = ((dst_bat_idx * o_b_s) +
                  (dst_row_idx * o_r_s) +
                  (dst_col_idx * o_c_s))
-    blk_o_msk = (blk_o_idx >= 0 and blk_o_idx < o_b * o_b_s)
+    blk_o_msk = (blk_o_idx >= 0 and
+                 blk_o_idx < o_b * o_b_s)
     tl.store(o + blk_o_idx, blk_v, mask=blk_o_msk)

blksprs/layouting/sparsity_layout.py

@@ -3,21 +3,23 @@ import math
 import torch
 import triton
 from torch import Tensor
+from torch._library.triton import wrap_triton, triton_op
 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, \
+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_dimensions, validate_device, \
     validate_contiguous, validate_sparsity, validate_sparsity_block_size
 
 
-def build_sparsity_layout(x: Tensor, sparsity_block_size: int, triton_block_size: int = None) -> Tensor:
+@torch.amp.custom_fwd(device_type="cuda", cast_inputs=torch.float16)
+def build_sparsity_layout(x: Tensor, sparsity_block_size: int) -> Tensor:
     """Builds the sparsity layout of a dense tensor in regular form covering its sparse blocks.
 
     Args:
         x (Tensor): A block-sparse (or dense) tensor in regular form.
         sparsity_block_size (int): The size of the sparsity blocks.
-        triton_block_size (int, optional): The block size to use for the triton kernel (default ``None``).
 
     Returns:
         Tensor: The sparsity layout of the input block-sparse (or dense) tensor.
@@ -27,41 +29,47 @@ def build_sparsity_layout(x: Tensor, sparsity_block_size: int, triton_block_size
     validate_contiguous(x)
     validate_device(x)
 
-    output = torch.zeros(x.size(0), x.size(1) // sparsity_block_size, x.size(2) // sparsity_block_size,
-                         dtype=torch.bool, device=x.device)
+    return build_sparsity_layout_operation(x, sparsity_block_size)
 
-    x_b, x_r, x_c = x.size()
-    x_b_s, x_r_s, x_c_s = stride(x)
-    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_op("blksprs::build_sparsity_layout", mutates_args={})
+def build_sparsity_layout_operation(x: Tensor, sparsity_block_size: int) -> Tensor:
+    with torch.no_grad():
+        output = torch.zeros(x.size(0), x.size(1) // sparsity_block_size, x.size(2) // sparsity_block_size,
+                             dtype=torch.bool, device=x.device)
 
-    validate_triton_block_size(triton_block_size, sparsity_block_size)
+        x_b, x_r, x_c = x.size()
+        x_b_s, x_r_s, x_c_s = stride(x)
+        o_b, o_r, o_c = output.size()
+        o_b_s, o_r_s, o_c_s = stride(output)
 
-    triton_grid = lambda meta: [x_b,
-                                triton.cdiv(x_r, meta["TRITON_BLOCK_SIZE"]),
-                                triton.cdiv(x_c, meta["TRITON_BLOCK_SIZE"])]
+        triton_grid = lambda meta: [x_b,
+                                    triton.cdiv(x_r, meta["TRITON_BLOCK_SIZE"]),
+                                    triton.cdiv(x_c, meta["TRITON_BLOCK_SIZE"])]
 
-    (kernel_sparsity_layout[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,
-      sparsity_block_size,
-      triton_block_size))
+        (wrap_triton(build_sparsity_layout_kernel)[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,
+          sparsity_block_size))
 
-    return output
+        return output
 
 
+@triton.autotune(
+    configs=get_autotune_configs(),
+    key=["sparsity_block_size"],
+    prune_configs_by={"early_config_prune": prune_autotune_configs},
+    reset_to_zero=["o"]
+)
 @triton.jit
-def kernel_sparsity_layout(x,
-                           x_b, x_b_s, x_r_s, x_c_s,
-                           o,
-                           o_b, o_b_s, o_r_s, o_c_s,
-                           sparsity_block_size,
-                           TRITON_BLOCK_SIZE: tl.constexpr) -> None:
+def build_sparsity_layout_kernel(x,
+                                 x_b, x_b_s, x_r_s, x_c_s,
+                                 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_bat = tl.program_id(axis=0)
     pid_row = tl.program_id(axis=1)
@@ -71,7 +79,8 @@ def kernel_sparsity_layout(x,
     blk_x_idx = (pid_bat * 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_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)
 
     # Store sparsity layout value
@@ -83,9 +92,9 @@ def kernel_sparsity_layout(x,
         tl.store(o + blk_o_idx, 1, mask=blk_o_msk)
 
 
+@torch.amp.custom_fwd(device_type="cuda", cast_inputs=torch.float16)
 def build_sparsity_layout_adaption(x: BlksprsTensor, sparsity_layout_from: Tensor,
-                                   sparsity_block_size_from: int, sparsity_block_size_to: int,
-                                   triton_block_size: int = None) -> Tensor:
+                                   sparsity_block_size_from: int, sparsity_block_size_to: int) -> Tensor:
     """Builds the sparsity layout of a block-sparse tensor in compressed form if a different sparsity block size were
         used.
         
@@ -94,7 +103,6 @@ def build_sparsity_layout_adaption(x: BlksprsTensor, sparsity_layout_from: Tenso
         sparsity_layout_from (Tensor): The sparsity layout of the input block-sparse tensor.
         sparsity_block_size_from (int): The size of the sparsity blocks of the input tensor.
         sparsity_block_size_to (int): The desired size of the sparsity blocks for the resulting layout.
-        triton_block_size (int, optional): The block size to use for the triton kernel (default ``None``).
 
     Returns:
         Tensor: The sparsity layout in regular form using the new sparsity block size of the input block-sparse tensor
@@ -107,54 +115,62 @@ def build_sparsity_layout_adaption(x: BlksprsTensor, sparsity_layout_from: Tenso
     validate_sparsity(sparsity_block_size_from, (x, sparsity_layout_from))
     validate_sparsity_block_size(sparsity_block_size_from, x)
     validate_sparsity_block_size(sparsity_block_size_to)
-    min_sparsity_block_size = min(sparsity_block_size_from, sparsity_block_size_to)
-    validate_triton_block_size(triton_block_size, min_sparsity_block_size)
 
     sparsity_lut = torch.nonzero(sparsity_layout_from).contiguous()
 
     validate_contiguous(sparsity_layout_from, sparsity_lut)
 
-    o_b = sparsity_layout_from.size(0)
-    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)
+    return build_sparsity_layout_adaption_operation(x, sparsity_layout_from, sparsity_lut,
+                                                    sparsity_block_size_from, sparsity_block_size_to)
 
-    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 = stride(x)
-    s_lut_r, s_lut_c = sparsity_lut.size()
-    s_lut_r_s, s_lut_c_s = stride(sparsity_lut)
-    o_b_s, o_r_s, o_c_s = stride(output)
+@triton_op("blksprs::build_sparsity_layout_adaption", mutates_args={})
+def build_sparsity_layout_adaption_operation(x: Tensor, sparsity_layout_from: Tensor, sparsity_lut: Tensor,
+                                             sparsity_block_size_from: int, sparsity_block_size_to: int) -> Tensor:
+    with torch.no_grad():
+        o_b = sparsity_layout_from.size(0)
+        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)
 
-    if triton_block_size is None:
-        triton_block_size = get_triton_block_size(sparsity_block_size_from)
+        output = torch.zeros(o_b, o_r, o_c, dtype=torch.bool, device=x.device)
 
-    triton_grid = lambda meta: [x_b,
-                                triton.cdiv(x_r, meta["TRITON_BLOCK_SIZE"]),
-                                triton.cdiv(x_c, meta["TRITON_BLOCK_SIZE"])]
+        x_b, x_r, x_c = x.size()
+        x_b_s, x_r_s, x_c_s = stride(x)
+        s_lut_r, s_lut_c = sparsity_lut.size()
+        s_lut_r_s, s_lut_c_s = stride(sparsity_lut)
+        o_b_s, o_r_s, o_c_s = stride(output)
 
-    (kernel_sparsity_layout_adaption[triton_grid]
-     (x,
-      x_b, x_b_s, x_r_s, x_c_s,
-      sparsity_lut, s_lut_r, s_lut_r_s, s_lut_c_s,
-      output,
-      o_b, o_b_s, o_r_s, o_c_s,
-      sparsity_block_size_from,
-      sparsity_block_size_to,
-      triton_block_size))
+        triton_grid = lambda meta: [x_b,
+                                    triton.cdiv(x_r, meta["TRITON_BLOCK_SIZE"]),
+                                    triton.cdiv(x_c, meta["TRITON_BLOCK_SIZE"])]
 
-    return output
+        (wrap_triton(build_sparsity_layout_adaption_kernel)[triton_grid]
+         (x,
+          x_b, x_b_s, x_r_s, x_c_s,
+          sparsity_lut, s_lut_r, s_lut_r_s, s_lut_c_s,
+          output,
+          o_b, o_b_s, o_r_s, o_c_s,
+          sparsity_block_size_from,
+          sparsity_block_size_to))
 
+        return output
 
+
+@triton.autotune(
+    configs=get_autotune_configs(),
+    key=["sparsity_block_size_from", "sparsity_block_size_to"],
+    prune_configs_by={"early_config_prune": prune_autotune_configs_conversion},
+    reset_to_zero=["o"]
+)
 @triton.jit
-def kernel_sparsity_layout_adaption(x,
-                                    x_b, x_b_s, x_r_s, x_c_s,
-                                    s_lut, s_lut_r, s_lut_r_s, s_lut_c_s,
-                                    o,
-                                    o_b, o_b_s, o_r_s, o_c_s,
-                                    sparsity_block_size_from,
-                                    sparsity_block_size_to,
-                                    TRITON_BLOCK_SIZE: tl.constexpr) -> None:
+def build_sparsity_layout_adaption_kernel(x,
+                                          x_b, x_b_s, x_r_s, x_c_s,
+                                          s_lut, s_lut_r, s_lut_r_s, s_lut_c_s,
+                                          o,
+                                          o_b, o_b_s, o_r_s, o_c_s,
+                                          sparsity_block_size_from,
+                                          sparsity_block_size_to,
+                                          TRITON_BLOCK_SIZE: tl.constexpr) -> None:
     # Get triton block indices
     pid_blk = tl.program_id(axis=0)
     pid_row = tl.program_id(axis=1)
@@ -177,21 +193,23 @@ def kernel_sparsity_layout_adaption(x,
     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_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)
 
     # Store sparsity layout value
     if tl.min(blk_x) != 0 or tl.max(blk_x) != 0:
         blk_o_idx = ((spa_bat * o_b_s) +
-                     (((spa_row * sparsity_block_size_from + pid_row * TRITON_BLOCK_SIZE)
+                     (((pid_row * TRITON_BLOCK_SIZE + spa_row * sparsity_block_size_from)
                        // sparsity_block_size_to) * o_r_s) +
-                     (((spa_col * sparsity_block_size_from + pid_col * TRITON_BLOCK_SIZE)
+                     (((pid_col * TRITON_BLOCK_SIZE + spa_col * sparsity_block_size_from)
                        // sparsity_block_size_to) * o_c_s))
         blk_o_msk = (blk_o_idx >= 0 and blk_o_idx < o_b * o_b_s)
         tl.store(o + blk_o_idx, 1, mask=blk_o_msk)
 
 
-def build_sparsity_layout_matmul(sparsity_layout_x: Tensor, sparsity_layout_y: Tensor):
+@torch.amp.custom_fwd(device_type="cuda", cast_inputs=torch.float16)
+def build_sparsity_layout_matmul(sparsity_layout_x: Tensor, sparsity_layout_y: Tensor) -> Tensor:
     """Builds the precise sparsity layout of the result of a matrix multiplication between the two input tensors.
 
     Args:
@@ -205,6 +223,7 @@ def build_sparsity_layout_matmul(sparsity_layout_x: Tensor, sparsity_layout_y: T
     return torch.matmul(sparsity_layout_x.to(torch.float), sparsity_layout_y.to(torch.float)).to(torch.bool)
 
 
+@torch.amp.custom_fwd(device_type="cuda", cast_inputs=torch.float16)
 def build_sparsity_layout_matmul_fast(sparsity_layout_x: Tensor, sparsity_layout_y: Tensor):
     """Builds the approximate sparsity layout of the result of a matrix multiplication between the two input tensors.
 
@@ -225,3 +244,8 @@ def build_sparsity_layout_matmul_fast(sparsity_layout_x: Tensor, sparsity_layout
     sparsity_layout_y_slice = torch.max(sparsity_layout_y, dim=-2).values.unsqueeze(1)
 
     return torch.logical_or(sparsity_layout_x_slice, sparsity_layout_y_slice)
+
+
+def build_sparsity_layout_full(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)