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

paddle/fluid/operators/bincount_op.cc

-/* Copyright (c) 2020 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 <string>
-#include <unordered_map>
-#include <vector>
-
-#include "paddle/fluid/framework/infershape_utils.h"
-#include "paddle/fluid/framework/op_registry.h"
-#include "paddle/phi/core/infermeta_utils.h"
-#include "paddle/phi/infermeta/binary.h"
-
-namespace paddle {
-namespace operators {
-
-class BincountOp : public framework::OperatorWithKernel {
- public:
-  using framework::OperatorWithKernel::OperatorWithKernel;
-
-  phi::KernelKey GetExpectedKernelType(
-      const framework::ExecutionContext &ctx) const {
-    auto data_type =
-        ctx.HasInput("Weights")
-            ? OperatorWithKernel::IndicateVarDataType(ctx, "Weights")
-            : OperatorWithKernel::IndicateVarDataType(ctx, "X");
-    return phi::KernelKey(data_type, ctx.device_context().GetPlace());
-  }
-};
-
-class BincountOpMaker : public framework::OpProtoAndCheckerMaker {
- public:
-  void Make() override {
-    AddInput("X", "(Tensor) The input tensor of Bincount op,");
-    AddInput("Weights", "(Tensor) The weights tensor of Bincount op,")
-        .AsDispensable();
-    AddOutput("Out", "(Tensor) The output tensor of Bincount op,");
-    AddAttr<int>("minlength", "(int) The minimal numbers of bins")
-        .SetDefault(0)
-        .EqualGreaterThan(0)
-        .SupportTensor();
-    AddComment(R"DOC(
-          Bincount Operator.
-          Computes frequency of each value in the input tensor.
-          Elements of input tensor should be non-negative ints.
-      )DOC");
-  }
-};
-
-}  // namespace operators
-}  // namespace paddle
-
-namespace ops = paddle::operators;
-DECLARE_INFER_SHAPE_FUNCTOR(bincount,
-                            BincountInferShapeFunctor,
-                            PD_INFER_META(phi::BincountInferMeta));
-REGISTER_OPERATOR(
-    bincount,
-    ops::BincountOp,
-    ops::BincountOpMaker,
-    paddle::framework::EmptyGradOpMaker<paddle::framework::OpDesc>,
-    paddle::framework::EmptyGradOpMaker<paddle::imperative::OpBase>,
-    BincountInferShapeFunctor);

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

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

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

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}

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)

paddle/phi/ops/compat/bincount_sig.cc

-// Copyright (c) 2022 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/phi/core/compat/op_utils.h"
-
-namespace phi {
-
-KernelSignature BincountOpArgumentMapping(
-    const ArgumentMappingContext& ctx UNUSED) {
-  return KernelSignature("bincount", {"X", "Weights"}, {"minlength"}, {"Out"});
-}
-
-}  // namespace phi
-
-PD_REGISTER_ARG_MAPPING_FN(bincount, phi::BincountOpArgumentMapping);