import hipdnn
import torch


def build_rope_backward_graph(
    hipdnn_handle, torch_tensor_dy, torch_tensor_cos, torch_tensor_sin, hipdnn_data_type
):
    graph = hipdnn.pygraph(
        handle=hipdnn_handle,
        io_data_type=hipdnn_data_type,
        intermediate_data_type=hipdnn.data_type.FLOAT,
        compute_data_type=hipdnn.data_type.FLOAT,
        name="rope_backward_inference",
    )
    hipdnn_tensor_dy = graph.tensor_like(torch_tensor_dy)
    hipdnn_tensor_cos = graph.tensor_like(torch_tensor_cos)
    hipdnn_tensor_sin = graph.tensor_like(torch_tensor_sin)
    hipdnn_tensor_dx = graph.rope_backward(
        dy=hipdnn_tensor_dy, cos=hipdnn_tensor_cos, sin=hipdnn_tensor_sin, name="rope_backward"
    )
    hipdnn_tensor_dx.set_output(True)
    graph.build(hipdnn_handle)
    return (graph, hipdnn_tensor_dy, hipdnn_tensor_cos, hipdnn_tensor_sin, hipdnn_tensor_dx)


if __name__ == "__main__":
    batch, seq_len, dim = 2, 4, 8
    hipdnn_data_type = hipdnn.data_type.FLOAT
    torch_data_type = torch.float32
    torch_tensor_dy = torch.rand(batch, seq_len, dim, dtype=torch_data_type, device="cuda")
    torch_tensor_cos = torch.rand(batch, seq_len, dim, dtype=torch_data_type, device="cuda")
    torch_tensor_sin = torch.rand(batch, seq_len, dim, dtype=torch_data_type, device="cuda")
    hipdnn_handle = hipdnn.create_handle()
    graph, hipdnn_tensor_dy, hipdnn_tensor_cos, hipdnn_tensor_sin, hipdnn_tensor_dx = (
        build_rope_backward_graph(
            hipdnn_handle, torch_tensor_dy, torch_tensor_cos, torch_tensor_sin, hipdnn_data_type
        )
    )
    torch_tensor_dx = torch.empty(batch, seq_len, dim, dtype=torch_data_type, device="cuda")
    variant_pack = {
        hipdnn_tensor_dy: torch_tensor_dy.data_ptr(),
        hipdnn_tensor_cos: torch_tensor_cos.data_ptr(),
        hipdnn_tensor_sin: torch_tensor_sin.data_ptr(),
        hipdnn_tensor_dx: torch_tensor_dx.data_ptr(),
    }
    workspace = torch.empty(graph.get_workspace_size(), dtype=torch.uint8, device="cuda")
    graph.exec(variant_pack=variant_pack, workspace=workspace.data_ptr())
    print("rope_backward graph execution complete.")