diff --git a/mindspore/ccsrc/backend/kernel_compiler/cpu/scatter_arithmetic_cpu_kernel.cc b/mindspore/ccsrc/backend/kernel_compiler/cpu/scatter_arithmetic_cpu_kernel.cc index ddd56bee9c..012e04e3ad 100644 --- a/mindspore/ccsrc/backend/kernel_compiler/cpu/scatter_arithmetic_cpu_kernel.cc +++ b/mindspore/ccsrc/backend/kernel_compiler/cpu/scatter_arithmetic_cpu_kernel.cc @@ -83,118 +83,97 @@ bool ScatterArithmeticCPUKernel::Launch(const std::vector template void ScatterArithmeticCPUKernel::ScatterAdd(T *input, const int *indices, const T *updates) { - auto task = [this, input, indices, updates](size_t start, size_t end) { - for (size_t i = start; i < end; i++) { - auto base_index_updates = i * inner_size_; - auto base_index_input = indices[i] * inner_size_; - for (size_t j = 0; j < inner_size_; j++) { - input[base_index_input + j] += updates[base_index_updates + j]; - } + for (size_t i = 0; i < indices_size_; i++) { + auto base_index_updates = i * inner_size_; + auto base_index_input = indices[i] * inner_size_; + for (size_t j = 0; j < inner_size_; j++) { + input[base_index_input + j] += updates[base_index_updates + j]; } - }; - CPUKernelUtils::ParallelFor(task, indices_size_); + } } template void ScatterArithmeticCPUKernel::ScatterSub(T *input, const int *indices, const T *updates) { - auto task = [this, input, indices, updates](size_t start, size_t end) { - for (size_t i = start; i < end; i++) { - auto base_index_updates = i * inner_size_; - auto base_index_input = indices[i] * inner_size_; - for (size_t j = 0; j < inner_size_; j++) { - input[base_index_input + j] -= updates[base_index_updates + j]; - } + for (size_t i = 0; i < indices_size_; i++) { + auto base_index_updates = i * inner_size_; + auto base_index_input = indices[i] * inner_size_; + for (size_t j = 0; j < inner_size_; j++) { + input[base_index_input + j] -= updates[base_index_updates + j]; } - }; - CPUKernelUtils::ParallelFor(task, indices_size_); + } } template void ScatterArithmeticCPUKernel::ScatterMul(T *input, const int *indices, const T *updates) { - auto task = [this, input, indices, updates](size_t start, size_t end) { - for (size_t i = start; i < end; i++) { - auto base_index_updates = i * inner_size_; - auto base_index_input = indices[i] * inner_size_; - for (size_t j = 0; j < inner_size_; j++) { - input[base_index_input + j] *= updates[base_index_updates + j]; - } + for (size_t i = 0; i < indices_size_; i++) { + auto base_index_updates = i * inner_size_; + auto base_index_input = indices[i] * inner_size_; + for (size_t j = 0; j < inner_size_; j++) { + input[base_index_input + j] *= updates[base_index_updates + j]; } - }; - CPUKernelUtils::ParallelFor(task, indices_size_); + } } template void ScatterArithmeticCPUKernel::ScatterDiv(T *input, const int *indices, const T *updates) { - auto task = [this, input, indices, updates](size_t start, size_t end) { - for (size_t i = start; i < end; i++) { - for (size_t j = 0; j < inner_size_; j++) { - auto dividend = input[indices[i] * inner_size_ + j]; - auto divisor = updates[i * inner_size_ + j]; - if (divisor == 0) { - if (dividend == 0) { - input[indices[i] * inner_size_ + j] = std::numeric_limits::quiet_NaN(); - continue; - } - if (std::numeric_limits::has_infinity) { - input[indices[i] * inner_size_ + j] = - dividend > 0 ? std::numeric_limits::infinity() : -std::numeric_limits::infinity(); - } else { - input[indices[i] * inner_size_ + j] = - dividend > 0 ? std::numeric_limits::max() : std::numeric_limits::min(); - } + for (size_t i = 0; i < indices_size_; i++) { + for (size_t j = 0; j < inner_size_; j++) { + auto dividend = input[indices[i] * inner_size_ + j]; + auto divisor = updates[i * inner_size_ + j]; + if (divisor == 0) { + if (dividend == 0) { + input[indices[i] * inner_size_ + j] = std::numeric_limits::quiet_NaN(); continue; } - input[indices[i] * inner_size_ + j] = dividend / divisor; + if (std::numeric_limits::has_infinity) { + input[indices[i] * inner_size_ + j] = + dividend > 0 ? std::numeric_limits::infinity() : -std::numeric_limits::infinity(); + } else { + input[indices[i] * inner_size_ + j] = + dividend > 0 ? std::numeric_limits::max() : std::numeric_limits::min(); + } + continue; } + input[indices[i] * inner_size_ + j] = dividend / divisor; } - }; - CPUKernelUtils::ParallelFor(task, indices_size_); + } } template void ScatterArithmeticCPUKernel::ScatterMax(T *input, const int *indices, const T *updates) { - auto task = [this, input, indices, updates](size_t start, size_t end) { - for (size_t i = start; i < end; i++) { - auto base_index_updates = i * inner_size_; - auto base_index_input = indices[i] * inner_size_; - for (size_t j = 0; j < inner_size_; j++) { - input[base_index_input + j] = input[base_index_input + j] > updates[base_index_updates + j] - ? input[base_index_input + j] - : updates[base_index_updates + j]; - } + for (size_t i = 0; i < indices_size_; i++) { + auto base_index_updates = i * inner_size_; + auto base_index_input = indices[i] * inner_size_; + for (size_t j = 0; j < inner_size_; j++) { + input[base_index_input + j] = input[base_index_input + j] > updates[base_index_updates + j] + ? input[base_index_input + j] + : updates[base_index_updates + j]; } - }; - CPUKernelUtils::ParallelFor(task, indices_size_); + } } template void ScatterArithmeticCPUKernel::ScatterMin(T *input, const int *indices, const T *updates) { - auto task = [this, input, indices, updates](size_t start, size_t end) { - for (size_t i = start; i < end; i++) { - auto base_index_updates = i * inner_size_; - auto base_index_input = indices[i] * inner_size_; - for (size_t j = 0; j < inner_size_; j++) { - input[base_index_input + j] = input[base_index_input + j] < updates[base_index_updates + j] - ? input[base_index_input + j] - : updates[base_index_updates + j]; - } + for (size_t i = 0; i < indices_size_; i++) { + auto base_index_updates = i * inner_size_; + auto base_index_input = indices[i] * inner_size_; + for (size_t j = 0; j < inner_size_; j++) { + input[base_index_input + j] = input[base_index_input + j] < updates[base_index_updates + j] + ? input[base_index_input + j] + : updates[base_index_updates + j]; } - }; - CPUKernelUtils::ParallelFor(task, indices_size_); + } } template void ScatterArithmeticCPUKernel::ScatterUpdate(T *input, const int *indices, const T *updates) { - auto task = [this, input, indices, updates](size_t start, size_t end) { - for (size_t i = start; i < end; i++) { - auto base_index_updates = i * inner_size_; - auto base_index_input = indices[i] * inner_size_; - for (size_t j = 0; j < inner_size_; j++) { - input[base_index_input + j] = updates[base_index_updates + j]; - } + for (size_t i = 0; i < indices_size_; i++) { + auto base_index_updates = i * inner_size_; + auto base_index_input = indices[i] * inner_size_; + for (size_t j = 0; j < inner_size_; j++) { + input[base_index_input + j] = updates[base_index_updates + j]; } - }; - CPUKernelUtils::ParallelFor(task, indices_size_); + } } } // namespace kernel } // namespace mindspore