further additions of dynamic benchmarking

Differential Revision: D83537801

Pull Request resolved: #496

Author: adamomainz

tritonbench/utils/dynamic_runner.py

@@ -0,0 +1,85 @@
+from typing import Callable, Dict, Generator, List, Optional
+
+from tritonbench.utils.triton_op import BenchmarkOperator
+
+
+def dynamic_run(
+    benchmarks: Dict[str, Callable],
+    input_iter: Optional[Generator],
+    **kwargs,
+) -> None:
+    """
+    Run a list of benchmarks with a given set of inputs and kwargs.
+    Kwargs in this case are the command-line arguments available in tritonbench.
+
+    Example:
+
+    def triton_add(x, y):
+        ...
+
+    def input_iter():
+        size = 2**12
+        x = torch.rand(size, device="cuda")
+        y = torch.rand(size, device="cuda")
+        yield x,y
+
+    benchmarks = {
+        "triton_add": triton_add,
+        "triton_add2": triton_add,
+    }
+
+    dynamic_run(benchmarks=benchmarks, input_iter=input_iter, benchmark_name="vector_add")
+    """
+
+    # Convert kwargs into a list of command-line arguments
+    arg_list = []
+    for k, v in kwargs.items():
+        key = f"--{k.replace('_', '-')}"
+        arg_list.append(key)
+        arg_list.append(str(v))
+
+    op = BenchmarkOperator(extra_args=arg_list)
+
+    op.set_input_iter(input_iter)
+
+    for k, v in benchmarks.items():
+        op.add_benchmark(bm_func_name=k, bm_callable=v)
+
+    op.run()
+    print(op.output)
+    return op.output
+
+
+def dynamic_run_once(
+    benchmarks: Dict[str, Callable], single_input: Optional[List], **kwargs
+):
+    """
+    Run a list of benchmarks with a given set of inputs and kwargs.
+    Kwargs in this case are the command-line arguments available in tritonbench.
+
+    The single_input is a list of arguments that will be passed to the benchmark function all together
+
+    Example:
+
+    def triton_add(x, y):
+        ...
+
+    benchmarks = {
+        "triton_add": triton_add,
+        "triton_add2": triton_add,
+    }
+    size = 2**12
+    x = torch.rand(size, device="cuda")
+    y = torch.rand(size, device="cuda")
+    dynamic_run_once(benchmarks=benchmarks, single_input=[x, y], benchmark_name="vector_add")
+    """
+
+    def input_iterator(*args):
+        def generator():
+            yield args
+
+        return generator
+
+    input_generator = input_iterator(*single_input)
+    output = dynamic_run(benchmarks, input_generator, **kwargs)
+    return output

tritonbench/utils/triton_op.py

@@ -939,18 +939,22 @@ def input_callable(self):
         )
         self.input_iter = input_iter
         self._available_num_inputs = sum(1 for _ in self.get_input_iter())
-        self._num_inputs = self._available_num_inputs - self._input_id
+        self._num_inputs = self._available_num_inputs - len(self._input_ids)
+        self._input_ids = [i for i in range(0, self._num_inputs)]
 
     def add_benchmark(self, bm_func_name: str, bm_callable: Callable):
+        def _inner(self, *args, **kwargs):
+            return bm_callable(*args, **kwargs)
+
         decorator_kwargs = {
             "operator_name": self.name,
             "func_name": bm_func_name,
             "enabled": True,
         }
-        decorated_func = register_benchmark(**decorator_kwargs)(bm_callable)
+        decorated_func = register_benchmark(**decorator_kwargs)(_inner)
         bound_method = types.MethodType(decorated_func, self)
         setattr(self, bm_func_name or bm_callable.__name__, bound_method)
-        REGISTERED_BENCHMARKS[bm_func_name] = bm_callable
+        REGISTERED_BENCHMARKS[bm_func_name] = _inner
 
     def run(
         self,