Nodes.h

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.
 */
 
#ifndef __LOCO_IR_NODES_H__
#define __LOCO_IR_NODES_H__
 
#include "loco/IR/Node.h"
#include "loco/IR/Use.h"
#include "loco/IR/Domain.h"
#include "loco/IR/DataType.h"
#include "loco/IR/DataTypeTraits.h"
#include "loco/IR/Dimension.h"
#include "loco/IR/Window.h"
#include "loco/IR/Stride.h"
#include "loco/IR/Padding2D.h"
#include "loco/IR/PaddingND.h"
#include "loco/IR/TensorAxis.h"
#include "loco/IR/TensorAxisSet.h"
#include "loco/IR/FeatureCodec.h"
#include "loco/IR/FilterCodec.h"
#include "loco/IR/DepthwiseFilterCodec.h"
#include "loco/IR/MatrixCodec.h"
#include "loco/IR/NodeMixins.h"
#include "loco/IR/CanonicalNodeDecl.h"
#include "loco/IR/GraphInputIndex.h"
#include "loco/IR/GraphOutputIndex.h"
 
namespace loco
{
 
class Graph;
class GraphInput;
class GraphOutput;
 
class Push /* to user */ final
  : public CanonicalNodeDef<CanonicalOpcode::Push, FixedArity<1>::Mixin>
{
public:
  Push() = default;
 
public:
  Node *from(void) const { return at(0)->node(); }
  void from(Node *node) { at(0)->node(node); }
 
public:
  void index(const GraphOutputIndex &index);
 
  GraphOutputIndex index(void) const;
 
  bool indexed(void) const { return _index != -1; }
 
private:
  int64_t _index = -1; // Uninitialized
};
 
void link(GraphOutput *, Push *push);
 
Push *push_node(Graph *g, const GraphOutputIndex &index);
 
class Pull /* from user */ final
  : public CanonicalNodeDef<CanonicalOpcode::Pull, FixedArity<0>::Mixin,
                            With<NodeTrait::TensorShape>::Mixin>
{
public:
  Pull() = default;
 
public:
  void index(const GraphInputIndex &index);
 
  GraphInputIndex index(void) const;
 
  bool indexed(void) const { return _index != -1; }
 
public:
  void dtype(const DataType &d);
  DataType dtype(void) const;
 
private:
  int64_t _index = -1; // Uninitialized
 
  DataType _dtype = DataType::Unknown;
};
 
void link(GraphInput *, Pull *pull);
 
Pull *pull_node(Graph *g, const GraphInputIndex &index);
 
class Forward final : public CanonicalNodeDef<CanonicalOpcode::Forward, FixedArity<1>::Mixin>
{
public:
  Forward() = default;
 
public:
  Node *input(void) const { return at(0)->node(); }
  void input(Node *node) { at(0)->node(node); }
};
 
class ReLU final : public CanonicalNodeDef<CanonicalOpcode::ReLU, FixedArity<1>::Mixin>
{
public:
  ReLU() = default;
 
public:
  Node *input(void) const { return at(0)->node(); }
  void input(Node *node) { at(0)->node(node); }
};
 
class ReLU6 final : public CanonicalNodeDef<CanonicalOpcode::ReLU6, FixedArity<1>::Mixin>
{
public:
  ReLU6() = default;
 
public:
  Node *input(void) const { return at(0)->node(); }
  void input(Node *node) { at(0)->node(node); }
};
 
class Tanh final : public CanonicalNodeDef<CanonicalOpcode::Tanh, FixedArity<1>::Mixin>
{
public:
  Tanh() = default;
 
public:
  Node *input(void) const { return at(0)->node(); }
  void input(Node *node) { at(0)->node(node); }
};
 
class ConstGen final
  : public CanonicalNodeDef<CanonicalOpcode::ConstGen, FixedArity<0>::Mixin,
                            With<NodeTrait::DataType>::Mixin, With<NodeTrait::TensorShape>::Mixin>
{
public:
  ConstGen() = default;
 
public:
  template <DataType DT> uint32_t size(void) const;
 
  template <DataType DT> void size(uint32_t size);
 
  template <DataType DT> const typename DataTypeImpl<DT>::Type &at(uint32_t n) const;
 
  template <DataType DT> typename DataTypeImpl<DT>::Type &at(uint32_t n);
 
private:
  std::vector<uint8_t> _data;
};
 
class MaxPool2D final : public CanonicalNodeDef<CanonicalOpcode::MaxPool2D, FixedArity<1>::Mixin>
{
public:
  Node *ifm(void) const { return at(0)->node(); }
  void ifm(Node *node) { at(0)->node(node); }
 
public:
  const Padding2D *pad(void) const { return &_pad; }
  Padding2D *pad(void) { return &_pad; }
 
public:
  const Window<2> *window(void) const { return &_window; }
  Window<2> *window(void) { return &_window; }
 
public:
  const Stride<2> *stride(void) const { return &_stride; }
  Stride<2> *stride(void) { return &_stride; }
 
private:
  // Pad
  Padding2D _pad;
  // Window
  Window<2> _window;
  // Stride
  Stride<2> _stride;
};
 
class AvgPool2D final : public CanonicalNodeDef<CanonicalOpcode::AvgPool2D, FixedArity<1>::Mixin>
{
public:
  enum class Convention
  {
    Unknown,
    // Use the number of elements in each receptive field as a divisor
    Full,
    // Use the number of valid (non-padding) elements in each receptive field as a divisor
    Valid
  };
 
public:
  Node *ifm(void) const { return at(0)->node(); }
  void ifm(Node *node) { at(0)->node(node); }
 
public:
  Convention convention(void) const { return _convention; }
  void convention(const Convention &convention) { _convention = convention; }
 
public:
  const Padding2D *pad(void) const { return &_pad; }
  Padding2D *pad(void) { return &_pad; }
 
public:
  const Window<2> *window(void) const { return &_window; }
  Window<2> *window(void) { return &_window; }
 
public:
  const Stride<2> *stride(void) const { return &_stride; }
  Stride<2> *stride(void) { return &_stride; }
 
private:
  Convention _convention = Convention::Unknown;
  Padding2D _pad;
  Window<2> _window;
  Stride<2> _stride;
};
 
class FeatureEncode final
  : public CanonicalNodeDef<CanonicalOpcode::FeatureEncode, FixedArity<1>::Mixin>
{
public:
  Node *input(void) const { return at(0)->node(); }
  void input(Node *node) { at(0)->node(node); }
 
public:
  FeatureEncoder *encoder(void) const { return _enc.get(); }
  void encoder(std::unique_ptr<FeatureEncoder> &&enc) { _enc = std::move(enc); }
 
private:
  std::unique_ptr<FeatureEncoder> _enc{nullptr};
};
 
class FeatureDecode final
  : public CanonicalNodeDef<CanonicalOpcode::FeatureDecode, FixedArity<1>::Mixin>
{
public:
  Node *input(void) const { return at(0)->node(); }
  void input(Node *node) { at(0)->node(node); }
 
public:
  FeatureDecoder *decoder(void) const { return _dec.get(); }
  void decoder(std::unique_ptr<FeatureDecoder> &&dec) { _dec = std::move(dec); }
 
private:
  std::unique_ptr<FeatureDecoder> _dec{nullptr};
};
 
class FilterEncode final
  : public CanonicalNodeDef<CanonicalOpcode::FilterEncode, FixedArity<1>::Mixin>
{
public:
  Node *input(void) const { return at(0)->node(); }
  void input(Node *node) { at(0)->node(node); }
 
public:
  FilterEncoder *encoder(void) const { return _enc.get(); }
  void encoder(std::unique_ptr<FilterEncoder> &&enc) { _enc = std::move(enc); }
 
private:
  std::unique_ptr<FilterEncoder> _enc{nullptr};
};
 
class FilterDecode final
  : public CanonicalNodeDef<CanonicalOpcode::FilterDecode, FixedArity<1>::Mixin>
{
public:
  Node *input(void) const { return at(0)->node(); }
  void input(Node *node) { at(0)->node(node); }
 
public:
  FilterDecoder *decoder(void) const { return _dec.get(); }
  void decoder(std::unique_ptr<FilterDecoder> &&dec) { _dec = std::move(dec); }
 
private:
  std::unique_ptr<FilterDecoder> _dec{nullptr};
};
 
class DepthwiseFilterEncode final
  : public CanonicalNodeDef<CanonicalOpcode::DepthwiseFilterEncode, FixedArity<1>::Mixin>
{
public:
  Node *input(void) const { return at(0)->node(); }
  void input(Node *node) { at(0)->node(node); }
 
public:
  DepthwiseFilterEncoder *encoder(void) const { return _enc.get(); }
  void encoder(std::unique_ptr<DepthwiseFilterEncoder> &&enc) { _enc = std::move(enc); }
 
private:
  std::unique_ptr<DepthwiseFilterEncoder> _enc{nullptr};
};
 
class DepthwiseFilterDecode final
  : public CanonicalNodeDef<CanonicalOpcode::DepthwiseFilterDecode, FixedArity<1>::Mixin>
{
public:
  Node *input(void) const { return at(0)->node(); }
  void input(Node *node) { at(0)->node(node); }
 
public:
  DepthwiseFilterDecoder *decoder(void) const { return _dec.get(); }
  void decoder(std::unique_ptr<DepthwiseFilterDecoder> &&dec) { _dec = std::move(dec); }
 
private:
  std::unique_ptr<DepthwiseFilterDecoder> _dec{nullptr};
};
 
enum class ReshapeType
{
  Fixed, // shape is known at compile time
  // Add another type for a case when shape is not known at compile time
};
 
template <ReshapeType RT> class Reshape;
 
template <>
class Reshape<ReshapeType::Fixed> final
  : public CanonicalNodeDef<CanonicalOpcode::FixedReshape, FixedArity<1>::Mixin,
                            With<NodeTrait::TensorShape>::Mixin>
{
public:
  Node *input(void) const { return at(0)->node(); }
  void input(Node *node) { at(0)->node(node); }
};
 
using FixedReshape = Reshape<ReshapeType::Fixed>;
 
class TensorConcat final
  : public CanonicalNodeDef<CanonicalOpcode::TensorConcat, FixedArity<2>::Mixin>
{
public:
  Node *lhs(void) const { return at(0)->node(); }
  void lhs(Node *node) { at(0)->node(node); }
 
  Node *rhs(void) const { return at(1)->node(); }
  void rhs(Node *node) { at(1)->node(node); }
 
public:
  uint32_t axis(void) const { return _axis; }
  void axis(uint32_t val) { _axis = val; }
 
private:
  // Axis
  uint32_t _axis{0};
};
 
class Conv2D final : public CanonicalNodeDef<CanonicalOpcode::Conv2D, FixedArity<2>::Mixin>
{
public:
  Node *ifm(void) const { return at(0)->node(); }
  void ifm(Node *node) { at(0)->node(node); }
 
  Node *ker(void) const { return at(1)->node(); }
  void ker(Node *node) { at(1)->node(node); }
 
public:
  const Padding2D *pad(void) const { return &_pad; }
  Padding2D *pad(void) { return &_pad; }
 
public:
  const Stride<2> *stride(void) const { return &_stride; }
  Stride<2> *stride(void) { return &_stride; }
 
private:
  Padding2D _pad;
  Stride<2> _stride;
 
  // TODO Support "Dilation"
};
 
class DepthwiseConv2D final
  : public CanonicalNodeDef<CanonicalOpcode::DepthwiseConv2D, FixedArity<2>::Mixin>
{
public:
  Node *ifm(void) const { return at(0)->node(); }
  void ifm(Node *node) { at(0)->node(node); }
 
  Node *ker(void) const { return at(1)->node(); }
  void ker(Node *node) { at(1)->node(node); }
 
public:
  const Padding2D *pad(void) const { return &_pad; }
  Padding2D *pad(void) { return &_pad; }
 
public:
  const Stride<2> *stride(void) const { return &_stride; }
  Stride<2> *stride(void) { return &_stride; }
 
private:
  Padding2D _pad;
  Stride<2> _stride;
 
  // TODO Support "Dilation"
};
 
enum class ReduceFunc
{
  Mean, // ReduceMean
  // TODO Support other reduce operations
};
 
class TensorReduce final
  : public CanonicalNodeDef<CanonicalOpcode::TensorReduce, FixedArity<1>::Mixin>
{
public:
  Node *input(void) const { return at(0)->node(); }
  void input(Node *node) { at(0)->node(node); }
 
public:
  const TensorAxisSet *axes(void) const { return &_axes; }
  TensorAxisSet *axes(void) { return &_axes; }
 
public:
  ReduceFunc func(void) const { return _func; }
  void func(ReduceFunc func) { _func = func; }
 
private:
  TensorAxisSet _axes;
  ReduceFunc _func{ReduceFunc::Mean};
};
 
class TransposedConv2D final
  : public CanonicalNodeDef<CanonicalOpcode::TransposedConv2D, FixedArity<2>::Mixin>
{
public:
  Node *ifm(void) const { return at(0)->node(); }
  void ifm(Node *node) { at(0)->node(node); }
 
  Node *ker(void) const { return at(1)->node(); }
  void ker(Node *node) { at(1)->node(node); }
 
public:
  const Padding2D *pad(void) const { return &_pad; }
  Padding2D *pad(void) { return &_pad; }
 
public:
  const Stride<2> *stride(void) const { return &_stride; }
  Stride<2> *stride(void) { return &_stride; }
 
private:
  Padding2D _pad;
  Stride<2> _stride;
 
  // TODO Support "Dilation"
};
 
template <Domain D> class Softmax;
 
template <>
class Softmax<Domain::Tensor> final
  : public CanonicalNodeDef<CanonicalOpcode::TensorSoftmax, FixedArity<1>::Mixin>
{
public:
  Softmax() = default;
 
public:
  Node *input(void) const { return at(0)->node(); }
  void input(Node *node) { return at(0)->node(node); }
 
  uint32_t axis(void) const { return _axis; }
  void axis(uint32_t axis) { _axis = axis; }
 
private:
  uint32_t _axis = 0;
};
 
using TensorSoftmax = Softmax<Domain::Tensor>;
 
class BiasDecode final : public CanonicalNodeDef<CanonicalOpcode::BiasDecode, FixedArity<1>::Mixin>
{
public:
  BiasDecode() = default;
 
public:
  Node *input(void) const { return at(0)->node(); }
  void input(Node *node) { at(0)->node(node); }
};
 
class BiasEncode final : public CanonicalNodeDef<CanonicalOpcode::BiasEncode, FixedArity<1>::Mixin>
{
public:
  BiasEncode() = default;
 
public:
  Node *input(void) const { return at(0)->node(); }
  void input(Node *node) { at(0)->node(node); }
};
 
template <Domain D> class BiasAdd;
 
template <>
class BiasAdd<Domain::Tensor> final
  : public CanonicalNodeDef<CanonicalOpcode::TensorBiasAdd, FixedArity<2>::Mixin>
{
public:
  BiasAdd() = default;
 
public:
  Node *value(void) const { return at(0)->node(); }
  void value(Node *node) { return at(0)->node(node); }
 
  Node *bias(void) const { return at(1)->node(); }
  void bias(Node *node) { return at(1)->node(node); }
 
  uint32_t axis(void) const { return _axis; }
  void axis(uint32_t axis) { _axis = axis; }
 
private:
  uint32_t _axis = 0;
};
 
//
// Alias for external users
//
// loco::TensorBiasAdd
//        vs.
// loco::BiasAdd<loco::Domain::Tensor>
//
using TensorBiasAdd = BiasAdd<Domain::Tensor>;
 
template <>
class BiasAdd<Domain::Feature> final
  : public CanonicalNodeDef<CanonicalOpcode::FeatureBiasAdd, FixedArity<2>::Mixin>
{
public:
  BiasAdd() = default;
 
public:
  Node *value(void) const { return at(0)->node(); }
  void value(Node *node) { return at(0)->node(node); }
 
  Node *bias(void) const { return at(1)->node(); }
  void bias(Node *node) { return at(1)->node(node); }
};
 
using FeatureBiasAdd = BiasAdd<Domain::Feature>;
 
class TensorConstantPad final
  : public CanonicalNodeDef<CanonicalOpcode::TensorConstantPad, FixedArity<2>::Mixin>
{
public:
  Node *input(void) const { return at(0)->node(); }
  void input(Node *node) { at(0)->node(node); }
 
  Node *constant(void) const { return at(1)->node(); }
  void constant(Node *node) { at(1)->node(node); }
 
public:
  const PaddingND *padding(void) const { return &_padding; }
  PaddingND *padding(void) { return &_padding; }
 
private:
  PaddingND _padding;
};
 
class EltwiseAdd final : public CanonicalNodeDef<CanonicalOpcode::EltwiseAdd, FixedArity<2>::Mixin>
{
public:
  EltwiseAdd() = default;
 
public:
  Node *lhs(void) const { return at(0)->node(); }
  void lhs(Node *node) { return at(0)->node(node); }
 
  Node *rhs(void) const { return at(1)->node(); }
  void rhs(Node *node) { return at(1)->node(node); }
};
 
class EltwiseMax final : public CanonicalNodeDef<CanonicalOpcode::EltwiseMax, FixedArity<2>::Mixin>
{
public:
  EltwiseMax() = default;
 
public:
  Node *lhs(void) const { return at(0)->node(); }
  void lhs(Node *node) { return at(0)->node(node); }
 
  Node *rhs(void) const { return at(1)->node(); }
  void rhs(Node *node) { return at(1)->node(node); }
};
 
class EltwiseMul final : public CanonicalNodeDef<CanonicalOpcode::EltwiseMul, FixedArity<2>::Mixin>
{
public:
  EltwiseMul() = default;
 
public:
  Node *lhs(void) const { return at(0)->node(); }
  void lhs(Node *node) { return at(0)->node(node); }
 
  Node *rhs(void) const { return at(1)->node(); }
  void rhs(Node *node) { return at(1)->node(node); }
};
 
class EltwiseSub final : public CanonicalNodeDef<CanonicalOpcode::EltwiseSub, FixedArity<2>::Mixin>
{
public:
  EltwiseSub() = default;
 
public:
  Node *lhs(void) const { return at(0)->node(); }
  void lhs(Node *node) { return at(0)->node(node); }
 
  Node *rhs(void) const { return at(1)->node(); }
  void rhs(Node *node) { return at(1)->node(node); }
};
 
class EltwiseDiv final : public CanonicalNodeDef<CanonicalOpcode::EltwiseDiv, FixedArity<2>::Mixin>
{
public:
  EltwiseDiv() = default;
 
public:
  Node *lhs(void) const { return at(0)->node(); }
  void lhs(Node *node) { return at(0)->node(node); }
 
  Node *rhs(void) const { return at(1)->node(); }
  void rhs(Node *node) { return at(1)->node(node); }
};
 
class EltwiseSqrt final
  : public CanonicalNodeDef<CanonicalOpcode::EltwiseSqrt, FixedArity<1>::Mixin>
{
public:
  EltwiseSqrt() = default;
 
public:
  Node *input(void) const { return at(0)->node(); }
  void input(Node *node) { at(0)->node(node); }
};
 
class TensorBroadcast final
  : public CanonicalNodeDef<CanonicalOpcode::TensorBroadcast, FixedArity<1>::Mixin>
{
public:
  TensorBroadcast() = default;
 
public:
  Node *input(void) const { return at(0)->node(); }
  void input(Node *node) { at(0)->node(node); }
 
public:
  class Mapping final
  {
  public:
    Mapping() = default;
 
  public:
    bool defined(const TensorAxis &axis) const;
 
    const Dimension &dim(const TensorAxis &axis) const;
    Dimension &dim(const TensorAxis &axis);
 
  private:
    std::map<TensorAxis, Dimension> _content;
  };
 
  Mapping *mapping(void) { return &_mapping; }
  const Mapping *mapping(void) const { return &_mapping; }
 
private:
  Mapping _mapping;
};
 
class MatrixEncode final
  : public CanonicalNodeDef<CanonicalOpcode::MatrixEncode, FixedArity<1>::Mixin>
{
public:
  MatrixEncode() = default;
 
public:
  Node *input(void) const { return at(0)->node(); }
  void input(Node *node) { at(0)->node(node); }
 
public:
  MatrixEncoder *encoder(void) const { return _enc.get(); }
  void encoder(std::unique_ptr<MatrixEncoder> &&enc) { _enc = std::move(enc); }
 
private:
  std::unique_ptr<MatrixEncoder> _enc{nullptr};
};
 
class MatrixDecode final
  : public CanonicalNodeDef<CanonicalOpcode::MatrixDecode, FixedArity<1>::Mixin>
{
public:
  MatrixDecode() = default;
 
public:
  Node *input(void) const { return at(0)->node(); }
  void input(Node *node) { at(0)->node(node); }
 
public:
  MatrixDecoder *decoder(void) const { return _dec.get(); }
  void decoder(std::unique_ptr<MatrixDecoder> &&dec) { _dec = std::move(dec); }
 
private:
  std::unique_ptr<MatrixDecoder> _dec{nullptr};
};
 
class MatMul final : public CanonicalNodeDef<CanonicalOpcode::MatMul, FixedArity<2>::Mixin>
{
public:
  MatMul() = default;
 
public:
  Node *lhs(void) const { return at(0)->node(); }
  void lhs(Node *node) { return at(0)->node(node); }
 
  Node *rhs(void) const { return at(1)->node(); }
  void rhs(Node *node) { return at(1)->node(node); }
};
 
class TensorTranspose final
  : public CanonicalNodeDef<CanonicalOpcode::TensorTranspose, FixedArity<1>::Mixin>
{
public:
  TensorTranspose() = default;
 
public:
  Node *input(void) const { return at(0)->node(); }
  void input(Node *node) { return at(0)->node(node); }
 
  class Perm final
  {
  public:
    Perm() = default;
 
  public:
    uint32_t size() const { return _vals.size(); }
    void size(uint32_t size) { _vals.resize(size); }
 
    const TensorAxis &axis(TensorAxis n) const { return _vals[n]; }
    TensorAxis &axis(TensorAxis n) { return _vals[n]; }
 
  private:
    std::vector<TensorAxis> _vals;
  };
 
  Perm *perm(void) { return &_perm; }
  const Perm *perm(void) const { return &_perm; }
 
private:
  Perm _perm;
};
 
} // namespace loco
 
#endif // __LOCO_IR_NODES_H__