【Hackathon 6th No.54】move fake_quantize_moving_average_abs_max to phi - part (#63732)

zeroRains · web-flow · commit 0bc130a4e54d · 2024-04-26T00:24:17.000+08:00
diff --git a/paddle/fluid/operators/fake_quantize_op.cc b/paddle/fluid/operators/fake_quantize_op.cc
@@ -316,32 +316,6 @@ struct FindRangeAbsMaxFunctor<phi::CPUContext, T> {
 
 template struct FindRangeAbsMaxFunctor<phi::CPUContext, float>;
 
-template <typename T>
-struct FindMovingAverageAbsMaxFunctor<phi::CPUContext, T> {
-  void operator()(const phi::CPUContext &ctx,
-                  const phi::DenseTensor &in_accum,
-                  const phi::DenseTensor &in_state,
-                  const T *cur_scale,
-                  const float rate,
-                  phi::DenseTensor *out_state,
-                  phi::DenseTensor *out_accum,
-                  phi::DenseTensor *out_scale) {
-    T accum = in_accum.data<T>()[0];
-    T state = in_state.data<T>()[0];
-    T scale = cur_scale[0];
-
-    state = rate * state + 1;
-    accum = rate * accum + scale;
-    scale = accum / state;
-
-    out_state->mutable_data<T>(ctx.GetPlace())[0] = state;
-    out_accum->mutable_data<T>(ctx.GetPlace())[0] = accum;
-    out_scale->mutable_data<T>(ctx.GetPlace())[0] = scale;
-  }
-};
-
-template struct FindMovingAverageAbsMaxFunctor<phi::CPUContext, float>;
-
 class FakeQuantOrWithDequantAbsMaxOp : public framework::OperatorWithKernel {
  public:
   FakeQuantOrWithDequantAbsMaxOp(const std::string &type,
@@ -858,18 +832,6 @@ PD_REGISTER_STRUCT_KERNEL(fake_quantize_range_abs_max,
                           ops::FakeQuantizeRangeAbsMaxKernel,
                           float) {}
 
-REGISTER_OPERATOR(
-    fake_quantize_moving_average_abs_max,
-    ops::FakeQuantOrWithDequantMovingAverageAbsMaxOp,
-    ops::FakeQuantOrWithDequantMovingAverageAbsMaxOpMaker,
-    paddle::framework::EmptyGradOpMaker<paddle::framework::OpDesc>,
-    paddle::framework::EmptyGradOpMaker<paddle::imperative::OpBase>);
-PD_REGISTER_STRUCT_KERNEL(fake_quantize_moving_average_abs_max,
-                          CPU,
-                          ALL_LAYOUT,
-                          ops::FakeQuantizeMovingAverageAbsMaxKernel,
-                          float) {}
-
 REGISTER_OPERATOR(
     fake_quantize_dequantize_moving_average_abs_max,
     ops::FakeQuantOrWithDequantMovingAverageAbsMaxOp,
diff --git a/paddle/fluid/operators/fake_quantize_op.cu b/paddle/fluid/operators/fake_quantize_op.cu
@@ -36,12 +36,6 @@ PD_REGISTER_STRUCT_KERNEL(fake_quantize_range_abs_max,
                           ops::FakeQuantizeRangeAbsMaxKernel,
                           float,
                           float16) {}
-PD_REGISTER_STRUCT_KERNEL(fake_quantize_moving_average_abs_max,
-                          GPU,
-                          ALL_LAYOUT,
-                          ops::FakeQuantizeMovingAverageAbsMaxKernel,
-                          float,
-                          float16) {}
 PD_REGISTER_STRUCT_KERNEL(moving_average_abs_max_scale,
                           GPU,
                           ALL_LAYOUT,
diff --git a/paddle/fluid/operators/fake_quantize_op.cu.h b/paddle/fluid/operators/fake_quantize_op.cu.h
@@ -494,51 +494,8 @@ struct FindRangeAbsMaxFunctor<phi::GPUContext, T> {
   }
 };
 
-template <typename T>
-__global__ void FindMovingAverageAbsMaxKernel(const T *in_state,
-                                              const T *in_accum,
-                                              const T *cur_scale,
-                                              const T rate,
-                                              T *out_state,
-                                              T *out_accum,
-                                              T *out_scale) {
-  T state = rate * (*in_state) + T(1.0f);
-  T accum = rate * (*in_accum) + (*cur_scale);
-  *out_state = state;
-  *out_accum = accum;
-  *out_scale = accum / state;
-}
-
 template struct FindRangeAbsMaxFunctor<phi::GPUContext, float>;
 
-template <typename T>
-struct FindMovingAverageAbsMaxFunctor<phi::GPUContext, T> {
-  void operator()(const phi::GPUContext &ctx,
-                  const phi::DenseTensor &in_accum,
-                  const phi::DenseTensor &in_state,
-                  const T *cur_scale,
-                  const float rate,
-                  phi::DenseTensor *out_state,
-                  phi::DenseTensor *out_accum,
-                  phi::DenseTensor *out_scale) {
-    const auto gpu_place = ctx.GetPlace();
-
-    T rate_t = static_cast<T>(rate);
-    T *out_state_data = out_state->mutable_data<T>(gpu_place);
-    T *out_accum_data = out_accum->mutable_data<T>(gpu_place);
-    T *out_scale_data = out_scale->mutable_data<T>(gpu_place);
-
-    FindMovingAverageAbsMaxKernel<T>
-        <<<1, 1, 0, ctx.stream()>>>(in_state.data<T>(),
-                                    in_accum.data<T>(),
-                                    cur_scale,
-                                    rate_t,
-                                    out_state_data,
-                                    out_accum_data,
-                                    out_scale_data);
-  }
-};
-
 // ChannelClipAndQuantDequantKernel for quant_axis is 0
 template <typename T>
 __global__ void ChannelClipAndQuantDequantKernelQuantAxis0(const T *in,
diff --git a/paddle/fluid/operators/fake_quantize_op.h b/paddle/fluid/operators/fake_quantize_op.h
@@ -78,18 +78,6 @@ struct ChannelClipFakeQuantDequantFunctor {
                   phi::DenseTensor *out);
 };
 
-template <typename DeviceContext, typename T>
-struct FindMovingAverageAbsMaxFunctor {
-  void operator()(const DeviceContext &ctx,
-                  const phi::DenseTensor &in_accum,
-                  const phi::DenseTensor &in_state,
-                  const T *cur_scale,
-                  const float rate,
-                  phi::DenseTensor *out_state,
-                  phi::DenseTensor *out_accum,
-                  phi::DenseTensor *out_scale);
-};
-
 template <typename DeviceContext, typename T>
 class FakeAbsMaxKernelBase : public framework::OpKernel<T> {
  public:
@@ -275,14 +263,15 @@ class FakeMovingAverageAbsMaxKernelBase : public framework::OpKernel<T> {
     out_scale->mutable_data<T>(context.GetPlace());
     float moving_rate = context.Attr<float>("moving_rate");
 
-    FindMovingAverageAbsMaxFunctor<DeviceContext, T>()(dev_ctx,
-                                                       *in_accum,
-                                                       *in_state,
-                                                       cur_scale_data,
-                                                       moving_rate,
-                                                       out_state,
-                                                       out_accum,
-                                                       out_scale);
+    phi::funcs::FindMovingAverageAbsMaxFunctor<DeviceContext, T>()(
+        dev_ctx,
+        *in_accum,
+        *in_state,
+        cur_scale_data,
+        moving_rate,
+        out_state,
+        out_accum,
+        out_scale);
 
     RunClipFunctor(dev_ctx, *in, *out_scale, bin_cnt, round_type, out);
   }
@@ -298,21 +287,6 @@ class FakeMovingAverageAbsMaxKernelBase : public framework::OpKernel<T> {
                               phi::DenseTensor *out) const = 0;
 };
 
-template <typename T, typename DeviceContext>
-class FakeQuantizeMovingAverageAbsMaxKernel
-    : public FakeMovingAverageAbsMaxKernelBase<T, DeviceContext> {
- protected:
-  void RunClipFunctor(const DeviceContext &dev_ctx,
-                      const phi::DenseTensor &in,
-                      const phi::DenseTensor &in_scale,
-                      int bin_cnt,
-                      int round_type,
-                      phi::DenseTensor *out) const override {
-    phi::funcs::ClipAndFakeQuantFunctor<DeviceContext, T>()(
-        dev_ctx, in, in_scale, bin_cnt, round_type, out);
-  }
-};
-
 template <typename T, typename DeviceContext>
 class FakeQuantizeDequantizeMovingAverageAbsMaxKernel
     : public FakeMovingAverageAbsMaxKernelBase<T, DeviceContext> {
@@ -365,14 +339,15 @@ class MovingAverageAbsMaxScaleKernel : public framework::OpKernel<T> {
     out_scale->mutable_data<T>(context.GetPlace());
     float moving_rate = context.Attr<float>("moving_rate");
 
-    FindMovingAverageAbsMaxFunctor<DeviceContext, T>()(dev_ctx,
-                                                       *in_accum,
-                                                       *in_state,
-                                                       cur_scale_data,
-                                                       moving_rate,
-                                                       out_state,
-                                                       out_accum,
-                                                       out_scale);
+    phi::funcs::FindMovingAverageAbsMaxFunctor<DeviceContext, T>()(
+        dev_ctx,
+        *in_accum,
+        *in_state,
+        cur_scale_data,
+        moving_rate,
+        out_state,
+        out_accum,
+        out_scale);
   }
 };
 
diff --git a/paddle/fluid/operators/quantize_linear_op.h b/paddle/fluid/operators/quantize_linear_op.h
@@ -84,14 +84,15 @@ class QuantizeLinearKernel : public framework::OpKernel<T> {
         out_scale->mutable_data<T>(context.GetPlace());
         float moving_rate = context.Attr<float>("moving_rate");
 
-        FindMovingAverageAbsMaxFunctor<DeviceContext, T>()(dev_ctx,
-                                                           *in_accum,
-                                                           *in_state,
-                                                           cur_scale_data,
-                                                           moving_rate,
-                                                           out_state,
-                                                           out_accum,
-                                                           out_scale);
+        phi::funcs::FindMovingAverageAbsMaxFunctor<DeviceContext, T>()(
+            dev_ctx,
+            *in_accum,
+            *in_state,
+            cur_scale_data,
+            moving_rate,
+            out_state,
+            out_accum,
+            out_scale);
         if (only_observer) {
           framework::TensorCopy(*in, context.GetPlace(), dev_ctx, out);
         } else {
diff --git a/paddle/phi/api/yaml/op_compat.yaml b/paddle/phi/api/yaml/op_compat.yaml
@@ -1114,8 +1114,16 @@
     attrs : [int round_type = 1]
 
 - op : fake_quantize_moving_average_abs_max
-  extra :
-    attrs : [int round_type = 1]
+  inputs :
+    x : X
+    in_scale : InScale
+    in_accum : InAccum
+    in_state : InState
+  outputs :
+    out : Out
+    out_scale : OutScale
+    out_state : OutState
+    out_accum : OutAccum
 
 - op : fake_quantize_range_abs_max
   extra :
diff --git a/paddle/phi/api/yaml/ops.yaml b/paddle/phi/api/yaml/ops.yaml
@@ -985,6 +985,16 @@
     func : fake_quantize_abs_max
     data_type : x
 
+- op : fake_quantize_moving_average_abs_max
+  args : (Tensor x, Tensor in_scale, Tensor in_accum, Tensor in_state, float moving_rate = 0.9, int bit_length = 8, bool is_test = false, int round_type = 1)
+  output : Tensor(out), Tensor(out_scale), Tensor(out_state), Tensor(out_accum)
+  infer_meta :
+    func : FakeQuantOrWithDequantMovingAverageAbsMaxInferMeta
+  kernel :
+    func : fake_quantize_moving_average_abs_max
+    data_type : x
+  optional : in_accum, in_state, out_state, out_accum
+
 - op : fft_c2c
   args : (Tensor x, int64_t[] axes, str normalization, bool forward)
   output : Tensor
diff --git a/paddle/phi/infermeta/multiary.cc b/paddle/phi/infermeta/multiary.cc
@@ -1577,6 +1577,36 @@ void EditDistanceInferMeta(const MetaTensor& hyps,
   sequencenum->set_dtype(DataType::FLOAT32);
 }
 
+void FakeQuantOrWithDequantMovingAverageAbsMaxInferMeta(
+    const MetaTensor& x,
+    const MetaTensor& in_scale,
+    const MetaTensor& in_accum,
+    const MetaTensor& in_state,
+    float moving_rate,
+    int bit_length,
+    bool is_test,
+    int round_type,
+    MetaTensor* out,
+    MetaTensor* out_scale,
+    MetaTensor* out_state,
+    MetaTensor* out_accum) {
+  PADDLE_ENFORCE_EQ(bit_length >= 1 && bit_length <= 16,
+                    true,
+                    phi::errors::InvalidArgument(
+                        "'bit_length' should be between 1 and 16, but "
+                        "the received is %d",
+                        bit_length));
+  if (out_state) {
+    out_state->set_dims({1});
+  }
+  if (out_accum) {
+    out_accum->set_dims({1});
+  }
+  out->set_dims(x.dims());
+  out_scale->set_dims({1});
+  out->share_lod(x);
+}
+
 void FtrlInferMeta(const MetaTensor& param,
                    const MetaTensor& squared_accumulator,
                    const MetaTensor& linear_accumulator,
diff --git a/paddle/phi/infermeta/multiary.h b/paddle/phi/infermeta/multiary.h
@@ -321,6 +321,20 @@ void EditDistanceInferMeta(const MetaTensor& hyps,
                            MetaTensor* sequencenum,
                            MetaTensor* out);
 
+void FakeQuantOrWithDequantMovingAverageAbsMaxInferMeta(
+    const MetaTensor& x,
+    const MetaTensor& in_scale,
+    const MetaTensor& in_accum,
+    const MetaTensor& in_state,
+    float moving_rate,
+    int bit_length,
+    bool is_test,
+    int round_type,
+    MetaTensor* out,
+    MetaTensor* out_scale,
+    MetaTensor* out_state,
+    MetaTensor* out_accum);
+
 void FtrlInferMeta(const MetaTensor& param,
                    const MetaTensor& squared_accumulator,
                    const MetaTensor& linear_accumulator,
diff --git a/paddle/phi/kernels/cpu/fake_quantize_kernel.cc b/paddle/phi/kernels/cpu/fake_quantize_kernel.cc
@@ -20,3 +20,9 @@ PD_REGISTER_KERNEL(fake_quantize_abs_max,
                    ALL_LAYOUT,
                    phi::FakeQuantizeAbsMaxKernel,
                    float) {}
+
+PD_REGISTER_KERNEL(fake_quantize_moving_average_abs_max,
+                   CPU,
+                   ALL_LAYOUT,
+                   phi::FakeQuantOrWithDequantMovingAverageAbsMaxKernel,
+                   float) {}
diff --git a/paddle/phi/kernels/fake_quantize_kernel.h b/paddle/phi/kernels/fake_quantize_kernel.h
@@ -26,4 +26,20 @@ void FakeQuantizeAbsMaxKernel(const Context& dev_ctx,
                               DenseTensor* out,
                               DenseTensor* out_scale);
 
+template <typename T, typename Context>
+void FakeQuantOrWithDequantMovingAverageAbsMaxKernel(
+    const Context& dev_ctx,
+    const DenseTensor& x,
+    const DenseTensor& in_scale,
+    const paddle::optional<DenseTensor>& in_accum,
+    const paddle::optional<DenseTensor>& in_state,
+    float moving_rate,
+    int bit_length,
+    bool is_test,
+    int round_type,
+    DenseTensor* out,
+    DenseTensor* out_scale,
+    DenseTensor* out_state,
+    DenseTensor* out_accum);
+
 }  // namespace phi
diff --git a/paddle/phi/kernels/funcs/fake_quantize_functor.cc b/paddle/phi/kernels/funcs/fake_quantize_functor.cc
@@ -52,8 +52,36 @@ void ClipAndFakeQuantFunctor<Context, T>::operator()(const Context &ctx,
   }
 }
 
+template <typename Context, typename T>
+void FindMovingAverageAbsMaxFunctor<Context, T>::operator()(
+    const Context &ctx,
+    const DenseTensor &in_accum,
+    const DenseTensor &in_state,
+    const T *cur_scale,
+    const float rate,
+    DenseTensor *out_state,
+    DenseTensor *out_accum,
+    DenseTensor *out_scale) {
+  T accum = in_accum.data<T>()[0];
+  T state = in_state.data<T>()[0];
+  T scale = cur_scale[0];
+
+  state = rate * state + 1;
+  accum = rate * accum + scale;
+  scale = accum / state;
+
+  T *out_state_data = ctx.template Alloc<T>(out_state);
+  T *out_accum_data = ctx.template Alloc<T>(out_accum);
+  T *out_scale_data = ctx.template Alloc<T>(out_scale);
+
+  out_state_data[0] = state;
+  out_accum_data[0] = accum;
+  out_scale_data[0] = scale;
+}
+
 template class FindAbsMaxFunctor<CPUContext, float>;
 template class ClipAndFakeQuantFunctor<CPUContext, float>;
+template class FindMovingAverageAbsMaxFunctor<CPUContext, float>;
 
 }  // namespace funcs
 }  // namespace phi
diff --git a/paddle/phi/kernels/funcs/fake_quantize_functor.cu b/paddle/phi/kernels/funcs/fake_quantize_functor.cu
diff --git a/paddle/phi/kernels/funcs/fake_quantize_functor.h b/paddle/phi/kernels/funcs/fake_quantize_functor.h
diff --git a/paddle/phi/kernels/gpu/fake_quantize_kernel.cu b/paddle/phi/kernels/gpu/fake_quantize_kernel.cu
diff --git a/paddle/phi/kernels/impl/fake_quantize_kernel_impl.h b/paddle/phi/kernels/impl/fake_quantize_kernel_impl.h
