From ad0e8481f9f9459e40fc543658c7e8186aebaef9 Mon Sep 17 00:00:00 2001
From: jiaopu <jiaopu@cambricon.com>
Date: Fri, 10 Apr 2020 17:55:53 +0800
Subject: [PATCH] fix CHECK problem

---
 lite/core/mir/mlu_postprocess_pass.cc | 23 +++++++++++++++++------
 1 file changed, 17 insertions(+), 6 deletions(-)

diff --git a/lite/core/mir/mlu_postprocess_pass.cc b/lite/core/mir/mlu_postprocess_pass.cc
index 23303d6f60..3b1a03d364 100644
--- a/lite/core/mir/mlu_postprocess_pass.cc
+++ b/lite/core/mir/mlu_postprocess_pass.cc
@@ -541,12 +541,23 @@ void MLUPostprocessPass::ModifyInputOutputDataType(SSAGraph* graph) {
               << "MLU subgraph unexpected persistent input type!";
           out_node->AsArg().type = LiteType::GetTensorTy(
               TARGET(kMLU), PRECISION(kAny), DATALAYOUT(kNHWC));
+        } else if (out_node_type->precision() == PRECISION(kAny) &&
+                   out_node->outlinks.empty()) {
+          out_node->AsArg().is_persist = true;
+          out_node->AsArg().type = LiteType::GetTensorTy(
+              TARGET(kMLU), PRECISION(kAny), DATALAYOUT(kNHWC));
         } else {
           CHECK(out_node_type->precision() == PRECISION(kFloat))
               << "MLU subgraph unexpected common output type!";
           out_node->AsArg().type = LiteType::GetTensorTy(
               TARGET(kHost), PRECISION(kFloat), DATALAYOUT(kNCHW));
         }
+        const auto target = out_node->AsArg().type->target();
+        const auto precision = out_node->AsArg().type->precision();
+        const auto layout = out_node->AsArg().type->layout();
+        VLOG(4) << "arg name: " << out_node->AsArg().name
+                << " type: " << TargetToStr(target) << ", "
+                << PrecisionToStr(precision) << ", " << DataLayoutToStr(layout);
       }
     }
   }
@@ -597,13 +608,13 @@ void MLUPostprocessPass::ModifyLayout(SSAGraph* graph) {
 }
 
 void MLUPostprocessPass::Apply(const std::unique_ptr<SSAGraph>& graph) {
-  // currently for non-persistent input and output args, mlu subgraph op
-  // only support float16/float32 data type
+// currently for non-persistent input and output args, mlu subgraph op
+// only support float16/float32 data type
 
-  // in two situations as folllows:
-  // 1: feed->arg_in->subgraph->... 2: ...->subgraph->arg_out->fetch;
-  // arg_in and arg_out are assumed to be NHWC which user should be aware of.
-  // Thus here we change these args' layout to NHWC
+// in two situations as folllows:
+// 1: feed->arg_in->subgraph->... 2: ...->subgraph->arg_out->fetch;
+// arg_in and arg_out are assumed to be NHWC which user should be aware of.
+// Thus here we change these args' layout to NHWC
 #ifdef LITE_WITH_MLU
   ModifyInputOutputDataType(graph.get());
 
-- 
GitLab