diff --git a/compiler_opt/benchmark/benchmark_chromium.py b/compiler_opt/benchmark/benchmark_chromium.py
index 351084bb..9850d81e 100644
--- a/compiler_opt/benchmark/benchmark_chromium.py
+++ b/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,
diff --git a/compiler_opt/benchmark/benchmark_report_test.py b/compiler_opt/benchmark/benchmark_report_test.py
index 5b5f0844..706a0945 100644
--- a/compiler_opt/benchmark/benchmark_report_test.py
+++ b/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))
diff --git a/compiler_opt/benchmark/filter_tests.py b/compiler_opt/benchmark/filter_tests.py
index d938915a..01d1d999 100644
--- a/compiler_opt/benchmark/filter_tests.py
+++ b/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
diff --git a/compiler_opt/benchmark/gtest_executable_utils.py b/compiler_opt/benchmark/gtest_executable_utils.py
index 44180f8a..b99c0464 100644
--- a/compiler_opt/benchmark/gtest_executable_utils.py
+++ b/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)
diff --git a/compiler_opt/distributed/local/local_worker_manager.py b/compiler_opt/distributed/local/local_worker_manager.py
index 021f478c..1a336f97 100644
--- a/compiler_opt/distributed/local/local_worker_manager.py
+++ b/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
diff --git a/compiler_opt/es/blackbox_learner.py b/compiler_opt/es/blackbox_learner.py
index e6252215..1d3a9844 100644
--- a/compiler_opt/es/blackbox_learner.py
+++ b/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)
diff --git a/compiler_opt/es/regalloc_trace/regalloc_trace_worker.py b/compiler_opt/es/regalloc_trace/regalloc_trace_worker.py
index 22a045d1..2e53c5dd 100644
--- a/compiler_opt/es/regalloc_trace/regalloc_trace_worker.py
+++ b/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:
diff --git a/compiler_opt/rl/corpus.py b/compiler_opt/rl/corpus.py
index 3285611c..a828a537 100644
--- a/compiler_opt/rl/corpus.py
+++ b/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)):
diff --git a/compiler_opt/rl/data_reader.py b/compiler_opt/rl/data_reader.py
index 4f4cb6c4..5656ce70 100644
--- a/compiler_opt/rl/data_reader.py
+++ b/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,
diff --git a/compiler_opt/rl/imitation_learning/generate_bc_trajectories_lib.py b/compiler_opt/rl/imitation_learning/generate_bc_trajectories_lib.py
index 377ed3ca..6eb13fe9 100644
--- a/compiler_opt/rl/imitation_learning/generate_bc_trajectories_lib.py
+++ b/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
diff --git a/compiler_opt/rl/imitation_learning/generate_bc_trajectories_test.py b/compiler_opt/rl/imitation_learning/generate_bc_trajectories_test.py
index f87551e1..79c99186 100644
--- a/compiler_opt/rl/imitation_learning/generate_bc_trajectories_test.py
+++ b/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__':
diff --git a/compiler_opt/rl/imitation_learning/weighted_bc_trainer_lib.py b/compiler_opt/rl/imitation_learning/weighted_bc_trainer_lib.py
index 76cfcb8d..6765fd6d 100644
--- a/compiler_opt/rl/imitation_learning/weighted_bc_trainer_lib.py
+++ b/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()
diff --git a/compiler_opt/rl/inlining/config.py b/compiler_opt/rl/inlining/config.py
index a8bd6927..47fb4dac 100644
--- a/compiler_opt/rl/inlining/config.py
+++ b/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(
diff --git a/compiler_opt/rl/inlining/imitation_learning_config.py b/compiler_opt/rl/inlining/imitation_learning_config.py
index 60ffbe49..f8976827 100644
--- a/compiler_opt/rl/inlining/imitation_learning_config.py
+++ b/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)
diff --git a/compiler_opt/rl/log_reader.py b/compiler_opt/rl/log_reader.py
index aadda004..0dda75bb 100644
--- a/compiler_opt/rl/log_reader.py
+++ b/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:
diff --git a/compiler_opt/rl/policy_saver.py b/compiler_opt/rl/policy_saver.py
index f314d585..f77fdc89 100644
--- a/compiler_opt/rl/policy_saver.py
+++ b/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
diff --git a/compiler_opt/rl/regalloc/config.py b/compiler_opt/rl/regalloc/config.py
index 18a0663a..a3fdb012 100644
--- a/compiler_opt/rl/regalloc/config.py
+++ b/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)
diff --git a/compiler_opt/rl/regalloc_priority/config.py b/compiler_opt/rl/regalloc_priority/config.py
index 68795512..b912711f 100644
--- a/compiler_opt/rl/regalloc_priority/config.py
+++ b/compiler_opt/rl/regalloc_priority/config.py
@@ -23,9 +23,10 @@
@gin.configurable()
def get_regalloc_signature_spec():
- observation_spec = dict(
- (key, tf.TensorSpec(dtype=tf.int64, shape=(), name=key))
- for key in ('li_size', 'stage'))
+ observation_spec = {
+ key: tf.TensorSpec(dtype=tf.int64, shape=(), name=key)
+ for key in ('li_size', 'stage')
+ }
observation_spec['weight'] = tf.TensorSpec(
dtype=tf.float32, shape=(), name='weight')
diff --git a/compiler_opt/testing/model_test_utils.py b/compiler_opt/testing/model_test_utils.py
index 860f4537..556602e5 100644
--- a/compiler_opt/testing/model_test_utils.py
+++ b/compiler_opt/testing/model_test_utils.py
@@ -42,9 +42,10 @@ def gen_test_model(outdir: str):
def get_input_signature():
"""Returns (time_step_spec, action_spec) for LLVM register allocation."""
- inputs = dict(
- (key, tf.TensorSpec(dtype=tf.int64, shape=(num_registers), name=key))
- for key in per_register_feature_list)
+ inputs = {
+ key: tf.TensorSpec(dtype=tf.int64, shape=(num_registers), name=key)
+ for key in per_register_feature_list
+ }
return inputs
module = tf.Module()
diff --git a/compiler_opt/tools/feature_importance.py b/compiler_opt/tools/feature_importance.py
index f899e2bb..09b022fc 100644
--- a/compiler_opt/tools/feature_importance.py
+++ b/compiler_opt/tools/feature_importance.py
@@ -102,7 +102,7 @@ def main(_):
observation, total_size)
flattened_input = numpy.expand_dims(flattened_input, axis=0)
dataset = numpy.empty((_NUM_EXAMPLES.value, total_size))
- for i in range(0, _NUM_EXAMPLES.value):
+ for i in range(_NUM_EXAMPLES.value):
raw_trajectory = next(dataset_iter)
observation = feature_importance_utils.process_raw_trajectory(
raw_trajectory)
diff --git a/compiler_opt/tools/feature_importance_utils.py b/compiler_opt/tools/feature_importance_utils.py
index 35b5593c..6a32a019 100644
--- a/compiler_opt/tools/feature_importance_utils.py
+++ b/compiler_opt/tools/feature_importance_utils.py
@@ -139,7 +139,7 @@ def collapse_values(input_signature: SignatureType,
shap_values: A numpy array of shap values that need to be processed.
"""
output_shap_values = numpy.empty((num_examples, len(input_signature)))
- for i in range(0, num_examples):
+ for i in range(num_examples):
current_index = 0
current_feature = 0
for input_key in input_signature:
diff --git a/compiler_opt/tools/generate_vocab.py b/compiler_opt/tools/generate_vocab.py
index 59c0cea0..38fbd213 100644
--- a/compiler_opt/tools/generate_vocab.py
+++ b/compiler_opt/tools/generate_vocab.py
@@ -171,7 +171,7 @@ def main(_) -> None:
data_list = data_list.swapaxes(0, 1)
with mp.Pool(FLAGS.parallelism) as pool:
- feature_names = list(sorted(sequence_features))
+ feature_names = sorted(sequence_features)
for i, feature_values_arrays in enumerate(data_list):
pool.apply_async(_generate_vocab, (
feature_values_arrays,