# SPDX-License-Identifier: MIT
 

import torch
import aiter
from aiter.test_common import checkAllclose, perftest
import itertools
from enum import IntEnum
import argparse
from aiter import dtypes


@perftest()
def hip_rope_fwd(
    input, freqs, rotate_style, reuse_freqs_front_part, nope_first, transpose_output
):
    return aiter.rope_fwd(
        input, freqs, rotate_style, reuse_freqs_front_part, nope_first, transpose_output
    )


@perftest()
def hip_rope_bwd(
    output_grads,
    freqs,
    rotate_style,
    reuse_freqs_front_part,
    nope_first,
    transpose_output,
):
    return aiter.rope_bwd(
        output_grads,
        freqs,
        rotate_style,
        reuse_freqs_front_part,
        nope_first,
        transpose_output,
    )


@perftest()
def hip_rope_2c_fwd(
    input_x,
    input_y,
    freqs,
    rotate_style,
    reuse_freqs_front_part,
    nope_first,
    transpose_output,
):
    return aiter.rope_2c_fwd(
        input_x,
        input_y,
        freqs,
        rotate_style,
        reuse_freqs_front_part,
        nope_first,
        transpose_output,
    )


@perftest()
def hip_rope_2c_bwd(
    output_grads_x,
    output_grads_y,
    freqs,
    rotate_style,
    reuse_freqs_front_part,
    nope_first,
    transpose_output,
):
    return aiter.rope_2c_bwd(
        output_grads_x,
        output_grads_y,
        freqs,
        rotate_style,
        reuse_freqs_front_part,
        nope_first,
        transpose_output,
    )


@perftest()
def hip_rope_cached_fwd(
    input, cos, sin, rotate_style, reuse_freqs_front_part, nope_first, transpose_output
):
    return aiter.rope_cached_fwd(
        input,
        cos,
        sin,
        rotate_style,
        reuse_freqs_front_part,
        nope_first,
        transpose_output,
    )


@perftest()
def hip_rope_cached_bwd(
    output_grads,
    cos,
    sin,
    rotate_style,
    reuse_freqs_front_part,
    nope_first,
    transpose_output,
):
    return aiter.rope_cached_bwd(
        output_grads,
        cos,
        sin,
        rotate_style,
        reuse_freqs_front_part,
        nope_first,
        transpose_output,
    )


@perftest()
def hip_rope_cached_2c_fwd(
    input_x,
    input_y,
    cos,
    sin,
    rotate_style,
    reuse_freqs_front_part,
    nope_first,
    transpose_output,
):
    return aiter.rope_cached_2c_fwd(
        input_x,
        input_y,
        cos,
        sin,
        rotate_style,
        reuse_freqs_front_part,
        nope_first,
        transpose_output,
    )


@perftest()
def hip_rope_cached_2c_bwd(
    output_grads_x,
    output_grads_y,
    cos,
    sin,
    rotate_style,
    reuse_freqs_front_part,
    nope_first,
    transpose_output,
):
    return aiter.rope_cached_2c_bwd(
        output_grads_x,
        output_grads_y,
        cos,
        sin,
        rotate_style,
        reuse_freqs_front_part,
        nope_first,
        transpose_output,
    )


@perftest()
def hip_rope_cached_positions_fwd(
    input,
    cos,
    sin,
    positions,
    rotate_style,
    reuse_freqs_front_part,
    nope_first,
    transpose_output,
):
    return aiter.rope_cached_positions_fwd(
        input,
        cos,
        sin,
        positions,
        rotate_style,
        reuse_freqs_front_part,
        nope_first,
        transpose_output,
    )


@perftest()
def hip_rope_cached_positions_offsets_fwd(
    input,
    cos,
    sin,
    positions,
    offsets,
    rotate_style,
    reuse_freqs_front_part,
    nope_first,
    transpose_output,
):
    return aiter.rope_cached_positions_offsets_fwd(
        input,
        cos,
        sin,
        positions,
        offsets,
        rotate_style,
        reuse_freqs_front_part,
        nope_first,
        transpose_output,
    )


@perftest()
def hip_rope_cached_positions_fwd_inplace(
    input, cos, sin, positions, rotate_style, reuse_freqs_front_part, nope_first
):
    return aiter.rope_cached_positions_fwd_inplace(
        input, cos, sin, positions, rotate_style, reuse_freqs_front_part, nope_first
    )


@perftest()
def hip_rope_cached_positions_offsets_fwd_inplace(
    input,
    cos,
    sin,
    positions,
    offsets,
    rotate_style,
    reuse_freqs_front_part,
    nope_first,
):
    return aiter.rope_cached_positions_offsets_fwd_inplace(
        input,
        cos,
        sin,
        positions,
        offsets,
        rotate_style,
        reuse_freqs_front_part,
        nope_first,
    )


@perftest()
def hip_rope_cached_positions_2d_fwd(
    input_x,
    input_y,
    cos,
    sin,
    positions,
    rotate_style,
    reuse_freqs_front_part,
    nope_first,
    transpose_output,
):
    return aiter.rope_cached_positions_2c_fwd(
        input_x,
        input_y,
        cos,
        sin,
        positions,
        rotate_style,
        reuse_freqs_front_part,
        nope_first,
        transpose_output,
    )


@perftest()
def hip_rope_cached_positions_offsets_2d_fwd(
    input_x,
    input_y,
    cos,
    sin,
    positions,
    offsets,
    rotate_style,
    reuse_freqs_front_part,
    nope_first,
    transpose_output,
):
    return aiter.rope_cached_positions_offsets_2c_fwd(
        input_x,
        input_y,
        cos,
        sin,
        positions,
        offsets,
        rotate_style,
        reuse_freqs_front_part,
        nope_first,
        transpose_output,
    )


@perftest()
def hip_rope_cached_positions_2d_fwd_inplace(
    input_x,
    input_y,
    cos,
    sin,
    positions,
    rotate_style,
    reuse_freqs_front_part,
    nope_first,
):
    return aiter.rope_cached_positions_2c_fwd_inplace(
        input_x,
        input_y,
        cos,
        sin,
        positions,
        rotate_style,
        reuse_freqs_front_part,
        nope_first,
    )


@perftest()
def hip_rope_cached_positions_offsets_2d_fwd_inplace(
    input_x,
    input_y,
    cos,
    sin,
    positions,
    offsets,
    rotate_style,
    reuse_freqs_front_part,
    nope_first,
):
    return aiter.rope_cached_positions_offsets_2c_fwd_inplace(
        input_x,
        input_y,
        cos,
        sin,
        positions,
        offsets,
        rotate_style,
        reuse_freqs_front_part,
        nope_first,
    )


@perftest()
def hip_rope_thd_fwd(
    input, cu_seqlens, freqs, rotate_style, reuse_freqs_front_part, nope_first
):
    return aiter.rope_thd_fwd(
        input, cu_seqlens, freqs, rotate_style, reuse_freqs_front_part, nope_first
    )


@perftest()
def hip_rope_thd_bwd(
    output_grads, cu_seqlens, freqs, rotate_style, reuse_freqs_front_part, nope_first
):
    return aiter.rope_thd_bwd(
        output_grads,
        cu_seqlens,
        freqs,
        rotate_style,
        reuse_freqs_front_part,
        nope_first,
    )


@perftest()
def hip_rope_2d_fwd(
    input,
    height,
    width,
    cos_h,
    sin_h,
    cos_w,
    sin_w,
    rotate_style,
    reuse_freqs_front_part,
    nope_first,
):
    return aiter.rope_2d_fwd(
        input,
        height,
        width,
        cos_h,
        sin_h,
        cos_w,
        sin_w,
        rotate_style,
        reuse_freqs_front_part,
        nope_first,
    )


@perftest()
def hip_rope_2d_bwd(
    output_grads,
    height,
    width,
    cos_h,
    sin_h,
    cos_w,
    sin_w,
    rotate_style,
    reuse_freqs_front_part,
    nope_first,
):
    return aiter.rope_2d_bwd(
        output_grads,
        height,
        width,
        cos_h,
        sin_h,
        cos_w,
        sin_w,
        rotate_style,
        reuse_freqs_front_part,
        nope_first,
    )


@perftest()
def legacy_rope_cached_positions_2d_fwd(
    input_x, input_y, cos, sin, positions, rotate_style, nope_first
):
    aiter.rotary_embedding_fwd(
        positions,
        input_x,
        input_y,
        d,
        cos,
        sin,
        rotate_style is RotateStyle.NEOX,
        nope_first,
    )
    return input_x, input_y


@perftest()
def legacy_rope_cached_positions_offsets_2d_fwd(
    input_x, input_y, cos, sin, positions, offsets, rotate_style, nope_first
):
    rotate_dim = sin.size(-1) * 2
    aiter.batched_rotary_embedding(
        positions,
        input_x,
        input_y,
        d,
        cos,
        sin,
        rotate_style is RotateStyle.NEOX,
        nope_first,
        rotate_dim,
        offsets.view(-1),
    )
    return input_x, input_y


class RotateStyle(IntEnum):
    NEOX = (0,)
    GPTJ = 1


def rotate_half_neox(x):
    x1 = x[..., : x.shape[-1] // 2]
    x2 = x[..., x.shape[-1] // 2 :]
    return torch.cat((-x2, x1), dim=-1)


def rotate_half_gptj(x):
    x1 = x[..., ::2]
    x2 = x[..., 1::2]
    x = torch.stack((-x2, x1), dim=-1)
    return x.flatten(-2)


def ref_rope_sbhd_fwd(x, freqs, rotate_style, reuse_freqs_front_part, nope_first):
    rotate_half = (
        rotate_half_neox if rotate_style == RotateStyle.NEOX else rotate_half_gptj
    )
    rotate_dim = freqs.shape[-1] * (2 if reuse_freqs_front_part else 1)
    if nope_first:
        d = x.shape[-1]
        x, x_forward = x[..., d - rotate_dim :], x[..., : d - rotate_dim]
    else:
        x, x_forward = x[..., :rotate_dim], x[..., rotate_dim:]
    if reuse_freqs_front_part:
        if rotate_style == RotateStyle.NEOX:
            freqs = freqs.repeat([1] * (freqs.dim() - 1) + [2])
        elif rotate_style == RotateStyle.GPTJ:
            freqs = freqs.repeat_interleave(2, dim=-1)
    x_embed = (x * torch.cos(freqs)) + (rotate_half(x) * torch.sin(freqs))
    return (
        torch.cat((x_forward, x_embed.to(dtype=x.dtype)), dim=-1)
        if nope_first
        else torch.cat((x_embed.to(dtype=x.dtype), x_forward), dim=-1)
    )


def ref_rope_thd_fwd(
    x, cu_seqlens, freqs, rotate_style, reuse_freqs_front_part, nope_first
):
    seqlens = (cu_seqlens[1:] - cu_seqlens[:-1]).tolist()
    x_embed = torch.cat(
        [
            ref_rope_sbhd_fwd(
                xi.unsqueeze(1),
                freqs[: xi.size(0)],
                rotate_style,
                reuse_freqs_front_part,
                nope_first,
            )
            for xi in torch.split(x, seqlens)
        ]
    )
    return x_embed.squeeze(1)


def ref_rope_2d_fwd(x, size_h, size_w, cos_h, sin_h, cos_w, sin_w, rotate_style):
    rotate_half = (
        rotate_half_neox if rotate_style == RotateStyle.NEOX else rotate_half_gptj
    )
    s, b, h, d = x.shape
    x = x.view(s, size_h, size_w, h, d)
    x1, x2 = x.chunk(2, dim=-1)
    cos_h = cos_h[:, :size_h].unsqueeze(2)  # [1, H, 1, 1, D//2]
    sin_h = sin_h[:, :size_h].unsqueeze(2)  # [1, H, 1, 1, D//2]
    x1 = (x1 * cos_h) + (rotate_half(x1) * sin_h)
    cos_w = cos_w[:, :size_w].unsqueeze(1)  # [1, 1, W, 1, D//2]
    sin_w = sin_w[:, :size_w].unsqueeze(1)  # [1, 1, W, 1, D//2]
    x2 = (x2 * cos_w) + (rotate_half(x2) * sin_w)
    return torch.cat([x1, x2], dim=-1).view(s, b, h, d).to(dtype=x.dtype)


def test_rope_sbhd(
    input,
    freqs,
    grad,
    rotate_style,
    reuse_freqs_front_part,
    nope_first,
    transpose_output,
):
    input_msg = f"""
dtype: {input.dtype}, \
freq_dtype: {freqs.dtype}, \
dim_input: {str(input.shape):<20}, \
dim_freqs: {str(freqs.shape):<20}, \
rotate_style: {rotate_style.value}, \
reuse_freqs_front_part: {reuse_freqs_front_part}, \
nope_first: {nope_first}, \
transpose_output: {transpose_output}
"""

    ref = ref_rope_sbhd_fwd(
        input, freqs, rotate_style, reuse_freqs_front_part, nope_first
    )
    ref.backward(grad)

    cos = torch.cos(freqs)
    sin = torch.sin(freqs)

    hip_fwd, hip_fwd_avg = hip_rope_fwd(
        input, freqs, rotate_style, reuse_freqs_front_part, nope_first, transpose_output
    )
    hip_bwd, hip_bwd_avg = hip_rope_bwd(
        grad, freqs, rotate_style, reuse_freqs_front_part, nope_first, transpose_output
    )
    hip_cached_fwd, hip_cached_fwd_avg = hip_rope_cached_fwd(
        input,
        cos,
        sin,
        rotate_style,
        reuse_freqs_front_part,
        nope_first,
        transpose_output,
    )
    hip_cached_bwd, hip_cached_bwd_avg = hip_rope_cached_bwd(
        grad,
        cos,
        sin,
        rotate_style,
        reuse_freqs_front_part,
        nope_first,
        transpose_output,
    )

    checkAllclose(
        ref, hip_fwd, msg=f"rope_fwd - avg: {hip_fwd_avg:<8.2f} us - {input_msg}\n"
    )
    checkAllclose(
        input.grad,
        hip_bwd,
        msg=f"rope_bwd - avg: {hip_bwd_avg:<8.2f} us - {input_msg}\n",
    )
    checkAllclose(
        ref,
        hip_cached_fwd,
        msg=f"rope_cached_fwd - avg: {hip_cached_fwd_avg:<8.2f} us - {input_msg}\n",
    )
    checkAllclose(
        input.grad,
        hip_cached_bwd,
        msg=f"rope_cached_bwd - avg: {hip_cached_bwd_avg:<8.2f} us - {input_msg}\n",
    )


def test_rope_sbhd_2c(
    input_x,
    input_y,
    freqs,
    grad_x,
    grad_y,
    rotate_style,
    reuse_freqs_front_part,
    nope_first,
    transpose_output,
):
    assert (
        input_x.shape[0:2] == input_y.shape[0:2]
        and input_x.shape[3] == input_y.shape[3]
    )
    assert input_x.dtype == input_y.dtype

    input_msg = f"""
dtype: {input_x.dtype}, \
freq_dtype: {freqs.dtype}, \
dim_input: {str(input_x.shape):<20} - {str(input_y.shape):<20}, \
dim_freqs: {str(freqs.shape):<20}, \
rotate_style: {rotate_style.value}, \
reuse_freqs_front_part: {reuse_freqs_front_part}, \
nope_first: {nope_first}, \
transpose_output: {transpose_output}
"""

    ref_x = ref_rope_sbhd_fwd(
        input_x, freqs, rotate_style, reuse_freqs_front_part, nope_first
    )
    ref_y = ref_rope_sbhd_fwd(
        input_y, freqs, rotate_style, reuse_freqs_front_part, nope_first
    )
    ref_x.backward(grad_x)
    ref_y.backward(grad_y)

    cos = torch.cos(freqs)
    sin = torch.sin(freqs)

    (hip_fwd_x, hip_fwd_y), hip_fwd_avg = hip_rope_2c_fwd(
        input_x,
        input_y,
        freqs,
        rotate_style,
        reuse_freqs_front_part,
        nope_first,
        transpose_output,
    )
    (hip_bwd_x, hip_bwd_y), hip_bwd_avg = hip_rope_2c_bwd(
        grad_x,
        grad_y,
        freqs,
        rotate_style,
        reuse_freqs_front_part,
        nope_first,
        transpose_output,
    )
    (hip_cached_fwd_x, hip_cached_fwd_y), hip_cached_fwd_avg = hip_rope_cached_2c_fwd(
        input_x,
        input_y,
        cos,
        sin,
        rotate_style,
        reuse_freqs_front_part,
        nope_first,
        transpose_output,
    )
    (hip_cached_bwd_x, hip_cached_bwd_y), hip_cached_bwd_avg = hip_rope_cached_2c_bwd(
        grad_x,
        grad_y,
        cos,
        sin,
        rotate_style,
        reuse_freqs_front_part,
        nope_first,
        transpose_output,
    )

    checkAllclose(
        ref_x,
        hip_fwd_x,
        msg=f"rope_2c_fwd_x - avg: {hip_fwd_avg:<8.2f} us - {input_msg}\n",
    )
    checkAllclose(
        ref_y,
        hip_fwd_y,
        msg=f"rope_2c_fwd_y - avg: {hip_fwd_avg:<8.2f} us - {input_msg}\n",
    )
    checkAllclose(
        input_x.grad,
        hip_bwd_x,
        msg=f"rope_2c_bwd_x - avg: {hip_bwd_avg:<8.2f} us - {input_msg}\n",
    )
    checkAllclose(
        input_y.grad,
        hip_bwd_y,
        msg=f"rope_2c_bwd_y - avg: {hip_bwd_avg:<8.2f} us - {input_msg}\n",
    )
    checkAllclose(
        ref_x,
        hip_cached_fwd_x,
        msg=f"rope_cached_2c_fwd_x - avg: {hip_cached_fwd_avg:<8.2f} us - {input_msg}\n",
    )
    checkAllclose(
        ref_y,
        hip_cached_fwd_y,
        msg=f"rope_cached_2c_fwd_y - avg: {hip_cached_fwd_avg:<8.2f} us - {input_msg}\n",
    )
    checkAllclose(
        input_x.grad,
        hip_cached_bwd_x,
        msg=f"rope_cached_2c_bwd_x - avg: {hip_cached_bwd_avg:<8.2f} us - {input_msg}\n",
    )
    checkAllclose(
        input_y.grad,
        hip_cached_bwd_y,
        msg=f"rope_cached_2c_bwd_y - avg: {hip_cached_bwd_avg:<8.2f} us - {input_msg}\n",
    )


def test_rope_sbhd_1c_positions(
    input,
    freqs,
    grad,
    positions,
    offsets,
    rotate_style,
    reuse_freqs_front_part,
    nope_first,
    transpose_output,
):
    input_msg = f"""
dtype: {input.dtype}, \
freq_dtype: {freqs.dtype}, \
dim_input: {str(input.shape):<20}, \
dim_freqs: {str(freqs.shape):<20}, \
dim_positions: {str(positions.shape):<20}, \
dim_offsets: {str(offsets.shape) if offsets is not None else 'None'}, \
rotate_style: {rotate_style.value}, \
reuse_freqs_front_part: {reuse_freqs_front_part}, \
nope_first: {nope_first}, \
transpose_output: {transpose_output}
"""

    ref = ref_rope_sbhd_fwd(
        input,
        freqs[positions if offsets is None else torch.add(positions, offsets)].squeeze(
            -2
        ),
        rotate_style,
        reuse_freqs_front_part,
        nope_first,
    )

    cos = torch.cos(freqs)
    sin = torch.sin(freqs)

    if offsets is None:
        (hip_cached_fwd), hip_cached_fwd_avg = hip_rope_cached_positions_fwd(
            input,
            cos,
            sin,
            positions,
            rotate_style,
            reuse_freqs_front_part,
            nope_first,
            transpose_output,
        )
    else:
        (hip_cached_fwd), hip_cached_fwd_avg = hip_rope_cached_positions_offsets_fwd(
            input,
            cos,
            sin,
            positions,
            offsets,
            rotate_style,
            reuse_freqs_front_part,
            nope_first,
            transpose_output,
        )

    checkAllclose(
        ref,
        hip_cached_fwd,
        msg=f"rope_cached_position_fwd - avg: {hip_cached_fwd_avg:<8.2f} us - {input_msg}\n",
    )


def test_rope_sbhd_2c_positions(
    input_x,
    input_y,
    freqs,
    grad_x,
    grad_y,
    positions,
    offsets,
    rotate_style,
    reuse_freqs_front_part,
    nope_first,
    transpose_output,
):
    input_msg = f"""
dtype: {input_x.dtype}, \
freq_dtype: {freqs.dtype}, \
dim_input: {str(input_x.shape):<20} - {str(input_y.shape):<20}, \
dim_freqs: {str(freqs.shape):<20}, \
dim_positions: {str(positions.shape):<20}, \
dim_offsets: {str(offsets.shape) if offsets is not None else 'None'}, \
rotate_style: {rotate_style.value}, \
reuse_freqs_front_part: {reuse_freqs_front_part}, \
nope_first: {nope_first}, \
transpose_output: {transpose_output}
"""

    ref_x = ref_rope_sbhd_fwd(
        input_x,
        freqs[positions if offsets is None else torch.add(positions, offsets)].squeeze(
            -2
        ),
        rotate_style,
        reuse_freqs_front_part,
        nope_first,
    )
    ref_y = ref_rope_sbhd_fwd(
        input_y,
        freqs[positions if offsets is None else torch.add(positions, offsets)].squeeze(
            -2
        ),
        rotate_style,
        reuse_freqs_front_part,
        nope_first,
    )

    cos = torch.cos(freqs)
    sin = torch.sin(freqs)

    if offsets is None:
        (hip_cached_fwd_x, hip_cached_fwd_y), hip_cached_fwd_avg = (
            hip_rope_cached_positions_2d_fwd(
                input_x,
                input_y,
                cos,
                sin,
                positions,
                rotate_style,
                reuse_freqs_front_part,
                nope_first,
                transpose_output,
            )
        )
    else:
        (hip_cached_fwd_x, hip_cached_fwd_y), hip_cached_fwd_avg = (
            hip_rope_cached_positions_offsets_2d_fwd(
                input_x,
                input_y,
                cos,
                sin,
                positions,
                offsets,
                rotate_style,
                reuse_freqs_front_part,
                nope_first,
                transpose_output,
            )
        )

    checkAllclose(
        ref_x,
        hip_cached_fwd_x,
        msg=f"rope_cached_position_2d_fwd_x - avg: {hip_cached_fwd_avg:<8.2f} us - {input_msg}\n",
    )
    checkAllclose(
        ref_y,
        hip_cached_fwd_y,
        msg=f"rope_cached_position_2d_fwd_y - avg: {hip_cached_fwd_avg:<8.2f} us - {input_msg}\n",
    )


def compare_rope_sbhd_2c_positions_with_legacy(
    input_x,
    input_y,
    freqs,
    positions,
    offsets,
    rotate_style,
    nope_first,
    check_correction=False,
):
    input_msg = f"""dtype: {input_x.dtype}, \
freq_dtype: {freqs.dtype}, \
dim_input: {str(input_x.shape):<20} - {str(input_y.shape):<20}, \
dim_freqs: {str(freqs.shape):<20}, \
dim_positions: {str(positions.shape):<20}, \
dim_offsets: {str(offsets.shape) if offsets is not None else 'None'}, \
rotate_style: {rotate_style.value}, \
nope_first: {nope_first}
"""

    s, b, h_x, d = input_x.shape

    cos = torch.cos(freqs)
    sin = torch.sin(freqs)

    # perftest cannot test correction of inplace operators
    if check_correction:
        ref_x = ref_rope_sbhd_fwd(
            input_x,
            freqs[
                positions if offsets is None else torch.add(positions, offsets)
            ].squeeze(-2),
            rotate_style,
            True,
            nope_first,
        )
        ref_y = ref_rope_sbhd_fwd(
            input_y,
            freqs[
                positions if offsets is None else torch.add(positions, offsets)
            ].squeeze(-2),
            rotate_style,
            True,
            nope_first,
        )
        h_y = input_y.shape[2]
        hip_input_x, hip_input_y = input_x, input_y
        leg_input_x, leg_input_y = input_x.clone().view(
            s * b, -1
        ), input_y.clone().view(s * b, -1)
        if offsets is None:
            aiter.rope_cached_positions_2c_fwd_inplace(
                hip_input_x,
                hip_input_y,
                cos,
                sin,
                positions,
                rotate_style,
                True,
                nope_first,
            )
            aiter.rotary_embedding_fwd(
                positions,
                leg_input_x,
                leg_input_y,
                d,
                cos,
                sin,
                rotate_style is RotateStyle.NEOX,
                nope_first,
            )
        else:
            aiter.rope_cached_positions_offsets_2c_fwd_inplace(
                hip_input_x,
                hip_input_y,
                cos,
                sin,
                positions,
                offsets,
                rotate_style,
                True,
                nope_first,
            )
            aiter.batched_rotary_embedding(
                positions,
                leg_input_x,
                leg_input_y,
                d,
                cos,
                sin,
                rotate_style is RotateStyle.NEOX,
                nope_first,
                cos.size(-1) * 2,
                offsets.view(-1),
            )

        checkAllclose(ref_x, hip_input_x, msg=f"correction: hip_fwd_x - {input_msg}\n")
        checkAllclose(ref_y, hip_input_y, msg=f"correction: hip_fwd_y - {input_msg}\n")
        checkAllclose(
            ref_x,
            leg_input_x.view(s, b, h_x, d),
            msg=f"correction: leg_fwd_x - {input_msg}\n",
        )
        checkAllclose(
            ref_y,
            leg_input_y.view(s, b, h_y, d),
            msg=f"correction: leg_fwd_y - {input_msg}\n",
        )

    leg_cached_fwd_avg = 0.0001
    hip_cached_fwd_avg = 0.0001
    if offsets is None:
        _, leg_cached_fwd_avg = legacy_rope_cached_positions_2d_fwd(
            input_x.view(s * b, -1),
            input_y.view(s * b, -1),
            cos,
            sin,
            positions,
            rotate_style,
            nope_first,
        )
        _, hip_cached_fwd_avg = hip_rope_cached_positions_2d_fwd_inplace(
            input_x, input_y, cos, sin, positions, rotate_style, True, nope_first
        )
    else:
        _, leg_cached_fwd_avg = legacy_rope_cached_positions_offsets_2d_fwd(
            input_x.view(s * b, -1),
            input_y.view(s * b, -1),
            cos,
            sin,
            positions,
            offsets,
            rotate_style,
            nope_first,
        )
        _, hip_cached_fwd_avg = hip_rope_cached_positions_offsets_2d_fwd_inplace(
            input_x,
            input_y,
            cos,
            sin,
            positions,
            offsets,
            rotate_style,
            True,
            nope_first,
        )

    color = (
        "\033[91m"
        if hip_cached_fwd_avg > leg_cached_fwd_avg
        else (
            "\033[92m" if hip_cached_fwd_avg < leg_cached_fwd_avg * 0.75 else "\033[93m"
        )
    )
    endc = "\033[0m"
    print(
        f"{color}{input_msg}hip: {hip_cached_fwd_avg:<8.2f} us. leg: {leg_cached_fwd_avg:<8.2f} us. diff: {100*hip_cached_fwd_avg/leg_cached_fwd_avg}%.\n{endc}"
    )


def test_rope_thd(
    input, cu_seqlens, freqs, grad, rotate_style, reuse_freqs_front_part, nope_first
):
    torch.set_printoptions(profile="full")
    input_msg = f"""
dtype: {input.dtype}, \
freq_dtype: {freqs.dtype}, \
dim_input: {str(input.shape):<20}, \
dim_freqs: {str(freqs.shape):<20}, \
rotate_style: {rotate_style.value}, \
reuse_freqs_front_part: {reuse_freqs_front_part}, \
nope_first: {nope_first}, \
cu_seqlens: {cu_seqlens}
"""
    torch.set_printoptions(profile="default")

    ref = ref_rope_thd_fwd(
        input, cu_seqlens, freqs, rotate_style, reuse_freqs_front_part, nope_first
    )
    ref.backward(grad)

    hip_fwd, hip_fwd_avg = hip_rope_thd_fwd(
        input, cu_seqlens, freqs, rotate_style, reuse_freqs_front_part, nope_first
    )
    hip_bwd, hip_bwd_avg = hip_rope_thd_bwd(
        grad, cu_seqlens, freqs, rotate_style, reuse_freqs_front_part, nope_first
    )

    checkAllclose(
        ref, hip_fwd, msg=f"rope_thd_fwd - avg: {hip_fwd_avg:<8.2f} us - {input_msg}\n"
    )
    checkAllclose(
        input.grad,
        hip_bwd,
        msg=f"rope_thd_bwd - avg: {hip_bwd_avg:<8.2f} us - {input_msg}\n",
    )


def test_rope_2d(input, height, width, freqs_h, freqs_w, grad):
    input_msg = f"""
dtype: {input.dtype}, \
freq_dtype: {freqs_h.dtype}, \
dim_input: {str(input.shape):<20}, \
dim_freqs: {str(freqs_h.shape):<20}
"""

    cos_h = freqs_h.cos()
    sin_h = freqs_h.sin()
    cos_w = freqs_w.cos()
    sin_w = freqs_w.sin()

    ref = ref_rope_2d_fwd(
        input, height, width, cos_h, sin_h, cos_w, sin_w, RotateStyle.NEOX
    )
    ref.backward(grad)

    hip_fwd, hip_fwd_avg = hip_rope_2d_fwd(
        input, cos_h, sin_h, cos_w, sin_w, height, width, RotateStyle.NEOX, False, False
    )
    hip_bwd, hip_bwd_avg = hip_rope_2d_bwd(
        grad, cos_h, sin_h, cos_w, sin_w, height, width, RotateStyle.NEOX, False, False
    )

    checkAllclose(
        ref, hip_fwd, msg=f"rope_2d_fwd - avg: {hip_fwd_avg:<8.2f} us - {input_msg}\n"
    )
    checkAllclose(
        input.grad,
        hip_bwd,
        msg=f"rope_2d_bwd - avg: {hip_bwd_avg:<8.2f} us - {input_msg}\n",
    )


if __name__ == "__main__":
    parser = argparse.ArgumentParser()
    parser.add_argument(
        "--no_check",
        action="store_true",
        help="Do not check correctness of ops. Default: False.",
    )
    parser.add_argument(
        "--compare",
        action="store_true",
        help="Compare with legacy implementation. Default: False",
    )
    parser.add_argument(
        "--compare_check",
        action="store_true",
        help="Check correctness when compare with legacy implementation. Default: False",
    )
    args = parser.parse_args()

    # dtype_ = (dtypes.fp32, dtypes.fp16, dtypes.bf16)
    dtype_ = (dtypes.fp16, dtypes.bf16)
    transpose_output_ = (False, True)
    batch_size_ = (1, 2, 4)
    seq_size_ = (1024, 2048, 4096)
    head_size_ = (32, 64)
    hidden_dim_ = (128, 256)
    # [0]: rotary percentage, [1]: reuse front part, [2]: nope first
    rotary_percent_and_reuse_ = (
        (1.0, True, False),
        (1.0, False, False),
        (0.5, False, False),
        (0.5, True, False),
        (0.5, True, True),
        (0.5, False, True),
    )
    height_ = (32, 64)
    width_ = (32, 64)
    margin_ = (0, 3)
    rotate_style_ = (RotateStyle.NEOX, RotateStyle.GPTJ)

    # Test sbhd format for both cached and uncached
    if not args.no_check:
        for (
            dtype,
            fdtype,
            transpose_output,
            rotate_style,
            rotary_percent_and_reuse,
            b,
            s,
            h,
            d,
        ) in itertools.product(
            dtype_,
            dtype_,
            transpose_output_,
            rotate_style_,
            rotary_percent_and_reuse_,
            batch_size_[-1:],
            seq_size_[1:2],
            head_size_[-1:],
            hidden_dim_[-1:],
        ):
            rotary_percent = rotary_percent_and_reuse[0]
            reuse_freqs_front_part = rotary_percent_and_reuse[1]
            nope_first = (rotary_percent >= 1.0) and rotary_percent_and_reuse[2]
            freqs_ratio = 2 if reuse_freqs_front_part else 1
            input = torch.randn(
                (s, b, h, d), dtype=dtype, device="cuda", requires_grad=True
            )
            freqs = torch.randn(
                (s, 1, 1, int(d * rotary_percent) // freqs_ratio),
                dtype=fdtype,
                device="cuda",
            )
            grad = torch.randn((s, b, h, d), dtype=dtype, device="cuda")
            test_rope_sbhd(
                input,
                freqs,
                grad,
                rotate_style,
                reuse_freqs_front_part,
                nope_first,
                transpose_output,
            )
            input_x = torch.randn(
                (s, b, h, d), dtype=dtype, device="cuda", requires_grad=True
            )
            input_y = torch.randn(
                (s, b, h, d), dtype=dtype, device="cuda", requires_grad=True
            )
            grad_y = torch.randn((s, b, h, d), dtype=dtype, device="cuda")
            test_rope_sbhd_2c(
                input_x,
                input_y,
                freqs,
                grad,
                grad_y,
                rotate_style,
                reuse_freqs_front_part,
                nope_first,
                transpose_output,
            )

    # Test sbhd format for cached with position (and offsets)
    if not args.no_check:
        for (
            dtype,
            fdtype,
            transpose_output,
            rotate_style,
            rotary_percent_and_reuse,
            has_offsets,
            b,
            s,
            h_x,
            h_y,
            d,
        ) in itertools.product(
            dtype_,
            dtype_,
            transpose_output_,
            rotate_style_,
            rotary_percent_and_reuse_,
            (False, True),
            batch_size_[-1:],
            seq_size_[1:2],
            head_size_[-1:],
            head_size_[-1:],
            hidden_dim_[-1:],
        ):
            rotary_percent = rotary_percent_and_reuse[0]
            reuse_freqs_front_part = rotary_percent_and_reuse[1]
            nope_first = (rotary_percent >= 1.0) and rotary_percent_and_reuse[2]
            freqs_ratio = 2 if reuse_freqs_front_part else 1
            freqs = torch.randn(
                (s * 2, 1, 1, int(d * rotary_percent) // freqs_ratio),
                dtype=fdtype,
                device="cuda",
            )
            positions = torch.randint(
                int(s * 0.25) if has_offsets else 0,
                int(s * 0.75) if has_offsets else s,
                (
                    s,
                    b,
                ),
                device="cuda",
            )
            offsets = (
                torch.randint(
                    int(s * -0.25),
                    int(s * 0.25),
                    (
                        s,
                        b,
                    ),
                    device="cuda",
                )
                if has_offsets
                else None
            )
            input_x = torch.randn((s, b, h_x, d), dtype=dtype, device="cuda")
            input_y = torch.randn((s, b, h_y, d), dtype=dtype, device="cuda")
            grad_x = torch.randn((s, b, h_x, d), dtype=dtype, device="cuda")
            grad_y = torch.randn((s, b, h_y, d), dtype=dtype, device="cuda")
            # Note that the below tests cannot run together if backward is enabled due to grad info in inputs are not reset.
            test_rope_sbhd_1c_positions(
                input_x,
                freqs,
                grad_x,
                positions,
                offsets,
                rotate_style,
                reuse_freqs_front_part,
                nope_first,
                transpose_output,
            )
            test_rope_sbhd_2c_positions(
                input_x,
                input_y,
                freqs,
                grad_x,
                grad_y,
                positions,
                offsets,
                rotate_style,
                reuse_freqs_front_part,
                nope_first,
                transpose_output,
            )

    # Compare new with legacy
    if args.compare:
        # [0]: rotary percentage, [1]: reuse front part, [2]: nope first
        # reuse front part should always be True here since legacy implementation doesn't support the opposite setting.
        rotary_percent_and_reuse_compare_ = (
            (1.0, True, False),
            (0.5, True, False),
        )
        for (
            dtype,
            rotate_style,
            rotary_percent_and_reuse,
            has_offsets,
            b,
            s,
            h_x,
            h_y,
            d,
        ) in itertools.product(
            dtype_,  # legacy implementation doesn't support different scalar type between input/output and freqs/sin/cos
            rotate_style_,
            rotary_percent_and_reuse_compare_,
            (False, True),
            batch_size_[-1:],
            seq_size_[1:2],
            head_size_[-1:],
            head_size_[-1:],
            hidden_dim_[-1:],
        ):
            color, endc = "\033[95m", "\033[0m"
            print(
                f"{color}dtype: {dtype}, rotate_style: {rotate_style}, rpar: {rotary_percent_and_reuse}, (s,b,hx,hy,d): {s, b, h_x, h_y, d}, has_offsets: {has_offsets}{endc}"
            )
            rotary_percent = rotary_percent_and_reuse[0]
            reuse_freqs_front_part = rotary_percent_and_reuse[1]
            nope_first = (rotary_percent >= 1.0) and rotary_percent_and_reuse[2]
            freqs_ratio = 2 if reuse_freqs_front_part else 1
            freqs = torch.randn(
                (s * 2, 1, 1, int(d * rotary_percent) // freqs_ratio),
                dtype=dtype,
                device="cuda",
            )
            positions = torch.randint(
                int(s * 0.25) if has_offsets else 0,
                int(s * 0.75) if has_offsets else s,
                (
                    s,
                    b,
                ),
                device="cuda",
            )
            offsets = (
                torch.randint(
                    int(s * -0.25),
                    int(s * 0.25),
                    (
                        s,
                        b,
                    ),
                    device="cuda",
                )
                if has_offsets
                else None
            )
            input_x = torch.randn((s, b, h_x, d), dtype=dtype, device="cuda")
            input_y = torch.randn((s, b, h_y, d), dtype=dtype, device="cuda")
            compare_rope_sbhd_2c_positions_with_legacy(
                input_x,
                input_y,
                freqs,
                positions,
                offsets,
                rotate_style,
                nope_first,
                args.compare_check,
            )

    # Test thd format for uncached
    if not args.no_check:
        cu_seqlens = torch.tensor(
            [
                0,
                100,
                102,
                128,
                233,
                456,
                460,
                711,
                1024,
                1536,
                1739,
                1888,
                2000,
                2001,
                2048,
            ],
            dtype=dtypes.i32,
            device="cuda",
        )
        for (
            dtype,
            fdtype,
            rotate_style,
            rotary_percent_and_reuse,
            h,
            d,
        ) in itertools.product(
            dtype_,
            dtype_,
            rotate_style_,
            rotary_percent_and_reuse_,
            head_size_[-1:],
            hidden_dim_[-1:],
        ):
            rotary_percent = rotary_percent_and_reuse[0]
            reuse_freqs_front_part = rotary_percent_and_reuse[1]
            nope_first = (rotary_percent >= 1.0) and rotary_percent_and_reuse[2]
            freqs_ratio = 2 if reuse_freqs_front_part else 1
            input = torch.randn(
                (cu_seqlens[-1], h, d), dtype=dtype, device="cuda", requires_grad=True
            )
            freqs = torch.randn(
                (cu_seqlens[-1], 1, 1, int(d * rotary_percent) // freqs_ratio),
                dtype=fdtype,
                device="cuda",
            )
            grad = torch.randn((cu_seqlens[-1], h, d), dtype=dtype, device="cuda")
            test_rope_thd(
                input,
                cu_seqlens,
                freqs,
                grad,
                rotate_style,
                reuse_freqs_front_part,
                nope_first,
            )

    # Test 2d image format for cached
    if not args.no_check:
        for dtype, fdtype, b, h, d, height, width, margin in itertools.product(
            dtype_,
            dtype_,
            batch_size_[-1:],
            head_size_[-1:],
            hidden_dim_[-1:],
            height_[-1:],
            width_[-1:],
            margin_,
        ):
            input = torch.randn(
                (b, height * width, h, d),
                dtype=dtype,
                device="cuda",
                requires_grad=True,
            )
            freqs_h = torch.randn(
                (1, height + margin, 1, d // 2), dtype=fdtype, device="cuda"
            )
            freqs_w = torch.randn(
                (1, width + margin, 1, d // 2), dtype=fdtype, device="cuda"
            )
            grad = torch.randn((b, height * width, h, d), dtype=dtype, device="cuda")
            test_rope_2d(input, height, width, freqs_h, freqs_w, grad)
