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: ae3123b315
Branches Tags
${ item.name }
Create branch ${ searchTerm }
from 'ae3123b315'
${ noResults }
MegEngine/src/plugin/impl/opr_footprint.cpp

899 lines
35 kB
Raw Blame History 

  1. /**

  2.  * \file src/plugin/impl/opr_footprint.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/plugin/opr_footprint.h"

  13. #include "megbrain/opr/basic_arith.h"

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

  15. #include "megbrain/opr/dnn/convolution.h"

  16. #include "megbrain/opr/dnn/images2neibs.h"

  17. #include "megbrain/opr/dnn/local.h"

  18. #include "megbrain/opr/dnn/lrn.h"

  19. #include "megbrain/opr/dnn/pooling.h"

  20. #include "megbrain/opr/dnn/adaptive_pooling.h"

  21. #include "megbrain/opr/dnn/roi_pooling.h"

  22. #include "megbrain/opr/dnn/roi_align.h"

  23. #include "megbrain/opr/imgproc.h"

  24. #include "megbrain/opr/standalone/nms_opr.h"

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

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

  27. #include "megbrain/opr/rand.h"

  28. #include "megbrain/opr/dnn/batch_norm.h"

  29. #include "megbrain/opr/misc.h"

  30. #include "megbrain/opr/indexing.h"

  31. #include "megbrain/opr/internal/indexing_helper.h"

  32. #include "megbrain/opr/nn_int.h"

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

  34. #if MGB_ENABLE_JSON

  35. #include "megdnn/opr_param_json.h"

  36. #endif

  37. 

  38. #include "megbrain/utils/hash_ct.h"

  39. #include "midout.h"

  40. 

  41. MIDOUT_DECL(megbrain_opr_footprint)

  42. #define MIDOUT_B(...) \

  43.     MIDOUT_BEGIN(megbrain_opr_footprint, __VA_ARGS__) {

  44. #define MIDOUT_E \

  45.     }            \

  46.     MIDOUT_END();

  47. 

  48. using namespace mgb;

  49. 

  50. namespace {

  51. 

  52. template <class T>

  53. uint64_t opr_footprint_func(cg::OperatorNodeBase* opr);

  54. 

  55. // Elemwise

  56. template <>

  57. uint64_t opr_footprint_func<opr::Elemwise>(cg::OperatorNodeBase* opr) {

  58.     return opr->output()[0]->shape().total_nr_elems() *

  59.            (std::max<size_t>(opr->input().size(), 2) - 1);

  60. }

  61. 

  62. // AddUpdate

  63. template <>

  64. uint64_t opr_footprint_func<opr::AddUpdate>(cg::OperatorNodeBase* opr) {

  65.     mgb_assert(opr->input().size() == 2,

  66.                "AddUpdate opr should have two inputs");

  67.     auto&& out_shape = opr->output()[0]->shape();

  68.     return out_shape.total_nr_elems() * 3;

  69. }

  70. 

  71. template <class Conv>

  72. uint64_t eval_conv_computation(const TensorShape& src_shape,

  73.                                const TensorShape& filter_shape,

  74.                                const TensorShape& dst_shape,

  75.                                cg::OperatorNodeBase* opr) {

  76.     using Param = opr::ConvolutionForward::Param;

  77.     auto&& param = opr->cast_final_safe<Conv>().param();

  78. 

  79.     if (param.format == Param::Format::NHWCD4) {

  80.         size_t fh, fw;

  81.         size_t group = 1;

  82.         if (param.sparse == Param::Sparse::DENSE) {

  83.             fh = filter_shape[1];

  84.             fw = filter_shape[2];

  85.             group = 1;

  86.         } else {

  87.             // chanwise conv

  88.             mgb_assert(param.sparse == Param::Sparse::GROUP);

  89.             fh = filter_shape[2];

  90.             fw = filter_shape[3];

  91.             group = filter_shape[0];

  92. 

  93.             if (filter_shape.ndim == 5) {

  94.                 group *= 4;

  95.             }

  96.         }

  97.         return dst_shape.total_nr_elems() * fh * fw *

  98.             src_shape[2] * 4 / group * 2;

  99.     }

  100.     auto eval_conv_computation_nchwx = [&param, &src_shape, &filter_shape,

  101.                                         &dst_shape]() -> uint64_t {

  102.         size_t fh, fw;

  103.         bool hybird_nchwx = false;

  104.         size_t group = 1;

  105.         if (param.sparse == Param::Sparse::DENSE) {

  106.             //! if nchwxx mode src is nchw output is nchwxx

  107.             if (dst_shape.ndim == 5 && src_shape.ndim == 4) {

  108.                 fh = filter_shape[1];

  109.                 fw = filter_shape[2];

  110.                 hybird_nchwx = true;

  111.             } else {

  112.                 fh = filter_shape[2];

  113.                 fw = filter_shape[3];

  114.             }

  115.             group = 1;

  116.         } else {

  117.             mgb_assert(param.sparse == Param::Sparse::GROUP);

  118.             fh = filter_shape[3];

  119.             fw = filter_shape[4];

  120.             group = filter_shape[0];

  121.         }

  122.         if (param.format == Param::Format::NCHW88) {

  123.             //! if channel wise weight layout is {group/8, FH, FW, 1, 1, 8}

  124.             if (filter_shape[1] == 1 && filter_shape[2] == 1) {

  125.                 group *= 8;

  126.             }

  127.             size_t computation = dst_shape.total_nr_elems() * fh * fw *

  128.                                  src_shape[1] / group * 2;

  129.             return hybird_nchwx ? computation : computation * 8;

  130.         }

  131.         if (param.format == Param::Format::NCHW44 ||

  132.             param.format == Param::Format::NCHW44_DOT) {

  133.             //! if channel wise weight layout is {group/4, FH, FW, 1, 1, 4}

  134.             if (filter_shape[1] == 1 && filter_shape[2] == 1) {

  135.                 group *= 4;

  136.             }

  137.             size_t computation = dst_shape.total_nr_elems() * fh * fw *

  138.                                  src_shape[1] / group * 2;

  139.             return hybird_nchwx ? computation : computation * 4;

  140.         }

  141.         if (param.format == Param::Format::NCHW32 ||

  142.             param.format == Param::Format::NCHW32_NCHW4) {

  143.             return dst_shape.total_nr_elems() * fh * fw * src_shape[1] * 32 /

  144.                    group * 2;

  145.         }

  146.         mgb_assert(param.format == Param::Format::NCHW4 ||

  147.                            param.format == Param::Format::NCHW4_NCHW ||

  148.                            param.format == Param::Format::NCHW4_NCHW32,

  149.                    "format should be NCHW4/NCHW4_NCHW/NCHW4_NCHW32");

  150.         return dst_shape.total_nr_elems() * fh * fw * src_shape[1] * 4 / group *

  151.                2;

  152.     };

  153.     auto eval_conv_computation_chwn4 = [&param, &src_shape, &filter_shape,

  154.                                         &dst_shape]() -> uint64_t {

  155.         size_t fh, fw;

  156.         size_t group = 1;

  157.         if (param.sparse == Param::Sparse::DENSE) {

  158.             fh = filter_shape[1];

  159.             fw = filter_shape[2];

  160.             group = 1;

  161.         } else {

  162.             mgb_assert(param.sparse == Param::Sparse::GROUP);

  163.             fh = filter_shape[2];

  164.             fw = filter_shape[3];

  165.             group = filter_shape[0];

  166.         }

  167.         return dst_shape.total_nr_elems() * fh * fw * src_shape[0] * 4 / group *

  168.                2;

  169.     };

  170.     if (param.format == Param::Format::NCHW4 ||

  171.         param.format == Param::Format::NCHW4_NCHW ||

  172.         param.format == Param::Format::NCHW4_NCHW32 ||

  173.         param.format == Param::Format::NCHW88 ||

  174.         param.format == Param::Format::NCHW44 ||

  175.         param.format == Param::Format::NCHW44_DOT ||

  176.         param.format == Param::Format::NCHW32 ||

  177.         param.format == Param::Format::NCHW32_NCHW4) {

  178.         return eval_conv_computation_nchwx();

  179.     }

  180.     if (param.format == Param::Format::CHWN4) {

  181.         return eval_conv_computation_chwn4();

  182.     }

  183.     size_t cpos;

  184.     size_t spatial_start;

  185.     size_t group = 1;

  186.     switch (param.format) {

  187.         case Param::Format::NCHW:

  188.             cpos = 1;

  189.             spatial_start = 2;

  190.             break;

  191.         case Param::Format::NHWC:

  192.             cpos = 3;

  193.             spatial_start = 1;

  194.             break;

  195.         default:

  196.             mgb_assert(false, "Unknown CONV Param::Format type");

  197.     }

  198.     switch (param.sparse) {

  199.         case Param::Sparse::DENSE:

  200.             mgb_assert(filter_shape.ndim == 4 || filter_shape.ndim == 6,

  201.                        "DENSE conv filter shape dimension should be "

  202.                        "4/6(winograd mk4)");

  203.             break;

  204.         case Param::Sparse::GROUP:

  205.             mgb_assert(filter_shape.ndim == 5 || filter_shape.ndim == 7,

  206.                        "GROUP conv filter shape dimension should be "

  207.                        "5/7(winograd mk4)");

  208.             spatial_start++;

  209.             group = filter_shape[0];

  210.             break;

  211.         default:

  212.             mgb_assert(false, "Unkown CONV Param::Sparse type");

  213.     }

  214. 

  215.     uint64_t fh = static_cast<uint64_t>(filter_shape[spatial_start]);

  216.     uint64_t fw = static_cast<uint64_t>(filter_shape[spatial_start + 1]);

  217.     

  218.     // mul and add are counted as 2 operations

  219.     

  220.     return dst_shape.total_nr_elems() * fh * fw *

  221.            static_cast<uint64_t>(src_shape[cpos]) / group * 2;

  222. }

  223. 

  224. // ConvolutionForward

  225. template <>

  226. uint64_t opr_footprint_func<opr::ConvolutionForward>(

  227.         cg::OperatorNodeBase* opr) {

  228.     mgb_assert(opr->input().size() == 2,

  229.                "ConvolutionFwd opr should have two inputs");

  230.     auto&& out_shape = opr->output()[0]->shape();

  231.     auto&& src_shape = opr->input()[0]->shape();

  232.     auto&& filter_shape = opr->input()[1]->shape();

  233.     return eval_conv_computation<opr::ConvolutionForward>(

  234.             src_shape, filter_shape, out_shape, opr);

  235. }

  236. template <>

  237. uint64_t opr_footprint_func<opr::ConvBiasForward>(

  238.         cg::OperatorNodeBase* opr) {

  239.     mgb_assert(opr->input().size() == 2 || opr->input().size() == 3 ||

  240.                        opr->input().size() == 4,

  241.                "ConvBiasForward opr should have two/three/four inputs");

  242.     auto&& out_shape = opr->output()[0]->shape();

  243.     auto&& src_shape = opr->input()[0]->shape();

  244.     auto&& filter_shape = opr->input()[1]->shape();

  245.     uint64_t res = eval_conv_computation<opr::ConvBiasForward>(

  246.             src_shape, filter_shape, out_shape, opr);

  247.     if (opr->input().size() == 3) {

  248.         res += out_shape.total_nr_elems();

  249.     }

  250.     return res;

  251. }

  252. 

  253. // ConvolutionBackwardData

  254. template <>

  255. uint64_t opr_footprint_func<opr::ConvolutionBackwardData>(

  256.         cg::OperatorNodeBase* opr) {

  257.     mgb_assert(opr->input().size() == 2 || opr->input().size() == 3,

  258.                "ConvolutionBackwardData opr should have two or three inputs");

  259.     auto&& filter_shape = opr->input()[0]->shape();

  260.     auto&& diff_shape = opr->input()[1]->shape();

  261.     auto&& grad_shape = opr->output()[0]->shape();

  262.     return eval_conv_computation<opr::ConvolutionBackwardData>(

  263.             grad_shape, filter_shape, diff_shape, opr);

  264. }

  265. 

  266. // ConvolutionBackwardFilter

  267. template <>

  268. uint64_t opr_footprint_func<opr::ConvolutionBackwardFilter>(

  269.         cg::OperatorNodeBase* opr) {

  270.     mgb_assert(opr->input().size() == 3,

  271.                "ConvolutionBackwardData opr should have three inputs");

  272.     auto&& filter_shape = opr->input()[2]->shape();

  273.     auto&& diff_shape = opr->input()[1]->shape();

  274.     auto&& src_shape = opr->input()[0]->shape();

  275.     return eval_conv_computation<opr::ConvolutionBackwardFilter>(

  276.             src_shape, filter_shape, diff_shape, opr);

  277. }

  278. 

  279. // MatrixMul

  280. template <>

  281. uint64_t opr_footprint_func<opr::MatrixMul>(cg::OperatorNodeBase* opr) {

  282.     auto&& mopr = opr->cast_final_safe<opr::MatrixMul>();

  283.     auto &&i0 = opr->input(0)->shape(), &&i1 = opr->input(1)->shape();

  284.     mgb_assert(i0.ndim == 2 && i1.ndim == 2);

  285.     auto m = i0[0], k0 = i0[1], k1 = i1[0], n = i1[1];

  286.     if (mopr.param().transposeA) {

  287.         std::swap(m, k0);

  288.     }

  289.     if (mopr.param().transposeB) {

  290.         std::swap(k1, n);

  291.     }

  292.     mgb_assert(k0 == k1);

  293.     // mul and add are counted as 2 operations

  294.     return m * k0 * n * 2;

  295. }

  296. 

  297. template <>

  298. uint64_t opr_footprint_func<opr::LocalShareForward>(cg::OperatorNodeBase* opr) {

  299.     mgb_assert(opr->input().size() == 2,

  300.                "LocalShare opr should have two inputs");

  301.     auto&& out_shape = opr->output()[0]->shape();

  302.     auto&& src_shape = opr->input()[0]->shape();

  303.     auto&& filter_shape = opr->input()[1]->shape();

  304.     using Param = opr::LocalShareForward::Param;

  305.     auto&& param = opr->cast_final_safe<opr::LocalShareForward>().param();

  306.     mgb_assert(param.format == Param::Format::NCHW);

  307.     size_t groups = 1;

  308.     size_t kern_spatial_pos = 3;

  309.     if (param.sparse == Param::Sparse::GROUP) {

  310.         groups = filter_shape[0];

  311.         kern_spatial_pos = 4;

  312.     }

  313.     size_t fh = filter_shape[kern_spatial_pos],

  314.            fw = filter_shape[kern_spatial_pos + 1];

  315.     return out_shape.total_nr_elems() * fh * fw * src_shape[1] * 2 / groups;

  316. }

  317. 

  318. template <>

  319. uint64_t opr_footprint_func<opr::LocalShareBackwardData>(cg::OperatorNodeBase* opr) {

  320.     mgb_assert(opr->input().size() == 3,

  321.                "LocalShareBackwardData opr should have three inputs");

  322.     auto&& filter_shape = opr->input()[0]->shape();

  323.     auto&& diff_shape = opr->input()[1]->shape();

  324.     auto&& grad_shape = opr->output()[0]->shape();

  325.     using Param = opr::LocalShareForward::Param;

  326.     auto&& param = opr->cast_final_safe<opr::LocalShareBackwardData>().param();

  327.     mgb_assert(param.format == Param::Format::NCHW);

  328.     size_t groups = 1;

  329.     size_t kern_spatial_pos = 3;

  330.     if (param.sparse == Param::Sparse::GROUP) {

  331.         groups = filter_shape[0];

  332.         kern_spatial_pos = 4;

  333.     }

  334.     size_t fh = filter_shape[kern_spatial_pos],

  335.            fw = filter_shape[kern_spatial_pos + 1];

  336.     return diff_shape.total_nr_elems() * fh * fw * grad_shape[1] * 2 / groups;

  337. }

  338. 

  339. template <>

  340. uint64_t opr_footprint_func<opr::LocalShareBackwardFilter>(cg::OperatorNodeBase* opr) {

  341.     mgb_assert(opr->input().size() == 3,

  342.                "LocalShareBackwardFilter opr should have three inputs");

  343.     auto&& src_shape = opr->input()[0]->shape();

  344.     auto&& diff_shape = opr->input()[1]->shape();

  345.     auto&& grad_shape = opr->output()[0]->shape();

  346.     using Param = opr::LocalShareForward::Param;

  347.     auto&& param = opr->cast_final_safe<opr::LocalShareBackwardFilter>().param();

  348.     mgb_assert(param.format == Param::Format::NCHW);

  349.     size_t groups = 1;

  350.     size_t kern_spatial_pos = 3;

  351.     if (param.sparse == Param::Sparse::GROUP) {

  352.         groups = grad_shape[0];

  353.         kern_spatial_pos = 4;

  354.     }

  355.     size_t fh = grad_shape[kern_spatial_pos],

  356.            fw = grad_shape[kern_spatial_pos + 1];

  357.     return diff_shape.total_nr_elems() * fh * fw * src_shape[1] * 2 / groups;

  358. }

  359. 

  360. template <>

  361. uint64_t opr_footprint_func<opr::DeformableConvForward>(

  362.         cg::OperatorNodeBase* opr) {

  363.     mgb_assert(opr->input().size() == 4,

  364.                "DeformableConvForward opr should have four inputs");

  365.     auto&& out_shape = opr->output()[0]->shape();

  366.     auto&& filter_shape = opr->input()[1]->shape();

  367.     using Param = opr::DeformableConvForward::Param;

  368.     auto&& param = opr->cast_final_safe<opr::DeformableConvForward>().param();

  369.     size_t fh, fw, icpg;

  370.     mgb_assert(param.format == Param::Format::NCHW);

  371.     if (param.sparse == Param::Sparse::GROUP) {

  372.         icpg = filter_shape[2];

  373.         fh = filter_shape[3], fw = filter_shape[4];

  374.     } else {

  375.         icpg = filter_shape[1];

  376.         fh = filter_shape[2], fw = filter_shape[3];

  377.     }

  378.     //! conv(1 mul), mask(1, mul), accumulate(1 add)

  379.     return out_shape.total_nr_elems() * fh * fw * icpg * 3;

  380. }

  381. 

  382. template <>

  383. uint64_t opr_footprint_func<opr::DeformableConvBackwardFilter>(

  384.         cg::OperatorNodeBase* opr) {

  385.     mgb_assert(opr->input().size() == 5,

  386.                "DeformableConvBackwardFilter opr should have four inputs");

  387.     auto&& out_shape = opr->output()[0]->shape();

  388.     auto&& filter_shape = opr->input()[1]->shape();

  389.     using Param = opr::DeformableConvBackwardFilter::Param;

  390.     auto&& param = opr->cast_final_safe<opr::Convolution>().param();

  391.     size_t fh, fw, icpg;

  392.     mgb_assert(param.format == Param::Format::NCHW);

  393.     if (param.sparse == Param::Sparse::GROUP) {

  394.         icpg = filter_shape[2];

  395.         fh = filter_shape[3], fw = filter_shape[4];

  396.     } else {

  397.         icpg = filter_shape[1];

  398.         fh = filter_shape[2], fw = filter_shape[3];

  399.     }

  400.     //! deconv(1 mul), mask(1 mul), accumulate(1 add), bilinear(4 add, 4mul,

  401.     //! skip)

  402.     return out_shape.total_nr_elems() * fh * fw * icpg * 3;

  403. }

  404. 

  405. template <>

  406. uint64_t opr_footprint_func<opr::DeformableConvBackwardData>(

  407.         cg::OperatorNodeBase* opr) {

  408.     mgb_assert(opr->input().size() == 5,

  409.                "DeformableConvBackwardData opr should have four inputs");

  410.     auto&& out_shape = opr->output()[0]->shape();

  411.     auto&& filter_shape = opr->input()[1]->shape();

  412.     using Param = opr::DeformableConvForward::Param;

  413.     auto&& param = opr->cast_final_safe<opr::Convolution>().param();

  414.     size_t fh, fw, icpg;

  415.     mgb_assert(param.format == Param::Format::NCHW);

  416.     if (param.sparse == Param::Sparse::GROUP) {

  417.         icpg = filter_shape[2];

  418.         fh = filter_shape[3], fw = filter_shape[4];

  419.     } else {

  420.         icpg = filter_shape[1];

  421.         fh = filter_shape[2], fw = filter_shape[3];

  422.     }

  423.     //! deconv(1 mul), mask(1 mul), accumulate(1 add), grad_weight(1 mul, skip),

  424.     //! grad_coord(4mul, 4 add)

  425.     return out_shape.total_nr_elems() * fh * fw * icpg * 12;

  426. }

  427. 

  428. template <>

  429. uint64_t opr_footprint_func<opr::BatchConvBiasForward>(

  430.         cg::OperatorNodeBase* opr) {

  431.     mgb_assert(opr->input().size() == 2 || opr->input().size() == 3 ||

  432.                        opr->input().size() == 4,

  433.                "BatchConvBias opr should have two/three/four inputs");

  434.     auto&& out_shape = opr->output()[0]->shape();

  435.     auto&& src_shape = opr->input()[0]->shape();

  436.     auto&& filter_shape = opr->input()[1]->shape();

  437.     using Param = opr::BatchConvBiasForward::Param;

  438.     auto&& param = opr->cast_final_safe<opr::BatchConvBiasForward>().param();

  439.     size_t packed_channels = 1;

  440.     size_t kern_spatial_pos = 3;

  441.     if (param.format == Param::Format::NCHW4) {

  442.         packed_channels = 4;

  443.     }

  444.     size_t fh = filter_shape[kern_spatial_pos],

  445.            fw = filter_shape[kern_spatial_pos + 1];

  446.     return out_shape.total_nr_elems() * fh * fw * src_shape[1] *

  447.            packed_channels * 2;

  448. }

  449. 

  450. // Pooling

  451. template <>

  452. uint64_t opr_footprint_func<opr::PoolingForward>(cg::OperatorNodeBase* opr) {

  453.     auto&& param = opr->cast_final_safe<opr::PoolingForward>().param();

  454.     auto area = param.window_h * param.window_w;

  455.     return opr->output(0)->shape().total_nr_elems() * area;

  456. }

  457. 

  458. // Concat

  459. template <>

  460. uint64_t opr_footprint_func<opr::Concat>(cg::OperatorNodeBase* opr) {

  461.     auto&& out_shape = opr->output()[0]->shape();

  462.     return out_shape.total_nr_elems();

  463. }

  464. 

  465. // Dimshuffle

  466. template <>

  467. uint64_t opr_footprint_func<opr::Dimshuffle>(cg::OperatorNodeBase* opr) {

  468.     auto&& out = opr->output()[0];

  469.     return out->shape().total_nr_elems();

  470. }

  471. 

  472. // Reduce

  473. template <>

  474. uint64_t opr_footprint_func<opr::Reduce>(cg::OperatorNodeBase* opr) {

  475.     return opr->input()[0]->shape().total_nr_elems();

  476. }

  477. 

  478. // Host2DeviceCopy

  479. template <>

  480. uint64_t opr_footprint_func<opr::Host2DeviceCopy>(cg::OperatorNodeBase* opr) {

  481.     auto&& out_shape = opr->output()[0]->shape();

  482.     return out_shape.total_nr_elems();

  483. }

  484. 

  485. /******************* Registe Param Json Functions *************************/

  486. #if MGB_ENABLE_JSON

  487. template <class T>

  488. std::shared_ptr<json::Value> opr_param_json_func(cg::OperatorNodeBase* opr);

  489. 

  490. #define REGISTE_PARAM_JSON_FUNC(cls)                            \

  491.     template <>                                                 \

  492.     std::shared_ptr<json::Value> opr_param_json_func<opr::cls>( \

  493.             cg::OperatorNodeBase * opr) {                       \

  494.         return opr::opr_param_to_json(                          \

  495.                 opr->cast_final_safe<opr::cls>().param());      \

  496.     }

  497. 

  498. REGISTE_PARAM_JSON_FUNC(Elemwise)

  499. REGISTE_PARAM_JSON_FUNC(ConvolutionForward)

  500. REGISTE_PARAM_JSON_FUNC(Convolution3D)

  501. REGISTE_PARAM_JSON_FUNC(ConvBiasForward)

  502. REGISTE_PARAM_JSON_FUNC(ConvolutionBackwardData)

  503. REGISTE_PARAM_JSON_FUNC(Convolution3DBackwardData)

  504. REGISTE_PARAM_JSON_FUNC(ConvolutionBackwardFilter)

  505. REGISTE_PARAM_JSON_FUNC(MatrixMul)

  506. REGISTE_PARAM_JSON_FUNC(BatchedMatrixMul)

  507. REGISTE_PARAM_JSON_FUNC(Dot)

  508. REGISTE_PARAM_JSON_FUNC(MatrixInverse)

  509. REGISTE_PARAM_JSON_FUNC(PoolingForward)

  510. REGISTE_PARAM_JSON_FUNC(SVD)

  511. REGISTE_PARAM_JSON_FUNC(MaskConvolution)

  512. REGISTE_PARAM_JSON_FUNC(Images2Neibs)

  513. REGISTE_PARAM_JSON_FUNC(Local)

  514. REGISTE_PARAM_JSON_FUNC(GroupLocal)

  515. REGISTE_PARAM_JSON_FUNC(LRN)

  516. REGISTE_PARAM_JSON_FUNC(Concat)

  517. REGISTE_PARAM_JSON_FUNC(Reduce)

  518. REGISTE_PARAM_JSON_FUNC(LocalShareForward)

  519. REGISTE_PARAM_JSON_FUNC(LocalShareBackwardData)

  520. REGISTE_PARAM_JSON_FUNC(LocalShareBackwardFilter)

  521. REGISTE_PARAM_JSON_FUNC(DeformableConvForward)

  522. REGISTE_PARAM_JSON_FUNC(DeformableConvBackwardFilter)

  523. REGISTE_PARAM_JSON_FUNC(DeformableConvBackwardData)

  524. REGISTE_PARAM_JSON_FUNC(DeformablePSROIPoolingForward)

  525. REGISTE_PARAM_JSON_FUNC(BatchConvBiasForward)

  526. REGISTE_PARAM_JSON_FUNC(BatchNormForward)

  527. REGISTE_PARAM_JSON_FUNC(ElemwiseMultiType)

  528. REGISTE_PARAM_JSON_FUNC(Argsort)

  529. REGISTE_PARAM_JSON_FUNC(Argmax)

  530. REGISTE_PARAM_JSON_FUNC(Argmin)

  531. REGISTE_PARAM_JSON_FUNC(AdaptivePooling)

  532. REGISTE_PARAM_JSON_FUNC(ROIPooling)

  533. REGISTE_PARAM_JSON_FUNC(ROIAlign)

  534. REGISTE_PARAM_JSON_FUNC(WarpPerspective)

  535. REGISTE_PARAM_JSON_FUNC(WarpAffine)

  536. REGISTE_PARAM_JSON_FUNC(Remap)

  537. REGISTE_PARAM_JSON_FUNC(Resize)

  538. REGISTE_PARAM_JSON_FUNC(IndexingOneHot)

  539. REGISTE_PARAM_JSON_FUNC(IndexingSetOneHot)

  540. REGISTE_PARAM_JSON_FUNC(TopK)

  541. REGISTE_PARAM_JSON_FUNC(UniformRNG)

  542. REGISTE_PARAM_JSON_FUNC(GaussianRNG)

  543. REGISTE_PARAM_JSON_FUNC(Linspace)

  544. REGISTE_PARAM_JSON_FUNC(Eye)

  545. REGISTE_PARAM_JSON_FUNC(CvtColor)

  546. 

  547. 

  548. template <>

  549. std::shared_ptr<json::Value> opr_param_json_func<opr::Dimshuffle>(

  550.     cg::OperatorNodeBase * opr) {

  551.         auto param = opr->cast_final_safe<opr::Dimshuffle>().param();

  552. 

  553.         auto pattern = json::Array::make();

  554.         for (size_t i = 0; i < param.pattern_len; i++)

  555.             pattern->add(json::NumberInt::make(param.pattern[i]));

  556. 

  557.         return json::Object::make({

  558.             {"ndim", json::NumberInt::make(param.ndim)},

  559.             {"pattern", pattern},

  560.         });

  561.     }

  562. 

  563. template <>

  564. std::shared_ptr<json::Value> opr_param_json_func<opr::AxisAddRemove>(

  565.     cg::OperatorNodeBase * opr) {

  566.         auto param = opr->cast_final_safe<opr::AxisAddRemove>().param();

  567. 

  568.         auto desc = json::Array::make();

  569.         for (size_t i = 0; i < param.nr_desc; i++) {

  570.             auto axisdesc = param.desc[i];

  571.             desc->add(

  572.                 json::Object::make({

  573.                     {"method", json::NumberInt::make(

  574.                         static_cast<int32_t>(axisdesc.method))},

  575.                     {"axisnum", json::NumberInt::make(axisdesc.axis.get_raw())},

  576.                 }));

  577.         }

  578. 

  579.         return json::Object::make({

  580.             {"nr_desc", json::NumberInt::make(param.nr_desc)},

  581.             {"desc", desc},

  582.         });

  583.     }

  584. 

  585. std::shared_ptr<json::Value> indexing_param_to_json(

  586.         const std::vector<opr::indexing::AxisIndexer>& indices) {

  587.     auto desc = json::Array::make();

  588.     for (auto& index : indices) {

  589.         desc->add(json::Object::make({

  590.                 {"axis", json::NumberInt::make(index.axis.get_raw())},

  591.                 {"begin",

  592.                     json::NumberInt::make(index.begin.node() != nullptr)},

  593.                 {"end", json::NumberInt::make(index.end.node() != nullptr)},

  594.                 {"step",

  595.                     json::NumberInt::make(index.step.node() != nullptr)},

  596.                 {"idx", json::NumberInt::make(index.idx.node() != nullptr)},

  597.         }));

  598.     }

  599.     return desc;

  600. }

  601. 

  602. #define REGISTE_INDEXING_PARAM_JSON_FUNC(cls)                         \

  603.     template <>                                                       \

  604.     std::shared_ptr<json::Value> opr_param_json_func<opr::cls>(       \

  605.             cg::OperatorNodeBase * opr) {                             \

  606.         auto indices = opr->cast_final_safe<opr::cls>().index_desc(); \

  607.         return indexing_param_to_json(indices);                  \

  608.     }

  609. 

  610. REGISTE_INDEXING_PARAM_JSON_FUNC(Subtensor);

  611. REGISTE_INDEXING_PARAM_JSON_FUNC(SetSubtensor);

  612. REGISTE_INDEXING_PARAM_JSON_FUNC(IncrSubtensor);

  613. REGISTE_INDEXING_PARAM_JSON_FUNC(IndexingMultiAxisVec);

  614. REGISTE_INDEXING_PARAM_JSON_FUNC(IndexingSetMultiAxisVec);

  615. REGISTE_INDEXING_PARAM_JSON_FUNC(IndexingIncrMultiAxisVec);

  616. REGISTE_INDEXING_PARAM_JSON_FUNC(MeshIndexing);

  617. REGISTE_INDEXING_PARAM_JSON_FUNC(IncrMeshIndexing);

  618. REGISTE_INDEXING_PARAM_JSON_FUNC(SetMeshIndexing);

  619. REGISTE_INDEXING_PARAM_JSON_FUNC(BatchedMeshIndexing);

  620. REGISTE_INDEXING_PARAM_JSON_FUNC(BatchedIncrMeshIndexing);

  621. REGISTE_INDEXING_PARAM_JSON_FUNC(BatchedSetMeshIndexing);

  622. 

  623. template <>

  624. std::shared_ptr<json::Value> opr_param_json_func<opr::Reshape>(

  625.     cg::OperatorNodeBase * opr) {

  626.         auto desc = json::Array::make();

  627.         auto axis_param = opr->cast_final_safe<opr::Reshape>().param();

  628.         if (axis_param.axis != axis_param.MAX_NDIM){

  629.             return json::Object::make({

  630.                 {"axis", json::NumberInt::make(axis_param.axis)},

  631.             });

  632.         } else {

  633.             return json::Object::make();

  634.         }

  635.     }

  636. 

  637. template <>

  638. std::shared_ptr<json::Value> opr_param_json_func<opr::GetVarShape>(

  639.     cg::OperatorNodeBase * opr) {

  640.         auto desc = json::Array::make();

  641.         auto axis_param = opr->cast_final_safe<opr::GetVarShape>().param();

  642.         if (axis_param.axis != axis_param.MAX_NDIM){

  643.             return json::Object::make({

  644.                 {"axis", json::NumberInt::make(axis_param.axis)},

  645.             });

  646.         } else {

  647.             return json::Object::make();

  648.         }

  649.     }

  650. 

  651. template <>

  652. std::shared_ptr<json::Value> opr_param_json_func<opr::standalone::NMSKeep>(

  653.     cg::OperatorNodeBase * opr) {

  654.         auto nms_param = opr->cast_final_safe<opr::standalone::NMSKeep>().param();

  655.         return json::Object::make({

  656.                 {"iou_thresh", json::Number::make(nms_param.iou_thresh)},

  657.                 {"max_output", json::Number::make(nms_param.max_output)},

  658.             });

  659.     }

  660.     

  661. 

  662. #endif // MGB_ENABLE_JSON

  663. 

  664. }  // namespace

  665. 

  666. template <class OprType>

  667. void OprFootprint::add_single_comp_footprint() {

  668.     MIDOUT_B(OprType,

  669.              midout_iv(MGB_HASH_STR("OprFootprint::add_single_comp_footprint")))

  670.     auto&& record = m_type2comp_footprint.emplace(OprType::typeinfo(),

  671.                                                   opr_footprint_func<OprType>);

  672.     mgb_assert(record.second, "duplicate opr typeinfo");

  673.     MIDOUT_E

  674. }

  675. 

  676. #if MGB_ENABLE_JSON

  677. template <class OprType>

  678. void OprFootprint::add_single_param_json() {

  679.     auto&& record = m_type2param_json.emplace(OprType::typeinfo(),

  680.                                               opr_param_json_func<OprType>);

  681.     mgb_assert(record.second, "duplicate opr typeinfo");

  682. }

  683. #endif

  684. 

  685. void OprFootprint::init_all_footprints() {

  686.     add_single_comp_footprint<opr::Elemwise>();

  687.     add_single_comp_footprint<opr::AddUpdate>();

  688.     add_single_comp_footprint<opr::ConvolutionForward>();

  689.     add_single_comp_footprint<opr::ConvBiasForward>();

  690.     add_single_comp_footprint<opr::ConvolutionBackwardData>();

  691.     add_single_comp_footprint<opr::ConvolutionBackwardFilter>();

  692.     add_single_comp_footprint<opr::MatrixMul>();

  693.     add_single_comp_footprint<opr::PoolingForward>();

  694.     add_single_comp_footprint<opr::Concat>();

  695.     add_single_comp_footprint<opr::Dimshuffle>();

  696.     add_single_comp_footprint<opr::Reduce>();

  697.     add_single_comp_footprint<opr::Host2DeviceCopy>();

  698.     add_single_comp_footprint<opr::LocalShareForward>();

  699.     add_single_comp_footprint<opr::LocalShareBackwardData>();

  700.     add_single_comp_footprint<opr::LocalShareBackwardFilter>();

  701.     add_single_comp_footprint<opr::DeformableConvForward>();

  702.     add_single_comp_footprint<opr::DeformableConvBackwardFilter>();

  703.     add_single_comp_footprint<opr::DeformableConvBackwardData>();

  704.     add_single_comp_footprint<opr::BatchConvBiasForward>();

  705. 

  706. #if MGB_ENABLE_JSON

  707.     add_single_param_json<opr::Elemwise>();

  708.     add_single_param_json<opr::ConvolutionForward>();

  709.     add_single_param_json<opr::Convolution3D>();

  710.     add_single_param_json<opr::ConvBiasForward>();

  711.     add_single_param_json<opr::ConvolutionBackwardData>();

  712.     add_single_param_json<opr::Convolution3DBackwardData>();

  713.     add_single_param_json<opr::ConvolutionBackwardFilter>();

  714.     add_single_param_json<opr::MatrixMul>();

  715.     add_single_param_json<opr::BatchedMatrixMul>();

  716.     add_single_param_json<opr::Dot>();

  717.     add_single_param_json<opr::MatrixInverse>();

  718.     add_single_param_json<opr::PoolingForward>();

  719.     add_single_param_json<opr::SVD>();

  720.     add_single_param_json<opr::MaskConvolution>();

  721.     add_single_param_json<opr::Images2Neibs>();

  722.     add_single_param_json<opr::Local>();

  723.     add_single_param_json<opr::GroupLocal>();

  724.     add_single_param_json<opr::LRN>();

  725.     add_single_param_json<opr::Concat>();

  726.     add_single_param_json<opr::Dimshuffle>();

  727.     add_single_param_json<opr::AxisAddRemove>();

  728.     add_single_param_json<opr::Subtensor>();

  729.     add_single_param_json<opr::SetSubtensor>();

  730.     add_single_param_json<opr::IncrSubtensor>();

  731.     add_single_param_json<opr::IndexingMultiAxisVec>();

  732.     add_single_param_json<opr::IndexingSetMultiAxisVec>();

  733.     add_single_param_json<opr::IndexingIncrMultiAxisVec>();

  734.     add_single_param_json<opr::MeshIndexing>();

  735.     add_single_param_json<opr::SetMeshIndexing>();

  736.     add_single_param_json<opr::IncrMeshIndexing>();

  737.     add_single_param_json<opr::BatchedMeshIndexing>();

  738.     add_single_param_json<opr::BatchedSetMeshIndexing>();

  739.     add_single_param_json<opr::BatchedIncrMeshIndexing>();

  740.     add_single_param_json<opr::Reduce>();

  741.     add_single_param_json<opr::LocalShareForward>();

  742.     add_single_param_json<opr::LocalShareBackwardData>();

  743.     add_single_param_json<opr::LocalShareBackwardFilter>();

  744.     add_single_param_json<opr::DeformableConvForward>();

  745.     add_single_param_json<opr::DeformableConvBackwardFilter>();

  746.     add_single_param_json<opr::DeformableConvBackwardData>();

  747.     add_single_param_json<opr::DeformablePSROIPoolingForward>();

  748.     add_single_param_json<opr::BatchConvBiasForward>();

  749.     add_single_param_json<opr::BatchNormForward>();

  750.     add_single_param_json<opr::Reshape>();

  751.     add_single_param_json<opr::GetVarShape>();

  752.     add_single_param_json<opr::Argsort>();

  753.     add_single_param_json<opr::Argmin>();

  754.     add_single_param_json<opr::Argmax>();

  755.     add_single_param_json<opr::ElemwiseMultiType>();

  756.     add_single_param_json<opr::AdaptivePooling>();

  757.     add_single_param_json<opr::ROIPooling>();

  758.     add_single_param_json<opr::ROIAlign>();

  759.     add_single_param_json<opr::WarpPerspective>();

  760.     add_single_param_json<opr::Remap>();

  761.     add_single_param_json<opr::Resize>();

  762.     add_single_param_json<opr::IndexingOneHot>();

  763.     add_single_param_json<opr::IndexingSetOneHot>();

  764.     add_single_param_json<opr::WarpAffine>();

  765.     add_single_param_json<opr::TopK>();

  766.     add_single_param_json<opr::UniformRNG>();

  767.     add_single_param_json<opr::GaussianRNG>();

  768.     add_single_param_json<opr::Linspace>();

  769.     add_single_param_json<opr::Eye>();

  770.     add_single_param_json<opr::standalone::NMSKeep>();

  771.     add_single_param_json<opr::CvtColor>();

  772. 

  773. #endif

  774. }

  775. 

  776. OprFootprint::Result OprFootprint::calc_footprint(cg::OperatorNodeBase* opr) {

  777.     Result rst;

  778.     auto&& dep_map = opr->node_prop().dep_map();

  779.     for (auto&& inp : opr->input()) {

  780.         if (inp->mem_plan().valid())

  781.             rst.inp_layout.push_back(inp->layout());

  782.         else

  783.             rst.inp_layout.push_back({inp->shape(), inp->dtype()});

  784.         if (cg::OperatorNodeBase::NodeProp::is_device_value_dep(

  785.                     dep_map.at(inp))) {

  786.             rst.memory += inp->dtype().size(inp->shape().total_nr_elems());

  787.         }

  788.     }

  789.     for (auto&& out : opr->output()) {

  790.         if (out->contain_flag(VarNode::Flag::VOLATILE_CONTENT))

  791.             continue;

  792.         rst.out_shape.push_back(out->shape());

  793.         rst.memory += out->dtype().size(out->shape().total_nr_elems());

  794.     }

  795.     rst.computation = get_computation(opr);

  796. #if MGB_ENABLE_JSON

  797.     rst.param = get_param_json(opr);

  798. #endif

  799.     rst.opr_type = opr->dyn_typeinfo();

  800.     return rst;

  801. }

  802. 

  803. uint64_t OprFootprint::get_computation(cg::OperatorNodeBase* opr) {

  804.     auto comp_trait = m_type2comp_footprint.find(opr->dyn_typeinfo());

  805.     if (comp_trait != m_type2comp_footprint.end()) {

  806.         return (comp_trait->second)(opr);

  807.     }

  808.     return 0;

  809. }

  810. 

  811. #if MGB_ENABLE_JSON

  812. std::shared_ptr<json::Value> OprFootprint::get_param_json(

  813.         cg::OperatorNodeBase* opr) {

  814.     auto param_trait = m_type2param_json.find(opr->dyn_typeinfo());

  815.     if (param_trait != m_type2param_json.end()) {

  816.         return (param_trait->second)(opr);

  817.     }

  818.     return json::Object::make();

  819. }

  820. 

  821. std::shared_ptr<json::Value> OprFootprint::Result::to_json() const {

  822.     using namespace json;

  823.     std::shared_ptr<Value> comp;

  824.     if (computation) {

  825.         comp = NumberInt::make(computation);

  826.     } else {

  827.         comp = Null::make();

  828.     }

  829.     auto format_shape_arr = [](const TensorShapeArray& arr) {

  830.         auto ret = Array::make();

  831.         for (auto&& shp : arr) {

  832.             auto cur = Array::make();

  833.             for (size_t i = 0; i < shp.ndim; ++i) {

  834.                 cur->add(NumberInt::make(shp[i]));

  835.             }

  836.             ret->add(std::move(cur));

  837.         }

  838.         return ret;

  839.     };

  840.     auto format_layout_arr =

  841.             [](const TensorLayoutArray& arr) -> std::shared_ptr<Value> {

  842.         auto ret = Array::make();

  843.         bool have_non_contig = false;

  844.         for (auto&& item : arr) {

  845.             if (item.is_contiguous()) {

  846.                 ret->add(json::Null::make());

  847.             } else {

  848.                 have_non_contig = true;

  849.                 auto cur = Array::make();

  850.                 for (size_t i = 0; i < item.ndim; ++i) {

  851.                     cur->add(NumberInt::make(item.stride[i]));

  852.                 }

  853.                 ret->add(std::move(cur));

  854.             }

  855.         }

  856.         if (!have_non_contig) {

  857.             ret.reset();

  858.         }

  859.         return ret;

  860.     };

  861. 

  862.     TensorShapeArray inp_shape;

  863.     for (auto&& i : inp_layout)

  864.         inp_shape.push_back(i);

  865.     auto ret = Object::make({{"computation", std::move(comp)},

  866.                              {"memory", NumberInt::make(memory)},

  867.                              {"in_shapes", format_shape_arr(inp_shape)},

  868.                              {"out_shapes", format_shape_arr(out_shape)},

  869.                              {"param", param}});

  870.     if (auto inp_layout_json = format_layout_arr(inp_layout)) {

  871.         ret->operator[]("in_layouts") = std::move(inp_layout_json);

  872.     }

  873.     return ret;

  874. }

  875. 

  876. std::shared_ptr<json::Value> OprFootprint::get_opr_fp_graph_exec(

  877.         cg::ComputingGraph& graph, const SymbolVarArray& outputs) {

  878.     OprFootprint m_opr_footprint;

  879.     ComputingGraph::OutputSpec out_spec;

  880.     for (auto i : outputs) {

  881.         out_spec.emplace_back(i, nullptr);

  882.     }

  883.     graph.options().allocate_static_mem_after_graph_compile = true;

  884.     auto async_exec = graph.compile(out_spec);

  885.     std::vector<std::pair<json::String, std::shared_ptr<json::Value>>> rst_vals;

  886.     auto on_opr = [&m_opr_footprint, &rst_vals](cg::OperatorNodeBase* opr) {

  887.         Result trait(m_opr_footprint.calc_footprint(opr));

  888.         rst_vals.emplace_back(json::String(opr->id_str()), trait.to_json());

  889.         return true;

  890.     };

  891.     async_exec->iter_opr_seq(on_opr);

  892.     auto opr_fp = json::Object::make(rst_vals);

  893.     return json::Object::make(

  894.             {{"opr_footprint", opr_fp}, {"graph_exec", async_exec->to_json()}});

  895. }

  896. #endif

  897. 

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