Skip to content

Commit

Permalink
migrate reduce_amin,reduce_amax kernel to phi
Browse files Browse the repository at this point in the history
  • Loading branch information
@cxxly
cxxly committed Jul 29, 2022
1 parent 72b65d6 commit 8e1c893
Show file tree
Hide file tree
Showing 28 changed files with 953 additions and 348 deletions.
36 changes: 22 additions & 14 deletions paddle/fluid/operators/reduce_ops/reduce_amax_op.cc
View file
Original file line number Diff line number Diff line change
Expand Up @@ -11,20 +11,28 @@
// 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 "paddle/fluid/framework/infershape_utils.h"
#include "paddle/fluid/operators/reduce_ops/reduce_min_max_op.h"
#include "paddle/phi/core/infermeta_utils.h"
#include "paddle/phi/infermeta/unary.h"

namespace ops = paddle::operators;

class ReduceAMaxOpMaker : public ops::ReduceOpMaker {
protected:
virtual std::string GetName() const { return "reduce_amax"; }
virtual std::string GetOpType() const { return "Reduce reduce_amax"; }
};

DECLARE_INFER_SHAPE_FUNCTOR(reduce_amax,
ReduceAMaxInferShapeFunctor,
PD_INFER_META(phi::ReduceInferMetaBase));

REGISTER_REDUCE_OP(reduce_amax);
REGISTER_OP_CPU_KERNEL(
REGISTER_OPERATOR(
reduce_amax,
ops::ReduceKernel<phi::CPUContext, float, ops::MaxFunctor>,
ops::ReduceKernel<phi::CPUContext, double, ops::MaxFunctor>,
ops::ReduceKernel<phi::CPUContext, int, ops::MaxFunctor>,
ops::ReduceKernel<phi::CPUContext, int64_t, ops::MaxFunctor>);
REGISTER_OP_CPU_KERNEL(
reduce_amax_grad,
ops::ReduceGradKernel<phi::CPUContext, float, ops::AMaxOrAMinGradFunctor>,
ops::ReduceGradKernel<phi::CPUContext, double, ops::AMaxOrAMinGradFunctor>,
ops::ReduceGradKernel<phi::CPUContext, int, ops::AMaxOrAMinGradFunctor>,
ops::
ReduceGradKernel<phi::CPUContext, int64_t, ops::AMaxOrAMinGradFunctor>);
ops::ReduceOp,
ReduceAMaxOpMaker,
paddle::framework::DefaultGradOpMaker<paddle::framework::OpDesc, true>,
paddle::framework::DefaultGradOpMaker<paddle::imperative::OpBase, true>,
ReduceAMaxInferShapeFunctor);
REGISTER_OPERATOR(reduce_amax_grad, ops::ReduceGradOp)
36 changes: 0 additions & 36 deletions paddle/fluid/operators/reduce_ops/reduce_amax_op.kps
View file

This file was deleted.

24 changes: 0 additions & 24 deletions paddle/fluid/operators/reduce_ops/reduce_amax_op.part.cu
View file

This file was deleted.

36 changes: 22 additions & 14 deletions paddle/fluid/operators/reduce_ops/reduce_amin_op.cc
View file
Original file line number Diff line number Diff line change
Expand Up @@ -11,20 +11,28 @@
// 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 "paddle/fluid/framework/infershape_utils.h"
#include "paddle/fluid/operators/reduce_ops/reduce_min_max_op.h"
#include "paddle/phi/core/infermeta_utils.h"
#include "paddle/phi/infermeta/unary.h"

namespace ops = paddle::operators;

class ReduceAMinOpMaker : public ops::ReduceOpMaker {
protected:
virtual std::string GetName() const { return "reduce_amin"; }
virtual std::string GetOpType() const { return "Reduce reduce_amin"; }
};

DECLARE_INFER_SHAPE_FUNCTOR(reduce_amin,
ReduceAMinInferShapeFunctor,
PD_INFER_META(phi::ReduceInferMetaBase));

REGISTER_REDUCE_OP(reduce_amin);
REGISTER_OP_CPU_KERNEL(
REGISTER_OPERATOR(
reduce_amin,
ops::ReduceKernel<phi::CPUContext, float, ops::MinFunctor>,
ops::ReduceKernel<phi::CPUContext, double, ops::MinFunctor>,
ops::ReduceKernel<phi::CPUContext, int, ops::MinFunctor>,
ops::ReduceKernel<phi::CPUContext, int64_t, ops::MinFunctor>);
REGISTER_OP_CPU_KERNEL(
reduce_amin_grad,
ops::ReduceGradKernel<phi::CPUContext, float, ops::AMaxOrAMinGradFunctor>,
ops::ReduceGradKernel<phi::CPUContext, double, ops::AMaxOrAMinGradFunctor>,
ops::ReduceGradKernel<phi::CPUContext, int, ops::AMaxOrAMinGradFunctor>,
ops::
ReduceGradKernel<phi::CPUContext, int64_t, ops::AMaxOrAMinGradFunctor>);
ops::ReduceOp,
ReduceAMinOpMaker,
paddle::framework::DefaultGradOpMaker<paddle::framework::OpDesc, true>,
paddle::framework::DefaultGradOpMaker<paddle::imperative::OpBase, true>,
ReduceAMinInferShapeFunctor);
REGISTER_OPERATOR(reduce_amin_grad, ops::ReduceGradOp)
36 changes: 0 additions & 36 deletions paddle/fluid/operators/reduce_ops/reduce_amin_op.kps
View file

This file was deleted.

24 changes: 0 additions & 24 deletions paddle/fluid/operators/reduce_ops/reduce_amin_op.part.cu
View file

This file was deleted.

115 changes: 0 additions & 115 deletions paddle/fluid/operators/reduce_ops/reduce_min_max_op.h
View file
Original file line number Diff line number Diff line change
Expand Up @@ -55,120 +55,5 @@ struct MaxOrMinGradFunctor {
}
};

#define HANDLE_AXIS_DIM(BROADCAST_DIM, AXIS_DIM) \
if (broadcast_dim_size == BROADCAST_DIM && rank == AXIS_DIM) { \
AMaxOrAMinAxisIsListGradFunctor<DeviceContext, \
X, \
Y, \
DX, \
DY, \
Dim, \
BROADCAST_DIM, \
AXIS_DIM>( \
place, x, y, dx, dy, dim, axis_dim); \
}

template <typename DeviceContext,
typename X,
typename Y,
typename DX,
typename DY,
typename Dim,
int R,
int D>
void AMaxOrAMinAxisIsListGradFunctor(const DeviceContext& place,
X* x,
Y* y,
DX* dx,
DY* dy,
const Dim& dim,
const std::vector<int>& axis_dim) {
// R is x->dimensions().size();
// D is axis_dim->dimensions().size();
auto axis = Eigen::array<int, D>();
auto reshape_x = Eigen::array<int, R>();
auto reshape_y = Eigen::array<int, R>();

for (int i = 0; i < D; i++) axis[i] = axis_dim[i];
for (int i = 0; i < R; i++) {
reshape_x[i] = x->dimensions()[i];
reshape_y[i] = y->dimensions()[i];
}

auto equals = (*x) == y->broadcast(dim);
auto ones = dx->constant(1);
auto zeros = dx->constant(0);
auto mask = equals.select(ones, zeros);
dx->device(place) =
dy->broadcast(dim) * mask /
mask.reshape(reshape_x).sum(axis).reshape(reshape_y).broadcast(dim);
}

struct AMaxOrAMinGradFunctor {
template <typename DeviceContext,
typename X,
typename Y,
typename DX,
typename DY,
typename Dim>
void operator()(const DeviceContext& place,
X* x,
Y* y,
DX* dx,
DY* dy,
const Dim& dim,
int size) {
auto equals = (*x) == y->broadcast(dim);
auto ones = dx->constant(1);
auto zeros = dx->constant(0);
auto mask = equals.select(ones, zeros);

// If there are multiple minimum or maximum elements,
// we evenly distribute gradient between these equal values
size_t x_numel = 1;
for (size_t i = 0; i < x->dimensions().size(); i++)
x_numel *= x->dimensions()[i];
// reduce_all
if (size == static_cast<int>(x_numel)) {
auto equal_number = mask.sum()
.reshape(Eigen::array<int, 1>({1}))
.broadcast(Eigen::array<int, 1>({size}));
dx->device(place) = dy->broadcast(dim) * mask / equal_number;
return;
}

// compute forward reduce axis_dim by dim (which is broadcast_dim)
std::vector<int> axis_dim;
int broadcast_dim_size = static_cast<int>(dim.size());
for (int i = 0; i < broadcast_dim_size; i++) {
if (dim[i] > 1) {
axis_dim.push_back(i);
}
}

int rank = static_cast<int>(axis_dim.size());
// axis is a int element
if (rank == 1) {
auto axis = Eigen::array<int, 1>({axis_dim[0]});
dx->device(place) =
dy->broadcast(dim) * mask /
mask.sum(axis).reshape(dy->dimensions()).broadcast(dim);
return;
}
// axis is list, HANDLE_AXIS_DIM(broadcast_dim_size, rank)
HANDLE_AXIS_DIM(3, 2);
HANDLE_AXIS_DIM(4, 2);
HANDLE_AXIS_DIM(4, 3);
// comments for accelerating compiling temporarily.
// HANDLE_AXIS_DIM(5, 2);
// HANDLE_AXIS_DIM(5, 3);
// HANDLE_AXIS_DIM(5, 4);
// HANDLE_AXIS_DIM(6, 2);
// HANDLE_AXIS_DIM(6, 3);
// HANDLE_AXIS_DIM(6, 4);
// HANDLE_AXIS_DIM(6, 5);
}
};

} // namespace operators
} // namespace paddle
Loading

0 comments on commit 8e1c893

Please sign in to comment.