Extend support to varying block sizes on both dimensions for 2D matri…

…ces (#2302)

Summary:
Pull Request resolved: #2302

Extend support for reducing across individual dimensions on 2-dimensional matrices by allowing for varying block sizes on both the `M` (first) and `N` (second) dimensions.

The existing kernel performed a simplified reduction, assuming that the entire reduction dimension fit within one thread block. The new kernel implementation removes the need for this assumption, allowing both the reduction and the non-reduction dimensions to fit in multiple thread blocks. This implementation also enables autotuning on block sizes for both the `M` and `N` dimensions.

For 1D results, add a `sum_then_buffer` configuration which decides which kernel configuration to run. `Sum_then_buffer` sums individual blocks of input and adds these sums into a buffer. `Buffer_then_sum` adds blocks of raw input into a buffer, then reduces the buffer.

Reviewed By: davidberard98

Differential Revision: D58313958

fbshipit-source-id: 639ea6b7d7b92f478c0f5669a1cdc0dcb68004e3

torchbenchmark/operators/sum/kernels.py

@@ -44,7 +44,10 @@ def triton_sum_kernel_scalar_result(
 @triton.autotune(
     configs=[
         triton.Config(
-            {"BLOCK_SIZE_NON_REDUCE_DIM": b},
+            {
+                "BLOCK_SIZE_NON_REDUCE_DIM": b,
+                "BLOCK_SIZE_REDUCE_DIM": b,
+            },
             num_warps=w,
         )
         for b, w in itertools.product(
@@ -54,7 +57,7 @@ def triton_sum_kernel_scalar_result(
     key=["M", "N"],
 )
 @triton.jit
-def triton_sum_kernel_1D_result(
+def triton_sum_kernel_1D_result_sum_then_buffer(
     input_ptr,  # pointer to input matrix
     output_ptr,  # pointer to output matrix
     # matrix dimensions (input)
@@ -66,49 +69,128 @@ def triton_sum_kernel_1D_result(
     # reduction dimension
     dim: tl.constexpr,  # dimension along which to sum
 ):
+    """
+    Sum blocks of input using Triton and store in buffer
+    """
+
     pid = tl.program_id(axis=0)  # i-th block of input
 
-    block_start_m, block_start_n = 0, 0
-    offsets_m, offsets_n = None, None
-    if dim == 0:
-        block_start_n = pid * BLOCK_SIZE_REDUCE_DIM
-        # offsets have shape equal to input shape
-        offsets_m = block_start_m + tl.arange(
+    reduce_dim_len = M if dim == 0 else N
+    non_reduce_dim_len = N if dim == 0 else M
+
+    buffer = tl.zeros(
+        (1, BLOCK_SIZE_NON_REDUCE_DIM), dtype=tl.float32
+    )  # create buffer as a row tensor
+
+    block_start_non_reduce_dim = pid * BLOCK_SIZE_NON_REDUCE_DIM
+    offsets_non_reduce_dim = block_start_non_reduce_dim + tl.arange(
+        0, BLOCK_SIZE_NON_REDUCE_DIM
+    )
+    mask_non_reduce_dim = offsets_non_reduce_dim < non_reduce_dim_len
+
+    for block_start_reduce_dim in range(0, reduce_dim_len, BLOCK_SIZE_REDUCE_DIM):
+        offsets_reduce_dim = block_start_reduce_dim + tl.arange(
             0, BLOCK_SIZE_REDUCE_DIM
-        )  # create 1D vector for offsets on M-th dimension
-        offsets_n = block_start_n + tl.arange(
-            0, BLOCK_SIZE_NON_REDUCE_DIM
-        )  # create 1D vector for offsets on N-th dimension
-    elif dim == 1:
-        block_start_m = pid * BLOCK_SIZE_REDUCE_DIM
-        # offsets have shape equal to input shape
-        offsets_m = block_start_m + tl.arange(
-            0, BLOCK_SIZE_NON_REDUCE_DIM
-        )  # create 1D vector for offsets on M-th dimension
-        offsets_n = block_start_n + tl.arange(
+        )
+        mask_reduce_dim = offsets_reduce_dim < reduce_dim_len
+
+        idxs, mask = None, None
+        if dim == 0:
+            idxs = (
+                offsets_reduce_dim[:, None] * non_reduce_dim_len
+            ) + offsets_non_reduce_dim
+            mask = mask_reduce_dim[:, None] & mask_non_reduce_dim
+        elif dim == 1:
+            idxs = (
+                offsets_non_reduce_dim[:, None] * reduce_dim_len
+            ) + offsets_reduce_dim
+            mask = mask_non_reduce_dim[:, None] & mask_reduce_dim
+
+        input = tl.load(input_ptr + idxs, mask=mask, other=mask)
+
+        buffer += tl.sum(input, axis=dim)
+
+    buffer_view = buffer.reshape(
+        (BLOCK_SIZE_NON_REDUCE_DIM,), can_reorder=True
+    )  # reshape buffer to 1D, as tl.sum may return a 2D tensor
+
+    tl.store(output_ptr + offsets_non_reduce_dim, buffer_view, mask=mask_non_reduce_dim)
+
+
+@triton.autotune(
+    configs=[
+        triton.Config(
+            {
+                "BLOCK_SIZE_NON_REDUCE_DIM": b,
+                "BLOCK_SIZE_REDUCE_DIM": b,
+            },
+            num_warps=w,
+        )
+        for b, w in itertools.product(
+            [2, 4, 8, 16], [2, 4, 8]  # block sizes  # number of warps
+        )
+    ],
+    key=["M", "N"],
+)
+@triton.jit
+def triton_sum_kernel_1D_result_buffer_then_sum(
+    input_ptr,  # pointer to input matrix
+    output_ptr,  # pointer to output matrix
+    # matrix dimensions (input)
+    M,  # number of rows
+    N,  # number of columns
+    # block sizes (input)
+    BLOCK_SIZE_NON_REDUCE_DIM: tl.constexpr,  # number of elements in non-reduction dimension per block
+    BLOCK_SIZE_REDUCE_DIM: tl.constexpr,  # number of elements in reduction dimension per block
+    # reduction dimension
+    dim: tl.constexpr,  # dimension along which to sum
+):
+    """
+    Add blocks of input to a buffer and sum the buffer using Triton
+    """
+
+    pid = tl.program_id(axis=0)  # i-th block of input
+
+    reduce_dim_len = M if dim == 0 else N
+    non_reduce_dim_len = N if dim == 0 else M
+
+    buffer = tl.zeros(
+        (BLOCK_SIZE_REDUCE_DIM, BLOCK_SIZE_NON_REDUCE_DIM), dtype=tl.float32
+    )  # create buffer as a 2D tensor
+
+    block_start_non_reduce_dim = pid * BLOCK_SIZE_NON_REDUCE_DIM
+    offsets_non_reduce_dim = block_start_non_reduce_dim + tl.arange(
+        0, BLOCK_SIZE_NON_REDUCE_DIM
+    )
+    mask_non_reduce_dim = offsets_non_reduce_dim < non_reduce_dim_len
+
+    for block_start_reduce_dim in range(0, reduce_dim_len, BLOCK_SIZE_REDUCE_DIM):
+        offsets_reduce_dim = block_start_reduce_dim + tl.arange(
             0, BLOCK_SIZE_REDUCE_DIM
-        )  # create 1D vector for offsets on N-th dimension
+        )
+        mask_reduce_dim = offsets_reduce_dim < reduce_dim_len
 
-    # mask has shape equal to input shape
-    mask_m = offsets_m < M
-    mask_n = offsets_n < N
+        idxs, mask = None, None
+        if dim == 0:
+            idxs = (
+                offsets_reduce_dim[:, None] * non_reduce_dim_len
+            ) + offsets_non_reduce_dim
+            mask = mask_reduce_dim[:, None] & mask_non_reduce_dim
+        elif dim == 1:
+            idxs = (
+                offsets_non_reduce_dim[:, None] * reduce_dim_len
+            ) + offsets_reduce_dim
+            mask = mask_non_reduce_dim[:, None] & mask_reduce_dim
 
-    # create 2D matrices of pointers and masks, using above M and N vectors
-    idxs = (offsets_m[:, None] * N) + offsets_n
-    mask = mask_m[:, None] & mask_n
+        buffer += tl.load(input_ptr + idxs, mask=mask, other=mask)
 
-    # loaded pointers have shape equal to input shape
-    input = tl.load(
-        input_ptr + idxs, mask=mask, other=mask
-    )  # other=mask zeros out masked values from input
+    buffer_sum = tl.sum(buffer, axis=dim)
 
-    output = tl.sum(input, axis=dim)
+    buffer_view = buffer_sum.reshape(
+        (BLOCK_SIZE_NON_REDUCE_DIM,), can_reorder=True
+    )  # reshape buffer to 1D, as tl.sum may return a 2D tensor
 
-    # stored pointers have shape equal to output shape
-    if dim == 0:  # store output along N-th dimension
-        tl.store(output_ptr + offsets_n, output, mask=mask_n)
-    elif dim == 1:  # store output along M-th dimension
-        tl.store(output_ptr + offsets_m, output, mask=mask_m)
+    tl.store(output_ptr + offsets_non_reduce_dim, buffer_view, mask=mask_non_reduce_dim)
 
 
 @triton.autotune(

torchbenchmark/operators/sum/operator.py

@@ -13,7 +13,8 @@
 )
 
 from .kernels import (
-    triton_sum_kernel_1D_result,
+    triton_sum_kernel_1D_result_buffer_then_sum,
+    triton_sum_kernel_1D_result_sum_then_buffer,
     triton_sum_kernel_2D_result_dim_1,
     triton_sum_kernel_scalar_result,
 )
@@ -28,6 +29,12 @@ def parse_op_args(args: List[str]):
         default=None,
         help="[Optional] Dimension(s) on which kernel performs reduction; e.g. --reduce-dim 0, --reduce-dim 0 1",
     )
+    parser.add_argument(
+        "--sum-then-buffer",
+        type=int,  # 1: sum then buffer, 0: buffer then sum
+        default=1,
+        help="[Optional] For 1D results, 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)
 
 
@@ -41,6 +48,7 @@ def __init__(self, mode: str, device: str, extra_args: Optional[List[str]] = Non
         self.reduce_dim = (
             args.reduce_dim if args.reduce_dim else None
         )  # for 2D case, guaranteed to be a list with 1 integer
+        self.sum_then_buffer = args.sum_then_buffer
         self.sizes = range(1, 11)
 
     @register_benchmark()
@@ -61,13 +69,12 @@ def triton_sum(self, x: torch.Tensor):
             )  # race condition in cases where BLOCK_SIZE < n_elements^2
         elif x.dim() == 2 and num_output_dims == 1:
             M, N = x.shape
-            BLOCK_SIZE_M, BLOCK_SIZE_N = triton.next_power_of_2(
-                M
-            ), triton.next_power_of_2(N)
             grid = lambda meta: (
                 max(
                     triton.cdiv(M, meta["BLOCK_SIZE_REDUCE_DIM"]),
                     triton.cdiv(N, meta["BLOCK_SIZE_NON_REDUCE_DIM"]),
+                    triton.cdiv(M, meta["BLOCK_SIZE_NON_REDUCE_DIM"]),
+                    triton.cdiv(N, meta["BLOCK_SIZE_REDUCE_DIM"]),
                 ),
             )
         elif x.dim() == 3 and num_output_dims == 2 and self.reduce_dim[0] == 1:
@@ -94,23 +101,29 @@ def _inner():
             elif kernel_input.dim() == 2 and num_output_dims == 1:
                 if self.reduce_dim[0] == 0:
                     kernel_output = torch.empty(N, device=self.device)
-                    BLOCK_SIZE_REDUCE_DIM = BLOCK_SIZE_M
                 elif self.reduce_dim[0] == 1:
                     kernel_output = torch.empty(M, device=self.device)
-                    BLOCK_SIZE_REDUCE_DIM = BLOCK_SIZE_N
                 else:
                     raise Exception(
                         f"Existing sum Triton kernels do not support reducing input with shape {kernel_input.size} along dimension(s) {self.reduce_dim}"
                     )
 
-                triton_sum_kernel_1D_result[grid](
-                    kernel_input,
-                    kernel_output,
-                    M=M,
-                    N=N,
-                    BLOCK_SIZE_REDUCE_DIM=BLOCK_SIZE_REDUCE_DIM,
-                    dim=self.reduce_dim[0],
-                )
+                if self.sum_then_buffer:
+                    triton_sum_kernel_1D_result_sum_then_buffer[grid](
+                        kernel_input,
+                        kernel_output,
+                        M=M,
+                        N=N,
+                        dim=self.reduce_dim[0],
+                    )
+                else:
+                    triton_sum_kernel_1D_result_buffer_then_sum[grid](
+                        kernel_input,
+                        kernel_output,
+                        M=M,
+                        N=N,
+                        dim=self.reduce_dim[0],
+                    )
             elif (
                 kernel_input.dim() == 3
                 and num_output_dims == 2
@@ -201,6 +214,10 @@ def best_config(
         if example_inputs[0].dim() == 3 and self.reduce_dim and self.reduce_dim[0] == 1:
             return dump_autotuner_best_config(triton_sum_kernel_2D_result_dim_1)
         elif self.reduce_dim and len(self.reduce_dim) < example_inputs[0].dim():
-            return dump_autotuner_best_config(triton_sum_kernel_1D_result)
+            if self.sum_then_buffer:
+                return dump_autotuner_best_config(
+                    triton_sum_kernel_1D_result_sum_then_buffer
+                )
+            return dump_autotuner_best_config(triton_sum_kernel_1D_result_buffer_then_sum)
         else:
             return ""