temp_storage.hpp

// Copyright (c) 2017-2023 Advanced Micro Devices, Inc. All rights reserved.
//
// Permission is hereby granted, free of charge, to any person obtaining a copy
// of this software and associated documentation files (the "Software"), to deal
// in the Software without restriction, including without limitation the rights
// to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
// copies of the Software, and to permit persons to whom the Software is
// furnished to do so, subject to the following conditions:
//
// The above copyright notice and this permission notice shall be included in
// all copies or substantial portions of the Software.
//
// THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
// IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
// FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT.  IN NO EVENT SHALL THE
// AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
// LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
// OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN
// THE SOFTWARE.

#ifndef ROCPRIM_DETAIL_TEMP_STORAGE_HPP_
#define ROCPRIM_DETAIL_TEMP_STORAGE_HPP_

#include <cstddef>

#include "../config.hpp"
#include "../types.hpp"
#include "various.hpp"

BEGIN_ROCPRIM_NAMESPACE
namespace detail
{

namespace temp_storage
{

constexpr size_t default_alignment = 256;

constexpr size_t minimum_allocation_size = 4;

struct layout
{
    size_t size;

    size_t alignment = default_alignment;
};

template<typename T>
struct simple_partition
{
    T** dest;
    layout storage_layout;

    layout get_layout()
    {
        return this->storage_layout;
    }

    void set_storage(void* const storage)
    {
        *this->dest = this->storage_layout.size == 0 ? nullptr : static_cast<T*>(storage);
    }
};

template<typename T>
simple_partition<T> make_partition(T** dest, layout storage_layout)
{
    return simple_partition<T>{dest, storage_layout};
}

template<typename T>
simple_partition<T> make_partition(T** dest, size_t size, size_t alignment = default_alignment)
{
    return make_partition(dest, {size, alignment});
}

template<typename T>
simple_partition<T> ptr_aligned_array(T** dest, size_t elements)
{
    return make_partition(dest, elements * sizeof(T), alignof(T));
}

template<typename... Ts>
struct linear_partition
{
    ::rocprim::tuple<Ts...> sub_partitions;

    linear_partition(Ts... sub_partitions) : sub_partitions{sub_partitions...} {}

    layout get_layout()
    {
        size_t required_alignment = 1;
        size_t required_size      = 0;

        for_each_in_tuple(this->sub_partitions,
                          [&](auto& sub_partition)
                          {
                              const auto sub_layout = sub_partition.get_layout();

                              required_alignment
                                  = std::max(required_alignment, sub_layout.alignment);

                              if(sub_layout.size > 0)
                                  required_size = align_size(required_size, sub_layout.alignment)
                                                  + sub_layout.size;
                          });

        return {required_size, required_alignment};
    }

    void set_storage(void* const storage)
    {
        size_t offset = 0;
        for_each_in_tuple(this->sub_partitions,
                          [&](auto& sub_partition)
                          {
                              const auto sub_layout = sub_partition.get_layout();

                              if(sub_layout.size > 0)
                                  offset = align_size(offset, sub_layout.alignment);

                              sub_partition.set_storage(
                                  static_cast<void*>(static_cast<char*>(storage) + offset));
                              offset += sub_layout.size;
                          });
    }
};

template<typename... Ts>
linear_partition<Ts...> make_linear_partition(Ts... ts)
{
    return linear_partition<Ts...>(ts...);
}

template<typename... Ts>
struct union_partition
{
    ::rocprim::tuple<Ts...> sub_partitions;

    union_partition(Ts... sub_partitions) : sub_partitions{sub_partitions...} {}

    layout get_layout()
    {
        size_t required_alignment = 1;
        size_t required_size      = 0;

        for_each_in_tuple(this->sub_partitions,
                          [&](auto& sub_partition)
                          {
                              const auto sub_layout = sub_partition.get_layout();

                              required_alignment
                                  = std::max(required_alignment, sub_layout.alignment);
                              required_size = std::max(required_size, sub_layout.size);
                          });

        return {required_size, required_alignment};
    }

    void set_storage(void* const storage)
    {
        for_each_in_tuple(this->sub_partitions,
                          [&](auto& sub_partition) { sub_partition.set_storage(storage); });
    }
};

template<typename... Ts>
union_partition<Ts...> make_union_partition(Ts... ts)
{
    return union_partition<Ts...>(ts...);
}

template<typename TempStoragePartition>
hipError_t
    partition(void* const temporary_storage, size_t& storage_size, TempStoragePartition partition)
{
    const auto layout = partition.get_layout();
    // Make sure the user wont try to allocate 0 bytes of memory.
    const size_t required_size = std::max(layout.size, minimum_allocation_size);

    if(temporary_storage == nullptr)
    {
        storage_size = required_size;
        return hipSuccess;
    }
    else if(storage_size < required_size)
    {
        return hipErrorInvalidValue;
    }

    partition.set_storage(temporary_storage);

    return hipSuccess;
}
} // namespace temp_storage

} // namespace detail
END_ROCPRIM_NAMESPACE

#endif // ROCPRIM_DETAIL_TEMP_STORAGE_HPP_