/home/docs/checkouts/readthedocs.org/user_builds/advanced-micro-devices-composable-kernel/checkouts/develop/include/ck_tile/core/container/map.hpp Source File

/home/docs/checkouts/readthedocs.org/user_builds/advanced-micro-devices-composable-kernel/checkouts/develop/include/ck_tile/core/container/map.hpp Source File#

Composable Kernel: /home/docs/checkouts/readthedocs.org/user_builds/advanced-micro-devices-composable-kernel/checkouts/develop/include/ck_tile/core/container/map.hpp Source File

Go to the documentation of this 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