【Hackathon 6th Fundable Projects 3 No.241】 [fluid_ops] straight_throu…

…gh_estimator_grad (#67299)

paddle/fluid/operators/fake_quantize_op.cc

@@ -147,19 +147,9 @@ REGISTER_OPERATOR(
     ops::MovingAverageAbsMaxScaleOpMaker,
     ops::StraightThroughEstimatorMaker<paddle::framework::OpDesc>,
     ops::StraightThroughEstimatorMaker<paddle::imperative::OpBase>);
-PD_REGISTER_STRUCT_KERNEL(moving_average_abs_max_scale,
-                          CPU,
-                          ALL_LAYOUT,
-                          ops::MovingAverageAbsMaxScaleKernel,
-                          float) {}
 
 REGISTER_OPERATOR(straight_through_estimator_grad,
                   ops::StraightThroughEstimatorGradOp);
-PD_REGISTER_STRUCT_KERNEL(straight_through_estimator_grad,
-                          CPU,
-                          ALL_LAYOUT,
-                          ops::StraightThroughEstimatorGradKernel,
-                          float) {}
 
 REGISTER_OP_VERSION(fake_channel_wise_quantize_abs_max)
     .AddCheckpoint(

paddle/fluid/operators/ops_signature/moving_average_abs_max_scale_sig.cc

@@ -0,0 +1,31 @@
+// Copyright (c) 2024 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/phi/core/compat/op_utils.h"
+
+namespace phi {
+
+// we have to return every specific KernelSignature for inference now
+KernelSignature MovingAverageAbsMaxScaleOpArgumentMapping(
+    const ArgumentMappingContext& ctx) {
+  return KernelSignature("moving_average_abs_max_scale",
+                         {"X", "InAccum", "InState"},
+                         {"moving_rate", "is_test"},
+                         {"Out", "OutScale", "OutState", "OutAccum"});
+}
+
+}  // namespace phi
+
+PD_REGISTER_ARG_MAPPING_FN(moving_average_abs_max_scale,
+                           phi::MovingAverageAbsMaxScaleOpArgumentMapping);

paddle/fluid/operators/ops_signature/straight_through_estimator_grad_sig.cc

@@ -0,0 +1,28 @@
+// Copyright (c) 2024 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/phi/core/compat/op_utils.h"
+
+namespace phi {
+
+KernelSignature StraightThroughEstimatorGradOpArgumentMapping(
+    const ArgumentMappingContext& ctx) {
+  return KernelSignature(
+      "straight_through_estimator_grad", {"Out@GRAD"}, {}, {"X@GRAD"});
+}
+
+}  // namespace phi
+
+PD_REGISTER_ARG_MAPPING_FN(straight_through_estimator_grad,
+                           phi::StraightThroughEstimatorGradOpArgumentMapping);

paddle/fluid/pir/dialect/operator/utils/utils.cc

@@ -61,8 +61,6 @@ const std::unordered_set<std::string> LegacyOpList = {
     PushDenseOp::name(),
     SoftReluOp::name(),
     SoftReluGradOp::name(),
-    MovingAverageAbsMaxScaleOp::name(),
-    MovingAverageAbsMaxScale_Op::name(),
     CReduceAvgOp::name(),
     CReduceAvg_Op::name(),
     CReduceMaxOp::name(),

paddle/phi/infermeta/unary.cc

@@ -5990,6 +5990,12 @@ void StridedUnChangedInferMeta(const MetaTensor& x, MetaTensor* out) {
   out->set_strides(x.strides());
 }
 
+void StraightThroughEstimatorInferMeta(const MetaTensor& out_grad,
+                                       MetaTensor* x_grad) {
+  x_grad->set_dims(out_grad.dims());
+  x_grad->set_dtype(out_grad.dtype());
+}
+
 void NumberCountInferMeta(const MetaTensor& x,
                           int upper_range,
                           MetaTensor* out) {

paddle/phi/infermeta/unary.h

@@ -956,6 +956,9 @@ void NumberCountInferMeta(const MetaTensor& x,
 
 void StridedUnChangedInferMeta(const MetaTensor& x, MetaTensor* out);
 
+void StraightThroughEstimatorInferMeta(const MetaTensor& out_grad,
+                                       MetaTensor* x_grad);
+
 void LrnInferMeta(const MetaTensor& x,
                   int n,
                   MetaTensor* out,

paddle/phi/kernels/cpu/moving_average_abs_max_scale_kernel.cc

@@ -0,0 +1,22 @@
+// Copyright (c) 2024 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/phi/core/kernel_registry.h"
+#include "paddle/phi/kernels/impl/moving_average_abs_max_scale_kernel_impl.h"
+
+PD_REGISTER_KERNEL(moving_average_abs_max_scale,
+                   CPU,
+                   ALL_LAYOUT,
+                   phi::MovingAverageAbsMaxScaleKernel,
+                   float) {}

paddle/phi/kernels/cpu/straight_through_estimator_grad_kernel.cc

@@ -0,0 +1,21 @@
+// Copyright (c) 2024 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/phi/kernels/impl/straight_through_estimator_grad_kernel_impl.h"
+
+PD_REGISTER_KERNEL(straight_through_estimator_grad,
+                   CPU,
+                   ALL_LAYOUT,
+                   phi::StraightThroughEstimatorGradKernel,
+                   float) {}

paddle/phi/kernels/gpu/moving_average_abs_max_scale_kernel.cu

@@ -0,0 +1,23 @@
+// Copyright (c) 2024 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/phi/core/kernel_registry.h"
+#include "paddle/phi/kernels/impl/moving_average_abs_max_scale_kernel_impl.h"
+
+PD_REGISTER_KERNEL(moving_average_abs_max_scale,
+                   GPU,
+                   ALL_LAYOUT,
+                   phi::MovingAverageAbsMaxScaleKernel,
+                   float,
+                   phi::dtype::float16) {}

paddle/phi/kernels/gpu/straight_through_estimator_grad_kernel.cu

@@ -0,0 +1,22 @@
+// Copyright (c) 2024 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/phi/kernels/impl/straight_through_estimator_grad_kernel_impl.h"
+
+PD_REGISTER_KERNEL(straight_through_estimator_grad,
+                   GPU,
+                   ALL_LAYOUT,
+                   phi::StraightThroughEstimatorGradKernel,
+                   float,
+                   phi::dtype::float16) {}