Fedor Indutny's Blog

Allocating numbers

JIT #

This is the second blog post in the series about JIT compiling. The previous post was an introduction into the Just-In-Time code generation and, in particular, jit.js usage. If you haven't read it yet - I recommend you to familiarize yourself with it first.

Objectives #

Previously, we created a JIT compiler, supporting a very limited subset of JavaScript: integer numbers, math binary operators (+, -, *, /), and - unary operator. This time, we will extend it by adding floating point number support, and, to make the process funnier and to spice things up, we will allocate and store these numbers in the heap.

Though, because we are doing things one step at a time, our heap won't have Garbage Collection, and will live inside fixed sized memory chunk (say "yay" to simplicity!).

Stubs #

Knowing what we aim to do, we can now set up internal structures for these features. Essentially, what we'll need is a memory allocation procedure, that generates and returns memory addresses suitable for our goals.

This allocation code could be generated for every AST node using series of inlined assembly instructions, which works great and, more importantly, is incredibly fast for concise operations. But due to the relatively big code's size of this procedure, the resulting machine code output may become too big to be fit entirely into the CPU's cache, causing potential performance problems to the whole system.

Generally, this is considered a bad practice. A better approach would be parameterizing such code blocks into shared procedures called stubs (I picked that naming from v8's source and, perhaps, it is how these things are named in other VMs too). For even better optimization these procedures could be lazily compiled, i.e. we should not compile those ones that are not used by generated code. This technique is good for both compilation time and executable code size (and therefore CPU caches too).

Fortunately, jit.js lets you generate stubs easily:

var stubs = jit.stubs();

stubs.define('Allocate', function() {
  // Our code here
  // ....

  // Returning back to caller
  this.Return();
});

Simple, isn't it? Now, to use it in our JIT compiler we'll need to pass it in an options argument:

jit.compile(function() {
  // Compiler code generation happens in this context

  // Explanation:
  // Read address of 'Allocate' stub into 'rax' register and
  // call it.
  this.stub('rax', 'Allocate');

  this.Return();
}, { stubs: stubs });

As mentioned above, only stubs that were used during compilation process will actually be generated and reused between all callers.

Heap #

With this knowledge, we can proceed to the memory allocation phase. But first, lets take a short look at the structure and organization of the heap.

The heap is the place where JavaScript (and many other) VMs create and store objects (usually, ones that can't be fit into CPU registers). Some heap objects may contain references to other objects (in other words, can reference them). All live objects and their references create a directed graph, starting at so called roots (which are usually global variables and pointers on stack).

Although, it is usually used in VMs with JIT compilation, Garbage Collection is not required for the Heap. Indeed, many VMs and languages choose to use unmanaged memory instead (C/C++ as a banal example). In such cases you (as the language user) will generally need to explicitly free unused resources to not run out of the memory.

But for obvious reasons, the JavaScript subset compiler that we're implementing, should support both managed memory and Garbage Collection (which will be implemented later).

There are tons of books that may give you an advanced introduction into the heap allocation and garbage collection (my recommendation is The Garbage Collection Handbook), and considerably many ways to allocate and collect memory in the heap.

Usually, you will need to choose between the allocation speed and memory fragmentation. But, since we are not covering this very deeply, I would recommend to stick with the method called "bump allocation" for now.

Bump allocation #

Fixed-page bump allocation works in a following way.

  1. Take the memory chunk of fixed size (a page)
  2. Give away consequent slices of it as a return value of the allocation procedure.
  3. When running low on memory, perform the Garbage Collection and free all unused space, by either compacting live objects or evacuating them to the new memory chunk (replacing references to live objects in both cases).

In terms of jit.js and stubs API, this procedure may look as following:

// Create fixed size memory chunk
var page = new Buffer(1024);

// Set-up pointers to page start and page end
var offset = jit.ptr(page);
var end = jit.ptr(page, page.length);

stubs.define('Alloc', function() {

  // Save 'rbx' and 'rcx' registers
  this.spill(['rbx', 'rcx'], function() {
    // Load `offset`
    //
    // NOTE: We'll use pointer to `offset` variable, to be able to update
    // it below
    this.mov('rax', this.ptr(offset));
    this.mov('rax', ['rax']);

    // Load end
    //
    // NOTE: Same applies to end, though, we're not updating it right now
    this.mov('rbx', this.ptr(end));
    this.mov('rbx', ['rbx']);

    // Calculate new `offset`
    this.mov('rcx', 'rax');

    // We'll assume that all allocations are 16 bytes = two 64bit pointers
    this.add('rcx', 16);

    // Check if we won't overflow our fixed size buffer
    this.cmp('rcx', 'rbx');

    // this.j() performs conditional jump to the specified label.
    // 'g' stands for 'greater'
    // 'overflow' is a label name, bound below
    this.j('g', 'overflow');

    // Ok, we're good to go, update offset
    this.mov('rbx', this.ptr(offset));
    this.mov(['rbx'], 'rcx');

    // The first 64bit pointer is reserved for 'tag',
    // the second one is a `double` value
    this.mov(['rax'], 1);

    // Return 'rax'
    this.Return();

    // Overflowed :(
    this.bind('overflow')

    // Invoke javascript function!
    // NOTE: This is really funky stuff, but I'm not going to dive deep
    // into it right now
    this.runtime(function() {
      console.log('GC is needed, but not implemented');
    });

    // Crash
    this.int3();

    this.Return();
  });
});

That's it! Not totally straightforward, but not really complicated either!

This procedure will give away consequent slices of the page, and even tag them! (I'll cover tagging in one of the next posts. Basically, they're used to distinguish different kinds of heap objects).

Few things to note here:

  1. jit.ptr(buf, offset) returns a Buffer, containing a pointer to the given buf with offset added to it.
  2. this.spill() is a routine for saving and restoring registers to/from the memory (this process is usually called spilling). It takes list of the registers and the closure. These registers will be saved before entering the closure, and restored right after leaving it. NOTE: The restore code will be generated before each this.Return() too.
  3. this.mov(['rbx'], 'rcx') - stores rcx register into the memory location, pointed by the value of rbx register. NOTE: you can also specify an offset here: this.mov(['rbx', 8], 'rcx').
  4. jit.js supports branching primitives: this.cmp(a, b), this.j(condition, labelName), this.j(labelName), this.bind(labelName).

Floating point #

Now as we have a presumably working allocation procedure, let's recall what should be stored inside of this heap chunks. In the allocation procedure, we create chunks with the 8 byte tag value, and the 8 byte contents. This is enough to store double (as C type) floating point numbers.

There are plenty of assembly instructions to load/store/work with such numbers. But note that to work with them - you'll need to store them in the different register set: xmm0, xmm1, ... xmm15. Although, 64-bit floating numbers could be stored in the general purpose registers: rax, rbx, ... Performing math operations is possible only with a xmm register set. Here are some instructions, that are present in jit.js and should be useful for our compiler:

  1. movq('xmm', 'gp') or movq('gp', 'xmm') to move 64bits from the general purpose register (or memory pointed by it) to xmm, or the other way around.
  2. movsd('xmm', 'xmm') to move the value from one xmm to another.
  3. addsd, mulsd, subsd, divsd - addition, multiplication, subtraction, division.
  4. cvtsi2sd('xmm', 'gp'), cvts2si('gp', 'xmm') - converts integer into double, and double into integer, respectively.
  5. roundsd('mode', 'xmm', 'xmm') - round the src register using specified mode (which is one of: nearest, down, up, zero) and place the result into the dst register.

Using this sacred knowledge we can patch our existing code to make it work with the floating point numbers (yeah, we will remove the integer support for now):

// Compile
var fn = jit.compile(function() {
  // This will generate default entry boilerplate
  this.Proc(function() {
    visit.call(this, ast);

    // The result should be in 'rax' at this point
    //
    // This will generate default exit boilerplate
    this.Return();
  });
}, { stubs: stubs });

// Execute
console.log(fn());

function visit(ast) {
  if (ast.type === 'Program')
    visitProgram.call(this, ast);
  else if (ast.type === 'Literal')
    visitLiteral.call(this, ast);
  else if (ast.type === 'UnaryExpression')
    visitUnary.call(this, ast);
  else if (ast.type === 'BinaryExpression')
    visitBinary.call(this, ast);
  else
    throw new Error('Unknown ast node: ' + ast.type);
}

function visitProgram(ast) {
  assert.equal(ast.body.length,
               1,
               'Only one statement programs are supported');
  assert.equal(ast.body[0].type, 'ExpressionStatement');

  // We've a pointer in 'rax', convert it to integer
  visit.call(this, ast.body[0].expression);

  // Get floating point number out of heap number
  this.movq('xmm1', ['rax', 8]);

  // Round it towards zero
  this.roundsd('zero', 'xmm1', 'xmm1');

  // Convert double to integer
  this.cvtsd2si('rax', 'xmm1');
}

function visitLiteral(ast) {
  assert.equal(typeof ast.value, 'number');

  // Allocate new heap number
  this.stub('rax', 'Alloc');

  // Save 'rbx' register
  this.spill('rbx', function() {
    this.loadDouble('rbx', ast.value);
    this.mov(['rax', 8], 'rbx');
  });
}

function visitBinary(ast) {
  // Preserve 'rbx' after leaving the AST node
  this.spill('rbx', function() {
    // Visit right side of expresion
    visit.call(this, ast.right);

    // Move it to 'rbx'
    this.mov('rbx', 'rax');

    // Visit left side of expression (the result is in 'rax')
    visit.call(this, ast.left);

    //
    // So, to conclude, we've left side in 'rax' and right in 'rbx'
    //

    // Let's load their double values
    this.movq('xmm1', ['rax', 8]);
    this.movq('xmm2', ['rbx', 8]);

    // Execute binary operation
    if (ast.operator === '+') {
      this.addsd('xmm1', 'xmm2');
    } else if (ast.operator === '-') {
      this.subsd('xmm1', 'xmm2');
    } else if (ast.operator === '*') {
      this.mulsd('xmm1', 'xmm2');
    } else if (ast.operator === '/') {
      this.divsd('xmm1', 'xmm2');
    } else {
      throw new Error('Unsupported binary operator: ' + ast.operator);
    }

    // Allocate new number, and put value in it
    this.stub('rax', 'Alloc');
    this.movq(['rax', 8], 'xmm1');
  });
}

function visitUnary(ast) {
  if (ast.operator === '-') {
    // Negate argument by emulating binary expression
    visit.call(this, {
      type: 'BinaryExpression',
      operator: '*',
      left: ast.argument,
      right: { type: 'Literal', value: -1 }
    })
  } else {
    throw new Error('Unsupported unary operator: ' + ast.operator);
  }
}

To be continued #

So, that's all I have to say to you for now. On a more social theme, you may want subscribe to my twitter or watch my blog on github. Don't miss the next post!