【Infer Symbolic Shape】【BUAA】Add sparse_attention op (#67462)

Guanhuachen2003 · web-flow · commit ee1eb4853084 · 2024-08-23T14:50:32.000+08:00
* add sparseattention

* modify codestyle

* Update multiary_infer_sym.cc

* add constraints

* Update multiary_infer_sym.h

* Update multiary_infer_sym.cc

* modify CODESTYLE

* Update multiary_infer_sym.cc

* Update multiary_infer_sym.cc

* modify codestyle

* Update multiary_infer_sym.h

* remove unused variables
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
@@ -2067,12 +2067,95 @@ bool RmsNormOpInferSymbolicShape(
 //   return true;
 // }
 
-// bool SparseAttentionOpInferSymbolicShape(pir::Operation *op,
-//                                          pir::InferSymbolicShapeContext
-//                                          *infer_context) {
-//   // pass
-//   return true;
-// }
+bool SparseAttentionOpInferSymbolicShape(
+    pir::Operation *op, pir::InferSymbolicShapeContext *infer_context) {
+  // Get symbolic shapes of the input tensors
+  const auto &q_shape_or_data =
+      infer_context->GetShapeOrDataForValue(op->operand_source(0));
+  const auto &k_shape_or_data =
+      infer_context->GetShapeOrDataForValue(op->operand_source(1));
+  const auto &v_shape_or_data =
+      infer_context->GetShapeOrDataForValue(op->operand_source(2));
+  const auto &offset_shape_or_data =
+      infer_context->GetShapeOrDataForValue(op->operand_source(3));
+  const auto &columns_shape_or_data =
+      infer_context->GetShapeOrDataForValue(op->operand_source(4));
+
+  std::vector<symbol::DimExpr> q_dims = q_shape_or_data.shape();
+  std::vector<symbol::DimExpr> k_dims = k_shape_or_data.shape();
+  std::vector<symbol::DimExpr> v_dims = v_shape_or_data.shape();
+  std::vector<symbol::DimExpr> offset_dims = offset_shape_or_data.shape();
+  std::vector<symbol::DimExpr> columns_dims = columns_shape_or_data.shape();
+
+  // Ensure the input tensors have the expected rank
+  PADDLE_ENFORCE_EQ(q_dims.size(),
+                    4UL,
+                    common::errors::InvalidArgument(
+                        "Dimension in query's shapes should be 4."));
+  PADDLE_ENFORCE_EQ(k_dims.size(),
+                    4UL,
+                    common::errors::InvalidArgument(
+                        "Dimension in key's shapes should be 4."));
+  PADDLE_ENFORCE_EQ(v_dims.size(),
+                    4UL,
+                    common::errors::InvalidArgument(
+                        "Dimension in value's shapes should be 4."));
+  PADDLE_ENFORCE_EQ(offset_dims.size(),
+                    3UL,
+                    common::errors::InvalidArgument(
+                        "Dimension in offset's shapes should be 3."));
+  PADDLE_ENFORCE_EQ(columns_dims.size(),
+                    3UL,
+                    common::errors::InvalidArgument(
+                        "Dimension in columns' shapes should be 3."));
+
+  // Add equality constraints between corresponding dimensions
+  infer_context->AddEqualCstr(q_dims[0], k_dims[0]);  // batch_size
+  infer_context->AddEqualCstr(q_dims[1], k_dims[1]);  // num_heads
+  infer_context->AddEqualCstr(q_dims[2], k_dims[2]);  // M (seq_len)
+  infer_context->AddEqualCstr(q_dims[3], k_dims[3]);  // N (head_dim)
+  infer_context->AddEqualCstr(v_dims[0], k_dims[0]);  // batch_size
+  infer_context->AddEqualCstr(v_dims[1], k_dims[1]);  // num_heads
+  infer_context->AddEqualCstr(v_dims[3], k_dims[3]);  // head_dim
+
+  // Ensure that offset and columns dimensions are consistent with the input
+  // tensors
+  infer_context->AddEqualCstr(offset_dims[0], q_dims[0]);      // batch_size
+  infer_context->AddEqualCstr(offset_dims[1], q_dims[1]);      // num_heads
+  infer_context->AddEqualCstr(offset_dims[2], q_dims[2] + 1);  // seq_len + 1
+
+  infer_context->AddEqualCstr(columns_dims[0], q_dims[0]);  // batch_size
+  infer_context->AddEqualCstr(columns_dims[1], q_dims[1]);  // num_heads
+
+  // Prepare the output tensor shapes
+  auto batch_size = q_dims[0];
+  auto num_heads = q_dims[1];
+  auto M = q_dims[2];
+  auto N = q_dims[3];
+  auto sparse_nnz = columns_dims[2];
+
+  // Set output tensor shapes
+  std::vector<symbol::DimExpr> out_dims = {batch_size, num_heads, M, N};
+  infer_context->SetShapeOrDataForValue(
+      op->result(0),
+      symbol::ShapeOrDataDimExprs{symbol::TensorShapeOrDataDimExprs(out_dims)});
+
+  std::vector<symbol::DimExpr> sparse_dot_sdd_dims = {
+      batch_size, num_heads, sparse_nnz};
+  infer_context->SetShapeOrDataForValue(
+      op->result(1),
+      symbol::ShapeOrDataDimExprs{
+          symbol::TensorShapeOrDataDimExprs(sparse_dot_sdd_dims)});
+
+  std::vector<symbol::DimExpr> softmax_dims = {
+      batch_size, num_heads, sparse_nnz};
+  infer_context->SetShapeOrDataForValue(
+      op->result(2),
+      symbol::ShapeOrDataDimExprs{
+          symbol::TensorShapeOrDataDimExprs(softmax_dims)});
+
+  return true;
+}
 
 // bool SpectralNormOpInferSymbolicShape(pir::Operation *op,
 //                                       pir::InferSymbolicShapeContext
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
@@ -103,7 +103,7 @@ OP_DECLARE_INFER_SYMBOLIC_SHAPE(RoiPool)
 OP_DECLARE_INFER_SYMBOLIC_SHAPE(RoiAlign)
 OP_DECLARE_INFER_SYMBOLIC_SHAPE(SpectralNorm)
 // OP_DECLARE_INFER_SYMBOLIC_SHAPE(SequenceConv)
-// OP_DECLARE_INFER_SYMBOLIC_SHAPE(SparseAttention)
+OP_DECLARE_INFER_SYMBOLIC_SHAPE(SparseAttention)
 OP_DECLARE_INFER_SYMBOLIC_SHAPE(Stack)
 OP_DECLARE_INFER_SYMBOLIC_SHAPE(SigmoidCrossEntropyWithLogits)
 OP_DECLARE_INFER_SYMBOLIC_SHAPE(SigmoidCrossEntropyWithLogits_)
diff --git a/paddle/phi/ops/yaml/ops.yaml b/paddle/phi/ops/yaml/ops.yaml
@@ -4427,7 +4427,7 @@
   optional: key_padding_mask, attn_mask
   intermediate: sparse_dot_sdd, softmax
   backward: sparse_attention_grad
-  # interfaces : paddle::dialect::InferSymbolicShapeInterface
+  interfaces : paddle::dialect::InferSymbolicShapeInterface
 
 - op : spectral_norm
   args : (Tensor weight, Tensor u, Tensor v, int dim = 0, int power_iters = 1, float eps = 1e-12f)
