p60975123
/
mindspore
forked from mindspore-Ecosystem/mindspore
1
0
Fork 0
Code Issues Pull Requests Projects Releases Wiki Activity

!28537 detach ScatterUpdate to ScatterNdUpdate and TensorScatterNdUpdate for cpu backend 

Merge pull request !28537 from zhuzhongrui/pub_master
This commit is contained in:
i-robot 2022-01-05 07:27:57 +00:00 committed by Gitee
parent cbab78f824 23c4089013
commit c596735d03
No known key found for this signature in database
GPG Key ID: 173E9B9CA92EEF8F
2 changed files with 47 additions and 20 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

34
mindspore/ccsrc/backend/kernel_compiler/cpu/scatter_nd_update_cpu_kernel.cc
View File

@ -58,7 +58,7 @@ void Compute(const ComputeParams<T> *params, const size_t start, const size_t en
}
} // namespace
void ScatterNdUpdateCPUKernel::InitKernel(const CNodePtr &kernel_node) {
void ScatterUpdateCPUKernel::InitKernel(const CNodePtr &kernel_node) {
MS_EXCEPTION_IF_NULL(kernel_node);
kernel_name_ = AnfAlgo::GetCNodeName(kernel_node);
auto shape = AnfAlgo::GetPrevNodeOutputInferShape(kernel_node, 0);
@ -109,9 +109,9 @@ void ScatterNdUpdateCPUKernel::InitKernel(const CNodePtr &kernel_node) {
dtype_ = AnfAlgo::GetInputDeviceDataType(kernel_node, 0);
}
bool ScatterNdUpdateCPUKernel::Launch(const std::vector<kernel::AddressPtr> &inputs,
const std::vector<kernel::AddressPtr> &,
const std::vector<kernel::AddressPtr> &outputs) {
bool ScatterUpdateCPUKernel::Launch(const std::vector<kernel::AddressPtr> &inputs,
const std::vector<kernel::AddressPtr> &,
const std::vector<kernel::AddressPtr> &outputs) {
CHECK_KERNEL_INPUTS_NUM(inputs.size(), kScatterNdUpdateInputsNum, kernel_name_);
CHECK_KERNEL_OUTPUTS_NUM(outputs.size(), kScatterNdUpdateOutputsNum, kernel_name_);
switch (dtype_) {
@ -139,13 +139,9 @@ bool ScatterNdUpdateCPUKernel::Launch(const std::vector<kernel::AddressPtr> &inp
}
template <typename T>
void ScatterNdUpdateCPUKernel::LaunchKernel(const std::vector<AddressPtr> &inputs,
const std::vector<kernel::AddressPtr> &outputs) {
auto x = reinterpret_cast<T *>(outputs[0]->addr);
auto ret = memcpy_s(x, outputs[0]->size, inputs[0]->addr, inputs[0]->size);
if (ret != 0) {
MS_LOG(EXCEPTION) << "For '" << kernel_name_ << "', memcpy_s error. Error no: " << ret;
}
void ScatterUpdateCPUKernel::LaunchKernel(const std::vector<AddressPtr> &inputs,
const std::vector<kernel::AddressPtr> &outputs) {
T *x = reinterpret_cast<T *>(ScatterUpdateRealData(inputs, outputs));
ComputeParams<T> params;
params.x_ = x;
params.indices_ = reinterpret_cast<int *>(inputs[1]->addr);
@ -169,6 +165,22 @@ void ScatterNdUpdateCPUKernel::LaunchKernel(const std::vector<AddressPtr> &input
start += once_compute_size;
}
ParallelLaunch(tasks);
(void)memcpy_s(outputs[0]->addr, outputs[0]->size, x, inputs[0]->size);
}
void *ScatterNdUpdateCPUKernel::ScatterUpdateRealData(const std::vector<AddressPtr> &inputs,
const std::vector<kernel::AddressPtr> &) {
return inputs[0]->addr;
}
void *TensorScatterUpdateCPUKernel::ScatterUpdateRealData(const std::vector<AddressPtr> &inputs,
const std::vector<kernel::AddressPtr> &outputs) {
void *x = outputs[0]->addr;
auto ret = memcpy_s(x, outputs[0]->size, inputs[0]->addr, inputs[0]->size);
if (ret != 0) {
MS_LOG(EXCEPTION) << "For '" << kernel_name_ << "', memcpy_s error. Error no: " << ret;
}
return x;
}
} // namespace kernel
} // namespace mindspore

33
mindspore/ccsrc/backend/kernel_compiler/cpu/scatter_nd_update_cpu_kernel.h
View File

@ -36,16 +36,19 @@ struct ComputeParams {
size_t x_mem_size_{0};
};
class ScatterNdUpdateCPUKernel : public CPUKernel {
class ScatterUpdateCPUKernel : public CPUKernel {
public:
ScatterNdUpdateCPUKernel() = default;
~ScatterNdUpdateCPUKernel() override = default;
ScatterUpdateCPUKernel() = default;
~ScatterUpdateCPUKernel() 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;
virtual void *ScatterUpdateRealData(const std::vector<AddressPtr> &inputs,
const std::vector<kernel::AddressPtr> &outputs) = 0;
private:
template <typename T>
void LaunchKernel(const std::vector<AddressPtr> &inputs, const std::vector<kernel::AddressPtr> &outputs);
@ -57,13 +60,25 @@ class ScatterNdUpdateCPUKernel : public CPUKernel {
std::vector<int> out_strides_;
};
class ScatterNdUpdateCPUKernel : public ScatterUpdateCPUKernel {
protected:
void *ScatterUpdateRealData(const std::vector<AddressPtr> &inputs,
const std::vector<kernel::AddressPtr> &outputs) override;
};
class TensorScatterUpdateCPUKernel : public ScatterUpdateCPUKernel {
protected:
void *ScatterUpdateRealData(const std::vector<AddressPtr> &inputs,
const std::vector<kernel::AddressPtr> &outputs) override;
};
MS_REG_CPU_KERNEL(ScatterNdUpdate,
KernelAttr()
.AddInputAttr(kNumberTypeFloat32)
.AddInputAttr(kNumberTypeInt32)
.AddInputAttr(kNumberTypeFloat32)
.AddOutputAttr(kNumberTypeFloat32),
ScatterNdUpdateCPUKernel);
ScatterNdUpdateCPUKernel)
MS_REG_CPU_KERNEL(TensorScatterUpdate,
KernelAttr()
@ -71,7 +86,7 @@ MS_REG_CPU_KERNEL(TensorScatterUpdate,
.AddInputAttr(kNumberTypeInt32)
.AddInputAttr(kNumberTypeFloat32)
.AddOutputAttr(kNumberTypeFloat32),
ScatterNdUpdateCPUKernel);
TensorScatterUpdateCPUKernel)
MS_REG_CPU_KERNEL(ScatterNdUpdate,
KernelAttr()
@ -79,7 +94,7 @@ MS_REG_CPU_KERNEL(ScatterNdUpdate,
.AddInputAttr(kNumberTypeInt32)
.AddInputAttr(kNumberTypeFloat64)
.AddOutputAttr(kNumberTypeFloat64),
ScatterNdUpdateCPUKernel);
ScatterNdUpdateCPUKernel)
MS_REG_CPU_KERNEL(TensorScatterUpdate,
KernelAttr()
@ -87,7 +102,7 @@ MS_REG_CPU_KERNEL(TensorScatterUpdate,
.AddInputAttr(kNumberTypeInt32)
.AddInputAttr(kNumberTypeFloat64)
.AddOutputAttr(kNumberTypeFloat64),
ScatterNdUpdateCPUKernel);
TensorScatterUpdateCPUKernel)
MS_REG_CPU_KERNEL(ScatterNdUpdate,
KernelAttr()
@ -103,7 +118,7 @@ MS_REG_CPU_KERNEL(TensorScatterUpdate,
.AddInputAttr(kNumberTypeInt32)
.AddInputAttr(kNumberTypeInt32)
.AddOutputAttr(kNumberTypeInt32),
ScatterNdUpdateCPUKernel);
TensorScatterUpdateCPUKernel)
MS_REG_CPU_KERNEL(ScatterNdUpdate,
KernelAttr()
@ -119,7 +134,7 @@ MS_REG_CPU_KERNEL(TensorScatterUpdate,
.AddInputAttr(kNumberTypeInt32)
.AddInputAttr(kNumberTypeInt64)
.AddOutputAttr(kNumberTypeInt64),
ScatterNdUpdateCPUKernel);
TensorScatterUpdateCPUKernel)
} // namespace kernel
} // namespace mindspore