【Infer Symbolic Shape No.29/No.31】Add Eigh/EigvalshOpInferSymbolicShape implementation (#66582)

AndPuQing · web-flow · commit 9eef74f72ab0 · 2024-07-30T15:04:16.000+08:00
* add EighOpInferSymbolicShape implementation

* add eigvalsh
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
@@ -339,6 +339,26 @@ bool DistributeFpnProposalsOpInferSymbolicShape(
   return true;
 }
 
+bool EighOpInferSymbolicShape(pir::Operation *op,
+                              pir::InferSymbolicShapeContext *infer_context) {
+  const auto &x_shape =
+      infer_context->GetShapeOrDataForValue(op->operand_source(0)).shape();
+  std::vector<symbol::DimExpr> out_shape;
+  for (size_t i = 0; i < x_shape.size() - 1; ++i) {
+    out_shape.push_back(x_shape.at(i));
+  }
+  infer_context->SetShapeOrDataForValue(
+      op->result(0), symbol::TensorShapeOrDataDimExprs(out_shape));
+  infer_context->SetShapeOrDataForValue(
+      op->result(1), symbol::TensorShapeOrDataDimExprs(x_shape));
+  return true;
+}
+
+bool EigvalshOpInferSymbolicShape(
+    pir::Operation *op, pir::InferSymbolicShapeContext *infer_context) {
+  return EighOpInferSymbolicShape(op, infer_context);
+}
+
 bool FftC2cOpInferSymbolicShape(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
@@ -34,6 +34,8 @@ OP_DECLARE_INFER_SYMBOLIC_SHAPE(Cumsum_)
 OP_DECLARE_INFER_SYMBOLIC_SHAPE(DiagEmbed)
 OP_DECLARE_INFER_SYMBOLIC_SHAPE(Diagonal)
 OP_DECLARE_INFER_SYMBOLIC_SHAPE(DistributeFpnProposals)
+OP_DECLARE_INFER_SYMBOLIC_SHAPE(Eigh)
+OP_DECLARE_INFER_SYMBOLIC_SHAPE(Eigvalsh)
 OP_DECLARE_INFER_SYMBOLIC_SHAPE(FftC2c)
 OP_DECLARE_INFER_SYMBOLIC_SHAPE(FftC2r)
 OP_DECLARE_INFER_SYMBOLIC_SHAPE(FftR2c)
diff --git a/paddle/phi/ops/yaml/ops.yaml b/paddle/phi/ops/yaml/ops.yaml
@@ -1463,6 +1463,7 @@
   kernel :
     func : eigh
   backward : eigh_grad
+  interfaces : paddle::dialect::InferSymbolicShapeInterface
 
 - op : eigvals
   args : (Tensor x)
@@ -1481,6 +1482,7 @@
     func : eigvalsh
     data_type : x
   backward : eigvalsh_grad
+  interfaces : paddle::dialect::InferSymbolicShapeInterface
 
 - op : elu
   args : (Tensor x, float alpha = 1.0f)
diff --git a/test/legacy_test/test_eigh_op.py b/test/legacy_test/test_eigh_op.py
@@ -89,8 +89,12 @@ def init_input(self):
         self.x_type = np.float64
         self.x_np = np.random.random(self.x_shape).astype(self.x_type)
 
-    # def test_check_output(self):
-    #     self.check_output(no_check_set=['Eigenvectors'])
+    def test_check_output(self):
+        self.check_output(
+            no_check_set=['Eigenvectors'],
+            check_symbol_infer=True,
+            check_pir=True,
+        )
 
     def test_grad(self):
         self.check_grad(["X"], ["Eigenvalues"], check_pir=True)
@@ -100,6 +104,11 @@ class TestEighUPLOCase(TestEighOp):
     def init_config(self):
         self.UPLO = 'U'
 
+    def init_input(self):
+        self.x_shape = (20, 10, 10)
+        self.x_type = np.float64
+        self.x_np = np.random.random(self.x_shape).astype(self.x_type)
+
 
 class TestEighGPUCase(unittest.TestCase):
     def setUp(self):
diff --git a/test/legacy_test/test_eigvalsh_op.py b/test/legacy_test/test_eigvalsh_op.py
@@ -73,7 +73,11 @@ def init_input(self):
 
     def test_check_output(self):
         # Vectors in posetive or negative is equivalent
-        self.check_output(no_check_set=['Eigenvectors'], check_pir=True)
+        self.check_output(
+            no_check_set=['Eigenvectors'],
+            check_pir=True,
+            check_symbol_infer=True,
+        )
 
     def test_grad(self):
         self.check_grad(["X"], ["Eigenvalues"], check_pir=True)
