Program Listing for File DevicePool.hpp

↰ Return to documentation for file (pennylane_lightning/core/src/utils/cuda_utils/DevicePool.hpp)

// Copyright 2022-2023 Xanadu Quantum Technologies Inc.

// 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.

#pragma once

#include "cuda.h"
#include <mutex>
#include <unordered_set>
#include <vector>

#include "TSQueue.hpp"
#include "cuError.hpp"

namespace Pennylane::LightningGPU {

template <typename DeviceIndexType = int> class DevicePool {
  public:
    DevicePool() {
        for (std::size_t i = 0; i < DevicePool::getTotalDevices(); i++) {
            available_devices_.push(static_cast<DeviceIndexType>(i));
        }
    }
    virtual ~DevicePool() = default;

    auto getActiveDevices() -> const std::unordered_set<DeviceIndexType> & {
        return active_devices_;
    }

    inline bool isActive(const DeviceIndexType &index) {
        return active_devices_.find(index) != active_devices_.end();
    }
    inline bool isInactive(const DeviceIndexType &index) {
        return !isActive(index);
    }

    static std::size_t getTotalDevices() {
        int result;
        PL_CUDA_IS_SUCCESS(cudaGetDeviceCount(&result));
        return static_cast<std::size_t>(result);
    }

    int acquireDevice() {
        int dev_id;
        available_devices_.wait_and_pop(dev_id);
        {
            std::lock_guard<std::mutex> lg(m_);
            active_devices_.insert(dev_id);
        }
        return dev_id;
    }

    void releaseDevice(DeviceIndexType dev_id) {
        available_devices_.push(dev_id);
        {
            std::lock_guard<std::mutex> lg(m_);
            active_devices_.erase(dev_id);
        }
    }

    void syncDevice() { cudaDeviceSynchronize(); }

    static std::vector<cudaUUID_t> getDeviceUIDs() {
        int deviceCount;
        cudaGetDeviceCount(&deviceCount);
        int device;
        std::vector<cudaUUID_t> dev_uid(deviceCount);
        for (device = 0; device < deviceCount; device++) {
            cudaDeviceProp deviceProp;
            cudaGetDeviceProperties(&deviceProp, device);
            dev_uid[device] = deviceProp.uuid;
        }
        return dev_uid;
    }

    static void setDeviceIdx(int device_id) {
        PL_CUDA_IS_SUCCESS(cudaSetDevice(device_id));
    }

  private:
    std::unordered_set<DeviceIndexType> active_devices_;
    std::mutex m_;
    TSQueue<DeviceIndexType> available_devices_;
};

} // namespace Pennylane::LightningGPU

api/program_listing_file_pennylane_lightning_core_src_utils_cuda_utils_DevicePool.hpp

 

Download Python script

 

Download Notebook

 

View on GitHub