【Infer Symbolic Shape No.123】【BUAA】Add Bincount (#67435)

Author: Whsjrczr

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

@@ -234,12 +234,42 @@ bool Binomial_OpInferSymbolicShape(
   return BinomialOpInferSymbolicShape(op, infer_context);
 }
 
-// bool BincountOpInferSymbolicShape(pir::Operation *op,
-//                                   pir::InferSymbolicShapeContext
-//                                   *infer_context) {
-//   // pass
-//   return true;
-// }
+bool BincountOpInferSymbolicShape(
+    pir::Operation *op, pir::InferSymbolicShapeContext *infer_context) {
+  const auto &x_shape_or_data =
+      infer_context->GetShapeOrDataForValue(op->operand_source(0));
+  const std::vector<symbol::DimExpr> &x_dims = x_shape_or_data.shape();
+
+  PADDLE_ENFORCE_EQ(x_dims.size(),
+                    1,
+                    common::errors::InvalidArgument(
+                        "The 'shape' of Input(X) must be 1-D tensor. But the "
+                        "dimension of Input(X) is [%d]",
+                        x_dims.size()));
+
+  if (op->operand_source(1)) {
+    const auto &weights_shape_or_data =
+        infer_context->GetShapeOrDataForValue(op->operand_source(1));
+    const std::vector<symbol::DimExpr> &weights_dims =
+        weights_shape_or_data.shape();
+
+    PADDLE_ENFORCE_EQ(weights_dims.size(),
+                      1,
+                      common::errors::InvalidArgument(
+                          "The 'shape' of Input(Weights) must be 1-D tensor. "
+                          "But the dimension of Input(Weights) is [%d]",
+                          weights_dims.size()));
+    infer_context->AddEqualCstr(weights_dims[0], x_dims[0]);
+  }
+
+  symbol::DimExpr out_unknown = infer_context->GetNextSymName();
+  const std::vector<symbol::DimExpr> out_dims = {out_unknown};
+  symbol::ShapeOrDataDimExprs output_dims{
+      symbol::TensorShapeOrDataDimExprs(out_dims)};
+  infer_context->SetShapeOrDataForValue(op->result(0), output_dims);
+
+  return true;
+}
 
 // bool BmmOpInferSymbolicShape(pir::Operation *op,
 //                              pir::InferSymbolicShapeContext *infer_context) {

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

@@ -26,7 +26,7 @@ OP_DECLARE_INFER_SYMBOLIC_SHAPE(BceLoss_)
 OP_DECLARE_INFER_SYMBOLIC_SHAPE(BoxClip)
 OP_DECLARE_INFER_SYMBOLIC_SHAPE(Binomial)
 OP_DECLARE_INFER_SYMBOLIC_SHAPE(Binomial_)
-// OP_DECLARE_INFER_SYMBOLIC_SHAPE(Bincount)
+OP_DECLARE_INFER_SYMBOLIC_SHAPE(Bincount)
 // OP_DECLARE_INFER_SYMBOLIC_SHAPE(Bmm)
 // OP_DECLARE_INFER_SYMBOLIC_SHAPE(CholeskySolve)
 OP_DECLARE_INFER_SYMBOLIC_SHAPE(CtcAlign)

paddle/phi/ops/yaml/ops.yaml

@@ -549,6 +549,7 @@
   kernel:
     func: bincount
   optional: weights
+  interfaces : paddle::dialect::InferSymbolicShapeInterface
 
 - op : binomial
   args : (Tensor count, Tensor prob)

test/legacy_test/test_bincount_op.py

@@ -156,7 +156,7 @@ def init_test_case(self):
         self.Out = np.bincount(self.np_input, minlength=self.minlength)
 
     def test_check_output(self):
-        self.check_output(check_pir=True)
+        self.check_output(check_pir=True, check_symbol_infer=False)
 
 
 class TestCase1(TestBincountOp):