[Custom OP] support double attribute type (#68824)

BeingGod · web-flow · commit 7d9c687a7e54 · 2024-10-23T13:33:03.000+08:00
* [Custom OP] support double attribute type

* fix

* fix
diff --git a/paddle/fluid/framework/custom_operator.cc b/paddle/fluid/framework/custom_operator.cc
@@ -177,7 +177,7 @@ static void RunKernelFunc(
     } else {
       PADDLE_THROW(common::errors::Unimplemented(
           "Unsupported `%s` type value as custom attribute now. "
-          "Supported data types include `bool`, `int`, `float`, "
+          "Supported data types include `bool`, `int`, `float`, `double`, "
           "`int64_t`, `std::string`, `std::vector<int>`, "
           "`std::vector<float>`, `std::vector<int64_t>`, "
           "`std::vector<std::string>`, Please check whether "
@@ -1369,6 +1369,8 @@ void PD_RegisterOperator(const char* kernel_name_cstr,
         op_attrs.push_back("Attr_" + std::to_string(i) + ":int");
       } else if (attr_type == PD_KernelArgumentType::PD_ARG_TYPE_FLOAT32) {
         op_attrs.push_back("Attr_" + std::to_string(i) + ":float");
+      } else if (attr_type == PD_KernelArgumentType::PD_ARG_TYPE_FLOAT64) {
+        op_attrs.push_back("Attr_" + std::to_string(i) + ":double");
       } else if (attr_type == PD_KernelArgumentType::PD_ARG_TYPE_INT64) {
         op_attrs.push_back("Attr_" + std::to_string(i) + ":int64_t");
       } else if (attr_type == PD_KernelArgumentType::PD_ARG_TYPE_STRING) {
diff --git a/paddle/fluid/framework/custom_operator.h b/paddle/fluid/framework/custom_operator.h
@@ -64,6 +64,9 @@ class CustomOpMaker : public OpProtoAndCheckerMaker {
       } else if (attr_type_str == "float") {
         AddAttr<float>(attr_name, "custom operator float attribute.")
             .SetDefault(1.0f);
+      } else if (attr_type_str == "double") {
+        AddAttr<double>(attr_name, "custom operator double attribute.")
+            .SetDefault(1.0f);
       } else if (attr_type_str == "int64_t") {
         AddAttr<int64_t>(attr_name, "custom operator int64_t attribute.")
             .SetDefault(1);
@@ -89,7 +92,7 @@ class CustomOpMaker : public OpProtoAndCheckerMaker {
       } else {
         PADDLE_THROW(common::errors::Unimplemented(
             "Unsupported `%s` type value as custom attribute now. "
-            "Supported data types include `bool`, `int`, `float`, "
+            "Supported data types include `bool`, `int`, `float`, `double`"
             "`int64_t`, `std::string`, `std::vector<int>`, "
             "`std::vector<float>`, `std::vector<int64_t>`, "
             "`std::vector<std::string>`, Please check whether "
diff --git a/paddle/fluid/pybind/eager_functions.cc b/paddle/fluid/pybind/eager_functions.cc
@@ -636,6 +636,9 @@ PyObject* eager_api_run_custom_op(PyObject* self,
     } else if (attr_type_str == "float") {
       ctx.EmplaceBackAttr(
           CastPyArg2AttrFloat(obj, attr_start_idx + i));  // NOLINT
+    } else if (attr_type_str == "double") {
+      ctx.EmplaceBackAttr(
+          CastPyArg2AttrDouble(obj, attr_start_idx + i));  // NOLINT
     } else if (attr_type_str == "int64_t") {
       ctx.EmplaceBackAttr(
           CastPyArg2Long(obj, op_type, attr_start_idx + i));  // NOLINT
diff --git a/paddle/fluid/pybind/eager_utils.cc b/paddle/fluid/pybind/eager_utils.cc
@@ -270,6 +270,18 @@ float CastPyArg2AttrFloat(PyObject* obj, ssize_t arg_pos) {
   }
 }
 
+double CastPyArg2AttrDouble(PyObject* obj, ssize_t arg_pos) {
+  if (PyObject_CheckFloat(obj)) {
+    return PyObject_ToDouble(obj);
+  } else {
+    PADDLE_THROW(common::errors::InvalidType(
+        "argument (position %d) must be "
+        "float, but got %s",
+        arg_pos + 1,
+        (reinterpret_cast<PyTypeObject*>(obj->ob_type))->tp_name));
+  }
+}
+
 std::string CastPyArg2AttrString(PyObject* obj, ssize_t arg_pos) {
   if (PyObject_CheckStr(obj)) {
     Py_ssize_t size = 0;
diff --git a/paddle/fluid/pybind/eager_utils.h b/paddle/fluid/pybind/eager_utils.h
@@ -76,6 +76,7 @@ int CastPyArg2AttrInt(PyObject* obj, ssize_t arg_pos);
 int64_t CastPyArg2AttrLong(PyObject* obj, ssize_t arg_pos);
 size_t CastPyArg2AttrSize_t(PyObject* obj, ssize_t arg_pos);
 float CastPyArg2AttrFloat(PyObject* obj, ssize_t arg_pos);
+double CastPyArg2AttrDouble(PyObject* obj, ssize_t arg_pos);
 std::string CastPyArg2AttrString(PyObject* obj, ssize_t arg_pos);
 std::shared_ptr<imperative::VarBase> CastPyArg2VarBase(PyObject* obj,
                                                        ssize_t arg_pos);
diff --git a/paddle/phi/api/ext/op_meta_info.h b/paddle/phi/api/ext/op_meta_info.h
@@ -326,6 +326,7 @@ struct KernelFuncImpl<Return (*)(Args...), impl_fn> {
   PD_SPECIALIZE_ComputeCallHelper(bool);
   PD_SPECIALIZE_ComputeCallHelper(int);
   PD_SPECIALIZE_ComputeCallHelper(float);
+  PD_SPECIALIZE_ComputeCallHelper(double);
   PD_SPECIALIZE_ComputeCallHelper(int64_t);
   PD_SPECIALIZE_ComputeCallHelper(const std::string&);
   PD_SPECIALIZE_ComputeCallHelper(const std::vector<int>&);
@@ -337,6 +338,7 @@ struct KernelFuncImpl<Return (*)(Args...), impl_fn> {
   // - paddle::blank, std::vector<bool> and std::vector<double>
   //   are not used in op
   // - BlockDesc* and std::vector<BlockDesc*> are used in framework
+  // NOTE(beinggod): Now we need to support double.
 
   // NOTE(chenweihang): Used to be compatible with the 2.0.1 released
   // interface, and will be deprecated in the future
@@ -616,6 +618,9 @@ struct InferShapeFuncImpl<Return (*)(Args...), impl_fn> {
   PD_SPECIALIZE_InferShapeCallHelper_FOR_ATTR(std::vector<float>);
   PD_SPECIALIZE_InferShapeCallHelper_FOR_ATTR(std::vector<std::string>);
 
+  // NOTE(beinggod): extent attribute for custom op
+  PD_SPECIALIZE_InferShapeCallHelper_FOR_ATTR(double);
+
   // end: base template
   template <typename T>
   struct InferShapeCallHelper<TypeTag<T>> {
diff --git a/paddle/phi/api/lib/op_meta_info.cc b/paddle/phi/api/lib/op_meta_info.cc
@@ -514,6 +514,7 @@ OpMetaInfoBuilder& OpMetaInfoBuilder::Attrs(std::vector<std::string>&& attrs) {
       {"bool",
        "int",
        "float",
+       "double",
        "int64_t",
        "std::string",
        "std::vector<int>",
@@ -525,7 +526,7 @@ OpMetaInfoBuilder& OpMetaInfoBuilder::Attrs(std::vector<std::string>&& attrs) {
     if (custom_attrs_type.find(attr_type_str) == custom_attrs_type.end()) {
       PADDLE_THROW(common::errors::Unimplemented(
           "Unsupported `%s` type value as custom attribute now. "
-          "Supported data types include `bool`, `int`, `float`, "
+          "Supported data types include `bool`, `int`, `float`, `double`,  "
           "`int64_t`, `std::string`, `std::vector<int>`, "
           "`std::vector<float>`, `std::vector<int64_t>`, "
           "`std::vector<std::string>`, "
diff --git a/test/custom_op/CMakeLists.txt b/test/custom_op/CMakeLists.txt
@@ -59,6 +59,8 @@ if(WITH_TESTING)
   py_test(test_dispatch_jit SRCS test_dispatch_jit.py)
   py_test(test_multi_out_jit SRCS test_multi_out_jit.py)
   py_test(test_custom_attrs_jit SRCS test_custom_attrs_jit.py)
+  py_test(test_custom_extend_attrs_jit SRCS test_custom_extend_attrs_jit.py)
+  set_tests_properties(test_custom_extend_attrs_jit PROPERTIES TIMEOUT 180)
   py_test(test_custom_concat SRCS test_custom_concat.py)
   set_tests_properties(
     test_custom_concat PROPERTIES ENVIRONMENT
diff --git a/test/custom_op/extend_attr_test_op.cc b/test/custom_op/extend_attr_test_op.cc
@@ -0,0 +1,62 @@
+// 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
+//
+//     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 <cstdlib>
+#include <iostream>
+#include <vector>
+
+#include "paddle/extension.h"
+
+template <typename data_t>
+void assign_cpu_kernel(const data_t* x_data,
+                       data_t* out_data,
+                       int64_t x_numel) {
+  for (int i = 0; i < x_numel; ++i) {
+    out_data[i] = x_data[i];
+  }
+}
+
+void CheckAllForwardAttrs(const double& double_attr) {
+  if (std::abs(double_attr - 3.14) > 1e-10) {
+    throw std::runtime_error("double_attr value error.");
+  }
+}
+
+std::vector<paddle::Tensor> ExtentAttrTestForward(const paddle::Tensor& x,
+                                                  double double_attr) {
+  auto out = paddle::empty_like(x);
+
+  PD_DISPATCH_FLOATING_TYPES(
+      x.type(), "assign_cpu_kernel", ([&] {
+        assign_cpu_kernel<data_t>(
+            x.data<data_t>(), out.mutable_data<data_t>(), x.size());
+      }));
+
+  // Check attrs value
+  CheckAllForwardAttrs(double_attr);
+
+  return {out};
+}
+
+std::vector<std::vector<int64_t>> ExtentAttrTestInferShape(
+    const std::vector<int64_t>& x_shape, double double_attr) {
+  return {x_shape};
+}
+
+PD_BUILD_OP(extend_attr_test)
+    .Inputs({"X"})
+    .Outputs({"Out"})
+    .Attrs({"double_attr: double"})
+    .SetKernelFn(PD_KERNEL(ExtentAttrTestForward))
+    .SetInferShapeFn(PD_INFER_SHAPE(ExtentAttrTestInferShape));
diff --git a/test/custom_op/test_custom_extend_attrs_jit.py b/test/custom_op/test_custom_extend_attrs_jit.py
@@ -0,0 +1,62 @@
+# 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
+#
+#     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.
+
+import os
+import unittest
+
+import numpy as np
+from utils import extra_cc_args, extra_nvcc_args, paddle_includes
+
+import paddle
+from paddle.utils.cpp_extension import get_build_directory, load
+from paddle.utils.cpp_extension.extension_utils import run_cmd
+
+# Because Windows don't use docker, the shared lib already exists in the
+# cache dir, it will not be compiled again unless the shared lib is removed.
+file = f'{get_build_directory()}\\custom_extend_attrs_jit\\custom_extend_attrs_jit.pyd'
+if os.name == 'nt' and os.path.isfile(file):
+    cmd = f'del {file}'
+    run_cmd(cmd, True)
+
+# Compile and load custom op Just-In-Time.
+extend_custom_attrs = load(
+    name='custom_extend_attrs_jit',
+    sources=['extend_attr_test_op.cc'],
+    extra_include_paths=paddle_includes,  # add for Coverage CI
+    extra_cxx_cflags=extra_cc_args,  # test for cflags
+    extra_cuda_cflags=extra_nvcc_args,  # test for cflags
+    verbose=True,
+)
+
+
+class TestJitCustomAttrs(unittest.TestCase):
+    def setUp(self):
+        paddle.set_device('cpu')
+        # prepare test value
+        self.double_attr = 3.14
+
+    def test_func_attr_value(self):
+        x = paddle.ones([2, 2], dtype='float32')
+        x.stop_gradient = False
+        out = extend_custom_attrs.extend_attr_test(
+            x,
+            self.double_attr,
+        )
+        out.stop_gradient = True
+
+        np.testing.assert_array_equal(x.numpy(), out.numpy())
+
+
+if __name__ == '__main__':
+    unittest.main()
