PaddlePaddle
diff --git a/‎paddle/fluid/pir/dialect/op_generator/op_gen.py
Copy file name to clipboard
+33-1Lines changed: 33 additions & 1 deletion b/‎paddle/fluid/pir/dialect/op_generator/op_gen.py
Copy file name to clipboard
+33-1Lines changed: 33 additions & 1 deletion
diff --git a/‎paddle/fluid/pir/dialect/op_generator/reify_infer_shape_gen.py
Copy file name to clipboard
+26Lines changed: 26 additions & 0 deletions b/‎paddle/fluid/pir/dialect/op_generator/reify_infer_shape_gen.py
Copy file name to clipboard
+26Lines changed: 26 additions & 0 deletions
diff --git a/‎paddle/fluid/pir/dialect/operator/interface/interface.cc
Copy file name to clipboard
-38Lines changed: 0 additions & 38 deletions b/‎paddle/fluid/pir/dialect/operator/interface/interface.cc
Copy file name to clipboard
-38Lines changed: 0 additions & 38 deletions
diff --git a/‎paddle/fluid/pir/dialect/operator/interface/parse_kernel_key.cc
Copy file name to clipboard
+33Lines changed: 33 additions & 0 deletions b/‎paddle/fluid/pir/dialect/operator/interface/parse_kernel_key.cc
Copy file name to clipboard
+33Lines changed: 33 additions & 0 deletions
diff --git a/‎paddle/fluid/pir/dialect/operator/interface/reify_infer_shape.cc
Copy file name to clipboard
+153Lines changed: 153 additions & 0 deletions b/‎paddle/fluid/pir/dialect/operator/interface/reify_infer_shape.cc
Copy file name to clipboard
+153Lines changed: 153 additions & 0 deletions
@@ -30,6 +30,7 @@
 from op_member_func_gen import gen_op_get_inputs_outputs_str
 from op_verify_gen import gen_verify_func_str
 from parse_kernel_key_gen import gen_parse_kernel_key_str
+from reify_infer_shape_gen import gen_reify_infer_shape_str
 from vjp_interface_black_list import vjp_interface_black_list
 
 # import from paddle/fluid/primitive/code_gen/gen.py
@@ -51,6 +52,7 @@
 
 #include "paddle/pir/core/builder.h"
 #include "paddle/pir/core/operation_utils.h"
+#include "paddle/fluid/pir/dialect/operator/interface/reify_infer_shape.h"
 #include "paddle/pir/core/op_base.h"
 #include "paddle/pir/core/op_trait.h"
 #include "paddle/fluid/pir/dialect/operator/utils/utils.h"
@@ -101,6 +103,7 @@ class {op_name} : public pir::Op<{op_name}{interfaces}{traits}> {{
   void VerifySig();
 {get_kernel_type_for_var_declare}
 {parse_kernel_key_declare}
+{reify_infer_shape_declare}
 {get_inputs_and_outputs}
 {exclusive_interface}
 }};
@@ -123,6 +126,12 @@ class {op_name} : public pir::Op<{op_name}{interfaces}{traits}> {{
   static std::tuple<phi::DataType, phi::Backend> ParseKernelKey(pir::Operation *op);
 """
 
+reify_infer_shape_template = """
+  static bool ReifyInferShape(pir::Builder &builder,
+                              const std::vector<pir::OpOperand> &operands,
+                              std::vector<pir::Value> &reified_return_shapes);
+"""
+
 # =====================================
 # String Template for cc file code gen
 # =====================================
@@ -410,6 +419,15 @@ def __init__(self, op_yaml_item, op_compat_item):
         else:
             self.infer_meta_func = None
 
+        # print(self.op_yaml_item)
+        # exit(0)
+        if 'infer_shaped_type_op_interface' in self.op_yaml_item:
+            self.infer_shaped_type_op_interface_func = self.op_yaml_item[
+                'infer_shaped_type_op_interface'
+            ]["func"]
+        else:
+            self.infer_shaped_type_op_interface_func = None
+
         # parse backward name
         self.backward_name = self.parse_backward_name()
 
@@ -1252,6 +1270,10 @@ def OpGenerator(
                 if "paddle::dialect::ParseKernelKeyInterface" in op_interfaces:
                     parse_kernel_key_str = parse_kernel_key_template
 
+                reify_infer_shape_str = ""
+                if "paddle::dialect::ReifyInferShapeInterface" in op_interfaces:
+                    reify_infer_shape_str = reify_infer_shape_template
+
                 if op_infer_meta_map is not None:
                     (
                         build_args_with_muta_attr_not_input_for_declare,
@@ -1386,6 +1408,7 @@ def OpGenerator(
                         exclusive_interface=exclusive_interface_str,
                         get_kernel_type_for_var_declare=get_kernel_type_for_var_declare_str,
                         parse_kernel_key_declare=parse_kernel_key_str,
+                        reify_infer_shape_declare=reify_infer_shape_str,
                     )
                     op_defined_str = ""
                 else:
@@ -1408,6 +1431,7 @@ def OpGenerator(
                         exclusive_interface=exclusive_interface_str,
                         get_kernel_type_for_var_declare=get_kernel_type_for_var_declare_str,
                         parse_kernel_key_declare=parse_kernel_key_str,
+                        reify_infer_shape_declare=reify_infer_shape_str,
                     )
                     attribute_names_str = (
                         '"'
@@ -1579,13 +1603,20 @@ def OpGenerator(
                         op_output_optional_list,
                     )
 
-                # generate op GetKernelKeyForVar function str
+                # generate op ParseKernelKeyInterface function str
                 parse_kernel_key_define_str = ''
                 if "paddle::dialect::ParseKernelKeyInterface" in op_interfaces:
                     parse_kernel_key_define_str = gen_parse_kernel_key_str(
                         op_class_name
                     )
 
+                # generate op GetKernelKeyForVar function str
+                reify_infer_shape_define_str = ''
+                if "paddle::dialect::ReifyInferShapeInterface" in op_interfaces:
+                    reify_infer_shape_define_str = gen_reify_infer_shape_str(
+                        op_class_name
+                    )
+
                 # generate op GetKernelKeyForVar function str
                 op_get_kernel_type_for_var_str = ''
                 if dialect_name == "pd_op":
@@ -1644,6 +1675,7 @@ def OpGenerator(
                     ops_defined_list.append(op_infer_meta_str)
                     ops_defined_list.append(op_get_kernel_type_for_var_str)
                     ops_defined_list.append(parse_kernel_key_define_str)
+                    ops_defined_list.append(reify_infer_shape_define_str)
 
                     # NOTE(chenxi67)skip if dialect_name==cinn
                     if dialect_name == "cinn":
 
@@ -0,0 +1,26 @@
+# Copyright (c) 2023 PaddlePaddle Authors. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     https://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+OP_GET_KERNEL_TYPE_FOR_VAR_TEMPLATE = """
+bool {op_name}::ReifyInferShape(pir::Builder &builder,
+                                const std::vector<pir::OpOperand> &operands,
+                                std::vector<pir::Value> &reified_return_shapes) {{
+  VLOG(4) << "Reify infer shape for op: {op_name}";
+  return {op_name}ReifyInferShape(builder, operands, reified_return_shapes);
+}}
+"""
+
+
+def gen_reify_infer_shape_str(op_class_name):
+    return OP_GET_KERNEL_TYPE_FOR_VAR_TEMPLATE.format(op_name=op_class_name)
@@ -16,46 +16,8 @@
 #include "paddle/fluid/pir/dialect/operator/interface/get_kernel_type_for_var.h"
 #include "paddle/fluid/pir/dialect/operator/interface/infermeta.h"
 #include "paddle/fluid/pir/dialect/operator/interface/op_yaml_info.h"
-#include "paddle/fluid/pir/dialect/operator/interface/parse_kernel_key.h"
-#include "paddle/fluid/pir/dialect/operator/interface/vjp.h"
-
-namespace paddle {
-namespace dialect {
-std::vector<std::vector<pir::OpResult>> VjpInterface::Vjp(
-    pir::Operation* op,
-    const std::vector<std::vector<pir::Value>>& inputs,
-    const std::vector<std::vector<pir::OpResult>>& outputs,
-    const std::vector<std::vector<pir::OpResult>>& out_grads,
-    const std::vector<std::vector<bool>>& stop_gradients) {
-  std::vector<std::vector<pir::Value>> out_grads_value;
-  for (const auto& grad : out_grads) {
-    std::vector<pir::Value> grad_value;
-    for (auto op_result : grad) {
-      grad_value.emplace_back(op_result);
-    }
-    out_grads_value.emplace_back(std::move(grad_value));
-  }
-  return impl_->vjp_(op, inputs, outputs, out_grads_value, stop_gradients);
-}
-
-KernelKeyTuple UniqueOpParseKernelKey(pir::Operation* op) {
-  DenseTensorType x_type =
-      op->operand_source(0).type().dyn_cast<paddle::dialect::DenseTensorType>();
-  phi::DataType dtype = TransToPhiDataType(x_type.dtype());
-  pir::BoolAttribute is_sort = op->attribute<pir::BoolAttribute>("is_sorted");
-  phi::Backend backend = phi::Backend::UNDEFINED;
-  if (is_sort.data()) {
-    backend = phi::Backend::CPU;
-  }
-  return {dtype, backend};
-}
-
-}  // namespace dialect
-}  // namespace paddle
 
 IR_DEFINE_EXPLICIT_TYPE_ID(paddle::dialect::InferMetaInterface)
 IR_DEFINE_EXPLICIT_TYPE_ID(paddle::dialect::OpYamlInfoInterface)
-IR_DEFINE_EXPLICIT_TYPE_ID(paddle::dialect::VjpInterface)
 IR_DEFINE_EXPLICIT_TYPE_ID(paddle::dialect::DecompInterface)
 IR_DEFINE_EXPLICIT_TYPE_ID(paddle::dialect::GetKernelTypeForVarInterface)
-IR_DEFINE_EXPLICIT_TYPE_ID(paddle::dialect::ParseKernelKeyInterface)
@@ -0,0 +1,33 @@
+// Copyright (c) 2023 PaddlePaddle Authors. All Rights Reserved.
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     https://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+#include "paddle/fluid/pir/dialect/operator/interface/parse_kernel_key.h"
+
+namespace paddle::dialect {
+
+KernelKeyTuple UniqueOpParseKernelKey(pir::Operation* op) {
+  DenseTensorType x_type =
+      op->operand_source(0).type().dyn_cast<paddle::dialect::DenseTensorType>();
+  phi::DataType dtype = TransToPhiDataType(x_type.dtype());
+  pir::BoolAttribute is_sort = op->attribute<pir::BoolAttribute>("is_sorted");
+  phi::Backend backend = phi::Backend::UNDEFINED;
+  if (is_sort.data()) {
+    backend = phi::Backend::CPU;
+  }
+  return {dtype, backend};
+}
+
+}  // namespace paddle::dialect
+
+IR_DEFINE_EXPLICIT_TYPE_ID(paddle::dialect::ParseKernelKeyInterface)
@@ -0,0 +1,153 @@
+// Copyright (c) 2023 PaddlePaddle Authors. All Rights Reserved.
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     https://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+#include "paddle/fluid/pir/dialect/operator/interface/reify_infer_shape.h"
+#include "paddle/pir/dialect/shape/ir/shape_op.h"
+
+namespace paddle::dialect {
+
+bool ReifyInferShapeInterface::ReifyInferShape(
+    pir::Builder &builder,
+    const std::vector<pir::OpOperand> &operands,
+    std::vector<pir::Value> &reified_return_shapes) {
+  return impl_->reify_return_type_shapes(
+      operation(), builder, operands, reified_return_shapes);
+}
+}  // namespace paddle::dialect
+
+namespace paddle::dialect {
+
+namespace {
+
+bool DeriveShapeFromOperand(pir::Builder *builder,
+                            pir::Value operand,
+                            std::vector<pir::Value> *reified_return_shapes) {
+  auto shaped_type = operand.type().dyn_cast<pir::ShapedTypeInterface>();
+  if (!shaped_type) return false;
+  reified_return_shapes->assign(
+      {builder->Build<pir::shape::ShapeOfOp>(operand).result(0)});
+  return true;
+}
+
+// Returns a new scalar integer value having type `type`.
+//  Here `type` must be an integer or index type.
+pir::Value MaybeCastTo(pir::Builder &builder,  // NOLINT
+                       pir::Value value,
+                       pir::Type type) {
+  if (type == value.type()) return value;
+  // if (!type.IsIndex() && !value.type().IsIndex()) {
+  //   Value casted =
+  //       builder.Build<shape::IndexCastOp>(builder.index_type(), value)
+  //           .result(0);
+  //   return builder.Build<shape::IndexCastOp>(type, casted).result(0);
+  // }
+  // return builder.Build<shape::IndexCastOp>(type, value).result(0);
+}
+}  // namespace
+
+bool AbsOpReifyInferShape(
+    pir::Builder &builder,  // NOLINT
+    const std::vector<pir::OpOperand> &operands,
+    std::vector<pir::Value> &reified_return_shapes) {  // NOLINT
+  return DeriveShapeFromOperand(
+      &builder, operands.front().source(), &reified_return_shapes);
+}
+
+bool Abs_OpReifyInferShape(
+    pir::Builder &builder,  // NOLINT
+    const std::vector<pir::OpOperand> &operands,
+    std::vector<pir::Value> &reified_return_shapes) {  // NOLINT
+  return DeriveShapeFromOperand(
+      &builder, operands.front().source(), &reified_return_shapes);
+}
+
+bool TransposeOpReifyInferShape(
+    pir::Builder &builder,  // NOLINT
+    const std::vector<pir::OpOperand> &operands,
+    std::vector<pir::Value> &reified_return_shapes) {  // NOLINT
+  // auto operand_type = operands[0].type().dyn_cast<DenseTensorType>();
+  // // Currently not support unranked type.
+  // if (!operand_type) return false;
+  // std::vector<int64_t> permutation = this->permutation();
+  // std::vector<Value> shape_values(permutation.size());
+  // Type shape_scalar_type = builder.index_type();
+  // auto to_shape_scalar_type = [&](Value v) {
+  //   return MaybeCastTo(builder, v, shape_scalar_type);
+  // };
+  // auto shaped_type = operand_type.dyn_cast<ShapedTypeInterface>();
+  // auto shape_vector = shaped_type.GetDyShape();
+  // for (auto [idx, element] = std::tuple{0, shape_vector.begin()};
+  //      element != shape_vector.end();
+  //      ++idx, ++element) {
+  //   auto it = std::find(permutation.begin(), permutation.end(), idx);
+  //   // TODO(zhangbopd): Need BuildOrFold
+  //   Value value_dim = to_shape_scalar_type(
+  //       builder.Build<shape::TensorDimOp>(operands[0].source(),
+  //       idx).result(0));
+  //   shape_values[std::distance(permutation.begin(), it)] = value_dim;
+  // }
+  // Value output_shape =
+  //     builder.Build<shape::FromElementsOp>(shape_values).result(0);
+  // reified_return_shapes.push_back(output_shape);
+
+  return true;
+}
+
+bool ConcatOpReifyInferShape(
+    pir::Builder &builder,  // NOLINT
+    const std::vector<pir::OpOperand> &operands,
+    std::vector<pir::Value> &reified_return_shapes) {  // NOLINT
+  // std::vector<Value> inputs = {x()};
+  // auto operand_type = inputs[0].type().dyn_cast<DenseTensorType>();
+  // // Currently not support unranked type.
+  // if (!operand_type) return false;
+  // Type shapeScalarType = builder.index_type();
+  // auto to_shape_scalar_type = [&](Value v) {
+  //   return MaybeCastTo(builder, v, shapeScalarType);
+  // };
+  // std::vector<std::vector<Value>> all_shape_values;
+  // for (size_t inputId = 0; inputId < inputs.size(); ++inputId) {
+  //   Value operand = inputs[inputId];
+  //   auto operand_type = operand.type().dyn_cast<DenseTensorType>();
+  //   if (!operand_type) return false;
+  //   std::vector<Value> shape_values;
+  //   auto shaped_type = operand_type.dyn_cast<ShapedTypeInterface>();
+  //   auto shape_vector = shaped_type.GetDyShape();
+  //   for (auto [idx, element] = std::tuple{0, shape_vector.begin()};
+  //        element != shape_vector.end();
+  //        ++idx, ++element) {
+  //     Value value_dim = to_shape_scalar_type(
+  //         builder.Build<shape::TensorDimOp>(operand, idx).result(0));
+  //     shape_values.push_back(value_dim);
+  //   }
+  //   all_shape_values.emplace_back(std::move(shape_values));
+  // }
+  // [[maybe_unused]] int axis = this->dimension();
+  // auto &shape_values = all_shape_values[0];
+  // for (size_t vecId = 1; vecId < all_shape_values.size(); ++vecId) {
+  //   auto &otherShapeValues = all_shape_values[vecId];
+  //   if (otherShapeValues.size() != shape_values.size()) return false;
+  // TODO(zhangbopd): AddIOp
+  // shape_values[axis] =
+  //     builder.Build<arith::AddIOp>(shape_values[axis],
+  //     otherShapeValues[axis]);
+  // }
+  // Value output_shape =
+  //     builder.Build<shape::FromElementsOp>(shape_values).result(0);
+  // reified_return_shapes.push_back(output_shape);
+  return true;
+}
+
+}  // namespace paddle::dialect
+IR_DEFINE_EXPLICIT_TYPE_ID(paddle::dialect::ReifyInferShapeInterface)