Modify the DataType directly when the Operand of whileOp is FakeFull. (#69939)

zhanghonggeng · web-flow · commit fe607df42d8d · 2024-12-10T14:36:11.000+08:00
diff --git a/paddle/fluid/pir/transforms/general/auto_mixed_precision_pass.cc b/paddle/fluid/pir/transforms/general/auto_mixed_precision_pass.cc
@@ -705,6 +705,10 @@ class AutoMixedPrecisionPass : public pir::Pass {
           auto operand_phi_dtype = GetPhiDataTypeFromOpOperand(operand);
           if (IsPhiDataTypeFloat(operand_phi_dtype) &&
               operand_phi_dtype == precision_mode_) {
+            if (IsFakeFullOperandWhileOp(op, operand)) {
+              SetFakeFullDataType(operand, builder);
+              continue;
+            }
             DoInsertCastOp(op, operand, phi_dtype, builder);
           }
         } else if (IsOperandHasDenseTensorVectorType(operand)) {
@@ -735,6 +739,56 @@ class AutoMixedPrecisionPass : public pir::Pass {
     }
   }
 
+  std::vector<int64_t> ConvertDDimToVector(const common::DDim& ddim) {
+    std::vector<int64_t> dims;
+    for (int i = 0; i < ddim.size(); ++i) {
+      dims.push_back(ddim[i]);
+    }
+    return dims;
+  }
+
+  bool IsFakeFullOp(pir::OpOperand operand) {
+    auto defining_op_ = operand.source().defining_op();
+    if (defining_op_->isa<paddle::dialect::FullOp>()) {
+      auto full_op = defining_op_->dyn_cast<paddle::dialect::FullOp>();
+      auto shape_attr = full_op.attribute("shape")
+                            .dyn_cast<paddle::dialect::IntArrayAttribute>();
+      auto shape_dims = shape_attr.data().GetData();
+      auto result_dims = full_op.out()
+                             .type()
+                             .dyn_cast<paddle::dialect::DenseTensorType>()
+                             .dims();
+      std::vector<int64_t> result_dims_vec = ConvertDDimToVector(result_dims);
+
+      if (shape_dims.size() != result_dims_vec.size()) {
+        return true;
+      }
+      for (size_t i = 0; i < shape_dims.size(); ++i) {
+        if (shape_dims[i] != result_dims_vec[i]) {
+          return true;
+        }
+      }
+      return false;
+    }
+    return false;
+  }
+
+  bool IsFakeFullOperandWhileOp(pir::Operation* op, pir::OpOperand operand) {
+    return op->isa<paddle::dialect::WhileOp>() && IsFakeFullOp(operand);
+  }
+
+  void SetFakeFullDataType(pir::OpOperand operand,
+                           pir::Builder& builder) {  // NOLINT
+    auto defining_op_ = operand.source().defining_op();
+    auto full_op = defining_op_->dyn_cast<paddle::dialect::FullOp>();
+    pir::Attribute attr_dtype = paddle::dialect::DataTypeAttribute::get(
+        builder.ir_context(), phi::DataType::FLOAT32);
+    full_op->set_attribute("dtype", attr_dtype);
+
+    SetResultDataType(
+        full_op.out(), phi::DataType::FLOAT32, builder.ir_context());
+  }
+
   void SubOpRun(pir::Operation* op) {
     for (auto& region : *op) {
       for (auto& block : region) {
