UT_ConcurrentClassifier.h

Go to the documentation of this file.

/*
 * PROPRIETARY INFORMATION.  This software is proprietary to
 * Side Effects Software Inc., and is not to be reproduced,
 * transmitted, or disclosed in any way without written permission.
 *
 * NAME:        UT_ConcurrentClassifier.h (UT Library, C++)
 *
 * Lock-free parallel disjoint set data structure (aka UNION-FIND)
 * with path compression and union by rank
 *
 * Supports concurrent find(), same() and unite() calls as described
 * in the paper (renamed here to classRoot(), areInSameClass(), and
 * unionClasses(), respectively.
 *
 * "Wait-free Parallel Algorithms for the Union-Find Problem"
 * by Richard J. Anderson and Heather Woll
 *
 * In addition, this class supports optimistic locking (try_lock/unlock)
 * of disjoint sets and a combined unite+unlock operation.
 *
 * Based on the freely available implementation by Wenzel Jakob on GitHub.
 */


#ifndef __UT_ConcurrentClassifier__
#define __UT_ConcurrentClassifier__

#include <SYS/SYS_Types.h>
#include <SYS/SYS_TypeDecorate.h>
#include "UT_Array.h"

#if defined(ARM64)
#include <sse2neon.h>
#else
#include <immintrin.h>
#endif

SYS_DECLARE_LEGACY_TR(std::atomic<uint64>)

class UT_ConcurrentClassifier 
{
public:
                         UT_ConcurrentClassifier(int32 size = 0)
                            { if (size > 0) makeClass(size); }
    
    uint32               makeClass(uint32 num_classes = 1)
    {
        uint32 s = myData.size();
        myData.setSizeNoInit(s + num_classes);
        for (int i = 0, ie = num_classes; i < ie; ++i)
            myData(s + i) = (uint32) i;
        return s;
    }

    uint32               classRoot(uint32 id) const
    {
        while (id != parent(id)) {
            uint64 value = myData(id);
            uint32 new_parent = parent((uint32) value);
            uint64 new_value =
                (value & 0xFFFFFFFF00000000ULL) | new_parent;

            /* Try to update parent (may fail, that's ok) */
            if (value != new_value)
                myData(id).compare_exchange_weak(value, new_value);
            id = new_parent;
        }
        return id;
    }

    bool                 areInSameClass(uint32 id1, uint32 id2) const
    {
        for ( ; ; ) 
        {
            id1 = classRoot(id1);
            id2 = classRoot(id2);
            if (id1 == id2)
                return true;
            if (parent(id1) == id1)
                return false;
        }
    }

    uint32               unionClasses(uint32 id1, uint32 id2)
    {
        for ( ; ; ) 
        {
            id1 = classRoot(id1);
            id2 = classRoot(id2);

            if (id1 == id2)
                return id1;

            uint32 r1 = rank(id1), r2 = rank(id2);

            if (r1 > r2 || (r1 == r2 && id1 < id2)) 
            {
                UTswap(r1, r2);
                UTswap(id1, id2);
            }

            uint64 old_entry = ((uint64_t) r1 << 32) | id1;
            uint64 new_entry = ((uint64_t) r1 << 32) | id2;

            if (!myData(id1).compare_exchange_strong(old_entry, new_entry))
                continue;

            if (r1 == r2) {
                old_entry = ((uint64_t) r2 << 32) | id2;
                new_entry = ((uint64_t) (r2+1) << 32) | id2;

                // Try to update the rank (may fail, that's ok)
                myData(id2).compare_exchange_weak(old_entry, new_entry);
            }

            break;
        }
        return id2;
    }

    // Try to lock a disjoint union identified by one of its elements (this can
    // occasionally fail when there are concurrent operations). The parameter
    // 'id' will be updated to store the current representative ID of the union.
    bool                 try_lock(uint32 &id) 
    {
        const uint64 lock_flag = 1ULL << 63;
        id = classRoot(id);
        uint64 value = myData(id);
        if ((value & lock_flag) || (uint32) value != id)
            return false;

        // On IA32/x64, a PAUSE instruction is recommended for CAS busy loops.
#if 0
        #if defined(__i386__) || defined(__amd64__)
            __asm__ __volatile__ ("pause\n");
        #endif
#endif
        _mm_pause();
        return myData(id).compare_exchange_strong(value, value | lock_flag);
    }

    void                 unlock(uint32 id) 
    {
        const uint64 lock_flag = 1ULL << 63;
        myData(id) &= ~lock_flag;
    }

    
    // Return the representative index of the set that results from merging
    // locked disjoint sets 'id1' and 'id2'.
    uint32               unite_index_locked(uint32 id1, uint32 id2) const 
    {
        uint32 r1 = rank(id1), r2 = rank(id2);
        return (r1 > r2 || (r1 == r2 && id1 < id2)) ? id1 : id2;
    }

    
    // Atomically unite two locked disjoint sets and unlock them. Assumes
    // that here are no other concurrent unite() involving the same sets.
    uint32               unite_unlock(uint32 id1, uint32 id2) 
    {
        uint32 r1 = rank(id1), r2 = rank(id2);

        if (r1 > r2 || (r1 == r2 && id1 < id2)) 
        {
            UTswap(r1, r2);
            UTswap(id1, id2);
        }

        myData(id1) = ((uint64) r1 << 32) | id2;
        myData(id2) = ((uint64) (r2 + ((r1 == r2) ? 1 : 0)) << 32) | id2;

        return id2;
    }
    
    uint32               size() const { return (uint32) myData.size(); }
    uint32               entries() const  { return size(); }
    
    uint32               numClasses()
    {
        return tbb::parallel_reduce(tbb::blocked_range<uint32>(0, size()),
                    uint32(0),
                    [&](tbb::blocked_range<uint32> r, uint32 total)
                    {
                        for (uint32 i = r.begin(); i < r.end();
                             ++i)
                            total += (i == classRoot(i));
    
                        return total;
                    }, std::plus<uint32>());
    }

private:
    SYS_FORCE_INLINE
    uint32               rank(uint32 id) const 
                            { return ((uint32) (myData[id] >> 32)) & 0x7FFFFFFFu; }

    SYS_FORCE_INLINE
    uint32               parent(uint32_t id) const 
                            { return (uint32_t) myData(id); }
    
    using DataEntry = std::atomic<uint64>;
    using DataArray = UT_Array<DataEntry>;
    mutable DataArray    myData;
};

#endif