Skip to content

Commit b83ff45

Browse files
committed
Register reduce_op_kernel with multiple data types
1 parent c73f00f commit b83ff45

File tree

3 files changed

+33
-9
lines changed

3 files changed

+33
-9
lines changed

paddle/operators/reduce_op.cc

Lines changed: 13 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -190,10 +190,22 @@ REGISTER_OP(reduce_min, ops::ReduceOp, ops::ReduceMinOpMaker, reduce_min_grad,
190190
#define REGISTER_REDUCE_CPU_KERNEL(reduce_type, functor, grad_functor) \
191191
REGISTER_OP_CPU_KERNEL(reduce_type, \
192192
ops::ReduceKernel<paddle::platform::CPUDeviceContext, \
193-
float, ops::functor>); \
193+
float, ops::functor>, \
194+
ops::ReduceKernel<paddle::platform::CPUDeviceContext, \
195+
double, ops::functor>, \
196+
ops::ReduceKernel<paddle::platform::CPUDeviceContext, \
197+
int, ops::functor>, \
198+
ops::ReduceKernel<paddle::platform::CPUDeviceContext, \
199+
int64_t, ops::functor>); \
194200
REGISTER_OP_CPU_KERNEL( \
195201
reduce_type##_grad, \
196202
ops::ReduceGradKernel<paddle::platform::CPUDeviceContext, float, \
203+
ops::grad_functor>, \
204+
ops::ReduceGradKernel<paddle::platform::CPUDeviceContext, double, \
205+
ops::grad_functor>, \
206+
ops::ReduceGradKernel<paddle::platform::CPUDeviceContext, int, \
207+
ops::grad_functor>, \
208+
ops::ReduceGradKernel<paddle::platform::CPUDeviceContext, int64_t, \
197209
ops::grad_functor>);
198210

199211
FOR_EACH_KERNEL_FUNCTOR(REGISTER_REDUCE_CPU_KERNEL);

paddle/operators/reduce_op.cu

Lines changed: 13 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -20,10 +20,22 @@ namespace ops = paddle::operators;
2020
#define REGISTER_REDUCE_GPU_KERNEL(reduce_type, functor, grad_functor) \
2121
REGISTER_OP_CUDA_KERNEL( \
2222
reduce_type, ops::ReduceKernel<paddle::platform::CUDADeviceContext, \
23-
float, ops::functor>); \
23+
float, ops::functor>, \
24+
ops::ReduceKernel<paddle::platform::CUDADeviceContext, double, \
25+
ops::functor>, \
26+
ops::ReduceKernel<paddle::platform::CUDADeviceContext, int, \
27+
ops::functor>, \
28+
ops::ReduceKernel<paddle::platform::CUDADeviceContext, int64_t, \
29+
ops::functor>); \
2430
REGISTER_OP_CUDA_KERNEL( \
2531
reduce_type##_grad, \
2632
ops::ReduceGradKernel<paddle::platform::CUDADeviceContext, float, \
33+
ops::grad_functor>, \
34+
ops::ReduceGradKernel<paddle::platform::CUDADeviceContext, double, \
35+
ops::grad_functor>, \
36+
ops::ReduceGradKernel<paddle::platform::CUDADeviceContext, int, \
37+
ops::grad_functor>, \
38+
ops::ReduceGradKernel<paddle::platform::CUDADeviceContext, int64_t, \
2739
ops::grad_functor>);
2840

2941
FOR_EACH_KERNEL_FUNCTOR(REGISTER_REDUCE_GPU_KERNEL);

python/paddle/v2/fluid/tests/test_reduce_op.py

Lines changed: 7 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -19,7 +19,7 @@
1919
class TestSumOp(OpTest):
2020
def setUp(self):
2121
self.op_type = "reduce_sum"
22-
self.inputs = {'X': np.random.random((5, 6, 10)).astype("float32")}
22+
self.inputs = {'X': np.random.random((5, 6, 10)).astype("float64")}
2323
self.outputs = {'Out': self.inputs['X'].sum(axis=0)}
2424

2525
def test_check_output(self):
@@ -32,7 +32,7 @@ def test_check_grad(self):
3232
class TestMeanOp(OpTest):
3333
def setUp(self):
3434
self.op_type = "reduce_mean"
35-
self.inputs = {'X': np.random.random((5, 6, 2, 10)).astype("float32")}
35+
self.inputs = {'X': np.random.random((5, 6, 2, 10)).astype("float64")}
3636
self.attrs = {'dim': 1}
3737
self.outputs = {'Out': self.inputs['X'].mean(axis=self.attrs['dim'])}
3838

@@ -48,7 +48,7 @@ class TestMaxOp(OpTest):
4848

4949
def setUp(self):
5050
self.op_type = "reduce_max"
51-
self.inputs = {'X': np.random.random((5, 6, 10)).astype("float32")}
51+
self.inputs = {'X': np.random.random((5, 6, 10)).astype("float64")}
5252
self.attrs = {'dim': -1}
5353
self.outputs = {'Out': self.inputs['X'].max(axis=self.attrs['dim'])}
5454

@@ -61,7 +61,7 @@ class TestMinOp(OpTest):
6161

6262
def setUp(self):
6363
self.op_type = "reduce_min"
64-
self.inputs = {'X': np.random.random((5, 6, 10)).astype("float32")}
64+
self.inputs = {'X': np.random.random((5, 6, 10)).astype("float64")}
6565
self.attrs = {'dim': 2}
6666
self.outputs = {'Out': self.inputs['X'].min(axis=self.attrs['dim'])}
6767

@@ -72,7 +72,7 @@ def test_check_output(self):
7272
class TestKeepDimReduce(OpTest):
7373
def setUp(self):
7474
self.op_type = "reduce_sum"
75-
self.inputs = {'X': np.random.random((5, 6, 10)).astype("float32")}
75+
self.inputs = {'X': np.random.random((5, 6, 10)).astype("float64")}
7676
self.attrs = {'dim': -2, 'keep_dim': True}
7777
self.outputs = {
7878
'Out': self.inputs['X'].sum(axis=self.attrs['dim'], keepdims=True)
@@ -88,7 +88,7 @@ def test_check_grad(self):
8888
class Test1DReduce(OpTest):
8989
def setUp(self):
9090
self.op_type = "reduce_sum"
91-
self.inputs = {'X': np.random.random(20).astype("float32")}
91+
self.inputs = {'X': np.random.random(20).astype("float64")}
9292
self.outputs = {'Out': self.inputs['X'].sum(axis=0)}
9393

9494
def test_check_output(self):
@@ -101,7 +101,7 @@ def test_check_grad(self):
101101
class TestReduceAll(OpTest):
102102
def setUp(self):
103103
self.op_type = "reduce_sum"
104-
self.inputs = {'X': np.random.random((5, 6, 2, 10)).astype("float32")}
104+
self.inputs = {'X': np.random.random((5, 6, 2, 10)).astype("float64")}
105105
self.attrs = {'reduce_all': True}
106106
self.outputs = {'Out': self.inputs['X'].sum()}
107107

0 commit comments

Comments
 (0)