From 3fe99ad5c1fcd5775945ab56a329572860c66330 Mon Sep 17 00:00:00 2001
From: Qi Li <qili93@qq.com>
Date: Tue, 1 Jun 2021 23:43:15 +0800
Subject: [PATCH] [ROCM] add is_compiled_with_rocm api, test=develop (#33043)
 (#33228)

---
 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 ee4dcaa8979..7bac330376c 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 035d240e713..85b813a7f51 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 3ca16b66675..bc8a06cb1ed 100644
--- a/python/paddle/fluid/framework.py
+++ b/python/paddle/fluid/framework.py
@@ -52,6 +52,7 @@ __all__ = [
     'cuda_pinned_places',
     'in_dygraph_mode',
     'is_compiled_with_cuda',
+    'is_compiled_with_rocm',
     'is_compiled_with_xpu',
     'Variable',
     'require_version',
@@ -397,6 +398,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 7d6fae3ad77..dcaa1ca15e5 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