compiler/locomotiv/src/Node/TensorConcat.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 "NodeExecution.h"
  18
  19 #include "NodeDataImpl.h"
  20 #include "NodeDomain.h"
  21 #include "Validation.h"
  22
  23 #include <nncc/core/ADT/tensor/Shape.h>
  24 #include <nncc/core/ADT/tensor/Buffer.h>
  25 #include <nncc/core/ADT/tensor/Index.h>
  26 #include <nncc/core/ADT/tensor/IndexEnumerator.h>
  27 #include <nncc/core/ADT/tensor/LexicalLayout.h>
  28
  29 using nncc::core::ADT::tensor::Index;
  30 using nncc::core::ADT::tensor::IndexEnumerator;
  31 using nncc::core::ADT::tensor::LexicalLayout;
  32 using nncc::core::ADT::tensor::make_buffer;
  33 using nncc::core::ADT::tensor::Shape;
  34
  35 #include <cassert>
  36 #include <stdexcept>
  37
  38 namespace
  39 {
  40
  41 using namespace locomotiv;
  42
  43 void execute_node(loco::TensorConcat *tensor_concat)
  44 {
  45   validate(tensor_concat, "TensorConcat is nullptr");
  46
  47   auto lhs_data = annot_data(tensor_concat->lhs());
  48   auto rhs_data = annot_data(tensor_concat->rhs());
  49   auto axis = tensor_concat->axis();
  50
  51   validate(lhs_data && rhs_data, "Ingredient not ready");
  52   validate(lhs_data->dtype() == rhs_data->dtype(), "lhs and rhs of Concat should have same dtype");
  53
  54   validate(annot_domain(tensor_concat->lhs()) == loco::Domain::Tensor &&
  55                annot_domain(tensor_concat->rhs()) == loco::Domain::Tensor,
  56            "Some ingredients of TensorConcat is not Tensor");
  57
  58   // Calculate output shape
  59   Shape lhs_shape = *lhs_data->shape();
  60   Shape rhs_shape = *rhs_data->shape();
  61   Shape concat_shape;
  62
  63   assert(lhs_shape.rank() == rhs_shape.rank());
  64   concat_shape.resize(lhs_shape.rank());
  65   for (uint32_t index = 0; index < lhs_shape.rank(); ++index)
  66   {
  67     if (index == axis)
  68       concat_shape.dim(index) = lhs_shape.dim(index) + rhs_shape.dim(index);
  69     else
  70     {
  71       assert(lhs_shape.dim(index) == rhs_shape.dim(index));
  72       concat_shape.dim(index) = lhs_shape.dim(index);
  73     }
  74   }
  75   auto left_dim_size = lhs_shape.dim(axis);
  76
  77   // Copy data from two inputs LHS and RHS to Concat
  78   std::unique_ptr<NodeData> concat_data = nullptr;
  79   switch (lhs_data->dtype())
  80   {
  81     case loco::DataType::FLOAT32:
  82     {
  83       auto lhs_bufptr = lhs_data->as_f32_bufptr();
  84       auto rhs_bufptr = rhs_data->as_f32_bufptr();
  85       auto concat_buf = make_buffer<float, LexicalLayout>(concat_shape);
  86
  87       for (IndexEnumerator e{concat_shape}; e.valid(); e.advance())
  88       {
  89         const auto &e_index = e.current();
  90
  91         if (e_index.at(axis) < left_dim_size)
  92         {
  93           // Left index is same as output index
  94           concat_buf.at(e_index) = lhs_bufptr->at(e_index);
  95         }
  96         else
  97         {
  98           // Adjust right index to valid range
  99           Index r_index = e_index;
 100           r_index.at(axis) -= left_dim_size;
 101           concat_buf.at(e_index) = rhs_bufptr->at(r_index);
 102         }
 103       }
 104
 105       concat_data = make_data(concat_buf);
 106       break;
 107     }
 108     default:
 109       throw std::runtime_error("NYI for this DataType");
 110   }
 111
 112   assert(concat_data != nullptr);
 113   annot_data(tensor_concat, std::move(concat_data));
 114   annot_domain(tensor_concat, loco::Domain::Tensor);
 115 }
 116
 117 } // namespace
 118
 119 namespace locomotiv
 120 {
 121
 122 void NodeExecution::execute(loco::TensorConcat *tensor_concat) { execute_node(tensor_concat); }
 123
 124 } // namespace locomotiv