sunsuodong
5 years ago
parent
10076ffe1a
commit
ade90be427
@ -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,56 @@
|
||||
/**
|
||||
* 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().AddInputAttr(kNumberTypeFloat32).AddInputAttr(kNumberTypeFloat32).AddOutputAttr(kNumberTypeFloat32),
|
||||
AddNCPUKernel);
|
||||
MS_REG_CPU_KERNEL(AddN,
|
||||
KernelAttr()
|
||||
.AddInputAttr(kNumberTypeFloat32)
|
||||
.AddInputAttr(kNumberTypeFloat32)
|
||||
.AddInputAttr(kNumberTypeFloat32)
|
||||
.AddOutputAttr(kNumberTypeFloat32),
|
||||
AddNCPUKernel);
|
||||
} // namespace kernel
|
||||
} // namespace mindspore
|
||||
|
||||
#endif // MINDSPORE_CCSRC_KERNEL_CPU_ADDN_CPU_KERNEL_H_
|
||||
@ -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()
|
||||
Loading…
Reference in new issue