[PIR][DynamicShape] Add nullary_infer_sym and binary nullary_infer_sym

paddle/fluid/pir/dialect/operator/interface/infer_symbolic_shape/cinn_op_infer_sym.cc

@@ -76,7 +76,7 @@ bool ConcatOpInferSymbolicShape(
       out_dims[axis] = out_dims[axis] + operand_shape_or_data.shape()[axis];
     }
 
-    for (size_t i = 1; i < rank; ++i) {
+    for (size_t i = 0; i < rank; ++i) {
       if (i == static_cast<size_t>(axis)) continue;
       paddle::dialect::details::BuildCstrEqForTensorListAlongAxis(
           shape_analysis, input_values, i);
@@ -85,6 +85,9 @@ bool ConcatOpInferSymbolicShape(
     return out_dims;
   };
 
+  VLOG(3) << "constraints size:"
+          << shape_analysis->CreateDimExprBuilder().constraints().size();
+
   symbol::ShapeOrDataDimExprs shape_data{
       symbol::TensorShapeOrDataDimExprs(GetOutDimExprs())};
 

paddle/fluid/pir/dialect/operator/interface/infer_symbolic_shape/cinn_op_infer_sym.h

@@ -13,7 +13,6 @@
 // limitations under the License.
 
 #pragma once
-#include "paddle/fluid/pir/dialect/operator/utils/utils.h"
 #include "paddle/pir/include/dialect/shape/utils/shape_analysis.h"
 
 namespace cinn::dialect {

paddle/fluid/pir/dialect/operator/interface/infer_symbolic_shape/infer_sym_element_wise_binary.h

@@ -14,7 +14,6 @@
 
 #pragma once
 
-#include "paddle/fluid/pir/dialect/operator/utils/utils.h"
 #include "paddle/pir/include/dialect/shape/utils/shape_analysis.h"
 
 namespace paddle::dialect {

paddle/fluid/pir/dialect/operator/interface/infer_symbolic_shape/infer_symbolic_shape.h

@@ -16,6 +16,7 @@
 
 #include "paddle/fluid/pir/dialect/operator/interface/infer_symbolic_shape/cinn_op_infer_sym.h"
 #include "paddle/fluid/pir/dialect/operator/interface/infer_symbolic_shape/infer_sym_element_wise_binary.h"
+#include "paddle/fluid/pir/dialect/operator/interface/infer_symbolic_shape/nullary_infer_sym.h"
 #include "paddle/fluid/pir/dialect/operator/interface/infer_symbolic_shape/paddle_op_infer_sym.h"
 #include "paddle/fluid/pir/dialect/operator/interface/infer_symbolic_shape/same_operands_and_result.h"
 #include "paddle/fluid/pir/dialect/operator/interface/infer_symbolic_shape/unary_infer_sym.h"

paddle/fluid/pir/dialect/operator/interface/infer_symbolic_shape/nullary_infer_sym.cc

@@ -0,0 +1,74 @@
+// Copyright (c) 2024 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
+//
+//     http://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/infer_symbolic_shape/nullary_infer_sym.h"
+#include "paddle/fluid/pir/dialect/operator/interface/infer_symbolic_shape/infer_sym_utils.h"
+
+namespace paddle::dialect {
+
+bool EmptyOpInferSymbolicShape(pir::Operation *op,
+                               pir::ShapeConstraintIRAnalysis *shape_analysis) {
+  const auto &shape_gen_op = op->operand_source(0).defining_op();
+  if (shape_gen_op->isa<paddle::dialect::FullIntArrayOp>()) {
+    std::vector<int64_t> shape = details::GetVectorAttr(
+        shape_gen_op->dyn_cast<paddle::dialect::FullIntArrayOp>(), "value");
+    std::vector<symbol::DimExpr> sym_dims;
+    sym_dims.reserve(shape.size());
+    for (const int64_t &dim : shape) {
+      sym_dims.emplace_back(symbol::DimExpr(dim));
+    }
+
+    symbol::ShapeOrDataDimExprs shape_data{
+        symbol::TensorShapeOrDataDimExprs(sym_dims)};
+    shape_analysis->SetShapeOrDataForValue(op->result(0), shape_data);
+    return true;
+
+  } else {
+    pir::Value operand_source = op->operand_source(0);
+    const symbol::ShapeOrDataDimExprs &operand_shape_or_data =
+        shape_analysis->GetShapeOrDataForValue(operand_source);
+
+    shape_analysis->SetShapeOrDataForValue(op->result(0),
+                                           operand_shape_or_data);
+    return true;
+  }
+}
+
+bool GaussianOpInferSymbolicShape(
+    pir::Operation *op, pir::ShapeConstraintIRAnalysis *shape_analysis) {
+  const auto &shape_gen_op = op->operand_source(0).defining_op();
+
+  if (shape_gen_op->isa<paddle::dialect::FullIntArrayOp>()) {
+    std::vector<int64_t> shape = details::GetVectorAttr(
+        shape_gen_op->dyn_cast<paddle::dialect::FullIntArrayOp>(), "value");
+    std::vector<symbol::DimExpr> sym_dims;
+    sym_dims.reserve(shape.size());
+    for (const int64_t &dim : shape) {
+      sym_dims.emplace_back(symbol::DimExpr(dim));
+    }
+
+    symbol::ShapeOrDataDimExprs shape_data{
+        symbol::TensorShapeOrDataDimExprs(sym_dims)};
+    shape_analysis->SetShapeOrDataForValue(op->result(0), shape_data);
+    return true;
+
+  } else {
+    PADDLE_THROW(phi::errors::Unimplemented(
+        op->name() +
+        " 's InferSymbolicShape interface is NOT implemented now."));
+    return true;
+  }
+}
+
+}  // namespace paddle::dialect

paddle/fluid/pir/dialect/operator/interface/infer_symbolic_shape/nullary_infer_sym.h

@@ -0,0 +1,22 @@
+// Copyright (c) 2024 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
+//
+//     http://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.
+
+#pragma once
+
+#include "paddle/pir/include/dialect/shape/utils/shape_analysis.h"
+
+namespace paddle::dialect {
+OP_DECLARE_INFER_SYMBOLIC_SHAPE(Empty)
+OP_DECLARE_INFER_SYMBOLIC_SHAPE(Gaussian)
+}  // namespace paddle::dialect

paddle/fluid/pir/dialect/operator/interface/infer_symbolic_shape/paddle_op_infer_sym.cc

@@ -97,7 +97,6 @@ bool StackOpInferSymbolicShape(pir::Operation *op,
           static_cast<std::int64_t>(shape_data_list.size()));
     } else {
       for (int i = 0; i < rank; ++i) {
-        if (i == axis) continue;
         details::BuildCstrEqForTensorListAlongAxis(
             shape_analysis, shape_data_list, i);
       }
@@ -931,26 +930,6 @@ bool SplitOpInferSymbolicShape(pir::Operation *op,
 }
 
 //  Not Implemented Ops.
-
-bool DiagEmbedOpInferSymbolicShape(
-    pir::Operation *op, pir::ShapeConstraintIRAnalysis *shape_analysis) {
-  PADDLE_THROW(phi::errors::Unimplemented(
-      op->name() + " 's InferSymbolicShape interface is NOT implemented now."));
-  return true;
-}
-bool DiagonalOpInferSymbolicShape(
-    pir::Operation *op, pir::ShapeConstraintIRAnalysis *shape_analysis) {
-  PADDLE_THROW(phi::errors::Unimplemented(
-      op->name() + " 's InferSymbolicShape interface is NOT implemented now."));
-  return true;
-}
-bool DirichletOpInferSymbolicShape(
-    pir::Operation *op, pir::ShapeConstraintIRAnalysis *shape_analysis) {
-  PADDLE_THROW(phi::errors::Unimplemented(
-      op->name() + " 's InferSymbolicShape interface is NOT implemented now."));
-  return true;
-}
-
 bool GatherOpInferSymbolicShape(
     pir::Operation *op, pir::ShapeConstraintIRAnalysis *shape_analysis) {
   const auto &input_shape_or_data =
@@ -1020,17 +999,33 @@ bool GatherOpInferSymbolicShape(
 
 bool KronOpInferSymbolicShape(pir::Operation *op,
                               pir::ShapeConstraintIRAnalysis *shape_analysis) {
-  PADDLE_THROW(phi::errors::Unimplemented(
-      op->name() + " 's InferSymbolicShape interface is NOT implemented now."));
-  return true;
-}
-bool KthvalueOpInferSymbolicShape(
-    pir::Operation *op, pir::ShapeConstraintIRAnalysis *shape_analysis) {
-  PADDLE_THROW(phi::errors::Unimplemented(
-      op->name() + " 's InferSymbolicShape interface is NOT implemented now."));
+  const auto &x_shape_or_data =
+      shape_analysis->GetShapeOrDataForValue(op->operand_source(0)).shape();
+  const auto &y_shape_or_data =
+      shape_analysis->GetShapeOrDataForValue(op->operand_source(1)).shape();
+  const int rank_x = x_shape_or_data.size();
+  const int rank_y = y_shape_or_data.size();
+  const int rank = (rank_x > rank_y) ? rank_x : rank_y;
+
+  std::vector<symbol::DimExpr> dim_out;
+  dim_out.reserve(rank);
+  const auto one = symbol::DimExpr{1};
+  const auto minus_one = symbol::DimExpr{-1};
+  for (int i = 0; i < rank; i++) {
+    symbol::DimExpr dim_xi =
+        (i < rank - rank_x) ? one : x_shape_or_data.at(i - (rank - rank_x));
+    symbol::DimExpr dim_yi =
+        (i < rank - rank_y) ? one : y_shape_or_data.at(i - (rank - rank_y));
+    dim_out.push_back(dim_xi * dim_yi);
+  }
+  symbol::ShapeOrDataDimExprs shape_data{
+      symbol::TensorShapeOrDataDimExprs(dim_out)};
+  pir::Value res = op->result(0);
+  shape_analysis->SetShapeOrDataForValue(res, shape_data);
   return true;
 }
 
+//  Not Impelmented Ops.
 bool LogcumsumexpOpInferSymbolicShape(
     pir::Operation *op, pir::ShapeConstraintIRAnalysis *shape_analysis) {
   PADDLE_THROW(phi::errors::Unimplemented(
@@ -1095,32 +1090,6 @@ bool UniqueConsecutiveOpInferSymbolicShape(
       op->name() + " 's InferSymbolicShape interface is NOT implemented now."));
   return true;
 }
-
-bool EinsumOpInferSymbolicShape(
-    pir::Operation *op, pir::ShapeConstraintIRAnalysis *shape_analysis) {
-  PADDLE_THROW(phi::errors::Unimplemented(
-      op->name() + " 's InferSymbolicShape interface is NOT implemented now."));
-  return true;
-}
-bool EmptyOpInferSymbolicShape(pir::Operation *op,
-                               pir::ShapeConstraintIRAnalysis *shape_analysis) {
-  PADDLE_THROW(phi::errors::Unimplemented(
-      op->name() + " 's InferSymbolicShape interface is NOT implemented now."));
-  return true;
-}
-bool Exponential_OpInferSymbolicShape(
-    pir::Operation *op, pir::ShapeConstraintIRAnalysis *shape_analysis) {
-  PADDLE_THROW(phi::errors::Unimplemented(
-      op->name() + " 's InferSymbolicShape interface is NOT implemented now."));
-  return true;
-}
-bool GaussianOpInferSymbolicShape(
-    pir::Operation *op, pir::ShapeConstraintIRAnalysis *shape_analysis) {
-  PADDLE_THROW(phi::errors::Unimplemented(
-      op->name() + " 's InferSymbolicShape interface is NOT implemented now."));
-  return true;
-}
-
 bool LinspaceOpInferSymbolicShape(
     pir::Operation *op, pir::ShapeConstraintIRAnalysis *shape_analysis) {
   PADDLE_THROW(phi::errors::Unimplemented(

paddle/fluid/pir/dialect/operator/interface/infer_symbolic_shape/paddle_op_infer_sym.h

@@ -14,7 +14,6 @@
 
 #pragma once
 
-#include "paddle/fluid/pir/dialect/operator/utils/utils.h"
 #include "paddle/pir/include/dialect/shape/utils/shape_analysis.h"
 
 namespace paddle::dialect {
@@ -51,12 +50,8 @@ OP_DECLARE_INFER_SYMBOLIC_SHAPE(Split)
 
 //  Not Impelmented Ops.
 
-OP_DECLARE_INFER_SYMBOLIC_SHAPE(DiagEmbed)
-OP_DECLARE_INFER_SYMBOLIC_SHAPE(Diagonal)
-OP_DECLARE_INFER_SYMBOLIC_SHAPE(Dirichlet)
 OP_DECLARE_INFER_SYMBOLIC_SHAPE(Gather)
 OP_DECLARE_INFER_SYMBOLIC_SHAPE(Kron)
-OP_DECLARE_INFER_SYMBOLIC_SHAPE(Kthvalue)
 OP_DECLARE_INFER_SYMBOLIC_SHAPE(Logcumsumexp)
 OP_DECLARE_INFER_SYMBOLIC_SHAPE(MaskedSelect)
 OP_DECLARE_INFER_SYMBOLIC_SHAPE(Poisson)
@@ -67,10 +62,6 @@ OP_DECLARE_INFER_SYMBOLIC_SHAPE(TakeAlongAxis)
 OP_DECLARE_INFER_SYMBOLIC_SHAPE(Topk)
 OP_DECLARE_INFER_SYMBOLIC_SHAPE(Unbind)
 OP_DECLARE_INFER_SYMBOLIC_SHAPE(UniqueConsecutive)
-OP_DECLARE_INFER_SYMBOLIC_SHAPE(Einsum)
-OP_DECLARE_INFER_SYMBOLIC_SHAPE(Empty)
-OP_DECLARE_INFER_SYMBOLIC_SHAPE(Exponential_)
-OP_DECLARE_INFER_SYMBOLIC_SHAPE(Gaussian)
 OP_DECLARE_INFER_SYMBOLIC_SHAPE(Linspace)
 OP_DECLARE_INFER_SYMBOLIC_SHAPE(Logspace)
 OP_DECLARE_INFER_SYMBOLIC_SHAPE(Logsumexp)

paddle/fluid/pir/dialect/operator/interface/infer_symbolic_shape/same_operands_and_result.cc

@@ -154,6 +154,10 @@ bool Digamma_OpInferSymbolicShape(
     pir::Operation *op, pir::ShapeConstraintIRAnalysis *shape_analysis) {
   return SameOperandsAndResultShape(op, shape_analysis);
 }
+bool DirichletOpInferSymbolicShape(
+    pir::Operation *op, pir::ShapeConstraintIRAnalysis *shape_analysis) {
+  return SameOperandsAndResultShape(op, shape_analysis);
+}
 bool EqualOpInferSymbolicShape(pir::Operation *op,
                                pir::ShapeConstraintIRAnalysis *shape_analysis) {
   return SameOperandsAndResultShape(op, shape_analysis);
@@ -194,6 +198,10 @@ bool Expm1_OpInferSymbolicShape(
     pir::Operation *op, pir::ShapeConstraintIRAnalysis *shape_analysis) {
   return SameOperandsAndResultShape(op, shape_analysis);
 }
+bool Exponential_OpInferSymbolicShape(
+    pir::Operation *op, pir::ShapeConstraintIRAnalysis *shape_analysis) {
+  return SameOperandsAndResultShape(op, shape_analysis);
+}
 bool FetchOpInferSymbolicShape(pir::Operation *op,
                                pir::ShapeConstraintIRAnalysis *shape_analysis) {
   return SameOperandsAndResultShape(op, shape_analysis);

paddle/fluid/pir/dialect/operator/interface/infer_symbolic_shape/same_operands_and_result.h

@@ -14,7 +14,6 @@
 
 #pragma once
 
-#include "paddle/fluid/pir/dialect/operator/utils/utils.h"
 #include "paddle/pir/include/dialect/shape/utils/shape_analysis.h"
 
 namespace paddle::dialect {
@@ -50,6 +49,7 @@ OP_DECLARE_INFER_SYMBOLIC_SHAPE(Cosh)
 OP_DECLARE_INFER_SYMBOLIC_SHAPE(Cosh_)
 OP_DECLARE_INFER_SYMBOLIC_SHAPE(Digamma)
 OP_DECLARE_INFER_SYMBOLIC_SHAPE(Digamma_)
+OP_DECLARE_INFER_SYMBOLIC_SHAPE(Dirichlet)
 OP_DECLARE_INFER_SYMBOLIC_SHAPE(Equal)
 OP_DECLARE_INFER_SYMBOLIC_SHAPE(Equal_)
 OP_DECLARE_INFER_SYMBOLIC_SHAPE(Erf)
@@ -60,6 +60,7 @@ OP_DECLARE_INFER_SYMBOLIC_SHAPE(Exp)
 OP_DECLARE_INFER_SYMBOLIC_SHAPE(Exp_)
 OP_DECLARE_INFER_SYMBOLIC_SHAPE(Expm1)
 OP_DECLARE_INFER_SYMBOLIC_SHAPE(Expm1_)
+OP_DECLARE_INFER_SYMBOLIC_SHAPE(Exponential_)
 OP_DECLARE_INFER_SYMBOLIC_SHAPE(Fetch)
 OP_DECLARE_INFER_SYMBOLIC_SHAPE(Flip)
 OP_DECLARE_INFER_SYMBOLIC_SHAPE(Floor)