【Infer Symbolic Shape】【BUAA】Add trans_layout op (#67665)

Guanhuachen2003 · web-flow · commit 4af843e20874 · 2024-08-26T11:11:24.000+08:00
* Update unary_infer_sym.cc

* Update unary_infer_sym.h

* Update ops.yaml

* Update unary_infer_sym.cc

* 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
@@ -2884,11 +2884,7 @@ bool TransposeOpInferSymbolicShape(
     std::vector<symbol::DimExpr> dims;
     const auto &x_shape_or_data =
         infer_context->GetShapeOrDataForValue(op->operand_source(0));
-    if (x_shape_or_data.data().has_value()) {
-      dims = x_shape_or_data.data().value();
-    } else {
-      dims = x_shape_or_data.shape();
-    }
+    dims = x_shape_or_data.shape();
     return dims;
   }();
 
@@ -2930,6 +2926,10 @@ bool Transpose_OpInferSymbolicShape(
   return TransposeOpInferSymbolicShape(op, infer_context);
 }
 
+bool TransLayoutOpInferSymbolicShape(
+    pir::Operation *op, pir::InferSymbolicShapeContext *infer_context) {
+  return TransposeOpInferSymbolicShape(op, infer_context);
+}
 bool SqueezeOpInferSymbolicShape(
     pir::Operation *op, pir::InferSymbolicShapeContext *infer_context) {
   PADDLE_ENFORCE_EQ(
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
@@ -141,6 +141,7 @@ OP_DECLARE_INFER_SYMBOLIC_SHAPE(Topk)
 OP_DECLARE_INFER_SYMBOLIC_SHAPE(TopkV1)
 OP_DECLARE_INFER_SYMBOLIC_SHAPE(Transpose)
 OP_DECLARE_INFER_SYMBOLIC_SHAPE(Transpose_)
+OP_DECLARE_INFER_SYMBOLIC_SHAPE(TransLayout)
 OP_DECLARE_INFER_SYMBOLIC_SHAPE(Unbind)
 OP_DECLARE_INFER_SYMBOLIC_SHAPE(UniformInplace)
 OP_DECLARE_INFER_SYMBOLIC_SHAPE(UniformInplace_)
diff --git a/paddle/phi/ops/yaml/ops.yaml b/paddle/phi/ops/yaml/ops.yaml
@@ -4771,7 +4771,7 @@
   kernel :
     func : transpose
   backward : trans_layout_grad
-  # interfaces : paddle::dialect::InferSymbolicShapeInterface
+  interfaces : paddle::dialect::InferSymbolicShapeInterface
 
 - op : transpose
   args : (Tensor x, int[] perm)
