Tools/Scripts/pgo-profile

#!/usr/bin/env python3

import argparse
import logging
import os
import subprocess

from webkitpy.llvm_profile_utils import (LLVMProfileData, merge_raw_profiles_in_directory,
                                         resolve_profiled_frameworks, weighted_profiles_for_framework)

logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(levelname)s: %(message)s', datefmt='%Y-%m-%d %H:%M:%S')
logger = logging.getLogger()


BENCHMARK_GROUP_WEIGHTS = [('speedometer3', 0.6), ('jetstream3', 0.2), ('motionmark', 0.2)]


def pad(string, max_length):
    if len(string) > max_length:
        return string[:max_length - 1] + u'…'
    return string.ljust(max_length)


def shortened(name):
    return name if len(name) < 200 else name[0:99] + u'…' + name[-98:]


def assert_directory(directory_path):
    assert os.path.isdir(directory_path), f'No directory at {directory_path}'
    return directory_path


def assert_file(file_path):
    assert os.path.isfile(file_path), f'No file at {file_path}'
    return file_path


def summarize_parser(subparsers, parent_parser):
    parser = subparsers.add_parser('summarize', parents=[parent_parser],
                                   help='Dumps function names in a given .profraw file, '
                                        'sorted in descending order by function count')
    parser.add_argument('file', help='Path to the .profraw file')
    parser.set_defaults(func=summarize)
    return parser


def summarize(args):
    file = assert_file(args.file)

    show_process = LLVMProfileData.show(file)
    show_process.check_returncode()

    lines = show_process.stdout.splitlines()

    counts_and_functions = []

    for line_number in range(len(lines)):
        line = lines[line_number].strip()
        if line.startswith('Function count: '):
            count = int(line.split()[-1])
            symbol = lines[line_number - 3].strip()[:-1]
            counts_and_functions.append((count, symbol))

    counts_and_functions.sort(reverse=True)
    for count, name in counts_and_functions:
        print(pad(str(count), 15), shortened(name))


def merge_parser(subparsers, parent_parser):
    parser = subparsers.add_parser('merge', parents=[parent_parser],
                                   help='Merge a pile of *.profraw files into the *.profdata files we can build with.')
    parser.add_argument('directory', help='Path to the directory containing the *.profraw files')
    parser.set_defaults(func=merge)
    return parser


def merge(args):
    directory = assert_directory(args.directory)
    frameworks = resolve_profiled_frameworks([directory], suffix='.profraw')
    merge_raw_profiles_in_directory(frameworks, directory)


def combine_parser(subparsers, parent_parser):
    parser = subparsers.add_parser('combine', parents=[parent_parser],
                                   help='Combine directories containing *.profdata files from different platforms together.')
    for i in range(0, len(BENCHMARK_GROUP_WEIGHTS)):
        group, _ = BENCHMARK_GROUP_WEIGHTS[i]
        parser.add_argument(f'--{group}', default=None,
                            help=f'Path to the directory containing the *.profdata files from a {group} run.')
    parser.add_argument('--output', help='Path to the directory where the output will be placed.')
    parser.set_defaults(func=combine)
    return parser


def combine(args):
    assert args.output, 'Must specify output directory.'

    out = assert_directory(args.output)
    args = vars(args)
    directory_weight_pairs = []

    for i in range(0, len(BENCHMARK_GROUP_WEIGHTS)):
        group, weight = BENCHMARK_GROUP_WEIGHTS[i]
        if args[group]:
            profile_group = assert_directory(args[group])
            directory_weight_pairs.append((profile_group, weight))

    assert len(directory_weight_pairs) > 0, 'Must specify at least one group.'

    frameworks = resolve_profiled_frameworks([directory for directory, _ in directory_weight_pairs])

    for framework in frameworks:
        weighted_profiles = weighted_profiles_for_framework(framework, directory_weight_pairs)

        logger.info(f'Merging {framework.name} with simplified weights: {weighted_profiles}')
        output_file = os.path.join(out, f'{framework.name}.profdata')
        merge_process = LLVMProfileData.merge(output_file, weighted_profiles=weighted_profiles)
        logger.info(f'stdout: {merge_process.stdout}')
        logger.info(f'stderr: {merge_process.stderr}')
        merge_process.check_returncode()
        logger.info(f'{framework.name} is successfully merged')


def compress_parser(subparsers, parent_parser):
    parser = subparsers.add_parser('compress', parents=[parent_parser],
                                   help='Compress *.profdata files so that they can be checked in.')
    parser.add_argument('--input', help='Path to the directory containing the input *.profdata files.')
    parser.add_argument('--output', help='Path to the directory where the output will be placed.')
    parser.set_defaults(func=compress)
    return parser


def decompress_parser(subparsers, parent_parser):
    parser = subparsers.add_parser('decompress', parents=[parent_parser],
                                   help='Decompress checked-in *.profdata.compressed files back into '
                                        'the *.profdata files we can build with.')
    parser.add_argument('--input', help='Path to the directory containing the input *.profdata.compressed files.')
    parser.add_argument('--output', help='Path to the directory where the output will be placed.')
    parser.set_defaults(func=decompress)
    return parser


def compress(args):
    out = assert_directory(args.output)
    input_directory = assert_directory(args.input)
    frameworks = resolve_profiled_frameworks([input_directory])

    for framework in frameworks:
        logger.info(f'Compressing {framework.name}')
        input_file = os.path.join(input_directory, f'{framework.name}.profdata')
        output_file = os.path.join(out, f'{framework.name}.profdata.compressed')

        LLVMProfileData.compress(input_file, output_file)
        logger.info(f'{framework.name} is successfully compressed')


def decompress(args):
    out = assert_directory(args.output)
    input_directory = assert_directory(args.input)
    frameworks = resolve_profiled_frameworks([input_directory], suffix='.profdata.compressed')

    for framework in frameworks:
        logger.info(f'Decompressing {framework.name}')
        input_file = os.path.join(input_directory, f'{framework.name}.profdata.compressed')
        output_file = os.path.join(out, f'{framework.name}.profdata')

        LLVMProfileData.decompress(input_file, output_file)
        logger.info(f'{framework.name} is successfully decompressed')


if __name__ == '__main__':
    parser = argparse.ArgumentParser(prog='pgo-profile')
    verbose_parser = argparse.ArgumentParser(add_help=False)
    verbose_parser.add_argument('-v', '--verbose', action='store_true', default=False, help='Turn on debug logging')
    subparsers = parser.add_subparsers(help='valid sub-commands', required=True, dest='sub command')
    merge_parser(subparsers, verbose_parser)
    summarize_parser(subparsers, verbose_parser)
    combine_parser(subparsers, verbose_parser)
    compress_parser(subparsers, verbose_parser)
    decompress_parser(subparsers, verbose_parser)

    args = parser.parse_args()
    if args.verbose:
        logger.setLevel(logging.DEBUG)
    try:
        args.func(args)
    except subprocess.CalledProcessError as e:
        if e.stdout:
            logger.error(e.stdout)
        if e.stderr:
            logger.error(e.stderr)
        raise e