Close performance gap between lookup of freshly built and deserialized algorithms #12161

"""Measure cuda.compute merge-sort workspace-query host overhead.

Compare a factory-built algorithm with one restored from its serialized blob.
Only algorithm(temp_storage=None, ...) is repeated in the measured loops.
No workspace allocation, sorting execution, or stream synchronization occurs
inside those loops. Each sample is the average time per query in one batch.

Examples (run outside the CuPy checkout, using the installed wheel):
  python profile_merge_sort_query.py --trials 101 --output query-results.json
  python -m profiling.sampling run --mode cpu --native --flamegraph -r 10khz \\
      -o query-built.html profile_merge_sort_query.py --case built --duration 60
  python -m profiling.sampling run --mode cpu --native --flamegraph -r 10khz \\
      -o query-deserialized.html profile_merge_sort_query.py \\
      --case deserialized --duration 60

Uses the installed cuda.compute and CuPy, with no source overlay. The comparator
is the built-in OpKind.LESS; this isolates query overhead without custom-op JIT.
"""

import argparse
import json
import math
from pathlib import Path
import platform
import sys
import time


def positive_int(text):
    value = int(text)
    if value <= 0:
        raise argparse.ArgumentTypeError('must be positive')
    return value


# Keep separate named loops to distinguish cases in sampling flamegraphs.
def query_built_batch(algorithm, arguments, count):
    for _ in range(count):
        size = algorithm(**arguments)
    return size


def query_deserialized_batch(algorithm, arguments, count):
    for _ in range(count):
        size = algorithm(**arguments)
    return size


def main():
    parser = argparse.ArgumentParser(
        description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter)
    parser.add_argument('--size', type=positive_int, default=2048)
    parser.add_argument('--dtype', choices=('int32', 'int64', 'float32', 'float64'),
                        default='float32')
    parser.add_argument('--device', type=int, default=0)
    parser.add_argument('--case', choices=('both', 'built', 'deserialized'),
                        default='both')
    parser.add_argument('--trials', type=positive_int, default=101,
                        help='trials per case, unless --duration is supplied')
    parser.add_argument('--batch-size', type=positive_int, default=1000,
                        help='queries per measured trial')
    parser.add_argument('--warmup', type=positive_int, default=20)
    parser.add_argument('--duration', type=float,
                        help='run for this many seconds instead of --trials; '
                             'complete each built/deserialized pair')
    parser.add_argument('--output', type=Path,
                        help='save metadata and raw trials as JSON (new file)')
    args = parser.parse_args()
    if args.duration is not None and (
            not math.isfinite(args.duration) or args.duration <= 0):
        parser.error('--duration must be finite and positive')
    if args.output is not None and args.output.exists():
        parser.error(f'output already exists: {args.output}')

    import numpy as np
    import cupy as cp
    import cuda.cccl
    import cuda.compute as cc

    if cp.cuda.runtime.is_hip:
        parser.error('requires a CUDA build of CuPy')
    metadata = dict(python=platform.python_version(), executable=sys.executable,
                    cupy_version=cp.__version__, cupy_path=cp.__file__,
                    cuda_cccl_version=cuda.cccl.__version__,
                    compute_path=cc.__file__, options={
                        **vars(args), 'output': str(args.output)
                        if args.output is not None else None},
                    comparator='OpKind.LESS')
    print(json.dumps(metadata, indent=2), flush=True)

    with cp.cuda.Device(args.device):
        d_in = cp.ones(args.size, dtype=args.dtype)
        d_out = cp.full_like(d_in, -7)
        stream = cp.cuda.get_current_stream()
        stream.synchronize()
        arguments = dict(temp_storage=None, d_in_keys=d_in, d_out_keys=d_out,
                         d_in_values=None, d_out_values=None,
                         num_items=args.size, op=cc.OpKind.LESS, stream=stream)
        phases = {}
        start = time.perf_counter_ns()
        built = cc.make_merge_sort(d_in_keys=d_in, d_out_keys=d_out,
                                   d_in_values=None, d_out_values=None,
                                   op=cc.OpKind.LESS)
        phases['factory_ms'] = (time.perf_counter_ns() - start) / 1e6
        algorithms = {'built': built}
        if args.case != 'built':
            start = time.perf_counter_ns()
            blob = built.serialize()
            phases['serialize_ms'] = (time.perf_counter_ns() - start) / 1e6
            start = time.perf_counter_ns()
            algorithms['deserialized'] = cc.deserialize(blob)
            phases['deserialize_ms'] = (time.perf_counter_ns() - start) / 1e6
        cases = (['built', 'deserialized'] if args.case == 'both'
                 else [args.case])
        expected_size = None
        for case in cases:
            start = time.perf_counter_ns()
            size = algorithms[case](**arguments)
            phases[case + '_first_query_us'] = (
                time.perf_counter_ns() - start) / 1e3
            if expected_size is None:
                expected_size = size
            if size != expected_size:
                raise RuntimeError('workspace queries disagree')
            for _ in range(args.warmup):
                if algorithms[case](**arguments) != expected_size:
                    raise RuntimeError('workspace size changed during warmup')
        print('Setup / first queries:', json.dumps(phases), flush=True)
        print(f'Workspace: {expected_size} bytes; warming complete.', flush=True)
        loops = {'built': query_built_batch,
                 'deserialized': query_deserialized_batch}
        trials = []
        started = time.perf_counter()
        round_index = 0
        while (round_index < args.trials if args.duration is None else
               round_index == 0 or time.perf_counter() - started < args.duration):
            # Reverse order on alternating rounds to reduce ordering bias.
            order = cases if round_index % 2 == 0 else cases[::-1]
            for case in order:
                start = time.perf_counter_ns()
                size = loops[case](algorithms[case], arguments, args.batch_size)
                elapsed_ns = time.perf_counter_ns() - start
                if size != expected_size:
                    raise RuntimeError('workspace size changed during timing')
                trials.append(dict(case=case, round=round_index + 1,
                                   elapsed_ns=elapsed_ns,
                                   queries=args.batch_size,
                                   us_per_query=elapsed_ns / args.batch_size / 1e3))
            round_index += 1
        elapsed = time.perf_counter() - started
        # Check query-only behavior outside the measured region.
        if not bool(cp.all(d_in == 1)) or not bool(cp.all(d_out == -7)):
            raise RuntimeError('workspace query modified input/output arrays')

    percentiles = [5, 25, 50, 75, 95]
    summary = {}
    print('Microseconds/query: P5 P25 P50 P75 P95 (batch averages)')
    for case in cases:
        values = [t['us_per_query'] for t in trials if t['case'] == case]
        quantiles = np.percentile(values, percentiles).tolist()
        summary[case] = dict(trials=len(values), queries=len(values)*args.batch_size,
                             percentiles_us=dict(zip(map(str, percentiles), quantiles)))
        print(f'{case:>12}: ' + ' '.join(f'{v:.3f}' for v in quantiles)
              + f' ({len(values)} trials)', flush=True)
    if len(cases) == 2:
        paired_delta = []
        for i in range(0, len(trials), 2):
            pair = {t['case']: t['us_per_query'] for t in trials[i:i+2]}
            paired_delta.append(pair['deserialized'] - pair['built'])
        summary['paired_deserialized_minus_built_median_us'] = float(
            np.median(paired_delta))
        print('Median paired difference (deserialized - built): '
              f'{np.median(paired_delta):.3f} us/query')
    if args.output is not None:
        with args.output.open('x') as f:
            json.dump(dict(metadata=metadata, phases=phases,
                           workspace_bytes=expected_size, elapsed_s=elapsed,
                           summary=summary, trials=trials), f, indent=2)
        print(f'Saved {args.output}')


if __name__ == '__main__':
    main()