Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Jump to bottom

migrate reduce_amin,reduce_amax kernel to phi #44698

Merged
merged 1 commit into from
Aug 1, 2022
Merged

migrate reduce_amin,reduce_amax kernel to phi #44698

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
36 changes: 22 additions & 14 deletions paddle/fluid/operators/reduce_ops/reduce_amax_op.cc
View file
Open in desktop
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
Open in desktop

This file was deleted.

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

This file was deleted.

36 changes: 22 additions & 14 deletions paddle/fluid/operators/reduce_ops/reduce_amin_op.cc
View file
Open in desktop
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
Open in desktop

This file was deleted.

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

This file was deleted.

115 changes: 0 additions & 115 deletions paddle/fluid/operators/reduce_ops/reduce_min_max_op.h
View file
Open in desktop
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