scheduler add fitkernel

mindspore/lite/nnacl/constant_of_shape.h → mindspore/lite/nnacl/constant_of_shape_parameter.h

@@ -1,5 +1,5 @@
 /**
- * Copyright 2020 Huawei Technologies Co., Ltd
+ * Copyright 2020-2021 Huawei Technologies Co., Ltd
  *
  * Licensed under the Apache License, Version 2.0 (the "License");
  * you may not use this file except in compliance with the License.
@@ -13,14 +13,10 @@
  * See the License for the specific language governing permissions and
  * limitations under the License.
  */
-#ifndef MINDSPORE_LITE_NNACL_CONSTANT_OF_SHAPE_H_
-#define MINDSPORE_LITE_NNACL_CONSTANT_OF_SHAPE_H_
+#ifndef MINDSPORE_LITE_NNACL_CONSTANT_OF_SHAPE_PARAMETER_H_
+#define MINDSPORE_LITE_NNACL_CONSTANT_OF_SHAPE_PARAMETER_H_
 
-#ifdef ENABLE_NEON
-#include <arm_neon.h>
-#endif
 #include "nnacl/op_base.h"
-#include "nnacl/errorcode.h"
 
 typedef struct ConstantOfShapeParameter {
   OpParameter op_parameter_;
@@ -32,13 +28,4 @@ typedef struct ConstantOfShapeParameter {
   int element_size_;
 } ConstantOfShapeParameter;
 
-#ifdef __cplusplus
-extern "C" {
-#endif
-int ConstantOfShapeFp32(float *output, int start, int end, float value);
-int ConstantOfShapeInt32(int32_t *output, int start, int end, int32_t value);
-#ifdef __cplusplus
-}
-#endif
-
-#endif  // MINDSPORE_LITE_NNACL_CONSTANT_OF_SHAPE_H_
+#endif  // MINDSPORE_LITE_NNACL_CONSTANT_OF_SHAPE_PARAMETER_H_

mindspore/lite/nnacl/fp16/constant_of_shape_fp16.h

@@ -0,0 +1,41 @@
+/**
+ * Copyright 2020-2021 Huawei Technologies Co., Ltd
+ *
+ * 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 MINDSPORE_LITE_NNACL_FP16_CONSTANT_OF_SHAPE_FP16_H_
+#define MINDSPORE_LITE_NNACL_FP16_CONSTANT_OF_SHAPE_FP16_H_
+
+#ifdef ENABLE_NEON
+#include <arm_neon.h>
+#endif
+#include "nnacl/op_base.h"
+#include "nnacl/errorcode.h"
+#include "nnacl/constant_of_shape_parameter.h"
+
+#ifdef __cplusplus
+extern "C" {
+#endif
+#ifdef __cplusplus
+#ifdef ENABLE_NEON
+inline int ConstantOfShapeFp16(float16_t *output, int start, int end, float16_t value) {
+  for (int i = start; i < end; i++) {
+    output[i] = value;
+  }
+  return NNACL_OK;
+}
+#endif
+}
+#endif
+
+#endif  // MINDSPORE_LITE_NNACL_FP16_CONSTANT_OF_SHAPE_FP16_H_

mindspore/lite/nnacl/constant_of_shape.c → mindspore/lite/nnacl/fp32/constant_of_shape_fp32.h

@@ -1,5 +1,5 @@
 /**
- * Copyright 2020 Huawei Technologies Co., Ltd
+ * Copyright 2020-2021 Huawei Technologies Co., Ltd
  *
  * Licensed under the Apache License, Version 2.0 (the "License");
  * you may not use this file except in compliance with the License.
@@ -13,19 +13,33 @@
  * See the License for the specific language governing permissions and
  * limitations under the License.
  */
+#ifndef MINDSPORE_LITE_NNACL_FP32_CONSTANT_OF_SHAPE_FP32_H_
+#define MINDSPORE_LITE_NNACL_FP32_CONSTANT_OF_SHAPE_FP32_H_
+#include <memory.h>
+#include <float.h>
+#include "nnacl/op_base.h"
+#include "nnacl/errorcode.h"
+#include "nnacl/constant_of_shape_parameter.h"
 
-#include "nnacl/constant_of_shape.h"
-
-int ConstantOfShapeInt32(int32_t *output, int start, int end, int32_t value) {
+#ifdef __cplusplus
+extern "C" {
+#endif
+inline int ConstantOfShapeInt32(int32_t *output, int start, int end, int32_t value) {
   for (int i = start; i < end; i++) {
     output[i] = value;
   }
   return NNACL_OK;
 }
 
-int ConstantOfShapeFp32(float *output, int start, int end, float value) {
+inline int ConstantOfShapeFp32(float *output, int start, int end, float value) {
   for (int i = start; i < end; i++) {
     output[i] = value;
   }
   return NNACL_OK;
 }
+
+#ifdef __cplusplus
+}
+#endif
+
+#endif  // MINDSPORE_LITE_NNACL_FP32_CONSTANT_OF_SHAPE_FP32_H_

mindspore/lite/nnacl/infer/constant_of_shape_infer.h

@@ -17,7 +17,7 @@
 #define MINDSPORE_LITE_NNACL_CONSTANT_OF_SHAPE_INFER_H
 
 #include "nnacl/infer/common_infer.h"
-#include "nnacl/constant_of_shape.h"
+#include "nnacl/constant_of_shape_parameter.h"
 
 #ifdef __cplusplus
 extern "C" {

mindspore/lite/src/ops/populate/constant_of_shape_populate.cc

@@ -14,7 +14,7 @@
  * limitations under the License.
  */
 #include "src/ops/populate/populate_register.h"
-#include "nnacl/constant_of_shape.h"
+#include "nnacl/constant_of_shape_parameter.h"
 
 namespace mindspore::lite {
 namespace {

mindspore/lite/src/ops/populate/v0/constant_of_shape_populate_v0.cc

@@ -16,7 +16,7 @@
 
 #include "schema/model_v0_generated.h"
 #include "src/ops/populate/populate_register.h"
-#include "nnacl/constant_of_shape.h"
+#include "nnacl/constant_of_shape_parameter.h"
 
 namespace mindspore::lite {
 namespace {

mindspore/lite/src/runtime/kernel/arm/base/constant_of_shape.cc

@@ -52,6 +52,12 @@ int ConstantOfShapeCPUKernel::DoExecute(int task_id) {
       ConstantOfShapeInt32(reinterpret_cast<int32_t *>(output_ptr_), start, start + current_stride,
                            param_->value_.int32_value_);
       break;
+#ifdef ENABLE_NEON
+    case kNumberTypeFloat16:
+      ConstantOfShapeFp16(reinterpret_cast<float16_t *>(output_ptr_), start, start + current_stride,
+                          param_->value_.f32_value_);
+      break;
+#endif
     default:
       MS_LOG(ERROR) << "Invalid datatype in ConstantOfShapeRun";
       return RET_ERROR;

mindspore/lite/src/runtime/kernel/arm/base/constant_of_shape.h

@@ -20,7 +20,9 @@
 #include "include/errorcode.h"
 #include "src/lite_kernel.h"
 #include "include/context.h"
-#include "nnacl/constant_of_shape.h"
+#include "nnacl/constant_of_shape_parameter.h"
+#include "nnacl/fp32/constant_of_shape_fp32.h"
+#include "nnacl/fp16/constant_of_shape_fp16.h"
 
 namespace mindspore::kernel {
 class ConstantOfShapeCPUKernel : public LiteKernel {

mindspore/lite/src/scheduler.cc

@@ -387,6 +387,38 @@ int Scheduler::ScheduleSubGraphToKernels(size_t subgraph_index, std::vector<kern
   return RET_OK;
 }
 
+bool Scheduler::KernelFitCurrentSubGraph(const kernel::SubGraphType subgraph_type, const kernel::LiteKernel &kernel) {
+  switch (subgraph_type) {
+    case kernel::SubGraphType::kNotSubGraph:
+    case kernel::SubGraphType::kApuSubGraph:
+      return false;
+    case kernel::SubGraphType::kGpuSubGraph:
+      return kernel.desc().arch == kGPU;
+    case kernel::SubGraphType::kNpuSubGraph:
+      return kernel.desc().arch == kNPU;
+    case kernel::SubGraphType::kCpuFP16SubGraph: {
+      auto desc = kernel.desc();
+      if (desc.arch != kCPU) {
+        return false;
+      }
+      return (desc.data_type == kNumberTypeFloat16 || desc.data_type == kNumberTypeInt32 ||
+              desc.data_type == kNumberTypeInt || desc.data_type == kNumberTypeBool);
+    }
+    case kernel::SubGraphType::kCpuFP32SubGraph: {
+      auto desc = kernel.desc();
+      if (desc.arch != kCPU) {
+        return false;
+      }
+      return (desc.data_type == kNumberTypeFloat32 || desc.data_type == kNumberTypeFloat ||
+              desc.data_type == kNumberTypeInt8 || desc.data_type == kNumberTypeInt ||
+              desc.data_type == kNumberTypeInt32 || desc.data_type == kNumberTypeInt64 ||
+              desc.data_type == kNumberTypeUInt8 || desc.data_type == kNumberTypeBool);
+    }
+    default:
+      return false;
+  }
+}
+
 std::vector<kernel::LiteKernel *> Scheduler::FindAllSubGraphKernels(
   kernel::LiteKernel *head_kernel, std::map<const kernel::LiteKernel *, bool> *sinked_kernel_map) {
   MS_ASSERT(head_kernel != nullptr);

mindspore/lite/src/scheduler.h

@@ -68,6 +68,8 @@ class Scheduler {
 
   bool MergeOpIsReady(const kernel::LiteKernel *kernel, std::map<const kernel::LiteKernel *, bool> is_kernel_finish);
 
+  bool KernelFitCurrentSubGraph(const kernel::SubGraphType subgraph_type, const kernel::LiteKernel &kernel);
+
   std::vector<kernel::LiteKernel *> FindAllSubGraphKernels(
     kernel::LiteKernel *head_kernel, std::map<const kernel::LiteKernel *, bool> *sinked_kernel_map);