Team Ai
Datasetpublic

codekingpro/portable-devtools

sourceHugging Faceupdated 5mo agoView on Hugging Face
1likes14kdownloads
util_allocator.cuh881 linesDownload Raw Back to cub
1/******************************************************************************2 * Copyright (c) 2011, Duane Merrill.  All rights reserved.3 * Copyright (c) 2011-2018, NVIDIA CORPORATION.  All rights reserved.4 *5 * Redistribution and use in source and binary forms, with or without6 * modification, are permitted provided that the following conditions are met:7 *     * Redistributions of source code must retain the above copyright8 *       notice, this list of conditions and the following disclaimer.9 *     * Redistributions in binary form must reproduce the above copyright10 *       notice, this list of conditions and the following disclaimer in the11 *       documentation and/or other materials provided with the distribution.12 *     * Neither the name of the NVIDIA CORPORATION nor the13 *       names of its contributors may be used to endorse or promote products14 *       derived from this software without specific prior written permission.15 *16 * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" AND17 * ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE IMPLIED18 * WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE19 * DISCLAIMED. IN NO EVENT SHALL NVIDIA CORPORATION BE LIABLE FOR ANY20 * DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES21 * (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES;22 * LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND23 * ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT24 * (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE OF THIS25 * SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.26 *27 ******************************************************************************/28 29/******************************************************************************30 * Simple caching allocator for device memory allocations. The allocator is31 * thread-safe and capable of managing device allocations on multiple devices.32 ******************************************************************************/33 34#pragma once35 36#include <cub/config.cuh>37 38#if defined(_CCCL_IMPLICIT_SYSTEM_HEADER_GCC)39#  pragma GCC system_header40#elif defined(_CCCL_IMPLICIT_SYSTEM_HEADER_CLANG)41#  pragma clang system_header42#elif defined(_CCCL_IMPLICIT_SYSTEM_HEADER_MSVC)43#  pragma system_header44#endif // no system header45 46#include <cub/util_debug.cuh>47#include <cub/util_namespace.cuh>48 49#include <map>50#include <mutex>51#include <set>52 53#include <math.h>54 55CUB_NAMESPACE_BEGIN56 57 58/**59 * @addtogroup UtilMgmt60 * @{61 */62 63 64/******************************************************************************65 * CachingDeviceAllocator (host use)66 ******************************************************************************/67 68/**69 * @brief A simple caching allocator for device memory allocations.70 *71 * @par Overview72 * The allocator is thread-safe and stream-safe and is capable of managing cached73 * device allocations on multiple devices.  It behaves as follows:74 *75 * @par76 * - Allocations from the allocator are associated with an @p active_stream. Once freed,77 *   the allocation becomes available immediately for reuse within the @p active_stream78 *   with which it was associated with during allocation, and it becomes available for79 *   reuse within other streams when all prior work submitted to @p active_stream has completed.80 * - Allocations are categorized and cached by bin size. A new allocation request of81 *   a given size will only consider cached allocations within the corresponding bin.82 * - Bin limits progress geometrically in accordance with the growth factor83 *   @p bin_growth provided during construction. Unused device allocations within84 *   a larger bin cache are not reused for allocation requests that categorize to85 *   smaller bin sizes.86 * - Allocation requests below ( @p bin_growth ^ @p min_bin ) are rounded up to87 *   ( @p bin_growth ^ @p min_bin ).88 * - Allocations above ( @p bin_growth ^ @p max_bin ) are not rounded up to the nearest89 *   bin and are simply freed when they are deallocated instead of being returned90 *   to a bin-cache.91 * - If the total storage of cached allocations on a given device will exceed92 *   @p max_cached_bytes, allocations for that device are simply freed when they are93 *   deallocated instead of being returned to their bin-cache.94 *95 * @par96 * For example, the default-constructed CachingDeviceAllocator is configured with:97 * - @p bin_growth          = 898 * - @p min_bin             = 399 * - @p max_bin             = 7100 * - @p max_cached_bytes    = 6MB - 1B101 *102 * @par103 * which delineates five bin-sizes: 512B, 4KB, 32KB, 256KB, and 2MB104 * and sets a maximum of 6,291,455 cached bytes per device105 *106 */107struct CachingDeviceAllocator108{109 110    //---------------------------------------------------------------------111    // Constants112    //---------------------------------------------------------------------113 114    /// Out-of-bounds bin115    static constexpr unsigned int INVALID_BIN = (unsigned int) -1;116 117    /// Invalid size118    static constexpr size_t INVALID_SIZE = (size_t) -1;119 120#ifndef DOXYGEN_SHOULD_SKIP_THIS    // Do not document121 122    /// Invalid device ordinal123    static constexpr int INVALID_DEVICE_ORDINAL = -1;124 125    //---------------------------------------------------------------------126    // Type definitions and helper types127    //---------------------------------------------------------------------128 129    /**130     * Descriptor for device memory allocations131     */132    struct BlockDescriptor133    {134        // Device pointer135        void *d_ptr;136 137        // Size of allocation in bytes138        size_t bytes;139 140        // Bin enumeration141        unsigned int bin;142 143        // device ordinal144        int device;145 146        // Associated associated_stream147        cudaStream_t associated_stream;148 149        // Signal when associated stream has run to the point at which this block was freed150        cudaEvent_t ready_event;151 152        // Constructor (suitable for searching maps for a specific block, given its pointer and153        // device)154        BlockDescriptor(void *d_ptr, int device)155            : d_ptr(d_ptr)156            , bytes(0)157            , bin(INVALID_BIN)158            , device(device)159            , associated_stream(0)160            , ready_event(0)161        {}162 163        // Constructor (suitable for searching maps for a range of suitable blocks, given a device)164        BlockDescriptor(int device)165            : d_ptr(NULL)166            , bytes(0)167            , bin(INVALID_BIN)168            , device(device)169            , associated_stream(0)170            , ready_event(0)171        {}172 173        // Comparison functor for comparing device pointers174        static bool PtrCompare(const BlockDescriptor &a, const BlockDescriptor &b)175        {176          if (a.device == b.device)177            return (a.d_ptr < b.d_ptr);178          else179            return (a.device < b.device);180        }181 182        // Comparison functor for comparing allocation sizes183        static bool SizeCompare(const BlockDescriptor &a, const BlockDescriptor &b)184        {185            if (a.device == b.device)186                return (a.bytes < b.bytes);187            else188                return (a.device < b.device);189        }190    };191 192    /// BlockDescriptor comparator function interface193    typedef bool (*Compare)(const BlockDescriptor &, const BlockDescriptor &);194 195    class TotalBytes {196    public:197        size_t free;198        size_t live;199        TotalBytes() { free = live = 0; }200    };201 202    /// Set type for cached blocks (ordered by size)203    typedef std::multiset<BlockDescriptor, Compare> CachedBlocks;204 205    /// Set type for live blocks (ordered by ptr)206    typedef std::multiset<BlockDescriptor, Compare> BusyBlocks;207 208    /// Map type of device ordinals to the number of cached bytes cached by each device209    typedef std::map<int, TotalBytes> GpuCachedBytes;210 211 212    //---------------------------------------------------------------------213    // Utility functions214    //---------------------------------------------------------------------215 216    /**217     * Integer pow function for unsigned base and exponent218     */219    static unsigned int IntPow(220        unsigned int base,221        unsigned int exp)222    {223        unsigned int retval = 1;224        while (exp > 0)225        {226            if (exp & 1) {227                retval = retval * base;        // multiply the result by the current base228            }229            base = base * base;                // square the base230            exp = exp >> 1;                    // divide the exponent in half231        }232        return retval;233    }234 235 236    /**237     * Round up to the nearest power-of238     */239    void NearestPowerOf(240        unsigned int    &power,241        size_t          &rounded_bytes,242        unsigned int    base,243        size_t          value)244    {245        power = 0;246        rounded_bytes = 1;247 248        if (value * base < value)249        {250            // Overflow251            power = sizeof(size_t) * 8;252            rounded_bytes = size_t(0) - 1;253            return;254        }255 256        while (rounded_bytes < value)257        {258            rounded_bytes *= base;259            power++;260        }261    }262 263    //---------------------------------------------------------------------264    // Fields265    //---------------------------------------------------------------------266 267    /// Mutex for thread-safety268    std::mutex mutex;269 270    /// Geometric growth factor for bin-sizes271    unsigned int bin_growth;272 273    /// Minimum bin enumeration274    unsigned int min_bin;275 276    /// Maximum bin enumeration277    unsigned int max_bin;278 279    /// Minimum bin size280    size_t min_bin_bytes;281 282    /// Maximum bin size283    size_t max_bin_bytes;284 285    /// Maximum aggregate cached bytes per device286    size_t max_cached_bytes;287 288    /// Whether or not to skip a call to FreeAllCached() when destructor is called.289    /// (The CUDA runtime may have already shut down for statically declared allocators)290    const bool skip_cleanup;291 292    /// Whether or not to print (de)allocation events to stdout293    bool debug;294 295    /// Map of device ordinal to aggregate cached bytes on that device296    GpuCachedBytes cached_bytes;297 298    /// Set of cached device allocations available for reuse299    CachedBlocks cached_blocks;300 301    /// Set of live device allocations currently in use302    BusyBlocks live_blocks;303 304#endif // DOXYGEN_SHOULD_SKIP_THIS305 306    //---------------------------------------------------------------------307    // Methods308    //---------------------------------------------------------------------309 310    /**311     * @brief Constructor.312     *313     * @param bin_growth314     *   Geometric growth factor for bin-sizes315     *316     * @param min_bin317     *   Minimum bin (default is bin_growth ^ 1)318     *319     * @param max_bin320     *   Maximum bin (default is no max bin)321     *322     * @param max_cached_bytes323     *   Maximum aggregate cached bytes per device (default is no limit)324     *325     * @param skip_cleanup326     *   Whether or not to skip a call to @p FreeAllCached() when the destructor is called (default327     *   is to deallocate)328     *329     * @param debug330     *   Whether or not to print (de)allocation events to stdout (default is no stderr output)331     */332    CachingDeviceAllocator(unsigned int bin_growth,333                           unsigned int min_bin    = 1,334                           unsigned int max_bin    = INVALID_BIN,335                           size_t max_cached_bytes = INVALID_SIZE,336                           bool skip_cleanup       = false,337                           bool debug              = false)338        : bin_growth(bin_growth)339        , min_bin(min_bin)340        , max_bin(max_bin)341        , min_bin_bytes(IntPow(bin_growth, min_bin))342        , max_bin_bytes(IntPow(bin_growth, max_bin))343        , max_cached_bytes(max_cached_bytes)344        , skip_cleanup(skip_cleanup)345        , debug(debug)346        , cached_blocks(BlockDescriptor::SizeCompare)347        , live_blocks(BlockDescriptor::PtrCompare)348    {}349 350 351    /**352     * @brief Default constructor.353     *354     * Configured with:355     * @par356     * - @p bin_growth          = 8357     * - @p min_bin             = 3358     * - @p max_bin             = 7359     * - @p max_cached_bytes    = ( @p bin_growth ^ @p max_bin) * 3 ) - 1 = 6,291,455 bytes360     *361     * which delineates five bin-sizes: 512B, 4KB, 32KB, 256KB, and 2MB and362     * sets a maximum of 6,291,455 cached bytes per device363     */364    CachingDeviceAllocator(365        bool skip_cleanup = false,366        bool debug = false)367    :368        bin_growth(8),369        min_bin(3),370        max_bin(7),371        min_bin_bytes(IntPow(bin_growth, min_bin)),372        max_bin_bytes(IntPow(bin_growth, max_bin)),373        max_cached_bytes((max_bin_bytes * 3) - 1),374        skip_cleanup(skip_cleanup),375        debug(debug),376        cached_blocks(BlockDescriptor::SizeCompare),377        live_blocks(BlockDescriptor::PtrCompare)378    {}379 380 381    /**382     * @brief Sets the limit on the number bytes this allocator is allowed to cache per device.383     *384     * Changing the ceiling of cached bytes does not cause any allocations (in-use or385     * cached-in-reserve) to be freed.  See \p FreeAllCached().386     */387    cudaError_t SetMaxCachedBytes(size_t max_cached_bytes_)388    {389        // Lock390        mutex.lock();391 392        if (debug) _CubLog("Changing max_cached_bytes (%lld -> %lld)\n", (long long) this->max_cached_bytes, (long long) max_cached_bytes_);393 394        this->max_cached_bytes = max_cached_bytes_;395 396        // Unlock397        mutex.unlock();398 399        return cudaSuccess;400    }401 402    /**403     * @brief Provides a suitable allocation of device memory for the given size on the specified404     *        device.405     *406     * Once freed, the allocation becomes available immediately for reuse within the @p407     * active_stream with which it was associated with during allocation, and it becomes available408     * for reuse within other streams when all prior work submitted to @p active_stream has409     * completed.410     *411     * @param[in] device412     *   Device on which to place the allocation413     *414     * @param[out] d_ptr415     *   Reference to pointer to the allocation416     *417     * @param[in] bytes418     *   Minimum number of bytes for the allocation419     *420     * @param[in] active_stream421     *   The stream to be associated with this allocation422     */423    cudaError_t424    DeviceAllocate(int device, void **d_ptr, size_t bytes, cudaStream_t active_stream = 0)425    {426        *d_ptr                          = NULL;427        int entrypoint_device           = INVALID_DEVICE_ORDINAL;428        cudaError_t error               = cudaSuccess;429 430        if (device == INVALID_DEVICE_ORDINAL)431        {432            error = CubDebug(cudaGetDevice(&entrypoint_device));433            if (cudaSuccess != error)434            {435                return error;436            }437 438            device = entrypoint_device;439        }440 441        // Create a block descriptor for the requested allocation442        bool found = false;443        BlockDescriptor search_key(device);444        search_key.associated_stream = active_stream;445        NearestPowerOf(search_key.bin, search_key.bytes, bin_growth, bytes);446 447        if (search_key.bin > max_bin)448        {449            // Bin is greater than our maximum bin: allocate the request450            // exactly and give out-of-bounds bin.  It will not be cached451            // for reuse when returned.452            search_key.bin      = INVALID_BIN;453            search_key.bytes    = bytes;454        }455        else456        {457            // Search for a suitable cached allocation: lock458            mutex.lock();459 460            if (search_key.bin < min_bin)461            {462                // Bin is less than minimum bin: round up463                search_key.bin      = min_bin;464                search_key.bytes    = min_bin_bytes;465            }466 467            // Iterate through the range of cached blocks on the same device in the same bin468            CachedBlocks::iterator block_itr = cached_blocks.lower_bound(search_key);469            while ((block_itr != cached_blocks.end())470                    && (block_itr->device == device)471                    && (block_itr->bin == search_key.bin))472            {473                // To prevent races with reusing blocks returned by the host but still474                // in use by the device, only consider cached blocks that are475                // either (from the active stream) or (from an idle stream)476                bool is_reusable = false;477                if (active_stream == block_itr->associated_stream)478                {479                    is_reusable = true;480                }481                else482                {483                    const cudaError_t event_status = cudaEventQuery(block_itr->ready_event);484                    if(event_status != cudaErrorNotReady)485                    {486                        CubDebug(event_status);487                        is_reusable = true;488                    }489                }490 491                if(is_reusable)492                {493                    // Reuse existing cache block.  Insert into live blocks.494                    found = true;495                    search_key = *block_itr;496                    search_key.associated_stream = active_stream;497                    live_blocks.insert(search_key);498 499                    // Remove from free blocks500                    cached_bytes[device].free -= search_key.bytes;501                    cached_bytes[device].live += search_key.bytes;502 503                    if (debug) _CubLog("\tDevice %d reused cached block at %p (%lld bytes) for stream %lld (previously associated with stream %lld).\n",504                        device, search_key.d_ptr, (long long) search_key.bytes, (long long) search_key.associated_stream, (long long)  block_itr->associated_stream);505 506                    cached_blocks.erase(block_itr);507 508                    break;509                }510                block_itr++;511            }512 513            // Done searching: unlock514            mutex.unlock();515        }516 517        // Allocate the block if necessary518        if (!found)519        {520            // Set runtime's current device to specified device (entrypoint may not be set)521            if (device != entrypoint_device)522            {523                error = CubDebug(cudaGetDevice(&entrypoint_device));524                if (cudaSuccess != error)525                {526                    return error;527                }528 529                error = CubDebug(cudaSetDevice(device));530                if (cudaSuccess != error)531                {532                    return error;533                }534            }535 536            // Attempt to allocate537            error = CubDebug(cudaMalloc(&search_key.d_ptr, search_key.bytes));538            if (error == cudaErrorMemoryAllocation)539            {540                // The allocation attempt failed: free all cached blocks on device and retry541                if (debug) _CubLog("\tDevice %d failed to allocate %lld bytes for stream %lld, retrying after freeing cached allocations",542                      device, (long long) search_key.bytes, (long long) search_key.associated_stream);543 544                error = cudaSuccess;    // Reset the error we will return545                cudaGetLastError();     // Reset CUDART's error546 547                // Lock548                mutex.lock();549 550                // Iterate the range of free blocks on the same device551                BlockDescriptor free_key(device);552                CachedBlocks::iterator block_itr = cached_blocks.lower_bound(free_key);553 554                while ((block_itr != cached_blocks.end()) && (block_itr->device == device))555                {556                    // No need to worry about synchronization with the device: cudaFree is557                    // blocking and will synchronize across all kernels executing558                    // on the current device559 560                    // Free device memory and destroy stream event.561                    error = CubDebug(cudaFree(block_itr->d_ptr));562                    if (cudaSuccess != error)563                    {564                        break;565                    }566 567                    error = CubDebug(cudaEventDestroy(block_itr->ready_event));568                    if (cudaSuccess != error)569                    {570                        break;571                    }572 573                    // Reduce balance and erase entry574                    cached_bytes[device].free -= block_itr->bytes;575 576                    if (debug) _CubLog("\tDevice %d freed %lld bytes.\n\t\t  %lld available blocks cached (%lld bytes), %lld live blocks (%lld bytes) outstanding.\n",577                        device, (long long) block_itr->bytes, (long long) cached_blocks.size(), (long long) cached_bytes[device].free, (long long) live_blocks.size(), (long long) cached_bytes[device].live);578 579                    block_itr = cached_blocks.erase(block_itr);580                }581 582                // Unlock583                mutex.unlock();584 585                // Return under error586                if (error) return error;587 588                // Try to allocate again589                error = CubDebug(cudaMalloc(&search_key.d_ptr, search_key.bytes));590                if (cudaSuccess != error)591                {592                    return error;593                }594            }595 596            // Create ready event597            error =598              CubDebug(cudaEventCreateWithFlags(&search_key.ready_event, cudaEventDisableTiming));599 600            if (cudaSuccess != error)601            {602                return error;603            }604 605            // Insert into live blocks606            mutex.lock();607            live_blocks.insert(search_key);608            cached_bytes[device].live += search_key.bytes;609            mutex.unlock();610 611            if (debug) _CubLog("\tDevice %d allocated new device block at %p (%lld bytes associated with stream %lld).\n",612                      device, search_key.d_ptr, (long long) search_key.bytes, (long long) search_key.associated_stream);613 614            // Attempt to revert back to previous device if necessary615            if ((entrypoint_device != INVALID_DEVICE_ORDINAL) && (entrypoint_device != device))616            {617                error = CubDebug(cudaSetDevice(entrypoint_device));618                if (cudaSuccess != error)619                {620                    return error;621                }622            }623        }624 625        // Copy device pointer to output parameter626        *d_ptr = search_key.d_ptr;627 628        if (debug) _CubLog("\t\t%lld available blocks cached (%lld bytes), %lld live blocks outstanding(%lld bytes).\n",629            (long long) cached_blocks.size(), (long long) cached_bytes[device].free, (long long) live_blocks.size(), (long long) cached_bytes[device].live);630 631        return error;632    }633 634    /**635     * @brief Provides a suitable allocation of device memory for the given size on the current636     *        device.637     *638     * Once freed, the allocation becomes available immediately for reuse within the @p639     * active_stream with which it was associated with during allocation, and it becomes available640     * for reuse within other streams when all prior work submitted to @p active_stream has641     * completed.642     *643     * @param[out] d_ptr644     *   Reference to pointer to the allocation645     *646     * @param[in] bytes647     *   Minimum number of bytes for the allocation648     *649     * @param[in] active_stream650     *   The stream to be associated with this allocation651     */652    cudaError_t DeviceAllocate(void **d_ptr, size_t bytes, cudaStream_t active_stream = 0)653    {654        return DeviceAllocate(INVALID_DEVICE_ORDINAL, d_ptr, bytes, active_stream);655    }656 657    /**658     * @brief Frees a live allocation of device memory on the specified device, returning it to the659     *        allocator.660     *661     * Once freed, the allocation becomes available immediately for reuse within the662     * @p active_stream with which it was associated with during allocation, and it becomes663     * available for reuse within other streams when all prior work submitted to @p active_stream664     * has completed.665     */666    cudaError_t DeviceFree(667        int             device,668        void*           d_ptr)669    {670        int entrypoint_device           = INVALID_DEVICE_ORDINAL;671        cudaError_t error               = cudaSuccess;672 673        if (device == INVALID_DEVICE_ORDINAL)674        {675            error = CubDebug(cudaGetDevice(&entrypoint_device));676            if (cudaSuccess != error)677            {678                return error;679            }680            device = entrypoint_device;681        }682 683        // Lock684        mutex.lock();685 686        // Find corresponding block descriptor687        bool recached = false;688        BlockDescriptor search_key(d_ptr, device);689        BusyBlocks::iterator block_itr = live_blocks.find(search_key);690        if (block_itr != live_blocks.end())691        {692            // Remove from live blocks693            search_key = *block_itr;694            live_blocks.erase(block_itr);695            cached_bytes[device].live -= search_key.bytes;696 697            // Keep the returned allocation if bin is valid and we won't exceed the max cached threshold698            if ((search_key.bin != INVALID_BIN) && (cached_bytes[device].free + search_key.bytes <= max_cached_bytes))699            {700                // Insert returned allocation into free blocks701                recached = true;702                cached_blocks.insert(search_key);703                cached_bytes[device].free += search_key.bytes;704 705                if (debug) _CubLog("\tDevice %d returned %lld bytes from associated stream %lld.\n\t\t %lld available blocks cached (%lld bytes), %lld live blocks outstanding. (%lld bytes)\n",706                    device, (long long) search_key.bytes, (long long) search_key.associated_stream, (long long) cached_blocks.size(),707                    (long long) cached_bytes[device].free, (long long) live_blocks.size(), (long long) cached_bytes[device].live);708            }709        }710 711        // Unlock712        mutex.unlock();713 714        // First set to specified device (entrypoint may not be set)715        if (device != entrypoint_device)716        {717            error = CubDebug(cudaGetDevice(&entrypoint_device));718            if (cudaSuccess != error)719            {720                return error;721            }722 723            error = CubDebug(cudaSetDevice(device));724            if (cudaSuccess != error)725            {726                return error;727            }728        }729 730        if (recached)731        {732            // Insert the ready event in the associated stream (must have current device set properly)733            error = CubDebug(cudaEventRecord(search_key.ready_event, search_key.associated_stream));734            if (cudaSuccess != error)735            {736                return error;737            }738        }739 740        if (!recached)741        {742            // Free the allocation from the runtime and cleanup the event.743            error = CubDebug(cudaFree(d_ptr));744            if (cudaSuccess != error)745            {746                return error;747            }748 749            error = CubDebug(cudaEventDestroy(search_key.ready_event));750            if (cudaSuccess != error)751            {752                return error;753            }754 755            if (debug) _CubLog("\tDevice %d freed %lld bytes from associated stream %lld.\n\t\t  %lld available blocks cached (%lld bytes), %lld live blocks (%lld bytes) outstanding.\n",756                device, (long long) search_key.bytes, (long long) search_key.associated_stream, (long long) cached_blocks.size(), (long long) cached_bytes[device].free, (long long) live_blocks.size(), (long long) cached_bytes[device].live);757        }758 759        // Reset device760        if ((entrypoint_device != INVALID_DEVICE_ORDINAL) && (entrypoint_device != device))761        {762            error = CubDebug(cudaSetDevice(entrypoint_device));763            if (cudaSuccess != error)764            {765                return error;766            }767        }768 769        return error;770    }771 772    /**773     * @brief Frees a live allocation of device memory on the current device, returning it to the774     *        allocator.775     *776     * Once freed, the allocation becomes available immediately for reuse within the @p777     * active_stream with which it was associated with during allocation, and it becomes available778     * for reuse within other streams when all prior work submitted to @p active_stream has779     * completed.780     */781    cudaError_t DeviceFree(782        void*           d_ptr)783    {784        return DeviceFree(INVALID_DEVICE_ORDINAL, d_ptr);785    }786 787 788    /**789     * @brief Frees all cached device allocations on all devices790     */791    cudaError_t FreeAllCached()792    {793        cudaError_t error         = cudaSuccess;794        int entrypoint_device     = INVALID_DEVICE_ORDINAL;795        int current_device        = INVALID_DEVICE_ORDINAL;796 797        mutex.lock();798 799        while (!cached_blocks.empty())800        {801            // Get first block802            CachedBlocks::iterator begin = cached_blocks.begin();803 804            // Get entry-point device ordinal if necessary805            if (entrypoint_device == INVALID_DEVICE_ORDINAL)806            {807                error = CubDebug(cudaGetDevice(&entrypoint_device));808                if (cudaSuccess != error)809                {810                    break;811                }812            }813 814            // Set current device ordinal if necessary815            if (begin->device != current_device)816            {817                error = CubDebug(cudaSetDevice(begin->device));818                if (cudaSuccess != error)819                {820                    break;821                }822                current_device = begin->device;823            }824 825            // Free device memory826            error = CubDebug(cudaFree(begin->d_ptr));827            if (cudaSuccess != error)828            {829                break;830            }831 832            error = CubDebug(cudaEventDestroy(begin->ready_event));833            if (cudaSuccess != error)834            {835                break;836            }837 838            // Reduce balance and erase entry839            const size_t block_bytes = begin->bytes;840            cached_bytes[current_device].free -= block_bytes;841            cached_blocks.erase(begin);842 843            if (debug) _CubLog("\tDevice %d freed %lld bytes.\n\t\t  %lld available blocks cached (%lld bytes), %lld live blocks (%lld bytes) outstanding.\n",844                current_device, (long long) block_bytes, (long long) cached_blocks.size(), (long long) cached_bytes[current_device].free, (long long) live_blocks.size(), (long long) cached_bytes[current_device].live);845 846        }847 848        mutex.unlock();849 850        // Attempt to revert back to entry-point device if necessary851        if (entrypoint_device != INVALID_DEVICE_ORDINAL)852        {853            error = CubDebug(cudaSetDevice(entrypoint_device));854            if (cudaSuccess != error)855            {856                return error;857            }858        }859 860        return error;861    }862 863 864    /**865     * @brief Destructor866     */867    virtual ~CachingDeviceAllocator()868    {869        if (!skip_cleanup)870            FreeAllCached();871    }872 873};874 875 876 877 878/** @} */       // end group UtilMgmt879 880CUB_NAMESPACE_END881 
codekingpro/portable-devtools · Team Ai