【Infer Symbolic Shape No.230】【BUAA】Add stft, changed 3 files

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

@@ -1546,11 +1546,68 @@ bool SequenceMaskOpInferSymbolicShape(
 //   return true;
 // }
 
-// bool StftOpInferSymbolicShape(
-//     pir::Operation *op, pir::InferSymbolicShapeContext *infer_context) {
-//   // pass
-//   return true;
-// }
+bool StftOpInferSymbolicShape(pir::Operation *op,
+                              pir::InferSymbolicShapeContext *infer_context) {
+  const auto &x_shape_or_data =
+      infer_context->GetShapeOrDataForValue(op->operand_source(0));
+  const auto &window_shape_or_data =
+      infer_context->GetShapeOrDataForValue(op->operand_source(1));
+  const std::vector<symbol::DimExpr> &x_shape = x_shape_or_data.shape();
+  const std::vector<symbol::DimExpr> &window_shape =
+      window_shape_or_data.shape();
+
+  int n_fft = op->attribute<pir::Int32Attribute>("n_fft").data();
+  int hop_length = op->attribute<pir::Int32Attribute>("hop_length").data();
+  bool onesided = op->attribute<pir::BoolAttribute>("onesided").data();
+
+  const int x_rank = x_shape.size();
+
+  PADDLE_ENFORCE_EQ(
+      x_rank,
+      2,
+      common::errors::InvalidArgument(
+          "Input(X) of StftOp should be a tensor with shape [N, T], "
+          "but got rank %s.",
+          x_rank));
+
+  PADDLE_ENFORCE_GT(
+      hop_length,
+      0,
+      common::errors::InvalidArgument(
+          "Attribute(hop_length) should be greater than 0, but got %s.",
+          hop_length));
+
+  infer_context->AddEqualCstr(window_shape[0], symbol::DimExpr{n_fft});
+  const symbol::DimExpr seq_length = x_shape[x_rank - 1];
+  const symbol::DimExpr n_frames =
+      (symbol::DimExpr{1}) +
+      (seq_length - symbol::DimExpr{n_fft}) / symbol::DimExpr{hop_length};
+
+  if (seq_length.isa<int64_t>()) {
+    PADDLE_ENFORCE_LE(n_fft,
+                      seq_length.Get<std::int64_t>(),
+                      common::errors::InvalidArgument(
+                          "Attribute(frame_length) should be less equal than "
+                          "sequence length, but got (%s) > (%s).",
+                          n_fft,
+                          seq_length.Get<std::int64_t>()));
+  }
+
+  std::vector<symbol::DimExpr> output_shape;
+  output_shape.push_back(x_shape[0]);
+  if (onesided) {
+    output_shape.push_back(symbol::DimExpr{n_fft / 2 + 1});
+  } else {
+    output_shape.push_back(symbol::DimExpr{n_fft});
+  }
+  output_shape.push_back(n_frames);
+
+  infer_context->SetShapeOrDataForValue(
+      op->result(0),
+      symbol::ShapeOrDataDimExprs{
+          symbol::TensorShapeOrDataDimExprs(output_shape)});
+  return true;
+}
 
 bool SwigluOpInferSymbolicShape(pir::Operation *op,
                                 pir::InferSymbolicShapeContext *infer_context) {

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

@@ -85,7 +85,7 @@ OP_DECLARE_INFER_SYMBOLIC_SHAPE(SegmentPool)
 OP_DECLARE_INFER_SYMBOLIC_SHAPE(SequenceMask)
 // OP_DECLARE_INFER_SYMBOLIC_SHAPE(ShuffleBatch)
 // OP_DECLARE_INFER_SYMBOLIC_SHAPE(Solve)
-// OP_DECLARE_INFER_SYMBOLIC_SHAPE(Stft)
+OP_DECLARE_INFER_SYMBOLIC_SHAPE(Stft)
 OP_DECLARE_INFER_SYMBOLIC_SHAPE(Swiglu)
 OP_DECLARE_INFER_SYMBOLIC_SHAPE(TakeAlongAxis)
 OP_DECLARE_INFER_SYMBOLIC_SHAPE(TopPSampling)

paddle/phi/ops/yaml/ops.yaml

@@ -4592,7 +4592,7 @@
     func: stft
     data_type: x
   backward: stft_grad
-  # interfaces : paddle::dialect::InferSymbolicShapeInterface
+  interfaces : paddle::dialect::InferSymbolicShapeInterface
 
 - op : strided_slice
   args : (Tensor x, int[] axes, IntArray starts, IntArray ends, IntArray strides)