Manual reification on the JVM

In this post, we’re going to implement a small CLI utility for data pipelines over tabular data. We’re going to get the JVM to turn a pipeline like seqgen --stop 1000000 --field A then put '$A = $A * 1024' then stats1 --field A --accumulators min,sum into this hot loop:

; unrolling disabled where possible for clarity
        ↗  vpslld           zmm5, zmm4, 0xa   ; A * 1024
 1.16%  │  vpminsd          zmm0, zmm0, zmm5  ; min reduction
47.89%  │  vpaddd           zmm4, zmm1, zmm4  ; A = seqgen counter induction variable
 2.67%  │  vpaddd           zmm2, zmm2, zmm5  ; sum reduction
44.88%  │  lea              ebx, [rbx + 0x10]
        │  cmp              edi, ebx
        ╰  jg               loop

There is a lot of prior art in the JVM ecosystem, like Trino and Spark. These systems approach the problem by generating JVM bytecode or Java code at runtime, which is a lot of work and not easy to write. I’ve seen early proposals to add explicit runtime specialization to the JVM, and earlier Valhalla discussion of how specialization might be expressed, but my crystal ball says that’s still a decade away.

My goal with this project is to get there with normal-looking Java code. No custom bytecode generation to implement the logic, no embedded Java compilers, and no switching between precompiled kernels. Finding closed-form solutions or doing complex execution plan analysis is also out of scope. The focus is on the quality of the generated code for this Zen 4 AVX-512 target machine I’m on. I’ll first establish the reification primitives that make this possible, then use them to build a general batch protocol. Runtime respecialization and whole-pipeline fusion come later as extensions. Fusing stages together is what reaches the exact loop presented above, but it also goes past the comfort zone and runs into serious obstacles.

Background and motivation

I’ve been using Miller a lot lately. It’s an all-in-one replacement for awk, sed, cut, join, and sort for tabular data. Here’s a quick taste:

$ mlr --csv \
  # regex rename columns - from space to underscore
  rename -g -r ' ,_' then \
  # rename a specific column
  rename 'predicted_render_time,prt' then \
  # DSL for arbitrary transformations
  put '
    # $ references specific columns
    $rt = $render_end - $render_start;
    $delta = $prt - $rt;
  ' then \
  # univariate stats for each of the three columns listed
  stats1 -f prt,rt,delta -a p50,p75,p95 then \
  # the DSL has many useful functions
  # this script formats each column's nanos into millis
  put 'for (k, v in $*) { $[k] = fmtnum(v / 1e6, "%.2f") . " ms" }' \
  # input file comes last; don't ask me why
  "$HOME/kwin perf statistics HDMI-A-1.csv"

prt_p50,prt_p75,prt_p95,rt_p50,rt_p75,rt_p95,delta_p50,delta_p75,delta_p95
7.51 ms,9.72 ms,427.79 ms,2.07 ms,2.10 ms,2.23 ms,6.41 ms,8.31 ms,425.72 ms

You build small data pipelines: usually by reading from an input file, then transforming it with various verbs along the way, and finally outputting the results in your preferred format. It’s simple to start with and rich in capability.

I contributed a minor thing to Miller recently, which exposed me to the internals. I started thinking about Miller’s dynamic nature: at its core, every invocation can use a drastically different pipeline. The set of operations can vary, and so can the input schema.

Miller supports this by using dynamic dispatch through Go’s interfaces:

13
type RecordTransformer interface {
14
  Transform(
15
    inrecAndContext *types.RecordAndContext,
16
    outputRecordsAndContexts *[]*types.RecordAndContext, // list of *types.RecordAndContext
17
    inputDownstreamDoneChannel <-chan bool,
18
    outputDownstreamDoneChannel chan<- bool,
19
  )
20
}

Each record, or row of data, is itself a heterogeneous hash map with the data for each cell ultimately stored in an interface{}:

54
type Mlrval struct {
55
  printrep      string
56
  intf          interface{}
57
  err           error // Payload for MT_ERROR types
58
  printrepValid bool
59
  // Enumeration for string / int / float / boolean / etc.
60
  // I would call this "type" not "mvtype" but "type" is a keyword in Go.
61
  mvtype MVType
62
}

And invoking the per-record transformer adds another indirect function call:

242
func runSingleTransformerBatch(
243
  inputRecordsAndContexts []*types.RecordAndContext, // list of types.RecordAndContext
244
  recordTransformer RecordTransformer,
245
  isFirstInChain bool,
246
  outputRecordChannel chan<- []*types.RecordAndContext, // list of *types.RecordAndContext
247
  inputDownstreamDoneChannel <-chan bool,
248
  outputDownstreamDoneChannel chan<- bool,
249
  dataProcessingErrorChannel chan<- error,
250
  options *cli.TOptions,
251
) (bool, error) {
252
  outputRecordsAndContexts := make([]*types.RecordAndContext, 0, len(inputRecordsAndContexts))
253
  done := false
254

255
  for _, inputRecordAndContext := range inputRecordsAndContexts {
27 collapsed lines
256
    // --nr-progress-mod
257
    // TODO: function-pointer this away to reduce instruction count in the
258
    // normal case which it isn't used at all. No need to test if {static thing} != 0
259
    // on every record.
260
    if options.NRProgressMod != 0 {
261
      if isFirstInChain && inputRecordAndContext.Record != nil {
262
        context := &inputRecordAndContext.Context
263
        if context.NR%options.NRProgressMod == 0 {
264
          fmt.Fprintf(os.Stderr, "NR=%d FNR=%d FILENAME=%s\n", context.NR, context.FNR, context.FILENAME)
265
        }
266
      }
267
    }
268

269
    // Three things can come through:
270
    //
271
    // * End-of-stream marker
272
    // * Non-nil records to be printed
273
    // * Strings to be printed from put/filter DSL print/dump/etc
274
    //   statements. They are handled here rather than fmt.Println directly
275
    //   in the put/filter handlers since we want all print statements and
276
    //   record-output to be in the same goroutine, for deterministic
277
    //   output ordering.
278
    //
279
    // The first two are passed to the transformer. The third we send along
280
    // the output channel without involving the record-transformer, since
281
    // there is no record to be transformed.
282

283
    if inputRecordAndContext.EndOfStream || inputRecordAndContext.Record != nil {
284

285
      err := recordTransformer.Transform(
286
        inputRecordAndContext,
287
        &outputRecordsAndContexts,
288
        inputDownstreamDoneChannel,
289
        outputDownstreamDoneChannel,
290
      )

Records are passed over channels, in batches, and every transformer is a separate goroutine. Batch size is entirely dependent on individual transformers. For example, seqgen, which is a verb that generates an incrementing series, materializes all results in a single batch passed forward in a single channel element:

$ systemd-run --user --pty --wait --collect -p MemoryMax=20G -p MemorySwapMax=0 \
  mlr \
    seqgen -f counter --start 1 --stop 100000000 then \
    stats1 -f counter -a min
Running as unit: run-p1264700-i1286487.service
Press ^] three times within 1s to disconnect TTY.
          Finished with result: oom-kill                  
Main processes terminated with: code=killed, status=9/KILL
               Service runtime: 5.586ms
             CPU time consumed: 7.315s
                   Memory peak: 20G (swap: 0B)

Still, even if batch sizes were controlled, there’s no good way of predicting the size of each row. Knowing the schema ahead of time and planning batch sizes to have them fit in L3 cache (to maximize bandwidth in the downstream goroutine) would be worthwhile. Better yet, it would probably be faster to avoid the goroutine and stay in L1 cache by using smaller batches. This is usually called cache blocking, or tiling.

For a statically compiled program, it would be impossible to specialize the operations to all potential schemas the user might throw at it at runtime. Polars comes to mind and it does attempt exactly that by bundling many kernels and then dispatching to them at runtime, depending on the exact query plan.

# just a small benchmark wrapper to invoke Polars from the CLI and print its plan too
$ POLARS_MAX_THREADS=1 ./polars '(
    pl.LazyFrame()
      .select(pl.int_range(0, 1000000, dtype=pl.Int32).alias("A"))
      .with_columns(A=pl.col("A") * 2)
      .min()
 )'

<~~ plan [in-memory] ~~>
SELECT [col("A").min()]
   WITH_COLUMNS:
   [[(col("A")) * (2)]]
    SELECT [0.int_range([1000000]).alias("A")]
      DF []; PROJECT */0 COLUMNS
benchmark results over 100 runs (20 warmup)
polars collect [engine=in-memory]: min 0.170  p25 0.172  p50 0.173  p75 0.178  max 0.532 ms
polars collect [engine=streaming]: min 0.256  p25 0.260  p50 0.264  p75 0.269  max 0.357 ms

$ POLARS_MAX_THREADS=1 ./polars '(
    pl.LazyFrame()
      .select(pl.int_range(0, 1000000, dtype=pl.Int32).alias("A"))
      .with_columns(A=pl.col("A") * 2 * 3)
      .min()
 )'
<~~ plan [in-memory] ~~>
SELECT [col("A").min()]
   WITH_COLUMNS:
   [[([(col("A")) * (2)]) * (3)]]
    SELECT [0.int_range([1000000]).alias("A")]
      DF []; PROJECT */0 COLUMNS
benchmark results over 100 runs (20 warmup)
polars collect [engine=in-memory]: min 0.228  p25 0.230  p50 0.232  p75 0.234  max 0.254 ms
polars collect [engine=streaming]: min 0.284  p25 0.295  p50 0.302  p75 0.331  max 1.322 ms

A profiling run reveals key insights:

$ perf annotate --stdio --disassembler-style=intel --stdio-color always
# ...
    0.00 :   1d12610:        vmovdqu ymm2,YMMWORD PTR [rdi+r8*4]
    2.45 :   1d12616:        vmovdqu ymm3,YMMWORD PTR [rdi+r8*4+0x20]
    0.00 :   1d1261d:        vmovdqu ymm4,YMMWORD PTR [rdi+r8*4+0x40]
    5.30 :   1d12624:        vmovdqu ymm5,YMMWORD PTR [rdi+r8*4+0x60]
    0.00 :   1d1262b:        vpsllvd ymm2,ymm2,ymm1
    0.00 :   1d12630:        vpsllvd ymm3,ymm3,ymm1
    8.57 :   1d12635:        vpsllvd ymm4,ymm4,ymm1
    8.12 :   1d1263a:        vpsllvd ymm5,ymm5,ymm1
    2.58 :   1d1263f:        vmovdqu YMMWORD PTR [rsi+r8*4],ymm2
    6.34 :   1d12645:        vmovdqu YMMWORD PTR [rsi+r8*4+0x20],ymm3
   20.17 :   1d1264c:        vmovdqu YMMWORD PTR [rsi+r8*4+0x40],ymm4
   28.16 :   1d12653:        vmovdqu YMMWORD PTR [rsi+r8*4+0x60],ymm5
   17.15 :   1d1265a:        add    r8,0x20
    1.15 :   1d1265e:        cmp    rax,r8
    0.00 :   1d12661:        jne    1d12610 <polars_compute::arity::ptr_apply_unary_kernel::<i32, i32, <i32 as polars_compute::arithmetic::PrimitiveArithmeticKernelImpl>::prim_wrapping_mul_scalar::{closure#0}>+0x50>

• Polars is compiled for AVX2 here and doesn’t bundle a wider kernel, so it can’t make the best use of my CPU.
• In this run, the input and output pointers differ, so the kernel copies to a new location instead of transforming values in place.
• It is manually specialized by the Polars developers to detect a power-of-two scalar and strength-reduce the multiplication into a left shift. That’s pretty impressive, but they don’t decompose code like x * 5 into (x << 2) + x. Your production compiler would do that for static code.
• Polars launches two mult kernels back to back and doesn’t fuse operations.
• Each kernel seems to operate on entire columns at a time — no tiling is attempted to stay in CPU caches and reuse operands between back-to-back kernels.

Reification primitives

Before we can implement our own data pipeline, we need to figure out how to hit the JIT fast paths. Since the processing is controlled through runtime CLI arguments, we need to take care to assemble the computation so it optimizes as well as static code. Let’s start small by building a very simple interpreter. Something that will eventually implement a DSL similar to what Miller supports in put "$out = $in * 5".

@BenchmarkMode(Mode.AverageTime)
@OutputTimeUnit(TimeUnit.NANOSECONDS)
@OperationsPerInvocation(NaiveInterpreterBenchmark.ROWS)
public class NaiveInterpreterBenchmark {
  static final int ROWS = 1024;

  @State(Scope.Benchmark)
  public static class BenchState {
    int[] data = IntStream.range(0, ROWS).toArray();
    Expr interpreter = new ExprAdd(new ExprConst(2), new ExprRead());
  }

  // static, handwritten implementation:
  @Benchmark
  public void direct(BenchState state) {
    final int[] data = state.data;
    for (int i = 0; i < data.length; i++) {
      data[i] = data[i] + 2;
    }
  }

  // data-driven interpreter:
  public static interface Expr {
    public int eval(int input);
  }

  public static final record ExprConst(int val) implements Expr {
    @Override
    public int eval(int input) { return val; }
  }

  public static final record ExprAdd(Expr lhs, Expr rhs) implements Expr {
    @Override
    public int eval(int input) { return lhs.eval(input) + rhs.eval(input); }
  }

  public static final record ExprRead() implements Expr {
    @Override
    public int eval(int input) { return input; }
  }

  @Benchmark
  public void interpreted(BenchState state) {
    final int[] data = state.data;
    final var interpreter = state.interpreter;
    for (int i = 0; i < data.length; i++) {
      data[i] = interpreter.eval(data[i]);
    }
  }
}

Both of these optimize to the same hot code on OpenJDK 26:

Benchmark                              Mode  Cnt  Score    Error  Units
NaiveInterpreterBenchmark.direct       avgt    9  0.025 ±  0.001  ns/op
NaiveInterpreterBenchmark.interpreted  avgt    9  0.025 ±  0.001  ns/op

          ↗  mov         ecx, edx
          │  vpaddd      zmm1, zmm0, zmmword ptr [rsi + rcx*4 + 0xc]
          │  vmovdqu32   zmmword ptr [rsi + rcx*4 + 0xc], zmm1
  18.69%  │  vpaddd      zmm1, zmm0, zmmword ptr [rsi + rcx*4 + 0x4c]
   0.02%  │  vmovdqu32   zmmword ptr [rsi + rcx*4 + 0x4c], zmm1
  23.28%  │  vpaddd      zmm1, zmm0, zmmword ptr [rsi + rcx*4 + 0x8c]
          │  vmovdqu32   zmmword ptr [rsi + rcx*4 + 0x8c], zmm1
  20.86%  │  vpaddd      zmm1, zmm0, zmmword ptr [rsi + rcx*4 + 0xcc]
          │  vmovdqu32   zmmword ptr [rsi + rcx*4 + 0xcc], zmm1
  20.94%  │  lea         edx, [rcx + 0x40]
          │  cmp         edx, r11d
          ╰  jl          loop

Let’s try something just a bit more complicated:

public class NaiveInterpreterBenchmark {
  public static class BenchState {
    // ...

    Expr interpreter2 = new ExprAdd(
      new ExprConst(2),
      new ExprAdd(new ExprConst(3), new ExprRead())
    );
  }

  @Benchmark
  public void direct2(BenchState state) {
    final int[] data = state.data;
    for (int i = 0; i < data.length; i++) {
      data[i] = data[i] + 2 + 3;
    }
  }

  @Benchmark
  public void interpreted2(BenchState state) {
5 collapsed lines
    final int[] data = state.data;
    final var interpreter = state.interpreter2;
    for (int i = 0; i < data.length; i++) {
      data[i] = interpreter.eval(data[i]);
    }
  }
}

Benchmark                               Mode  Cnt  Score    Error  Units
NaiveInterpreterBenchmark.direct2       avgt    9  0.025 ±  0.001  ns/op
NaiveInterpreterBenchmark.interpreted2  avgt    9  1.211 ±  0.037  ns/op

The interpreter performance immediately falls off a cliff! Disassembly shows that vectorization is no longer present, and the scalar code is full of branches for type checks and other nastiness. I’ll spare you from having to look at it. Instead, let’s see the output of -XX:CompileCommand="print *ExprAdd::eval":

dev.miller4j.NaiveInterpreterBenchmark$ExprAdd::eval(I)I
  interpreter_invocation_count:      695088
  invocation_counter:                695088
  backedge_counter:                       0
  decompile_count:                        0
  mdo size: 560 bytes

   0 fast_aaccess_0
   1 fast_agetfield 7 <dev/miller4j/NaiveInterpreterBenchmark$ExprAdd.lhs:Ldev/miller4j/NaiveInterpreterBenchmark$Expr;>
   4 iload_1
   5 invokeinterface 16 <dev/miller4j/NaiveInterpreterBenchmark$Expr.eval(I)I>
  0    bci: 5    VirtualCallData    count(0) entries(1)
                                    'dev/miller4j/NaiveInterpreterBenchmark$ExprConst'(694577 1.00)
  10 fast_aaccess_0
  11 fast_agetfield 13 <dev/miller4j/NaiveInterpreterBenchmark$ExprAdd.rhs:Ldev/miller4j/NaiveInterpreterBenchmark$Expr;>
  14 iload_1
  15 invokeinterface 16 <dev/miller4j/NaiveInterpreterBenchmark$Expr.eval(I)I>
  144  bci: 15   VirtualCallData    count(0) entries(2)
                                    'dev/miller4j/NaiveInterpreterBenchmark$ExprRead'(347289 0.50)
                                    'dev/miller4j/NaiveInterpreterBenchmark$ExprAdd'(347288 0.50)
  20 iadd
  21 ireturn

Profiling data for ExprAdd::eval confirms that rhs.eval gets called with two different receivers — the call site is bimorphic, with dispatch split evenly between ExprAdd::eval and ExprRead::eval. This doesn’t necessarily prevent inlining, but introducing control flow prevents the loop from being vectorized.

We can recover the lost performance by storing the interpreter in a static final field:

  public static class BenchState {
    // ...

    Expr interpreter2 = new ExprAdd(
      new ExprConst(2),
      new ExprAdd(new ExprConst(3), new ExprRead())
    );

    static final Expr interpreter3 =
      new ExprAdd(new ExprConst(2), new ExprAdd(new ExprConst(3), new ExprRead()));
  }

  @Benchmark
  public void interpreted3(BenchState state) {
5 collapsed lines
    final int[] data = state.data;
    final var interpreter = BenchState.interpreter3;
    for (int i = 0; i < data.length; i++) {
      data[i] = interpreter.eval(data[i]);
    }
  }

Although the call site has the same VirtualCallData as before, the loop vectorizes again to reach peak performance:

Benchmark                               Mode  Cnt  Score    Error  Units
NaiveInterpreterBenchmark.interpreted3  avgt    9  0.025 ±  0.001  ns/op

When the receivers are provably constant — and they are when we store a tree of immutable Records in a static final field — the compiler can fully inline and devirtualize the call graph. There is just one caveat in that optimization decisions are specific to a call site at the bytecode level. So, when we complicate the example a bit more:

  public static class BenchState {
    // ...

    // unmodifiable List, with immutability enforced and understood by the VM
    static final List<Expr> interpreters = List.of(
      new ExprAdd(new ExprConst(2), new ExprAdd(new ExprConst(3), new ExprRead())),
      // same result, flipped operands:
      new ExprAdd(new ExprAdd(new ExprConst(3), new ExprRead()), new ExprConst(2))
    );
  }

  @Benchmark
  public void twoInterpreters(BenchState state) {
    final int[] data = state.data;
    final var interpreters = BenchState.interpreters;
    for (int i = 0; i < data.length; i++) {
      for (int j = 0; j < interpreters.size(); j++) {

        data[i] = interpreters.get(j).eval(data[i]);
      }
    }
  }

  @Benchmark
  public void twoInterpretersSplit(BenchState state) {
    // ...
    for (int j = 0; j < interpreters.size(); j++) {

      if (j == 0)
        data[i] = interpreters.get(j).eval(data[i]);
      if (j == 1) // for reasons unknown, this cannot be an else if
        data[i] = interpreters.get(j).eval(data[i]);
    }
  }

Both list elements have the same receiver class, ExprAdd, so the outer call site in @Benchmark twoInterpreters is monomorphic. The polymorphism appears deeper, in ExprAdd::eval — its lhs.eval and rhs.eval call sites see different receiver classes. Splitting moves each call into the constant-receiver regime, where inlining follows the object graph instead of the profile, and a profile only records receiver classes, so it can never tell our two ExprAdd instances apart. Our split call site variant still routes to a specific, constant object graph and the compiler can inline through it:

Benchmark                                       Mode  Cnt  Score    Error  Units
NaiveInterpreterBenchmark.twoInterpreters       avgt    9  9.975 ±  0.918  ns/op
NaiveInterpreterBenchmark.twoInterpretersSplit  avgt    9  0.025 ±  0.001  ns/op

The takeaway is that invoking interface methods such as Expr::eval is acceptable, so long as the receiver is constant enough for the VM. That means starting from a static final field and tracing the receiver all the way through immutable fields.

static final dynamic

All the examples we’ve seen so far used static code. The exact data structure or code had to have been written directly or produced in a class initializer and written to a static final field. This is obviously not enough to optimize the enormous surface area formed by CLI arguments for the real mlr command. We can’t jump this hoop with Java-the-language, but Java-the-virtual-machine can load new code at runtime, and that code can contain static finals. It’s not much different from how our compiled .java code gets loaded. To close the loop, we need to generate a class file at runtime and load it into the VM instance we are running on.

I explicitly did not want to hand-write bytecode for pipeline stages. Instead, an annotation processor turns normal Java functions into reification factories, in a generic way: annotating a method in Foo gets you a companion Foo_Reified holding a factory for it. At runtime, those factories use Cojen Maker to generate thin binding classes: they hoist runtime values into static final fields and materialize fresh call sites that invoke my existing Java methods. No logic is ever emitted as bytecode; it stays in the Java I wrote.

public class NaiveInterpreterBenchmark {

  @Reified
  // @Constant arguments become `static final` fields in the generated class
  static int reifyExpr(@Constant Expr expr, int input) {
    return expr.eval(input);
  }

  public static class BenchState {
    // from earlier
    Expr interpreter2 = new ExprAdd(
      new ExprConst(2),
      new ExprAdd(new ExprConst(3), new ExprRead())
    );

    // This generates and loads a new class into the current JVM, giving us an
    // instance handle that implements the interface.
    NaiveInterpreterBenchmark_Reified.ReifyExpr interpreter2Reified =
      // @Constant arguments are passed in at creation time:
      NaiveInterpreterBenchmark_Reified.reifyExpr(interpreter2);
  }

  // ...

  @Benchmark
  public void reified2(BenchState state) {
    final int[] data = state.data;
    final var interpreter = state.interpreter2Reified;
    for (int i = 0; i < data.length; i++) {
      // only non-@Constant arguments are supplied in the invocation
      data[i] = interpreter.apply(data[i]);
    }
  }
}

One of those runtime binding classes decompiles to this:

public final class GeneratedReifiedStatic44
    implements NaiveInterpreterBenchmark_Reified.ReifyExpr {
  private static final NaiveInterpreterBenchmark.Expr BOUND86 =
      (NaiveInterpreterBenchmark.Expr) Specializers.boundReference(86);

  @Override
  public int apply(int n) {
    // The annotation processor also generates a companion CompileCommand directive.
    // We nudge the compiler to inline this function and the call inside so it can
    // reliably fold the constant.
    //
    // inline dev.miller4j.NaiveInterpreterBenchmark::reifyExpr
    // inline dev.miller4j.GeneratedReifiedStatic*::apply
    return NaiveInterpreterBenchmark.reifyExpr(BOUND86, n);
  }
}

The generated bytecode only binds constants and creates the call site. This reification approach is different from monomorphization in static languages — we’re only duplicating the glue that calls our shared functions. Of course, once these functions get optimized to native code, we end up with a similar result: every instance is specialized and optimized independently. The main logic stays in reifyExpr, typechecks with the rest of the program, and therefore feels like ordinary Java. There’s not a separate stage of compilation that would defer errors until runtime. We’re not generating complicated blobs of bytecode and depending on the JVM’s verifier to surface issues at load time.

Code from the new class can link to (read: use) anything else that we’ve already loaded in. To pass data into the new class, we make it reference a simple global registry that stores arbitrary objects:

static final AtomicInteger REFERENCE_IDS = new AtomicInteger();
static final Object NULL_REFERENCE = new Object();
static final ConcurrentMap<Integer, Object> REFERENCES = new ConcurrentHashMap<>();

static int registerReference(@Nullable Object value) {
  final int id = REFERENCE_IDS.incrementAndGet();
  REFERENCES.put(id, value == null ? NULL_REFERENCE : value);
  return id;
}

static @Nullable Object boundReference(int id) {
  final Object value = REFERENCES.remove(id);
  if (value == null) throw new IllegalStateException("Missing bound reference: " + id);
  return value == NULL_REFERENCE ? null : value;
}

// later, when generating the binding class with Cojen Maker:
classMaker.addField(type, name).private_().static_().final_();
final var clinit = classMaker.addClinit();
clinit
    .field(name)
    .set(clinit.var(Specializers.class).invoke("boundReference", referenceId).cast(type));

This approach regains all of the performance without requiring a source-level static final field for every runtime expression:

Benchmark                               Mode  Cnt  Score    Error  Units
NaiveInterpreterBenchmark.interpreted2  avgt    9  1.227 ±  0.010  ns/op
NaiveInterpreterBenchmark.reified2      avgt    9  0.025 ±  0.001  ns/op

We’ll have plenty of @Constant candidates just in the DSL script alone, but generally speaking, every parameter from CLI verbs makes for great initial candidate. There are --start and --stop from the seqgen verb, and -n from head and tail. When we analyze the schema along the entire pipeline of verbs, that’s worth making constant too. We’ll even inline pointers to our memory segments into the instruction stream. But we have some more primitives to walk through first before we get ahead of ourselves.

Loops and Object[] registerFiles

We’ve looked at the importance of separating call sites, keeping profiling data isolated, and ideally make the entire call tree go through static final receivers. Composing computations into a bigger hierarchy with @Reified enables us to specialize what otherwise would be branches or dynamic dispatch into straight-line code instead. Let’s now do the same for loops, which are the other source of shared call sites that we need a separate solution for.

Speaking of composition, we need a good way to share state between the different functions we reify. Ideally, we’d store our data in local variables — specifically primitive types whenever possible — but since we’re not emitting Java code nor bytecode functions, this is not an option for us. JVM functions can only return up to one value, and anything that is passed as an argument gets copied. So how can we represent pipeline state and all the data for every column in the schema to be able to pass that between miller verbs?

The answer I came up with is using an Object[] registerFile on the semantic level while ensuring its elements get scalar-replaced by the compiler.

This is only possible with the Graal JIT compiler and not C2, so the rest of the post will use that unless noted otherwise. There are a couple basic requirements to make scalar-replacement reliable:

1. Use constant indices with this array. We’ll often use @Constant int slotIx coupled with registerFile[slotIx] = ....
2. Never mix or confuse types of a given slot. The patterns that work well are: (T) registerFile[ix] to read, registerFile[ix] = (T) value to store. This way, we can imagine having an @Reified function which simultaneously accepts @Constant int myStorageIndex and the Object[] state.
3. Initialize all the slots you need at the start, before first use.
4. (situational) Zero-out slots after last use — this will become relevant much later.

To show this off, let’s construct something that looks a bit more like the stats1 function from Miller. It’ll need to store the state of its accumulators calculate the requested reductions like min, max, sum, for every row of data it ingests.

@OperationsPerInvocation(4096)
public class PipelineBenchmark {

  @FunctionalInterface
  interface Processor {
    void apply(Object[] registerFile, int value);
  }

  // `extending` selects a specific interface instead of an auto-generated one
  @Reified(extending = Processor.class)
  public static void process(
      @Constant int registerIx,
      @Constant String reduction,
      Object[] registerFile,
      int value) {
    int accumulator = (int) registerFile[registerIx];
    accumulator =
        switch (reduction) {
          case "min" -> Math.min(accumulator, value);
          case "max" -> Math.max(accumulator, value);
          case "sum" -> accumulator + value;
          default -> throw new UnsupportedOperationException("unknown reduction");
        };
    registerFile[registerIx] = accumulator;
  }

  // This annotation is a little more involved: it accepts a list of operations
  // and unrolls to give each operation its own call site.
  // Runtime arguments are broadcast to each step and we use Object[] for its
  // internal mutability.
  //
  // Returns a SAM interface implementation with an `invoke` method accepting
  // the runtime arguments.
  @ReifiedLoop(step = Processor.class)
  static void pipelineSum(Object[] registerFile, int value) {}

  @State(Scope.Benchmark)
  public static class BenchState {
    int[] data = IntStream.range(0, 4096).toArray();

    List<Processor> reducers = List.of(
      PipelineBenchmark_Reified.process(0, "sum"),
      PipelineBenchmark_Reified.process(1, "min"),
      PipelineBenchmark_Reified.process(2, "max")
    );

    // Lambdas are more concise and can optimize just as well, but depend on
    // normal inlining heuristics. @Reified produces force-inlined call sites
    // that bypass these limits.
    PipelineBenchmark_Reified.PipelineSum pipelineReified =
        PipelineBenchmark_Reified.pipelineSum(reducers);
  }

  @Benchmark
  public long handwritten(BenchState state) {
10 collapsed lines
    final var arr = state.data;
    int sum = 0;
    int min = Integer.MAX_VALUE;
    int max = Integer.MIN_VALUE;
    for (final var el : arr) {
      sum = sum + el;
      min = Math.min(min, el);
      max = Math.max(max, el);
    }
    return sum + min + max;
  }

  @Benchmark
  public long reified(BenchState state) {
    final var arr = state.data;
    final var pipeline = state.pipelineReified;
    // initial state for each accumulator: sum, min, max
    Object[] registerFile = new Object[] {0, Integer.MAX_VALUE, Integer.MIN_VALUE};
    for (final var el : arr) {
      pipeline.invoke(registerFile, el);
    }
    return (int) registerFile[0] + (int) registerFile[1] + (int) registerFile[2];
  }
}

Benchmark                      Mode  Cnt  Score   Error  Units
PipelineBenchmark.handwritten  avgt    9  0.462 ± 0.005  ns/op
PipelineBenchmark.reified      avgt    9  0.461 ± 0.002  ns/op

The assembly output is equivalent for the two implementations and looks like this:

mov           eax, edx
mov           edi, r10d
mov           edx, r14d
add           edx, r13d
mov           edx, dword ptr [r8 + rdx*4 + 0xc] ; load element
add           ebp, edx                          ; sum reducer
cmp           edx, edi
mov           r10d, edx
cmovle        r10d, edi                         ; max reducer
cmp           edx, eax
cmovge        edx, eax                          ; min reducer
inc           r14d
cmp           r11d, r14d
jg            loop                              ; back to top

Notably, switch (reduction) and other operations on @Constant values have been partially evaluated when the function was compiled. The Object[] and its boxed accumulators are gone as well — all three accumulators live in registers for the whole loop. This is a tough register file. One of the most allocation-free we’ve ever seen from the standpoint of heap memory.

The generated class decompiles to this:

public final class GeneratedReifiedLoop4
implements PipelineBenchmark_Reified.PipelineSum {
    private static final PipelineBenchmark.Processor BOUND1 = (PipelineBenchmark.Processor)Specializers.boundReference((int)1);
    private static final PipelineBenchmark.Processor BOUND2 = (PipelineBenchmark.Processor)Specializers.boundReference((int)2);
    private static final PipelineBenchmark.Processor BOUND3 = (PipelineBenchmark.Processor)Specializers.boundReference((int)3);

    public void invoke(Object[] objectArray, int n) {
        BOUND1.apply(objectArray, n);
        BOUND2.apply(objectArray, n);
        BOUND3.apply(objectArray, n);
    }
}

Each step becomes its own static final field and gets a dedicated call site, allowing for peak optimization. The annotation processor emits method-specific CompileCommand directives. Together with a handful of static rules, we get:

inline dev.miller4j.GeneratedReifiedLoop*::invoke
inline dev.miller4j.PipelineBenchmark::pipelineSum

Forced inlining is critical to get scalar-replacement of Object[]. If we miss a single function, or enter a safepoint poll, or throw an exception when this object is still live, the compiler may bail out of the optimization.

Implementing the Miller pipeline

We’re now ready to implement the pipeline abstraction that assembles reified steps from CLI arguments and executes them end-to-end. To recap the driving principle: reification lets us lift the point of dynamic dispatch higher in the call tree, and keep everything below that statically dispatched. This unlocks inlining and scalarization of values we pass around, which in turn unlocks a whole host of further optimizations. Reification will be used to generate an executor loop with separate call sites for each stage of the pipeline, with these being maximally specialized with @Constant values read from the CLI or otherwise derived during planning. We’re only implementing a small subset of Miller verbs, and skipping file inputs for now. The overall architecture looks something like this:

1. Pipeline stages are assembled from CLI arguments,
2. A planning pre-pass calls into each stage to perform schema evolution and gathers required allocations from each stage. Each stage can express:
   1. what columns it wants to add/drop from the schema
   2. what private scratch data it needs to allocate — it will be preserved for the entire pipeline run
3. The final allocation and batch size are calculated,
4. Every stage gets reified into an instance of BatchOp, with its configuration, schema, and memory allocations available as @Constants,
5. A list of BatchOps gets reified into a driver loop,
6. The driver executes in a loop, with batches of rows passed forward between stages,
7. Execution finishes when the last stage returns an EOF marker.

The iteration protocol is batch-oriented. Every stage can receive a small number of rows (on the order of hundreds to a couple thousand), transform them, and emit the same or a different number of rows to a downstream step. The control value describes row count of the batch and an EOF bit — I did not want to rely on scalarization of a compound object, so this is a bit-packed integer. When a stage reports EOF early, the pipeline runner will skip past that stage in its main loop. The run ends after the last stage reports completion.

This is the public API I came up with for each stage to plug into:

interface PipelineStage {
  void plan(Planner planner, Schema inputSchema);

  BatchOp compileBatchOp(
    // Schema snapshot of the upstream step
    BatchLayout in,
    // Schema snapshot of what this step produces
    BatchLayout out,
    // Memory used to pass the input & output data for a batch
    MemorySegment data,
    // Scratch memory that persists throughout the full pipeline run
    MemorySegment scratch
  );

  // Reset mutable state kept by a compiled source stage
  default void resetCompiledBatch() {}

  // Initialize scratch memory, if needed
  default void resetScratch(BatchLayout layout, MemorySegment scratch) {}
}

@FunctionalInterface
interface BatchOp {
  // The control word encodes the number of rows being passed in a batch,
  // as well as bit flags like EOF.
  int apply(int control);
}

final class Planner {
  void addColumn(String name, ColumnType ty);
  void addScratch(String name, MemoryLayout layout);

  // Sets an optional hint which helps the planner's automatic sizing
  void setOutputRowCapacity(int maxRowCount);

  // Meant for transformations that can't be done in place: repoints a column
  // at freshly allocated storage so the prior (input) buffer and the new
  // (output) buffer don't alias. The previous allocation gets queued for
  // deallocation after this step.
  void dupStorage(String name);

  void dropColumn(String name);

  // Drops all columns to start from a clean slate.
  void reset();
}

final record BatchLayout(...) {
  int maxRowCount();

  // Returns a slice of the `base` segment for the given column
  MemorySegment resolveColumn(MemorySegment base, String name);

  // Returns a slice of the `base` segment for the given scratch allocation
  MemorySegment resolveScratch(MemorySegment base, String name);
}

Generating batches with seqgen

seqgen is the pipeline’s source. It creates batches containing one integer column, filled with consecutive values from start through stop. A simplified implementation of seqgen --start 1 --stop 500000000 looks like this:

final class SeqGen implements PipelineStage {
16 collapsed lines
  private static final VectorSpecies<Integer> SPECIES = IntVector.SPECIES_PREFERRED;
  private static final IntVector LANE_OFFSETS = IntVector.zero(SPECIES).addIndex(1);

  final String name;
  final int start;
  final int stop;
  int current;

  SeqGen(String name, int start, int stop) {
    if (start > stop) throw new IllegalArgumentException("start must not exceed stop");
    this.name = name;
    this.start = start;
    this.stop = stop;
    this.current = start;
  }

  @Override
  public void plan(Planner planner, Schema ignoredInput) {
    // This is a source: discard the input schema and create its output column.
    planner.reset();
    planner.addColumn(name, new ColumnInt());
  }

  @Override
  public BatchOp compileBatchOp(
      BatchLayout in, BatchLayout out, MemorySegment data, MemorySegment scratch) {
    // Resolve the column once, while constructing this specialized operator.
    final var output = out.resolveColumn(data, name);
    return SeqGen_Reified.batchLocal(stop, out.maxRowCount(), output, this);
  }

  @Reified(extending = BatchOp.class)
  static int batchLocal(
      @Constant int stop,
      @Constant int capacity,
      @Constant MemorySegment output,
      @Constant SeqGen state,
      int ignoredControl) {
    final int current = state.current;
    final int count = Math.min(Math.max(0, stop - current + 1), capacity);

    // Build [current, current + 1, ...] once, then advance a whole vector at a time.
    var values = LANE_OFFSETS.add(current);
    final int fullCount = SPECIES.loopBound(count);
    int i = 0;
    for (; i < fullCount; i += SPECIES.length()) {
      values.intoMemorySegment(
          output, (long) i * ValueLayout.JAVA_INT.byteSize(), ByteOrder.nativeOrder());
      values = values.add(SPECIES.length());
    }

    // The final batch can have a partial vector.
    if (i < count) {
      final var mask = SPECIES.indexInRange(i, count);
      values.intoMemorySegment(
          output, (long) i * ValueLayout.JAVA_INT.byteSize(), ByteOrder.nativeOrder(), mask);
    }

    state.current += count;
    return control(

      state.current > stop,

      count
    );
  }

  @Override
  public void resetCompiledBatch() {
    current = start;
  }
}

The SeqGen reference is constant in the compiled operator, but its current field remains mutable between calls. In a real pipeline, the full-vector hot path compiles to the following code:

movabs r11, 0x790ad5000000      ; output column's native address
movabs r8,  0x711b15c00         ; output MemorySegment's session
movabs r9,  0x7118a2140         ; this SeqGen instance
vmovdqu32 zmm0, [rip-0x113]     ; constant [16, 17, ... 31]
vmovdqu32 zmm1, [rip-0xdd]      ; constant [0, 1, ... 15]
vmovdqu32 zmm2, [rip-0xc5]      ; sixteen copies of 16
; bounds checks omitted
vpbroadcastd zmm3, edx          ; [current, current, ...]
vpaddd       zmm4, zmm3, zmm1   ; [current, current + 1, ...]

test        eax, 0xfffffff0
je          partial_vector_path ; code omitted

; MemorySegment liveness check hoisted before the loop
; as of this writing, needs a GraalVM EA build that includes this fix:
; https://github.com/oracle/graal/pull/13917
mov         r13d, dword ptr [r8+8]
test        r13d, r13d
jl          closed_session

; first iteration got peeled
vmovdqu32   [r11], zmm4
vpaddd      zmm3, zmm3, zmm0 ; [current + 16, ... current + 31]
mov         ecx, 0x10

vector_loop:
vmovdqu32   [r11+rcx*4], zmm3
vpaddd      zmm3, zmm3, zmm2
lea         ecx, [rcx+0x10]
cmp         ebp, ecx
ja          vector_loop

mov    ecx, eax
add    ecx, edx
mov    dword ptr [r9+0x10], ecx ; state.current += count
; control word packing omitted

On my AVX-512 CPU, each loop iteration emits 16 integers with one 64-byte store and advances all lanes with one vpaddd. The MemorySegment addresses, SeqGen object reference, and configuration parameters like --stop all became immediate values in the compiled instruction stream. Graal hoisted both the spatial bounds check and the temporal MemorySegment liveness check out of the full-vector loop. The hot loop itself is as good as it gets: just the store, vector addition, index increment, comparison and branch.

Scripting

For a more dynamic stage that leans even more heavily on reification, let’s look at the scripting as implemented for the put verb. It’s again just a small subset of what Miller offers: we have basic arithmetic, static column references and a ternary operator. No dynamic indexing, control flow, or built-in functions for now.

The implementation makes use of a small, additional reification helper - @Derived. When the function is reified, the builder method referred to by this annotation will be invoked, and its result will be stored the same way an @Constant value is stored. This lets us structure the tree of reified functions in a top-down way. Without @Derived, we’d have to build up the computations in a bottom-up fashion and then pass them to the final reified function that invokes them.

The builder method specified in @Derived("builder") can accept any subset of @Constant arguments, matched by name, that the parent function receives. This method is usually itself marked @Reified, or is a plain static function that assembles loop operations and calls a @ReifiedLoop generator function. I find that structuring the reification tree in a top-down way makes it much more composable.

final class Script implements PipelineStage {
  @Override
  public BatchOp compileBatchOp(
      BatchLayout in, BatchLayout out, MemorySegment data, MemorySegment scratch) {
    final var evaluator = Script_Reified.batchEvaluator(this, out, data);
    return evaluator::apply;
  }

  @Reified
  static int batchEvaluator(
      @Constant Script step,
      @Constant BatchLayout out,
      @Constant MemorySegment data,
      @Derived("buildStatementLoop") StatementLoop statements,
      int control) {
    final var rowCount = controlRowCount(control);
    for (var i = 0; i < rowCount; i++) {
      statements.invoke(i);
    }
    return control(controlOutputComplete(control), rowCount);
  }

  @Reified(extending = StatementEvaluator.class)
  static void statementEvaluator(
      @Constant BatchLayout out,
      @Constant MemorySegment data,
      @Constant Scripting.Expr e,
      @Derived("compileEvaluator") Evaluator eval,
      @Constant MemorySegment dataRef,
      int rowIx) {
    dataRef.set(ValueLayout.JAVA_INT, (long) rowIx * Integer.BYTES, eval.apply(rowIx));
  }

  @ReifiedLoop(step = StatementEvaluator.class)
  static void statementLoop(int rowIx) {}
}

compileEvaluator, the key @Derived builder method, recursively matches on the expression AST. Each branch invokes another generated reification factory, so every expression node becomes a reified Evaluator.

Here are a few interesting cases:

@Reified(extending = Evaluator.class)
static int constEvaluator(@Constant int value, int rowIx) {
  return value;
}

@Reified(extending = Evaluator.class)
static int fetchEvaluator(@Constant MemorySegment column, int rowIx) {
  return column.get(ValueLayout.JAVA_INT, (long) rowIx * Integer.BYTES);
}

@Reified(extending = Evaluator.class)
static int binaryEvaluator(
    @Constant Evaluator left,
    @Constant Evaluator right,
    @Constant String operator,
    int rowIx) {
  return switch (operator) {
    case "*" -> left.apply(rowIx) * right.apply(rowIx);
    // ...
    default -> throw new UnsupportedOperationException(operator);
  };
}

@Reified(extending = Evaluator.class)
static int ternaryEvaluator(
    @Constant Evaluator condition,
    @Constant Evaluator ifTrue,
    @Constant Evaluator ifFalse,
    int rowIx) {
  return condition.apply(rowIx) != 0 ? ifTrue.apply(rowIx) : ifFalse.apply(rowIx);
}

static Evaluator compileEvaluator(
    BatchLayout out, MemorySegment data, Scripting.Expr e) {
  return switch (e.kind()) {
    case Scripting.Kind.NUMBER ->
      Script_Reified.constEvaluator((int) e.number());

    case Scripting.Kind.COLUMN ->
      Script_Reified.fetchEvaluator(out.resolveColumn(data, e.text()));

    case Scripting.Kind.BINARY -> {
      final var left = compileEvaluator(out, data, e.operands().get(0));
      final var right = compileEvaluator(out, data, e.operands().get(1));
      yield Script_Reified.binaryEvaluator(left, right, e.text());
    }

    case Scripting.Kind.TERNARY -> {
      final var condition = compileEvaluator(out, data, e.operands().get(0));
      final var ifTrue = compileEvaluator(out, data, e.operands().get(1));
      final var ifFalse = compileEvaluator(out, data, e.operands().get(2));
      yield Script_Reified.ternaryEvaluator(condition, ifTrue, ifFalse);
    }

    // ...
  };
}

This is a little excessive in the number of @Reified functions used, and one could use lambdas for leaf computations, but I didn’t want to rely on Graal’s inlining heuristics that are bound to exhaust at some point.

Each assignment becomes its own reified statement: evaluate the RHS, then write the result to the destination column. We utterly rely on Graal to optimize this stage. Autovectorization turns our naive scalar code that operates on individual rows into a faster SIMD implementation. The compiler also deduplicates our memory accesses by caching repeated reads and only letting the last write in sequence to the same location survive.

A script like $B = $A * 16; $C = $B * 3 is therefore two vectorized statement evaluators that avoid round-tripping through memory:

vmovdqu32 zmm0, zmmword ptr [r11 + r13*4] ; read 16 values from $A
vpslld    zmm1, zmm0, 5                   ; $A * 32
vpslld    zmm0, zmm0, 4                   ; $B = $A * 16
vpaddd    zmm1, zmm0, zmm1                ; $C
vmovdqu32 zmmword ptr [rdx + r13*4], zmm0 ; store $B
vmovdqu32 zmmword ptr [r10 + r13*4], zmm1 ; store $C

The $B read in the second statement disappears completely: Graal strength-reduces $C = $B * 3 into ($A << 4) + ($A << 5). It’s faster to compute $B * 3 this way than to compute ($B << 1) + $B. By rebasing the computation on $A, it produces shorter dependency chains and more instruction-level parallelism. Overall, this example comes closest to the stated goal introduced at the start of the post. We get to write straightforward code that, as long as it’s reified appropriately, generates optimal code by taking advantage of the existing JIT compiler.

There is one small issue that I’m hoping gets addressed — oracle/graal#13958. When Graal improves aliasing analysis on MemorySegments, this should eliminate some column reloads that follow earlier stores to another column, but this only shows up in longer scripts.

Reductions

stats1 is a verb that reduces incoming values and only outputs the results once the upstream stage fully completes. Here’s a quick example of it:

$ mlr --tsv seqgen -f A --start 10 --stop 30 then stats1 -f A -a min,sum
A_min  A_sum
10     420

We implement it with reified loops over the requested fields and accumulators. Each accumulator occupies one slot of an Object[]. StatGrid is the field × accumulator grid representing the configuration parameters.

@Reified(extending = BatchOp.class)
static int kernel(
    @Constant StatGrid grid,
    @Constant BatchLayout in,
    @Constant BatchLayout out,
    @Constant MemorySegment data,
    @Constant MemorySegment scratch,

    // each of these generates a @ReifiedLoop over a varying number of steps,
    // all depending on StatGrid configuration
    @Derived("buildInit") Stats1_Reified.InitLoop init,
    @Derived("buildChunk") Stats1_Reified.ChunkLoop chunk,
    @Derived("buildMaskedChunk") Stats1_Reified.MaskedChunkLoop maskedChunk,
    @Derived("buildFlush") Stats1_Reified.FlushLoop flush,

    int control) {
  final boolean eof = controlOutputComplete(control);
  final int rowCount = controlRowCount(control);
  final var species = IntVector.SPECIES_PREFERRED;
  final int lanes = species.length();
  final int fullCount = species.loopBound(rowCount);
  final Object[] state = new Object[grid.slots()];

  init.invoke(state);
  for (int r = 0; r < fullCount; r += lanes) {

    chunk.invoke(state, r);
  }

  if (fullCount < rowCount) {
    maskedChunk.invoke(state, fullCount, species.indexInRange(fullCount, rowCount));
  }

  flush.invoke(state, control);
  return control(eof, eof ? 1 : 0);
}

// one slot of the chunk.invoke() loop seen above
@Reified
static void chunkStep(
    @Constant int slot,
    @Constant VectorOperators.Associative reducer,
    @Constant MemorySegment input,
    Object[] state,
    int row) {
  final var values =
      IntVector.fromMemorySegment(
          IntVector.SPECIES_PREFERRED,
          input,
          (long) row * Integer.BYTES,
          ByteOrder.nativeOrder());
  state[slot] = ((IntVector) state[slot]).lanewise(reducer, values);
}

Testing it on input fed from seqgen, stats1 --field A --accumulators min,max compiles to:

mov         r14d, edx
add         r14d, edi
vmovdqu32   zmm2, zmmword ptr [rbp + r14*4] ; load 16 ints from column A
vpmaxsd     zmm1, zmm1, zmm2                ; max reduction
vpminsd     zmm0, zmm0, zmm2                ; min reduction
lea         edx, [rdx + 0x10]
cmp         r13d, edx
jg          loop                            ; back to top

One load feeds both reductions and there are no remnants of Java objects or the heap to be seen anywhere in the hot loop.