[CINN] Add group_norm_grad symbol infer (#71031)

ooooo-create · web-flow · commit 065cf6b9ac14 · 2025-02-13T11:02:22.000+08:00
* Add group_norm_grad symbol infer

* add GroupNormGrad_
diff --git a/paddle/fluid/pir/dialect/operator/interface/infer_symbolic_shape/backward_infer_sym.cc b/paddle/fluid/pir/dialect/operator/interface/infer_symbolic_shape/backward_infer_sym.cc
@@ -17,6 +17,16 @@
 
 namespace paddle::dialect {
 
+void SameShapeInfer(pir::InferSymbolicShapeContext *infer_context,
+                    pir::Value &&dst,
+                    pir::Value &&src) {
+  auto src_shape = infer_context->GetShapeOrDataForValue(src).shape();
+  infer_context->SetShapeOrDataForValue(
+      dst,
+      symbol::ShapeOrDataDimExprs{
+          symbol::TensorShapeOrDataDimExprs(src_shape)});
+}
+
 bool FusedAttentionGradOpInferSymbolicShape(
     pir::Operation *op, pir::InferSymbolicShapeContext *infer_context) {
   bool is_test = op->attribute<pir::BoolAttribute>("is_test").data();
@@ -26,71 +36,81 @@ bool FusedAttentionGradOpInferSymbolicShape(
                         "GradOp is only callable when is_test is false"));
   bool pre_layer_norm =
       op->attribute<pir::BoolAttribute>("pre_layer_norm").data();
-  auto same_shape_infer = [&](pir::Value &&dst, pir::Value &&src) {
-    auto src_shape = infer_context->GetShapeOrDataForValue(src).shape();
-    infer_context->SetShapeOrDataForValue(
-        dst,
-        symbol::ShapeOrDataDimExprs{
-            symbol::TensorShapeOrDataDimExprs(src_shape)});
-  };
   if (!pre_layer_norm) {
     if (!paddle::dialect::details::IsFakeValue(op->result(6)) &&
         op->operand_source(11)) {
-      same_shape_infer(op->result(6), op->operand_source(11));
+      SameShapeInfer(infer_context, op->result(6), op->operand_source(11));
     }
     if (!paddle::dialect::details::IsFakeValue(op->result(7)) &&
         op->operand_source(12)) {
-      same_shape_infer(op->result(7), op->operand_source(12));
+      SameShapeInfer(infer_context, op->result(7), op->operand_source(12));
     }
   }
   if (pre_layer_norm && op->operand_source(9)) {
     if (!paddle::dialect::details::IsFakeValue(op->result(4))) {
-      same_shape_infer(op->result(4), op->operand_source(9));
+      SameShapeInfer(infer_context, op->result(4), op->operand_source(9));
     }
     if (!paddle::dialect::details::IsFakeValue(op->result(5)) &&
         op->operand_source(10)) {
-      same_shape_infer(op->result(5), op->operand_source(10));
+      SameShapeInfer(infer_context, op->result(5), op->operand_source(10));
     }
   }
-  same_shape_infer(op->result(8), op->operand_source(1));
+  SameShapeInfer(infer_context, op->result(8), op->operand_source(1));
   if (!paddle::dialect::details::IsFakeValue(op->result(3)) &&
       op->operand_source(13)) {
-    same_shape_infer(op->result(3), op->operand_source(8));
+    SameShapeInfer(infer_context, op->result(3), op->operand_source(8));
   }
   if (!paddle::dialect::details::IsFakeValue(op->result(10))) {
-    same_shape_infer(op->result(10), op->operand_source(7));
+    SameShapeInfer(infer_context, op->result(10), op->operand_source(7));
   }
-  same_shape_infer(op->result(9), op->operand_source(2));
+  SameShapeInfer(infer_context, op->result(9), op->operand_source(2));
   if (!paddle::dialect::details::IsFakeValue(op->result(0))) {
-    same_shape_infer(op->result(0), op->operand_source(3));
+    SameShapeInfer(infer_context, op->result(0), op->operand_source(3));
   }
   if (pre_layer_norm) {
     if (!paddle::dialect::details::IsFakeValue(op->result(11)) &&
         op->operand_source(13)) {
-      same_shape_infer(op->result(11), op->operand_source(13));
+      SameShapeInfer(infer_context, op->result(11), op->operand_source(13));
     }
   } else {
     if (!paddle::dialect::details::IsFakeValue(op->result(12)) &&
         op->operand_source(18)) {
-      same_shape_infer(op->result(12), op->operand_source(18));
+      SameShapeInfer(infer_context, op->result(12), op->operand_source(18));
     }
   }
-  same_shape_infer(op->result(19), op->operand_source(26));
-  same_shape_infer(op->result(14), op->operand_source(22));
-  same_shape_infer(op->result(15), op->operand_source(20));
-  same_shape_infer(op->result(16), op->operand_source(21));
-  same_shape_infer(op->result(17), op->operand_source(23));
-  same_shape_infer(op->result(18), op->operand_source(25));
+  SameShapeInfer(infer_context, op->result(19), op->operand_source(26));
+  SameShapeInfer(infer_context, op->result(14), op->operand_source(22));
+  SameShapeInfer(infer_context, op->result(15), op->operand_source(20));
+  SameShapeInfer(infer_context, op->result(16), op->operand_source(21));
+  SameShapeInfer(infer_context, op->result(17), op->operand_source(23));
+  SameShapeInfer(infer_context, op->result(18), op->operand_source(25));
   if (!paddle::dialect::details::IsFakeValue(op->result(2)) &&
       op->operand_source(6)) {
-    same_shape_infer(op->result(2), op->operand_source(6));
+    SameShapeInfer(infer_context, op->result(2), op->operand_source(6));
   }
-  same_shape_infer(op->result(13), op->operand_source(19));
+  SameShapeInfer(infer_context, op->result(13), op->operand_source(19));
   if (!paddle::dialect::details::IsFakeValue(op->result(1)) &&
       op->operand_source(4)) {
-    same_shape_infer(op->result(1), op->operand_source(4));
+    SameShapeInfer(infer_context, op->result(1), op->operand_source(4));
   }
-  same_shape_infer(op->result(20), op->operand_source(27));
+  SameShapeInfer(infer_context, op->result(20), op->operand_source(27));
   return true;
 }
+
+bool GroupNormGradOpInferSymbolicShape(
+    pir::Operation *op, pir::InferSymbolicShapeContext *infer_context) {
+  SameShapeInfer(infer_context, op->result(0), op->operand_source(3));
+  if (!paddle::dialect::details::IsFakeValue(op->operand_source(1))) {
+    SameShapeInfer(infer_context, op->result(1), op->operand_source(1));
+  }
+  if (!paddle::dialect::details::IsFakeValue(op->operand_source(2))) {
+    SameShapeInfer(infer_context, op->result(2), op->operand_source(2));
+  }
+  return true;
+}
+
+bool GroupNormGrad_OpInferSymbolicShape(
+    pir::Operation *op, pir::InferSymbolicShapeContext *infer_context) {
+  return GroupNormOpInferSymbolicShape(op, infer_context);
+}
 }  // namespace paddle::dialect
diff --git a/paddle/fluid/pir/dialect/operator/interface/infer_symbolic_shape/backward_infer_sym.h b/paddle/fluid/pir/dialect/operator/interface/infer_symbolic_shape/backward_infer_sym.h
@@ -18,4 +18,6 @@
 
 namespace paddle::dialect {
 OP_DECLARE_INFER_SYMBOLIC_SHAPE(FusedAttentionGrad)
+OP_DECLARE_INFER_SYMBOLIC_SHAPE(GroupNormGrad)
+OP_DECLARE_INFER_SYMBOLIC_SHAPE(GroupNormGrad_)
 }  // namespace paddle::dialect
diff --git a/paddle/phi/ops/yaml/backward.yaml b/paddle/phi/ops/yaml/backward.yaml
@@ -1414,6 +1414,7 @@
   composite : group_norm_grad(x, scale, bias, y, mean, variance, y_grad, epsilon, groups, data_format, x_grad, scale_grad, bias_grad)
   optional: scale, bias
   inplace : (y_grad -> x_grad)
+  interfaces : paddle::dialect::InferSymbolicShapeInterface
 
 - backward_op : gru_grad
   forward: gru (Tensor input, Tensor h0, Tensor weight, Tensor bias, str activation = "tanh",
