[CINN] add fused_attention infersymbolicShape (#69414)

ooooo-create · web-flow · commit 94959773827d · 2024-11-19T20:43:57.000+08:00
* fused_attention

* Fix compile

* refine code

* refine code
diff --git a/paddle/fluid/pir/dialect/operator/interface/infer_symbolic_shape/multiary_infer_sym.cc b/paddle/fluid/pir/dialect/operator/interface/infer_symbolic_shape/multiary_infer_sym.cc
@@ -1658,12 +1658,268 @@ bool FusedFeedforwardOpInferSymbolicShape(
   return true;
 }
 
-// bool FusedAttentionOpInferSymbolicShape(pir::Operation *op,
-//                                         pir::InferSymbolicShapeContext
-//                                         *infer_context) {
-//   // pass
-//   return true;
-// }
+bool FusedAttentionOpInferSymbolicShape(
+    pir::Operation *op, pir::InferSymbolicShapeContext *infer_context) {
+  const auto &x_shape_or_data =
+      infer_context->GetShapeOrDataForValue(op->operand_source(0));
+  const auto &qkv_weight_shape_or_data =
+      infer_context->GetShapeOrDataForValue(op->operand_source(3));
+  const auto &cache_kv_shape_or_data =
+      infer_context->GetShapeOrDataForValue(op->operand_source(5));
+  const auto &src_mask_shape_or_data =
+      infer_context->GetShapeOrDataForValue(op->operand_source(6));
+  const auto &qkv_bias_shape_or_data =
+      infer_context->GetShapeOrDataForValue(op->operand_source(4));
+  const std::vector<symbol::DimExpr> &x_shape = x_shape_or_data.shape();
+  const std::vector<symbol::DimExpr> &qkv_weight_shape =
+      qkv_weight_shape_or_data.shape();
+  symbol::DimExpr dim_head = 0;
+  symbol::DimExpr hidden_size = 0;
+  symbol::DimExpr nranks = 1;
+  bool transpose_qkv_wb =
+      op->attribute<pir::BoolAttribute>("transpose_qkv_wb").data();
+  int num_heads_ = op->attribute<pir::Int32Attribute>("num_heads").data();
+  symbol::DimExpr num_heads = symbol::DimExpr(num_heads_);
+  int ring_id = op->attribute<pir::Int32Attribute>("ring_id").data();
+  bool pre_layer_norm =
+      op->attribute<pir::BoolAttribute>("pre_layer_norm").data();
+  bool is_test = op->attribute<pir::BoolAttribute>("is_test").data();
+  if (transpose_qkv_wb) {
+    PADDLE_ENFORCE_EQ(qkv_weight_shape.size(),
+                      2,
+                      common::errors::InvalidArgument(
+                          "The dimensions of qkv_weight must be 2 if enable"
+                          "transpose_qkv_wb: (dim_embed, 3 * dim_embed),"
+                          "but received dimensions of"
+                          "Input is [%d]",
+                          qkv_weight_shape.size()));
+    PADDLE_ENFORCE_GT(num_heads_,
+                      0,
+                      common::errors::InvalidArgument(
+                          "The num_heads must be provided and greater than 0 "
+                          "if enable transpose_qkv_wb, but we got %d.",
+                          num_heads));
+    infer_context->AddEqualCstr((qkv_weight_shape[0] / num_heads) * num_heads,
+                                qkv_weight_shape[0]);
+    if (ring_id == -1) {
+      infer_context->AddEqualCstr(qkv_weight_shape[0] * symbol::DimExpr(3),
+                                  qkv_weight_shape[1]);
+    } else {
+      nranks = (qkv_weight_shape[0] * symbol::DimExpr(3)) / qkv_weight_shape[1];
+    }
+    dim_head = qkv_weight_shape[0] / (num_heads * nranks);
+    hidden_size = qkv_weight_shape[0];
+  } else {
+    PADDLE_ENFORCE_EQ(qkv_weight_shape.size(),
+                      4,
+                      common::errors::InvalidArgument(
+                          "The dimensions of qkv_weight must be 4 if not"
+                          "enable transpose_qkv_wb: (3, num_head, dim_head, "
+                          "dim_embed), but received [%d]",
+                          qkv_weight_shape.size()));
+    infer_context->AddEqualCstr(qkv_weight_shape[0], symbol::DimExpr(3));
+    if (ring_id == -1) {
+      infer_context->AddEqualCstr(qkv_weight_shape[1] * qkv_weight_shape[2],
+                                  qkv_weight_shape[3]);
+    }
+    num_heads = qkv_weight_shape[1];
+    dim_head = qkv_weight_shape[2];
+    hidden_size = qkv_weight_shape[3];
+  }
+  PADDLE_ENFORCE_EQ(x_shape.size(),
+                    3,
+                    common::errors::InvalidArgument(
+                        "The dimensions of x must be 3 (batch_size, seq_len, "
+                        "dim_embed), but received dimensions of Input is [%d]",
+                        x_shape.size()));
+  infer_context->AddEqualCstr(x_shape[2], hidden_size);
+  if (pre_layer_norm) {
+    infer_context->SetShapeOrDataForValue(
+        op->result(0),
+        symbol::ShapeOrDataDimExprs{
+            symbol::TensorShapeOrDataDimExprs({x_shape[0] * x_shape[1]})});
+    infer_context->SetShapeOrDataForValue(
+        op->result(1),
+        symbol::ShapeOrDataDimExprs{
+            symbol::TensorShapeOrDataDimExprs({x_shape[0] * x_shape[1]})});
+    infer_context->SetShapeOrDataForValue(
+        op->result(2),
+        symbol::ShapeOrDataDimExprs{
+            symbol::TensorShapeOrDataDimExprs(x_shape)});
+  } else {
+    // The follwing three code used to set unoptional output value.
+    // Now it's result related to the infermeta.
+    infer_context->SetSymbolForValueByStaticShape(op->result(0));
+    infer_context->SetSymbolForValueByStaticShape(op->result(1));
+    infer_context->SetSymbolForValueByStaticShape(op->result(2));
+    if (paddle::dialect::details::IsFakeValue(op->result(15))) {
+      infer_context->SetSymbolForValueByStaticShape(op->result(15));
+    } else {
+      infer_context->SetShapeOrDataForValue(
+          op->result(15),
+          symbol::ShapeOrDataDimExprs{
+              symbol::TensorShapeOrDataDimExprs({x_shape[0] * x_shape[1]})});
+    }
+    if (paddle::dialect::details::IsFakeValue(op->result(16))) {
+      infer_context->SetSymbolForValueByStaticShape(op->result(16));
+    } else {
+      infer_context->SetShapeOrDataForValue(
+          op->result(16),
+          symbol::ShapeOrDataDimExprs{
+              symbol::TensorShapeOrDataDimExprs({x_shape[0] * x_shape[1]})});
+    }
+    if (paddle::dialect::details::IsFakeValue(op->result(17))) {
+      infer_context->SetSymbolForValueByStaticShape(op->result(17));
+    } else {
+      infer_context->SetShapeOrDataForValue(
+          op->result(17),
+          symbol::ShapeOrDataDimExprs{
+              symbol::TensorShapeOrDataDimExprs(x_shape)});
+    }
+  }
+  if (transpose_qkv_wb) {
+    // [batch_size, seq_len, 3 * num_heads * dim_head]
+    infer_context->SetShapeOrDataForValue(
+        op->result(3),
+        symbol::ShapeOrDataDimExprs{symbol::TensorShapeOrDataDimExprs(
+            {x_shape[0],
+             x_shape[1],
+             symbol::DimExpr(3) * num_heads * dim_head})});
+    if (!qkv_bias_shape_or_data.isa<symbol::NullShapeOrDataDimExpr>()) {
+      infer_context->SetShapeOrDataForValue(
+          op->result(4),
+          symbol::ShapeOrDataDimExprs{symbol::TensorShapeOrDataDimExprs(
+              {x_shape[0],
+               x_shape[1],
+               symbol::DimExpr(3) * num_heads * dim_head})});
+    } else {
+      infer_context->SetSymbolForValueByStaticShape(op->result(4));
+    }
+  } else {
+    // [batch_size, seq_len, 3, num_head, head_size]
+    infer_context->SetShapeOrDataForValue(
+        op->result(3),
+        symbol::ShapeOrDataDimExprs{
+            symbol::TensorShapeOrDataDimExprs({x_shape[0],
+                                               x_shape[1],
+                                               symbol::DimExpr(3),
+                                               num_heads,
+                                               dim_head})});
+    if (!qkv_bias_shape_or_data.isa<symbol::NullShapeOrDataDimExpr>()) {
+      infer_context->SetShapeOrDataForValue(
+          op->result(4),
+          symbol::ShapeOrDataDimExprs{
+              symbol::TensorShapeOrDataDimExprs({x_shape[0],
+                                                 x_shape[1],
+                                                 symbol::DimExpr(3),
+                                                 num_heads,
+                                                 dim_head})});
+    } else {
+      infer_context->SetSymbolForValueByStaticShape(op->result(4));
+    }
+  }
+  // [3, batch_size, num_head, seq_len, head_size]
+  infer_context->SetShapeOrDataForValue(
+      op->result(5),
+      symbol::ShapeOrDataDimExprs{symbol::TensorShapeOrDataDimExprs(
+          {symbol::DimExpr(3), x_shape[0], num_heads, x_shape[1], dim_head})});
+  // cache_seq_len + seq_len if cache else seq_len
+  symbol::DimExpr out_seq_len = x_shape[1];
+  if (!cache_kv_shape_or_data.isa<symbol::NullShapeOrDataDimExpr>()) {
+    const std::vector<symbol::DimExpr> &cache_kv_shape =
+        cache_kv_shape_or_data.shape();
+    PADDLE_ENFORCE_EQ(
+        cache_kv_shape.size(),
+        5,
+        common::errors::InvalidArgument(
+            "The CacheKV must be 5 dims, but got %d", cache_kv_shape.size()));
+    infer_context->AddEqualCstr(cache_kv_shape[0], symbol::DimExpr(2));
+    infer_context->AddEqualCstr(cache_kv_shape[1], x_shape[0]);
+    infer_context->AddEqualCstr(cache_kv_shape[2], num_heads);
+    infer_context->AddEqualCstr(cache_kv_shape[4], dim_head);
+    out_seq_len = out_seq_len + cache_kv_shape[3];
+    // [3, batch_size, num_head, cache_seq_len + seq_len, head_size]
+    if (paddle::dialect::details::IsFakeValue(op->result(18))) {
+      infer_context->SetSymbolForValueByStaticShape(op->result(18));
+    } else {
+      infer_context->SetShapeOrDataForValue(
+          op->result(18),
+          symbol::ShapeOrDataDimExprs{
+              symbol::TensorShapeOrDataDimExprs({cache_kv_shape[0],
+                                                 cache_kv_shape[1],
+                                                 cache_kv_shape[2],
+                                                 out_seq_len,
+                                                 cache_kv_shape[4]})});
+    }
+  }
+  // [batch, num_head, seq_len, out_seq_len]
+  infer_context->SetShapeOrDataForValue(
+      op->result(6),
+      symbol::ShapeOrDataDimExprs{symbol::TensorShapeOrDataDimExprs(
+          {x_shape[0], num_heads, x_shape[1], out_seq_len})});
+  if (!src_mask_shape_or_data.isa<symbol::NullShapeOrDataDimExpr>()) {
+    infer_context->SetShapeOrDataForValue(
+        op->result(11),
+        symbol::ShapeOrDataDimExprs{symbol::TensorShapeOrDataDimExprs(
+            {x_shape[0], num_heads, x_shape[1], out_seq_len})});
+  } else {
+    // The follwing code used to set unoptional output value.
+    // Now it's result related to the infermeta.
+    infer_context->SetSymbolForValueByStaticShape(op->result(11));
+  }
+  // the same as QKOut's shape
+  infer_context->SetShapeOrDataForValue(
+      op->result(10),
+      symbol::ShapeOrDataDimExprs{symbol::TensorShapeOrDataDimExprs(
+          {x_shape[0], num_heads, x_shape[1], out_seq_len})});
+  if (!is_test) {
+    infer_context->SetShapeOrDataForValue(
+        op->result(9),
+        symbol::ShapeOrDataDimExprs{symbol::TensorShapeOrDataDimExprs(
+            {x_shape[0], num_heads, x_shape[1], out_seq_len})});
+  } else {
+    // The follwing code used to set unoptional output value.
+    // Now it's result related to the infermeta.
+    infer_context->SetSymbolForValueByStaticShape(op->result(9));
+  }
+
+  infer_context->SetShapeOrDataForValue(
+      op->result(8),
+      symbol::ShapeOrDataDimExprs{symbol::TensorShapeOrDataDimExprs(
+          {x_shape[0], num_heads, x_shape[1], out_seq_len})});
+
+  // [batch_size, num_heads, seq_len, head_dim]
+  infer_context->SetShapeOrDataForValue(
+      op->result(7),
+      symbol::ShapeOrDataDimExprs{symbol::TensorShapeOrDataDimExprs(
+          {x_shape[0], num_heads, x_shape[1], dim_head})});
+
+  // [batch_size, seq_len, number of heads*head size]
+  infer_context->SetShapeOrDataForValue(
+      op->result(12),
+      symbol::ShapeOrDataDimExprs{symbol::TensorShapeOrDataDimExprs(
+          {x_shape[0], x_shape[1], num_heads, dim_head})});
+
+  infer_context->SetShapeOrDataForValue(
+      op->result(13),
+      symbol::ShapeOrDataDimExprs{symbol::TensorShapeOrDataDimExprs(x_shape)});
+
+  if (!is_test) {
+    infer_context->SetShapeOrDataForValue(
+        op->result(14),
+        symbol::ShapeOrDataDimExprs{
+            symbol::TensorShapeOrDataDimExprs(x_shape)});
+  } else {
+    // The follwing code used to set unoptional output value.
+    // Now it's result related to the infermeta.
+    infer_context->SetSymbolForValueByStaticShape(op->result(14));
+  }
+
+  infer_context->SetShapeOrDataForValue(
+      op->result(19),
+      symbol::ShapeOrDataDimExprs{symbol::TensorShapeOrDataDimExprs(x_shape)});
+  return true;
+}
 
 bool FlashAttnVarlenQkvpackedOpInferSymbolicShape(
     pir::Operation *op, pir::InferSymbolicShapeContext *infer_context) {
diff --git a/paddle/fluid/pir/dialect/operator/interface/infer_symbolic_shape/multiary_infer_sym.h b/paddle/fluid/pir/dialect/operator/interface/infer_symbolic_shape/multiary_infer_sym.h
@@ -54,7 +54,7 @@ OP_DECLARE_INFER_SYMBOLIC_SHAPE(FakeQuantizeRangeAbsMax_)
 OP_DECLARE_INFER_SYMBOLIC_SHAPE(FullWithTensor)
 OP_DECLARE_INFER_SYMBOLIC_SHAPE(FlashAttn)
 OP_DECLARE_INFER_SYMBOLIC_SHAPE(FusedFeedforward)
-// OP_DECLARE_INFER_SYMBOLIC_SHAPE(FusedAttention)
+OP_DECLARE_INFER_SYMBOLIC_SHAPE(FusedAttention)
 OP_DECLARE_INFER_SYMBOLIC_SHAPE(FlashAttnQkvpacked)
 OP_DECLARE_INFER_SYMBOLIC_SHAPE(FlashAttnVarlenQkvpacked)
 // OP_DECLARE_INFER_SYMBOLIC_SHAPE(FlashAttnUnpadded)
diff --git a/paddle/phi/ops/yaml/inconsistent/static_ops.yaml b/paddle/phi/ops/yaml/inconsistent/static_ops.yaml
@@ -1028,7 +1028,7 @@
     func: FusedAttentionInferMeta
   optional: cache_kv, ln_scale, ln_bias, qkv_bias, src_mask, out_linear_bias, ln_scale_2, ln_bias_2, ln_mean_2, ln_var_2, bias_dropout_residual_out, cache_kv_out
   backward: fused_attention_grad
-  # interfaces : paddle::dialect::InferSymbolicShapeInterface
+  interfaces : paddle::dialect::InferSymbolicShapeInterface
 
 - op: fused_feedforward
   args: (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)
diff --git a/paddle/pir/src/dialect/shape/utils/shape_analysis.cc b/paddle/pir/src/dialect/shape/utils/shape_analysis.cc
@@ -559,13 +559,17 @@ void ShapeConstraintIRAnalysis::InferShapeOrDataForValue(Value val) {
         op->dyn_cast<pir::InferSymbolicShapeInterface>();
     if (infer_symbolic_shape_interface) {
       infer_symbolic_shape_interface.InferSymbolicShape(&context_);
+      int index = -1;
       for (auto& result_value : op->results()) {
+        index++;
         if (!result_value || !result_value.type()) {
           continue;
         }
         if (!context_.HasShapeOrDataForValue(result_value)) {
           PADDLE_THROW(common::errors::Fatal(
-              op->name() + " HAS ERROR on InferSymbolicShape!"));
+              op->name() +
+              " HAS ERROR on InferSymbolicShape! The result value with index " +
+              std::to_string(index) + " don't has shape or data."));
         }
       }
     } else {
diff --git a/test/legacy_test/test_fused_attention_op_api.py b/test/legacy_test/test_fused_attention_op_api.py

Original file line number	Diff line number	Diff line change
`@@ -559,13 +559,17 @@ void ShapeConstraintIRAnalysis::InferShapeOrDataForValue(Value val) {`
`559`	`559`	`op->dyn_cast<pir::InferSymbolicShapeInterface>();`
`560`	`560`	`if (infer_symbolic_shape_interface) {`
`561`	`561`	`infer_symbolic_shape_interface.InferSymbolicShape(&context_);`
	`562`	`+ int index = -1;`
`562`	`563`	`for (auto& result_value : op->results()) {`
	`564`	`+ index++;`
`563`	`565`	`if (!result_value \|\| !result_value.type()) {`
`564`	`566`	`continue;`
`565`	`567`	`}`
`566`	`568`	`if (!context_.HasShapeOrDataForValue(result_value)) {`
`567`	`569`	`PADDLE_THROW(common::errors::Fatal(`
`568`		`- op->name() + " HAS ERROR on InferSymbolicShape!"));`
	`570`	`+ op->name() +`
	`571`	`+ " HAS ERROR on InferSymbolicShape! The result value with index " +`
	`572`	`+ std::to_string(index) + " don't has shape or data."));`
`569`	`573`	`}`
`570`	`574`	`}`
`571`	`575`	`} else {`