[AMP][Custom Device] add fp16 op detection for custom device (#56053)

* [AMP] add fp16 op detection for custom device * resolve conflicts
PaddlePaddle · Aug 14, 2023 · 9d40da3 · 9d40da3
1 parent 066097e
commit 9d40da3
Show file tree

Hide file tree

Showing 2 changed files with 10 additions and 7 deletions.
diff --git a/paddle/fluid/imperative/amp_auto_cast.cc b/paddle/fluid/imperative/amp_auto_cast.cc
@@ -52,13 +52,16 @@ OpSupportedInfos(const std::string& place,
       {"CPU", &platform::is_cpu_place},
       {"XPU", &platform::is_xpu_place},
       {"CUSTOM_DEVICE", &platform::is_custom_place},
+#ifdef PADDLE_WITH_CUSTOM_DEVICE
+      {query_place, &platform::is_custom_place},
+#endif
   };
-  PADDLE_ENFORCE_NE(
-      is_target_place.count(query_place),
-      0,
-      platform::errors::InvalidArgument(
-          "The argument `place` should be 'GPU', 'CPU', 'XPU', but got '%s'.",
-          place));
+  PADDLE_ENFORCE_NE(is_target_place.count(query_place),
+                    0,
+                    platform::errors::InvalidArgument(
+                        "The argument `place` should be 'GPU', 'CPU', 'XPU' or "
+                        "other Custom Device, but got '%s'.",
+                        place));
 
   std::unordered_set<std::string> all_ops;
   const auto& op_info = framework::OpInfoMap::Instance().map();

diff --git a/python/paddle/static/amp/fp16_lists.py b/python/paddle/static/amp/fp16_lists.py
@@ -98,7 +98,7 @@ def _get_sys_unsupported_list(dtype):
     elif isinstance(
         paddle.framework._current_expected_place(), paddle.CustomPlace
     ):
-        device = 'CUSTOM_DEVICE'
+        device = paddle.framework._current_expected_place().get_device_type()
     else:
         device = 'GPU'
     all_ops, _, sys_unsupported_list = core.op_supported_infos(device, var_type)