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

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

/*******************************************************************************************************
 * IMIX tool
 *
 * Count instruction mix - dynamic count of each instruction in the trace and
 *                         provide dynamic opcode histograms per data type
 *
 * NOTE: the tool callbacks might be called from different threads.
 */
#include <stdio.h>
#include <string.h>
#include <vector>

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

Knob<bool>  gKnobNoHist("no-hist", false, "Forbids creating per data type dynamic opcode histograms");
Knob<bool>  gKnobShowIp("show-ip", false, "Shows IPs of the instructions");

// Global variables 
uint32_t    gMaxNumOfHwThreads = 0;
uint32_t    gMaxNumOfTiles     = 0;
uint32_t    gNumOfInstructions = 0;
uint32_t    gNumOfOpcodes      = 0;
uint32_t    gNumOfDataTypes    = 0;

struct PerTileInstCounter {
    PerTileInstCounter()
    {
        count.resize(gMaxNumOfHwThreads);
        for (uint32_t i = 0; i < gMaxNumOfHwThreads; i++)
        {
            count[i].resize(gNumOfInstructions, 0);
        }
    }
    std::vector<std::vector<uint64_t>> count;
};

std::vector<PerTileInstCounter>    icount;
std::vector<uint64_t>              totalIcount;
std::vector<std::vector<uint64_t>> opcodes;
uint64_t    total_icount = 0;
std::string kernelName;

/*
 * 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(tileId < gMaxNumOfTiles&& tid < gMaxNumOfHwThreads);

    // Obtain the instruction ID within the kernel
    uint32_t id = GTReplay_InsId(ins);
    // Update the instruction counter corresponding to the current HW thread and current instruction
    icount[tileId].count[tid][id]++;
}

/*
 * 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)
{
    // accumulated icounts over HW threads 
    for (uint32_t i = 0; i < gNumOfInstructions; i++)
    {
        totalIcount[i] = 0;

        for (uint32_t tile = 0; tile < gMaxNumOfTiles; tile++)
        {
            for (uint32_t t = 0; t < gMaxNumOfHwThreads; t++)
            {
                totalIcount[i] += icount[tile].count[t][i];
            }
        }

        total_icount += totalIcount[i];
    }

    // Create opcode histograms
    for (uint32_t i = 0; i < gNumOfInstructions; i++)
    {
        GTReplayIns ins = GTReplay_Ins(kernel, i);
        uint32_t opcode = GTReplay_Opcode(ins);
        uint32_t dataType = GTReplay_DataType(ins);

        opcodes[opcode][dataType] += totalIcount[i];
    }

    std::cout << "\n\n=================\n";
    std::cout << "IMIX TOOL\n";
    std::cout << "=================\n\n";

    std::cout << "Kernel: " << kernelName << "\n\n";
    std::cout << "TOTAL ICOUNT = " << total_icount << "\n";

   
    for (uint32_t i = 0; i < gNumOfInstructions; i++)
    {
        GTReplayIns ins = GTReplay_Ins(kernel, i);
        if (gKnobShowIp)
        {
            uint32_t ip = GTReplay_GetInsIp(ins);
            std::cout << std::hex << ip << ":  ";
        }
        std::cout << "[" << std::dec << std::setw(15) << totalIcount[i] << "] " << std::string(GTReplay_Disasm(ins)) << std::endl;
    }

    if (!gKnobNoHist)
    {
        // Print opcode histograms per data type  
        for (uint32_t t = 0; t < gNumOfDataTypes; t++)
        {
            bool found = false;

            for (uint32_t o = 0; o < gNumOfOpcodes; o++)
            {
                if (opcodes[o][t])
                {
                    found = true;
                    break;
                }
            }
        
            if (!found)
            {
                continue;
            }

            std::cout << "\nDATA TYPE: " << std::string(GTReplay_DataTypeName(t)) << "\n\n";
        
            for (uint32_t o = 0; o < gNumOfOpcodes; o++)
            {
                if (opcodes[o][t])
                {
                    std::cout << std::setw(10) << std::string(GTReplay_OpcodeName(o)) << "   " << opcodes[o][t] << "\n";
                }
            }
        }
    }
}

/*
 * 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);

    gMaxNumOfHwThreads = GTReplay_MaxNumOfHWThreads(gModelId);

    gMaxNumOfTiles = GTReplay_MaxNumOfTiles(kernel);
    GTREPLAY_ASSERT(gMaxNumOfTiles);
	
    gNumOfInstructions = GTReplay_NumOfInstructions(kernel);

    // Traverse all the basic blocks 
    for (GTReplayBbl bbl = GTReplay_BblHead(kernel); GTReplay_BblValid(bbl); bbl = GTReplay_BblNext(bbl))
    {
        // Traverse all the instruction within the basic blocks 
        for (GTReplayIns ins = GTReplay_InsHead(bbl); GTReplay_InsValid(ins); ins = GTReplay_InsNext(ins))
        {
            // Register callback to be called before instruction execution
            GTReplay_RegisterCallbackBeforeIns(kernel, ins, BeforeInsCallback, NULL);
        }
    }

    // Allocate and initialize buffers
    icount.resize(gMaxNumOfTiles);
    totalIcount.resize(gNumOfInstructions, 0);

    gNumOfOpcodes = GTReplay_NumOpcodes();
    gNumOfDataTypes = GTReplay_NumDataTypes();

    opcodes.resize(gNumOfOpcodes);

    for (uint32_t i = 0; i < gNumOfOpcodes; i++)
    {
        opcodes[i].resize(gNumOfDataTypes+1, 0);
    }

    uint32_t kernelNameSize = 0;
    GTReplay_GetKernelName(kernel, &kernelNameSize, nullptr);

    char* buf = new char[kernelNameSize + 1]();
    GTReplay_GetKernelName(kernel, &kernelNameSize, buf);

    kernelName = std::string(buf);
    
    delete[] buf;
}

/*
 * 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();
}
