eprint

import argparse
import json
import logging
import os
import re
import subprocess
import sys
import time
from enum import Enum
from typing import Any, Dict, List, NamedTuple, Optional, Set, Pattern

IS_WINDOWS: bool = os.name == "nt"

def eprint(*args: Any, **kwargs: Any) -> None:
    print(*args, file=sys.stderr, flush=True, **kwargs)

class LintSeverity(str, Enum):
    ERROR = "error"
    WARNING = "warning"
    ADVICE = "advice"
    DISABLED = "disabled"

class LintMessage(NamedTuple):
    path: Optional[str]
    line: Optional[int]
    char: Optional[int]
    code: str
    severity: LintSeverity
    name: str
    original: Optional[str]
    replacement: Optional[str]
    description: Optional[str]


LintSeverity

import argparse
import json
import logging
import os
import re
import subprocess
import sys
import time
from enum import Enum
from typing import Any, Dict, List, NamedTuple, Optional, Set, Pattern

IS_WINDOWS: bool = os.name == "nt"

def eprint(*args: Any, **kwargs: Any) -> None:
    print(*args, file=sys.stderr, flush=True, **kwargs)

class LintSeverity(str, Enum):
    ERROR = "error"
    WARNING = "warning"
    ADVICE = "advice"
    DISABLED = "disabled"

class LintMessage(NamedTuple):
    path: Optional[str]
    line: Optional[int]
    char: Optional[int]
    code: str
    severity: LintSeverity
    name: str
    original: Optional[str]
    replacement: Optional[str]
    description: Optional[str]

def as_posix(name: str) -> str:
    return name.replace("\\", "/") if IS_WINDOWS else name

# fmt: off
# https://www.flake8rules.com/


LintMessage

import sys
import time
from enum import Enum
from typing import Any, Dict, List, NamedTuple, Optional, Set, Pattern

IS_WINDOWS: bool = os.name == "nt"

def eprint(*args: Any, **kwargs: Any) -> None:
    print(*args, file=sys.stderr, flush=True, **kwargs)

class LintSeverity(str, Enum):
    ERROR = "error"
    WARNING = "warning"
    ADVICE = "advice"
    DISABLED = "disabled"

class LintMessage(NamedTuple):
    path: Optional[str]
    line: Optional[int]
    char: Optional[int]
    code: str
    severity: LintSeverity
    name: str
    original: Optional[str]
    replacement: Optional[str]
    description: Optional[str]

def as_posix(name: str) -> str:
    return name.replace("\\", "/") if IS_WINDOWS else name

# fmt: off
# https://www.flake8rules.com/
DOCUMENTED_IN_FLAKE8RULES: Set[str] = {
    "E101", "E111", "E112", "E113", "E114", "E115", "E116", "E117",
    "E121", "E122", "E123", "E124", "E125", "E126", "E127", "E128", "E129",
    "E131", "E133",
    "E201", "E202", "E203",
    "E211",
    "E221", "E222", "E223", "E224", "E225", "E226", "E227", "E228",
    "E231",
    "E241", "E242",
    "E251",
    "E261", "E262", "E265", "E266",
    "E271", "E272", "E273", "E274", "E275",


# fmt: on

# stdin:2: W802 undefined name 'foo'
# stdin:3:6: T484 Name 'foo' is not defined
# stdin:3:-100: W605 invalid escape sequence '\/'
# stdin:3:1: E302 expected 2 blank lines, found 1
RESULTS_RE: Pattern[str] = re.compile(
    r"""(?mx)
    ^
    (?P<file>.*?):
    (?P<line>\d+):
    (?:(?P<column>-?\d+):)?
    \s(?P<code>\S+?):?
    \s(?P<message>.*)
    $
    """
)

def _test_results_re() -> None:
    """
    >>> def t(s): return RESULTS_RE.search(s).groupdict()

    >>> t(r"file.py:80:1: E302 expected 2 blank lines, found 1")
    ... # doctest: +NORMALIZE_WHITESPACE
    {'file': 'file.py', 'line': '80', 'column': '1', 'code': 'E302',
     'message': 'expected 2 blank lines, found 1'}

    >>> t(r"file.py:7:1: P201: Resource `stdout` is acquired but not always released.")
    ... # doctest: +NORMALIZE_WHITESPACE
    {'file': 'file.py', 'line': '7', 'column': '1', 'code': 'P201',
     'message': 'Resource `stdout` is acquired but not always released.'}

    >>> t(r"file.py:8:-10: W605 invalid escape sequence '/'")
    ... # doctest: +NORMALIZE_WHITESPACE
    {'file': 'file.py', 'line': '8', 'column': '-10', 'code': 'W605',
     'message': "invalid escape sequence '/'"}
    """
    pass

def _run_command(
    args: List[str],
    *,
    extra_env: Optional[Dict[str, str]],
) -> "subprocess.CompletedProcess[str]":
    logging.debug(
        "$ %s",
        " ".join(
            ([f"{k}={v}" for (k, v) in extra_env.items()] if extra_env else []) + args
        ),
    )
    start_time = time.monotonic()
    try:
        return subprocess.run(
            args,
            stdout=subprocess.PIPE,
            stderr=subprocess.PIPE,
            check=True,


    >>> def t(s): return RESULTS_RE.search(s).groupdict()

    >>> t(r"file.py:80:1: E302 expected 2 blank lines, found 1")
    ... # doctest: +NORMALIZE_WHITESPACE
    {'file': 'file.py', 'line': '80', 'column': '1', 'code': 'E302',
     'message': 'expected 2 blank lines, found 1'}

    >>> t(r"file.py:7:1: P201: Resource `stdout` is acquired but not always released.")
    ... # doctest: +NORMALIZE_WHITESPACE
    {'file': 'file.py', 'line': '7', 'column': '1', 'code': 'P201',
     'message': 'Resource `stdout` is acquired but not always released.'}

    >>> t(r"file.py:8:-10: W605 invalid escape sequence '/'")
    ... # doctest: +NORMALIZE_WHITESPACE
    {'file': 'file.py', 'line': '8', 'column': '-10', 'code': 'W605',
     'message': "invalid escape sequence '/'"}
    """
    pass

def _run_command(
    args: List[str],
    *,
    extra_env: Optional[Dict[str, str]],
) -> "subprocess.CompletedProcess[str]":
    logging.debug(
        "$ %s",
        " ".join(
            ([f"{k}={v}" for (k, v) in extra_env.items()] if extra_env else []) + args
        ),
    )
    start_time = time.monotonic()
    try:
        return subprocess.run(
            args,
            stdout=subprocess.PIPE,
            stderr=subprocess.PIPE,
            check=True,
            encoding="utf-8",
        )
    finally:
        end_time = time.monotonic()
        logging.debug("took %dms", (end_time - start_time) * 1000)

def run_command(
    args: List[str],
    *,
    extra_env: Optional[Dict[str, str]],
    retries: int,
) -> "subprocess.CompletedProcess[str]":
    remaining_retries = retries
    while True:
        try:
            return _run_command(args, extra_env=extra_env)
        except subprocess.CalledProcessError as err:
            if remaining_retries == 0 or not re.match(
                r"^ERROR:1:1: X000 linting with .+ timed out after \d+ seconds",
                err.stdout,
            ):
                raise err
            remaining_retries -= 1
            logging.warning(


    logging.debug(
        "$ %s",
        " ".join(
            ([f"{k}={v}" for (k, v) in extra_env.items()] if extra_env else []) + args
        ),
    )
    start_time = time.monotonic()
    try:
        return subprocess.run(
            args,
            stdout=subprocess.PIPE,
            stderr=subprocess.PIPE,
            check=True,
            encoding="utf-8",
        )
    finally:
        end_time = time.monotonic()
        logging.debug("took %dms", (end_time - start_time) * 1000)

def run_command(
    args: List[str],
    *,
    extra_env: Optional[Dict[str, str]],
    retries: int,
) -> "subprocess.CompletedProcess[str]":
    remaining_retries = retries
    while True:
        try:
            return _run_command(args, extra_env=extra_env)
        except subprocess.CalledProcessError as err:
            if remaining_retries == 0 or not re.match(
                r"^ERROR:1:1: X000 linting with .+ timed out after \d+ seconds",
                err.stdout,
            ):
                raise err
            remaining_retries -= 1
            logging.warning(
                "(%s/%s) Retrying because command failed with: %r",
                retries - remaining_retries,
                retries,
                err,
            )
            time.sleep(1)

def get_issue_severity(code: str) -> LintSeverity:
    # "B901": `return x` inside a generator
    # "B902": Invalid first argument to a method
    # "B903": __slots__ efficiency
    # "B950": Line too long
    # "C4": Flake8 Comprehensions
    # "C9": Cyclomatic complexity
    # "E2": PEP8 horizontal whitespace "errors"
    # "E3": PEP8 blank line "errors"
    # "E5": PEP8 line length "errors"
    # "F401": Name imported but unused
    # "F403": Star imports used
    # "F405": Name possibly from star imports
    # "T400": type checking Notes
    # "T49": internal type checker errors or unmatched messages
    if any(
        code.startswith(x)
        for x in [


    remaining_retries = retries
    while True:
        try:
            return _run_command(args, extra_env=extra_env)
        except subprocess.CalledProcessError as err:
            if remaining_retries == 0 or not re.match(
                r"^ERROR:1:1: X000 linting with .+ timed out after \d+ seconds",
                err.stdout,
            ):
                raise err
            remaining_retries -= 1
            logging.warning(
                "(%s/%s) Retrying because command failed with: %r",
                retries - remaining_retries,
                retries,
                err,
            )
            time.sleep(1)

def get_issue_severity(code: str) -> LintSeverity:
    # "B901": `return x` inside a generator
    # "B902": Invalid first argument to a method
    # "B903": __slots__ efficiency
    # "B950": Line too long
    # "C4": Flake8 Comprehensions
    # "C9": Cyclomatic complexity
    # "E2": PEP8 horizontal whitespace "errors"
    # "E3": PEP8 blank line "errors"
    # "E5": PEP8 line length "errors"
    # "F401": Name imported but unused
    # "F403": Star imports used
    # "F405": Name possibly from star imports
    # "T400": type checking Notes
    # "T49": internal type checker errors or unmatched messages
    if any(
        code.startswith(x)
        for x in [
            "B9",
            "C4",
            "C9",
            "E2",
            "E3",
            "E5",
            "F401",
            "F403",
            "F405",
            "T400",
            "T49",
        ]
    ):
        return LintSeverity.ADVICE

    # "F821": Undefined name
    # "E999": syntax error
    if any(code.startswith(x) for x in ["F821", "E999"]):
        return LintSeverity.ERROR

    # "F": PyFlakes Error
    # "B": flake8-bugbear Error
    # "E": PEP8 "Error"
    # "W": PEP8 Warning
    # possibly other plugins...
    return LintSeverity.WARNING

def get_issue_documentation_url(code: str) -> str:
    if code in DOCUMENTED_IN_FLAKE8RULES:
        return f"https://www.flake8rules.com/rules/{code}.html"

    if code in DOCUMENTED_IN_FLAKE8COMPREHENSIONS:
        return "https://pypi.org/project/flake8-comprehensions/#rules"

    if code in DOCUMENTED_IN_BUGBEAR:
        return "https://github.com/PyCQA/flake8-bugbear#list-of-warnings"

    return ""

def check_files(
    filenames: List[str],
    flake8_plugins_path: Optional[str],
    severities: Dict[str, LintSeverity],
    retries: int,


            "F405",
            "T400",
            "T49",
        ]
    ):
        return LintSeverity.ADVICE

    # "F821": Undefined name
    # "E999": syntax error
    if any(code.startswith(x) for x in ["F821", "E999"]):
        return LintSeverity.ERROR

    # "F": PyFlakes Error
    # "B": flake8-bugbear Error
    # "E": PEP8 "Error"
    # "W": PEP8 Warning
    # possibly other plugins...
    return LintSeverity.WARNING

def get_issue_documentation_url(code: str) -> str:
    if code in DOCUMENTED_IN_FLAKE8RULES:
        return f"https://www.flake8rules.com/rules/{code}.html"

    if code in DOCUMENTED_IN_FLAKE8COMPREHENSIONS:
        return "https://pypi.org/project/flake8-comprehensions/#rules"

    if code in DOCUMENTED_IN_BUGBEAR:
        return "https://github.com/PyCQA/flake8-bugbear#list-of-warnings"

    return ""

def check_files(
    filenames: List[str],
    flake8_plugins_path: Optional[str],
    severities: Dict[str, LintSeverity],
    retries: int,
) -> List[LintMessage]:
    try:
        proc = run_command(
            [sys.executable, "-mflake8", "--exit-zero"] + filenames,
            extra_env={"FLAKE8_PLUGINS_PATH": flake8_plugins_path}
            if flake8_plugins_path
            else None,
            retries=retries,
        )
    except (OSError, subprocess.CalledProcessError) as err:
        return [
            LintMessage(
                path=None,


    # "B": flake8-bugbear Error
    # "E": PEP8 "Error"
    # "W": PEP8 Warning
    # possibly other plugins...
    return LintSeverity.WARNING

def get_issue_documentation_url(code: str) -> str:
    if code in DOCUMENTED_IN_FLAKE8RULES:
        return f"https://www.flake8rules.com/rules/{code}.html"

    if code in DOCUMENTED_IN_FLAKE8COMPREHENSIONS:
        return "https://pypi.org/project/flake8-comprehensions/#rules"

    if code in DOCUMENTED_IN_BUGBEAR:
        return "https://github.com/PyCQA/flake8-bugbear#list-of-warnings"

    return ""

def check_files(
    filenames: List[str],
    flake8_plugins_path: Optional[str],
    severities: Dict[str, LintSeverity],
    retries: int,
) -> List[LintMessage]:
    try:
        proc = run_command(
            [sys.executable, "-mflake8", "--exit-zero"] + filenames,
            extra_env={"FLAKE8_PLUGINS_PATH": flake8_plugins_path}
            if flake8_plugins_path
            else None,
            retries=retries,
        )
    except (OSError, subprocess.CalledProcessError) as err:
        return [
            LintMessage(
                path=None,
                line=None,
                char=None,
                code="FLAKE8",
                severity=LintSeverity.ERROR,
                name="command-failed",
                original=None,
                replacement=None,
                description=(
                    f"Failed due to {err.__class__.__name__}:\n{err}"
                    if not isinstance(err, subprocess.CalledProcessError)
                    else (
                        "COMMAND (exit code {returncode})\n"
                        "{command}\n\n"
                        "STDERR\n{stderr}\n\n"
                        "STDOUT\n{stdout}"
                    ).format(
                        returncode=err.returncode,
                        command=" ".join(as_posix(x) for x in err.cmd),
                        stderr=err.stderr.strip() or "(empty)",
                        stdout=err.stdout.strip() or "(empty)",
                    )
                ),
            )
        ]

    return [
        LintMessage(
            path=match["file"],
            name=match["code"],
            description="{}\nSee {}".format(
                match["message"],
                get_issue_documentation_url(match["code"]),
            ),
            line=int(match["line"]),
            char=int(match["column"])
            if match["column"] is not None and not match["column"].startswith("-")
            else None,
            code="FLAKE8",
            severity=severities.get(match["code"]) or get_issue_severity(match["code"]),
            original=None,
            replacement=None,
        )
        for match in RESULTS_RE.finditer(proc.stdout)
    ]

def main() -> None:
    parser = argparse.ArgumentParser(
        description="Flake8 wrapper linter.",
        fromfile_prefix_chars="@",
    )
    parser.add_argument(
        "--flake8-plugins-path",
        help="FLAKE8_PLUGINS_PATH env value",
    )
    parser.add_argument(
        "--severity",
        action="append",
        help="map code to severity (e.g. `B950:advice`)",
    )
    parser.add_argument(
        "--retries",
        default=3,
        type=int,


main

        LintMessage(
            path=match["file"],
            name=match["code"],
            description="{}\nSee {}".format(
                match["message"],
                get_issue_documentation_url(match["code"]),
            ),
            line=int(match["line"]),
            char=int(match["column"])
            if match["column"] is not None and not match["column"].startswith("-")
            else None,
            code="FLAKE8",
            severity=severities.get(match["code"]) or get_issue_severity(match["code"]),
            original=None,
            replacement=None,
        )
        for match in RESULTS_RE.finditer(proc.stdout)
    ]

def main() -> None:
    parser = argparse.ArgumentParser(
        description="Flake8 wrapper linter.",
        fromfile_prefix_chars="@",
    )
    parser.add_argument(
        "--flake8-plugins-path",
        help="FLAKE8_PLUGINS_PATH env value",
    )
    parser.add_argument(
        "--severity",
        action="append",
        help="map code to severity (e.g. `B950:advice`)",
    )
    parser.add_argument(
        "--retries",
        default=3,
        type=int,
        help="times to retry timed out flake8",
    )
    parser.add_argument(
        "--verbose",
        action="store_true",
        help="verbose logging",
    )
    parser.add_argument(
        "filenames",
        nargs="+",
        help="paths to lint",
    )
    args = parser.parse_args()

    logging.basicConfig(
        format="<%(threadName)s:%(levelname)s> %(message)s",
        level=logging.NOTSET
        if args.verbose
        else logging.DEBUG
        if len(args.filenames) < 1000
        else logging.INFO,
        stream=sys.stderr,
    )

    flake8_plugins_path = (
        None
        if args.flake8_plugins_path is None
        else os.path.realpath(args.flake8_plugins_path)
    )

    severities: Dict[str, LintSeverity] = {}
    if args.severity:
        for severity in args.severity:
            parts = severity.split(":", 1)
            assert len(parts) == 2, f"invalid severity `{severity}`"
            severities[parts[0]] = LintSeverity(parts[1])

    lint_messages = check_files(
        args.filenames, flake8_plugins_path, severities, args.retries
    )
    for lint_message in lint_messages:
        print(json.dumps(lint_message._asdict()), flush=True)

if __name__ == "__main__":
    main()


"""
NEWLINE: Checks files to make sure there are no trailing newlines.
"""
import argparse
import json
import logging
import os
import sys

from enum import Enum
from typing import NamedTuple, Optional

NEWLINE = 10  # ASCII "\n"
LINTER_CODE = "NEWLINE"

class LintSeverity(str, Enum):
    ERROR = "error"
    WARNING = "warning"
    ADVICE = "advice"
    DISABLED = "disabled"

class LintMessage(NamedTuple):
    path: Optional[str]
    line: Optional[int]
    char: Optional[int]
    code: str
    severity: LintSeverity
    name: str
    original: Optional[str]
    replacement: Optional[str]
    description: Optional[str]

def correct_trailing_newlines(filename: str) -> bool:
    with open(filename, "rb") as f:
        a = len(f.read(2))
        if a == 0:
            return True
        elif a == 1:
            # file is wrong whether or not the only byte is a newline
            return False
        else:
            f.seek(-2, os.SEEK_END)
            b, c = f.read(2)
            # no ASCII byte is part of any non-ASCII character in UTF-8
            return b != NEWLINE and c == NEWLINE

def check_file(filename: str) -> Optional[LintMessage]:
    logging.debug("Checking file %s", filename)

    with open(filename, "rb") as f:
        a = len(f.read(2))
        if a == 0:
            # File is empty, just leave it alone.
            return None

        if a == 1:
            # file is wrong whether or not the only byte is a newline
            return LintMessage(
                path=filename,
                line=None,
                char=None,
                code=LINTER_CODE,
                severity=LintSeverity.ERROR,
                name="testestTrailing newline",
                original=None,
                replacement=None,
                description="Trailing newline found. Run `lintrunner --take NEWLINE -a` to apply changes.",
            )

        # Read the last two bytes
        f.seek(-2, os.SEEK_END)
        b, c = f.read(2)
        # no ASCII byte is part of any non-ASCII character in UTF-8
        if b != NEWLINE and c == NEWLINE:
            return None

        f.seek(0)
        try:
            original = f.read().decode("utf-8")
        except Exception as err:
            return None

        return LintMessage(
            path=filename,
            line=None,
            char=None,
            code=LINTER_CODE,
            severity=LintSeverity.ERROR,
            name="Trailing newline",
            original=original,
            replacement=original.rstrip("\n") + "\n",
            description="Trailing newline found. Run `lintrunner --take NEWLINE -a` to apply changes.",
        )

if __name__ == "__main__":
    parser = argparse.ArgumentParser(
        description="native functions linter",
        fromfile_prefix_chars="@",
    )
    parser.add_argument(
        "--verbose",
        action="store_true",
        help="location of native_functions.yaml",
    )
    parser.add_argument(
        "filenames",
        nargs="+",
        help="paths to lint",
    )

    args = parser.parse_args()

    logging.basicConfig(
        format="<%(threadName)s:%(levelname)s> %(message)s",
        level=logging.NOTSET
        if args.verbose
        else logging.DEBUG
        if len(args.filenames) < 1000
        else logging.INFO,
        stream=sys.stderr,
    )

    lint_messages = []
    for filename in args.filenames:
        lint_message = check_file(filename)
        if lint_message is not None:
            lint_messages.append(lint_message)

    for lint_message in lint_messages:
        print(json.dumps(lint_message._asdict()), flush=True)


LintSeverity

"""
NEWLINE: Checks files to make sure there are no trailing newlines.
"""
import argparse
import json
import logging
import os
import sys

from enum import Enum
from typing import NamedTuple, Optional

NEWLINE = 10  # ASCII "\n"
LINTER_CODE = "NEWLINE"

class LintSeverity(str, Enum):
    ERROR = "error"
    WARNING = "warning"
    ADVICE = "advice"
    DISABLED = "disabled"

class LintMessage(NamedTuple):
    path: Optional[str]
    line: Optional[int]
    char: Optional[int]
    code: str
    severity: LintSeverity
    name: str
    original: Optional[str]
    replacement: Optional[str]
    description: Optional[str]

def correct_trailing_newlines(filename: str) -> bool:
    with open(filename, "rb") as f:
        a = len(f.read(2))
        if a == 0:
            return True
        elif a == 1:


LintMessage

import argparse
import json
import logging
import os
import sys

from enum import Enum
from typing import NamedTuple, Optional

NEWLINE = 10  # ASCII "\n"
LINTER_CODE = "NEWLINE"

class LintSeverity(str, Enum):
    ERROR = "error"
    WARNING = "warning"
    ADVICE = "advice"
    DISABLED = "disabled"

class LintMessage(NamedTuple):
    path: Optional[str]
    line: Optional[int]
    char: Optional[int]
    code: str
    severity: LintSeverity
    name: str
    original: Optional[str]
    replacement: Optional[str]
    description: Optional[str]

def correct_trailing_newlines(filename: str) -> bool:
    with open(filename, "rb") as f:
        a = len(f.read(2))
        if a == 0:
            return True
        elif a == 1:
            # file is wrong whether or not the only byte is a newline
            return False
        else:
            f.seek(-2, os.SEEK_END)
            b, c = f.read(2)
            # no ASCII byte is part of any non-ASCII character in UTF-8
            return b != NEWLINE and c == NEWLINE

def check_file(filename: str) -> Optional[LintMessage]:
    logging.debug("Checking file %s", filename)


class LintSeverity(str, Enum):
    ERROR = "error"
    WARNING = "warning"
    ADVICE = "advice"
    DISABLED = "disabled"

class LintMessage(NamedTuple):
    path: Optional[str]
    line: Optional[int]
    char: Optional[int]
    code: str
    severity: LintSeverity
    name: str
    original: Optional[str]
    replacement: Optional[str]
    description: Optional[str]

def correct_trailing_newlines(filename: str) -> bool:
    with open(filename, "rb") as f:
        a = len(f.read(2))
        if a == 0:
            return True
        elif a == 1:
            # file is wrong whether or not the only byte is a newline
            return False
        else:
            f.seek(-2, os.SEEK_END)
            b, c = f.read(2)
            # no ASCII byte is part of any non-ASCII character in UTF-8
            return b != NEWLINE and c == NEWLINE

def check_file(filename: str) -> Optional[LintMessage]:
    logging.debug("Checking file %s", filename)

    with open(filename, "rb") as f:
        a = len(f.read(2))
        if a == 0:
            # File is empty, just leave it alone.
            return None

        if a == 1:
            # file is wrong whether or not the only byte is a newline
            return LintMessage(
                path=filename,
                line=None,
                char=None,
                code=LINTER_CODE,
                severity=LintSeverity.ERROR,
                name="testestTrailing newline",


    original: Optional[str]
    replacement: Optional[str]
    description: Optional[str]

def correct_trailing_newlines(filename: str) -> bool:
    with open(filename, "rb") as f:
        a = len(f.read(2))
        if a == 0:
            return True
        elif a == 1:
            # file is wrong whether or not the only byte is a newline
            return False
        else:
            f.seek(-2, os.SEEK_END)
            b, c = f.read(2)
            # no ASCII byte is part of any non-ASCII character in UTF-8
            return b != NEWLINE and c == NEWLINE

def check_file(filename: str) -> Optional[LintMessage]:
    logging.debug("Checking file %s", filename)

    with open(filename, "rb") as f:
        a = len(f.read(2))
        if a == 0:
            # File is empty, just leave it alone.
            return None

        if a == 1:
            # file is wrong whether or not the only byte is a newline
            return LintMessage(
                path=filename,
                line=None,
                char=None,
                code=LINTER_CODE,
                severity=LintSeverity.ERROR,
                name="testestTrailing newline",
                original=None,
                replacement=None,
                description="Trailing newline found. Run `lintrunner --take NEWLINE -a` to apply changes.",
            )

        # Read the last two bytes
        f.seek(-2, os.SEEK_END)
        b, c = f.read(2)
        # no ASCII byte is part of any non-ASCII character in UTF-8
        if b != NEWLINE and c == NEWLINE:
            return None

        f.seek(0)
        try:
            original = f.read().decode("utf-8")
        except Exception as err:
            return None

        return LintMessage(
            path=filename,
            line=None,
            char=None,
            code=LINTER_CODE,
            severity=LintSeverity.ERROR,
            name="Trailing newline",
            original=original,
            replacement=original.rstrip("\n") + "\n",
            description="Trailing newline found. Run `lintrunner --take NEWLINE -a` to apply changes.",
        )

if __name__ == "__main__":
    parser = argparse.ArgumentParser(
        description="native functions linter",
        fromfile_prefix_chars="@",
    )
    parser.add_argument(
        "--verbose",
        action="store_true",
        help="location of native_functions.yaml",
    )
    parser.add_argument(
        "filenames",
        nargs="+",
        help="paths to lint",
    )

    args = parser.parse_args()


# SPDX-FileCopyrightText: Copyright (c) 2024-present NVIDIA CORPORATION & AFFILIATES.
# All rights reserved.
# SPDX-License-Identifier: BSD-3-Clause
import pytest
import torch
from .global_params import generate_input_sizes, FLOAT_DTYPES
from .normalization import norm_bwd_nvf_benchmark, norm_bwd_baseline_benchmark
from .core import DEFAULT_EXECUTORS

@pytest.mark.parametrize("size", generate_input_sizes(dims=4))
@pytest.mark.parametrize("dtype", FLOAT_DTYPES)
@pytest.mark.parametrize(
    "channels_last",
    [
        pytest.param(True, marks=pytest.mark.outer_persistent),
        pytest.param(False, marks=pytest.mark.inner_persistent),
    ],
)
def test_batchnorm_bwd_nvf_benchmark(
    benchmark,
    size: tuple,
    dtype: torch.dtype,
    channels_last: bool,
    disable_validation: bool,
    disable_benchmarking: bool,
    eps: float = 1e-5,
):
    norm_bwd_nvf_benchmark(
        benchmark,
        size,
        dtype,
        "batch_norm",
        channels_last,
        disable_validation,
        disable_benchmarking,
        eps,
    )

@pytest.mark.parametrize("executor", DEFAULT_EXECUTORS)
@pytest.mark.parametrize("size", generate_input_sizes(dims=4))
@pytest.mark.parametrize("dtype", FLOAT_DTYPES)
@pytest.mark.parametrize(
    "channels_last",
    [
        pytest.param(True, marks=pytest.mark.outer_persistent),
        pytest.param(False, marks=pytest.mark.inner_persistent),
    ],
)
def test_batchnorm_bwd_baseline_benchmark(
    benchmark, size: tuple, dtype: torch.dtype, channels_last: bool, executor: str
):
    norm_bwd_baseline_benchmark(
        benchmark, size, dtype, channels_last, executor, "batch_norm"
    )


        size,
        dtype,
        "batch_norm",
        channels_last,
        disable_validation,
        disable_benchmarking,
        eps,
    )

@pytest.mark.parametrize("executor", DEFAULT_EXECUTORS)
@pytest.mark.parametrize("size", generate_input_sizes(dims=4))
@pytest.mark.parametrize("dtype", FLOAT_DTYPES)
@pytest.mark.parametrize(
    "channels_last",
    [
        pytest.param(True, marks=pytest.mark.outer_persistent),
        pytest.param(False, marks=pytest.mark.inner_persistent),
    ],
)
def test_batchnorm_bwd_baseline_benchmark(
    benchmark, size: tuple, dtype: torch.dtype, channels_last: bool, executor: str
):
    norm_bwd_baseline_benchmark(
        benchmark, size, dtype, channels_last, executor, "batch_norm"
    )


# SPDX-FileCopyrightText: Copyright (c) 2024-present NVIDIA CORPORATION & AFFILIATES.
# All rights reserved.
# SPDX-License-Identifier: BSD-3-Clause
import pytest
from nvfuser_direct import FusionDefinition, DataType
from nvfuser_direct.pytorch_utils import torch_dtype_to_nvfuser_dtype
from .core import run_benchmark, clear_dynamo_cache, with_executor, DEFAULT_EXECUTORS
import torch
from .global_params import generate_input_sizes, FLOAT_DTYPES, PROMOTE_DTYPES
from .torch_ops import silu_mul

def silu_mul_fwd_fusion(fd: FusionDefinition, dtype: DataType):
    T0 = fd.define_tensor(
        shape=[-1, -1], contiguity=[True, True], dtype=dtype, is_cpu=False
    )
    T1 = fd.define_tensor(
        shape=[-1, -1], contiguity=[True, True], dtype=dtype, is_cpu=False
    )
    if dtype in PROMOTE_DTYPES:
        T0 = fd.ops.cast(T0, dtype=DataType.Float)
        T1 = fd.ops.cast(T1, dtype=DataType.Float)
    T2 = fd.ops.neg(T0)
    T3 = fd.ops.exp(T2)
    S4 = fd.define_scalar(1.00000, dtype=DataType.Double)
    T5 = fd.ops.add(S4, T3)
    T6 = fd.ops.reciprocal(T5)
    T7 = fd.ops.mul(T0, T6)
    T8 = fd.ops.mul(T7, T1)
    if dtype in PROMOTE_DTYPES:
        T8 = fd.ops.cast(T8, dtype=dtype)
    fd.add_output(T8)

@pytest.mark.parametrize("size", generate_input_sizes(dims=2))
@pytest.mark.parametrize("dtype", FLOAT_DTYPES)
@pytest.mark.pointwise
def test_silu_mul_fwd_nvf_benchmark(
    benchmark,
    size: tuple,
    dtype: torch.dtype,
    disable_validation: bool,
    disable_benchmarking: bool,
):
    inputs = [torch.randn(*size, device="cuda", dtype=dtype) for _ in range(2)]

    with FusionDefinition() as fd:
        silu_mul_fwd_fusion(fd, torch_dtype_to_nvfuser_dtype(dtype))
    if not disable_validation:
        eager_output = silu_mul(inputs)
        fd.validate(inputs, [eager_output])

    if not disable_benchmarking:
        run_benchmark(benchmark, fd.execute, inputs)

@pytest.mark.parametrize("executor", DEFAULT_EXECUTORS)
@pytest.mark.parametrize("size", generate_input_sizes(dims=2))
@pytest.mark.parametrize("dtype", FLOAT_DTYPES)
@pytest.mark.pointwise
def test_silu_mul_fwd_baseline_benchmark(
    benchmark,
    size: tuple,
    dtype: torch.dtype,
    executor: str,
):
    if executor == "torchcompile":
        clear_dynamo_cache()
    inputs = [
        torch.randn(*size, device="cuda", dtype=dtype, requires_grad=True)
        for _ in range(2)
    ]

    benchmark_fn = with_executor(executor, silu_mul)

    # Inputs and outputs are same as nvFuser, no need for manual IOByte computation
    run_benchmark(
        benchmark,
        benchmark_fn,
        inputs,
    )


# SPDX-FileCopyrightText: Copyright (c) 2024-present NVIDIA CORPORATION & AFFILIATES.
# All rights reserved.
# SPDX-License-Identifier: BSD-3-Clause
import pytest
from nvfuser_direct import FusionDefinition, DataType
from nvfuser_direct.pytorch_utils import torch_dtype_to_nvfuser_dtype
from .core import run_benchmark, clear_dynamo_cache, with_executor, DEFAULT_EXECUTORS
import torch
from .global_params import generate_input_sizes, FLOAT_DTYPES, PROMOTE_DTYPES
from .torch_ops import silu_mul

def silu_mul_fwd_fusion(fd: FusionDefinition, dtype: DataType):
    T0 = fd.define_tensor(
        shape=[-1, -1], contiguity=[True, True], dtype=dtype, is_cpu=False
    )
    T1 = fd.define_tensor(
        shape=[-1, -1], contiguity=[True, True], dtype=dtype, is_cpu=False
    )
    if dtype in PROMOTE_DTYPES:
        T0 = fd.ops.cast(T0, dtype=DataType.Float)
        T1 = fd.ops.cast(T1, dtype=DataType.Float)
    T2 = fd.ops.neg(T0)
    T3 = fd.ops.exp(T2)
    S4 = fd.define_scalar(1.00000, dtype=DataType.Double)
    T5 = fd.ops.add(S4, T3)
    T6 = fd.ops.reciprocal(T5)
    T7 = fd.ops.mul(T0, T6)
    T8 = fd.ops.mul(T7, T1)
    if dtype in PROMOTE_DTYPES:
        T8 = fd.ops.cast(T8, dtype=dtype)
    fd.add_output(T8)

@pytest.mark.parametrize("size", generate_input_sizes(dims=2))
@pytest.mark.parametrize("dtype", FLOAT_DTYPES)
@pytest.mark.pointwise
def test_silu_mul_fwd_nvf_benchmark(
    benchmark,
    size: tuple,
    dtype: torch.dtype,
    disable_validation: bool,
    disable_benchmarking: bool,
):
    inputs = [torch.randn(*size, device="cuda", dtype=dtype) for _ in range(2)]

    with FusionDefinition() as fd:
        silu_mul_fwd_fusion(fd, torch_dtype_to_nvfuser_dtype(dtype))
    if not disable_validation:
        eager_output = silu_mul(inputs)
        fd.validate(inputs, [eager_output])


        shape=[-1, -1], contiguity=[True, True], dtype=dtype, is_cpu=False
    )
    if dtype in PROMOTE_DTYPES:
        T0 = fd.ops.cast(T0, dtype=DataType.Float)
        T1 = fd.ops.cast(T1, dtype=DataType.Float)
    T2 = fd.ops.neg(T0)
    T3 = fd.ops.exp(T2)
    S4 = fd.define_scalar(1.00000, dtype=DataType.Double)
    T5 = fd.ops.add(S4, T3)
    T6 = fd.ops.reciprocal(T5)
    T7 = fd.ops.mul(T0, T6)
    T8 = fd.ops.mul(T7, T1)
    if dtype in PROMOTE_DTYPES:
        T8 = fd.ops.cast(T8, dtype=dtype)
    fd.add_output(T8)

@pytest.mark.parametrize("size", generate_input_sizes(dims=2))
@pytest.mark.parametrize("dtype", FLOAT_DTYPES)
@pytest.mark.pointwise
def test_silu_mul_fwd_nvf_benchmark(
    benchmark,
    size: tuple,
    dtype: torch.dtype,
    disable_validation: bool,
    disable_benchmarking: bool,
):
    inputs = [torch.randn(*size, device="cuda", dtype=dtype) for _ in range(2)]

    with FusionDefinition() as fd:
        silu_mul_fwd_fusion(fd, torch_dtype_to_nvfuser_dtype(dtype))
    if not disable_validation:
        eager_output = silu_mul(inputs)
        fd.validate(inputs, [eager_output])

    if not disable_benchmarking:
        run_benchmark(benchmark, fd.execute, inputs)

@pytest.mark.parametrize("executor", DEFAULT_EXECUTORS)
@pytest.mark.parametrize("size", generate_input_sizes(dims=2))
@pytest.mark.parametrize("dtype", FLOAT_DTYPES)
@pytest.mark.pointwise
def test_silu_mul_fwd_baseline_benchmark(
    benchmark,
    size: tuple,
    dtype: torch.dtype,
    executor: str,
):
    if executor == "torchcompile":
        clear_dynamo_cache()
    inputs = [
        torch.randn(*size, device="cuda", dtype=dtype, requires_grad=True)
        for _ in range(2)
    ]

    benchmark_fn = with_executor(executor, silu_mul)


    dtype: torch.dtype,
    disable_validation: bool,
    disable_benchmarking: bool,
):
    inputs = [torch.randn(*size, device="cuda", dtype=dtype) for _ in range(2)]

    with FusionDefinition() as fd:
        silu_mul_fwd_fusion(fd, torch_dtype_to_nvfuser_dtype(dtype))
    if not disable_validation:
        eager_output = silu_mul(inputs)
        fd.validate(inputs, [eager_output])

    if not disable_benchmarking:
        run_benchmark(benchmark, fd.execute, inputs)

@pytest.mark.parametrize("executor", DEFAULT_EXECUTORS)
@pytest.mark.parametrize("size", generate_input_sizes(dims=2))
@pytest.mark.parametrize("dtype", FLOAT_DTYPES)
@pytest.mark.pointwise
def test_silu_mul_fwd_baseline_benchmark(
    benchmark,
    size: tuple,
    dtype: torch.dtype,
    executor: str,
):
    if executor == "torchcompile":
        clear_dynamo_cache()
    inputs = [
        torch.randn(*size, device="cuda", dtype=dtype, requires_grad=True)
        for _ in range(2)
    ]

    benchmark_fn = with_executor(executor, silu_mul)

    # Inputs and outputs are same as nvFuser, no need for manual IOByte computation
    run_benchmark(
        benchmark,
        benchmark_fn,
        inputs,
    )


# SPDX-FileCopyrightText: Copyright (c) 2024-present NVIDIA CORPORATION & AFFILIATES.
# All rights reserved.
# SPDX-License-Identifier: BSD-3-Clause
import pytest
from nvfuser_direct import FusionDefinition, DataType
from nvfuser_direct.pytorch_utils import torch_dtype_to_nvfuser_dtype
from .core import run_benchmark, clear_dynamo_cache, with_executor, DEFAULT_EXECUTORS
import torch
from .global_params import generate_input_sizes, FLOAT_DTYPES, PROMOTE_DTYPES
from .torch_ops import gelu

def gelu_fwd_fusion(
    fd: FusionDefinition,
    dtype: DataType,
) -> None:
    input = fd.define_tensor(
        shape=[-1, -1], contiguity=[True, True], dtype=dtype, is_cpu=False
    )
    bias = fd.define_tensor(shape=[-1], contiguity=[True], dtype=dtype, is_cpu=False)
    if dtype in PROMOTE_DTYPES:
        input = fd.ops.cast(input, dtype=DataType.Float)
        bias = fd.ops.cast(bias, dtype=DataType.Float)
    S_079 = fd.define_scalar(0.79788456)
    S_004 = fd.define_scalar(0.044715)
    bias = fd.ops.broadcast_in_dim(bias, shape=[1, input.size(-1)], broadcast_dims=[1])
    T1 = fd.ops.add(input, bias)
    T2 = fd.ops.mul(S_079, T1)
    T3 = fd.ops.mul(S_004, T1)
    T4 = fd.ops.mul(T3, T1)
    S1 = fd.define_scalar(1.0)
    T5 = fd.ops.add(T4, S1)
    T6 = fd.ops.mul(T2, T5)
    T7 = fd.ops.tanh(T6)
    T8 = fd.ops.add(S1, T7)
    T9 = fd.ops.mul(T8, T1)
    S2 = fd.define_scalar(0.50)
    T10 = fd.ops.mul(S2, T9)
    if dtype in PROMOTE_DTYPES:
        T10 = fd.ops.cast(T10, dtype=dtype)
    fd.add_output(T10)

@pytest.mark.parametrize("size", generate_input_sizes(dims=2))
@pytest.mark.parametrize("dtype", FLOAT_DTYPES)
@pytest.mark.pointwise
def test_gelu_fwd_nvf_benchmark(
    benchmark,
    size: tuple,
    dtype: torch.dtype,
    disable_validation: bool,
    disable_benchmarking: bool,
):
    inputs = [
        torch.randn(size, device="cuda", dtype=dtype, requires_grad=True),  # in_tensor
        torch.ones(size[-1], device="cuda", dtype=dtype),  # bias
    ]
    with FusionDefinition() as fd:
        gelu_fwd_fusion(fd, torch_dtype_to_nvfuser_dtype(dtype))
    if not disable_validation:
        eager_output = gelu(inputs)
        fd.validate(inputs, [eager_output])
    if not disable_benchmarking:
        run_benchmark(benchmark, fd.execute, inputs)

@pytest.mark.parametrize("executor", DEFAULT_EXECUTORS)
@pytest.mark.parametrize("size", generate_input_sizes(dims=2))
@pytest.mark.parametrize("dtype", FLOAT_DTYPES)
@pytest.mark.pointwise
def test_gelu_fwd_baseline_benchmark(
    benchmark,
    size: tuple,
    dtype: torch.dtype,
    executor: str,
):
    if executor == "torchcompile":
        clear_dynamo_cache()
    inputs = [
        torch.randn(size, device="cuda", dtype=dtype, requires_grad=True),  # in_tensor
        torch.ones(size[-1], device="cuda", dtype=dtype),  # bias
    ]

    benchmark_fn = with_executor(executor, gelu)

    # Inputs and outputs are same as nvFuser, no need for manual IOByte computation
    run_benchmark(benchmark, benchmark_fn, inputs)


# SPDX-FileCopyrightText: Copyright (c) 2024-present NVIDIA CORPORATION & AFFILIATES.
# All rights reserved.
# SPDX-License-Identifier: BSD-3-Clause
import pytest
from nvfuser_direct import FusionDefinition, DataType
from nvfuser_direct.pytorch_utils import torch_dtype_to_nvfuser_dtype
from .core import run_benchmark, clear_dynamo_cache, with_executor, DEFAULT_EXECUTORS
import torch
from .global_params import generate_input_sizes, FLOAT_DTYPES, PROMOTE_DTYPES
from .torch_ops import gelu

def gelu_fwd_fusion(
    fd: FusionDefinition,
    dtype: DataType,
) -> None:
    input = fd.define_tensor(
        shape=[-1, -1], contiguity=[True, True], dtype=dtype, is_cpu=False
    )
    bias = fd.define_tensor(shape=[-1], contiguity=[True], dtype=dtype, is_cpu=False)
    if dtype in PROMOTE_DTYPES:
        input = fd.ops.cast(input, dtype=DataType.Float)
        bias = fd.ops.cast(bias, dtype=DataType.Float)
    S_079 = fd.define_scalar(0.79788456)
    S_004 = fd.define_scalar(0.044715)
    bias = fd.ops.broadcast_in_dim(bias, shape=[1, input.size(-1)], broadcast_dims=[1])
    T1 = fd.ops.add(input, bias)
    T2 = fd.ops.mul(S_079, T1)
    T3 = fd.ops.mul(S_004, T1)
    T4 = fd.ops.mul(T3, T1)
    S1 = fd.define_scalar(1.0)
    T5 = fd.ops.add(T4, S1)
    T6 = fd.ops.mul(T2, T5)
    T7 = fd.ops.tanh(T6)
    T8 = fd.ops.add(S1, T7)
    T9 = fd.ops.mul(T8, T1)
    S2 = fd.define_scalar(0.50)
    T10 = fd.ops.mul(S2, T9)
    if dtype in PROMOTE_DTYPES:
        T10 = fd.ops.cast(T10, dtype=dtype)
    fd.add_output(T10)

@pytest.mark.parametrize("size", generate_input_sizes(dims=2))
@pytest.mark.parametrize("dtype", FLOAT_DTYPES)
@pytest.mark.pointwise
def test_gelu_fwd_nvf_benchmark(
    benchmark,
    size: tuple,
    dtype: torch.dtype,
    disable_validation: bool,
    disable_benchmarking: bool,
):
    inputs = [
        torch.randn(size, device="cuda", dtype=dtype, requires_grad=True),  # in_tensor
        torch.ones(size[-1], device="cuda", dtype=dtype),  # bias
    ]
    with FusionDefinition() as fd:
        gelu_fwd_fusion(fd, torch_dtype_to_nvfuser_dtype(dtype))
    if not disable_validation:
        eager_output = gelu(inputs)


    T1 = fd.ops.add(input, bias)
    T2 = fd.ops.mul(S_079, T1)
    T3 = fd.ops.mul(S_004, T1)
    T4 = fd.ops.mul(T3, T1)
    S1 = fd.define_scalar(1.0)
    T5 = fd.ops.add(T4, S1)
    T6 = fd.ops.mul(T2, T5)
    T7 = fd.ops.tanh(T6)
    T8 = fd.ops.add(S1, T7)
    T9 = fd.ops.mul(T8, T1)
    S2 = fd.define_scalar(0.50)
    T10 = fd.ops.mul(S2, T9)
    if dtype in PROMOTE_DTYPES:
        T10 = fd.ops.cast(T10, dtype=dtype)
    fd.add_output(T10)

@pytest.mark.parametrize("size", generate_input_sizes(dims=2))
@pytest.mark.parametrize("dtype", FLOAT_DTYPES)
@pytest.mark.pointwise
def test_gelu_fwd_nvf_benchmark(
    benchmark,
    size: tuple,
    dtype: torch.dtype,
    disable_validation: bool,
    disable_benchmarking: bool,
):
    inputs = [
        torch.randn(size, device="cuda", dtype=dtype, requires_grad=True),  # in_tensor
        torch.ones(size[-1], device="cuda", dtype=dtype),  # bias
    ]
    with FusionDefinition() as fd:
        gelu_fwd_fusion(fd, torch_dtype_to_nvfuser_dtype(dtype))
    if not disable_validation:
        eager_output = gelu(inputs)
        fd.validate(inputs, [eager_output])
    if not disable_benchmarking:
        run_benchmark(benchmark, fd.execute, inputs)

@pytest.mark.parametrize("executor", DEFAULT_EXECUTORS)
@pytest.mark.parametrize("size", generate_input_sizes(dims=2))
@pytest.mark.parametrize("dtype", FLOAT_DTYPES)
@pytest.mark.pointwise
def test_gelu_fwd_baseline_benchmark(
    benchmark,
    size: tuple,
    dtype: torch.dtype,
    executor: str,
):
    if executor == "torchcompile":
        clear_dynamo_cache()
    inputs = [
        torch.randn(size, device="cuda", dtype=dtype, requires_grad=True),  # in_tensor
        torch.ones(size[-1], device="cuda", dtype=dtype),  # bias
    ]

    benchmark_fn = with_executor(executor, gelu)


    disable_validation: bool,
    disable_benchmarking: bool,
):
    inputs = [
        torch.randn(size, device="cuda", dtype=dtype, requires_grad=True),  # in_tensor
        torch.ones(size[-1], device="cuda", dtype=dtype),  # bias
    ]
    with FusionDefinition() as fd:
        gelu_fwd_fusion(fd, torch_dtype_to_nvfuser_dtype(dtype))
    if not disable_validation:
        eager_output = gelu(inputs)
        fd.validate(inputs, [eager_output])
    if not disable_benchmarking:
        run_benchmark(benchmark, fd.execute, inputs)

@pytest.mark.parametrize("executor", DEFAULT_EXECUTORS)
@pytest.mark.parametrize("size", generate_input_sizes(dims=2))
@pytest.mark.parametrize("dtype", FLOAT_DTYPES)
@pytest.mark.pointwise
def test_gelu_fwd_baseline_benchmark(
    benchmark,
    size: tuple,
    dtype: torch.dtype,
    executor: str,
):
    if executor == "torchcompile":
        clear_dynamo_cache()
    inputs = [
        torch.randn(size, device="cuda", dtype=dtype, requires_grad=True),  # in_tensor
        torch.ones(size[-1], device="cuda", dtype=dtype),  # bias
    ]

    benchmark_fn = with_executor(executor, gelu)

    # Inputs and outputs are same as nvFuser, no need for manual IOByte computation
    run_benchmark(benchmark, benchmark_fn, inputs)


# SPDX-FileCopyrightText: Copyright (c) 2024-present NVIDIA CORPORATION & AFFILIATES.
# All rights reserved.
# SPDX-License-Identifier: BSD-3-Clause
import pytest
from nvfuser_direct import FusionDefinition, DataType
from nvfuser_direct.pytorch_utils import torch_dtype_to_nvfuser_dtype
from .core import (
    run_benchmark,
    clear_dynamo_cache,
    unary_bwd_torch,
    compute_total_iobytes,
    with_executor,
    DEFAULT_EXECUTORS,
)
import torch
from .global_params import generate_input_sizes, FLOAT_DTYPES, PROMOTE_DTYPES
from .torch_ops import dropout_layernorm

def dropout_layernorm_bwd_fusion(
    fd: FusionDefinition, dtype: DataType, dropout_p: float
) -> None:
    """
    Backward pass fusion definition for computing:
        output = layernorm (input2 + dropout (input1 p=dropout_p))

    Fusion inputs: inputs, dropout_mask, rms, grads, weights
    Fusion outputs: grad_input, grad_weights, grad_bias
    """
    T1 = fd.define_tensor(
        shape=[-1, -1], contiguity=[True, True], dtype=DataType.Bool, is_cpu=False
    )  # mask
    T2 = fd.define_tensor(
        shape=[-1], contiguity=[True], dtype=DataType.Float, is_cpu=False
    )  # mean
    T3 = fd.define_tensor(
        shape=[-1, 1], contiguity=[True, None], dtype=DataType.Float, is_cpu=False
    )  # invstd
    T4 = fd.define_tensor(
        shape=[-1, -1], contiguity=[True, True], dtype=dtype, is_cpu=False
    )  # grads
    T5 = fd.define_tensor(
        shape=[-1], contiguity=[True], dtype=dtype, is_cpu=False
    )  # weights
    T6 = fd.define_tensor(
        shape=[-1, -1], contiguity=[True, True], dtype=dtype, is_cpu=False
    )  # input1

    T7 = fd.define_tensor(
        shape=[-1, -1], contiguity=[True, True], dtype=dtype, is_cpu=False
    )  # input2

    if dtype in PROMOTE_DTYPES:
        T1 = fd.ops.cast(T1, dtype=DataType.Float)
        T4 = fd.ops.cast(T4, dtype=DataType.Float)
        T5 = fd.ops.cast(T5, dtype=DataType.Float)
        T6 = fd.ops.cast(T6, dtype=DataType.Float)
        T7 = fd.ops.cast(T7, dtype=DataType.Float)

    T9 = fd.ops.mul(T6, T1)
    S10 = fd.define_scalar(1 / (1 - dropout_p), dtype=DataType.Double)
    T11 = fd.ops.mul(T9, S10)
    T12 = fd.ops.add(T7, T11)

    T16 = fd.ops.broadcast_in_dim(T2, shape=[T6.size(0), 1], broadcast_dims=[0])
    V19 = T6.shape()
    T20 = fd.ops.broadcast_in_dim(T16, shape=V19, broadcast_dims=[0, 1])
    T21 = fd.ops.sub(T12, T20)
    T25 = fd.ops.broadcast_in_dim(T3, shape=V19, broadcast_dims=[0, 1])
    T26 = fd.ops.mul(T21, T25)
    T30 = fd.ops.broadcast_in_dim(T5, shape=V19, broadcast_dims=[1])
    T35 = fd.ops.sum(T4, dims=[0], keepdim=False, dtype=DataType.Null)

    T37 = fd.ops.mul(T4, T30)
    T38 = fd.ops.mul(T4, T26)
    T39 = fd.ops.sum(T38, dims=[0], keepdim=False, dtype=DataType.Null)

    T41 = fd.ops.mul(T37, T25)
    T42 = fd.ops.mul(T37, T21)
    T43 = fd.ops.sum(T42, dims=[1], keepdim=False, dtype=DataType.Null)
    T47 = fd.ops.broadcast_in_dim(T43, shape=[T6.size(0), 1], broadcast_dims=[0])
    T48 = fd.ops.neg(T41)
    T49 = fd.ops.sum(T48, dims=[1], keepdim=False, dtype=DataType.Null)
    T53 = fd.ops.broadcast_in_dim(T49, shape=[T6.size(0), 1], broadcast_dims=[0])
    S54 = fd.define_scalar(-0.500000, dtype=DataType.Double)
    T55 = fd.ops.mul(S54, T47)
    S56 = fd.define_scalar(3.00000, dtype=DataType.Double)
    T57 = fd.ops.pow(T3, S56)
    T58 = fd.ops.mul(T55, T57)
    T61 = fd.ops.sum(T53, dims=[1], keepdim=False, dtype=DataType.Null)
    T62 = fd.ops.sum(T58, dims=[1], keepdim=False, dtype=DataType.Null)
    T66 = fd.ops.broadcast_in_dim(T62, shape=[T6.size(0), 1], broadcast_dims=[0])
    T70 = fd.ops.broadcast_in_dim(T66, shape=V19, broadcast_dims=[0, 1])
    T74 = fd.ops.broadcast_in_dim(T2, shape=[T6.size(0), 1], broadcast_dims=[0])
    T78 = fd.ops.broadcast_in_dim(T74, shape=V19, broadcast_dims=[0, 1])
    S79 = fd.define_scalar(2.00000, dtype=DataType.Double)
    T80 = fd.ops.mul(S79, T70)
    T81 = fd.ops.sub(T12, T78)
    T82 = fd.ops.mul(T80, T81)
    S84 = fd.ops.reciprocal(T6.size(1))
    T85 = fd.ops.mul(T82, S84)
    T89 = fd.ops.broadcast_in_dim(T61, shape=[T6.size(0), 1], broadcast_dims=[0])
    T93 = fd.ops.broadcast_in_dim(T89, shape=V19, broadcast_dims=[0, 1])
    T95 = fd.ops.mul(S84, T93)
    T96 = fd.ops.add(T85, T95)
    T97 = fd.ops.add(T41, T96)

    T100 = fd.ops.mul(T97, S10)
    T101 = fd.ops.mul(T100, T1)

    if dtype in PROMOTE_DTYPES:
        T35 = fd.ops.cast(T35, dtype=dtype)
        T39 = fd.ops.cast(T39, dtype=dtype)
        T97 = fd.ops.cast(T97, dtype=dtype)
        T101 = fd.ops.cast(T101, dtype=dtype)

    fd.add_output(T101)
    fd.add_output(T97)
    fd.add_output(T39)
    fd.add_output(T35)

def dropout_layernorm_bwd_iobytes(size: tuple, dtype: torch.dtype):
    # Manual IOByte computation is required since nvFuser input/outputs differ from baseline outputs (output, grad_out).
    nvf_inp_out = {
        # Inputs
        "input1": (size, dtype),
        "input2": (size, dtype),
        "weights": (size[1], dtype),
        "mean": (size[0], torch.float),
        "invstd": (size[0], torch.float),
        "grad_out": (size, dtype),
        "dropout_mask": (size, torch.bool),
        # Outputs
        "grad_in1": (size, dtype),
        "grad_in2": (size, dtype),
        "grad_weights": (size[1], dtype),
        "grad_bias": (size[1], dtype),
    }
    return compute_total_iobytes(nvf_inp_out)

@pytest.mark.parametrize("size", generate_input_sizes(dims=2))
@pytest.mark.parametrize("dtype", FLOAT_DTYPES)
@pytest.mark.inner_outer_persistent
def test_dropout_layernorm_bwd_nvf_benchmark(
    benchmark,
    size: tuple,
    dtype: torch.dtype,
    disable_validation: bool,
    disable_benchmarking: bool,
    eps: float = 1e-5,
):
    input1 = torch.randn(size, device="cuda", dtype=dtype, requires_grad=True)
    input2 = torch.randn(size, device="cuda", dtype=dtype, requires_grad=True)
    grads = torch.randn(size, device="cuda", dtype=dtype)
    weights = torch.randn(size[1], device="cuda", dtype=dtype, requires_grad=True)
    bias = torch.randn(size[1], device="cuda", dtype=dtype, requires_grad=True)

    dropout_p = 0.2
    dropout_mask = torch.lt(torch.rand(*size, device="cuda"), 1 - dropout_p)
    # Manually compute dropout for validation
    x = input2 + 1 / (1 - dropout_p) * dropout_mask * input1

    mean = x.to(torch.float).mean(dim=-1)
    variance = x.to(torch.float).var(dim=-1, unbiased=False)
    invstd = (1.0 / torch.sqrt(variance + eps)).unsqueeze(1)

    with FusionDefinition() as fd:
        dropout_layernorm_bwd_fusion(
            fd, torch_dtype_to_nvfuser_dtype(dtype), dropout_p=dropout_p
        )
    if not disable_validation:
        eager_output = torch.nn.functional.layer_norm(
            x.to(torch.double),
            input1.shape[1:],
            weight=weights.to(torch.double),
            bias=bias.to(torch.double),
        )

        eager_output.backward(grads.to(torch.double))
        fd.validate(
            [dropout_mask, mean, invstd, grads, weights, input1, input2],
            [input1.grad, input2.grad, weights.grad, bias.grad],
        )
    if not disable_benchmarking:
        run_benchmark(
            benchmark,
            fd.execute,
            [dropout_mask, mean, invstd, grads, weights, input1, input2],
        )

@pytest.mark.parametrize("executor", DEFAULT_EXECUTORS)
@pytest.mark.parametrize("size", generate_input_sizes(dims=2))
@pytest.mark.parametrize("dtype", FLOAT_DTYPES)
@pytest.mark.inner_outer_persistent
@pytest.mark.inner_persistent
def test_dropout_layernorm_bwd_baseline_benchmark(
    benchmark,
    size: tuple,
    dtype: torch.dtype,
    executor: str,
):
    if executor == "torchcompile":
        clear_dynamo_cache()

    dropout_p = 0.2
    input1 = torch.randn(size, device="cuda", dtype=dtype, requires_grad=True)
    input2 = torch.randn(size, device="cuda", dtype=dtype, requires_grad=True)
    grads = torch.randn(size, device="cuda", dtype=dtype)
    weights = torch.randn(size[1], device="cuda", dtype=dtype, requires_grad=True)
    bias = torch.randn(size[1], device="cuda", dtype=dtype, requires_grad=True)

    # Compile the fwd fn for torchcompile
    fwd_fn = with_executor(executor, dropout_layernorm)
    fwd_inputs = [input1, input2, weights, bias, dropout_p]
    outputs = fwd_fn(fwd_inputs)

    # Manually compute IOBytes: See PR #1725
    run_benchmark(
        benchmark,
        unary_bwd_torch,
        [outputs, grads, *fwd_inputs],
        iobytes=dropout_layernorm_bwd_iobytes(size, dtype),
    )


# SPDX-FileCopyrightText: Copyright (c) 2024-present NVIDIA CORPORATION & AFFILIATES.
# All rights reserved.
# SPDX-License-Identifier: BSD-3-Clause
import pytest
from nvfuser_direct import FusionDefinition, DataType
from nvfuser_direct.pytorch_utils import torch_dtype_to_nvfuser_dtype
from .core import (
    run_benchmark,
    clear_dynamo_cache,
    unary_bwd_torch,
    compute_total_iobytes,
    with_executor,
    DEFAULT_EXECUTORS,
)
import torch
from .global_params import generate_input_sizes, FLOAT_DTYPES, PROMOTE_DTYPES
from .torch_ops import dropout_layernorm

def dropout_layernorm_bwd_fusion(
    fd: FusionDefinition, dtype: DataType, dropout_p: float
) -> None:
    """
    Backward pass fusion definition for computing:
        output = layernorm (input2 + dropout (input1 p=dropout_p))

    Fusion inputs: inputs, dropout_mask, rms, grads, weights
    Fusion outputs: grad_input, grad_weights, grad_bias
    """
    T1 = fd.define_tensor(
        shape=[-1, -1], contiguity=[True, True], dtype=DataType.Bool, is_cpu=False
    )  # mask
    T2 = fd.define_tensor(
        shape=[-1], contiguity=[True], dtype=DataType.Float, is_cpu=False
    )  # mean
    T3 = fd.define_tensor(
        shape=[-1, 1], contiguity=[True, None], dtype=DataType.Float, is_cpu=False
    )  # invstd
    T4 = fd.define_tensor(
        shape=[-1, -1], contiguity=[True, True], dtype=dtype, is_cpu=False
    )  # grads
    T5 = fd.define_tensor(
        shape=[-1], contiguity=[True], dtype=dtype, is_cpu=False
    )  # weights
    T6 = fd.define_tensor(
        shape=[-1, -1], contiguity=[True, True], dtype=dtype, is_cpu=False
    )  # input1

    T7 = fd.define_tensor(
        shape=[-1, -1], contiguity=[True, True], dtype=dtype, is_cpu=False
    )  # input2

    if dtype in PROMOTE_DTYPES:
        T1 = fd.ops.cast(T1, dtype=DataType.Float)
        T4 = fd.ops.cast(T4, dtype=DataType.Float)
        T5 = fd.ops.cast(T5, dtype=DataType.Float)
        T6 = fd.ops.cast(T6, dtype=DataType.Float)
        T7 = fd.ops.cast(T7, dtype=DataType.Float)

    T9 = fd.ops.mul(T6, T1)
    S10 = fd.define_scalar(1 / (1 - dropout_p), dtype=DataType.Double)
    T11 = fd.ops.mul(T9, S10)
    T12 = fd.ops.add(T7, T11)

    T16 = fd.ops.broadcast_in_dim(T2, shape=[T6.size(0), 1], broadcast_dims=[0])
    V19 = T6.shape()
    T20 = fd.ops.broadcast_in_dim(T16, shape=V19, broadcast_dims=[0, 1])
    T21 = fd.ops.sub(T12, T20)
    T25 = fd.ops.broadcast_in_dim(T3, shape=V19, broadcast_dims=[0, 1])
    T26 = fd.ops.mul(T21, T25)
    T30 = fd.ops.broadcast_in_dim(T5, shape=V19, broadcast_dims=[1])
    T35 = fd.ops.sum(T4, dims=[0], keepdim=False, dtype=DataType.Null)

    T37 = fd.ops.mul(T4, T30)
    T38 = fd.ops.mul(T4, T26)
    T39 = fd.ops.sum(T38, dims=[0], keepdim=False, dtype=DataType.Null)

    T41 = fd.ops.mul(T37, T25)
    T42 = fd.ops.mul(T37, T21)
    T43 = fd.ops.sum(T42, dims=[1], keepdim=False, dtype=DataType.Null)
    T47 = fd.ops.broadcast_in_dim(T43, shape=[T6.size(0), 1], broadcast_dims=[0])
    T48 = fd.ops.neg(T41)
    T49 = fd.ops.sum(T48, dims=[1], keepdim=False, dtype=DataType.Null)
    T53 = fd.ops.broadcast_in_dim(T49, shape=[T6.size(0), 1], broadcast_dims=[0])
    S54 = fd.define_scalar(-0.500000, dtype=DataType.Double)
    T55 = fd.ops.mul(S54, T47)
    S56 = fd.define_scalar(3.00000, dtype=DataType.Double)
    T57 = fd.ops.pow(T3, S56)
    T58 = fd.ops.mul(T55, T57)
    T61 = fd.ops.sum(T53, dims=[1], keepdim=False, dtype=DataType.Null)
    T62 = fd.ops.sum(T58, dims=[1], keepdim=False, dtype=DataType.Null)
    T66 = fd.ops.broadcast_in_dim(T62, shape=[T6.size(0), 1], broadcast_dims=[0])
    T70 = fd.ops.broadcast_in_dim(T66, shape=V19, broadcast_dims=[0, 1])
    T74 = fd.ops.broadcast_in_dim(T2, shape=[T6.size(0), 1], broadcast_dims=[0])
    T78 = fd.ops.broadcast_in_dim(T74, shape=V19, broadcast_dims=[0, 1])
    S79 = fd.define_scalar(2.00000, dtype=DataType.Double)
    T80 = fd.ops.mul(S79, T70)
    T81 = fd.ops.sub(T12, T78)
    T82 = fd.ops.mul(T80, T81)
    S84 = fd.ops.reciprocal(T6.size(1))
    T85 = fd.ops.mul(T82, S84)
    T89 = fd.ops.broadcast_in_dim(T61, shape=[T6.size(0), 1], broadcast_dims=[0])
    T93 = fd.ops.broadcast_in_dim(T89, shape=V19, broadcast_dims=[0, 1])
    T95 = fd.ops.mul(S84, T93)
    T96 = fd.ops.add(T85, T95)
    T97 = fd.ops.add(T41, T96)

    T100 = fd.ops.mul(T97, S10)
    T101 = fd.ops.mul(T100, T1)

    if dtype in PROMOTE_DTYPES:
        T35 = fd.ops.cast(T35, dtype=dtype)
        T39 = fd.ops.cast(T39, dtype=dtype)
        T97 = fd.ops.cast(T97, dtype=dtype)
        T101 = fd.ops.cast(T101, dtype=dtype)

    fd.add_output(T101)
    fd.add_output(T97)
    fd.add_output(T39)
    fd.add_output(T35)

def dropout_layernorm_bwd_iobytes(size: tuple, dtype: torch.dtype):
    # Manual IOByte computation is required since nvFuser input/outputs differ from baseline outputs (output, grad_out).
    nvf_inp_out = {
        # Inputs
        "input1": (size, dtype),
        "input2": (size, dtype),
        "weights": (size[1], dtype),
        "mean": (size[0], torch.float),
        "invstd": (size[0], torch.float),
        "grad_out": (size, dtype),
        "dropout_mask": (size, torch.bool),
        # Outputs
        "grad_in1": (size, dtype),
        "grad_in2": (size, dtype),
        "grad_weights": (size[1], dtype),
        "grad_bias": (size[1], dtype),
    }
    return compute_total_iobytes(nvf_inp_out)


    T93 = fd.ops.broadcast_in_dim(T89, shape=V19, broadcast_dims=[0, 1])
    T95 = fd.ops.mul(S84, T93)
    T96 = fd.ops.add(T85, T95)
    T97 = fd.ops.add(T41, T96)

    T100 = fd.ops.mul(T97, S10)
    T101 = fd.ops.mul(T100, T1)

    if dtype in PROMOTE_DTYPES:
        T35 = fd.ops.cast(T35, dtype=dtype)
        T39 = fd.ops.cast(T39, dtype=dtype)
        T97 = fd.ops.cast(T97, dtype=dtype)
        T101 = fd.ops.cast(T101, dtype=dtype)

    fd.add_output(T101)
    fd.add_output(T97)
    fd.add_output(T39)
    fd.add_output(T35)

def dropout_layernorm_bwd_iobytes(size: tuple, dtype: torch.dtype):
    # Manual IOByte computation is required since nvFuser input/outputs differ from baseline outputs (output, grad_out).
    nvf_inp_out = {
        # Inputs
        "input1": (size, dtype),
        "input2": (size, dtype),
        "weights": (size[1], dtype),
        "mean": (size[0], torch.float),
        "invstd": (size[0], torch.float),
        "grad_out": (size, dtype),
        "dropout_mask": (size, torch.bool),
        # Outputs
        "grad_in1": (size, dtype),
        "grad_in2": (size, dtype),
        "grad_weights": (size[1], dtype),
        "grad_bias": (size[1], dtype),
    }
    return compute_total_iobytes(nvf_inp_out)

@pytest.mark.parametrize("size", generate_input_sizes(dims=2))
@pytest.mark.parametrize("dtype", FLOAT_DTYPES)
@pytest.mark.inner_outer_persistent
def test_dropout_layernorm_bwd_nvf_benchmark(
    benchmark,
    size: tuple,
    dtype: torch.dtype,
    disable_validation: bool,
    disable_benchmarking: bool,
    eps: float = 1e-5,
):
    input1 = torch.randn(size, device="cuda", dtype=dtype, requires_grad=True)
    input2 = torch.randn(size, device="cuda", dtype=dtype, requires_grad=True)
    grads = torch.randn(size, device="cuda", dtype=dtype)
    weights = torch.randn(size[1], device="cuda", dtype=dtype, requires_grad=True)
    bias = torch.randn(size[1], device="cuda", dtype=dtype, requires_grad=True)

    dropout_p = 0.2


        # Inputs
        "input1": (size, dtype),
        "input2": (size, dtype),
        "weights": (size[1], dtype),
        "mean": (size[0], torch.float),
        "invstd": (size[0], torch.float),
        "grad_out": (size, dtype),
        "dropout_mask": (size, torch.bool),
        # Outputs
        "grad_in1": (size, dtype),
        "grad_in2": (size, dtype),
        "grad_weights": (size[1], dtype),
        "grad_bias": (size[1], dtype),
    }
    return compute_total_iobytes(nvf_inp_out)

@pytest.mark.parametrize("size", generate_input_sizes(dims=2))
@pytest.mark.parametrize("dtype", FLOAT_DTYPES)
@pytest.mark.inner_outer_persistent
def test_dropout_layernorm_bwd_nvf_benchmark(
    benchmark,
    size: tuple,
    dtype: torch.dtype,
    disable_validation: bool,
    disable_benchmarking: bool,
    eps: float = 1e-5,
):
    input1 = torch.randn(size, device="cuda", dtype=dtype, requires_grad=True)
    input2 = torch.randn(size, device="cuda", dtype=dtype, requires_grad=True)
    grads = torch.randn(size, device="cuda", dtype=dtype)
    weights = torch.randn(size[1], device="cuda", dtype=dtype, requires_grad=True)
    bias = torch.randn(size[1], device="cuda", dtype=dtype, requires_grad=True)

    dropout_p = 0.2
    dropout_mask = torch.lt(torch.rand(*size, device="cuda"), 1 - dropout_p)
    # Manually compute dropout for validation
    x = input2 + 1 / (1 - dropout_p) * dropout_mask * input1

    mean = x.to(torch.float).mean(dim=-1)
    variance = x.to(torch.float).var(dim=-1, unbiased=False)
    invstd = (1.0 / torch.sqrt(variance + eps)).unsqueeze(1)

    with FusionDefinition() as fd:
        dropout_layernorm_bwd_fusion(
            fd, torch_dtype_to_nvfuser_dtype(dtype), dropout_p=dropout_p
        )
    if not disable_validation:
        eager_output = torch.nn.functional.layer_norm(
            x.to(torch.double),
            input1.shape[1:],
            weight=weights.to(torch.double),
            bias=bias.to(torch.double),
        )

        eager_output.backward(grads.to(torch.double))
        fd.validate(
            [dropout_mask, mean, invstd, grads, weights, input1, input2],
            [input1.grad, input2.grad, weights.grad, bias.grad],
        )
    if not disable_benchmarking:
        run_benchmark(
            benchmark,
            fd.execute,
            [dropout_mask, mean, invstd, grads, weights, input1, input2],
        )

@pytest.mark.parametrize("executor", DEFAULT_EXECUTORS)
@pytest.mark.parametrize("size", generate_input_sizes(dims=2))
@pytest.mark.parametrize("dtype", FLOAT_DTYPES)
@pytest.mark.inner_outer_persistent
@pytest.mark.inner_persistent
def test_dropout_layernorm_bwd_baseline_benchmark(
    benchmark,
    size: tuple,
    dtype: torch.dtype,
    executor: str,
):
    if executor == "torchcompile":
        clear_dynamo_cache()

    dropout_p = 0.2
    input1 = torch.randn(size, device="cuda", dtype=dtype, requires_grad=True)
    input2 = torch.randn(size, device="cuda", dtype=dtype, requires_grad=True)
    grads = torch.randn(size, device="cuda", dtype=dtype)


        )

        eager_output.backward(grads.to(torch.double))
        fd.validate(
            [dropout_mask, mean, invstd, grads, weights, input1, input2],
            [input1.grad, input2.grad, weights.grad, bias.grad],
        )
    if not disable_benchmarking:
        run_benchmark(
            benchmark,
            fd.execute,
            [dropout_mask, mean, invstd, grads, weights, input1, input2],
        )

@pytest.mark.parametrize("executor", DEFAULT_EXECUTORS)
@pytest.mark.parametrize("size", generate_input_sizes(dims=2))
@pytest.mark.parametrize("dtype", FLOAT_DTYPES)
@pytest.mark.inner_outer_persistent
@pytest.mark.inner_persistent
def test_dropout_layernorm_bwd_baseline_benchmark(
    benchmark,
    size: tuple,
    dtype: torch.dtype,
    executor: str,
):
    if executor == "torchcompile":
        clear_dynamo_cache()

    dropout_p = 0.2
    input1 = torch.randn(size, device="cuda", dtype=dtype, requires_grad=True)
    input2 = torch.randn(size, device="cuda", dtype=dtype, requires_grad=True)
    grads = torch.randn(size, device="cuda", dtype=dtype)
    weights = torch.randn(size[1], device="cuda", dtype=dtype, requires_grad=True)
    bias = torch.randn(size[1], device="cuda", dtype=dtype, requires_grad=True)

    # Compile the fwd fn for torchcompile
    fwd_fn = with_executor(executor, dropout_layernorm)
    fwd_inputs = [input1, input2, weights, bias, dropout_p]
    outputs = fwd_fn(fwd_inputs)

    # Manually compute IOBytes: See PR #1725
    run_benchmark(
        benchmark,
        unary_bwd_torch,
        [outputs, grads, *fwd_inputs],
        iobytes=dropout_layernorm_bwd_iobytes(size, dtype),
    )


# SPDX-FileCopyrightText: Copyright (c) 2025-present NVIDIA CORPORATION & AFFILIATES.
# All rights reserved.
# SPDX-License-Identifier: BSD-3-Clause

import os

def test_env_vars_are_set():
    for var in [
        "TORCHINDUCTOR_COORDINATE_DESCENT_TUNING",
        "TORCHINDUCTOR_COORDINATE_DESCENT_CHECK_ALL_DIRECTIONS",
    ]:
        assert os.environ[var] == "1"


# SPDX-FileCopyrightText: Copyright (c) 2024-present NVIDIA CORPORATION & AFFILIATES.
# All rights reserved.
# SPDX-License-Identifier: BSD-3-Clause
import pytest
from .core import (
    run_benchmark,
    clear_dynamo_cache,
    with_executor,
    DEFAULT_EXECUTORS,
    check_module_available,
)
import torch
from .global_params import FLOAT_DTYPES, BENCHMARK_CONFIG
from .torch_ops import rmsnorm_add
import itertools
from random import sample
from typing import List, Tuple

if check_module_available("flashinfer"):
    from flashinfer import fused_add_rmsnorm

def generate_input_sizes_rmsnorm_add(dims: int = 2) -> List[Tuple]:
    """
    Local version of generate_input_sizes specifically for rmsnorm_add_fwd tests.
    Uses batch sizes [1, 2, 4, 8, 16, 32, ..., 16384] (powers of 2).
    Uses hidden sizes [256, 512, 768, 1024, ..., 16384] (step size 256).
    """
    inputs = []

    if dims == 2:
        input_ranges = []

        step_size = 256
        # Use powers of 2 from 1 to 16384 for batch sizes
        batch_range = [2**i for i in range(0, 15)]  # [1, 2, 4, 8, 16, 32, ..., 16384]

        # Hidden range is always [256, 512, 768, 1024, ..., 16384]
        hidden_range = list(range(step_size, 16384 + 1, step_size))

        # Use the same hidden range for all batch sizes since it's already capped at 16384
        input_ranges.append((batch_range, hidden_range))

        for batch_range_subset, hidden_range_subset in input_ranges:
            inputs.extend(
                list(itertools.product(batch_range_subset, hidden_range_subset))
            )
    else:
        raise NotImplementedError(
            f"Generating input sizes of dimension {dims} is not implemented for rmsnorm_add"
        )

    if BENCHMARK_CONFIG["num_inputs"] is not None:
        inputs = sample(inputs, BENCHMARK_CONFIG["num_inputs"])

    return inputs

# flash infer does inplace update of inputs and residual tensors
# needs to explicitly return the updated tensors to correctly compute IO bytes
# https://github.com/flashinfer-ai/flashinfer/blob/ba2b4aa636c4ecf99981794767ffbf89267720cd/flashinfer/norm.py#L106
def flashinfer_rmsnorm_add_wrapper(inputs_list):
    inputs, weights, residual = inputs_list
    fused_add_rmsnorm(inputs, residual, weights)
    return inputs, residual

@pytest.mark.parametrize(
    "executor",
    DEFAULT_EXECUTORS
    + [
        pytest.param(
            "flashinfer",
            marks=pytest.mark.skipif(
                not check_module_available("flashinfer"),
                reason="flashinfer executor is not available on this device",
            ),
        )
    ],
)
@pytest.mark.parametrize("size", generate_input_sizes_rmsnorm_add(dims=2))
@pytest.mark.parametrize("dtype", FLOAT_DTYPES)
@pytest.mark.inner_persistent
def test_rmsnorm_add_fwd_baseline_benchmark(
    benchmark,
    size: tuple,
    dtype: torch.dtype,
    executor: str,
):
    if executor == "torchcompile":
        clear_dynamo_cache()
    inputs = torch.randn(size, device="cuda", dtype=dtype, requires_grad=False)
    weights = torch.randn(size[1], device="cuda", dtype=dtype, requires_grad=False)
    residual = torch.randn(size, device="cuda", dtype=dtype, requires_grad=False)
    # requires flashinfer to be installed
    # `pip install flashinfer-python` works in pjnl container
    if executor == "flashinfer":
        benchmark_fn = flashinfer_rmsnorm_add_wrapper
    else:
        benchmark_fn = with_executor(executor, rmsnorm_add)
    run_benchmark(
        benchmark,
        benchmark_fn,
        [inputs, weights, residual],
    )


# SPDX-License-Identifier: BSD-3-Clause
import pytest
from .core import (
    run_benchmark,
    clear_dynamo_cache,
    with_executor,
    DEFAULT_EXECUTORS,
    check_module_available,
)
import torch
from .global_params import FLOAT_DTYPES, BENCHMARK_CONFIG
from .torch_ops import rmsnorm_add
import itertools
from random import sample
from typing import List, Tuple

if check_module_available("flashinfer"):
    from flashinfer import fused_add_rmsnorm

def generate_input_sizes_rmsnorm_add(dims: int = 2) -> List[Tuple]:
    """
    Local version of generate_input_sizes specifically for rmsnorm_add_fwd tests.
    Uses batch sizes [1, 2, 4, 8, 16, 32, ..., 16384] (powers of 2).
    Uses hidden sizes [256, 512, 768, 1024, ..., 16384] (step size 256).
    """
    inputs = []

    if dims == 2:
        input_ranges = []

        step_size = 256
        # Use powers of 2 from 1 to 16384 for batch sizes
        batch_range = [2**i for i in range(0, 15)]  # [1, 2, 4, 8, 16, 32, ..., 16384]

        # Hidden range is always [256, 512, 768, 1024, ..., 16384]
        hidden_range = list(range(step_size, 16384 + 1, step_size))

        # Use the same hidden range for all batch sizes since it's already capped at 16384
        input_ranges.append((batch_range, hidden_range))

        for batch_range_subset, hidden_range_subset in input_ranges:
            inputs.extend(
                list(itertools.product(batch_range_subset, hidden_range_subset))
            )
    else:
        raise NotImplementedError(
            f"Generating input sizes of dimension {dims} is not implemented for rmsnorm_add"
        )

    if BENCHMARK_CONFIG["num_inputs"] is not None:
        inputs = sample(inputs, BENCHMARK_CONFIG["num_inputs"])

    return inputs

# flash infer does inplace update of inputs and residual tensors
# needs to explicitly return the updated tensors to correctly compute IO bytes
# https://github.com/flashinfer-ai/flashinfer/blob/ba2b4aa636c4ecf99981794767ffbf89267720cd/flashinfer/norm.py#L106
def flashinfer_rmsnorm_add_wrapper(inputs_list):
    inputs, weights, residual = inputs_list
    fused_add_rmsnorm(inputs, residual, weights)
    return inputs, residual

@pytest.mark.parametrize(
    "executor",
    DEFAULT_EXECUTORS
    + [
        pytest.param(
            "flashinfer",
            marks=pytest.mark.skipif(
                not check_module_available("flashinfer"),
                reason="flashinfer executor is not available on this device",


        input_ranges.append((batch_range, hidden_range))

        for batch_range_subset, hidden_range_subset in input_ranges:
            inputs.extend(
                list(itertools.product(batch_range_subset, hidden_range_subset))
            )
    else:
        raise NotImplementedError(
            f"Generating input sizes of dimension {dims} is not implemented for rmsnorm_add"
        )

    if BENCHMARK_CONFIG["num_inputs"] is not None:
        inputs = sample(inputs, BENCHMARK_CONFIG["num_inputs"])

    return inputs

# flash infer does inplace update of inputs and residual tensors
# needs to explicitly return the updated tensors to correctly compute IO bytes
# https://github.com/flashinfer-ai/flashinfer/blob/ba2b4aa636c4ecf99981794767ffbf89267720cd/flashinfer/norm.py#L106
def flashinfer_rmsnorm_add_wrapper(inputs_list):
    inputs, weights, residual = inputs_list
    fused_add_rmsnorm(inputs, residual, weights)
    return inputs, residual

@pytest.mark.parametrize(
    "executor",
    DEFAULT_EXECUTORS
    + [
        pytest.param(
            "flashinfer",
            marks=pytest.mark.skipif(
                not check_module_available("flashinfer"),
                reason="flashinfer executor is not available on this device",
            ),
        )
    ],
)
@pytest.mark.parametrize("size", generate_input_sizes_rmsnorm_add(dims=2))
@pytest.mark.parametrize("dtype", FLOAT_DTYPES)
@pytest.mark.inner_persistent
def test_rmsnorm_add_fwd_baseline_benchmark(
    benchmark,


    fused_add_rmsnorm(inputs, residual, weights)
    return inputs, residual

@pytest.mark.parametrize(
    "executor",
    DEFAULT_EXECUTORS
    + [
        pytest.param(
            "flashinfer",
            marks=pytest.mark.skipif(
                not check_module_available("flashinfer"),
                reason="flashinfer executor is not available on this device",
            ),
        )
    ],
)
@pytest.mark.parametrize("size", generate_input_sizes_rmsnorm_add(dims=2))
@pytest.mark.parametrize("dtype", FLOAT_DTYPES)
@pytest.mark.inner_persistent
def test_rmsnorm_add_fwd_baseline_benchmark(
    benchmark,
    size: tuple,
    dtype: torch.dtype,
    executor: str,
):
    if executor == "torchcompile":
        clear_dynamo_cache()
    inputs = torch.randn(size, device="cuda", dtype=dtype, requires_grad=False)
    weights = torch.randn(size[1], device="cuda", dtype=dtype, requires_grad=False)
    residual = torch.randn(size, device="cuda", dtype=dtype, requires_grad=False)
    # requires flashinfer to be installed
    # `pip install flashinfer-python` works in pjnl container
    if executor == "flashinfer":
        benchmark_fn = flashinfer_rmsnorm_add_wrapper
    else:
        benchmark_fn = with_executor(executor, rmsnorm_add)
    run_benchmark(
        benchmark,
        benchmark_fn,
        [inputs, weights, residual],
    )


benchmarks.python.conftest

# SPDX-FileCopyrightText: Copyright (c) 2024-present NVIDIA CORPORATION & AFFILIATES.
# All rights reserved.
# SPDX-License-Identifier: BSD-3-Clause
import pytest
from .core import BENCHMARK_CONFIG
from nvfuser_direct.pytorch_utils import DEVICE_PROPERTIES
import os

ORIGINAL_ENV_VARS = {}

def pytest_sessionstart(session):
    for var, value in [
        ("TORCHINDUCTOR_COORDINATE_DESCENT_TUNING", "1"),
        ("TORCHINDUCTOR_COORDINATE_DESCENT_CHECK_ALL_DIRECTIONS", "1"),
    ]:
        ORIGINAL_ENV_VARS[var] = os.environ.get(var)
        os.environ[var] = value

def pytest_sessionfinish(session):
    for var, value in ORIGINAL_ENV_VARS.items():
        if value is not None:
            os.environ[var] = value
        else:
            os.environ.pop(var, None)

def pytest_addoption(parser):
    parser.addoption(
        "--disable-validation",
        action="store_true",
        help="Disable output validation in benchmarks.",
    )
    parser.addoption(
        "--disable-benchmarking",
        action="store_true",
        help="Disable benchmarking.",
    )
    parser.addoption(
        "--benchmark-eager",
        action="store_true",
        help="Benchmarks torch eager mode.",
    )
    parser.addoption(
        "--benchmark-thunder",
        action="store_true",
        help="Benchmarks thunder jit.",
    )
    parser.addoption(
        "--benchmark-torchcompile",
        action="store_true",
        help="Benchmarks torch.compile mode.",
    )
    parser.addoption(
        "--benchmark-thunder-torchcompile",
        action="store_true",
        help="Benchmarks torch.compile mode.",
    )
    parser.addoption(
        "--benchmark-flashinfer",
        action="store_true",
        help="Benchmarks flashinfer mode.",
    )
    parser.addoption(
        "--benchmark-quack",
        action="store_true",
        help="Benchmarks quack mode.",
    )
    # pytest-benchmark does not have CLI options to set rounds/warmup_rounds for benchmark.pedantic.
    # The following two options are used to overwrite the default values through CLI.
    parser.addoption(
        "--benchmark-rounds",
        action="store",
        default=10,
        help="Number of rounds for each benchmark.",
    )

    parser.addoption(
        "--benchmark-warmup-rounds",
        action="store",
        default=1,
        help="Number of warmup rounds for each benchmark.",
    )

    parser.addoption(
        "--benchmark-num-inputs",
        action="store",
        default=None,
        help="Number of inputs to randomly sample for each benchmark.",
    )

    parser.addoption(
        "--with-nsys",
        action="store_true",
        default=False,
        help="Run benchmark scripts with nsys. Disable all other profilers.",
    )

@pytest.fixture
def disable_validation(request):
    return request.config.getoption("--disable-validation")

@pytest.fixture
def disable_benchmarking(request):
    return request.config.getoption("--disable-benchmarking")

def pytest_make_parametrize_id(val, argname):
    if isinstance(val, tuple):
        return f'{argname}=[{"_".join(str(v) for v in val)}]'
    return f"{argname}={repr(val)}"

def pytest_benchmark_update_machine_info(config, machine_info):
    machine_info.update(DEVICE_PROPERTIES)

def pytest_configure(config):
    BENCHMARK_CONFIG["rounds"] = int(config.getoption("--benchmark-rounds"))
    BENCHMARK_CONFIG["warmup_rounds"] = int(
        config.getoption("--benchmark-warmup-rounds")
    )
    BENCHMARK_CONFIG["with_nsys"] = config.getoption("--with-nsys")

    if config.getoption("--benchmark-num-inputs"):
        BENCHMARK_CONFIG["num_inputs"] = int(config.getoption("--benchmark-num-inputs"))

    # Scheduler markers may become stale and are not 100% accurate.
    config.addinivalue_line(
        "markers",
        "inner_outer_persistent: mark tests using inner_outer_persistent scheduler if not being segmented.",
    )
    config.addinivalue_line(
        "markers",
        "inner_persistent: mark tests using inner_persistent scheduler if not being segmented.",
    )
    config.addinivalue_line(
        "markers",
        "outer_persistent: mark tests using outer_persistent scheduler if not being segmented.",
    )
    config.addinivalue_line(
        "markers",
        "reduction: mark tests using reduction scheduler if not being segmented.",
    )
    config.addinivalue_line(
        "markers",
        "matmul: mark tests using matmul scheduler if not being segmented.",
    )
    config.addinivalue_line(
        "markers",
        "resize: mark tests using resize scheduler if not being segmented.",
    )
    config.addinivalue_line(
        "markers",
        "transpose: mark tests using transpose scheduler if not being segmented.",
    )
    config.addinivalue_line(
        "markers",
        "pointwise: mark tests using pointwise scheduler if not being segmented.",
    )

def pytest_collection_modifyitems(session, config, items):
    """
    The baseline benchmarks use `executor` parameter with
    values ["eager", "torchcompile", "thunder", "thunder-torchcompile"] that are optionally
    run using `--benchmark-{executor}` flag. They are skipped by
    default.
    """

    from nvfuser_direct.pytorch_utils import retry_on_oom_or_skip_test

    executors = [
        "eager",
        "torchcompile",
        "thunder",
        "thunder-torchcompile",
        "flashinfer",
        "quack",
    ]

    def get_test_executor(item) -> str | None:
        if hasattr(item, "callspec") and "executor" in item.callspec.params:
            test_executor = item.callspec.params["executor"]
            assert (
                test_executor in executors
            ), f"Expected executor to be one of 'eager', 'torchcompile', 'thunder', 'thunder-torchcompile', found {test_executor}."
            return test_executor
        return None

    executors_to_skip = []

    for executor in executors:
        if not config.getoption(f"--benchmark-{executor}"):
            executors_to_skip.append(executor)

    for item in items:
        item.obj = retry_on_oom_or_skip_test(item.obj)

        test_executor = get_test_executor(item)

        if test_executor is not None and test_executor in executors_to_skip:
            item.add_marker(
                pytest.mark.skip(
                    reason=f"need --benchmark-{test_executor} option to run."
                )
            )


# SPDX-FileCopyrightText: Copyright (c) 2024-present NVIDIA CORPORATION & AFFILIATES.
# All rights reserved.
# SPDX-License-Identifier: BSD-3-Clause
import pytest
from .core import BENCHMARK_CONFIG
from nvfuser_direct.pytorch_utils import DEVICE_PROPERTIES
import os

ORIGINAL_ENV_VARS = {}

def pytest_sessionstart(session):
    for var, value in [
        ("TORCHINDUCTOR_COORDINATE_DESCENT_TUNING", "1"),
        ("TORCHINDUCTOR_COORDINATE_DESCENT_CHECK_ALL_DIRECTIONS", "1"),
    ]:
        ORIGINAL_ENV_VARS[var] = os.environ.get(var)
        os.environ[var] = value

def pytest_sessionfinish(session):
    for var, value in ORIGINAL_ENV_VARS.items():
        if value is not None:
            os.environ[var] = value
        else:
            os.environ.pop(var, None)

def pytest_addoption(parser):
    parser.addoption(
        "--disable-validation",
        action="store_true",
        help="Disable output validation in benchmarks.",
    )
    parser.addoption(
        "--disable-benchmarking",
        action="store_true",
        help="Disable benchmarking.",


# SPDX-FileCopyrightText: Copyright (c) 2024-present NVIDIA CORPORATION & AFFILIATES.
# All rights reserved.
# SPDX-License-Identifier: BSD-3-Clause
import pytest
from .core import BENCHMARK_CONFIG
from nvfuser_direct.pytorch_utils import DEVICE_PROPERTIES
import os

ORIGINAL_ENV_VARS = {}

def pytest_sessionstart(session):
    for var, value in [
        ("TORCHINDUCTOR_COORDINATE_DESCENT_TUNING", "1"),
        ("TORCHINDUCTOR_COORDINATE_DESCENT_CHECK_ALL_DIRECTIONS", "1"),
    ]:
        ORIGINAL_ENV_VARS[var] = os.environ.get(var)
        os.environ[var] = value

def pytest_sessionfinish(session):
    for var, value in ORIGINAL_ENV_VARS.items():
        if value is not None:
            os.environ[var] = value
        else:
            os.environ.pop(var, None)

def pytest_addoption(parser):
    parser.addoption(
        "--disable-validation",
        action="store_true",
        help="Disable output validation in benchmarks.",
    )
    parser.addoption(
        "--disable-benchmarking",
        action="store_true",
        help="Disable benchmarking.",
    )
    parser.addoption(
        "--benchmark-eager",
        action="store_true",
        help="Benchmarks torch eager mode.",
    )
    parser.addoption(
        "--benchmark-thunder",


ORIGINAL_ENV_VARS = {}

def pytest_sessionstart(session):
    for var, value in [
        ("TORCHINDUCTOR_COORDINATE_DESCENT_TUNING", "1"),
        ("TORCHINDUCTOR_COORDINATE_DESCENT_CHECK_ALL_DIRECTIONS", "1"),
    ]:
        ORIGINAL_ENV_VARS[var] = os.environ.get(var)
        os.environ[var] = value

def pytest_sessionfinish(session):
    for var, value in ORIGINAL_ENV_VARS.items():
        if value is not None:
            os.environ[var] = value
        else:
            os.environ.pop(var, None)

def pytest_addoption(parser):
    parser.addoption(
        "--disable-validation",
        action="store_true",
        help="Disable output validation in benchmarks.",
    )
    parser.addoption(
        "--disable-benchmarking",
        action="store_true",
        help="Disable benchmarking.",
    )
    parser.addoption(
        "--benchmark-eager",
        action="store_true",
        help="Benchmarks torch eager mode.",
    )
    parser.addoption(
        "--benchmark-thunder",
        action="store_true",
        help="Benchmarks thunder jit.",
    )
    parser.addoption(
        "--benchmark-torchcompile",
        action="store_true",
        help="Benchmarks torch.compile mode.",
    )
    parser.addoption(
        "--benchmark-thunder-torchcompile",
        action="store_true",
        help="Benchmarks torch.compile mode.",
    )
    parser.addoption(
        "--benchmark-flashinfer",
        action="store_true",
        help="Benchmarks flashinfer mode.",
    )
    parser.addoption(
        "--benchmark-quack",
        action="store_true",
        help="Benchmarks quack mode.",
    )
    # pytest-benchmark does not have CLI options to set rounds/warmup_rounds for benchmark.pedantic.
    # The following two options are used to overwrite the default values through CLI.
    parser.addoption(
        "--benchmark-rounds",
        action="store",
        default=10,
        help="Number of rounds for each benchmark.",
    )

    parser.addoption(
        "--benchmark-warmup-rounds",
        action="store",
        default=1,
        help="Number of warmup rounds for each benchmark.",
    )

    parser.addoption(
        "--benchmark-num-inputs",
        action="store",
        default=None,
        help="Number of inputs to randomly sample for each benchmark.",
    )

    parser.addoption(
        "--with-nsys",
        action="store_true",
        default=False,
        help="Run benchmark scripts with nsys. Disable all other profilers.",
    )

@pytest.fixture
def disable_validation(request):
    return request.config.getoption("--disable-validation")

@pytest.fixture
def disable_benchmarking(request):
    return request.config.getoption("--disable-benchmarking")

def pytest_make_parametrize_id(val, argname):
    if isinstance(val, tuple):
        return f'{argname}=[{"_".join(str(v) for v in val)}]'
    return f"{argname}={repr(val)}"

def pytest_benchmark_update_machine_info(config, machine_info):
    machine_info.update(DEVICE_PROPERTIES)


        default=1,
        help="Number of warmup rounds for each benchmark.",
    )

    parser.addoption(
        "--benchmark-num-inputs",
        action="store",
        default=None,
        help="Number of inputs to randomly sample for each benchmark.",
    )

    parser.addoption(
        "--with-nsys",
        action="store_true",
        default=False,
        help="Run benchmark scripts with nsys. Disable all other profilers.",
    )

@pytest.fixture
def disable_validation(request):
    return request.config.getoption("--disable-validation")

@pytest.fixture
def disable_benchmarking(request):
    return request.config.getoption("--disable-benchmarking")

def pytest_make_parametrize_id(val, argname):
    if isinstance(val, tuple):
        return f'{argname}=[{"_".join(str(v) for v in val)}]'
    return f"{argname}={repr(val)}"

def pytest_benchmark_update_machine_info(config, machine_info):
    machine_info.update(DEVICE_PROPERTIES)

def pytest_configure(config):
    BENCHMARK_CONFIG["rounds"] = int(config.getoption("--benchmark-rounds"))
    BENCHMARK_CONFIG["warmup_rounds"] = int(


        "--benchmark-num-inputs",
        action="store",
        default=None,
        help="Number of inputs to randomly sample for each benchmark.",
    )

    parser.addoption(
        "--with-nsys",
        action="store_true",
        default=False,
        help="Run benchmark scripts with nsys. Disable all other profilers.",
    )

@pytest.fixture
def disable_validation(request):
    return request.config.getoption("--disable-validation")

@pytest.fixture
def disable_benchmarking(request):
    return request.config.getoption("--disable-benchmarking")

def pytest_make_parametrize_id(val, argname):
    if isinstance(val, tuple):
        return f'{argname}=[{"_".join(str(v) for v in val)}]'
    return f"{argname}={repr(val)}"

def pytest_benchmark_update_machine_info(config, machine_info):
    machine_info.update(DEVICE_PROPERTIES)

def pytest_configure(config):
    BENCHMARK_CONFIG["rounds"] = int(config.getoption("--benchmark-rounds"))
    BENCHMARK_CONFIG["warmup_rounds"] = int(
        config.getoption("--benchmark-warmup-rounds")
    )
    BENCHMARK_CONFIG["with_nsys"] = config.getoption("--with-nsys")

    if config.getoption("--benchmark-num-inputs"):


    )

    parser.addoption(
        "--with-nsys",
        action="store_true",
        default=False,
        help="Run benchmark scripts with nsys. Disable all other profilers.",
    )

@pytest.fixture
def disable_validation(request):
    return request.config.getoption("--disable-validation")

@pytest.fixture
def disable_benchmarking(request):
    return request.config.getoption("--disable-benchmarking")

def pytest_make_parametrize_id(val, argname):
    if isinstance(val, tuple):
        return f'{argname}=[{"_".join(str(v) for v in val)}]'
    return f"{argname}={repr(val)}"

def pytest_benchmark_update_machine_info(config, machine_info):
    machine_info.update(DEVICE_PROPERTIES)

def pytest_configure(config):
    BENCHMARK_CONFIG["rounds"] = int(config.getoption("--benchmark-rounds"))
    BENCHMARK_CONFIG["warmup_rounds"] = int(
        config.getoption("--benchmark-warmup-rounds")
    )
    BENCHMARK_CONFIG["with_nsys"] = config.getoption("--with-nsys")

    if config.getoption("--benchmark-num-inputs"):
        BENCHMARK_CONFIG["num_inputs"] = int(config.getoption("--benchmark-num-inputs"))

    # Scheduler markers may become stale and are not 100% accurate.
    config.addinivalue_line(
        "markers",
        "inner_outer_persistent: mark tests using inner_outer_persistent scheduler if not being segmented.",


        help="Run benchmark scripts with nsys. Disable all other profilers.",
    )

@pytest.fixture
def disable_validation(request):
    return request.config.getoption("--disable-validation")

@pytest.fixture
def disable_benchmarking(request):
    return request.config.getoption("--disable-benchmarking")

def pytest_make_parametrize_id(val, argname):
    if isinstance(val, tuple):
        return f'{argname}=[{"_".join(str(v) for v in val)}]'
    return f"{argname}={repr(val)}"

def pytest_benchmark_update_machine_info(config, machine_info):
    machine_info.update(DEVICE_PROPERTIES)

def pytest_configure(config):
    BENCHMARK_CONFIG["rounds"] = int(config.getoption("--benchmark-rounds"))
    BENCHMARK_CONFIG["warmup_rounds"] = int(
        config.getoption("--benchmark-warmup-rounds")
    )
    BENCHMARK_CONFIG["with_nsys"] = config.getoption("--with-nsys")

    if config.getoption("--benchmark-num-inputs"):
        BENCHMARK_CONFIG["num_inputs"] = int(config.getoption("--benchmark-num-inputs"))

    # Scheduler markers may become stale and are not 100% accurate.
    config.addinivalue_line(
        "markers",
        "inner_outer_persistent: mark tests using inner_outer_persistent scheduler if not being segmented.",
    )
    config.addinivalue_line(
        "markers",
        "inner_persistent: mark tests using inner_persistent scheduler if not being segmented.",


@pytest.fixture
def disable_validation(request):
    return request.config.getoption("--disable-validation")

@pytest.fixture
def disable_benchmarking(request):
    return request.config.getoption("--disable-benchmarking")

def pytest_make_parametrize_id(val, argname):
    if isinstance(val, tuple):
        return f'{argname}=[{"_".join(str(v) for v in val)}]'
    return f"{argname}={repr(val)}"

def pytest_benchmark_update_machine_info(config, machine_info):
    machine_info.update(DEVICE_PROPERTIES)

def pytest_configure(config):
    BENCHMARK_CONFIG["rounds"] = int(config.getoption("--benchmark-rounds"))
    BENCHMARK_CONFIG["warmup_rounds"] = int(
        config.getoption("--benchmark-warmup-rounds")
    )
    BENCHMARK_CONFIG["with_nsys"] = config.getoption("--with-nsys")

    if config.getoption("--benchmark-num-inputs"):
        BENCHMARK_CONFIG["num_inputs"] = int(config.getoption("--benchmark-num-inputs"))

    # Scheduler markers may become stale and are not 100% accurate.
    config.addinivalue_line(
        "markers",
        "inner_outer_persistent: mark tests using inner_outer_persistent scheduler if not being segmented.",
    )
    config.addinivalue_line(
        "markers",
        "inner_persistent: mark tests using inner_persistent scheduler if not being segmented.",
    )
    config.addinivalue_line(
        "markers",
        "outer_persistent: mark tests using outer_persistent scheduler if not being segmented.",
    )
    config.addinivalue_line(
        "markers",
        "reduction: mark tests using reduction scheduler if not being segmented.",
    )
    config.addinivalue_line(
        "markers",
        "matmul: mark tests using matmul scheduler if not being segmented.",
    )
    config.addinivalue_line(
        "markers",
        "resize: mark tests using resize scheduler if not being segmented.",
    )
    config.addinivalue_line(
        "markers",
        "transpose: mark tests using transpose scheduler if not being segmented.",
    )
    config.addinivalue_line(
        "markers",
        "pointwise: mark tests using pointwise scheduler if not being segmented.",
    )

def pytest_collection_modifyitems(session, config, items):
    """
    The baseline benchmarks use `executor` parameter with
    values ["eager", "torchcompile", "thunder", "thunder-torchcompile"] that are optionally
    run using `--benchmark-{executor}` flag. They are skipped by
    default.
    """

    from nvfuser_direct.pytorch_utils import retry_on_oom_or_skip_test

    executors = [
        "eager",
        "torchcompile",
        "thunder",
        "thunder-torchcompile",
        "flashinfer",
        "quack",
    ]


        "reduction: mark tests using reduction scheduler if not being segmented.",
    )
    config.addinivalue_line(
        "markers",
        "matmul: mark tests using matmul scheduler if not being segmented.",
    )
    config.addinivalue_line(
        "markers",
        "resize: mark tests using resize scheduler if not being segmented.",
    )
    config.addinivalue_line(
        "markers",
        "transpose: mark tests using transpose scheduler if not being segmented.",
    )
    config.addinivalue_line(
        "markers",
        "pointwise: mark tests using pointwise scheduler if not being segmented.",
    )

def pytest_collection_modifyitems(session, config, items):
    """
    The baseline benchmarks use `executor` parameter with
    values ["eager", "torchcompile", "thunder", "thunder-torchcompile"] that are optionally
    run using `--benchmark-{executor}` flag. They are skipped by
    default.
    """

    from nvfuser_direct.pytorch_utils import retry_on_oom_or_skip_test

    executors = [
        "eager",
        "torchcompile",
        "thunder",
        "thunder-torchcompile",
        "flashinfer",
        "quack",
    ]

    def get_test_executor(item) -> str | None:
        if hasattr(item, "callspec") and "executor" in item.callspec.params:
            test_executor = item.callspec.params["executor"]
            assert (
                test_executor in executors
            ), f"Expected executor to be one of 'eager', 'torchcompile', 'thunder', 'thunder-torchcompile', found {test_executor}."
            return test_executor
        return None

    executors_to_skip = []

    for executor in executors:
        if not config.getoption(f"--benchmark-{executor}"):
            executors_to_skip.append(executor)

    for item in items:
        item.obj = retry_on_oom_or_skip_test(item.obj)

        test_executor = get_test_executor(item)

        if test_executor is not None and test_executor in executors_to_skip:
            item.add_marker(
                pytest.mark.skip(
                    reason=f"need --benchmark-{test_executor} option to run."
                )
            )


def pytest_collection_modifyitems(session, config, items):
    """
    The baseline benchmarks use `executor` parameter with
    values ["eager", "torchcompile", "thunder", "thunder-torchcompile"] that are optionally
    run using `--benchmark-{executor}` flag. They are skipped by
    default.
    """

    from nvfuser_direct.pytorch_utils import retry_on_oom_or_skip_test

    executors = [
        "eager",
        "torchcompile",
        "thunder",
        "thunder-torchcompile",
        "flashinfer",
        "quack",
    ]

    def get_test_executor(item) -> str | None:
        if hasattr(item, "callspec") and "executor" in item.callspec.params:
            test_executor = item.callspec.params["executor"]
            assert (
                test_executor in executors
            ), f"Expected executor to be one of 'eager', 'torchcompile', 'thunder', 'thunder-torchcompile', found {test_executor}."
            return test_executor
        return None

    executors_to_skip = []

    for executor in executors:
        if not config.getoption(f"--benchmark-{executor}"):
            executors_to_skip.append(executor)

    for item in items:
        item.obj = retry_on_oom_or_skip_test(item.obj)

        test_executor = get_test_executor(item)

        if test_executor is not None and test_executor in executors_to_skip:
            item.add_marker(
                pytest.mark.skip(
                    reason=f"need --benchmark-{test_executor} option to run."
                )
            )