[PIR][OneDNN] Add cpu bf16 placement pass check data type (#68021)

zhanglirong1999 · web-flow · commit 1d268a36a288 · 2024-09-14T21:36:13.000+08:00
* [PIR][OneDNN] Fix cpu bf16 placement pass

* skip 2 pattern if mkldnn not replace
diff --git a/paddle/fluid/pir/transforms/onednn/cpu_bfloat16_placement_pass.cc b/paddle/fluid/pir/transforms/onednn/cpu_bfloat16_placement_pass.cc
@@ -96,6 +96,21 @@ class OneDNNBf16PlacementPattern : public pir::RewritePattern {
         return false;
       }
     }
+
+    int i = 0;
+    for (auto& value : op->operands_source()) {
+      pir::Type type = op->operand_type(i++);
+      if (!type.isa<paddle::dialect::DenseTensorType>()) {
+        // We skip pir::VectorType
+        // TODO(Lirong, Xinyi): Support pir::VectorType in bf16
+        return false;
+      }
+      pir::Type op_dtype = pir::GetDataTypeFromValue(value);
+      // Only float input can be converted to bfloat16
+      if (!op_dtype.isa<pir::Float32Type>()) {
+        return false;
+      }
+    }
     return true;
   }
 
@@ -176,6 +191,15 @@ class RemoveOrphanedPattern : public pir::RewritePattern {
         !op->isa<paddle::onednn::dialect::FusedMatmulOp>()) {
       return false;
     }
+    auto op_attr = op->attributes();
+    if (op_attr.find("mkldnn_data_type") != op_attr.end()) {
+      auto mkldnn_data_type = op_attr.at("mkldnn_data_type")
+                                  .dyn_cast<pir::StrAttribute>()
+                                  .AsString();
+      if (mkldnn_data_type != "bfloat16") {
+        return false;
+      }
+    }
 
     bool prev_fp32 = false;
     bool next_fp32 = false;
@@ -322,6 +346,15 @@ class RemoveUnsupportedOpPattern : public pir::RewritePattern {
         !op->isa<paddle::onednn::dialect::FusedMatmulOp>()) {
       return false;
     }
+    auto op_attr = op->attributes();
+    if (op_attr.find("mkldnn_data_type") != op_attr.end()) {
+      auto mkldnn_data_type = op_attr.at("mkldnn_data_type")
+                                  .dyn_cast<pir::StrAttribute>()
+                                  .AsString();
+      if (mkldnn_data_type != "bfloat16") {
+        return false;
+      }
+    }
 
     uint32_t num_operands = op->num_operands();
     for (uint32_t i = 0; i < num_operands; i++) {
@@ -332,14 +365,20 @@ 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;
     }
