llvm/tools/opt-viewer/optpmap.py

   1 import sys
   2 import multiprocessing
   3
   4
   5 _current = None
   6 _total = None
   7
   8
   9 def _init(current, total):
  10     global _current
  11     global _total
  12     _current = current
  13     _total = total
  14
  15
  16 def _wrapped_func(func_and_args):
  17     func, argument, should_print_progress, filter_ = func_and_args
  18
  19     if should_print_progress:
  20         with _current.get_lock():
  21             _current.value += 1
  22         sys.stdout.write("\r\t{} of {}".format(_current.value, _total.value))
  23         sys.stdout.flush()
  24
  25     return func(argument, filter_)
  26
  27
  28 def pmap(
  29     func, iterable, processes, should_print_progress, filter_=None, *args, **kwargs
  30 ):
  31     """
  32     A parallel map function that reports on its progress.
  33
  34     Applies `func` to every item of `iterable` and return a list of the
  35     results. If `processes` is greater than one, a process pool is used to run
  36     the functions in parallel. `should_print_progress` is a boolean value that
  37     indicates whether a string 'N of M' should be printed to indicate how many
  38     of the functions have finished being run.
  39     """
  40     global _current
  41     global _total
  42     _current = multiprocessing.Value("i", 0)
  43     _total = multiprocessing.Value("i", len(iterable))
  44
  45     func_and_args = [(func, arg, should_print_progress, filter_) for arg in iterable]
  46     if processes == 1:
  47         result = list(map(_wrapped_func, func_and_args, *args, **kwargs))
  48     else:
  49         pool = multiprocessing.Pool(
  50             initializer=_init,
  51             initargs=(
  52                 _current,
  53                 _total,
  54             ),
  55             processes=processes,
  56         )
  57         result = pool.map(_wrapped_func, func_and_args, *args, **kwargs)
  58         pool.close()
  59         pool.join()
  60
  61     if should_print_progress:
  62         sys.stdout.write("\r")
  63     return result