【Infer Symbolic Shape No.153】add symbol_infer_interface for lstm (#68907)

ooooo-create · web-flow · commit d25eca11198e · 2024-11-04T11:18:13.000+08:00
* add symbol_infer_interface for lstm

* test

* Clear

* refine
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
@@ -3039,13 +3039,74 @@ bool SpectralNormOpInferSymbolicShape(
 
   return true;
 }
-// bool LstmOpInferSymbolicShape(pir::Operation *op,
-//                               pir::InferSymbolicShapeContext
-//                               *infer_context)
-//                               {
-//   // pass
-//   return true;
-// }
+bool LstmOpInferSymbolicShape(pir::Operation *op,
+                              pir::InferSymbolicShapeContext *infer_context) {
+  const auto &input_shape =
+      infer_context->GetShapeOrDataForValue(op->operand_source(0)).shape();
+  const auto &weight_shape =
+      infer_context->GetShapeOrDataForValue(op->operand_source(3)).shape();
+  const symbol::ShapeOrDataDimExprs &bias_shape_or_data =
+      infer_context->GetShapeOrDataForValue(op->operand_source(4));
+  const auto &bias_shape = bias_shape_or_data.shape();
+  bool use_peepholes =
+      op->attribute<pir::BoolAttribute>("use_peepholes").data();
+  bool is_test = op->attribute<pir::BoolAttribute>("is_test").data();
+  PADDLE_ENFORCE_EQ(
+      input_shape.size(),
+      2,
+      common::errors::InvalidArgument(
+          "Input(X)'s rank must be 2, but received %d.", input_shape.size()));
+  if (op->operand_source(1)) {
+    const auto &h0_shape =
+        infer_context->GetShapeOrDataForValue(op->operand_source(1)).shape();
+    const auto &c0_shape =
+        infer_context->GetShapeOrDataForValue(op->operand_source(2)).shape();
+    size_t dims = h0_shape.size();
+    for (size_t i = 0; i < dims; i++) {
+      infer_context->AddEqualCstr(h0_shape[i], c0_shape[i]);
+    }
+  }
+
+  PADDLE_ENFORCE_EQ(
+      weight_shape.size(),
+      2,
+      common::errors::InvalidArgument(
+          "The rank of Input(Weight) should be 2, but received %d.",
+          weight_shape.size()));
+  symbol::DimExpr frame_size = input_shape[1] / symbol::DimExpr(4);
+  infer_context->AddEqualCstr(weight_shape[0], frame_size);
+  infer_context->AddEqualCstr(weight_shape[1], frame_size * symbol::DimExpr(4));
+  PADDLE_ENFORCE_EQ(bias_shape.size(),
+                    2,
+                    common::errors::InvalidArgument(
+                        "The rank of Input(Bias) should be 2, but received %d.",
+                        bias_shape.size()));
+  infer_context->AddEqualCstr(bias_shape[0], symbol::DimExpr{1});
+  if (use_peepholes) {
+    infer_context->AddEqualCstr(bias_shape[1], frame_size * symbol::DimExpr{7});
+  } else {
+    infer_context->AddEqualCstr(bias_shape[1], frame_size * symbol::DimExpr{4});
+  }
+  std::vector<symbol::DimExpr> out_shape = {input_shape[0], frame_size};
+  symbol::ShapeOrDataDimExprs out_shape_or_data =
+      symbol::ShapeOrDataDimExprs{symbol::TensorShapeOrDataDimExprs(out_shape)};
+  infer_context->SetShapeOrDataForValue(op->result(0), out_shape_or_data);
+  infer_context->SetShapeOrDataForValue(op->result(1), out_shape_or_data);
+  if (!is_test) {
+    infer_context->SetShapeOrDataForValue(
+        op->result(2),
+        symbol::ShapeOrDataDimExprs{
+            symbol::TensorShapeOrDataDimExprs(input_shape)});
+    infer_context->SetShapeOrDataForValue(op->result(3), out_shape_or_data);
+  } else {
+    infer_context->SetShapeOrDataForValue(
+        op->result(2),
+        symbol::ShapeOrDataDimExprs{
+            symbol::TensorShapeOrDataDimExprs(input_shape)});
+    infer_context->SetShapeOrDataForValue(op->result(3), out_shape_or_data);
+  }
+  return true;
+}
 
 // bool MergedAdamOpInferSymbolicShape(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
@@ -78,7 +78,7 @@ OP_DECLARE_INFER_SYMBOLIC_SHAPE(Linspace)
 OP_DECLARE_INFER_SYMBOLIC_SHAPE(LinearInterp)
 OP_DECLARE_INFER_SYMBOLIC_SHAPE(LlmInt8Linear)
 OP_DECLARE_INFER_SYMBOLIC_SHAPE(Logspace)
-// OP_DECLARE_INFER_SYMBOLIC_SHAPE(Lstm)
+OP_DECLARE_INFER_SYMBOLIC_SHAPE(Lstm)
 OP_DECLARE_INFER_SYMBOLIC_SHAPE(MatchMatrixTensor)
 OP_DECLARE_INFER_SYMBOLIC_SHAPE(MaskedMultiheadAttention_)
 // OP_DECLARE_INFER_SYMBOLIC_SHAPE(MergedAdam)
diff --git a/paddle/phi/ops/yaml/ops.yaml b/paddle/phi/ops/yaml/ops.yaml
@@ -3157,6 +3157,7 @@
   optional: h0, c0
   intermediate: batch_gate, batch_cell_pre_act
   backward: lstm_grad
+  interfaces : paddle::dialect::InferSymbolicShapeInterface
 
 - op : lstsq
   args : (Tensor x, Tensor y, Scalar rcond=0.0f, str driver="gels")
