rl/log_reader: use NumPy view in _add_feature by Arjunmehta312 · Pull Request #572 · google/ml-compiler-opt

Arjunmehta312 · 2026-08-06T08:37:38Z

Use the zero-copy NumPy view exposed by LogReaderTensorValue when populating protobuf repeated fields in _add_feature, making read_log_as_sequence_examples ~15% faster (~13-17% across runs, medians, disjoint 95% CIs) on representative synthetic workloads, with byte-for-byte identical serialized output. Adds compiler_opt/rl/log_reader_benchmark.py to reproduce the measurement.

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

mtrofin · 2026-08-17T18:52:40Z

"Benchmarks on representative synthetic workloads measured approximately a 15% improvement in read_log_as_sequence_examples" what are the representative synthetic workloads?

Arjunmehta312 · 2026-08-18T03:09:00Z

The benchmark used synthetic logs in the same "simple log format" that read_log_as_sequence_examples consumes: a JSON header followed by raw binary tensor buffers, mirroring the feature mix used by the supported problem configs (inlining, regalloc) — one float32 and one int64 feature tensor per observation, plus a float32 score tensor. Feature tensors were large (multi-million elements, multi-MB buffers) to reflect production logs.

Measurement: read_log_as_sequence_examples was run end-to-end on such a log, with the old and new _add_feature interleaved (25 samples each) and compared by median with 95% confidence intervals (disjoint). Depending on the run the improvement was ~13-17% (~15%). In every run the serialized SequenceExample output (21,971,863 bytes) was byte-for-byte identical between the two variants.

mtrofin · 2026-08-19T05:32:41Z

The benchmark used synthetic logs in the same "simple log format" that read_log_as_sequence_examples consumes: a JSON header followed by raw binary tensor buffers, mirroring the feature mix used by the supported problem configs (inlining, regalloc) — one float32 and one int64 feature tensor per observation, plus a float32 score tensor. Feature tensors were large (multi-million elements, multi-MB buffers) to reflect production logs.

Measurement: read_log_as_sequence_examples was run end-to-end on such a log, with the old and new _add_feature interleaved (25 samples each) and compared by median with 95% confidence intervals (disjoint). Depending on the run the improvement was ~13-17% (~15%). In every run the serialized SequenceExample output (21,971,863 bytes) was byte-for-byte identical between the two variants.

Can you contribute the benchmark?

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.

Arjunmehta312 · 2026-08-19T06:59:51Z

Yes, added compiler_opt/rl/log_reader_benchmark.py in commit 9f1ebf0.

It generates a synthetic log in the same simple log format (JSON header + raw tensor buffers, float32 and int64 feature tensors per observation, float32 score tensor), measures read_log_as_sequence_examples for the current _add_feature against a reference copy of the previous implementation in interleaved order, and reports median/mean/p95 with a 95% confidence interval. It also asserts the serialized output of the two variants is byte-for-byte identical.

Usage, e.g.:

python compiler_opt/rl/log_reader_benchmark.py --observations=8 --elems=2100000 --samples=10

The default parameters are close to the workload the ~15% figure was measured on; the size is configurable so it can be scaled down for quick checks. It is a standalone script (not collected by pytest) and does not add a CI dependency.

mtrofin

can you also clarify if you used AI for this patch, and either way, simplify the commit message. It doesn't need all the flowery headings, which just pull on the attention of a reader. It can just say it makes (which) performance improvement and adds a benchmark demonstrating it, plus the numbers. I'd be surprised if you need more than 2 sentences.

mtrofin · 2026-08-19T13:59:40Z

+
+from compiler_opt.rl import log_reader
+
+flags.DEFINE_integer("observations", 8, "Number of observations to log.")


can you use the _OBSERVATIONS = flags.DEFINE_integer(... style see e.g. compiler_opt/tools/generate_default_trace.py

Done - switched to the _OBSERVATIONS = flags.DEFINE_integer(...) style as in compiler_opt/tools/generate_default_trace.py (commit 1d5db98).

mtrofin · 2026-08-19T14:00:33Z

+FLAGS = flags.FLAGS
+
+
+def _write_log(fname: str, observations: int, elems: int) -> None:


elem_count or num_elem, otherwise it reads like these are the actual elements

Done - renamed to elem_count (commit 1d5db98).

mtrofin · 2026-08-19T14:04:11Z

+    lst.extend(value)
+
+
+def _time_parse(fname: str, add_feature) -> tuple[float, dict[str, tf.train.SequenceExample]]:


would timeit be more canonical?

Done - the interleaved measurements now use timeit (commit 1d5db98).

mtrofin · 2026-08-19T14:34:45Z

+    return elapsed, result
+
+
+def _stats(samples: list[float]) -> tuple[float, float, float, tuple[float, float]]:


the repo already depends on scipy, you can use that for e.g. p95, etc.

Done - stats now use numpy/scipy (np.percentile for p95, scipy.stats.sem + t.interval for the 95% CI) in commit 1d5db98.

mtrofin · 2026-08-19T14:40:55Z

+machine-dependent; results are reported as median, mean, p95 and a 95%
+confidence interval.
+"""
+


can you add a usage example as a comment - if the default flag values are fine, perfect (i.e. the goal is to get a reader confident in how to use this "as intended". They can figure out later if they want to adjust flags)

Done - added a usage example to the module docstring; the default flag values are the intended workload (commit 1d5db98).

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.

Arjunmehta312 · 2026-08-20T04:14:59Z

Thanks for the review.

Regarding AI: this patch was developed with the help of an LLM coding assistant (similar to Cursor); the change itself is small and was reviewed and validated before submission.

I've also simplified the message as suggested - it's now two sentences: the change makes read_log_as_sequence_examples ~15% faster (~13-17% across runs, medians with disjoint 95% CIs) and adds compiler_opt/rl/log_reader_benchmark.py to reproduce the measurement.

All benchmark comments are addressed in commit 1d5db98 (flag style, elem_count rename, timeit, numpy/scipy stats, usage example).

mtrofin · 2026-08-20T15:15:32Z

see the license header error, otherwise lgtm.

mtrofin reviewed Aug 19, 2026

View reviewed changes

Address review feedback in log reader benchmark

1d5db98

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.


		from compiler_opt.rl import log_reader

		flags.DEFINE_integer("observations", 8, "Number of observations to log.")

		FLAGS = flags.FLAGS


		def _write_log(fname: str, observations: int, elems: int) -> None:

		lst.extend(value)


		def _time_parse(fname: str, add_feature) -> tuple[float, dict[str, tf.train.SequenceExample]]:

		return elapsed, result


		def _stats(samples: list[float]) -> tuple[float, float, float, tuple[float, float]]:

Uh oh!

Conversation

Arjunmehta312 commented Aug 6, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

mtrofin commented Aug 17, 2026

Uh oh!

Arjunmehta312 commented Aug 18, 2026

Uh oh!

mtrofin commented Aug 19, 2026

Uh oh!

Arjunmehta312 commented Aug 19, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

mtrofin left a comment

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Arjunmehta312 commented Aug 20, 2026

Uh oh!

mtrofin commented Aug 20, 2026

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

2 participants

Arjunmehta312 commented Aug 6, 2026 •

edited

Loading

Arjunmehta312 commented Aug 19, 2026 •

edited

Loading