Changes to Code

1

There is a moderate amount of dissatisfaction with Code[T]. I think there are 2 essential issues:

  • The need to split up large functions creates a lot of complexity in the emitter. The solution is to hide function splitting inside Code. This doesn’t require an interface change but is involved.
  • Code[T] looks like a value of type T, but isn’t. Sometimes you want to be able to refer to values.

I also think code should be emitted in a stateful way. This is the classical design. Therefore, I propose the following:

There are actually two types of values in JVM worth distinguishing: simple or primitive values like constants, variable reads, etc. and large, complex compound (but functional) computations that produce a value. The latter thing you probably don’t want to duplicate in the code, and there’s no point in memoizing something like like the former. This, Code[T] becomes SimpleValue[T] and Value[T] for non-statement code.

For statement code, I want something like cb: CodeBuilder which builds up code via a series of emit statements:

cb.emitGoto(L)
cb.emitTarget(L)
cb.emitInvoke(...)
cb.emitIfX(...)
cb.emit(v)

where v: (Simple)Value[T] emits the code for the value computation. CodeBuilder would also force the existence of a context that could support things like memoization. Memoization would return, and be idempotent, on SimpleValues.

CodeBuilder feels like a significant change, so I propose as an intermediate step Code[T] becomes Code (no type parameter) for statements (Unit value), and we continue to build up code blocks functionally. So the above code looks like:

Code(
  Code.goto(L),
  Code.target(L),
  Code.invoke(L),
  ...)

In this design with Code, EmitTriplet becomes:

case class EmitTriplet[T](
  val setup: Code,
  val m: Value[Bool],
  val v: Value[T]

Eventually, with CodeBuilder in the context, it just becomes EmitValue with m and v.

With a backend that automatically splits large functions, I think this will allow us to clean up higher levels of the code.

edit: I would appreciate feedback from people in the Code trenches on this.

2

There are actually two types of values in JVM worth distinguishing: simple or primitive values like constants, variable reads, etc. and large, complex compound (but functional) computations that produce a value. The latter thing you probably don’t want to duplicate in the code, and there’s no point in memoizing something like like the former. This, Code[T] becomes SimpleValue[T] and Value[T] for non-statement code.

I started making this change and ran into a snag, before switching to the stream emitter. I agree completely.

With a backend that automatically splits large functions, I think this will allow us to clean up higher levels of the code.

As I see it, the difficulty in automatically splitting large functions is knowing what local variables are actually local to a sub-method, and which need to be lifted to class fields (or passed as function arguments). Unless we make everything be class fields, and don’t use locals at all, but I suspect we would hit class size issues doing that.

As you suggest, I think Code needs to handle binding values to locals. To support method splitting, I thought Code could abstract over whether values are bound to locals or class fields. Also, I think Code would either need to let the value binding interface (memoization, as you called it) express the scope over which the binding should be valid, or to do a liveness analysis pass before method splitting.

Overall, I was imagining Code would need to become more of a low-level IR, that we could inspect after Emit to make decisions like where to split into methods and whether to bind values as locals or class fields.

3

I think this addresses my essential problems with Code[T]. The lack of a distinction currently between SimpleValue and Value is annoying but not insurmountable (you just have to be careful to keep binding things to names). I think handling locals vs fields correctly is the higher priority thing. I think we are leaving performance on the table by not emitting new methods. Now that I’m emitting new methods for ndarrays, they perform much better, but it’s possible to get an error about referencing a local in the wrong place if you write a complicated expression.

4

The lack of a distinction currently between SimpleValue and Value is annoying but not insurmountable

I went back and forth a lot on this. The problem is that the JVM is an annoying mixture between a (pure) stack machine and a register machine (locals). Value corresponds to computations in stack land, and SimpleValue corresponds to the register machine picture. One option is to eschew the stack machine and only put things in locals. I’m very tempted by this, it would simplify things, but the main problem I see is that it breaks missingness optimizations in EmitTriplet that uses branch targets (see CodeConditional) that are broken by going to/from local variable. It might be that this doesn’t actually matter. It also increases code size (~3x?) which may exacerbate our code size issues.

Another option, more work, but we could switch to a pure register machine, and have Code module both break up large functions and optimize register machine instructions to take advantage of the stack if possible.

If we’re concerned about the missingness optimization in EmitTriplet, we can also m: Code[Bool] into an explicit conditional with a pair of labels like CodeConditional.

It might be worth benchmarking the current code with code conditional optimization disabled.

5

John, a few more responses:

I think handling locals vs fields correctly is the higher priority thing.

An option before we split is to always use fields. I feel like that isn’t going to perform well, but I don’t think I have a good intuition for how the bytecode impacts performance. I think we can do an experiment where we make newLocal just return a new field.

I think we are leaving performance on the table by not emitting new methods.

Did you feel like the interface made this harder than it needed to be?

it’s possible to get an error about referencing a local in the wrong place if you write a complicated expression.

Where the local is in one function but you reference it in another? What’s the wrong place? In the functional code model it is hard to report that at the use because you don’t know where a code expression is going to end up. One option is to put the function in the code. This makes primitive expressions (e.g. constants) more complicated because you need to say what function the constant is going to go into when you construct it. With CodeBuilder could check on emit, which would be immediate for writes but deferred for complex values.

I’ll think about this. I am generally in strong favor of runtime checking (asserts) to catch bugs.

6

As I see it, the difficulty in automatically splitting large functions is knowing what local variables are actually local to a sub-method

You want to a liveness dataflow analysis on the function local variables. This is classical, any compiler textbook will explain this. Essentially, this tells you, for each edge in control flow graph, which values are “live”, that is, in this context, which values need to be transmitted if you cut this edge.

It seems harder to me to decide where to cut. Maybe you can think about some kind of min-cut/max-flow, but you want to balance that with an even cut (or cut into k even pieces). You also need to compute edge weights. This has a few components: the cost of converting the live variables at the edge to fields (or as arguments if you don’t need to transmit changes back), the cost of introducing the cut (function call) both of which require estimating the execution frequency of different parts of the code. If you cut at multiple points, you also need a switch on either end to branch to the correct point, so the real cost function is not linear. Doable, but feels like it will require some tuning.

Overall, I was imagining Code would need to become more of a low-level IR

Yes, I agree. More precisely, the Code interface will build a lower-level IR which will then be lowered to bytecode by splitting.

7

I started making this change and ran into a snag

What was the snag?

8

Did you feel like the interface made this harder than it needed to be?

Emitting a new method isn’t very hard, thought that was fine. All I had to do was:

val eVti = typeToTypeInfo(TVoid)
val innerMethod = mb.fb.newMethod(eVti)
innerMethod.emit(loops)
innerMethod.invoke()

Instead of just returning loops. The only problem is that by doing this I’m getting errors about locals.

What’s the wrong place?

I have a Code expression that contains a reference to a local in the current method. When I emit a new method whose body is that Code expression, I now can’t use that local.

9

I think this is improper use of Code – you should pass that Code in as an argument. This is hard to do when you’re trying to generate code for, say, NDArrayMap2 or something, since you could have external variable references (to local variables) referenced inside the map IR. Is that where the problem lies?

10

It isn’t super easy to even tell which local is being referenced, but I suspect it’s coming from something like you describe.