Support static graph code-gen for bincount (#54686)

PaddlePaddle · Jun 26, 2023 · b547c4a · b547c4a
1 parent 733eca8
commit b547c4a
Show file tree

Hide file tree

Showing 7 changed files with 34 additions and 108 deletions.
diff --git a/paddle/fluid/operators/bincount_op.cc b/paddle/fluid/operators/bincount_op.cc
diff --git a/paddle/fluid/operators/generator/get_expected_kernel_func.cc b/paddle/fluid/operators/generator/get_expected_kernel_func.cc
@@ -410,5 +410,14 @@ phi::KernelKey GetConvExpectedKernelType(
   return phi::KernelKey(input_data_type, ctx.GetPlace());
 }
 
+phi::KernelKey GetBincountExpectedKernelType(
+    const framework::ExecutionContext& ctx,
+    const framework::OperatorWithKernel* op_ptr) {
+  auto data_type = ctx.HasInput("Weights")
+                       ? op_ptr->IndicateVarDataType(ctx, "Weights")
+                       : op_ptr->IndicateVarDataType(ctx, "X");
+  return phi::KernelKey(data_type, ctx.device_context().GetPlace());
+}
+
 }  // namespace operators
 }  // namespace paddle
diff --git a/paddle/fluid/operators/generator/get_expected_kernel_func.h b/paddle/fluid/operators/generator/get_expected_kernel_func.h
@@ -100,5 +100,9 @@ phi::KernelKey GetConvExpectedKernelType(
     const framework::ExecutionContext& ctx,
     const framework::OperatorWithKernel* op_ptr);
 
+phi::KernelKey GetBincountExpectedKernelType(
+    const framework::ExecutionContext& ctx,
+    const framework::OperatorWithKernel* op_ptr);
+
 }  // namespace operators
 }  // namespace paddle
diff --git a/paddle/phi/api/yaml/legacy_ops.yaml b/paddle/phi/api/yaml/legacy_ops.yaml
@@ -121,15 +121,6 @@
   view : (mean -> mean_out), (variance -> variance_out)
   backward : batch_norm_grad
 
-- op : bincount
-  args: (Tensor x, Tensor weights, Scalar(int) minlength = 0)
-  output: Tensor(out)
-  infer_meta:
-    func: BincountInferMeta
-  kernel:
-    func: bincount
-  optional: weights
-
 - op : cast
   args : (Tensor x, DataType dtype)
   output : Tensor

diff --git a/paddle/phi/api/yaml/op_compat.yaml b/paddle/phi/api/yaml/op_compat.yaml
@@ -372,6 +372,18 @@
   extra :
     attrs : [bool use_mkldnn = false]
 
+- op : bincount
+  inputs :
+    {x : X, weights : Weights}
+  outputs :
+    out : Out
+  scalar:
+    minlength:
+      data_type : int
+      support_tensor : true
+  get_expected_kernel_type :
+    bincount : GetBincountExpectedKernelType
+
 - op : bitwise_and
   inputs :
     {x : X, y : Y}

diff --git a/paddle/phi/api/yaml/ops.yaml b/paddle/phi/api/yaml/ops.yaml
@@ -294,6 +294,15 @@
   data_transform :
     skip_transform : out_size, size_tensor, scale_tensor
 
+- op : bincount
+  args: (Tensor x, Tensor weights, Scalar(int) minlength = 0)
+  output: Tensor(out)
+  infer_meta:
+    func: BincountInferMeta
+  kernel:
+    func: bincount
+  optional: weights
+
 - op : bitwise_and
   args : (Tensor x, Tensor y)
   output : Tensor(out)

diff --git a/paddle/phi/ops/compat/bincount_sig.cc b/paddle/phi/ops/compat/bincount_sig.cc