map.hpp Source File

map.hpp Source File#

Composable Kernel: map.hpp Source File
// SPDX-License-Identifier: MIT
// Copyright (c) 2018-2024, Advanced Micro Devices, Inc. All rights reserved.
 
#pragma once
 
#include "ck_tile/core/config.hpp"
#include "ck_tile/core/container/array.hpp"
#include "ck_tile/core/container/sequence.hpp"
#include "ck_tile/core/container/tuple.hpp"
 
namespace ck_tile {
 
// naive map
template <typename key, typename data, index_t max_size = 128>
struct map
{
    using pair_type = tuple<key, data>;
    using impl_type = array<pair_type, max_size>;
 
    impl_type impl_;
    index_t size_;
 
    struct iterator
    {
        impl_type& impl_;
        index_t pos_;
 
        CK_TILE_HOST_DEVICE constexpr iterator(impl_type& impl, index_t pos)
            : impl_{impl}, pos_{pos}
        {
        }
 
        CK_TILE_HOST_DEVICE constexpr iterator& operator++()
        {
            pos_++;
            return *this;
        }
 
        CK_TILE_HOST_DEVICE constexpr bool operator!=(const iterator& other) const
        {
            return other.pos_ != pos_;
        }
 
        CK_TILE_HOST_DEVICE constexpr pair_type& operator*() { return impl_.at(pos_); }
    };
 
    struct const_iterator
    {
        const impl_type& impl_;
        index_t pos_;
 
        CK_TILE_HOST_DEVICE constexpr const_iterator(const impl_type& impl, index_t pos)
            : impl_{impl}, pos_{pos}
        {
        }
 
        CK_TILE_HOST_DEVICE constexpr const_iterator& operator++()
        {
            pos_++;
 
            return *this;
        }
 
        CK_TILE_HOST_DEVICE constexpr bool operator!=(const const_iterator& other) const
        {
            return other.pos_ != pos_;
        }
 
        CK_TILE_HOST_DEVICE constexpr const pair_type& operator*() const { return impl_.at(pos_); }
    };
 
    CK_TILE_HOST_DEVICE constexpr map() : impl_{}, size_{0} {}
 
    CK_TILE_HOST_DEVICE constexpr index_t size() const { return size_; }
 
    CK_TILE_HOST_DEVICE void clear() { size_ = 0; }
 
    CK_TILE_HOST_DEVICE constexpr index_t find_position(const key& k) const
    {
        for(index_t i = 0; i < size(); i++)
        {
            if(impl_[i].template at<0>() == k)
            {
                return i;
            }
        }
 
        return size_;
    }
 
    CK_TILE_HOST_DEVICE constexpr const_iterator find(const key& k) const
    {
        return const_iterator{impl_, find_position(k)};
    }
 
    CK_TILE_HOST_DEVICE constexpr iterator find(const key& k)
    {
        return iterator{impl_, find_position(k)};
    }
 
    CK_TILE_HOST_DEVICE constexpr const data& operator[](const key& k) const
    {
        const auto it = find(k);
 
        // FIXME
        // assert(it.pos_ < size());
 
        return impl_[it.pos_].template at<1>();
    }
 
    CK_TILE_HOST_DEVICE constexpr data& operator()(const key& k)
    {
        auto it = find(k);
 
        // if entry not found
        if(it.pos_ == size())
        {
            impl_(it.pos_).template at<0>() = k;
            size_++;
        }
 
        // FIXME
        // assert(size_ <= max_size);
 
        return impl_(it.pos_).template at<1>();
    }
 
    // WARNING: needed by compiler for C++ range-based for loop only, don't use this function!
    CK_TILE_HOST_DEVICE constexpr const_iterator begin() const { return const_iterator{impl_, 0}; }
 
    // WARNING: needed by compiler for C++ range-based for loop only, don't use this function!
    CK_TILE_HOST_DEVICE constexpr const_iterator end() const
    {
        return const_iterator{impl_, size_};
    }
 
    // WARNING: needed by compiler for C++ range-based for loop only, don't use this function!
    CK_TILE_HOST_DEVICE constexpr iterator begin() { return iterator{impl_, 0}; }
 
    // WARNING: needed by compiler for C++ range-based for loop only, don't use this function!
    CK_TILE_HOST_DEVICE constexpr iterator end() { return iterator{impl_, size_}; }
};
 
template <typename key, typename data, index_t max_size>
CK_TILE_HOST_DEVICE static void print(const map<key, data, max_size>& m)
{
    printf("map{size_: %d, impl_: [", m.size_);
    for(const auto& [k, d] : m)
    {
        printf("{key: ");
        print(k);
        printf(", data: ");
        print(d);
        printf("}, ");
    }
    printf("]}");
}
 
} // namespace ck_tile