/*========================== begin_copyright_notice ============================
Copyright (C) 2021-2022 Intel Corporation

SPDX-License-Identifier: MIT
============================= end_copyright_notice ===========================*/

/*******************************************************************************************************
 * BBLTRACE tool
 *
 * Collect dynamic control flow trace.
 *
 * NOTE: the tool callbacks might be called from different threads.
 */
#include <vector>
#include <list>
#include <stdio.h>
#include <string.h>

#include "gtreplay_assert.h"
#include "gtreplay_client.h"
#include "knob_parser.h"


// Global variables 
uint32_t maxNumOfHwThreads = 0;
uint32_t numOfInstructions = 0;

std::vector<std::list<uint32_t>> bbltrace;

/*
 * BeforeInsCallback - callback called before instruction execution
 *
 * @params[in] tid - the ID of the GPU HW thread for which the callback is called
 * @params[in] ins - a handle to the current instruction
 * @params[in] state - a handle to the HW Thread state corresponding to tid
 */
void BeforeInsCallback(uint32_t tileId, uint32_t tid, GTReplayIns ins, GTReplayState state, void*)
{
    GTREPLAY_ASSERT(tid < maxNumOfHwThreads);
     // Add current BBL ID to the trace
	 bbltrace[tid].push_back(GTReplay_BblIdFromIns(ins));
}

/*
 * AfterInsCallback - callback called after instruction execution
 *                    An illustration - this tool doesn't need to register after instruction callbacks
 *
 * @params[in] tid - the ID of the GPU HW thread for which the callback is called
 * @params[in] ins - a handle to the current instruction
 * @params[in] state - a handle to the HW Thread state corresponding to tid
 */
void AfterInsCallback(uint32_t tileId, uint32_t tid, GTReplayIns ins, GTReplayState state, void*)
{
}

/*
 * OnKernelComplete - callback called upon kernel completion
 *
 * @params[in] kernel - a handle to the kernel
 */
void OnKernelComplete(GTReplayKernel kernel)
{
    std::cout << "\n\n=================\n";
    std::cout << "BBLTRACE TOOL\n";
    std::cout << "=================\n\n";

    // Iterate over all available HW threads
    for (uint32_t tid = 0; tid < maxNumOfHwThreads; tid++)
    {
	    // If the trace is not empty
        if (!bbltrace[tid].empty())
        {
		   // Print the trace as a sequence of BBL IDs
            std::cout << "===========TID = " << tid << " ===========\n";

            bool first = true;
            for (auto& b : bbltrace[tid])
            {
                if ((b == 0) && !first) { std::cout << "EOT;\n"; } first = false;
                std::cout << b << "->";
            }
            std::cout << "EOT;\n\n";
        }

        // Clean the trace
        bbltrace[tid].clear();
    }

    bbltrace.clear();
}

/*
 * OnKernelBuild - callback called before kernel execution
 *                 The purpose of this callback is to traverse the kernel binary and instrument callbacks
 *
 * @params[in] kernel - a handle to the kernel
 */
void OnKernelBuild(GTReplayKernel kernel)
{
    uint32_t gModelId = GTReplay_GetModel(kernel);

    maxNumOfHwThreads = GTReplay_MaxNumOfHWThreads(gModelId);

    // Initialize the data structure
    bbltrace.resize(maxNumOfHwThreads);

    // Traverse all the basic blocks 
    for (GTReplayBbl bbl = GTReplay_BblHead(kernel); GTReplay_BblValid(bbl); bbl = GTReplay_BblNext(bbl))
    {
        GTReplayIns ins = GTReplay_InsHead(bbl);

        // Register callback to be called before the head instruction of the basic block
        GTReplay_RegisterCallbackBeforeIns(kernel, ins, BeforeInsCallback, NULL);
    }
}

/*
 * GTReplay_Entry - tool entry point
 */
extern "C"
DLLEXP void FASTCALL GTReplay_Entry(int argc, const char *argv[])
{
    // configure GTReplay
    ConfigureGTReplay(argc, argv);
    
    // register OnKernelBuild and OnKernelComplete callbacks
    GTReplay_RegisterOnKernelBuildCallback(OnKernelBuild);
    GTReplay_RegisterOnKernelCompleteCallback(OnKernelComplete);

    // Start GTReplay
    GTReplay_Start();
}
