From cf99c0d5f8f03e95fd65aff1c65ddf74c4d061ef Mon Sep 17 00:00:00 2001
From: Linjie Chen <40840292+linjieccc@users.noreply.github.com>
Date: Mon, 23 Aug 2021 11:00:30 +0800
Subject: [PATCH] Add cuda.device_count api (#34811)

* Add cuda device count api

* update coda format

* fix unittest error

* update code format

* update comment
---
 python/paddle/device/cuda/__init__.py         | 23 ++++++++++++++++
 .../tests/unittests/test_cuda_device_count.py | 26 +++++++++++++++++++
 2 files changed, 49 insertions(+)
 create mode 100644 python/paddle/fluid/tests/unittests/test_cuda_device_count.py

diff --git a/python/paddle/device/cuda/__init__.py b/python/paddle/device/cuda/__init__.py
index 8b5879fe9a..834cda71fd 100644
--- a/python/paddle/device/cuda/__init__.py
+++ b/python/paddle/device/cuda/__init__.py
@@ -22,6 +22,7 @@ __all__ = [
     'Event',
     'current_stream',
     'synchronize',
+    'device_count',
 ]
 
 
@@ -94,3 +95,25 @@ def synchronize(device=None):
             raise ValueError("device type must be int or paddle.CUDAPlace")
 
     return core._device_synchronize(device_id)
+
+
+def device_count():
+    '''
+    Return the number of GPUs available.
+    
+    Returns:
+        int: the number of GPUs available.
+
+    Examples:
+        .. code-block:: python
+
+            import paddle
+
+            paddle.device.cuda.device_count()
+
+    '''
+
+    num_gpus = core.get_cuda_device_count() if hasattr(
+        core, 'get_cuda_device_count') else 0
+
+    return num_gpus
diff --git a/python/paddle/fluid/tests/unittests/test_cuda_device_count.py b/python/paddle/fluid/tests/unittests/test_cuda_device_count.py
new file mode 100644
index 0000000000..f4114c9d45
--- /dev/null
+++ b/python/paddle/fluid/tests/unittests/test_cuda_device_count.py
@@ -0,0 +1,26 @@
+# Copyright (c) 2021 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.
+
+import paddle
+import unittest
+
+
+class TestDeviceCount(unittest.TestCase):
+    def test_device_count(self):
+        s = paddle.device.cuda.device_count()
+        self.assertIsNotNone(s)
+
+
+if __name__ == "__main__":
+    unittest.main()
-- 
GitLab