[run_utils] Apply run options to config (#449)

Author: xuzhao9

benchmarks/dump_ir/run.py

@@ -3,21 +3,27 @@
 For autotuned operators, we save the IRs of the best kernels.
 """
 
+import argparse
 import os
 from pathlib import Path
-import argparse
-from tritonbench.utils.env_utils import is_fbcode
-from tritonbench.utils.run_utils import run_in_task, run_one_operator
-from tritonbench.operators import list_custom_triton_operators
 
-from typing import List, Dict
+from typing import Dict, List
+
 from libfb.py import parutil
+from tritonbench.operators import list_custom_triton_operators
+from tritonbench.utils.env_utils import is_fbcode
+from tritonbench.utils.run_utils import run_in_task, run_one_operator
 
-METADATA_DIR = parutil.get_file_path("tritonbench/metadata") if is_fbcode() \
+METADATA_DIR = (
+    parutil.get_file_path("tritonbench/metadata")
+    if is_fbcode()
     else Path(__file__).parent.parent.parent.joinpath("tritonbench/metadata")
+)
 
 OSS_CUSTOM_TRITON_YAML = os.path.join(METADATA_DIR, "oss_triton_operators.yaml")
-INTERNAL_CUSTOM_TRITON_YAML = os.path.join(METADATA_DIR, "fb/internal_triton_operators.yaml")
+INTERNAL_CUSTOM_TRITON_YAML = os.path.join(
+    METADATA_DIR, "fb/internal_triton_operators.yaml"
+)
 
 
 def get_parser():
@@ -29,25 +35,38 @@ def get_parser():
         help="Output directory to save the IRs",
     )
     parser.add_argument(
-        "--run-in-task",
-        action="store_true",
-        help="indicate running in task."
+        "--run-in-task", action="store_true", help="indicate running in task."
     )
     return parser
 
 
 def run_operator(op: str, subop: List[str], output_dir: str):
     """Run a Tritonbench operator and save its IR to the specified directory"""
-    opbench_args = ["--run-in-task", "--op", op, "--only", ",".join(subop), "--dump-ir", output_dir]
+    opbench_args = [
+        "--run-in-task",
+        "--op",
+        op,
+        "--only",
+        ",".join(subop),
+        "--dump-ir",
+        output_dir,
+    ]
     run_in_task(op, opbench_args)
 
+
 if __name__ == "__main__":
     parser = get_parser()
     args, extra_args = parser.parse_known_args()
     if args.run_in_task:
         run_one_operator(extra_args, with_bwd=True)
         exit(0)
-    custom_triton_op_yamls = [OSS_CUSTOM_TRITON_YAML, INTERNAL_CUSTOM_TRITON_YAML] if is_fbcode() else [OSS_CUSTOM_TRITON_YAML]
-    operators: Dict[str, List[str]] = list_custom_triton_operators(custom_triton_op_yamls)
+    custom_triton_op_yamls = (
+        [OSS_CUSTOM_TRITON_YAML, INTERNAL_CUSTOM_TRITON_YAML]
+        if is_fbcode()
+        else [OSS_CUSTOM_TRITON_YAML]
+    )
+    operators: Dict[str, List[str]] = list_custom_triton_operators(
+        custom_triton_op_yamls
+    )
     [run_operator(op, operators[op].keys(), args.output_dir) for op in operators]
     print(f"[tritonbench][dump_ir] Result saved to {args.output_dir}")

run.py

@@ -102,7 +102,7 @@ def run(args: List[str] = []):
     if args == []:
         args = sys.argv[1:]
     if config := os.environ.get("TRITONBENCH_RUN_CONFIG", None):
-        run_config(config)
+        run_config(config, args)
         return
 
     # Log the tool usage

tritonbench/operators/layer_norm/operator.py

@@ -43,6 +43,7 @@ def parse_op_args(args: List[str]):
 
 try:
     from quack.quack_layernorm import layernorm as quack_layernorm
+
     HAS_QUACK_KERNEL = True
 except ModuleNotFoundError:
     HAS_QUACK_KERNEL = False

tritonbench/utils/run_utils.py

@@ -98,15 +98,19 @@ def get_github_env() -> Dict[str, str]:
     return out
 
 
-def run_config(config_file: str):
+def run_config(config_file: str, args: List[str]):
     assert Path(config_file).exists(), f"Config file {config_file} must exist."
     with open(config_file, "r") as fp:
         config = yaml.safe_load(fp)
     for benchmark_name in config:
         benchmark_config = config[benchmark_name]
         op_name = benchmark_config["op"]
-        op_args = benchmark_config["args"].split(" ")
+        op_args = benchmark_config["args"].split(" ") + args
         env_string = benchmark_config.get("envs", None)
+        disabled = benchmark_config.get("disabled", False)
+        if disabled:
+            logger.info(f"Skipping disabled benchmark {benchmark_name}.")
+            continue
         extra_envs = {}
         if env_string:
             for env_part in env_string.split(" "):