Load input shapes from highway_self_gating

Summary: Load input shapes from highway_self_gating traces.

Differential Revision: D83660588

Author: xuzhao9

tritonbench/data/__init__.py

@@ -1 +1,6 @@
-from .loader import get_input_loader
+from .loader import (
+    get_input_loader,
+    INPUT_CONFIG_DIR,
+    INTERNAL_INPUT_CONFIG_DIR,
+    SUPPORTED_INPUT_OPS,
+)

tritonbench/data/loader.py

@@ -1,16 +1,29 @@
 import importlib
+from pathlib import Path
 from typing import Any
 
+from tritonbench.utils.env_utils import is_fbcode
+
+SUPPORTED_INPUT_OPS = ["highway_self_gating"]
+
+INPUT_CONFIG_DIR = Path(__file__).parent.parent.joinpath("input_configs")
+INTERNAL_INPUT_CONFIG_DIR = (
+    importlib.resources.files("tritonbench.data.input_configs.fb")
+    if is_fbcode()
+    else None
+)
+
 
 def get_input_loader(tritonbench_op: Any, op: str, input: str):
     if hasattr(tritonbench_op, "aten_op_name"):
-        loader_type = "aten"
+        generator_module = importlib.import_module(
+            ".input_loaders.aten", package=__package__
+        )
+    elif op in SUPPORTED_INPUT_OPS:
+        op_module = ".".join(tritonbench_op.__module__.split(".")[:-1])
+        generator_module = importlib.import_module(f"{op_module}.input_loader")
     else:
         raise RuntimeError(f"Unsupported op: {op}")
-
-    generator_module = importlib.import_module(
-        f".input_loaders.{loader_type}", package=__package__
-    )
     input_iter_getter = generator_module.get_input_iter
     input_iter = input_iter_getter(tritonbench_op, op, input)
     return input_iter

tritonbench/operator_loader/aten/loader.py

@@ -22,8 +22,7 @@ def __init__(
             self.tb_args.input_loader = self.aten_op_input
 
     def get_input_iter(self) -> Generator:
-        for inp in self._get_input_iter():
-            yield inp
+        raise NotImplementedError("get_input_iter is not implemented for AtenOperator.")
 
     def eager(self, *input):
         args, kwargs = input

tritonbench/utils/triton_op.py

@@ -33,6 +33,7 @@
 from tritonbench.components.export import export_data
 
 from tritonbench.components.power.chart import power_chart_begin, power_chart_end
+from tritonbench.data import SUPPORTED_INPUT_OPS
 from tritonbench.utils.constants import (
     DEFAULT_QUANTILES,
     DEFAULT_REP,
@@ -777,7 +778,11 @@ def __init__(
     # Run the post initialization
     def __post__init__(self):
         if self.tb_args.input_loader:
-            if is_fbcode() and not hasattr(self, "aten_op_name"):
+            if (
+                is_fbcode()
+                and not hasattr(self, "aten_op_name")
+                and self.name not in SUPPORTED_INPUT_OPS
+            ):
                 from tritonbench.data.fb.input_loader import get_input_loader
 
                 self.get_input_iter = get_input_loader(
@@ -786,7 +791,7 @@ def __post__init__(self):
             else:
                 from tritonbench.data import get_input_loader
 
-                self._get_input_iter = get_input_loader(
+                self.get_input_iter = get_input_loader(
                     self, self.name, self.tb_args.input_loader
                 )
         # Count total available inputs directly