[Infer Symbolic Shape No.225] Add symbol_infer_interface for shuffle_batch (#68978)

* Add symbol_infer_interface for shuffle_batch

* Fix

* Refine logic

* apply review

* revert irrelevant change

Author: ooooo-create

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

@@ -1960,12 +1960,30 @@ bool SequenceMaskOpInferSymbolicShape(
   return true;
 }
 
-// bool ShuffleBatchOpInferSymbolicShape(pir::Operation *op,
-//                                       pir::InferSymbolicShapeContext
-//                                       *infer_context) {
-//   // pass
-//   return true;
-// }
+bool ShuffleBatchOpInferSymbolicShape(
+    pir::Operation *op, pir::InferSymbolicShapeContext *infer_context) {
+  const auto &x_shape =
+      infer_context->GetShapeOrDataForValue(op->operand_source(0)).shape();
+  infer_context->SetShapeOrDataForValue(
+      op->result(0),
+      symbol::ShapeOrDataDimExprs{symbol::TensorShapeOrDataDimExprs(x_shape)});
+  infer_context->SetShapeOrDataForValue(
+      op->result(2),
+      symbol::ShapeOrDataDimExprs{symbol::TensorShapeOrDataDimExprs({1})});
+  const symbol::DimExpr shuffleidx = [&] {
+    symbol::DimExpr shuffleidx{1};
+    for (const auto &i : x_shape) {
+      shuffleidx = shuffleidx * i;
+    }
+    return shuffleidx;
+  }();
+
+  infer_context->SetShapeOrDataForValue(
+      op->result(1),
+      symbol::ShapeOrDataDimExprs{
+          symbol::TensorShapeOrDataDimExprs({shuffleidx})});
+  return true;
+}
 
 bool StftOpInferSymbolicShape(pir::Operation *op,
                               pir::InferSymbolicShapeContext *infer_context) {

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

@@ -88,7 +88,7 @@ OP_DECLARE_INFER_SYMBOLIC_SHAPE(ReduceAs)
 OP_DECLARE_INFER_SYMBOLIC_SHAPE(Searchsorted)
 OP_DECLARE_INFER_SYMBOLIC_SHAPE(SegmentPool)
 OP_DECLARE_INFER_SYMBOLIC_SHAPE(SequenceMask)
-// OP_DECLARE_INFER_SYMBOLIC_SHAPE(ShuffleBatch)
+OP_DECLARE_INFER_SYMBOLIC_SHAPE(ShuffleBatch)
 OP_DECLARE_INFER_SYMBOLIC_SHAPE(Solve)
 OP_DECLARE_INFER_SYMBOLIC_SHAPE(Stft)
 OP_DECLARE_INFER_SYMBOLIC_SHAPE(Swiglu)

paddle/phi/ops/yaml/ops.yaml

@@ -4508,7 +4508,7 @@
   traits : pir::SideEffectTrait
   data_transform :
     skip_transform : seed
-  # interfaces : paddle::dialect::InferSymbolicShapeInterface
+  interfaces : paddle::dialect::InferSymbolicShapeInterface
 
 - op : shuffle_channel
   args : (Tensor x, int group = 1)