• Home
  • Line#
  • Scopes#
  • Navigate#
  • Raw
  • Download
1# Copyright 2020 Huawei Technologies Co., Ltd
2#
3# Licensed under the Apache License, Version 2.0 (the "License");
4# you may not use this file except in compliance with the License.
5# You may obtain a copy of the License at
6#
7# http://www.apache.org/licenses/LICENSE-2.0
8#
9# Unless required by applicable law or agreed to in writing, software
10# distributed under the License is distributed on an "AS IS" BASIS,
11# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12# See the License for the specific language governing permissions and
13# limitations under the License.
14# ============================================================================
15
16import numpy as np
17import pytest
18import mindspore.context as context
19from mindspore import Tensor
20from mindspore.nn import Cell
21import mindspore.ops.operations as P
22
23
24class SumOutNet(Cell):
25    def __init__(self):
26        super(SumOutNet, self).__init__()
27        self.square = P.Square()
28        self.sum = P.ReduceSum()
29
30    def construct(self, x):
31        mul_res = self.square(x)
32        return self.sum(mul_res, (0,))
33
34
35class SingleOutNet(Cell):
36    def __init__(self):
37        super(SingleOutNet, self).__init__()
38        self.add = P.Add()
39        self.mul = P.Mul()
40        self.sum = P.ReduceSum()
41
42    def construct(self, x, y):
43        mul_res = self.mul(x, y)
44        sum_res = self.sum(mul_res, ())
45        return self.add(sum_res, x)
46
47
48class MultiOutNet(Cell):
49    def __init__(self):
50        super(MultiOutNet, self).__init__()
51        self.add = P.Add()
52        self.mul = P.Mul()
53        self.sum = P.ReduceSum()
54
55    def construct(self, x, y):
56        add_res = self.add(x, y)
57        mul_res = self.mul(add_res, add_res)
58        sum_res = self.sum(mul_res, ())
59        return self.add(add_res, sum_res)
60
61
62def atomic_add_sum_output():
63    np.random.seed(0)
64    input_x = np.random.normal(0, 1, [2, 3, 4, 3]).astype(np.float32)
65
66    expect = np.sum(np.square(input_x), axis=(0,))
67
68    net = SumOutNet()
69    result = net(Tensor(input_x))
70
71    res = np.allclose(expect, result.asnumpy(), rtol=1.e-4, atol=1.e-7, equal_nan=True)
72    assert res
73
74
75def atomic_add_single_output():
76    np.random.seed(0)
77    input_x = np.random.normal(0, 1, [2, 2, 2, 256]).astype(np.float32)
78    input_y = np.random.normal(0, 1, [2, 2, 2, 256]).astype(np.float32)
79
80    expect = np.sum(input_x * input_y) + input_x
81
82    net = SingleOutNet()
83    result = net(Tensor(input_x), Tensor(input_y))
84
85    res = np.allclose(expect, result.asnumpy(), rtol=1.e-4, atol=1.e-7, equal_nan=True)
86    assert res
87
88
89def atomic_add_multi_output():
90    np.random.seed(0)
91    input_x = np.random.normal(0, 1, [2, 2, 2, 256]).astype(np.float32)
92    input_y = np.random.normal(0, 1, [2, 2, 2, 256]).astype(np.float32)
93
94    expect = np.sum(np.square(input_x + input_y)) + (input_x + input_y)
95
96    net = MultiOutNet()
97    result = net(Tensor(input_x), Tensor(input_y))
98
99    res = np.allclose(expect, result.asnumpy(), rtol=1.e-4, atol=1.e-7, equal_nan=True)
100    assert res
101
102
103@pytest.mark.level0
104@pytest.mark.platform_x86_gpu_training
105@pytest.mark.env_onecard
106def test_atomic_add_sum_output_gpu():
107    context.set_context(mode=context.GRAPH_MODE, enable_graph_kernel=True, device_target="GPU")
108    atomic_add_sum_output()
109
110
111@pytest.mark.level0
112@pytest.mark.platform_x86_gpu_training
113@pytest.mark.env_onecard
114def test_atomic_add_single_output_gpu():
115    context.set_context(mode=context.GRAPH_MODE, enable_graph_kernel=True, device_target="GPU")
116    atomic_add_single_output()
117
118
119@pytest.mark.level0
120@pytest.mark.platform_x86_gpu_training
121@pytest.mark.env_onecard
122def test_atomic_add_multi_output_gpu():
123    context.set_context(mode=context.GRAPH_MODE, enable_graph_kernel=True, device_target="GPU")
124    atomic_add_multi_output()
125
126
127@pytest.mark.level0
128@pytest.mark.platform_arm_ascend_training
129@pytest.mark.platform_x86_ascend_training
130@pytest.mark.env_onecard
131def test_atomic_add_sum_output_ascend():
132    context.set_context(mode=context.GRAPH_MODE, enable_graph_kernel=True, device_target="Ascend")
133    atomic_add_sum_output()
134
135
136@pytest.mark.level0
137@pytest.mark.platform_arm_ascend_training
138@pytest.mark.platform_x86_ascend_training
139@pytest.mark.env_onecard
140def test_atomic_add_single_output_ascend():
141    context.set_context(mode=context.GRAPH_MODE, enable_graph_kernel=True, device_target="Ascend")
142    atomic_add_single_output()
143
144
145@pytest.mark.level0
146@pytest.mark.platform_arm_ascend_training
147@pytest.mark.platform_x86_ascend_training
148@pytest.mark.env_onecard
149def test_atomic_add_multi_output_ascend():
150    context.set_context(mode=context.GRAPH_MODE, enable_graph_kernel=True, device_target="Ascend")
151    atomic_add_multi_output()
152