import operator_benchmark as op_bench

import torch


"""Microbenchmarks for point-wise unary operator."""


# Configs for pointwise unary ops
unary_ops_configs_short = op_bench.config_list(
    attr_names=["M", "N"],
    attrs=[
        [512, 512],
    ],
    cross_product_configs={
        "device": ["cpu", "cuda"],
    },
    tags=["short"],
)

unary_ops_configs_long = op_bench.cross_product_configs(
    M=[256, 1024], N=[256, 1024], device=["cpu", "cuda"], tags=["long"]
)


class UnaryOpBenchmark(op_bench.TorchBenchmarkBase):
    def init(self, M, N, device, op_func):
        self.inputs = {"input": torch.rand(M, N, device=device)}
        self.op_func = op_func

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


def bernoulli_(input):
    return input.bernoulli_()


def cauchy_(input):
    return input.cauchy_()


def digamma_(input):
    return input.digamma_()


def exponential_(input):
    return input.exponential_()


def normal_(input):
    return input.normal_()


def random_(input):
    return input.random_()


def sign_(input):
    return input.sign_()


def uniform_(input):
    return input.uniform_()


def half_(input):
    return input.half()


def long_(input):
    return input.long()


unary_ops_list = op_bench.op_list(
    attr_names=["op_name", "op_func"],
    attrs=[
        ["abs", torch.abs],
        ["abs_", torch.abs_],
        ["acos", torch.acos],
        ["acos_", torch.acos_],
        ["argsort", torch.argsort],
        ["asin", torch.asin],
        ["asin_", torch.asin_],
        ["atan", torch.atan],
        ["atan_", torch.atan_],
        ["ceil", torch.ceil],
        ["ceil_", torch.ceil_],
        ["clone", torch.clone],
        ["cos", torch.cos],
        ["cos_", torch.cos_],
        ["cosh", torch.cosh],
        ["digamma", torch.digamma],
        ["erf", torch.erf],
        ["erf_", torch.erf_],
        ["erfc", torch.erfc],
        ["erfc_", torch.erfc_],
        ["erfinv", torch.erfinv],
        ["exp", torch.exp],
        ["exp_", torch.exp_],
        ["expm1", torch.expm1],
        ["expm1_", torch.expm1_],
        ["floor", torch.floor],
        ["floor_", torch.floor_],
        ["frac", torch.frac],
        ["frac_", torch.frac_],
        ["hardshrink", torch.hardshrink],
        ["lgamma", torch.lgamma],
        ["log", torch.log],
        ["log10", torch.log10],
        ["log10_", torch.log10_],
        ["log1p", torch.log1p],
        ["log1p_", torch.log1p_],
        ["log2", torch.log2],
        ["log2_", torch.log2_],
        ["log_", torch.log_],
        ["logit", torch.logit],
        ["logit_", torch.logit_],
        ["neg", torch.neg],
        ["neg_", torch.neg_],
        ["reciprocal", torch.reciprocal],
        ["reciprocal_", torch.reciprocal_],
        ["relu", torch.relu],
        ["relu_", torch.relu_],
        ["round", torch.round],
        ["round_", torch.round_],
        ["rsqrt", torch.rsqrt],
        ["rsqrt_", torch.rsqrt_],
        ["sigmoid", torch.sigmoid],
        ["sigmoid_", torch.sigmoid_],
        ["sign", torch.sign],
        ["sgn", torch.sgn],
        ["sin", torch.sin],
        ["sin_", torch.sin_],
        ["sinh", torch.sinh],
        ["sqrt", torch.sqrt],
        ["sqrt_", torch.sqrt_],
        ["square", torch.square],
        ["square_", torch.square_],
        ["tan", torch.tan],
        ["tan_", torch.tan_],
        ["tanh", torch.tanh],
        ["tanh_", torch.tanh_],
        ["trunc", torch.trunc],
        ["trunc_", torch.trunc_],
        ["unique", torch.functional._return_output],
        ["zero_", torch.zero_],
        ["bernoulli_", bernoulli_],
        ["cauchy_", cauchy_],
        ["digamma_", digamma_],
        ["exponential_", exponential_],
        ["normal_", normal_],
        ["random_", random_],
        ["sign_", sign_],
        ["uniform_", uniform_],
        ["half", half_],
        ["long", long_],
    ],
)


op_bench.generate_pt_tests_from_op_list(
    unary_ops_list, unary_ops_configs_short + unary_ops_configs_long, UnaryOpBenchmark
)


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