Proposal: low-level IR

1

This is a follow up to Changes to Code on the evolution of my thinking with respect to Code[_].

One of the challenges we’re facing in the backend is handling JVM limits (e.g. on method bytecode size). The current wrapping logic mostly works, although it complicates the code, but more importantly, it is an issue that cuts across the entire code generator and new code generation (like ndarrays and emit streams) needs to engage with it but don’t.

Thus, we need to decouple method splitting from the code generator. My proposal is what I’m calling the low-level IR (lir). In this model, Code[_] will be a wrapper that will build lir.

lir has:

  • Classx (x to not conflict with the Java builtin Class)
  • Method
  • (basic) Block
  • X (for eXpression): StmtX, ControlX, and ValueX.

A Classx is a set of methods and fields.

A Method is a set of locals and basic blocks, with a designated entry block.

A basic block is a list of StmtX. It is straight line code. The last statement in a block must be a ControlX (a switch, if, goto or return).

Non-control statements include loading and storing fields, locals and invocations that return void. Currently, function calls that return values and new instance are ValueX, although they might perform mutation. I think we can revisit that in the future and make them control that writes a local.

Note, the lir has no explicit representation of the JVM stack. ValueX can be nested, so (add (sub ...) (neg ...) will be compiled to bytecode where the intermediates are put on the stack, but that isn’t visible at the lir level.

What’s more, the JVM stack is not used between basic blocks. This makes splitting at basic blocks simpler. It conflicts with joinpoint’s use of the stack and ParameterPack.push. Those both blocks for this moving forward.

Then plan is for Emit to not wrap, and then post-process the lir to satisfy JVM limits. I have a splitter underway that looks like:

  • move all locals to fields, including inputs to control statements
  • split large statements
  • move body of basic blocks into separate a method(s)

Then each basic block turns into a function call and a control statement (if, branch, etc.) For this simple strategy to work, the control flow of the original method must fit into a single method. I’m not 100% sure that is the case. (@tpoterba?) This will let us simplify the exiting emitter code and give us a place to drop in a smarter splitter (e.g. that minimizes spilling to fields).

How does Code[_] change in this picture? My model of Code[T] is a computation and a value of type T that can be used after that computation is run. The value is represented by a ValueX. So Code looks like:

class Code[+T](
  val start: lir.Block,
  val end: lir.Block,
  val v: lir.ValueX)

start is a block that represents the beginning of the computation (jump to start if you want to execute it). end represents where the computation ends up (append things there if you want to do something after, possibly using v). v is the value you can use after end.

There are two exceptional cases to consider: Code[Unit] and non-returning computations, e.g. what to do about L.goto.

A Code[Unit] has v = null since there is no value.

In this setup, you can’t write Code(L.goto, ...) because a basic block must have one and exactly one control statement at the end.

I think we can represent non-returning computations with Code[Nothing]. For a Code[Nothing], end = v = null, since there is no end, nor a value. Then you can’t say Code(L.goto, ...) but you can say b.mux(s, t): Code[Nothing] where b: Code[Boolean] and s, t: Code[Nothing].

I look into peeling off the Nothing business in the current code. Until then, non-returning computations should be Code[Unit] and we won’t statically check putting values after non-returning computations.

lir.X has a parent pointer, so values that are reused will immediately throw errors. For generating control instructions (which basically involves appending to a basic block), I’ll make Code[_] consumers clear out the producer so they can’t be reused (e.g. move semantics).

I have a draft for most of this and I’ll aim for a PR beginning of next week.

1 Like
2

I think we can represent non-returning computations with Code[Nothing] .

Hmm, fleshing this out, I do wonder if this is more trouble than it is worth. At the very least, we’ll need an additional construct which does something like “emit this Code[Nothing] and then resume at this block over here”, Code[T](start: Code[Nothing], continue: Code[T]) where continue emits a label referenced by start and then it doesn’t really look any different than non-terminating code having type Code[Unit].

3

I like this a lot. I agree completely on the need for an IR below Code.

I had some thoughts over the weekend that relate to the question of how to represent non-returning computations, and how other things like CodeConditional, EmitTriplet, and Stream relate to Code.

I’ll use CodeConditional as a motivating example. I’d argue CodeConditional doesn’t really fit into your Code[T] picture, because it doesn’t have a single terminating block, and doesn’t have a final value. Of course it can be converted to a Code[Boolean], but we want to do that as late as possible, if at all. You could reasonably argue that CodeConditional should be an implementation detail of Code, not exposed in the interface, but Stream, EmitTriplet, and non-returning computations definitely do need to be exposed. They are currently working as standalone classes, but one motivation for unifying them under a shared interface is that Stream and EmitTriplet both need to be allowed element types of streams, and codifying the common operations would simplify the stream emitter, which I’ll say more about below.

The main commonality I see between those is that we can define an operation "prepend with setup: Code[Unit]" for all of them. Lets call these “computations”, where Code[T] is the special kind of computation that returns a value of type T. Concretely, I propose adding a superclass to Code[_]:

  • CodeComputation
    • Code[T]
    • CodeControl
    • CodeConditional
    • EmitTriplet (or even EmitTriplet[T <: CodeComputation])
    • Stream[T <: CodeComputation]

A CodeComputation should support the operations

object Code {
  def apply[T <: CodeComputation](c1: Code[Unit], c2: Code[Unit], c3: T): T
  // etc...

  def mux[T <: CodeComputation](cond: Code[Boolean], then: T, else: T): T
}

I think in your proposed model, a CodeControl is like a Code[T] but with only the start block, to which a Code[Unit] can be prepended. The basic example is a single jump instruction.

I believe these operations are sufficient to allow me to specialize the element type of streams to CodeComputation. Right now the type is completely unconstrained, to allow for Stream[Stream[T]] and Stream[EmitTriplet]. Supporting that has added complexity, mostly in the form of CPS versions of operators like map (mapCPS(stream: Stream[A], f: (A, B => Code[Ctrl]) => Code[Ctrl]): Stream[B] instead of map(stream: Stream[A], f: A => B): Stream[B])). I believe specializing the element type to CodeComputation, and using the generic "prepend with Code[Unit]", I could get rid of all the CPS, which would be a significant simplification.

I think this would also allow a uniform interface that would simplify some common patterns. For example, the common

EmitTriplet(Code(extraSetup, et.setup), et.m, et.v)

could be written

Code(extraSetup, et)
4

A Code[Nothing] needs to be more general than straightline code that ends in a jump. In particular “resume at this block over here” need not be well defined, because Code[Nothing] can have many control flow paths ending in jumps to several different labels. I think the semantics needs to be: a Code[Nothing] is the root of a CFG where all blocks end in jumps to previously defined labels (either defined internally to the Code[Nothing] or externally).

5

I realize I owe you a response on this thread. I’m still cogitating.

6

I realized there’s an important part of how I think about Code[Nothing] (or Code[Ctrl] in the current codebase) that I haven’t explained. I think the discussion will be easier if I call these continuations, to keep them more distinct from Code[T], let’s call the type Cont.

There’s a duality between continuations Cont and computations Code[T]. This duality is tighter if we consider continuations that take parameters, Cont[T], like Milo’s joinpoints, so I’ll start by considering those, but afterwards only consider Cont = Cont[Unit].

  • A Code[T] represents a computation that goes from some indeterminate point in the past to now, producing a value of type T. A Code[T] should only be used once. If it needs to be used multiple times, its result can be stored in a local variable, giving a Value[T], which can then be used arbitrarily.
  • A Cont[T] represents a computation that goes from now to some indeterminate point in the future, consuming a value of type T. A Cont[T] should only be used once. If it needs to be used multiple times, its body can be assigned a label, giving a Label[T], which can then be jumped to arbitrarily.

We can implement Cont[T] using Cont[Unit] by passing the parameter in an externally declared variable, so we need only consider Cont = Cont[Unit] and Label = Label[Unit].

Integrating this into your lir proposal, Label would just be a label, and Cont would be

class Cont {
  val body: lir.Block
}

except that the body has not yet been given a label, and instructions may be added at the front, e.g. with Code.apply(c: Code[Unit], k: Cont): Cont. Think of it as the root of a CFG we are building bottom up.

The interface for working with Cont and Label would be completely parallel to Code and Value: a Cont can be memoized, which assigns it a label, adds its blocks to the underlying lir Method, and returns the Label, unless the Cont was already just a jump to a label, in which case the already existing label is returned.

As an example, in this model CodeConditional could be

trait CodeConditional { self =>
  def emit(m: Method, ctrue: Cont, cfalse: Cont): Cont

  def &&(rhs: CodeConditional): CodeConditional = new CodeConditional {
    def emit(m: Method, ctrue: Cont, cfalse: Cont): Cont = {
      val lfalse = m.memoizeCont(cfalse)
      self.emit(m,
        rhs.emit(m, ctrue, lfalse),
        lfalse)
    }
  }

  // etc.
}

Of course, we could get away with only having Label and not Cont, at the expense of creating a lot of

label1:
  instruction
  goto label2
label2:
  instruction
  goto label3
...

and

label1:
  goto label2

If the JIT runs on the method, those will get cleaned up, and maybe you prefer that approach. But I think streams in particular would create a lot of those redundant labels, and we don’t completely understand when or if those methods will get compiled by the JIT. Also, I think Cont would only be used in CodeConditional and Stream, so it would be invisible to most users of Code.