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