【Infer Symbolic Shape No.15】Add LruOpInferSymbolicShape implementation (#68096)

enkilee · web-flow · commit 79b4a869f548 · 2024-09-18T15:10:11.000+08:00
* fix

* fix
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
@@ -1615,6 +1615,38 @@ bool LogsumexpOpInferSymbolicShape(
   return details::ReduceInferDim(op, infer_context, axis, keepdim, reduce_all);
 }
 
+bool LrnOpInferSymbolicShape(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> &x_shape = x_shape_or_data.shape();
+  int x_size = x_shape.size();
+  PADDLE_ENFORCE_EQ(
+      x_size,
+      4,
+      common::errors::InvalidArgument("Input(input) rank should be 4, "
+                                      "but received input rank (%d) != 4",
+                                      x_size));
+  int n_value = op->attribute<pir::Int32Attribute>("n").data();
+  PADDLE_ENFORCE_GT(
+      n_value,
+      0UL,
+      common::errors::InvalidArgument("Argument(n) should be positive, "
+                                      "but received n(%d) not greater than 0",
+                                      n_value));
+  PADDLE_ENFORCE_EQ(
+      n_value % 2,
+      1UL,
+      common::errors::InvalidArgument("Argument(n) should be odd value, "
+                                      "but received n(%d) is not an odd value",
+                                      n_value));
+  infer_context->SetShapeOrDataForValue(
+      op->result(0), symbol::TensorShapeOrDataDimExprs(x_shape));
+  infer_context->SetShapeOrDataForValue(
+      op->result(1), symbol::TensorShapeOrDataDimExprs(x_shape));
+  return true;
+}
+
 bool LuOpInferSymbolicShape(pir::Operation *op,
                             pir::InferSymbolicShapeContext *infer_context) {
   const auto &x_shape_or_data =
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
@@ -85,6 +85,7 @@ OP_DECLARE_INFER_SYMBOLIC_SHAPE(L1Norm_)
 OP_DECLARE_INFER_SYMBOLIC_SHAPE(LpPool2d)
 OP_DECLARE_INFER_SYMBOLIC_SHAPE(Logcumsumexp)
 OP_DECLARE_INFER_SYMBOLIC_SHAPE(Logsumexp)
+OP_DECLARE_INFER_SYMBOLIC_SHAPE(Lrn)
 OP_DECLARE_INFER_SYMBOLIC_SHAPE(Lu)
 OP_DECLARE_INFER_SYMBOLIC_SHAPE(Lu_)
 OP_DECLARE_INFER_SYMBOLIC_SHAPE(Mode)
diff --git a/paddle/phi/ops/yaml/inconsistent/static_ops.yaml b/paddle/phi/ops/yaml/inconsistent/static_ops.yaml
@@ -527,6 +527,7 @@
     func: lrn
     data_type: x
   backward: lrn_grad
+  interfaces : paddle::dialect::InferSymbolicShapeInterface
 
 - op : matmul
   args : (Tensor x, Tensor y, bool transpose_x = false, bool transpose_y = false)
