[PIR] Support act param in scale and fix BatchNorm(act="hard_swish") (#68036)

SigureMo · web-flow · commit 8f58190d297f · 2024-09-06T10:35:39.000+08:00
diff --git a/python/paddle/base/layer_helper.py b/python/paddle/base/layer_helper.py
@@ -17,7 +17,6 @@
 from typing import TYPE_CHECKING, Any, Generator
 
 import paddle
-from paddle import _C_ops
 
 from . import unique_name
 from .dygraph_utils import _append_activation_in_dygraph
@@ -176,20 +175,9 @@ def append_activation(self, input_var: Tensor) -> Tensor:
             res = _append_activation_in_dygraph(input_var, act_type, use_cudnn)
             return res
         elif in_pir_mode():
-
-            def _append_activation_in_pir(input, act=None, use_cudnn=None):
-                if act is None:
-                    return input
-
-                attrs = ()
-                if use_cudnn:
-                    attrs = ('use_cudnn', use_cudnn)
-                act_op = getattr(_C_ops, act)
-                if act == 'softmax':
-                    return act_op(input, -1)
-                return act_op(input, *attrs)
-
-            return _append_activation_in_pir(input_var, act_type, use_cudnn)
+            return paddle.pir_utils.append_activation_in_pir(
+                input_var, act_type, use_cudnn
+            )
         else:
             tmp = self.create_variable_for_type_inference(dtype=input_var.dtype)
             self.append_op(
diff --git a/python/paddle/nn/layer/norm.py b/python/paddle/nn/layer/norm.py
@@ -33,7 +33,7 @@
 
 import numpy as np
 
-from paddle import _C_ops, in_dynamic_mode
+from paddle import _C_ops, in_dynamic_mode, pir_utils
 from paddle.device import get_all_custom_device_type
 
 from ...base import dygraph_utils
@@ -1081,11 +1081,7 @@ def forward(self, input: Tensor) -> Tensor:
                 self._use_global_stats,
                 self._trainable_statistics,
             )
-            if self._act is None:
-                return batch_norm_out
-
-            act_op = getattr(_C_ops, self._act)
-            return act_op(batch_norm_out)
+            return pir_utils.append_activation_in_pir(batch_norm_out, self._act)
         else:
             # create output
             # mean and mean_out share the same memory
diff --git a/python/paddle/pir_utils.py b/python/paddle/pir_utils.py
@@ -254,3 +254,21 @@ def analysis_io(program: paddle.pir.Program):
             total_io += get_memory(value)
 
     return total_io / 1024 / 1024 / 1024
+
+
+def append_activation_in_pir(input, act=None, use_cudnn=None):
+    if act is None:
+        return input
+
+    act_name_mapping = {
+        "hard_swish": "hardswish",
+    }
+    act = act_name_mapping.get(act, act)
+
+    attrs = ()
+    if use_cudnn:
+        attrs = ('use_cudnn', use_cudnn)
+    act_op = getattr(paddle._C_ops, act)
+    if act == 'softmax':
+        return act_op(input, -1)
+    return act_op(input, *attrs)
diff --git a/python/paddle/tensor/math.py b/python/paddle/tensor/math.py
@@ -294,9 +294,8 @@ def scale(
         out = _C_ops.scale(x, scale, float(bias), bias_after_scale)
         return dygraph_utils._append_activation_in_dygraph(out, act)
     elif in_pir_mode():
-        if act is None:
-            return _C_ops.scale(x, scale, float(bias), bias_after_scale)
-        raise ValueError("act is not implement in pir of scale api.")
+        out = _C_ops.scale(x, scale, float(bias), bias_after_scale)
+        return paddle.pir_utils.append_activation_in_pir(out, act)
     else:
         check_variable_and_dtype(
             x,
