Op Input Generation via FACTO (#62)

Author: jiannanWang

.github/workflows/smoke-test.yml

@@ -24,8 +24,17 @@ jobs:
     - name: Install package and dependencies
       run: uv sync --dev
 
+    - name: Clone FACTO source
+      run: git clone https://github.com/pytorch-labs/FACTO.git
+    
+    - name: Build and install FACTO
+      run: cd FACTO && uv pip install .
+    
     - name: Run smoke test
       run: uv run python -m BackendBench.scripts.main --suite smoke --backend aten 
 
+    - name: Run FACTO test
+      run: uv run python -m BackendBench.scripts.main --suite facto --backend aten --ops "add.Tensor" 
+    
     - name: Run pytest tests
       run: uv run pytest test/

BackendBench/backends/flag_gems.py

@@ -1,5 +1,7 @@
 import torch
 
+from BackendBench.opregistry import register_operator
+
 from .base import Backend
 
 try:
@@ -284,6 +286,9 @@ def __init__(self) -> None:
             torch.ops.aten.eye.m: flag_gems.ops.eye_m,
             torch.ops.aten.to.dtype: flag_gems.ops.to_dtype,
         }
+        # Register all operators in the global registry to ensure consistent object identity
+        for op in self.ops.keys():
+            register_operator(op)
 
     def __getitem__(self, key):
         return self.ops[key]

BackendBench/facto_suite.py

@@ -0,0 +1,143 @@
+import logging
+from collections import defaultdict
+
+import torch
+from torch.utils._python_dispatch import TorchDispatchMode
+
+try:
+    from facto.inputgen.argtuple.gen import ArgumentTupleGenerator
+    from facto.inputgen.utils.config import TensorConfig
+    from facto.specdb.db import SpecDictDB
+except ImportError:
+    ArgumentTupleGenerator = None
+    TensorConfig = None
+    SpecDictDB = None
+
+
+from .eval import allclose
+from .opregistry import get_operator
+from .suite import OpTest, TestSuite
+
+logger = logging.getLogger(__name__)
+
+
+class FactoTest:
+    def __init__(self, *args, **kwargs):
+        self.args = args
+        self.kwargs = kwargs
+
+
+class FactoOpTest(OpTest):
+    def __init__(self, op, correctness_tests):
+        self.op = op
+        self._correctness_tests = correctness_tests
+        self.performance_tests = []
+
+    @property
+    def correctness_tests(self):
+        for test in self._correctness_tests:
+            yield FactoTest(*test.args, **test.kwargs)
+
+
+class OpTracerMode(TorchDispatchMode):
+    def __init__(self):
+        self.ops = []
+        self.args = []
+        self.kwargs = []
+
+    def __torch_dispatch__(self, fn, types, args=(), kwargs={}):
+        self.ops.append(fn)
+        self.args.append(args)
+        self.kwargs.append(kwargs)
+        return fn(*args, **kwargs)
+
+
+def build_facto_op_tests(device, dtype, filter=None, num_runs=10, empty=False, probability=1.0):
+    facto_op_tests = []
+    failed = []
+    for spec_name in SpecDictDB:
+        try:
+            if filter and spec_name not in filter:
+                continue
+
+            # Get canonical operator from registry
+            op = get_operator(spec_name)
+            if op is None:
+                logger.debug(f"Skipping {spec_name}: operator resolution failed")
+                continue
+
+            config = TensorConfig(
+                empty=empty,
+            ).set_probability(probability)
+
+            spec = SpecDictDB[spec_name]
+            generator = ArgumentTupleGenerator(spec, config)
+
+            op_tests = defaultdict(list)
+
+            for idx, (posargs, inkwargs, outargs) in enumerate(generator.gen()):
+                if idx >= num_runs:
+                    break
+
+                # Filter arguments to target device/dtype
+                filtered_posargs = []
+                for arg in posargs:
+                    if isinstance(arg, torch.Tensor):
+                        arg = arg.to(device=device, dtype=dtype)
+                    filtered_posargs.append(arg)
+
+                filtered_inkwargs = {}
+                for k, v in inkwargs.items():
+                    if isinstance(v, torch.Tensor):
+                        v = v.to(device=device, dtype=dtype)
+                    filtered_inkwargs[k] = v
+
+                filtered_outargs = {}
+                for k, v in outargs.items():
+                    if isinstance(v, torch.Tensor):
+                        v = v.to(device=device, dtype=dtype)
+                    filtered_outargs[k] = v
+
+                all_kwargs = {**filtered_inkwargs, **filtered_outargs}
+
+                try:
+                    # Trace execution to find underlying PyTorch ops
+                    with OpTracerMode() as tracer:
+                        ref = op(*filtered_posargs, **all_kwargs)
+                except Exception:
+                    logger.debug(f"FACTO spec {spec_name} couldn't run underlying op {op}")
+                    continue
+
+                # Check if we captured exactly one op (clean mapping)
+                if len(tracer.ops) == 1:
+                    try:
+                        # Verify the traced op produces the same result
+                        res = tracer.ops[0](*filtered_posargs, **all_kwargs)
+                        if allclose(ref, res):
+                            op_tests[tracer.ops[0]].append(
+                                FactoTest(*filtered_posargs, **all_kwargs)
+                            )
+                    except Exception:
+                        logger.debug(
+                            f"FACTO spec {spec_name} couldn't run underlying op {tracer.ops[0]}"
+                        )
+                else:
+                    logger.debug(f"FACTO spec {spec_name} has {len(tracer.ops)} ops")
+
+            for traced_op, tests in op_tests.items():
+                if len(tests) > 0:
+                    facto_op_tests.append(FactoOpTest(traced_op, tests))
+        except Exception:
+            logger.debug(f"FACTO spec {spec_name} failed")
+            failed.append(spec_name)
+
+    logger.debug(f"Failed specs: {failed}")
+
+    return facto_op_tests
+
+
+class FactoTestSuite(TestSuite):
+    def __init__(self, name, device, dtype, filter=None, num_runs=10, empty=False, probability=1.0):
+        super().__init__(
+            name, build_facto_op_tests(device, dtype, filter, num_runs, empty, probability)
+        )

BackendBench/opregistry.py

@@ -0,0 +1,104 @@
+import logging
+
+import torch
+
+logger = logging.getLogger(__name__)
+
+
+def _extract_spec_name_from_op(op_obj):
+    try:
+        # PyTorch operator objects have _name attribute that contains the full name
+        if hasattr(op_obj, "_name"):
+            full_name = op_obj._name
+            # full_name is typically like "aten::add.Tensor"
+            if "::" in full_name:
+                # Remove the "aten::" prefix
+                spec_name = full_name.split("::", 1)[1]
+                return spec_name
+        return None
+
+    except Exception as e:
+        logger.debug(f"Failed to extract spec name from operator {op_obj}: {e}")
+        return None
+
+
+class OpRegistry:
+    def __init__(self):
+        self._registry = {}
+
+    def get_operator(self, input_obj):
+        if isinstance(input_obj, str):
+            return self._get_operator_from_spec_name(input_obj)
+        else:
+            return self._get_operator_from_object(input_obj)
+
+    def _get_operator_from_spec_name(self, spec_name):
+        # Return cached operator if available
+        if spec_name in self._registry:
+            return self._registry[spec_name]
+
+        # Parse spec name
+        op_parts = spec_name.split(".")
+        op_name = op_parts[0]
+        overload = op_parts[1] if len(op_parts) > 1 else "default"
+
+        try:
+            # Resolve operator using PyTorch's API
+            op = getattr(torch.ops.aten, op_name).__getattr__(overload)
+
+            # Cache the resolved operator
+            self._registry[spec_name] = op
+            # logger.debug(f"Registered operator: {spec_name} -> {op}")
+            return op
+
+        except AttributeError as e:
+            logger.warning(f"Failed to resolve operator {spec_name}: {e}")
+            return None
+
+    def _get_operator_from_object(self, op_obj):
+        # Extract spec name from the operator object
+        spec_name = _extract_spec_name_from_op(op_obj)
+
+        # Check if we already have this operator registered
+        if spec_name in self._registry:
+            return self._registry[spec_name]
+
+        # Register the provided operator object
+        self._registry[spec_name] = op_obj
+        # logger.debug(f"Registered operator from object: {spec_name} -> {op_obj}")
+        return op_obj
+
+    def register_operator(self, op_obj):
+        return self._get_operator_from_object(op_obj)
+
+    def get_all_registered_ops(self):
+        return self._registry.copy()
+
+    def clear(self):
+        self._registry.clear()
+
+    def __len__(self):
+        return len(self._registry)
+
+    def __contains__(self, spec_name):
+        """Check if operator is registered."""
+        return spec_name in self._registry
+
+    def __repr__(self):
+        return f"OpRegistry({len(self._registry)} ops)"
+
+
+# Global operator registry instance
+_op_registry = OpRegistry()
+
+
+def get_operator(input_obj):
+    return _op_registry.get_operator(input_obj)
+
+
+def register_operator(op_obj):
+    return _op_registry.register_operator(op_obj)
+
+
+def get_registry():
+    return _op_registry

BackendBench/scripts/main.py

@@ -7,6 +7,7 @@
 import BackendBench.eval as eval
 import click
 import torch
+from BackendBench.facto_suite import FactoTestSuite
 from BackendBench.llm_client import ClaudeKernelGenerator
 from BackendBench.opinfo_suite import OpInfoTestSuite
 from BackendBench.suite import SmokeTestSuite
@@ -38,7 +39,7 @@ def setup_logging(log_level):
 @click.option(
     "--suite",
     default="smoke",
-    type=click.Choice(["smoke", "opinfo", "torchbench"]),
+    type=click.Choice(["smoke", "opinfo", "torchbench", "facto"]),
     help="Which suite to run",
 )
 @click.option(
@@ -128,6 +129,12 @@ def cli(
             filter=ops,
             topn=topn_inputs,
         ),
+        "facto": lambda: FactoTestSuite(
+            "facto_cuda_bfloat16",
+            "cuda",
+            torch.bfloat16,
+            filter=ops,
+        ),
     }[suite]()
 
     # For LLM backend, we need to generate kernels first

BackendBench/suite.py

@@ -1,5 +1,7 @@
 import torch
 
+from BackendBench.opregistry import get_operator
+
 
 def randn(*args, **kwargs):
     return lambda: torch.randn(*args, **kwargs)
@@ -40,7 +42,7 @@ def __iter__(self):
     "smoke",
     [
         OpTest(
-            torch.ops.aten.relu.default,
+            get_operator(torch.ops.aten.relu.default),
             [
                 Test(randn(2, device="cpu")),
             ],

pyproject.toml

@@ -30,6 +30,9 @@ dependencies = [
 flaggems = [
     # flag_gems must be installed from source: https://github.com/FlagOpen/FlagGems
 ]
+facto = [
+    # facto must be installed from source: https://github.com/pytorch-labs/FACTO
+]
 
 [project.scripts]
 backendbench = "BackendBench.scripts.main:cli"