onert-micro/luci-interpreter/src/kernels/SplitV.cpp

   1 /*
   2  * Copyright (c) 2020 Samsung Electronics Co., Ltd. All Rights Reserved
   3  * Copyright 2023 The TensorFlow Authors. All Rights Reserved.
   4  *
   5  * Licensed under the Apache License, Version 2.0 (the "License");
   6  * you may not use this file except in compliance with the License.
   7  * You may obtain a copy of the License at
   8  *
   9  *    http://www.apache.org/licenses/LICENSE-2.0
  10  *
  11  * Unless required by applicable law or agreed to in writing, software
  12  * distributed under the License is distributed on an "AS IS" BASIS,
  13  * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
  14  * See the License for the specific language governing permissions and
  15  * limitations under the License.
  16  */
  17
  18 #include "Builders.h"
  19 #include "Utils.h"
  20 #include "Split.h"
  21
  22 namespace luci_interpreter
  23 {
  24
  25 void configure_kernel_CircleSplitV(const circle::Operator *cur_op, BaseRuntimeGraph *runtime_graph)
  26 {
  27   const auto axis_index = cur_op->inputs()->operator[](2);
  28   LUCI_INTERPRETER_CHECK(axis_index != -1);
  29
  30   const auto axis = runtime_graph->getCircleTensorByIndex(axis_index);
  31   LUCI_INTERPRETER_CHECK(axis != nullptr);
  32
  33   // Dynamic output tensors are needed if axis tensor is not constant
  34   // Now doesn't support dynamic shapes for SplitV
  35   LUCI_INTERPRETER_CHECK(runtime_graph->getConstDataByTensor(axis) != nullptr);
  36 }
  37
  38 void execute_kernel_CircleSplitV(const circle::Operator *cur_op, BaseRuntimeGraph *runtime_graph)
  39 {
  40   const auto input_index = cur_op->inputs()->operator[](0);
  41   const auto axis_index = cur_op->inputs()->operator[](2);
  42
  43   assert(input_index != -1);
  44   assert(axis_index != -1);
  45
  46   const auto input = runtime_graph->getCircleTensorByIndex(input_index);
  47   const auto axis = runtime_graph->getCircleTensorByIndex(axis_index);
  48
  49   assert(input != nullptr);
  50   assert(axis != nullptr);
  51
  52   const auto *axis_data = runtime_graph->getDataByTensor(axis);
  53   if (axis_data == nullptr)
  54     axis_data = runtime_graph->getConstDataByTensor(axis);
  55
  56   assert(axis_data);
  57
  58   int32_t axis_value = (kernels::getTensorData<int32_t>(axis_data))[0];
  59   if (axis_value < 0)
  60     axis_value += Tensor::num_dims(input);
  61
  62   assert(axis_value >= 0);
  63   assert(axis_value < Tensor::num_dims(input));
  64
  65   switch (Tensor::element_type(input))
  66   {
  67 #ifndef DIS_FLOAT
  68     case DataType::FLOAT32:
  69     {
  70       return splitImpl<float>(cur_op, input, axis_value, runtime_graph);
  71     }
  72 #endif // DIS_FLOAT
  73 #ifndef DIS_QUANT
  74     case DataType::S8:
  75     {
  76       return splitImpl<int8_t>(cur_op, input, axis_value, runtime_graph);
  77     }
  78     case DataType::S16:
  79     {
  80       return splitImpl<int16_t>(cur_op, input, axis_value, runtime_graph);
  81     }
  82 #endif // DIS_QUANT
  83     case DataType::S32:
  84     {
  85       return splitImpl<int32_t>(cur_op, input, axis_value, runtime_graph);
  86     }
  87     default:
  88       assert(false && "Unsupported type");
  89   }
  90 }
  91
  92 } // namespace luci_interpreter