#!/bin/bash
set -u

DIR="$(cd "$(dirname "$(readlink -f "${BASH_SOURCE[0]}")")" && pwd)"
MYSQLITE3="$DIR"/sqlite3.26.0

# check number of arguments
if [ $# -ne 1 ]; then
    echo " usage: cudaKernelStats [SQLite DB file exported from NSys QDREP file]"
    exit 1
fi

# check if file exists
if [ ! -f "$1" ]
then
    echo "$1 file not found. Exiting"
    exit 1
fi

# check if file opened is DB; if not, exit.
# The sqlite3 file format is defined at https://www.sqlite.org/draft/fileformat.html
HEADER=$(head -c 16 "$1")
if [ "$HEADER" != "SQLite format 3" ]
then
    echo "$1 is not an SQlite DB file. Exiting."
    exit 1
fi

# check if DB contains cuda kernel data; if not, exit
COUNT=$("$MYSQLITE3" "$1" "SELECT COUNT(*) FROM CUPTI_ACTIVITY_KIND_KERNEL");
if [ $COUNT -eq 0 ]
then
    echo "CUDA trace data was not collected."
    exit 0
fi

# check if cudaKernelStats table exists
if [ "$("$MYSQLITE3" "$1" "SELECT name FROM sqlite_master WHERE type='table' AND name='cudaKernelStats';")" != "cudaKernelStats" ]
then
    printf "\nGenerating cuda Kernel and Memory Operation Statistics...\n"

    "$MYSQLITE3" "$1" <<GenerateSummary
    PRAGMA SYNCHRONOUS=OFF;

    CREATE TABLE MemcpyOperationStrings (id INTEGER, name TEXT);
    INSERT INTO MemcpyOperationStrings (id, name) VALUES (0, '[CUDA memcpy Unknown]'),
        (1, '[CUDA memcpy HtoD]'), (2, '[CUDA memcpy DtoH]'), (3, '[CUDA memcpy HtoA]'),
        (4, '[CUDA memcpy AtoH]'), (5, '[CUDA memcpy AtoA]'), (6, '[CUDA memcpy AtoD]'),
        (7, '[CUDA memcpy DtoA]'), (8, '[CUDA memcpy DtoD]'), (9, '[CUDA memcpy HtoH]'),
        (10, '[CUDA memcpy PtoP]'), (11, '[CUDA Unified Memory memcpy HtoD]'),
        (12, '[CUDA Unified Memory memcpy DtoH]'), (13, '[CUDA Unified Memory memcpy DtoD]');

    CREATE TABLE cudaMemoryOperationTimeStats (num INTEGER, min INTEGER, max INTEGER, avg INTEGER, total INTEGER, name TEXT);

    INSERT INTO cudaMemoryOperationTimeStats SELECT count(copyKind), min(end-start), max(end-start), avg(end-start), sum(end-start), name as Name
        FROM CUPTI_ACTIVITY_KIND_MEMCPY
        INNER JOIN MemcpyOperationStrings ON MemcpyOperationStrings.id = CUPTI_ACTIVITY_KIND_MEMCPY.copyKind
        GROUP BY copyKind;

    INSERT INTO cudaMemoryOperationTimeStats SELECT count(*), min(end-start), max(end-start), avg(end-start), sum(end-start), '[CUDA memset]'
        FROM CUPTI_ACTIVITY_KIND_MEMSET;

    CREATE TABLE cudaMemoryOperationByteStats (num INTEGER, min INTEGER, max INTEGER, avg INTEGER, total INTEGER, name TEXT);
    INSERT INTO cudaMemoryOperationByteStats SELECT count(copyKind), min(bytes), max(bytes), avg(bytes), sum(bytes), name as Name
        FROM CUPTI_ACTIVITY_KIND_MEMCPY
        INNER JOIN MemcpyOperationStrings ON MemcpyOperationStrings.id = CUPTI_ACTIVITY_KIND_MEMCPY.copyKind
        GROUP BY copyKind;

    INSERT INTO cudaMemoryOperationByteStats SELECT count(*), min(bytes), max(bytes), avg(bytes), sum(bytes), '[CUDA memset]'
        FROM CUPTI_ACTIVITY_KIND_MEMSET;

    CREATE TABLE cudaKernelStats (shortName INTEGER, num INTEGER, min INTEGER, max INTEGER, avg INTEGER, total INTEGER);
    INSERT INTO cudaKernelStats SELECT shortName, count(shortName), min(end-start), max(end-start), avg(end-start), sum(end-start)
        FROM CUPTI_ACTIVITY_KIND_KERNEL GROUP BY shortName;
GenerateSummary
fi

printf "cuda Kernel Statistics\n\n"
TOTALTIME=$("$MYSQLITE3" "$1" "SELECT sum(total) FROM cudaKernelStats");
echo -e ".width 0 12 0 14 12 12 80 \n SELECT round((total*100.0)/$TOTALTIME,1) as 'Time(%)',
    total as 'Time (ns)', num as Instances, round(avg,1) as 'Avg (ns)', min as 'Min (ns)', max as 'Max (ns)', value as Name
    FROM cudaKernelStats INNER JOIN StringIds ON StringIds.id = cudaKernelStats.shortName ORDER BY total DESC;" | "$MYSQLITE3" -column -header "$1"

printf "\n\ncuda Memory Operation Statistics (time)\n\n"
TOTALTIME=$("$MYSQLITE3" "$1" "SELECT sum(total) FROM cudaMemoryOperationTimeStats");
echo -e ".width 0 12 0 14 12 12 80 \n SELECT round((total*100.0)/$TOTALTIME,1) as 'Time(%)',
    total as 'Time (ns)', num as Operations, round(avg,1) as 'Avg (ns)', min as 'Min (ns)', max as 'Max (ns)', name as Name
    FROM cudaMemoryOperationTimeStats WHERE num > 0 ORDER BY total DESC;" | "$MYSQLITE3" -column -header "$1"

printf "\n\ncuda Memory Operation Statistics (bytes)\n\n"
echo -e ".width 0 14 16 16 16 80 \n SELECT
    round(total/1024,4) as 'Total Bytes (KB)', num as Operations, round(avg/1024,4) as 'Avg (KB)', min as 'Min (bytes)',
    round(max/1024,4) as 'Max (KB)', name as Name
    FROM cudaMemoryOperationByteStats WHERE num > 0 ORDER BY total DESC;" | "$MYSQLITE3" -column -header "$1"

printf "\n\n"
