1 files changed, 45 insertions, 21 deletions

diff --git a/llvm/utils/opt-viewer/opt-viewer.py b/llvm/utils/opt-viewer/opt-viewer.py
index 394a472b15e..5e5daf7feb0 100755
--- a/llvm/utils/opt-viewer/opt-viewer.py
+++ b/llvm/utils/opt-viewer/opt-viewer.py
@@ -2,24 +2,28 @@
 
 from __future__ import print_function
 
-desc = '''Generate HTML output to visualize optimization records from the YAML files
-generated with -fsave-optimization-record and -fdiagnostics-show-hotness.
-
-The tools requires PyYAML and Pygments Python packages.'''
-
-import optrecord
-import functools
-from multiprocessing import Pool
-from multiprocessing import Lock, cpu_count
-import errno
 import argparse
+import cgi
+import errno
+import functools
+from multiprocessing import cpu_count
 import os.path
 import re
 import shutil
+
 from pygments import highlight
 from pygments.lexers.c_cpp import CppLexer
 from pygments.formatters import HtmlFormatter
-import cgi
+
+import optpmap
+import optrecord
+
+
+desc = '''Generate HTML output to visualize optimization records from the YAML files
+generated with -fsave-optimization-record and -fdiagnostics-show-hotness.
+
+The tools requires PyYAML and Pygments Python packages.'''
+
 
 # This allows passing the global context to the child processes.
 class Context:
@@ -177,7 +181,13 @@ def map_remarks(all_remarks):
                     context.caller_loc[caller] = arg['DebugLoc']
 
 
-def generate_report(pmap, all_remarks, file_remarks, source_dir, output_dir, should_display_hotness):
+def generate_report(all_remarks,
+                    file_remarks,
+                    source_dir,
+                    output_dir,
+                    should_display_hotness,
+                    num_jobs,
+                    should_print_progress):
     try:
         os.makedirs(output_dir)
     except OSError as e:
@@ -187,7 +197,12 @@ def generate_report(pmap, all_remarks, file_remarks, source_dir, output_dir, sho
             raise
 
     _render_file_bound = functools.partial(_render_file, source_dir, output_dir, context)
-    pmap(_render_file_bound, file_remarks.items())
+    if should_print_progress:
+        print('Rendering HTML files...')
+    optpmap.pmap(_render_file_bound,
+                 file_remarks.items(),
+                 num_jobs,
+                 should_print_progress)
 
     if should_display_hotness:
         sorted_remarks = sorted(optrecord.itervalues(all_remarks), key=lambda r: (r.Hotness, r.File, r.Line, r.Column, r.PassWithDiffPrefix, r.yaml_tag, r.Function), reverse=True)
@@ -220,16 +235,25 @@ if __name__ == '__main__':
         '-s',
         default='',
         help='set source directory')
+    parser.add_argument(
+        '--no-progress-indicator',
+        '-n',
+        action='store_true',
+        default=False,
+        help='Do not display any indicator of how many YAML files were read '
+             'or rendered into HTML.')
     args = parser.parse_args()
 
-    if args.jobs == 1:
-        pmap = map
-    else:
-        pool = Pool(processes=args.jobs)
-        pmap = pool.map
-
-    all_remarks, file_remarks, should_display_hotness = optrecord.gather_results(pmap, args.yaml_files)
+    print_progress = not args.no_progress_indicator
+    all_remarks, file_remarks, should_display_hotness = \
+        optrecord.gather_results(args.yaml_files, args.jobs, print_progress)
 
     map_remarks(all_remarks)
 
-    generate_report(pmap, all_remarks, file_remarks, args.source_dir, args.output_dir, should_display_hotness)
+    generate_report(all_remarks,
+                    file_remarks,
+                    args.source_dir,
+                    args.output_dir,
+                    should_display_hotness,
+                    args.jobs,
+                    print_progress)