Skip to content

Commit ad41fce

Browse files
Silv3SSilv3S
authored
[PHI] Migrate squeeze and squeeze_grad kernels (#48634)
* squeeze kernel * squeze fwd * whitespace
1 parent 4aad4dc commit ad41fce

File tree

3 files changed

+147
-36
lines changed

3 files changed

+147
-36
lines changed

‎paddle/fluid/operators/mkldnn/reshape_mkldnn_op.cc‎

Lines changed: 3 additions & 36 deletions
Original file line numberDiff line numberDiff line change
@@ -21,7 +21,6 @@ enum class ReshapeKernelOpName {
2121
reshape,
2222
reshape2,
2323
squeeze,
24-
squeeze2,
2524
flatten,
2625
flatten2,
2726
};
@@ -106,9 +105,6 @@ class ReshapeMKLDNNKernel : public framework::OpKernel<T> {
106105
case ReshapeKernelOpName::squeeze:
107106
InferShapeSqueezeOp(ctx, x_dims, out_dims);
108107
break;
109-
case ReshapeKernelOpName::squeeze2:
110-
InferShapeSqueeze2Op(ctx, x_dims, out_dims);
111-
break;
112108
case ReshapeKernelOpName::flatten:
113109
InferShapeFlattenOp(ctx, x_dims, out_dims);
114110
break;
@@ -172,16 +168,6 @@ class ReshapeMKLDNNKernel : public framework::OpKernel<T> {
172168
out_dims = GetOutputShape(axes, x_dims, true);
173169
}
174170

175-
void InferShapeSqueeze2Op(const framework::ExecutionContext& ctx,
176-
framework::DDim& x_dims, // NOLINT
177-
framework::DDim& out_dims) const { // NOLINT
178-
auto* out = ctx.Output<phi::DenseTensor>("Out");
179-
auto* xshape = ctx.Output<phi::DenseTensor>("XShape");
180-
auto xshape_dims = xshape->dims();
181-
x_dims = phi::slice_ddim(xshape_dims, 1, xshape_dims.size());
182-
out_dims = out->dims();
183-
}
184-
185171
void InferShapeFlattenOp(const framework::ExecutionContext& ctx,
186172
framework::DDim& x_dims, // NOLINT
187173
framework::DDim& out_dims) const { // NOLINT
@@ -342,19 +328,16 @@ class ReshapeGradMKLDNNKernel : public ReshapeMKLDNNKernel<T, op_name> {
342328
InferShapeReshapeSqueezeGradOp(ctx, x_dims);
343329
break;
344330
case ReshapeKernelOpName::reshape2:
345-
InferShapeReshape2Squeeze2Flatten2GradOp(ctx, x_dims);
331+
InferShapeReshape2Flatten2GradOp(ctx, x_dims);
346332
break;
347333
case ReshapeKernelOpName::squeeze:
348334
InferShapeReshapeSqueezeGradOp(ctx, x_dims);
349335
break;
350-
case ReshapeKernelOpName::squeeze2:
351-
InferShapeReshape2Squeeze2Flatten2GradOp(ctx, x_dims);
352-
break;
353336
case ReshapeKernelOpName::flatten:
354337
InferShapeFlattenGradOp(ctx, x_dims);
355338
break;
356339
case ReshapeKernelOpName::flatten2:
357-
InferShapeReshape2Squeeze2Flatten2GradOp(ctx, x_dims);
340+
InferShapeReshape2Flatten2GradOp(ctx, x_dims);
358341
break;
359342
default:
360343
PADDLE_THROW(paddle::platform::errors::OutOfRange(
@@ -369,7 +352,7 @@ class ReshapeGradMKLDNNKernel : public ReshapeMKLDNNKernel<T, op_name> {
369352
dx_dims = dx->dims();
370353
}
371354

372-
void InferShapeReshape2Squeeze2Flatten2GradOp(
355+
void InferShapeReshape2Flatten2GradOp(
373356
const framework::ExecutionContext& ctx,
374357
framework::DDim& dx_dims) const { // NOLINT
375358
auto xshape_dims = ctx.Input<phi::DenseTensor>("XShape")->dims();
@@ -401,22 +384,6 @@ REGISTER_OP_KERNEL(
401384
ops::ReshapeGradMKLDNNKernel<paddle::platform::bfloat16,
402385
ReshapeKernelOpName::squeeze>);
403386

404-
REGISTER_OP_KERNEL(
405-
squeeze2,
406-
MKLDNN,
407-
paddle::platform::CPUPlace,
408-
ops::ReshapeMKLDNNKernel<float, ReshapeKernelOpName::squeeze2>,
409-
ops::ReshapeMKLDNNKernel<paddle::platform::bfloat16,
410-
ReshapeKernelOpName::squeeze2>);
411-
412-
REGISTER_OP_KERNEL(
413-
squeeze2_grad,
414-
MKLDNN,
415-
paddle::platform::CPUPlace,
416-
ops::ReshapeGradMKLDNNKernel<float, ReshapeKernelOpName::squeeze2>,
417-
ops::ReshapeGradMKLDNNKernel<paddle::platform::bfloat16,
418-
ReshapeKernelOpName::squeeze2>);
419-
420387
REGISTER_OP_KERNEL(
421388
reshape,
422389
MKLDNN,

‎paddle/phi/kernels/onednn/squeeze_grad_kernel.cc‎

Lines changed: 59 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,59 @@
1+
// Copyright (c) 2022 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+
#include "paddle/phi/kernels/squeeze_grad_kernel.h"
16+
17+
#include "paddle/phi/backends/onednn/onednn_reuse.h"
18+
#include "paddle/phi/core/kernel_registry.h"
19+
20+
namespace phi {
21+
22+
template <typename T, typename Context>
23+
void SqueezeGradKernel(const Context& dev_ctx,
24+
const DenseTensor& xshape,
25+
const DenseTensor& dout,
26+
const IntArray& axes,
27+
DenseTensor* dx) {
28+
auto dout_vec_dims = vectorize(dout.dims());
29+
auto dout_type = funcs::ToOneDNNDataType(dout.dtype());
30+
31+
funcs::ReorderOneDNNHandler reorder_handler(
32+
dout_vec_dims, dout.dtype(), dout_type, dev_ctx.GetEngine());
33+
34+
auto reorder_src_memory_p = reorder_handler.AcquireSrcMemory(
35+
dout.mem_desc(), funcs::to_void_cast(dout.data<T>()));
36+
auto reorder_dst_memory_p = reorder_handler.AcquireDstMemory(
37+
dx,
38+
funcs::GetPlainOneDNNFormat(dout_vec_dims.size()),
39+
dev_ctx.GetPlace());
40+
auto reorder_p = reorder_handler.AcquireReorder(reorder_dst_memory_p,
41+
reorder_src_memory_p);
42+
43+
auto& astream = OneDNNContext::tls().get_stream();
44+
reorder_p->execute(astream, *reorder_src_memory_p, *reorder_dst_memory_p);
45+
astream.wait();
46+
47+
auto dx_dims = slice_ddim(xshape.dims(), 1, xshape.dims().size());
48+
dx->Resize(dx_dims);
49+
reorder_dst_memory_p->get_desc().reshape(vectorize(dx_dims));
50+
}
51+
52+
} // namespace phi
53+
54+
PD_REGISTER_KERNEL(squeeze_grad,
55+
OneDNN,
56+
ONEDNN,
57+
phi::SqueezeGradKernel,
58+
float,
59+
phi::dtype::bfloat16) {}

‎paddle/phi/kernels/onednn/squeeze_kernel.cc‎

Lines changed: 85 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,85 @@
1+
// Copyright (c) 2022 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+
#include "paddle/phi/kernels/squeeze_kernel.h"
16+
17+
#include "paddle/phi/backends/onednn/onednn_reuse.h"
18+
#include "paddle/phi/core/kernel_registry.h"
19+
#include "paddle/phi/kernels/funcs/unsqueeze.h"
20+
21+
namespace phi {
22+
23+
template <typename T, typename Context>
24+
void ExecuteSqueeze(const Context& dev_ctx,
25+
const DenseTensor& x,
26+
const DDim& x_dims,
27+
const DDim& out_dims,
28+
DenseTensor* out) {
29+
auto x_vec_dims = vectorize(x_dims);
30+
31+
funcs::ReorderOneDNNHandler reorder_handler(
32+
x_vec_dims,
33+
x.dtype(),
34+
funcs::ToOneDNNDataType(x.dtype()),
35+
dev_ctx.GetEngine());
36+
37+
auto reorder_src_memory_p = reorder_handler.AcquireSrcMemory(
38+
x.mem_desc(), funcs::to_void_cast(x.data<T>()));
39+
out->Resize(x_dims); // to match x numel, format is changed later
40+
// reorder is done into a plain tag to allow usage with blocked formats
41+
auto reorder_dst_memory_p = reorder_handler.AcquireDstMemory(
42+
out, funcs::GetPlainOneDNNFormat(x_dims.size()), dev_ctx.GetPlace());
43+
auto reorder_p = reorder_handler.AcquireReorder(reorder_dst_memory_p,
44+
reorder_src_memory_p);
45+
auto& astream = OneDNNContext::tls().get_stream();
46+
reorder_p->execute(astream, *reorder_src_memory_p, *reorder_dst_memory_p);
47+
astream.wait();
48+
49+
out->Resize(out_dims);
50+
out->set_mem_desc(
51+
reorder_dst_memory_p->get_desc().reshape(vectorize(out_dims)));
52+
}
53+
54+
template <typename T, typename Context>
55+
void SqueezeKernel(const Context& dev_ctx,
56+
const DenseTensor& x,
57+
const IntArray& axes,
58+
DenseTensor* out) {
59+
auto x_dims = x.dims();
60+
std::vector<int32_t> tmp(axes.GetData().begin(), axes.GetData().end());
61+
auto out_dims = funcs::GetOutputSqueezeShape(tmp, x_dims, true);
62+
ExecuteSqueeze<T, Context>(dev_ctx, x, x_dims, out_dims, out);
63+
}
64+
65+
template <typename T, typename Context>
66+
void SqueezeWithXShapeKernel(const Context& dev_ctx,
67+
const DenseTensor& x,
68+
const IntArray& axes,
69+
DenseTensor* out,
70+
DenseTensor* xshape) {
71+
auto x_dims = slice_ddim(xshape->dims(), 1, xshape->dims().size());
72+
auto out_dims = out->dims();
73+
ExecuteSqueeze<T, Context>(dev_ctx, x, x_dims, out_dims, out);
74+
}
75+
} // namespace phi
76+
77+
PD_REGISTER_KERNEL(
78+
squeeze, OneDNN, ONEDNN, phi::SqueezeKernel, float, phi::dtype::bfloat16) {}
79+
80+
PD_REGISTER_KERNEL(squeeze_with_xshape,
81+
OneDNN,
82+
ONEDNN,
83+
phi::SqueezeWithXShapeKernel,
84+
float,
85+
phi::dtype::bfloat16) {}

0 commit comments

Comments
 (0)