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

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

/*******************************************************************************************************
 * TOGGLE tool
 *
 * Count dynamic amount of toggling bits - the ones that changed their values 0->1 and 1->0
 *
 * NOTE: the tool callbacks might be called from different threads.
 */
#include <stdio.h>
#include <string.h>
#include <vector>
#ifdef TARGET_WINDOWS
#include "intrin.h"
#endif
#ifdef TARGET_LINUX
#include "x86intrin.h"
#endif

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

// Structure definitions
typedef union {
    uint8_t  byte[32];
    uint16_t word[32];
    uint32_t dword[32];
    uint64_t qword[32];
    int8_t   sbyte[32];
    int16_t  sword[32];
    int32_t  sdword[32];
    int64_t  sqword[32];
    float    spfloat[32];
    double   dpfloat[32];
} Operand;

typedef struct {
    uint32_t dword[8];
} FullReg32;

typedef struct {
    uint32_t dword[16];
} FullReg64;

typedef union {
    FullReg32 reg32[16];
    FullReg64 reg64[16];
} SendDest;

// Global variables 
uint32_t   gMaxNumOfHwThreads = 0;
uint32_t   gMaxNumOfTiles = 0;
uint32_t   gRegWidth = 32;

uint64_t  total_icount = 0;
uint64_t  total_toggle_bits = 0;
std::vector<std::vector<uint64_t>> icount;
std::vector<std::vector<SendDest>> sendRegsBefore;
std::vector<std::vector<Operand>>  dstBefore;
std::vector<std::vector<uint32_t>> execMask;
std::vector<std::vector<uint64_t>> toggledBits;
std::string kernelName;

void HandleSendBefore(uint32_t tileId, uint32_t tid, GTReplayIns ins, GTReplayState state)
{
    uint32_t numOfElements = 0;
    uint32_t elementWidth = 0;

    // Obtain and save the registers
    GTReplay_GetSendDestination(ins, state, (uint8_t*)&sendRegsBefore[tileId][tid], &numOfElements, &elementWidth);
}

void HandleSendAfter(uint32_t tileId, uint32_t tid, GTReplayIns ins, GTReplayState state)
{
    SendDest dst;
    uint32_t numOfElements = 0;
    uint32_t elementWidth = 0;

    // Obtain the registers
    GTReplay_GetSendDestination(ins, state, (uint8_t*)&dst, &numOfElements, &elementWidth);

    if (numOfElements == 0)
    {
        return;
    }

    uint32_t count = 0;

    if (gRegWidth == 32)
    {
        // go over all registers
        for (uint32_t i = 0; i < numOfElements; i++)
        {
            FullReg32 afreg = dst.reg32[i], bfreg = sendRegsBefore[tileId][tid].reg32[i];

            // go over all elements
            for (uint32_t j = 0; j < 8; j++)
            {
                // compute the amount of toggled bits
                uint32_t tmp = bfreg.dword[j] ^ afreg.dword[j];

                count += _mm_popcnt_u32(tmp);
            }
        }
    }
    else
    {
        // go over all registers
        for (uint32_t i = 0; i < numOfElements; i++)
        {
            FullReg64 afreg = dst.reg64[i], bfreg = sendRegsBefore[tileId][tid].reg64[i];

            // go over all elements
            for (uint32_t j = 0; j < 16; j++)
            {
                // compute the amount of toggled bits
                uint32_t tmp = bfreg.dword[j] ^ afreg.dword[j];

                count += _mm_popcnt_u32(tmp);
            }
        }
    }

    toggledBits[tileId][tid] += count;
}

/*
 * 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);
    // Update the instruction counter corresponding to the current HW thread
    icount[tileId][tid]++;

    // Check whether the instruction has destination
    if (!GTReplay_HasDestination(ins))
    {
	    // If not, there is nothing to do
        return;
    }

    // Check whether the instruction is a SEND instruction
    if (GTReplay_IsSend(ins))
    {
        // Is yes, handle SEND instruction separately
        HandleSendBefore(tileId, tid, ins, state);
        return;
    }

    // Obtain exec mask
    execMask[tileId][tid] = GTReplay_DynamicExecMask(ins, state);
    
    uint32_t numOfElements = 0;
    uint32_t elementWidth = 0;

    // Obtain and save destination before
    GTReplay_GetDestination(ins, state, execMask[tileId][tid], (uint8_t*)&dstBefore[tileId][tid], &numOfElements, &elementWidth);
}

/*
 * AfterInsCallback - callback called after 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 AfterInsCallback(uint32_t tileId, uint32_t tid, GTReplayIns ins, GTReplayState state, void*)
{
    GTREPLAY_ASSERT(tileId < gMaxNumOfTiles&& tid < gMaxNumOfHwThreads);

    // Check whether the instruction has destination
    if (!GTReplay_HasDestination(ins))
    {
	    // If not, there is nothing to do
        return;
    }

    // Check whether the instruction is a SEND instruction
    if (GTReplay_IsSend(ins))
    {
        // Is yes, handle SEND instruction separately
        HandleSendAfter(tileId, tid, ins, state);
        return;
    }

    Operand  dstAfter = {};
    uint32_t numOfElements = 0;
    uint32_t elementWidth = 0;

    // Obtain destination after
    GTReplay_GetDestination(ins, state, execMask[tileId][tid], (uint8_t*)&dstAfter, &numOfElements, &elementWidth);

    uint32_t count = 0;

    // Iterate over all elements
    for (uint32_t i = 0; i < numOfElements; i++)
    {
        // Compute the amount of toggled bits
        switch (elementWidth) {
        case 1: count += _mm_popcnt_u32(dstBefore[tileId][tid].byte[i] ^ dstAfter.byte[i]); break;
        case 2: count += _mm_popcnt_u32(dstBefore[tileId][tid].word[i] ^ dstAfter.word[i]); break;
        case 4: count += _mm_popcnt_u32(dstBefore[tileId][tid].dword[i] ^ dstAfter.dword[i]); break;
        case 8:
        {
            uint64_t tmp = dstBefore[tileId][tid].qword[i] ^ dstAfter.qword[i];

            count += _mm_popcnt_u32((uint32_t)(tmp & 0xFFFFFFFF));
            count += _mm_popcnt_u32((uint32_t)(tmp >> 32));

            break;
        }
        default: break;
        }
    }

    toggledBits[tileId][tid] += count;
}

/*
 * OnKernelComplete - callback called upon kernel completion
 *
 * @params[in] kernel - a handle to the kernel
 */
 void OnKernelComplete(GTReplayKernel kernel)
{
    total_icount = 0;
    total_toggle_bits = 0;

    // Accumulate counters from all HW threads
    for (uint32_t tileId = 0; tileId < gMaxNumOfTiles; tileId++)
    {
        for (uint32_t t = 0; t < gMaxNumOfHwThreads; t++)
        {
            total_icount += icount[tileId][t];
            total_toggle_bits += toggledBits[tileId][t];
        }
    }

    // Print the results
    std::cout << "\n\n=================\n";
    std::cout << "BIT TOGGLING TOOL\n";
    std::cout << "=================\n\n";
    std::cout.imbue(std::locale(""));
    std::cout << "Kernel: " << kernelName << "\n\n";
    std::cout << "TOTAL ICOUNT       = " << total_icount << "\n\n";
    std::cout << "TOTAL TOGGLED BITS = " << total_toggle_bits << "\n\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);

    gRegWidth = GTReplay_RegisterWidth(gModelId);

    // 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);
            // Register callback to be called after instruction execution
            GTReplay_RegisterCallbackAfterIns(kernel, ins, AfterInsCallback, NULL);
        }
    }

    // Allocate and initialize buffers
    icount.resize(gMaxNumOfTiles);
    toggledBits.resize(gMaxNumOfTiles);
    execMask.resize(gMaxNumOfTiles);
    sendRegsBefore.resize(gMaxNumOfTiles);
    dstBefore.resize(gMaxNumOfTiles);
    for (uint32_t i = 0; i < gMaxNumOfTiles; i++)
    {
        icount[i].resize(gMaxNumOfHwThreads, 0);
        toggledBits[i].resize(gMaxNumOfHwThreads, 0);
        execMask[i].resize(gMaxNumOfHwThreads, 0);
        sendRegsBefore[i].resize(gMaxNumOfHwThreads);
        dstBefore[i].resize(gMaxNumOfHwThreads);
    }

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