format code and fix bugs from input loader (#272)

Summary:
The Aten op diff broke OSS CI, fixing it.

Remove the old `operator_loader` files.

Pull Request resolved: #272

Reviewed By: FindHao

Differential Revision: D78094345

Pulled By: xuzhao9

fbshipit-source-id: 23c8a6fd311bfc654de9502c067dd3fdccac8152

Author: xuzhao9

docker/tritonbench-nightly.dockerfile

@@ -51,7 +51,7 @@ RUN echo "\
 conda activate base\n\
 export CONDA_HOME=/workspace/miniconda3\n\
 export CUDA_HOME=/usr/local/cuda\n\
-export PATH=/home/runner/bin\${PATH:+:\${PATH}}\n\
+export PATH=\${CUDA_HOME}/bin:/home/runner/bin\${PATH:+:\${PATH}}\n\
 export LD_LIBRARY_PATH=\${CUDA_HOME}/lib64\${LD_LIBRARY_PATH:+:\${LD_LIBRARY_PATH}}\n\
 export LIBRARY_PATH=\${CUDA_HOME}/lib64\${LIBRARY_PATHPATH:+:\${LIBRARY_PATHPATH}}\n" >> /workspace/setup_instance.sh
 

test/test_gpu/main.py

@@ -44,7 +44,11 @@
 # Ops that require special arguments in backwards
 BWD_ARGS_OPS: Dict[str, List[str]] = skip_tests.get("bwd_args", {})
 
-TEST_OPERATORS = set(list_operators_by_collection(op_collection="buck"))
+TEST_OPERATORS = (
+    set(list_operators_by_collection(op_collection="buck"))
+    if is_fbcode()
+    else set(list_operators_by_collection(op_collection="default"))
+)
 
 
 def check_ci_output(op):

test/test_gpu/skip_tests_h100_pytorch.yaml

@@ -26,7 +26,7 @@ fp8_fused_quant_gemm_rowwise:
 gemm:
   # internal only kernels
   - hstu_triton_matmul
-  # TODO: fix PT2 cutlass
+  # pt2 cutlass kernel
   - pt2_cutlass_matmul
 # jagged tests are slow, so disable them in OSS
 jagged_layer_norm:

tritonbench/__init__.py

@@ -1,5 +0,0 @@
-from .operators import list_operators, load_opbench_by_name
-from .operators_collection import (
-    list_operator_collections,
-    list_operators_by_collection,
-)

tritonbench/operator_loader/__init__.py

@@ -1 +1 @@
-from .utils import is_loader_op, get_op_loader_bench_cls_by_name, list_loader_operators
+from .utils import get_op_loader_bench_cls_by_name, is_loader_op, list_loader_operators

tritonbench/operator_loader/aten/__init__.py

@@ -1 +1 @@
-from .loader import list_aten_ops, get_aten_loader_cls_by_name
+from .loader import get_aten_loader_cls_by_name, list_aten_ops

tritonbench/operator_loader/aten/loader.py

@@ -1,18 +1,21 @@
 import argparse
 import os
+import types
+from typing import Dict, Generator, List, Optional
+
 import torch
 import yaml
-import types
-from typing import List, Optional, Generator, Dict
 from tritonbench.utils.triton_op import BenchmarkOperator, register_benchmark
 
 # The config file defines available ATen operators and their corresponding input shapes.
 ATEN_CONFIG_YAML = os.path.join(os.path.dirname(__file__), "config.yaml")
 aten = torch.ops.aten
 
-class AtenOperator(BenchmarkOperator):
 
-    def __init__(self, tb_args: argparse.Namespace, extra_args: Optional[List[str]] = None):
+class AtenOperator(BenchmarkOperator):
+    def __init__(
+        self, tb_args: argparse.Namespace, extra_args: Optional[List[str]] = None
+    ):
         super().__init__(tb_args, extra_args)
         self.aten_op = eval(self.aten_op_name)
         if not self.tb_args.input_loader:
@@ -42,7 +45,8 @@ def get_default_aten_op_input(aten_op_name: str) -> str:
         config = yaml.safe_load(f)
     return config[aten_op_name]
 
-def get_aten_loader_cls_by_name(aten_op_name: str, aten_op_input: Optional[str]=None):
+
+def get_aten_loader_cls_by_name(aten_op_name: str, aten_op_input: Optional[str] = None):
     """
     Return a class generated from the given aten op name and input.
     If input is not provided, use the default input from the config file.
@@ -55,12 +59,19 @@ def get_aten_loader_cls_by_name(aten_op_name: str, aten_op_input: Optional[str]=
     op_name_module.Operator = op_class
     op_class.name = op_cls_name
     op_class.aten_op_name = aten_op_name
-    op_class.aten_op_input = aten_op_input if aten_op_input else get_default_aten_op_input(aten_op_name)
+    op_class.aten_op_input = (
+        aten_op_input if aten_op_input else get_default_aten_op_input(aten_op_name)
+    )
     # register two backends for each aten op: eager and inductor
-    register_benchmark(operator_name=op_cls_name, func_name="eager", baseline=True)(op_class.eager)
-    register_benchmark(operator_name=op_cls_name, func_name="inductor", baseline=False)(op_class.inductor)
+    register_benchmark(operator_name=op_cls_name, func_name="eager", baseline=True)(
+        op_class.eager
+    )
+    register_benchmark(operator_name=op_cls_name, func_name="inductor", baseline=False)(
+        op_class.inductor
+    )
     return op_class
 
+
 def list_aten_ops() -> Dict[str, str]:
     """
     Load all ATen operators from the config file.