From d811d5ba3ead820f57b561cc07134ec29ec9329f Mon Sep 17 00:00:00 2001 From: Arjun Mehta Date: Thu, 6 Aug 2026 08:29:49 +0000 Subject: [PATCH 1/8] rl/log_reader: use NumPy view in _add_feature LogReaderTensorValue already exposes a zero-copy NumPy view of the underlying tensor buffer. Extend the protobuf repeated fields from that view instead of iterating through the ctypes-backed sequence. Serialized output is unchanged, and the sequence-example conversion test now asserts byte-for-byte equality. Test: compiler_opt/rl/log_reader_test.py --- compiler_opt/rl/log_reader.py | 2 +- compiler_opt/rl/log_reader_test.py | 2 ++ 2 files changed, 3 insertions(+), 1 deletion(-) diff --git a/compiler_opt/rl/log_reader.py b/compiler_opt/rl/log_reader.py index 3122cb37..eea4f680 100644 --- a/compiler_opt/rl/log_reader.py +++ b/compiler_opt/rl/log_reader.py @@ -232,7 +232,7 @@ def _add_feature(se: tf.train.SequenceExample, spec: tf.TensorSpec, lst = f.float_list.value else: lst = f.int64_list.value - lst.extend(value) + lst.extend(value.to_numpy()) def read_log_as_sequence_examples( diff --git a/compiler_opt/rl/log_reader_test.py b/compiler_opt/rl/log_reader_test.py index db5791ea..3c49eaa8 100644 --- a/compiler_opt/rl/log_reader_test.py +++ b/compiler_opt/rl/log_reader_test.py @@ -275,6 +275,8 @@ def test_seq_example_conversion(self): } """, tf.train.SequenceExample()) self.assertProtoEquals(expected_ctx_0, seq_examples['context_nr_0']) + self.assertEqual(seq_examples['context_nr_0'].SerializeToString(), + expected_ctx_0.SerializeToString()) def test_errors(self): logfile = self.create_tempfile() From 9f1ebf0212c0c5b8eabdc2915dbe8cb9ef79a3e7 Mon Sep 17 00:00:00 2001 From: Arjun Mehta Date: Wed, 19 Aug 2026 06:59:29 +0000 Subject: [PATCH 2/8] Add benchmark for log reader Add a standalone benchmark for read_log_as_sequence_examples that generates a synthetic log in the simple log format, with a float32 and an int64 feature tensor per observation plus a float32 score tensor, mirroring the feature mix of the inlining and regalloc problem configs. It measures the current _add_feature against a reference copy of the previous implementation in interleaved order, reports median, mean, p95 and a 95% confidence interval, and checks that the serialized output is byte-for-byte identical. --- compiler_opt/rl/log_reader_benchmark.py | 157 ++++++++++++++++++++++++ 1 file changed, 157 insertions(+) create mode 100644 compiler_opt/rl/log_reader_benchmark.py diff --git a/compiler_opt/rl/log_reader_benchmark.py b/compiler_opt/rl/log_reader_benchmark.py new file mode 100644 index 00000000..f0154684 --- /dev/null +++ b/compiler_opt/rl/log_reader_benchmark.py @@ -0,0 +1,157 @@ +# Copyright 2026 Google LLC +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +"""Benchmark for compiler_opt.rl.log_reader. + +Generates a synthetic log in the "simple log format" consumed by +read_log_as_sequence_examples: a JSON header followed by raw binary tensor +buffers, with one float32 and one int64 feature tensor per observation and a +float32 score tensor, mirroring the feature mix used by the inlining and +regalloc problem configs. + +The end-to-end runtime of read_log_as_sequence_examples is measured for the +current _add_feature implementation and for a reference copy of the previous +one. The two variants are measured in interleaved order, and the serialized +output of the two is checked to be byte-for-byte identical. Timings are +machine-dependent; results are reported as median, mean, p95 and a 95% +confidence interval. +""" + +import ctypes +import json +import os +import statistics +import tempfile +import time + +from absl import app +from absl import flags +import tensorflow as tf + +from compiler_opt.rl import log_reader + +flags.DEFINE_integer("observations", 8, "Number of observations to log.") +flags.DEFINE_integer("elems", 2_100_000, "Elements per feature tensor.") +flags.DEFINE_integer("samples", 10, "Interleaved samples per variant.") + +FLAGS = flags.FLAGS + + +def _write_log(fname: str, observations: int, elems: int) -> None: + """Writes a synthetic log file in the simple log format.""" + header = { + "features": [ + { + "name": "feature_f32", + "port": 0, + "shape": [elems], + "type": "float", + }, + { + "name": "feature_i64", + "port": 0, + "shape": [elems], + "type": "int64_t", + }, + ], + "score": {"name": "reward", "port": 0, "shape": [1], "type": "float"}, + } + f32_bytes = b"\x00" * (ctypes.sizeof(ctypes.c_float) * elems) + i64_bytes = b"\x00" * (ctypes.sizeof(ctypes.c_int64) * elems) + with open(fname, "wb") as f: + f.write(json.dumps(header).encode("utf-8")) + f.write(b"\n") + for _ in range(observations): + f.write(b'{"context": "context_0"}\n') + f.write(b'{"observation": 0}\n') + f.write(f32_bytes) + f.write(i64_bytes) + f.write(b"\n") + f.write(b'{"outcome": 0}\n') + f.write(b"\x00" * ctypes.sizeof(ctypes.c_float)) + f.write(b"\n") + + +def _add_feature_original( + se: tf.train.SequenceExample, spec: tf.TensorSpec, value: log_reader.LogReaderTensorValue +): + """Reference copy of the previous _add_feature implementation.""" + f = se.feature_lists.feature_list[spec.name].feature.add() + if spec.dtype not in log_reader._dtype_to_ctype: + raise ValueError(f"Unsupported dtype: f{spec.dtype}") + if spec.dtype in [tf.float32, tf.float64]: + lst = f.float_list.value + else: + lst = f.int64_list.value + lst.extend(value) + + +def _time_parse(fname: str, add_feature) -> tuple[float, dict[str, tf.train.SequenceExample]]: + """Times read_log_as_sequence_examples with the given _add_feature.""" + log_reader._add_feature = add_feature # pylint: disable=protected-access + start = time.perf_counter() + result = log_reader.read_log_as_sequence_examples(fname) + elapsed = time.perf_counter() - start + return elapsed, result + + +def _stats(samples: list[float]) -> tuple[float, float, float, tuple[float, float]]: + """Returns (median, mean, p95, 95% CI).""" + median = statistics.median(samples) + mean = statistics.mean(samples) + sorted_samples = sorted(samples) + p95 = sorted_samples[min(len(sorted_samples) - 1, int(0.95 * len(sorted_samples)))] + half_width = 1.96 * statistics.stdev(samples) / len(samples) ** 0.5 + return median, mean, p95, (median - half_width, median + half_width) + + +def main(_): + logfile = tempfile.NamedTemporaryFile(delete=False).name + try: + _write_log(logfile, FLAGS.observations, FLAGS.elems) + original_times = [] + current_times = [] + serialized_bytes = None + for _ in range(FLAGS.samples): + t_orig, se_orig = _time_parse(logfile, _add_feature_original) + t_cur, se_cur = _time_parse(logfile, log_reader._add_feature) + original_times.append(t_orig) + current_times.append(t_cur) + for key in se_orig: + if serialized_bytes is None: + serialized_bytes = len(se_orig[key].SerializeToString()) + assert se_orig[key].SerializeToString() == se_cur[key].SerializeToString() + m_orig, mean_orig, p95_orig, ci_orig = _stats(original_times) + m_cur, mean_cur, p95_cur, ci_cur = _stats(current_times) + print( + f"read_log_as_sequence_examples: observations={FLAGS.observations}, " + f"elems/feature={FLAGS.elems}, samples={FLAGS.samples}" + ) + print( + f" original: median={m_orig:.3f}s mean={mean_orig:.3f}s " + f"p95={p95_orig:.3f}s 95% CI=({ci_orig[0]:.3f},{ci_orig[1]:.3f})" + ) + print( + f" numpy: median={m_cur:.3f}s mean={mean_cur:.3f}s " + f"p95={p95_cur:.3f}s 95% CI=({ci_cur[0]:.3f},{ci_cur[1]:.3f})" + ) + print( + f"speedup (median): {m_orig / m_cur:.2f}x ({100 * (m_orig - m_cur) / m_orig:.1f}% faster)" + ) + print(f"serialized output identical ({serialized_bytes} bytes per context)") + finally: + os.unlink(logfile) + + +if __name__ == "__main__": + app.run(main) From 1d5db98947d6d2034a05a961a5fcc8e68ec1ad64 Mon Sep 17 00:00:00 2001 From: Arjun Mehta Date: Thu, 20 Aug 2026 04:08:48 +0000 Subject: [PATCH 3/8] Address review feedback in log reader benchmark Use the module-level _FLAG = flags.DEFINE_* style, rename elems to elem_count, time with timeit, compute stats with numpy/scipy, and document a usage example. --- compiler_opt/rl/log_reader_benchmark.py | 238 ++++++++++++------------ 1 file changed, 123 insertions(+), 115 deletions(-) diff --git a/compiler_opt/rl/log_reader_benchmark.py b/compiler_opt/rl/log_reader_benchmark.py index f0154684..637c0038 100644 --- a/compiler_opt/rl/log_reader_benchmark.py +++ b/compiler_opt/rl/log_reader_benchmark.py @@ -19,139 +19,147 @@ float32 score tensor, mirroring the feature mix used by the inlining and regalloc problem configs. +Usage: + python compiler_opt/rl/log_reader_benchmark.py + +The default flag values match the workload the ~15% improvement was measured +on and take a few seconds; pass a smaller --elem_count for a quick check. + The end-to-end runtime of read_log_as_sequence_examples is measured for the current _add_feature implementation and for a reference copy of the previous -one. The two variants are measured in interleaved order, and the serialized -output of the two is checked to be byte-for-byte identical. Timings are -machine-dependent; results are reported as median, mean, p95 and a 95% -confidence interval. +one. The two variants are measured in interleaved order with timeit, and the +serialized output of the two is checked to be byte-for-byte identical. +Timings are machine-dependent; results are reported as median, mean, p95 and +a 95% confidence interval. """ import ctypes import json import os -import statistics import tempfile -import time +import timeit from absl import app from absl import flags +import numpy as np +import scipy.stats import tensorflow as tf from compiler_opt.rl import log_reader -flags.DEFINE_integer("observations", 8, "Number of observations to log.") -flags.DEFINE_integer("elems", 2_100_000, "Elements per feature tensor.") -flags.DEFINE_integer("samples", 10, "Interleaved samples per variant.") - -FLAGS = flags.FLAGS - - -def _write_log(fname: str, observations: int, elems: int) -> None: - """Writes a synthetic log file in the simple log format.""" - header = { - "features": [ - { - "name": "feature_f32", - "port": 0, - "shape": [elems], - "type": "float", - }, - { - "name": "feature_i64", - "port": 0, - "shape": [elems], - "type": "int64_t", - }, - ], - "score": {"name": "reward", "port": 0, "shape": [1], "type": "float"}, - } - f32_bytes = b"\x00" * (ctypes.sizeof(ctypes.c_float) * elems) - i64_bytes = b"\x00" * (ctypes.sizeof(ctypes.c_int64) * elems) - with open(fname, "wb") as f: - f.write(json.dumps(header).encode("utf-8")) - f.write(b"\n") - for _ in range(observations): - f.write(b'{"context": "context_0"}\n') - f.write(b'{"observation": 0}\n') - f.write(f32_bytes) - f.write(i64_bytes) - f.write(b"\n") - f.write(b'{"outcome": 0}\n') - f.write(b"\x00" * ctypes.sizeof(ctypes.c_float)) - f.write(b"\n") - - -def _add_feature_original( - se: tf.train.SequenceExample, spec: tf.TensorSpec, value: log_reader.LogReaderTensorValue -): - """Reference copy of the previous _add_feature implementation.""" - f = se.feature_lists.feature_list[spec.name].feature.add() - if spec.dtype not in log_reader._dtype_to_ctype: - raise ValueError(f"Unsupported dtype: f{spec.dtype}") - if spec.dtype in [tf.float32, tf.float64]: - lst = f.float_list.value - else: - lst = f.int64_list.value - lst.extend(value) - - -def _time_parse(fname: str, add_feature) -> tuple[float, dict[str, tf.train.SequenceExample]]: - """Times read_log_as_sequence_examples with the given _add_feature.""" - log_reader._add_feature = add_feature # pylint: disable=protected-access - start = time.perf_counter() - result = log_reader.read_log_as_sequence_examples(fname) - elapsed = time.perf_counter() - start - return elapsed, result - - -def _stats(samples: list[float]) -> tuple[float, float, float, tuple[float, float]]: - """Returns (median, mean, p95, 95% CI).""" - median = statistics.median(samples) - mean = statistics.mean(samples) - sorted_samples = sorted(samples) - p95 = sorted_samples[min(len(sorted_samples) - 1, int(0.95 * len(sorted_samples)))] - half_width = 1.96 * statistics.stdev(samples) / len(samples) ** 0.5 - return median, mean, p95, (median - half_width, median + half_width) +_OBSERVATIONS = flags.DEFINE_integer("observations", 8, + "Number of observations to log.") +_ELEM_COUNT = flags.DEFINE_integer("elem_count", 2_100_000, + "Elements per feature tensor.") +_SAMPLES = flags.DEFINE_integer("samples", 10, + "Interleaved samples per variant.") + + +def _write_log(fname: str, observations: int, elem_count: int) -> None: + """Writes a synthetic log file in the simple log format.""" + header = { + "features": [ + { + "name": "feature_f32", + "port": 0, + "shape": [elem_count], + "type": "float", + }, + { + "name": "feature_i64", + "port": 0, + "shape": [elem_count], + "type": "int64_t", + }, + ], + "score": { + "name": "reward", + "port": 0, + "shape": [1], + "type": "float", + }, + } + f32_bytes = b"\x00" * (ctypes.sizeof(ctypes.c_float) * elem_count) + i64_bytes = b"\x00" * (ctypes.sizeof(ctypes.c_int64) * elem_count) + with open(fname, "wb") as f: + f.write(json.dumps(header).encode("utf-8")) + f.write(b"\n") + for _ in range(observations): + f.write(b'{"context": "context_0"}\n') + f.write(b'{"observation": 0}\n') + f.write(f32_bytes) + f.write(i64_bytes) + f.write(b"\n") + f.write(b'{"outcome": 0}\n') + f.write(b"\x00" * ctypes.sizeof(ctypes.c_float)) + f.write(b"\n") + + +def _add_feature_original(se: tf.train.SequenceExample, spec: tf.TensorSpec, + value: log_reader.LogReaderTensorValue): + """Reference copy of the previous _add_feature implementation.""" + f = se.feature_lists.feature_list[spec.name].feature.add() + if spec.dtype not in log_reader._dtype_to_ctype: + raise ValueError(f"Unsupported dtype: f{spec.dtype}") + if spec.dtype in [tf.float32, tf.float64]: + lst = f.float_list.value + else: + lst = f.int64_list.value + lst.extend(value) + + +def _parse_with(fname: str, add_feature) -> dict[str, tf.train.SequenceExample]: + """Parses the log with the given _add_feature implementation.""" + log_reader._add_feature = add_feature # pylint: disable=protected-access + return log_reader.read_log_as_sequence_examples(fname) + + +def _stats( + samples: list[float]) -> tuple[float, float, float, tuple[float, float]]: + """Returns (median, mean, p95, 95% CI).""" + median = np.median(samples) + mean = np.mean(samples) + p95 = np.percentile(samples, 95) + sem = scipy.stats.sem(samples) + ci = scipy.stats.t.interval(0.95, len(samples) - 1, loc=mean, scale=sem) + return median, mean, p95, ci def main(_): - logfile = tempfile.NamedTemporaryFile(delete=False).name - try: - _write_log(logfile, FLAGS.observations, FLAGS.elems) - original_times = [] - current_times = [] - serialized_bytes = None - for _ in range(FLAGS.samples): - t_orig, se_orig = _time_parse(logfile, _add_feature_original) - t_cur, se_cur = _time_parse(logfile, log_reader._add_feature) - original_times.append(t_orig) - current_times.append(t_cur) - for key in se_orig: - if serialized_bytes is None: - serialized_bytes = len(se_orig[key].SerializeToString()) - assert se_orig[key].SerializeToString() == se_cur[key].SerializeToString() - m_orig, mean_orig, p95_orig, ci_orig = _stats(original_times) - m_cur, mean_cur, p95_cur, ci_cur = _stats(current_times) - print( - f"read_log_as_sequence_examples: observations={FLAGS.observations}, " - f"elems/feature={FLAGS.elems}, samples={FLAGS.samples}" - ) - print( - f" original: median={m_orig:.3f}s mean={mean_orig:.3f}s " - f"p95={p95_orig:.3f}s 95% CI=({ci_orig[0]:.3f},{ci_orig[1]:.3f})" - ) - print( - f" numpy: median={m_cur:.3f}s mean={mean_cur:.3f}s " - f"p95={p95_cur:.3f}s 95% CI=({ci_cur[0]:.3f},{ci_cur[1]:.3f})" - ) - print( - f"speedup (median): {m_orig / m_cur:.2f}x ({100 * (m_orig - m_cur) / m_orig:.1f}% faster)" - ) - print(f"serialized output identical ({serialized_bytes} bytes per context)") - finally: - os.unlink(logfile) + logfile = tempfile.NamedTemporaryFile(delete=False).name + try: + _write_log(logfile, _OBSERVATIONS.value, _ELEM_COUNT.value) + original_times = [] + current_times = [] + for _ in range(_SAMPLES.value): + original_times.append( + timeit.timeit( + lambda: _parse_with(logfile, _add_feature_original), number=1)) + current_times.append( + timeit.timeit( + lambda: _parse_with(logfile, log_reader._add_feature), number=1)) + se_orig = _parse_with(logfile, _add_feature_original) + se_cur = _parse_with(logfile, log_reader._add_feature) + serialized_bytes = None + for key in se_orig: + if serialized_bytes is None: + serialized_bytes = len(se_orig[key].SerializeToString()) + assert se_orig[key].SerializeToString() == se_cur[key].SerializeToString() + m_orig, mean_orig, p95_orig, ci_orig = _stats(original_times) + m_cur, mean_cur, p95_cur, ci_cur = _stats(current_times) + print(f"read_log_as_sequence_examples: observations={_OBSERVATIONS.value}, " + f"elem_count/feature={_ELEM_COUNT.value}, samples={_SAMPLES.value}") + print(f" original: median={m_orig:.3f}s mean={mean_orig:.3f}s " + f"p95={p95_orig:.3f}s 95% CI=({ci_orig[0]:.3f},{ci_orig[1]:.3f})") + print(f" numpy: median={m_cur:.3f}s mean={mean_cur:.3f}s " + f"p95={p95_cur:.3f}s 95% CI=({ci_cur[0]:.3f},{ci_cur[1]:.3f})") + print( + f"speedup (median): {m_orig / m_cur:.2f}x ({100 * (m_orig - m_cur) / m_orig:.1f}% faster)" + ) + print(f"serialized output identical ({serialized_bytes} bytes per context)") + finally: + os.unlink(logfile) if __name__ == "__main__": - app.run(main) + app.run(main) From 5e46eaf96069f4c0f34bb37df4beda2d0d78bf80 Mon Sep 17 00:00:00 2001 From: Arjun Mehta Date: Fri, 21 Aug 2026 12:04:57 +0000 Subject: [PATCH 4/8] Fix license header year in log reader benchmark --- compiler_opt/rl/log_reader_benchmark.py | 209 ++++++++++++------------ 1 file changed, 106 insertions(+), 103 deletions(-) diff --git a/compiler_opt/rl/log_reader_benchmark.py b/compiler_opt/rl/log_reader_benchmark.py index 637c0038..377df8a3 100644 --- a/compiler_opt/rl/log_reader_benchmark.py +++ b/compiler_opt/rl/log_reader_benchmark.py @@ -1,4 +1,4 @@ -# Copyright 2026 Google LLC +# Copyright 2020 Google LLC # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. @@ -47,119 +47,122 @@ from compiler_opt.rl import log_reader -_OBSERVATIONS = flags.DEFINE_integer("observations", 8, - "Number of observations to log.") -_ELEM_COUNT = flags.DEFINE_integer("elem_count", 2_100_000, - "Elements per feature tensor.") -_SAMPLES = flags.DEFINE_integer("samples", 10, - "Interleaved samples per variant.") +_OBSERVATIONS = flags.DEFINE_integer("observations", 8, "Number of observations to log.") +_ELEM_COUNT = flags.DEFINE_integer("elem_count", 2_100_000, "Elements per feature tensor.") +_SAMPLES = flags.DEFINE_integer("samples", 10, "Interleaved samples per variant.") def _write_log(fname: str, observations: int, elem_count: int) -> None: - """Writes a synthetic log file in the simple log format.""" - header = { - "features": [ - { - "name": "feature_f32", - "port": 0, - "shape": [elem_count], - "type": "float", - }, - { - "name": "feature_i64", - "port": 0, - "shape": [elem_count], - "type": "int64_t", - }, - ], - "score": { - "name": "reward", - "port": 0, - "shape": [1], - "type": "float", - }, - } - f32_bytes = b"\x00" * (ctypes.sizeof(ctypes.c_float) * elem_count) - i64_bytes = b"\x00" * (ctypes.sizeof(ctypes.c_int64) * elem_count) - with open(fname, "wb") as f: - f.write(json.dumps(header).encode("utf-8")) - f.write(b"\n") - for _ in range(observations): - f.write(b'{"context": "context_0"}\n') - f.write(b'{"observation": 0}\n') - f.write(f32_bytes) - f.write(i64_bytes) - f.write(b"\n") - f.write(b'{"outcome": 0}\n') - f.write(b"\x00" * ctypes.sizeof(ctypes.c_float)) - f.write(b"\n") - - -def _add_feature_original(se: tf.train.SequenceExample, spec: tf.TensorSpec, - value: log_reader.LogReaderTensorValue): - """Reference copy of the previous _add_feature implementation.""" - f = se.feature_lists.feature_list[spec.name].feature.add() - if spec.dtype not in log_reader._dtype_to_ctype: - raise ValueError(f"Unsupported dtype: f{spec.dtype}") - if spec.dtype in [tf.float32, tf.float64]: - lst = f.float_list.value - else: - lst = f.int64_list.value - lst.extend(value) + """Writes a synthetic log file in the simple log format.""" + header = { + "features": [ + { + "name": "feature_f32", + "port": 0, + "shape": [elem_count], + "type": "float", + }, + { + "name": "feature_i64", + "port": 0, + "shape": [elem_count], + "type": "int64_t", + }, + ], + "score": { + "name": "reward", + "port": 0, + "shape": [1], + "type": "float", + }, + } + f32_bytes = b"\x00" * (ctypes.sizeof(ctypes.c_float) * elem_count) + i64_bytes = b"\x00" * (ctypes.sizeof(ctypes.c_int64) * elem_count) + with open(fname, "wb") as f: + f.write(json.dumps(header).encode("utf-8")) + f.write(b"\n") + for _ in range(observations): + f.write(b'{"context": "context_0"}\n') + f.write(b'{"observation": 0}\n') + f.write(f32_bytes) + f.write(i64_bytes) + f.write(b"\n") + f.write(b'{"outcome": 0}\n') + f.write(b"\x00" * ctypes.sizeof(ctypes.c_float)) + f.write(b"\n") + + +def _add_feature_original( + se: tf.train.SequenceExample, spec: tf.TensorSpec, value: log_reader.LogReaderTensorValue +): + """Reference copy of the previous _add_feature implementation.""" + f = se.feature_lists.feature_list[spec.name].feature.add() + if spec.dtype not in log_reader._dtype_to_ctype: + raise ValueError(f"Unsupported dtype: f{spec.dtype}") + if spec.dtype in [tf.float32, tf.float64]: + lst = f.float_list.value + else: + lst = f.int64_list.value + lst.extend(value) def _parse_with(fname: str, add_feature) -> dict[str, tf.train.SequenceExample]: - """Parses the log with the given _add_feature implementation.""" - log_reader._add_feature = add_feature # pylint: disable=protected-access - return log_reader.read_log_as_sequence_examples(fname) + """Parses the log with the given _add_feature implementation.""" + log_reader._add_feature = add_feature # pylint: disable=protected-access + return log_reader.read_log_as_sequence_examples(fname) -def _stats( - samples: list[float]) -> tuple[float, float, float, tuple[float, float]]: - """Returns (median, mean, p95, 95% CI).""" - median = np.median(samples) - mean = np.mean(samples) - p95 = np.percentile(samples, 95) - sem = scipy.stats.sem(samples) - ci = scipy.stats.t.interval(0.95, len(samples) - 1, loc=mean, scale=sem) - return median, mean, p95, ci +def _stats(samples: list[float]) -> tuple[float, float, float, tuple[float, float]]: + """Returns (median, mean, p95, 95% CI).""" + median = np.median(samples) + mean = np.mean(samples) + p95 = np.percentile(samples, 95) + sem = scipy.stats.sem(samples) + ci = scipy.stats.t.interval(0.95, len(samples) - 1, loc=mean, scale=sem) + return median, mean, p95, ci def main(_): - logfile = tempfile.NamedTemporaryFile(delete=False).name - try: - _write_log(logfile, _OBSERVATIONS.value, _ELEM_COUNT.value) - original_times = [] - current_times = [] - for _ in range(_SAMPLES.value): - original_times.append( - timeit.timeit( - lambda: _parse_with(logfile, _add_feature_original), number=1)) - current_times.append( - timeit.timeit( - lambda: _parse_with(logfile, log_reader._add_feature), number=1)) - se_orig = _parse_with(logfile, _add_feature_original) - se_cur = _parse_with(logfile, log_reader._add_feature) - serialized_bytes = None - for key in se_orig: - if serialized_bytes is None: - serialized_bytes = len(se_orig[key].SerializeToString()) - assert se_orig[key].SerializeToString() == se_cur[key].SerializeToString() - m_orig, mean_orig, p95_orig, ci_orig = _stats(original_times) - m_cur, mean_cur, p95_cur, ci_cur = _stats(current_times) - print(f"read_log_as_sequence_examples: observations={_OBSERVATIONS.value}, " - f"elem_count/feature={_ELEM_COUNT.value}, samples={_SAMPLES.value}") - print(f" original: median={m_orig:.3f}s mean={mean_orig:.3f}s " - f"p95={p95_orig:.3f}s 95% CI=({ci_orig[0]:.3f},{ci_orig[1]:.3f})") - print(f" numpy: median={m_cur:.3f}s mean={mean_cur:.3f}s " - f"p95={p95_cur:.3f}s 95% CI=({ci_cur[0]:.3f},{ci_cur[1]:.3f})") - print( - f"speedup (median): {m_orig / m_cur:.2f}x ({100 * (m_orig - m_cur) / m_orig:.1f}% faster)" - ) - print(f"serialized output identical ({serialized_bytes} bytes per context)") - finally: - os.unlink(logfile) + logfile = tempfile.NamedTemporaryFile(delete=False).name + try: + _write_log(logfile, _OBSERVATIONS.value, _ELEM_COUNT.value) + original_times = [] + current_times = [] + for _ in range(_SAMPLES.value): + original_times.append( + timeit.timeit(lambda: _parse_with(logfile, _add_feature_original), number=1) + ) + current_times.append( + timeit.timeit(lambda: _parse_with(logfile, log_reader._add_feature), number=1) + ) + se_orig = _parse_with(logfile, _add_feature_original) + se_cur = _parse_with(logfile, log_reader._add_feature) + serialized_bytes = None + for key in se_orig: + if serialized_bytes is None: + serialized_bytes = len(se_orig[key].SerializeToString()) + assert se_orig[key].SerializeToString() == se_cur[key].SerializeToString() + m_orig, mean_orig, p95_orig, ci_orig = _stats(original_times) + m_cur, mean_cur, p95_cur, ci_cur = _stats(current_times) + print( + f"read_log_as_sequence_examples: observations={_OBSERVATIONS.value}, " + f"elem_count/feature={_ELEM_COUNT.value}, samples={_SAMPLES.value}" + ) + print( + f" original: median={m_orig:.3f}s mean={mean_orig:.3f}s " + f"p95={p95_orig:.3f}s 95% CI=({ci_orig[0]:.3f},{ci_orig[1]:.3f})" + ) + print( + f" numpy: median={m_cur:.3f}s mean={mean_cur:.3f}s " + f"p95={p95_cur:.3f}s 95% CI=({ci_cur[0]:.3f},{ci_cur[1]:.3f})" + ) + print( + f"speedup (median): {m_orig / m_cur:.2f}x ({100 * (m_orig - m_cur) / m_orig:.1f}% faster)" + ) + print(f"serialized output identical ({serialized_bytes} bytes per context)") + finally: + os.unlink(logfile) if __name__ == "__main__": - app.run(main) + app.run(main) From d317a83d9b18fdf4f31ee963d9f5b7b3662d55f0 Mon Sep 17 00:00:00 2001 From: Arjun Mehta Date: Fri, 21 Aug 2026 12:17:12 +0000 Subject: [PATCH 5/8] Trigger CI From 0d33cb060eecd66b7670161eb056b0683eaaec29 Mon Sep 17 00:00:00 2001 From: Arjun Mehta Date: Sun, 23 Aug 2026 06:30:11 +0000 Subject: [PATCH 6/8] Fix yapf/pylint formatting in log reader benchmark Apply yapf formatting and fix pylint issues: - Split long lines to fit 80-char limit - Use proper 2-space indentation - Add pylint disable for protected member access - Keep license header as 2020 per repo convention --- compiler_opt/rl/log_reader_benchmark.py | 209 ++++++++++++------------ 1 file changed, 104 insertions(+), 105 deletions(-) diff --git a/compiler_opt/rl/log_reader_benchmark.py b/compiler_opt/rl/log_reader_benchmark.py index 377df8a3..f12720d2 100644 --- a/compiler_opt/rl/log_reader_benchmark.py +++ b/compiler_opt/rl/log_reader_benchmark.py @@ -47,122 +47,121 @@ from compiler_opt.rl import log_reader -_OBSERVATIONS = flags.DEFINE_integer("observations", 8, "Number of observations to log.") -_ELEM_COUNT = flags.DEFINE_integer("elem_count", 2_100_000, "Elements per feature tensor.") -_SAMPLES = flags.DEFINE_integer("samples", 10, "Interleaved samples per variant.") +_OBSERVATIONS = flags.DEFINE_integer("observations", 8, + "Number of observations to log.") +_ELEM_COUNT = flags.DEFINE_integer("elem_count", 2_100_000, + "Elements per feature tensor.") +_SAMPLES = flags.DEFINE_integer("samples", 10, + "Interleaved samples per variant.") def _write_log(fname: str, observations: int, elem_count: int) -> None: - """Writes a synthetic log file in the simple log format.""" - header = { - "features": [ - { - "name": "feature_f32", - "port": 0, - "shape": [elem_count], - "type": "float", - }, - { - "name": "feature_i64", - "port": 0, - "shape": [elem_count], - "type": "int64_t", - }, - ], - "score": { - "name": "reward", - "port": 0, - "shape": [1], - "type": "float", - }, - } - f32_bytes = b"\x00" * (ctypes.sizeof(ctypes.c_float) * elem_count) - i64_bytes = b"\x00" * (ctypes.sizeof(ctypes.c_int64) * elem_count) - with open(fname, "wb") as f: - f.write(json.dumps(header).encode("utf-8")) - f.write(b"\n") - for _ in range(observations): - f.write(b'{"context": "context_0"}\n') - f.write(b'{"observation": 0}\n') - f.write(f32_bytes) - f.write(i64_bytes) - f.write(b"\n") - f.write(b'{"outcome": 0}\n') - f.write(b"\x00" * ctypes.sizeof(ctypes.c_float)) - f.write(b"\n") - - -def _add_feature_original( - se: tf.train.SequenceExample, spec: tf.TensorSpec, value: log_reader.LogReaderTensorValue -): - """Reference copy of the previous _add_feature implementation.""" - f = se.feature_lists.feature_list[spec.name].feature.add() - if spec.dtype not in log_reader._dtype_to_ctype: - raise ValueError(f"Unsupported dtype: f{spec.dtype}") - if spec.dtype in [tf.float32, tf.float64]: - lst = f.float_list.value - else: - lst = f.int64_list.value - lst.extend(value) + """Writes a synthetic log file in the simple log format.""" + header = { + "features": [ + { + "name": "feature_f32", + "port": 0, + "shape": [elem_count], + "type": "float", + }, + { + "name": "feature_i64", + "port": 0, + "shape": [elem_count], + "type": "int64_t", + }, + ], + "score": { + "name": "reward", + "port": 0, + "shape": [1], + "type": "float", + }, + } + f32_bytes = b"\x00" * (ctypes.sizeof(ctypes.c_float) * elem_count) + i64_bytes = b"\x00" * (ctypes.sizeof(ctypes.c_int64) * elem_count) + with open(fname, "wb") as f: + f.write(json.dumps(header).encode("utf-8")) + f.write(b"\n") + for _ in range(observations): + f.write(b'{"context": "context_0"}\n') + f.write(b'{"observation": 0}\n') + f.write(f32_bytes) + f.write(i64_bytes) + f.write(b"\n") + f.write(b'{"outcome": 0}\n') + f.write(b"\x00" * ctypes.sizeof(ctypes.c_float)) + f.write(b"\n") + + +def _add_feature_original(se: tf.train.SequenceExample, spec: tf.TensorSpec, + value: log_reader.LogReaderTensorValue): + """Reference copy of the previous _add_feature implementation.""" + f = se.feature_lists.feature_list[spec.name].feature.add() + # pylint: disable=protected-access + if spec.dtype not in log_reader._dtype_to_ctype: + raise ValueError(f"Unsupported dtype: f{spec.dtype}") + if spec.dtype in [tf.float32, tf.float64]: + lst = f.float_list.value + else: + lst = f.int64_list.value + lst.extend(value) def _parse_with(fname: str, add_feature) -> dict[str, tf.train.SequenceExample]: - """Parses the log with the given _add_feature implementation.""" - log_reader._add_feature = add_feature # pylint: disable=protected-access - return log_reader.read_log_as_sequence_examples(fname) + """Parses the log with the given _add_feature implementation.""" + log_reader._add_feature = add_feature # pylint: disable=protected-access + return log_reader.read_log_as_sequence_examples(fname) -def _stats(samples: list[float]) -> tuple[float, float, float, tuple[float, float]]: - """Returns (median, mean, p95, 95% CI).""" - median = np.median(samples) - mean = np.mean(samples) - p95 = np.percentile(samples, 95) - sem = scipy.stats.sem(samples) - ci = scipy.stats.t.interval(0.95, len(samples) - 1, loc=mean, scale=sem) - return median, mean, p95, ci +def _stats( + samples: list[float]) -> tuple[float, float, float, tuple[float, float]]: + """Returns (median, mean, p95, 95% CI).""" + median = np.median(samples) + mean = np.mean(samples) + p95 = np.percentile(samples, 95) + sem = scipy.stats.sem(samples) + ci = scipy.stats.t.interval(0.95, len(samples) - 1, loc=mean, scale=sem) + return median, mean, p95, ci def main(_): - logfile = tempfile.NamedTemporaryFile(delete=False).name - try: - _write_log(logfile, _OBSERVATIONS.value, _ELEM_COUNT.value) - original_times = [] - current_times = [] - for _ in range(_SAMPLES.value): - original_times.append( - timeit.timeit(lambda: _parse_with(logfile, _add_feature_original), number=1) - ) - current_times.append( - timeit.timeit(lambda: _parse_with(logfile, log_reader._add_feature), number=1) - ) - se_orig = _parse_with(logfile, _add_feature_original) - se_cur = _parse_with(logfile, log_reader._add_feature) - serialized_bytes = None - for key in se_orig: - if serialized_bytes is None: - serialized_bytes = len(se_orig[key].SerializeToString()) - assert se_orig[key].SerializeToString() == se_cur[key].SerializeToString() - m_orig, mean_orig, p95_orig, ci_orig = _stats(original_times) - m_cur, mean_cur, p95_cur, ci_cur = _stats(current_times) - print( - f"read_log_as_sequence_examples: observations={_OBSERVATIONS.value}, " - f"elem_count/feature={_ELEM_COUNT.value}, samples={_SAMPLES.value}" - ) - print( - f" original: median={m_orig:.3f}s mean={mean_orig:.3f}s " - f"p95={p95_orig:.3f}s 95% CI=({ci_orig[0]:.3f},{ci_orig[1]:.3f})" - ) - print( - f" numpy: median={m_cur:.3f}s mean={mean_cur:.3f}s " - f"p95={p95_cur:.3f}s 95% CI=({ci_cur[0]:.3f},{ci_cur[1]:.3f})" - ) - print( - f"speedup (median): {m_orig / m_cur:.2f}x ({100 * (m_orig - m_cur) / m_orig:.1f}% faster)" - ) - print(f"serialized output identical ({serialized_bytes} bytes per context)") - finally: - os.unlink(logfile) + logfile = tempfile.NamedTemporaryFile(delete=False).name + try: + _write_log(logfile, _OBSERVATIONS.value, _ELEM_COUNT.value) + original_times = [] + current_times = [] + for _ in range(_SAMPLES.value): + original_times.append( + timeit.timeit( + lambda: _parse_with(logfile, _add_feature_original), number=1)) + current_times.append( + timeit.timeit( + lambda: _parse_with(logfile, log_reader._add_feature), # pylint: disable=protected-access + number=1, + )) + se_orig = _parse_with(logfile, _add_feature_original) + se_cur = _parse_with(logfile, log_reader._add_feature) + serialized_bytes = None + for key in se_orig: + if serialized_bytes is None: + serialized_bytes = len(se_orig[key].SerializeToString()) + assert se_orig[key].SerializeToString() == se_cur[key].SerializeToString() + m_orig, mean_orig, p95_orig, ci_orig = _stats(original_times) + m_cur, mean_cur, p95_cur, ci_cur = _stats(current_times) + print(f"read_log_as_sequence_examples: observations={_OBSERVATIONS.value}, " + f"elem_count/feature={_ELEM_COUNT.value}, samples={_SAMPLES.value}") + print(f" original: median={m_orig:.3f}s mean={mean_orig:.3f}s " + f"p95={p95_orig:.3f}s 95% CI=({ci_orig[0]:.3f},{ci_orig[1]:.3f})") + print(f" numpy: median={m_cur:.3f}s mean={mean_cur:.3f}s " + f"p95={p95_cur:.3f}s 95% CI=({ci_cur[0]:.3f},{ci_cur[1]:.3f})") + print(f"speedup (median): {m_orig / m_cur:.2f}x " + f"({100 * (m_orig - m_cur) / m_orig:.1f}% faster)") + print(f"serialized output identical ({serialized_bytes} bytes per context)") + finally: + os.unlink(logfile) if __name__ == "__main__": - app.run(main) + app.run(main) From 6ab77a8af90f499bf5c50c95224b9386aa3f5f54 Mon Sep 17 00:00:00 2001 From: Arjun Mehta Date: Sun, 23 Aug 2026 06:35:28 +0000 Subject: [PATCH 7/8] Trigger CI re-run From a3352f80476527b1f66f9bbb0a6013b0736af32f Mon Sep 17 00:00:00 2001 From: Arjun Mehta Date: Sun, 23 Aug 2026 16:39:11 +0000 Subject: [PATCH 8/8] Fix pylint warnings in log reader benchmark --- compiler_opt/rl/log_reader_benchmark.py | 11 ++++++----- 1 file changed, 6 insertions(+), 5 deletions(-) diff --git a/compiler_opt/rl/log_reader_benchmark.py b/compiler_opt/rl/log_reader_benchmark.py index f12720d2..086eec11 100644 --- a/compiler_opt/rl/log_reader_benchmark.py +++ b/compiler_opt/rl/log_reader_benchmark.py @@ -127,6 +127,7 @@ def _stats( def main(_): + # pylint: disable=consider-using-with logfile = tempfile.NamedTemporaryFile(delete=False).name try: _write_log(logfile, _OBSERVATIONS.value, _ELEM_COUNT.value) @@ -138,10 +139,9 @@ def main(_): lambda: _parse_with(logfile, _add_feature_original), number=1)) current_times.append( timeit.timeit( - lambda: _parse_with(logfile, log_reader._add_feature), # pylint: disable=protected-access - number=1, - )) + lambda: _parse_with(logfile, log_reader._add_feature), number=1)) # pylint: disable=protected-access se_orig = _parse_with(logfile, _add_feature_original) + # pylint: disable=protected-access se_cur = _parse_with(logfile, log_reader._add_feature) serialized_bytes = None for key in se_orig: @@ -156,8 +156,9 @@ def main(_): f"p95={p95_orig:.3f}s 95% CI=({ci_orig[0]:.3f},{ci_orig[1]:.3f})") print(f" numpy: median={m_cur:.3f}s mean={mean_cur:.3f}s " f"p95={p95_cur:.3f}s 95% CI=({ci_cur[0]:.3f},{ci_cur[1]:.3f})") - print(f"speedup (median): {m_orig / m_cur:.2f}x " - f"({100 * (m_orig - m_cur) / m_orig:.1f}% faster)") + speedup_str = (f"speedup (median): {m_orig / m_cur:.2f}x " + f"({100 * (m_orig - m_cur) / m_orig:.1f}% faster)") + print(speedup_str) print(f"serialized output identical ({serialized_bytes} bytes per context)") finally: os.unlink(logfile)