cuMemcpy2d_trace.cpp

cuMemcpy2d_trace.cpp

#include <stdio.h>
#include "pin.H"

static FILE *g_trace;
static BOOL g_enabled;
static ADDRINT g_startAddr;
static ADDRINT g_endAddrLow;
static ADDRINT g_endAddrHeight;
static ADDRINT g_ret;

static void printip(ADDRINT v, char *dis, BOOL taken, ADDRINT dest, ADDRINT next)
{
    if (dest == g_startAddr)
    {
        fprintf(g_trace, "ENABLED: 0x%lx: %s (dest: 0x%lx, next: 0x%lx)\n",
            (unsigned long)v, dis, (unsigned long)dest, (unsigned long)next);
        g_enabled = true;
        g_ret = next;
    }
    else if (dest == g_ret)
    {
        fprintf(g_trace, "DISABLED: 0x%lx: %s (dest: 0x%lx, next: 0x%lx)\n",
            (unsigned long)v, dis, (unsigned long)dest, (unsigned long)next);
        g_enabled = false;
        g_ret = -1;
    }

    if (!g_enabled) return;

    fprintf(g_trace, "0x%lx: %s (dest: 0x%lx, taken: %d)\n",
        (unsigned long)v, dis, dest, taken);
}

static void Instruction(INS ins, VOID *v)
{
    auto *st = new std::string(INS_Disassemble(ins));

    if (!INS_IsControlFlow(ins))
    {
        //INS_InsertCall(ins, IPOINT_BEFORE, (AFUNPTR)allInstructions,
        //    IARG_INST_PTR, IARG_PTR, st->c_str(), IARG_RETURN_IP , IARG_END);

        return;
    }

    ADDRINT next = INS_NextAddress(ins);

    INS_InsertCall(ins, IPOINT_BEFORE, (AFUNPTR)printip, IARG_INST_PTR,
       IARG_PTR, st->c_str(), IARG_BRANCH_TAKEN, IARG_BRANCH_TARGET_ADDR, IARG_ADDRINT, next, IARG_END);
}


static void ImageLoad(IMG img, void *v)
{
    std::string searchFor("C:\\Windows\\SYSTEM32\\nvcuda.dll");
    std::string current(IMG_Name(img));

    if (IMG_IsMainExecutable(img))
    {
        g_endAddrLow = IMG_LowAddress(img);
        g_endAddrHeight = IMG_LowAddress(img);

        fprintf(g_trace, "Base image: %s\n", current.c_str());
        return;
    }

    if (searchFor.compare(current))
    {
        //fprintf(g_trace, "Ignoring image: %s\n", current.c_str());
        return;
    }

    std::string symSearchFor("cuMemcpy2D_v2");

    fprintf(g_trace, "Parsing image: %s\n", current.c_str());

    for (SYM sym = IMG_RegsymHead(img); SYM_Valid(sym); sym = SYM_Next(sym))
    {
        std::string symName = SYM_Name(sym);
        auto symAddress = SYM_Address(sym);

        //fprintf(g_trace, "Symbol %s address 0x%lx\n", symName.c_str(), (unsigned long)symAddress);

        if (!symSearchFor.compare(symName))
        {
            g_startAddr = symAddress;
            fprintf(g_trace, "Symbol start set to %s address 0x%lx\n", symName.c_str(), (unsigned long)symAddress);
        }

        RTN rtn = RTN_FindByName(img, SYM_Name(sym).c_str());
        if (!RTN_Valid(rtn))
        {
            fprintf(g_trace, "Routine not found, continue...\n");
            continue;
        }

        //fprintf(g_trace, "Routine %s address 0x%lx, size 0x%lx\n",
        //    RTN_Name(rtn).c_str(), (unsigned long)RTN_Address(rtn), (unsigned long)RTN_Size(rtn));
    }
}

static void Fini(INT32 code, VOID *v)
{
    fprintf(g_trace, "#eof\n");
    fclose(g_trace);
}

static INT32 Usage()
{
    PIN_ERROR("This Pintool prints the IPs of every instruction executed\n"
              + KNOB_BASE::StringKnobSummary() + "\n");
    return -1;
}

int main(int argc, char * argv[])
{
    g_trace = fopen("ig_trace.out", "w");

    if (PIN_Init(argc, argv)) return Usage();

    PIN_InitSymbols();

    INS_AddInstrumentFunction(Instruction, 0);
    IMG_AddInstrumentFunction(ImageLoad, 0);
    PIN_AddFiniFunction(Fini, 0);

    PIN_StartProgram();

    return 0;
}