lstm layout inference & Added unidirectional lstm layout inference (#392)

Signed-off-by: Chen Xin <jack.chen@verisilicon.com>

Co-authored-by: Chen Xin <jack.chen@verisilicon.com>

include/tim/vx/ops/unidirectional_sequence_lstm.h

@@ -47,7 +47,7 @@ namespace ops {
       UnidirectionalSequenceLstm(
           Graph* graph, float cell_clip, float proj_clip,
           ActivationType act_type, float forget_bias, bool time_major = false,
-          ActivationType recurrent_act_type = ActivationType::kNONE,
+          ActivationType recurrent_act_type = ActivationType::kSIGMOID,
           bool return_sequences = false /*False: only return last state*/
       );
 

src/tim/transform/layout_inference.cc

@@ -61,6 +61,7 @@
 #include "ops/conv3d_layout_inference.h"
 #include "ops/default_layout_inference.h"
 #include "ops/transpose_layout_inference.h"
+#include "ops/unidirectional_lstm_layout_inference.h"
 
 #include <algorithm>
 #include <deque>
@@ -120,7 +121,7 @@ bool LayoutInferContext::IsVisited(const std::shared_ptr<vx::Operation>& op) con
 bool LayoutInferContext::IsReadyForInfer(
     const std::shared_ptr<vx::Operation>& op) const {
   for (const auto& tensor : op->impl()->InputsTensor()) {
-    if (!tensor->IsConstTensor() &&
+    if (!tensor->IsConstTensor() && tensor->GetId() != (uint32_t)-1 &&
         (tensor_pv_.end() == tensor_pv_.find(tensor))) {
       return false;
     }
@@ -261,6 +262,7 @@ std::vector<std::shared_ptr<vx::Tensor>> HandleLayoutInfer(
     REGIST_LAYOUT_INFERENCE(VSI_NN_OP_BATCH_NORM, BatchNorm);
     REGIST_LAYOUT_INFERENCE(VSI_NN_OP_PERMUTE, Transpose);
     REGIST_LAYOUT_INFERENCE(VSI_NN_OP_CONV3D, Conv3d);
+    REGIST_LAYOUT_INFERENCE(VSI_NN_OP_LSTM_OVXLIB, UnidirectionalLstm);
     REGIST_LOGICAL_LAYOUT_INFERENCE(VSI_NN_OP_LOGICAL_OPS);
     REGIST_REDUCE_LAYOUT_INFERENCE(VSI_NN_OP_REDUCE);
     // use default layout inference

src/tim/transform/ops/op_layout_inference.cc

@@ -78,19 +78,20 @@ std::shared_ptr<vx::Tensor> OpLayoutInfer::InsertPermute(
         MapAxis(perm->AsStdVec(), out_spec.quantization_.ChannelDim()));
   }
   auto out_tensor = context_->infer_graph_->CreateTensor(out_spec);
-  auto perm_op =
+  auto perm_op = context_->infer_graph_->CreateOperation<vx::ops::Transpose>(
-      context_->infer_graph_->CreateOperation<vx::ops::Transpose>(perm->AsStdVec());
+      perm->AsStdVec());
   (*perm_op).BindInput(input).BindOutput(out_tensor);
   return out_tensor;
 }
 
 std::vector<std::shared_ptr<vx::Tensor>> OpLayoutInfer::CreateOutputsTensor(
-      std::shared_ptr<IPermuteVector> required_pv) {
+    std::shared_ptr<IPermuteVector> required_pv) {
   std::vector<std::shared_ptr<vx::Tensor>> outputs_tensor;
 
   if (op_->impl()->OutputsTensor().size() > 1) {
     // todo(sven): potential bug here if node have multi-output and require layout inference
-    std::cout <<"warning at "<< __FUNCTION__ << ", #" << __LINE__ << std::endl;
+    std::cout << "warning at " << __FUNCTION__ << ", #" << __LINE__
+              << std::endl;
   }
 
   uint32_t i = 0;
@@ -206,8 +207,8 @@ OpLayoutInfer::AlignPermuteVectorForMutilInputs() {
       std::shared_ptr<vx::Tensor> perm_out;
       if (i_src->IsConstTensor()) {
         required_pv->IsAligned()
-            ? perm_out = context_->infer_graph_->CreateTensor(i_src->GetSpec(),
+            ? perm_out = context_->infer_graph_->CreateTensor(
-                                                              i_src->GetDataRef())
+                  i_src->GetSpec(), i_src->GetDataRef())
             : perm_out = PermuteConstTensor(i_src, required_pv);
       } else {
         auto final_pv =
@@ -265,25 +266,25 @@ OpLayoutInfer::AlignPermuteVectorForElementWise() {
   return required_pv;
 }
 
-
 void OpLayoutInfer::ReverseInputsPermuteVector() {
   for (const auto& i_src : op_->impl()->InputsTensor()) {
     std::shared_ptr<vx::Tensor> perm_out;
     std::shared_ptr<IPermuteVector> input_pv;
-    if (i_src->IsConstTensor()) {
+    if (i_src->GetId() != (uint32_t)-1) {
-      perm_out = context_->infer_graph_->CreateTensor(i_src->GetSpec(),
+      if (i_src->IsConstTensor()) {
-                                                      i_src->GetDataRef());
+        perm_out = context_->infer_graph_->CreateTensor(i_src->GetSpec(),
-      input_pv = MakeShared(i_src->GetShape().size());
+                                                        i_src->GetDataRef());
-    } else {
+        input_pv = MakeShared(i_src->GetShape().size());
-      perm_out = context_->GetMapedTensor(i_src);
+      } else {
-      input_pv = context_->GetPermuteVector(i_src);
+        perm_out = context_->GetMapedTensor(i_src);
-      if (!input_pv->IsAligned()) {
+        input_pv = context_->GetPermuteVector(i_src);
-        perm_out =
+        if (!input_pv->IsAligned()) {
-            InsertPermute(perm_out, input_pv->Reverse());
+          perm_out = InsertPermute(perm_out, input_pv->Reverse());
+        }
       }
+      context_->UpdateTensorMap(i_src, perm_out);
+      context_->SetPermuteVector(i_src, MakeShared(input_pv->Rank()));
     }
-    context_->UpdateTensorMap(i_src, perm_out);
-    context_->SetPermuteVector(i_src, MakeShared(input_pv->Rank()));
   }
 }
 

src/tim/transform/ops/unidirectional_lstm_layout_inference.h

@@ -0,0 +1,94 @@
+/****************************************************************************
+ *
+ *    Copyright (c) 2022 Vivante Corporation
+ *
+ *    Permission is hereby granted, free of charge, to any person obtaining a
+ *    copy of this software and associated documentation files (the "Software"),
+ *    to deal in the Software without restriction, including without limitation
+ *    the rights to use, copy, modify, merge, publish, distribute, sublicense,
+ *    and/or sell copies of the Software, and to permit persons to whom the
+ *    Software is furnished to do so, subject to the following conditions:
+ *
+ *    The above copyright notice and this permission notice shall be included in
+ *    all copies or substantial portions of the Software.
+ *
+ *    THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+ *    IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+ *    FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+ *    AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+ *    LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING
+ *    FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER
+ *    DEALINGS IN THE SOFTWARE.
+ *
+ *****************************************************************************/
+#ifndef TIM_LAYOUT_INFER_UNIDIRECTIONAL_LSTM_LAYOUT_INFERENCE_H_
+#define TIM_LAYOUT_INFER_UNIDIRECTIONAL_LSTM_LAYOUT_INFERENCE_H_
+
+#include "tim/vx/ops/reshape.h"
+#include "tim/vx/ops/nbg.h"
+#include "tim/vx/ops/transpose.h"
+#include "tim/vx/ops/batchnorm.h"
+#include "tim/vx/ops/clip.h"
+
+#include "ops/op_layout_inference.h"
+#include "permute_vector.h"
+#include "direct_map_op_impl.h"
+
+namespace tim {
+namespace transform {
+
+class UnidirectionalLstmLayoutInfer : public OpLayoutInfer {
+ public:
+  UnidirectionalLstmLayoutInfer(
+      const std::shared_ptr<vx::Operation> op,
+      std::shared_ptr<layout_inference_impl::LayoutInferContext>& context)
+      : OpLayoutInfer(op, context) {}
+
+  // reverse any applied permute on it's input tensor
+  void OnInputs(
+      std::vector<std::shared_ptr<vx::Tensor>>& next_tensors) override {
+    ReverseInputsPermuteVector();
+
+    auto cloned_op = op_->Clone(context_->infer_graph_);
+
+    for (const auto& i_src : op_->impl()->InputsTensor()) {
+      std::shared_ptr<vx::Tensor> infer_tensor;
+      std::shared_ptr<IPermuteVector> required_pv;
+      if ((i_src->IsConstTensor() &&
+              !(i_src->GetSpec().attr_ & vx::TensorAttribute::INPUT))) {
+        infer_tensor = context_->infer_graph_->CreateTensor(
+            i_src->GetSpec(), i_src->GetDataRef());
+        context_->UpdateTensorMap(i_src, infer_tensor);
+      }
+      if (i_src->GetId() == (uint32_t)-1) {
+        infer_tensor = context_->infer_graph_->CreateTensorPlaceHolder();
+        context_->UpdateTensorMap(i_src, infer_tensor);
+      }
+      required_pv = MakeShared(i_src->GetShape().size());
+      context_->SetPermuteVector(i_src, required_pv);
+    }
+
+
+    for (const auto& i_src : op_->impl()->InputsTensor()) {
+      (*cloned_op).BindInput(context_->GetMapedTensor(i_src));
+    }
+
+    std::vector<std::shared_ptr<IPermuteVector>> required_pv_lst;
+    for (auto out_tensor : op_->impl()->OutputsTensor()) {
+      required_pv_lst.push_back(MakeShared(out_tensor->GetShape().size()));
+    }
+    auto out_infer = CreateOutputsTensor(required_pv_lst);
+
+    (*cloned_op).BindOutputs(out_infer);
+    uint32_t i = 0;
+    for (auto out_tensor : op_->impl()->OutputsTensor()) {
+      context_->SetPermuteVector(out_tensor, required_pv_lst[i++]);
+      next_tensors.push_back(out_tensor);
+    }
+  }
+};
+
+}  // namespace transform
+}  // namespace tim
+
+#endif

src/tim/vx/ops/unidirectional_sequence_lstm_test.cc

@@ -144,7 +144,7 @@ TEST(LSTM_CELL, shape_in_2_cell_4_out_4_float32) {
 
     auto lstm_cell_op = g->CreateOperation<tim::vx::ops::UnidirectionalSequenceLstm>(
         0.0, 0.0, tim::vx::ops::UnidirectionalSequenceLstm::ActivationType::kTANH, 0.0, false,
-        tim::vx::ops::UnidirectionalSequenceLstm::kNONE, true);
+        tim::vx::ops::UnidirectionalSequenceLstm::kSIGMOID, true);
 
     (*lstm_cell_op)
         .BindInputs({