From 6a5b7e5987d4454123ae1a5d400b880db41b0855 Mon Sep 17 00:00:00 2001
From: Qi Li <qili93@qq.com>
Date: Thu, 27 May 2021 16:28:18 +0800
Subject: [PATCH] [ROCM] add is_compiled_with_rocm api, test=develop (#33043)

---
 python/paddle/__init__.py                        |  2 ++
 python/paddle/device.py                          |  2 ++
 python/paddle/fluid/framework.py                 | 16 ++++++++++++++++
 .../paddle/utils/cpp_extension/cpp_extension.py  |  4 ++--
 4 files changed, 22 insertions(+), 2 deletions(-)

diff --git a/python/paddle/__init__.py b/python/paddle/__init__.py
index ee4dcaa897..7bac330376 100755
--- a/python/paddle/__init__.py
+++ b/python/paddle/__init__.py
@@ -258,6 +258,7 @@ from .device import get_cudnn_version  # noqa: F401
 from .device import set_device  # noqa: F401
 from .device import get_device  # noqa: F401
 from .fluid.framework import is_compiled_with_cuda  # noqa: F401
+from .fluid.framework import is_compiled_with_rocm  # noqa: F401
 from .device import is_compiled_with_xpu  # noqa: F401
 from .device import is_compiled_with_npu  # noqa: F401
 from .device import XPUPlace  # noqa: F401
@@ -384,6 +385,7 @@ __all__ = [     #noqa
            'less_equal',
            'triu',
            'is_compiled_with_cuda',
+           'is_compiled_with_rocm',
            'sin',
            'dist',
            'unbind',
diff --git a/python/paddle/device.py b/python/paddle/device.py
index 803d54e11b..fce01d0d67 100644
--- a/python/paddle/device.py
+++ b/python/paddle/device.py
@@ -19,6 +19,7 @@ from paddle.fluid import core
 from paddle.fluid import framework
 from paddle.fluid.dygraph.parallel import ParallelEnv
 from paddle.fluid.framework import is_compiled_with_cuda  #DEFINE_ALIAS
+from paddle.fluid.framework import is_compiled_with_rocm  #DEFINE_ALIAS
 
 __all__ = [
     'get_cudnn_version',
@@ -33,6 +34,7 @@ __all__ = [
     #            'CUDAPinnedPlace',
     #            'CUDAPlace',
     'is_compiled_with_cuda',
+    'is_compiled_with_rocm',
     'is_compiled_with_npu'
 ]
 
diff --git a/python/paddle/fluid/framework.py b/python/paddle/fluid/framework.py
index a858ba7834..bffeaf2c6c 100644
--- a/python/paddle/fluid/framework.py
+++ b/python/paddle/fluid/framework.py
@@ -53,6 +53,7 @@ __all__ = [
     'cuda_pinned_places',
     'in_dygraph_mode',
     'is_compiled_with_cuda',
+    'is_compiled_with_rocm',
     'is_compiled_with_xpu',
     'Variable',
     'require_version',
@@ -398,6 +399,21 @@ def is_compiled_with_cuda():
     return core.is_compiled_with_cuda()
 
 
+def is_compiled_with_rocm():
+    """
+    Whether this whl package can be used to run the model on AMD or Hygon GPU(ROCm).
+
+    Returns (bool): `True` if ROCm is currently available, otherwise `False`.
+
+    Examples:
+        .. code-block:: python
+
+            import paddle
+            support_gpu = paddle.is_compiled_with_rocm()
+    """
+    return core.is_compiled_with_rocm()
+
+
 def cuda_places(device_ids=None):
     """
     **Note**:
diff --git a/python/paddle/utils/cpp_extension/cpp_extension.py b/python/paddle/utils/cpp_extension/cpp_extension.py
index 6045ac7d1e..da3e838915 100644
--- a/python/paddle/utils/cpp_extension/cpp_extension.py
+++ b/python/paddle/utils/cpp_extension/cpp_extension.py
@@ -42,10 +42,10 @@ if IS_WINDOWS and six.PY3:
     from unittest.mock import Mock
     _du_build_ext.get_export_symbols = Mock(return_value=None)
 
+CUDA_HOME = find_cuda_home()
 if core.is_compiled_with_rocm():
     ROCM_HOME = find_rocm_home()
-else:
-    CUDA_HOME = find_cuda_home()
+    CUDA_HOME = ROCM_HOME
 
 
 def setup(**attr):
-- 
GitLab