rocm-systems/projects/rocprofiler-sdk/source/lib/rocprofiler-sdk/buffer.hpp

// MIT License
//
// Copyright (c) 2023-2025 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.

#pragma once

#include <rocprofiler-sdk/buffer.h>
#include <rocprofiler-sdk/fwd.h>

#include "lib/common/container/record_header_buffer.hpp"
#include "lib/common/container/stable_vector.hpp"
#include "lib/common/demangle.hpp"

#include <fmt/format.h>

#include <sys/types.h>
#include <array>
#include <atomic>
#include <cstdint>
#include <optional>

namespace rocprofiler
{
namespace buffer
{
struct instance
{
    using buffer_t                       = common::container::record_header_buffer;
    static constexpr auto size           = 2;  // double buffering
    static constexpr auto sync_wait_usec = std::chrono::microseconds{10};

    mutable std::array<buffer_t, size>         buffers       = {};
    mutable std::array<std::atomic_flag, size> syncer        = {false, false};  // r/w lock
    mutable std::atomic<uint32_t>              buffer_idx    = {};              // array index
    mutable std::atomic<uint64_t>              drop_count    = {};
    uint64_t                                   watermark     = 0;
    uint64_t                                   context_id    = 0;  // rocprofiler_context_id_t value
    uint64_t                                   buffer_id     = 0;  // rocprofiler_buffer_id_t value
    uint64_t                                   task_group_id = 0;  // thread-pool assignment
    rocprofiler_buffer_tracing_cb_t            callback      = nullptr;
    void*                                      callback_data = nullptr;
    rocprofiler_buffer_policy_t                policy        = ROCPROFILER_BUFFER_POLICY_NONE;

    template <typename Tp>
    bool emplace(uint32_t, uint32_t, Tp&);

    buffer_t& get_internal_buffer();
    buffer_t& get_internal_buffer(size_t);
};

using unique_buffer_vec_t = common::container::stable_vector<std::unique_ptr<instance>, 4>;

bool
is_valid_buffer_id(rocprofiler_buffer_id_t id);

std::optional<rocprofiler_buffer_id_t>
allocate_buffer();

unique_buffer_vec_t*
get_buffers();

instance*
get_buffer(rocprofiler_buffer_id_t buffer_id);

instance*
get_buffer(uint64_t buffer_idx);

rocprofiler_status_t
flush(rocprofiler_buffer_id_t buffer_id, bool wait);

rocprofiler_status_t
flush(uint64_t buffer_idx, bool wait);
}  // namespace buffer
}  // namespace rocprofiler

inline rocprofiler::buffer::instance::buffer_t&
rocprofiler::buffer::instance::get_internal_buffer()
{
    auto idx = buffer_idx.load() % buffers.size();
    return buffers.at(idx);
}

inline rocprofiler::buffer::instance::buffer_t&
rocprofiler::buffer::instance::get_internal_buffer(size_t idx)
{
    return buffers.at(idx % buffers.size());
}

inline rocprofiler::buffer::instance*
rocprofiler::buffer::get_buffer(uint64_t buffer_idx)
{
    return get_buffer(rocprofiler_buffer_id_t{buffer_idx});
}

inline rocprofiler_status_t
rocprofiler::buffer::flush(uint64_t buffer_idx, bool wait)
{
    return flush(rocprofiler_buffer_id_t{buffer_idx}, wait);
}

template <typename Tp>
inline bool
rocprofiler::buffer::instance::emplace(uint32_t category, uint32_t kind, Tp& value)
{
    struct local_sync
    {
        local_sync(uint64_t buffer_id, uint64_t idx, std::atomic_flag& flag)
        : m_flag{flag}
        {
            if(m_flag.test_and_set())
            {
                ROCP_INFO << fmt::format(
                    "waiting for buffer flush to complete [id={}, index={}]...", buffer_id, idx);
                while(m_flag.test_and_set())
                {
                    ROCP_TRACE << fmt::format(
                        "waiting for buffer flush to complete [id={}, index={}]...",
                        buffer_id,
                        idx);
                    std::this_thread::yield();
                    std::this_thread::sleep_for(sync_wait_usec);
                }
            }
        }

        ~local_sync() { m_flag.clear(); }

        std::atomic_flag& m_flag;
    };

    // get the index of the current buffer
    auto get_idx = [this]() { return buffer_idx.load(std::memory_order_acquire) % buffers.size(); };
    auto idx     = get_idx();
    auto success = false;
    {
        auto _syncer = local_sync{buffer_id, idx, syncer.at(idx)};  // ensure buffer not flushing
        success      = buffers.at(idx).emplace(category, kind, value);
    }

    if(!success)
    {
        if(buffers.at(idx).capacity() < sizeof(value))
        {
            ROCP_CI_LOG(ERROR) << fmt::format(
                "buffer {} too small (size={}) to hold an object of type {} with size {}",
                buffer_id,
                buffers.at(idx).capacity(),
                common::cxx_demangle(typeid(value).name()),
                sizeof(value));
            return false;
        }

        if(policy == ROCPROFILER_BUFFER_POLICY_LOSSLESS)
        {
            // blocks until buffer is flushed
            do
            {
                buffer::flush(buffer_id, true);
                idx          = get_idx();
                auto _syncer = local_sync{buffer_id, idx, syncer.at(idx)};  // wait for flush
                success      = buffers.at(idx).emplace(category, kind, value);
            } while(!success);
        }
        else
        {
            ++drop_count;
        }
    }

    if(buffers.at(idx).count() >= watermark)
    {
        // flush without syncing
        buffer::flush(buffer_id, false);
    }

    return success;
}