/**
 * \file dnn/test/x86/elemwise_bmark.cpp
 * MegEngine is Licensed under the Apache License, Version 2.0 (the "License")
 *
 * Copyright (c) 2014-2021 Megvii Inc. All rights reserved.
 *
 * Unless required by applicable law or agreed to in writing,
 * software distributed under the License is distributed on an
 * "AS IS" BASIS, WITHOUT ARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 */
#include "test/x86/fixture.h"

#include "megdnn/oprs.h"
#include "test/common/benchmarker.h"
#include "test/common/checker.h"
#include "test/common/rng.h"

using namespace megdnn;
using namespace test;

#define TEST_IN_DIFF_DISTRUBUTION(proportion_of_inf, dataset_number) \
    max_val = 88.3762626647949f / (1 - proportion_of_inf);           \
    UniformFloatRNG rng##dataset_number(0.f, max_val);               \
    B.set_rng(0, &rng##dataset_number);                              \
    B.execs({{355600}, {}});

TEST_F(X86, BENCHMARK_ELEM_EXP_BASED_OPTRS) {
    Benchmarker<ElemwiseForward> B(handle());
    using Mode = ElemwiseForward::Param::Mode;
    // UniformFloatWithZeroRNG rng(80, 100, 0.1);
    printf("Test Optr exp(x)\n");
    B.set_param(Mode::EXP);
    B.execs({{355600}, {}});
    float max_val = 0;
    TEST_IN_DIFF_DISTRUBUTION(0.25, 1)
    TEST_IN_DIFF_DISTRUBUTION(0.5, 2)
    TEST_IN_DIFF_DISTRUBUTION(0.75, 3)
    TEST_IN_DIFF_DISTRUBUTION(0.9999, 4)

    printf("Test Optr tanh(x)\n");
    B.set_param(Mode::TANH);
    B.execs({{355600}, {}});
    max_val = 0;
    TEST_IN_DIFF_DISTRUBUTION(0.25, 5)
    TEST_IN_DIFF_DISTRUBUTION(0.5, 6)
    TEST_IN_DIFF_DISTRUBUTION(0.75, 7)
    TEST_IN_DIFF_DISTRUBUTION(0.9999, 8)

    printf("Test Optr fast_tanh(x)\n");
    B.set_param(Mode::FAST_TANH);
    B.execs({{355600}, {}});

    printf("Test Optr sigmoid(x)\n");
    B.set_param(Mode::SIGMOID);
    B.execs({{355600}, {}});
    max_val = 0;
    TEST_IN_DIFF_DISTRUBUTION(0.25, 13)
    TEST_IN_DIFF_DISTRUBUTION(0.5, 14)
    TEST_IN_DIFF_DISTRUBUTION(0.75, 15)
    TEST_IN_DIFF_DISTRUBUTION(0.9999, 16)

    printf("Test Optr tanh_grad(x)\n");
    B.set_param(Mode::TANH_GRAD);
    B.execs({{355600}, {355600}, {}});

    printf("Test Optr fast_tanh_grad(x)\n");
    B.set_param(Mode::FAST_TANH_GRAD);
    B.execs({{355600}, {355600}, {}});
}

// 1. Unary
#define BENCHMARK_UNARY(Optr, size)  \
    printf("Test for %s \n", #Optr); \
    B.set_param(Mode::Optr);         \
    B.execs(                         \
            {{                       \
                     4,              \
                     4,              \
                     4,              \
                     1 + size / 64,  \
             },                      \
             {}});

// 2. Binary
#define BENCHMARK_BINARY(Optr, size) \
    B.set_param(Mode::Optr);         \
    B.execs({{size}, {size}, {}});

#define BENCHMARK_BINARY_SCALAR(Optr, size) \
    B.set_param(Mode::Optr);                \
    B.execs({{size}, {1}, {}});

#define BENCHMARK_BINARY_1C11(Optr, chan) \
    B.set_param(Mode::Optr);              \
    B.execs({{9, chan, 33, 127}, {1, chan, 1, 1}, {}});

#define BENCHMARK_BINARY_ALL_KINDS(Optr, size) \
    printf("Test for %s \n", #Optr);           \
    BENCHMARK_BINARY(Optr, size)               \
    BENCHMARK_BINARY_SCALAR(Optr, size)        \
    BENCHMARK_BINARY_1C11(Optr, (1 + size / 37719))

// 3. Ternary
#define BENCHMARK_TERNARY(Optr, size) \
    B.set_param(Mode::Optr);          \
    B.execs({{size}, {size}, {size}, {}});

#define BENCHMARK_TERNARY_SCALAR(Optr, size) \
    B.set_param(Mode::Optr);                 \
    B.execs({{size}, {size}, {1}, {}});

#define BENCHMARK_TERNARY_1C11(Optr, chan) \
    B.set_param(Mode::Optr);               \
    B.execs({{1, chan, 1, 1}, {9, chan, 33, 127}, {1, chan, 1, 1}, {}});

#define BENCHMARK_TERNARY_ALL_KINDS(Optr, size) \
    printf("Test for %s \n", #Optr);            \
    BENCHMARK_TERNARY(Optr, size)               \
    BENCHMARK_TERNARY_SCALAR(Optr, size)        \
    BENCHMARK_TERNARY_1C11(Optr, (size / 37719))

#define BENCHMARK_CASE_INT(size)                    \
    BENCHMARK_BINARY_ALL_KINDS(ADD, size)           \
    BENCHMARK_BINARY_ALL_KINDS(SUB, size)           \
    BENCHMARK_BINARY_ALL_KINDS(MUL, size)           \
    BENCHMARK_BINARY_ALL_KINDS(TRUE_DIV, size)      \
    BENCHMARK_BINARY_ALL_KINDS(MIN, size)           \
    BENCHMARK_BINARY_ALL_KINDS(MAX, size)           \
    BENCHMARK_UNARY(RELU, size)                     \
    BENCHMARK_UNARY(ABS, size)                      \
    BENCHMARK_BINARY_ALL_KINDS(FUSE_ADD_RELU, size) \
    BENCHMARK_TERNARY_ALL_KINDS(FUSE_MUL_ADD3, size)

#define BENCHMARK_CASE_FLOAT(size)                  \
    BENCHMARK_CASE_INT(size)                        \
    BENCHMARK_BINARY_ALL_KINDS(FUSE_ADD_TANH, size) \
    BENCHMARK_BINARY_ALL_KINDS(FUSE_ADD_SIGMOID, size)

TEST_F(X86, BENCHMARK_ELEM_EVERY_DTYPE) {
    Benchmarker<ElemwiseForward> B(handle());
    using Mode = ElemwiseForward::Param::Mode;

    printf("\nTest case float32:\n");
    B.set_dtype(0, dtype::Float32());
    B.set_dtype(1, dtype::Float32());
    B.set_dtype(2, dtype::Float32());
    BENCHMARK_CASE_FLOAT(1556011)

    // printf("\nTest case int32:\n");
    // B.set_dtype(0, dtype::Int32());
    // B.set_dtype(1, dtype::Int32());
    // B.set_dtype(2, dtype::Int32());
    // BENCHMARK_CASE_INT(1556011)

    // printf("\nTest case int16:\n");
    // B.set_dtype(0, dtype::Int16());
    // B.set_dtype(1, dtype::Int16());
    // B.set_dtype(2, dtype::Int16());
    // BENCHMARK_CASE_INT(1556011)

    // printf("\nTest case int8:\n");
    // B.set_dtype(0, dtype::Int8());
    // B.set_dtype(1, dtype::Int8());
    // B.set_dtype(2, dtype::Int8());
    // BENCHMARK_CASE_INT(1556011)
}