From 60167dae5b3619a6988db710502c92a01d946197 Mon Sep 17 00:00:00 2001
From: fangzehua <fangzehua1@huawei.com>
Date: Fri, 13 Nov 2020 17:29:12 +0800
Subject: [PATCH] add unique dynamic

---
 .../backend/kernel_compiler/cpu/unique_cpu_kernel.cc  | 11 +++++++++++
 .../backend/kernel_compiler/cpu/unique_cpu_kernel.h   |  1 +
 .../kernel_compiler/cpu/unique_with_pad_cpu_kernel.cc |  4 +++-
 3 files changed, 15 insertions(+), 1 deletion(-)

diff --git a/mindspore/ccsrc/backend/kernel_compiler/cpu/unique_cpu_kernel.cc b/mindspore/ccsrc/backend/kernel_compiler/cpu/unique_cpu_kernel.cc
index e367bbba3d..cd030c7180 100644
--- a/mindspore/ccsrc/backend/kernel_compiler/cpu/unique_cpu_kernel.cc
+++ b/mindspore/ccsrc/backend/kernel_compiler/cpu/unique_cpu_kernel.cc
@@ -22,6 +22,7 @@ namespace kernel {
 const size_t kUseBucketUniqueSize = 100000;
 const size_t kUniqueThreadNum = 23;
 void UniqueCPUKernel::InitKernel(const CNodePtr &kernel_node) {
+  node_ = kernel_node;
   CheckParam(kernel_node);
   auto input_shape = AnfAlgo::GetPrevNodeOutputInferShape(kernel_node, 0);
   input_size_ = input_shape[0];
@@ -45,6 +46,16 @@ bool UniqueCPUKernel::Launch(const std::vector<kernel::AddressPtr> &inputs,
   } else if (dtype_ == kNumberTypeFloat32) {
     LaunchKernel<float, int>(inputs, workspace, outputs);
   }
+  if (node_ != nullptr) {
+    std::vector<size_t> out_shape;
+    out_shape.emplace_back(output_size_);
+    std::vector<TypeId> dtypes;
+    size_t output_num = AnfAlgo::GetOutputTensorNum(node_);
+    for (size_t i = 0; i < output_num; i++) {
+      dtypes.push_back(AnfAlgo::GetOutputInferDataType(node_, i));
+    }
+    AnfAlgo::SetOutputInferTypeAndShape(dtypes, {out_shape, AnfAlgo::GetOutputInferShape(node_, 1)}, node_.get());
+  }
   return true;
 }
 
diff --git a/mindspore/ccsrc/backend/kernel_compiler/cpu/unique_cpu_kernel.h b/mindspore/ccsrc/backend/kernel_compiler/cpu/unique_cpu_kernel.h
index 339c87131b..be46c08537 100644
--- a/mindspore/ccsrc/backend/kernel_compiler/cpu/unique_cpu_kernel.h
+++ b/mindspore/ccsrc/backend/kernel_compiler/cpu/unique_cpu_kernel.h
@@ -59,6 +59,7 @@ class UniqueCPUKernel : public CPUKernel {
   size_t input_size_{0};
   TypeId dtype_{kTypeUnknown};
   size_t output_size_{0};
+  CNodePtr node_ = nullptr;
 
   template <typename DataType>
   static size_t BucketId(DataType data, size_t bucket_num) {
diff --git a/mindspore/ccsrc/backend/kernel_compiler/cpu/unique_with_pad_cpu_kernel.cc b/mindspore/ccsrc/backend/kernel_compiler/cpu/unique_with_pad_cpu_kernel.cc
index dbefe6ed90..96cc8715ae 100644
--- a/mindspore/ccsrc/backend/kernel_compiler/cpu/unique_with_pad_cpu_kernel.cc
+++ b/mindspore/ccsrc/backend/kernel_compiler/cpu/unique_with_pad_cpu_kernel.cc
@@ -22,12 +22,14 @@ namespace kernel {
 bool UniqueWithPadCPUKernel::Launch(const std::vector<kernel::AddressPtr> &inputs,
                                     const std::vector<kernel::AddressPtr> &workspace,
                                     const std::vector<kernel::AddressPtr> &outputs) {
-  UniqueCPUKernel::Launch(inputs, workspace, outputs);
   if (dtype_ == kNumberTypeInt32) {
+    UniqueCPUKernel::LaunchKernel<int, int>(inputs, workspace, outputs);
     PadOutput<int>(inputs, outputs);
   } else if (dtype_ == kNumberTypeInt64) {
+    UniqueCPUKernel::LaunchKernel<int64_t, int>(inputs, workspace, outputs);
     PadOutput<int64_t>(inputs, outputs);
   } else if (dtype_ == kNumberTypeFloat32) {
+    UniqueCPUKernel::LaunchKernel<float, int>(inputs, workspace, outputs);
     PadOutput<float>(inputs, outputs);
   }
   return true;