【Infer Symbolic Shape BUAA No.12】Add channelshuffle op (#66946)

Guanhuachen2003 · web-flow · commit 7d25f8c16bc9 · 2024-08-05T13:01:09.000+08:00
* Update unary_infer_sym.cc

* Update unary_infer_sym.h

* Update ops.yaml

* Update unary_infer_sym.cc
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
@@ -488,6 +488,51 @@ bool Cumsum_OpInferSymbolicShape(
     pir::Operation *op, pir::InferSymbolicShapeContext *infer_context) {
   return CumsumOpInferSymbolicShape(op, infer_context);
 }
+bool ChannelShuffleOpInferSymbolicShape(
+    pir::Operation *op, pir::InferSymbolicShapeContext *infer_context) {
+  const auto &x_shape_or_data =
+      infer_context->GetShapeOrDataForValue(op->operand_source(0));
+  const std::vector<symbol::DimExpr> &input_dims = x_shape_or_data.shape();
+
+  int groups = op->attribute<pir::Int32Attribute>("groups").data();
+  std::string data_format =
+      op->attribute<pir::StrAttribute>("data_format").AsString();
+
+  PADDLE_ENFORCE_EQ(
+      input_dims.size(),
+      4,
+      phi::errors::InvalidArgument("Input should be a 4-D tensor of format [N, "
+                                   "C, H, W] or [N, H, W, C], but got %u.",
+                                   input_dims.size()));
+  PADDLE_ENFORCE_GE(
+      groups,
+      1,
+      phi::errors::InvalidArgument("groups should be larger than 0."));
+  PADDLE_ENFORCE_EQ(
+      data_format == "NCHW" || data_format == "NHWC",
+      true,
+      phi::errors::InvalidArgument("data_format must be one of NCHW and NHWC. "
+                                   "But received data_format: %s",
+                                   data_format));
+
+  const bool channel_last = (data_format == "NHWC");
+
+  symbol::DimExpr channels;
+  if (!channel_last) {
+    channels = input_dims[1];
+  } else {
+    channels = input_dims[3];
+  }
+
+  symbol::DimExpr groups_expr = symbol::DimExpr(groups);
+  symbol::DimExpr expected_channels = groups_expr * (channels / groups_expr);
+
+  infer_context->AddEqualCstr(channels, expected_channels);
+
+  infer_context->SetShapeOrDataForValue(op->result(0), x_shape_or_data);
+
+  return true;
+}
 
 bool DiagEmbedOpInferSymbolicShape(
     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
@@ -39,6 +39,7 @@ OP_DECLARE_INFER_SYMBOLIC_SHAPE(Cumprod)
 OP_DECLARE_INFER_SYMBOLIC_SHAPE(Cumprod_)
 OP_DECLARE_INFER_SYMBOLIC_SHAPE(Cumsum)
 OP_DECLARE_INFER_SYMBOLIC_SHAPE(Cumsum_)
+OP_DECLARE_INFER_SYMBOLIC_SHAPE(ChannelShuffle)
 OP_DECLARE_INFER_SYMBOLIC_SHAPE(DiagEmbed)
 OP_DECLARE_INFER_SYMBOLIC_SHAPE(Diagonal)
 OP_DECLARE_INFER_SYMBOLIC_SHAPE(DistributeFpnProposals)
diff --git a/paddle/phi/ops/yaml/ops.yaml b/paddle/phi/ops/yaml/ops.yaml
@@ -832,6 +832,7 @@
   kernel :
     func : channel_shuffle
   backward : channel_shuffle_grad
+  interfaces : paddle::dialect::InferSymbolicShapeInterface
 
 - op : check_finite_and_unscale_
   args : (Tensor[] x, Tensor scale)
