ruff rules for comprehensions and performance (#420)

* ruff rules for comprehensions and performance

* uv tool run yapf -irp .

---------

Co-authored-by: Aiden Grossman <[email protected]>

compiler_opt/benchmark/benchmark_chromium.py

@@ -206,9 +206,9 @@ def main(_):
     with open(test_description, encoding='UTF-8') as test_description_file:
       print(test_description)
       test_descriptions.append(json.load(test_description_file))
-  test_executables = []
-  for test_description in test_descriptions:
-    test_executables.append(test_description['executable'])
+  test_executables = [
+      test_description['executable'] for test_description in test_descriptions
+  ]
 
   if FLAGS.compile_llvm:
     benchmarking_utils.build_llvm(FLAGS.model_path, FLAGS.llvm_use_incremental,

compiler_opt/benchmark/benchmark_report_test.py

@@ -81,9 +81,8 @@ def test_loading(self):
                 'PerfCounter_1': [50],
             }
         })
-    self.assertSetEqual(report.names(), set(['BM_A', 'BM_B']))
-    self.assertSetEqual(report.counters(),
-                        set(['PerfCounter_0', 'PerfCounter_1']))
+    self.assertSetEqual(report.names(), {'BM_A', 'BM_B'})
+    self.assertSetEqual(report.counters(), {'PerfCounter_0', 'PerfCounter_1'})
     self.assertEqual(
         report.counter_means('BM_A', 'PerfCounter_0'),
         (10.488088481701517, 0.7071067811865476))

compiler_opt/benchmark/filter_tests.py

@@ -66,9 +66,7 @@ def main(_):
     test_suite_description = json.load(test_description_file)
     test_outputs = gtest_executable_utils.run_test_suite(
         test_suite_description, FLAGS.executable_path, [], FLAGS.num_threads)
-    test_list = []
-    for test_output in test_outputs:
-      test_list.append(test_output['name'])
+    test_list = [test_output['name'] for test_output in test_outputs]
     # copy the old test suite and just replace the tests array
     new_test_suite_description = test_suite_description
     new_test_suite_description['tests'] = test_list

compiler_opt/benchmark/gtest_executable_utils.py

@@ -125,9 +125,8 @@ def run_test_suite(test_suite_description: Dict[str, List[str]],
   if num_threads is None:
     num_threads = 1
 
-  test_descriptions = []
-  for test in test_suite_description['tests']:
-    test_descriptions.append((test_executable, test, perf_counters))
+  test_descriptions = [(test_executable, test, perf_counters)
+                       for test in test_suite_description['tests']]
 
   test_data_output = Parallel(n_jobs=num_threads)(
       delayed(run_and_parse)(test_description)

compiler_opt/distributed/local/local_worker_manager.py

@@ -184,7 +184,7 @@ def _msg_pump(self):
       # clear out pending futures and mark ourselves as "stopped" by null-ing
       # the map
       with self._lock:
-        for _, v in self._map.items():
+        for v in self._map.values():
           v.set_exception(concurrent.futures.CancelledError())
         self._map = None
 

compiler_opt/es/blackbox_learner.py

@@ -168,13 +168,12 @@ def __init__(self,
 
   def _get_perturbations(self) -> List[npt.NDArray[np.float32]]:
     """Get perturbations for the model weights."""
-    perturbations = []
     rng = np.random.default_rng(seed=self._seed)
-    for _ in range(self._config.total_num_perturbations):
-      perturbations.append(
-          rng.normal(size=len(self._model_weights)) *
-          self._config.precision_parameter)
-    return perturbations
+    return [
+        rng.normal(size=len(self._model_weights)) *
+        self._config.precision_parameter
+        for _ in range(self._config.total_num_perturbations)
+    ]
 
   def _update_model(self, perturbations: List[npt.NDArray[np.float32]],
                     rewards: List[float]) -> None:
@@ -276,10 +275,10 @@ def run_step(self, pool: FixedWorkerPool) -> None:
           p for p in initial_perturbations for p in (p, -p)
       ]
 
-    perturbations_as_policies = []
-    for perturbation in initial_perturbations:
-      perturbations_as_policies.append(
-          self._get_policy_from_perturbation(perturbation))
+    perturbations_as_policies = [
+        self._get_policy_from_perturbation(perturbation)
+        for perturbation in initial_perturbations
+    ]
 
     results = self._evaluator.get_results(pool, perturbations_as_policies)
     rewards = self._evaluator.get_rewards(results)

compiler_opt/es/regalloc_trace/regalloc_trace_worker.py

@@ -110,13 +110,12 @@ def _build_corpus(self, modules: Collection[corpus.ModuleSpec],
       else:
         tflite_policy_dir = None
 
-    compile_futures = []
     with concurrent.futures.ThreadPoolExecutor(
         max_workers=self._thread_count) as thread_pool:
-      for module in modules:
-        compile_futures.append(
-            thread_pool.submit(self._compile_module, module, output_directory,
-                               tflite_policy_dir))
+      compile_futures = [
+          thread_pool.submit(self._compile_module, module, output_directory,
+                             tflite_policy_dir) for module in modules
+      ]
 
       for future in compile_futures:
         if future.exception() is not None:

compiler_opt/rl/corpus.py

@@ -342,7 +342,7 @@ def __init__(self,
                         '{context.module_full_path}') + additional_flags
 
     # don't use add/remove for replace
-    add_keys = set(k.split('=', maxsplit=1)[0] for k in additional_flags)
+    add_keys = {k.split('=', maxsplit=1)[0] for k in additional_flags}
     if add_keys.intersection(
         set(replace_flags)) or set(delete_flags).intersection(
             set(replace_flags)) or add_keys.intersection(set(delete_flags)):

compiler_opt/rl/data_reader.py

@@ -44,11 +44,12 @@ def _parser_fn(serialized_proto):
     # and stored in the feature list.
     context_features = {}
     # pylint: disable=g-complex-comprehension
-    sequence_features = dict(
-        (tensor_spec.name,
-         tf.io.FixedLenSequenceFeature(
-             shape=tensor_spec.shape, dtype=tensor_spec.dtype))
-        for tensor_spec in agent_cfg.time_step_spec.observation.values())
+    sequence_features = {
+        tensor_spec.name:
+            tf.io.FixedLenSequenceFeature(
+                shape=tensor_spec.shape, dtype=tensor_spec.dtype)
+        for tensor_spec in agent_cfg.time_step_spec.observation.values()
+    }
     sequence_features[
         agent_cfg.action_spec.name] = tf.io.FixedLenSequenceFeature(
             shape=agent_cfg.action_spec.shape,

compiler_opt/rl/imitation_learning/generate_bc_trajectories_lib.py

@@ -561,10 +561,7 @@ def explore_at_state_generator(
       yield base_seq, base_policy
 
   def _build_replay_prefix_list(self, seq_ex):
-    ret_list = []
-    for int_list in seq_ex:
-      ret_list.append(int_list.int64_list.value[0])
-    return ret_list
+    return [int_list.int64_list.value[0] for int_list in seq_ex]
 
   def _create_timestep(self, curr_obs_dict: env.TimeStep):
     curr_obs = curr_obs_dict.obs

compiler_opt/rl/imitation_learning/generate_bc_trajectories_test.py

@@ -711,11 +711,11 @@ def test_gen_trajectories(self):
             'horizon': 10
         } for name in module_names]
         self.assertEqual(
-            set(frozenset(dict_i) for dict_i in max_profiles_dict_list),
-            set(frozenset(dict_i) for dict_i in comp_dict_list))
+            {frozenset(dict_i) for dict_i in max_profiles_dict_list},
+            {frozenset(dict_i) for dict_i in comp_dict_list})
         self.assertEqual(
-            set(frozenset(dict_i) for dict_i in pol_profiles_dict_list),
-            set(frozenset(dict_i) for dict_i in comp_dict_list))
+            {frozenset(dict_i) for dict_i in pol_profiles_dict_list},
+            {frozenset(dict_i) for dict_i in comp_dict_list})
 
 
 if __name__ == '__main__':

compiler_opt/rl/imitation_learning/weighted_bc_trainer_lib.py

@@ -54,7 +54,7 @@
 @gin.configurable
 class TrainingWeights:
   """Class for computing weights for training.
-  
+
   To use, create an instance by specifying the partitions used in
   collecting the data with generate_bc_trajectories. Next, run multiple steps
   of update_weights with collected profiles from generate_bc_trajectories,
@@ -79,7 +79,7 @@ def _bucket_by_feature(
       self, data: List[ProfilingDictValueType],
       feature_name: str) -> List[List[ProfilingDictValueType]]:
     """Partitions the profiles according to the feature name.
-    
+
     Partitions the profiles according to the feature name and the
     buckets defined by self._partitions.
 
@@ -100,11 +100,11 @@ def _bucket_by_feature(
 
   def _get_exp_gradient_step(self, loss, step_size) -> np.ndarray:
     """Exponentiated gradient step.
-    
+
     Args:
       loss: observed losses to update the weights
       step_size: step size for the update
-    
+
     Returns:
       probability distribution from the updated weights
     """
@@ -117,7 +117,7 @@ def create_new_profile(self,
                          data_eval: List[ProfilingDictValueType],
                          eps: float = 1e-5) -> List[ProfilingDictValueType]:
     """Create a new profile which contains the regret and relative reward.
-    
+
     The regret is measured as the difference between the loss of the data_eval
     profiles and of the data_comparator profiles. The reward is the negative
     regret normalized by the loss of the data_comparator profiles.
@@ -162,11 +162,11 @@ def update_weights(
       self, comparator_profile: List[ProfilingDictValueType],
       policy_profile: List[ProfilingDictValueType]) -> np.ndarray:
     """Constructs a new profile and uses the loss to update self._probs with EG.
-    
+
     Args:
       comparator_profile: baseline profiles to measure improvement against
       policy_profile: profiles to evaluate for improvement
-    
+
     Returns:
       Updated probabilities to use as weights in training.
     """
@@ -199,7 +199,7 @@ def get_weights(self) -> np.ndarray:
 @gin.configurable
 class ImitationLearningTrainer:
   """Implements one iteration of the BC-Max algorithm.
-  
+
   BC-Max can be found at https://arxiv.org/pdf/2403.19462."""
 
   def __init__(
@@ -234,11 +234,12 @@ def __init__(
     self._global_step = 0
 
     observation_spec, action_spec = config.get_inlining_signature_spec()
-    sequence_features = dict(
-        (tensor_spec.name,
-         tf.io.FixedLenSequenceFeature(
-             shape=tensor_spec.shape, dtype=tensor_spec.dtype))
-        for tensor_spec in observation_spec[-1].values())
+    sequence_features = {
+        tensor_spec.name:
+            tf.io.FixedLenSequenceFeature(
+                shape=tensor_spec.shape, dtype=tensor_spec.dtype)
+        for tensor_spec in observation_spec[-1].values()
+    }
     sequence_features.update({
         action_spec.name:
             tf.io.FixedLenSequenceFeature(
@@ -277,7 +278,7 @@ def _parse_func(self, raw_record, sequence_features):
 
   def _make_feature_label(self, parsed_example, num_processors):
     """Function to pre-process the parsed examples from dataset.
-    
+
     Removes certein features not used for training and reshapes
     features appropriately."""
     concat_arr = []
@@ -307,10 +308,10 @@ def _make_feature_label(self, parsed_example, num_processors):
 
   def load_dataset(self, filepaths: List[str]) -> tf.data.TFRecordDataset:
     """Load datasets from specified filepaths for training.
-    
+
     Args:
       filepaths: paths to dataset files
-      
+
     Returns:
       dataset: loaded tf dataset"""
     ignore_order = tf.data.Options()

compiler_opt/rl/inlining/config.py

@@ -25,8 +25,8 @@
 @gin.configurable()
 def get_inlining_signature_spec():
   """Returns (time_step_spec, action_spec) for LLVM inlining."""
-  observation_spec = dict(
-      (key, tf.TensorSpec(dtype=tf.int64, shape=(), name=key)) for key in (
+  observation_spec = {
+      key: tf.TensorSpec(dtype=tf.int64, shape=(), name=key) for key in (
           # Base features
           'caller_basic_block_count',
           'caller_conditionally_executed_blocks',
@@ -64,7 +64,8 @@ def get_inlining_signature_spec():
           'is_multiple_blocks',
           'nested_inlines',
           'nested_inline_cost_estimate',
-          'threshold'))
+          'threshold')
+  }
   reward_spec = tf.TensorSpec(dtype=tf.float32, shape=(), name='reward')
   time_step_spec = time_step.time_step_spec(observation_spec, reward_spec)
   action_spec = tensor_spec.BoundedTensorSpec(

compiler_opt/rl/inlining/imitation_learning_config.py

@@ -33,14 +33,15 @@ def get_inlining_signature_spec():
   """Returns (time_step_spec, action_spec) for collecting IL trajectories."""
   time_step_spec, _ = config.get_inlining_signature_spec()
   observation_spec = time_step_spec.observation
-  observation_spec.update(
-      dict((key, tf.TensorSpec(dtype=tf.int64, shape=(), name=key)) for key in (
+  observation_spec.update({
+      key: tf.TensorSpec(dtype=tf.int64, shape=(), name=key) for key in (
           'is_callee_avail_external',
           'is_caller_avail_external',
           # following features are not used in training.
           'inlining_default',
           SequenceExampleFeatureNames.label_name,
-          'policy_label')))  # testing only
+          'policy_label')
+  })  # testing only
 
   observation_spec[SequenceExampleFeatureNames.module_name] = tf.TensorSpec(
       dtype=tf.string, shape=(), name=SequenceExampleFeatureNames.module_name)
@@ -63,11 +64,12 @@ def get_input_signature():
   """Returns (time_step_spec, action_spec) wrapping a trained policy."""
   time_step_spec, action_spec = config.get_inlining_signature_spec()
   observation_spec = time_step_spec.observation
-  observation_spec.update(
-      dict((key, tf.TensorSpec(dtype=tf.int64, shape=(), name=key)) for key in (
+  observation_spec.update({
+      key: tf.TensorSpec(dtype=tf.int64, shape=(), name=key) for key in (
           'is_callee_avail_external',
           'is_caller_avail_external',
-      )))
+      )
+  })
 
   time_step_spec = time_step.time_step_spec(observation_spec,
                                             time_step_spec.reward)

compiler_opt/rl/log_reader.py

@@ -181,9 +181,7 @@ def _enumerate_log_from_stream(
       context = event['context']
       continue
     observation_id = int(event['observation'])
-    features = []
-    for ts in tensor_specs:
-      features.append(_read_tensor(f, ts))
+    features = [_read_tensor(f, ts) for ts in tensor_specs]
     f.readline()
     score = None
     if score_spec is not None:

compiler_opt/rl/policy_saver.py

@@ -180,8 +180,9 @@ def _write_output_signature(
         tf.nest.flatten(saved_model.signatures['action'].structured_outputs))
 
     # Map spec name to index in flattened outputs.
-    sm_action_indices = dict(
-        (k.name.lower(), i) for i, k in enumerate(sm_action_signature))
+    sm_action_indices = {
+        k.name.lower(): i for i, k in enumerate(sm_action_signature)
+    }
 
     # List mapping flattened structured outputs to tensors.
     sm_action_tensors = saved_model.signatures['action'].outputs

compiler_opt/rl/regalloc/config.py

@@ -31,26 +31,28 @@ def get_regalloc_signature_spec():
   """Returns (time_step_spec, action_spec) for LLVM register allocation."""
   num_registers = get_num_registers()
 
-  observation_spec = dict(
-      (key, tf.TensorSpec(dtype=tf.int64, shape=(num_registers), name=key))
-      for key in ('mask', 'is_hint', 'is_local', 'is_free'))
-  observation_spec.update(
-      dict((key,
-            tensor_spec.BoundedTensorSpec(
-                dtype=tf.int64,
-                shape=(num_registers),
-                name=key,
-                minimum=0,
-                maximum=6)) for key in ('max_stage', 'min_stage')))
-  observation_spec.update(
-      dict((key,
-            tf.TensorSpec(dtype=tf.float32, shape=(num_registers), name=key))
-           for key in ('weighed_reads_by_max', 'weighed_writes_by_max',
-                       'weighed_read_writes_by_max', 'weighed_indvars_by_max',
-                       'hint_weights_by_max', 'start_bb_freq_by_max',
-                       'end_bb_freq_by_max', 'hottest_bb_freq_by_max',
-                       'liverange_size', 'use_def_density', 'nr_defs_and_uses',
-                       'nr_broken_hints', 'nr_urgent', 'nr_rematerializable')))
+  observation_spec = {
+      key: tf.TensorSpec(dtype=tf.int64, shape=(num_registers), name=key)
+      for key in ('mask', 'is_hint', 'is_local', 'is_free')
+  }
+  observation_spec.update({
+      key:
+          tensor_spec.BoundedTensorSpec(
+              dtype=tf.int64,
+              shape=(num_registers),
+              name=key,
+              minimum=0,
+              maximum=6) for key in ('max_stage', 'min_stage')
+  })
+  observation_spec.update({
+      key: tf.TensorSpec(dtype=tf.float32, shape=(num_registers), name=key)
+      for key in ('weighed_reads_by_max', 'weighed_writes_by_max',
+                  'weighed_read_writes_by_max', 'weighed_indvars_by_max',
+                  'hint_weights_by_max', 'start_bb_freq_by_max',
+                  'end_bb_freq_by_max', 'hottest_bb_freq_by_max',
+                  'liverange_size', 'use_def_density', 'nr_defs_and_uses',
+                  'nr_broken_hints', 'nr_urgent', 'nr_rematerializable')
+  })
   observation_spec['progress'] = tensor_spec.BoundedTensorSpec(
       dtype=tf.float32, shape=(), name='progress', minimum=0, maximum=1)