[CINN] Add fuesd_attention_grad symbol infer

paddle/fluid/pir/dialect/operator/interface/infer_symbolic_shape/backward_infer_sym.cc

@@ -0,0 +1,96 @@
+// Copyright (c) 2024 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
+//
+//     http://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.
+
+#include "paddle/fluid/pir/dialect/operator/interface/infer_symbolic_shape/backward_infer_sym.h"
+#include "paddle/fluid/pir/dialect/operator/interface/infer_symbolic_shape/infer_sym_utils.h"
+
+namespace paddle::dialect {
+
+bool FusedAttentionGradOpInferSymbolicShape(
+    pir::Operation *op, pir::InferSymbolicShapeContext *infer_context) {
+  bool is_test = op->attribute<pir::BoolAttribute>("is_test").data();
+  PADDLE_ENFORCE_EQ(is_test,
+                    false,
+                    common::errors::InvalidArgument(
+                        "GradOp is only callable when is_test is false"));
+  bool pre_layer_norm =
+      op->attribute<pir::BoolAttribute>("pre_layer_norm").data();
+  auto same_shape_infer = [&](pir::Value &&dst, pir::Value &&src) {
+    auto src_shape = infer_context->GetShapeOrDataForValue(src).shape();
+    infer_context->SetShapeOrDataForValue(
+        dst,
+        symbol::ShapeOrDataDimExprs{
+            symbol::TensorShapeOrDataDimExprs(src_shape)});
+  };
+  if (!pre_layer_norm) {
+    if (!paddle::dialect::details::IsFakeValue(op->result(6)) &&
+        op->operand_source(11)) {
+      same_shape_infer(op->result(6), op->operand_source(11));
+    }
+    if (!paddle::dialect::details::IsFakeValue(op->result(7)) &&
+        op->operand_source(12)) {
+      same_shape_infer(op->result(7), op->operand_source(12));
+    }
+  }
+  if (pre_layer_norm && op->operand_source(9)) {
+    if (!paddle::dialect::details::IsFakeValue(op->result(4))) {
+      same_shape_infer(op->result(4), op->operand_source(9));
+    }
+    if (!paddle::dialect::details::IsFakeValue(op->result(5)) &&
+        op->operand_source(10)) {
+      same_shape_infer(op->result(5), op->operand_source(10));
+    }
+  }
+  same_shape_infer(op->result(8), op->operand_source(1));
+  if (!paddle::dialect::details::IsFakeValue(op->result(3)) &&
+      op->operand_source(13)) {
+    same_shape_infer(op->result(3), op->operand_source(8));
+  }
+  if (!paddle::dialect::details::IsFakeValue(op->result(10))) {
+    same_shape_infer(op->result(10), op->operand_source(7));
+  }
+  same_shape_infer(op->result(9), op->operand_source(2));
+  if (!paddle::dialect::details::IsFakeValue(op->result(0))) {
+    same_shape_infer(op->result(0), op->operand_source(3));
+  }
+  if (pre_layer_norm) {
+    if (!paddle::dialect::details::IsFakeValue(op->result(11)) &&
+        op->operand_source(13)) {
+      same_shape_infer(op->result(11), op->operand_source(13));
+    }
+  } else {
+    if (!paddle::dialect::details::IsFakeValue(op->result(12)) &&
+        op->operand_source(18)) {
+      same_shape_infer(op->result(12), op->operand_source(18));
+    }
+  }
+  same_shape_infer(op->result(19), op->operand_source(26));
+  same_shape_infer(op->result(14), op->operand_source(22));
+  same_shape_infer(op->result(15), op->operand_source(20));
+  same_shape_infer(op->result(16), op->operand_source(21));
+  same_shape_infer(op->result(17), op->operand_source(23));
+  same_shape_infer(op->result(18), op->operand_source(25));
+  if (!paddle::dialect::details::IsFakeValue(op->result(2)) &&
+      op->operand_source(6)) {
+    same_shape_infer(op->result(2), op->operand_source(6));
+  }
+  same_shape_infer(op->result(13), op->operand_source(19));
+  if (!paddle::dialect::details::IsFakeValue(op->result(1)) &&
+      op->operand_source(4)) {
+    same_shape_infer(op->result(1), op->operand_source(4));
+  }
+  same_shape_infer(op->result(20), op->operand_source(27));
+  return true;
+}
+}  // namespace paddle::dialect

paddle/fluid/pir/dialect/operator/interface/infer_symbolic_shape/backward_infer_sym.h

@@ -0,0 +1,21 @@
+// Copyright (c) 2024 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
+//
+//     http://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/dialect/shape/utils/shape_analysis.h"
+
+namespace paddle::dialect {
+OP_DECLARE_INFER_SYMBOLIC_SHAPE(FusedAttentionGrad)
+}  // namespace paddle::dialect

paddle/fluid/pir/dialect/operator/interface/infer_symbolic_shape/infer_symbolic_shape.h

@@ -14,6 +14,7 @@
 
 #pragma once
 
+#include "paddle/fluid/pir/dialect/operator/interface/infer_symbolic_shape/backward_infer_sym.h"
 #include "paddle/fluid/pir/dialect/operator/interface/infer_symbolic_shape/binary_infer_sym.h"
 #include "paddle/fluid/pir/dialect/operator/interface/infer_symbolic_shape/cinn_op_infer_sym.h"
 #include "paddle/fluid/pir/dialect/operator/interface/infer_symbolic_shape/element_wise_binary.h"

paddle/phi/ops/yaml/inconsistent/static_backward.yaml

@@ -203,6 +203,7 @@
     data_type : x
   optional: ln_scale, ln_bias, qkv_bias, src_mask, out_linear_bias, ln_scale_2, ln_bias_2, qkv_bias_grad, qkv_bias_out_grad, src_mask_out_grad, out_linear_bias_grad, ln_scale_grad, ln_bias_grad, ln_scale_2_grad, ln_bias_2_grad, ln_out_grad, bias_dropout_residual_out_grad, ln_out, ln_mean, ln_var,  ln_mean_2, ln_var_2, bias_dropout_residual_out, qkv_bias, qkv_bias_out, src_mask, src_mask_out, out_linear_bias
   no_need_buffer: qkv_bias_out, qkv_out, qk_out, qktv_out, out_linear_out, src_mask
+  interfaces : paddle::dialect::InferSymbolicShapeInterface
 
 - backward_op : fused_feedforward_grad
   forward: fused_feedforward (Tensor x, Tensor dropout1_seed, Tensor dropout2_seed, Tensor linear1_weight, Tensor linear1_bias, Tensor linear2_weight, Tensor linear2_bias, Tensor ln1_scale, Tensor ln1_bias, Tensor ln2_scale, Tensor ln2_bias, bool pre_layer_norm, float ln1_epsilon, float ln2_epsilon, str act_method, float dropout1_prob, float dropout2_prob, str dropout1_implementation, str dropout2_implementation, bool is_test, bool dropout1_fix_seed, bool dropout2_fix_seed, int dropout1_seed_val, int dropout2_seed_val, bool add_residual, int ring_id) -> Tensor(out), Tensor(dropout1_mask), Tensor(dropout2_mask), Tensor(ln1_mean), Tensor(ln1_variance), Tensor(ln2_mean), Tensor(ln2_variance), Tensor(linear1_out), Tensor(ln1_out), Tensor(dropout1_out), Tensor(dropout2_out)