#!/bin/bash
set -u

function print_usage()
{
    cat << EOF
$(basename "$0") -- GPU Memory Operations Summary (by Size)

    No arguments.

    Output: All memory values given in KiB
        Total : Total number of KiB utilized by this operation
        Operations : Number of executions of this operation
        Average : The average memory size of this operation
        Minimum : The smallest memory size of this operation
        Maximum : The largest memory size of this operation
        Name : The name of the operation

    This report provides a summary of GPU memory operations and
    the amount of memory they utilize.
EOF
}

### BEGIN include inc_setup ###

EXIT_HELP=25
EXIT_DB=26
EXIT_NODATA=27

# Verify number of params
if [ $# -lt 1 ]
then
    print_usage ${BASH_SOURCE[0]}
    exit ${EXIT_HELP}
fi

# Set DB file
DATABASE="$1"

# Verify DB file exists
if [ ! -f "${DATABASE}" ]
then
    exit ${EXIT_DB}
fi

# Verify DB file contents
# The sqlite3 file format is defined at https://sqlite.org/fileformat.html
DB_FILE_HEADER=$(head -c 16 "$DATABASE" | tr '\0' '\n')
if [ "${DB_FILE_HEADER}" != "SQLite format 3" ]
then
    exit ${EXIT_DB}
fi

# Helper function for error messages
function echoerr() # accepts multiple args
{
    echo "$@" >&2
}

# Setup standard vars

# If we were run by nsys, the path to the preferred sqlite3 should have been
# passed as an env-var.  If not, hope the user has it in their path.
SQLITE3="${NSYS_STATS_SCRIPTS_SQLITE:-sqlite3}"
SQLITE3OPTS="-header -csv -readonly"

RUN_SQLITE="eval \"${SQLITE3}\" ${SQLITE3OPTS} \"${DATABASE}\""

### END include inc_setup ###

### BEGIN: include from inc_table_exists ###

TABLE_EXISTS_TABLES=( )

function table_exists()
{
    local TABLE_NAME=$1

    if [ "${#TABLE_EXISTS_TABLES[@]}" -eq 0 ]
    then
        TABLE_EXISTS_TABLES=( $("${SQLITE3}" ${SQLITE3OPTS} "${DATABASE}" \
                "SELECT name FROM sqlite_master WHERE type = 'table' OR type = 'view'") )
    fi

    for TABLE in "${TABLE_EXISTS_TABLES[@]}"
    do
        if [ "${TABLE}" = "${TABLE_NAME}" ]
        then
            echo "true"
            return 1
        fi
    done
    echo "false"
    return 0
}

### END: include from inc_table_exists ###

### BEGIN include inc_helper_cte ###

MemKindStrsCTE="""
    MemKindStrs (id, name) AS (
    VALUES
        (0,     'Pageable'),
        (1,     'Pinned'),
        (2,     'Device'),
        (3,     'Array'),
        (4,     'Unknown')
    ),
"""

MemcpyOperStrsCTE="""
    MemcpyOperStrs (id, name) AS (
    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]')
    ),
"""

### END include inc_helper_cte ###


QUERY=()

if $(table_exists "CUPTI_ACTIVITY_KIND_MEMCPY")
then
    if [ ${#QUERY[@]} -gt 0 ]
    then
        QUERY+=("UNION ALL")
    fi

    Q="""
        SELECT
            mos.name AS name,
            mcpy.bytes AS size
        FROM
            CUPTI_ACTIVITY_KIND_MEMCPY as mcpy
        INNER JOIN
            MemcpyOperStrs AS mos
            ON mos.id = mcpy.copyKind
    """
    QUERY+=("$Q")
fi

if $(table_exists "CUPTI_ACTIVITY_KIND_MEMSET")
then
    if [ ${#QUERY[@]} -gt 0 ]
    then
        QUERY+=("UNION ALL")
    fi

    Q="""
        SELECT
            '[CUDA memset]' AS name,
            bytes AS size
        FROM
            CUPTI_ACTIVITY_KIND_MEMSET
    """
    QUERY+=("$Q")
fi

if [ ${#QUERY[@]} -eq 0 ]
then
    echoerr "$DATABASE does not contain CUDA memory data."
    exit ${EXIT_NODATA}
fi

${RUN_SQLITE} << EOF

WITH
    ${MemcpyOperStrsCTE}
    memops AS (
        ${QUERY[@]}
    ),
    summary AS (
        SELECT
            name AS name,
            sum(size) AS total,
            count(*) AS num,
            avg(size) AS avg,
            min(size) AS min,
            max(size) AS max
        FROM memops
        GROUP BY 1
    )
SELECT
    printf('%.3f', summary.total/1024.0, 3) AS "Total",
    summary.num AS "Operations",
    printf('%.3f', summary.avg/1024.0, 3) AS "Average",
    printf('%.3f', summary.min/1024.0, 3) AS "Minimum",
    printf('%.3f', summary.max/1024.0, 3) AS "Maximum",
    summary.name AS "Operation"
FROM
    summary
ORDER BY 1 DESC
;

EOF
