import hipdnn
import torch


def build_reduction_graph(hipdnn_handle, torch_tensor_x, mode, y_dims, hipdnn_data_type):
    # Create graph
    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="reduction_inference",
    )

    # Create hipdnn tensors
    hipdnn_tensor_x = graph.tensor_like(torch_tensor_x)

    # Create op
    hipdnn_tensor_y = graph.reduction(
        input=hipdnn_tensor_x,
        mode=mode,
        compute_data_type=hipdnn.data_type.FLOAT,
        name="reduction",
    )
    hipdnn_tensor_y.set_dim(y_dims).set_output(True)
    graph.build(hipdnn_handle)

    return (graph, hipdnn_tensor_x, hipdnn_tensor_y)


if __name__ == "__main__":
    # Input dimensions
    batch = 2  # Batch size
    seq_len = 1024  # Number of seq
    embedding_dim = 768  # Number of feature
    mode = hipdnn.reduction_mode.ADD  # Mode

    hipdnn_data_type = hipdnn.data_type.FLOAT
    torch_data_type = torch.float32

    torch_tensor_x = torch.rand(batch, seq_len, embedding_dim, dtype=torch_data_type, device="cuda")

    hipdnn_handle = hipdnn.create_handle()

    graph, hipdnn_tensor_x, hipdnn_tensor_y = build_reduction_graph(
        hipdnn_handle, torch_tensor_x, mode, [batch, seq_len, 1], hipdnn_data_type
    )

    torch_tensor_y = torch.empty(hipdnn_tensor_y.get_dim(), dtype=torch_data_type, device="cuda")
    variant_pack = {
        hipdnn_tensor_x: torch_tensor_x.data_ptr(),
        hipdnn_tensor_y: torch_tensor_y.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("reduction graph execution complete.")