Browse Source

!8583 add unique dynamic

From: @fangzehua
Reviewed-by: 
Signed-off-by:
tags/v1.1.0
mindspore-ci-bot Gitee 5 years ago
parent
commit
c09f2cb53f
3 changed files with 15 additions and 1 deletions
  1. +11
    -0
      mindspore/ccsrc/backend/kernel_compiler/cpu/unique_cpu_kernel.cc
  2. +1
    -0
      mindspore/ccsrc/backend/kernel_compiler/cpu/unique_cpu_kernel.h
  3. +3
    -1
      mindspore/ccsrc/backend/kernel_compiler/cpu/unique_with_pad_cpu_kernel.cc

+ 11
- 0
mindspore/ccsrc/backend/kernel_compiler/cpu/unique_cpu_kernel.cc View File

@@ -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;
}



+ 1
- 0
mindspore/ccsrc/backend/kernel_compiler/cpu/unique_cpu_kernel.h View File

@@ -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) {


+ 3
- 1
mindspore/ccsrc/backend/kernel_compiler/cpu/unique_with_pad_cpu_kernel.cc View File

@@ -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;


Loading…
Cancel
Save