【BUAA】【Infer Symbolic Shape】add set_value_with_tensor, sigmoid_cross_entropy_with_logits,swiglu (#67098)

uanu2002 · web-flow · commit 245e28922c5f · 2024-08-12T18:16:26.000+08:00
diff --git a/paddle/fluid/pir/dialect/operator/interface/infer_symbolic_shape/binary_infer_sym.cc b/paddle/fluid/pir/dialect/operator/interface/infer_symbolic_shape/binary_infer_sym.cc
@@ -914,12 +914,30 @@ bool SearchsortedOpInferSymbolicShape(
 //   return true;
 // }
 
-// bool SwigluOpInferSymbolicShape(pir::Operation *op,
-//                                 pir::InferSymbolicShapeContext
-//                                 *infer_context) {
-//   // pass
-//   return true;
-// }
+bool SwigluOpInferSymbolicShape(pir::Operation *op,
+                                pir::InferSymbolicShapeContext *infer_context) {
+  const auto &x_shape_or_data =
+      infer_context->GetShapeOrDataForValue(op->operand_source(0));
+  size_t rank = x_shape_or_data.shape().size();
+  if (op->operand_source(1)) {
+    const auto &y_shape_or_data =
+        infer_context->GetShapeOrDataForValue(op->operand_source(1));
+    for (size_t i = 0; i < rank; ++i) {
+      infer_context->AddEqualCstr(x_shape_or_data.shape()[i],
+                                  y_shape_or_data.shape()[i]);
+    }
+    infer_context->SetShapeOrDataForValue(op->result(0), x_shape_or_data);
+  } else {
+    std::vector<symbol::DimExpr> x_shape = x_shape_or_data.shape();
+    // TODO(CINN): Add distribute constraint
+    x_shape[rank - 1] = x_shape[rank - 1] / symbol::DimExpr{2};
+    infer_context->SetShapeOrDataForValue(
+        op->result(0),
+        symbol::ShapeOrDataDimExprs{
+            symbol::TensorShapeOrDataDimExprs(x_shape)});
+  }
+  return true;
+}
 
 bool IscloseOpInferSymbolicShape(
     pir::Operation *op, pir::InferSymbolicShapeContext *infer_context) {
diff --git a/paddle/fluid/pir/dialect/operator/interface/infer_symbolic_shape/binary_infer_sym.h b/paddle/fluid/pir/dialect/operator/interface/infer_symbolic_shape/binary_infer_sym.h
@@ -62,7 +62,7 @@ OP_DECLARE_INFER_SYMBOLIC_SHAPE(Mv)
 OP_DECLARE_INFER_SYMBOLIC_SHAPE(ReduceAs)
 OP_DECLARE_INFER_SYMBOLIC_SHAPE(Searchsorted)
 // OP_DECLARE_INFER_SYMBOLIC_SHAPE(SequenceMask)
-// OP_DECLARE_INFER_SYMBOLIC_SHAPE(Swiglu)
+OP_DECLARE_INFER_SYMBOLIC_SHAPE(Swiglu)
 OP_DECLARE_INFER_SYMBOLIC_SHAPE(TakeAlongAxis)
 OP_DECLARE_INFER_SYMBOLIC_SHAPE(TopPSampling)
 // OP_DECLARE_INFER_SYMBOLIC_SHAPE(TdmChild)
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
@@ -1507,19 +1507,46 @@ bool StackOpInferSymbolicShape(pir::Operation *op,
 //   return true;
 // }
 
-// bool SigmoidCrossEntropyWithLogitsOpInferSymbolicShape(pir::Operation *op,
-//                                                        pir::InferSymbolicShapeContext
-//                                                        *infer_context) {
-//   // pass
-//   return true;
-// }
+bool SigmoidCrossEntropyWithLogitsOpInferSymbolicShape(
+    pir::Operation *op, pir::InferSymbolicShapeContext *infer_context) {
+  const auto &input_shape_or_data =
+      infer_context->GetShapeOrDataForValue(op->operand_source(0));
+  const auto &label_shape_or_data =
+      infer_context->GetShapeOrDataForValue(op->operand_source(1));
 
-// bool SigmoidCrossEntropyWithLogits_OpInferSymbolicShape(pir::Operation *op,
-//                                                         pir::InferSymbolicShapeContext
-//                                                         *infer_context) {
-//   return SigmoidCrossEntropyWithLogitsOpInferSymbolicShape(op,
-//   infer_context);
-// }
+  size_t rank = input_shape_or_data.shape().size();
+  PADDLE_ENFORCE_EQ(rank,
+                    label_shape_or_data.shape().size(),
+                    common::errors::InvalidArgument(
+                        "Input(X) and Input(Label) shall have the same rank."
+                        "But received: the rank of Input(X) is [%d], "
+                        "the rank of Input(Label) is [%d].",
+                        rank,
+                        label_shape_or_data.shape().size()));
+
+  for (size_t i = 0; i < rank; ++i) {
+    infer_context->AddEqualCstr(input_shape_or_data.shape()[i],
+                                label_shape_or_data.shape()[i]);
+  }
+  if (op->operand_source(2)) {
+    const auto &pos_shape_or_data =
+        infer_context->GetShapeOrDataForValue(op->operand_source(2));
+    for (size_t i = 0; i < rank; ++i) {
+      infer_context->AddEqualCstr(input_shape_or_data.shape()[i],
+                                  pos_shape_or_data.shape()[i]);
+    }
+  }
+  infer_context->SetShapeOrDataForValue(
+      op->result(0),
+      symbol::ShapeOrDataDimExprs{
+          symbol::TensorShapeOrDataDimExprs(input_shape_or_data.shape())});
+  return true;
+}
+
+bool SigmoidCrossEntropyWithLogits_OpInferSymbolicShape(
+    pir::Operation *op, pir::InferSymbolicShapeContext *infer_context) {
+  return SigmoidCrossEntropyWithLogitsOpInferSymbolicShape(op, infer_context);
+}
 
 // bool SyncBatchNormOpInferSymbolicShape(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,8 +78,8 @@ OP_DECLARE_INFER_SYMBOLIC_SHAPE(NearestInterp)
 // OP_DECLARE_INFER_SYMBOLIC_SHAPE(RoiPool)
 OP_DECLARE_INFER_SYMBOLIC_SHAPE(RoiAlign)
 OP_DECLARE_INFER_SYMBOLIC_SHAPE(Stack)
-// OP_DECLARE_INFER_SYMBOLIC_SHAPE(SigmoidCrossEntropyWithLogits)
-// OP_DECLARE_INFER_SYMBOLIC_SHAPE(SigmoidCrossEntropyWithLogits_)
+OP_DECLARE_INFER_SYMBOLIC_SHAPE(SigmoidCrossEntropyWithLogits)
+OP_DECLARE_INFER_SYMBOLIC_SHAPE(SigmoidCrossEntropyWithLogits_)
 // OP_DECLARE_INFER_SYMBOLIC_SHAPE(SyncBatchNorm)
 // OP_DECLARE_INFER_SYMBOLIC_SHAPE(SyncBatchNorm_)
 // OP_DECLARE_INFER_SYMBOLIC_SHAPE(SaveCombine)
diff --git a/paddle/fluid/pir/dialect/operator/interface/infer_symbolic_shape/unary_infer_sym.cc b/paddle/fluid/pir/dialect/operator/interface/infer_symbolic_shape/unary_infer_sym.cc
@@ -1908,12 +1908,36 @@ bool SumOpInferSymbolicShape(pir::Operation *op,
 
   return true;
 }
-// bool SetValueWithTensorOpInferSymbolicShape(pir::Operation *op,
-//                                             pir::InferSymbolicShapeContext
-//                                             *infer_context) {
-//   // pass
-//   return true;
-// }
+
+bool SetValueOpInferSymbolicShape(
+    pir::Operation *op, pir::InferSymbolicShapeContext *infer_context) {
+  const auto &input_shape_or_data =
+      infer_context->GetShapeOrDataForValue(op->operand_source(0));
+  PADDLE_ENFORCE_LT(
+      input_shape_or_data.shape().size(),
+      7,
+      common::errors::InvalidArgument("The SetValueOp's rank of input should "
+                                      "be less than 7, but received %d.",
+                                      input_shape_or_data.shape().size()));
+
+  infer_context->SetShapeOrDataForValue(op->result(0), input_shape_or_data);
+  return true;
+}
+
+bool SetValue_OpInferSymbolicShape(
+    pir::Operation *op, pir::InferSymbolicShapeContext *infer_context) {
+  return SetValueOpInferSymbolicShape(op, infer_context);
+}
+
+bool SetValueWithTensorOpInferSymbolicShape(
+    pir::Operation *op, pir::InferSymbolicShapeContext *infer_context) {
+  return SetValueOpInferSymbolicShape(op, infer_context);
+}
+
+bool SetValueWithTensor_OpInferSymbolicShape(
+    pir::Operation *op, pir::InferSymbolicShapeContext *infer_context) {
+  return SetValueOpInferSymbolicShape(op, infer_context);
+}
 
 // bool TraceOpInferSymbolicShape(pir::Operation *op,
 //                                pir::InferSymbolicShapeContext *infer_context)
diff --git a/paddle/fluid/pir/dialect/operator/interface/infer_symbolic_shape/unary_infer_sym.h b/paddle/fluid/pir/dialect/operator/interface/infer_symbolic_shape/unary_infer_sym.h
@@ -106,7 +106,10 @@ OP_DECLARE_INFER_SYMBOLIC_SHAPE(SplitWithNum)
 OP_DECLARE_INFER_SYMBOLIC_SHAPE(Squeeze)
 OP_DECLARE_INFER_SYMBOLIC_SHAPE(Squeeze_)
 OP_DECLARE_INFER_SYMBOLIC_SHAPE(Sum)
-// OP_DECLARE_INFER_SYMBOLIC_SHAPE(SetValueWithTensor)
+OP_DECLARE_INFER_SYMBOLIC_SHAPE(SetValue)
+OP_DECLARE_INFER_SYMBOLIC_SHAPE(SetValue_)
+OP_DECLARE_INFER_SYMBOLIC_SHAPE(SetValueWithTensor)
+OP_DECLARE_INFER_SYMBOLIC_SHAPE(SetValueWithTensor_)
 // OP_DECLARE_INFER_SYMBOLIC_SHAPE(Trace)
 OP_DECLARE_INFER_SYMBOLIC_SHAPE(Tile)
 OP_DECLARE_INFER_SYMBOLIC_SHAPE(Topk)
diff --git a/paddle/phi/ops/yaml/ops.yaml b/paddle/phi/ops/yaml/ops.yaml
@@ -4134,6 +4134,7 @@
   kernel:
     func: set_value_with_tensor
   backward: set_value_with_tensor_grad
+  interfaces : paddle::dialect::InferSymbolicShapeInterface
 
 - op : sgd_
   args : (Tensor param, Tensor learning_rate, Tensor grad, Tensor master_param, bool multi_precision=false)
@@ -4225,6 +4226,7 @@
   inplace : (x -> out)
   backward : sigmoid_cross_entropy_with_logits_grad
   optional : pos_weight
+  interfaces : paddle::dialect::InferSymbolicShapeInterface
 
 - op : sign
   args : (Tensor x)
@@ -4507,6 +4509,7 @@
      func : swiglu
   optional : y
   backward: swiglu_grad
+  interfaces : paddle::dialect::InferSymbolicShapeInterface
 
 - op : swish
   args : (Tensor x)
