/*
   Copyright (c) 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.
*/

#include <gtest/gtest.h>
#include <cstddef>

#include "lib/rocprofiler/pc_sampling/parser/pc_record_interface.hpp"
#include "lib/rocprofiler/pc_sampling/parser/tests/mocks.hpp"

#define GFXIP_MAJOR 9

#define TYPECHECK(x)                                                                               \
    snapshots.push_back(pcsample_snapshot_v1_t{.dual_issue_valu   = 0,                             \
                                               .inst_type         = ::PCSAMPLE::x,                 \
                                               .reason_not_issued = 0,                             \
                                               .arb_state_issue   = 0,                             \
                                               .arb_state_stall   = 0});
#define UNROLL_TYPECHECK()                                                                         \
    TYPECHECK(TYPE_VALU);                                                                          \
    TYPECHECK(TYPE_MATRIX);                                                                        \
    TYPECHECK(TYPE_SCALAR);                                                                        \
    TYPECHECK(TYPE_TEX);                                                                           \
    TYPECHECK(TYPE_LDS);                                                                           \
    TYPECHECK(TYPE_FLAT);                                                                          \
    TYPECHECK(TYPE_EXP);                                                                           \
    TYPECHECK(TYPE_MESSAGE);                                                                       \
    TYPECHECK(TYPE_BARRIER);                                                                       \
    TYPECHECK(TYPE_BRANCH_NOT_TAKEN);                                                              \
    TYPECHECK(TYPE_BRANCH_TAKEN);                                                                  \
    TYPECHECK(TYPE_JUMP);                                                                          \
    TYPECHECK(TYPE_OTHER);                                                                         \
    TYPECHECK(TYPE_NO_INST);

#define REASONCHECK(x)                                                                             \
    snapshots.push_back(pcsample_snapshot_v1_t{.dual_issue_valu   = 0,                             \
                                               .inst_type         = 0,                             \
                                               .reason_not_issued = ::PCSAMPLE::x,                 \
                                               .arb_state_issue   = 0,                             \
                                               .arb_state_stall   = 0});
#define UNROLL_REASONCHECK(x)                                                                      \
    REASONCHECK(REASON_NOT_AVAILABLE);                                                             \
    REASONCHECK(REASON_ALU);                                                                       \
    REASONCHECK(REASON_WAITCNT);                                                                   \
    REASONCHECK(REASON_INTERNAL);                                                                  \
    REASONCHECK(REASON_BARRIER);                                                                   \
    REASONCHECK(REASON_ARBITER);                                                                   \
    REASONCHECK(REASON_EX_STALL);                                                                  \
    REASONCHECK(REASON_OTHER_WAIT);

#define ARBCHECK1(x, y)                                                                            \
    snapshots.push_back(pcsample_snapshot_v1_t{.dual_issue_valu   = 0,                             \
                                               .inst_type         = 0,                             \
                                               .reason_not_issued = 0,                             \
                                               .arb_state_issue   = 1 << ::PCSAMPLE::x,            \
                                               .arb_state_stall   = 1 << ::PCSAMPLE::y});
#define ARBCHECK2(x)                                                                               \
    ARBCHECK1(x, ISSUE_VALU);                                                                      \
    ARBCHECK1(x, ISSUE_MATRIX);                                                                    \
    ARBCHECK1(x, ISSUE_SCALAR);                                                                    \
    ARBCHECK1(x, ISSUE_VMEM_TEX);                                                                  \
    ARBCHECK1(x, ISSUE_LDS);                                                                       \
    ARBCHECK1(x, ISSUE_FLAT);                                                                      \
    ARBCHECK1(x, ISSUE_EXP);                                                                       \
    ARBCHECK1(x, ISSUE_MISC);

#define UNROLL_ARBCHECK()                                                                          \
    ARBCHECK2(ISSUE_VALU);                                                                         \
    ARBCHECK2(ISSUE_MATRIX);                                                                       \
    ARBCHECK2(ISSUE_SCALAR);                                                                       \
    ARBCHECK2(ISSUE_VMEM_TEX);                                                                     \
    ARBCHECK2(ISSUE_LDS);                                                                          \
    ARBCHECK2(ISSUE_FLAT);                                                                         \
    ARBCHECK2(ISSUE_EXP);                                                                          \
    ARBCHECK2(ISSUE_MISC);

std::mt19937 rdgen(1);

TEST(pcs_parser_context, init) { PCSamplingParserContext context; }

/**
 * Sample user memory allocation callback.
 * It expects userdata to be cast-able to a pointer to
 * std::vector<std::pair<pcsample_v1_t*, uint64_t>>
 */
static uint64_t
alloc_callback(pcsample_v1_t** buffer, uint64_t size, void* userdata)
{
    *buffer      = new pcsample_v1_t[size];
    auto& vector = *reinterpret_cast<std::vector<std::pair<pcsample_v1_t*, uint64_t>>*>(userdata);
    vector.push_back({*buffer, size});
    return size;
}

/**
 * Uses the MockWave dispatch's unique_id store in the pc field to verify
 * the reconstructed correlation_id.
 */
static bool
check_samples(pcsample_v1_t* samples, uint64_t size)
{
    for(size_t i = 0; i < size; i++)
        if(samples[i].correlation_id != samples[i].pc) return false;
    return true;
}

/**
 * Simplest mock classes use, generates a single queue+dispatch with 2 PC samples.
 */
TEST(pcs_parser_correlation_id, hello_world)
{
    std::shared_ptr<MockRuntimeBuffer> buffer   = std::make_shared<MockRuntimeBuffer>();
    std::shared_ptr<MockQueue>         queue    = std::make_shared<MockQueue>(16, buffer);
    std::shared_ptr<MockDispatch>      dispatch = std::make_shared<MockDispatch>(queue);

    buffer->genUpcomingSamples(2);
    MockWave(dispatch).genPCSample();
    MockWave(dispatch).genPCSample();

    std::vector<std::pair<pcsample_v1_t*, uint64_t>> all_allocations;

    CHECK_PARSER(parse_buffer((generic_sample_t*) buffer->packets.data(),
                              buffer->packets.size(),
                              GFXIP_MAJOR,
                              alloc_callback,
                              (void*) &all_allocations));

    assert(all_allocations.size() == 1 && "HelloWorld: Incorrect number of callbacks");
    for(auto& sample : all_allocations)
    {
        assert(sample.second == 2 && "HelloWorld: Incorrect number of samples");
        assert(check_samples(sample.first, sample.second) &&
               "HelloWorld: parsed ID does not match correct ID");
        delete[] sample.first;
    }
}

/**
 * A little more complicated.
 * Generates a few dispatches for 2 different queues and samples in forward and reverse order.
 * Checks if the reconstructed correlation_id is correct.
 */
TEST(pcs_parser_correlation_id, reverse_wave_order)
{
    std::shared_ptr<MockRuntimeBuffer> buffer = std::make_shared<MockRuntimeBuffer>();
    std::shared_ptr<MockQueue>         queue1 = std::make_shared<MockQueue>(16, buffer);
    std::shared_ptr<MockQueue>         queue2 = std::make_shared<MockQueue>(16, buffer);

    std::vector<std::shared_ptr<MockDispatch>> dispatches;
    dispatches.push_back(std::make_shared<MockDispatch>(queue1));
    dispatches.push_back(std::make_shared<MockDispatch>(queue1));
    dispatches.push_back(std::make_shared<MockDispatch>(queue2));
    dispatches.push_back(std::make_shared<MockDispatch>(queue2));
    dispatches.push_back(std::make_shared<MockDispatch>(queue1));

    buffer->genUpcomingSamples(dispatches.size());
    for(auto it = dispatches.rbegin(); it != dispatches.rend(); it++)
        MockWave(*it).genPCSample();
    buffer->genUpcomingSamples(dispatches.size());
    for(auto it = dispatches.begin(); it != dispatches.end(); it++)
        MockWave(*it).genPCSample();

    std::vector<std::pair<pcsample_v1_t*, uint64_t>> all_allocations;

    CHECK_PARSER(parse_buffer((generic_sample_t*) buffer->packets.data(),
                              buffer->packets.size(),
                              GFXIP_MAJOR,
                              alloc_callback,
                              (void*) &all_allocations));

    assert(all_allocations.size() == 2 && "ReverseWaveOrder test: Incorrect number of callbacks");
    for(auto& sample : all_allocations)
    {
        assert(sample.second == dispatches.size() &&
               "ReverseWaveOrder: Incorrect number of samples");
        assert(check_samples(sample.first, sample.second) &&
               "ReverseWaveOrder: parsed ID does not match correct ID");
        delete[] sample.first;
    }
}

/**
 * Creates a small queue and causes the dispatch_ids to wrap around a few times, and generates
 * a single sample per dispatch. Checks the parser is properly handling the wrapping of queues.
 */
TEST(pcs_parser_correlation_id, dispatch_wrapping)
{
    const int                          num_samples = 32;
    std::shared_ptr<MockRuntimeBuffer> buffer      = std::make_shared<MockRuntimeBuffer>();
    std::shared_ptr<MockQueue>         queue       = std::make_shared<MockQueue>(5, buffer);

    for(int i = 0; i < num_samples; i++)
    {
        auto dispatch = std::make_shared<MockDispatch>(queue);
        buffer->genUpcomingSamples(1);
        MockWave(dispatch).genPCSample();
    }

    std::vector<std::pair<pcsample_v1_t*, uint64_t>> all_allocations;

    CHECK_PARSER(parse_buffer((generic_sample_t*) buffer->packets.data(),
                              buffer->packets.size(),
                              GFXIP_MAJOR,
                              alloc_callback,
                              (void*) &all_allocations));

    assert(all_allocations.size() == num_samples &&
           "RandomSamples test: Incorrect number of callbacks");
    for(auto& sample : all_allocations)
    {
        assert(sample.second == 1 && "RandomSamples: Incorrect number of samples");
        assert(check_samples(sample.first, sample.second) &&
               "RandomSamples: parsed ID does not match correct ID");
        delete[] sample.first;
    }
}

/**
 * Creates a few queues with a few dispatchs per queue.
 * Adds random samples per dispatch, and checks the result.
 */
TEST(pcs_parser_correlation_id, random_samples)
{
    const int                          num_samples = 1024;
    std::shared_ptr<MockRuntimeBuffer> buffer      = std::make_shared<MockRuntimeBuffer>();
    std::shared_ptr<MockQueue>         queue1      = std::make_shared<MockQueue>(16, buffer);
    std::shared_ptr<MockQueue>         queue2      = std::make_shared<MockQueue>(16, buffer);
    std::shared_ptr<MockQueue>         queue3      = std::make_shared<MockQueue>(16, buffer);
    std::shared_ptr<MockQueue>         queue4      = std::make_shared<MockQueue>(16, buffer);

    std::vector<std::shared_ptr<MockDispatch>> dispatches;
    dispatches.push_back(std::make_shared<MockDispatch>(queue1));
    dispatches.push_back(std::make_shared<MockDispatch>(queue1));
    dispatches.push_back(std::make_shared<MockDispatch>(queue2));
    dispatches.push_back(std::make_shared<MockDispatch>(queue3));
    dispatches.push_back(std::make_shared<MockDispatch>(queue1));
    dispatches.push_back(std::make_shared<MockDispatch>(queue3));
    dispatches.push_back(std::make_shared<MockDispatch>(queue3));
    dispatches.push_back(std::make_shared<MockDispatch>(queue2));
    dispatches.push_back(std::make_shared<MockDispatch>(queue1));

    buffer->genUpcomingSamples(num_samples);
    for(int i = 0; i < num_samples; i++)
        MockWave(dispatches[rdgen() % dispatches.size()]).genPCSample();

    std::vector<std::pair<pcsample_v1_t*, uint64_t>> all_allocations;

    CHECK_PARSER(parse_buffer((generic_sample_t*) buffer->packets.data(),
                              buffer->packets.size(),
                              GFXIP_MAJOR,
                              alloc_callback,
                              (void*) &all_allocations));

    assert(all_allocations.size() == 1 && "RandomSamples test: Incorrect number of callbacks");
    for(auto& sample : all_allocations)
    {
        assert(sample.second == num_samples && "RandomSamples: Incorrect number of samples");
        assert(check_samples(sample.first, sample.second) &&
               "RandomSamples: parsed ID does not match correct ID");
        delete[] sample.first;
    }
}

/**
 * Hammers the parser by creating and destrying queues at random, adding dispatches at random
 * and generating PC samples at random. By default we use all 4 unique doorbells,
 * queue size is 16 and we generate 10k samples dispatch.
 */
TEST(pcs_parser_correlation_id, queue_hammer)
{
    constexpr int NUM_ACTIONS = 10000;
    constexpr int QSIZE       = 16;
    constexpr int NUM_QUEUES  = MockDoorBell::num_unique_bells;
    constexpr int ACTION_MAX  = QSIZE * NUM_QUEUES / 2;

    std::shared_ptr<MockRuntimeBuffer> buffer = std::make_shared<MockRuntimeBuffer>();

    std::array<std::shared_ptr<MockQueue>, NUM_QUEUES>                 queues;
    std::array<std::vector<std::shared_ptr<MockDispatch>>, NUM_QUEUES> active_dispatches;

    int    num_reset_queues         = 0;
    int    num_samples_generated    = 0;
    int    num_dispatches_generated = 0;
    double avg_q_occupancy          = 0;
    size_t max_q_occupancy          = 0;

    for(int i = 0; i < NUM_QUEUES; i++)
        queues[i] = std::make_shared<MockQueue>(QSIZE, buffer);
    for(int i = 0; i < NUM_QUEUES; i++)
        active_dispatches[i].push_back(std::make_shared<MockDispatch>(queues[i]));

    for(int i = 0; i < NUM_ACTIONS; i++)
    {
        int q      = rdgen() % NUM_QUEUES;
        int action = rdgen() % ACTION_MAX;
        if(action == 0)
        {
            // Delete queue and create new one
            active_dispatches[q] = {};
            queues[q].reset();
            queues[q] = std::make_shared<MockQueue>(QSIZE, buffer);
            num_reset_queues++;
        }
        else if(action > ACTION_MAX / 2 && active_dispatches[q].size() > 1)
        {
            // Delete dispatch
            active_dispatches[q].erase(active_dispatches[q].begin(),
                                       active_dispatches[q].begin() + 1);
        }

        // Add new dispatch
        if(active_dispatches[q].size() < QSIZE)
        {
            active_dispatches[q].push_back(std::make_shared<MockDispatch>(queues[q]));
            num_dispatches_generated += 1;
        }

        // Generate one "pc" sample for each queue
        buffer->genUpcomingSamples(NUM_QUEUES);
        for(auto& queue : active_dispatches)
        {
            assert(queue.size() > 0);
            std::shared_ptr<MockDispatch> rand_dispatch = queue[rdgen() % queue.size()];
            MockWave(rand_dispatch).genPCSample();
            num_samples_generated += 1;
            avg_q_occupancy += queue.size();
            max_q_occupancy = std::max(max_q_occupancy, queue.size());
        }
    }

    std::cout << "Hammer Stats: " << std::endl;
    std::cout << "num_reset_queues: " << num_reset_queues << std::endl;
    std::cout << "num_samples_generated: " << num_samples_generated << std::endl;
    std::cout << "num_dispatches_generated: " << num_dispatches_generated << std::endl;
    std::cout << "Avg queue occupancy: " << avg_q_occupancy / (NUM_ACTIONS * NUM_QUEUES)
              << std::endl;
    std::cout << "Max queue occupancy: " << max_q_occupancy << "\n\n" << std::endl;

    std::vector<std::pair<pcsample_v1_t*, uint64_t>> all_allocations;

    CHECK_PARSER(parse_buffer((generic_sample_t*) buffer->packets.data(),
                              buffer->packets.size(),
                              GFXIP_MAJOR,
                              alloc_callback,
                              (void*) &all_allocations));

    assert(all_allocations.size() == NUM_ACTIONS &&
           "QueueHammer test: Incorrect number of callbacks");
    for(auto sb = 0ul; sb < all_allocations.size(); sb++)
    {
        pcsample_v1_t* samples     = all_allocations[sb].first;
        size_t         num_samples = all_allocations[sb].second;

        assert(num_samples == NUM_QUEUES && "QueueHammer: Incorrect number of samples");
        assert(check_samples(samples, num_samples) &&
               "QueueHammer: parsed ID does not match correct ID");
        delete[] samples;
    }
}

TEST(pcs_parser_correlation_id, multi_buffer)
{
    std::shared_ptr<MockRuntimeBuffer> firstBuffer = std::make_shared<MockRuntimeBuffer>();
    std::shared_ptr<MockQueue>         queue       = std::make_shared<MockQueue>(16, firstBuffer);
    std::shared_ptr<MockDispatch>      dispatch1   = std::make_shared<MockDispatch>(queue);
    std::shared_ptr<MockDispatch>      dispatch2   = std::make_shared<MockDispatch>(queue);

    firstBuffer->genUpcomingSamples(4);
    MockWave(dispatch1).genPCSample();
    MockWave(dispatch2).genPCSample();
    MockWave(dispatch1).genPCSample();
    MockWave(dispatch2).genPCSample();

    std::shared_ptr<MockRuntimeBuffer> secondBuffer = std::make_shared<MockRuntimeBuffer>();
    const auto&                        packets      = firstBuffer->packets;
    secondBuffer->packets = std::vector<packet_union_t>(packets.begin() + 2, packets.end());

    std::vector<std::pair<pcsample_v1_t*, uint64_t>> all_allocations;

    CHECK_PARSER(parse_buffer((generic_sample_t*) firstBuffer->packets.data(),
                              firstBuffer->packets.size(),
                              GFXIP_MAJOR,
                              alloc_callback,
                              (void*) &all_allocations));
    CHECK_PARSER(parse_buffer((generic_sample_t*) secondBuffer->packets.data(),
                              secondBuffer->packets.size(),
                              GFXIP_MAJOR,
                              alloc_callback,
                              (void*) &all_allocations));

    assert(all_allocations.size() == 2 && "MultiBuffer: Incorrect number of callbacks");
    auto& sample = all_allocations[1];
    assert(sample.second == 4 && "MultiBuffer: Incorrect number of samples");
    assert(check_samples(sample.first, sample.second) &&
           "MultiBuffer: parsed ID does not match correct ID");

    delete[] all_allocations[0].first;
    delete[] all_allocations[1].first;
};

/**
 * Benchmarks how fast the parser can process samples on a single threaded case
 * Current: 5600X with -Ofast, up to >140 million samples/s or ~9GB/s R/W (18GB/s bidirectional)
 */
static void
Benchmark(bool bWarmup)
{
    constexpr size_t SAMPLE_PER_DISPATCH = 8192;
    constexpr size_t DISP_PER_QUEUE      = 12;
    constexpr size_t NUM_QUEUES          = MockDoorBell::num_unique_bells;

    std::shared_ptr<MockRuntimeBuffer> buffer = std::make_shared<MockRuntimeBuffer>();
    std::array<std::vector<std::shared_ptr<MockDispatch>>, NUM_QUEUES> active_dispatches;

    for(size_t q = 0; q < NUM_QUEUES; q++)
    {
        std::shared_ptr<MockQueue> queue = std::make_shared<MockQueue>(DISP_PER_QUEUE * 2, buffer);
        for(size_t d = 0; d < DISP_PER_QUEUE; d++)
            active_dispatches[q].push_back(std::make_shared<MockDispatch>(queue));
    }

    constexpr size_t TOTAL_NUM_SAMPLES = NUM_QUEUES * DISP_PER_QUEUE * SAMPLE_PER_DISPATCH;
    buffer->genUpcomingSamples(TOTAL_NUM_SAMPLES);

    for(auto& queue : active_dispatches)
        for(auto& dispatch : queue)
            for(size_t i = 0; i < SAMPLE_PER_DISPATCH; i++)
                MockWave(dispatch).genPCSample();

    std::pair<pcsample_v1_t*, size_t> userdata;
    userdata.first  = new pcsample_v1_t[TOTAL_NUM_SAMPLES];
    userdata.second = TOTAL_NUM_SAMPLES;

    auto t0 = std::chrono::system_clock::now();
    CHECK_PARSER(parse_buffer((generic_sample_t*) buffer->packets.data(),
                              buffer->packets.size(),
                              GFXIP_MAJOR,
                              [](pcsample_v1_t** sample, uint64_t size, void* userdata_) {
                                  auto* pair = reinterpret_cast<std::pair<pcsample_v1_t*, size_t>*>(
                                      userdata_);
                                  assert(TOTAL_NUM_SAMPLES == pair->second);
                                  *sample = pair->first;
                                  return size;
                              },
                              &userdata));
    auto  t1             = std::chrono::system_clock::now();
    float samples_per_us = float(TOTAL_NUM_SAMPLES) / (t1 - t0).count() * 1E3f;

    if(!bWarmup)
    {
        std::cout << "Benchmark: Parsed " << int(samples_per_us * 1E3f + 0.5f) * 1E-3f
                  << " Msample/s (";
        std::cout << int(sizeof(pcsample_v1_t) * samples_per_us) << " MB/s)" << std::endl;
    }

    delete[] userdata.first;
}

TEST(pcs_parser, benchmark)
{
    Benchmark(true);
    Benchmark(false);
    Benchmark(false);
    Benchmark(false);
}

class WaveSnapTest
{
public:
    WaveSnapTest()
    {
        buffer   = std::make_shared<MockRuntimeBuffer>();
        queue    = std::make_shared<MockQueue>(16, buffer);
        dispatch = std::make_shared<MockDispatch>(queue);
    }

    void Test()
    {
        FillBuffers();
        CheckBuffers();
    }

    virtual void FillBuffers()  = 0;
    virtual void CheckBuffers() = 0;

    void genPCSample(int wave_cnt, int inst_type, int reason, int arb_issue, int arb_stall)
    {
        wave_cnt &= 0x3F;
        inst_type &= 0xF;
        reason &= 0x7;
        arb_issue &= 0xFF;
        arb_stall &= 0xFF;

        perf_sample_snapshot_v1 snap;
        ::memset(&snap, 0, sizeof(snap));
        snap.pc             = dispatch->unique_id;
        snap.correlation_id = dispatch->getMockId();

        snap.perf_snapshot_data = (inst_type << 3) | (reason << 7);
        snap.perf_snapshot_data |= (arb_issue << 10) | (arb_stall << 18);
        snap.perf_snapshot_data1 = wave_cnt;

        assert(dispatch.get());
        dispatch->submit(packet_union_t{.snap = snap});
    };

    std::shared_ptr<MockRuntimeBuffer> buffer;
    std::shared_ptr<MockQueue>         queue;
    std::shared_ptr<MockDispatch>      dispatch;
};

class WaveCntTest : public WaveSnapTest
{
public:
    void FillBuffers() override
    {
        // Loop over all possible wave_cnt
        buffer->genUpcomingSamples(max_wave_number);
        for(size_t i = 0; i < max_wave_number; i++)
            genPCSample(i, GFX9::TYPE_LDS, GFX9::REASON_ALU, GFX9::ISSUE_VALU, GFX9::ISSUE_VALU);
    }

    void CheckBuffers() override
    {
        auto parsed = buffer->get_parsed_buffer(9);  // GFXIP==9
        assert(parsed.size() == 1);
        assert(parsed[0].size() == max_wave_number);

        for(size_t i = 0; i < max_wave_number; i++)
            assert(parsed[0][i].wave_count == i);
    }

    const size_t                        max_wave_number = 64;
    std::vector<pcsample_snapshot_v1_t> snapshots;
};

class InstTypeTest : public WaveSnapTest
{
public:
    void FillBuffers() override
    {
        // Loop over inst_type_issued
        UNROLL_TYPECHECK();
        buffer->genUpcomingSamples(GFX9::TYPE_LAST);
        for(int i = 0; i < GFX9::TYPE_LAST; i++)
            genPCSample(i, i, GFX9::REASON_ALU, GFX9::ISSUE_MATRIX, GFX9::ISSUE_MATRIX);
    }

    void CheckBuffers() override
    {
        auto parsed = buffer->get_parsed_buffer(9);  // GFXIP==9
        assert(parsed.size() == 1);
        assert(parsed[0].size() == GFX9::TYPE_LAST);
        assert(snapshots.size() == GFX9::TYPE_LAST);

        for(size_t i = 0; i < GFX9::TYPE_LAST; i++)
            assert(snapshots[i].inst_type == parsed[0][i].snapshot.inst_type);
    }

    std::vector<pcsample_snapshot_v1_t> snapshots;
};

class StallReasonTest : public WaveSnapTest
{
public:
    void FillBuffers() override
    {
        // Loop over reason_not_issued
        UNROLL_REASONCHECK();
        buffer->genUpcomingSamples(GFX9::REASON_LAST);
        for(int i = 0; i < GFX9::REASON_LAST; i++)
            genPCSample(i, GFX9::TYPE_MATRIX, i, GFX9::ISSUE_MATRIX, GFX9::ISSUE_MATRIX);
    }

    void CheckBuffers() override
    {
        auto parsed = buffer->get_parsed_buffer(9);  // GFXIP==9
        assert(parsed.size() == 1);
        assert(parsed[0].size() == GFX9::REASON_LAST);
        assert(snapshots.size() == GFX9::REASON_LAST);

        for(size_t i = 0; i < GFX9::REASON_LAST; i++)
            assert(snapshots[i].reason_not_issued == parsed[0][i].snapshot.reason_not_issued);
    }

    std::vector<pcsample_snapshot_v1_t> snapshots;
};

class ArbStateTest : public WaveSnapTest
{
public:
    void FillBuffers() override
    {
        // Loop over arb_state_issue
        UNROLL_ARBCHECK();
        buffer->genUpcomingSamples(GFX9::ISSUE_LAST * GFX9::ISSUE_LAST);
        for(int i = 0; i < GFX9::ISSUE_LAST; i++)
            for(int j = 0; j < GFX9::ISSUE_LAST; j++)
                genPCSample(i, GFX9::TYPE_MATRIX, GFX9::REASON_ALU, 1 << i, 1 << j);
    }

    void CheckBuffers() override
    {
        auto parsed = buffer->get_parsed_buffer(9);  // GFXIP==9
        assert(parsed.size() == 1);
        assert(parsed[0].size() == GFX9::ISSUE_LAST * GFX9::ISSUE_LAST);
        assert(snapshots.size() == GFX9::ISSUE_LAST * GFX9::ISSUE_LAST);

        for(size_t i = 0; i < GFX9::ISSUE_LAST * GFX9::ISSUE_LAST; i++)
        {
            auto& snap = snapshots[i];
            assert(snap.arb_state_issue == parsed[0][i].snapshot.arb_state_issue);
            assert(snap.arb_state_stall == parsed[0][i].snapshot.arb_state_stall);
        }
    }

    std::vector<pcsample_snapshot_v1_t> snapshots;
};

class WaveIssueAndErrorTest : public WaveSnapTest
{
    void FillBuffers() override
    {
        buffer->genUpcomingSamples(16);
        for(int valid = 0; valid <= 1; valid++)
            for(int issued = 0; issued <= 1; issued++)
                for(int dual = 0; dual <= 1; dual++)
                    for(int error = 0; error <= 1; error++)
                        genPCSample(valid, issued, dual, error);
    }

    void CheckBuffers() override
    {
        const int num_combinations = 16;
        auto      parsed           = buffer->get_parsed_buffer(9);  // GFXIP==9
        assert(parsed.size() == 1);
        assert(parsed[0].size() == num_combinations);
        assert(compare.size() == num_combinations);

        for(size_t i = 0; i < num_combinations; i++)
        {
            assert(compare[i].flags.valid == parsed[0][i].flags.valid);
            assert(compare[i].wave_issued == parsed[0][i].wave_issued);
            assert(compare[i].snapshot.dual_issue_valu == parsed[0][i].snapshot.dual_issue_valu);
        }
    }

    union trap_snapshot_v1
    {
        struct
        {
            uint32_t valid     : 1;
            uint32_t issued    : 1;
            uint32_t dual      : 1;
            uint32_t reserved  : 23;
            uint32_t error     : 1;
            uint32_t reserved2 : 5;
        };
        uint32_t raw;
    };

    void genPCSample(bool valid, bool issued, bool dual, bool error)
    {
        pcsample_v1_t sample;
        ::memset(&sample, 0, sizeof(sample));
        sample.pc             = dispatch->unique_id;
        sample.correlation_id = dispatch->getMockId();

        sample.flags.valid              = valid && !error;
        sample.wave_issued              = issued;
        sample.snapshot.dual_issue_valu = dual;

        assert(dispatch.get());

        compare.push_back(sample);

        trap_snapshot_v1 snap;
        snap.valid  = valid;
        snap.issued = issued;
        snap.dual   = dual;
        snap.error  = error;

        perf_sample_snapshot_v1 pss;
        pss.perf_snapshot_data = snap.raw;
        pss.correlation_id     = dispatch->getMockId();
        dispatch->submit(std::move(pss));
    };

    std::vector<pcsample_v1_t> compare;
};

class WaveOtherFieldsTest : public WaveSnapTest
{
    void FillBuffers() override
    {
        buffer->genUpcomingSamples(3);
        genPCSample(1, 2, 3, 4, 5, 6, 7, 8);       // Counting
        genPCSample(3, 5, 7, 11, 13, 17, 19, 23);  // Some prime numbers
        genPCSample(23, 19, 17, 13, 11, 7, 5, 3);  // Some reversed primes
    }

    void CheckBuffers() override
    {
        auto parsed = buffer->get_parsed_buffer(9);  // GFXIP==9
        assert(parsed.size() == 1);
        assert(parsed[0].size() == 3);
        assert(compare.size() == 3);

        for(size_t i = 0; i < 3; i++)
        {
            assert(parsed[0][i].flags.has_stall_reason == true);
            assert(parsed[0][i].flags.has_wave_cnt == true);
            assert(parsed[0][i].flags.has_memory_counter == false);

            assert(compare[i].exec_mask == parsed[0][i].exec_mask);
            assert(compare[i].workgroud_id_x == parsed[0][i].workgroud_id_x);
            assert(compare[i].workgroud_id_y == parsed[0][i].workgroud_id_y);
            assert(compare[i].workgroud_id_z == parsed[0][i].workgroud_id_z);

            assert(compare[i].chiplet == parsed[0][i].chiplet);
            assert(compare[i].wave_id == parsed[0][i].wave_id);
            assert(compare[i].hw_id == parsed[0][i].hw_id);
            assert(compare[i].correlation_id == parsed[0][i].correlation_id);
        }
    }

    void genPCSample(int pc, int exec, int blkx, int blky, int blkz, int chip, int wave, int hwid)
    {
        pcsample_v1_t sample;
        ::memset(&sample, 0, sizeof(sample));

        sample.exec_mask      = exec;
        sample.workgroud_id_x = blkx;
        sample.workgroud_id_y = blky;
        sample.workgroud_id_z = blkz;

        sample.chiplet        = chip;
        sample.wave_id        = wave;
        sample.hw_id          = hwid;
        sample.correlation_id = dispatch->unique_id;

        compare.push_back(sample);

        perf_sample_snapshot_v1 snap;
        ::memset(&snap, 0, sizeof(snap));
        snap.exec_mask = exec;

        snap.workgroud_id_x      = blkx;
        snap.workgroud_id_y      = blky;
        snap.workgroud_id_z      = blkz;
        snap.chiplet_and_wave_id = (chip << 8) | (wave & 0x3F);
        snap.hw_id               = hwid;
        snap.correlation_id      = dispatch->getMockId();

        assert(dispatch.get());
        dispatch->submit(snap);

        (void) pc;
    };

    std::vector<pcsample_v1_t> compare;
};

// FIXME (vladimir): For some reason, the test can stochastically fail.
// Did not have time to get into details.
TEST(pcs_parser, gfx9)
{
    WaveCntTest{}.Test();
    InstTypeTest{}.Test();
    StallReasonTest{}.Test();
    ArbStateTest{}.Test();
    WaveIssueAndErrorTest{}.Test();
    // FIXME: this might crash some time.
    // WaveOtherFieldsTest{}.Test();

    std::cout << "GFX9 Test Done." << std::endl;
}

// TODO: refactor the tests, modularize them and extract unit tests
// from the integration f