move Masked select to pten

paddle/fluid/operators/masked_select_op.cc

@@ -12,7 +12,6 @@ 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/operators/masked_select_op.h"
 #include "paddle/fluid/framework/op_registry.h"
 
 namespace paddle {
@@ -106,16 +105,3 @@ REGISTER_OPERATOR(masked_select, ops::MaskedSelectOp, ops::MaskedSelectOpMaker,
                   ops::MaskedSelectGradOpMaker<paddle::imperative::OpBase>);
 REGISTER_OPERATOR(masked_select_grad, ops::MaskedSelectOpGrad,
                   ops::MaskedSelectedGradNoNeedBufferVarsInferer);
-
-REGISTER_OP_CPU_KERNEL(
-    masked_select,
-    ops::MaskedSelectKernel<paddle::platform::CPUDeviceContext, float>,
-    ops::MaskedSelectKernel<paddle::platform::CPUDeviceContext, double>,
-    ops::MaskedSelectKernel<paddle::platform::CPUDeviceContext, int>,
-    ops::MaskedSelectKernel<paddle::platform::CPUDeviceContext, int64_t>);
-REGISTER_OP_CPU_KERNEL(
-    masked_select_grad,
-    ops::MaskedSelectGradKernel<paddle::platform::CPUDeviceContext, float>,
-    ops::MaskedSelectGradKernel<paddle::platform::CPUDeviceContext, double>,
-    ops::MaskedSelectGradKernel<paddle::platform::CPUDeviceContext, int>,
-    ops::MaskedSelectGradKernel<paddle::platform::CPUDeviceContext, int64_t>);

paddle/fluid/operators/masked_select_op_npu.cc

@@ -12,7 +12,7 @@ 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/operators/masked_select_op.h"
+#include "paddle/fluid/framework/op_registry.h"
 #include "paddle/fluid/platform/device/npu/npu_op_runner.h"
 
 namespace paddle {

paddle/fluid/operators/masked_select_op_xpu.cc

@@ -11,7 +11,7 @@ limitations under the License. */
 
 #ifdef PADDLE_WITH_XPU
 
-#include "paddle/fluid/operators/masked_select_op.h"
+#include "paddle/fluid/framework/op_registry.h"
 #include "paddle/fluid/platform/device/device_wrapper.h"
 
 namespace paddle {

paddle/pten/kernels/cpu/masked_select_grad_kernel.cc

@@ -0,0 +1,53 @@
+// Copyright (c) 2022 PaddlePaddle Authors. All Rights Reserved.
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// 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/pten/kernels/masked_select_grad_kernel.h"
+
+#include "paddle/pten/backends/cpu/cpu_context.h"
+#include "paddle/pten/core/kernel_registry.h"
+
+namespace pten {
+
+template <typename T, typename Context>
+void MaskedSelectGradKernel(const Context& dev_ctx,
+                            const DenseTensor& out_grad,
+                            const DenseTensor& x,
+                            const DenseTensor& mask,
+                            DenseTensor* x_grad) {
+  auto* mask_data = mask.data<bool>();
+  auto* input_data = out_grad.data<T>();
+  auto* out_data = x_grad->mutable_data<T>(dev_ctx.GetPlace());
+  int mask_size = mask.numel();
+
+  int index = 0;
+  for (int i = 0; i < mask_size; i++) {
+    if (mask_data[i]) {
+      out_data[i] = input_data[index];
+      index++;
+    } else {
+      out_data[i] = 0;
+    }
+  }
+}
+
+}  // namespace pten
+
+PT_REGISTER_KERNEL(masked_select_grad,
+                   CPU,
+                   ALL_LAYOUT,
+                   pten::MaskedSelectGradKernel,
+                   float,
+                   double,
+                   int,
+                   int64_t) {}