    out = fd.ops.matmul(a, b)
    fd.add_output(out)

def load_matmul_problems():
    with open(os.path.join(os.path.dirname(__file__), "matmul_problems.csv"), "r") as f:
        reader = csv.reader(f)
        next(reader, None)  # skip header row
        rows = list((int(m), int(n), int(k), layout) for m, n, k, layout in reader)

        def row_mem(row):
            """Compute gmem bytes used in a half-precision GEMM

            This computes only the space required for operands and output,
            ignoring intermediates like split-K buffers and assuming no bias
            terms.
            """
            m, n, k, _ = row
            return ((m + n) * k + m * n) * 2

        def three_way_cmp(a, b) -> int:
            """Perform a three-way comparison like the Python2 cmp() function

            This returns 0 if a == b, -1 if a < b, and 1 if a > b
            """
            return int(a > b) - int(a < b)

        def mem_cmp(row1, row2):
            """Compare two rows based on memory use"""
            return three_way_cmp(row_mem(row1), row_mem(row2))

        # Reverse sort by expected memory use to avoid fragmentation
        rows.sort(key=functools.cmp_to_key(mem_cmp), reverse=True)

        return rows

def maybe_skip_oom_case(m: int, n: int, k: int):
    expected_mem = (m * k + n * k + m * n) * 2  # operands plus output
    expected_mem *= 2  # account for multiple runs/deferred frees

    _, total = torch.cuda.mem_get_info()
    max_mem = total * 0.9
    if expected_mem > max_mem:
        pytest.skip(
            f"Case takes more than {max_mem / (2 ** 30): .2f} GiB. Skipping to avoid OOM"


        next(reader, None)  # skip header row
        rows = list((int(m), int(n), int(k), layout) for m, n, k, layout in reader)

        def row_mem(row):
            """Compute gmem bytes used in a half-precision GEMM

            This computes only the space required for operands and output,
            ignoring intermediates like split-K buffers and assuming no bias
            terms.
            """
            m, n, k, _ = row
            return ((m + n) * k + m * n) * 2

        def three_way_cmp(a, b) -> int:
            """Perform a three-way comparison like the Python2 cmp() function

            This returns 0 if a == b, -1 if a < b, and 1 if a > b
            """
            return int(a > b) - int(a < b)

        def mem_cmp(row1, row2):
            """Compare two rows based on memory use"""
            return three_way_cmp(row_mem(row1), row_mem(row2))

        # Reverse sort by expected memory use to avoid fragmentation
        rows.sort(key=functools.cmp_to_key(mem_cmp), reverse=True)

        return rows

def maybe_skip_oom_case(m: int, n: int, k: int):
    expected_mem = (m * k + n * k + m * n) * 2  # operands plus output
    expected_mem *= 2  # account for multiple runs/deferred frees

    _, total = torch.cuda.mem_get_info()
    max_mem = total * 0.9
    if expected_mem > max_mem:
        pytest.skip(
            f"Case takes more than {max_mem / (2 ** 30): .2f} GiB. Skipping to avoid OOM"
        )

@pytest.mark.parametrize("half_reduction", [False, True], ids=["fullred", "halfred"])


# 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
import numpy as np
from .torch_ops import layernorm

def layernorm_fwd_fusion(
    fd: FusionDefinition,
    dtype: DataType,
    eps: float = 1e-5,
) -> None:
    T0 = fd.define_tensor(
        shape=[-1, -1], contiguity=[True, True], dtype=dtype, is_cpu=False
    )
    T1 = fd.define_tensor(shape=[-1], contiguity=[True], dtype=dtype, is_cpu=False)
    T2 = fd.define_tensor(shape=[-1], contiguity=[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.cast(T2, dtype=DataType.Float)

    T3, T4 = fd.ops.var_mean(T0, dims=[1], correction=0, keepdim=False)

    T7 = fd.ops.broadcast_in_dim(T3, shape=[T0.size(0), 1], broadcast_dims=[0])
    T11 = fd.ops.broadcast_in_dim(T4, shape=[T0.size(0), 1], broadcast_dims=[0])

    S12 = fd.define_scalar(eps, dtype=DataType.Double)
    T13 = fd.ops.add(T7, S12)
    T14 = fd.ops.rsqrt(T13)

    V17 = T0.shape()
    T18 = fd.ops.broadcast_in_dim(T11, shape=V17, broadcast_dims=[0, 1])
    T19 = fd.ops.sub(T0, T18)
    T23 = fd.ops.broadcast_in_dim(T14, shape=V17, broadcast_dims=[0, 1])
    T24 = fd.ops.mul(T19, T23)

    T25 = fd.ops.broadcast_in_dim(T1, shape=V17, broadcast_dims=[1])
    T26 = fd.ops.mul(T24, T25)
    T27 = fd.ops.broadcast_in_dim(T2, shape=V17, broadcast_dims=[1])
    T28 = fd.ops.add(T26, T27)

    if dtype in PROMOTE_DTYPES:
        T28 = fd.ops.cast(T28, dtype=dtype)

    fd.add_output(T28)
    fd.add_output(T4)
    fd.add_output(T14)

def layernorm_fwd_iobytes(size: tuple, dtype: torch.dtype):
    # Manual IOBytes computation required since nvFuser outputs (out, mean, invstd) differs from baselines (out)
    # Total IO bytes = in_tensor (size, dtype) + weights (size[1], dtype) + bias (size[1], dtype) +
    #       mean (size[0], float) + invstd (size[0], float) + outputs (size, dtype)
    return int(
        dtype.itemsize * 2 * (np.prod(size) + size[1])
        + torch.float.itemsize * 2 * size[0]
    )

@pytest.mark.parametrize("size", generate_input_sizes(dims=2))
@pytest.mark.parametrize("dtype", FLOAT_DTYPES)
@pytest.mark.inner_persistent
def test_layernorm_fwd_nvf_benchmark(
    benchmark,
    size: tuple,
    dtype: torch.dtype,
    disable_validation: bool,
    disable_benchmarking: bool,
    eps: float = 1e-5,
):
    batch_size, hidden_size = size
    inputs = [
        torch.randn(size, device="cuda", dtype=dtype),
        torch.randn(hidden_size, device="cuda", dtype=dtype),
        torch.randn(hidden_size, device="cuda", dtype=dtype),
    ]

    with FusionDefinition() as fd:
        layernorm_fwd_fusion(fd, torch_dtype_to_nvfuser_dtype(dtype))

    if not disable_validation:
        eager_output = layernorm(inputs)
        mean = inputs[0].to(torch.float).mean(dim=-1)
        variance = inputs[0].to(torch.float).var(dim=-1, unbiased=False)
        invstd = (1.0 / torch.sqrt(variance + eps)).unsqueeze(1)

        fd.validate(inputs, [eager_output, mean, invstd])

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

import os

@pytest.mark.parametrize("executor", DEFAULT_EXECUTORS)
@pytest.mark.parametrize("size", generate_input_sizes(dims=2))
@pytest.mark.parametrize("dtype", FLOAT_DTYPES)
@pytest.mark.inner_persistent
def test_layernorm_fwd_baseline_benchmark(
    benchmark,
    size: tuple,
    dtype: torch.dtype,
    executor: str,
):
    if executor == "torchcompile":
        clear_dynamo_cache()
        assert os.environ.get("TORCHINDUCTOR_COORDINATE_DESCENT_TUNING") == "1"
        assert (
            os.environ.get("TORCHINDUCTOR_COORDINATE_DESCENT_CHECK_ALL_DIRECTIONS")
            == "1"
        )
    batch_size, hidden_size = size
    inputs = [
        torch.randn(size, device="cuda", dtype=dtype, requires_grad=True),
        torch.randn(hidden_size, device="cuda", dtype=dtype, requires_grad=True),
        torch.randn(hidden_size, device="cuda", dtype=dtype, requires_grad=True),
    ]

    benchmark_fn = with_executor(executor, layernorm)

    # Manually compute IOBytes: See PR #1725
    run_benchmark(
        benchmark,
        benchmark_fn,
        inputs,
        iobytes=layernorm_fwd_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, with_executor, DEFAULT_EXECUTORS
import torch
from .global_params import generate_input_sizes, FLOAT_DTYPES, PROMOTE_DTYPES
import numpy as np
from .torch_ops import layernorm

def layernorm_fwd_fusion(
    fd: FusionDefinition,
    dtype: DataType,
    eps: float = 1e-5,
) -> None:
    T0 = fd.define_tensor(
        shape=[-1, -1], contiguity=[True, True], dtype=dtype, is_cpu=False
    )
    T1 = fd.define_tensor(shape=[-1], contiguity=[True], dtype=dtype, is_cpu=False)
    T2 = fd.define_tensor(shape=[-1], contiguity=[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.cast(T2, dtype=DataType.Float)

    T3, T4 = fd.ops.var_mean(T0, dims=[1], correction=0, keepdim=False)

    T7 = fd.ops.broadcast_in_dim(T3, shape=[T0.size(0), 1], broadcast_dims=[0])
    T11 = fd.ops.broadcast_in_dim(T4, shape=[T0.size(0), 1], broadcast_dims=[0])

    S12 = fd.define_scalar(eps, dtype=DataType.Double)
    T13 = fd.ops.add(T7, S12)
    T14 = fd.ops.rsqrt(T13)

    V17 = T0.shape()
    T18 = fd.ops.broadcast_in_dim(T11, shape=V17, broadcast_dims=[0, 1])
    T19 = fd.ops.sub(T0, T18)
    T23 = fd.ops.broadcast_in_dim(T14, shape=V17, broadcast_dims=[0, 1])
    T24 = fd.ops.mul(T19, T23)

    T25 = fd.ops.broadcast_in_dim(T1, shape=V17, broadcast_dims=[1])
    T26 = fd.ops.mul(T24, T25)
    T27 = fd.ops.broadcast_in_dim(T2, shape=V17, broadcast_dims=[1])
    T28 = fd.ops.add(T26, T27)

    if dtype in PROMOTE_DTYPES:
        T28 = fd.ops.cast(T28, dtype=dtype)

    fd.add_output(T28)
    fd.add_output(T4)
    fd.add_output(T14)

def layernorm_fwd_iobytes(size: tuple, dtype: torch.dtype):
    # Manual IOBytes computation required since nvFuser outputs (out, mean, invstd) differs from baselines (out)
    # Total IO bytes = in_tensor (size, dtype) + weights (size[1], dtype) + bias (size[1], dtype) +
    #       mean (size[0], float) + invstd (size[0], float) + outputs (size, dtype)
    return int(
        dtype.itemsize * 2 * (np.prod(size) + size[1])
        + torch.float.itemsize * 2 * size[0]
    )

@pytest.mark.parametrize("size", generate_input_sizes(dims=2))
@pytest.mark.parametrize("dtype", FLOAT_DTYPES)
@pytest.mark.inner_persistent
def test_layernorm_fwd_nvf_benchmark(
    benchmark,
    size: tuple,
    dtype: torch.dtype,
    disable_validation: bool,


    V17 = T0.shape()
    T18 = fd.ops.broadcast_in_dim(T11, shape=V17, broadcast_dims=[0, 1])
    T19 = fd.ops.sub(T0, T18)
    T23 = fd.ops.broadcast_in_dim(T14, shape=V17, broadcast_dims=[0, 1])
    T24 = fd.ops.mul(T19, T23)

    T25 = fd.ops.broadcast_in_dim(T1, shape=V17, broadcast_dims=[1])
    T26 = fd.ops.mul(T24, T25)
    T27 = fd.ops.broadcast_in_dim(T2, shape=V17, broadcast_dims=[1])
    T28 = fd.ops.add(T26, T27)

    if dtype in PROMOTE_DTYPES:
        T28 = fd.ops.cast(T28, dtype=dtype)

    fd.add_output(T28)
    fd.add_output(T4)
    fd.add_output(T14)

def layernorm_fwd_iobytes(size: tuple, dtype: torch.dtype):
    # Manual IOBytes computation required since nvFuser outputs (out, mean, invstd) differs from baselines (out)
    # Total IO bytes = in_tensor (size, dtype) + weights (size[1], dtype) + bias (size[1], dtype) +
    #       mean (size[0], float) + invstd (size[0], float) + outputs (size, dtype)
    return int(
        dtype.itemsize * 2 * (np.prod(size) + size[1])
        + torch.float.itemsize * 2 * size[0]
    )

@pytest.mark.parametrize("size", generate_input_sizes(dims=2))
@pytest.mark.parametrize("dtype", FLOAT_DTYPES)
@pytest.mark.inner_persistent
def test_layernorm_fwd_nvf_benchmark(
    benchmark,
    size: tuple,
    dtype: torch.dtype,
    disable_validation: bool,
    disable_benchmarking: bool,
    eps: float = 1e-5,
):
    batch_size, hidden_size = size
    inputs = [
        torch.randn(size, device="cuda", dtype=dtype),
        torch.randn(hidden_size, device="cuda", dtype=dtype),
        torch.randn(hidden_size, device="cuda", dtype=dtype),
    ]


        T28 = fd.ops.cast(T28, dtype=dtype)

    fd.add_output(T28)
    fd.add_output(T4)
    fd.add_output(T14)

def layernorm_fwd_iobytes(size: tuple, dtype: torch.dtype):
    # Manual IOBytes computation required since nvFuser outputs (out, mean, invstd) differs from baselines (out)
    # Total IO bytes = in_tensor (size, dtype) + weights (size[1], dtype) + bias (size[1], dtype) +
    #       mean (size[0], float) + invstd (size[0], float) + outputs (size, dtype)
    return int(
        dtype.itemsize * 2 * (np.prod(size) + size[1])
        + torch.float.itemsize * 2 * size[0]
    )

@pytest.mark.parametrize("size", generate_input_sizes(dims=2))
@pytest.mark.parametrize("dtype", FLOAT_DTYPES)
@pytest.mark.inner_persistent
def test_layernorm_fwd_nvf_benchmark(
    benchmark,
    size: tuple,
    dtype: torch.dtype,
    disable_validation: bool,
    disable_benchmarking: bool,
    eps: float = 1e-5,
):
    batch_size, hidden_size = size
    inputs = [
        torch.randn(size, device="cuda", dtype=dtype),
        torch.randn(hidden_size, device="cuda", dtype=dtype),
        torch.randn(hidden_size, device="cuda", dtype=dtype),
    ]

    with FusionDefinition() as fd:
        layernorm_fwd_fusion(fd, torch_dtype_to_nvfuser_dtype(dtype))

    if not disable_validation:
        eager_output = layernorm(inputs)
        mean = inputs[0].to(torch.float).mean(dim=-1)
        variance = inputs[0].to(torch.float).var(dim=-1, unbiased=False)
        invstd = (1.0 / torch.sqrt(variance + eps)).unsqueeze(1)

        fd.validate(inputs, [eager_output, mean, invstd])

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

import os

@pytest.mark.parametrize("executor", DEFAULT_EXECUTORS)
@pytest.mark.parametrize("size", generate_input_sizes(dims=2))
@pytest.mark.parametrize("dtype", FLOAT_DTYPES)
@pytest.mark.inner_persistent
def test_layernorm_fwd_baseline_benchmark(
    benchmark,
    size: tuple,
    dtype: torch.dtype,
    executor: str,
):
    if executor == "torchcompile":
        clear_dynamo_cache()
        assert os.environ.get("TORCHINDUCTOR_COORDINATE_DESCENT_TUNING") == "1"
        assert (
            os.environ.get("TORCHINDUCTOR_COORDINATE_DESCENT_CHECK_ALL_DIRECTIONS")


    if not disable_validation:
        eager_output = layernorm(inputs)
        mean = inputs[0].to(torch.float).mean(dim=-1)
        variance = inputs[0].to(torch.float).var(dim=-1, unbiased=False)
        invstd = (1.0 / torch.sqrt(variance + eps)).unsqueeze(1)

        fd.validate(inputs, [eager_output, mean, invstd])

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

import os

@pytest.mark.parametrize("executor", DEFAULT_EXECUTORS)
@pytest.mark.parametrize("size", generate_input_sizes(dims=2))
@pytest.mark.parametrize("dtype", FLOAT_DTYPES)
@pytest.mark.inner_persistent
def test_layernorm_fwd_baseline_benchmark(
    benchmark,
    size: tuple,
    dtype: torch.dtype,
    executor: str,
):
    if executor == "torchcompile":
        clear_dynamo_cache()
        assert os.environ.get("TORCHINDUCTOR_COORDINATE_DESCENT_TUNING") == "1"
        assert (
            os.environ.get("TORCHINDUCTOR_COORDINATE_DESCENT_CHECK_ALL_DIRECTIONS")
            == "1"
        )
    batch_size, hidden_size = size
    inputs = [
        torch.randn(size, device="cuda", dtype=dtype, requires_grad=True),
        torch.randn(hidden_size, device="cuda", dtype=dtype, requires_grad=True),
        torch.randn(hidden_size, device="cuda", dtype=dtype, requires_grad=True),
    ]

    benchmark_fn = with_executor(executor, layernorm)

    # Manually compute IOBytes: See PR #1725
    run_benchmark(
        benchmark,
        benchmark_fn,
        inputs,
        iobytes=layernorm_fwd_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,
    with_executor,
    DEFAULT_EXECUTORS,
)
import torch
from .global_params import generate_input_sizes, FLOAT_DTYPES, PROMOTE_DTYPES
import numpy as np
from .torch_ops import rmsnorm

def rmsnorm_bwd_fusion(
    fd: FusionDefinition,
    dtype: DataType,
):
    T4 = fd.define_tensor(
        shape=[-1, -1], contiguity=[True, True], dtype=dtype, is_cpu=False
    )
    T5 = fd.define_tensor(
        shape=[-1, 1], contiguity=[True, None], dtype=DataType.Float, is_cpu=False
    )
    T6 = fd.define_tensor(
        shape=[-1, -1], contiguity=[True, True], dtype=dtype, is_cpu=False
    )
    T7 = fd.define_tensor(shape=[-1], contiguity=[True], dtype=dtype, is_cpu=False)

    S0 = fd.define_scalar(2.0, dtype=DataType.Double)

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

    T14 = fd.ops.broadcast_in_dim(T5, shape=T4.shape(), broadcast_dims=[0, 1])
    T15 = fd.ops.reciprocal(T14)
    T16 = fd.ops.mul(T4, T15)
    T20 = fd.ops.broadcast_in_dim(T7, shape=T4.shape(), broadcast_dims=[1])

    T23 = fd.ops.mul(T6, T16)
    T24 = fd.ops.mul(T6, T20)
    T25 = fd.ops.sum(T23, dims=[0], keepdim=False, dtype=DataType.Null)

    T28 = fd.ops.mul(T24, T15)
    T29 = fd.ops.neg(T24)
    T30 = fd.ops.mul(T29, T4)
    T32 = fd.ops.pow(T14, S0)
    T33 = fd.ops.reciprocal(T32)
    T34 = fd.ops.mul(T30, T33)
    T35 = fd.ops.sum(T34, dims=[1], keepdim=False, dtype=DataType.Null)
    T41 = fd.ops.broadcast_in_dim(T35, shape=[T4.size(0), 1], broadcast_dims=[0])
    T43 = fd.ops.mul(S0, T5)
    T44 = fd.ops.reciprocal(T43)
    T45 = fd.ops.mul(T41, T44)
    S48 = fd.ops.reciprocal(T4.size(1))
    T49 = fd.ops.mul(T45, S48)
    T50 = fd.ops.sum(T49, dims=[1], keepdim=False, dtype=DataType.Null)
    T54 = fd.ops.broadcast_in_dim(T50, shape=[T4.size(0), 1], broadcast_dims=[0])
    T58 = fd.ops.broadcast_in_dim(T54, shape=T4.shape(), broadcast_dims=[0, 1])
    T59 = fd.ops.mul(T58, S0)
    T62 = fd.ops.mul(T59, T4)
    T63 = fd.ops.add(T28, T62)

    if dtype in PROMOTE_DTYPES:
        T63 = fd.ops.cast(T63, dtype=dtype)
        T25 = fd.ops.cast(T25, dtype=dtype)

    fd.add_output(T63)
    fd.add_output(T25)

def rmsnorm_bwd_iobytes(size: tuple, dtype: torch.dtype):
    # Manual IOBytes computation since nvfuser input/outputs (in_tensor, grad_out, rms, weigts, grad_in, grad_weights) differ from baselines (out, grad_out)
    # Total IO bytes = in_tensor (size, dtype) + grad_out (size, dtype) + rms_eps(size[0], float) + weights (size[1], dtype) +
    #       grad_in (size, dtype) + grad_weights (size[1], dtype)
    return int(
        dtype.itemsize * (3 * np.prod(size) + 2 * size[1])
        + torch.float.itemsize * size[0]
    )

@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_rmsnorm_bwd_nvf_benchmark(
    benchmark,
    size: tuple,
    dtype: torch.dtype,
    disable_validation: bool,
    disable_benchmarking: bool,
    eps: float = 1e-5,
):
    inputs = 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)

    squared_mean = (inputs.to(torch.float) ** 2).mean(1, keepdim=True)
    rms_eps = torch.sqrt(squared_mean + eps)

    with FusionDefinition() as fd:
        rmsnorm_bwd_fusion(fd, torch_dtype_to_nvfuser_dtype(dtype))

    if not disable_validation:
        eager_output = weights.to(torch.double) * (
            inputs.to(torch.double) / rms_eps.to(torch.double)
        )
        eager_output.backward(grads.to(torch.double))
        fd.validate([inputs, rms_eps, grads, weights], [inputs.grad, weights.grad])

    if not disable_benchmarking:
        run_benchmark(benchmark, fd.execute, [inputs, rms_eps, grads, weights])

@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_rmsnorm_bwd_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)
    grads = torch.randn(size, device="cuda", dtype=dtype)
    weights = torch.randn(size[1], device="cuda", dtype=dtype, requires_grad=True)

    # Compile the fwd fn for torchcompile
    fwd_fn = with_executor(executor, rmsnorm)
    fwd_inputs = [inputs, weights]
    outputs = fwd_fn(fwd_inputs)

    # Manually compute IOBytes: See PR #1725
    run_benchmark(
        benchmark,
        unary_bwd_torch,
        [outputs, grads, *fwd_inputs],
        iobytes=rmsnorm_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,
    with_executor,
    DEFAULT_EXECUTORS,
)
import torch
from .global_params import generate_input_sizes, FLOAT_DTYPES, PROMOTE_DTYPES
import numpy as np
from .torch_ops import rmsnorm

def rmsnorm_bwd_fusion(
    fd: FusionDefinition,
    dtype: DataType,
):
    T4 = fd.define_tensor(
        shape=[-1, -1], contiguity=[True, True], dtype=dtype, is_cpu=False
    )
    T5 = fd.define_tensor(
        shape=[-1, 1], contiguity=[True, None], dtype=DataType.Float, is_cpu=False
    )
    T6 = fd.define_tensor(
        shape=[-1, -1], contiguity=[True, True], dtype=dtype, is_cpu=False
    )
    T7 = fd.define_tensor(shape=[-1], contiguity=[True], dtype=dtype, is_cpu=False)

    S0 = fd.define_scalar(2.0, dtype=DataType.Double)

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

    T14 = fd.ops.broadcast_in_dim(T5, shape=T4.shape(), broadcast_dims=[0, 1])
    T15 = fd.ops.reciprocal(T14)
    T16 = fd.ops.mul(T4, T15)
    T20 = fd.ops.broadcast_in_dim(T7, shape=T4.shape(), broadcast_dims=[1])

    T23 = fd.ops.mul(T6, T16)
    T24 = fd.ops.mul(T6, T20)
    T25 = fd.ops.sum(T23, dims=[0], keepdim=False, dtype=DataType.Null)

    T28 = fd.ops.mul(T24, T15)
    T29 = fd.ops.neg(T24)
    T30 = fd.ops.mul(T29, T4)
    T32 = fd.ops.pow(T14, S0)
    T33 = fd.ops.reciprocal(T32)
    T34 = fd.ops.mul(T30, T33)
    T35 = fd.ops.sum(T34, dims=[1], keepdim=False, dtype=DataType.Null)
    T41 = fd.ops.broadcast_in_dim(T35, shape=[T4.size(0), 1], broadcast_dims=[0])
    T43 = fd.ops.mul(S0, T5)
    T44 = fd.ops.reciprocal(T43)
    T45 = fd.ops.mul(T41, T44)
    S48 = fd.ops.reciprocal(T4.size(1))
    T49 = fd.ops.mul(T45, S48)
    T50 = fd.ops.sum(T49, dims=[1], keepdim=False, dtype=DataType.Null)
    T54 = fd.ops.broadcast_in_dim(T50, shape=[T4.size(0), 1], broadcast_dims=[0])
    T58 = fd.ops.broadcast_in_dim(T54, shape=T4.shape(), broadcast_dims=[0, 1])
    T59 = fd.ops.mul(T58, S0)
    T62 = fd.ops.mul(T59, T4)
    T63 = fd.ops.add(T28, T62)

    if dtype in PROMOTE_DTYPES:
        T63 = fd.ops.cast(T63, dtype=dtype)
        T25 = fd.ops.cast(T25, dtype=dtype)

    fd.add_output(T63)
    fd.add_output(T25)

def rmsnorm_bwd_iobytes(size: tuple, dtype: torch.dtype):
    # Manual IOBytes computation since nvfuser input/outputs (in_tensor, grad_out, rms, weigts, grad_in, grad_weights) differ from baselines (out, grad_out)
    # Total IO bytes = in_tensor (size, dtype) + grad_out (size, dtype) + rms_eps(size[0], float) + weights (size[1], dtype) +
    #       grad_in (size, dtype) + grad_weights (size[1], dtype)
    return int(
        dtype.itemsize * (3 * np.prod(size) + 2 * size[1])
        + torch.float.itemsize * size[0]
    )

@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_rmsnorm_bwd_nvf_benchmark(
    benchmark,
    size: tuple,
    dtype: torch.dtype,


    T43 = fd.ops.mul(S0, T5)
    T44 = fd.ops.reciprocal(T43)
    T45 = fd.ops.mul(T41, T44)
    S48 = fd.ops.reciprocal(T4.size(1))
    T49 = fd.ops.mul(T45, S48)
    T50 = fd.ops.sum(T49, dims=[1], keepdim=False, dtype=DataType.Null)
    T54 = fd.ops.broadcast_in_dim(T50, shape=[T4.size(0), 1], broadcast_dims=[0])
    T58 = fd.ops.broadcast_in_dim(T54, shape=T4.shape(), broadcast_dims=[0, 1])
    T59 = fd.ops.mul(T58, S0)
    T62 = fd.ops.mul(T59, T4)
    T63 = fd.ops.add(T28, T62)

    if dtype in PROMOTE_DTYPES:
        T63 = fd.ops.cast(T63, dtype=dtype)
        T25 = fd.ops.cast(T25, dtype=dtype)

    fd.add_output(T63)
    fd.add_output(T25)

def rmsnorm_bwd_iobytes(size: tuple, dtype: torch.dtype):
    # Manual IOBytes computation since nvfuser input/outputs (in_tensor, grad_out, rms, weigts, grad_in, grad_weights) differ from baselines (out, grad_out)
    # Total IO bytes = in_tensor (size, dtype) + grad_out (size, dtype) + rms_eps(size[0], float) + weights (size[1], dtype) +
    #       grad_in (size, dtype) + grad_weights (size[1], dtype)
    return int(
        dtype.itemsize * (3 * np.prod(size) + 2 * size[1])
        + torch.float.itemsize * size[0]
    )

@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_rmsnorm_bwd_nvf_benchmark(
    benchmark,
    size: tuple,
    dtype: torch.dtype,
    disable_validation: bool,
    disable_benchmarking: bool,
    eps: float = 1e-5,
):
    inputs = 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)

    squared_mean = (inputs.to(torch.float) ** 2).mean(1, keepdim=True)
    rms_eps = torch.sqrt(squared_mean + eps)


        T25 = fd.ops.cast(T25, dtype=dtype)

    fd.add_output(T63)
    fd.add_output(T25)

def rmsnorm_bwd_iobytes(size: tuple, dtype: torch.dtype):
    # Manual IOBytes computation since nvfuser input/outputs (in_tensor, grad_out, rms, weigts, grad_in, grad_weights) differ from baselines (out, grad_out)
    # Total IO bytes = in_tensor (size, dtype) + grad_out (size, dtype) + rms_eps(size[0], float) + weights (size[1], dtype) +
    #       grad_in (size, dtype) + grad_weights (size[1], dtype)
    return int(
        dtype.itemsize * (3 * np.prod(size) + 2 * size[1])
        + torch.float.itemsize * size[0]
    )

@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_rmsnorm_bwd_nvf_benchmark(
    benchmark,
    size: tuple,
    dtype: torch.dtype,
    disable_validation: bool,
    disable_benchmarking: bool,
    eps: float = 1e-5,
):
    inputs = 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)

    squared_mean = (inputs.to(torch.float) ** 2).mean(1, keepdim=True)
    rms_eps = torch.sqrt(squared_mean + eps)

    with FusionDefinition() as fd:
        rmsnorm_bwd_fusion(fd, torch_dtype_to_nvfuser_dtype(dtype))

    if not disable_validation:
        eager_output = weights.to(torch.double) * (
            inputs.to(torch.double) / rms_eps.to(torch.double)
        )
        eager_output.backward(grads.to(torch.double))
        fd.validate([inputs, rms_eps, grads, weights], [inputs.grad, weights.grad])

    if not disable_benchmarking:
        run_benchmark(benchmark, fd.execute, [inputs, rms_eps, grads, weights])

@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_rmsnorm_bwd_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)
    grads = torch.randn(size, device="cuda", dtype=dtype)
    weights = torch.randn(size[1], device="cuda", dtype=dtype, requires_grad=True)

    # Compile the fwd fn for torchcompile


    with FusionDefinition() as fd:
        rmsnorm_bwd_fusion(fd, torch_dtype_to_nvfuser_dtype(dtype))

    if not disable_validation:
        eager_output = weights.to(torch.double) * (
            inputs.to(torch.double) / rms_eps.to(torch.double)
        )
        eager_output.backward(grads.to(torch.double))
        fd.validate([inputs, rms_eps, grads, weights], [inputs.grad, weights.grad])

    if not disable_benchmarking:
        run_benchmark(benchmark, fd.execute, [inputs, rms_eps, grads, weights])

@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_rmsnorm_bwd_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)
    grads = torch.randn(size, device="cuda", dtype=dtype)
    weights = torch.randn(size[1], device="cuda", dtype=dtype, requires_grad=True)

    # Compile the fwd fn for torchcompile
    fwd_fn = with_executor(executor, rmsnorm)
    fwd_inputs = [inputs, weights]
    outputs = fwd_fn(fwd_inputs)

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


# 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_fwd_nvf_benchmark, norm_fwd_baseline_benchmark

@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_instancenorm_fwd_nvf_benchmark(
    benchmark,
    size: tuple,
    dtype: torch.dtype,
    channels_last: bool,
    disable_validation: bool,
    disable_benchmarking: bool,
):
    norm_fwd_nvf_benchmark(
        benchmark,
        size,
        dtype,
        "instance_norm",
        channels_last,
        disable_validation,
        disable_benchmarking,
    )

@pytest.mark.parametrize("executor", ["eager", "torchcompile"])
@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_instancenorm_fwd_baseline_benchmark(
    benchmark, size: tuple, dtype: torch.dtype, channels_last: bool, executor: str
):
    norm_fwd_baseline_benchmark(
        benchmark, size, dtype, channels_last, executor, "instance_norm"
    )


        benchmark,
        size,
        dtype,
        "instance_norm",
        channels_last,
        disable_validation,
        disable_benchmarking,
    )

@pytest.mark.parametrize("executor", ["eager", "torchcompile"])
@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_instancenorm_fwd_baseline_benchmark(
    benchmark, size: tuple, dtype: torch.dtype, channels_last: bool, executor: str
):
    norm_fwd_baseline_benchmark(
        benchmark, size, dtype, channels_last, executor, "instance_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,
    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_fwd_fusion(
    fd: FusionDefinition, dtype: DataType, dropout_p: float, eps: float = 1e-5
) -> None:
    """
    Forward pass fusion definition for computing:
        output = layernorm (input2 + dropout (input1, p=dropout_p))

    Fusion inputs: input1, input2, weights, bias
    Fusion outputs: output, mean, invstd, dropout_mask
    """
    T2 = fd.define_tensor(
        shape=[-1, -1], contiguity=[True, True], dtype=dtype, is_cpu=False
    )
    T3 = fd.define_tensor(
        shape=[-1, -1], contiguity=[True, True], dtype=dtype, is_cpu=False
    )
    T1 = fd.define_tensor(shape=[-1], contiguity=[True], dtype=dtype, is_cpu=False)
    T0 = fd.define_tensor(shape=[-1], contiguity=[True], dtype=dtype, is_cpu=False)
    S3 = fd.define_scalar(0.00000, dtype=DataType.Double)
    S4 = fd.define_scalar(1.00000, dtype=DataType.Double)
    T8 = fd.ops.uniform(S3, S4, shape=T2.shape(), dtype=DataType.Float)
    S9 = fd.define_scalar(1 - dropout_p, dtype=DataType.Double)
    T10 = fd.ops.lt(T8, S9)
    T11 = fd.ops.cast(T10, dtype=DataType.Float)
    if dtype in PROMOTE_DTYPES:
        T0 = fd.ops.cast(T0, dtype=DataType.Float)
        T1 = fd.ops.cast(T1, dtype=DataType.Float)
        T2 = fd.ops.cast(T2, dtype=DataType.Float)
        T3 = fd.ops.cast(T3, dtype=DataType.Float)

    # Dropout + Add
    T13 = fd.ops.mul(T2, T11)
    S14 = fd.define_scalar(1 / (1 - dropout_p), dtype=DataType.Double)
    T15 = fd.ops.mul(T13, S14)
    T16 = fd.ops.add(T3, T15)
    # Layernorm
    T17, T18 = fd.ops.var_mean(T16, dims=[1], correction=0, keepdim=False)
    T22 = fd.ops.broadcast_in_dim(T17, shape=[T2.size(0), 1], broadcast_dims=[0])
    T26 = fd.ops.broadcast_in_dim(T18, shape=[T2.size(0), 1], broadcast_dims=[0])
    S27 = fd.define_scalar(eps, dtype=DataType.Double)
    T28 = fd.ops.add(T22, S27)
    T29 = fd.ops.rsqrt(T28)
    T33 = fd.ops.broadcast_in_dim(T26, shape=T2.shape(), broadcast_dims=[0, 1])
    T34 = fd.ops.sub(T16, T33)
    T38 = fd.ops.broadcast_in_dim(T29, shape=T2.shape(), broadcast_dims=[0, 1])
    T39 = fd.ops.mul(T34, T38)
    T43 = fd.ops.broadcast_in_dim(T1, shape=T2.shape(), broadcast_dims=[1])
    T45 = fd.ops.mul(T39, T43)
    T49 = fd.ops.broadcast_in_dim(T0, shape=T2.shape(), broadcast_dims=[1])
    T51 = fd.ops.add(T45, T49)
    if dtype in PROMOTE_DTYPES:
        T51 = fd.ops.cast(T51, dtype=dtype)

    fd.add_output(T51)
    fd.add_output(T18)
    fd.add_output(T29)
    fd.add_output(T10)

def dropout_layernorm_fwd_iobytes(size: tuple, dtype: torch.dtype):
    # Manual IOByte computation is required since nvFuser outputs differ from baseline outputs (output).
    nvf_inp_out = {
        # Inputs
        "input1": (size, dtype),
        "input2": (size, dtype),
        "weights": (size[1], dtype),
        "bias": (size[1], dtype),
        # Outputs
        "mean": (size[0], torch.float),
        "invstd": (size[0], torch.float),
        "outputs": (size, dtype),
        "dropout_mask": (size, torch.bool),
    }
    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_persistent
def test_dropout_layernorm_fwd_nvf_benchmark(
    benchmark,
    size: tuple,
    dtype: torch.dtype,
    disable_validation: bool,
    disable_benchmarking: bool,
    eps: float = 1e-5,
):
    inputs = [
        torch.randn(size, device="cuda", dtype=dtype),
        torch.randn(size, device="cuda", dtype=dtype),
        torch.ones(size[1], device="cuda", dtype=dtype),
        torch.zeros(size[1], device="cuda", dtype=dtype),
    ]

    dropout_p = 0.2
    with FusionDefinition() as fd:
        dropout_layernorm_fwd_fusion(fd, torch_dtype_to_nvfuser_dtype(dtype), dropout_p)

    if not disable_validation:
        # For validating use a fusion definition with dropout_p=0.0
        with FusionDefinition() as val_fd:
            dropout_layernorm_fwd_fusion(
                val_fd, torch_dtype_to_nvfuser_dtype(dtype), dropout_p=0.0, eps=eps
            )

        dropout_mask = torch.ones(size, dtype=torch.bool, device="cuda")

        # dropout + add
        x = inputs[1] + torch.nn.functional.dropout(inputs[0], p=0.0)
        # layernorm
        eager_output = torch.nn.functional.layer_norm(
            x.to(torch.float),
            inputs[0].shape[1:],
            weight=inputs[2].to(torch.float),
            bias=inputs[3].to(torch.float),
        )

        # mean and invstd are computed for the output of dropout + add
        mean = x.to(torch.double).mean(dim=-1)
        variance = x.to(torch.double).var(dim=-1, unbiased=False)
        invstd = (1.0 / torch.sqrt(variance + eps)).unsqueeze(1)

        val_fd.validate(
            inputs,
            [
                eager_output.to(dtype),
                mean.to(torch.float),
                invstd.to(torch.float),
                dropout_mask,
            ],
        )

    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.inner_persistent
def test_dropout_layernorm_fwd_baseline_benchmark(
    benchmark,
    size: tuple,
    dtype: torch.dtype,
    executor: str,
):
    if executor == "torchcompile":
        clear_dynamo_cache()

    dropout_p = 0.2
    inputs = [
        torch.randn(size, device="cuda", dtype=dtype, requires_grad=True),
        torch.randn(size, device="cuda", dtype=dtype, requires_grad=True),
        torch.ones(size[1], device="cuda", dtype=dtype, requires_grad=True),
        torch.zeros(size[1], device="cuda", dtype=dtype, requires_grad=True),
        dropout_p,
    ]

    benchmark_fn = with_executor(executor, dropout_layernorm)

    # Manually compute IOBytes: See PR #1725
    run_benchmark(
        benchmark,
        benchmark_fn,
        inputs,
        iobytes=dropout_layernorm_fwd_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,
    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_fwd_fusion(
    fd: FusionDefinition, dtype: DataType, dropout_p: float, eps: float = 1e-5
) -> None:
    """
    Forward pass fusion definition for computing:
        output = layernorm (input2 + dropout (input1, p=dropout_p))

    Fusion inputs: input1, input2, weights, bias
    Fusion outputs: output, mean, invstd, dropout_mask
    """
    T2 = fd.define_tensor(
        shape=[-1, -1], contiguity=[True, True], dtype=dtype, is_cpu=False
    )
    T3 = fd.define_tensor(
        shape=[-1, -1], contiguity=[True, True], dtype=dtype, is_cpu=False
    )
    T1 = fd.define_tensor(shape=[-1], contiguity=[True], dtype=dtype, is_cpu=False)
    T0 = fd.define_tensor(shape=[-1], contiguity=[True], dtype=dtype, is_cpu=False)
    S3 = fd.define_scalar(0.00000, dtype=DataType.Double)
    S4 = fd.define_scalar(1.00000, dtype=DataType.Double)
    T8 = fd.ops.uniform(S3, S4, shape=T2.shape(), dtype=DataType.Float)
    S9 = fd.define_scalar(1 - dropout_p, dtype=DataType.Double)
    T10 = fd.ops.lt(T8, S9)
    T11 = fd.ops.cast(T10, dtype=DataType.Float)
    if dtype in PROMOTE_DTYPES:
        T0 = fd.ops.cast(T0, dtype=DataType.Float)
        T1 = fd.ops.cast(T1, dtype=DataType.Float)
        T2 = fd.ops.cast(T2, dtype=DataType.Float)
        T3 = fd.ops.cast(T3, dtype=DataType.Float)

    # Dropout + Add
    T13 = fd.ops.mul(T2, T11)
    S14 = fd.define_scalar(1 / (1 - dropout_p), dtype=DataType.Double)
    T15 = fd.ops.mul(T13, S14)
    T16 = fd.ops.add(T3, T15)
    # Layernorm
    T17, T18 = fd.ops.var_mean(T16, dims=[1], correction=0, keepdim=False)
    T22 = fd.ops.broadcast_in_dim(T17, shape=[T2.size(0), 1], broadcast_dims=[0])
    T26 = fd.ops.broadcast_in_dim(T18, shape=[T2.size(0), 1], broadcast_dims=[0])
    S27 = fd.define_scalar(eps, dtype=DataType.Double)
    T28 = fd.ops.add(T22, S27)
    T29 = fd.ops.rsqrt(T28)
    T33 = fd.ops.broadcast_in_dim(T26, shape=T2.shape(), broadcast_dims=[0, 1])
    T34 = fd.ops.sub(T16, T33)
    T38 = fd.ops.broadcast_in_dim(T29, shape=T2.shape(), broadcast_dims=[0, 1])
    T39 = fd.ops.mul(T34, T38)
    T43 = fd.ops.broadcast_in_dim(T1, shape=T2.shape(), broadcast_dims=[1])
    T45 = fd.ops.mul(T39, T43)
    T49 = fd.ops.broadcast_in_dim(T0, shape=T2.shape(), broadcast_dims=[1])
    T51 = fd.ops.add(T45, T49)
    if dtype in PROMOTE_DTYPES:
        T51 = fd.ops.cast(T51, dtype=dtype)

    fd.add_output(T51)
    fd.add_output(T18)
    fd.add_output(T29)
    fd.add_output(T10)

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

@pytest.mark.parametrize("size", generate_input_sizes(dims=2))


    S27 = fd.define_scalar(eps, dtype=DataType.Double)
    T28 = fd.ops.add(T22, S27)
    T29 = fd.ops.rsqrt(T28)
    T33 = fd.ops.broadcast_in_dim(T26, shape=T2.shape(), broadcast_dims=[0, 1])
    T34 = fd.ops.sub(T16, T33)
    T38 = fd.ops.broadcast_in_dim(T29, shape=T2.shape(), broadcast_dims=[0, 1])
    T39 = fd.ops.mul(T34, T38)
    T43 = fd.ops.broadcast_in_dim(T1, shape=T2.shape(), broadcast_dims=[1])
    T45 = fd.ops.mul(T39, T43)
    T49 = fd.ops.broadcast_in_dim(T0, shape=T2.shape(), broadcast_dims=[1])
    T51 = fd.ops.add(T45, T49)
    if dtype in PROMOTE_DTYPES:
        T51 = fd.ops.cast(T51, dtype=dtype)

    fd.add_output(T51)
    fd.add_output(T18)
    fd.add_output(T29)
    fd.add_output(T10)

def dropout_layernorm_fwd_iobytes(size: tuple, dtype: torch.dtype):
    # Manual IOByte computation is required since nvFuser outputs differ from baseline outputs (output).
    nvf_inp_out = {
        # Inputs
        "input1": (size, dtype),
        "input2": (size, dtype),
        "weights": (size[1], dtype),
        "bias": (size[1], dtype),
        # Outputs
        "mean": (size[0], torch.float),
        "invstd": (size[0], torch.float),
        "outputs": (size, dtype),
        "dropout_mask": (size, torch.bool),
    }
    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_persistent
def test_dropout_layernorm_fwd_nvf_benchmark(
    benchmark,
    size: tuple,
    dtype: torch.dtype,
    disable_validation: bool,
    disable_benchmarking: bool,
    eps: float = 1e-5,
):
    inputs = [
        torch.randn(size, device="cuda", dtype=dtype),
        torch.randn(size, device="cuda", dtype=dtype),
        torch.ones(size[1], device="cuda", dtype=dtype),
        torch.zeros(size[1], device="cuda", dtype=dtype),
    ]


def dropout_layernorm_fwd_iobytes(size: tuple, dtype: torch.dtype):
    # Manual IOByte computation is required since nvFuser outputs differ from baseline outputs (output).
    nvf_inp_out = {
        # Inputs
        "input1": (size, dtype),
        "input2": (size, dtype),
        "weights": (size[1], dtype),
        "bias": (size[1], dtype),
        # Outputs
        "mean": (size[0], torch.float),
        "invstd": (size[0], torch.float),
        "outputs": (size, dtype),
        "dropout_mask": (size, torch.bool),
    }
    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_persistent
def test_dropout_layernorm_fwd_nvf_benchmark(
    benchmark,
    size: tuple,
    dtype: torch.dtype,
    disable_validation: bool,
    disable_benchmarking: bool,
    eps: float = 1e-5,
):
    inputs = [
        torch.randn(size, device="cuda", dtype=dtype),
        torch.randn(size, device="cuda", dtype=dtype),
        torch.ones(size[1], device="cuda", dtype=dtype),
        torch.zeros(size[1], device="cuda", dtype=dtype),
    ]

    dropout_p = 0.2
    with FusionDefinition() as fd:
        dropout_layernorm_fwd_fusion(fd, torch_dtype_to_nvfuser_dtype(dtype), dropout_p)

    if not disable_validation:
        # For validating use a fusion definition with dropout_p=0.0
        with FusionDefinition() as val_fd:
            dropout_layernorm_fwd_fusion(
                val_fd, torch_dtype_to_nvfuser_dtype(dtype), dropout_p=0.0, eps=eps
            )

        dropout_mask = torch.ones(size, dtype=torch.bool, device="cuda")

        # dropout + add
        x = inputs[1] + torch.nn.functional.dropout(inputs[0], p=0.0)
        # layernorm
        eager_output = torch.nn.functional.layer_norm(
            x.to(torch.float),
            inputs[0].shape[1:],
            weight=inputs[2].to(torch.float),
            bias=inputs[3].to(torch.float),
        )

        # mean and invstd are computed for the output of dropout + add
        mean = x.to(torch.double).mean(dim=-1)
        variance = x.to(torch.double).var(dim=-1, unbiased=False)
        invstd = (1.0 / torch.sqrt(variance + eps)).unsqueeze(1)

        val_fd.validate(
            inputs,
            [
                eager_output.to(dtype),
                mean.to(torch.float),
                invstd.to(torch.float),
                dropout_mask,
            ],
        )

    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.inner_persistent
def test_dropout_layernorm_fwd_baseline_benchmark(
    benchmark,
    size: tuple,
    dtype: torch.dtype,
    executor: str,
):
    if executor == "torchcompile":
        clear_dynamo_cache()

    dropout_p = 0.2
    inputs = [
        torch.randn(size, device="cuda", dtype=dtype, requires_grad=True),
        torch.randn(size, device="cuda", dtype=dtype, requires_grad=True),
        torch.ones(size[1], device="cuda", dtype=dtype, requires_grad=True),


        invstd = (1.0 / torch.sqrt(variance + eps)).unsqueeze(1)

        val_fd.validate(
            inputs,
            [
                eager_output.to(dtype),
                mean.to(torch.float),
                invstd.to(torch.float),
                dropout_mask,
            ],
        )

    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.inner_persistent
def test_dropout_layernorm_fwd_baseline_benchmark(
    benchmark,
    size: tuple,
    dtype: torch.dtype,
    executor: str,
):
    if executor == "torchcompile":
        clear_dynamo_cache()

    dropout_p = 0.2
    inputs = [
        torch.randn(size, device="cuda", dtype=dtype, requires_grad=True),
        torch.randn(size, device="cuda", dtype=dtype, requires_grad=True),
        torch.ones(size[1], device="cuda", dtype=dtype, requires_grad=True),
        torch.zeros(size[1], device="cuda", dtype=dtype, requires_grad=True),
        dropout_p,
    ]

    benchmark_fn = with_executor(executor, dropout_layernorm)

    # Manually compute IOBytes: See PR #1725
    run_benchmark(
        benchmark,
        benchmark_fn,
        inputs,
        iobytes=dropout_layernorm_fwd_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,
    with_executor,
    DEFAULT_EXECUTORS,
)
import torch
from .global_params import generate_attn_inputs, FLOAT_DTYPES, PROMOTE_DTYPES
from .torch_ops import nanogpt_attn

# Fusion from nanogpt attention module
# The nvFuser defintion only includes the non-matmul computation (masked_fill + softmax + dropout)
def nanogpt_attn_bwd_fusion(
    fd: FusionDefinition, dtype: DataType, head_size: int, dropout_p: float
):
    S0 = fd.define_scalar(1 / (1 - dropout_p), dtype=DataType.Double)
    S1 = fd.define_scalar(1 / head_size**0.5, dtype=DataType.Double)

    T2 = fd.define_tensor(
        shape=[-1, -1, -1, -1],
        contiguity=[True, True, True, True],
        dtype=dtype,
        is_cpu=False,
    )
    T3 = fd.define_tensor(
        shape=[-1, -1, -1, -1],
        contiguity=[True, True, True, True],
        dtype=dtype,
        is_cpu=False,
    )
    T4 = fd.define_tensor(
        shape=[-1, -1, -1, -1],
        contiguity=[True, True, True, True],
        dtype=DataType.Bool,
        is_cpu=False,
    )
    T5 = fd.define_tensor(
        shape=[-1, -1, -1, -1],
        contiguity=[None, None, True, True],
        dtype=DataType.Bool,
        is_cpu=False,
    )

    if dtype in PROMOTE_DTYPES:
        T2 = fd.ops.cast(T2, dtype=DataType.Float)
        T3 = fd.ops.cast(T3, dtype=DataType.Float)
        T4 = fd.ops.cast(T4, dtype=DataType.Float)

    T7 = fd.ops.mul(T2, S0)
    T8 = fd.ops.mul(T7, T4)
    T9 = fd.ops.mul(T3, T8)
    T10 = fd.ops.sum(T9, dims=[3], keepdim=False, dtype=DataType.Null)

    T16 = fd.ops.broadcast_in_dim(
        T10, shape=[T2.size(0), T2.size(1), T2.size(2), 1], broadcast_dims=[0, 1, 2]
    )

    V21 = T2.shape()
    T22 = fd.ops.broadcast_in_dim(T16, shape=V21, broadcast_dims=[0, 1, 2, 3])

    T23 = fd.ops.sub(T8, T22)
    T24 = fd.ops.mul(T3, T23)
    S25 = fd.define_scalar(0.00000, dtype=DataType.Double)
    T26 = fd.ops.where(T5, S25, T24)
    T27 = fd.ops.mul(T26, S1)

    if dtype in PROMOTE_DTYPES:
        T27 = fd.ops.cast(T27, dtype=dtype)

    fd.add_output(T27)

def nanogpt_attn_bwd_iobytes(size: tuple, dtype: torch.dtype):
    # Manual IOByte computation is required since nvFuser input/outputs (grad_out, attn, dropout_mask, bias_mask, grad_input]) differ from baseline input/outputs (output, grad_output).

    # Total IO bytes = grad_out ([bs, nh, seq_len, seq_len], dtype) + attn ([bs, nh, seq_len, seq_len], dtype) +
    #       dropout_mask ([bs, nh, seq_len, seq_len], bool) + bias_mask ([bs, nh, seq_len, seq_len], bool) + grad_in ([bs, nh, seq_len, seq_len], dtype)
    bs, seq_len, nh, n_embd = size

    return int(
        bs * nh * seq_len * seq_len * (3 * dtype.itemsize + 2 * torch.bool.itemsize)
    )

@pytest.mark.parametrize("size", generate_attn_inputs())
@pytest.mark.parametrize("dtype", FLOAT_DTYPES)
@pytest.mark.inner_persistent
def test_nanogpt_attn_bwd_nvf_benchmark(
    benchmark,
    size: tuple,
    dtype: torch.dtype,
    disable_validation: bool,
    disable_benchmarking: bool,
):
    batch_size, seq_len, nh, n_embd = size
    hs = n_embd // nh
    dropout_p = 0.2
    inputs = torch.randn(
        batch_size, nh, seq_len, seq_len, device="cuda", dtype=dtype, requires_grad=True
    )
    bias = torch.tril(torch.ones(seq_len, seq_len, device="cuda")).view(
        1, 1, seq_len, seq_len
    )
    dropout_mask = torch.lt(
        torch.rand(batch_size, nh, seq_len, seq_len, device="cuda"), 1 - dropout_p
    )
    bias_mask = bias[:, :, :seq_len, :seq_len] == 0
    bias_mask = bias_mask.broadcast_to(batch_size, nh, seq_len, seq_len)
    grads = torch.randn(batch_size, nh, seq_len, seq_len, device="cuda", dtype=dtype)

    attn = inputs / (hs**0.5)
    attn = attn.masked_fill(bias[:, :, :seq_len, :seq_len] == 0, float("-inf"))
    attn = torch.nn.functional.softmax(attn, dim=-1)

    with FusionDefinition() as fd:
        nanogpt_attn_bwd_fusion(fd, torch_dtype_to_nvfuser_dtype(dtype), hs, dropout_p)

    if not disable_validation:
        # Use dropout_mask instead of torch.nn.functional.dropout for validating results.
        out = attn * dropout_mask * 1 / (1 - dropout_p)
        out.backward(grads)
        fd.validate([grads, attn, dropout_mask, bias_mask], [inputs.grad])

    if not disable_benchmarking:
        run_benchmark(benchmark, fd.execute, [grads, attn, dropout_mask, bias_mask])

@pytest.mark.parametrize("executor", DEFAULT_EXECUTORS)
@pytest.mark.parametrize("size", generate_attn_inputs())
@pytest.mark.parametrize("dtype", FLOAT_DTYPES)
@pytest.mark.inner_persistent
def test_nanogpt_attn_bwd_baseline_benchmark(
    benchmark,
    size: tuple,
    dtype: torch.dtype,
    executor: str,
):
    if executor == "torchcompile":
        clear_dynamo_cache()
    batch_size, seq_len, nh, n_embd = size
    dropout_p = 0.2
    inputs = torch.randn(
        batch_size, nh, seq_len, seq_len, device="cuda", dtype=dtype, requires_grad=True
    )
    bias = torch.tril(torch.ones(seq_len, seq_len, device="cuda")).view(
        1, 1, seq_len, seq_len
    )

    # Compile the fwd fn for torchcompile
    fwd_fn = with_executor(executor, nanogpt_attn)
    fwd_inputs = [inputs, bias, size, dropout_p]
    outputs = fwd_fn(fwd_inputs)

    grads = torch.randn(batch_size, nh, seq_len, seq_len, device="cuda", dtype=dtype)

    # Manually compute IOBytes: See PR #1725
    run_benchmark(
        benchmark,
        unary_bwd_torch,
        [outputs, grads, *fwd_inputs],
        iobytes=nanogpt_attn_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,
    with_executor,
    DEFAULT_EXECUTORS,
)
import torch
from .global_params import generate_attn_inputs, FLOAT_DTYPES, PROMOTE_DTYPES
from .torch_ops import nanogpt_attn

# Fusion from nanogpt attention module
# The nvFuser defintion only includes the non-matmul computation (masked_fill + softmax + dropout)
def nanogpt_attn_bwd_fusion(
    fd: FusionDefinition, dtype: DataType, head_size: int, dropout_p: float
):
    S0 = fd.define_scalar(1 / (1 - dropout_p), dtype=DataType.Double)
    S1 = fd.define_scalar(1 / head_size**0.5, dtype=DataType.Double)

    T2 = fd.define_tensor(
        shape=[-1, -1, -1, -1],
        contiguity=[True, True, True, True],
        dtype=dtype,
        is_cpu=False,
    )
    T3 = fd.define_tensor(
        shape=[-1, -1, -1, -1],
        contiguity=[True, True, True, True],
        dtype=dtype,
        is_cpu=False,
    )
    T4 = fd.define_tensor(
        shape=[-1, -1, -1, -1],
        contiguity=[True, True, True, True],
        dtype=DataType.Bool,
        is_cpu=False,
    )
    T5 = fd.define_tensor(
        shape=[-1, -1, -1, -1],
        contiguity=[None, None, True, True],
        dtype=DataType.Bool,
        is_cpu=False,
    )

    if dtype in PROMOTE_DTYPES:
        T2 = fd.ops.cast(T2, dtype=DataType.Float)
        T3 = fd.ops.cast(T3, dtype=DataType.Float)
        T4 = fd.ops.cast(T4, dtype=DataType.Float)

    T7 = fd.ops.mul(T2, S0)
    T8 = fd.ops.mul(T7, T4)
    T9 = fd.ops.mul(T3, T8)
    T10 = fd.ops.sum(T9, dims=[3], keepdim=False, dtype=DataType.Null)

    T16 = fd.ops.broadcast_in_dim(
        T10, shape=[T2.size(0), T2.size(1), T2.size(2), 1], broadcast_dims=[0, 1, 2]
    )

    V21 = T2.shape()
    T22 = fd.ops.broadcast_in_dim(T16, shape=V21, broadcast_dims=[0, 1, 2, 3])

    T23 = fd.ops.sub(T8, T22)
    T24 = fd.ops.mul(T3, T23)
    S25 = fd.define_scalar(0.00000, dtype=DataType.Double)
    T26 = fd.ops.where(T5, S25, T24)
    T27 = fd.ops.mul(T26, S1)

    if dtype in PROMOTE_DTYPES:
        T27 = fd.ops.cast(T27, dtype=dtype)

    fd.add_output(T27)

def nanogpt_attn_bwd_iobytes(size: tuple, dtype: torch.dtype):
    # Manual IOByte computation is required since nvFuser input/outputs (grad_out, attn, dropout_mask, bias_mask, grad_input]) differ from baseline input/outputs (output, grad_output).

    # Total IO bytes = grad_out ([bs, nh, seq_len, seq_len], dtype) + attn ([bs, nh, seq_len, seq_len], dtype) +
    #       dropout_mask ([bs, nh, seq_len, seq_len], bool) + bias_mask ([bs, nh, seq_len, seq_len], bool) + grad_in ([bs, nh, seq_len, seq_len], dtype)
    bs, seq_len, nh, n_embd = size

    return int(
        bs * nh * seq_len * seq_len * (3 * dtype.itemsize + 2 * torch.bool.itemsize)
    )

@pytest.mark.parametrize("size", generate_attn_inputs())
@pytest.mark.parametrize("dtype", FLOAT_DTYPES)
@pytest.mark.inner_persistent
def test_nanogpt_attn_bwd_nvf_benchmark(
    benchmark,
    size: tuple,


    T16 = fd.ops.broadcast_in_dim(
        T10, shape=[T2.size(0), T2.size(1), T2.size(2), 1], broadcast_dims=[0, 1, 2]
    )

    V21 = T2.shape()
    T22 = fd.ops.broadcast_in_dim(T16, shape=V21, broadcast_dims=[0, 1, 2, 3])

    T23 = fd.ops.sub(T8, T22)
    T24 = fd.ops.mul(T3, T23)
    S25 = fd.define_scalar(0.00000, dtype=DataType.Double)
    T26 = fd.ops.where(T5, S25, T24)
    T27 = fd.ops.mul(T26, S1)

    if dtype in PROMOTE_DTYPES:
        T27 = fd.ops.cast(T27, dtype=dtype)

    fd.add_output(T27)

def nanogpt_attn_bwd_iobytes(size: tuple, dtype: torch.dtype):
    # Manual IOByte computation is required since nvFuser input/outputs (grad_out, attn, dropout_mask, bias_mask, grad_input]) differ from baseline input/outputs (output, grad_output).

    # Total IO bytes = grad_out ([bs, nh, seq_len, seq_len], dtype) + attn ([bs, nh, seq_len, seq_len], dtype) +
    #       dropout_mask ([bs, nh, seq_len, seq_len], bool) + bias_mask ([bs, nh, seq_len, seq_len], bool) + grad_in ([bs, nh, seq_len, seq_len], dtype)
    bs, seq_len, nh, n_embd = size

    return int(
        bs * nh * seq_len * seq_len * (3 * dtype.itemsize + 2 * torch.bool.itemsize)
    )

@pytest.mark.parametrize("size", generate_attn_inputs())
@pytest.mark.parametrize("dtype", FLOAT_DTYPES)
@pytest.mark.inner_persistent
def test_nanogpt_attn_bwd_nvf_benchmark(
    benchmark,
    size: tuple,
    dtype: torch.dtype,
    disable_validation: bool,
    disable_benchmarking: bool,
):
    batch_size, seq_len, nh, n_embd = size
    hs = n_embd // nh
    dropout_p = 0.2
    inputs = torch.randn(
        batch_size, nh, seq_len, seq_len, device="cuda", dtype=dtype, requires_grad=True
    )
    bias = torch.tril(torch.ones(seq_len, seq_len, device="cuda")).view(
        1, 1, seq_len, seq_len


        T27 = fd.ops.cast(T27, dtype=dtype)

    fd.add_output(T27)

def nanogpt_attn_bwd_iobytes(size: tuple, dtype: torch.dtype):
    # Manual IOByte computation is required since nvFuser input/outputs (grad_out, attn, dropout_mask, bias_mask, grad_input]) differ from baseline input/outputs (output, grad_output).

    # Total IO bytes = grad_out ([bs, nh, seq_len, seq_len], dtype) + attn ([bs, nh, seq_len, seq_len], dtype) +
    #       dropout_mask ([bs, nh, seq_len, seq_len], bool) + bias_mask ([bs, nh, seq_len, seq_len], bool) + grad_in ([bs, nh, seq_len, seq_len], dtype)
    bs, seq_len, nh, n_embd = size

    return int(
        bs * nh * seq_len * seq_len * (3 * dtype.itemsize + 2 * torch.bool.itemsize)
    )

@pytest.mark.parametrize("size", generate_attn_inputs())
@pytest.mark.parametrize("dtype", FLOAT_DTYPES)
@pytest.mark.inner_persistent
def test_nanogpt_attn_bwd_nvf_benchmark(
    benchmark,
    size: tuple,
    dtype: torch.dtype,
    disable_validation: bool,
    disable_benchmarking: bool,
):
    batch_size, seq_len, nh, n_embd = size
    hs = n_embd // nh
    dropout_p = 0.2
    inputs = torch.randn(
        batch_size, nh, seq_len, seq_len, device="cuda", dtype=dtype, requires_grad=True
    )
    bias = torch.tril(torch.ones(seq_len, seq_len, device="cuda")).view(
        1, 1, seq_len, seq_len
    )
    dropout_mask = torch.lt(
        torch.rand(batch_size, nh, seq_len, seq_len, device="cuda"), 1 - dropout_p
    )
    bias_mask = bias[:, :, :seq_len, :seq_len] == 0
    bias_mask = bias_mask.broadcast_to(batch_size, nh, seq_len, seq_len)
    grads = torch.randn(batch_size, nh, seq_len, seq_len, device="cuda", dtype=dtype)

    attn = inputs / (hs**0.5)
    attn = attn.masked_fill(bias[:, :, :seq_len, :seq_len] == 0, float("-inf"))
    attn = torch.nn.functional.softmax(attn, dim=-1)

    with FusionDefinition() as fd:
        nanogpt_attn_bwd_fusion(fd, torch_dtype_to_nvfuser_dtype(dtype), hs, dropout_p)

    if not disable_validation:
        # Use dropout_mask instead of torch.nn.functional.dropout for validating results.
        out = attn * dropout_mask * 1 / (1 - dropout_p)
        out.backward(grads)
        fd.validate([grads, attn, dropout_mask, bias_mask], [inputs.grad])

    if not disable_benchmarking:
        run_benchmark(benchmark, fd.execute, [grads, attn, dropout_mask, bias_mask])

@pytest.mark.parametrize("executor", DEFAULT_EXECUTORS)
@pytest.mark.parametrize("size", generate_attn_inputs())
@pytest.mark.parametrize("dtype", FLOAT_DTYPES)
@pytest.mark.inner_persistent
def test_nanogpt_attn_bwd_baseline_benchmark(
    benchmark,
    size: tuple,
    dtype: torch.dtype,
    executor: str,
):
    if executor == "torchcompile":
        clear_dynamo_cache()
    batch_size, seq_len, nh, n_embd = size
    dropout_p = 0.2
    inputs = torch.randn(
        batch_size, nh, seq_len, seq_len, device="cuda", dtype=dtype, requires_grad=True
    )
    bias = torch.tril(torch.ones(seq_len, seq_len, device="cuda")).view(


    attn = attn.masked_fill(bias[:, :, :seq_len, :seq_len] == 0, float("-inf"))
    attn = torch.nn.functional.softmax(attn, dim=-1)

    with FusionDefinition() as fd:
        nanogpt_attn_bwd_fusion(fd, torch_dtype_to_nvfuser_dtype(dtype), hs, dropout_p)

    if not disable_validation:
        # Use dropout_mask instead of torch.nn.functional.dropout for validating results.
        out = attn * dropout_mask * 1 / (1 - dropout_p)
        out.backward(grads)
        fd.validate([grads, attn, dropout_mask, bias_mask], [inputs.grad])

    if not disable_benchmarking:
        run_benchmark(benchmark, fd.execute, [grads, attn, dropout_mask, bias_mask])

@pytest.mark.parametrize("executor", DEFAULT_EXECUTORS)
@pytest.mark.parametrize("size", generate_attn_inputs())
@pytest.mark.parametrize("dtype", FLOAT_DTYPES)
@pytest.mark.inner_persistent
def test_nanogpt_attn_bwd_baseline_benchmark(
    benchmark,
    size: tuple,
    dtype: torch.dtype,
    executor: str,
):
    if executor == "torchcompile":
        clear_dynamo_cache()
    batch_size, seq_len, nh, n_embd = size
    dropout_p = 0.2
    inputs = torch.randn(
        batch_size, nh, seq_len, seq_len, device="cuda", dtype=dtype, requires_grad=True
    )
    bias = torch.tril(torch.ones(seq_len, seq_len, device="cuda")).view(
        1, 1, seq_len, seq_len
    )

    # Compile the fwd fn for torchcompile
    fwd_fn = with_executor(executor, nanogpt_attn)
    fwd_inputs = [inputs, bias, size, dropout_p]
    outputs = fwd_fn(fwd_inputs)

    grads = torch.randn(batch_size, nh, seq_len, seq_len, device="cuda", dtype=dtype)

    # Manually compute IOBytes: See PR #1725
    run_benchmark(
        benchmark,
        unary_bwd_torch,
        [outputs, grads, *fwd_inputs],
        iobytes=nanogpt_attn_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,
    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_rmsnorm

def dropout_rmsnorm_fwd_fusion(
    fd: FusionDefinition,
    dtype: DataType,
    dropout_p: float,
    eps: float = 1e-5,
) -> None:
    """
    Forward pass fusion definition for computing:
        output = rmsnorm (input2 + dropout (input1, p=dropout_p))

    Fusion inputs: input, weights
    Fusion outputs: output, dropout_mask, rms
    """
    T0 = fd.define_tensor(
        shape=[-1, -1], contiguity=[True, True], dtype=dtype, is_cpu=False
    )
    T2 = fd.define_tensor(
        shape=[-1, -1], contiguity=[True, True], dtype=dtype, is_cpu=False
    )
    T1 = fd.define_tensor(shape=[-1], contiguity=[True], dtype=dtype, is_cpu=False)

    S2 = fd.define_scalar(0.00000, dtype=DataType.Double)
    S3 = fd.define_scalar(1.00000, dtype=DataType.Double)

    V6 = T0.shape()
    T7 = fd.ops.uniform(S2, S3, shape=V6, dtype=DataType.Float)
    S8 = fd.define_scalar(1 - dropout_p, dtype=DataType.Double)
    T9 = fd.ops.lt(T7, S8)
    T10 = fd.ops.cast(T9, dtype=DataType.Float)

    if dtype in PROMOTE_DTYPES:
        T0 = fd.ops.cast(T0, dtype=DataType.Float)
        T1 = fd.ops.cast(T1, dtype=DataType.Float)
        T2 = fd.ops.cast(T2, dtype=DataType.Float)

    T12 = fd.ops.mul(T0, T10)
    S13 = fd.define_scalar(1 / (1 - dropout_p), dtype=DataType.Double)
    T14 = fd.ops.mul(T12, S13)
    T15 = fd.ops.add(T2, T14)
    S16 = fd.define_scalar(2.00000, dtype=DataType.Double)
    T17 = fd.ops.pow(T15, S16)
    T18 = fd.ops.sum(T17, dims=[1], keepdim=False, dtype=DataType.Null)

    T22 = fd.ops.broadcast_in_dim(T18, shape=[T0.size(0), 1], broadcast_dims=[0])

    S24 = fd.ops.reciprocal(T0.size(1))
    T25 = fd.ops.mul(T22, S24)
    S26 = fd.define_scalar(eps, dtype=DataType.Double)
    T27 = fd.ops.add(T25, S26)
    T28 = fd.ops.sqrt(T27)

    T33 = fd.ops.broadcast_in_dim(T28, shape=V6, broadcast_dims=[0, 1])

    T35 = fd.ops.reciprocal(T33)
    T36 = fd.ops.mul(T15, T35)
    T40 = fd.ops.broadcast_in_dim(T1, shape=V6, broadcast_dims=[1])
    T42 = fd.ops.mul(T40, T36)

    if dtype in PROMOTE_DTYPES:
        T42 = fd.ops.cast(T42, dtype=dtype)

    fd.add_output(T42)
    fd.add_output(T9)
    fd.add_output(T28)

def dropout_rmsnorm_fwd_iobytes(size: tuple, dtype: torch.dtype):
    # Manual IOByte computation is required since nvFuser input/outputs differ from baseline outputs (output).
    nvf_inp_out = {
        # Inputs
        "input1": (size, dtype),
        "input2": (size, dtype),
        "weights": (size[1], dtype),
        # Outputs
        "rms": (size[0], torch.float),
        "output": (size, dtype),
        "dropout_mask": (size, torch.bool),
    }
    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_persistent
def test_dropout_rmsnorm_fwd_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)
    input2 = torch.randn(size, device="cuda", dtype=dtype)
    weights = torch.randn(size[1], device="cuda", dtype=dtype)

    dropout_p = 0.2

    with FusionDefinition() as fd:
        dropout_rmsnorm_fwd_fusion(
            fd, torch_dtype_to_nvfuser_dtype(dtype), dropout_p, eps
        )

    if not disable_validation:
        # For validating use a fusion definition with dropout_p=0.0
        with FusionDefinition() as val_fd:
            dropout_rmsnorm_fwd_fusion(
                val_fd, torch_dtype_to_nvfuser_dtype(dtype), dropout_p=0.0, eps=eps
            )

        dropout_mask = torch.ones(size, dtype=torch.bool, device="cuda")

        x = input2.to(torch.double) + torch.nn.functional.dropout(
            input1.to(torch.double), p=0.0
        )
        rms_eps = torch.sqrt(x.pow(2).mean(-1, keepdim=True) + eps)
        eager_output = weights.to(torch.double) * (x / rms_eps)
        val_fd.validate(
            [input1, input2, weights],
            [eager_output.to(dtype), dropout_mask, rms_eps.to(torch.float)],
        )

    if not disable_benchmarking:
        run_benchmark(benchmark, fd.execute, [input1, input2, weights])

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

    inputs = [
        torch.randn(size, device="cuda", dtype=dtype, requires_grad=True),
        torch.randn(size, device="cuda", dtype=dtype, requires_grad=True),
        torch.ones(size[1], device="cuda", dtype=dtype, requires_grad=True),
        dropout_p,
    ]

    benchmark_fn = with_executor(executor, dropout_rmsnorm)

    # Manually compute IOBytes: See PR #1725
    run_benchmark(
        benchmark,
        benchmark_fn,
        inputs,
        iobytes=dropout_rmsnorm_fwd_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,
    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_rmsnorm

def dropout_rmsnorm_fwd_fusion(
    fd: FusionDefinition,
    dtype: DataType,
    dropout_p: float,
    eps: float = 1e-5,
) -> None:
    """
    Forward pass fusion definition for computing:
        output = rmsnorm (input2 + dropout (input1, p=dropout_p))

    Fusion inputs: input, weights
    Fusion outputs: output, dropout_mask, rms
    """
    T0 = fd.define_tensor(
        shape=[-1, -1], contiguity=[True, True], dtype=dtype, is_cpu=False
    )
    T2 = fd.define_tensor(
        shape=[-1, -1], contiguity=[True, True], dtype=dtype, is_cpu=False
    )
    T1 = fd.define_tensor(shape=[-1], contiguity=[True], dtype=dtype, is_cpu=False)

    S2 = fd.define_scalar(0.00000, dtype=DataType.Double)
    S3 = fd.define_scalar(1.00000, dtype=DataType.Double)

    V6 = T0.shape()
    T7 = fd.ops.uniform(S2, S3, shape=V6, dtype=DataType.Float)
    S8 = fd.define_scalar(1 - dropout_p, dtype=DataType.Double)
    T9 = fd.ops.lt(T7, S8)
    T10 = fd.ops.cast(T9, dtype=DataType.Float)

    if dtype in PROMOTE_DTYPES:
        T0 = fd.ops.cast(T0, dtype=DataType.Float)
        T1 = fd.ops.cast(T1, dtype=DataType.Float)
        T2 = fd.ops.cast(T2, dtype=DataType.Float)

    T12 = fd.ops.mul(T0, T10)
    S13 = fd.define_scalar(1 / (1 - dropout_p), dtype=DataType.Double)
    T14 = fd.ops.mul(T12, S13)
    T15 = fd.ops.add(T2, T14)
    S16 = fd.define_scalar(2.00000, dtype=DataType.Double)
    T17 = fd.ops.pow(T15, S16)
    T18 = fd.ops.sum(T17, dims=[1], keepdim=False, dtype=DataType.Null)

    T22 = fd.ops.broadcast_in_dim(T18, shape=[T0.size(0), 1], broadcast_dims=[0])

    S24 = fd.ops.reciprocal(T0.size(1))
    T25 = fd.ops.mul(T22, S24)
    S26 = fd.define_scalar(eps, dtype=DataType.Double)
    T27 = fd.ops.add(T25, S26)
    T28 = fd.ops.sqrt(T27)

    T33 = fd.ops.broadcast_in_dim(T28, shape=V6, broadcast_dims=[0, 1])

    T35 = fd.ops.reciprocal(T33)
    T36 = fd.ops.mul(T15, T35)
    T40 = fd.ops.broadcast_in_dim(T1, shape=V6, broadcast_dims=[1])
    T42 = fd.ops.mul(T40, T36)

    if dtype in PROMOTE_DTYPES:
        T42 = fd.ops.cast(T42, dtype=dtype)

    fd.add_output(T42)
    fd.add_output(T9)
    fd.add_output(T28)

def dropout_rmsnorm_fwd_iobytes(size: tuple, dtype: torch.dtype):
    # Manual IOByte computation is required since nvFuser input/outputs differ from baseline outputs (output).
    nvf_inp_out = {
        # Inputs
        "input1": (size, dtype),
        "input2": (size, dtype),
        "weights": (size[1], dtype),
        # Outputs
        "rms": (size[0], torch.float),
        "output": (size, dtype),
        "dropout_mask": (size, torch.bool),
    }
    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_persistent


    T25 = fd.ops.mul(T22, S24)
    S26 = fd.define_scalar(eps, dtype=DataType.Double)
    T27 = fd.ops.add(T25, S26)
    T28 = fd.ops.sqrt(T27)

    T33 = fd.ops.broadcast_in_dim(T28, shape=V6, broadcast_dims=[0, 1])

    T35 = fd.ops.reciprocal(T33)
    T36 = fd.ops.mul(T15, T35)
    T40 = fd.ops.broadcast_in_dim(T1, shape=V6, broadcast_dims=[1])
    T42 = fd.ops.mul(T40, T36)

    if dtype in PROMOTE_DTYPES:
        T42 = fd.ops.cast(T42, dtype=dtype)

    fd.add_output(T42)
    fd.add_output(T9)
    fd.add_output(T28)

def dropout_rmsnorm_fwd_iobytes(size: tuple, dtype: torch.dtype):
    # Manual IOByte computation is required since nvFuser input/outputs differ from baseline outputs (output).
    nvf_inp_out = {
        # Inputs
        "input1": (size, dtype),
        "input2": (size, dtype),
        "weights": (size[1], dtype),
        # Outputs
        "rms": (size[0], torch.float),
        "output": (size, dtype),
        "dropout_mask": (size, torch.bool),
    }
    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_persistent
def test_dropout_rmsnorm_fwd_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)
    input2 = torch.randn(size, device="cuda", dtype=dtype)
    weights = torch.randn(size[1], device="cuda", dtype=dtype)

    dropout_p = 0.2

    with FusionDefinition() as fd:


def dropout_rmsnorm_fwd_iobytes(size: tuple, dtype: torch.dtype):
    # Manual IOByte computation is required since nvFuser input/outputs differ from baseline outputs (output).
    nvf_inp_out = {
        # Inputs
        "input1": (size, dtype),
        "input2": (size, dtype),
        "weights": (size[1], dtype),
        # Outputs
        "rms": (size[0], torch.float),
        "output": (size, dtype),
        "dropout_mask": (size, torch.bool),
    }
    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_persistent
def test_dropout_rmsnorm_fwd_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)
    input2 = torch.randn(size, device="cuda", dtype=dtype)
    weights = torch.randn(size[1], device="cuda", dtype=dtype)

    dropout_p = 0.2

    with FusionDefinition() as fd:
        dropout_rmsnorm_fwd_fusion(
            fd, torch_dtype_to_nvfuser_dtype(dtype), dropout_p, eps
        )

    if not disable_validation:
        # For validating use a fusion definition with dropout_p=0.0
        with FusionDefinition() as val_fd:
            dropout_rmsnorm_fwd_fusion(
                val_fd, torch_dtype_to_nvfuser_dtype(dtype), dropout_p=0.0, eps=eps
            )

        dropout_mask = torch.ones(size, dtype=torch.bool, device="cuda")

        x = input2.to(torch.double) + torch.nn.functional.dropout(
            input1.to(torch.double), p=0.0
        )
        rms_eps = torch.sqrt(x.pow(2).mean(-1, keepdim=True) + eps)
        eager_output = weights.to(torch.double) * (x / rms_eps)
        val_fd.validate(
            [input1, input2, weights],
            [eager_output.to(dtype), dropout_mask, rms_eps.to(torch.float)],
        )

    if not disable_benchmarking:
        run_benchmark(benchmark, fd.execute, [input1, input2, weights])

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

    inputs = [
        torch.randn(size, device="cuda", dtype=dtype, requires_grad=True),
        torch.randn(size, device="cuda", dtype=dtype, requires_grad=True),
        torch.ones(size[1], device="cuda", dtype=dtype, requires_grad=True),


        dropout_mask = torch.ones(size, dtype=torch.bool, device="cuda")

        x = input2.to(torch.double) + torch.nn.functional.dropout(
            input1.to(torch.double), p=0.0
        )
        rms_eps = torch.sqrt(x.pow(2).mean(-1, keepdim=True) + eps)
        eager_output = weights.to(torch.double) * (x / rms_eps)
        val_fd.validate(
            [input1, input2, weights],
            [eager_output.to(dtype), dropout_mask, rms_eps.to(torch.float)],
        )

    if not disable_benchmarking:
        run_benchmark(benchmark, fd.execute, [input1, input2, weights])

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

    inputs = [
        torch.randn(size, device="cuda", dtype=dtype, requires_grad=True),
        torch.randn(size, device="cuda", dtype=dtype, requires_grad=True),
        torch.ones(size[1], device="cuda", dtype=dtype, requires_grad=True),
        dropout_p,
    ]

    benchmark_fn = with_executor(executor, dropout_rmsnorm)

    # Manually compute IOBytes: See PR #1725
    run_benchmark(
        benchmark,
        benchmark_fn,
        inputs,
        iobytes=dropout_rmsnorm_fwd_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,
    with_executor,
    DEFAULT_EXECUTORS,
)
import torch
from .global_params import generate_attn_inputs, FLOAT_DTYPES, PROMOTE_DTYPES
from .torch_ops import huggingface_attn

# Fusion from huggingface attention implementation
# The nvFuser defintion only includes the non-matmul computation (add + reshape + softmax + dropout)
def huggingface_attn_bwd_fusion(
    fd: FusionDefinition,
    dtype: DataType,
    batch_size: int,
    num_heads: int,
    dropout_p: float,
):
    S0 = fd.define_scalar(1 / (1 - dropout_p), dtype=DataType.Double)
    T5 = fd.define_tensor(
        shape=[-1, -1, -1], contiguity=[True, True, True], dtype=dtype, is_cpu=False
    )
    T6 = fd.define_tensor(
        shape=[-1, -1, -1], contiguity=[True, True, True], dtype=dtype, is_cpu=False
    )
    T7 = fd.define_tensor(
        shape=[-1, -1, -1],
        contiguity=[True, True, True],
        dtype=DataType.Bool,
        is_cpu=False,
    )

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

    T10 = fd.ops.mul(T5, S0)
    T11 = fd.ops.mul(T10, T7)
    T13 = fd.ops.mul(T6, T11)
    T14 = fd.ops.sum(T13, dims=[2], keepdim=False, dtype=DataType.Null)

    T19 = fd.ops.broadcast_in_dim(
        T14, shape=[T5.size(0), T5.size(1), 1], broadcast_dims=[0, 1]
    )
    V24 = T5.shape()
    T25 = fd.ops.broadcast_in_dim(T19, shape=V24, broadcast_dims=[0, 1, 2])
    T27 = fd.ops.sub(T11, T25)
    T28 = fd.ops.mul(T6, T27)
    T35 = fd.ops.reshape(T28, new_shape=[batch_size, num_heads, T5.size(1), T5.size(2)])

    if dtype in PROMOTE_DTYPES:
        T35 = fd.ops.cast(T35, dtype=dtype)
    fd.add_output(T35)

def huggingface_attn_bwd_iobytes(size: tuple, dtype: torch.dtype):
    # Manual IOByte computation is required since nvFuser input/outputs (grad_out, attn, dropout_mask, grad_input]) differ from baseline input/outputs (output, grad_output).

    # Total IO bytes = grad_output ([bs, nh, seq_len, seq_len], dtype) + attn ([bs*nh, seq_len, seq_len], dtype) + dropout_mask ([bs*nh, seq_len, seq_len], bool) + grad_input ([bs, nh, seq_len, seq_len], dtype)
    bs, seq_len, nh, n_embd = size
    return int(bs * nh * seq_len * seq_len * (dtype.itemsize * 3 + torch.bool.itemsize))

@pytest.mark.parametrize("size", generate_attn_inputs())
@pytest.mark.parametrize("dtype", FLOAT_DTYPES)
@pytest.mark.inner_persistent
def test_huggingface_attn_bwd_nvf_benchmark(
    benchmark,
    size: tuple,
    dtype: torch.dtype,
    disable_validation: bool,
    disable_benchmarking: bool,
):
    batch_size, seq_len, nh, n_embd = size

    dropout_p = 0.2
    inputs = torch.randn(
        batch_size, nh, seq_len, seq_len, device="cuda", dtype=dtype, requires_grad=True
    )
    dropout_mask = torch.lt(
        torch.rand(batch_size * nh, seq_len, seq_len, device="cuda"), 1 - dropout_p
    )
    attention_mask = torch.zeros(
        batch_size, nh, seq_len, seq_len, device="cuda", dtype=dtype
    )
    grads = torch.randn(batch_size * nh, seq_len, seq_len, device="cuda", dtype=dtype)

    attn = (inputs + attention_mask).view(batch_size * nh, seq_len, seq_len)
    attn = torch.nn.functional.softmax(attn, dim=-1)

    with FusionDefinition() as fd:
        huggingface_attn_bwd_fusion(
            fd, torch_dtype_to_nvfuser_dtype(dtype), batch_size, nh, dropout_p
        )

    if not disable_validation:
        # Use dropout_mask instead of torch.nn.functional.dropout for validating results.
        out = attn * dropout_mask * 1 / (1 - dropout_p)
        out.backward(grads)
        fd.validate([grads, attn, dropout_mask], [inputs.grad])

    if not disable_benchmarking:
        run_benchmark(benchmark, fd.execute, [grads, attn, dropout_mask])

@pytest.mark.parametrize("executor", DEFAULT_EXECUTORS)
@pytest.mark.parametrize("size", generate_attn_inputs())
@pytest.mark.parametrize("dtype", FLOAT_DTYPES)
@pytest.mark.inner_persistent
def test_huggingface_attn_bwd_baseline_benchmark(
    benchmark,
    size: tuple,
    dtype: torch.dtype,
    executor: str,
):
    if executor == "torchcompile":
        clear_dynamo_cache()
    batch_size, seq_len, nh, n_embd = size
    dropout_p = 0.2
    inputs = torch.randn(
        batch_size, nh, seq_len, seq_len, device="cuda", dtype=dtype, requires_grad=True
    )
    attention_mask = torch.zeros(
        batch_size, nh, seq_len, seq_len, device="cuda", dtype=dtype
    )

    # Compile the fwd fn for torchcompile
    fwd_fn = with_executor(executor, huggingface_attn)
    fwd_inputs = [inputs, attention_mask, size, dropout_p]
    outputs = fwd_fn(fwd_inputs)
    grads = torch.randn(batch_size * nh, seq_len, seq_len, device="cuda", dtype=dtype)

    # Manually compute IOBytes: See PR #1725
    run_benchmark(
        benchmark,
        unary_bwd_torch,
        [outputs, grads, *fwd_inputs],
        iobytes=huggingface_attn_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,
    with_executor,
    DEFAULT_EXECUTORS,
)
import torch
from .global_params import generate_attn_inputs, FLOAT_DTYPES, PROMOTE_DTYPES
from .torch_ops import huggingface_attn

# Fusion from huggingface attention implementation
# The nvFuser defintion only includes the non-matmul computation (add + reshape + softmax + dropout)
def huggingface_attn_bwd_fusion(
    fd: FusionDefinition,
    dtype: DataType,
    batch_size: int,
    num_heads: int,
    dropout_p: float,
):
    S0 = fd.define_scalar(1 / (1 - dropout_p), dtype=DataType.Double)
    T5 = fd.define_tensor(
        shape=[-1, -1, -1], contiguity=[True, True, True], dtype=dtype, is_cpu=False
    )
    T6 = fd.define_tensor(
        shape=[-1, -1, -1], contiguity=[True, True, True], dtype=dtype, is_cpu=False
    )
    T7 = fd.define_tensor(
        shape=[-1, -1, -1],
        contiguity=[True, True, True],
        dtype=DataType.Bool,
        is_cpu=False,
    )

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

    T10 = fd.ops.mul(T5, S0)
    T11 = fd.ops.mul(T10, T7)
    T13 = fd.ops.mul(T6, T11)
    T14 = fd.ops.sum(T13, dims=[2], keepdim=False, dtype=DataType.Null)

    T19 = fd.ops.broadcast_in_dim(
        T14, shape=[T5.size(0), T5.size(1), 1], broadcast_dims=[0, 1]
    )
    V24 = T5.shape()
    T25 = fd.ops.broadcast_in_dim(T19, shape=V24, broadcast_dims=[0, 1, 2])
    T27 = fd.ops.sub(T11, T25)
    T28 = fd.ops.mul(T6, T27)
    T35 = fd.ops.reshape(T28, new_shape=[batch_size, num_heads, T5.size(1), T5.size(2)])

    if dtype in PROMOTE_DTYPES:
        T35 = fd.ops.cast(T35, dtype=dtype)
    fd.add_output(T35)

def huggingface_attn_bwd_iobytes(size: tuple, dtype: torch.dtype):
    # Manual IOByte computation is required since nvFuser input/outputs (grad_out, attn, dropout_mask, grad_input]) differ from baseline input/outputs (output, grad_output).

    # Total IO bytes = grad_output ([bs, nh, seq_len, seq_len], dtype) + attn ([bs*nh, seq_len, seq_len], dtype) + dropout_mask ([bs*nh, seq_len, seq_len], bool) + grad_input ([bs, nh, seq_len, seq_len], dtype)
    bs, seq_len, nh, n_embd = size
    return int(bs * nh * seq_len * seq_len * (dtype.itemsize * 3 + torch.bool.itemsize))

@pytest.mark.parametrize("size", generate_attn_inputs())
@pytest.mark.parametrize("dtype", FLOAT_DTYPES)
@pytest.mark.inner_persistent
def test_huggingface_attn_bwd_nvf_benchmark(
    benchmark,
    size: tuple,
    dtype: torch.dtype,
    disable_validation: bool,
    disable_benchmarking: bool,
):


    T10 = fd.ops.mul(T5, S0)
    T11 = fd.ops.mul(T10, T7)
    T13 = fd.ops.mul(T6, T11)
    T14 = fd.ops.sum(T13, dims=[2], keepdim=False, dtype=DataType.Null)

    T19 = fd.ops.broadcast_in_dim(
        T14, shape=[T5.size(0), T5.size(1), 1], broadcast_dims=[0, 1]
    )
    V24 = T5.shape()
    T25 = fd.ops.broadcast_in_dim(T19, shape=V24, broadcast_dims=[0, 1, 2])
    T27 = fd.ops.sub(T11, T25)
    T28 = fd.ops.mul(T6, T27)
    T35 = fd.ops.reshape(T28, new_shape=[batch_size, num_heads, T5.size(1), T5.size(2)])

    if dtype in PROMOTE_DTYPES:
        T35 = fd.ops.cast(T35, dtype=dtype)
    fd.add_output(T35)

def huggingface_attn_bwd_iobytes(size: tuple, dtype: torch.dtype):
    # Manual IOByte computation is required since nvFuser input/outputs (grad_out, attn, dropout_mask, grad_input]) differ from baseline input/outputs (output, grad_output).

    # Total IO bytes = grad_output ([bs, nh, seq_len, seq_len], dtype) + attn ([bs*nh, seq_len, seq_len], dtype) + dropout_mask ([bs*nh, seq_len, seq_len], bool) + grad_input ([bs, nh, seq_len, seq_len], dtype)
    bs, seq_len, nh, n_embd = size
    return int(bs * nh * seq_len * seq_len * (dtype.itemsize * 3 + torch.bool.itemsize))

@pytest.mark.parametrize("size", generate_attn_inputs())
@pytest.mark.parametrize("dtype", FLOAT_DTYPES)
@pytest.mark.inner_persistent
def test_huggingface_attn_bwd_nvf_benchmark(
    benchmark,
    size: tuple,
    dtype: torch.dtype,
    disable_validation: bool,
    disable_benchmarking: bool,
):
    batch_size, seq_len, nh, n_embd = size

    dropout_p = 0.2
    inputs = torch.randn(
        batch_size, nh, seq_len, seq_len, device="cuda", dtype=dtype, requires_grad=True
    )
    dropout_mask = torch.lt(
        torch.rand(batch_size * nh, seq_len, seq_len, device="cuda"), 1 - dropout_p


    T27 = fd.ops.sub(T11, T25)
    T28 = fd.ops.mul(T6, T27)
    T35 = fd.ops.reshape(T28, new_shape=[batch_size, num_heads, T5.size(1), T5.size(2)])

    if dtype in PROMOTE_DTYPES:
        T35 = fd.ops.cast(T35, dtype=dtype)
    fd.add_output(T35)

def huggingface_attn_bwd_iobytes(size: tuple, dtype: torch.dtype):
    # Manual IOByte computation is required since nvFuser input/outputs (grad_out, attn, dropout_mask, grad_input]) differ from baseline input/outputs (output, grad_output).

    # Total IO bytes = grad_output ([bs, nh, seq_len, seq_len], dtype) + attn ([bs*nh, seq_len, seq_len], dtype) + dropout_mask ([bs*nh, seq_len, seq_len], bool) + grad_input ([bs, nh, seq_len, seq_len], dtype)
    bs, seq_len, nh, n_embd = size
    return int(bs * nh * seq_len * seq_len * (dtype.itemsize * 3 + torch.bool.itemsize))

@pytest.mark.parametrize("size", generate_attn_inputs())
@pytest.mark.parametrize("dtype", FLOAT_DTYPES)
@pytest.mark.inner_persistent
def test_huggingface_attn_bwd_nvf_benchmark(
    benchmark,
    size: tuple,
    dtype: torch.dtype,
    disable_validation: bool,
    disable_benchmarking: bool,
):
    batch_size, seq_len, nh, n_embd = size

    dropout_p = 0.2
    inputs = torch.randn(
        batch_size, nh, seq_len, seq_len, device="cuda", dtype=dtype, requires_grad=True
    )
    dropout_mask = torch.lt(
        torch.rand(batch_size * nh, seq_len, seq_len, device="cuda"), 1 - dropout_p
    )
    attention_mask = torch.zeros(
        batch_size, nh, seq_len, seq_len, device="cuda", dtype=dtype
    )
    grads = torch.randn(batch_size * nh, seq_len, seq_len, device="cuda", dtype=dtype)

    attn = (inputs + attention_mask).view(batch_size * nh, seq_len, seq_len)
    attn = torch.nn.functional.softmax(attn, dim=-1)

    with FusionDefinition() as fd:
        huggingface_attn_bwd_fusion(
            fd, torch_dtype_to_nvfuser_dtype(dtype), batch_size, nh, dropout_p
        )

    if not disable_validation:
        # Use dropout_mask instead of torch.nn.functional.dropout for validating results.
        out = attn * dropout_mask * 1 / (1 - dropout_p)
        out.backward(grads)
        fd.validate([grads, attn, dropout_mask], [inputs.grad])

    if not disable_benchmarking:
        run_benchmark(benchmark, fd.execute, [grads, attn, dropout_mask])

@pytest.mark.parametrize("executor", DEFAULT_EXECUTORS)
@pytest.mark.parametrize("size", generate_attn_inputs())
@pytest.mark.parametrize("dtype", FLOAT_DTYPES)
@pytest.mark.inner_persistent
def test_huggingface_attn_bwd_baseline_benchmark(
    benchmark,
    size: tuple,
    dtype: torch.dtype,
    executor: str,
):
    if executor == "torchcompile":
        clear_dynamo_cache()
    batch_size, seq_len, nh, n_embd = size
    dropout_p = 0.2
    inputs = torch.randn(
        batch_size, nh, seq_len, seq_len, device="cuda", dtype=dtype, requires_grad=True
    )
    attention_mask = torch.zeros(


    with FusionDefinition() as fd:
        huggingface_attn_bwd_fusion(
            fd, torch_dtype_to_nvfuser_dtype(dtype), batch_size, nh, dropout_p
        )

    if not disable_validation:
        # Use dropout_mask instead of torch.nn.functional.dropout for validating results.
        out = attn * dropout_mask * 1 / (1 - dropout_p)
        out.backward(grads)
        fd.validate([grads, attn, dropout_mask], [inputs.grad])

    if not disable_benchmarking:
        run_benchmark(benchmark, fd.execute, [grads, attn, dropout_mask])

@pytest.mark.parametrize("executor", DEFAULT_EXECUTORS)
@pytest.mark.parametrize("size", generate_attn_inputs())
@pytest.mark.parametrize("dtype", FLOAT_DTYPES)
@pytest.mark.inner_persistent
def test_huggingface_attn_bwd_baseline_benchmark(
    benchmark,
    size: tuple,
    dtype: torch.dtype,
    executor: str,
):
    if executor == "torchcompile":
        clear_dynamo_cache()
    batch_size, seq_len, nh, n_embd = size
    dropout_p = 0.2
    inputs = torch.randn(
        batch_size, nh, seq_len, seq_len, device="cuda", dtype=dtype, requires_grad=True
    )
    attention_mask = torch.zeros(
        batch_size, nh, seq_len, seq_len, device="cuda", dtype=dtype
    )

    # Compile the fwd fn for torchcompile
    fwd_fn = with_executor(executor, huggingface_attn)
    fwd_inputs = [inputs, attention_mask, size, dropout_p]
    outputs = fwd_fn(fwd_inputs)
    grads = torch.randn(batch_size * nh, seq_len, seq_len, device="cuda", dtype=dtype)

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