Add simple fused Triton kernel for jagged_sum operator (#2322)

Summary:
Pull Request resolved: #2322

Add Triton kernel implementation to `jagged_sum` operator in TritonBench. This Triton kernel performs a sum along the ragged dimension of a nested tensor of logical dimensions `(B, *, M)`, where `*` is the ragged dimension.  It loads in blocks of the `values` tensor along its last dimension `M`, reduces each block of variable length along its first dimension `*`, and stores each of `B` reductions in an output tensor of shape `(B, M)`.

This Triton kernel is benchmarked against two PyTorch implementations, one which does not pad blocks of variable length and one which does pad.

Reviewed By: davidberard98

Differential Revision: D58549297

fbshipit-source-id: 247f81ed65e010de490114963aac3326e91acb8c

torchbenchmark/operators/jagged_sum/kernels.py

@@ -0,0 +1,155 @@
+import itertools
+
+import triton
+import triton.language as tl
+
+
+BLOCK_SIZES = [2**n for n in range(2, 11, 3)]
+NUM_WARPS = [2, 4, 8]
+NUM_STAGES = [2, 4, 8]
+
+
+@triton.autotune(
+    configs=[
+        triton.Config(
+            {
+                "BLOCK_SIZE_RAGGED": b_r,
+                "BLOCK_SIZE_M": b_m,
+            },
+            num_warps=w,
+            num_stages=s,
+        )
+        for b_r, b_m, w, s in itertools.product(
+            BLOCK_SIZES,  # block sizes on non-reduction dimension
+            BLOCK_SIZES,  # block sizes on reduction dimension
+            NUM_WARPS,  # number of warps
+            NUM_STAGES,  # number of stages
+        )
+    ],
+    key=["M"],
+)
+@triton.jit
+def triton_jagged_sum_kernel_simple_fused_sum_then_buffer(
+    input_ptr_values,  # pointer to input values (2D tensor)
+    input_ptr_offsets,  # pointer to input offsets (1D tensor)
+    output_ptr,  # pointer to output tensor (2D tensor)
+    # matrix dimensions (input)
+    M,  # number of elements in M-th dimension, with logical dimensions (B, *, M)
+    MAX_SEQLEN,  # max length of ragged dimension
+    # block sizes (input)
+    BLOCK_SIZE_RAGGED: tl.constexpr,  # number of elements in ragged dimension per block, with logical dimensions (B, *, M)
+    BLOCK_SIZE_M: tl.constexpr,  # number of elements in M-th dimension per block, with logical dimensions (B, *, M)
+):
+    pid = tl.program_id(axis=0)  # i-th tensor in nested tensor
+    pid_ragged = pid // tl.cdiv(M, BLOCK_SIZE_M)
+    pid_m = pid % tl.cdiv(M, BLOCK_SIZE_M)
+
+    buffer = tl.zeros(
+        (1, BLOCK_SIZE_M), dtype=tl.float32
+    )  # create buffer as a row tensor
+
+    block_start_m = pid_m * BLOCK_SIZE_M
+    offsets_m = block_start_m + tl.arange(0, BLOCK_SIZE_M)
+    mask_m = offsets_m < M
+
+    ragged_start, ragged_end = tl.load(input_ptr_offsets + pid_ragged), tl.load(
+        input_ptr_offsets + (pid_ragged + 1)
+    )  # load start and end offsets for current program, similar to offsets[i] and offsets[i + 1]
+
+    for block_pos in range(
+        0, MAX_SEQLEN, BLOCK_SIZE_RAGGED
+    ):  # loop over ragged dimension, ranging until maximum seqlen
+        block_start_ragged = ragged_start + block_pos  # offset block position by start of current program
+        offsets_ragged = block_start_ragged + tl.arange(0, BLOCK_SIZE_RAGGED)
+        mask_ragged = offsets_ragged < ragged_end
+
+        idxs = (offsets_ragged[:, None] * M) + offsets_m
+        mask = mask_ragged[:, None] & mask_m
+
+        input = tl.load(input_ptr_values + idxs, mask=mask, other=0)
+
+        buffer += tl.sum(input, axis=0)
+
+    buffer_view = buffer.reshape(
+        (BLOCK_SIZE_M,),
+    )  # reshape buffer to 1D, as tl.sum may return a 2D tensor
+
+    output_offsets = offsets_m + (
+        pid_ragged * M
+    )  # output is offset by both ragged dimension and M-th dimension
+    output_mask = output_offsets < (M * (pid_ragged + 1))
+
+    tl.store(output_ptr + output_offsets, buffer_view, mask=output_mask)
+
+
+@triton.autotune(
+    configs=[
+        triton.Config(
+            {
+                "BLOCK_SIZE_RAGGED": b_r,
+                "BLOCK_SIZE_M": b_m,
+            },
+            num_warps=w,
+            num_stages=s,
+        )
+        for b_r, b_m, w, s in itertools.product(
+            BLOCK_SIZES,  # block sizes on non-reduction dimension
+            BLOCK_SIZES,  # block sizes on reduction dimension
+            NUM_WARPS,  # number of warps
+            NUM_STAGES,  # number of stages
+        )
+    ],
+    key=["M"],
+)
+@triton.jit
+def triton_jagged_sum_kernel_simple_fused_buffer_then_sum(
+    input_ptr_values,  # pointer to input values (2D tensor)
+    input_ptr_offsets,  # pointer to input offsets (1D tensor)
+    output_ptr,  # pointer to output tensor (2D tensor)
+    # matrix dimensions (input)
+    M,  # number of elements in M-th dimension, with logical dimensions (B, *, M)
+    MAX_SEQLEN,  # max length of ragged dimension
+    # block sizes (input)
+    BLOCK_SIZE_RAGGED: tl.constexpr,  # number of elements in ragged dimension per block, with logical dimensions (B, *, M)
+    BLOCK_SIZE_M: tl.constexpr,  # number of elements in M-th dimension per block, with logical dimensions (B, *, M)
+):
+    pid = tl.program_id(axis=0)  # i-th tensor in nested tensor
+    pid_ragged = pid // tl.cdiv(M, BLOCK_SIZE_M)
+    pid_m = pid % tl.cdiv(M, BLOCK_SIZE_M)
+
+    buffer = tl.zeros(
+        (BLOCK_SIZE_RAGGED, BLOCK_SIZE_M), dtype=tl.float32
+    )  # create buffer as a row tensor
+
+    block_start_m = pid_m * BLOCK_SIZE_M
+    offsets_m = block_start_m + tl.arange(0, BLOCK_SIZE_M)
+    mask_m = offsets_m < M
+
+    ragged_start, ragged_end = tl.load(input_ptr_offsets + pid_ragged), tl.load(
+        input_ptr_offsets + (pid_ragged + 1)
+    )  # load start and end offsets for current program, similar to offsets[i] and offsets[i + 1]
+
+    for block_pos in range(
+        0, MAX_SEQLEN, BLOCK_SIZE_RAGGED
+    ):  # loop over ragged dimension, ranging until maximum seqlen
+        block_start_ragged = ragged_start + block_pos  # offset block position by start of current program
+        offsets_ragged = block_start_ragged + tl.arange(0, BLOCK_SIZE_RAGGED)
+        mask_ragged = offsets_ragged < ragged_end
+
+        idxs = (offsets_ragged[:, None] * M) + offsets_m
+        mask = mask_ragged[:, None] & mask_m
+
+        buffer += tl.load(input_ptr_values + idxs, mask=mask, other=0)
+
+    buffer_sum = tl.sum(buffer, axis=0)
+
+    buffer_view = buffer_sum.reshape(
+        (BLOCK_SIZE_M,),
+    )  # reshape buffer to 1D, as tl.sum may return a 2D tensor
+
+    output_offsets = offsets_m + (
+        pid_ragged * M
+    )  # output is offset by both ragged dimension and M-th dimension
+    output_mask = output_offsets < (M * (pid_ragged + 1))
+
+    tl.store(output_ptr + output_offsets, buffer_view, mask=output_mask)

torchbenchmark/operators/jagged_sum/operator.py

@@ -15,6 +15,11 @@
     register_metric,
 )
 
+from .kernels import (
+    triton_jagged_sum_kernel_simple_fused_buffer_then_sum,
+    triton_jagged_sum_kernel_simple_fused_sum_then_buffer,
+)
+
 seed = 16
 random.seed(seed)
 torch.manual_seed(seed)
@@ -38,21 +43,57 @@ def parse_op_args(args: List[str]):
         default=0.5,
         help="Average sparsity for nested tensor (float, (0.0-1.0))",
     )
+    parser.add_argument(
+        "--sum-then-buffer",
+        type=int,  # 1: sum then buffer, 0: buffer then sum
+        default=1,
+        help="[Optional] For Triton kernels, determines whether to sum individual blocks then add to a buffer or add to a buffer then sum; 1: sum then buffer, 0: buffer then sum",
+    )
     return parser.parse_args(args)
 
 
+def execute_kernel_simple_fused(x, max_seqlen, sum_then_buffer):
+    B, M = x.shape[0], x.shape[2]
+    grid = lambda meta: ((len(x.offsets()) - 1) * triton.cdiv(M, meta["BLOCK_SIZE_M"]),)
+    kernel_output = torch.zeros((B, M), device=x.device)
+
+    if sum_then_buffer:
+        triton_jagged_sum_kernel_simple_fused_sum_then_buffer[grid](
+            x.values(),
+            x.offsets(),
+            kernel_output,
+            M=M,
+            MAX_SEQLEN=max_seqlen,
+        )
+    else:
+        triton_jagged_sum_kernel_simple_fused_buffer_then_sum[grid](
+            x.values(),
+            x.offsets(),
+            kernel_output,
+            M=M,
+            MAX_SEQLEN=max_seqlen,
+        )
+
+    return kernel_output
+
+
 class Operator(BenchmarkOperator):
 
     DEFAULT_METRICS = ["latency", "accuracy"]
-    use_cuda_graphs = False  # enables GPU/CPU sync (for methods like NestedTensor unbind)
+    use_cuda_graphs = (
+        False  # enables GPU/CPU sync (for methods like NestedTensor unbind)
+    )
 
     def __init__(self, mode: str, device: str, extra_args: Optional[List[str]] = None):
         super().__init__(mode=mode, device=device, extra_args=extra_args)
-        self.sizes = range(4, 10, 2)
+        self.sizes = list(range(2, 8, 2)) + list(
+            range(8, 12)
+        )  # bias towards larger sizes, which are more representative of real-world shapes
 
         args = parse_op_args(self.extra_args)
         self.seqlen = args.seqlen
         self.sparsity = args.sparsity
+        self.sum_then_buffer = args.sum_then_buffer
 
     @register_benchmark(baseline=True)
     def torch_jagged_sum_no_pad(self, x: torch.Tensor):
@@ -75,6 +116,13 @@ def torch_jagged_sum_pad(self, x: torch.Tensor):
             dim=1,
         )  # sum along ragged dimension (dim == 1)
 
+    @register_benchmark()
+    def triton_jagged_sum_no_pad(self, x: torch.Tensor):
+        def _inner():
+            return execute_kernel_simple_fused(x, self.seqlen, self.sum_then_buffer)
+
+        return _inner
+
     def get_x_val(self, example_inputs):
         return len(example_inputs[0])
 
@@ -156,7 +204,7 @@ def gbps(self, fn_name, example_inputs, metrics: BenchmarkOperatorMetrics):
             / metrics.latency
             * GIGABYTES_PER_BYTE
         )
-    
+
     @register_metric(x_only=True)
     def input_shape(
         self, fn_name: str, example_inputs, metrics: BenchmarkOperatorMetrics
@@ -166,3 +214,15 @@ def input_shape(
             "*",
             example_inputs[0].shape[2],
         )  # return (B, '*', M) for each example input
+
+    @register_metric(skip_baseline=True)
+    def best_config(
+        self, fn_name, example_inputs, metrics: BenchmarkOperatorMetrics
+    ) -> str:
+        if self.sum_then_buffer:
+            return dump_autotuner_best_config(
+                triton_jagged_sum_kernel_simple_fused_sum_then_buffer
+            )
+        return dump_autotuner_best_config(
+            triton_jagged_sum_kernel_simple_fused_buffer_then_sum
+        )