From 016927b57cb4073438d026932514d1f580cf2dc2 Mon Sep 17 00:00:00 2001
From: hong19860320 <9973393+hong19860320@users.noreply.github.com>
Date: Wed, 15 Apr 2020 16:38:18 +0800
Subject: [PATCH] [OPT] Add RK and MTK support for model optimize tool (#3410)

---
 lite/api/opt.cc                  | 12 ++++++------
 lite/api/opt_base.cc             | 11 ++++++++++-
 lite/api/python/pybind/pybind.cc |  1 +
 3 files changed, 17 insertions(+), 7 deletions(-)

diff --git a/lite/api/opt.cc b/lite/api/opt.cc
index db4f8b7c3c..a6ad7cff6f 100644
--- a/lite/api/opt.cc
+++ b/lite/api/opt.cc
@@ -108,17 +108,17 @@ std::vector<Place> ParserValidPlaces() {
       valid_places.emplace_back(Place{TARGET(kX86), PRECISION(kInt64)});
     } else if (target_repr == "npu") {
       valid_places.emplace_back(TARGET(kNPU));
-    } else if (target_repr == "apu") {
-      valid_places.emplace_back(
-          Place{TARGET(kAPU), PRECISION(kInt8), DATALAYOUT(kNCHW)});
     } else if (target_repr == "xpu") {
       valid_places.emplace_back(TARGET(kXPU));
+    } else if (target_repr == "mlu") {
+      valid_places.emplace_back(TARGET(kMLU));
     } else if (target_repr == "rknpu") {
       valid_places.emplace_back(TARGET(kRKNPU));
       valid_places.emplace_back(
           TARGET(kRKNPU), PRECISION(kInt8), DATALAYOUT(kNCHW));
-    } else if (target_repr == "mlu") {
-      valid_places.emplace_back(TARGET(kMLU));
+    } else if (target_repr == "apu") {
+      valid_places.emplace_back(
+          Place{TARGET(kAPU), PRECISION(kInt8), DATALAYOUT(kNCHW)});
     } else {
       LOG(FATAL) << lite::string_format(
           "Wrong target '%s' found, please check the command flag "
@@ -194,9 +194,9 @@ void PrintOpsInfo(std::set<std::string> valid_ops = {}) {
                                       "kOpenCL",
                                       "kFPGA",
                                       "kNPU",
-                                      "kAPU",
                                       "kXPU",
                                       "kRKNPU",
+                                      "kAPU",
                                       "kAny",
                                       "kUnk"};
   int maximum_optype_length = 0;
diff --git a/lite/api/opt_base.cc b/lite/api/opt_base.cc
index bd86f48624..14c1ca4a4e 100644
--- a/lite/api/opt_base.cc
+++ b/lite/api/opt_base.cc
@@ -63,6 +63,13 @@ void OptBase::SetValidPlaces(const std::string& valid_places) {
       valid_places_.emplace_back(TARGET(kNPU));
     } else if (target_repr == "xpu") {
       valid_places_.emplace_back(TARGET(kXPU));
+    } else if (target_repr == "rknpu") {
+      valid_places_.emplace_back(TARGET(kRKNPU));
+      valid_places_.emplace_back(
+          TARGET(kRKNPU), PRECISION(kInt8), DATALAYOUT(kNCHW));
+    } else if (target_repr == "apu") {
+      valid_places_.emplace_back(
+          Place{TARGET(kAPU), PRECISION(kInt8), DATALAYOUT(kNCHW)});
     } else {
       LOG(FATAL) << lite::string_format(
           "Wrong target '%s' found, please check the command flag "
@@ -183,7 +190,7 @@ void OptBase::PrintHelpInfo() {
       "        `set_param_file(param_file_path)`\n"
       "        `set_model_type(protobuf|naive_buffer)`\n"
       "        `set_optimize_out(output_optimize_model_dir)`\n"
-      "        `set_valid_places(arm|opencl|x86|npu|xpu)`\n"
+      "        `set_valid_places(arm|opencl|x86|npu|xpu|rknpu|apu)`\n"
       "        `run_optimize(false|true)`\n"
       "        `  ----fasle&true refer to whether to record ops info for "
       "tailoring lib, false by default`\n"
@@ -208,6 +215,8 @@ void OptBase::PrintOpsInfo(const std::set<std::string>& valid_ops) {
                                                      "kFPGA",
                                                      "kNPU",
                                                      "kXPU",
+                                                     "kRKNPU",
+                                                     "kAPU",
                                                      "kAny",
                                                      "kUnk"};
   // Get the lengh of the first column: maximum length of the op_type
diff --git a/lite/api/python/pybind/pybind.cc b/lite/api/python/pybind/pybind.cc
index b78f564e4a..06d1c607fd 100644
--- a/lite/api/python/pybind/pybind.cc
+++ b/lite/api/python/pybind/pybind.cc
@@ -183,6 +183,7 @@ void BindLitePlace(py::module *m) {
       .value("FPGA", TargetType::kFPGA)
       .value("NPU", TargetType::kNPU)
       .value("MLU", TargetType::kMLU)
+      .value("RKNPU", TargetType::kRKNPU)
       .value("APU", TargetType::kAPU)
       .value("Any", TargetType::kAny);
 
-- 
GitLab