1 files changed, 139 insertions, 0 deletions
diff --git a/collect_profile.py b/collect_profile.py
new file mode 100755
index 00000000..27d5e5d6
--- /dev/null
+++ b/collect_profile.py
@@ -0,0 +1,139 @@
+#!/usr/bin/python
+
+# Script to generate and collect PGO data based on benchmark
+from __future__ import print_function
+
+import argparse
+import config
+import logging
+import os
+import subprocess
+import sys
+import tempfile
+
+# Turn the logging level to INFO before importing other code, to avoid having
+# failed import logging messages confuse the user.
+logging.basicConfig(level=logging.INFO)
+
+def _parse_arguments_internal(argv):
+    """
+    Parse command line arguments
+
+    @param argv: argument list to parse
+
+    @returns:    tuple of parsed arguments and argv suitable for remote runs
+
+    @raises SystemExit if arguments are malformed, or required arguments
+            are not present.
+    """
+
+    parser = argparse.ArgumentParser(description='Run this script to collect '
+                                                 'PGO data.')
+
+    parser.add_argument('-b', '--bench',
+                        help='Select which benchmark to collect profdata.')
+
+    parser.add_argument('-d', '--pathDUT', default='/data/local/tmp',
+                        help='Specify where to generate PGO data on device, '
+                             'set to /data/local/tmp by default.')
+
+    parser.add_argument('-p', '--path', default=config.bench_suite_dir,
+                        help='Specify the location to put the profdata, set '
+                             ' to bench_suite_dir by default.')
+
+    parser.add_argument('-s', '--serial',
+                        help='Device serial number.')
+
+    parser.add_argument('-r', '--remote', default='localhost',
+                        help='hostname[:port] if the ADB device is connected '
+                             'to a remote machine. Ensure this workstation '
+                             'is configured for passwordless ssh access as '
+                             'users "root" or "adb"')
+    return parser.parse_args(argv)
+
+# Call run.py to build benchmark with -fprofile-generate flags and run on DUT
+def run_suite(bench, serial, remote, pathDUT):
+    logging.info('Build and run instrumented benchmark...')
+    run_cmd = ['./run.py', '-b=' + bench]
+    if serial:
+        run_cmd.append('-s=' + serial)
+    run_cmd.append('-r=' + remote)
+    run_cmd.append('-f=-fprofile-generate=%s' % pathDUT)
+    run_cmd.append('--ldflags=-fprofile-generate=%s' % pathDUT)
+    try:
+        subprocess.check_call(run_cmd)
+    except subprocess.CalledProcessError:
+        logging.error('Error running %s.', run_cmd)
+        raise
+
+# Pull profraw data from device using pull_device.py script in autotest utils.
+def pull_result(bench, serial, remote, pathDUT, path):
+    logging.info('Pulling profraw data from device to local')
+    pull_cmd = [os.path.join(config.android_home,
+                             config.autotest_dir,
+                             'site_utils/pull_device.py')]
+    pull_cmd.append('-b=' + bench)
+    pull_cmd.append('-r=' + remote)
+    if serial:
+        pull_cmd.append('-s=' + serial)
+    pull_cmd.append('-p=' + path)
+    pull_cmd.append('-d=' + pathDUT)
+    try:
+        subprocess.check_call(pull_cmd)
+    except:
+        logging.error('Error while pulling profraw data.')
+        raise
+
+# Use llvm-profdata tool to convert profraw data to the format llvm can
+# recgonize.
+def merge(bench, pathDUT, path):
+    logging.info('Generate profdata for PGO...')
+    # Untar the compressed rawdata file collected from device
+    tmp_dir = tempfile.mkdtemp()
+    untar_cmd = ['tar',
+                 '-xf',
+                 os.path.join(path, bench + '_profraw.tar'),
+                 '-C',
+                 tmp_dir]
+
+    # call llvm-profdata to merge the profraw data
+    profdata = os.path.join(path, bench + '.profdata')
+    merge_cmd = ['llvm-profdata',
+                 'merge',
+                 '-output=' + profdata,
+                 tmp_dir + pathDUT]
+    try:
+        subprocess.check_call(untar_cmd)
+        subprocess.check_call(merge_cmd)
+        logging.info('Profdata is generated successfully, located at %s',
+                     profdata)
+    except:
+        logging.error('Error while merging profraw data.')
+        raise
+    finally:
+        subprocess.check_call(['rm', '-rf', tmp_dir])
+
+def main(argv):
+    """
+    Entry point for nightly_run script.
+
+    @param argv: arguments list
+    """
+    arguments = _parse_arguments_internal(argv)
+
+    bench = arguments.bench
+    serial = arguments.serial
+    path = arguments.path
+    remote = arguments.remote
+
+    # Create a profraw directory to collect data
+    pathDUT = os.path.join(arguments.pathDUT, bench + '_profraw')
+
+    run_suite(bench, serial, remote, pathDUT)
+
+    pull_result(bench, serial, remote, pathDUT, path)
+
+    merge(bench, pathDUT, path)
+
+if __name__ == '__main__':
+    main(sys.argv[1:])