Cleanup gbps metrics

Summary: gbps is a user-custom metric, not built-in metric, we are unifying its API here.

Reviewed By: FindHao, adamomainz

Differential Revision: D78163498

fbshipit-source-id: 3bdca98094b2ad1b2e1041cfb593c538562a8dc8

Author: xuzhao9

tritonbench/operators/gather_gemv/operator.py

@@ -5,14 +5,9 @@
 """
 
 import argparse
-import csv
-import os
-import statistics
-from typing import Any, Callable, Generator, List, Optional
+from typing import Callable, Generator, List, Optional
 
-import numpy
 import torch
-import triton
 from torch._dynamo.testing import rand_strided
 
 from tritonbench.utils.triton_op import (
@@ -27,7 +22,7 @@
 
 class Operator(BenchmarkOperator):
     @register_metric()
-    def gbps(self, fn_name, example_inputs, metrics: BenchmarkOperatorMetrics):
+    def gbps(self, fn, example_inputs, metrics: BenchmarkOperatorMetrics):
         arg0_1, arg1_1, arg2_1 = example_inputs
         return (
             2

tritonbench/operators/jagged_layer_norm/operator.py

@@ -180,7 +180,7 @@ def _get_accuracy(self, fn: Callable, baseline_fn: Callable) -> bool:
         )
 
     @register_metric(skip_baseline=True)
-    def gbps(self, fn_name, example_inputs, metrics: BenchmarkOperatorMetrics):
+    def gbps(self, fn, example_inputs, metrics: BenchmarkOperatorMetrics):
         return (
             example_inputs[0].element_size()
             * example_inputs[0].numel()

tritonbench/operators/jagged_mean/operator.py

@@ -1,8 +1,5 @@
 import argparse
-import itertools
-import math
 import os
-import random
 from typing import Callable, Generator, List, Optional, Tuple
 
 import torch
@@ -229,7 +226,7 @@ def _get_accuracy(self, fn: Callable, baseline_fn: Callable) -> bool:
         )
 
     @register_metric()
-    def gbps(self, fn_name, example_inputs, metrics: BenchmarkOperatorMetrics):
+    def gbps(self, fn, example_inputs, metrics: BenchmarkOperatorMetrics):
         return (
             example_inputs[0].element_size()
             * example_inputs[0].numel()

tritonbench/operators/jagged_softmax/operator.py

@@ -213,7 +213,7 @@ def _get_accuracy(self, fn: Callable, baseline_fn: Callable) -> bool:
         )
 
     @register_metric(skip_baseline=True)
-    def gbps(self, fn_name, example_inputs, metrics: BenchmarkOperatorMetrics):
+    def gbps(self, fn, example_inputs, metrics: BenchmarkOperatorMetrics):
         return (
             example_inputs[0].element_size()
             * example_inputs[0].numel()

tritonbench/operators/jagged_sum/operator.py

@@ -227,7 +227,7 @@ def accuracy(self, fn: Callable, baseline_fn: Callable) -> bool:
         )
 
     @register_metric()
-    def gbps(self, fn_name, example_inputs, metrics: BenchmarkOperatorMetrics):
+    def gbps(self, fn, example_inputs, metrics: BenchmarkOperatorMetrics):
         return (
             example_inputs[0].element_size()
             * example_inputs[0].numel()

tritonbench/operators/layer_norm/operator.py

@@ -88,7 +88,7 @@ def get_x_val(self, args):
         return N
 
     @register_metric()
-    def gbps(self, fn_name, args, metrics: BenchmarkOperatorMetrics) -> float:
+    def gbps(self, fn, args, metrics: BenchmarkOperatorMetrics) -> float:
         x = args[0]
         base = x.numel() * x.element_size() / metrics.latency * 1e-6
         return {

tritonbench/operators/low_mem_dropout/operator.py

@@ -15,7 +15,7 @@
 
 class Operator(BenchmarkOperator):
     @register_metric()
-    def gbps(self, fn_name, example_inputs, metrics: BenchmarkOperatorMetrics):
+    def gbps(self, fn, example_inputs, metrics: BenchmarkOperatorMetrics):
         return (
             3
             * example_inputs[1].element_size()

tritonbench/operators/softmax/operator.py

@@ -122,7 +122,7 @@ def get_x_val(self, example_inputs):
         return [shape[0], shape[1]]
 
     @register_metric()
-    def gbps(self, fn_name, example_inputs, metrics: BenchmarkOperatorMetrics) -> float:
+    def gbps(self, fn, example_inputs, metrics: BenchmarkOperatorMetrics) -> float:
         return (
             2
             * example_inputs[0].nelement()

tritonbench/operators/sum/operator.py

@@ -295,7 +295,7 @@ def _get_accuracy(self, fn: Callable, baseline_fn: Callable) -> bool:
         )
 
     @register_metric()
-    def gbps(self, fn_name, example_inputs, metrics: BenchmarkOperatorMetrics):
+    def gbps(self, fn, example_inputs, metrics: BenchmarkOperatorMetrics):
         return (
             example_inputs[0].element_size()
             * example_inputs[0].numel()

tritonbench/operators/vector_add/operator.py

@@ -17,7 +17,7 @@ class Operator(BenchmarkOperator):
     DEFAULT_METRICS = ["latency", "gbps"]
 
     @register_metric()
-    def gbps(self, fn_name, example_inputs, metrics: BenchmarkOperatorMetrics):
+    def gbps(self, fn, example_inputs, metrics: BenchmarkOperatorMetrics):
         def normalize(lat):
             return (
                 3