            # Each call to uniform with DataType::Float will advance the offset by one
            # See Note [Divide offset by 4] in rng.cpp for more information
            stateless_sequence = [
                fd_det.execute([inputs[0], seed, rng_offset])
                for rng_offset in range(10)
            ]

            for i, (sful, sless) in enumerate(
                zip(stateful_sequence, stateless_sequence)
            ):
                nvfuser_direct_test.assertEqual(sful[0], sless[0])
        # Restore the RNG state
        torch.cuda.set_rng_state(state)

# Test that the range of generated uniform values spans the proper range
# https://github.com/NVIDIA/Fuser/issues/1653
def test_uniform_range(nvfuser_direct_test):
    dtypes = [DataType.Double, DataType.Float, DataType.Half, DataType.BFloat16]

    def run_test(left: float, right: float, dtype: DataType):
        samples_per_run = 2**29

        def fusion_fn(fd: FusionDefinition):
            # Generate enough values to reasonably expect to sample the ends of the range
            S0 = fd.define_scalar(left, dtype=DataType.Double)
            S1 = fd.define_scalar(right, dtype=DataType.Double)
            output = fd.ops.uniform(S0, S1, shape=[samples_per_run], dtype=dtype)
            fd.add_output(output)

        with FusionDefinition() as fd:
            fusion_fn(fd)

        output = fd.execute([])[0]

        x = output.amax()
        m = output.amin()
        mu = output.type(torch.float64).mean()
        # Repeat to improve chances of sampling extreme values
        num_runs = 100
        num_samples = num_runs * samples_per_run
        for i in range(num_runs):
            u = fd.execute([])[0]
            x = torch.maximum(x, u.amax())
            m = torch.minimum(m, u.amin())
            mu = mu + (u.type(torch.float64).mean() - mu) / (i + 2)

        # round-trip cast to find expected min
        theomin = torch.tensor(left, dtype=output.dtype).item()
        theomu = 0.5 * (right + left)
        theomax = torch.nextafter(
            torch.tensor(right, dtype=output.dtype),
            torch.tensor(left, dtype=output.dtype),
        )

        assert (
            m.item() >= theomin
        ), f"{output.dtype} expected min generated value {theomin} but found {m.item()}"
        assert (
            m.item() <= theomax
        ), f"{output.dtype} expected max generated value {theomax} but found {x.item()}"

        # uniform distribution on [0, 1) has mean 0.5 and variance 1/12
        # The standard error of the mean is then 1/sqrt(12 *
        # num_samples). We use the precision at 1.0 as a surrogate for
        # the contribution of rounding to the standard error of the
        # finite-precision mean.
        assert abs(mu.item() - theomu) < (right - left) * max(
            right - x.item(), 3.0 / math.sqrt(12 * num_samples)
        ), f"{output.dtype} expected mean generated value {theomu} but found {mu.item()}"

        if dtype not in [DataType.Float, DataType.Double]:
            # For reduced precision types, check that we sample the extreme
            # values. We don't do this for full precision types since the
            # amount of samples required would be too large.
            assert (
                m.item() == theomin
            ), f"{output.dtype} expected min generated value {theomin} but found {m.item()}"
            assert (
                x.item() == theomax
            ), f"{output.dtype} expected max generated value {theomax} but found {x.item()}"

    # test standard and non-standard uniform
    for left, right in [[0.0, 1.0], [-1.5, 3.7]]:
        for dtype in dtypes:
            run_test(left, right, dtype)

def test_cat_qwen2_v2(nvfuser_direct_test):
    def qwen2_cat_fusion_2(fd: FusionDefinition) -> None:
        T0 = fd.define_tensor(
            shape=[2048, 512],
            contiguity=[True, True],
            dtype=DataType.BFloat16,
            is_cpu=False,
            stride_order=[1, 0],
        )
        S1 = fd.define_scalar(None, dtype=DataType.Int)
        S2 = fd.define_scalar(None, dtype=DataType.Int)
        T3 = fd.define_tensor(
            shape=[1, 4, 2048, 128],


        ), f"{output.dtype} expected mean generated value {theomu} but found {mu.item()}"

        if dtype not in [DataType.Float, DataType.Double]:
            # For reduced precision types, check that we sample the extreme
            # values. We don't do this for full precision types since the
            # amount of samples required would be too large.
            assert (
                m.item() == theomin
            ), f"{output.dtype} expected min generated value {theomin} but found {m.item()}"
            assert (
                x.item() == theomax
            ), f"{output.dtype} expected max generated value {theomax} but found {x.item()}"

    # test standard and non-standard uniform
    for left, right in [[0.0, 1.0], [-1.5, 3.7]]:
        for dtype in dtypes:
            run_test(left, right, dtype)

def test_cat_qwen2_v2(nvfuser_direct_test):
    def qwen2_cat_fusion_2(fd: FusionDefinition) -> None:
        T0 = fd.define_tensor(
            shape=[2048, 512],
            contiguity=[True, True],
            dtype=DataType.BFloat16,
            is_cpu=False,
            stride_order=[1, 0],
        )
        S1 = fd.define_scalar(None, dtype=DataType.Int)
        S2 = fd.define_scalar(None, dtype=DataType.Int)
        T3 = fd.define_tensor(
            shape=[1, 4, 2048, 128],
            contiguity=[None, True, True, True],
            dtype=DataType.BFloat16,
            is_cpu=False,
            stride_order=[3, 1, 2, 0],
        )
        T4 = fd.define_tensor(
            shape=[1, 2048, 128],
            contiguity=[None, True, True],
            dtype=DataType.BFloat16,
            is_cpu=False,
            stride_order=[2, 0, 1],
        )
        T5 = fd.define_tensor(
            shape=[1, 2048, 128],
            contiguity=[None, True, True],
            dtype=DataType.BFloat16,
            is_cpu=False,
            stride_order=[2, 0, 1],
        )
        T6 = fd.define_tensor(
            shape=[1, 28, 2048, 128],
            contiguity=[None, True, True, True],
            dtype=DataType.BFloat16,
            is_cpu=False,
            stride_order=[3, 1, 2, 0],
        )
        T7 = fd.define_tensor(
            shape=[1, 2048, 512],
            contiguity=[None, True, True],
            dtype=DataType.BFloat16,
            is_cpu=False,
            stride_order=[2, 1, 0],
        )
        T12 = fd.ops.reshape(T0, new_shape=[1, 2048, 512])
        T13 = fd.ops.cast(T12, dtype=DataType.Float)
        S14 = fd.define_scalar(0.00000, dtype=DataType.Double)
        S15 = fd.define_scalar(1.00000, dtype=DataType.Double)
        S16 = fd.define_scalar(1, dtype=DataType.Int)
        S17 = fd.define_scalar(2048, dtype=DataType.Int)
        S18 = fd.define_scalar(512, dtype=DataType.Int)
        T20 = fd.ops.uniform(
            S14,
            S15,
            shape=[S16, S17, S18],
            rng_seed=S2,
            rng_offset=S1,
            dtype=DataType.BFloat16,
        )
        S21 = fd.define_scalar(4.00000, dtype=DataType.Double)
        T22 = fd.ops.mul(T13, S21)
        S23 = fd.define_scalar(0.900000, dtype=DataType.Double)
        T24 = fd.ops.lt(T20, S23)
        T25 = fd.ops.cast(T24, dtype=DataType.Float)
        T41 = fd.ops.slice(
            T3,
            start_indices=[0, 0, 0, 64],
            end_indices=[1, 4, 2048, 128],
            strides=[1, 1, 1, 1],
            manual_normalization=0,
        )
        T42 = fd.ops.mul(T22, T25)
        T43 = fd.ops.cast(T41, dtype=DataType.Float)
        T44 = fd.ops.neg(T43)
        T50 = fd.ops.broadcast_in_dim(
            T4, shape=[1, 1, 2048, 128], broadcast_dims=[0, 2, 3]
        )
        T66 = fd.ops.slice(
            T3,
            start_indices=[0, 0, 0, 0],
            end_indices=[1, 4, 2048, 64],
            strides=[1, 1, 1, 1],
            manual_normalization=0,
        )
        T67 = fd.ops.cast(T44, dtype=DataType.BFloat16)
        T73 = fd.ops.broadcast_in_dim(
            T5, shape=[1, 1, 2048, 128], broadcast_dims=[0, 2, 3]
        )
        T89 = fd.ops.slice(
            T6,
            start_indices=[0, 0, 0, 64],
            end_indices=[1, 28, 2048, 128],
            strides=[1, 1, 1, 1],
            manual_normalization=0,
        )
        S90 = fd.define_scalar(1.11111, dtype=DataType.Double)
        T91 = fd.ops.mul(T42, S90)
        T97 = fd.ops.broadcast_in_dim(
            T50, shape=[1, 4, 2048, 128], broadcast_dims=[0, 1, 2, 3]
        )
        T98 = fd.ops.cat([T67, T66], dim=-1, manual_padding=0)
        T104 = fd.ops.broadcast_in_dim(
            T73, shape=[1, 4, 2048, 128], broadcast_dims=[0, 1, 2, 3]
        )
        T105 = fd.ops.cast(T89, dtype=DataType.Float)
        T106 = fd.ops.cast(T97, dtype=DataType.Float)
        T107 = fd.ops.cast(T98, dtype=DataType.Float)
        T108 = fd.ops.cast(T104, dtype=DataType.Float)
        T109 = fd.ops.cast(T3, dtype=DataType.Float)
        T110 = fd.ops.neg(T105)
        T111 = fd.ops.cast(T7, dtype=DataType.Float)
        T112 = fd.ops.mul(T107, T106)
        T113 = fd.ops.mul(T109, T108)
        T129 = fd.ops.slice(
            T6,
            start_indices=[0, 0, 0, 0],
            end_indices=[1, 28, 2048, 64],
            strides=[1, 1, 1, 1],
            manual_normalization=0,
        )
        T130 = fd.ops.cast(T110, dtype=DataType.BFloat16)
        T131 = fd.ops.add(T111, T91)
        T137 = fd.ops.broadcast_in_dim(
            T50, shape=[1, 28, 2048, 128], broadcast_dims=[0, 1, 2, 3]
        )
        T138 = fd.ops.cat([T130, T129], dim=-1, manual_padding=0)
        T144 = fd.ops.broadcast_in_dim(
            T73, shape=[1, 28, 2048, 128], broadcast_dims=[0, 1, 2, 3]
        )
        T145 = fd.ops.cast(T131, dtype=DataType.BFloat16)
        T146 = fd.ops.cast(T137, dtype=DataType.Float)
        T147 = fd.ops.cast(T138, dtype=DataType.Float)
        T148 = fd.ops.cast(T144, dtype=DataType.Float)
        T149 = fd.ops.cast(T6, dtype=DataType.Float)
        T155 = fd.ops.reshape(T145, new_shape=[1, 2048, 4, 128])
        T156 = fd.ops.add(T113, T112)
        T157 = fd.ops.mul(T147, T146)
        T158 = fd.ops.mul(T149, T148)
        T159 = fd.ops.permute(T155, dims=[0, 2, 1, 3])
        T160 = fd.ops.cast(T156, dtype=DataType.BFloat16)
        T167 = fd.ops.broadcast_in_dim(
            T159, shape=[1, 4, 1, 2048, 128], broadcast_dims=[0, 1, 3, 4]
        )
        T174 = fd.ops.broadcast_in_dim(
            T160, shape=[1, 4, 1, 2048, 128], broadcast_dims=[0, 1, 3, 4]
        )
        T181 = fd.ops.broadcast_in_dim(
            T167, shape=[1, 4, 7, 2048, 128], broadcast_dims=[0, 1, 2, 3, 4]
        )
        T188 = fd.ops.broadcast_in_dim(
            T174, shape=[1, 4, 7, 2048, 128], broadcast_dims=[0, 1, 2, 3, 4]
        )
        T189 = fd.ops.add(T158, T157)
        T195 = fd.ops.reshape(T181, new_shape=[1, 28, 2048, 128])
        T201 = fd.ops.reshape(T188, new_shape=[1, 28, 2048, 128])
        T202 = fd.ops.cast(T189, dtype=DataType.BFloat16)
        T203 = fd.ops.stride_order(T195, stride_order=[3, 2, 1, 0])
        T204 = fd.ops.stride_order(T201, stride_order=[3, 2, 1, 0])
        T205 = fd.ops.stride_order(T202, stride_order=[3, 2, 1, 0])
        fd.add_output(T159)
        fd.add_output(T160)
        fd.add_output(T203)
        fd.add_output(T204)
        fd.add_output(T205)

    inputs = [
        torch.testing.make_tensor((2048, 512), dtype=torch.bfloat16, device="cuda:0"),
        25546,
        1400552702872758,
        torch.randn(1048576, dtype=torch.bfloat16, device="cuda:0").as_strided(
            (1, 4, 2048, 128), (1048576, 128, 512, 1)
        ),
        torch.randn(262144, dtype=torch.bfloat16, device="cuda:0").as_strided(
            (1, 2048, 128), (262144, 1, 2048)
        ),
        torch.randn(262144, dtype=torch.bfloat16, device="cuda:0").as_strided(
            (1, 2048, 128), (262144, 1, 2048)
        ),
        torch.randn(7340032, dtype=torch.bfloat16, device="cuda:0").as_strided(
            (1, 28, 2048, 128), (7340032, 128, 3584, 1)
        ),
        torch.testing.make_tensor(
            (1, 2048, 512), dtype=torch.bfloat16, device="cuda:0"
        ),


def test_prim_rms_norm_fwd(nvfuser_direct_test):
    input_size = [64, 128, 1024]
    dtype = torch.float32
    device = "cuda"
    inputs = [
        torch.randn(*input_size, device=device, requires_grad=True),
        torch.nn.Parameter(torch.randn(input_size[2], dtype=dtype, device=device)),
    ]

    def primitive_definition(
        inputs: torch.Tensor,
        weight: torch.Tensor,
        normalization_axis: int,
        keepdim: bool,
    ) -> torch.Tensor:
        var = inputs.mul(inputs).mean(normalization_axis, keepdim)
        pre_shift_scale_norm_output = inputs / torch.sqrt(var + 1e-12)
        norm_output = weight * pre_shift_scale_norm_output
        return norm_output

    def nvfuser_fusion(
        fd: FusionDefinition,
        normalization_axis: int,
        norm_size: int,
        input_shape: List[int],
        eps: float,
        keepDim: bool,
    ) -> None:
        inputs = fd.define_tensor(
            shape=[-1, -1, -1],
            contiguity=[True, True, True],
            dtype=DataType.Float,
        )
        weights = fd.define_tensor(shape=[-1], contiguity=[True], dtype=DataType.Float)
        inputs_sq = fd.ops.mul(inputs, inputs)
        sum0 = fd.ops.sum(inputs_sq, dims=[normalization_axis], keepdim=keepDim)
        norm_const = fd.define_scalar(norm_size)
        var = fd.ops.div(sum0, norm_const)
        eps_const = fd.define_scalar(eps)
        var_eps = fd.ops.add(var, eps_const)
        invstd = fd.ops.rsqrt(var_eps)
        pre_scale = fd.ops.mul(inputs, invstd)
        weights_bcast = fd.ops.broadcast_in_dim(
            weights, shape=input_shape, broadcast_dims=[2]
        )
        out = fd.ops.mul(pre_scale, weights_bcast)
        fd.add_output(out)
        fd.add_output(invstd)

    fusion_func = partial(
        nvfuser_fusion,
        normalization_axis=2,
        norm_size=inputs[0].size()[2],
        input_shape=inputs[0].size(),
        eps=1e-12,
        keepDim=True,
    )
    nvf_out, _ = nvfuser_direct_test.exec_nvfuser(fusion_func, inputs)

    eager_out = primitive_definition(inputs[0], inputs[1], 2, True)

    nvfuser_direct_test.assertEqual(eager_out, nvf_out[0])

def test_litgpt_variants_gpt_neox_like(nvfuser_direct_test):
    """
    This test hits a segfault in nvfuser::preseg_passes::MovePadPass::replaceCat.
    The old optimization pass updates the fusion within the filterByType generator


        )
        T23 = fd.ops.broadcast_in_dim(
            T22, shape=[16, 16, 128, 128], broadcast_dims=[0, 1, 2, 3]
        )
        T24 = fd.ops.div(T17, T23)
        S25 = fd.define_scalar(16, dtype=DataType.Int)
        S26 = fd.define_scalar(16, dtype=DataType.Int)
        S27 = fd.define_scalar(128, dtype=DataType.Int)
        S28 = fd.define_scalar(128, dtype=DataType.Int)
        S29 = fd.define_scalar(0.00000, dtype=DataType.Double)
        S30 = fd.define_scalar(1.00000, dtype=DataType.Double)
        T31 = fd.ops.uniform(S29, S30, shape=[S25, S26, S27, S28], dtype=DataType.Float)
        S32 = fd.define_scalar(1.0 - prob, dtype=DataType.Double)
        T33 = fd.ops.lt(T31, S32)
        T34 = fd.ops.cast(T33, dtype=DataType.Float)
        T35 = fd.ops.mul(T24, T34)
        S36 = fd.define_scalar(1.0 / (1.0 - prob), dtype=DataType.Double)
        T37 = fd.ops.mul(T35, S36)
        fd.add_output(T37)

    def torch_def(acts, bias, n_seq_len, n_head_dim, prob):
        att = acts * (1.0 / math.sqrt(n_head_dim))
        att = att.masked_fill(bias[:, :, :n_seq_len, :n_seq_len] == 0, float("-inf"))
        att = torch.nn.functional.softmax(att, dim=-1)
        att = torch.nn.functional.dropout(att, p=prob)
        return att

    # NOTE: The dropout probabilities need to be set to 0 elements zeroed out
    # in order to match implementations as eager and nvFuser do not have matching
    # blocking.
    nvf_out, _ = nvfuser_direct_test.exec_nvfuser(
        partial(nvfuser_fusion, prob=0.0), inputs
    )
    eager_out = torch_def(inputs[0], inputs[1], 128, 64, 0.0)

    for idx in range(len(nvf_out)):
        nvfuser_direct_test.assertEqual(eager_out, nvf_out[idx])

def test_nanogpt_split_mha_linears(nvfuser_direct_test):
    inputs = [
        torch.randn(16, 128, 3072, device="cuda"),
    ]

    def nvfuser_fusion_0(fd: FusionDefinition) -> None:
        T0 = fd.from_pytorch(inputs[0])


        att = torch.nn.functional.dropout(att, p=prob)
        return att

    # NOTE: The dropout probabilities need to be set to 0 elements zeroed out
    # in order to match implementations as eager and nvFuser do not have matching
    # blocking.
    nvf_out, _ = nvfuser_direct_test.exec_nvfuser(
        partial(nvfuser_fusion, prob=0.0), inputs
    )
    eager_out = torch_def(inputs[0], inputs[1], 128, 64, 0.0)

    for idx in range(len(nvf_out)):
        nvfuser_direct_test.assertEqual(eager_out, nvf_out[idx])

def test_nanogpt_split_mha_linears(nvfuser_direct_test):
    inputs = [
        torch.randn(16, 128, 3072, device="cuda"),
    ]

    def nvfuser_fusion_0(fd: FusionDefinition) -> None:
        T0 = fd.from_pytorch(inputs[0])
        T0_slice1 = fd.ops.slice(T0, [0, 0, 0], [16, 128, 1024], [1, 1, 1])
        T0_slice2 = fd.ops.slice(T0, [0, 0, 1024], [16, 128, 2048], [1, 1, 1])
        T0_slice3 = fd.ops.slice(T0, [0, 0, 2048], [16, 128, 3072], [1, 1, 1])
        T1_slice1 = fd.ops.reshape(T0_slice1, [16, 128, 16, 64])
        T1_slice2 = fd.ops.reshape(T0_slice2, [16, 128, 16, 64])
        T1_slice3 = fd.ops.reshape(T0_slice3, [16, 128, 16, 64])
        T2_slice1 = fd.ops.permute(T1_slice1, [0, 2, 1, 3])
        T2_slice2 = fd.ops.permute(T1_slice2, [0, 2, 1, 3])
        T2_slice3 = fd.ops.permute(T1_slice3, [0, 2, 1, 3])
        fd.add_output(T2_slice1)
        fd.add_output(T2_slice2)
        fd.add_output(T2_slice3)

    def torch_def_0(acts, n_embd, n_head):
        B, T, C = acts.size()
        q, k, v = acts.split(n_embd, dim=2)
        k = k.view(B, T, n_head, (C // 3) // n_head).transpose(1, 2)  # (B, nh, T, hs)
        q = q.view(B, T, n_head, (C // 3) // n_head).transpose(1, 2)  # (B, nh, T, hs)
        v = v.view(B, T, n_head, (C // 3) // n_head).transpose(1, 2)  # (B, nh, T, hs)
        return (
            q,
            k,
            v,
        )

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


def test_nanogpt_split_mha_linears(nvfuser_direct_test):
    inputs = [
        torch.randn(16, 128, 3072, device="cuda"),
    ]

    def nvfuser_fusion_0(fd: FusionDefinition) -> None:
        T0 = fd.from_pytorch(inputs[0])
        T0_slice1 = fd.ops.slice(T0, [0, 0, 0], [16, 128, 1024], [1, 1, 1])
        T0_slice2 = fd.ops.slice(T0, [0, 0, 1024], [16, 128, 2048], [1, 1, 1])
        T0_slice3 = fd.ops.slice(T0, [0, 0, 2048], [16, 128, 3072], [1, 1, 1])
        T1_slice1 = fd.ops.reshape(T0_slice1, [16, 128, 16, 64])
        T1_slice2 = fd.ops.reshape(T0_slice2, [16, 128, 16, 64])
        T1_slice3 = fd.ops.reshape(T0_slice3, [16, 128, 16, 64])
        T2_slice1 = fd.ops.permute(T1_slice1, [0, 2, 1, 3])
        T2_slice2 = fd.ops.permute(T1_slice2, [0, 2, 1, 3])
        T2_slice3 = fd.ops.permute(T1_slice3, [0, 2, 1, 3])
        fd.add_output(T2_slice1)
        fd.add_output(T2_slice2)
        fd.add_output(T2_slice3)

    def torch_def_0(acts, n_embd, n_head):
        B, T, C = acts.size()
        q, k, v = acts.split(n_embd, dim=2)
        k = k.view(B, T, n_head, (C // 3) // n_head).transpose(1, 2)  # (B, nh, T, hs)
        q = q.view(B, T, n_head, (C // 3) // n_head).transpose(1, 2)  # (B, nh, T, hs)
        v = v.view(B, T, n_head, (C // 3) // n_head).transpose(1, 2)  # (B, nh, T, hs)
        return (
            q,
            k,
            v,
        )

    def nvfuser_fusion_1(fd: FusionDefinition) -> None:
        T0 = fd.define_tensor(
            shape=[-1, -1, -1],
            contiguity=[True, True, True],
            dtype=DataType.Float,
            is_cpu=False,
        )
        T1 = fd.ops.slice(
            T0,
            start_indices=[0, 0, 0],
            end_indices=[16, 128, 1024],
            strides=[1, 1, 1],
        )
        T2 = fd.ops.slice(
            T0,
            start_indices=[0, 0, 1024],
            end_indices=[16, 128, 2048],
            strides=[1, 1, 1],
        )


        T1_slice3 = fd.ops.reshape(T0_slice3, [16, 128, 16, 64])
        T2_slice1 = fd.ops.permute(T1_slice1, [0, 2, 1, 3])
        T2_slice2 = fd.ops.permute(T1_slice2, [0, 2, 1, 3])
        T2_slice3 = fd.ops.permute(T1_slice3, [0, 2, 1, 3])
        fd.add_output(T2_slice1)
        fd.add_output(T2_slice2)
        fd.add_output(T2_slice3)

    def torch_def_0(acts, n_embd, n_head):
        B, T, C = acts.size()
        q, k, v = acts.split(n_embd, dim=2)
        k = k.view(B, T, n_head, (C // 3) // n_head).transpose(1, 2)  # (B, nh, T, hs)
        q = q.view(B, T, n_head, (C // 3) // n_head).transpose(1, 2)  # (B, nh, T, hs)
        v = v.view(B, T, n_head, (C // 3) // n_head).transpose(1, 2)  # (B, nh, T, hs)
        return (
            q,
            k,
            v,
        )

    def nvfuser_fusion_1(fd: FusionDefinition) -> None:
        T0 = fd.define_tensor(
            shape=[-1, -1, -1],
            contiguity=[True, True, True],
            dtype=DataType.Float,
            is_cpu=False,
        )
        T1 = fd.ops.slice(
            T0,
            start_indices=[0, 0, 0],
            end_indices=[16, 128, 1024],
            strides=[1, 1, 1],
        )
        T2 = fd.ops.slice(
            T0,
            start_indices=[0, 0, 1024],
            end_indices=[16, 128, 2048],
            strides=[1, 1, 1],
        )
        T3 = fd.ops.slice(
            T0,
            start_indices=[0, 0, 2048],
            end_indices=[16, 128, 3072],
            strides=[1, 1, 1],
        )
        fd.add_output(T1)
        fd.add_output(T2)
        fd.add_output(T3)

    def torch_def_1(acts, n_embd, n_head):
        B, T, C = acts.size()
        q, k, v = acts.split(n_embd, dim=2)
        return (
            q,
            k,
            v,
        )

    tests = [
        (nvfuser_fusion_0, torch_def_0),
        (nvfuser_fusion_1, torch_def_1),
    ]

    for nvf_func, torch_func in tests:
        nvf_out, _ = nvfuser_direct_test.exec_nvfuser(nvf_func, inputs)
        eager_out = torch_func(*inputs, 1024, 16)
        for idx in range(len(eager_out)):
            nvfuser_direct_test.assertEqual(eager_out[idx], nvf_out[idx])


            start_indices=[0, 0, 0],
            end_indices=[16, 128, 1024],
            strides=[1, 1, 1],
        )
        T2 = fd.ops.slice(
            T0,
            start_indices=[0, 0, 1024],
            end_indices=[16, 128, 2048],
            strides=[1, 1, 1],
        )
        T3 = fd.ops.slice(
            T0,
            start_indices=[0, 0, 2048],
            end_indices=[16, 128, 3072],
            strides=[1, 1, 1],
        )
        fd.add_output(T1)
        fd.add_output(T2)
        fd.add_output(T3)

    def torch_def_1(acts, n_embd, n_head):
        B, T, C = acts.size()
        q, k, v = acts.split(n_embd, dim=2)
        return (
            q,
            k,
            v,
        )

    tests = [
        (nvfuser_fusion_0, torch_def_0),
        (nvfuser_fusion_1, torch_def_1),
    ]

    for nvf_func, torch_func in tests:
        nvf_out, _ = nvfuser_direct_test.exec_nvfuser(nvf_func, inputs)
        eager_out = torch_func(*inputs, 1024, 16)
        for idx in range(len(eager_out)):
            nvfuser_direct_test.assertEqual(eager_out[idx], nvf_out[idx])

def test_prim_layer_norm_fwd(nvfuser_direct_test):
    input_size = [64, 128, 1024]
    dtype = torch.float32
    device = "cuda"
    inputs = [
        torch.randn(*input_size, device=device, requires_grad=True),
        torch.nn.Parameter(torch.randn(input_size[2], dtype=dtype, device=device)),


        eps=1e-12,
        keepDim=True,
    )
    nvf_var_mean_out, _ = nvfuser_direct_test.exec_nvfuser(fusion_func_2, inputs)

    eager_out = primitive_definition(inputs[0], inputs[1], inputs[2], 2, True)

    nvfuser_direct_test.assertEqual(eager_out, nvf_out[0])
    nvfuser_direct_test.assertEqual(eager_out, nvf_var_mean_out[0])

def test_prim_rms_norm_fwd(nvfuser_direct_test):
    input_size = [64, 128, 1024]
    dtype = torch.float32
    device = "cuda"
    inputs = [
        torch.randn(*input_size, device=device, requires_grad=True),
        torch.nn.Parameter(torch.randn(input_size[2], dtype=dtype, device=device)),
    ]

    def primitive_definition(
        inputs: torch.Tensor,
        weight: torch.Tensor,
        normalization_axis: int,
        keepdim: bool,
    ) -> torch.Tensor:
        var = inputs.mul(inputs).mean(normalization_axis, keepdim)
        pre_shift_scale_norm_output = inputs / torch.sqrt(var + 1e-12)
        norm_output = weight * pre_shift_scale_norm_output
        return norm_output

    def nvfuser_fusion(
        fd: FusionDefinition,
        normalization_axis: int,
        norm_size: int,
        input_shape: List[int],
        eps: float,
        keepDim: bool,
    ) -> None:
        inputs = fd.define_tensor(
            shape=[-1, -1, -1],
            contiguity=[True, True, True],
            dtype=DataType.Float,
        )
        weights = fd.define_tensor(shape=[-1], contiguity=[True], dtype=DataType.Float)
        inputs_sq = fd.ops.mul(inputs, inputs)
        sum0 = fd.ops.sum(inputs_sq, dims=[normalization_axis], keepdim=keepDim)
        norm_const = fd.define_scalar(norm_size)
        var = fd.ops.div(sum0, norm_const)
        eps_const = fd.define_scalar(eps)


        diff = fd.ops.sub(inputs, mean)
        diff_sq = fd.ops.mul(diff, diff)
        sum1 = fd.ops.sum(diff_sq, dims=[normalization_axis], keepdim=keepDim)
        var = fd.ops.div(sum1, norm_const)
        eps_const = fd.define_scalar(eps)
        var_eps = fd.ops.add(var, eps_const)
        invstd = fd.ops.rsqrt(var_eps)
        pre_scale_bias = fd.ops.mul(diff, invstd)
        weights_bcast = fd.ops.broadcast_in_dim(
            weights, shape=input_shape, broadcast_dims=[2]
        )
        scale = fd.ops.mul(pre_scale_bias, weights_bcast)
        bias_bcast = fd.ops.broadcast_in_dim(
            bias, shape=input_shape, broadcast_dims=[2]
        )
        out = fd.ops.add(scale, bias_bcast)
        fd.add_output(out)
        fd.add_output(mean)
        fd.add_output(invstd)

    def nvfuser_fusion_var_mean(
        fd: FusionDefinition,
        normalization_axis: int,
        norm_size: int,
        input_shape: List[int],
        eps: float,
        keepDim: bool,
    ) -> None:
        inputs = fd.define_tensor(
            shape=[-1, -1, -1],
            contiguity=[True, True, True],
            dtype=DataType.Float,
        )
        weights = fd.define_tensor(shape=[-1], contiguity=[True], dtype=DataType.Float)
        bias = fd.define_tensor(shape=[-1], contiguity=[True], dtype=DataType.Float)
        var, mean = fd.ops.var_mean(
            inputs, dims=[normalization_axis], correction=0, keepdim=keepDim
        )
        eps_const = fd.define_scalar(eps)
        var_eps = fd.ops.add(var, eps_const)
        invstd = fd.ops.rsqrt(var_eps)
        diff = fd.ops.sub(inputs, mean)
        pre_scale_bias = fd.ops.mul(diff, invstd)
        weights_bcast = fd.ops.broadcast_in_dim(
            weights, shape=input_shape, broadcast_dims=[2]
        )
        scale = fd.ops.mul(pre_scale_bias, weights_bcast)
        bias_bcast = fd.ops.broadcast_in_dim(
            bias, shape=input_shape, broadcast_dims=[2]
        )
        out = fd.ops.add(scale, bias_bcast)
        fd.add_output(out)
        fd.add_output(mean)
        fd.add_output(invstd)

    fusion_func_1 = partial(
        nvfuser_fusion,
        normalization_axis=2,
        norm_size=inputs[0].size()[2],
        input_shape=inputs[0].size(),
        eps=1e-12,
        keepDim=True,
    )
    nvf_out, _ = nvfuser_direct_test.exec_nvfuser(fusion_func_1, inputs)

    fusion_func_2 = partial(
        nvfuser_fusion_var_mean,
        normalization_axis=2,
        norm_size=inputs[0].size()[2],
        input_shape=inputs[0].size(),
        eps=1e-12,
        keepDim=True,
    )
    nvf_var_mean_out, _ = nvfuser_direct_test.exec_nvfuser(fusion_func_2, inputs)


        normalization_axis=2,
        norm_size=inputs[0].size()[2],
        input_shape=inputs[0].size(),
        eps=1e-12,
        keepDim=True,
    )
    nvf_out, _ = nvfuser_direct_test.exec_nvfuser(fusion_func, inputs)

    eager_out = primitive_definition(inputs[0], inputs[1], 2, True)

    nvfuser_direct_test.assertEqual(eager_out, nvf_out[0])

def test_litgpt_variants_gpt_neox_like(nvfuser_direct_test):
    """
    This test hits a segfault in nvfuser::preseg_passes::MovePadPass::replaceCat.
    The old optimization pass updates the fusion within the filterByType generator
    This causes a dynamic cast on a dangling pointer.
    """

    def nvfuser_fusion_id4(fd: FusionDefinition) -> None:
        T0 = fd.define_tensor(
            shape=[128, 16],
            contiguity=[True, True],
            dtype=DataType.Float,
            is_cpu=False,
            stride_order=[1, 0],
        )
        T1 = fd.define_tensor(
            shape=[128, 16],
            contiguity=[True, True],
            dtype=DataType.Float,
            is_cpu=False,
            stride_order=[1, 0],
        )
        T2 = fd.define_tensor(
            shape=[5, 5, 192],
            contiguity=[True, True, True],
            dtype=DataType.Float,
            is_cpu=False,
            stride_order=[2, 1, 0],
        )
        T12 = fd.ops.slice(
            T0, start_indices=[0, 0], end_indices=[5, 16], strides=[1, 1]
        )
        T22 = fd.ops.slice(
            T1, start_indices=[0, 0], end_indices=[5, 16], strides=[1, 1]
        )
        T29 = fd.ops.reshape(T2, new_shape=[5, 5, 4, 3, 16])
        T30 = fd.ops.permute(T29, dims=[0, 2, 3, 1, 4])
        T49 = fd.ops.slice(
            T30,
            start_indices=[0, 0, 0, 0, 0],
            end_indices=[5, 4, 1, 5, 16],
            strides=[1, 1, 1, 1, 1],
        )
        T68 = fd.ops.slice(
            T30,
            start_indices=[0, 0, 1, 0, 0],
            end_indices=[5, 4, 2, 5, 16],
            strides=[1, 1, 1, 1, 1],
        )
        T87 = fd.ops.slice(
            T30,
            start_indices=[0, 0, 2, 0, 0],
            end_indices=[5, 4, 3, 5, 16],
            strides=[1, 1, 1, 1, 1],
        )
        T93 = fd.ops.reshape(T49, new_shape=[5, 4, 5, 16])
        T99 = fd.ops.reshape(T68, new_shape=[5, 4, 5, 16])
        T105 = fd.ops.reshape(T87, new_shape=[5, 4, 5, 16])
        T121 = fd.ops.slice(
            T93,
            start_indices=[0, 0, 0, 0],
            end_indices=[5, 4, 5, 8],
            strides=[1, 1, 1, 1],
        )
        T137 = fd.ops.slice(
            T93,
            start_indices=[0, 0, 0, 8],
            end_indices=[5, 4, 5, 16],
            strides=[1, 1, 1, 1],
        )
        T138 = fd.ops.neg(T137)
        T139 = fd.ops.cat([T138, T121], dim=-1)
        S140 = fd.define_scalar(5, dtype=DataType.Int)
        S141 = fd.define_scalar(4, dtype=DataType.Int)
        S142 = fd.define_scalar(5, dtype=DataType.Int)
        S143 = fd.define_scalar(16, dtype=DataType.Int)
        T145 = fd.ops.broadcast_in_dim(
            T12, shape=[S140, S141, S142, S143], broadcast_dims=[2, 3]
        )
        T146 = fd.ops.mul(T93, T145)
        S147 = fd.define_scalar(5, dtype=DataType.Int)
        S148 = fd.define_scalar(4, dtype=DataType.Int)
        S149 = fd.define_scalar(5, dtype=DataType.Int)
        S150 = fd.define_scalar(16, dtype=DataType.Int)
        T152 = fd.ops.broadcast_in_dim(
            T22, shape=[S147, S148, S149, S150], broadcast_dims=[2, 3]
        )
        T153 = fd.ops.mul(T139, T152)
        T154 = fd.ops.add(T146, T153)
        T170 = fd.ops.slice(
            T99,
            start_indices=[0, 0, 0, 0],
            end_indices=[5, 4, 5, 8],
            strides=[1, 1, 1, 1],
        )
        T186 = fd.ops.slice(
            T99,
            start_indices=[0, 0, 0, 8],
            end_indices=[5, 4, 5, 16],
            strides=[1, 1, 1, 1],
        )
        T187 = fd.ops.neg(T186)
        T188 = fd.ops.cat([T187, T170], dim=-1)
        T189 = fd.ops.mul(T99, T145)
        T190 = fd.ops.mul(T188, T152)
        T191 = fd.ops.add(T189, T190)
        T207 = fd.ops.slice(
            T93,
            start_indices=[0, 0, 0, 0],
            end_indices=[5, 4, 5, 0],
            strides=[1, 1, 1, 1],
        )
        T208 = fd.ops.cat([T154, T207], dim=-1)
        T224 = fd.ops.slice(
            T99,
            start_indices=[0, 0, 0, 0],
            end_indices=[5, 4, 5, 0],
            strides=[1, 1, 1, 1],
        )
        T225 = fd.ops.cat([T191, T224], dim=-1)
        S226 = fd.define_scalar(0.500000, dtype=DataType.Double)
        T227 = fd.ops.mul(T208, S226)
        T228 = fd.ops.permute(T225, dims=[0, 1, 3, 2])
        S229 = fd.define_scalar(0.500000, dtype=DataType.Double)
        T230 = fd.ops.mul(T228, S229)
        fd.add_output(T105)
        fd.add_output(T145)
        fd.add_output(T152)
        fd.add_output(T227)
        fd.add_output(T230)

    with FusionDefinition() as fd:
        nvfuser_fusion_id4(fd)

    inputs = [
        torch.testing.make_tensor((128, 16), dtype=torch.float32, device="cuda:0"),
        torch.testing.make_tensor((128, 16), dtype=torch.float32, device="cuda:0"),
        torch.testing.make_tensor((5, 5, 192), dtype=torch.float32, device="cuda:0"),
    ]
    fd.validate(inputs)

    # check python reproduction
    nvf_out, _ = nvfuser_direct_test.exec_nvfuser(nvfuser_fusion_id4, inputs)


                fd_det.execute([inputs[0], seed, rng_offset])
                for rng_offset in range(10)
            ]

            for i, (sful, sless) in enumerate(
                zip(stateful_sequence, stateless_sequence)
            ):
                nvfuser_direct_test.assertEqual(sful[0], sless[0])
        # Restore the RNG state
        torch.cuda.set_rng_state(state)

# Test that the range of generated uniform values spans the proper range
# https://github.com/NVIDIA/Fuser/issues/1653
def test_uniform_range(nvfuser_direct_test):
    dtypes = [DataType.Double, DataType.Float, DataType.Half, DataType.BFloat16]

    def run_test(left: float, right: float, dtype: DataType):
        samples_per_run = 2**29

        def fusion_fn(fd: FusionDefinition):
            # Generate enough values to reasonably expect to sample the ends of the range
            S0 = fd.define_scalar(left, dtype=DataType.Double)
            S1 = fd.define_scalar(right, dtype=DataType.Double)
            output = fd.ops.uniform(S0, S1, shape=[samples_per_run], dtype=dtype)
            fd.add_output(output)

        with FusionDefinition() as fd:
            fusion_fn(fd)

        output = fd.execute([])[0]

        x = output.amax()
        m = output.amin()
        mu = output.type(torch.float64).mean()
        # Repeat to improve chances of sampling extreme values
        num_runs = 100
        num_samples = num_runs * samples_per_run
        for i in range(num_runs):
            u = fd.execute([])[0]
            x = torch.maximum(x, u.amax())
            m = torch.minimum(m, u.amin())
            mu = mu + (u.type(torch.float64).mean() - mu) / (i + 2)

        # round-trip cast to find expected min
        theomin = torch.tensor(left, dtype=output.dtype).item()


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

import torch

from nvfuser_direct import FusionDefinition, DataType

def test_contiguous():
    with FusionDefinition() as fd:
        inp = fd.define_tensor([-1, -1], dtype=DataType.Float, contiguity=True)
        out = fd.ops.sum(inp, [-1])
        fd.add_output(out)

    # It might not be obvious but this tensor is indeed contiguous by
    # definition. inp.size(0) == 1 so inp.stride(0) shouldn't matter.
    #
    # I ran into such a tensor when I shard a reference tensor for a particular
    # device using slicing:
    # ```
    # inp = torch.arange(4, device="cuda").view(2, 2)
    # inp = inp[0:1, 0:1].contiguous()
    # ```
    # Despite `.contiguous()`, inp.stride(0) remains to be 2.
    inp = torch.as_strided(torch.ones([1], device="cuda"), [1, 1], [2, 1])
    (out,) = fd.execute([inp])
    torch.testing.assert_close(out.cpu(), torch.ones([1]))


# SPDX-FileCopyrightText: Copyright (c) 2025-present NVIDIA CORPORATION & AFFILIATES.
# All rights reserved.
# SPDX-License-Identifier: BSD-3-Clause
# Owner(s): ["module: nvfuser"]

from nvfuser_direct import (
    FusionDefinition,
    PythonProfiler,
    DataType,
    version,
    LruFusionCache,
    LRUCache,
)
import torch
import pytest
import io
import re
from contextlib import redirect_stdout, redirect_stderr
from typing import Callable

def test_python_version_API():
    from nvfuser_direct.nvfuser_direct_version import Version

    assert version() > "0.0.0"
    assert version() > Version("0.0.0")

def test_fusion_not_defined():
    inputs = [
        torch.randn(4, 4, device="cpu"),
    ]

    # A FusionDefinition object is constructed but not defined, should trip an error
    try:
        fd = FusionDefinition()
        out = fd.execute(inputs)
        raise RuntimeError("Expecting an error for an empty FusionDefinition!")
    except NotImplementedError as e:
        assert (
            "Fusion does not exist! Use `with FusionDefinition() as fd: ...` to define a fusion."
            in str(e)
        )

def test_fusion_empty():
    inputs = [
        torch.randn(4, 4, device="cpu"),
    ]

    # A FusionDefinition object is constructed but not defined, should trip an error
    with pytest.raises(NotImplementedError, match="Fusion is empty!"):
        with FusionDefinition() as fd:
            pass
        out = fd.execute(inputs)

def test_from_pytorch_fails_on_cpu_tensor():
    inputs = [
        torch.randn(4, 4, device="cpu"),
    ]

    stdout_capture = io.StringIO()
    stderr_capture = io.StringIO()
    error_msg = (
        "Found unsupported device cpu, only scalar CPU or CUDA tensors are supported"
    )
    with pytest.raises(ValueError, match=error_msg), redirect_stdout(
        stdout_capture
    ), redirect_stderr(stderr_capture):
        with FusionDefinition() as fd:
            t0 = fd.from_pytorch(inputs[0])
            t1 = fd.ops.relu(t0)
            fd.add_output(t1)

def test_fusion_definition_print():
    with FusionDefinition() as fd:
        tv0 = fd.define_tensor(
            shape=[2, 4, 8],
        )
        tv1 = fd.define_tensor(
            shape=[2, 4, 8],
        )
        tv2 = fd.ops.add(tv0, tv1)
        fd.add_output(tv2)

    # Test fusion math representation
    prescheduled_fusion_definition = """Inputs:
  T0_g_float[iS0{2}, iS1{4}, iS2{8}]
  T1_g_float[iS3{2}, iS4{4}, iS5{8}]
Outputs:
  T2_g_float[iS6{2}, iS7{4}, iS8{8}]

%kernel_math {
T2_g_float[iS6{2}, iS7{4}, iS8{8}]
   = T0_g_float[iS0{2}, iS1{4}, iS2{8}]
   + T1_g_float[iS3{2}, iS4{4}, iS5{8}];
} // %kernel_math \n\n"""
    assert fd.fusion.print_math() == prescheduled_fusion_definition

    # Test CUDA kernel representation
    cuda_kernel = """// Codegen generated code
__global__ void CUDAGeneratedKernel(Tensor<float, 3, 3> T0, Tensor<float, 3, 3> T1, Tensor<float, 3, 3> T2) {
  NVFUSER_DEFINE_MAGIC_ZERO;
  #pragma unroll
  for(nvfuser_index_t i0 = 0LL; i0 < 2LL; ++i0) {
    nvfuser_index_t i1;
    i1 = T0.alloc_stride[0LL] * i0;
    nvfuser_index_t i2;
    i2 = T1.alloc_stride[0LL] * i0;
    nvfuser_index_t i3;
    i3 = 32LL * i0;
    #pragma unroll
    for(nvfuser_index_t i4 = 0LL; i4 < 4LL; ++i4) {
      nvfuser_index_t i5;
      i5 = i1 + (T0.alloc_stride[1LL] * i4);
      nvfuser_index_t i6;
      i6 = i2 + (T1.alloc_stride[1LL] * i4);
      nvfuser_index_t i7;
      i7 = i3 + (8LL * i4);
      #pragma unroll
      for(nvfuser_index_t i8 = 0LL; i8 < 8LL; ++i8) {
        nvfuser_index_t i9;
        i9 = i8 + nvfuser_zero;
        T2[(i7 + i9)]
          = T0[(i5 + (T0.alloc_stride[2LL] * i9))]
          + T1[(i6 + (T1.alloc_stride[2LL] * i9))];
      }
    }
  }
  NVFUSER_UPDATE_MAGIC_ZERO;
}
"""
    assert fd.fusion.print_kernel() == cuda_kernel

    # Test TensorView string representation
    assert str(tv0) == r"T0_g_float[iS0{2}, iS1{4}, iS2{8}]"
    assert str(tv1) == r"T1_g_float[iS3{2}, iS4{4}, iS5{8}]"
    assert str(tv2) == r"T2_g_float[iS6{2}, iS7{4}, iS8{8}]"

    # Test TensorDomain string representation
    assert str(tv0.domain()) == r"[iS0{2}, iS1{4}, iS2{8}]"
    assert str(tv1.domain()) == r"[iS3{2}, iS4{4}, iS5{8}]"
    assert str(tv2.domain()) == r"[iS6{2}, iS7{4}, iS8{8}]"

    # Test IterDomain string representation
    assert str(tv0.axis(0)) == r"iS0{2}"
    assert str(tv1.axis(0)) == r"iS3{2}"
    assert str(tv2.axis(0)) == r"iS6{2}"

    # Test axis extents
    assert str(tv0.axis(0).extent()) == r"2"
    assert str(tv1.axis(0).extent()) == r"2"
    assert str(tv2.axis(0).extent()) == r"2"

def test_fusion_execution_cache():
    with FusionDefinition() as fd:
        tv0 = fd.define_tensor(
            shape=[2, 4, 8],
        )
        tv1 = fd.define_tensor(
            shape=[2, 4, 8],
        )
        tv2 = fd.ops.add(tv0, tv1)
        fd.add_output(tv2)

    # Test fusion execution
    inputs = [
        torch.randn(2, 4, 8, device="cuda"),
        torch.randn(2, 4, 8, device="cuda"),
    ]
    results = fd.execute(inputs)
    torch.testing.assert_close(results[0], inputs[0] + inputs[1])

    # Check that TensorViews are accessible after running fusion
    assert str(tv0.domain()) == "[iS0{2}, iS1{4}, iS2{8}]"
    assert str(tv1.domain()) == "[iS3{2}, iS4{4}, iS5{8}]"
    assert str(tv2.domain()) == "[iS6{2}, iS7{4}, iS8{8}]"

    with pytest.raises(
        AssertionError,
        match=r"FusionExecutorCache already exists! Use execute\(\) to execute the fusion.",
    ):
        fd.manual_execute(inputs)

    # Test fusion math representation after compilation
    prescheduled_fusion_definition = """Inputs:
  T0_g_float[iS0{2}, iS1{4}, iS2{8}]
  T1_g_float[iS3{2}, iS4{4}, iS5{8}]
Outputs:
  T2_g_float[iS6{2}, iS7{4}, iS8{8}]

%kernel_math {
T2_g_float[iS6{2}, iS7{4}, iS8{8}]
   = T0_g_float[iS0{2}, iS1{4}, iS2{8}]
   + T1_g_float[iS3{2}, iS4{4}, iS5{8}];
} // %kernel_math \n\n"""
    assert fd.fusion.print_math() == prescheduled_fusion_definition

    # Test compilation status
    assert fd.fec.is_compiled(inputs)

    # Test scheduled IR representation
    scheduled_ir = """Inputs:
  T0_g_float[iS58{1}, iS59{1}, iS57{128}]
  T1_g_float[iS46{1}, iS47{1}, iS45{128}]
Outputs:
  T2_g_float[iblockIdx.x28{1}, iUS29{1}, ithreadIdx.x27{128}] ca_pos( 2 ) produce_pos( 3 )

%kernel {
T3_l_float[iblockIdx.x52{1}, iUS53{1}, ithreadIdx.x51{128}] ca_pos( 2 )
   = Set( T0_g_float[iS58{1}, iS59{1}, iS57{128}], cache_op=Streaming )
T4_l_float[iblockIdx.x40{1}, iUS41{1}, ithreadIdx.x39{128}] ca_pos( 2 )
   = Set( T1_g_float[iS46{1}, iS47{1}, iS45{128}], cache_op=Streaming )
T5_l_float[iblockIdx.x34{1}, iUS35{1}, ithreadIdx.x33{128}] ca_pos( 3 ) produce_pos( 2 )
   = T3_l_float[iblockIdx.x52{1}, iUS53{1}, ithreadIdx.x51{128}] ca_pos( 2 )
   + T4_l_float[iblockIdx.x40{1}, iUS41{1}, ithreadIdx.x39{128}] ca_pos( 2 );
T2_g_float[iblockIdx.x28{1}, iUS29{1}, ithreadIdx.x27{128}] ca_pos( 2 ) produce_pos( 3 )
   = Set( T5_l_float[iblockIdx.x34{1}, iUS35{1}, ithreadIdx.x33{128}] ca_pos( 3 ) produce_pos( 2 ), cache_op=Streaming )
} // %kernel\n"""
    assert fd.fec.get_scheduled_ir(inputs) == scheduled_ir
    assert fd.fec.get_most_recent_scheduled_ir() == scheduled_ir

    # Test CUDA kernel representation
    cuda_kernel = """// Codegen generated code
__global__ void nvfuser_pointwise_f0_c1_r0_g0(Tensor<float, 3, 3> T0, Tensor<float, 3, 3> T1, Tensor<float, 3, 3> T2) {
  nvfuser_index_t i0;
  i0 = ((nvfuser_index_t)threadIdx.x) % 32;
  nvfuser_index_t i1;
  i1 = ((nvfuser_index_t)threadIdx.x) / 32;
  nvfuser_index_t i2;
  i2 = i0 / 8;
  nvfuser_index_t i3;
  i3 = i0 % 8;
  if ((((nvfuser_index_t)threadIdx.x) < 64)) {
    Array<float, 1, 1> T4;
    T4[0] = 0;
    T4[0]
       = T1[(((T1.alloc_stride[0LL] * i1) + (T1.alloc_stride[1LL] * i2)) + (T1.alloc_stride[2LL] * i3))];
    Array<float, 1, 1> T3;
    T3[0] = 0;
    T3[0]
       = T0[(((T0.alloc_stride[0LL] * i1) + (T0.alloc_stride[1LL] * i2)) + (T0.alloc_stride[2LL] * i3))];
    Array<float, 1, 1> T5;
    T5[0]
      = T3[0]
      + T4[0];
    T2[((nvfuser_index_t)threadIdx.x)]
       = T5[0];
  }
}\n"""
    assert fd.fec.get_cuda_kernel(inputs) == cuda_kernel

def test_kernel_executor():
    with FusionDefinition() as fd:
        tv0 = fd.define_tensor(
            shape=[2, 4, 8],
        )
        tv1 = fd.define_tensor(
            shape=[2, 4, 8],
        )
        tv2 = fd.ops.add(tv0, tv1)
        fd.add_output(tv2)

    # Test fusion execution
    inputs = [
        torch.randn(2, 4, 8, device="cuda"),
        torch.randn(2, 4, 8, device="cuda"),
    ]
    results = fd.manual_execute(inputs)
    torch.testing.assert_close(results[0], inputs[0] + inputs[1])
    assert fd.ke.is_compiled()

    with pytest.raises(
        AssertionError,
        match=r"KernelExecutor already exists! Use manual_execute\(\) to execute the fusion.",
    ):
        fd.execute(inputs)

def test_repro_script_for():
    inputs = [
        torch.ones(2, 4, 8, device="cuda"),
        torch.ones(2, 4, 8, device="cuda"),
    ]

    with FusionDefinition() as fd:
        t0 = fd.from_pytorch(inputs[0])
        t1 = fd.from_pytorch(inputs[1])
        c0 = fd.define_scalar(3.0)

        t2 = fd.ops.add(t0, t1)
        t3 = fd.ops.mul(t2, c0)
        t4 = fd.ops.sum(t3, [1], False, DataType.Float)

        fd.add_output(t4)

    expected_repro = """
import torch
from nvfuser_direct import FusionDefinition, DataType
def nvfuser_fusion(fd : FusionDefinition) -> None :
    tv0 = fd.define_tensor(shape=[-1, -1, -1], contiguity=[True, True, True], dtype=DataType.Float, is_cpu=False)
    tv1 = fd.define_tensor(shape=[-1, -1, -1], contiguity=[True, True, True], dtype=DataType.Float, is_cpu=False)
    tv2 = fd.ops.add(tv0, tv1)
    tv3 = fd.ops.mul(tv2, 3.00000)
    tv4 = fd.ops.sum(tv3, dims=[1], dtype=DataType.Float)
    fd.add_output(tv4)
with FusionDefinition() as fd:
    nvfuser_fusion(fd)
"""

    expected_inputs = """
inputs = [
    torch.testing.make_tensor((2, 4, 8), dtype=torch.float32, device='cuda:0'),
    torch.testing.make_tensor((2, 4, 8), dtype=torch.float32, device='cuda:0'),
]
fd.execute(inputs)\n"""

    # fd.repro_script_for(inputs) should have input arguments
    repro_with_inputs = fd.repro_script_for(inputs)
    assert expected_repro in repro_with_inputs
    assert expected_inputs in repro_with_inputs

    # fd.repro_script_for() should NOT have input arguments
    repro_without_inputs = fd.repro_script_for()
    assert expected_repro in repro_without_inputs
    assert expected_inputs not in repro_without_inputs

    # Check last_repro_script fails gracefully.
    with pytest.raises(
        AssertionError,
        match=r"fd.last_repro_script\(\) cannot provide a repro because fd.execute\(inputs, save_repro_inputs=True\) was not executed!",
    ):
        fd.last_repro_script()

    # Test fd.execute(inputs, save_repro_inputs=True) ; fd.last_repro_script()
    fd.execute(inputs, save_repro_inputs=True)
    last_repro = fd.last_repro_script()
    assert repro_with_inputs == last_repro

def test_define_tensor():
    with FusionDefinition() as fd:
        tv0 = fd.define_tensor(
            shape=[-1, -1, -1],
            contiguity=[True, True, True],
            dtype=DataType.Float,
            is_cpu=False,
            stride_order=[2, 1, 0],
        )
        tv1 = fd.define_tensor(
            shape=[-1, -1, -1],
            contiguity=[True, True, True],
            dtype=DataType.Float,
            is_cpu=False,
            stride_order=[2, 1, 0],
        )
        tv2 = fd.ops.add(tv0, tv1)
        fd.add_output(tv2)

    # Test fusion execution with dynamic shapes
    inputs = [
        torch.randn(2, 4, 8, device="cuda"),
        torch.randn(2, 4, 8, device="cuda"),
    ]
    outputs = fd.execute(inputs)
    torch.testing.assert_close(outputs[0], inputs[0] + inputs[1])

@pytest.mark.skipif(
    torch.cuda.device_count() < 2,
    reason="test_selected_device requires multiple GPUs",
)
def test_execute_with_different_device():
    inputs = [
        torch.ones(2, 4, 8, device="cuda:1"),
        torch.ones(2, 4, 8, device="cuda:1"),
    ]

    with FusionDefinition() as fd:
        t0 = fd.from_pytorch(inputs[0])
        t1 = fd.from_pytorch(inputs[1])
        c0 = fd.define_scalar(3.0)

        t2 = fd.ops.add(t0, t1)
        t3 = fd.ops.mul(t2, c0)
        t4 = fd.ops.sum(t3, [1], False, DataType.Float)

        fd.add_output(t4)

    outputs = fd.execute(inputs, device="cuda:1")
    assert len(outputs) == 1
    assert outputs[0].device.index == 1

def test_enable_disable_options():
    m = 24
    n = 16
    k = 8
    inps = [
        torch.randn(m, k, device="cuda", dtype=torch.float),
        torch.randn(k, n, device="cuda", dtype=torch.float),
    ]

    def fusion_func(fd: FusionDefinition, inps) -> None:
        t0 = fd.from_pytorch(inps[0])
        t1 = fd.from_pytorch(inps[1])
        t2 = fd.ops.matmul(t0, t1)
        fd.add_output(t2)

    with FusionDefinition() as fd:
        fusion_func(fd, inps=inps)

    # By default, matmul will be be run through expr_eval scheduler.
    # Through setting the enable and disable options as below,
    # we can execute it through matmul scheduler. The above fusion will not
    # be accepted by the matmul scheduler since the outputs are of type Float and raises a RuntimeError.
    # Note: We use this error-based test since for compatible dtypes (float16/bfloat16),
    # the matmul scheduler ran into a scheduling error on H100. This test might be more robust against
    # changes in matmul scheduler in the interim.

    with pytest.raises(
        RuntimeError, match="Can not find a scheduler to schedule fusion segment"
    ):
        fd.execute(
            inps, _enable_options=["fuse_matmul"], _disable_options=["matmul_expr_eval"]
        )

# Test that we properly raise an error when passing inputs with the wrong types
def test_mismatched_input_types():
    scalar_inp = 2.0
    tensor_inp = torch.rand((15,), dtype=torch.float32, device="cuda:0")

    def fusion_func(fd: FusionDefinition):
        T0 = fd.define_tensor(
            shape=[-1],
            contiguity=[True],
            dtype=DataType.Float,
            is_cpu=False,
            stride_order=[0],
        )
        s0 = fd.define_scalar()
        T1 = fd.ops.mul(T0, s0)
        fd.add_output(T1)

    with FusionDefinition() as fd:
        fusion_func(fd)

    with pytest.raises(
        Exception,
        match="Expected input 0, .*, to be an at::Tensor but got scalar 2",
    ):
        nvf_out = fd.execute([scalar_inp, scalar_inp])

    with pytest.raises(
        Exception,
        match=re.escape(
            "Expected input 1, d2, to be a scalar but got float tensor of rank 1"
        ),
    ):
        nvf_out = fd.execute([tensor_inp, tensor_inp])

    with pytest.raises(
        Exception,
        match="Expected input 0, .*, to be bound to a tensor of dtype float, but got a tensor of dtype __half",
    ):
        wrong_tensor_inp = torch.rand((15,), dtype=torch.float16, device="cuda:0")
        nvf_out = fd.execute([wrong_tensor_inp, 2.0])

    with pytest.raises(
        Exception,
        match=re.escape(
            "Scalar value (2,1) is not compatible with the expected data type: double."
        ),
    ):
        nvf_out = fd.execute([tensor_inp, 2.0 + 1.0j])

def test_fusion_profiler_auto_scheduled():
    inputs = [
        torch.randn((2, 5), dtype=torch.float, device="cuda:0"),
        torch.randn((2, 5), dtype=torch.float, device="cuda:0"),
    ]

    def fusion_func(fd: FusionDefinition) -> None:
        T0 = fd.from_pytorch(inputs[0])
        T1 = fd.from_pytorch(inputs[1])
        T2 = fd.ops.add(T0, T1)
        T3 = fd.ops.sum(T2, dim=-1)
        T4 = fd.ops.sum(T3, dim=-1)
        fd.add_output(T4)

    with FusionDefinition() as fd:
        fusion_func(fd)

    # Testing that the profile returns 2 segments
    try:
        with PythonProfiler() as prof:
            fd.execute(inputs)
        assert prof.profile.segments == 2
        assert len(prof.profile.kernel_profiles) == 2
    except Exception as e:
        print(e)
        raise RuntimeError(
            "FusionDefinition did not run correctly with profile enabled! Error: "
            + str(e)
        )

def test_fusion_profiler_user_scheduled():
    inputs = [
        torch.randn((2, 5), dtype=torch.float, device="cuda:0"),
        torch.randn((2, 5), dtype=torch.float, device="cuda:0"),
    ]

    def fusion_func(fd: FusionDefinition) -> None:
        T0 = fd.from_pytorch(inputs[0])
        T1 = fd.from_pytorch(inputs[1])
        T2 = fd.ops.add(T0, T1)
        fd.add_output(T2)

    with FusionDefinition() as fd:
        fusion_func(fd)

    try:
        with PythonProfiler(auto_scheduled=False) as prof:
            fd.manual_execute(inputs)
        assert prof.profile.fusion_id >= 0
        assert prof.profile.kernel_profiles[0].scheduler == "user"
    except Exception as e:
        raise RuntimeError(
            "FusionDefinition did not run correctly with profile enabled! Error: "
            + str(e)
        )

def test_fusion_profiler_with_noncodegen_kernels():
    inputs = [
        torch.randn((2, 4, 16), dtype=torch.bfloat16, device="cuda:0"),
        torch.randn((2, 4, 16), dtype=torch.bfloat16, device="cuda:0"),
        torch.randn((16, 16), dtype=torch.bfloat16, device="cuda:0"),
    ]

    def fusion_func(fd: FusionDefinition) -> None:
        T0 = fd.from_pytorch(inputs[0])
        T1 = fd.from_pytorch(inputs[1])
        T2 = fd.from_pytorch(inputs[2])
        T3 = fd.ops.linear(T0, T2)
        T4 = fd.ops.add(T3, T1)
        fd.add_output(T4)

    with FusionDefinition() as fd:
        fusion_func(fd)

    try:
        with PythonProfiler() as prof:
            fd.execute(inputs)
        assert len(prof.profile.kernel_profiles) == 2
        assert len(prof.profile.kernel_profiles[0].name) > 0
        assert len(prof.profile.kernel_profiles[1].name) > 0
    except Exception as e:
        raise RuntimeError(
            "FusionDefinition did not run correctly with profile enabled! Error: "
            + str(e)
        )

def test_lru_fusion_cache_decorator():
    inputs = [
        torch.randn(2, 4, device="cuda"),
        torch.randn(2, 4, device="cuda"),
    ]

    def create_fd1(fd: FusionDefinition):
        t0 = fd.from_pytorch(inputs[0])
        t1 = fd.from_pytorch(inputs[1])
        t2 = fd.ops.div(t0, t1)
        fd.add_output(t2)

    def create_fd2(fd: FusionDefinition):
        t0 = fd.from_pytorch(inputs[0])
        t1 = fd.from_pytorch(inputs[1])
        t2 = fd.ops.div(t1, t0)
        fd.add_output(t2)

    @LruFusionCache(max_fusions=10)
    def create_fusion(create_fd: Callable):
        with FusionDefinition() as fd:
            create_fd(fd)
        return fd

    empty_stats = """Max Fusions Allowed: 10
The fusion cache is empty.\n"""
    assert create_fusion.stats() == empty_stats

    # Test LRU cache compilation
    for i in range(5):
        fd1 = create_fusion(create_fd1)
        outputs = fd1.execute(inputs)
        assert torch.allclose(outputs[0], inputs[0] / inputs[1])

        fd2 = create_fusion(create_fd2)
        outputs = fd2.execute(inputs)
        assert torch.allclose(outputs[0], inputs[1] / inputs[0])
        assert fd1.fusion == fd1.fusion
        assert fd2.fusion == fd2.fusion
        assert fd1.fusion != fd2.fusion

    expected_stats = """Max Fusions Allowed: 10
Total Fusions in Cache: 2
Total Unique Fusions Compiled: 2
Cache Hits by LRU ordering:
\t0 -> 4 hits
\t1 -> 4 hits
Cache Lookups: 10
Cache Hits: 8
Hit Rate: 80%\n"""
    assert create_fusion.stats() == expected_stats

def test_lru_cache():
    bcast_input = torch.randn(1, 4, device="cuda")
    full_input = torch.randn(2, 4, device="cuda")

    def nvfuser_fusion_id0(fd: FusionDefinition) -> None:
        T0 = fd.define_tensor(
            shape=[1, -1],
            contiguity=[None, True],
            dtype=DataType.Float,
# ... truncated ...


# All rights reserved.
# SPDX-License-Identifier: BSD-3-Clause
# Owner(s): ["module: nvfuser"]

from nvfuser_direct import (
    FusionDefinition,
    PythonProfiler,
    DataType,
    version,
    LruFusionCache,
    LRUCache,
)
import torch
import pytest
import io
import re
from contextlib import redirect_stdout, redirect_stderr
from typing import Callable

def test_python_version_API():
    from nvfuser_direct.nvfuser_direct_version import Version

    assert version() > "0.0.0"
    assert version() > Version("0.0.0")

def test_fusion_not_defined():
    inputs = [
        torch.randn(4, 4, device="cpu"),
    ]

    # A FusionDefinition object is constructed but not defined, should trip an error
    try:
        fd = FusionDefinition()
        out = fd.execute(inputs)
        raise RuntimeError("Expecting an error for an empty FusionDefinition!")
    except NotImplementedError as e:
        assert (
            "Fusion does not exist! Use `with FusionDefinition() as fd: ...` to define a fusion."
            in str(e)
        )

def test_fusion_empty():


    DataType,
    version,
    LruFusionCache,
    LRUCache,
)
import torch
import pytest
import io
import re
from contextlib import redirect_stdout, redirect_stderr
from typing import Callable

def test_python_version_API():
    from nvfuser_direct.nvfuser_direct_version import Version

    assert version() > "0.0.0"
    assert version() > Version("0.0.0")

def test_fusion_not_defined():
    inputs = [
        torch.randn(4, 4, device="cpu"),
    ]

    # A FusionDefinition object is constructed but not defined, should trip an error
    try:
        fd = FusionDefinition()
        out = fd.execute(inputs)
        raise RuntimeError("Expecting an error for an empty FusionDefinition!")
    except NotImplementedError as e:
        assert (
            "Fusion does not exist! Use `with FusionDefinition() as fd: ...` to define a fusion."
            in str(e)
        )

def test_fusion_empty():
    inputs = [
        torch.randn(4, 4, device="cpu"),
    ]

    # A FusionDefinition object is constructed but not defined, should trip an error
    with pytest.raises(NotImplementedError, match="Fusion is empty!"):
        with FusionDefinition() as fd:
            pass
        out = fd.execute(inputs)

def test_from_pytorch_fails_on_cpu_tensor():
    inputs = [
        torch.randn(4, 4, device="cpu"),
    ]

    stdout_capture = io.StringIO()


    assert version() > Version("0.0.0")

def test_fusion_not_defined():
    inputs = [
        torch.randn(4, 4, device="cpu"),
    ]

    # A FusionDefinition object is constructed but not defined, should trip an error
    try:
        fd = FusionDefinition()
        out = fd.execute(inputs)
        raise RuntimeError("Expecting an error for an empty FusionDefinition!")
    except NotImplementedError as e:
        assert (
            "Fusion does not exist! Use `with FusionDefinition() as fd: ...` to define a fusion."
            in str(e)
        )

def test_fusion_empty():
    inputs = [
        torch.randn(4, 4, device="cpu"),
    ]

    # A FusionDefinition object is constructed but not defined, should trip an error
    with pytest.raises(NotImplementedError, match="Fusion is empty!"):
        with FusionDefinition() as fd:
            pass
        out = fd.execute(inputs)

def test_from_pytorch_fails_on_cpu_tensor():
    inputs = [
        torch.randn(4, 4, device="cpu"),
    ]

    stdout_capture = io.StringIO()
    stderr_capture = io.StringIO()
    error_msg = (
        "Found unsupported device cpu, only scalar CPU or CUDA tensors are supported"
    )
    with pytest.raises(ValueError, match=error_msg), redirect_stdout(
        stdout_capture
    ), redirect_stderr(stderr_capture):
        with FusionDefinition() as fd:
            t0 = fd.from_pytorch(inputs[0])
            t1 = fd.ops.relu(t0)
            fd.add_output(t1)


        raise RuntimeError("Expecting an error for an empty FusionDefinition!")
    except NotImplementedError as e:
        assert (
            "Fusion does not exist! Use `with FusionDefinition() as fd: ...` to define a fusion."
            in str(e)
        )

def test_fusion_empty():
    inputs = [
        torch.randn(4, 4, device="cpu"),
    ]

    # A FusionDefinition object is constructed but not defined, should trip an error
    with pytest.raises(NotImplementedError, match="Fusion is empty!"):
        with FusionDefinition() as fd:
            pass
        out = fd.execute(inputs)

def test_from_pytorch_fails_on_cpu_tensor():
    inputs = [
        torch.randn(4, 4, device="cpu"),
    ]

    stdout_capture = io.StringIO()
    stderr_capture = io.StringIO()
    error_msg = (
        "Found unsupported device cpu, only scalar CPU or CUDA tensors are supported"
    )
    with pytest.raises(ValueError, match=error_msg), redirect_stdout(
        stdout_capture
    ), redirect_stderr(stderr_capture):
        with FusionDefinition() as fd:
            t0 = fd.from_pytorch(inputs[0])
            t1 = fd.ops.relu(t0)
            fd.add_output(t1)

def test_fusion_definition_print():
    with FusionDefinition() as fd:
        tv0 = fd.define_tensor(
            shape=[2, 4, 8],
        )
        tv1 = fd.define_tensor(
            shape=[2, 4, 8],
        )
        tv2 = fd.ops.add(tv0, tv1)
        fd.add_output(tv2)

    # Test fusion math representation
    prescheduled_fusion_definition = """Inputs:
  T0_g_float[iS0{2}, iS1{4}, iS2{8}]
  T1_g_float[iS3{2}, iS4{4}, iS5{8}]
Outputs:
  T2_g_float[iS6{2}, iS7{4}, iS8{8}]


def test_from_pytorch_fails_on_cpu_tensor():
    inputs = [
        torch.randn(4, 4, device="cpu"),
    ]

    stdout_capture = io.StringIO()
    stderr_capture = io.StringIO()
    error_msg = (
        "Found unsupported device cpu, only scalar CPU or CUDA tensors are supported"
    )
    with pytest.raises(ValueError, match=error_msg), redirect_stdout(
        stdout_capture
    ), redirect_stderr(stderr_capture):
        with FusionDefinition() as fd:
            t0 = fd.from_pytorch(inputs[0])
            t1 = fd.ops.relu(t0)
            fd.add_output(t1)

def test_fusion_definition_print():
    with FusionDefinition() as fd:
        tv0 = fd.define_tensor(
            shape=[2, 4, 8],
        )
        tv1 = fd.define_tensor(
            shape=[2, 4, 8],
        )
        tv2 = fd.ops.add(tv0, tv1)
        fd.add_output(tv2)

    # Test fusion math representation
    prescheduled_fusion_definition = """Inputs:
  T0_g_float[iS0{2}, iS1{4}, iS2{8}]
  T1_g_float[iS3{2}, iS4{4}, iS5{8}]
Outputs:
  T2_g_float[iS6{2}, iS7{4}, iS8{8}]

%kernel_math {
T2_g_float[iS6{2}, iS7{4}, iS8{8}]
   = T0_g_float[iS0{2}, iS1{4}, iS2{8}]
   + T1_g_float[iS3{2}, iS4{4}, iS5{8}];
} // %kernel_math \n\n"""
    assert fd.fusion.print_math() == prescheduled_fusion_definition

    # Test CUDA kernel representation
    cuda_kernel = """// Codegen generated code
__global__ void CUDAGeneratedKernel(Tensor<float, 3, 3> T0, Tensor<float, 3, 3> T1, Tensor<float, 3, 3> T2) {
  NVFUSER_DEFINE_MAGIC_ZERO;
  #pragma unroll
  for(nvfuser_index_t i0 = 0LL; i0 < 2LL; ++i0) {
    nvfuser_index_t i1;
    i1 = T0.alloc_stride[0LL] * i0;
    nvfuser_index_t i2;
    i2 = T1.alloc_stride[0LL] * i0;
    nvfuser_index_t i3;
    i3 = 32LL * i0;
    #pragma unroll
    for(nvfuser_index_t i4 = 0LL; i4 < 4LL; ++i4) {
      nvfuser_index_t i5;
      i5 = i1 + (T0.alloc_stride[1LL] * i4);
      nvfuser_index_t i6;
      i6 = i2 + (T1.alloc_stride[1LL] * i4);
      nvfuser_index_t i7;
      i7 = i3 + (8LL * i4);
      #pragma unroll
      for(nvfuser_index_t i8 = 0LL; i8 < 8LL; ++i8) {
        nvfuser_index_t i9;
        i9 = i8 + nvfuser_zero;
        T2[(i7 + i9)]
          = T0[(i5 + (T0.alloc_stride[2LL] * i9))]
          + T1[(i6 + (T1.alloc_stride[2LL] * i9))];
      }
    }
  }
  NVFUSER_UPDATE_MAGIC_ZERO;
}
"""
    assert fd.fusion.print_kernel() == cuda_kernel

    # Test TensorView string representation
    assert str(tv0) == r"T0_g_float[iS0{2}, iS1{4}, iS2{8}]"
    assert str(tv1) == r"T1_g_float[iS3{2}, iS4{4}, iS5{8}]"
    assert str(tv2) == r"T2_g_float[iS6{2}, iS7{4}, iS8{8}]"

    # Test TensorDomain string representation
    assert str(tv0.domain()) == r"[iS0{2}, iS1{4}, iS2{8}]"
    assert str(tv1.domain()) == r"[iS3{2}, iS4{4}, iS5{8}]"
    assert str(tv2.domain()) == r"[iS6{2}, iS7{4}, iS8{8}]"

    # Test IterDomain string representation
    assert str(tv0.axis(0)) == r"iS0{2}"
    assert str(tv1.axis(0)) == r"iS3{2}"
    assert str(tv2.axis(0)) == r"iS6{2}"

    # Test axis extents
    assert str(tv0.axis(0).extent()) == r"2"
    assert str(tv1.axis(0).extent()) == r"2"
    assert str(tv2.axis(0).extent()) == r"2"

def test_fusion_execution_cache():
    with FusionDefinition() as fd:
        tv0 = fd.define_tensor(
            shape=[2, 4, 8],
        )
        tv1 = fd.define_tensor(
            shape=[2, 4, 8],
        )
        tv2 = fd.ops.add(tv0, tv1)
        fd.add_output(tv2)

    # Test fusion execution
    inputs = [
        torch.randn(2, 4, 8, device="cuda"),
        torch.randn(2, 4, 8, device="cuda"),
    ]
    results = fd.execute(inputs)
    torch.testing.assert_close(results[0], inputs[0] + inputs[1])


    assert str(tv0) == r"T0_g_float[iS0{2}, iS1{4}, iS2{8}]"
    assert str(tv1) == r"T1_g_float[iS3{2}, iS4{4}, iS5{8}]"
    assert str(tv2) == r"T2_g_float[iS6{2}, iS7{4}, iS8{8}]"

    # Test TensorDomain string representation
    assert str(tv0.domain()) == r"[iS0{2}, iS1{4}, iS2{8}]"
    assert str(tv1.domain()) == r"[iS3{2}, iS4{4}, iS5{8}]"
    assert str(tv2.domain()) == r"[iS6{2}, iS7{4}, iS8{8}]"

    # Test IterDomain string representation
    assert str(tv0.axis(0)) == r"iS0{2}"
    assert str(tv1.axis(0)) == r"iS3{2}"
    assert str(tv2.axis(0)) == r"iS6{2}"

    # Test axis extents
    assert str(tv0.axis(0).extent()) == r"2"
    assert str(tv1.axis(0).extent()) == r"2"
    assert str(tv2.axis(0).extent()) == r"2"

def test_fusion_execution_cache():
    with FusionDefinition() as fd:
        tv0 = fd.define_tensor(
            shape=[2, 4, 8],
        )
        tv1 = fd.define_tensor(
            shape=[2, 4, 8],
        )
        tv2 = fd.ops.add(tv0, tv1)
        fd.add_output(tv2)

    # Test fusion execution
    inputs = [
        torch.randn(2, 4, 8, device="cuda"),
        torch.randn(2, 4, 8, device="cuda"),
    ]
    results = fd.execute(inputs)
    torch.testing.assert_close(results[0], inputs[0] + inputs[1])

    # Check that TensorViews are accessible after running fusion
    assert str(tv0.domain()) == "[iS0{2}, iS1{4}, iS2{8}]"
    assert str(tv1.domain()) == "[iS3{2}, iS4{4}, iS5{8}]"
    assert str(tv2.domain()) == "[iS6{2}, iS7{4}, iS8{8}]"

    with pytest.raises(
        AssertionError,
        match=r"FusionExecutorCache already exists! Use execute\(\) to execute the fusion.",
    ):
        fd.manual_execute(inputs)

    # Test fusion math representation after compilation
    prescheduled_fusion_definition = """Inputs:
  T0_g_float[iS0{2}, iS1{4}, iS2{8}]
  T1_g_float[iS3{2}, iS4{4}, iS5{8}]
Outputs:
  T2_g_float[iS6{2}, iS7{4}, iS8{8}]

%kernel_math {
T2_g_float[iS6{2}, iS7{4}, iS8{8}]
   = T0_g_float[iS0{2}, iS1{4}, iS2{8}]
   + T1_g_float[iS3{2}, iS4{4}, iS5{8}];
} // %kernel_math \n\n"""
    assert fd.fusion.print_math() == prescheduled_fusion_definition

    # Test compilation status
    assert fd.fec.is_compiled(inputs)

    # Test scheduled IR representation
    scheduled_ir = """Inputs:
  T0_g_float[iS58{1}, iS59{1}, iS57{128}]
  T1_g_float[iS46{1}, iS47{1}, iS45{128}]
Outputs:
  T2_g_float[iblockIdx.x28{1}, iUS29{1}, ithreadIdx.x27{128}] ca_pos( 2 ) produce_pos( 3 )

%kernel {
T3_l_float[iblockIdx.x52{1}, iUS53{1}, ithreadIdx.x51{128}] ca_pos( 2 )
   = Set( T0_g_float[iS58{1}, iS59{1}, iS57{128}], cache_op=Streaming )
T4_l_float[iblockIdx.x40{1}, iUS41{1}, ithreadIdx.x39{128}] ca_pos( 2 )
   = Set( T1_g_float[iS46{1}, iS47{1}, iS45{128}], cache_op=Streaming )
T5_l_float[iblockIdx.x34{1}, iUS35{1}, ithreadIdx.x33{128}] ca_pos( 3 ) produce_pos( 2 )
   = T3_l_float[iblockIdx.x52{1}, iUS53{1}, ithreadIdx.x51{128}] ca_pos( 2 )
   + T4_l_float[iblockIdx.x40{1}, iUS41{1}, ithreadIdx.x39{128}] ca_pos( 2 );
T2_g_float[iblockIdx.x28{1}, iUS29{1}, ithreadIdx.x27{128}] ca_pos( 2 ) produce_pos( 3 )
   = Set( T5_l_float[iblockIdx.x34{1}, iUS35{1}, ithreadIdx.x33{128}] ca_pos( 3 ) produce_pos( 2 ), cache_op=Streaming )
} // %kernel\n"""
    assert fd.fec.get_scheduled_ir(inputs) == scheduled_ir
    assert fd.fec.get_most_recent_scheduled_ir() == scheduled_ir

    # Test CUDA kernel representation
    cuda_kernel = """// Codegen generated code
__global__ void nvfuser_pointwise_f0_c1_r0_g0(Tensor<float, 3, 3> T0, Tensor<float, 3, 3> T1, Tensor<float, 3, 3> T2) {
  nvfuser_index_t i0;
  i0 = ((nvfuser_index_t)threadIdx.x) % 32;
  nvfuser_index_t i1;
  i1 = ((nvfuser_index_t)threadIdx.x) / 32;
  nvfuser_index_t i2;
  i2 = i0 / 8;
  nvfuser_index_t i3;
  i3 = i0 % 8;
  if ((((nvfuser_index_t)threadIdx.x) < 64)) {
    Array<float, 1, 1> T4;
    T4[0] = 0;
    T4[0]
       = T1[(((T1.alloc_stride[0LL] * i1) + (T1.alloc_stride[1LL] * i2)) + (T1.alloc_stride[2LL] * i3))];
    Array<float, 1, 1> T3;
    T3[0] = 0;
    T3[0]
       = T0[(((T0.alloc_stride[0LL] * i1) + (T0.alloc_stride[1LL] * i2)) + (T0.alloc_stride[2LL] * i3))];
    Array<float, 1, 1> T5;
    T5[0]
      = T3[0]
      + T4[0];
    T2[((nvfuser_index_t)threadIdx.x)]
       = T5[0];
  }
}\n"""
    assert fd.fec.get_cuda_kernel(inputs) == cuda_kernel

def test_kernel_executor():
    with FusionDefinition() as fd:
        tv0 = fd.define_tensor(
            shape=[2, 4, 8],
        )
        tv1 = fd.define_tensor(
            shape=[2, 4, 8],
        )
        tv2 = fd.ops.add(tv0, tv1)
        fd.add_output(tv2)

    # Test fusion execution
    inputs = [
        torch.randn(2, 4, 8, device="cuda"),
        torch.randn(2, 4, 8, device="cuda"),
    ]
    results = fd.manual_execute(inputs)
    torch.testing.assert_close(results[0], inputs[0] + inputs[1])


  if ((((nvfuser_index_t)threadIdx.x) < 64)) {
    Array<float, 1, 1> T4;
    T4[0] = 0;
    T4[0]
       = T1[(((T1.alloc_stride[0LL] * i1) + (T1.alloc_stride[1LL] * i2)) + (T1.alloc_stride[2LL] * i3))];
    Array<float, 1, 1> T3;
    T3[0] = 0;
    T3[0]
       = T0[(((T0.alloc_stride[0LL] * i1) + (T0.alloc_stride[1LL] * i2)) + (T0.alloc_stride[2LL] * i3))];
    Array<float, 1, 1> T5;
    T5[0]
      = T3[0]
      + T4[0];
    T2[((nvfuser_index_t)threadIdx.x)]
       = T5[0];
  }
}\n"""
    assert fd.fec.get_cuda_kernel(inputs) == cuda_kernel

def test_kernel_executor():
    with FusionDefinition() as fd:
        tv0 = fd.define_tensor(
            shape=[2, 4, 8],
        )
        tv1 = fd.define_tensor(
            shape=[2, 4, 8],
        )
        tv2 = fd.ops.add(tv0, tv1)
        fd.add_output(tv2)

    # Test fusion execution
    inputs = [
        torch.randn(2, 4, 8, device="cuda"),
        torch.randn(2, 4, 8, device="cuda"),
    ]
    results = fd.manual_execute(inputs)
    torch.testing.assert_close(results[0], inputs[0] + inputs[1])
    assert fd.ke.is_compiled()

    with pytest.raises(
        AssertionError,
        match=r"KernelExecutor already exists! Use manual_execute\(\) to execute the fusion.",
    ):
        fd.execute(inputs)

def test_repro_script_for():
    inputs = [
        torch.ones(2, 4, 8, device="cuda"),
        torch.ones(2, 4, 8, device="cuda"),
    ]

    with FusionDefinition() as fd:
        t0 = fd.from_pytorch(inputs[0])
        t1 = fd.from_pytorch(inputs[1])
        c0 = fd.define_scalar(3.0)

        t2 = fd.ops.add(t0, t1)
        t3 = fd.ops.mul(t2, c0)
        t4 = fd.ops.sum(t3, [1], False, DataType.Float)

        fd.add_output(t4)

    expected_repro = """


        )
        tv2 = fd.ops.add(tv0, tv1)
        fd.add_output(tv2)

    # Test fusion execution
    inputs = [
        torch.randn(2, 4, 8, device="cuda"),
        torch.randn(2, 4, 8, device="cuda"),
    ]
    results = fd.manual_execute(inputs)
    torch.testing.assert_close(results[0], inputs[0] + inputs[1])
    assert fd.ke.is_compiled()

    with pytest.raises(
        AssertionError,
        match=r"KernelExecutor already exists! Use manual_execute\(\) to execute the fusion.",
    ):
        fd.execute(inputs)

def test_repro_script_for():
    inputs = [
        torch.ones(2, 4, 8, device="cuda"),
        torch.ones(2, 4, 8, device="cuda"),
    ]

    with FusionDefinition() as fd:
        t0 = fd.from_pytorch(inputs[0])
        t1 = fd.from_pytorch(inputs[1])
        c0 = fd.define_scalar(3.0)

        t2 = fd.ops.add(t0, t1)
        t3 = fd.ops.mul(t2, c0)
        t4 = fd.ops.sum(t3, [1], False, DataType.Float)

        fd.add_output(t4)

    expected_repro = """
import torch
from nvfuser_direct import FusionDefinition, DataType
def nvfuser_fusion(fd : FusionDefinition) -> None :
    tv0 = fd.define_tensor(shape=[-1, -1, -1], contiguity=[True, True, True], dtype=DataType.Float, is_cpu=False)
    tv1 = fd.define_tensor(shape=[-1, -1, -1], contiguity=[True, True, True], dtype=DataType.Float, is_cpu=False)
    tv2 = fd.ops.add(tv0, tv1)
    tv3 = fd.ops.mul(tv2, 3.00000)
    tv4 = fd.ops.sum(tv3, dims=[1], dtype=DataType.Float)
    fd.add_output(tv4)
with FusionDefinition() as fd:
    nvfuser_fusion(fd)
"""

    expected_inputs = """
inputs = [
    torch.testing.make_tensor((2, 4, 8), dtype=torch.float32, device='cuda:0'),
    torch.testing.make_tensor((2, 4, 8), dtype=torch.float32, device='cuda:0'),
]
fd.execute(inputs)\n"""

    # fd.repro_script_for(inputs) should have input arguments
    repro_with_inputs = fd.repro_script_for(inputs)
    assert expected_repro in repro_with_inputs
    assert expected_inputs in repro_with_inputs

    # fd.repro_script_for() should NOT have input arguments
    repro_without_inputs = fd.repro_script_for()
    assert expected_repro in repro_without_inputs
    assert expected_inputs not in repro_without_inputs

    # Check last_repro_script fails gracefully.
    with pytest.raises(
        AssertionError,
        match=r"fd.last_repro_script\(\) cannot provide a repro because fd.execute\(inputs, save_repro_inputs=True\) was not executed!",
    ):
        fd.last_repro_script()

    # Test fd.execute(inputs, save_repro_inputs=True) ; fd.last_repro_script()
    fd.execute(inputs, save_repro_inputs=True)
    last_repro = fd.last_repro_script()
    assert repro_with_inputs == last_repro

def test_define_tensor():
    with FusionDefinition() as fd:
        tv0 = fd.define_tensor(
            shape=[-1, -1, -1],
            contiguity=[True, True, True],
            dtype=DataType.Float,
            is_cpu=False,
            stride_order=[2, 1, 0],
        )
        tv1 = fd.define_tensor(
            shape=[-1, -1, -1],
            contiguity=[True, True, True],
            dtype=DataType.Float,
            is_cpu=False,
            stride_order=[2, 1, 0],
        )
        tv2 = fd.ops.add(tv0, tv1)
        fd.add_output(tv2)


    assert expected_inputs in repro_with_inputs

    # fd.repro_script_for() should NOT have input arguments
    repro_without_inputs = fd.repro_script_for()
    assert expected_repro in repro_without_inputs
    assert expected_inputs not in repro_without_inputs

    # Check last_repro_script fails gracefully.
    with pytest.raises(
        AssertionError,
        match=r"fd.last_repro_script\(\) cannot provide a repro because fd.execute\(inputs, save_repro_inputs=True\) was not executed!",
    ):
        fd.last_repro_script()

    # Test fd.execute(inputs, save_repro_inputs=True) ; fd.last_repro_script()
    fd.execute(inputs, save_repro_inputs=True)
    last_repro = fd.last_repro_script()
    assert repro_with_inputs == last_repro

def test_define_tensor():
    with FusionDefinition() as fd:
        tv0 = fd.define_tensor(
            shape=[-1, -1, -1],
            contiguity=[True, True, True],
            dtype=DataType.Float,
            is_cpu=False,
            stride_order=[2, 1, 0],
        )
        tv1 = fd.define_tensor(
            shape=[-1, -1, -1],
            contiguity=[True, True, True],
            dtype=DataType.Float,
            is_cpu=False,
            stride_order=[2, 1, 0],
        )
        tv2 = fd.ops.add(tv0, tv1)
        fd.add_output(tv2)

    # Test fusion execution with dynamic shapes
    inputs = [
        torch.randn(2, 4, 8, device="cuda"),
        torch.randn(2, 4, 8, device="cuda"),
    ]
    outputs = fd.execute(inputs)
    torch.testing.assert_close(outputs[0], inputs[0] + inputs[1])

@pytest.mark.skipif(
    torch.cuda.device_count() < 2,
    reason="test_selected_device requires multiple GPUs",
)
def test_execute_with_different_device():
    inputs = [
        torch.ones(2, 4, 8, device="cuda:1"),
        torch.ones(2, 4, 8, device="cuda:1"),
    ]

    with FusionDefinition() as fd:
        t0 = fd.from_pytorch(inputs[0])
        t1 = fd.from_pytorch(inputs[1])
        c0 = fd.define_scalar(3.0)

        t2 = fd.ops.add(t0, t1)
        t3 = fd.ops.mul(t2, c0)
        t4 = fd.ops.sum(t3, [1], False, DataType.Float)


            dtype=DataType.Float,
            is_cpu=False,
            stride_order=[2, 1, 0],
        )
        tv2 = fd.ops.add(tv0, tv1)
        fd.add_output(tv2)

    # Test fusion execution with dynamic shapes
    inputs = [
        torch.randn(2, 4, 8, device="cuda"),
        torch.randn(2, 4, 8, device="cuda"),
    ]
    outputs = fd.execute(inputs)
    torch.testing.assert_close(outputs[0], inputs[0] + inputs[1])

@pytest.mark.skipif(
    torch.cuda.device_count() < 2,
    reason="test_selected_device requires multiple GPUs",
)
def test_execute_with_different_device():
    inputs = [
        torch.ones(2, 4, 8, device="cuda:1"),
        torch.ones(2, 4, 8, device="cuda:1"),
    ]

    with FusionDefinition() as fd:
        t0 = fd.from_pytorch(inputs[0])
        t1 = fd.from_pytorch(inputs[1])
        c0 = fd.define_scalar(3.0)

        t2 = fd.ops.add(t0, t1)
        t3 = fd.ops.mul(t2, c0)
        t4 = fd.ops.sum(t3, [1], False, DataType.Float)

        fd.add_output(t4)

    outputs = fd.execute(inputs, device="cuda:1")
    assert len(outputs) == 1
    assert outputs[0].device.index == 1

def test_enable_disable_options():
    m = 24
    n = 16
    k = 8
    inps = [
        torch.randn(m, k, device="cuda", dtype=torch.float),
        torch.randn(k, n, device="cuda", dtype=torch.float),
    ]

    def fusion_func(fd: FusionDefinition, inps) -> None:
        t0 = fd.from_pytorch(inps[0])
        t1 = fd.from_pytorch(inps[1])
        t2 = fd.ops.matmul(t0, t1)
        fd.add_output(t2)

    with FusionDefinition() as fd:
        fusion_func(fd, inps=inps)


        torch.ones(2, 4, 8, device="cuda:1"),
        torch.ones(2, 4, 8, device="cuda:1"),
    ]

    with FusionDefinition() as fd:
        t0 = fd.from_pytorch(inputs[0])
        t1 = fd.from_pytorch(inputs[1])
        c0 = fd.define_scalar(3.0)

        t2 = fd.ops.add(t0, t1)
        t3 = fd.ops.mul(t2, c0)
        t4 = fd.ops.sum(t3, [1], False, DataType.Float)

        fd.add_output(t4)

    outputs = fd.execute(inputs, device="cuda:1")
    assert len(outputs) == 1
    assert outputs[0].device.index == 1

def test_enable_disable_options():
    m = 24
    n = 16
    k = 8
    inps = [
        torch.randn(m, k, device="cuda", dtype=torch.float),
        torch.randn(k, n, device="cuda", dtype=torch.float),
    ]

    def fusion_func(fd: FusionDefinition, inps) -> None:
        t0 = fd.from_pytorch(inps[0])
        t1 = fd.from_pytorch(inps[1])
        t2 = fd.ops.matmul(t0, t1)
        fd.add_output(t2)

    with FusionDefinition() as fd:
        fusion_func(fd, inps=inps)

    # By default, matmul will be be run through expr_eval scheduler.
    # Through setting the enable and disable options as below,
    # we can execute it through matmul scheduler. The above fusion will not
    # be accepted by the matmul scheduler since the outputs are of type Float and raises a RuntimeError.
    # Note: We use this error-based test since for compatible dtypes (float16/bfloat16),
    # the matmul scheduler ran into a scheduling error on H100. This test might be more robust against
    # changes in matmul scheduler in the interim.

    with pytest.raises(
        RuntimeError, match="Can not find a scheduler to schedule fusion segment"
    ):
        fd.execute(
            inps, _enable_options=["fuse_matmul"], _disable_options=["matmul_expr_eval"]
        )

# Test that we properly raise an error when passing inputs with the wrong types
def test_mismatched_input_types():
    scalar_inp = 2.0
    tensor_inp = torch.rand((15,), dtype=torch.float32, device="cuda:0")

    def fusion_func(fd: FusionDefinition):
        T0 = fd.define_tensor(
            shape=[-1],
            contiguity=[True],
            dtype=DataType.Float,
            is_cpu=False,
            stride_order=[0],
        )
        s0 = fd.define_scalar()
        T1 = fd.ops.mul(T0, s0)
        fd.add_output(T1)

    with FusionDefinition() as fd:


    with FusionDefinition() as fd:
        fusion_func(fd, inps=inps)

    # By default, matmul will be be run through expr_eval scheduler.
    # Through setting the enable and disable options as below,
    # we can execute it through matmul scheduler. The above fusion will not
    # be accepted by the matmul scheduler since the outputs are of type Float and raises a RuntimeError.
    # Note: We use this error-based test since for compatible dtypes (float16/bfloat16),
    # the matmul scheduler ran into a scheduling error on H100. This test might be more robust against
    # changes in matmul scheduler in the interim.

    with pytest.raises(
        RuntimeError, match="Can not find a scheduler to schedule fusion segment"
    ):
        fd.execute(
            inps, _enable_options=["fuse_matmul"], _disable_options=["matmul_expr_eval"]
        )

# Test that we properly raise an error when passing inputs with the wrong types
def test_mismatched_input_types():
    scalar_inp = 2.0
    tensor_inp = torch.rand((15,), dtype=torch.float32, device="cuda:0")

    def fusion_func(fd: FusionDefinition):
        T0 = fd.define_tensor(
            shape=[-1],
            contiguity=[True],
            dtype=DataType.Float,
            is_cpu=False,
            stride_order=[0],
        )
        s0 = fd.define_scalar()
        T1 = fd.ops.mul(T0, s0)
        fd.add_output(T1)

    with FusionDefinition() as fd:
        fusion_func(fd)

    with pytest.raises(
        Exception,
        match="Expected input 0, .*, to be an at::Tensor but got scalar 2",
    ):
        nvf_out = fd.execute([scalar_inp, scalar_inp])

    with pytest.raises(
        Exception,
        match=re.escape(
            "Expected input 1, d2, to be a scalar but got float tensor of rank 1"
        ),
    ):
        nvf_out = fd.execute([tensor_inp, tensor_inp])

    with pytest.raises(
        Exception,
        match="Expected input 0, .*, to be bound to a tensor of dtype float, but got a tensor of dtype __half",
    ):
        wrong_tensor_inp = torch.rand((15,), dtype=torch.float16, device="cuda:0")
        nvf_out = fd.execute([wrong_tensor_inp, 2.0])

    with pytest.raises(
        Exception,
        match=re.escape(
            "Scalar value (2,1) is not compatible with the expected data type: double."
        ),
    ):
        nvf_out = fd.execute([tensor_inp, 2.0 + 1.0j])

def test_fusion_profiler_auto_scheduled():
    inputs = [
        torch.randn((2, 5), dtype=torch.float, device="cuda:0"),
        torch.randn((2, 5), dtype=torch.float, device="cuda:0"),
    ]

    def fusion_func(fd: FusionDefinition) -> None:
        T0 = fd.from_pytorch(inputs[0])
        T1 = fd.from_pytorch(inputs[1])
        T2 = fd.ops.add(T0, T1)
        T3 = fd.ops.sum(T2, dim=-1)
        T4 = fd.ops.sum(T3, dim=-1)
        fd.add_output(T4)

    with FusionDefinition() as fd:
        fusion_func(fd)

    # Testing that the profile returns 2 segments


        ),
    ):
        nvf_out = fd.execute([tensor_inp, tensor_inp])

    with pytest.raises(
        Exception,
        match="Expected input 0, .*, to be bound to a tensor of dtype float, but got a tensor of dtype __half",
    ):
        wrong_tensor_inp = torch.rand((15,), dtype=torch.float16, device="cuda:0")
        nvf_out = fd.execute([wrong_tensor_inp, 2.0])

    with pytest.raises(
        Exception,
        match=re.escape(
            "Scalar value (2,1) is not compatible with the expected data type: double."
        ),
    ):
        nvf_out = fd.execute([tensor_inp, 2.0 + 1.0j])

def test_fusion_profiler_auto_scheduled():
    inputs = [
        torch.randn((2, 5), dtype=torch.float, device="cuda:0"),
        torch.randn((2, 5), dtype=torch.float, device="cuda:0"),
    ]

    def fusion_func(fd: FusionDefinition) -> None:
        T0 = fd.from_pytorch(inputs[0])
        T1 = fd.from_pytorch(inputs[1])
        T2 = fd.ops.add(T0, T1)
        T3 = fd.ops.sum(T2, dim=-1)
        T4 = fd.ops.sum(T3, dim=-1)
        fd.add_output(T4)

    with FusionDefinition() as fd:
        fusion_func(fd)

    # Testing that the profile returns 2 segments
    try:
        with PythonProfiler() as prof:
            fd.execute(inputs)
        assert prof.profile.segments == 2
        assert len(prof.profile.kernel_profiles) == 2
    except Exception as e:
        print(e)
        raise RuntimeError(
            "FusionDefinition did not run correctly with profile enabled! Error: "
            + str(e)
        )

def test_fusion_profiler_user_scheduled():
    inputs = [
        torch.randn((2, 5), dtype=torch.float, device="cuda:0"),
        torch.randn((2, 5), dtype=torch.float, device="cuda:0"),
    ]

    def fusion_func(fd: FusionDefinition) -> None:
        T0 = fd.from_pytorch(inputs[0])
        T1 = fd.from_pytorch(inputs[1])
        T2 = fd.ops.add(T0, T1)
        fd.add_output(T2)

    with FusionDefinition() as fd:
        fusion_func(fd)

    try:
        with PythonProfiler(auto_scheduled=False) as prof:
            fd.manual_execute(inputs)


        T4 = fd.ops.sum(T3, dim=-1)
        fd.add_output(T4)

    with FusionDefinition() as fd:
        fusion_func(fd)

    # Testing that the profile returns 2 segments
    try:
        with PythonProfiler() as prof:
            fd.execute(inputs)
        assert prof.profile.segments == 2
        assert len(prof.profile.kernel_profiles) == 2
    except Exception as e:
        print(e)
        raise RuntimeError(
            "FusionDefinition did not run correctly with profile enabled! Error: "
            + str(e)
        )

def test_fusion_profiler_user_scheduled():
    inputs = [
        torch.randn((2, 5), dtype=torch.float, device="cuda:0"),
        torch.randn((2, 5), dtype=torch.float, device="cuda:0"),
    ]

    def fusion_func(fd: FusionDefinition) -> None:
        T0 = fd.from_pytorch(inputs[0])
        T1 = fd.from_pytorch(inputs[1])
        T2 = fd.ops.add(T0, T1)
        fd.add_output(T2)

    with FusionDefinition() as fd:
        fusion_func(fd)

    try:
        with PythonProfiler(auto_scheduled=False) as prof:
            fd.manual_execute(inputs)
        assert prof.profile.fusion_id >= 0
        assert prof.profile.kernel_profiles[0].scheduler == "user"
    except Exception as e:
        raise RuntimeError(
            "FusionDefinition did not run correctly with profile enabled! Error: "
            + str(e)
        )

def test_fusion_profiler_with_noncodegen_kernels():
    inputs = [
        torch.randn((2, 4, 16), dtype=torch.bfloat16, device="cuda:0"),
        torch.randn((2, 4, 16), dtype=torch.bfloat16, device="cuda:0"),
        torch.randn((16, 16), dtype=torch.bfloat16, device="cuda:0"),
    ]

    def fusion_func(fd: FusionDefinition) -> None:
        T0 = fd.from_pytorch(inputs[0])
        T1 = fd.from_pytorch(inputs[1])
        T2 = fd.from_pytorch(inputs[2])
        T3 = fd.ops.linear(T0, T2)
        T4 = fd.ops.add(T3, T1)
        fd.add_output(T4)

    with FusionDefinition() as fd:
        fusion_func(fd)


        T0 = fd.from_pytorch(inputs[0])
        T1 = fd.from_pytorch(inputs[1])
        T2 = fd.ops.add(T0, T1)
        fd.add_output(T2)

    with FusionDefinition() as fd:
        fusion_func(fd)

    try:
        with PythonProfiler(auto_scheduled=False) as prof:
            fd.manual_execute(inputs)
        assert prof.profile.fusion_id >= 0
        assert prof.profile.kernel_profiles[0].scheduler == "user"
    except Exception as e:
        raise RuntimeError(
            "FusionDefinition did not run correctly with profile enabled! Error: "
            + str(e)
        )

def test_fusion_profiler_with_noncodegen_kernels():
    inputs = [
        torch.randn((2, 4, 16), dtype=torch.bfloat16, device="cuda:0"),
        torch.randn((2, 4, 16), dtype=torch.bfloat16, device="cuda:0"),
        torch.randn((16, 16), dtype=torch.bfloat16, device="cuda:0"),
    ]

    def fusion_func(fd: FusionDefinition) -> None:
        T0 = fd.from_pytorch(inputs[0])
        T1 = fd.from_pytorch(inputs[1])
        T2 = fd.from_pytorch(inputs[2])
        T3 = fd.ops.linear(T0, T2)
        T4 = fd.ops.add(T3, T1)
        fd.add_output(T4)

    with FusionDefinition() as fd:
        fusion_func(fd)

    try:
        with PythonProfiler() as prof:
            fd.execute(inputs)
        assert len(prof.profile.kernel_profiles) == 2
        assert len(prof.profile.kernel_profiles[0].name) > 0
        assert len(prof.profile.kernel_profiles[1].name) > 0
    except Exception as e:
        raise RuntimeError(
            "FusionDefinition did not run correctly with profile enabled! Error: "
            + str(e)
        )

def test_lru_fusion_cache_decorator():
    inputs = [
        torch.randn(2, 4, device="cuda"),
        torch.randn(2, 4, device="cuda"),
    ]

    def create_fd1(fd: FusionDefinition):
        t0 = fd.from_pytorch(inputs[0])
        t1 = fd.from_pytorch(inputs[1])
        t2 = fd.ops.div(t0, t1)
        fd.add_output(t2)

    def create_fd2(fd: FusionDefinition):
        t0 = fd.from_pytorch(inputs[0])
        t1 = fd.from_pytorch(inputs[1])
        t2 = fd.ops.div(t1, t0)
        fd.add_output(t2)


        T3 = fd.ops.linear(T0, T2)
        T4 = fd.ops.add(T3, T1)
        fd.add_output(T4)

    with FusionDefinition() as fd:
        fusion_func(fd)

    try:
        with PythonProfiler() as prof:
            fd.execute(inputs)
        assert len(prof.profile.kernel_profiles) == 2
        assert len(prof.profile.kernel_profiles[0].name) > 0
        assert len(prof.profile.kernel_profiles[1].name) > 0
    except Exception as e:
        raise RuntimeError(
            "FusionDefinition did not run correctly with profile enabled! Error: "
            + str(e)
        )

def test_lru_fusion_cache_decorator():
    inputs = [
        torch.randn(2, 4, device="cuda"),
        torch.randn(2, 4, device="cuda"),
    ]

    def create_fd1(fd: FusionDefinition):
        t0 = fd.from_pytorch(inputs[0])
        t1 = fd.from_pytorch(inputs[1])
        t2 = fd.ops.div(t0, t1)
        fd.add_output(t2)

    def create_fd2(fd: FusionDefinition):
        t0 = fd.from_pytorch(inputs[0])
        t1 = fd.from_pytorch(inputs[1])
        t2 = fd.ops.div(t1, t0)
        fd.add_output(t2)

    @LruFusionCache(max_fusions=10)
    def create_fusion(create_fd: Callable):
        with FusionDefinition() as fd:
            create_fd(fd)
        return fd

    empty_stats = """Max Fusions Allowed: 10
The fusion cache is empty.\n"""
    assert create_fusion.stats() == empty_stats

    # Test LRU cache compilation
    for i in range(5):
        fd1 = create_fusion(create_fd1)
        outputs = fd1.execute(inputs)
        assert torch.allclose(outputs[0], inputs[0] / inputs[1])

        fd2 = create_fusion(create_fd2)
        outputs = fd2.execute(inputs)
        assert torch.allclose(outputs[0], inputs[1] / inputs[0])
        assert fd1.fusion == fd1.fusion
        assert fd2.fusion == fd2.fusion
        assert fd1.fusion != fd2.fusion

    expected_stats = """Max Fusions Allowed: 10
Total Fusions in Cache: 2
Total Unique Fusions Compiled: 2
Cache Hits by LRU ordering:
\t0 -> 4 hits
\t1 -> 4 hits
Cache Lookups: 10
Cache Hits: 8
Hit Rate: 80%\n"""
    assert create_fusion.stats() == expected_stats

def test_lru_cache():
    bcast_input = torch.randn(1, 4, device="cuda")
    full_input = torch.randn(2, 4, device="cuda")

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


        fd2 = create_fusion(create_fd2)
        outputs = fd2.execute(inputs)
        assert torch.allclose(outputs[0], inputs[1] / inputs[0])
        assert fd1.fusion == fd1.fusion
        assert fd2.fusion == fd2.fusion
        assert fd1.fusion != fd2.fusion

    expected_stats = """Max Fusions Allowed: 10
Total Fusions in Cache: 2
Total Unique Fusions Compiled: 2
Cache Hits by LRU ordering:
\t0 -> 4 hits
\t1 -> 4 hits
Cache Lookups: 10
Cache Hits: 8
Hit Rate: 80%\n"""
    assert create_fusion.stats() == expected_stats

def test_lru_cache():
    bcast_input = torch.randn(1, 4, device="cuda")
    full_input = torch.randn(2, 4, device="cuda")

    def nvfuser_fusion_id0(fd: FusionDefinition) -> None:
        T0 = fd.define_tensor(
            shape=[1, -1],
            contiguity=[None, True],
            dtype=DataType.Float,
            is_cpu=False,
            stride_order=[1, 0],
        )
        T1 = fd.define_tensor(
            shape=[-1, -1],
            contiguity=[True, True],
            dtype=DataType.Float,
            is_cpu=False,
            stride_order=[1, 0],
        )
        T2 = fd.ops.sub(T0, T1)
        fd.add_output(T2)

    # nvfuser_fusion_id1 reverses the input argument order of nvfuser_fusion_id0
    def nvfuser_fusion_id1(fd: FusionDefinition) -> None:
        T1 = fd.define_tensor(
            shape=[-1, -1],
            contiguity=[True, True],
            dtype=DataType.Float,
            is_cpu=False,
            stride_order=[1, 0],
        )
        T0 = fd.define_tensor(
            shape=[1, -1],
            contiguity=[None, True],
            dtype=DataType.Float,
            is_cpu=False,
            stride_order=[1, 0],
        )
        T2 = fd.ops.sub(T0, T1)
        fd.add_output(T2)

    cache = LRUCache(max_fusions=10)

    with FusionDefinition() as fd0:
        nvfuser_fusion_id0(fd0)

    assert cache.num_fusions() == 0
    fd0.fec = cache.cache_compile(fd0.fusion)
    del fd0._fusion
    # Check that LRUCache caches the first fusion
    assert cache.num_fusions() == 1

    with FusionDefinition() as fd1:
        nvfuser_fusion_id1(fd1)

    assert cache.num_fusions() == 1
    fd1.fec = cache.cache_compile(fd1.fusion)
    del fd1._fusion
    # Check that LRUCache creates a separate entry for second fusion
    assert cache.num_fusions() == 2

    assert torch.allclose(
        fd0.execute([bcast_input, full_input])[0], bcast_input - full_input
    )
    assert torch.allclose(
        fd1.execute([full_input, bcast_input])[0], bcast_input - full_input
    )


    try:
        with PythonProfiler(auto_scheduled=False) as prof:
            fd.manual_execute(inputs)
        assert prof.profile.fusion_id >= 0
        assert prof.profile.kernel_profiles[0].scheduler == "user"
    except Exception as e:
        raise RuntimeError(
            "FusionDefinition did not run correctly with profile enabled! Error: "
            + str(e)
        )

def test_fusion_profiler_with_noncodegen_kernels():
    inputs = [
        torch.randn((2, 4, 16), dtype=torch.bfloat16, device="cuda:0"),
        torch.randn((2, 4, 16), dtype=torch.bfloat16, device="cuda:0"),
        torch.randn((16, 16), dtype=torch.bfloat16, device="cuda:0"),
    ]

    def fusion_func(fd: FusionDefinition) -> None:
        T0 = fd.from_pytorch(inputs[0])
        T1 = fd.from_pytorch(inputs[1])
        T2 = fd.from_pytorch(inputs[2])
        T3 = fd.ops.linear(T0, T2)
        T4 = fd.ops.add(T3, T1)
        fd.add_output(T4)

    with FusionDefinition() as fd:
        fusion_func(fd)

    try:
        with PythonProfiler() as prof:
            fd.execute(inputs)
        assert len(prof.profile.kernel_profiles) == 2
        assert len(prof.profile.kernel_profiles[0].name) > 0
        assert len(prof.profile.kernel_profiles[1].name) > 0
    except Exception as e:
        raise RuntimeError(
            "FusionDefinition did not run correctly with profile enabled! Error: "
            + str(e)
        )

def test_lru_fusion_cache_decorator():
    inputs = [
        torch.randn(2, 4, device="cuda"),


    try:
        with PythonProfiler() as prof:
            fd.execute(inputs)
        assert len(prof.profile.kernel_profiles) == 2
        assert len(prof.profile.kernel_profiles[0].name) > 0
        assert len(prof.profile.kernel_profiles[1].name) > 0
    except Exception as e:
        raise RuntimeError(
            "FusionDefinition did not run correctly with profile enabled! Error: "
            + str(e)
        )

def test_lru_fusion_cache_decorator():
    inputs = [
        torch.randn(2, 4, device="cuda"),
        torch.randn(2, 4, device="cuda"),
    ]

    def create_fd1(fd: FusionDefinition):
        t0 = fd.from_pytorch(inputs[0])
        t1 = fd.from_pytorch(inputs[1])
        t2 = fd.ops.div(t0, t1)
        fd.add_output(t2)

    def create_fd2(fd: FusionDefinition):
        t0 = fd.from_pytorch(inputs[0])
        t1 = fd.from_pytorch(inputs[1])
        t2 = fd.ops.div(t1, t0)
        fd.add_output(t2)

    @LruFusionCache(max_fusions=10)
    def create_fusion(create_fd: Callable):
        with FusionDefinition() as fd:
            create_fd(fd)
        return fd

    empty_stats = """Max Fusions Allowed: 10
The fusion cache is empty.\n"""
    assert create_fusion.stats() == empty_stats

    # Test LRU cache compilation
    for i in range(5):
        fd1 = create_fusion(create_fd1)