PaddlePaddle
diff --git a/‎paddle/fluid/pir/dialect/op_generator/op_build_gen.py
Copy file name to clipboard
+1-2Lines changed: 1 addition & 2 deletions b/‎paddle/fluid/pir/dialect/op_generator/op_build_gen.py
Copy file name to clipboard
+1-2Lines changed: 1 addition & 2 deletions
diff --git a/‎paddle/fluid/pir/dialect/op_generator/op_gen.py
Copy file name to clipboard
+1-12Lines changed: 1 addition & 12 deletions b/‎paddle/fluid/pir/dialect/op_generator/op_gen.py
Copy file name to clipboard
+1-12Lines changed: 1 addition & 12 deletions
diff --git a/‎paddle/fluid/pir/dialect/operator/ir/manual_op.h
Copy file name to clipboard
+25-12Lines changed: 25 additions & 12 deletions b/‎paddle/fluid/pir/dialect/operator/ir/manual_op.h
Copy file name to clipboard
+25-12Lines changed: 25 additions & 12 deletions
diff --git a/‎paddle/fluid/pir/dialect/operator/trait/forward_only.h
Copy file name to clipboard
+30Lines changed: 30 additions & 0 deletions b/‎paddle/fluid/pir/dialect/operator/trait/forward_only.h
Copy file name to clipboard
+30Lines changed: 30 additions & 0 deletions
diff --git a/‎paddle/fluid/pir/dialect/operator/trait/trait.cc
Copy file name to clipboard
+2Lines changed: 2 additions & 0 deletions b/‎paddle/fluid/pir/dialect/operator/trait/trait.cc
Copy file name to clipboard
+2Lines changed: 2 additions & 0 deletions
diff --git a/‎paddle/fluid/pybind/pybind.cc
Copy file name to clipboard
+15Lines changed: 15 additions & 0 deletions b/‎paddle/fluid/pybind/pybind.cc
Copy file name to clipboard
+15Lines changed: 15 additions & 0 deletions
diff --git a/‎paddle/phi/ops/yaml/inconsistent/dygraph_ops.yaml
Copy file name to clipboard
+9Lines changed: 9 additions & 0 deletions b/‎paddle/phi/ops/yaml/inconsistent/dygraph_ops.yaml
Copy file name to clipboard
+9Lines changed: 9 additions & 0 deletions
diff --git a/‎paddle/phi/ops/yaml/inconsistent/onednn_ops_extra.yaml
Copy file name to clipboard
+3Lines changed: 3 additions & 0 deletions b/‎paddle/phi/ops/yaml/inconsistent/onednn_ops_extra.yaml
Copy file name to clipboard
+3Lines changed: 3 additions & 0 deletions
@@ -819,12 +819,11 @@ def gen_build_func_str(
         op_non_mutable_attribute_name_list,
         op_non_mutable_attribute_type_list,
     )
-    str_is_forward_only = "true" if op_info.is_forward_only else "false"
     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, {str_is_forward_only});"""
+  ::pir::PassStopGradientsDefaultly(argument);"""
 
     GET_ATTRIBUTES_FROM_MAP_TEMPLATE = """
   PADDLE_ENFORCE_NE(
 
@@ -97,6 +97,7 @@
 #include "paddle/fluid/pir/dialect/operator/interface/parse_kernel_key.h"
 #include "paddle/fluid/pir/dialect/operator/interface/vjp.h"
 #include "paddle/fluid/pir/dialect/operator/trait/inplace.h"
+#include "paddle/fluid/pir/dialect/operator/trait/forward_only.h"
 #include "paddle/fluid/pir/dialect/operator/utils/op_yaml_info_util.h"
 #include "paddle/fluid/pir/dialect/operator/utils/utils.h"
 #include "paddle/pir/include/core/builder.h"
@@ -543,18 +544,6 @@ def __init__(self, op_yaml_item, op_compat_item, yaml_file):
         else:
             self.dynamic_fallback = False
 
-        self.is_forward_only = (
-            False
-            if (
-                (
-                    'backward' in self.op_yaml_item
-                    and self.op_yaml_item['backward'] is not None
-                )
-                or 'forward' in self.op_yaml_item
-            )
-            else True
-        )
-
     def parse_op_traits(self):
         if 'traits' in self.op_yaml_item:
             return self.op_yaml_item['traits']
 
@@ -22,6 +22,7 @@
 #include "paddle/fluid/pir/dialect/operator/interface/infermeta.h"
 #include "paddle/fluid/pir/dialect/operator/interface/op_yaml_info.h"
 #include "paddle/fluid/pir/dialect/operator/interface/vjp.h"
+#include "paddle/fluid/pir/dialect/operator/trait/forward_only.h"
 #include "paddle/fluid/pir/dialect/operator/trait/inplace.h"
 #include "paddle/fluid/pir/dialect/operator/utils/op_yaml_info_util.h"
 #include "paddle/fluid/pir/dialect/operator/utils/utils.h"
@@ -202,8 +203,10 @@ class SplitGradOp : public pir::Op<SplitGradOp, OpYamlInfoInterface> {
       pir::AttributeMap *p_attributes);
 };
 
-class CreateArrayOp
-    : public pir::Op<CreateArrayOp, OpYamlInfoInterface, InferMetaInterface> {
+class CreateArrayOp : public pir::Op<CreateArrayOp,
+                                     OpYamlInfoInterface,
+                                     InferMetaInterface,
+                                     paddle::dialect::ForwardOnlyTrait> {
  public:
   using Op::Op;
   static const char *name() { return "pd_op.create_array"; }
@@ -223,7 +226,8 @@ class CreateArrayOp
 
 class CreateArrayLikeOp : public pir::Op<CreateArrayLikeOp,
                                          OpYamlInfoInterface,
-                                         InferMetaInterface> {
+                                         InferMetaInterface,
+                                         paddle::dialect::ForwardOnlyTrait> {
  public:
   using Op::Op;
   static const char *name() { return "pd_op.create_array_like"; }
@@ -243,8 +247,10 @@ class CreateArrayLikeOp : public pir::Op<CreateArrayLikeOp,
       pir::AttributeMap *p_attributes);
 };
 
-class ArrayLengthOp
-    : public pir::Op<ArrayLengthOp, OpYamlInfoInterface, InferMetaInterface> {
+class ArrayLengthOp : public pir::Op<ArrayLengthOp,
+                                     OpYamlInfoInterface,
+                                     InferMetaInterface,
+                                     paddle::dialect::ForwardOnlyTrait> {
  public:
   using Op::Op;
   static const char *name() { return "pd_op.array_length"; }
@@ -266,7 +272,8 @@ class ArrayLengthOp
 class ArrayReadOp : public pir::Op<ArrayReadOp,
                                    OpYamlInfoInterface,
                                    paddle::dialect::VjpInterface,
-                                   InferMetaInterface> {
+                                   InferMetaInterface,
+                                   paddle::dialect::ForwardOnlyTrait> {
  public:
   using Op::Op;
   static const char *name() { return "pd_op.array_read"; }
@@ -301,7 +308,8 @@ class ArrayWrite_Op : public pir::Op<ArrayWrite_Op,
                                      OpYamlInfoInterface,
                                      paddle::dialect::VjpInterface,
                                      InferMetaInterface,
-                                     InplaceTrait> {
+                                     InplaceTrait,
+                                     paddle::dialect::ForwardOnlyTrait> {
  public:
   using Op::Op;
   static const char *name() { return "pd_op.array_write_"; }
@@ -389,7 +397,8 @@ class SliceArrayOp
     : public pir::Op<SliceArrayOp,
                      paddle::dialect::OpYamlInfoInterface,
                      paddle::dialect::InferMetaInterface,
-                     paddle::dialect::GetKernelTypeForVarInterface> {
+                     paddle::dialect::GetKernelTypeForVarInterface,
+                     paddle::dialect::ForwardOnlyTrait> {
  public:
   using Op::Op;
   static const char *name() { return "pd_op.slice_array"; }
@@ -423,7 +432,8 @@ class SliceArrayDenseOp
     : public pir::Op<SliceArrayDenseOp,
                      paddle::dialect::OpYamlInfoInterface,
                      paddle::dialect::InferMetaInterface,
-                     paddle::dialect::GetKernelTypeForVarInterface> {
+                     paddle::dialect::GetKernelTypeForVarInterface,
+                     paddle::dialect::ForwardOnlyTrait> {
  public:
   using Op::Op;
   static const char *name() { return "pd_op.slice_array_dense"; }
@@ -455,7 +465,8 @@ class AssignArrayOp
     : public pir::Op<AssignArrayOp,
                      paddle::dialect::OpYamlInfoInterface,
                      paddle::dialect::InferMetaInterface,
-                     paddle::dialect::GetKernelTypeForVarInterface> {
+                     paddle::dialect::GetKernelTypeForVarInterface,
+                     paddle::dialect::ForwardOnlyTrait> {
  public:
   using Op::Op;
   static const char *name() { return "pd_op.assign_array"; }
@@ -486,7 +497,8 @@ class AssignArray_Op
     : public pir::Op<AssignArray_Op,
                      paddle::dialect::OpYamlInfoInterface,
                      paddle::dialect::InferMetaInterface,
-                     paddle::dialect::GetKernelTypeForVarInterface> {
+                     paddle::dialect::GetKernelTypeForVarInterface,
+                     paddle::dialect::ForwardOnlyTrait> {
  public:
   using Op::Op;
   static const char *name() { return "pd_op.assign_array_"; }
@@ -769,7 +781,8 @@ class ArrayPopOp : public pir::Op<ArrayPopOp,
                                   paddle::dialect::OpYamlInfoInterface,
                                   paddle::dialect::InferMetaInterface,
                                   paddle::dialect::GetKernelTypeForVarInterface,
-                                  InplaceTrait> {
+                                  InplaceTrait,
+                                  paddle::dialect::ForwardOnlyTrait> {
  public:
   using Op::Op;
   static const char *name() { return "pd_op.array_pop"; }
 
@@ -0,0 +1,30 @@
+// Copyright (c) 2023 PaddlePaddle Authors. All Rights Reserved.
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     https://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+#pragma once
+
+#include "paddle/pir/include/core/op_base.h"
+
+namespace paddle {
+namespace dialect {
+class ForwardOnlyTrait : public pir::OpTraitBase<ForwardOnlyTrait> {
+ public:
+  explicit ForwardOnlyTrait(pir::Operation *op)
+      : pir::OpTraitBase<ForwardOnlyTrait>(op) {}
+};
+
+}  // namespace dialect
+}  // namespace paddle
+
+IR_DECLARE_EXPLICIT_TYPE_ID(paddle::dialect::ForwardOnlyTrait)
@@ -13,12 +13,14 @@
 // limitations under the License.
 
 #include "paddle/fluid/pir/dialect/operator/trait/custom_vjp.h"
+#include "paddle/fluid/pir/dialect/operator/trait/forward_only.h"
 #include "paddle/fluid/pir/dialect/operator/trait/inplace.h"
 #ifdef PADDLE_WITH_DNNL
 #include "paddle/fluid/pir/dialect/operator/trait/onednn.h"
 #endif
 IR_DEFINE_EXPLICIT_TYPE_ID(paddle::dialect::InplaceTrait)
 IR_DEFINE_EXPLICIT_TYPE_ID(paddle::dialect::CustomVjpTrait)
+IR_DEFINE_EXPLICIT_TYPE_ID(paddle::dialect::ForwardOnlyTrait)
 
 #ifdef PADDLE_WITH_DNNL
 IR_DEFINE_EXPLICIT_TYPE_ID(paddle::dialect::OneDNNTrait)
 
@@ -209,6 +209,7 @@ limitations under the License. */
 #include "paddle/fluid/pir/dialect/operator/ir/api_builder.h"
 #include "paddle/fluid/pir/dialect/operator/ir/manual_pylayer_op.h"
 #include "paddle/fluid/pir/dialect/operator/trait/custom_vjp.h"
+#include "paddle/fluid/pir/dialect/operator/trait/forward_only.h"
 #include "paddle/fluid/prim/utils/eager/eager_tensor_operants.h"
 #include "paddle/fluid/prim/utils/static/static_tensor_operants.h"
 #include "paddle/fluid/primitive/base/decomp_trans.h"
@@ -969,6 +970,20 @@ void BindVjp(pybind11::module *m) {
            Returns:
                out (bool): True means that the op has custom vjp rules, False means it does not.
            )DOC");
+  m->def(
+      "is_forward_only",
+      [](pir::Operation &op) -> py::bool_ {
+        return op.info().HasTrait<paddle::dialect::ForwardOnlyTrait>();
+      },
+      R"DOC(
+           Return whether an op is forward only op.
+
+           Args:
+               op (pir::Operation): op to be checked
+
+           Returns:
+               out (bool): True means that the op is forward only op, False means it does not.
+           )DOC");
 }
 
 void BindDecomp(pybind11::module *m) {
 
@@ -75,6 +75,7 @@
     func : distribute_fpn_proposals
     data_type : fpn_rois
   optional : rois_num
+  traits : paddle::dialect::ForwardOnlyTrait
 
 - op : divide
   args : (Tensor x, Tensor y)
@@ -140,6 +141,7 @@
   kernel :
     func : equal
   inplace: (x -> out)
+  traits : paddle::dialect::ForwardOnlyTrait
 
 - op : floor_divide
   args : (Tensor x, Tensor y)
@@ -149,6 +151,7 @@
   kernel :
     func : floor_divide
   inplace: (x -> out)
+  traits : paddle::dialect::ForwardOnlyTrait
 
 - op : fused_adam_
   args : (Tensor[] params, Tensor[] grads, Tensor learning_rate, Tensor[] moments1, Tensor[] moments2, Tensor[] beta1_pows, Tensor[] beta2_pows, Tensor[] master_params, Tensor skip_update, Scalar beta1, Scalar beta2, Scalar epsilon, int chunk_size, float weight_decay, bool use_adamw, bool multi_precision, bool use_global_beta_pow)
@@ -176,6 +179,7 @@
   kernel :
     func : greater_equal
   inplace: (x -> out)
+  traits : paddle::dialect::ForwardOnlyTrait
 
 - op : greater_than
   args : (Tensor x, Tensor y)
@@ -185,6 +189,7 @@
   kernel :
     func : greater_than
   inplace: (x -> out)
+  traits : paddle::dialect::ForwardOnlyTrait
 
 - op : hardswish
   args : (Tensor x)
@@ -204,6 +209,7 @@
   kernel :
     func : less_equal
   inplace: (x -> out)
+  traits : paddle::dialect::ForwardOnlyTrait
 
 - op : less_than
   args : (Tensor x, Tensor y)
@@ -213,6 +219,7 @@
   kernel :
     func : less_than
   inplace: (x -> out)
+  traits : paddle::dialect::ForwardOnlyTrait
 
 - op : matmul
   args : (Tensor x, Tensor y, bool transpose_x = false, bool transpose_y = false)
@@ -273,6 +280,7 @@
   kernel :
     func : not_equal
   inplace: (x -> out)
+  traits : paddle::dialect::ForwardOnlyTrait
 
 - op : remainder
   args : (Tensor x, Tensor y)
@@ -282,6 +290,7 @@
   kernel :
     func : remainder
   inplace : (x -> out)
+  traits : paddle::dialect::ForwardOnlyTrait
 
 - op : reshape
   args : (Tensor x, IntArray shape)
 
@@ -136,6 +136,7 @@
   extra_args : str mkldnn_data_type="float32"
 
 - op : gaussian
+  traits : paddle::dialect::ForwardOnlyTrait
 
 - op : gelu
   extra_args : str mkldnn_data_type="float32"
@@ -228,6 +229,7 @@
 
 - op : prior_box
   extra_args : bool use_quantizer=false, str mkldnn_data_type="float32"
+  traits : paddle::dialect::ForwardOnlyTrait
 
 - op : relu
 
@@ -253,6 +255,7 @@
 
 - op : shape
   extra_args : str mkldnn_data_type="float32"
+  traits : paddle::dialect::ForwardOnlyTrait
 
 - op : shuffle_channel