MegEngine/src/opr/test/misc.cpp

/**
 * \file src/opr/test/misc.cpp
 * MegEngine is Licensed under the Apache License, Version 2.0 (the "License")
 *
 * Copyright (c) 2014-2020 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 "megbrain/opr/misc.h"
#include "megbrain/opr/basic_arith_wrapper.h"
#include "megbrain/opr/blas.h"
#include "megbrain/opr/io.h"
#include "megbrain/opr/tensor_manip.h"
#include "megbrain/opr/utility.h"
#include "megbrain/test/autocheck.h"
#include "megbrain/test/helper.h"
#include "megbrain/test/megdnn_helper.h"

#include <numeric>
#include <random>

using namespace mgb;

namespace {
    void shape_abc(const TensorShape &shape, size_t axis,
            size_t &A, size_t &B, size_t &C) {
        auto acc_mul = [](const size_t *first, const size_t *last) {
            return std::accumulate(
                    first, last, 1u, std::multiplies<size_t>());
        };
        A = acc_mul(shape.shape, shape.shape+axis);
        B = shape.shape[axis];
        C = acc_mul(shape.shape+axis+1, shape.shape+shape.ndim);
    }

    void argsort_data_gen(HostTensorND& dest) {
        mgb_assert(dest.layout().ndim == 2 && dest.layout().is_contiguous());
        size_t m = dest.layout()[0], n = dest.layout()[1];
        auto ptr = dest.ptr<float>();
        RNGxorshf rng{next_rand_seed()};
        std::uniform_real_distribution<float> dist_base{-10.f, 10.f},
                dist_delta{0.1f, 1.2f};
        for (size_t i = 0; i < m; ++i) {
            auto v = dist_base(rng);
            for (size_t j = 0; j < n; ++j) {
                ptr[j] = v;
                v += dist_delta(rng);
            }
            std::shuffle(ptr, ptr + n, rng);
            ptr += n;
        }
    }
}

TEST(TestOprMisc, Argmxx) {
    auto run = [](bool is_max, int32_t axis, TensorShape sshape) {
        auto dshape = sshape;
        dshape.shape[axis] = 1;
        using Checker = AutoOprChecker<1, 1>;
        auto make_graph = [&](const Checker::SymInpArray &inputs) ->
            Checker::SymOutArray {
            if (is_max)
                return {opr::Argmax::make(inputs[0], {axis})};
            else
                return {opr::Argmin::make(inputs[0], {axis})};
        };
        auto better_than = [&](float curr, float best)
        {
            if (is_max)
                return curr > best;
            else
                return curr < best;
        };
        auto fwd = [&](Checker::NumOutArray &out, Checker::NumInpArray inp) {
            out[0].dtype(dtype::Int32()).resize(dshape);
            size_t A, B, C;
            shape_abc(sshape, axis, A, B, C);
            for (size_t a = 0; a < A; ++a) for (size_t c = 0; c < C; ++c) {
                float best_val;
                size_t best_arg = -1;
                if (is_max)
                    best_val = std::numeric_limits<float>::lowest();
                else
                    best_val = std::numeric_limits<float>::max();
                for (size_t b = 0; b < B; ++b) {
                    float curr_val = inp[0]->ptr<float>()[(a*B+b)*C+c];
                    if (better_than(curr_val, best_val)) {
                        best_val = curr_val;
                        best_arg = b;
                    }
                }
                out[0].ptr<int>()[a*C+c] = best_arg;
            }
        };
        Checker{make_graph, fwd}.
            set_input_allow_grad(0, false).
            set_output_allow_grad(0, false).
            run({sshape}).
            run({sshape}).
            run({sshape});
    };
    run(true, 0, {5});
    run(true, 1, {2, 3, 4, 5});
    run(true, 2, {2, 3, 4, 5});
    run(true, 3, {2, 3, 4, 5});
    run(false, 0, {3, 4, 5});
    run(false, 1, {2, 3, 4, 5});
    run(false, 2, {2, 3, 4, 5});
    run(false, 3, {2, 3, 4, 5});
}

TEST(TestOprMisc, Argsort) {
    using Order = opr::Argsort::Param::Order;
    auto run = [](Order order) {
        using Checker = AutoOprChecker<1, 2>;
        auto make_graph = [&](const Checker::SymInpArray& inputs)
                -> Checker::SymOutArray {
            return opr::Argsort::make(inputs[0], order);
        };
        auto fwd = [&](Checker::NumOutArray& out, Checker::NumInpArray inp) {
            size_t m = inp[0]->shape()[0], n = inp[0]->shape()[1];
            auto pi = inp[0]->ptr<float>();
            auto poval = out[0].resize({m, n}).ptr<float>();
            auto poidx = out[1].resize({m, n}).ptr<int>();

            using KV = std::pair<float, int>;
            std::vector<KV> row(n);
            for (size_t i = 0; i < m; ++i) {
                for (size_t j = 0; j < n; ++j) {
                    row[j].first = pi[i * n + j];
                    row[j].second = j;
                }
                if (order == Order::ASCENDING) {
                    std::sort(row.begin(), row.end());
                } else {
                    std::sort(row.begin(), row.end(), std::greater<KV>{});
                }

                for (size_t j = 0; j < n; ++j) {
                    poval[i * n + j] = row[j].first;
                    poidx[i * n + j] = row[j].second;
                }
            }
        };
        Checker::RunOptions opt;
        opt.numdiff_eps = 0.045;
        Checker{make_graph, fwd}
                .set_input_generator(0, argsort_data_gen)
                .set_output_allow_grad(1, false)
                .run({TensorShape{1, 1}}, opt)
                .run({TensorShape{5, 3}}, opt)
                .run({TensorShape{10, 24}}, opt);
    };
    run(Order::ASCENDING);
    run(Order::DESCENDING);
}

TEST(TestOprMisc, Cumsum) {
    using Param = opr::Cumsum::Param;
    auto run = [](const Param &param) {
        using Checker = AutoOprChecker<1, 1>;
        auto make_graph = [&](const Checker::SymInpArray &inputs) ->
            Checker::SymOutArray {
                return {opr::Cumsum::make(inputs[0], param)};
            };
        auto fwd = [&](Checker::NumOutArray &out, Checker::NumInpArray inp) {
            out[0].resize(inp[0]->shape());

            auto pin = inp[0]->ptr<float>(), pout = out[0].ptr<float>();
            size_t A, B, C;
            int real_axis = param.axis;
            if (real_axis < 0) real_axis += 3;
            shape_abc(inp[0]->shape(), real_axis, A, B, C);
            ptrdiff_t stride = C;
            if (param.reverse)
                stride = -stride;
            for (size_t i = 0; i < A; ++ i) {
                for (size_t k = 0; k < C; ++ k) {
                    auto pi = pin + i * B * C + k,
                         po = pout + i * B * C + k;
                    if (param.reverse) {
                        pi += (B - 1) * C;
                        po += (B - 1) * C;
                    }
                    if (param.exclusive) {
                        *po = 0;
                        po += stride;
                    }
                    float sum = 0;
                    for (size_t j = 0; j < B - 1; ++ j) {
                        sum += pi[j * stride];
                        po[j * stride] = sum;
                    }
                    if (!param.exclusive) {
                        po[(B - 1) * stride] = sum + pi[(B - 1) * stride];
                    }
                }
            }
        };
        Checker{make_graph, fwd}.
            run({TensorShape{2, 3, 4}}).
            run({TensorShape{3, 1, 2}}).
            run({TensorShape{4, 2, 3}});
    };

    // test negative axis
    for (int32_t axis = -3; axis < 3; ++axis)
        for (int mask = 0; mask < 4; ++mask)
            run({axis, bool(mask >> 1), bool(mask & 1)});
}

TEST(TestOprMisc, CondTake) {
    using Param = opr::CondTake::Param;
    using Checker = AutoOprChecker<2, 1>;
    auto make_graph = [&](const Checker::SymInpArray &inputs) ->
            Checker::SymOutArray {
        return {opr::CondTake::make(
                inputs[0], inputs[1], {Param::Mode::LT})[0]};
    };

    auto fwd = [&](Checker::NumOutArray &out, Checker::NumInpArray inp) {
        std::vector<float> values;
        auto data = inp[0]->ptr<float>(), mask = inp[1]->ptr<float>();
        auto isize = inp[0]->shape().total_nr_elems();
        for (size_t i = 0; i < isize; ++ i) {
            if (mask[i] < 0) {
                values.push_back(data[i]);
            }
        }
        out[0].resize({values.size()});
        memcpy(out[0].ptr<float>(),
                values.data(), sizeof(float) * values.size());
    };

    auto ensure_nonempty = [](Checker::NumInpArray inp) {
        auto mask = inp[1]->ptr<float>();
        auto isize = inp[1]->shape().total_nr_elems();
        for (size_t i = 0; i < isize; ++ i) {
            if (mask[i] < 0)
                return;
        }
        mask[isize - 1] = -1;
    };

    auto mki = [](const TensorShape &shp) -> Checker::ShapeInpArray {
        return {shp, shp};
    };
    Checker{make_graph, fwd}.
        set_input_allow_grad(1, false).
        set_input_coordinator(ensure_nonempty).
        run(mki({2})).
        run(mki({3, 5, 8})).
        run(mki({100}));
}

TEST(TestOprMisc, CondTakeEmptyOut) {
    using Param = opr::CondTake::Param;
    HostTensorGenerator<> gen;
    auto host_x = gen({1});
    host_x->ptr<float>()[0] = 1;
    auto graph = ComputingGraph::make();
    auto x = opr::Host2DeviceCopy::make(*graph, host_x);
    auto out = opr::CondTake::make(x, x, {Param::Mode::LT});
    HostTensorND host_out0, host_out1;
    auto func = graph->compile({make_callback_copy(out[0], host_out0),
            make_callback_copy(out[1], host_out1)});
    func->execute();
    ASSERT_EQ(TensorShape{0}, host_out0.shape());
    ASSERT_EQ(TensorShape{0}, host_out1.shape());
}

TEST(TestOprMisc, TopKValueOnly) {
    auto run = [](bool dyn_k, bool non_contig) {
        using Checker = AutoOprChecker<1, 1>;
        std::shared_ptr<HostTensorND> host_k;

        SymbolVar var_x0, var_x1;

        auto make_graph = [&](const Checker::SymInpArray& inputs)
                -> Checker::SymOutArray {
            auto k = opr::Host2DeviceCopy::make(
                    *inputs[0].node()->owner_graph(), host_k);
            if (dyn_k) {
                k = opr::MarkDynamicVar::make(k);
            }
            auto x = inputs[0];
            if (non_contig) {
                var_x0 = x;
                x = opr::Subtensor::make(
                        x, {opr::Subtensor::AxisIndexer::make_interval(
                                   1, None, opr::GetVarShape::make(x, 1) / 2,
                                   None)});
                var_x1 = x;
            }
            auto outs = opr::TopK::make(x, k, opr::TopK::Param::Mode::KTH_ONLY);
            return {outs[0]};
        };
        auto fwd = [&](Checker::NumOutArray& out, Checker::NumInpArray inp) {
            auto opr = megdnn_naive_handle()->create_operator<megdnn::TopK>();
            int k = host_k->ptr<int>()[0];
            HostTensorND x = *inp[0];
            if (non_contig) {
                auto layout = x.layout();
                layout.shape[1] /= 2;
                x = x.sub(SubTensorSpec::make_from_layout(layout));
            }

            TensorLayout outl0, outl1;
            opr->deduce_layout(k, x.layout(), outl0, outl1);

            size_t wk_size =
                    opr->get_workspace_in_bytes(k, x.layout(), outl0, outl1);
            std::unique_ptr<dt_byte[]> wk_store{new dt_byte[wk_size]};
            opr->exec(k, x.as_megdnn(), out[0].resize(outl0).as_megdnn(), {},
                      {wk_store.get(), wk_size});
        };
        Checker checker{make_graph, fwd};
        checker.set_input_generator(0, argsort_data_gen);

        host_k = std::make_shared<HostTensorND>(checker.comp_node(),
                                                TensorShape{1}, dtype::Int32{});
        host_k->ptr<int>()[0] = 1;
        Checker::RunOptions opt;
        opt.numdiff_eps = 0.047;
        auto invoke = [&](int k, size_t m, size_t n) {

            host_k->ptr<int>()[0] = k;
            checker.run({TensorShape{m, n}}, opt);
        };

        if (!non_contig) {
            invoke(1, 1, 1);
        }
        invoke(-2, 3, 2);
        invoke(-1, 4, 5);
        invoke(3, 10, 33);
        invoke(-8, 23, 35);

        if (non_contig) {
            ASSERT_EQ(prev_dev_ptr(var_x0), prev_dev_ptr(var_x1));
        }
    };

    for (auto i : {false, true}) {
        for (auto j : {false, true}) {
            run(i, j);
        }
    }
}

TEST(TestOprMisc, TopKSorted) {
    using Checker = AutoOprChecker<1, 2>;
    std::shared_ptr<HostTensorND> host_k;
    auto constexpr mode = opr::TopK::Param::Mode::VALUE_IDX_SORTED;

    auto make_graph =
            [&](const Checker::SymInpArray& inputs) -> Checker::SymOutArray {
        auto k = opr::Host2DeviceCopy::make(*inputs[0].node()->owner_graph(),
                                            host_k);
        auto x = inputs[0];
        return opr::TopK::make(x, k, mode);
    };
    auto fwd = [&](Checker::NumOutArray& out, Checker::NumInpArray inp) {
        auto opr = megdnn_naive_handle()->create_operator<megdnn::TopK>();
        opr->param().mode = mode;
        int k = host_k->ptr<int>()[0];
        TensorLayout outl0, outl1;
        opr->deduce_layout(k, inp[0]->layout(), outl0, outl1);

        size_t wk_size =
                opr->get_workspace_in_bytes(k, inp[0]->layout(), outl0, outl1);
        std::unique_ptr<dt_byte[]> wk_store{new dt_byte[wk_size]};
        opr->exec(k, inp[0]->as_megdnn(), out[0].resize(outl0).as_megdnn(),
                  out[1].resize(outl1).as_megdnn(), {wk_store.get(), wk_size});
    };
    Checker checker{make_graph, fwd};
    checker.set_input_generator(0, argsort_data_gen)
            .set_output_allow_grad(1, false);

    host_k = std::make_shared<HostTensorND>(checker.comp_node(), TensorShape{1},
                                            dtype::Int32{});
    host_k->ptr<int>()[0] = 1;
    Checker::RunOptions opt;
    opt.numdiff_eps = 0.047;
    auto invoke = [&](int k, size_t m, size_t n) {

        host_k->ptr<int>()[0] = k;
        checker.run({TensorShape{m, n}}, opt);
    };

    invoke(1, 1, 1);
    invoke(-1, 3, 5);
    invoke(5, 13, 23);
    invoke(-8, 35, 4);
}

TEST(TestOprMisc, TopKSortedIdxOnly) {
    HostTensorGenerator<> gen;
    auto graph = ComputingGraph::make();
    std::shared_ptr<HostTensorND> host_x = gen({2, 5});
    std::shared_ptr<HostTensorND> host_y = gen({2, 5});
    for (size_t i = 0; i < 10; ++i) {
        host_y->ptr<float>()[i] = 0.0f;
    }
    auto x = opr::Host2DeviceCopy::make(*graph, host_x),
         idx = opr::TopK::make(x, x.make_scalar(3),
                               opr::TopK::Param::Mode::VALUE_IDX_SORTED)[1],
         y = opr::TypeCvt::make(idx, dtype::Float32{}),
         gx = cg::grad(opr::reduce_sum(y, y.make_scalar(1)), x);
    HostTensorND host_gx;
    auto func = graph->compile({make_callback_copy(gx, host_gx)});
    func->execute();
    MGB_ASSERT_TENSOR_EQ(host_gx, *host_y);
}

// vim: syntax=cpp.doxygen foldmethod=marker foldmarker=f{{{,f}}}