[Inference] Process instance_norm/layer_norm/group_norm input/output data type specially (#63007)

yuanlehome · web-flow · commit f32ce8be9673 · 2024-03-26T19:21:31.000+08:00
* process instance_norm/layer_norm/group_norm input/output data type specially

* fix
diff --git a/paddle/fluid/framework/ir/auto_mixed_precision_pass.cc b/paddle/fluid/framework/ir/auto_mixed_precision_pass.cc
@@ -669,7 +669,8 @@ bool AutoMixedPrecisionPass::InputVarsNotConvert(
     if (std::find(vecs.begin(), vecs.end(), var_name) != vecs.end()) {
       return true;
     }
-  } else if (GetOpOriginalType(op_desc->Type()) == "instance_norm") {
+  } else if (GetOpOriginalType(op_desc->Type()) == "instance_norm" ||
+             GetOpOriginalType(op_desc->Type()) == "layer_norm") {
     auto vecs = op_desc->Input("Bias");
     if (std::find(vecs.begin(), vecs.end(), var_name) != vecs.end()) {
       return true;
@@ -705,37 +706,15 @@ bool AutoMixedPrecisionPass::InputVarsNotConvert(
     if (std::find(vecs.begin(), vecs.end(), var_name) != vecs.end()) {
       return true;
     }
-  }
-
-  if (backend_ == phi::Backend::XPU) {
-    if (GetOpOriginalType(op_desc->Type()) == "layer_norm") {
-      auto vecs = op_desc->Input("Bias");
-      if (std::find(vecs.begin(), vecs.end(), var_name) != vecs.end()) {
-        return true;
-      }
-      vecs = op_desc->Input("Scale");
-      if (std::find(vecs.begin(), vecs.end(), var_name) != vecs.end()) {
-        return true;
-      }
-    } else if (GetOpOriginalType(op_desc->Type()) == "instance_norm") {
-      auto vecs = op_desc->Input("Bias");
-      if (std::find(vecs.begin(), vecs.end(), var_name) != vecs.end()) {
-        return true;
-      }
-      vecs = op_desc->Input("Scale");
-      if (std::find(vecs.begin(), vecs.end(), var_name) != vecs.end()) {
-        return true;
-      }
-    } else if (GetOpOriginalType(op_desc->Type()) == "quantize_linear" ||
-               GetOpOriginalType(op_desc->Type()) == "dequantize_linear") {
-      auto vecs = op_desc->Input("Scale");
-      if (std::find(vecs.begin(), vecs.end(), var_name) != vecs.end()) {
-        return true;
-      }
-      vecs = op_desc->Input("ZeroPoint");
-      if (std::find(vecs.begin(), vecs.end(), var_name) != vecs.end()) {
-        return true;
-      }
+  } else if (GetOpOriginalType(op_desc->Type()) == "quantize_linear" ||
+             GetOpOriginalType(op_desc->Type()) == "dequantize_linear") {
+    auto vecs = op_desc->Input("Scale");
+    if (std::find(vecs.begin(), vecs.end(), var_name) != vecs.end()) {
+      return true;
+    }
+    vecs = op_desc->Input("ZeroPoint");
+    if (std::find(vecs.begin(), vecs.end(), var_name) != vecs.end()) {
+      return true;
     }
   }
 
@@ -784,18 +763,24 @@ bool AutoMixedPrecisionPass::OutputVarsNotConvert(
     if (std::find(vecs.begin(), vecs.end(), var_name) != vecs.end()) {
       return true;
     }
-  }
-
-  if (backend_ == phi::Backend::XPU) {
-    if (GetOpOriginalType(op_desc->Type()) == "layer_norm") {
-      auto vecs = op_desc->Output("Mean");
-      if (std::find(vecs.begin(), vecs.end(), var_name) != vecs.end()) {
-        return true;
-      }
-      vecs = op_desc->Output("Variance");
-      if (std::find(vecs.begin(), vecs.end(), var_name) != vecs.end()) {
-        return true;
-      }
+  } else if (GetOpOriginalType(op_desc->Type()) == "layer_norm" ||
+             GetOpOriginalType(op_desc->Type()) == "group_norm") {
+    auto vecs = op_desc->Output("Mean");
+    if (std::find(vecs.begin(), vecs.end(), var_name) != vecs.end()) {
+      return true;
+    }
+    vecs = op_desc->Output("Variance");
+    if (std::find(vecs.begin(), vecs.end(), var_name) != vecs.end()) {
+      return true;
+    }
+  } else if (GetOpOriginalType(op_desc->Type()) == "instance_norm") {
+    auto vecs = op_desc->Output("SavedMean");
+    if (std::find(vecs.begin(), vecs.end(), var_name) != vecs.end()) {
+      return true;
+    }
+    vecs = op_desc->Output("SavedVariance");
+    if (std::find(vecs.begin(), vecs.end(), var_name) != vecs.end()) {
+      return true;
     }
   }
 
