    @property
    def communicator(self):
        return self._communicator

    @property
    def size(self):
        return self._communicator.size()

    @property
    def rank(self):
        return self._communicator.rank()

    @property
    def local_size(self):
        return self._communicator.local_size()

    @property
    def local_rank(self):
        return self._communicator.local_rank()

    def shard_tensor(
        self, t: torch.Tensor, dim: int, mesh: nvfuser.multidevice.DeviceMesh
    ) -> torch.Tensor:
        assert t.is_cpu, (
            "This is not strictly required but it's a general good practice "
            "for unit tests to create unsharded data on CPU to reduce GPU "
            "memory footprint."
        )
        return mesh.shard_tensor(t, dim, self.rank).cuda(self.rank)

@pytest.fixture
def multidevice_settings():
    return MultideviceSettings()

@pytest.mark.mpi
@pytest.mark.parametrize(
    "backend_type", [CommunicatorBackend.ucc, CommunicatorBackend.nccl]
)
@pytest.mark.parametrize("s", [1, 8])
@pytest.mark.parametrize("m", [2**16])
@pytest.mark.parametrize("k", [2**12, 2**16])
@pytest.mark.parametrize("n", [2**10])
@pytest.mark.parametrize("dtype", [torch.float16, torch.float32])
def test_overlap_allgather_matmul_stream_outermost(
    benchmark,
    multidevice_settings,
    backend_type,


    m,
    k,
    n,
    dtype,
    validate_output=False,
):
    """Test overlapping all-gather with matrix multiplication using stream parallelism.

    Args:
        benchmark: pytest-benchmark fixture
        multidevice_settings: Settings for multi-device execution
        backend_type: Communication backend to use
        s: Number of streams
        m: Matrix dimension m
        k: Matrix dimension k
        n: Matrix dimension n
        dtype: Data type for computation
        validate_output: Whether to validate output against reference
    """

    def fusion_definition(fd, m, k, n, s, d) -> None:
        """Fusion definition for overlapping all-gather with matrix multiplication.

        This fusion implements a matrix multiplication operation with overlapping
        all-gather communication, using stream parallelism for the outermost dimension.
        """
        x = fd.define_tensor(
            shape=[s, d, m // (s * d), k],
            contiguity=True,
            dtype=torch_dtype_to_nvfuser_dtype(dtype),
        )
        weight = fd.define_tensor(
            shape=[n, k],
            contiguity=True,
            dtype=torch_dtype_to_nvfuser_dtype(dtype),
        )
        bias = fd.define_tensor(
            shape=[n], contiguity=True, dtype=torch_dtype_to_nvfuser_dtype(dtype)
        )
        out = fd.ops.linear(x, weight, bias)
        fd.add_output(out)
        return x, weight, bias, out

    def multidevice_schedule(fd, tensors, num_devices):
        mesh = nvfuser.multidevice.DeviceMesh(range(num_devices))
        for tv in tensors:
            tv.set_device_mesh(mesh)

        x, weight, bias, out = tensors
        x.axis(1).parallelize(nvfuser.ParallelType.mesh_x)
        out.axis(0).parallelize(nvfuser.ParallelType.stream)

    d = multidevice_settings.size
    assert m % (s * d) == 0
    os.environ["UCC_CL_BASIC_TLS"] = "nccl"
    torch.cuda.set_device(multidevice_settings.local_rank)

    # Create input tensors
    x_unsharded = torch.testing.make_tensor(
        s, d, m // (s * d), k, dtype=dtype, device="cpu"
    )
    x = multidevice_settings.shard_tensor(


        This fusion implements a matrix multiplication operation with overlapping
        all-gather communication, using stream parallelism for the outermost dimension.
        """
        x = fd.define_tensor(
            shape=[s, d, m // (s * d), k],
            contiguity=True,
            dtype=torch_dtype_to_nvfuser_dtype(dtype),
        )
        weight = fd.define_tensor(
            shape=[n, k],
            contiguity=True,
            dtype=torch_dtype_to_nvfuser_dtype(dtype),
        )
        bias = fd.define_tensor(
            shape=[n], contiguity=True, dtype=torch_dtype_to_nvfuser_dtype(dtype)
        )
        out = fd.ops.linear(x, weight, bias)
        fd.add_output(out)
        return x, weight, bias, out

    def multidevice_schedule(fd, tensors, num_devices):
        mesh = nvfuser.multidevice.DeviceMesh(range(num_devices))
        for tv in tensors:
            tv.set_device_mesh(mesh)

        x, weight, bias, out = tensors
        x.axis(1).parallelize(nvfuser.ParallelType.mesh_x)
        out.axis(0).parallelize(nvfuser.ParallelType.stream)

    d = multidevice_settings.size
    assert m % (s * d) == 0
    os.environ["UCC_CL_BASIC_TLS"] = "nccl"
    torch.cuda.set_device(multidevice_settings.local_rank)

    # Create input tensors
    x_unsharded = torch.testing.make_tensor(
        s, d, m // (s * d), k, dtype=dtype, device="cpu"
    )
    x = multidevice_settings.shard_tensor(
        x_unsharded, 1, nvfuser.multidevice.DeviceMesh(range(multidevice_settings.size))
    )
    weight = torch.testing.make_tensor(n, k, dtype=dtype, device="cuda")
    bias = torch.testing.make_tensor(n, dtype=dtype, device="cuda")
    inputs = [x, weight, bias]

    # Create fusion definition
    with FusionDefinition() as fd:
        tensors = fusion_definition(fd, m, k, n, s, d)


    inputs = [x, weight, bias]

    # Create fusion definition
    with FusionDefinition() as fd:
        tensors = fusion_definition(fd, m, k, n, s, d)
        multidevice_schedule(fd, tensors, d)

    params = nvfuser.multidevice.MultiDeviceExecutorParams()
    params.backend_type = backend_type
    params.use_allocation_cache = True
    multidevice_executor = nvfuser.multidevice.MultiDeviceExecutor(fd.fusion, params)

    if validate_output:
        outputs = multidevice_executor.run(inputs)
        out = outputs[0].cpu()
        assert out.dtype == dtype
        assert out.shape == torch.Size([s, d, m // (s * d), n])
        out_ref = torch.nn.functional.linear(x_unsharded, weight.cpu(), bias.cpu())
        torch.testing.assert_close(out, out_ref, rtol=float("inf"), atol=1e-1)

    def benchmark_fn(*args):
        outputs = multidevice_executor.run(inputs)
        return outputs[0]

    benchmark_cuda_events_pedantic(
        benchmark,
        benchmark_fn,
        [inputs],
        warmup_rounds=BENCHMARK_CONFIG["warmup_rounds"],
        rounds=BENCHMARK_CONFIG["rounds"],
    )


# 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 gelu

def gelu_bwd_fusion(
    fd: FusionDefinition,
    dtype: DataType,
) -> None:
    input = fd.define_tensor(
        shape=[-1, -1], contiguity=[True, True], dtype=dtype, is_cpu=False
    )
    grad = fd.define_tensor(
        shape=[-1, -1], contiguity=[True, True], dtype=dtype, is_cpu=False
    )
    bias = fd.define_tensor(shape=[-1], contiguity=[True], dtype=dtype, is_cpu=False)
    if dtype in PROMOTE_DTYPES:
        input = fd.ops.cast(input, dtype=DataType.Float)
        grad = fd.ops.cast(grad, dtype=DataType.Float)
        bias = fd.ops.cast(bias, dtype=DataType.Float)
    S_079 = fd.define_scalar(0.79788456)
    S_004 = fd.define_scalar(0.044715)
    S_010 = fd.define_scalar(0.1070322243)
    bias = fd.ops.broadcast_in_dim(bias, shape=[1, input.size(-1)], broadcast_dims=[1])
    T1 = fd.ops.add(input, bias)
    T2 = fd.ops.mul(T1, S_079)
    T3 = fd.ops.mul(T1, S_004)
    T4 = fd.ops.mul(T3, T1)
    S1 = fd.define_scalar(1.0)
    T5 = fd.ops.add(S1, T4)
    T6 = fd.ops.mul(T2, T5)
    T7 = fd.ops.tanh(T6)
    S2 = fd.define_scalar(0.50)
    T8 = fd.ops.mul(T1, S2)
    T9 = fd.ops.mul(T7, T7)
    T10 = fd.ops.neg(T9)
    T11 = fd.ops.add(T10, S1)
    T12 = fd.ops.mul(T1, S_010)
    T13 = fd.ops.mul(T12, T1)
    T14 = fd.ops.add(T13, S_079)
    T15 = fd.ops.mul(T11, T14)
    T16 = fd.ops.mul(T8, T15)
    T17 = fd.ops.add(T7, S1)
    T18 = fd.ops.mul(T17, S2)
    T19 = fd.ops.add(T16, T18)
    T20 = fd.ops.mul(grad, T19)
    if dtype in PROMOTE_DTYPES:
        T20 = fd.ops.cast(T20, dtype=dtype)
    fd.add_output(T20)

def gelu_bwd_iobytes(size: tuple, dtype: torch.dtype):
    # Total IO bytes = in_tensor + grad_out + bias + grad_input
    return int(dtype.itemsize * (3 * np.prod(size) + size[1]))

@pytest.mark.parametrize("size", generate_input_sizes(dims=2))
@pytest.mark.parametrize("dtype", FLOAT_DTYPES)
@pytest.mark.pointwise
def test_gelu_bwd_nvf_benchmark(
    benchmark,
    size: tuple,
    dtype: torch.dtype,
    disable_validation: bool,
    disable_benchmarking: bool,
):
    inputs = torch.randn(size, device="cuda", dtype=dtype, requires_grad=True)
    grads = torch.randn(size, device="cuda", dtype=dtype)
    bias = torch.ones(size[-1], device="cuda", dtype=dtype)
    with FusionDefinition() as fd:
        gelu_bwd_fusion(fd, torch_dtype_to_nvfuser_dtype(dtype))

    if not disable_validation:
        eager_output = torch.nn.functional.gelu(
            inputs.to(torch.double) + bias.to(torch.double), approximate="tanh"
        )
        eager_output.backward(grads.to(torch.double))
        fd.validate([inputs, grads, bias], [inputs.grad])

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

@pytest.mark.parametrize("executor", DEFAULT_EXECUTORS)
@pytest.mark.parametrize("size", generate_input_sizes(dims=2))
@pytest.mark.parametrize("dtype", FLOAT_DTYPES)
@pytest.mark.pointwise
def test_gelu_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)
    bias = torch.ones(size[-1], device="cuda", dtype=dtype)
    grads = torch.randn(size, device="cuda", dtype=dtype)

    fwd_fn = with_executor(executor, gelu)
    fwd_inputs = [inputs, bias]
    outputs = fwd_fn(fwd_inputs)

    run_benchmark(
        benchmark,
        unary_bwd_torch,
        [outputs, grads, *fwd_inputs],
        iobytes=gelu_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 gelu

def gelu_bwd_fusion(
    fd: FusionDefinition,
    dtype: DataType,
) -> None:
    input = fd.define_tensor(
        shape=[-1, -1], contiguity=[True, True], dtype=dtype, is_cpu=False
    )
    grad = fd.define_tensor(
        shape=[-1, -1], contiguity=[True, True], dtype=dtype, is_cpu=False
    )
    bias = fd.define_tensor(shape=[-1], contiguity=[True], dtype=dtype, is_cpu=False)
    if dtype in PROMOTE_DTYPES:
        input = fd.ops.cast(input, dtype=DataType.Float)
        grad = fd.ops.cast(grad, dtype=DataType.Float)
        bias = fd.ops.cast(bias, dtype=DataType.Float)
    S_079 = fd.define_scalar(0.79788456)
    S_004 = fd.define_scalar(0.044715)
    S_010 = fd.define_scalar(0.1070322243)
    bias = fd.ops.broadcast_in_dim(bias, shape=[1, input.size(-1)], broadcast_dims=[1])
    T1 = fd.ops.add(input, bias)
    T2 = fd.ops.mul(T1, S_079)
    T3 = fd.ops.mul(T1, S_004)
    T4 = fd.ops.mul(T3, T1)
    S1 = fd.define_scalar(1.0)
    T5 = fd.ops.add(S1, T4)
    T6 = fd.ops.mul(T2, T5)
    T7 = fd.ops.tanh(T6)
    S2 = fd.define_scalar(0.50)
    T8 = fd.ops.mul(T1, S2)
    T9 = fd.ops.mul(T7, T7)
    T10 = fd.ops.neg(T9)
    T11 = fd.ops.add(T10, S1)
    T12 = fd.ops.mul(T1, S_010)
    T13 = fd.ops.mul(T12, T1)
    T14 = fd.ops.add(T13, S_079)
    T15 = fd.ops.mul(T11, T14)
    T16 = fd.ops.mul(T8, T15)
    T17 = fd.ops.add(T7, S1)
    T18 = fd.ops.mul(T17, S2)
    T19 = fd.ops.add(T16, T18)
    T20 = fd.ops.mul(grad, T19)
    if dtype in PROMOTE_DTYPES:
        T20 = fd.ops.cast(T20, dtype=dtype)
    fd.add_output(T20)

def gelu_bwd_iobytes(size: tuple, dtype: torch.dtype):
    # Total IO bytes = in_tensor + grad_out + bias + grad_input
    return int(dtype.itemsize * (3 * np.prod(size) + size[1]))

@pytest.mark.parametrize("size", generate_input_sizes(dims=2))
@pytest.mark.parametrize("dtype", FLOAT_DTYPES)
@pytest.mark.pointwise
def test_gelu_bwd_nvf_benchmark(
    benchmark,
    size: tuple,
    dtype: torch.dtype,
    disable_validation: bool,
    disable_benchmarking: bool,
):
    inputs = torch.randn(size, device="cuda", dtype=dtype, requires_grad=True)
    grads = torch.randn(size, device="cuda", dtype=dtype)
    bias = torch.ones(size[-1], device="cuda", dtype=dtype)


    T7 = fd.ops.tanh(T6)
    S2 = fd.define_scalar(0.50)
    T8 = fd.ops.mul(T1, S2)
    T9 = fd.ops.mul(T7, T7)
    T10 = fd.ops.neg(T9)
    T11 = fd.ops.add(T10, S1)
    T12 = fd.ops.mul(T1, S_010)
    T13 = fd.ops.mul(T12, T1)
    T14 = fd.ops.add(T13, S_079)
    T15 = fd.ops.mul(T11, T14)
    T16 = fd.ops.mul(T8, T15)
    T17 = fd.ops.add(T7, S1)
    T18 = fd.ops.mul(T17, S2)
    T19 = fd.ops.add(T16, T18)
    T20 = fd.ops.mul(grad, T19)
    if dtype in PROMOTE_DTYPES:
        T20 = fd.ops.cast(T20, dtype=dtype)
    fd.add_output(T20)

def gelu_bwd_iobytes(size: tuple, dtype: torch.dtype):
    # Total IO bytes = in_tensor + grad_out + bias + grad_input
    return int(dtype.itemsize * (3 * np.prod(size) + size[1]))

@pytest.mark.parametrize("size", generate_input_sizes(dims=2))
@pytest.mark.parametrize("dtype", FLOAT_DTYPES)
@pytest.mark.pointwise
def test_gelu_bwd_nvf_benchmark(
    benchmark,
    size: tuple,
    dtype: torch.dtype,
    disable_validation: bool,
    disable_benchmarking: bool,
):
    inputs = torch.randn(size, device="cuda", dtype=dtype, requires_grad=True)
    grads = torch.randn(size, device="cuda", dtype=dtype)
    bias = torch.ones(size[-1], device="cuda", dtype=dtype)
    with FusionDefinition() as fd:
        gelu_bwd_fusion(fd, torch_dtype_to_nvfuser_dtype(dtype))

    if not disable_validation:
        eager_output = torch.nn.functional.gelu(


    T14 = fd.ops.add(T13, S_079)
    T15 = fd.ops.mul(T11, T14)
    T16 = fd.ops.mul(T8, T15)
    T17 = fd.ops.add(T7, S1)
    T18 = fd.ops.mul(T17, S2)
    T19 = fd.ops.add(T16, T18)
    T20 = fd.ops.mul(grad, T19)
    if dtype in PROMOTE_DTYPES:
        T20 = fd.ops.cast(T20, dtype=dtype)
    fd.add_output(T20)

def gelu_bwd_iobytes(size: tuple, dtype: torch.dtype):
    # Total IO bytes = in_tensor + grad_out + bias + grad_input
    return int(dtype.itemsize * (3 * np.prod(size) + size[1]))

@pytest.mark.parametrize("size", generate_input_sizes(dims=2))
@pytest.mark.parametrize("dtype", FLOAT_DTYPES)
@pytest.mark.pointwise
def test_gelu_bwd_nvf_benchmark(
    benchmark,
    size: tuple,
    dtype: torch.dtype,
    disable_validation: bool,
    disable_benchmarking: bool,
):
    inputs = torch.randn(size, device="cuda", dtype=dtype, requires_grad=True)
    grads = torch.randn(size, device="cuda", dtype=dtype)
    bias = torch.ones(size[-1], device="cuda", dtype=dtype)
    with FusionDefinition() as fd:
        gelu_bwd_fusion(fd, torch_dtype_to_nvfuser_dtype(dtype))

    if not disable_validation:
        eager_output = torch.nn.functional.gelu(
            inputs.to(torch.double) + bias.to(torch.double), approximate="tanh"
        )
        eager_output.backward(grads.to(torch.double))
        fd.validate([inputs, grads, bias], [inputs.grad])

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

@pytest.mark.parametrize("executor", DEFAULT_EXECUTORS)
@pytest.mark.parametrize("size", generate_input_sizes(dims=2))
@pytest.mark.parametrize("dtype", FLOAT_DTYPES)
@pytest.mark.pointwise
def test_gelu_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)
    bias = torch.ones(size[-1], device="cuda", dtype=dtype)
    grads = torch.randn(size, device="cuda", dtype=dtype)

    fwd_fn = with_executor(executor, gelu)
    fwd_inputs = [inputs, bias]


    grads = torch.randn(size, device="cuda", dtype=dtype)
    bias = torch.ones(size[-1], device="cuda", dtype=dtype)
    with FusionDefinition() as fd:
        gelu_bwd_fusion(fd, torch_dtype_to_nvfuser_dtype(dtype))

    if not disable_validation:
        eager_output = torch.nn.functional.gelu(
            inputs.to(torch.double) + bias.to(torch.double), approximate="tanh"
        )
        eager_output.backward(grads.to(torch.double))
        fd.validate([inputs, grads, bias], [inputs.grad])

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

@pytest.mark.parametrize("executor", DEFAULT_EXECUTORS)
@pytest.mark.parametrize("size", generate_input_sizes(dims=2))
@pytest.mark.parametrize("dtype", FLOAT_DTYPES)
@pytest.mark.pointwise
def test_gelu_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)
    bias = torch.ones(size[-1], device="cuda", dtype=dtype)
    grads = torch.randn(size, device="cuda", dtype=dtype)

    fwd_fn = with_executor(executor, gelu)
    fwd_inputs = [inputs, bias]
    outputs = fwd_fn(fwd_inputs)

    run_benchmark(
        benchmark,
        unary_bwd_torch,
        [outputs, grads, *fwd_inputs],
        iobytes=gelu_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, 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 softmax

def softmax_fwd_fusion(
    fd: FusionDefinition, dtype: DataType, reduction_axis: int
) -> None:
    T0 = fd.define_tensor(
        shape=[-1, -1],
        contiguity=[True, True],
        dtype=dtype,
        is_cpu=False,
    )
    if dtype in PROMOTE_DTYPES:
        T0 = fd.ops.cast(T0, dtype=DataType.Float)
    T2 = fd.ops.max(T0, dims=[reduction_axis], keepdim=False, dtype=DataType.Null)

    bcast_dims = [False, False]
    bcast_dims[reduction_axis] = True

    T7 = fd.ops.broadcast(T2, is_broadcast_dim=bcast_dims)
    T13 = fd.ops.sub(T0, T7)
    T14 = fd.ops.exp(T13)
    T15 = fd.ops.sum(T14, dims=[reduction_axis], keepdim=False, dtype=DataType.Null)

    T20 = fd.ops.broadcast(T15, is_broadcast_dim=bcast_dims)

    T26 = fd.ops.reciprocal(T20)
    T27 = fd.ops.mul(T14, T26)

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

def softmax_fwd_iobytes(size: tuple, dtype: torch.dtype):
    # Total IO bytes = input + output
    return int(np.prod(size) * dtype.itemsize * 2)

@pytest.mark.parametrize("size", generate_input_sizes(dims=2))
@pytest.mark.parametrize("dtype", FLOAT_DTYPES)
@pytest.mark.parametrize(
    "reduction_axis",
    [
        pytest.param(0, marks=pytest.mark.outer_persistent),
        pytest.param(1, marks=pytest.mark.inner_persistent),
    ],
)
def test_softmax_fwd_nvf_benchmark(
    benchmark,
    size: tuple,
    dtype: torch.dtype,
    reduction_axis: int,
    disable_validation: bool,
    disable_benchmarking: bool,
):
    inputs = [torch.randn(size, device="cuda", dtype=dtype)]

    with FusionDefinition() as fd:
        softmax_fwd_fusion(fd, torch_dtype_to_nvfuser_dtype(dtype), reduction_axis)

    if not disable_validation:
        eager_output = softmax([inputs[0], reduction_axis])
        fd.validate(inputs, [eager_output])

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

@pytest.mark.parametrize("executor", DEFAULT_EXECUTORS)
@pytest.mark.parametrize("size", generate_input_sizes(dims=2))
@pytest.mark.parametrize("dtype", FLOAT_DTYPES)
@pytest.mark.parametrize(
    "reduction_axis",
    [
        pytest.param(0, marks=pytest.mark.outer_persistent),
        pytest.param(1, marks=pytest.mark.inner_persistent),
    ],
)
def test_softmax_fwd_baseline_benchmark(
    benchmark,
    size: tuple,
    dtype: torch.dtype,
    reduction_axis: int,
    executor: str,
):
    if executor == "torchcompile":
        clear_dynamo_cache()
    inputs = torch.randn(size, device="cuda", dtype=dtype, requires_grad=True)

    benchmark_fn = with_executor(executor, softmax)
    run_benchmark(
        benchmark,
        benchmark_fn,
        [inputs, reduction_axis],
        iobytes=softmax_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 softmax

def softmax_fwd_fusion(
    fd: FusionDefinition, dtype: DataType, reduction_axis: int
) -> None:
    T0 = fd.define_tensor(
        shape=[-1, -1],
        contiguity=[True, True],
        dtype=dtype,
        is_cpu=False,
    )
    if dtype in PROMOTE_DTYPES:
        T0 = fd.ops.cast(T0, dtype=DataType.Float)
    T2 = fd.ops.max(T0, dims=[reduction_axis], keepdim=False, dtype=DataType.Null)

    bcast_dims = [False, False]
    bcast_dims[reduction_axis] = True

    T7 = fd.ops.broadcast(T2, is_broadcast_dim=bcast_dims)
    T13 = fd.ops.sub(T0, T7)
    T14 = fd.ops.exp(T13)
    T15 = fd.ops.sum(T14, dims=[reduction_axis], keepdim=False, dtype=DataType.Null)

    T20 = fd.ops.broadcast(T15, is_broadcast_dim=bcast_dims)

    T26 = fd.ops.reciprocal(T20)
    T27 = fd.ops.mul(T14, T26)

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

def softmax_fwd_iobytes(size: tuple, dtype: torch.dtype):
    # Total IO bytes = input + output
    return int(np.prod(size) * dtype.itemsize * 2)

@pytest.mark.parametrize("size", generate_input_sizes(dims=2))
@pytest.mark.parametrize("dtype", FLOAT_DTYPES)
@pytest.mark.parametrize(
    "reduction_axis",
    [
        pytest.param(0, marks=pytest.mark.outer_persistent),
        pytest.param(1, marks=pytest.mark.inner_persistent),
    ],
)
def test_softmax_fwd_nvf_benchmark(
    benchmark,
    size: tuple,
    dtype: torch.dtype,


    T2 = fd.ops.max(T0, dims=[reduction_axis], keepdim=False, dtype=DataType.Null)

    bcast_dims = [False, False]
    bcast_dims[reduction_axis] = True

    T7 = fd.ops.broadcast(T2, is_broadcast_dim=bcast_dims)
    T13 = fd.ops.sub(T0, T7)
    T14 = fd.ops.exp(T13)
    T15 = fd.ops.sum(T14, dims=[reduction_axis], keepdim=False, dtype=DataType.Null)

    T20 = fd.ops.broadcast(T15, is_broadcast_dim=bcast_dims)

    T26 = fd.ops.reciprocal(T20)
    T27 = fd.ops.mul(T14, T26)

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

def softmax_fwd_iobytes(size: tuple, dtype: torch.dtype):
    # Total IO bytes = input + output
    return int(np.prod(size) * dtype.itemsize * 2)

@pytest.mark.parametrize("size", generate_input_sizes(dims=2))
@pytest.mark.parametrize("dtype", FLOAT_DTYPES)
@pytest.mark.parametrize(
    "reduction_axis",
    [
        pytest.param(0, marks=pytest.mark.outer_persistent),
        pytest.param(1, marks=pytest.mark.inner_persistent),
    ],
)
def test_softmax_fwd_nvf_benchmark(
    benchmark,
    size: tuple,
    dtype: torch.dtype,
    reduction_axis: int,
    disable_validation: bool,
    disable_benchmarking: bool,
):
    inputs = [torch.randn(size, device="cuda", dtype=dtype)]


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

def softmax_fwd_iobytes(size: tuple, dtype: torch.dtype):
    # Total IO bytes = input + output
    return int(np.prod(size) * dtype.itemsize * 2)

@pytest.mark.parametrize("size", generate_input_sizes(dims=2))
@pytest.mark.parametrize("dtype", FLOAT_DTYPES)
@pytest.mark.parametrize(
    "reduction_axis",
    [
        pytest.param(0, marks=pytest.mark.outer_persistent),
        pytest.param(1, marks=pytest.mark.inner_persistent),
    ],
)
def test_softmax_fwd_nvf_benchmark(
    benchmark,
    size: tuple,
    dtype: torch.dtype,
    reduction_axis: int,
    disable_validation: bool,
    disable_benchmarking: bool,
):
    inputs = [torch.randn(size, device="cuda", dtype=dtype)]

    with FusionDefinition() as fd:
        softmax_fwd_fusion(fd, torch_dtype_to_nvfuser_dtype(dtype), reduction_axis)

    if not disable_validation:
        eager_output = softmax([inputs[0], reduction_axis])
        fd.validate(inputs, [eager_output])

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

@pytest.mark.parametrize("executor", DEFAULT_EXECUTORS)
@pytest.mark.parametrize("size", generate_input_sizes(dims=2))
@pytest.mark.parametrize("dtype", FLOAT_DTYPES)
@pytest.mark.parametrize(
    "reduction_axis",
    [
        pytest.param(0, marks=pytest.mark.outer_persistent),
        pytest.param(1, marks=pytest.mark.inner_persistent),
    ],
)
def test_softmax_fwd_baseline_benchmark(
    benchmark,
    size: tuple,
    dtype: torch.dtype,
    reduction_axis: int,
    executor: str,
):
    if executor == "torchcompile":


        softmax_fwd_fusion(fd, torch_dtype_to_nvfuser_dtype(dtype), reduction_axis)

    if not disable_validation:
        eager_output = softmax([inputs[0], reduction_axis])
        fd.validate(inputs, [eager_output])

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

@pytest.mark.parametrize("executor", DEFAULT_EXECUTORS)
@pytest.mark.parametrize("size", generate_input_sizes(dims=2))
@pytest.mark.parametrize("dtype", FLOAT_DTYPES)
@pytest.mark.parametrize(
    "reduction_axis",
    [
        pytest.param(0, marks=pytest.mark.outer_persistent),
        pytest.param(1, marks=pytest.mark.inner_persistent),
    ],
)
def test_softmax_fwd_baseline_benchmark(
    benchmark,
    size: tuple,
    dtype: torch.dtype,
    reduction_axis: int,
    executor: str,
):
    if executor == "torchcompile":
        clear_dynamo_cache()
    inputs = torch.randn(size, device="cuda", dtype=dtype, requires_grad=True)

    benchmark_fn = with_executor(executor, softmax)
    run_benchmark(
        benchmark,
        benchmark_fn,
        [inputs, reduction_axis],
        iobytes=softmax_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
import torch
from .global_params import generate_input_sizes, FLOAT_DTYPES, PROMOTE_DTYPES
import numpy as np

def gelu_bwd_reduction_fusion(
    fd: FusionDefinition, dtype: DataType, reduction_axis: int
) -> None:
    input = fd.define_tensor(
        shape=[-1, -1], contiguity=[True, True], dtype=dtype, is_cpu=False
    )
    grad = fd.define_tensor(
        shape=[-1, -1], contiguity=[True, True], dtype=dtype, is_cpu=False
    )
    bias = fd.define_tensor(shape=[-1], contiguity=[True], dtype=dtype, is_cpu=False)
    if dtype in PROMOTE_DTYPES:
        input = fd.ops.cast(input, dtype=DataType.Float)
        grad = fd.ops.cast(grad, dtype=DataType.Float)
        bias = fd.ops.cast(bias, dtype=DataType.Float)
    S_079 = fd.define_scalar(0.79788456)
    S_004 = fd.define_scalar(0.044715)
    S_010 = fd.define_scalar(0.1070322243)
    bias = fd.ops.broadcast_in_dim(bias, shape=[1, input.size(-1)], broadcast_dims=[1])
    T1 = fd.ops.add(input, bias)
    T2 = fd.ops.mul(T1, S_079)
    T3 = fd.ops.mul(T1, S_004)
    T4 = fd.ops.mul(T3, T1)
    S1 = fd.define_scalar(1.0)
    T5 = fd.ops.add(S1, T4)
    T6 = fd.ops.mul(T2, T5)
    T7 = fd.ops.tanh(T6)
    S2 = fd.define_scalar(0.50)
    T8 = fd.ops.mul(T1, S2)
    T9 = fd.ops.mul(T7, T7)
    T10 = fd.ops.neg(T9)
    T11 = fd.ops.add(T10, S1)
    T12 = fd.ops.mul(T1, S_010)
    T13 = fd.ops.mul(T12, T1)
    T14 = fd.ops.add(T13, S_079)
    T15 = fd.ops.mul(T11, T14)
    T16 = fd.ops.mul(T8, T15)
    T17 = fd.ops.add(T7, S1)
    T18 = fd.ops.mul(T17, S2)
    T19 = fd.ops.add(T16, T18)
    T20 = fd.ops.mul(grad, T19)
    T21 = fd.ops.sum(T20, dims=[reduction_axis], keepdim=False)
    if dtype in PROMOTE_DTYPES:
        T21 = fd.ops.cast(T21, dtype=dtype)
    fd.add_output(T21)

def gelu_bwd_reduction_torch(
    inputs: list,
):  # [output, grad_out, in_tensor, reduction_axis]
    eager_output, grad_out, in_tensor, reduction_axis = inputs
    eager_output.backward(grad_out, retain_graph=True)
    return in_tensor.grad.sum(reduction_axis)

def gelu_bwd_reduction_iobytes(size: tuple, dtype: torch.dtype, reduction_axis: int):
    # Total IO bytes = in_tensor + grad_out + bias + grad_input
    return int(
        dtype.itemsize * (2 * np.prod(size) + size[1] + size[1 - reduction_axis])
    )

@pytest.mark.parametrize("size", generate_input_sizes(dims=2))
@pytest.mark.parametrize("dtype", FLOAT_DTYPES)
@pytest.mark.parametrize("reduction_axis", [0, 1])
@pytest.mark.reduction
def test_gelu_bwd_reduction_nvf_benchmark(
    benchmark,
    size: tuple,
    dtype: torch.dtype,
    reduction_axis: int,
    disable_validation: bool,
    disable_benchmarking: bool,
):
    inputs = torch.randn(*size, device="cuda", dtype=dtype, requires_grad=True)
    grads = torch.randn(*size, device="cuda", dtype=dtype)
    bias = torch.ones(size[-1], device="cuda", dtype=dtype)
    with FusionDefinition() as fd:
        gelu_bwd_reduction_fusion(
            fd, torch_dtype_to_nvfuser_dtype(dtype), reduction_axis
        )

    if not disable_validation:
        eager_output = torch.nn.functional.gelu(
            inputs.to(torch.double) + bias.to(torch.double), approximate="tanh"
        )
        eager_output.backward(grads.to(torch.double))
        reduction_out = inputs.grad.to(torch.double).sum(reduction_axis)
        fd.validate([inputs, grads, bias], [reduction_out.to(dtype)])

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

@pytest.mark.parametrize("executor", ["eager", "torchcompile"])
@pytest.mark.parametrize("size", generate_input_sizes(dims=2))
@pytest.mark.parametrize("dtype", FLOAT_DTYPES)
@pytest.mark.parametrize("reduction_axis", [0, 1])
@pytest.mark.reduction
def test_gelu_bwd_reduction_baseline_benchmark(
    benchmark,
    size: tuple,
    dtype: torch.dtype,
    reduction_axis: int,
    executor: str,
):
    if executor == "torchcompile":
        clear_dynamo_cache()
    inputs = torch.randn(size, device="cuda", dtype=dtype, requires_grad=True)
    bias = torch.ones(size[-1], device="cuda", dtype=dtype)
    grads = torch.randn(size, device="cuda", dtype=dtype)
    eager_output = torch.nn.functional.gelu(inputs + bias, approximate="tanh")

    benchmark_fn = with_executor(executor, gelu_bwd_reduction_torch)

    run_benchmark(
        benchmark,
        benchmark_fn,
        [eager_output, grads, inputs, reduction_axis],
        iobytes=gelu_bwd_reduction_iobytes(size, dtype, reduction_axis),
    )


# 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
import torch
from .global_params import generate_input_sizes, FLOAT_DTYPES, PROMOTE_DTYPES
import numpy as np

def gelu_bwd_reduction_fusion(
    fd: FusionDefinition, dtype: DataType, reduction_axis: int
) -> None:
    input = fd.define_tensor(
        shape=[-1, -1], contiguity=[True, True], dtype=dtype, is_cpu=False
    )
    grad = fd.define_tensor(
        shape=[-1, -1], contiguity=[True, True], dtype=dtype, is_cpu=False
    )
    bias = fd.define_tensor(shape=[-1], contiguity=[True], dtype=dtype, is_cpu=False)
    if dtype in PROMOTE_DTYPES:
        input = fd.ops.cast(input, dtype=DataType.Float)
        grad = fd.ops.cast(grad, dtype=DataType.Float)
        bias = fd.ops.cast(bias, dtype=DataType.Float)
    S_079 = fd.define_scalar(0.79788456)
    S_004 = fd.define_scalar(0.044715)
    S_010 = fd.define_scalar(0.1070322243)
    bias = fd.ops.broadcast_in_dim(bias, shape=[1, input.size(-1)], broadcast_dims=[1])
    T1 = fd.ops.add(input, bias)
    T2 = fd.ops.mul(T1, S_079)
    T3 = fd.ops.mul(T1, S_004)
    T4 = fd.ops.mul(T3, T1)
    S1 = fd.define_scalar(1.0)
    T5 = fd.ops.add(S1, T4)
    T6 = fd.ops.mul(T2, T5)
    T7 = fd.ops.tanh(T6)
    S2 = fd.define_scalar(0.50)
    T8 = fd.ops.mul(T1, S2)
    T9 = fd.ops.mul(T7, T7)
    T10 = fd.ops.neg(T9)
    T11 = fd.ops.add(T10, S1)
    T12 = fd.ops.mul(T1, S_010)
    T13 = fd.ops.mul(T12, T1)
    T14 = fd.ops.add(T13, S_079)
    T15 = fd.ops.mul(T11, T14)
    T16 = fd.ops.mul(T8, T15)
    T17 = fd.ops.add(T7, S1)
    T18 = fd.ops.mul(T17, S2)
    T19 = fd.ops.add(T16, T18)
    T20 = fd.ops.mul(grad, T19)
    T21 = fd.ops.sum(T20, dims=[reduction_axis], keepdim=False)
    if dtype in PROMOTE_DTYPES:
        T21 = fd.ops.cast(T21, dtype=dtype)
    fd.add_output(T21)

def gelu_bwd_reduction_torch(
    inputs: list,
):  # [output, grad_out, in_tensor, reduction_axis]
    eager_output, grad_out, in_tensor, reduction_axis = inputs
    eager_output.backward(grad_out, retain_graph=True)
    return in_tensor.grad.sum(reduction_axis)

def gelu_bwd_reduction_iobytes(size: tuple, dtype: torch.dtype, reduction_axis: int):
    # Total IO bytes = in_tensor + grad_out + bias + grad_input
    return int(
        dtype.itemsize * (2 * np.prod(size) + size[1] + size[1 - reduction_axis])
    )

@pytest.mark.parametrize("size", generate_input_sizes(dims=2))
@pytest.mark.parametrize("dtype", FLOAT_DTYPES)
@pytest.mark.parametrize("reduction_axis", [0, 1])


    S2 = fd.define_scalar(0.50)
    T8 = fd.ops.mul(T1, S2)
    T9 = fd.ops.mul(T7, T7)
    T10 = fd.ops.neg(T9)
    T11 = fd.ops.add(T10, S1)
    T12 = fd.ops.mul(T1, S_010)
    T13 = fd.ops.mul(T12, T1)
    T14 = fd.ops.add(T13, S_079)
    T15 = fd.ops.mul(T11, T14)
    T16 = fd.ops.mul(T8, T15)
    T17 = fd.ops.add(T7, S1)
    T18 = fd.ops.mul(T17, S2)
    T19 = fd.ops.add(T16, T18)
    T20 = fd.ops.mul(grad, T19)
    T21 = fd.ops.sum(T20, dims=[reduction_axis], keepdim=False)
    if dtype in PROMOTE_DTYPES:
        T21 = fd.ops.cast(T21, dtype=dtype)
    fd.add_output(T21)

def gelu_bwd_reduction_torch(
    inputs: list,
):  # [output, grad_out, in_tensor, reduction_axis]
    eager_output, grad_out, in_tensor, reduction_axis = inputs
    eager_output.backward(grad_out, retain_graph=True)
    return in_tensor.grad.sum(reduction_axis)

def gelu_bwd_reduction_iobytes(size: tuple, dtype: torch.dtype, reduction_axis: int):
    # Total IO bytes = in_tensor + grad_out + bias + grad_input
    return int(
        dtype.itemsize * (2 * np.prod(size) + size[1] + size[1 - reduction_axis])
    )

@pytest.mark.parametrize("size", generate_input_sizes(dims=2))
@pytest.mark.parametrize("dtype", FLOAT_DTYPES)
@pytest.mark.parametrize("reduction_axis", [0, 1])
@pytest.mark.reduction
def test_gelu_bwd_reduction_nvf_benchmark(
    benchmark,
    size: tuple,
    dtype: torch.dtype,
    reduction_axis: int,
    disable_validation: bool,
    disable_benchmarking: bool,


    T15 = fd.ops.mul(T11, T14)
    T16 = fd.ops.mul(T8, T15)
    T17 = fd.ops.add(T7, S1)
    T18 = fd.ops.mul(T17, S2)
    T19 = fd.ops.add(T16, T18)
    T20 = fd.ops.mul(grad, T19)
    T21 = fd.ops.sum(T20, dims=[reduction_axis], keepdim=False)
    if dtype in PROMOTE_DTYPES:
        T21 = fd.ops.cast(T21, dtype=dtype)
    fd.add_output(T21)

def gelu_bwd_reduction_torch(
    inputs: list,
):  # [output, grad_out, in_tensor, reduction_axis]
    eager_output, grad_out, in_tensor, reduction_axis = inputs
    eager_output.backward(grad_out, retain_graph=True)
    return in_tensor.grad.sum(reduction_axis)

def gelu_bwd_reduction_iobytes(size: tuple, dtype: torch.dtype, reduction_axis: int):
    # Total IO bytes = in_tensor + grad_out + bias + grad_input
    return int(
        dtype.itemsize * (2 * np.prod(size) + size[1] + size[1 - reduction_axis])
    )

@pytest.mark.parametrize("size", generate_input_sizes(dims=2))
@pytest.mark.parametrize("dtype", FLOAT_DTYPES)
@pytest.mark.parametrize("reduction_axis", [0, 1])
@pytest.mark.reduction
def test_gelu_bwd_reduction_nvf_benchmark(
    benchmark,
    size: tuple,
    dtype: torch.dtype,
    reduction_axis: int,
    disable_validation: bool,
    disable_benchmarking: bool,
):
    inputs = torch.randn(*size, device="cuda", dtype=dtype, requires_grad=True)
    grads = torch.randn(*size, device="cuda", dtype=dtype)
    bias = torch.ones(size[-1], device="cuda", dtype=dtype)
    with FusionDefinition() as fd:
        gelu_bwd_reduction_fusion(
            fd, torch_dtype_to_nvfuser_dtype(dtype), reduction_axis


def gelu_bwd_reduction_torch(
    inputs: list,
):  # [output, grad_out, in_tensor, reduction_axis]
    eager_output, grad_out, in_tensor, reduction_axis = inputs
    eager_output.backward(grad_out, retain_graph=True)
    return in_tensor.grad.sum(reduction_axis)

def gelu_bwd_reduction_iobytes(size: tuple, dtype: torch.dtype, reduction_axis: int):
    # Total IO bytes = in_tensor + grad_out + bias + grad_input
    return int(
        dtype.itemsize * (2 * np.prod(size) + size[1] + size[1 - reduction_axis])
    )

@pytest.mark.parametrize("size", generate_input_sizes(dims=2))
@pytest.mark.parametrize("dtype", FLOAT_DTYPES)
@pytest.mark.parametrize("reduction_axis", [0, 1])
@pytest.mark.reduction
def test_gelu_bwd_reduction_nvf_benchmark(
    benchmark,
    size: tuple,
    dtype: torch.dtype,
    reduction_axis: int,
    disable_validation: bool,
    disable_benchmarking: bool,
):
    inputs = torch.randn(*size, device="cuda", dtype=dtype, requires_grad=True)
    grads = torch.randn(*size, device="cuda", dtype=dtype)
    bias = torch.ones(size[-1], device="cuda", dtype=dtype)
    with FusionDefinition() as fd:
        gelu_bwd_reduction_fusion(
            fd, torch_dtype_to_nvfuser_dtype(dtype), reduction_axis
        )

    if not disable_validation:
        eager_output = torch.nn.functional.gelu(
            inputs.to(torch.double) + bias.to(torch.double), approximate="tanh"
        )
        eager_output.backward(grads.to(torch.double))
        reduction_out = inputs.grad.to(torch.double).sum(reduction_axis)
        fd.validate([inputs, grads, bias], [reduction_out.to(dtype)])

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

@pytest.mark.parametrize("executor", ["eager", "torchcompile"])
@pytest.mark.parametrize("size", generate_input_sizes(dims=2))
@pytest.mark.parametrize("dtype", FLOAT_DTYPES)
@pytest.mark.parametrize("reduction_axis", [0, 1])
@pytest.mark.reduction
def test_gelu_bwd_reduction_baseline_benchmark(
    benchmark,
    size: tuple,
    dtype: torch.dtype,
    reduction_axis: int,
    executor: str,
):
    if executor == "torchcompile":
        clear_dynamo_cache()
    inputs = torch.randn(size, device="cuda", dtype=dtype, requires_grad=True)
    bias = torch.ones(size[-1], device="cuda", dtype=dtype)
    grads = torch.randn(size, device="cuda", dtype=dtype)
    eager_output = torch.nn.functional.gelu(inputs + bias, approximate="tanh")


            fd, torch_dtype_to_nvfuser_dtype(dtype), reduction_axis
        )

    if not disable_validation:
        eager_output = torch.nn.functional.gelu(
            inputs.to(torch.double) + bias.to(torch.double), approximate="tanh"
        )
        eager_output.backward(grads.to(torch.double))
        reduction_out = inputs.grad.to(torch.double).sum(reduction_axis)
        fd.validate([inputs, grads, bias], [reduction_out.to(dtype)])

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

@pytest.mark.parametrize("executor", ["eager", "torchcompile"])
@pytest.mark.parametrize("size", generate_input_sizes(dims=2))
@pytest.mark.parametrize("dtype", FLOAT_DTYPES)
@pytest.mark.parametrize("reduction_axis", [0, 1])
@pytest.mark.reduction
def test_gelu_bwd_reduction_baseline_benchmark(
    benchmark,
    size: tuple,
    dtype: torch.dtype,
    reduction_axis: int,
    executor: str,
):
    if executor == "torchcompile":
        clear_dynamo_cache()
    inputs = torch.randn(size, device="cuda", dtype=dtype, requires_grad=True)
    bias = torch.ones(size[-1], device="cuda", dtype=dtype)
    grads = torch.randn(size, device="cuda", dtype=dtype)
    eager_output = torch.nn.functional.gelu(inputs + bias, approximate="tanh")

    benchmark_fn = with_executor(executor, gelu_bwd_reduction_torch)

    run_benchmark(
        benchmark,
        benchmark_fn,
        [eager_output, grads, inputs, reduction_axis],
        iobytes=gelu_bwd_reduction_iobytes(size, dtype, reduction_axis),
    )


# SPDX-FileCopyrightText: Copyright (c) 2024-present NVIDIA CORPORATION & AFFILIATES.
# All rights reserved.
# SPDX-License-Identifier: BSD-3-Clause
from functools import partial

def llama_hf_cfg(config_str):
    class Config:
        def __init__(
            self, n_head, head_size, n_query_groups, rope_n_elem, batches, seq_length
        ):
            self.n_head = n_head
            self.head_size = head_size
            self.n_query_groups = n_query_groups
            self.rope_n_elem = rope_n_elem
            self.batches = batches
            self.seq_length = seq_length

    configs = {}
    configs["llama_2_7b_hf"] = Config(
        n_head=32,
        head_size=128,
        n_query_groups=32,
        rope_n_elem=128,
        batches=2,
        seq_length=4096,
    )
    configs["llama_3_8B"] = Config(
        n_head=32,
        head_size=128,
        n_query_groups=8,
        rope_n_elem=128,
        batches=2,
        seq_length=8192,
    )

    return configs[config_str]

def hf_qwen2_cfg():
    from transformers import AutoConfig

    config = AutoConfig.from_pretrained("Qwen/Qwen2.5-7B-Instruct")
    config.batch_size = 1
    config.seq_len = 4096
    config._attn_implementation = "sdpa"
    return config

def hf_phi3_cfg():
    from transformers import AutoConfig

    config = AutoConfig.from_pretrained("microsoft/Phi-3.5-mini-instruct")
    config.batch_size = 1
    config.seq_len = 8192
    config._attn_implementation = "sdpa"
    return config

def hf_mistral_nemo_cfg():
    import json
    from transformers.models.mistral import MistralConfig

    mistral_cfg_str = r"""{
      "_name_or_path": "mistralai/Mistral-Nemo-Base-2407",
      "architectures": [
        "MistralForCausalLM"
      ],
      "attention_dropout": 0.0,
      "bos_token_id": 1,
      "eos_token_id": 2,
      "head_dim": 128,
      "hidden_act": "silu",
      "hidden_size": 5120,
      "initializer_range": 0.02,
      "intermediate_size": 14336,
      "max_position_embeddings": 128000,
      "model_type": "mistral",
      "num_attention_heads": 32,
      "num_hidden_layers": 40,
      "num_key_value_heads": 8,
      "rms_norm_eps": 1e-05,
      "rope_theta": 1000000.0,
      "sliding_window": null,
      "tie_word_embeddings": false,
      "torch_dtype": "bfloat16",
      "transformers_version": "4.43.3",
      "use_cache": true,
      "vocab_size": 131072
    }
    """

    cfg = MistralConfig.from_dict(json.loads(mistral_cfg_str))
    cfg.batch_size = 1
    cfg.seq_len = 4096
    cfg._attn_implementation = "sdpa"

    return cfg

def litgpt_cfg(model_name):
    import litgpt

    cfg = litgpt.Config.from_name(model_name)
    cfg.batch_size = 1
    cfg.seq_len = 4096
    cfg.name_or_path = model_name

    return cfg

configs = {
    "llama_2_7b_hf": partial(llama_hf_cfg, config_str="llama_2_7b_hf"),
    "llama_3_8B": partial(llama_hf_cfg, config_str="llama_3_8B"),
    "hf_qwen2": hf_qwen2_cfg,
    "hf_phi3": hf_phi3_cfg,
    "hf_mistral_nemo": hf_mistral_nemo_cfg,
    "litgpt": litgpt_cfg,
}


# SPDX-FileCopyrightText: Copyright (c) 2024-present NVIDIA CORPORATION & AFFILIATES.
# All rights reserved.
# SPDX-License-Identifier: BSD-3-Clause
from functools import partial

def llama_hf_cfg(config_str):
    class Config:
        def __init__(
            self, n_head, head_size, n_query_groups, rope_n_elem, batches, seq_length
        ):
            self.n_head = n_head
            self.head_size = head_size
            self.n_query_groups = n_query_groups
            self.rope_n_elem = rope_n_elem
            self.batches = batches
            self.seq_length = seq_length

    configs = {}
    configs["llama_2_7b_hf"] = Config(
        n_head=32,
        head_size=128,
        n_query_groups=32,
        rope_n_elem=128,
        batches=2,
        seq_length=4096,
    )
    configs["llama_3_8B"] = Config(
        n_head=32,
        head_size=128,
        n_query_groups=8,
        rope_n_elem=128,
        batches=2,
        seq_length=8192,
    )

    return configs[config_str]

def hf_qwen2_cfg():
    from transformers import AutoConfig

    config = AutoConfig.from_pretrained("Qwen/Qwen2.5-7B-Instruct")
    config.batch_size = 1
    config.seq_len = 4096
    config._attn_implementation = "sdpa"
    return config

def hf_phi3_cfg():
    from transformers import AutoConfig

    config = AutoConfig.from_pretrained("microsoft/Phi-3.5-mini-instruct")
    config.batch_size = 1
    config.seq_len = 8192
    config._attn_implementation = "sdpa"
    return config


    configs["llama_2_7b_hf"] = Config(
        n_head=32,
        head_size=128,
        n_query_groups=32,
        rope_n_elem=128,
        batches=2,
        seq_length=4096,
    )
    configs["llama_3_8B"] = Config(
        n_head=32,
        head_size=128,
        n_query_groups=8,
        rope_n_elem=128,
        batches=2,
        seq_length=8192,
    )

    return configs[config_str]

def hf_qwen2_cfg():
    from transformers import AutoConfig

    config = AutoConfig.from_pretrained("Qwen/Qwen2.5-7B-Instruct")
    config.batch_size = 1
    config.seq_len = 4096
    config._attn_implementation = "sdpa"
    return config

def hf_phi3_cfg():
    from transformers import AutoConfig

    config = AutoConfig.from_pretrained("microsoft/Phi-3.5-mini-instruct")
    config.batch_size = 1
    config.seq_len = 8192
    config._attn_implementation = "sdpa"
    return config

def hf_mistral_nemo_cfg():
    import json
    from transformers.models.mistral import MistralConfig

    mistral_cfg_str = r"""{
      "_name_or_path": "mistralai/Mistral-Nemo-Base-2407",
      "architectures": [
        "MistralForCausalLM"


        head_size=128,
        n_query_groups=8,
        rope_n_elem=128,
        batches=2,
        seq_length=8192,
    )

    return configs[config_str]

def hf_qwen2_cfg():
    from transformers import AutoConfig

    config = AutoConfig.from_pretrained("Qwen/Qwen2.5-7B-Instruct")
    config.batch_size = 1
    config.seq_len = 4096
    config._attn_implementation = "sdpa"
    return config

def hf_phi3_cfg():
    from transformers import AutoConfig

    config = AutoConfig.from_pretrained("microsoft/Phi-3.5-mini-instruct")
    config.batch_size = 1
    config.seq_len = 8192
    config._attn_implementation = "sdpa"
    return config

def hf_mistral_nemo_cfg():
    import json
    from transformers.models.mistral import MistralConfig

    mistral_cfg_str = r"""{
      "_name_or_path": "mistralai/Mistral-Nemo-Base-2407",
      "architectures": [
        "MistralForCausalLM"
      ],
      "attention_dropout": 0.0,
      "bos_token_id": 1,
      "eos_token_id": 2,
      "head_dim": 128,
      "hidden_act": "silu",
      "hidden_size": 5120,
      "initializer_range": 0.02,
      "intermediate_size": 14336,
      "max_position_embeddings": 128000,


def hf_qwen2_cfg():
    from transformers import AutoConfig

    config = AutoConfig.from_pretrained("Qwen/Qwen2.5-7B-Instruct")
    config.batch_size = 1
    config.seq_len = 4096
    config._attn_implementation = "sdpa"
    return config

def hf_phi3_cfg():
    from transformers import AutoConfig

    config = AutoConfig.from_pretrained("microsoft/Phi-3.5-mini-instruct")
    config.batch_size = 1
    config.seq_len = 8192
    config._attn_implementation = "sdpa"
    return config

def hf_mistral_nemo_cfg():
    import json
    from transformers.models.mistral import MistralConfig

    mistral_cfg_str = r"""{
      "_name_or_path": "mistralai/Mistral-Nemo-Base-2407",
      "architectures": [
        "MistralForCausalLM"
      ],
      "attention_dropout": 0.0,
      "bos_token_id": 1,
      "eos_token_id": 2,
      "head_dim": 128,
      "hidden_act": "silu",
      "hidden_size": 5120,
      "initializer_range": 0.02,
      "intermediate_size": 14336,
      "max_position_embeddings": 128000,
      "model_type": "mistral",
      "num_attention_heads": 32,
      "num_hidden_layers": 40,
      "num_key_value_heads": 8,
      "rms_norm_eps": 1e-05,
      "rope_theta": 1000000.0,
      "sliding_window": null,
      "tie_word_embeddings": false,
      "torch_dtype": "bfloat16",
      "transformers_version": "4.43.3",
      "use_cache": true,
      "vocab_size": 131072
    }
    """

    cfg = MistralConfig.from_dict(json.loads(mistral_cfg_str))
    cfg.batch_size = 1
    cfg.seq_len = 4096
    cfg._attn_implementation = "sdpa"

    return cfg

def litgpt_cfg(model_name):
    import litgpt

    cfg = litgpt.Config.from_name(model_name)
    cfg.batch_size = 1
    cfg.seq_len = 4096
    cfg.name_or_path = model_name

    return cfg

configs = {
    "llama_2_7b_hf": partial(llama_hf_cfg, config_str="llama_2_7b_hf"),
    "llama_3_8B": partial(llama_hf_cfg, config_str="llama_3_8B"),
    "hf_qwen2": hf_qwen2_cfg,
    "hf_phi3": hf_phi3_cfg,
    "hf_mistral_nemo": hf_mistral_nemo_cfg,
    "litgpt": litgpt_cfg,


      "num_key_value_heads": 8,
      "rms_norm_eps": 1e-05,
      "rope_theta": 1000000.0,
      "sliding_window": null,
      "tie_word_embeddings": false,
      "torch_dtype": "bfloat16",
      "transformers_version": "4.43.3",
      "use_cache": true,
      "vocab_size": 131072
    }
    """

    cfg = MistralConfig.from_dict(json.loads(mistral_cfg_str))
    cfg.batch_size = 1
    cfg.seq_len = 4096
    cfg._attn_implementation = "sdpa"

    return cfg

def litgpt_cfg(model_name):
    import litgpt

    cfg = litgpt.Config.from_name(model_name)
    cfg.batch_size = 1
    cfg.seq_len = 4096
    cfg.name_or_path = model_name

    return cfg

configs = {
    "llama_2_7b_hf": partial(llama_hf_cfg, config_str="llama_2_7b_hf"),
    "llama_3_8B": partial(llama_hf_cfg, config_str="llama_3_8B"),
    "hf_qwen2": hf_qwen2_cfg,
    "hf_phi3": hf_phi3_cfg,
    "hf_mistral_nemo": hf_mistral_nemo_cfg,
    "litgpt": litgpt_cfg,
}


Config

# SPDX-FileCopyrightText: Copyright (c) 2024-present NVIDIA CORPORATION & AFFILIATES.
# All rights reserved.
# SPDX-License-Identifier: BSD-3-Clause
from functools import partial

def llama_hf_cfg(config_str):
    class Config:
        def __init__(
            self, n_head, head_size, n_query_groups, rope_n_elem, batches, seq_length
        ):
            self.n_head = n_head
            self.head_size = head_size
            self.n_query_groups = n_query_groups
            self.rope_n_elem = rope_n_elem
            self.batches = batches
            self.seq_length = seq_length

    configs = {}
    configs["llama_2_7b_hf"] = Config(
        n_head=32,
        head_size=128,
        n_query_groups=32,
        rope_n_elem=128,
        batches=2,
        seq_length=4096,
    )
    configs["llama_3_8B"] = Config(
        n_head=32,
        head_size=128,
        n_query_groups=8,
        rope_n_elem=128,
        batches=2,
        seq_length=8192,
    )

    return configs[config_str]


# 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 layernorm

def layernorm_bwd_fusion(
    fd: FusionDefinition,
    dtype: DataType,
) -> None:
    T0 = fd.define_tensor(
        shape=[-1, -1], contiguity=[True, True], dtype=dtype, is_cpu=False
    )
    T1 = fd.define_tensor(
        shape=[-1, -1], contiguity=[True, True], dtype=dtype, is_cpu=False
    )

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

    T4 = 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)
        T4 = fd.ops.cast(T4, dtype=DataType.Float)

    T9 = fd.ops.broadcast_in_dim(T2, shape=[T0.size(0), 1], broadcast_dims=[0])
    V12 = T0.shape()
    T13 = fd.ops.broadcast_in_dim(T9, shape=V12, broadcast_dims=[0, 1])
    T14 = fd.ops.sub(T0, T13)

    T18 = fd.ops.broadcast_in_dim(T3, shape=V12, broadcast_dims=[0, 1])
    T19 = fd.ops.mul(T14, T18)

    T23 = fd.ops.broadcast_in_dim(T4, shape=V12, broadcast_dims=[1])
    T28 = fd.ops.sum(T1, dims=[0], keepdim=False, dtype=DataType.Null)

    T30 = fd.ops.mul(T1, T23)
    T31 = fd.ops.mul(T1, T19)
    T32 = fd.ops.sum(T31, dims=[0], keepdim=False, dtype=DataType.Null)

    T34 = fd.ops.mul(T30, T18)
    T35 = fd.ops.mul(T30, T14)
    T36 = fd.ops.sum(T35, dims=[1], keepdim=False, dtype=DataType.Null)

    T40 = fd.ops.broadcast_in_dim(T36, shape=[T0.size(0), 1], broadcast_dims=[0])
    T41 = fd.ops.neg(T34)
    T42 = fd.ops.sum(T41, dims=[1], keepdim=False, dtype=DataType.Null)
    T46 = fd.ops.broadcast_in_dim(T42, shape=[T0.size(0), 1], broadcast_dims=[0])
    S47 = fd.define_scalar(-0.500000, dtype=DataType.Double)
    T48 = fd.ops.mul(S47, T40)
    S49 = fd.define_scalar(3.00000, dtype=DataType.Double)
    T50 = fd.ops.pow(T3, S49)
    T51 = fd.ops.mul(T48, T50)
    T54 = fd.ops.sum(T46, dims=[1], keepdim=False, dtype=DataType.Null)
    T55 = fd.ops.sum(T51, dims=[1], keepdim=False, dtype=DataType.Null)

    T59 = fd.ops.broadcast_in_dim(T55, shape=[T0.size(0), 1], broadcast_dims=[0])
    T63 = fd.ops.broadcast_in_dim(T59, shape=V12, broadcast_dims=[0, 1])
    T67 = fd.ops.broadcast_in_dim(T2, shape=[T0.size(0), 1], broadcast_dims=[0])
    T71 = fd.ops.broadcast_in_dim(T67, shape=V12, broadcast_dims=[0, 1])

    S72 = fd.define_scalar(2.00000, dtype=DataType.Double)
    T73 = fd.ops.mul(S72, T63)
    T74 = fd.ops.sub(T0, T71)
    T75 = fd.ops.mul(T73, T74)

    S77 = fd.ops.reciprocal(T0.size(1))
    T78 = fd.ops.mul(T75, S77)
    T82 = fd.ops.broadcast_in_dim(T54, shape=[T0.size(0), 1], broadcast_dims=[0])
    T86 = fd.ops.broadcast_in_dim(T82, shape=V12, broadcast_dims=[0, 1])
    T88 = fd.ops.mul(S77, T86)
    T89 = fd.ops.add(T78, T88)
    T90 = fd.ops.add(T34, T89)

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

    fd.add_output(T90)
    fd.add_output(T32)
    fd.add_output(T28)

def layernorm_bwd_iobytes(size: tuple, dtype: torch.dtype):
    # Manual IOBytes computation since nvfuser input/outputs (in_tensor, grad_out, mean, invstd, weigts) differ from baselines (out, grad_out)
    # Total IO bytes = in_tensor (size, dtype) + grad_out (size, dtype) + mean (size[0], float) +
    #       invstd (size[0], float) + weights (size[1], dtype) +
    #       grad_in (size, dtype) + grad_weights (size[1], dtype)+ grad_bias (size[1], dtype)
    return int(
        dtype.itemsize * 3 * (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_outer_persistent
@pytest.mark.inner_persistent
def test_layernorm_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)
    bias = torch.randn(size[1], device="cuda", dtype=dtype, requires_grad=True)

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

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

    if not disable_validation:
        eager_output = torch.nn.functional.layer_norm(
            inputs.to(torch.double),
            inputs.shape[1:],
            weight=weights.to(torch.double),
            bias=bias.to(torch.double),
        )
        eager_output.backward(grads.to(torch.double))
        fd.validate(
            [inputs, grads, mean, invstd, weights],
            [inputs.grad, weights.grad, bias.grad],
        )

    if not disable_benchmarking:
        run_benchmark(benchmark, fd.execute, [inputs, grads, mean, invstd, 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_layernorm_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)
    bias = torch.randn(size[1], device="cuda", dtype=dtype, requires_grad=True)

    fwd_fn = with_executor(executor, layernorm)
    fwd_inputs = [inputs, weights, bias]
    outputs = fwd_fn(fwd_inputs)

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


# SPDX-FileCopyrightText: Copyright (c) 2024-present NVIDIA CORPORATION & AFFILIATES.
# All rights reserved.
# SPDX-License-Identifier: BSD-3-Clause
import pytest
from nvfuser_direct import FusionDefinition, DataType
from nvfuser_direct.pytorch_utils import torch_dtype_to_nvfuser_dtype
from .core import (
    run_benchmark,
    clear_dynamo_cache,
    unary_bwd_torch,
    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_bwd_fusion(
    fd: FusionDefinition,
    dtype: DataType,
) -> None:
    T0 = fd.define_tensor(
        shape=[-1, -1], contiguity=[True, True], dtype=dtype, is_cpu=False
    )
    T1 = fd.define_tensor(
        shape=[-1, -1], contiguity=[True, True], dtype=dtype, is_cpu=False
    )

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

    T4 = 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)
        T4 = fd.ops.cast(T4, dtype=DataType.Float)

    T9 = fd.ops.broadcast_in_dim(T2, shape=[T0.size(0), 1], broadcast_dims=[0])
    V12 = T0.shape()
    T13 = fd.ops.broadcast_in_dim(T9, shape=V12, broadcast_dims=[0, 1])
    T14 = fd.ops.sub(T0, T13)

    T18 = fd.ops.broadcast_in_dim(T3, shape=V12, broadcast_dims=[0, 1])
    T19 = fd.ops.mul(T14, T18)

    T23 = fd.ops.broadcast_in_dim(T4, shape=V12, broadcast_dims=[1])
    T28 = fd.ops.sum(T1, dims=[0], keepdim=False, dtype=DataType.Null)

    T30 = fd.ops.mul(T1, T23)
    T31 = fd.ops.mul(T1, T19)
    T32 = fd.ops.sum(T31, dims=[0], keepdim=False, dtype=DataType.Null)

    T34 = fd.ops.mul(T30, T18)
    T35 = fd.ops.mul(T30, T14)
    T36 = fd.ops.sum(T35, dims=[1], keepdim=False, dtype=DataType.Null)

    T40 = fd.ops.broadcast_in_dim(T36, shape=[T0.size(0), 1], broadcast_dims=[0])
    T41 = fd.ops.neg(T34)
    T42 = fd.ops.sum(T41, dims=[1], keepdim=False, dtype=DataType.Null)
    T46 = fd.ops.broadcast_in_dim(T42, shape=[T0.size(0), 1], broadcast_dims=[0])
    S47 = fd.define_scalar(-0.500000, dtype=DataType.Double)
    T48 = fd.ops.mul(S47, T40)
    S49 = fd.define_scalar(3.00000, dtype=DataType.Double)
    T50 = fd.ops.pow(T3, S49)
    T51 = fd.ops.mul(T48, T50)
    T54 = fd.ops.sum(T46, dims=[1], keepdim=False, dtype=DataType.Null)
    T55 = fd.ops.sum(T51, dims=[1], keepdim=False, dtype=DataType.Null)

    T59 = fd.ops.broadcast_in_dim(T55, shape=[T0.size(0), 1], broadcast_dims=[0])
    T63 = fd.ops.broadcast_in_dim(T59, shape=V12, broadcast_dims=[0, 1])
    T67 = fd.ops.broadcast_in_dim(T2, shape=[T0.size(0), 1], broadcast_dims=[0])
    T71 = fd.ops.broadcast_in_dim(T67, shape=V12, broadcast_dims=[0, 1])

    S72 = fd.define_scalar(2.00000, dtype=DataType.Double)
    T73 = fd.ops.mul(S72, T63)
    T74 = fd.ops.sub(T0, T71)
    T75 = fd.ops.mul(T73, T74)

    S77 = fd.ops.reciprocal(T0.size(1))
    T78 = fd.ops.mul(T75, S77)
    T82 = fd.ops.broadcast_in_dim(T54, shape=[T0.size(0), 1], broadcast_dims=[0])
    T86 = fd.ops.broadcast_in_dim(T82, shape=V12, broadcast_dims=[0, 1])
    T88 = fd.ops.mul(S77, T86)
    T89 = fd.ops.add(T78, T88)
    T90 = fd.ops.add(T34, T89)

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

    fd.add_output(T90)
    fd.add_output(T32)
    fd.add_output(T28)

def layernorm_bwd_iobytes(size: tuple, dtype: torch.dtype):
    # Manual IOBytes computation since nvfuser input/outputs (in_tensor, grad_out, mean, invstd, weigts) differ from baselines (out, grad_out)
    # Total IO bytes = in_tensor (size, dtype) + grad_out (size, dtype) + mean (size[0], float) +
    #       invstd (size[0], float) + weights (size[1], dtype) +
    #       grad_in (size, dtype) + grad_weights (size[1], dtype)+ grad_bias (size[1], dtype)
    return int(
        dtype.itemsize * 3 * (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_outer_persistent
@pytest.mark.inner_persistent
def test_layernorm_bwd_nvf_benchmark(
    benchmark,
    size: tuple,


    T75 = fd.ops.mul(T73, T74)

    S77 = fd.ops.reciprocal(T0.size(1))
    T78 = fd.ops.mul(T75, S77)
    T82 = fd.ops.broadcast_in_dim(T54, shape=[T0.size(0), 1], broadcast_dims=[0])
    T86 = fd.ops.broadcast_in_dim(T82, shape=V12, broadcast_dims=[0, 1])
    T88 = fd.ops.mul(S77, T86)
    T89 = fd.ops.add(T78, T88)
    T90 = fd.ops.add(T34, T89)

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

    fd.add_output(T90)
    fd.add_output(T32)
    fd.add_output(T28)

def layernorm_bwd_iobytes(size: tuple, dtype: torch.dtype):
    # Manual IOBytes computation since nvfuser input/outputs (in_tensor, grad_out, mean, invstd, weigts) differ from baselines (out, grad_out)
    # Total IO bytes = in_tensor (size, dtype) + grad_out (size, dtype) + mean (size[0], float) +
    #       invstd (size[0], float) + weights (size[1], dtype) +
    #       grad_in (size, dtype) + grad_weights (size[1], dtype)+ grad_bias (size[1], dtype)
    return int(
        dtype.itemsize * 3 * (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_outer_persistent
@pytest.mark.inner_persistent
def test_layernorm_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)
    bias = torch.randn(size[1], device="cuda", dtype=dtype, requires_grad=True)

    mean = inputs.to(torch.float).mean(dim=-1)


    fd.add_output(T90)
    fd.add_output(T32)
    fd.add_output(T28)

def layernorm_bwd_iobytes(size: tuple, dtype: torch.dtype):
    # Manual IOBytes computation since nvfuser input/outputs (in_tensor, grad_out, mean, invstd, weigts) differ from baselines (out, grad_out)
    # Total IO bytes = in_tensor (size, dtype) + grad_out (size, dtype) + mean (size[0], float) +
    #       invstd (size[0], float) + weights (size[1], dtype) +
    #       grad_in (size, dtype) + grad_weights (size[1], dtype)+ grad_bias (size[1], dtype)
    return int(
        dtype.itemsize * 3 * (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_outer_persistent
@pytest.mark.inner_persistent
def test_layernorm_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)
    bias = torch.randn(size[1], device="cuda", dtype=dtype, requires_grad=True)

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

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

    if not disable_validation:
        eager_output = torch.nn.functional.layer_norm(
            inputs.to(torch.double),
            inputs.shape[1:],
            weight=weights.to(torch.double),
            bias=bias.to(torch.double),
        )
        eager_output.backward(grads.to(torch.double))
        fd.validate(
            [inputs, grads, mean, invstd, weights],
            [inputs.grad, weights.grad, bias.grad],
        )

    if not disable_benchmarking:
        run_benchmark(benchmark, fd.execute, [inputs, grads, mean, invstd, 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_layernorm_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)
    bias = torch.randn(size[1], device="cuda", dtype=dtype, requires_grad=True)


            inputs.to(torch.double),
            inputs.shape[1:],
            weight=weights.to(torch.double),
            bias=bias.to(torch.double),
        )
        eager_output.backward(grads.to(torch.double))
        fd.validate(
            [inputs, grads, mean, invstd, weights],
            [inputs.grad, weights.grad, bias.grad],
        )

    if not disable_benchmarking:
        run_benchmark(benchmark, fd.execute, [inputs, grads, mean, invstd, 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_layernorm_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)
    bias = torch.randn(size[1], device="cuda", dtype=dtype, requires_grad=True)

    fwd_fn = with_executor(executor, layernorm)
    fwd_inputs = [inputs, weights, bias]
    outputs = fwd_fn(fwd_inputs)

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


# SPDX-FileCopyrightText: Copyright (c) 2026-present NVIDIA CORPORATION & AFFILIATES.
# All rights reserved.
# SPDX-License-Identifier: BSD-3-Clause
import pytest
from .benchmark_inference import (
    InferenceBenchmarkConfig,
    InferenceBenchmark,
    _register_nvfp4_ops,
)
from nvfuser_direct.pytorch_utils import DEVICE_PROPERTIES

# NOTE: for some reason, changing the order of nvfp4 and cudagraph parameter breaks thunder benchmark. I suspect there's something with thunder's cache.
@pytest.mark.parametrize("input_length", [4096])
@pytest.mark.parametrize("output_length", [4])
@pytest.mark.parametrize("mode", ["thunder", "inductor"])
@pytest.mark.parametrize("enable_nvfp4", [True, False])
@pytest.mark.parametrize("enable_cudagraph", [False, True])
def test_llama4_inference_benchmark(
    benchmark,
    input_length: int,
    output_length: int,
    mode: str,
    enable_nvfp4: bool,
    enable_cudagraph: bool,
):
    if mode == "inductor" and enable_nvfp4:
        pytest.skip("nvfp4 is not supported by inductor yet.")

    if mode == "thunder" and enable_nvfp4 and enable_cudagraph:
        pytest.skip("FIXME: nvfp4 and cudagraph doesn't work together.")

    if DEVICE_PROPERTIES["gpu_compute_capability_major"] < 10:
        if enable_nvfp4:
            pytest.skip("nvfp4 support requires compute_capability >= 10.0")
        if mode == "thunder" and enable_cudagraph:
            pytest.skip("cudagraph doesn't support grouped matmul")

    if enable_nvfp4:
        _register_nvfp4_ops()

    config = InferenceBenchmarkConfig(
        model_name="meta-llama/Llama-4-Maverick-17B-128E",
        batch_size=1,
        input_length=input_length,
        output_length=output_length,
        num_layers=2,
        num_iterations=5,
        warmup_iterations=5,
        mode=mode,
        enable_nvfp4=enable_nvfp4,
        fx_report_folder=None,
        enable_nv_linear=True,
        disable_moe_replacement=False,
        attn_implementation=None,
        thunder_cache=None,
        enable_cudagraph=enable_cudagraph,
        debug_moe=False,
        use_hardcoded_model=True,
    )
    inference_benchmark = InferenceBenchmark(config)

    inference_benchmark.run_benchmark()
    inference_benchmark.print_results()


# 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 scale_bias_relu

def sbr_fwd_fusion(
    fd: FusionDefinition,
    dtype: DataType,
):
    T0 = fd.define_tensor(shape=[-1], contiguity=[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, -1],
        contiguity=[True, True],
        dtype=dtype,
        is_cpu=False,
    )

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

    V7 = T2.shape()
    T8 = fd.ops.broadcast_in_dim(T1, shape=V7, broadcast_dims=[1])
    T11 = fd.ops.mul(T2, T8)

    T18 = fd.ops.broadcast_in_dim(T0, shape=V7, broadcast_dims=[1])
    T20 = fd.ops.add(T11, T18)

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

    S22 = fd.define_scalar(0.00000, dtype=DataType.Double)
    T23 = fd.ops.gt(T20, S22)
    T25 = fd.ops.where(T23, T20, S22)

    fd.add_output(T23)
    fd.add_output(T25)

def sbr_fwd_iobytes(size: tuple, dtype: torch.dtype):
    # Total IO bytes = input (dtype, size) + bias (dtype, size[-1]) + scale (dtype, size[-1])+ output (dtype, size) + bool_mask (bool, size)
    return int(
        dtype.itemsize * 2 * (np.prod(size) + size[-1])
        + torch.bool.itemsize * np.prod(size)
    )

@pytest.mark.parametrize("size", generate_input_sizes(dims=2))
@pytest.mark.parametrize("dtype", FLOAT_DTYPES)
@pytest.mark.pointwise
def test_sbr_fwd_nvf_benchmark(
    benchmark,
    size: tuple,
    dtype: torch.dtype,
    disable_validation: bool,
    disable_benchmarking: bool,
):
    inputs = torch.randn(*size, device="cuda", dtype=dtype, requires_grad=True)
    bias = torch.ones(size[-1], device="cuda", dtype=dtype)
    scale = torch.ones(size[-1], device="cuda", dtype=dtype)

    with FusionDefinition() as fd:
        sbr_fwd_fusion(fd, torch_dtype_to_nvfuser_dtype(dtype))
    if not disable_validation:
        eager_output = torch.nn.functional.relu(inputs * scale + bias)
        bool_mask = torch.gt(inputs * scale + bias, 0.0)
        fd.validate([bias, scale, inputs], [bool_mask, eager_output])

    if not disable_benchmarking:
        run_benchmark(benchmark, fd.execute, [bias, scale, inputs])

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

    benchmark_fn = with_executor(executor, scale_bias_relu)

    run_benchmark(
        benchmark,
        benchmark_fn,
        [inputs, scale, bias],
        iobytes=sbr_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 scale_bias_relu

def sbr_fwd_fusion(
    fd: FusionDefinition,
    dtype: DataType,
):
    T0 = fd.define_tensor(shape=[-1], contiguity=[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, -1],
        contiguity=[True, True],
        dtype=dtype,
        is_cpu=False,
    )

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

    V7 = T2.shape()
    T8 = fd.ops.broadcast_in_dim(T1, shape=V7, broadcast_dims=[1])
    T11 = fd.ops.mul(T2, T8)

    T18 = fd.ops.broadcast_in_dim(T0, shape=V7, broadcast_dims=[1])
    T20 = fd.ops.add(T11, T18)

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

    S22 = fd.define_scalar(0.00000, dtype=DataType.Double)
    T23 = fd.ops.gt(T20, S22)
    T25 = fd.ops.where(T23, T20, S22)

    fd.add_output(T23)
    fd.add_output(T25)

def sbr_fwd_iobytes(size: tuple, dtype: torch.dtype):
    # Total IO bytes = input (dtype, size) + bias (dtype, size[-1]) + scale (dtype, size[-1])+ output (dtype, size) + bool_mask (bool, size)
    return int(
        dtype.itemsize * 2 * (np.prod(size) + size[-1])
        + torch.bool.itemsize * np.prod(size)
    )

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


        T2 = fd.ops.cast(T2, dtype=DataType.Float)

    V7 = T2.shape()
    T8 = fd.ops.broadcast_in_dim(T1, shape=V7, broadcast_dims=[1])
    T11 = fd.ops.mul(T2, T8)

    T18 = fd.ops.broadcast_in_dim(T0, shape=V7, broadcast_dims=[1])
    T20 = fd.ops.add(T11, T18)

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

    S22 = fd.define_scalar(0.00000, dtype=DataType.Double)
    T23 = fd.ops.gt(T20, S22)
    T25 = fd.ops.where(T23, T20, S22)

    fd.add_output(T23)
    fd.add_output(T25)

def sbr_fwd_iobytes(size: tuple, dtype: torch.dtype):
    # Total IO bytes = input (dtype, size) + bias (dtype, size[-1]) + scale (dtype, size[-1])+ output (dtype, size) + bool_mask (bool, size)
    return int(
        dtype.itemsize * 2 * (np.prod(size) + size[-1])
        + torch.bool.itemsize * np.prod(size)
    )

@pytest.mark.parametrize("size", generate_input_sizes(dims=2))
@pytest.mark.parametrize("dtype", FLOAT_DTYPES)
@pytest.mark.pointwise
def test_sbr_fwd_nvf_benchmark(
    benchmark,
    size: tuple,
    dtype: torch.dtype,
    disable_validation: bool,
    disable_benchmarking: bool,
):
    inputs = torch.randn(*size, device="cuda", dtype=dtype, requires_grad=True)
    bias = torch.ones(size[-1], device="cuda", dtype=dtype)
    scale = torch.ones(size[-1], device="cuda", dtype=dtype)

    with FusionDefinition() as fd:
        sbr_fwd_fusion(fd, torch_dtype_to_nvfuser_dtype(dtype))
    if not disable_validation:
        eager_output = torch.nn.functional.relu(inputs * scale + bias)


    S22 = fd.define_scalar(0.00000, dtype=DataType.Double)
    T23 = fd.ops.gt(T20, S22)
    T25 = fd.ops.where(T23, T20, S22)

    fd.add_output(T23)
    fd.add_output(T25)

def sbr_fwd_iobytes(size: tuple, dtype: torch.dtype):
    # Total IO bytes = input (dtype, size) + bias (dtype, size[-1]) + scale (dtype, size[-1])+ output (dtype, size) + bool_mask (bool, size)
    return int(
        dtype.itemsize * 2 * (np.prod(size) + size[-1])
        + torch.bool.itemsize * np.prod(size)
    )

@pytest.mark.parametrize("size", generate_input_sizes(dims=2))
@pytest.mark.parametrize("dtype", FLOAT_DTYPES)
@pytest.mark.pointwise
def test_sbr_fwd_nvf_benchmark(
    benchmark,
    size: tuple,
    dtype: torch.dtype,
    disable_validation: bool,
    disable_benchmarking: bool,
):
    inputs = torch.randn(*size, device="cuda", dtype=dtype, requires_grad=True)
    bias = torch.ones(size[-1], device="cuda", dtype=dtype)
    scale = torch.ones(size[-1], device="cuda", dtype=dtype)

    with FusionDefinition() as fd:
        sbr_fwd_fusion(fd, torch_dtype_to_nvfuser_dtype(dtype))
    if not disable_validation:
        eager_output = torch.nn.functional.relu(inputs * scale + bias)
        bool_mask = torch.gt(inputs * scale + bias, 0.0)
        fd.validate([bias, scale, inputs], [bool_mask, eager_output])

    if not disable_benchmarking:
        run_benchmark(benchmark, fd.execute, [bias, scale, inputs])

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

    benchmark_fn = with_executor(executor, scale_bias_relu)


):
    inputs = torch.randn(*size, device="cuda", dtype=dtype, requires_grad=True)
    bias = torch.ones(size[-1], device="cuda", dtype=dtype)
    scale = torch.ones(size[-1], device="cuda", dtype=dtype)

    with FusionDefinition() as fd:
        sbr_fwd_fusion(fd, torch_dtype_to_nvfuser_dtype(dtype))
    if not disable_validation:
        eager_output = torch.nn.functional.relu(inputs * scale + bias)
        bool_mask = torch.gt(inputs * scale + bias, 0.0)
        fd.validate([bias, scale, inputs], [bool_mask, eager_output])

    if not disable_benchmarking:
        run_benchmark(benchmark, fd.execute, [bias, scale, inputs])

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

    benchmark_fn = with_executor(executor, scale_bias_relu)

    run_benchmark(
        benchmark,
        benchmark_fn,
        [inputs, scale, bias],
        iobytes=sbr_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
from .core import run_benchmark
import torch

import csv
import functools
import os

def matmul_fusion(fd: FusionDefinition, inputs: list[torch.Tensor]) -> None:
    a = fd.from_pytorch(inputs[0])
    b = fd.from_pytorch(inputs[1])
    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"
        )

@pytest.mark.parametrize("half_reduction", [False, True], ids=["fullred", "halfred"])
@pytest.mark.parametrize("executor", ["eager"])
@pytest.mark.parametrize("dtype", [torch.float16, torch.bfloat16], ids=["fp16", "bf16"])
@pytest.mark.parametrize(
    "config", load_matmul_problems(), ids=lambda val: "-".join(str(v) for v in val)
)
@pytest.mark.matmul
def test_matmul_baseline_benchmark(
    benchmark,
    executor: str,
    config: tuple,
    dtype: torch.dtype,
    half_reduction: bool,
    disable_validation: bool,
    disable_benchmarking: bool,
):
    m, n, k, layout = config

    maybe_skip_oom_case(m, n, k)

    torch.backends.cuda.matmul.allow_fp16_reduced_precision_reduction = half_reduction
    torch.backends.cuda.matmul.allow_bf16_reduced_precision_reduction = half_reduction

    a = torch.randn(m, k, device="cuda", dtype=dtype)
    b = torch.randn(k, n, device="cuda", dtype=dtype)

    if layout == "NT" or layout == "NN":
        a = a.as_strided(size=[m, k], stride=[1, m])
    if layout == "TN" or layout == "NN":
        b = b.as_strided(size=[k, n], stride=[1, k])

    # NOTE: we never need to validate eager, as it is our baseline
    run_benchmark(
        benchmark,
        lambda ab: torch.matmul(*ab),
        [a, b],
    )

@pytest.mark.parametrize("half_reduction", [False, True], ids=["fullred", "halfred"])
@pytest.mark.parametrize("dtype", [torch.float16, torch.bfloat16], ids=["fp16", "bf16"])
@pytest.mark.parametrize(
    "config", load_matmul_problems(), ids=lambda val: "-".join(str(v) for v in val)
)
@pytest.mark.matmul
def test_matmul_nvf_benchmark(
    benchmark,
    config: tuple,
    dtype: torch.dtype,
    half_reduction: bool,
    disable_validation: bool,
    disable_benchmarking: bool,
):
    m, n, k, layout = config

    maybe_skip_oom_case(m, n, k)

    torch.backends.cuda.matmul.allow_fp16_reduced_precision_reduction = half_reduction
    torch.backends.cuda.matmul.allow_bf16_reduced_precision_reduction = half_reduction

    if half_reduction:
        # See https://github.com/NVIDIA/Fuser/pull/1719
        pytest.skip("Reduced precision reduction not implemented in nvFuser")

    a = torch.randn(m, k, device="cuda", dtype=dtype)
    b = torch.randn(k, n, device="cuda", dtype=dtype)

    if layout == "NT" or layout == "NN":
        a = a.as_strided(size=[m, k], stride=[1, m])
    if layout == "TN" or layout == "NN":
        b = b.as_strided(size=[k, n], stride=[1, k])

    with FusionDefinition() as fd:
        matmul_fusion(fd, [a, b])

    kwargs = dict(
        _enable_options=["fuse_matmul"], _disable_options=["matmul_expr_eval"]
    )

    if not disable_validation:
        eager_output = torch.matmul(a, b)
        fd.validate([a, b], [eager_output], **kwargs)

    if not disable_benchmarking:
        run_benchmark(benchmark, lambda *args: fd.execute(*args, **kwargs), [a, b])


# 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
from .core import run_benchmark
import torch

import csv
import functools
import os

def matmul_fusion(fd: FusionDefinition, inputs: list[torch.Tensor]) -> None:
    a = fd.from_pytorch(inputs[0])
    b = fd.from_pytorch(inputs[1])
    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


# 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
from .core import run_benchmark
import torch

import csv
import functools
import os

def matmul_fusion(fd: FusionDefinition, inputs: list[torch.Tensor]) -> None:
    a = fd.from_pytorch(inputs[0])
    b = fd.from_pytorch(inputs[1])
    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"
        )

@pytest.mark.parametrize("half_reduction", [False, True], ids=["fullred", "halfred"])
@pytest.mark.parametrize("executor", ["eager"])
@pytest.mark.parametrize("dtype", [torch.float16, torch.bfloat16], ids=["fp16", "bf16"])
@pytest.mark.parametrize(
    "config", load_matmul_problems(), ids=lambda val: "-".join(str(v) for v in val)
)


            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"])
@pytest.mark.parametrize("executor", ["eager"])
@pytest.mark.parametrize("dtype", [torch.float16, torch.bfloat16], ids=["fp16", "bf16"])
@pytest.mark.parametrize(
    "config", load_matmul_problems(), ids=lambda val: "-".join(str(v) for v in val)
)
@pytest.mark.matmul
def test_matmul_baseline_benchmark(
    benchmark,
    executor: str,
    config: tuple,
    dtype: torch.dtype,
    half_reduction: bool,
    disable_validation: bool,
    disable_benchmarking: bool,
):
    m, n, k, layout = config


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"])
@pytest.mark.parametrize("executor", ["eager"])
@pytest.mark.parametrize("dtype", [torch.float16, torch.bfloat16], ids=["fp16", "bf16"])
@pytest.mark.parametrize(
    "config", load_matmul_problems(), ids=lambda val: "-".join(str(v) for v in val)
)
@pytest.mark.matmul
def test_matmul_baseline_benchmark(
    benchmark,
    executor: str,
    config: tuple,
    dtype: torch.dtype,
    half_reduction: bool,
    disable_validation: bool,
    disable_benchmarking: bool,
):
    m, n, k, layout = config

    maybe_skip_oom_case(m, n, k)

    torch.backends.cuda.matmul.allow_fp16_reduced_precision_reduction = half_reduction
    torch.backends.cuda.matmul.allow_bf16_reduced_precision_reduction = half_reduction

    a = torch.randn(m, k, device="cuda", dtype=dtype)
    b = torch.randn(k, n, device="cuda", dtype=dtype)

    if layout == "NT" or layout == "NN":
        a = a.as_strided(size=[m, k], stride=[1, m])
    if layout == "TN" or layout == "NN":
        b = b.as_strided(size=[k, n], stride=[1, k])

    # NOTE: we never need to validate eager, as it is our baseline
    run_benchmark(
        benchmark,
        lambda ab: torch.matmul(*ab),
        [a, b],
    )

@pytest.mark.parametrize("half_reduction", [False, True], ids=["fullred", "halfred"])
@pytest.mark.parametrize("dtype", [torch.float16, torch.bfloat16], ids=["fp16", "bf16"])
@pytest.mark.parametrize(
    "config", load_matmul_problems(), ids=lambda val: "-".join(str(v) for v in val)
)
@pytest.mark.matmul
def test_matmul_nvf_benchmark(
    benchmark,
    config: tuple,
    dtype: torch.dtype,
    half_reduction: bool,
    disable_validation: bool,
    disable_benchmarking: bool,
):
    m, n, k, layout = config

    maybe_skip_oom_case(m, n, k)


    if layout == "NT" or layout == "NN":
        a = a.as_strided(size=[m, k], stride=[1, m])
    if layout == "TN" or layout == "NN":
        b = b.as_strided(size=[k, n], stride=[1, k])

    # NOTE: we never need to validate eager, as it is our baseline
    run_benchmark(
        benchmark,
        lambda ab: torch.matmul(*ab),
        [a, b],
    )

@pytest.mark.parametrize("half_reduction", [False, True], ids=["fullred", "halfred"])
@pytest.mark.parametrize("dtype", [torch.float16, torch.bfloat16], ids=["fp16", "bf16"])
@pytest.mark.parametrize(
    "config", load_matmul_problems(), ids=lambda val: "-".join(str(v) for v in val)
)
@pytest.mark.matmul
def test_matmul_nvf_benchmark(
    benchmark,
    config: tuple,
    dtype: torch.dtype,
    half_reduction: bool,
    disable_validation: bool,
    disable_benchmarking: bool,
):
    m, n, k, layout = config

    maybe_skip_oom_case(m, n, k)

    torch.backends.cuda.matmul.allow_fp16_reduced_precision_reduction = half_reduction
    torch.backends.cuda.matmul.allow_bf16_reduced_precision_reduction = half_reduction

    if half_reduction:
        # See https://github.com/NVIDIA/Fuser/pull/1719
        pytest.skip("Reduced precision reduction not implemented in nvFuser")

    a = torch.randn(m, k, device="cuda", dtype=dtype)
    b = torch.randn(k, n, device="cuda", dtype=dtype)

    if layout == "NT" or layout == "NN":
        a = a.as_strided(size=[m, k], stride=[1, m])
    if layout == "TN" or layout == "NN":
        b = b.as_strided(size=[k, n], stride=[1, k])

    with FusionDefinition() as fd:
        matmul_fusion(fd, [a, b])

    kwargs = dict(
        _enable_options=["fuse_matmul"], _disable_options=["matmul_expr_eval"]
    )

    if not disable_validation:
        eager_output = torch.matmul(a, b)
        fd.validate([a, b], [eager_output], **kwargs)

    if not disable_benchmarking:
        run_benchmark(benchmark, lambda *args: fd.execute(*args, **kwargs), [a, b])


import torch

import csv
import functools
import os

def matmul_fusion(fd: FusionDefinition, inputs: list[torch.Tensor]) -> None:
    a = fd.from_pytorch(inputs[0])
    b = fd.from_pytorch(inputs[1])
    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