OpenI
/
MegEngine
Not watched
1
0
Fork 0
Code
Releases 31 Wiki Activity
Issues 0 Pull Requests 0 Datasets Model Cloudbrain
You can not select more than 25 topics Topics must start with a chinese character,a letter or number, can include dashes ('-') and can be up to 35 characters long.
Tree: 4b08e79da8
Branches Tags
${ item.name }
Create branch ${ searchTerm }
from '4b08e79da8'
${ noResults }
MegEngine/src/opr/test/blas.cpp

930 lines
32 kB
Raw Blame History 

  1. /**

  2.  * \file src/opr/test/blas.cpp

  3.  * MegEngine is Licensed under the Apache License, Version 2.0 (the "License")

  4.  *

  5.  * Copyright (c) 2014-2021 Megvii Inc. All rights reserved.

  6.  *

  7.  * Unless required by applicable law or agreed to in writing,

  8.  * software distributed under the License is distributed on an

  9.  * "AS IS" BASIS, WITHOUT ARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.

  10.  */

  11. 

  12. #include "megbrain/opr/blas.h"

  13. #include "megbrain/test/helper.h"

  14. #include "megbrain/test/autocheck.h"

  15. #include "megbrain/test/megdnn_helper.h"

  16. #include "megbrain/opr/io.h"

  17. #include "megbrain/opr/tensor_manip.h"

  18. #include "megbrain/comp_node_env.h"

  19. #include "megbrain/opr/basic_arith_wrapper.h"

  20. #include "megbrain/opr/tensor_gen.h"

  21. #include <random>

  22. 

  23. using namespace mgb;

  24. 

  25. namespace {

  26. template <typename dt_src, typename dt_dst>

  27. void brute_force_gemm(size_t M, size_t N, size_t K, bool transa, bool transb,

  28.                       const dt_src* x, const dt_src* y, dt_dst* z) {

  29.     for (size_t m = 0; m < M; ++m)

  30.         for (size_t n = 0; n < N; ++n) {

  31.             dt_dst cur = dt_dst(0);

  32.             for (size_t k = 0; k < K; ++k) {

  33.                 cur += x[transa ? (k * M + m) : (m * K + k)] *

  34.                        y[transb ? (n * K + k) : (k * N + n)];

  35.             }

  36.             z[m * N + n] = cur;

  37.         }

  38. }

  39. 

  40. float brute_force_dot(const HostTensorND& a, const HostTensorND& b) {

  41.     auto sz = std::max(a.shape(0), b.shape(0));

  42.     size_t ap = 0, bp = 0;

  43.     float ret = 0;

  44.     auto pa = a.ptr<float>(), pb = b.ptr<float>();

  45.     auto as = a.layout().stride[0], bs = b.layout().stride[0];

  46.     if (a.shape(0) != sz)

  47.         as = 0;

  48.     if (b.shape(0) != sz)

  49.         bs = 0;

  50.     for (size_t i = 0; i < sz; ++i) {

  51.         ret += pa[ap] * pb[bp];

  52.         ap += as;

  53.         bp += bs;

  54.     }

  55.     return ret;

  56. }

  57. 

  58. // (m,k) * (k,n) = (m,n)

  59. void run_sgemm_test(bool transa, bool transb) {

  60.     using Checker = AutoOprChecker<2, 1>;

  61.     auto make_graph =

  62.             [&](const Checker::SymInpArray& inputs) -> Checker::SymOutArray {

  63.         auto param = opr::MatrixMul::Param{transa, transb};

  64.         return {opr::MatrixMul::make(inputs[0], inputs[1], param)};

  65.     };

  66.     auto fwd = [&](Checker::NumOutArray& dest, Checker::NumInpArray inp) {

  67.         size_t M, N, K;

  68.         M = inp[0]->shape().shape[0];

  69.         K = inp[0]->shape().shape[1];

  70.         if (transa)

  71.             std::swap(M, K);

  72.         N = inp[1]->shape().shape[transb ? 0 : 1];

  73. 

  74.         auto z = dest[0].comp_node(inp[0]->comp_node())

  75.                          .resize({M, N})

  76.                          .ptr<float>();

  77.         // brute-force gemm

  78.         brute_force_gemm(M, N, K, transa, transb, inp[0]->ptr<float>(),

  79.                          inp[1]->ptr<float>(), z);

  80.     };

  81. 

  82.     auto mkshp = [](bool trans, size_t m, size_t k) {

  83.         TensorShape rst{m, k};

  84.         if (trans)

  85.             std::swap(rst.shape[0], rst.shape[1]);

  86.         return rst;

  87.     };

  88.     using namespace std::placeholders;

  89.     auto mkx = std::bind(mkshp, transa, _1, _2);

  90.     auto mky = std::bind(mkshp, transb, _1, _2);

  91. 

  92.     Checker::RunOptions opt;

  93.     opt.numdiff_eps = 1;

  94.     Checker(make_graph, fwd)

  95.             .run({mkx(4, 6), mky(6, 2)}, opt)

  96.             .run({mkx(2, 3), mky(3, 100)}, opt)

  97.             .run({mkx(20, 3), mky(3, 20)}, opt);

  98. }

  99. 

  100. #define FWD_BATCH_GEMM(dt_src, dt_dst)                                       \

  101.     [transa, transb](Checker::NumOutArray& dest, Checker::NumInpArray inp) { \

  102.         bool ta(transa), tb(transb);                                         \

  103.         HostTensorND a, b;                                                   \

  104.         size_t B, M, N, K;                                                   \

  105.         a.copy_from(*(inp[0]));                                              \

  106.         b.copy_from(*(inp[1]));                                              \

  107.         B = a.shape().shape[0];                                              \

  108.         M = a.shape().shape[1];                                              \

  109.         K = a.shape().shape[2];                                              \

  110.         N = b.shape().shape[tb ? 1 : 2];                                     \

  111.         if (ta)                                                              \

  112.             std::swap(M, K);                                                 \

  113.         auto x = a.ptr<dt_src>(), y = b.ptr<dt_src>();                       \

  114.         auto z = dest[0].resize({B, M, N}).ptr<dt_dst>();                    \

  115.         for (size_t b = 0; b < B; ++b) {                                     \

  116.             brute_force_gemm(M, N, K, ta, tb, x + b * M * K, y + b * K * N,  \

  117.                              z + b * M * N);                                 \

  118.         }                                                                    \

  119.     }

  120. 

  121. void run_batched_sgemm_test(bool transa, bool transb) {

  122.     using Checker = AutoOprChecker<2, 1>;

  123.     auto make_graph =

  124.             [&](const Checker::SymInpArray& inputs) -> Checker::SymOutArray {

  125.         return {opr::BatchedMatrixMul::make(inputs[0], inputs[1],

  126.                                             {transa, transb})};

  127.     };

  128. 

  129.     auto fwd = FWD_BATCH_GEMM(float, float);

  130. 

  131.     auto mkshp = [](bool trans, size_t b, size_t m, size_t k) {

  132.         TensorShape rst{b, m, k};

  133.         if (trans)

  134.             std::swap(rst.shape[1], rst.shape[2]);

  135.         return rst;

  136.     };

  137.     using namespace std::placeholders;

  138.     auto mkx = std::bind(mkshp, transa, _1, _2, _3);

  139.     auto mky = std::bind(mkshp, transb, _1, _2, _3);

  140. 

  141.     Checker::RunOptions opt;

  142.     opt.numdiff_eps = 1;

  143.     Checker(make_graph, fwd)

  144.             .run({mkx(3, 5, 7), mky(3, 7, 2)}, opt)

  145.             .run({mkx(64, 1, 2), mky(64, 2, 1)}, opt)

  146.             .run({mkx(1, 2, 3), mky(1, 3, 4)}, opt);

  147. }

  148. 

  149. auto gen_fp16 = [](HostTensorND& dest) {

  150.     RNGxorshf rng{next_rand_seed()};

  151.     auto rand_real = [&rng]() {

  152.         std::uniform_real_distribution<float> dist(-1, 1);

  153.         return dist(rng);

  154.     };

  155.     auto ptr = dest.ptr<dt_float16>();

  156.     size_t elems = dest.shape().total_nr_elems();

  157.     for (size_t i = 0; i < elems; i++) {

  158.         ptr[i] = dt_float16(rand_real());

  159.     }

  160. };

  161. 

  162. auto gen_int8 = [](HostTensorND& dest) {

  163.     HostTensorGenerator<dtype::Int8, RandomDistribution::UNIFORM>

  164.             int8_generator{-128, 127};

  165.     dest = *int8_generator(dest.shape(), dest.comp_node());

  166. };

  167. 

  168. void run_batched_hgemm_test(bool transa, bool transb) {

  169.     using Checker = AutoOprChecker<2, 1>;

  170.     auto make_graph =

  171.             [&](const Checker::SymInpArray& inputs) -> Checker::SymOutArray {

  172.         return {opr::BatchedMatrixMul::make(inputs[0], inputs[1],

  173.                                             {transa, transb})};

  174.     };

  175.     auto fwd = FWD_BATCH_GEMM(dt_float16, dt_float16);

  176.     auto mkshp = [](bool trans, size_t b, size_t m, size_t k) {

  177.         TensorShape rst{b, m, k};

  178.         if (trans)

  179.             std::swap(rst.shape[1], rst.shape[2]);

  180.         return rst;

  181.     };

  182. 

  183.     using namespace std::placeholders;

  184.     auto mkx = std::bind(mkshp, transa, _1, _2, _3);

  185.     auto mky = std::bind(mkshp, transb, _1, _2, _3);

  186. 

  187.     Checker checker(make_graph, fwd);

  188.     Checker::RunOptions opt;

  189.     opt.outputs_max_err = 1e-2;

  190. 

  191.     checker.set_input_dtype(0, dtype::Float16())

  192.             .set_input_dtype(1, dtype::Float16())

  193.             .set_input_generator(0, gen_fp16)

  194.             .set_input_generator(1, gen_fp16)

  195.             .set_input_allow_grad(0, false)

  196.             .set_input_allow_grad(1, false)

  197.             .set_output_allow_grad(0, false);

  198. 

  199.     checker.run({mkx(3, 5, 7), mky(3, 7, 2)}, opt)

  200.             .run({mkx(64, 1, 2), mky(64, 2, 1)}, opt)

  201.             .run({mkx(1, 2, 3), mky(1, 3, 4)}, opt);

  202. }

  203. 

  204. void run_batched_igemm_test(bool transa, bool transb) {

  205.     using Checker = AutoOprChecker<2, 1>;

  206.     auto make_graph =

  207.             [&](const Checker::SymInpArray& inputs) -> Checker::SymOutArray {

  208.         return {opr::BatchedMatrixMul::make(inputs[0], inputs[1],

  209.                                             {transa, transb})};

  210.     };

  211. 

  212.     auto fwd = FWD_BATCH_GEMM(int8_t, int32_t);

  213. 

  214.     auto mkshp = [](bool trans, size_t b, size_t m, size_t k) {

  215.         TensorShape rst{b, m, k};

  216.         if (trans)

  217.             std::swap(rst.shape[1], rst.shape[2]);

  218.         return rst;

  219.     };

  220. 

  221.     using namespace std::placeholders;

  222.     auto mkx = std::bind(mkshp, transa, _1, _2, _3);

  223.     auto mky = std::bind(mkshp, transb, _1, _2, _3);

  224. 

  225.     Checker::RunOptions opt;

  226.     opt.numdiff_eps = 1;

  227.     Checker checker(make_graph, fwd);

  228. 

  229.     checker.set_input_dtype(0, dtype::Int8())

  230.             .set_input_dtype(1, dtype::Int8())

  231.             .set_input_generator(0, gen_int8)

  232.             .set_input_generator(1, gen_int8)

  233.             .set_input_allow_grad(0, false)

  234.             .set_input_allow_grad(1, false)

  235.             .set_output_allow_grad(0, false);

  236. 

  237.     checker.run({mkx(3, 5, 7), mky(3, 7, 2)}, opt)

  238.             .run({mkx(64, 1, 2), mky(64, 2, 1)}, opt)

  239.             .run({mkx(1, 2, 3), mky(1, 3, 4)}, opt);

  240. }

  241. 

  242. template <typename ctype>

  243. float getter(ctype val) {

  244.     return val;

  245. }

  246. 

  247. template <>

  248. float getter<dt_qint32>(dt_qint32 val) {

  249.     return (float)val.as_int32();

  250. }

  251. 

  252. template <typename dt_src, typename dt_dst>

  253. void run_trans_inp_test_case(bool trans_a, bool trans_b) {

  254.     HostTensorGenerator<typename DTypeTrait<dt_src>::dtype> gen;

  255.     std::shared_ptr<HostTensorND> host_x = gen({1, 1}), host_y = gen({1, 1});

  256.     auto graph = ComputingGraph::make();

  257.     auto do_trans = [](SymbolVar x) {

  258.         return opr::Dimshuffle::make(x, {1, 0});

  259.     };

  260.     auto x = opr::Host2DeviceCopy::make(*graph, host_x),

  261.          y = opr::Host2DeviceCopy::make(*graph, host_y);

  262.     if (trans_a) {

  263.         x = do_trans(x);

  264.     }

  265.     if (trans_b) {

  266.         y = do_trans(y);

  267.     }

  268.     OperatorNodeConfig config;

  269.     if (DTypeTrait<dt_dst>::enumv == DTypeEnum::Int16) {

  270.         config.output_dtype(dtype::Int16());

  271.     }

  272.     auto z = opr::MatrixMul::make(x, y, {}, {}, config);

  273.     HostTensorND host_z;

  274.     auto func = graph->compile({make_callback_copy(z, host_z)});

  275. 

  276.     auto run = [&](size_t M, size_t K, size_t N) {

  277.         *host_x = *(trans_a ? gen({K, M}) : gen({M, K}));

  278.         *host_y = *(trans_b ? gen({N, K}) : gen({K, N}));

  279.         func->execute();

  280.         ASSERT_EQ(TensorShape({M, N}), host_z.shape());

  281.         ASSERT_EQ(!trans_a, x.node()->dev_tensor().layout().is_contiguous());

  282.         ASSERT_EQ(!trans_b, y.node()->dev_tensor().layout().is_contiguous());

  283. 

  284.         auto px = host_x->ptr<dt_src>(), py = host_y->ptr<dt_src>();

  285.         auto pz = host_z.ptr<dt_dst>();

  286.         auto make_strd = [](bool trans, int h, int w, int* dst) {

  287.             if (trans) {

  288.                 dst[0] = 1;

  289.                 dst[1] = h;

  290.             } else {

  291.                 dst[0] = w;

  292.                 dst[1] = 1;

  293.             }

  294.         };

  295.         int strd_x[2], strd_y[2];

  296.         make_strd(trans_a, M, K, strd_x);

  297.         make_strd(trans_b, K, N, strd_y);

  298.         for (size_t i = 0; i < M; ++i) {

  299.             for (size_t j = 0; j < N; ++j) {

  300.                 dt_dst sum = 0;

  301.                 for (size_t k = 0; k < K; ++k) {

  302.                     dt_dst xv = px[i * strd_x[0] + k * strd_x[1]],

  303.                            yv = py[k * strd_y[0] + j * strd_y[1]];

  304.                     sum += xv * yv;

  305.                 }

  306.                 MGB_ASSERT_FLOAT_EQ(getter(sum), getter(pz[i * N + j]))

  307.                         << trans_a << ' ' << trans_b;

  308.             }

  309.         }

  310.     };

  311.     run(4, 8, 12);

  312.     run(8, 12, 16);

  313. }

  314. 

  315. template <typename dt_src, typename dt_dst>

  316. void run_trans_inp_test() {

  317.     for (bool ta : {false, true}) {

  318.         for (bool tb : {false, true}) {

  319.             run_trans_inp_test_case<dt_src, dt_dst>(ta, tb);

  320.         }

  321.     }

  322. }

  323. 

  324. template <typename dt_src, typename dt_dst>

  325. void inline mul_add(dt_src& a, dt_src& b, dt_dst& c) {

  326.     c += dt_dst(a) * dt_dst(b);

  327. }

  328. 

  329. template <>

  330. void inline mul_add(dt_qint8& a, dt_qint8& b, dt_qint32& c) {

  331.     c += dt_qint32(a.as_int8()) * dt_qint32(b.as_int8());

  332. }

  333. 

  334. template <typename dt_gen>

  335. std::shared_ptr<HostTensorND> bgemm_gen(const TensorShape& shp) {

  336.     HostTensorGenerator<typename DTypeTrait<dt_gen>::dtype> gen;

  337.     return gen(shp);

  338. }

  339. 

  340. template <>

  341. std::shared_ptr<HostTensorND> bgemm_gen<dt_float16>(const TensorShape& shp) {

  342.     CompNode cn = CompNode::load("xpu0");

  343.     std::shared_ptr<HostTensorND> ret =

  344.             std::make_shared<HostTensorND>(cn, dtype::Float16{});

  345.     (*ret).resize(shp);

  346.     gen_fp16(*ret);

  347.     return ret;

  348. }

  349. 

  350. template <typename dt_src, typename dt_dst>

  351. void run_bgemm_trans_inp_test_case(bool trans_a, bool trans_b) {

  352.     std::shared_ptr<HostTensorND> host_x = bgemm_gen<dt_src>({1, 1, 1}),

  353.                                   host_y = bgemm_gen<dt_src>({1, 1, 1});

  354. 

  355.     auto graph = ComputingGraph::make();

  356.     auto x = opr::Host2DeviceCopy::make(*graph, host_x),

  357.          y = opr::Host2DeviceCopy::make(*graph, host_y);

  358. 

  359.     trans_a ? (x = opr::Dimshuffle::make(x, {0, 2, 1})) : 0;

  360.     trans_b ? (y = opr::Dimshuffle::make(y, {0, 2, 1})) : 0;

  361. 

  362.     auto z = opr::BatchedMatrixMul::make(x, y, {}, {}, OperatorNodeConfig{});

  363.     HostTensorND host_z;

  364.     auto func = graph->compile({make_callback_copy(z, host_z)});

  365.     auto run = [&](size_t B, size_t M, size_t K, size_t N) {

  366.         *host_x = *(trans_a ? bgemm_gen<dt_src>({B, K, M})

  367.                             : bgemm_gen<dt_src>({B, M, K}));

  368.         *host_y = *(trans_b ? bgemm_gen<dt_src>({B, N, K})

  369.                             : bgemm_gen<dt_src>({B, K, N}));

  370.         func->execute();

  371.         ASSERT_EQ(TensorShape({B, M, N}), host_z.shape());

  372.         ASSERT_EQ(!trans_a, x.node()->dev_tensor().layout().is_contiguous());

  373.         ASSERT_EQ(!trans_b, y.node()->dev_tensor().layout().is_contiguous());

  374. 

  375.         int strd_x[3], strd_y[3];

  376.         auto px = host_x->ptr<dt_src>(), py = host_y->ptr<dt_src>();

  377.         auto pz = host_z.ptr<dt_dst>();

  378.         auto make_strd = [](bool trans, int h, int w, int* dst) {

  379.             dst[0] = h * w;

  380.             dst[1] = trans ? 1 : w;

  381.             dst[2] = trans ? h : 1;

  382.         };

  383.         make_strd(trans_a, M, K, strd_x);

  384.         make_strd(trans_b, K, N, strd_y);

  385.         for (size_t b = 0; b < B; ++b)

  386.             for (size_t i = 0; i < M; ++i)

  387.                 for (size_t j = 0; j < N; ++j) {

  388.                     dt_dst sum = dt_dst(0);

  389.                     for (size_t k = 0; k < K; ++k) {

  390.                         dt_src xv = px[b * strd_x[0] + i * strd_x[1] +

  391.                                        k * strd_x[2]],

  392.                                yv = py[b * strd_y[0] + k * strd_y[1] +

  393.                                        j * strd_y[2]];

  394.                         mul_add(xv, yv, sum);

  395.                     }

  396.                     MGB_ASSERT_FLOAT_NEAR(getter(sum),

  397.                                           getter(pz[(b * M + i) * N + j]), 5e-3)

  398.                             << trans_a << ' ' << trans_b;

  399.                 }

  400.     };

  401.     run(2, 4, 8, 12);

  402.     run(2, 8, 12, 16);

  403. }

  404. 

  405. }  // anonymous namespace

  406. 

  407. TEST(TestOprBlas, MatrixMul_NN) {

  408.     run_sgemm_test(false, false);

  409. }

  410. 

  411. TEST(TestOprBlas, MatrixMul_NT) {

  412.     run_sgemm_test(false, true);

  413. }

  414. 

  415. TEST(TestOprBlas, MatrixMul_TN) {

  416.     run_sgemm_test(true, false);

  417. }

  418. 

  419. TEST(TestOprBlas, MatrixMul_TT) {

  420.     run_sgemm_test(true, true);

  421. }

  422. 

  423. TEST(TestOprDNN, MatrixMulExePolicy) {

  424.     using Param = opr::MatrixMul::Param;

  425.     Param param;

  426.     using Policy = opr::MatrixMul::ExecutionPolicy;

  427.     using S = Policy::Strategy;

  428. 

  429.     auto cn = CompNode::load("cpux");

  430. 

  431. #if MGB_ENABLE_FASTRUN

  432.     for (auto strategy :

  433.          SmallVector<S>{S::PROFILE, S::HEURISTIC, S::PROFILE | S::REPRODUCIBLE,

  434.                         S::PROFILE | S::HEURISTIC}) {

  435. #else

  436.     for (auto strategy: {S:HEURISTIC, S::PROFILE | S::HEURISTIC}) {

  437. #endif

  438. 

  439.         auto graph = ComputingGraph::make();

  440.         HostTensorGenerator<> gen;

  441. 

  442.         auto mkvar = [&](const char* name, const TensorShape& shp) {

  443.             return opr::Host2DeviceCopy::make(*graph, gen(shp), cn)

  444.                     .rename(name);

  445.         };

  446. 

  447.         auto A = mkvar("A", {32, 64});

  448.         auto B = mkvar("B", {64, 32});

  449. 

  450.         Policy policy;

  451.         policy.strategy = strategy;

  452. 

  453.         auto C = opr::MatrixMul::make(A, B, param, policy);

  454.         HostTensorND host_c;

  455.         auto func = graph->compile({make_callback_copy(C, host_c)});

  456.         func->execute();

  457.     }

  458. }

  459. 

  460. 

  461. TEST(TestOprBlas, BatchedMatrixMulFp32_NN) {

  462.     run_batched_sgemm_test(false, false);

  463. }

  464. 

  465. TEST(TestOprBlas, BatchedMatrixMulFp32_NT) {

  466.     run_batched_sgemm_test(false, true);

  467. }

  468. 

  469. TEST(TestOprBlas, BatchedMatrixMulFp32_TN) {

  470.     run_batched_sgemm_test(true, false);

  471. }

  472. 

  473. TEST(TestOprBlas, BatchedMatrixMulFp32_TT) {

  474.     run_batched_sgemm_test(true, true);

  475. }

  476. 

  477. TEST(TestOprBlas, BatchedMatrixMulFp16_NN) {

  478.     run_batched_hgemm_test(false, false);

  479. }

  480. 

  481. TEST(TestOprBlas, BatchedMatrixMulFp16_NT) {

  482.     run_batched_hgemm_test(false, true);

  483. }

  484. 

  485. TEST(TestOprBlas, BatchedMatrixMulFp16_TN) {

  486.     run_batched_hgemm_test(true, false);

  487. }

  488. 

  489. TEST(TestOprBlas, BatchedMatrixMulFp16_TT) {

  490.     run_batched_hgemm_test(true, true);

  491. }

  492. 

  493. TEST(TestOprBlas, BatchedMatrixMulInt8_NN) {

  494.     if (CompNode::load("xpux").device_type() == CompNode::DeviceType::CUDA &&

  495.         !check_compute_capability(6, 1)) {

  496.         return;

  497.     }

  498.     run_batched_igemm_test(false, false);

  499. }

  500. 

  501. TEST(TestOprBlas, BatchedMatrixMulInt8_NT) {

  502.     if (CompNode::load("xpux").device_type() == CompNode::DeviceType::CUDA &&

  503.         !check_compute_capability(6, 1)) {

  504.         return;

  505.     }

  506.     run_batched_igemm_test(false, true);

  507. }

  508. 

  509. TEST(TestOprBlas, BatchedMatrixMulInt8_TN) {

  510.     if (CompNode::load("xpux").device_type() == CompNode::DeviceType::CUDA &&

  511.         !check_compute_capability(6, 1)) {

  512.         return;

  513.     }

  514.     run_batched_igemm_test(true, false);

  515. }

  516. 

  517. TEST(TestOprBlas, BatchedMatrixMulInt8_TT) {

  518.     if (CompNode::load("xpux").device_type() == CompNode::DeviceType::CUDA &&

  519.         !check_compute_capability(6, 1)) {

  520.         return;

  521.     }

  522.     run_batched_igemm_test(true, true);

  523. }

  524. 

  525. TEST(TestOprBlas, TransBatchedMatrixMulFp32_NN) {

  526.     run_bgemm_trans_inp_test_case<float, float>(false, false);

  527. }

  528. 

  529. TEST(TestOprBlas, TransBatchedMatrixMulFp32_NT) {

  530.     run_bgemm_trans_inp_test_case<float, float>(false, true);

  531. }

  532. 

  533. TEST(TestOprBlas, TransBatchedMatrixMulFp32_TN) {

  534.     run_bgemm_trans_inp_test_case<float, float>(true, false);

  535. }

  536. 

  537. TEST(TestOprBlas, TransBatchedMatrixMulFp32_TT) {

  538.     run_bgemm_trans_inp_test_case<float, float>(true, true);

  539. }

  540. 

  541. TEST(TestOprBlas, TransBatchedMatrixMulInt8_NN) {

  542.     if (CompNode::load("xpux").device_type() == CompNode::DeviceType::CUDA &&

  543.         !check_compute_capability(6, 1)) {

  544.         return;

  545.     }

  546.     run_bgemm_trans_inp_test_case<int8_t, int32_t>(false, false);

  547. }

  548. 

  549. TEST(TestOprBlas, TransBatchedMatrixMulInt8_NT) {

  550.     if (CompNode::load("xpux").device_type() == CompNode::DeviceType::CUDA &&

  551.         !check_compute_capability(6, 1)) {

  552.         return;

  553.     }

  554.     run_bgemm_trans_inp_test_case<int8_t, int32_t>(false, true);

  555. }

  556. 

  557. TEST(TestOprBlas, TransBatchedMatrixMulInt8_TN) {

  558.     if (CompNode::load("xpux").device_type() == CompNode::DeviceType::CUDA &&

  559.         !check_compute_capability(6, 1)) {

  560.         return;

  561.     }

  562.     run_bgemm_trans_inp_test_case<int8_t, int32_t>(true, false);

  563. }

  564. 

  565. TEST(TestOprBlas, TransBatchedMatrixMulInt8_TT) {

  566.     if (CompNode::load("xpux").device_type() == CompNode::DeviceType::CUDA &&

  567.         !check_compute_capability(6, 1)) {

  568.         return;

  569.     }

  570.     run_bgemm_trans_inp_test_case<int8_t, int32_t>(true, true);

  571. }

  572. 

  573. TEST(TestOprBlas, TransBatchedMatrixMulFp16_NN) {

  574.     run_bgemm_trans_inp_test_case<dt_float16, dt_float16>(false, false);

  575. }

  576. 

  577. TEST(TestOprBlas, TransBatchedMatrixMulFp16_NT) {

  578.     run_bgemm_trans_inp_test_case<dt_float16, dt_float16>(false, true);

  579. }

  580. 

  581. TEST(TestOprBlas, TransBatchedMatrixMulFp16_TN) {

  582.     run_bgemm_trans_inp_test_case<dt_float16, dt_float16>(true, false);

  583. }

  584. 

  585. TEST(TestOprBlas, TransBatchedMatrixMulFp16_TT) {

  586.     run_bgemm_trans_inp_test_case<dt_float16, dt_float16>(true, true);

  587. }

  588. 

  589. TEST(TestOprBlas, TransBatchedMatrixMulQS8_NN) {

  590.     if (CompNode::load("xpux").device_type() == CompNode::DeviceType::CUDA &&

  591.         !check_compute_capability(6, 1)) {

  592.         return;

  593.     }

  594.     run_bgemm_trans_inp_test_case<dt_qint8, dt_qint32>(false, false);

  595. }

  596. 

  597. TEST(TestOprBlas, TransBatchedMatrixMulQS8_NT) {

  598.     if (CompNode::load("xpux").device_type() == CompNode::DeviceType::CUDA &&

  599.         !check_compute_capability(6, 1)) {

  600.         return;

  601.     }

  602.     run_bgemm_trans_inp_test_case<dt_qint8, dt_qint32>(false, true);

  603. }

  604. 

  605. TEST(TestOprBlas, TransBatchedMatrixMulQS8_TN) {

  606.     if (CompNode::load("xpux").device_type() == CompNode::DeviceType::CUDA &&

  607.         !check_compute_capability(6, 1)) {

  608.         return;

  609.     }

  610.     run_bgemm_trans_inp_test_case<dt_qint8, dt_qint32>(true, false);

  611. }

  612. 

  613. TEST(TestOprBlas, TransBatchedMatrixMulQS8_TT) {

  614.     if (CompNode::load("xpux").device_type() == CompNode::DeviceType::CUDA &&

  615.         !check_compute_capability(6, 1)) {

  616.         return;

  617.     }

  618.     run_bgemm_trans_inp_test_case<dt_qint8, dt_qint32>(true, true);

  619. }

  620. 

  621. TEST(TestOprBlas, DotBasic) {

  622.     HostTensorGenerator<> gen;

  623.     auto host_x = gen({123}), host_y = gen({123});

  624.     auto graph = ComputingGraph::make();

  625.     auto x = opr::Host2DeviceCopy::make(*graph, host_x),

  626.          y = opr::Host2DeviceCopy::make(*graph, host_y),

  627.          z = opr::Dot::make(x, y);

  628.     HostTensorND host_z;

  629.     auto func = graph->compile({make_callback_copy(z, host_z)});

  630.     func->execute();

  631.     MGB_ASSERT_FLOAT_EQ(brute_force_dot(*host_x, *host_y),

  632.                         *host_z.ptr<float>());

  633. }

  634. 

  635. TEST(TestOprBlas, Dot) {

  636.     using Checker = AutoOprChecker<2, 1>;

  637. 

  638.     auto make_graph =

  639.             [&](const Checker::SymInpArray& inputs) -> Checker::SymOutArray {

  640.         return {opr::Dot::make(inputs[0], inputs[1])};

  641.     };

  642. 

  643.     auto fwd = [](Checker::NumOutArray& dest, Checker::NumInpArray inp) {

  644.         auto &&i0 = *inp[0], &&i1 = *inp[1];

  645.         auto&& out = dest[0].resize({1});

  646.         *out.ptr<float>() = brute_force_dot(i0, i1);

  647.     };

  648. 

  649.     Checker(make_graph, fwd)

  650.             .run({TensorShape{15}, TensorShape{1}})

  651.             .run({TensorShape{1}, TensorShape{16}})

  652.             .run({TensorShape{23}, TensorShape{23}})

  653.             .run({TensorShape{1000}, TensorShape{1000}});

  654. }

  655. 

  656. TEST(TestOprBlas, TransMatMul) {

  657.     run_trans_inp_test<float, float>();

  658. }

  659. 

  660. TEST(TestOprBlas, TransMatMul8x8x16) {

  661.     if (CompNode::load("xpux").device_type() != CompNode::DeviceType::CUDA) {

  662.         run_trans_inp_test<dt_int8, dt_int16>();

  663.     } else {

  664.         printf("testcase skipped on unsupported arch\n");

  665.     }

  666. }

  667. 

  668. TEST(TestOprBlas, TransMatMul8x8x32) {

  669.     if (CompNode::load("xpux").device_type() == CompNode::DeviceType::CUDA &&

  670.         !check_compute_capability(6, 1)) {

  671.         return;

  672.     }

  673.     run_trans_inp_test<dt_int8, dt_int32>();

  674. }

  675. 

  676. TEST(TestOprBlas, NonContigMatmul) {

  677.     using Checker = AutoOprChecker<2, 1>;

  678.     auto make_graph =

  679.             [](const Checker::SymInpArray& inputs) -> Checker::SymOutArray {

  680. 

  681.         using Ad = opr::Subtensor::AxisIndexer;

  682.         auto x = inputs[0],

  683.              xsub = opr::Subtensor::make(

  684.                      x, {Ad::make_interval(0, None, None, x.make_scalar(2))}),

  685.              y = inputs[1],

  686.              ysub = opr::Subtensor::make(

  687.                      y, {Ad::make_interval(1, None, None, x.make_scalar(3))});

  688.         return {opr::MatrixMul::make(xsub, ysub)};

  689.     };

  690.     auto fwd = [](Checker::NumOutArray& dest, Checker::NumInpArray inp) {

  691.         auto &&shp0 = inp[0]->shape(), &&shp1 = inp[1]->shape();

  692.         size_t m = (shp0.shape[0] + 1) / 2, k = shp0.shape[1],

  693.                n = (shp1.shape[1] + 2) / 3;

  694.         auto dptr = dest[0].resize({m, n}).ptr<float>();

  695.         memset(dptr, 0, sizeof(float) * m * n);

  696.         for (size_t i = 0; i < m; ++i) {

  697.             auto ptr_a = inp[0]->ptr<float>({i * 2}),

  698.                  ptr_c = dest[0].ptr<float>({i});

  699.             for (size_t kk = 0; kk < k; ++kk) {

  700.                 auto va = ptr_a[kk];

  701.                 auto ptr_b = inp[1]->ptr<float>({kk});

  702.                 for (size_t j = 0; j < n; ++j) {

  703.                     ptr_c[j] += va * ptr_b[j * 3];

  704.                 }

  705.             }

  706.         }

  707.     };

  708. 

  709.     Checker(make_graph, fwd)

  710.             .run({TensorShape{2, 1}, TensorShape{1, 3}})

  711.             .run({TensorShape{5, 2}, TensorShape{2, 6}})

  712.             .run({TensorShape{6, 3}, TensorShape{3, 8}});

  713. }

  714. 

  715. TEST(TestOprBlas, MatrixInverse) {

  716.     using Checker = AutoOprChecker<1, 1>;

  717.     auto make_graph =

  718.             [=](const Checker::SymInpArray& inputs) -> Checker::SymOutArray {

  719.         return {opr::MatrixInverse::make(inputs[0])};

  720.     };

  721.     auto fwd = [=](Checker::NumOutArray& dest, Checker::NumInpArray inp) {

  722.         auto opr =

  723.                 megdnn_naive_handle()->create_operator<megdnn::MatrixInverse>();

  724. 

  725.         auto wk_size =

  726.                 opr->get_workspace_in_bytes(inp[0]->layout(), inp[0]->layout());

  727.         std::unique_ptr<dt_byte[]> wk{new dt_byte[wk_size]};

  728.         opr->exec(inp[0]->as_megdnn(),

  729.                   dest[0].resize(inp[0]->shape()).as_megdnn(),

  730.                   {wk.get(), wk_size});

  731.     };

  732.     // ensure low condition number for generated matrices

  733.     auto input_coord = [](const Checker::NumInpArray& inp) {

  734.         auto shp = inp[0]->shape();

  735.         size_t n = shp[shp.ndim - 1];

  736.         size_t batch = 1;

  737.         for (size_t i = 0; i < shp.ndim - 2; ++i) {

  738.             batch *= shp[i];

  739.         }

  740.         std::vector<int> perm(n);

  741.         for (size_t i = 0; i < n; ++i) {

  742.             perm[i] = i;

  743.         }

  744.         auto ptr = inp[0]->ptr<float>();

  745.         for (size_t i = 0; i < batch; ++i, ptr += n * n) {

  746. #if __cplusplus >= 201703L

  747.             std::default_random_engine rng_engine;

  748.             std::shuffle(perm.begin(), perm.end(), rng_engine);

  749. #else

  750.             std::random_shuffle(perm.begin(), perm.end());

  751. #endif

  752.             for (size_t j = 0; j < n; ++j) {

  753.                 ptr[j * n + perm[j]] += 5;

  754.             }

  755.         }

  756.     };

  757. 

  758.     Checker{make_graph, fwd}

  759.             .set_input_coordinator(input_coord)

  760.             .run({TensorShape{5, 5}})

  761.             .run({TensorShape{2, 5, 5}})

  762.             .run({TensorShape{2, 6, 3, 3}});

  763. }

  764. 

  765. namespace {

  766. 

  767. void gen_svd_input(HostTensorND& dest) {

  768.     auto ptr = dest.ptr<float>();

  769.     auto dim = dest.layout().ndim;

  770.     size_t n = dest.layout().shape[dim - 2], m = dest.layout().shape[dim - 1];

  771.     size_t j = 0, k = 0;

  772.     float batch_off = 0;

  773.     float max_val = std::min(m, n) * std::min(m, n) + 0.99;

  774.     for (size_t i = 0, it = dest.layout().total_nr_elems(); i < it; ++i) {

  775.         if (i % (n * m) == 0) {

  776.             batch_off += 0.32;

  777.             j = k = 0;

  778.         }

  779.         if (!((i % (n * m)) % (m + 1)))

  780.             ptr[i] = (j++) + ((++k / 10.0));

  781.         else

  782.             ptr[i] = (j++);

  783.         ptr[i] += batch_off;

  784.         ptr[i] = std::fmod(ptr[i], max_val);

  785.     }

  786. }

  787. 

  788. template <int have_u, int have_s, int have_v>

  789. void run_svd_empty_grad_test() {

  790.     using Checker = AutoOprChecker<1, have_u + have_s + have_v>;

  791.     auto make_graph = [=](const typename Checker::SymInpArray& inputs) {

  792.         auto out = opr::SVD::make(inputs[0], opr::SVD::Param{false, true});

  793.         typename Checker::SymOutArray ret;

  794.         int idx = 0;

  795.         if (have_u) {

  796.             ret[idx++] = out[0];

  797.         }

  798.         if (have_s) {

  799.             ret[idx++] = out[1];

  800.         }

  801.         if (have_v) {

  802.             ret[idx++] = out[2];

  803.         }

  804.         return ret;

  805.     };

  806.     auto fwd = [=](typename Checker::NumOutArray& dest,

  807.                    typename Checker::NumInpArray inp) {

  808.         auto opr = megdnn_naive_handle()->create_operator<megdnn::SVDForward>();

  809.         opr->param().compute_uv = true;

  810.         TensorLayout ul, sl, vtl;

  811.         opr->deduce_layout(inp[0]->layout(), ul, sl, vtl);

  812.         HostTensorND tmp_u{dest[0].comp_node(), ul},

  813.                 tmp_s{dest[0].comp_node(), sl}, tmp_v{dest[0].comp_node(), vtl};

  814.         auto wk_size =

  815.                 opr->get_workspace_in_bytes(inp[0]->layout(), ul, sl, vtl);

  816.         auto wk = std::make_unique<dt_byte[]>(wk_size);

  817.         auto out0 = tmp_u.as_megdnn(), out1 = tmp_s.as_megdnn(),

  818.              out2 = tmp_v.as_megdnn();

  819.         int idx = 0;

  820.         if (have_u) {

  821.             out0 = dest[idx++].resize(ul).as_megdnn();

  822.         }

  823.         if (have_s) {

  824.             out1 = dest[idx++].resize(sl).as_megdnn();

  825.         }

  826.         if (have_v) {

  827.             out2 = dest[idx++].resize(vtl).as_megdnn();

  828.         }

  829.         opr->exec(inp[0]->as_megdnn(), out0, out1, out2, {wk.get(), wk_size});

  830.     };

  831.     Checker checker{make_graph, fwd};

  832.     checker.set_input_generator(0, gen_svd_input);

  833.     if (have_u) {

  834.         checker.set_output_allow_check(0, false);

  835.     }

  836.     if (have_v) {

  837.         checker.set_output_allow_check(have_u + have_s, false);

  838.     }

  839.     checker.run({TensorShape{3, 3}})

  840.             .run({TensorShape{2, 3, 3}})

  841.             .run({TensorShape{2, 4, 2}})

  842.             .run({TensorShape{3, 1, 2, 4}})

  843.             .run({TensorShape{2, 3, 2, 3}});

  844. }

  845. 

  846. }  // anonymous namespace

  847. 

  848. TEST(TestOprBlas, SingularValueDecomposition) {

  849.     using Checker = AutoOprChecker<1, 3>;

  850.     auto make_graph =

  851.             [=](const Checker::SymInpArray& inputs) -> Checker::SymOutArray {

  852.         auto out = opr::SVD::make(inputs[0], opr::SVD::Param{false, true});

  853.         return {out[0], out[1], out[2]};

  854.     };

  855.     auto fwd = [=](Checker::NumOutArray& dest, Checker::NumInpArray inp) {

  856.         auto opr = megdnn_naive_handle()->create_operator<megdnn::SVDForward>();

  857.         opr->param().compute_uv = true;

  858.         TensorLayout ul, sl, vtl;

  859.         opr->deduce_layout(inp[0]->layout(), ul, sl, vtl);

  860.         auto wk_size =

  861.                 opr->get_workspace_in_bytes(inp[0]->layout(), ul, sl, vtl);

  862.         auto wk = std::make_unique<dt_byte[]>(wk_size);

  863.         opr->exec(inp[0]->as_megdnn(), dest[0].resize(ul).as_megdnn(),

  864.                   dest[1].resize(sl).as_megdnn(),

  865.                   dest[2].resize(vtl).as_megdnn(), {wk.get(), wk_size});

  866.     };

  867.     Checker{make_graph, fwd}

  868.             .set_input_generator(0, gen_svd_input)

  869.             .set_output_allow_check(0, false)

  870.             .set_output_allow_check(2, false)

  871.             .run({TensorShape{3, 3}})

  872.             .run({TensorShape{2, 3, 3}})

  873.             .run({TensorShape{2, 4, 2}})

  874.             .run({TensorShape{3, 1, 2, 4}})

  875.             .run({TensorShape{2, 3, 2, 3}});

  876. }

  877. 

  878. TEST(TestOprBlas, SingularValueDecompositionZeroGrad) {

  879.     run_svd_empty_grad_test<0, 0, 1>();

  880.     run_svd_empty_grad_test<0, 1, 0>();

  881.     run_svd_empty_grad_test<0, 1, 1>();

  882.     run_svd_empty_grad_test<1, 0, 0>();

  883.     run_svd_empty_grad_test<1, 0, 1>();

  884.     run_svd_empty_grad_test<1, 1, 0>();

  885.     run_svd_empty_grad_test<1, 1, 1>();

  886. }

  887. 

  888. #if MGB_ENABLE_FASTRUN

  889. TEST(TestOprBlas, MatrixMulExePolicy) {

  890.     using Param = opr::MatrixMul::Param;

  891.     Param param;

  892.     using Policy = opr::MatrixMul::ExecutionPolicy;

  893.     using S = Policy::Strategy;

  894.     Policy policy;

  895.     policy.strategy = S::PROFILE;

  896. 

  897.     auto cn = CompNode::load("cpux");

  898. 

  899.     int nr_get = 0;

  900.     auto on_get = [&nr_get](const std::string&, const void*, size_t,

  901.                             const void*, size_t) { ++nr_get; };

  902.     PersistentCacheHook cache_hook{on_get};

  903. 

  904.     auto graph = ComputingGraph::make();

  905.     HostTensorGenerator<> gen;

  906. 

  907.     auto mkvar = [&](const char* name, const TensorShape& shp) {

  908.         return opr::Host2DeviceCopy::make(*graph, gen(shp), cn).rename(name);

  909.     };

  910. 

  911.     auto a = mkvar("a", {20, 50});

  912.     auto b = mkvar("b", {50, 40});

  913.     auto matmul = opr::MatrixMul::make(a, b, param, policy, {});

  914. 

  915.     HostTensorND host_y;

  916.     graph->options().no_profiling_on_shape_change = true;

  917.     auto func = graph->compile({make_callback_copy(matmul, host_y)});

  918.     func->execute();

  919.     ASSERT_GT(nr_get, 0);

  920.     int nr = nr_get;

  921.     graph->options().no_profiling_on_shape_change = false;

  922.     func = graph->compile({make_callback_copy(matmul, host_y)});

  923.     func->execute();

  924.     ASSERT_GT(nr_get, nr);

  925. }

  926. #endif

  927. 

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

  929. //