Support batched bb traces in TraceBlackboxEvaluator

This patch makes TraceBlackboxEvaluator support batched bb traces.
Essentially this means we can split all of the bb traces obtained from a
single memtrace into a bunch of chunks. We can then evaluate only one of
these chunks per perturbation. This adds some randomness to the training
which can be good and also massively speeds up evaluation.

Reviewers: svkeerthy, mtrofin

Reviewed By: mtrofin

Pull Request: #510

Author: boomanaiden154

compiler_opt/es/blackbox_evaluator.py

@@ -15,9 +15,12 @@
 
 import abc
 import concurrent.futures
+import os
+import random
 
 from absl import logging
 import gin
+import tensorflow as tf
 
 from compiler_opt.distributed.worker import FixedWorkerPool
 from compiler_opt.rl import corpus
@@ -114,20 +117,33 @@ def __init__(self, train_corpus: corpus.Corpus,
                bb_trace_path: str, function_index_path: str):
     self._train_corpus = train_corpus
     self._estimator_type = estimator_type
-    self._bb_trace_path = bb_trace_path
+    self._bb_trace_paths = []
+    if tf.io.gfile.isdir(bb_trace_path):
+      self._bb_trace_paths.extend([
+          os.path.join(bb_trace_path, bb_trace)
+          for bb_trace in tf.io.gfile.listdir(bb_trace_path)
+      ])
+    else:
+      self._bb_trace_paths.append(bb_trace_path)
     self._function_index_path = function_index_path
 
-    self._baseline: float | None = None
+    self._baselines: list[float] | None = None
 
   def get_results(
       self, pool: FixedWorkerPool,
       perturbations: list[bytes]) -> list[concurrent.futures.Future]:
-    job_args = [{
-        'modules': self._train_corpus.module_specs,
-        'function_index_path': self._function_index_path,
-        'bb_trace_path': self._bb_trace_path,
-        'policy_as_bytes': perturbation,
-    } for perturbation in perturbations]
+    job_args = []
+    self._current_baselines = []
+    for perturbation in perturbations:
+      bb_trace_path_index = random.randrange(len(self._bb_trace_paths))
+      bb_trace_path = self._bb_trace_paths[bb_trace_path_index]
+      self._current_baselines.append(self._baselines[bb_trace_path_index])
+      job_args.append({
+          'modules': self._train_corpus.module_specs,
+          'function_index_path': self._function_index_path,
+          'bb_trace_path': bb_trace_path,
+          'policy_as_bytes': perturbation
+      })
 
     _, futures = buffered_scheduler.schedule_on_worker_pool(
         action=lambda w, args: w.compile_corpus_and_evaluate(**args),
@@ -138,15 +154,15 @@ def get_results(
     return futures
 
   def set_baseline(self, pool: FixedWorkerPool) -> None:
-    if self._baseline is not None:
+    if self._baselines is not None:
       raise RuntimeError('The baseline has already been set.')
 
     job_args = [{
         'modules': self._train_corpus.module_specs,
         'function_index_path': self._function_index_path,
-        'bb_trace_path': self._bb_trace_path,
+        'bb_trace_path': bb_trace_path,
         'policy_as_bytes': None,
-    }]
+    } for bb_trace_path in self._bb_trace_paths]
 
     _, futures = buffered_scheduler.schedule_on_worker_pool(
         action=lambda w, args: w.compile_corpus_and_evaluate(**args),
@@ -155,21 +171,22 @@ def set_baseline(self, pool: FixedWorkerPool) -> None:
 
     concurrent.futures.wait(
         futures, return_when=concurrent.futures.ALL_COMPLETED)
-    if len(futures) != 1:
-      raise ValueError('Expected to have one result for setting the baseline,'
-                       f' got {len(futures)}')
+    if len(futures) != len(self._bb_trace_paths):
+      raise ValueError(
+          f'Expected to have {len(self._bb_trace_paths)} results for setting,'
+          f'the baseline, got {len(futures)}.')
 
-    self._baseline = futures[0].result()
+    self._baselines = [future.result() for future in futures]
 
   def get_rewards(
       self, results: list[concurrent.futures.Future]) -> list[float | None]:
     rewards = []
 
-    for result in results:
+    for result, baseline in zip(results, self._current_baselines):
       if result.exception() is not None:
         raise result.exception()
 
       rewards.append(
-          compilation_runner.calculate_reward(result.result(), self._baseline))
+          compilation_runner.calculate_reward(result.result(), baseline))
 
     return rewards

compiler_opt/es/blackbox_evaluator_test.py

@@ -67,6 +67,9 @@ def test_trace_get_results(self):
       evaluator = blackbox_evaluator.TraceBlackboxEvaluator(
           test_corpus, blackbox_optimizers.EstimatorType.FORWARD_FD,
           'fake_bb_trace_path', 'fake_function_index_path')
+      # pylint: disable=protected-access
+      evaluator._baselines = [1]
+      # pylint: enable=protected-access
       results = evaluator.get_results(pool, perturbations)
       self.assertSequenceAlmostEqual([result.result() for result in results],
                                      [1.0, 1.0, 1.0])
@@ -85,7 +88,9 @@ def test_trace_set_baseline(self):
           'fake_bb_trace_path', 'fake_function_index_path')
       evaluator.set_baseline(pool)
       # pylint: disable=protected-access
-      self.assertAlmostEqual(evaluator._baseline, 10)
+      self.assertLen(evaluator._baselines, 1)
+      self.assertAlmostEqual(evaluator._baselines[0], 10)
+      # pylint: enable=protected-access
 
   def test_trace_get_rewards(self):
     f1 = concurrent.futures.Future()
@@ -101,10 +106,56 @@ def test_trace_get_rewards(self):
         'fake_bb_trace_path', 'fake_function_index_path')
 
     # pylint: disable=protected-access
-    evaluator._baseline = 2
+    evaluator._current_baselines = [2, 3]
+    # pylint: enable=protected-access
     rewards = evaluator.get_rewards(results)
 
     # Only check for two decimal places as the reward calculation uses a
     # reasonably large delta (0.01) when calculating the difference to
     # prevent division by zero.
-    self.assertSequenceAlmostEqual(rewards, [0, -0.5], 2)
+    self.assertSequenceAlmostEqual(rewards, [0, 0], 2)
+
+  def test_trace_multiple_get_results(self):
+    with local_worker_manager.LocalWorkerPoolManager(
+        blackbox_test_utils.ESTraceWorker,
+        count=3,
+        worker_args=(),
+        worker_kwargs={}) as pool:
+      perturbations = [b'00', b'01', b'10']
+      test_corpus = corpus.create_corpus_for_testing(
+          location=self.create_tempdir().full_path,
+          elements=[corpus.ModuleSpec(name='name1', size=1)])
+      bb_trace_dir = self.create_tempdir()
+      bb_trace_dir.create_file('bb_trace1.pb')
+      bb_trace_dir.create_file('bb_trace2.pb')
+      evaluator = blackbox_evaluator.TraceBlackboxEvaluator(
+          test_corpus, blackbox_optimizers.EstimatorType.FORWARD_FD,
+          bb_trace_dir.full_path, 'fake_function_index_path')
+      # pylint: disable=protected-access
+      evaluator._baselines = [1, 2]
+      # pylint: enable=protected-access
+      results = evaluator.get_results(pool, perturbations)
+      self.assertSequenceAlmostEqual([result.result() for result in results],
+                                     [1.0, 1.0, 1.0])
+
+  def test_trace_multiple_set_baseline(self):
+    with local_worker_manager.LocalWorkerPoolManager(
+        blackbox_test_utils.ESTraceWorker,
+        count=1,
+        worker_args=(),
+        worker_kwargs={}) as pool:
+      test_corpus = corpus.create_corpus_for_testing(
+          location=self.create_tempdir().full_path,
+          elements=[corpus.ModuleSpec(name='name1', size=1)])
+      bb_trace_dir = self.create_tempdir()
+      bb_trace_dir.create_file('bb_trace1.pb')
+      bb_trace_dir.create_file('bb_trace2.pb')
+      evaluator = blackbox_evaluator.TraceBlackboxEvaluator(
+          test_corpus, blackbox_optimizers.EstimatorType.FORWARD_FD,
+          bb_trace_dir.full_path, 'fake_function_index_path')
+      evaluator.set_baseline(pool)
+      # pylint: disable=protected-access
+      self.assertLen(evaluator._baselines, 2)
+      self.assertAlmostEqual(evaluator._baselines[0], 10)
+      self.assertAlmostEqual(evaluator._baselines[1], 10)
+      # pylint: enable=protected-access