| # Licensed to the Apache Software Foundation (ASF) under one |
| # or more contributor license agreements. See the NOTICE file |
| # distributed with this work for additional information |
| # regarding copyright ownership. The ASF licenses this file |
| # to you 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. |
| # ============================================================================= |
| ''' |
| This script includes Device class and its subclasses for python users |
| to call singa::Device and its methods. |
| |
| TODO(wangwei) implement py CudaGPU class. |
| ''' |
| |
| # from builtins import object |
| from . import singa_wrap as singa |
| |
| |
| def get_num_gpus(): |
| assert singa.USE_CUDA, 'SINGA has not been compiled with CUDA enabled.' |
| return singa.Platform.GetNumGPUs() |
| |
| |
| def get_gpu_ids(): |
| assert singa.USE_CUDA, 'SINGA has not been compiled with CUDA enabled.' |
| return singa.Platform.GetGPUIDs() |
| |
| |
| def get_gpu_mem_size(id): |
| assert singa.USE_CUDA, 'SINGA has not been compiled with CUDA enabled.' |
| return singa.Platform.GetGPUMemSize(id) |
| |
| |
| def get_num_opencl_platforms(): |
| assert singa.USE_OPENCL, 'SINGA has not been compiled with OpenCL enabled.' |
| return singa.Platform.GetNumOpenclPlatforms() |
| |
| |
| def get_num_opencl_devices(): |
| assert singa.USE_OPENCL, 'SINGA has not been compiled with OpenCL enabled.' |
| return singa.Platform.GetNumOpenclDevices() |
| |
| |
| def device_query(id, verbose=False): |
| assert singa.USE_CUDA, 'SINGA has not been compiled with CUDA enabled.' |
| return singa.Platform.DeviceQuery(id, verbose) |
| |
| |
| def create_cpu_device(): |
| '''Create the default CPU device. |
| |
| Returns: |
| a swig converted CPU device. |
| ''' |
| return singa.Platform.GetDefaultDevice() |
| |
| |
| def create_cuda_gpus(num): |
| '''Create a list of CudaGPU devices. |
| |
| Args: |
| num (int): number of device to create. |
| Returns: |
| a list of swig converted CudaGPU devices. |
| ''' |
| assert singa.USE_CUDA, 'SINGA has not been compiled with CUDA enabled.' |
| return singa.Platform.CreateCudaGPUs(num) |
| |
| |
| def create_cuda_gpu(): |
| '''Create a single CudaGPU device. |
| |
| Returns: |
| a swig converted CudaGPU device. |
| ''' |
| assert singa.USE_CUDA, 'SINGA has not been compiled with CUDA enabled.' |
| return create_cuda_gpu_on(0) |
| |
| |
| def create_cuda_gpus_on(device_ids): |
| '''Create a list of CudaGPU devices. |
| |
| Args: |
| device_ids (list): a list of GPU card IDs. |
| |
| Returns: |
| a list of swig converted CudaGPU devices. |
| ''' |
| assert singa.USE_CUDA, 'SINGA has not been compiled with CUDA enabled.' |
| return singa.Platform.CreateCudaGPUsOn(device_ids) |
| |
| |
| def create_cuda_gpu_on(device_id): |
| '''Create a CudaGPU device on the given device ID. |
| |
| Args: |
| device_id (int): GPU card ID. |
| |
| Returns: |
| a swig converted CudaGPU device. |
| ''' |
| assert singa.USE_CUDA, 'SINGA has not been compiled with CUDA enabled.' |
| devices = create_cuda_gpus_on([device_id]) |
| return devices[0] |
| |
| |
| def create_opencl_device(): |
| '''Create the default OpenCL device. |
| |
| Returns: |
| a swig converted OpenCL device. |
| ''' |
| assert singa.USE_OPENCL, 'SINGA has not been compiled with OpenCL enabled.' |
| return singa.Platform.GetDefaultOpenclDevice() |
| |
| |
| default_device = singa.Platform.GetDefaultDevice() |
| |
| |
| def get_default_device(): |
| '''Get the default host device which is a CppCPU device''' |
| return default_device |
| |
| |
| def enable_lazy_alloc(enable): |
| singa.Device.EnableLazyAlloc(enable) |