forked from mindspore-Ecosystem/mindspore
sync lstm ops code from master to r0.3
This commit is contained in:
parent
e5c45bd339
commit
ba39d53c22
|
@ -85,7 +85,7 @@ bool IsInputFormatDtypeMatched(const KernelAttr &kernel_attr, const std::vector<
|
|||
const std::vector<TypeId> &input_types,
|
||||
const std::vector<size_t> &input_not_cnode_indexes) {
|
||||
if (kernel_attr.GetInputSize() != input_types.size()) {
|
||||
MS_LOG(ERROR) << "required input num:" << kernel_attr.GetInputSize() << ", actual input num:" << input_types.size();
|
||||
MS_LOG(DEBUG) << "required input num:" << kernel_attr.GetInputSize() << ", actual input num:" << input_types.size();
|
||||
return false;
|
||||
}
|
||||
auto input_num = input_types.size();
|
||||
|
@ -109,6 +109,21 @@ bool IsInputFormatDtypeMatched(const KernelAttr &kernel_attr, const std::vector<
|
|||
}
|
||||
return true;
|
||||
}
|
||||
|
||||
void ExpandKernelAttr(const CNodePtr &kernel_node, KernelAttr *kernel_attr) {
|
||||
MS_EXCEPTION_IF_NULL(kernel_attr);
|
||||
TypeId input_dtype = kernel_attr->GetInputAttr(0).first;
|
||||
size_t input_num = AnfAlgo::GetInputTensorNum(kernel_node);
|
||||
for (size_t i = 1; i < input_num; ++i) {
|
||||
kernel_attr->AddInputAttr(input_dtype);
|
||||
}
|
||||
|
||||
TypeId output_dtype = kernel_attr->GetOutputAttr(0).first;
|
||||
size_t output_num = AnfAlgo::GetOutputTensorNum(kernel_node);
|
||||
for (size_t i = 1; i < output_num; ++i) {
|
||||
kernel_attr->AddOutputAttr(output_dtype);
|
||||
}
|
||||
}
|
||||
} // namespace
|
||||
|
||||
void SetKernelInfo(const CNodePtr &kernel_node) {
|
||||
|
@ -125,12 +140,16 @@ void SetKernelInfo(const CNodePtr &kernel_node) {
|
|||
kernel::CPUKernelFactory::GetInstance().GetSupportedKernelAttrList(AnfAlgo::GetCNodeName(kernel_node));
|
||||
|
||||
for (size_t index = 0; index < kernel_attrs.size(); ++index) {
|
||||
if (IsInputFormatDtypeMatched(kernel_attrs[index], input_formats, input_types, input_not_cnode_indexes)) {
|
||||
auto kernel_attr = kernel_attrs[index];
|
||||
if (kernel_attr.GetAllSame()) {
|
||||
ExpandKernelAttr(kernel_node, &kernel_attr);
|
||||
}
|
||||
if (IsInputFormatDtypeMatched(kernel_attr, input_formats, input_types, input_not_cnode_indexes)) {
|
||||
MS_LOG(INFO) << "Input format and dtype is matched, index: " << index;
|
||||
GetOutputFormatsAndDtypes(kernel_node, kernel_attrs[index], &output_formats, &output_types);
|
||||
UpdatePrevNotCNodeFormatDtype(kernel_attrs[index], input_not_cnode_indexes, kernel_node);
|
||||
GetOutputFormatsAndDtypes(kernel_node, kernel_attr, &output_formats, &output_types);
|
||||
UpdatePrevNotCNodeFormatDtype(kernel_attr, input_not_cnode_indexes, kernel_node);
|
||||
for (auto &input_index : input_not_cnode_indexes) {
|
||||
input_types[input_index] = kernel_attrs[index].GetInputAttr(input_index).first;
|
||||
input_types[input_index] = kernel_attr.GetInputAttr(input_index).first;
|
||||
}
|
||||
break;
|
||||
}
|
||||
|
|
|
@ -46,8 +46,14 @@ class KernelAttr {
|
|||
return *this;
|
||||
}
|
||||
|
||||
KernelAttr &SetAllSameAttr(bool all_same) {
|
||||
all_same_ = all_same;
|
||||
return *this;
|
||||
}
|
||||
|
||||
const DataType &GetInputAttr(const size_t index) const { return input_type_[index]; }
|
||||
const DataType &GetOutputAttr(const size_t index) const { return output_type_[index]; }
|
||||
bool GetAllSame() const { return all_same_; }
|
||||
|
||||
size_t GetInputSize() const { return input_type_.size(); }
|
||||
size_t GetOutputSize() const { return output_type_.size(); }
|
||||
|
@ -55,6 +61,7 @@ class KernelAttr {
|
|||
private:
|
||||
std::vector<DataType> input_type_;
|
||||
std::vector<DataType> output_type_;
|
||||
bool all_same_;
|
||||
};
|
||||
} // namespace cpu
|
||||
} // namespace device
|
||||
|
|
|
@ -0,0 +1,66 @@
|
|||
/**
|
||||
* Copyright 2020 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.
|
||||
*/
|
||||
|
||||
#include "kernel/cpu/addn_cpu_kernel.h"
|
||||
#include "device/cpu/cpu_device_address.h"
|
||||
#include "ir/primitive.h"
|
||||
|
||||
namespace mindspore {
|
||||
namespace kernel {
|
||||
void AddNCPUKernel::InitKernel(const CNodePtr &kernel_node) {
|
||||
CheckParam(kernel_node);
|
||||
input_num_ = AnfAlgo::GetInputTensorNum(kernel_node);
|
||||
output_shape_ = AnfAlgo::GetOutputInferShape(kernel_node, 0);
|
||||
CPUKernelUtils::ExpandDimsTo4(&output_shape_);
|
||||
}
|
||||
|
||||
bool AddNCPUKernel::Launch(const std::vector<kernel::AddressPtr> &inputs,
|
||||
const std::vector<kernel::AddressPtr> & /*workspace*/,
|
||||
const std::vector<kernel::AddressPtr> &outputs) {
|
||||
auto output_addr = reinterpret_cast<float *>(outputs[0]->addr);
|
||||
|
||||
for (size_t i = 0; i < output_shape_[0]; ++i) {
|
||||
for (size_t j = 0; j < output_shape_[1]; ++j) {
|
||||
for (size_t k = 0; k < output_shape_[2]; ++k) {
|
||||
for (size_t m = 0; m < output_shape_[3]; ++m) {
|
||||
auto offset = CPUKernelUtils::CalcOffset(output_shape_, i, j, k, m);
|
||||
float sum = 0;
|
||||
for (size_t index = 0; index < input_num_; ++index) {
|
||||
auto input_addr = reinterpret_cast<float *>(inputs[index]->addr);
|
||||
sum += input_addr[offset];
|
||||
}
|
||||
output_addr[offset] = sum;
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
return true;
|
||||
}
|
||||
|
||||
void AddNCPUKernel::CheckParam(const CNodePtr &kernel_node) {
|
||||
auto input_shape = AnfAlgo::GetPrevNodeOutputInferShape(kernel_node, 0);
|
||||
if (input_shape.size() > 4) {
|
||||
MS_LOG(EXCEPTION) << "Input dims is " << input_shape.size() << ", but AddNCPUKernel olny support 4d or lower.";
|
||||
}
|
||||
|
||||
size_t output_num = AnfAlgo::GetOutputTensorNum(kernel_node);
|
||||
if (output_num != 1) {
|
||||
MS_LOG(EXCEPTION) << "Output number is " << output_num << ", but AddNCPUKernel needs 1 output.";
|
||||
}
|
||||
}
|
||||
} // namespace kernel
|
||||
} // namespace mindspore
|
|
@ -0,0 +1,48 @@
|
|||
/**
|
||||
* Copyright 2020 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_CCSRC_KERNEL_CPU_ADDN_CPU_KERNEL_H_
|
||||
#define MINDSPORE_CCSRC_KERNEL_CPU_ADDN_CPU_KERNEL_H_
|
||||
#include <vector>
|
||||
#include <memory>
|
||||
#include "kernel/cpu/cpu_kernel.h"
|
||||
#include "kernel/cpu/cpu_kernel_factory.h"
|
||||
|
||||
namespace mindspore {
|
||||
namespace kernel {
|
||||
class AddNCPUKernel : public CPUKernel {
|
||||
public:
|
||||
AddNCPUKernel() : input_num_(0) {}
|
||||
~AddNCPUKernel() override = default;
|
||||
|
||||
void InitKernel(const CNodePtr &kernel_node) override;
|
||||
|
||||
bool Launch(const std::vector<AddressPtr> &inputs, const std::vector<AddressPtr> &workspace,
|
||||
const std::vector<AddressPtr> &outputs) override;
|
||||
|
||||
private:
|
||||
void CheckParam(const CNodePtr &kernel_node);
|
||||
size_t input_num_;
|
||||
std::vector<size_t> output_shape_;
|
||||
};
|
||||
|
||||
MS_REG_CPU_KERNEL(AddN,
|
||||
KernelAttr().SetAllSameAttr(true).AddInputAttr(kNumberTypeFloat32).AddOutputAttr(kNumberTypeFloat32),
|
||||
AddNCPUKernel);
|
||||
} // namespace kernel
|
||||
} // namespace mindspore
|
||||
|
||||
#endif // MINDSPORE_CCSRC_KERNEL_CPU_ADDN_CPU_KERNEL_H_
|
|
@ -45,6 +45,7 @@ bool ConcatCPUKernel::Launch(const std::vector<kernel::AddressPtr> &inputs,
|
|||
const std::vector<kernel::AddressPtr> & /*workspace*/,
|
||||
const std::vector<kernel::AddressPtr> &outputs) {
|
||||
auto output_addr = reinterpret_cast<float *>(outputs[0]->addr);
|
||||
auto buff_size = outputs[0]->size;
|
||||
size_t dim0 = output_shape_[0];
|
||||
size_t dim1 = output_shape_[1];
|
||||
size_t dim2 = output_shape_[2];
|
||||
|
@ -53,28 +54,28 @@ bool ConcatCPUKernel::Launch(const std::vector<kernel::AddressPtr> &inputs,
|
|||
for (size_t i = 0; i < dim0; ++i) {
|
||||
for (size_t j = 0; j < dim1; ++j) {
|
||||
for (size_t k = 0; k < dim2; ++k) {
|
||||
CopyDataToOutput(inputs, i, j, k, &output_addr);
|
||||
CopyDataToOutput(inputs, i, j, k, &output_addr, &buff_size);
|
||||
}
|
||||
}
|
||||
}
|
||||
} else if (axis_ == 2) {
|
||||
for (size_t i = 0; i < dim0; ++i) {
|
||||
for (size_t j = 0; j < dim1; ++j) {
|
||||
CopyDataToOutput(inputs, i, j, 0, &output_addr);
|
||||
CopyDataToOutput(inputs, i, j, 0, &output_addr, &buff_size);
|
||||
}
|
||||
}
|
||||
} else if (axis_ == 1) {
|
||||
for (size_t i = 0; i < dim0; ++i) {
|
||||
CopyDataToOutput(inputs, i, 0, 0, &output_addr);
|
||||
CopyDataToOutput(inputs, i, 0, 0, &output_addr, &buff_size);
|
||||
}
|
||||
} else if (axis_ == 0) {
|
||||
CopyDataToOutput(inputs, 0, 0, 0, &output_addr);
|
||||
CopyDataToOutput(inputs, 0, 0, 0, &output_addr, &buff_size);
|
||||
}
|
||||
return true;
|
||||
}
|
||||
|
||||
void ConcatCPUKernel::CopyDataToOutput(const std::vector<kernel::AddressPtr> &inputs, size_t dim0, size_t dim1,
|
||||
size_t dim2, float **output_addr) {
|
||||
size_t dim2, float **output_addr, size_t *buff_size) {
|
||||
for (size_t i = 0; i < input_shape_list_.size(); ++i) {
|
||||
auto input_i_shape = input_shape_list_[i];
|
||||
auto input_i_addr = reinterpret_cast<float *>(inputs[i]->addr);
|
||||
|
@ -82,11 +83,12 @@ void ConcatCPUKernel::CopyDataToOutput(const std::vector<kernel::AddressPtr> &in
|
|||
size_t num = CPUKernelUtils::GetElementNumOnAxis(input_i_shape, axis_);
|
||||
num *= input_i_shape[axis_];
|
||||
auto pos = CPUKernelUtils::CalcOffset(input_i_shape, dim0, dim1, dim2, 0);
|
||||
auto ret = memcpy_s(*output_addr, num * sizeof(float), input_i_addr + pos, num * sizeof(float));
|
||||
auto ret = memcpy_s(*output_addr, *buff_size, input_i_addr + pos, num * sizeof(float));
|
||||
if (ret != EOK) {
|
||||
MS_LOG(EXCEPTION) << "memcpy failed.";
|
||||
}
|
||||
*output_addr += num;
|
||||
*buff_size -= num * sizeof(float);
|
||||
}
|
||||
}
|
||||
|
||||
|
|
|
@ -24,7 +24,7 @@ namespace mindspore {
|
|||
namespace kernel {
|
||||
class ConcatCPUKernel : public CPUKernel {
|
||||
public:
|
||||
ConcatCPUKernel() = default;
|
||||
ConcatCPUKernel() : axis_(0) {}
|
||||
~ConcatCPUKernel() override = default;
|
||||
|
||||
void InitKernel(const CNodePtr &kernel_node) override;
|
||||
|
@ -35,16 +35,15 @@ class ConcatCPUKernel : public CPUKernel {
|
|||
private:
|
||||
void CheckParam(const CNodePtr &kernel_node);
|
||||
void CopyDataToOutput(const std::vector<kernel::AddressPtr> &inputs, size_t dim0, size_t dim1, size_t dim2,
|
||||
float **output_addr);
|
||||
float **output_addr, size_t *buff_size);
|
||||
int axis_;
|
||||
std::vector<std::vector<size_t>> input_shape_list_;
|
||||
std::vector<size_t> output_shape_;
|
||||
};
|
||||
|
||||
MS_REG_CPU_KERNEL(
|
||||
Concat,
|
||||
KernelAttr().AddInputAttr(kNumberTypeFloat32).AddInputAttr(kNumberTypeFloat32).AddOutputAttr(kNumberTypeFloat32),
|
||||
ConcatCPUKernel);
|
||||
MS_REG_CPU_KERNEL(Concat,
|
||||
KernelAttr().SetAllSameAttr(true).AddInputAttr(kNumberTypeFloat32).AddOutputAttr(kNumberTypeFloat32),
|
||||
ConcatCPUKernel);
|
||||
} // namespace kernel
|
||||
} // namespace mindspore
|
||||
|
||||
|
|
|
@ -42,7 +42,7 @@ std::shared_ptr<CPUKernel> CPUKernelFactory::Create(const std::string &kernel_na
|
|||
MS_EXCEPTION_IF_NULL(kernel_info);
|
||||
const KernelBuildInfo *kernel_build_Info = kernel_info->select_kernel_build_info();
|
||||
MS_EXCEPTION_IF_NULL(kernel_build_Info);
|
||||
std::pair<bool, size_t> ret_pair = CPUKernelAttrCheck(kernel_name, kernel_build_Info);
|
||||
std::pair<bool, size_t> ret_pair = CPUKernelAttrCheck(kernel_name, *kernel_build_Info);
|
||||
if (ret_pair.first) {
|
||||
return (name_to_attr_creator_.find(kernel_name)->second)[ret_pair.second].second();
|
||||
}
|
||||
|
@ -50,7 +50,7 @@ std::shared_ptr<CPUKernel> CPUKernelFactory::Create(const std::string &kernel_na
|
|||
}
|
||||
|
||||
std::pair<bool, size_t> CPUKernelFactory::CPUKernelAttrCheck(const std::string &kernel_name,
|
||||
const KernelBuildInfo *kernel_info) {
|
||||
const KernelBuildInfo &kernel_info) {
|
||||
auto iter = name_to_attr_creator_.find(kernel_name);
|
||||
if (iter == name_to_attr_creator_.end()) {
|
||||
MS_LOG(INFO) << "Not registered CPU kernel: op[" << kernel_name << "]!";
|
||||
|
@ -59,27 +59,33 @@ std::pair<bool, size_t> CPUKernelFactory::CPUKernelAttrCheck(const std::string &
|
|||
auto creators = iter->second;
|
||||
for (size_t index = 0; index < creators.size(); ++index) {
|
||||
auto attr_creator = creators[index];
|
||||
for (size_t i = 0; i < kernel_info->GetInputNum(); ++i) {
|
||||
if (kernel_info->GetInputDeviceType(i) != attr_creator.first.GetInputAttr(i).first) {
|
||||
MS_LOG(WARNING) << "cpu kernel attr check failed. input index: " << i << ".";
|
||||
MS_LOG(WARNING) << "kernel info type:" << kernel_info->GetInputDeviceType(i) << ", "
|
||||
<< "register type:" << attr_creator.first.GetInputAttr(i).first;
|
||||
return std::make_pair(false, 0);
|
||||
}
|
||||
if (CPUKernelSingleAttrCheck(attr_creator.first, kernel_info)) {
|
||||
return std::make_pair(true, index);
|
||||
}
|
||||
for (size_t i = 0; i < kernel_info->GetOutputNum(); ++i) {
|
||||
if (kernel_info->GetOutputDeviceType(i) != attr_creator.first.GetOutputAttr(i).first) {
|
||||
MS_LOG(WARNING) << "cpu kernel attr check failed. output index: " << i << ".";
|
||||
MS_LOG(WARNING) << "kernel info type:" << kernel_info->GetOutputDeviceType(i) << ", "
|
||||
<< "register type:" << attr_creator.first.GetOutputAttr(i).first;
|
||||
return std::make_pair(false, 0);
|
||||
}
|
||||
}
|
||||
return std::make_pair(true, index);
|
||||
}
|
||||
return std::make_pair(false, 0);
|
||||
}
|
||||
|
||||
bool CPUKernelFactory::CPUKernelSingleAttrCheck(const KernelAttr &kernel_attr, const KernelBuildInfo &kernel_info) {
|
||||
for (size_t i = 0; i < kernel_info.GetInputNum(); ++i) {
|
||||
auto dtype = kernel_attr.GetAllSame() ? kernel_attr.GetInputAttr(0).first : kernel_attr.GetInputAttr(i).first;
|
||||
if (kernel_info.GetInputDeviceType(i) != dtype) {
|
||||
MS_LOG(DEBUG) << "input index:" << i << ", kernel info type:" << kernel_info.GetInputDeviceType(i)
|
||||
<< ", register type:" << dtype;
|
||||
return false;
|
||||
}
|
||||
}
|
||||
for (size_t i = 0; i < kernel_info.GetOutputNum(); ++i) {
|
||||
auto dtype = kernel_attr.GetAllSame() ? kernel_attr.GetOutputAttr(0).first : kernel_attr.GetOutputAttr(i).first;
|
||||
if (kernel_info.GetOutputDeviceType(i) != dtype) {
|
||||
MS_LOG(DEBUG) << "output index:" << i << ", kernel info type:" << kernel_info.GetOutputDeviceType(i)
|
||||
<< ", register type:" << dtype;
|
||||
return false;
|
||||
}
|
||||
}
|
||||
return true;
|
||||
}
|
||||
|
||||
std::vector<KernelAttr> CPUKernelFactory::GetSupportedKernelAttrList(const std::string &kernel_name) {
|
||||
std::vector<KernelAttr> result;
|
||||
auto iter = name_to_attr_creator_.find(kernel_name);
|
||||
|
|
|
@ -35,7 +35,6 @@ class CPUKernelFactory {
|
|||
public:
|
||||
static CPUKernelFactory &GetInstance();
|
||||
void Register(const std::string &kernel_name, const KernelAttr &kernel_attr, CPUKernelCreator &&kernel_creator);
|
||||
std::shared_ptr<CPUKernel> Create(const std::string &kernel_name);
|
||||
std::shared_ptr<CPUKernel> Create(const std::string &kernel_name, const CNodePtr &apply_kernel);
|
||||
std::vector<KernelAttr> GetSupportedKernelAttrList(const std::string &kernel_name);
|
||||
|
||||
|
@ -43,7 +42,8 @@ class CPUKernelFactory {
|
|||
CPUKernelFactory() = default;
|
||||
~CPUKernelFactory() = default;
|
||||
DISABLE_COPY_AND_ASSIGN(CPUKernelFactory)
|
||||
std::pair<bool, size_t> CPUKernelAttrCheck(const std::string &kernel_name, const KernelBuildInfo *kernel_info);
|
||||
std::pair<bool, size_t> CPUKernelAttrCheck(const std::string &kernel_name, const KernelBuildInfo &kernel_info);
|
||||
bool CPUKernelSingleAttrCheck(const KernelAttr &kernel_attr, const KernelBuildInfo &kernel_info);
|
||||
std::map<std::string, std::vector<std::pair<KernelAttr, CPUKernelCreator>>> name_to_attr_creator_;
|
||||
};
|
||||
|
||||
|
|
|
@ -40,7 +40,7 @@ bool GatherV2CPUKernel::Launch(const std::vector<kernel::AddressPtr> &inputs,
|
|||
const std::vector<kernel::AddressPtr> & /*workspace*/,
|
||||
const std::vector<kernel::AddressPtr> &outputs) {
|
||||
auto output_addr = reinterpret_cast<float *>(outputs[0]->addr);
|
||||
|
||||
auto buff_size = outputs[0]->size;
|
||||
size_t dim0 = input_shape_[0];
|
||||
size_t dim1 = input_shape_[1];
|
||||
size_t dim2 = input_shape_[2];
|
||||
|
@ -49,29 +49,29 @@ bool GatherV2CPUKernel::Launch(const std::vector<kernel::AddressPtr> &inputs,
|
|||
for (size_t i = 0; i < dim0; ++i) {
|
||||
for (size_t j = 0; j < dim1; ++j) {
|
||||
for (size_t k = 0; k < dim2; ++k) {
|
||||
CopyDataToOutput(inputs, i, j, k, &output_addr);
|
||||
CopyDataToOutput(inputs, i, j, k, &output_addr, &buff_size);
|
||||
}
|
||||
}
|
||||
}
|
||||
} else if (axis_ == 2) {
|
||||
for (size_t i = 0; i < dim0; ++i) {
|
||||
for (size_t j = 0; j < dim1; ++j) {
|
||||
CopyDataToOutput(inputs, i, j, 0, &output_addr);
|
||||
CopyDataToOutput(inputs, i, j, 0, &output_addr, &buff_size);
|
||||
}
|
||||
}
|
||||
} else if (axis_ == 1) {
|
||||
for (size_t i = 0; i < dim0; ++i) {
|
||||
CopyDataToOutput(inputs, i, 0, 0, &output_addr);
|
||||
CopyDataToOutput(inputs, i, 0, 0, &output_addr, &buff_size);
|
||||
}
|
||||
} else if (axis_ == 0) {
|
||||
CopyDataToOutput(inputs, 0, 0, 0, &output_addr);
|
||||
CopyDataToOutput(inputs, 0, 0, 0, &output_addr, &buff_size);
|
||||
}
|
||||
|
||||
return true;
|
||||
}
|
||||
|
||||
void GatherV2CPUKernel::CopyDataToOutput(const std::vector<kernel::AddressPtr> &inputs, size_t dim0, size_t dim1,
|
||||
size_t dim2, float **output_addr) {
|
||||
size_t dim2, float **output_addr, size_t *buff_size) {
|
||||
auto input_addr = reinterpret_cast<float *>(inputs[0]->addr);
|
||||
auto indices_addr = reinterpret_cast<int *>(inputs[1]->addr);
|
||||
|
||||
|
@ -88,11 +88,12 @@ void GatherV2CPUKernel::CopyDataToOutput(const std::vector<kernel::AddressPtr> &
|
|||
pos = CPUKernelUtils::CalcOffset(input_shape_, index, 0, 0, 0);
|
||||
}
|
||||
size_t num = CPUKernelUtils::GetElementNumOnAxis(input_shape_, axis_);
|
||||
auto ret = memcpy_s(*output_addr, num * sizeof(float), input_addr + pos, num * sizeof(float));
|
||||
auto ret = memcpy_s(*output_addr, *buff_size, input_addr + pos, num * sizeof(float));
|
||||
if (ret != EOK) {
|
||||
MS_LOG(EXCEPTION) << "memcpy failed.";
|
||||
}
|
||||
*output_addr += num;
|
||||
*buff_size -= num * sizeof(float);
|
||||
}
|
||||
}
|
||||
|
||||
|
|
|
@ -24,7 +24,7 @@ namespace mindspore {
|
|||
namespace kernel {
|
||||
class GatherV2CPUKernel : public CPUKernel {
|
||||
public:
|
||||
GatherV2CPUKernel() = default;
|
||||
GatherV2CPUKernel() : axis_(0) {}
|
||||
~GatherV2CPUKernel() override = default;
|
||||
|
||||
void InitKernel(const CNodePtr &kernel_node) override;
|
||||
|
@ -34,7 +34,7 @@ class GatherV2CPUKernel : public CPUKernel {
|
|||
|
||||
private:
|
||||
void CopyDataToOutput(const std::vector<kernel::AddressPtr> &inputs, size_t dim0, size_t dim1, size_t dim2,
|
||||
float **output_addr);
|
||||
float **output_addr, size_t *buff_size);
|
||||
void CheckParam(const CNodePtr &kernel_node);
|
||||
std::vector<size_t> input_shape_;
|
||||
std::vector<size_t> indices_shape_;
|
||||
|
|
|
@ -0,0 +1,78 @@
|
|||
# Copyright 2020 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.
|
||||
# ============================================================================
|
||||
|
||||
import numpy as np
|
||||
import pytest
|
||||
|
||||
import mindspore.context as context
|
||||
import mindspore.nn as nn
|
||||
from mindspore import Tensor
|
||||
from mindspore.common import dtype as mstype
|
||||
from mindspore.ops import operations as P
|
||||
|
||||
context.set_context(mode=context.GRAPH_MODE, device_target='CPU')
|
||||
|
||||
class Net2I(nn.Cell):
|
||||
def __init__(self):
|
||||
super(Net2I, self).__init__()
|
||||
self.addn = P.AddN()
|
||||
|
||||
def construct(self, x, y):
|
||||
return self.addn((x, y))
|
||||
|
||||
@pytest.mark.level0
|
||||
@pytest.mark.platform_x86_cpu
|
||||
@pytest.mark.env_onecard
|
||||
def test_net_2Input():
|
||||
x = np.arange(2 * 3 * 2).reshape(2, 3, 2).astype(np.float32)
|
||||
y = np.arange(2 * 3 * 2).reshape(2, 3, 2).astype(np.float32)
|
||||
addn = Net2I()
|
||||
output = addn(Tensor(x, mstype.float32), Tensor(y, mstype.float32))
|
||||
print("output:\n", output)
|
||||
expect_result = [[[0., 2.],
|
||||
[4., 6.],
|
||||
[8., 10.]],
|
||||
[[12., 14.],
|
||||
[16., 18.],
|
||||
[20., 22.]]]
|
||||
|
||||
assert (output.asnumpy() == expect_result).all()
|
||||
|
||||
class Net3I(nn.Cell):
|
||||
def __init__(self):
|
||||
super(Net3I, self).__init__()
|
||||
self.addn = P.AddN()
|
||||
|
||||
def construct(self, x, y, z):
|
||||
return self.addn((x, y, z))
|
||||
|
||||
@pytest.mark.level0
|
||||
@pytest.mark.platform_x86_cpu
|
||||
@pytest.mark.env_onecard
|
||||
def test_net_3Input():
|
||||
x = np.arange(2 * 3).reshape(2, 3).astype(np.float32)
|
||||
y = np.arange(2 * 3).reshape(2, 3).astype(np.float32)
|
||||
z = np.arange(2 * 3).reshape(2, 3).astype(np.float32)
|
||||
addn = Net3I()
|
||||
output = addn(Tensor(x, mstype.float32), Tensor(y, mstype.float32), Tensor(z, mstype.float32))
|
||||
print("output:\n", output)
|
||||
expect_result = [[0., 3., 6.],
|
||||
[9., 12., 15]]
|
||||
|
||||
assert (output.asnumpy() == expect_result).all()
|
||||
|
||||
if __name__ == '__main__':
|
||||
test_net_2Input()
|
||||
test_net_3Input()
|
|
@ -71,13 +71,13 @@ def test_in2_axis1():
|
|||
assert np.all(diff < error)
|
||||
assert np.all(-diff < error)
|
||||
|
||||
class Concat_Axis2(nn.Cell):
|
||||
class Concat_in3_Axis2(nn.Cell):
|
||||
def __init__(self):
|
||||
super(Concat_Axis2, self).__init__()
|
||||
super(Concat_in3_Axis2, self).__init__()
|
||||
self.cat = P.Concat(axis=-1)
|
||||
|
||||
def construct(self, x1, x2):
|
||||
return self.cat((x1, x2))
|
||||
def construct(self, x1, x2, x3):
|
||||
return self.cat((x1, x2, x3))
|
||||
|
||||
@pytest.mark.level0
|
||||
@pytest.mark.platform_x86_cpu
|
||||
|
@ -86,10 +86,10 @@ def test_in3_axis2():
|
|||
x1 = Tensor(np.arange(2 * 2 * 1).reshape(2, 2, 1), mstype.float32)
|
||||
x2 = Tensor(np.arange(2 * 2 * 2).reshape(2, 2, 2), mstype.float32)
|
||||
x3 = Tensor(np.arange(2 * 2 * 3).reshape(2, 2, 3), mstype.float32)
|
||||
cat = Concat_Axis2()
|
||||
output_ms = cat(x1, x2)
|
||||
cat = Concat_in3_Axis2()
|
||||
output_ms = cat(x1, x2, x3)
|
||||
print("output:\n", output_ms)
|
||||
output_np = np.concatenate((x1.asnumpy(), x2.asnumpy()), axis=-1)
|
||||
output_np = np.concatenate((x1.asnumpy(), x2.asnumpy(), x3.asnumpy()), axis=-1)
|
||||
|
||||
error = np.ones(shape=output_np.shape) * 10e-6
|
||||
diff = output_ms.asnumpy() - output_np
|
||||
|
|
|
@ -1,4 +1,4 @@
|
|||
# Copyright 2019 Huawei Technologies Co., Ltd
|
||||
# Copyright 2020 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.
|
||||
|
|
Loading…
Reference in New Issue