        channel_dim = 1 if not channels_last else num_dims - 1
        reduction_axes = [i for i in range(len(size)) if i != channel_dim]
        if norm == "instance_norm":
            reduction_axes.remove(batch_dim)

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

        fd.validate(
            [inputs, weight, bias, running_mean, running_var],
            [eager_output, mean, invstd],
        )

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

def norm_bwd_nvf_benchmark(
    benchmark,
    size: tuple,
    dtype: torch.dtype,
    norm: str,
    channels_last: bool,
    disable_validation: bool,
    disable_benchmarking: bool,
    eps: float = 1e-5,
):
    """
    Common benchmark setup for batchnorm/instance forward call in training mode.
    """

    assert norm in ["batch_norm", "instance_norm"], NotImplementedError

    # Size is assumed to be in the order N, C, ...

    num_dims = len(size)

    at_inputs = torch.randn(size, device="cuda", dtype=dtype, requires_grad=True)
    at_grads = torch.randn(size, device="cuda", dtype=dtype)
    weight = torch.randn(size[1], device="cuda", dtype=dtype, requires_grad=True)
    bias = torch.randn(size[1], device="cuda", dtype=dtype, requires_grad=True)

    # CPP benchmarks assume mean and variance to be of type Float
    running_mean = torch.zeros(size[1], device="cuda", dtype=torch.float)
    running_var = torch.ones(size[1], device="cuda", dtype=torch.float)

    if channels_last:
        at_inputs = at_inputs.to(memory_format=torch.channels_last)
        at_inputs.retain_grad()
        at_grads = at_grads.to(memory_format=torch.channels_last)

        inputs = at_inputs.clone().detach().permute((0, *range(2, num_dims), 1))
        grads = at_grads.clone().detach().permute((0, *range(2, num_dims), 1))

    else:
        inputs = at_inputs
        grads = at_grads

    batch_dim = 0
    channel_dim = 1 if not channels_last else num_dims - 1
    reduction_axes = [i for i in range(len(size)) if i != channel_dim]
    if norm == "instance_norm":
        reduction_axes.remove(batch_dim)
    mean = inputs.to(torch.float).mean(dim=reduction_axes)
    var = inputs.to(torch.float).var(dim=reduction_axes, unbiased=False)
    invstd = 1.0 / torch.sqrt(var + eps)

    with FusionDefinition() as fd:
        norm_bwd_fusion(
            fd=fd,
            dtype=torch_dtype_to_nvfuser_dtype(dtype),
            norm=norm,
            num_dims=num_dims,
            channels_last=channels_last,
        )

    if not disable_validation:
        # PyTorch expects running mean and variance to be of same type as input.
        if norm == "batch_norm":
            eager_output = torch.nn.functional.batch_norm(
                at_inputs.to(torch.double),
                running_mean.to(torch.double),
                running_var.to(torch.double),
                weight=weight.to(torch.double),
                bias=bias.to(torch.double),
                training=True,
            )
        elif norm == "instance_norm":
            eager_output = torch.nn.functional.instance_norm(
                at_inputs.to(torch.double),
                running_mean.to(torch.double),
                running_var.to(torch.double),
                weight=weight.to(torch.double),
                bias=bias.to(torch.double),
            )

        eager_output.backward(at_grads.to(torch.double))

        if channels_last:
            eager_grad = at_inputs.grad.permute((0, *range(2, num_dims), 1))
        else:
            eager_grad = at_inputs.grad

        fd.validate(
            [inputs, grads, weight, running_mean, running_var, mean, invstd],
            [eager_grad, weight.grad, bias.grad],
        )

    if not disable_benchmarking:
        run_benchmark(
            benchmark,
            fd.execute,
            [inputs, grads, weight, running_mean, running_var, mean, invstd],
        )

def batchnorm_fwd_fn(inputs: list):
    input, weight, bias, running_mean, running_var = inputs
    return torch.nn.functional.batch_norm(
        input,
        running_mean,
        running_var,
        weight=weight,
        bias=bias,
        training=True,
    )

def instancenorm_fwd_fn(inputs: list):
    input, weight, bias, running_mean, running_var = inputs
    return torch.nn.functional.instance_norm(
        input,
        running_mean,
        running_var,


        eager_output.backward(at_grads.to(torch.double))

        if channels_last:
            eager_grad = at_inputs.grad.permute((0, *range(2, num_dims), 1))
        else:
            eager_grad = at_inputs.grad

        fd.validate(
            [inputs, grads, weight, running_mean, running_var, mean, invstd],
            [eager_grad, weight.grad, bias.grad],
        )

    if not disable_benchmarking:
        run_benchmark(
            benchmark,
            fd.execute,
            [inputs, grads, weight, running_mean, running_var, mean, invstd],
        )

def batchnorm_fwd_fn(inputs: list):
    input, weight, bias, running_mean, running_var = inputs
    return torch.nn.functional.batch_norm(
        input,
        running_mean,
        running_var,
        weight=weight,
        bias=bias,
        training=True,
    )

def instancenorm_fwd_fn(inputs: list):
    input, weight, bias, running_mean, running_var = inputs
    return torch.nn.functional.instance_norm(
        input,
        running_mean,
        running_var,
        weight=weight,
        bias=bias,
    )

def norm_fwd_iobytes(size: tuple, dtype: torch.dtype, norm: str):
    # Manual IOBytes computation required since nvFuser outputs (out, mean, invstd) differs from baselines (out)
    # size = [N, C, H, W]
    # Total IO bytes = in_tensor (size, dtype) + weight (size[1], dtype) + bias (size[1], dtype) +
    #           running_mean (size[1], float) + running_var (size[1], float) + output (size, dtype) +
    #           mean ([C]/[N, C] , float) + invstd ([C]/[N, C] , float)
    stat_size = (


    if not disable_benchmarking:
        run_benchmark(
            benchmark,
            fd.execute,
            [inputs, grads, weight, running_mean, running_var, mean, invstd],
        )

def batchnorm_fwd_fn(inputs: list):
    input, weight, bias, running_mean, running_var = inputs
    return torch.nn.functional.batch_norm(
        input,
        running_mean,
        running_var,
        weight=weight,
        bias=bias,
        training=True,
    )

def instancenorm_fwd_fn(inputs: list):
    input, weight, bias, running_mean, running_var = inputs
    return torch.nn.functional.instance_norm(
        input,
        running_mean,
        running_var,
        weight=weight,
        bias=bias,
    )

def norm_fwd_iobytes(size: tuple, dtype: torch.dtype, norm: str):
    # Manual IOBytes computation required since nvFuser outputs (out, mean, invstd) differs from baselines (out)
    # size = [N, C, H, W]
    # Total IO bytes = in_tensor (size, dtype) + weight (size[1], dtype) + bias (size[1], dtype) +
    #           running_mean (size[1], float) + running_var (size[1], float) + output (size, dtype) +
    #           mean ([C]/[N, C] , float) + invstd ([C]/[N, C] , float)
    stat_size = (
        size[1] if norm == "batch_norm" else size[0] * size[1]
    )  # size of mean/invstd
    return int(
        dtype.itemsize * 2 * (np.prod(size) + size[1])
        + torch.float.itemsize * 2 * (size[1] + stat_size)
    )

def norm_bwd_iobytes(size: tuple, dtype: torch.dtype, norm: str):
    # Manual IOBytes computation since nvfuser input/outputs (in_tensor, grad_out, mean, invstd, weigts, grad_in, grad_weight, grad_bias) differ from baselines (out, grad_out)
    # size = [N, C, H, W]


        input,
        running_mean,
        running_var,
        weight=weight,
        bias=bias,
        training=True,
    )

def instancenorm_fwd_fn(inputs: list):
    input, weight, bias, running_mean, running_var = inputs
    return torch.nn.functional.instance_norm(
        input,
        running_mean,
        running_var,
        weight=weight,
        bias=bias,
    )

def norm_fwd_iobytes(size: tuple, dtype: torch.dtype, norm: str):
    # Manual IOBytes computation required since nvFuser outputs (out, mean, invstd) differs from baselines (out)
    # size = [N, C, H, W]
    # Total IO bytes = in_tensor (size, dtype) + weight (size[1], dtype) + bias (size[1], dtype) +
    #           running_mean (size[1], float) + running_var (size[1], float) + output (size, dtype) +
    #           mean ([C]/[N, C] , float) + invstd ([C]/[N, C] , float)
    stat_size = (
        size[1] if norm == "batch_norm" else size[0] * size[1]
    )  # size of mean/invstd
    return int(
        dtype.itemsize * 2 * (np.prod(size) + size[1])
        + torch.float.itemsize * 2 * (size[1] + stat_size)
    )

def norm_bwd_iobytes(size: tuple, dtype: torch.dtype, norm: str):
    # Manual IOBytes computation since nvfuser input/outputs (in_tensor, grad_out, mean, invstd, weigts, grad_in, grad_weight, grad_bias) differ from baselines (out, grad_out)
    # size = [N, C, H, W]
    # Total IO bytes = in_tensor (size, dtype) + weight (size[1], dtype) +
    #           running_mean (size[1], float) + running_var (size[1], float) +
    #           mean ([C]/[N, C] , float) + invstd ([C]/[N, C] , float) + grad_out (size, dtype) +
    #           grad_in (size, dtype) + grad_weight (size[1], dtype) + grad_bias (size[1], dtype)
    stat_size = size[1] if norm == "batch_norm" else size[0] * size[1]
    return int(
        dtype.itemsize * 3 * (np.prod(size) + size[1])
        + torch.float.itemsize * 2 * (size[1] + stat_size)
    )

def norm_fwd_baseline_benchmark(
    benchmark,
    size: tuple,
    dtype: torch.dtype,


        weight=weight,
        bias=bias,
    )

def norm_fwd_iobytes(size: tuple, dtype: torch.dtype, norm: str):
    # Manual IOBytes computation required since nvFuser outputs (out, mean, invstd) differs from baselines (out)
    # size = [N, C, H, W]
    # Total IO bytes = in_tensor (size, dtype) + weight (size[1], dtype) + bias (size[1], dtype) +
    #           running_mean (size[1], float) + running_var (size[1], float) + output (size, dtype) +
    #           mean ([C]/[N, C] , float) + invstd ([C]/[N, C] , float)
    stat_size = (
        size[1] if norm == "batch_norm" else size[0] * size[1]
    )  # size of mean/invstd
    return int(
        dtype.itemsize * 2 * (np.prod(size) + size[1])
        + torch.float.itemsize * 2 * (size[1] + stat_size)
    )

def norm_bwd_iobytes(size: tuple, dtype: torch.dtype, norm: str):
    # Manual IOBytes computation since nvfuser input/outputs (in_tensor, grad_out, mean, invstd, weigts, grad_in, grad_weight, grad_bias) differ from baselines (out, grad_out)
    # size = [N, C, H, W]
    # Total IO bytes = in_tensor (size, dtype) + weight (size[1], dtype) +
    #           running_mean (size[1], float) + running_var (size[1], float) +
    #           mean ([C]/[N, C] , float) + invstd ([C]/[N, C] , float) + grad_out (size, dtype) +
    #           grad_in (size, dtype) + grad_weight (size[1], dtype) + grad_bias (size[1], dtype)
    stat_size = size[1] if norm == "batch_norm" else size[0] * size[1]
    return int(
        dtype.itemsize * 3 * (np.prod(size) + size[1])
        + torch.float.itemsize * 2 * (size[1] + stat_size)
    )

def norm_fwd_baseline_benchmark(
    benchmark,
    size: tuple,
    dtype: torch.dtype,
    channels_last: bool,
    executor: str,
    norm: str,
):
    if executor == "torchcompile":
        clear_dynamo_cache()

    assert norm in ["batch_norm", "instance_norm"], NotImplementedError

    # Size is assumed to be in the order N, C, ...
    inputs = torch.randn(size, device="cuda", dtype=dtype, requires_grad=True)
    weight = torch.randn(size[1], device="cuda", dtype=dtype, requires_grad=True)
    bias = torch.randn(size[1], device="cuda", dtype=dtype, requires_grad=True)


    return int(
        dtype.itemsize * 2 * (np.prod(size) + size[1])
        + torch.float.itemsize * 2 * (size[1] + stat_size)
    )

def norm_bwd_iobytes(size: tuple, dtype: torch.dtype, norm: str):
    # Manual IOBytes computation since nvfuser input/outputs (in_tensor, grad_out, mean, invstd, weigts, grad_in, grad_weight, grad_bias) differ from baselines (out, grad_out)
    # size = [N, C, H, W]
    # Total IO bytes = in_tensor (size, dtype) + weight (size[1], dtype) +
    #           running_mean (size[1], float) + running_var (size[1], float) +
    #           mean ([C]/[N, C] , float) + invstd ([C]/[N, C] , float) + grad_out (size, dtype) +
    #           grad_in (size, dtype) + grad_weight (size[1], dtype) + grad_bias (size[1], dtype)
    stat_size = size[1] if norm == "batch_norm" else size[0] * size[1]
    return int(
        dtype.itemsize * 3 * (np.prod(size) + size[1])
        + torch.float.itemsize * 2 * (size[1] + stat_size)
    )

def norm_fwd_baseline_benchmark(
    benchmark,
    size: tuple,
    dtype: torch.dtype,
    channels_last: bool,
    executor: str,
    norm: str,
):
    if executor == "torchcompile":
        clear_dynamo_cache()

    assert norm in ["batch_norm", "instance_norm"], NotImplementedError

    # Size is assumed to be in the order N, C, ...
    inputs = torch.randn(size, device="cuda", dtype=dtype, requires_grad=True)
    weight = torch.randn(size[1], device="cuda", dtype=dtype, requires_grad=True)
    bias = torch.randn(size[1], device="cuda", dtype=dtype, requires_grad=True)

    running_mean = torch.zeros(size[1], device="cuda", dtype=dtype)
    running_var = torch.ones(size[1], device="cuda", dtype=dtype)
    if channels_last:
        inputs = inputs.to(memory_format=torch.channels_last)

    norm_fwd_fn = batchnorm_fwd_fn if norm == "batch_norm" else instancenorm_fwd_fn

    benchmark_fn = with_executor(executor, norm_fwd_fn)

    # Manually compute IOBytes: See PR #1725
    run_benchmark(
        benchmark,
        benchmark_fn,
        [inputs, weight, bias, running_mean, running_var],
        iobytes=norm_fwd_iobytes(size, dtype, norm),
    )

def norm_bwd_baseline_benchmark(
    benchmark,
    size: tuple,
    dtype: torch.dtype,
    channels_last: bool,
    executor: str,
    norm: str,
):
    if executor == "torchcompile":
        clear_dynamo_cache()

    assert norm in ["batch_norm", "instance_norm"], NotImplementedError

    # Size is assumed to be in the order N, C, ...
    inputs = torch.randn(size, device="cuda", dtype=dtype, requires_grad=True)
    weight = torch.randn(size[1], device="cuda", dtype=dtype, requires_grad=True)
    bias = torch.randn(size[1], device="cuda", dtype=dtype, requires_grad=True)


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

    running_mean = torch.zeros(size[1], device="cuda", dtype=dtype)
    running_var = torch.ones(size[1], device="cuda", dtype=dtype)
    if channels_last:
        inputs = inputs.to(memory_format=torch.channels_last)

    norm_fwd_fn = batchnorm_fwd_fn if norm == "batch_norm" else instancenorm_fwd_fn

    benchmark_fn = with_executor(executor, norm_fwd_fn)

    # Manually compute IOBytes: See PR #1725
    run_benchmark(
        benchmark,
        benchmark_fn,
        [inputs, weight, bias, running_mean, running_var],
        iobytes=norm_fwd_iobytes(size, dtype, norm),
    )

def norm_bwd_baseline_benchmark(
    benchmark,
    size: tuple,
    dtype: torch.dtype,
    channels_last: bool,
    executor: str,
    norm: str,
):
    if executor == "torchcompile":
        clear_dynamo_cache()

    assert norm in ["batch_norm", "instance_norm"], NotImplementedError

    # Size is assumed to be in the order N, C, ...
    inputs = torch.randn(size, device="cuda", dtype=dtype, requires_grad=True)
    weight = torch.randn(size[1], device="cuda", dtype=dtype, requires_grad=True)
    bias = torch.randn(size[1], device="cuda", dtype=dtype, requires_grad=True)

    running_mean = torch.zeros(size[1], device="cuda", dtype=dtype)
    running_var = torch.ones(size[1], device="cuda", dtype=dtype)
    grads = torch.randn(size, device="cuda", dtype=dtype)

    if channels_last:
        inputs = inputs.to(memory_format=torch.channels_last)
        grads = grads.to(memory_format=torch.channels_last)

    norm_fwd_fn = batchnorm_fwd_fn if norm == "batch_norm" else instancenorm_fwd_fn

    fwd_fn = with_executor(executor, norm_fwd_fn)
    fwd_inputs = [inputs, weight, bias, running_mean, running_var]
    outputs = fwd_fn(fwd_inputs)

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


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

def pointwise_mul_fusion(
    fd: FusionDefinition,
    dtype: DataType,
) -> 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.mul(T0, T0)
    if dtype in PROMOTE_DTYPES:
        T2 = fd.ops.cast(T2, dtype=dtype)
    fd.add_output(T2)

def pointwise_mul_fwd_fn(inputs: list):  # in_tensor
    return torch.mul(inputs[0], inputs[0])

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

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

    if not disable_validation:
        eager_output = torch.mul(inputs[0].to(torch.double), inputs[0].to(torch.double))
        fd.validate(inputs, [eager_output.to(dtype)])

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

@pytest.mark.parametrize("executor", DEFAULT_EXECUTORS)
@pytest.mark.parametrize("size", generate_input_sizes(dims=2))
@pytest.mark.parametrize("dtype", FLOAT_DTYPES)
@pytest.mark.pointwise
def test_pointwise_mul_baseline_benchmark(
    benchmark,
    size: tuple,
    dtype: torch.dtype,
    executor: str,
):
    if executor == "torchcompile":
        clear_dynamo_cache()
    input = torch.randn(size, device="cuda", dtype=dtype)

    benchmark_fn = with_executor(executor, pointwise_mul_fwd_fn)

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


# 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

def pointwise_mul_fusion(
    fd: FusionDefinition,
    dtype: DataType,
) -> 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.mul(T0, T0)
    if dtype in PROMOTE_DTYPES:
        T2 = fd.ops.cast(T2, dtype=dtype)
    fd.add_output(T2)

def pointwise_mul_fwd_fn(inputs: list):  # in_tensor
    return torch.mul(inputs[0], inputs[0])

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

    with FusionDefinition() as fd:
        pointwise_mul_fusion(fd, torch_dtype_to_nvfuser_dtype(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

def pointwise_mul_fusion(
    fd: FusionDefinition,
    dtype: DataType,
) -> 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.mul(T0, T0)
    if dtype in PROMOTE_DTYPES:
        T2 = fd.ops.cast(T2, dtype=dtype)
    fd.add_output(T2)

def pointwise_mul_fwd_fn(inputs: list):  # in_tensor
    return torch.mul(inputs[0], inputs[0])

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

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

    if not disable_validation:
        eager_output = torch.mul(inputs[0].to(torch.double), inputs[0].to(torch.double))
        fd.validate(inputs, [eager_output.to(dtype)])


    dtype: DataType,
) -> 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.mul(T0, T0)
    if dtype in PROMOTE_DTYPES:
        T2 = fd.ops.cast(T2, dtype=dtype)
    fd.add_output(T2)

def pointwise_mul_fwd_fn(inputs: list):  # in_tensor
    return torch.mul(inputs[0], inputs[0])

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

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

    if not disable_validation:
        eager_output = torch.mul(inputs[0].to(torch.double), inputs[0].to(torch.double))
        fd.validate(inputs, [eager_output.to(dtype)])

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

@pytest.mark.parametrize("executor", DEFAULT_EXECUTORS)
@pytest.mark.parametrize("size", generate_input_sizes(dims=2))
@pytest.mark.parametrize("dtype", FLOAT_DTYPES)
@pytest.mark.pointwise
def test_pointwise_mul_baseline_benchmark(
    benchmark,
    size: tuple,
    dtype: torch.dtype,
    executor: str,
):
    if executor == "torchcompile":
        clear_dynamo_cache()
    input = torch.randn(size, device="cuda", dtype=dtype)

    benchmark_fn = with_executor(executor, pointwise_mul_fwd_fn)

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


    disable_validation: bool,
    disable_benchmarking: bool,
):
    inputs = [torch.randn(size, device="cuda", dtype=dtype)]

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

    if not disable_validation:
        eager_output = torch.mul(inputs[0].to(torch.double), inputs[0].to(torch.double))
        fd.validate(inputs, [eager_output.to(dtype)])

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

@pytest.mark.parametrize("executor", DEFAULT_EXECUTORS)
@pytest.mark.parametrize("size", generate_input_sizes(dims=2))
@pytest.mark.parametrize("dtype", FLOAT_DTYPES)
@pytest.mark.pointwise
def test_pointwise_mul_baseline_benchmark(
    benchmark,
    size: tuple,
    dtype: torch.dtype,
    executor: str,
):
    if executor == "torchcompile":
        clear_dynamo_cache()
    input = torch.randn(size, device="cuda", dtype=dtype)

    benchmark_fn = with_executor(executor, pointwise_mul_fwd_fn)

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


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

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

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


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

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


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

def reduction_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.sum(T0, dims=[reduction_axis], keepdim=False)
    if dtype in PROMOTE_DTYPES:
        T2 = fd.ops.cast(T2, dtype=dtype)
    fd.add_output(T2)

def reduction_fwd_fn(inputs: list):  # in_tensor, reduction_axis
    return torch.sum(inputs[0], dim=inputs[1])

@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_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)]

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

    if not disable_validation:
        eager_output = torch.sum(inputs[0].to(torch.double), dim=reduction_axis)
        fd.validate(inputs, [eager_output.to(dtype)])

    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", [0, 1])
@pytest.mark.reduction
def test_reduction_baseline_benchmark(
    benchmark,
    size: tuple,
    dtype: torch.dtype,
    reduction_axis: int,
    executor: str,
):
    if executor == "torchcompile":
        clear_dynamo_cache()
    input = torch.randn(size, device="cuda", dtype=dtype)

    benchmark_fn = with_executor(executor, reduction_fwd_fn)

    # Inputs and outputs are same as nvFuser, no need for manual IOByte computation
    run_benchmark(
        benchmark,
        benchmark_fn,
        [input, 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, DEFAULT_EXECUTORS
import torch
from .global_params import generate_input_sizes, FLOAT_DTYPES, PROMOTE_DTYPES

def reduction_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.sum(T0, dims=[reduction_axis], keepdim=False)
    if dtype in PROMOTE_DTYPES:
        T2 = fd.ops.cast(T2, dtype=dtype)
    fd.add_output(T2)

def reduction_fwd_fn(inputs: list):  # in_tensor, reduction_axis
    return torch.sum(inputs[0], dim=inputs[1])

@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_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)]


import torch
from .global_params import generate_input_sizes, FLOAT_DTYPES, PROMOTE_DTYPES

def reduction_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.sum(T0, dims=[reduction_axis], keepdim=False)
    if dtype in PROMOTE_DTYPES:
        T2 = fd.ops.cast(T2, dtype=dtype)
    fd.add_output(T2)

def reduction_fwd_fn(inputs: list):  # in_tensor, reduction_axis
    return torch.sum(inputs[0], dim=inputs[1])

@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_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)]

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

    if not disable_validation:


) -> 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.sum(T0, dims=[reduction_axis], keepdim=False)
    if dtype in PROMOTE_DTYPES:
        T2 = fd.ops.cast(T2, dtype=dtype)
    fd.add_output(T2)

def reduction_fwd_fn(inputs: list):  # in_tensor, reduction_axis
    return torch.sum(inputs[0], dim=inputs[1])

@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_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)]

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

    if not disable_validation:
        eager_output = torch.sum(inputs[0].to(torch.double), dim=reduction_axis)
        fd.validate(inputs, [eager_output.to(dtype)])

    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", [0, 1])
@pytest.mark.reduction
def test_reduction_baseline_benchmark(
    benchmark,
    size: tuple,
    dtype: torch.dtype,
    reduction_axis: int,
    executor: str,
):
    if executor == "torchcompile":
        clear_dynamo_cache()
    input = torch.randn(size, device="cuda", dtype=dtype)

    benchmark_fn = with_executor(executor, reduction_fwd_fn)


    disable_benchmarking: bool,
):
    inputs = [torch.randn(*size, device="cuda", dtype=dtype)]

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

    if not disable_validation:
        eager_output = torch.sum(inputs[0].to(torch.double), dim=reduction_axis)
        fd.validate(inputs, [eager_output.to(dtype)])

    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", [0, 1])
@pytest.mark.reduction
def test_reduction_baseline_benchmark(
    benchmark,
    size: tuple,
    dtype: torch.dtype,
    reduction_axis: int,
    executor: str,
):
    if executor == "torchcompile":
        clear_dynamo_cache()
    input = torch.randn(size, device="cuda", dtype=dtype)

    benchmark_fn = with_executor(executor, reduction_fwd_fn)

    # Inputs and outputs are same as nvFuser, no need for manual IOByte computation
    run_benchmark(
        benchmark,
        benchmark_fn,
        [input, 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, DEFAULT_EXECUTORS
import torch
from .global_params import generate_input_sizes, FLOAT_DTYPES, PROMOTE_DTYPES

def get_n_groups(C):
    # start num_groups from 1 and increase to 32 at max
    # 32 is a widely used value for num_groups
    # it doesn't make sense to use num_groups > C
    num_groups = 1
    while num_groups * 2 <= 32 and C % (num_groups * 2) == 0:
        num_groups *= 2
    return num_groups

# This version is based on requires_grad = False.
# When requires_grad = True, two additional tensors (T4, T14) are added to outputs.
def groupnorm_fwd_fusion(
    fd: FusionDefinition,
    dtype: DataType,
    n_groups: int = 32,
    eps: float = 1e-5,
) -> None:
    # inputs, T0-x, T1-weight, T2-bias
    T0 = fd.define_tensor(
        shape=[-1, -1, -1, -1],
        contiguity=[True, True, True, True],
        dtype=dtype,
        is_cpu=False,
    )
    T1 = fd.define_tensor(shape=[-1], contiguity=[True], dtype=dtype, is_cpu=False)
    T2 = fd.define_tensor(shape=[-1], contiguity=[True], dtype=dtype, is_cpu=False)

    # initial input shape: [N, C, H, W]
    # reshape input to [N, n_groups, C//n_groups, H, W] and do normalization with scale and bias
    # reshape back to [N, C, H, W]
    V0 = T0.shape()
    G0 = fd.define_scalar(n_groups, dtype=DataType.Int)
    C0 = fd.ops.div(T0.size(1), G0)
    T0 = fd.ops.reshape(
        T0, new_shape=[T0.size(0), n_groups, C0, T0.size(2), T0.size(3)]
    )
    if dtype in PROMOTE_DTYPES:
        T0 = fd.ops.cast(T0, dtype=DataType.Float)
        T1 = fd.ops.cast(T1, dtype=DataType.Float)
        T2 = fd.ops.cast(T2, dtype=DataType.Float)

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

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

    S12 = fd.define_scalar(eps, dtype=DataType.Double)
    T13 = fd.ops.add(T7, S12)
    T14 = fd.ops.rsqrt(T13)
    # N, G, C//G, H, W
    V3 = T0.shape()
    T18 = fd.ops.broadcast_in_dim(T11, shape=V3, broadcast_dims=[0, 1, 2, 3, 4])
    T19 = fd.ops.sub(T0, T18)
    T23 = fd.ops.broadcast_in_dim(T14, shape=V3, broadcast_dims=[0, 1, 2, 3, 4])
    T24 = fd.ops.mul(T19, T23)

    # reshape weights and bias to [1, n_groups, C//n_groups, 1, 1]
    # due to https://github.com/NVIDIA/Fuser/issues/2671 must define C1 and C2
    # using T1.size(0) and T2.size(0), can't directly reuse C0 which is based on T0.size(1)
    C1 = fd.ops.div(T1.size(0), G0)
    T1 = fd.ops.reshape(T1, new_shape=[1, n_groups, C1, 1, 1])

    C2 = fd.ops.div(T2.size(0), G0)
    T2 = fd.ops.reshape(T2, new_shape=[1, n_groups, C2, 1, 1])

    # broadcast weights and bias to [N, n_groups, C//n_groups, H, W]
    T25 = fd.ops.broadcast_in_dim(T1, shape=V3, broadcast_dims=[0, 1, 2, 3, 4])
    T26 = fd.ops.mul(T24, T25)
    T27 = fd.ops.broadcast_in_dim(T2, shape=V3, broadcast_dims=[0, 1, 2, 3, 4])
    T28 = fd.ops.add(T26, T27)

    # convert back to original dtype and shape
    if dtype in PROMOTE_DTYPES:
        T28 = fd.ops.cast(T28, dtype=dtype)
    T28 = fd.ops.reshape(T28, new_shape=V0)
    fd.add_output(T28)

def groupnorm_fwd(inputs: list):  # [in_tensor, weights, bias, n_groups]
    return torch.nn.functional.group_norm(
        inputs[0], num_groups=inputs[3], weight=inputs[1], bias=inputs[2], eps=1e-05
    )

@pytest.mark.parametrize("size", generate_input_sizes(dims=4))
@pytest.mark.parametrize("dtype", FLOAT_DTYPES)
@pytest.mark.inner_persistent
def test_groupnorm_fwd_nvf_benchmark(
    benchmark,
    size: tuple,
    dtype: torch.dtype,
    disable_validation: bool,
    disable_benchmarking: bool,
    eps: float = 1e-5,
):
    N, C, H, W = size
    x = torch.randn(size, device="cuda", dtype=dtype)
    weight = torch.randn(C, device="cuda", dtype=dtype)
    bias = torch.randn(C, device="cuda", dtype=dtype)
    num_groups = get_n_groups(C)

    with FusionDefinition() as fd:
        groupnorm_fwd_fusion(fd, torch_dtype_to_nvfuser_dtype(dtype), num_groups)

    if not disable_validation:
        eager_output = groupnorm_fwd([x, weight, bias, num_groups])
        fd.validate([x, weight, bias], [eager_output])

    if not disable_benchmarking:
        run_benchmark(benchmark, fd.execute, [x, weight, bias])

@pytest.mark.parametrize("executor", DEFAULT_EXECUTORS)
@pytest.mark.parametrize("size", generate_input_sizes(dims=4))
@pytest.mark.parametrize("dtype", FLOAT_DTYPES)
@pytest.mark.inner_persistent
def test_groupnorm_fwd_baseline_benchmark(
    benchmark,
    size: tuple,
    dtype: torch.dtype,
    executor: str,
):
    if executor == "torchcompile":
        clear_dynamo_cache()
    N, C, H, W = size
    x = torch.randn(size, device="cuda", dtype=dtype)
    weight = torch.randn(C, device="cuda", dtype=dtype)
    bias = torch.randn(C, device="cuda", dtype=dtype)
    num_groups = get_n_groups(C)

    benchmark_fn = with_executor(executor, groupnorm_fwd)

    run_benchmark(
        benchmark,
        benchmark_fn,
        [x, weight, bias, num_groups],
    )


# 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

def get_n_groups(C):
    # start num_groups from 1 and increase to 32 at max
    # 32 is a widely used value for num_groups
    # it doesn't make sense to use num_groups > C
    num_groups = 1
    while num_groups * 2 <= 32 and C % (num_groups * 2) == 0:
        num_groups *= 2
    return num_groups

# This version is based on requires_grad = False.
# When requires_grad = True, two additional tensors (T4, T14) are added to outputs.
def groupnorm_fwd_fusion(
    fd: FusionDefinition,
    dtype: DataType,
    n_groups: int = 32,
    eps: float = 1e-5,
) -> None:
    # inputs, T0-x, T1-weight, T2-bias
    T0 = fd.define_tensor(
        shape=[-1, -1, -1, -1],
        contiguity=[True, True, True, True],
        dtype=dtype,
        is_cpu=False,
    )
    T1 = fd.define_tensor(shape=[-1], contiguity=[True], dtype=dtype, is_cpu=False)
    T2 = fd.define_tensor(shape=[-1], contiguity=[True], dtype=dtype, is_cpu=False)


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

def get_n_groups(C):
    # start num_groups from 1 and increase to 32 at max
    # 32 is a widely used value for num_groups
    # it doesn't make sense to use num_groups > C
    num_groups = 1
    while num_groups * 2 <= 32 and C % (num_groups * 2) == 0:
        num_groups *= 2
    return num_groups

# This version is based on requires_grad = False.
# When requires_grad = True, two additional tensors (T4, T14) are added to outputs.
def groupnorm_fwd_fusion(
    fd: FusionDefinition,
    dtype: DataType,
    n_groups: int = 32,
    eps: float = 1e-5,
) -> None:
    # inputs, T0-x, T1-weight, T2-bias
    T0 = fd.define_tensor(
        shape=[-1, -1, -1, -1],
        contiguity=[True, True, True, True],
        dtype=dtype,
        is_cpu=False,
    )
    T1 = fd.define_tensor(shape=[-1], contiguity=[True], dtype=dtype, is_cpu=False)
    T2 = fd.define_tensor(shape=[-1], contiguity=[True], dtype=dtype, is_cpu=False)

    # initial input shape: [N, C, H, W]
    # reshape input to [N, n_groups, C//n_groups, H, W] and do normalization with scale and bias
    # reshape back to [N, C, H, W]
    V0 = T0.shape()
    G0 = fd.define_scalar(n_groups, dtype=DataType.Int)
    C0 = fd.ops.div(T0.size(1), G0)
    T0 = fd.ops.reshape(
        T0, new_shape=[T0.size(0), n_groups, C0, T0.size(2), T0.size(3)]
    )
    if dtype in PROMOTE_DTYPES:
        T0 = fd.ops.cast(T0, dtype=DataType.Float)
        T1 = fd.ops.cast(T1, dtype=DataType.Float)
        T2 = fd.ops.cast(T2, dtype=DataType.Float)

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

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

    S12 = fd.define_scalar(eps, dtype=DataType.Double)
    T13 = fd.ops.add(T7, S12)
    T14 = fd.ops.rsqrt(T13)
    # N, G, C//G, H, W
    V3 = T0.shape()
    T18 = fd.ops.broadcast_in_dim(T11, shape=V3, broadcast_dims=[0, 1, 2, 3, 4])
    T19 = fd.ops.sub(T0, T18)
    T23 = fd.ops.broadcast_in_dim(T14, shape=V3, broadcast_dims=[0, 1, 2, 3, 4])
    T24 = fd.ops.mul(T19, T23)

    # reshape weights and bias to [1, n_groups, C//n_groups, 1, 1]
    # due to https://github.com/NVIDIA/Fuser/issues/2671 must define C1 and C2
    # using T1.size(0) and T2.size(0), can't directly reuse C0 which is based on T0.size(1)
    C1 = fd.ops.div(T1.size(0), G0)
    T1 = fd.ops.reshape(T1, new_shape=[1, n_groups, C1, 1, 1])

    C2 = fd.ops.div(T2.size(0), G0)
    T2 = fd.ops.reshape(T2, new_shape=[1, n_groups, C2, 1, 1])

    # broadcast weights and bias to [N, n_groups, C//n_groups, H, W]
    T25 = fd.ops.broadcast_in_dim(T1, shape=V3, broadcast_dims=[0, 1, 2, 3, 4])
    T26 = fd.ops.mul(T24, T25)
    T27 = fd.ops.broadcast_in_dim(T2, shape=V3, broadcast_dims=[0, 1, 2, 3, 4])
    T28 = fd.ops.add(T26, T27)

    # convert back to original dtype and shape
    if dtype in PROMOTE_DTYPES:
        T28 = fd.ops.cast(T28, dtype=dtype)
    T28 = fd.ops.reshape(T28, new_shape=V0)
    fd.add_output(T28)

def groupnorm_fwd(inputs: list):  # [in_tensor, weights, bias, n_groups]
    return torch.nn.functional.group_norm(
        inputs[0], num_groups=inputs[3], weight=inputs[1], bias=inputs[2], eps=1e-05
    )

@pytest.mark.parametrize("size", generate_input_sizes(dims=4))
@pytest.mark.parametrize("dtype", FLOAT_DTYPES)
@pytest.mark.inner_persistent
def test_groupnorm_fwd_nvf_benchmark(
    benchmark,
    size: tuple,
    dtype: torch.dtype,
    disable_validation: bool,
    disable_benchmarking: bool,
    eps: float = 1e-5,
):
    N, C, H, W = size


    # using T1.size(0) and T2.size(0), can't directly reuse C0 which is based on T0.size(1)
    C1 = fd.ops.div(T1.size(0), G0)
    T1 = fd.ops.reshape(T1, new_shape=[1, n_groups, C1, 1, 1])

    C2 = fd.ops.div(T2.size(0), G0)
    T2 = fd.ops.reshape(T2, new_shape=[1, n_groups, C2, 1, 1])

    # broadcast weights and bias to [N, n_groups, C//n_groups, H, W]
    T25 = fd.ops.broadcast_in_dim(T1, shape=V3, broadcast_dims=[0, 1, 2, 3, 4])
    T26 = fd.ops.mul(T24, T25)
    T27 = fd.ops.broadcast_in_dim(T2, shape=V3, broadcast_dims=[0, 1, 2, 3, 4])
    T28 = fd.ops.add(T26, T27)

    # convert back to original dtype and shape
    if dtype in PROMOTE_DTYPES:
        T28 = fd.ops.cast(T28, dtype=dtype)
    T28 = fd.ops.reshape(T28, new_shape=V0)
    fd.add_output(T28)

def groupnorm_fwd(inputs: list):  # [in_tensor, weights, bias, n_groups]
    return torch.nn.functional.group_norm(
        inputs[0], num_groups=inputs[3], weight=inputs[1], bias=inputs[2], eps=1e-05
    )

@pytest.mark.parametrize("size", generate_input_sizes(dims=4))
@pytest.mark.parametrize("dtype", FLOAT_DTYPES)
@pytest.mark.inner_persistent
def test_groupnorm_fwd_nvf_benchmark(
    benchmark,
    size: tuple,
    dtype: torch.dtype,
    disable_validation: bool,
    disable_benchmarking: bool,
    eps: float = 1e-5,
):
    N, C, H, W = size
    x = torch.randn(size, device="cuda", dtype=dtype)
    weight = torch.randn(C, device="cuda", dtype=dtype)
    bias = torch.randn(C, device="cuda", dtype=dtype)
    num_groups = get_n_groups(C)

    with FusionDefinition() as fd:


    T26 = fd.ops.mul(T24, T25)
    T27 = fd.ops.broadcast_in_dim(T2, shape=V3, broadcast_dims=[0, 1, 2, 3, 4])
    T28 = fd.ops.add(T26, T27)

    # convert back to original dtype and shape
    if dtype in PROMOTE_DTYPES:
        T28 = fd.ops.cast(T28, dtype=dtype)
    T28 = fd.ops.reshape(T28, new_shape=V0)
    fd.add_output(T28)

def groupnorm_fwd(inputs: list):  # [in_tensor, weights, bias, n_groups]
    return torch.nn.functional.group_norm(
        inputs[0], num_groups=inputs[3], weight=inputs[1], bias=inputs[2], eps=1e-05
    )

@pytest.mark.parametrize("size", generate_input_sizes(dims=4))
@pytest.mark.parametrize("dtype", FLOAT_DTYPES)
@pytest.mark.inner_persistent
def test_groupnorm_fwd_nvf_benchmark(
    benchmark,
    size: tuple,
    dtype: torch.dtype,
    disable_validation: bool,
    disable_benchmarking: bool,
    eps: float = 1e-5,
):
    N, C, H, W = size
    x = torch.randn(size, device="cuda", dtype=dtype)
    weight = torch.randn(C, device="cuda", dtype=dtype)
    bias = torch.randn(C, device="cuda", dtype=dtype)
    num_groups = get_n_groups(C)

    with FusionDefinition() as fd:
        groupnorm_fwd_fusion(fd, torch_dtype_to_nvfuser_dtype(dtype), num_groups)

    if not disable_validation:
        eager_output = groupnorm_fwd([x, weight, bias, num_groups])
        fd.validate([x, weight, bias], [eager_output])

    if not disable_benchmarking:
        run_benchmark(benchmark, fd.execute, [x, weight, bias])

@pytest.mark.parametrize("executor", DEFAULT_EXECUTORS)
@pytest.mark.parametrize("size", generate_input_sizes(dims=4))
@pytest.mark.parametrize("dtype", FLOAT_DTYPES)
@pytest.mark.inner_persistent
def test_groupnorm_fwd_baseline_benchmark(
    benchmark,
    size: tuple,
    dtype: torch.dtype,
    executor: str,
):
    if executor == "torchcompile":
        clear_dynamo_cache()
    N, C, H, W = size
    x = torch.randn(size, device="cuda", dtype=dtype)
    weight = torch.randn(C, device="cuda", dtype=dtype)
    bias = torch.randn(C, device="cuda", dtype=dtype)
    num_groups = get_n_groups(C)


    x = torch.randn(size, device="cuda", dtype=dtype)
    weight = torch.randn(C, device="cuda", dtype=dtype)
    bias = torch.randn(C, device="cuda", dtype=dtype)
    num_groups = get_n_groups(C)

    with FusionDefinition() as fd:
        groupnorm_fwd_fusion(fd, torch_dtype_to_nvfuser_dtype(dtype), num_groups)

    if not disable_validation:
        eager_output = groupnorm_fwd([x, weight, bias, num_groups])
        fd.validate([x, weight, bias], [eager_output])

    if not disable_benchmarking:
        run_benchmark(benchmark, fd.execute, [x, weight, bias])

@pytest.mark.parametrize("executor", DEFAULT_EXECUTORS)
@pytest.mark.parametrize("size", generate_input_sizes(dims=4))
@pytest.mark.parametrize("dtype", FLOAT_DTYPES)
@pytest.mark.inner_persistent
def test_groupnorm_fwd_baseline_benchmark(
    benchmark,
    size: tuple,
    dtype: torch.dtype,
    executor: str,
):
    if executor == "torchcompile":
        clear_dynamo_cache()
    N, C, H, W = size
    x = torch.randn(size, device="cuda", dtype=dtype)
    weight = torch.randn(C, device="cuda", dtype=dtype)
    bias = torch.randn(C, device="cuda", dtype=dtype)
    num_groups = get_n_groups(C)

    benchmark_fn = with_executor(executor, groupnorm_fwd)

    run_benchmark(
        benchmark,
        benchmark_fn,
        [x, weight, bias, num_groups],
    )


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

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

    Fusion inputs: input, dropout_mask, rms, grad_output, weights
    Fusion outputs: grad_input, grad_weights
    """
    T5 = fd.define_tensor(
        shape=[-1, -1], contiguity=[True, True], dtype=dtype, is_cpu=False
    )  # input1
    T4 = fd.define_tensor(
        shape=[-1, -1], contiguity=[True, True], dtype=dtype, is_cpu=False
    )  # input2
    T6 = fd.define_tensor(
        shape=[-1, -1], contiguity=[True, True], dtype=DataType.Bool, is_cpu=False
    )  # dropout_mask
    T7 = fd.define_tensor(
        shape=[-1, 1], contiguity=[True, None], dtype=DataType.Float, is_cpu=False
    )  # rms_eps
    T8 = fd.define_tensor(
        shape=[-1, -1], contiguity=[True, True], dtype=dtype, is_cpu=False
    )  # grads
    T9 = fd.define_tensor(
        shape=[-1], contiguity=[True], dtype=dtype, is_cpu=False
    )  # weights

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

    T12 = fd.ops.mul(T5, T6)
    S13 = fd.define_scalar(1 / (1 - dropout_p), dtype=DataType.Double)
    T14 = fd.ops.mul(T12, S13)
    T15 = fd.ops.add(T4, T14)

    V19 = T5.shape()
    T20 = fd.ops.broadcast_in_dim(T7, shape=V19, broadcast_dims=[0, 1])
    T22 = fd.ops.reciprocal(T20)
    T23 = fd.ops.mul(T15, T22)

    T27 = fd.ops.broadcast_in_dim(T9, shape=V19, broadcast_dims=[1])

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

    T35 = fd.ops.mul(T31, T22)
    T36 = fd.ops.neg(T31)
    T37 = fd.ops.mul(T36, T15)
    S38 = fd.define_scalar(2.00000, dtype=DataType.Double)
    T39 = fd.ops.pow(T20, S38)
    T40 = fd.ops.reciprocal(T39)
    T41 = fd.ops.mul(T37, T40)
    T42 = fd.ops.sum(T41, dims=[1], keepdim=False, dtype=DataType.Null)

    T47 = fd.ops.broadcast_in_dim(T42, shape=[T5.size(0), 1], broadcast_dims=[0])

    T50 = fd.ops.mul(S38, T7)
    T51 = fd.ops.reciprocal(T50)
    T52 = fd.ops.mul(T47, T51)
    S55 = fd.ops.reciprocal(T5.size(1))
    T56 = fd.ops.mul(T52, S55)
    T57 = fd.ops.sum(T56, dims=[1], keepdim=False, dtype=DataType.Null)

    T61 = fd.ops.broadcast_in_dim(T57, shape=[T5.size(0), 1], broadcast_dims=[0])
    T65 = fd.ops.broadcast_in_dim(T61, shape=V19, broadcast_dims=[0, 1])
    T66 = fd.ops.mul(T65, S38)
    T69 = fd.ops.mul(T66, T15)
    T70 = fd.ops.add(T35, T69)

    T73 = fd.ops.mul(T70, S13)
    T74 = fd.ops.mul(T73, T6)

    if dtype in PROMOTE_DTYPES:
        T70 = fd.ops.cast(T70, dtype=dtype)
        T74 = fd.ops.cast(T74, dtype=dtype)
        T32 = fd.ops.cast(T32, dtype=dtype)

    fd.add_output(T74)
    fd.add_output(T70)
    fd.add_output(T32)

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

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

    dropout_p = 0.2
    dropout_mask = torch.lt(torch.rand(*size, device="cuda"), 1 - dropout_p)

    # Manually compute dropout for validation
    x = input2 + 1 / (1 - dropout_p) * dropout_mask * input1
    squared_mean = (x.to(torch.float) ** 2).mean(1, keepdim=True)
    rms_eps = torch.sqrt(squared_mean + eps)

    with FusionDefinition() as fd:
        dropout_rmsnorm_bwd_fusion(fd, torch_dtype_to_nvfuser_dtype(dtype), dropout_p)

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

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

@pytest.mark.parametrize("executor", DEFAULT_EXECUTORS)
@pytest.mark.parametrize("size", generate_input_sizes(dims=2))
@pytest.mark.parametrize("dtype", FLOAT_DTYPES)
@pytest.mark.inner_outer_persistent
@pytest.mark.inner_persistent
def test_dropout_rmsnorm_bwd_baseline_benchmark(
    benchmark,
    size: tuple,
    dtype: torch.dtype,
    executor: str,
):
    if executor == "torchcompile":
        clear_dynamo_cache()
    dropout_p = 0.2
    input1 = torch.randn(size, device="cuda", dtype=dtype, requires_grad=True)
    input2 = torch.randn(size, device="cuda", dtype=dtype, requires_grad=True)
    grads = torch.randn(size, device="cuda", dtype=dtype)
    weights = torch.randn(size[1], device="cuda", dtype=dtype, requires_grad=True)

    fwd_fn = with_executor(executor, dropout_rmsnorm)
    fwd_inputs = [input1, input2, weights, dropout_p]
    outputs = fwd_fn(fwd_inputs)

    run_benchmark(
        benchmark,
        unary_bwd_torch,
        [outputs, grads, *fwd_inputs],
        iobytes=dropout_rmsnorm_bwd_iobytes(size, dtype),
    )


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

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

    Fusion inputs: input, dropout_mask, rms, grad_output, weights
    Fusion outputs: grad_input, grad_weights
    """
    T5 = fd.define_tensor(
        shape=[-1, -1], contiguity=[True, True], dtype=dtype, is_cpu=False
    )  # input1
    T4 = fd.define_tensor(
        shape=[-1, -1], contiguity=[True, True], dtype=dtype, is_cpu=False
    )  # input2
    T6 = fd.define_tensor(
        shape=[-1, -1], contiguity=[True, True], dtype=DataType.Bool, is_cpu=False
    )  # dropout_mask
    T7 = fd.define_tensor(
        shape=[-1, 1], contiguity=[True, None], dtype=DataType.Float, is_cpu=False
    )  # rms_eps
    T8 = fd.define_tensor(
        shape=[-1, -1], contiguity=[True, True], dtype=dtype, is_cpu=False
    )  # grads
    T9 = fd.define_tensor(
        shape=[-1], contiguity=[True], dtype=dtype, is_cpu=False
    )  # weights

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

    T12 = fd.ops.mul(T5, T6)
    S13 = fd.define_scalar(1 / (1 - dropout_p), dtype=DataType.Double)
    T14 = fd.ops.mul(T12, S13)
    T15 = fd.ops.add(T4, T14)

    V19 = T5.shape()
    T20 = fd.ops.broadcast_in_dim(T7, shape=V19, broadcast_dims=[0, 1])
    T22 = fd.ops.reciprocal(T20)
    T23 = fd.ops.mul(T15, T22)

    T27 = fd.ops.broadcast_in_dim(T9, shape=V19, broadcast_dims=[1])

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

    T35 = fd.ops.mul(T31, T22)
    T36 = fd.ops.neg(T31)
    T37 = fd.ops.mul(T36, T15)
    S38 = fd.define_scalar(2.00000, dtype=DataType.Double)
    T39 = fd.ops.pow(T20, S38)
    T40 = fd.ops.reciprocal(T39)
    T41 = fd.ops.mul(T37, T40)
    T42 = fd.ops.sum(T41, dims=[1], keepdim=False, dtype=DataType.Null)

    T47 = fd.ops.broadcast_in_dim(T42, shape=[T5.size(0), 1], broadcast_dims=[0])

    T50 = fd.ops.mul(S38, T7)
    T51 = fd.ops.reciprocal(T50)
    T52 = fd.ops.mul(T47, T51)
    S55 = fd.ops.reciprocal(T5.size(1))
    T56 = fd.ops.mul(T52, S55)
    T57 = fd.ops.sum(T56, dims=[1], keepdim=False, dtype=DataType.Null)

    T61 = fd.ops.broadcast_in_dim(T57, shape=[T5.size(0), 1], broadcast_dims=[0])
    T65 = fd.ops.broadcast_in_dim(T61, shape=V19, broadcast_dims=[0, 1])
    T66 = fd.ops.mul(T65, S38)
    T69 = fd.ops.mul(T66, T15)
    T70 = fd.ops.add(T35, T69)

    T73 = fd.ops.mul(T70, S13)
    T74 = fd.ops.mul(T73, T6)

    if dtype in PROMOTE_DTYPES:
        T70 = fd.ops.cast(T70, dtype=dtype)
        T74 = fd.ops.cast(T74, dtype=dtype)
        T32 = fd.ops.cast(T32, dtype=dtype)

    fd.add_output(T74)
    fd.add_output(T70)
    fd.add_output(T32)

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


    T61 = fd.ops.broadcast_in_dim(T57, shape=[T5.size(0), 1], broadcast_dims=[0])
    T65 = fd.ops.broadcast_in_dim(T61, shape=V19, broadcast_dims=[0, 1])
    T66 = fd.ops.mul(T65, S38)
    T69 = fd.ops.mul(T66, T15)
    T70 = fd.ops.add(T35, T69)

    T73 = fd.ops.mul(T70, S13)
    T74 = fd.ops.mul(T73, T6)

    if dtype in PROMOTE_DTYPES:
        T70 = fd.ops.cast(T70, dtype=dtype)
        T74 = fd.ops.cast(T74, dtype=dtype)
        T32 = fd.ops.cast(T32, dtype=dtype)

    fd.add_output(T74)
    fd.add_output(T70)
    fd.add_output(T32)

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

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

    dropout_p = 0.2


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

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

    dropout_p = 0.2
    dropout_mask = torch.lt(torch.rand(*size, device="cuda"), 1 - dropout_p)

    # Manually compute dropout for validation
    x = input2 + 1 / (1 - dropout_p) * dropout_mask * input1
    squared_mean = (x.to(torch.float) ** 2).mean(1, keepdim=True)
    rms_eps = torch.sqrt(squared_mean + eps)

    with FusionDefinition() as fd:
        dropout_rmsnorm_bwd_fusion(fd, torch_dtype_to_nvfuser_dtype(dtype), dropout_p)

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

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

@pytest.mark.parametrize("executor", DEFAULT_EXECUTORS)
@pytest.mark.parametrize("size", generate_input_sizes(dims=2))
@pytest.mark.parametrize("dtype", FLOAT_DTYPES)
@pytest.mark.inner_outer_persistent
@pytest.mark.inner_persistent
def test_dropout_rmsnorm_bwd_baseline_benchmark(
    benchmark,
    size: tuple,
    dtype: torch.dtype,
    executor: str,
):
    if executor == "torchcompile":
        clear_dynamo_cache()
    dropout_p = 0.2
    input1 = torch.randn(size, device="cuda", dtype=dtype, requires_grad=True)
    input2 = torch.randn(size, device="cuda", dtype=dtype, requires_grad=True)
    grads = torch.randn(size, device="cuda", dtype=dtype)
    weights = torch.randn(size[1], device="cuda", dtype=dtype, requires_grad=True)


        )
        eager_output.backward(grads.to(torch.double))
        fd.validate(
            [input1, input2, dropout_mask, rms_eps, grads, weights],
            [input1.grad, input2.grad, weights.grad],
        )

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

@pytest.mark.parametrize("executor", DEFAULT_EXECUTORS)
@pytest.mark.parametrize("size", generate_input_sizes(dims=2))
@pytest.mark.parametrize("dtype", FLOAT_DTYPES)
@pytest.mark.inner_outer_persistent
@pytest.mark.inner_persistent
def test_dropout_rmsnorm_bwd_baseline_benchmark(
    benchmark,
    size: tuple,
    dtype: torch.dtype,
    executor: str,
):
    if executor == "torchcompile":
        clear_dynamo_cache()
    dropout_p = 0.2
    input1 = torch.randn(size, device="cuda", dtype=dtype, requires_grad=True)
    input2 = torch.randn(size, device="cuda", dtype=dtype, requires_grad=True)
    grads = torch.randn(size, device="cuda", dtype=dtype)
    weights = torch.randn(size[1], device="cuda", dtype=dtype, requires_grad=True)

    fwd_fn = with_executor(executor, dropout_rmsnorm)
    fwd_inputs = [input1, input2, weights, dropout_p]
    outputs = fwd_fn(fwd_inputs)

    run_benchmark(
        benchmark,
        unary_bwd_torch,
        [outputs, grads, *fwd_inputs],
        iobytes=dropout_rmsnorm_bwd_iobytes(size, dtype),
    )


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

import os
import pytest
import torch
import nvfuser_direct as nvfuser
from nvfuser_direct import FusionDefinition, CommunicatorBackend
from nvfuser_direct.pytorch_utils import torch_dtype_to_nvfuser_dtype
from .core import BENCHMARK_CONFIG, clear_l2_cache

# Run command:
# export NVFUSER_BUILD_WITH_UCC=1
# pip install --no-build-isolation -e python -v
# mpirun -np 1 pytest benchmarks/python/benchmark_overlap.py --with-mpi

class CUDAEventTimer:
    """Custom CUDA event-based timer for accurate GPU timing.

    This timer uses CUDA events to measure elapsed time between operations,
    providing more accurate GPU timing than CPU-based timing methods.
    """

    def __init__(self):
        self.start_event = torch.cuda.Event(enable_timing=True)
        self.end_event = torch.cuda.Event(enable_timing=True)
        self.is_running = False

    def __call__(self):
        """Record timing events and compute elapsed time.

        Returns:
            float: Elapsed time in seconds
        """
        if self.is_running:
            self.end_event.record()
            torch.cuda.synchronize()
            elapsed_ms = self.start_event.elapsed_time(self.end_event)
            self.is_running = False
            return elapsed_ms / 1000.0  # Convert ms to seconds
        else:
            self.start_event.record()
            self.is_running = True
            return 0.0

    def cleanup(self):
        """Ensure timer is not running and synchronize CUDA."""
        if self.is_running:
            self.end_event.record()
            torch.cuda.synchronize()
            self.is_running = False

def benchmark_cuda_events_pedantic(
    benchmark, benchmark_fn, inputs, rounds, warmup_rounds
):
    """Wrapper for benchmark_cuda_events that uses pytest's pedantic method with CUDA events.

    Args:
        benchmark: pytest-benchmark fixture
        benchmark_fn: Function to benchmark
        inputs: List of inputs to pass to benchmark_fn
        rounds: Number of rounds to run
        warmup_rounds: Number of warmup rounds
    """

    def setup():
        clear_l2_cache()
        return inputs, {}

    def wrapped_fn(*args):
        benchmark_fn(*args[0])
        return None

    # Set our custom CUDA event timer
    benchmark._timer = CUDAEventTimer()

    benchmark.pedantic(
        wrapped_fn,
        setup=setup,
        rounds=rounds,
        warmup_rounds=warmup_rounds,
        iterations=1,
    )

class MultideviceSettings:
    """Settings and utilities for multi-device execution."""

    def __init__(self):
        self._communicator = nvfuser.multidevice.Communicator.instance()
        torch.manual_seed(0)

    @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,
    s,
    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(
        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)
        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"],
    )


CUDAEventTimer

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

import os
import pytest
import torch
import nvfuser_direct as nvfuser
from nvfuser_direct import FusionDefinition, CommunicatorBackend
from nvfuser_direct.pytorch_utils import torch_dtype_to_nvfuser_dtype
from .core import BENCHMARK_CONFIG, clear_l2_cache

# Run command:
# export NVFUSER_BUILD_WITH_UCC=1
# pip install --no-build-isolation -e python -v
# mpirun -np 1 pytest benchmarks/python/benchmark_overlap.py --with-mpi

class CUDAEventTimer:
    """Custom CUDA event-based timer for accurate GPU timing.

    This timer uses CUDA events to measure elapsed time between operations,
    providing more accurate GPU timing than CPU-based timing methods.
    """

    def __init__(self):
        self.start_event = torch.cuda.Event(enable_timing=True)
        self.end_event = torch.cuda.Event(enable_timing=True)
        self.is_running = False

    def __call__(self):
        """Record timing events and compute elapsed time.

        Returns:
            float: Elapsed time in seconds
        """
        if self.is_running:
            self.end_event.record()
            torch.cuda.synchronize()
            elapsed_ms = self.start_event.elapsed_time(self.end_event)
            self.is_running = False
            return elapsed_ms / 1000.0  # Convert ms to seconds
        else:
            self.start_event.record()
            self.is_running = True
            return 0.0

    def cleanup(self):
        """Ensure timer is not running and synchronize CUDA."""
        if self.is_running:
            self.end_event.record()
            torch.cuda.synchronize()
            self.is_running = False

def benchmark_cuda_events_pedantic(
    benchmark, benchmark_fn, inputs, rounds, warmup_rounds
):
    """Wrapper for benchmark_cuda_events that uses pytest's pedantic method with CUDA events.

    Args:
        benchmark: pytest-benchmark fixture
        benchmark_fn: Function to benchmark
        inputs: List of inputs to pass to benchmark_fn
        rounds: Number of rounds to run
        warmup_rounds: Number of warmup rounds
    """

    def setup():
        clear_l2_cache()
        return inputs, {}

    def wrapped_fn(*args):


        """
        if self.is_running:
            self.end_event.record()
            torch.cuda.synchronize()
            elapsed_ms = self.start_event.elapsed_time(self.end_event)
            self.is_running = False
            return elapsed_ms / 1000.0  # Convert ms to seconds
        else:
            self.start_event.record()
            self.is_running = True
            return 0.0

    def cleanup(self):
        """Ensure timer is not running and synchronize CUDA."""
        if self.is_running:
            self.end_event.record()
            torch.cuda.synchronize()
            self.is_running = False

def benchmark_cuda_events_pedantic(
    benchmark, benchmark_fn, inputs, rounds, warmup_rounds
):
    """Wrapper for benchmark_cuda_events that uses pytest's pedantic method with CUDA events.

    Args:
        benchmark: pytest-benchmark fixture
        benchmark_fn: Function to benchmark
        inputs: List of inputs to pass to benchmark_fn
        rounds: Number of rounds to run
        warmup_rounds: Number of warmup rounds
    """

    def setup():
        clear_l2_cache()
        return inputs, {}

    def wrapped_fn(*args):
        benchmark_fn(*args[0])
        return None

    # Set our custom CUDA event timer
    benchmark._timer = CUDAEventTimer()

    benchmark.pedantic(
        wrapped_fn,
        setup=setup,
        rounds=rounds,
        warmup_rounds=warmup_rounds,
        iterations=1,
    )

class MultideviceSettings:
    """Settings and utilities for multi-device execution."""

    def __init__(self):
        self._communicator = nvfuser.multidevice.Communicator.instance()
        torch.manual_seed(0)

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

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

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


MultideviceSettings

    def setup():
        clear_l2_cache()
        return inputs, {}

    def wrapped_fn(*args):
        benchmark_fn(*args[0])
        return None

    # Set our custom CUDA event timer
    benchmark._timer = CUDAEventTimer()

    benchmark.pedantic(
        wrapped_fn,
        setup=setup,
        rounds=rounds,
        warmup_rounds=warmup_rounds,
        iterations=1,
    )

class MultideviceSettings:
    """Settings and utilities for multi-device execution."""

    def __init__(self):
        self._communicator = nvfuser.multidevice.Communicator.instance()
        torch.manual_seed(0)

    @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,


    @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,
    s,
    m,
    k,
    n,
    dtype,


            "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,
    s,
    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(
        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)
        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"],
    )


    def wrapped_fn(*args):
        benchmark_fn(*args[0])
        return None

    # Set our custom CUDA event timer
    benchmark._timer = CUDAEventTimer()

    benchmark.pedantic(
        wrapped_fn,
        setup=setup,
        rounds=rounds,
        warmup_rounds=warmup_rounds,
        iterations=1,
    )

class MultideviceSettings:
    """Settings and utilities for multi-device execution."""

    def __init__(self):
        self._communicator = nvfuser.multidevice.Communicator.instance()
        torch.manual_seed(0)

    @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()


from .core import BENCHMARK_CONFIG, clear_l2_cache

# Run command:
# export NVFUSER_BUILD_WITH_UCC=1
# pip install --no-build-isolation -e python -v
# mpirun -np 1 pytest benchmarks/python/benchmark_overlap.py --with-mpi

class CUDAEventTimer:
    """Custom CUDA event-based timer for accurate GPU timing.

    This timer uses CUDA events to measure elapsed time between operations,
    providing more accurate GPU timing than CPU-based timing methods.
    """

    def __init__(self):
        self.start_event = torch.cuda.Event(enable_timing=True)
        self.end_event = torch.cuda.Event(enable_timing=True)
        self.is_running = False

    def __call__(self):
        """Record timing events and compute elapsed time.

        Returns:
            float: Elapsed time in seconds
        """
        if self.is_running:
            self.end_event.record()
            torch.cuda.synchronize()
            elapsed_ms = self.start_event.elapsed_time(self.end_event)
            self.is_running = False
            return elapsed_ms / 1000.0  # Convert ms to seconds
        else:
            self.start_event.record()
            self.is_running = True
            return 0.0

    def cleanup(self):
        """Ensure timer is not running and synchronize CUDA."""
        if self.is_running:
            self.end_event.record()
            torch.cuda.synchronize()
            self.is_running = False

def benchmark_cuda_events_pedantic(
    benchmark, benchmark_fn, inputs, rounds, warmup_rounds
):
    """Wrapper for benchmark_cuda_events that uses pytest's pedantic method with CUDA events.

    Args:
        benchmark: pytest-benchmark fixture
        benchmark_fn: Function to benchmark
        inputs: List of inputs to pass to benchmark_fn
        rounds: Number of rounds to run
        warmup_rounds: Number of warmup rounds


cleanup

        self.end_event = torch.cuda.Event(enable_timing=True)
        self.is_running = False

    def __call__(self):
        """Record timing events and compute elapsed time.

        Returns:
            float: Elapsed time in seconds
        """
        if self.is_running:
            self.end_event.record()
            torch.cuda.synchronize()
            elapsed_ms = self.start_event.elapsed_time(self.end_event)
            self.is_running = False
            return elapsed_ms / 1000.0  # Convert ms to seconds
        else:
            self.start_event.record()
            self.is_running = True
            return 0.0

    def cleanup(self):
        """Ensure timer is not running and synchronize CUDA."""
        if self.is_running:
            self.end_event.record()
            torch.cuda.synchronize()
            self.is_running = False

def benchmark_cuda_events_pedantic(
    benchmark, benchmark_fn, inputs, rounds, warmup_rounds
):
    """Wrapper for benchmark_cuda_events that uses pytest's pedantic method with CUDA events.

    Args:
        benchmark: pytest-benchmark fixture
        benchmark_fn: Function to benchmark
        inputs: List of inputs to pass to benchmark_fn
        rounds: Number of rounds to run
        warmup_rounds: Number of warmup rounds
    """

    def setup():
        clear_l2_cache()
        return inputs, {}

    def wrapped_fn(*args):


setup

        """Ensure timer is not running and synchronize CUDA."""
        if self.is_running:
            self.end_event.record()
            torch.cuda.synchronize()
            self.is_running = False

def benchmark_cuda_events_pedantic(
    benchmark, benchmark_fn, inputs, rounds, warmup_rounds
):
    """Wrapper for benchmark_cuda_events that uses pytest's pedantic method with CUDA events.

    Args:
        benchmark: pytest-benchmark fixture
        benchmark_fn: Function to benchmark
        inputs: List of inputs to pass to benchmark_fn
        rounds: Number of rounds to run
        warmup_rounds: Number of warmup rounds
    """

    def setup():
        clear_l2_cache()
        return inputs, {}

    def wrapped_fn(*args):
        benchmark_fn(*args[0])
        return None

    # Set our custom CUDA event timer
    benchmark._timer = CUDAEventTimer()

    benchmark.pedantic(
        wrapped_fn,
        setup=setup,
        rounds=rounds,
        warmup_rounds=warmup_rounds,
        iterations=1,
    )

class MultideviceSettings:
    """Settings and utilities for multi-device execution."""


            self.is_running = False

def benchmark_cuda_events_pedantic(
    benchmark, benchmark_fn, inputs, rounds, warmup_rounds
):
    """Wrapper for benchmark_cuda_events that uses pytest's pedantic method with CUDA events.

    Args:
        benchmark: pytest-benchmark fixture
        benchmark_fn: Function to benchmark
        inputs: List of inputs to pass to benchmark_fn
        rounds: Number of rounds to run
        warmup_rounds: Number of warmup rounds
    """

    def setup():
        clear_l2_cache()
        return inputs, {}

    def wrapped_fn(*args):
        benchmark_fn(*args[0])
        return None

    # Set our custom CUDA event timer
    benchmark._timer = CUDAEventTimer()

    benchmark.pedantic(
        wrapped_fn,
        setup=setup,
        rounds=rounds,
        warmup_rounds=warmup_rounds,
        iterations=1,
    )

class MultideviceSettings:
    """Settings and utilities for multi-device execution."""

    def __init__(self):
        self._communicator = nvfuser.multidevice.Communicator.instance()
        torch.manual_seed(0)


communicator

    # Set our custom CUDA event timer
    benchmark._timer = CUDAEventTimer()

    benchmark.pedantic(
        wrapped_fn,
        setup=setup,
        rounds=rounds,
        warmup_rounds=warmup_rounds,
        iterations=1,
    )

class MultideviceSettings:
    """Settings and utilities for multi-device execution."""

    def __init__(self):
        self._communicator = nvfuser.multidevice.Communicator.instance()
        torch.manual_seed(0)

    @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:


size

        wrapped_fn,
        setup=setup,
        rounds=rounds,
        warmup_rounds=warmup_rounds,
        iterations=1,
    )

class MultideviceSettings:
    """Settings and utilities for multi-device execution."""

    def __init__(self):
        self._communicator = nvfuser.multidevice.Communicator.instance()
        torch.manual_seed(0)

    @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."


rank

        iterations=1,
    )

class MultideviceSettings:
    """Settings and utilities for multi-device execution."""

    def __init__(self):
        self._communicator = nvfuser.multidevice.Communicator.instance()
        torch.manual_seed(0)

    @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)


class MultideviceSettings:
    """Settings and utilities for multi-device execution."""

    def __init__(self):
        self._communicator = nvfuser.multidevice.Communicator.instance()
        torch.manual_seed(0)

    @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()


        self._communicator = nvfuser.multidevice.Communicator.instance()
        torch.manual_seed(0)

    @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]