mem_trace.cpp

/*
 * mem_trace.cpp — CS204 PIN Tool
 * IIT Ropar | 2025-26
 *
 * Instruments a target binary and records every memory access to a
 * .trace file in the format expected by cache_sim:
 *
 *   R 0x<hex address>
 *   W 0x<hex address>
 *
 * Optionally with access size:   R 0x<addr> <size>
 *
 * For cross-program comparison, by default only accesses from the
 * MAIN EXECUTABLE are recorded — libc, libstdc++, ld-linux, etc. are
 * skipped. This makes traces comparable across benchmarks because the
 * runtime/loader noise (which dominates raw traces) is excluded.
 *
 * Usage:
 *   pin -t obj-intel64/mem_trace.so -o bench.trace -- ./bench
 *
 * Options:
 *   -o <file>      Output trace file              (default: mem.trace)
 *   -max <N>       Stop after N accesses          (default: 0 = unlimited)
 *   -rw <r|w|rw>   Record reads / writes / both   (default: rw)
 *   -flush <N>     Flush every N accesses         (default: 100000)
 *   -main_only <0|1>  1 = only main image, 0 = everything (default: 1)
 *   -size <0|1>    1 = append access size to each line (default: 0)
 *   -roi <0|1>     1 = require user ROI markers to start tracing (default: 0)
 *                  (call __trace_start() / __trace_stop() in your benchmark)
 */

#include "pin.H"

#include <iostream>
#include <fstream>
#include <cstdint>
#include <string>

// ─────────────────────────────────────────────
//  Command-line knobs
// ─────────────────────────────────────────────
KNOB<std::string> KnobOutputFile(
    KNOB_MODE_WRITEONCE, "pintool",
    "o", "mem.trace", "Output trace file path");

KNOB<UINT64> KnobMaxAccesses(
    KNOB_MODE_WRITEONCE, "pintool",
    "max", "0", "Stop after this many accesses (0 = unlimited)");

KNOB<std::string> KnobMode(
    KNOB_MODE_WRITEONCE, "pintool",
    "rw", "rw", "Record: r = reads only, w = writes only, rw = both");

KNOB<UINT64> KnobFlushEvery(
    KNOB_MODE_WRITEONCE, "pintool",
    "flush", "100000", "Flush trace file every N accesses");

KNOB<BOOL> KnobMainOnly(
    KNOB_MODE_WRITEONCE, "pintool",
    "main_only", "0",
    "Only record accesses from the main executable image");

KNOB<BOOL> KnobIncludeSize(
    KNOB_MODE_WRITEONCE, "pintool",
    "size", "0", "Append access size to each trace line");

KNOB<BOOL> KnobUseROI(
    KNOB_MODE_WRITEONCE, "pintool",
    "roi", "0",
    "Only trace between __trace_start() and __trace_stop() calls");

// ─────────────────────────────────────────────
//  Globals
// ─────────────────────────────────────────────
static std::ofstream TraceFile;
static UINT64        AccessCount  = 0;
static UINT64        MaxAccesses  = 0;
static UINT64        FlushEvery   = 100000;
static bool          RecordReads  = true;
static bool          RecordWrites = true;
static bool          MainOnly     = true;
static bool          IncludeSize  = false;
static bool          UseROI       = false;

// ROI state: if -roi 1, tracing only happens when TracingOn == true.
// Otherwise, TracingOn is forced true from the start.
static bool          TracingOn    = true;

// Cached main-image address range. Instruction is recorded only if its
// PC lies within [MainLow, MainHigh). Filled by ImageLoad callback.
static ADDRINT       MainLow      = 0;
static ADDRINT       MainHigh     = 0;

// ─────────────────────────────────────────────
//  ROI hooks — benchmark can define empty functions
//      extern "C" void __trace_start() {}
//      extern "C" void __trace_stop()  {}
//  We intercept them and flip TracingOn.
// ─────────────────────────────────────────────
static VOID OnTraceStart() { TracingOn = true;  }
static VOID OnTraceStop()  { TracingOn = false; }

// ─────────────────────────────────────────────
//  Analysis routines
// ─────────────────────────────────────────────
static VOID RecordRead(ADDRINT addr, UINT32 size) {
    if (!TracingOn) return;
    if (MaxAccesses && AccessCount >= MaxAccesses) return;
    TraceFile << "R 0x" << std::hex << (UINT64)addr;
    if (IncludeSize) TraceFile << " " << std::dec << size;
    TraceFile << "\n";
    ++AccessCount;
    if (FlushEvery && (AccessCount % FlushEvery) == 0) TraceFile.flush();
}

static VOID RecordWrite(ADDRINT addr, UINT32 size) {
    if (!TracingOn) return;
    if (MaxAccesses && AccessCount >= MaxAccesses) return;
    TraceFile << "W 0x" << std::hex << (UINT64)addr;
    if (IncludeSize) TraceFile << " " << std::dec << size;
    TraceFile << "\n";
    ++AccessCount;
    if (FlushEvery && (AccessCount % FlushEvery) == 0) TraceFile.flush();
}

// ─────────────────────────────────────────────
//  Image load — find main executable's range and hook ROI markers
// ─────────────────────────────────────────────
static VOID ImageLoad(IMG img, VOID* /* v */) {
    if (IMG_IsMainExecutable(img)) {
        MainLow  = IMG_LowAddress(img);
        MainHigh = IMG_HighAddress(img);
        std::cerr << "[mem_trace] Main image: " << IMG_Name(img)
                  << "  [0x" << std::hex << MainLow
                  << " - 0x" << MainHigh << "]\n" << std::dec;

        if (UseROI) {
            RTN startRtn = RTN_FindByName(img, "__trace_start");
            if (RTN_Valid(startRtn)) {
                RTN_Open(startRtn);
                RTN_InsertCall(startRtn, IPOINT_BEFORE,
                               (AFUNPTR)OnTraceStart, IARG_END);
                RTN_Close(startRtn);
                std::cerr << "[mem_trace] ROI: hooked __trace_start\n";
            }
            RTN stopRtn = RTN_FindByName(img, "__trace_stop");
            if (RTN_Valid(stopRtn)) {
                RTN_Open(stopRtn);
                RTN_InsertCall(stopRtn, IPOINT_BEFORE,
                               (AFUNPTR)OnTraceStop, IARG_END);
                RTN_Close(stopRtn);
                std::cerr << "[mem_trace] ROI: hooked __trace_stop\n";
            }
        }
    }
}

// ─────────────────────────────────────────────
//  Per-instruction instrumentation
// ─────────────────────────────────────────────
static VOID Instruction(INS ins, VOID* /* v */) {
    // Image filter: skip instructions outside the main executable
    if (MainOnly) {
        ADDRINT pc = INS_Address(ins);
        if (pc < MainLow || pc >= MainHigh) return;
    }

    if (MaxAccesses && AccessCount >= MaxAccesses) return;

    UINT32 memOperands = INS_MemoryOperandCount(ins);
    for (UINT32 op = 0; op < memOperands; ++op) {
        UINT32 size = INS_MemoryOperandSize(ins, op);

        if (RecordReads && INS_MemoryOperandIsRead(ins, op)) {
            INS_InsertPredicatedCall(
                ins, IPOINT_BEFORE, (AFUNPTR)RecordRead,
                IARG_MEMORYOP_EA, op,
                IARG_UINT32, size,
                IARG_END);
        }
        if (RecordWrites && INS_MemoryOperandIsWritten(ins, op)) {
            INS_InsertPredicatedCall(
                ins, IPOINT_BEFORE, (AFUNPTR)RecordWrite,
                IARG_MEMORYOP_EA, op,
                IARG_UINT32, size,
                IARG_END);
        }
    }
}

// ─────────────────────────────────────────────
//  Fini callback
// ─────────────────────────────────────────────
static VOID Fini(INT32 /* code */, VOID* /* v */) {
    TraceFile.flush();
    TraceFile.close();
    std::cerr << std::dec
              << "[mem_trace] Done.\n"
              << "  Accesses recorded : " << AccessCount << "\n"
              << "  Output file       : " << KnobOutputFile.Value() << "\n";
}

// ─────────────────────────────────────────────
//  Entry point
// ─────────────────────────────────────────────
INT32 Usage() {
    std::cerr << "CS204 Memory Trace PIN Tool\n\n"
              << KNOB_BASE::StringKnobSummary() << "\n";
    return -1;
}

int main(int argc, char* argv[]) {
    // PIN_InitSymbols is REQUIRED for RTN_FindByName (ROI hooks) to work
    PIN_InitSymbols();

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

    MaxAccesses  = KnobMaxAccesses.Value();
    FlushEvery   = KnobFlushEvery.Value();
    MainOnly     = KnobMainOnly.Value();
    IncludeSize  = KnobIncludeSize.Value();
    UseROI       = KnobUseROI.Value();

    std::string mode = KnobMode.Value();
    RecordReads  = (mode.find('r') != std::string::npos);
    RecordWrites = (mode.find('w') != std::string::npos);

    // If ROI is on, start with tracing OFF (benchmark turns it on)
    TracingOn = !UseROI;

    TraceFile.open(KnobOutputFile.Value().c_str());
    if (!TraceFile.is_open()) {
        std::cerr << "[mem_trace] ERROR: cannot open output file: "
                  << KnobOutputFile.Value() << "\n";
        return 1;
    }

    TraceFile << "# CS204 Memory Trace | IIT Ropar\n"
              << "# Format: <R|W> 0x<address>"
              << (IncludeSize ? " <size>" : "") << "\n"
              << "# main_only=" << MainOnly
              << "  roi=" << UseROI
              << "  mode=" << mode << "\n";
    TraceFile.flush();

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

    std::cerr << "[mem_trace] Tracing started -> "
              << KnobOutputFile.Value()
              << "  (main_only=" << MainOnly
              << ", roi=" << UseROI
              << ", size=" << IncludeSize << ")\n";

    PIN_StartProgram();   // never returns
    return 0;
}