Add asymmetric shapes to test sum Triton kernels (#2310)

Summary:
Pull Request resolved: #2310

Add asymmetric shapes to `get_input_iter()` in order to test accuracy and performance of `sum` Triton kernel implementations against PyTorch.

This diff generates tensors with dimensions of different sizes. For example, a 2D asymmetric tensor would have shape `(n, n + 3)`; a 3D asymmetric tensor would have shape `(n, n + 3, n + 6)`.

Reviewed By: jbschlosser

Differential Revision: D58509022

fbshipit-source-id: 98e0eacbef6f1cdd09ebc8289582a2eb70118506

torchbenchmark/operators/sum/operator.py

@@ -1,4 +1,5 @@
 import argparse
+import itertools
 from typing import Callable, Generator, List, Optional, Tuple
 
 import torch
@@ -19,6 +20,9 @@
     triton_sum_kernel_scalar_result,
 )
 
+GIGABYTES_PER_BYTE = 1e-6
+ABSOLUTE_TOLERANCE = 1e-3
+
 
 def parse_op_args(args: List[str]):
     parser = argparse.ArgumentParser()
@@ -132,7 +136,7 @@ def __init__(self, mode: str, device: str, extra_args: Optional[List[str]] = Non
         self.input_dim = args.input_dim
         self.reduce_dim = args.reduce_dim
         self.sum_then_buffer = args.sum_then_buffer
-        self.sizes = range(1, 11)
+        self.sizes = range(1, 11, 2)
 
     @register_benchmark()
     def triton_sum(self, x: torch.Tensor):
@@ -191,11 +195,10 @@ def get_input_iter(self) -> Generator:
             self.input_dim <= 3
         ), f"Existing sum Triton kernels do not support input dimension {self.input_dim}"
 
-        for size in self.get_x_vals():
+        sizes = itertools.product(self.get_x_vals(), repeat=self.input_dim)
+        for size in sizes:
             input_tensor = torch.randn(
-                tuple(
-                    [size for _ in range(self.input_dim)]
-                ),  # tuple with self.input_dim dimensions
+                size,  # tuple with self.input_dim dimensions
                 device=self.device,
                 dtype=self.dtype,
             )
@@ -204,15 +207,15 @@ def get_input_iter(self) -> Generator:
     def _get_accuracy(self, fn: Callable, baseline_fn: Callable) -> bool:
         output = fn()
         baseline_output = baseline_fn()
-        return torch.allclose(output, baseline_output, atol=1e-3)
+        return torch.allclose(output, baseline_output, atol=ABSOLUTE_TOLERANCE)
 
     @register_metric()
     def gbps(self, fn_name, example_inputs, metrics: BenchmarkOperatorMetrics):
         return (
             example_inputs[0].element_size()
             * example_inputs[0].numel()
             / metrics.latency
-            * 1e-6
+            * GIGABYTES_PER_BYTE
         )
 
     @register_metric(skip_baseline=True)
@@ -231,3 +234,9 @@ def best_config(
             return dump_autotuner_best_config(triton_sum_kernel_2D_result_dim_1)
         else:
             return ""
+
+    @register_metric(x_only=True)
+    def input_shape(
+        self, fn_name: str, example_inputs, metrics: BenchmarkOperatorMetrics
+    ):
+        return example_inputs[0].shape  # return (B, M) for example input