【pir cinn】modify some bug of Dpa model (#70123)

xiaoguoguo626807 · web-flow · commit 0aa87d154873 · 2024-12-11T19:23:49.000+08:00
* modify DPA model

* modify ir_printer
diff --git a/paddle/cinn/hlir/dialect/operator/transforms/pd_to_cinn_pass.cc b/paddle/cinn/hlir/dialect/operator/transforms/pd_to_cinn_pass.cc
@@ -23,6 +23,7 @@
 #include "paddle/fluid/pir/dialect/operator/ir/pd_op.h"
 #include "paddle/fluid/pir/dialect/operator/utils/utils.h"
 #include "paddle/fluid/pir/drr/include/drr_pattern_base.h"
+#include "paddle/fluid/pir/utils/general_functions.h"
 #include "paddle/pir/include/core/builtin_dialect.h"
 #include "paddle/pir/include/core/builtin_op.h"
 #include "paddle/pir/include/pass/pass.h"
@@ -522,11 +523,11 @@ class PowOpPattern : public pir::OpRewritePattern<paddle::dialect::PowOp> {
   void Rewrite(paddle::dialect::PowOp op,
                pir::PatternRewriter &rewriter) const override {
     auto factor = op->attribute("y").dyn_cast<pir::FloatAttribute>().data();
-    auto full_op =
-        rewriter.Build<paddle::dialect::FullOp>(std::vector<int64_t>({1}),
-                                                factor,
-                                                phi::DataType::FLOAT32,
-                                                phi::CPUPlace());
+    auto full_op = rewriter.Build<paddle::dialect::FullOp>(
+        std::vector<int64_t>({1}),
+        factor,
+        pir::GetValueDtype(op->operand_source(0)),
+        phi::CPUPlace());
 
     auto elementwise_pow = rewriter.Build<paddle::dialect::ElementwisePowOp>(
         op->operand_source(0), full_op->result(0));
diff --git a/paddle/fluid/pir/utils/general_functions.cc b/paddle/fluid/pir/utils/general_functions.cc
@@ -279,4 +279,35 @@ bool ValueIsPersistable(const pir::Value& value) {
   return true;
 }
 
+phi::DataType GetTensorDtype(pir::Type type) {
+  if (!type) {
+    PADDLE_THROW(
+        common::errors::InvalidArgument("The type of value is nullptr."));
+  }
+  if (auto dense_tensor_type = type.dyn_cast<pir::DenseTensorType>()) {
+    return paddle::dialect::TransToPhiDataType(dense_tensor_type.dtype());
+  } else if (auto sparse_coo_tensor_type =
+                 type.dyn_cast<paddle::dialect::SparseCooTensorType>()) {
+    return paddle::dialect::TransToPhiDataType(sparse_coo_tensor_type.dtype());
+  } else if (auto sparse_csr_tensor_type =
+                 type.dyn_cast<paddle::dialect::SparseCsrTensorType>()) {
+    return paddle::dialect::TransToPhiDataType(sparse_csr_tensor_type.dtype());
+  } else if (auto select_rows =
+                 type.dyn_cast<paddle::dialect::SelectedRowsType>()) {
+    return paddle::dialect::TransToPhiDataType(select_rows.dtype());
+  } else if (auto dense_array =
+                 type.dyn_cast<paddle::dialect::DenseTensorArrayType>()) {
+    return paddle::dialect::TransToPhiDataType(dense_array.dtype());
+  } else {
+    PADDLE_THROW(common::errors::InvalidArgument(
+        "Currently, we can only get phi::DataType from DenseTensorType and "
+        "SelectedRowsType, DenseTensorArrayType,SparseCooTensorType or "
+        "SparseCsrTensorType."));
+  }
+}
+
+phi::DataType GetValueDtype(const pir::Value& val) {
+  return GetTensorDtype(val.type());
+}
+
 }  // namespace pir
diff --git a/paddle/fluid/pir/utils/general_functions.h b/paddle/fluid/pir/utils/general_functions.h
@@ -20,6 +20,7 @@
 #include "paddle/phi/common/data_type.h"
 #include "paddle/phi/common/place.h"
 #include "paddle/pir/include/core/type.h"
+#include "paddle/pir/include/core/value.h"
 #include "paddle/pir/include/pass/pass.h"
 
 namespace paddle {
@@ -247,4 +248,7 @@ std::vector<Value> GetUsedExternalValue(const Block& block);
  */
 bool ValueIsPersistable(const pir::Value& value);
 
+phi::DataType GetTensorDtype(pir::Type type);
+phi::DataType GetValueDtype(const pir::Value& val);
+
 }  // namespace pir
diff --git a/paddle/fluid/pybind/pir.cc b/paddle/fluid/pybind/pir.cc
@@ -13,6 +13,7 @@
 // limitations under the License.
 
 #include "paddle/fluid/pybind/pir.h"
+#include "paddle/fluid/pybind/pir_utils.h"
 
 #include <Python.h>
 #include <algorithm>
@@ -210,35 +211,6 @@ std::string GetValueInfo(Value v) {
   return ss.str();
 }
 
-phi::DataType GetTensorDtype(Type type) {
-  if (!type) {
-    PADDLE_THROW(
-        common::errors::InvalidArgument("The type of value is nullptr."));
-  }
-  if (auto dense_tensor_type = type.dyn_cast<DenseTensorType>()) {
-    return dialect::TransToPhiDataType(dense_tensor_type.dtype());
-  } else if (auto sparse_coo_tensor_type =
-                 type.dyn_cast<SparseCooTensorType>()) {
-    return dialect::TransToPhiDataType(sparse_coo_tensor_type.dtype());
-  } else if (auto sparse_csr_tensor_type =
-                 type.dyn_cast<SparseCsrTensorType>()) {
-    return dialect::TransToPhiDataType(sparse_csr_tensor_type.dtype());
-  } else if (auto select_rows = type.dyn_cast<SelectedRowsType>()) {
-    return dialect::TransToPhiDataType(select_rows.dtype());
-  } else if (auto dense_array = type.dyn_cast<DenseTensorArrayType>()) {
-    return dialect::TransToPhiDataType(dense_array.dtype());
-  } else {
-    PADDLE_THROW(common::errors::InvalidArgument(
-        "Currently, we can only get phi::DataType from DenseTensorType and "
-        "SelectedRowsType, DenseTensorArrayType,SparseCooTensorType or "
-        "SparseCsrTensorType."));
-  }
-}
-
-phi::DataType GetValueDtype(Value value) {
-  return GetTensorDtype(value.type());
-}
-
 py::object Clone(const Program &self, IrMapping *p_mapper = nullptr) {
   IrMapping mapper;
   if (p_mapper == nullptr) {
@@ -274,7 +246,7 @@ pir::Value AppendDataOp(pir::Block *block,
        paddle::dialect::IntArrayAttribute::get(
            ctx, phi::IntArray(phi::vectorize(GetValueDims(value))))},
       {"dtype",
-       paddle::dialect::DataTypeAttribute::get(ctx, GetValueDtype(value))},
+       paddle::dialect::DataTypeAttribute::get(ctx, pir::GetValueDtype(value))},
       {"place", PlaceAttribute::get(ctx, phi::Place())}};
   std::vector<pir::Type> output_types{value.type()};
   pir::Operation *operation =
@@ -1427,7 +1399,7 @@ void BindValue(py::module *m) {
           })
       .def_property(
           "dtype",
-          [](Value self) { return GetValueDtype(self); },
+          [](Value self) { return pir::GetValueDtype(self); },
           [](Value self, phi::DataType dtype) {
             PADDLE_THROW(common::errors::InvalidArgument(
                 "can't set dtype when building static graph"));
@@ -2299,7 +2271,7 @@ static void inline CreateVariableIfNotExist(
       phi::DeviceContextPool &pool = phi::DeviceContextPool::Instance();
       const phi::DeviceContext *dev_ctx = nullptr;
       dev_ctx = pool.Get(exe->GetPlace());
-      dev_ctx->Alloc(tensor_temp, GetValueDtype(value));
+      dev_ctx->Alloc(tensor_temp, pir::GetValueDtype(value));
     }
   }
   return;
diff --git a/paddle/pir/src/core/ir_printer.cc b/paddle/pir/src/core/ir_printer.cc
@@ -389,10 +389,9 @@ void IrPrinter::PrintOpReturnType(const Operation& op) {
 
 void IrPrinter::AddValueAlias(Value v, const std::string& alias) {
   const void* key = v.impl();
-  PADDLE_ENFORCE_EQ(aliases_.find(key),
-                    aliases_.end(),
-                    common::errors::InvalidArgument("Value already has alias"));
-  aliases_[key] = alias;
+  if (aliases_.find(key) == aliases_.end()) {
+    aliases_[key] = alias;
+  }
 }
 
 class CustomPrinter : public IrPrinter {
diff --git a/python/paddle/autograd/backward_utils.py b/python/paddle/autograd/backward_utils.py
@@ -330,7 +330,11 @@ def copy(self, new_block):
 def _check_vjp_dynamic_shape(op, inputs):
     for items in inputs:
         for item in items:
-            if item.initialized() and -1 in item.shape:
+            if (
+                item.is_dense_tensor_type()
+                and item.initialized()
+                and -1 in item.shape
+            ):
                 return True
 
 
diff --git a/python/paddle/autograd/ir_backward.py b/python/paddle/autograd/ir_backward.py
@@ -696,24 +696,29 @@ def append_yield(
 
                 if op.name() == "cf.tuple_push":
                     stackop = op.operand_source(0).get_defining_op()
-                    with dynamic_shape_prim_vjp_guard(op, inputs):
-                        copy_out = paddle.framework.core.call_vjp(
-                            op,
-                            inputs,
-                            outputs,
-                            output_grads,
-                            input_grad_stopgradients,
-                        )
+                    if stackop.result(2).use_empty():
+                        with dynamic_shape_prim_vjp_guard(op, inputs):
+                            copy_out = paddle.framework.core.call_vjp(
+                                op,
+                                inputs,
+                                outputs,
+                                output_grads,
+                                input_grad_stopgradients,
+                            )
 
-                    pop_op = bwd_block.ops[-1]
-                    while_tuple_ops.append(pop_op)
-                    while_tuple_ops.append(op)
-                    while_tuple_ops.append(stackop)
-                    bwd_ops = [pop_op]
-                    for output, copy_output in zip(inputs[1:], copy_out[1:]):
-                        control_flow_value_to_copyvalue_map[output[0]] = (
-                            copy_output[0]
-                        )
+                        pop_op = bwd_block.ops[-1]
+                        while_tuple_ops.append(pop_op)
+                        while_tuple_ops.append(op)
+                        while_tuple_ops.append(stackop)
+                        bwd_ops = [pop_op]
+                        for output, copy_output in zip(
+                            inputs[1:], copy_out[1:]
+                        ):
+                            control_flow_value_to_copyvalue_map[output[0]] = (
+                                copy_output[0]
+                            )
+                    else:
+                        bwd_ops = [stackop.result(2).first_use().owner()]
                 else:
                     # all(zero_flag) support this op has no contribution for grad
                     # should be delete (prune sub_graph)
