【Infer Symbolic Shape BUAA No.64】Add inverse op (#66633)

crazyxiaoxi · web-flow · commit 0a8c6b7c3302 · 2024-08-07T10:23:24.000+08:00
* inverse

* fix some problem

* testop

* testop

* cstr

* rerun

* fix conflict

* delet for
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
@@ -1006,6 +1006,25 @@ bool KthvalueOpInferSymbolicShape(
   return true;
 }
 
+bool InverseOpInferSymbolicShape(
+    pir::Operation *op, pir::InferSymbolicShapeContext *infer_context) {
+  const auto &input_shape =
+      infer_context->GetShapeOrDataForValue(op->operand_source(0));
+  std::vector<symbol::DimExpr> input_dims = input_shape.shape();
+  int input_rank = input_dims.size();
+
+  infer_context->AddEqualCstr(input_dims[input_rank - 2],
+                              input_dims[input_rank - 1]);
+
+  std::vector<symbol::DimExpr> output_dims = input_dims;
+
+  infer_context->SetShapeOrDataForValue(
+      op->result(0),
+      symbol::ShapeOrDataDimExprs{
+          symbol::TensorShapeOrDataDimExprs(output_dims)});
+  return true;
+}
+
 bool LpPool2dOpInferSymbolicShape(
     pir::Operation *op, pir::InferSymbolicShapeContext *infer_context) {
   const auto &kernel_size = [&]() -> std::vector<symbol::DimExpr> {
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
@@ -62,6 +62,7 @@ OP_DECLARE_INFER_SYMBOLIC_SHAPE(FillDiagonal)
 OP_DECLARE_INFER_SYMBOLIC_SHAPE(FillDiagonal_)
 OP_DECLARE_INFER_SYMBOLIC_SHAPE(Flatten)
 OP_DECLARE_INFER_SYMBOLIC_SHAPE(Flatten_)
+OP_DECLARE_INFER_SYMBOLIC_SHAPE(Inverse)
 // OP_DECLARE_INFER_SYMBOLIC_SHAPE(GumbelSoftmax)
 OP_DECLARE_INFER_SYMBOLIC_SHAPE(IdentityLoss)
 OP_DECLARE_INFER_SYMBOLIC_SHAPE(IdentityLoss_)
diff --git a/paddle/phi/ops/yaml/ops.yaml b/paddle/phi/ops/yaml/ops.yaml
@@ -2548,6 +2548,7 @@
   kernel :
     func : inverse
   backward : inverse_grad
+  interfaces : paddle::dialect::InferSymbolicShapeInterface
 
 - op : is_empty
   args : (Tensor x)
