xref: /aosp_15_r20/external/pytorch/benchmarks/fastrnns/conftest.py (revision da0073e96a02ea20f0ac840b70461e3646d07c45)

import pytest  # noqa: F401


default_rnns = [
    "cudnn",
    "aten",
    "jit",
    "jit_premul",
    "jit_premul_bias",
    "jit_simple",
    "jit_multilayer",
    "py",
]
default_cnns = ["resnet18", "resnet18_jit", "resnet50", "resnet50_jit"]
all_nets = default_rnns + default_cnns


def pytest_generate_tests(metafunc):
    # This creates lists of tests to generate, can be customized
    if metafunc.cls.__name__ == "TestBenchNetwork":
        metafunc.parametrize("net_name", all_nets, scope="class")
        metafunc.parametrize(
            "executor", [metafunc.config.getoption("executor")], scope="class"
        )
        metafunc.parametrize(
            "fuser", [metafunc.config.getoption("fuser")], scope="class"
        )


def pytest_addoption(parser):
    parser.addoption("--fuser", default="old", help="fuser to use for benchmarks")
    parser.addoption(
        "--executor", default="legacy", help="executor to use for benchmarks"
    )