def main(argv):
if len(argv) > 1:
raise app.UsageError('Too many command-line arguments.')
# Initialize runner and file_suffix according to compile_task.
if FLAGS.compile_task == 'inlining':
runner = inlining_runner.InliningRunner(
clang_path=FLAGS.clang_path, llvm_size_path=FLAGS.llvm_size_path)
file_suffix = ['.bc', '.cmd']
with open(os.path.join(FLAGS.data_path, 'module_paths'), 'r') as f:
module_paths = [
os.path.join(FLAGS.data_path, name.rstrip('\n')) for name in f
]
file_paths = [
tuple([p + suffix for suffix in file_suffix]) for p in module_paths
]
# Sampling if needed.
if FLAGS.sampling_rate < 1:
sampled_modules = int(len(file_paths) * FLAGS.sampling_rate)
file_paths = random.sample(file_paths, k=sampled_modules)
ctx = multiprocessing.get_context()
pool = ctx.Pool(FLAGS.num_workers)
index = 0
total_successful_examples = 0
with tf.io.TFRecordWriter(FLAGS.output_path) as file_writer:
while index < len(file_paths):
# Shard data collection and sink to tfrecord periodically to avoid OOM.
next_index = min(index + _BATCH_SIZE, len(file_paths))
sharded_file_paths = file_paths[index:next_index]
index = next_index
results = [
pool.apply_async(runner.collect_data, (path, '', None))
for path in sharded_file_paths
]
# Wait till all jobs finish.
waiting_time = 0
while True:
if sum([not r.ready() for r in results]) == 0:
break
logging.info('%d/%d: %d of %d modules finished in %d seconds.', index,
len(file_paths), sum([r.ready() for r in results]),
len(sharded_file_paths), waiting_time)
time.sleep(1)
waiting_time += 1
# Write successful examples to tfrecord.
successful_count = len(
[file_writer.write(r.get()[0]) for r in results if r.successful()])
logging.info('%d/%d: %d of %d modules succeeded.', index, len(file_paths),
successful_count, len(sharded_file_paths))
total_successful_examples += successful_count
logging.info('%d of %d modules succeeded in total.',
total_successful_examples, len(file_paths))
def test_exception_handling(self, mock_run_inlining):
mock_run_inlining.side_effect = subprocess.CalledProcessError(
returncode=1, cmd='error')
inliner = inlining_runner.InliningRunner('', '')
with self.assertRaisesRegexp(subprocess.CalledProcessError, 'error'):
_, _ = inliner.collect_data(
file_paths=('bc', 'cmd'),
tf_policy_path='policy_path',
default_policy_size=None)
self.assertEqual(1, mock_run_inlining.call_count)
def test_default(self, mock_run_inlining):
mock_run_inlining.side_effect = _inlining
inliner = inlining_runner.InliningRunner('', '')
example, default_size = inliner.collect_data(
file_paths=('bc', 'cmd'), tf_policy_path='', default_policy_size=None)
self.assertEqual(2, mock_run_inlining.call_count)
expected_example = _get_sequence_example_with_reward(
_DEFAULT_FEATURE_VALUE, 0)
self.assertProtoEquals(expected_example,
tf.train.SequenceExample.FromString(example))
self.assertEqual(_DEFAULT_NATIVE_SIZE, default_size)
def train_eval(agent_name='ppo',
warmstart_policy_dir=None,
num_policy_iterations=0,
num_iterations=100,
batch_size=64,
train_sequence_length=1,
deploy_policy_name='saved_policy'):
"""Train for LLVM inliner."""
root_dir = FLAGS.root_dir
# Initialize trainer and policy saver.
time_step_spec, action_spec = config.create_signature_specs(config.CONFIG)
tf_agent = agent_creators.create_agent(agent_name, time_step_spec,
action_spec)
llvm_trainer = trainer.Trainer(root_dir=root_dir, agent=tf_agent)
policy_dict = {
'saved_policy': tf_agent.policy,
'saved_collect_policy': tf_agent.collect_policy,
}
saver = policy_saver.PolicySaver(policy_dict=policy_dict)
if warmstart_policy_dir:
warmstart_policy = policy_loader.load(warmstart_policy_dir)
tf_agent.policy.update(policy=warmstart_policy,
tau=1.0,
tau_non_trainable=None,
sort_variables_by_name=False)
with open(os.path.join(FLAGS.data_path, 'module_paths'), 'r') as f:
module_paths = [
os.path.join(FLAGS.data_path, name.rstrip('\n')) for name in f
]
file_paths = [(path + '.bc', path + '.cmd') for path in module_paths]
runner = inlining_runner.InliningRunner(
clang_path=FLAGS.clang_path, llvm_size_path=FLAGS.llvm_size_path)
sequence_example_iterator_fn = (
data_reader.create_sequence_example_iterator_fn(
agent_name=agent_name,
config=config.CONFIG,
batch_size=batch_size,
train_sequence_length=train_sequence_length))
data_collector = local_data_collector.LocalDataCollector(
file_paths=file_paths,
num_workers=FLAGS.num_workers,
num_modules=FLAGS.num_modules,
runner=runner.collect_data,
parser=sequence_example_iterator_fn)
for policy_iteration in range(num_policy_iterations):
policy_path = os.path.join(root_dir, 'policy', str(policy_iteration))
saver.save(policy_path)
dataset_iter = data_collector.collect_data(
policy_path=os.path.join(policy_path, deploy_policy_name))
llvm_trainer.train(dataset_iter, num_iterations)
data_collector.on_dataset_consumed(dataset_iter)
# Save final policy.
saver.save(root_dir)