import operator_benchmark as op_bench

import torch
import torch.nn as nn


"""
Microbenchmarks for MaxPool1d and AvgPool1d operators.
"""

# Configs for pool-1d ops
pool_1d_configs_short = op_bench.config_list(
    attr_names=["kernel", "stride", "N", "C", "L"],
    attrs=[
        [3, 1, 8, 256, 256],
    ],
    cross_product_configs={
        "device": ["cpu", "cuda"],
    },
    tags=["short"],
)

pool_1d_configs_long = op_bench.cross_product_configs(
    kernel=[3],
    stride=[1, 2],
    N=[8, 16],
    C=[3],
    L=[128, 256],
    device=["cpu", "cuda"],
    tags=["long"],
)

pool_1d_ops_list = op_bench.op_list(
    attr_names=["op_name", "op_func"],
    attrs=[
        ["MaxPool1d", nn.MaxPool1d],
        ["AvgPool1d", nn.AvgPool1d],
    ],
)


class Pool1dBenchmark(op_bench.TorchBenchmarkBase):
    def init(self, kernel, stride, N, C, L, device, op_func):
        self.inputs = {"input": torch.rand(N, C, L, device=device)}
        self.op_func = op_func(kernel, stride=stride)

    def forward(self, input):
        return self.op_func(input)


op_bench.generate_pt_tests_from_op_list(
    pool_1d_ops_list, pool_1d_configs_short + pool_1d_configs_long, Pool1dBenchmark
)


"""
Microbenchmarks for MaxPool2d and AvgPool2d operators.
"""


# Configs for pool-2d ops
pool_2d_configs_short = op_bench.config_list(
    attr_names=["kernel", "stride", "N", "C", "H", "W"],
    attrs=[
        [[3, 1], [2, 1], 1, 16, 32, 32],
    ],
    cross_product_configs={
        "device": ["cpu", "cuda"],
    },
    tags=["short"],
)

pool_2d_configs_long = op_bench.cross_product_configs(
    kernel=[[3, 2], [3, 3]],
    stride=[[2, 2]],
    N=[8, 16],
    C=[32],
    H=[32, 64],
    W=[32, 64],
    device=["cpu", "cuda"],
    tags=["long"],
)

pool_2d_ops_list = op_bench.op_list(
    attr_names=["op_name", "op_func"],
    attrs=[
        ["MaxPool2d", nn.MaxPool2d],
        ["AvgPool2d", nn.AvgPool2d],
        ["AdaptiveMaxPool2d", lambda kernel, stride: nn.AdaptiveMaxPool2d(kernel)],
        [
            "FractionalMaxPool2d",
            lambda kernel, stride: nn.FractionalMaxPool2d(kernel, output_size=2),
        ],
    ],
)


class Pool2dBenchmark(op_bench.TorchBenchmarkBase):
    def init(self, kernel, stride, N, C, H, W, device, op_func):
        self.inputs = {"input": torch.rand(N, C, H, W, device=device)}
        self.op_func = op_func(kernel, stride=stride)

    def forward(self, input):
        return self.op_func(input)


op_bench.generate_pt_tests_from_op_list(
    pool_2d_ops_list, pool_2d_configs_short + pool_2d_configs_long, Pool2dBenchmark
)


"""
Microbenchmarks for MaxPool3d and AvgPool3d operators.
"""


# Configs for pool-3d ops
pool_3d_configs_short = op_bench.config_list(
    attr_names=["kernel", "stride", "N", "C", "D", "H", "W"],
    attrs=[
        [[3, 1, 3], [2, 1, 2], 1, 16, 16, 32, 32],
    ],
    cross_product_configs={
        "device": ["cpu", "cuda"],
    },
    tags=["short"],
)

pool_3d_configs_long = op_bench.cross_product_configs(
    kernel=[[3, 2, 3], [3, 3, 3]],
    stride=[[2, 2, 2]],
    N=[8, 16],
    C=[32],
    D=[32],
    H=[32, 64],
    W=[32, 64],
    device=["cpu", "cuda"],
    tags=["long"],
)


pool_3d_ops_list = op_bench.op_list(
    attr_names=["op_name", "op_func"],
    attrs=[
        ["MaxPool3d", nn.MaxPool3d],
        ["AvgPool3d", nn.AvgPool3d],
        ["AdaptiveMaxPool3d", lambda kernel, stride: nn.AdaptiveMaxPool3d(kernel)],
        [
            "FractionalMaxPool3d",
            lambda kernel, stride: nn.FractionalMaxPool3d(kernel, output_size=2),
        ],
    ],
)


class Pool3dBenchmark(op_bench.TorchBenchmarkBase):
    def init(self, kernel, stride, N, C, D, H, W, device, op_func):
        self.inputs = {"input": torch.rand(N, C, D, H, W, device=device)}
        self.op_func = op_func(kernel, stride=stride)

    def forward(self, input):
        return self.op_func(input)


op_bench.generate_pt_tests_from_op_list(
    pool_3d_ops_list, pool_3d_configs_short + pool_3d_configs_long, Pool3dBenchmark
)


if __name__ == "__main__":
    op_bench.benchmark_runner.main()