modify logic op stopgradient (#69716)

xiaoguoguo626807 · web-flow · commit 77bd1e9171f4 · 2024-11-27T10:36:23.000+08:00
diff --git a/paddle/fluid/pir/dialect/op_generator/op_build_gen.py b/paddle/fluid/pir/dialect/op_generator/op_build_gen.py
@@ -136,6 +136,15 @@
 
 _PREPARE_DATA_WITH_VECTOR_INT64_MTTABLE_ATTRIBUTE = {'FrobeniusNormOp'}
 
+LOGIC_OP_LIST = {
+    'LogicalAndOp',
+    'LogicalAnd_Op',
+    'LogicalOrOp',
+    'LogicalOr_Op',
+    'LogicalNotOp',
+    'LogicalNot_Op',
+    'LogicalXorOp',
+}
 OP_BUILD_TEMPLATE = """
 void {op_name}::Build({build_args}) {{
 {build_info}
@@ -841,8 +850,11 @@ def gen_build_func_str(
     build_outputs_str = f"""
   std::vector<pir::Type> argument_outputs = {op_info.class_name}::InferMeta(argument_inputs, &argument_attributes);
   argument.AddAttributes(argument_attributes);
-  argument.AddOutputs(argument_outputs.begin(), argument_outputs.end());
-  ::pir::PassStopGradientsDefaultly(argument);"""
+  argument.AddOutputs(argument_outputs.begin(), argument_outputs.end());\n"""
+    if op_info.class_name in LOGIC_OP_LIST:
+        build_outputs_str += "::pir::TrueStopGradientsDefaultly(argument);\n"
+    else:
+        build_outputs_str += "::pir::PassStopGradientsDefaultly(argument);" ""
 
     GET_ATTRIBUTES_FROM_MAP_TEMPLATE = """
   PADDLE_ENFORCE_NE(
diff --git a/paddle/pir/include/core/builtin_op.h b/paddle/pir/include/core/builtin_op.h
@@ -258,6 +258,7 @@ class IR_API ConstantTensorOp : public ConstantOp {
 };
 
 void PassStopGradientsDefaultly(OperationArgument &argument);  // NOLINT
+void TrueStopGradientsDefaultly(OperationArgument &argument);  // NOLINT
 void RefreshStopGradientsDefaultly(Operation *Op);
 }  // namespace pir
 
diff --git a/paddle/pir/src/core/builtin_op.cc b/paddle/pir/src/core/builtin_op.cc
@@ -63,6 +63,17 @@ void PassStopGradientsDefaultly(OperationArgument &argument) {  // NOLINT
       pir::ArrayAttribute::get(pir::IrContext::Instance(), outs_stop_gradient));
 }
 
+void TrueStopGradientsDefaultly(OperationArgument &argument) {  // NOLINT
+  VLOG(10) << "Builder construction stop gradient as True for OpResults.";
+  bool stop_gradient = true;
+  std::vector<pir::Attribute> outs_stop_gradient(
+      argument.output_types.size(),
+      pir::BoolAttribute::get(pir::IrContext::Instance(), stop_gradient));
+  argument.AddAttribute(
+      kStopGradientAttrName,
+      pir::ArrayAttribute::get(pir::IrContext::Instance(), outs_stop_gradient));
+}
+
 void RefreshStopGradientsDefaultly(Operation *op) {
   bool stop_gradient = true;
   for (auto value : op->operands_source()) {
