Add unit tests on CPU for TritonBench features (#2323)

Summary:
Pull Request resolved: #2323

Add unit tests that run on the CPU to verify the behavior of the following:
- `x_only = True` for metric registration in [`register_metric()`](https://www.internalfb.com/code/fbsource/[731f07681fbbb38750aee3b165137e39fa6cee50]/fbcode/pytorch/benchmark/torchbenchmark/util/triton_op.py?lines=337)
- custom `label` argument for benchmark registration in [`register_benchmark()`](https://www.internalfb.com/code/fbsource/[731f07681fbbb38750aee3b165137e39fa6cee50]/fbcode/pytorch/benchmark/torchbenchmark/util/triton_op.py?lines=316)

Reviewed By: xuzhao9

Differential Revision: D58558868

fbshipit-source-id: 50e35ec3359db02ce86f45b1d44c34dac9c1a03b

torchbenchmark/operators/test_op/__init__.py

@@ -0,0 +1 @@
+from .operator import Operator

torchbenchmark/operators/test_op/operator.py

@@ -0,0 +1,44 @@
+from typing import Generator, List, Optional
+
+import torch
+
+from torchbenchmark.util.triton_op import (
+    BenchmarkOperator,
+    BenchmarkOperatorMetrics,
+    register_benchmark,
+    register_metric,
+)
+
+
+class Operator(BenchmarkOperator):
+
+    DEFAULT_METRICS = ["test_metric"]
+
+    def __init__(self, mode: str, device: str, extra_args: Optional[List[str]] = None):
+        super().__init__(mode=mode, device=device, extra_args=extra_args)
+
+    @register_benchmark(label="new_op_label")
+    def test_op(self, x: torch.Tensor):
+        return lambda: x
+
+    def get_x_val(self, example_inputs):
+        return example_inputs[0].shape
+
+    def get_x_vals(self) -> List[int]:
+        return [2**n for n in [1, 2, 3]]
+
+    def get_input_iter(self) -> Generator:
+        for x in self.get_x_vals():
+            yield (torch.Tensor(torch.randn(x, device=self.device, dtype=self.dtype)),)
+
+    @register_metric(x_only=True)
+    def test_metric(
+        self, fn_name: str, example_inputs, metrics: BenchmarkOperatorMetrics
+    ):
+        return [ex.shape[0] + 2 for ex in example_inputs]
+
+    @register_metric()
+    def test_metric_per_benchmark(
+        self, fn_name: str, example_inputs, metrics: BenchmarkOperatorMetrics
+    ):
+        return [ex.shape[0] + 3 for ex in example_inputs]

torchbenchmark/util/triton_op.py

@@ -468,6 +468,7 @@ def __init__(self, mode: str, device: str, extra_args: Optional[List[str]]=None)
         self._only = _split_params_by_comma(self.tb_args.only)
         self._input_id = self.tb_args.input_id
         self._num_inputs = self.tb_args.num_inputs
+        self.device = device
 
     # Run the post initialization
     def __post__init__(self):