!11618 [MS_LITE] control flow encoder

From: @YeFeng_24
Reviewed-by: 
Signed-off-by:

mindspore/lite/src/ops/tensorlist_getitem.cc

@@ -120,8 +120,17 @@ int TensorListGetItem::InferShape(std::vector<lite::Tensor *> inputs_, std::vect
   if (!infer_flag()) {
     return RET_INFER_INVALID;
   }
-  auto input0 = reinterpret_cast<TensorList *>(inputs_[0]);
-  auto get_index = inputs_[1];
+  MS_ASSERT(inputs_.size() >= 3);
+  MS_ASSERT(inputs_.at(0) != nullptr);
+  MS_ASSERT(inputs_.at(1) != nullptr);
+  MS_ASSERT(inputs_.at(2) != nullptr);
+  auto input0 = reinterpret_cast<TensorList *>(inputs_.at(0));
+  if (input0->tensors_data_type() != GetElementDType()) {
+    MS_LOG(ERROR) << "op dtype: " << GetElementDType()
+                  << " is not equal in_tensor[0] dtype: " << input0->tensors_data_type();
+    return RET_ERROR;
+  }
+  auto get_index = inputs_.at(1);
   MS_ASSERT(get_index != nullptr);
   if (get_index->ElementsNum() != 1) {
     MS_LOG(ERROR) << "get_index->ElementsNum():" << get_index->ElementsNum() << " must be equal to 1!";
@@ -178,6 +187,7 @@ int TensorListGetItem::InferShape(std::vector<lite::Tensor *> inputs_, std::vect
     output->set_data_type(GetElementDType());
     output->set_shape(element_shape_);
   }
+  output->set_format(input0->GetTensor(index_)->format());
   return RET_OK;
 }
 }  // namespace lite

mindspore/lite/src/runtime/kernel/arm/fp32/tensorlist_getitem_fp32.cc

@@ -29,35 +29,38 @@ using mindspore::schema::PrimitiveType_TensorListGetItem;
 namespace mindspore::kernel {
 
 int TensorListGetItemCPUKernel::Init() {
-  auto input0 = reinterpret_cast<lite::TensorList *>(in_tensors_[0]);
+  MS_ASSERT(in_tensors_.size() >= 2);
+  MS_ASSERT(in_tensors_.at(0) != nullptr);
+  auto input0 = reinterpret_cast<lite::TensorList *>(in_tensors_.at(0));
   if (dtype_ != input0->tensors_data_type()) {
-    MS_LOG(ERROR) << "op dtype:" << dtype_ << " is not equal in_tensors[0] dtype:" << input0->tensors_data_type();
+    MS_LOG(ERROR) << "op dtype: " << dtype_ << " is not equal in_tensor[0] dtype: " << input0->tensors_data_type();
     return RET_ERROR;
   }
-  if (in_tensors_[1]->ElementsNum() != 1) {
-    MS_LOG(ERROR) << "in_tensors_[1]->ElementsNum():" << in_tensors_[1]->ElementsNum() << " must be equal to 1!";
-    return RET_ERROR;
-  }
-  index_ = reinterpret_cast<int *>(in_tensors_[1]->data_c())[0];
+  return RET_OK;
+}
+
+int TensorListGetItemCPUKernel::Run() {
+  MS_ASSERT(in_tensors_.size() >= 2);
+  MS_ASSERT(in_tensors_.at(0) != nullptr);
+  MS_ASSERT(in_tensors_.at(1) != nullptr);
+  MS_ASSERT(out_tensors_.at(0) != nullptr);
+  auto input0 = reinterpret_cast<lite::TensorList *>(in_tensors_.at(0));
+  MS_ASSERT(in_tensors_.at(1)->data_c() != nullptr);
+  index_ = reinterpret_cast<int *>(in_tensors_.at(1)->data_c())[0];
   int dim0 = input0->ElementsNum() - 1;
   if (index_ < 0 || index_ > dim0) {
     MS_LOG(ERROR) << "index tensor:[" << index_ << "] must be in [0, " << dim0 << "]!";
     return RET_ERROR;
   }
-  return RET_OK;
-}
-
-int TensorListGetItemCPUKernel::Run() {
-  auto input0 = reinterpret_cast<lite::TensorList *>(in_tensors_[0]);
   auto src_ptr = input0->GetTensor(index_);
   MS_ASSERT(src_ptr != nullptr);
   if (src_ptr->data_type() != kTypeUnknown) {
-    if (src_ptr->ElementsNum() != out_tensors_[0]->ElementsNum()) {
+    if (src_ptr->ElementsNum() != out_tensors_.at(0)->ElementsNum()) {
       MS_LOG(ERROR) << "src_ptr->ElementsNum():" << src_ptr->ElementsNum()
-                    << " must be equal to out_tensors_[0]->ElementsNum():" << out_tensors_[0]->ElementsNum();
+                    << " must be equal to out_tensors_[0]->ElementsNum():" << out_tensors_.at(0)->ElementsNum();
       return RET_ERROR;
     }
-    auto status = lite::Tensor::CopyTensorData(*src_ptr, out_tensors_[0]);
+    auto status = lite::Tensor::CopyTensorData(*src_ptr, out_tensors_.at(0));
     if (status == RET_ERROR) {
       MS_LOG(ERROR) << "copy tensor data failed!";
       return RET_ERROR;
@@ -65,19 +68,17 @@ int TensorListGetItemCPUKernel::Run() {
   } else {
     // reset 0 and dtype = dtype_
     // TODO(DT_VARIANT): dtype = DT_VARIANT is not handle
-    memset(out_tensors_[0]->MutableData(), 0, out_tensors_[0]->Size());
+    auto out_data = out_tensors_[0]->MutableData();
+    if (out_data == nullptr) {
+      MS_LOG(ERROR) << "data of out_tensors_[0] is nullptr";
+      return RET_ERROR;
+    }
+    memset(out_data, 0, out_tensors_[0]->Size());
   }
   return RET_OK;
 }
 
-int TensorListGetItemCPUKernel::ReSize() {
-  auto ret = this->Init();
-  if (ret != RET_OK) {
-    MS_LOG(ERROR) << "Init kernel failed!";
-    return ret;
-  }
-  return RET_OK;
-}
+int TensorListGetItemCPUKernel::ReSize() { return RET_OK; }
 
 kernel::LiteKernel *CpuTensorListGetItemFp32KernelCreator(const std::vector<lite::Tensor *> &inputs,
                                                           const std::vector<lite::Tensor *> &outputs,

mindspore/lite/src/runtime/kernel/arm/fp32/tensorlist_setitem_fp32.cc

@@ -30,20 +30,6 @@ namespace mindspore::kernel {
 
 int TensorListSetItemCPUKernel::Init() { return RET_OK; }
 
-int TensorListSetItemCPUKernel::IncrementOutputSize(int origin_size) {
-  output0_ = reinterpret_cast<lite::TensorList *>(out_tensors_[0]);
-  int new_tensors_size = origin_size + 1;
-  output0_->set_shape({new_tensors_size});
-  std::vector<std::vector<int>> out_shape;
-  out_shape.resize(new_tensors_size, in_tensors_[2]->shape());
-  auto ret = output0_->MallocTensorListData(in_tensors_[2]->data_type(), out_shape);
-  if (ret != RET_OK) {
-    MS_LOG(ERROR) << "increment output size malloc tensorlist data error";
-    return ret;
-  }
-  return RET_OK;
-}
-
 int TensorListSetItemCPUKernel::Run() {
   input0_ = reinterpret_cast<lite::TensorList *>(in_tensors_[0]);
   if (dtype_ != kTypeUnknown && dtype_ != input0_->tensors_data_type()) {
@@ -61,10 +47,8 @@ int TensorListSetItemCPUKernel::Run() {
   }
   index_ = reinterpret_cast<int *>(in_tensors_[1]->data_c())[0];
   if (index_ < 0 || index_ > dim0) {
-    if (IncrementOutputSize(output0_->shape()[0]) != RET_OK) {
-      MS_LOG(ERROR) << "Resizeoutput Error ,index tensor:[" << index_ << "] must be in [0, " << dim0 << "]!";
-      return RET_ERROR;
-    }
+    MS_LOG(ERROR) << "index tensor:[" << index_ << "] must be in [0, " << dim0 << "]!";
+    return RET_ERROR;
   }
   input2_ = in_tensors_[2];
   MS_ASSERT(input2_ != nullptr);
@@ -73,13 +57,6 @@ int TensorListSetItemCPUKernel::Run() {
   }
   output0_ = reinterpret_cast<lite::TensorList *>(out_tensors_[0]);
   MS_ASSERT(output0_ != nullptr);
-  // new loop count
-  if (output0_->ElementsNum() != static_cast<int>(output0_->tensors().size()) && output0_->tensors().empty()) {
-    if (IncrementOutputSize(0) != RET_OK) {
-      MS_LOG(ERROR) << "Resizeoutput Error!";
-      return RET_ERROR;
-    }
-  }
   // copy each tensor in tensors_
   for (int i = 0; i < output0_->ElementsNum(); ++i) {
     if (i == index_) {
@@ -115,6 +92,10 @@ int TensorListSetItemCPUKernel::Run() {
       }
 
       if (src->data_type() != kTypeUnknown) {
+        if (src->Size() != dst->Size()) {
+          MS_LOG(ERROR) << "src->Size():" << src->Size() << " must be equal to dst->Size():" << dst->Size();
+          return RET_ERROR;
+        }
         auto ret = lite::Tensor::CopyTensorData(*src, dst);
         if (ret != RET_OK) {
           MS_LOG(ERROR) << "CopyTensorData[" << i << "] is failed!";

mindspore/lite/src/runtime/kernel/arm/fp32/tensorlist_setitem_fp32.h

@@ -36,7 +36,6 @@ class TensorListSetItemCPUKernel : public LiteKernel {
   int Init() override;
   int ReSize() override;
   int Run() override;
-  int IncrementOutputSize(int origin_size);
 
  private:
   lite::TensorList *input0_ = nullptr;

mindspore/lite/test/models_tf.cfg

@@ -3,4 +3,5 @@ decoder_step_201217_modified.pb 5
 unet_model_reconstruct.pb 1;1,256,256,3
 encoder_201228.pb 3;1:1,22:1
 female_model_step2_int16_noiseout.pb 66
-fasterrcnn_crop.pb 1;236,190,3
+encoder_0111_control_flow.pb 4;1:1,44:1:1
+encoder_0111.pb 4;1:1,44:1:1

mindspore/lite/test/models_tflite_posttraining.cfg

@@ -1,3 +1,3 @@
 mobilenet.tflite 0.5
-transformer_20200831_encoder_fp32.tflite 69
+transformer_20200831_encoder_fp32.tflite 70
 transformer_20200831_decoder_fp32.tflite 35

mindspore/lite/test/models_with_several_inputs_or_without_outputs.cfg

@@ -35,4 +35,4 @@ ml_video_edit_img_segment_adaptise.pb;2
 ml_video_edit_img_segment_adaptise_pb2tflite.tflite;2
 ml_video_edit_video_segment_gauss_adaptis_part2.pb;2
 ml_video_edit_video_segment_gauss_adaptis_part2_pb2tflite.tflite;2
-tiny-yolov3-11.onnx;2;1,416,416,3:1,2
+# tiny-yolov3-11.onnx;2;1,416,416,3:1,2