From 6c4188a52057c9fd6ec903b2ee50349aba9eeb55 Mon Sep 17 00:00:00 2001
From: Alessandro Comodi <acomodi@antmicro.com>
Date: Fri, 4 Mar 2022 11:05:55 +0100
Subject: [PATCH] utils: parallelize all roi parts generation

Signed-off-by: Alessandro Comodi <acomodi@antmicro.com>
---
 utils/roi_all.py | 28 +++++++++++++++++++++++-----
 1 file changed, 23 insertions(+), 5 deletions(-)

diff --git a/utils/roi_all.py b/utils/roi_all.py
index 660f9521..d3d3cea4 100755
--- a/utils/roi_all.py
+++ b/utils/roi_all.py
@@ -9,13 +9,26 @@
 #
 # SPDX-License-Identifier: ISC
 import argparse
-import yaml
-import subprocess
+import multiprocessing as mp
 import os
 import re
+import subprocess
+import yaml
+
 from prjxray import util
 
 
+def worker(arglist):
+    part, cwd = arglist
+    cmd = "make roi_only"
+
+    env = os.environ.copy()
+    env['XRAY_PART'] = part
+
+    print("running subprocess")
+    subprocess.run(cmd.split(' '), check=True, env=env, cwd=cwd)
+
+
 def main():
     """Rois all parts for a family by calling "make roi_only" over all parts
     with the same device as XRAY_PART.
@@ -36,12 +49,17 @@ def main():
         if device['fabric'] == information['device']:
             valid_devices.append(name)
 
+    tasks = []
+
     for part, data in util.get_parts(db_root).items():
         if data['device'] in valid_devices:
-            command = "make roi_only"
-            env['XRAY_PART'] = part
             cwd = os.getenv('XRAY_FUZZERS_DIR')
-            subprocess.run(command.split(' '), check=True, env=env, cwd=cwd)
+
+            tasks.append((part, cwd))
+
+    with mp.Pool() as pool:
+        for _ in pool.imap_unordered(worker, tasks):
+            pass
 
 
 if __name__ == '__main__':