CircleTensorExporter.cpp

Go to the documentation of this file.

/*
 * Copyright (c) 2019 Samsung Electronics Co., Ltd. All Rights Reserved
 *
 * Licensed under the Apache License, Version 2.0 (the "License");
 * you may not use this file except in compliance with the License.
 * You may obtain a copy of the License at
 *
 *      http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 */
 
#include "CircleTensorExporter.h"
#include "CircleTypeInference.h"
#include "ShapeInference.h"
 
// TODO Fix include style
#include "loco/IR/Algorithm.h"
#include "loco/IR/CanonicalNode.h"
#include "loco/IR/CanonicalNodeVisitor.h"
#include "loco/IR/DataTypeTraits.h"
 
#include "Dialect/IR/TFLNodes.h"
 
#include <oops/InternalExn.h>
 
using namespace circle;
using namespace flatbuffers;
 
namespace
{
 
using namespace exo;
using namespace exo::circle_detail;
 
class TFLTensorInfo
{
public:
  TFLTensorInfo() = default;
 
public:
  void name(const std::string &name) { _name = name; }
  const std::string &name(void) const { return _name; }
 
public:
  const circle::TensorType &dtype(void) const { return _dtype; }
  void dtype(const circle::TensorType &dtype) { _dtype = dtype; }
 
  const ShapeDescription &shape(void) const { return _shape; }
  void shape(const ShapeDescription &shape) { _shape = shape; }
 
public:
  locoex::TFLConst *tfl_content(void) const { return _tfl_content; }
  void tfl_content(locoex::TFLConst *c) { _tfl_content = c; }
 
private:
  std::string _name;
 
  circle::TensorType _dtype{circle::TensorType_FLOAT32};
  ShapeDescription _shape{};
 
  // TODO Find a better design
  loco::ConstGen *_content = nullptr; // TODO deprecate
  locoex::TFLConst *_tfl_content = nullptr;
};
 
using TFLTensorContext = std::vector<TFLTensorInfo>;
 
struct NoOpDetector final : public loco::CanonicalNodeMutableVisitor<bool>
{
  bool visit(loco::BiasEncode *) final
  {
    // BiasEncode is always noop
    return true;
  }
 
  bool visit(loco::FilterEncode *node) final
  {
    auto encoder = dynamic_cast<loco::PermutingEncoder<loco::Domain::Filter> *>(node->encoder());
    if (encoder != nullptr)
    {
      auto perm = encoder->perm();
      return isNHWC(perm);
    }
    return false;
  }
 
  bool visit(loco::FeatureEncode *node) final
  {
    auto encoder = dynamic_cast<loco::PermutingEncoder<loco::Domain::Feature> *>(node->encoder());
    if (encoder != nullptr)
    {
      auto perm = encoder->perm();
      return isNHWC(perm);
    }
    return false;
  }
 
  bool visit(loco::FeatureDecode *node) final
  {
    auto decoder = dynamic_cast<loco::PermutingDecoder<loco::Domain::Feature> *>(node->decoder());
    if (decoder != nullptr)
    {
      auto perm = decoder->perm();
      return isNHWC(perm);
    }
    return false;
  }
 
  // Return false by default
  bool visit(loco::Node *) final { return false; }
};
 
bool isNoOp(loco::Node *node)
{
  if (auto canonical_node = dynamic_cast<loco::CanonicalNode *>(node))
  {
    NoOpDetector d;
    return canonical_node->accept(&d);
  }
  return false;
}
 
void allocateCircleTensor(loco::Node *node, TFLTensorContext &ctx)
{
  if (isNoOp(node))
  {
    assert(node->arity() == 1 && node->arg(0) != nullptr);
    set_tensor_index(node, get_tensor_index(node->arg(0)));
    return;
  }
 
  auto tensor_index = static_cast<TFLTensorIndex>(ctx.size());
  // TODO Use Graph-level metadata for Input & Output
  auto tensor_name = "t_" + std::to_string(tensor_index);
 
  TFLTensorInfo tensor_info;
 
  tensor_info.name(tensor_name);
  tensor_info.dtype(TypeInference::get(node));
  tensor_info.shape(ShapeInference::get(node));
 
  tensor_info.tfl_content(dynamic_cast<locoex::TFLConst *>(node));
 
  set_tensor_index(node, tensor_index);
 
  ctx.emplace_back(tensor_info);
}
 
} // namespace
 
namespace
{
 
flatbuffers::Offset<Vector<int32_t>> encodeShape(FlatBufferBuilder &builder,
                                                 const ShapeDescription &shape)
{
  assert(shape._rank_known && "unknown number of dimensions is not supported");
  return builder.CreateVector(shape._dims);
}
 
flatbuffers::Offset<circle::Buffer> encodeOpBuffer(FlatBufferBuilder &builder)
{
  return CreateBuffer(builder);
}
 
template <typename NodeT>
flatbuffers::Offset<circle::Buffer> encodeOpBuffer(FlatBufferBuilder &builder, NodeT *)
{
  return CreateBuffer(builder);
}
 
template <loco::DataType DT>
flatbuffers::Offset<circle::Buffer> encodeOpBufferByDType(FlatBufferBuilder &builder,
                                                          locoex::TFLConst *c)
{
  using NativeType = typename loco::DataTypeImpl<DT>::Type;
 
  std::vector<NativeType> raw_data;
  const uint32_t size = c->size<DT>();
  raw_data.reserve(size);
  for (uint32_t i = 0; i < size; ++i)
  {
    raw_data.push_back(c->at<DT>(i));
  }
  const size_t raw_size = size * sizeof(NativeType);
  auto array_offset = builder.CreateVector(reinterpret_cast<uint8_t *>(raw_data.data()), raw_size);
  return CreateBuffer(builder, array_offset);
}
 
template <>
flatbuffers::Offset<circle::Buffer> encodeOpBuffer(FlatBufferBuilder &builder, locoex::TFLConst *c)
{
  if (c->dtype() == loco::DataType::FLOAT32)
  {
    return encodeOpBufferByDType<loco::DataType::FLOAT32>(builder, c);
  }
  else if (c->dtype() == loco::DataType::S32)
  {
    return encodeOpBufferByDType<loco::DataType::S32>(builder, c);
  }
 
  INTERNAL_EXN_V("Unsupported datatype", oops::to_uint32(c->dtype()));
}
 
} // namespace
 
namespace exo
{
namespace circle_detail
{
 
void exportOpDefinedTensor(const TFLTensorInfo &info, FlatBufferBuilder &builder,
                           SerializedModelData &gd)
{
  // Create and register output tensor shape
  auto shape_offset = encodeShape(builder, info.shape());
 
  // encode and register output tensor buffer
  auto buffer = info.tfl_content() == nullptr ? encodeOpBuffer(builder)
                                              : encodeOpBuffer(builder, info.tfl_content());
 
  auto buffer_id = static_cast<uint32_t>(gd._buffers.size());
  gd._buffers.push_back(buffer);
 
  auto name_offset = builder.CreateString(info.name());
  auto tensor_offset = CreateTensor(builder, shape_offset, info.dtype(), buffer_id, name_offset,
                                    /*quantization*/ 0, /*is_variable*/ false);
  gd._tensors.push_back(tensor_offset);
}
 
void exportOpDefinedTensors(loco::Graph *g, FlatBufferBuilder &builder, SerializedModelData &gd)
{
  TFLTensorContext tensor_ctx;
 
  for (auto node : loco::postorder_traversal(loco::output_nodes(g)))
  {
    allocateCircleTensor(node, tensor_ctx);
  }
 
  // add one empty buffer
  //   note: this follows TFLite
  //   note: there's a comment in tflite fbs file
  //   - Note the 0th entry of this array must be an empty buffer (sentinel).
  //   - This is a convention so that tensors without a buffer can provide 0 as
  //   - their buffer.
  auto buffer = encodeOpBuffer(builder);
  gd._buffers.push_back(buffer);
 
  for (const auto &tensor_info : tensor_ctx)
  {
    exportOpDefinedTensor(tensor_info, builder, gd);
  }
}
 
} // namespace circle_detail
} // namespace exo