compiler/luci/pass/src/ResolveCustomOpBatchMatMulPass.cpp

   1 /*
   2  * Copyright (c) 2020 Samsung Electronics Co., Ltd. All Rights Reserved
   3  *
   4  * Licensed under the Apache License, Version 2.0 (the "License");
   5  * you may not use this file except in compliance with the License.
   6  * You may obtain a copy of the License at
   7  *
   8  *    http://www.apache.org/licenses/LICENSE-2.0
   9  *
  10  * Unless required by applicable law or agreed to in writing, software
  11  * distributed under the License is distributed on an "AS IS" BASIS,
  12  * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
  13  * See the License for the specific language governing permissions and
  14  * limitations under the License.
  15  */
  16
  17 #include "luci/Pass/ResolveCustomOpBatchMatMulPass.h"
  18
  19 #include <luci/IR/CircleNodes.h>
  20 #include <luci/Profile/CircleNodeOrigin.h>
  21
  22 #include <flatbuffers/flexbuffers.h>
  23
  24 namespace
  25 {
  26
  27 bool resolve_custom_op(luci::CircleCustom *cop)
  28 {
  29   const std::string custom_code = cop->custom_code();
  30   const std::vector<uint8_t> custom_options = cop->custom_options();
  31
  32   if (custom_code == "BatchMatMulV2")
  33   {
  34     auto name = cop->name();
  35     assert(name.length() > 0);
  36
  37     auto batch_matmul = cop->graph()->nodes()->create<luci::CircleBatchMatMul>();
  38     // input
  39     batch_matmul->x(cop->inputs(0));
  40     batch_matmul->y(cop->inputs(1));
  41     // TODO find much better way of parsing custom_options
  42     // adj
  43     auto map = flexbuffers::GetRoot(custom_options).AsMap();
  44     batch_matmul->adj_x(map["adj_x"].AsBool());
  45     batch_matmul->adj_y(map["adj_y"].AsBool());
  46     batch_matmul->name(name + "/BatchMatMul");
  47     luci::add_origin(batch_matmul, luci::get_origin(cop));
  48
  49     auto customOut = loco::succs(cop);
  50     assert(customOut.size() == 1);
  51     replace(*customOut.begin()).with(batch_matmul);
  52
  53     return true;
  54   }
  55
  56   return false;
  57 }
  58
  59 } // namespace
  60
  61 namespace luci
  62 {
  63
  64 /**
  65  *  BEFORE
  66  *         |             |
  67  *    [CircleNode]  [CircleNode]
  68  *          \           /
  69  *         [CircleCustom]("BatchMatMulV2")
  70  *               |
  71  *        [CircleCustomOut]
  72  *               |
  73  *          [CircleNode]
  74  *               |
  75  *
  76  *  AFTER
  77  *         |             |
  78  *    [CircleNode]  [CircleNode]
  79  *          \           /
  80  *       [CircleBatchMatMul]
  81  *               |
  82  *          [CircleNode]
  83  *               |
  84  */
  85 bool ResolveCustomOpBatchMatMulPass::run(loco::Graph *g)
  86 {
  87   bool changed = false;
  88   for (auto node : loco::active_nodes(loco::output_nodes(g)))
  89   {
  90     auto cop = dynamic_cast<luci::CircleCustom *>(node);
  91     if (not cop)
  92       continue;
  93
  94     if (resolve_custom_op(cop))
  95       changed = true;
  96   }
  97
  98   return changed;
  99 }
 100
 101 } // namespace luci