device.py 7.8 KB
Newer Older
1 2 3 4 5 6 7 8 9 10 11 12 13 14
#   Copyright (c) 2020 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.

15
# TODO: define the functions to manipulate devices 
16
import re
T
taixiurong 已提交
17
import os
18 19
from paddle.fluid import core
from paddle.fluid import framework
20
from paddle.fluid.dygraph.parallel import ParallelEnv
21 22
from paddle.fluid.framework import is_compiled_with_cuda  # noqa: F401
from paddle.fluid.framework import is_compiled_with_rocm  # noqa: F401
23

24

25
__all__ = [  # noqa
26
    'get_cudnn_version',
27
    'set_device',
28 29
    'get_device',
    'XPUPlace',
W
Wenyu 已提交
30
    'is_compiled_with_xpu',
31
    'is_compiled_with_cuda',
32
    'is_compiled_with_rocm',
33
    'is_compiled_with_npu'
34 35
]

36 37 38
_cudnn_version = None


39 40
# TODO: WITH_ASCEND_CL may changed to WITH_NPU or others in the future
# for consistent.
41 42
def is_compiled_with_npu():
    """
43
    Whether paddle was built with WITH_ASCEND_CL=ON to support Ascend NPU.
44 45 46 47 48 49 50 51 52 53 54 55

    Returns (bool): `True` if NPU is supported, otherwise `False`.

    Examples:
        .. code-block:: python

            import paddle
            support_npu = paddle.is_compiled_with_npu()
    """
    return core.is_compiled_with_npu()


56 57 58 59 60 61 62 63 64 65
def is_compiled_with_xpu():
    """
    Whether paddle was built with WITH_XPU=ON to support Baidu Kunlun

    Returns (bool): whether paddle was built with WITH_XPU=ON

    Examples:
        .. code-block:: python

            import paddle
66
            support_xpu = paddle.is_compiled_with_xpu()
67 68 69 70 71 72 73 74 75 76 77 78 79
    """
    return core.is_compiled_with_xpu()


def XPUPlace(dev_id):
    """
    Return a Baidu Kunlun Place

    Parameters:
        dev_id(int): Baidu Kunlun device id

    Examples:
        .. code-block:: python
80 81
            # required: xpu
            
82
            import paddle
83
            place = paddle.XPUPlace(0)
84 85 86 87
    """
    return core.XPUPlace(dev_id)


88 89 90 91 92 93 94 95 96 97 98 99 100
def get_cudnn_version():
    """
    This funciton return the version of cudnn. the retuen value is int which represents the 
    cudnn version. For example, if it return 7600, it represents the version of cudnn is 7.6.
    
    Returns:
        int: A int value which represents the cudnn version. If cudnn version is not installed, it return None.

    Examples:
        .. code-block:: python
            
            import paddle

101
            cudnn_version = paddle.get_cudnn_version()
102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118



    """
    global _cudnn_version
    if not core.is_compiled_with_cuda():
        return None
    if _cudnn_version is None:
        cudnn_version = int(core.cudnn_version())
        _cudnn_version = cudnn_version
        if _cudnn_version < 0:
            return None
        else:
            return cudnn_version
    else:
        return _cudnn_version

119

C
chentianyu03 已提交
120
def _convert_to_place(device):
121 122 123
    lower_device = device.lower()
    if lower_device == 'cpu':
        place = core.CPUPlace()
124 125
    elif lower_device == 'gpu':
        if not core.is_compiled_with_cuda():
126 127
            raise ValueError("The device should not be 'gpu', "
                             "since PaddlePaddle is not compiled with CUDA")
128
        place = core.CUDAPlace(ParallelEnv().dev_id)
129 130
    elif lower_device == 'xpu':
        if not core.is_compiled_with_xpu():
131 132
            raise ValueError("The device should not be 'xpu', "
                             "since PaddlePaddle is not compiled with XPU")
T
taixiurong 已提交
133 134 135
        selected_xpus = os.getenv("FLAGS_selected_xpus", "0").split(",")
        device_id = int(selected_xpus[0])
        place = core.XPUPlace(device_id)
H
houj04 已提交
136 137 138 139 140 141 142
    elif lower_device == 'npu':
        if not core.is_compiled_with_npu():
            raise ValueError("The device should not be 'npu', "
                             "since PaddlePaddle is not compiled with NPU")
        selected_npus = os.getenv("FLAGS_selected_npus", "0").split(",")
        device_id = int(selected_npus[0])
        place = core.NPUPlace(device_id)
143
    else:
144 145
        avaliable_gpu_device = re.match(r'gpu:\d+', lower_device)
        avaliable_xpu_device = re.match(r'xpu:\d+', lower_device)
H
houj04 已提交
146 147
        avaliable_npu_device = re.match(r'npu:\d+', lower_device)
        if not avaliable_gpu_device and not avaliable_xpu_device and not avaliable_npu_device:
148
            raise ValueError(
H
houj04 已提交
149
                "The device must be a string which is like 'cpu', 'gpu', 'gpu:x', 'xpu', 'xpu:x', 'npu' or 'npu:x'"
150
            )
151 152 153
        if avaliable_gpu_device:
            if not core.is_compiled_with_cuda():
                raise ValueError(
154
                    "The device should not be {}, since PaddlePaddle is "
155 156 157 158 159 160 161 162
                    "not compiled with CUDA".format(avaliable_gpu_device))
            device_info_list = device.split(':', 1)
            device_id = device_info_list[1]
            device_id = int(device_id)
            place = core.CUDAPlace(device_id)
        if avaliable_xpu_device:
            if not core.is_compiled_with_xpu():
                raise ValueError(
163
                    "The device should not be {}, since PaddlePaddle is "
164 165 166 167 168
                    "not compiled with XPU".format(avaliable_xpu_device))
            device_info_list = device.split(':', 1)
            device_id = device_info_list[1]
            device_id = int(device_id)
            place = core.XPUPlace(device_id)
H
houj04 已提交
169 170 171 172 173 174 175 176 177
        if avaliable_npu_device:
            if not core.is_compiled_with_npu():
                raise ValueError(
                    "The device should not be {}, since PaddlePaddle is "
                    "not compiled with NPU".format(avaliable_npu_device))
            device_info_list = device.split(':', 1)
            device_id = device_info_list[1]
            device_id = int(device_id)
            place = core.NPUPlace(device_id)
C
chentianyu03 已提交
178
    return place
179

C
chentianyu03 已提交
180 181 182

def set_device(device):
    """
H
houj04 已提交
183
    Paddle supports running calculations on various types of devices, including CPU, GPU, XPU and NPU.
C
chentianyu03 已提交
184 185 186 187 188
    They are represented by string identifiers. This function can specify the global device
    which the OP will run.

    Parameters:
        device(str): This parameter determines the specific running device.
H
houj04 已提交
189 190
            It can be ``cpu``, ``gpu``, ``xpu``, ``npu``, ``gpu:x``, ``xpu:x`` and ``npu:x``,
            where ``x`` is the index of the GPUs, XPUs or NPUs.
C
chentianyu03 已提交
191 192 193 194 195 196 197 198 199 200 201 202 203

    Examples:

     .. code-block:: python
            
        import paddle

        paddle.set_device("cpu")
        x1 = paddle.ones(name='x1', shape=[1, 2], dtype='int32')
        x2 = paddle.zeros(name='x2', shape=[1, 2], dtype='int32')
        data = paddle.stack([x1,x2], axis=1)
    """
    place = _convert_to_place(device)
204 205
    framework._set_expected_place(place)
    return place
206 207 208 209 210


def get_device():
    """
    This funciton can get the current global device of the program is running.
H
houj04 已提交
211
    It's a string which is like 'cpu', 'gpu:x', 'xpu:x' and 'npu:x'. if the global device is not
212
    set, it will return a string which is 'gpu:x' when cuda is avaliable or it 
213 214 215 216 217 218 219
    will return a string which is 'cpu' when cuda is not avaliable.

    Examples:

     .. code-block:: python
            
        import paddle
220
        device = paddle.get_device()
221 222 223 224 225 226 227 228 229

    """
    device = ''
    place = framework._current_expected_place()
    if isinstance(place, core.CPUPlace):
        device = 'cpu'
    elif isinstance(place, core.CUDAPlace):
        device_id = place.get_device_id()
        device = 'gpu:' + str(device_id)
230 231 232
    elif isinstance(place, core.XPUPlace):
        device_id = place.get_device_id()
        device = 'xpu:' + str(device_id)
H
houj04 已提交
233 234 235
    elif isinstance(place, core.NPUPlace):
        device_id = place.get_device_id()
        device = 'npu:' + str(device_id)
236 237

    return device