#!/usr/bin/env python

import gpustats
import time

class GPULowUtilization(gpustats.GPUOperation):

    ROW_LIMIT = 50
    THRESHOLD = 50
    NUM_CHUNKS = 30

    usage = f"""{{SCRIPT}}[:rows=<limit>][:start=<time>][:end=<time>][:nvtx=<range[@domain]>][:threshold=<percent>][:chunks=<number>] -- GPU Low Utilization

    Options:
        rows=<limit> - Maximum number of rows returned by the query.
            Default is {ROW_LIMIT}.

        start=<time> - Start time used for filtering in nanoseconds.

        end=<time> - End time used for filtering in nanoseconds.

        nvtx=<range[@domain]> - NVTX range text and domain used for filtering.
            Do not specify the domain for ranges in the default domain. Note
            that only the first matching record will be considered.
            If this option is used along with the 'start' and/or 'end' options,
            the explicit start/end times will override the NVTX range times.

        threshold=<percent> - Maximum percentage of time the GPU is being used.
            Default is {THRESHOLD}%.

        chunks=<number> - Number of equal-duration chunks in which the GPU time
            utilization is calculated for each device in a process. If the profile
            duration cannot be evenly divided by <number>, the duration of every
            chunk is rounded up to the nearest integer nanosecond, such that all
            chunks still have the same duration. Due to this rounding:
                - The last chunk may overlap the end of the profiling duration,
                  effectively making the active chunk duration smaller than the
                  other chunks. This difference is accounted for in the in-use
                  percent calculation.
                - In extreme cases, the actual number of active chunks can be
                  smaller than <number>.
            Possible values are integers between 1 and 1000.
            Default is {NUM_CHUNKS}.

    Output: All time values default to nanoseconds
        Row# : Row number of chunk
        In-Use : Percentage of time the GPU is being used
        Duration : Duration of chunk
        Start : Start time of chunk
        PID : Process identifier
        Device ID : GPU device identifier

    This rule identifies time regions with low GPU utilization. For each
    process, each GPU device is examined, and a time range is created that
    starts with the beginning of the first GPU operation on that device and
    ends with the end of the last GPU operation on that device. This time range
    is then divided into equal chunks, and the GPU utilization is calculated
    for each chunk. The utilization includes all GPU operations as well as
    profiling overheads that the user cannot address.

    Note that the utilization refers to the "time" utilization and not the
    "resource" utilization. This script attempts to find time gaps when the GPU
    is or isn't being used, but does not take into account how many GPU
    resources are being used. Therefore, a single running memcpy is considered
    the same amount of "utilization" as a huge kernel that takes over all the
    cores. If multiple operations run concurrently in the same chunk, their
    utilization will be added up and may exceed 100%.

    Chunks with an in-use percentage less than the threshold value are displayed.
    If consecutive chunks have a low in-use percentage, the individual chunks are
    coalesced into a single display record, keeping the weighted average of
    percentages. This is why returned chunks may have different durations.
"""

    create_chunk_table = """
    CREATE TEMP TABLE CHUNK (
        rangeId   INTEGER PRIMARY KEY   NOT NULL
    )
"""

    insert_chunk_table = """
    INSERT INTO temp.CHUNK
    WITH RECURSIVE
        range AS (
            SELECT
                0 AS rangeId
            UNION ALL
            SELECT
                rangeId + 1 AS rangeId
            FROM
                range
            LIMIT {NUM_CHUNKS}
        )
    SELECT rangeId FROM range
"""

    query_format_columns = """
    SELECT
        ROW_NUMBER() OVER(ORDER BY average, duration) AS "Row#",
        average AS "In-Use:ratio_%",
        duration AS "Duration:dur_ns",
        start AS "Start:ts_ns",
        pid AS "PID",
        deviceId AS "Device ID",
        contextId AS "Context ID",
        globalId AS "_Global ID",
        api AS "_API"
    FROM
        ({GPU_UNION_TABLE})
    LIMIT {ROW_LIMIT}
"""

# Return chunks that have an average GPU utilization below the given threshold.
# 1. CTE "range": Define the range being analyzed for each deviceId/PID with
#    the corresponding chunk size.
# 2. CTE "chunk": Duplicate chunks for each deviceId/PID with the appropriate
#    start and end.
# 3. CTE "utilization": Find all ranges being run in each chunk and keep only
#    the ones that have a percentage of GPU utilization lower than the threshold.
#    If there are multiple streams, the utilizations are added up.
# 4. CTE "chunkgroup": Give a groupId that will be used to define consecutive
#    chunks.
# 5. Coalesce chunks with same groupId and calculate the weighted average.
    query_chunk = """
    WITH
        ops AS (
            {{GPU_TABLE}}
        ),
        range AS (
            SELECT
                min(start) AS start,
                max(end) AS end,
                ceil(CAST(max(end) - min(start) AS FLOAT) / {NUM_CHUNKS}) AS chunkSize,
                pid,
                globalId,
                deviceId,
                contextId,
                api
            FROM
                ops
            GROUP BY deviceId, pid
        ),
        chunk AS (
            SELECT
                chunk.rangeId,
                chunk.rangeId * range.chunkSize + range.start AS cstart,
                min(chunk.rangeId * range.chunkSize + range.start + range.chunkSize, range.end) AS cend,
                chunkSize,
                range.pid,
                range.globalId,
                range.deviceId,
                range.contextId,
                range.api
            FROM
                temp.CHUNK AS chunk
            JOIN
                range
                ON cstart < cend
        ),
        utilization AS (
            SELECT
                chunk.rangeId,
                chunk.cstart AS start,
                chunk.cend AS end,
                chunk.cend - chunk.cstart AS size,
                sum(CAST(coalesce(min(ops.end, chunk.cend) - max(ops.start, chunk.cstart), 0) AS FLOAT)) / (chunk.cend - chunk.cstart) * 100 AS timePercentage,
                chunk.pid,
                chunk.globalId,
                chunk.deviceId,
                chunk.contextId,
                chunk.api
            FROM
                chunk
            LEFT JOIN
                ops
                ON      ops.deviceId == chunk.deviceId
                    AND ops.pid == chunk.pid
                    AND ops.start < chunk.cend
                    AND ops.end > chunk.cstart
            GROUP BY
                chunk.rangeId, ops.deviceId, ops.pid
            HAVING
                timePercentage < {THRESHOLD}
        ),
        chunkgroup AS
        (
            SELECT
                *,
                rangeId - ROW_NUMBER() OVER (PARTITION BY pid, deviceId ORDER BY rangeId) AS groupId
            FROM
                utilization
        )
    SELECT
        min(start) AS start,
        max(end) - min(start) AS duration,
        round(sum(size * timePercentage) / sum(size), 1) AS average,
        pid,
        globalId,
        deviceId,
        contextId,
        api
    FROM
        chunkgroup
    GROUP BY groupId, deviceId, pid
    LIMIT {ROW_LIMIT}
"""

    def setup(self):
        err = super().setup()
        if err != None:
            return err

        parser = self.ArgumentParser(allow_abbrev=False)
        parser.add_optional_arg('rows', default=self.ROW_LIMIT, type=int, \
            help='maximum number of rows returned by the query')
        parser.add_optional_arg('start', type=int, \
            help='start time used for filtering')
        parser.add_optional_arg('end', type=int, \
            help='end time used for filtering')
        parser.add_optional_arg('nvtx', type=str, \
            help='NVTX range and domain used for filtering')
        parser.add_optional_arg('threshold', default=self.THRESHOLD, type=int, \
            help='maximum percentage of time the GPU is being used')
        parser.add_optional_arg('chunks', default=self.NUM_CHUNKS, type=int, \
            help='number of equal-duration chunks')
        args = parser.parse_optional_dashless_args(self.args)

        if args.chunks and not 1 <= args.chunks <= 1000:
            parser.error("argument --chunks: value must be between 1 and 1000")

        err = self.filter_time_range(args.start, args.end, args.nvtx)
        if err != None:
            return err

        self.statements = [
            self.create_chunk_table,
            self.insert_chunk_table.format(NUM_CHUNKS = args.chunks)]

        err = self.create_gpu_ops_view(self.query_chunk.format(
            NUM_CHUNKS = args.chunks,
            THRESHOLD = args.threshold,
            ROW_LIMIT = args.rows))
        if err != None:
            return err

        self.query = self.query_format_columns.format(
            GPU_UNION_TABLE = self.query_gpu_ops_union(),
            ROW_LIMIT = args.rows)

if __name__ == "__main__":
    GPULowUtilization.Main()
