[PIR][oneDNN] Fix bf16 placement for Reshape (#68716)

LLee233 · web-flow · commit cea0e0e208b9 · 2024-10-16T16:49:38.000+08:00
* Generalize bf16 placement condition

* add input name for Fused_Matmul
diff --git a/paddle/fluid/pir/transforms/onednn/cpu_bfloat16_placement_pass.cc b/paddle/fluid/pir/transforms/onednn/cpu_bfloat16_placement_pass.cc
@@ -97,9 +97,29 @@ class OneDNNBf16PlacementPattern : public pir::RewritePattern {
       }
     }
 
-    int i = 0;
-    for (auto& value : op->operands_source()) {
-      pir::Type type = op->operand_type(i++);
+    const std::vector<std::string> permitted_input_names = {
+        "x", "y", "input", "residual_param", "residual_data"};
+    auto op_name = op->name();
+    auto op_info = pir::IrContext::Instance()->GetRegisteredOpInfo(op_name);
+    if (!op_info) return false;
+    paddle::dialect::OpYamlInfoParser yaml_parser(
+        op_info.GetInterfaceImpl<paddle::dialect::OpYamlInfoInterface>()
+            ->get_op_info_(op_name),
+        paddle::dialect::IsLegacyOp(op_name));
+    auto input_names = yaml_parser.InputNames();
+
+    for (size_t i = 0; i < op->num_operands(); i++) {
+      pir::Value value = op->operand_source(i);
+      if (!value) continue;
+      std::string input_name = input_names[i];
+      auto iter = std::find(permitted_input_names.begin(),
+                            permitted_input_names.end(),
+                            input_name);
+      if (iter == permitted_input_names.end()) {
+        continue;
+      }
+      pir::Type type = op->operand_type(i);
+      if (!type) continue;
       if (!type.isa<paddle::dialect::DenseTensorType>()) {
         // We skip pir::VectorType
         // TODO(Lirong, Xinyi): Support pir::VectorType in bf16
@@ -364,24 +384,6 @@ class RemoveUnsupportedOpPattern : public pir::RewritePattern {
       }
     }
 
-    bool unsupported_op = false;
-    int i = 0;
-    for (auto& value : op->operands_source()) {
-      pir::Type type = op->operand_type(i++);
-      if (!type.isa<paddle::dialect::DenseTensorType>()) {
-        return false;
-      }
-      pir::Type op_dtype = pir::GetDataTypeFromValue(value);
-      // Only float input can be converted to bfloat16
-      if (!op_dtype.isa<pir::Float32Type>()) {
-        unsupported_op = true;
-        break;
-      }
-    }
-
-    if (!unsupported_op) {
-      return false;
-    }
     return true;
   }
 
diff --git a/test/ir/pir/fused_pass/onednn/test_cpu_bfloat16_pir_pass.py b/test/ir/pir/fused_pass/onednn/test_cpu_bfloat16_pir_pass.py
@@ -456,13 +456,8 @@ def build_ir_program(self):
                 x = paddle.static.data(
                     name='x', shape=[5, 5, 5, 5], dtype='float32'
                 )
-                y = paddle.static.data(
-                    name='y', shape=[5, 5, 5, 5], dtype='float32'
-                )
-                reshape_x = paddle.reshape(x, [5, 125])
-                reshape_y = paddle.reshape(y, [125, 5])
 
-                out = paddle.matmul(reshape_x, reshape_y)
+                out = paddle.reshape(x, [5, 125])
                 out = paddle.assign(out)
                 self.pass_attr_list = [
                     {'onednn_placement_pass': {}},
@@ -472,12 +467,12 @@ def build_ir_program(self):
                 ]
                 self.feeds = {
                     "x": np.random.random((5, 5, 5, 5)).astype("float32"),
-                    "y": np.random.random((5, 5, 5, 5)).astype("float32"),
                 }
                 self.fetch_list = [out]
                 self.valid_op_map = {
-                    "onednn_op.matmul": 1,
-                    "pd_op.add": 0,
+                    "onednn_op.reshape": 1,
+                    "onednn_op.quantize": 1,
+                    "onednn_op.dequantize": 1,
                 }
                 return [main_prog, start_prog]
 
