Skip to content

Commit faf40da

Browse files
oyxuan-11oyxuan-11
authored
[NPU] Support NPU kernel of stack op (#31711)
1 parent d55120d commit faf40da

File tree

2 files changed

+259
-0
lines changed

2 files changed

+259
-0
lines changed

‎paddle/fluid/operators/stack_op_npu.cc‎

Lines changed: 106 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,106 @@
1+
/* Copyright (c) 2021 PaddlePaddle Authors. All Rights Reserved.
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+
#ifdef PADDLE_WITH_ASCEND_CL
16+
#include <memory>
17+
#include <string>
18+
#include <vector>
19+
20+
#include "paddle/fluid/operators/activation_op.h"
21+
#include "paddle/fluid/operators/npu_op_runner.h"
22+
#include "paddle/fluid/operators/stack_op.h"
23+
#include "paddle/fluid/operators/unsqueeze_op.h"
24+
25+
namespace paddle {
26+
namespace operators {
27+
28+
using Tensor = framework::Tensor;
29+
30+
template <typename DeviceContext, typename T>
31+
class StackNPUKernel : public framework::OpKernel<T> {
32+
public:
33+
void Compute(const framework::ExecutionContext& ctx) const override {
34+
auto x = ctx.MultiInput<Tensor>("X");
35+
int32_t N = x.size();
36+
37+
PADDLE_ENFORCE_GT(
38+
N, 0, platform::errors::InvalidArgument("number of input Tensor <= 0"));
39+
40+
std::vector<paddle::framework::Tensor> x_list;
41+
for (int i = 0; i < N; i++) {
42+
x_list.push_back(*x[i]);
43+
}
44+
45+
int axis = ctx.Attr<int>("axis");
46+
47+
if (axis < 0) {
48+
axis = axis + x_list[0].dims().size() + 1;
49+
}
50+
auto* out = ctx.Output<Tensor>("Y");
51+
52+
auto place = ctx.GetPlace();
53+
54+
auto stream =
55+
ctx.template device_context<paddle::platform::NPUDeviceContext>()
56+
.stream();
57+
58+
out->mutable_data<T>(place);
59+
60+
if (axis != 0) {
61+
auto x_dim = x_list[0].dims();
62+
std::vector<int> vec_dim_tmp;
63+
vec_dim_tmp.push_back(N);
64+
for (auto i = 0; i < x_dim.size(); ++i) {
65+
vec_dim_tmp.push_back(x_dim[i]);
66+
}
67+
68+
Tensor tmp_stack(out->type());
69+
tmp_stack.Resize(framework::make_ddim(vec_dim_tmp));
70+
tmp_stack.mutable_data<T>(ctx.GetPlace());
71+
72+
auto runner =
73+
NpuOpRunner("Pack", {x_list}, {tmp_stack}, {{"axis", 0}, {"N", N}});
74+
runner.Run(stream);
75+
76+
std::vector<int64_t> vec_trans;
77+
for (auto i = 1; i <= x_dim.size(); ++i) {
78+
vec_trans.push_back(i);
79+
if (i == axis) {
80+
vec_trans.push_back(0);
81+
}
82+
}
83+
84+
auto runner_trans_final =
85+
NpuOpRunner("TransposeD", {tmp_stack}, {*out}, {{"perm", vec_trans}});
86+
runner_trans_final.Run(stream);
87+
88+
} else {
89+
auto runner =
90+
NpuOpRunner("Pack", {x_list}, {*out}, {{"axis", axis}, {"N", N}});
91+
runner.Run(stream);
92+
}
93+
}
94+
};
95+
96+
} // namespace operators
97+
} // namespace paddle
98+
99+
namespace ops = paddle::operators;
100+
101+
REGISTER_OP_NPU_KERNEL(
102+
stack, ops::StackNPUKernel<paddle::platform::NPUDeviceContext, float>,
103+
ops::StackNPUKernel<paddle::platform::NPUDeviceContext,
104+
paddle::platform::float16>);
105+
106+
#endif

‎python/paddle/fluid/tests/unittests/npu/test_stack_op_npu.py‎

Lines changed: 153 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,153 @@
1+
# Copyright (c) 2021 PaddlePaddle Authors. All Rights Reserved.
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+
from __future__ import print_function
16+
17+
import numpy as np
18+
import unittest
19+
import sys
20+
sys.path.append("..")
21+
from op_test import OpTest
22+
import paddle
23+
import paddle.fluid as fluid
24+
import paddle.fluid.core as core
25+
26+
paddle.enable_static()
27+
SEED = 2021
28+
29+
30+
@unittest.skipIf(not paddle.is_compiled_with_npu(),
31+
"core is not compiled with NPU")
32+
class TestStack1(OpTest):
33+
def initDefaultParameters(self):
34+
self.num_inputs = 4
35+
self.input_dim = (5, 6, 7)
36+
self.axis = 0
37+
self.dtype = 'float32'
38+
39+
def get_x_names(self):
40+
x_names = []
41+
for i in range(self.num_inputs):
42+
x_names.append('x{}'.format(i))
43+
return x_names
44+
45+
def setUp(self):
46+
self.initDefaultParameters()
47+
self.set_npu()
48+
self.op_type = "stack"
49+
self.place = paddle.NPUPlace(0)
50+
51+
self.x = []
52+
for i in range(self.num_inputs):
53+
self.x.append(
54+
np.random.random(size=self.input_dim).astype(self.dtype))
55+
56+
tmp = []
57+
x_names = self.get_x_names()
58+
for i in range(self.num_inputs):
59+
tmp.append((x_names[i], self.x[i]))
60+
61+
self.inputs = {'X': tmp}
62+
self.outputs = {'Y': np.stack(self.x, axis=self.axis)}
63+
self.attrs = {'axis': self.axis}
64+
65+
def set_npu(self):
66+
self.__class__.use_npu = True
67+
68+
def test_check_output(self):
69+
self.check_output_with_place(self.place, check_dygraph=False)
70+
71+
72+
class TestStack2(OpTest):
73+
def initDefaultParameters(self):
74+
self.num_inputs = 4
75+
self.input_dim = (2, 3, 4)
76+
self.axis = -1
77+
self.dtype = 'float32'
78+
79+
def get_x_names(self):
80+
x_names = []
81+
for i in range(self.num_inputs):
82+
x_names.append('x{}'.format(i))
83+
return x_names
84+
85+
def setUp(self):
86+
self.initDefaultParameters()
87+
self.set_npu()
88+
self.op_type = "stack"
89+
self.place = paddle.NPUPlace(0)
90+
91+
self.x = []
92+
for i in range(self.num_inputs):
93+
self.x.append(
94+
np.random.random(size=self.input_dim).astype(self.dtype))
95+
96+
tmp = []
97+
x_names = self.get_x_names()
98+
for i in range(self.num_inputs):
99+
tmp.append((x_names[i], self.x[i]))
100+
101+
self.inputs = {'X': tmp}
102+
self.outputs = {'Y': np.stack(self.x, axis=self.axis)}
103+
self.attrs = {'axis': self.axis}
104+
105+
def set_npu(self):
106+
self.__class__.use_npu = True
107+
108+
def test_check_output(self):
109+
self.check_output_with_place(self.place, check_dygraph=False)
110+
111+
112+
class TestStack3(OpTest):
113+
def initDefaultParameters(self):
114+
self.num_inputs = 4
115+
self.input_dim = (2, 3, 4)
116+
self.axis = 1
117+
self.dtype = 'float32'
118+
119+
def get_x_names(self):
120+
x_names = []
121+
for i in range(self.num_inputs):
122+
x_names.append('x{}'.format(i))
123+
return x_names
124+
125+
def setUp(self):
126+
self.initDefaultParameters()
127+
self.set_npu()
128+
self.op_type = "stack"
129+
self.place = paddle.NPUPlace(0)
130+
131+
self.x = []
132+
for i in range(self.num_inputs):
133+
self.x.append(
134+
np.random.random(size=self.input_dim).astype(self.dtype))
135+
136+
tmp = []
137+
x_names = self.get_x_names()
138+
for i in range(self.num_inputs):
139+
tmp.append((x_names[i], self.x[i]))
140+
141+
self.inputs = {'X': tmp}
142+
self.outputs = {'Y': np.stack(self.x, axis=self.axis)}
143+
self.attrs = {'axis': self.axis}
144+
145+
def set_npu(self):
146+
self.__class__.use_npu = True
147+
148+
def test_check_output(self):
149+
self.check_output_with_place(self.place, check_dygraph=False)
150+
151+
152+
if __name__ == '__main__':
153+
unittest.main()

0 commit comments

Comments
 (0)