compiler/moco-tf/src/Canonicalization/PadCanonicalizer.cpp

   1 /*
   2  * Copyright (c) 2019 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 "PadCanonicalizer.h"
  18
  19 #include <moco/IR/TFDialect.h>
  20
  21 #include "loco/Service/TypeInference.h"
  22
  23 #include <stdex/Memory.h>
  24
  25 namespace
  26 {
  27
  28 bool canonicalize_pad(loco::Graph *graph, moco::TFPad *node)
  29 {
  30   /**
  31    * @note This will replace TFPad node with Canonical TensorConstantPad
  32    *
  33    *       Before
  34    *                 input --- TFPad -- C
  35    *                 paddings --/
  36    *       After
  37    *                 paddings  ------- TFPad --
  38    *                                  /
  39    *                 input ----------- TensorConstantPad -- C
  40    *                 ConstGen --------/
  41    *       Where
  42    *                 input : input of TFPad
  43    *                 paddings : paddings of TFPad. it becomes TensorConstantPad's attribute.
  44    *                 C : a node that uses TFPad as an input. TFPad is disconnected from C.
  45    *                 ConstGen : constant value of Pad. TFPad has zero value by default.
  46    */
  47
  48   auto pad_node = graph->nodes()->create<loco::TensorConstantPad>();
  49
  50   auto constant_node = graph->nodes()->create<loco::ConstGen>();
  51
  52   auto input_node = node->input();
  53   // TODO: support other dtype.
  54   assert(loco::dtype_get(input_node) == loco::DataType::FLOAT32);
  55   constant_node->dtype(loco::DataType::FLOAT32);
  56   // TODO: constant node changes to scalar when it is implemented.
  57   constant_node->shape({1});
  58   constant_node->size<loco::DataType::FLOAT32>(1);
  59   constant_node->at<loco::DataType::FLOAT32>(0) = 0.0f;
  60
  61   auto const_paddings_node = loco::must_cast<loco::ConstGen *>(node->paddings());
  62   // TODO: support S64 type.
  63   assert(const_paddings_node->dtype() == loco::DataType::S32);
  64   assert(const_paddings_node->rank() == 2);
  65   assert(const_paddings_node->dim(1).value() == 2);
  66
  67   auto padding = pad_node->padding();
  68   uint32_t padding_rank = const_paddings_node->dim(0).value();
  69   padding->rank(padding_rank);
  70
  71   for (uint32_t i = 0; i < padding_rank; i++)
  72   {
  73     padding->front(i) = const_paddings_node->at<loco::DataType::S32>(i << 1);
  74     padding->back(i) = const_paddings_node->at<loco::DataType::S32>((i << 1) + 1);
  75   }
  76
  77   // update connections
  78   pad_node->input(input_node);
  79   pad_node->constant(constant_node);
  80
  81   // replace node
  82   replace(node).with(pad_node);
  83
  84   return true;
  85 }
  86
  87 } // namespace
  88
  89 namespace moco
  90 {
  91 namespace tf
  92 {
  93
  94 bool PadCanonicalizer::transform(TFPad *node) const
  95 {
  96   return canonicalize_pad(node->graph(), node);
  97 }
  98
  99 } // namespace tf
 100 } // namespace moco