vendor/llvm/llvm-trunk-r308421

1 files changed, 12 insertions, 2 deletions

diff --git a/tools/opt-viewer/opt-viewer.py b/tools/opt-viewer/opt-viewer.py
index e6dd6a0286fe..69bcaedb7669 100755
--- a/tools/opt-viewer/opt-viewer.py
+++ b/tools/opt-viewer/opt-viewer.py
@@ -219,7 +219,11 @@ def generate_report(all_remarks,
 
 if __name__ == '__main__':
     parser = argparse.ArgumentParser(description=desc)
-    parser.add_argument('yaml_files', nargs='+')
+    parser.add_argument(
+        'yaml_dirs_or_files',
+        nargs='+',
+        help='List of optimization record files or directories searched '
+             'for optimization record files.')
     parser.add_argument(
         '--output-dir',
         '-o',
@@ -248,8 +252,14 @@ if __name__ == '__main__':
     args = parser.parse_args()
 
     print_progress = not args.no_progress_indicator
+
+    files = optrecord.find_opt_files(args.yaml_dirs_or_files)
+    if not files:
+        parser.error("No *.opt.yaml files found")
+        sys.exit(1)
+
     all_remarks, file_remarks, should_display_hotness = \
-        optrecord.gather_results(args.yaml_files, args.jobs, print_progress)
+        optrecord.gather_results(files, args.jobs, print_progress)
 
     map_remarks(all_remarks)