Add --baseline flag to override baseline

Summary: Per discussion in https://www.internalfb.com/diff/D58503025?transaction_fbid=1683162169162931

Reviewed By: xuzhao9, sijiac

Differential Revision: D58609789

fbshipit-source-id: 22701148ac1e1645c6dab40d1ef021a8e627a734

torchbenchmark/util/triton_op.py

@@ -382,6 +382,12 @@ def parse_args(
  default=None,
  help="Specify one or multiple operator implementations to run."
  )
+ parser.add_argument(
+ "--baseline",
+ type=str,
+ default=None,
+ help="Override default baseline."
+ )
  parser.add_argument(
  "--num-inputs",
  type=int,
@@ -451,6 +457,8 @@ def __init__(self, mode: str, device: str, extra_args: Optional[List[str]]=None)
  self.DEFAULT_METRICS,
  unprocessed_args
  )
+ if self.tb_args.baseline:
+ BASELINE_BENCHMARKS[self.name] = self.tb_args.baseline
  self.required_metrics = list(set(self.tb_args.metrics.split(",")))
  self._only = _split_params_by_comma(self.tb_args.only)
  self._input_id = self.tb_args.input_id