# mypy: ignore-errors

import functools
import unittest

import torch
from functorch.experimental.control_flow import map
from torch._higher_order_ops.flex_attention import (
    flex_attention as flex_attention_hop,
)
from torch.nn.attention.flex_attention import (
    _create_empty_block_mask,
    create_block_mask,
    flex_attention,
)
from torch.testing import make_tensor
from torch._higher_order_ops.inline_asm_elementwise import inline_asm_elementwise
from torch.testing._internal.common_device_type import onlyCUDA
from torch.testing._internal.common_dtype import all_types_and, custom_types
from torch.testing._internal.opinfo.core import DecorateInfo, OpInfo, SampleInput
from torch._higher_order_ops.invoke_subgraph import mark_compile_region
from torch._higher_order_ops import InvokeQuant, invoke_quant_packed


def sample_inputs_map(opinfo, device, dtype, requires_grad, **kwargs):
    make_arg = functools.partial(
        make_tensor, device=device, dtype=dtype, requires_grad=requires_grad
    )
    yield SampleInput(
        [make_arg(2, 2, 2, low=0.1, high=2), make_arg(2, 2, 2, low=0.1, high=2)],
        args=(make_arg(1, low=0.1, high=2), make_arg(1, low=0.1, high=2)),
    )


def inner_f(x, y0, y1):
    return [x[0].cos().add_(1.0) * y0, (x[1] + y1.sin()).cos_().view(x[1].size())]


def simple_map(xs, y0, y1):
    def f(x, y0, y1):
        return inner_f(x, y0, y1)

    return map(f, xs, y0, y1)


def nested_map(xs, y0, y1):
    def f1(xx, y0, y1):
        def f2(x, y0, y1):
            return inner_f(x, y0, y1)

        return map(f2, xx, y0, y1)

    return map(f1, xs, y0, y1)


def triple_nested_map(xs, y0, y1):
    def f0(xs, y0, y1):
        def f1(xx, y0, y1):
            def f2(x, y0, y1):
                return inner_f(x, y0, y1)

            return map(f2, xx, y0, y1)

        return map(f1, xs, y0, y1)

    return map(f0, xs, y0, y1)


# PLEASE DON'T ADD ANYTHING NEW TO THIS LIST,
# and do add an OpInfo for your HOP.
# The OpInfo lets us do automated testing for the HOP to check that
# your HOP will work correctly with PyTorch!
#
# Your new HOP may fail some automated testing. That's OK. If you don't
# care about certain features (like torch.export), it's fine to xfail those
# failing tests. It is less fine to xfail a more critical check (like checking
# if torch.compile works with your HOP, or if your HOP has a docstring).
# If you don't know if a test is fine to xfail, please ask.
#
# There are legitimate reasons why something cannot be added to this list
# (e.g. it uses executorch which is not in PyTorch). If that's the case then
# please leave a comment.
FIXME_hop_that_doesnt_have_opinfo_test_allowlist = [
    "custom_function_call",
    "autograd_function_apply",
    "run_and_save_rng_state",
    "run_with_rng_state",
    "run_dtensor_rng_op",
    "graphsafe_run_with_rng_state",
    "out_dtype",
    "trace_wrapped",
    'tag_activation_checkpoint',
    'executorch_call_delegate',
    'wrap',
    'wrap_with_set_grad_enabled',
    'auto_functionalized_v2',
    'associative_scan',
    'flat_apply',  # is WIP, doesn't pass any of the tests yet
    'wrap_with_autocast',
    'wrap_activation_checkpoint',
    'run_const_graph',
    'auto_functionalized',
    "map",  # T183144629
    "map_impl",
    "with_effects",
    "strict_mode",
    "_export_tracepoint",
    "call_torchbind",
    "triton_kernel_wrapper_mutation",
    "triton_kernel_wrapper_functional",
    "hints_wrapper",
    "dynamo_bypassing_wrapper",  # TODO(soulitzer)
    "foreach_map",
    "aoti_call_delegate",
    "print",
    "inductor_compiled_code",  # Tested separately in test_inductor_wrap_inductor_compile_regions
    "invoke_leaf_function",  # Needs torch.compile, tested separately in test_leaf_function*
]

torch.library.define(
    "testlib::mutating_custom_op",
    "(Tensor(a!) x, Tensor(b!) z) -> (Tensor, Tensor, Tensor)",
    tags=torch.Tag.pt2_compliant_tag,
)


@torch.library.impl("testlib::mutating_custom_op", "cpu")
def foo_impl_cpu(x, z):
    x.add_(5)
    z.add_(5)
    return x.clone(), z.clone(), x + z


@torch.library.impl("testlib::mutating_custom_op", "cuda")
def foo_impl_cuda(x, z):
    x.add_(5)
    z.add_(5)
    return x.clone(), z.clone(), x + z


@torch.library.impl("testlib::mutating_custom_op", "xpu")
def foo_impl_xpu(x, z):
    x.add_(5)
    z.add_(5)
    return x.clone(), z.clone(), x + z


@torch.library.register_fake("testlib::mutating_custom_op")
def foo_impl_abstract(x, z):
    return x.clone(), z.clone(), x + z


def sample_inputs_cond(opinfo, device, dtype, requires_grad, **kwargs):
    make_arg = functools.partial(
        make_tensor, device=device, dtype=dtype, requires_grad=requires_grad
    )
    yield SampleInput(make_arg(2, 2, 2, low=0.1, high=2))


def simple_cond(x):
    return torch.cond(x.sum() > 2, lambda x: (x.cos(),), lambda x: (x.sin(),), [x])


def sample_inputs_invoke_subgraph(opinfo, device, dtype, requires_grad, **kwargs):
    make_arg = functools.partial(
        make_tensor, device=device, dtype=dtype, requires_grad=requires_grad
    )
    yield SampleInput(make_arg(2, 2, 2, low=0.1, high=2))


@mark_compile_region
def fn_for_invoke_subgraph(x):
    return torch.sin(x)


def simple_invoke_subgraph(x):
    return fn_for_invoke_subgraph(x)


def sample_inputs_auto_functionalize(opinfo, device, dtype, requires_grad, **kwargs):
    make_arg = functools.partial(
        make_tensor, device=device, dtype=dtype, requires_grad=False
    )
    yield SampleInput(
        make_arg(2, 2, 2, low=0.1, high=2), make_arg(2, 2, 2, low=0.1, high=2)
    )


def simple_auto_functionalize(x, z):
    return torch.ops.testlib.mutating_custom_op(x, z)


def sample_inputs_flex_attention(opinfo, device, dtype, requires_grad, **kwargs):
    make_arg = functools.partial(
        make_tensor, device=device, dtype=dtype, requires_grad=requires_grad
    )

    def score_mod(score, b, h, m, n):
        return score + h

    q, k, v = (make_arg(2, 2, 128, 8, low=0.1, high=2) for _ in range(3))
    block_mask = _create_empty_block_mask(q, k)
    yield SampleInput(q, k, v, score_mod, block_mask)


def sample_inputs_flex_attention_backward(
    opinfo, device, dtype, requires_grad, **kwargs
):
    make_arg = functools.partial(
        make_tensor, device=device, dtype=dtype, requires_grad=False
    )

    def score_mod(score, b, h, m, n):
        return score

    def mask_mod(b, h, m, n):
        return m >= n

    q, k, v = (make_arg(2, 2, 128, 16, low=0.1, high=2) for _ in range(3))
    block_mask = create_block_mask(mask_mod, B=2, H=2, Q_LEN=128, KV_LEN=128, device=device)
    scale = 1.0 / q.size(-1) ** 0.5
    out, logsumexp, _ = flex_attention_hop(
        q, k, v, score_mod, block_mask.as_tuple(), scale, {},
    )
    yield SampleInput(
        q,
        args=(
            k, v, out.detach(), logsumexp.detach(), torch.rand_like(out), None,
            score_mod, None, block_mask.as_tuple(),
            scale, {}, (), (),
        ),
    )


def sample_inputs_flex_attention_backward_explicit_buffers(
    opinfo, device, dtype, requires_grad, **kwargs
):
    make_arg = functools.partial(
        make_tensor, device=device, dtype=dtype, requires_grad=False
    )
    mask_offset = torch.full((), 128, device=device, dtype=torch.int32)

    def score_mod(score, b, h, m, n):
        return score

    def mask_mod(b, h, m, n):
        return m + mask_offset >= n

    q, k, v = (make_arg(2, 2, 128, 16, low=0.1, high=2) for _ in range(3))
    block_mask = create_block_mask(mask_mod, B=2, H=2, Q_LEN=128, KV_LEN=128, device=device)
    scale = 1.0 / q.size(-1) ** 0.5
    out, logsumexp, _ = flex_attention_hop(
        q, k, v, score_mod, block_mask.as_tuple(), scale, {},
    )
    yield SampleInput(
        q,
        args=(
            k, v, out.detach(), logsumexp.detach(), torch.rand_like(out), None,
            score_mod, None, block_mask.as_tuple(),
            scale, {}, (), (),
        ),
    )


def simple_flex_attention_backward(
    query,
    key,
    value,
    out,
    logsumexp,
    grad_out,
    grad_logsumexp,
    fw_graph,
    joint_graph,
    block_mask,
    scale,
    kernel_options,
    score_mod_other_buffers,
    mask_mod_other_buffers,
):
    return torch.ops.higher_order.flex_attention_backward(
        query,
        key,
        value,
        out,
        logsumexp,
        grad_out,
        grad_logsumexp,
        fw_graph,
        joint_graph,
        block_mask,
        scale,
        kernel_options,
        score_mod_other_buffers,
        mask_mod_other_buffers,
    )


def sample_inputs_while_loop(opinfo, device, dtype, requires_grad, **kwargs):
    make_arg = functools.partial(
        make_tensor, device=device, dtype=dtype, requires_grad=False
    )
    yield SampleInput(
        torch.tensor(3),
        make_arg(2, 3, 4, low=0.1, high=2),
    )


def simple_while_loop(iter_t, x):
    def cond_fn(iter_t, x):
        return iter_t > 0

    def body_fn(iter_t, x):
        return iter_t - 1, x.cos()

    return torch._higher_order_ops.while_loop(cond_fn, body_fn, (iter_t, x))


def simple_while_loop_stack_output(iter_t, x):
    def cond_fn(iter_t, x):
        return iter_t > 0

    def body_fn(iter_t, x):
        return iter_t - 1, x.cos()

    return torch._higher_order_ops.while_loop_stack_output(
        cond_fn, body_fn, (iter_t, x), tuple()
    )


def sample_inputs_local_map_hop(opinfo, device, dtype, requires_grad, **kwargs):
    # TODO: once HOPs support DTensor inputs, we should also test DTensors
    make_arg = functools.partial(
        make_tensor, device=device, dtype=dtype, requires_grad=False
    )
    yield SampleInput(
        make_arg(2, 3, 4, low=0.1, high=2),
        make_arg(2, 3, 4, low=0.1, high=2),
    )


def simple_local_map_hop(inp1, inp2):
    def body_gm(inp1, inp2):
        return inp1.cos() + inp2.sin()

    gm = torch.fx.symbolic_trace(body_gm)

    if not torch.distributed.is_available():
        raise AssertionError("Expected torch.distributed to be available")
    from torch.distributed.tensor.placement_types import Replicate

    gm.meta["local_map_kwargs"] = {
        "in_placements": (Replicate(), Replicate(), Replicate()),
        "out_placements": ((Replicate(), Replicate(), Replicate()),),
    }

    # TODO: Dynamo would rewrite this op differently
    return torch._higher_order_ops.local_map_hop(gm, inp1, inp2)


def sample_inputs_scan(opinfo, device, dtype, requires_grad, **kwargs):
    make_arg = functools.partial(
        make_tensor, device=device, dtype=dtype, requires_grad=requires_grad
    )
    yield SampleInput(
        make_arg(2, 2, low=0.1, high=2),
        make_arg(2, 2, 2, low=0.1, high=2),
    )


def simple_scan(init, xs):
    def combine_fn(carry, x):
        result = carry @ x + x
        return result, carry.clone()

    return torch._higher_order_ops.scan(combine_fn, init, xs)


quant_tracer = InvokeQuant()


def simple_invoke_quant(x):
    def fn(x, y):
        return (torch.sin(x) * y,)

    return quant_tracer(fn, x, x)[0] * 2.0


def simple_invoke_quant_packed(x):
    def fn(x):
        return (torch.sin(x),)

    return invoke_quant_packed(fn, x)[0] * 2.0


def sample_inputs_inline_asm(opinfo, device, dtype, requires_grad, **kwargs):
    make_arg = functools.partial(
        make_tensor, device=device, dtype=dtype, requires_grad=requires_grad
    )
    yield SampleInput(make_arg(2, 2, 2, low=0.1, high=2))


def simple_inline_asm(x):
    if torch.version.hip:
        return inline_asm_elementwise(
            x,
            asm_str="v_mov_b32_e32 $0, $1",
            constraints="=v, v",
            dtype=torch.float32,
        )

    return inline_asm_elementwise(
        x, asm_str="mov.f32 $0, $1;", constraints="=f,f", dtype=torch.float32
    )


hop_db = [
    OpInfo(
        name="scan",
        variant_test_name="simple",
        op=simple_scan,
        sample_inputs_func=sample_inputs_scan,
        dtypes=all_types_and(torch.bool, torch.half),
        supports_out=False,
        check_batched_grad=False,
        check_batched_gradgrad=False,
        check_batched_forward_grad=False,
        check_inplace_batched_forward_grad=False,
        supports_autograd=False,
        # "torch.compile with aot_autograd does not currently support double backward."
        supports_gradgrad=False,
    ),
    OpInfo(
        name="invoke_subgraph",
        variant_test_name="simple",
        op=simple_invoke_subgraph,
        sample_inputs_func=sample_inputs_invoke_subgraph,
        dtypes=all_types_and(torch.bool, torch.half),
        supports_out=False,
        check_batched_grad=False,
        check_batched_gradgrad=False,
        check_batched_forward_grad=False,
        check_inplace_batched_forward_grad=False,
        supports_autograd=True,
        # "torch.compile with aot_autograd does not currently support double backward."
        supports_gradgrad=False,
    ),
    OpInfo(
        name="map",
        variant_test_name="simple",
        op=simple_map,
        sample_inputs_func=sample_inputs_map,
        dtypes=all_types_and(torch.bool, torch.half),
        supports_out=False,
        check_batched_grad=False,
        check_batched_gradgrad=False,
        check_batched_forward_grad=False,
        check_inplace_batched_forward_grad=False,
    ),
    OpInfo(
        name="map",
        variant_test_name="nested",
        op=nested_map,
        sample_inputs_func=sample_inputs_map,
        dtypes=all_types_and(torch.bool, torch.half),
        supports_out=False,
        check_batched_grad=False,
        check_batched_gradgrad=False,
        check_batched_forward_grad=False,
        check_inplace_batched_forward_grad=False,
    ),
    OpInfo(
        name="map",
        variant_test_name="triple_nested",
        op=triple_nested_map,
        sample_inputs_func=sample_inputs_map,
        dtypes=all_types_and(torch.bool, torch.half),
        supports_out=False,
        check_batched_grad=False,
        check_batched_gradgrad=False,
        check_batched_forward_grad=False,
        check_inplace_batched_forward_grad=False,
    ),
    OpInfo(
        name="cond",
        variant_test_name="simple",
        op=simple_cond,
        sample_inputs_func=sample_inputs_cond,
        dtypes=all_types_and(torch.bool, torch.half),
        supports_out=False,
        check_batched_grad=False,
        check_batched_gradgrad=False,
        check_batched_forward_grad=False,
        check_inplace_batched_forward_grad=False,
        supports_autograd=True,
        # "torch.compile with aot_autograd does not currently support double backward."
        supports_gradgrad=False,
    ),
    OpInfo(
        name="invoke_quant",
        variant_test_name="simple",
        op=simple_invoke_quant,
        sample_inputs_func=sample_inputs_invoke_subgraph,
        dtypes=all_types_and(torch.bool, torch.half),
        supports_out=False,
        check_batched_grad=False,
        check_batched_gradgrad=False,
        check_batched_forward_grad=False,
        check_inplace_batched_forward_grad=False,
        supports_autograd=True,
        # "torch.compile with aot_autograd does not currently support double backward."
        skips=(
            DecorateInfo(unittest.expectedFailure, "TestHOP", "test_aot_export"),
            DecorateInfo(
                unittest.expectedFailure, "TestHOP", "test_pre_dispatch_export"
            ),
            DecorateInfo(unittest.expectedFailure, "TestHOP", "test_serialize_export"),
            DecorateInfo(unittest.expectedFailure, "TestHOP", "test_retrace_export"),
        ),
        # "torch.compile with aot_autograd does not currently support double backward."
        supports_gradgrad=False,
    ),
    OpInfo(
        name="invoke_quant_packed",
        variant_test_name="simple",
        op=simple_invoke_quant_packed,
        sample_inputs_func=sample_inputs_invoke_subgraph,
        dtypes=all_types_and(torch.bool, torch.half),
        supports_out=False,
        check_batched_grad=False,
        check_batched_gradgrad=False,
        check_batched_forward_grad=False,
        check_inplace_batched_forward_grad=False,
        supports_autograd=True,
        # "torch.compile with aot_autograd does not currently support double backward."
        supports_gradgrad=False,
    ),
    OpInfo(
        name="while_loop",
        variant_test_name="simple",
        op=simple_while_loop,
        sample_inputs_func=sample_inputs_while_loop,
        dtypes=all_types_and(torch.bool, torch.half),
        supports_out=False,
        check_batched_grad=False,
        check_batched_gradgrad=False,
        check_batched_forward_grad=False,
        check_inplace_batched_forward_grad=False,
        supports_autograd=False,
    ),
    OpInfo(
        name="while_loop_stack_output",
        variant_test_name="simple",
        op=simple_while_loop_stack_output,
        sample_inputs_func=sample_inputs_while_loop,
        dtypes=all_types_and(torch.bool, torch.half),
        supports_out=False,
        check_batched_grad=False,
        check_batched_gradgrad=False,
        check_batched_forward_grad=False,
        check_inplace_batched_forward_grad=False,
        supports_autograd=False,
    ),
    OpInfo(
        name="auto_functionalize",
        variant_test_name="simple",
        op=simple_auto_functionalize,
        sample_inputs_func=sample_inputs_auto_functionalize,
        dtypes=all_types_and(torch.bool, torch.half),
        supports_out=False,
        check_batched_grad=False,
        check_batched_gradgrad=False,
        check_batched_forward_grad=False,
        check_inplace_batched_forward_grad=False,
        supports_autograd=False,
    ),
    OpInfo(
        name="flex_attention",
        variant_test_name="simple",
        op=flex_attention,
        sample_inputs_func=sample_inputs_flex_attention,
        dtypes=custom_types(torch.float16, torch.float32),
        supports_out=False,
        check_batched_grad=False,
        check_batched_gradgrad=False,
        check_batched_forward_grad=False,
        check_inplace_batched_forward_grad=False,
        skips=(
            DecorateInfo(unittest.expectedFailure, "TestHOP", "test_aot_export"),
            DecorateInfo(
                unittest.expectedFailure, "TestHOP", "test_pre_dispatch_export"
            ),
            DecorateInfo(unittest.expectedFailure, "TestHOP", "test_serialize_export"),
            DecorateInfo(unittest.expectedFailure, "TestHOP", "test_retrace_export"),
        ),
        decorators=[onlyCUDA],
    ),
    OpInfo(
        name="flex_attention_backward",
        variant_test_name="simple",
        op=simple_flex_attention_backward,
        sample_inputs_func=sample_inputs_flex_attention_backward,
        dtypes=custom_types(torch.float16, torch.float32),
        supports_out=False,
        check_batched_grad=False,
        check_batched_gradgrad=False,
        check_batched_forward_grad=False,
        check_inplace_batched_forward_grad=False,
        supports_autograd=False,
        supports_gradgrad=False,
        skips=(
            DecorateInfo(unittest.expectedFailure, "TestHOP", "test_aot_export"),
            DecorateInfo(
                unittest.expectedFailure, "TestHOP", "test_pre_dispatch_export"
            ),
            DecorateInfo(unittest.expectedFailure, "TestHOP", "test_serialize_export"),
            DecorateInfo(unittest.expectedFailure, "TestHOP", "test_retrace_export"),
        ),
        decorators=[onlyCUDA],
    ),
    OpInfo(
        name="flex_attention_backward",
        variant_test_name="explicit_buffers",
        op=simple_flex_attention_backward,
        sample_inputs_func=sample_inputs_flex_attention_backward_explicit_buffers,
        dtypes=custom_types(torch.float16, torch.float32),
        supports_out=False,
        check_batched_grad=False,
        check_batched_gradgrad=False,
        check_batched_forward_grad=False,
        check_inplace_batched_forward_grad=False,
        supports_autograd=False,
        supports_gradgrad=False,
        skips=(
            DecorateInfo(unittest.expectedFailure, "TestHOP", "test_aot_export"),
            DecorateInfo(
                unittest.expectedFailure, "TestHOP", "test_pre_dispatch_export"
            ),
            DecorateInfo(unittest.expectedFailure, "TestHOP", "test_serialize_export"),
            DecorateInfo(unittest.expectedFailure, "TestHOP", "test_retrace_export"),
        ),
        decorators=[onlyCUDA],
    ),
    OpInfo(
        name="local_map_hop",
        variant_test_name="simple",
        op=simple_local_map_hop,
        sample_inputs_func=sample_inputs_local_map_hop,
        dtypes=custom_types(torch.float16, torch.float32),
        supports_out=False,
        check_batched_grad=False,
        check_batched_gradgrad=False,
        check_batched_forward_grad=False,
        check_inplace_batched_forward_grad=False,
        skips=(
            DecorateInfo(unittest.expectedFailure, "TestHOP", "test_aot_export"),
            DecorateInfo(
                unittest.expectedFailure, "TestHOP", "test_pre_dispatch_export"
            ),
            DecorateInfo(unittest.expectedFailure, "TestHOP", "test_serialize_export"),
            DecorateInfo(unittest.expectedFailure, "TestHOP", "test_retrace_export"),
        ),
        decorators=[
            onlyCUDA,
            unittest.skipIf(
                not torch.distributed.is_available(), "requires distributed build"
            ),
        ],
    ),
    OpInfo(
        name="inline_asm_elementwise",
        variant_test_name="simple",
        op=simple_inline_asm,
        sample_inputs_func=sample_inputs_inline_asm,
        dtypes=custom_types(torch.float32),
        supports_out=False,
        check_batched_grad=False,
        check_batched_gradgrad=False,
        check_batched_forward_grad=False,
        check_inplace_batched_forward_grad=False,
        supports_autograd=False,
        decorators=[onlyCUDA],
    ),
]
